Repository navigation
frost(gated_attention_block): the block-sparse decode split rule removes the wave tail -- where the fill rule leaves a second wave the split minimizes waves x (tiles per chunk + 1); every one-wave schedule keeps its measured split - #1802
Conversation
…ves the wave tail -- where the fill rule leaves a second wave (B x H_kv x S_q x split > SMs) the split minimizes waves x (tiles per chunk + 1) (T 128: 1 -> 3; B 32 x S_q 1: 4 -> 3; B 4 x S_q 4: 7 -> 6); every one-wave schedule keeps its measured split The fill rule (ceil(SMs / items), capped at the item's tile count and at 16) puts one chunk-item per SM but does not see the tail wave it leaves where the chunk-items overflow one wave of SMs: a CTA streams its chunk serially, so 256 equal items on 212 SMs cost two item durations, not 1.2. The block's end-to-end census on a Rubin cc 10.7 part (212 SMs at the 2376 MHz lock, CUDA-graph replay, the 24 / 2 geometry at top_k 512) shows the three schedules it leaves in two waves: the T 128 step's 256 unsplit 17-tile items (64.7 / 69.0 / 70.9 us at KV 8K / 32K / 128K), B 32 x S_q 1 at split 4 (256 chunk-items in 1.2 waves) and B 4 x S_q 4 at split 7 (224 chunk-items in 1.06 waves). sparse_decode_split_choice gains the wave rule: where the fill rule's pick overflows one wave (units x split > SMs) the split in [1, min(ceil((top_k + 1) / 32), SPARSE_DECODE_SPLIT_CAP)] that minimizes ceil(units x split / SMs) x (ceil(tiles / split) + SPARSE_DECODE_SPLIT_ITEM_COST_TILES) (sparse_decode_split_rounds: the waves the launch takes, times the tiles a chunk streams plus the chunk-item's fixed cost -- one tile, a MODEL constant: the bare waves x tiles count would cut the 17-tile T 128 item into nine two-tile chunks), ties to the smaller split. T 128 1 -> 3 (768 chunk-items of 6 / 6 / 5 tiles), B 32 x S_q 1 4 -> 3 (192, one wave), B 4 x S_q 4 7 -> 6 (192, one wave); every launch the fill rule keeps in one wave -- every B <= 4 x S_q 1 cell and B 1 x S_q 4, the measured 16s -- is untouched, bitwise. The combine form a moved cell takes is the pinned one (the lean group walk at B 32 x 24 heads at split 3 as at 4; the combine module's own forms at B 4 x S_q 4 at split 6 as at 7). A performance choice only: the same function under every value; decode_split_kv= pins any split in [1, the item's tile count]. Tests (RED-first on the pre-change tree: the constant's ImportError and `4 == 3` on the B 32 x S_q 1 class): the rule table (one-wave rows unchanged, the second-wave rows, the undo row, the one-wave invariant, the bound) and the 19 measured decode cells at 212 and at 204 SMs (the ten one-wave cells bitwise the former rule, the six two-wave cells and the three T 128 cells moved, the combine form per moved grid). Docs: the decode-mode paragraph names the wave rule; the published requires row too. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
RomanAnders90
added
cat-feature
|
Re-verification of the rebased head (
No kernel timeout on any module. The rule table rows (the one-wave invariant, the second-wave rows, the undo row, the bound) and the 19 measured decode cells at 212 and at 204 SMs are inside the first module's count. |
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info
📝 Walkthrough
Merge Risk: ⚪ Minimal · up to
|
| Check name | Status | Explanation |
|---|---|---|
| Title check |  |
The title clearly describes the main change: a wave-aware default block-sparse decode split rule that preserves one-wave choices. It is longer than preferred but remains specific and related to the ch… |
| Description check | |
The description provides detailed summary, rationale, API impact, testing results, benchmark data, limitations, and open items. It omits several template checklist sections, but the core required chan… |
| Docstring Coverage | |
Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 files. (2 skipped: 1 … |
| Linked Issues check | |
Check skipped because no linked issues were found for this pull request. |
| Out of Scope Changes check | |
Check skipped because no linked issues were found for this pull request. |
Full details: Title check
-
Explanation
The title clearly describes the main change: a wave-aware default block-sparse decode split rule that preserves one-wave choices. It is longer than preferred but remains specific and related to the changeset.
Full details: Description check
-
Explanation
The description provides detailed summary, rationale, API impact, testing results, benchmark data, limitations, and open items. It omits several template checklist sections, but the core required change and validation information is substantially complete.
Full details: Docstring Coverage
-
Explanation
Docstring coverage is 100.00% which is sufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff. Analyzed 3 functions across 1 files. (2 skipped: 1 unsupported, 1 too large.)
✨ Finishing Touches: 💡 1
-
🛠️ Fix failing CI checks: 💡
-
- Commit to this branch
- Create a new PR
🧪 Generate unit tests (beta)
-
- Create a new PR
-
- Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Comment @coderabbitai help to get the list of available commands.
|
@cudnn-ci-bot run python/*/cutedsl |
|
🏁 Pipeline finished SHA: 23 passed, 1 failed, 2 skipped, 11 manual
|
Anerudhan
left a comment
Anerudhan
left a comment
There was a problem hiding this comment.
Codex bot review · model: gpt-6-astra
Reviewed 6187508ea6a7 against base 9c9db1ef311b.
No actionable finding. The new default stays within the existing split domain, preserves one-wave choices and explicit overrides, and is computed once during declaration. The sparse adapter sizes its partial-output workspace and selects its combine path from the resulting split. This is a bounded scheduling change; performance outside the measured device and workload remains unqualified.
Validation: Exact-head host/declaration suite on SM100 with DSL 4.7.1: 22 passed, 56 native SM107 cases deselected. Checked one-wave preservation, multi-wave choices, the split-undo case, bounds, explicit pins, gate placement and typed input rejection. Traced the selected split through block declaration, sparse-adapter construction, workspace sizing and combine selection. The existing kernel partitions each item's tiles into contiguous chunks with the maximum chunk length used by this model. The new search examines at most 16 choices and adds no execute-time work. No kernel body, barrier or public signature changes. Read the author's current-head native validation: 78 decode cases, 140 serving-decode cases and 23 capture cases passed. These are author results, not independently repeated native runs in this review. Reused pipeline 76026654. All three changed files match its integration tree. The SM107 wildcard target collects the affected module and is still running; SM80/SM90/SM100 lanes have passed, but their native Rubin cases skip. Current develop differs from the PR base only in separate SDPA D128/prepared-path changes. The sparse decoder and gated-block contracts consumed here are unchanged. Precheck found no candidates. Added the existing repository project and preserved the appropriate labels.
Limitations: Both available SM107 test GPUs were occupied, so I did not independently rerun native correctness, capture or timing. Host checks do not validate native execution. CI follow-up remains active. The fixed per-item cost is a scheduling estimate, not a guarantee for other SM counts, shorter live lists or unseen shapes. The explicit split override remains available; the existing measurements do not establish a general end-to-end speedup.
Approved: no P0 or high-risk P1 found. Remaining findings stay with the owner; merge timing stays with the owner.
|
Pipeline 76026654 finished with SM80, SM90 and SM100 passing and SM107 failing. All three changed files match the integration tree. SM107 JUnit confirms all 78 cases in the changed The SM107 lane reports five failures: an illegal-address error and worker loss in Codex bot review — model |
On
develop(33bee40: #1789's in-block indexer decode step and #1792's DSL 4.7.0 fixes merged): ONE commit overdevelop, theblock-sparse DECODE mode's default split rule. No binding / tracker / Capabilities change (a knob-domain default only, Rule S2 exempt:
decode_split_kvstays the one served knob, any split in[1, ceil((top_k + 1) / 32)]pins the former choice).Summary
The block-sparse decode mode (
kv_source="paged"on aQsaSpecgeometry) chooses its split-KV when the declaration leavesdecode_split_kv=None. The former FILL rule --ceil(SMs / (B x H_kv x S_q))chunks per work item, capped at the item's tile count andat 16 -- puts one chunk-item per SM but does not see the tail wave it leaves where the chunk-items overflow one wave of SMs: a CTA streams
its chunk serially, so 256 equal items on 212 SMs cost two item durations, not 1.2. The block's end-to-end census on a Rubin cc 10.7 part
(212 SMs at the 2376 MHz lock, CUDA-graph replay, the 24 / 2 geometry at
top_k512) showed three schedules left in two waves: theT 128 step's (B 32 x S_q 4) 256 unsplit 17-tile items, B 32 x S_q 1 at split 4 (256 chunk-items in 1.2 waves) and B 4 x S_q 4 at split 7
(224 chunk-items in 1.06 waves).
sparse_decode_split_choicegains the WAVE rule: where the fill rule's pick overflows one wave (units x split > SMs) the split in[1, min(ceil((top_k + 1) / 32), SPARSE_DECODE_SPLIT_CAP)]that minimizesceil(units x split / SMs) x (ceil(tiles / split) + SPARSE_DECODE_SPLIT_ITEM_COST_TILES)-- the waves the launch takes, times the tiles a chunk streams plus the chunk-item's fixed cost(
sparse_decode_split_rounds; the cost constant is ONE tile, a model term: the barewaves x tilescount would cut the 17-tile T 128 iteminto nine two-tile chunks) -- ties to the smaller split. On a 212-SM part: T 128 1 -> 3 (768 chunk-items of 6 / 6 / 5 tiles), B 32 x
S_q 1 4 -> 3 (192 chunk-items, one wave), B 4 x S_q 4 7 -> 6 (192, one wave). Every launch the fill rule keeps in one wave -- every
B <= 4 x S_q 1 cell and B 1 x S_q 4, the measured 16s -- is untouched, bitwise. The combine form a moved cell takes is the pinned one
(the lean group walk at B 32 x 24 heads at split 3 as at 4; the combine module's own forms at B 4 x S_q 4 at split 6 as at 7). A
performance choice only: the same function under every value.
Tests
test_block_qsa_decode.py(RED-first on the pre-change tree: the constant'sImportErrorand4 == 3on the B 32 x S_q 1 class): therule table (one-wave rows unchanged, the second-wave rows, the undo row, the one-wave invariant, the bound) and the 19 measured decode cells
at 212 and at 204 SMs (the ten one-wave cells bitwise the former rule, the six two-wave cells and the three T 128 cells moved, the combine
form per moved grid). Docs: the decode-mode paragraph and the published requires row name the wave rule.
Re-verified after the rebase onto
developon a SHARED Rubin cc 10.7 development part (204 SMs; functional only, never a performancenumber):
test_block_qsa_decode.pyat every tier,test_block_serving_decode.pyandtest_block_serving_capture.pyat L0 -- the countsin the first comment below.
Measured (on the same single commit before the rebase onto
develop; the rebased patch is byte-identical): the block-sparse decode split rule (the wave rule) vs the former fill-rule split, pinned in the same process, with the vLLM QSA attention sub-layer as the baselineSetup. One exclusive GPU of a Rubin cc 10.7 performance node, 212 SMs at the 2376 MHz lock (busy clock sampled per cell, power-cap counter delta 0 on every cell, Xid count unchanged, no foreign processes); HBM SOL % against that GPU's copy bandwidth measured in place (13180.6 GB/s), never the datasheet. Tree
43d11dcbfonb4664a2d4(the e2e-bind stack base + this one commit), the extension built from4f175f3fe, import attested per process. Protocol: ONE process per cell with every arm inside it, A/B/A x 3 rounds x 100 iterations, slots shuffled per iteration, CUDA events, per-round medians; the CUDA-graph REPLAY regime is the quoted column.vllm_qsa= the whole vLLM layer per step (SM100-family code untuned for cc 10.7, its own CUDA top-k library);frost_e2e= the whole FROST block in ONE execute under the NEW rule (decode_split_kv=None);FROST at the former split= the same block declared at the former fill-rule split in the same process (the lever column); page 16, bf16 pools and a bf16 compressed-key cache, uniformkv_lens. Speed-ups are POSITIVEbase / new - 1.Geomeans (POSITIVE = faster): the lever (the former split / the new rule - 1) over the 7 moved cells +10.2 % -- the three T 128 cells +17.2 %, B 32 x S_q 1 +4.7 %, B 4 x S_q 4 +6.7 %; vs vLLM over the 11 cells +10.9 %, the three T 128 cells +12.9 % (the model's band under this rule: parity to about +8 %; the rule is not the T 128 fix), B 32 x S_q 1 +17.9 %, B <= 4 +5.7 %, the quick-gate seven +9.7 %.
Bar: no cell loses beyond its control band (|Delta| <= max(1 %, the pair)); the moved cells read in-process against the pinned former split; the unmoved cells bitwise. Cells losing beyond their band: 0; the unmoved cells bitwise (out and list): True; the selection bitwise between the two splits on every cell: True; every control pair within 1 %: True. Verdict: PASS (every cell within its band; the unmoved cells bitwise; every arm admitted).
Oracles (never widened; the admission of every timed FROST arm): the selection == the fp64 top-k of the kernel's OWN operands up to ties, cos(out) > 0.999 vs the fp32 reference fed THAT list, the gated O within 2e-2, two launches bitwise, the captured closure bitwise incl. the list; the pinned block's list bitwise the new rule's on every cell and its out bitwise on the unmoved cells; the caller block fed FROST's list == the e2e block bitwise:
Per-stage census of ONE replay, B 32 x 8K x S_q 4,
frost_e2eat split 3 (7 kernels; device us summed 80.8; first start to last end 79.9 us; the slot's replay median 82.7 us):cudnn_kernel_frost_sm100_matmul_128x128x128_128x128x32_cluster1x1_1ctamma_swapAB__i64div32__TensorMacudnn_kernel_frost_qsa_prepare_tensorptrbf16gmemalign16o242562561_tensorptrbf16gmemalign16o22562561_cudnn_kernel_frost_qsa_score_paged_tensorptrbf16gmemalign16o4128i64i641_tensorptrbf16gmemalign16o412cudnn_kernel_frost_qsa_topk_decode_tensorptri32gmemalign16o1_tensorptri32gmemo1_tensorptri32gmemaligcudnn_kernel__kernel_TensorMapswizzleTensorMapSwizzles128b_3_dtypeclass_cutlassbase_dsltypingBFloat1cudnn_kernel__combine_lean_walk_kernel_tensorptrf32gmemoi64i64i64i641_tensorptrf32gmemoi64i64i641_tecudnn_kernel_frost_sm100_matmul_64x32x128_64x32x32_cluster1x1_1ctamma__i64div16__TensorMapswizzleTenPer-stage census of ONE replay, B 32 x 8K x S_q 4,
frost_e2e_pinnedat split 1 (6 kernels; device us summed 97.1; first start to last end 96.5 us; the slot's replay median 99.2 us):cudnn_kernel_frost_sm100_matmul_128x128x128_128x128x32_cluster1x1_1ctamma_swapAB__i64div32__TensorMacudnn_kernel_frost_qsa_prepare_tensorptrbf16gmemalign16o242562561_tensorptrbf16gmemalign16o22562561_cudnn_kernel_frost_qsa_score_paged_tensorptrbf16gmemalign16o4128i64i641_tensorptrbf16gmemalign16o412cudnn_kernel_frost_qsa_topk_decode_tensorptri32gmemalign16o1_tensorptri32gmemo1_tensorptri32gmemaligcudnn_kernel__kernel_TensorMapswizzleTensorMapSwizzles128b_3_dtypeclass_cutlassbase_dsltypingBFloat1cudnn_kernel_frost_sm100_matmul_64x32x128_64x32x32_cluster1x1_1ctamma__i64div16__TensorMapswizzleTenVerdict: PASS. Every arm admitted by its own oracle (never widened); every control pair within 1 % (the widest -0.79 %); no cell loses beyond its band -- every moved cell gains: +6.7 % at B 4 x 32K x S_q 4 (split 7 -> 6), +5.7 / +4.9 / +3.6 % at B 32 x S_q 1 (4 -> 3; 8K / 32K / 128K), +20.0 / +16.8 / +14.8 % at the three T 128 cells (1 -> 3); geomean +10.2 % over the 7 moved cells. The four unmoved cells (split 16 -> 16) are bitwise identical between the two declarations (out and selection).
Finding (reported, not a bar). The T 128 cells were modelled at parity to about +8 % vs vLLM under this rule and read +15.0 / +11.9 / +11.9 %. The census above says which kernel moved: at B 32 x 8K x S_q 4 the sparse tile kernel 64.8 us at split 1 -> 44.7 us at split 3, plus a 3.9-us lean-walk combine that split 1 did not need (-16.2 us net on an 82.7-us step); the GEMMs, prepare, scorer and top-k are identical within 0.2 us. The schedule lever (two waves of 17-tile items -> four waves at split 3) is worth more than the model's ceiling; this rule is still not the T 128 fix (the per-row re-gather of a sequence's listed KV is a separate change), so the T 128 numbers here are this rule's share alone.
Limitations. (1) One GPU of one node class; a second node class gets its own table. (2) On the moved cells the output differs from the former split by the combine's fp32 summation order only (cos 0.999996-0.999997, max abs 9.8e-4 = one bf16 ulp at the output's magnitude; the selection bitwise); the unmoved cells are bitwise. (3) The rule also moves classes outside these 11 cells (B 8 x S_q 1 / B 2 x S_q 4: 14 -> 9; B 16 x S_q 1: 7 -> 6; B 8 x S_q 4: 4 -> 3), unmeasured here, inside the knob's contract (the same function). (4) PDL on the sparse chain is not in the measured stack; bf16 pools only. (5) The vLLM arm is admitted by its own oracle on every cell; it is not bitwise across two launches on 10 of 11 cells (its own property), the FROST arms are.
Open items
SPARSE_DECODE_SPLIT_ITEM_COST_TILES = 1is a model term; the measurement above agrees with its direction on everymoved cell, so it is NOT refit here. The T 128 cells read above the model's band (parity to about +8 % vs the vLLM layer): the census
attributes the gain to the sparse tile kernel alone (64.8 -> 44.7 us at split 3, plus a 3.9-us combine); this rule is still not the
T 128 fix (the per-row re-gather of a sequence's listed KV is a separate change).
inside the knob's contract (the same function); a second node class (a different SM count) gets its own table.
🤖 Generated with Claude Code
Summary by CodeRabbit
Improvements
Documentation