Conversation
…tten variant scaffolding
Address PR review feedback (ROCm#4136): ChunkK was hardcoded to _BLOCK_KV in the public flydsl_fp8_paged_mqa_logits call despite compile_fp8_paged_mqa_logits already accepting block_kv. Now ChunkK is forwarded, defaulting to _BLOCK_KV (128) and validated to be a multiple of MFMA_N before dispatch.
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
|
fhuizing
marked this pull request as ready for review
July 14, 2026 12:41
Contributor
|
Yields a consistent ~1.5% E2E improvement for GLM5.2 MXFP4 without accuracy impact. 
|
Kernel: extend fp8_paged_mqa_logits from KVBlockSize==1 to KVBlockSize>=1
(block-flat co-packed layout: KVB fp8 key rows grouped first, then KVB f32
scales). Per-lane gather resolves each logical KV column independently
through the block table (kv_indices[b, col // KVB] -> block, col % KVB ->
token), so KVB is a compile-time constant that folds to the per-token slot
layout at KVB==1. i32 gather-offset ceiling now accounts for KVB*index_dim.
Test: add --kv-block-size sweep axis (1, 64) and reference block-flat
dequant; convert to a pure correctness test -- drop run_perftest / the
fn() wrapper for a direct kernel call, and remove the timing-derived
us/TFLOPS output columns (gate stays exact -inf mask + calc_diff < 1e-3).
Verified on gfx950 from a fresh compile cache across KVB {1,64}, dtypes
{fnuz,fn}, SplitKV {auto,1,4}, heads/head_dim {64,128}, short/long context.
…ion bench Support the production shuffle_weight(16x16) KV data layout in the FlyDSL paged fp8 MQA-logits kernel (Preshuffle=True, KVBlockSize % 16 == 0), making it a true drop-in for deepgemm_fp8_paged_mqa_logits(Preshuffle=True). Only the intra-block key-byte offset changes; the block-table gather, co-packed/token-ordered scale gather, MFMA, and causal mask are unchanged. - _mqa_logits_common: add load_pack_v8i32_preshuffle (closed-form offset into the shuffle_weight(16x16) data section; 8-byte halves stay contiguous). - fp8_paged_mqa_logits: thread a compile-time preshuffle flag through the kernel builder (with a _ps name tag), compile_* cache key, and the public API; replace the NotImplementedError guard with the KVBlockSize % 16 == 0 assert. - test: add a preshuffle path (oracle from the unshuffled cache, kernel fed the shuffled copy) and a small dedicated case set (kvb 16/64, D 64/128, FN patch). - bench_flydsl_vs_gluon: note the A/B is fixed at KVBlockSize=1 (Gluon faults at KVBlockSize>1); KVBlockSize>1 / Preshuffle coverage lives in the flydsl test. - add bench_flydsl_paged_mqa_logits.py: standalone single-config bench (no Gluon) for rocprofv3 isolation profiling, covering both block-flat and preshuffle paths. Validated on gfx950: block-flat and preshuffle correctness pass the unchanged gate (exact -inf mask + calc_diff < 1e-3); rocprofv3 traces the _ps kernel.
# Conflicts: # aiter/ops/flydsl/__init__.py # aiter/ops/flydsl/kernels/fp8_mqa_logits.py
…kage Upstream relocated the mqa-logits kernels into aiter/ops/flydsl/kernels/mqa_logits/. Move the paged FP8 kernel and its shared helper module there too for consistency: kernels/fp8_paged_mqa_logits.py -> kernels/mqa_logits/fp8_paged_mqa_logits.py kernels/_mqa_logits_common.py -> kernels/mqa_logits/_mqa_logits_common.py Fix the sibling .tensor_shim imports to ..tensor_shim (now one level up) and update the public __init__ import path. No behavior change; public API (flydsl_fp8_paged_mqa_logits) is unchanged.
Both paged benchmarks sized the block pool from max_model_len and then handed blocks out with `pool[counter % num_blocks]`, so the block table wrapped and every sequence reused the same physical blocks. At B=16, kv_len=32768 that is a 1024-block pool serving 8192 block references: an 8.65 MB KV working set standing in for 69 MB, small enough to live in cache. Every number measured against it was therefore a cache benchmark. That flattered the kernel by ~13% on its own, and distorted the head-to-head against production Gluon unevenly -- given the honest footprint FlyDSL loses 31-51% while the reference loses only 1-4%, since the reference is slow enough that memory is not its constraint. The reported speedup falls from 1.3-4.2x to 1.7-3.1x. Size the pool from the blocks the batch actually needs, build the block table vectorized rather than one element per GPU write, and report the touched footprint alongside the requested bytes. --pool-blocks shrinks the pool deliberately for cache-residency studies, and the vs-gluon bench now fails with a clear message when a shape would exceed the kernel's 2^31 i32 gather-offset limit instead of tripping an assert inside the launch path. Co-authored-by: Cursor <cursoragent@cursor.com>
…apes) SplitKV was auto-computed with the production Triton host formula, which ignores how many waves a CTA contains and so drifts with shape: it lands 2.5 waves per SIMD at B=16 but 10 at B=128. Resident wave count is the lever that matters here -- this kernel is bound by memory throughput rather than per-wave latency, so too few waves leaves the memory system idle and too many just queue. Measured across six shapes (batch 4..128, kv_len 4k..128k) the curve is flat between 2.0 and 2.5 waves/SIMD and rises steeply either side: 1.0/SIMD costs 33.9us and 4.0/SIMD 29.4us against 27.6us at the optimum. Target WavePerEU waves per SIMD directly, accounting for the CTA's wave count, and default to paged_w1. At equal resident wave count, splitting a CTA across 2 or 4 waves is slower on every shape measured: a 1-wave CTA schedules more freely, and the waves have no work to share anyway since each owns a disjoint set of KV columns with no barrier or LDS between them. Worth 4.6-14.4% per shape, 9.3% aggregate, and it removes the shape sensitivity -- all six shapes now land within 1% of 27.7us for the same total KV work, where the old defaults spread them over 29.2-32.5us. Verified against the full 870-case correctness sweep (max error 5.3e-4, tolerance 1e-3); the block-flat KVBlockSize=1 path is unchanged. Two public-surface notes for reviewers. This deliberately diverges from the deepgemm_fp8_paged_mqa_logits host formula it previously mirrored. And WavePerEU changes meaning, from a multiplier in the old formula to the target waves per SIMD; its default of 2 is where the curve bottoms out. FP8_PAGED_MQA_LOGITS_DEFAULT_VARIANT changes paged_w4 -> paged_w1. Co-authored-by: Cursor <cursoragent@cursor.com>
The kernel was written against raw MLIR builders -- 58 _to_raw() calls and 65 arith.* calls -- for arithmetic FlyDSL already overloads. The worst case was the store predicate, 17 lines of nested arith.andi/arith.cmpi to say `(lane_div_N == 0) & (col <= q_limit)`. Route unsigned division and modulo through new udiv/umod helpers rather than the bare `//` and `%` operators: fx.Int32 is signed, so `//` lowers to floordivsi, which expands to several instructions to get negative-operand rounding right. Every quantity divided here is non-negative, so going via fx.Uint32 keeps the single-instruction divui/remui the hand-written calls emitted. Verified codegen-neutral: the ATT-decoded ISA is byte-identical to the previous commit, 655 instructions in the same order. 870/870 correctness cases pass at 5.3e-4, timing is unchanged at 28.0us, and ruff reports the same findings as before. Co-authored-by: Cursor <cursoragent@cursor.com>
_mqa_logits_common carried three helpers nothing referenced: fn_to_fnuz_i64 (45 lines -- the dense kernel defines its own local copy), load_pack_i64 (likewise), and i32_add. Removing them is the bulk of the shrink. Comments and docstrings were 30% of the two files, much of it restating the code or duplicating the PR description. Trimmed to what a reader cannot derive: the co-pack layout contract, the ReLU scale-hoist identity, the i32 gather ceiling, the D-reg-to-head mapping, and the warning that a Python `if` on KVB would become a runtime scf.if and drop branch-local addresses. 932 -> 770 lines across the two files. ISA byte-identical to the previous commit (655 instructions), 870/870 correctness at 5.3e-4, 27.8us unchanged, no new ruff findings. Co-authored-by: Cursor <cursoragent@cursor.com>
Replaces the hand-built scf.ForOp plus `with ir.InsertionPoint(body)` and trailing scf.YieldOp with the scf.for_ generator, which owns the insertion point and terminator. Same pattern the fmha_gfx1250 kernel already uses. Only two lines, but it removes the last raw MLIR region plumbing from the kernel's main loop. The masked store keeps scf.IfOp: FlyDSL exposes no `if_` helper, and switching to a predicated store would change codegen. ISA byte-identical (655 instructions), 870/870 correctness at 5.3e-4, 27.7us. Also investigated expressing the preshuffle as a FlyDSL layout, as examples/04-preshuffle_gemm.py does in two lines. Not adopted: that example indexes a whole contiguous tensor, whereas our access is a paged gather off a per-lane runtime base. GTensor._linear_offset is element-granular, and while every preshuffle stride happens to be dword-aligned, a 5-D view would just replace one multiply-add chain with another for about two lines. It only pays off alongside FlyDSL-native tensors and copy atoms. Co-authored-by: Cursor <cursoragent@cursor.com>
The paged gather hand-computed byte offsets and branched on Preshuffle to
pick between two of them. Both layouts are actually the same coordinate
space -- (block, token//16, token%16, K-step, sub-group, lane half) -- and
differ only in the stride vector:
block-flat (blk, 4D, D/4, 16, 4, 2) token c at c*D
shuffled (blk, 4D, 4, 256, 64, 2) token c at (c//16)*16D + (c%16)*16
So the choice moves into make_kv_key_view and the kernel loses its `if
preshuffle` branch entirely; one load_pack_kv serves both. The co-packed
scale tail becomes a (block, token) view whose base absorbs the KVB*D key
offset, replacing the manual scale_byte arithmetic.
load_pack_kv slices the view down to the `sub` axis once rather than
indexing all six coordinates per load. Without that the invariant part of
the address is re-derived in each of the four loads and the kernel grows 49
instructions; with it, 14.
Trade, measured rather than assumed: the kernel file drops 15 lines and a
branch, the shared module gains 21 for the two view constructors, so the
pair is +6 overall. Codegen is 14 instructions larger (669 vs 655) because
the hand-written version exploited disjoint bit-fields to fuse addresses
into v_or3_b32/v_add3_u32 that the general dot product does not reach for.
Speed holds on every path (3 reps each, layout vs previous): preshuffle
27.58 vs 27.83us, block-flat KVB=64 9.78 vs 9.98us, KVB=1 10.20 vs 10.53us.
870/870 correctness at 5.3e-4, no new ruff findings.
Co-authored-by: Cursor <cursoragent@cursor.com>
The sweep crossed seven axes into 864 cases that compiled only 12 distinct kernels, re-running each about 72 times. The kernel specialises at compile time on (heads, head_dim, KVBlockSize, Preshuffle); batch, next_n, context length, SplitKV and the Q dtype are runtime values, and none of the axes interact -- nothing couples head_dim to SplitKV or the dtype to KVBlockSize -- so the product was testing interactions that cannot exist. q_dtype alone cost 432 cases while selecting no different code path on gfx950, where convert_q_fn is always False. Now: a full cross of the compile-time axes (16 cases, one per distinct kernel -- more preshuffle coverage than before, which only hand-picked 4), each runtime axis varied alone off a base, and four interaction cases where the tile-range arithmetic genuinely couples. --exhaustive still runs the original 870 for release validation. Validated by mutation testing rather than by the suite passing, since a suite cut too far also passes. Eleven deliberate kernel bugs -- swapped preshuffle and block-flat strides, causal off-by-one, dropped tail clamp, wrong lane half, wrong block-table index, dropped kv_scale, wrong head weight, floor-instead-of-ceil split (needs SplitKV>1), batch-blind output row (needs batch>1), next_n-blind causal bound (needs next_n>1). The full suite catches 11/11; the curated matrix also catches 11/11. 29 cases in 8.7s against 870 in 79s. Harness in att_profiles/mutation_check.py. Co-authored-by: Cursor <cursoragent@cursor.com>
Neither was ever exercised: every case ran the default variant at ChunkK=128, so paged_w2 and paged_w4 had no coverage at all, and the Tier-1 change of DEFAULT_VARIANT silently moved which one CI touched. That hole was real, not theoretical. A mutation of the wave-to-column-tile assignment (`wave * N_TILES_PER_WAVE` -> `wave * (N_TILES_PER_WAVE - 1)`) is correct by construction at one wave per CTA, where `wave` is always 0, and wrong for any multi-wave variant. It survived both the curated suite *and* the original 870-case exhaustive sweep. Adds variant and chunk_k parameters plumbed through to the kernel, reported in the results table, and six cases: paged_w2 and paged_w4 at the default tile width, ChunkK 64 and 256 against the variants that divide them, and a preshuffle x paged_w4 cross. 35 cases in 9.0s, and the mutation set now runs 12/12 where the exhaustive suite manages 11/12. Co-authored-by: Cursor <cursoragent@cursor.com>
Docstrings trimmed to what a reader cannot derive -- the oracle's formula and mask rule, the co-pack byte layout, what the preshuffled copy is for -- dropping the restatement and the cross-references to other files. _build_inputs returned an 8-tuple that was unpacked over ten lines at the call site; it now returns a NamedTuple, so the test body reads inp.q_fp8 rather than positional unpacking, and the contract is documented by the type. Its block table is also built vectorized instead of one element per GPU write, which is why the suite went from 9.0s to 7.5s. Same values: the per-case errors are unchanged to the digit. The pool there is still sized from max_model_len, so sequences share blocks once the batch outgrows it. Left alone deliberately, with a comment saying why: harmless for correctness and arguably extra coverage of block reuse, unlike in the benchmarks where it shrank the measured working set into cache. Still 35 cases, still 12/12 on the mutation set. Co-authored-by: Cursor <cursoragent@cursor.com>
Revert buffer load/store aux to the upstream IntegerAttr form expected by flydsl 0.3.1, and route scaled MFMA through the expr wrapper so cbsz/blgp attributes lower correctly during paged MQA kernel JIT. Co-authored-by: Cursor <cursoragent@cursor.com>
…MQA logits Prefetch kv_indices one tile ahead (skip reload when ind_off unchanged) and carry K-step-0 across the inner column loop to hide VMEM latency on the production kvb64/preshuffle decode path. Co-authored-by: Cursor <cursoragent@cursor.com>
Author
|
Update: ATT-guided prefetch on the production FlyDSL path (KVB=64 + preshuffle, B16/kv32k): ~27.6 → ~25.8 µs (−6.4%), ATT total stall −17%. From WaveScope dev work using |
Bind loop SSA values via default args (B023) and apply Black formatting. Co-authored-by: Cursor <cursoragent@cursor.com>
samremes
reviewed
Aug 25, 2026
| """Load 32 fp8 bytes as the vector<8xi32> A/B operand of the 32x32x64 MFMA. | ||
|
|
||
| Each lane holds 4 K-groups of 8 bytes at ``lane8 + kk*16``. Reads go out as | ||
| 2-dword loads because a ``v8i8`` buffer_load fails to lower on gfx942. |
Contributor
There was a problem hiding this comment.
This kernel is using 32x32x64 mfma but this comment talks about gfx942, is this relevant?
| try: | ||
| return torch.cuda.get_device_properties(device_index).multi_processor_count | ||
| except Exception: # noqa: BLE001 | ||
| return 304 |
Contributor
There was a problem hiding this comment.
should be 256 if this is a gfx950 targeting kernel?
| @@ -0,0 +1,335 @@ | |||
| # SPDX-License-Identifier: MIT | |||
Contributor
There was a problem hiding this comment.
I think benchmarks are usually part of the main tests, not separate. The test will serve both as validation and benchmark.
| @@ -0,0 +1,427 @@ | |||
| # SPDX-License-Identifier: MIT | |||
Contributor
There was a problem hiding this comment.
Separate tests/benchmarks for kernel comparisons probably shouldn't be created.
| ) | ||
|
|
||
|
|
||
| @benchmark() |
Contributor
There was a problem hiding this comment.
should this be decorating some other function, now this includes all input building too?
| a_pack, b_col, w_frag, kv_scale, m_tiles=M_TILES, k_steps=K_STEPS | ||
| ) | ||
| is_writer = _to_raw((lane_div_N == 0) & (col <= q_limit)) | ||
| with ir.InsertionPoint(scf.IfOp(is_writer).then_block): |
Contributor
There was a problem hiding this comment.
ir/scf should not be usually needed
Comment on lines
+20
to
+21
| from flydsl.expr import arith, range_constexpr, rocdl | ||
| from flydsl.expr.numeric import ArithValue |
Contributor
There was a problem hiding this comment.
Check if arith or ArithValue are really needed.
| ``fp8_mqa_logits`` carries its own copies and does not import this. | ||
| """ | ||
|
|
||
| # No `from __future__ import annotations`: FlyDSL arg typing needs real |
Comment on lines
+8
to
+10
| Split out of ``fp8_paged_mqa_logits`` to keep the gather layouts, the MFMA | ||
| head-reduce and the output view separable from the kernel body. The dense | ||
| ``fp8_mqa_logits`` carries its own copies and does not import this. |
Contributor
There was a problem hiding this comment.
Could you remove copies from elsewhere? And remove this comment.
| while len(ops) < 6: | ||
| ops.append(0) | ||
| ops.extend([neutral, 0, neutral]) | ||
| return rocdl.mfma_scale_f32_32x32x64_f8f6f4(result_type, ops) |
Contributor
There was a problem hiding this comment.
check if this is already supported more natively in flydsl mma
Author
There was a problem hiding this comment.
every higher level flydsl way explored seems to result in more sloc.
Apply Black formatting and resolve Ruff PLR0124/I001 so CI pre-checks pass. Co-authored-by: Cursor <cursoragent@cursor.com>
amd-sriram
added a commit
to amd-sriram/vllm
that referenced
this pull request
Aug 26, 2026
The decode half of the sparse-attention indexer runs on aiter's Triton/Gluon deepgemm_fp8_paged_mqa_logits. ROCm/aiter#4221 adds a gfx950 FlyDSL implementation of the same kernel, reported at 1.3-4.2x over Gluon on MI355X, and the decode logits are the indexer's dominant per-step cost at long context. Route the gfx950 branch of rocm_fp8_paged_mqa_logits at it when VLLM_ROCM_USE_AITER_FLYDSL_PAGED_MQA_LOGITS is set, so one build can run both arms and any aiter without the PR falls back to the existing path. Off by default while the PR is unmerged. This is a second call site rather than the module swap used for the prefill kernel in the preceding commit. The two launchers take the same tensors in the same order, but the FlyDSL one's remaining knobs are keyword-only with its own tuned defaults, so sharing a caller would mean forcing the Gluon ChunkK=256 on it. KVBlockSize is not optional: the kernel asserts it against the KV cache's own block dimension, so it is passed through along with Preshuffle, exactly as the Gluon call does. The FlyDSL path also skips the -inf prefill and the NaN scrub the Gluon path needs. Gluon writes the full row and leaves NaN in the padding; FlyDSL writes every column up to its causal bound `context_lens[b] - next_n + n` and nothing past it (`is_writer = ... & (col <= q_limit)`). That bound is exactly the decode top-k's read bound, so the untouched tail of the workspace is never read. Its docstring asks callers for an -inf prefill because a caller that reads the whole row would need one; this one does not. The FlyDSL kernel is written against gfx950, so _flydsl_paged_mqa_logits_kernel() returns None on every other architecture and the Triton/Gluon kernel is used regardless of the flag. Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com>
amd-sriram
added a commit
to amd-sriram/vllm
that referenced
this pull request
Aug 26, 2026
The decode half of the sparse-attention indexer runs on aiter's Triton/Gluon deepgemm_fp8_paged_mqa_logits. ROCm/aiter#4221 adds a gfx950 FlyDSL implementation of the same kernel, reported at 1.3-4.2x over Gluon on MI355X, and the decode logits are the indexer's dominant per-step cost at long context. Route the gfx950 branch of rocm_fp8_paged_mqa_logits at it when VLLM_ROCM_USE_AITER_FLYDSL_PAGED_MQA_LOGITS is set, so one build can run both arms and any aiter without the PR falls back to the existing path. Off by default while the PR is unmerged. This is a second call site rather than the module swap used for the prefill kernel in the preceding commit. The two launchers take the same tensors in the same order, but the FlyDSL one's remaining knobs are keyword-only with its own tuned defaults, so sharing a caller would mean forcing the Gluon ChunkK=256 on it. KVBlockSize is not optional: the kernel asserts it against the KV cache's own block dimension, so it is passed through along with Preshuffle, exactly as the Gluon call does. The FlyDSL path skips the -inf prefill the kernel's docstring asks for. That request is for a caller that reads the whole row; this one does not. FlyDSL writes every column up to its causal bound `context_lens[b] - next_n + n` and nothing past it (`is_writer = ... & (col <= q_limit)`), and that bound is the decode top-k's read bound, so the tail the fill would cover is never written and never read. The sanitize is kept, and is the same bounded sanitize_decode_logits() the Gluon path calls. Its window is that same causal bound, so it covers exactly the columns FlyDSL wrote: no more work than the Gluon path already does, and both paths hand the top-k a workspace in the same state. Keeping it means the flag toggles the kernel and nothing else, which is what makes the two arms comparable. The FlyDSL kernel is written against gfx950, so _flydsl_paged_mqa_logits_kernel() returns None on every other architecture and the Triton/Gluon kernel is used regardless of the flag. Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com>
1 task
Land the production indexer mapping (resident Q, depth-4 NT B-ring, SplitKV) without the v2/v3 experiments or local profiling artifacts. Co-authored-by: Cursor <cursoragent@cursor.com>
Compile a separate 16x16x128 kernel per head count so the indexer matches gluon's ChunkQ=heads set without changing the FlyDSL page pipeline. Co-authored-by: Cursor <cursoragent@cursor.com>
Drop the unused second query row on next_n=1 so MFMA/VALU and store waits match the actual decode shape. Co-authored-by: Cursor <cursoragent@cursor.com>
Keep next-page K in fixed register banks instead of scf.for yield copies, store one wave-wide page of logits, and prefetch the page table so gather latency stays overlapped. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
…path
Route H32/H64 D128 KVB64 preshuffle next_n in {1,2} through the 16x16x128 kernel, keep the 32x32x64 implementation as the generic fallback, and satisfy Black/Ruff.
Co-authored-by: Cursor <cursoragent@cursor.com>
Drop the generic 32x32x64 kernel and fold helpers into the gfx950 decoder so H32/H64 stays on the tuned 16x16x128 pipeline, with i64 output-row addressing and idiomatic FlyDSL MMA/scheduling that still matches the previous ISA. Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Cursor <cursoragent@cursor.com>
…QA logits Match the Gluon contract so next_n can be 1..8 without padding Q to 8 rows; optional next_n_lens still selects live rows per sequence. Co-authored-by: Cursor <cursoragent@cursor.com>
amd-sriram
added a commit
to amd-sriram/vllm
that referenced
this pull request
Sep 14, 2026
The decode half of the sparse-attention indexer runs on aiter's Triton/Gluon deepgemm_fp8_paged_mqa_logits. ROCm/aiter#4221 adds a gfx950 FlyDSL implementation of the same kernel, reported at 1.3-4.2x over Gluon on MI355X, and the decode logits are the indexer's dominant per-step cost at long context. Route the gfx950 branch of rocm_fp8_paged_mqa_logits at it when VLLM_ROCM_USE_AITER_FLYDSL_PAGED_MQA_LOGITS is set, so one build can run both arms and any aiter without the PR falls back to the existing path. Off by default while the PR is unmerged. This is a second call site rather than the module swap used for the prefill kernel in the preceding commit. The two launchers take the same tensors in the same order, but the FlyDSL one's remaining knobs are keyword-only with its own tuned defaults, so sharing a caller would mean forcing the Gluon ChunkK=256 on it. KVBlockSize is not optional: the kernel asserts it against the KV cache's own block dimension, so it is passed through along with Preshuffle, exactly as the Gluon call does. The FlyDSL path skips the -inf prefill the kernel's docstring asks for. That request is for a caller that reads the whole row; this one does not. FlyDSL writes every column up to its causal bound `context_lens[b] - next_n + n` and nothing past it (`is_writer = ... & (col <= q_limit)`), and that bound is the decode top-k's read bound, so the tail the fill would cover is never written and never read. The sanitize is kept, and is the same bounded sanitize_decode_logits() the Gluon path calls. Its window is that same causal bound, so it covers exactly the columns FlyDSL wrote: no more work than the Gluon path already does, and both paths hand the top-k a workspace in the same state. Keeping it means the flag toggles the kernel and nothing else, which is what makes the two arms comparable. The FlyDSL kernel is written against gfx950, so _flydsl_paged_mqa_logits_kernel() returns None on every other architecture and the Triton/Gluon kernel is used regardless of the flag. Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Motivation
Follow-up to PR #3913 (dense FP8 MQA-logits FlyDSL). This PR adds the paged / decode lightning-indexer path for gfx950 (CDNA4).
The DeepSeek indexer produces sparse-attention selection logits during decode. For each query row
mand KV positionn:QandKare fp8 e4m3;weights,kv_scale, and the output are f32. K and its per-token scale are co-packed in a preshuffled paged cache (KVBlockSize=64) and gathered through a block table.Ragged
next_nfor confidence-scheduled verificationThe motivation for per-sequence
next_nis the same problem vLLM describes for DSpark adaptive verification [1]: a staticnum_speculative_tokensis not optimal across concurrencies, because “per-position acceptance decays fast” and a low-probability tail token “costs a slot in every verification batch.”From that post: DSpark’s confidence head “scores each drafted token’s chance of surviving verification,” and the scheduler keeps a contiguous prefix per request. “Slots compete across requests: position 5 of a confident request can outrank position 1 of a low-confidence one.” The decode step is therefore varlen. vLLM captures graphs with a promised
max_query_len = num_speculative_tokens + 1so “one graph serves any mix of 1 tonum_speculative_tokens + 1tokens per request.” On SM100 they note that this needs a varlen indexer (DeepGEMM, landed with vLLM PR #47808).The indexer on gfx950 must follow that mix. Production Gluon
deepgemm_fp8_paged_mqa_logitstakes one launch-widenext_n = Q.shape[1], so a trimmed batch is either padded to the longest prefix (wasted Q rows and stores) or split into homogeneous launches. This kernel accepts:Q[B, next_n, 32, 128]withnext_n ∈ {1..8}(same packed layout as Gluon; no forced pad to 8)next_n_lens[B], each in{1..next_n}: live rows for sequencebare the leadingmin(next_n_lens[b], next_n)rows-infThe compiled kernel still reserves LDS for eight rows (one binary, graph-capture friendly). The compute loop early-exits when
row >= nn, so a confidence-trimmed sequence does not replay MFMA / stores for discarded draft slots. That is the gfx950 counterpart of the varlen indexer vLLM cites for SM100 [1].Scope: gfx942 / KVBlockSize=1 / H=64 were explored earlier in the branch and dropped. The landed kernel is gfx950 H=32 D=128 KVB=64 only.
Technical details
Kernel:
aiter/ops/flydsl/kernels/mqa_logits/fp8_paged_mqa_logits_gfx950.pyPublic entry:
flydsl_fp8_paged_mqa_logits(re-exported frommqa_logits/fp8_paged_mqa_logits.py)next_n:rows_per_batchfromQ.shape[1];nn = rows_per_batchifnext_n_lensis omitted, elsemin(next_n_lens[b], rows_per_batch). Q/W LDS fill and logit stores are indexed withpid_batch * rows_per_batch + row.aux=2).v_mfma_scale_f32_16x16x128_f8f6f4viafx.make_mma_atom+fx.gemm. H=32 is two M-atoms. ReLU·w,shuffle_xor(16/32), lanes 0–15 NT-store.min(real_pages, ceil(total_cu * 3 / B))targeting that occupancy. ExplicitSplitKVremains graph-capture safe.Test plan
op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits.py— gfx950 correctness vs torch port of vLLMfp8_paged_mqa_logits_torch. Gate: exact-infmask +calc_diff < 1e-3.op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_gfx950.py— long-context / SplitKV / wide-output cases, including raggednext_n_lens.op_tests/flydsl_tests/test_flydsl_fp8_paged_mqa_logits_ragged_nn.py— compact uniformnext_n ∈ {1,2,4,8}and mixednext_n_lensin one batch (unused rows stay-inf; mutating padded Q rows must not change live logits).Performance (MI355X / gfx950)
Workload is the Probe_1079 decode histogram: 212 calls, 137 unique Nq (1..256, mean 117.1), scattered
randpermpages, KV pool 44362 blocks,max_model_len=258048. 10 iters / 3 warmup. Times in µs.Ragged mix (FlyDSL-only). Per-sequence
next_n_lenssampled in 1..8 (mean 4.56, mean live Q rows 528). This is the confidence-trim shape: one launch, heterogeneous prefixes. Gluon has nonext_n_lensand would have to pad to 8 or split launches.Uniform compact vs Gluon (
deepgemm_fp8_paged_mqa_logits, Triton 3.7.0,ChunkK=256,WavePerEU=2). Same Nq mix; both backends getQ[:, next_n]. Ratio is Gluon / FlyDSL (>1 means FlyDSL is faster).Crossover is nn≈2 at 8k and between nn=2 and nn=3 at 128k. Gluon wins at a uniform
next_n=1; FlyDSL pulls ahead as more Q rows reuse each fetched page. Under the adaptive-verification mix described in [1], meannext_nis not the pad width: our ragged numbers (mean nn=4.56) sit between the uniform nn=4 and nn=5 columns rather than at padded nn=8.References