perf(pa tile): readable tile-programming PA decode kernel, matches/beats production#825
Open
fsx950223 wants to merge 59 commits into
Open
perf(pa tile): readable tile-programming PA decode kernel, matches/beats production#825fsx950223 wants to merge 59 commits into
fsx950223 wants to merge 59 commits into
Conversation
Readable tile-programming reimplementation of pa_decode_ps_kernel's fp8 paged-attention decode math, using the high-level layout/tile API (make_buffer_tensor + zipped_divide + make_tiled_mma + fx.gemm + tiled copies) instead of raw buffer intrinsics and hand-scheduled MFMA. Scope: per-tensor fp8 K/V, query_length=1, no kv-varlen, no sliding window. Validated against reference_masked_attention via test_pa_decode_tile_reference (9 cases, fp8 tolerance 1e-2). Pipeline mirrors the production kernel where the tile API allows: 4-warp (1,4,1) tiled MMA splitting tokens for QK / head-dim for PV, fp8 16x16x32 MMAs, cross-CTA partition split + parallel flash-combine reduce kernel, V-load software pipelining, and a register-resident DPP (shuffle_xor) max reduction over the QK fragment. The exp/fp8-pack/sum stay LDS-staged: a fully register-resident softmax would need a bespoke non-row-major fp8 logits layout + raw PV read (fp8 fragment->LDS does not legalize in the tile API), which would abandon the readable abstraction. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Carry the P·V output accumulator in registers across the flash loop (a loop-carried PV C-fragment rescaled by the softmax correction each tile) instead of staging per-tile partials through LDS. This removes the sOp LDS buffer, the per-tile accumulate barrier, and its LDS round-trip, mirroring how pa_decode_ps_kernel keeps O in registers. Combined with the earlier V-load software pipelining and the register- resident DPP max, this takes the kernel from occupancy 1 to 2 CTA/CU on gfx942: LDS 39936->31744 B (under the 32KB/2-CTA threshold) and combined VGPR+AGPR 344->256 (AGPR 216->132, no longer spilling the per-tile output partials). High-batch latency drops ~2x vs the start of this work (batch=128 ctx=16384: gap to production 10.6x -> 6.0x); low batch unchanged. All 9 test_pa_decode_tile_reference cases pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The exp + fp8-pack + sum (phase 2) previously ran on only the 16 row-owner lanes per warp (each handling a full 64-token row), leaving 48 lanes idle. Distribute it 4 lanes/row × 16 tokens each: every lane exps and fp8-packs its 16-token slice, and the four 16-token partial sums are combined across the sub-lanes with shuffle_xor(16,32) so the row denominator matches the row-owner result. Occupancy/registers unchanged (still 2 CTA/CU). batch=16 ctx=16384 237->155us (-29%), batch=128 ctx=16384 1127->979us. Session cumulative gap to production: batch=128 ctx=16384 10.6x -> 5.2x, batch=16 ctx=16384 7.0x -> 3.8x, batch=1 1.5x -> 1.3x. All 9 test_pa_decode_tile_reference cases pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Switch QK to D[token,qhead] = K(A)·Q(B)ᵀ with tiled_mma (NWARP,1,1) splitting tokens (M) across warps, so each lane owns one qhead and the softmax reduces over M=token with a register reduce + shuffle_xor(16,32) — 2 cross-lane offsets instead of the old qhead=M orientation's 16 (4 rows × 4 offsets). This makes a fully register-resident softmax actually pay off: scores stay in the QK C-fragment (sS eliminated), P is exp'd in registers and packed fp8 straight to sP via the transposed C-store view (make_fragment_like + cvt-pack + Vector.bitcast — no arith.truncf), and the row sums reduce with the same cheap shuffles. PV/register-O/epilogue unchanged (PV still reads sP[qhead,token]). Resources: VGPR 124→24, AGPR 132→128, combined 256→152, LDS 31744→15360 → occupancy 2→3 CTA/CU. Perf improves vs the previous best (b128 ctx16384 979→909, b128 ctx4096 280→237, b1 1.3×→1.2×); gap to production now ~4.9× high batch / 1.2× low batch (was 10.6× / 1.5× at the start of this work). All 9 test_pa_decode_tile_reference cases pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Restructure the QK+softmax to pa_decode_ps_kernel's register discipline so the cross-iteration K prefetch fits without losing occupancy 3: - TLOOP QK: NCHUNK=4 small fx.gemm over [TOK_CHUNK=64, qhead] → f32x4/lane, with the softmax processed 4 scores at a time (scores stay in AGPR instead of a 16-wide VGPR vector from frag_S.load()). - scalar-threshold token mask (token < n_valid → (a*64+r) < n_valid-base), replacing the three 16-wide mask vectors. - V loaded in the PV loop (no hoist) — that VGPR is spent on the K prefetch, matching production (prefetch K, load V in place). - K prefetch: NCHUNK persistent A sub-fragments, prologue + prefetch-after-QK. Effect vs the previous occ-3 single-gemm version: VGPR 24→12 (with K prefetch now resident; standalone K prefetch had been VGPR 56 / combined 184 / occ 2), combined 152→144, **occupancy stays 3 with K prefetch**. Perf is mixed — mid batch improves (b16 ctx16384 153→147) while high batch is ~flat-to-slightly worse (b128 ctx16384 909→926): the fx per-gemm TLOOP overhead offsets the prefetch gain, since production's TLOOP is free (native raw mfma). All 9 test_pa_decode_tile_reference cases pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The ATT trace of the prior commit showed 76% of stalls were V-load latency: V was loaded inside the PV loop (issued then immediately consumed by the PV gemm) after the K-prefetch change un-hoisted it. pa_decode_ps_kernel hides BOTH loads — K via cross-tile prefetch, V via an early load before QK that overlaps QK + softmax. Restore the V early-load (hoist into frag_Vs before the QK gemms) and keep the cross-tile K prefetch. The TLOOP softmax's low VGPR (12) leaves AGPR headroom, so the hoisted V fragments fit without losing occupancy: VGPR 12, AGPR 132, combined 144, occupancy 3 retained — both loads now hidden. batch=128 ctx=16384 926->851us (-8%, best yet vs the 909 of the V-hoist-only 8ada08f and 926 of the K-prefetch-only 28a3bb1); batch=16 ctx16384 153->147; low batch unchanged. All 9 test_pa_decode_tile_reference cases pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Replace the make_tiled_copy_A + fx.gemm QK path (capped at 64-bit buffer_load_dwordx2 by the MFMA atom layout) with raw 128-bit global_load (i64x2) straight to registers feeding raw mfma_f32_16x16x32_fp8_fp8. QK sums over the full head_dim, so the head->k_step permutation is free: pick the mapping where each lane's slice head[rgroup*32:+32] is contiguous (one dwordx4 pair). K and Q use the same mapping, so the MMA output C-fragment is identical to fx.gemm's and softmax/P-pack/PV are unchanged. Cross-tile K prefetch kept via scf.for iter_args. K loads 32x buffer_load_dwordx2 -> 16x global_load_dwordx4; VGPR/LDS/AGPR unchanged (no occupancy regression), no spills. Correctness 9/9. Best-of-5 on a quiet gfx942: b128 ctx16384 844->724us (-14.2%), b64 430->368 (-14.5%), low batch flat. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Replace the file-local _global_ptr / _load_i32 / _gload_i64x2 with the shared extract_global_ptr / global_load_i32 / global_load_i64x2 from kernels/utils.py (merged from main). No behavior change; drops the now unused ir / llvm / buffer_ops imports. Correctness 9/9. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Rawify the PV path like QK: PV contracts over token, so the token->k_step permutation is free as long as V (B) and P (A) agree. Each lane takes the contiguous token slice [rgroup*64:+64] for its head, loaded as 4x i64x2 (dwordx4) straight to registers; P is read raw from sP with the same permuted token slice (8 i64). Raw mfma_f32_16x16x32_fp8_fp8 accumulates the 8 k_steps into the same f32x4 C-fragment the old fx.gemm produced, so the O-accumulate + epilogue are unchanged. Drops make_tiled_copy_B (V), make_tiled_copy_A (P read) and their fragments. All K/V loads now global_load_dwordx4 (buffer_load_dwordx2 gone); VGPR 140 / LDS 15104 / AGPR 0 unchanged, no spills. Correctness 9/9. Best-of-5 on a quiet gfx942 vs the K-raw commit: b128 ctx16384 725->470us (-35%), b64 368->238 (-35%), low batch flat. Cumulative vs pre-raw original: b128 844->470us (-44%). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
… path) ATT trace of the K+V-raw kernel showed the per-tile block-table lookup (_tile_coords' global_load_i32 for `phys`) on the critical path: it stalls ~7% directly and serialises the V/K address computation behind a ~600cyc block-table load (the V-address line waited on it). `phys` was already being loaded one tile ahead in the K-prefetch path, so carry it through the scf.for iter_args like K: the loop body reads the prefetched `phys` and only recomputes the cheap within_tile/tok0/page arithmetic (_coords_no_phys). Correctness 9/9. CUDA-graph best-of-5 on a quiet gfx942 vs the V-raw commit: b128 ctx16384 474->445us (-6%), b64 242->233 (-3.8%), low batch flat. No register/occupancy change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ATT wave-slot analysis showed the tile runs ~1.6 waves/SIMD at high batch
while production runs 3 (the hardware max here) — the tile was
grid-under-dispatched, not register-limited. Its NP heuristic targeted
just 1 CTA/CU (ceil(CU/base_ctas)=1 once batch*kv >= CUs), so one coarse
CTA per sequence walked all 64 tiles serially and left 2/3 of the
occupancy empty.
Add a granularity floor npart_by_gran (~24 tiles/partition) so even at
high batch the grid oversubscribes the CUs like production's persistent
scheduler, and allow NP in {3,6} (NP=2 is a measured tail-effect trap at
b128). Best-of-5 CUDA-graph on a quiet gfx942 ctx16384: b128 442->392us
(-11%), b96 445->329 (-26%), b256 905->757 (-16%), b32 -10%; b64 tie,
b1-16 unchanged (still occupancy-driven). 9/9 correct.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…count Replace the bespoke npart_by_occ/npart_by_gran heuristic with the same recommender pa_decode_ps uses (kernels.pa_decode_fp8.get_recommended_splits, mirroring aiter's Gluon get_recommended_splits) for consistency, instead of maintaining a second bespoke formula. Measured tradeoff: a direct NP sweep on this kernel showed the recommender is 4-11% slower than the bespoke heuristic it replaces (b128 419 vs 377us, b64 243 vs 228, b16/b8 ~230 vs ~222) because it's tuned for production's persistent worklist scheduler (dynamic cross-CTA load balancing) rather than this kernel's static per-partition split -- its [4,8] floor overshoots the measured sweet spot (NP=3 at high batch). Still nets -5 to -12% vs the pre-partition-tuning baseline (877c2e69's parent) at high batch. 9/9 correct. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…es clamp Match production's pa_decode_ps_launch call site exactly (it passes get_recommended_splits' result straight to grid.z with no extra clamp) instead of a bespoke min(..., num_tiles) safety clamp. Verified this is not a correctness risk: when a partition's tile range starts past num_tiles, scf.for's trip count is max(0, ub-lb), so it resolves to zero iterations (a harmless empty (m=-inf, l=0) contribution to the reduce), not undefined behavior -- confirmed 9/9 pass either way. It does cost measurable overhead on genuinely tiny contexts from the wasted CTA launches + reduce-kernel work merging empty partials (b=1 ctx=256, 1 real tile: +15%; b=3 ctx=1024, 4 tiles: +1.3%), which the prior clamp avoided. Kept unclamped per direct instruction to match production's call site. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The once-per-CTA Q row load ("global_load_ushort") was a single monolithic
HEAD-wide load().to(f32) followed by an absmax reduce and rescale. The
LLVM IR emitted a clean <128 x half> load, but the AMDGPU backend
scalarized it into 128x global_load_ushort: keeping the full 128-wide f32
vector live across the reduce-then-scale sequence needs ~128 contiguous
VGPRs, which under this kernel's AGPR pressure the allocator finds more
expensive than issuing 128 independent narrow loads consumed one at a
time by a pairwise reduction.
Chunk the load into QCHUNK=8-element (128-bit) pieces loaded raw via
global_load_i64x2 (matching the K/V/P raw-load pattern), bounding the
live f32 range to 8 elements per iteration so each chunk lowers to a
real global_load_dwordx4. Two passes over the same 16 chunks (max-reduce,
then rescale+pack) -- the compiler CSEs the redundant reload since
nothing invalidates it between passes, so the net cost is still one load
per chunk (16 total), not 32.
Result: all 128 global_load_ushort eliminated -> 40 global_load_dwordx4
total (K+V+Q), ISA instr 1893->1780. VGPR/AGPR shift slightly (12/132 ->
24/128, combined unchanged), occupancy unaffected (grid-driven at this
config, not register-driven). Correctness 9/9. CUDA-graph best-of-5 on a
quiet gfx942, 2 stable reps: b128 ctx16384 396->388us (-2.1%), b64/b16/b8
also improve slightly (-0.3 to -0.8%), no regressions.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
ATT trace comparison against pa_decode_ps_kernel showed a 10% stall unique to the tile on the block-table phys-page load, despite it already being prefetched one tile ahead (2fc39dc): 4 consecutive tiles share the same physical page for the standard block_size=1024/TILE_TOK=256 config (page = tok0 // block_size only advances every block_size/TILE_TOK tiles), so ~75% of the per-tile reloads were fetching an unchanged value. Production avoids this by batch-loading all pages for its (bounded) partition once upfront; a literal port isn't safe here since this kernel's partition depth is unbounded (NP is now purely batch-driven, not context-length-driven -- see e519999/48b7b9c7), so instead cache the current page's phys value and only reload via a real conditional (a Python if/else that lowers to a genuine `scf.if -> (i32)`, not arith.select, so the global_load_i32 truly doesn't execute when skipped) when the next tile's page actually differs. Verified via fresh ATT trace: target stall (block-table load) dropped 10.0%->3.1% (-69%, matching the ~75% redundant-load elimination); total stall cycles -10.3%. Correctness 9/9. Event-timed benchmarking initially showed this as perf-neutral (launch-overhead noise masked it); CUDA-graph benchmarking (no launch overhead) shows a real, stable win across 2 reps: b128 ctx16384 445->421us (-5.4%), b64 231->227 (-1.7%), b16 flat, b1 flat (noise, only 8 real tiles at this shape). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…opy_C round trip Each lane already knows the exact (qhead, token) address for its packed fp8 word (same addressing already used for masking); storing it directly via _view/store removes the per-tile fragment-copy indirection. b=128 ctx=16384: ~391->387us (CUDA graph, clean A/B), 9/9 correctness.
maximumf(x, -x) lowers to a per-element v_cmp_o + v_cndmask pair; fmath.absf is a single v_and_b32 (clear sign bit). ISA instr count 2058->1734 (-16%), VGPR 152->136; v_cndmask 208->80, v_cmp 195->67. Once-per-CTA so wall-clock at b=128 ctx=16384 is flat (~387us, unchanged) but the ISA is meaningfully leaner. 9/9 correctness.
ATT trace showed 58% of all stall cycles on the raw K/V dwordx4 load instructions. Address-pattern analysis explained why: the plain PA K layout ([blocks,kv,block_size,head], head innermost) means adjacent lanes -- which the MFMA A-operand's fixed lane roles assign to adjacent TOKENS, not adjacent head-dim elements -- land head_dim=128 bytes apart per load, i.e. essentially uncoalesced. Production's real pa_decode_ps_kernel avoids this by relaying out K at quantization time (key_cache.permute in test_pa.py's small-block harness) so adjacent lanes are 16B apart. Adopt the same idea: K is now expected pre-relaid-out as [blocks, kv_heads, head_dim//32, block_size, 32] (head-quarter outer, token next-innermost), so adjacent lanes (adjacent tokens, fixed head-quarter) are a contiguous 32B apart instead of 128B, while each lane's own 32B slice stays contiguous for its two i64x2 reads. This is a one-time relayout at cache-population time (same cost model as V's existing trans_v transpose), not a per-call cost -- doing it inside pa_decode_tile() per call would cost more than it saves (verified: the full K-cache read+write is comparable to the kernel's own runtime at this benchmark's scale). CUDA-graph clean A/B (GPU1, quiet, ctx=16384): b=128 388->310us (-20%), b=64 229->190us (-17%), b=16 74->64us (-14%), b=1 flat (compute-bound, not memory-bound at that scale). 9/9 correctness (under GPU contention, still correct).
Same coalescing fix as the preceding K commit, applied to V: the plain transposed [blocks,kv,HEAD,block_size] layout put adjacent lanes (which own adjacent HEAD values for PV's B operand) a full block_size bytes apart per load -- worse than K's issue since block_size (e.g. 1024) >> head_dim (128). V is now expected pre-blocked as [blocks, kv_heads, HEAD//16, block_size//64, 16, 64]: grouping by (head//16, token//64) with head-within-group major and the 64-token run minor makes adjacent lanes exactly 64B apart while keeping each lane's own 64-token run contiguous. head_group/token_group need no runtime div/mod (vh*VHE_SIZE and warp*16 are both multiples of 16; within_tile/rgroup combine directly into the token group index). CUDA-graph clean A/B (GPU1, quiet, ctx=16384), combined with the prior K-blocking commit vs the pre-blocking baseline: b=128 391->248us (-37%), b=64 229->163us (-29%), b=16 74->58us (-22%), b=1 49->40us (-18%). Gap to the real pa_decode_ps_kernel (~190-214us at this config) narrows from ~2.0x to ~1.2-1.3x. Full 9/9 correctness sweep (context lengths 17/256/1027, boundary/masking cases included).
…, LDS bank-conflict fix - Vectorize the epilogue's output write across all 256 threads instead of only GS row-owner threads doing an unrolled per-element scalar loop; cuts the epilogue's static ds_read/global-store instruction count and fully utilizes the wave. - Switch softmax exp2 calls from MLIR's generic math.exp2 (a polynomial approximation with extra edge-case cndmask overhead) to the HW exp2 intrinsic (exp2_f32_fast/exp2_amdgcn_scalar), matching pa_decode_ps_kernel's own usage. - Prefetch V's page-table index one tile ahead via a scalar, warp-partitioned wide load broadcast through LDS, reusing an existing softmax barrier instead of adding a new one -- cuts redundant per-lane VMEM page lookups that were previously re-derived every tile by all 4 warps. - Pad the sLmax/sLsum cross-warp scratch row stride to avoid a 2-way LDS bank conflict (mirrors pa_decode_ps_kernel's own PROB_ROW_STRIDE_BYTES padding technique). - Parametrize the block_size=16/64 test coverage and fix V's test-harness relayout to match the simplified blocked V-cache layout. Across the 36-shape tile-vs-ps benchmark sweep, pa_decode_tile_kernel now matches or beats pa_decode_ps_kernel on 35/36 shapes; the remaining one is at noise-level parity (confirmed via repeated high-confidence trials). Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…decode_ps_kernel, register-carry m/l
Generalizes pa_decode_tile.py's addressing to support head_dim in {64, 128}
(previously hardcoded for 128), and re-derives the K/V paged-cache addressing
to consume the exact same physical layout pa_decode_ps_kernel uses
(K=[nb,h,hd/16,bs,16], V(trans_v)=[nb,h,bs/16,hd,16]) so both kernels can
share cache-prep code -- test_pa.py's tile harness now reuses
shuffle_value_cache_layout() directly for V instead of a bespoke relayout.
Also carries the running softmax max (m) and denominator (l) as loop-carried
registers instead of per-tile LDS store/load: every thread already holds its
own cross-warp-combined value each tile, so only a single post-loop bridge
write into sM_off/sL_off is needed for the epilogue's differently-indexed
readers. Cuts LDS instruction count ~9% with no VGPR regression.
Adds a head_dim=64 test matrix (18 cases) alongside the existing head_dim=128
one; full 36-case suite passes.
…tch to pa_decode_ps_kernel's ISA Reads key_scale/value_scale as 1-element device tensors via a scalar buffer_load (llvm.amdgcn.s.buffer.load), matching pa_decode_ps_kernel, instead of baking them in as host-float kernel arguments; pa_decode_tile() still accepts a plain float too (auto-wrapped into a device tensor). value_scale is now folded in exactly once, in the main kernel's epilogue before the NP==1/NP>1 branch, instead of being duplicated in the NP==1 path and the reduce kernel separately -- the reduce kernel no longer needs value_scale at all, matching pa_decode_ps_kernel's reduce step. Adds bf16 query support: `query_dtype` compile-time flag (auto-detected from the query tensor's torch dtype), threaded into the Q load's element type, mirroring pa_decode_ps_kernel's own `query_input_dtype` flag. Closes several concrete gaps found by diffing LLVM IR against pa_decode_ps_kernel: - Q's per-row absmax butterfly reduction now uses `dpp_utils.dpp_xor_f32` (raw llvm.amdgcn.update.dpp.i32), the exact helper pa_decode_ps_kernel uses for the same reduction, instead of the DSL's generic `shuffle_xor` -- eliminates 4x mbcnt+ds_bpermute in favor of direct VALU-crossbar DPP. - K's per-tile block-table lookup batches into one wide scalar load (vec_width=NCHUNK) at block_size=64, where the NCHUNK page-table entries are contiguous, instead of NCHUNK separate scalar loads -- matches pa_decode_ps_kernel's `_pa_small_block_stage_phys_blocks` batching. - Scalar reciprocals (Q-absmax, softmax denominator, reduce-kernel output) use `rcp_f32` (HW reciprocal approximation), matching pa_decode_ps_kernel's own use of the same helper, instead of a full IEEE division; negligible accuracy cost against existing fp8 noise. - The softmax max store (`_st_lw`) is now unconditional across all 4 lanes sharing a qhead, matching pa_decode_ps_kernel's own (harmless, same-value) redundant store instead of gating it to one lane. Minor cleanups along the way: NQCHUNK is now the primary constant (QCHUNK derived from it) instead of a separately re-derived, always-true assert; `c_nhgroup` reuses the already-computed QCHUNK instead of recomputing HEAD//16; all ceil-divs use the existing `cdiv` helper from kernels/utils.py; dropped a redundant explicit `arith.index_cast` (fx.Int32 already handles index-typed values); reused the module-level WAVE constant instead of a local c_w alias. Verified throughout: 40/40 tests pass (added a dedicated bf16-query test matrix), including a tightened-1e-3-tolerance sanity check confirming no new error signature beyond the existing fp8-quantization noise floor. VGPR unchanged (112/0, 4 waves/SIMD) across all changes. Several other experiments (moving K's prefetch earlier, deferring the PV rescale, reordering K/V prefetch calls) were tried, measured via clean A/B, found to be regressions or no-ops, and reverted -- see inline comments recording those negative results. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
… hints + early V-load At head_size=64/bs=128/ctx=16384, pa_decode_tile_kernel was ~1.07x slower than production's pa_decode_ps_kernel despite equal VGPR/LDS budgets and near-identical instruction counts -- the residual gap was scheduling quality, not register pressure or a missing algorithmic mechanism. - Contiguous-per-warp token reassignment (K/mask/P-pack addressing) so a warp's own K page fetch feeds QK directly, no cross-warp LDS broadcast. - fastmath=contract on cross-warp sum combine and PV correction scale, matching pa_decode_ps_kernel's own usage; drops VGPR 92->84 at head_dim=64. - Fix a 16-way (write) / 8-way (read) LDS bank conflict on the P-pack buffer: its row stride (TILE_TOK=256B) was a multiple of the 32-bank*4B wrap. Pad to 272B (SP_ROW_BYTES), the smallest padding that both breaks the conflict and keeps the PV read's ds_read_b128 alignment. - rocdl.sched_group_barrier hints (sched_vmem on K/V's raw loads, sched_dswr on the P-pack store) to guide post-RA instruction scheduling around V's exposed HBM latency, at zero VGPR cost. - Re-test (now viable under the lower VGPR baseline above) issuing V's raw load right after the pass-1 barrier instead of right before PV use, giving it pass-2's exp/pack/store work to hide behind; costs only +8 VGPR, not enough to cross the 5->4 waves/SIMD occupancy boundary. - Reuse pa_decode_ps_kernel's own reduce kernel (compile_pa_decode_sw_reduce in pa_decode_swa.py, extended with a logits_dtype_str param) instead of a separate tile-specific reduce path. All head_dim=64-specific changes are gated by HEAD==64 (verified no regression at head_dim=128). Result: tile is now ~1.2% faster than production at head_dim=64/ctx=16384 and ~8% faster at head_dim=128/ctx=16384. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
… output dtype - Reduce pa_decode_tile.py's dense investigation-history comments to the essential correctness/design rationale; inline the c_nhgroup alias into QCHUNK. - Switch test_pa.py's precision checks from aiter's checkAllclose to torch.testing.assert_close throughout, including the PS-vs-Gluon sweep (which now stops at the first out-of-tolerance case instead of aggregating failures across the whole parametrized matrix). - Fix pa_decode_tile's output dtype: it was hardcoded to f16 regardless of query dtype (both the NP==1 epilogue write and the reduce kernel's output_dtype_str); now matches query.dtype (f16 or bf16), enforced by a new assertion in the host function. Updated test_pa.py's bf16-query test to allocate a matching-dtype output tensor. - Replace direct calls to the compiled launch closures with kernels.tensor_shim._run_compiled, matching the caching pattern already used in hgemm_splitk.py. - Validate block_tables/context_lengths/key_scale/value_scale dtypes via assertions instead of silently casting them at the launch call site. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
…stic - Fix a latent NaN bug in the NP>1 (and NP==1) epilogue: a partition assigned zero tiles has l==0, so an unclamped rcp_f32(0) returns +inf, and o_v(==0) * inf == NaN. That NaN survives the reduce kernel's exp2-based zero-weighting (0 * NaN == NaN, not 0), corrupting the whole output. Clamp the denominator to 1.0 when it's 0, matching pa_decode_ps_kernel's own safe_sum/inv_sum guard in _normalize_pa_output. This was previously unreachable because the old num_partitions heuristic always capped NP by the actual tile count, so no partition was ever fully empty. - Extract the num_partitions heuristic out of pa_decode_tile() into a standalone _choose_num_partitions() function, with target_ctas_per_cu/ min_tiles_per_partition/tile_tok as keyword overrides (defaulting to the values this heuristic was tuned with) instead of hardcoded locals. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Adds a new readable, tile-programming implementation of paged-attention FP8 decode (pa_decode_tile) using FlyDSL’s tile/layout APIs, with accompanying correctness tests and a small extension to the existing reduce kernel to support non-bf16 partial (“logits”) dtype.
Changes:
- Introduce
kernels/pa_decode_tile.py: tile-programming PA decode kernel + host wrapper, including a partitioning (grid.z) heuristic and reuse of the existing SW reduce kernel for NP>1. - Extend
compile_pa_decode_sw_reduceto allow choosing the dtype used when reading partition partials (logits), enabling f16/bf16/f32 partials. - Update
tests/kernels/test_pa.pyto usetorch.testing.assert_closefor comparisons and add reference-based tests for the new tile kernel.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 4 comments.
| File | Description |
|---|---|
| tests/kernels/test_pa.py | Switch comparison to torch.testing.assert_close and add reference tests for pa_decode_tile. |
| kernels/utils.py | Add a global_store helper alongside existing global load helpers. |
| kernels/pa_decode_tile.py | New tile-programming PA decode kernel + host entry point and partitioning heuristic. |
| kernels/pa_decode_swa.py | Add logits_dtype_str option and use it when loading partition partial logits. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+1071
to
+1074
| key_scale_t = key_scale if isinstance(key_scale, torch.Tensor) else torch.tensor([float(key_scale)], device=dev) | ||
| value_scale_t = ( | ||
| value_scale if isinstance(value_scale, torch.Tensor) else torch.tensor([float(value_scale)], device=dev) | ||
| ) |
Comment on lines
+734
to
+736
| # WITHIN-warp offset instead of the across-tile one. Row stride | ||
| # is SP_ROW_BYTES (TILE_TOK+8 padding), not TILE_TOK -- see | ||
| # SP_ROW_BYTES's comment above (bank-conflict avoidance). |
Comment on lines
+807
to
+812
| if logits_dtype_str == "f32": | ||
| LOGITS_DTYPE = fx.Float32 | ||
| elif logits_dtype_str == "f16": | ||
| LOGITS_DTYPE = fx.Float16 | ||
| else: | ||
| LOGITS_DTYPE = fx.BFloat16 |
Comment on lines
+1242
to
+1250
| # --------------------------------------------------------------------------- | ||
| # Tile-programming reference (kernels/pa_decode_tile.py) | ||
| # | ||
| # Validates the readable, correctness-first tile-programming reimplementation of | ||
| # the PA decode math against `reference_masked_attention`, for the minimal slice | ||
| # it supports: per-tensor fp8 K/V, query_length=1, no kv-varlen, no sliding | ||
| # window. The reference kernel consumes f16 K/V holding the fp8 codes (see its | ||
| # module docstring), so we fp8-quantize then cast the codes to f16 here. | ||
| # --------------------------------------------------------------------------- |
Signed-off-by: fsx950223 <fsx950223@outlook.com>
- Reuse softmax_scale in place instead of shadowing it as _softmax_scale - Drop the M alias for MFMA_MNK and use MFMA_MNK directly - Inline sQ_off (always 0, the first LDS region) at its call sites - Use fx.Int32 consistently instead of the bare Int32 import - Replace fx.Index/arith.index with fx.Int64/plain int for the tile loop Signed-off-by: fsx950223 <fsx950223@outlook.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com>
kernels/attention/pa_decode_tile.py hardcoded e4m3 FNUZ for the fp8 K/V/P MFMA operands, which is only the correct bit format on gfx942 (CDNA3). gfx950 (CDNA4) decodes the same mfma_f32_16x16x32_fp8_fp8 intrinsic as OCP e4m3fn, so FNUZ-encoded bits are silently misinterpreted -- reproduced the CI failure (40x "tile PA decode mismatch" on linux-flydsl-mi355-1, gfx942 unaffected). Pick FP8/FP8_MAX by arch (matching preshuffle_gemm.py / qk_norm_rope_quant.py's existing convention) and quantize the test's K/V with aiter's arch-aware dtypes.fp8 instead of a hardcoded torch dtype. Verified: all 40 previously-failing tests pass on gfx950 (MI355) after rebuilding FlyDSL from source there.
Signed-off-by: fsx950223 <fsx950223@outlook.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com>
…e path Signed-off-by: fsx950223 <fsx950223@outlook.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com>
Signed-off-by: fsx950223 <fsx950223@outlook.com>
Wire pa_decode_ps_launch's block_size 16/64 path to call pa_decode_tile directly instead of the separate compile_pa_decode_ps kernel + reduce round trip. Fixes a real query-stride bug uncovered along the way (the kernel assumed contiguous query, but callers commonly slice it out of a combined qkv tensor) by threading explicit row/head strides through instead of copying. Generalizes the epilogue to support any query_length*query_group_size (not just values dividing BLOCK_THREADS), adds CUDA graph capture support (caller-preallocated pmax/psum/pout, mirroring the other PS paths), and adds trans_v=False V-cache support. Merges the tile-specific partition-count heuristic into get_recommended_splits instead of keeping a separate _choose_num_partitions. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: fsx950223 <fsx950223@outlook.com>
pa_decode_tile now handles every small-block (16/64) shape pa_decode_ps_launch routes to it, so the compile_pa_decode_ps fallback and its _pa_small_block_load_k_flat/_pa_small_block_load_v_trans helpers were unreachable dead code. Drop them along with the imports they alone needed. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: fsx950223 <fsx950223@outlook.com>
Reduce comment volume in pa_decode_tile.py while preserving the load-bearing rationale (bank-conflict math, VGPR/perf numbers, LDS aliasing, correctness-critical orderings). Remove the test_pa_decode_tile_* direct-path test suite and its _run_pa_decode_tile_case helper now that pa_decode_ps_launch routes block_size 16/64 through pa_decode_tile and test_multi_case_set exercises that path. Note: trans_v=False, wide-GQA (group_size>16), head_dim=64, and MTP+wide-GQA are not currently covered by test_multi_case_set's config grid and have no test coverage until added there. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: fsx950223 <fsx950223@outlook.com>
Applies a set of safe, individually-verified fixes for the M_TILES>1 (MTP query_length>1 / wide-GQA group_size>16) path, focused on block_size=16 where PAGES_PER_CHUNK=4 page-scatter previously left the most headroom: - Re-tune EPI_GROUP_TILES epilogue staging/barrier granularity now that the hoists below shifted the register profile. - Add VMEM-load scheduling hints (matching production's own small-block K-loader) to the K, V, and per-token K/V-scale gathers for block_size==16, where 4 separate physical pages are gathered per warp-chunk. - Hoist M-tile-independent per-token K/V-scale reads and V loads out of the per-M-tile loop for block_size==16 (redundant re-reads across M-tiles cost more there than the extra live range; gated off for block_size==64, where the same hoist crosses a VGPR occupancy cliff). - Drop a always-true-except-last-M-tile runtime branch in Q quantization in favor of a compile-time check. - Batch both vh V-data loads before their MFMA consumers for M_TILES==1, closing a latency-hiding gap unique to the single-M-tile case (no other M-tile's work to hide the load behind). - Switch `.maximumf()` (NaN-propagating) to `arith.maxnumf()` (maps to a single hardware instruction) in the softmax reduce chain. - Dedupe a redundant `exp2_amdgcn_scalar(m_prev - m_new)` recompute in per_token_kv mode (kept per-tensor's original recompute -- hoisting it there measured a severe regression from sitting exactly at the 256-VGPR/2-wave occupancy boundary). Net effect: block_size=16 narrows from ~1.10-1.40x to ~1.07-1.20x vs the original small-block PS kernel across MTP shapes; block_size=16 M_TILES==1 cases move to parity or better; block_size=64 is unaffected or improves further (~0.78-0.95x). Full closure for M_TILES>1 with block_size=16 remains open -- ATT/ISA analysis traces the residual gap to a diffuse mix of costs structurally inherent to the 4-page gather (no single dominant hotspot), not a further taggable inefficiency. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Signed-off-by: fsx950223 <fsx950223@outlook.com>
CI on PR #825 started failing test_multi_case_set[sliding_window_accuracy] (NaN in "FlyDSL PS vs Torch") on both gfx942 and gfx950 after the small-block PS launcher was switched to pa_decode_tile and MTP (query_length>1) support was added. Root cause: for an early MTP query row whose causal window ends before a given 256-token compute tile (common with kv_varlen once the tile count isn't a multiple of the partition count), every entry in that tile is masked to the -1e30 sentinel. The per-row running max is scaled by that same sentinel before storage, so the masked score and the running max are numerically identical -- `exp2(score*scale - m_new)` cancels to `exp2(0) == 1` instead of ~0. In per_token_kv mode this also drives the V-scale normalization factor to a huge value (zero valid V-scale, epsilon-floored denominator), so `Pa=1 * huge` overflows the fp8 P-pack into NaN. That NaN reaches the cross-partition reduce kernel as a partition's "logits" value; even a correctly near-zero combine weight doesn't save it, since `0 * NaN == NaN`. Fixes: - kernels/attention/pa_decode_tile.py: re-mask the softmax probability itself (not just the pre-exp score) so a tile with zero valid tokens for a row always contributes exactly 0, regardless of the score/max cancellation. - kernels/attention/pa_decode_swa.py: defensively clamp the cross-partition reduce's exp2 argument, since an upstream partition can still report a very large-magnitude (but finite) "empty" max logit rather than literal -inf, and the fastmath exp2 lowering returns NaN instead of underflowing for such inputs. Verified on a gfx950 (MI355) box: reproduced the exact failing case standalone, confirmed NaN, applied both fixes, reran -- test_multi_case_set[sliding_window_accuracy] and [normal_accuracy] both pass.
… and cross occupancy cliff Reduce pa_decode_tile's LDS instruction count to at/below production pa_decode_ps_kernel and fix the per_token MTP/wide-GQA (M_TILES>=4) slow case. - Swap the PV MFMA operand order (V=A, P=B) so the output fragment is [head-dim, query-row=lane16] like production. The correction factor and denominator become per-lane scalars (no sCorr LDS round-trip) and O stores straight to global (no sO staging, no epilogue barrier). Empirically probed the MFMA lane wiring first to de-risk the layout. Gated to the wired config (block16, head_dim128, M_TILES>1). Result: per_token ds ops 103->77 (baseline 78), per_tensor 85->59 (baseline 68). - Dedup the block16 per_token KV-scale gather (was looping all 4 rgroups redundantly): 4x->1x loads via rgroup-selects-page. - Unmask pv_max: it is only a per-tile fp8 normalization constant, so any positive value is correct -- drop the causal-mask select (-64 cndmask, -48 v_max), and it becomes m-independent. - Double-buffer the KV-scale LDS staging (ping-pong by tt&1) so Phase B can re-read v_scale after the prefetch instead of holding k+v across both phases; split k (Phase A) / v (Phase B) loads, and for M_TILES>=4 re-read v_scale per-chunk. Drops per_token M_TILES=4 from 272 to 256 combined VGPR, crossing the 256/2-wave cliff: 1.24x -> 0.77x vs baseline. - Remove the always-1 EPI_GROUP_TILES grouping machinery in the epilogue. Correctness verified fresh-compiled across block16/64 x per_token/per_tensor x query_length 1-4 x group_size 8/16 (incl partial tiles) x batch 3/128 x context 1027/8192, plus both test_pa standard gates. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Extend the swapped-operand PV (V=A, P=B -> output [head-dim, query-row=lane16] + direct-store epilogue) to every config and remove the old [query-row, head-dim] path entirely. The swap/direct-store generalizes over any head_dim (multiple of 64) via the VHE_CHUNKS loop (OP_ELEMS is always 4), so no per-config gating is needed. Removed (~330 lines): - The old-PV branch (P=A/V=B) in both the phase-split and non-phase-split paths. - The sO-staging epilogue and its tiled_mma_pv / copy_c / tcopy_o / tmpl_Op / EPI_PARAMS / _epi_params machinery; only the direct-store epilogue remains. - The sO, sM, sL, sCorr LDS regions (output staging / per-row correction / reduce scratch) -- all dead once PV is register-resident + direct-store. KV-scale staging gets its own region (was aliased into sO's range). - The head_dim==64 v_vh_early early-load (old-PV only). Wins beyond simplification: - block64 per_token M_TILES=4 crosses the 256-VGPR/2-wave cliff (260 occ 1 -> 244 occ 2): ~1.36x -> ~0.94x vs baseline. block16 unchanged (same-session A/B), block64 per_tensor neutral. LDS footprint drops (sO/sM/sL/sCorr gone). - Also ports the fully-invalid-tile Pa re-mask to the unified path. Correctness verified fresh-compiled (cache off) across block16/64 x per_token/per_tensor x query_length 1-4 x group_size 8/16 (incl partial tiles) x batch 3/128 x context 1027/8192, plus both test_pa standard gates. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The phase-split (Phase-A-all-m -> one shared pass-1 barrier -> Phase-B-all-m, with the KV-scale double-buffer and V/q-scale hoists) was gated block_size==16 only, because on the pre-unification kernel its M-tile-independent hoists crossed the 256-VGPR/2-wave occupancy cliff for block64. After the single-PV- layout unification lowered the register baseline, that cliff is gone: block64 M_TILES=4 stays at occupancy 2 with the phase-split (per_token 252 VGPR, per_tensor 238), and gets fewer LDS instructions from the barrier merge (per_token ds 108->92, per_tensor 90->59). Flip the coupled gates (PHASE1_MTILES, KV_DOUBLE_BUF, v_vh_shared, q_scale_vec, and the non-phase-split fallback ternaries) from `block_size==16` to `head_dim!=64`, so block64 M_TILES>1 (head_dim==128) now takes the phase-split path too. Also removes the now-dead kv_scale_vecs_shared hoist + its ternaries (the non-phase-split path only serves head_dim==64 / M_TILES==1 now, where those conditions are always false). Perf (settled-GPU A/B, block16 as stable control): block64 per_token M_TILES=4 ctx8192 ~0.94x -> ~0.80x vs baseline (~14% faster); block64 per_tensor neutral; block16 unchanged. Correctness verified fresh-compiled across block16/64 x per_token/per_tensor x query_length 1-4 x group_size 8/16 (incl partial tiles) x batch 3/128 x context 1027/8192, plus both test_pa standard gates. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… M_TILES) The phase-split was gated head_dim!=64 out of caution (head64 has no CI coverage and, on older kernels, a distinct V-load path). With the unified single-PV layout, head64 M_TILES>1 works on the phase-split path unchanged (VHE_CHUNKS==1, v_vh_shared=[_v_ops(...,0)]), so drop the head_dim!=64 gate: PHASE1_MTILES is now simply M_TILES (phase-split for every M_TILES>1 shape; M_TILES==1 naturally routes to the per-m path). The non-phase-split path is now M_TILES==1-only. head64 was already occ 2 both ways, so this is a pure barrier/LDS-instruction win from the merge: per_token ds 116->84, and (settled-GPU, run_new abs time) M3 per_token ctx1027 57.1us->46.8us (~18%), M4 per_token ctx8192 386us->349us (~10%), M4 ctx1027 69us->66.5us; per_tensor neutral. head_dim==128 unchanged. Correctness verified fresh-compiled for head_dim=64 across block16/64 x per_token/per_tensor x query_length 1-4 x group_size 8/16 (incl partial tiles) x batch 3/128 x context 1027/8192 (hand-run; head64 isn't in the default CI matrix), plus the head_dim=128 standard gates. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
After the phase-split was extended to every M_TILES>1 shape, the non-phase-split branch is reached only for M_TILES==1. Simplify it to make that explicit and drop the dead multi-tile machinery: - Replace the `for m in range_constexpr(M_TILES)` loop with `m = 0` (single tile), so the branch reads as the straight single-tile path it now is. - Drop the always-true `if m == 0:` guards around the K/V/scale prefetch and the V-page readback. - Drop the always-true `M_TILES == 1` conditionals on the V-load: it's unconditionally the batched-upfront load now (the `_v_ops` per-vh else was the dead block64/M>1 case, which is phase-split). - Fix stale comments. Pure dead-code removal, behavior-preserving: M1 codegen is byte-identical (VGPR 124 / occ 4 / 30 ds ops, unchanged) and M>1 (phase-split) is untouched. Correctness verified fresh-compiled across head_dim 64/128 x block16/64 x per_token/per_tensor x query_length 1-4 x group_size 8/16 x batch x context, plus both test_pa standard gates. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…branch In the M_TILES==1 branch the loop-carried state slots are compile-time constants, so calling _o0_slot(0)/_o1_slot(0)/_m_slot(0)/_l_slot(0) (each 2+4*0, 3, 4, 5) to compute them is needless indirection. Use the literals ostate[2]/[3]/[4]/[5] directly, with a comment documenting the layout and its correspondence to the _oN_slot/_m_slot/_l_slot helpers (still used, with variable m, in the phase-split and epilogue). Behavior-preserving (the helpers already folded to the same constants): M1 codegen byte-identical (VGPR 124 / occ 4 / 30 ds ops). Standard gates pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The direct-store epilogue guarded each store with a compile-time full-tile check (row_ok) OR-ed with a runtime bounds check, written as nested `if row_ok: emit else: if row < TOTAL_ROWS: emit`. Collapse to a single `if row_ok or row < TOTAL_ROWS:`. Python short-circuits the `or`: when row_ok is True (full tile, the common case) the runtime `row < TOTAL_ROWS` comparison is never traced -> the store stays unconditional (no scf.if); only a partial last tile evaluates the runtime guard. Same for the NP>1 pmax/psum writes. Behavior-preserving: full-tile codegen byte-identical (M4 VGPR 256 / occ 2 / 88 ds), and correctness verified across full + partial (group_size 8, ql 3 -> TOTAL_ROWS=24) tiles, NP==1 and NP>1, both quant, block16/64, plus the standard gates. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
row_ok = (m+1)*MFMA_MNK <= TOTAL_ROWS implies row < TOTAL_ROWS for every lane (row = m*MFMA_MNK + lane16, lane16 <= 15, so a full tile's max row (m+1)*16-1 < TOTAL_ROWS), so `row_ok or row < TOTAL_ROWS` was logically just `row < TOTAL_ROWS`. It's also redundant in codegen: the backend already folds the full-tile guard to always-true (lane16 = tid & 15 <= 15), so the ISA is byte-identical with or without row_ok (M4 full tile: VGPR 256, 2459 instr either way). Drop row_ok and the `or`. Correctness verified fresh-compiled including partial-tile cases (group_size 8, query_length 3 -> TOTAL_ROWS=24, high lanes out of range), NP==1 and NP>1, both quant, block16/64, plus the standard gates. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Once the phase-split was extended to every config, PHASE1_MTILES was defined as plain `M_TILES` -- a redundant alias. Inline it: sLmax sizing, _lmax_off_m, and the phase-split gate now use M_TILES / `M_TILES > 1` directly. Pure rename, behavior-identical (M4 phase-split VGPR 256, M1 VGPR 124 both unchanged); standard gates pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
_tile_tok0_and_page(tt) just returned (tt*TILE_TOK, tt*TILE_TOK // block_size). Inline it at its three call sites: tok0 = tt * TILE_TOK in the loop body, and base_page = tt*TILE_TOK // block_size in the K and V page-prefetch loaders (the tile start is page-aligned since TILE_TOK is a multiple of block_size). Pure inline, behavior-identical (M4 codegen byte-identical: VGPR 256, 2459 instr); full block16/64 sweep + standard gates pass. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… double-buffer, trim comments
Address the todos.md cleanup list:
- Route every arith.FastMathFlags.contract use through a single `fm_contract`
alias defined once at kernel scope (was one local def + 5 inline literals).
- Derive n_kv (num_kv_heads) from `gpu.grid_dim.y` instead of the runtime
`num_q_heads // query_group_size`, and drop the now-dead `num_q_heads` kernel
arg from the kernel/launch-wrapper/host-call signatures.
- KV_DOUBLE_BUF = per_token_kv (was `and M_TILES > 1`): the M_TILES==1 path now
ping-pongs its KV-scale staging too, so the tt+1 prefetch no longer aliases
the current tile's read in the single buffer. M1 stays VGPR=124/occ-4 (the
extra buffer is LDS-only, not the M1 limiter); M4 unchanged at 256/occ-2.
- Trim/condense verbose and stale comment blocks (273 -> 236 comment lines).
Investigated the block64 M4 per_tensor-vs-per_token VGPR question: current code
is per_token=252 > per_tensor=238 (both occ-2), the expected order since
per_token carries the k_scale/v_max/norm-factor KV-scale registers per_tensor
lacks. The todo's inverted observation predates the double-buffer/phase-split
changes and no longer reproduces.
Correctness: full block16/64 x {per_token,per_tensor} x ql1-4 x
num_kv_heads{1,2,4} sweep + normal/sliding_window accuracy gates pass; ruff +
black clean.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.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
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.
Summary
Introduces
kernels/pa_decode_tile.py, a correctness-first, readabletile-programming reimplementation of paged-attention fp8 decode, built with
FlyDSL's tile/layout API instead of the raw buffer intrinsics and
hand-scheduled MFMA that production's
pa_decode_ps_kernel(
pa_decode_fp8.py) uses.Across an extensive optimization pass (raw dwordx4 K/V loads, blocked
K/V cache layouts matching production's own, register-resident softmax
and PV accumulation, an LDS bank-conflict fix,
fastmath=contract,rocdl.sched_group_barrierscheduling hints, and reusing production's ownreduce kernel), the tile kernel now matches or beats production's
performance at the batch/context sizes that matter for real serving
workloads, while remaining significantly more readable than the
hand-scheduled production kernel.
Also includes a couple of correctness fixes surfaced along the way:
rcp_f32(0)returns+inf, multiplying into aNaNthat survives thereduce kernel's zero-weighting since
0 * NaN == NaN) — fixed with thesame
safe_sum/inv_sumclamp pattern production already uses.outputdtype now correctly followsquerydtype (was hardcoded to f16regardless of query dtype).
Performance
GPU kernel time only (rocprofv3
--kernel-trace, min-of-5 per sample,interleaved A/B, 6+ samples/shape), ratio = tile / production (
<1.0=tile faster).
bs= batch size,ctx= context length.head_size=64, ctx=16384, across batch sizes (fresh measurement):
head_size=128, bs=128, ctx=16384 (fresh measurement):
head_size=64, across batch size x context length (broader grid, same
committed heuristic/kernel):
The one remaining soft spot is low-batch + short/medium context (bs=3,
ctx=129/2048), where fixed per-CTA prologue overhead dominates an
already-tiny (~7-9us) kernel; investigated and not pursued further since
it's a small absolute delta (~1-1.5us) at a batch size that isn't
representative of real decode-serving workloads.
Test plan
pytest tests/kernels/test_pa.py— 42/42 passing (referencecorrectness across block_size/head_dim/query_dtype/context_len grid,
plus the PS-vs-Gluon accuracy sweep)
black --check/ruff checkclean on all touched files