From dbd9809fa0f350bc9cee4cec91c01cdd44bc146c Mon Sep 17 00:00:00 2001 From: vivienfanghuagood <89012307+vivienfanghuagood@users.noreply.github.com> Date: Fri, 30 Jan 2026 11:18:20 +0800 Subject: [PATCH] Extend CK fmha_batch_prefill kernel coverage to head_dim=256 (#3328) Co-authored-by: Po Yen Chen Co-authored-by: Illia Silin <98187287+illsilin@users.noreply.github.com> [ROCm/composable_kernel commit: f3d8b7210fb99827bcb1d1bdaf9672b3ae8fb209] --- 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 9a2d727253..42f686e0c0 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 @@ -630,6 +630,7 @@ class KernelComponentFactory: if dtype in ["fp16", "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 elif dtype in ["fp8bf16"]: return {