diff --git a/megatron/core/ssm/gated_delta_net.py b/megatron/core/ssm/gated_delta_net.py deleted file mode 100644 index 7e28691c15c..00000000000 --- a/megatron/core/ssm/gated_delta_net.py +++ /dev/null @@ -1,1453 +0,0 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. -# Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh. - -# Some of this code was adopted from https://github.com/huggingface/transformers -# This source code is licensed under the Apache license found in the -# LICENSE file in the root directory of this source tree. - -import logging -from dataclasses import dataclass -from functools import lru_cache -from typing import Optional, Union - -import torch -import torch.nn as nn -import torch.nn.functional as F -from torch import Tensor - -from megatron.core import tensor_parallel -from megatron.core.context_parallel_layout import ( - contiguous_to_zigzag_chunks, - zigzag_to_contiguous_chunks, -) -from megatron.core.fp8_utils import get_fp8_align_size -from megatron.core.inference.contexts import BaseInferenceContext -from megatron.core.jit import jit_fuser -from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group -from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.ssm.mamba_context_parallel import ( - _all_to_all_cp2hp, - _all_to_all_hp2cp, - _redo_attention_load_balancing, - _undo_attention_load_balancing, -) -from megatron.core.ssm.utils import _split_tensor_factory -from megatron.core.tensor_parallel import get_cuda_rng_tracker -from megatron.core.transformer import TransformerConfig -from megatron.core.transformer.identity_op import IdentityOp -from megatron.core.transformer.module import MegatronModule -from megatron.core.transformer.spec_utils import ModuleSpec, build_module -from megatron.core.transformer.utils import ( - ensure_metadata_has_dp_cp_group, - make_sharded_tensors_for_checkpoint, - sharded_state_dict_default, -) -from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push - -try: - from fla.modules.convolution import causal_conv1d - from fla.modules.l2norm import l2norm - from fla.ops.cp import build_cp_context - from fla.ops.gated_delta_rule import chunk_gated_delta_rule - - HAVE_FLA = True -except ImportError: - causal_conv1d = None - l2norm = None - chunk_gated_delta_rule = None - - HAVE_FLA = False - - -logger = logging.getLogger(__name__) - - -@dataclass -class GatedDeltaNetSubmodules: - """ - Contains the module specs for the input linear, output norm, and output linear layers. - """ - - in_proj: Union[ModuleSpec, type] = IdentityOp - out_norm: Union[ModuleSpec, type] = IdentityOp - out_proj: Union[ModuleSpec, type] = IdentityOp - - -class GatedDeltaNet(MegatronModule): - """Gated Delta Net (GDN) layer class - - GDN layer takes input with size [s, b, h] - and returns output of the same size. - """ - - def __init__( - self, - config: TransformerConfig, - submodules: GatedDeltaNetSubmodules, - layer_number: int = None, - bias: bool = False, - conv_bias: bool = False, - conv_init: Optional[float] = None, - use_qk_l2norm: bool = True, - A_init_range: tuple[float, float] = (1, 16), - pg_collection: ProcessGroupCollection = None, - name: str | None = None, - **kwargs, - ): - """ - Args: - config: The config of the model. - submodules: Contains the module specs for the input and output linear layers. - layer_number: The layer number of this GDN layer. - bias: Whether to use bias in the linear layers. - conv_bias: Whether to use bias in the causal convolution. - conv_init: The initialization range for the causal convolution weights. - use_qk_l2norm: Whether to use L2 normalization in the kernel of the gated delta rule. - A_init_range: The initialization range for the attention weights. - pg_collection: The required process groups to use for tensor model parallel and context - parallel. - name (str | None): module instance name passed top-down from its paranet module - """ - - if not HAVE_FLA: - raise ImportError( - "FLA is not installed. Please install it with `pip install flash-linear-attention`." - ) - - super().__init__(config) - - # Attributes from arguments - self.layer_number = layer_number - self.bias = bias - self.conv_bias = conv_bias - self.conv_init = conv_init - assert A_init_range[0] >= 0 and A_init_range[1] >= A_init_range[0] - self.A_init_range = A_init_range - self.use_qk_l2norm = use_qk_l2norm - assert pg_collection is not None, "pg_collection must be provided for GatedDeltaNet" - self.pg_collection = pg_collection - self.tp_group = pg_collection.tp - self.cp_size = self.pg_collection.cp.size() - self.tp_size = self.pg_collection.tp.size() - self.sp_size = self.tp_size if config.sequence_parallel else 1 - self.gdn_pre_gated_delta_rule_fusion = config.gdn_pre_gated_delta_rule_fusion - - # Attributes from config - self.config = config - if self.config.deterministic_mode: - if self.gdn_pre_gated_delta_rule_fusion: - raise ValueError( - "Pre-GDR fusion is non-deterministic, but deterministic_mode=True. " - "Disable gdn_pre_gated_delta_rule_fusion or deterministic_mode." - ) - self.hidden_size = config.hidden_size - self.act_fn = config.activation_func - self.activation = self.act_fn.__name__ - self.conv_kernel_dim = config.linear_conv_kernel_dim - self.key_head_dim = config.linear_key_head_dim - self.value_head_dim = config.linear_value_head_dim - self.num_key_heads = config.linear_num_key_heads - self.num_value_heads = config.linear_num_value_heads - self.qk_dim = self.key_head_dim * self.num_key_heads - self.v_dim = self.value_head_dim * self.num_value_heads - self.qk_dim_local_tp = self.qk_dim // self.tp_size - self.v_dim_local_tp = self.v_dim // self.tp_size - - # Headwise CP uses head-parallel layout: each CP rank handles a slice of - # heads. The static cp_size (== max dynamic cp_size) must evenly divide - # the per-TP head counts so that every possible runtime cp_size also - # divides. Chunkwise CP keeps heads local and does not need this split. - if self.config.linear_cp_mode == "headwise": - num_key_heads_per_tp = self.num_key_heads // self.tp_size - num_value_heads_per_tp = self.num_value_heads // self.tp_size - assert num_key_heads_per_tp % self.cp_size == 0, ( - f"GDN head-parallel CP requires the static (max) cp_size ({self.cp_size}) " - f"to evenly divide num_key_heads per TP rank ({num_key_heads_per_tp}); " - f"all runtime dynamic cp_size values divide the static one and so will also divide." - ) - assert num_value_heads_per_tp % self.cp_size == 0, ( - f"GDN head-parallel CP requires the static (max) cp_size ({self.cp_size}) " - f"to evenly divide num_value_heads per TP rank ({num_value_heads_per_tp}); " - f"all runtime dynamic cp_size values divide the static one and so will also divide." - ) - - # Input projection (hidden_states -> q, k, v, gate, beta, alpha) - # TODO: for now, output gate is forced for GDN. - # We may remove this restriction in the future. - self.in_proj_dim = self.qk_dim * 2 + self.v_dim * 2 + self.num_value_heads * 2 - if self.config.fp8: - fp8_align_size = get_fp8_align_size(self.config.fp8_recipe) - assert self.in_proj_dim % fp8_align_size == 0, ( - "For FP8, the innermost dimension of the GDN layer " - "input projection output tensor must be a multiple of 16." - ) - self.in_proj = build_module( - submodules.in_proj, - self.hidden_size, - self.in_proj_dim, - config=self.config, - init_method=self.config.init_method, - gather_output=False, - bias=bias, - skip_bias_add=False, - is_expert=False, - tp_comm_buffer_name="fc1", - tp_group=self.pg_collection.tp, - name=(name + ".in_proj") if name is not None else None, - ) - - # Conv1d for QKV - self.conv_dim = self.qk_dim * 2 + self.v_dim - self.conv_dim_local_tp = self.conv_dim // self.tp_size - - # weight shape: [conv_dim, 1, d_conv] - # bias shape: [conv_dim] - self.conv1d = nn.Conv1d( - in_channels=self.conv_dim_local_tp, - out_channels=self.conv_dim_local_tp, - bias=conv_bias, - kernel_size=self.conv_kernel_dim, - groups=self.conv_dim_local_tp, - padding=self.conv_kernel_dim - 1, - device=torch.cuda.current_device(), - dtype=config.params_dtype, - ) - setattr(self.conv1d.weight, "tensor_model_parallel", True) - setattr(self.conv1d.weight, "partition_dim", 0) - if conv_bias: - setattr(self.conv1d.bias, "tensor_model_parallel", True) - setattr(self.conv1d.bias, "partition_dim", 0) - - # Time step projection (discretization) - self.num_v_heads_local_tp = self.num_value_heads // self.tp_size - # dt_bias parameter - self.dt_bias = nn.Parameter( - torch.empty( - self.num_v_heads_local_tp, - dtype=config.params_dtype, - device=torch.cuda.current_device(), - ) - ) - setattr(self.dt_bias, "tensor_model_parallel", True) - setattr(self.dt_bias, "partition_dim", 0) - # A_log parameter - self.A_log = nn.Parameter( - torch.empty( - self.num_v_heads_local_tp, - dtype=config.params_dtype, - device=torch.cuda.current_device(), - ) - ) - setattr(self.A_log, "tensor_model_parallel", True) - setattr(self.A_log, "partition_dim", 0) - - if self.config.deterministic_mode: - self.gated_delta_rule = torch_chunk_gated_delta_rule - else: - self.gated_delta_rule = chunk_gated_delta_rule - - # Output layernorm before projection - self.out_norm = build_module( - submodules.out_norm, - config=self.config, - hidden_size=self.value_head_dim, - eps=self.config.layernorm_epsilon, - ) - - self.out_proj = build_module( - submodules.out_proj, - self.v_dim, - self.hidden_size, - config=self.config, - init_method=self.config.output_layer_init_method, - bias=bias, - input_is_parallel=True, - skip_bias_add=True, - is_expert=False, - tp_comm_buffer_name="fc2", - tp_group=self.pg_collection.tp, - name=(name + ".out_proj") if name is not None else None, - ) - - # Whole-module recompute: when "gdn" is in recompute_modules (selective granularity), - # the entire GatedDeltaNet compute is wrapped in a normal checkpoint and recomputed - # in the backward pass. - self.recompute_gdn = False - if self.config.recompute_granularity == "selective" and self.config.recompute_modules: - self.recompute_gdn = "gdn" in self.config.recompute_modules - - # Cache for CP context objects consumed by FLA kernels. Rebuilding these per-forward - # is unsafe under CUDA graph capture because build_cp_context allocates - # fresh tensors whose memory pointers are baked into the captured graph; - # on the next call those tensors are reallocated, leaving the replayed - # graph pointing at stale memory. For non-packed (SBHD) input the - # cu_seqlens is fully determined by the (static) global sequence length - # and batch size, so we cache the (cu_seqlens, cp_context) pair keyed on - # both values. - self._chunkwise_cp_context_cache = {} - - self.reset_parameters() - - def reset_parameters(self): - """Reset the parameters.""" - if self.config.perform_initialization: - with get_cuda_rng_tracker().fork(): - # conv1d.weight - if self.conv_init is not None: - nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init) - # dt_bias - torch.ones( - self.num_v_heads_local_tp, - out=self.dt_bias.data, - dtype=self.config.params_dtype, - device=torch.cuda.current_device(), - ) - # A_log - A = torch.empty( - self.num_v_heads_local_tp, - dtype=self.config.params_dtype, - device=torch.cuda.current_device(), - ).uniform_(*self.A_init_range) - self.A_log.data.copy_(torch.log(A)) - - def forward( - self, - hidden_states: Tensor, - attention_mask: Tensor, - inference_context: Optional[BaseInferenceContext] = None, - packed_seq_params: Optional[PackedSeqParams] = None, - sequence_len_offset: Optional[int] = None, - *, - pg_collection: Optional[ProcessGroupCollection] = None, - inference_params: Optional[BaseInferenceContext] = None, - **kwargs, - ): - """ - Perform a forward pass through the GDN module. - - Args: - hidden_states (Tensor): Hidden states. - attention_mask (Tensor): Attention mask. - inference_context (Optional[BaseInferenceContext]): Inference context that manages - KV cache. - packed_seq_params (Optional[PackedSeqparams]): Parameters used for THD format. - sequence_len_offset (Optional[int]): Sequence length offset used for - inference CUDA graphs. - - Return: - (tuple[Tensor, Tensor]) GDN output and bias. - - """ - # TODO: Deal with attention_mask - - inference_context = deprecate_inference_params(inference_context, inference_params) - - # Route the CP group to either the headwise (Ulysses-style) path or the - # chunkwise CP path according to config.linear_cp_mode. The two paths - # are mutually exclusive — whichever one is active owns the full CP - # group, and the other is given a size-1 group (None). The unused-path - # helpers already treat a None group as size 1, avoiding a costly and - # CUDA-graph-unsafe `torch.distributed.new_group` on every forward. - base_cp_group = pg_collection.cp if pg_collection is not None else self.pg_collection.cp - cp_group = resolve_cp_group(base_cp_group, packed_seq_params) - if self.config.linear_cp_mode == "chunkwise": - cp_group_chunkwise = cp_group - cp_group_headwise = None - elif self.config.linear_cp_mode == "headwise": - cp_group_chunkwise = None - cp_group_headwise = cp_group - elif cp_group.size() == 1: - cp_group_chunkwise = None - cp_group_headwise = None - else: - raise ValueError( - f"Unsupported linear_cp_mode {self.config.linear_cp_mode!r}; " - "expected 'headwise' or 'chunkwise'." - ) - cp_size_chunkwise = cp_group_chunkwise.size() if cp_group_chunkwise is not None else 1 - cp_size_headwise = cp_group_headwise.size() if cp_group_headwise is not None else 1 - - seq_len_local, batch, _ = hidden_states.shape - seq_len_post_headwise = seq_len_local * self.sp_size * cp_size_headwise - seq_len_global = seq_len_post_headwise * cp_size_chunkwise - - if inference_context is not None: - assert ( - inference_context.is_static_batching() - ), "GDN does not currently support dynamic inference batching." - assert not self.config.sequence_parallel - # TODO: support inference - raise NotImplementedError("GDN does not support inference for now.") - - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - assert batch == 1, "Packed sequence expects batch dimension to be 1" - assert ( - not self.config.deterministic_mode - ), "Packed sequence does not support deterministic mode." - - # Resolve cu_seqlens with alignment padding handling. - # cu_seqlens in packed_seq_params is the global (pre-CP-split) cu_seqlens, so we - # validate against the global sequence length. - cu_seqlens_q = self._resolve_cu_seqlens( - packed_seq_params.cu_seqlens_q_padded, - packed_seq_params.cu_seqlens_q, - seq_len_global, - "cu_seqlens_q", - cp_size=self.cp_size, - ) - cu_seqlens_kv = self._resolve_cu_seqlens( - packed_seq_params.cu_seqlens_kv_padded, - packed_seq_params.cu_seqlens_kv, - seq_len_global, - "cu_seqlens_kv", - cp_size=self.cp_size, - ) - assert torch.equal(cu_seqlens_q, cu_seqlens_kv), ( - "Currently only support cu_seqlens_q equals to cu_seqlens_kv, " - f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}" - ) - num_packed_seqs = cu_seqlens_q.shape[0] - 1 - assert num_packed_seqs > 0, ( - "Number of packed sequences must be greater than 0, " - f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}" - ) - else: - cu_seqlens_q = None - cu_seqlens_kv = None - - if cp_size_chunkwise > 1: - if cu_seqlens_q is None: - # Non-packed input: the only sources of cu_seqlens are the static - # global sequence length and batch size. Cache both the cu_seqlens - # tensor and the resulting chunkwise CP context so we don't - # reallocate them on every forward — those reallocations break - # CUDA graph capture. - cache_key = (seq_len_global, batch) - cached = self._chunkwise_cp_context_cache.get(cache_key) - if cached is None: - cached_cu_seqlens = ( - torch.arange( - batch + 1, device=torch.cuda.current_device(), dtype=torch.long - ) - * seq_len_global - ) - cached_ctx = build_cp_context( - cu_seqlens=cached_cu_seqlens, - group=cp_group_chunkwise, - conv1d_kernel_size=self.conv_kernel_dim, - ) - cached = (cached_cu_seqlens, cached_ctx) - self._chunkwise_cp_context_cache[cache_key] = cached - cu_seqlens_q, chunkwise_cp_context = cached - else: - chunkwise_cp_context = build_cp_context( - cu_seqlens=cu_seqlens_q, - group=cp_group_chunkwise, - conv1d_kernel_size=self.conv_kernel_dim, - ) - else: - chunkwise_cp_context = None - - if self.recompute_gdn and self.training: - - def _checkpointed_compute(hidden_states): - return self._forward_compute( - hidden_states, - batch, - seq_len_post_headwise, - cp_size_headwise, - cp_group_headwise, - cp_size_chunkwise, - cp_group_chunkwise, - cu_seqlens_q, - packed_seq_params, - chunkwise_cp_context, - ) - - out, out_bias = tensor_parallel.checkpoint(_checkpointed_compute, False, hidden_states) - else: - out, out_bias = self._forward_compute( - hidden_states, - batch, - seq_len_post_headwise, - cp_size_headwise, - cp_group_headwise, - cp_size_chunkwise, - cp_group_chunkwise, - cu_seqlens_q, - packed_seq_params, - chunkwise_cp_context, - ) - - return out, out_bias - - def _forward_compute( - self, - hidden_states, - batch, - seq_len_post_headwise, - cp_size_headwise, - cp_group_headwise, - cp_size_chunkwise, - cp_group_chunkwise, - cu_seqlens_q, - packed_seq_params, - chunkwise_cp_context, - ): - """Core GDN computation (in_proj -> conv1d -> gated_delta_rule -> gated norm -> out_proj). - - Extracted from ``forward`` so the entire module can be wrapped in a recompute - checkpoint when ``recompute_modules`` contains ``"gdn"`` (selective full-module - recompute, normal checkpointing). - - Returns: - Tuple of (output, output_bias). - """ - # Input projection - nvtx_range_push(suffix="in_proj") - qkvzba, _ = self.in_proj(hidden_states) - nvtx_range_pop(suffix="in_proj") - - # Chunkwise CP expects the contiguous-time chunk layout (rank r holds chunks - # [2r, 2r+1]) inside conv1d / chunk_gated_delta_rule. Megatron attention CP - # feeds us the zigzag attention-load-balanced layout (rank r holds - # [r, 2*cp-r-1]), so reshuffle chunks over the CP group with a single - # all-to-all — no full-sequence gather required. - # TODO: Move CP layout ownership to a model/region-level scheduler so hybrid models can - # enter contiguous layout before GDN regions instead of paying module-local conversions. - if cp_size_chunkwise > 1: - nvtx_range_push(suffix="zigzag_to_contiguous") - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - qkvzba = zigzag_to_contiguous_chunks( - qkvzba, cp_group_chunkwise, seq_dim=0, cu_seqlens=cu_seqlens_q - ) - else: - qkvzba = zigzag_to_contiguous_chunks(qkvzba, cp_group_chunkwise, seq_dim=0) - nvtx_range_pop(suffix="zigzag_to_contiguous") - - qkvzba, thd_cp_a2a_inv = self._a2a_cp_to_hp( - qkvzba, - cp_size_headwise, - cp_group_headwise, - cu_seqlens_q, - seq_len_post_headwise, - packed_seq_params, - ) - - if self.gdn_pre_gated_delta_rule_fusion: - if cp_size_chunkwise > 1 and batch > 1: - raise ValueError( - "GDN chunkwise CP with SBHD inputs currently requires micro_batch_size == 1 " - "because the FLA gated delta rule backend requires a single batch dimension " - "when cp_context is used. Use packed THD input or micro_batch_size=1." - ) - if cp_size_chunkwise > 1 and self.config.gdn_conv_pad_alignment is not None: - raise ValueError( - "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " - "chunk-local causal-conv inputs can change later chunk numerics." - ) - nvtx_range_push(suffix="fused_streamed_pre_gated_delta_rule") - seq_idx = ( - packed_seq_params.seq_idx - if packed_seq_params is not None - and packed_seq_params.qkv_format == 'thd' - and cp_size_chunkwise == 1 - else None - ) - fused_cu_seqlens_q = ( - cu_seqlens_q - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' - else None - ) - query, key, value, gate, beta, g = self._fused_streamed_pre_gated_delta_rule( - qkvzba, - cu_seqlens_q=fused_cu_seqlens_q, - seq_idx=seq_idx, - cp_group=cp_group_chunkwise if cp_size_chunkwise > 1 else None, - cp_size_headwise=cp_size_headwise, - cp_group_headwise=cp_group_headwise, - ) - nvtx_range_pop(suffix="fused_streamed_pre_gated_delta_rule") - else: - nvtx_range_push(suffix="pre_gated_delta_rule") - if cp_size_chunkwise > 1 and packed_seq_params is None and batch > 1: - # TODO: If additional gated delta rule backends are added, handle this - # SBHD + chunkwise CP + batch>1 case per backend instead of - # unconditionally rejecting it. - raise ValueError( - "GDN chunkwise CP with SBHD inputs currently requires micro_batch_size == 1 " - "because the FLA gated delta rule backend requires a single batch dimension " - "when cp_context is used. Use packed THD input or micro_batch_size=1." - ) - if cp_size_chunkwise > 1 and self.config.gdn_conv_pad_alignment is not None: - raise ValueError( - "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " - "chunk-local causal-conv inputs can change later chunk numerics." - ) - query, key, value, gate, beta, g = self.pre_gated_delta_rule( - qkvzba, - batch, - seq_len_post_headwise, - cp_size_headwise, - cp_group_headwise, - cu_seqlens_q, - chunkwise_cp_context, - packed_seq_params=packed_seq_params, - ) - nvtx_range_pop(suffix="pre_gated_delta_rule") - - nvtx_range_push(suffix="gated_delta_rule") - core_attn_out, _ = self.gated_delta_rule( - query, - key, - value, - g=g, - beta=beta, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=False, - cu_seqlens=cu_seqlens_q, - cp_context=chunkwise_cp_context, - ) - nvtx_range_pop(suffix="gated_delta_rule") - - # RMSNorm - nvtx_range_push(suffix="gated_norm") - norm_out = self._apply_gated_norm(core_attn_out, gate) - nvtx_range_pop(suffix="gated_norm") - - # Transpose: b s x --> s b x - # From bshd back to sbhd format - norm_out = norm_out.reshape(batch, seq_len_post_headwise, -1) - norm_out = norm_out.transpose(0, 1).contiguous() - - # Inverse of the zigzag -> contiguous reshuffle performed before conv1d. - # Restores the Megatron attention-load-balanced layout that downstream - # layers and loss computation expect. - # TODO: The planned CP layout refactor should keep consecutive GDN layers contiguous and - # restore zigzag only at SDPA/canonical-layout boundaries. - if cp_size_chunkwise > 1: - nvtx_range_push(suffix="contiguous_to_zigzag") - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - norm_out = contiguous_to_zigzag_chunks( - norm_out, cp_group=cp_group_chunkwise, seq_dim=0, cu_seqlens=cu_seqlens_q - ) - else: - norm_out = contiguous_to_zigzag_chunks( - norm_out, cp_group=cp_group_chunkwise, seq_dim=0 - ) - nvtx_range_pop(suffix="contiguous_to_zigzag") - - norm_out = self._a2a_hp_to_cp( - norm_out, cp_size_headwise, cp_group_headwise, packed_seq_params, thd_cp_a2a_inv - ) - - # Output projection - nvtx_range_push(suffix="out_proj") - out, out_bias = self.out_proj(norm_out) - nvtx_range_pop(suffix="out_proj") - - return out, out_bias - - def pre_gated_delta_rule( - self, - qkvzba, - batch, - seq_len, - cp_size_headwise, - cp_group_headwise, - cu_seqlens_q=None, - chunkwise_cp_context=None, - packed_seq_params=None, - ): - """Prepare QKV, gate, beta, and decay tensors before the gated delta rule.""" - - # Transpose: s b x --> b s x - # From sbhd to bshd format - qkvzba = qkvzba.transpose(0, 1) - - # Split, reorder, and reshape the tensor into q, k, v, gate, beta, alpha - qkv, gate, beta, alpha = torch.split( - qkvzba, - [ - (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // cp_size_headwise, - self.v_dim_local_tp // cp_size_headwise, - self.num_value_heads // self.tp_size // cp_size_headwise, - self.num_value_heads // self.tp_size // cp_size_headwise, - ], - dim=-1, - ) - gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) - beta = beta.reshape(batch, seq_len, -1) - alpha = alpha.reshape(batch, seq_len, -1) - - kernel_batch = batch - kernel_seq_len = seq_len - - # Convolution on qkv - nvtx_range_push(suffix="conv1d") - assert ( - qkv.shape[1] == kernel_seq_len - ), f"Shape mismatch: {qkv.shape[1]=} != {kernel_seq_len=}" - qkv_channels_split_sections = [ - self.qk_dim_local_tp, - self.qk_dim_local_tp, - self.v_dim_local_tp, - ] - conv1d_weight = get_parameter_local_cp_headwise( - self.conv1d.weight, - dim=0, - cp_group=cp_group_headwise, - split_sections=qkv_channels_split_sections, - ) - conv1d_bias = ( - get_parameter_local_cp_headwise( - self.conv1d.bias, - dim=0, - cp_group=cp_group_headwise, - split_sections=qkv_channels_split_sections, - ) - if self.conv_bias - else None - ) - if self.config.deterministic_mode: - qkv = qkv.transpose(1, 2).contiguous() # b, s, d -> b, d, s - conv_out = F.conv1d( - input=qkv, # Torch-native only accept [b, d, s] format input - weight=conv1d_weight, - bias=conv1d_bias, - stride=self.conv1d.stride, - padding=self.conv1d.padding, - dilation=self.conv1d.dilation, - groups=self.conv_dim_local_tp // cp_size_headwise, - ) - qkv = self.act_fn(conv_out[..., :kernel_seq_len]) - qkv = qkv.transpose(1, 2) # b, d, s -> b, s, d - else: - assert self.activation in ["silu", "swish"] - _orig_seq = qkv.shape[1] - _pad_n = 0 - _conv_input = qkv.contiguous() - _conv_cu_seqlens = cu_seqlens_q - _conv_cp_context = chunkwise_cp_context - if self.config.gdn_conv_pad_alignment is not None: - if packed_seq_params is None or cu_seqlens_q is None: - raise ValueError( - "gdn_conv_pad_alignment is only supported with packed sequence " - "parameters in THD format. SBHD inputs do not need causal-conv padding." - ) - if chunkwise_cp_context is not None: - raise ValueError( - "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " - "chunk-local causal-conv inputs can change later chunk numerics." - ) - _pad_n = -_orig_seq % self.config.gdn_conv_pad_alignment - if _pad_n > 0: - _conv_input = torch.nn.functional.pad(_conv_input, (0, 0, 0, _pad_n)) - # cu_seqlens_q is None in non-packed-sequence mode; only the - # last-segment offset needs to grow to cover the padding tail. - if cu_seqlens_q is not None: - _conv_cu_seqlens = cu_seqlens_q.clone() - _conv_cu_seqlens[-1] += _pad_n - qkv, _ = causal_conv1d( - x=_conv_input, # FLA conv1d accepts [b, s, d] format input - weight=conv1d_weight.squeeze(1), # d, 1, w -> d, w - bias=conv1d_bias, - activation=self.activation, - initial_state=None, - output_final_state=False, - cu_seqlens=_conv_cu_seqlens, - cp_context=_conv_cp_context, - ) - if _pad_n > 0: - qkv = qkv[:, :_orig_seq, :] - nvtx_range_pop(suffix="conv1d") - - # Prepare QKV tensors (split, reshape, L2 norm, repeat_interleave, contiguous) - nvtx_range_push(suffix="prepare_qkv_for_gated_delta_rule") - query, key, value, gate, beta, alpha = self._prepare_qkv_for_gated_delta_rule( - qkv, gate, beta, alpha, kernel_batch, kernel_seq_len, cp_size_headwise=cp_size_headwise - ) - nvtx_range_pop(suffix="prepare_qkv_for_gated_delta_rule") - - # Calculate g and beta - nvtx_range_push(suffix="g_and_beta") - A_log_local_cp = get_parameter_local_cp_headwise( - self.A_log, dim=0, cp_group=cp_group_headwise - ) - dt_bias_local_cp = get_parameter_local_cp_headwise( - self.dt_bias, dim=0, cp_group=cp_group_headwise - ) - g, beta = self._compute_g_and_beta(A_log_local_cp, dt_bias_local_cp, alpha, beta) - nvtx_range_pop(suffix="g_and_beta") - - return query, key, value, gate, beta, g - - def _fused_streamed_pre_gated_delta_rule( - self, - qkvzba, - cu_seqlens_q=None, - seq_idx=None, - cp_group=None, - cp_size_headwise=1, - cp_group_headwise=None, - ): - """Call the streamed fused pre-GDR wrapper.""" - - try: - from megatron.core.fusions.fused_pre_gated_delta_rule import ( - fused_streamed_pre_gated_delta_rule, - ) - except ImportError as exc: - raise ImportError( - "gdn_pre_gated_delta_rule_fusion requires the streamed pre-GDR fusion " - "dependencies, including causal-conv1d." - ) from exc - - qkv_channels_split_sections = [ - self.qk_dim_local_tp, - self.qk_dim_local_tp, - self.v_dim_local_tp, - ] - conv1d_weight = get_parameter_local_cp_headwise( - self.conv1d.weight, - dim=0, - cp_group=cp_group_headwise, - split_sections=qkv_channels_split_sections, - ) - conv1d_bias = ( - get_parameter_local_cp_headwise( - self.conv1d.bias, - dim=0, - cp_group=cp_group_headwise, - split_sections=qkv_channels_split_sections, - ) - if self.conv_bias - else None - ) - A_log = get_parameter_local_cp_headwise(self.A_log, dim=0, cp_group=cp_group_headwise) - dt_bias = get_parameter_local_cp_headwise(self.dt_bias, dim=0, cp_group=cp_group_headwise) - num_key_heads = self.qk_dim_local_tp // self.key_head_dim // cp_size_headwise - num_value_heads = self.v_dim_local_tp // self.value_head_dim // cp_size_headwise - - return fused_streamed_pre_gated_delta_rule( - qkvzba, - conv1d_weight, - conv1d_bias, - A_log, - dt_bias, - num_key_heads=num_key_heads, - num_value_heads=num_value_heads, - key_head_dim=self.key_head_dim, - value_head_dim=self.value_head_dim, - use_qk_l2norm=self.use_qk_l2norm, - cu_seqlens=cu_seqlens_q, - seq_idx=seq_idx, - cp_group=cp_group, - ) - - def _a2a_cp_to_hp( - self, - qkvzba: torch.Tensor, - cp_size: int, - cp_group: torch.distributed.ProcessGroup, - cu_seqlens_q: Optional[torch.Tensor], - seq_len: int, - packed_seq_params: Optional[PackedSeqParams], - ) -> tuple[torch.Tensor, Optional[torch.Tensor]]: - """Run GDN context-parallel to hidden-parallel A2A and return its inverse context.""" - if cp_size > 1: - # Pre-permute head dim so a single unsectioned a2a is equivalent to per-section a2a. - head_perm = _build_head_perm_for_split_sections( - ( - self.qk_dim_local_tp, - self.qk_dim_local_tp, - self.v_dim_local_tp, - self.v_dim_local_tp, - self.num_value_heads // self.tp_size, - self.num_value_heads // self.tp_size, - ), - cp_size, - qkvzba.device, - ) - qkvzba = qkvzba.index_select(-1, head_perm) - - thd_cp_a2a_inv = None - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - qkvzba = tensor_a2a_cp2hp( - qkvzba, - seq_dim=0, - head_dim=-1, - cp_group=cp_group, - undo_attention_load_balancing=False, - ) - if cp_size > 1: - # Permute at the seq dim so that a single unsectioned a2a - # is equivalent to per-sequence a2a. - # This also folds the ``_undo_attention_load_balancing`` step. - thd_cp_a2a_idx, thd_cp_a2a_inv = _build_thd_cp_a2a_perm( - cu_seqlens_q, cp_size, seq_len - ) - qkvzba = qkvzba.index_select(0, thd_cp_a2a_idx) - else: - qkvzba = tensor_a2a_cp2hp(qkvzba, seq_dim=0, head_dim=-1, cp_group=cp_group) - - return qkvzba, thd_cp_a2a_inv - - def _a2a_hp_to_cp( - self, - norm_out: torch.Tensor, - cp_size: int, - cp_group: torch.distributed.ProcessGroup, - packed_seq_params: Optional[PackedSeqParams], - thd_cp_a2a_inv: Optional[torch.Tensor], - ) -> torch.Tensor: - """Run GDN hidden-parallel to context-parallel A2A using CP-to-HP context.""" - if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': - if cp_size > 1: - assert thd_cp_a2a_inv is not None - norm_out = norm_out.index_select(0, thd_cp_a2a_inv) - norm_out = tensor_a2a_hp2cp( - norm_out, - seq_dim=0, - head_dim=-1, - cp_group=cp_group, - redo_attention_load_balancing=False, - ) - else: - norm_out = tensor_a2a_hp2cp(norm_out, seq_dim=0, head_dim=-1, cp_group=cp_group) - - return norm_out - - @jit_fuser - def _apply_gated_norm(self, x, gate): - # Output Norm - x_dtype = x.dtype - x = x.reshape(-1, x.shape[-1]) - y = self.out_norm(x) - # Output gate - gate = gate.reshape(-1, gate.shape[-1]) - y = y * self.act_fn(gate.float()) - y = y.to(x_dtype) - return y - - @jit_fuser - def _prepare_qkv_for_gated_delta_rule( - self, qkv, gate, beta, alpha, batch, seq_len, cp_size_headwise - ): - """ - Prepare query, key, value, gate, beta, alpha tensors for gated delta rule. - Fuses split, reshape, L2 norm, repeat_interleave, and contiguous operations. - """ - # Split qkv into query_key and value - query_key, value = torch.split( - qkv, - [2 * self.qk_dim_local_tp // cp_size_headwise, self.v_dim_local_tp // cp_size_headwise], - dim=-1, - ) - - # Reshape query_key and value - query_key = query_key.reshape(batch, seq_len, -1, self.key_head_dim) - value = value.reshape(batch, seq_len, -1, self.value_head_dim) - - # Apply L2 norm to query and key - if self.use_qk_l2norm: - query_key = l2norm(query_key.contiguous()) - - # Split query and key - split_size = self.qk_dim_local_tp // self.key_head_dim // cp_size_headwise - query, key = torch.split(query_key, [split_size, split_size], dim=2) - - # Expand query and key if needed (grouped query attention) - if self.num_value_heads // self.num_key_heads > 1: - repeat_factor = self.num_value_heads // self.num_key_heads - query = query.repeat_interleave(repeat_factor, dim=2) - key = key.repeat_interleave(repeat_factor, dim=2) - - # Make all tensors contiguous - query = query.contiguous() - key = key.contiguous() - value = value.contiguous() - gate = gate.contiguous() - beta = beta.contiguous() - alpha = alpha.contiguous() - - return query, key, value, gate, beta, alpha - - @jit_fuser - def _compute_g_and_beta(self, A_log_local_cp, dt_bias_local_cp, alpha, beta): - """ - Compute g (decay) and beta (sigmoid) for gated delta rule. - Fuses exp, softplus, mul, neg, and sigmoid operations. - """ - g = -A_log_local_cp.exp() * F.softplus(alpha.float() + dt_bias_local_cp) # In fp32 - beta = beta.sigmoid() - return g, beta - - def _resolve_cu_seqlens( - self, cu_seqlens_padded, cu_seqlens_actual, total_seq_len, name, cp_size: int = 1 - ) -> torch.Tensor: - """Resolve cu_seqlens for packed sequence all-to-all, handling alignment padding.""" - if cu_seqlens_padded is not None: - cu_seqlens = cu_seqlens_padded - else: - cu_seqlens = cu_seqlens_actual - - total_cu = cu_seqlens[-1].cpu().item() - if total_cu != total_seq_len: - raise ValueError( - f"GDN: {name}[-1]={total_cu} does not match " - f"total_sequence_length={total_seq_len}. " - f"({cu_seqlens_padded=}, {cu_seqlens_actual=})." - ) - - seq_lengths = cu_seqlens[1:] - cu_seqlens[:-1] - if (seq_lengths % cp_size != 0).any(): - raise ValueError( - f"All per-sequence lengths in cu_seqlens must be divisible by cp_size={cp_size}, " - f"but got lengths: {seq_lengths.tolist()}" - ) - - return cu_seqlens - - def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_group=None): - """Provide a sharded state dictionary for distributed checkpointing.""" - # Guard for cases metadata is not provided - metadata = ensure_metadata_has_dp_cp_group(metadata) - - sharded_state_dict = {} - # Parameters - self._save_to_state_dict(sharded_state_dict, "", keep_vars=True) - sharded_state_dict = make_sharded_tensors_for_checkpoint( - sharded_state_dict, - prefix, - tensor_parallel_layers_axis_map={ - "A_log": 0, - "dt_bias": 0, - }, # parameters sharded across TP - sharded_offsets=sharded_offsets, - tp_group=(tp_group if tp_group is not None else self.pg_collection.tp), - dp_cp_group=metadata['dp_cp_group'], - ) - # Submodules - tp_group = tp_group if tp_group is not None else self.pg_collection.tp - for name, module in self.named_children(): - if name == "conv1d": - # Add TP sharding for Conv1d - module_sd = module.state_dict(prefix="", keep_vars=True) - tp_sharding_map = {f"weight": 0} - if self.conv_bias: - tp_sharding_map[f"bias"] = 0 - module_sharded_sd = make_sharded_tensors_for_checkpoint( - module_sd, - f"{prefix}{name}.", - tp_sharding_map, - sharded_offsets, - tp_group=tp_group, - dp_cp_group=metadata['dp_cp_group'], - ) - else: - module_sharded_sd = sharded_state_dict_default( - module, f"{prefix}{name}.", sharded_offsets, metadata, tp_group=tp_group - ) - - sharded_state_dict.update(module_sharded_sd) - - # At this point the TP sharding is correctly defined for each tensor, but some of the - # tensors must be additionally split into separate parts - in_proj_dim_local_tp = self.in_proj_dim // self.tp_size - assert sharded_state_dict[f"{prefix}in_proj.weight"].data.size(0) == in_proj_dim_local_tp, ( - in_proj_dim_local_tp, - sharded_state_dict[f"{prefix}in_proj.weight"], - ) - - sharded_state_dict[f"{prefix}in_proj.weight"] = _split_tensor_factory( - sharded_state_dict[f"{prefix}in_proj.weight"], - [ - self.qk_dim_local_tp, - self.qk_dim_local_tp, - self.v_dim_local_tp, - self.v_dim_local_tp, - self.num_value_heads // self.tp_size, - self.num_value_heads // self.tp_size, - ], - ["query", "key", "value", "z", "beta", "alpha"], - 0, - ) - - conv_layer_name_list = ["conv1d.weight"] - assert ( - sharded_state_dict[f"{prefix}conv1d.weight"].data.size(0) == self.conv_dim_local_tp - ), (self.conv_dim_local_tp, sharded_state_dict[f"{prefix}conv1d.weight"]) - if self.conv_bias: - conv_layer_name_list.append("conv1d.bias") - assert ( - sharded_state_dict[f"{prefix}conv1d.bias"].data.size(0) == self.conv_dim_local_tp - ), (self.conv_dim_local_tp, sharded_state_dict[f"{prefix}conv1d.bias"]) - for conv_layer_name in conv_layer_name_list: - sharded_state_dict[f"{prefix}{conv_layer_name}"] = _split_tensor_factory( - sharded_state_dict[f"{prefix}{conv_layer_name}"], - [self.qk_dim_local_tp, self.qk_dim_local_tp, self.v_dim_local_tp], - ["query", "key", "value"], - 0, - ) - - return sharded_state_dict - - def backward_dw(self): - """Execute weight gradient computation for all linear layers.""" - self._backward_in_proj() - self._backward_out_proj() - - def _backward_in_proj(self): - """Computes weight gradients of input projection layer.""" - self.in_proj.backward_dw() - - def _backward_out_proj(self): - """Computes weight gradients of output projection layer.""" - self.out_proj.backward_dw() - - -def _build_thd_cp_a2a_perm( - cu_seqlens: torch.Tensor, cp_size: int, t_global: int -) -> tuple[torch.Tensor, torch.Tensor]: - cu = cu_seqlens.to(dtype=torch.long) - t_local = t_global // cp_size - - positions = torch.arange(t_global, device=cu.device) - seq_idx = torch.bucketize(positions, cu[1:], right=True) - seq_lens = torch.diff(cu) - halves = seq_lens // (2 * cp_size) # per-sequence half-chunk size - local_starts = cu[:-1] // cp_size - global_starts = cu[:-1] - - half_i = halves[seq_idx] - pos_in_seq = positions - global_starts[seq_idx] - - natural_chunk = pos_in_seq // half_i # in [0, 2*cp) - offset = pos_in_seq - natural_chunk * half_i - - # Invert the ordering produced by `_undo_attention_load_balancing`: - # natural_chunk < cp: load_balanced = 2 * natural_chunk - # natural_chunk >= cp: load_balanced = 4*cp - 2*natural_chunk - 1 - lb_chunk = torch.where( - natural_chunk < cp_size, 2 * natural_chunk, 4 * cp_size - 2 * natural_chunk - 1 - ) - - # In the per-sequence load-balanced layout each rank owns load-balanced - # chunks (2r) and (2r+1), in that order, of every sequence. - rank = lb_chunk // 2 - half_within_rank = lb_chunk - 2 * rank - k = half_within_rank * half_i + offset - - idx = rank * t_local + local_starts[seq_idx] + k - - inv = torch.empty_like(idx) - inv[idx] = positions - - return idx, inv - - -@lru_cache(maxsize=8) -def _build_head_perm_for_split_sections( - split_sections: tuple[int, ...], cp_size: int, device: torch.device -) -> torch.Tensor: - assert all( - s % cp_size == 0 for s in split_sections - ), f"split_sections {split_sections} must be divisible by cp_size {cp_size} for GDN" - offset = 0 - parts = [] - for s in split_sections: - parts.append( - torch.arange(offset, offset + s, device=device, dtype=torch.long).view(cp_size, -1) - ) - offset += s - - return torch.cat(parts, dim=-1).view(-1) - - -#################### -# Context parallel utilities -#################### -def get_parameter_local_cp_headwise( - param: torch.Tensor, - dim: int, - cp_group: torch.distributed.ProcessGroup, - split_sections: Optional[list[int]] = None, -) -> torch.Tensor: - """Get the local parameter for the current context parallel rank. - - Args: - param (torch.Tensor): The entire parameter to get the local parameter for. - dim (int): The dimension to split the parameter along. Usually the dimension of head. - cp_group (torch.distributed.ProcessGroup): The context parallel group. - split_sections (Optional[list[int]]): If not None, - first split the parameter along the dimension dim into sections, - then get the local hidden parallel weights separately, - finally concatenate the local hidden parallel weights along the dimension dim. - - Returns: - torch.Tensor: The local parameter for the current context parallel rank. - """ - - cp_size = cp_group.size() if cp_group is not None else 1 - - # No need to split if CP size is 1. - if cp_size == 1: - return param - - cp_rank = cp_group.rank() - - # Split first if needed. - if split_sections is not None: - inputs = torch.split(param, split_sections, dim=dim) - outputs = [] - for p in inputs: - p = get_parameter_local_cp_headwise(p, dim, cp_group) - outputs.append(p) - return torch.cat(outputs, dim=dim) - - # Slice the parameter. - slices = [slice(None)] * param.dim() - dim_size = param.size(dim=dim) - slices[dim] = slice(cp_rank * dim_size // cp_size, (cp_rank + 1) * dim_size // cp_size) - param = param[slices] - return param - - -def tensor_a2a_cp2hp( - tensor: torch.Tensor, - seq_dim: int, - head_dim: int, - cp_group: torch.distributed.ProcessGroup, - split_sections: Optional[list[int]] = None, - undo_attention_load_balancing: bool = True, -): - """All-to-all context parallel to hidden parallel. - - This communication primitive is used by GDN headwise CP mode. - - Args: - tensor (torch.Tensor): The tensor to all-to-all. - Currently only support (seq_len, batch, head_dim) shaped tensor. - seq_dim (int): The dimension of sequence length. Currently only supports seq_dim == 0. - head_dim (int): The dimension of head. Currently only supports head_dim == -1 or 2. - cp_group (torch.distributed.ProcessGroup): The context parallel group. - split_sections (Optional[list[int]]): If not None, split the tensor along the dimension - head_dim into sections first, then do all-to-all for each section separately, - finally concatenate the separated tensors along the dimension head_dim. - undo_attention_load_balancing (bool): Whether to undo the attention load balancing of CP. - - Returns: - torch.Tensor: The all-to-all tensor. - """ - - cp_size = cp_group.size() if cp_group is not None else 1 - - # No need to all-to-all if CP size is 1. - if cp_size == 1: - return tensor - - # Limitations of mamba_context_parallel._all_to_all_cp2hp. - assert seq_dim == 0, f"tensor_a2a_cp2hp only supports seq_dim == 0 for now, but got {seq_dim=}" - assert ( - head_dim == -1 or head_dim == 2 - ), f"tensor_a2a_cp2hp only supports head_dim == -1 or 2 for now, but got {head_dim=}" - assert ( - tensor.dim() == 3 - ), f"tensor_a2a_cp2hp only supports 3-d input tensor for now, but got {tensor.dim()=}" - - # Split first if needed. - if split_sections is not None: - inputs = torch.split(tensor, split_sections, dim=head_dim) - outputs = [] - for x in inputs: - x = tensor_a2a_cp2hp( - x, - seq_dim=seq_dim, - head_dim=head_dim, - cp_group=cp_group, - undo_attention_load_balancing=False, - ) - outputs.append(x) - tensor = torch.cat(outputs, dim=head_dim) - else: - tensor = _all_to_all_cp2hp(tensor, cp_group) - - # Undo attention load balancing last if needed. - if undo_attention_load_balancing: - tensor = _undo_attention_load_balancing(tensor, cp_size) - return tensor - - -def tensor_a2a_hp2cp( - tensor: torch.Tensor, - seq_dim: int, - head_dim: int, - cp_group: torch.distributed.ProcessGroup, - split_sections: Optional[list[int]] = None, - redo_attention_load_balancing: bool = True, -): - """All-to-all hidden parallel to context parallel. - - This communication primitive is used by GDN headwise CP mode. - - Args: - tensor (torch.Tensor): The tensor to all-to-all. - Currently only support (seq_len, batch, head_dim) shaped tensor. - seq_dim (int): The dimension of sequence length. Currently only supports seq_dim == 0. - head_dim (int): The dimension of head. Currently only supports head_dim == -1 or 2. - cp_group (torch.distributed.ProcessGroup): The context parallel group. - split_sections (Optional[list[int]]): If not None, first split the tensor along the - dimension head_dim into sections, then do all-to-all for each section separately, - finally concatenate the separated tensors along the dimension head_dim. - redo_attention_load_balancing (bool): Whether to redo the attention load balancing of HP. - - Returns: - torch.Tensor: The all-to-all tensor. - """ - - cp_size = cp_group.size() if cp_group is not None else 1 - - # No need to all-to-all if CP size is 1. - if cp_size == 1: - return tensor - - # Limitations of mamba_context_parallel._all_to_all_hp2cp. - assert seq_dim == 0, f"tensor_a2a_hp2cp only supports seq_dim == 0 for now, but got {seq_dim=}" - assert ( - head_dim == -1 or head_dim == 2 - ), f"tensor_a2a_hp2cp only supports head_dim == -1 or 2 for now, but got {head_dim=}" - assert ( - tensor.dim() == 3 - ), f"tensor_a2a_hp2cp only supports 3-d input tensor for now, but got {tensor.dim()=}" - - # Redo attention load balancing first if needed. - if redo_attention_load_balancing: - tensor = _redo_attention_load_balancing(tensor, cp_size) - - # Split first if needed. - if split_sections is not None: - inputs = torch.split(tensor, split_sections, dim=head_dim) - outputs = [] - for x in inputs: - x = tensor_a2a_hp2cp( - x, - seq_dim=seq_dim, - head_dim=head_dim, - cp_group=cp_group, - redo_attention_load_balancing=False, - ) - outputs.append(x) - tensor = torch.cat(outputs, dim=head_dim) - else: - tensor = _all_to_all_hp2cp(tensor, cp_group) - - return tensor - - -#################### -# Torch native gated delta rule -#################### -def torch_chunk_gated_delta_rule( - query, - key, - value, - g, - beta, - chunk_size=64, - initial_state=None, - output_final_state=False, - use_qk_l2norm_in_kernel=False, - cu_seqlens=None, - cp_context=None, -): - # pylint: disable=line-too-long - ''' - Torch-native implementation of chunked gated delta rule for deterministic mode. - Need this because FLA is not deterministic. - - Reference: https://github.com/huggingface/transformers/blob/144c8ce2809a2e21914017652700e1ecb450501e/src/transformers/models/qwen3_next/modeling_qwen3_next.py#L470-L547 - ''' - - assert ( - cu_seqlens is None - ), "cu_seqlens is not supported for torch_chunk_gated_delta_rule for now." - assert ( - cp_context is None - ), "cp_context is not supported for torch_chunk_gated_delta_rule for now." - - initial_dtype = query.dtype - if use_qk_l2norm_in_kernel: - query = l2norm(query, dim=-1, eps=1e-6) - key = l2norm(key, dim=-1, eps=1e-6) - query, key, value, beta, g = [ - x.transpose(1, 2).contiguous().to(torch.float32) for x in (query, key, value, beta, g) - ] - - batch_size, num_heads, sequence_length, k_head_dim = key.shape - v_head_dim = value.shape[-1] - pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size - query = F.pad(query, (0, 0, 0, pad_size)) - key = F.pad(key, (0, 0, 0, pad_size)) - value = F.pad(value, (0, 0, 0, pad_size)) - beta = F.pad(beta, (0, pad_size)) - g = F.pad(g, (0, pad_size)) - total_sequence_length = sequence_length + pad_size - scale = 1 / (query.shape[-1] ** 0.5) - query = query * scale - - v_beta = value * beta.unsqueeze(-1) - k_beta = key * beta.unsqueeze(-1) - # reshape to chunks - query, key, value, k_beta, v_beta = [ - x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) - for x in (query, key, value, k_beta, v_beta) - ] - g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) - mask = torch.triu( - torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=0 - ) - - # chunk decay - g = g.cumsum(dim=-1) - decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() - attn = -((k_beta @ key.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0) - for i in range(1, chunk_size): - row = attn[..., i, :i].clone() - sub = attn[..., :i, :i].clone() - attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) - attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) - value = attn @ v_beta - k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) - last_recurrent_state = ( - torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(value) - if initial_state is None - else initial_state.to(value) - ) - core_attn_out = torch.zeros_like(value) - mask = torch.triu( - torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=query.device), diagonal=1 - ) - - # for each chunk - for i in range(0, total_sequence_length // chunk_size): - q_i, k_i, v_i = query[:, :, i], key[:, :, i], value[:, :, i] - attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) - v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state - v_new = v_i - v_prime - attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state - core_attn_out[:, :, i] = attn_inter + attn @ v_new - last_recurrent_state = ( - last_recurrent_state * g[:, :, i, -1, None, None].exp() - + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new - ) - - if not output_final_state: - last_recurrent_state = None - core_attn_out = core_attn_out.reshape( - core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1] - ) - core_attn_out = core_attn_out[:, :, :sequence_length] - core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype) - return core_attn_out, last_recurrent_state diff --git a/megatron/core/ssm/gated_delta_net/__init__.py b/megatron/core/ssm/gated_delta_net/__init__.py new file mode 100644 index 00000000000..fdde9796490 --- /dev/null +++ b/megatron/core/ssm/gated_delta_net/__init__.py @@ -0,0 +1,36 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +"""Gated Delta Net (GDN) layer exports. + +This package replaces the former ``megatron/core/ssm/gated_delta_net.py`` module +at the same import path; the names below preserve that module's public surface. +""" + +from megatron.core.ssm.gated_delta_net.common import ( + HAVE_FLA, + GatedDeltaNetSubmodules, + _build_head_perm_for_split_sections, + _build_thd_cp_a2a_perm, + _split_tensor_factory, + causal_conv1d, + chunk_gated_delta_rule, + get_parameter_local_cp, + l2norm, + tensor_a2a_cp2hp, + tensor_a2a_hp2cp, +) +from megatron.core.ssm.gated_delta_net.gdn import GatedDeltaNet, torch_chunk_gated_delta_rule + +__all__ = [ + "HAVE_FLA", + "GatedDeltaNet", + "GatedDeltaNetSubmodules", + "_split_tensor_factory", + "causal_conv1d", + "chunk_gated_delta_rule", + "get_parameter_local_cp", + "l2norm", + "tensor_a2a_cp2hp", + "tensor_a2a_hp2cp", + "torch_chunk_gated_delta_rule", +] diff --git a/megatron/core/ssm/gated_delta_net/common.py b/megatron/core/ssm/gated_delta_net/common.py new file mode 100644 index 00000000000..9d7055cf67a --- /dev/null +++ b/megatron/core/ssm/gated_delta_net/common.py @@ -0,0 +1,917 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh. + +# Some of this code was adopted from https://github.com/huggingface/transformers +# This source code is licensed under the Apache license found in the +# LICENSE file in the root directory of this source tree. + +from dataclasses import dataclass +from functools import lru_cache +from typing import Optional, Protocol, Union + +import torch +import torch.nn as nn + +from megatron.core.fp8_utils import get_fp8_align_size +from megatron.core.inference.contexts import BaseInferenceContext +from megatron.core.jit import jit_fuser +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.mamba_context_parallel import ( + _all_to_all_cp2hp, + _all_to_all_hp2cp, + _redo_attention_load_balancing, + _undo_attention_load_balancing, +) +from megatron.core.ssm.utils import _split_tensor_factory +from megatron.core.tensor_parallel import get_cuda_rng_tracker +from megatron.core.transformer import TransformerConfig +from megatron.core.transformer.identity_op import IdentityOp +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.utils import ( + ensure_metadata_has_dp_cp_group, + make_sharded_tensors_for_checkpoint, + sharded_state_dict_default, +) + +try: + from fla.modules.convolution import causal_conv1d + from fla.modules.l2norm import l2norm + from fla.ops.cp import build_cp_context + from fla.ops.gated_delta_rule import chunk_gated_delta_rule + + HAVE_FLA = True +except ImportError: + build_cp_context = None + causal_conv1d = None + l2norm = None + chunk_gated_delta_rule = None + + HAVE_FLA = False + +__all__ = [ + "HAVE_FLA", + "GatedDeltaNetSubmodules", + "_GDNBase", + "_build_head_perm_for_split_sections", + "_build_thd_cp_a2a_perm", + "_split_tensor_factory", + "a2a_cp_to_hp", + "a2a_hp_to_cp", + "build_cp_context", + "causal_conv1d", + "chunk_gated_delta_rule", + "get_parameter_local_cp", + "l2norm", + "tensor_a2a_cp2hp", + "tensor_a2a_hp2cp", +] + + +@dataclass +class GatedDeltaNetSubmodules: + """ + Contains the module specs for the input linear, output norm, and output linear layers. + """ + + in_proj: Union[ModuleSpec, type] = IdentityOp + out_norm: Union[ModuleSpec, type] = IdentityOp + out_proj: Union[ModuleSpec, type] = IdentityOp + + +class GatedDeltaRuleInterface(Protocol): + """ + Unified typing protocol for GDN core computation interfaces. + """ + + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + *, + scale: float | None = None, + initial_state: torch.Tensor | None = None, + output_final_state: bool = False, + use_qk_l2norm_in_kernel: bool = False, + cu_seqlens: torch.LongTensor | None = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor | None]: ... + + +class _GDNBase(MegatronModule): + """Shared implementation for the Gated Delta Net (GDN) layer. + + Hosts the fused input projection, causal convolution on q/k/v, CP all-to-all + plumbing, kernel-input preparation, gated output norm + projection, and + sharded checkpointing. + """ + + dt_bias_dim: int + a_log_dim: int + in_proj_extra_dim: int + in_proj_dim: int + + dt_bias: nn.Parameter + A_log: nn.Parameter + + gated_delta_rule: GatedDeltaRuleInterface + + def __init__( + self, + config: TransformerConfig, + submodules: GatedDeltaNetSubmodules, + layer_number: int = None, + bias: bool = False, + conv_bias: bool = False, + conv_init: float | None = None, + use_qk_l2norm: bool = True, + A_init_range: tuple[float, float] = (1, 16), + pg_collection: ProcessGroupCollection = None, + *, + name: str | None = None, + cp_comm_type: str | None = None, + pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, + ): + """ + Args: + config: The config of the model. + submodules: Contains the module specs for the input and output linear layers. + layer_number: The layer number of this GDN layer. + bias: Whether to use bias in the linear layers. + conv_bias: Whether to use bias in the causal convolution. + conv_init: The initialization range for the causal convolution weights. + use_qk_l2norm: Whether to use L2 normalization in the kernel of the gated delta rule. + A_init_range: The initialization range for the attention weights. + pg_collection: The required process groups to use for tensor model parallel and context + parallel. + name (str | None): Optional module path prefix used for child module names. + cp_comm_type (Optional[str]): Accepted for TransformerLayer compatibility and + ignored; GDN implements context parallelism with its own all-to-alls rather + than the attention CP communication schemes. + pp_layer_offset (Optional[int]): Pipeline layer offset forwarded by + TransformerLayer. Stored for MTP/TransformerLayer API compatibility. + is_mtp_layer (bool): Whether this module is inside an MTP prediction depth. + """ + if not HAVE_FLA: + raise ImportError( + "FLA is not installed. Please install it with " + "`pip install flash-linear-attention[cuda]`." + ) + + super().__init__(config) + + # Attributes from arguments + self.layer_number = layer_number + self._pp_layer_offset = pp_layer_offset + self.is_mtp_layer = is_mtp_layer + self.bias = bias + self.conv_bias = conv_bias + self.conv_init = conv_init + assert A_init_range[0] >= 0 and A_init_range[1] >= A_init_range[0] + self.A_init_range = A_init_range + self.use_qk_l2norm = use_qk_l2norm + assert pg_collection is not None, "pg_collection must be provided for GatedDeltaNet" + self.pg_collection = pg_collection + self.tp_group = pg_collection.tp + # Static/max CP size from model construction. Runtime dynamic CP paths must resolve + # the effective group from packed_seq_params instead of using this value. + self.cp_size = self.pg_collection.cp.size() + self.tp_size = self.pg_collection.tp.size() + self.sp_size = self.tp_size if config.sequence_parallel else 1 + + # Attributes from config + self.config = config + self.hidden_size = config.hidden_size + self.act_fn = config.activation_func + self.activation = self.act_fn.__name__ + self.conv_kernel_dim = config.linear_conv_kernel_dim + self.key_head_dim = config.linear_key_head_dim + self.value_head_dim = config.linear_value_head_dim + self.num_key_heads = config.linear_num_key_heads + self.num_value_heads = config.linear_num_value_heads + self.qk_dim = self.key_head_dim * self.num_key_heads + self.v_dim = self.value_head_dim * self.num_value_heads + self.qk_dim_local_tp = self.qk_dim // self.tp_size + self.v_dim_local_tp = self.v_dim // self.tp_size + + # Headwise CP shards heads over the CP group; chunkwise CP keeps heads local. + if self.config.linear_cp_mode == "headwise": + num_key_heads_per_tp = self.num_key_heads // self.tp_size + num_value_heads_per_tp = self.num_value_heads // self.tp_size + assert num_key_heads_per_tp % self.cp_size == 0, ( + f"GDN head-parallel CP requires the static (max) cp_size ({self.cp_size}) " + f"to evenly divide num_key_heads per TP rank ({num_key_heads_per_tp}); " + f"all runtime dynamic cp_size values divide the static one and so will also divide." + ) + assert num_value_heads_per_tp % self.cp_size == 0, ( + f"GDN head-parallel CP requires the static (max) cp_size ({self.cp_size}) " + f"to evenly divide num_value_heads per TP rank ({num_value_heads_per_tp}); " + f"all runtime dynamic cp_size values divide the static one and so will also divide." + ) + + self.num_v_heads_local_tp = self.num_value_heads // self.tp_size + + attrs_to_check = ( + "dt_bias_dim", + "a_log_dim", + "in_proj_extra_dim", + "in_proj_split_names", + "in_proj_split_sections", + "gated_delta_rule", + ) + self._setup_variant_attrs() + for attr in attrs_to_check: + assert hasattr(self, attr), f"Attribute {attr} for GDN is not set" + assert getattr(self, attr) is not None, f"Attribute {attr} for GDN is not set" + # Full input projection width: q, k, v, output gate, and variant-specific gate features. + self.in_proj_dim = self.qk_dim * 2 + self.v_dim * 2 + self.in_proj_extra_dim + + if self.config.fp8: + fp8_align_size = get_fp8_align_size(self.config.fp8_recipe) + assert self.in_proj_dim % fp8_align_size == 0, ( + "For FP8, the innermost dimension of the GDN layer " + "input projection output tensor must be a multiple of 16." + ) + self.in_proj = build_module( + submodules.in_proj, + self.hidden_size, + self.in_proj_dim, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=bias, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name="fc1", + tp_group=self.pg_collection.tp, + name=(name + ".in_proj") if name is not None else None, + ) + + # Conv1d for QKV + self.conv_dim = self.qk_dim * 2 + self.v_dim + self.conv_dim_local_tp = self.conv_dim // self.tp_size + + self.conv1d = nn.Conv1d( + in_channels=self.conv_dim_local_tp, + out_channels=self.conv_dim_local_tp, + bias=conv_bias, + kernel_size=self.conv_kernel_dim, + groups=self.conv_dim_local_tp, + padding=self.conv_kernel_dim - 1, + device=torch.cuda.current_device(), + dtype=config.params_dtype, + ) + setattr(self.conv1d.weight, "tensor_model_parallel", True) + setattr(self.conv1d.weight, "partition_dim", 0) + if conv_bias: + setattr(self.conv1d.bias, "tensor_model_parallel", True) + setattr(self.conv1d.bias, "partition_dim", 0) + + self.dt_bias = nn.Parameter( + torch.empty( + self.dt_bias_dim, dtype=self.config.params_dtype, device=torch.cuda.current_device() + ) + ) + setattr(self.dt_bias, "tensor_model_parallel", True) + setattr(self.dt_bias, "partition_dim", 0) + + self.A_log = nn.Parameter( + torch.empty( + self.a_log_dim, dtype=self.config.params_dtype, device=torch.cuda.current_device() + ) + ) + setattr(self.A_log, "tensor_model_parallel", True) + setattr(self.A_log, "partition_dim", 0) + + # Output layernorm before projection + self.out_norm = build_module( + submodules.out_norm, + config=self.config, + hidden_size=self.value_head_dim, + eps=self.config.layernorm_epsilon, + ) + self.recompute_norm_out = False + self.norm_out_checkpoint = None + self.recompute_gdn = False + if self.config.recompute_granularity == "selective" and self.config.recompute_modules: + self.recompute_norm_out = "gdn_norm_out" in self.config.recompute_modules + self.recompute_gdn = "gdn" in self.config.recompute_modules + + self.out_proj = build_module( + submodules.out_proj, + self.v_dim, + self.hidden_size, + config=self.config, + init_method=self.config.output_layer_init_method, + bias=bias, + input_is_parallel=True, + skip_bias_add=True, + is_expert=False, + tp_comm_buffer_name="fc2", + tp_group=self.pg_collection.tp, + name=(name + ".out_proj") if name is not None else None, + ) + # TODO: Packed sequence cu_seqlens can vary per batch; cache only static SBHD + # cp_context entries here and revisit routing metadata lifetime in the CP layout refactor. + self._chunkwise_cp_context_cache: dict[tuple[int, int], tuple[torch.Tensor, object]] = {} + + self.reset_parameters() + + def _setup_variant_attrs(self): + """Set GDN projection sections, gate parameter sizes, and kernel callable. + + Must set: + - ``in_proj_extra_dim`` (the in_proj sections beyond q/k/v/z; the base + class derives ``in_proj_dim`` from it) + - ``in_proj_split_names`` + - ``in_proj_split_sections`` + - ``dt_bias_dim`` / ``a_log_dim`` (sizes of the gate parameters, which the + base class creates after the conv1d module to preserve the original + parameter registration order) + - ``gated_delta_rule`` (the kernel callable). + """ + raise NotImplementedError + + def _reset_dt_bias(self): + """Initialize ``dt_bias``. Called from ``reset_parameters`` under the RNG tracker. + + Defaults to ones; subclasses can override this if their kernel expects a + different step-size parametrization. + """ + torch.ones( + self.dt_bias_dim, + dtype=self.config.params_dtype, + device=torch.cuda.current_device(), + out=self.dt_bias.data, + ) + + def reset_parameters(self): + """Reset the parameters.""" + if self.config.perform_initialization: + with get_cuda_rng_tracker().fork(): + if self.conv_init is not None: + nn.init.uniform_(self.conv1d.weight, -self.conv_init, self.conv_init) + self._reset_dt_bias() + A = torch.empty( + self.A_log.shape[0], + dtype=self.config.params_dtype, + device=torch.cuda.current_device(), + ).uniform_(*self.A_init_range) + self.A_log.data.copy_(torch.log(A)) + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + inference_context: Optional[BaseInferenceContext] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + sequence_len_offset: Optional[int] = None, + *, + inference_params: Optional[BaseInferenceContext] = None, + **kwargs, + ) -> tuple[torch.Tensor, torch.Tensor]: + # pylint: disable=missing-function-docstring + raise NotImplementedError + + @jit_fuser + def _apply_gated_norm(self, x, gate): + # Output Norm + x_dtype = x.dtype + x = x.reshape(-1, x.shape[-1]) + y = self.out_norm(x) + # Output gate + gate = gate.reshape(-1, gate.shape[-1]) + y = y * self.act_fn(gate.float()) + y = y.to(x_dtype) + return y + + @jit_fuser + def _prepare_input_for_gated_delta_rule( + self, + qkv: torch.Tensor, + gate: torch.Tensor, + A_log_local_cp: torch.Tensor, + dt_bias_local_cp: torch.Tensor, + batch: int, + seq_len: int, + *gate_feats: tuple[torch.Tensor], + cp_size_headwise: int | None = None, + ) -> dict[str, torch.Tensor]: + """ + Prepare all gated delta rule kernel inputs. + + Fuses split, reshape, L2 norm, decay/gate activations, repeat_interleave, and + contiguous operations. ``gate_feats`` holds the in_proj sections after qkv + and gate, which ``_compute_gates`` turns into the decay and gating tensors. + + Returns: + (dict[str, Tensor]): Kernel inputs keyed by kernel argument name (``q``, + ``k``, ``v``, ``g``, and ``beta``), and the output + gate (z) tensor under the ``gate`` key, which is not a kernel input. + """ + cp_size = 1 if cp_size_headwise is None else cp_size_headwise + + # Split qkv into query_key and value + query_key, value = torch.split( + qkv, [2 * self.qk_dim_local_tp // cp_size, self.v_dim_local_tp // cp_size], dim=-1 + ) + + # Reshape query_key and value + query_key = query_key.reshape(batch, seq_len, -1, self.key_head_dim) + value = value.reshape(batch, seq_len, -1, self.value_head_dim) + + # Apply L2 norm to query and key + if self.use_qk_l2norm: + query_key = l2norm(query_key.contiguous()) + + # Split query and key + split_size = self.qk_dim_local_tp // self.key_head_dim // cp_size + query, key = torch.split(query_key, [split_size, split_size], dim=2) + + # Expand query and key if needed (grouped query attention) + repeat_factor = self.num_value_heads // self.num_key_heads + if repeat_factor > 1: + query = query.repeat_interleave(repeat_factor, dim=2) + key = key.repeat_interleave(repeat_factor, dim=2) + + g, variant_kernel_inputs = self._compute_gates( + A_log_local_cp, dt_bias_local_cp, batch, seq_len, *gate_feats + ) + + kernel_inputs = { + "q": query.contiguous(), + "k": key.contiguous(), + "v": value.contiguous(), + "g": g.contiguous(), + "gate": gate.contiguous(), + **variant_kernel_inputs, + } + return kernel_inputs + + def _compute_gates( + self, + A_log_local_cp: torch.Tensor, + dt_bias_local_cp: torch.Tensor, + batch: int, + seq_len: int, + *gate_feats: tuple[torch.Tensor], + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """ + Compute the log-decay ``g`` and remaining kernel inputs. + + Args: + A_log_local_cp: CP-local slice of ``A_log``. + dt_bias_local_cp: CP-local slice of ``dt_bias``. + batch: Batch size. + seq_len: Sequence length. + gate_feats: The in_proj output sections after the qkv and output-gate sections. + + Returns: + (tuple[Tensor, dict[str, Tensor]]): The log-decay ``g`` and a dict of the + remaining kernel inputs keyed by kernel argument name. + """ + raise NotImplementedError + + def _resolve_cu_seqlens( + self, cu_seqlens_padded, cu_seqlens_actual, total_seq_len, name, cp_size: int = 1 + ) -> torch.Tensor: + """Resolve cu_seqlens for packed sequence all-to-all, handling alignment padding.""" + if cu_seqlens_padded is not None: + cu_seqlens = cu_seqlens_padded + else: + cu_seqlens = cu_seqlens_actual + + total_cu = cu_seqlens[-1].cpu().item() + if total_cu != total_seq_len: + raise ValueError( + f"GDN: {name}[-1]={total_cu} does not match " + f"total_sequence_length={total_seq_len}. " + f"({cu_seqlens_padded=}, {cu_seqlens_actual=})." + ) + + seq_lengths = cu_seqlens[1:] - cu_seqlens[:-1] + if (seq_lengths % cp_size != 0).any(): + raise ValueError( + f"All per-sequence lengths in cu_seqlens must be divisible by cp_size={cp_size}, " + f"but got lengths: {seq_lengths.tolist()}" + ) + + return cu_seqlens + + def sharded_state_dict(self, prefix="", sharded_offsets=(), metadata=None, tp_group=None): + """Provide a sharded state dictionary for distributed checkpointing.""" + # Guard for cases metadata is not provided + metadata = ensure_metadata_has_dp_cp_group(metadata) + + sharded_state_dict = {} + # Parameters + self._save_to_state_dict(sharded_state_dict, "", keep_vars=True) + sharded_state_dict = make_sharded_tensors_for_checkpoint( + sharded_state_dict, + prefix, + tensor_parallel_layers_axis_map={ + "A_log": 0, + "dt_bias": 0, + }, # parameters sharded across TP + sharded_offsets=sharded_offsets, + tp_group=(tp_group if tp_group is not None else self.pg_collection.tp), + dp_cp_group=metadata['dp_cp_group'], + ) + # Submodules + tp_group = tp_group if tp_group is not None else self.pg_collection.tp + for name, module in self.named_children(): + if name == "conv1d": + # Add TP sharding for Conv1d + module_sd = module.state_dict(prefix="", keep_vars=True) + tp_sharding_map = {f"weight": 0} + if self.conv_bias: + tp_sharding_map[f"bias"] = 0 + module_sharded_sd = make_sharded_tensors_for_checkpoint( + module_sd, + f"{prefix}{name}.", + tp_sharding_map, + sharded_offsets, + tp_group=tp_group, + dp_cp_group=metadata['dp_cp_group'], + ) + else: + module_sharded_sd = sharded_state_dict_default( + module, f"{prefix}{name}.", sharded_offsets, metadata, tp_group=tp_group + ) + + sharded_state_dict.update(module_sharded_sd) + + # At this point the TP sharding is correctly defined for each tensor, but some of the + # tensors must be additionally split into separate parts + in_proj_dim_local_tp = self.in_proj_dim // self.tp_size + assert sharded_state_dict[f"{prefix}in_proj.weight"].data.size(0) == in_proj_dim_local_tp, ( + in_proj_dim_local_tp, + sharded_state_dict[f"{prefix}in_proj.weight"], + ) + + sharded_state_dict[f"{prefix}in_proj.weight"] = _split_tensor_factory( + sharded_state_dict[f"{prefix}in_proj.weight"], + list(self.in_proj_split_sections), + self.in_proj_split_names, + 0, + ) + + conv_layer_name_list = ["conv1d.weight"] + assert ( + sharded_state_dict[f"{prefix}conv1d.weight"].data.size(0) == self.conv_dim_local_tp + ), (self.conv_dim_local_tp, sharded_state_dict[f"{prefix}conv1d.weight"]) + if self.conv_bias: + conv_layer_name_list.append("conv1d.bias") + assert ( + sharded_state_dict[f"{prefix}conv1d.bias"].data.size(0) == self.conv_dim_local_tp + ), (self.conv_dim_local_tp, sharded_state_dict[f"{prefix}conv1d.bias"]) + for conv_layer_name in conv_layer_name_list: + sharded_state_dict[f"{prefix}{conv_layer_name}"] = _split_tensor_factory( + sharded_state_dict[f"{prefix}{conv_layer_name}"], + [self.qk_dim_local_tp, self.qk_dim_local_tp, self.v_dim_local_tp], + ["query", "key", "value"], + 0, + ) + + return sharded_state_dict + + def backward_dw(self): + """Execute weight gradient computation for all linear layers.""" + self._backward_in_proj() + self._backward_out_proj() + + def _backward_in_proj(self): + """Computes weight gradients of input projection layer.""" + self.in_proj.backward_dw() + + def _backward_out_proj(self): + """Computes weight gradients of output projection layer.""" + self.out_proj.backward_dw() + + +def _build_thd_cp_a2a_perm( + cu_seqlens: torch.Tensor, cp_size: int, t_global: int +) -> tuple[torch.Tensor, torch.Tensor]: + cu = cu_seqlens.to(dtype=torch.long) + t_local = t_global // cp_size + + positions = torch.arange(t_global, device=cu.device) + seq_idx = torch.bucketize(positions, cu[1:], right=True) + seq_lens = torch.diff(cu) + halves = seq_lens // (2 * cp_size) # per-sequence half-chunk size + local_starts = cu[:-1] // cp_size + global_starts = cu[:-1] + + half_i = halves[seq_idx] + pos_in_seq = positions - global_starts[seq_idx] + + natural_chunk = pos_in_seq // half_i # in [0, 2*cp) + offset = pos_in_seq - natural_chunk * half_i + + # Invert the ordering produced by `_undo_attention_load_balancing`: + # natural_chunk < cp: load_balanced = 2 * natural_chunk + # natural_chunk >= cp: load_balanced = 4*cp - 2*natural_chunk - 1 + lb_chunk = torch.where( + natural_chunk < cp_size, 2 * natural_chunk, 4 * cp_size - 2 * natural_chunk - 1 + ) + + # In the per-sequence load-balanced layout each rank owns load-balanced + # chunks (2r) and (2r+1), in that order, of every sequence. + rank = lb_chunk // 2 + half_within_rank = lb_chunk - 2 * rank + k = half_within_rank * half_i + offset + + idx = rank * t_local + local_starts[seq_idx] + k + + inv = torch.empty_like(idx) + inv[idx] = positions + + return idx, inv + + +@lru_cache(maxsize=8) +def _build_head_perm_for_split_sections( + split_sections: tuple[int, ...], cp_size: int, device: torch.device +) -> torch.Tensor: + assert all( + s % cp_size == 0 for s in split_sections + ), f"split_sections {split_sections} must be divisible by cp_size {cp_size} for GDN" + offset = 0 + parts = [] + for s in split_sections: + parts.append( + torch.arange(offset, offset + s, device=device, dtype=torch.long).view(cp_size, -1) + ) + offset += s + + return torch.cat(parts, dim=-1).view(-1) + + +#################### +# Context parallel utilities +#################### +def get_parameter_local_cp( + param: torch.Tensor, + dim: int, + cp_group: torch.distributed.ProcessGroup | None, + split_sections: Optional[list[int]] = None, +) -> torch.Tensor: + """Get the local parameter for the current context parallel rank. + + Args: + param (torch.Tensor): The entire parameter to get the local parameter for. + dim (int): The dimension to split the parameter along. Usually the dimension of head. + cp_group (torch.distributed.ProcessGroup): The context parallel group. + split_sections (Optional[list[int]]): If not None, + first split the parameter along the dimension dim into sections, + then get the local hidden parallel weights separately, + finally concatenate the local hidden parallel weights along the dimension dim. + + Returns: + torch.Tensor: The local parameter for the current context parallel rank. + """ + + cp_size = cp_group.size() if cp_group is not None else 1 + + # No need to split if CP size is 1. + if cp_size == 1: + return param + + assert cp_group is not None + cp_rank = cp_group.rank() + + # Split first if needed. + if split_sections is not None: + inputs = torch.split(param, split_sections, dim=dim) + outputs = [] + for p in inputs: + p = get_parameter_local_cp(p, dim, cp_group) + outputs.append(p) + return torch.cat(outputs, dim=dim) + + # Slice the parameter. + slices = [slice(None)] * param.dim() + dim_size = param.size(dim=dim) + slices[dim] = slice(cp_rank * dim_size // cp_size, (cp_rank + 1) * dim_size // cp_size) + param = param[slices] + return param + + +def tensor_a2a_cp2hp( + tensor: torch.Tensor, + seq_dim: int, + head_dim: int, + cp_group: torch.distributed.ProcessGroup | None, + split_sections: Optional[list[int]] = None, + undo_attention_load_balancing: bool = True, +): + """All-to-all context parallel to hidden parallel. + + Args: + tensor (torch.Tensor): The tensor to all-to-all. + Currently only support (seq_len, batch, head_dim) shaped tensor. + seq_dim (int): The dimension of sequence length. Currently only supports seq_dim == 0. + head_dim (int): The dimension of head. Currently only supports head_dim == -1 or 2. + cp_group (torch.distributed.ProcessGroup): The context parallel group. + split_sections (Optional[list[int]]): If not None, split the tensor along the dimension + head_dim into sections first, then do all-to-all for each section separately, + finally concatenate the separated tensors along the dimension head_dim. + undo_attention_load_balancing (bool): Whether to undo the attention load balancing of CP. + + Returns: + torch.Tensor: The all-to-all tensor. + """ + + cp_size = cp_group.size() if cp_group is not None else 1 + + # No need to all-to-all if CP size is 1. + if cp_size == 1: + return tensor + + assert cp_group is not None + + # Limitations of mamba_context_parallel._all_to_all_cp2hp. + assert seq_dim == 0, f"tensor_a2a_cp2hp only supports seq_dim == 0 for now, but got {seq_dim=}" + assert ( + head_dim == -1 or head_dim == 2 + ), f"tensor_a2a_cp2hp only supports head_dim == -1 or 2 for now, but got {head_dim=}" + assert ( + tensor.dim() == 3 + ), f"tensor_a2a_cp2hp only supports 3-d input tensor for now, but got {tensor.dim()=}" + + # Split first if needed. + if split_sections is not None: + inputs = torch.split(tensor, split_sections, dim=head_dim) + outputs = [] + for x in inputs: + x = tensor_a2a_cp2hp( + x, + seq_dim=seq_dim, + head_dim=head_dim, + cp_group=cp_group, + undo_attention_load_balancing=False, + ) + outputs.append(x) + tensor = torch.cat(outputs, dim=head_dim) + else: + tensor = _all_to_all_cp2hp(tensor, cp_group) + + # Undo attention load balancing last if needed. + if undo_attention_load_balancing: + tensor = _undo_attention_load_balancing(tensor, cp_size) + return tensor + + +def tensor_a2a_hp2cp( + tensor: torch.Tensor, + seq_dim: int, + head_dim: int, + cp_group: torch.distributed.ProcessGroup | None, + split_sections: Optional[list[int]] = None, + redo_attention_load_balancing: bool = True, +): + """All-to-all hidden parallel to context parallel. + + Args: + tensor (torch.Tensor): The tensor to all-to-all. + Currently only support (seq_len, batch, head_dim) shaped tensor. + seq_dim (int): The dimension of sequence length. Currently only supports seq_dim == 0. + head_dim (int): The dimension of head. Currently only supports head_dim == -1 or 2. + cp_group (torch.distributed.ProcessGroup): The context parallel group. + split_sections (Optional[list[int]]): If not None, first split the tensor along the + dimension head_dim into sections, then do all-to-all for each section separately, + finally concatenate the separated tensors along the dimension head_dim. + redo_attention_load_balancing (bool): Whether to redo the attention load balancing of HP. + + Returns: + torch.Tensor: The all-to-all tensor. + """ + + cp_size = cp_group.size() if cp_group is not None else 1 + + # No need to all-to-all if CP size is 1. + if cp_size == 1: + return tensor + + assert cp_group is not None + + # Limitations of mamba_context_parallel._all_to_all_hp2cp. + assert seq_dim == 0, f"tensor_a2a_hp2cp only supports seq_dim == 0 for now, but got {seq_dim=}" + assert ( + head_dim == -1 or head_dim == 2 + ), f"tensor_a2a_hp2cp only supports head_dim == -1 or 2 for now, but got {head_dim=}" + assert ( + tensor.dim() == 3 + ), f"tensor_a2a_hp2cp only supports 3-d input tensor for now, but got {tensor.dim()=}" + + # Redo attention load balancing first if needed. + if redo_attention_load_balancing: + tensor = _redo_attention_load_balancing(tensor, cp_size) + + # Split first if needed. + if split_sections is not None: + inputs = torch.split(tensor, split_sections, dim=head_dim) + outputs = [] + for x in inputs: + x = tensor_a2a_hp2cp( + x, + seq_dim=seq_dim, + head_dim=head_dim, + cp_group=cp_group, + redo_attention_load_balancing=False, + ) + outputs.append(x) + tensor = torch.cat(outputs, dim=head_dim) + else: + tensor = _all_to_all_hp2cp(tensor, cp_group) + + return tensor + + +def a2a_cp_to_hp( + qkvzba: torch.Tensor, + in_proj_split_sections: tuple[int, ...], + cp_size: int, + cp_group: torch.distributed.ProcessGroup, + cu_seqlens_q: torch.Tensor | None, + seq_len: int, + packed_seq_params: PackedSeqParams | None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + """Run GDN context-parallel to hidden-parallel A2A and return its inverse context. + + Args: + qkvzba: in_proj output in sbhd format, sharded along the sequence dim over CP. + in_proj_split_sections: per-section sizes of the in_proj output, local to this + TP rank, used to build the pre-a2a head permutation. + cp_size: context-parallel world size. + cp_group: context-parallel process group. + cu_seqlens_q: cumulative sequence lengths, required for the ``thd`` path. + seq_len: global (unsharded) sequence length. + packed_seq_params: packed-sequence params; the ``thd`` path is taken when its + ``qkv_format`` is ``'thd'``. + + Returns: + The hidden-parallel tensor and the sequence-dim inverse permutation to hand to + :func:`a2a_hp_to_cp` (``None`` outside the ``thd`` + CP>1 case). + """ + if cp_size > 1: + # Pre-permute head dim so a single unsectioned a2a is equivalent to per-section a2a. + head_perm = _build_head_perm_for_split_sections( + in_proj_split_sections, cp_size, qkvzba.device + ) + qkvzba = qkvzba.index_select(-1, head_perm) + + thd_cp_a2a_inv = None + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + qkvzba = tensor_a2a_cp2hp( + qkvzba, seq_dim=0, head_dim=-1, cp_group=cp_group, undo_attention_load_balancing=False + ) + if cp_size > 1: + # Permute at the seq dim so that a single unsectioned a2a + # is equivalent to per-sequence a2a. + # This also folds the ``_undo_attention_load_balancing`` step. + thd_cp_a2a_idx, thd_cp_a2a_inv = _build_thd_cp_a2a_perm(cu_seqlens_q, cp_size, seq_len) + qkvzba = qkvzba.index_select(0, thd_cp_a2a_idx) + else: + qkvzba = tensor_a2a_cp2hp(qkvzba, seq_dim=0, head_dim=-1, cp_group=cp_group) + + return qkvzba, thd_cp_a2a_inv + + +def a2a_hp_to_cp( + norm_out: torch.Tensor, + cp_size: int, + cp_group: torch.distributed.ProcessGroup, + packed_seq_params: PackedSeqParams | None, + thd_cp_a2a_inv: torch.Tensor | None, +) -> torch.Tensor: + """Run GDN hidden-parallel to context-parallel A2A using CP-to-HP context. + + Args: + norm_out: gated-norm output in sbhd format, sharded along the head dim over CP. + cp_size: context-parallel world size. + cp_group: context-parallel process group. + packed_seq_params: packed-sequence params; the ``thd`` path is taken when its + ``qkv_format`` is ``'thd'``. + thd_cp_a2a_inv: sequence-dim inverse permutation returned by + :func:`a2a_cp_to_hp`, required on the ``thd`` path when ``cp_size > 1``. + + Returns: + The context-parallel tensor, matching the layout of the GDN module input. + """ + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + if cp_size > 1: + assert thd_cp_a2a_inv is not None + norm_out = norm_out.index_select(0, thd_cp_a2a_inv) + norm_out = tensor_a2a_hp2cp( + norm_out, seq_dim=0, head_dim=-1, cp_group=cp_group, redo_attention_load_balancing=False + ) + else: + norm_out = tensor_a2a_hp2cp(norm_out, seq_dim=0, head_dim=-1, cp_group=cp_group) + + return norm_out diff --git a/megatron/core/ssm/gated_delta_net/gdn.py b/megatron/core/ssm/gated_delta_net/gdn.py new file mode 100644 index 00000000000..79c58a41710 --- /dev/null +++ b/megatron/core/ssm/gated_delta_net/gdn.py @@ -0,0 +1,723 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2025, Songlin Yang, Jan Kautz, Ali Hatamizadeh. + +# Some of this code was adopted from https://github.com/huggingface/transformers +# This source code is licensed under the Apache license found in the +# LICENSE file in the root directory of this source tree. + +from functools import partial +from typing import Optional + +import torch +import torch.nn.functional as F + +from megatron.core import tensor_parallel +from megatron.core.context_parallel_layout import ( + contiguous_to_zigzag_chunks, + zigzag_to_contiguous_chunks, +) +from megatron.core.inference.contexts import BaseInferenceContext +from megatron.core.jit import jit_fuser +from megatron.core.packed_seq_params import PackedSeqParams, resolve_cp_group +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.gated_delta_net.common import ( + _GDNBase, + a2a_cp_to_hp, + a2a_hp_to_cp, + build_cp_context, + causal_conv1d, + chunk_gated_delta_rule, + get_parameter_local_cp, + l2norm, +) +from megatron.core.utils import deprecate_inference_params, nvtx_range_pop, nvtx_range_push + + +class GatedDeltaNet(_GDNBase): + # pylint: disable=missing-class-docstring + def _setup_variant_attrs(self): + """Set the GDN in_proj sizing, split tables, gate parameter dims, and kernel.""" + self.gdn_pre_gated_delta_rule_fusion = self.config.gdn_pre_gated_delta_rule_fusion + if self.config.deterministic_mode and self.gdn_pre_gated_delta_rule_fusion: + raise ValueError( + "Pre-GDR fusion is non-deterministic, but deterministic_mode=True. " + "Disable gdn_pre_gated_delta_rule_fusion or deterministic_mode." + ) + + # alpha, beta + self.in_proj_extra_dim = self.num_value_heads * 2 + + # Per-section sizes (and names) of the in_proj output, local to this TP rank. + # Used for the CP head permutation (pre-a2a), for splitting the projection + # output (post-a2a), and for the sharded checkpoint split of in_proj.weight. + self.in_proj_split_names = ["query", "key", "value", "z", "beta", "alpha"] + self.in_proj_split_sections = ( + self.qk_dim_local_tp, # q + self.qk_dim_local_tp, # k + self.v_dim_local_tp, # v + self.v_dim_local_tp, # gate (z) + self.num_value_heads // self.tp_size, # beta + self.num_value_heads // self.tp_size, # alpha + ) + self.dt_bias_dim = self.num_v_heads_local_tp + self.a_log_dim = self.num_v_heads_local_tp + + if self.config.deterministic_mode: + self.gated_delta_rule = torch_chunk_gated_delta_rule + else: + self.gated_delta_rule = chunk_gated_delta_rule + + def _get_feat_dim_split(self, cp_size_headwise: int) -> tuple[int, int, int, int]: + """Return GDN1 qkv/z/beta/alpha split sizes for a runtime headwise CP size.""" + return ( + (self.qk_dim_local_tp * 2 + self.v_dim_local_tp) // cp_size_headwise, + self.v_dim_local_tp // cp_size_headwise, + self.num_value_heads // self.tp_size // cp_size_headwise, + self.num_value_heads // self.tp_size // cp_size_headwise, + ) + + @jit_fuser + def _compute_gates( + self, + A_log_local_cp: torch.Tensor, + dt_bias_local_cp: torch.Tensor, + batch: int, + seq_len: int, + *gate_feats: tuple[torch.Tensor], + ) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Compute the per-head log-decay g and the write strength beta.""" + # ``gate_feats`` arrives in ``in_proj_split_names`` order: beta, then alpha. + beta, alpha = gate_feats + g = -A_log_local_cp.exp() * F.softplus(alpha.float() + dt_bias_local_cp) # In fp32 + beta = beta.sigmoid() + return g, {"beta": beta.contiguous()} + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + inference_context: Optional[BaseInferenceContext] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + sequence_len_offset: Optional[int] = None, + *, + pg_collection: Optional[ProcessGroupCollection] = None, + inference_params: Optional[BaseInferenceContext] = None, + **kwargs, + ): + """ + Perform a forward pass through the GDN module. + + Return: + (tuple[torch.Tensor, torch.Tensor]) GDN output and bias. + """ + + inference_context = deprecate_inference_params(inference_context, inference_params) + + base_cp_group = pg_collection.cp if pg_collection is not None else self.pg_collection.cp + cp_group = resolve_cp_group(base_cp_group, packed_seq_params) + if self.config.linear_cp_mode == "chunkwise": + cp_group_chunkwise = cp_group + cp_group_headwise = None + elif self.config.linear_cp_mode == "headwise": + cp_group_chunkwise = None + cp_group_headwise = cp_group + elif cp_group.size() == 1: + cp_group_chunkwise = None + cp_group_headwise = None + else: + raise ValueError( + f"Unsupported linear_cp_mode {self.config.linear_cp_mode!r}; " + "expected 'headwise' or 'chunkwise'." + ) + cp_size_chunkwise = cp_group_chunkwise.size() if cp_group_chunkwise is not None else 1 + cp_size_headwise = cp_group_headwise.size() if cp_group_headwise is not None else 1 + cp_size_runtime = cp_group.size() + + seq_len_local, batch, _ = hidden_states.shape + seq_len_post_headwise = seq_len_local * self.sp_size * cp_size_headwise + seq_len_global = seq_len_post_headwise * cp_size_chunkwise + + if inference_context is not None: + assert ( + inference_context.is_static_batching() + ), "GDN does not currently support dynamic inference batching." + assert not self.config.sequence_parallel + # TODO: support inference + raise NotImplementedError("GDN does not support inference for now.") + + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + assert batch == 1, "Packed sequence expects batch dimension to be 1" + assert ( + not self.config.deterministic_mode + ), "Packed sequence does not support deterministic mode." + + # Resolve cu_seqlens with alignment padding handling. + cu_seqlens_q = self._resolve_cu_seqlens( + packed_seq_params.cu_seqlens_q_padded, + packed_seq_params.cu_seqlens_q, + seq_len_global, + "cu_seqlens_q", + cp_size=cp_size_runtime, + ) + cu_seqlens_kv = self._resolve_cu_seqlens( + packed_seq_params.cu_seqlens_kv_padded, + packed_seq_params.cu_seqlens_kv, + seq_len_global, + "cu_seqlens_kv", + cp_size=cp_size_runtime, + ) + assert torch.equal(cu_seqlens_q, cu_seqlens_kv), ( + "Currently only support cu_seqlens_q equals to cu_seqlens_kv, " + f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}" + ) + num_packed_seqs = cu_seqlens_q.shape[0] - 1 + assert num_packed_seqs > 0, ( + "Number of packed sequences must be greater than 0, " + f"but got {cu_seqlens_q=} and {cu_seqlens_kv=}" + ) + else: + cu_seqlens_q = None + cu_seqlens_kv = None + + if cp_size_chunkwise > 1: + if cu_seqlens_q is None: + cache_key = (seq_len_global, batch) + cached = self._chunkwise_cp_context_cache.get(cache_key) + if cached is None: + cached_cu_seqlens = ( + torch.arange( + batch + 1, device=torch.cuda.current_device(), dtype=torch.long + ) + * seq_len_global + ) + cached_ctx = build_cp_context( + cu_seqlens=cached_cu_seqlens, + group=cp_group_chunkwise, + conv1d_kernel_size=self.conv_kernel_dim, + ) + cached = (cached_cu_seqlens, cached_ctx) + self._chunkwise_cp_context_cache[cache_key] = cached + cu_seqlens_q, chunkwise_cp_context = cached + else: + chunkwise_cp_context = build_cp_context( + cu_seqlens=cu_seqlens_q, + group=cp_group_chunkwise, + conv1d_kernel_size=self.conv_kernel_dim, + ) + else: + chunkwise_cp_context = None + + if self.recompute_gdn and self.training: + + def _checkpointed_compute(hidden_states): + return self._forward_compute( + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ) + + out, out_bias = tensor_parallel.checkpoint(_checkpointed_compute, False, hidden_states) + else: + out, out_bias = self._forward_compute( + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ) + + return out, out_bias + + def _forward_compute( + self, + hidden_states, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + packed_seq_params, + chunkwise_cp_context, + ): + """Core GDN computation (in_proj -> conv1d -> gated_delta_rule -> norm -> out_proj).""" + # Input projection + nvtx_range_push(suffix="in_proj") + qkvzba, _ = self.in_proj(hidden_states) + nvtx_range_pop(suffix="in_proj") + + # TODO: Move CP layout ownership to a model/region-level scheduler so hybrid models can + # enter contiguous layout before GDN regions instead of paying module-local conversions. + if cp_size_chunkwise > 1: + nvtx_range_push(suffix="zigzag_to_contiguous") + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + qkvzba = zigzag_to_contiguous_chunks( + qkvzba, cp_group_chunkwise, seq_dim=0, cu_seqlens=cu_seqlens_q + ) + else: + qkvzba = zigzag_to_contiguous_chunks(qkvzba, cp_group_chunkwise, seq_dim=0) + nvtx_range_pop(suffix="zigzag_to_contiguous") + + qkvzba, thd_cp_a2a_inv = a2a_cp_to_hp( + qkvzba, + self.in_proj_split_sections, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + seq_len_post_headwise, + packed_seq_params, + ) + + if self.gdn_pre_gated_delta_rule_fusion: + if cp_size_chunkwise > 1 and batch > 1: + raise ValueError( + "GDN chunkwise CP with SBHD inputs currently requires micro_batch_size == 1 " + "when cp_context is used. Use packed THD input or micro_batch_size=1." + ) + if cp_size_chunkwise > 1 and self.config.gdn_conv_pad_alignment is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) + nvtx_range_push(suffix="fused_streamed_pre_gated_delta_rule") + seq_idx = ( + packed_seq_params.seq_idx + if packed_seq_params is not None + and packed_seq_params.qkv_format == 'thd' + and cp_size_chunkwise == 1 + else None + ) + fused_cu_seqlens_q = ( + cu_seqlens_q + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + else None + ) + query, key, value, gate, beta, g = self._fused_streamed_pre_gated_delta_rule( + qkvzba, + cu_seqlens_q=fused_cu_seqlens_q, + seq_idx=seq_idx, + cp_group=cp_group_chunkwise if cp_size_chunkwise > 1 else None, + cp_group_headwise=cp_group_headwise, + ) + kernel_inputs = {"q": query, "k": key, "v": value, "g": g, "beta": beta} + nvtx_range_pop(suffix="fused_streamed_pre_gated_delta_rule") + else: + nvtx_range_push(suffix="pre_gated_delta_rule") + if cp_size_chunkwise > 1 and packed_seq_params is None and batch > 1: + # TODO: If additional gated delta rule backends are added, handle this + # SBHD + chunkwise CP + batch>1 case per backend instead of + # unconditionally rejecting it. + raise ValueError( + "GDN chunkwise CP with SBHD inputs currently requires micro_batch_size == 1 " + "when cp_context is used. Use packed THD input or micro_batch_size=1." + ) + if cp_size_chunkwise > 1 and self.config.gdn_conv_pad_alignment is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) + query, key, value, gate, beta, g = self.pre_gated_delta_rule( + qkvzba, + batch, + seq_len_post_headwise, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q, + chunkwise_cp_context, + packed_seq_params=packed_seq_params, + ) + kernel_inputs = {"q": query, "k": key, "v": value, "g": g, "beta": beta} + nvtx_range_pop(suffix="pre_gated_delta_rule") + + nvtx_range_push(suffix="gated_delta_rule") + core_attn_out, _ = self.gated_delta_rule( + **kernel_inputs, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + cu_seqlens=cu_seqlens_q, + cp_context=chunkwise_cp_context, + ) + nvtx_range_pop(suffix="gated_delta_rule") + + if self.recompute_norm_out and self.training: + self.norm_out_checkpoint = tensor_parallel.CheckpointWithoutOutput() + norm_func = partial( + self._gated_norm_and_layout_restore, + thd_cp_a2a_inv=thd_cp_a2a_inv, + batch=batch, + seq_len=seq_len_post_headwise, + packed_seq_params=packed_seq_params, + cp_size_headwise=cp_size_headwise, + cp_group_headwise=cp_group_headwise, + cp_size_chunkwise=cp_size_chunkwise, + cp_group_chunkwise=cp_group_chunkwise, + cu_seqlens_q=cu_seqlens_q, + ) + norm_out = self.norm_out_checkpoint.checkpoint(norm_func, core_attn_out, gate) + else: + norm_out = self._gated_norm_and_layout_restore( + core_attn_out, + gate, + thd_cp_a2a_inv, + batch, + seq_len_post_headwise, + packed_seq_params, + cp_size_headwise, + cp_group_headwise, + cp_size_chunkwise, + cp_group_chunkwise, + cu_seqlens_q, + ) + + # Output projection + nvtx_range_push(suffix="out_proj") + out, out_bias = self.out_proj(norm_out) + nvtx_range_pop(suffix="out_proj") + + if self.recompute_norm_out and self.training: + self.norm_out_checkpoint.discard_output_and_register_recompute(out) + + return out, out_bias + + def _gated_norm_and_layout_restore( + self, + core_attn_out: torch.Tensor, + gate: torch.Tensor, + thd_cp_a2a_inv: torch.Tensor | None, + batch: int, + seq_len: int, + packed_seq_params: PackedSeqParams | None, + cp_size_headwise: int, + cp_group_headwise: torch.distributed.ProcessGroup | None, + cp_size_chunkwise: int, + cp_group_chunkwise: torch.distributed.ProcessGroup | None, + cu_seqlens_q: torch.Tensor | None, + ) -> torch.Tensor: + nvtx_range_push(suffix="gated_norm") + norm_out = self._apply_gated_norm(core_attn_out, gate) + nvtx_range_pop(suffix="gated_norm") + + norm_out = norm_out.reshape(batch, seq_len, -1) + norm_out = norm_out.transpose(0, 1).contiguous() + + # TODO: The planned CP layout refactor should keep consecutive GDN layers contiguous and + # restore zigzag only at SDPA/canonical-layout boundaries. + if cp_size_chunkwise > 1: + nvtx_range_push(suffix="contiguous_to_zigzag") + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + norm_out = contiguous_to_zigzag_chunks( + norm_out, cp_group=cp_group_chunkwise, seq_dim=0, cu_seqlens=cu_seqlens_q + ) + else: + norm_out = contiguous_to_zigzag_chunks( + norm_out, cp_group=cp_group_chunkwise, seq_dim=0 + ) + nvtx_range_pop(suffix="contiguous_to_zigzag") + + return a2a_hp_to_cp( + norm_out, cp_size_headwise, cp_group_headwise, packed_seq_params, thd_cp_a2a_inv + ) + + def pre_gated_delta_rule( + self, + qkvzba, + batch, + seq_len, + cp_size_headwise, + cp_group_headwise, + cu_seqlens_q=None, + chunkwise_cp_context=None, + packed_seq_params=None, + ): + """Prepare QKV, gate, beta, and decay tensors before the gated delta rule.""" + + qkvzba = qkvzba.transpose(0, 1) + qkv, gate, beta, alpha = torch.split( + qkvzba, self._get_feat_dim_split(cp_size_headwise), dim=-1 + ) + gate = gate.reshape(batch, seq_len, -1, self.value_head_dim) + + nvtx_range_push(suffix="conv1d") + kernel_seq_len = qkv.shape[1] + qkv_channels_split_sections = [ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + ] + conv1d_weight = get_parameter_local_cp( + self.conv1d.weight, + dim=0, + cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, + ) + conv1d_bias = ( + get_parameter_local_cp( + self.conv1d.bias, + dim=0, + cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, + ) + if self.conv_bias + else None + ) + if self.config.deterministic_mode: + qkv = qkv.transpose(1, 2).contiguous() + conv_out = F.conv1d( + input=qkv, + weight=conv1d_weight, + bias=conv1d_bias, + stride=self.conv1d.stride, + padding=self.conv1d.padding, + dilation=self.conv1d.dilation, + groups=self.conv_dim_local_tp // cp_size_headwise, + ) + qkv = self.act_fn(conv_out[..., :kernel_seq_len]) + qkv = qkv.transpose(1, 2) + else: + assert self.activation in ["silu", "swish"] + orig_seq = qkv.shape[1] + pad_n = 0 + conv_input = qkv.contiguous() + conv_cu_seqlens = cu_seqlens_q + conv_cp_context = chunkwise_cp_context + if self.config.gdn_conv_pad_alignment is not None: + if packed_seq_params is None or cu_seqlens_q is None: + raise ValueError( + "gdn_conv_pad_alignment is only supported with packed sequence " + "parameters in THD format. SBHD inputs do not need causal-conv padding." + ) + if chunkwise_cp_context is not None: + raise ValueError( + "gdn_conv_pad_alignment is incompatible with GDN chunkwise CP. Padding " + "chunk-local causal-conv inputs can change later chunk numerics." + ) + pad_n = -orig_seq % self.config.gdn_conv_pad_alignment + if pad_n > 0: + conv_input = torch.nn.functional.pad(conv_input, (0, 0, 0, pad_n)) + conv_cu_seqlens = cu_seqlens_q.clone() + conv_cu_seqlens[-1] += pad_n + qkv, _ = causal_conv1d( + x=conv_input, + weight=conv1d_weight.squeeze(1), + bias=conv1d_bias, + activation=self.activation, + initial_state=None, + output_final_state=False, + cu_seqlens=conv_cu_seqlens, + cp_context=conv_cp_context, + ) + if pad_n > 0: + qkv = qkv[:, :orig_seq, :] + nvtx_range_pop(suffix="conv1d") + + A_log_local_cp = get_parameter_local_cp(self.A_log, dim=0, cp_group=cp_group_headwise) + dt_bias_local_cp = get_parameter_local_cp(self.dt_bias, dim=0, cp_group=cp_group_headwise) + + nvtx_range_push(suffix="prepare_input_for_gated_delta_rule") + kernel_inputs = self._prepare_input_for_gated_delta_rule( + qkv, + gate, + A_log_local_cp, + dt_bias_local_cp, + batch, + kernel_seq_len, + beta, + alpha, + cp_size_headwise=cp_size_headwise, + ) + nvtx_range_pop(suffix="prepare_input_for_gated_delta_rule") + + gate = kernel_inputs.pop("gate") + + return ( + kernel_inputs["q"], + kernel_inputs["k"], + kernel_inputs["v"], + gate, + kernel_inputs["beta"], + kernel_inputs["g"], + ) + + def _fused_streamed_pre_gated_delta_rule( + self, qkvzba, cu_seqlens_q=None, seq_idx=None, cp_group=None, cp_group_headwise=None + ): + """Call the streamed fused pre-GDR wrapper.""" + + try: + from megatron.core.fusions.fused_pre_gated_delta_rule import ( + fused_streamed_pre_gated_delta_rule, + ) + except ImportError as exc: + raise ImportError( + "gdn_pre_gated_delta_rule_fusion requires the streamed pre-GDR fusion " + "dependencies, including causal-conv1d." + ) from exc + + qkv_channels_split_sections = [ + self.qk_dim_local_tp, + self.qk_dim_local_tp, + self.v_dim_local_tp, + ] + conv1d_weight = get_parameter_local_cp( + self.conv1d.weight, + dim=0, + cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, + ) + conv1d_bias = ( + get_parameter_local_cp( + self.conv1d.bias, + dim=0, + cp_group=cp_group_headwise, + split_sections=qkv_channels_split_sections, + ) + if self.conv_bias + else None + ) + A_log = get_parameter_local_cp(self.A_log, dim=0, cp_group=cp_group_headwise) + dt_bias = get_parameter_local_cp(self.dt_bias, dim=0, cp_group=cp_group_headwise) + num_value_heads = A_log.numel() + num_key_heads = (conv1d_weight.shape[0] - num_value_heads * self.value_head_dim) // ( + 2 * self.key_head_dim + ) + + return fused_streamed_pre_gated_delta_rule( + qkvzba, + conv1d_weight, + conv1d_bias, + A_log, + dt_bias, + num_key_heads=num_key_heads, + num_value_heads=num_value_heads, + key_head_dim=self.key_head_dim, + value_head_dim=self.value_head_dim, + use_qk_l2norm=self.use_qk_l2norm, + cu_seqlens=cu_seqlens_q, + seq_idx=seq_idx, + cp_group=cp_group, + ) + + +#################### +# Torch native gated delta rule +#################### +def torch_chunk_gated_delta_rule( + q, + k, + v, + g, + beta, + chunk_size=64, + initial_state=None, + output_final_state=False, + use_qk_l2norm_in_kernel=False, + cu_seqlens=None, + cp_context=None, + scale=None, +) -> tuple[torch.Tensor, torch.Tensor | None]: + # pylint: disable=line-too-long + ''' + Torch-native implementation of chunked gated delta rule for deterministic mode. + Need this because FLA is not deterministic. + + ``scale`` defaults to ``1 / sqrt(K)``, matching the FLA kernel. + + Reference: https://github.com/huggingface/transformers/blob/144c8ce2809a2e21914017652700e1ecb450501e/src/transformers/models/qwen3_next/modeling_qwen3_next.py#L470-L547 + ''' + + assert ( + cu_seqlens is None + ), "cu_seqlens is not supported for torch_chunk_gated_delta_rule for now." + assert ( + cp_context is None + ), "cp_context is not supported for torch_chunk_gated_delta_rule for now." + + initial_dtype = q.dtype + if use_qk_l2norm_in_kernel: + q = l2norm(q, dim=-1, eps=1e-6) + k = l2norm(k, dim=-1, eps=1e-6) + q, k, v, beta, g = [ + x.transpose(1, 2).contiguous().to(torch.float32) for x in (q, k, v, beta, g) + ] + + batch_size, num_heads, sequence_length, k_head_dim = k.shape + v_head_dim = v.shape[-1] + pad_size = (chunk_size - sequence_length % chunk_size) % chunk_size + q = F.pad(q, (0, 0, 0, pad_size)) + k = F.pad(k, (0, 0, 0, pad_size)) + v = F.pad(v, (0, 0, 0, pad_size)) + beta = F.pad(beta, (0, pad_size)) + g = F.pad(g, (0, pad_size)) + total_sequence_length = sequence_length + pad_size + if scale is None: + scale = 1 / (q.shape[-1] ** 0.5) + q = q * scale + + v_beta = v * beta.unsqueeze(-1) + k_beta = k * beta.unsqueeze(-1) + # reshape to chunks + q, k, v, k_beta, v_beta = [ + x.reshape(x.shape[0], x.shape[1], -1, chunk_size, x.shape[-1]) + for x in (q, k, v, k_beta, v_beta) + ] + g = g.reshape(g.shape[0], g.shape[1], -1, chunk_size) + mask = torch.triu( + torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=0 + ) + + # chunk decay + g = g.cumsum(dim=-1) + decay_mask = ((g.unsqueeze(-1) - g.unsqueeze(-2)).tril().exp().float()).tril() + attn = -((k_beta @ k.transpose(-1, -2)) * decay_mask).masked_fill(mask, 0) + for i in range(1, chunk_size): + row = attn[..., i, :i].clone() + sub = attn[..., :i, :i].clone() + attn[..., i, :i] = row + (row.unsqueeze(-1) * sub).sum(-2) + attn = attn + torch.eye(chunk_size, dtype=attn.dtype, device=attn.device) + v = attn @ v_beta + k_cumdecay = attn @ (k_beta * g.exp().unsqueeze(-1)) + last_recurrent_state = ( + torch.zeros(batch_size, num_heads, k_head_dim, v_head_dim).to(v) + if initial_state is None + else initial_state.to(v) + ) + core_attn_out = torch.zeros_like(v) + mask = torch.triu( + torch.ones(chunk_size, chunk_size, dtype=torch.bool, device=q.device), diagonal=1 + ) + + # for each chunk + for i in range(0, total_sequence_length // chunk_size): + q_i, k_i, v_i = q[:, :, i], k[:, :, i], v[:, :, i] + attn = (q_i @ k_i.transpose(-1, -2) * decay_mask[:, :, i]).masked_fill_(mask, 0) + v_prime = (k_cumdecay[:, :, i]) @ last_recurrent_state + v_new = v_i - v_prime + attn_inter = (q_i * g[:, :, i, :, None].exp()) @ last_recurrent_state + core_attn_out[:, :, i] = attn_inter + attn @ v_new + last_recurrent_state = ( + last_recurrent_state * g[:, :, i, -1, None, None].exp() + + (k_i * (g[:, :, i, -1, None] - g[:, :, i]).exp()[..., None]).transpose(-1, -2) @ v_new + ) + + if not output_final_state: + last_recurrent_state = None + core_attn_out = core_attn_out.reshape( + core_attn_out.shape[0], core_attn_out.shape[1], -1, core_attn_out.shape[-1] + ) + core_attn_out = core_attn_out[:, :, :sequence_length] + core_attn_out = core_attn_out.transpose(1, 2).contiguous().to(initial_dtype) + return core_attn_out, last_recurrent_state diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 1f11d8d2d7b..841521950fb 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -598,7 +598,7 @@ class TransformerConfig(ModelParallelConfig): recompute_modules: Optional[List[str]] = None """The submodules to recompute. choices: "core_attn", "moe_act", "layernorm", "mla_up_proj", "mlp", "moe", - "shared_experts", "mhc", "gdn". + "shared_experts", "mhc", "gdn", "gdn_norm_out". default: ["core_attn"]. "core_attn": recompute the core attention part of the transformer layer. "moe_act": recompute the MoE MLP activation function. @@ -613,7 +613,10 @@ class TransformerConfig(ModelParallelConfig): "gdn": recompute the entire GatedDeltaNet module (in_proj, conv1d, gated delta rule, gated norm, CP all-to-all and out_proj). Requires experimental_attention_variant="gated_delta_net". - "moe_act", "layernorm", "mla_up_proj", and "mhc" use output-discarding checkpointing, + "gdn_norm_out": recompute the GatedDeltaNet gated normalization output via + CheckpointWithoutOutput. Requires experimental_attention_variant="gated_delta_net". + "moe_act", "layernorm", "mla_up_proj", "mhc", and "gdn_norm_out" use + output-discarding checkpointing, "core_attn", "mlp", "moe", "shared_experts", and "gdn" use normal checkpointing. """ @@ -2184,6 +2187,7 @@ def __post_init__(self): "shared_experts", "mhc", "gdn", + "gdn_norm_out", } invalid_modules = set(self.recompute_modules) - allowed_modules assert not invalid_modules, ( @@ -2202,6 +2206,15 @@ def __post_init__(self): "multi_latent_attention." ) + if ( + "gdn_norm_out" in self.recompute_modules + and self.experimental_attention_variant != "gated_delta_net" + ): + raise ValueError( + "gdn_norm_out in recompute_modules is only supported with " + "experimental_attention_variant='gated_delta_net'." + ) + if ( "gdn" in self.recompute_modules and self.experimental_attention_variant != "gated_delta_net" @@ -2211,6 +2224,12 @@ def __post_init__(self): "experimental_attention_variant='gated_delta_net'." ) + if "gdn" in self.recompute_modules and "gdn_norm_out" in self.recompute_modules: + raise ValueError( + "'gdn' and 'gdn_norm_out' in recompute_modules cannot be used together. " + "'gdn' recomputes the full GatedDeltaNet module, including gated norm." + ) + if "core_attn" in self.recompute_modules: warnings.warn( "If you are using transformer_engine as the transformer implementation, " diff --git a/pyproject.toml b/pyproject.toml index 5405a2edab5..ff718cfabb2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -94,7 +94,7 @@ dev = [ "opentelemetry-api~=1.33.1", "mamba-ssm~=2.2", "causal-conv1d~=1.5", - "flash-linear-attention>=0.4.2,<0.5", + "flash-linear-attention[tilelang]==0.5.1", "megatron-energon[av_decode]~=7.0", "av", "flashinfer-python>=0.5.0,<0.7.0", @@ -163,7 +163,7 @@ build = [ ] linting = [ "ruff~=0.9.0", - "black==24.4.2", + "black==26.3.0", "isort==5.13.2", "flake8==7.1.0", "pylint==3.2.6", @@ -190,8 +190,34 @@ override-dependencies = [ "torch; sys_platform == 'never'", "torchvision; sys_platform == 'never'", "triton; sys_platform == 'never'", + # TorchX 0.7.0 caps urllib3 below 1.27; require the fix for GHSA-38jv-5279-wg99. + "urllib3>=2.6.3", ] +[[tool.uv.dependency-metadata]] +name = "flash-mla" +version = "1.0.0+b7643bd" +requires-dist = [] + +[[tool.uv.dependency-metadata]] +name = "transformer-engine" +version = "2.14.0+f031cf87" +requires-dist = [ + "einops", + "importlib-metadata", + "nvdlfw-inspect", + "onnx", + "onnxscript", + "packaging", + "pydantic", + "torch", +] + +[[tool.uv.dependency-metadata]] +name = "fast-hadamard-transform" +version = "1.0.4.post1" +requires-dist = ["torch", "packaging", "ninja"] + [tool.uv.sources] flash_mla = [ @@ -218,7 +244,7 @@ line_length = 100 skip_string_normalization = true # recognized by future versions, disallows to reformat code with incompatible versions # Matches NeMO version so people working on both codebases don't need two different version of black installed -required_version = "24" +required_version = "26" skip_magic_trailing_comma = true include = '\.pyi?$' exclude = ''' diff --git a/tests/test_utils/recipes/h100/unit-tests.yaml b/tests/test_utils/recipes/h100/unit-tests.yaml index 054b4a33f3c..52e573f79d6 100644 --- a/tests/test_utils/recipes/h100/unit-tests.yaml +++ b/tests/test_utils/recipes/h100/unit-tests.yaml @@ -172,7 +172,21 @@ products: scope: [unit-tests] n_repeat: [1] time_limit: [1800] - - test_case: [tests/unit_tests/ssm/test_gated_delta_net.py] + - test_case: [tests/unit_tests/ssm/gated_delta_net/test_gdn.py] + products: + - environment: [lts, dev] + tag: [latest] + scope: [unit-tests] + n_repeat: [1] + time_limit: [1800] + - test_case: [tests/unit_tests/ssm/gated_delta_net/test_gdn_fusion.py] + products: + - environment: [lts, dev] + tag: [latest] + scope: [unit-tests] + n_repeat: [1] + time_limit: [1800] + - test_case: [tests/unit_tests/ssm/gated_delta_net/test_gdn_parallel.py] products: - environment: [lts, dev] tag: [latest] diff --git a/tests/unit_tests/ssm/gated_delta_net/test_gdn.py b/tests/unit_tests/ssm/gated_delta_net/test_gdn.py new file mode 100644 index 00000000000..93d91ff05ec --- /dev/null +++ b/tests/unit_tests/ssm/gated_delta_net/test_gdn.py @@ -0,0 +1,809 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import copy +import inspect +import os +from unittest import mock + +import pytest +import torch +import torch.nn.functional as F + +from megatron.core import parallel_state +from megatron.core.dist_checkpointing.mapping import ShardedTensorFactory +from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_experimental_attention_variant_module_spec, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.ssm.gated_delta_net import GatedDeltaNet, torch_chunk_gated_delta_rule +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from tests.unit_tests.test_utilities import Utils +from tests.unit_tests.transformer.test_multi_latent_attention import ( + make_test_packed_seq_params, + make_test_packed_seq_params_with_padding, +) + +try: + import fla + + HAVE_FLA = True +except ImportError: + HAVE_FLA = False + +# https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#nccl-multi-rank-gpu-enable +# NVLS doesn't support one single GPU to be shared by multiple ranks, so disable this in test +os.environ.update({"NCCL_NVLS_ENABLE": "0"}) + +try: + from causal_conv1d.cpp_functions import causal_conv1d_bwd_function +except ImportError: + HAVE_FUSED_PRE_GDR = False +else: + HAVE_FUSED_PRE_GDR = callable(causal_conv1d_bwd_function) + + +def _make_gdn_config(**overrides): + config_kwargs = { + "hidden_size": 128, + "linear_conv_kernel_dim": 2, + "linear_key_head_dim": 32, + "linear_value_head_dim": 32, + "linear_num_key_heads": 4, + "linear_num_value_heads": 8, + "num_layers": 1, + "normalization": "RMSNorm", + "use_cpu_initialization": True, + "layernorm_zero_centered_gamma": True, + "num_attention_heads": 8, + "activation_func": F.silu, + "bf16": True, + "experimental_attention_variant": "gated_delta_net", + "linear_attention_freq": [1], + "transformer_impl": "transformer_engine", + } + config_kwargs.update(overrides) + return TransformerConfig(**config_kwargs) + + +def test_gdn_pre_gated_delta_rule_fusion_defaults_to_disabled(): + config = _make_gdn_config() + assert not config.gdn_pre_gated_delta_rule_fusion + + +def test_gdn_pre_gated_delta_rule_fusion_accepts_gdn_variant(): + config = _make_gdn_config(gdn_pre_gated_delta_rule_fusion=True) + assert config.gdn_pre_gated_delta_rule_fusion + + +def test_gdn_pre_gated_delta_rule_fusion_requires_gdn_variant(): + with pytest.raises(ValueError, match="experimental_attention_variant='gated_delta_net'"): + _make_gdn_config( + experimental_attention_variant=None, + linear_attention_freq=None, + gdn_pre_gated_delta_rule_fusion=True, + ) + + +def test_gdn_norm_out_recompute_accepts_gdn_variant(): + config = _make_gdn_config(recompute_granularity="selective", recompute_modules=["gdn_norm_out"]) + assert "gdn_norm_out" in config.recompute_modules + + +def test_gdn_and_norm_out_recompute_are_mutually_exclusive(): + with pytest.raises(ValueError, match="'gdn' and 'gdn_norm_out'"): + _make_gdn_config( + recompute_granularity="selective", recompute_modules=["gdn", "gdn_norm_out"] + ) + + +def test_gdn_norm_out_recompute_requires_gdn_variant(): + with pytest.raises(ValueError, match="experimental_attention_variant='gated_delta_net'"): + _make_gdn_config( + experimental_attention_variant=None, + linear_attention_freq=None, + recompute_granularity="selective", + recompute_modules=["gdn_norm_out"], + ) + + +def test_gdn_conv_pad_alignment_rejects_chunkwise_cp(): + with pytest.raises(AssertionError, match="gdn_conv_pad_alignment is incompatible"): + _make_gdn_config( + context_parallel_size=2, linear_cp_mode="chunkwise", gdn_conv_pad_alignment=4096 + ) + + +def test_gdn_chunkwise_cp_head_divisibility_ignores_cp_size(): + config = _make_gdn_config( + tensor_model_parallel_size=2, + context_parallel_size=4, + linear_cp_mode="chunkwise", + linear_num_key_heads=4, + linear_num_value_heads=8, + ) + assert config.linear_cp_mode == "chunkwise" + + +def test_torch_chunk_gated_delta_rule_preserves_public_signature(): + signature = inspect.signature(torch_chunk_gated_delta_rule) + assert tuple(signature.parameters) == ( + "q", + "k", + "v", + "g", + "beta", + "chunk_size", + "initial_state", + "output_final_state", + "use_qk_l2norm_in_kernel", + "cu_seqlens", + "cp_context", + "scale", + ) + + +def test_gdn_headwise_cp_head_divisibility_includes_cp_size(): + with pytest.raises(AssertionError, match="linear_head_parallel_size"): + _make_gdn_config( + tensor_model_parallel_size=2, + context_parallel_size=4, + linear_cp_mode="headwise", + linear_num_key_heads=4, + linear_num_value_heads=8, + ) + + +@pytest.mark.parametrize( + ("tp_size", "sp", "cp_size", "linear_cp_mode"), + [ + # cp_size=1: the CP path is inactive, so linear_cp_mode choice is irrelevant. + # Cover the "chunkwise" default and skip the "headwise" variants for brevity. + (1, False, 1, None), + (2, False, 1, None), + (2, True, 1, None), + # cp_size=2: exercise both CP paths. + (1, False, 2, "headwise"), + (2, False, 2, "headwise"), + (2, True, 2, "headwise"), + (1, False, 2, "chunkwise"), + (2, False, 2, "chunkwise"), + (2, True, 2, "chunkwise"), + ], +) +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.internal +class TestGatedDeltaNet: + + @pytest.fixture(scope='function', autouse=True) + def setup_method(self, tp_size, sp, cp_size, linear_cp_mode): + # Initialize parallel and random seed + Utils.initialize_model_parallel( + tensor_model_parallel_size=tp_size, + pipeline_model_parallel_size=1, + context_parallel_size=cp_size, + ) + model_parallel_cuda_manual_seed(123) + self.tp_size = tp_size + self.cp_size = cp_size + self.sp_size = tp_size if sp else 1 + self.linear_cp_mode = linear_cp_mode + if self.linear_cp_mode == "headwise": + self.cp_size_chunkwise = 1 + self.cp_size_headwise = self.cp_size + elif self.linear_cp_mode == "chunkwise": + self.cp_size_chunkwise = self.cp_size + self.cp_size_headwise = 1 + elif self.cp_size == 1: + self.cp_size_chunkwise = 1 + self.cp_size_headwise = 1 + else: + raise ValueError(f"Invalid linear CP mode: {self.linear_cp_mode}") + + # Get TP and CP process groups from device mesh + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group) + + # Initialize model, with the same config as Qwen Next except `num_layers` + self.transformer_config = TransformerConfig( + hidden_size=2048, + linear_conv_kernel_dim=4, + linear_key_head_dim=128, + linear_value_head_dim=128, + linear_num_key_heads=16, + linear_num_value_heads=32, + num_layers=1, + normalization="RMSNorm", + use_cpu_initialization=True, + layernorm_zero_centered_gamma=True, + num_attention_heads=16, + num_query_groups=2, + activation_func=F.silu, + bf16=True, + tensor_model_parallel_size=tp_size, + sequence_parallel=sp, + context_parallel_size=cp_size, + experimental_attention_variant="gated_delta_net", + linear_attention_freq=[1], + linear_cp_mode=self.linear_cp_mode, + transformer_impl="transformer_engine", + ) + gdn_submodules = get_experimental_attention_variant_module_spec( + config=self.transformer_config + ).submodules + + self.gdn = GatedDeltaNet( + self.transformer_config, + submodules=gdn_submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + self.gdn = self.gdn.cuda().bfloat16() + + def teardown_method(self): + Utils.destroy_model_parallel() + + def test_gpu_forward(self): + gdn = self.gdn + + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 + seq_length = 64 + hidden_states = torch.ones( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + attention_mask = None + + output, bias = gdn(hidden_states, attention_mask) + + assert output.dim() == 3, f"Output too many dimensions ({output.shape=})" + assert output.shape[0] == seq_length // self.sp_size // self.cp_size, ( + f"Output shape {output.shape[0]=} mismatch with " + f" {seq_length=} // {self.sp_size=} // {self.cp_size=}." + ) + assert ( + output.shape[1] == micro_batch_size + ), f"Output shape {output.shape[1]=} mismatch with {micro_batch_size=}" + assert ( + output.shape[2] == gdn.config.hidden_size + ), f"Output shape {output.shape[2]=} mismatch with {gdn.config.hidden_size=}" + assert ( + output.dtype == hidden_states.dtype + ), f"Output dtype {output.dtype=} mismatch with {hidden_states.dtype=}" + + @pytest.mark.flaky_in_dev # Issue #5473 + def test_selective_recompute_gdn(self): + """Whole-module 'gdn' recompute must match the non-recompute forward and gradients. + + The same module/input is run twice (recompute off, then on); the forward output and + all parameter / input gradients must agree within a tight tolerance (rtol/atol=1e-4). + The recompute path is run-to-run deterministic on these kernels (empirically bitwise), + so a tolerance well below the bf16 floor is expected to hold. + """ + gdn = self.gdn + gdn.train() + + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 + seq_length = 64 + torch.manual_seed(1234) + base_input = torch.randn( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + def run(recompute): + gdn.recompute_gdn = recompute + gdn.zero_grad(set_to_none=True) + hidden_states = base_input.clone().detach().requires_grad_(True) + output, _ = gdn(hidden_states, None) + output.float().square().mean().backward() + param_grads = { + name: param.grad.detach().clone() + for name, param in gdn.named_parameters() + if param.grad is not None + } + return output.detach().clone(), hidden_states.grad.detach().clone(), param_grads + + try: + out_ref, dinput_ref, pgrad_ref = run(recompute=False) + out_rc, dinput_rc, pgrad_rc = run(recompute=True) + finally: + gdn.recompute_gdn = False + + torch.testing.assert_close(out_rc, out_ref, rtol=1e-4, atol=1e-4) + torch.testing.assert_close(dinput_rc, dinput_ref, rtol=1e-4, atol=1e-4) + assert pgrad_ref.keys() == pgrad_rc.keys(), "recompute changed the set of grad params" + assert len(pgrad_ref) > 0, "expected at least one parameter gradient" + for name in pgrad_ref: + torch.testing.assert_close( + pgrad_rc[name], + pgrad_ref[name], + rtol=1e-4, + atol=1e-4, + msg=lambda m, n=name: f"gradient mismatch for parameter '{n}': {m}", + ) + + @pytest.mark.flaky_in_dev # Issue #5473 + def test_selective_recompute_gdn_norm_out(self): + """Output-discarding 'gdn_norm_out' recompute must preserve outputs and gradients.""" + gdn = self.gdn + gdn.train() + + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 + seq_length = 64 + torch.manual_seed(1234) + base_input = torch.randn( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + def run(recompute_norm_out): + gdn.recompute_gdn = False + gdn.recompute_norm_out = recompute_norm_out + gdn.norm_out_checkpoint = None + gdn.zero_grad(set_to_none=True) + hidden_states = base_input.clone().detach().requires_grad_(True) + output, _ = gdn(hidden_states, None) + output.float().square().mean().backward() + param_grads = { + name: param.grad.detach().clone() + for name, param in gdn.named_parameters() + if param.grad is not None + } + return output.detach().clone(), hidden_states.grad.detach().clone(), param_grads + + try: + out_ref, dinput_ref, pgrad_ref = run(recompute_norm_out=False) + out_rc, dinput_rc, pgrad_rc = run(recompute_norm_out=True) + finally: + gdn.recompute_norm_out = False + gdn.norm_out_checkpoint = None + + torch.testing.assert_close(out_rc, out_ref, rtol=1e-4, atol=1e-4) + torch.testing.assert_close(dinput_rc, dinput_ref, rtol=1e-4, atol=1e-4) + assert pgrad_ref.keys() == pgrad_rc.keys(), "recompute changed the set of grad params" + assert len(pgrad_ref) > 0, "expected at least one parameter gradient" + for name in pgrad_ref: + torch.testing.assert_close( + pgrad_rc[name], + pgrad_ref[name], + rtol=1e-4, + atol=1e-4, + msg=lambda m, n=name: f"gradient mismatch for parameter '{n}': {m}", + ) + + def test_gpu_forward_rejects_sbhd_chunkwise_cp_batch_gt_one(self): + if not (self.linear_cp_mode == "chunkwise" and self.cp_size > 1): + pytest.skip("Only chunkwise CP with CP>1 uses the FLA CP batch guard.") + + gdn = self.gdn + + micro_batch_size = 2 + seq_length = 64 + hidden_states = torch.ones( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + with pytest.raises(ValueError, match="requires micro_batch_size == 1"): + gdn(hidden_states, None) + + def test_gpu_forward_rejects_sbhd_conv_padding(self): + gdn = self.gdn + gdn.config.gdn_conv_pad_alignment = 4096 + + micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 + seq_length = 64 + hidden_states = torch.ones( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + expected_error = ( + "incompatible with GDN chunkwise CP" + if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 + else "only supported with packed sequence" + ) + with pytest.raises(ValueError, match=expected_error): + gdn(hidden_states, None) + + def test_deterministic_mode(self): + if self.cp_size > 1: + pytest.skip( + "deterministic_mode uses torch_chunk_gated_delta_rule, which does not support CP." + ) + + tp_group = parallel_state.get_tensor_model_parallel_group() + cp_group = parallel_state.get_context_parallel_group() + pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group) + + det_config = copy.deepcopy(self.transformer_config) + det_config.deterministic_mode = True + + gdn_submodules = get_experimental_attention_variant_module_spec( + config=det_config + ).submodules + + model_parallel_cuda_manual_seed(42) + torch.manual_seed(42) + gdn = ( + GatedDeltaNet( + det_config, + submodules=gdn_submodules, + layer_number=1, + bias=False, + conv_bias=False, + conv_init=1.0, + use_qk_l2norm=True, + A_init_range=(1, 16), + pg_collection=pg_collection, + ) + .cuda() + .bfloat16() + ) + + # deterministic_mode must select the torch-native kernel, not FLA. + assert gdn.gated_delta_rule is torch_chunk_gated_delta_rule + + micro_batch_size = 2 + seq_length = 64 + torch.manual_seed(0) + base_input = torch.randn( + (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + + def run(): + hidden_states = base_input.clone().requires_grad_(True) + output, _ = gdn(hidden_states, None) + output.float().sum().backward() + grads = { + name: param.grad.detach().clone() + for name, param in gdn.named_parameters() + if param.grad is not None + } + gdn.zero_grad(set_to_none=True) + return output.detach().clone(), grads, hidden_states.grad.detach().clone() + + out1, grads1, input_grad1 = run() + out2, grads2, input_grad2 = run() + + rank = torch.distributed.get_rank() + assert torch.equal(out1, out2), f"Output not reproducible ({rank=})" + assert torch.equal(input_grad1, input_grad2), f"Input grad not reproducible ({rank=})" + assert set(grads1.keys()) == set(grads2.keys()) + for name in grads1: + assert torch.equal( + grads1[name], grads2[name] + ), f"Grad not reproducible for {name} ({rank=})" + + def test_module_construction(self): + gdn = self.gdn + assert gdn.in_proj_dim == 2 * gdn.qk_dim + 2 * gdn.v_dim + 2 * gdn.num_value_heads + assert gdn.A_log.shape == (gdn.num_value_heads // self.tp_size,) + assert gdn.dt_bias.shape == (gdn.num_value_heads // self.tp_size,) + + def test_sharded_state_dict_splits_gdn_parameters(self): + sharded_sd = self.gdn.sharded_state_dict(prefix="gdn.") + + in_proj_weight = sharded_sd["gdn.in_proj.weight"] + conv1d_weight = sharded_sd["gdn.conv1d.weight"] + assert isinstance(in_proj_weight, ShardedTensorFactory) + assert isinstance(conv1d_weight, ShardedTensorFactory) + + in_proj_chunks = in_proj_weight.build() + assert tuple(chunk.key for chunk in in_proj_chunks) == tuple( + f"gdn.in_proj.weight.{name}" for name in self.gdn.in_proj_split_names + ) + assert sum(chunk.data.numel() for chunk in in_proj_chunks) == in_proj_weight.data.numel() + + conv1d_chunks = conv1d_weight.build() + assert tuple(chunk.key for chunk in conv1d_chunks) == ( + "gdn.conv1d.weight.query", + "gdn.conv1d.weight.key", + "gdn.conv1d.weight.value", + ) + assert sum(chunk.data.numel() for chunk in conv1d_chunks) == conv1d_weight.data.numel() + + def test_jit_compiled_helpers(self): + import torch._dynamo + + gdn = self.gdn + batch = 2 + seq_len = 16 + + device = torch.cuda.current_device() + num_v_heads_local = gdn.num_value_heads // gdn.tp_size // self.cp_size_headwise + qk_dim_local = gdn.qk_dim_local_tp // self.cp_size_headwise + v_dim_local = gdn.v_dim_local_tp // self.cp_size_headwise + + qkv = torch.randn( + batch, seq_len, 2 * qk_dim_local + v_dim_local, device=device, dtype=torch.bfloat16 + ) + gate = torch.randn( + batch, + seq_len, + num_v_heads_local, + gdn.value_head_dim, + device=device, + dtype=torch.bfloat16, + ) + gate_feats = ( + torch.randn(batch, seq_len, num_v_heads_local, device=device, dtype=torch.bfloat16), + torch.randn(batch, seq_len, num_v_heads_local, device=device, dtype=torch.bfloat16), + ) # beta, alpha + + # Disable dynamo so coverage.py can trace through the method bodies, + # which are normally wrapped by @jit_fuser (torch.compile). + A_log_mock = torch.randn(num_v_heads_local, device=device, dtype=torch.bfloat16) + dt_bias_mock = torch.randn(num_v_heads_local, device=device, dtype=torch.bfloat16) + + with torch._dynamo.config.patch(disable=True): + kernel_inputs = gdn._prepare_input_for_gated_delta_rule( + qkv, + gate, + A_log_mock, + dt_bias_mock, + batch, + seq_len, + *gate_feats, + cp_size_headwise=self.cp_size_headwise, + ) + + query = kernel_inputs["q"] + key = kernel_inputs["k"] + value = kernel_inputs["v"] + g = kernel_inputs["g"] + gate_out = kernel_inputs["gate"] + beta_out = kernel_inputs["beta"] + + assert query.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim) + assert key.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim) + assert value.shape == (batch, seq_len, num_v_heads_local, gdn.value_head_dim) + for t in (query, key, value, gate_out, beta_out): + assert t.is_contiguous() + + assert g.dtype == torch.float32 + assert g.shape == (batch, seq_len, num_v_heads_local) + assert beta_out.shape == (batch, seq_len, num_v_heads_local) + + def test_fused_pre_gated_delta_rule_headwise_cp_uses_cp_local_parameters(self): + if not HAVE_FUSED_PRE_GDR: + pytest.skip("causal-conv1d fused backward is not installed.") + if not (self.linear_cp_mode == "headwise" and self.cp_size > 1): + pytest.skip("Only headwise CP with CP>1 needs CP-local fused pre-GDR params.") + + gdn = self.gdn + batch = 2 + seq_len = 16 + qk_channels = gdn.qk_dim_local_tp // self.cp_size_headwise + v_channels = gdn.v_dim_local_tp // self.cp_size_headwise + num_key_heads = qk_channels // gdn.key_head_dim + num_value_heads = v_channels // gdn.value_head_dim + qkvzba_dim = 2 * qk_channels + 2 * v_channels + 2 * num_value_heads + qkvzba = torch.randn( + seq_len, batch, qkvzba_dim, device=torch.cuda.current_device(), dtype=torch.bfloat16 + ) + captured = {} + + def fake_fused_streamed_pre_gated_delta_rule( + qkvzba_arg, + conv1d_weight, + conv1d_bias, + A_log, + dt_bias, + *, + num_key_heads, + num_value_heads, + **kwargs, + ): + captured.update( + { + "qkvzba": qkvzba_arg, + "conv1d_weight": conv1d_weight, + "conv1d_bias": conv1d_bias, + "A_log": A_log, + "dt_bias": dt_bias, + "num_key_heads": num_key_heads, + "num_value_heads": num_value_heads, + "cp_group": kwargs["cp_group"], + } + ) + return tuple(torch.empty(0, device=qkvzba_arg.device) for _ in range(6)) + + with mock.patch( + "megatron.core.fusions.fused_pre_gated_delta_rule." + "fused_streamed_pre_gated_delta_rule", + side_effect=fake_fused_streamed_pre_gated_delta_rule, + ): + gdn._fused_streamed_pre_gated_delta_rule(qkvzba, cp_group_headwise=gdn.pg_collection.cp) + + assert captured["qkvzba"] is qkvzba + assert captured["conv1d_weight"].shape == ( + 2 * qk_channels + v_channels, + 1, + gdn.conv_kernel_dim, + ) + assert captured["conv1d_bias"] is None + assert captured["A_log"].shape == (num_value_heads,) + assert captured["dt_bias"].shape == (num_value_heads,) + assert captured["num_key_heads"] == num_key_heads + assert captured["num_value_heads"] == num_value_heads + assert captured["cp_group"] is None + + def test_gpu_forward_thd_correctness(self): + if self.sp_size > 1: + pytest.skip("Sequence parallel is not supported for this test case.") + if self.cp_size > 1 and self.linear_cp_mode == "chunkwise": + pytest.skip("Chunkwise CP is not supported for this test case.") + + atol, rtol = 3e-4, 3e-4 + + # Input shape + sequence_length = 32 + micro_batch_size = 4 + cu_seqlens = [0, 32, 64, 96, 128] + # sbhd input shape: [sequence length, batch size, hidden size] + sub_sequence_length = sequence_length // self.cp_size + hidden_states_sbhd = torch.rand( + (sub_sequence_length, micro_batch_size, self.gdn.config.hidden_size) + ) + attention_mask_sbhd = None + hidden_states_sbhd = hidden_states_sbhd.cuda().bfloat16() + # thd input shape: [sequence length * batch size, 1, hidden size] + hidden_states_thd = hidden_states_sbhd.transpose(0, 1).contiguous() + hidden_states_thd = hidden_states_thd.view(-1, 1, self.gdn.config.hidden_size) + attention_mask_thd = None + packed_seq_params = make_test_packed_seq_params(cu_seqlens=cu_seqlens) + + # THD format + output_thd, _ = self.gdn( + hidden_states_thd, attention_mask_thd, packed_seq_params=packed_seq_params + ) + # SBHD format + output_sbhd, _ = self.gdn(hidden_states_sbhd, attention_mask_sbhd) + output_sbhd_T = output_sbhd.transpose(0, 1).contiguous().view(*output_thd.shape) + + rank = torch.distributed.get_rank() + assert output_thd.shape[0] == sub_sequence_length * micro_batch_size + assert output_thd.shape[1] == 1 + assert output_thd.shape[2] == self.gdn.config.hidden_size + torch.testing.assert_close( + output_sbhd_T, + output_thd, + atol=atol, + rtol=rtol, + msg=lambda msg: f"Output mismatch ({rank=}): {msg}", + ) + + def test_gpu_forward_thd_padding_correctness(self): + if self.sp_size > 1: + pytest.skip("Sequence parallel is not supported for this test case.") + if self.cp_size > 1 and self.linear_cp_mode == "chunkwise": + pytest.skip("Chunkwise CP is not supported for this test case.") + + atol, rtol = 3e-4, 3e-4 + sequence_length = 32 + micro_batch_size = 4 + + # sbhd input shape: [sequence length, batch size, hidden size] + sub_sequence_length = sequence_length // self.cp_size + hidden_states_sbhd = torch.rand( + (sub_sequence_length, micro_batch_size, self.gdn.config.hidden_size), + device=torch.cuda.current_device(), + dtype=torch.bfloat16, + ) + output_sbhd, _ = self.gdn(hidden_states_sbhd, None) + + # thd input shape: [sequence length * batch size, 1, hidden size] + hidden_states_thd = hidden_states_sbhd.transpose(0, 1).contiguous() + hidden_states_thd = hidden_states_thd.view(-1, 1, self.gdn.config.hidden_size) + output_bshd = output_sbhd.transpose(0, 1).contiguous() + + rank = torch.distributed.get_rank() + + # A) padded branch: prefer *_padded when available. + padded_params = make_test_packed_seq_params_with_padding( + cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 128] + ) + output_thd_padded, _ = self.gdn(hidden_states_thd, None, packed_seq_params=padded_params) + output_thd2bshd = output_thd_padded.view(*output_bshd.shape) + torch.testing.assert_close( + output_bshd[:, :30, :], + output_thd2bshd[:, :30, :], + atol=atol, + rtol=rtol, + msg=lambda msg: f"THD padded output mismatch ({rank=}): {msg}", + ) + + # B) no-padded branch: use actual cu_seqlens when it matches total_sequence_length. + no_padding_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 128]) + output_thd_no_padding, _ = self.gdn( + hidden_states_thd, None, packed_seq_params=no_padding_params + ) + assert output_thd_no_padding.shape == output_thd_padded.shape + + # C) explicit causal-conv padding is only applied to packed inputs and + # should not affect the original unpadded token outputs. + self.gdn.config.gdn_conv_pad_alignment = 48 + output_thd_conv_pad, _ = self.gdn( + hidden_states_thd, None, packed_seq_params=no_padding_params + ) + self.gdn.config.gdn_conv_pad_alignment = None + assert output_thd_conv_pad.shape == output_thd_no_padding.shape + torch.testing.assert_close( + output_thd_conv_pad, + output_thd_no_padding, + atol=atol, + rtol=rtol, + msg=lambda msg: f"THD conv-padded output mismatch ({rank=}): {msg}", + ) + + # D) padded mismatch branch: if *_padded[-1] mismatches total_sequence_length, should raise. + padded_mismatch_params = make_test_packed_seq_params_with_padding( + cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 126] + ) + with pytest.raises(ValueError, match="does not match"): + self.gdn(hidden_states_thd, None, packed_seq_params=padded_mismatch_params) + + # E) actual mismatch branch without *_padded: should raise. + actual_mismatch_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 129]) + with pytest.raises(ValueError, match="does not match"): + self.gdn(hidden_states_thd, None, packed_seq_params=actual_mismatch_params) + + +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +@pytest.mark.internal +class TestGDNCuSeqlensResolve: + + @pytest.fixture + def mock_gdn(self): + class MockGDN: + _resolve_cu_seqlens = GatedDeltaNet._resolve_cu_seqlens + + return MockGDN() + + def test_padded_preferred_when_available(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + padded = torch.tensor([0, 504, 1008], dtype=torch.int32) + result = mock_gdn._resolve_cu_seqlens(padded, actual, 1008, "cu_seqlens_q", cp_size=2) + assert torch.equal(result, padded) + + def test_actual_used_when_no_padding(self, mock_gdn): + actual = torch.tensor([0, 504, 1008], dtype=torch.int32) + result = mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) + assert torch.equal(result, actual) + + def test_raises_when_padding_mismatch(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) + + def test_raises_when_padded_mismatches_total(self, mock_gdn): + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + padded = torch.tensor([0, 504, 1004], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(padded, actual, 1008, "cu_seqlens_q", cp_size=2) + + def test_raises_when_not_divisible_by_cp_size(self, mock_gdn): + actual = torch.tensor([0, 505, 1008], dtype=torch.int32) + with pytest.raises(ValueError, match="must be divisible by cp_size"): + mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) + + def test_cp1_still_validates_total(self, mock_gdn): + mock_gdn.cp_size = 1 + actual = torch.tensor([0, 500, 1000], dtype=torch.int32) + with pytest.raises(ValueError, match="does not match"): + mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=1) diff --git a/tests/unit_tests/ssm/test_gated_delta_net.py b/tests/unit_tests/ssm/gated_delta_net/test_gdn_fusion.py similarity index 55% rename from tests/unit_tests/ssm/test_gated_delta_net.py rename to tests/unit_tests/ssm/gated_delta_net/test_gdn_fusion.py index 06e4c136d66..6b12ec6a1c8 100644 --- a/tests/unit_tests/ssm/test_gated_delta_net.py +++ b/tests/unit_tests/ssm/gated_delta_net/test_gdn_fusion.py @@ -1,49 +1,21 @@ -# Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import os -from functools import partial -from unittest import mock import pytest import torch import torch.nn.functional as F from megatron.core import parallel_state -from megatron.core.models.common.embeddings.rope_utils import ( - get_pos_emb_on_this_cp_rank as get_tensor_on_this_cp_rank, -) from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( get_experimental_attention_variant_module_spec, - get_transformer_block_with_experimental_attention_variant_spec, ) -from megatron.core.models.gpt.gpt_model import GPTModel from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.ssm.gated_delta_net import ( - GatedDeltaNet, - _build_head_perm_for_split_sections, - _build_thd_cp_a2a_perm, - tensor_a2a_cp2hp, - tensor_a2a_hp2cp, -) +from megatron.core.ssm.gated_delta_net import GatedDeltaNet from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed from megatron.core.transformer import TransformerConfig -from megatron.core.utils import unwrap_model -from megatron.training.arguments import parse_args -from megatron.training.checkpointing import load_checkpoint, save_checkpoint -from megatron.training.global_vars import set_args -from megatron.training.training import get_model -from tests.unit_tests.dist_checkpointing import ( - TempNamedDir, - init_basic_mock_args, - init_checkpointing_mock_args, -) from tests.unit_tests.test_utilities import Utils -from tests.unit_tests.transformer.test_attention import _test_parallel_attention_correctness -from tests.unit_tests.transformer.test_multi_latent_attention import ( - make_test_packed_seq_params, - make_test_packed_seq_params_with_padding, -) try: import fla @@ -56,19 +28,6 @@ # NVLS doesn't support one single GPU to be shared by multiple ranks, so disable this in test os.environ.update({"NCCL_NVLS_ENABLE": "0"}) - -def _unpack_sequence(x: torch.Tensor, cu_seqlens: torch.Tensor, dim=1) -> list[torch.Tensor]: - unpacked_x = [] - cu_seqlens_list = cu_seqlens.tolist() - num_seqs = len(cu_seqlens_list) - 1 - for i in range(num_seqs): - idx_start = cu_seqlens_list[i] - idx_end = cu_seqlens_list[i + 1] - chunked_index = [slice(None)] * dim + [slice(idx_start, idx_end)] - unpacked_x.append(x[tuple(chunked_index)]) - return unpacked_x - - try: from causal_conv1d.cpp_functions import causal_conv1d_bwd_function except ImportError: @@ -77,66 +36,6 @@ def _unpack_sequence(x: torch.Tensor, cu_seqlens: torch.Tensor, dim=1) -> list[t HAVE_FUSED_PRE_GDR = callable(causal_conv1d_bwd_function) -def _make_gdn_config(**overrides): - config_kwargs = { - "hidden_size": 128, - "linear_conv_kernel_dim": 2, - "linear_key_head_dim": 32, - "linear_value_head_dim": 32, - "linear_num_key_heads": 4, - "linear_num_value_heads": 8, - "num_layers": 1, - "normalization": "RMSNorm", - "use_cpu_initialization": True, - "layernorm_zero_centered_gamma": True, - "num_attention_heads": 8, - "activation_func": F.silu, - "bf16": True, - "experimental_attention_variant": "gated_delta_net", - "linear_attention_freq": [1], - "transformer_impl": "transformer_engine", - } - config_kwargs.update(overrides) - return TransformerConfig(**config_kwargs) - - -def test_gdn_pre_gated_delta_rule_fusion_defaults_to_disabled(): - config = _make_gdn_config() - assert not config.gdn_pre_gated_delta_rule_fusion - - -def test_gdn_pre_gated_delta_rule_fusion_accepts_gdn_variant(): - config = _make_gdn_config(gdn_pre_gated_delta_rule_fusion=True) - assert config.gdn_pre_gated_delta_rule_fusion - - -def test_gdn_pre_gated_delta_rule_fusion_requires_gdn_variant(): - with pytest.raises(ValueError, match="experimental_attention_variant='gated_delta_net'"): - _make_gdn_config( - experimental_attention_variant=None, - linear_attention_freq=None, - gdn_pre_gated_delta_rule_fusion=True, - ) - - -def test_gdn_conv_pad_alignment_rejects_chunkwise_cp(): - with pytest.raises(AssertionError, match="gdn_conv_pad_alignment is incompatible"): - _make_gdn_config( - context_parallel_size=2, linear_cp_mode="chunkwise", gdn_conv_pad_alignment=4096 - ) - - -def test_gdn_chunkwise_cp_head_divisibility_ignores_cp_size(): - config = _make_gdn_config( - tensor_model_parallel_size=2, - context_parallel_size=4, - linear_cp_mode="chunkwise", - linear_num_key_heads=4, - linear_num_value_heads=8, - ) - assert config.linear_cp_mode == "chunkwise" - - def test_fused_pre_gdr_split_batched_recv_send_works(): from megatron.core.fusions.fused_pre_gated_delta_rule import _split_batched_recv_send_works @@ -160,486 +59,6 @@ def test_fused_pre_gdr_split_batched_recv_send_works(): _split_batched_recv_send_works([], ["recv"]) -def test_gdn_headwise_cp_head_divisibility_includes_cp_size(): - with pytest.raises(AssertionError, match="linear_head_parallel_size"): - _make_gdn_config( - tensor_model_parallel_size=2, - context_parallel_size=4, - linear_cp_mode="headwise", - linear_num_key_heads=4, - linear_num_value_heads=8, - ) - - -@pytest.mark.parametrize( - ("tp_size", "sp", "cp_size", "linear_cp_mode"), - [ - # cp_size=1: the CP path is inactive, so linear_cp_mode choice is irrelevant. - # Cover the "chunkwise" default and skip the "headwise" variants for brevity. - (1, False, 1, None), - (2, False, 1, None), - (2, True, 1, None), - # cp_size=2: exercise both CP paths. - (1, False, 2, "headwise"), - (2, False, 2, "headwise"), - (2, True, 2, "headwise"), - (1, False, 2, "chunkwise"), - (2, False, 2, "chunkwise"), - (2, True, 2, "chunkwise"), - ], -) -@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") -@pytest.mark.internal -class TestGatedDeltaNet: - - @pytest.fixture(scope='function', autouse=True) - def setup_method(self, tp_size, sp, cp_size, linear_cp_mode): - # Initialize parallel and random seed - Utils.initialize_model_parallel( - tensor_model_parallel_size=tp_size, - pipeline_model_parallel_size=1, - context_parallel_size=cp_size, - ) - model_parallel_cuda_manual_seed(123) - self.tp_size = tp_size - self.cp_size = cp_size - self.sp_size = tp_size if sp else 1 - self.linear_cp_mode = linear_cp_mode - if self.linear_cp_mode == "headwise": - self.cp_size_chunkwise = 1 - self.cp_size_headwise = self.cp_size - elif self.linear_cp_mode == "chunkwise": - self.cp_size_chunkwise = self.cp_size - self.cp_size_headwise = 1 - elif self.cp_size == 1: - self.cp_size_chunkwise = 1 - self.cp_size_headwise = 1 - else: - raise ValueError(f"Invalid linear CP mode: {self.linear_cp_mode}") - - # Get TP and CP process groups from device mesh - tp_group = parallel_state.get_tensor_model_parallel_group() - cp_group = parallel_state.get_context_parallel_group() - pg_collection = ProcessGroupCollection(tp=tp_group, cp=cp_group) - - # Initialize model, with the same config as Qwen Next except `num_layers` - self.transformer_config = TransformerConfig( - hidden_size=2048, - linear_conv_kernel_dim=4, - linear_key_head_dim=128, - linear_value_head_dim=128, - linear_num_key_heads=16, - linear_num_value_heads=32, - num_layers=1, - normalization="RMSNorm", - use_cpu_initialization=True, - layernorm_zero_centered_gamma=True, - num_attention_heads=16, - num_query_groups=2, - activation_func=F.silu, - bf16=True, - tensor_model_parallel_size=tp_size, - sequence_parallel=sp, - context_parallel_size=cp_size, - experimental_attention_variant="gated_delta_net", - linear_attention_freq=[1], - linear_cp_mode=self.linear_cp_mode, - transformer_impl="transformer_engine", - ) - gdn_submodules = get_experimental_attention_variant_module_spec( - config=self.transformer_config - ).submodules - - self.gdn = GatedDeltaNet( - self.transformer_config, - submodules=gdn_submodules, - layer_number=1, - bias=False, - conv_bias=False, - conv_init=1.0, - use_qk_l2norm=True, - A_init_range=(1, 16), - pg_collection=pg_collection, - ) - self.gdn = self.gdn.cuda().bfloat16() - - def teardown_method(self): - Utils.destroy_model_parallel() - - def test_gpu_forward(self): - gdn = self.gdn - - micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 - seq_length = 64 - hidden_states = torch.ones( - (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), - device=torch.cuda.current_device(), - dtype=torch.bfloat16, - ) - attention_mask = None - - output, bias = gdn(hidden_states, attention_mask) - - assert output.dim() == 3, f"Output too many dimensions ({output.shape=})" - assert output.shape[0] == seq_length // self.sp_size // self.cp_size, ( - f"Output shape {output.shape[0]=} mismatch with " - f" {seq_length=} // {self.sp_size=} // {self.cp_size=}." - ) - assert ( - output.shape[1] == micro_batch_size - ), f"Output shape {output.shape[1]=} mismatch with {micro_batch_size=}" - assert ( - output.shape[2] == gdn.config.hidden_size - ), f"Output shape {output.shape[2]=} mismatch with {gdn.config.hidden_size=}" - assert ( - output.dtype == hidden_states.dtype - ), f"Output dtype {output.dtype=} mismatch with {hidden_states.dtype=}" - - @pytest.mark.flaky_in_dev # Issue #5473 - def test_selective_recompute_gdn(self): - """Whole-module 'gdn' recompute must match the non-recompute forward and gradients. - - The same module/input is run twice (recompute off, then on); the forward output and - all parameter / input gradients must agree within a tight tolerance (rtol/atol=1e-4). - The recompute path is run-to-run deterministic on these kernels (empirically bitwise), - so a tolerance well below the bf16 floor is expected to hold. - """ - gdn = self.gdn - gdn.train() - - micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 - seq_length = 64 - torch.manual_seed(1234) - base_input = torch.randn( - (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), - device=torch.cuda.current_device(), - dtype=torch.bfloat16, - ) - - def run(recompute): - gdn.recompute_gdn = recompute - gdn.zero_grad(set_to_none=True) - hidden_states = base_input.clone().detach().requires_grad_(True) - output, _ = gdn(hidden_states, None) - output.float().square().mean().backward() - param_grads = { - name: param.grad.detach().clone() - for name, param in gdn.named_parameters() - if param.grad is not None - } - return output.detach().clone(), hidden_states.grad.detach().clone(), param_grads - - try: - out_ref, dinput_ref, pgrad_ref = run(recompute=False) - out_rc, dinput_rc, pgrad_rc = run(recompute=True) - finally: - gdn.recompute_gdn = False - - torch.testing.assert_close(out_rc, out_ref, rtol=1e-4, atol=1e-4) - torch.testing.assert_close(dinput_rc, dinput_ref, rtol=1e-4, atol=1e-4) - assert pgrad_ref.keys() == pgrad_rc.keys(), "recompute changed the set of grad params" - assert len(pgrad_ref) > 0, "expected at least one parameter gradient" - for name in pgrad_ref: - torch.testing.assert_close( - pgrad_rc[name], - pgrad_ref[name], - rtol=1e-4, - atol=1e-4, - msg=lambda m, n=name: f"gradient mismatch for parameter '{n}': {m}", - ) - - def test_gpu_forward_rejects_sbhd_chunkwise_cp_batch_gt_one(self): - if not (self.linear_cp_mode == "chunkwise" and self.cp_size > 1): - pytest.skip("Only chunkwise CP with CP>1 uses the FLA CP batch guard.") - - gdn = self.gdn - - micro_batch_size = 2 - seq_length = 64 - hidden_states = torch.ones( - (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), - device=torch.cuda.current_device(), - dtype=torch.bfloat16, - ) - - with pytest.raises(ValueError, match="requires micro_batch_size == 1"): - gdn(hidden_states, None) - - def test_gpu_forward_rejects_sbhd_conv_padding(self): - gdn = self.gdn - gdn.config.gdn_conv_pad_alignment = 4096 - - micro_batch_size = 1 if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 else 2 - seq_length = 64 - hidden_states = torch.ones( - (seq_length // self.sp_size // self.cp_size, micro_batch_size, gdn.config.hidden_size), - device=torch.cuda.current_device(), - dtype=torch.bfloat16, - ) - - expected_error = ( - "incompatible with GDN chunkwise CP" - if self.linear_cp_mode == "chunkwise" and self.cp_size > 1 - else "only supported with packed sequence" - ) - with pytest.raises(ValueError, match=expected_error): - gdn(hidden_states, None) - - def test_jit_compiled_helpers(self): - import torch._dynamo - - gdn = self.gdn - batch = 2 - seq_len = 16 - - num_v_heads_local = gdn.num_value_heads // gdn.tp_size // self.cp_size_headwise - - qkv_last_dim = (2 * gdn.qk_dim_local_tp + gdn.v_dim_local_tp) // self.cp_size_headwise - qkv = torch.randn( - batch, seq_len, qkv_last_dim, device=torch.cuda.current_device(), dtype=torch.bfloat16 - ) - gate = torch.randn( - batch, - seq_len, - num_v_heads_local, - gdn.value_head_dim, - device=torch.cuda.current_device(), - dtype=torch.bfloat16, - ) - beta = torch.randn( - batch, - seq_len, - num_v_heads_local, - device=torch.cuda.current_device(), - dtype=torch.bfloat16, - ) - alpha = torch.randn( - batch, - seq_len, - num_v_heads_local, - device=torch.cuda.current_device(), - dtype=torch.bfloat16, - ) - - # Disable dynamo so coverage.py can trace through the method bodies, - # which are normally wrapped by @jit_fuser (torch.compile). - with torch._dynamo.config.patch(disable=True): - query, key, value, gate_out, beta_out, alpha_out = ( - gdn._prepare_qkv_for_gated_delta_rule( - qkv, gate, beta, alpha, batch, seq_len, cp_size_headwise=self.cp_size_headwise - ) - ) - - assert query.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim) - assert key.shape == (batch, seq_len, num_v_heads_local, gdn.key_head_dim) - assert value.shape == (batch, seq_len, num_v_heads_local, gdn.value_head_dim) - assert query.is_contiguous() - assert key.is_contiguous() - assert value.is_contiguous() - - A_log_mock = torch.randn( - num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16 - ) - dt_bias_mock = torch.randn( - num_v_heads_local, device=torch.cuda.current_device(), dtype=torch.bfloat16 - ) - - with torch._dynamo.config.patch(disable=True): - g, beta_sig = gdn._compute_g_and_beta(A_log_mock, dt_bias_mock, alpha, beta) - - assert g.dtype == torch.float32 - assert g.shape == alpha.shape - assert beta_sig.shape == beta.shape - - def test_fused_pre_gated_delta_rule_headwise_cp_uses_cp_local_parameters(self): - if not HAVE_FUSED_PRE_GDR: - pytest.skip("causal-conv1d fused backward is not installed.") - if not (self.linear_cp_mode == "headwise" and self.cp_size > 1): - pytest.skip("Only headwise CP with CP>1 needs CP-local fused pre-GDR params.") - - gdn = self.gdn - batch = 2 - seq_len = 16 - qk_channels = gdn.qk_dim_local_tp // self.cp_size_headwise - v_channels = gdn.v_dim_local_tp // self.cp_size_headwise - num_key_heads = qk_channels // gdn.key_head_dim - num_value_heads = v_channels // gdn.value_head_dim - qkvzba_dim = 2 * qk_channels + 2 * v_channels + 2 * num_value_heads - qkvzba = torch.randn( - seq_len, batch, qkvzba_dim, device=torch.cuda.current_device(), dtype=torch.bfloat16 - ) - captured = {} - - def fake_fused_streamed_pre_gated_delta_rule( - qkvzba_arg, - conv1d_weight, - conv1d_bias, - A_log, - dt_bias, - *, - num_key_heads, - num_value_heads, - **kwargs, - ): - captured.update( - { - "qkvzba": qkvzba_arg, - "conv1d_weight": conv1d_weight, - "conv1d_bias": conv1d_bias, - "A_log": A_log, - "dt_bias": dt_bias, - "num_key_heads": num_key_heads, - "num_value_heads": num_value_heads, - "cp_group": kwargs["cp_group"], - } - ) - return tuple(torch.empty(0, device=qkvzba_arg.device) for _ in range(6)) - - with mock.patch( - "megatron.core.fusions.fused_pre_gated_delta_rule." - "fused_streamed_pre_gated_delta_rule", - side_effect=fake_fused_streamed_pre_gated_delta_rule, - ): - gdn._fused_streamed_pre_gated_delta_rule( - qkvzba, - cp_size_headwise=self.cp_size_headwise, - cp_group_headwise=gdn.pg_collection.cp, - ) - - assert captured["qkvzba"] is qkvzba - assert captured["conv1d_weight"].shape == ( - 2 * qk_channels + v_channels, - 1, - gdn.conv_kernel_dim, - ) - assert captured["conv1d_bias"] is None - assert captured["A_log"].shape == (num_value_heads,) - assert captured["dt_bias"].shape == (num_value_heads,) - assert captured["num_key_heads"] == num_key_heads - assert captured["num_value_heads"] == num_value_heads - assert captured["cp_group"] is None - - def test_gpu_forward_thd_correctness(self): - if self.sp_size > 1: - pytest.skip("Sequence parallel is not supported for this test case.") - if self.cp_size > 1 and self.linear_cp_mode == "chunkwise": - pytest.skip("Chunkwise CP is not supported for this test case.") - - atol, rtol = 3e-4, 3e-4 - - # Input shape - sequence_length = 32 - micro_batch_size = 4 - cu_seqlens = [0, 32, 64, 96, 128] - # sbhd input shape: [sequence length, batch size, hidden size] - sub_sequence_length = sequence_length // self.cp_size - hidden_states_sbhd = torch.rand( - (sub_sequence_length, micro_batch_size, self.gdn.config.hidden_size) - ) - attention_mask_sbhd = None - hidden_states_sbhd = hidden_states_sbhd.cuda().bfloat16() - # thd input shape: [sequence length * batch size, 1, hidden size] - hidden_states_thd = hidden_states_sbhd.transpose(0, 1).contiguous() - hidden_states_thd = hidden_states_thd.view(-1, 1, self.gdn.config.hidden_size) - attention_mask_thd = None - packed_seq_params = make_test_packed_seq_params(cu_seqlens=cu_seqlens) - - # THD format - output_thd, _ = self.gdn( - hidden_states_thd, attention_mask_thd, packed_seq_params=packed_seq_params - ) - # SBHD format - output_sbhd, _ = self.gdn(hidden_states_sbhd, attention_mask_sbhd) - output_sbhd_T = output_sbhd.transpose(0, 1).contiguous().view(*output_thd.shape) - - rank = torch.distributed.get_rank() - assert output_thd.shape[0] == sub_sequence_length * micro_batch_size - assert output_thd.shape[1] == 1 - assert output_thd.shape[2] == self.gdn.config.hidden_size - torch.testing.assert_close( - output_sbhd_T, - output_thd, - atol=atol, - rtol=rtol, - msg=lambda msg: f"Output mismatch ({rank=}): {msg}", - ) - - def test_gpu_forward_thd_padding_correctness(self): - if self.sp_size > 1: - pytest.skip("Sequence parallel is not supported for this test case.") - if self.cp_size > 1 and self.linear_cp_mode == "chunkwise": - pytest.skip("Chunkwise CP is not supported for this test case.") - - atol, rtol = 3e-4, 3e-4 - sequence_length = 32 - micro_batch_size = 4 - - # sbhd input shape: [sequence length, batch size, hidden size] - sub_sequence_length = sequence_length // self.cp_size - hidden_states_sbhd = torch.rand( - (sub_sequence_length, micro_batch_size, self.gdn.config.hidden_size), - device=torch.cuda.current_device(), - dtype=torch.bfloat16, - ) - output_sbhd, _ = self.gdn(hidden_states_sbhd, None) - - # thd input shape: [sequence length * batch size, 1, hidden size] - hidden_states_thd = hidden_states_sbhd.transpose(0, 1).contiguous() - hidden_states_thd = hidden_states_thd.view(-1, 1, self.gdn.config.hidden_size) - output_bshd = output_sbhd.transpose(0, 1).contiguous() - - rank = torch.distributed.get_rank() - - # A) padded branch: prefer *_padded when available. - padded_params = make_test_packed_seq_params_with_padding( - cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 128] - ) - output_thd_padded, _ = self.gdn(hidden_states_thd, None, packed_seq_params=padded_params) - output_thd2bshd = output_thd_padded.view(*output_bshd.shape) - torch.testing.assert_close( - output_bshd[:, :30, :], - output_thd2bshd[:, :30, :], - atol=atol, - rtol=rtol, - msg=lambda msg: f"THD padded output mismatch ({rank=}): {msg}", - ) - - # B) no-padded branch: use actual cu_seqlens when it matches total_sequence_length. - no_padding_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 128]) - output_thd_no_padding, _ = self.gdn( - hidden_states_thd, None, packed_seq_params=no_padding_params - ) - assert output_thd_no_padding.shape == output_thd_padded.shape - - # C) explicit causal-conv padding is only applied to packed inputs and - # should not affect the original unpadded token outputs. - self.gdn.config.gdn_conv_pad_alignment = 48 - output_thd_conv_pad, _ = self.gdn( - hidden_states_thd, None, packed_seq_params=no_padding_params - ) - self.gdn.config.gdn_conv_pad_alignment = None - assert output_thd_conv_pad.shape == output_thd_no_padding.shape - torch.testing.assert_close( - output_thd_conv_pad, - output_thd_no_padding, - atol=atol, - rtol=rtol, - msg=lambda msg: f"THD conv-padded output mismatch ({rank=}): {msg}", - ) - - # D) padded mismatch branch: if *_padded[-1] mismatches total_sequence_length, should raise. - padded_mismatch_params = make_test_packed_seq_params_with_padding( - cu_seqlens=[0, 30, 60, 90, 120], cu_seqlens_padded=[0, 32, 64, 96, 126] - ) - with pytest.raises(ValueError, match="does not match"): - self.gdn(hidden_states_thd, None, packed_seq_params=padded_mismatch_params) - - # E) actual mismatch branch without *_padded: should raise. - actual_mismatch_params = make_test_packed_seq_params(cu_seqlens=[0, 32, 64, 96, 129]) - with pytest.raises(ValueError, match="does not match"): - self.gdn(hidden_states_thd, None, packed_seq_params=actual_mismatch_params) - - @pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") @pytest.mark.skipif(not HAVE_FUSED_PRE_GDR, reason="causal-conv1d fused backward is not installed.") @pytest.mark.internal @@ -1620,288 +1039,3 @@ def test_fused_and_unfused_packed_partial_boundary_chunkwise_cp_match(self): f"({rank=}): {msg}" ), ) - - -@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") -@pytest.mark.internal -class TestGDNCuSeqlensResolve: - - @pytest.fixture - def mock_gdn(self): - class MockGDN: - _resolve_cu_seqlens = GatedDeltaNet._resolve_cu_seqlens - - return MockGDN() - - def test_padded_preferred_when_available(self, mock_gdn): - actual = torch.tensor([0, 500, 1000], dtype=torch.int32) - padded = torch.tensor([0, 504, 1008], dtype=torch.int32) - result = mock_gdn._resolve_cu_seqlens(padded, actual, 1008, "cu_seqlens_q", cp_size=2) - assert torch.equal(result, padded) - - def test_actual_used_when_no_padding(self, mock_gdn): - actual = torch.tensor([0, 504, 1008], dtype=torch.int32) - result = mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) - assert torch.equal(result, actual) - - def test_raises_when_padding_mismatch(self, mock_gdn): - actual = torch.tensor([0, 500, 1000], dtype=torch.int32) - with pytest.raises(ValueError, match="does not match"): - mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) - - def test_raises_when_padded_mismatches_total(self, mock_gdn): - actual = torch.tensor([0, 500, 1000], dtype=torch.int32) - padded = torch.tensor([0, 504, 1004], dtype=torch.int32) - with pytest.raises(ValueError, match="does not match"): - mock_gdn._resolve_cu_seqlens(padded, actual, 1008, "cu_seqlens_q", cp_size=2) - - def test_raises_when_not_divisible_by_cp_size(self, mock_gdn): - actual = torch.tensor([0, 505, 1008], dtype=torch.int32) - with pytest.raises(ValueError, match="must be divisible by cp_size"): - mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=2) - - def test_cp1_still_validates_total(self, mock_gdn): - mock_gdn.cp_size = 1 - actual = torch.tensor([0, 500, 1000], dtype=torch.int32) - with pytest.raises(ValueError, match="does not match"): - mock_gdn._resolve_cu_seqlens(None, actual, 1008, "cu_seqlens_q", cp_size=1) - - -@pytest.mark.parametrize("sequence_packing", [False, True]) -@pytest.mark.parametrize( - ("tp", "sp", "cp", "linear_cp_mode"), - [ - (4, False, 1, None), # TP w/o SP - (4, True, 1, None), # TP w/ SP - (1, False, 2, "headwise"), # Headwise CP - (2, False, 2, "headwise"), # TP w/o SP + Headwise CP - (2, True, 2, "headwise"), # TP w/ SP + Headwise CP - (1, False, 2, "chunkwise"), # Chunkwise CP - (2, False, 2, "chunkwise"), # TP w/o SP + chunkwise CP - (2, True, 2, "chunkwise"), # TP w/ SP + chunkwise CP - ], -) -@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") -def test_parallel_gated_delta_net_correctness( - tmp_path_dist_ckpt, sequence_packing, tp, sp, cp, linear_cp_mode -): - transformer_config = TransformerConfig( - hidden_size=128, - linear_conv_kernel_dim=2, - linear_key_head_dim=32, - linear_value_head_dim=32, - linear_num_key_heads=4, - linear_num_value_heads=8, - num_layers=1, - normalization="RMSNorm", - use_cpu_initialization=True, - layernorm_zero_centered_gamma=True, - num_attention_heads=8, - activation_func=F.silu, - bf16=True, - experimental_attention_variant="gated_delta_net", - linear_attention_freq=[1], - linear_cp_mode=linear_cp_mode, - transformer_impl="transformer_engine", - ) - - transformer_layer_spec = get_transformer_block_with_experimental_attention_variant_spec( - config=transformer_config, vp_stage=None, pp_rank=0 - ) - - cosine_similarity_threshold = None - if cp > 1: - atol, rtol = 2e-3, 1e-2 - cosine_similarity_threshold = 0.9999 - else: - atol, rtol = 2e-4, 2e-3 - cosine_similarity_threshold = 0.99999 - - is_chunkwise_cp = linear_cp_mode == "chunkwise" and cp > 1 - micro_batch_size = 1 if is_chunkwise_cp and not sequence_packing else 4 - - _test_parallel_attention_correctness( - transformer_config=transformer_config, - transformer_layer_spec=transformer_layer_spec, - tmp_path_dist_ckpt=tmp_path_dist_ckpt, - atol=atol, - rtol=rtol, - cosine_similarity_threshold=cosine_similarity_threshold, - tp=tp, - sp=sp, - cp=cp, - seed=123, - sequence_length=256, - micro_batch_size=micro_batch_size, - sequence_packing=sequence_packing, - ) - - -@pytest.mark.parametrize("cp_size", [2, 4], scope="class") -@pytest.mark.internal -class TestFusedThdAllToAll: - """Verify fused 1 AllToAll + permute matches the per-sequence, per-channel loop in GDN.""" - - @pytest.fixture(scope='class', autouse=True) - def setup_method(self, request, cp_size): - Utils.initialize_model_parallel( - tensor_model_parallel_size=1, - pipeline_model_parallel_size=1, - context_parallel_size=cp_size, - ) - model_parallel_cuda_manual_seed(123) - # Attach on the class so every test method can read self.cp_*. - request.cls.cp_size = cp_size - request.cls.cp_group = parallel_state.get_context_parallel_group() - yield - Utils.destroy_model_parallel() - - @staticmethod - def _per_seq_a2a_cp2hp(local_t, cu_seqlens, cp_group, split_sections=None): - cp_size = cp_group.size() - unpacked = _unpack_sequence(local_t, cu_seqlens // cp_size, dim=0) - outputs = [] - for x in unpacked: - outputs.append( - tensor_a2a_cp2hp( - x, - seq_dim=0, - head_dim=-1, - cp_group=cp_group, - split_sections=split_sections, - undo_attention_load_balancing=True, - ) - ) - return torch.cat(outputs, dim=0) - - @staticmethod - def _per_seq_a2a_hp2cp(global_t, cu_seqlens, cp_group, split_sections=None): - unpacked = _unpack_sequence(global_t, cu_seqlens, dim=0) - outputs = [] - for x in unpacked: - outputs.append( - tensor_a2a_hp2cp( - x, - seq_dim=0, - head_dim=-1, - cp_group=cp_group, - split_sections=split_sections, - redo_attention_load_balancing=True, - ) - ) - return torch.cat(outputs, dim=0) - - # ---- Optimized: single a2a + production permutation helper ---- - - @staticmethod - def _batched_a2a_cp2hp(local_t, cu_seqlens, cp_group, split_sections=None): - cp_size = cp_group.size() - t_global = int(cu_seqlens[-1].item()) - if split_sections is not None and cp_size > 1: - head_perm = _build_head_perm_for_split_sections(split_sections, cp_size, local_t.device) - local_t = local_t.index_select(-1, head_perm) - naive = tensor_a2a_cp2hp( - local_t, - seq_dim=0, - head_dim=-1, - cp_group=cp_group, - split_sections=None, # always single fused a2a - undo_attention_load_balancing=False, - ) - idx, _ = _build_thd_cp_a2a_perm(cu_seqlens, cp_size, t_global) - return naive.index_select(0, idx) - - @staticmethod - def _batched_a2a_hp2cp(global_t, cu_seqlens, cp_group, split_sections=None): - cp_size = cp_group.size() - t_global = int(cu_seqlens[-1].item()) - _, inv = _build_thd_cp_a2a_perm(cu_seqlens, cp_size, t_global) - permuted = global_t.index_select(0, inv) - return tensor_a2a_hp2cp( - permuted, - seq_dim=0, - head_dim=-1, - cp_group=cp_group, - split_sections=split_sections, - redo_attention_load_balancing=False, - ) - - @pytest.mark.parametrize( - "cu_seqlens", - [ - (0, 32, 64), # 2 equal sequences - (0, 32, 64, 96, 128), # 4 equal sequences (matches existing THD test) - (0, 16, 48, 80), # 3 unequal sequences - ], - ) - @pytest.mark.parametrize("split_sections", [(8, 8, 4, 16, 32, 4)]) - def test_cp2hp_batched_matches_per_seq(self, cu_seqlens, split_sections): - cu = torch.tensor(cu_seqlens, dtype=torch.long, device=torch.cuda.current_device()) - if (torch.diff(cu) % self.cp_size != 0).any(): - pytest.skip(f"cu_seqlens {cu_seqlens} not divisible by cp_size {self.cp_size}") - - T_global = cu_seqlens[-1] - T_local = T_global // self.cp_size - hidden = sum(split_sections) - torch.manual_seed(42) - local_t = ( - torch.rand(T_local, 1, hidden, device=torch.cuda.current_device()) - .bfloat16() - .contiguous() - ) - - out_ref = self._per_seq_a2a_cp2hp(local_t, cu, self.cp_group, split_sections=split_sections) - out_fused = self._batched_a2a_cp2hp( - local_t, cu, self.cp_group, split_sections=split_sections - ) - - rank = torch.distributed.get_rank() - assert torch.equal(out_fused, out_ref), ( - f"Batched CP->HP mismatch on rank={rank} " f"(split_sections={split_sections})" - ) - - @pytest.mark.parametrize("cu_seqlens", [(0, 32, 64), (0, 32, 64, 96, 128), (0, 16, 48, 80)]) - def test_hp2cp_batched_matches_per_seq(self, cu_seqlens): - cu = torch.tensor(cu_seqlens, dtype=torch.long, device=torch.cuda.current_device()) - if ((cu[1:] - cu[:-1]) % self.cp_size != 0).any(): - pytest.skip(f"cu_seqlens {cu_seqlens} not divisible by cp_size {self.cp_size}") - - T_global = cu_seqlens[-1] - hidden = 32 - # Hidden must be divisible by cp_size for the HP-sharded input layout. - assert hidden % self.cp_size == 0 - h_local = hidden // self.cp_size - torch.manual_seed(42) - global_t = ( - torch.rand(T_global, 1, h_local, device=torch.cuda.current_device()) - .bfloat16() - .contiguous() - ) - - out_ref = self._per_seq_a2a_hp2cp(global_t, cu, self.cp_group) - out_fused = self._batched_a2a_hp2cp(global_t, cu, self.cp_group) - - rank = torch.distributed.get_rank() - assert torch.equal(out_fused, out_ref), f"Batched HP->CP mismatch on rank={rank}" - - @pytest.mark.parametrize("cu_seqlens", [(0, 32, 64, 96, 128)]) - def test_cp2hp_hp2cp_round_trip(self, cu_seqlens): - """cp2hp followed by hp2cp on the batched path should be the identity.""" - cu = torch.tensor(cu_seqlens, dtype=torch.long, device=torch.cuda.current_device()) - if ((cu[1:] - cu[:-1]) % self.cp_size != 0).any(): - pytest.skip(f"cu_seqlens {cu_seqlens} not divisible by cp_size {self.cp_size}") - - T_global = cu_seqlens[-1] - T_local = T_global // self.cp_size - hidden = 32 - torch.manual_seed(7) - local_t = ( - torch.rand(T_local, 1, hidden, device=torch.cuda.current_device()) - .bfloat16() - .contiguous() - ) - - mid = self._batched_a2a_cp2hp(local_t, cu, self.cp_group) - back = self._batched_a2a_hp2cp(mid, cu, self.cp_group) - - assert torch.equal(back, local_t), "Batched cp2hp -> hp2cp not identity" diff --git a/tests/unit_tests/ssm/gated_delta_net/test_gdn_parallel.py b/tests/unit_tests/ssm/gated_delta_net/test_gdn_parallel.py new file mode 100644 index 00000000000..781326e1fdd --- /dev/null +++ b/tests/unit_tests/ssm/gated_delta_net/test_gdn_parallel.py @@ -0,0 +1,285 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + +import os + +import pytest +import torch +import torch.nn.functional as F + +from megatron.core import parallel_state +from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_transformer_block_with_experimental_attention_variant_spec, +) +from megatron.core.ssm.gated_delta_net.common import ( + _build_head_perm_for_split_sections, + _build_thd_cp_a2a_perm, + tensor_a2a_cp2hp, + tensor_a2a_hp2cp, +) +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer import TransformerConfig +from tests.unit_tests.test_utilities import Utils +from tests.unit_tests.transformer.test_attention import _test_parallel_attention_correctness + +try: + import fla + + HAVE_FLA = True +except ImportError: + HAVE_FLA = False + +# https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/env.html#nccl-multi-rank-gpu-enable +# NVLS doesn't support one single GPU to be shared by multiple ranks, so disable this in test +os.environ.update({"NCCL_NVLS_ENABLE": "0"}) + + +def _unpack_sequence(x: torch.Tensor, cu_seqlens: torch.Tensor, dim=1) -> list[torch.Tensor]: + unpacked_x = [] + cu_seqlens_list = cu_seqlens.tolist() + num_seqs = len(cu_seqlens_list) - 1 + for i in range(num_seqs): + idx_start = cu_seqlens_list[i] + idx_end = cu_seqlens_list[i + 1] + chunked_index = [slice(None)] * dim + [slice(idx_start, idx_end)] + unpacked_x.append(x[tuple(chunked_index)]) + return unpacked_x + + +@pytest.mark.parametrize("sequence_packing", [False, True]) +@pytest.mark.parametrize( + ("tp", "sp", "cp", "linear_cp_mode"), + [ + (4, False, 1, None), # TP w/o SP + (4, True, 1, None), # TP w/ SP + (1, False, 2, "headwise"), # Headwise CP + (2, False, 2, "headwise"), # TP w/o SP + Headwise CP + (2, True, 2, "headwise"), # TP w/ SP + Headwise CP + (1, False, 2, "chunkwise"), # Chunkwise CP + (2, False, 2, "chunkwise"), # TP w/o SP + chunkwise CP + (2, True, 2, "chunkwise"), # TP w/ SP + chunkwise CP + ], +) +@pytest.mark.skipif(not HAVE_FLA, reason="FLA is not installed.") +def test_parallel_gated_delta_net_correctness( + tmp_path_dist_ckpt, sequence_packing, tp, sp, cp, linear_cp_mode +): + transformer_config = TransformerConfig( + hidden_size=128, + linear_conv_kernel_dim=2, + linear_key_head_dim=32, + linear_value_head_dim=32, + linear_num_key_heads=4, + linear_num_value_heads=8, + num_layers=1, + normalization="RMSNorm", + use_cpu_initialization=True, + layernorm_zero_centered_gamma=True, + num_attention_heads=8, + activation_func=F.silu, + bf16=True, + experimental_attention_variant="gated_delta_net", + linear_attention_freq=[1], + linear_cp_mode=linear_cp_mode, + transformer_impl="transformer_engine", + ) + + transformer_layer_spec = get_transformer_block_with_experimental_attention_variant_spec( + config=transformer_config, vp_stage=None, pp_rank=0 + ) + + cosine_similarity_threshold = None + if cp > 1: + atol, rtol = 2e-3, 1e-2 + cosine_similarity_threshold = 0.9999 + else: + atol, rtol = 2e-4, 2e-3 + cosine_similarity_threshold = 0.99999 + + is_chunkwise_cp = linear_cp_mode == "chunkwise" and cp > 1 + micro_batch_size = 1 if is_chunkwise_cp and not sequence_packing else 4 + + _test_parallel_attention_correctness( + transformer_config=transformer_config, + transformer_layer_spec=transformer_layer_spec, + tmp_path_dist_ckpt=tmp_path_dist_ckpt, + atol=atol, + rtol=rtol, + cosine_similarity_threshold=cosine_similarity_threshold, + tp=tp, + sp=sp, + cp=cp, + seed=123, + sequence_length=256, + micro_batch_size=micro_batch_size, + sequence_packing=sequence_packing, + ) + + +@pytest.mark.parametrize("cp_size", [2, 4], scope="class") +@pytest.mark.internal +class TestFusedThdAllToAll: + """Verify fused 1 AllToAll + permute matches the per-sequence, per-channel loop in GDN.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request, cp_size): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=cp_size, + ) + model_parallel_cuda_manual_seed(123) + # Attach on the class so every test method can read self.cp_*. + request.cls.cp_size = cp_size + request.cls.cp_group = parallel_state.get_context_parallel_group() + yield + Utils.destroy_model_parallel() + + @staticmethod + def _per_seq_a2a_cp2hp(local_t, cu_seqlens, cp_group, split_sections=None): + cp_size = cp_group.size() + unpacked = _unpack_sequence(local_t, cu_seqlens // cp_size, dim=0) + outputs = [] + for x in unpacked: + outputs.append( + tensor_a2a_cp2hp( + x, + seq_dim=0, + head_dim=-1, + cp_group=cp_group, + split_sections=split_sections, + undo_attention_load_balancing=True, + ) + ) + return torch.cat(outputs, dim=0) + + @staticmethod + def _per_seq_a2a_hp2cp(global_t, cu_seqlens, cp_group, split_sections=None): + unpacked = _unpack_sequence(global_t, cu_seqlens, dim=0) + outputs = [] + for x in unpacked: + outputs.append( + tensor_a2a_hp2cp( + x, + seq_dim=0, + head_dim=-1, + cp_group=cp_group, + split_sections=split_sections, + redo_attention_load_balancing=True, + ) + ) + return torch.cat(outputs, dim=0) + + # ---- Optimized: single a2a + production permutation helper ---- + + @staticmethod + def _batched_a2a_cp2hp(local_t, cu_seqlens, cp_group, split_sections=None): + cp_size = cp_group.size() + t_global = int(cu_seqlens[-1].item()) + if split_sections is not None and cp_size > 1: + head_perm = _build_head_perm_for_split_sections(split_sections, cp_size, local_t.device) + local_t = local_t.index_select(-1, head_perm) + naive = tensor_a2a_cp2hp( + local_t, + seq_dim=0, + head_dim=-1, + cp_group=cp_group, + split_sections=None, # always single fused a2a + undo_attention_load_balancing=False, + ) + idx, _ = _build_thd_cp_a2a_perm(cu_seqlens, cp_size, t_global) + return naive.index_select(0, idx) + + @staticmethod + def _batched_a2a_hp2cp(global_t, cu_seqlens, cp_group, split_sections=None): + cp_size = cp_group.size() + t_global = int(cu_seqlens[-1].item()) + _, inv = _build_thd_cp_a2a_perm(cu_seqlens, cp_size, t_global) + permuted = global_t.index_select(0, inv) + return tensor_a2a_hp2cp( + permuted, + seq_dim=0, + head_dim=-1, + cp_group=cp_group, + split_sections=split_sections, + redo_attention_load_balancing=False, + ) + + @pytest.mark.parametrize( + "cu_seqlens", + [ + (0, 32, 64), # 2 equal sequences + (0, 32, 64, 96, 128), # 4 equal sequences (matches existing THD test) + (0, 16, 48, 80), # 3 unequal sequences + ], + ) + @pytest.mark.parametrize("split_sections", [(8, 8, 4, 16, 32, 4)]) + def test_cp2hp_batched_matches_per_seq(self, cu_seqlens, split_sections): + cu = torch.tensor(cu_seqlens, dtype=torch.long, device=torch.cuda.current_device()) + if (torch.diff(cu) % self.cp_size != 0).any(): + pytest.skip(f"cu_seqlens {cu_seqlens} not divisible by cp_size {self.cp_size}") + + T_global = cu_seqlens[-1] + T_local = T_global // self.cp_size + hidden = sum(split_sections) + torch.manual_seed(42) + local_t = ( + torch.rand(T_local, 1, hidden, device=torch.cuda.current_device()) + .bfloat16() + .contiguous() + ) + + out_ref = self._per_seq_a2a_cp2hp(local_t, cu, self.cp_group, split_sections=split_sections) + out_fused = self._batched_a2a_cp2hp( + local_t, cu, self.cp_group, split_sections=split_sections + ) + + rank = torch.distributed.get_rank() + assert torch.equal(out_fused, out_ref), ( + f"Batched CP->HP mismatch on rank={rank} " f"(split_sections={split_sections})" + ) + + @pytest.mark.parametrize("cu_seqlens", [(0, 32, 64), (0, 32, 64, 96, 128), (0, 16, 48, 80)]) + def test_hp2cp_batched_matches_per_seq(self, cu_seqlens): + cu = torch.tensor(cu_seqlens, dtype=torch.long, device=torch.cuda.current_device()) + if ((cu[1:] - cu[:-1]) % self.cp_size != 0).any(): + pytest.skip(f"cu_seqlens {cu_seqlens} not divisible by cp_size {self.cp_size}") + + T_global = cu_seqlens[-1] + hidden = 32 + # Hidden must be divisible by cp_size for the HP-sharded input layout. + assert hidden % self.cp_size == 0 + h_local = hidden // self.cp_size + torch.manual_seed(42) + global_t = ( + torch.rand(T_global, 1, h_local, device=torch.cuda.current_device()) + .bfloat16() + .contiguous() + ) + + out_ref = self._per_seq_a2a_hp2cp(global_t, cu, self.cp_group) + out_fused = self._batched_a2a_hp2cp(global_t, cu, self.cp_group) + + rank = torch.distributed.get_rank() + assert torch.equal(out_fused, out_ref), f"Batched HP->CP mismatch on rank={rank}" + + @pytest.mark.parametrize("cu_seqlens", [(0, 32, 64, 96, 128)]) + def test_cp2hp_hp2cp_round_trip(self, cu_seqlens): + """cp2hp followed by hp2cp on the batched path should be the identity.""" + cu = torch.tensor(cu_seqlens, dtype=torch.long, device=torch.cuda.current_device()) + if ((cu[1:] - cu[:-1]) % self.cp_size != 0).any(): + pytest.skip(f"cu_seqlens {cu_seqlens} not divisible by cp_size {self.cp_size}") + + T_global = cu_seqlens[-1] + T_local = T_global // self.cp_size + hidden = 32 + torch.manual_seed(7) + local_t = ( + torch.rand(T_local, 1, hidden, device=torch.cuda.current_device()) + .bfloat16() + .contiguous() + ) + + mid = self._batched_a2a_cp2hp(local_t, cu, self.cp_group) + back = self._batched_a2a_hp2cp(mid, cu, self.cp_group) + + assert torch.equal(back, local_t), "Batched cp2hp -> hp2cp not identity" diff --git a/tests/unit_tests/ssm/test_split_tensor_factory.py b/tests/unit_tests/ssm/test_split_tensor_factory.py index abb668e16a8..359d18642c8 100644 --- a/tests/unit_tests/ssm/test_split_tensor_factory.py +++ b/tests/unit_tests/ssm/test_split_tensor_factory.py @@ -1,4 +1,4 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. import logging from unittest import mock @@ -7,7 +7,7 @@ import torch from megatron.core.dist_checkpointing import ShardedTensor -from megatron.core.ssm.gated_delta_net import ( +from megatron.core.ssm.gated_delta_net.common import ( _split_tensor_factory as gated_delta_split_tensor_factory, ) from megatron.core.ssm.mamba_mixer import _split_tensor_factory as mamba_split_tensor_factory diff --git a/uv.lock b/uv.lock index 55dc1bd308b..ae4e074bd42 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'", @@ -78,8 +78,23 @@ overrides = [ { name = "torch", marker = "sys_platform == 'never'" }, { name = "torchvision", marker = "sys_platform == 'never'" }, { name = "triton", marker = "sys_platform == 'never'" }, + { name = "urllib3", specifier = ">=2.6.3" }, ] +[[manifest.dependency-metadata]] +name = "fast-hadamard-transform" +version = "1.0.4.post1" +requires-dist = ["torch", "packaging", "ninja"] + +[[manifest.dependency-metadata]] +name = "flash-mla" +version = "1.0.0+b7643bd" + +[[manifest.dependency-metadata]] +name = "transformer-engine" +version = "2.14.0+f031cf87" +requires-dist = ["einops", "importlib-metadata", "nvdlfw-inspect", "onnx", "onnxscript", "packaging", "pydantic", "torch"] + [[package]] name = "absl-py" version = "2.4.0" @@ -164,7 +179,7 @@ source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "aiohappyeyeballs" }, { name = "aiosignal" }, - { name = "async-timeout", marker = "python_full_version < '3.11'" }, + { name = "async-timeout", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, { name = "attrs" }, { name = "frozenlist" }, { name = "multidict" }, @@ -304,7 +319,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 = [ @@ -358,10 +373,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 = [ @@ -666,7 +681,7 @@ wheels = [ [[package]] name = "black" -version = "24.4.2" +version = "26.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "click" }, @@ -674,24 +689,38 @@ dependencies = [ { name = "packaging" }, { name = "pathspec" }, { name = "platformdirs" }, + { name = "pytokens" }, { name = "tomli", marker = "python_full_version < '3.11' 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/a2/47/c9997eb470a7f48f7aaddd3d9a828244a2e4199569e38128715c48059ac1/black-24.4.2.tar.gz", hash = "sha256:c872b53057f000085da66a19c55d68f6f8ddcac2642392ad3a355878406fbd4d", size = 642299, upload-time = "2024-04-26T00:32:15.305Z" } -wheels = [ - { url = "https://files.pythonhosted.org/packages/40/f6/3adc48c210527a7b651aaed43824a9b8bd04b3fb361a5227bad046e1c876/black-24.4.2-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:dd1b5a14e417189db4c7b64a6540f31730713d173f0b63e55fabd52d61d8fdce", size = 1631487, upload-time = "2024-04-26T00:40:28.969Z" }, - { url = "https://files.pythonhosted.org/packages/a2/25/70aa1bec12c841a03e333e312daa0cf2fee50ea6336ac4851c93c0e2b411/black-24.4.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:8e537d281831ad0e71007dcdcbe50a71470b978c453fa41ce77186bbe0ed6021", size = 1456317, upload-time = "2024-04-26T00:39:10.333Z" }, - { url = "https://files.pythonhosted.org/packages/e0/7d/7f8df0fdbbbefc4362d3eca6b69b7a8a4249a8a88dabc00a207d31fddcd7/black-24.4.2-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:eaea3008c281f1038edb473c1aa8ed8143a5535ff18f978a318f10302b254063", size = 1822765, upload-time = "2024-04-26T00:34:56.436Z" }, - { url = "https://files.pythonhosted.org/packages/5c/21/1ee97841c469c1551133cbe47448cdba9628c7d9431f74f114f02e3b233c/black-24.4.2-cp310-cp310-win_amd64.whl", hash = "sha256:7768a0dbf16a39aa5e9a3ded568bb545c8c2727396d063bbaf847df05b08cd96", size = 1409336, upload-time = "2024-04-26T00:35:30.392Z" }, - { url = "https://files.pythonhosted.org/packages/9b/f7/591d601c3046ceb65b97291dfe87fa25124cffac3d97aaaba89d0f0d7bdf/black-24.4.2-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:257d724c2c9b1660f353b36c802ccece186a30accc7742c176d29c146df6e474", size = 1615013, upload-time = "2024-04-26T00:39:49.415Z" }, - { url = "https://files.pythonhosted.org/packages/c9/17/5e0036b265bbf6bc44970d93d48febcbc03701b671db3c9603fd43ebc616/black-24.4.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:bdde6f877a18f24844e381d45e9947a49e97933573ac9d4345399be37621e26c", size = 1436163, upload-time = "2024-04-26T00:40:20.267Z" }, - { url = "https://files.pythonhosted.org/packages/c5/48/34176b522e8cff4620a5d96c2e323ff2413f574870eb25efa8025885e028/black-24.4.2-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e151054aa00bad1f4e1f04919542885f89f5f7d086b8a59e5000e6c616896ffb", size = 1803382, upload-time = "2024-04-26T00:34:38.665Z" }, - { url = "https://files.pythonhosted.org/packages/74/ce/e8eec1a77edbfa982bee3b5460dcdd4fe0e4e3165fc15d8ec44d04da7776/black-24.4.2-cp311-cp311-win_amd64.whl", hash = "sha256:7e122b1c4fb252fd85df3ca93578732b4749d9be076593076ef4d07a0233c3e1", size = 1417802, upload-time = "2024-04-26T00:35:08.804Z" }, - { url = "https://files.pythonhosted.org/packages/f4/75/3a29de3bda4006cc280d833b5d961cf7df3810a21f49e7a63a7e551fb351/black-24.4.2-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:accf49e151c8ed2c0cdc528691838afd217c50412534e876a19270fea1e28e2d", size = 1645176, upload-time = "2024-04-26T00:42:35.606Z" }, - { url = "https://files.pythonhosted.org/packages/be/b8/9c152301774fa62a265b035a8ede4d6280827904ea1af8c3be10a28d3187/black-24.4.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:88c57dc656038f1ab9f92b3eb5335ee9b021412feaa46330d5eba4e51fe49b04", size = 1446227, upload-time = "2024-04-26T00:40:35.195Z" }, - { url = "https://files.pythonhosted.org/packages/25/6d/eb15a1b155f755f43766cc473618c6e1de6555d6a1764965643f486dcf01/black-24.4.2-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:be8bef99eb46d5021bf053114442914baeb3649a89dc5f3a555c88737e5e98fc", size = 1832011, upload-time = "2024-04-26T00:34:37.825Z" }, - { url = "https://files.pythonhosted.org/packages/43/24/942b22571b0171be7c6f701cdc3e3b7221f5b522ef02cf82503a547a657b/black-24.4.2-cp312-cp312-win_amd64.whl", hash = "sha256:415e686e87dbbe6f4cd5ef0fbf764af7b89f9057b97c908742b6008cc554b9c0", size = 1428800, upload-time = "2024-04-26T00:35:55.838Z" }, - { url = "https://files.pythonhosted.org/packages/0f/89/294c9a6b6c75a08da55e9d05321d0707e9418735e3062b12ef0f54c33474/black-24.4.2-py3-none-any.whl", hash = "sha256:d36ed1124bb81b32f8614555b34cc4259c3fbc7eec17870e8ff8ded335b58d8c", size = 205925, upload-time = "2024-04-26T00:32:12.495Z" }, +sdist = { url = "https://files.pythonhosted.org/packages/11/5f/25b7b149b8b7d3b958efa4faa56446560408c0f2651108a517526de0320a/black-26.3.0.tar.gz", hash = "sha256:4d438dfdba1c807c6c7c63c4f15794dda0820d2222e7c4105042ac9ddfc5dd0b", size = 664127, upload-time = "2026-03-06T17:42:33.7Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e9/45/0df73428226c2197b8b1e2ca15654f85cece1efe5f060c910b641a35de4a/black-26.3.0-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:135bf8a352e35b3bfba4999c256063d8d86514654599eca7635e914a55d60ec3", size = 1866623, upload-time = "2026-03-06T17:46:07.622Z" }, + { url = "https://files.pythonhosted.org/packages/40/e1/7467fcccf3532853b013bee22c9cdef6aa3314a58ccc73eb5a8a2750e50e/black-26.3.0-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:6024a2959b6c62c311c564ce23ce0eaa977a50ed52a53f7abc83d2c9eb62b8d8", size = 1703733, upload-time = "2026-03-06T17:46:09.334Z" }, + { url = "https://files.pythonhosted.org/packages/e8/72/ceb0a5091b6dff654f77ee6488b91d45fbea1385338798935eb83090d27e/black-26.3.0-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:264144203ea3374542a1591b6fb317561662d074bce5d91ad6afa8d8d3e4ec3d", size = 1768094, upload-time = "2026-03-06T17:46:11.182Z" }, + { url = "https://files.pythonhosted.org/packages/49/cc/6af7e15fb728f30f3e3d4257d2f3d3fe5c5f4ada30b0e8feb92f50118d5c/black-26.3.0-cp310-cp310-win_amd64.whl", hash = "sha256:1a15d1386dce3af3993bf9baeb68d3e492cbb003dae05c3ecf8530a9b75edf85", size = 1413004, upload-time = "2026-03-06T17:46:12.867Z" }, + { url = "https://files.pythonhosted.org/packages/c4/04/7f5ffd40078ab54efa738797e1d547a3fce893f1de212a7a2e65b4a36254/black-26.3.0-cp310-cp310-win_arm64.whl", hash = "sha256:d86a70bf048235aff62a79e229fe5d9e7809c7a05a3dd12982e7ccdc2678e096", size = 1219839, upload-time = "2026-03-06T17:46:14.133Z" }, + { url = "https://files.pythonhosted.org/packages/f9/ec/e4db9f2b2db8226ae20d48b589c69fd64477657bf241c8ccaea3bc4feafa/black-26.3.0-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:3da07abe65732483e915ab7f9c7c50332c293056436e9519373775d62539607c", size = 1851905, upload-time = "2026-03-06T17:46:15.447Z" }, + { url = "https://files.pythonhosted.org/packages/62/2c/ccecfcbd6a0610ecf554e852a146f053eaeb5b281dd9cb634338518c765e/black-26.3.0-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:fc9fd683ccabc3dc9791b93db494d93b5c6c03b105453b76d71e5474e9dfa6e7", size = 1689299, upload-time = "2026-03-06T17:46:17.396Z" }, + { url = "https://files.pythonhosted.org/packages/1a/53/8dcb860242012d6da9c6b1b930c3e4c947eb42feb1fc70f2a4e7332c90c5/black-26.3.0-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8e2c7e2c5ee09ff575869258b2c07064c952637918fc5e15f6ebd45e45eae0aa", size = 1753902, upload-time = "2026-03-06T17:46:19.592Z" }, + { url = "https://files.pythonhosted.org/packages/5d/21/f37b3efcc8cf2d01ec9eb5466598aa53bed2292db236723ac4571e24c4de/black-26.3.0-cp311-cp311-win_amd64.whl", hash = "sha256:a849286bfc3054eaeb233b6df9056fcf969ee18bf7ecb71b0257e838a0f05e6d", size = 1413841, upload-time = "2026-03-06T17:46:20.981Z" }, + { url = "https://files.pythonhosted.org/packages/eb/74/e70f5f2a74301d8f10276b90715699d51d7db1c3dd79cf13966d32ba7b18/black-26.3.0-cp311-cp311-win_arm64.whl", hash = "sha256:c93c83af43cda73ed8265d001214779ab245fa7a861a75b3e43828f4fb1f5657", size = 1220105, upload-time = "2026-03-06T17:46:23.269Z" }, + { url = "https://files.pythonhosted.org/packages/1d/76/b21711045b7f4c4f1774048d0b34dd10a265c42255658b251ce3303ae3c7/black-26.3.0-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:c2b1e5eec220b419e3591a0aaa6351bd3a9c01fe6291fbaf76d84308eb7a2ede", size = 1895944, upload-time = "2026-03-06T17:46:24.841Z" }, + { url = "https://files.pythonhosted.org/packages/f2/c3/8c56e73283326bc92a36101c660228fff09a2403a57a03cacf3f7f84cf62/black-26.3.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:1bab64de70bccc992432bee56cdffbe004ceeaa07352127c386faa87e81f9261", size = 1718669, upload-time = "2026-03-06T17:46:26.639Z" }, + { url = "https://files.pythonhosted.org/packages/7b/8b/712a3ae8f17c1f3cd6f9ac2fffb167a27192f5c7aba68724e8c4ab8474ad/black-26.3.0-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:5b6c5f734290803b7b26493ffd734b02b72e6c90d82d45ac4d5b862b9bdf7720", size = 1794844, upload-time = "2026-03-06T17:46:28.334Z" }, + { url = "https://files.pythonhosted.org/packages/ba/5b/ee955040e446df86473287dd24dc69c80dd05e02cc358bca90e22059f7b1/black-26.3.0-cp312-cp312-win_amd64.whl", hash = "sha256:7c767396af15b54e1a6aae99ddf241ae97e589f666b1d22c4b6618282a04e4ca", size = 1420461, upload-time = "2026-03-06T17:46:29.965Z" }, + { url = "https://files.pythonhosted.org/packages/12/77/40b8bd44f032bb34c9ebf47ffc5bb47a2520d29e0a4b8a780ab515223b5a/black-26.3.0-cp312-cp312-win_arm64.whl", hash = "sha256:765fd6ddd00f35c55250fdc6b790c272d54ac3f44da719cc42df428269b45980", size = 1229667, upload-time = "2026-03-06T17:46:31.654Z" }, + { url = "https://files.pythonhosted.org/packages/28/c3/21a834ce3de02c64221243f2adac63fa3c3f441efdb3adbf4136b33dfeb0/black-26.3.0-cp313-cp313-macosx_10_13_x86_64.whl", hash = "sha256:59754fd8f43ef457be190594c07a52c999e22cb1534dc5344bff1d46fdf1027d", size = 1895195, upload-time = "2026-03-06T17:46:33.12Z" }, + { url = "https://files.pythonhosted.org/packages/1c/f9/212d9697dd78362dadb778d4616b74c8c2cf7f2e4a55aac2adeb0576f2e9/black-26.3.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:1fd94cfee67b8d336761a0b08629a25938e4a491c440951ce517a7209c99b5ff", size = 1718472, upload-time = "2026-03-06T17:46:34.576Z" }, + { url = "https://files.pythonhosted.org/packages/a2/dd/da980b2f512441375b73cb511f38a2c3db4be83ccaa1302b8d39c9fa2dff/black-26.3.0-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:6f7b3e653a90ca1ef4e821c20f8edaee80b649c38d2532ed2e9073a9534b14a7", size = 1793741, upload-time = "2026-03-06T17:46:36.261Z" }, + { url = "https://files.pythonhosted.org/packages/93/11/cd69ae8826fe3bc6eaf525c8c557266d522b258154a2968eb46d6d25fac7/black-26.3.0-cp313-cp313-win_amd64.whl", hash = "sha256:f8fb9d7c2496adc83614856e1f6e55a9ce4b7ae7fc7f45b46af9189ddb493464", size = 1422522, upload-time = "2026-03-06T17:46:37.607Z" }, + { url = "https://files.pythonhosted.org/packages/75/f5/647cf50255203eb286be197925e86eedc101d5409147505db3e463229228/black-26.3.0-cp313-cp313-win_arm64.whl", hash = "sha256:e8618c1d06838f56afbcb3ffa1aa16436cec62b86b38c7b32ca86f53948ffb91", size = 1231807, upload-time = "2026-03-06T17:46:39.072Z" }, + { url = "https://files.pythonhosted.org/packages/ff/77/b197e701f15fd694d20d8ee0001efa2e29eba917aa7c3610ff7b10ae0f88/black-26.3.0-cp314-cp314-macosx_10_15_x86_64.whl", hash = "sha256:d0c6f64ead44f4369c66f1339ecf68e99b40f2e44253c257f7807c5a3ef0ca32", size = 1889209, upload-time = "2026-03-06T17:46:40.453Z" }, + { url = "https://files.pythonhosted.org/packages/93/85/b4d4924ac898adc2e39fc7a923bed99797535bc16dea4bc63944c3903c2b/black-26.3.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:ed6f0809134e51ec4a7509e069cdfa42bf996bd0fd1df6d3146b907f36e28893", size = 1720830, upload-time = "2026-03-06T17:46:42.009Z" }, + { url = "https://files.pythonhosted.org/packages/00/b1/5c0bf29fe5b43fcc6f3e8480c6566d21a02d4e702b3846944e7daa06dea9/black-26.3.0-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:cc6ac0ea5dd5fa6311ca82edfa3620cba0ed0426022d10d2d5d39aedbf3e1958", size = 1787676, upload-time = "2026-03-06T17:46:43.382Z" }, + { url = "https://files.pythonhosted.org/packages/b8/ce/cc8cf14806c144d6a16512272c537d5450f50675d3e8c038705430e90fd9/black-26.3.0-cp314-cp314-win_amd64.whl", hash = "sha256:884bc0aefa96adabcba0b77b10e9775fd52d4b766e88c44dc6f41f7c82787fc8", size = 1445406, upload-time = "2026-03-06T17:46:44.948Z" }, + { url = "https://files.pythonhosted.org/packages/cf/bb/049ea0fad9f8bdec7b647948adcf74bb720bd71dcb213decd553e05b2699/black-26.3.0-cp314-cp314-win_arm64.whl", hash = "sha256:be3bd02aab5c4ab03703172f5530ddc8fc8b5b7bb8786230e84c9e011cee9ca1", size = 1257945, upload-time = "2026-03-06T17:46:46.432Z" }, + { url = "https://files.pythonhosted.org/packages/39/d7/7360654ba4f8b41afcaeb5aca973cfea5591da75aff79b0a8ae0bb8883f6/black-26.3.0-py3-none-any.whl", hash = "sha256:e825d6b121910dff6f04d7691f826d2449327e8e71c26254c030c4f3d2311985", size = 206848, upload-time = "2026-03-06T17:42:31.133Z" }, ] [[package]] @@ -769,7 +798,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 = [ @@ -956,7 +985,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 = [ @@ -1185,7 +1214,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" }, @@ -1238,37 +1267,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]] @@ -1538,8 +1567,8 @@ name = "emerging-optimizers" version = "0.3.0" source = { git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0#b309e2f01cda75dc96a6dc1a2355a7b3b64b5e16" } 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]] @@ -1547,7 +1576,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 = [ @@ -1630,15 +1659,14 @@ wheels = [ [[package]] name = "fla-core" -version = "0.4.2" +version = "0.5.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "einops" }, - { name = "torch", marker = "sys_platform == 'never'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/53/f9/9e05c48f92b1388a8a357141eb557ed0dd6d4bb936e1d05d35f01976657f/fla_core-0.4.2.tar.gz", hash = "sha256:e9fef6fcdf122029f9feb7dccfeb85eb9650e6aabc72d2a65b36558e9c590edd", size = 377722, upload-time = "2026-03-12T14:45:46.101Z" } +sdist = { url = "https://files.pythonhosted.org/packages/4e/62/99e149f19a447ce809d7f4fa64ae61c073da337505b9db7f389502470820/fla_core-0.5.1.tar.gz", hash = "sha256:7f3cf56edfbaa9115f4937d1181372e5c7b11809ad8eb2e411fffc3caf729f48", size = 498800, upload-time = "2026-06-18T18:17:15.377Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/ee/36/3c303f92bafea7c3f97d68bbb83d18cc42e30cd0bfb1b7cfe589360f11d6/fla_core-0.4.2-py3-none-any.whl", hash = "sha256:cba3db29380002da3cbfc0db94d6efac19aaf528900d19c05c2765e8f3cc485b", size = 510239, upload-time = "2026-03-12T14:45:43.708Z" }, + { url = "https://files.pythonhosted.org/packages/ce/78/a55ee7a62515dcb9220770dd99dfe59ac6599da8af84ad1d20f9f407df4d/fla_core-0.5.1-py3-none-any.whl", hash = "sha256:02150d34aa1e37f6b8ed9b2feec5d29af93680573e5077ed981238179c11fb06", size = 702955, upload-time = "2026-06-18T18:17:12.229Z" }, ] [[package]] @@ -1657,15 +1685,20 @@ wheels = [ [[package]] name = "flash-linear-attention" -version = "0.4.2" +version = "0.5.1" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "fla-core" }, { name = "transformers" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/d1/cb/46cc27a829a10b308927c5dbc99176906a021bb0770253699e93f3cd81a0/flash_linear_attention-0.4.2.tar.gz", hash = "sha256:f97c01ebe7cf390323af07dd3fb65ade07da16724339bf70c78607bc0c007c34", size = 148464, upload-time = "2026-03-12T14:45:46.945Z" } +sdist = { url = "https://files.pythonhosted.org/packages/4c/e8/8f115be585a046795e4a1f7ade727889219bb66b8d96cc668b8c9e437c0e/flash_linear_attention-0.5.1.tar.gz", hash = "sha256:8840fd4c37de8b0612dc8fd493867f3d330672ba2f17c024a2ce37239634e247", size = 221733, upload-time = "2026-06-18T18:17:16.661Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/60/ee/a3cba17965482b35c4990af90bad108e82c32edcb59911c37f318b5f4198/flash_linear_attention-0.4.2-py3-none-any.whl", hash = "sha256:c08be006ce4dbe1be81f54938ee8e6fc7968cfba397c8d06c7669e97b8c44c0d", size = 284661, upload-time = "2026-03-12T14:45:44.905Z" }, + { url = "https://files.pythonhosted.org/packages/68/3c/5819fb19dc071302ca818616a4e64d4454e1f1193929eb8365c8d38e9052/flash_linear_attention-0.5.1-py3-none-any.whl", hash = "sha256:9022862f0a238752372c81290694b8b2ed1cb2d13fc40d340ffeeb554d6bee5c", size = 403096, upload-time = "2026-06-18T18:17:14.025Z" }, +] + +[package.optional-dependencies] +tilelang = [ + { name = "tilelang" }, ] [[package]] @@ -2276,7 +2309,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" }, @@ -2714,7 +2747,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 = [ @@ -2788,7 +2821,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 = [ @@ -2929,7 +2962,7 @@ dev = [ { name = "einops" }, { name = "emerging-optimizers", marker = "python_full_version >= '3.12'" }, { name = "fastapi" }, - { name = "flash-linear-attention" }, + { name = "flash-linear-attention", extra = ["tilelang"], marker = "extra == 'extra-13-megatron-core-dev'" }, { name = "flashinfer-python" }, { name = "hypercorn" }, { name = "mamba-ssm" }, @@ -3057,7 +3090,7 @@ requires-dist = [ { name = "emerging-optimizers", marker = "python_full_version >= '3.12' and extra == 'lts'", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.3.0" }, { name = "fastapi", marker = "extra == 'dev'", specifier = "~=0.50" }, { name = "fastapi", marker = "extra == 'lts'", specifier = "~=0.50" }, - { name = "flash-linear-attention", marker = "extra == 'dev'", specifier = ">=0.4.2,<0.5" }, + { name = "flash-linear-attention", extras = ["tilelang"], marker = "extra == 'dev'", specifier = "==0.5.1" }, { name = "flashinfer-python", marker = "extra == 'dev'", specifier = ">=0.5.0,<0.7.0" }, { name = "flashinfer-python", marker = "extra == 'lts'", specifier = ">=0.5.0,<0.7.0" }, { name = "flask-restful", marker = "extra == 'mlm'" }, @@ -3125,7 +3158,7 @@ docs = [ { name = "sphinx-copybutton" }, ] linting = [ - { name = "black", specifier = "==24.4.2" }, + { name = "black", specifier = "==26.3.0" }, { name = "flake8", specifier = "==7.1.0" }, { name = "isort", specifier = "==5.13.2" }, { name = "pylint", specifier = "==3.2.6" }, @@ -3434,7 +3467,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 = [ @@ -4095,7 +4128,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" }, @@ -4203,7 +4236,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" }, @@ -4235,9 +4268,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" }, @@ -4250,7 +4283,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" }, @@ -4930,10 +4963,10 @@ resolution-markers = [ "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' and extra == 'extra-13-megatron-core-lts')" }, + { name = "python-dateutil", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "pytz", marker = "python_full_version < '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "tzdata", 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/33/01/d40b85317f86cf08d853a4f495195c73815fdf205eef3993821720274518/pandas-2.3.3.tar.gz", hash = "sha256:e05e1af93b977f7eafa636d043f9f94c7ee3ac81af99c13508215942e64c993b", size = 4495223, upload-time = "2025-09-29T23:34:51.853Z" } wheels = [ @@ -5053,9 +5086,9 @@ 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 = "numpy", version = "2.4.4", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" }, - { name = "python-dateutil", marker = "python_full_version >= '3.11'" }, - { name = "tzdata", marker = "(python_full_version >= '3.11' and sys_platform == 'emscripten') or (python_full_version >= '3.11' and sys_platform == 'win32')" }, + { name = "numpy", version = "2.4.4", 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 = "python-dateutil", marker = "python_full_version >= '3.11' or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "tzdata", 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' 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')" }, ] sdist = { url = "https://files.pythonhosted.org/packages/da/99/b342345300f13440fe9fe385c3c481e2d9a595ee3bab4d3219247ac94e9a/pandas-3.0.2.tar.gz", hash = "sha256:f4753e73e34c8d83221ba58f232433fca2748be8b18dbca02d242ed153945043", size = 4645855, upload-time = "2026-03-31T06:48:30.816Z" } wheels = [ @@ -5998,6 +6031,45 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/a3/73/89930efabd4da63cea44a3f438aeb753d600123570e6d6264e763617a9ce/python_multipart-0.0.24-py3-none-any.whl", hash = "sha256:9b110a98db707df01a53c194f0af075e736a770dc5058089650d70b4a182f950", size = 24420, upload-time = "2026-04-05T20:49:12.555Z" }, ] +[[package]] +name = "pytokens" +version = "0.4.1" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/b6/34/b4e015b99031667a7b960f888889c5bd34ef585c85e1cb56a594b92836ac/pytokens-0.4.1.tar.gz", hash = "sha256:292052fe80923aae2260c073f822ceba21f3872ced9a68bb7953b348e561179a", size = 23015, upload-time = "2026-01-30T01:03:45.924Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/42/24/f206113e05cb8ef51b3850e7ef88f20da6f4bf932190ceb48bd3da103e10/pytokens-0.4.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:2a44ed93ea23415c54f3face3b65ef2b844d96aeb3455b8a69b3df6beab6acc5", size = 161522, upload-time = "2026-01-30T01:02:50.393Z" }, + { url = "https://files.pythonhosted.org/packages/d4/e9/06a6bf1b90c2ed81a9c7d2544232fe5d2891d1cd480e8a1809ca354a8eb2/pytokens-0.4.1-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:add8bf86b71a5d9fb5b89f023a80b791e04fba57960aa790cc6125f7f1d39dfe", size = 246945, upload-time = "2026-01-30T01:02:52.399Z" }, + { url = "https://files.pythonhosted.org/packages/69/66/f6fb1007a4c3d8b682d5d65b7c1fb33257587a5f782647091e3408abe0b8/pytokens-0.4.1-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:670d286910b531c7b7e3c0b453fd8156f250adb140146d234a82219459b9640c", size = 259525, upload-time = "2026-01-30T01:02:53.737Z" }, + { url = "https://files.pythonhosted.org/packages/04/92/086f89b4d622a18418bac74ab5db7f68cf0c21cf7cc92de6c7b919d76c88/pytokens-0.4.1-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:4e691d7f5186bd2842c14813f79f8884bb03f5995f0575272009982c5ac6c0f7", size = 262693, upload-time = "2026-01-30T01:02:54.871Z" }, + { url = "https://files.pythonhosted.org/packages/b4/7b/8b31c347cf94a3f900bdde750b2e9131575a61fdb620d3d3c75832262137/pytokens-0.4.1-cp310-cp310-win_amd64.whl", hash = "sha256:27b83ad28825978742beef057bfe406ad6ed524b2d28c252c5de7b4a6dd48fa2", size = 103567, upload-time = "2026-01-30T01:02:56.414Z" }, + { url = "https://files.pythonhosted.org/packages/3d/92/790ebe03f07b57e53b10884c329b9a1a308648fc083a6d4a39a10a28c8fc/pytokens-0.4.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:d70e77c55ae8380c91c0c18dea05951482e263982911fc7410b1ffd1dadd3440", size = 160864, upload-time = "2026-01-30T01:02:57.882Z" }, + { url = "https://files.pythonhosted.org/packages/13/25/a4f555281d975bfdd1eba731450e2fe3a95870274da73fb12c40aeae7625/pytokens-0.4.1-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:4a58d057208cb9075c144950d789511220b07636dd2e4708d5645d24de666bdc", size = 248565, upload-time = "2026-01-30T01:02:59.912Z" }, + { url = "https://files.pythonhosted.org/packages/17/50/bc0394b4ad5b1601be22fa43652173d47e4c9efbf0044c62e9a59b747c56/pytokens-0.4.1-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b49750419d300e2b5a3813cf229d4e5a4c728dae470bcc89867a9ad6f25a722d", size = 260824, upload-time = "2026-01-30T01:03:01.471Z" }, + { url = "https://files.pythonhosted.org/packages/4e/54/3e04f9d92a4be4fc6c80016bc396b923d2a6933ae94b5f557c939c460ee0/pytokens-0.4.1-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:d9907d61f15bf7261d7e775bd5d7ee4d2930e04424bab1972591918497623a16", size = 264075, upload-time = "2026-01-30T01:03:04.143Z" }, + { url = "https://files.pythonhosted.org/packages/d1/1b/44b0326cb5470a4375f37988aea5d61b5cc52407143303015ebee94abfd6/pytokens-0.4.1-cp311-cp311-win_amd64.whl", hash = "sha256:ee44d0f85b803321710f9239f335aafe16553b39106384cef8e6de40cb4ef2f6", size = 103323, upload-time = "2026-01-30T01:03:05.412Z" }, + { url = "https://files.pythonhosted.org/packages/41/5d/e44573011401fb82e9d51e97f1290ceb377800fb4eed650b96f4753b499c/pytokens-0.4.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:140709331e846b728475786df8aeb27d24f48cbcf7bcd449f8de75cae7a45083", size = 160663, upload-time = "2026-01-30T01:03:06.473Z" }, + { url = "https://files.pythonhosted.org/packages/f0/e6/5bbc3019f8e6f21d09c41f8b8654536117e5e211a85d89212d59cbdab381/pytokens-0.4.1-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:6d6c4268598f762bc8e91f5dbf2ab2f61f7b95bdc07953b602db879b3c8c18e1", size = 255626, upload-time = "2026-01-30T01:03:08.177Z" }, + { url = "https://files.pythonhosted.org/packages/bf/3c/2d5297d82286f6f3d92770289fd439956b201c0a4fc7e72efb9b2293758e/pytokens-0.4.1-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:24afde1f53d95348b5a0eb19488661147285ca4dd7ed752bbc3e1c6242a304d1", size = 269779, upload-time = "2026-01-30T01:03:09.756Z" }, + { url = "https://files.pythonhosted.org/packages/20/01/7436e9ad693cebda0551203e0bf28f7669976c60ad07d6402098208476de/pytokens-0.4.1-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:5ad948d085ed6c16413eb5fec6b3e02fa00dc29a2534f088d3302c47eb59adf9", size = 268076, upload-time = "2026-01-30T01:03:10.957Z" }, + { url = "https://files.pythonhosted.org/packages/2e/df/533c82a3c752ba13ae7ef238b7f8cdd272cf1475f03c63ac6cf3fcfb00b6/pytokens-0.4.1-cp312-cp312-win_amd64.whl", hash = "sha256:3f901fe783e06e48e8cbdc82d631fca8f118333798193e026a50ce1b3757ea68", size = 103552, upload-time = "2026-01-30T01:03:12.066Z" }, + { url = "https://files.pythonhosted.org/packages/cb/dc/08b1a080372afda3cceb4f3c0a7ba2bde9d6a5241f1edb02a22a019ee147/pytokens-0.4.1-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:8bdb9d0ce90cbf99c525e75a2fa415144fd570a1ba987380190e8b786bc6ef9b", size = 160720, upload-time = "2026-01-30T01:03:13.843Z" }, + { url = "https://files.pythonhosted.org/packages/64/0c/41ea22205da480837a700e395507e6a24425151dfb7ead73343d6e2d7ffe/pytokens-0.4.1-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:5502408cab1cb18e128570f8d598981c68a50d0cbd7c61312a90507cd3a1276f", size = 254204, upload-time = "2026-01-30T01:03:14.886Z" }, + { url = "https://files.pythonhosted.org/packages/e0/d2/afe5c7f8607018beb99971489dbb846508f1b8f351fcefc225fcf4b2adc0/pytokens-0.4.1-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:29d1d8fb1030af4d231789959f21821ab6325e463f0503a61d204343c9b355d1", size = 268423, upload-time = "2026-01-30T01:03:15.936Z" }, + { url = "https://files.pythonhosted.org/packages/68/d4/00ffdbd370410c04e9591da9220a68dc1693ef7499173eb3e30d06e05ed1/pytokens-0.4.1-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:970b08dd6b86058b6dc07efe9e98414f5102974716232d10f32ff39701e841c4", size = 266859, upload-time = "2026-01-30T01:03:17.458Z" }, + { url = "https://files.pythonhosted.org/packages/a7/c9/c3161313b4ca0c601eeefabd3d3b576edaa9afdefd32da97210700e47652/pytokens-0.4.1-cp313-cp313-win_amd64.whl", hash = "sha256:9bd7d7f544d362576be74f9d5901a22f317efc20046efe2034dced238cbbfe78", size = 103520, upload-time = "2026-01-30T01:03:18.652Z" }, + { url = "https://files.pythonhosted.org/packages/8f/a7/b470f672e6fc5fee0a01d9e75005a0e617e162381974213a945fcd274843/pytokens-0.4.1-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:4a14d5f5fc78ce85e426aa159489e2d5961acf0e47575e08f35584009178e321", size = 160821, upload-time = "2026-01-30T01:03:19.684Z" }, + { url = "https://files.pythonhosted.org/packages/80/98/e83a36fe8d170c911f864bfded690d2542bfcfacb9c649d11a9e6eb9dc41/pytokens-0.4.1-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:97f50fd18543be72da51dd505e2ed20d2228c74e0464e4262e4899797803d7fa", size = 254263, upload-time = "2026-01-30T01:03:20.834Z" }, + { url = "https://files.pythonhosted.org/packages/0f/95/70d7041273890f9f97a24234c00b746e8da86df462620194cef1d411ddeb/pytokens-0.4.1-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:dc74c035f9bfca0255c1af77ddd2d6ae8419012805453e4b0e7513e17904545d", size = 268071, upload-time = "2026-01-30T01:03:21.888Z" }, + { url = "https://files.pythonhosted.org/packages/da/79/76e6d09ae19c99404656d7db9c35dfd20f2086f3eb6ecb496b5b31163bad/pytokens-0.4.1-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:f66a6bbe741bd431f6d741e617e0f39ec7257ca1f89089593479347cc4d13324", size = 271716, upload-time = "2026-01-30T01:03:23.633Z" }, + { url = "https://files.pythonhosted.org/packages/79/37/482e55fa1602e0a7ff012661d8c946bafdc05e480ea5a32f4f7e336d4aa9/pytokens-0.4.1-cp314-cp314-win_amd64.whl", hash = "sha256:b35d7e5ad269804f6697727702da3c517bb8a5228afa450ab0fa787732055fc9", size = 104539, upload-time = "2026-01-30T01:03:24.788Z" }, + { url = "https://files.pythonhosted.org/packages/30/e8/20e7db907c23f3d63b0be3b8a4fd1927f6da2395f5bcc7f72242bb963dfe/pytokens-0.4.1-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:8fcb9ba3709ff77e77f1c7022ff11d13553f3c30299a9fe246a166903e9091eb", size = 168474, upload-time = "2026-01-30T01:03:26.428Z" }, + { url = "https://files.pythonhosted.org/packages/d6/81/88a95ee9fafdd8f5f3452107748fd04c24930d500b9aba9738f3ade642cc/pytokens-0.4.1-cp314-cp314t-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:79fc6b8699564e1f9b521582c35435f1bd32dd06822322ec44afdeba666d8cb3", size = 290473, upload-time = "2026-01-30T01:03:27.415Z" }, + { url = "https://files.pythonhosted.org/packages/cf/35/3aa899645e29b6375b4aed9f8d21df219e7c958c4c186b465e42ee0a06bf/pytokens-0.4.1-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d31b97b3de0f61571a124a00ffe9a81fb9939146c122c11060725bd5aea79975", size = 303485, upload-time = "2026-01-30T01:03:28.558Z" }, + { url = "https://files.pythonhosted.org/packages/52/a0/07907b6ff512674d9b201859f7d212298c44933633c946703a20c25e9d81/pytokens-0.4.1-cp314-cp314t-musllinux_1_2_x86_64.whl", hash = "sha256:967cf6e3fd4adf7de8fc73cd3043754ae79c36475c1c11d514fc72cf5490094a", size = 306698, upload-time = "2026-01-30T01:03:29.653Z" }, + { url = "https://files.pythonhosted.org/packages/39/2a/cbbf9250020a4a8dd53ba83a46c097b69e5eb49dd14e708f496f548c6612/pytokens-0.4.1-cp314-cp314t-win_amd64.whl", hash = "sha256:584c80c24b078eec1e227079d56dc22ff755e0ba8654d8383b2c549107528918", size = 116287, upload-time = "2026-01-30T01:03:30.912Z" }, + { url = "https://files.pythonhosted.org/packages/c6/78/397db326746f0a342855b81216ae1f0a32965deccfd7c830a2dbc66d2483/pytokens-0.4.1-py3-none-any.whl", hash = "sha256:26cef14744a8385f35d0e095dc8b3a7583f6c953c2e3d269c7f82484bf5ad2de", size = 13729, upload-time = "2026-01-30T01:03:45.029Z" }, +] + [[package]] name = "pytz" version = "2026.1.post1" @@ -6252,7 +6324,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 = [ @@ -7311,7 +7383,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 = [ @@ -7323,7 +7395,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 (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')" }, ] 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 = [ @@ -7580,6 +7652,31 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/af/df/c7891ef9d2712ad774777271d39fdef63941ffba0a9d59b7ad1fd2765e57/tiktoken-0.12.0-cp314-cp314t-win_amd64.whl", hash = "sha256:f61c0aea5565ac82e2ec50a05e02a6c44734e91b51c10510b084ea1b8e633a71", size = 920667, upload-time = "2025-10-06T20:22:34.444Z" }, ] +[[package]] +name = "tilelang" +version = "0.1.9" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "apache-tvm-ffi" }, + { name = "cloudpickle" }, + { name = "ml-dtypes" }, + { 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 (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-dev') or (extra == 'extra-13-megatron-core-dev' and extra == 'extra-13-megatron-core-lts')" }, + { name = "psutil" }, + { name = "setuptools", marker = "sys_platform == 'darwin'" }, + { name = "torch", marker = "sys_platform == 'never'" }, + { name = "torch-c-dlpack-ext", marker = "python_full_version < '3.14'" }, + { name = "tqdm" }, + { name = "typing-extensions" }, + { name = "z3-solver" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/56/70/5051f65821baa30a3d61fc48f8ba10c776490315e8c90f82559b92089756/tilelang-0.1.9.tar.gz", hash = "sha256:287f727c913bb648fcf6c1968809ba3390e55eeed257a5c6bb9a80bc05966af4", size = 93395292, upload-time = "2026-04-22T09:19:11.988Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/90/db/4dd76da8c8585c605639a21bc098d504e317fe324a72f01ce3c7370250b4/tilelang-0.1.9-cp38-abi3-macosx_11_0_arm64.whl", hash = "sha256:00ed594fdeb229c5505b9ffa895c3c5daeb28641c78f783fa1f724cf1e08cecd", size = 36599020, upload-time = "2026-04-22T09:14:39.366Z" }, + { url = "https://files.pythonhosted.org/packages/f7/8a/1cbeee79d62abaa02441c2d00621554e41aa62dbf3b94a4feb3867184b01/tilelang-0.1.9-cp38-abi3-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:4bbccfe9035aed775ffafb6dc25a5994504b24e2c5d95d0f39643edfafa7bf12", size = 45419374, upload-time = "2026-04-22T09:15:56.014Z" }, + { url = "https://files.pythonhosted.org/packages/c6/a7/f4bfb86f87e107703146e703204cec2c0eae2492b633e0052b0ace3febb6/tilelang-0.1.9-cp38-abi3-manylinux_2_34_aarch64.whl", hash = "sha256:77ab0ee2f40f66ea015b6b21426d482751e28cbc635ef9d1198cbd6502454a7c", size = 42110365, upload-time = "2026-04-22T09:17:18.292Z" }, +] + [[package]] name = "tokenizers" version = "0.22.2" @@ -7687,21 +7784,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" }, @@ -7791,7 +7888,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 = [ @@ -7925,11 +8022,11 @@ wheels = [ [[package]] name = "urllib3" -version = "1.26.20" +version = "2.7.0" source = { registry = "https://pypi.org/simple" } -sdist = { url = "https://files.pythonhosted.org/packages/e4/e8/6ff5e6bc22095cfc59b6ea711b687e2b7ed4bdb373f7eeec370a97d7392f/urllib3-1.26.20.tar.gz", hash = "sha256:40c2dc0c681e47eb8f90e7e27bf6ff7df2e677421fd46756da1161c39ca70d32", size = 307380, upload-time = "2024-08-29T15:43:11.37Z" } +sdist = { url = "https://files.pythonhosted.org/packages/53/0c/06f8b233b8fd13b9e5ee11424ef85419ba0d8ba0b3138bf360be2ff56953/urllib3-2.7.0.tar.gz", hash = "sha256:231e0ec3b63ceb14667c67be60f2f2c40a518cb38b03af60abc813da26505f4c", size = 433602, upload-time = "2026-05-07T16:13:18.596Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/33/cf/8435d5a7159e2a9c83a95896ed596f68cf798005fe107cc655b5c5c14704/urllib3-1.26.20-py2.py3-none-any.whl", hash = "sha256:0ed14ccfbf1c30a9072c7ca157e4319b70d65f623e91e7b32fadb2853431016e", size = 144225, upload-time = "2024-08-29T15:43:08.921Z" }, + { url = "https://files.pythonhosted.org/packages/7f/3e/5db95bcf282c52709639744ca2a8b149baccf648e39c8cc87553df9eae0c/urllib3-2.7.0-py3-none-any.whl", hash = "sha256:9fb4c81ebbb1ce9531cce37674bbc6f1360472bc18ca9a553ede278ef7276897", size = 131087, upload-time = "2026-05-07T16:13:17.151Z" }, ] [[package]] @@ -8634,6 +8731,20 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/69/68/c8739671f5699c7dc470580a4f821ef37c32c4cb0b047ce223a7f115757f/yarl-1.23.0-py3-none-any.whl", hash = "sha256:a2df6afe50dea8ae15fa34c9f824a3ee958d785fd5d089063d960bae1daa0a3f", size = 48288, upload-time = "2026-03-01T22:07:51.388Z" }, ] +[[package]] +name = "z3-solver" +version = "4.15.4.0" +source = { registry = "https://pypi.org/simple" } +sdist = { url = "https://files.pythonhosted.org/packages/8a/8e/0c8f17309549d2e5cde9a3ccefa6365437f1e7bafe71878eaf9478e47b18/z3_solver-4.15.4.0.tar.gz", hash = "sha256:928c29b58c4eb62106da51c1914f6a4a55d0441f8f48a81b9da07950434a8946", size = 5018600, upload-time = "2025-10-29T18:12:03.062Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/63/33/a3d5d2eaeb0f7b3174d57d405437eabb2075d4d50bd9ea0957696c435c7b/z3_solver-4.15.4.0-py3-none-macosx_13_0_arm64.whl", hash = "sha256:407e825cc9211f95ef46bdc8d151bf630e7ab2d62a21d24cd74c09cc5b73f3aa", size = 37052538, upload-time = "2025-10-29T18:11:46.233Z" }, + { url = "https://files.pythonhosted.org/packages/47/84/fd7ffac1551cd9f8d44fe41358f738be670fc4c24dfd514fab503f2cf3e7/z3_solver-4.15.4.0-py3-none-macosx_13_0_x86_64.whl", hash = "sha256:00bd10c5a6a5f6112d3a9a810d0799227e52f76caa860dafa5e00966bb47eb13", size = 39807925, upload-time = "2025-10-29T18:11:49.81Z" }, + { url = "https://files.pythonhosted.org/packages/21/c9/bb51a96af0091324c81b803f16c49f719f9f6ea0b0bb52200f5c97ec4892/z3_solver-4.15.4.0-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e103a6f203f505b8b8b8e5c931cc407c95b61556512d4921c1ddc0b3f41b08e", size = 29268352, upload-time = "2025-10-29T18:11:53.032Z" }, + { url = "https://files.pythonhosted.org/packages/bf/2e/0b49f7e4e53817cfb09a0f6585012b782dfe0b666e8abefcb4fac0570606/z3_solver-4.15.4.0-py3-none-manylinux_2_34_aarch64.whl", hash = "sha256:62c7e9cbdd711932301f29919ad9158de9b2f58b4d281dd259bbcd0a2f408ba1", size = 27226534, upload-time = "2025-10-29T18:11:55.59Z" }, + { url = "https://files.pythonhosted.org/packages/26/91/33de49538444d4aafbe47415c450c2f9abab1733e1226f276b496672f46c/z3_solver-4.15.4.0-py3-none-win32.whl", hash = "sha256:be3bc916545c96ffbf89e00d07104ff14f78336e55db069177a1bfbcc01b269d", size = 13191672, upload-time = "2025-10-29T18:11:58.424Z" }, + { url = "https://files.pythonhosted.org/packages/03/d6/a0b135e4419df475177ae78fc93c422430b0fd8875649486f9a5989772e6/z3_solver-4.15.4.0-py3-none-win_amd64.whl", hash = "sha256:00e35b02632ed085ea8199fb230f6015e6fc40554a6680c097bd5f060e827431", size = 16259597, upload-time = "2025-10-29T18:12:01.14Z" }, +] + [[package]] name = "zipp" version = "3.23.0"