CSA: add fused Compressor forward+backward CuTe-DSL kernels (ported from Megatron-LM)

docs/fe-oss-apis/csa.md

@@ -0,0 +1,201 @@
+# CSA Fused Compressor
+
+**This is an experimental API and subject to change.**
+
+## Overview
+
+The CSA module hosts CuTe-DSL kernels for the CSA/HCA experimental attention variants
+(the components that are not shared with the [DSA module](dsa.md)). Its first operation
+is the **fused Compressor**: one forward and one backward kernel for the `Compressor`
+gated-softmax pooling region (THD packed layout) used by CSA/HCA in Megatron-LM.
+
+The kernels were ported from Megatron-LM at the maintainers' request
+([Megatron-LM PR #5984](https://github.com/NVIDIA/Megatron-LM/pull/5984); measurements
+and numerics in
+[Megatron-LM issue #5968](https://github.com/NVIDIA/Megatron-LM/issues/5968)). The eager
+region they replace decomposes into ~39 forward and ~51 backward kernel launches per
+call (at `compress_ratio = 4`) and materializes `(total_comp, 2*ratio, 1, head_dim)`
+window intermediates; the fused path is 1 + 1 kernels (plus one fp32 `dAPE`-buffer
+zero-fill in backward — the backward kernel writes `dKV`/`dScore` in full, including
+exact zeros to never-consumed positions, so those buffers need no fills).
+
+### Semantics
+
+For each THD segment `s` (`cu_seqlens[s]..cu_seqlens[s+1]`) and each output block `b` of
+`ratio` tokens, with the overlapping window form (`coff == 2`, window size `2 * ratio`):
+
+- `k in [0, ratio)`: previous block's token, first-half projection column, APE row `k`
+  — invalid for the segment's first block (score `-inf`, kv `0`);
+- `k in [ratio, 2*ratio)`: own block's token, second-half projection column, APE row
+  `k - ratio`.
+
+```text
+out[b, j] = sum_k kv[w(b,k), c(k,j)] * softmax_k(score[w(b,k), c(k,j)] + ape[k % ratio, c(k,j)])
+```
+
+computed in fp32 with a single final bf16 rounding. Per-segment tail tokens
+(`seqlen % ratio`) are dropped, as in the eager code. Output rows beyond
+`cu_seqlens_comp[-1]` (a static CUDA-graph capacity) are computed with first-in-segment
+semantics from token 0, exactly like the eager gather; the backward ignores incoming
+gradients on such padding rows.
+
+### Numerics
+
+All arithmetic is fp32 with one final bf16 rounding; `mul.rn.f32` / `fma.rn.f32` are
+pinned in PTX so results do not depend on compiler FMA contraction. Against an
+fp32-intermediate eager reference (same op order, fp32 throughout), `dKV`/`dScore` are
+**bit-identical** and the forward matches within one bf16 rounding step on a tiny
+fraction of elements. Forward, `dKV` and `dScore` are bitwise run-to-run deterministic.
+`dAPE` is reduced with one fp32 atomic per `(k, dim)` per CTA and is **not** bitwise
+run-to-run deterministic; the backward APIs raise under
+`torch.use_deterministic_algorithms(True)`.
+
+### Support surface (`check_support`)
+
+- Compute capability **10.0** (the only validated architecture so far; the kernels use
+  no arch-specific features, wider enablement is possible after validation)
+- `ratio == 4`, `coff == 2` (the production CSA/HCA configuration; the kernels are
+  generic over `(ratio, head_dim, coff in {1, 2})` and the gate can be lifted once
+  validated)
+- BF16 `kv` / `score` / `out`, FP32 `ape`, int32 `cu_seqlens` / `cu_seqlens_comp`
+- int32 flat offsets: `total_tokens * coff * head_dim < 2**31` and
+  `total_comp * head_dim < 2**31`
+- `total_comp > 0` requires `total_tokens >= ratio` (each compressed row gathers a
+  window of `ratio` tokens)
+- `head_dim <= 8388480` (forward launch `gridDim.y` bound)
+- 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
+
+## Installation
+
+```bash
+pip install nvidia-cudnn-frontend[cutedsl]
+```
+
+## API Usage
+
+### High-level wrappers
+
+```python
+from cudnn import CSA
+
+# forward: (total_tokens, coff*head_dim) BF16 kv/score, (ratio, coff*head_dim) FP32 ape,
+# (B+1,) int32 cu_seqlens / cu_seqlens_comp
+result = CSA.csa_compressor_forward_wrapper(
+    kv, score, ape, cu_seqlens, cu_seqlens_comp,
+    ratio=4, head_dim=128, coff=2,
+    total_comp=None,  # defaults to cu_seqlens_comp[-1] (synchronizes); pass a static
+                      # capacity explicitly to stay CUDA-graph capture-safe
+    stream=None,
+)
+out = result["out"]  # (total_comp, head_dim) BF16
+
+grads = CSA.csa_compressor_backward_wrapper(
+    kv, score, ape, cu_seqlens, cu_seqlens_comp, grad_out,
+    ratio=4, head_dim=128, coff=2, stream=None,
+)
+grad_kv, grad_score, grad_ape = grads  # BF16, BF16, FP32
+```
+
+The wrappers cache compiled API instances; the underlying JIT is shared per
+`(ratio, head_dim, coff, device)`, so runtime shape changes never recompile.
+
+### Class API
+
+```python
+from cudnn import CSACompressorForward, CSACompressorBackward
+
+op = CSACompressorForward(
+    sample_kv, sample_score, sample_ape, sample_cu_seqlens, sample_cu_seqlens_comp,
+    sample_out, ratio=4, coff=2,
+)
+op.check_support()
+op.compile()
+op.execute(kv, score, ape, cu_seqlens, cu_seqlens_comp, out, current_stream=None)
+```
+
+`CSACompressorBackward.execute` additionally takes `grad_out` and the
+`grad_kv` / `grad_score` / `grad_ape` buffers. `grad_kv` / `grad_score` may be
+**uninitialized**: the kernel writes every position (disjoint, atomic-free stores;
+never-consumed positions — segment tails, the last block's first-half columns, segments
+shorter than `ratio`, token-capacity padding beyond `cu_seqlens[-1]` — get exact zeros
+from their unique owning CTA, matching autograd). When `total_comp == 0` the kernel is
+not launched and the buffers are left untouched (zero them yourself if you need
+autograd's exact zeros; the high-level wrapper does). `grad_ape` must be
+**zero-initialized** (fp32 atomic accumulation).
+
+### CUDA graphs
+
+The launch path is capture-compatible once the kernels for a `(ratio, head_dim, coff)`
+configuration are compiled: run one warmup call (or `compile()`) per configuration
+before capturing, and pass `total_comp` explicitly. A call that would JIT under capture
+raises a `RuntimeError` instead of corrupting the capture.
+
+### Environment variables
+
+- `CUDNNFE_CSA_COMPRESSOR_FAST_LAUNCH=0` — disable the cached-launch host optimization
+  (a per-config snapshot of the CuTe-DSL launch state, replayed with in-place argument
+  mutation; it removes tens of microseconds of per-call host overhead for these
+  microsecond-scale kernels). The snapshot construction introspects
+  private-but-stable DSL internals; on any structural mismatch (e.g. a future
+  `nvidia-cutlass-dsl` upgrade) it falls back to the regular launch path automatically.
+
+## Performance
+
+Measured on 1x B200 (CC 10.0, driver 590.48.01); BF16 `kv`/`score`, FP32 `ape`; `ratio = 4`, `coff = 2`;
+THD packs of 8192-token sequences; eager baseline = the exact replaced region of
+Megatron-LM `Compressor._forward_thd` on identical inputs.
+
+*Isolated GPU kernel time* (nsys, sum of kernel durations per iteration, 50 iterations
+after 20 warmup; no launch/host overhead; backward includes its `dAPE` zero-fill):
+
+| THD pack | head_dim | eager fwd | fused fwd | fwd | eager bwd | fused bwd | bwd |
+|---|---|---|---|---|---|---|---|
+| 1 x 8192 | 128 | 117.8 us | 4.5 us | **26.5x** | 187.2 us | 12.8 us | **14.6x** |
+| 3 x 8192 | 128 | 229.8 us | 10.0 us | **23.0x** | 352.7 us | 22.2 us | **15.9x** |
+| 1 x 8192 | 512 | 263.3 us | 12.4 us | **21.2x** | 425.0 us | 22.8 us | **18.6x** |
+| 3 x 8192 | 512 | 664.3 us | 35.0 us | **19.0x** | 1155.8 us | 66.0 us | **17.5x** |
+
+*End-to-end wall clock of the same region* (CUDA events, median of 100; includes launch
+overhead; eager backward goes through torch autograd, fused backward is the explicit
+backward wrapper call — see the measurement-basis note above; not comparable to the
+kernel-time numbers):
+
+| THD pack | head_dim | eager fwd | fused fwd | fwd | eager bwd | fused bwd | bwd |
+|---|---|---|---|---|---|---|---|
+| 1 x 8192 | 128 | 333.7 us | 37.8 us | **8.8x** | 506.7 us | 47.2 us | **10.7x** |
+| 3 x 8192 | 128 | 373.9 us | 35.1 us | **10.7x** | 561.9 us | 54.6 us | **10.3x** |
+| 1 x 8192 | 512 | 409.2 us | 39.7 us | **10.3x** | 646.3 us | 57.3 us | **11.3x** |
+| 3 x 8192 | 512 | 823.7 us | 61.4 us | **13.4x** | 1475.5 us | 101.8 us | **14.5x** |
+
+Environment: driver 590.48.01, PyTorch 2.13.0 (CUDA 13.3), `nvidia-cutlass-dsl` 4.6.1.
+Measurement basis: identical inputs over exactly the replaced region for both
+implementations; eager backward = torch autograd of the recorded eager graph; fused
+backward = the backward wrapper (kernel + the fp32 `dAPE` zero-fill + host validation,
+no autograd engine).
+
+An ncu hardware-ceiling audit of the **ported kernels (prior to the two optimization
+commits)** — cache-flushed, `--set full`, all four benchmark shapes — showed measured
+DRAM read volume matching the algorithmically necessary bytes within 1% (the THD gather
+adds no over-read; stores fully coalesced at 32/32 bytes per sector, loads 29-30/32),
+with neither L2 (<27% of peak) nor DRAM (<33%) 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 optimizations that audit identified are folded in here: 32-bit
+vectorized forward accesses and backward kernel-side zero-writes replacing the two bf16
+grad-buffer fills. They do not change the bytes the kernels must read, so the
+traffic-optimality conclusion carries over; the utilization percentages above predate
+them.
+
+## Testing
+
+```bash
+(cd test/python && pytest fe_api/csa/test_CSA_compressor.py)
+```
+
+The tests validate numerics against an fp32-intermediate eager reference (bitwise
+`dKV`/`dScore`), the upstream eager numerics, and an fp64 oracle, plus ragged packs,
+static-capacity padding, kernel-side zero-writes into uninitialized gradient buffers
+(NaN-canary, including the `total_comp == 0` zeros fallback), run-to-run determinism,
+CUDA-graph capture/replay, and `check_support` boundaries.

docs/fe-oss-apis/overview.md

@@ -22,6 +22,7 @@ This folder documents the Python FE APIs implemented under `python/cudnn`. For d
 - [Grouped GEMM + Wgrad](gemm_fusions/grouped_gemm_wgrad.md)
 - [Block Sparse Attention (BSA)](bsa.md)
 - [Native Sparse Attention (NSA)](nsa.md)
+- [CSA Fused Compressor](csa.md)
 - [RMSNorm + RHT + Amax](rmsnorm_rht_amax.md)
 - [SDPA Forward FE OSS API (SM100, D=256)](https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/operations/Attention.html#sdpa-forward-fe-oss-sm100-d256)
 - [SDPA Backward FE OSS API (SM100, D=256)](https://docs.nvidia.com/deeplearning/cudnn/frontend/latest/operations/Attention.html#sdpa-backward-fe-oss-sm100-d256)

python/cudnn/__init__.py

@@ -282,6 +282,11 @@ def _dlopen_cudnn():
     "block_sparse_attention_forward": (".block_sparse_attention", "block_sparse_attention_forward"),
     "block_sparse_attention_backward": (".block_sparse_attention", "block_sparse_attention_backward"),
     "DSA": (".deepseek_sparse_attention", "DSA"),
+    "CSA": (".csa", "CSA"),
+    "CSACompressorForward": (".csa", "CSACompressorForward"),
+    "CSACompressorBackward": (".csa", "CSACompressorBackward"),
+    "csa_compressor_forward_wrapper": (".csa", "csa_compressor_forward_wrapper"),
+    "csa_compressor_backward_wrapper": (".csa", "csa_compressor_backward_wrapper"),
     "NSA": (".native_sparse_attention", "NSA"),
     "GemmSwigluSm100": (".gemm_swiglu", "GemmSwigluSm100"),
     "gemm_swiglu_wrapper_sm100": (".gemm_swiglu", "gemm_swiglu_wrapper_sm100"),

python/cudnn/csa/README.md

@@ -0,0 +1,18 @@
+# CSA module
+
+Fused CuTe-DSL kernels for the CSA/HCA experimental attention variants (the components
+that are not shared with the DSA module, which lives in
+`python/cudnn/deepseek_sparse_attention/`).
+
+- **Compressor**: fused forward+backward kernels for the `Compressor` gated-softmax
+  pooling (THD packed layout): gather -> `+ APE` -> overlap-window transform -> fp32
+  softmax -> gated weighted sum -> bf16 cast, as one kernel per direction. Ported from
+  Megatron-LM ([PR #5984](https://github.com/NVIDIA/Megatron-LM/pull/5984), measurements
+  in [issue #5968](https://github.com/NVIDIA/Megatron-LM/issues/5968)). See
+  [docs/fe-oss-apis/csa.md](../../../docs/fe-oss-apis/csa.md).
+
+## Acknowledgements
+
+The fused Compressor kernels were contributed by the GLM training-performance team
+(Zhipu AI). The CSA/HCA attention variants and the surrounding DSA/CSA kernel family are
+by Hongxiao Bai, Jiayu Sun and Jie Fang.

python/cudnn/csa/__init__.py

@@ -0,0 +1,42 @@
+from importlib import import_module
+
+_SYMBOLS = {
+    "CSACompressorForward": (".compressor", "CSACompressorForward"),
+    "CSACompressorBackward": (".compressor", "CSACompressorBackward"),
+    "csa_compressor_forward_wrapper": (".compressor", "csa_compressor_forward_wrapper"),
+    "csa_compressor_backward_wrapper": (".compressor", "csa_compressor_backward_wrapper"),
+}
+
+
+def _load_symbol(name):
+    module_name, symbol_name = _SYMBOLS[name]
+    module = import_module(module_name, package=__name__)
+    symbol = getattr(module, symbol_name)
+    globals()[name] = symbol
+    return symbol
+
+
+def __getattr__(name):
+    if name == "CSA":
+        return CSA
+    if name in _SYMBOLS:
+        return _load_symbol(name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+class CSANamespace:
+    def __getattr__(self, name):
+        if name in _SYMBOLS:
+            return _load_symbol(name)
+        raise AttributeError(f"CSA has no attribute {name!r}")
+
+
+CSA = CSANamespace()
+
+__all__ = [
+    "CSA",
+    "CSACompressorBackward",
+    "CSACompressorForward",
+    "csa_compressor_backward_wrapper",
+    "csa_compressor_forward_wrapper",
+]

python/cudnn/csa/compressor/__init__.py

@@ -0,0 +1,13 @@
+from .api import (
+    CSACompressorForward,
+    CSACompressorBackward,
+    csa_compressor_forward_wrapper,
+    csa_compressor_backward_wrapper,
+)
+
+__all__ = [
+    "CSACompressorBackward",
+    "CSACompressorForward",
+    "csa_compressor_backward_wrapper",
+    "csa_compressor_forward_wrapper",
+]