mirror of
https://github.com/ROCm/composable_kernel.git
synced 2026-07-12 18:17:58 +00:00
2645149bbf
The decode pipeline's Tier-2 LDS page-table cache bulk-loads
block_tables_ptr_[block_table_offset + i] into block_tables_lds[i] at
kernel entry so subsequent refresh_*_offsets calls can resolve the
phys_page via a one-cycle ds_read_b32 broadcast instead of a scoreboarded
s_load_dword. Before this commit the bulk load covered an *absolute*
prefix [0, num_blocks) of the batch's page table — i.e. every page index
this CTA could ever produce, even the ones earlier splits skipped over.
That conservative load (1) wasted bulk-load bytes on splits 1+ and (2)
made the 4096-entry static cap a *total-sequence* limit rather than a
per-split limit, so the long-context decode
b=64 sq=1 sk=128000 hq=64 hk=8 d=128 bf16 page=16
(total_kv_pages = 8000, num_splits = 4, last split window = [6000, 8000))
tripped `assert(split_end_page <= kPageTableLdsEntries)` on splits 2/3.
The wrapper's split-KV scheduler would otherwise keep each CTA's working
set well under the cap (2000 pages here vs 4096 cap) — the assert was
firing on data the kernel doesn't actually need.
Make the cache per-split-window: bulk-load only
[split_start_page, split_end_page) where
split_start_page = ⌊num_blocks_start · kPageBlockSize / page_size⌋,
and shift refresh_{k,v}_offsets' lookup by split_start_page so the LDS
index stays in [0, split_window_pages). split_start_page is a
kernel-entry constant, so the subtract folds into the s_load_dword's
immediate offset on every refresh call — no per-iteration cost.
On the prefill path (num_blocks_start == 0) split_start_page == 0 and
the change collapses to the original absolute-indexed load, so prefill
codegen is bit-identical.
Verified on gfx950 (256 CUs):
* Originally failing shape now passes correctness and runs at
5713 GB/s (vs Triton 4402 GB/s).
* Default regression (b=128 sq=1 sk=16384 hq=64 hk=8, 3 runs × 4
d/dtype combos) all PASS at both block_size=32 (baseline path)
and block_size=16; bandwidth unchanged vs pre-patch baseline.
* Long-context decode shapes that previously couldn't run at
page_size=16 — sk ∈ {65K, 262K, 524K} — now all PASS at
~5350-5690 GB/s.
* Prefill (b=1 sq=4096) PASS, perf unchanged.
Co-authored-by: Cursor <cursoragent@cursor.com>