CSA: add fused Compressor forward+backward CuTe-DSL kernels (ported from Megatron-LM)#427
Open
zkyue wants to merge 4 commits into
Open
CSA: add fused Compressor forward+backward CuTe-DSL kernels (ported from Megatron-LM)#427zkyue wants to merge 4 commits into
zkyue wants to merge 4 commits into
Conversation
Port the fused CSA/HCA Compressor gated-pooling kernels from Megatron-LM (NVIDIA/Megatron-LM#5984) into the FE-OSS Python API, per the maintainer request in NVIDIA/Megatron-LM#5984 (comment). One forward and one backward CuTe-DSL kernel fuse the THD gated-softmax pooling region (gather -> +APE -> overlap-window transform -> fp32 softmax -> gated weighted sum -> bf16 cast). Kernel math is unchanged from the Megatron original (verified bitwise on forward/dKV/dScore); the interface is reshaped to the APIBase pattern: CSACompressorForward / CSACompressorBackward classes, csa_compressor_forward_wrapper / csa_compressor_backward_wrapper, a new python/cudnn/csa package, docs, and fe_api tests (numerics vs fp32/upstream eager references and an fp64 oracle, ragged packs, static-capacity padding, run-to-run determinism, CUDA-graph capture/replay, check_support boundaries). Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
…opy) Widen every bf16 load/store in the fused Compressor forward kernel from scalar (16-bit) to paired 32-bit accesses: each thread now owns 2 adjacent head dims and moves its window slices through 32-bit universal copies (cute.autovec_copy over register fragments; cute.assume makes the alignment provable). The per-lane fp32 math is unchanged and in the same order, so the forward output stays bitwise identical to the previous kernel (re-verified against the Megatron-LM original on the same 5 THD shape-mix gate as the original port). Odd head_dims keep a scalar (vec == 1) instantiation of the same kernel; the head_dim 128/512 production shapes take vec == 2. The launch schedule moves from (rows_per_cta=4, threads=128) to one row per CTA with 64-thread column groups, which widens the sub-wave grids that limit the small shapes. Measured (nsys pure kernel time, B200, ratio=4 coff=2, THD packs of 8192-token sequences, forward kernel only): shape before after speedup 1x8192 d128 6.0 us 4.5 us 1.36x 3x8192 d128 12.6 us 10.0 us 1.26x 1x8192 d512 16.0 us 12.4 us 1.29x 3x8192 d512 42.1 us 35.0 us 1.20x Alternatives measured and rejected on the same matrix (all bitwise-equal): a PTX-unpack vec2 prototype (slower than the pure-DSL version: 4.7/10.5/12.9/35.7 us), and 64/128/256-bit per-thread vectors, which cut executed instructions but lose to register pressure and occupancy (80/147/255 regs vs 48; the 256-bit variant spills). Registers stay spill-free at 48 (previous kernel: 40). Tests: numerics shape matrix extended with an odd head_dim case to cover the scalar-layout instantiation, and a check_support rejection for head_dims beyond the launch bound. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
The backward previously required zero-initialized dKV/dScore buffers: the
kernel only wrote consumed positions, and two tensor-wide bf16 fill
kernels (torch.zeros_like) supplied the zeros everywhere else. Those
fills cost as much DRAM traffic as the kernel's own stores.
The kernel now writes every position itself: consumed positions get their
gradients as before (disjoint, atomic-free stores), and each
never-consumed slot class gets an exact zero from a unique natural owner,
keeping all stores disjoint and dKV/dScore bitwise run-to-run
deterministic:
- first-half columns of each segment's last block's own tokens
(coff == 2; no next block consumes them) -> that last block;
- per-segment tail tokens (seqlen % ratio, both halves) -> the
segment's last block;
- all tokens of segments with zero output blocks (seqlen < ratio) ->
CTA column bidx == 0;
- tokens beyond cu_seqlens[-1] (static token-capacity padding of the
gradient buffers, the CUDA-graph case) -> grid-strided across CTA
columns.
The backward wrapper allocates grad_kv/grad_score with torch.empty_like
instead of torch.zeros_like (the fp32 dAPE fill stays: it is the atomic
accumulator). When total_comp == 0 the kernel cannot launch, so the
wrapper falls back to zeroed allocations to preserve autograd's
exact-zero semantics; the class API documents the same caveat.
Measured on the backward region (kernel + remaining fills vs kernel +
three fills before), B200, ratio=4 coff=2, THD packs of 8192-token
sequences; nsys = sum of kernel durations per iteration, wall = CUDA
events around the wrapper call:
shape nsys before -> after wall before -> after
1x8192 d128 15.5 -> 12.8 us (1.21x) 56.2 -> 47.2 us (1.19x)
3x8192 d128 29.9 -> 22.2 us (1.35x) 61.8 -> 54.6 us (1.13x)
1x8192 d512 34.2 -> 22.8 us (1.50x) 67.9 -> 57.3 us (1.19x)
3x8192 d512 90.9 -> 66.0 us (1.38x) 111.6 -> 101.8 us (1.10x)
The zero-writes cost the kernel 8 registers (64 -> 72, occupancy 44% ->
41% at 3x8192 d512) and a small amount of extra store traffic, repaid
several times over by dropping the two whole-tensor fills.
dKV/dScore remain bitwise identical to the Megatron-LM original on the
5-shape port gate and to the fp32 eager reference on every test shape.
Tests: NaN-canary suite proving the kernel fully overwrites uninitialized
buffers (ragged, degenerate short segments, all-tiny packs, head_dim 512,
static-capacity padded rows) bitwise against zero-initialized runs and
the eager reference; a zeros-fallback test for packs where no segment
reaches ratio tokens. The existing CUDA-graph replay test covers the
token-capacity padding class (it fails without the grid-strided sweep).
Docs: csa.md updated (buffer contract, fill count, perf tables refreshed
for both this commit and the forward vectorization).
Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
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 (4)
🚧 Files skipped from review as they are similar to previous changes (4)
📝 WalkthroughWalkthroughAdds an experimental CSA fused Compressor API with lazy exports, Python wrappers, SM100 CuTe-DSL forward/backward kernels, cached launches, CUDA-graph support, documentation, and extensive correctness and runtime tests. ChangesCSA Fused Compressor
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant Caller
participant ForwardWrapper
participant API
participant Kernel
participant Output
Caller->>ForwardWrapper: submit packed CSA tensors
ForwardWrapper->>API: infer shape and fetch compiled configuration
API->>Kernel: launch fused compressor
Kernel->>Output: write compressed bf16 output
Output-->>Caller: return output
Possibly related issues
Suggested labels: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
🧹 Nitpick comments (1)
python/cudnn/csa/compressor/__init__.py (1)
8-13: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winSort the
__all__entries.Ruff RUF022 reports this list as unsorted. Place the backward symbols before the forward symbols in both class and wrapper groups.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@python/cudnn/csa/compressor/__init__.py` around lines 8 - 13, Sort the __all__ entries in the compressor module by placing CSACompressorBackward before CSACompressorForward, and csa_compressor_backward_wrapper before csa_compressor_forward_wrapper.Source: Linters/SAST tools
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@docs/fe-oss-apis/csa.md`:
- Around line 61-62: Update the support-bounds list in the CSA documentation to
include the enforced constraints from the compressor API: total_comp * head_dim
must be less than 2**31, and total_comp must be zero when total_tokens is less
than ratio. Preserve the existing bounds.
- Around line 190-192: Update the documented pytest command near the CSA
compressor test instructions to run from the test/python working directory,
using the fe_api/csa/test_CSA_compressor.py path so pytest.ini and conftest.py
are applied consistently.
In `@python/cudnn/csa/__init__.py`:
- Line 36: Update the module-level __all__ declaration to a literal list of
explicit string names, including CSA and every public symbol currently provided
by _SYMBOLS.keys(); remove the dynamic starred mapping-key expansion while
preserving the same export set.
In `@python/cudnn/csa/compressor/api.py`:
- Around line 109-118: Update the warnings.warn call in the
deterministic-algorithm handling block to pass stacklevel=2, so the
RuntimeWarning points to the immediate caller rather than this internal helper.
Leave the existing message, warning category, and exception path unchanged.
---
Nitpick comments:
In `@python/cudnn/csa/compressor/__init__.py`:
- Around line 8-13: Sort the __all__ entries in the compressor module by placing
CSACompressorBackward before CSACompressorForward, and
csa_compressor_backward_wrapper before csa_compressor_forward_wrapper.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: 64d2d78a-480d-43f2-bb95-7cfce924c524
📒 Files selected for processing (9)
docs/fe-oss-apis/csa.mddocs/fe-oss-apis/overview.mdpython/cudnn/__init__.pypython/cudnn/csa/README.mdpython/cudnn/csa/__init__.pypython/cudnn/csa/compressor/__init__.pypython/cudnn/csa/compressor/api.pypython/cudnn/csa/compressor/compressor_sm100.pytest/python/fe_api/csa/test_CSA_compressor.py
- docs: list the total_comp * head_dim < 2**31 and total_comp > 0 with total_tokens < ratio support boundaries already enforced in check_support - docs: run the compressor tests from test/python per repo convention - csa/__init__.py: spell out __all__ as an explicit literal (Ruff PLE0604) - compressor/__init__.py: sort __all__ (Ruff RUF022) - compressor/api.py: add stacklevel=2 to the deterministic-mode RuntimeWarning Signed-off-by: zky <51477259+zkyue@users.noreply.github.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
1 participant
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
Ports the fused CSA/HCA
Compressorgated-pooling kernels from Megatron-LM into theFE-OSS Python API, following the maintainer guidance on Megatron-LM PR
#5984 that the kernels belong in
cudnn-frontend like the other CSA/HCA fused kernels
(comment).
Measurements and numerics analysis for the original kernels are in Megatron-LM issue
#5968.
The
Compressorused by the CSA/HCA experimental attention variants implements itsgated-softmax pooling in eager PyTorch. For the THD packed layout, the region between
the two projection GEMMs and the RMSNorm — gather-index build, gather,
+ APE,overlap-window transform (
coff == 2), fp32 softmax over the window, gated weightedsum — decomposes into ~39 forward and ~51 backward kernel launches per call (at
compress_ratio = 4) and materializes(total_comp, 2*ratio, 1, head_dim)windowintermediates. This PR delivers that region as one forward and one backward CuTe-DSL
kernel behind the repo's
APIBasepattern.On top of the port (first commit, kernels byte-for-byte output-identical to the
Megatron original), two optimization commits — both identified by an ncu
hardware-ceiling audit of the ported kernels and both leaving forward/
dKV/dScorebitwise unchanged — are stacked separately for reviewability:
cute.autovec_copyover register fragments,2 head dims per thread): forward kernel 1.20–1.36x across the benchmark shapes.
every never-consumed
dKV/dScoreslot itself, so the wrapper allocates thosebuffers with
empty_likeand the two whole-tensor bf16 zero-fills disappear:backward region 1.21–1.50x.
Once this lands, Megatron-LM PR #5984 will be reworked into a thin dispatch that calls
these wrappers when
cudnn(with thecutedslextra) is available, keeping its eagerfallback — the autograd wiring stays on the framework side; the APIs here are pure
kernels-plus-validation.
Files
python/cudnn/csa/compressor/compressor_sm100.py— the two kernels plus launchmachinery (JIT cache per
(ratio, head_dim, coff, device), capture-safe JIT guard,cached-launch fast path). Kernel math is unchanged from the Megatron original
(bitwise-verified per commit); the optimization commits change the memory access
pattern (forward) and add the zero-write ownership scheme (backward).
python/cudnn/csa/compressor/api.py—CSACompressorForward/CSACompressorBackward(
APIBase:check_support/compile/execute) and the high-levelcsa_compressor_forward_wrapper/csa_compressor_backward_wrapper(allocateoutputs, return
TupleDict, LRU-cache compiled instances).python/cudnn/csa/— new package (CSAlazy namespace, README), registered inpython/cudnn/__init__.pylazy imports next toDSA/NSA. Happy to move/rename thepackage (e.g. under
deepseek_sparse_attention/) if you prefer a different home forCSA-side kernels.
docs/fe-oss-apis/csa.md(+overview.mdlink) — semantics, numerics contract,support surface, usage, perf tables.
test/python/fe_api/csa/test_CSA_compressor.py— 40 tests, see below.Support surface (
check_support)Compute capability 10.0 (the only validated architecture; the kernels use no
arch-specific features, so wider enablement is likely straightforward but stays opt-in
until validated),
ratio == 4/coff == 2(the production CSA/HCA configuration; thekernels are generic over
(ratio, head_dim, coff in {1, 2})and the gate can be liftedonce validated), BF16
kv/score/out, FP32ape, int32cu_seqlens/cu_seqlens_comp, int32 flat offsets (total_tokens * coff * head_dim < 2**31),head_dim <= 8388480(the forward launchgridDim.ybound; the pre-vectorizationschedule had the same envelope, just unchecked),
contiguous tensors on one CUDA device with 16-byte-aligned base pointers (4-byte for the
int32
cu_seqlens; contiguity does not imply base alignment for storage-offset views,so this is checked per call). Launches are anchored in the tensors' device context
(tensors on a non-current device work correctly), and explicit-stream calls
record_streamtheir operands so caller-released storages are not recycled while thekernel is pending.
Numerics
All arithmetic is fp32 with a single final bf16 rounding;
mul.rn.f32/fma.rn.f32arepinned in PTX so results do not depend on compiler FMA contraction.
verified on-device on 5 THD shape mixes (including ragged packs, segments shorter than
ratio, and static-capacity padding): forward,dKV,dScorealltorch.equalagainst the original module on identical inputs (
dAPEagrees within the fp32-atomicenvelope, max |diff| 2.5e-5; bitwise is not defined for either implementation there).
The gate was re-run after each optimization commit.
dKV/dScorebit-identical on every tested shape; forward differs on <0.006% ofelements within one bf16 rounding step.
output (the eager path itself rounds softmax weights to bf16 and multiplies in bf16).
dKV/dScore(disjoint stores). The kernel writes thosebuffers in full: consumed positions get their gradients, and every never-consumed
position (segment-tail tokens; the last block's first-half columns; segments shorter
than
ratio; token-capacity padding beyondcu_seqlens[-1]) gets an exact zero froma unique owning CTA — matching autograd without zero-initialized buffers (when
total_comp == 0nothing launches and the wrapper falls back to zeroed allocation).Forward,
dKV,dScorereplay bitwise identically run to run.dAPEisaccumulated with fp32 atomics and is not bitwise run-to-run deterministic; the
backward APIs therefore raise under strict
torch.use_deterministic_algorithms(True)and warn-and-run in warn-only mode (a deterministic per-CTA partials reduction is
possible follow-up work).
CUDA graphs & static-capacity padding
compile()) per(ratio, head_dim, coff)configuration; a call that would JIT undercapture raises a clear
RuntimeErrorinstead of corrupting the capture. Tested:capture fwd+bwd in one graph, replay with new data and a smaller device-side true
row count (static capacity, changed
cu_seqlenscontents), bitwisedKV/dScoreagainst the eager reference.
cu_seqlens_comp[-1]exactly like the eager code (windowfrom token 0 with first-in-segment semantics); backward ignores incoming gradients on
such padding rows — tested, including nonzero padding-row gradients and a pack whose
leading segment is shorter than
ratio.Performance
Measured on 1x B200 (CC 10.0, driver 590.48.01); PyTorch 2.13.0 (CUDA 13.3),
nvidia-cutlass-dsl4.6.1; BF16kv/score, FP32ape;ratio = 4,coff = 2; THDpacks of 8192-token sequences. Both implementations run on identical inputs over exactly
the replaced region (projection GEMMs, RMSNorm, RoPE excluded — unchanged); the eager
baseline is the verbatim replaced region of Megatron-LM
Compressor._forward_thd.Isolated GPU kernel time (nsys, sum of kernel durations per iteration, 50 iterations
after 20 warmup; no launch/host overhead; backward includes its
dAPEzero-fill):End-to-end wall clock of the same region (CUDA events, median of 100; includes launch
overhead; eager backward goes through torch autograd on the recorded eager graph, fused
backward is the explicit backward wrapper call — no autograd engine; not comparable to
the kernel-time rows):
How close to the hardware ceiling: an ncu audit of the ported kernels, prior to
the two optimization commits (cache-flushed,
--set full, all four benchmark shapes,forward and backward) showed measured DRAM read volume matching the algorithmically
necessary bytes within 1% — the THD gather adds no over-read; stores fully coalesced
(32/32 bytes per sector) and loads nearly so (29–30/32). Neither L2 (<27% of peak) nor
DRAM (<33%) was close to saturation at these microsecond-scale sizes: the gap to a pure
DRAM-floor time was a mix of sub-wave grid width / occupancy, memory latency and (at
the largest shape) issue pressure — not wasted traffic. The two audit follow-ups are
the two optimization commits in this PR (forward 32-bit vectorization; backward
zero-fill elimination); they do not change the bytes the kernels must read, so the
traffic-optimality conclusion carries over, while the utilization percentages quoted
predate them. Registers stay spill-free throughout (forward 48, backward 72 per
thread); wider forward vectors (64/128/256-bit) were measured and rejected —
instruction count drops but register pressure and occupancy losses dominate.
The wall-clock numbers rely on the cached-launch fast path
(
CUDNNFE_CSA_COMPRESSOR_FAST_LAUNCH=0disables it): a per-config snapshot of theCuTe-DSL launch state replayed with in-place argument mutation, which removes tens of
microseconds of per-call host overhead for these microsecond-scale kernels. The snapshot
construction introspects private-but-stable DSL internals; any structural mismatch on a
future
nvidia-cutlass-dslupgrade is caught at build time and the wrapper permanentlyfalls back to the regular launch path (it never launches from a partially built
snapshot). The isolated kernel times are launch-path independent.
Tests
test/python/fe_api/csa/test_CSA_compressor.py(40 tests, allL0; skip cleanlywithout
cudnn[cutedsl]or on non-CC-10.0 GPUs — the repo's pytest conftest itselfrequires CUDA). The eager reference is a self-contained
mirror of the replaced Megatron-LM region (gather ->
+APE-> overlap transform ->softmax -> weighted sum) with selectable numerics (verbatim upstream / fp32
intermediates / fp64 oracle):
dKV/dScore), vs upstream eager (tolerance), vsfp64 oracle (fused error <= eager error), over ragged multi-segment packs including
segments shorter than
ratio,head_dim128 and 512 plus an oddhead_dim(65)covering the scalar-layout forward instantiation;
segment shorter than
ratio); empty-output edge case;dKV/dScorebuffers staybitwise-equal to zero-initialized runs and to the eager reference across ragged /
degenerate / all-tiny / padded shapes — covering every never-consumed slot class
directly, including grad buffers with token-capacity padding beyond
cu_seqlens[-1]— and packs where no segment reaches
ratiotokens (total_comp == 0) fall back toexact-zero allocation;
dKV/dScore; deterministic-mode rejection of thebackward (strict) and warn-and-run semantics (warn-only);
which also exercises the token-capacity zero sweep) + the loud JIT-under-capture
error;
check_supportacceptance on meta samples and 13 rejection boundaries (ratio/coffenvelope, dtypes, shape mismatches, int32 flat-offset and head_dim launch bounds,
total_comp > 0withtotal_tokens < ratio, contiguity, CPU tensors);(
record_stream), multi-device launch correctness (tensors on a non-current device);black --line-length 160(26.3.1, per.pre-commit-config.yaml) clean.Summary by CodeRabbit
New Features
Documentation
Tests