add more instances, manage splitkv verifications

2026-07-12 02:05:50 +00:00 · 2026-04-17 13:20:45 +00:00
parent a30e0c5cce
commit 00514bed47
3 changed files with 46 additions and 51 deletions
--- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py
+++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py
@@ -206,6 +206,15 @@ float {F_func_name}([[maybe_unused]] fmha_fwd_traits t, [[maybe_unused]] fmha_fw
 """
 FMHA_FWD_API_FOOTER_TEMPLATE = """
 float fmha_fwd(fmha_fwd_traits traits, fmha_fwd_args args, const ck_tile::stream_config& config) {{
+    // Force-kernel override: dispatch a specific kernel via fmha_fwd_all()
+    {{ const char* fk = std::getenv("CK_TILE_FMHA_FWD_FORCE_KERNEL");
+       if(fk && fk[0]) {{
+           auto results = fmha_fwd_all(traits, args, config);
+           if(!results.empty()) return results[0].second;
+           return -1;
+       }}
+    }}
+
    const std::string device_name = ck_tile::get_device_name();

    const bool is_swa = (traits.mask_type != mask_enum::no_mask) and
@@ -230,6 +239,12 @@ FMHA_FWD_ALL_API_FUNC_TEMPLATE = """
 std::vector<std::pair<std::string, float>> {F_func_name}([[maybe_unused]] fmha_fwd_traits t, [[maybe_unused]] fmha_fwd_args a, [[maybe_unused]] const ck_tile::stream_config& s) {{
    std::vector<std::pair<std::string, float>> results;

+    // Force-kernel filter: when set, only the matching kernel is dispatched.
+    // Used by per-kernel verification (-v=1 -run_all_kernels=1).
+    std::string __fmha_force_kernel;
+    {{ const char* __fk = std::getenv("CK_TILE_FMHA_FWD_FORCE_KERNEL");
+       if(__fk && __fk[0]) __fmha_force_kernel = __fk; }}
+
    [[maybe_unused]] const float min_cu_util_rate = 0.8; // minimum CU utilization rate

    unsigned num_cus;
@@ -267,7 +282,8 @@ FMHA_FWD_ALL_API_PER_HDIM_CASE = """{F_if}(t.hdim_q <= {F_hdim} && t.hdim_v <= {
 # 1. Always uses "if" (not "else if") — doesn't skip after first match
 # 2. No seqtune heuristic — runs all tile sizes
 # 3. Pushes result into vector instead of returning
-FMHA_FWD_ALL_API_INNER_DISPATCH = """if((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_lse == {F_lse})  && (t.has_dropout == {F_dropout}) && (t.qscale_type == {F_qscale_check}) && (t.skip_min_seqlen_q == {F_skip}) &&(t.has_sink == {F_sink}) &&
+FMHA_FWD_ALL_API_INNER_DISPATCH = """if((__fmha_force_kernel.empty() || __fmha_force_kernel == \"{F_kname}\") &&
+   (t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.has_lse == {F_lse})  && (t.has_dropout == {F_dropout}) && (t.qscale_type == {F_qscale_check}) && (t.skip_min_seqlen_q == {F_skip}) &&(t.has_sink == {F_sink}) &&
        ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck}) && ({F_constraint})) {{
    using trait_ = fmha_fwd_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, {F_lse}, {F_dropout}, {F_qscale}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}, {F_trload}, {F_skip}, {F_sink}>;
    float t_ = fmha_fwd_<trait_, {F_arch.tag}>(s, a);
@@ -1138,6 +1154,9 @@ class KernelComponentFactoryGfx9(CompatibilityRuleFactoryGfx9):
            ]  # fmt: skip
            extra[(256, 256)] = [
                FmhaFwdTileSize(128, 128,  64, 256,  32, 256,  4, 1, 1,  4, 1, 1,  32, 32, 16,  32, 32, 16,  -1),
+                FmhaFwdTileSize(128, 128, 128, 256,  64, 256,  4, 1, 1,  4, 1, 1,  32, 32, 16,  32, 32, 16,  -1),
+                FmhaFwdTileSize(256, 128,  32, 256,  32, 256,  8, 1, 1,  8, 1, 1,  32, 32, 16,  32, 32, 16,  -1),
+                FmhaFwdTileSize(256, 128,  64, 256,  32, 256,  8, 1, 1,  8, 1, 1,  32, 32, 16,  32, 32, 16,  -1),
            ]  # fmt: skip
        return extra

--- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py
+++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py
@@ -270,6 +270,15 @@ float fmha_fwd_splitkv_(const ck_tile::stream_config& s, fmha_fwd_splitkv_args a
 }}

 float fmha_fwd_splitkv(fmha_fwd_splitkv_traits t, fmha_fwd_splitkv_args a, const ck_tile::stream_config& s) {{
+    // Force-kernel override: dispatch a specific kernel via fmha_fwd_splitkv_all()
+    {{ const char* fk = std::getenv("CK_TILE_FMHA_FWD_SPLITKV_FORCE_KERNEL");
+       if(fk && fk[0]) {{
+           auto results = fmha_fwd_splitkv_all(t, a, s);
+           if(!results.empty()) return results[0].second;
+           return -1;
+       }}
+    }}
+
    float r = -1;

    [[maybe_unused]] const std::string device_name = ck_tile::get_device_name();
@@ -313,6 +322,12 @@ FMHA_FWD_SPLITKV_ALL_API_FUNC_TEMPLATE = """
 std::vector<std::pair<std::string, float>> {F_func_name}([[maybe_unused]] fmha_fwd_splitkv_traits t, [[maybe_unused]] fmha_fwd_splitkv_args a, [[maybe_unused]] const ck_tile::stream_config& s) {{
    std::vector<std::pair<std::string, float>> results;

+    // Force-kernel filter: when set, only the matching kernel is dispatched.
+    // Used by per-kernel verification (-v=1 -run_all_kernels=1).
+    std::string __fmha_force_kernel;
+    {{ const char* __fk = std::getenv("CK_TILE_FMHA_FWD_SPLITKV_FORCE_KERNEL");
+       if(__fk && __fk[0]) __fmha_force_kernel = __fk; }}
+
    [[maybe_unused]] const std::string device_name = ck_tile::get_device_name();

 {F_dispatch}
@@ -323,7 +338,8 @@ std::vector<std::pair<std::string, float>> {F_func_name}([[maybe_unused]] fmha_f
 # Key differences from FMHA_FWD_SPLITKV_API_INNER_DISPATCH:
 # 1. Always uses "if" (not "else if") — doesn't skip after first match
 # 2. Pushes result into vector instead of returning
-FMHA_FWD_SPLITKV_ALL_API_INNER_DISPATCH = """if((t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.do_fp8_static_quant == {F_squant}) &&
+FMHA_FWD_SPLITKV_ALL_API_INNER_DISPATCH = """if((__fmha_force_kernel.empty() || __fmha_force_kernel == \"{F_kname}\") &&
+   (t.is_group_mode == {F_mode}) && (t.is_v_rowmajor == {F_vlayout}) && (t.has_logits_soft_cap == {F_logits}) && ({F_mask_check}) && (t.bias_type == {F_bias_check}) && (t.do_fp8_static_quant == {F_squant}) &&
        ((a.block_table_ptr != nullptr) == {F_pagedkv}) && (t.has_sink == {F_sink}) && ({F_scheck}) && ({F_skcheck}) && ({F_dcheck}) && ({F_dvcheck})) {{
    using traits_ = fmha_fwd_splitkv_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bm0}, {F_bn0}, {F_bk0}, {F_bn1}, {F_bk1}, {F_bk0max}, {F_vlayout}, {F_pipeline_enum}, {F_logits}, {F_mask}, {F_bias}, true, {F_squant}, {F_pagedkv},{F_sink}, {F_spad}, {F_skpad}, {F_dpad}, {F_dvpad}>;

@@ -993,7 +1009,11 @@ class KernelComponentFactoryGfx9(KernelComponentFactoryBase):
                FmhaFwdTileSize( 32, 128, 32, 128, 32, 128, 2, 1, 1, 2, 1, 1, 16, 16, 16, 16, 16, 16, -1),
            ]  # fmt: skip
            extra["256"] = [
+                FmhaFwdTileSize(128, 128,  32, 256,  32, 256,  4, 1, 1,  4, 1, 1,  32, 32, 16,  32, 32, 16,  -1),
+                FmhaFwdTileSize(128, 128,  64, 256,  32, 256,  4, 1, 1,  4, 1, 1,  32, 32, 16,  32, 32, 16,  -1),
                FmhaFwdTileSize(128, 128,  64, 256, 128, 256,  4, 1, 1,  4, 1, 1,  32, 32, 16,  32, 32, 16,  -1),
+                FmhaFwdTileSize(128, 128, 128, 256,  32, 256,  4, 1, 1,  4, 1, 1,  32, 32, 16,  32, 32, 16,  -1),
+                FmhaFwdTileSize(128, 128, 128, 256,  64, 256,  4, 1, 1,  4, 1, 1,  32, 32, 16,  32, 32, 16,  -1),
            ]  # fmt: skip
        return extra

--- a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp
+++ b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp
@@ -1784,58 +1784,14 @@ fwd_result fmha_fwd_run(mode_enum mode,

        if(do_validation != 0)
        {
+            const char* force_kernel_env =
 #if CK_TILE_FMHA_FWD_SPLITKV_API
-            if(use_splitkv_all)
-            {
-                std::cout << "[run_all_kernels] per-kernel verification is currently supported "
-                             "for fwd kernels only (not splitkv path); running heuristic "
-                             "verification pass instead (-run_all_kernels=0, -v="
-                          << do_validation << ")"
-                          << std::endl;
-
-                return fmha_fwd_run<DataTypeConfig>(mode,
-                                                    batch,
-                                                    nhead,
-                                                    nhead_k,
-                                                    seqlen_qs,
-                                                    seqlen_ks,
-                                                    hdim_q,
-                                                    hdim_v,
-                                                    seqlen_knew,
-                                                    seqlen_qpads,
-                                                    seqlen_kpads,
-                                                    q_eff_lens_per_batch,
-                                                    kv_eff_lens_per_batch,
-                                                    rotary_dim,
-                                                    i_perm,
-                                                    o_perm,
-                                                    scale_s,
-                                                    logits_soft_cap,
-                                                    is_v_rowmajor,
-                                                    lse,
-                                                    page_block_size,
-                                                    use_cache_batch_idx,
-                                                    bias_str,
-                                                    p_drop,
-                                                    drop_seed,
-                                                    drop_offset,
-                                                    drop_prefs,
-                                                    mask_str,
-                                                    qscale_str,
-                                                    is_rotary_interleaved,
-                                                    num_splits,
-                                                    init_method,
-                                                    seed,
-                                                    do_validation,
-                                                    init_sink_value,
-                                                    stream_config,
-                                                    false,
-                                                    json);
-            }
+                use_splitkv_all ? "CK_TILE_FMHA_FWD_SPLITKV_FORCE_KERNEL"
+                                : "CK_TILE_FMHA_FWD_FORCE_KERNEL";
+#else
+                "CK_TILE_FMHA_FWD_FORCE_KERNEL";
 #endif

-            constexpr const char* force_kernel_env = "CK_TILE_FMHA_FWD_FORCE_KERNEL";
-
            ck_tile::stream_config verify_sc{stream_config.stream_id_,
                                             false,
                                             0,