CSA: add fused Compressor forward+backward CuTe-DSL kernels (ported from Megatron-LM) - #427
CSA: add fused Compressor forward+backward CuTe-DSL kernels (ported from Megatron-LM)#427zkyue wants to merge 10 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>
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds an experimental CSA fused Compressor API with lazy exports, Python wrappers, SM100 CuTe-DSL forward/backward kernels, ratio-specific dispatch, cached launches, CUDA-graph support, benchmarking, 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 Scheduler
participant Kernel
Caller->>ForwardWrapper: submit packed CSA tensors
ForwardWrapper->>API: infer shape and fetch compiled configuration
API->>Scheduler: select ratio-specific kernel and schedule
Scheduler->>Kernel: launch or replay cached arguments
Kernel-->>Caller: write compressed bf16 output
Suggested labels: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
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>
|
Thanks @zkyue for the contribution. |
Anerudhan
added
orig-external
|
Hi @zkyue , thanks for your contribution! I have some comments:
|
…raphs Address review feedback on the PR NVIDIA#427 fairness benchmark (benchmark/csa/bench_csa_compressor.py). - Drop the subtraction-derived backward-graph estimate. Graph variants are now two directly measured, symmetric columns for both eager and fused: a forward-only graph (eager under no_grad) and a forward+backward total graph. - Capture the eager fwd+bwd graph with stable, pre-allocated zero .grad buffers (the supported torch CUDA-graph pattern): the captured region zeros them in place, then runs forward + autograd backward, so every replay accumulates into a zeroed buffer -- numerically identical to a single fresh backward. A per-shape graph-vs-fresh-backward cross-check is printed (all shapes bitwise-equal). - The fused total graph captures the forward wrapper immediately followed by the backward wrapper. - Emit both the per-call and graph tables from a single run; graph speedups are eager/fused of the displayed us, truncated to one decimal (never rounded up). The backward replay is reported only as ~total-fwd (a reference), never as a column. Not collected by pytest; black -l160 clean. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
Address review feedback on the CSA compressor performance tables (docs/fe-oss-apis/csa.md). - Replace the CUDA-graph table (which published a subtraction-derived backward column) with two directly measured columns: eager/fused forward-only graph and eager/fused forward+backward total graph. No backward-graph column; backward is noted only as ~total-fwd, never as a measured value. - Refresh both the per-call and graph tables from a single run of the benchmark harness; graph speedups are eager/fused of the displayed us, truncated to one decimal. - State that capture collapses each side's per-op launches into a single replay rather than removes launch/host overhead; say less launch-bound rather than compute-bound. - Note how the re-run reproduces the previously published per-call wall clock, and document the eager-backward capture basis (stable zero .grad buffers). Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
|
Thanks for both questions — answers below with measurements. Q1 — extending to
Q2 — eager baseline and CUDA Graphs. You're right — the eager baseline in the harness was not CUDA-graphed; it timed each The graph variants are captured symmetrically as two directly measured columns for each Headline numbers (µs, B200, Graphing both sides is the fair comparison: capturing each side into a graph collapses its |
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
benchmark/csa/bench_csa_compressor.py (2)
271-277: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDuplicated leaf-tensor construction across three sites.
The
kv/score/ape→.clone().requires_grad_(True)triple is repeated verbatim in_capture_eager_total(twice) and inmeasure's_fwd_for_bwd. Extracting a small helper reduces the chance of the three copies drifting (e.g. if a dtype/device nuance needs to change later).♻️ Suggested refactor
+def _make_leaves(kv, score, ape): + return kv.clone().requires_grad_(True), score.clone().requires_grad_(True), ape.clone().requires_grad_(True) + + def _capture_eager_total(kv, score, ape, cu, cuc, total_comp, go, ratio, d, coff, warmup): ... - kvl = kv.clone().requires_grad_(True) - scl = score.clone().requires_grad_(True) - apl = ape.clone().requires_grad_(True) + kvl, scl, apl = _make_leaves(kv, score, ape) ... - ref_k = kv.clone().requires_grad_(True) - ref_s = score.clone().requires_grad_(True) - ref_a = ape.clone().requires_grad_(True) + ref_k, ref_s, ref_a = _make_leaves(kv, score, ape)Also applies to: 304-306, 340-346
🤖 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 `@benchmark/csa/bench_csa_compressor.py` around lines 271 - 277, Extract the repeated kv/score/ape leaf-tensor construction into a small helper and reuse it at the three sites in _capture_eager_total and measure’s _fwd_for_bwd. The helper should clone each input and enable gradients, while preserving the existing stable zero grad buffer initialization where required.
117-117: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueMinor lint cleanups flagged by Ruff.
[0] + list(...)can be[0, *list(...)](RUF005) at Lines 117 and 119, and the loop variabledis unused in the loop bodies at Lines 424, 432, 445, 458 (B007, rename to_).Also applies to: 119-119, 424-424, 432-432, 445-445, 458-458
🤖 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 `@benchmark/csa/bench_csa_compressor.py` at line 117, Apply Ruff’s minor lint fixes in the benchmark code: replace the list concatenations used to build `cu` at both affected locations with iterable unpacking, and rename the unused loop variable `d` to `_` in the loops around the four reported locations. Preserve the existing computation and loop behavior.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 `@benchmark/csa/bench_csa_compressor.py`:
- Around line 384-392: Guard all displayed-timing divisions against zero in _x
and _xtrunc1, and in the inline ef / ff and eb / fb calculations near the
reported rows. Preserve the existing speedup output for nonzero denominators,
while returning a safe fallback or otherwise skipping only the affected speedup
value when a rounded denominator is 0.0 so collected rows still print.
---
Nitpick comments:
In `@benchmark/csa/bench_csa_compressor.py`:
- Around line 271-277: Extract the repeated kv/score/ape leaf-tensor
construction into a small helper and reuse it at the three sites in
_capture_eager_total and measure’s _fwd_for_bwd. The helper should clone each
input and enable gradients, while preserving the existing stable zero grad
buffer initialization where required.
- Line 117: Apply Ruff’s minor lint fixes in the benchmark code: replace the
list concatenations used to build `cu` at both affected locations with iterable
unpacking, and rename the unused loop variable `d` to `_` in the loops around
the four reported locations. Preserve the existing computation and loop
behavior.
🪄 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: 5c82cdfa-7f83-42fd-97cf-ef9d92022166
📒 Files selected for processing (2)
benchmark/csa/bench_csa_compressor.pydocs/fe-oss-apis/csa.md
🚧 Files skipped from review as they are similar to previous changes (1)
- docs/fe-oss-apis/csa.md
Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
The kernels are already generic over coff — for coff=1 the window is the
block's own ratio tokens (no overlap transform, always valid); only the
check_support gate restricted the APIs to the production coff=2 form.
Widen the gate to coff in {1, 2} (ratio stays gated at 4), update the
error message and the envelope prose in api.py / docs / README, and
parametrize the tests over both coff forms: numerics vs the
fp32/upstream/fp64 eager references (the eager reference already
implements coff=1 faithfully — it skips the overlap transform), static-
capacity padding, replay determinism, empty outputs, NaN-canary
zero-writes, deterministic-mode error paths, CUDA-graph capture,
class-vs-wrapper equivalence, and check_support accept/reject
boundaries (new coff=0 / coff=3 rejects; a new empty-segment shape
covers both coff forms). No kernel changes.
Addresses hxbai's coff=1 question on the PR.
Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
Bring the PR's Python docstring coverage over CodeRabbit's 80% pre-merge gate: module docstrings for the csa package/compressor __init__ files and one-line docstrings for the remaining undocumented helpers in api.py (cache/ctor/compile plumbing), compressor_sm100.py (launcher ctors), the test helpers, and the benchmark helpers. interrogate 1.7.0 over the PR's six CSA Python files: 69.4% -> 100%. Docstrings only — no logic, kernel, or launch-path changes. Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
|
LGTM too |
…th ratio dispatch
The generic CSA compressor kernels keep the whole coff * ratio pooling
window in per-thread registers; past ratio ~32 that hits the
255-register wall and spills kilobytes of local memory per thread, so
ratio=128 needs dedicated kernels rather than a lifted gate. This adds
them behind the same public API: check_support now admits
ratio in {4, 128} x coff in {1, 2} (ratio=128 additionally gated to
head_dim in {128, 512}, the hardware-validated set) and the wrappers
and class APIs route to the matching kernel family by ratio
transparently.
Design: the forward streams the window with a chunked online softmax
(one CTA per output row, per-column (max, denom, acc) triples merged in
one fixed smem order); its launch schedule is bucketed by output-row
count (small/default/large, all precompiled for CUDA-graph safety) and
per bucket selects two-phase accumulation and an ex2.approx fast exp
where measured faster. The backward stages each row's window through
shared memory chunk-parallel, fuses the per-chunk den/S partial sums
into the same pass, merges partials in a fixed chunk order, and stores
gradients with a hoisted 1/den multiply — keeping kernel-side
zero-writes to never-consumed dKV/dScore slots and fp32-atomic dAPE
exactly as at ratio=4. ptxas (sm_100a): forward 32-51 registers,
backward 48-128, 0 spill / 0 stack across all 16 shipped
(config, schedule) kernels (reproducer:
benchmark/csa/reg_probe_csa_compressor_r128.py).
Numerics contract, per ratio family (docs/fe-oss-apis/csa.md): ratio=4
keeps its bitwise contract UNCHANGED — dKV/dScore bit-identical to the
fp32-intermediate eager reference, at coff 1 and 2. ratio=128 is
deterministic + faithful to that same fp32-intermediate eager
reference: bitwise run-to-run determinism of out/dKV/dScore
(NaN-prefill replay tested); the same values within final-bf16 rounding
at the gate tolerances (differing elements <= max(1, 0.1%), max_abs <=
1.6e-2, calibrated on the gate's documented input distribution — bf16
deviations scale 2^k with 2^k inputs while the fp32 intermediates stay
finite); the eager reference's non-finite propagation (finite inputs
that overflow its fp32 intermediates poison both sides alike, with the
stated caveat that the fused order saturates earlier near fp32 max —
un-normalized chunk partials); and, on inputs whose fp32 intermediates
stay finite in both evaluation orders, fp64-oracle parity (at least as
close to an fp64 oracle as the eager reference). The reduction reorders
and fast-exp buckets were each adopted on a measured win and are
covered by that contract.
Performance (nsys isolated GPU kernel time, 1x B200, vs the
fp32-intermediate eager reference region on identical inputs): forward
13.2-44.4x, backward 7.4-21.9x across coff {1, 2} x head_dim {128, 512}
x 8k-131k tokens; e.g. coff=2, d=128, 65k tokens: fwd 726.3 -> 16.4 us
(44.4x), bwd 960.0 -> 61.2 us (15.7x). Full tables with methodology in
docs/fe-oss-apis/csa.md.
Validation (B200, CC 10.0): the committed 21-case contract gate
(benchmark/csa/gate_csa_compressor_r128.py) passes 21/21 with shipped
schedule coverage asserted (forward 10/10, backward 6/6 buckets);
pytest fe_api/csa/test_CSA_compressor.py: 88 passed, 1 skipped at the
default L0 level and 98 passed, 1 skipped with -m "L0 or L1" (both
ratio families; the ratio=4 rows, including the coff=1 envelope, stay
on their bitwise assertions). New tests cover the ratio=128 dispatch
envelope (schedule selection at every nb_total bucket boundary; every
shipped (config, schedule) kernel executed against the full contract at
L1), exact-zero assertions on every never-consumed dKV/dScore slot
class in the NaN-canary, CUDA-graph capture/replay for the new family,
and grad_ape zeroing-ownership regressions.
Signed-off-by: zky <51477259+zkyue@users.noreply.github.com>
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
python/cudnn/csa/compressor/compressor_sm100_r128.py (1)
857-916: 🚀 Performance & Scalability | 🔵 Trivial | 💤 Low valueOptional: the
grad_outvector is loaded twice per row.
fr_go2(phase 2) andfr_go(phase 4) load the samemGOslice for the same(bb, cvec). Registers survivecute.arch.barrier(), so phase 4 could reuse the phase-2 value and drop one global load per row. Only worth it if it doesn't push the register budget past the ptxas 0-spill target you documented.🤖 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/compressor_sm100_r128.py` around lines 857 - 916, Optionally reuse the phase-2 grad_out register tensor `fr_go2` in phase 4 instead of reloading the same `mGO` slice into `fr_go`, preserving correctness across the barrier. Verify ptxas remains within the documented zero-spill register budget; retain the existing phase-4 load if reuse causes spills or exceeds that budget.benchmark/csa/gate_csa_compressor_r128.py (1)
130-213: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winConsider extracting the shared eager reference instead of duplicating it.
_batch_of_row,_overlap_transform_thd,_eager_pool,run_eager_bwd,_make_inputsand_make_goare verbatim copies oftest/python/fe_api/csa/test_CSA_compressor.py(lines ~73-163). The module docstring stakes the gate's authority on them being "identical", but nothing enforces that — a fix to one copy silently invalidates the other. A single importable helper (e.g.benchmark/csa/_csa_eager_ref.pyconsumed by both) would keep the guarantee mechanical.🤖 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 `@benchmark/csa/gate_csa_compressor_r128.py` around lines 130 - 213, The eager reference helpers are duplicated between the benchmark gate and the CSA compressor test, so the claimed identical behavior is not enforced. Extract _batch_of_row, _overlap_transform_thd, _eager_pool, run_eager_bwd, _make_inputs, and _make_go into one importable shared helper module, then update both consumers to import and use those definitions while preserving their current behavior and interfaces.
🤖 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 `@benchmark/csa/reg_probe_csa_compressor_r128.py`:
- Around line 75-85: Update the kernel probing loop over M._COMPILED in the
relevant probe function to use the documented compiled-handle PTX accessor
fn.ptx instead of fn.artifacts.PTX. If PTX retention is disabled and fn.ptx is
unavailable, guard that case and skip the kernel without aborting the probe,
while preserving the existing all_clean and count behavior for processed
kernels.
---
Nitpick comments:
In `@benchmark/csa/gate_csa_compressor_r128.py`:
- Around line 130-213: The eager reference helpers are duplicated between the
benchmark gate and the CSA compressor test, so the claimed identical behavior is
not enforced. Extract _batch_of_row, _overlap_transform_thd, _eager_pool,
run_eager_bwd, _make_inputs, and _make_go into one importable shared helper
module, then update both consumers to import and use those definitions while
preserving their current behavior and interfaces.
In `@python/cudnn/csa/compressor/compressor_sm100_r128.py`:
- Around line 857-916: Optionally reuse the phase-2 grad_out register tensor
`fr_go2` in phase 4 instead of reloading the same `mGO` slice into `fr_go`,
preserving correctness across the barrier. Verify ptxas remains within the
documented zero-spill register budget; retain the existing phase-4 load if reuse
causes spills or exceeds that budget.
🪄 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: c0694b0f-5805-4238-8542-7555a2d89388
📒 Files selected for processing (11)
benchmark/csa/bench_csa_compressor.pybenchmark/csa/gate_csa_compressor_r128.pybenchmark/csa/reg_probe_csa_compressor_r128.pydocs/fe-oss-apis/csa.mdpython/cudnn/csa/README.mdpython/cudnn/csa/__init__.pypython/cudnn/csa/compressor/__init__.pypython/cudnn/csa/compressor/api.pypython/cudnn/csa/compressor/compressor_sm100.pypython/cudnn/csa/compressor/compressor_sm100_r128.pytest/python/fe_api/csa/test_CSA_compressor.py
🚧 Files skipped from review as they are similar to previous changes (6)
- python/cudnn/csa/README.md
- python/cudnn/csa/compressor/init.py
- python/cudnn/csa/init.py
- docs/fe-oss-apis/csa.md
- benchmark/csa/bench_csa_compressor.py
- python/cudnn/csa/compressor/compressor_sm100.py
| for key, fn in sorted(M._COMPILED.items(), key=str): | ||
| kind, _ratio, d, coff, sched, _dev = key | ||
| if kind == "r128fwd": | ||
| vec, tchunks, threads_x, twophase, fastexp = sched | ||
| tag = f"fwd_c{coff}d{d}_v{vec}t{tchunks}x{threads_x}" + ("_2ph" if twophase else "") + ("_fexp" if fastexp else "") | ||
| else: | ||
| vec, tchunks, threads_x, fastexp = sched | ||
| tag = f"bwd_c{coff}d{d}_v{vec}t{tchunks}x{threads_x}" + ("_fexp" if fastexp else "") | ||
| all_clean = ptxas_one(tag, fn.artifacts.PTX, args.arch, out_dir) and all_clean | ||
| n += 1 | ||
| print(f"{n} kernels probed ({args.arch}); {'ALL 0 spill / 0 stack' if all_clean else 'SPILL/STACK OR PTXAS FAILURE DETECTED'}", flush=True) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🌐 Web query:
nvidia-cutlass-dsl cute.compile compiled artifacts PTX attribute CUTE_DSL_KEEP
💡 Result:
In NVIDIA CUTLASS CuTe DSL, CUTE_DSL_KEEP is an environment variable used to control the dumping of generated artifacts (such as IR, PTX, CUBIN, and SASS) to files for offline inspection [1][2]. Key details regarding artifacts and compilation options: Environment Variable CUTE_DSL_KEEP You can set CUTE_DSL_KEEP to a comma-separated list of tokens to save specific artifacts [1][3]. For example: - export CUTE_DSL_KEEP=ptx (Saves PTX assembly) [1][2] - export CUTE_DSL_KEEP=cubin (Saves CUBIN binary) [1][2] - export CUTE_DSL_KEEP=all (Saves all supported artifacts) [1][2] Programmatic Access to Compiled Artifacts When a kernel is compiled using cute.compile, the resulting object allows programmatic access to the generated artifacts via specific attributes [1][2]. Importantly, these attributes are only populated if the corresponding environment variable (e.g., CUTE_DSL_KEEP_PTX=1 or the relevant CUTE_DSL_KEEP token) is set [4]: - compiled_func.ptx: Accesses the generated PTX code [1][2]. - compiled_func.cubin: Accesses the generated CUBIN data [1][2]. - compiled_func.mlir: Accesses the generated IR code [1][2]. Compilation Options in cute.compile Alternatively, you can control the generation of these files directly within your Python code when calling cute.compile by using specific compilation options [5]: - KeepPTX: Equivalent to setting a flag to keep the generated PTX file [5]. - KeepCUBIN: Equivalent to setting a flag to keep the generated CUBIN file [5]. These options are typically imported from cutlass.cute and passed into cute.compile to manage artifact retention without relying solely on environment variables [5].
Citations:
- 1: https://docs.nvidia.com/cutlass/4.5.3/media/docs/pythonDSL/cute_dsl_general/debugging.html
- 2: https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/cute_dsl_general/debugging.html
- 3: https://docs.nvidia.com/cutlass/4.5.2/media/docs/pythonDSL/cute_dsl_general/debugging.html
- 4: docs: note when DSL dumps are populated cutlass#2891
- 5: https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.html
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Map the benchmark file and inspect the relevant section.
ast-grep outline benchmark/csa/reg_probe_csa_compressor_r128.py --view expanded || true
echo '---'
cat -n benchmark/csa/reg_probe_csa_compressor_r128.py | sed -n '1,140p'
echo '---'
# Search for artifact access patterns in the repo.
rg -n "artifacts\.PTX|\.ptx\b|KeepPTX|CUTE_DSL_KEEP|artifacts\." benchmark -S || trueRepository: NVIDIA/cudnn-frontend
Length of output: 5956
🏁 Script executed:
#!/bin/bash
set -euo pipefail
python3 - <<'PY'
from pathlib import Path
p = Path('benchmark/csa/reg_probe_csa_compressor_r128.py')
print(p.exists(), p.stat().st_size if p.exists() else None)
PYRepository: NVIDIA/cudnn-frontend
Length of output: 169
🌐 Web query:
"fn.artifacts.PTX" "cutlass" "cute.compile" OR "compiled_func.ptx"
💡 Result:
In NVIDIA CUTLASS, when using the CuTe Python DSL, cute.compile generates a compiled JIT executor object [1][2]. To access the generated PTX artifacts from this compiled function (e.g., compiled_func.ptx), you must explicitly enable the retention of these artifacts during the compilation process [1][3]. By default, these attributes are not populated [3]. You can enable them using one of the following methods: 1. Environment Variables: Set the CUTE_DSL_KEEP environment variable to include ptx (e.g., export CUTE_DSL_KEEP=ptx) [1][4]. This allows you to access the generated PTX code programmatically via the ptx attribute of the compiled function object [1][3]. 2. Compilation Options: Pass KeepPTX as a compilation option to cute.compile [5][6][7]. This can be done using the Pythonic composition interface: python from cutlass.cute import KeepPTX compiled_func = cute.compile[KeepPTX](my_kernel, ...) print(compiled_func.ptx) Alternatively, you can pass it as a string option if supported by your version of CUTLASS (e.g., options="--keep-ptx") [8][7]. Note that simply referring to compiled_func.ptx without having set these configuration flags will typically return None [3]. Accessing these artifacts is primarily intended for debugging, profiling, and offline inspection of generated assembly [1][4][9].
Citations:
- 1: https://docs.nvidia.com/cutlass/4.5.3/media/docs/pythonDSL/cute_dsl_general/debugging.html
- 2: https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/cute_dsl_general/dsl_jit_caching.html
- 3: docs: note when DSL dumps are populated cutlass#2891
- 4: https://docs.nvidia.com/cutlass/4.5.1/media/docs/pythonDSL/cute_dsl_general/debugging.html
- 5: https://docs.nvidia.com/cutlass/latest/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.html
- 6: https://docs.nvidia.com/cutlass/4.5.3/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.html
- 7: https://docs.nvidia.com/cutlass/4.4.2/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.html
- 8: https://docs.nvidia.com/cutlass/4.3.2/media/docs/pythonDSL/cute_dsl_general/dsl_jit_compilation_options.html
- 9: https://github.com/NVIDIA/cutlass/blob/master/CHANGELOG.md
Use the documented PTX handle fn.artifacts.PTX isn’t the public CuTe DSL surface; the compiled handle exposes fn.ptx when PTX retention is enabled. Switching to that accessor — or guarding it and skipping the kernel — avoids an AttributeError that aborts the whole probe.
🤖 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 `@benchmark/csa/reg_probe_csa_compressor_r128.py` around lines 75 - 85, Update
the kernel probing loop over M._COMPILED in the relevant probe function to use
the documented compiled-handle PTX accessor fn.ptx instead of fn.artifacts.PTX.
If PTX retention is disabled and fn.ptx is unavailable, guard that case and skip
the kernel without aborting the probe, while preserving the existing all_clean
and count behavior for processed kernels.
|
@hxbai two updates on this branch since your last pass (cc @jiayus-nvidia — thanks for the LGTM, and apologies for moving the target right after it): coff=1: delivered as promised, same bitwise contract. ratio=4 now admits ratio=128: delivered in this PR (one commit,
Branch state validated end-to-end: pytest 88 passed / 1 skipped at L0 (98 / 1 at |
|
@zkyue thanks for your contribution! What about putting these kernels into https://github.com/NVIDIA/cudnn-frontend/tree/develop/python/cudnn/deepseek_sparse_attention rather than creating a new |
|
Absolutely, happy to relocate. Two shapes, whichever you both prefer: (a) flat into |
|
Thanks @zkyue I do prefer the current directory structure. Also: We are planning to move to Apache-2 License from MIT license. Pinging here for your approval for the files you authored. |
4 participants
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