From c53e12be46bd941eb04b6a472c10a6da070e4894 Mon Sep 17 00:00:00 2001 From: vivienfanghuagood <89012307+vivienfanghuagood@users.noreply.github.com> Date: Fri, 28 Nov 2025 16:11:45 +0800 Subject: [PATCH] Extend CK fmha_batch_prefill kernel coverage to head_dim=256 --- example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py | 1 + 1 file changed, 1 insertion(+) diff --git a/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py b/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py index bd63c5d78f..7a45687476 100644 --- a/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py +++ b/example/ck_tile/01_fmha/codegen/ops/fmha_batch_prefill.py @@ -579,6 +579,7 @@ def get_hdim_tile_size_dict(dtype: str) -> Optional[dict]: if dtype == "fp16" or dtype == "bf16": return { 128 : [FmhaFwdTileSize(128, 128, 32, 128, 32, 128, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1)], + 256 : [FmhaFwdTileSize(128, 128, 32, 256, 32, 256, 4, 1, 1, 4, 1, 1, 32, 32, 16, 32, 32, 16, -1)], } # fmt: skip else: return None