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

benchmark/csa/reg_probe_csa_compressor_r128.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+"""ptxas register/spill probe for every shipped ratio=128 CSA compressor kernel.
+
+JIT-compiles the full shipped dispatch envelope — every (config, schedule) pair the
+``nb_total`` bucket tables can select, forward and backward, over coff {1, 2} x
+head_dim {128, 512} (16 kernels) — then runs ``ptxas -v`` on each kernel's PTX and
+prints a table of registers / spill bytes / stack bytes / ex2.approx count. This
+reproduces the register table published in docs/fe-oss-apis/csa.md.
+
+Exits nonzero if any kernel spills, uses stack, or fails ptxas.
+
+Requires a CC 10.0 GPU (the JIT needs a device), ``ptxas`` on PATH, and the
+``cudnn[cutedsl]`` install. Not collected by pytest. Run, e.g.::
+
+    CUDA_VISIBLE_DEVICES=0 python benchmark/csa/reg_probe_csa_compressor_r128.py
+"""
+
+import argparse
+import os
+import re
+import shutil
+import subprocess
+import sys
+import tempfile
+
+os.environ.setdefault("CUTE_DSL_KEEP", "ptx")  # keep PTX artifacts on the compiled handles
+
+import torch  # noqa: E402
+
+
+def ptxas_one(tag, ptx_text, arch, out_dir):
+    path = os.path.join(out_dir, f"{tag}.ptx")
+    with open(path, "w") as f:
+        f.write(ptx_text)
+    r = subprocess.run(["ptxas", "-v", f"-arch={arch}", "-o", os.devnull, path], capture_output=True, text=True)
+    out = r.stderr + r.stdout
+    regs = re.search(r"Used (\d+) registers", out)
+    spill = re.search(r"(\d+) bytes spill stores, (\d+) bytes spill loads", out)
+    stack = re.search(r"(\d+) bytes stack frame", out)
+    n_ex2 = ptx_text.count("ex2.approx")
+    print(
+        f"{tag:36} regs={regs.group(1) if regs else '?':>3} "
+        f"spill={(spill.group(1) + '/' + spill.group(2)) if spill else '?'} "
+        f"stack={stack.group(1) if stack else '?'} ex2.approx={n_ex2:3d} rc={r.returncode}",
+        flush=True,
+    )
+    clean = r.returncode == 0 and spill is not None and stack is not None and spill.group(1) == spill.group(2) == "0" and stack.group(1) == "0"
+    return clean
+
+
+def main():
+    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    ap.add_argument("--arch", default="sm_100a", help="ptxas target architecture (default: sm_100a)")
+    ap.add_argument("--keep-ptx", default=None, help="directory to keep the per-kernel .ptx files in (default: temporary)")
+    args = ap.parse_args()
+    # CUTE_DSL_KEEP=ptx makes the DSL drop each kernel's PTX into the current
+    # directory; run the compiles from the (possibly temporary) output directory so
+    # the repository tree stays clean.
+    out_dir = os.path.abspath(args.keep_ptx) if args.keep_ptx else tempfile.mkdtemp(prefix="csa_r128_ptx_")
+    os.makedirs(out_dir, exist_ok=True)
+    os.chdir(out_dir)
+    from cudnn.csa.compressor import compressor_sm100_r128 as M
+
+    assert torch.cuda.is_available() and torch.cuda.get_device_capability() == (10, 0), "requires a CC 10.0 GPU"
+    dev = torch.device("cuda", torch.cuda.current_device())
+    # Compile every schedule bucket each shipped config can select at runtime
+    # (precompile with nb_total=None walks the small/default/large tables).
+    for coff, d in [(1, 128), (2, 128), (1, 512), (2, 512)]:
+        M.precompile_fwd_r128(128, d, coff, dev)
+        M.precompile_bwd_r128(128, d, coff, dev)
+
+    all_clean = True
+    n = 0
+    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)
+    if args.keep_ptx:
+        print(f"PTX kept in {out_dir}", flush=True)
+    else:
+        os.chdir(tempfile.gettempdir())
+        shutil.rmtree(out_dir, ignore_errors=True)
+    sys.exit(0 if all_clean else 1)
+
+
+if __name__ == "__main__":
+    main()

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,19 @@
+# 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` -> optional overlap-window transform
+  (`coff == 2`) -> 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,53 @@
+"""``cudnn.csa`` — CuTe-DSL kernels for the CSA/HCA experimental attention variants.
+
+Symbols (the fused ``Compressor`` APIs) resolve lazily on first attribute access, so
+importing ``cudnn`` never pulls in the optional ``[cutedsl]`` dependency stack.
+"""
+
+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):
+    """Import the symbol behind lazy attribute ``name`` and cache it in module globals."""
+    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):
+    """Resolve the lazily exported symbols and the ``CSA`` namespace (PEP 562)."""
+    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:
+    """Namespace object mirroring the package's lazy symbols (``cudnn.CSA.<symbol>``)."""
+
+    def __getattr__(self, name):
+        """Lazily resolve ``CSA.<name>`` through the package's symbol table."""
+        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,15 @@
+"""Public surface for the fused CSA/HCA Compressor kernels (re-exports from ``.api``)."""
+
+from .api import (
+    CSACompressorForward,
+    CSACompressorBackward,
+    csa_compressor_forward_wrapper,
+    csa_compressor_backward_wrapper,
+)
+
+__all__ = [
+    "CSACompressorBackward",
+    "CSACompressorForward",
+    "csa_compressor_backward_wrapper",
+    "csa_compressor_forward_wrapper",
+]