diff --git a/docs/fe-oss-apis/csa.md b/docs/fe-oss-apis/csa.md new file mode 100644 index 000000000..1e08ac0f7 --- /dev/null +++ b/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. diff --git a/docs/fe-oss-apis/overview.md b/docs/fe-oss-apis/overview.md index 2bc3ccfb2..3c25c11cd 100644 --- a/docs/fe-oss-apis/overview.md +++ b/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) diff --git a/python/cudnn/__init__.py b/python/cudnn/__init__.py index 6e8b99829..a5256c866 100644 --- a/python/cudnn/__init__.py +++ b/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"), diff --git a/python/cudnn/csa/README.md b/python/cudnn/csa/README.md new file mode 100644 index 000000000..e42de24b3 --- /dev/null +++ b/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. diff --git a/python/cudnn/csa/__init__.py b/python/cudnn/csa/__init__.py new file mode 100644 index 000000000..4319adadb --- /dev/null +++ b/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", +] diff --git a/python/cudnn/csa/compressor/__init__.py b/python/cudnn/csa/compressor/__init__.py new file mode 100644 index 000000000..60ce1f9dc --- /dev/null +++ b/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", +] diff --git a/python/cudnn/csa/compressor/api.py b/python/cudnn/csa/compressor/api.py new file mode 100644 index 000000000..139e58a97 --- /dev/null +++ b/python/cudnn/csa/compressor/api.py @@ -0,0 +1,580 @@ +"""APIBase wrappers for the fused CSA/HCA Compressor gated-pooling CuTe-DSL kernels. + +``CSACompressorForward`` / ``CSACompressorBackward`` wrap the forward and backward +kernels in ``compressor_sm100.py`` (ported from Megatron-LM, see +https://github.com/NVIDIA/Megatron-LM/pull/5984 and +https://github.com/NVIDIA/Megatron-LM/issues/5968). The kernels fuse the gated-softmax +pooling region of the CSA/HCA ``Compressor`` for the THD packed layout: + + 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)]) + +over the per-block window ``w`` (``2 * ratio`` entries for the overlapping ``coff == 2`` +form; the previous block's half-window is invalid for each segment's first block). The +framework-side autograd wiring stays in the caller (e.g. a ``torch.autograd.Function`` +that calls the forward wrapper in ``forward()`` and the backward wrapper in +``backward()``); these APIs are pure kernels-plus-validation. + +Validated envelope (``check_support``): compute capability 10.0, ``ratio == 4``, +``coff == 2`` (the production CSA/HCA configuration), BF16 ``kv``/``score``/``out``, +FP32 ``ape``, int32 ``cu_seqlens``/``cu_seqlens_comp``, and int32 flat offsets +(``total_tokens * coff * head_dim < 2**31``). The kernels themselves are generic over +``(ratio, head_dim, coff in {1, 2})``; wider gates can be lifted once validated. + +Numerics contract (see ``compressor_sm100.py`` for details): fp32 arithmetic with one +final bf16 rounding, ``mul.rn``/``fma.rn`` pinned in PTX. Forward, ``dKV`` and ``dScore`` +are bitwise run-to-run deterministic; ``dAPE`` uses one fp32 atomic per ``(k, dim)`` per +CTA and is not (the backward APIs refuse to run under +``torch.use_deterministic_algorithms(True)``). +""" + +from __future__ import annotations + +import threading +import warnings +from collections import OrderedDict +from contextlib import contextmanager +from typing import Iterator, Optional + +import torch +import cuda.bindings.driver as cuda + +from cudnn.api_base import APIBase, TupleDict + +from .compressor_sm100 import ( + CU_ALIGN_BYTES, + PTR_ALIGN_BYTES, + SUPPORTED_COMPUTE_CAPABILITY, + precompile_bwd, + precompile_fwd, + run_bwd, + run_fwd, +) + +# int32 flat offsets: every element offset the kernels compute must fit in int32. +_INT32_LIMIT = 2**31 +# Forward launch schedule gridDim.y bound: at 128 threads per column CTA the largest +# launchable head_dim is 128 * 65535 (identical for the 64-thread vec == 2 path, which +# halves the column count). +_MAX_HEAD_DIM = 128 * 65535 +# Bound + eviction policy follow python/cudnn/graph.py's graph_cache precedent. +_API_CACHE_MAXSIZE = 256 + + +class _LruDict: + """Bounded thread-safe LRU mapping (per ``graph.py``'s ``graph_cache`` precedent).""" + + def __init__(self, maxsize: int = _API_CACHE_MAXSIZE): + self._data: OrderedDict = OrderedDict() + self._lock = threading.Lock() + self._maxsize = maxsize + + def get(self, key, default=None): + with self._lock: + if key not in self._data: + return default + self._data.move_to_end(key) + return self._data[key] + + def put(self, key, value) -> None: + with self._lock: + self._data[key] = value + self._data.move_to_end(key) + while len(self._data) > self._maxsize: + self._data.popitem(last=False) + + +def _resolve_stream_handle(current_stream: Optional[cuda.CUstream]) -> Optional[int]: + """Integer stream handle for the launch path (None -> torch current stream).""" + if current_stream is None: + return None + return int(current_stream) + + +@contextmanager +def _torch_stream_context(current_stream: Optional[cuda.CUstream], device: torch.device) -> Iterator[None]: + """Run torch work on ``current_stream`` (device-tagged) when one is given.""" + if current_stream is None: + yield + return + with torch.cuda.stream(torch.cuda.get_stream_from_external(int(current_stream), device)): + yield + + +def _reject_deterministic_backward() -> None: + """The backward accumulates ``dAPE`` with fp32 atomics and is not deterministic. + + Mirrors torch's deterministic-mode semantics: strict mode raises, warn-only mode + warns and runs. + """ + if torch.are_deterministic_algorithms_enabled(): + message = ( + "CSA compressor backward accumulates dAPE with fp32 atomics and is not " + "deterministic; torch.use_deterministic_algorithms(True) is set. Use a " + "deterministic (eager) implementation instead." + ) + if torch.is_deterministic_algorithms_warn_only_enabled(): + warnings.warn(message, RuntimeWarning, stacklevel=2) + else: + raise RuntimeError(message) + + +class _CSACompressorBase(APIBase): + """Shared descriptor plumbing and ``check_support`` for forward and backward.""" + + def __init__( + self, + sample_kv: torch.Tensor, # (total_tokens, coff * head_dim) BF16 + sample_score: torch.Tensor, # (total_tokens, coff * head_dim) BF16 + sample_ape: torch.Tensor, # (ratio, coff * head_dim) FP32 + sample_cu_seqlens: torch.Tensor, # (B + 1,) INT32 token offsets + sample_cu_seqlens_comp: torch.Tensor, # (B + 1,) INT32 compressed-block offsets + sample_out: torch.Tensor, # (total_comp, head_dim) BF16 (forward output / backward grad_out) + ratio: int = 4, + coff: int = 2, + ): + super().__init__() + self._warn_experimental_api() + + self.kv_desc = self._make_tensor_desc(sample_kv, name="sample_kv") + self.score_desc = self._make_tensor_desc(sample_score, name="sample_score") + self.ape_desc = self._make_tensor_desc(sample_ape, name="sample_ape") + self.cu_desc = self._make_tensor_desc(sample_cu_seqlens, name="sample_cu_seqlens") + self.cuc_desc = self._make_tensor_desc(sample_cu_seqlens_comp, name="sample_cu_seqlens_comp") + self.out_desc = self._make_tensor_desc(sample_out, name="sample_out") + + self.ratio = int(ratio) + self.coff = int(coff) + + self.total_tokens = None + self.total_comp = None + self.head_dim = None + self.n_seg = None + self.target_device: Optional[torch.device] = None + + def check_support(self) -> bool: + """Validate the configuration against the kernels' validated envelope. + + Malformed inputs and configurations outside the envelope raise ``ValueError`` + (device-capability failures raise ``RuntimeError``), mirroring the other FE-OSS + APIs; there is no soft fallback path inside this API. + """ + self._logger.debug("Entering check_support") + self._value_error_if( + self.ratio != 4 or self.coff != 2, + f"CSA compressor is validated for ratio=4, coff=2 only (the production CSA/HCA configuration), got ratio={self.ratio}, coff={self.coff}", + ) + self._value_error_if( + self.kv_desc.ndim != 2, + f"kv must be 2-D (total_tokens, coff * head_dim), got {self.kv_desc.shape}", + ) + self._value_error_if( + self.out_desc.ndim != 2, + f"out/grad_out must be 2-D (total_comp, head_dim), got {self.out_desc.shape}", + ) + total_tokens, width = self.kv_desc.shape + total_comp, head_dim = self.out_desc.shape + self._value_error_if( + head_dim < 1 or width != self.coff * head_dim, + f"kv width must equal coff * head_dim = {self.coff} * {head_dim}, got {width}", + ) + self._value_error_if( + self.score_desc.shape != self.kv_desc.shape, + f"score shape {self.score_desc.shape} != kv shape {self.kv_desc.shape}", + ) + self._value_error_if( + self.ape_desc.shape != (self.ratio, width), + f"ape must be (ratio, coff * head_dim) = ({self.ratio}, {width}), got {self.ape_desc.shape}", + ) + self._value_error_if( + self.cu_desc.ndim != 1 or self.cuc_desc.ndim != 1, + "cu_seqlens and cu_seqlens_comp must be 1-D", + ) + self._value_error_if( + self.cu_desc.shape != self.cuc_desc.shape or self.cu_desc.shape[0] < 2, + f"cu_seqlens and cu_seqlens_comp must both have B + 1 >= 2 entries, got {self.cu_desc.shape} and {self.cuc_desc.shape}", + ) + + self._check_dtype(self.kv_desc, torch.bfloat16, name="kv") + self._check_dtype(self.score_desc, torch.bfloat16, name="score") + self._check_dtype(self.ape_desc, torch.float32, name="ape") + self._check_dtype(self.cu_desc, torch.int32, name="cu_seqlens") + self._check_dtype(self.cuc_desc, torch.int32, name="cu_seqlens_comp") + self._check_dtype(self.out_desc, torch.bfloat16, name="out/grad_out") + + # int32 flat offsets: the kernels index flat views with int32 arithmetic. + self._value_error_if( + total_tokens * width >= _INT32_LIMIT, + f"total_tokens * coff * head_dim must be < 2**31 for int32 flat offsets, got {total_tokens} * {width}", + ) + self._value_error_if( + total_comp * head_dim >= _INT32_LIMIT, + f"total_comp * head_dim must be < 2**31 for int32 flat offsets, got {total_comp} * {head_dim}", + ) + # gridDim.y bound of the forward launch schedule (64/128-thread column groups): + # head_dims beyond this cannot be launched (the pre-vectorization schedule had + # the same 128 * 65535 envelope, just unchecked). + self._value_error_if( + head_dim > _MAX_HEAD_DIM, + f"head_dim must be <= {_MAX_HEAD_DIM} (forward launch gridDim.y bound), got {head_dim}", + ) + # Rows (including static-capacity padding rows) gather a window of `ratio` + # tokens; the eager gather has the same requirement. + self._value_error_if( + total_comp > 0 and total_tokens < self.ratio, + f"total_comp={total_comp} > 0 requires at least ratio={self.ratio} tokens, got {total_tokens}", + ) + for desc, name in ( + (self.kv_desc, "kv"), + (self.score_desc, "score"), + (self.ape_desc, "ape"), + (self.cu_desc, "cu_seqlens"), + (self.cuc_desc, "cu_seqlens_comp"), + (self.out_desc, "out/grad_out"), + ): + self._value_error_if(not desc.is_contiguous(), f"{name} must be contiguous") + + # Device resolution: all runtime (CUDA) descriptors on one device; meta + # descriptors are metadata-only stand-ins and pin nothing. + all_descs = (self.kv_desc, self.score_desc, self.ape_desc, self.cu_desc, self.cuc_desc, self.out_desc) + devices = {desc.device for desc in all_descs} + self._value_error_if( + any(dev.type not in ("cuda", "meta") for dev in devices), + f"all tensors must be CUDA tensors, got devices {sorted(str(dev) for dev in devices)}", + ) + cuda_devices = {dev for dev in devices if dev.type == "cuda"} + self._value_error_if( + len(cuda_devices) > 1, + f"all tensors must share one CUDA device, got {sorted(str(dev) for dev in cuda_devices)}", + ) + self._runtime_error_if(not torch.cuda.is_available(), "CSA compressor requires CUDA") + if cuda_devices: + target = next(iter(cuda_devices)) + if target.index is None: + target = torch.device("cuda", torch.cuda.current_device()) + else: + target = torch.device("cuda", torch.cuda.current_device()) + + capability = torch.cuda.get_device_capability(target) + self._runtime_error_if( + capability != SUPPORTED_COMPUTE_CAPABILITY, + f"CSA compressor requires compute capability {SUPPORTED_COMPUTE_CAPABILITY} (the only validated architecture so far), found SM{capability[0]}.{capability[1]} on {target}", + ) + + self.total_tokens = total_tokens + self.total_comp = total_comp + self.head_dim = head_dim + self.n_seg = self.cu_desc.shape[0] - 1 + self.target_device = target + self._is_supported = True + return True + + def _validate_runtime_tensor(self, tensor, name, shape, dtype, device, align): + """Cheap per-call validation of one runtime tensor.""" + if tuple(tensor.shape) != shape: + raise ValueError(f"{name} must have shape {shape}, got {tuple(tensor.shape)}") + if tensor.dtype != dtype: + raise ValueError(f"{name} must have dtype {dtype}, got {tensor.dtype}") + if not tensor.is_cuda or tensor.device != device: + raise ValueError(f"{name} must be a CUDA tensor on {device}, got {tensor.device}") + if not tensor.is_contiguous(): + raise ValueError(f"{name} must be contiguous, got stride {tuple(tensor.stride())}") + # Contiguity does not imply base-pointer alignment (storage-offset views); the + # kernels' pointer wrappers assume it. + if tensor.data_ptr() % align: + raise ValueError(f"{name} base pointer must be {align}-byte aligned, got 0x{tensor.data_ptr():x}") + + @staticmethod + def _record_streams(tensors, current_stream: Optional[cuda.CUstream], device: torch.device) -> None: + """Keep tensor storages alive for work enqueued on an explicit external stream. + + The launch path takes raw pointers, so PyTorch's caching allocator does not know + the kernel on ``current_stream`` still reads/writes these tensors: without + ``record_stream`` it may recycle a storage freed by the caller while the kernel + is pending. Only needed for explicit streams — with ``current_stream=None`` the + launch lands on torch's current stream and ordinary stream semantics apply. + """ + if current_stream is None: + return + consumer = torch.cuda.get_stream_from_external(int(current_stream), device) + for t in tensors: + t.record_stream(consumer) + + +class CSACompressorForward(_CSACompressorBase): + """Fused CSA compressor forward: one kernel over the whole THD pack. + + Rows in ``[cu_seqlens_comp[-1], total_comp)`` are static-capacity padding (for + CUDA-graph static shapes) and are computed with first-in-segment semantics from + token 0, exactly like the eager gather. + """ + + def compile(self) -> None: + self._logger.debug("Entering compile") + self._ensure_support_checked() + if self._compiled_kernel is not None: + return + precompile_fwd(self.ratio, self.head_dim, self.coff, self.target_device) + + ratio, head_dim, coff = self.ratio, self.head_dim, self.coff + + def tensor_api(kv, score, ape, cu_seqlens, cu_seqlens_comp, out, stream_handle): + run_fwd(kv, score, ape, cu_seqlens, cu_seqlens_comp, out, out.shape[0], ratio, head_dim, coff, stream_handle=stream_handle) + + self._compiled_kernel = tensor_api + self._logger.debug("Kernel compiled successfully") + + def execute( + self, + kv: torch.Tensor, # (total_tokens, coff * head_dim) BF16, contiguous + score: torch.Tensor, # (total_tokens, coff * head_dim) BF16, contiguous + ape: torch.Tensor, # (ratio, coff * head_dim) FP32, contiguous + cu_seqlens: torch.Tensor, # (B + 1,) INT32, contiguous + cu_seqlens_comp: torch.Tensor, # (B + 1,) INT32, contiguous + out: torch.Tensor, # (total_comp, head_dim) BF16, contiguous + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + """Run the compiled forward kernel; ``out`` is fully overwritten.""" + self._logger.debug("Entering execute") + if self._compiled_kernel is None: + raise ValueError("CSACompressorForward kernel not compiled") + device = self.target_device + width = self.coff * self.head_dim + self._validate_runtime_tensor(kv, "kv", (self.total_tokens, width), torch.bfloat16, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(score, "score", (self.total_tokens, width), torch.bfloat16, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(ape, "ape", (self.ratio, width), torch.float32, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(cu_seqlens, "cu_seqlens", (self.n_seg + 1,), torch.int32, device, CU_ALIGN_BYTES) + self._validate_runtime_tensor(cu_seqlens_comp, "cu_seqlens_comp", (self.n_seg + 1,), torch.int32, device, CU_ALIGN_BYTES) + self._validate_runtime_tensor(out, "out", (self.total_comp, self.head_dim), torch.bfloat16, device, PTR_ALIGN_BYTES) + if out.numel() == 0: + return + self._compiled_kernel(kv, score, ape, cu_seqlens, cu_seqlens_comp, out, _resolve_stream_handle(current_stream)) + self._record_streams((kv, score, ape, cu_seqlens, cu_seqlens_comp, out), current_stream, device) + + +class CSACompressorBackward(_CSACompressorBase): + """Fused CSA compressor backward: recompute window probs, write grads in one kernel. + + ``grad_kv``/``grad_score`` may be UNINITIALIZED: the kernel writes every position — + consumed positions get their gradients (disjoint, atomic-free stores), and every + never-consumed position (segment-tail tokens; for ``coff == 2`` the first-half + columns of each segment's last block; tokens of segments shorter than ``ratio``; + tokens beyond ``cu_seqlens[-1]`` when the buffers carry static token-capacity + padding) gets an exact zero from its unique owning CTA — matching autograd without + separate zero-fill kernels. Exception: when ``total_comp == 0`` the kernel is not + launched and the buffers are left untouched, so a caller that needs autograd-exact + zeros in that case must zero them itself (the high-level wrapper does). + ``grad_ape`` must always be zero-initialized: it is accumulated with fp32 atomics + and is not bitwise run-to-run deterministic (``grad_kv``/``grad_score`` are); + ``execute`` raises under ``torch.use_deterministic_algorithms(True)``. Incoming + gradients on static-capacity padding rows (``[cu_seqlens_comp[-1], total_comp)``) + are ignored. + """ + + def compile(self) -> None: + self._logger.debug("Entering compile") + self._ensure_support_checked() + if self._compiled_kernel is not None: + return + precompile_bwd(self.ratio, self.head_dim, self.coff, self.target_device) + + ratio, head_dim, coff = self.ratio, self.head_dim, self.coff + + def tensor_api(kv, score, ape, cu_seqlens, cu_seqlens_comp, grad_out, grad_kv, grad_score, grad_ape, stream_handle): + run_bwd( + kv, + score, + ape, + cu_seqlens, + cu_seqlens_comp, + grad_out, + grad_kv, + grad_score, + grad_ape, + grad_out.shape[0], + ratio, + head_dim, + coff, + stream_handle=stream_handle, + ) + + self._compiled_kernel = tensor_api + self._logger.debug("Kernel compiled successfully") + + def execute( + self, + kv: torch.Tensor, # (total_tokens, coff * head_dim) BF16, contiguous + score: torch.Tensor, # (total_tokens, coff * head_dim) BF16, contiguous + ape: torch.Tensor, # (ratio, coff * head_dim) FP32, contiguous + cu_seqlens: torch.Tensor, # (B + 1,) INT32, contiguous + cu_seqlens_comp: torch.Tensor, # (B + 1,) INT32, contiguous + grad_out: torch.Tensor, # (total_comp, head_dim) BF16, contiguous + grad_kv: torch.Tensor, # (total_tokens, coff * head_dim) BF16 (may be uninitialized; fully written when total_comp > 0) + grad_score: torch.Tensor, # (total_tokens, coff * head_dim) BF16 (may be uninitialized; fully written when total_comp > 0) + grad_ape: torch.Tensor, # (ratio, coff * head_dim) FP32, ZERO-INITIALIZED + current_stream: Optional[cuda.CUstream] = None, + ) -> None: + """Run the compiled backward kernel into the gradient buffers (see class docs).""" + self._logger.debug("Entering execute") + if self._compiled_kernel is None: + raise ValueError("CSACompressorBackward kernel not compiled") + _reject_deterministic_backward() + device = self.target_device + width = self.coff * self.head_dim + self._validate_runtime_tensor(kv, "kv", (self.total_tokens, width), torch.bfloat16, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(score, "score", (self.total_tokens, width), torch.bfloat16, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(ape, "ape", (self.ratio, width), torch.float32, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(cu_seqlens, "cu_seqlens", (self.n_seg + 1,), torch.int32, device, CU_ALIGN_BYTES) + self._validate_runtime_tensor(cu_seqlens_comp, "cu_seqlens_comp", (self.n_seg + 1,), torch.int32, device, CU_ALIGN_BYTES) + self._validate_runtime_tensor(grad_out, "grad_out", (self.total_comp, self.head_dim), torch.bfloat16, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(grad_kv, "grad_kv", (self.total_tokens, width), torch.bfloat16, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(grad_score, "grad_score", (self.total_tokens, width), torch.bfloat16, device, PTR_ALIGN_BYTES) + self._validate_runtime_tensor(grad_ape, "grad_ape", (self.ratio, width), torch.float32, device, PTR_ALIGN_BYTES) + if grad_out.numel() == 0: + return + self._compiled_kernel(kv, score, ape, cu_seqlens, cu_seqlens_comp, grad_out, grad_kv, grad_score, grad_ape, _resolve_stream_handle(current_stream)) + self._record_streams((kv, score, ape, cu_seqlens, cu_seqlens_comp, grad_out, grad_kv, grad_score, grad_ape), current_stream, device) + + +# module-level bounded LRU cache of compiled API instances (thread-safe): +# (kind, ratio, head_dim, coff, total_tokens, total_comp, n_seg, device) -> api +# The compiled kernel underneath is shared per (ratio, head_dim, coff, device) through +# compressor_sm100's compile cache, so shape changes only rebuild the cheap wrapper +# object, never the JIT. +_api_cache = _LruDict() +# serializes verdict construction + JIT so concurrent same-config callers cannot +# compile the same kernel twice +_api_build_lock = threading.Lock() + + +def _get_api(kind, kv, score, ape, cu_seqlens, cu_seqlens_comp, out_shape, ratio, coff): + """Build (or fetch) a compiled forward/backward API instance for these tensors.""" + key = (kind, int(ratio), int(coff), out_shape[1], tuple(kv.shape), out_shape[0], cu_seqlens.shape[0], kv.device.index) + api = _api_cache.get(key) + if api is not None: + return api + with _api_build_lock: + api = _api_cache.get(key) + if api is not None: + return api + sample_out = torch.empty(out_shape, dtype=torch.bfloat16, device="meta") + cls = CSACompressorForward if kind == "fwd" else CSACompressorBackward + api = cls( + sample_kv=kv, + sample_score=score, + sample_ape=ape, + sample_cu_seqlens=cu_seqlens, + sample_cu_seqlens_comp=cu_seqlens_comp, + sample_out=sample_out, + ratio=ratio, + coff=coff, + ) + api.check_support() + api.compile() + _api_cache.put(key, api) + return api + + +def _infer_head_dim(kv: torch.Tensor, head_dim: Optional[int], coff: int) -> int: + """Infer ``head_dim`` from the packed kv width when not given explicitly.""" + if kv.ndim != 2: + raise ValueError(f"kv must be 2-D (total_tokens, coff * head_dim), got {tuple(kv.shape)}") + if head_dim is not None: + return int(head_dim) + width = kv.shape[1] + if coff < 1 or width % coff != 0: + raise ValueError(f"cannot infer head_dim from kv width {width} and coff {coff}") + return width // coff + + +def csa_compressor_forward_wrapper( + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_comp: torch.Tensor, + ratio: int = 4, + head_dim: Optional[int] = None, + coff: int = 2, + total_comp: Optional[int] = None, + stream: Optional[cuda.CUstream] = None, +) -> TupleDict: + """High-level forward wrapper. Allocates and returns the pooled output. + + Args: + kv: ``(total_tokens, coff * head_dim)`` BF16 gate values (THD packed). + score: ``(total_tokens, coff * head_dim)`` BF16 gate scores. + ape: ``(ratio, coff * head_dim)`` FP32 additive position embedding. + cu_seqlens: ``(B + 1,)`` int32 cumulative token counts per segment. + cu_seqlens_comp: ``(B + 1,)`` int32 cumulative compressed-block counts, + ``cu_seqlens_comp[b + 1] - cu_seqlens_comp[b] == seqlen_b // ratio``. + ratio: compression ratio (tokens per output block); validated envelope: 4. + head_dim: output feature dimension; inferred from ``kv`` width when omitted. + coff: 2 for the overlapping window form; validated envelope: 2. + total_comp: output row count. Defaults to ``cu_seqlens_comp[-1]`` (synchronizes); + pass it explicitly (e.g. a static CUDA-graph capacity, which must be + ``>= cu_seqlens_comp[-1]``) to stay capture-safe. + stream: CUDA stream for allocation and kernel launch (None -> current stream). + + Returns: + ``{'out': (total_comp, head_dim) BF16}`` pooled output (pre-RMSNorm). + """ + head_dim = _infer_head_dim(kv, head_dim, coff) + if total_comp is None: + if cu_seqlens_comp.numel() < 1: + raise ValueError("cu_seqlens_comp must have B + 1 >= 2 entries") + total_comp = int(cu_seqlens_comp[-1].item()) + api = _get_api("fwd", kv, score, ape, cu_seqlens, cu_seqlens_comp, (int(total_comp), head_dim), ratio, coff) + with torch.cuda.device(kv.device), _torch_stream_context(stream, kv.device): + out = torch.empty(int(total_comp), head_dim, dtype=torch.bfloat16, device=kv.device) + with torch.cuda.nvtx.range("csa_compressor_fwd_kernel"): + api.execute(kv, score, ape, cu_seqlens, cu_seqlens_comp, out, current_stream=stream) + return TupleDict(out=out) + + +def csa_compressor_backward_wrapper( + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_comp: torch.Tensor, + grad_out: torch.Tensor, + ratio: int = 4, + head_dim: Optional[int] = None, + coff: int = 2, + stream: Optional[cuda.CUstream] = None, +) -> TupleDict: + """High-level backward wrapper. Allocates the grad buffers and fills them. + + ``grad_out`` is ``(total_comp, head_dim)`` BF16 (the incoming gradient of the + forward wrapper's ``out``); gradients on static-capacity padding rows are ignored. + ``grad_kv``/``grad_score`` are allocated UNINITIALIZED — the kernel writes every + position, storing exact zeros to never-consumed positions itself, so no zero-fill + kernels run; ``grad_ape`` is allocated zeroed (fp32 atomic accumulation). When + ``total_comp == 0`` the kernel is not launched and all three grads are allocated as + zeros instead, preserving autograd's exact-zero semantics. + Raises ``RuntimeError`` under ``torch.use_deterministic_algorithms(True)`` because + ``grad_ape`` is accumulated with fp32 atomics (``grad_kv``/``grad_score`` are + deterministic and bitwise reproducible). + + Returns: + ``{'grad_kv': (total_tokens, coff * head_dim) BF16, + 'grad_score': (total_tokens, coff * head_dim) BF16, + 'grad_ape': (ratio, coff * head_dim) FP32}`` + """ + head_dim = _infer_head_dim(kv, head_dim, coff) + _reject_deterministic_backward() + if grad_out.ndim != 2: + raise ValueError(f"grad_out must be 2-D (total_comp, head_dim), got {tuple(grad_out.shape)}") + api = _get_api("bwd", kv, score, ape, cu_seqlens, cu_seqlens_comp, tuple(grad_out.shape), ratio, coff) + with torch.cuda.device(kv.device), _torch_stream_context(stream, kv.device): + if grad_out.shape[0] == 0: + # No kernel launch below -> the buffers must carry autograd's exact zeros. + grad_kv = torch.zeros_like(kv) + grad_score = torch.zeros_like(score) + else: + grad_kv = torch.empty_like(kv) + grad_score = torch.empty_like(score) + grad_ape = torch.zeros_like(ape, dtype=torch.float32) + with torch.cuda.nvtx.range("csa_compressor_bwd_kernel"): + api.execute(kv, score, ape, cu_seqlens, cu_seqlens_comp, grad_out, grad_kv, grad_score, grad_ape, current_stream=stream) + return TupleDict(grad_kv=grad_kv, grad_score=grad_score, grad_ape=grad_ape) diff --git a/python/cudnn/csa/compressor/compressor_sm100.py b/python/cudnn/csa/compressor/compressor_sm100.py new file mode 100644 index 000000000..51753db2f --- /dev/null +++ b/python/cudnn/csa/compressor/compressor_sm100.py @@ -0,0 +1,996 @@ +"""Fused CuTe-DSL forward+backward kernels for the CSA/HCA ``Compressor`` gated pooling. + +Ported from Megatron-LM (https://github.com/NVIDIA/Megatron-LM/pull/5984, see also +https://github.com/NVIDIA/Megatron-LM/issues/5968 for measurements and numerics); the +kernel math is unchanged. This module holds the kernels and the launch machinery; the +APIBase wrappers live in ``api.py``. + +The kernels fuse the gated-softmax pooling region of the CSA/HCA ``Compressor`` (THD +packed layout): the chain + + gather-index build -> gather -> ``+ APE`` -> overlap-window transform (``coff == 2``) + -> fp32 softmax over the window -> gated weighted sum -> bf16 cast + +is ONE forward kernel and ONE backward kernel (JIT-compiled per ``(ratio, head_dim, +coff)`` configuration). + +Semantics (ground truth = the eager region in Megatron-LM +``Compressor._forward_thd``): + For each segment ``s`` (``cu_seqlens[s]..cu_seqlens[s+1]``) and each output block ``b`` + of ``ratio`` tokens, the ``2 * ratio`` window (``coff == 2``) is + + - ``k in [0, ratio)``: previous block's token ``tok0 - ratio + k``, first-half + projection column ``j``, APE row ``k`` -> invalid for the segment's first block + (score ``-inf``, kv ``0``, exactly like the eager overlap-window transform). + - ``k in [ratio, 2 * ratio)``: own token ``tok0 + k - ratio``, second-half projection + column ``d + j``, APE row ``k - ratio``. + + ``out[b, j] = sum_k kv_k * softmax_k(score_k + ape_k)`` (fp32, single final bf16 + rounding). ``coff == 1`` (no overlap): the window is the block's own ``ratio`` tokens, + column ``j``, always valid. Per-segment tail tokens (``seqlen % ratio``) are dropped, + as in the eager code. + +Numerics: + All arithmetic is fp32 with a single final bf16 rounding. The fp32 accumulation + structure mirrors the eager ops (serial max, serial sum of exp, ``p = e / denom``, + serial sum ``kv * p``), with ``mul.rn.f32``/``fma.rn.f32`` pinned in PTX so results do + not depend on compiler FMA contraction. Against an fp32-intermediate eager reference, + ``dKV``/``dScore`` are bit-identical and the forward matches to within one bf16 + rounding step (see the tests and the Megatron-LM issue for data). + +Backward: + Atomic-free for ``dKV``/``dScore``: every consumed input element belongs to exactly one + pooling window (for ``coff == 2``, first-half columns are consumed by the NEXT block's + window and second-half columns by the OWN block's window), so gradient stores are + disjoint. Elements never consumed (segment-tail tokens; for ``coff == 2`` the + first-half projection columns of each segment's last block; all tokens of segments + shorter than ``ratio``; tokens beyond ``cu_seqlens[-1]`` when the gradient buffers + carry static token-capacity padding) are written as exact zeros by the kernel + itself — each such slot has a unique natural owner (see the kernel docstring) — so + ``dKV``/``dScore`` buffers need no zero-initialization and no separate fill kernels, + matching autograd output exactly. ``dAPE`` is accumulated in registers over ``rows_per_cta`` blocks and + then reduced with one fp32 atomic per ``(k, dim)`` per CTA into a buffer the caller + must still zero-initialize; ``dAPE`` is therefore not bitwise run-to-run deterministic + (forward, ``dKV`` and ``dScore`` are). When ``total_comp == 0`` no kernel is launched + and the buffers are left untouched (the wrapper falls back to allocating zeros). + +Static-capacity padding (``total_comp > cu_seqlens_comp[-1]``): + Forward computes the padding rows exactly like the eager code: they gather the window + from token 0 with first-in-segment semantics (requires ``total_tokens >= ratio``, like + the eager gather). Backward ignores incoming gradients on padding rows (they are tail + padding for CUDA-graph static shapes and are not consumed downstream). + +CUDA graphs: + The launch path is capture-compatible once the kernels for a given + ``(ratio, head_dim, coff)`` configuration have been JIT-compiled; compile (or run one + eager step) per configuration before capture. A call that would JIT under capture + raises a ``RuntimeError`` instead of corrupting the capture. + +``CUDNNFE_CSA_COMPRESSOR_FAST_LAUNCH=0`` disables only the cached-launch optimization +(see ``_FastLauncher``). +""" + +from __future__ import annotations + +import ctypes +import os +import threading + +import torch +import cuda.bindings.driver as cuda_driver + +import cutlass +import cutlass.base_dsl.typing as _cutlass_typing +import cutlass.cute as cute +import cutlass.cute.arch as cute_arch +import cutlass.cute.math as cute_math +from cutlass._mlir.dialects import llvm as _llvm +from cutlass.cute.runtime import make_ptr + +_ENV_FAST_LAUNCH = "CUDNNFE_CSA_COMPRESSOR_FAST_LAUNCH" + +# The only compute capability the kernels have been validated on so far. The kernels use +# no architecture-specific features (plain loads/stores, fp32 atomics, pinned mul/fma +# PTX), but wider coverage stays opt-in until validated per architecture. The +# ``cute.compile`` default arch resolution maps (10, 0) to ``sm_100a``, so the CC gate +# also pins the generated SASS target. +SUPPORTED_COMPUTE_CAPABILITY = (10, 0) + + +# ============================================================================= +# Cached fast-path launcher +# ============================================================================= +# Each steady-state CuTe-DSL call spends tens of microseconds of pure host Python +# (rebuilding pointer/scalar/stream argument objects, adapter lookups, fresh ctypes +# allocations, re-packing the void** array) to end at a ~3-4 us C launch call. For +# microsecond-scale kernels that overhead dominates the wall clock, so the launch state +# is snapshotted ONCE per (kernel, config, device, thread) and replayed with in-place +# mutation of the argument storages. +# ============================================================================= + +# torch's raw current-stream query (~0.5 us) vs `torch.cuda.current_stream` object +# construction (~2-3 us). Same handle the slow path ends up passing. Private API: guard +# the bind so module import survives torch builds that do not expose it. +_raw_stream = getattr(torch._C, "_cuda_getCurrentRawStream", None) +if _raw_stream is None: # pragma: no cover - older/stripped torch builds + + def _raw_stream(device_index=None): + """Fallback raw-stream query via the public torch API.""" + return torch.cuda.current_stream(device_index).cuda_stream + + +def _fast_launch_enabled() -> bool: + """Return True unless the cached-launch optimization is disabled via environment.""" + return os.environ.get(_ENV_FAST_LAUNCH, "1") == "1" + + +def _view_for_arg(arg, addr): + """Build a ctypes view over the storage backing one execution-args slot.""" + if isinstance(arg, _cutlass_typing.Numeric): + if isinstance(arg, _cutlass_typing.Boolean): + return ctypes.c_bool.from_address(addr) + if isinstance(arg, _cutlass_typing.Integer): + width = type(arg).width + signed = getattr(type(arg), "signed", True) + ctype = getattr(ctypes, f"c_{'int' if signed else 'uint'}{width}") + return ctype.from_address(addr) + if isinstance(arg, _cutlass_typing.Float32): + return ctypes.c_float.from_address(addr) + if isinstance(arg, _cutlass_typing.Float64): + return ctypes.c_double.from_address(addr) + raise TypeError(f"unsupported numeric scalar {type(arg)!r}") + # A cute runtime Pointer (make_ptr) stores its address in a c_void_p `_desc`; + # CUstream's storage is its own pointer-sized handle. + if hasattr(arg, "_desc") and isinstance(arg._desc, ctypes.c_void_p): + return ctypes.c_void_p.from_address(addr) + if type(arg).__name__ == "CUstream": + return ctypes.c_void_p.from_address(addr) + raise TypeError(f"unsupported argument type {type(arg)!r}") + + +class _FastLauncher: + """Replayable launch state for one compiled CuTe-DSL function. + + ``slots[i]`` is a ctypes view over the storage feeding runtime argument ``i`` (same + order as the tuple passed to ``fn(*args)``); write ``.value`` then call ``launch()``. + + Guards: + - Only flat argument tuples of cute runtime pointers (``make_ptr``), cutlass + scalars, and ``CUstream`` are eligible; anything else raises during construction + and the wrapper stays on its slow path. + - Construction introspects private-but-stable DSL internals + (``_default_executor``, ``_get_invoke_packed_args``); any structural mismatch on + a future ``nvidia-cutlass-dsl`` upgrade raises during construction, and the + wrapper permanently falls back to the regular (slow) launch path rather than + attempting a launch from a partially built snapshot. + """ + + __slots__ = ("slots", "_capi", "_packed", "_res", "_has_res", "_keep") + + def __init__(self, fn, args): + # Must run after the wrapper's first real `fn(*args)` call so the default + # executor (device context, loaded modules) exists. + exe_args, adapted = fn.execution_args.generate_execution_args(args, {}) + executor = fn._default_executor + if executor is None: + raise RuntimeError("build _FastLauncher after the first fn() call") + if len(exe_args) != len(args): + # struct/dlpack args expand to multiple slots -> unsupported. + raise TypeError(f"non-flat exe_args ({len(exe_args)} slots for {len(args)} args)") + # Private copy of the packed void** array: the executor's own is a shared + # thread-local scratch buffer that any interleaved slow-path call would repoint + # to its (dead) per-call storages. + tls_packed = executor._get_invoke_packed_args(exe_args) + total = len(exe_args) + executor._num_extra_args + packed = (ctypes.c_void_p * total)() + for i in range(total): + packed[i] = tls_packed[i] + views = [] + for arg, exe_arg in zip(args, exe_args): + addr = exe_arg.value if isinstance(exe_arg, ctypes.c_void_p) else int(exe_arg) + views.append(_view_for_arg(arg, addr)) + self.slots = views + self._capi = executor.capi_func + self._has_res = executor._has_cuda_result + if self._has_res: + # Private result storage: the executor's own `cuda_result` is shared by + # every launcher built from this compiled function, so concurrent launches + # from different threads would overwrite one another's CUDA status. The + # result address is the first extra slot after the base arguments (see + # jit_executor._get_invoke_packed_args). + self._res = type(executor.cuda_result)() + packed[len(exe_args)] = ctypes.addressof(self._res) + else: + self._res = None + self._packed = packed + # Keep every object owning a storage referenced by `packed` alive. + self._keep = (args, exe_args, adapted, fn, executor) + + def launch(self): + """Replay the snapshotted launch with the current slot values.""" + self._capi(self._packed) + if self._has_res: + result = self._res.value + if result != 0: + raise RuntimeError(f"CUDA error {result} in CuTe-DSL fast launch (set {_ENV_FAST_LAUNCH}=0 to fall back to the slow launch path)") + + +class _FastCache: + """Thread-local ``{key: _FastLauncher}`` with build-once semantics. + + ``get`` returns a launcher or None (not built / build failed / disabled). ``put`` + attempts to build; a failed build is remembered so the wrapper pays the (cheap) + attempt exactly once per thread and stays on its slow path afterwards. The cache is + thread-local because callers may run forward on the main thread while backward runs + on an autograd thread, and slot mutation is not thread-safe. + """ + + def __init__(self): + self._tls = threading.local() + + def _map(self): + """Return this thread's key -> launcher map.""" + cache_map = getattr(self._tls, "m", None) + if cache_map is None: + cache_map = {} + self._tls.m = cache_map + return cache_map + + def get(self, key): + """Return the cached launcher for ``key`` or None.""" + launcher = self._map().get(key) + return launcher if launcher is not None and launcher is not False else None + + def put(self, key, fn, args): + """Try to build and cache a launcher for ``key``; never fails the call.""" + if not _fast_launch_enabled(): + return + cache_map = self._map() + if key in cache_map: + return + try: + cache_map[key] = _FastLauncher(fn, args) + except Exception: # pylint: disable=broad-except + # Structural mismatch (DSL upgrade, exotic arg): remember and keep the + # wrapper on its slow path. Never fail the call. + cache_map[key] = False + + +_FAST = _FastCache() + + +# ============================================================================= +# CuTe-DSL kernel definitions +# ============================================================================= + +_NEG_INF = float("-inf") + + +@cutlass.dsl_user_op +def _fmul_rn(a, b, *, loc=None, ip=None): + """fp32 multiply pinned to ``mul.rn.f32`` (opaque to FMA contraction). + + The eager ``(kv * weights).sum(dim=1)`` and the softmax-backward inner sum both + accumulate ROUNDED products serially; letting the compiler contract mul+add into + FMA breaks bit-exactness against the fp32 eager reference. + """ + return cutlass.Float32( + _llvm.inline_asm( + cutlass.Float32.mlir_type, + [ + cutlass.Float32(a).ir_value(loc=loc, ip=ip), + cutlass.Float32(b).ir_value(loc=loc, ip=ip), + ], + "mul.rn.f32 $0, $1, $2;", + "=f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cutlass.dsl_user_op +def _ffma_rn(a, b, c, *, loc=None, ip=None): + """fp32 fma pinned to ``fma.rn.f32``. + + The eager softmax-backward epilogue is ``ds = fma(p, -S, round(p * dp))``; pinning + removes any dependence on compiler contraction choices. + """ + return cutlass.Float32( + _llvm.inline_asm( + cutlass.Float32.mlir_type, + [ + cutlass.Float32(a).ir_value(loc=loc, ip=ip), + cutlass.Float32(b).ir_value(loc=loc, ip=ip), + cutlass.Float32(c).ir_value(loc=loc, ip=ip), + ], + "fma.rn.f32 $0, $1, $2, $3;", + "=f,f,f,f", + has_side_effects=False, + is_align_stack=False, + asm_dialect=_llvm.AsmDialect.AD_ATT, + loc=loc, + ip=ip, + ) + ) + + +@cute.kernel +def _compressor_fwd_kernel( + mKV: cute.Tensor, # flat [T * W] bf16, W = coff * d + mScore: cute.Tensor, # flat [T * W] bf16 + mAPE: cute.Tensor, # flat [ratio * W] fp32 + mCu: cute.Tensor, # [n_seq + 1] int32 (token cu_seqlens) + mCuComp: cute.Tensor, # [n_seq + 1] int32 (block cu_seqlens) + mOut: cute.Tensor, # flat [nb_total * d] bf16 + nb_total: cutlass.Int32, + n_seq: cutlass.Int32, + ratio: cutlass.Constexpr, + d: cutlass.Constexpr, + coff: cutlass.Constexpr, + vec: cutlass.Constexpr, + rows_per_cta: cutlass.Constexpr, + threads: cutlass.Constexpr, +): + """Forward: one thread per (output block, ``vec`` adjacent head dims). + + ``vec == 2`` widens every bf16 access to one 32-bit load/store (``vec == 1`` is the + scalar layout for odd ``head_dim``). The per-thread window slices are contiguous and + ``vec``-aligned by construction (``W``, ``d`` and the thread's first column are all + multiples of ``vec``), which ``cute.assume`` makes provable so ``autovec_copy`` + lowers each slice to a single ``vec * 16``-bit universal copy. Wider vectors were + measured and rejected: 64/128/256-bit variants cut instructions but blow up + registers (80/147/255 per thread) and occupancy, losing to ``vec == 2`` on every + production shape. + + The per-lane fp32 math is IDENTICAL to the scalar kernel (same op order per output + element, one lane per head dim), so the output stays bitwise stable across the + ``vec`` configurations. + """ + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + ncol: cutlass.Constexpr = d // vec # thread-column count per output row + col = bidy * threads + tidx + W: cutlass.Constexpr = coff * d + win: cutlass.Constexpr = 2 * ratio if coff == 2 else ratio + + if col < ncol: + cvec = col * vec # first head-dim column of this thread's lane group + + # Hoist APE loads: constant per (k, lane) across all rows. + ape_k = [] + for k in cutlass.range_constexpr(win): + if cutlass.const_expr(coff == 2 and k < ratio): + colbase = cvec + else: + colbase = (d + cvec) if cutlass.const_expr(coff == 2) else cvec + fr_a = cute.make_rmem_tensor((vec,), cutlass.Float32) + aoff = cute.assume((k % ratio) * W + colbase, divby=vec) + gA = cute.make_tensor(mAPE.iterator + aoff, cute.make_layout(vec)) + cute.autovec_copy(gA, fr_a) + for j in cutlass.range_constexpr(vec): + ape_k.append(cutlass.Float32(fr_a[j])) + + # True compressed row count; rows in [nb_valid, nb_total) are static-capacity + # padding and gather the window from token 0 with first-in-segment semantics, + # like the eager code. + nb_valid = mCuComp[n_seq] + + for rr in cutlass.range_constexpr(rows_per_cta): + bb = bidx * rows_per_cta + rr + if bb < nb_total: + # Per-segment boundary scan (n_seq is small). + seq_idx = cutlass.Int32(0) + bis = cutlass.Int32(0) + if bb < nb_valid: + bis = cutlass.Int32(bb) + for s in cutlass.range(n_seq): + cs = mCuComp[s] + ce = mCuComp[s + 1] + if bb >= cs: + if bb < ce: + seq_idx = s + bis = bb - cs + tok0 = mCu[seq_idx] + bis * ratio + + sv = [] + kvv = [] + for k in cutlass.range_constexpr(win): + fr_s = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + fr_k = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + if cutlass.const_expr(coff == 2 and k < ratio): + if bis > 0: + off = cute.assume((tok0 - ratio + k) * W + cvec, divby=vec) + gS = cute.make_tensor(mScore.iterator + off, cute.make_layout(vec)) + gK = cute.make_tensor(mKV.iterator + off, cute.make_layout(vec)) + cute.autovec_copy(gS, fr_s) + cute.autovec_copy(gK, fr_k) + # Same value construction as the scalar kernel: the invalid + # window contributes the CONSTANT -inf score (no APE add — APE + # values are not required to be finite) and a zero kv lane. + for j in cutlass.range_constexpr(vec): + v = cutlass.Float32(_NEG_INF) + u = cutlass.Float32(0.0) + if bis > 0: + v = cutlass.Float32(fr_s[j]) + ape_k[k * vec + j] + u = cutlass.Float32(fr_k[j]) + sv.append(v) + kvv.append(u) + else: + if cutlass.const_expr(coff == 2): + off = cute.assume((tok0 + k - ratio) * W + d + cvec, divby=vec) + else: + off = cute.assume((tok0 + k) * W + cvec, divby=vec) + gS = cute.make_tensor(mScore.iterator + off, cute.make_layout(vec)) + gK = cute.make_tensor(mKV.iterator + off, cute.make_layout(vec)) + cute.autovec_copy(gS, fr_s) + cute.autovec_copy(gK, fr_k) + for j in cutlass.range_constexpr(vec): + sv.append(cutlass.Float32(fr_s[j]) + ape_k[k * vec + j]) + kvv.append(cutlass.Float32(fr_k[j])) + + fr_o = cute.make_rmem_tensor((vec,), cutlass.BFloat16) + for j in cutlass.range_constexpr(vec): + mx = sv[j] + for k in cutlass.range_constexpr(1, win): + if sv[k * vec + j] > mx: + mx = sv[k * vec + j] + den = cutlass.Float32(0.0) + ex = [] + for k in cutlass.range_constexpr(win): + e = cute_math.exp(sv[k * vec + j] - mx) + den = den + e + ex.append(e) + acc = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(win): + acc = acc + _fmul_rn(kvv[k * vec + j], ex[k] / den) + fr_o[j] = cutlass.BFloat16(acc) + ooff = cute.assume(bb * d + cvec, divby=vec) + gO = cute.make_tensor(mOut.iterator + ooff, cute.make_layout(vec)) + cute.autovec_copy(fr_o, gO) + + +@cute.kernel +def _compressor_bwd_kernel( + mKV: cute.Tensor, # flat [T * W] bf16 + mScore: cute.Tensor, # flat [T * W] bf16 + mAPE: cute.Tensor, # flat [ratio * W] fp32 + mCu: cute.Tensor, # [n_seq + 1] int32 + mCuComp: cute.Tensor, # [n_seq + 1] int32 + mGO: cute.Tensor, # flat [nb_total * d] bf16 + mGKV: cute.Tensor, # flat [T * W] bf16 (fully written; may be uninitialized) + mGS: cute.Tensor, # flat [T * W] bf16 (fully written; may be uninitialized) + mGAPE: cute.Tensor, # flat [ratio * W] fp32 (zero-initialized) + nb_total: cutlass.Int32, + n_seq: cutlass.Int32, + total_tokens: cutlass.Int32, + ratio: cutlass.Constexpr, + d: cutlass.Constexpr, + coff: cutlass.Constexpr, + rows_per_cta: cutlass.Constexpr, + threads: cutlass.Constexpr, +): + """Backward: recompute window probs, disjoint ``dKV``/``dScore`` stores, ``dAPE`` atomics. + + ``dKV``/``dScore`` are FULLY written by the kernel: consumed positions get their + gradients, and every never-consumed position gets an exact zero from its unique + natural owner (see below), so the caller can pass uninitialized buffers instead of + paying two tensor-wide zero-fills. The zero-write ownership keeps all stores + disjoint (no atomics, bitwise run-to-run deterministic): + + - for ``coff == 2``, the first-half columns of each segment's LAST block's own + tokens (no next block consumes them) — written by that last block; + - per-segment tail tokens (``seqlen % ratio``, both halves) — written by the + segment's last block; + - all tokens of segments with zero output blocks (``seqlen < ratio``) — written + by the CTA column ``bidx == 0``; + - tokens beyond ``cu_seqlens[-1]`` (static token-capacity padding of the + gradient buffers) — grid-strided across the CTA columns. + + ``dAPE`` is still accumulated into a caller-zero-initialized buffer with one fp32 + atomic per ``(k, dim)`` per CTA. Rows in ``[cu_seqlens_comp[-1], nb_total)`` are + static-capacity padding; their incoming gradients are ignored. + """ + tidx, _, _ = cute.arch.thread_idx() + bidx, bidy, _ = cute.arch.block_idx() + dim = bidy * threads + tidx + W: cutlass.Constexpr = coff * d + win: cutlass.Constexpr = 2 * ratio if coff == 2 else ratio + ZERO_BF16 = cutlass.BFloat16(0.0) + + if dim < d: + ape_k = [] + dape = [] + for k in cutlass.range_constexpr(win): + if cutlass.const_expr(coff == 2 and k < ratio): + col = dim + else: + col = (d + dim) if cutlass.const_expr(coff == 2) else dim + ape_k.append(mAPE[(k % ratio) * W + col]) + dape.append(cutlass.Float32(0.0)) + + nb_valid = mCuComp[n_seq] + + # CTA column (0, bidy) zeroes both halves of every token in segments that have + # zero output blocks (seqlen < ratio): those tokens are never consumed by any + # pooling window, so no block-owning CTA would otherwise write them. + if bidx == 0: + for s in cutlass.range(n_seq): + if mCuComp[s + 1] == mCuComp[s]: + t0 = mCu[s] + t1 = mCu[s + 1] + for tt in cutlass.range(t1 - t0): + mGKV[(t0 + tt) * W + dim] = ZERO_BF16 + mGS[(t0 + tt) * W + dim] = ZERO_BF16 + if cutlass.const_expr(coff == 2): + mGKV[(t0 + tt) * W + d + dim] = ZERO_BF16 + mGS[(t0 + tt) * W + d + dim] = ZERO_BF16 + + # Tokens in [cu_seqlens[-1], total_tokens) are static token-capacity padding of + # the gradient buffers (CUDA-graph static shapes): no segment owns them, so the + # CTA columns zero them in a grid-strided sweep. count == 0 in the common + # exact-size case. The quotient/remainder split keeps every intermediate within + # int32 for any count < 2**31. + gdimx, _, _ = cute.arch.grid_dim() + pad0 = mCu[n_seq] + pad_count = total_tokens - pad0 + if bidx < pad_count: + my_count = pad_count // gdimx + if bidx < pad_count % gdimx: + my_count = my_count + 1 + for i in cutlass.range(my_count): + t = pad0 + bidx + i * gdimx + mGKV[t * W + dim] = ZERO_BF16 + mGS[t * W + dim] = ZERO_BF16 + if cutlass.const_expr(coff == 2): + mGKV[t * W + d + dim] = ZERO_BF16 + mGS[t * W + d + dim] = ZERO_BF16 + + for rr in cutlass.range_constexpr(rows_per_cta): + bb = bidx * rows_per_cta + rr + if bb < nb_valid: + seq_idx = cutlass.Int32(0) + bis = cutlass.Int32(bb) + for s in cutlass.range(n_seq): + cs = mCuComp[s] + ce = mCuComp[s + 1] + if bb >= cs: + if bb < ce: + seq_idx = s + bis = bb - cs + tok0 = mCu[seq_idx] + bis * ratio + + # Recompute window probs (same order as forward). + sv = [] + kvv = [] + offs = [] + for k in cutlass.range_constexpr(win): + if cutlass.const_expr(coff == 2 and k < ratio): + off = (tok0 - ratio + k) * W + dim + v = cutlass.Float32(_NEG_INF) + u = cutlass.Float32(0.0) + if bis > 0: + v = cutlass.Float32(mScore[off]) + ape_k[k] + u = cutlass.Float32(mKV[off]) + else: + if cutlass.const_expr(coff == 2): + off = (tok0 + k - ratio) * W + d + dim + else: + off = (tok0 + k) * W + dim + v = cutlass.Float32(mScore[off]) + ape_k[k] + u = cutlass.Float32(mKV[off]) + sv.append(v) + kvv.append(u) + offs.append(off) + + mx = sv[0] + for k in cutlass.range_constexpr(1, win): + if sv[k] > mx: + mx = sv[k] + den = cutlass.Float32(0.0) + ex = [] + for k in cutlass.range_constexpr(win): + e = cute_math.exp(sv[k] - mx) + den = den + e + ex.append(e) + + go = cutlass.Float32(mGO[bb * d + dim]) + + # Same expression tree as torch's softmax_backward_data: + # dp_k = go * kv_k ; S = serial sum of ROUNDED dp_k * p_k ; + # ds_k = fma(p_k, -S, round(dp_k * p_k)) ; dkv_k = go * p_k. + p = [] + dp = [] + S = cutlass.Float32(0.0) + for k in cutlass.range_constexpr(win): + pk = ex[k] / den + dpk = go * kvv[k] + S = S + _fmul_rn(dpk, pk) + p.append(pk) + dp.append(dpk) + + for k in cutlass.range_constexpr(win): + if cutlass.const_expr(coff == 2 and k < ratio): + if bis > 0: + ds = _ffma_rn(p[k], -S, _fmul_rn(dp[k], p[k])) + mGKV[offs[k]] = cutlass.BFloat16(go * p[k]) + mGS[offs[k]] = cutlass.BFloat16(ds) + dape[k] = dape[k] + ds + else: + ds = _ffma_rn(p[k], -S, _fmul_rn(dp[k], p[k])) + mGKV[offs[k]] = cutlass.BFloat16(go * p[k]) + mGS[offs[k]] = cutlass.BFloat16(ds) + dape[k] = dape[k] + ds + + # The segment's last block additionally zeroes the never-consumed slots + # it is the unique natural owner of: (a) for coff == 2 the first-half + # columns of its own tokens (there is no next block to consume them), + # (b) the segment's tail tokens (seqlen % ratio, both halves). + is_last = bb + 1 == mCuComp[seq_idx + 1] + if is_last: + if cutlass.const_expr(coff == 2): + for k in cutlass.range_constexpr(ratio): + mGKV[(tok0 + k) * W + dim] = ZERO_BF16 + mGS[(tok0 + k) * W + dim] = ZERO_BF16 + tail0 = tok0 + ratio + tail1 = mCu[seq_idx + 1] + for tt in cutlass.range(tail1 - tail0): + mGKV[(tail0 + tt) * W + dim] = ZERO_BF16 + mGS[(tail0 + tt) * W + dim] = ZERO_BF16 + if cutlass.const_expr(coff == 2): + mGKV[(tail0 + tt) * W + d + dim] = ZERO_BF16 + mGS[(tail0 + tt) * W + d + dim] = ZERO_BF16 + + # One fp32 atomic per (k, dim) per CTA (amortized over rows_per_cta rows). + for k in cutlass.range_constexpr(win): + if cutlass.const_expr(coff == 2 and k < ratio): + col = dim + else: + col = (d + dim) if cutlass.const_expr(coff == 2) else dim + cute_arch.atomic_add(mGAPE.iterator + ((k % ratio) * W + col), dape[k]) + + +_EXT = (1 << 31) - 1 # flat extent placeholder (int32 offsets, no bounds checks) + + +@cute.jit +def _compressor_fwd_launch( + kv_ptr: cute.Pointer, + score_ptr: cute.Pointer, + ape_ptr: cute.Pointer, + cu_ptr: cute.Pointer, + cuc_ptr: cute.Pointer, + out_ptr: cute.Pointer, + nb_total: cutlass.Int32, + n_seq: cutlass.Int32, + stream: cuda_driver.CUstream, + ratio: cutlass.Constexpr, + d: cutlass.Constexpr, + coff: cutlass.Constexpr, + vec: cutlass.Constexpr, + rows_per_cta: cutlass.Constexpr, + threads: cutlass.Constexpr, +): + """JIT entry point that wraps raw pointers into tensors and launches forward.""" + lay = cute.make_layout(_EXT) + mKV = cute.make_tensor(kv_ptr, lay) + mScore = cute.make_tensor(score_ptr, lay) + mAPE = cute.make_tensor(ape_ptr, lay) + mCu = cute.make_tensor(cu_ptr, lay) + mCuComp = cute.make_tensor(cuc_ptr, lay) + mOut = cute.make_tensor(out_ptr, lay) + ncol = d // vec + gx = (nb_total + rows_per_cta - 1) // rows_per_cta + gy = (ncol + threads - 1) // threads + _compressor_fwd_kernel( + mKV, + mScore, + mAPE, + mCu, + mCuComp, + mOut, + nb_total, + n_seq, + ratio, + d, + coff, + vec, + rows_per_cta, + threads, + ).launch(grid=(gx, gy, 1), block=(threads, 1, 1), stream=stream) + + +@cute.jit +def _compressor_bwd_launch( + kv_ptr: cute.Pointer, + score_ptr: cute.Pointer, + ape_ptr: cute.Pointer, + cu_ptr: cute.Pointer, + cuc_ptr: cute.Pointer, + go_ptr: cute.Pointer, + gkv_ptr: cute.Pointer, + gs_ptr: cute.Pointer, + gape_ptr: cute.Pointer, + nb_total: cutlass.Int32, + n_seq: cutlass.Int32, + total_tokens: cutlass.Int32, + stream: cuda_driver.CUstream, + ratio: cutlass.Constexpr, + d: cutlass.Constexpr, + coff: cutlass.Constexpr, + rows_per_cta: cutlass.Constexpr, + threads: cutlass.Constexpr, +): + """JIT entry point that wraps raw pointers into tensors and launches backward.""" + lay = cute.make_layout(_EXT) + mKV = cute.make_tensor(kv_ptr, lay) + mScore = cute.make_tensor(score_ptr, lay) + mAPE = cute.make_tensor(ape_ptr, lay) + mCu = cute.make_tensor(cu_ptr, lay) + mCuComp = cute.make_tensor(cuc_ptr, lay) + mGO = cute.make_tensor(go_ptr, lay) + mGKV = cute.make_tensor(gkv_ptr, lay) + mGS = cute.make_tensor(gs_ptr, lay) + mGAPE = cute.make_tensor(gape_ptr, lay) + gx = (nb_total + rows_per_cta - 1) // rows_per_cta + gy = (d + threads - 1) // threads + _compressor_bwd_kernel( + mKV, + mScore, + mAPE, + mCu, + mCuComp, + mGO, + mGKV, + mGS, + mGAPE, + nb_total, + n_seq, + total_tokens, + ratio, + d, + coff, + rows_per_cta, + threads, + ).launch(grid=(gx, gy, 1), block=(threads, 1, 1), stream=stream) + + +_COMPILED = {} +# Serializes JIT compilation so concurrent same-config callers cannot compile the same +# kernel twice (the compiled-function cache itself is a plain dict guarded by the GIL). +_COMPILE_LOCK = threading.Lock() +_BWD_ROWS, _BWD_THREADS = 8, 128 + + +def _fwd_schedule(d): + """Forward launch schedule ``(vec, rows_per_cta, threads)`` for ``head_dim == d``. + + ``vec = 2`` (32-bit paired bf16 accesses) whenever ``d`` is even, else the scalar + ``vec = 1`` layout. One output row per CTA with 64-thread column groups: measured + optimum across the production shapes (1x/3x 8192-token packs, head_dim 128/512) — + smaller CTAs raise the sub-wave grid width that limits the small shapes, and wider + per-thread vectors trade instructions for registers/occupancy at a loss (see the + kernel docstring). For enormous head_dims whose column count would overflow the + 65535 ``gridDim.y`` limit at 64 threads, fall back to 128-thread CTAs (the previous + schedule's capability envelope). + """ + vec = 2 if d % 2 == 0 else 1 + ncol = d // vec + threads = 64 if ncol >= 64 else ncol + if (ncol + threads - 1) // threads > 65535: + threads = 128 + return vec, 1, threads + + +# make_ptr assumed alignments below. Contiguity does NOT imply base-pointer alignment +# (storage-offset views), so the API layer checks every runtime tensor's data_ptr() +# against these before launching. +PTR_ALIGN_BYTES = 16 # bf16 / fp32 operands +CU_ALIGN_BYTES = 4 # int32 cu_seqlens operands + + +def _bf16_ptr(t): + """Wrap a bf16 tensor's data pointer for the DSL.""" + return make_ptr(cutlass.BFloat16, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + + +def _f32_ptr(t): + """Wrap an fp32 tensor's data pointer for the DSL.""" + return make_ptr(cutlass.Float32, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=16) + + +def _i32_ptr(t): + """Wrap an int32 tensor's data pointer for the DSL.""" + return make_ptr(cutlass.Int32, t.data_ptr(), cute.AddressSpace.gmem, assumed_align=4) + + +def _compile_fwd(key, args, ratio, d, coff): + """JIT-compile the forward launch entry for ``key`` (capture-guarded).""" + with _COMPILE_LOCK: + fn = _COMPILED.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"CSA compressor: first call for config {key} happened under CUDA " + "graph capture (JIT compilation is not capture-safe); compile() or " + "run one eager step for this configuration before capturing." + ) + fn = cute.compile(_compressor_fwd_launch, *args, ratio, d, coff, *_fwd_schedule(d)) + _COMPILED[key] = fn + return fn + + +def _compile_bwd(key, args, ratio, d, coff): + """JIT-compile the backward launch entry for ``key`` (capture-guarded).""" + with _COMPILE_LOCK: + fn = _COMPILED.get(key) + if fn is None: + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + f"CSA compressor: first call for config {key} happened under CUDA " + "graph capture (JIT compilation is not capture-safe); compile() or " + "run one eager step for this configuration before capturing." + ) + fn = cute.compile(_compressor_bwd_launch, *args, ratio, d, coff, _BWD_ROWS, _BWD_THREADS) + _COMPILED[key] = fn + return fn + + +def precompile_fwd(ratio, d, coff, device): + """Ensure the forward kernel for ``(ratio, d, coff, device)`` is JIT-compiled. + + Compilation only traces types (pointers/scalars/stream are runtime arguments), so + tiny scratch buffers stand in for the real tensors; nothing is launched. + """ + key = ("fwd", ratio, d, coff, device.index) + if key in _COMPILED: + return + with torch.cuda.device(device): + scratch_bf16 = torch.zeros(16, device=device, dtype=torch.bfloat16) + scratch_f32 = torch.zeros(16, device=device, dtype=torch.float32) + scratch_i32 = torch.zeros(16, device=device, dtype=torch.int32) + stream = cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) + args = ( + _bf16_ptr(scratch_bf16), + _bf16_ptr(scratch_bf16), + _f32_ptr(scratch_f32), + _i32_ptr(scratch_i32), + _i32_ptr(scratch_i32), + _bf16_ptr(scratch_bf16), + cutlass.Int32(0), + cutlass.Int32(1), + stream, + ) + _compile_fwd(key, args, ratio, d, coff) + + +def precompile_bwd(ratio, d, coff, device): + """Ensure the backward kernel for ``(ratio, d, coff, device)`` is JIT-compiled.""" + key = ("bwd", ratio, d, coff, device.index) + if key in _COMPILED: + return + with torch.cuda.device(device): + scratch_bf16 = torch.zeros(16, device=device, dtype=torch.bfloat16) + scratch_f32 = torch.zeros(16, device=device, dtype=torch.float32) + scratch_i32 = torch.zeros(16, device=device, dtype=torch.int32) + stream = cuda_driver.CUstream(torch.cuda.current_stream(device).cuda_stream) + args = ( + _bf16_ptr(scratch_bf16), + _bf16_ptr(scratch_bf16), + _f32_ptr(scratch_f32), + _i32_ptr(scratch_i32), + _i32_ptr(scratch_i32), + _bf16_ptr(scratch_bf16), + _bf16_ptr(scratch_bf16), + _bf16_ptr(scratch_bf16), + _f32_ptr(scratch_f32), + cutlass.Int32(0), + cutlass.Int32(1), + cutlass.Int32(0), + stream, + ) + _compile_bwd(key, args, ratio, d, coff) + + +def run_fwd(kv, score, ape, cu_i, cuc_i, out, nb_total, ratio, d, coff, stream_handle=None): + """Launch the forward kernel (cached fast path -> compiled slow path -> JIT). + + ``stream_handle`` is an integer CUDA stream handle; None uses torch's current + stream on ``kv``'s device. The launch is anchored in ``kv``'s device context: the + compiled module and the default-stream query are per-device, and launching from a + foreign current device silently misbehaves. + """ + dev = kv.device.index + key = ("fwd", ratio, d, coff, dev) + if stream_handle is None: + stream_handle = _raw_stream(dev) + with torch.cuda.device(dev): + launcher = _FAST.get(key) + if launcher is not None: + # Cached launch: mutate the snapshotted argument storages in place; this is + # the same launch the slow path below performs. + slots = launcher.slots + slots[0].value = kv.data_ptr() + slots[1].value = score.data_ptr() + slots[2].value = ape.data_ptr() + slots[3].value = cu_i.data_ptr() + slots[4].value = cuc_i.data_ptr() + slots[5].value = out.data_ptr() + slots[6].value = nb_total + slots[7].value = cu_i.numel() - 1 + slots[8].value = stream_handle + launcher.launch() + return + stream = cuda_driver.CUstream(stream_handle) + args = ( + _bf16_ptr(kv), + _bf16_ptr(score), + _f32_ptr(ape), + _i32_ptr(cu_i), + _i32_ptr(cuc_i), + _bf16_ptr(out), + cutlass.Int32(nb_total), + cutlass.Int32(cu_i.numel() - 1), + stream, + ) + fn = _COMPILED.get(key) + if fn is None: + fn = _compile_fwd(key, args, ratio, d, coff) + fn(*args) + _FAST.put(key, fn, args) + + +def run_bwd(kv, score, ape, cu_i, cuc_i, go, gkv, gs, gape, nb_total, ratio, d, coff, stream_handle=None): + """Launch the backward kernel (cached fast path -> compiled slow path -> JIT). + + Device-context anchoring as in :func:`run_fwd`. The gradient-buffer token capacity + (for the kernel's padding-token zero sweep) is derived from ``kv``'s element count. + """ + dev = kv.device.index + key = ("bwd", ratio, d, coff, dev) + total_tokens = kv.numel() // (coff * d) + if stream_handle is None: + stream_handle = _raw_stream(dev) + with torch.cuda.device(dev): + launcher = _FAST.get(key) + if launcher is not None: + slots = launcher.slots + slots[0].value = kv.data_ptr() + slots[1].value = score.data_ptr() + slots[2].value = ape.data_ptr() + slots[3].value = cu_i.data_ptr() + slots[4].value = cuc_i.data_ptr() + slots[5].value = go.data_ptr() + slots[6].value = gkv.data_ptr() + slots[7].value = gs.data_ptr() + slots[8].value = gape.data_ptr() + slots[9].value = nb_total + slots[10].value = cu_i.numel() - 1 + slots[11].value = total_tokens + slots[12].value = stream_handle + launcher.launch() + return + stream = cuda_driver.CUstream(stream_handle) + args = ( + _bf16_ptr(kv), + _bf16_ptr(score), + _f32_ptr(ape), + _i32_ptr(cu_i), + _i32_ptr(cuc_i), + _bf16_ptr(go), + _bf16_ptr(gkv), + _bf16_ptr(gs), + _f32_ptr(gape), + cutlass.Int32(nb_total), + cutlass.Int32(cu_i.numel() - 1), + cutlass.Int32(total_tokens), + stream, + ) + fn = _COMPILED.get(key) + if fn is None: + fn = _compile_bwd(key, args, ratio, d, coff) + fn(*args) + _FAST.put(key, fn, args) diff --git a/test/python/fe_api/csa/test_CSA_compressor.py b/test/python/fe_api/csa/test_CSA_compressor.py new file mode 100644 index 000000000..8e5293782 --- /dev/null +++ b/test/python/fe_api/csa/test_CSA_compressor.py @@ -0,0 +1,850 @@ +"""Tests for the fused CSA/HCA Compressor gated-pooling kernels (``cudnn.csa``). + +Ported with the kernels from Megatron-LM (https://github.com/NVIDIA/Megatron-LM/pull/5984, +measurements and numerics in https://github.com/NVIDIA/Megatron-LM/issues/5968). Covers: + + - numerics of the fused region vs an fp32-intermediate eager reference + (``dKV``/``dScore`` bit-identical, forward within one bf16 rounding step), vs the + upstream eager numerics (tolerance), and vs an fp64 oracle (fused error <= eager + error), over ragged THD packs including segments shorter than ``ratio``; + - static-capacity padding rows (``total_comp > cu_seqlens_comp[-1]``); + - kernel-side zero-writes to never-consumed ``dKV``/``dScore`` slots: NaN-canary + (uninitialized) gradient buffers stay bitwise-equal to zero-initialized runs and to + the eager reference, and the ``total_comp == 0`` host fallback still hands back + exact zeros; + - run-to-run determinism of forward / ``dKV`` / ``dScore`` (``dAPE`` uses fp32 atomics + and is exempt by design; the backward refuses to run under + ``torch.use_deterministic_algorithms(True)``); + - CUDA graph capture: warmup -> capture fwd+bwd -> replay (including replay with new + data and a smaller device-side true row count), and the loud error when the first + call for a configuration would JIT under capture; + - ``check_support`` boundaries (validated envelope: CC 10.0, ratio 4, coff 2, BF16 + kv/score, FP32 ape, int32 cu_seqlens and int32 flat-offset bounds). + +The eager reference below mirrors the exact region of Megatron-LM +``Compressor._forward_thd`` (non-pre-grouped THD path) that the fused kernels replace: +gather-index build -> gather -> ``+ APE`` -> overlap-window transform (``coff == 2``, +``Compressor._overlap_transform_thd``) -> fp32 softmax -> gated weighted sum -> bf16 +cast. ``mode`` selects the numerics: "upstream" reproduces the eager code exactly +(softmax weights rounded to bf16, bf16 multiply); "fp32" keeps all intermediates fp32 +with a single final bf16 rounding (the fused kernels' numerics); "fp64" is an oracle. +""" + +import pytest +import torch + + +def _import_compressor(): + # Skip only when the optional cutedsl dependency stack is missing; a broken + # cudnn.csa package itself must fail the tests, not skip them. + pytest.importorskip("cutlass", reason="Environment not supported: cudnn[cutedsl] not installed") + pytest.importorskip("cuda.bindings", reason="Environment not supported: cuda-python not installed") + from cudnn.csa import compressor + + return compressor + + +def _require_sm100(): + if not torch.cuda.is_available(): + pytest.skip("CUDA GPU required") + if torch.cuda.get_device_capability() != (10, 0): + pytest.skip("compute capability 10.0 GPU required") + + +# --------------------------------------------------------------------------- +# Eager reference (self-contained mirror of the Megatron-LM eager region) +# --------------------------------------------------------------------------- + + +def _batch_of_row(cu_seqlens, total): + """Segment index owning each packed row (mirror of Megatron-LM ``batch_of_row``).""" + n_seg = cu_seqlens.shape[0] - 1 + row_idx = torch.arange(total, device=cu_seqlens.device, dtype=torch.int64) + return torch.bucketize(row_idx, cu_seqlens[1:], right=True).clamp(max=max(n_seg - 1, 0)) + + +def _overlap_transform_thd(tensor, is_first_in_seg, head_dim, fill_value=0): + """Mirror of Megatron-LM ``Compressor._overlap_transform_thd``. + + Input shape: [total_comp, ratio, b, coff * head_dim] + Output shape: [total_comp, 2 * ratio, b, head_dim] + """ + n, ratio, b_dim, _ = tensor.size() + d = head_dim + new_tensor = tensor.new_full((n, 2 * ratio, b_dim, d), fill_value) + new_tensor[:, ratio:] = tensor[:, :, :, d:] + # Previous group's first-half data -- shift by 1 along dim-0. + prev_data = torch.roll(tensor[:, :, :, :d], shifts=1, dims=0) + # Zero-fill (or fill_value-fill) segment boundaries. + prev_data[is_first_in_seg] = fill_value + new_tensor[:, :ratio] = prev_data + return new_tensor + + +def _eager_pool(kv, score, ape, cu_seqlens, cu_seqlens_comp, total_comp, ratio, d, coff, mode): + device = kv.device + row_idx = torch.arange(total_comp, device=device, dtype=cu_seqlens_comp.dtype) + batch_ids = _batch_of_row(cu_seqlens_comp, total_comp) + valid_comp = row_idx < cu_seqlens_comp[-1] + local_pos = row_idx - cu_seqlens_comp[batch_ids] + local_pos = torch.where(valid_comp, local_pos, torch.zeros_like(local_pos)) + base = cu_seqlens[batch_ids].unsqueeze(1) + local_pos.unsqueeze(1) * ratio + base = torch.where(valid_comp.unsqueeze(1), base, torch.zeros_like(base)) + offsets = torch.arange(ratio, device=device, dtype=base.dtype).unsqueeze(0) + gather_idx = base + offsets # (total_comp, ratio) + + if mode == "fp32": + kv = kv.float() + score = score.float() + elif mode == "fp64": + kv = kv.double() + score = score.double() + ape = ape.double() + + kv_grouped = kv[gather_idx] # (total_comp, ratio, 1, coff * d) + score_grouped = score[gather_idx] + score_grouped = score_grouped + ape.view(1, ratio, 1, -1) + + if coff == 2: + is_first = local_pos == 0 + kv_grouped = _overlap_transform_thd(kv_grouped, is_first, d, fill_value=0) + score_grouped = _overlap_transform_thd(score_grouped, is_first, d, fill_value=float("-inf")) + + if mode == "upstream": + weights = torch.softmax(score_grouped, dim=1, dtype=torch.float32).to(kv_grouped.dtype) + out = (kv_grouped * weights).sum(dim=1) + elif mode == "fp32": + weights = torch.softmax(score_grouped, dim=1, dtype=torch.float32) + out = (kv_grouped * weights).sum(dim=1).to(torch.bfloat16) + else: # fp64 oracle + weights = torch.softmax(score_grouped, dim=1, dtype=torch.float64) + out = (kv_grouped * weights).sum(dim=1) + return out # (total_comp, 1, d) + + +# --------------------------------------------------------------------------- +# Input construction and runners +# --------------------------------------------------------------------------- + + +def _make_inputs(lens, d, ratio, coff, seed=1234, device="cuda"): + total = sum(lens) + w = coff * d + gen = torch.Generator(device="cpu").manual_seed(seed) + kv = torch.randn(total, 1, w, generator=gen, dtype=torch.float32).to(torch.bfloat16) + score = (torch.randn(total, 1, w, generator=gen, dtype=torch.float32).mul_(1.5)).to(torch.bfloat16) + ape = torch.randn(ratio, w, generator=gen, dtype=torch.float32).mul_(0.25) + cu = torch.tensor([0] + list(torch.tensor(lens).cumsum(0)), dtype=torch.int32, device=device) + seg_comp = torch.tensor([seg_len // ratio for seg_len in lens]) + cuc = torch.tensor([0] + list(seg_comp.cumsum(0)), dtype=torch.int32, device=device) + total_comp = int(cuc[-1].item()) + go = torch.randn(total_comp, 1, d, generator=gen, dtype=torch.float32).to(torch.bfloat16) + return kv.to(device), score.to(device), ape.to(device), cu, cuc, total_comp, go.to(device) + + +def _run_eager(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go, mode): + """Forward + backward through the eager reference; returns (out, dKV, dScore, dAPE).""" + dtype = torch.float64 if mode == "fp64" else None + kv_l = (kv.to(dtype) if dtype else kv.clone()).requires_grad_(True) + score_l = (score.to(dtype) if dtype else score.clone()).requires_grad_(True) + ape_l = (ape.to(dtype) if dtype else ape.clone()).requires_grad_(True) + out = _eager_pool(kv_l, score_l, ape_l, cu, cuc, total_comp, ratio, d, coff, mode) + out.backward(go.to(out.dtype)) + torch.cuda.synchronize() + return out.detach(), kv_l.grad.detach(), score_l.grad.detach(), ape_l.grad.detach() + + +def _run_fused(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go): + """Forward + backward through the fused wrappers; returns (out, dKV, dScore, dAPE).""" + compressor = _import_compressor() + total = kv.shape[0] + out = compressor.csa_compressor_forward_wrapper( + kv.view(total, -1), + score.view(total, -1), + ape, + cu, + cuc, + ratio=ratio, + head_dim=d, + coff=coff, + total_comp=total_comp, + )["out"] + grads = compressor.csa_compressor_backward_wrapper( + kv.view(total, -1), + score.view(total, -1), + ape, + cu, + cuc, + go.view(total_comp, d), + ratio=ratio, + head_dim=d, + coff=coff, + ) + torch.cuda.synchronize() + return ( + out.view(total_comp, 1, d), + grads["grad_kv"].view_as(kv), + grads["grad_score"].view_as(score), + grads["grad_ape"], + ) + + +_SHAPES = [ + # (lens, head_dim, ratio, coff) + pytest.param([2048], 128, 4, 2, id="b1-d128-r4"), + pytest.param([1023, 2048, 509], 128, 4, 2, id="ragged3-d128-r4"), + pytest.param([2048], 512, 4, 2, id="b1-d512-r4"), + pytest.param([3, 515, 1024, 129], 128, 4, 2, id="short-seg-d128-r4"), + # odd head_dim exercises the scalar (vec == 1) forward layout; even head_dims all + # take the vectorized (vec == 2) one. + pytest.param([260], 65, 4, 2, id="b1-d65-odd-r4"), +] + + +# --------------------------------------------------------------------------- +# Numerics +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +@pytest.mark.parametrize("lens,d,ratio,coff", _SHAPES) +def test_numerics_vs_references(lens, d, ratio, coff): + """Fused fwd+bwd vs fp32-eager (bitwise dKV/dScore), upstream eager, and fp64 oracle.""" + _require_sm100() + kv, score, ape, cu, cuc, total_comp, go = _make_inputs(lens, d, ratio, coff) + + r_fused = _run_fused(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go) + r_fp32 = _run_eager(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go, mode="fp32") + r_up = _run_eager(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go, mode="upstream") + r_fp64 = _run_eager(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go, mode="fp64") + + # vs fp32-intermediate eager reference (the fused kernels' numerics contract): + # dKV / dScore bit-identical; forward within one bf16 rounding step on a tiny + # fraction of elements; dAPE within fp32 atomics reorder noise. + assert torch.equal(r_fused[1], r_fp32[1]), "dKV must be bit-identical to the fp32 reference" + assert torch.equal(r_fused[2], r_fp32[2]), "dScore must be bit-identical to the fp32 reference" + fwd_diff = (r_fused[0].float() - r_fp32[0].float()).abs() + n_diff = (r_fused[0] != r_fp32[0]).sum().item() + assert n_diff <= max(1, int(0.001 * r_fused[0].numel())), n_diff + assert fwd_diff.max().item() <= 1.6e-2 + assert (r_fused[3] - r_fp32[3]).abs().max().item() <= 1e-3 + + # vs the verbatim upstream eager numerics: not bit-identical (the eager path rounds + # softmax weights to bf16 and multiplies in bf16), but close. + for fused_t, up_t in zip(r_fused, r_up): + assert torch.allclose(fused_t.float(), up_t.float(), rtol=0, atol=0.1) + + # vs the fp64 oracle: the fused kernel is at least as accurate as the eager code on + # every output. + for i in range(4): + err_fused = (r_fused[i].double() - r_fp64[i].double()).abs().max().item() + err_up = (r_up[i].double() - r_fp64[i].double()).abs().max().item() + assert err_fused <= err_up * (1 + 1e-6) + 1e-4, (i, err_fused, err_up) + + +@pytest.mark.L0 +def test_replay_determinism(): + """Forward, dKV and dScore replay bitwise identically run to run (dAPE is exempt).""" + _require_sm100() + kv, score, ape, cu, cuc, total_comp, go = _make_inputs([1023, 2048, 509], 128, 4, 2) + runs = [_run_fused(kv, score, ape, cu, cuc, total_comp, 4, 128, 2, go) for _ in range(3)] + for other in runs[1:]: + assert torch.equal(runs[0][0], other[0]) + assert torch.equal(runs[0][1], other[1]) + assert torch.equal(runs[0][2], other[2]) + # dAPE is accumulated with fp32 atomics; equality is not guaranteed, closeness is. + assert torch.allclose(runs[0][3], other[3], rtol=0, atol=1e-3) + + +_PADDING_SHAPES = [ + # (lens, head_dim, ratio, coff, pad); the first case has a LEADING segment shorter + # than ratio (0 compressed blocks), so padding rows gather tokens [0, ratio) that + # span a segment boundary -- exactly like the eager gather. + ([3, 515, 1024, 129], 128, 4, 2, 8), + ([1023, 2048, 509], 128, 4, 2, 8), +] + + +@pytest.mark.L0 +@pytest.mark.parametrize("lens,d,ratio,coff,pad", _PADDING_SHAPES) +def test_static_capacity_padding(lens, d, ratio, coff, pad): + """Static-capacity padding rows: eager-matching forward, ignored padding gradients.""" + _require_sm100() + kv, score, ape, cu, cuc, total_true, _ = _make_inputs(lens, d, ratio, coff) + capacity = total_true + pad + gen = torch.Generator(device="cpu").manual_seed(7) + go = torch.randn(capacity, 1, d, generator=gen, dtype=torch.float32) + go = go.to(torch.bfloat16).cuda() + go_zero_pad = go.clone() + go_zero_pad[total_true:] = 0 + + r_fused = _run_fused(kv, score, ape, cu, cuc, capacity, ratio, d, coff, go) + r_fp32 = _run_eager(kv, score, ape, cu, cuc, capacity, ratio, d, coff, go_zero_pad, mode="fp32") + + # Forward: padding rows replicate row 0's window exactly like the eager code, so the + # full padded output (valid + padding rows) obeys the same criteria as the unpadded + # comparison. + assert (r_fused[0] != r_fp32[0]).sum().item() <= max(1, int(0.001 * r_fused[0].numel())) + assert (r_fused[0].float() - r_fp32[0].float()).abs().max().item() <= 1.6e-2 + + # Backward: incoming gradients on padding rows are ignored by design -- the fused + # gradients (computed with NONZERO padding-row grads) match the eager reference run + # with zeroed padding-row grads bit-for-bit on dKV/dScore. + assert torch.equal(r_fused[1], r_fp32[1]) + assert torch.equal(r_fused[2], r_fp32[2]) + assert (r_fused[3] - r_fp32[3]).abs().max().item() <= 1e-3 + + # And explicitly: nonzero vs zero padding-row grads produce identical fused grads. + r_fused_zero = _run_fused(kv, score, ape, cu, cuc, capacity, ratio, d, coff, go_zero_pad) + assert torch.equal(r_fused[1], r_fused_zero[1]) + assert torch.equal(r_fused[2], r_fused_zero[2]) + + +@pytest.mark.L0 +def test_empty_output(): + """total_comp == 0 launches nothing and returns well-formed empty/zero tensors.""" + _require_sm100() + compressor = _import_compressor() + d, ratio, coff = 128, 4, 2 + w = coff * d + kv = torch.randn(2, w, device="cuda").to(torch.bfloat16) + score = torch.randn(2, w, device="cuda").to(torch.bfloat16) + ape = torch.randn(ratio, w, device="cuda") + cu = torch.tensor([0, 2], dtype=torch.int32, device="cuda") + cuc = torch.tensor([0, 0], dtype=torch.int32, device="cuda") + out = compressor.csa_compressor_forward_wrapper(kv, score, ape, cu, cuc, ratio=ratio, head_dim=d, coff=coff, total_comp=0)["out"] + assert out.shape == (0, d) and out.dtype == torch.bfloat16 + grads = compressor.csa_compressor_backward_wrapper(kv, score, ape, cu, cuc, out, ratio=ratio, head_dim=d, coff=coff) + assert grads["grad_kv"].abs().sum().item() == 0 + assert grads["grad_score"].abs().sum().item() == 0 + assert grads["grad_ape"].abs().sum().item() == 0 + + +@pytest.mark.L0 +def test_backward_wrapper_zeros_when_no_blocks(): + """Multi-segment packs where NO segment reaches ratio tokens: total_comp == 0, so + the kernel cannot launch and the wrapper must hand back exact-zero grads (the + host-side zeros fallback behind the uninitialized-buffer optimization).""" + _require_sm100() + compressor = _import_compressor() + d, ratio, coff = 128, 4, 2 + w = coff * d + lens = [3, 2, 1] + total = sum(lens) + kv = torch.randn(total, w, device="cuda").to(torch.bfloat16) + score = torch.randn(total, w, device="cuda").to(torch.bfloat16) + ape = torch.randn(ratio, w, device="cuda") + cu = torch.tensor([0, 3, 5, 6], dtype=torch.int32, device="cuda") + cuc = torch.tensor([0, 0, 0, 0], dtype=torch.int32, device="cuda") + go = torch.empty(0, d, dtype=torch.bfloat16, device="cuda") + grads = compressor.csa_compressor_backward_wrapper(kv, score, ape, cu, cuc, go, ratio=ratio, head_dim=d, coff=coff) + assert grads["grad_kv"].shape == kv.shape and grads["grad_kv"].abs().sum().item() == 0 + assert grads["grad_score"].shape == score.shape and grads["grad_score"].abs().sum().item() == 0 + assert grads["grad_ape"].abs().sum().item() == 0 + + +_CANARY_SHAPES = [ + # (lens, head_dim, pad, tok_pad) — every never-consumed dKV/dScore slot class must + # be hit: segment tails (seqlen % ratio), the last block's first-half columns, whole + # segments shorter than ratio (zero blocks), static-capacity padding rows + # (pad > 0 extra grad_out rows), and static token-capacity padding of the gradient + # buffers themselves (tok_pad > 0 tokens beyond cu_seqlens[-1]). + pytest.param([2048], 128, 0, 0, id="b1-d128"), + pytest.param([1023, 2048, 509], 128, 0, 0, id="ragged3-d128"), + pytest.param([3, 515, 1024, 129], 128, 0, 0, id="short-seg-d128"), + pytest.param([5, 6, 7], 128, 0, 0, id="all-tiny-d128"), + pytest.param([1023, 2048, 509], 512, 0, 0, id="ragged3-d512"), + pytest.param([3, 515, 1024, 129], 128, 8, 0, id="short-seg-d128-padded"), + pytest.param([1023, 2048, 509], 128, 8, 0, id="ragged3-d128-padded"), + pytest.param([1023, 2048, 509], 128, 0, 37, id="ragged3-d128-tokpad"), + pytest.param([3, 515, 1024, 129], 128, 8, 21, id="short-seg-d128-padded-tokpad"), +] + + +@pytest.mark.L0 +@pytest.mark.parametrize("lens,d,pad,tok_pad", _CANARY_SHAPES) +def test_backward_fills_uninitialized_buffers(lens, d, pad, tok_pad): + """NaN-canary: the backward kernel fully overwrites garbage dKV/dScore buffers. + + The kernel writes exact zeros to every never-consumed slot itself (there are no + separate zero-fill kernels anymore), so running it into NaN-poisoned buffers must + produce bitwise the same dKV/dScore as running it into zero-initialized buffers. + """ + _require_sm100() + compressor = _import_compressor() + ratio, coff = 4, 2 + kv, score, ape, cu, cuc, total_true, _ = _make_inputs(lens, d, ratio, coff) + total_comp = total_true + pad + gen = torch.Generator(device="cpu").manual_seed(11) + go = torch.randn(total_comp, d, generator=gen, dtype=torch.float32).to(torch.bfloat16).cuda() + total = kv.shape[0] + tok_pad + kv2 = torch.cat([kv.view(kv.shape[0], -1), torch.randn(tok_pad, coff * d, generator=gen, dtype=torch.float32).to(torch.bfloat16).cuda()]) + score2 = torch.cat([score.view(score.shape[0], -1), torch.randn(tok_pad, coff * d, generator=gen, dtype=torch.float32).to(torch.bfloat16).cuda()]) + + bwd = compressor.CSACompressorBackward( + sample_kv=kv2, + sample_score=score2, + sample_ape=ape, + sample_cu_seqlens=cu, + sample_cu_seqlens_comp=cuc, + sample_out=torch.empty(total_comp, d, dtype=torch.bfloat16, device="meta"), + ratio=ratio, + coff=coff, + ) + assert bwd.check_support() + bwd.compile() + + def run(poison): + grad_kv = torch.empty_like(kv2) + grad_score = torch.empty_like(score2) + if poison: + grad_kv.fill_(float("nan")) + grad_score.fill_(float("nan")) + else: + grad_kv.zero_() + grad_score.zero_() + grad_ape = torch.zeros_like(ape) + bwd.execute(kv2, score2, ape, cu, cuc, go, grad_kv, grad_score, grad_ape) + torch.cuda.synchronize() + return grad_kv, grad_score, grad_ape + + gkv_ref, gs_ref, gape_ref = run(poison=False) + gkv_nan, gs_nan, gape_nan = run(poison=True) + assert not torch.isnan(gkv_nan).any(), "unwritten dKV slots survived (NaN canary)" + assert not torch.isnan(gs_nan).any(), "unwritten dScore slots survived (NaN canary)" + assert torch.equal(gkv_nan, gkv_ref) + assert torch.equal(gs_nan, gs_ref) + assert torch.allclose(gape_nan, gape_ref, rtol=0, atol=1e-3) + + # And the zero-slot pattern matches autograd: never-consumed slots are exact zeros, + # bitwise as the fp32 eager reference computes them. (The fused backward ignores + # incoming gradients on static-capacity padding rows by design, so the eager + # reference runs with those rows zeroed.) + go_ref = go.clone() + go_ref[total_true:] = 0 + r_fp32 = _run_eager( + kv2.view(total, 1, -1), + score2.view(total, 1, -1), + ape, + cu, + cuc, + total_comp, + ratio, + d, + coff, + go_ref.view(total_comp, 1, d), + mode="fp32", + ) + assert torch.equal(gkv_nan.view_as(r_fp32[1]), r_fp32[1]) + assert torch.equal(gs_nan.view_as(r_fp32[2]), r_fp32[2]) + + +# --------------------------------------------------------------------------- +# Deterministic mode +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_backward_rejects_deterministic_mode(): + """The backward raises under torch.use_deterministic_algorithms (dAPE fp32 atomics).""" + _require_sm100() + compressor = _import_compressor() + kv, score, ape, cu, cuc, total_comp, go = _make_inputs([512, 256], 128, 4, 2) + total = kv.shape[0] + # Forward is deterministic and keeps working. + prev_det = torch.are_deterministic_algorithms_enabled() + prev_warn = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True, warn_only=False) + try: + out = compressor.csa_compressor_forward_wrapper( + kv.view(total, -1), score.view(total, -1), ape, cu, cuc, ratio=4, head_dim=128, coff=2, total_comp=total_comp + )["out"] + assert out.shape == (total_comp, 128) + with pytest.raises(RuntimeError, match="not deterministic"): + compressor.csa_compressor_backward_wrapper( + kv.view(total, -1), score.view(total, -1), ape, cu, cuc, go.view(total_comp, 128), ratio=4, head_dim=128, coff=2 + ) + finally: + torch.use_deterministic_algorithms(prev_det, warn_only=prev_warn) + + +@pytest.mark.L0 +def test_backward_warns_in_warn_only_deterministic_mode(): + """warn_only deterministic mode warns (torch semantics) and still runs the backward.""" + _require_sm100() + compressor = _import_compressor() + kv, score, ape, cu, cuc, total_comp, go = _make_inputs([512, 256], 128, 4, 2) + r_ref = _run_fused(kv, score, ape, cu, cuc, total_comp, 4, 128, 2, go) + prev_det = torch.are_deterministic_algorithms_enabled() + prev_warn = torch.is_deterministic_algorithms_warn_only_enabled() + torch.use_deterministic_algorithms(True, warn_only=True) + try: + with pytest.warns(RuntimeWarning, match="not deterministic"): + grads = compressor.csa_compressor_backward_wrapper( + kv.view(kv.shape[0], -1), score.view(score.shape[0], -1), ape, cu, cuc, go.view(total_comp, 128), ratio=4, head_dim=128, coff=2 + ) + finally: + torch.use_deterministic_algorithms(prev_det, warn_only=prev_warn) + assert torch.equal(grads["grad_kv"].view_as(r_ref[1]), r_ref[1]) + assert torch.equal(grads["grad_score"].view_as(r_ref[2]), r_ref[2]) + + +# --------------------------------------------------------------------------- +# CUDA graph capture +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +@pytest.mark.filterwarnings("ignore::UserWarning") +def test_cuda_graph_capture(): + """Warmup -> capture fwd+bwd -> replay; JIT under capture raises a clear error.""" + _require_sm100() + compressor = _import_compressor() + ratio, d, coff = 4, 128, 2 + lens = [512, 256] + kv, score, ape, cu, cuc, total_true, _ = _make_inputs(lens, d, ratio, coff) + capacity = total_true + 8 # static capacity, as with CUDA-graph static shapes + total = kv.shape[0] + + kv_s = kv.view(total, -1).clone() + score_s = score.view(total, -1).clone() + ape_s = ape.clone() + go_s = torch.zeros(capacity, d, device="cuda", dtype=torch.bfloat16) + go_s[:total_true] = torch.randn(total_true, d, device="cuda").to(torch.bfloat16) + + def _fused_fwd_bwd(): + out = compressor.csa_compressor_forward_wrapper(kv_s, score_s, ape_s, cu, cuc, ratio=ratio, head_dim=d, coff=coff, total_comp=capacity)["out"] + grads = compressor.csa_compressor_backward_wrapper(kv_s, score_s, ape_s, cu, cuc, go_s, ratio=ratio, head_dim=d, coff=coff) + return out, grads + + # One warmup per configuration on a side stream (JIT-compiles both kernels). + side = torch.cuda.Stream() + side.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side): + _fused_fwd_bwd() + torch.cuda.current_stream().wait_stream(side) + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + out_c, grads_c = _fused_fwd_bwd() + gkv_c, gscore_c, gape_c = grads_c["grad_kv"], grads_c["grad_score"], grads_c["grad_ape"] + + # Replay on the same data must reproduce the direct (non-captured) fused results + # bitwise on forward/dKV/dScore. + graph.replay() + torch.cuda.synchronize() + ref_out, ref_grads = _fused_fwd_bwd() + torch.cuda.synchronize() + assert torch.equal(out_c, ref_out) + assert torch.equal(gkv_c, ref_grads["grad_kv"]) + assert torch.equal(gscore_c, ref_grads["grad_score"]) + assert torch.allclose(gape_c, ref_grads["grad_ape"], rtol=0, atol=1e-3) + + # Replay with new data and a SMALLER device-side true row count (the fixed capacity + # stays static, cu/cuc contents change) -- the graph-replayed gradients must match + # the fp32 eager reference bitwise on dKV/dScore. + lens2 = [384, 128] + kv2, score2, ape2, cu2, cuc2, total2, _ = _make_inputs(lens2, d, ratio, coff, seed=99) + n2 = sum(lens2) + kv_s.zero_() + score_s.zero_() + kv_s[:n2] = kv2.view(n2, -1) + score_s[:n2] = score2.view(n2, -1) + ape_s.copy_(ape2) + cu.copy_(cu2) + cuc.copy_(cuc2) + go_s.zero_() + go_s[:total2] = torch.randn(total2, d, device="cuda").to(torch.bfloat16) + graph.replay() + torch.cuda.synchronize() + r_fp32 = _run_eager( + kv_s.view(-1, 1, coff * d), + score_s.view(-1, 1, coff * d), + ape_s, + cu, + cuc, + capacity, + ratio, + d, + coff, + go_s.view(capacity, 1, d), + mode="fp32", + ) + assert torch.equal(gkv_c.view_as(r_fp32[1]), r_fp32[1]) + assert torch.equal(gscore_c.view_as(r_fp32[2]), r_fp32[2]) + + # A first call for a NEW configuration under capture must raise loudly instead of + # JIT-compiling (which is not capture-safe). head_dim 192 is used by no other test + # in this module, so this configuration is guaranteed to be uncompiled regardless of + # test execution order. + d_new = 192 + kv3 = torch.randn(256, coff * d_new, device="cuda").to(torch.bfloat16) + score3 = torch.randn(256, coff * d_new, device="cuda").to(torch.bfloat16) + ape3 = torch.randn(ratio, coff * d_new, device="cuda") + cu3 = torch.tensor([0, 256], dtype=torch.int32, device="cuda") + cuc3 = torch.tensor([0, 64], dtype=torch.int32, device="cuda") + graph2 = torch.cuda.CUDAGraph() + with pytest.raises(RuntimeError, match="CUDA graph capture"): + with torch.cuda.graph(graph2): + compressor.csa_compressor_forward_wrapper(kv3, score3, ape3, cu3, cuc3, ratio=ratio, head_dim=d_new, coff=coff, total_comp=64) + # The CUDA context must remain usable after the aborted capture. + torch.cuda.synchronize() + probe = torch.ones(8, device="cuda") + assert probe.sum().item() == 8 + + +# --------------------------------------------------------------------------- +# check_support boundaries +# --------------------------------------------------------------------------- + + +def _meta(shape, dtype, stride=None): + """Metadata-only sample tensor (meta device) for check_support tests.""" + if stride is None: + stride = [] + acc = 1 + for s in reversed(shape): + stride.append(acc) + acc *= s + stride = tuple(reversed(stride)) + return torch.empty_strided(shape, stride, dtype=dtype, device="meta") + + +def _meta_samples( + total=512, + d=128, + ratio=4, + coff=2, + n_seg=2, + total_comp=None, + kv_dtype=torch.bfloat16, + ape_dtype=torch.float32, + cu_dtype=torch.int32, + out_dtype=torch.bfloat16, + score_shape=None, + cuc_len=None, + kv_stride=None, +): + w = coff * d + if total_comp is None: + total_comp = total // ratio + kv = _meta((total, w), kv_dtype, stride=kv_stride) + score = _meta(score_shape or (total, w), kv_dtype) + ape = _meta((ratio, w), ape_dtype) + cu = _meta((n_seg + 1,), cu_dtype) + cuc = _meta((cuc_len or (n_seg + 1),), cu_dtype) + out = _meta((total_comp, d), out_dtype) + return dict( + sample_kv=kv, + sample_score=score, + sample_ape=ape, + sample_cu_seqlens=cu, + sample_cu_seqlens_comp=cuc, + sample_out=out, + ) + + +@pytest.mark.L0 +def test_check_support_accepts_envelope(): + """Metadata-only samples inside the validated envelope pass check_support.""" + _require_sm100() + compressor = _import_compressor() + for cls in (compressor.CSACompressorForward, compressor.CSACompressorBackward): + api = cls(**_meta_samples(), ratio=4, coff=2) + assert api.check_support() is True + assert api.head_dim == 128 and api.total_tokens == 512 and api.total_comp == 128 + + +@pytest.mark.L0 +@pytest.mark.parametrize( + "kwargs,ctor,match", + [ + (dict(), dict(ratio=128, coff=1), "validated for ratio=4, coff=2"), + (dict(), dict(ratio=8, coff=2), "validated for ratio=4, coff=2"), + (dict(kv_dtype=torch.float16), dict(), "kv"), + (dict(ape_dtype=torch.bfloat16), dict(), "ape"), + (dict(cu_dtype=torch.int64), dict(), "cu_seqlens"), + (dict(out_dtype=torch.float32), dict(), "out"), + (dict(score_shape=(512, 128)), dict(), "score shape"), + (dict(cuc_len=4), dict(), "B \\+ 1"), + (dict(total=2**25, d=128), dict(), "int32 flat offsets"), + (dict(total=4, d=8388482, total_comp=1), dict(), "head_dim"), + (dict(total=2, total_comp=1), dict(), "requires at least ratio"), + (dict(kv_stride=(512, 2)), dict(), "contiguous"), + ], +) +def test_check_support_rejects(kwargs, ctor, match): + """check_support raises ValueError for configurations outside the envelope.""" + _require_sm100() + compressor = _import_compressor() + samples = _meta_samples(**kwargs) + api = compressor.CSACompressorForward(**samples, ratio=ctor.get("ratio", 4), coff=ctor.get("coff", 2)) + with pytest.raises(ValueError, match=match): + api.check_support() + + +@pytest.mark.L0 +def test_check_support_rejects_cpu_tensors(): + """CPU sample tensors are rejected with a clear error.""" + _require_sm100() + compressor = _import_compressor() + samples = _meta_samples() + samples["sample_kv"] = torch.empty(512, 256, dtype=torch.bfloat16, device="cpu") + api = compressor.CSACompressorForward(**samples, ratio=4, coff=2) + with pytest.raises(ValueError, match="CUDA"): + api.check_support() + + +# --------------------------------------------------------------------------- +# Class API vs wrapper equivalence +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_class_api_matches_wrapper(): + """The explicit class API produces bitwise-identical results to the wrappers.""" + _require_sm100() + compressor = _import_compressor() + ratio, d, coff = 4, 128, 2 + kv, score, ape, cu, cuc, total_comp, go = _make_inputs([1023, 509], d, ratio, coff) + total = kv.shape[0] + kv2, score2, go2 = kv.view(total, -1), score.view(total, -1), go.view(total_comp, d) + + r_wrapped = _run_fused(kv, score, ape, cu, cuc, total_comp, ratio, d, coff, go) + + fwd = compressor.CSACompressorForward( + sample_kv=kv2, + sample_score=score2, + sample_ape=ape, + sample_cu_seqlens=cu, + sample_cu_seqlens_comp=cuc, + sample_out=torch.empty(total_comp, d, dtype=torch.bfloat16, device="meta"), + ratio=ratio, + coff=coff, + ) + assert fwd.check_support() + fwd.compile() + out = torch.empty(total_comp, d, dtype=torch.bfloat16, device="cuda") + fwd.execute(kv2, score2, ape, cu, cuc, out) + + bwd = compressor.CSACompressorBackward( + sample_kv=kv2, + sample_score=score2, + sample_ape=ape, + sample_cu_seqlens=cu, + sample_cu_seqlens_comp=cuc, + sample_out=torch.empty(total_comp, d, dtype=torch.bfloat16, device="meta"), + ratio=ratio, + coff=coff, + ) + assert bwd.check_support() + bwd.compile() + grad_kv = torch.zeros_like(kv2) + grad_score = torch.zeros_like(score2) + grad_ape = torch.zeros_like(ape) + bwd.execute(kv2, score2, ape, cu, cuc, go2, grad_kv, grad_score, grad_ape) + torch.cuda.synchronize() + + assert torch.equal(out.view_as(r_wrapped[0]), r_wrapped[0]) + assert torch.equal(grad_kv.view_as(r_wrapped[1]), r_wrapped[1]) + assert torch.equal(grad_score.view_as(r_wrapped[2]), r_wrapped[2]) + assert torch.allclose(grad_ape, r_wrapped[3], rtol=0, atol=1e-3) + + +# --------------------------------------------------------------------------- +# Runtime hazards: alignment, explicit streams, multi-device +# --------------------------------------------------------------------------- + + +@pytest.mark.L0 +def test_misaligned_base_pointer_rejected(): + """Contiguous storage-offset views with unaligned base pointers raise ValueError.""" + _require_sm100() + compressor = _import_compressor() + d, ratio, coff = 128, 4, 2 + w = coff * d + total = 512 + buf = torch.randn(total * w + 8, device="cuda").to(torch.bfloat16) + kv = buf[2 : 2 + total * w].view(total, w) # contiguous, 4-byte-aligned base + assert kv.is_contiguous() and kv.data_ptr() % 16 != 0 + score = torch.randn(total, w, device="cuda").to(torch.bfloat16) + ape = torch.randn(ratio, w, device="cuda") + cu = torch.tensor([0, total], dtype=torch.int32, device="cuda") + cuc = torch.tensor([0, total // ratio], dtype=torch.int32, device="cuda") + with pytest.raises(ValueError, match="aligned"): + compressor.csa_compressor_forward_wrapper(kv, score, ape, cu, cuc, ratio=ratio, head_dim=d, coff=coff, total_comp=total // ratio) + + +@pytest.mark.L0 +def test_explicit_stream_input_lifetime(): + """Inputs released right after an explicit-stream call are not recycled early. + + The launch path takes raw pointers, so the wrapper must ``record_stream`` its + operands on the external stream; otherwise the caching allocator can hand the freed + input storage to allocations on another stream while the kernel is still pending. + """ + _require_sm100() + compressor = _import_compressor() + if not hasattr(torch.cuda, "_sleep"): + pytest.skip("torch.cuda._sleep not available") + import cuda.bindings.driver as cuda_driver + + d, ratio, coff = 128, 4, 2 + w = coff * d + lens = [4096] + kv, score, ape, cu, cuc, total_comp, _ = _make_inputs(lens, d, ratio, coff) + total = kv.shape[0] + kv2, score2 = kv.view(total, w).contiguous(), score.view(total, w).contiguous() + + # Ground truth on the default stream (from private clones). + expected = compressor.csa_compressor_forward_wrapper( + kv2.clone(), score2.clone(), ape.clone(), cu.clone(), cuc.clone(), ratio=ratio, head_dim=d, coff=coff, total_comp=total_comp + )["out"] + torch.cuda.synchronize() + + side = torch.cuda.Stream() + ext = cuda_driver.CUstream(side.cuda_stream) + with torch.cuda.stream(side): + torch.cuda._sleep(int(5e8)) # block the side stream so the kernel stays pending + out = compressor.csa_compressor_forward_wrapper(kv2, score2, ape, cu, cuc, ratio=ratio, head_dim=d, coff=coff, total_comp=total_comp, stream=ext)["out"] + # Drop every caller reference while the kernel is still queued behind the sleep, + # then try hard to get the freed storages reallocated and scribbled on the default + # (idle) stream. + del kv, score, kv2, score2, ape, cu, cuc + junk = [torch.full((total, w), 7.0, device="cuda", dtype=torch.bfloat16) for _ in range(4)] + junk.append(torch.full((ratio, w), 7.0, device="cuda", dtype=torch.float32)) + junk.append(torch.full((64,), 7, device="cuda", dtype=torch.int32)) + torch.cuda.synchronize() + assert torch.equal(out, expected), "explicit-stream inputs were recycled before the kernel consumed them" + + +@pytest.mark.L0 +def test_multi_device_launch(): + """Tensors on a non-current device produce correct results (device anchoring).""" + _require_sm100() + compressor = _import_compressor() + if torch.cuda.device_count() < 2: + pytest.skip("needs >= 2 visible GPUs") + if torch.cuda.get_device_capability(1) != (10, 0): + pytest.skip("second GPU is not CC 10.0") + d, ratio, coff = 128, 4, 2 + kv, score, ape, cu, cuc, total_comp, go = _make_inputs([1024], d, ratio, coff, device="cuda:1") + total = kv.shape[0] + kv2, score2, go2 = kv.view(total, -1), score.view(total, -1), go.view(total_comp, d) + + assert torch.cuda.current_device() == 0 # launch with a FOREIGN current device + out_foreign = compressor.csa_compressor_forward_wrapper(kv2, score2, ape, cu, cuc, ratio=ratio, head_dim=d, coff=coff, total_comp=total_comp)["out"] + grads_foreign = compressor.csa_compressor_backward_wrapper(kv2, score2, ape, cu, cuc, go2, ratio=ratio, head_dim=d, coff=coff) + torch.cuda.synchronize(torch.device("cuda", 1)) + + with torch.cuda.device(1): + out_native = compressor.csa_compressor_forward_wrapper(kv2, score2, ape, cu, cuc, ratio=ratio, head_dim=d, coff=coff, total_comp=total_comp)["out"] + grads_native = compressor.csa_compressor_backward_wrapper(kv2, score2, ape, cu, cuc, go2, ratio=ratio, head_dim=d, coff=coff) + torch.cuda.synchronize() + + assert out_foreign.device == torch.device("cuda", 1) + assert out_foreign.abs().sum().item() > 0 # the historic failure mode was all-zeros + assert torch.equal(out_foreign, out_native) + assert torch.equal(grads_foreign["grad_kv"], grads_native["grad_kv"]) + assert torch.equal(grads_foreign["grad_score"], grads_native["grad_score"])