From a30e0c5ccef09597ab97fb87fd75429de69fd853 Mon Sep 17 00:00:00 2001 From: Mohsen Saffari Date: Fri, 17 Apr 2026 12:03:58 +0000 Subject: [PATCH] add run_all_kernels benchmarking mode with extended tuning tiles --- .../ck_tile/01_fmha/codegen/ops/fmha_fwd.py | 193 +++++++++- .../01_fmha/codegen/ops/fmha_fwd_splitkv.py | 186 ++++++++- example/ck_tile/01_fmha/example_fmha_fwd.cpp | 7 +- example/ck_tile/01_fmha/fmha_fwd.hpp | 5 + example/ck_tile/01_fmha/fmha_fwd_runner.hpp | 353 ++++++++++++++++++ 5 files changed, 723 insertions(+), 21 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 4063c4b120..3c272699a1 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py @@ -225,6 +225,56 @@ float fmha_fwd(fmha_fwd_traits traits, fmha_fwd_args args, const ck_tile::stream }} """ +# --- Templates for fmha_fwd_all() "run all kernels" benchmarking mode --- +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; + + [[maybe_unused]] const float min_cu_util_rate = 0.8; // minimum CU utilization rate + + unsigned num_cus; + if(!get_num_cus(num_cus)) {{ + return results; + }} + + [[maybe_unused]] auto get_num_blocks = [&](unsigned kM0) {{ + return get_num_thread_blocks(a.batch, a.nhead_q, a.max_seqlen_q, kM0); + }}; + + [[maybe_unused]] const std::string device_name = ck_tile::get_device_name(); + +{F_dispatch} + return results; +}} +""" + +FMHA_FWD_ALL_API_PER_ARCH = """{F_if}({F_arch.device_name_check}) {{ +{F_dtype_case} +}} +""" + +FMHA_FWD_ALL_API_PER_DTYPE = """{F_if}(t.data_type.compare(\"{F_dtype}\") == 0) {{ +{F_hdim_case} +}} +""" + +FMHA_FWD_ALL_API_PER_HDIM_CASE = """{F_if}(t.hdim_q <= {F_hdim} && t.hdim_v <= {F_hdim_v}) {{ +{F_inner_dispatch} +}} +""" + +# Key differences from FMHA_FWD_API_INNER_DISPATCH: +# 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}) && + ({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); + if(t_ >= 0) results.push_back({{ \"{F_kname}\", t_ }}); +}} +""" + FMHA_FWD_API_PER_ARCH = """{F_if}({F_arch.device_name_check}) {{ {F_dtype_case} }} @@ -291,6 +341,7 @@ class FmhaFwdApiTrait: tr_load: str sink: str constraint: CppConstraint + is_tuning_extra: bool = False # True for extended tiles from get_tuning_extra_tiles() @property def name(self) -> str: @@ -598,6 +649,111 @@ class FmhaFwdApiPool: F_func_name=func_name, F_dispatch=indent(per_arch) ) + def render_all( + self, func_name, filter_fn: Optional[Callable[[FmhaFwdApiTrait], bool]] = None + ) -> str: + """Render a function that runs ALL matching kernel instances (no heuristic tile selection).""" + if filter_fn is None: + + def accept_all(trait: FmhaFwdApiTrait) -> bool: + return True + + filter_fn = accept_all + + def has_traits(node) -> bool: + if isinstance(node, list): + return any(filter_fn(elem) for elem in node) + elif isinstance(node, OrderedDict): + return any(has_traits(val) for val in node.values()) + return False + + per_arch = str() + for i_arch, (arch, pool_by_arch) in enumerate( + item for item in self.pool.items() if has_traits(item[1]) + ): + per_dtypes = str() + for i_dtype, (dtype, pool_by_dtype) in enumerate( + item for item in pool_by_arch.items() if has_traits(item[1]) + ): + per_hdim_case = str() + for i_hdim, ((hdim, hdim_v), pool_by_hdim) in enumerate( + item for item in pool_by_dtype.items() if has_traits(item[1]) + ): + inners = str() + for trait in (trait for trait in pool_by_hdim if filter_fn(trait)): + # Build the kernel name string matching FmhaFwdKernel.name format + pad_suffix = "" + for flag, tag in [ + (trait.spad, "s"), + (trait.skpad, "sk"), + (trait.dpad, "d"), + (trait.dvpad, "dv"), + ]: + if flag == "t": + pad_suffix += tag + pad_suffix = f"_p{pad_suffix}" if pad_suffix else "_npad" + kname = ( + f"fmha_fwd_d{hdim}_{dtype}" + f"_{'group' if trait.mode == 'group' else 'batch'}" + f"_b{trait.bm0}x{trait.bn0}x{trait.bk0}x{trait.bn1}x{trait.bk1}x{trait.bk0max}" + f"{pad_suffix}" + ) + if trait.is_tuning_extra: + kname += " [ext]" + inners += FMHA_FWD_ALL_API_INNER_DISPATCH.format( + F_arch=arch, + F_mode=MODE_MAP[trait.mode], + F_vlayout=LAYOUT_MAP[trait.vlayout], + F_pipeline_enum=PIPELINE_ENUM_MAP[trait.pipeline_tag], + F_logits=BOOL_MAP[trait.logits], + F_mask=get_mask_cpp_type(trait.mask), + F_mask_check=get_mask_cpp_check_expr(trait.mask), + F_bias_check=BIAS_CHECK_MAP[trait.bias], + F_bias=BIAS_MAP[trait.bias], + F_lse=BOOL_MAP[trait.lse], + F_dropout=BOOL_MAP[trait.dropout], + F_skip=BOOL_MAP[trait.skip], + F_trload=BOOL_MAP[trait.tr_load], + F_qscale_check=QSCALE_CHECK_MAP[trait.qscale], + F_qscale=QSCALE_MAP[trait.qscale], + F_sink=BOOL_MAP[trait.sink], + F_scheck=trait.scheck, + F_skcheck=trait.skcheck, + F_dcheck=trait.dcheck, + F_dvcheck=trait.dvcheck, + F_constraint=trait.constraint, + F_spad=BOOL_MAP[trait.spad], + F_skpad=BOOL_MAP[trait.skpad], + F_dpad=BOOL_MAP[trait.dpad], + F_dvpad=BOOL_MAP[trait.dvpad], + F_bm0=trait.bm0, + F_bn0=trait.bn0, + F_bk0=trait.bk0, + F_bn1=trait.bn1, + F_bk1=trait.bk1, + F_bk0max=trait.bk0max, + F_hdim=hdim, + F_dtype=FWD_DTYPE_MAP[dtype], + F_kname=kname, + ) + per_hdim_case += FMHA_FWD_ALL_API_PER_HDIM_CASE.format( + F_if=if_(i_hdim), + F_hdim=hdim, + F_hdim_v=hdim_v, + F_inner_dispatch=indent(inners), + ) + per_dtypes += FMHA_FWD_ALL_API_PER_DTYPE.format( + F_if=if_(i_dtype), F_dtype=dtype, F_hdim_case=indent(per_hdim_case) + ) + per_arch += FMHA_FWD_ALL_API_PER_ARCH.format( + F_if=if_(i_arch), + F_arch=arch, + F_dtype_case=indent(per_dtypes), + ) + return FMHA_FWD_ALL_API_FUNC_TEMPLATE.format( + F_func_name=func_name, F_dispatch=indent(per_arch) + ) + @dataclass class FmhaFwdTileSize: @@ -972,10 +1128,14 @@ class KernelComponentFactoryGfx9(CompatibilityRuleFactoryGfx9): @classmethod def get_tuning_extra_tiles(cls, dtype: str) -> dict: - """Additional tile sizes only available via tuning receipts (150, 250). - These tiles are NOT used by the heuristic dispatch path.""" + """Additional tile sizes merged for tuning/benchmarking receipts (0, 3, 150, 250). + For CK standalone (0/3) these are tagged is_tuning_extra and excluded from heuristic. + For AITER tuning (150/250) they participate in heuristic, selected via CSV.""" extra = {} if dtype in cls._DT_FP16_BF16: + extra[(128, 128)] = [ + FmhaFwdTileSize( 16, 128, 64, 128, 32, 128, 1, 1, 1, 1, 1, 1, 16, 16, 32, 16, 16, 32, -1), + ] # 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), ] # fmt: skip @@ -1467,17 +1627,25 @@ def get_fwd_blobs( factories = get_factories_for_targets(targets, get_factory) - # Tuning receipts (150, 250) include extended tile sizes for CSV-driven selection - _TUNING_RECEIPTS = frozenset({150, 250}) + # Receipts that include extended tuning tiles from get_tuning_extra_tiles(). + # - 150/250: AITER tuning receipts (tiles are part of the heuristic, selected via CSV) + # - 0/3: CK standalone build receipts (tiles are compiled for fmha_fwd_all() + # benchmarking but excluded from the heuristic via is_tuning_extra flag) + _EXTRA_TILE_RECEIPTS = frozenset({0, 3, 150, 250}) + # For standalone CK receipts, mark extended tiles so the heuristic excludes them. + _MARK_AS_TUNING_EXTRA = frozenset({0, 3}) for factory, dtype in ((f, t) for f in factories for t in f.supported_dtypes()): d = factory.get_hdim_tile_size_dict(dtype) - if receipt in _TUNING_RECEIPTS and hasattr(factory, 'get_tuning_extra_tiles'): + # Track which tiles are standard so we can tag extras with is_tuning_extra. + standard_tile_ids = {key: set(id(t) for t in tiles) for key, tiles in d.items()} + if receipt in _EXTRA_TILE_RECEIPTS and hasattr(factory, 'get_tuning_extra_tiles'): for key, extra_tiles in factory.get_tuning_extra_tiles(dtype).items(): if key in d: - d[key] = d[key] + extra_tiles + d[key] = sorted(d[key] + extra_tiles, key=lambda t: t.F_bm0) else: - d[key] = list(extra_tiles) + d[key] = sorted(extra_tiles, key=lambda t: t.F_bm0) + standard_tile_ids[key] = set() # all tiles in new key are extra # for hdim_str, mode, mask, bias, lse in itertools.product(d.keys(), MODE_MAP.keys(), MASK_MAP.keys(), ["t", "f"], ["t", "f"]): for ((hdim, hdim_v), tiles), mode in itertools.product( d.items(), MODE_MAP.keys() @@ -1510,7 +1678,10 @@ def get_fwd_blobs( if not fnmatch.fnmatch(k.name, kernel_filter): continue - api_pool.register_traits(k.api_trait()) + trait = k.api_trait() + if receipt in _MARK_AS_TUNING_EXTRA: + trait.is_tuning_extra = id(tile) not in standard_tile_ids.get((hdim, hdim_v), set()) + api_pool.register_traits(trait) gen.append(k) return (api_pool, gen) @@ -1525,10 +1696,10 @@ def write_fwd_api( autogen_dir: Path, ) -> None: def accept_only_v3(trait: FmhaFwdApiTrait) -> bool: - return trait.pipeline_tag == "qr_async_trload_v3" + return trait.pipeline_tag == "qr_async_trload_v3" and not trait.is_tuning_extra def accept_only_v2(trait: FmhaFwdApiTrait) -> bool: - return not accept_only_v3(trait) + return trait.pipeline_tag != "qr_async_trload_v3" and not trait.is_tuning_extra content = "".join( [ @@ -1542,6 +1713,8 @@ def write_fwd_api( False ] ), + # --- additive: fmha_fwd_all() for "run all kernels" benchmarking --- + api_pool.render_all("fmha_fwd_all"), ] ) update_file(autogen_dir / FMHA_FWD_API_FILENAME, content) 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 5b5ade1b69..1ac8f3cf60 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 @@ -7,7 +7,7 @@ import itertools from collections import OrderedDict from dataclasses import dataclass from pathlib import Path -from typing import List, Optional, Union +from typing import Callable, List, Optional, Union from codegen.arch import ArchTrait, get_factories_for_targets from codegen.cmake_config import GEN_DIR @@ -308,6 +308,45 @@ FMHA_FWD_SPLITKV_API_INNER_DISPATCH = """{F_if}((t.is_group_mode == {F_mode}) && """ +# --- Templates for fmha_fwd_splitkv_all() "run all kernels" benchmarking mode --- +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; + + [[maybe_unused]] const std::string device_name = ck_tile::get_device_name(); + +{F_dispatch} + return results; +}} +""" + +# 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}) && + ((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}>; + + using OaccDataType = typename FmhaFwdTypeConfig<{F_dtype}>::OaccDataType; + constexpr ck_tile::index_t kM0 = ck_tile::BlockFmhaSplitKVCombinePipelineTileSizes::kM0; + static_assert({F_bm0} % kM0 == 0); + static_assert({F_bn1} % {F_bn1comb} == 0); + + if (t.has_lse) {{ + if constexpr (!std::is_same_v<{F_dtype}, FmhaFwdFp8>) {{ + using traits2_ = fmha_fwd_splitkv_combine_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bn1comb}, true, {F_squant}, {F_spad}, {F_dvpad}>; + float t_ = fmha_fwd_splitkv_(s, a); + if(t_ >= 0) results.push_back({{ \"{F_kname}\", t_ }}); + }} + }} else {{ + using traits2_ = fmha_fwd_splitkv_combine_traits_<{F_hdim}, {F_dtype}, {F_mode}, {F_bn1comb}, false, {F_squant}, {F_spad}, {F_dvpad}>; + float t_ = fmha_fwd_splitkv_(s, a); + if(t_ >= 0) results.push_back({{ \"{F_kname}\", t_ }}); + }} +}} +""" + + @dataclass class FmhaFwdSplitKVApiTrait: arch: ArchTrait @@ -335,6 +374,7 @@ class FmhaFwdSplitKVApiTrait: pagedkv: str sink: str # sink or not bn1comb: int # tile size along v head_dim of combine kernel + is_tuning_extra: bool = False # True for extended tiles from get_tuning_extra_tiles() @property def name(self) -> str: @@ -558,8 +598,10 @@ class FmhaFwdSplitKVApiPool: for i_dtype, (dtype, pool_by_dtype) in enumerate(pool_by_arch.items()): per_hdim_case = str() for i_hdim, (hdim, pool_by_hdim) in enumerate(pool_by_dtype.items()): + # Exclude tuning-extra tiles from the heuristic dispatch + heuristic_traits = [t for t in pool_by_hdim if not t.is_tuning_extra] inners = str() - for i_trait, trait in enumerate(pool_by_hdim): + for i_trait, trait in enumerate(heuristic_traits): inners += FMHA_FWD_SPLITKV_API_INNER_DISPATCH.format( F_if=if_(i_trait), F_arch=arch, @@ -614,6 +656,110 @@ class FmhaFwdSplitKVApiPool: F_dispatch=indent(per_arch) ) + def render_all( + self, + func_name, + filter_fn: Optional[Callable[[FmhaFwdSplitKVApiTrait], bool]] = None, + ) -> str: + """Render a function that runs ALL matching splitkv kernel instances (no heuristic).""" + if filter_fn is None: + + def accept_all(trait: FmhaFwdSplitKVApiTrait) -> bool: + return True + + filter_fn = accept_all + + def has_traits(node) -> bool: + if isinstance(node, list): + return any(filter_fn(elem) for elem in node) + elif isinstance(node, OrderedDict): + return any(has_traits(val) for val in node.values()) + return False + + per_arch = str() + for i_arch, (arch, pool_by_arch) in enumerate( + item for item in self.pool.items() if has_traits(item[1]) + ): + per_dtypes = str() + for i_dtype, (dtype, pool_by_dtype) in enumerate( + item for item in pool_by_arch.items() if has_traits(item[1]) + ): + per_hdim_case = str() + for i_hdim, (hdim, pool_by_hdim) in enumerate( + item for item in pool_by_dtype.items() if has_traits(item[1]) + ): + inners = str() + for trait in (t for t in pool_by_hdim if filter_fn(t)): + # Build the kernel name string with padding suffix + pad_suffix = "" + for flag, tag in [ + (trait.spad, "s"), + (trait.skpad, "sk"), + (trait.dpad, "d"), + (trait.dvpad, "dv"), + ]: + if flag == "t": + pad_suffix += tag + pad_suffix = f"_p{pad_suffix}" if pad_suffix else "_npad" + kname = ( + f"fmha_fwd_splitkv_d{hdim}_{dtype}" + f"_{'group' if trait.mode == 'group' else 'batch'}" + f"_b{trait.bm0}x{trait.bn0}x{trait.bk0}x{trait.bn1}x{trait.bk1}x{trait.bk0max}" + f"{pad_suffix}" + ) + if trait.is_tuning_extra: + kname += " [ext]" + inners += FMHA_FWD_SPLITKV_ALL_API_INNER_DISPATCH.format( + F_arch=arch, + F_mode=MODE_MAP[trait.mode], + F_vlayout=LAYOUT_MAP[trait.vlayout], + F_pipeline_enum=PIPELINE_ENUM_MAP[trait.pipeline_tag], + F_logits=BOOL_MAP[trait.logits], + F_mask=get_mask_map(self.mask_impl)[trait.mask], + F_mask_check=get_mask_check_map(self.mask_impl)[trait.mask], + F_bias_check=BIAS_CHECK_MAP[trait.bias], + F_bias=BIAS_MAP[trait.bias], + F_lse=BOOL_MAP[trait.lse], + F_squant=BOOL_MAP[trait.squant], + F_pagedkv=BOOL_MAP[trait.pagedkv], + F_sink=BOOL_MAP[trait.sink], + F_scheck=trait.scheck, + F_skcheck=trait.skcheck, + F_dcheck=trait.dcheck, + F_dvcheck=trait.dvcheck, + F_spad=BOOL_MAP[trait.spad], + F_skpad=BOOL_MAP[trait.skpad], + F_dpad=BOOL_MAP[trait.dpad], + F_dvpad=BOOL_MAP[trait.dvpad], + F_bm0=trait.bm0, + F_bn0=trait.bn0, + F_bk0=trait.bk0, + F_bn1=trait.bn1, + F_bk1=trait.bk1, + F_bk0max=trait.bk0max, + F_hdim=hdim, + F_dtype=FWD_DTYPE_MAP[dtype], + F_bn1comb=trait.bn1comb, + F_kname=kname, + ) + per_hdim_case += FMHA_FWD_API_PER_HDIM_CASE.format( + F_if=if_(i_hdim), + F_hdim=hdim, + F_hdim_v=hdim, + F_inner_dispatch=indent(inners), + ) + per_dtypes += FMHA_FWD_API_PER_DTYPE.format( + F_if=if_(i_dtype), F_dtype=dtype, F_hdim_case=indent(per_hdim_case) + ) + per_arch += FMHA_FWD_API_PER_ARCH.format( + F_if=if_(i_arch), + F_arch=arch, + F_dtype_case=indent(per_dtypes), + ) + return FMHA_FWD_SPLITKV_ALL_API_FUNC_TEMPLATE.format( + F_func_name=func_name, F_dispatch=indent(per_arch) + ) + @dataclass class FmhaFwdSplitKVCombineTileSize: @@ -837,10 +983,15 @@ class KernelComponentFactoryGfx9(KernelComponentFactoryBase): @staticmethod def get_tuning_extra_tiles(dtype: str) -> dict: - """Additional tile sizes only available via tuning receipts (150, 250). - These tiles are NOT used by the heuristic dispatch path.""" + """Additional tile sizes merged for tuning/benchmarking receipts (0, 3, 150, 250). + For CK standalone (0/3) these are tagged is_tuning_extra and excluded from heuristic. + For AITER tuning (150/250) they participate in heuristic, selected via CSV.""" extra = {} if dtype in ["fp16", "bf16"]: + extra["128"] = [ + FmhaFwdTileSize( 16, 128, 32, 128, 32, 128, 1, 1, 1, 1, 1, 1, 16, 16, 16, 16, 16, 16, -1), + 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, 64, 256, 128, 256, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1), ] # fmt: skip @@ -910,21 +1061,30 @@ def get_fwd_splitkv_blobs( factories = get_factories_for_targets(targets, get_factory) - # Tuning receipts (150, 250) include extended tile sizes for CSV-driven selection - _TUNING_RECEIPTS = frozenset({150, 250}) - for factory, dtype in itertools.product(factories, FWD_DTYPE_MAP.keys()): d = factory.get_hdim_tile_size_dict(dtype) if d is None: continue # Build per-hdim tile lists: original (single) tile + optional tuning extras + # Receipts that include extended tuning tiles from get_tuning_extra_tiles(). + # - 150/250: AITER tuning receipts (tiles are part of heuristic, selected via CSV) + # - 0/3: CK standalone build (tiles compiled for fmha_fwd_splitkv_all() + # benchmarking but excluded from heuristic via is_tuning_extra flag) + _EXTRA_TILE_RECEIPTS = frozenset({0, 3, 150, 250}) + _MARK_AS_TUNING_EXTRA = frozenset({0, 3}) + # Track which tiles are standard so we can tag extras with is_tuning_extra. hdim_tiles = {} + standard_tile_ids = {} # hdim_str -> set of id() for standard tiles for hdim_str in d.keys(): hdim_tiles[hdim_str] = [d[hdim_str]] - if receipt in _TUNING_RECEIPTS and hasattr(factory, 'get_tuning_extra_tiles'): + standard_tile_ids[hdim_str] = {id(d[hdim_str])} + if receipt in _EXTRA_TILE_RECEIPTS and hasattr(factory, 'get_tuning_extra_tiles'): for hdim_str, extra_list in factory.get_tuning_extra_tiles(dtype).items(): if hdim_str in hdim_tiles: - hdim_tiles[hdim_str].extend(extra_list) + hdim_tiles[hdim_str] = sorted(hdim_tiles[hdim_str] + extra_list, key=lambda t: t.F_bm0) + else: + hdim_tiles[hdim_str] = sorted(extra_list, key=lambda t: t.F_bm0) + standard_tile_ids[hdim_str] = set() # all tiles are extra # for hdim_str, mode, mask, bias, lse in itertools.product(d.keys(), MODE_MAP.keys(), MASK_MAP.keys(), ["t", "f"], ["t", "f"]): for hdim_str, mode in itertools.product(d.keys(), MODE_MAP.keys()): hdim = int(hdim_str) @@ -940,6 +1100,7 @@ def get_fwd_splitkv_blobs( or pipeline.F_logits == "f" ): continue + is_extra = (receipt in _MARK_AS_TUNING_EXTRA) and (id(tile) not in standard_tile_ids.get(hdim_str, set())) k = Kernel( F_arch=factory.arch, F_idx=0, @@ -950,6 +1111,7 @@ def get_fwd_splitkv_blobs( F_pipeline=pipeline, mask_impl=mask_impl, ) + k._is_tuning_extra = is_extra if kernel_filter != "": if not fnmatch.fnmatch(k.name, kernel_filter): continue @@ -1070,7 +1232,10 @@ def write_single_kernel( def write_fwd_splitkv_api(api_pool: FmhaFwdSplitKVApiPool, autogen_dir: Path) -> None: - update_file(autogen_dir / FMHA_FWD_SPLITKV_API_FILENAME, api_pool.api) + update_file( + autogen_dir / FMHA_FWD_SPLITKV_API_FILENAME, + api_pool.api + api_pool.render_all("fmha_fwd_splitkv_all"), + ) def write_blobs( @@ -1139,6 +1304,7 @@ def write_blobs( dpad=kernel.F_pipeline.F_dpad, dvpad=kernel.F_pipeline.F_dvpad, bn1comb=combine_kernel.F_tile.F_bn1, + is_tuning_extra=getattr(kernel, '_is_tuning_extra', False), ) ) write_fwd_splitkv_api(api_pool, output_dir) diff --git a/example/ck_tile/01_fmha/example_fmha_fwd.cpp b/example/ck_tile/01_fmha/example_fmha_fwd.cpp index 122d232a1c..b41b8d14f6 100644 --- a/example/ck_tile/01_fmha/example_fmha_fwd.cpp +++ b/example/ck_tile/01_fmha/example_fmha_fwd.cpp @@ -119,7 +119,10 @@ auto create_args(int argc, char* argv[]) "", "Batch-mode only: per-batch effective seqlen for KV (exclude PAD).\n" "Comma-separated list of length 'b'. If empty, no override.") - .insert("init_sink", "0", "value to init the output tensor sink value for validation"); + .insert("init_sink", "0", "value to init the output tensor sink value for validation") + .insert("run_all_kernels", + "0", + "benchmark ALL kernel instances and print sorted results"); bool result = arg_parser.parse(argc, argv); return std::make_tuple(result, arg_parser); @@ -163,6 +166,7 @@ auto run(const ck_tile::ArgParser& arg_parser) std::string init_method = arg_parser.get_str("init"); uint32_t seed = arg_parser.get_uint32("seed"); int init_sink_value = arg_parser.get_int("init_sink"); + bool run_all_kernels = arg_parser.get_bool("run_all_kernels"); ck_tile::stream_config stream_config{nullptr, true, @@ -211,6 +215,7 @@ auto run(const ck_tile::ArgParser& arg_parser) do_validation, init_sink_value, stream_config, + run_all_kernels, json); } diff --git a/example/ck_tile/01_fmha/fmha_fwd.hpp b/example/ck_tile/01_fmha/fmha_fwd.hpp index 521f1e4738..1810015100 100644 --- a/example/ck_tile/01_fmha/fmha_fwd.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd.hpp @@ -17,6 +17,7 @@ #include #include #include +#include struct FmhaFwdFp32 { @@ -1674,6 +1675,7 @@ struct fmha_fwd_traits // TODO: padding check is inside this api }; float fmha_fwd(fmha_fwd_traits, fmha_fwd_args, const ck_tile::stream_config&); +std::vector> fmha_fwd_all(fmha_fwd_traits, fmha_fwd_args, const ck_tile::stream_config&); struct fmha_fwd_pagedkv_traits { @@ -1715,6 +1717,9 @@ struct fmha_fwd_splitkv_traits float fmha_fwd_splitkv(fmha_fwd_splitkv_traits, fmha_fwd_splitkv_args, const ck_tile::stream_config&); +std::vector> fmha_fwd_splitkv_all(fmha_fwd_splitkv_traits, + fmha_fwd_splitkv_args, + const ck_tile::stream_config&); struct fmha_fwd_appendkv_traits { diff --git a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp index 40b8006381..80fb24d204 100644 --- a/example/ck_tile/01_fmha/fmha_fwd_runner.hpp +++ b/example/ck_tile/01_fmha/fmha_fwd_runner.hpp @@ -10,14 +10,17 @@ #include "utils.hpp" #include "ck_tile/utility/json_dump.hpp" +#include #include #include #include #include #include +#include #include #include #include +#include #include #include #include @@ -249,6 +252,7 @@ fwd_result fmha_fwd_run(mode_enum mode, int do_validation, int init_sink_value, const ck_tile::stream_config& stream_config, + bool run_all_kernels = false, std::optional json = std::nullopt) { using TypeConfig = FmhaFwdTypeConfig; @@ -1564,6 +1568,355 @@ fwd_result fmha_fwd_run(mode_enum mode, return fmha_fwd(fmha_traits, fmha_args, sc); }; + // RAII guard for std::cout.rdbuf() redirect — restores original buffer + // even if the called function throws. + struct rdbuf_guard + { + std::ostream& os; + std::streambuf* orig; + rdbuf_guard(std::ostream& s, std::streambuf* buf) : os(s), orig(s.rdbuf(buf)) {} + ~rdbuf_guard() { os.rdbuf(orig); } + }; + + // --- run_all_kernels path: benchmark every matching instance --- + if(run_all_kernels) + { + std::vector> all_results; + std::string heuristic_full_kname; + std::string captured_kernel_log; // full kernel names from log_level=1 + + // Use the user's log_level but capture stdout so we can reformat it. + // This way log_level=1 (kname=1) produces full kernel names we can + // print one-per-line instead of comma-separated on one line. + ck_tile::stream_config all_sc{stream_config.stream_id_, + stream_config.time_kernel_, + stream_config.log_level_, + stream_config.cold_niters_, + stream_config.nrepeat_}; + +#if CK_TILE_FMHA_FWD_SPLITKV_API + const bool use_splitkv_all = (1 < num_splits || use_kvcache); + if(use_splitkv_all) + { + fmha_fwd_splitkv_traits fmha_splitkv_traits; + init_traits(fmha_splitkv_traits); + + fmha_fwd_splitkv_args fmha_splitkv_args; + init_args(fmha_splitkv_args); + + { + std::ostringstream log_oss; + rdbuf_guard guard(std::cout, log_oss.rdbuf()); + all_results = + fmha_fwd_splitkv_all(fmha_splitkv_traits, fmha_splitkv_args, all_sc); + captured_kernel_log = log_oss.str(); + } + + // Identify which kernel the heuristic would select + { + std::ostringstream oss; + rdbuf_guard guard(std::cout, oss.rdbuf()); + ck_tile::stream_config hsc{nullptr, + false, + /*log_level=*/1, + /*warmup=*/0, + /*repeat=*/1, + false}; + fmha_fwd_splitkv(fmha_splitkv_traits, fmha_splitkv_args, hsc); + std::string captured = oss.str(); + auto pos = captured.find("fmha_fwd_splitkv_"); + if(pos != std::string::npos) + { + // extract just the main kernel name (before the combine kernel name) + heuristic_full_kname = captured.substr(pos); + // the output has ", " after the main name + auto comma_pos = heuristic_full_kname.find(", "); + if(comma_pos != std::string::npos) + heuristic_full_kname.resize(comma_pos); + auto end = heuristic_full_kname.find_last_not_of(" \n\r\t"); + if(end != std::string::npos) + heuristic_full_kname.resize(end + 1); + } + } + } + else +#endif // CK_TILE_FMHA_FWD_SPLITKV_API + { + fmha_fwd_traits fmha_traits; + init_traits(fmha_traits); + + fmha_fwd_args fmha_args; + init_args(fmha_args); + + { + std::ostringstream log_oss; + rdbuf_guard guard(std::cout, log_oss.rdbuf()); + all_results = fmha_fwd_all(fmha_traits, fmha_args, all_sc); + captured_kernel_log = log_oss.str(); + } + + // Identify which kernel the heuristic would select + { + std::ostringstream oss; + rdbuf_guard guard(std::cout, oss.rdbuf()); + ck_tile::stream_config hsc{nullptr, + false, + /*log_level=*/1, + /*warmup=*/0, + /*repeat=*/1, + false}; + fmha_fwd(fmha_traits, fmha_args, hsc); + std::string captured = oss.str(); + auto pos = captured.find("fmha_fwd_"); + if(pos != std::string::npos) + { + heuristic_full_kname = captured.substr(pos); + auto end = heuristic_full_kname.find_last_not_of(" \n\r\t"); + if(end != std::string::npos) + heuristic_full_kname.resize(end + 1); + } + } + } + + if(all_results.empty()) + { + std::cout << " no matching instances" << std::flush << std::endl; + return fwd_result::no_instance; + } + + // Match short kname against the heuristic's full kname + // Short: "fmha_fwd_d128_bf16_batch_b32x32x128x128x32x128_npad" + // or: "fmha_fwd_d128_bf16_batch_b32x32x128x128x32x128_npad [ext]" + // Full: "fmha_fwd_d128_bf16_batch_b32x32x128x128x32x128_r1x1x1_..._vr_npad_nlogits_..." + auto is_heuristic = [&](const std::string& raw_name) -> bool { + if(heuristic_full_kname.empty()) + return false; + // Strip [ext] tag if present + std::string short_name = raw_name; + auto bracket = short_name.find(" ["); + if(bracket != std::string::npos) + short_name.resize(bracket); + auto last_sep = short_name.rfind('_'); + if(last_sep == std::string::npos) + return false; + std::string tile_prefix = short_name.substr(0, last_sep); + std::string pad_suffix = short_name.substr(last_sep); // e.g. "_npad" or "_ps" + // Check tile prefix matches + if(heuristic_full_kname.find(tile_prefix) == std::string::npos) + return false; + // Match pad suffix as a delimited token: suffix must be followed by '_' or + // end-of-string to avoid "_ps" matching "_psk" or "_psskddv" + auto pos = heuristic_full_kname.find(pad_suffix); + while(pos != std::string::npos) + { + auto end_pos = pos + pad_suffix.size(); + if(end_pos == heuristic_full_kname.size() || heuristic_full_kname[end_pos] == '_') + return true; + pos = heuristic_full_kname.find(pad_suffix, pos + 1); + } + return false; + }; + + std::cout << std::endl; + // sort by time (fastest first) + std::sort(all_results.begin(), all_results.end(), [](const auto& a, const auto& b) { + return a.second < b.second; + }); + std::cout << "[run_all_kernels] " << all_results.size() << " instance(s) benchmarked:" +#if CK_TILE_FMHA_FWD_SPLITKV_API + << (use_splitkv_all ? " (num_splits=" + std::to_string(num_splits) + ")" : "") +#endif + << std::endl; + // find max kernel name length for alignment + size_t max_kname_len = 0; + for(const auto& [kname, t] : all_results) + max_kname_len = std::max(max_kname_len, kname.size()); + + // compute index width for alignment + int idx_width = 1; + { + size_t n = all_results.size(); + while(n >= 10) + { + ++idx_width; + n /= 10; + } + } + + for(size_t i = 0; i < all_results.size(); i++) + { + const auto& [kname, t] = all_results[i]; + const float total_t = appendkv_ave_time + t; + const float tf = static_cast(flop) / 1.E9 / total_t; + const float bw = num_byte / 1.E6 / total_t; + std::cout << std::fixed << " [" << std::setw(idx_width) << (i + 1) << "] " << std::left + << std::setw(max_kname_len) << kname << std::right << ", " << std::setw(7) + << std::setprecision(3) << total_t << " ms" << ", " << std::setw(8) + << std::setprecision(2) << tf << " TFlops" << ", " << std::setw(8) + << std::setprecision(2) << bw << " GB/s" + << (is_heuristic(kname) ? " <-- heuristic" : "") + << std::endl; + } + + // When kname=1 (log_level >= 1), print full kernel names one per line + if(stream_config.log_level_ >= 1 && !captured_kernel_log.empty()) + { + std::cout << std::endl << "[run_all_kernels] full kernel names:" << std::endl; + // The captured log is comma-separated kernel names on one line. + // For splitkv, it alternates: main_kernel, combine_kernel, main_kernel, ... + // We skip combine kernels (contain "_combine_"). + std::string token; + std::istringstream iss(captured_kernel_log); + while(std::getline(iss, token, ',')) + { + auto start = token.find_first_not_of(" \n\r\t"); + if(start == std::string::npos) + continue; + auto end = token.find_last_not_of(" \n\r\t"); + auto display = token.substr(start, end - start + 1); + // Skip combine kernel names + if(display.find("_combine_") != std::string::npos) + continue; + if(!display.empty()) + std::cout << " " << display << std::endl; + } + } + + if(do_validation != 0) + { +#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); + } +#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, + 0, + 1, + false}; + + bool all_pass = true; + for(size_t i = 0; i < all_results.size(); i++) + { + const auto& [kname, _] = all_results[i]; + + if(setenv(force_kernel_env, kname.c_str(), 1) != 0) + { + std::cout << "[run_all_kernels] failed to set " << force_kernel_env + << " for kernel: " << kname << std::endl; + all_pass = false; + break; + } + + std::cout << "[run_all_kernels] verifying [" << (i + 1) << "/" << all_results.size() + << "] " << kname << std::endl; + + const auto verify_result = 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, + verify_sc, + false, + std::nullopt); + + if(verify_result != fwd_result::success) + all_pass = false; + } + + unsetenv(force_kernel_env); + + if(!all_pass) + return fwd_result::failure; + + std::cout << "[run_all_kernels] per-kernel verification passed for " + << all_results.size() << " kernel(s)" << std::endl; + + return fwd_result::success; + } + + return fwd_result::success; + } + // --- normal (heuristic) path --- + float fwd_ave_time = -1.0f; #if CK_TILE_FMHA_ENABLE_HEAD_GROUPING const bool allow_head_grouping = !i_perm && !use_kvcache && (num_splits <= 1) &&