feat(ck_tile/fmha): add gfx1100 batch_prefill_gfx11 pipeline policy

gfx1100 had no batch_prefill pipeline: the arch supported qr, qr_pagedkv,
qr_nwarp_sshuffle, appendkv and bwd only. qr_async cannot be reused
because gfx11 has no parent async dwordx2 K copy, and the fallback gemm1
(ARegBSmem) cannot consume a gfx11 WMMA P tile.

Add BlockFmhaBatchPrefillPipelineQRKSVSAsyncGfx11Policy:

  - synchronous K staging (kUseSyncKLoad) instead of async copy
  - vec4 K packing on N0=32 / K0=32 linear tiles
  - Independent-V gemm1 (ARegBReg plus a tiled V LDS layout)
  - P permuted from WMMA C layout to A layout before gemm1

Enabled surface is deliberately narrow: block size 256, tile 128x32x32,
hdim 128, linear page KV and no dropout. All other tiles, head dims,
dtypes and architectures stay on the existing pipeline, unchanged.
Codegen emits exactly that set (192 instances), validation rejects
anything outside it with a named error, and the pipeline carries a
device-side static_assert so the two gates cannot drift apart.

The tile engine harness needed three fixes to exercise the path: the
ctypes runner zeroed the paged K/V buffers instead of copying the host
pages, it hardcoded window_size_right=-1 so causal masks did not mask,
and the Python side passed dense BHSD without the linear page pack or
the Q/O transpose. fmha_benchmark.py now also exits non-zero when
kernels fail to build or verification fails; previously both printed the
failure and exited 0.

Author: JinHe-102

projects/composablekernel/dispatcher/bindings/ctypes/fmha_ctypes_lib.cpp

@@ -1552,8 +1552,16 @@ int fmha_dispatcher_run_batch_prefill(const void* q_host,
     }
 
     HIP_CHECK(hipMemcpy(q_dev, q_host, q_bytes, hipMemcpyHostToDevice));
-    HIP_CHECK(hipMemset(k_dev, 0, kv_page_bytes));
-    HIP_CHECK(hipMemset(v_dev, 0, kv_page_bytes));
+    {
+        const int64_t k_src_bytes =
+            static_cast<int64_t>(total_pages) * nhead_k * page_block_size * hdim_q * in_bytes;
+        const int64_t v_src_bytes =
+            static_cast<int64_t>(total_pages) * nhead_k * page_block_size * hdim_v * in_bytes;
+        HIP_CHECK(hipMemset(k_dev, 0, kv_page_bytes));
+        HIP_CHECK(hipMemset(v_dev, 0, kv_page_bytes));
+        HIP_CHECK(hipMemcpy(k_dev, k_host, k_src_bytes, hipMemcpyHostToDevice));
+        HIP_CHECK(hipMemcpy(v_dev, v_host, v_src_bytes, hipMemcpyHostToDevice));
+    }
     HIP_CHECK(hipMemset(o_dev, 0, o_bytes));
 
     args.q_ptr           = q_dev;
