Add fused CuTe-DSL forward+backward kernels for the CSA Compressor gated pooling (THD)

[zkyue]

What

Addresses #5968 (contribution invited by the maintainers in
#5968 (comment)).

The Compressor used by the CSA/HCA experimental attention variants implements its gated
softmax pooling in eager PyTorch. In _forward_thd, 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 weighted sum — decomposes into ~39
forward and ~51 backward kernel launches per call (at compress_ratio = 4), materializes
(total_comp, 2*ratio, 1, head_dim) window intermediates (also saved for autograd), and
its backward pays an indexing_backward plus a radix sort for the gather.

This PR adds an additive fused fast path for exactly that region: one forward and one
backward kernel written in the CuTe Python DSL (nvidia-cutlass-dsl, JIT-compiled at
first call), keeping the existing eager implementation as the semantic reference and
fallback. Forward+backward replace the ~90 eager launches with 1 + 1 kernels (plus 3
grad-buffer zero-fills in backward).

Files

• megatron/core/transformer/experimental_attention_variant/csa_fused_compressor.py — new,
  self-contained: import-guarded DSL imports, the two kernels, a cached-launch fast path,
  the autograd wrapper, the explicit op API (compress_thd_fused) and the dispatch helper
  (maybe_compress_thd_fused).
• megatron/core/transformer/experimental_attention_variant/csa.py — minimal insertion in
  _forward_thd: try the fused path, fall into the (unchanged, re-indented) eager region
  when it returns None.
• tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_compressor.py — unit tests.
• tests/unit_tests/transformer/experimental_attention_variant/bench_csa_fused_compressor.py —
  the standalone benchmark/correctness harness behind the numbers below (not collected by pytest).

Dispatch / gating (initial)

The fast path engages only for: THD packed non-pre-grouped path, compress_ratio == 4
(coff == 2), bf16 kv/score, fp32 ape, compute capability 10.0 (the only
validated architecture; the kernels use no TMA/tcgen05/MMA features, so wider coverage is
likely straightforward but stays opt-in until validated), and int32 flat offsets
(total_tokens * coff * head_dim < 2**31). Everything else — SBHD, the pre-grouped CP-prep
input path, compress_ratio == 128, missing DSL, other devices — keeps the eager
implementation. MCORE_CSA_FUSED_COMPRESSOR=0 disables the dispatch entirely.
compress_ratio == 128 (coff == 1) is functionally supported by the explicit op API and
covered by the tests, but stays on eager until tuned (currently only ~3.1× fwd / ~1.2× bwd
wall-clock there). The dispatch also keeps the eager path under
torch.use_deterministic_algorithms(True) (see dAPE below) and under torch.compile
tracing (the raw-pointer launch path is not traceable; eager lets the compiler fuse the
region itself).

Numerics

All arithmetic is fp32 with a single final bf16 rounding; mul.rn.f32/fma.rn.f32 are
pinned in PTX so results do not depend on compiler FMA contraction. Fresh numbers from the
included harness on the ported code (B200, details below):

• vs an fp32-intermediate eager reference (same ops, fp32 throughout, one final bf16
  rounding): dKV/dScore bit-identical on every tested shape; forward differs on
  <0.006% of elements with max abs difference 3.9e-3–7.8e-3.
• vs the current eager code the outputs are not bit-identical, because the eager path
  itself rounds intermediates to bf16 (torch.softmax(...).to(kv.dtype), bf16 multiply).
  Against an fp64 oracle the fused kernels' max abs error is equal to or lower than
  eager's on every tested output (table in the harness output / issue).
• Backward is atomic-free for dKV/dScore (disjoint stores; unconsumed elements keep
  exact zeros, matching autograd). Forward, dKV, dScore replay bitwise identically run
  to run. dAPE is accumulated with fp32 atomics and is not bitwise run-to-run
  deterministic (max run-to-run deviation 1.5e-5–6.9e-5 over 10 replays per shape in the
  issue measurements, fp32); the dispatch therefore falls back to eager under
  torch.use_deterministic_algorithms(True), and the explicit op raises in backward in
  that mode. A deterministic per-CTA-partials reduction is possible follow-up work if
  required.

CUDA graphs & fixed_total_comp (new vs the issue)

Both items the issue left open are implemented and tested here:

• Graph capture: the launch path is capture-compatible after a one-step eager warmup
  per (ratio, head_dim, coff) configuration; a first call that would JIT under capture
  raises a clear RuntimeError instead of corrupting the capture. The unit test captures
  fwd+bwd in one graph and replays with new data and a smaller device-side true row
  count (static capacity, changed cu_seqlens contents), checking bitwise dKV/dScore
  against the eager reference.
• Static-capacity padding rows: forward computes padding rows exactly like the eager
  code (they gather the window from token 0 with first-in-segment semantics, as the eager
  gather does); backward ignores incoming gradients on padding rows (they are tail
  padding, not consumed downstream) — tested, including that nonzero padding-row gradients
  change nothing and including a pack whose leading segment is shorter than ratio.

Performance

Measured on 1×B200 (CC 10.0, driver 590.48.01); PyTorch 2.13.0a0 (CUDA 13.3),
nvidia-cutlass-dsl 4.6.1; bf16 kv/score (total, 1, coff*head_dim), fp32 ape;
compress_ratio = 4, coff = 2; THD packs of 8192-token sequences. Both implementations
run on identical inputs over exactly the replaced region (projection GEMMs, RMSNorm, RoPE
excluded — unchanged). Numbers re-measured on this PR's code with the included harness.

