CSA compressor: extend the validated ratio envelope to ratio=2 - #710
Merged
Conversation
The generic compressor kernels (compressor_sm100.py) are JIT-specialized
per (ratio, head_dim, coff) and carry no ratio-dependent assumption below
the documented register bound at ratio ~ 32; the {4, 128} envelope in
api.py's check_support encoded validation scope, not a kernel limitation.
Widen the generic-family branch to ratio in {2, 4} -- ratio=2 (coff=2) is
the compressor configuration used in production training of the model
family this operation serves. No kernel changes; the ratio=128 family and
its routing are untouched.
Validation on B200 (CC 10.0, torch 2.13, cutlass-dsl 4.5.1):
- ratio=2 test cases added mirroring every ratio=4 pattern (numerics vs
the fp32-intermediate eager reference with the bitwise dKV/dScore
contract now asserted at ratio in {2, 4}, ragged/short/empty-segment
packs, static-capacity padding, NaN-canary zero-write ownership, replay
determinism, check_support acceptance/rejection): module suite 145
passed, 1 skipped (multi-GPU test on a single-GPU runner); the new
cases fail with the envelope ValueError on the unpatched gate.
- at ratio=2, out/dKV/dScore are bitwise-identical to the fp32 eager
reference in every tested config; dAPE differs at fp32 atomic-order
level only (~1e-6 relative), per the documented contract.
- nsys pure-kernel at ratio=2 coff=2 (50-iter mean): fwd 4.7/10.4/13.5/
39.6 us, bwd 11.4/25.6/27.7/69.2 us for d128/d512 x 1/3x8192-token
packs -- 18-21x the eager reference region.
Signed-off-by: zky <kaiyue.zhou@z.ai>
Contributor
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: 📒 Files selected for processing (3)
Included review availability: Your plan provides up to 12 included reviews per hour; 11 remain after this review. 📝 WalkthroughWalkthroughThe CSA compressor now supports generic ratios 2 and 4 with ChangesCSA ratio support
Estimated code review effort: 3 (Moderate) | ~20 minutes Merge Risk: ⚪ Minimal · up to This PR widens the supported CSA compressor ratio envelope to include ratio=2, with accompanying tests and documentation updates; no actionable merge-blocking risk remains after normal checks and review. Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
Collaborator
|
@cudnn-ci-bot run oss |
Anerudhan
approved these changes
Aug 24, 2026
|
🏁 Pipeline finished SHA: |
Anerudhan
added
orig-external
Anerudhan
approved these changes
Aug 24, 2026
Merged
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
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Widen
check_supportinpython/cudnn/csa/compressor/api.pyto acceptratio=2in thegeneric-kernel family: the validated envelope becomes
ratio in {2, 4, 128}(withcoff in {1, 2}at every ratio, as before).No kernel code changes. The generic kernels in
compressor_sm100.pyareJIT-specialized per
(ratio, head_dim, coff)and are ratio-generic; the{4, 128}gateencoded validation scope, not a kernel limitation. The dedicated ratio=128 family
(
compressor_sm100_r128.py) and the ratio-based routing are untouched.Why
ratio=2(withcoff=2) is the compressor configuration used in production training ofthe model family this operation serves (the kernels were ported from Megatron-LM in
#427). With the gate at
{4, 128}, integrating the module at ratio=2 requires bypassingthe validation layer and calling
compressor_sm100.run_fwd/run_bwddirectly; wideningthe envelope lets that configuration go through the supported, validated API.
Why ratio=2 needs no kernel work (verified by reading the kernels and by the test
additions below):
win = coff * ratiois a compile-time constant; register pressureshrinks with ratio (the family is documented register-bound only beyond
ratio ~ 32,and ratio=4 is already validated);
k % ratio), the overlap-window first-block invalidation, the backwardzero-write ownership classes (
seqlen % ratiotails,seqlen < ratiosegments, thelast block's first-half columns at
coff=2, token-capacity padding) and the dAPEatomic targets are all expressed in
ratiowith no lower bound;_fwd_scheduleis a function ofhead_dimonly; the backward uses fixedrows_per_cta/threads);check_supportare checked independently ontotal_tokens * coff * head_dim,total_comp * head_dimandratio * coff * head_dim, so the (relatively larger)total_compat ratio=2 isalready covered.
Numerics at ratio=2
B200 (CC 10.0), torch 2.13.0 / CUDA 13.3 / nvidia-cutlass-dsl 4.5.1, seeded inputs from
this suite's generator; fused wrappers vs the suite's fp32-intermediate eager reference
(the numerics-contract comparator):
dKV/dScorecarry the generic family's bitwise contract at ratio=2, exactly as atratio=4. On every tested ratio=2 config the forward is additionally fully bitwise (the
contract only requires "within one bf16 rounding step on a tiny fraction of elements";
the ratio=4 controls show that documented behavior).
dAPEdiffers at fp32atomic-order level only (~1e-6 relative), per the documented contract (
dAPEis notrun-to-run deterministic in either implementation; the backward still refuses strict
torch.use_deterministic_algorithms(True)).Performance at ratio=2
nsys pure-kernel times (sum of kernel durations inside a
cudaProfilerApi-gated50-iteration unsynced loop / 50; JIT + fast-launch warmup before the window), ratio=2
coff=2, single B200; eager = the fp32-intermediate reference region (fwd + autogradbwd, ~93 kernel launches per step):
(The fused ratio=2 kernels run at the same throughput class as the validated ratio=4
configs — same kernel bodies, smaller unrolled window.)
Tests
test/python/fe_api/csa/test_CSA_compressor.pyadds ratio=2 cases following the exactratio=4 patterns, both window forms (
coff1 and 2):oracle (12 new
_SHAPEScases: single/ragged/short-segment/empty-segment packs, d128/512 and the odd
d=65scalar layout), with the bitwisedKV/dScoreassertionnow applied at
ratio in {2, 4};the tiny-segment packs use 1-token segments so the
seqlen < ratioclass is stillexercised at ratio=2);
check_supportacceptance at ratio=2 and rejection cases (coff0/3 at ratio=2, atotal_tokens < ratiocase at ratio=2, and the updated{2, 4, 128}rejectionmessage).
Module suite (
-m "L0 or L1", single B200): 145 passed, 1 skipped(
test_multi_device_launch, single-GPU runner). On the unpatched gate the new ratio=2cases fail with the expected envelope rejection
(
ValueError: CSA compressor is validated for ratio in {4, 128} only, got ratio=2, ...),so the tests pin the widening in both directions.
Docs:
docs/fe-oss-apis/csa.mdsupport surface and numerics-contract sections updated(
ratio in {2, 4}for the generic family; ratio=128 family unchanged).Summary by CodeRabbit
New Features
Bug Fixes
Documentation