@@ -1613,14 +1621,15 @@ int fmha_dispatcher_run_batch_prefill(const void* q_host,
     args.batch_stride_randval = 0;
     args.batch_stride_lse     = static_cast<int64_t>(nhead_q) * seqlen_q;
     args.batch_stride_o       = 0;
-    args.window_size_left     = -1;
-    args.window_size_right    = -1;
-    args.sink_size            = 0;
-    args.mask_type            = mask_type_int;
-    args.p_drop               = has_dropout ? 0.2f : 0.0f;
-    args.s_randval            = false;
-    args.drop_seed_offset     = has_dropout ? std::make_pair(uint64_t(1), uint64_t(0))
-                                            : std::make_pair(uint64_t(0), uint64_t(0));
+    // right=-1 is a full row, so causal (top_left=1 / bottom_right=2) needs right=0.
+    args.window_size_left  = -1;
+    args.window_size_right = (mask_type_int == 1 || mask_type_int == 2) ? 0 : -1;
+    args.sink_size         = 0;
+    args.mask_type         = mask_type_int;
+    args.p_drop            = has_dropout ? 0.2f : 0.0f;
+    args.s_randval         = false;
+    args.drop_seed_offset  = has_dropout ? std::make_pair(uint64_t(1), uint64_t(0))
+                                         : std::make_pair(uint64_t(0), uint64_t(0));
 
     try
     {

projects/composablekernel/dispatcher/codegen/fmha/codegen.py

@@ -753,6 +753,16 @@ def _batch_prefill_kernel_body(name: str, config: dict) -> str:
     warp = alg["warp"]
     pad = alg["padding"]
     ns = f"ns_{name}"
+    pipeline_policy = (
+        ", ck_tile::BlockFmhaBatchPrefillPipelineQRKSVSAsyncGfx11Policy"
+        if alg["pipeline"] == "batch_prefill_gfx11"
+        else ""
+    )
+    pipeline_enum_key = (
+        "batch_prefill_gfx11"
+        if alg["pipeline"] == "batch_prefill_gfx11"
+        else "batch_prefill_async"
+    )
     return f"""// SPDX-License-Identifier: MIT
 #pragma once
 
@@ -808,7 +818,7 @@ def _batch_prefill_kernel_body(name: str, config: dict) -> str:
     false,
     {sig["page_size"]},
     fmha_trait>;
-using fmha_pipeline = ck_tile::BlockFmhaBatchPrefillPipelineQRKSVSAsync<fmha_pipeline_problem>;
+using fmha_pipeline = ck_tile::BlockFmhaBatchPrefillPipelineQRKSVSAsync<fmha_pipeline_problem{pipeline_policy}>;
 using fmha_epilogue = ck_tile::Default2DEpilogue<
     ck_tile::Default2DEpilogueProblem<typename FmhaFwdTypeConfig<fmha_dtype>::OaccDataType,
                                       typename FmhaFwdTypeConfig<fmha_dtype>::ODataType,
@@ -821,7 +831,7 @@ def _batch_prefill_kernel_body(name: str, config: dict) -> str:
                                              {mode_cpp},
                                              {tile[0]}, {tile[1]}, {tile[2]}, {tile[3]}, {tile[4]}, {tile[5]},
                                              {vlayout_cpp},
-                                             {PIPELINE_ENUM_TO_CPP["batch_prefill_async"]},
+                                             {PIPELINE_ENUM_TO_CPP[pipeline_enum_key]},
                                              {_bool_cpp(sig["logits"])},
                                              fmha_mask,
                                              {_bias_cpp(sig["bias"])},

projects/composablekernel/dispatcher/codegen/fmha/fmha_arch_specs.json

@@ -31,7 +31,7 @@
       "family": "rdna3",
       "arch_tag": "ck_tile::gfx11_t",
       "supported_dtypes": ["fp16", "bf16"],
-      "supported_pipelines": ["qr", "qr_pagedkv", "qr_nwarp_sshuffle", "appendkv", "bwd"],
+      "supported_pipelines": ["qr", "qr_pagedkv", "qr_nwarp_sshuffle", "appendkv", "batch_prefill_gfx11", "bwd"],
       "supports_trload": false,
       "supports_v3": false
     },

projects/composablekernel/dispatcher/codegen/fmha/instance_gen.py

@@ -1457,6 +1457,8 @@ def bwd_dq_wave_warp(tile, hq, trload=False):
 _MASK_MAP = {"no": "no", "causal": "top_left", "generic": "generic"}
 _BIAS_MAP = {"no": "no", "bias": "bias", "alibi": "alibi"}
 
+GFX11_WARP_SIZE = 32
+
 
 def _pad_val(s: str) -> int:
     if s == "f":
@@ -2392,17 +2394,32 @@ def _bp_bk1(bm0, bn0, bk0, hq):
         if restrict_hdims is not None:
             hdims = [hv for hv in hdims if hv in restrict_hdims]
         for hq, hv in hdims:
+            # gfx1100 must stay inside Gfx11Policy::UseIndependentVBuffer; the
+            # fallback gemm1 does not compile against gfx11 WMMA.
+            if arch == "gfx1100" and (hq, hv) != (128, 128):
+                continue
             tiles = generate_splitkv_tiles(arch, dtype, hq, hv)
             bp_specs = get_batch_prefill_pipelines(dtype, hq, receipt)
             for tc in tiles:
                 bk1 = _bp_bk1(tc.bm0, tc.bn0, tc.bk0, hq)
+                if arch == "gfx1100" and (
+                    tc.rm0 * GFX11_WARP_SIZE != 256
+                    or tc.bn0 != 32
+                    or tc.bn1 != 128
+                    or bk1 != 32
+                ):
+                    continue
 
                 # qr_async stages K into LDS through a bk1-major descriptor while the gemm0
                 # loop reads bk0 chunks, therefore the pipeline requires bk0 == bk1
                 if tc.bk0 != bk1:
                     continue
 
                 for spec in bp_specs:
+                    if arch == "gfx1100" and (
+                        spec.kv_memory_layout != "linear" or spec.dropout == "t"
+                    ):
+                        continue
                     mm = _MASK_MAP.get(spec.mask, spec.mask)
                     mb = _BIAS_MAP.get(spec.bias, spec.bias)
                     if allowed_masks is not None and mm not in allowed_masks:
@@ -2421,7 +2438,9 @@ def _bp_bk1(bm0, bn0, bk0, hq):
                                 mode="group",
                                 hdim_q=hq,
                                 hdim_v=hv,
-                                pipeline="qr_async",
+                                pipeline="batch_prefill_gfx11"
+                                if arch == "gfx1100"
+                                else "qr_async",
                                 tile_m0=tc.bm0,
                                 tile_n0=tc.bn0,
                                 tile_k0=tc.bk0,

projects/composablekernel/dispatcher/codegen/fmha/symbol_map.py

@@ -240,6 +240,7 @@
     "qr_nwarp_sshuffle": "ck_tile::BlockFmhaFwdSplitKVPipelineNWarpSShuffleQRKSVS",
     "appendkv": "ck_tile::BlockFmhaFwdAppendKVPipeline",
     "batch_prefill_async": "ck_tile::BlockFmhaBatchPrefillPipelineQRKSVSAsync",
+    "batch_prefill_gfx11": "ck_tile::BlockFmhaBatchPrefillPipelineQRKSVSAsync",
 }
 
 # Pipeline tag → C++ pipeline enum value.
@@ -254,6 +255,7 @@
     "qr_pagedkv": "ck_tile::BlockFmhaPipelineEnum::QRKSVS",
     "qr_nwarp_sshuffle": "ck_tile::BlockFmhaPipelineEnum::QRKSVS",
     "batch_prefill_async": "ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC",
+    "batch_prefill_gfx11": "ck_tile::BlockFmhaPipelineEnum::QRKSVS_ASYNC",
 }
 
 BOOL_MAP = {

projects/composablekernel/dispatcher/codegen/fmha/validation.py

@@ -881,8 +881,34 @@ def validate_config(
             result.add_error("batch_prefill page_size must be a positive power of two")
         if sig.get("mode", "batch") != "group":
             result.add_error("batch_prefill requires group mode")
-        if pipeline != "qr_async":
-            result.add_error("batch_prefill currently uses qr_async pipeline")
+        if pipeline not in {"qr_async", "batch_prefill_gfx11"}:
+            result.add_error(
+                "batch_prefill currently uses qr_async or batch_prefill_gfx11 pipeline"
+            )
+        if arch == "gfx1100" and pipeline != "batch_prefill_gfx11":
+            result.add_error(
+                "gfx1100 batch_prefill must use batch_prefill_gfx11 policy"
+            )
+        if pipeline == "batch_prefill_gfx11":
+            # Mirror of Gfx11Policy::UseIndependentVBuffer. Outside this
+            # predicate gemm1 falls back to a shape gfx11 WMMA cannot consume.
+            if arch != "gfx1100":
+                result.add_error("batch_prefill_gfx11 is only supported on gfx1100")
+            if sig.get("kv_memory_layout") != "linear":
+                result.add_error(
+                    "batch_prefill_gfx11 currently supports only linear KV layout"
+                )
+            if (hdim_q, hdim_v) != (128, 128):
+                result.add_error("batch_prefill_gfx11 currently supports only hdim 128")
+            if sig.get("dropout", False):
+                result.add_error("batch_prefill_gfx11 does not support dropout")
+            wave = alg["wave"]
+            if len(wave) >= 3 and wave[0] * wave[1] * wave[2] * 32 != 256:
+                result.add_error("batch_prefill_gfx11 requires block size 256")
+            if len(tile) >= 5 and (tile[1], tile[3], tile[4]) != (32, 128, 32):
+                result.add_error(
+                    "batch_prefill_gfx11 requires tile N0=32, N1=128, K1=32"
+                )
 
     if family == "fwd_appendkv":
         if sig.get("mode", "batch") != "batch":

projects/composablekernel/dispatcher/python/fmha_utils.py

@@ -840,6 +840,79 @@ def run(
             Q_c = np.ascontiguousarray(Q.astype(in_dt))
             K_c = np.ascontiguousarray(K.astype(in_dt))
             V_c = np.ascontiguousarray(V.astype(in_dt))
+        if api_family == "batch_prefill":
+            page_size = kwargs.get("page_size", 64)
+            kv_layout = kwargs.get("kv_layout", 0)
+
+            def _pack_linear_paged(x):
+                bsz, heads, seqlen, hdim = x.shape
+                pages_per_seq = (seqlen + page_size - 1) // page_size
+                packed = np.zeros(
+                    (bsz * pages_per_seq, heads, page_size, hdim), dtype=x.dtype
+                )
+                for batch_id in range(bsz):
+                    for page_id in range(pages_per_seq):
+                        begin = page_id * page_size
+                        end = min(begin + page_size, seqlen)
+                        packed[
+                            batch_id * pages_per_seq + page_id, :, : end - begin, :
+                        ] = x[batch_id, :, begin:end, :]
+                return np.ascontiguousarray(packed.reshape(-1))
+
+            if kv_layout != 1:
+                return FmhaResult(
+                    success=False,
+                    error="batch_prefill runner packs linear KV only",
+                )
+            K_c = _pack_linear_paged(K_c)
+            V_c = _pack_linear_paged(V_c)
+            Q_api = np.ascontiguousarray(np.transpose(Q_c, (0, 2, 1, 3)))
+            O_api = np.zeros(
+                (prob.batch, prob.seqlen_q, prob.nhead_q, prob.hdim_v), dtype=out_dt
+            )
+            time_ms = ctypes.c_float(0.0)
+            rc = self._lib._lib.fmha_dispatcher_run_batch_prefill(
+                ctypes.c_void_p(Q_api.ctypes.data),
+                ctypes.c_void_p(K_c.ctypes.data),
+                ctypes.c_void_p(V_c.ctypes.data),
+                ctypes.c_void_p(O_api.ctypes.data),
+                prob.batch,
+                prob.nhead_q,
+                prob.nhead_k,
+                prob.seqlen_q,
+                prob.seqlen_k,
+                prob.hdim_q,
+                prob.hdim_v,
+                ctypes.c_float(prob.scale),
+                mask_type,
+                bias_type,
+                page_size,
+                kv_layout,
+                kwargs.get("kv_lookup", 0),
+                kwargs.get("is_v_rowmajor", 1),
+                data_type.encode(),
+                has_lse,
+                has_dropout,
+                has_logits,
+                has_sink,
+                has_skip,
+                ctypes.byref(time_ms),
+            )
+            if rc != 0:
+                return FmhaResult(success=False, error=f"Kernel failed (rc={rc})")
+            ops = prob.num_ops
+            tflops = (
+                ops / (time_ms.value * 1e-3) / 1e12
+                if time_ms.value > 0 and ops > 0
+                else 0.0
+            )
+            O_c = np.ascontiguousarray(np.transpose(O_api, (0, 2, 1, 3)))
+            if data_type == "bf16":
+                O_c = _bf16_to_float32(O_c)
+            return FmhaResult(
+                success=True, output=O_c, time_ms=time_ms.value, tflops=tflops
+            )
+
         O_c = np.zeros(prob.o_shape(), dtype=out_dt)
 
         d_q, d_k, d_v, d_o = (ctypes.c_void_p() for _ in range(4))

projects/composablekernel/dispatcher/tests/test_fmha_codegen.py

@@ -158,6 +158,52 @@ def test_batch_prefill_requires_group_mode(self):
         self.assertFalse(result.valid)
         self.assertTrue(any("group mode" in error for error in result.errors))
 
+    def test_gfx1100_batch_prefill_codegen_emits_gfx11_policy(self):
+        config = sample_config(
+            arch="gfx1100",
+            signature={
+                "family": "batch_prefill",
+                "mode": "group",
+                "paged_kv": True,
+                "page_size": 16,
+                "kv_memory_layout": "linear",
+                "kv_lookup_table": "vllm",
+            },
+            algorithm={
+                "pipeline": "batch_prefill_gfx11",
+                "tile": [128, 32, 32, 128, 32, 128],
+                "wave": [8, 1, 1, 8, 1, 1, 1, 1, 1],
+                "warp": [16, 16, 16, 16, 16, 16, 16, 16, 16],
+            },
+        )
+        result = validate_config(config)
+        self.assertTrue(result.valid, result.errors)
+        with tempfile.TemporaryDirectory() as tmpdir:
+            cmd = [
+                sys.executable,
+                str(CODEGEN),
+                "--output-dir",
+                tmpdir,
+                "--gpu-target",
+                "gfx1100",
+                "--config-json",
+                json.dumps(config),
+            ]
+            proc = subprocess.run(
+                cmd, capture_output=True, text=True, cwd=str(ROOT / "codegen")
+            )
+            self.assertEqual(proc.returncode, 0, msg=proc.stderr or proc.stdout)
+            generated = list(Path(tmpdir).glob("fmha_*.hpp"))
+            self.assertEqual(len(generated), 1)
+            text = generated[0].read_text()
+            self.assertIn("BlockFmhaBatchPrefillPipelineQRKSVSAsyncGfx11Policy", text)
+            self.assertNotIn("QRKSVSAsyncDefaultPolicy", text)
+            # kHasSink occupies the bool slot immediately before kPageBlockSize.
+            self.assertRegex(
+                text,
+                r"TileFmhaBatchPrefillTraits<[\s\S]*false,\s*16,",
+            )
+
     def _batch_prefill_traits_args(self, sink=False):
         config = sample_config(
             arch="gfx950",

projects/composablekernel/dispatcher/tests/test_fmha_dispatcher.cpp

@@ -159,6 +159,52 @@ TEST(FmhaDispatcherTest, PlansSingleStageFwd)
     EXPECT_EQ(plan.stages[0].family, FmhaKernelFamily::Fwd);
 }
 
+TEST(FmhaDispatcherTest, PlansSingleStageBatchPrefillOnGfx1100)
+{
+    FmhaRegistry registry;
+    registry.register_kernel(std::make_shared<MockFmhaKernel>(
+        make_key(FmhaKernelFamily::BatchPrefill, "prefill"), "prefill"));
+
+    FmhaDispatcher dispatcher(&registry);
+
+    fmha_batch_prefill_traits traits{};
+    traits.hdim_q           = 128;
+    traits.hdim_v           = 128;
+    traits.data_type        = "fp16";
+    traits.is_group_mode    = true;
+    traits.is_v_rowmajor    = true;
+    traits.mask_type        = mask_enum::no_mask;
+    traits.bias_type        = bias_enum::no_bias;
+    traits.has_lse          = true;
+    traits.kv_memory_layout = ck_tile::BlockAttentionKVCacheMemoryLayoutEnum::LINEAR_LAYOUT;
+    traits.kv_lookup_table  = ck_tile::BlockAttentionKVCacheLookupTableEnum::VLLM_BLOCK_TABLE_2D;
+    traits.page_size        = 16;
+
+    fmha_batch_prefill_args args{};
+    args.batch             = 1;
+    args.seqlen_q          = 128;
+    args.seqlen_k          = 1024;
+    args.max_seqlen_q      = 128;
+    args.hdim_q            = 128;
+    args.hdim_v            = 128;
+    args.nhead_q           = 8;
+    args.nhead_k           = 8;
+    args.num_total_pages   = 1024;
+    args.page_block_size   = 16;
+    args.kv_memory_layout  = ck_tile::BlockAttentionKVCacheMemoryLayoutEnum::LINEAR_LAYOUT;
+    args.kv_lookup_table   = ck_tile::BlockAttentionKVCacheLookupTableEnum::VLLM_BLOCK_TABLE_2D;
+    args.kv_indptr         = reinterpret_cast<void*>(0x1);
+    args.kv_page_indices   = reinterpret_cast<void*>(0x1);
+    args.kv_last_page_lens = reinterpret_cast<void*>(0x1);
+    args.seqstart_q_ptr    = reinterpret_cast<void*>(0x1);
+
+    auto problem = FmhaProblem::from_invocation(FmhaInvocation::make(traits, args), "gfx1100");
+    auto plan    = dispatcher.plan(problem);
+    ASSERT_TRUE(plan.is_valid());
+    ASSERT_EQ(plan.stages.size(), 1u);
+    EXPECT_EQ(plan.stages[0].family, FmhaKernelFamily::BatchPrefill);
+}
+
 TEST(FmhaDispatcherTest, PlansSingleStagePagedKv)
 {
     FmhaRegistry registry;