mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-06-29 11:16:59 +00:00
c9bc5350c8
Add a `bool kEnablePaging_` non-type template parameter on
UnifiedAttentionPipeline (default true preserves the paged behaviour).
When false, `refresh_*_offsets` collapses to a single per-row
`logical_token * row_stride` imad — no block_tables fetch, no
/ % page_size arithmetic, no Tier 0 scalar-promote, no Tier 2 LDS-cache
populate. The host selects between paths via a new
`args.kv_contiguous` runtime flag plumbed through dispatch_variant<V>.
Twelve new prefill instances pin EnablePaging=false:
prefill_d{64,128} × {fp16, bf16, fp8} × {mask, nmask}
Decode variants stay on the paged path — callers without a KV cache
don't have decode workloads, and the binary-size cost isn't justified.
Measured impact on the same physical K/V memory (sq=1×4096, causal,
page_size=32 paged baseline, MI355, n=30 iters):
variant sk paged contig Δ
prefill_d64 bf16 4096 0.274 0.227 -17.1 %
prefill_d64 bf16 16384 1.529 1.198 -21.6 %
prefill_d64 bf16 32768 3.218 2.505 -22.1 %
prefill_d64 fp8 4096 0.299 0.235 -21.4 %
prefill_d64 fp8 16384 1.489 1.150 -22.7 %
prefill_d64 fp8 32768 3.054 2.386 -21.9 %
prefill_d128 bf16 4096 0.493 0.397 -19.3 %
prefill_d128 bf16 16384 2.638 2.224 -15.7 %
prefill_d128 bf16 32768 5.731 4.598 -19.8 %
prefill_d128 fp8 4096 0.476 0.341 -28.3 %
prefill_d128 fp8 16384 2.416 1.792 -25.8 %
prefill_d128 fp8 32768 4.973 3.727 -25.0 %
prefill_d128 fp8 at -28 % is the single biggest UA optimisation
measured to date — bigger than Tier 0 (-12 %), Tier 2 (-5 %), and the
Tier-3 d=64 fp8 win (-16 %).
Correctness validated by bit-exact comparison against the paged
instance with page_size=32 and identity block_tables on 48 shape ×
dtype × mask combinations.
Co-authored-by: Cursor <cursoragent@cursor.com>