You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
{{ message }}
Repository navigation
perf(cake_gdn_det_kcp): chain-aware walk widths in the deterministic KCP GDN package - #8538
Cake's deterministic context-parallel GDN training kernel has one data-type variant: BF16 q/k/v/out/dout with FP32 scalar log-decay g and beta. The grouped Q/K head counts (H_qk, HV) are runtime parameters with production shape (16, 48). The kernel has three partition layouts: contiguous (C), whole-buffer zigzag (Z) and per-document zigzag + ragged (D). This PR delivers the same variant. It changes the host-side rule that picks the value-head tile width for the cross-rank state-summary walk launches. A walk runs either 128 value columns per CTA or 64, and the 64 form runs two CTAs per segment and head. Forward walks still use the 64 form when the grid fits one wave. They now also use it when an LPT schedule of the walk's segment chain lengths shows that the two-CTA grid finishes within its longest chain. Backward offset walks always use the 64 form. Backward tape walks use it on sm_100a and sm_103a, while sm_107a keeps the previous two-wave rule. One exception applies on sm_100a: a backward walk keeps the previous two-wave rule when its longest chain is at least 400 chunks and its two-CTA schedule would take longer than the one-CTA schedule. The rule reads only host plan integers (segment lengths, head count, SM count and device architecture), never tensor contents. It takes the same decision in the forward, the checkpoint recompute and the recompute-handle replay. The non-deterministic KCP package (PR #8521) uses the same thresholds. The change is host-side Python only and the kernels are unchanged. Only launch widths differ, so outputs and gradients are bitwise identical to before on every measured row, and peak memory is unchanged. A is the previous Cake version of this package, as shipped after PR #8385. B is this PR on the same base. B/A < 1 means faster. Each GPU was measured on 66 rows: packs P1-P5 x CP 2/4/8 x layouts C/Z/D at B1 and B3. On 33 of these rows the launch form changes, and the other 33 are unchanged controls. On GB300 (sm_103a), changed rows range from 0.984 to 1.004 at B3 (median 0.989) and from 0.976 to 1.001 at B1 (median 0.999), though only 6 B1 rows change. Unchanged GB300 rows range from 0.994 to 1.013 (median 1.000). On GB200 (sm_100a), changed rows range from 0.987 to 1.016 at B3 (median 0.992) and from 0.978 to 1.002 at B1 (median 1.000). Unchanged GB200 rows range from 0.992 to 1.003 (median 1.000). Taking the median over both sessions, changed rows at CP 2 / 4 / 8 come to 0.989 / 0.990 / 0.999 on GB300 and 0.991 / 0.994 / 1.000 on GB200. That is about a 1 % gain at CP 2 and CP 4 and about parity at CP 8. The largest gain is CP 8 P3 C B1 on GB300 at 0.977 / 0.976. The slowest changed row is CP 8 P3 C B3 on GB200, at 1.016 in session 1 and 1.006 in session 2. R200 (sm_107a) rows were not measured.
ms = slowest-rank median; B1 = module forward with autograd recording; B3 = forward + non-reentrant checkpoint replay + backward, including all communication and copies; Walk forms = changed (this PR launches a different width) or unchanged (same launches as A); s1/s2 = two sessions.
Cake also ships these separately: a non-deterministic variant of this kernel, single-GPU GDN training packages without context parallelism (deterministic and non-deterministic), GDN inference kernels for chunked prefill and decode, and the KDA deterministic context-parallel training sibling, which covers the per-channel-gate form; this kernel ships on sm_100a, sm_103a and sm_107a.
Summary by CodeRabbit
Performance Improvements
Improved walk-width selection for forward and backward operations using launch-chain length, GPU architecture, and execution scheduling characteristics.
Added architecture-specific handling for long walks to help select an appropriate execution strategy.
Testing
Expanded coverage for width selection across architectures, walk types, and execution paths.
…KCP GDN package
The forward state walk (#2) also takes the two-CTA vh = 64 form when its
grid finishes within the longest walked segment (LPT of 2 * HV CTAs per
segment over two CTAs per SM). The adjoint offset walk (#7) always takes
vh = 64, and the adjoint tape walk takes vh = 64 on sm_100a and sm_103a;
sm_107a keeps the two-wave rule. Both forms of each program write
bitwise-identical results. The decision reads host plan integers only
(walk list, segment chunk table, HV, SM count, device architecture), so
the forward, the checkpoint recompute and the recompute-handle replay
take the same form. This is the rule of the non-deterministic KCP
package.
On sm_100a, a #7 walk (offset and tape) whose longest segment has at
least 400 chunks and whose two-CTA grid has a larger LPT makespan than
its one-CTA grid (both over 2 x SM slots) takes the two-wave rule again
(vh = 64 iff 2 * n_walk * HV <= 4 * SM). The decision reads host plan
integers only (walk_dhu_long_chain, memoized per walk list in _Tables)
and is passed with each #7 launch. Other architectures keep their
rules. This is the clause of the non-deterministic KCP package.
No actionable comments were generated in the recent review.
Walkthrough
KCP walk-width selection now considers LPT chain-fit and long-chain decisions in addition to walk count and SM count. DHU selection also uses walk type and architecture. The plan derives and caches these decisions, then passes them into launch resolution.
Changes: KCP walk-width selection
Layer / File(s)
Summary
Width rules and conditions flashinfer/experimental/cake_gdn_det_train_kcp/cake_kcp.py, flashinfer/experimental/cake_gdn_det_train_kcp/README.md, flashinfer/experimental/cake_gdn_det_train_kcp/cake_options.py, tests/experimental/test_cake_gdn_det_train_kcp_api.py
Forward width 64 is selected when the one-wave condition or the LPT chain-fit condition holds. DHU width selection uses architecture, walk type, and the SM100 long-chain condition. Tests cover the selection boundaries and pinned vectors.
Plan metadata and launch selection flashinfer/experimental/cake_gdn_det_train_kcp/cake_kcp.py, tests/experimental/test_cake_gdn_det_train_kcp_api.py
The plan memoizes chain-fit and long-chain decisions and passes them, with architecture data, to launch resolution. Tests check launch records and grids across forward, checkpoint recompute, and recompute-handle paths.
sequenceDiagram
participant CallDetKcp
participant Tables
participant LaunchResolver
participant WidthRules
CallDetKcp->>Tables: query chain-fit or long-chain decision
Tables-->>CallDetKcp: return walk decision
CallDetKcp->>LaunchResolver: pass walk decision and launch data
LaunchResolver->>WidthRules: select fwdh_vh or dhu_vh width
Loading
Possibly related issues
flashinfer-ai/flashinfer#7818 — The change applies width selection across recompute and replay launch paths, but does not remove communication.
flashinfer-ai/flashinfer#7251 — The change directly updates the host-side vh64 offset and tape walk selection rules.
Suggested labels:op: linear attention, run-ci
Suggested reviewers:yzh119
Merge Risk: Minimal · up to 77aa4
The chain-aware launch-width changes have no demonstrated actionable defect in the supplied evidence and are mergeable subject to normal checks. Performance on unmeasured shapes remains unproven, not a demonstrated regression.
Pre-merge checks: 3 · 2
Failed checks: 2 warnings
Check name
Status
Explanation
Resolution
Description check
The description gives a detailed and relevant summary, implementation rationale, benchmark data, and test scope. However, it does not follow the required template: it omits the Related Issues, Pull Re…
Add the missing template sections. Complete the pre-commit and test checkboxes, provide the experimental tracking issue and required experimental metadata, and add the tests under the experimental-tests fence.
Docstring Coverage
Docstring coverage is 73.91% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documenta…
Write docstrings for the functions missing them to satisfy the coverage threshold.
Passed checks: 3 passed
Check name
Status
Explanation
Title check
The title clearly identifies the performance change and the affected deterministic KCP GDN package.
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: Description check
Explanation
The description gives a detailed and relevant summary, implementation rationale, benchmark data, and test scope. However, it does not follow the required template: it omits the Related Issues, Pull Request Checklist, Pre-commit Checks, Tests checklist, and required Experimental Track section with a tracking issue and experimental-tests targets.
Resolution
Add the missing template sections. Complete the pre-commit and test checkboxes, provide the experimental tracking issue and required experimental metadata, and add the tests under the experimental-tests fence.
Full details: Docstring Coverage
Explanation
Docstring coverage is 73.91% which is insufficient. The required threshold is 80.00%. Docstring coverage is scoped to functions touched by this diff with cyclomatic complexity of at least 2. Documentation-only edits and deletions within surviving functions are included; unchanged diff context and fully deleted functions are excluded. Analyzed 23 functions across 3 files. (1 skipped: 1 unsupported.)
Resolution
Write docstrings for the functions missing them to satisfy the coverage threshold.
Finishing Touches
Generate unit tests (beta)
Create a new PR
Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.
Included review availability: This review used your included allowance. Your plan provides up to 8 included reviews per hour; 4 remain after this review.
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
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.
flashinfer-bot
added
the
op: linear attention
Cake's deterministic context-parallel GDN training kernel has one data-type variant: BF16 q/k/v/out/dout with FP32 scalar log-decay g and beta. The grouped Q/K head counts
(H_qk, HV)are runtime parameters with production shape (16, 48). The kernel has three partition layouts: contiguous (C), whole-buffer zigzag (Z) and per-document zigzag + ragged (D). This PR delivers the same variant. It changes the host-side rule that picks the value-head tile width for the cross-rank state-summary walk launches. A walk runs either 128 value columns per CTA or 64, and the 64 form runs two CTAs per segment and head. Forward walks still use the 64 form when the grid fits one wave. They now also use it when an LPT schedule of the walk's segment chain lengths shows that the two-CTA grid finishes within its longest chain. Backward offset walks always use the 64 form. Backward tape walks use it on sm_100a and sm_103a, while sm_107a keeps the previous two-wave rule. One exception applies on sm_100a: a backward walk keeps the previous two-wave rule when its longest chain is at least 400 chunks and its two-CTA schedule would take longer than the one-CTA schedule. The rule reads only host plan integers (segment lengths, head count, SM count and device architecture), never tensor contents. It takes the same decision in the forward, the checkpoint recompute and the recompute-handle replay. The non-deterministic KCP package (PR #8521) uses the same thresholds. The change is host-side Python only and the kernels are unchanged. Only launch widths differ, so outputs and gradients are bitwise identical to before on every measured row, and peak memory is unchanged. A is the previous Cake version of this package, as shipped after PR #8385. B is this PR on the same base. B/A < 1 means faster. Each GPU was measured on 66 rows: packs P1-P5 x CP 2/4/8 x layouts C/Z/D at B1 and B3. On 33 of these rows the launch form changes, and the other 33 are unchanged controls. On GB300 (sm_103a), changed rows range from 0.984 to 1.004 at B3 (median 0.989) and from 0.976 to 1.001 at B1 (median 0.999), though only 6 B1 rows change. Unchanged GB300 rows range from 0.994 to 1.013 (median 1.000). On GB200 (sm_100a), changed rows range from 0.987 to 1.016 at B3 (median 0.992) and from 0.978 to 1.002 at B1 (median 1.000). Unchanged GB200 rows range from 0.992 to 1.003 (median 1.000). Taking the median over both sessions, changed rows at CP 2 / 4 / 8 come to 0.989 / 0.990 / 0.999 on GB300 and 0.991 / 0.994 / 1.000 on GB200. That is about a 1 % gain at CP 2 and CP 4 and about parity at CP 8. The largest gain is CP 8 P3 C B1 on GB300 at 0.977 / 0.976. The slowest changed row is CP 8 P3 C B3 on GB200, at 1.016 in session 1 and 1.006 in session 2. R200 (sm_107a) rows were not measured.ms = slowest-rank median; B1 = module forward with autograd recording; B3 = forward + non-reentrant checkpoint replay + backward, including all communication and copies; Walk forms = changed (this PR launches a different width) or unchanged (same launches as A); s1/s2 = two sessions.
Packs: P1 241,817 tokens / 6 documents, P2 261,987 / 3, P3 259,174 / 2, P4 254,976 / 1 (one document spans the whole group), P5 249,856 / 4. Heads (16, 48); causal conv1d width 4 over 10,240 channels.
Software: CUDA 13.4 toolkit, torch 2.14.1+cu130.
NVIDIA GB300 (sm_103a)
NVIDIA GB200 (sm_100a)
Cake also ships these separately: a non-deterministic variant of this kernel, single-GPU GDN training packages without context parallelism (deterministic and non-deterministic), GDN inference kernels for chunked prefill and decode, and the KDA deterministic context-parallel training sibling, which covers the per-channel-gate form; this kernel ships on sm_100a, sm_103a and sm_107a.
Summary by CodeRabbit