Isolated GPU kernel time (nsys, sum of kernel durations per iteration, 50 iterations
after 20 warmup; no launch/host overhead; backward includes its grad-buffer fills):

THD pack	head_dim	eager fwd	fused fwd	fwd	eager bwd	fused bwd	bwd
1×8192	128	119.4 µs	6.0 µs	19.8×	193.1 µs	16.0 µs	12.1×
3×8192	128	231.2 µs	12.6 µs	18.3×	359.7 µs	29.8 µs	12.1×
1×8192	512	265.9 µs	15.9 µs	16.7×	430.9 µs	34.1 µs	12.6×
3×8192	512	665.7 µs	42.3 µs	15.8×	1164.1 µs	91.0 µs	12.8×

End-to-end wall clock of the same region (CUDA events, median of 100, includes launch
overhead and, for backward, autograd engine overhead — not comparable to the kernel-time
numbers above):

THD pack	head_dim	eager fwd	fused fwd	fwd	eager bwd	fused bwd	bwd
1×8192	128	336.4 µs	34.5 µs	9.8×	544.4 µs	110.0 µs	4.9×
3×8192	128	374.3 µs	37.6 µs	10.0×	568.8 µs	111.4 µs	5.1×
1×8192	512	418.4 µs	40.0 µs	10.5×	649.9 µs	111.1 µs	5.9×
3×8192	512	825.4 µs	66.9 µs	12.3×	1488.6 µs	142.6 µs	10.4×

(The fused wall numbers go through the explicit op API of the harness and therefore
include its host-side argument validation, a few µs per call; the production dispatch
path performs its cheaper gating checks instead. They are slightly below the issue's
wall numbers, which were measured on the pre-review version of the wrapper without the
added validation; the isolated kernel times above are unchanged.)

Tests

tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_compressor.py
(all skip cleanly without CUDA / DSL / CC 10.0):

• numerics vs fp32-eager (bitwise dKV/dScore), vs verbatim upstream eager (tolerance),
  vs fp64 oracle (fused error ≤ eager error), over ragged multi-segment packs including
  segments shorter than ratio;
• coff == 1 (compress_ratio == 128) functional correctness;
• fixed_total_comp static-capacity padding (fwd semantics + ignored padding grads,
  including a pack with a leading segment shorter than ratio);
• run-to-run determinism of forward/dKV/dScore;
• CUDA-graph capture/replay + loud JIT-under-capture error;
• dispatch gating and eager fallback (maybe_compress_thd_fused), including the
  deterministic-mode fallback (and the explicit op raising in backward under
  torch.use_deterministic_algorithms(True)), plus Compressor._forward_thd-level
  integration (fused engages, matches eager, gradients flow).

The module carries pytestmark = pytest.mark.launch_on_gb200 so the CC 10.0 CI lane
exercises it; everywhere else the tests skip cleanly (no CUDA / no DSL / other archs).

Also ran the existing related suites with the fused dispatch live:
test_attention_variant_csa.py, test_csa_cp_layout_kernels.py, test_csa_cp_utils.py —
136 passed, 1 skipped (needs ≥2 ranks), 0 failed.

Notes

• The fast path is currently on by default for the supported configuration (with the
  MCORE_CSA_FUSED_COMPRESSOR=0 kill switch), matching the scope proposed in [ENHANCEMENT] Fused forward+backward CUDA kernel for the CSA/HCA Compressor gated pooling (THD path) #5968. It
  does change training numerics relative to the current eager code (more accurate against
  the fp64 oracle, but not bit-identical) — happy to flip this to opt-in if you prefer.
• The cached-launch fast path introspects private-but-stable DSL internals to remove
  per-call host overhead; any structural mismatch on a future nvidia-cutlass-dsl
  upgrade raises during construction and the wrapper permanently falls back to the
  regular launch path (it never launches from a partially built snapshot).
  MCORE_CSA_FUSED_COMPRESSOR_FAST_LAUNCH=0 disables just this optimization.
• Follow-up candidates (out of scope here): deterministic dAPE reduction, tuning for
  compress_ratio == 128, wider arch enablement after validation.

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[hxbai]

Hi @zkyue Thanks for your contribution!

I think it is better to put the kernels in cudnn-frontend like other CSA/HCA fused kernels. Is it convenient for you to do so, or do you need our help to port the kernel?

[zkyue]

@hxbai Sure, happy to do that — the kernels are self-contained CuTe DSL and should map cleanly onto the cudnn-frontend python API structure (we've contributed there before). Plan: I'll open a PR against cudnn-frontend porting the fwd+bwd kernels (APIBase-style wrapper, THD semantics, tests), then rework this PR into a thin dispatch that uses the cudnn-frontend implementation when available, keeping the eager fallback. Will link the kernel PR here once it's up.

[zkyue]

@hxbai The cudnn-frontend kernel PR is up: NVIDIA/cudnn-frontend#427 (fwd+bwd kernels, APIBase-style wrapper, THD semantics, 40 tests; outputs bit-identical to the kernels in this PR). Once it lands I'll rework this PR into the thin dispatch that uses the cudnn-frontend implementation when available, keeping the eager fallback.