diff --git a/megatron/core/quantization/indexer_quantization.py b/megatron/core/quantization/indexer_quantization.py new file mode 100644 index 00000000000..f9f6488a2ce --- /dev/null +++ b/megatron/core/quantization/indexer_quantization.py @@ -0,0 +1,559 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Quantization utilities for DSA indexer inputs. + +The current SM100 path uses Transformer Engine for the expensive +BF16-to-MXFP8 conversion. Its +rowwise E8M0 scales are logical ``(flattened_rows, head_dim / 32)`` scales. +The compact cuDNN Indexer instead consumes the scales in its THD/GQA-aware +Blackwell 128x4 physical layout. A small Triton kernel performs only that +byte reordering and writes every padded byte into caller-owned storage. + +Precision-specific helpers remain explicitly named so this module can also +host SM90 FP8 and future indexer quantization paths without ambiguity. +""" + +from __future__ import annotations + +import math +from dataclasses import dataclass +from typing import Any + +import torch +from torch import Tensor + +try: + import transformer_engine_torch as tex + from transformer_engine.pytorch.tensor.mxfp8_tensor import MXFP8Quantizer + + HAVE_TE_MXFP8 = True +except ImportError: + tex = None + MXFP8Quantizer = None + HAVE_TE_MXFP8 = False + +try: + import triton + import triton.language as tl + + HAVE_TRITON = True +except ImportError: + from unittest.mock import MagicMock + + from megatron.core.utils import null_decorator + + triton = MagicMock() + triton.jit = null_decorator + tl = MagicMock() + HAVE_TRITON = False + + +def _ceil_div(a: int, b: int) -> int: + return (a + b - 1) // b + + +def indexer_mxfp8_scale_shape( + batch_size: int, max_seqlen: int, num_heads: int, head_dim: int, sf_vec_size: int = 32 +) -> tuple[int, int, int]: + """Return the physical BSHD E8M0 scale shape expected by the DSA kernel.""" + if min(batch_size, max_seqlen, num_heads, head_dim) <= 0: + raise ValueError("MXFP8 indexer scale dimensions must all be positive") + if sf_vec_size != 32: + raise ValueError(f"MXFP8 indexer only supports sf_vec_size=32, got {sf_vec_size}") + if head_dim % sf_vec_size != 0: + raise ValueError(f"MXFP8 indexer head_dim ({head_dim}) must be divisible by {sf_vec_size}") + + scale_groups = head_dim // sf_vec_size + packed_rows = _ceil_div(max_seqlen * num_heads, 128) * 128 + packed_groups = _ceil_div(scale_groups, 4) * 4 + return batch_size, packed_rows, packed_groups + + +def make_indexer_mxfp8_scale_cu_seqlens(cu_seqlens: Tensor, num_heads: int) -> Tensor: + """Return compact THD scale prefixes with independently padded sequences. + + ``num_heads`` is the number of logical scale rows per token. Each returned + sequence span is minimally padded so that its packed scale rows are a + multiple of the Blackwell 128-row atom. + """ + if ( + cu_seqlens.dtype != torch.int32 + or cu_seqlens.ndim != 1 + or cu_seqlens.numel() < 2 + or not cu_seqlens.is_contiguous() + ): + raise ValueError("cu_seqlens must be a contiguous int32 tensor with at least two elements") + if num_heads <= 0: + raise ValueError("num_heads must be positive") + + token_alignment = 128 // math.gcd(128, num_heads) + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + padded_lengths = ((lengths + token_alignment - 1) // token_alignment) * token_alignment + out = torch.zeros_like(cu_seqlens) + torch.cumsum(padded_lengths, dim=0, out=out[1:]) + return out + + +def indexer_mxfp8_thd_scale_capacity(total_tokens: int, num_sequences: int, num_heads: int) -> int: + """Return the maximum padded-token scale capacity for static THD geometry.""" + if min(total_tokens, num_sequences, num_heads) <= 0: + raise ValueError("THD MXFP8 scale capacity dimensions must all be positive") + + token_alignment = 128 // math.gcd(128, num_heads) + positive_sequences = min(total_tokens, num_sequences) + return ( + (total_tokens + positive_sequences * (token_alignment - 1)) // token_alignment + ) * token_alignment + + +@torch.compile +def _refresh_indexer_mxfp8_scale_cu_seqlens( + destination: Tensor, cu_seqlens: Tensor, token_alignment: int +) -> None: + """Write padded THD scale prefixes into caller-owned storage.""" + lengths = cu_seqlens[1:] - cu_seqlens[:-1] + padded_lengths = ( + torch.div(lengths + token_alignment - 1, token_alignment, rounding_mode="floor") + * token_alignment + ) + destination.zero_() + torch.cumsum(padded_lengths, dim=0, out=destination[1:]) + + +def refresh_indexer_mxfp8_scale_cu_seqlens( + destination: Tensor, cu_seqlens: Tensor, num_heads: int +) -> None: + """Refresh caller-owned compact THD scale prefixes from live sequence boundaries.""" + if ( + cu_seqlens.dtype != torch.int32 + or cu_seqlens.ndim != 1 + or cu_seqlens.numel() < 2 + or not cu_seqlens.is_contiguous() + ): + raise ValueError("cu_seqlens must be a contiguous int32 tensor with at least two elements") + if ( + destination.device != cu_seqlens.device + or destination.dtype != torch.int32 + or destination.ndim != 1 + or destination.numel() != cu_seqlens.numel() + or not destination.is_contiguous() + ): + raise ValueError( + "destination must be contiguous int32 storage matching cu_seqlens shape and device" + ) + if num_heads <= 0: + raise ValueError("num_heads must be positive") + + token_alignment = 128 // math.gcd(128, num_heads) + _refresh_indexer_mxfp8_scale_cu_seqlens(destination, cu_seqlens, token_alignment) + + +def indexer_mxfp8_thd_scale_shape( + padded_tokens: int, num_heads: int, head_dim: int, sf_vec_size: int = 32 +) -> tuple[int, int, int]: + """Return the compact THD E8M0 scale shape expected by the DSA kernel.""" + if min(padded_tokens, num_heads, head_dim) <= 0: + raise ValueError("MXFP8 indexer scale dimensions must all be positive") + if sf_vec_size != 32: + raise ValueError(f"MXFP8 indexer only supports sf_vec_size=32, got {sf_vec_size}") + if head_dim % sf_vec_size != 0: + raise ValueError(f"MXFP8 indexer head_dim ({head_dim}) must be divisible by {sf_vec_size}") + + packed_rows = padded_tokens * num_heads + if packed_rows % 128 != 0: + raise ValueError("THD MXFP8 scale rows must be a multiple of 128") + scale_groups = head_dim // sf_vec_size + packed_groups = _ceil_div(scale_groups, 4) * 4 + return 1, packed_rows, packed_groups + + +@dataclass +class IndexerMXFP8QuantizationBuffers: + """Preallocated TE destination and optional input padding for one tensor.""" + + quantizer: Any + quantized: Any + data: Tensor + logical_scale: Tensor + padded_input: Tensor | None + input_shape: tuple[int, ...] + num_rows: int + + def matches(self, x: Tensor) -> bool: + """Return whether these buffers can quantize ``x`` without allocation.""" + return all( + ( + tuple(x.shape) == self.input_shape, + x.device == self.data.device, + x.dtype == torch.bfloat16, + x.is_contiguous(), + self.data.dtype == torch.float8_e4m3fn, + tuple(self.data.shape) == self.input_shape, + self.data.is_contiguous(), + self.logical_scale.device == x.device, + self.logical_scale.dtype == torch.uint8, + self.logical_scale.is_contiguous(), + self.padded_input is None + or ( + self.padded_input.device == x.device + and self.padded_input.dtype == x.dtype + and self.padded_input.is_contiguous() + ), + ) + ) + + +def create_indexer_mxfp8_quantization_buffers(x: Tensor) -> IndexerMXFP8QuantizationBuffers: + """Allocate an unswizzled TE MXFP8 destination outside CUDA graph capture.""" + if not HAVE_TE_MXFP8: + raise RuntimeError("MXFP8 indexer quantization requires Transformer Engine MXFP8 support") + if not x.is_cuda or x.dtype != torch.bfloat16 or not x.is_contiguous() or x.ndim < 2: + raise ValueError("MXFP8 indexer input must be a contiguous CUDA BF16 tensor with ndim >= 2") + + head_dim = x.shape[-1] + if head_dim % 32 != 0: + raise ValueError(f"MXFP8 indexer head_dim ({head_dim}) must be divisible by 32") + num_rows = x.numel() // head_dim + padded_rows = _ceil_div(num_rows, 32) * 32 + + # TE requires the product of leading dimensions to be divisible by 32. + # Avoid an input copy for the common aligned Q/K shapes. Only small, + # ragged K tensors need the padded staging buffer. + padded_input = None + quantized_shape: tuple[int, ...] + if padded_rows == num_rows: + quantized_shape = tuple(x.shape) + else: + quantized_shape = (padded_rows, head_dim) + padded_input = torch.zeros(quantized_shape, dtype=x.dtype, device=x.device) + + quantizer = MXFP8Quantizer(fp8_dtype=tex.DType.kFloat8E4M3, rowwise=True, columnwise=False) + # The Indexer has its own THD/GQA packing. TE must expose logical scales. + quantizer.optimize_for_gemm = False + quantized = quantizer.make_empty(quantized_shape, dtype=x.dtype, device=x.device) + data = ( + quantized._rowwise_data.view(torch.float8_e4m3fn) + .reshape(padded_rows, head_dim)[:num_rows] + .reshape(x.shape) + ) + logical_scale = quantized._rowwise_scale_inv + return IndexerMXFP8QuantizationBuffers( + quantizer=quantizer, + quantized=quantized, + data=data, + logical_scale=logical_scale, + padded_input=padded_input, + input_shape=tuple(x.shape), + num_rows=num_rows, + ) + + +@triton.jit +def _pack_indexer_mxfp8_scale_bshd_kernel( + out_ptr, + logical_scale_ptr, + seqlen, + total_out_bytes, + NUM_HEADS: tl.constexpr, + REAL_GROUPS: tl.constexpr, + LOGICAL_PADDED_GROUPS: tl.constexpr, + PADDED_ROWS: tl.constexpr, + PADDED_GROUPS: tl.constexpr, + BLOCK: tl.constexpr, +): + """Pack logical BSHD TE scale bytes into the Indexer physical layout.""" + out_linear = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + in_bounds = out_linear < total_out_bytes + bytes_per_batch = PADDED_ROWS * PADDED_GROUPS + batch = out_linear // bytes_per_batch + logical_linear = out_linear - batch * bytes_per_batch + packed_row = logical_linear // PADDED_GROUPS + scale_group = logical_linear - packed_row * PADDED_GROUPS + local_token = packed_row // NUM_HEADS + head = packed_row - local_token * NUM_HEADS + + global_token = batch * seqlen + local_token + valid = in_bounds & (local_token < seqlen) & (scale_group < REAL_GROUPS) + source_row = global_token * NUM_HEADS + head + source_offset = source_row * LOGICAL_PADDED_GROUPS + scale_group + value = tl.load(logical_scale_ptr + source_offset, mask=valid, other=0) + + # Map logical (row, scale_group) to NVIDIA's packed F8_128x4 layout. + tile_idx = (packed_row // 128) * (PADDED_GROUPS // 4) + scale_group // 4 + physical_offset = ( + batch * bytes_per_batch + + tile_idx * 512 + + (packed_row % 32) * 16 + + ((packed_row % 128) // 32) * 4 + + scale_group % 4 + ) + tl.store(out_ptr + physical_offset, value, mask=in_bounds) + + +@triton.jit +def _pack_indexer_mxfp8_scale_thd_kernel( + out_ptr, + logical_scale_ptr, + cu_seqlens_ptr, + cu_seqlens_scale_padded_ptr, + NUM_HEADS: tl.constexpr, + REAL_GROUPS: tl.constexpr, + LOGICAL_PADDED_GROUPS: tl.constexpr, + PADDED_GROUPS: tl.constexpr, + BATCH_SIZE: tl.constexpr, + SEARCH_STEPS: tl.constexpr, +): + """Pack logical THD TE scale bytes into concatenated padded scale spans.""" + tile_idx = tl.program_id(0) + scale_tiles = PADDED_GROUPS // 4 + mn_tile = tile_idx // scale_tiles + scale_tile = tile_idx - mn_tile * scale_tiles + packed_row_start = mn_tile * 128 + + batch_lo = 0 + batch_hi = BATCH_SIZE + for _ in range(SEARCH_STEPS): + batch_mid = (batch_lo + batch_hi) // 2 + next_packed_row = tl.load(cu_seqlens_scale_padded_ptr + batch_mid + 1) * NUM_HEADS + in_lower_half = packed_row_start < next_packed_row + batch_hi = tl.where(in_lower_half, batch_mid, batch_hi) + batch_lo = tl.where(in_lower_half, batch_lo, batch_mid + 1) + batch = tl.minimum(batch_lo, BATCH_SIZE - 1) + + logical = tl.arange(0, 512) + row_in_tile = logical // 4 + packed_row = packed_row_start + row_in_tile + scale_group = scale_tile * 4 + logical % 4 + + scale_row_start = tl.load(cu_seqlens_scale_padded_ptr + batch) * NUM_HEADS + local_row = packed_row - scale_row_start + local_token = local_row // NUM_HEADS + head = local_row - local_token * NUM_HEADS + seq_start = tl.load(cu_seqlens_ptr + batch) + seq_end = tl.load(cu_seqlens_ptr + batch + 1) + seq_len = seq_end - seq_start + + valid = (local_token < seq_len) & (scale_group < REAL_GROUPS) + source_row = (seq_start + local_token) * NUM_HEADS + head + source_offset = source_row * LOGICAL_PADDED_GROUPS + scale_group + value = tl.load(logical_scale_ptr + source_offset, mask=valid, other=0) + + physical_offset = ( + tile_idx * 512 + (row_in_tile % 32) * 16 + ((row_in_tile % 128) // 32) * 4 + logical % 4 + ) + tl.store(out_ptr + physical_offset, value) + + +def pack_indexer_mxfp8_scale( + logical_scale: Tensor, + out_scale: Tensor, + *, + num_heads: int, + real_groups: int, + cu_seqlens: Tensor | None = None, + cu_seqlens_scale_padded: Tensor | None = None, + seqlen: int = 0, +) -> Tensor: + """Pack TE logical E8M0 bytes into caller-owned Indexer scale storage.""" + if not HAVE_TRITON: + raise RuntimeError("MXFP8 indexer scale packing requires Triton") + if ( + logical_scale.dtype != torch.uint8 + or logical_scale.ndim != 2 + or not logical_scale.is_cuda + or not logical_scale.is_contiguous() + ): + raise ValueError("logical_scale must be a contiguous CUDA uint8 matrix") + if ( + out_scale.dtype != torch.float8_e8m0fnu + or out_scale.ndim != 3 + or out_scale.device != logical_scale.device + or not out_scale.is_contiguous() + ): + raise ValueError("out_scale must be contiguous CUDA E8M0 storage on logical_scale.device") + if num_heads <= 0 or real_groups <= 0: + raise ValueError("num_heads and real_groups must be positive") + if out_scale.shape[1] % 128 != 0 or out_scale.shape[2] % 4 != 0: + raise ValueError("out_scale must have Blackwell 128x4 padded dimensions") + if real_groups > logical_scale.shape[1] or real_groups > out_scale.shape[2]: + raise ValueError("scale group count exceeds logical or packed scale storage") + + is_thd = cu_seqlens is not None + if is_thd: + if ( + cu_seqlens.device != logical_scale.device + or cu_seqlens.dtype != torch.int32 + or cu_seqlens.ndim != 1 + or cu_seqlens.numel() < 2 + or not cu_seqlens.is_contiguous() + ): + raise ValueError("cu_seqlens must be contiguous CUDA int32 storage") + if ( + cu_seqlens_scale_padded is None + or cu_seqlens_scale_padded.device != logical_scale.device + or cu_seqlens_scale_padded.dtype != torch.int32 + or cu_seqlens_scale_padded.ndim != 1 + or cu_seqlens_scale_padded.numel() != cu_seqlens.numel() + or not cu_seqlens_scale_padded.is_contiguous() + ): + raise ValueError( + "cu_seqlens_scale_padded must be contiguous CUDA int32 storage " + "matching cu_seqlens" + ) + if out_scale.shape[0] != 1 or out_scale.shape[1] % num_heads != 0: + raise ValueError("THD out_scale must have one L dimension and whole-token rows") + elif seqlen <= 0: + raise ValueError("BSHD scale packing requires a positive seqlen") + elif cu_seqlens_scale_padded is not None: + raise ValueError("cu_seqlens_scale_padded is only valid for THD scale packing") + + if is_thd: + _pack_indexer_mxfp8_scale_thd_kernel[(out_scale.numel() // 512,)]( + out_scale.view(torch.uint8), + logical_scale, + cu_seqlens, + cu_seqlens_scale_padded, + NUM_HEADS=num_heads, + REAL_GROUPS=real_groups, + LOGICAL_PADDED_GROUPS=logical_scale.shape[1], + PADDED_GROUPS=out_scale.shape[2], + BATCH_SIZE=cu_seqlens.numel() - 1, + SEARCH_STEPS=(cu_seqlens.numel() - 1).bit_length(), + ) + else: + total_out_bytes = out_scale.numel() + block = 256 + _pack_indexer_mxfp8_scale_bshd_kernel[(triton.cdiv(total_out_bytes, block),)]( + out_scale.view(torch.uint8), + logical_scale, + seqlen, + total_out_bytes, + NUM_HEADS=num_heads, + REAL_GROUPS=real_groups, + LOGICAL_PADDED_GROUPS=logical_scale.shape[1], + PADDED_ROWS=out_scale.shape[1], + PADDED_GROUPS=out_scale.shape[2], + BLOCK=block, + ) + return out_scale + + +def quantize_indexer_mxfp8( + x: Tensor, + *, + cu_seqlens: Tensor | None = None, + cu_seqlens_scale_padded: Tensor | None = None, + buffers: IndexerMXFP8QuantizationBuffers | None = None, + out_scale: Tensor | None = None, + sf_vec_size: int = 32, +) -> tuple[Tensor, Tensor]: + """Quantize BF16 BSHD/THD input with TE and pack scales for the Indexer. + + BSHD Q is ``(B, S, H, D)`` and K is ``(B, S, D)``. Packed THD Q is + ``(T, H, D)`` and K is ``(T, D)``. Supplying ``buffers`` and + ``out_scale`` makes the operation allocation-free and capture-safe. + """ + if sf_vec_size != 32: + raise ValueError(f"MXFP8 indexer only supports sf_vec_size=32, got {sf_vec_size}") + if not x.is_cuda or x.dtype != torch.bfloat16 or not x.is_contiguous(): + raise ValueError("MXFP8 indexer input must be a contiguous CUDA BF16 tensor") + + is_thd = cu_seqlens is not None + if is_thd: + if x.ndim == 3: + _, num_heads, head_dim = x.shape + elif x.ndim == 2: + _, head_dim = x.shape + num_heads = 1 + else: + raise ValueError(f"Packed THD MXFP8 input must be 2D or 3D, got shape {x.shape}") + if cu_seqlens_scale_padded is None: + raise ValueError("Packed THD MXFP8 quantization requires padded scale cu_seqlens") + seqlen = 0 + else: + if x.ndim == 4: + batch_size, seqlen, num_heads, head_dim = x.shape + elif x.ndim == 3: + batch_size, seqlen, head_dim = x.shape + num_heads = 1 + else: + raise ValueError(f"BSHD MXFP8 input must be 3D or 4D, got shape {x.shape}") + if cu_seqlens_scale_padded is not None: + raise ValueError("Padded scale cu_seqlens are only valid for packed THD input") + + if buffers is None: + buffers = create_indexer_mxfp8_quantization_buffers(x) + elif not buffers.matches(x): + raise ValueError("MXFP8 quantization buffers do not match the input tensor") + + source = x + if buffers.padded_input is not None: + buffers.padded_input[: buffers.num_rows].copy_(x.reshape(buffers.num_rows, head_dim)) + source = buffers.padded_input + buffers.quantizer.update_quantized(source, buffers.quantized) + + if is_thd: + if out_scale is None and torch.cuda.is_current_stream_capturing(): + raise RuntimeError("THD MXFP8 CUDA graph capture requires preallocated scale storage") + expected_scale_shape = ( + indexer_mxfp8_thd_scale_shape( + int(cu_seqlens_scale_padded[-1].item()), num_heads, head_dim, sf_vec_size + ) + if out_scale is None + else None + ) + else: + expected_scale_shape = indexer_mxfp8_scale_shape( + batch_size, seqlen, num_heads, head_dim, sf_vec_size + ) + if out_scale is None: + assert expected_scale_shape is not None + out_scale = torch.empty(expected_scale_shape, dtype=torch.float8_e8m0fnu, device=x.device) + elif is_thd and ( + out_scale.device != x.device + or out_scale.dtype != torch.float8_e8m0fnu + or out_scale.ndim != 3 + or out_scale.shape[0] != 1 + or out_scale.shape[1] % 128 != 0 + or out_scale.shape[2] != _ceil_div(head_dim // sf_vec_size, 4) * 4 + or not out_scale.is_contiguous() + ): + raise ValueError( + "THD out_scale must be contiguous E8M0 storage with shape " + "(1, multiple_of_128, padded_scale_groups)" + ) + elif not is_thd and ( + out_scale.device != x.device + or out_scale.dtype != torch.float8_e8m0fnu + or tuple(out_scale.shape) != expected_scale_shape + or not out_scale.is_contiguous() + ): + raise ValueError( + f"out_scale must be contiguous E8M0 storage with shape {expected_scale_shape}" + ) + + pack_indexer_mxfp8_scale( + buffers.logical_scale, + out_scale, + num_heads=num_heads, + real_groups=head_dim // sf_vec_size, + cu_seqlens=cu_seqlens, + cu_seqlens_scale_padded=cu_seqlens_scale_padded, + seqlen=seqlen, + ) + return buffers.data, out_scale + + +__all__ = [ + "HAVE_TE_MXFP8", + "HAVE_TRITON", + "IndexerMXFP8QuantizationBuffers", + "create_indexer_mxfp8_quantization_buffers", + "indexer_mxfp8_scale_shape", + "indexer_mxfp8_thd_scale_capacity", + "indexer_mxfp8_thd_scale_shape", + "make_indexer_mxfp8_scale_cu_seqlens", + "pack_indexer_mxfp8_scale", + "quantize_indexer_mxfp8", + "refresh_indexer_mxfp8_scale_cu_seqlens", +] diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index e2b409290b4..eb95a7a1bae 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -18,12 +18,20 @@ 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_kernels import ( + BSHDCompactIndexerWorkspace, FusedCSAIndexerSparseAttnFromTopkFunc, + THDCompactIndexerWorkspace, batch_of_row, + bshd_compact_indexer_available, build_flat_topk_idxs, + build_thd_compact_k_layout, csa_sparse_attn, fused_csa_indexer_sparse_attn, indexer_topk, + pack_thd_compact_k, + prepare_bshd_compact_indexer_workspace, + prepare_thd_compact_indexer_workspace, + thd_compact_indexer_available, ) from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexerLossAutoScaler, @@ -1113,9 +1121,11 @@ def _forward_thd( if total_comp == 0: return None, cu_seqlens_compressed - # Token-wise projections on the FULL flat input — no boundary issue. - kv, _ = self.linear_wkv(x) # (total, 1, coff * head_dim) - score, _ = self.linear_wgate(x) # (total, 1, coff * head_dim) + # Run the compressor GEMMs in high precision (BF16) even under FP8 training. + with get_fp8_disabled_context(self.config): + # Token-wise projections on the FULL flat input — no boundary issue. + 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)``. @@ -1628,6 +1638,16 @@ def __init__( else: self.indexer = None + # Compact CUDA graphs reference caller-owned cuDNN buffers by + # address. Retain every warmed-up static geometry for the lifetime of + # this module so later graph captures cannot invalidate an earlier + # graph's storage. ``_active_*`` identifies the immediately preceding + # warmup. + self._bshd_compact_indexer_workspaces: list[BSHDCompactIndexerWorkspace] = [] + self._active_bshd_compact_indexer_workspace: BSHDCompactIndexerWorkspace | None = None + self._thd_compact_indexer_workspaces: list[THDCompactIndexerWorkspace] = [] + self._active_thd_compact_indexer_workspace: THDCompactIndexerWorkspace | None = None + def backward_dw(self): """Compute the deferred weight gradients of the optional compressor/indexer submodules. @@ -1643,6 +1663,154 @@ def backward_dw(self): # Private helpers – each owns one logical slice of the forward pass. # ------------------------------------------------------------------ + def _get_bshd_compact_indexer_workspace( + self, + q: torch.Tensor, + k: torch.Tensor, + *, + topk: int, + ratio: int, + return_softmax: bool = False, + ) -> BSHDCompactIndexerWorkspace | None: + """Return persistent compact storage for a warmed-up BSHD geometry. + + ``q`` and ``k`` arrive in the public SBHD/SBD layouts. Workspace + metadata records the BSHD/BSD shapes produced inside the fused + wrapper, avoiding an extra permutation during graph capture. + """ + precision = self.config.dsa_indexer_precision + if self.config.cuda_graph_impl == "none" or not bshd_compact_indexer_available( + q, k, precision + ): + return None + if q.ndim != 4 or k.ndim != 3: + raise ValueError("BSHD compact workspace expects SBHD q and SBD k inputs") + + q_shape = (q.shape[1], q.shape[0], q.shape[2], q.shape[3]) + k_shape = (k.shape[1], k.shape[0], k.shape[2]) + match_kwargs = dict( + q_shape=q_shape, + k_shape=k_shape, + device=q.device, + topk=topk, + ratio=ratio, + return_softmax=return_softmax, + precision=precision, + ) + capturing = torch.cuda.is_current_stream_capturing() + workspace = self._active_bshd_compact_indexer_workspace + if capturing: + if workspace is not None and workspace.matches_shape(**match_kwargs): + return workspace + raise ValueError( + "BSHD compact CUDA graph capture requires an eagerly prepared compact " + "workspace for the active static geometry. Run eager warmup before capture." + ) + + for workspace in self._bshd_compact_indexer_workspaces: + if workspace.matches_shape(**match_kwargs): + self._active_bshd_compact_indexer_workspace = workspace + return workspace + + q_bshd = q.permute(1, 0, 2, 3).contiguous() + k_bsd = k.permute(1, 0, 2).contiguous() + workspace = prepare_bshd_compact_indexer_workspace( + q_bshd, + k_bsd, + topk=topk, + ratio=ratio, + return_softmax=return_softmax, + precision=precision, + ) + self._active_bshd_compact_indexer_workspace = workspace + if workspace is not None: + self._bshd_compact_indexer_workspaces.append(workspace) + return workspace + + def _get_thd_compact_indexer_workspace( + self, + q: torch.Tensor, + k: torch.Tensor, + *, + topk: int, + ratio: int, + cu_seqlens_q: torch.Tensor, + cu_seqlens_k: torch.Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + q_causal_offsets: torch.Tensor | None = None, + return_softmax: bool = False, + ) -> THDCompactIndexerWorkspace | None: + """Return persistent compact storage for a warmed-up THD graph geometry. + + Workspace sizing synchronizes with the host, so it is performed only + during eager CUDA-graph warmup. Capture reuses the active workspace; + when warmup was skipped or static shapes changed, capture must fail + clearly. Unsupported devices/frontends retain the non-compact path. + """ + precision = self.config.dsa_indexer_precision + if self.config.cuda_graph_impl == "none" or not thd_compact_indexer_available( + q, k, precision + ): + return None + + capturing = torch.cuda.is_current_stream_capturing() + workspace = self._active_thd_compact_indexer_workspace + if capturing: + if workspace is not None and workspace.matches( + q=q, + k=k, + topk=topk, + ratio=ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + q_causal_offsets=q_causal_offsets, + return_softmax=return_softmax, + precision=precision, + ): + return workspace + raise ValueError( + "THD compact CUDA graph capture requires an eagerly prepared compact " + "workspace for the active packed geometry. Run eager warmup before capture." + ) + + for workspace in self._thd_compact_indexer_workspaces: + if workspace.matches( + q=q, + k=k, + topk=topk, + ratio=ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + q_causal_offsets=q_causal_offsets, + return_softmax=return_softmax, + precision=precision, + ): + self._active_thd_compact_indexer_workspace = workspace + return workspace + + workspace = prepare_thd_compact_indexer_workspace( + q, + k, + topk=topk, + ratio=ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + q_causal_offsets=q_causal_offsets, + return_softmax=return_softmax, + precision=precision, + ) + self._active_thd_compact_indexer_workspace = workspace + if workspace is not None: + self._thd_compact_indexer_workspaces.append(workspace) + return workspace + def _build_kv_full( self, kv: torch.Tensor, x: torch.Tensor ) -> Tuple[torch.Tensor, Optional[torch.Tensor], int]: @@ -1815,6 +1983,9 @@ def _forward_fused_indexer_inference( q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk( x_det, qr_det, packed_seq_params ) + compact_workspace = self._get_bshd_compact_indexer_workspace( + q_indexer, k_indexer, topk=self.indexer.index_topk, ratio=self.compress_ratio + ) topk_indices_cmp, _ = indexer_topk( q_indexer, k_indexer, @@ -1822,6 +1993,9 @@ def _forward_fused_indexer_inference( self.indexer.index_topk, self.compress_ratio, indexer_softmax_scale=self.indexer.softmax_scale, + compact_workspace=compact_workspace, + precision=self.config.dsa_indexer_precision, + deterministic=self.config.deterministic_mode, ) compress_topk_idxs = torch.where(topk_indices_cmp >= 0, topk_indices_cmp + offset, -1) flat_idxs, flat_tlen = build_flat_topk_idxs( @@ -1865,6 +2039,14 @@ def _forward_fused_indexer_training( nvtx_range_pop("compressed_indices") indexer_loss_coeff = self.config.dsa_indexer_loss_coeff or 0.0 + sparse_loss = getattr(self.config, "dsa_indexer_use_sparse_loss", True) + compact_workspace = self._get_bshd_compact_indexer_workspace( + q_indexer, + k_indexer, + topk=self.indexer.index_topk, + ratio=self.compress_ratio, + return_softmax=indexer_loss_coeff > 0 and sparse_loss, + ) nvtx_range_push("sparse_attn_kernel") output, indexer_loss = fused_csa_indexer_sparse_attn( @@ -1880,9 +2062,12 @@ def _forward_fused_indexer_training( self.softmax_scale, self.indexer.softmax_scale, indexer_loss_coeff, - sparse_loss=getattr(self.config, "dsa_indexer_use_sparse_loss", True), + sparse_loss=sparse_loss, kv_offset=offset, calculate_per_token_loss=self.config.calculate_per_token_loss, + compact_workspace=compact_workspace, + indexer_precision=self.config.dsa_indexer_precision, + deterministic=self.config.deterministic_mode, ) nvtx_range_pop("sparse_attn_kernel") @@ -2222,17 +2407,40 @@ def _forward_fused_indexer_inference_thd( topk_indices_cmp = torch.full((total_q, 0), -1, dtype=torch.int32, device=query.device) else: k_thd = k_indexer.squeeze(1) + topk_k_thd = k_thd + topk_cu_seqlens_k = cu_seqlens_compressed_idx + topk_max_seqlen_k = max_seqlen_compressed_idx + if thd_compact_indexer_available(q_thd, k_thd, self.config.dsa_indexer_precision): + topk_cu_seqlens_k, source_row_map = build_thd_compact_k_layout( + cu_seqlens_q, cu_seqlens_compressed_idx, k_thd.shape[0], self.compress_ratio + ) + topk_k_thd = pack_thd_compact_k(k_thd, source_row_map) + topk_max_seqlen_k += 2 + + compact_workspace = self._get_thd_compact_indexer_workspace( + q_thd, + topk_k_thd, + topk=self.indexer.index_topk, + ratio=self.compress_ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=topk_cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=topk_max_seqlen_k, + ) topk_indices_cmp, _ = indexer_topk( q_thd, - k_thd, + topk_k_thd, w_thd, topk=self.indexer.index_topk, ratio=self.compress_ratio, indexer_softmax_scale=self.indexer.softmax_scale, cu_seqlens_q=cu_seqlens_q, - cu_seqlens_kv=cu_seqlens_compressed_idx, + cu_seqlens_kv=topk_cu_seqlens_k, max_seqlen_q=max_seqlen_q, - max_seqlen_kv=max_seqlen_compressed_idx, + max_seqlen_kv=topk_max_seqlen_k, + compact_workspace=compact_workspace, + precision=self.config.dsa_indexer_precision, + deterministic=self.config.deterministic_mode, ) # Shift into per-segment full-KV index space. @@ -2290,7 +2498,7 @@ def _forward_fused_indexer_training_thd( ``(total_q, 1, np * hn)``. """ sparse_loss = getattr(self.config, "dsa_indexer_use_sparse_loss", True) - indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) + indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) or 0.0 x_det = x.detach() qr_det = qr.detach() @@ -2308,6 +2516,30 @@ def _forward_fused_indexer_training_thd( w_thd = weights_indexer.squeeze(1) k_thd = k_indexer.squeeze(1) + workspace_k = k_thd + workspace_cu_seqlens_k = cu_seqlens_compressed_idx + workspace_max_seqlen_k = max_seqlen_compressed_idx + if self.config.cuda_graph_impl != "none" and thd_compact_indexer_available( + q_thd, k_thd, self.config.dsa_indexer_precision + ): + workspace_cu_seqlens_k, source_row_map = build_thd_compact_k_layout( + cu_seqlens_q, cu_seqlens_compressed_idx, k_thd.shape[0], self.compress_ratio + ) + workspace_k = pack_thd_compact_k(k_thd, source_row_map) + workspace_max_seqlen_k += 2 + + compact_workspace = self._get_thd_compact_indexer_workspace( + q_thd, + workspace_k, + topk=self.indexer.index_topk, + ratio=self.compress_ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=workspace_cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=workspace_max_seqlen_k, + return_softmax=indexer_loss_coeff > 0 and sparse_loss, + ) + # Supply unpadded cu_seqlens so padding rows are excluded from # the indexer KL loss (mirrors the unfused path's cu_seqlens_q_for_loss). # Only pass when they actually differ (by reference or storage) to avoid @@ -2345,6 +2577,9 @@ def _forward_fused_indexer_training_thd( compressed_kv=compressed_kv, calculate_per_token_loss=self.config.calculate_per_token_loss, cu_seqlens_q_unpadded=cu_seqlens_q_unpadded, + compact_workspace=compact_workspace, + indexer_precision=self.config.dsa_indexer_precision, + deterministic=self.config.deterministic_mode, ) if indexer_loss_coeff > 0: @@ -2404,6 +2639,7 @@ def _forward_thd_cp( # within its sequence. ``seq_to_rank_row`` maps each compressed block to # the row where its K and KV are stored in the all-gathered buffer. compressed_topk = seq_to_rank_row = None + compact_indexer_predict = None ratio = self.compress_ratio indexer = self.indexer indexer_loss_coeff = self.config.dsa_indexer_loss_coeff or 0.0 @@ -2496,20 +2732,65 @@ def _forward_thd_cp( k_indexer_seq_major = torch.index_select( k_indexer_rank_major, 0, seq_to_rank_row.clamp_min(0) ) - # Each top-k entry is still a logical compressed id within that - # query's sequence here. - compressed_topk, indexer_layout = cp_utils.compute_cp_indexer_topk( - q_indexer_cp, - weights_indexer_cp, - k_indexer_seq_major, + indexer_layout = cp_utils.build_cp_indexer_layout( cu_seqlens, cu_seqlens_compressed, global_start, - ratio, - indexer.index_topk, - indexer.softmax_scale, - max_seqlen_q=max_seqlen_q, - use_fused=self.use_fused_kernels, + l_local, + k_indexer_seq_major.shape[0], + ) + topk_indexer_layout = indexer_layout + k_indexer_for_topk = k_indexer_seq_major + if self.use_fused_kernels: + topk_indexer_layout, source_row_map = cp_utils.build_cp_compact_indexer_layout( + indexer_layout, cu_seqlens_compressed, k_indexer_seq_major.shape[0], ratio + ) + k_indexer_for_topk = cp_utils.pack_cp_compact_indexer_k( + k_indexer_seq_major, source_row_map + ) + return_indexer_softmax = ( + self.use_fused_kernels + and training_with_grad + and sparse_indexer_loss + and indexer_loss_coeff > 0 + ) + compact_workspace = None + if self.use_fused_kernels: + compact_workspace = self._get_thd_compact_indexer_workspace( + q_indexer_cp, + k_indexer_for_topk, + topk=indexer.index_topk, + ratio=ratio, + cu_seqlens_q=topk_indexer_layout[0], + cu_seqlens_k=topk_indexer_layout[1], + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_q // ratio + 2, + q_causal_offsets=topk_indexer_layout[2], + return_softmax=return_indexer_softmax, + ) + # Each top-k entry is still a logical compressed id within that + # query's sequence here. Only the ids are remapped later; the + # compact softmax slots remain aligned with their selected keys. + (compressed_topk, indexer_layout, compact_indexer_predict) = ( + cp_utils.compute_cp_indexer_topk( + q_indexer_cp, + weights_indexer_cp, + k_indexer_for_topk, + cu_seqlens, + cu_seqlens_compressed, + global_start, + ratio, + indexer.index_topk, + indexer.softmax_scale, + max_seqlen_q=max_seqlen_q, + use_fused=self.use_fused_kernels, + deterministic=self.config.deterministic_mode, + precision=self.config.dsa_indexer_precision, + compact_workspace=compact_workspace, + return_softmax=return_indexer_softmax, + indexer_layout=topk_indexer_layout, + logical_indexer_layout=indexer_layout, + ) ) # ---- Step 5: attention compressed KV path ------------------------- @@ -2583,11 +2864,7 @@ def _forward_thd_cp( real_seqlens = cu_seqlens_q_unpadded[1:] - cu_seqlens_q_unpadded[:-1] positions = global_rows - cu_seqlens[batch_ids] q_padding_mask = positions >= real_seqlens[batch_ids] - output, indexer_loss = ( - FusedCSAIndexerSparseAttnFromTopkFunc.apply - if self.use_fused_kernels - else _unfused_indexer_sparse_attn_from_topk - )( + sparse_attn_args = ( query, kv_full_thd, self.attn_sink.float(), @@ -2607,6 +2884,12 @@ def _forward_thd_cp( indexer_layout, q_padding_mask, ) + if self.use_fused_kernels: + sparse_attn = FusedCSAIndexerSparseAttnFromTopkFunc.apply + sparse_attn_args += (compact_indexer_predict,) + else: + sparse_attn = _unfused_indexer_sparse_attn_from_topk + output, indexer_loss = sparse_attn(*sparse_attn_args) if indexer_loss_coeff > 0: DSAIndexerLossLoggingHelper.save_loss_to_tracker( loss=indexer_loss, diff --git a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py index 1a85466d4c1..d1e9a074024 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_cp_utils.py @@ -15,7 +15,14 @@ from megatron.core.fusions.fused_mla_yarn_rope_apply import fused_mla_rope_inplace from megatron.core.models.common.embeddings.rope_utils import _apply_rotary_pos_emb_bshd from megatron.core.transformer.experimental_attention_variant import csa_cp_layout_kernels -from megatron.core.transformer.experimental_attention_variant.csa_kernels import indexer_topk +from megatron.core.transformer.experimental_attention_variant.csa_kernels import ( + THDCompactIndexerWorkspace, + build_thd_compact_k_layout, + indexer_topk, + pack_thd_compact_k, +) + +CPIndexerLayout = Tuple[torch.Tensor, torch.Tensor, torch.Tensor] # ============================================================================= # RoPE Wrappers @@ -274,17 +281,19 @@ def prepare_cp_compressor_input( @torch.compile -def _build_cp_indexer_layout( +def build_cp_indexer_layout( cu_seqlens_q: torch.Tensor, cu_seqlens_compressed: torch.Tensor, global_start: int, local_rows: int, + total_k_rows: int, ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Build the indexer's packed local-Q/full-K metadata.""" # Each real Q segment intersects its sequence with this rank's row interval, # while K keeps the sequence's full compressed segment. The final synthetic - # segment holds CP capacity padding and has zero K rows. Causal offsets - # restore each non-empty local Q segment's position in the original sequence. + # segment owns the fixed-capacity K tail so the logical endpoint matches the + # physical buffer. Causal offsets restore each non-empty local Q segment's + # position in the original sequence. global_end = global_start + local_rows zero = torch.zeros((1,), dtype=cu_seqlens_q.dtype, device=cu_seqlens_q.device) local_starts = cu_seqlens_q[:-1].clamp_min(global_start) @@ -293,13 +302,38 @@ def _build_cp_indexer_layout( q_prefix = torch.cumsum(q_lens, dim=0, dtype=torch.int32) padding_q = (global_end - cu_seqlens_q[-1].clamp_min(global_start)).clamp_min(0) cu_q_topk = torch.cat((zero, q_prefix, (q_prefix[-1] + padding_q).view(1))) - cu_k_topk = torch.cat((cu_seqlens_compressed, cu_seqlens_compressed[-1:])) + # The gathered K buffer has fixed CP capacity, which can exceed the number + # of valid compressed rows. Assign its tail to the synthetic padding + # segment so cu_k[-1] still matches the physical K tensor length. + k_capacity_end = cu_seqlens_compressed.new_full((1,), total_k_rows) + cu_k_topk = torch.cat((cu_seqlens_compressed, k_capacity_end)) q_causal_offsets = torch.cat( (torch.where(q_lens > 0, local_starts - cu_seqlens_q[:-1], 0), zero) ) return cu_q_topk, cu_k_topk, q_causal_offsets +def build_cp_compact_indexer_layout( + logical_layout: CPIndexerLayout, + cu_seqlens_compressed: torch.Tensor, + total_k_rows: int, + ratio: int, +) -> Tuple[CPIndexerLayout, torch.Tensor]: + """Pad each THD K segment with an invisible row for compact Top-K.""" + cu_q_topk, _, q_causal_offsets = logical_layout + cu_k_topk, source_row_map = build_thd_compact_k_layout( + cu_q_topk, cu_seqlens_compressed, total_k_rows, ratio + ) + return (cu_q_topk, cu_k_topk, q_causal_offsets), source_row_map + + +def pack_cp_compact_indexer_k( + k_indexer_seq_major: torch.Tensor, source_row_map: torch.Tensor +) -> torch.Tensor: + """Insert zero-valued, causally invisible K rows for compact THD.""" + return pack_thd_compact_k(k_indexer_seq_major, source_row_map) + + def compute_cp_indexer_topk( q_indexer_local: torch.Tensor, weights_indexer_local: torch.Tensor, @@ -312,14 +346,20 @@ def compute_cp_indexer_topk( indexer_softmax_scale: float, max_seqlen_q: int, use_fused: bool, -) -> Tuple[Optional[torch.Tensor], Optional[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]]: - """Return local top-k and its local-Q/full-K packed layout.""" + deterministic: bool = False, + precision: str = "bf16", + compact_workspace: Optional[THDCompactIndexerWorkspace] = None, + return_softmax: bool = False, + indexer_layout: Optional[CPIndexerLayout] = None, + logical_indexer_layout: Optional[CPIndexerLayout] = None, +) -> Tuple[Optional[torch.Tensor], Optional[CPIndexerLayout], Optional[torch.Tensor]]: + """Return local top-k, packed layout, and optional compact Top-K softmax.""" topk_width = int(topk_width) if topk_width == 0 or k_indexer_seq_major.shape[0] == 0: - return None, None + return None, None, None max_seqlen_kv = int(max_seqlen_q) // int(ratio) if max_seqlen_kv == 0: - return None, None + return None, None, None global_start = int(global_start) l_local = q_indexer_local.shape[0] @@ -329,9 +369,22 @@ def compute_cp_indexer_topk( f"{l_local}, got {weights_indexer_local.shape[0]}." ) - cu_q_topk, cu_k_topk, q_causal_offsets = _build_cp_indexer_layout( - cu_seqlens_q, cu_seqlens_compressed, global_start, l_local - ) + if logical_indexer_layout is None: + logical_indexer_layout = build_cp_indexer_layout( + cu_seqlens_q, cu_seqlens_compressed, global_start, l_local, k_indexer_seq_major.shape[0] + ) + if not use_fused: + indexer_layout = logical_indexer_layout + elif indexer_layout is None: + indexer_layout, source_row_map = build_cp_compact_indexer_layout( + logical_indexer_layout, cu_seqlens_compressed, k_indexer_seq_major.shape[0], ratio + ) + k_indexer_seq_major = pack_cp_compact_indexer_k(k_indexer_seq_major, source_row_map) + + if use_fused: + # Each real segment has one or two invisible K padding rows. + max_seqlen_kv += 2 + cu_q_topk, cu_k_topk, q_causal_offsets = indexer_layout if not use_fused: global_rows = torch.arange( @@ -381,9 +434,9 @@ def compute_cp_indexer_topk( values, rows = torch.topk(scores, selected_width, dim=-1) local_rows = k_positions[rows].to(torch.int32) output[start:end, :selected_width] = torch.where(torch.isfinite(values), local_rows, -1) - return output, (cu_q_topk, cu_k_topk, q_causal_offsets) + return output, logical_indexer_layout, None - topk, _ = indexer_topk( + topk_result = indexer_topk( q_indexer_local, k_indexer_seq_major, weights_indexer_local, @@ -395,5 +448,14 @@ def compute_cp_indexer_topk( max_seqlen_q=int(max_seqlen_q), max_seqlen_kv=int(max_seqlen_kv), q_causal_offsets=q_causal_offsets, + compact_workspace=compact_workspace, + precision=precision, + deterministic=deterministic, + return_softmax=return_softmax, ) - return topk, (cu_q_topk, cu_k_topk, q_causal_offsets) + compact_predict = None + if return_softmax: + topk, _, compact_predict = topk_result + else: + topk, _ = topk_result + return topk, logical_indexer_layout, compact_predict diff --git a/megatron/core/transformer/experimental_attention_variant/csa_kernels.py b/megatron/core/transformer/experimental_attention_variant/csa_kernels.py index 06f2345fe0a..181f36b16f6 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa_kernels.py +++ b/megatron/core/transformer/experimental_attention_variant/csa_kernels.py @@ -23,12 +23,27 @@ from __future__ import annotations +import warnings +from dataclasses import dataclass from functools import lru_cache from typing import Optional, Tuple import torch from torch import Tensor +from megatron.core.quantization.indexer_quantization import ( + HAVE_TE_MXFP8, + HAVE_TRITON, + IndexerMXFP8QuantizationBuffers, + create_indexer_mxfp8_quantization_buffers, + indexer_mxfp8_scale_shape, + indexer_mxfp8_thd_scale_capacity, + indexer_mxfp8_thd_scale_shape, + make_indexer_mxfp8_scale_cu_seqlens, + quantize_indexer_mxfp8, + refresh_indexer_mxfp8_scale_cu_seqlens, +) + # --------------------------------------------------------------------------- # Lazy kernel imports # --------------------------------------------------------------------------- @@ -38,6 +53,553 @@ _DSA = None +@dataclass(frozen=True) +class _BSHDCompactIndexerGeometry: + """Static BSHD geometry for a caller-owned compact workspace.""" + + q_shape: tuple[int, ...] + k_shape: tuple[int, ...] + ratio: int + precision: str = "bf16" + + +@dataclass(frozen=True) +class _THDCompactIndexerGeometry: + """Static THD geometry for a caller-owned compact workspace.""" + + q_shape: tuple[int, ...] + k_shape: tuple[int, ...] + batch_size: int + has_q_causal_offsets: bool + ratio: int + max_seqlen_q: int + max_seqlen_k: int + precision: str = "bf16" + + +@dataclass +class MXFP8IndexerWorkspace: + """Caller-owned TE quantization destinations and packed Indexer scales.""" + + q_buffers: IndexerMXFP8QuantizationBuffers + k_buffers: IndexerMXFP8QuantizationBuffers + q_scale: Tensor + k_scale: Tensor + cu_seqlens_q_scale_padded: Tensor | None = None + cu_seqlens_k_scale_padded: Tensor | None = None + + +@dataclass +class BSHDCompactIndexerWorkspace: + """Caller-owned buffers for compact BSHD CUDA-graph capture.""" + + cand_buffer: Tensor + out_indices: Tensor + out_logits: Tensor + softmax_out: Tensor | None + geometry: _BSHDCompactIndexerGeometry + mxfp8: MXFP8IndexerWorkspace | None = None + + def matches_shape( + self, + *, + q_shape: tuple[int, ...], + k_shape: tuple[int, ...], + device: torch.device, + topk: int, + ratio: int, + return_softmax: bool, + precision: str = "bf16", + ) -> bool: + """Return whether static BSHD metadata matches this workspace.""" + if len(q_shape) != 4 or len(k_shape) != 3: + return False + batch_size, seqlen_q, num_heads, head_dim = q_shape + k_batch, seqlen_k, k_head_dim = k_shape + out_shape = (batch_size, seqlen_q, topk) + softmax_valid = not return_softmax or ( + self.softmax_out is not None + and self.softmax_out.device == device + and self.softmax_out.dtype == torch.float32 + and tuple(self.softmax_out.shape) == out_shape + and self.softmax_out.is_contiguous() + ) + mxfp8_valid = self.mxfp8 is None + if precision == "mxfp8": + mxfp8 = self.mxfp8 + q_scale_shape = indexer_mxfp8_scale_shape(batch_size, seqlen_q, num_heads, head_dim) + k_scale_shape = indexer_mxfp8_scale_shape(batch_size, seqlen_k, 1, k_head_dim) + mxfp8_valid = mxfp8 is not None and all( + ( + tuple(mxfp8.q_buffers.input_shape) == q_shape, + mxfp8.q_buffers.data.device == device, + tuple(mxfp8.k_buffers.input_shape) == k_shape, + mxfp8.k_buffers.data.device == device, + mxfp8.q_scale.device == device, + mxfp8.q_scale.dtype == torch.float8_e8m0fnu, + tuple(mxfp8.q_scale.shape) == q_scale_shape, + mxfp8.q_scale.is_contiguous(), + mxfp8.k_scale.device == device, + mxfp8.k_scale.dtype == torch.float8_e8m0fnu, + tuple(mxfp8.k_scale.shape) == k_scale_shape, + mxfp8.k_scale.is_contiguous(), + mxfp8.cu_seqlens_q_scale_padded is None, + mxfp8.cu_seqlens_k_scale_padded is None, + ) + ) + return all( + ( + precision in ("bf16", "mxfp8"), + self.geometry.precision == precision, + self.geometry.q_shape == q_shape, + self.geometry.k_shape == k_shape, + self.geometry.ratio == ratio, + batch_size == k_batch, + head_dim == k_head_dim, + self.cand_buffer.device == device, + self.cand_buffer.dtype == torch.float32, + self.cand_buffer.is_contiguous(), + self.out_indices.device == device, + self.out_indices.dtype == torch.int32, + tuple(self.out_indices.shape) == out_shape, + self.out_indices.is_contiguous(), + self.out_logits.device == device, + self.out_logits.dtype == torch.float32, + tuple(self.out_logits.shape) == out_shape, + self.out_logits.is_contiguous(), + softmax_valid, + mxfp8_valid, + ) + ) + + def matches( + self, + *, + q: Tensor, + k: Tensor, + topk: int, + ratio: int, + return_softmax: bool, + precision: str = "bf16", + ) -> bool: + """Return whether this workspace can serve the current BSHD call.""" + inputs_valid = all( + ( + q.dtype == torch.bfloat16, + k.dtype == torch.bfloat16, + q.device == k.device, + q.is_contiguous(), + k.is_contiguous(), + ) + ) + if not inputs_valid or not self.matches_shape( + q_shape=tuple(q.shape), + k_shape=tuple(k.shape), + device=q.device, + topk=topk, + ratio=ratio, + return_softmax=return_softmax, + precision=precision, + ): + return False + if precision == "mxfp8": + assert self.mxfp8 is not None + return self.mxfp8.q_buffers.matches(q) and self.mxfp8.k_buffers.matches(k) + return True + + def validate( + self, + *, + q: Tensor, + k: Tensor, + topk: int, + ratio: int, + return_softmax: bool, + precision: str = "bf16", + ) -> None: + """Validate caller-owned BSHD storage before dispatch.""" + if not self.matches( + q=q, k=k, topk=topk, ratio=ratio, return_softmax=return_softmax, precision=precision + ): + raise ValueError( + "BSHD compact_workspace does not match the current buffers or static geometry; " + "prepare a workspace for this exact call during eager warmup." + ) + + +@dataclass +class THDCompactIndexerWorkspace: + """Caller-owned buffers and geometry for compact THD CUDA-graph capture.""" + + cand_batch_offsets: Tensor + cand_buffer: Tensor + out_indices: Tensor + out_logits: Tensor + softmax_out: Tensor | None + geometry: _THDCompactIndexerGeometry + mxfp8: MXFP8IndexerWorkspace | None = None + + def matches( + self, + *, + q: Tensor, + k: Tensor, + topk: int, + ratio: int, + cu_seqlens_q: Tensor, + cu_seqlens_k: Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + q_causal_offsets: Tensor | None, + return_softmax: bool, + precision: str = "bf16", + ) -> bool: + """Return whether this workspace can serve the current static indexer geometry.""" + geometry = self.geometry + out_shape = (q.shape[0], topk) + batch_size = geometry.batch_size + inputs_valid = all( + ( + q.dtype == torch.bfloat16, + k.dtype == torch.bfloat16, + tuple(q.shape) == geometry.q_shape, + tuple(k.shape) == geometry.k_shape, + q.device == k.device, + q.is_contiguous(), + k.is_contiguous(), + ) + ) + cu_seqlens_valid = all( + tensor.device == q.device + and tensor.dtype == torch.int32 + and tensor.ndim == 1 + and tensor.numel() == batch_size + 1 + and tensor.is_contiguous() + for tensor in (cu_seqlens_q, cu_seqlens_k) + ) + q_causal_offsets_valid = ( + q_causal_offsets is None + if not geometry.has_q_causal_offsets + else ( + q_causal_offsets is not None + and q_causal_offsets.device == q.device + and q_causal_offsets.dtype == torch.int32 + and q_causal_offsets.ndim == 1 + and q_causal_offsets.numel() == batch_size + and q_causal_offsets.is_contiguous() + ) + ) + softmax_valid = not return_softmax or ( + self.softmax_out is not None + and self.softmax_out.device == q.device + and self.softmax_out.dtype == torch.float32 + and tuple(self.softmax_out.shape) == out_shape + and self.softmax_out.is_contiguous() + ) + mxfp8_valid = self.mxfp8 is None + if precision == "mxfp8": + mxfp8 = self.mxfp8 + q_scale_shape = indexer_mxfp8_thd_scale_shape( + indexer_mxfp8_thd_scale_capacity(q.shape[0], batch_size, q.shape[1]), + q.shape[1], + q.shape[2], + ) + k_scale_shape = indexer_mxfp8_thd_scale_shape( + indexer_mxfp8_thd_scale_capacity(k.shape[0], batch_size, 1), 1, k.shape[1] + ) + scale_prefixes_valid = mxfp8 is not None and all( + prefix is not None + and prefix.device == q.device + and prefix.dtype == torch.int32 + and prefix.ndim == 1 + and prefix.numel() == batch_size + 1 + and prefix.is_contiguous() + for prefix in (mxfp8.cu_seqlens_q_scale_padded, mxfp8.cu_seqlens_k_scale_padded) + ) + mxfp8_valid = ( + mxfp8 is not None + and scale_prefixes_valid + and all( + ( + mxfp8.q_scale.device == q.device, + mxfp8.q_scale.dtype == torch.float8_e8m0fnu, + tuple(mxfp8.q_scale.shape) == q_scale_shape, + mxfp8.q_scale.is_contiguous(), + mxfp8.k_scale.device == q.device, + mxfp8.k_scale.dtype == torch.float8_e8m0fnu, + tuple(mxfp8.k_scale.shape) == k_scale_shape, + mxfp8.k_scale.is_contiguous(), + mxfp8.q_buffers.matches(q), + mxfp8.k_buffers.matches(k), + ) + ) + ) + static_valid = ( + precision in ("bf16", "mxfp8") + and inputs_valid + and geometry.precision == precision + and geometry.ratio == ratio + and geometry.max_seqlen_q == max_seqlen_q + and geometry.max_seqlen_k == max_seqlen_k + and cu_seqlens_valid + and q_causal_offsets_valid + and self.cand_batch_offsets.device == q.device + and self.cand_batch_offsets.dtype == torch.int64 + and self.cand_batch_offsets.ndim == 1 + and self.cand_batch_offsets.numel() == batch_size + 1 + and self.cand_batch_offsets.is_contiguous() + and self.cand_buffer.device == q.device + and self.cand_buffer.dtype == torch.float32 + and self.cand_buffer.is_contiguous() + and self.out_indices.device == q.device + and self.out_indices.dtype == torch.int32 + and tuple(self.out_indices.shape) == out_shape + and self.out_indices.is_contiguous() + and self.out_logits.device == q.device + and self.out_logits.dtype == torch.float32 + and tuple(self.out_logits.shape) == out_shape + and self.out_logits.is_contiguous() + and softmax_valid + and mxfp8_valid + ) + return static_valid + + def validate( + self, + *, + q: Tensor, + k: Tensor, + topk: int, + ratio: int, + cu_seqlens_q: Tensor, + cu_seqlens_k: Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + q_causal_offsets: Tensor | None, + return_softmax: bool, + precision: str = "bf16", + ) -> None: + """Validate static workspace metadata without synchronizing.""" + if not self.matches( + q=q, + k=k, + topk=topk, + ratio=ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + q_causal_offsets=q_causal_offsets, + return_softmax=return_softmax, + precision=precision, + ): + raise ValueError( + "THD compact_workspace does not match the current buffers or static geometry; " + "prepare a compatible workspace during eager warmup." + ) + + +def _compact_indexer_available(q: Tensor, k: Tensor, precision: str) -> bool: + """Return whether this device/frontend can dispatch compact Top-K.""" + _ensure_dsa_namespace() + try: + compact_wrapper = getattr(_DSA, "indexer_forward_top_k_wrapper", None) + except (AttributeError, ImportError): + compact_wrapper = None + return ( + precision in ("bf16", "mxfp8") + and (precision != "mxfp8" or (HAVE_TE_MXFP8 and HAVE_TRITON)) + and callable(compact_wrapper) + and q.dtype == torch.bfloat16 + and k.dtype == torch.bfloat16 + and torch.cuda.get_device_capability(q.device)[0] >= 10 + ) + + +def bshd_compact_indexer_available(q: Tensor, k: Tensor, precision: str = "bf16") -> bool: + """Return whether this device/frontend can dispatch compact BSHD Top-K.""" + return _compact_indexer_available(q, k, precision) + + +def thd_compact_indexer_available(q: Tensor, k: Tensor, precision: str = "bf16") -> bool: + """Return whether this device/frontend can dispatch compact THD Top-K.""" + return _compact_indexer_available(q, k, precision) + + +def prepare_bshd_compact_indexer_workspace( + q: Tensor, + k: Tensor, + *, + topk: int, + ratio: int, + return_softmax: bool = False, + precision: str = "bf16", +) -> BSHDCompactIndexerWorkspace | None: + """Preallocate compact BSHD buffers before CUDA graph capture. + + ``q`` and ``k`` use the BSHD/BSD layouts consumed by + :func:`_indexer_topk_core`. The workspace forces the upstream compact + wrapper's single-launch mode so its candidate sizing is identical for + BF16 and MXFP8 and no per-window scratch is created during capture. + """ + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("BSHD compact indexer workspace must be prepared before CUDA capture") + if q.ndim != 4 or k.ndim != 3: + raise ValueError( + f"BSHD workspace expects q (B,S,H,D) and k (B,S,D), got {q.shape}, {k.shape}" + ) + if not q.is_contiguous() or not k.is_contiguous(): + raise ValueError("BSHD workspace inputs must be contiguous") + + _ensure_dsa_namespace() + size_helper = getattr(_DSA, "compress_topk_cand_buffer_size", None) + if not bshd_compact_indexer_available(q, k, precision) or not callable(size_helper): + return None + + batch_size, seqlen_q, num_heads, head_dim = q.shape + k_batch, seqlen_k, k_head_dim = k.shape + if batch_size != k_batch or head_dim != k_head_dim: + raise ValueError("BSHD workspace q/k batch size and head dimension must match") + cand_floats = size_helper(batch_size, seqlen_q, seqlen_k, ratio, microbatch_rows=0) + device = q.device + out_shape = (batch_size, seqlen_q, topk) + mxfp8_workspace = None + if precision == "mxfp8": + mxfp8_workspace = MXFP8IndexerWorkspace( + q_buffers=create_indexer_mxfp8_quantization_buffers(q), + k_buffers=create_indexer_mxfp8_quantization_buffers(k), + q_scale=torch.empty( + indexer_mxfp8_scale_shape(batch_size, seqlen_q, num_heads, head_dim), + dtype=torch.float8_e8m0fnu, + device=device, + ), + k_scale=torch.empty( + indexer_mxfp8_scale_shape(batch_size, seqlen_k, 1, k_head_dim), + dtype=torch.float8_e8m0fnu, + device=device, + ), + ) + return BSHDCompactIndexerWorkspace( + cand_buffer=torch.empty(cand_floats, dtype=torch.float32, device=device), + out_indices=torch.empty(out_shape, dtype=torch.int32, device=device), + out_logits=torch.empty(out_shape, dtype=torch.float32, device=device), + softmax_out=( + torch.empty(out_shape, dtype=torch.float32, device=device) if return_softmax else None + ), + geometry=_BSHDCompactIndexerGeometry( + q_shape=tuple(q.shape), k_shape=tuple(k.shape), ratio=ratio, precision=precision + ), + mxfp8=mxfp8_workspace, + ) + + +@torch.compile +def _refresh_thd_compact_cand_batch_offsets( + destination: Tensor, cu_seqlens_q: Tensor, ratio: int, q_causal_offsets: Tensor | None +) -> None: + """Refresh caller-owned THD candidate offsets from graph inputs.""" + cu_q = cu_seqlens_q.to(torch.int64) + q_lengths = cu_q[1:] - cu_q[:-1] + if q_causal_offsets is None: + q_starts = torch.zeros_like(q_lengths) + else: + q_starts = q_causal_offsets.to(torch.int64) + + def prefix_candidate_count(length: Tensor) -> Tensor: + quotient = torch.div(length, ratio, rounding_mode="floor") + remainder = length - quotient * ratio + return ratio * quotient * (quotient - 1) // 2 + quotient * (remainder + 1) + + per_sequence = prefix_candidate_count(q_starts + q_lengths) - prefix_candidate_count(q_starts) + destination.zero_() + torch.cumsum(per_sequence, dim=0, out=destination[1:]) + + +def prepare_thd_compact_indexer_workspace( + q: Tensor, + k: Tensor, + *, + topk: int, + ratio: int, + cu_seqlens_q: Tensor, + cu_seqlens_k: Tensor, + max_seqlen_q: int, + max_seqlen_k: int, + q_causal_offsets: Tensor | None = None, + return_softmax: bool = False, + precision: str = "bf16", +) -> THDCompactIndexerWorkspace | None: + """Preallocate compact THD buffers before CUDA graph capture. + + Returns None when the installed cuDNN Frontend or GPU does not expose + the SM100 compact API. The sizing helper performs a GPU-to-host sync and + therefore this function must never be called from inside graph capture. + Eager warmup must prepare this workspace before any THD compact capture. + """ + if torch.cuda.is_current_stream_capturing(): + raise RuntimeError("THD compact indexer workspace must be prepared before CUDA capture") + + _ensure_dsa_namespace() + size_helper = getattr(_DSA, "compress_topk_cand_buffer_size_thd", None) + if not thd_compact_indexer_available(q, k, precision) or not callable(size_helper): + return None + + cand_batch_offsets, cand_floats = size_helper( + cu_seqlens_q, cu_seqlens_k, ratio, q_causal_offsets=q_causal_offsets + ) + device = q.device + # Sequence boundaries may change between graph capture and replay while + # tensor shapes and maxima remain static. Reserve the geometry-wide upper + # bound so refreshed candidate offsets always address valid storage. + cand_floats = max(cand_floats, q.shape[0] * int(max_seqlen_k)) + out_shape = (q.shape[0], topk) + mxfp8_workspace = None + if precision == "mxfp8": + q_buffers = create_indexer_mxfp8_quantization_buffers(q) + k_buffers = create_indexer_mxfp8_quantization_buffers(k) + cu_seqlens_q_scale_padded = make_indexer_mxfp8_scale_cu_seqlens(cu_seqlens_q, q.shape[1]) + cu_seqlens_k_scale_padded = make_indexer_mxfp8_scale_cu_seqlens(cu_seqlens_k, 1) + batch_size = cu_seqlens_q.numel() - 1 + q_scale_capacity = indexer_mxfp8_thd_scale_capacity(q.shape[0], batch_size, q.shape[1]) + k_scale_capacity = indexer_mxfp8_thd_scale_capacity(k.shape[0], batch_size, 1) + mxfp8_workspace = MXFP8IndexerWorkspace( + q_buffers=q_buffers, + k_buffers=k_buffers, + q_scale=torch.empty( + indexer_mxfp8_thd_scale_shape(q_scale_capacity, q.shape[1], q.shape[2]), + dtype=torch.float8_e8m0fnu, + device=device, + ), + k_scale=torch.empty( + indexer_mxfp8_thd_scale_shape(k_scale_capacity, 1, k.shape[1]), + dtype=torch.float8_e8m0fnu, + device=device, + ), + cu_seqlens_q_scale_padded=cu_seqlens_q_scale_padded, + cu_seqlens_k_scale_padded=cu_seqlens_k_scale_padded, + ) + return THDCompactIndexerWorkspace( + cand_batch_offsets=cand_batch_offsets, + cand_buffer=torch.empty(cand_floats, dtype=torch.float32, device=device), + out_indices=torch.empty(out_shape, dtype=torch.int32, device=device), + out_logits=torch.empty(out_shape, dtype=torch.float32, device=device), + softmax_out=( + torch.empty(out_shape, dtype=torch.float32, device=device) if return_softmax else None + ), + geometry=_THDCompactIndexerGeometry( + q_shape=tuple(q.shape), + k_shape=tuple(k.shape), + batch_size=cu_seqlens_q.numel() - 1, + has_q_causal_offsets=q_causal_offsets is not None, + ratio=ratio, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + precision=precision, + ), + mxfp8=mxfp8_workspace, + ) + + def _ensure_flash_mla(): """Lazily import the FlashMLA sparse-forward kernel. @@ -184,6 +746,76 @@ def batch_of_row(cu_seqlens_q: Tensor, total_q: Optional[int] = None) -> Tensor: ) +@torch.compile +def build_thd_compact_k_layout( + cu_seqlens_q: Tensor, cu_seqlens_k: Tensor, total_k_rows: int, ratio: int +) -> Tuple[Tensor, Tensor]: + """Build compact-THD K metadata and a physical-to-logical row map. + + The compact wrapper requires ``cu_seqlens_k[-1] == k.shape[0]`` and + ``seqlen_q[b] <= seqlen_k[b] * ratio`` for every packed segment. A + floor-compressed sequence violates the latter whenever it has a tail. + Insert one zero-valued row after every real K segment and distribute any + fixed-capacity K tail across those segments. The compact causal mask can + never expose these appended rows, so returned local K ids are unchanged. + + ``cu_seqlens_q`` may contain one additional synthetic padding segment, as + used by the CP path. Its K rows consume the corresponding fixed-capacity + tail instead of being assigned to a real sequence. + """ + num_sequences = cu_seqlens_k.shape[0] - 1 + zero = torch.zeros((1,), dtype=cu_seqlens_k.dtype, device=cu_seqlens_k.device) + valid_k_lens = cu_seqlens_k[1:] - cu_seqlens_k[:-1] + + padding_q = (cu_seqlens_q[-1] - cu_seqlens_q[num_sequences]).clamp_min(0) + padding_k = torch.div(padding_q + int(ratio) - 1, int(ratio), rounding_mode="floor") + capacity_gap = (int(total_k_rows) - cu_seqlens_k[-1]).clamp_min(0) + # The static compressed capacity reserves enough rows for a synthetic Q + # segment. All remaining capacity can be spread over real segments. + remaining_extra = (capacity_gap - padding_k).clamp_min(0) + sequence_ids = torch.arange(num_sequences, dtype=cu_seqlens_k.dtype, device=cu_seqlens_k.device) + base_extra = torch.div(remaining_extra, num_sequences, rounding_mode="floor") + extra_remainder = remaining_extra - base_extra * num_sequences + extra_per_sequence = base_extra + (sequence_ids < extra_remainder).to(cu_seqlens_k.dtype) + + kernel_k_lens = valid_k_lens + 1 + extra_per_sequence + kernel_k_prefix = torch.cumsum(kernel_k_lens, dim=0, dtype=torch.int32) + if cu_seqlens_q.shape[0] == cu_seqlens_k.shape[0]: + compact_cu_seqlens_k = torch.cat((zero, kernel_k_prefix)) + else: + compact_cu_seqlens_k = torch.cat( + (zero, kernel_k_prefix, (kernel_k_prefix[-1] + padding_k).view(1)) + ) + + kernel_rows = torch.arange( + int(total_k_rows) + num_sequences, dtype=cu_seqlens_k.dtype, device=cu_seqlens_k.device + ) + # Avoid ``torch.bucketize`` here: ``torch.compile`` can miscompile CUDA + # bucketize when the boundaries are produced inside the same graph. + kernel_sequence_ids = ( + (kernel_rows.unsqueeze(1) >= compact_cu_seqlens_k[1:].unsqueeze(0)) + .sum(dim=1, dtype=torch.int64) + .clamp_max(num_sequences) + ) + safe_sequence_ids = kernel_sequence_ids.clamp_max(num_sequences - 1) + kernel_positions = kernel_rows - compact_cu_seqlens_k[kernel_sequence_ids] + valid_source = (kernel_sequence_ids < num_sequences) & ( + kernel_positions < valid_k_lens[safe_sequence_ids] + ) + source_rows = cu_seqlens_k[safe_sequence_ids] + kernel_positions + source_row_map = torch.where( + valid_source, source_rows, torch.full_like(source_rows, int(total_k_rows)) + ).to(torch.int64) + return compact_cu_seqlens_k, source_row_map + + +@torch.compile +def pack_thd_compact_k(k: Tensor, source_row_map: Tensor) -> Tensor: + """Pack THD K with zero-valued rows that remain causally unreachable.""" + k_with_padding = torch.cat((k, torch.zeros_like(k[:1]))) + return torch.index_select(k_with_padding, 0, source_row_map) + + def local_to_global_flat( local_idxs: Tensor, batch_size: int, @@ -480,33 +1112,39 @@ def _indexer_topk_core( max_seqlen_q: Optional[int] = None, max_seqlen_kv: Optional[int] = None, q_causal_offsets: Optional[Tensor] = None, -) -> Tuple[Tensor, Tensor, Tensor]: + use_compact: bool = False, + return_softmax: bool = False, + compact_workspace: BSHDCompactIndexerWorkspace | THDCompactIndexerWorkspace | None = None, + precision: str = "bf16", + deterministic: bool = False, +) -> Tuple[Tensor, Tensor, Optional[Tensor], Optional[Tensor]]: """Layout-agnostic core for :func:`indexer_topk`. - Wraps cuDNN Frontend's CuTe-DSL indexer-forward kernel. - The pipeline (forward → per-row valid lengths → radix top-K → pad-to-``topk`` → ``topk_length``) - is the same for both layouts; only the input shape glue, valid-length derivation, - and output reshape differ. Selected by ``cu_seqlens_q``. + Wraps cuDNN Frontend's CuTe-DSL indexer-forward kernels. On SM10x, + ``use_compact=True`` selects the combined forward + Top-K wrapper when it + is available. Otherwise this falls back to the existing forward → per-row + valid lengths → radix Top-K pipeline. Compact BSHD and THD CUDA-graph + capture requires a layout-matching workspace prepared during eager warmup. BSHD layout (``cu_seqlens_q is None``): q: ``(b, sq, idx_nh, idx_hd)`` bf16, C-contiguous. k: ``(b, sk, idx_hd)`` bf16, C-contiguous. w: ``(b, sq, idx_nh)`` bf16, C-contiguous, **already ``indexer_softmax_scale``-scaled by the caller**. - Returns: - ``(topk_indices (b, sq, topk) int32, - topk_length (b, sq) int32)`` — invalid slots ``-1``. + Returns local ``topk_indices`` and ``topk_length``. The third return + is the dense score tensor for the fallback path (otherwise ``None``); + the fourth is the compact kernel's Top-K softmax when requested + (otherwise ``None``). THD packed layout (``cu_seqlens_q is not None``): q: ``(total_q, idx_nh, idx_hd)`` bf16. k: ``(total_k, idx_hd)`` bf16. w: ``(total_q, idx_nh)`` bf16, already scaled. cu_seqlens_q/kv, max_seqlen_q/kv: standard packed args. - Returns: - ``(topk_indices (total_q, topk) int32, - topk_length (total_q,) int32)`` — per-batch LOCAL ids - in ``[0, seqlen_kv[batch])``; use :func:`local_to_global_flat` - (with ``cu_seqlens_q/kv``) to promote to flat-global ids. + The first two returns are ``topk_indices (total_q, topk)`` and + ``topk_length (total_q,)``. Indices are per-batch LOCAL ids in + ``[0, seqlen_kv[batch])``; use :func:`local_to_global_flat` (with + ``cu_seqlens_q/kv``) to promote them to flat-global ids. Two internal entry points besides :func:`indexer_topk`: @@ -514,6 +1152,11 @@ def _indexer_topk_core( so the SBHD→BSHD permute can be performed once and reused across the indexer forward and the score-recompute backward kernels. """ + if precision not in ("bf16", "mxfp8"): + raise ValueError(f"Unsupported DSA indexer precision: {precision!r}") + if precision == "mxfp8" and not use_compact: + raise ValueError("MXFP8 indexer precision requires the compact forward + Top-K path") + is_thd = cu_seqlens_q is not None device = q.device @@ -527,8 +1170,195 @@ def _indexer_topk_core( raise ValueError(f"THD w must be (total_q, idx_nh), got {w.shape}") if max_seqlen_kv == 0 or k.shape[0] == 0: raise ValueError("indexer_topk requires at least one K row.") + elif k.shape[1] == 0: + raise ValueError("indexer_topk requires at least one K row.") - _ensure_dsa_namespace() + if precision == "mxfp8" and ( + q.shape[-2:] != (64, 128) or k.shape[-1] != 128 or w.shape[-1] != 64 + ): + raise ValueError( + "MXFP8 compact indexer requires 64 Q heads, head_dim=128, one K head, and 64 weights" + ) + + _ensure_dsa_namespace() + + # Symbol detection preserves compatibility with cuDNN Frontend versions + # that do not expose the compact forward + Top-K wrapper. + try: + compact_wrapper = getattr(_DSA, "indexer_forward_top_k_wrapper", None) + except (AttributeError, ImportError): + compact_wrapper = None + compact_available = ( + use_compact + and (precision != "mxfp8" or (HAVE_TE_MXFP8 and HAVE_TRITON)) + and callable(compact_wrapper) + and all(t.dtype == torch.bfloat16 for t in (q, k, w)) + and torch.cuda.get_device_capability(device)[0] >= 10 + ) + if precision == "mxfp8" and not compact_available: + raise RuntimeError( + "MXFP8 compact indexer requires Transformer Engine MXFP8, Triton, a cuDNN " + "Frontend compact wrapper with MXFP8 support, BF16 source tensors, and SM100+" + ) + if precision == "bf16" and use_compact and not compact_available: + warnings.warn( + "Compact indexer forward + Top-K was requested but is unavailable; " + "falling back to dense indexer forward + standalone Top-K.", + RuntimeWarning, + stacklevel=2, + ) + capturing = compact_available and torch.cuda.is_current_stream_capturing() + if compact_available and capturing and compact_workspace is None: + layout = "THD" if is_thd else "BSHD" + raise ValueError( + f"{layout} compact CUDA graph capture requires a preallocated compact_workspace. " + "Prepare it during eager warmup before capture." + ) + + if compact_available: + compact_kwargs = dict( + ratio=ratio, + precision=precision, + return_softmax=return_softmax, + topk_indices_global=False, + deterministic=deterministic, + ) + if is_thd: + if compact_workspace is not None: + if not isinstance(compact_workspace, THDCompactIndexerWorkspace): + raise ValueError("THD compact dispatch requires a THDCompactIndexerWorkspace") + compact_workspace.validate( + q=q, + k=k, + topk=topk, + ratio=ratio, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_kv, + max_seqlen_q=int(max_seqlen_q), + max_seqlen_k=int(max_seqlen_kv), + q_causal_offsets=q_causal_offsets, + return_softmax=return_softmax, + precision=precision, + ) + _refresh_thd_compact_cand_batch_offsets( + compact_workspace.cand_batch_offsets, cu_seqlens_q, ratio, q_causal_offsets + ) + compact_kwargs.update( + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_k=cu_seqlens_kv, + max_seqlen_q=int(max_seqlen_q), + max_seqlen_k=int(max_seqlen_kv), + ) + if compact_workspace is not None: + compact_kwargs.update( + cand_buffer=compact_workspace.cand_buffer, + cand_batch_offsets=compact_workspace.cand_batch_offsets, + out_indices=compact_workspace.out_indices, + out_logits=compact_workspace.out_logits, + ) + if return_softmax: + assert compact_workspace.softmax_out is not None + compact_kwargs["softmax_out"] = compact_workspace.softmax_out + if q_causal_offsets is not None: + compact_kwargs["q_causal_offsets"] = q_causal_offsets + elif compact_workspace is not None: + if not isinstance(compact_workspace, BSHDCompactIndexerWorkspace): + raise ValueError("BSHD compact dispatch requires a BSHDCompactIndexerWorkspace") + compact_workspace.validate( + q=q, k=k, topk=topk, ratio=ratio, return_softmax=return_softmax, precision=precision + ) + compact_kwargs.update( + cand_buffer=compact_workspace.cand_buffer, + out_indices=compact_workspace.out_indices, + out_logits=compact_workspace.out_logits, + microbatch_rows=0, + ) + if return_softmax: + assert compact_workspace.softmax_out is not None + compact_kwargs["softmax_out"] = compact_workspace.softmax_out + + kernel_q, kernel_k = q, k + if precision == "mxfp8": + mxfp8_workspace = compact_workspace.mxfp8 if compact_workspace is not None else None + cu_seqlens_q_scale_padded = None + cu_seqlens_k_scale_padded = None + if is_thd: + if mxfp8_workspace is not None: + cu_seqlens_q_scale_padded = mxfp8_workspace.cu_seqlens_q_scale_padded + cu_seqlens_k_scale_padded = mxfp8_workspace.cu_seqlens_k_scale_padded + assert cu_seqlens_q_scale_padded is not None + assert cu_seqlens_k_scale_padded is not None + refresh_indexer_mxfp8_scale_cu_seqlens( + cu_seqlens_q_scale_padded, cu_seqlens_q, q.shape[1] + ) + refresh_indexer_mxfp8_scale_cu_seqlens( + cu_seqlens_k_scale_padded, cu_seqlens_kv, 1 + ) + else: + cu_seqlens_q_scale_padded = make_indexer_mxfp8_scale_cu_seqlens( + cu_seqlens_q, q.shape[1] + ) + cu_seqlens_k_scale_padded = make_indexer_mxfp8_scale_cu_seqlens( + cu_seqlens_kv, 1 + ) + + kernel_q, q_scale = quantize_indexer_mxfp8( + q, + cu_seqlens=cu_seqlens_q, + cu_seqlens_scale_padded=cu_seqlens_q_scale_padded, + buffers=(mxfp8_workspace.q_buffers if mxfp8_workspace is not None else None), + out_scale=(mxfp8_workspace.q_scale if mxfp8_workspace is not None else None), + ) + kernel_k, k_scale = quantize_indexer_mxfp8( + k, + cu_seqlens=cu_seqlens_kv, + cu_seqlens_scale_padded=cu_seqlens_k_scale_padded, + buffers=(mxfp8_workspace.k_buffers if mxfp8_workspace is not None else None), + out_scale=(mxfp8_workspace.k_scale if mxfp8_workspace is not None else None), + ) + compact_kwargs.update(q_scale=q_scale, k_scale=k_scale, sf_vec_size=32) + if is_thd: + compact_kwargs.update( + cu_seqlens_q_scale_padded=cu_seqlens_q_scale_padded, + cu_seqlens_k_scale_padded=cu_seqlens_k_scale_padded, + ) + + if is_thd: + compact_result = compact_wrapper( + kernel_q, kernel_k.unsqueeze(1), w, top_k=topk, **compact_kwargs + ) + else: + compact_result = compact_wrapper( + kernel_q, kernel_k.unsqueeze(2), w, top_k=topk, **compact_kwargs + ) + + topk_indices = compact_result["indices"] + compact_logits = compact_result["logits"] + compact_softmax = compact_result["softmax"] if return_softmax else None + if compact_workspace is not None: + returned_buffers = ( + ("indices", topk_indices, compact_workspace.out_indices), + ("logits", compact_logits, compact_workspace.out_logits), + ) + if return_softmax: + assert compact_workspace.softmax_out is not None + returned_buffers += (("softmax", compact_softmax, compact_workspace.softmax_out),) + for name, actual, expected in returned_buffers: + if actual.data_ptr() != expected.data_ptr(): + layout = "THD" if is_thd else "BSHD" + raise RuntimeError( + f"cuDNN compact {layout} {name} did not alias " + "the caller-owned workspace buffer" + ) + + topk_indices = topk_indices.int() + topk_length = (topk_indices >= 0).sum(dim=-1).int() + if is_thd: + return topk_indices, topk_length, None, compact_softmax + b, sq = q.shape[:2] + return (topk_indices.view(b, sq, topk), topk_length.view(b, sq), None, compact_softmax) + + if is_thd: # Kernel wants k as 3-D ``(total_k, h_kv, idx_hd)``. forward_kwargs = dict( cu_seqlens_q=cu_seqlens_q, @@ -559,10 +1389,6 @@ def _indexer_topk_core( ) seq_lens = torch.where(row_valid, seq_lens, torch.zeros_like(seq_lens)) else: - if k.shape[1] == 0: - raise ValueError("indexer_topk requires at least one K row.") - - _ensure_dsa_namespace() # Kernel wants k as 4-D ``(b, sk, h_kv, idx_hd)``. scores = _DSA.indexer_forward_wrapper(q, k.unsqueeze(2), w, ratio=ratio)[ "scores" @@ -602,8 +1428,8 @@ def _indexer_topk_core( # ---------------- Layout-specific output reshape -------------------- if is_thd: - return topk_indices.int(), topk_length, scores - return (topk_indices.view(b, sq, topk).int(), topk_length.view(b, sq), scores) + return topk_indices.int(), topk_length, scores, None + return (topk_indices.view(b, sq, topk).int(), topk_length.view(b, sq), scores, None) def indexer_topk( @@ -619,11 +1445,16 @@ def indexer_topk( max_seqlen_q: Optional[int] = None, max_seqlen_kv: Optional[int] = None, q_causal_offsets: Optional[Tensor] = None, -) -> Tuple[Tensor, Tensor]: + compact_workspace: BSHDCompactIndexerWorkspace | THDCompactIndexerWorkspace | None = None, + precision: str = "bf16", + deterministic: bool = False, + return_softmax: bool = False, +) -> Tuple[Tensor, Tensor] | Tuple[Tensor, Tensor, Optional[Tensor]]: """Score + top-K selection for inference (no KL loss, no backward). - Built on cuDNN Frontend's CuTe-DSL indexer forward kernel followed by - TRT-LLM's radix top-K kernel. + Uses cuDNN Frontend's combined forward + Top-K kernel on SM10x when + available, and otherwise falls back to the dense indexer forward followed + by the standalone radix Top-K kernel. Args: q_indexer: SBHD ``(sq, b, idx_nh, idx_hd)`` / @@ -642,6 +1473,14 @@ def indexer_topk( max_seqlen_kv: THD only — per-batch max KV length. q_causal_offsets: THD only — optional ``(B,)`` int32 CUDA tensor. Entry ``b`` is the sequence-relative position of that segment's first Q. + compact_workspace: caller-owned compact BSHD or THD buffers prepared + during eager warmup. A matching workspace is required while + capturing the compact path. + deterministic: resolve exact-value ties at the K-th boundary toward + the smallest local KV indices. The output slot order remains + unspecified. + return_softmax: also return the compact kernel's Top-K softmax. The + third return is ``None`` when compact dispatch is unavailable. Returns: SBHD: ``(topk_indices (b, sq, topk), topk_length (b, sq))`` int32 @@ -674,7 +1513,7 @@ def indexer_topk( k = k_indexer.permute(1, 0, 2).contiguous() w = weights.permute(1, 0, 2).contiguous() - topk_indices, topk_length, _ = _indexer_topk_core( + topk_indices, topk_length, _, compact_softmax = _indexer_topk_core( q, k, w, @@ -685,7 +1524,14 @@ def indexer_topk( max_seqlen_q=int(max_seqlen_q) if max_seqlen_q is not None else None, max_seqlen_kv=int(max_seqlen_kv) if max_seqlen_kv is not None else None, q_causal_offsets=q_causal_offsets, + use_compact=True, + return_softmax=return_softmax, + compact_workspace=compact_workspace, + precision=precision, + deterministic=deterministic, ) + if return_softmax: + return topk_indices, topk_length, compact_softmax return topk_indices, topk_length @@ -1052,6 +1898,9 @@ def forward( max_seqlen_compressed_idx: Optional[int], # indexer K max compressed_kv: Optional[Tensor] = None, # THD only — pre-packed compressed KV cu_seqlens_q_unpadded: Optional[Tensor] = None, # THD only — unpadded Q cu_seqlens + compact_workspace: BSHDCompactIndexerWorkspace | THDCompactIndexerWorkspace | None = None, + indexer_precision: str = "bf16", + deterministic: bool = False, ) -> Tuple[Tensor, Tensor]: """Fused forward: indexer scoring, sparse attention, KL loss, and indexer backward.""" _ensure_dsa_namespace() @@ -1089,24 +1938,45 @@ def forward( else: w_indexer_scaled = w_indexer - # ---- 2. Indexer scoring + top-K (with scores retained). --------------- + # ---- 2. Indexer scoring + top-K. ------------------------------------- + # Compact THD requires physical K coverage and enough per-sequence K + # rows for floor-compressed tails. Keep the original K/layout for the + # loss and its backward, and pad only the non-differentiable Top-K call. + topk_k_indexer = k_indexer_flat + topk_cu_seqlens_k = cu_seqlens_compressed_idx + topk_max_seqlen_k = max_seqlen_compressed_idx + if is_thd and thd_compact_indexer_available( + q_indexer_flat, k_indexer_flat, indexer_precision + ): + assert cu_seqlens_q is not None + assert cu_seqlens_compressed_idx is not None + compact_cu_seqlens_k, source_row_map = build_thd_compact_k_layout( + cu_seqlens_q, cu_seqlens_compressed_idx, k_indexer_flat.shape[0], ratio + ) + topk_k_indexer = pack_thd_compact_k(k_indexer_flat, source_row_map) + topk_cu_seqlens_k = compact_cu_seqlens_k + topk_max_seqlen_k = int(max_seqlen_compressed_idx) + 2 + # Pass the original ``indexer_topk`` (not min(indexer_topk, n_comp)) so - # that the output is always padded to a fixed size. flash_mla_sparse_fwd - # requires a consistent TopK dimension; _indexer_topk_core handles the - # case where sk < topk internally (selects min(topk, sk) values, then - # pads to topk with -1). - topk_indices_cmp, _, indexer_scores = _indexer_topk_core( + # that the output is always padded to a fixed size. Top-K dispatch is + # independent of the configured loss: SM100 always uses the compact + # wrapper, while an enabled dense KL loss separately recomputes the + # full score tensor it needs below. + topk_indices_cmp, _, indexer_scores, compact_predict = _indexer_topk_core( q_indexer_flat, - k_indexer_flat, + topk_k_indexer, w_indexer_scaled, indexer_topk, ratio, cu_seqlens_q=cu_seqlens_q, - cu_seqlens_kv=cu_seqlens_compressed_idx, + cu_seqlens_kv=topk_cu_seqlens_k, max_seqlen_q=int(max_seqlen_q) if max_seqlen_q is not None else None, - max_seqlen_kv=( - int(max_seqlen_compressed_idx) if max_seqlen_compressed_idx is not None else None - ), + max_seqlen_kv=(int(topk_max_seqlen_k) if topk_max_seqlen_k is not None else None), + use_compact=True, + return_softmax=loss_coeff > 0 and sparse_loss, + compact_workspace=compact_workspace, + precision=indexer_precision, + deterministic=deterministic, ) # ---- 3. Combine indices (indexer first, then window) + globalize. ---- @@ -1162,7 +2032,7 @@ def forward( # The caller only passes cu_seqlens_q_unpadded when it differs from # cu_seqlens_q (checked via data_ptr), so no GPU→CPU sync is needed. padding_row_mask: Optional[Tensor] = None # True = padding (excluded from loss) - if is_thd and cu_seqlens_q_unpadded is not None: + if loss_coeff > 0 and is_thd and cu_seqlens_q_unpadded is not None: real_seg_lens = cu_seqlens_q_unpadded[1:] - cu_seqlens_q_unpadded[:-1] row_idx = torch.arange(total_q, device=query.device, dtype=torch.int32) row_batch_ids = batch_of_row(cu_seqlens_q, total_q=total_q) @@ -1174,92 +2044,123 @@ def forward( padding_row_mask = pos_in_seg >= real_len_per_row # ---- 5. Derive predict from indexer_scores, compute target. ---------- - # Layout-specific attn tensors (detached — loss is not differentiable - # through them). - if is_thd: - assert compressed_kv is not None, "compressed_kv is required for THD" - q_attn_det = query.detach() - k_attn_compressed_det = compressed_kv.detach() - lse_indexer_det = lse_indexer.detach() + if loss_coeff <= 0: + indexer_loss = torch.zeros((), device=query.device, dtype=torch.float32) else: - q_attn_det = query.detach().permute(1, 0, 2, 3).contiguous() - k_attn_compressed_det = kv_full[kv_offset:].detach().permute(1, 0, 2).contiguous() - lse_indexer_det = lse_indexer.reshape(sq, b, np_).permute(1, 0, 2) - - # Invalidate padding rows for the loss/backward path. The sparse - # attention (steps 3-4) has already built global_idxs from the - # original topk_indices_cmp, so this mutation only affects steps 5-7. - if padding_row_mask is not None: - topk_indices_cmp = topk_indices_cmp.clone() - topk_indices_cmp[padding_row_mask] = -1 - indexer_scores = indexer_scores.clone() - indexer_scores[padding_row_mask] = float('-inf') - - if sparse_loss: - # Derive predict: gather topk scores from indexer_scores → softmax. - safe_indices = topk_indices_cmp.clamp(min=0).long() - gathered_scores = torch.gather(indexer_scores, dim=-1, index=safe_indices) - gathered_scores = torch.where( - topk_indices_cmp >= 0, gathered_scores, torch.finfo(torch.float32).min - ) - predict = torch.softmax(gathered_scores, dim=-1) - - # THD: _compute_attn_target's kernel addresses K by flat ids over - # the packed (total_k, D) buffer, so promote per-segment-local - # indices to flat-global against cu_seqlens_compressed_idx. + # Layout-specific attention tensors are detached because the loss + # is not differentiable through the attention-score target. if is_thd: - topk_for_target = local_to_global_flat( - topk_indices_cmp, - batch_size=-1, - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_kv=cu_seqlens_compressed_idx, - ) + assert compressed_kv is not None, "compressed_kv is required for THD" + q_attn_det = query.detach() + k_attn_compressed_det = compressed_kv.detach() + lse_indexer_det = lse_indexer.detach() else: - topk_for_target = topk_indices_cmp + q_attn_det = query.detach().permute(1, 0, 2, 3).contiguous() + k_attn_compressed_det = kv_full[kv_offset:].detach().permute(1, 0, 2).contiguous() + lse_indexer_det = lse_indexer.reshape(sq, b, np_).permute(1, 0, 2) - target = _compute_attn_target( - q_attn_det, - k_attn_compressed_det, - lse_indexer_det, - topk_for_target, - softmax_scale, - qhead_per_kv_head=np_, - topk_indices_global=is_thd, - ) + # Invalidate padding rows for the loss/backward path. The sparse + # attention has already consumed the unmodified selected indices. + if padding_row_mask is not None: + topk_indices_cmp = topk_indices_cmp.clone() + topk_indices_cmp[padding_row_mask] = -1 + if indexer_scores is not None: + indexer_scores = indexer_scores.clone() + indexer_scores[padding_row_mask] = float('-inf') + if compact_predict is not None: + compact_predict = compact_predict.clone() + compact_predict[padding_row_mask] = 0 + + if sparse_loss: + if compact_predict is not None: + predict = compact_predict + else: + # Fallback: gather Top-K scores from the dense score tensor. + assert indexer_scores is not None + safe_indices = topk_indices_cmp.clamp(min=0).long() + gathered_scores = torch.gather(indexer_scores, dim=-1, index=safe_indices) + gathered_scores = torch.where( + topk_indices_cmp >= 0, gathered_scores, torch.finfo(torch.float32).min + ) + predict = torch.softmax(gathered_scores, dim=-1) + + # THD: _compute_attn_target's kernel addresses K by flat ids over + # the packed (total_k, D) buffer, so promote per-segment-local + # indices to flat-global against cu_seqlens_compressed_idx. + if is_thd: + topk_for_target = local_to_global_flat( + topk_indices_cmp, + batch_size=-1, + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_compressed_idx, + ) + else: + topk_for_target = topk_indices_cmp + + target = _compute_attn_target( + q_attn_det, + k_attn_compressed_det, + lse_indexer_det, + topk_for_target, + softmax_scale, + qhead_per_kv_head=np_, + topk_indices_global=is_thd, + ) - if loss_coeff > 0: indexer_loss = _kl_loss_from_target_predict( target, predict, topk_indices_cmp, loss_coeff, calculate_per_token_loss ) else: - indexer_loss = torch.zeros((), device=query.device, dtype=torch.float32) - else: - index_score = indexer_scores - index_lse = torch.logsumexp(indexer_scores, dim=-1) + if indexer_scores is None: + k_unsqueeze_dim = 1 if is_thd else 2 + dense_indexer_kwargs = {} + if is_thd: + dense_indexer_kwargs = dict( + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_compressed_idx, + max_seqlen_q=int(max_seqlen_q), + max_seqlen_kv=int(max_seqlen_compressed_idx), + ) + index_score, index_lse = _compute_dense_indexer_score( + q_indexer_flat, + k_indexer_flat.unsqueeze(k_unsqueeze_dim), + w_indexer, + qhead_per_kv_head=idx_nh, + indexer_softmax_scale=indexer_softmax_scale, + ratio=ratio, + **dense_indexer_kwargs, + ) + else: + index_score = indexer_scores + index_lse = torch.logsumexp(indexer_scores, dim=-1) + if padding_row_mask is not None: + index_score = index_score.masked_fill( + padding_row_mask.unsqueeze(-1), float("-inf") + ) + index_lse = index_lse.masked_fill(padding_row_mask, float("-inf")) - k_unsqueeze_dim = 1 if is_thd else 2 - dense_attn_kwargs = {} - if is_thd: - dense_attn_kwargs = dict( - cu_seqlens_q=cu_seqlens_q, - cu_seqlens_kv=cu_seqlens_compressed_idx, - max_seqlen_q=int(max_seqlen_q), - max_seqlen_kv=int(max_seqlen_compressed_idx), + k_unsqueeze_dim = 1 if is_thd else 2 + dense_attn_kwargs = {} + if is_thd: + dense_attn_kwargs = dict( + cu_seqlens_q=cu_seqlens_q, + cu_seqlens_kv=cu_seqlens_compressed_idx, + max_seqlen_q=int(max_seqlen_q), + max_seqlen_kv=int(max_seqlen_compressed_idx), + ) + attn_score, attn_l1norm = _compute_dense_attn_score( + q_attn_det, + k_attn_compressed_det.unsqueeze(k_unsqueeze_dim), + lse_indexer_det, + qhead_per_kv_head=np_, + softmax_scale=softmax_scale, + ratio=ratio, + **dense_attn_kwargs, ) - attn_score, attn_l1norm = _compute_dense_attn_score( - q_attn_det, - k_attn_compressed_det.unsqueeze(k_unsqueeze_dim), - lse_indexer_det, - qhead_per_kv_head=np_, - softmax_scale=softmax_scale, - ratio=ratio, - **dense_attn_kwargs, - ) - if padding_row_mask is not None: - attn_score = attn_score.masked_fill(padding_row_mask.unsqueeze(-1), 0) - attn_l1norm = attn_l1norm.masked_fill(padding_row_mask, 0) + if padding_row_mask is not None: + attn_score = attn_score.masked_fill(padding_row_mask.unsqueeze(-1), 0) + attn_l1norm = attn_l1norm.masked_fill(padding_row_mask, 0) - if loss_coeff > 0: indexer_loss = _kl_loss_from_dense_scores( attn_score, attn_l1norm, @@ -1268,8 +2169,6 @@ def forward( loss_coeff, calculate_per_token_loss, ) - else: - indexer_loss = torch.zeros((), device=query.device, dtype=torch.float32) # ---- 6. Eagerly compute indexer backward (grad_loss=1). ------------ # The actual grad_loss scaling is deferred to backward (when @@ -1279,13 +2178,12 @@ def forward( # internally; since masked rows contribute 0, multiplying back by # total_q still yields the correct real-token sum — and avoids a # GPU→CPU sync that would break CUDA graph capture. - indexer_loss_coeff = loss_coeff - if calculate_per_token_loss: - indexer_loss_coeff = loss_coeff * (total_q if is_thd else b * sq) - - unit_grad_loss = torch.ones((), device=query.device, dtype=torch.float32) - if loss_coeff > 0: + indexer_loss_coeff = loss_coeff + if calculate_per_token_loss: + indexer_loss_coeff = loss_coeff * (total_q if is_thd else b * sq) + unit_grad_loss = torch.ones((), device=query.device, dtype=torch.float32) + if sparse_loss: attn_score_for_bwd = target.clone() index_score_for_bwd = predict.clone() @@ -1469,7 +2367,8 @@ def backward(ctx, grad_output, grad_loss): # cu_seqlens_q, cu_seqlens_kv, cu_seqlens_kv_full, # cu_seqlens_compressed_idx, # max_seqlen_q, max_seqlen_compressed_idx, - # compressed_kv, cu_seqlens_q_unpadded + # compressed_kv, cu_seqlens_q_unpadded, compact_workspace, + # indexer_precision, deterministic return ( grad_query, grad_kv_full, @@ -1494,6 +2393,9 @@ def backward(ctx, grad_output, grad_loss): None, None, None, + None, + None, + None, ) @@ -1525,6 +2427,7 @@ def forward( max_seqlen_q: int, indexer_layout: Tuple[Tensor, Tensor, Tensor], q_padding_mask: Optional[Tensor] = None, + compact_predict: Optional[Tensor] = None, ) -> Tuple[Tensor, Tensor]: """Run fused sparse attention using caller-supplied top-k indices.""" _ensure_dsa_namespace() @@ -1553,15 +2456,25 @@ def forward( indexer_topk_idxs_for_loss = indexer_topk_idxs.masked_fill( q_padding_mask.unsqueeze(-1), -1 ) - weights_scaled = weights - if indexer_softmax_scale != 1.0: - weights_scaled = (weights.float() * indexer_softmax_scale).to(weights.dtype) - q_bshd, k_bsd, w_bsh, topk_bst = _thd_to_fake_bshd( - q_indexer, k_indexer, weights_scaled, indexer_topk_idxs_for_loss - ) - predict = _DSA.sparse_indexer_score_recompute_wrapper( - q_bshd, k_bsd, w_bsh, topk_bst, qhead_per_kv_head=idx_nh, topk_indices_global=True - )["predict"].squeeze(0) + if compact_predict is not None: + predict = compact_predict + if q_padding_mask is not None: + predict = predict.masked_fill(q_padding_mask.unsqueeze(-1), 0) + else: + weights_scaled = weights + if indexer_softmax_scale != 1.0: + weights_scaled = (weights.float() * indexer_softmax_scale).to(weights.dtype) + q_bshd, k_bsd, w_bsh, topk_bst = _thd_to_fake_bshd( + q_indexer, k_indexer, weights_scaled, indexer_topk_idxs_for_loss + ) + predict = _DSA.sparse_indexer_score_recompute_wrapper( + q_bshd, + k_bsd, + w_bsh, + topk_bst, + qhead_per_kv_head=idx_nh, + topk_indices_global=True, + )["predict"].squeeze(0) target = _compute_attn_target( query.detach(), compressed_kv.detach(), @@ -1736,6 +2649,7 @@ def backward(ctx, grad_output, grad_loss): None, None, None, + None, ) @@ -1764,6 +2678,9 @@ def fused_csa_indexer_sparse_attn( max_seqlen_compressed_idx: Optional[int] = None, compressed_kv: Optional[Tensor] = None, cu_seqlens_q_unpadded: Optional[Tensor] = None, + compact_workspace: BSHDCompactIndexerWorkspace | THDCompactIndexerWorkspace | None = None, + indexer_precision: str = "bf16", + deterministic: bool = False, ) -> Tuple[Tensor, Tensor]: """Path B (training): fused indexer (+KL loss) + sparse attention. @@ -1842,7 +2759,17 @@ def fused_csa_indexer_sparse_attn( so padding rows are excluded from the indexer KL loss and backward gradients. Ignored when ``None`` or when it equals ``cu_seqlens_q``. + compact_workspace: optional caller-owned BSHD or THD compact buffers. + A matching workspace is required during compact CUDA-graph capture; + prepare it outside capture with the matching + ``prepare_*_compact_indexer_workspace`` helper. + deterministic: resolve exact-value compact Top-K ties toward the + smallest local KV indices. This is normally sourced from + ``TransformerConfig.deterministic_mode``. """ + if indexer_precision == "mxfp8" and not sparse_loss and loss_coeff > 0: + raise ValueError("MXFP8 indexer loss supports only sparse indexer loss") + if cu_seqlens_q is not None: missing = [ name @@ -1884,14 +2811,23 @@ def fused_csa_indexer_sparse_attn( max_seqlen_compressed_idx, compressed_kv, cu_seqlens_q_unpadded, + compact_workspace, + indexer_precision, + deterministic, ) __all__ = [ + "BSHDCompactIndexerWorkspace", + "THDCompactIndexerWorkspace", "batch_of_row", + "build_thd_compact_k_layout", "build_flat_topk_idxs", "local_to_global_flat", "csa_sparse_attn", "indexer_topk", "fused_csa_indexer_sparse_attn", + "prepare_bshd_compact_indexer_workspace", + "prepare_thd_compact_indexer_workspace", + "pack_thd_compact_k", ] diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index edc21996907..7d65a071f14 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -341,6 +341,9 @@ class TransformerConfig(ModelParallelConfig): """Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using the top-k indices.""" + dsa_indexer_precision: Literal["bf16", "mxfp8"] = "bf16" + """Precision used only by the fused compact DSA indexer forward and Top-K.""" + dsa_kernel_backend: Literal["none", "tilelang", "cudnn"] = "none" """Optional fused ordinary-DSA kernel backend. Unsupported layouts use PyTorch fallback.""" @@ -1702,6 +1705,11 @@ def __post_init__(self): "cp_comm_type=allgather only." ) elif self.experimental_attention_variant == "dsv4_hybrid": + if self.dsa_indexer_precision not in ("bf16", "mxfp8"): + raise ValueError( + "dsa_indexer_precision must be 'bf16' or 'mxfp8', " + f"got {self.dsa_indexer_precision!r}" + ) assert self.multi_latent_attention, "DSv4 Hybrid requires multi_latent_attention." assert self.csa_compress_ratios is not None, "csa_compress_ratios must be set" mtp_layers = self.mtp_num_layers or 0 @@ -1724,6 +1732,12 @@ def __post_init__(self): assert not self.qk_clip, "QK clipping is not supported with DSv4 Hybrid Attention." self.hetereogenous_dist_checkpoint = True + uses_ratio4_indexer = 4 in self.csa_compress_ratios and not self.csa_dense_mode + indexer_loss_enabled = (self.dsa_indexer_loss_coeff or 0.0) > 0 + uses_mxfp8_indexer = uses_ratio4_indexer and self.dsa_indexer_precision == "mxfp8" + if uses_mxfp8_indexer and self.dsa_kernel_backend != "cudnn": + raise ValueError("MXFP8 DSA indexer precision requires dsa_kernel_backend='cudnn'") + if self.dsa_kernel_backend == "tilelang": raise ValueError( "dsv4_hybrid does not support dsa_kernel_backend='tilelang'; use 'cudnn' " @@ -1737,8 +1751,20 @@ def __post_init__(self): f"dsa_kernel_backend='cudnn' requires SM90+ (Hopper or later), " f"but current device has compute capability {sm[0]}.{sm[1]}." ) - uses_ratio4_indexer = 4 in self.csa_compress_ratios and not self.csa_dense_mode - indexer_loss_enabled = (self.dsa_indexer_loss_coeff or 0.0) > 0 + if uses_mxfp8_indexer: + if sm[0] < 10: + raise ValueError("MXFP8 compact DSA indexer requires SM100 or later") + if self.dsa_indexer_n_heads != 64 or self.dsa_indexer_head_dim != 128: + raise ValueError( + "MXFP8 compact DSA indexer requires dsa_indexer_n_heads=64 and " + "dsa_indexer_head_dim=128" + ) + if indexer_loss_enabled and not self.dsa_indexer_use_sparse_loss: + raise ValueError( + "MXFP8 DSA indexer loss supports only sparse loss; set " + "dsa_indexer_use_sparse_loss=True" + ) + if ( sm[0] == 9 and uses_ratio4_indexer @@ -1753,6 +1779,33 @@ def __post_init__(self): from cudnn import DSA + if sm[0] >= 10 and uses_ratio4_indexer: + compact_wrapper = getattr(DSA, "indexer_forward_top_k_wrapper", None) + required_parameters = {"deterministic"} + if uses_mxfp8_indexer: + required_parameters.update( + { + "precision", + "q_scale", + "k_scale", + "cu_seqlens_q_scale_padded", + "cu_seqlens_k_scale_padded", + "sf_vec_size", + } + ) + wrapper_parameters = ( + set(inspect.signature(compact_wrapper).parameters) + if callable(compact_wrapper) + else set() + ) + missing_parameters = required_parameters - wrapper_parameters + if missing_parameters: + raise ValueError( + "Fused DSA indexer requires a compatible cuDNN Frontend compact " + "wrapper; " + f"missing parameters: {', '.join(sorted(missing_parameters))}" + ) + if ( self.context_parallel_size > 1 or self.dynamic_context_parallel ) and uses_ratio4_indexer: diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index a168381eaa7..7a65d8164b7 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -4925,7 +4925,7 @@ def _add_experimental_attention_variant_args(parser): 'experimental_attention_variant="dsv4_hybrid" launches and to "none" ' 'otherwise.', ) - # Note: --dsa-indexer-{n-heads,head-dim,topk,loss-coeff,use-sparse-loss}, + # Note: --dsa-indexer-{n-heads,head-dim,topk,loss-coeff,use-sparse-loss,precision}, # --csa-window-size, --csa-compress-rotary-base, --csa-dense-mode are # auto-generated by ArgumentGroupFactory from TransformerConfig fields # (none of them are in the exclude list at line 2500-2576). diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py index 39bd99624db..9a0e53334ee 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -1,6 +1,6 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest import torch @@ -687,6 +687,57 @@ def test_constructor(self, compress_ratio): elif compress_ratio == 128: assert csa.indexer is None + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_bshd_compact_workspace_prepared_in_warmup_and_reused_in_capture(self, compress_ratio): + """CUDA capture reuses the BSHD workspace created during eager warmup.""" + if compress_ratio != 4: + pytest.skip("compact indexer workspace applies only to ratio-4 layers") + + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=self._get_layer_number(compress_ratio), + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, + ).cuda() + previous_cuda_graph_impl = csa.config.cuda_graph_impl + csa.config.cuda_graph_impl = "local" + q = torch.empty(8, 1, 64, 128, dtype=torch.bfloat16, device="cuda") + k = torch.empty(2, 1, 128, dtype=torch.bfloat16, device="cuda") + workspace = MagicMock() + workspace.matches_shape.return_value = True + + try: + with ( + patch( + "megatron.core.transformer.experimental_attention_variant.csa." + "bshd_compact_indexer_available", + return_value=True, + ), + patch( + "megatron.core.transformer.experimental_attention_variant.csa." + "prepare_bshd_compact_indexer_workspace", + return_value=workspace, + ) as prepare_workspace, + patch.object(torch.cuda, "is_current_stream_capturing", side_effect=[False, True]), + ): + warmup_workspace = csa._get_bshd_compact_indexer_workspace(q, k, topk=8, ratio=4) + capture_workspace = csa._get_bshd_compact_indexer_workspace(q, k, topk=8, ratio=4) + + assert warmup_workspace is workspace + assert capture_workspace is workspace + assert csa._bshd_compact_indexer_workspaces == [workspace] + prepare_workspace.assert_called_once() + prepared_q, prepared_k = prepare_workspace.call_args.args + assert prepared_q.shape == (1, 8, 64, 128) + assert prepared_k.shape == (1, 2, 128) + workspace.matches_shape.assert_called_once() + finally: + csa.config.cuda_graph_impl = previous_cuda_graph_impl + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_forward(self, compress_ratio): """Test forward pass with compressed attention.""" @@ -1104,7 +1155,7 @@ def _cu_seqlens(seg_lens, device='cpu'): class TestCsaThdIndexHelpers: """CSA THD index helpers — pure-Python, no GPU. Mirrors the - organisation of ``TestThdPureHelpers`` in ``test_dsa_kernels.py``: + organisation of ``TestThdPureHelpers`` in ``test_csa_kernels.py``: one mega-class with section comments per helper, since each helper only needs 2–3 tests and they share no fixtures. @@ -1324,7 +1375,7 @@ def test_thd_matches_sbhd_b1_equivalent(self): # # These integration tests exercise the THD branches of the full # Compressor / CSAIndexer / CompressedSparseAttention modules — the -# layer above the kernel-level THD tests in test_dsa_kernels.py and the +# layer above the kernel-level THD tests in test_csa_kernels.py and the # autograd-Function tests in test_attention_variant_dsa.py. # # Strategy: most tests use a B=1 single-segment THD input and compare @@ -1729,6 +1780,149 @@ def _make_thd_inputs(self, seg_lens): packed = _make_packed_seq_params_thd(seg_lens) return query, key, value, x, qr, packed + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_compact_workspace_reused_across_dynamic_boundaries(self): + """Dynamic sequence boundaries share one workspace for the same static geometry.""" + csa = self._build_csa(compress_ratio=4) + previous_cuda_graph_impl = csa.config.cuda_graph_impl + csa.config.cuda_graph_impl = "local" + q = torch.empty(64, 64, 128, dtype=torch.bfloat16, device="cuda") + k = torch.empty(16, 128, dtype=torch.bfloat16, device="cuda") + cu_q_variants = [ + torch.tensor([0, split, 64], dtype=torch.int32, device="cuda") for split in range(1, 33) + ] + cu_k = torch.tensor([0, 8, 16], dtype=torch.int32, device="cuda") + workspace = MagicMock() + workspace.matches.return_value = True + + try: + with ( + patch( + "megatron.core.transformer.experimental_attention_variant.csa." + "thd_compact_indexer_available", + return_value=True, + ), + patch( + "megatron.core.transformer.experimental_attention_variant.csa." + "prepare_thd_compact_indexer_workspace", + return_value=workspace, + ) as prepare_workspace, + patch.object( + torch.cuda, + "is_current_stream_capturing", + side_effect=[False] * len(cu_q_variants) + [True], + ), + ): + warmup_workspaces = [ + csa._get_thd_compact_indexer_workspace( + q, + k, + topk=8, + ratio=4, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=64, + max_seqlen_k=8, + ) + for cu_q in cu_q_variants + ] + capture_workspace = csa._get_thd_compact_indexer_workspace( + q, + k, + topk=8, + ratio=4, + cu_seqlens_q=cu_q_variants[-1], + cu_seqlens_k=cu_k, + max_seqlen_q=64, + max_seqlen_k=8, + ) + + assert all(warmup_workspace is workspace for warmup_workspace in warmup_workspaces) + assert capture_workspace is workspace + assert csa._thd_compact_indexer_workspaces == [workspace] + prepare_workspace.assert_called_once() + assert workspace.matches.call_count == len(cu_q_variants) + assert all( + "check_sequence_values" not in call.kwargs + for call in workspace.matches.call_args_list + ) + finally: + csa.config.cuda_graph_impl = previous_cuda_graph_impl + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_compact_workspace_capture_requires_matching_warmup(self): + """Capture must fail clearly when eager warmup did not seed a matching workspace.""" + csa = self._build_csa(compress_ratio=4) + previous_cuda_graph_impl = csa.config.cuda_graph_impl + csa.config.cuda_graph_impl = "local" + q = torch.empty(8, 64, 128, dtype=torch.bfloat16, device="cuda") + k = torch.empty(2, 128, dtype=torch.bfloat16, device="cuda") + cu_q = torch.tensor([0, 8], dtype=torch.int32, device="cuda") + cu_k = torch.tensor([0, 2], dtype=torch.int32, device="cuda") + + try: + with ( + patch( + "megatron.core.transformer.experimental_attention_variant.csa." + "thd_compact_indexer_available", + return_value=True, + ), + patch.object(torch.cuda, "is_current_stream_capturing", return_value=True), + ): + with pytest.raises(ValueError, match="Run eager warmup before capture"): + csa._get_thd_compact_indexer_workspace( + q, + k, + topk=8, + ratio=4, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=8, + max_seqlen_k=2, + ) + finally: + csa.config.cuda_graph_impl = previous_cuda_graph_impl + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_thd_compact_workspace_unavailable_keeps_dense_fallback(self): + """Capture needs no compact workspace when the compact kernel cannot dispatch.""" + csa = self._build_csa(compress_ratio=4) + previous_cuda_graph_impl = csa.config.cuda_graph_impl + csa.config.cuda_graph_impl = "local" + q = torch.empty(8, 64, 128, dtype=torch.bfloat16, device="cuda") + k = torch.empty(2, 128, dtype=torch.bfloat16, device="cuda") + cu_q = torch.tensor([0, 8], dtype=torch.int32, device="cuda") + cu_k = torch.tensor([0, 2], dtype=torch.int32, device="cuda") + + try: + with ( + patch( + "megatron.core.transformer.experimental_attention_variant.csa." + "thd_compact_indexer_available", + return_value=False, + ), + patch( + "megatron.core.transformer.experimental_attention_variant.csa." + "prepare_thd_compact_indexer_workspace" + ) as prepare_workspace, + patch.object(torch.cuda, "is_current_stream_capturing", return_value=True), + ): + workspace = csa._get_thd_compact_indexer_workspace( + q, + k, + topk=8, + ratio=4, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=8, + max_seqlen_k=2, + ) + + assert workspace is None + prepare_workspace.assert_not_called() + finally: + csa.config.cuda_graph_impl = previous_cuda_graph_impl + # ---- Path A (compress_ratio=128: indexer disabled, all-compressed) ---- @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py index 7a189e59ca6..11dd6e8a2a4 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_cp_utils.py @@ -201,26 +201,53 @@ def fake_apply(hidden, boundary, cu, global_start, ratio, d_comp, c_cap): assert torch.equal(rank_rows, torch.tensor([0, 1, 2, 3, 10, 11, 12, 13], dtype=torch.int32)) -def test_compute_cp_indexer_topk_passes_offsets_without_repacking_k(monkeypatch): +def test_compute_cp_indexer_topk_pads_compact_k_contract(monkeypatch): topk_calls = [] def fake_indexer_topk( - q, k, _weights, *, topk, cu_seqlens_q, cu_seqlens_kv, q_causal_offsets, **_ + q, + k, + _weights, + *, + topk, + cu_seqlens_q, + cu_seqlens_kv, + q_causal_offsets, + deterministic, + precision, + compact_workspace, + return_softmax, + **_, ): topk_calls.append( - (k.clone(), cu_seqlens_q.clone(), cu_seqlens_kv.clone(), q_causal_offsets.clone()) + { + "k": k.clone(), + "cu_q": cu_seqlens_q.clone(), + "cu_k": cu_seqlens_kv.clone(), + "offsets": q_causal_offsets.clone(), + "deterministic": deterministic, + "precision": precision, + "workspace": compact_workspace, + "return_softmax": return_softmax, + } ) - return torch.full((q.shape[0], int(topk)), len(topk_calls), dtype=torch.int32), None + indices = torch.full((q.shape[0], int(topk)), len(topk_calls), dtype=torch.int32) + if return_softmax: + predict = torch.full((q.shape[0], int(topk)), 0.5, dtype=torch.float32) + return indices, None, predict + return indices, None monkeypatch.setattr(csa_cp_utils, "indexer_topk", fake_indexer_topk) q = torch.randn(8, 2) weights = torch.randn(8, 1) - k_seq = torch.arange(8, dtype=torch.float32).reshape(4, 2) + # The final row models fixed CP capacity beyond the valid compressed rows. + k_seq = torch.arange(10, dtype=torch.float32).reshape(5, 2) cu_q = torch.tensor([0, 5, 13, 20], dtype=torch.int32) cu_comp = torch.tensor([0, 1, 3, 4], dtype=torch.int32) + workspace = object() - out, metadata = compute_cp_indexer_topk( + out, metadata, compact_predict = compute_cp_indexer_topk( q, weights, k_seq, @@ -232,23 +259,36 @@ def fake_indexer_topk( indexer_softmax_scale=0.5, max_seqlen_q=8, use_fused=True, + deterministic=True, + precision="mxfp8", + compact_workspace=workspace, + return_softmax=True, ) + call = topk_calls[0] assert torch.equal(out, torch.ones(8, 2, dtype=torch.int32)) - assert torch.equal(topk_calls[0][0], k_seq) - assert torch.equal(topk_calls[0][1], torch.tensor([0, 0, 6, 8, 8], dtype=torch.int32)) - assert torch.equal(topk_calls[0][2], torch.tensor([0, 1, 3, 4, 4], dtype=torch.int32)) - assert torch.equal(topk_calls[0][3], torch.tensor([0, 2, 0, 0], dtype=torch.int32)) + torch.testing.assert_close(compact_predict, torch.full((8, 2), 0.5)) + expected_k = torch.tensor( + [[0, 1], [0, 0], [0, 0], [2, 3], [4, 5], [0, 0], [6, 7], [0, 0]], dtype=torch.float32 + ) + assert torch.equal(call["k"], expected_k) + assert torch.equal(call["cu_q"], torch.tensor([0, 0, 6, 8, 8], dtype=torch.int32)) + assert torch.equal(call["cu_k"], torch.tensor([0, 3, 6, 8, 8], dtype=torch.int32)) + assert torch.equal(call["offsets"], torch.tensor([0, 2, 0, 0], dtype=torch.int32)) + assert call["deterministic"] is True + assert call["precision"] == "mxfp8" + assert call["workspace"] is workspace + assert call["return_softmax"] is True metadata_q, metadata_k, metadata_offsets = metadata - assert torch.equal(metadata_q, topk_calls[0][1]) - assert torch.equal(metadata_k, topk_calls[0][2]) - assert torch.equal(metadata_offsets, topk_calls[0][3]) + assert torch.equal(metadata_q, call["cu_q"]) + assert torch.equal(metadata_k, torch.tensor([0, 1, 3, 4, 5], dtype=torch.int32)) + assert torch.equal(metadata_offsets, call["offsets"]) assert compute_cp_indexer_topk( q, weights, k_seq[:0], cu_q, cu_comp, 2, 4, 2, 1.0, 10, True - ) == (None, None) + ) == (None, None, None) assert compute_cp_indexer_topk( q, weights, k_seq, cu_q, torch.zeros_like(cu_comp), 2, 4, 2, 1.0, 3, True - ) == (None, None) + ) == (None, None, None) assert len(topk_calls) == 1 @@ -267,7 +307,7 @@ def fail_if_fused(*_args, **_kwargs): topk_width = 3 scale = 0.7 - actual, _ = compute_cp_indexer_topk( + actual, _, compact_predict = compute_cp_indexer_topk( q, weights, k, @@ -280,6 +320,7 @@ def fail_if_fused(*_args, **_kwargs): max_seqlen_q=8, use_fused=False, ) + assert compact_predict is None expected = torch.full((q.shape[0], topk_width), -1, dtype=torch.int32) for local_row, global_row in enumerate(range(7, 15)): diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_kernels.py index 796c64d44f1..529d7f72c86 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_csa_kernels.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_csa_kernels.py @@ -20,6 +20,7 @@ from __future__ import annotations +import inspect import math import sys import types @@ -45,6 +46,8 @@ fused_csa_indexer_sparse_attn, indexer_topk, local_to_global_flat, + prepare_bshd_compact_indexer_workspace, + prepare_thd_compact_indexer_workspace, ) # --------------------------------------------------------------------------- @@ -287,6 +290,28 @@ def fake_compactify(global_idxs): ), "(b) length tensor differs between CPU fallback and cuDNN kernel" +# --------------------------------------------------------------------------- +# build_thd_compact_k_layout +# --------------------------------------------------------------------------- + + +class TestBuildThdCompactKLayout: + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compiled_source_row_map_preserves_k_and_appends_padding(self): + cu_seqlens_q = torch.tensor([0, 128], dtype=torch.int32, device='cuda') + cu_seqlens_k = torch.tensor([0, 512], dtype=torch.int32, device='cuda') + + compact_cu_seqlens_k, source_row_map = dk.build_thd_compact_k_layout( + cu_seqlens_q, cu_seqlens_k, total_k_rows=512, ratio=4 + ) + + expected_source_row_map = torch.arange(513, dtype=torch.int64, device='cuda') + assert torch.equal( + compact_cu_seqlens_k, torch.tensor([0, 513], dtype=torch.int32, device='cuda') + ) + assert torch.equal(source_row_map, expected_source_row_map) + + # --------------------------------------------------------------------------- # _kl_loss_from_target_predict # --------------------------------------------------------------------------- @@ -706,6 +731,7 @@ def fake_filtered_topk(scores_flat, seq_lens, top_k, next_n, return_val): return {'indices': torch.zeros(n_rows, top_k, dtype=torch.int32, device='cuda')} fake_dsa = MagicMock() + fake_dsa.indexer_forward_top_k_wrapper = None fake_dsa.indexer_forward_wrapper.side_effect = fake_indexer_forward fake_dsa.indexer_top_k_wrapper.side_effect = fake_filtered_topk dk._DSA = fake_dsa @@ -758,6 +784,7 @@ def fake_filtered_topk(scores_flat, seq_lens, top_k, next_n, return_val): kernel_indices2 = torch.zeros(b2 * sq2, sk2, dtype=torch.int32, device='cuda') fake_dsa_b = MagicMock() + fake_dsa_b.indexer_forward_top_k_wrapper = None fake_dsa_b.indexer_forward_wrapper.return_value = {'scores': scores2} fake_dsa_b.indexer_top_k_wrapper.return_value = {'indices': kernel_indices2} dk._DSA = fake_dsa_b @@ -782,6 +809,7 @@ def fake_indexer_forward_c(q_bshd, k_bshd, w_bsh, ratio): return {'scores': torch.zeros(b3, sq3, sk3, dtype=torch.float32, device='cuda')} fake_dsa_c = MagicMock() + fake_dsa_c.indexer_forward_top_k_wrapper = None fake_dsa_c.indexer_forward_wrapper.side_effect = fake_indexer_forward_c fake_dsa_c.indexer_top_k_wrapper.return_value = { 'indices': torch.zeros(b3 * sq3, sk3, dtype=torch.int32, device='cuda') @@ -796,6 +824,52 @@ def fake_indexer_forward_c(q_bshd, k_bshd, w_bsh, ratio): captured_w['w'].float(), expected_w.float(), atol=1e-2, rtol=1e-2 ), "(c) weights were not pre-scaled by indexer_softmax_scale" + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_indexer_topk_uses_compact_bf16_wrapper_on_sm10x(self, reset_lazy_kernel_state): + sq, b, idx_nh, idx_hd = 4, 2, 4, 64 + sk, topk, ratio = 3, 5, 4 + scale = 0.125 + q = torch.randn(sq, b, idx_nh, idx_hd, dtype=torch.bfloat16, device='cuda') + k = torch.randn(sk, b, idx_hd, dtype=torch.bfloat16, device='cuda') + w = torch.full((sq, b, idx_nh), 8.0, dtype=torch.bfloat16, device='cuda') + + fake_dsa = MagicMock(name='_DSA_compact_topk_stub') + + def fake_compact(q_bshd, k_bshd, w_bsh, top_k, **kwargs): + indices = torch.full((b, sq, top_k), -1, dtype=torch.int32, device=q_bshd.device) + indices[..., :2] = torch.tensor([0, 1], dtype=torch.int32, device=q_bshd.device) + return { + 'indices': indices, + 'logits': torch.zeros(b, sq, top_k, dtype=torch.float32, device=q_bshd.device), + } + + fake_dsa.indexer_forward_top_k_wrapper.side_effect = fake_compact + dk._DSA = fake_dsa + + with patch.object(torch.cuda, 'get_device_capability', return_value=(10, 0)): + indices, lengths = indexer_topk( + q, k, w, topk=topk, ratio=ratio, indexer_softmax_scale=scale + ) + + assert indices.shape == (b, sq, topk) + assert torch.all(indices[..., :2] == torch.tensor([0, 1], device='cuda')) + assert torch.all(indices[..., 2:] == -1) + assert torch.all(lengths == 2) + fake_dsa.indexer_forward_wrapper.assert_not_called() + fake_dsa.indexer_top_k_wrapper.assert_not_called() + + compact_call = fake_dsa.indexer_forward_top_k_wrapper.call_args + q_bshd, k_bshd, w_bsh = compact_call.args[:3] + assert q_bshd.shape == (b, sq, idx_nh, idx_hd) + assert k_bshd.shape == (b, sk, 1, idx_hd) + expected_w = (w.float() * scale).to(torch.bfloat16).permute(1, 0, 2).contiguous() + assert torch.equal(w_bsh, expected_w) + assert compact_call.kwargs['top_k'] == topk + assert compact_call.kwargs['ratio'] == ratio + assert compact_call.kwargs['precision'] == 'bf16' + assert compact_call.kwargs['return_softmax'] is False + assert compact_call.kwargs['topk_indices_global'] is False + # --------------------------------------------------------------------------- # csa_sparse_attn / CSASparseAttnFunc forward (mocked) @@ -908,6 +982,7 @@ def _install_full_dsa_mock( target_fn = predict_fn fake_dsa = MagicMock(name='_DSA_full_stub') + fake_dsa.indexer_forward_top_k_wrapper = None def fake_indexer_forward(q_bshd, k_bshd, w_bsh, ratio): return {'scores': torch.zeros(b, sq, n_comp, dtype=torch.float32, device=q_bshd.device)} @@ -1025,6 +1100,7 @@ def _install_full_dsa_mock_dense( ) fake_dsa = MagicMock(name='_DSA_full_dense_stub') + fake_dsa.indexer_forward_top_k_wrapper = None def fake_indexer_forward(q_bshd, k_bshd, w_bsh, ratio): return {'scores': torch.zeros(b, sq, n_comp, dtype=torch.float32, device=q_bshd.device)} @@ -1194,6 +1270,97 @@ def test_indexer_loss_formula(self, loss_coeff, target_kind, expected, reset_laz indexer_loss, torch.tensor(expected, device='cuda'), rtol=1e-5, atol=1e-5 ), f"got {indexer_loss.item()}, expected {expected}" + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_sparse_loss_uses_compact_topk_softmax_on_sm10x(self, reset_lazy_kernel_state): + s = self.SHAPES + topk = 2 + inputs = self._make_inputs() + fake_dsa, _ = _install_full_dsa_mock( + b=s['b'], + sq=s['sq'], + np_=s['np_'], + d=s['d'], + n_comp=s['n_comp'], + idx_nh=s['idx_nh'], + target_fn=lambda B, S, K, dev: _peaked_dist(B, S, K, dev, peak_idx=0), + ) + + def fake_compact(q_bshd, k_bshd, w_bsh, top_k, **kwargs): + indices = torch.arange(top_k, dtype=torch.int32, device=q_bshd.device) + indices = indices.view(1, 1, top_k).expand(s['b'], s['sq'], top_k).contiguous() + predict = torch.tensor([0.25, 0.75], dtype=torch.float32, device=q_bshd.device) + predict = predict.view(1, 1, top_k).expand(s['b'], s['sq'], top_k).contiguous() + return {'indices': indices, 'logits': predict.log(), 'softmax': predict} + + fake_dsa.indexer_forward_top_k_wrapper = MagicMock(side_effect=fake_compact) + with patch.object(torch.cuda, 'get_device_capability', return_value=(10, 0)): + _, indexer_loss = fused_csa_indexer_sparse_attn( + **inputs, + indexer_topk=topk, + ratio=4, + softmax_scale=0.5, + loss_coeff=1.0, + sparse_loss=True, + kv_offset=s['skv'] - s['n_comp'], + deterministic=True, + ) + + # target = delta_0 and compact predict[0] = 0.25, so KL = log(4). + assert torch.allclose( + indexer_loss, torch.tensor(math.log(4), device='cuda'), rtol=1e-5, atol=1e-5 + ) + fake_dsa.indexer_forward_wrapper.assert_not_called() + fake_dsa.indexer_top_k_wrapper.assert_not_called() + compact_call = fake_dsa.indexer_forward_top_k_wrapper.call_args + assert compact_call.kwargs['precision'] == 'bf16' + assert compact_call.kwargs['return_softmax'] is True + assert compact_call.kwargs['topk_indices_global'] is False + assert compact_call.kwargs['deterministic'] is True + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_disabled_loss_uses_compact_without_softmax(self, reset_lazy_kernel_state): + """A disabled dense loss must not force dense-score materialization.""" + s = self.SHAPES + topk = 2 + inputs = self._make_inputs() + fake_dsa, _ = _install_full_dsa_mock( + b=s['b'], sq=s['sq'], np_=s['np_'], d=s['d'], n_comp=s['n_comp'], idx_nh=s['idx_nh'] + ) + + def fake_compact(q_bshd, _k_bshd, _w_bsh, top_k, **kwargs): + return { + 'indices': torch.zeros( + q_bshd.shape[0], q_bshd.shape[1], top_k, dtype=torch.int32, device='cuda' + ), + 'logits': torch.zeros( + q_bshd.shape[0], q_bshd.shape[1], top_k, dtype=torch.float32, device='cuda' + ), + } + + fake_dsa.indexer_forward_top_k_wrapper = MagicMock(side_effect=fake_compact) + with patch.object(torch.cuda, 'get_device_capability', return_value=(10, 0)): + _, indexer_loss = fused_csa_indexer_sparse_attn( + **inputs, + indexer_topk=topk, + ratio=4, + softmax_scale=0.5, + loss_coeff=0.0, + sparse_loss=False, + kv_offset=s['skv'] - s['n_comp'], + deterministic=True, + ) + + assert torch.equal(indexer_loss, torch.zeros_like(indexer_loss)) + fake_dsa.indexer_forward_wrapper.assert_not_called() + fake_dsa.indexer_top_k_wrapper.assert_not_called() + fake_dsa.sparse_indexer_score_recompute_wrapper.assert_not_called() + fake_dsa.sparse_attn_score_recompute_wrapper.assert_not_called() + fake_dsa.dense_attn_score_recompute_wrapper.assert_not_called() + fake_dsa.indexer_backward_wrapper.assert_not_called() + compact_call = fake_dsa.indexer_forward_top_k_wrapper.call_args + assert compact_call.kwargs['return_softmax'] is False + assert compact_call.kwargs['deterministic'] is True + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") def test_sparse_path_fwd_output_bwd_grads_and_topk_clamp(self, reset_lazy_kernel_state): """Combined coverage for the sparse-loss path's three non-numerical @@ -1342,6 +1509,47 @@ def make(*shape, dtype, rg=False): weights=weights, ) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dense_loss_uses_compact_topk_and_recomputes_scores(self, reset_lazy_kernel_state): + """Dense KL selection is independent from the compact Top-K dispatch.""" + s = self.SHAPES + inputs = self._make_inputs() + fake_dsa, _ = _install_full_dsa_mock_dense( + b=s['b'], sq=s['sq'], np_=s['np_'], d=s['d'], n_comp=s['n_comp'], idx_nh=s['idx_nh'] + ) + + def fake_compact(q_bshd, _k_bshd, _w_bsh, top_k, **kwargs): + return { + 'indices': torch.zeros( + q_bshd.shape[0], q_bshd.shape[1], top_k, dtype=torch.int32, device='cuda' + ), + 'logits': torch.zeros( + q_bshd.shape[0], q_bshd.shape[1], top_k, dtype=torch.float32, device='cuda' + ), + } + + fake_dsa.indexer_forward_top_k_wrapper = MagicMock(side_effect=fake_compact) + with patch.object(torch.cuda, 'get_device_capability', return_value=(10, 0)): + _, loss = fused_csa_indexer_sparse_attn( + **inputs, + indexer_topk=2, + ratio=4, + softmax_scale=0.5, + indexer_softmax_scale=0.125, + loss_coeff=1.0, + sparse_loss=False, + kv_offset=s['skv'] - s['n_comp'], + deterministic=True, + ) + + assert torch.equal(loss, torch.zeros_like(loss)) + fake_dsa.indexer_forward_wrapper.assert_not_called() + fake_dsa.indexer_forward_top_k_wrapper.assert_called_once() + fake_dsa.dense_indexer_score_recompute_wrapper.assert_called_once() + compact_call = fake_dsa.indexer_forward_top_k_wrapper.call_args + assert compact_call.kwargs['return_softmax'] is False + assert compact_call.kwargs['deterministic'] is True + @pytest.mark.parametrize( "loss_coeff, target_kind, expected", [ @@ -1611,6 +1819,46 @@ def _ref_indexer_full_score( return torch.where(valid, s, torch.full_like(s, float('-inf'))) +def _dequantize_thd_indexer_mxfp8( + data: torch.Tensor, + packed_scale: torch.Tensor, + lengths: list[int], + cu_seqlens_scale_padded: torch.Tensor, +) -> torch.Tensor: + """Dequantize the packed THD MXFP8 representation for reference checks.""" + has_heads = data.ndim == 3 + num_heads = data.shape[-2] if has_heads else 1 + head_dim = data.shape[-1] + scale_groups = head_dim // 32 + _, _, padded_groups = packed_scale.shape + + logical_rows = packed_scale.shape[1] + rows = torch.arange(logical_rows, device=data.device).view(-1, 1) + groups = torch.arange(scale_groups, device=data.device).view(1, -1) + tile_idx = (rows // 128) * (padded_groups // 4) + groups // 4 + offsets = tile_idx * 512 + (rows % 32) * 16 + ((rows % 128) // 32) * 4 + groups % 4 + logical_scale = ( + packed_scale.view(torch.uint8) + .flatten()[offsets] + .contiguous() + .view(torch.float8_e8m0fnu) + .float() + ) + + result = torch.empty_like(data, dtype=torch.float32) + start = 0 + for batch, length in enumerate(lengths): + values = data[start : start + length].float().reshape(length, num_heads, scale_groups, 32) + scale_start = int(cu_seqlens_scale_padded[batch].item()) * num_heads + scale = logical_scale[scale_start : scale_start + length * num_heads].reshape( + length, num_heads, scale_groups, 1 + ) + dequantized = (values * scale).reshape(length, num_heads, head_dim) + result[start : start + length] = dequantized if has_heads else dequantized.squeeze(1) + start += length + return result + + def _ref_attn_full_score( q_bshd_fp32: torch.Tensor, # (B, Sq, H, D) k_bsd_fp32: torch.Tensor, # (B, Sk, D) — MQA @@ -2127,6 +2375,184 @@ def test_real_indexer_topk_set_matches_reference(self, dummy, reset_lazy_kernel_ f"actual {sorted(actual_set)} vs ref {sorted(ref_set)}" ) + @pytest.mark.parametrize("precision", ["bf16", "mxfp8"]) + def test_real_thd_compact_topk_cuda_graph_capture(self, precision, reset_lazy_kernel_state): + """THD compact Top-K follows the warmup/buffer/replay contract.""" + _skip_if_real_kernels_unavailable() + if torch.cuda.get_device_capability()[0] < 10: + pytest.skip("compact THD indexer forward + Top-K requires SM100+") + + from cudnn import DSA + + if not hasattr(DSA, 'compress_topk_cand_buffer_size_thd'): + pytest.skip("installed cuDNN Frontend lacks the compact THD workspace helper") + compact_wrapper = getattr(DSA, "indexer_forward_top_k_wrapper", None) + mxfp8_parameters = {"q_scale", "cu_seqlens_q_scale_padded", "cu_seqlens_k_scale_padded"} + if precision == "mxfp8" and mxfp8_parameters - set( + inspect.signature(compact_wrapper).parameters if callable(compact_wrapper) else () + ): + pytest.skip("installed cuDNN Frontend lacks compact MXFP8 indexer support") + + ratio, topk, idx_nh, idx_hd = 4, 16, 64, 128 + q_lens, k_lens = [64, 96], [16, 24] + max_seqlen_q, max_seqlen_k = 96, 24 + cu_q = _make_cu_seqlens(q_lens, device='cuda') + cu_k = _make_cu_seqlens(k_lens, device='cuda') + total_q, total_k = sum(q_lens), sum(k_lens) + torch.manual_seed(0) + q = torch.randn(total_q, idx_nh, idx_hd, dtype=torch.bfloat16, device='cuda') + k = torch.randn(total_k, idx_hd, dtype=torch.bfloat16, device='cuda') + w_raw = torch.randn(total_q, idx_nh, dtype=torch.bfloat16, device='cuda') + sm_scale = idx_hd**-0.5 + w = (w_raw.float() * sm_scale).to(torch.bfloat16) + + workspace = prepare_thd_compact_indexer_workspace( + q, + k, + topk=topk, + ratio=ratio, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_k=max_seqlen_k, + return_softmax=True, + precision=precision, + ) + assert workspace is not None + assert workspace.softmax_out is not None + assert (workspace.mxfp8 is not None) == (precision == "mxfp8") + if precision == "mxfp8": + assert workspace.mxfp8 is not None + assert torch.equal( + workspace.mxfp8.cu_seqlens_q_scale_padded, + torch.tensor([0, 64, 160], dtype=torch.int32, device="cuda"), + ) + assert torch.equal( + workspace.mxfp8.cu_seqlens_k_scale_padded, + torch.tensor([0, 128, 256], dtype=torch.int32, device="cuda"), + ) + assert workspace.mxfp8.q_scale.shape == (1, 162 * idx_nh, 4) + assert workspace.mxfp8.k_scale.shape == (1, 256, 4) + + def run(): + return dk._indexer_topk_core( + q, + k, + w, + topk=topk, + ratio=ratio, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_k, + max_seqlen_q=max_seqlen_q, + max_seqlen_kv=max_seqlen_k, + use_compact=True, + return_softmax=True, + compact_workspace=workspace, + precision=precision, + ) + + # Match cuDNN Frontend's documented prerequisite: three eager calls on + # a side stream perform value validation and JIT compilation. + side_stream = torch.cuda.Stream() + side_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side_stream): + for _ in range(3): + eager_indices, eager_lengths, _, eager_softmax = run() + torch.cuda.current_stream().wait_stream(side_stream) + eager_indices = eager_indices.clone() + eager_lengths = eager_lengths.clone() + eager_logits = workspace.out_logits.clone() + eager_softmax = eager_softmax.clone() + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_indices, captured_lengths, _, captured_softmax = run() + + assert captured_indices.data_ptr() == workspace.out_indices.data_ptr() + assert captured_softmax.data_ptr() == workspace.softmax_out.data_ptr() + + # Prove replay writes all caller-owned outputs rather than observing + # values left by warmup/capture. + captured_indices.fill_(-123) + captured_lengths.fill_(-123) + workspace.out_logits.fill_(float('nan')) + captured_softmax.fill_(float('nan')) + if precision == "mxfp8": + q_lens = [65, 95] + cu_q.copy_(_make_cu_seqlens(q_lens, device='cuda')) + graph.replay() + torch.cuda.synchronize() + + if precision == "mxfp8": + assert workspace.mxfp8 is not None + assert torch.equal( + workspace.mxfp8.cu_seqlens_q_scale_padded, + torch.tensor([0, 66, 162], dtype=torch.int32, device="cuda"), + ) + else: + assert torch.equal(captured_indices, eager_indices) + assert torch.equal(captured_lengths, eager_lengths) + torch.testing.assert_close(workspace.out_logits, eager_logits) + torch.testing.assert_close(captured_softmax, eager_softmax) + + # Validate the replayed local-id Top-K sets and aligned logits against + # the matching dense packed-segment reference. For MXFP8 this must use + # dequantized Q/K and the already-scaled BF16 weights; comparing against + # the original BF16 inputs would measure expected quantization error. + if precision == "mxfp8": + assert workspace.mxfp8 is not None + assert workspace.mxfp8.cu_seqlens_q_scale_padded is not None + assert workspace.mxfp8.cu_seqlens_k_scale_padded is not None + ref_q = _dequantize_thd_indexer_mxfp8( + workspace.mxfp8.q_buffers.data, + workspace.mxfp8.q_scale, + q_lens, + workspace.mxfp8.cu_seqlens_q_scale_padded, + ) + ref_k = _dequantize_thd_indexer_mxfp8( + workspace.mxfp8.k_buffers.data, + workspace.mxfp8.k_scale, + k_lens, + workspace.mxfp8.cu_seqlens_k_scale_padded, + ) + ref_w, ref_sm_scale = w.float(), 1.0 + else: + ref_q, ref_k = q.float(), k.float() + ref_w, ref_sm_scale = w_raw.float(), sm_scale + + # Radix emission order is arbitrary. + q_start = k_start = 0 + for q_len, k_len in zip(q_lens, k_lens): + ref_scores = _ref_indexer_full_score( + ref_q[q_start : q_start + q_len].unsqueeze(0), + ref_k[k_start : k_start + k_len].unsqueeze(0), + ref_w[q_start : q_start + q_len].unsqueeze(0), + sm_scale=ref_sm_scale, + ratio=ratio, + )[0] + for row in range(q_len): + actual_indices = captured_indices[q_start + row] + valid = actual_indices >= 0 + expected_length = min(topk, (row + 1) // ratio, k_len) + assert int(captured_lengths[q_start + row]) == expected_length + if expected_length == 0: + assert not bool(valid.any()) + continue + + expected_set = set(torch.topk(ref_scores[row], k=expected_length).indices.tolist()) + actual_set = set(actual_indices[valid].tolist()) + assert len(actual_set & expected_set) >= max(1, expected_length - 1) + torch.testing.assert_close( + workspace.out_logits[q_start + row, valid], + ref_scores[row, actual_indices[valid].long()], + atol=2e-2, + rtol=2e-2, + ) + + q_start += q_len + k_start += k_len + # --------------------------------------------------------------------------- # Real ``csa_sparse_attn``: forward + backward parity vs PyTorch reference. @@ -2335,7 +2761,7 @@ def test_real_fused_dense_loss_matches_reference(self, reset_lazy_kernel_state): w_bsh_scaled_bf = (w_bsh_bf.float() * s['indexer_softmax_scale']).to(w_bsh_bf.dtype) else: w_bsh_scaled_bf = w_bsh_bf - topk_indices_cmp, _, _ = _indexer_topk_core( + topk_indices_cmp, _, _, _ = _indexer_topk_core( q_idx_bshd_bf, k_idx_bsd_bf, w_bsh_scaled_bf, effective_topk, s['ratio'] ) compress_topk_idxs = torch.where(topk_indices_cmp >= 0, topk_indices_cmp + kv_offset, -1) @@ -2371,6 +2797,105 @@ def test_real_fused_dense_loss_matches_reference(self, reset_lazy_kernel_state): f"abs diff = {(indexer_loss - loss_ref).abs().item():.3e}" ) + @pytest.mark.parametrize("precision", ["bf16", "mxfp8"]) + def test_full_bshd_cuda_graph_capture(self, precision, reset_lazy_kernel_state): + """Capture and replay the complete zero-loss BSHD fused forward.""" + _skip_if_real_kernels_unavailable(need_flash_mla=True) + if torch.cuda.get_device_capability()[0] < 10: + pytest.skip("compact BSHD indexer forward + Top-K requires SM100+") + + from cudnn import DSA + + compact_wrapper = getattr(DSA, "indexer_forward_top_k_wrapper", None) + if not hasattr(DSA, "compress_topk_cand_buffer_size"): + pytest.skip("installed cuDNN Frontend lacks the compact BSHD workspace helper") + if ( + compact_wrapper is None + or "deterministic" not in inspect.signature(compact_wrapper).parameters + ): + pytest.skip("installed cuDNN Frontend lacks deterministic compact Top-K") + if precision == "mxfp8" and "q_scale" not in inspect.signature(compact_wrapper).parameters: + pytest.skip("installed cuDNN Frontend lacks compact MXFP8 indexer support") + + s = self.SHAPES + torch.manual_seed(31) + dev = 'cuda' + query = torch.randn(s['sq'], s['b'], s['np_'], s['d'], dtype=torch.bfloat16, device=dev) + kv_full = torch.randn(s['skv'], s['b'], s['d'], dtype=torch.bfloat16, device=dev) + attn_sink = torch.zeros(s['np_'], dtype=torch.float32, device=dev) + window_idxs = torch.randint( + 0, s['sq'], (s['b'], s['sq'], s['win_topk']), dtype=torch.int32, device=dev + ) + q_indexer = torch.randn( + s['sq'], s['b'], s['idx_nh'], s['idx_hd'], dtype=torch.bfloat16, device=dev + ) + k_indexer = torch.randn(s['n_comp'], s['b'], s['idx_hd'], dtype=torch.bfloat16, device=dev) + weights = torch.randn(s['sq'], s['b'], s['idx_nh'], dtype=torch.bfloat16, device=dev) + kv_offset = s['skv'] - s['n_comp'] + q_bshd = q_indexer.permute(1, 0, 2, 3).contiguous() + k_bsd = k_indexer.permute(1, 0, 2).contiguous() + workspace = prepare_bshd_compact_indexer_workspace( + q_bshd, + k_bsd, + topk=s['indexer_topk'], + ratio=s['ratio'], + return_softmax=False, + precision=precision, + ) + assert workspace is not None + assert workspace.softmax_out is None + assert (workspace.mxfp8 is not None) == (precision == "mxfp8") + + def run(): + return fused_csa_indexer_sparse_attn( + query, + kv_full, + attn_sink, + window_idxs, + q_indexer, + k_indexer, + weights, + indexer_topk=s['indexer_topk'], + ratio=s['ratio'], + softmax_scale=s['softmax_scale'], + indexer_softmax_scale=s['indexer_softmax_scale'], + loss_coeff=0.0, + sparse_loss=False, + kv_offset=kv_offset, + compact_workspace=workspace, + indexer_precision=precision, + deterministic=True, + ) + + side_stream = torch.cuda.Stream() + side_stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(side_stream): + for _ in range(3): + eager_output, eager_loss = run() + torch.cuda.current_stream().wait_stream(side_stream) + eager_output = eager_output.clone() + eager_loss = eager_loss.clone() + eager_indices = workspace.out_indices.clone() + eager_logits = workspace.out_logits.clone() + torch.cuda.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_output, captured_loss = run() + + workspace.out_indices.fill_(-123) + workspace.out_logits.fill_(float('nan')) + captured_output.fill_(float('nan')) + captured_loss.fill_(float('nan')) + graph.replay() + torch.cuda.synchronize() + + assert torch.equal(workspace.out_indices, eager_indices) + torch.testing.assert_close(workspace.out_logits, eager_logits) + torch.testing.assert_close(captured_output, eager_output) + assert torch.equal(captured_loss, eager_loss) + assert torch.equal(captured_loss, torch.zeros_like(captured_loss)) + # =========================================================================== # THD packed-sequence path @@ -2566,6 +3091,7 @@ def test_indexer_topk_thd_dispatch_calls_thd_kernel_path(self, reset_lazy_kernel q_causal_offsets = torch.tensor([5, 0], dtype=torch.int32, device=q.device) fake_dsa = MagicMock(name='_DSA_thd_stub') + fake_dsa.indexer_forward_top_k_wrapper = None def fake_indexer_forward(q_thd, k_thd, w_thd, ratio, **kwargs): # Verify THD kwargs were forwarded. @@ -2584,19 +3110,21 @@ def fake_indexer_forward(q_thd, k_thd, w_thd, ratio, **kwargs): } dk._DSA = fake_dsa - topk_idxs, topk_len = indexer_topk( - q, - k, - w, - topk=2, - ratio=4, - indexer_softmax_scale=128**-0.5, - cu_seqlens_q=cu_q, - cu_seqlens_kv=cu_kv, - max_seqlen_q=3, - max_seqlen_kv=2, - q_causal_offsets=q_causal_offsets, - ) + with pytest.warns(RuntimeWarning, match="Compact indexer.*falling back"): + topk_idxs, topk_len = indexer_topk( + q, + k, + w, + topk=2, + ratio=4, + indexer_softmax_scale=128**-0.5, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + max_seqlen_q=3, + max_seqlen_kv=2, + q_causal_offsets=q_causal_offsets, + deterministic=True, + ) # THD return shape: (total_q, topk) + (total_q,). assert topk_idxs.shape == (total_q, 2) assert topk_len.shape == (total_q,) @@ -2605,6 +3133,232 @@ def fake_indexer_forward(q_thd, k_thd, w_thd, ratio, **kwargs): seq_lens = fake_dsa.indexer_top_k_wrapper.call_args.args[1] assert torch.equal(seq_lens, torch.tensor([1, 1, 2, 0, 0], device=q.device)) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_indexer_topk_thd_compact_dispatch(self, reset_lazy_kernel_state): + q, k, w, cu_q, cu_kv = self._make_indexer_topk_thd_inputs() + topk = 2 + q_causal_offsets = torch.tensor([5, 0], dtype=torch.int32, device=q.device) + fake_dsa = MagicMock(name='_DSA_thd_compact_stub') + fake_dsa.indexer_forward_top_k_wrapper.return_value = { + 'indices': torch.zeros(q.shape[0], topk, dtype=torch.int32, device=q.device), + 'logits': torch.zeros(q.shape[0], topk, dtype=torch.float32, device=q.device), + 'softmax': torch.full((q.shape[0], topk), 0.25, device=q.device), + } + dk._DSA = fake_dsa + + with ( + patch.object(torch.cuda, 'get_device_capability', return_value=(10, 0)), + patch.object(torch.cuda, 'is_current_stream_capturing', return_value=False), + ): + indices, lengths, softmax = indexer_topk( + q, + k, + w, + topk=topk, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + max_seqlen_q=3, + max_seqlen_kv=2, + q_causal_offsets=q_causal_offsets, + deterministic=True, + return_softmax=True, + ) + + assert indices.shape == (q.shape[0], topk) + assert torch.all(lengths == topk) + assert torch.equal(softmax, torch.full((q.shape[0], topk), 0.25, device=q.device)) + fake_dsa.indexer_forward_wrapper.assert_not_called() + fake_dsa.indexer_top_k_wrapper.assert_not_called() + compact_call = fake_dsa.indexer_forward_top_k_wrapper.call_args + assert compact_call.args[1].shape == (k.shape[0], 1, k.shape[1]) + assert compact_call.kwargs['cu_seqlens_q'] is cu_q + assert compact_call.kwargs['cu_seqlens_k'] is cu_kv + assert compact_call.kwargs['max_seqlen_q'] == 3 + assert compact_call.kwargs['max_seqlen_k'] == 2 + assert compact_call.kwargs['q_causal_offsets'] is q_causal_offsets + assert compact_call.kwargs['precision'] == 'bf16' + assert compact_call.kwargs['topk_indices_global'] is False + assert compact_call.kwargs['deterministic'] is True + assert compact_call.kwargs['return_softmax'] is True + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_indexer_topk_mxfp8_compact_dispatch(self, reset_lazy_kernel_state): + sq, b, heads, head_dim, sk, topk = 4, 1, 64, 128, 2, 2 + q = torch.randn(sq, b, heads, head_dim, dtype=torch.bfloat16, device="cuda") + k = torch.randn(sk, b, head_dim, dtype=torch.bfloat16, device="cuda") + w = torch.randn(sq, b, heads, dtype=torch.bfloat16, device="cuda") + q_fp8 = torch.empty(b, sq, heads, head_dim, dtype=torch.float8_e4m3fn, device="cuda") + k_fp8 = torch.empty(b, sk, head_dim, dtype=torch.float8_e4m3fn, device="cuda") + q_scale = torch.empty(b, 256, 4, dtype=torch.float8_e8m0fnu, device="cuda") + k_scale = torch.empty(b, 128, 4, dtype=torch.float8_e8m0fnu, device="cuda") + + fake_dsa = MagicMock(name="_DSA_mxfp8_compact_stub") + fake_dsa.indexer_forward_top_k_wrapper.return_value = { + "indices": torch.zeros(b * sq, topk, dtype=torch.int32, device="cuda"), + "logits": torch.zeros(b * sq, topk, dtype=torch.float32, device="cuda"), + } + dk._DSA = fake_dsa + + with ( + patch.object(torch.cuda, "get_device_capability", return_value=(10, 0)), + patch.object( + dk, "quantize_indexer_mxfp8", side_effect=[(q_fp8, q_scale), (k_fp8, k_scale)] + ) as quantize, + ): + indices, lengths = indexer_topk( + q, k, w, topk=topk, ratio=4, precision="mxfp8", deterministic=True + ) + + assert indices.shape == (b, sq, topk) + assert torch.all(lengths == topk) + assert quantize.call_count == 2 + fake_dsa.indexer_forward_wrapper.assert_not_called() + fake_dsa.indexer_top_k_wrapper.assert_not_called() + compact_call = fake_dsa.indexer_forward_top_k_wrapper.call_args + assert compact_call.args[0] is q_fp8 + assert compact_call.args[1].data_ptr() == k_fp8.data_ptr() + assert compact_call.kwargs["precision"] == "mxfp8" + assert compact_call.kwargs["q_scale"] is q_scale + assert compact_call.kwargs["k_scale"] is k_scale + assert compact_call.kwargs["sf_vec_size"] == 32 + assert compact_call.kwargs["deterministic"] is True + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_indexer_topk_mxfp8_never_falls_back(self, reset_lazy_kernel_state): + sq, b, heads, head_dim, sk = 4, 1, 64, 128, 2 + q = torch.zeros(sq, b, heads, head_dim, dtype=torch.bfloat16, device="cuda") + k = torch.zeros(sk, b, head_dim, dtype=torch.bfloat16, device="cuda") + w = torch.zeros(sq, b, heads, dtype=torch.bfloat16, device="cuda") + fake_dsa = MagicMock(name="_DSA_without_compact_stub") + fake_dsa.indexer_forward_top_k_wrapper = None + dk._DSA = fake_dsa + + with patch.object(torch.cuda, "get_device_capability", return_value=(10, 0)): + with pytest.raises(RuntimeError, match="MXFP8 compact indexer requires"): + indexer_topk(q, k, w, topk=2, precision="mxfp8") + + fake_dsa.indexer_forward_wrapper.assert_not_called() + fake_dsa.indexer_top_k_wrapper.assert_not_called() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_indexer_topk_thd_capture_requires_preallocated_compact_workspace( + self, reset_lazy_kernel_state + ): + q, k, w, cu_q, cu_kv = self._make_indexer_topk_thd_inputs() + fake_dsa = MagicMock(name='_DSA_thd_capture_stub') + fake_dsa.indexer_forward_wrapper.return_value = { + 'scores': torch.zeros(q.shape[0], 2, dtype=torch.float32, device=q.device) + } + fake_dsa.indexer_top_k_wrapper.side_effect = lambda scores, seq_lens, **kw: { + 'indices': torch.zeros( + scores.shape[0], kw['top_k'], dtype=torch.int32, device=scores.device + ) + } + dk._DSA = fake_dsa + + with ( + patch.object(torch.cuda, 'get_device_capability', return_value=(10, 0)), + patch.object(torch.cuda, 'is_current_stream_capturing', return_value=True), + ): + with pytest.raises(ValueError, match='compact_workspace'): + indexer_topk( + q, + k, + w, + topk=2, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + max_seqlen_q=3, + max_seqlen_kv=2, + ) + + fake_dsa.indexer_forward_top_k_wrapper.assert_not_called() + fake_dsa.indexer_forward_wrapper.assert_not_called() + fake_dsa.indexer_top_k_wrapper.assert_not_called() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_indexer_topk_thd_capture_uses_preallocated_compact_workspace( + self, reset_lazy_kernel_state + ): + q, k, w, cu_q, cu_kv = self._make_indexer_topk_thd_inputs() + topk = 2 + fake_dsa = MagicMock(name='_DSA_thd_capture_workspace_stub') + fake_dsa.compress_topk_cand_buffer_size_thd.return_value = ( + torch.tensor([0, 4, 7], dtype=torch.int64, device=q.device), + 7, + ) + dk._DSA = fake_dsa + + with ( + patch.object(torch.cuda, 'get_device_capability', return_value=(10, 0)), + patch.object(torch.cuda, 'is_current_stream_capturing', return_value=False), + ): + workspace = prepare_thd_compact_indexer_workspace( + q, + k, + topk=topk, + ratio=4, + cu_seqlens_q=cu_q, + cu_seqlens_k=cu_kv, + max_seqlen_q=3, + max_seqlen_k=2, + ) + assert workspace is not None + assert workspace.softmax_out is None + changed_cu_q = _make_cu_seqlens([2, 3], device="cuda") + assert workspace.matches( + q=q, + k=k, + topk=topk, + ratio=4, + cu_seqlens_q=changed_cu_q, + cu_seqlens_k=cu_kv, + max_seqlen_q=3, + max_seqlen_k=2, + q_causal_offsets=None, + return_softmax=False, + ) + assert not workspace.matches( + q=q, + k=k, + topk=topk, + ratio=4, + cu_seqlens_q=torch.tensor([0, 1, 3, 5], dtype=torch.int32, device=q.device), + cu_seqlens_k=torch.tensor([0, 1, 2, 4], dtype=torch.int32, device=q.device), + max_seqlen_q=3, + max_seqlen_k=2, + q_causal_offsets=None, + return_softmax=False, + ) + fake_dsa.indexer_forward_top_k_wrapper.return_value = { + 'indices': workspace.out_indices.zero_(), + 'logits': workspace.out_logits.zero_(), + } + + with ( + patch.object(torch.cuda, 'get_device_capability', return_value=(10, 0)), + patch.object(torch.cuda, 'is_current_stream_capturing', return_value=True), + ): + indexer_topk( + q, + k, + w, + topk=topk, + cu_seqlens_q=cu_q, + cu_seqlens_kv=cu_kv, + max_seqlen_q=3, + max_seqlen_kv=2, + compact_workspace=workspace, + ) + + fake_dsa.indexer_forward_wrapper.assert_not_called() + compact_call = fake_dsa.indexer_forward_top_k_wrapper.call_args + assert compact_call.kwargs['cand_buffer'] is workspace.cand_buffer + assert compact_call.kwargs['cand_batch_offsets'] is workspace.cand_batch_offsets + assert compact_call.kwargs['out_indices'] is workspace.out_indices + assert compact_call.kwargs['out_logits'] is workspace.out_logits + assert 'softmax_out' not in compact_call.kwargs + # ===================================================================== # csa_sparse_attn(is_thd=True) # ===================================================================== @@ -2742,8 +3496,14 @@ class TestRealKernelFusedIndexerSparseAttnThd: indexer_softmax_scale=128**-0.5, ) - @pytest.mark.parametrize('sparse_loss', [False, True], ids=['dense_loss', 'sparse_loss']) - def test_thd_single_segment_matches_sbhd_b1(self, sparse_loss, reset_lazy_kernel_state): + @pytest.mark.parametrize( + 'sparse_loss,indexer_precision', + [(False, 'bf16'), (True, 'bf16'), (True, 'mxfp8')], + ids=['dense_loss', 'sparse_loss', 'sparse_loss_mxfp8'], + ) + def test_thd_single_segment_matches_sbhd_b1( + self, sparse_loss, indexer_precision, reset_lazy_kernel_state + ): """B=1 THD invocation should match the equivalent SBHD-b=1 call on the same input tensors (just reshaped), for both dense-loss and sparse-loss Path B. @@ -2786,6 +3546,7 @@ def test_thd_single_segment_matches_sbhd_b1(self, sparse_loss, reset_lazy_kernel loss_coeff=loss_coeff, sparse_loss=sparse_loss, kv_offset=kv_offset, + indexer_precision=indexer_precision, ) # ---- THD equivalent -------------------------------------------------- @@ -2831,6 +3592,7 @@ def test_thd_single_segment_matches_sbhd_b1(self, sparse_loss, reset_lazy_kernel max_seqlen_q=s['sq'], max_seqlen_compressed_idx=s['n_comp'], compressed_kv=compressed_kv_thd, + indexer_precision=indexer_precision, ) # SBHD and THD share the same underlying kernels; for B=1 the @@ -2838,7 +3600,7 @@ def test_thd_single_segment_matches_sbhd_b1(self, sparse_loss, reset_lazy_kernel # indexer's radix top-K, which can shift a few scores at the # boundary. Use the same tolerance as the SBHD-vs-PyTorch test. assert torch.allclose(loss_thd, loss_sbhd, atol=5e-2, rtol=1e-1), ( - f"sparse_loss={sparse_loss}: thd = {loss_thd.item():.6f}, " + f"sparse_loss={sparse_loss}, precision={indexer_precision}: thd = {loss_thd.item():.6f}, " f"sbhd = {loss_sbhd.item():.6f}, " f"abs diff = {(loss_thd - loss_sbhd).abs().item():.3e}" ) @@ -3334,7 +4096,7 @@ def _sbhd_to_bshd(q_sbhd, k_sbd, w_sbh, sm_scale): q_idx_bshd_bf, k_idx_bsd_bf, _, w_bsh_scaled_bf = _sbhd_to_bshd( q_idx_init, k_idx_init, w_init, s['indexer_softmax_scale'] ) - topk_indices_cmp, _, _ = _indexer_topk_core( + topk_indices_cmp, _, _, _ = _indexer_topk_core( q_idx_bshd_bf, k_idx_bsd_bf, w_bsh_scaled_bf, effective_topk, s['ratio'] ) compress_topk_idxs = torch.where( @@ -3381,7 +4143,7 @@ def _sbhd_to_bshd(q_sbhd, k_sbd, w_sbh, sm_scale): q_idx_bshd_k, k_idx_bsd_k, _, w_bsh_scaled_k = _sbhd_to_bshd( q_idx_init, k_idx_init, w_init, s['indexer_softmax_scale'] ) - _, _, kernel_indexer_scores = _indexer_topk_core( + _, _, kernel_indexer_scores, _ = _indexer_topk_core( q_idx_bshd_k, k_idx_bsd_k, w_bsh_scaled_k, effective_topk, s['ratio'] ) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py index 7f2378df896..8ca17a1016c 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention_cp.py @@ -1,6 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import gc +import inspect import os import statistics from contextlib import contextmanager, nullcontext @@ -16,6 +17,7 @@ from megatron.core.extensions.transformer_engine import HAVE_TE from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.quantization.indexer_quantization import HAVE_TE_MXFP8, HAVE_TRITON from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from tests.unit_tests.test_utilities import Utils from tests.unit_tests.transformer.experimental_attention_variant.test_dsv4_hybrid_attention import ( @@ -93,6 +95,38 @@ def _dsv4_cp_fused_kernels_available(): "DSv4 fused DSA cases require SM90+, flash_mla, and cudnn.DSA" ) +_DSV4_CP_MXFP8_KERNELS_UNAVAILABLE_REASON = ( + "DSv4 MXFP8 CP cases require SM100+, Transformer Engine MXFP8, Triton, " + "flash_mla, and a compatible cudnn.DSA compact wrapper" +) + + +def _dsv4_cp_mxfp8_kernels_available(): + """Return whether this host exposes the full MXFP8 compact CP stack.""" + if ( + not _dsv4_cp_fused_kernels_available() + or torch.cuda.get_device_capability()[0] < 10 + or not HAVE_TE_MXFP8 + or not HAVE_TRITON + ): + return False + from cudnn import DSA + + compact_wrapper = getattr(DSA, "indexer_forward_top_k_wrapper", None) + if not callable(compact_wrapper): + return False + required_parameters = { + "deterministic", + "precision", + "q_scale", + "k_scale", + "cu_seqlens_q_scale_padded", + "cu_seqlens_k_scale_padded", + "sf_vec_size", + "q_causal_offsets", + } + return required_parameters <= set(inspect.signature(compact_wrapper).parameters) + class _ReferenceCPGroup: def rank(self): @@ -270,6 +304,8 @@ def _make_dsv4_cp_config( dsa_indexer_use_sparse_loss=True, use_fused_kernels=True, apply_rope_fusion=True, + dsa_indexer_precision="bf16", + cuda_graph_impl="none", ): """Build the DSv4 flash attention config used by CP tests.""" shape = _DSV4_VARIANTS[_DSV4_CP_TEST_VARIANT] @@ -300,6 +336,21 @@ def _make_dsv4_cp_config( expert_model_parallel_size=1, dsa_kernel_backend="cudnn" if use_fused_kernels else "none", apply_rope_fusion=apply_rope_fusion, + dsa_indexer_precision=dsa_indexer_precision, + cuda_graph_impl=cuda_graph_impl, + thd_max_packed_sequences=( + len(_DSV4_CP_RAGGED_PADDED_SEG_LENS) + if cuda_graph_impl != "none" and context_parallel_size > 1 + else None + ), + max_seqlen_per_dp_cp_rank=( + sum(_DSV4_CP_RAGGED_PADDED_SEG_LENS) // context_parallel_size + if cuda_graph_impl != "none" and context_parallel_size > 1 + else None + ), + pad_packed_seq_alignment=( + "max" if cuda_graph_impl != "none" and context_parallel_size > 1 else None + ), ) @@ -537,6 +588,7 @@ def _measure_cp1_cuda_graph_time(self, layer_number, packed, padded_tokens): dsa_indexer_loss_coeff=1.0, dsa_indexer_use_sparse_loss=True, apply_rope_fusion=True, + cuda_graph_impl="local", ) attn = _build_attention(config, layer_number=layer_number, pg_collection=self.ref_pg).cuda() hidden_values, grad_values = _make_hidden_and_grad(padded_tokens, config.hidden_size) @@ -646,6 +698,59 @@ def test_thd_cp_matches_full_reference_forward_backward( del cp_attn, ref_attn, full_hidden, local_hidden, ref_hidden, local_out, ref_out, grad _clear_cuda_test_state() + def test_thd_cp_mxfp8_matches_full_reference_forward_backward(self): + """MXFP8 CP top-k, loss, and gradients match the CP1 THD path.""" + if not _dsv4_cp_mxfp8_kernels_available(): + pytest.skip(_DSV4_CP_MXFP8_KERNELS_UNAVAILABLE_REASON) + + packed, padded_tokens, local_idx = _make_ragged_cp_case(self.cp_size, self.cp_rank) + torch.manual_seed(_SEED + 1600) + model_parallel_cuda_manual_seed(_SEED + 1600) + config_cp = _make_dsv4_cp_config( + context_parallel_size=self.cp_size, dsa_indexer_precision="mxfp8" + ) + config_ref = _make_dsv4_cp_config(context_parallel_size=1, dsa_indexer_precision="mxfp8") + cp_attn = _build_attention(config_cp, layer_number=2, pg_collection=self.pg).cuda() + ref_attn = _build_attention(config_ref, layer_number=2, pg_collection=self.ref_pg).cuda() + _copy_module_parameters(cp_attn, ref_attn) + + full_hidden = torch.randn( + padded_tokens, 1, config_cp.hidden_size, dtype=torch.bfloat16, device="cuda" + ) + local_hidden = full_hidden.index_select(0, local_idx).detach().clone().requires_grad_(True) + ref_hidden = full_hidden.detach().clone().requires_grad_(True) + + local_out, _ = cp_attn( + hidden_states=local_hidden, attention_mask=None, packed_seq_params=packed + ) + ref_out, _ = ref_attn( + hidden_states=ref_hidden, attention_mask=None, packed_seq_params=packed + ) + _assert_cp_tensor_match( + local_out.detach(), ref_out.detach().index_select(0, local_idx), "layer=2:mxfp8:output" + ) + + grad = torch.randn_like(ref_out) + local_out.backward(grad.index_select(0, local_idx)) + ref_out.backward(grad) + _assert_cp_tensor_match( + local_hidden.grad.detach(), + ref_hidden.grad.index_select(0, local_idx), + "layer=2:mxfp8:hidden_grad", + ) + + ref_params = dict(ref_attn.named_parameters()) + for name, param in cp_attn.named_parameters(): + ref_grad = ref_params[name].grad + assert param.grad is not None, f"Missing CP grad for {name}" + assert ref_grad is not None, f"Missing reference grad for {name}" + grad_sum = param.grad.detach().clone() + dist.all_reduce(grad_sum, group=self.pg.cp) + _assert_cp_tensor_match(grad_sum, ref_grad, f"layer=2:mxfp8:param_grad:{name}") + + del cp_attn, ref_attn, full_hidden, local_hidden, ref_hidden, local_out, ref_out, grad + _clear_cuda_test_state() + def test_thd_cp_zero_indexer_loss_keeps_indexer_grads(self): """Zero indexer loss must still mark indexer grads ready for overlapped DDP.""" packed, padded_tokens, local_idx = _make_ragged_cp_case(self.cp_size, self.cp_rank) @@ -813,6 +918,7 @@ def test_thd_cp_cuda_graph_matches_eager_forward_backward( dsa_indexer_use_sparse_loss=True, use_fused_kernels=dsa_fused, apply_rope_fusion=rope_fused, + cuda_graph_impl="local", ) graph_attn = _build_attention( config, layer_number=layer_number, pg_collection=self.pg @@ -882,6 +988,75 @@ def test_thd_cp_cuda_graph_matches_eager_forward_backward( del eager_attn, test_grad, eager_hidden, graph_out, graph_hidden_grad _clear_cuda_test_state() + def test_thd_cp_mxfp8_cuda_graph_matches_eager_forward_backward(self): + """Workspace-backed MXFP8 CP capture matches eager forward/backward.""" + if not _dsv4_cp_mxfp8_kernels_available(): + pytest.skip(_DSV4_CP_MXFP8_KERNELS_UNAVAILABLE_REASON) + + packed, padded_tokens, local_idx = _make_ragged_cp_case(self.cp_size, self.cp_rank) + torch.manual_seed(_SEED + 1700) + model_parallel_cuda_manual_seed(_SEED + 1700) + config = _make_dsv4_cp_config( + context_parallel_size=self.cp_size, + dsa_indexer_precision="mxfp8", + cuda_graph_impl="local", + ) + graph_attn = _build_attention(config, layer_number=2, pg_collection=self.pg).cuda() + eager_attn = _build_attention(config, layer_number=2, pg_collection=self.pg).cuda() + graph_attn.train() + eager_attn.train() + _copy_module_parameters(graph_attn, eager_attn) + + full_hidden = torch.randn( + padded_tokens, 1, config.hidden_size, dtype=torch.bfloat16, device="cuda" + ) + test_hidden = full_hidden.index_select(0, local_idx).detach().clone() + test_grad = torch.randn_like(test_hidden) + static_hidden = test_hidden.detach().clone().requires_grad_(True) + eager_hidden = test_hidden.detach().clone().requires_grad_(True) + static_grad = test_grad.detach().clone() + + graph, graph_output = _capture_dsv4_attention_forward_backward( + graph_attn, static_hidden, static_grad, packed + ) + workspace = graph_attn.core_attention._active_thd_compact_indexer_workspace + assert workspace is not None + assert workspace.geometry.precision == "mxfp8" + assert workspace.mxfp8 is not None + assert workspace.softmax_out is not None + + with torch.no_grad(): + static_hidden.copy_(test_hidden) + static_grad.copy_(test_grad) + _zero_existing_grads(graph_attn, static_hidden) + graph.replay() + torch.cuda.synchronize() + graph_out = graph_output.detach().clone() + graph_hidden_grad = static_hidden.grad.detach().clone() + graph_param_grads = { + name: param.grad.detach().clone() + for name, param in graph_attn.named_parameters() + if param.grad is not None + } + + eager_out, eager_hidden_grad, eager_param_grads = _run_dsv4_attention_forward_backward( + eager_attn, eager_hidden, test_grad, packed + ) + torch.cuda.synchronize() + assert graph_param_grads.keys() == eager_param_grads.keys() + _assert_cp_graph_bitwise_match(graph_out, eager_out, "layer=2:mxfp8_graph:output") + _assert_cp_graph_fused_grad_match( + graph_hidden_grad, eager_hidden_grad, "layer=2:mxfp8_graph:hidden_grad" + ) + for name, graph_grad in graph_param_grads.items(): + _assert_cp_graph_fused_grad_match( + graph_grad, eager_param_grads[name], f"layer=2:mxfp8_graph:param_grad:{name}" + ) + + del graph, graph_output, graph_attn, eager_attn, full_hidden, test_hidden, test_grad + del static_hidden, eager_hidden, static_grad, graph_out, graph_hidden_grad + _clear_cuda_test_state() + @pytest.mark.parametrize( "layer_number", [2, 3], ids=["ratio_4_indexer", "ratio_128_compressor"] ) @@ -926,6 +1101,7 @@ def test_thd_cp_cuda_graph_replay_accepts_changed_padded_boundaries(self, layer_ dsa_indexer_use_sparse_loss=True, use_fused_kernels=True, apply_rope_fusion=True, + cuda_graph_impl="local", ) graph_attn = _build_attention( config, layer_number=layer_number, pg_collection=self.pg @@ -1082,6 +1258,7 @@ def test_thd_cp_cuda_graph_time_scales_vs_cp1(self, layer_number): dsa_indexer_loss_coeff=1.0, dsa_indexer_use_sparse_loss=True, apply_rope_fusion=True, + cuda_graph_impl="local", ) cp_attn = _build_attention( config_cp, layer_number=layer_number, pg_collection=self.pg diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py index 255115c307e..f146c20b082 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py @@ -1,6 +1,7 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import gc +import inspect import logging import math import os @@ -122,6 +123,7 @@ def _make_config( use_fused_kernels: bool = False, calculate_per_token_loss: bool = False, dsa_indexer_use_sparse_loss: bool = False, + dsa_indexer_precision: str = "bf16", legacy_kernel_fusion: bool | None = None, kernel_backend: str | None = None, use_legacy_attention_type: bool = False, @@ -153,6 +155,7 @@ def _make_config( dsa_indexer_topk=shape["dsa_indexer_topk"], dsa_indexer_loss_coeff=0.01, dsa_indexer_use_sparse_loss=dsa_indexer_use_sparse_loss, + dsa_indexer_precision=dsa_indexer_precision, calculate_per_token_loss=calculate_per_token_loss, add_bias_linear=False, bf16=True, @@ -1001,6 +1004,68 @@ def teardown_method(self): gc.collect() torch.cuda.empty_cache() + def test_mxfp8_indexer_attention_matches_native_reference(self): + """MXFP8 compact forward keeps the BF16 sparse-loss backward contract.""" + _skip_if_real_kernels_unavailable() + if torch.cuda.get_device_capability()[0] < 10: + pytest.skip("MXFP8 compact indexer requires SM100+") + + from cudnn import DSA + + compact_wrapper = getattr(DSA, "indexer_forward_top_k_wrapper", None) + required_parameters = {"q_scale", "cu_seqlens_q_scale_padded", "cu_seqlens_k_scale_padded"} + if not callable(compact_wrapper) or required_parameters - set( + inspect.signature(compact_wrapper).parameters + ): + pytest.skip("installed cuDNN Frontend lacks MXFP8 compact indexer support") + + config = _make_config( + "flash", + 4, + use_fused_kernels=True, + dsa_indexer_use_sparse_loss=True, + dsa_indexer_precision="mxfp8", + ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp"]) + spec = get_dsv4_hybrid_module_spec_for_backend(config=config, backend=TESpecProvider()) + real_layer = build_module( + spec, config=config, layer_number=1, cp_comm_type=None, pg_collection=pg_collection + ).cuda() + native_layer = NativeDSv4HybridAttention(config, 4).cuda() + real_params = _copy_real_params_to_native(real_layer, native_layer) + + seqlen = 512 + hidden_states = torch.randn( + seqlen, 1, config.hidden_size, dtype=torch.bfloat16, device="cuda", requires_grad=True + ) + hidden_states_native = hidden_states.detach().clone().requires_grad_(True) + grad = torch.randn_like(hidden_states) + + real_out, _ = real_layer(hidden_states=hidden_states, attention_mask=None) + native_out, native_indexer_loss = native_layer(hidden_states_native, pg_collection) + _assert_similarity(real_out.detach(), native_out.detach(), "mxfp8-indexer:out", eps=5e-3) + + real_out.backward(grad) + native_out.backward(grad) + assert native_indexer_loss is not None + native_indexer_loss.backward() + _assert_similarity( + hidden_states.grad, hidden_states_native.grad, "mxfp8-indexer:hidden_grad", eps=3e-2 + ) + + for name, native_param in native_layer.named_parameters(): + real_param = real_params[name] + assert native_param.grad is not None, f"Missing native grad for {name}" + assert real_param.grad is not None, f"Missing real grad for {name}" + _assert_similarity( + real_param.grad, native_param.grad, f"mxfp8-indexer:param_grad:{name}", eps=3e-2 + ) + + del real_layer, native_layer, real_params + del hidden_states, hidden_states_native, real_out, native_out, grad, native_indexer_loss + gc.collect() + torch.cuda.empty_cache() + @pytest.mark.parametrize(("backend", "use_fused_kernels"), _DSA_BACKENDS) @pytest.mark.parametrize("variant", ["flash", "pro"]) @pytest.mark.parametrize("compress_ratio", [1, 4, 128]) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_indexer_quantization.py b/tests/unit_tests/transformer/experimental_attention_variant/test_indexer_quantization.py new file mode 100644 index 00000000000..afbd75e3041 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_indexer_quantization.py @@ -0,0 +1,164 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Tests for DSA indexer quantization utilities.""" + +import pytest +import torch + +from megatron.core.quantization.indexer_quantization import ( + HAVE_TE_MXFP8, + HAVE_TRITON, + create_indexer_mxfp8_quantization_buffers, + indexer_mxfp8_scale_shape, + indexer_mxfp8_thd_scale_capacity, + indexer_mxfp8_thd_scale_shape, + make_indexer_mxfp8_scale_cu_seqlens, + quantize_indexer_mxfp8, + refresh_indexer_mxfp8_scale_cu_seqlens, +) + +pytestmark = [ + pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available"), + pytest.mark.skipif( + torch.cuda.is_available() and torch.cuda.get_device_capability()[0] < 10, + reason="MXFP8 indexer quantization requires SM100+", + ), + pytest.mark.skipif(not HAVE_TE_MXFP8, reason="Transformer Engine MXFP8 not available"), + pytest.mark.skipif(not HAVE_TRITON, reason="Triton not available"), +] + + +def _unpack_scales(scale: torch.Tensor, logical_rows: int) -> torch.Tensor: + """Undo the Blackwell 128x4 scale swizzle into logical ``(L, M, G)``.""" + l_size, _, padded_groups = scale.shape + scale_groups = 4 + rows = torch.arange(logical_rows, device=scale.device).view(-1, 1) + groups = torch.arange(scale_groups, device=scale.device).view(1, -1) + tile_idx = (rows // 128) * (padded_groups // 4) + groups // 4 + offsets = tile_idx * 512 + (rows % 32) * 16 + ((rows % 128) // 32) * 4 + groups % 4 + bytes_per_l = scale.shape[1] * padded_groups + l_offsets = torch.arange(l_size, device=scale.device).view(-1, 1, 1) + offsets = offsets.unsqueeze(0) + l_offsets * bytes_per_l + scale_bytes = scale.view(torch.uint8).flatten()[offsets] + return scale_bytes.contiguous().view(torch.float8_e8m0fnu).float() + + +def _reference_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """PyTorch reference for E4M3 values and round-up E8M0 scales.""" + grouped = x.float().reshape(*x.shape[:-1], x.shape[-1] // 32, 32) + amax = grouped.abs().amax(dim=-1) + raw_scale = amax / 448.0 + scale = torch.where( + raw_scale == 0, + torch.zeros_like(raw_scale), + torch.pow(2.0, torch.ceil(torch.log2(raw_scale))), + ) + quant_scale = torch.where(scale == 0, torch.zeros_like(scale), scale.reciprocal()) + data = (grouped * quant_scale.unsqueeze(-1)).reshape_as(x).to(torch.float8_e4m3fn) + return data, scale + + +def test_bshd_quantization_matches_reference(): + torch.manual_seed(123) + batch_size, seqlen, num_heads, head_dim = 2, 3, 64, 128 + x = torch.randn(batch_size, seqlen, num_heads, head_dim, dtype=torch.bfloat16, device="cuda") + + data, packed_scale = quantize_indexer_mxfp8(x) + ref_data, ref_scale = _reference_quantize(x) + unpacked_scale = _unpack_scales(packed_scale, seqlen * num_heads).reshape_as(ref_scale) + + assert data.dtype == torch.float8_e4m3fn + assert packed_scale.dtype == torch.float8_e8m0fnu + assert packed_scale.shape == indexer_mxfp8_scale_shape(batch_size, seqlen, num_heads, head_dim) + assert torch.equal(data.float(), ref_data.float()) + torch.testing.assert_close(unpacked_scale, ref_scale, rtol=0, atol=0) + + +def test_thd_quantization_uses_concatenated_padded_scale_spans(): + torch.manual_seed(456) + q_lens = [1, 2, 3, 4, 5] + num_heads, head_dim = 64, 128 + cu_seqlens = torch.tensor([0, 1, 3, 6, 10, 15], dtype=torch.int32, device="cuda") + cu_seqlens_scale_padded = make_indexer_mxfp8_scale_cu_seqlens(cu_seqlens, num_heads) + x = torch.randn(sum(q_lens), num_heads, head_dim, dtype=torch.bfloat16, device="cuda") + + data, packed_scale = quantize_indexer_mxfp8( + x, cu_seqlens=cu_seqlens, cu_seqlens_scale_padded=cu_seqlens_scale_padded + ) + ref_data, ref_scale = _reference_quantize(x) + unpacked_scale = _unpack_scales(packed_scale, packed_scale.shape[1])[0] + + assert torch.equal(data.float(), ref_data.float()) + assert torch.equal( + cu_seqlens_scale_padded, + torch.tensor([0, 2, 4, 8, 12, 18], dtype=torch.int32, device="cuda"), + ) + assert packed_scale.shape == indexer_mxfp8_thd_scale_shape(18, num_heads, head_dim) + start = 0 + for batch, length in enumerate(q_lens): + scale_start = int(cu_seqlens_scale_padded[batch].item()) * num_heads + actual = unpacked_scale[scale_start : scale_start + length * num_heads] + expected = ref_scale[start : start + length].reshape(length * num_heads, -1) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + padded_end = int(cu_seqlens_scale_padded[batch + 1].item()) * num_heads + padding = unpacked_scale[scale_start + length * num_heads : padded_end] + # E8M0 has no numerical zero: a zero storage byte decodes to its minimum value. + assert torch.all(padding == torch.finfo(torch.float8_e8m0fnu).tiny) + start += length + + +@pytest.mark.parametrize("num_heads", [1, 64]) +def test_thd_preallocated_quantization_cuda_graph_replay(num_heads): + torch.manual_seed(789) + capture_lens = [128, 128] + replay_lens = [1, 255] + head_dim = 128 + total_tokens = sum(capture_lens) + cu_seqlens = torch.tensor([0, 128, 256], dtype=torch.int32, device="cuda") + cu_seqlens_scale_padded = make_indexer_mxfp8_scale_cu_seqlens(cu_seqlens, num_heads) + shape = (total_tokens, num_heads, head_dim) if num_heads > 1 else (total_tokens, head_dim) + x = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + buffers = create_indexer_mxfp8_quantization_buffers(x) + scale_capacity = indexer_mxfp8_thd_scale_capacity(total_tokens, len(capture_lens), num_heads) + assert int(cu_seqlens_scale_padded[-1].item()) < scale_capacity + out_scale = torch.zeros( + indexer_mxfp8_thd_scale_shape(scale_capacity, num_heads, head_dim), + dtype=torch.float8_e8m0fnu, + device="cuda", + ) + + def run(): + refresh_indexer_mxfp8_scale_cu_seqlens(cu_seqlens_scale_padded, cu_seqlens, num_heads) + return quantize_indexer_mxfp8( + x, + cu_seqlens=cu_seqlens, + cu_seqlens_scale_padded=cu_seqlens_scale_padded, + buffers=buffers, + out_scale=out_scale, + ) + + for _ in range(3): + run() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_data, captured_scale = run() + + assert captured_data.data_ptr() == buffers.data.data_ptr() + assert captured_scale.data_ptr() == out_scale.data_ptr() + first_data = captured_data.float().clone() + + x.copy_(torch.randn_like(x)) + cu_seqlens.copy_( + torch.tensor([0, replay_lens[0], sum(replay_lens)], dtype=torch.int32, device="cuda") + ) + graph.replay() + torch.cuda.synchronize() + expected_scale_prefix = make_indexer_mxfp8_scale_cu_seqlens(cu_seqlens, num_heads) + expected_data, expected_scale = quantize_indexer_mxfp8( + x, cu_seqlens=cu_seqlens, cu_seqlens_scale_padded=expected_scale_prefix + ) + assert int(expected_scale_prefix[-1].item()) == scale_capacity + assert torch.equal(cu_seqlens_scale_padded, expected_scale_prefix) + assert not torch.equal(captured_data.float(), first_data) + assert torch.equal(captured_data.float(), expected_data.float()) + assert torch.equal(captured_scale.view(torch.uint8), expected_scale.view(torch.uint8))