ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-14 11:07:44 +00:00
Code Issues Packages Projects Releases Wiki Activity

add run_all_kernels benchmarking mode with extended tuning tiles

Browse Source
This commit is contained in:
Mohsen Saffari
2026-04-17 12:03:58 +00:00
parent cc2963c884
commit a30e0c5cce
5 changed files with 723 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

193
example/ck_tile/01_fmha/codegen/ops/fmha_fwd.py
View File

@@ -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<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;
[[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_<trait_, {F_arch.tag}>(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)

186
example/ck_tile/01_fmha/codegen/ops/fmha_fwd_splitkv.py
View File

@@ -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<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;
[[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<OaccDataType, {F_bn1comb}>::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_<traits_, traits2_, {F_arch.tag}>(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_<traits_, traits2_, {F_arch.tag}>(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)

7
example/ck_tile/01_fmha/example_fmha_fwd.cpp
View File

@@ -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);
}

5
example/ck_tile/01_fmha/fmha_fwd.hpp
View File

@@ -17,6 +17,7 @@
#include <type_traits>
#include <utility>
#include <variant>
#include <vector>
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<std::pair<std::string, float>> 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<std::pair<std::string, float>> fmha_fwd_splitkv_all(fmha_fwd_splitkv_traits,
fmha_fwd_splitkv_args,
const ck_tile::stream_config&);
struct fmha_fwd_appendkv_traits
{

353
example/ck_tile/01_fmha/fmha_fwd_runner.hpp
View File

@@ -10,14 +10,17 @@
#include "utils.hpp"
#include "ck_tile/utility/json_dump.hpp"
#include <algorithm>
#include <array>
#include <cstdlib>
#include <cstring>
#include <functional>
#include <cmath>
#include <iomanip>
#include <numeric>
#include <optional>
#include <ostream>
#include <sstream>
#include <string>
#include <tuple>
#include <utility>
@@ -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<std::string> json = std::nullopt)
{
using TypeConfig = FmhaFwdTypeConfig<DataTypeConfig>;
@@ -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<std::pair<std::string, float>> 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 ", <combine_name>" 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<float>(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<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);
}
#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<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,
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) &&