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

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
 
@@ -782,6 +792,7 @@ def _batch_prefill_kernel_body(name: str, config: dict) -> str:
                                                        {_qscale_cpp(sig["qscale"])},
                                                        {alg["block_per_cu"]},
                                                        false,
+                                                       {_bool_cpp(sig["sink"])},
                                                        {sig["page_size"]},
                                                        {_kv_memory_cpp(sig["kv_memory_layout"])},
                                                        {_kv_lookup_cpp(sig["kv_lookup_table"])}>;
@@ -807,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,
@@ -820,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"])},
@@ -833,6 +844,7 @@ def _batch_prefill_kernel_body(name: str, config: dict) -> str:
                                              {_bool_cpp(pad[3])},
                                              false,
                                              false,
+                                             {_bool_cpp(sig["sink"])},
                                              {sig["page_size"]},
                                              {_kv_memory_cpp(sig["kv_memory_layout"])},
                                              {_kv_lookup_cpp(sig["kv_lookup_table"])}>;

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))