Repository navigation
Conversation
… load scheduling, lean FP8 MTP4 gfx950-only tile-kernel changes, each behind a traits flag that requires is_gfx950: - BF16_MTP_FUSED: BF16 multi-query (MTP4 verify) page-128 D128 tiles keep Q in padded LDS rows, give every M-tile a private P/Lsum slot (one publication barrier), carry V across M-tiles and pair K16 packs into gfx950's K32 BF16 MFMA. The resident set fits 256 VGPRs (2 CTAs per CU). The task's last tile is peeled and carries the V tail mask; unowned KV rows stay masked. - BF16_DECODE_NT_KV / DECODE_Q_EARLY / DECODE_LAST_V_AFTER_QK: BF16 single-M decode streams K and V nontemporally, issues Q ahead of the page ids, and loads the last tile's V right after QK. - FP8_SCALES_AFTER_Q / FP8_LOOP_SCALES_LATE: per-token FP8 page-128 tiles stage scales so their vmcnt waits stop draining K and V ahead of the softmax. - FP8_MTP4_LEAN / FP8_MTP4_Q_LDS / FP8_MTP4_PAGES_AHEAD: FP8 MTP4 verify folds the K scale and the V normalization into exp2 and the PV MFMA's P scale, rereads Q from LDS and carries the next tile's page ids, with waves_per_eu=2. - The gfx950 BF16-KV tile launch preloads its leading kernel arguments. gfx942 and gfx1250 code paths are unchanged (gfx942: every compiled PA kernel bit-identical). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…in one rung On gfx950 the dual-mode reducer replaces the (1, 2, 4, 8, 16, 32) rung ladder with one 32-partition rung. LLVM hoisted the ladder's shared leading loads above the count branch, waited on them there and trickled the rest in small batches: 32 live partitions took several serial L2 round trips. One rung issues every load at once; partitions past the live count read zeros from the per-sequence buffer bound. Above 32 live partitions, each group's loads now issue before the weight barrier, as on gfx1250. Trade-off: sequences with few live partitions now pay the 32-load rung (see the PR description). gfx942 and gfx1250 are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…e-16 tiles gfx950 scalar-scale (per-tensor) FP8 page-16 single-M tiles now take the existing V-prefetch pipe (PAGE16_VPIPE) that per-token FP8 page 16 already uses; before, only page 128 qualified, so page-16 tiles issued V after the softmax barriers. Single-CTA page-16 cost drops from ~2.7 us per 256-token tile. Paired in-process A/B vs the previous head on MI355X (FlyDSL 0.5.0): 32 page-16 review shapes 1.035x geomean (single-CTA 30K context 1.53x, H8 B32 1.18x), outputs bit-exact; other paths unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
gfx950 BF16 page-16 single-M decode tiles take the V-prefetch pipe as well (they had no V prefetch: ~3.7 us per 256-token tile in a single CTA). Paired A/B on MI355X: single-CTA 30K context 1.58x, random-length B16 1.04x, B1 short contexts +1-2%, large grids -0.4..-0.8%; outputs bit-exact. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…128 decode gfx950 FP8 single-M page-128 decode loads V with nontemporal flat loads (FP8_DECODE_NT_V), as BF16 decode already streams K and V. K stays cached: streaming K as well slowed the short single-tile tasks, and on page 16 streaming fights the V-prefetch pipe (single-CTA 30K -21%). Paired A/B on MI355X over the FP8 page-128 decode shapes: 1.033x geomean (b64 32K 1.18x, b4 32K 1.13x, b32 8K 1.05x), no shape below 0.998x; outputs bit-exact. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…d of rebuilding them pa_decode built a fresh one-element scale tensor on every call for scalar FP8 KV scales (torch.ones for None, torch.tensor for a Python float). Under a CUDA graph each is a captured fill or host copy: a scale-less FP8 call ran 4 dispatches instead of 2. The scale is now one persistent tensor per (device, value), cached only when built outside graph capture (a captured fill has not run yet and lives in that graph's pool), so captured-first calls keep working. Host-only; kernels unchanged on every architecture. Paired A/B on MI355X over the 51-shape review set: 1.245x geomean; short scalar-FP8 calls 9.9 -> 6.6 us; outputs bit-exact. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… from the tile kernel A plan with max_partitions == 1 gives every sequence at most one task (capacity == batch), yet gfx950 FP8/BF16 still wrote packed partials and launched the reducer to copy them out. The tile kernel now stores the normalized rows itself, as gfx1250 and MXFP4 already did: slot s runs packed task s and zero-fills sequence s when it owns no task (reduce_info count 0), so the second dispatch and the partial scratch round trip go away. Sinks and strided outputs keep the reducer. The direct store takes the query's row/head strides for any query length. Paired A/B on MI355X: W6 fp8 page-16 B64 mp1 1k 10.55 -> 8.65 us (1.21x), W5 bf16 B64H8 mp1 1.07x, R7 single-CTA 30k 1.01x; outputs bit-exact; other shapes compile the same kernels. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… V; FP8 issues next K after QK gfx950 scalar-scale page-16 tiles issue this tile's V at tile start, then load the next tile's page ids from the block table and stage them in LDS. In-order vmcnt makes that staging wait drain V as well, so V overlapped only QK, and the next K went out only after the P barrier, with PV left to hide it. The next page ids now load ahead of V (they are masked past the task and staged only under the existing next-tile guard), and FP8 issues the next K from them right after QK. BF16 keeps K after the P barrier: its next K needs twice the registers and was 4% slower per tile. Single-CTA cost per 256-token tile on MI355X: FP8 1.693 -> 1.412 us, BF16 2.227 -> 2.152 us. Paired A/B: R7 single-CTA 30k 1.20x, R5b budget=CU 32k 1.16x, W6 B64 mp1 1.10x, R1 128k 1.04x, P2 H8 B32 1.03x; outputs bit-exact; other shapes within +-1% (re-measured candidate-first). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…e bound The planner gives every nonempty sequence one task, then spreads the other slots in proportion to tiles. With mixed lengths and about two slots per sequence the long sequences keep the longest tasks: 16 tiles on a P2 H8 B32 random 1k-8k batch whose average is 8.5, so the kernel ends on that tail. On gfx950, single-block plans now give a sequence of t tiles min(ceil(t / T), max_partitions) tasks, with T the smallest bound from ceil(total / capacity) up that fits the capacity. A binary search over [ceil(total / capacity), ceil(total / (capacity - rows))] finds it; one block reduction per halving, int32 math, skipped for full plans. Equal-length batches keep their task size and may use fewer slots, which also shortens the reducer. Paired A/B on MI355X (decode only, plans built outside the graph): H2 D256 H8 B32 1.16x, P2 B32 1.07x, H1 B32 1.05x, MiniMax b3 225k four shapes 1.02-1.05x, P4 B8 1.03x; P1 B128 0.99x; 87-shape geomean 1.006-1.008x. Planner refresh +0.06..0.51 us. test_pa_decode_plan asserts the min-max counts when the planner compiles the balanced variant and the proportional split otherwise. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ssion The one-partition direct output generalized the direct store's address to query lengths above 1. Single-query kernels, including the MXFP4 plan-free ones that already stored directly, now keep the original expression, so their generated code stays as it was; only longer queries take the position/head split. Paired A/B: W3/W4 MXFP4 1.000x/1.001x, W5/W6/R7 within noise, outputs bit-exact. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… gfx942 On gfx942 the D1024 FP8/BF16 and BF16 D512 MTP tile kernels hold the whole next-tile K prefetch and the tile's whole V gather in registers. They hit 512 VGPRs and spill 350-2000 registers to scratch. Where a lane's K row reaches 384 bytes (BF16 D192+, FP8 D512+; with several M-tiles, 1 KiB), the tile now carries only page ids. K is streamed in pieces of four head groups (two in flight) and V in head chunks (four in flight), with a scheduling fence after each chunk. M-tiles pack P into private slots behind one barrier and share every K/V chunk. Spills drop to zero. gfx950 and gfx1250 binaries are unchanged. Paired in-process A/B vs main on MI308X (gfx942, 80 CU): - 70 D1024 cases: 1.87-6.14x, all bit-exact - BF16 D512 MTP4: 2.30-4.18x, same fp32-reference error - BF16 D256 page16: 1.28-1.41x at batch >= 16, 0.81-0.94x at batch 1 - BF16 D192 page64: 1.21-1.27x at batch >= 16 107 cases: geomean 2.48x. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…nes use DPP The reducer's softmax weights reduce each row's partition max and sum across lanes with shuffle_xor. On gfx950 the 8/4/2/1-lane steps lowered to ds_swizzle, an LDS crossbar round trip with a full lgkmcnt wait each: eight serial ~44-cycle waits per reducer wave, on the critical path of every short call. Those steps now use the 16-lane DPP xor patterns (the same partner lane, so results are bit-identical); the 16/32-lane steps keep permlane swaps. Other architectures are unchanged. Paired A/B on MI355X over 87 shapes: 1.0069x, all outputs bit-exact; short calls 1-3% (MiniMax fp8/bf16 b1 decode, R1/R5b/H1 page 16), the lowest shapes within +-0.4% when re-measured alone. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The gfx9 tile issues a single M-tile's V loads only after the P barrier, and PV consumes them right away. The load latency therefore sits on every tile's critical path. ATT on MI308X showed the PV MFMAs waiting on vmcnt one V load at a time. gfx950 already prefetches V at the tile start for its tuned FP8 shapes; gfx942 never did. On gfx942, single-M-tile FP8 (up to D256, any page size) and BF16 page16 D128 now take the existing prefetch_v schedule. V is loaded at the tile start and overlaps QK and softmax. BF16 page128 and BF16 D64 measured neutral and are left as they were; BF16 D192+ streams instead. gfx950 binaries are unchanged. Paired in-process A/B vs c52a931 on MI308X (gfx942, 80 CU): - 135 single-M-tile grid cases: geomean 1.052x - FP8 page16 D128: 1.02-1.42x - BF16 page16 D128: 1.04-1.21x - FP8 page128 vectorized V at batch 1: 0.98-0.99x - 111 unit-test cases: geomean 1.024x, worst 0.96x Per-token outputs move by at most 1e-6 relative; the max error against the fp32 reference is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…rids When more CTAs than CUs hold work, occupancy matters more than one CTA's latency. When fewer do, the reverse holds. The schedule now records whether the host-bounded busy-CTA count (the existing split_workgroups bound times query splits) exceeds the CU count, and uses that on gfx942: - FP8 D128 MTP tiles (2-4 M-tiles) need 260-272 VGPRs, i.e. one wave per SIMD. On busy grids they compile with waves_per_eu=2. That caps them at 256 VGPRs at the cost of about ten spills, and two CTAs share each CU. - Single M-tile K/V streaming between 384-byte and 1 KiB K rows (BF16 D192/D256) now runs only on busy grids. On small grids the K prefetch it gives up was the latency that mattered: BF16 D256 batch 1 was 0.81x vs main. From 1 KiB rows, kernels always stream, because the unstreamed kernels spill. gfx950 and gfx1250 binaries are unchanged. Paired in-process A/B vs a5a617e on MI308X (gfx942, 80 CU), 310 cases: geomean 1.077x. - FP8 D128 MTP4/5/8 at batch >= 16 (or batch 64 ctx1K): 1.13-1.40x - BF16 D256 page16 batch 1 ctx 1K/8K: 1.23x, back to 1.0x vs main - Small grids: 0.97-1.00x. The worst is a long-row UT case (mtp4-g16-page128-varied), which is busy by the host bound. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…tition decode A one-partition plan gives each sequence a single CTA that walks all of its tiles; on sparse grids that is a latency chain of four waves, one per SIMD. On gfx950, when the grid has at most one CTA per CU, single-query decode tiles (scalar FP8 or BF16, no window, not the head-256 pipe) now launch eight waves as two four-wave groups. Group g walks tiles start + g, + 2, ... with its own copy of the LDS layout; both run the same iteration count (group 1's last tile may lie past the task and is fully masked), so their CTA barriers pair up. After the loop group 1 publishes its (O, max, sum) per lane and group 0 rescales both to the shared maximum and stores, directly or as the task's partial. The grid condition joins the schedule key ahead of its last two fields, which tests read. Single-CTA cost per 256-token tile on MI355X: FP8 page 16 1.42 -> 1.15 us, BF16 page 16 2.15 -> 1.54 us, FP8 page 128 1.69 -> 1.41 us. Paired A/B: R7 single-CTA 30k 1.19x, W6 B64 mp1 1.04x; all other shapes compile the same kernels. FP8 P is quantized against each group's running maximum, so those outputs move within FP8 rounding (max 3.4e-2 of the output max). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… when both LDS copies fit MXFP4 tiles have no per-tile CTA barrier: each wave owns 64 tokens and the four waves merge once after the loop, so a plan-free sequence is one CTA's latency chain. The dual-group CTA now covers it when two copies of the MXFP4 layout (BF16 V staging and staged page ids) fit gfx950's 160 KiB of LDS. The groups merge lane-wise (O, max, sum) before group 0's four-wave merge, and a barrier keeps that merge's LDS writes behind every lane's read of the merge slots, which share the space. Paired A/B on MI355X: W3 MXFP4 plan-free 64k 316.5 -> 226.2 us (1.40x); W4 (a 512k table stages 4096 page ids, so it keeps one group) unchanged; every other shape compiles the same kernels. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… max Every FP8 tile quantizes Q to e4m3 per row, first reducing |q| over each lane's 8-element BF16/FP16 unit. LLVM widened every element to f32 for the max: about 30 VALU per unit on the prologue's critical path, 120 per lane for an MTP4 tile. A 16-bit float's magnitude orders like its sign-cleared bit pattern, so on gfx950 the unit now reduces as a packed 16-bit integer max over sign-cleared words, with one conversion of the winner. The result is the same value. Paired A/B on MI355X over 87 shapes: 1.0045x, all outputs bit-exact; FP8 MTP4 verify b3 32k 1.03x, short page-16 FP8 calls ~1-2%. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…de tiles Single-query decode without a sliding window can only see tokens past a row's bound in a task's last tile, which the pipeline already peels for the V tail mask. The head-256 pipe still compared and selected every score chunk of every tile. Interior tiles now skip that mask (single_scores and pack_single take the tile's mask flag); the peeled tile and every other configuration keep it. On the page-16 V pipe the masks are free and the change measured 1.5% slower, so it is left out there. Paired A/B on MI355X: X2 D256 page-64 per-token B8 mp1 1.12x, X1 B1/B8 1.06x/1.05x, X3 page 16 1.04x; 87-shape geomean 1.0015x. Without the select, the multiply contracts into the exp's FMA, so head-256 outputs move within rounding (max 2.4e-3). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…h integer DPP max After each lane's |q| max, the row's 16 lanes combined theirs with four xor DPP steps of maxnumf. Each step cost a copy, one or two v_mov_b32_dpp, an s_nop and a canonicalizing v_max_f32 besides the max: about 22 VALU per M-tile on the prologue's critical path. The magnitudes are non-negative, so on gfx950 the reduction now runs on their bits with integer max. Rotate-by-8 and rotate-by-4 within the row and the two quad swaps cover all 16 lanes, each a single full-mask DPP pattern that LLVM folds into the max. The result is the same value. Paired A/B on MI355X over 87 shapes: 1.0024x (FP8 shapes 1.0037x), all outputs bit-exact; FP8 MTP4 verify b3 32k 1.025x. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…es' page ids Plan-free MXFP4 stages the task's page ids in LDS, up to 4096 of them when the table is long. Two copies of that next to two BF16 V staging areas exceeded gfx950's 160 KiB, so long tables kept one tile group. A dual group only walks tiles part_start + 2j, so it now stages just those ids (each tile's pair contiguous) and looks a tile up at ((tt - part_start) / 2) * pages_per_tile; the page region of each copy halves and both copies fit. Paired A/B on MI355X: W4 MXFP4 plan-free 2k context in a 512k table 14.10 -> 12.53 us (1.13x); W3 and every other shape unchanged and bit-exact. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…head of the page ids Single-M FP8 page-128 decode with a full 16-row tile already issues Q before the page ids and stages it while K is in flight; scalar page-16 tiles still loaded Q after them. Page 16 now takes the early path too, including partial tiles (GQA 8): a padded row's early load reads the last real row so it stays inside the query, and the existing runtime row guard never stages it. Paired A/B on MI355X over 87 shapes: 1.0046x, all outputs bit-exact; FP8 page-16 shapes 1.010x (short calls up to 1.03x). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…6 page-16 decode The scalar head-256 pipe and BF16 page-16 single-M decode tiles now issue Q ahead of the page ids and stage it while K is in flight, as FP8 head-128 decode does (partial tiles load the last real row for their padded rows). Per-token head-256 keeps Q after its scale loads: its prologue orders the scales ahead of K, and early Q cost it 10%. Paired A/B on MI355X over 87 shapes: 1.0023x, all outputs bit-exact; short head-256 calls 1.04-1.05x (P3 B1/B8, W1 B1, R5 4k), BF16 page-16 R2 up to 1.025x. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ruff 0.16.0, the version CI pins, flags the dpp_utils import placed after tensor_shim (I001). Import order only; no behavior change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Black wrapped the condition, which left the noqa on the closing line; ruff 0.16.0, the version CI pins, then reports SIM102 and an unused noqa. Comment placement only; no behavior change. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Brings in #7388 (FlyDSL-dtype tile kernel keys, no batch-first plan grid). This branch's traits and kernels now use its names: kv_dtype/query_dtype FlyDSL types, OUT_DTYPE, KV_COMPUTE_BLOCK and MXFP4_KV_DTYPE in place of MXFP4/is_bf16_kv/Q_IN_DTYPE/Q_DTYPE/KV_DTYPE/TILE_TOK; the gfx950 one-partition direct-output allowance moves into the schedule's direct_output check with #7388's. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The gfx9 prologue issued a CTA's Q loads after the first tile's page-table and K loads. Q staging (quantize, LDS write, barrier) therefore waited on the in-order vmcnt chain for every K load: ATT on MI308X showed a 1056-cycle wait in a 9500-cycle single-tile task. gfx942 now issues Q loads first, as gfx950 already does for MTP4-fused and MXFP4 kernels. Padding rows of a partial last M-tile load the last real row (clamped index) and staging still zero-fills them. Streamed kernels start with K pieces and keep their order. gfx950 binaries are unchanged. Paired in-process A/B vs 9428a09 on MI308X (gfx942), 111 unit-test cases: geomean 1.016x (FP8 1.01-1.05x, BF16 1.00-1.03x, streamed D1024 0.98-0.99x before excluding them). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… gfx942 Tiles that load V at the tile start (V prefetch, MTP M-tiles) issued the next tile's block-table read only later, behind those V loads. Its pin-to-context select then waited on vmcnt(0), i.e. for all of V; ATT on MI308X showed about 1000 cycles of a 4400-cycle FP8 page16 tile spent there. The page-table read now issues before the V loads, fenced so the scheduler keeps that order, and the pin happens at the use. The wait drops to vmcnt(8). Applied on gfx942 to FP8 per-tensor single M-tiles and to MTP tiles, except BF16 with two M-tiles. gfx950 binaries are unchanged. Paired in-process A/B vs the previous head on MI308X (gfx942): - FP8 per-tensor single M-tile: up to 1.18x (page128 plain V); page16 1.03-1.07x at batch 64 - MTP M-tiles: 1.00-1.12x - Excluded: FP8 per-token and BF16 page16 single tiles (0.95-0.98x) and BF16 two-M-tile MTP (0.95-0.98x) Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…busy gfx942 grids V prefetch at the tile start keeps the V registers live through QK and softmax. For FP8 per-tensor page128 with a vectorized V cache, that cost more than the hidden latency on small grids: 0.955-0.967x on the latency-bound unit-test shapes. It still gains 1.11-1.15x on busy grids. Such tiles now prefetch only when more CTAs than CUs hold work. Per-token, plain-V and page16 tiles are unchanged. Paired in-process A/B on MI308X (gfx942): small_query/probability_tail page128 UT cases 1.04x, busy grids unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…D256 on gfx942 A streamed BF16 D256 page16 tile needs just over 256 VGPRs, so it ran one wave per SIMD. Capped at 256 with waves_per_eu=2 (254 VGPRs, no spills), two CTAs share each CU. That only matters on busy grids, where this path is selected. Paired in-process A/B on MI308X (gfx942): BF16 D256 page16 B64 ctx32K 1.38x (main 0.94x before, now 1.30x vs main); other sizes unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… grids MTP tiles with two M-tiles and K rows of 256 bytes or more need 272-324 VGPRs (FP8 D256 GQA8 MTP4, BF16 D128 GQA8 MTP4). That leaves one wave per SIMD. On busy grids they now take the streamed multi-M-tile schedule from c52a931: K and V stream per chunk and both M-tiles share every chunk. That frees the registers held by the next-tile K prefetch and the tile's V gather, and two CTAs fit per CU. Small grids keep the prefetching schedule. Paired in-process A/B on MI308X (gfx942), busy sizes (batch >= 16, or batch 1 with a 32K context): - FP8 D256 page128 GQA8 MTP4 (2 KV heads): 1.18-1.28x - FP8 D256 page64 GQA8 MTP4 per-token: 1.23-1.52x - BF16 D128 page64 GQA8 MTP4: 1.17-1.20x Max error against the fp32 reference is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… undefined fp8_words packs four f32 values into one word with two cvt_pk_fp8_f32: the first writes the low half and the second the high half. The first conversion's tied old word was a literal 0, so every packed word cost a v_mov_b32 0 ahead of it (22 per FP8 MTP4 verify tile, on the P path). On gfx950 that old word is now undef; the result is bit-identical. gfx942 keeps the zero, so its kernels are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ts first loads' and preload them The reducer's first load is reduce_info, whose pointer sat at byte 0x50 of the kernel arguments; on gfx950 nothing was preloaded, so every call waited on a kernarg round trip before it could fetch its partition range. gfx950 now orders the reducer's arguments as reduce_info, the statistics and partials pointers, their strides, query_rows and num_kv_heads, which are the 14 dwords that preload into user SGPRs, and enables the preload. reduce_info is then fetched at once, with the remaining arguments in the same round trip. The kernel body is shared; gfx942 and gfx1250 keep their argument order. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
A CTA runs all query positions of an MTP row together. On small grids (few sequences, short contexts, sliding windows) most CUs are idle, and the CTA's serial latency (four M-tiles of QK, softmax and PV per tile) is the decode time. On gfx942 the schedule now splits the query positions over more CTAs when they fit; each re-reads the row's K/V, mostly from L2: - FP8 up to D128: one M-tile per CTA (4-way) while 4x the host-bounded busy CTAs fit about three per CU, else 2-way while two per CU fit. - BF16: 2-way while one per CU fits. Busy grids keep the unsplit occupancy schedules. Outputs and their fp32- reference error are unchanged, since rows are independent. Paired in-process A/B vs 9428a09 on MI308X (gfx942): - 111 unit-test cases: geomean 1.19x (MTP4/8 UT cases 1.4-2.4x) - 180 MTP grid cases: geomean 1.056x (batch 1 up to 1.45x; busy grids 1.00x) FP8 D256 (0.93-1.05x) is excluded. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…page64 on gfx942 FP8 per-token D256 page64 single tiles with V prefetch take the per-token M1 schedule, which was tuned on gfx950. On busy gfx942 grids it was 0.88-0.93x vs main (B16 ctx32K, B64 ctx8K/32K). On small grids it gains 1.01-1.07x. Busy grids keep the unprefetched schedule. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…sk slot per CU A multi-query (MTP) tile quantizes and stages all of its query rows and publishes their maxima before any KV work, a cost paid once per task. With the default budget of two slots per CU, MTP4 verify ran twice as many, shorter tasks, and the reducer merged twice as many partials. On gfx950, plans with query_length > 1 and no explicit workgroup_budget now default to one slot per CU. Paired over the 14 MiniMax-M3 MTP4 verify workloads: 12 faster (up to 1.14x), FP8 b3 225K 0.972x and b16 32K 0.993x; all 87 shapes 1.0084x. Explicit budgets, other architectures and single-query plans are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The tile grid is (plan capacity, KV heads x query splits). Dispatch walks x first, so split 1 of task 0 dispatches only after every slot of split 0, padding slots included. On small grids, where the plan has many more slots than tasks, the splits ran in sequential waves of mostly empty CTAs. On large grids a task's splits ran far apart, each re-reading the task's K/V from HBM. On gfx942, split kernels now launch (KV heads x query splits, capacity), so a task's splits dispatch together. They run concurrently and share the K/V they re-read. The cache key records the choice. BF16 splits forced by the LDS limit at large head dims only use it when tasks own at least four tiles; with fewer, the split-major grid measured 0.93-0.97x. Paired in-process A/B vs 8f9fe39 on MI308X (gfx942): - 180 MTP grid cases: geomean 1.092x - FP8 MTP4/8 batch 1: 1.04-1.45x - FP8 sliding-window MTP4 batch 1/16: 1.28-1.63x - BF16 D1024 MTP batch >= 16 with ctx >= 8K: 1.10-1.26x - 111 unit-test cases: geomean 1.018x (MTP4/8 and window cases 1.16-1.20x) gfx950 binaries are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…or on small grids The same dispatch-order effect as query splits: with grid (capacity, KV heads), a second KV head's tasks dispatch only after every slot of the first head, including the plan's empty padding slots. On small gfx942 grids the planned capacity far exceeds the tasks, so heads now launch head-major there. Busy grids keep the slot-major order, which measured slightly better for them (0.93-0.98x head-major). Paired in-process A/B vs ece6342 on MI308X (gfx942), two-KV-head geometries: batch 1 ctx 1K/8K 1.11-1.29x (FP8 D64 MTP3, FP8 D256 page16, BF16 D64 MTP4); other grid sizes and the unit-test cases unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…ngs after the weight barrier Above 32 live partitions the gfx950 reducer issued each group's partial loads before the weight barrier and held them in registers until then. With one partition group, which large batches get, that rung spans all 256 partitions: the kernel compiled to 266 VGPRs, one wave per SIMD, and every partition count ran at that occupancy. Batches of 128-256 sequences with 2-8 partitions each were 4-9% slower than main. gfx950 now issues loads ahead of the barrier only when a group's rung is at most 64 partitions; a single group keeps main's order and reuses the LDS allocated up front. That kernel drops to 51 VGPRs. Outputs are bit-identical. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…rids A sliding window bounds every MTP task to a few KV tiles. The task's serial latency therefore stays the decode time even on grids the idle-CU rule counts as busy. Windowed MTP tiles on gfx942 now split FP8 into single-M-tile CTAs on any grid, and BF16 two-way while four split CTAs fit per CU. Unwindowed rows keep the idle-CU thresholds. Paired in-process A/B vs db5c85c on MI308X (gfx942), 70 windowed MTP4 cases: geomean 1.045x (grid 1.050x) - BF16 window 257 batch 16: 1.32-1.34x - BF16 window 700 batch 64: 1.12-1.18x - FP8 per-tensor vectorized V batch 64: 1.18-1.20x Unit-test cases unchanged. The max error against the fp32 reference is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
BF16 D64 rows are 128 bytes, the same as FP8 D128, and their split CTAs stay small. BF16 D64 MTP rows now split two-way while the unsplit busy CTAs fit one per CU, as FP8 does, rather than only while the split ones do. Paired in-process A/B vs 922fde0 on MI308X (gfx942), BF16 D64 MTP geometries: - UT mtp4-g16-page128-cluster-d64: 1.38x - B1 ctx32K and B16 ctx1K: 1.12-1.13x - everything else 0.99-1.01x 41 cases: geomean 1.017x. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… per CU on gfx942 waves_per_eu=2 trades spills (FP8 D128 MTP) or nothing (streamed BF16 D256) for a second resident CTA per CU. That only pays when the host-bounded busy CTAs fill two per CU. Between one and two per CU it only costs spills: the long-row UT case mtp4-g16-page128-varied (6 rows up to 4100 tokens) ran at 0.94-0.97x vs main. The cache key gains the two-per-CU flag (before D256_PIPE). Paired in-process A/B vs 26f6401 on MI308X (gfx942): mtp4-g16-page128-varied 1.01-1.10x; grid MTP and BF16 D256 cases unchanged (0.99-1.01x). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Four BF16 D128 M-tiles per CTA hold 360 VGPRs: one wave per SIMD. With split-major launches (ece6342) a two-way split no longer re-reads K/V from HBM, since a task's two CTAs dispatch together and share it. Each CTA then holds two M-tiles and, on busy grids, takes the streamed two-M-tile schedule at two CTAs per CU. Busy BF16 D128 MTP rows with at least four M-tiles now split two-way. D192 (0.76-0.83x) and two-M-tile rows (0.85-0.96x) measured slower and keep their schedules. Paired in-process A/B vs df92172 on MI308X (gfx942): BF16 D128 MTP4/MTP8 page128 at batch >= 16: 1.03-1.10x; other BF16 MTP cases unchanged. 92 cases: geomean 1.020x. Max error against the fp32 reference is unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
… on gfx950 The previous commit gave gfx950 multi-query plans one task slot per CU by default. With 128-256 sequences that leaves one or two tasks per sequence, so uneven lengths can no longer be split: on MiniMax-M3 MTP4 verify with lengths uniform in 1K-8K, FP8 B256 ran 0.96x of main and BF16 B256 1.25x instead of 1.66x. The default budget is now four slots per sequence and KV head, clamped to one to two slots per CU. Batches up to 64 sequences keep one slot per CU (all 14 fixed-length MTP4 verify workloads plan exactly as before); larger batches return to two. Random-length MTP4 verify at B64-B256: 1.17x-1.61x of main (FP8 B256 1.32x). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…pa-gfx950-tile-opt)
Union of both sides in the pa_decode tile kernel. #7709's policies stay gfx942-only and #7700's paths gfx950-only: - traits: the cache key carries #7700's SPARSE_GRID and #7709's busy, two-per-CU and split-major flags ahead of D256_PIPE; prefetch_v keeps both arches' rules; P_BUFFERS gives each M-tile a slot for MTP4_FUSED, D256_PIPE, BF16_MTP_FUSED or STREAM_KV. - pipeline: tile() issues both next-page lookaheads (PAGE16_PAGES_FIRST, PAGES_EARLY). single_query_tile takes #7709's STREAM_KV K/V branch, else #7700's late-scale / V-after-QK path with PAGES_EARLY ahead of REUSE_KV_PAGES. - op_lds: early Q covers MTP4_FUSED, DECODE_Q_EARLY, GFX942_Q_EARLY and MXFP4; the padded-row clamp takes #7700's form. - launch: the split-major grid with CTA_THREADS blocks, the gfx950 kernarg preload, and waves_per_eu=2 for FP8_MTP4_LEAN or WAVES_PER_EU_2. Compiled binaries: gfx950 and gfx1250 match #7700+#7701; gfx942 matches #7709 except the padded-row kernels, whose clamp compares against CTA_ROWS - 1 instead of CTA_ROWS (same rows selected). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
…re restaging STREAM_KV tiles stage the next tile's page ids in LDS at the start of the tile, so the page-table wait does not drain the K stream. The first tile's ids are staged in the prologue and read back by every wave after its barrier, but nothing ordered that read before the first tile's restaging: a fast wave could overwrite the ids while a slow wave had yet to read them, and the slow wave then gathered the first tile's V from the second tile's pages. Later tiles are ordered by the tile barriers. Streamed tasks now take one barrier after the read. The race showed as run-to-run output differences of up to 1e-3 on FP8 per-tensor D256 MTP4 two-head busy grids (1-2 of 40 runs on gfx942; 8-16 of 20 on gfx950 with streaming enabled). After the fix, every streamed case repeats bit-exactly; the cost is 1.0005x geomean over 161 streamed gfx942 cases. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
The barrier follows the first tile's page-id read in the merged prologue, ahead of #7700's Q-operand selection. STREAM_KV is gfx942-only in this stack, so gfx950 and gfx1250 binaries are unchanged. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
gfx950 pa_decode has the same register problem #7709 solved on gfx942: a lane holding a whole next-tile K prefetch plus the tile's V gather spills once its K row reaches 1 KiB (BF16 D512+, FP8 D1024), and single or paired M-tiles with 384/256-byte rows drop to one wave per SIMD on busy grids. Let STREAM_KV select on gfx950 too, outside the fused BF16 MTP and head-256 pipes, which keep their own schedules; streamed tiles skip the dual-group CTA. The busy-grid flags it reads are now arch-neutral (renamed busy_grid / two_per_cu); every other consumer stays gfx942-only. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
#7709's gfx942 query split also pays on gfx950: when few CTAs hold work, per-CTA latency dominates, and splitting the MTP query positions over more CTAs (re-reading K/V from L2) shortens the critical path. Apply it on gfx950 with two differences found by measurement: - per-token FP8 keeps its fused MTP4 tiles (splitting them ran 0.75x); - gfx950 MTP plans fix capacity at one or two slots per CU regardless of load, so a capacity bound no longer means idle CUs. The FP8 split also requires at most eight KV tiles per CU (tile_demand, the uncapped sequence x head x tile count); B64/32K per-tensor ran 0.87x without it. The BF16 D128 split for busy grids stays gfx942-only (0.6-0.9x on gfx950). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
#7709 launches gfx942 query splits and multi-head grids split-major, so a slot's splits (or a small grid's heads) dispatch together instead of queueing behind every empty plan slot of the previous split. On gfx950 the order only pays while at most a quarter of the CUs hold work; on fuller grids slot-major order measured up to 1.25x faster, so gfx950 takes it only below that bound, for FP8 and for BF16 up to D128 (LDS-forced BF16 D192/D256 splits ran 0.91-0.93x). MXFP4 keeps slot-major. Two gfx950-only paths assumed the slot-major axes and now follow the grid: the BF16 D128 planned output's slot count (grid_dim.x; with a split-major grid it read heads x splits and stored partials at wrong offsets) and the one-partition plan's empty-row zero fill (block_id("y") as the KV head). Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: Reviews: |
fsx950223
marked this pull request as ready for review
November 12, 2026 07:44
NOT READY FOR HUMAN REVIEWPR Self-ReviewRevision Author self-review accepted. Readiness report: Author answers
Show the author's self-review answersBehavior and intentValidation and unverified scopePR-specific questionChanged op testsPR state
Copilot findings
Code and op-test hygiene
Required CI
Shadow mode reports what would block but keeps the status successful. Set |
This branch has not been deployed
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
Brings three of #7709's gfx942
pa_decodeschedule policies to gfx950, each re-tuned on MI355X:gfx942 and gfx1250 kernels are bit-for-bit unchanged.
On MI355X, over all 817 cases of the unit-test-geometry matrix described below: geomean 1.19x against the merged dependencies
03894b767e, i.e. #7700+#7701+#7709 without this PR's commits (unit-test cases 1.22x, grid cases 1.19x). The 333 cases whose kernels change run 1.54x faster, and none is slower than 0.98x.This PR depends on #7700, #7701 and #7709.
372b2bc843,50e1d38283,cbffdd201b.03894b767e...cbffdd201b.fsx950223/pa-stack-7700-7701-7709, not intomain.59d523033c, [FlyDSL] [Perf] pa_decode: gfx950 reducer issues all partition loads in one rung #770102dc5d11b8and [FlyDSL] [Perf] pa_decode: gfx942 schedule tuning (K/V streaming, latency hiding, grid-size policies) #770965eed8399b. Since then [FlyDSL] [Perf] pa_decode: gfx950 fused BF16 MTP4 verify tile, decode load scheduling and lean FP8 MTP4; four paths extended to gfx942 #7700 has moved tobdabd08dad: it extends four gfx950 paths to gfx942, adopts a named cache key, and keeps main's Q absmax above D256 on gfx942. [FlyDSL] [Perf] pa_decode: gfx950 reducer issues all partition loads in one rung #7701 has moved toa8b9f032b9. The stack will be rebuilt and re-measured once those heads settle.Motivation
#7709 found four schedule problems on gfx942. To see which also hold on gfx950, I opened each of #7709's gfx942-only gates on gfx950 and measured it on its own, against main and against #7700+#7701. Three hold:
Three gfx942 policies did not carry over and stay gfx942-only:
waves_per_eu=2on busy grids: [FlyDSL] [Perf] pa_decode: gfx950 fused BF16 MTP4 verify tile, decode load scheduling and lean FP8 MTP4; four paths extended to gfx942 #7700's gfx950 MTP plan holds capacity at one or two slots per CU, so it changes no case.Changes
All changes are in
aiter/ops/flydsl/kernels/pa_decode/{traits,op_epilog}.pyandpa_decode_kernel.py, one commit each:372b2bc843K/V streaming on gfx950 (STREAM_KV).gfx942_busy->busy_grid,gfx942_two_per_cu->two_per_cu. All their other consumers stay gfx942-only.50e1d38283MTP query split over idle gfx950 CUs, using [FlyDSL] [Perf] pa_decode: gfx942 schedule tuning (K/V streaming, latency hiding, grid-size policies) #7709's rules with two gfx950 differences:tile_demand, the sequence x head x tile count before the plan's capacity cap. gfx950 MTP plans fix capacity regardless of load, so a capacity bound no longer means idle CUs (B64/32K per-tensor ran 0.87x without the cap).cbffdd201bSplit-major launches on sparse gfx950 grids.grid_dim.x. With a split-major grid that is heads x splits, so partials were stored at wrong offsets.block_id("y")as the KV head.Not changed:
plan_pa_decodeand its budgets;Merged dependencies
The commits before the last three (up to
03894b767e) are built from:59d523033c;02dc5d11b8;8274dfbc42;65eed8399b.The #7709 merge resolves 19 conflict hunks in
traits.py,pipeline.py,op_lds.pyandpa_decode_kernel.pyas a union of both sides: #7709's policies stay gfx942-only and #7700's paths gfx950-only. A compile-only replay of every matrix case gives:While porting streaming, I found a first-tile LDS race in #7709's streamed path. It is fixed in #7709 itself (
65eed8399b), not here.Performance (if applicable)
How these numbers were measured:
03894b767e(the merged dependencies), after =cbffdd201b.op_tests/test_flydsl_pa_decode.pythat runs on gfx950 at its unit-test lengths (115 cases), plus each geometry at batch {1, 16, 64} x context {1K, 8K, 32K} with the default plan (702 cases).mtp4-g16-page128-direct...-direct-d64Per commit, each against its parent, over the cases it changes:
By family over all 817 cases:
Known regressions: no changed case is below 0.98x. Three unchanged cases read 0.974-0.976x; their binaries are identical, so those readings are noise.
Testing
Results:
op_tests/test_flydsl_pa_decode.pypasses in full atcbffdd201b: 355 passed, 37 skipped.-direct,-direct-d64and BF16-clustercases), so the existing tests cover them.No unit test is added: the existing cases reach every new gfx950 specialization. Not tested on MI250X (gfx90a is unsupported by this kernel), MI300X (304 CUs) or MI325X.
Documentation
The API, the plan and the set of supported shapes are unchanged; the change is internal to the gfx950 tile schedule.
Dependencies
Breaking Changes
None
PR Self-Review
Reviewed revision
cbffdd2
Behavior and intent
On gfx950 the pa_decode tile schedule now uses three policies #7709 introduced for gfx942:
Each is re-tuned on MI355X where the gfx942 rule lost: per-token FP8 keeps its fused MTP4 tiles, the FP8 split is bounded by the uncapped tile count, and split-major is limited to sparse grids and to FP8 or BF16 up to D128. Without these policies, gfx950 spilled at D512/D1024 and left CUs idle under MTP (1.19x geomean over the matrix). gfx942 and gfx1250 binaries are unchanged.
Validation and unverified scope
Evidence on MI355X:
pytest op_tests/test_flydsl_pa_decode.py: 355 passed, 37 skipped. The negative control with commit 3's axis fixes reverted fails 9 cases.det_check.py).Evidence on gfx942:
COMPILE_ONLY=1) of every matrix case: gfx942 and gfx1250 binaries identical to03894b767e.Unverified:
max_context_lengthget no tile bound, so their FP8 MTP rows do not split on gfx950.bdabd08dad/ [FlyDSL] [Perf] pa_decode: gfx950 reducer issues all partition loads in one rung #7701a8b9f032b9; every number above is for the dependency SHAs listed in the summary.PR-specific question
op_tests/test_flydsl_pa_decode.pyis the only changed op test in this diff, and its change comes from #7700 (0a4638a355). It is not one of this PR's three commits.This PR's own commits change no test. Their main boundary is the split-major grid on gfx950: it swaps which block-id axis names the plan slot. Any gfx950-only path that still assumes slot-major axes silently reads the wrong slot or head. is the split-major grid on gfx950: it swaps which block-id axis names the plan slot. Any gfx950-only path that still assumes slot-major axes silently reads the wrong slot or head.
Two such paths existed:
grid_dim.x);block_id("y")).Opening split-major on gfx950 without fixing them produced relative errors of 1-19 on BF16 D128 two-head shapes. Both now select the axis from
SPLIT_MAJOR, and their IR is unchanged when it is off. The remaining directblock_id/grid_dimreads are gfx1250-only pipelines, or the MXFP4 direct path, which keeps slot-major. Evidence: with the two fixes reverted, the existing UT fails 9 cases; with them, all pass; every changed case is deterministic and matches the base's fp32-reference error.The second boundary is the cache key. The renamed busy flags keep their positions ahead of
D256_PIPE, socreate()'s positional unpacking and the test'scache_key[-2:]are unchanged; the pipe-policy tests pass.Changed op tests
op_tests/test_flydsl_pa_decode.py: changed by #7700 (0a4638a355), not by this PR's three commits. Its plan-structure check now accepts gfx950's balanced min-max planner, asserting each sequence's task count equalsmin(ceil(tiles / T), max_partitions)for the smallest feasible boundT. It is not redundant, because no other test checks the plan's task layout.This PR's commits add no test, because the existing
test_pa_decodecases reach every new gfx950 specialization: thehead1024/head512/head192streamed rows, the MTP4/MTP8-direct,-merge,-probability_tailand window cases for the query split, and the-direct/-clustercases for split-major. The negative control shows they fail when the split-major axis fixes are missing.🤖 Generated with Claude Code