From 241b7cec48cbbcafbe10ae2d33818851dc3b2e44 Mon Sep 17 00:00:00 2001 From: zky Date: Fri, 24 Jul 2026 04:47:51 +0000 Subject: [PATCH 1/2] Dispatch Compressor._forward_thd to the cudnn-frontend fused CSA compressor The fused forward+backward kernels for the gated-pooling region of Compressor._forward_thd now live in the cudnn-frontend Python package (cudnn.csa.compressor, cudnn-frontend PR #427). Megatron-side this is a thin additive dispatch: maybe_compress_thd_fused() returns the pooled tensor when the frontend is importable and the configuration is supported (THD non-pre-grouped path, compress_ratio == 4, bf16 kv/score, fp32 ape, compute capability 10.0, int32 flat offsets) and None otherwise, in which case the eager region runs unchanged. SBHD and the pre-grouped CP-prep path are untouched. Frontend availability is probed by importing the concrete entry points (csa_compressor_forward_wrapper / csa_compressor_backward_wrapper), not by version comparison: nvidia-cudnn-frontend installs that predate the CSA compressor API lack cudnn.csa and keep the eager path, as does any import failure (e.g. the DSL extra missing). The dispatch also keeps the eager path under torch.use_deterministic_algorithms(True) (dAPE fp32 atomics) and under torch.compile tracing (raw-pointer launch path is not traceable). MCORE_CSA_FUSED_COMPRESSOR=0 disables the dispatch entirely. compress_ratio == 128 is functionally supported by the frontend kernels but stays on the eager path until tuned (see issue #5968). Addresses #5968. Signed-off-by: zky --- .../experimental_attention_variant/csa.py | 96 ++++--- .../csa_fused_compressor.py | 241 ++++++++++++++++++ 2 files changed, 300 insertions(+), 37 deletions(-) create mode 100644 megatron/core/transformer/experimental_attention_variant/csa_fused_compressor.py diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index e2b409290b4..a50585aa1cb 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -17,6 +17,9 @@ from megatron.core.transformer.enums import AttnMaskType from megatron.core.transformer.experimental_attention_variant import csa_cp_layout_kernels from megatron.core.transformer.experimental_attention_variant import csa_cp_utils as cp_utils +from megatron.core.transformer.experimental_attention_variant.csa_fused_compressor import ( + maybe_compress_thd_fused, +) from megatron.core.transformer.experimental_attention_variant.csa_kernels import ( FusedCSAIndexerSparseAttnFromTopkFunc, batch_of_row, @@ -1117,47 +1120,66 @@ def _forward_thd( kv, _ = self.linear_wkv(x) # (total, 1, coff * head_dim) score, _ = self.linear_wgate(x) # (total, 1, coff * head_dim) - if pre_grouped: - # Compressor-prep already groups rows as ``[g * ratio, (g + 1) * ratio)``. - kv_grouped = kv.reshape(total_comp, ratio, 1, -1) - score_grouped = score.reshape(total_comp, ratio, 1, -1) - local_pos = None - else: - # Build gather index: (total_comp, ratio). ``total_comp`` can be a - # static capacity for CUDA graph capture, so rows beyond the true - # ``cu_seqlens_compressed[-1]`` are mapped to a safe source row and left - # as tail padding by downstream index lowering. - row_idx = torch.arange(total_comp, device=device, dtype=cu_seqlens_compressed.dtype) - batch_ids = batch_of_row(cu_seqlens_compressed, total_q=total_comp) - valid_comp = row_idx < cu_seqlens_compressed[-1] - local_pos = row_idx - cu_seqlens_compressed[batch_ids] - local_pos = torch.where(valid_comp, local_pos, torch.zeros_like(local_pos)) - # (total_comp, 1) + (1, ratio) → (total_comp, ratio) - 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) - - kv_grouped = kv[gather_idx] # (total_comp, ratio, 1, coff * d) - score_grouped = score[gather_idx] - - # APE: (ratio, coff * d) → broadcast (1, ratio, 1, coff * d). - score_grouped = score_grouped + self.ape.view(1, ratio, 1, -1) + # Additive fused fast path (CuTe DSL, one kernel per direction) for the + # gather / overlap-window / softmax / weighted-sum region below. Returns None + # (-> keep the eager region) for any unsupported configuration; see + # csa_fused_compressor.py for the gating rules. + compressed_thd = None + if not pre_grouped: + compressed_thd = maybe_compress_thd_fused( + kv, + score, + self.ape, + cu_seqlens, + cu_seqlens_compressed, + total_comp, + ratio=ratio, + head_dim=self.head_dim, + coff=self.coff, + ) - if self.overlap: + if compressed_thd is None: if pre_grouped: - is_first = compressed_group_ids[:total_comp] == 0 + # Compressor-prep already groups rows as ``[g * ratio, (g + 1) * ratio)``. + kv_grouped = kv.reshape(total_comp, ratio, 1, -1) + score_grouped = score.reshape(total_comp, ratio, 1, -1) + local_pos = None else: - is_first = local_pos == 0 # (total_comp,) - kv_grouped = self._overlap_transform_thd(kv_grouped, is_first, fill_value=0) - score_grouped = self._overlap_transform_thd( - score_grouped, is_first, fill_value=float("-inf") - ) + # Build gather index: (total_comp, ratio). ``total_comp`` can be a + # static capacity for CUDA graph capture, so rows beyond the true + # ``cu_seqlens_compressed[-1]`` are mapped to a safe source row and left + # as tail padding by downstream index lowering. + row_idx = torch.arange(total_comp, device=device, dtype=cu_seqlens_compressed.dtype) + batch_ids = batch_of_row(cu_seqlens_compressed, total_q=total_comp) + valid_comp = row_idx < cu_seqlens_compressed[-1] + local_pos = row_idx - cu_seqlens_compressed[batch_ids] + local_pos = torch.where(valid_comp, local_pos, torch.zeros_like(local_pos)) + # (total_comp, 1) + (1, ratio) → (total_comp, ratio) + 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) + + kv_grouped = kv[gather_idx] # (total_comp, ratio, 1, coff * d) + score_grouped = score[gather_idx] + + # APE: (ratio, coff * d) → broadcast (1, ratio, 1, coff * d). + score_grouped = score_grouped + self.ape.view(1, ratio, 1, -1) + + if self.overlap: + if pre_grouped: + is_first = compressed_group_ids[:total_comp] == 0 + else: + is_first = local_pos == 0 # (total_comp,) + kv_grouped = self._overlap_transform_thd(kv_grouped, is_first, fill_value=0) + score_grouped = self._overlap_transform_thd( + score_grouped, is_first, fill_value=float("-inf") + ) - # Batched softmax + weighted sum — single kernel for all entries. - # (total_comp, [2*]ratio, 1, [coff*]d) → (total_comp, 1, head_dim) - weights = torch.softmax(score_grouped, dim=1, dtype=torch.float32).to(kv_grouped.dtype) - compressed_thd = (kv_grouped * weights).sum(dim=1) + # Batched softmax + weighted sum — single kernel for all entries. + # (total_comp, [2*]ratio, 1, [coff*]d) → (total_comp, 1, head_dim) + weights = torch.softmax(score_grouped, dim=1, dtype=torch.float32).to(kv_grouped.dtype) + compressed_thd = (kv_grouped * weights).sum(dim=1) compressed_thd = self.norm(compressed_thd.to(dtype)) diff --git a/megatron/core/transformer/experimental_attention_variant/csa_fused_compressor.py b/megatron/core/transformer/experimental_attention_variant/csa_fused_compressor.py new file mode 100644 index 00000000000..d4680422f86 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa_fused_compressor.py @@ -0,0 +1,241 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Dispatch shim for the fused CSA/HCA ``Compressor`` gated-pooling kernels. + +The fused forward+backward kernels live in the cudnn-frontend Python package +(``cudnn.csa.compressor``, added in https://github.com/NVIDIA/cudnn-frontend/pull/427, +following maintainer guidance on https://github.com/NVIDIA/Megatron-LM/pull/5984). +This module contains only the framework-side wiring: + + - an import guard that probes for the frontend's CSA compressor API by importing the + concrete entry points (capability detection, no version comparisons — installs that + predate the API simply lack ``cudnn.csa`` and keep the eager path); + - a ``torch.autograd.Function`` connecting the frontend's forward/backward wrappers + to autograd; + - the dispatch helper :func:`maybe_compress_thd_fused` used by + ``Compressor._forward_thd``: the fused result when the configuration is supported, + or ``None`` (the caller keeps its eager implementation, which remains the semantic + reference and the fallback everywhere). + +What the fused path replaces (see the frontend's ``docs/fe-oss-apis/csa.md`` for kernel +details): for the THD packed, non-pre-grouped path, the chain gather-index build -> +gather -> ``+ APE`` -> overlap-window transform (``coff == 2``) -> fp32 softmax -> +gated weighted sum -> bf16 cast, i.e. one fused compute kernel per direction (backward +adds a small ``dAPE`` zero-init) instead of ~40 forward / ~50 backward eager launches. +Numerics are fp32 with a single final bf16 rounding: not +bit-identical to the eager region (which rounds the softmax weights to bf16 and +multiplies in bf16) but at least as accurate against an fp64 oracle; forward, ``dKV`` +and ``dScore`` are bitwise run-to-run deterministic, ``dAPE`` is accumulated with fp32 +atomics and is not — the dispatch therefore keeps the eager path under +``torch.use_deterministic_algorithms(True)``. Measurements and numerics analysis: +cudnn-frontend PR #427 and https://github.com/NVIDIA/Megatron-LM/issues/5968. + +Dispatch gating (initial; everything else keeps eager): + - ``MCORE_CSA_FUSED_COMPRESSOR=0`` kill switch (disables the dispatch entirely); + - cudnn-frontend with the CSA compressor API importable, CUDA device with compute + capability 10.0 (the frontend's validated envelope); + - ``compress_ratio == 4`` / ``coff == 2`` (the production CSA/HCA configuration; + ``compress_ratio == 128`` stays on eager pending tuning), bf16 ``kv``/``score``, + fp32 ``ape``, int32 flat offsets (``total_tokens * coff * head_dim < 2**31``); + - eager under deterministic mode (``dAPE`` atomics, above) and under + ``torch.compile`` tracing (the frontend launch path takes raw pointers; eager lets + the compiler fuse the region itself). + +CUDA graphs: the frontend launch path is capture-compatible once the kernel for a given +``(ratio, head_dim, coff)`` configuration has been JIT-compiled; run one eager step per +configuration before capture (a first call that would JIT under capture raises a +``RuntimeError`` instead of corrupting the capture). The dispatch passes the caller's +static ``total_comp`` capacity through, so no device synchronization is introduced. +""" + +import os +from typing import Optional + +import torch + +_ENV_ENABLE = "MCORE_CSA_FUSED_COMPRESSOR" + +# The compute capability the frontend kernels are validated for (its ``check_support`` +# raises on anything else); the dispatch pre-filters so other devices keep eager +# silently. Mirrors ``cudnn.csa.compressor``'s envelope — widen together with it. +_SUPPORTED_COMPUTE_CAPABILITY = (10, 0) + +_UNINITIALIZED = object() +# Lazily resolved ``cudnn.csa.compressor`` module: ``_UNINITIALIZED`` -> not probed yet, +# ``None`` -> probed and unavailable (reason kept in ``_frontend_error``). +_frontend = _UNINITIALIZED +_frontend_error: Optional[Exception] = None + +_DEVICE_SUPPORTED_CACHE = {} + + +def _get_frontend(): + """Return cudnn-frontend's ``cudnn.csa.compressor`` module, or None if unavailable. + + Probed lazily (on the first dispatch call, not at Megatron import) by importing the + concrete entry points rather than comparing version numbers: any install that has + them can serve the dispatch, and ``nvidia-cudnn-frontend`` releases that predate the + CSA compressor API (cudnn-frontend #427) lack ``cudnn.csa`` and fall back to eager + naturally. Any import-time failure — missing package, missing ``nvidia-cutlass-dsl`` + (the frontend's ``cutedsl`` extra), a partially installed or ABI-incompatible + frontend — also degrades to eager instead of breaking ``csa.py``'s import chain; the + error is kept in ``_frontend_error`` for debugging. + """ + global _frontend, _frontend_error + if _frontend is _UNINITIALIZED: + try: + from cudnn.csa import compressor as fe_compressor + + fe_compressor.csa_compressor_forward_wrapper # pylint: disable=pointless-statement + fe_compressor.csa_compressor_backward_wrapper # pylint: disable=pointless-statement + _frontend = fe_compressor + # A partially installed frontend or DSL can raise more than ImportError. + except Exception as e: # pylint: disable=broad-except + _frontend = None + _frontend_error = e + return _frontend + + +def _dispatch_enabled() -> bool: + """Return True unless the fused compressor is disabled via environment variable.""" + return os.environ.get(_ENV_ENABLE, "1").lower() not in ("0", "false", "off") + + +def fused_compressor_available(device: Optional[torch.device] = None) -> bool: + """Return True when the fused kernels can run: frontend importable + supported device.""" + if _get_frontend() is None: + return False + try: + if not torch.cuda.is_available(): + return False + if device is not None and device.type != "cuda": + return False + if device is None or device.index is None: + index = torch.cuda.current_device() + else: + index = device.index + except (RuntimeError, AssertionError): # pragma: no cover - no CUDA context + return False + supported = _DEVICE_SUPPORTED_CACHE.get(index) + if supported is None: + supported = torch.cuda.get_device_capability(index) == _SUPPORTED_COMPUTE_CAPABILITY + _DEVICE_SUPPORTED_CACHE[index] = supported + return supported + + +class _CompressThdFused(torch.autograd.Function): + """Autograd wiring around the cudnn-frontend forward/backward wrappers. + + The wrappers allocate their outputs and validate their inputs; gradient semantics + (exact zeros for never-consumed elements, incoming gradients on static-capacity + padding rows ignored) match autograd on the eager region — see the frontend docs. + The backward wrapper raises under strict ``torch.use_deterministic_algorithms(True)`` + (``dAPE`` fp32 atomics); the dispatch already keeps eager in that mode, so this only + triggers when deterministic mode is enabled between forward and backward. + """ + + @staticmethod + def forward( + ctx, kv, score, ape, cu_seqlens, cu_seqlens_comp, ratio, head_dim, coff, total_comp + ): + """Run the frontend forward wrapper; saves inputs for the fused backward.""" + kv = kv.contiguous() + score = score.contiguous() + ape_c = ape.contiguous() + cu_i = cu_seqlens.to(dtype=torch.int32).contiguous() + cuc_i = cu_seqlens_comp.to(dtype=torch.int32).contiguous() + out = _get_frontend().csa_compressor_forward_wrapper( + kv, + score, + ape_c, + cu_i, + cuc_i, + ratio=ratio, + head_dim=head_dim, + coff=coff, + total_comp=total_comp, + )["out"] + ctx.save_for_backward(kv, score, ape_c, cu_i, cuc_i) + ctx.dims = (ratio, head_dim, coff) + return out + + @staticmethod + def backward(ctx, grad_out): + """Run the frontend backward wrapper; returns (dKV, dScore, dAPE, None...).""" + kv, score, ape, cu_i, cuc_i = ctx.saved_tensors + ratio, head_dim, coff = ctx.dims + grad_kv, grad_score, grad_ape = _get_frontend().csa_compressor_backward_wrapper( + kv, + score, + ape, + cu_i, + cuc_i, + grad_out.contiguous(), + ratio=ratio, + head_dim=head_dim, + coff=coff, + ) + return grad_kv, grad_score, grad_ape, None, None, None, None, None, None + + +def maybe_compress_thd_fused( + kv: torch.Tensor, + score: torch.Tensor, + ape: torch.Tensor, + cu_seqlens: torch.Tensor, + cu_seqlens_comp: torch.Tensor, + total_comp: int, + ratio: int, + head_dim: int, + coff: int, +) -> Optional[torch.Tensor]: + """Dispatch helper for ``Compressor._forward_thd``: fused result or None (use eager). + + Returns the pooled ``(total_comp, 1, head_dim)`` bf16 tensor when the cudnn-frontend + fused fast path supports the configuration, or None when the caller should keep the + eager implementation. See the module docstring for the gating rules; the frontend's + ``check_support`` re-validates the envelope and raises (instead of falling back) on + anything the gates below let through. + """ + if not _dispatch_enabled(): + return None + if kv.device.type != "cuda": + return None + if not fused_compressor_available(kv.device): + return None + if ratio != 4 or coff != 2: + return None + if kv.dtype != torch.bfloat16 or score.dtype != torch.bfloat16: + return None + if ape.dtype != torch.float32: + return None + # The backward is not deterministic for dAPE (fp32 atomics); respect torch's + # deterministic mode by keeping the (deterministic) eager path. + if torch.are_deterministic_algorithms_enabled(): + return None + # The frontend launch path uses raw pointers; keep eager (which torch.compile can + # trace and fuse itself) when compiling. + is_compiling = getattr(getattr(torch, "compiler", None), "is_compiling", None) + if is_compiling is not None and is_compiling(): + return None + total = kv.shape[0] + if kv.dim() != 3 or kv.shape[1] != 1 or kv.shape[2] != coff * head_dim: + return None + if score.shape != kv.shape: + return None + if total_comp <= 0 or total < ratio: + return None + if total * coff * head_dim >= 2**31: + return None + out = _CompressThdFused.apply( + kv.reshape(total, coff * head_dim), + score.reshape(total, coff * head_dim), + ape, + cu_seqlens, + cu_seqlens_comp, + ratio, + head_dim, + coff, + total_comp, + ) + return out.unsqueeze(1) From e073b2822b183859a692636f666e1356689ba33c Mon Sep 17 00:00:00 2001 From: zky Date: Fri, 24 Jul 2026 04:47:51 +0000 Subject: [PATCH 2/2] Add unit tests for the cudnn-frontend fused CSA compressor dispatch The kernels themselves are validated in cudnn-frontend (PR #427); these tests cover the Megatron-side wiring: numerics of the dispatched fused region vs an fp32-intermediate eager reference (bitwise dKV/dScore, forward within one bf16 rounding step) and vs the verbatim upstream eager numerics (tolerance) over ragged THD packs including segments shorter than ratio; fixed_total_comp static-capacity padding rows (including that nonzero padding-row gradients are ignored); dispatch gating and eager fallback (kill switch, missing/old cudnn-frontend, deterministic mode, ratio 128, non-bf16, layout, empty output), plus Compressor._forward_thd-level integration (fused engages and matches eager, bitwise eager fallback without the frontend, gradients flow). The module is marked launch_on_gb200 for the CC 10.0 CI lane and skips cleanly without CUDA, without a cudnn-frontend that provides cudnn.csa, or off compute capability 10.0. Addresses #5968. Signed-off-by: zky --- .../test_csa_fused_compressor.py | 432 ++++++++++++++++++ 1 file changed, 432 insertions(+) create mode 100644 tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_compressor.py diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_compressor.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_compressor.py new file mode 100644 index 00000000000..6f717a101e5 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_fused_compressor.py @@ -0,0 +1,432 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Unit tests for the cudnn-frontend-backed fused CSA Compressor dispatch. + +The kernels themselves are validated in cudnn-frontend (PR #427: numerics vs fp32/fp64 +references, backward zero-write coverage, CUDA-graph capture, determinism). These tests +cover the Megatron-side wiring only: + + - numerics of the dispatched fused region vs the eager region it replaces, with the + original PR #5984 gates: ``dKV``/``dScore`` bitwise vs an fp32-intermediate eager + reference, forward within one bf16 rounding step, tolerance vs the verbatim + upstream eager numerics; + - static-capacity padding rows (``fixed_total_comp``) through the dispatch; + - dispatch gating / eager fallback of ``maybe_compress_thd_fused``: the + ``MCORE_CSA_FUSED_COMPRESSOR=0`` kill switch, deterministic mode, unsupported + configurations, and a missing/old cudnn-frontend (no ``cudnn.csa``); + - ``Compressor._forward_thd`` integration: the fused dispatch engages and matches + eager, gradients flow, and the module falls back to the bitwise-identical eager + path when the frontend is unavailable. + +Without a cudnn-frontend that provides ``cudnn.csa`` (or without CUDA / off compute +capability 10.0) every test skips. +""" + +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +import torch + +from megatron.core.transformer.experimental_attention_variant import csa as csa_module +from megatron.core.transformer.experimental_attention_variant import csa_fused_compressor as cfc +from megatron.core.transformer.experimental_attention_variant.csa import ( + Compressor, + CompressorSubmodules, + batch_of_row, +) + +# Run this module on GB200 hardware in CI (marker-driven selection, see +# tests/unit_tests/find_test_cases.py); everywhere else the tests skip via +# _require_fused(). +pytestmark = pytest.mark.launch_on_gb200 + + +def _require_fused(): + if not torch.cuda.is_available(): + pytest.skip("fused CSA compressor tests require CUDA") + if cfc._get_frontend() is None: + pytest.skip( + "cudnn-frontend with the CSA compressor API (cudnn.csa, cudnn-frontend #427) " + f"is not available: {cfc._frontend_error!r}" + ) + if not cfc.fused_compressor_available(): + pytest.skip("fused CSA compressor requires compute capability 10.0") + + +# --------------------------------------------------------------------------- +# Eager reference: verbatim replica of the region of ``Compressor._forward_thd`` +# (non-pre-grouped THD path) that the fused dispatch replaces, from the projection +# outputs (kv, score) to the pre-RMSNorm pooled output. ``mode`` selects the +# numerics: "upstream" reproduces the current 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). The overlap-window +# transform is the real upstream implementation. +# --------------------------------------------------------------------------- + + +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_q=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() + + 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 + stub = SimpleNamespace(head_dim=d) + kv_grouped = Compressor._overlap_transform_thd(stub, kv_grouped, is_first, fill_value=0) + score_grouped = Compressor._overlap_transform_thd( + stub, score_grouped, is_first, 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) + else: # fp32 intermediates, single final bf16 rounding + weights = torch.softmax(score_grouped, dim=1, dtype=torch.float32) + out = (kv_grouped * weights).sum(dim=1).to(torch.bfloat16) + return out # (total_comp, 1, d) + + +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).""" + kv_l = kv.clone().requires_grad_(True) + score_l = score.clone().requires_grad_(True) + ape_l = 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 dispatch; returns (out, dKV, dScore, dAPE).""" + kv_l = kv.clone().requires_grad_(True) + score_l = score.clone().requires_grad_(True) + ape_l = ape.clone().requires_grad_(True) + out = cfc.maybe_compress_thd_fused( + kv_l, score_l, ape_l, cu, cuc, total_comp, ratio=ratio, head_dim=d, coff=coff + ) + assert out is not None, "fused dispatch did not engage" + out.backward(go) + torch.cuda.synchronize() + return out.detach(), kv_l.grad.detach(), score_l.grad.detach(), ape_l.grad.detach() + + +_SHAPES = [ + # (lens, head_dim); ratio = 4, coff = 2 (the only dispatched configuration) + pytest.param([2048], 128, id="b1-d128"), + pytest.param([1023, 2048, 509], 128, id="ragged3-d128"), + pytest.param([2048], 512, id="b1-d512"), + pytest.param([3, 515, 1024, 129], 128, id="short-seg-d128"), +] + + +@pytest.mark.parametrize("lens,d", _SHAPES) +def test_numerics_vs_eager(lens, d): + """Dispatched fused fwd+bwd vs fp32-eager (bitwise dKV/dScore) and upstream eager.""" + _require_fused() + ratio, coff = 4, 2 + 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") + + # 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) + + +def test_fixed_total_comp_padding(): + """Static-capacity padding rows: eager-matching forward, ignored padding gradients.""" + _require_fused() + # 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. + lens, d, ratio, coff, pad = [3, 515, 1024, 129], 128, 4, 2, 8 + 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 + + +def test_dispatch_gating_and_fallback(monkeypatch): + """``maybe_compress_thd_fused`` returns None for every unsupported configuration.""" + _require_fused() + kv, score, ape, cu, cuc, total_comp, _ = _make_inputs([512, 256], 128, 4, 2) + kwargs = dict(ratio=4, head_dim=128, coff=2) + + supported = cfc.maybe_compress_thd_fused(kv, score, ape, cu, cuc, total_comp, **kwargs) + assert supported is not None and supported.shape == (total_comp, 1, 128) + + # Kill switch. + monkeypatch.setenv("MCORE_CSA_FUSED_COMPRESSOR", "0") + assert cfc.maybe_compress_thd_fused(kv, score, ape, cu, cuc, total_comp, **kwargs) is None + monkeypatch.delenv("MCORE_CSA_FUSED_COMPRESSOR") + + # Missing/old cudnn-frontend (no ``cudnn.csa``): the probe caches None and the + # dispatch keeps eager. + with pytest.MonkeyPatch.context() as mp: + mp.setattr(cfc, "_frontend", None) + assert not cfc.fused_compressor_available() + assert cfc.maybe_compress_thd_fused(kv, score, ape, cu, cuc, total_comp, **kwargs) is None + + # Deterministic mode keeps the (deterministic) eager path — dAPE uses fp32 atomics. + 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: + assert cfc.maybe_compress_thd_fused(kv, score, ape, cu, cuc, total_comp, **kwargs) is None + finally: + torch.use_deterministic_algorithms(prev_det, warn_only=prev_warn) + + # Enabling deterministic mode between forward and backward: the frontend backward + # raises instead of silently returning a nondeterministic dAPE. + kv_l = kv.clone().requires_grad_(True) + score_l = score.clone().requires_grad_(True) + ape_l = ape.clone().requires_grad_(True) + out = cfc.maybe_compress_thd_fused(kv_l, score_l, ape_l, cu, cuc, total_comp, **kwargs) + assert out is not None + torch.use_deterministic_algorithms(True, warn_only=False) + try: + with pytest.raises(RuntimeError, match="not deterministic"): + out.backward(torch.ones_like(out)) + finally: + torch.use_deterministic_algorithms(prev_det, warn_only=prev_warn) + + # compress_ratio 128 stays on eager (kernels support it in the frontend, not yet a + # wall-clock win at production sizes). + kv1, score1, ape1, cu1, cuc1, tc1, _ = _make_inputs([1024], 128, 128, 1) + assert ( + cfc.maybe_compress_thd_fused( + kv1, score1, ape1, cu1, cuc1, tc1, ratio=128, head_dim=128, coff=1 + ) + is None + ) + + # Non-bf16 inputs fall back. + assert ( + cfc.maybe_compress_thd_fused(kv.float(), score.float(), ape, cu, cuc, total_comp, **kwargs) + is None + ) + + # Unexpected layout falls back. + assert ( + cfc.maybe_compress_thd_fused( + kv.view(kv.shape[0], -1), + score.view(score.shape[0], -1), + ape, + cu, + cuc, + total_comp, + **kwargs, + ) + is None + ) + + # Empty output falls back (nothing to compute). + assert cfc.maybe_compress_thd_fused(kv, score, ape, cu, cuc, 0, **kwargs) is None + + +class TestCompressorFusedIntegration: + """``Compressor._forward_thd`` level: fused dispatch engages and matches eager.""" + + @pytest.fixture(scope='class', autouse=True) + def class_environment(self, request): + # Skip (do not crash) on machines without CUDA / the frontend / CC 10.0 before + # touching model-parallel state. + _require_fused() + + from megatron.core.process_groups_config import ProcessGroupCollection + from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + from tests.unit_tests.test_utilities import Utils + + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + from megatron.core.transformer.transformer_config import MLATransformerConfig + + cls.config = MLATransformerConfig( + num_layers=4, + hidden_size=256, + num_attention_heads=16, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + q_lora_rank=64, + kv_lora_rank=64, + qk_head_dim=32, + qk_pos_emb_head_dim=32, + v_head_dim=64, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + multi_latent_attention=True, + experimental_attention_variant='dsv4_hybrid', + csa_compress_ratios=[4, 128, 4, 128], + csa_window_size=8, + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=0.0, + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + def _make_compressor(self): + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + return Compressor( + config=self.config, + submodules=CompressorSubmodules( + linear_wkv=ModuleSpec(module=TELinear), + linear_wgate=ModuleSpec(module=TELinear), + norm=ModuleSpec(module=TENorm), + ), + compress_ratio=4, + head_dim=self.config.v_head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + def test_forward_thd_fused_matches_eager(self, monkeypatch): + """THD forward: the fused dispatch engages and matches the eager path closely.""" + _require_fused() + compressor = self._make_compressor() + lens = [255, 512, 129] + total = sum(lens) + x = torch.randn(total, 1, self.config.hidden_size, dtype=torch.bfloat16, device="cuda") + cu_seqlens = torch.tensor( + [0] + list(torch.tensor(lens).cumsum(0)), dtype=torch.int32, device="cuda" + ) + + returns = [] + orig = csa_module.maybe_compress_thd_fused + + def _spy(*args, **kwargs): + result = orig(*args, **kwargs) + returns.append(result) + return result + + with patch.object(csa_module, "maybe_compress_thd_fused", side_effect=_spy): + out_fused, cuc_fused = compressor._forward_thd( + x, cu_seqlens, max_seqlen_q=max(lens), fixed_total_comp=total // 4 + ) + assert len(returns) == 1 + assert returns[0] is not None, "fused fast path did not engage" + + monkeypatch.setenv("MCORE_CSA_FUSED_COMPRESSOR", "0") + out_eager, cuc_eager = compressor._forward_thd( + x, cu_seqlens, max_seqlen_q=max(lens), fixed_total_comp=total // 4 + ) + monkeypatch.delenv("MCORE_CSA_FUSED_COMPRESSOR") + + assert out_fused.shape == out_eager.shape + assert torch.equal(cuc_fused, cuc_eager) + assert torch.allclose(out_fused.float(), out_eager.float(), rtol=0, atol=0.1) + + # Missing/old cudnn-frontend: bitwise the same eager path as the kill switch. + with pytest.MonkeyPatch.context() as mp: + mp.setattr(cfc, "_frontend", None) + out_fb, cuc_fb = compressor._forward_thd( + x, cu_seqlens, max_seqlen_q=max(lens), fixed_total_comp=total // 4 + ) + assert torch.equal(out_fb, out_eager) + assert torch.equal(cuc_fb, cuc_eager) + + def test_forward_thd_gradients_flow(self): + """Gradients flow through the fused fast path to inputs and parameters.""" + _require_fused() + compressor = self._make_compressor() + lens = [256, 512] + total = sum(lens) + x = torch.randn( + total, 1, self.config.hidden_size, dtype=torch.bfloat16, device="cuda" + ).requires_grad_(True) + cu_seqlens = torch.tensor( + [0] + list(torch.tensor(lens).cumsum(0)), dtype=torch.int32, device="cuda" + ) + out, _ = compressor._forward_thd(x, cu_seqlens, max_seqlen_q=max(lens)) + out.sum().backward() + assert x.grad is not None + assert compressor.ape.grad is not None + assert compressor.ape.grad.abs().sum().item() > 0