ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-03 05:37:34 +00:00
Code Issues Packages Projects Releases Wiki Activity

Expand medium tier dispatch threshold for multi-seq prefill

Browse Source 
Raise the medium tier boundary from avg_q <= 32 to avg_q <= 128.
This routes prefill shapes with 76-253 sequences (avg query length
32-108) to the 4-warp kBlockM=128 kernel instead of the 8-warp
kBlockM=256 kernel.

Results on affected shapes (via aiter):
  76-seq prefill:  1.44x vs Triton -> 1.03x (near-tied)
  157-seq prefill: 1.61x -> 1.02x
  181-seq prefill: 1.74x -> 1.00x (tied)
  229-seq prefill: 1.77x -> 1.05x

Made-with: Cursor
This commit is contained in:
Amir Ghamarian
2026-03-28 11:28:21 +00:00
parent b3877ba60a
commit 5f9b03746d
1 changed files with 1 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
example/ck_tile/42_unified_attention/unified_attention.cpp
View File

@@ -50,7 +50,7 @@ static tile_tier select_tile_tier(const unified_attention_args& args)
return tile_tier::small; // pure decode: 2 warps, kBlockM=64
const index_t kBlockQ_medium = 128 / args.num_queries_per_kv; // kBlockQ for 4-warp kernel
if(avg_q <= kBlockQ_medium * 2)
if(avg_q <= kBlockQ_medium * 8)
return tile_tier::medium; // many short seqs: 4 warps, kBlockM=128
return tile_tier::large; // long prefill: 8 warps, kBlockM=256