Support vLLM paged KV cache layouts on the gfx950 attention kernel - #1056
Support vLLM paged KV cache layouts on the gfx950 attention kernel#1056akii96 wants to merge 5 commits into
Conversation
635108b to
943152d
Compare
jhinpan
left a comment
jhinpan
left a comment
There was a problem hiding this comment.
The page-128 addressing change holds up on an exact-base runtime: the 10 focused cases, full attention file (103/103), and additional partial-page, runtime-pitch, split-K, and ragged-varlen probes pass on gfx950. Causal LPT was bit-identical and 1.039x faster at B=1, Sq=8192, Skv=32768, H=16, Hkv=1 in 30 interleaved samples. Requesting changes for the remaining host-boundary, public-flag, and test-coverage defects below; the PR body's end-to-end gsm8k/TTFT TODOs remain unverified here.
| return None | ||
| block_stride, row_stride, head_stride, elem_stride = (int(s) for s in k.stride()) | ||
| dense_row = num_kv_heads * head_dim | ||
| if elem_stride != 1 or row_stride < dense_row or block_stride != page_size * row_stride: |
There was a problem hiding this comment.
[blocking] The direct-stride gate accepts pitches that the generated layout and buffer descriptor cannot represent. A synthetic valid stride tuple with page_size=128 and row_stride=2**24 is accepted here; downstream that becomes page_elems=2**31 in fx.Int32(...) and a 2**32-byte page descriptor. Please bound the runtime pitch before launch (signed-int32 element extent and unsigned-32-bit descriptor byte span, both exclusive), falling back to a copy or rejecting it, and add boundary tests.
There was a problem hiding this comment.
Thanks, good catch, and your arithmetic was right. Fixed in e5ed3c5.
page_elems = page_size * row_stride
if page_elems >= 2**31 or page_elems * k.element_size() >= 2**32:
return NoneNone means dense, which routes into the .contiguous() copy that already existed, so it falls back rather than rejects. Boundary tests at 2**24 - 1, 2**24 and 2**25.
The byte bound can't fire today since this path is bf16/f16 only, but kept for a future 1 or 4 byte dtype.
It also tests the copy itself: removing both .contiguous() calls left every pitch test green, so the fallback was untested.
| # Reverse the causal q-block grid axis so the heaviest blocks issue first, | ||
| # shortening the makespan tail. Bit-identical, being a permutation of | ||
| # workgroup -> q-block. See flash_attn_utils._init_causal_lpt_order. | ||
| causal_lpt: bool = True, |
There was a problem hiding this comment.
[blocking] causal_lpt=False is not honored by two dualwave routes: _build_splitk and long _build_varlen neither accept nor forward this flag, so the underlying builder default silently re-enables LPT. That breaks the public override and prevents a trustworthy A/B for mixed varlen workloads, where reversing global grid_dim.y can place a short/decode sequence's only active block at the end. Thread the flag through both cached builder signatures/calls and add route-specific off/on tests; given the reported TTFT TODO, gate mixed varlen unless its default-on result is measured.
There was a problem hiding this comment.
Thanks, both fixed in e5ed3c5. _build_splitk and _build_varlen now take and forward the flag, and both call sites pass it.
You were right to ask for the measurement, it does regress. Sweeping B as a fraction of one sequence's q-block count put the sign change at 0.75 to 0.875 of it, at 8k/16k/32k alike. causal_lpt now defaults to auto, on while one sequence spans at least twice the batch's sequence count in q-blocks. True and False still force. The factor of two matters: admitting the whole range up to 1.0 still includes a 0.985x row.
Not the mechanism you suggested, though: I tried a per-sequence reversal so a short row keeps q-block 0, and timings were identical. It tracks sequence count against q-block count.
Route-specific off/on tests are there too, but they can't catch this on their own, since a route that ignores the flag returns identical bytes. It's pinned instead by asserting the two settings land on distinct lru_cache entries.
| return flydsl_flash_attn_func(q, k, v, causal=False, dualwave_swp_xcd_swizzle=flag).clone() | ||
|
|
||
|
|
||
| @_requires_gfx950 |
There was a problem hiding this comment.
[blocking] Inserting the new test here ends test_xcd_swizzle_is_bit_identical immediately after its local run definition. Its previous off/on launches, synchronization, and equality assertion are gone, so all four parametrized XCD cases pass without launching a kernel; I reproduced launch_calls=[] by instrumenting the test. Restore those three lines before this decorator while keeping the new LPT test.
There was a problem hiding this comment.
Thanks, and confirmed exactly as you described. Restored in e5ed3c5, with the LPT test kept after it.
Two things while in there: two definitions of test_causal_lpt_splitk_off_on_bit_identical had accumulated and Python kept only the last, now deduped. And I swept the file for tests with no assertion at all; none.
| _PAGED_PAGE_SIZE = 64 | ||
| # KV tile length the dualwave pipeline consumes per step; a page holds one or more. | ||
| _PAGED_TILE = 64 | ||
| _PAGED_PAGE_SIZES = (64, 128) |
There was a problem hiding this comment.
The PR body and commit describe any positive BLOCK_N multiple, while the public entry point accepts only 64 and 128 even though the traits support larger multiples. Please either cover and expose at least a >2-tile page, or narrow the external claim to the implemented 64/128 contract.
There was a problem hiding this comment.
Went with exposing it, in e5ed3c5. _PAGED_PAGE_SIZES is gone and the check is just a multiple-of-64 test now. The three paged equivalence tests run 64/128/256, so a 4-tile page is covered against the dense, contiguous and vLLM-view references. Body updated to match.
Also added a real vectorized rejection test. The old one built a linear cache at page_size 32, so that path was never hit.
bce2084 to
e5ed3c5
Compare
|
Thanks so much @jhinpan, genuinely useful review!
One number worth flagging is that mixed prefill plus decode batches measure 1.15x here against 1.45x before, so that row is where a caller's routing threshold should be fitted. Happy to iterate further on any of it! |
Causal work grows with the q-block index, so the natural order issues the heaviest block last. The reversal already existed but was reachable only from the fp8 builder; this shares it and enables it for bf16 behind a CAUSAL_LPT trait. Measured 2.0-8.9% on causal paged prefill at 32k, bit-identical since it is a permutation. The shared helper reverses over grid_dim.y instead of recomputing from seq_len, which is a bijection on dispatched indices by construction. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
… layouts The paged path assumed a page holds exactly one KV tile and that the cache is contiguous in the layout the kernel derives from its shape. Both hold for the configurations swept so far and neither holds for a vLLM KV cache serving MiniMax-M3, which uses page_size 128 against a 64-token tile and hands the kernel two strided views of one interleaved allocation. Three changes make that cache addressable without repacking it: * page_size may be any multiple of BLOCK_N. The staged block table is indexed per KV tile rather than per page, so the LDS window check counts tiles; a page spanning several tiles expands to its constituent tiles on stage. * The KV row stride is read off the tensor instead of inferred from the head count and head dim. A cache whose K and V are strided views of one buffer is passed through as-is; calling .contiguous() on it would cost a cache-sized copy and drop the K/V interleave. * A single-KV-head cache whose heads are not exactly D apart is accepted. With Hkv == 1 the head stride is unobservable, so requiring it to equal D rejected valid layouts for a dimension the kernel never indexes. The vectorized KV layout keeps the one-tile-per-page assumption in its address arithmetic, so it now rejects multi-tile pages explicitly rather than reading the wrong rows. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
…CD test Addresses review feedback on the paged-KV and causal-LPT commits. The direct-stride gate accepted pitches the generated layout cannot represent. A page descriptor spans page_size * row_stride elements, emitted as a signed int32 extent and a 32-bit byte span, so page_size 128 with row_stride 2**24 reached exactly 2**31 and wrapped, reading the wrong rows with no error. _paged_kv_row_stride now rejects such a pitch on the host and reports the cache as dense, which routes it through the pre-existing .contiguous() copy. causal_lpt was not honored on two routes. _build_splitk and _build_varlen neither accepted nor forwarded it, so the builder default silently re-enabled LPT and the public override was a no-op there. Both now take and forward the flag. The reviewer asked for mixed varlen to be gated unless default-on was measured. Measuring it showed the concern was justified: varlen sizes grid_dim.y from the batch-wide max seqlen, and once the batch holds enough sequences relative to one sequence's q-block count, reversing that axis fights the scheduler rather than shortening the causal tail. causal_lpt now defaults to None, meaning auto, with True/False still forcing either way so an A/B stays trustworthy. The gate admits a batch only while one sequence spans at least twice as many q-blocks as the batch has sequences. The factor of two is not slack: sweeping B as a fraction of the q-block count at max_seqlen_q 8k/16k/32k puts the crossover at 0.75-0.875 of it, consistently across all three, so admitting the whole range would keep LPT on over a band measuring 0.974-0.998x. Half is the widest round threshold whose entire range is a win, measured 1.006-1.075x. Measured with an interleaved off/on/off harness carrying an off-vs-off control arm, so 1-2% effects sit above a 0.1-0.5% noise floor rather than inside it. Dense and split-K keep LPT on by default, where it measured neutral to positive across B and S, so only varlen consults the gate. A new test had been inserted directly after the local run() definition in test_xcd_swizzle_is_bit_identical, absorbing its three closing lines. All four parametrized cases were passing without launching a kernel. Restored. page_size is widened to any multiple of the 64-token KV tile, which the traits layer already supported, so the implementation now matches the documented contract instead of narrowing the text to 64/128. Tests: 121 pass cold-cache on gfx950, up from 103. New coverage for the pitch bounds at the exact 2**24 boundary, causal_lpt on the split-K and long-varlen routes, the auto-gate resolution, page_size 256 against the dense, contiguous and vLLM-view references, and the vectorized layout's multi-tile rejection, which the PR body claimed was covered but was not: that test built a linear cache at a sub-tile page_size. Two of those tests are written against mutations that the obvious version misses. Asserting that every cached dualwave builder takes causal_lpt does not catch a call site dropping the argument, since both builders default it to True and LPT is a permutation, so outputs are unchanged; the builders are lru_cached on the flag, so the surviving signal is that two calls differing only in it must produce two cache entries. Likewise the pitch bound is unit-tested on stride tuples that no real cache can reach, leaving the consequence untested, so the fallback is exercised by stubbing _paged_kv_row_stride to refuse and requiring a real packed vLLM-view cache to still match both references. Removing the .contiguous() calls fails that and nothing else. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com>
e5ed3c5 to
b83922b
Compare
jhinpan
left a comment
jhinpan
left a comment
There was a problem hiding this comment.
Reviewed b83922b in three passes: prior-finding closure; paged addressing/API/cache-key audit; and tests/CI. On a locked MI355X (gfx950), 34 focused paged/LPT/XCD cases passed, the full attention file passed 145/145, and additional page_size=192/320 partial-page probes passed for both single-split and split-K=3 (max abs error <= 5.13e-4 versus the fp32 reference). check_repo and Python style also pass. One descriptor-boundary hole remains: dense layouts return before the new page-span check, so the newly advertised unrestricted 64-multiple page_size can still create a >=231-element / >=232-byte page descriptor and silently wrap. Requesting changes until the dense effective pitch is bounded (or oversized pages are rejected) with a regression test.
| return None | ||
| if num_kv_heads > 1 and head_stride != head_dim: | ||
| return None | ||
| if row_stride == dense_row: |
There was a problem hiding this comment.
[blocking] This early return bypasses the descriptor-span check below for dense caches. For example, Hkv=64, D=128, page_size=262144 gives dense_row=8192 and page_elems=231 (232 bytes); this helper returns None here, the caller treats that as the dense/default stride, and k.contiguous() is a no-op for an already-contiguous cache, so init_descriptors still casts the overflowing page extent and builds the wrapped descriptor. This became reachable when the PR widened page_size to any multiple of 64. Please validate the effective dense pitch before this return (and reject when even a dense copy cannot fit), with a dense-boundary test in addition to the strided cases.
There was a problem hiding this comment.
Thanks @jhinpan , confirmed, and your example lands on exactly 2**31. Fixed in 05d6f49.
Checked at the entry point rather than before this return, since None here means "dense" and the helper can't refuse. Rejecting is right anyway: dense is already the narrowest pitch, so a copy has nothing to shrink into.
dense_page_elems = page_size * num_kv_heads * D
if not _paged_page_descriptor_fits(dense_page_elems, k.element_size()):
raise NotImplementedError(...)Both bounds share that predicate now. It has to be the resolved num_kv_heads, not k.shape[2]: DEFAULT_STRIDE_KV_N comes from the former, so bounding the cache's count leaves the same wrap reachable by asking for more KV heads than it holds.
Boundary test on the predicate, plus a launch-path one for an oversized page and head count.
The page-span check sat after the dense early return, so it only ever saw a pitch read off the tensor. Harmless while page_size was 64 or 128, but widening it left the dense pitch unbounded past 2**31 elements / 2**32 bytes, wrapping silently. Both bounds now share one predicate, consulted at the entry point too. Dense rejects rather than falling back, having no narrower layout to copy into, and uses the resolved num_kv_heads, which is what DEFAULT_STRIDE_KV_N is built from. 152 pass cold-cache on gfx950, up from 145. Signed-off-by: Aakif Nawaz <aakif.nawaz@amd.com> Co-authored-by: Cursor <cursoragent@cursor.com>
|
@yanguahe take a look? |
|
aiter | 0.9469 | 0.9477 accuracy drop 1%? |
There was a problem hiding this comment.
The paged addressing and the review follow-ups on pitch bounds, causal_lpt plumbing, and page_size look in good shape. The remaining ask is to pin down the gsm8k gap and to narrow the in-place layout claim.
gsm8k (−1.4pp)
The table is from an older e2e point (943152d + vllm-project/vllm#53335). Please re-run gsm8k on the kernel after the pitch / LPT / page_size review fixes, not that earlier commit.
Before treating the 1.4pp as a kernel issue, please also show:
- Routing. On the 8-shot gsm8k prompts, how many steps actually entered FlyDSL vs fell through to AITER? If almost none hit FlyDSL, this drop is not this kernel.
- Layer type. Full softmax attention layers may use dualwave; MiniMax sparse layers must stay on the existing AITER/sparse path. A sparse layer sent to dualwave would drop accuracy and would be independent of the paged layout work.
- Repeatability. Re-run the same gsm8k setup 2–3 times and check whether the gap stays above ~1pp; or run once more with
dualwave_swp_lazy_rescale=Falseonly.
A single run at ~0.6pp binomial stderr is not enough to decide. The near-tie token note in #53335 is useful, but it does not replace a re-run on current HEAD.
PR description: in-place KV contract
Please do not describe this as in-place for an arbitrary vLLM cache. The kernel can consume the cache without a copy only when:
- layout is
linearandpage_sizeis a multiple of 64 (vectorized remainspage_size=64only); - and either
Hkv == 1, or the heads are already packed NHD sohead_stride == D.
GQA + HND still goes through .contiguous(). That is fine for MiniMax-M3 TP=4 (Hkv=1); it is not a general vLLM zero-copy promise. Please state that bound in the PR body (and the flydsl_flash_attn_func docstring if you touch it).
4 participants
Motivation
vLLM cannot use the gfx950 attention kernel today. The kernel and the cache disagree about layout in three ways:
page_size128, which is two.Dapart. With one head there is nothing to space, so the check rejects valid layouts.Any one of them forces the caller to repack the cache before every call: a full copy, and it loses the K/V interleave, which is most of the reason to call this kernel.
Technical Details
The kernel now addresses a vLLM cache in place, along three axes.
Page size. A page may hold any multiple of
BLOCK_Ntokens. The staged block table is indexed per tile rather than per page, so a larger page expands to the tiles it covers. The vectorized layout hard codes one tile per page in its address arithmetic and rejects multi tile pages rather than reading the wrong rows.Stride. The KV row stride is read off the tensor rather than inferred from head count and head dim, so two strided views of one interleaved buffer are addressable as they are. The pitch is bounded on the host, since a page descriptor spans
page_size * row_strideelements as a signed int32 and wraps beyond it. A cache too wide to address is copied instead.Head stride. A single KV head no longer has to sit exactly
Dfrom the next, a dimension the kernel never indexes.Separately, causal q-blocks issue heaviest first. Work grows with the q-block index, so the natural order leaves the longest until last and everything waits on it. The reversal is a permutation, so output is bit identical. It only pays when one sequence dominates the dispatch, since varlen sizes
grid_dim.yfrom the batch wide max seqlen, socausal_lptdefaults to auto: on while one sequence spans twice the batch's sequence count in q-blocks, off past that.TrueandFalseforce. Dense and split-K keep it on.Test Plan
15 new tests in
tests/kernels/test_flash_attn_fwd.py, 27 cases after parametrization, gated on gfx950 like the rest of the file. Collected count goes 93 to 121.causal_lpt, and every call site passes itTwo of these are pinned by mutation rather than assertion. LPT is a permutation and the builders default it on, so a route that ignores the flag still returns identical bytes and ordinary off/on tests pass. Refusing a pitch likewise says nothing about the copy behind it. Deleting the call site argument fails only the builder cache test, and removing the fallback copy fails only the copy path test.
Separately, a standalone harness times this kernel against the AITER unified attention kernel vLLM uses today, on the shapes a MiniMax-M3 server produces at TP=4 (16 q-heads, 1 KV head, head dim 128, page_size 128, bf16). Both arms are checked against a float32 per sequence reference before anything is timed.
Test Result
Re-measured on MI355X,
flydsl 0.3.2, AITERf93e6cf, page_size 128, 16 q-heads, 1 KV head, D=128, bf16. Both arms are checked against a float32 reference before timing, and host enqueue is amortised over 20 calls:flydsl_flash_attn_funcspends about 0.13ms per call on the host against AITER's 0.03ms, so one event pair per call charges that idle GPU time to the kernel and understates this one by 25 to 35%.pre e5ed3c5in parentheses.Prefill, this kernel vs AITER UA:
Ragged and mixed batches at ctx 8k, which is what chunked prefill actually schedules. Times in ms:
Decode is the other side of that.
BLOCK_Mis 256, so a single token row fills 1/256 of its tile while still streaming that row's whole KV:This is a prefill kernel, and the gap widens with both batch size and context. A caller should route decode to AITER and keep this kernel for prefill shaped work. The same effect sets the floor on the prefill side: batches with many decode rows mixed in fall below 1.0x, measured 0.85x for a 4096 chunk sharing a batch with 16 decode rows. Below about 0.15ms of GPU work the host path dominates instead, since
flydsl_flash_attn_funcspends about 0.13ms per call enqueueing, which is why a full 1024 token prefill measures 0.28x and only crosses 1.0x at ctx 4096.The decode mixed rows are the one place the two runs disagree beyond noise, 1.15x against 1.45x, so a routing threshold should be fitted on them rather than on the pure prefill rows.
Full attention suite passes, 121 tests including the 27 new ones, from a cold JIT cache.
End to end vLLM serving on
amd/MiniMax-M3-MXFP4, TP=4, 133k max model len, against a server running only AITER UA, measured at943152dwith caller side routing from vllm-project/vllm#53335. Not re-run: the kernel level numbers above hold up after the review changes, so re-evaluating the vLLM side is deferred until this lands, and these uplifts are expected to hold at a similar level. Output token throughput, mean TPOT, P99 TTFT:The 8k rows are flat because the caller keeps this kernel out of those batches: at 8k a scheduler step is one chunk plus more decode rows than this kernel handles well.
gsm8k, full set, 8 shot, same commit: