diff --git a/megatron/core/pipeline_parallel/schedules.py b/megatron/core/pipeline_parallel/schedules.py index c6c6ed071ad..2b27f367b00 100644 --- a/megatron/core/pipeline_parallel/schedules.py +++ b/megatron/core/pipeline_parallel/schedules.py @@ -321,6 +321,23 @@ def forward_step_calc_loss( else: MTPLossAutoScaler.set_loss_scale(loss_scale / num_microbatches) + # Set the loss scale for the DSA indexer loss. + if hasattr(config, 'dsa_indexer_loss_coeff') and config.dsa_indexer_loss_coeff is not None: + from megatron.core.transformer.experimental_attention_variant.dsa import ( + DSAIndexerLossAutoScaler, + ) + + device = get_tensor_device(output_tensor) + loss_scale = ( + config.grad_scale_func(torch.ones(1, device=device)) + if config.grad_scale_func is not None + else torch.ones(1, device=device) + ) + if config.calculate_per_token_loss: + DSAIndexerLossAutoScaler.set_loss_scale(loss_scale) + else: + DSAIndexerLossAutoScaler.set_loss_scale(loss_scale / num_microbatches) + return output_tensor, num_tokens diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 547c1828a95..3346ea03017 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -20,6 +20,12 @@ fused_qk_topk_naive, rotate_activation, ) +from megatron.core.transformer.experimental_attention_variant.dsa_kernels import ( + build_flat_topk_idxs, + dsa_sparse_attn, + fused_indexer_sparse_attn, + indexer_topk, +) from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_config import TransformerConfig @@ -596,7 +602,7 @@ def __init__( softmax_scale = config.v_head_dim**-0.5 self.softmax_scale = softmax_scale - self.force_unfused_dsa = getattr(config, 'force_unfused_dsa', True) + self.apply_dsa_kernel_fusion = config.apply_dsa_kernel_fusion # Learnable attention sink per head self.attn_sink = nn.Parameter(torch.zeros(self.n_local_heads, dtype=torch.float32)) @@ -631,6 +637,265 @@ def __init__( else: self.indexer = None + # ------------------------------------------------------------------ + # Private helpers – each owns one logical slice of the forward pass. + # ------------------------------------------------------------------ + + def _build_kv_full( + self, kv: torch.Tensor, x: torch.Tensor + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], int]: + """Concatenate original KV with compressed KV (if applicable). + + Returns: + kv_full: [n_kv, b, v_head_dim] original + compressed KV. + compressed_kv: [n_compressed, b, v_head_dim] or None. + n_compressed: number of compressed positions (0 when unused). + """ + if self.compressor is not None and self.compress_ratio > 1: + compressed_kv = self.compressor(x) + if compressed_kv is not None: + kv_full = torch.cat([kv, compressed_kv], dim=0) + n_compressed = compressed_kv.size(0) + else: + kv_full = kv + compressed_kv = None + n_compressed = 0 + else: + kv_full = kv + compressed_kv = None + n_compressed = 0 + return kv_full, compressed_kv, n_compressed + + def _forward_unfused_csa( + self, + query: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + kv_full: torch.Tensor, + compressed_kv: Optional[torch.Tensor], + n_compressed: int, + offset: int, + window_idxs: torch.Tensor, + packed_seq_params: Optional[PackedSeqParams], + ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """PyTorch fallback path (no fused kernels). + + Returns ``(output, indexer_loss)``. + """ + sq, b, np, hn = query.size() + indexer_loss = None + + if self.compress_ratio > 1 and n_compressed > 0: + nvtx_range_push("compressed_indices") + if self.indexer is not None: + x_det = x.detach() + qr_det = qr.detach() + + causal_mask = ( + torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(sq, -1) + ) + positions = torch.arange(1, sq + 1, device=x.device).unsqueeze(1) + causal_mask = ( + torch.where(causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0) + .unsqueeze(0) + .expand(b, -1, -1) + ) # [b, sq, n_compressed] + + if self.training and torch.is_grad_enabled(): + q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk( + x_det, qr_det, packed_seq_params + ) + indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) + key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, np, -1) + # ``FusedDSAIndexerLoss`` does not accept a separate + # indexer_softmax_scale; apply it here via the + # weights-scaling trick so the effective weights match + # the pre-scale-split behaviour. + weights_for_unfused = weights_indexer.float() * self.indexer.softmax_scale + topk_indices_compressed, indexer_loss = FusedDSAIndexerLoss.apply( + q_indexer, + weights_for_unfused, + k_indexer, + query.detach(), + key_for_loss.detach(), + self.softmax_scale, + min(self.indexer.index_topk, n_compressed), + indexer_loss_coeff, + causal_mask, + getattr(self.config, "dsa_indexer_use_sparse_loss", True), + self.indexer.pg_collection, + self.config.calculate_per_token_loss, + ) + if indexer_loss_coeff > 0: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_layers + (self.config.mtp_num_layers or 0), + ) + else: + _, topk_indices_compressed = self.indexer( + x_det, qr_det, mask=causal_mask, packed_seq_params=packed_seq_params + ) + + n_valid_per_pos = positions // self.compress_ratio # [sq, 1] + valid = topk_indices_compressed < n_valid_per_pos + compress_topk_idxs = torch.where( + valid, topk_indices_compressed + offset, torch.tensor(-1, device=x.device) + ) + else: + compress_topk_idxs = get_compress_topk_idxs( + self.compress_ratio, b, sq, offset, query.device + ) + + topk_idxs = torch.cat([window_idxs, compress_topk_idxs], dim=-1) + nvtx_range_pop("compressed_indices") + else: + topk_idxs = window_idxs + + topk_idxs = topk_idxs.int() + + nvtx_range_push("sparse_attn_kernel") + output = unfused_compressed_sparse_attn( + query, kv_full, self.attn_sink.float(), topk_idxs, self.softmax_scale + ) + nvtx_range_pop("sparse_attn_kernel") + return output, indexer_loss + + def _forward_fused_no_indexer( + self, + query: torch.Tensor, + kv_full: torch.Tensor, + n_compressed: int, + offset: int, + window_idxs: torch.Tensor, + ) -> torch.Tensor: + """Path A: fused sparse attn with window or deterministic compressed indices.""" + sq, b, np, hn = query.size() + + nvtx_range_push("compressed_indices") + if self.compress_ratio > 1 and n_compressed > 0: + compress_topk_idxs = get_compress_topk_idxs( + self.compress_ratio, b, sq, offset, query.device + ) + flat_idxs, _ = build_flat_topk_idxs( + window_idxs, compress_topk_idxs, batch_size=b, seqlen_kv=kv_full.shape[0] + ) + else: + flat_idxs, _ = build_flat_topk_idxs( + window_idxs, batch_size=b, seqlen_kv=kv_full.shape[0] + ) + nvtx_range_pop("compressed_indices") + + nvtx_range_push("sparse_attn_kernel") + output = dsa_sparse_attn( + query, kv_full, self.attn_sink.float(), flat_idxs, self.softmax_scale + ) + nvtx_range_pop("sparse_attn_kernel") + return output + + def _forward_fused_indexer_inference( + self, + query: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + kv_full: torch.Tensor, + n_compressed: int, + offset: int, + window_idxs: torch.Tensor, + packed_seq_params: Optional[PackedSeqParams], + ) -> torch.Tensor: + """Path C: separate indexer forward (no loss) + fused sparse attn (compact).""" + b = query.size(1) + + nvtx_range_push("compressed_indices") + x_det = x.detach() + qr_det = qr.detach() + q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk( + x_det, qr_det, packed_seq_params + ) + topk_indices_cmp, _ = indexer_topk( + q_indexer, + k_indexer, + weights_indexer, + min(self.indexer.index_topk, n_compressed), + self.compress_ratio, + indexer_softmax_scale=self.indexer.softmax_scale, + ) + compress_topk_idxs = torch.where(topk_indices_cmp >= 0, topk_indices_cmp + offset, -1) + flat_idxs, flat_tlen = build_flat_topk_idxs( + window_idxs, compress_topk_idxs, batch_size=b, seqlen_kv=kv_full.shape[0], compact=True + ) + nvtx_range_pop("compressed_indices") + + nvtx_range_push("sparse_attn_kernel") + output = dsa_sparse_attn( + query, + kv_full, + self.attn_sink.float(), + flat_idxs, + self.softmax_scale, + topk_length=flat_tlen, + ) + nvtx_range_pop("sparse_attn_kernel") + return output + + def _forward_fused_indexer_training( + self, + query: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + kv_full: torch.Tensor, + n_compressed: int, + offset: int, + window_idxs: torch.Tensor, + packed_seq_params: Optional[PackedSeqParams], + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Path B: fused indexer (with loss) + fused sparse attn. + + Returns ``(output, indexer_loss)``. + """ + nvtx_range_push("compressed_indices") + x_det = x.detach() + qr_det = qr.detach() + q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk( + x_det, qr_det, packed_seq_params + ) + nvtx_range_pop("compressed_indices") + + indexer_loss_coeff = self.config.dsa_indexer_loss_coeff or 0.0 + + nvtx_range_push("sparse_attn_kernel") + output, indexer_loss = fused_indexer_sparse_attn( + query, + kv_full, + self.attn_sink.float(), + window_idxs, + q_indexer, + k_indexer, + weights_indexer, + min(self.indexer.index_topk, n_compressed), + self.compress_ratio, + self.softmax_scale, + self.indexer.softmax_scale, + indexer_loss_coeff, + sparse_loss=getattr(self.config, "dsa_indexer_use_sparse_loss", True), + kv_offset=offset, + calculate_per_token_loss=self.config.calculate_per_token_loss, + ) + nvtx_range_pop("sparse_attn_kernel") + + if indexer_loss_coeff > 0: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_layers + (self.config.mtp_num_layers or 0), + ) + return output, indexer_loss + + # ------------------------------------------------------------------ + # Public entry point + # ------------------------------------------------------------------ + def forward( self, query: torch.Tensor, @@ -663,115 +928,43 @@ def forward( sq, b, np, hn = query.size() - # --- Step 1: Prepare single-head KV (squeeze singleton head dim) --- kv = key.squeeze(-2) # [sq, b, 1, v_head_dim] -> [sq, b, v_head_dim] - - # --- Step 2: Compression --- - if self.compressor is not None and self.compress_ratio > 1: - compressed_kv = self.compressor(x) # [n_compressed, b, v_head_dim] - if compressed_kv is not None: - kv_full = torch.cat([kv, compressed_kv], dim=0) - n_compressed = compressed_kv.size(0) - else: - kv_full = kv - n_compressed = 0 - else: - kv_full = kv - n_compressed = 0 - + kv_full, compressed_kv, n_compressed = self._build_kv_full(kv, x) offset = sq # compressed indices start after original positions - - # --- Step 3: Window indices --- window_idxs = get_window_topk_idxs(self.window_size, b, sq, query.device) - # --- Step 4: Compressed indices --- - indexer_loss = None - - if self.force_unfused_dsa: - if self.compress_ratio > 1 and n_compressed > 0: - nvtx_range_push("compressed_indices") - if self.indexer is not None: - x_det = x.detach() - qr_det = qr.detach() - - causal_mask = ( - torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(sq, -1) - ) - positions = torch.arange(1, sq + 1, device=x.device).unsqueeze(1) - causal_mask = ( - torch.where( - causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0 - ) - .unsqueeze(0) - .expand(b, -1, -1) - ) # [b, sq, n_compressed] - - if self.training and torch.is_grad_enabled(): - q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk( - x_det, qr_det, packed_seq_params - ) - indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) - # compressed_kv is [n, b, hn]; expand to [n, b, np, hn] for loss - key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, np, -1) - # ``FusedDSAIndexerLoss`` does not accept a separate - # indexer_softmax_scale; apply it here via the - # weights-scaling trick so the effective weights match - # the pre-scale-split behaviour. - weights_for_unfused = weights_indexer * self.indexer.softmax_scale - topk_indices_compressed, indexer_loss = FusedDSAIndexerLoss.apply( - q_indexer, - weights_for_unfused, - k_indexer, - query.detach(), - key_for_loss.detach(), - self.softmax_scale, - min(self.indexer.index_topk, n_compressed), - indexer_loss_coeff, - causal_mask, - getattr(self.config, "dsa_indexer_use_sparse_loss", True), - self.indexer.pg_collection, - ) - if indexer_loss_coeff > 0: - DSAIndexerLossLoggingHelper.save_loss_to_tracker( - loss=indexer_loss, - layer_number=self.layer_number, - num_layers=self.config.num_layers - + (self.config.mtp_num_layers or 0), - ) - else: - _, topk_indices_compressed = self.indexer( - x_det, qr_det, mask=causal_mask, packed_seq_params=packed_seq_params - ) - - n_valid_per_pos = positions // self.compress_ratio # [sq, 1] - valid = topk_indices_compressed < n_valid_per_pos - compress_topk_idxs = torch.where( - valid, topk_indices_compressed + offset, torch.tensor(-1, device=x.device) - ) - else: - compress_topk_idxs = get_compress_topk_idxs( - self.compress_ratio, b, sq, offset, query.device - ) - - topk_idxs = torch.cat([window_idxs, compress_topk_idxs], dim=-1) - nvtx_range_pop("compressed_indices") - else: - topk_idxs = window_idxs + has_indexer_compressed = ( + self.compress_ratio > 1 and n_compressed > 0 and self.indexer is not None + ) - topk_idxs = topk_idxs.int() + indexer_loss = None - # --- Step 5: Sparse attention --- - nvtx_range_push("sparse_attn_kernel") - output = unfused_compressed_sparse_attn( - query, kv_full, self.attn_sink.float(), topk_idxs, self.softmax_scale + if not self.apply_dsa_kernel_fusion: + output, indexer_loss = self._forward_unfused_csa( + query, + x, + qr, + kv_full, + compressed_kv, + n_compressed, + offset, + window_idxs, + packed_seq_params, + ) + elif has_indexer_compressed and self.training and torch.is_grad_enabled(): + output, indexer_loss = self._forward_fused_indexer_training( + query, x, qr, kv_full, n_compressed, offset, window_idxs, packed_seq_params + ) + elif has_indexer_compressed: + output = self._forward_fused_indexer_inference( + query, x, qr, kv_full, n_compressed, offset, window_idxs, packed_seq_params ) - nvtx_range_pop("sparse_attn_kernel") - else: - raise ValueError("Fused path is not supported for CompressedSparseAttention") + output = self._forward_fused_no_indexer( + query, kv_full, n_compressed, offset, window_idxs + ) - # --- Step 6: Attach indexer loss --- - if indexer_loss is not None and self.training and torch.is_grad_enabled(): + if indexer_loss is not None: output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) nvtx_range_pop("compressed_sparse_attn") diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index 94f0fae781c..5ee9d07b886 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -196,6 +196,7 @@ def compute_dsa_indexer_loss( sparse_loss: bool, pg_collection: ProcessGroupCollection, causal_mask_override: Optional[torch.Tensor] = None, + calculate_per_token_loss: bool = False, ) -> torch.Tensor: """ Compute KL divergence loss between index_scores and true attention_scores. @@ -216,6 +217,10 @@ def compute_dsa_indexer_loss( sparse_loss: bool, whether to use sparse indexer loss. If True, only the topk indices will be used to compute the loss. pg_collection: Process group collection, must have TP process group. + causal_mask_override: Optional mask used by compressed KV paths. + calculate_per_token_loss: If True, return a raw local sum so the global + token divisor can be applied by finalize_model_grads. If False, keep + the historical local BSHD average over ``batch * seqlen`` rows. Returns: index_loss: KL divergence loss (scalar). @@ -308,7 +313,11 @@ def compute_dsa_indexer_loss( # [b, sq, sk] -> [b, sq] -> [1] # Each token has same weight in the loss. - kl_div = kl_per_element.sum(dim=-1).mean() + kl_per_row = kl_per_element.sum(dim=-1) + if calculate_per_token_loss: + kl_div = kl_per_row.sum() + else: + kl_div = kl_per_row.mean() # Scale by coefficient. indexer_loss = kl_div * loss_coeff @@ -388,7 +397,18 @@ def fused_qk_topk_naive( def fwd_fused_indexer_loss_naive( - q, weights, k, query, key, topk, softmax_scale, loss_coeff, mask, sparse_loss, pg_collection + q, + weights, + k, + query, + key, + topk, + softmax_scale, + loss_coeff, + mask, + sparse_loss, + pg_collection, + calculate_per_token_loss, ): """Naive implementation of forward pass for indexer loss.""" index_scores, topk_indices = fused_qk_topk_naive(q, k, weights, topk, mask) @@ -403,6 +423,7 @@ def fwd_fused_indexer_loss_naive( sparse_loss, pg_collection, causal_mask_override=mask, + calculate_per_token_loss=calculate_per_token_loss, ) return topk_indices, indexer_loss @@ -421,6 +442,7 @@ def bwd_fused_indexer_loss_naive( grad_loss, pg_collection, causal_mask_override=None, + calculate_per_token_loss=False, ): """Naive implementation of backward pass for indexer loss.""" index_scores = _compute_index_scores(q, weights, k) # [B, Sq, Sk] @@ -520,11 +542,15 @@ def bwd_fused_indexer_loss_naive( del attention_scores_sum # Backward through loss = kl_div * loss_coeff - # where kl_div = kl_per_element.sum(dim=-1).mean() + # where kl_div is either kl_per_element.sum(dim=-1).mean() or the raw + # local sum when calculate_per_token_loss=True. grad_kl_div = grad_loss * loss_coeff # scalar - # Backward through mean: distribute gradient equally - grad_kl_per_row = grad_kl_div / (b * sq) # scalar value for each row + if calculate_per_token_loss: + grad_kl_per_row = grad_kl_div + else: + # Backward through mean: distribute gradient equally + grad_kl_per_row = grad_kl_div / (b * sq) # scalar value for each row # Backward through sum(dim=-1): broadcast back to [b, sq, sk] # Each element in a row contributes to the sum, so gradient is same for all @@ -630,6 +656,7 @@ def forward( mask, sparse_loss, pg_collection, + calculate_per_token_loss, ): """ Fused forward: index_scores never materialized in full. @@ -646,6 +673,7 @@ def forward( mask, sparse_loss, pg_collection, + calculate_per_token_loss, ) # Save for backward (recomputation strategy) @@ -654,6 +682,7 @@ def forward( ctx.loss_coeff = loss_coeff ctx.sparse_loss = sparse_loss ctx.pg_collection = pg_collection + ctx.calculate_per_token_loss = calculate_per_token_loss return topk_indices, loss @@ -677,10 +706,11 @@ def backward(ctx, grad_topk_indices, grad_loss): grad_loss, ctx.pg_collection, causal_mask_override=mask, + calculate_per_token_loss=ctx.calculate_per_token_loss, ) # query and key are detached in forward, so return None for their gradients - return grad_q, grad_weights, grad_k, None, None, None, None, None, None, None, None + return grad_q, grad_weights, grad_k, None, None, None, None, None, None, None, None, None class DSAIndexerLossAutoScaler(torch.autograd.Function): @@ -1201,6 +1231,7 @@ def forward( float_mask, getattr(self.config, "dsa_indexer_use_sparse_loss", False), self.indexer.pg_collection, + self.config.calculate_per_token_loss, ) # Save indexer loss for logging if indexer_loss_coeff > 0: diff --git a/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py new file mode 100644 index 00000000000..50435d11592 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/dsa_kernels.py @@ -0,0 +1,1024 @@ +# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. + +""" +DSA kernel wrappers for Megatron's DSv4 sparse attention. + +Mirrors the three integration paths of the old standalone ``dsa_kernels`` +package, but built on top of + +* :mod:`cudnn.deepseek_sparse_attention` (a.k.a. ``DSA``) — CuTe-DSL backward + + indexer score kernels + TRT-LLM radix top-K, shipped as part of + cuDNN Frontend. +* :mod:`flash_mla` — production sparse-attention forward kernel, expected to + be available as a separate PyPI package. + +Public API (same shape as the old ``dsa_kernels`` package): + +* ``build_flat_topk_idxs`` / ``local_to_global_flat`` — index helpers. +* ``dsa_sparse_attn`` — Path A / Path C step 2, differentiable sparse attention. +* ``indexer_topk`` — Path C inference indexer scoring + top-K. +* ``fused_indexer_sparse_attn`` — Path B training, fused indexer loss + + sparse attention with shared backward. +""" + +from __future__ import annotations + +from typing import Optional, Tuple + +import torch +from torch import Tensor + +# --------------------------------------------------------------------------- +# Lazy kernel imports +# --------------------------------------------------------------------------- + + +_flash_mla_sparse_fwd = None +_DSA = None + + +def _ensure_flash_mla(): + """Lazily import the FlashMLA sparse-forward kernel. + + FlashMLA ships ``flash_mla_sparse_fwd`` with a multi-head-KV signature; + :func:`_dsa_fwd_flash_mla` below is a thin adapter that unbatches the + DSA-shape inputs and pads ``TopK`` to the alignment expected by + FlashMLA's SM90 / SM100 kernels. + """ + global _flash_mla_sparse_fwd + if _flash_mla_sparse_fwd is not None: + return + + try: + from flash_mla import flash_mla_sparse_fwd as _fwd + except ImportError as e: + raise ImportError( + "FlashMLA is required for DSA sparse attention forward. " + "Install from https://github.com/deepseek-ai/FlashMLA/tree/nv_dev " + "so that `from flash_mla import flash_mla_sparse_fwd` succeeds." + ) from e + _flash_mla_sparse_fwd = _fwd + + +def _get_topk_alignment() -> int: + """Minimum ``TopK`` alignment required by the current GPU architecture. + + * SM90 : dual-warpgroup loop steps by 2 blocks → ``2 * B_TOPK = 128`` + * SM100: single-pipeline loop steps by 1 block → ``B_TOPK`` (64 for + head64, 128 for head128). DSA uses ``D = 512`` which maps to the + head64 kernel path → 64. + """ + sm = torch.cuda.get_device_capability() + if sm[0] >= 10: + return 64 + return 128 + + +def _dsa_fwd_flash_mla( + q: Tensor, + kv: Tensor, + topk_idxs: Tensor, + softmax_scale: float, + d_v: int = 512, + attn_sink: Optional[Tensor] = None, + topk_length: Optional[Tensor] = None, + indexer_topk: int = 0, +) -> Tuple[Tensor, Tensor, Optional[Tensor]]: + """DSA-shaped adapter around :func:`flash_mla.flash_mla_sparse_fwd`. + + Accepts flat (unbatched) tensors with global indices; pads ``TopK`` to + the GPU-specific alignment; returns ``(out, lse, lse_indexer)``. + """ + assert not ( + indexer_topk > 0 and topk_length is not None + ), "indexer_topk > 0 requires non-compact mode (topk_length must be None)" + _ensure_flash_mla() + + _total_S_q, _H, _D = q.shape + TopK = topk_idxs.shape[-1] + topk_align = _get_topk_alignment() + TopK_padded = (TopK + topk_align - 1) // topk_align * topk_align + if TopK_padded != TopK: + pad_width = TopK_padded - TopK + topk_idxs = torch.nn.functional.pad(topk_idxs, (0, pad_width), value=-1) + + kv_3d = kv.unsqueeze(1) # (total_S_kv, 1, D) h_kv=1 + indices = topk_idxs.unsqueeze(1) # (total_S_q, 1, TopK_padded) h_kv=1 + + with torch.cuda.nvtx.range("flash_mla_sparse_fwd"): + res = _flash_mla_sparse_fwd( + q, + kv_3d, + indices, + softmax_scale, + d_v=d_v, + attn_sink=attn_sink, + topk_length=topk_length, + indexer_topk=indexer_topk, + ) + if indexer_topk > 0: + out, _max_logits, lse, lse_indexer = res + else: + out, _max_logits, lse = res + lse_indexer = None + + if indexer_topk > 0: + # When indexer_topk == total TopK, lse_indexer should equal lse but + # the kernel may not snapshot correctly; fall back to lse. + if indexer_topk >= TopK: + return out, lse, lse.clone() + return out, lse, lse_indexer + return out, lse, None + + +def _ensure_dsa_namespace(): + """Lazily import the cudnn-frontend DSA namespace.""" + global _DSA + if _DSA is not None: + return + try: + from cudnn import DSA as _ns + except ImportError as e: + raise ImportError( + "cudnn-frontend DSA namespace not available. Install with " + "`pip install nvidia-cudnn-frontend[cutedsl]`." + ) from e + _DSA = _ns + + +# --------------------------------------------------------------------------- +# Index helpers +# --------------------------------------------------------------------------- + + +def local_to_global_flat(local_idxs: Tensor, batch_size: int, seqlen_kv: int) -> Tensor: + """Convert local per-batch indices to global flat indices. + + Follows the convention used by FlashMLA / SparseAttentionBackward: + flat row order is SBHD ``row[s * B + b]``; global index is + ``local * B + b`` for valid entries and ``-1`` otherwise. + + Args: + local_idxs: ``(b, sq, topk)`` int, values in ``[0, seqlen_kv)`` or -1. + batch_size: ``B``. + seqlen_kv: KV sequence length per batch (used for shape assertions + only; callers compute the values). + + Returns: + ``(sq*b, topk)`` int32. + """ + b, sq, topk = local_idxs.shape + assert b == batch_size + + idxs_sb = local_idxs.permute(1, 0, 2).reshape(sq * b, topk) + valid = idxs_sb >= 0 + batch_ids = torch.arange(sq * b, device=local_idxs.device) % b + batch_ids_exp = batch_ids.unsqueeze(1).expand_as(idxs_sb) + idxs_sb = torch.where(valid, idxs_sb * b + batch_ids_exp, idxs_sb) + return idxs_sb.int() + + +def build_flat_topk_idxs( + *idx_groups: Tensor, batch_size: int, seqlen_kv: int, compact: bool = False +) -> Tuple[Tensor, Optional[Tensor]]: + """Combine local per-batch index groups and convert to flat global form. + + Each *idx_group* is ``(b, sq, topk_i)`` with local per-batch KV indices + (already in ``kv_full`` index space, i.e. with any compressed-position + offset applied). ``-1`` marks invalid positions. + + Args: + *idx_groups: one or more ``(b, sq, topk_i)`` int tensors. + batch_size: ``B``. + seqlen_kv: total KV sequence length per batch. + compact: if True, pack valid entries to the front of each row and + additionally return ``topk_length``; if False, leave as-is and + return ``None``. + + Returns: + ``(topk_idxs, topk_length)`` where + ``topk_idxs`` is ``(sq*b, total_topk)`` int32 (flat global) and + ``topk_length`` is ``(sq*b,)`` int32 when ``compact``, else ``None``. + """ + combined = torch.cat(idx_groups, dim=-1) # (b, sq, total_topk) + b, sq, total_topk = combined.shape + + # Globalize first, compact second. Both ops are element-wise + (-1)-preserving, + # so swapping the order is a no-op for correctness; the win is that the + # global indices come out already in (sq*b, total_topk) flat layout, which is + # exactly the row order the cuDNN compactify kernel returns its per-row + # ``length`` in — no extra permute on the length tensor. + global_idxs = local_to_global_flat(combined, b, seqlen_kv) + + topk_length_flat = None + if compact: + if global_idxs.is_cuda: + # Fast path: single warp-per-row CuTe DSL kernel from cuDNN's DSA + # namespace. Replaces a stable argsort + gather + sum + permute + # chain with one global-load + global-store per element. + _ensure_dsa_namespace() + res = _DSA.compactify_wrapper(global_idxs) + global_idxs, topk_length_flat = res["indices"], res["topk_length"] + else: + # CPU fallback so the unit tests that exercise this helper without + # CUDA still work. Production callers always go through the CUDA + # path above. + valid_mask = global_idxs >= 0 + sorted_indices = valid_mask.int().argsort(dim=-1, descending=True, stable=True) + global_idxs = global_idxs.gather(-1, sorted_indices) + topk_length_flat = valid_mask.sum(dim=-1).int() + + return global_idxs, topk_length_flat + + +# --------------------------------------------------------------------------- +# Path A + Path C step 2: differentiable sparse attention +# --------------------------------------------------------------------------- + + +class SparseAttnFunc(torch.autograd.Function): + """SM100 sparse attention fwd + bwd on flat tensors. + + Forward uses :mod:`flash_mla`; backward uses cuDNN Frontend's + :attr:`cudnn.DSA.sparse_attention_backward_wrapper`. + """ + + @staticmethod + def forward( + ctx, + q: Tensor, # (total_sq, H, D) bf16 + kv: Tensor, # (total_skv, D) bf16 + attn_sink: Tensor, # (H,) f32 + topk_idxs: Tensor, # (total_sq, TopK) int32 global + topk_length: Optional[Tensor], # (total_sq,) int32 or None + softmax_scale: float, + indexer_topk: int, + ) -> Tuple[Tensor, Tensor, Optional[Tensor]]: + """Run FlashMLA sparse-attention forward and save tensors for backward.""" + out, lse, lse_indexer = _dsa_fwd_flash_mla( + q, + kv, + topk_idxs, + softmax_scale, + attn_sink=attn_sink, + topk_length=topk_length, + indexer_topk=indexer_topk, + ) + + ctx.save_for_backward(q, kv, attn_sink, topk_idxs, out, lse) + ctx.softmax_scale = softmax_scale + ctx.topk_length = topk_length + return out, lse, lse_indexer + + @staticmethod + def backward(ctx, dO, d_lse, d_lse_indexer): + """Compute sparse-attention backward via cuDNN DSA wrapper.""" + _ensure_dsa_namespace() + + q, kv, attn_sink, topk_idxs, out, lse = ctx.saved_tensors + + result = _DSA.sparse_attention_backward_wrapper( + q, + kv, + out, + dO, + lse, + attn_sink, + topk_idxs, + softmax_scale=ctx.softmax_scale, + topk_length=ctx.topk_length, + ) + dq, dkv, d_sink = result["dq"], result["dkv"], result["d_sink"] + return dq, dkv, d_sink, None, None, None, None + + +def dsa_sparse_attn( + query: Tensor, + kv: Tensor, + attn_sink: Tensor, + topk_idxs: Tensor, + softmax_scale: float, + topk_length: Optional[Tensor] = None, + indexer_topk: int = 0, +) -> Tensor: + """Sparse attention (Path A / Path C step 2). + + Args: + query: ``(sq, b, np, d)`` bf16 SBHD. + kv: ``(skv, b, d)`` bf16 SBD (K=V). + attn_sink: ``(np,)`` f32. + topk_idxs: ``(sq*b, topk)`` int32 — **flat global** indices produced + by :func:`build_flat_topk_idxs`. + softmax_scale: scalar float. + topk_length: ``(sq*b,)`` int32 — optional compact fast-path. Must be + ``None`` when ``indexer_topk > 0`` (FlashMLA constraint). + indexer_topk: int; ``0`` for Paths A/C, positive for Path B to enable + FlashMLA's ``lse_indexer`` output. + + Returns: + ``(sq, b, np * d_v)`` bf16 output. + """ + sq, b, np_, d = query.shape + skv = kv.shape[0] + + q_flat = query.reshape(sq * b, np_, d) + kv_flat = kv.reshape(skv * b, d) + + out_flat, _lse, _lse_indexer = SparseAttnFunc.apply( + q_flat, kv_flat, attn_sink, topk_idxs, topk_length, softmax_scale, indexer_topk + ) + + d_v = out_flat.shape[-1] + return out_flat.reshape(sq, b, np_, d_v).reshape(sq, b, np_ * d_v) + + +# --------------------------------------------------------------------------- +# Path C inference: indexer scoring + top-K +# --------------------------------------------------------------------------- + + +def _indexer_topk_bshd( + q_bshd: Tensor, k_bsd: Tensor, w_bsh: Tensor, topk: int, ratio: int = 4 +) -> Tuple[Tensor, Tensor, Tensor]: + """BSHD-layout core for :func:`indexer_topk`. + + Internal entry point used by both the public SBHD wrapper and Path B's + ``FusedIndexerSparseAttnFunc.forward`` so the SBHD→BSHD permute can be + performed once at the call site and reused across both the indexer + forward and the score-backward kernels (predict / target). + + Args: + q_bshd: ``(b, sq, idx_nh, idx_hd)`` bf16, C-contiguous. + k_bsd: ``(b, sk, idx_hd)`` bf16, C-contiguous. + w_bsh: ``(b, sq, idx_nh)`` bf16, C-contiguous, **already + ``indexer_softmax_scale``-scaled** by the caller. + topk: number of top-K indices to return per query. + ratio: compression ratio for the kernel's causal mask. + + Returns: + ``(topk_indices, topk_length, scores)`` where: + + * ``topk_indices``: ``(b, sq, topk)`` int32, invalid slots ``-1``. + * ``topk_length``: ``(b, sq)`` int32, per-row valid count. + * ``scores``: ``(b, sq, sk)`` fp32, raw scores from + :attr:`cudnn.DSA.indexer_forward_wrapper` with ``-inf`` on + causally-masked positions. + """ + _ensure_dsa_namespace() + + b, sq, _idx_nh, _idx_hd = q_bshd.shape + sk = k_bsd.shape[1] + device = q_bshd.device + + k_bshd = k_bsd.unsqueeze(2) # (b, sk, 1, idx_hd) + + scores = _DSA.indexer_forward_wrapper(q_bshd, k_bshd, w_bsh, ratio=ratio)[ + "scores" + ] # (b, sq, sk) fp32, -inf on masked positions + + # Top-K selection via the TRT-LLM CuTe-DSL radix kernel. + n_rows = b * sq + scores_flat = scores.reshape(n_rows, sk).contiguous() + q_idx = torch.arange(sq, device=device) + valid_per_q = ((q_idx + 1) // ratio).clamp(max=sk).to(torch.int32) # (sq,) + seq_lens = valid_per_q.repeat(b) # (b*sq,), row-major over (b, sq) + + topk_k = min(topk, sk) + tk_result = _DSA.indexer_top_k_wrapper( + scores_flat, seq_lens, top_k=topk_k, next_n=1, return_val=False + ) + topk_indices = tk_result["indices"].view(b, sq, topk_k) + + if topk_k < topk: + pad = torch.full((b, sq, topk - topk_k), -1, dtype=torch.int32, device=device) + topk_indices = torch.cat([topk_indices, pad], dim=-1) + + topk_length = (topk_indices >= 0).sum(dim=-1).int() # (b, sq) + return topk_indices.int(), topk_length, scores + + +def _sbhd_to_bshd_indexer_inputs( + q_indexer: Tensor, k_indexer: Tensor, weights: Tensor, indexer_softmax_scale: float +) -> Tuple[Tensor, Tensor, Tensor, Tensor]: + """Permute the indexer inputs SBHD→BSHD once, returning both the raw + BSHD weights and (when needed) a separate scaled copy. + + The ``relu(c·x) = c·relu(x)`` trick lets us push the indexer softmax + scale onto ``W`` (``(B, S_q, H)``, small) instead of the score tensor + (``(B, S_q, S_k)``, big). The raw ``w_bsh`` is preserved for the + backward GEMM path, which takes ``sm_scale`` directly. When + ``indexer_softmax_scale == 1.0`` the two views alias each other. + + Returns ``(q_bshd, k_bsd, w_bsh, w_bsh_scaled)``. + """ + q_bshd = q_indexer.permute(1, 0, 2, 3).contiguous() + k_bsd = k_indexer.permute(1, 0, 2).contiguous() + w_bsh = weights.permute(1, 0, 2).contiguous() + + if indexer_softmax_scale != 1.0: + w_bsh_scaled = (w_bsh.float() * indexer_softmax_scale).to(w_bsh.dtype) + else: + w_bsh_scaled = w_bsh + + return q_bshd, k_bsd, w_bsh, w_bsh_scaled + + +def indexer_topk( + q_indexer: Tensor, + k_indexer: Tensor, + weights: Tensor, + topk: int, + ratio: int = 4, + indexer_softmax_scale: float = 1.0, +) -> Tuple[Tensor, 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. + + Args: + q_indexer: ``(sq, b, idx_nh, idx_hd)`` bf16 SBHD. + k_indexer: ``(sk, b, idx_hd)`` bf16 SBD. + weights: ``(sq, b, idx_nh)`` bf16 SBH — raw (unscaled) weights. + topk: number of top-K indices to select. + ratio: compression ratio for the causal mask. + indexer_softmax_scale: scale applied to the indexer ``Q @ K^T`` + scores (typically ``idx_hd ** -0.5``). Applied internally via + the weights-scaling trick (``relu(c·x) = c·relu(x)`` for + ``c > 0``) so the caller passes raw weights. Default ``1.0`` + means weights are treated as already-scaled. + + Returns: + topk_indices: ``(b, sq, topk)`` int32 — local per-batch indices into + ``k_indexer``; invalid positions are ``-1``. + topk_length: ``(b, sq)`` int32 — per-query valid count. + """ + q_bshd, k_bsd, _w_bsh_raw, w_bsh_scaled = _sbhd_to_bshd_indexer_inputs( + q_indexer, k_indexer, weights, indexer_softmax_scale + ) + topk_indices, topk_length, _ = _indexer_topk_bshd(q_bshd, k_bsd, w_bsh_scaled, topk, ratio) + return topk_indices, topk_length + + +# --------------------------------------------------------------------------- +# Path B: fused indexer + sparse attention (training) +# --------------------------------------------------------------------------- + + +_CLIP_PROB_MIN = torch.finfo(torch.float32).tiny # kept compatible w/ cudnn kernel + + +def _compute_indexer_predict( + q_indexer_bshd: Tensor, + k_indexer_bsd: Tensor, + weights_bsh: Tensor, + topk_indices: Tensor, + qhead_per_kv_head: int, +) -> Tensor: + """Compute ``predict`` distribution (softmax over top-K of indexer scores). + + Wraps :attr:`cudnn.DSA.sparse_indexer_score_recompute_wrapper`. + + Args: + q_indexer_bshd: ``(B, S_q, H_q, D)`` bf16. + k_indexer_bsd: ``(B, S_k, D)`` bf16. + weights_bsh: ``(B, S_q, H_q)`` bf16. + topk_indices: ``(B, S_q, topk)`` int32. + qhead_per_kv_head: ``H_q`` (MQA). + + Returns: + predict: ``(B, S_q, topk)`` fp32, softmax over the top-K axis. + """ + _ensure_dsa_namespace() + result = _DSA.sparse_indexer_score_recompute_wrapper( + q_indexer_bshd, + k_indexer_bsd, + weights_bsh, + topk_indices, + qhead_per_kv_head=qhead_per_kv_head, + ) + return result["predict"] + + +def _compute_attn_target( + q_attn_bshd: Tensor, + k_attn_bsd: Tensor, + lse: Tensor, + topk_indices: Tensor, + softmax_scale: float, + qhead_per_kv_head: int, +) -> Tensor: + """Compute ``target`` distribution (L1-normalised head-sum softmax). + + Wraps :attr:`cudnn.DSA.sparse_attn_score_recompute_wrapper`. + + Shapes match :func:`_compute_indexer_predict`; ``lse`` is + ``(B, S_q, H_q)`` FP32 (comes from the attention forward pass). + """ + _ensure_dsa_namespace() + result = _DSA.sparse_attn_score_recompute_wrapper( + q_attn_bshd, + k_attn_bsd, + lse, + topk_indices, + softmax_scale, + qhead_per_kv_head=qhead_per_kv_head, + ) + return result["target"] + + +def _kl_loss_from_target_predict( + target: Tensor, + predict: Tensor, + topk_indices: Tensor, + loss_coeff: float, + calculate_per_token_loss: bool = False, +) -> Tensor: + """KL(target || predict) reduced over ``(B, S_q)`` and scaled by loss_coeff. + + Rows with no valid top-K positions (early query rows with ratio causal + masking) contribute 0 to the loss — the sparse score kernels produce + garbage for those rows, mirroring ``compute_dsa_indexer_loss``'s + ``row_valid`` handling. The default mean is taken over all ``(B, S_q)`` + positions. Per-token-loss mode returns a raw local sum so finalize can + apply the global token divisor. + """ + eps = _CLIP_PROB_MIN + t = target.clamp(min=eps) + p = predict.clamp(min=eps) + kl_per_row = (t * (torch.log(t) - torch.log(p))).sum(dim=-1) # (B, S_q) + + row_valid = (topk_indices >= 0).any(dim=-1) # (B, S_q) + kl_per_row = torch.where(row_valid, kl_per_row, torch.zeros_like(kl_per_row)) + loss = kl_per_row.sum() if calculate_per_token_loss else kl_per_row.mean() + return loss_coeff * loss + + +# --------------------------------------------------------------------------- +# Dense path (``sparse_loss=False``) — full-KV indexer loss +# --------------------------------------------------------------------------- + + +def _compute_dense_indexer_score( + q_indexer_bshd: Tensor, + k_indexer_bshd: Tensor, + weights_bsh: Tensor, + qhead_per_kv_head: int, + indexer_softmax_scale: float, + ratio: int, +) -> Tuple[Tensor, Tensor]: + """Dense indexer score forward over the full ``S_k`` axis. + + Wraps :attr:`cudnn.DSA.dense_indexer_score_recompute_wrapper`. Returns + ``(out, denom)`` where + + * ``out`` : ``(B, S_q, S_k)`` fp32, the raw head-reduced score + ``S[b,q,k] = indexer_softmax_scale * sum_h ReLU(Q_h · K_k^T) · W_{b,q,h}`` + with the kernel's ``ratio``-causal mask applied to invalid columns. + * ``denom`` : ``(B, S_q)`` fp32, the LSE denom of ``out`` along + ``S_k`` — i.e. ``predict = exp(out - denom[..., None])`` is the + indexer softmax distribution over the full KV. + + Both outputs are forwarded into :func:`_kl_loss_from_dense_scores` + *and* saved for the dense-path backward, where the dense indexer-grad + kernel consumes them directly. + """ + _ensure_dsa_namespace() + result = _DSA.dense_indexer_score_recompute_wrapper( + q_indexer_bshd, + k_indexer_bshd, + weights_bsh, + qhead_per_kv_head=qhead_per_kv_head, + sm_scale=indexer_softmax_scale, + ratio=ratio, + ) + return result["out"], result["denom"] + + +def _compute_dense_attn_score( + q_attn_bshd: Tensor, + k_attn_bshd: Tensor, + lse: Tensor, + qhead_per_kv_head: int, + softmax_scale: float, + ratio: int, +) -> Tuple[Tensor, Tensor]: + """Dense attention score forward over the full ``S_k`` axis. + + Wraps :attr:`cudnn.DSA.dense_attn_score_recompute_wrapper`. Returns + ``(out, denom)`` where + + * ``out`` : ``(B, S_q, S_k)`` fp32, the head-summed unnormalized + attention probability ``S[b,q,k] = sum_h exp(Q_h · K_k^T · scale - LSE[b,q,h])`` + with ``ratio`` causal mask applied. + * ``denom`` : ``(B, S_q)`` fp32, the L1-norm denom ``sum_k S[b,q,:]``. + ``target = out / denom[..., None]`` is the L1-normalized + head-summed attention distribution. + """ + _ensure_dsa_namespace() + result = _DSA.dense_attn_score_recompute_wrapper( + q_attn_bshd, + k_attn_bshd, + lse, + softmax_scale, + qhead_per_kv_head=qhead_per_kv_head, + ratio=ratio, + ) + return result["out"], result["denom"] + + +def _kl_loss_from_dense_scores( + attn_score: Tensor, + attn_l1norm: Tensor, + index_score: Tensor, + index_lse: Tensor, + loss_coeff: float, + calculate_per_token_loss: bool = False, +) -> Tensor: + """KL(target || predict) over the **full** KV axis, averaged over ``(B, S_q)``. + + Derives ``target = attn_score / attn_l1norm`` (L1-normalised, matches + ``compute_dsa_indexer_loss``'s ``attention_scores / sum`` step) and + ``log_predict = index_score - index_lse`` (LSE-normalised log-softmax), + then computes ``KL = sum_k target * (log target - log predict)`` and + scales by ``loss_coeff``. + + Rows where the kernel's ``ratio`` causal mask leaves no valid KV + position have ``attn_l1norm <= 0`` (L1) or ``index_lse == -inf`` + (LSE); those rows contribute 0 to the loss — the same ``row_valid`` + semantics as the reference ``compute_dsa_indexer_loss``. + """ + eps = _CLIP_PROB_MIN + # row_valid: rows with at least one un-masked KV position. + row_valid = (attn_l1norm > eps) & torch.isfinite(index_lse) + + # Safe denoms: replace invalid rows with a finite value so target / + # log-predict don't produce NaN; the row mask zeroes their KL below. + safe_l1 = attn_l1norm.clamp(min=eps) + safe_lse = torch.where(row_valid, index_lse, torch.zeros_like(index_lse)) + + target = attn_score / safe_l1.unsqueeze(-1) + target_clamped = target.clamp(min=eps) + log_predict = index_score - safe_lse.unsqueeze(-1) + + kl_per_row = (target_clamped * (torch.log(target_clamped) - log_predict)).sum( + dim=-1 + ) # (B, S_q) + kl_per_row = torch.where(row_valid, kl_per_row, torch.zeros_like(kl_per_row)) + loss = kl_per_row.sum() if calculate_per_token_loss else kl_per_row.mean() + return loss_coeff * loss + + +class FusedIndexerSparseAttnFunc(torch.autograd.Function): + """Path B: fused indexer (+KL loss) + sparse attention in one autograd. + + Differentiable w.r.t. ``query``, ``kv_full``, ``attn_sink``, + ``q_indexer``, ``k_indexer``, ``weights``. + + Two indexer-loss variants, selected by the ``sparse_loss`` argument + (matches ``compute_dsa_indexer_loss`` in the reference ``dsa.py``): + + * **Sparse loss** (``sparse_loss=True``) — KL is computed only over + the top-K KV positions the indexer has selected. + * **Dense loss** (``sparse_loss=False``, the default) — KL is + computed over *all* causally valid KV positions. + + Both variants share the FlashMLA sparse-attention forward + the + cuDNN sparse-attn backward; only the indexer-loss path branches. + """ + + @staticmethod + def forward( + ctx, + # Sparse attn inputs (differentiable) + query: Tensor, # (sq, b, np, d) bf16 + kv_full: Tensor, # (skv, b, d) bf16 + attn_sink: Tensor, # (np,) f32 + # Window indices (not differentiable) + window_idxs: Tensor, # (b, sq, win_topk) int32 + # Indexer inputs (differentiable) + q_indexer: Tensor, # (sq, b, idx_nh, idx_hd) bf16 + k_indexer: Tensor, # (n_comp, b, idx_hd) bf16 + weights: Tensor, # (sq, b, idx_nh) bf16 — raw (unscaled) + # Scalars + indexer_topk: int, + ratio: int, + softmax_scale: float, + indexer_softmax_scale: float, + loss_coeff: float, + sparse_loss: bool, + kv_offset: int, + calculate_per_token_loss: bool, + ) -> Tuple[Tensor, Tensor]: + """Fused forward: indexer scoring, sparse attention, KL loss, and indexer backward.""" + _ensure_dsa_namespace() + + sq, b, np_, d = query.shape + skv = kv_full.shape[0] + n_comp = k_indexer.shape[0] + idx_nh, idx_hd = q_indexer.shape[2], q_indexer.shape[3] + + effective_topk = min(indexer_topk, n_comp) + + # ---- 1. Permute indexer inputs SBHD->BSHD ONCE. ------------------- + q_idx_bshd, k_idx_bsd, w_bsh, w_bsh_scaled = _sbhd_to_bshd_indexer_inputs( + q_indexer, k_indexer, weights, indexer_softmax_scale + ) + + # ---- 2. Indexer scoring + top-K (with scores retained). ------------- + topk_indices_cmp, _, indexer_scores = _indexer_topk_bshd( + q_idx_bshd, k_idx_bsd, w_bsh_scaled, effective_topk, ratio + ) # topk_indices_cmp: (b, sq, effective_topk) int32; indexer_scores: (b, sq, n_comp) fp32 + + # ---- 3. Combine indices (indexer first, then window). -------------- + compress_topk_idxs = torch.where(topk_indices_cmp >= 0, topk_indices_cmp + kv_offset, -1) + combined_local = torch.cat([compress_topk_idxs, window_idxs], dim=-1) + global_idxs = local_to_global_flat(combined_local, b, skv) + + # ---- 4. FlashMLA forward (non-compact, indexer_topk > 0). --------- + q_flat = query.reshape(sq * b, np_, d) + kv_flat = kv_full.reshape(skv * b, d) + out_flat, lse, lse_indexer = _dsa_fwd_flash_mla( + q_flat, + kv_flat, + global_idxs, + softmax_scale, + attn_sink=attn_sink, + topk_length=None, + indexer_topk=effective_topk, + ) + + # ---- 5. Derive predict from indexer_scores, compute target. -------- + # Attention-path tensors (detached — loss is not differentiable through them). + q_attn_bshd = query.detach().permute(1, 0, 2, 3).contiguous() + k_attn_compressed_bsd = kv_full[kv_offset:].detach().permute(1, 0, 2).contiguous() + lse_indexer_bsqh = lse_indexer.reshape(sq, b, np_).permute(1, 0, 2) + + 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=2, 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) # (b, sq, topk) fp32 + + target = _compute_attn_target( + q_attn_bshd, + k_attn_compressed_bsd, + lse_indexer_bsqh, + topk_indices_cmp, + softmax_scale, + qhead_per_kv_head=np_, + ) + + 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: + # Dense: use full indexer_scores directly + logsumexp. + index_score = indexer_scores # (b, sq, n_comp) fp32 + index_lse = torch.logsumexp(indexer_scores, dim=-1) # (b, sq) fp32 + + attn_score, attn_l1norm = _compute_dense_attn_score( + q_attn_bshd, + k_attn_compressed_bsd.unsqueeze(2), + lse_indexer_bsqh, + qhead_per_kv_head=np_, + softmax_scale=softmax_scale, + ratio=ratio, + ) + + if loss_coeff > 0: + indexer_loss = _kl_loss_from_dense_scores( + attn_score, + attn_l1norm, + index_score, + index_lse, + 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 + # DSAIndexerLossAutoScaler provides the correct scale). + indexer_loss_coeff = loss_coeff + if calculate_per_token_loss: + indexer_loss_coeff = loss_coeff * (b * sq) + + unit_grad_loss = torch.ones((), device=query.device, dtype=torch.float32) + + if loss_coeff > 0: + if sparse_loss: + attn_score_for_bwd = target.clone() + index_score_for_bwd = predict.clone() + ig = _DSA.indexer_backward_wrapper( + q_idx_bshd, + w_bsh, + k_idx_bsd, + attn_score_for_bwd, + index_score_for_bwd, + topk_indices_cmp, + sm_scale=indexer_softmax_scale, + loss_coeff=indexer_loss_coeff, + grad_loss=unit_grad_loss, + block_I=128, + ) + else: + attn_score_for_bwd = attn_score.clone() + index_score_for_bwd = index_score.clone() + ig = _DSA.dense_indexer_backward_wrapper( + q_idx_bshd, + w_bsh, + k_idx_bsd, + attn_score_for_bwd, + attn_l1norm, + index_score_for_bwd, + index_lse, + sm_scale=indexer_softmax_scale, + loss_coeff=indexer_loss_coeff, + grad_loss=unit_grad_loss, + ratio=ratio, + block_I=128, + ) + # BSHD -> SBHD (match input layout). + precomputed_grad_q_indexer = ig["d_index_q"].permute(1, 0, 2, 3).contiguous() + precomputed_grad_k_indexer = ig["d_index_k"].permute(1, 0, 2).contiguous() + precomputed_grad_weights = ig["d_weights"].permute(1, 0, 2).contiguous() + else: + precomputed_grad_q_indexer = torch.zeros_like(q_indexer) + precomputed_grad_k_indexer = torch.zeros_like(k_indexer) + precomputed_grad_weights = torch.zeros_like(weights) + + # ---- 7. Save context (only sparse-attn bwd tensors + indexer grads). + ctx.save_for_backward( + q_flat, + kv_flat, + attn_sink, + global_idxs, + out_flat, + lse, + precomputed_grad_q_indexer, + precomputed_grad_k_indexer, + precomputed_grad_weights, + ) + ctx.softmax_scale = softmax_scale + ctx.sq = sq + ctx.b = b + ctx.np_ = np_ + ctx.d = d + ctx.skv = skv + + # ---- 8. Return. --------------------------------------------------- + d_v = out_flat.shape[-1] + output = out_flat.reshape(sq, b, np_, d_v).reshape(sq, b, np_ * d_v) + return output, indexer_loss + + @staticmethod + def backward(ctx, grad_output, grad_loss): + """Backward: sparse attention bwd + scale pre-computed indexer grads.""" + ( + q_flat, + kv_flat, + attn_sink, + global_idxs, + out_flat, + lse, + precomputed_grad_q_indexer, + precomputed_grad_k_indexer, + precomputed_grad_weights, + ) = ctx.saved_tensors + + sq, b, np_, d = ctx.sq, ctx.b, ctx.np_, ctx.d + skv = ctx.skv + + # ---- 1. Sparse attn backward. ------------------------------------- + d_v = out_flat.shape[-1] + dO_flat = grad_output.reshape(sq * b, np_, d_v) + + attn_bwd = _DSA.sparse_attention_backward_wrapper( + q_flat, + kv_flat, + out_flat, + dO_flat, + lse, + attn_sink, + global_idxs, + softmax_scale=ctx.softmax_scale, + topk_length=None, + ) + grad_query = attn_bwd["dq"].reshape(sq, b, np_, d) + grad_kv_full = attn_bwd["dkv"].reshape(skv, b, d) + d_sink = attn_bwd["d_sink"] + + # ---- 2. Scale pre-computed indexer grads by grad_loss. ------------- + grad_q_indexer = precomputed_grad_q_indexer * grad_loss + grad_k_indexer = precomputed_grad_k_indexer * grad_loss + grad_weights = precomputed_grad_weights * grad_loss + + # Grads: query, kv_full, attn_sink, window_idxs, q_indexer, k_indexer, + # weights, indexer_topk, ratio, softmax_scale, indexer_softmax_scale, + # loss_coeff, sparse_loss, kv_offset, calculate_per_token_loss + return ( + grad_query, + grad_kv_full, + d_sink, + None, + grad_q_indexer, + grad_k_indexer, + grad_weights, + None, + None, + None, + None, + None, + None, + None, + None, + ) + + +def fused_indexer_sparse_attn( + query: Tensor, + kv_full: Tensor, + attn_sink: Tensor, + window_idxs: Tensor, + q_indexer: Tensor, + k_indexer: Tensor, + weights: Tensor, + indexer_topk: int, + ratio: int, + softmax_scale: float, + indexer_softmax_scale: float = 1.0, + loss_coeff: float = 0.0, + sparse_loss: bool = False, + kv_offset: int = 0, + calculate_per_token_loss: bool = False, +) -> Tuple[Tensor, Tensor]: + """Path B (training): fused indexer (+KL loss) + sparse attention. + + See :class:`FusedIndexerSparseAttnFunc` for the detailed data flow. + + Args: + query: ``(sq, b, np, d)`` bf16 SBHD — attention query. + kv_full: ``(skv, b, d)`` bf16 SBD — original + compressed KV. + attn_sink: ``(np,)`` f32 — learnable sink per head. + window_idxs: ``(b, sq, win_topk)`` int32 — local window indices. + q_indexer: ``(sq, b, idx_nh, idx_hd)`` bf16 — indexer query. + k_indexer: ``(n_comp, b, idx_hd)`` bf16 — indexer key (compressed). + weights: ``(sq, b, idx_nh)`` bf16 — raw indexer weights. + indexer_topk: number of top-K compressed positions to select. + ratio: compression ratio used for the causal mask. + softmax_scale: attention ``Q @ K^T`` scale, typically + ``1/sqrt(v_head_dim)``. + indexer_softmax_scale: indexer ``Q @ K^T`` scale, typically + ``1/sqrt(idx_hd)``. Applied internally — caller passes raw + (unscaled) ``weights``. + loss_coeff: coefficient scaling the KL divergence loss. + sparse_loss: if ``True``, KL is computed only over the top-K + positions (cheap, less informative); if ``False`` (the + default, matches ``transformer_config.dsa_indexer_use_sparse_loss``), + KL is computed over the full causally-valid KV (more + informative, matches the DeepSeek-V3.2 paper, larger + intermediate-tensor footprint). See + :class:`FusedIndexerSparseAttnFunc` for the full data flow + of each variant. + kv_offset: start of compressed region within ``kv_full``. + calculate_per_token_loss: if True, report raw local KL sum and + compensate the cuDNN backward wrappers' local averaging. + + Returns: + ``(output, indexer_loss)`` where ``output`` is ``(sq, b, np * d_v)`` + bf16 and ``indexer_loss`` is a scalar f32. + """ + return FusedIndexerSparseAttnFunc.apply( + query, + kv_full, + attn_sink, + window_idxs, + q_indexer, + k_indexer, + weights, + indexer_topk, + ratio, + softmax_scale, + indexer_softmax_scale, + loss_coeff, + sparse_loss, + kv_offset, + calculate_per_token_loss, + ) + + +__all__ = [ + "build_flat_topk_idxs", + "local_to_global_flat", + "dsa_sparse_attn", + "indexer_topk", + "fused_indexer_sparse_attn", +] diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index c54c7f58f48..794f5c6726b 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -337,6 +337,11 @@ class TransformerConfig(ModelParallelConfig): """Whether to use dense mode for compressed sparse attention. If True, the CSA indexer will be disabled.""" + apply_dsa_kernel_fusion: bool = False + """If True, use fused DSA sparse-attention kernels (FlashMLA forward + cuDNN DSA backward, + indexer scoring, and top-K selection). Requires ``flash_mla`` and ``nvidia-cudnn-frontend`` + with CuTe-DSL support. When False, falls back to unfused PyTorch implementations.""" + #################### # linear attention #################### @@ -1442,6 +1447,44 @@ def __post_init__(self): assert not self.qk_clip, "QK clipping is not supported with DSv4 Hybrid Attention." self.hetereogenous_dist_checkpoint = True + if self.apply_dsa_kernel_fusion: + assert ( + torch.cuda.is_available() + ), "apply_dsa_kernel_fusion requires a CUDA device, but none is available." + sm = torch.cuda.get_device_capability() + assert sm[0] >= 10, ( + f"apply_dsa_kernel_fusion requires SM100+ (Blackwell or later), " + f"but current device has compute capability {sm[0]}.{sm[1]}." + ) + + _flash_mla_available = True + try: + from flash_mla import flash_mla_sparse_fwd # noqa: F401 + except ImportError: + _flash_mla_available = False + + _cudnn_dsa_available = True + try: + from cudnn import DSA # noqa: F401 + except ImportError: + _cudnn_dsa_available = False + + if not _flash_mla_available or not _cudnn_dsa_available: + missing = [] + if not _flash_mla_available: + missing.append( + "flash_mla (install from " + "https://github.com/deepseek-ai/FlashMLA/tree/nv_dev)" + ) + if not _cudnn_dsa_available: + missing.append("cudnn-frontend DSA (nvidia-cudnn-frontend[cutedsl])") + raise ValueError( + f"apply_dsa_kernel_fusion requires fused DSA kernels, but the " + f"following packages are not available: {', '.join(missing)}. " + f"Install them or pass --no-dsa-kernel-fusion to use the unfused " + f"PyTorch fallback." + ) + if self.fp8: # cannot support first last layer bf16 with delayed scaling if self.first_last_layers_bf16 and self.fp8_recipe == Fp8Recipe.delayed: diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 853973b92cd..2dbee13856a 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -2566,6 +2566,7 @@ def _add_network_size_args(parser): "persist_layer_norm", "bias_dropout_fusion", "apply_rope_fusion", + "apply_dsa_kernel_fusion", ] transformer_factory = ArgumentGroupFactory(TransformerConfig, exclude=exclude) transformer_group = transformer_factory.build_group(parser, "transformer configuration") @@ -4601,6 +4602,13 @@ def _add_experimental_attention_variant_args(parser): 'transformer layer (valid values: 0, 4, 128). ' 'The list length must equal num_layers.', ) + group.add_argument( + '--no-dsa-kernel-fusion', + action='store_false', + help='Disable fused DSA sparse-attention kernels (FlashMLA + cuDNN DSA) ' + 'and fall back to unfused PyTorch implementations.', + dest='apply_dsa_kernel_fusion', + ) return parser diff --git a/pyproject.toml b/pyproject.toml index 9ccafd4094e..2fb61c6f36c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -98,6 +98,7 @@ dev = [ "megatron-energon[av_decode]~=6.0", "av", "flashinfer-python>=0.5.0,<0.7.0", + "nvidia-cudnn-frontend", "wget", "onnxscript", "fastapi~=0.50", # Forcing a little bit more recent version of fastapi to be compatible with pydantic 2.0 diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp2_dsv4_hybrid_fused/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp2_dsv4_hybrid_fused/model_config.yaml new file mode 100644 index 00000000000..13047bc9de1 --- /dev/null +++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp2_dsv4_hybrid_fused/model_config.yaml @@ -0,0 +1,72 @@ +ENV_VARS: + CUDA_DEVICE_MAX_CONNECTIONS: 1 + NVTE_ALLOW_NONDETERMINISTIC_ALGO: 0 + NCCL_ALGO: Ring + CUBLAS_WORKSPACE_CONFIG: :4096:8 + ENABLE_LIGHTWEIGHT_MODE: true +MODEL_ARGS: + --num-layers: 6 + --hidden-size: 512 + --num-attention-heads: 8 + --multi-latent-attention: true + --q-lora-rank: 192 + --kv-lora-rank: 64 + --qk-head-dim: 16 + --qk-pos-emb-head-dim: 8 + --v-head-dim: 16 + --experimental-attention-variant: dsv4_hybrid + --dsa-indexer-n-heads: 64 + --dsa-indexer-head-dim: 128 + --dsa-indexer-topk: 512 + --dsa-indexer-loss-coeff: 0.01 + --dsa-indexer-use-sparse-loss: true + --csa-window-size: 128 + --csa-compress-ratios: ([0,4,128,4,128,4]) + --csa-compress-rotary-base: 40000 + --attention-backend: fused + --pipeline-model-parallel-layout: "Et|tt|tt|tL" + --log-params-norm: true + --log-num-zeros-in-grad: true + --log-validation-ppl-to-tensorboard: true + --log-timers-to-tensorboard: true + --tensorboard-dir: ${TENSORBOARD_PATH} + --micro-batch-size: 4 + --global-batch-size: 32 + --seq-length: 1024 + --position-embedding-type: rope + --max-position-embeddings: 1024 + --train-iters: 50 + --timing-log-level: 0 + --lr-decay-iters: 320000 + --save: ${CHECKPOINT_SAVE_PATH} + --load: ${CHECKPOINT_LOAD_PATH} + --data-path: ${DATA_PATH}/text/the_pile/shard00/my-gpt3_00_text_document + --vocab-file: ${DATA_PATH}/text/the_pile/shard00/bpe/vocab.json + --merge-file: ${DATA_PATH}/text/the_pile/shard00/bpe/merges.txt + --split: 949,50,1 + --distributed-backend: nccl + --lr: 0.00015 + --lr-decay-style: cosine + --min-lr: 1.0e-5 + --weight-decay: 1e-2 + --clip-grad: 1.0 + --lr-warmup-fraction: .01 + --log-interval: 1 + --save-interval: 25 + --eval-interval: 1000 + --eval-iters: 10 + --transformer-impl: transformer_engine + --tensor-model-parallel-size: 1 + --pipeline-model-parallel-size: 2 + --sequence-parallel: true + --untie-embeddings-and-output-weights: true + --deterministic-mode: true + --no-gradient-accumulation-fusion: true + --attention-softmax-in-fp32: true + --use-mcore-models: true + --ckpt-format: torch_dist + --data-cache-path: ${DATA_CACHE_PATH} + --bf16: true + --attention-backend: unfused + --log-memory-to-tensorboard: true +TEST_TYPE: ckpt-resume diff --git a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml similarity index 98% rename from tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml rename to tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml index 6541e9d35cc..70efb998694 100644 --- a/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml +++ b/tests/functional_tests/test_cases/gpt/gpt3_mcore_te_tp1_pp2_dsv4_hybrid_mhc_mtp/model_config.yaml @@ -23,6 +23,7 @@ MODEL_ARGS: --csa-window-size: 128 --csa-compress-ratios: ([0,4,128,4,0]) --csa-compress-rotary-base: 40000 + --no-dsa-kernel-fusion: true --attention-backend: fused --enable-hyper-connections: true --num-residual-streams: 4 diff --git a/tests/test_utils/recipes/gb200/gpt.yaml b/tests/test_utils/recipes/gb200/gpt.yaml index e10cce0cc3c..0360470c18f 100644 --- a/tests/test_utils/recipes/gb200/gpt.yaml +++ b/tests/test_utils/recipes/gb200/gpt.yaml @@ -124,6 +124,11 @@ products: - environment: [dev] scope: [nightly] platforms: [dgx_gb200] + # - test_case: [gpt3_mcore_te_tp1_pp2_dsv4_hybrid_fused] + # products: + # - environment: [dev] + # scope: [mr, mr-github, mr-github-slim] + # platforms: [dgx_gb200] # - test_case: [gpt3_mcore_te_tp1_pp2_resume_torch_dist_cp4_a2a_p2p_nondeterministic] # products: # - environment: [dev] diff --git a/tests/test_utils/recipes/h100/gpt.yaml b/tests/test_utils/recipes/h100/gpt.yaml index 9e74d25a87c..be2371f1c4d 100644 --- a/tests/test_utils/recipes/h100/gpt.yaml +++ b/tests/test_utils/recipes/h100/gpt.yaml @@ -133,6 +133,11 @@ products: platforms: [dgx_h100] - environment: [lts] scope: [nightly] + - test_case: [gpt3_mcore_te_tp1_pp2_dsv4_hybrid_mhc_mtp] + products: + - environment: [dev] + scope: [mr, mr-github, mr-github-slim] + platforms: [dgx_h100] # - test_case: [gpt3_mcore_te_tp1_pp2_resume_torch_dist_cp4_a2a_p2p_nondeterministic] # products: # - environment: [dev] @@ -367,11 +372,6 @@ products: - environment: [dev] scope: [mr, mr-github, mr-github-slim] platforms: [dgx_h100] - - test_case: [gpt3_mcore_te_tp2_pp2_dsv4_hybrid_mhc_mtp] - products: - - environment: [dev] - scope: [mr, mr-github, mr-github-slim] - platforms: [dgx_h100] - test_case: [gpt3_mcore_te_tp2_pp2_resume_torch_dist_ddp_average_in_collective] products: - environment: [dev] diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index b7e8365804a..a84389a8057 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -36,6 +36,7 @@ "actual_vocab_size": 131072, "add_bias_linear": False, "add_qkv_bias": False, + "apply_dsa_kernel_fusion": True, "apply_query_key_layer_scaling": False, "apply_residual_connection_post_layernorm": False, "apply_rope_fusion": False, diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py index 757b9dd283a..5e1a6827e54 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_dsa.py @@ -212,6 +212,54 @@ def test_dsa_indexer_loss_sparse(self, seqlen_and_topk): assert loss_sparse >= 0 assert loss_dense >= 0 + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_indexer_loss_per_token_scale(self, seqlen_and_topk): + batch_size = 2 + seqlen = seqlen_and_topk[0] + num_heads = 4 + head_dim = 128 + index_topk = seqlen_and_topk[1] + + index_scores = torch.randn(batch_size, seqlen, seqlen, dtype=torch.float32).cuda() + causal_mask = torch.triu( + torch.full( + (seqlen, seqlen), float('-inf'), dtype=torch.float32, device=index_scores.device + ), + diagonal=1, + ) + masked_index_scores = index_scores + causal_mask + topk_k = min(index_topk, seqlen) + topk_indices = masked_index_scores.topk(topk_k, dim=-1)[1] + + query = torch.randn(seqlen, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda() + key = torch.randn(seqlen, batch_size, num_heads, head_dim, dtype=torch.bfloat16).cuda() + softmax_scale = head_dim**-0.5 + + for sparse_loss in [False, True]: + loss_mean = compute_dsa_indexer_loss( + index_scores=index_scores.clone(), + topk_indices=topk_indices, + query=query, + key=key, + softmax_scale=softmax_scale, + loss_coeff=1.0, + sparse_loss=sparse_loss, + pg_collection=self.pg_collection, + ) + loss_sum = compute_dsa_indexer_loss( + index_scores=index_scores.clone(), + topk_indices=topk_indices, + query=query, + key=key, + softmax_scale=softmax_scale, + loss_coeff=1.0, + sparse_loss=sparse_loss, + pg_collection=self.pg_collection, + calculate_per_token_loss=True, + ) + + assert torch.allclose(loss_sum, loss_mean * (batch_size * seqlen), rtol=1e-3, atol=1e-3) + class TestDSAIndexerLossAutoScaler: """Test DSAIndexerLossAutoScaler autograd function.""" @@ -274,6 +322,8 @@ def test_backward_pass(self): atol=0, ), f"Gradient should be scaled by loss scale, expected {expected_grad_per_element}, got {dummy_input.grad[0].item()}" + DSAIndexerLossAutoScaler.main_loss_backward_scale = None + class TestFusedDSAIndexerLossGradient: """Test that FusedDSAIndexerLoss manual backward matches autograd backward.""" @@ -290,7 +340,8 @@ def setup_method(self, request): Utils.destroy_model_parallel() @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") - def test_fused_indexer_loss_gradient_matches_autograd(self): + @pytest.mark.parametrize("calculate_per_token_loss", [False, True]) + def test_fused_indexer_loss_gradient_matches_autograd(self, calculate_per_token_loss): """ Test that the manually written backward in FusedDSAIndexerLoss produces the same gradients as PyTorch autograd on the unfused implementation. @@ -305,7 +356,10 @@ def test_fused_indexer_loss_gradient_matches_autograd(self): for seqlen, index_topk in [[16, 8], [32, 16], [64, 32]]: for sparse_loss in [False, True]: - tag = f"[seqlen={seqlen}, topk={index_topk}, sparse={sparse_loss}]" + tag = ( + f"[seqlen={seqlen}, topk={index_topk}, sparse={sparse_loss}, " + f"per_token={calculate_per_token_loss}]" + ) torch.manual_seed(42) q_ref = ( @@ -351,6 +405,7 @@ def test_fused_indexer_loss_gradient_matches_autograd(self): loss_coeff=loss_coeff, sparse_loss=sparse_loss, pg_collection=self.pg_collection, + calculate_per_token_loss=calculate_per_token_loss, ) loss_ref.backward() @@ -375,6 +430,7 @@ def test_fused_indexer_loss_gradient_matches_autograd(self): mask, sparse_loss, self.pg_collection, + calculate_per_token_loss, ) loss_fused.backward() @@ -472,6 +528,7 @@ def test_fused_indexer_loss_gradient_tp_consistency(self): mask, sparse_loss, pg_collection_tp1, + False, ) loss_tp1.backward() @@ -528,6 +585,7 @@ def test_fused_indexer_loss_gradient_tp_consistency(self): mask, sparse_loss, pg_collection_tpn, + False, ) loss_tpn.backward() diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py new file mode 100644 index 00000000000..974c90f83d8 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsa_kernels.py @@ -0,0 +1,2406 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Unit tests for ``megatron.core.transformer.experimental_attention_variant.dsa_kernels``. + +Coverage: + +* Pure-Python helpers: :func:`local_to_global_flat`, :func:`build_flat_topk_idxs`, + :func:`_kl_loss_from_target_predict` — full correctness checks; no GPU + kernels required (CPU is fine). +* Lazy-import gates: :func:`_ensure_flash_mla`, :func:`_ensure_dsa_namespace` + raise informative ``ImportError`` when the optional packages are missing. +* GPU helpers: :func:`_get_topk_alignment` — runs only on CUDA. +* Wrapper functions :func:`_dsa_fwd_flash_mla`, :func:`indexer_topk`, + :func:`dsa_sparse_attn`, :func:`fused_indexer_sparse_attn` — exercised with + ``unittest.mock`` stand-ins for the underlying ``flash_mla`` / + ``cudnn.DSA`` kernels so the data-marshalling logic (shape conversions, + TopK padding, predict/target/KL composition, autograd plumbing) is + validated without requiring the real CUDA kernels. +""" + +from __future__ import annotations + +import math +import sys +import types +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from megatron.core.transformer.experimental_attention_variant import dsa_kernels as dk +from megatron.core.transformer.experimental_attention_variant.dsa_kernels import ( + FusedIndexerSparseAttnFunc, + SparseAttnFunc, + _dsa_fwd_flash_mla, + _ensure_dsa_namespace, + _ensure_flash_mla, + _get_topk_alignment, + _kl_loss_from_dense_scores, + _kl_loss_from_target_predict, + build_flat_topk_idxs, + dsa_sparse_attn, + fused_indexer_sparse_attn, + indexer_topk, + local_to_global_flat, +) + +# --------------------------------------------------------------------------- +# Test fixtures / helpers +# --------------------------------------------------------------------------- + + +@pytest.fixture +def reset_lazy_kernel_state(): + """Reset the module-level lazy import slots before/after each test. + + The wrapper-function tests patch ``_flash_mla_sparse_fwd`` / ``_DSA`` + directly; we need to ensure each test starts from a clean slate so the + lazy ``_ensure_*`` calls are exercised consistently. + """ + saved_flash = dk._flash_mla_sparse_fwd + saved_dsa = dk._DSA + dk._flash_mla_sparse_fwd = None + dk._DSA = None + yield + dk._flash_mla_sparse_fwd = saved_flash + dk._DSA = saved_dsa + + +def _make_local_idxs(b: int, sq: int, topk: int, *, with_invalid: bool = False) -> torch.Tensor: + """Build a deterministic ``(b, sq, topk)`` int64 tensor of local indices. + + Values for batch ``i``, query ``s``, slot ``k`` are + ``i * 100 + s * 10 + k``. When ``with_invalid`` is True every other + slot is replaced with -1. + """ + base = ( + torch.arange(b, dtype=torch.int64).view(b, 1, 1) * 100 + + torch.arange(sq, dtype=torch.int64).view(1, sq, 1) * 10 + + torch.arange(topk, dtype=torch.int64).view(1, 1, topk) + ) + if with_invalid: + mask = torch.arange(topk).view(1, 1, topk) % 2 == 1 + base = torch.where(mask.expand(b, sq, topk), torch.full_like(base, -1), base) + return base + + +def _uniform_dist(B, S, K, dev): + """Uniform ``1/K`` distribution of shape ``(B, S, K)``.""" + return torch.full((B, S, K), 1.0 / max(K, 1), dtype=torch.float32, device=dev) + + +def _peaked_dist(B, S, K, dev, peak_idx=0): + """Distribution with all probability mass on ``peak_idx``.""" + out = torch.zeros(B, S, K, dtype=torch.float32, device=dev) + out[..., peak_idx] = 1.0 + return out + + +# --------------------------------------------------------------------------- +# local_to_global_flat +# --------------------------------------------------------------------------- + + +class TestLocalToGlobalFlat: + """Pure-Python index conversion (no GPU required).""" + + @pytest.mark.parametrize( + "b, sq, topk, with_invalid", + [ + (1, 4, 5, False), # b=1 identity case + (2, 3, 4, False), # basic multi-batch + (3, 5, 4, False), # larger batch (stresses the formula) + (2, 3, 6, True), # invalid entries interleaved with valid ones + ], + ids=['b1_identity', 'basic', 'larger_b', 'with_invalid'], + ) + def test_global_index_conversion(self, b, sq, topk, with_invalid): + """Shape, dtype, ``-1`` preservation, and the formula + ``global[s*b + bid, k] = local[bid, s, k] * b + bid`` (for valid entries) + in one fixture. Row ``r`` of the output corresponds to query ``s = r // b`` + and batch id ``bid = r % b``. + """ + local = _make_local_idxs(b, sq, topk, with_invalid=with_invalid) + out = local_to_global_flat(local, b, seqlen_kv=128) + + assert out.shape == (sq * b, topk) + assert out.dtype == torch.int32 + + permuted = local.permute(1, 0, 2).reshape(sq * b, topk) + batch_ids = (torch.arange(sq * b) % b).unsqueeze(1) + expected = torch.where( + permuted >= 0, permuted * b + batch_ids, torch.full_like(permuted, -1) + ).int() + assert torch.equal(out, expected) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_cpu_cuda_parity(self): + """CPU and CUDA execution paths produce identical results.""" + local = _make_local_idxs(b=2, sq=4, topk=3, with_invalid=True) + out_cpu = local_to_global_flat(local, 2, seqlen_kv=64) + out_cuda = local_to_global_flat(local.cuda(), 2, seqlen_kv=64) + assert torch.equal(out_cpu, out_cuda.cpu()) + + +# --------------------------------------------------------------------------- +# build_flat_topk_idxs +# --------------------------------------------------------------------------- + + +class TestBuildFlatTopkIdxs: + """Pure-Python multi-group concat + optional compaction.""" + + @pytest.mark.parametrize( + "group_specs", + [ + # Each spec is a list of (topk_i, with_invalid) for each group. + [(4, False)], # single group, all valid + [(2, False), (3, False)], # two groups, all valid + ], + ids=['single_group', 'two_groups'], + ) + def test_non_compact_concat_then_globalise(self, group_specs): + """Without ``compact`` the helper concatenates groups along ``topk`` + and applies the local→global conversion verbatim. + """ + b, sq = 2, 3 + groups = [ + _make_local_idxs(b, sq, t, with_invalid=inv) + 50 * i + for i, (t, inv) in enumerate(group_specs) + ] + total_topk = sum(t for t, _ in group_specs) + + flat, length = build_flat_topk_idxs(*groups, batch_size=b, seqlen_kv=256) + + expected = local_to_global_flat(torch.cat(groups, dim=-1), b, seqlen_kv=256) + assert flat.shape == (sq * b, total_topk) + assert flat.dtype == torch.int32 + assert torch.equal(flat, expected) + assert length is None + + @pytest.mark.parametrize( + "group_specs, expected_valid_per_row", + [ + # Single group: 6 slots with every odd slot invalid → 3 valid. + ([(6, True)], 3), + # Two groups: g1 has 2 valid out of 4, g2 fully valid (2) → 4 valid. + ([(4, True), (2, False)], 4), + ], + ids=['single_group', 'two_groups'], + ) + def test_compact_packs_valid_first(self, group_specs, expected_valid_per_row): + """With ``compact=True`` the helper packs valid entries to the front + of each row, fills the tail with ``-1``, and returns a per-row + ``topk_length`` that equals the count of valid entries. + """ + b, sq = 2, 3 + groups = [ + _make_local_idxs(b, sq, t, with_invalid=inv) + 100 * i + for i, (t, inv) in enumerate(group_specs) + ] + total_topk = sum(t for t, _ in group_specs) + + flat, length = build_flat_topk_idxs(*groups, batch_size=b, seqlen_kv=512, compact=True) + + assert flat.shape == (sq * b, total_topk) + assert flat.dtype == torch.int32 + assert length is not None + assert length.shape == (sq * b,) + assert length.dtype == torch.int32 + + # Per-row layout: valid global indices first, then -1 padding. + for row in range(sq * b): + n = int(length[row]) + assert n == expected_valid_per_row, f"row {row}: wrong length" + assert torch.all(flat[row, :n] >= 0), f"row {row}: leading entries should be valid" + assert torch.all(flat[row, n:] == -1), f"row {row}: trailing entries should be -1" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compact_cuda_path(self, reset_lazy_kernel_state): + """Combined coverage for the CUDA compact path (sub-blocks + self-label on failure): + + * (a) Dispatch + plumbing (mocked compactify): the wrapper is + called exactly once, with the already-globalised + ``(sq*b, total_topk)`` int32 tensor as input, and its + returned ``(indices, topk_length)`` flow back verbatim. + * (b) End-to-end parity (real cuDNN, skipped without it): the + cuDNN ``compactify`` kernel produces the same ``(flat, + length)`` pair as the pure-PyTorch CPU fallback. + """ + # ---- (a) dispatch via mocked compactify -------------------------- + b, sq, topk = 2, 3, 4 + local = _make_local_idxs(b, sq, topk, with_invalid=True).to(torch.int32, copy=False).cuda() + compact_indices = torch.full((sq * b, topk), 99, dtype=torch.int32, device='cuda') + topk_length = torch.full((sq * b,), 7, dtype=torch.int32, device='cuda') + + captured = {} + + def fake_compactify(global_idxs): + captured['input'] = global_idxs + return {'indices': compact_indices, 'topk_length': topk_length} + + fake_dsa = MagicMock(name='_DSA_compactify_stub') + fake_dsa.compactify_wrapper.side_effect = fake_compactify + dk._DSA = fake_dsa + + flat, length = build_flat_topk_idxs(local, batch_size=b, seqlen_kv=512, compact=True) + fake_dsa.compactify_wrapper.assert_called_once() + kernel_input = captured['input'] + assert kernel_input.shape == (sq * b, topk), "(a) wrapper input shape" + assert kernel_input.dtype == torch.int32, "(a) wrapper input dtype" + assert kernel_input.is_cuda, "(a) wrapper input not on CUDA" + expected_input = local_to_global_flat(local, b, seqlen_kv=512) + assert torch.equal( + kernel_input, expected_input + ), "(a) wrapper input != local_to_global_flat(local)" + assert flat is compact_indices, "(a) returned flat is not the kernel output" + assert length is topk_length, "(a) returned length is not the kernel output" + + # ---- (b) real-kernel parity vs CPU fallback ---------------------- + # Skipped when cuDNN is not installed; reset state so the real + # _DSA import happens on the next call inside build_flat_topk_idxs. + try: + cudnn = pytest.importorskip("cudnn") + except pytest.skip.Exception: + return # already passed (a); skip the parity sub-block silently + if not hasattr(cudnn, 'DSA'): + return + dk._DSA = None # force real lazy-import + + b2, sq2 = 4, 5 + local_a = _make_local_idxs(b2, sq2, 6, with_invalid=True) + local_b = _make_local_idxs(b2, sq2, 4, with_invalid=False) + 200 + + flat_cpu, len_cpu = build_flat_topk_idxs( + local_a, local_b, batch_size=b2, seqlen_kv=512, compact=True + ) + flat_cuda, len_cuda = build_flat_topk_idxs( + local_a.cuda(), local_b.cuda(), batch_size=b2, seqlen_kv=512, compact=True + ) + assert torch.equal( + flat_cpu, flat_cuda.cpu() + ), "(b) flat tensor differs between CPU fallback and cuDNN kernel" + assert torch.equal( + len_cpu, len_cuda.cpu() + ), "(b) length tensor differs between CPU fallback and cuDNN kernel" + + +# --------------------------------------------------------------------------- +# _kl_loss_from_target_predict +# --------------------------------------------------------------------------- + + +class TestKLLossFromTargetPredict: + """Pure-Python KL loss computation: combined assertions for all + properties (scalar/dtype, identity, non-negativity, coeff linearity, + invalid-row masking, analytical formula).""" + + def test_kl_loss_properties(self): + """All KL-loss invariants checked sequentially. Each block raises + an informative ``AssertionError`` so a failure pinpoints the + broken sub-property. + """ + torch.manual_seed(0) + b, sq, topk = 2, 3, 4 + topk_indices = torch.zeros(b, sq, topk, dtype=torch.int32) + + # ---- (a) scalar/dtype + identity: KL(p || p) == 0 ----------------- + identical = torch.softmax(torch.randn(b, sq, topk), dim=-1) + loss_identical = _kl_loss_from_target_predict( + identical, identical.clone(), topk_indices, loss_coeff=1.0 + ) + assert loss_identical.shape == torch.Size([]), "identity: not scalar" + assert loss_identical.dtype == torch.float32, "identity: not fp32" + assert torch.allclose( + loss_identical, torch.tensor(0.0), atol=1e-6 + ), f"identity: KL(p || p) != 0 (got {loss_identical.item()})" + + # ---- (b) non-negativity + linearity in loss_coeff ----------------- + target = torch.softmax(torch.randn(b, sq, topk), dim=-1) + predict = torch.softmax(torch.randn(b, sq, topk), dim=-1) + loss_1 = _kl_loss_from_target_predict(target, predict, topk_indices, loss_coeff=1.0) + loss_3 = _kl_loss_from_target_predict(target, predict, topk_indices, loss_coeff=3.0) + assert loss_1.item() >= 0.0, f"non-negativity: got {loss_1.item()}" + assert torch.allclose( + loss_3, 3.0 * loss_1, atol=1e-5, rtol=1e-5 + ), f"linearity: 3*loss_1 = {3*loss_1.item()} vs loss_3 = {loss_3.item()}" + + # ---- (c) invalid-row masking -------------------------------------- + # Construct deterministic distributions with strictly-positive per-row KL. + t_inv = torch.full((b, sq, topk), 0.1, dtype=torch.float32) + t_inv[..., 0] = 0.7 + p_inv = torch.full((b, sq, topk), 0.7, dtype=torch.float32) / topk + p_inv[..., -1] = 1.0 - p_inv[..., :-1].sum(dim=-1) + + idx_all_valid = torch.zeros(b, sq, topk, dtype=torch.int32) + loss_full = _kl_loss_from_target_predict(t_inv, p_inv, idx_all_valid, loss_coeff=1.0) + assert loss_full.item() > 0, "all-valid baseline must be positive" + + # Mark the first row of every batch invalid → fewer valid rows, + # smaller KL sum, same denominator (mean over all (B, S_q)). + idx_partial = idx_all_valid.clone() + idx_partial[:, 0, :] = -1 + loss_partial = _kl_loss_from_target_predict(t_inv, p_inv, idx_partial, loss_coeff=1.0) + assert ( + loss_partial.item() < loss_full.item() + ), f"partial-invalid: {loss_partial.item()} should be < {loss_full.item()}" + + # All-invalid → loss exactly 0. + idx_all_invalid = torch.full_like(idx_all_valid, -1) + loss_zero = _kl_loss_from_target_predict(t_inv, p_inv, idx_all_invalid, loss_coeff=1.0) + assert loss_zero.item() == 0.0, f"all-invalid: got {loss_zero.item()}" + + # ---- (d) analytical formula: target = δ_0, predict = uniform(1/K) - + # per-row KL = log(K); mean = log(K); loss = coeff * log(K). + target_d = _peaked_dist(b, sq, topk, 'cpu', peak_idx=0) + predict_d = torch.full((b, sq, topk), 1.0 / topk, dtype=torch.float32) + loss_d = _kl_loss_from_target_predict(target_d, predict_d, topk_indices, loss_coeff=2.5) + expected = 2.5 * math.log(topk) + assert torch.allclose( + loss_d, torch.tensor(expected), rtol=1e-5, atol=1e-5 + ), f"analytical: {loss_d.item()} vs expected {expected}" + + def test_per_token_loss_reports_raw_sum(self): + torch.manual_seed(1) + b, sq, topk = 2, 5, 4 + target = torch.softmax(torch.randn(b, sq, topk), dim=-1) + predict = torch.softmax(torch.randn(b, sq, topk), dim=-1) + topk_indices = torch.zeros(b, sq, topk, dtype=torch.int32) + + loss_mean = _kl_loss_from_target_predict(target, predict, topk_indices, loss_coeff=0.5) + loss_sum = _kl_loss_from_target_predict( + target, predict, topk_indices, loss_coeff=0.5, calculate_per_token_loss=True + ) + + assert torch.allclose(loss_sum, loss_mean * (b * sq), rtol=1e-5, atol=1e-5) + + +class TestKLLossFromDenseScores: + def test_per_token_loss_reports_raw_sum(self): + b, sq, sk = 2, 5, 4 + loss_coeff = 0.5 + + attn_score = _peaked_dist(b, sq, sk, 'cpu', peak_idx=0) + attn_l1norm = torch.ones(b, sq, dtype=torch.float32) + index_score = torch.zeros(b, sq, sk, dtype=torch.float32) + index_lse = torch.full((b, sq), math.log(sk), dtype=torch.float32) + + loss_mean = _kl_loss_from_dense_scores( + attn_score, attn_l1norm, index_score, index_lse, loss_coeff + ) + loss_sum = _kl_loss_from_dense_scores( + attn_score, + attn_l1norm, + index_score, + index_lse, + loss_coeff, + calculate_per_token_loss=True, + ) + + assert torch.allclose(loss_sum, loss_mean * (b * sq), rtol=1e-5, atol=1e-5) + + +# --------------------------------------------------------------------------- +# _ensure_flash_mla / _ensure_dsa_namespace +# --------------------------------------------------------------------------- + + +_LAZY_IMPORT_CASES = [ + pytest.param( + 'flash_mla', + 'flash_mla_sparse_fwd', + _ensure_flash_mla, + '_flash_mla_sparse_fwd', + "FlashMLA is required", + id='flash_mla', + ), + pytest.param( + 'cudnn', 'DSA', _ensure_dsa_namespace, '_DSA', "cudnn-frontend DSA", id='cudnn_dsa' + ), +] + + +@pytest.mark.parametrize( + "module_name, attr_name, ensure_fn, slot_name, error_match", _LAZY_IMPORT_CASES +) +class TestLazyKernelImports: + """Lazy-import behaviour shared by ``_ensure_flash_mla`` and + ``_ensure_dsa_namespace``: combined error-on-missing + caches-on-success + fixture (assertion blocks self-label on failure). + """ + + def test_lazy_import_raises_and_caches( + self, reset_lazy_kernel_state, module_name, attr_name, ensure_fn, slot_name, error_match + ): + # ---- (a) raises informative ImportError when the module is absent -- + # Setting ``sys.modules[name] = None`` makes ``import name`` fail. + with patch.dict(sys.modules, {module_name: None}): + with pytest.raises(ImportError, match=error_match): + ensure_fn() + + # ---- (b) caches the import on success ------------------------------ + sentinel = MagicMock(name=f"{attr_name}_sentinel") + fake_module = types.ModuleType(module_name) + setattr(fake_module, attr_name, sentinel) + + with patch.dict(sys.modules, {module_name: fake_module}): + ensure_fn() + assert ( + getattr(dk, slot_name) is sentinel + ), f"(b) {module_name}: ensure_fn() did not bind sentinel" + + # Second call must be a no-op — even after the sys.modules entry is gone. + with patch.dict(sys.modules, {}, clear=False): + sys.modules.pop(module_name, None) + ensure_fn() + assert ( + getattr(dk, slot_name) is sentinel + ), f"(b) {module_name}: cached sentinel was lost on 2nd call" + + +# --------------------------------------------------------------------------- +# _get_topk_alignment +# --------------------------------------------------------------------------- + + +class TestGetTopkAlignment: + """Architecture-dependent alignment for FlashMLA top-K padding.""" + + @pytest.mark.parametrize( + "sm_major, expected", [(7, 128), (8, 128), (9, 128), (10, 64), (12, 64), (13, 64)] + ) + def test_alignment_per_sm(self, sm_major, expected): + """SM10x and newer use 64-byte TopK alignment; older arches use 128.""" + with patch('torch.cuda.get_device_capability', return_value=(sm_major, 0)): + assert _get_topk_alignment() == expected + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_runs_on_real_gpu(self): + """On any real GPU the alignment must agree with the documented rule.""" + align = _get_topk_alignment() + sm = torch.cuda.get_device_capability() + expected = 64 if sm[0] >= 10 else 128 + assert align == expected + + +# --------------------------------------------------------------------------- +# _dsa_fwd_flash_mla — wrapper around flash_mla.flash_mla_sparse_fwd +# --------------------------------------------------------------------------- + + +def _make_flash_mla_stub(d_v: int = 512, *, lse_scalar: float = 0.0, out_fill: float = 0.0): + """Build a callable stand-in for ``flash_mla.flash_mla_sparse_fwd``. + + The real kernel signature is + ``(q, kv, indices, softmax_scale, d_v, attn_sink, topk_length, indexer_topk)`` + and returns ``(out, max_logits, lse)`` or ``(out, max_logits, lse, lse_indexer)`` + when ``indexer_topk > 0``. + + The stub returns deterministic, easily-distinguishable tensors so callers + can numerically verify the wrapper's reshape / split logic. The most + recent ``out`` and ``lse`` are stashed on ``stub.last_out`` / + ``stub.last_lse`` for direct equality checks. + """ + + stub = MagicMock(name='flash_mla_sparse_fwd_stub') + + def _impl(q, kv, indices, softmax_scale, d_v, attn_sink, topk_length, indexer_topk): + total_S_q, H, _D = q.shape + # Distinguishable per-element pattern: out[i, h, k] = out_fill + i + 0.001*h + 1e-6*k + # (works in bf16 at this magnitude, useful for verifying that the + # wrapper does not silently reshape across the wrong axes). + idx_i = torch.arange(total_S_q, dtype=torch.float32, device=q.device).view(-1, 1, 1) + idx_h = torch.arange(H, dtype=torch.float32, device=q.device).view(1, -1, 1) + idx_k = torch.arange(d_v, dtype=torch.float32, device=q.device).view(1, 1, -1) + out_f32 = out_fill + idx_i + 0.001 * idx_h + 1e-6 * idx_k + out = out_f32.to(q.dtype) + max_logits = torch.zeros(total_S_q, H, dtype=torch.float32, device=q.device) + # lse[i, h] = lse_scalar + i + 0.5*h — a deterministic pattern. + lse = lse_scalar + ( + torch.arange(total_S_q, dtype=torch.float32, device=q.device).view(-1, 1) + + 0.5 * torch.arange(H, dtype=torch.float32, device=q.device).view(1, -1) + ) + stub.last_out = out + stub.last_lse = lse + if indexer_topk > 0: + # Make lse_indexer distinct from lse so we can tell which one the + # wrapper returned. + lse_indexer = lse + 100.0 + stub.last_lse_indexer = lse_indexer + return out, max_logits, lse, lse_indexer + stub.last_lse_indexer = None + return out, max_logits, lse + + stub.side_effect = _impl + stub.last_out = None + stub.last_lse = None + stub.last_lse_indexer = None + return stub + + +class TestDsaFwdFlashMla: + """Adapter logic around FlashMLA: shape massaging, TopK padding, return + tuples — including numerical pass-through of the kernel outputs. + """ + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_fwd_flash_mla_adapter(self, reset_lazy_kernel_state): + """All adapter behaviours in one fixture (assertion blocks self-label + on failure): + + * (a) TopK is padded up to GPU-specific alignment; padded slots are + ``-1``; ``out`` / ``lse`` are passed through verbatim; + kernel arg shapes match the SBHD-flat → MQA-h_kv=1 contract. + * (b) ``indexer_topk == 0`` -> ``lse_indexer is None``. + * (c) ``0 < indexer_topk < TopK`` -> kernel's ``lse_indexer`` is + returned verbatim (not silently swapped for ``lse``). + * (d) ``indexer_topk == TopK`` -> fallback to ``lse.clone()`` (a + kernel snapshot quirk). + * (e) ``indexer_topk > 0 + topk_length`` is rejected with the + expected error. + """ + total_sq, H, D = 4, 2, 512 + align = _get_topk_alignment() + + q = torch.randn(total_sq, H, D, dtype=torch.bfloat16, device='cuda') + kv = torch.randn(8, D, dtype=torch.bfloat16, device='cuda') + + # ---- (a) padding + numerical pass-through ------------------------ + TopK_unpadded = 5 + expected_padded = ((TopK_unpadded + align - 1) // align) * align + topk_idxs = torch.arange(total_sq * TopK_unpadded, dtype=torch.int32, device='cuda').view( + total_sq, TopK_unpadded + ) + + stub = _make_flash_mla_stub(d_v=D) + dk._flash_mla_sparse_fwd = stub + + out, lse, lse_indexer = _dsa_fwd_flash_mla(q, kv, topk_idxs, softmax_scale=0.5, d_v=D) + assert lse_indexer is None, "(a) lse_indexer should be None when indexer_topk=0" + assert torch.equal(out, stub.last_out), "(a) out is not pass-through" + assert torch.equal(lse, stub.last_lse), "(a) lse is not pass-through" + + called_args = stub.call_args.args + kv_3d, indices_arg = called_args[1], called_args[2] + assert kv_3d.shape == (8, 1, D), f"(a) KV shape {tuple(kv_3d.shape)} != (8, 1, {D})" + assert indices_arg.shape == (total_sq, 1, expected_padded), ( + f"(a) indices shape {tuple(indices_arg.shape)} != " + f"({total_sq}, 1, {expected_padded})" + ) + if expected_padded > TopK_unpadded: + assert torch.all( + indices_arg[..., TopK_unpadded:] == -1 + ), "(a) padded slots should be -1" + assert torch.equal( + indices_arg[..., :TopK_unpadded].squeeze(1), topk_idxs + ), "(a) original entries should survive padding unchanged" + + # ---- (b–d) indexer_topk branches --------------------------------- + TopK = align # already aligned, no padding + topk_idxs_aligned = torch.zeros(total_sq, TopK, dtype=torch.int32, device='cuda') + + stub = _make_flash_mla_stub(d_v=D, lse_scalar=1.5) + dk._flash_mla_sparse_fwd = stub + + # (b) indexer_topk == 0 + _, _, lse_idx_b = _dsa_fwd_flash_mla(q, kv, topk_idxs_aligned, 0.5, indexer_topk=0) + assert lse_idx_b is None, "(b) indexer_topk=0 must yield lse_indexer=None" + + # (c) 0 < indexer_topk < TopK + _, lse_c, lse_idx_c = _dsa_fwd_flash_mla( + q, kv, topk_idxs_aligned, 0.5, indexer_topk=TopK // 2 + ) + assert lse_idx_c is not None, "(c) lse_indexer should be present" + assert torch.equal( + lse_idx_c, stub.last_lse_indexer + ), "(c) lse_indexer should be kernel pass-through" + assert not torch.equal(lse_idx_c, lse_c), "(c) wrapper silently swapped lse_indexer for lse" + + # (d) indexer_topk == TopK -> fallback to lse.clone() + _, lse_d, lse_idx_d = _dsa_fwd_flash_mla(q, kv, topk_idxs_aligned, 0.5, indexer_topk=TopK) + assert torch.equal( + lse_idx_d, lse_d + ), "(d) lse_indexer should equal lse on TopK-cap fallback" + assert ( + lse_idx_d.data_ptr() != lse_d.data_ptr() + ), "(d) fallback should be a clone, not an alias" + + # ---- (e) topk_length + indexer_topk > 0 is rejected -------------- + # Use CPU tensors here — the assert fires before any kernel call. + with pytest.raises(AssertionError, match="indexer_topk > 0 requires non-compact"): + _dsa_fwd_flash_mla( + torch.zeros(2, 2, 512, dtype=torch.bfloat16), + torch.zeros(4, 512, dtype=torch.bfloat16), + torch.zeros(2, 8, dtype=torch.int32), + softmax_scale=0.5, + topk_length=torch.zeros(2, dtype=torch.int32), + indexer_topk=4, + ) + + +# --------------------------------------------------------------------------- +# indexer_topk — cudnn DSA wrapper for inference +# --------------------------------------------------------------------------- + + +class TestIndexerTopk: + """Indexer scoring + radix top-K wrapper. All three properties combined + in a single fixture; sub-block names appear in failure messages. + """ + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_indexer_topk_wrapper(self, reset_lazy_kernel_state): + """Combined assertions for: + + * (a) basic call: shapes / dtypes of the returned (topk_indices, + topk_length); kernels are called with the right BSHD layouts + and the SBHD-flat (b*sq, sk) scores; indexer_top_k kwargs. + * (b) topk > sk clamping: kernel call uses ``sk`` keys, trailing + ``[sk:]`` slots are -1, ``topk_length == sk``. + * (c) ``indexer_softmax_scale`` pre-scales the weights via the + ``relu(c·x) = c·relu(x)`` trick before reaching the kernel. + """ + + # ---- (a) basic call ---------------------------------------------- + sq, b, idx_nh, idx_hd = 6, 2, 4, 64 + sk = 12 + topk = 5 + ratio = 4 + + q_indexer = torch.randn(sq, b, idx_nh, idx_hd, dtype=torch.bfloat16, device='cuda') + k_indexer = torch.randn(sk, b, idx_hd, dtype=torch.bfloat16, device='cuda') + weights = torch.randn(sq, b, idx_nh, dtype=torch.bfloat16, device='cuda') + + scores = torch.randn(b, sq, sk, dtype=torch.float32, device='cuda') + captured = {} + + def fake_indexer_forward(q_bshd, k_bshd, w_bsh, ratio): + captured['indexer_forward'] = { + 'q_shape': q_bshd.shape, + 'k_shape': k_bshd.shape, + 'w_shape': w_bsh.shape, + 'ratio': ratio, + } + return {'scores': scores} + + def fake_filtered_topk(scores_flat, seq_lens, top_k, next_n, return_val): + captured['filtered_topk'] = { + 'scores_shape': scores_flat.shape, + 'seq_lens_shape': seq_lens.shape, + 'top_k': top_k, + 'next_n': next_n, + 'return_val': return_val, + } + n_rows = scores_flat.shape[0] + return {'indices': torch.zeros(n_rows, top_k, dtype=torch.int32, device='cuda')} + + fake_dsa = MagicMock() + 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 + + topk_indices, topk_length = indexer_topk( + q_indexer, k_indexer, weights, topk=topk, ratio=ratio + ) + + assert topk_indices.shape == (b, sq, topk), "(a) topk_indices shape" + assert topk_indices.dtype == torch.int32, "(a) topk_indices dtype" + assert topk_length.shape == (b, sq), "(a) topk_length shape" + assert topk_length.dtype == torch.int32, "(a) topk_length dtype" + # BSHD / BSD layouts handed to the kernels. + assert captured['indexer_forward']['q_shape'] == ( + b, + sq, + idx_nh, + idx_hd, + ), "(a) indexer_forward q_shape" + assert captured['indexer_forward']['k_shape'] == ( + b, + sk, + 1, + idx_hd, + ), "(a) indexer_forward k_shape (must be unsqueezed h_kv=1)" + assert captured['indexer_forward']['w_shape'] == ( + b, + sq, + idx_nh, + ), "(a) indexer_forward w_shape" + assert captured['indexer_forward']['ratio'] == ratio, "(a) ratio kwarg" + assert captured['filtered_topk']['scores_shape'] == ( + b * sq, + sk, + ), "(a) topK scores_flat shape" + assert captured['filtered_topk']['seq_lens_shape'] == (b * sq,), "(a) topK seq_lens shape" + assert captured['filtered_topk']['top_k'] == min(topk, sk), "(a) top_k kwarg" + assert captured['filtered_topk']['next_n'] == 1, "(a) next_n kwarg" + assert captured['filtered_topk']['return_val'] is False, "(a) return_val kwarg" + + # ---- (b) topk > sk clamping -------------------------------------- + dk._DSA = None # force fresh mocks + sq2, b2, idx_nh2, idx_hd2 = 4, 1, 2, 32 + sk2 = 3 + topk2 = 8 # > sk + q2 = torch.randn(sq2, b2, idx_nh2, idx_hd2, dtype=torch.bfloat16, device='cuda') + k2 = torch.randn(sk2, b2, idx_hd2, dtype=torch.bfloat16, device='cuda') + w2 = torch.randn(sq2, b2, idx_nh2, dtype=torch.bfloat16, device='cuda') + scores2 = torch.zeros(b2, sq2, sk2, dtype=torch.float32, device='cuda') + kernel_indices2 = torch.zeros(b2 * sq2, sk2, dtype=torch.int32, device='cuda') + + fake_dsa_b = MagicMock() + 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 + + topk_indices2, topk_length2 = indexer_topk(q2, k2, w2, topk=topk2, ratio=4) + assert topk_indices2.shape == (b2, sq2, topk2), "(b) padded topk_indices shape" + assert torch.all(topk_indices2[..., sk2:] == -1), "(b) trailing slots not -1" + assert torch.all(topk_length2 == sk2), "(b) topk_length should equal sk" + + # ---- (c) indexer_softmax_scale pre-scales weights --------------- + dk._DSA = None + sq3, b3, idx_nh3, idx_hd3 = 2, 1, 2, 32 + sk3 = 4 + scale = 0.125 + q3 = torch.zeros(sq3, b3, idx_nh3, idx_hd3, dtype=torch.bfloat16, device='cuda') + k3 = torch.zeros(sk3, b3, idx_hd3, dtype=torch.bfloat16, device='cuda') + w3 = torch.full((sq3, b3, idx_nh3), 8.0, dtype=torch.bfloat16, device='cuda') + captured_w = {} + + def fake_indexer_forward_c(q_bshd, k_bshd, w_bsh, ratio): + captured_w['w'] = w_bsh.detach().clone() + return {'scores': torch.zeros(b3, sq3, sk3, dtype=torch.float32, device='cuda')} + + fake_dsa_c = MagicMock() + 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') + } + dk._DSA = fake_dsa_c + + indexer_topk(q3, k3, w3, topk=sk3, ratio=4, indexer_softmax_scale=scale) + expected_w = ( + (w3.float() * scale).to(torch.bfloat16).permute(1, 0, 2).reshape(b3, sq3, idx_nh3) + ) + assert torch.allclose( + captured_w['w'].float(), expected_w.float(), atol=1e-2, rtol=1e-2 + ), "(c) weights were not pre-scaled by indexer_softmax_scale" + + +# --------------------------------------------------------------------------- +# dsa_sparse_attn / SparseAttnFunc forward (mocked) +# --------------------------------------------------------------------------- + + +class TestDsaSparseAttn: + """Numerical fwd + bwd test for the public ``dsa_sparse_attn`` entry + point. The underlying kernels are mocked so the whole wrapper — including + the SBHD↔flat reshape on the forward and the autograd plumbing on the + backward — can be checked against deterministic ground truth. + """ + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dsa_sparse_attn_fwd_and_bwd(self, reset_lazy_kernel_state): + """Combined forward + backward fixture (assertion blocks self-label + on failure): + + * (a) Forward output equals the FlashMLA stub's ``out`` reshaped + from flat ``(sq*b, np_, d_v)`` back to ``(sq, b, np_ * d_v)``. + * (b) Backward maps kernel grads onto the right SBHD leaf tensors + with the correct shapes; the bwd kernel is invoked exactly + once. + """ + sq, b, np_, d = 4, 2, 2, 512 + skv = 6 + TopK = _get_topk_alignment() + + query = torch.randn(sq, b, np_, d, dtype=torch.bfloat16, device='cuda', requires_grad=True) + kv = torch.randn(skv, b, d, dtype=torch.bfloat16, device='cuda', requires_grad=True) + attn_sink = torch.zeros(np_, dtype=torch.float32, device='cuda', requires_grad=True) + topk_idxs = torch.zeros(sq * b, TopK, dtype=torch.int32, device='cuda') + + # Coordinated stubs: FlashMLA fwd + cuDNN sparse-attn bwd, both + # deterministic so every gradient slot is independently verifiable. + flash_stub = _make_flash_mla_stub(d_v=d) + dq_kernel = torch.full((sq * b, np_, d), 7.0, dtype=torch.bfloat16, device='cuda') + dkv_kernel = torch.full((skv * b, d), -3.0, dtype=torch.bfloat16, device='cuda') + d_sink_kernel = torch.full((np_,), 11.0, dtype=torch.float32, device='cuda') + fake_dsa = MagicMock() + fake_dsa.sparse_attention_backward_wrapper.return_value = { + 'dq': dq_kernel, + 'dkv': dkv_kernel, + 'd_sink': d_sink_kernel, + } + dk._flash_mla_sparse_fwd = flash_stub + dk._DSA = fake_dsa + + out = dsa_sparse_attn(query, kv, attn_sink, topk_idxs, softmax_scale=0.5) + + # ---- (a) forward ------------------------------------------------ + assert out.shape == (sq, b, np_ * d), "(a) forward shape" + assert out.dtype == torch.bfloat16, "(a) forward dtype" + expected_out = flash_stub.last_out.reshape(sq, b, np_, d).reshape(sq, b, np_ * d) + assert torch.equal(out, expected_out), "(a) forward value pass-through" + + # ---- (b) backward ----------------------------------------------- + out.sum().backward() + assert query.grad is not None, "(b) query.grad missing" + assert kv.grad is not None, "(b) kv.grad missing" + assert attn_sink.grad is not None, "(b) attn_sink.grad missing" + assert torch.equal( + query.grad, dq_kernel.reshape(sq, b, np_, d) + ), "(b) query.grad mis-reshaped" + assert torch.equal(kv.grad, dkv_kernel.reshape(skv, b, d)), "(b) kv.grad mis-reshaped" + assert torch.equal(attn_sink.grad, d_sink_kernel), "(b) attn_sink.grad mismatch" + fake_dsa.sparse_attention_backward_wrapper.assert_called_once() + + +# --------------------------------------------------------------------------- +# fused_indexer_sparse_attn — Path B autograd Function (mocked) +# --------------------------------------------------------------------------- + + +def _install_full_dsa_mock( + *, + b: int, + sq: int, + np_: int, + d: int, + n_comp: int, + idx_nh: int, + predict_fn=None, + target_fn=None, + dq_value: float = 7.0, + dkv_value: float = -3.0, + d_sink_value: float = 11.0, + d_index_q_value: float = 0.5, + d_weights_value: float = -0.25, + d_index_k_value: float = 1.5, +): + """Patch the module-level ``_DSA`` and ``_flash_mla_sparse_fwd`` slots + with a coordinated set of deterministic stubs covering every kernel + invoked by :class:`FusedIndexerSparseAttnFunc`. + + ``predict_fn`` / ``target_fn`` (if provided) build the per-row + distribution given ``(b, sq, topk, device)``. By default both return a + uniform ``1/topk`` distribution, which yields ``KL(target || predict) = 0`` + so the loss is exactly zero. + + All backward kernels return constant-filled tensors so each gradient slot + can be independently verified. + """ + + if predict_fn is None: + predict_fn = lambda B, S, K, dev: torch.full( + (B, S, K), 1.0 / max(K, 1), dtype=torch.float32, device=dev + ) + if target_fn is None: + target_fn = predict_fn + + fake_dsa = MagicMock(name='_DSA_full_stub') + + 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)} + + fake_dsa.indexer_forward_wrapper.side_effect = fake_indexer_forward + + def fake_filtered_topk(scores_flat, seq_lens, top_k, next_n, return_val): + return { + 'indices': torch.zeros( + scores_flat.shape[0], top_k, dtype=torch.int32, device=scores_flat.device + ) + } + + fake_dsa.indexer_top_k_wrapper.side_effect = fake_filtered_topk + + def fake_sparse_indexer_score_backward(q, k, w, topk_indices, qhead_per_kv_head): + topk = topk_indices.shape[-1] + return {'predict': predict_fn(b, sq, topk, q.device)} + + fake_dsa.sparse_indexer_score_recompute_wrapper.side_effect = fake_sparse_indexer_score_backward + + def fake_sparse_attn_score_backward(q, k, lse, topk_indices, sm_scale, qhead_per_kv_head): + topk = topk_indices.shape[-1] + return {'target': target_fn(b, sq, topk, q.device)} + + fake_dsa.sparse_attn_score_recompute_wrapper.side_effect = fake_sparse_attn_score_backward + + def fake_sparse_attn_backward(q, kv, out, dout, lse, attn_sink, topk_idxs, **kwargs): + return { + 'dq': torch.full_like(q, dq_value), + 'dkv': torch.full_like(kv, dkv_value), + 'd_sink': torch.full_like(attn_sink, d_sink_value), + } + + fake_dsa.sparse_attention_backward_wrapper.side_effect = fake_sparse_attn_backward + + def fake_indexer_grad_backward( + q_idx_bshd, + w_bsh, + k_idx_bsd, + attn_score, + index_score, + topk_indices, + sm_scale, + loss_coeff, + grad_loss, + block_I, + ): + return { + 'd_index_q': torch.full_like(q_idx_bshd, d_index_q_value), + 'd_weights': torch.full_like(w_bsh, d_weights_value), + 'd_index_k': torch.full_like(k_idx_bsd, d_index_k_value), + } + + fake_dsa.indexer_backward_wrapper.side_effect = fake_indexer_grad_backward + + flash_stub = _make_flash_mla_stub(d_v=d) + + dk._DSA = fake_dsa + dk._flash_mla_sparse_fwd = flash_stub + return fake_dsa, flash_stub + + +def _install_full_dsa_mock_dense( + *, + b: int, + sq: int, + np_: int, + d: int, + n_comp: int, + idx_nh: int, + target_score_fn=None, + target_l1norm_fn=None, + predict_score_fn=None, + predict_lse_fn=None, + dq_value: float = 7.0, + dkv_value: float = -3.0, + d_sink_value: float = 11.0, + d_index_q_value: float = 0.5, + d_weights_value: float = -0.25, + d_index_k_value: float = 1.5, +): + """Coordinated stubs covering the dense-loss (``sparse_loss=False``) path. + + Mirrors :func:`_install_full_dsa_mock` for the sparse path, but stubs + the four dense-only kernel wrappers: + + * ``dense_indexer_score_recompute_wrapper`` -> ``(out, denom=index_lse)`` + * ``dense_attn_score_recompute_wrapper`` -> ``(out, denom=attn_l1norm)`` + * ``dense_indexer_backward_wrapper`` -> ``{d_index_q, d_weights, d_index_k}`` + + Defaults make ``target == predict == uniform(1/n_comp)`` so KL == 0. + Override the four ``*_fn`` callables to drive the loss to known + analytical values; each callable receives ``(B, S_q, S_k, device)`` and + returns the score tensor (``S_k``-dim) or denom (no ``S_k`` dim). + """ + + if target_score_fn is None: + target_score_fn = lambda B, S, K, dev: torch.full( + (B, S, K), 1.0 / max(K, 1), dtype=torch.float32, device=dev + ) + if target_l1norm_fn is None: + target_l1norm_fn = lambda B, S, K, dev: torch.ones((B, S), dtype=torch.float32, device=dev) + if predict_score_fn is None: + predict_score_fn = lambda B, S, K, dev: torch.zeros( + (B, S, K), dtype=torch.float32, device=dev + ) + if predict_lse_fn is None: + predict_lse_fn = lambda B, S, K, dev: torch.full( + (B, S), float(math.log(max(K, 1))), dtype=torch.float32, device=dev + ) + + fake_dsa = MagicMock(name='_DSA_full_dense_stub') + + 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)} + + fake_dsa.indexer_forward_wrapper.side_effect = fake_indexer_forward + + def fake_filtered_topk(scores_flat, seq_lens, top_k, next_n, return_val): + return { + 'indices': torch.zeros( + scores_flat.shape[0], top_k, dtype=torch.int32, device=scores_flat.device + ) + } + + fake_dsa.indexer_top_k_wrapper.side_effect = fake_filtered_topk + + def fake_dense_indexer_score(q, k, w, qhead_per_kv_head, sm_scale, ratio): + dev = q.device + return { + 'out': predict_score_fn(b, sq, n_comp, dev), + 'denom': predict_lse_fn(b, sq, n_comp, dev), + } + + fake_dsa.dense_indexer_score_recompute_wrapper.side_effect = fake_dense_indexer_score + + def fake_dense_attn_score(q, k, lse, softmax_scale, qhead_per_kv_head, ratio): + dev = q.device + return { + 'out': target_score_fn(b, sq, n_comp, dev), + 'denom': target_l1norm_fn(b, sq, n_comp, dev), + } + + fake_dsa.dense_attn_score_recompute_wrapper.side_effect = fake_dense_attn_score + + def fake_sparse_attn_backward(q, kv, out, dout, lse, attn_sink, topk_idxs, **kwargs): + return { + 'dq': torch.full_like(q, dq_value), + 'dkv': torch.full_like(kv, dkv_value), + 'd_sink': torch.full_like(attn_sink, d_sink_value), + } + + fake_dsa.sparse_attention_backward_wrapper.side_effect = fake_sparse_attn_backward + + def fake_dense_indexer_grad_backward( + q_idx_bshd, + w_bsh, + k_idx_bsd, + attn_score, + attn_l1norm, + index_score, + index_lse, + sm_scale, + loss_coeff, + grad_loss, + ratio, + block_I, + ): + return { + 'd_index_q': torch.full_like(q_idx_bshd, d_index_q_value), + 'd_weights': torch.full_like(w_bsh, d_weights_value), + 'd_index_k': torch.full_like(k_idx_bsd, d_index_k_value), + } + + fake_dsa.dense_indexer_backward_wrapper.side_effect = fake_dense_indexer_grad_backward + + flash_stub = _make_flash_mla_stub(d_v=d) + + dk._DSA = fake_dsa + dk._flash_mla_sparse_fwd = flash_stub + return fake_dsa, flash_stub + + +class TestFusedIndexerSparseAttn: + """End-to-end numerical tests for the Path B autograd Function with all + underlying CUDA kernels mocked. + """ + + # Common shapes shared across the forward tests. + SHAPES = dict(sq=4, b=2, np_=2, d=512, skv=8, n_comp=4, idx_nh=4, idx_hd=64) + + def _make_inputs(self, *, requires_grad=False): + """Build the seven differentiable + one non-differentiable inputs.""" + s = self.SHAPES + win_topk = _get_topk_alignment() - 2 # exercise padding + torch.manual_seed(0) + + def make(*shape, dtype, rg=False): + t = torch.randn(*shape, dtype=dtype, device='cuda') + if requires_grad and rg: + t = t.detach().clone().requires_grad_(True) + return t + + query = make(s['sq'], s['b'], s['np_'], s['d'], dtype=torch.bfloat16, rg=True) + kv_full = make(s['skv'], s['b'], s['d'], dtype=torch.bfloat16, rg=True) + attn_sink = torch.zeros(s['np_'], dtype=torch.float32, device='cuda') + if requires_grad: + attn_sink = attn_sink.detach().clone().requires_grad_(True) + window_idxs = torch.zeros(s['b'], s['sq'], win_topk, dtype=torch.int32, device='cuda') + q_indexer = make(s['sq'], s['b'], s['idx_nh'], s['idx_hd'], dtype=torch.bfloat16, rg=True) + k_indexer = make(s['n_comp'], s['b'], s['idx_hd'], dtype=torch.bfloat16, rg=True) + weights = make(s['sq'], s['b'], s['idx_nh'], dtype=torch.bfloat16, rg=True) + return dict( + query=query, + kv_full=kv_full, + attn_sink=attn_sink, + window_idxs=window_idxs, + q_indexer=q_indexer, + k_indexer=k_indexer, + weights=weights, + ) + + @pytest.mark.parametrize( + "loss_coeff, target_kind, expected", + [ + # KL(target == predict) == 0 → loss == 0 regardless of coeff. + (1.0, 'uniform', 0.0), + # loss_coeff == 0 short-circuits even when target != predict. + (0.0, 'peaked', 0.0), + # target = δ_0, predict = uniform(1/K) → KL = log(K) per row, + # mean over rows = log(K), scaled by coeff = coeff * log(K). + (0.7, 'peaked', 0.7 * math.log(2)), + # Linearity in loss_coeff: doubling the coeff doubles the loss. + (2.0, 'peaked', 2.0 * math.log(2)), + ], + ids=['identical_dists', 'coeff_zero', 'analytical_kl', 'linearity_x2'], + ) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_indexer_loss_formula(self, loss_coeff, target_kind, expected, reset_lazy_kernel_state): + """All four loss-property cases share one fixture: + + * KL is zero when target == predict, + * ``loss_coeff == 0`` short-circuits to zero, + * for ``target = δ_0`` and ``predict = uniform(1/K)`` the per-row + KL is exactly ``log(K)`` so the mean is ``loss_coeff * log(K)``, + * the loss is linear in ``loss_coeff``. + """ + s = self.SHAPES + topk = 2 # = effective_topk = min(indexer_topk, n_comp); appears as K + target_fn = ( + _uniform_dist + if target_kind == 'uniform' + else (lambda B, S, K, dev: _peaked_dist(B, S, K, dev, peak_idx=0)) + ) + + inputs = self._make_inputs() + _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'], + predict_fn=_uniform_dist, + target_fn=target_fn, + ) + + _, indexer_loss = fused_indexer_sparse_attn( + **inputs, + indexer_topk=topk, + ratio=4, + softmax_scale=0.5, + loss_coeff=loss_coeff, + sparse_loss=True, + kv_offset=s['skv'] - s['n_comp'], + ) + + assert torch.allclose( + 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_path_fwd_output_bwd_grads_and_topk_clamp(self, reset_lazy_kernel_state): + """Combined coverage for the sparse-loss path's three non-numerical + properties (assertion blocks self-label on failure): + + * (a) ``output`` is exactly the FlashMLA stub's ``out`` reshaped + from ``(sq*b, np_, d_v)`` to ``(sq, b, np_ * d_v)``. + * (b) After backward, each leaf gradient equals the corresponding + mocked kernel output, with q/kv/attn_sink coming from the + sparse-attn bwd kernel and q_indexer/k_indexer/weights coming + from the indexer bwd kernel (BSHD → SBHD permute applied). + * (c) ``indexer_topk > n_comp`` is clamped to ``n_comp`` before the + radix TopK kernel is called. + """ + s = self.SHAPES + + # ---- (a) forward pass-through (no grads needed) ------------------ + inputs = self._make_inputs() + _, flash_stub_a = _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'] + ) + output_a, _ = fused_indexer_sparse_attn( + **inputs, + indexer_topk=2, + ratio=4, + softmax_scale=0.5, + indexer_softmax_scale=0.125, + loss_coeff=0.0, + sparse_loss=True, + kv_offset=s['skv'] - s['n_comp'], + ) + assert output_a.shape == (s['sq'], s['b'], s['np_'] * s['d']), "(a) shape" + assert output_a.dtype == torch.bfloat16, "(a) dtype" + expected_a = flash_stub_a.last_out.reshape(s['sq'], s['b'], s['np_'], s['d']).reshape( + s['sq'], s['b'], s['np_'] * s['d'] + ) + assert torch.equal(output_a, expected_a), "(a) forward value pass-through" + + # ---- (b) backward grad propagation ------------------------------- + dk._DSA = None # fresh mocks + dk._flash_mla_sparse_fwd = None + inputs_b = self._make_inputs(requires_grad=True) + _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'], + dq_value=7.0, + dkv_value=-3.0, + d_sink_value=11.0, + d_index_q_value=0.5, + d_weights_value=-0.25, + d_index_k_value=1.5, + ) + output_b, indexer_loss_b = fused_indexer_sparse_attn( + **inputs_b, + indexer_topk=2, + ratio=4, + softmax_scale=0.5, + indexer_softmax_scale=0.125, + loss_coeff=1.0, + sparse_loss=True, + kv_offset=s['skv'] - s['n_comp'], + ) + (output_b.sum() + indexer_loss_b).backward() + for name, value in [ + ('query', 7.0), + ('kv_full', -3.0), + ('attn_sink', 11.0), + ('q_indexer', 0.5), + ('k_indexer', 1.5), + ('weights', -0.25), + ]: + grad = inputs_b[name].grad + assert grad is not None, f"(b) {name}: missing grad" + assert torch.equal(grad, torch.full_like(inputs_b[name], value)), ( + f"(b) {name}: grad does not equal full({value}); " + f"got first elem = {grad.float().flatten()[0].item()}" + ) + + # ---- (c) indexer_topk > n_comp clamp ----------------------------- + dk._DSA = None + dk._flash_mla_sparse_fwd = None + inputs_c = self._make_inputs() + fake_dsa_c, _ = _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'] + ) + fused_indexer_sparse_attn( + **inputs_c, + indexer_topk=999, # > n_comp + ratio=4, + softmax_scale=0.5, + loss_coeff=0.0, + sparse_loss=True, + kv_offset=s['skv'] - s['n_comp'], + ) + topk_call = fake_dsa_c.indexer_top_k_wrapper.call_args + assert ( + topk_call.kwargs['top_k'] == s['n_comp'] + ), f"(c) top_k clamp: got {topk_call.kwargs['top_k']}, expected {s['n_comp']}" + + +# --------------------------------------------------------------------------- +# fused_indexer_sparse_attn — dense path (sparse_loss=False) +# --------------------------------------------------------------------------- + + +class TestDenseFusedIndexerSparseAttn: + """End-to-end tests for the dense-loss branch of Path B with all + underlying CUDA kernels mocked. Mirrors :class:`TestFusedIndexerSparseAttn` + but exercises the ``sparse_loss=False`` code path through + :class:`FusedIndexerSparseAttnFunc`. + """ + + SHAPES = dict(sq=4, b=2, np_=2, d=512, skv=8, n_comp=4, idx_nh=4, idx_hd=64) + + def _make_inputs(self, *, requires_grad=False): + s = self.SHAPES + win_topk = _get_topk_alignment() - 2 # exercise padding + torch.manual_seed(0) + + def make(*shape, dtype, rg=False): + t = torch.randn(*shape, dtype=dtype, device='cuda') + if requires_grad and rg: + t = t.detach().clone().requires_grad_(True) + return t + + query = make(s['sq'], s['b'], s['np_'], s['d'], dtype=torch.bfloat16, rg=True) + kv_full = make(s['skv'], s['b'], s['d'], dtype=torch.bfloat16, rg=True) + attn_sink = torch.zeros(s['np_'], dtype=torch.float32, device='cuda') + if requires_grad: + attn_sink = attn_sink.detach().clone().requires_grad_(True) + window_idxs = torch.zeros(s['b'], s['sq'], win_topk, dtype=torch.int32, device='cuda') + q_indexer = make(s['sq'], s['b'], s['idx_nh'], s['idx_hd'], dtype=torch.bfloat16, rg=True) + k_indexer = make(s['n_comp'], s['b'], s['idx_hd'], dtype=torch.bfloat16, rg=True) + weights = make(s['sq'], s['b'], s['idx_nh'], dtype=torch.bfloat16, rg=True) + return dict( + query=query, + kv_full=kv_full, + attn_sink=attn_sink, + window_idxs=window_idxs, + q_indexer=q_indexer, + k_indexer=k_indexer, + weights=weights, + ) + + @pytest.mark.parametrize( + "loss_coeff, target_kind, expected", + [ + # Identical dists: KL == 0 regardless of coeff. + (1.0, 'uniform', 0.0), + # loss_coeff == 0 short-circuits even when target != predict. + (0.0, 'peaked', 0.0), + # target = δ_0, predict = uniform(1/n_comp) + # per-row KL = log(n_comp); mean = log(n_comp); loss = coeff * log(n_comp). + # n_comp = 4 here. + (0.7, 'peaked', 0.7 * math.log(4)), + (2.0, 'peaked', 2.0 * math.log(4)), + ], + ids=['identical_dists', 'coeff_zero', 'analytical_kl', 'linearity_x2'], + ) + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dense_indexer_loss_formula( + self, loss_coeff, target_kind, expected, reset_lazy_kernel_state + ): + """``_kl_loss_from_dense_scores`` is the dense analogue of + ``_kl_loss_from_target_predict``. Verifies the same four KL + properties (zero, coeff-zero short-circuit, analytical formula, + linearity in coeff) over the dense ``(B, S_q, S_k)`` tensors. + + We drive the stub outputs so that: + + * predict = ``softmax(0)`` over S_k = uniform(1/n_comp). This is + encoded as ``index_score = 0`` everywhere, ``index_lse = log(n_comp)``; + ``predict = exp(score - lse) = 1/n_comp``. + * For ``target = uniform``: attn_score = 1/n_comp uniformly, attn_l1norm = 1. + * For ``target = δ_0``: attn_score peaked on slot 0 with sum 1, attn_l1norm = 1. + """ + s = self.SHAPES + + if target_kind == 'uniform': + target_score_fn = lambda B, S, K, dev: torch.full( + (B, S, K), 1.0 / max(K, 1), dtype=torch.float32, device=dev + ) + else: + + def target_score_fn(B, S, K, dev): + t = torch.zeros((B, S, K), dtype=torch.float32, device=dev) + t[..., 0] = 1.0 + return t + + target_l1norm_fn = lambda B, S, K, dev: torch.ones((B, S), dtype=torch.float32, device=dev) + + inputs = self._make_inputs() + _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'], + target_score_fn=target_score_fn, + target_l1norm_fn=target_l1norm_fn, + # predict_score_fn / predict_lse_fn defaults give uniform predict. + ) + + _, indexer_loss = fused_indexer_sparse_attn( + **inputs, + indexer_topk=2, + ratio=4, + softmax_scale=0.5, + loss_coeff=loss_coeff, + sparse_loss=False, + kv_offset=s['skv'] - s['n_comp'], + ) + + assert torch.allclose( + 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_dense_path_fwd_kernel_calls_and_bwd_grads(self, reset_lazy_kernel_state): + """Combined coverage for the dense-loss path's two non-numerical + properties (assertion blocks self-label on failure): + + * (a) The forward invokes ``dense_attn_score_recompute_wrapper`` + (NOT the sparse score kernels) with the right BSHD/4-D shapes + and ``ratio`` / scale args. The indexer predict is derived + directly from ``indexer_forward_wrapper`` scores (no separate + ``dense_indexer_score_recompute_wrapper`` call). + * (b) The forward eagerly invokes ``dense_indexer_backward_wrapper`` + (NOT the sparse one), threads ``ratio`` through, and the + resulting grads land on the right SBHD leaves (BSHD → SBHD + permute applied for the indexer-side grads, scaled by + ``grad_loss`` in the actual backward). + """ + s = self.SHAPES + ratio = 4 + softmax_scale = 0.5 + idx_scale = 0.125 + loss_coeff = 1.0 + + # ---- (a) forward kernel selection + arg shapes ------------------- + inputs_a = self._make_inputs() + fake_dsa_a, _ = _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'] + ) + fused_indexer_sparse_attn( + **inputs_a, + indexer_topk=2, + ratio=ratio, + softmax_scale=softmax_scale, + indexer_softmax_scale=idx_scale, + loss_coeff=loss_coeff, + sparse_loss=False, + kv_offset=s['skv'] - s['n_comp'], + ) + # Indexer predict is derived from indexer_forward_wrapper scores + # (gather + logsumexp), NOT from dense_indexer_score_recompute_wrapper. + fake_dsa_a.dense_indexer_score_recompute_wrapper.assert_not_called() + fake_dsa_a.dense_attn_score_recompute_wrapper.assert_called_once() + fake_dsa_a.sparse_indexer_score_recompute_wrapper.assert_not_called() + fake_dsa_a.sparse_attn_score_recompute_wrapper.assert_not_called() + + attn_call = fake_dsa_a.dense_attn_score_recompute_wrapper.call_args + q_attn, k_attn, lse_arg, sm_arg = attn_call.args + assert q_attn.shape == (s['b'], s['sq'], s['np_'], s['d']), "(a) dense attn score: q shape" + assert k_attn.shape == ( + s['b'], + s['n_comp'], + 1, + s['d'], + ), "(a) dense attn score: k shape (h_kv=1)" + assert lse_arg.shape == (s['b'], s['sq'], s['np_']), "(a) dense attn score: lse shape" + assert sm_arg == softmax_scale, "(a) dense attn score: positional softmax_scale" + assert attn_call.kwargs['qhead_per_kv_head'] == s['np_'] + assert attn_call.kwargs['ratio'] == ratio + + # ---- (b) forward-eager indexer backward + grad propagation -------- + dk._DSA = None + dk._flash_mla_sparse_fwd = None + inputs_b = self._make_inputs(requires_grad=True) + fake_dsa_b, _ = _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'], + dq_value=7.0, + dkv_value=-3.0, + d_sink_value=11.0, + d_index_q_value=0.5, + d_weights_value=-0.25, + d_index_k_value=1.5, + ) + output, indexer_loss = fused_indexer_sparse_attn( + **inputs_b, + indexer_topk=2, + ratio=ratio, + softmax_scale=softmax_scale, + indexer_softmax_scale=idx_scale, + loss_coeff=loss_coeff, + sparse_loss=False, + kv_offset=s['skv'] - s['n_comp'], + ) + (output.sum() + indexer_loss).backward() + + # dense_indexer_backward_wrapper is called eagerly during forward. + fake_dsa_b.dense_indexer_backward_wrapper.assert_called_once() + fake_dsa_b.indexer_backward_wrapper.assert_not_called() + ig_call = fake_dsa_b.dense_indexer_backward_wrapper.call_args + assert ig_call.kwargs['ratio'] == ratio, "(b) ratio not threaded through" + assert ig_call.kwargs['sm_scale'] == idx_scale, "(b) sm_scale not threaded" + assert ig_call.kwargs['loss_coeff'] == loss_coeff, "(b) loss_coeff not threaded" + + for name, value in [ + ('query', 7.0), + ('kv_full', -3.0), + ('attn_sink', 11.0), + ('q_indexer', 0.5), + ('k_indexer', 1.5), + ('weights', -0.25), + ]: + grad = inputs_b[name].grad + assert grad is not None, f"(b) {name}: missing grad" + assert torch.equal( + grad, torch.full_like(inputs_b[name], value) + ), f"(b) {name}: grad does not equal full({value})" + + +# --------------------------------------------------------------------------- +# Real-kernel parity tests (cuDNN + optional FlashMLA) +# --------------------------------------------------------------------------- +# +# Everything above this banner stubs ``cudnn.DSA`` and ``flash_mla`` with +# ``MagicMock``-based fakes; that exercises the Python plumbing of +# ``dsa_kernels.py`` (shape transforms, autograd wiring, KL composition) +# but does NOT verify that the cuDNN kernels themselves compute what +# ``dsa_kernels.py`` expects them to compute. +# +# The tests below close that gap by running each helper / public function +# end-to-end against a small PyTorch reference implementation. Numeric +# tolerances are bf16-friendly (atol/rtol ~ 5e-2 for raw scores, 1e-3 for +# normalized distributions, 5e-2 for backward grads). +# +# Skipped automatically when: +# * CUDA is unavailable; +# * cuDNN frontend is not installed (``import cudnn`` fails); +# * ``cudnn.DSA`` namespace is missing; +# * SM is too low (sparse: SM90+; dense: SM100+); +# * for FlashMLA-needing tests, ``flash_mla`` is not installed. +# --------------------------------------------------------------------------- + + +def _skip_if_real_kernels_unavailable(*, sm_min: int = 9, need_flash_mla: bool = False): + """Pytest-side gate for real-kernel tests. Raises ``pytest.skip`` if + any of the runtime dependencies are missing. + """ + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + sm_major = torch.cuda.get_device_capability()[0] + if sm_major < sm_min: + pytest.skip(f"requires SM{sm_min}+, found SM{sm_major}") + cudnn = pytest.importorskip("cudnn") + cudnn_frontend = pytest.importorskip("cudnn_frontend") + from packaging.version import Version + + if Version(cudnn_frontend.__version__) < Version("1.24.0"): + pytest.skip(f"requires cudnn_frontend>=1.24.0, found {cudnn_frontend.__version__}") + if not hasattr(cudnn, 'DSA'): + pytest.skip("cudnn.DSA namespace not available") + if need_flash_mla: + pytest.importorskip("flash_mla") + + +# --------------------------------------------------------------------------- +# PyTorch reference implementations +# --------------------------------------------------------------------------- + + +def _ratio_causal_valid_mask(sq: int, sk: int, ratio: int, device) -> torch.Tensor: + """``(Sq, Sk)`` bool: valid iff ``k_idx < min(Sk, (q_idx + 1) // ratio)``. + + Matches the cuDNN dense-score kernels' built-in causal mask + (``col_limit = min(S_k, (q + 1) // ratio)``) and ``csa.py``'s + ``compress_ratio`` mask formulation. + """ + q_idx = torch.arange(sq, device=device).unsqueeze(1) # (Sq, 1) + k_idx = torch.arange(sk, device=device).unsqueeze(0) # (1, Sk) + col_limit = ((q_idx + 1) // ratio).clamp(max=sk) + return k_idx < col_limit # (Sq, Sk) + + +def _ref_indexer_full_score( + q_bshd_fp32: torch.Tensor, # (B, Sq, H, D) + k_bsd_fp32: torch.Tensor, # (B, Sk, D) — MQA + w_bsh_fp32: torch.Tensor, # (B, Sq, H) + sm_scale: float, + ratio: int, +) -> torch.Tensor: + """Reference for ``_bwd_dense_indexer_score.out``. + + ``S[b,q,k] = sm_scale * sum_h ReLU(Q[b,q,h] @ K[b,k]^T) * W[b,q,h]``, + with the kernel's bottom-right ratio causal mask producing ``-inf`` + at masked positions. + """ + B, Sq, _, _ = q_bshd_fp32.shape + Sk = k_bsd_fp32.shape[1] + qk = torch.einsum('bqhd,bkd->bqhk', q_bshd_fp32, k_bsd_fp32) # (B, Sq, H, Sk) + relu_qk = torch.relu(qk) + s = (relu_qk * w_bsh_fp32.unsqueeze(-1)).sum(dim=2) * sm_scale # (B, Sq, Sk) + valid = _ratio_causal_valid_mask(Sq, Sk, ratio, s.device).unsqueeze(0) + return torch.where(valid, s, torch.full_like(s, float('-inf'))) + + +def _ref_attn_full_score( + q_bshd_fp32: torch.Tensor, # (B, Sq, H, D) + k_bsd_fp32: torch.Tensor, # (B, Sk, D) — MQA + lse_bshq_fp32: torch.Tensor, # (B, Sq, H) + softmax_scale: float, + ratio: int, +) -> torch.Tensor: + """Reference for ``_bwd_dense_attn_score.out``. + + ``out[b,q,k] = sum_h exp(Q[b,q,h] @ K[b,k]^T * scale - LSE[b,q,h])``, + with the ratio causal mask producing ``0`` at masked positions + (the per-head ``exp`` is zeroed out, contributing nothing to the sum). + """ + B, Sq, _, _ = q_bshd_fp32.shape + Sk = k_bsd_fp32.shape[1] + qk = torch.einsum('bqhd,bkd->bqhk', q_bshd_fp32, k_bsd_fp32) * softmax_scale + p = torch.exp(qk - lse_bshq_fp32.unsqueeze(-1)) + s = p.sum(dim=2) # (B, Sq, Sk) + valid = _ratio_causal_valid_mask(Sq, Sk, ratio, s.device).unsqueeze(0) + return torch.where(valid, s, torch.zeros_like(s)) + + +def _ref_indexer_predict_sparse(q_bshd_fp32, k_bsd_fp32, w_bsh_fp32, topk_indices, sm_scale): + """Reference for ``sparse_indexer_score_recompute_wrapper.predict``. + + Compute the full-KV indexer score, gather ``topk_indices``, softmax + over the topK axis. ``-1`` entries in topk are masked to ``-inf`` + so they contribute zero probability. + """ + qk = torch.einsum('bqhd,bkd->bqhk', q_bshd_fp32, k_bsd_fp32) + s = (torch.relu(qk) * w_bsh_fp32.unsqueeze(-1)).sum(dim=2) * sm_scale # (B, Sq, Sk) + valid = topk_indices >= 0 + safe = topk_indices.clamp(min=0).long() + s_topk = torch.gather(s, dim=-1, index=safe) + s_topk = torch.where(valid, s_topk, torch.full_like(s_topk, float('-inf'))) + return torch.softmax(s_topk, dim=-1) + + +def _ref_attn_target_sparse(q_bshd_fp32, k_bsd_fp32, lse_bsh_fp32, topk_indices, softmax_scale): + """Reference for ``sparse_attn_score_recompute_wrapper.target``. + + Per-head ``exp(QK*scale - LSE)``, sum over heads, gather topK, + L1-normalise over the topK axis. ``-1`` entries are zero-masked + pre-normalisation. + """ + qk = torch.einsum('bqhd,bkd->bqhk', q_bshd_fp32, k_bsd_fp32) * softmax_scale + p = torch.exp(qk - lse_bsh_fp32.unsqueeze(-1)) # (B, Sq, H, Sk) + s = p.sum(dim=2) # (B, Sq, Sk) + valid = topk_indices >= 0 + safe = topk_indices.clamp(min=0).long() + s_topk = torch.gather(s, dim=-1, index=safe) + s_topk = torch.where(valid, s_topk, torch.zeros_like(s_topk)) + denom = s_topk.sum(dim=-1, keepdim=True).clamp(min=1e-10) + return s_topk / denom + + +def _ref_dense_indexer_loss( + q_indexer_bshd_fp32, + k_indexer_bsd_fp32, + w_bsh_fp32, + q_attn_bshd_fp32, + k_attn_bsd_fp32, + lse_bshq_fp32, + indexer_softmax_scale: float, + attn_softmax_scale: float, + ratio: int, + loss_coeff: float, +) -> torch.Tensor: + """Reference dense KL loss (matches ``compute_dsa_indexer_loss(sparse_loss=False)`` + in ``dsa.py``). Uses the same ratio causal mask the kernel applies. + """ + eps = torch.finfo(torch.float32).tiny + # Per-(b,q,k) raw scores via the same formulas the kernels use. + attn_scores = _ref_attn_full_score( + q_attn_bshd_fp32, k_attn_bsd_fp32, lse_bshq_fp32, attn_softmax_scale, ratio + ) # (B, Sq, Sk) head-summed, ratio-masked, zeros at masked positions. + index_scores = _ref_indexer_full_score( + q_indexer_bshd_fp32, k_indexer_bsd_fp32, w_bsh_fp32, indexer_softmax_scale, ratio + ) # (B, Sq, Sk) ReLU·W, ratio-masked, -inf at masked positions. + + # L1-norm denom for target; LSE for predict. + attn_denom = attn_scores.sum(dim=-1) # (B, Sq) + index_lse = torch.logsumexp(index_scores, dim=-1) # (B, Sq), -inf for fully-masked rows + + row_valid = (attn_denom > eps) & torch.isfinite(index_lse) + + safe_l1 = attn_denom.clamp(min=eps) + safe_lse = torch.where(row_valid, index_lse, torch.zeros_like(index_lse)) + + target = attn_scores / safe_l1.unsqueeze(-1) + target_clamped = target.clamp(min=eps) + # Mask within-row: ratio-causal-masked positions have ``index_scores = + # -inf`` (from ``_ref_indexer_full_score``). Letting them flow into + # ``log_predict`` would make per-position contributions blow up to + # +inf (``target_clamped * (log(target) - (-inf)) = +inf``). They have + # zero mass under ``target`` (``_ref_attn_full_score`` zeros those + # positions) so their KL contribution should be 0; explicitly mask. + position_valid = torch.isfinite(index_scores) + log_predict = torch.where( + position_valid, index_scores - safe_lse.unsqueeze(-1), torch.zeros_like(index_scores) + ) + contributions = target_clamped * (torch.log(target_clamped) - log_predict) + contributions = torch.where(position_valid, contributions, torch.zeros_like(contributions)) + kl_per_row = contributions.sum(dim=-1) + kl_per_row = torch.where(row_valid, kl_per_row, torch.zeros_like(kl_per_row)) + return loss_coeff * kl_per_row.mean() + + +def _ref_sparse_attn_forward( + q_flat_bf16: torch.Tensor, # (total_Sq, H, D) + kv_flat_bf16: torch.Tensor, # (total_Skv, D) — K=V, MQA + attn_sink_fp32: torch.Tensor, # (H,) + topk_idxs: torch.Tensor, # (total_Sq, topk) int32, global + softmax_scale: float, + d_v: int, +): + """Pure-PyTorch reference for FlashMLA sparse-attn-fwd output. + + Mirrors the math FlashMLA implements: + * Scores ``S[i, h, k] = Q[i, h] @ K[topk[i, k]]^T * scale`` for valid ``k``; + * Append a per-head sink logit (``attn_sink``); + * ``softmax`` over the (topk + sink) axis; + * ``out[i, h] = sum_k softmax[i, h, k] * V[topk[i, k]]`` (excluding sink). + + Returns ``(out, lse)`` in the same shapes/dtype as the FlashMLA kernel. + Invalid ``-1`` topk entries contribute zero to the softmax (logit -inf). + """ + total_Sq, H, D = q_flat_bf16.shape + topk = topk_idxs.shape[-1] + device = q_flat_bf16.device + q_fp32 = q_flat_bf16.float() + kv_fp32 = kv_flat_bf16.float() + + valid = topk_idxs >= 0 # (total_Sq, topk) + safe = topk_idxs.clamp(min=0).long() + k_gathered = kv_fp32[safe] # (total_Sq, topk, D) + + qk = torch.einsum('ihd,ikd->ihk', q_fp32, k_gathered) * softmax_scale # (Sq, H, topk) + qk = torch.where(valid.unsqueeze(1).expand(-1, H, -1), qk, torch.full_like(qk, float('-inf'))) + sink = attn_sink_fp32.view(1, H, 1).expand(total_Sq, H, 1) # logit + logits = torch.cat([qk, sink], dim=-1) # (Sq, H, topk + 1) + probs = torch.softmax(logits, dim=-1) # numerically stable + probs_kv = probs[..., :topk] # exclude sink contribution from output + + v_gathered = k_gathered # K = V (MQA, head-broadcast) + out_fp32 = torch.einsum('ihk,ikd->ihd', probs_kv, v_gathered) # (Sq, H, D_v=D) + if d_v != D: + out_fp32 = out_fp32[..., :d_v] + out = out_fp32.to(q_flat_bf16.dtype) + + # FlashMLA's KV-only LSE excludes the sink term: + # lse_kv[i, h] = logsumexp_k(qk[i, h, k]) over valid k only. + lse_kv = torch.logsumexp(qk, dim=-1) # (Sq, H), -inf for fully-masked rows + return out, lse_kv + + +# --------------------------------------------------------------------------- +# Score-helper parity tests (sparse + dense): real cuDNN vs PyTorch reference +# --------------------------------------------------------------------------- + + +# Shared small shape across all real-kernel tests to maximize cuDNN compile-cache +# hits. ``ratio=1`` (standard upper-triangular causal) keeps the math simple +# and ensures every row has at least one valid KV position. +_REAL_SHAPES_SPARSE = dict( + b=2, + sq=128, + sk=128, + n_comp=128, + np_=32, + d=512, + idx_nh=32, + idx_hd=128, + # topk = lcm(64, 128) = 128 satisfies SparseScoreRecomputeSm100's + # `topk % n_block_size == 0` (64 for score_type=attention, 128 for indexer). + topk=128, + ratio=1, + softmax_scale=512**-0.5, + indexer_softmax_scale=128**-0.5, +) + + +def _build_real_score_inputs(s, *, with_lse: bool = True, with_topk: bool = True): + """Build a coherent set of bf16 BSHD inputs for the score-helper tests. + + Returns a dict with both bf16 (kernel-ready) and fp32 (reference-math) + views of every tensor, plus optional LSE / topk_indices. + """ + torch.manual_seed(0) + dev = 'cuda' + + q_idx = torch.randn(s['b'], s['sq'], s['idx_nh'], s['idx_hd'], dtype=torch.bfloat16, device=dev) + k_idx = torch.randn(s['b'], s['sk'], s['idx_hd'], dtype=torch.bfloat16, device=dev) + w = torch.randn(s['b'], s['sq'], s['idx_nh'], dtype=torch.bfloat16, device=dev) + + q_attn = torch.randn(s['b'], s['sq'], s['np_'], s['d'], dtype=torch.bfloat16, device=dev) + k_attn = torch.randn(s['b'], s['sk'], s['d'], dtype=torch.bfloat16, device=dev) + + out = dict(q_idx=q_idx, k_idx=k_idx, w=w, q_attn=q_attn, k_attn=k_attn) + + if with_lse: + # LSE = logsumexp(QK*scale, dim=Sk) with the kernel's ratio mask. + # Real LSE input avoids exp(-inf - finite) underflow during reference. + qk = torch.einsum('bqhd,bkd->bqhk', q_attn.float(), k_attn.float()) * s['softmax_scale'] + valid = _ratio_causal_valid_mask(s['sq'], s['sk'], s['ratio'], qk.device).view( + 1, s['sq'], 1, s['sk'] + ) + qk_masked = torch.where(valid, qk, torch.full_like(qk, float('-inf'))) + out['lse'] = torch.logsumexp(qk_masked, dim=-1).clamp(min=-1e30).contiguous() + + if with_topk: + # Pick distinct random valid indices per (b, sq), with a few -1s + # interleaved to exercise the invalid-slot path. + topk = s['topk'] + torch.manual_seed(123) + idxs = torch.randint(0, s['sk'], (s['b'], s['sq'], topk), dtype=torch.int32, device=dev) + # Mark a few slots invalid (-1) to test the topk_indices < 0 path. + invalid = torch.rand(s['b'], s['sq'], topk, device=dev) < 0.1 + idxs = torch.where(invalid, torch.full_like(idxs, -1), idxs) + out['topk'] = idxs + + return out + + +class TestRealKernelScoreHelpers: + """Real-kernel parity tests for the four ``_compute_*`` score helpers + against PyTorch reference implementations. Single parametrized test + covers all four; numeric tolerance is bf16-friendly (raw fp32 score + sums agree to ~5%, normalized distributions to ~5e-3). + """ + + # Each case: (id, sm_min, kernel_name, runner). The runner does the + # call + ref + assertion; it returns nothing on success. + @pytest.mark.parametrize( + "case", + ['sparse_indexer_predict', 'sparse_attn_target', 'dense_indexer_score', 'dense_attn_score'], + ) + def test_real_score_helper(self, case, reset_lazy_kernel_state): + _skip_if_real_kernels_unavailable(sm_min=10) + + s = _REAL_SHAPES_SPARSE + # Each case needs a different combination of the input fixture. + x = _build_real_score_inputs( + s, + with_lse=case.endswith('_attn_target') or case.endswith('_attn_score'), + with_topk=case.startswith('sparse_'), + ) + + from megatron.core.transformer.experimental_attention_variant import dsa_kernels as _dk + + if case == 'sparse_indexer_predict': + # The kernel takes sm_scale=1.0; scale is applied via weights + # pre-multiplication (relu(c·x)·W trick). Reference mirrors that. + scale = s['indexer_softmax_scale'] + w_scaled = (x['w'].float() * scale).to(x['w'].dtype) + out = _dk._compute_indexer_predict( + x['q_idx'], x['k_idx'], w_scaled, x['topk'], qhead_per_kv_head=s['idx_nh'] + ) + ref = _ref_indexer_predict_sparse( + x['q_idx'].float(), x['k_idx'].float(), w_scaled.float(), x['topk'], sm_scale=1.0 + ) + # Softmax outputs in [0, 1]; bf16 element-wise noise can break + # absolute tolerance, so compare directions via cosine similarity. + assert out.shape == ref.shape == (s['b'], s['sq'], s['topk']) + cos = torch.nn.functional.cosine_similarity( + out.flatten().unsqueeze(0).float(), ref.flatten().unsqueeze(0).float() + ).item() + assert cos > 0.99, ( + f"{case}: cos sim = {cos:.4f}, " + f"max abs diff = {(out - ref).abs().max().item():.3e}" + ) + + elif case == 'sparse_attn_target': + out = _dk._compute_attn_target( + x['q_attn'], + x['k_attn'], + x['lse'], + x['topk'], + softmax_scale=s['softmax_scale'], + qhead_per_kv_head=s['np_'], + ) + ref = _ref_attn_target_sparse( + x['q_attn'].float(), + x['k_attn'].float(), + x['lse'], + x['topk'], + softmax_scale=s['softmax_scale'], + ) + assert out.shape == ref.shape == (s['b'], s['sq'], s['topk']) + cos = torch.nn.functional.cosine_similarity( + out.flatten().unsqueeze(0).float(), ref.flatten().unsqueeze(0).float() + ).item() + assert cos > 0.99, ( + f"{case}: cos sim = {cos:.4f}, " + f"max abs diff = {(out - ref).abs().max().item():.3e}" + ) + + elif case == 'dense_indexer_score': + out, denom = _dk._compute_dense_indexer_score( + x['q_idx'], + x['k_idx'].unsqueeze(2), + x['w'], + qhead_per_kv_head=s['idx_nh'], + indexer_softmax_scale=s['indexer_softmax_scale'], + ratio=s['ratio'], + ) + ref_out = _ref_indexer_full_score( + x['q_idx'].float(), + x['k_idx'].float(), + x['w'].float(), + sm_scale=s['indexer_softmax_scale'], + ratio=s['ratio'], + ) + ref_denom = torch.logsumexp(ref_out, dim=-1) + assert out.shape == ref_out.shape == (s['b'], s['sq'], s['sk']) + assert denom.shape == ref_denom.shape == (s['b'], s['sq']) + # Raw fp32 score sums: relative tolerance dominates. Compare + # only valid positions (masked = -inf in both, NaN under sub). + valid = ( + _ratio_causal_valid_mask(s['sq'], s['sk'], s['ratio'], out.device) + .unsqueeze(0) + .expand_as(out) + ) + diff = torch.where(valid, (out - ref_out).abs(), torch.zeros_like(out)) + scale = ref_out.where(valid, torch.zeros_like(ref_out)).abs().max().item() + assert diff.max().item() <= max( + 5e-2, 5e-2 * scale + ), f"{case}: max abs diff = {diff.max().item():.3e}, scale = {scale:.3e}" + row_valid = torch.isfinite(ref_denom) + assert torch.allclose( + denom[row_valid], ref_denom[row_valid], atol=5e-3, rtol=5e-2 + ), f"{case}: LSE max abs diff = {(denom - ref_denom)[row_valid].abs().max().item():.3e}" + + elif case == 'dense_attn_score': + out, denom = _dk._compute_dense_attn_score( + x['q_attn'], + x['k_attn'].unsqueeze(2), + x['lse'], + qhead_per_kv_head=s['np_'], + softmax_scale=s['softmax_scale'], + ratio=s['ratio'], + ) + ref_out = _ref_attn_full_score( + x['q_attn'].float(), + x['k_attn'].float(), + x['lse'], + softmax_scale=s['softmax_scale'], + ratio=s['ratio'], + ) + ref_denom = ref_out.sum(dim=-1) + assert out.shape == ref_out.shape == (s['b'], s['sq'], s['sk']) + assert denom.shape == ref_denom.shape == (s['b'], s['sq']) + valid = ( + _ratio_causal_valid_mask(s['sq'], s['sk'], s['ratio'], out.device) + .unsqueeze(0) + .expand_as(out) + ) + diff = torch.where(valid, (out - ref_out).abs(), torch.zeros_like(out)) + # exp(QK*scale - LSE) outputs in (0, ~1]: absolute dominates. + assert diff.max().item() <= 5e-3, f"{case}: max abs diff = {diff.max().item():.3e}" + assert torch.allclose(denom, ref_denom, atol=5e-3, rtol=5e-2), ( + f"{case}: denom max abs diff = " f"{(denom - ref_denom).abs().max().item():.3e}" + ) + + else: + raise AssertionError(f"unknown case: {case}") + + +# --------------------------------------------------------------------------- +# KL loss reference parity (dense path; sparse already CPU-tested above). +# --------------------------------------------------------------------------- + + +class TestRealKernelKLLossDense: + """End-to-end parity for ``_kl_loss_from_dense_scores``: run the real + cuDNN dense score kernels, feed their outputs into the helper, and + compare the KL value to the all-PyTorch reference. + """ + + @pytest.mark.parametrize("dummy", [None]) + def test_real_dense_kl_loss_matches_reference(self, dummy, reset_lazy_kernel_state): + _skip_if_real_kernels_unavailable(sm_min=10) + from megatron.core.transformer.experimental_attention_variant.dsa_kernels import ( + _compute_dense_attn_score, + _compute_dense_indexer_score, + _kl_loss_from_dense_scores, + ) + + s = _REAL_SHAPES_SPARSE + x = _build_real_score_inputs(s, with_lse=True, with_topk=False) + loss_coeff = 0.5 + + index_score, index_lse = _compute_dense_indexer_score( + x['q_idx'], + x['k_idx'].unsqueeze(2), + x['w'], + qhead_per_kv_head=s['idx_nh'], + indexer_softmax_scale=s['indexer_softmax_scale'], + ratio=s['ratio'], + ) + attn_score, attn_l1norm = _compute_dense_attn_score( + x['q_attn'], + x['k_attn'].unsqueeze(2), + x['lse'], + qhead_per_kv_head=s['np_'], + softmax_scale=s['softmax_scale'], + ratio=s['ratio'], + ) + + loss_actual = _kl_loss_from_dense_scores( + attn_score, attn_l1norm, index_score, index_lse, loss_coeff + ) + loss_ref = _ref_dense_indexer_loss( + x['q_idx'].float(), + x['k_idx'].float(), + x['w'].float(), + x['q_attn'].float(), + x['k_attn'].float(), + x['lse'], + indexer_softmax_scale=s['indexer_softmax_scale'], + attn_softmax_scale=s['softmax_scale'], + ratio=s['ratio'], + loss_coeff=loss_coeff, + ) + assert torch.allclose(loss_actual, loss_ref, atol=1e-3, rtol=1e-2), ( + f"actual = {loss_actual.item():.6f}, ref = {loss_ref.item():.6f}, " + f"abs diff = {(loss_actual - loss_ref).abs().item():.3e}" + ) + + +# --------------------------------------------------------------------------- +# Real ``indexer_topk``: the top-K set should match the reference ranking. +# --------------------------------------------------------------------------- + + +class TestRealKernelIndexerTopk: + """Real-kernel parity for :func:`indexer_topk`: the SET of selected + top-K indices must match a PyTorch reference ranking. We compare sets + rather than ordered lists because BF16 ties may be broken differently. + """ + + @pytest.mark.parametrize("dummy", [None]) + def test_real_indexer_topk_set_matches_reference(self, dummy, reset_lazy_kernel_state): + _skip_if_real_kernels_unavailable(sm_min=10) # IndexerForward is SM100+ + from megatron.core.transformer.experimental_attention_variant.dsa_kernels import ( + indexer_topk, + ) + + # IndexerForward requires idx_hd=128 and qhpkv in (32, 64). Use an + # SBHD shape that matches what csa.py produces (tensors are SBHD, + # ratio is the indexer's compression ratio). b=2 exercises the + # batch-aware ``seq_lens.repeat(b)`` and the ``(b*sq, sk) → (b, sq, + # topk)`` reshape inside ``_indexer_topk_bshd``. + s = dict( + b=2, + sq=128, + idx_nh=32, + idx_hd=128, + sk=128, + indexer_topk=8, + ratio=4, + indexer_softmax_scale=128**-0.5, + ) + torch.manual_seed(0) + dev = 'cuda' + q_indexer = torch.randn( + s['sq'], s['b'], s['idx_nh'], s['idx_hd'], dtype=torch.bfloat16, device=dev + ) + k_indexer = torch.randn(s['sk'], 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) + + topk_indices, topk_length = indexer_topk( + q_indexer, + k_indexer, + weights, + topk=s['indexer_topk'], + ratio=s['ratio'], + indexer_softmax_scale=s['indexer_softmax_scale'], + ) + assert topk_indices.shape == (s['b'], s['sq'], s['indexer_topk']) + assert topk_indices.dtype == torch.int32 + + # Reference: full indexer score, ratio causal mask, take top-K per row + # by descending score. Score is sm_scale * sum_h ReLU(Q@K) * W. + q_bshd = q_indexer.permute(1, 0, 2, 3).contiguous().float() + k_bsd = k_indexer.permute(1, 0, 2).contiguous().float() + w_bsh = weights.permute(1, 0, 2).contiguous().float() + ref_scores = _ref_indexer_full_score( + q_bshd, k_bsd, w_bsh, sm_scale=s['indexer_softmax_scale'], ratio=s['ratio'] + ) # (B, Sq, Sk), -inf at masked positions + + # For each row, count valid positions (un-masked). topk_length should + # match min(indexer_topk, num_valid). + n_valid = (ref_scores > float('-inf')).sum(dim=-1) # (B, Sq) + expected_length = n_valid.clamp(max=s['indexer_topk']).int() + assert torch.equal(topk_length, expected_length) + + # Set comparison row-by-row. Skip rows with 0 valid (kernel returns + # all -1; reference picks arbitrary -inf positions). + ref_topk = torch.topk(ref_scores, k=s['indexer_topk'], dim=-1).indices + for bi in range(s['b']): + for qi in range(s['sq']): + n = int(expected_length[bi, qi].item()) + if n == 0: + # Kernel must report all -1. + assert torch.all(topk_indices[bi, qi] == -1) + continue + actual_set = set(topk_indices[bi, qi, :n].tolist()) + ref_set = set(ref_topk[bi, qi, :n].tolist()) + # BF16 ties may differ: allow up to ~10% mismatch on small K. + inter = actual_set & ref_set + assert len(inter) >= max(1, n - 1), ( + f"row (b={bi}, q={qi}): " + f"actual {sorted(actual_set)} vs ref {sorted(ref_set)}" + ) + + +# --------------------------------------------------------------------------- +# Real ``dsa_sparse_attn``: forward + backward parity vs PyTorch reference. +# --------------------------------------------------------------------------- + + +class TestRealKernelDsaSparseAttn: + """Real-kernel parity for :func:`dsa_sparse_attn`. Forward uses real + FlashMLA + the SBHD/flat reshape wrapper; backward uses the real cuDNN + sparse-attn-bwd kernel. Both checked in one test against the + pure-PyTorch sparse-attn reference (``_ref_sparse_attn_forward``). + """ + + SHAPES = dict(b=2, sq=128, np_=64, d=512, skv=128, topk=32, softmax_scale=512**-0.5) + + def _make_inputs(self, *, requires_grad: bool): + s = self.SHAPES + torch.manual_seed(0) + dev = 'cuda' + + def make_leaf(*shape, dtype): + t = torch.randn(*shape, dtype=dtype, device=dev) + return t.detach().clone().requires_grad_(True) if requires_grad else t + + query = make_leaf(s['sq'], s['b'], s['np_'], s['d'], dtype=torch.bfloat16) + kv = make_leaf(s['skv'], s['b'], s['d'], dtype=torch.bfloat16) + attn_sink = torch.zeros(s['np_'], dtype=torch.float32, device=dev) + if requires_grad: + attn_sink = attn_sink.detach().clone().requires_grad_(True) + + # Coherent valid global topk indices in SBHD-flat layout, with a + # standard causal mask (index <= q_idx). + torch.manual_seed(1) + topk_local = torch.randint( + 0, s['skv'], (s['b'], s['sq'], s['topk']), dtype=torch.int64, device=dev + ) + q_idx = torch.arange(s['sq'], device=dev).view(1, -1, 1) + topk_local = torch.minimum(topk_local, q_idx) + global_idxs = local_to_global_flat(topk_local, s['b'], s['skv']).contiguous() + return query, kv, attn_sink, global_idxs + + def test_real_dsa_sparse_attn_fwd_bwd_matches_reference(self, reset_lazy_kernel_state): + """Forward output AND backward gradients (dq, dkv, d_sink) must + match a pure-PyTorch sparse-attn reference. Combining both checks + in one test halves cuDNN compile time vs running them separately, + since they share the same kernel cache key. + """ + _skip_if_real_kernels_unavailable(sm_min=10, need_flash_mla=True) + s = self.SHAPES + + # ---- Real path: forward + backward via dsa_sparse_attn ---- + query, kv, attn_sink, global_idxs = self._make_inputs(requires_grad=True) + out = dsa_sparse_attn(query, kv, attn_sink, global_idxs, softmax_scale=s['softmax_scale']) + torch.manual_seed(7) + upstream = torch.randn_like(out) + (out * upstream).sum().backward() + dq_actual = query.grad.float().clone() + dkv_actual = kv.grad.float().clone() + dsink_actual = attn_sink.grad.float().clone() + out_actual = out.float().detach().clone() + + # ---- Reference: pure-PyTorch forward + autograd ---- + query_ref, kv_ref, attn_sink_ref, _ = self._make_inputs(requires_grad=True) + q_flat = query_ref.reshape(s['sq'] * s['b'], s['np_'], s['d']) + kv_flat = kv_ref.reshape(s['skv'] * s['b'], s['d']) + ref_out_flat, _ = _ref_sparse_attn_forward( + q_flat, + kv_flat, + attn_sink_ref, + global_idxs, + softmax_scale=s['softmax_scale'], + d_v=s['d'], + ) + ref_out = ref_out_flat.reshape(s['sq'], s['b'], s['np_'], s['d']).reshape( + s['sq'], s['b'], s['np_'] * s['d'] + ) + (ref_out * upstream).sum().backward() + + # ---- Forward + backward parity (cos sim) ---- + # bf16 GEMM accumulators in FlashMLA fwd / cuDNN sparse-attn-bwd + # make element-wise tolerances brittle (esp. dkv); compare each + # tensor's direction via cosine similarity instead. + def _cos(a, b): + return torch.nn.functional.cosine_similarity( + a.flatten().unsqueeze(0).float(), b.flatten().unsqueeze(0).float() + ).item() + + assert out_actual.shape == ref_out.shape + for name, actual, ref in [ + ('forward', out_actual, ref_out.float()), + ('dq', dq_actual, query_ref.grad.float()), + ('dkv', dkv_actual, kv_ref.grad.float()), + ('d_sink', dsink_actual, attn_sink_ref.grad.float()), + ]: + cos = _cos(actual, ref) + assert cos > 0.99, ( + f"{name}: cos sim = {cos:.4f}, " + f"max abs diff = {(actual - ref).abs().max().item():.3e}" + ) + + +# --------------------------------------------------------------------------- +# Real ``fused_indexer_sparse_attn``: dense-loss path end-to-end parity. +# --------------------------------------------------------------------------- + + +class TestRealKernelFusedIndexerSparseAttn: + """End-to-end parity for the dense loss path of + :func:`fused_indexer_sparse_attn`: real cuDNN dense kernels (forward + + backward) + real FlashMLA, compared to ``_ref_dense_indexer_loss``. + + Backward grad correctness for the indexer-grad kernel is established by + ``TestRealKernelKLLossDense`` (kernel-level math) and + ``TestDenseFusedIndexerSparseAttn::test_dense_backward_calls_dense_indexer_grad`` + (mock-based plumbing). This class only checks the loss SCALAR value. + """ + + # FlashMLA only accepts indexer_topk ∈ {0, 512, 1024, 2048} and a limited + # set of h_q values (np_=64 is the supported one used by the sibling + # DsaSparseAttn real-kernel test). n_comp must be ≥ indexer_topk; skv ≥ + # n_comp so kv_offset = skv - n_comp > 0 still exercises the offset path. + SHAPES = dict( + b=2, + sq=128, + np_=64, + d=512, + skv=640, + n_comp=512, + idx_nh=32, + idx_hd=128, + indexer_topk=512, + ratio=4, + win_topk=8, + softmax_scale=512**-0.5, + indexer_softmax_scale=128**-0.5, + ) + + def test_real_fused_dense_loss_matches_reference(self, reset_lazy_kernel_state): + """Real dense path's KL loss value matches the all-PyTorch reference + on the same inputs. The reference uses an analytical + ``logsumexp(QK*scale, ratio mask)`` for ``lse_indexer`` (FlashMLA + emits its own internal lse_indexer that differs slightly), so the + tolerance is wider than for the kernel-only ``KLLossDense`` test. + """ + _skip_if_real_kernels_unavailable(sm_min=10, need_flash_mla=True) + s = self.SHAPES + torch.manual_seed(0) + dev = 'cuda' + loss_coeff = 0.5 + + # Build inputs once; share between actual and reference. + 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) + torch.manual_seed(1) + win_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'] + + # Real path. + _, indexer_loss = fused_indexer_sparse_attn( + query, + kv_full, + attn_sink, + win_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=loss_coeff, + sparse_loss=False, + kv_offset=kv_offset, + ) + + # Reference: SBHD->BSHD once, build analytical lse_ref, compute KL. + q_idx_bshd = q_indexer.permute(1, 0, 2, 3).contiguous().float() + k_idx_bsd = k_indexer.permute(1, 0, 2).contiguous().float() + w_bsh = weights.permute(1, 0, 2).contiguous().float() + q_attn_bshd = query.permute(1, 0, 2, 3).contiguous().float() + k_attn_bsd = kv_full[kv_offset:].permute(1, 0, 2).contiguous().float() + + # PyTorch reference that mirrors the fused path's dense-loss math + # exactly. Two non-obvious requirements: + # * Use FlashMLA's emitted ``lse_indexer`` (logsumexp over the + # indexer-selected top-K positions, with the per-head sink term), + # not an analytical full-KV logsumexp. Otherwise the per-row LSE + # basis differs from the kernel by ~50x. + # * Do NOT apply the ratio-causal mask in the reference scores — + # the dense-score-recompute kernels emit values at every position + # (no internal masking). Masking the reference would shift the + # ``attn_score / attn_l1norm`` normalization and the indexer LSE + # basis, producing a different KL than the kernel's. + from megatron.core.transformer.experimental_attention_variant.dsa_kernels import ( + _dsa_fwd_flash_mla, + _indexer_topk_bshd, + _kl_loss_from_dense_scores, + _sbhd_to_bshd_indexer_inputs, + ) + + # Run indexer + FlashMLA to capture the same ``lse_indexer`` the fused + # path consumes internally. + effective_topk = min(s['indexer_topk'], s['n_comp']) + q_idx_bshd_bf, k_idx_bsd_bf, _, w_bsh_scaled_bf = _sbhd_to_bshd_indexer_inputs( + q_indexer, k_indexer, weights, s['indexer_softmax_scale'] + ) + topk_indices_cmp, _ = _indexer_topk_bshd( + 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) + combined_local = torch.cat([compress_topk_idxs, win_idxs], dim=-1) + global_idxs = local_to_global_flat(combined_local, s['b'], s['skv']) + q_flat = query.reshape(s['sq'] * s['b'], s['np_'], s['d']) + kv_flat = kv_full.reshape(s['skv'] * s['b'], s['d']) + _, _, lse_indexer = _dsa_fwd_flash_mla( + q_flat, + kv_flat, + global_idxs, + s['softmax_scale'], + attn_sink=attn_sink, + topk_length=None, + indexer_topk=effective_topk, + ) + lse_indexer_bsqh = lse_indexer.reshape(s['sq'], s['b'], s['np_']).permute(1, 0, 2) + + # Attention path: exp(QK*scale - lse_indexer), head-summed. No mask. + qk_attn = torch.einsum('bqhd,bkd->bqhk', q_attn_bshd, k_attn_bsd) * s['softmax_scale'] + attn_score_ref = torch.exp(qk_attn - lse_indexer_bsqh.unsqueeze(-1)).sum(dim=2) + attn_l1norm_ref = attn_score_ref.sum(dim=-1) + + # Indexer path: ReLU(QK_indexer) * W head-summed. The fused path + # calls ``_compute_dense_indexer_score`` with ``w_bsh_scaled`` (already + # multiplied by ``indexer_softmax_scale``) AND passes + # ``indexer_softmax_scale`` again as the kernel's ``sm_scale``, + # double-applying the factor (apparent bug in + # ``fused_indexer_sparse_attn`` at ``dsa_kernels.py:800-807``). Mirror + # that here so the reference matches the fused-path output; revisit + # if the upstream pre-scale + kernel-scale duplication is fixed. + qk_idx = torch.einsum('bqhd,bkd->bqhk', q_idx_bshd, k_idx_bsd) + idx_score_ref = (torch.relu(qk_idx) * w_bsh.unsqueeze(-1)).sum(dim=2) * ( + s['indexer_softmax_scale'] ** 2 + ) + idx_lse_ref = torch.logsumexp(idx_score_ref, dim=-1) + + loss_ref = _kl_loss_from_dense_scores( + attn_score_ref, attn_l1norm_ref, idx_score_ref, idx_lse_ref, loss_coeff + ) + assert torch.allclose(indexer_loss, loss_ref, atol=5e-2, rtol=1e-1), ( + f"actual = {indexer_loss.item():.6f}, ref = {loss_ref.item():.6f}, " + f"abs diff = {(indexer_loss - loss_ref).abs().item():.3e}" + ) + + +# --------------------------------------------------------------------------- +# Public surface +# --------------------------------------------------------------------------- + + +class TestPublicApi: + """The ``__all__`` list documents the public surface; verify that every + advertised symbol is importable, that the public free functions are + callable, and that the autograd Functions inherit from the right base. + """ + + def test_public_surface(self): + from megatron.core.transformer.experimental_attention_variant import dsa_kernels + + for name in dsa_kernels.__all__: + assert hasattr(dsa_kernels, name), f"__all__ lists {name!r} but it is missing" + + for fn in ( + build_flat_topk_idxs, + local_to_global_flat, + dsa_sparse_attn, + indexer_topk, + fused_indexer_sparse_attn, + ): + assert callable(fn) + + assert issubclass(SparseAttnFunc, torch.autograd.Function) + assert issubclass(FusedIndexerSparseAttnFunc, torch.autograd.Function) 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 new file mode 100644 index 00000000000..a45a00effc1 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_native_parity.py @@ -0,0 +1,799 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import gc +import math + +import pytest +import torch +import torch.nn as nn +import torch.nn.functional as F + +from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider +from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_dsv4_hybrid_module_spec_for_backend, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.experimental_attention_variant.dsa import DSAIndexerLossAutoScaler +from megatron.core.transformer.spec_utils import build_module +from megatron.core.transformer.transformer_config import MLATransformerConfig +from megatron.core.utils import init_method_normal, scaled_init_method_normal +from tests.unit_tests.test_utilities import Utils + +_SEED = 1234 +_FUSED_SIMILARITY_EPS = 1.5e-4 +_UNFUSED_SIMILARITY_EPS = 2.3e-5 + + +@torch.compile +def _native_q_rms_norm(query: torch.Tensor, eps: float) -> torch.Tensor: + return query * torch.rsqrt(query.square().mean(-1, keepdim=True) + eps) + + +_DSV4_VARIANTS = { + "flash": { + "hidden_size": 4096, + "num_attention_heads": 64, + "q_lora_rank": 1024, + "v_head_dim": 512, + "qk_pos_emb_head_dim": 64, + "o_groups": 8, + "o_lora_rank": 1024, + "csa_compress_rotary_base": 40000, + "dsa_indexer_topk": 512, + }, + "pro": { + "hidden_size": 7168, + "num_attention_heads": 128, + "q_lora_rank": 1536, + "v_head_dim": 512, + "qk_pos_emb_head_dim": 64, + "o_groups": 16, + "o_lora_rank": 1024, + "csa_compress_rotary_base": 160000, + "dsa_indexer_topk": 1024, + }, +} + +_DSA_BACKENDS = [ + pytest.param("fused", True, id="fused"), + pytest.param("unfused", False, id="unfused"), +] + +_CASE_SEQLENS = [2048, 4096, 8192] + + +def _make_config( + variant: str, + compress_ratio: int, + apply_dsa_kernel_fusion: bool = False, + calculate_per_token_loss: bool = False, +) -> MLATransformerConfig: + shape = _DSV4_VARIANTS[variant] + mcore_ratio = 0 if compress_ratio == 1 else compress_ratio + qk_head_dim = shape["v_head_dim"] - shape["qk_pos_emb_head_dim"] + config = MLATransformerConfig( + multi_latent_attention=True, + experimental_attention_variant="dsv4_hybrid", + num_layers=1, + hidden_size=shape["hidden_size"], + num_attention_heads=shape["num_attention_heads"], + q_lora_rank=shape["q_lora_rank"], + kv_lora_rank=qk_head_dim, + qk_head_dim=qk_head_dim, + qk_pos_emb_head_dim=shape["qk_pos_emb_head_dim"], + v_head_dim=shape["v_head_dim"], + o_groups=shape["o_groups"], + o_lora_rank=shape["o_lora_rank"], + csa_compress_ratios=[mcore_ratio], + csa_window_size=128, + csa_dense_mode=False, + dsa_indexer_n_heads=64, + dsa_indexer_head_dim=128, + dsa_indexer_topk=shape["dsa_indexer_topk"], + dsa_indexer_loss_coeff=0.01, + dsa_indexer_use_sparse_loss=True, + calculate_per_token_loss=calculate_per_token_loss, + add_bias_linear=False, + bf16=True, + params_dtype=torch.bfloat16, + layernorm_epsilon=1e-6, + normalization="RMSNorm", + qk_layernorm=True, + layernorm_zero_centered_gamma=False, + expert_model_parallel_size=1, + tensor_model_parallel_size=1, + sequence_parallel=False, + context_parallel_size=1, + apply_rope_fusion=False, + rope_type="rope", + rotary_base=10000, + rotary_percent=1.0, + csa_compress_rotary_base=shape["csa_compress_rotary_base"], + recompute_granularity=None, + recompute_modules=[], + fine_grained_activation_offloading=False, + gradient_accumulation_fusion=False, + fp8=False, + fp4=False, + init_method=init_method_normal(0.02), + output_layer_init_method=scaled_init_method_normal(0.02, 1, multiplier=2.0), + kv_channels=shape["v_head_dim"], + num_query_groups=shape["num_attention_heads"], + batch_invariant_mode=False, + cache_mla_latents=False, + use_cpu_initialization=True, + perform_initialization=True, + symmetric_ar_type=None, + disable_parameter_transpose_cache=False, + init_model_with_meta_device=False, + delay_wgrad_compute=False, + tp_comm_overlap=False, + softmax_scale=None, + apply_dsa_kernel_fusion=apply_dsa_kernel_fusion, + ) + return config + + +def _precompute_freqs_cis(dim: int, seqlen: int, device, base: float) -> torch.Tensor: + freqs = 1.0 / (base ** (torch.arange(0, dim, 2, dtype=torch.float32, device=device) / dim)) + t = torch.arange(seqlen, device=device) + freqs = torch.outer(t, freqs) + return torch.cat((freqs, freqs), dim=-1)[:, None, None, :] + + +def _apply_rotary_emb( + x: torch.Tensor, freqs_cis: torch.Tensor, inverse: bool = False +) -> torch.Tensor: + if x.numel() == 0: + return x + freqs = freqs_cis.to(x.device) + if freqs.dim() == x.dim() + 1 and freqs.size(-2) == 1: + freqs = freqs.squeeze(-2) + + rot_dim = freqs.size(-1) + x_rot, x_pass = x[..., :rot_dim], x[..., rot_dim:] + x1 = x_rot[..., 0::2] + x2 = x_rot[..., 1::2] + x_rot = torch.cat((x1, x2), dim=-1) + + cos = torch.cos(freqs).to(x_rot.dtype) + sin = torch.sin(freqs).to(x_rot.dtype) + if inverse: + sin = -sin + + rot_half_1, rot_half_2 = torch.chunk(x_rot, 2, dim=-1) + x_rotated = torch.cat((-rot_half_2, rot_half_1), dim=-1) + out = (x_rot * cos) + (x_rotated * sin) + + x1, x2 = torch.chunk(out, 2, dim=-1) + out = torch.stack((x1, x2), dim=-1).flatten(start_dim=-2) + return torch.cat((out, x_pass), dim=-1) + + +def _native_hadamard_transform(x: torch.Tensor) -> torch.Tensor: + n = x.size(-1) + if n <= 0 or n & (n - 1): + raise ValueError(f"Hadamard transform requires power-of-two last dim, got {n}") + dtype = x.dtype + y = x.float() + shape = y.shape + h = 1 + while h < n: + y = y.reshape(*shape[:-1], -1, 2, h) + a = y[..., 0, :] + b = y[..., 1, :] + y = torch.cat((a + b, a - b), dim=-1) + h *= 2 + return (y.reshape(shape) * (n**-0.5)).to(dtype) + + +def _get_window_topk_idxs( + window_size: int, batch_size: int, seqlen: int, device: torch.device +) -> torch.Tensor: + base = torch.arange(seqlen, device=device).unsqueeze(1) + offsets = torch.arange(window_size, device=device) + matrix = (base - window_size + 1).clamp(min=0) + offsets + matrix = torch.where(matrix > base, -1, matrix) + return matrix.unsqueeze(0).expand(batch_size, -1, -1) + + +def _get_compress_topk_idxs( + ratio: int, batch_size: int, seqlen: int, offset: int, device: torch.device +) -> torch.Tensor: + n_compressed = seqlen // ratio + matrix = torch.arange(n_compressed, device=device).repeat(seqlen, 1) + mask = matrix >= torch.arange(1, seqlen + 1, device=device).unsqueeze(1) // ratio + matrix = torch.where(mask, -1, matrix + offset) + return matrix.unsqueeze(0).expand(batch_size, -1, -1) + + +def _native_sparse_attn( + query: torch.Tensor, + kv_full: torch.Tensor, + attn_sink: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + sq, batch_size, num_heads, head_dim = query.size() + kv_t = kv_full.permute(1, 0, 2) + safe_indices = topk_indices.clamp(min=0).long() + gather_index = safe_indices.unsqueeze(-1).expand(-1, -1, -1, head_dim) + kv_gathered = torch.gather(kv_t.unsqueeze(1).expand(-1, sq, -1, -1), dim=2, index=gather_index) + + q = query.permute(1, 2, 0, 3).float() + scores = torch.einsum("bnsh,bskh->bnsk", q, kv_gathered.float()) * softmax_scale + scores = scores.masked_fill((topk_indices < 0).unsqueeze(1), float("-inf")) + + sink = attn_sink.view(1, num_heads, 1, 1).float() + scores_max = torch.max(scores.max(dim=-1, keepdim=True).values, sink) + exp_scores = torch.exp(scores - scores_max) + exp_sink = torch.exp(sink - scores_max) + attn_weights = exp_scores / (exp_scores.sum(dim=-1, keepdim=True) + exp_sink) + + output = torch.einsum("bnsk,bskh->bnsh", attn_weights, kv_gathered.float()) + output = output.to(query.dtype).permute(2, 0, 1, 3).contiguous() + return output.reshape(sq, batch_size, num_heads * head_dim) + + +def _native_fused_sparse_indexer_loss( + index_scores: torch.Tensor, + topk_indices: torch.Tensor, + query: torch.Tensor, + compressed_kv: torch.Tensor, + attn_sink: torch.Tensor, + softmax_scale: float, + loss_coeff: float, + calculate_per_token_loss: bool, +) -> torch.Tensor: + batch_size, seqlen, topk = topk_indices.size() + num_heads, head_dim = query.size(2), query.size(3) + safe_indices = topk_indices.clamp(min=0).long() + valid = topk_indices >= 0 + row_valid = valid.any(dim=-1, keepdim=True) + + predict_logits = torch.gather(index_scores, dim=-1, index=safe_indices) + predict_logits = predict_logits.masked_fill(~valid, float("-inf")) + predict_logits = predict_logits.masked_fill(~row_valid, 0.0) + predict = F.softmax(predict_logits, dim=-1, dtype=torch.float32) + predict = predict * row_valid.float() + + compressed_kv_t = compressed_kv.detach().permute(1, 0, 2) + selected_kv = torch.gather( + compressed_kv_t.unsqueeze(1).expand(-1, seqlen, -1, -1), + dim=2, + index=safe_indices.unsqueeze(-1).expand(-1, -1, -1, head_dim), + ) + q = query.detach().permute(1, 2, 0, 3).float() + attn_scores = torch.einsum("bhsd,bskd->bhsk", q, selected_kv.float()) + attn_scores = attn_scores * softmax_scale + attn_scores = attn_scores.masked_fill(~valid.unsqueeze(1), float("-inf")) + + sink = attn_sink.detach().view(1, num_heads, 1, 1).float() + score_max = torch.max(attn_scores.max(dim=-1, keepdim=True).values, sink) + exp_scores = torch.exp(attn_scores - score_max) + exp_sink = torch.exp(sink - score_max) + attn_probs = exp_scores / (exp_scores.sum(dim=-1, keepdim=True) + exp_sink) + target = attn_probs.sum(dim=1) + target = target / target.sum(dim=-1, keepdim=True).clamp(min=1e-10) + target = target * row_valid.float() + + eps = torch.finfo(torch.float32).tiny + target = target.clamp(min=eps) + predict = predict.clamp(min=eps) + kl_per_row = (target * (torch.log(target) - torch.log(predict))).sum(dim=-1) + kl_per_row = torch.where(row_valid.squeeze(-1), kl_per_row, torch.zeros_like(kl_per_row)) + loss = kl_per_row.sum() if calculate_per_token_loss else kl_per_row.mean() + return loss_coeff * loss + + +def _native_unfused_sparse_indexer_loss( + index_scores: torch.Tensor, + topk_indices: torch.Tensor, + query: torch.Tensor, + compressed_kv: torch.Tensor, + softmax_scale: float, + loss_coeff: float, + sparse_loss: bool, + causal_mask: torch.Tensor, + calculate_per_token_loss: bool, +) -> torch.Tensor: + sq, batch_size, num_heads, _ = query.size() + sk = compressed_kv.size(0) + mask = causal_mask.to(dtype=torch.float32) + + attention_scores = torch.einsum( + "sbhd,tbd->bhst", query.detach().float(), compressed_kv.detach().float() + ) + attention_scores = attention_scores * softmax_scale + attention_scores = attention_scores + mask.view(batch_size, 1, sq, sk) + index_scores = index_scores + mask + + if sparse_loss: + index_mask = torch.full( + (batch_size, sq, sk), float("-inf"), dtype=torch.float32, device=index_scores.device + ).scatter_(-1, topk_indices.clamp(min=0), 0) + attention_scores = attention_scores + index_mask.view(batch_size, 1, sq, sk) + index_scores = index_scores + index_mask + + row_valid = (mask > float("-inf")).any(dim=-1) + attn_row_mask = row_valid.view(batch_size, 1, sq, 1) + idx_row_mask = row_valid.view(batch_size, sq, 1) + + attention_scores = attention_scores.masked_fill(~attn_row_mask, 0.0) + index_scores = index_scores.masked_fill(~idx_row_mask, 0.0) + + attention_probs = F.softmax(attention_scores, dim=-1, dtype=torch.float32) + predict = F.softmax(index_scores, dim=-1, dtype=torch.float32) + attention_probs = attention_probs * attn_row_mask.float() + predict = predict * idx_row_mask.float() + + target = attention_probs.sum(dim=1) + target = target / target.sum(dim=-1, keepdim=True) + eps = torch.finfo(torch.float32).tiny + target = target.clamp(min=eps) + predict = predict.clamp(min=eps) + kl_per_row = (target * (torch.log(target) - torch.log(predict))).sum(dim=-1) + kl_per_row = torch.where(row_valid, kl_per_row, torch.zeros_like(kl_per_row)) + loss = kl_per_row.sum() if calculate_per_token_loss else kl_per_row.mean() + return loss * loss_coeff + + +class NativeCompressor(nn.Module): + def __init__( + self, config: MLATransformerConfig, compress_ratio: int, head_dim: int, rotate: bool + ): + super().__init__() + self.compress_ratio = compress_ratio + self.head_dim = head_dim + self.overlap = compress_ratio == 4 + self.coff = 1 + int(self.overlap) + self.rotate = rotate + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + self.rope_base = ( + config.csa_compress_rotary_base if compress_ratio > 1 else config.rotary_base + ) + + self.linear_wkv = nn.Linear(config.hidden_size, self.coff * head_dim, bias=False) + self.linear_wgate = nn.Linear(config.hidden_size, self.coff * head_dim, bias=False) + self.ape = nn.Parameter( + torch.empty(compress_ratio, self.coff * head_dim, dtype=torch.float32) + ) + self.norm = nn.RMSNorm(head_dim, eps=config.layernorm_epsilon) + + def _overlap_transform(self, tensor: torch.Tensor, fill_value: float = 0) -> torch.Tensor: + n_groups, ratio, batch_size, _ = tensor.size() + new_tensor = tensor.new_full((n_groups, 2 * ratio, batch_size, self.head_dim), fill_value) + new_tensor[:, ratio:] = tensor[:, :, :, self.head_dim :] + new_tensor[1:, :ratio] = tensor[:-1, :, :, : self.head_dim] + return new_tensor + + def forward(self, x: torch.Tensor) -> torch.Tensor | None: + sq, batch_size, _ = x.size() + ratio = self.compress_ratio + if sq < ratio: + return None + + kv = self.linear_wkv(x) + score = self.linear_wgate(x) + + cutoff = (sq // ratio) * ratio + kv = kv[:cutoff] + score = score[:cutoff] + n_compressed = cutoff // ratio + + kv = kv.view(n_compressed, ratio, batch_size, -1) + score = score.view(n_compressed, ratio, batch_size, -1) + score = score + self.ape.view(1, ratio, 1, -1) + + if self.overlap: + kv = self._overlap_transform(kv, fill_value=0) + score = self._overlap_transform(score, fill_value=float("-inf")) + + kv = (kv * torch.softmax(score, dim=1)).sum(dim=1) + kv = self.norm(kv.to(x.dtype)) + + pos_dim = self.qk_pos_emb_head_dim + content, rotary = torch.split(kv, [self.head_dim - pos_dim, pos_dim], dim=-1) + freqs_cis = _precompute_freqs_cis( + pos_dim, n_compressed * ratio, device=x.device, base=self.rope_base + ) + freqs_cis = freqs_cis[: n_compressed * ratio : ratio][:n_compressed] + rotary = _apply_rotary_emb(rotary, freqs_cis) + kv = torch.cat([content, rotary], dim=-1) + + if self.rotate: + kv = _native_hadamard_transform(kv) + return kv + + +class NativeCSAIndexer(nn.Module): + def __init__(self, config: MLATransformerConfig, compress_ratio: int): + super().__init__() + self.compress_ratio = compress_ratio + self.index_n_heads = config.dsa_indexer_n_heads + self.index_head_dim = config.dsa_indexer_head_dim + self.index_topk = config.dsa_indexer_topk + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + self.softmax_scale = self.index_head_dim**-0.5 + self.apply_dsa_kernel_fusion = config.apply_dsa_kernel_fusion + self.rope_base = config.csa_compress_rotary_base + + self.linear_wq_b = nn.Linear( + config.q_lora_rank, self.index_n_heads * self.index_head_dim, bias=False + ) + self.linear_weights_proj = nn.Linear(config.hidden_size, self.index_n_heads, bias=False) + self.compressor = NativeCompressor( + config=config, compress_ratio=compress_ratio, head_dim=self.index_head_dim, rotate=True + ) + + def forward_before_topk( + self, x: torch.Tensor, qr: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + sq, batch_size, _ = x.size() + q = self.linear_wq_b(qr).view(sq, batch_size, self.index_n_heads, self.index_head_dim) + pos_dim = self.qk_pos_emb_head_dim + q_content, q_rotary = torch.split(q, [self.index_head_dim - pos_dim, pos_dim], dim=-1) + freqs_cis = _precompute_freqs_cis(pos_dim, sq, device=x.device, base=self.rope_base) + q_rotary = _apply_rotary_emb(q_rotary, freqs_cis) + q = _native_hadamard_transform(torch.cat([q_content, q_rotary], dim=-1)) + + k = self.compressor(x) + weights = self.linear_weights_proj(x) * (self.index_n_heads**-0.5) + return q, k, weights + + def forward( + self, x: torch.Tensor, qr: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + q, k, weights = self.forward_before_topk(x, qr) + weights_scaled = weights.float() * self.softmax_scale + if self.apply_dsa_kernel_fusion: + weights_scaled = weights_scaled.to(weights.dtype).float() + scores = torch.einsum("sbhd,tbd->sbht", q.float(), k.float()) + scores = torch.relu(scores) * weights_scaled.unsqueeze(-1) + scores = scores.sum(dim=2).transpose(0, 1) + + sq = x.size(0) + n_compressed = k.size(0) + valid_per_query = ( + torch.arange(1, sq + 1, device=x.device).unsqueeze(0) // self.compress_ratio + ).clamp(max=n_compressed) + invalid = torch.arange(n_compressed, device=x.device).view( + 1, 1, -1 + ) >= valid_per_query.unsqueeze(-1) + scores = scores.masked_fill(invalid.expand_as(scores), float("-inf")) + + topk = min(self.index_topk, n_compressed) + topk_scores, topk_indices = scores.topk(topk, dim=-1) + topk_indices = torch.where(topk_scores.isneginf(), -1, topk_indices) + return q, k, weights, scores, topk_indices + + +class NativeCompressedSparseAttention(nn.Module): + def __init__(self, config: MLATransformerConfig, compress_ratio: int): + super().__init__() + self.compress_ratio = compress_ratio + self.window_size = config.csa_window_size + self.num_heads = config.num_attention_heads + self.head_dim = config.v_head_dim + self.softmax_scale = self.head_dim**-0.5 + self.indexer_loss_coeff = config.dsa_indexer_loss_coeff + self.indexer_use_sparse_loss = config.dsa_indexer_use_sparse_loss + self.calculate_per_token_loss = config.calculate_per_token_loss + self.apply_dsa_kernel_fusion = config.apply_dsa_kernel_fusion + + self.attn_sink = nn.Parameter(torch.zeros(self.num_heads, dtype=torch.float32)) + self.compressor = ( + NativeCompressor( + config=config, compress_ratio=compress_ratio, head_dim=self.head_dim, rotate=False + ) + if compress_ratio > 1 + else None + ) + self.indexer = NativeCSAIndexer(config, compress_ratio) if compress_ratio == 4 else None + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + x: torch.Tensor, + qr: torch.Tensor, + pg_collection, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + sq, batch_size, _, _ = query.size() + kv = key.squeeze(-2) + n_compressed = 0 + + if self.compressor is not None: + compressed_kv = self.compressor(x) + if compressed_kv is not None: + kv_full = torch.cat([kv, compressed_kv], dim=0) + n_compressed = compressed_kv.size(0) + else: + kv_full = kv + else: + compressed_kv = None + kv_full = kv + + window_idxs = _get_window_topk_idxs(self.window_size, batch_size, sq, query.device) + indexer_loss = None + if self.compress_ratio > 1 and n_compressed > 0: + offset = sq + if self.indexer is not None: + q_idx, k_idx, weights_idx, index_scores, topk_compressed = self.indexer( + x.detach(), qr.detach() + ) + topk_compressed_for_attn = torch.where( + topk_compressed >= 0, topk_compressed + offset, -1 + ) + + if not self.apply_dsa_kernel_fusion: + causal_mask = ( + torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(sq, -1) + ) + positions = torch.arange(1, sq + 1, device=x.device).unsqueeze(1) + causal_mask = ( + torch.where( + causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0 + ) + .unsqueeze(0) + .expand(batch_size, -1, -1) + ) + indexer_loss = _native_unfused_sparse_indexer_loss( + index_scores, + topk_compressed, + query.detach(), + compressed_kv.detach(), + self.softmax_scale, + self.indexer_loss_coeff, + self.indexer_use_sparse_loss, + causal_mask, + self.calculate_per_token_loss, + ) + else: + indexer_loss = _native_fused_sparse_indexer_loss( + index_scores, + topk_compressed, + query, + compressed_kv, + self.attn_sink, + self.softmax_scale, + self.indexer_loss_coeff, + self.calculate_per_token_loss, + ) + else: + topk_compressed_for_attn = _get_compress_topk_idxs( + self.compress_ratio, batch_size, sq, offset, query.device + ) + if self.indexer is not None and self.apply_dsa_kernel_fusion: + topk_idxs = torch.cat([topk_compressed_for_attn, window_idxs], dim=-1) + else: + topk_idxs = torch.cat([window_idxs, topk_compressed_for_attn], dim=-1) + else: + topk_idxs = window_idxs + + output = _native_sparse_attn(query, kv_full, self.attn_sink, topk_idxs, self.softmax_scale) + return output, indexer_loss + + +class NativeDSv4HybridAttention(nn.Module): + def __init__(self, config: MLATransformerConfig, compress_ratio: int): + super().__init__() + self.config = config + self.compress_ratio = compress_ratio + self.num_heads = config.num_attention_heads + self.head_dim = config.v_head_dim + self.pos_dim = config.qk_pos_emb_head_dim + self.nope_dim = config.v_head_dim - config.qk_pos_emb_head_dim + self.rope_base = ( + config.csa_compress_rotary_base if compress_ratio > 1 else config.rotary_base + ) + + self.linear_q_down_proj = nn.Linear(config.hidden_size, config.q_lora_rank, bias=False) + self.q_layernorm = nn.RMSNorm(config.q_lora_rank, eps=config.layernorm_epsilon) + self.linear_q_up_proj = nn.Linear( + config.q_lora_rank, config.num_attention_heads * config.v_head_dim, bias=False + ) + self.linear_kv_proj = nn.Linear(config.hidden_size, config.v_head_dim, bias=False) + self.kv_layernorm = nn.RMSNorm(config.v_head_dim, eps=config.layernorm_epsilon) + self.core_attention = NativeCompressedSparseAttention(config, compress_ratio) + group_in = (config.num_attention_heads * config.v_head_dim) // config.o_groups + self.linear_o_group_proj = nn.Parameter( + torch.empty(config.o_groups * config.o_lora_rank, group_in) + ) + self.linear_proj = nn.Linear( + config.o_groups * config.o_lora_rank, config.hidden_size, bias=False + ) + + def forward( + self, hidden_states: torch.Tensor, pg_collection + ) -> tuple[torch.Tensor, torch.Tensor | None]: + sq, batch_size, _ = hidden_states.size() + freqs_cis = _precompute_freqs_cis(self.pos_dim, sq, hidden_states.device, self.rope_base) + + qr = self.q_layernorm(self.linear_q_down_proj(hidden_states)) + query = self.linear_q_up_proj(qr).view(sq, batch_size, self.num_heads, self.head_dim) + query = _native_q_rms_norm(query, self.config.layernorm_epsilon) + q_content, q_rotary = torch.split(query, [self.nope_dim, self.pos_dim], dim=-1) + query = torch.cat([q_content, _apply_rotary_emb(q_rotary, freqs_cis)], dim=-1) + + key = self.kv_layernorm(self.linear_kv_proj(hidden_states)) + k_content, k_rotary = torch.split(key, [self.nope_dim, self.pos_dim], dim=-1) + key = torch.cat([k_content, _apply_rotary_emb(k_rotary, freqs_cis)], dim=-1) + key = key.unsqueeze(-2) + + core_out, indexer_loss = self.core_attention( + query=query, key=key, x=hidden_states, qr=qr, pg_collection=pg_collection + ) + + core_out = core_out.view(sq, batch_size, self.num_heads, self.head_dim) + out_content, out_rotary = torch.split(core_out, [self.nope_dim, self.pos_dim], dim=-1) + core_out = torch.cat( + [out_content, _apply_rotary_emb(out_rotary, freqs_cis, inverse=True)], dim=-1 + ) + core_out = core_out.view(sq, batch_size, -1) + + core_out = core_out.view(sq, batch_size, self.config.o_groups, -1) + wo_a = self.linear_o_group_proj.view(self.config.o_groups, self.config.o_lora_rank, -1) + core_out = torch.einsum("...gd,grd->...gr", core_out, wo_a) + core_out = core_out.reshape(sq, batch_size, -1) + return self.linear_proj(core_out), indexer_loss + + +def _cosine_sim(a: torch.Tensor, b: torch.Tensor) -> float: + return F.cosine_similarity( + a.flatten().double().unsqueeze(0), b.flatten().double().unsqueeze(0) + ).item() + + +def _tensor_sim(a: torch.Tensor, b: torch.Tensor) -> float: + a, b = a.double(), b.double() + denom = (a * a + b * b).sum() + return (2.0 * (a * b).sum() / denom).item() if denom else 1.0 + + +def _assert_similarity(a: torch.Tensor, b: torch.Tensor, label: str, eps: float): + assert torch.isfinite(a).all() + assert torch.isfinite(b).all() + cosine_sim = _cosine_sim(a, b) + tensor_sim = _tensor_sim(a, b) + assert cosine_sim > 1 - eps, f"{label}: cosine_sim={cosine_sim:.10f}, eps={eps}" + assert tensor_sim > 1 - eps, f"{label}: tensor_sim={tensor_sim:.10f}, eps={eps}" + + +def _copy_real_params_to_native(real_layer: nn.Module, native_layer: nn.Module): + real_params = dict(real_layer.named_parameters()) + for name, native_param in native_layer.named_parameters(): + assert name in real_params, f"Missing real parameter for native parameter {name}" + real_param = real_params[name] + assert ( + native_param.shape == real_param.shape + ), f"Shape mismatch for {name}: native={native_param.shape}, real={real_param.shape}" + native_param.data = real_param.data.to( + device=native_param.device, dtype=real_param.dtype + ).clone() + return real_params + + +def _skip_if_real_kernels_unavailable(*, sm_min: int = 9): + """Pytest-side gate for real-kernel tests. Raises ``pytest.skip`` if + any of the runtime dependencies are missing. + """ + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + sm_major = torch.cuda.get_device_capability()[0] + if sm_major < sm_min: + pytest.skip(f"requires SM{sm_min}+, found SM{sm_major}") + cudnn = pytest.importorskip("cudnn") + cudnn_frontend = pytest.importorskip("cudnn_frontend") + from packaging.version import Version + + if Version(cudnn_frontend.__version__) < Version("1.24.0"): + pytest.skip(f"requires cudnn_frontend>=1.24.0, found {cudnn_frontend.__version__}") + if not hasattr(cudnn, 'DSA'): + pytest.skip("cudnn.DSA namespace not available") + pytest.importorskip("flash_mla") + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +@pytest.mark.parametrize(("backend", "apply_dsa_kernel_fusion"), _DSA_BACKENDS) +@pytest.mark.parametrize("variant", ["flash", "pro"]) +@pytest.mark.parametrize("compress_ratio", [1, 4, 128]) +@pytest.mark.parametrize("seqlen", _CASE_SEQLENS) +@pytest.mark.parametrize("calculate_per_token_loss", [False, True]) +def test_dsv4_hybrid_attention_matches_native_reference( + variant: str, + compress_ratio: int, + seqlen: int, + backend: str, + apply_dsa_kernel_fusion: bool, + calculate_per_token_loss: bool, +): + if apply_dsa_kernel_fusion: + _skip_if_real_kernels_unavailable(sm_min=10) + major, _ = torch.cuda.get_device_capability() + if major < 10 and not apply_dsa_kernel_fusion and seqlen > 4096: + pytest.skip("seqlen > 4096 may OOM on Hopper with unfused DSA implementation") + + DSAIndexerLossAutoScaler.main_loss_backward_scale = None + + Utils.initialize_model_parallel(tensor_model_parallel_size=1, context_parallel_size=1) + try: + torch.manual_seed(_SEED) + torch.cuda.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + config = _make_config( + variant, + compress_ratio, + apply_dsa_kernel_fusion=apply_dsa_kernel_fusion, + calculate_per_token_loss=calculate_per_token_loss, + ) + similarity_eps = ( + _UNFUSED_SIMILARITY_EPS if not apply_dsa_kernel_fusion else _FUSED_SIMILARITY_EPS + ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=["tp", "cp"]) + spec = get_dsv4_hybrid_module_spec_for_backend(config=config, backend=TESpecProvider()) + + mcore_ratio = 0 if compress_ratio == 1 else compress_ratio + real_layer = build_module( + spec, config=config, layer_number=1, cp_comm_type=None, pg_collection=pg_collection + ).cuda() + native_layer = NativeDSv4HybridAttention(config, mcore_ratio).cuda() + real_params = _copy_real_params_to_native(real_layer, native_layer) + + bsz = 1 + for _ in range(1): + hidden_states = torch.randn( + seqlen, + bsz, + 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(), + f"{backend}-{variant}-{compress_ratio}-{seqlen}:out", + eps=similarity_eps, + ) + + real_out.backward(grad) + native_out.backward(grad) + if native_indexer_loss is not None: + native_indexer_loss.backward() + + _assert_similarity( + hidden_states.grad, + hidden_states_native.grad, + f"{backend}-{variant}-{compress_ratio}-{seqlen}:hidden_grad", + eps=similarity_eps, + ) + + for name, native_param in native_layer.named_parameters(): + real_param = real_params[name] + if compress_ratio != 4 and ".indexer." in name: + continue + 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"{backend}-{variant}-{compress_ratio}-{seqlen}:param_grad:{name}", + eps=similarity_eps, + ) + del real_layer, native_layer, real_params + del hidden_states, hidden_states_native, grad, real_out, native_out, native_indexer_loss + finally: + Utils.destroy_model_parallel() + gc.collect() + torch.cuda.empty_cache() diff --git a/uv.lock b/uv.lock index 62ea93ec3ac..a2628601c01 100644 --- a/uv.lock +++ b/uv.lock @@ -1,5 +1,5 @@ version = 1 -revision = 2 +revision = 3 requires-python = ">=3.10" resolution-markers = [ "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'win32' and extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts'", @@ -305,7 +305,7 @@ version = "1.4.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "frozenlist" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/61/62/06741b579156360248d1ec624842ad0edf697050bbaf7c3e46394e106ad1/aiosignal-1.4.0.tar.gz", hash = "sha256:f47eecd9468083c2029cc99945502cb7708b082c232f9aca65da147157b251c7", size = 25007, upload-time = "2025-07-03T22:54:43.528Z" } wheels = [ @@ -359,10 +359,10 @@ name = "anyio" version = "4.9.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "exceptiongroup", marker = "python_full_version < '3.11'" }, + { name = "exceptiongroup", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "idna" }, { name = "sniffio" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/95/7d/4c1bd541d4dffa1b52bd83fb8527089e097a106fc90b467a7313b105f840/anyio-4.9.0.tar.gz", hash = "sha256:673c0c244e15788651a4ff38710fea9675823028a6f08a5eda409e0c9840a028", size = 190949, upload-time = "2025-03-17T00:02:54.77Z" } wheels = [ @@ -797,7 +797,7 @@ name = "cffi" version = "2.0.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "pycparser", marker = "implementation_name != 'PyPy'" }, + { name = "pycparser", marker = "implementation_name != 'PyPy' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/eb/56/b1ba7935a17738ae8453301356628e8147c79dbb825bcbc73dc7401f9846/cffi-2.0.0.tar.gz", hash = "sha256:44d1b5909021139fe36001ae048dbdde8214afa20200eda0f64c068cac5d5529", size = 523588, upload-time = "2025-09-08T23:24:04.541Z" } wheels = [ @@ -984,7 +984,7 @@ name = "click" version = "8.3.2" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/57/75/31212c6bf2503fdf920d87fee5d7a86a2e3bcf444984126f13d8e4016804/click-8.3.2.tar.gz", hash = "sha256:14162b8b3b3550a7d479eafa77dfd3c38d9dc8951f6f69c78913a8f9a7540fd5", size = 302856, upload-time = "2026-04-03T19:14:45.118Z" } wheels = [ @@ -1153,8 +1153,8 @@ name = "cryptography" version = "46.0.6" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cffi", marker = "platform_python_implementation != 'PyPy'" }, - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "cffi", marker = "platform_python_implementation != 'PyPy' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "typing-extensions", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/a4/ba/04b1bd4218cbc58dc90ce967106d51582371b898690f3ae0402876cc4f34/cryptography-46.0.6.tar.gz", hash = "sha256:27550628a518c5c6c903d84f637fbecf287f6cb9ced3804838a1295dc1fd0759", size = 750542, upload-time = "2026-03-25T23:34:53.396Z" } wheels = [ @@ -1213,7 +1213,7 @@ name = "cuda-bindings" version = "13.2.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-pathfinder" }, + { name = "cuda-pathfinder", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-lts')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/1a/fe/7351d7e586a8b4c9f89731bfe4cf0148223e8f9903ff09571f78b3fb0682/cuda_bindings-13.2.0-cp310-cp310-manylinux_2_24_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:08b395f79cb89ce0cd8effff07c4a1e20101b873c256a1aeb286e8fd7bd0f556", size = 5744254, upload-time = "2026-03-11T00:12:29.798Z" }, @@ -1266,37 +1266,37 @@ wheels = [ [package.optional-dependencies] cublas = [ - { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] cudart = [ - { name = "nvidia-cuda-runtime", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cuda-runtime", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] cufft = [ - { name = "nvidia-cufft", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cufft", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] cufile = [ - { name = "nvidia-cufile", marker = "sys_platform == 'linux'" }, + { name = "nvidia-cufile", marker = "sys_platform == 'linux' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] cupti = [ - { name = "nvidia-cuda-cupti", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cuda-cupti", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] curand = [ - { name = "nvidia-curand", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-curand", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] cusolver = [ - { name = "nvidia-cusolver", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cusolver", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] cusparse = [ - { name = "nvidia-cusparse", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cusparse", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] nvjitlink = [ - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] nvrtc = [ - { name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-cuda-nvrtc", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] nvtx = [ - { name = "nvidia-nvtx", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or sys_platform == 'linux'" }, + { name = "nvidia-nvtx", marker = "(python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version >= '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or sys_platform == 'linux' or (sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] [[package]] @@ -1564,8 +1564,8 @@ name = "emerging-optimizers" version = "0.2.0" source = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0#1effa026ff096b7fa1063ca2fba19d98be6e6cdf" } dependencies = [ - { name = "absl-py", marker = "python_full_version >= '3.12'" }, - { name = "torch", marker = "python_full_version >= '3.12' and sys_platform == 'never'" }, + { name = "absl-py", marker = "python_full_version >= '3.12' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "torch", marker = "(python_full_version >= '3.12' and sys_platform == 'never') or (python_full_version < '3.12' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform != 'never' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] [[package]] @@ -1573,7 +1573,7 @@ name = "exceptiongroup" version = "1.3.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/50/79/66800aadf48771f6b62f7eb014e352e5d06856655206165d775e675a02c9/exceptiongroup-1.3.1.tar.gz", hash = "sha256:8b412432c6055b0b7d14c310000ae93352ed6754f70fa8f7c34141f91c4e3219", size = 30371, upload-time = "2025-11-21T23:01:54.787Z" } wheels = [ @@ -2371,7 +2371,7 @@ dependencies = [ { name = "filelock" }, { name = "fsspec", version = "2026.2.0", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.14' or sys_platform != 'win32' or extra == 'extra-13-megatron-core-dev' or extra == 'extra-13-megatron-core-lts'" }, { name = "fsspec", version = "2026.3.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.14' and sys_platform == 'win32' and extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "hf-xet", marker = "platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64'" }, + { name = "hf-xet", marker = "platform_machine == 'aarch64' or platform_machine == 'amd64' or platform_machine == 'arm64' or platform_machine == 'x86_64' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "packaging" }, { name = "pyyaml" }, { name = "requests" }, @@ -2941,7 +2941,7 @@ resolution-markers = [ "python_full_version < '3.11' and extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", ] dependencies = [ - { name = "mdurl", marker = "python_full_version < '3.11'" }, + { name = "mdurl", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/38/71/3b932df36c1a044d397a1f92d1cf91ee0a503d91e470cbd670aa66b07ed0/markdown-it-py-3.0.0.tar.gz", hash = "sha256:e3f60a94fa066dc52ec76661e37c851cb232d92f9886b15cb560aaada2df8feb", size = 74596, upload-time = "2023-06-03T06:41:14.443Z" } wheels = [ @@ -3015,7 +3015,7 @@ resolution-markers = [ "python_full_version == '3.11.*' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", ] dependencies = [ - { name = "mdurl", marker = "python_full_version >= '3.11'" }, + { name = "mdurl", marker = "python_full_version >= '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/5b/f5/4ec618ed16cc4f8fb3b701563655a69816155e79e24a17b651541804721d/markdown_it_py-4.0.0.tar.gz", hash = "sha256:cb0a2b4aa34f932c007117b194e945bd74e0ec24133ceb5bac59009cda1cb9f3", size = 73070, upload-time = "2025-08-11T12:57:52.854Z" } wheels = [ @@ -3187,6 +3187,7 @@ dev = [ { name = "mamba-ssm" }, { name = "megatron-energon", extra = ["av-decode"], marker = "extra == 'extra-13-megatron-core-dev'" }, { name = "multi-storage-client" }, + { name = "nvidia-cudnn-frontend" }, { name = "nvidia-modelopt", marker = "(sys_platform != 'darwin' and extra == 'extra-13-megatron-core-dev') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "nvidia-resiliency-ext" }, { name = "nvtx" }, @@ -3321,6 +3322,7 @@ requires-dist = [ { name = "multi-storage-client", marker = "extra == 'dev'", specifier = "~=0.27" }, { name = "multi-storage-client", marker = "extra == 'lts'", specifier = "~=0.27" }, { name = "numpy" }, + { name = "nvidia-cudnn-frontend", marker = "extra == 'dev'" }, { name = "nvidia-modelopt", extras = ["torch"], marker = "sys_platform != 'darwin' and extra == 'dev'" }, { name = "nvidia-resiliency-ext", marker = "extra == 'dev'", git = "https://github.com/NVIDIA/nvidia-resiliency-ext.git?rev=15a851565a4ce846c04431ecb0cf09903ab4837e" }, { name = "nvtx", marker = "extra == 'dev'", specifier = "~=0.2" }, @@ -3653,7 +3655,7 @@ name = "multidict" version = "6.7.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/1a/c2/c2d94cbe6ac1753f3fc980da97b3d930efe1da3af3c9f5125354436c073d/multidict-6.7.1.tar.gz", hash = "sha256:ec6652a1bee61c53a3e5776b6049172c53b6aaba34f18c9ad04f82712bac623d", size = 102010, upload-time = "2026-01-26T02:46:45.979Z" } wheels = [ @@ -4110,24 +4112,35 @@ name = "numpy" version = "2.0.2" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'win32'", - "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'win32'", - "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'emscripten'", - "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'emscripten'", - "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'win32'", - "python_full_version < '3.13' and platform_machine != 's390x' and sys_platform == 'win32'", - "python_full_version < '3.13' and platform_machine == 's390x' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", - "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", - "python_full_version < '3.13' and platform_machine != 's390x' and sys_platform == 'emscripten'", - "python_full_version < '3.13' and platform_machine == 's390x' and sys_platform == 'emscripten'", - "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version < '3.13' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version < '3.13' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.11' and platform_machine != 's390x' and extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts'", + "python_full_version < '3.11' and platform_machine == 's390x' and extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version < '3.11' and platform_machine != 's390x' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version < '3.11' and platform_machine == 's390x' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version < '3.11' and extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", ] sdist = { url = "https://files.pythonhosted.org/packages/a9/75/10dd1f8116a8b796cb2c737b674e02d02e80454bda953fa7e65d8c12b016/numpy-2.0.2.tar.gz", hash = "sha256:883c987dee1880e2a864ab0dc9892292582510604156762362d9326444636e78", size = 18902015, upload-time = "2024-08-26T20:19:40.945Z" } wheels = [ @@ -4280,7 +4293,6 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/04/74/f4c001f4714c3ad9ce037e18cf2b9c64871a84951eaa0baf683a9ca9301c/numpy-2.4.4-pp311-pypy311_pp73-win_amd64.whl", hash = "sha256:f2cf083b324a467e1ab358c105f6cad5ea950f50524668a80c486ff1db24e119", size = 12509075, upload-time = "2026-03-29T13:21:57.644Z" }, ] - [[package]] name = "nvdlfw-inspect" version = "0.2.2" @@ -4338,7 +4350,7 @@ name = "nvidia-cudnn-cu13" version = "9.19.0.56" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/f1/84/26025437c1e6b61a707442184fa0c03d083b661adf3a3eecfd6d21677740/nvidia_cudnn_cu13-9.19.0.56-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:6ed29ffaee1176c612daf442e4dd6cfeb6a0caa43ddcbeb59da94953030b1be4", size = 433781201, upload-time = "2026-02-03T20:40:53.805Z" }, @@ -4348,24 +4360,27 @@ wheels = [ [[package]] name = "nvidia-cudnn-frontend" -version = "1.22.0" -source = { registry = "https://pypi.org/simple" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/40/7d/28ab9cb9119fc6a3a383d943448ab310fe787daf784869b167dc7269969f/nvidia_cudnn_frontend-1.22.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:dbd3100ae212dd1f4691f8c096fe3aded46491f9a6cb258bfb802d07ca1a88fc", size = 2670597, upload-time = "2026-04-03T02:27:56.886Z" }, - { url = "https://files.pythonhosted.org/packages/8b/b4/976996f1ab721bbcae4b7379652949ddcd41803817d4b65b9bd0d726aa60/nvidia_cudnn_frontend-1.22.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:62bf9c8569caf4d9518dae0755507ad36a4e311726aa015fde104c38a1630f76", size = 2811815, upload-time = "2026-04-03T02:32:24.504Z" }, - { url = "https://files.pythonhosted.org/packages/2b/56/755412cf4ce5ad95bcb00be3144c8e1fa07cbbae073f31a7b75ddec96ca0/nvidia_cudnn_frontend-1.22.0-cp310-cp310-win_amd64.whl", hash = "sha256:22748b41049d02c029719467924ea20d928517dd8f35e204a390f97407298eb2", size = 2260435, upload-time = "2026-04-03T02:24:51.808Z" }, - { url = "https://files.pythonhosted.org/packages/bf/ff/e4955b6fdff929ddf04a1252facae6201b308e001c91c690e96f65c4e90a/nvidia_cudnn_frontend-1.22.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cdff54c945fbabf9da06fd64ded60cf1ec94d580474f5746786c0effd759fedc", size = 2672347, upload-time = "2026-04-03T02:28:51.106Z" }, - { url = "https://files.pythonhosted.org/packages/52/27/62fc6e2cddff7d6396be3685342ceec1c12fe2ee50e6f31d270887ecb5ad/nvidia_cudnn_frontend-1.22.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:bb50bd2758c6d47c6210451c5c1932ed16e7563d7629228f4cc97edc0e01d0c5", size = 2814387, upload-time = "2026-04-03T02:32:47.972Z" }, - { url = "https://files.pythonhosted.org/packages/d5/4f/de06583ec21313f31d8b83bc2164e88fc22f5b48d8eb5cb45490fcf7c262/nvidia_cudnn_frontend-1.22.0-cp311-cp311-win_amd64.whl", hash = "sha256:49f817377a19e10e4aafa5797cd68315739dfdb2fc6a67dd1052b64c805d24ec", size = 2261332, upload-time = "2026-04-03T02:25:17.241Z" }, - { url = "https://files.pythonhosted.org/packages/7e/f1/67681e585abd98f968298c771b72830ce984a90fd0d787098d2ea2ba55c7/nvidia_cudnn_frontend-1.22.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:bc9c12891d5427ef49b72b26df2b7889d623086d77c9e33b021c2de417d3e4dc", size = 2673215, upload-time = "2026-04-03T02:29:41.421Z" }, - { url = "https://files.pythonhosted.org/packages/0e/46/95b7779a2f71dfccce1783cc5ac210dda0124b93f8bf66cf62ed3d9ce0a5/nvidia_cudnn_frontend-1.22.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:98ffa05699d71795372f112fa2361c13be716fa3fda911c1e809903163ea5d11", size = 2815106, upload-time = "2026-04-03T02:33:11.473Z" }, - { url = "https://files.pythonhosted.org/packages/61/47/522e84a37eedb1f680e74df449d39fe6f8641779523313d1a8522d449766/nvidia_cudnn_frontend-1.22.0-cp312-cp312-win_amd64.whl", hash = "sha256:81fde93d9b86ad631e17da1e2c103c4a7a541ec7abcb7f9a121cbd018c8eff26", size = 2262120, upload-time = "2026-04-03T02:25:40.18Z" }, - { url = "https://files.pythonhosted.org/packages/c7/93/43541b581207024824cb740f429bf882aaf3bde3633bd4099393dd9c0c16/nvidia_cudnn_frontend-1.22.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c9bdf48cf989b2a77f8b52623fc31c078362fd34389207d11cdb0b5624a7b311", size = 2673259, upload-time = "2026-04-03T02:30:30.634Z" }, - { url = "https://files.pythonhosted.org/packages/9b/5b/af9da5a455064380e68a441b9cfa1f1212dd6363bd02b5aa696d319bd211/nvidia_cudnn_frontend-1.22.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d02c4b4aae3e243ddb08ad4eb939988bcf7b1aefe25f5d400f6858c7276a6631", size = 2815032, upload-time = "2026-04-03T02:33:34.171Z" }, - { url = "https://files.pythonhosted.org/packages/ba/1d/3a15b719817ca6241e5f3a7a38608af21a3259e550a5dee5520e29adac00/nvidia_cudnn_frontend-1.22.0-cp313-cp313-win_amd64.whl", hash = "sha256:4906a38954725e35bc8431874f4d9db60d50e0d9dbc40ecaf8e5f40df545350b", size = 2262156, upload-time = "2026-04-03T02:26:03.322Z" }, - { url = "https://files.pythonhosted.org/packages/27/ec/8c9b53a9174cca2d0062cbd8cb7c31403a38cb4c79984a9c554830cac5e9/nvidia_cudnn_frontend-1.22.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0f650058bda46a6542dfc3d021803021e7932e1cd6bb78cf46e81fa219717b5e", size = 2674887, upload-time = "2026-04-03T02:31:21.166Z" }, - { url = "https://files.pythonhosted.org/packages/89/bd/3464d181ec2d94085cab98fd5ea4d312478aa6cb16ff38994a9188ac9f05/nvidia_cudnn_frontend-1.22.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:90f30b0d6563d050ca1972efa594a31d5affe5c3eeb467542e715d7ee73e3b5b", size = 2815841, upload-time = "2026-04-03T02:33:56.66Z" }, - { url = "https://files.pythonhosted.org/packages/a5/fd/bdec32a32b44f52b60a03f43e8619552ea0eb90a61de06632a054bf17d6a/nvidia_cudnn_frontend-1.22.0-cp314-cp314-win_amd64.whl", hash = "sha256:5994400a7f76a1be5e327a9ac1a4a635ee734d2ac8a5875e52481c52cf2b0922", size = 2263464, upload-time = "2026-04-03T02:26:26.553Z" }, +version = "1.24.0" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/71/05/bf31134b6e5d41c5a1e4abc65b5bb5461a73be74013da273b18fef7a2244/nvidia_cudnn_frontend-1.24.0-cp310-cp310-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:e8833079f0283948cb5f99a2dc0c8fbff29d320e6a5635c4f77fa7eaf877043b", size = 3222069, upload-time = "2026-05-20T05:01:12.464Z" }, + { url = "https://files.pythonhosted.org/packages/50/4d/861c75757e688a5094da621871688d7583b52c7cb3ad75d0d5ab1dcdff68/nvidia_cudnn_frontend-1.24.0-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5736397ab8f29e06731960e055d27f65e8fda70c70f988e8bd91b671a3506999", size = 3370408, upload-time = "2026-05-20T05:01:36.098Z" }, + { url = "https://files.pythonhosted.org/packages/de/13/cee7f47acbb1d85a2019522f88238cbd4ac182a65fdea5d41e504b7c48be/nvidia_cudnn_frontend-1.24.0-cp310-cp310-win_amd64.whl", hash = "sha256:77bc9f3203c677f74b6cedf84125514b4881dc82f4177cc4ab33949693abe6aa", size = 2764295, upload-time = "2026-05-20T05:01:56.212Z" }, + { url = "https://files.pythonhosted.org/packages/81/75/767973a56b98d2a8fbc04c78fc28684cdb0df7c032ba4858090243df81cc/nvidia_cudnn_frontend-1.24.0-cp311-cp311-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:9c0f39f211bb105798c7a8617b1d674e01dbd538b97714025628ea7127bfaf8a", size = 3223392, upload-time = "2026-05-20T05:02:56.092Z" }, + { url = "https://files.pythonhosted.org/packages/f5/8c/f5226ab5163dfffbe82ff41b9e1fbf649c908c077a5416bb16a1c3634cfa/nvidia_cudnn_frontend-1.24.0-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dec07dcfea168792098b9a2652ec465d79e228d3d17a2f86a08404b487336530", size = 3372046, upload-time = "2026-05-20T05:03:20.091Z" }, + { url = "https://files.pythonhosted.org/packages/a7/bf/fa08a8bd953ad8db7ea3f64676b30f1cb2846f6518595a6faf7fb3558db4/nvidia_cudnn_frontend-1.24.0-cp311-cp311-win_amd64.whl", hash = "sha256:b461259b85b7a7e3a1c41b02c33ce4fde0dbcde7e0a227a968dedb74d311e2c3", size = 2764877, upload-time = "2026-05-20T05:04:21.965Z" }, + { url = "https://files.pythonhosted.org/packages/25/f4/a57efcbeb1dec0a047fa8f8aaa4defa4935db7d52e8afaf83224f1258ec6/nvidia_cudnn_frontend-1.24.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:3b4398cecbaa555baa73a9b8716233632e3c16259c6ab999d83c51ca3b8fd09c", size = 3224420, upload-time = "2026-05-20T05:04:50.051Z" }, + { url = "https://files.pythonhosted.org/packages/19/b7/c1b8de8292b8fe28b2ffc95601a0b69392536a9515315bcfcdb0b07d2af8/nvidia_cudnn_frontend-1.24.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:144bfe09098d681d4c793c867fff53b1fb7c49f845324d5d0c52d824b14b74f5", size = 3376971, upload-time = "2026-05-20T05:05:19.998Z" }, + { url = "https://files.pythonhosted.org/packages/dc/00/284b8de99fbc1e8fd91c292a024943ee61d3361aa669c435ee44e14b6498/nvidia_cudnn_frontend-1.24.0-cp312-cp312-win_amd64.whl", hash = "sha256:5476d6a51ebaf5ef04e462e0052f1d9bca1af6274f738cb509715b4cf443a8df", size = 2765023, upload-time = "2026-05-20T05:05:40.269Z" }, + { url = "https://files.pythonhosted.org/packages/41/3e/430941e91a0c5234c79aa0bd6dcdb27c0e3a443f66953f881f8cb3428e93/nvidia_cudnn_frontend-1.24.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:ccf1d2352f4b82fafbfb803512493c7a211224b424b6a78fbf45e42a94dcdb81", size = 3224661, upload-time = "2026-05-20T05:06:01.577Z" }, + { url = "https://files.pythonhosted.org/packages/d8/95/02691a6f0db4c2194899579a9c196df8990b140e0503fc6be1f09e7aa063/nvidia_cudnn_frontend-1.24.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:ec45b08e0ab511f61532bc980c343c15eb5eda9ff14a3c80e75bc0aa3776860c", size = 3376128, upload-time = "2026-05-20T05:06:21.613Z" }, + { url = "https://files.pythonhosted.org/packages/d0/c5/e6c9ef439e167675e64fd1ca025df79211cafdf70ec9edbf6b99f876cafd/nvidia_cudnn_frontend-1.24.0-cp313-cp313-win_amd64.whl", hash = "sha256:2d02744a46726d262d80ed54299fae6491e4385e7d580eb6a027fb5b3b2c1db8", size = 2765241, upload-time = "2026-05-20T05:06:45.6Z" }, + { url = "https://files.pythonhosted.org/packages/f0/fc/c5a1d386f22dc8e17304874368ba213ff8c093093e9b0ccc3b2bd81a1d12/nvidia_cudnn_frontend-1.24.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:190857577a11d22b62da1863cb1a4b72692f98913d1c9c2e6a72224d08685575", size = 3226124, upload-time = "2026-05-20T05:07:14.426Z" }, + { url = "https://files.pythonhosted.org/packages/2e/39/e7f12c1a640174bdebbbf87be19819483c9da9f8f3af948f5dcb46cf501d/nvidia_cudnn_frontend-1.24.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b7e9aca4b6ce4d4bd484f01a1c4f530a9bea317cd69d21760d582b47e05514d8", size = 3375383, upload-time = "2026-05-20T05:07:38.976Z" }, + { url = "https://files.pythonhosted.org/packages/16/f8/cae1cffd3ab944301e3a736395f85925a887fc5255c100f561638d6b2fba/nvidia_cudnn_frontend-1.24.0-cp314-cp314-win_amd64.whl", hash = "sha256:b8968eb9dd9a71fe3d64b55d1e9731cffb7272578a9a39c9bf816f5e27c3f14e", size = 2765850, upload-time = "2026-05-20T05:07:59.264Z" }, + { url = "https://files.pythonhosted.org/packages/9e/27/c04b542fd2a882fabed1cb4e6778a8a6159f1123e00c2f0da8383690c713/nvidia_cudnn_frontend-1.24.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:65044724acc5fcb92ab829c2d2fcefc7fd4419881030e04358a66a4bcbce5d43", size = 3230149, upload-time = "2026-05-20T05:08:21.883Z" }, + { url = "https://files.pythonhosted.org/packages/fc/17/89a0eccbf5be9455c395f977022e299c910f1d7d5b0820661908f552ac70/nvidia_cudnn_frontend-1.24.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:37c6ef88c7cbc41eab6e36d5c6715ff0d6c639f76ca12c65ce7a24b453e184eb", size = 3382455, upload-time = "2026-05-20T05:08:48.412Z" }, + { url = "https://files.pythonhosted.org/packages/d7/d2/dac91d6a6fa2e6c07b59239f01fb07472f36d8368c29840647cfcd1c7dd0/nvidia_cudnn_frontend-1.24.0-cp314-cp314t-win_amd64.whl", hash = "sha256:79e902e124123d84d52fa06c0931415cccfa71ff8550c1c156dd3539bffeda5c", size = 2791017, upload-time = "2026-05-20T05:09:07.477Z" }, ] [[package]] @@ -4373,7 +4388,7 @@ name = "nvidia-cufft" version = "12.0.0.61" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/8b/ae/f417a75c0259e85c1d2f83ca4e960289a5f814ed0cea74d18c353d3e989d/nvidia_cufft-12.0.0.61-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:2708c852ef8cd89d1d2068bdbece0aa188813a0c934db3779b9b1faa8442e5f5", size = 214053554, upload-time = "2025-09-04T08:31:38.196Z" }, @@ -4405,9 +4420,9 @@ name = "nvidia-cusolver" version = "12.0.4.66" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "nvidia-cusparse", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-cublas", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "nvidia-cusparse", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/c8/c3/b30c9e935fc01e3da443ec0116ed1b2a009bb867f5324d3f2d7e533e776b/nvidia_cusolver-12.0.4.66-py3-none-manylinux_2_27_aarch64.whl", hash = "sha256:02c2457eaa9e39de20f880f4bd8820e6a1cfb9f9a34f820eb12a155aa5bc92d2", size = 223467760, upload-time = "2025-09-04T08:33:04.222Z" }, @@ -4420,7 +4435,7 @@ name = "nvidia-cusparse" version = "12.6.3.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-nvjitlink", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/f8/94/5c26f33738ae35276672f12615a64bd008ed5be6d1ebcb23579285d960a9/nvidia_cusparse-12.6.3.3-py3-none-manylinux2014_aarch64.manylinux_2_17_aarch64.whl", hash = "sha256:80bcc4662f23f1054ee334a15c72b8940402975e0eab63178fc7e670aa59472c", size = 162155568, upload-time = "2025-09-04T08:33:42.864Z" }, @@ -4715,11 +4730,11 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 's390x'", ] dependencies = [ - { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" } }, + { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-lts') or (python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "protobuf" }, - { name = "typing-extensions" }, + { name = "protobuf", marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "typing-extensions", marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c5/93/942d2a0f6a70538eea042ce0445c8aefd46559ad153469986f29a743c01c/onnx-1.21.0.tar.gz", hash = "sha256:4d8b67d0aaec5864c87633188b91cc520877477ec0254eda122bef8be43cd764", size = 12074608, upload-time = "2026-03-27T21:33:36.118Z" } wheels = [ @@ -4814,12 +4829,12 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 's390x'", ] dependencies = [ - { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" } }, + { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-lts') or (python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "onnx", version = "1.21.0", source = { registry = "https://pypi.org/simple" } }, - { name = "sympy" }, - { name = "typing-extensions" }, + { name = "onnx", version = "1.21.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "sympy", marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "typing-extensions", marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/b2/a5/acc43c8fa6edbc584d127fb6bbd13ae9ebfc01b9675c74e0da2de15fa4a6/onnx_ir-0.2.0.tar.gz", hash = "sha256:8bad3906691987290789b26d05e0dbff467029a0b1e411e12e4cae02e43503e4", size = 141693, upload-time = "2026-02-24T02:31:10.998Z" } wheels = [ @@ -4890,13 +4905,13 @@ resolution-markers = [ "python_full_version < '3.11' and platform_machine == 's390x'", ] dependencies = [ - { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" } }, + { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-lts') or (python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "onnx", version = "1.21.0", source = { registry = "https://pypi.org/simple" } }, - { name = "onnx-ir", version = "0.2.0", source = { registry = "https://pypi.org/simple" } }, - { name = "packaging" }, - { name = "typing-extensions" }, + { name = "onnx", version = "1.21.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "onnx-ir", version = "0.2.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "packaging", marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "typing-extensions", marker = "(python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/e7/2b/538fdeb0e25bed5d7e0f954af5710543e2629499fb74381afc3333f8a8ae/onnxscript-0.6.2.tar.gz", hash = "sha256:abb2e6f464db40c9b8c7fbb3e64cca04cf3f4495e67c4eda5eac17b784191ce3", size = 590865, upload-time = "2026-02-10T22:53:39.638Z" } wheels = [ @@ -4994,7 +5009,7 @@ wheels = [ [[package]] name = "opentelemetry-api" -version = "1.41.0" +version = "1.42.0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", @@ -5012,12 +5027,11 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "importlib-metadata", marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, { name = "typing-extensions", marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/47/8e/3778a7e87801d994869a9396b9fc2a289e5f9be91ff54a27d41eace494b0/opentelemetry_api-1.41.0.tar.gz", hash = "sha256:9421d911326ec12dee8bc933f7839090cad7a3f13fcfb0f9e82f8174dc003c09", size = 71416, upload-time = "2026-04-09T14:38:34.544Z" } +sdist = { url = "https://files.pythonhosted.org/packages/86/ca/25288069c399be6769159d9fb7b1190b603537d82aad2fa2746a0cc2c8c6/opentelemetry_api-1.42.0.tar.gz", hash = "sha256:ea84c893ad177791d138e0349d6ceebd8d3bf006440900400ce220008dafc372", size = 72300, upload-time = "2026-05-19T09:46:29.885Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/58/ee/99ab786653b3bda9c37ade7e24a7b607a1b1f696063172768417539d876d/opentelemetry_api-1.41.0-py3-none-any.whl", hash = "sha256:0e77c806e6a89c9e4f8d372034622f3e1418a11bdbe1c80a50b3d3397ad0fa4f", size = 69007, upload-time = "2026-04-09T14:38:11.833Z" }, + { url = "https://files.pythonhosted.org/packages/1b/0b/be5daf659b82b525338fde371dfcfab09b606a19bb5620c37076964710ec/opentelemetry_api-1.42.0-py3-none-any.whl", hash = "sha256:558d88f88192a973579910ef6f2c13db47a268d5ec2e53e83e50e74a39a02922", size = 61310, upload-time = "2026-05-19T09:46:06.561Z" }, ] [[package]] @@ -5064,7 +5078,7 @@ wheels = [ [[package]] name = "opentelemetry-exporter-prometheus" -version = "0.62b0" +version = "0.63b0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", @@ -5082,13 +5096,13 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "opentelemetry-api", version = "1.41.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, - { name = "opentelemetry-sdk", version = "1.41.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, + { name = "opentelemetry-api", version = "1.42.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, + { name = "opentelemetry-sdk", version = "1.42.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, { name = "prometheus-client", marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/0f/ec/fa8a722199dc2e75dc582779d62207b00b0bdb014b5635594afa0cf3ee43/opentelemetry_exporter_prometheus-0.62b0.tar.gz", hash = "sha256:4d1106566a9b3e8dff028e69e9f2dc90723e6b431c900ff8c72982fcf11dbae5", size = 15441, upload-time = "2026-04-09T14:38:38.934Z" } +sdist = { url = "https://files.pythonhosted.org/packages/18/2c/0643113a5bef20e8242f7ae7915913fab61e8c901d391518a0aefa2da6fc/opentelemetry_exporter_prometheus-0.63b0.tar.gz", hash = "sha256:76b52078ee70131542e53d5cf1942cadd6d5628e7a1bf1f60047f29fa079e9b1", size = 15231, upload-time = "2026-05-19T09:46:35.1Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/d1/1e/43645fadd561471af2aec95906a3dd54af1a8e7782322310e802a810ad3a/opentelemetry_exporter_prometheus-0.62b0-py3-none-any.whl", hash = "sha256:cd7e8acae3be5f425ffa2e0864eea474fa7a40706f786de7a2d23846573d8f75", size = 13278, upload-time = "2026-04-09T14:38:19.367Z" }, + { url = "https://files.pythonhosted.org/packages/2b/48/18e2b0eec9242beb168b1100ed6c602c2a378d0ccb779d1c0a1b85b9ba89/opentelemetry_exporter_prometheus-0.63b0-py3-none-any.whl", hash = "sha256:0cfe4846bf5905f096a4d9678ffe25c7fe6f662f6c7282d2b191138d6bf487fb", size = 12466, upload-time = "2026-05-19T09:46:15.025Z" }, ] [[package]] @@ -5147,7 +5161,7 @@ wheels = [ [[package]] name = "opentelemetry-sdk" -version = "1.41.0" +version = "1.42.0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", @@ -5165,13 +5179,13 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "opentelemetry-api", version = "1.41.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, - { name = "opentelemetry-semantic-conventions", version = "0.62b0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, + { name = "opentelemetry-api", version = "1.42.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, + { name = "opentelemetry-semantic-conventions", version = "0.63b0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, { name = "typing-extensions", marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/f8/0e/a586df1186f9f56b5a0879d52653effc40357b8e88fc50fe300038c3c08b/opentelemetry_sdk-1.41.0.tar.gz", hash = "sha256:7bddf3961131b318fc2d158947971a8e37e38b1cd23470cfb72b624e7cc108bd", size = 230181, upload-time = "2026-04-09T14:38:47.225Z" } +sdist = { url = "https://files.pythonhosted.org/packages/b7/c9/dabaaf1c754a57b82b5a36aeca3806d92c1877ccfb12a697b65f88bf027c/opentelemetry_sdk-1.42.0.tar.gz", hash = "sha256:2479e462cc69357825c2c847ce4a601bc1b17e1279aa7f80d3490f0ae614d0e5", size = 239072, upload-time = "2026-05-19T09:46:42.992Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/2c/13/a7825118208cb32e6a4edcd0a99f925cbef81e77b3b0aedfd9125583c543/opentelemetry_sdk-1.41.0-py3-none-any.whl", hash = "sha256:a596f5687964a3e0d7f8edfdcf5b79cbca9c93c7025ebf5fb00f398a9443b0bd", size = 180214, upload-time = "2026-04-09T14:38:30.657Z" }, + { url = "https://files.pythonhosted.org/packages/7b/7d/16bf9a9d42ebbd1679e0cda018d57a0712f3b6f6f1e7ae5ef3c7ee5927c0/opentelemetry_sdk-1.42.0-py3-none-any.whl", hash = "sha256:ec4a4f69e15220b3d7bccd93217aac745682bb6435b9381f7bb44cb7e07b4f2b", size = 170879, upload-time = "2026-05-19T09:46:25.871Z" }, ] [[package]] @@ -5217,7 +5231,7 @@ wheels = [ [[package]] name = "opentelemetry-semantic-conventions" -version = "0.62b0" +version = "0.63b0" source = { registry = "https://pypi.org/simple" } resolution-markers = [ "python_full_version >= '3.14' and sys_platform == 'win32'", @@ -5235,12 +5249,12 @@ resolution-markers = [ "python_full_version < '3.11'", ] dependencies = [ - { name = "opentelemetry-api", version = "1.41.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, + { name = "opentelemetry-api", version = "1.42.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, { name = "typing-extensions", marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/a3/b0/c14f723e86c049b7bf8ff431160d982519b97a7be2857ed2247377397a24/opentelemetry_semantic_conventions-0.62b0.tar.gz", hash = "sha256:cbfb3c8fc259575cf68a6e1b94083cc35adc4a6b06e8cf431efa0d62606c0097", size = 145753, upload-time = "2026-04-09T14:38:48.274Z" } +sdist = { url = "https://files.pythonhosted.org/packages/20/f8/be4625838aae098c2f9fbdc062a1b3128ebb9e799b891b654ee8cad94897/opentelemetry_semantic_conventions-0.63b0.tar.gz", hash = "sha256:cfea295264654fa324fcef24aa56fb1836fdc0da27db128645dc6aa76115cc6c", size = 148333, upload-time = "2026-05-19T09:46:44.01Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/58/6c/5e86fa1759a525ef91c2d8b79d668574760ff3f900d114297765eb8786cb/opentelemetry_semantic_conventions-0.62b0-py3-none-any.whl", hash = "sha256:0ddac1ce59eaf1a827d9987ab60d9315fb27aea23304144242d1fcad9e16b489", size = 231619, upload-time = "2026-04-09T14:38:32.394Z" }, + { url = "https://files.pythonhosted.org/packages/8f/6f/8d0ce225b8fdbb72c97cf4130107d861eafcb3d8e5c3f5891e8556177316/opentelemetry_semantic_conventions-0.63b0-py3-none-any.whl", hash = "sha256:1f3962732b04f43e4fef28173c9a3615b8847b4b2d6386fdc085361b29875ab9", size = 203712, upload-time = "2026-05-19T09:46:27.569Z" }, ] [[package]] @@ -5328,27 +5342,7 @@ wheels = [ name = "packaging" version = "26.0" source = { registry = "https://pypi.org/simple" } -resolution-markers = [ - "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'win32'", - "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'win32'", - "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'emscripten'", - "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'emscripten'", - "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'win32'", - "python_full_version < '3.13' and platform_machine != 's390x' and sys_platform == 'win32'", - "python_full_version < '3.13' and platform_machine == 's390x' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", - "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", - "python_full_version < '3.13' and platform_machine != 's390x' and sys_platform == 'emscripten'", - "python_full_version < '3.13' and platform_machine == 's390x' and sys_platform == 'emscripten'", - "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version < '3.13' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version < '3.13' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", -] -sdist = { url = "https://files.pythonhosted.org/packages/a1/d4/1fc4078c65507b51b96ca8f8c3ba19e6a61c8253c72794544580a7b6c24d/packaging-25.0.tar.gz", hash = "sha256:d443872c98d677bf60f6a1f2f8c1cb748e8fe762d2bf9d3148b5599295b0fc4f", size = 165727, upload-time = "2025-04-19T11:48:59.673Z" } +sdist = { url = "https://files.pythonhosted.org/packages/65/ee/299d360cdc32edc7d2cf530f3accf79c4fca01e96ffc950d8a52213bd8e4/packaging-26.0.tar.gz", hash = "sha256:00243ae351a257117b6a241061796684b084ed1c516a08c48a3f7e147a9d80b4", size = 143416, upload-time = "2026-01-21T20:50:39.064Z" } wheels = [ { url = "https://files.pythonhosted.org/packages/b7/b9/c538f279a4e237a006a2c98387d081e9eb060d203d8ed34467cc0f0b9b53/packaging-26.0-py3-none-any.whl", hash = "sha256:b36f1fef9334a5588b4166f8bcd26a14e521f2b55e6b9de3aaa80d3ff7a37529", size = 74366, upload-time = "2026-01-21T20:50:37.788Z" }, ] @@ -5358,30 +5352,41 @@ name = "pandas" version = "2.3.3" source = { registry = "https://pypi.org/simple" } resolution-markers = [ - "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'win32'", - "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'win32'", - "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'emscripten'", - "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'emscripten'", - "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'win32'", - "python_full_version < '3.13' and platform_machine != 's390x' and sys_platform == 'win32'", - "python_full_version < '3.13' and platform_machine == 's390x' and sys_platform == 'win32'", - "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'emscripten'", - "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'emscripten'", - "python_full_version < '3.13' and platform_machine != 's390x' and sys_platform == 'emscripten'", - "python_full_version < '3.13' and platform_machine == 's390x' and sys_platform == 'emscripten'", - "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version < '3.13' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", - "python_full_version < '3.13' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", + "python_full_version < '3.11' and platform_machine != 's390x' and extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts'", + "python_full_version < '3.11' and platform_machine == 's390x' and extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version >= '3.14' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version >= '3.14' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.13.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.12.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.13.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.12.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.11.*' and platform_machine != 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version < '3.11' and platform_machine != 's390x' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version < '3.11' and platform_machine == 's390x' and extra == 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", + "python_full_version < '3.11' and extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts'", ] dependencies = [ - { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, - { name = "python-dateutil", marker = "python_full_version < '3.11'" }, - { name = "pytz", marker = "python_full_version < '3.11'" }, - { name = "tzdata", marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or extra == 'extra-13-megatron-core-dev'" }, + { name = "python-dateutil", marker = "python_full_version < '3.11' or extra == 'extra-13-megatron-core-dev'" }, + { name = "pytz", marker = "python_full_version < '3.11' or extra == 'extra-13-megatron-core-dev'" }, + { name = "tzdata", marker = "python_full_version < '3.11' or extra == 'extra-13-megatron-core-dev'" }, ] sdist = { url = "https://files.pythonhosted.org/packages/33/01/d40b85317f86cf08d853a4f495195c73815fdf205eef3993821720274518/pandas-2.3.3.tar.gz", hash = "sha256:e05e1af93b977f7eafa636d043f9f94c7ee3ac81af99c13508215942e64c993b", size = 4495223, upload-time = "2025-09-29T23:34:51.853Z" } wheels = [ @@ -6241,7 +6246,7 @@ name = "pyjwt" version = "2.12.1" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c2/27/a3b6e5bf6ff856d2509292e95c8f57f0df7017cf5394921fc4e4ef40308a/pyjwt-2.12.1.tar.gz", hash = "sha256:c74a7a2adf861c04d002db713dd85f84beb242228e671280bf709d765b03672b", size = 102564, upload-time = "2026-03-13T19:27:37.25Z" } wheels = [ @@ -6624,10 +6629,10 @@ default = [ { name = "grpcio" }, { name = "opencensus" }, { name = "opentelemetry-exporter-prometheus", version = "0.54b1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-13-megatron-core-dev' or extra == 'extra-13-megatron-core-lts'" }, - { name = "opentelemetry-exporter-prometheus", version = "0.62b0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, + { name = "opentelemetry-exporter-prometheus", version = "0.63b0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, { name = "opentelemetry-proto" }, { name = "opentelemetry-sdk", version = "1.33.1", source = { registry = "https://pypi.org/simple" }, marker = "extra == 'extra-13-megatron-core-dev' or extra == 'extra-13-megatron-core-lts'" }, - { name = "opentelemetry-sdk", version = "1.41.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, + { name = "opentelemetry-sdk", version = "1.42.0", source = { registry = "https://pypi.org/simple" }, marker = "(extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (extra != 'extra-13-megatron-core-dev' and extra != 'extra-13-megatron-core-lts')" }, { name = "prometheus-client" }, { name = "py-spy" }, { name = "pydantic" }, @@ -6643,7 +6648,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "attrs" }, { name = "rpds-py" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/22/f5/df4e9027acead3ecc63e50fe1e36aca1523e1719559c499951bb4b53188f/referencing-0.37.0.tar.gz", hash = "sha256:44aefc3142c5b842538163acb373e24cce6632bd54bdb01b21ad5863489f50d8", size = 78036, upload-time = "2025-10-13T15:30:48.871Z" } wheels = [ @@ -7859,7 +7864,7 @@ version = "0.52.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "anyio" }, - { name = "typing-extensions", marker = "python_full_version < '3.13'" }, + { name = "typing-extensions", marker = "python_full_version < '3.13' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/c4/68/79977123bb7be889ad680d79a40f339082c1978b5cfcf62c2d8d196873ac/starlette-0.52.1.tar.gz", hash = "sha256:834edd1b0a23167694292e94f597773bc3f89f362be6effee198165a35d62933", size = 2653702, upload-time = "2026-01-18T13:34:11.062Z" } wheels = [ @@ -7871,7 +7876,7 @@ name = "sympy" version = "1.14.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "mpmath" }, + { name = "mpmath", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (python_full_version < '3.13' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-lts') or (python_full_version < '3.13' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and sys_platform == 'win32' and extra == 'extra-13-megatron-core-lts') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/83/d3/803453b36afefb7c2bb238361cd4ae6125a569b4db67cd9e79846ba2d68c/sympy-1.14.0.tar.gz", hash = "sha256:d3d3fe8df1e5a0b42f0e7bdf50541697dbe7d23746e894990c030e2b05e72517", size = 7793921, upload-time = "2025-04-27T18:05:01.611Z" } wheels = [ @@ -7963,7 +7968,7 @@ resolution-markers = [ dependencies = [ { name = "ml-dtypes", version = "0.4.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.13' and extra == 'extra-13-megatron-core-dev') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11'" }, + { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-lts') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-dev') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/3c/b9/ea25aba62c688a87d7d7d9cc5926d602e2f9e84fa72586825486fb180b7e/tensorstore-0.1.74.tar.gz", hash = "sha256:a062875f27283d30ce4959c408c253ecb336fce8e3f9837c064e3d30cda79203", size = 6795605, upload-time = "2025-04-24T15:42:18.829Z" } wheels = [ @@ -8020,7 +8025,7 @@ resolution-markers = [ "python_full_version == '3.11.*' and platform_machine == 's390x' and sys_platform != 'emscripten' and sys_platform != 'win32'", ] dependencies = [ - { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, + { name = "ml-dtypes", version = "0.5.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-lts') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (python_full_version >= '3.11' and extra != 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.0.2", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and python_full_version < '3.13' and extra == 'extra-13-megatron-core-dev') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (python_full_version >= '3.13' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and extra == 'extra-13-megatron-core-lts') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] @@ -8258,21 +8263,21 @@ name = "torch" version = "2.11.0" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "cuda-bindings", marker = "sys_platform == 'linux'" }, + { name = "cuda-bindings", marker = "sys_platform == 'linux' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "cuda-toolkit", extra = ["cublas", "cudart", "cufft", "cufile", "cupti", "curand", "cusolver", "cusparse", "nvjitlink", "nvrtc", "nvtx"], marker = "sys_platform == 'linux' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "filelock", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "fsspec", version = "2026.2.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "jinja2", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "filelock", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "fsspec", version = "2026.2.0", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "jinja2", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "networkx", version = "3.4.2", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "networkx", version = "3.6.1", source = { registry = "https://pypi.org/simple" }, marker = "(python_full_version >= '3.11' and sys_platform != 'emscripten' and sys_platform != 'win32') or (python_full_version < '3.11' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, - { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux'" }, - { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux'" }, - { name = "setuptools", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "sympy", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, - { name = "triton", marker = "sys_platform == 'never'" }, - { name = "typing-extensions", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32')" }, + { name = "nvidia-cudnn-cu13", marker = "sys_platform == 'linux' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "nvidia-cusparselt-cu13", marker = "sys_platform == 'linux' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "nvidia-nccl-cu13", marker = "sys_platform == 'linux' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "nvidia-nvshmem-cu13", marker = "sys_platform == 'linux' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "setuptools", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "sympy", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "triton", marker = "sys_platform == 'never' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "typing-extensions", marker = "(python_full_version < '3.11' and sys_platform == 'emscripten') or (python_full_version < '3.11' and sys_platform == 'win32') or (sys_platform != 'emscripten' and sys_platform != 'win32') or (sys_platform == 'emscripten' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts') or (sys_platform == 'win32' and extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] wheels = [ { url = "https://files.pythonhosted.org/packages/ac/f2/c1690994afe461aae2d0cac62251e6802a703dec0a6c549c02ecd0de92a9/torch-2.11.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2c0d7fcfbc0c4e8bb5ebc3907cbc0c6a0da1b8f82b1fc6e14e914fa0b9baf74e", size = 80526521, upload-time = "2026-03-23T18:12:06.86Z" }, @@ -8331,7 +8336,7 @@ name = "tqdm" version = "4.67.3" source = { registry = "https://pypi.org/simple" } dependencies = [ - { name = "colorama", marker = "sys_platform == 'win32'" }, + { name = "colorama", marker = "sys_platform == 'win32' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/09/a9/6ba95a270c6f1fbcd8dac228323f2777d886cb206987444e4bce66338dd4/tqdm-4.67.3.tar.gz", hash = "sha256:7d825f03f89244ef73f1d4ce193cb1774a8179fd96f31d7e1dcde62092b960bb", size = 169598, upload-time = "2026-02-03T17:35:53.048Z" } wheels = [ @@ -8510,7 +8515,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, { name = "h11" }, - { name = "typing-extensions", marker = "python_full_version < '3.11'" }, + { name = "typing-extensions", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/62/f2/368268300fb8af33743508d738ef7bb4d56afdb46c6d9c0fa3dd515df171/uvicorn-0.43.0.tar.gz", hash = "sha256:ab1652d2fb23abf124f36ccc399828558880def222c3cb3d98d24021520dc6e8", size = 85686, upload-time = "2026-04-03T18:37:48.984Z" } wheels = [