From 00514bed4727d9eff3fe516d7168b7a22e2f967c Mon Sep 17 00:00:00 2001 From: Mohsen Saffari Date: Fri, 17 Apr 2026 13:20:45 +0000 Subject: [PATCH] add more instances, manage splitkv verifications --- .../ck_tile/01_fmha/codegen/ops/fmha_fwd.py | 21 +++++++- .../01_fmha/codegen/ops/fmha_fwd_splitkv.py | 22 +++++++- example/ck_tile/01_fmha/fmha_fwd_runner.hpp | 54 ++----------------- 3 files changed, 46 insertions(+), 51 deletions(-) diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py index 3c272699a1..06166607ca 100644 --- 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> {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> 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_(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 diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py index 1ac8f3cf60..e5e3f68022 100644 --- 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> {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> 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> {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 diff --git a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp index 80fb24d204..343e459fc5 100644 --- 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(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,