diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 24d45c9bbdb..412550b5d24 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -303,6 +303,7 @@ def __init__( cp_comm_type: str | None = None, pg_collection: ProcessGroupCollection | None = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ): """ @@ -314,6 +315,7 @@ def __init__( self.config = config self.layer_number = layer_number self._pp_layer_offset = pp_layer_offset + self.is_mtp_layer = is_mtp_layer self.attn_mask_type = attn_mask_type self.attention_type = attention_type @@ -1636,6 +1638,7 @@ def __init__( cp_comm_type: str | None = None, pg_collection: ProcessGroupCollection | None = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ): """ @@ -1651,6 +1654,7 @@ def __init__( cp_comm_type=cp_comm_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) @@ -2052,6 +2056,7 @@ def __init__( attn_mask_type: AttnMaskType = AttnMaskType.padding, cp_comm_type: str | None = None, pg_collection: ProcessGroupCollection | None = None, + is_mtp_layer: bool = False, name: str | None = None, ): """ @@ -2066,6 +2071,7 @@ def __init__( attention_type="cross", cp_comm_type=cp_comm_type, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, name=name, ) diff --git a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py index 2472d0ec2fb..cb03e3c5c94 100644 --- a/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py +++ b/megatron/core/transformer/experimental_attention_variant/absorbed_mla.py @@ -146,6 +146,7 @@ def __init__( pg_collection: ProcessGroupCollection = None, pp_layer_offset: Optional[int] = None, name: str | None = None, + is_mtp_layer: bool = False, ): if pg_collection is None: pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -160,6 +161,7 @@ def __init__( pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, name=name, + is_mtp_layer=is_mtp_layer, ) assert not config.add_bias_linear, "add_bias_linear is not supported for AbsorbedMLA" diff --git a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py new file mode 100644 index 00000000000..d14e6402b51 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -0,0 +1,694 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + +from dataclasses import dataclass +from typing import NoReturn, Optional, Union + +import torch + +from megatron.core import tensor_parallel +from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.models.common.embeddings import ( + RotaryEmbedding, + YarnRotaryEmbedding, + apply_rotary_pos_emb, +) +from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as off_interface, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.attention import Attention +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.torch_norm import LayerNormBuilder +from megatron.core.transformer.transformer_config import MLATransformerConfig +from megatron.core.typed_torch import apply_module +from megatron.core.utils import get_pg_size, is_te_min_version + +try: + from megatron.core.fusions.fused_mla_yarn_rope_apply import ( + fused_mla_rope_inplace, + fused_mla_rope_out_of_place, + ) +except Exception: + fused_mla_rope_inplace = None + fused_mla_rope_out_of_place = None + + +if HAVE_TE: + from megatron.core.extensions.transformer_engine import TELinear, set_save_original_input +else: + (TEColumnParallelLinear, TELinear, set_save_original_input) = (None, None, None) + + +@torch.compile +def _q_rms_norm(q: torch.Tensor, eps: float) -> torch.Tensor: + """Fused RMS normalization for query tensor (no learnable weight).""" + return q * torch.rsqrt(q.square().mean(-1, keepdim=True) + eps) + + +@dataclass +class DSv4HybridSelfAttentionSubmodules: + """Submodules for the DSv4HybridAttention layer.""" + + q_layernorm: LayerNormBuilder + kv_layernorm: LayerNormBuilder + + linear_q_down_proj: Union[ModuleSpec, type] = None + linear_q_up_proj: Union[ModuleSpec, type] = None + linear_kv_proj: Union[ModuleSpec, type] = None + core_attention: Union[ModuleSpec, type] = None + linear_proj: Union[ModuleSpec, type] = None + + +class DSv4HybridAttention(Attention): + """DeepSeek-v4 Hybrid Attention layer.""" + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSv4HybridSelfAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType, + attention_type: str, + cp_comm_type: Optional[str] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, + name: str | None = None, + ) -> None: + + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + attention_type=attention_type, + attn_mask_type=attn_mask_type, + cp_comm_type=cp_comm_type, + pg_collection=pg_collection, + pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, + name=name, + ) + self.config: MLATransformerConfig + + assert ( + get_pg_size(self.pg_collection.tp) == 1 + ), "DSv4 Hybrid Attention only supports TP size 1." + + assert ( + not self.checkpoint_core_attention + ), "Checkpoint core attention is not supported in DSv4 Hybrid Attention." + assert ( + not self.offload_qkv_linear + ), "Offload qkv linear is not supported in DSv4 Hybrid Attention." + + self.query_projection_size = self.config.v_head_dim * self.config.num_attention_heads + + self.q_head_dim = self.config.v_head_dim + + self.key_hidden_size = self.q_head_dim + self.val_hidden_size = self.config.v_head_dim + + self.recompute_up_proj = ( + self.config.recompute_granularity == 'selective' + and "mla_up_proj" in self.config.recompute_modules + ) + self.qkv_up_checkpoint = None + + self.softmax_scale = None + + if is_mtp_layer: + layer_idx = self.config.num_layers + layer_number - 1 + compress_ratio = self.config.csa_compress_ratios[layer_idx] + else: + compress_ratio = self.config.csa_compress_ratios[layer_number - 1] + use_compressed_yarn = compress_ratio > 1 + rope_base = ( + self.config.csa_compress_rotary_base if use_compressed_yarn else self.config.rotary_base + ) + self._dsv4_compress_ratio = compress_ratio + self._dsv4_rope_base = rope_base + self._dsv4_uses_yarn_rope = use_compressed_yarn + if not use_compressed_yarn: + self.rotary_pos_emb = RotaryEmbedding( + self.config.qk_pos_emb_head_dim, + rotary_percent=self.config.rotary_percent, + rotary_base=rope_base, + cp_group=self.pg_collection.cp, + ) + else: + self.rotary_pos_emb = YarnRotaryEmbedding( + self.config.qk_pos_emb_head_dim, + rotary_base=rope_base, + scaling_factor=self.config.rotary_scaling_factor, + original_max_position_embeddings=self.config.original_max_position_embeddings, + beta_fast=self.config.beta_fast, + beta_slow=self.config.beta_slow, + mscale=self.config.mscale, + mscale_all_dim=self.config.mscale_all_dim, + cp_group=self.pg_collection.cp, + ) + + core_attn_extra_kwargs = { + "rotary_pos_emb": self.rotary_pos_emb, + "compress_ratio": compress_ratio, + "is_mtp_layer": is_mtp_layer, + } + self.core_attention = build_module( + submodules.core_attention, + config=self.config, + layer_number=self.layer_number, + attn_mask_type=self.attn_mask_type, + attention_type=self.attention_type, + softmax_scale=self.softmax_scale, + k_channels=self.q_head_dim, + v_channels=self.config.v_head_dim, + cp_comm_type=cp_comm_type, + pg_collection=self.pg_collection, + **core_attn_extra_kwargs, + ) + + # Output. + self.o_local_groups = self.config.o_groups + assert ( + self.query_projection_size % self.config.o_groups == 0 + ), "num_attention_heads * v_head_dim must be divisible by o_groups" + group_proj_in_size = self.query_projection_size // self.config.o_groups + group_proj_out_size = self.config.o_groups * self.config.o_lora_rank + + _linear_o_group_proj = torch.empty( + group_proj_out_size, + group_proj_in_size, + device=torch.cuda.current_device(), + dtype=self.config.params_dtype, + ) + self.config.init_method(_linear_o_group_proj) + self.linear_o_group_proj = torch.nn.Parameter(_linear_o_group_proj) + + linear_proj_in_size = self.config.o_groups * self.config.o_lora_rank + + self.linear_proj = build_module( + submodules.linear_proj, + linear_proj_in_size, + self.config.hidden_size, + config=self.config, + init_method=self.config.output_layer_init_method, + bias=self.config.add_bias_linear, + input_is_parallel=True, + skip_bias_add=True, + is_expert=False, + tp_comm_buffer_name='proj', + tp_group=self.pg_collection.tp, + ) + + if ( + HAVE_TE + and isinstance(self.linear_proj, TELinear) + and ( + ( + self.config.fp8 + and self.config.fp8_recipe != 'delayed' + and is_te_min_version("2.6.0dev0") + ) + or (self.config.fp4 and is_te_min_version("2.7.0.dev0")) + ) + ): + # For fp8/fp4 training, the output of the fused core_attn is saved by itself, and + # linear_proj also saves the quantized tensor of this output. Here we set the + # linear_proj to save the original input tensors to avoid the extra memory usage of + # the quantized tensor. + set_save_original_input(self.linear_proj) + + def forward( + self, + hidden_states, + attention_mask, + key_value_states=None, + inference_context=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + rotary_pos_cos_sin=None, + attention_bias=None, + packed_seq_params=None, + position_ids=None, + sequence_len_offset=None, + *, + inference_params=None, + ): + """Forward pass for DeepSeek-v4 Hybrid Attention""" + assert ( + rotary_pos_emb is None + ), "Rotary position embeddings should not be passed into DSv4HybridAttention." + assert ( + attention_bias is None + ), "Attention bias should not be passed into DSv4HybridAttention." + assert ( + rotary_pos_cos is None and rotary_pos_sin is None + ), "DSv4HybridAttention does not support Flash Decoding" + assert ( + not rotary_pos_cos_sin + ), "Flash-infer rope has not been tested with DSv4HybridAttention." + assert ( + inference_context is None and inference_params is None + ), "Inference is not supported for DSv4HybridAttention." + assert ( + packed_seq_params is None + ), "Packed sequence is not supported for DSv4HybridAttention." + + # ===================== + # Query, Key, and Value + # ===================== + # Get the query, key and value tensors based on the type of attention - + # self or cross attn. + query, key, value, q_compressed, kv_compressed = self.get_query_key_value_tensors( + hidden_states, key_value_states, position_ids, None, inference_context=inference_context + ) + + # TODO: Currently, TE can only accept contiguous tensors for MLA + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + # ================================== + # core attention computation + # ================================== + # Need corresponding TE change + core_attn_manager = off_interface( + self.offload_core_attention and self.training, query, "core_attn" + ) + with core_attn_manager as query: + core_attn_out = self.core_attention( + query, + key, + value, + attention_mask, + packed_seq_params=None, + x=hidden_states, + qr=q_compressed, + ) + core_attn_out = core_attn_manager.group_offload( + core_attn_out, forced_released_tensors=[query, key, value] + ) + + if self.recompute_up_proj: + assert self.qkv_up_checkpoint is not None + self.qkv_up_checkpoint.discard_output_and_register_recompute(core_attn_out) + self.qkv_up_checkpoint = None + + # inverse RoPE on last qk_pos_emb_head_dim of each head + seq_len = core_attn_out.size(0) + n_heads = self.num_attention_heads_per_partition + pos_dim = self.config.qk_pos_emb_head_dim + nope_dim = self.config.v_head_dim - pos_dim + core_attn_out = core_attn_out.view(seq_len, core_attn_out.size(1), n_heads, -1) + rope_seqlen = seq_len + # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's + # concentration factor (mscale) is NOT part of the DSv4 model contract -- + # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0. + mscale = 1.0 + rotary_pos_cos = None + rotary_pos_sin = None + if self.config.apply_rope_fusion: + # ``mscale=1.0`` strips yarn's concentration factor from the + # cached cos/sin so the fused kernel matches the unfused + # path's forced ``mscale=1.0`` (DSv4 "pure rotation"). + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + rope_seqlen, dtype=hidden_states.dtype, packed_seq=False, mscale=mscale + ) + rotary_pos_emb = None + assert inference_context is None, "Inference with MLA RoPE fusion is not supported" + assert ( + fused_mla_rope_inplace is not None + ), "Fused MLA RoPE apply is not imported successfully" + elif self._dsv4_uses_yarn_rope: + rotary_pos_emb, _ = self.rotary_pos_emb(rope_seqlen, packed_seq=False) + else: + rotary_pos_emb = self.rotary_pos_emb(rope_seqlen, packed_seq=False) + if self.config.apply_rope_fusion: + # Fused DSA backward retains the raw attention output O. Applying + # inverse RoPE to its view in-place corrupts the retained O used by + # the softmax backward, so this call needs private storage. + assert fused_mla_rope_out_of_place is not None + core_attn_out = fused_mla_rope_out_of_place( + core_attn_out, + rotary_pos_cos, + rotary_pos_sin, + nope_dim, + pos_dim, + None, + self.pg_collection.cp.rank(), + self.pg_collection.cp.size(), + inverse=True, + remove_interleaving=True, + ) + else: + content_part, rot_part = torch.split( + core_attn_out, [core_attn_out.size(-1) - pos_dim, pos_dim], dim=-1 + ) + rot_part = apply_rotary_pos_emb( + rot_part, + rotary_pos_emb, + self.config, + cu_seqlens=None, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=True, + inverse=True, + mla_output_remove_interleaving=True, + ) + core_attn_out = torch.cat([content_part, rot_part], dim=-1) + core_attn_out = core_attn_out.view(seq_len, core_attn_out.size(1), -1) + + # Grouped output + core_attn_out = core_attn_out.view( + core_attn_out.size(0), core_attn_out.size(1), self.o_local_groups, -1 + ) + wo_a_weight = self.linear_o_group_proj.view( + self.o_local_groups, self.config.o_lora_rank, -1 + ) + core_attn_out = torch.einsum("...gd,grd->...gr", core_attn_out, wo_a_weight) + core_attn_out = core_attn_out.reshape(*core_attn_out.shape[:-2], -1) + + # ================= + # Output. [sq, b, h] + # ================= + attn_proj_manager = off_interface(self.offload_attn_proj, core_attn_out, "attn_proj") + with attn_proj_manager as core_attn_out: + output, bias = self.linear_proj(core_attn_out) + output = attn_proj_manager.group_offload(output, forced_released_tensors=[core_attn_out]) + + return output, bias + + +class DSv4HybridSelfAttention(DSv4HybridAttention): + """DSv4Hybrid Self-attention layer class + + Self-attention layer takes input with size [s, b, h] + and returns output of the same size. + """ + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSv4HybridSelfAttentionSubmodules, + layer_number: int, + attn_mask_type=AttnMaskType.padding, + cp_comm_type: Optional[str] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, + name: str | None = None, + ): + if pg_collection is None: + # Compatibility fallback for callers not yet passing process groups explicitly. + pg_collection = ProcessGroupCollection.use_mpu_process_groups() + + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + attn_mask_type=attn_mask_type, + attention_type="self", + cp_comm_type=cp_comm_type, + pg_collection=pg_collection, + pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, + name=name, + ) + + q_down_proj_kwargs = {} + if submodules.linear_q_down_proj in [TELinear]: + q_down_proj_kwargs['parallel_mode'] = 'duplicated' + else: + raise ValueError(f"Unsupported linear_q_down_proj: {submodules.linear_q_down_proj}") + + self.linear_q_down_proj = build_module( + submodules.linear_q_down_proj, + self.config.hidden_size, + self.config.q_lora_rank, + config=self.config, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name='q_down_proj', + skip_weight_param_allocation=False, + tp_group=None, + name=(name + ".linear_q_down_proj") if name is not None else None, + **q_down_proj_kwargs, + ) + + self.linear_q_up_proj = build_module( + submodules.linear_q_up_proj, + self.config.q_lora_rank, + self.config.num_attention_heads * self.q_head_dim, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name='q_up_proj', + tp_group=pg_collection.tp, + name=(name + ".linear_q_up_proj") if name is not None else None, + ) + + self.linear_kv_proj = build_module( + submodules.linear_kv_proj, + self.config.hidden_size, + self.config.v_head_dim, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name='kv_up_proj', + tp_group=pg_collection.tp, + name=(name + ".linear_kv_proj") if name is not None else None, + ) + self.kv_layernorm = submodules.kv_layernorm( + hidden_size=self.config.v_head_dim, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + + self.q_layernorm = submodules.q_layernorm( + hidden_size=self.config.q_lora_rank, + config=self.config, + eps=self.config.layernorm_epsilon, + ) + + def get_query_key_value_tensors( + self, + hidden_states, + key_value_states=None, + position_ids=None, + packed_seq_params=None, + inference_context=None, + *, + inference_params=None, + ): + """ + Derives `query`, `key` and `value` tensors from `hidden_states`. + """ + # s = sequence length, b = batch size, h = hidden size, n = num attention heads + # Attention heads [s, b, n*h] + assert ( + hidden_states.ndim == 3 + ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D" + assert ( + packed_seq_params is None + ), "Packed sequence is not supported for DSv4HybridAttention." + + assert ( + inference_context is None and inference_params is None + ), "Inference is not supported for DSv4HybridSelfAttention." + + # ========================================= + # Prepare RoPE and seqlen related params + # ========================================= + rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( + inference_context, None, hidden_states, self.config, None + ) + + # rotary_pos_emb:[s, b, 1, 64] + # DSv4 reference (DS-Inf) RoPE is pure rotation (norm-preserving). Yarn's + # concentration factor (mscale) is NOT part of the DSv4 model contract -- + # the model relies on Q/KV RMS-norm + unit-magnitude rotation. Force 1.0. + mscale = 1.0 + rotary_pos_cos = None + rotary_pos_sin = None + if self.config.apply_rope_fusion: + # ``mscale=1.0`` strips yarn's concentration factor from the + # cached cos/sin so the fused kernel matches the unfused + # path's forced ``mscale=1.0`` (DSv4 "pure rotation"). + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + rotary_seq_len, dtype=hidden_states.dtype, packed_seq=False, mscale=mscale + ) + rotary_pos_emb = None + assert inference_context is None, "Inference with MLA RoPE fusion is not supported" + assert ( + fused_mla_rope_inplace is not None + ), "Fused MLA RoPE apply is not imported successfully" + elif self._dsv4_uses_yarn_rope: + rotary_pos_emb, _ = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) + else: + rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=False) + + # ========================================= + # QKV down projection and layernorm + # ========================================= + # q_compressed: [s, b, q_lora_rank] + q_compressed, _ = self.linear_q_down_proj(hidden_states) + + kv_compressed = hidden_states + k_pos_emb = None + + # ========================================= + # Apply norm + # ========================================= + + if self.config.q_lora_rank is not None: + # q_compressed: [num_tokens, q_lora_rank] + q_compressed = apply_module(self.q_layernorm)(q_compressed) + + # ========================================= + # QKV up projection and RoPE apply + # ========================================= + + def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb): + """Apply the up projection and RoPE to the SBHD query and key.""" + # q_compressed: [s, b, q_lora_rank] + # q: [s, b, n * (qk_head_dim + qk_pos_emb_head_dim)] + q, _ = self.linear_q_up_proj(q_compressed) + + # q: [num_tokens, n, q_head_dim] + q = q.view(*q.size()[:-1], self.num_attention_heads_per_partition, self.q_head_dim) + q = _q_rms_norm(q, self.config.layernorm_epsilon) + + kv, _ = self.linear_kv_proj(kv_compressed) + kv = self.kv_layernorm(kv) + + # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] + if k_pos_emb is not None: + k_pos_emb = torch.unsqueeze(k_pos_emb, -2) + + if self.config.apply_rope_fusion: + cp_rank = self.pg_collection.cp.rank() + cp_size = self.pg_collection.cp.size() + query = fused_mla_rope_inplace( + q, + rotary_pos_cos, + rotary_pos_sin, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + None, + cp_rank, + cp_size, + remove_interleaving=True, + ) + kv = kv.unsqueeze(-2) + kv = fused_mla_rope_inplace( + kv, + rotary_pos_cos, + rotary_pos_sin, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + None, + cp_rank, + cp_size, + remove_interleaving=True, + ) + key = kv + value = kv + else: + q_len = q.size()[0] + # Keep direct forward calls with shorter sequences aligned to their inputs. + rotary_pos_emb = rotary_pos_emb[0:q_len] + + # q_no_pe: [num_tokens, n, qk_head_dim] + # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] + q_no_pe, q_pos_emb = torch.split( + q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1 + ) + + # RoPE and query (shared for wkv and latent) + # q_pos_emb: [num_tokens, n, qk_pos_emb_head_dim] + q_pos_emb = apply_rotary_pos_emb( + q_pos_emb, + rotary_pos_emb, + config=self.config, + cu_seqlens=None, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + ) + # query: [num_tokens, n, (qk_head_dim + v_head_dim)] + query = torch.cat([q_no_pe, q_pos_emb], dim=-1) + + pos_dim = self.config.qk_pos_emb_head_dim + kv_no_pe, k_pos_emb = torch.split(kv, [kv.size(-1) - pos_dim, pos_dim], dim=-1) + + # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] + k_pos_emb = apply_rotary_pos_emb( + k_pos_emb, + rotary_pos_emb, + config=self.config, + cu_seqlens=None, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + ) + + # Single head: key = value = [num_tokens, 1, v_head_dim] + kv = torch.cat([kv_no_pe, k_pos_emb], dim=-1).unsqueeze(-2) + key = kv + value = kv + + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + return query, key, value + + if self.recompute_up_proj: + quantization = self.config.fp8 or self.config.fp4 + self.qkv_up_checkpoint = tensor_parallel.CheckpointWithoutOutput(fp8=quantization) + query, key, value = self.qkv_up_checkpoint.checkpoint( + qkv_up_proj_and_rope_apply, q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) + else: + query, key, value = qkv_up_proj_and_rope_apply( + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) + + return query, key, value, q_compressed, kv_compressed + + def backward_dw(self) -> NoReturn: + """Execute weight gradient computation""" + self._backward_kv_proj() + self._backward_q_proj() + self._backward_output_proj() + + def _backward_kv_proj(self): + """Computes weight gradients of KV projection layers""" + self.linear_kv_proj.backward_dw() + + def _backward_q_proj(self): + """Computes weight gradients of Q projection layers""" + self.linear_q_down_proj.backward_dw() + self.linear_q_up_proj.backward_dw() + + def _backward_output_proj(self): + """Computes weight gradients of output projection layer""" + self.linear_proj.backward_dw() + + def set_for_recompute_input_layernorm(self): + """Set the attention layer for recompute input_layernorm. Only needed for fp8/fp4.""" + set_save_original_input(self.linear_q_down_proj) + set_save_original_input(self.linear_kv_proj) diff --git a/megatron/core/transformer/experimental_attention_variant/dsv4_module_specs.py b/megatron/core/transformer/experimental_attention_variant/dsv4_module_specs.py new file mode 100644 index 00000000000..e142413283e --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/dsv4_module_specs.py @@ -0,0 +1,95 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +"""Backend-neutral module specs for DeepSeek-V4 compressed sparse attention.""" + +from typing import Protocol + +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, + Compressor, + CompressorSubmodules, + CSAIndexer, + CSAIndexerSubmodules, +) +from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + DSv4HybridSelfAttentionSubmodules, +) +from megatron.core.transformer.identity_op import IdentityOp +from megatron.core.transformer.spec_utils import ModuleSpec +from megatron.core.transformer.transformer_config import TransformerConfig + + +class DSv4BackendSpecProvider(Protocol): + """Minimal backend interface needed to build the DSv4 attention spec.""" + + def linear(self) -> type: + """Return the backend's non-parallel linear module.""" + ... + + def column_parallel_linear(self) -> type: + """Return the backend's column-parallel linear module.""" + ... + + def row_parallel_linear(self) -> type: + """Return the backend's row-parallel linear module.""" + ... + + def layer_norm( + self, rms_norm: bool = False, for_qk: bool = False, has_residual: bool = False + ) -> type: + """Return the backend's normalization module.""" + ... + + +def get_dsv4_hybrid_module_spec_for_backend( + config: TransformerConfig, backend: DSv4BackendSpecProvider +) -> ModuleSpec: + """Build a DSv4 compressed sparse-attention spec for an explicit backend.""" + assert config.multi_latent_attention, "Currently only MLA supports sparse attention." + assert config.qk_l2_norm is False, "qk_l2_norm is not supported with MLA." + + rms_norm = config.normalization == "RMSNorm" + qk_norm = ( + backend.layer_norm(rms_norm=rms_norm, for_qk=True) if config.qk_layernorm else IdentityOp + ) + + compressor_spec = ModuleSpec( + module=Compressor, + submodules=CompressorSubmodules( + linear_wkv=backend.linear(), + linear_wgate=backend.linear(), + norm=backend.layer_norm(rms_norm=True, for_qk=False), + ), + ) + indexer_spec = ModuleSpec( + module=CSAIndexer, + submodules=CSAIndexerSubmodules( + linear_wq_b=backend.linear(), + linear_weights_proj=backend.linear(), + compressor=compressor_spec, + ), + ) + core_attention = ModuleSpec( + module=CompressedSparseAttention, + submodules=CompressedSparseAttentionSubmodules( + compressor=compressor_spec, indexer=indexer_spec + ), + ) + + return ModuleSpec( + module=DSv4HybridSelfAttention, + params={"attn_mask_type": AttnMaskType.causal}, + submodules=DSv4HybridSelfAttentionSubmodules( + linear_q_down_proj=backend.linear(), + linear_q_up_proj=backend.column_parallel_linear(), + linear_kv_proj=backend.column_parallel_linear(), + core_attention=core_attention, + linear_proj=backend.row_parallel_linear(), + q_layernorm=qk_norm, + kv_layernorm=qk_norm, + ), + metainfo={"fuse_input_layernorm": False}, + ) diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 2e19fefada6..be8b7bf9a7a 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -143,6 +143,7 @@ def __init__( cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ) -> None: # TODO(nschank): Restructure so that the Attention initializer knows which specific @@ -155,6 +156,7 @@ def __init__( attn_mask_type=attn_mask_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) self.config: MLATransformerConfig @@ -482,6 +484,7 @@ def __init__( cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, name: str | None = None, ): if pg_collection is None: @@ -496,6 +499,7 @@ def __init__( cp_comm_type=cp_comm_type, pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, + is_mtp_layer=is_mtp_layer, name=name, ) @@ -1232,6 +1236,7 @@ def __init__( attn_mask_type=AttnMaskType.padding, cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, + is_mtp_layer: bool = False, pp_layer_offset: Optional[int] = None, name: str | None = None, ): @@ -1247,6 +1252,7 @@ def __init__( attention_type="self", cp_comm_type=cp_comm_type, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, pp_layer_offset=pp_layer_offset, name=name, ) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index ee5870ba3ee..5d3d12c3486 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -281,8 +281,10 @@ class TransformerConfig(ModelParallelConfig): #################### # attention variant #################### - experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa']] = None - """Type of attention variant to use. Currently support gated_delta_net and dsa.""" + experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa', 'dsv4_hybrid']] = ( + None + ) + """Type of attention variant to use. Currently support gated_delta_net, dsa, and dsv4_hybrid.""" experimental_attention_variant_loss_scale_func: Optional[Callable[[torch.Tensor], None]] = None """Optional hook for experimental attention variants to receive the main loss scale.""" @@ -1392,6 +1394,26 @@ def __post_init__(self): "dsa_indexer_skip_topk_offset must be non-negative, got " f"{self.dsa_indexer_skip_topk_offset}." ) + elif self.experimental_attention_variant == "dsv4_hybrid": + assert self.multi_latent_attention, "DSv4 Hybrid requires multi_latent_attention." + assert self.csa_compress_ratios is not None, "csa_compress_ratios must be set" + mtp_layers = self.mtp_num_layers or 0 + expected_len = self.num_layers + mtp_layers + assert len(self.csa_compress_ratios) == expected_len, ( + f"csa_compress_ratios length ({len(self.csa_compress_ratios)}) must equal " + f"num_layers + mtp_num_layers ({self.num_layers} + {mtp_layers} = {expected_len})" + ) + assert all( + ratio in [0, 4, 128] for ratio in self.csa_compress_ratios + ), "csa_compress_ratios must be 0, 4, or 128" + assert ( + self.tensor_model_parallel_size == 1 + ), "DSv4 Hybrid Attention only supports TP size 1." + assert ( + self.context_parallel_size == 1 + ), "DSv4 Hybrid Attention does not support context parallelism yet." + assert not self.qk_clip, "QK clipping is not supported with DSv4 Hybrid Attention." + self.hetereogenous_dist_checkpoint = True if self.fp8: # cannot support first last layer bf16 with delayed scaling @@ -2853,10 +2875,12 @@ class MLATransformerConfig(TransformerConfig): """Rank of Query tensor's low rank representation.""" kv_lora_rank: int = 512 - """Rank of Key and Value tensors' low rank representation.""" + """Rank of Key and Value tensors' low rank representation. + This is not used for DSv4 Hybrid Attention and will be overridden automatically.""" qk_head_dim: int = 128 - """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim""" + """Dimension of the head in the QK projection. q_head_dim = qk_head_dim + qk_pos_emb_head_dim + This is not used for DSv4 Hybrid Attention and will be overridden automatically.""" qk_pos_emb_head_dim: int = 64 """Dimension of the position embedding in the QK projection.""" @@ -2894,9 +2918,15 @@ class MLATransformerConfig(TransformerConfig): mscale_all_dim: float = 0.0 """Mscale all dimensions for YaRN RoPE in Multi-Latent Attention, used by yarn.""" + o_groups: int = 8 + """Number of groups for grouped low-rank output projection (wo_a).""" + + o_lora_rank: int = 1024 + """Low-rank dimension per group for grouped output (wo_a). Used when o_groups > 0.""" + cache_mla_latents: bool = False """Cache the low dimensional tensors for MLA rather than full KV cache. - This is only for the dynamic inference backend and requires that + This is only for the dynamic inference backend and requires that Flash MLA is installed.""" mla_down_proj_fusion: bool = False @@ -2906,12 +2936,41 @@ class MLATransformerConfig(TransformerConfig): def __post_init__(self): super().__post_init__() - if self.multi_latent_attention and self.apply_rope_fusion and self.rope_type != "yarn": + if ( + self.multi_latent_attention + and self.apply_rope_fusion + and self.rope_type != "yarn" + and self.experimental_attention_variant != "dsv4_hybrid" + ): raise ValueError("apply_rope_fusion for MLA only works with YARN RoPE.") if self.attention_output_gate: raise NotImplementedError("Output gate is not supported for MLA yet.") + # DSv4 hybrid: derive qk_head_dim and kv_lora_rank from v_head_dim and qk_pos_emb_head_dim + if self.experimental_attention_variant == "dsv4_hybrid": + assert ( + not self.mla_down_proj_fusion + ), "MLA down projection fusion must be disabled for DSv4 hybrid mode." + assert self.q_lora_rank is not None, "DSv4 hybrid mode requires q_lora_rank." + assert self.o_groups > 0, "DSv4 hybrid mode requires o_groups to be positive." + assert self.o_lora_rank > 0, "DSv4 hybrid mode requires o_lora_rank to be positive." + assert ( + self.num_attention_heads * self.v_head_dim + ) % self.o_groups == 0, ( + "num_attention_heads * v_head_dim must be divisible by o_groups." + ) + log_single_rank( + logger, + logging.WARNING, + "DSv4 hybrid mode is enabled, deriving qk_head_dim and kv_lora_rank from " + "v_head_dim and qk_pos_emb_head_dim", + ) + derived = self.v_head_dim - self.qk_pos_emb_head_dim + assert derived > 0, "v_head_dim must be greater than qk_pos_emb_head_dim." + self.qk_head_dim = derived + self.kv_lora_rank = derived + if self.cache_mla_latents: assert ( self.apply_rope_fusion is False diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index f6ea382077e..d0de05c9eb1 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -368,6 +368,8 @@ def __init__( attention_optional_kwargs["pg_collection"] = pg_collection if pp_layer_offset is not None: attention_optional_kwargs["pp_layer_offset"] = pp_layer_offset + if is_mtp_layer: + attention_optional_kwargs["is_mtp_layer"] = True # [Module 2: SelfAttention] self.self_attention = build_module( diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 3929b69d606..a7f631f74ae 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -323,6 +323,15 @@ def no_rope_freq_type(x): # it's a single int but in str return int(x) + +def compress_ratios_type(x): + """Parse per-layer compression ratios for compressed sparse attention.""" + if isinstance(x, list): + return x + assert isinstance(x, str) + return _eval_pattern(x) + + def moe_freq_type(x): """Frequency between MoE layers and Dense layers. @@ -2071,6 +2080,7 @@ def _add_network_size_args(parser): "no_rope_freq", "moe_layer_freq", "linear_attention_freq", + "csa_compress_ratios", "moe_router_load_balancing_type", "moe_aux_loss_coeff", "cp_comm_type", @@ -3245,6 +3255,11 @@ def _add_mla_args(parser): help="Mscale for YaRN RoPE in multi-latent attention.") group.add_argument('--mscale-all-dim', type=float, default=0.0, help="Mscale all dimensions for YaRN RoPE in multi-latent attention.") + group.add_argument('--o-groups', type=int, default=8, + help="Number of groups for grouped low-rank output projection (wo_a).") + group.add_argument('--o-lora-rank', type=int, default=1024, + help="Low-rank dimension per group for grouped output (wo_a). " + "Used when o-groups > 0.") group.add_argument('--cache-mla-latents', action='store_true', default=False, help="If set caches the mla down projected latents with mla flash decode.") group.add_argument( @@ -3269,6 +3284,15 @@ def _add_experimental_attention_variant_args(parser): 'where 1 indicates an LA layer and 0 indicates a SDPA layer. ' 'Examples: "([0]+[1]*23)": 1 SDPA layer followed by 23 LA layers, ' '"([1]*3+[0]*2)*2": Three LA layers followed by two SDPA layers, repeated twice.') + group.add_argument( + '--csa-compress-ratios', + type=compress_ratios_type, + default=None, + help='Per-layer compress ratios for compressed sparse attention. ' + 'Accepts a Python list expression such as "[0,0,4,128,4,128]" or ' + '"([0]+[4,128]*2)*3". Valid values are 0, 4, and 128, and the ' + 'list length must equal num_layers plus mtp_num_layers.', + ) return parser def _add_heterogeneous_args(parser): diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py new file mode 100644 index 00000000000..47ab8f98650 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py @@ -0,0 +1,577 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from unittest.mock import patch + +import pytest +import torch + +from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.transformer_config import MLATransformerConfig +from tests.unit_tests.test_utilities import Utils + +try: + from fast_hadamard_transform import hadamard_transform as _hadamard_transform + + HAVE_HADAMARD = True +except ImportError: + HAVE_HADAMARD = False + _hadamard_transform = None + +_SEED = 42 + + +def _mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + return x * scale + + +@pytest.fixture(autouse=True) +def patch_hadamard_if_needed(): + """Patch hadamard_transform in dsa/csa modules if the library is not installed.""" + if not HAVE_HADAMARD: + with ( + patch( + 'megatron.core.transformer.experimental_attention_variant.dsa.hadamard_transform', + _mock_hadamard_transform, + ), + patch( + 'megatron.core.transformer.experimental_attention_variant.csa.rotate_activation', + lambda x: x * (x.size(-1) ** -0.5), + ), + ): + yield + else: + yield + + +# --------------------------------------------------------------------------- +# Config / spec helpers +# --------------------------------------------------------------------------- + + +def _make_config( + num_layers=4, + hidden_size=256, + num_attention_heads=16, + v_head_dim=64, + qk_pos_emb_head_dim=32, + q_lora_rank=64, + o_groups=8, + o_lora_rank=64, + csa_compress_ratios=None, + csa_window_size=8, + tensor_model_parallel_size=1, + sequence_parallel=False, + dsa_indexer_n_heads=8, + dsa_indexer_head_dim=64, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=0.0, + **extra_config_kwargs, +): + """Create an MLATransformerConfig for DSv4 hybrid attention tests.""" + if csa_compress_ratios is None: + csa_compress_ratios = [0, 4, 128, 4] + return MLATransformerConfig( + num_layers=num_layers, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + add_bias_linear=False, + tensor_model_parallel_size=tensor_model_parallel_size, + sequence_parallel=sequence_parallel, + q_lora_rank=q_lora_rank, + kv_lora_rank=v_head_dim - qk_pos_emb_head_dim, + qk_head_dim=v_head_dim - qk_pos_emb_head_dim, + qk_pos_emb_head_dim=qk_pos_emb_head_dim, + v_head_dim=v_head_dim, + o_groups=o_groups, + o_lora_rank=o_lora_rank, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + multi_latent_attention=True, + experimental_attention_variant='dsv4_hybrid', + csa_compress_ratios=csa_compress_ratios, + csa_window_size=csa_window_size, + dsa_indexer_n_heads=dsa_indexer_n_heads, + dsa_indexer_head_dim=dsa_indexer_head_dim, + dsa_indexer_topk=dsa_indexer_topk, + dsa_indexer_loss_coeff=dsa_indexer_loss_coeff, + **extra_config_kwargs, + ) + + +def _make_attention_spec(config): + """Build the full DSv4HybridSelfAttention ModuleSpec using the canonical spec builder.""" + from megatron.core.extensions.transformer_engine_spec_provider import TESpecProvider + from megatron.core.transformer.experimental_attention_variant.dsv4_module_specs import ( + get_dsv4_hybrid_module_spec_for_backend, + ) + + return get_dsv4_hybrid_module_spec_for_backend(config=config, backend=TESpecProvider()) + + +def test_module_spec_is_built_from_explicit_backend(): + """The neutral spec builder should use only its explicitly supplied backend.""" + from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + Compressor, + CSAIndexer, + ) + from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + ) + from megatron.core.transformer.experimental_attention_variant.dsv4_module_specs import ( + get_dsv4_hybrid_module_spec_for_backend, + ) + + class Linear: + pass + + class ColumnParallelLinear: + pass + + class RowParallelLinear: + pass + + class Norm: + pass + + class Backend: + def linear(self): + return Linear + + def column_parallel_linear(self): + return ColumnParallelLinear + + def row_parallel_linear(self): + return RowParallelLinear + + def layer_norm(self, rms_norm=False, for_qk=False, has_residual=False): + return Norm + + spec = get_dsv4_hybrid_module_spec_for_backend(_make_config(), Backend()) + + assert spec.module is DSv4HybridSelfAttention + assert spec.submodules.linear_q_down_proj is Linear + assert spec.submodules.linear_q_up_proj is ColumnParallelLinear + assert spec.submodules.linear_kv_proj is ColumnParallelLinear + assert spec.submodules.linear_proj is RowParallelLinear + assert spec.submodules.core_attention.module is CompressedSparseAttention + assert spec.submodules.core_attention.submodules.compressor.module is Compressor + assert spec.submodules.core_attention.submodules.indexer.module is CSAIndexer + + +def test_config_includes_mtp_ratio_and_derives_dimensions(): + """DSv4 config should account for MTP and derive its shared Q/KV content width.""" + config = _make_config(num_layers=2, mtp_num_layers=1, csa_compress_ratios=[0, 4, 128]) + + expected_content_dim = config.v_head_dim - config.qk_pos_emb_head_dim + assert config.qk_head_dim == expected_content_dim + assert config.kv_lora_rank == expected_content_dim + assert config.hetereogenous_dist_checkpoint is True + + +def test_config_rejects_context_parallelism(): + """The SBHD slice should fail early instead of silently accepting unsupported CP.""" + with pytest.raises(AssertionError, match="does not support context parallelism"): + _make_config(context_parallel_size=2) + + +def _build_attention(config, layer_number, pg_collection, **kwargs): + """Instantiate a DSv4HybridSelfAttention from config.""" + from megatron.core.transformer.spec_utils import build_module + + spec = _make_attention_spec(config) + return build_module( + spec, config=config, layer_number=layer_number, pg_collection=pg_collection, **kwargs + ) + + +# =========================================================================== +# Constructor tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridAttentionConstructor: + """Test construction of DSv4HybridSelfAttention in the supported TP=1 configuration.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + def test_basic_construction(self): + """Verify the layer builds and has the expected sub-modules.""" + from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + ) + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + config = _make_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=1, pg_collection=pg) + + assert isinstance(attn, DSv4HybridSelfAttention) + assert hasattr(attn, 'linear_q_down_proj') + assert hasattr(attn, 'linear_q_up_proj') + assert hasattr(attn, 'linear_kv_proj') + assert hasattr(attn, 'linear_proj') + assert hasattr(attn, 'linear_o_group_proj') + assert hasattr(attn, 'core_attention') + assert hasattr(attn, 'q_layernorm') + assert hasattr(attn, 'kv_layernorm') + + def test_q_head_dim_equals_v_head_dim(self): + """q_head_dim must equal v_head_dim for DSv4 hybrid.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + config = _make_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=1, pg_collection=pg) + + assert attn.q_head_dim == config.v_head_dim + + def test_current_main_constructor_kwargs(self): + """Current TransformerLayer forwards module names and pipeline offsets.""" + config = _make_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention( + config, + layer_number=1, + pg_collection=pg, + pp_layer_offset=0, + name="decoder.layers.0.self_attention", + ) + + assert attn._pp_layer_offset == 0 + + @pytest.mark.parametrize("layer_number", [1, 2, 3, 4]) + def test_rope_base_varies_with_compress_ratio(self, layer_number): + """Layers with compress_ratio > 1 should use csa_compress_rotary_base.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + ratios = [0, 4, 128, 4] + config = _make_config(csa_compress_ratios=ratios) + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=layer_number, pg_collection=pg) + + ratio = ratios[layer_number - 1] + if ratio > 1: + expected_base = config.csa_compress_rotary_base + else: + expected_base = config.rotary_base + + # inv_freq is derived from rotary_base; verify the correct base was used + dim = config.qk_pos_emb_head_dim + recomputed_inv_freq = 1.0 / ( + expected_base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) + ) + assert torch.allclose( + attn.rotary_pos_emb.inv_freq.cpu(), recomputed_inv_freq, rtol=1e-5, atol=1e-5 + ) + + +# =========================================================================== +# Forward / backward tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridAttentionForwardBackward: + """Test forward and backward passes of DSv4HybridSelfAttention.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.config = _make_config(dsa_indexer_loss_coeff=1.0) + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("layer_number", [1, 2, 3, 4]) + def test_forward_output_shape(self, layer_number): + """Forward should produce [sq, b, hidden_size] output.""" + seq_len = 256 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention( + self.config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + output, bias = attn(hidden_states=hidden, attention_mask=None) + + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + assert output.dtype == torch.bfloat16 + assert not torch.isnan(output).any() + + @pytest.mark.parametrize("layer_number", [1, 2]) + def test_backward_gradient_flow(self, layer_number): + """Backward should produce gradients for all trainable parameters.""" + seq_len = 256 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention( + self.config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + attn.train() + + hidden = ( + torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + + output, bias = attn(hidden_states=hidden, attention_mask=None) + loss = output.sum() + loss.backward() + + assert hidden.grad is not None, "No gradient on hidden_states" + for name, param in attn.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"No gradient for parameter {name}" + + def test_eval_mode(self): + """Forward should work in eval mode.""" + seq_len = 128 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + attn.eval() + + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + with torch.no_grad(): + output, bias = attn(hidden_states=hidden, attention_mask=None) + + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + assert not torch.isnan(output).any() + + def test_different_seq_lengths(self): + """Forward should handle various sequence lengths.""" + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=2, pg_collection=self.pg).cuda() + + for seq_len in [64, 128, 256]: + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + output, bias = attn(hidden_states=hidden, attention_mask=None) + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + + +# =========================================================================== +# get_query_key_value_tensors tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridQKV: + """Test get_query_key_value_tensors internals.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.config = _make_config() + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + def test_qkv_shapes(self): + """Query, key, value should have correct shapes.""" + seq_len = 64 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + q, k, v, q_compressed, kv_compressed = attn.get_query_key_value_tensors(hidden) + + n_heads = self.config.num_attention_heads + v_dim = self.config.v_head_dim + + assert q.shape == (seq_len, batch_size, n_heads, v_dim) + # key and value are single-head (MQA-style) with an extra head dim + assert k.shape[-1] == v_dim + assert v.shape[-1] == v_dim + + def test_key_equals_value(self): + """In the wkv path, key and value should be the same tensor.""" + seq_len = 64 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + q, k, v, _, _ = attn.get_query_key_value_tensors(hidden) + assert torch.equal(k, v), "key and value should be identical in wkv path" + + +# =========================================================================== +# Grouped output projection tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridGroupedOutput: + """Test that grouped output projection (wo_a) parameters are created.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + def test_o_group_proj_shape(self): + """linear_o_group_proj should have the correct shape.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + o_groups = 8 + o_lora_rank = 64 + config = _make_config(o_groups=o_groups, o_lora_rank=o_lora_rank) + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=1, pg_collection=pg) + + expected_out = o_groups * o_lora_rank + expected_in = (config.v_head_dim * config.num_attention_heads) // o_groups + assert attn.linear_o_group_proj.shape == (expected_out, expected_in) + assert attn.linear_o_group_proj.requires_grad + + +# =========================================================================== +# apply_rope_fusion tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridRopeFusion: + """Test that apply_rope_fusion=True works for both yarn and non-yarn layers. + + DSv4 Hybrid uses YarnRotaryEmbedding for layers with compress_ratio > 1 + and standard RotaryEmbedding for layers with compress_ratio <= 1. The + fused RoPE path must obtain cos/sin from both embedding classes via + get_cached_cos_sin. + + compress_ratios=[0, 4, 128, 4]: layer 1 has ratio 0 (standard + RotaryEmbedding), layers 2-4 have ratio > 1 (YarnRotaryEmbedding). + """ + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + def test_rope_fusion_forward_backward_parity(self): + """Fused RoPE forward/backward succeeds and matches the unfused path.""" + seq_len = 128 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + fused_config = _make_config(apply_rope_fusion=True) + attn_fused = _build_attention(fused_config, layer_number=4, pg_collection=self.pg).cuda() + attn_fused.train() + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + unfused_config = _make_config(apply_rope_fusion=False) + attn_unfused = _build_attention( + unfused_config, layer_number=4, pg_collection=self.pg + ).cuda() + attn_unfused.train() + + hidden = torch.randn( + seq_len, batch_size, fused_config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + out_fused, _ = attn_fused(hidden_states=hidden, attention_mask=None) + out_unfused, _ = attn_unfused(hidden_states=hidden, attention_mask=None) + + assert out_fused.shape == (seq_len, batch_size, fused_config.hidden_size) + assert torch.isfinite(out_fused).all() + # The remaining difference is bf16 accumulation order between the fused + # Triton kernel and eager PyTorch operations. + torch.testing.assert_close(out_fused, out_unfused, atol=3e-2, rtol=3e-2) + + hidden_fused = hidden.detach().clone().requires_grad_(True) + hidden_unfused = hidden.detach().clone().requires_grad_(True) + + attn_fused(hidden_states=hidden_fused, attention_mask=None)[0].sum().backward() + attn_unfused(hidden_states=hidden_unfused, attention_mask=None)[0].sum().backward() + + assert hidden_fused.grad is not None + for name, param in attn_fused.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"No gradient for parameter {name}"