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 kv_offset_array_transform to batch async for page size 16

Browse Source
This commit is contained in:
ltqin
2025-09-12 08:59:05 +00:00
parent f3e775c800
commit 667ed5825c
3 changed files with 136 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py
View File

@@ -63,6 +63,7 @@ using fmha_trait_{F_idx} = ck_tile::TileFmhaTraits<{F_spad},
{F_dropout},
{F_squant},
{F_occupancy},
false,
{F_sglang_layout}>;
using fmha_variant_{F_idx} = ck_tile::ComposedAttention<{F_logits} * ck_tile::LOGITS_SOFT_CAP, CK_TILE_FMHA_FWD_FAST_EXP2>;
@@ -517,8 +518,6 @@ class KernelComponentFactory:
for logits, mask, bias, lse, dropout, sglang in itertools.product(["t", "f"], get_mask_map(mask_impl).keys(), BIAS_MAP.keys(), ["t", "f"], ["t", "f"], ["t", "f"]):
pipelines.append(FmhaFwdPipeline('qr_async', 'row', 't', 'f', 't', 't', logits, bias, lse, dropout, squant, mask, sglang))
pipelines.append(FmhaFwdPipeline('qr_async', 'row', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask, sglang))
# pipelines.append(FmhaFwdPipeline('qr_async', 'col', 't', 'f', 't', 't', logits, bias, lse, dropout, squant, mask))
# pipelines.append(FmhaFwdPipeline('qr_async', 'col', 't', 't', 't', 't', logits, bias, lse, dropout, squant, mask))
else:
assert False
return pipelines