From 3efbd82ebebcd518234f201ff5be7b84579add03 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Fri, 24 Apr 2026 12:51:18 +0000 Subject: [PATCH 01/16] init commit --- .../core/fusions/fused_mla_yarn_rope_apply.py | 197 ++-- .../models/common/embeddings/rope_utils.py | 29 + ...rimental_attention_variant_module_specs.py | 73 ++ megatron/core/models/gpt/gpt_model.py | 4 +- .../experimental_attention_variant/csa.py | 787 ++++++++++++++++ .../deepseek_v4_hybrid_attention.py | 886 ++++++++++++++++++ .../experimental_attention_variant/dsa.py | 130 ++- .../core/transformer/transformer_config.py | 46 +- megatron/training/arguments.py | 39 + .../fusions/test_mla_yarn_rope_apply.py | 39 +- 10 files changed, 2096 insertions(+), 134 deletions(-) create mode 100644 megatron/core/transformer/experimental_attention_variant/csa.py create mode 100644 megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py index 1fd5dcfae37..90acee5814b 100644 --- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py +++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py @@ -65,11 +65,11 @@ def _get_thd_token_idx(cu_seqlens, pid_m, seq_num, cp_rank, cp_size): restore_value=["Q"], ) @triton.jit -def rotary_fwd_q_kernel( +def _mla_rope_fwd_inplace_kernel( Q, COS, SIN, - qk_head_dim, + nope_dim, emb_dim: tl.constexpr, head_num: tl.constexpr, batch_size, @@ -77,17 +77,20 @@ def rotary_fwd_q_kernel( cu_seqlens_q, stride_x_seq, stride_x_nheads, + stride_cos_seq, + stride_sin_seq, cp_rank, cp_size, + INVERSE: tl.constexpr, BLOCK_H: tl.constexpr, ): """ - Triton kernel of the forward pass for applying YARN RoPE to MLA's query. - This kernel inplace modifies the input tensor Q. + Forward pass: apply RoPE inplace to the trailing emb_dim elements. + Reads from interleaved layout, writes back to interleaved layout. Input: - Q: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + Q: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] COS/SIN: [max_seq_len, emb_dim] batch_size: batch size for sbhd format, not used for thd format @@ -102,10 +105,13 @@ def rotary_fwd_q_kernel( else: token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = tl.load(COS + token_idx * stride_cos_seq + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * stride_sin_seq + tl.arange(0, emb_dim // 2)) + cos_right = tl.load(COS + token_idx * stride_cos_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * stride_sin_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + if INVERSE: + sin_left = -sin_left + sin_right = -sin_right cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) @@ -113,7 +119,7 @@ def rotary_fwd_q_kernel( Q = Q + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads - x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + qk_head_dim + x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim mask = x_off < head_num * stride_x_nheads # x1 = t[..., 0::2], x2 = t[..., 1::2] x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 @@ -124,10 +130,8 @@ def rotary_fwd_q_kernel( x_left = x_1 * cos_left - x_2 * sin_left x_right = x_2 * cos_right + x_1 * sin_right - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - tl.store(Q + x_left_off, x_left, mask=mask) - tl.store(Q + x_right_off, x_right, mask=mask) + tl.store(Q + x_1_off, x_left, mask=mask) + tl.store(Q + x_2_off, x_right, mask=mask) @triton.autotune( @@ -145,11 +149,11 @@ def rotary_fwd_q_kernel( restore_value=["DO"], ) @triton.jit -def rotary_bwd_q_kernel( +def _mla_rope_bwd_inplace_kernel( DO, COS, SIN, - qk_head_dim, + nope_dim, emb_dim: tl.constexpr, head_num: tl.constexpr, batch_size, @@ -157,17 +161,20 @@ def rotary_bwd_q_kernel( cu_seqlens_q, stride_x_seq, stride_x_nheads, + stride_cos_seq, + stride_sin_seq, cp_rank, cp_size, + INVERSE: tl.constexpr, BLOCK_H: tl.constexpr, ): """ - Triton kernel of the backward pass for applying YARN RoPE to MLA's query. - This kernel inplace modifies the input tensor DO. + Backward pass: inverse RoPE inplace on the trailing emb_dim elements. + Reads from interleaved layout, writes to interleaved layout. Input: - DO: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + DO: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] COS/SIN: [max_seq_len, emb_dim] batch_size, seq_num, and cu_seqlens_q are the same as in the forward pass @@ -180,10 +187,13 @@ def rotary_bwd_q_kernel( else: token_idx = _get_thd_token_idx(cu_seqlens_q, pid_m, seq_num, cp_rank, cp_size) - cos_left = tl.load(COS + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - sin_left = tl.load(SIN + token_idx * emb_dim + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * emb_dim + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_left = tl.load(COS + token_idx * stride_cos_seq + tl.arange(0, emb_dim // 2)) + sin_left = tl.load(SIN + token_idx * stride_sin_seq + tl.arange(0, emb_dim // 2)) + cos_right = tl.load(COS + token_idx * stride_cos_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + sin_right = tl.load(SIN + token_idx * stride_sin_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + if INVERSE: + sin_left = -sin_left + sin_right = -sin_right cos_left = cos_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) sin_left = sin_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) cos_right = cos_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) @@ -191,25 +201,24 @@ def rotary_bwd_q_kernel( DO = DO + pid_m * stride_x_seq + pid_head * BLOCK_H * stride_x_nheads - x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + qk_head_dim + x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim mask = x_off < head_num * stride_x_nheads - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - x_left = tl.load(DO + x_left_off, mask=mask) - x_right = tl.load(DO + x_right_off, mask=mask) + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_left = tl.load(DO + x_1_off, mask=mask) + x_right = tl.load(DO + x_2_off, mask=mask) x_1 = x_left * cos_left + x_right * sin_right x_2 = -x_left * sin_left + x_right * cos_right - x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 - x_2_off = x_1_off + 1 tl.store(DO + x_1_off, x_1, mask=mask) tl.store(DO + x_2_off, x_2, mask=mask) -class ApplyMLARotaryEmbQ(torch.autograd.Function): +class _FusedMLARoPEInplace(torch.autograd.Function): """ - Autograd function for applying YARN RoPE to MLA's query. + Autograd function for applying RoPE inplace to the trailing emb_dim + elements of a multi-head tensor (leaving the first nope_dim elements unchanged). """ @staticmethod @@ -218,22 +227,24 @@ def forward( q, cos, sin, - qk_head_dim, + nope_dim, emb_dim, cu_seqlens_q, cp_rank, cp_size, rotary_interleaved=False, + inverse=False, ): """ - Forward function for ApplyMLARotaryEmbQ. + Forward function for _FusedMLARoPEInplace. Args: - q: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + q: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] cos/sin: [max_seq_len, 1, 1, emb_dim] cu_seqlens_q: [seq_num + 1] accumulated sequence lengths for thd format rotary_interleaved: whether to apply RoPE interleaved, only supports False for now + inverse: if True, negate sin inside the kernel to apply the inverse rotation """ assert not rotary_interleaved max_seqlen = None @@ -249,17 +260,17 @@ def forward( total_seqlen, nheads, headdim = q.shape seq_num = len(cu_seqlens_q) - 1 assert q.stride(-1) == 1 - assert cos.is_contiguous() - assert sin.is_contiguous() - assert headdim == qk_head_dim + emb_dim + assert cos.stride(-1) == 1 + assert sin.stride(-1) == 1 + assert headdim == nope_dim + emb_dim assert emb_dim % 4 == 0 grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_fwd_q_kernel[grid]( + _mla_rope_fwd_inplace_kernel[grid]( q, cos, sin, - qk_head_dim, + nope_dim, emb_dim, nheads, batch_size, @@ -267,14 +278,18 @@ def forward( cu_seqlens_q, q.stride(0), q.stride(1), + cos.stride(0), + sin.stride(0), cp_rank, cp_size, + INVERSE=inverse, ) ctx.save_for_backward(cos, sin) - ctx.qk_head_dim = qk_head_dim + ctx.nope_dim = nope_dim ctx.emb_dim = emb_dim ctx.cu_seqlens_q = cu_seqlens_q ctx.rotary_interleaved = rotary_interleaved + ctx.inverse = inverse ctx.cp_rank = cp_rank ctx.cp_size = cp_size if cu_seqlens_q is None: @@ -284,11 +299,11 @@ def forward( @staticmethod def backward(ctx, grad): """ - Backward function for ApplyMLARotaryEmbQ. + Backward function for _FusedMLARoPEInplace. Args: - grad: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + grad: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] """ cos, sin = ctx.saved_tensors max_seqlen = None @@ -304,11 +319,11 @@ def backward(ctx, grad): assert grad.stride(-1) == 1 grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_bwd_q_kernel[grid]( + _mla_rope_bwd_inplace_kernel[grid]( grad, cos, sin, - ctx.qk_head_dim, + ctx.nope_dim, ctx.emb_dim, nheads, batch_size, @@ -316,49 +331,54 @@ def backward(ctx, grad): ctx.cu_seqlens_q, grad.stride(0), grad.stride(1), + cos.stride(0), + sin.stride(0), ctx.cp_rank, ctx.cp_size, + INVERSE=ctx.inverse, ) if ctx.cu_seqlens_q is None: grad = grad.view(max_seqlen, batch_size, nheads, headdim) - return grad, None, None, None, None, None, None, None, None + return grad, None, None, None, None, None, None, None, None, None -def fused_apply_mla_rope_for_q( +def fused_mla_rope_inplace( t: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, - qk_head_dim: int, + nope_dim: int, emb_dim: int, cu_seqlens_q: Optional[torch.Tensor] = None, cp_rank: int = 0, cp_size: int = 1, rotary_interleaved: bool = False, + inverse: bool = False, ): """ - Fused function for applying YARN RoPE to MLA's query. - This function inplace modifies the input tensor t. - Along the last dimension of t, the last emb_dim elements are applied with RoPE. - The first qk_head_dim elements are not modified. - It is an experimental feature and may change in future versions. + Fused RoPE applied inplace to the trailing emb_dim elements of a tensor, + leaving the first nope_dim elements unchanged. It supports both sbhd and thd input formats. + When ``inverse=True`` the rotation is reversed, which is useful for + undoing RoPE on the attention output. + For the notations below, seq_len is the length of the sequence per batch for sbhd format, total_seq_len is the total length of the sequences for thd format. max_seq_len is the maximum length of the sequences in the input tensor. Args: - t: [seq_len, batch_size, head_num, qk_head_dim + emb_dim] - or [total_seq_len, head_num, qk_head_dim + emb_dim] + t: [seq_len, batch_size, head_num, nope_dim + emb_dim] + or [total_seq_len, head_num, nope_dim + emb_dim] cos/sin: [max_seq_len, 1, 1, emb_dim] cu_seqlens_q: [seq_num + 1] accumulated sequence lengths for thd format rotary_interleaved: whether to apply RoPE interleaved, only supports False for now + inverse: if True, apply the inverse rotation Returns: t: inplace modified input tensor """ - return ApplyMLARotaryEmbQ.apply( - t, cos, sin, qk_head_dim, emb_dim, cu_seqlens_q, cp_rank, cp_size, rotary_interleaved + return _FusedMLARoPEInplace.apply( + t, cos, sin, nope_dim, emb_dim, cu_seqlens_q, cp_rank, cp_size, rotary_interleaved, inverse ) @@ -376,7 +396,7 @@ def fused_apply_mla_rope_for_q( key=["emb_dim", "k_dim", "v_dim", "head_num"], ) @triton.jit -def rotary_fwd_kv_kernel( +def _mla_rope_fwd_kv_split_kernel( KV, K_POS_EMB, O_KEY, @@ -402,9 +422,8 @@ def rotary_fwd_kv_kernel( BLOCK_H: tl.constexpr, ): """ - Triton kernel of the forward pass for applying YARN RoPE to MLA's key and value. - It splits the input tensor KV into key and value, - and concatenates the processed RoPE to the key. + Forward pass: split KV into key and value, apply RoPE to k_pos_emb, + and concatenate the result onto key. Input: KV: [seq_len, batch_size, head_num, k_dim + v_dim] @@ -460,14 +479,14 @@ def rotary_fwd_kv_kernel( x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - x_left_off = ( + x_1_off = ( tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + k_dim - + tl.arange(0, emb_dim // 2)[None, :] + + tl.arange(0, emb_dim // 2)[None, :] * 2 ) - x_right_off = x_left_off + emb_dim // 2 - tl.store(K_ptr + x_left_off, x_left, mask=mask) - tl.store(K_ptr + x_right_off, x_right, mask=mask) + x_2_off = x_1_off + 1 + tl.store(K_ptr + x_1_off, x_left, mask=mask) + tl.store(K_ptr + x_2_off, x_right, mask=mask) @triton.autotune( @@ -484,7 +503,7 @@ def rotary_fwd_kv_kernel( key=["emb_dim", "k_dim", "v_dim", "head_num"], ) @triton.jit -def rotary_bwd_kv_kernel( +def _mla_rope_bwd_kv_split_kernel( dK, dV, dKV, @@ -510,7 +529,7 @@ def rotary_bwd_kv_kernel( BLOCK_H: tl.constexpr, ): """ - Triton kernel of the backward pass for applying YARN RoPE to MLA's key and value. + Backward pass for the KV-split RoPE. Input: dK: [seq_len, batch_size, head_num, emb_dim + k_dim] @@ -555,10 +574,10 @@ def rotary_bwd_kv_kernel( dK_ptr = dK + pid_m * stride_dk_seq + i * BLOCK_H * stride_dk_nheads x_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + k_dim mask = x_off < head_num * stride_dk_nheads - x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] - x_right_off = x_left_off + emb_dim // 2 - x_left = tl.load(dK_ptr + x_left_off, mask=mask) - x_right = tl.load(dK_ptr + x_right_off, mask=mask) + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_left = tl.load(dK_ptr + x_1_off, mask=mask) + x_right = tl.load(dK_ptr + x_2_off, mask=mask) x_left_accum += x_left x_right_accum += x_right x_left_accum = tl.sum(x_left_accum, axis=0) @@ -578,9 +597,10 @@ def rotary_bwd_kv_kernel( tl.store(dEMB_ptr + tl.arange(0, emb_dim // 2) * 2 + 1, x_2) -class ApplyMLARotaryEmbKV(torch.autograd.Function): +class _FusedMLARoPEKVSplit(torch.autograd.Function): """ - Autograd function for applying YARN RoPE to MLA's key and value. + Autograd function for applying RoPE to MLA's key and value. + Splits KV, applies RoPE to k_pos_emb, concatenates onto key. """ @staticmethod @@ -599,7 +619,7 @@ def forward( rotary_interleaved=False, ): """ - Forward function for ApplyMLARotaryEmbKV. + Forward function for _FusedMLARoPEKVSplit. Args: kv: [seq_len, batch_size, head_num, k_dim + v_dim] @@ -634,7 +654,7 @@ def forward( o_value = kv.new_empty(total_seqlen, nheads, v_dim) grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_fwd_kv_kernel[grid]( + _mla_rope_fwd_kv_split_kernel[grid]( kv, k_pos_emb, o_key, @@ -674,7 +694,7 @@ def forward( @staticmethod def backward(ctx, dk, dv): """ - Backward function for ApplyMLARotaryEmbKV. + Backward function for _FusedMLARoPEKVSplit. Args: dk: [seq_len, batch_size, head_num, emb_dim + k_dim] @@ -702,7 +722,7 @@ def backward(ctx, dk, dv): d_emb = dk.new_empty(total_seqlen, 1, ctx.emb_dim) grid = lambda META: (total_seqlen, triton.cdiv(nheads, META["BLOCK_H"])) - rotary_bwd_kv_kernel[grid]( + _mla_rope_bwd_kv_split_kernel[grid]( dk, dv, d_kv, @@ -732,7 +752,7 @@ def backward(ctx, dk, dv): return d_kv, d_emb, None, None, None, None, None, None, None, None, None -def fused_apply_mla_rope_for_kv( +def fused_mla_rope_kv_split( kv: torch.Tensor, k_pos_emb: torch.Tensor, cos: torch.Tensor, @@ -746,7 +766,7 @@ def fused_apply_mla_rope_for_kv( rotary_interleaved: bool = False, ): """ - Fused function for applying YARN RoPE to MLA's key and value. + Fused function for applying RoPE to MLA's key and value. It splits the input tensor kv into key and value, and concatenates the processed RoPE to the key. @@ -767,7 +787,7 @@ def fused_apply_mla_rope_for_kv( or [total_seq_len, head_num, emb_dim + k_dim] value: [seq_len, batch_size, head_num, v_dim] or [total_seq_len, head_num, v_dim] """ - return ApplyMLARotaryEmbKV.apply( + return _FusedMLARoPEKVSplit.apply( kv, k_pos_emb, cos, @@ -780,3 +800,10 @@ def fused_apply_mla_rope_for_kv( cp_size, rotary_interleaved, ) + + +# --------------------------------------------------------------------------- +# Backward-compatible aliases (deprecated, prefer the new names above) +# --------------------------------------------------------------------------- +fused_apply_mla_rope_for_q = fused_mla_rope_inplace +fused_apply_mla_rope_for_kv = fused_mla_rope_kv_split diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index b990615da29..0608abb040a 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -95,6 +95,7 @@ def _apply_rotary_pos_emb_bshd( rotary_interleaved: bool = False, mla_rotary_interleaved: bool = False, mscale: float = 1.0, + inverse: bool = False, multi_latent_attention: Optional[bool] = None, ) -> Tensor: """Apply rotary positional embedding to input tensor T. @@ -118,6 +119,13 @@ def _apply_rotary_pos_emb_bshd( ) mla_rotary_interleaved = multi_latent_attention + # Some callers may pass freqs with an extra singleton axis, e.g. + # t: [s, b, d] and freqs: [s, 1, 1, d]. In that case, broadcasting would + # accidentally expand to [s, s, b, d]. Squeeze the extra singleton axis to + # keep freqs rank aligned with t. + if freqs.dim() == t.dim() + 1 and freqs.size(-2) == 1: + freqs = freqs.squeeze(-2) + rot_dim = freqs.shape[-1] # ideally t_pass is empty so rotary pos embedding is applied to all tensor t @@ -132,8 +140,17 @@ def _apply_rotary_pos_emb_bshd( # second part is sine component, need to change signs with _rotate_half method cos_ = (torch.cos(freqs) * mscale).to(t.dtype) sin_ = (torch.sin(freqs) * mscale).to(t.dtype) + if inverse: + sin_ = -sin_ t = (t * cos_) + (_rotate_half(t, rotary_interleaved) * sin_) + + # Fallback to original permutation + # DSv4 applies rope on V and O, so we need to uninterleave the tensor. + if mla_rotary_interleaved: + x1, x2 = torch.chunk(t, 2, dim=-1) + t = torch.stack((x1, x2), dim=-1).flatten(start_dim=-2) + return torch.cat((t, t_pass), dim=-1) @@ -193,6 +210,7 @@ def _apply_rotary_pos_emb_thd( rotary_interleaved: bool = False, mla_rotary_interleaved: bool = False, mscale: float = 1.0, + inverse: bool = False, cp_group: torch.distributed.ProcessGroup = None, multi_latent_attention: Optional[bool] = None, ) -> Tensor: @@ -246,6 +264,7 @@ def _apply_rotary_pos_emb_thd( rotary_interleaved=rotary_interleaved, mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, + inverse=inverse, ).squeeze(1) else: # CASE 2: Traditional mapping without offsets @@ -262,6 +281,7 @@ def _apply_rotary_pos_emb_thd( rotary_interleaved=rotary_interleaved, mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, + inverse=inverse, ).squeeze(1) @@ -273,6 +293,7 @@ def apply_rotary_pos_emb( mscale: float = 1.0, cp_group: torch.distributed.ProcessGroup = None, mla_rotary_interleaved: bool = False, + inverse: bool = False, ): """ Reroute to the appropriate apply_rotary_pos_emb function depending on @@ -307,6 +328,12 @@ def apply_rotary_pos_emb( "Using unfused implementation." ) use_unfused = True + if inverse: + warnings.warn( + "inverse RoPE is not supported by TE's fused RoPE. " + "Using unfused implementation." + ) + use_unfused = True if not use_unfused: assert fused_apply_rotary_pos_emb is not None, "apply_rope_fusion is not available." return fused_apply_rotary_pos_emb(t, freqs, interleaved=config.rotary_interleaved) @@ -328,6 +355,7 @@ def apply_rotary_pos_emb( rotary_interleaved=config.rotary_interleaved, mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, + inverse=inverse, ) else: return _apply_rotary_pos_emb_thd( @@ -338,6 +366,7 @@ def apply_rotary_pos_emb( mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, cp_group=cp_group, + inverse=inverse, ) diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py index fa4e15db856..984a9d4ea7f 100644 --- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -12,6 +12,19 @@ DSAttention, DSAttentionSubmodules, ) +from megatron.core.transformer.experimental_attention_variant.csa import ( + Compressor, + CompressorSubmodules, + CSAIndexer, + CSAIndexerSubmodules, + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, +) +from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + DSv4HybridSelfAttentionSubmodules, +) + from megatron.core.transformer.hyper_connection import HyperConnectionModule from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.multi_latent_attention import ( @@ -128,6 +141,64 @@ def get_dsa_module_spec_for_backend( return attention +def get_dsv4_hybrid_module_spec_for_backend( + config: TransformerConfig, backend: BackendSpecProvider = None +) -> ModuleSpec: + """Helper function to get module spec for DSv4 Hybrid Sparse Attention.""" + 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." + + # Adjust for RMS norm. + rms_norm = config.normalization == "RMSNorm" + # DSA indexer requires normalized q as input, so here we cannot fuse qk layernorm + # with linear projection and have to use unfused qk layernorm. + 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, + ), + ) + + attention = 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_up_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}, + ) + return attention + + def get_experimental_attention_variant_module_spec( config: TransformerConfig, backend: BackendSpecProvider = None ) -> ModuleSpec: @@ -140,6 +211,8 @@ def get_experimental_attention_variant_module_spec( return get_gated_delta_net_module_spec(config=config, backend=backend) elif config.experimental_attention_variant == "dsa": return get_dsa_module_spec_for_backend(config=config, backend=backend) + elif config.experimental_attention_variant == "dsv4_hybrid": + return get_dsv4_hybrid_module_spec_for_backend(config=config, backend=backend) else: raise ValueError( f"Invalid experimental attention variant: {config.experimental_attention_variant}" diff --git a/megatron/core/models/gpt/gpt_model.py b/megatron/core/models/gpt/gpt_model.py index 7bc92069e3a..05cd4d4c3e1 100644 --- a/megatron/core/models/gpt/gpt_model.py +++ b/megatron/core/models/gpt/gpt_model.py @@ -176,7 +176,7 @@ def __init__( cp_group=self.pg_collection.cp, ) - elif self.position_embedding_type == 'yarn': + elif self.position_embedding_type == 'yarn' and not self.config.multi_latent_attention: self.rotary_pos_emb = YarnRotaryEmbedding( kv_channels=self.config.kv_channels, rotary_percent=rotary_percent, @@ -381,7 +381,7 @@ def _preprocess( and packed_seq_params.qkv_format == 'thd', cp_group=packed_seq_params.cp_group if packed_seq_params is not None else None, ) - elif self.position_embedding_type == 'yarn': + elif self.position_embedding_type == 'yarn' and not self.config.multi_latent_attention: if self.training or not self.config.flash_decode: rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( inference_context, self.decoder, decoder_input, self.config, packed_seq_params diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py new file mode 100644 index 00000000000..9493f8348e4 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -0,0 +1,787 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +import copy +from dataclasses import dataclass +from functools import lru_cache +from typing import Optional, Tuple, Union + +import torch +import torch.nn as nn + +from megatron.core.fusions.fused_mla_yarn_rope_apply import fused_mla_rope_inplace +from megatron.core.models.common.embeddings import ( + RotaryEmbedding, + apply_rotary_pos_emb, +) +from megatron.core.packed_seq_params import PackedSeqParams +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.experimental_attention_variant.dsa import ( + DSAIndexerLossAutoScaler, + DSAIndexerLossLoggingHelper, + FusedDSAIndexerLoss, + fused_qk_topk_naive, + rotate_activation, +) +from megatron.core.transformer.module import MegatronModule +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.transformer_config import TransformerConfig +from megatron.core.utils import get_pg_size, nvtx_range_pop, nvtx_range_push + + +# --------------------------------------------------------------------------- +# Helper functions for index computation +# --------------------------------------------------------------------------- + + +@lru_cache(maxsize=8) +def _get_window_topk_idxs_cached(window_size: int, seqlen: int, device_str: str) -> torch.Tensor: + """Compute sliding-window indices for a single sequence (cached). + + Returns: + indices: [seqlen, window_size] int tensor, -1 for invalid positions. + """ + base = torch.arange(seqlen, device=device_str).unsqueeze(1) + offsets = torch.arange(window_size, device=device_str) + matrix = (base - window_size + 1).clamp(min=0) + offsets + matrix = torch.where(matrix > base, -1, matrix) + return matrix + + +def get_window_topk_idxs( + window_size: int, batch_size: int, seqlen: int, device: torch.device +) -> torch.Tensor: + """Sliding-window indices [batch, seqlen, window_size].""" + matrix = _get_window_topk_idxs_cached(window_size, seqlen, str(device)) + return matrix.unsqueeze(0).expand(batch_size, -1, -1) + + +@lru_cache(maxsize=8) +def _get_compress_topk_idxs_cached( + ratio: int, seqlen: int, offset: int, device_str: str +) -> torch.Tensor: + """Compute all-compressed-positions indices for a single sequence (cached). + + Returns: + indices: [seqlen, seqlen // ratio] int tensor, -1 for future positions. + """ + n_compressed = seqlen // ratio + matrix = torch.arange(n_compressed, device=device_str).repeat(seqlen, 1) + mask = matrix >= torch.arange(1, seqlen + 1, device=device_str).unsqueeze(1) // ratio + matrix = torch.where(mask, -1, matrix + offset) + return matrix + + +def get_compress_topk_idxs( + ratio: int, batch_size: int, seqlen: int, offset: int, device: torch.device +) -> torch.Tensor: + """All-compressed-position indices [batch, seqlen, seqlen // ratio].""" + matrix = _get_compress_topk_idxs_cached(ratio, seqlen, offset, str(device)) + return matrix.unsqueeze(0).expand(batch_size, -1, -1) + + +# --------------------------------------------------------------------------- +# Helper functions for RoPE +# --------------------------------------------------------------------------- + + +def _apply_rope( + x: torch.Tensor, + nope_dim: int, + pos_dim: int, + rotary_pos_emb_module: RotaryEmbedding, + config: TransformerConfig, + rotary_seq_len: int, + ratio: int = 1, + cp_group: torch.distributed.ProcessGroup = None, +) -> torch.Tensor: + """Apply RoPE to the last ``qk_pos_emb_head_dim`` dims, leaving the rest unchanged. + + Accepts both 3-D ``[seq, batch, head_dim]`` and 4-D ``[seq, batch, heads, head_dim]`` + inputs. When the input is 3-D a temporary head dimension is inserted for + ``apply_rotary_pos_emb`` and removed before returning. + """ + if ratio == 1: + total_seq_len = rotary_seq_len + else: + total_seq_len = rotary_seq_len * ratio + mscale = 1.0 + rotary_pos_cos = None + rotary_pos_sin = None + if config.rope_type == "rope": + rotary_pos_emb = rotary_pos_emb_module(total_seq_len, packed_seq=False) + mscale = 1.0 + else: + if config.apply_rope_fusion: + rotary_pos_cos, rotary_pos_sin = rotary_pos_emb_module.get_cached_cos_sin( + total_seq_len, dtype=x.dtype, packed_seq=False + ) + rotary_pos_emb = None + assert ( + fused_mla_rope_inplace is not None + ), "Fused MLA RoPE apply is not imported successfully" + else: + rotary_pos_emb, mscale = rotary_pos_emb_module(total_seq_len, packed_seq=False) + if rotary_pos_emb is not None and ratio > 1: + rotary_pos_emb = rotary_pos_emb[:total_seq_len:ratio][:rotary_seq_len] + if rotary_pos_cos is not None and ratio > 1: + rotary_pos_cos = rotary_pos_cos[:total_seq_len:ratio][:rotary_seq_len] + if rotary_pos_sin is not None and ratio > 1: + rotary_pos_sin = rotary_pos_sin[:total_seq_len:ratio][:rotary_seq_len] + + squeeze_head = x.dim() == 3 + if squeeze_head: + x = x.unsqueeze(-2) + if config.apply_rope_fusion: + out = fused_mla_rope_inplace( + x, + rotary_pos_cos, + rotary_pos_sin, + nope_dim, + pos_dim, + None, + cp_group.rank(), + cp_group.size(), + ) + else: + x_nope, x_pe = torch.split(x, [nope_dim, pos_dim], dim=-1) + x_pe = apply_rotary_pos_emb( + x_pe, + rotary_pos_emb, + config=config, + cu_seqlens=None, + mscale=mscale, + cp_group=cp_group, + mla_rotary_interleaved=True, + ) + out = torch.cat([x_nope, x_pe], dim=-1) + if squeeze_head: + out = out.squeeze(-2) + return out + + +# --------------------------------------------------------------------------- +# Sparse attention kernel (unfused, differentiable) +# --------------------------------------------------------------------------- + + +def unfused_compressed_sparse_attn( + query: torch.Tensor, + kv_full: torch.Tensor, + attn_sink: torch.Tensor, + topk_indices: torch.Tensor, + softmax_scale: float, +) -> torch.Tensor: + """Differentiable sparse attention with MQA and attention sink. + + Args: + query: [sq, b, np, hn] multi-head query. + kv_full: [n_kv, b, hn] single-head KV (original + compressed). + attn_sink: [np] per-head learnable bias. + topk_indices: [b, sq, topk] indices into kv_full (int32, -1 = invalid). + softmax_scale: float + + Returns: + output: [sq, b, np * hn] + """ + sq, b, np_, hn = query.size() + + # --- Gather KV at topk positions --- + # kv_full: [n_kv, b, hn] -> [b, n_kv, hn] + kv_t = kv_full.permute(1, 0, 2) + + safe_indices = topk_indices.clamp(min=0).long() # [b, sq, topk] + safe_indices_exp = safe_indices.unsqueeze(-1).expand(-1, -1, -1, hn) # [b, sq, topk, hn] + # [b, n_kv, hn] -> [b, 1, n_kv, hn] -> gather -> [b, sq, topk, hn] + kv_gathered = torch.gather( + kv_t.unsqueeze(1).expand(-1, sq, -1, -1), dim=2, index=safe_indices_exp + ) + + # --- Attention scores --- + # query: [sq, b, np, hn] -> [b, np, sq, hn] + q = query.permute(1, 2, 0, 3).float() + kv_g = kv_gathered.float() # [b, sq, topk, hn] + + # [b, np, sq, topk] + scores = torch.einsum("bnsh,bskh->bnsk", q, kv_g) * softmax_scale + + # Mask invalid + invalid_mask = (topk_indices < 0).unsqueeze(1) # [b, 1, sq, topk] + scores = scores.masked_fill(invalid_mask, float("-inf")) + + # --- Softmax with attention sink --- + sink = attn_sink.view(1, np_, 1, 1).float() + scores_max = scores.max(dim=-1, keepdim=True).values # [b, np, sq, 1] + scores_max = torch.max(scores_max, sink) + + exp_scores = torch.exp(scores - scores_max) # [b, np, sq, topk] + exp_sink = torch.exp(sink - scores_max) # [1, np, 1, 1] + + sum_exp = exp_scores.sum(dim=-1, keepdim=True) + exp_sink + attn_weights = exp_scores / sum_exp # [b, np, sq, topk] + + # --- Weighted sum --- + output = torch.einsum("bnsk,bskh->bnsh", attn_weights, kv_g) + output = output.to(query.dtype) + + # [b, np, sq, hn] -> [sq, b, np, hn] -> [sq, b, np * hn] + output = output.permute(2, 0, 1, 3).contiguous() + output = output.reshape(sq, b, np_ * hn) + return output + + +# --------------------------------------------------------------------------- +# Compressor +# --------------------------------------------------------------------------- + + +@dataclass +class CompressorSubmodules: + """Submodule specs for CSA and HCA Compressor.""" + + linear_wkv: Union[ModuleSpec, type] = None + linear_wgate: Union[ModuleSpec, type] = None + norm: Union[ModuleSpec, type] = None + + +class Compressor(MegatronModule): + """Gated pooling compressor for CSA and HCA sparse attention. + + Compresses a sequence of tokens into a shorter sequence by pooling groups of + ``compress_ratio`` tokens using learned gated weights. + + For ``compress_ratio == 4``, overlapping compression is used (``coff = 2``). + For ``compress_ratio == 128``, non-overlapping compression is used (``coff = 1``). + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CompressorSubmodules, + compress_ratio: int, + head_dim: int, + rotate: bool = False, + rotary_pos_emb: nn.Module = None, + pg_collection: Optional[ProcessGroupCollection] = None, + ) -> None: + super().__init__(config=config) + + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp'] + ) + self.pg_collection = pg_collection + + self.compress_ratio = compress_ratio + self.head_dim = head_dim + self.overlap = compress_ratio == 4 + self.coff = 1 + int(self.overlap) + self.rotate = rotate + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + + self.rotary_pos_emb = rotary_pos_emb + + proj_out_dim = self.coff * head_dim + + self.linear_wkv = build_module( + submodules.linear_wkv, + config.hidden_size, + proj_out_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) + + self.linear_wgate = build_module( + submodules.linear_wgate, + config.hidden_size, + proj_out_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) + + # keep to high precision + _ape = torch.empty( + compress_ratio, + proj_out_dim, + device=torch.cuda.current_device(), + dtype=torch.float32 + ) + config.init_method(_ape) + self.ape = nn.Parameter(_ape) + + norm_config = copy.copy(config) + norm_config.normalization = "RMSNorm" + self.norm = build_module( + submodules.norm, config=norm_config, hidden_size=head_dim, eps=config.layernorm_epsilon + ) + + def _overlap_transform(self, tensor: torch.Tensor, fill_value: float = 0) -> torch.Tensor: + """Apply overlapping window transform for 4x compression. + + Input shape: [n_groups, ratio, b, coff * head_dim] + Output shape: [n_groups, 2 * ratio, b, head_dim] + """ + n_groups, ratio, b_dim, _ = tensor.size() + d = self.head_dim + new_tensor = tensor.new_full((n_groups, 2 * ratio, b_dim, d), fill_value) + new_tensor[:, ratio:] = tensor[:, :, :, d:] + new_tensor[1:, :ratio] = tensor[:-1, :, :, :d] + return new_tensor + + def forward(self, x: torch.Tensor) -> Optional[torch.Tensor]: + """Compress hidden states into shorter KV sequence. + + Args: + x: [sq, b, hidden_size] + + Returns: + compressed_kv [sq // ratio, b, head_dim] or None if too short. + """ + nvtx_range_push("compressor") + + if self.config.sequence_parallel and get_pg_size(self.pg_collection.tp) > 1: + x = gather_from_sequence_parallel_region(x, group=self.pg_collection.tp) + + sq, b, _ = x.size() + ratio = self.compress_ratio + + if sq < ratio: + nvtx_range_pop("compressor") + return None + + kv, _ = self.linear_wkv(x) # [sq, b, coff * head_dim] + score, _ = self.linear_wgate(x) # [sq, b, coff * head_dim] + + cutoff = (sq // ratio) * ratio + if cutoff < sq: + kv = kv[:cutoff] + score = score[:cutoff] + + n_compressed = cutoff // ratio + + # Reshape: [n_compressed, ratio, b, coff * head_dim] + kv = kv.view(n_compressed, ratio, b, -1) + score = score.view(n_compressed, ratio, b, -1) + + # APE: [ratio, coff * head_dim] -> [1, ratio, 1, coff * head_dim] + score = score + self.ape.view(1, ratio, 1, -1) + + if self.overlap: + kv = self._overlap_transform(kv, fill_value=0) + score = self._overlap_transform(score, fill_value=float("-inf")) + + kv = (kv * torch.softmax(score, dim=1)).sum(dim=1) # [n_compressed, b, head_dim] + + kv = self.norm(kv.to(x.dtype)) + + kv = _apply_rope( + kv, + self.head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + self.rotary_pos_emb, + self.config, + n_compressed, + ratio=ratio, + cp_group=self.pg_collection.cp, + ) + + if self.rotate: + kv = rotate_activation(kv) + + nvtx_range_pop("compressor") + return kv # [n_compressed, b, head_dim] + + +# --------------------------------------------------------------------------- +# CSAIndexer +# --------------------------------------------------------------------------- + + +@dataclass +class CSAIndexerSubmodules: + """Submodule specs for CSAIndexer.""" + + linear_wq_b: Union[ModuleSpec, type] = None + linear_weights_proj: Union[ModuleSpec, type] = None + compressor: Union[ModuleSpec, type] = None + + +class CSAIndexer(MegatronModule): + """Learned top-k retrieval over compressed positions for CSA sparse attention. + + Computes index scores to select the most relevant compressed KV positions for each + query. Reuses the scoring logic from ``DSAIndexer`` (einsum -> relu -> weight -> sum + -> topk) and ``rotate_activation`` (Hadamard transform) from ``dsa.py``. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CSAIndexerSubmodules, + compress_ratio: int, + rotary_pos_emb: nn.Module = None, + pg_collection: Optional[ProcessGroupCollection] = None, + ) -> None: + super().__init__(config=config) + + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp'] + ) + self.pg_collection = pg_collection + + self.compress_ratio = compress_ratio + self.hidden_size = config.hidden_size + self.qk_pos_emb_head_dim = config.qk_pos_emb_head_dim + self.q_lora_rank = ( + config.q_lora_rank if config.q_lora_rank is not None else config.hidden_size + ) + + self.index_n_heads = config.dsa_indexer_n_heads + self.index_head_dim = config.dsa_indexer_head_dim + self.index_topk = config.dsa_indexer_topk + + self.softmax_scale: float = self.index_head_dim ** -0.5 + + self.rotary_pos_emb = rotary_pos_emb + + # Q projection + self.linear_wq_b = build_module( + submodules.linear_wq_b, + self.q_lora_rank, + self.index_n_heads * self.index_head_dim, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) + + # Weights projection + self.linear_weights_proj = build_module( + submodules.linear_weights_proj, + self.hidden_size, + self.index_n_heads, + config=config, + init_method=config.init_method, + bias=False, + skip_bias_add=False, + skip_weight_param_allocation=False, + parallel_mode="duplicated", + ) + + # Own compressor (smaller head_dim, with Hadamard rotation) + self.compressor = build_module( + submodules.compressor, + config=config, + compress_ratio=compress_ratio, + head_dim=self.index_head_dim, + rotate=True, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + ) + + def forward_before_topk( + self, + x: torch.Tensor, + qr: torch.Tensor, + packed_seq_params: Optional[PackedSeqParams] = None, + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Compute Q, compressed K, and weights before top-k selection.""" + nvtx_range_push("indexer_before_topk") + + # Gather for SP + if self.config.sequence_parallel and get_pg_size(self.pg_collection.tp) > 1: + x = gather_from_sequence_parallel_region(x, group=self.pg_collection.tp) + qr = gather_from_sequence_parallel_region(qr, group=self.pg_collection.tp) + + sq, bsz, _ = x.size() + + # Q path + q, _ = self.linear_wq_b(qr) # [sq, b, n_heads * head_dim] + q = q.reshape(sq, bsz, self.index_n_heads, self.index_head_dim) + q = _apply_rope( + q, + self.index_head_dim - self.qk_pos_emb_head_dim, + self.qk_pos_emb_head_dim, + self.rotary_pos_emb, + self.config, + sq, + ratio=1, + cp_group=self.pg_collection.cp, + ) + q = rotate_activation(q) + + # K path: own compressor + k = self.compressor(x) # [sq//ratio, b, index_head_dim] + + weights, _ = self.linear_weights_proj(x) # [sq, b, n_heads] + weights = weights * (self.index_n_heads ** -0.5) + + nvtx_range_pop("indexer_before_topk") + return q, k, weights + + def forward( + self, + x: torch.Tensor, + qr: torch.Tensor, + mask: Optional[torch.Tensor] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Return (index_scores, topk_indices).""" + nvtx_range_push("indexer") + assert packed_seq_params is None, "Packed sequence not supported for CSAIndexer" + q, k, weights = self.forward_before_topk(x, qr, packed_seq_params) + nvtx_range_push("indexer_qk_topk") + effective_topk = min(self.index_topk, k.size(0)) + index_scores, topk_indices = fused_qk_topk_naive(q, k, weights, effective_topk, mask) + nvtx_range_pop("indexer_qk_topk") + nvtx_range_pop("indexer") + return index_scores, topk_indices + + +# --------------------------------------------------------------------------- +# CompressedSparseAttention (core attention) +# --------------------------------------------------------------------------- + + +@dataclass +class CompressedSparseAttentionSubmodules: + """Submodule specs for CompressedSparseAttention.""" + + compressor: Union[ModuleSpec, type] = None + indexer: Union[ModuleSpec, type] = None + + +class CompressedSparseAttention(MegatronModule): + """Sparse core attention for CompressedSparseAttention. + + Combines sliding window attention with compressed KV attention. The spec always + provides compressor and indexer submodule specs; this ``__init__`` inspects + ``config.csa_compress_ratios[layer_idx]`` and conditionally builds them: + + * ``ratio == 1``: window-only (compressor and indexer NOT built) + * ``ratio == 4``: window + 4x compressed + learned Indexer (both built) + * ``ratio == 128``: window + 128x compressed, attend to all (compressor built only) + """ + + def __init__( + self, + config: TransformerConfig, + submodules: CompressedSparseAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType, + attention_type: str, + attention_dropout: Optional[float] = None, + softmax_scale: Optional[float] = None, + k_channels: Optional[int] = None, + v_channels: Optional[int] = None, + cp_comm_type: str = "p2p", + pg_collection: Optional[ProcessGroupCollection] = None, + rotary_pos_emb: nn.Module = None, + ): + super().__init__(config=config) + + if pg_collection is None: + pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp'] + ) + self.pg_collection = pg_collection + + self.layer_number = layer_number + layer_idx = layer_number - 1 # layer_number is 1-indexed + self.compress_ratio = config.csa_compress_ratios[layer_idx] + self.window_size = config.csa_window_size + self.v_head_dim = config.v_head_dim + + n_local_heads = config.num_attention_heads // get_pg_size(pg_collection.tp) + self.n_local_heads = n_local_heads + + if softmax_scale is None: + softmax_scale = config.v_head_dim ** -0.5 + self.softmax_scale = softmax_scale + + self.force_unfused_dsa = getattr(config, 'force_unfused_dsa', True) + + # Learnable attention sink per head + self.attn_sink = nn.Parameter(torch.zeros(n_local_heads, dtype=torch.float32)) + + # Conditionally build Compressor (ratio > 1) + if self.compress_ratio > 1 and submodules.compressor is not None: + self.compressor = build_module( + submodules.compressor, + config=config, + compress_ratio=self.compress_ratio, + head_dim=config.v_head_dim, + rotate=False, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + ) + else: + self.compressor = None + + # Conditionally build Indexer (ratio == 4) + if self.compress_ratio == 4 and submodules.indexer is not None: + self.indexer = build_module( + submodules.indexer, + config=config, + compress_ratio=self.compress_ratio, + rotary_pos_emb=rotary_pos_emb, + pg_collection=pg_collection, + ) + else: + self.indexer = None + + def forward( + self, + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + attention_mask: torch.Tensor, + x: torch.Tensor = None, + qr: torch.Tensor = None, + attn_mask_type: AttnMaskType = None, + attention_bias: torch.Tensor = None, + packed_seq_params: PackedSeqParams = None, + ) -> torch.Tensor: + """Forward pass for CompressedSparseAttention. + + Args: + query: [sq, b, np, v_head_dim] + key: [sq, b, 1, v_head_dim] (single-head MQA; head dim squeezed internally) + value: unused (key == value in MQA) + attention_mask: attention mask (may be None for causal). + x: [sq, b, hidden_size] original hidden states. + qr: [sq, b, q_lora_rank] compressed query representation. + + Returns: + output: [sq, b, np * v_head_dim] + """ + nvtx_range_push("compressed_sparse_attn") + assert packed_seq_params is None, \ + "Packed sequence not supported for CompressedSparseAttention" + + sq, b, np, hn = query.size() + + # --- Step 1: Prepare single-head KV (squeeze singleton head dim) --- + kv = key.squeeze(-2) # [sq, b, 1, v_head_dim] -> [sq, b, v_head_dim] + + # --- Step 2: Compression --- + if self.compressor is not None and self.compress_ratio > 1: + compressed_kv = self.compressor(x) # [n_compressed, b, v_head_dim] + if compressed_kv is not None: + kv_full = torch.cat([kv, compressed_kv], dim=0) + n_compressed = compressed_kv.size(0) + else: + kv_full = kv + n_compressed = 0 + else: + kv_full = kv + n_compressed = 0 + + offset = sq # compressed indices start after original positions + + # --- Step 3: Window indices --- + window_idxs = get_window_topk_idxs(self.window_size, b, sq, query.device) + + # --- Step 4: Compressed indices --- + indexer_loss = None + + if self.force_unfused_dsa: + if self.compress_ratio > 1 and n_compressed > 0: + nvtx_range_push("compressed_indices") + if self.indexer is not None: + x_det = x.detach() + qr_det = qr.detach() + + causal_mask = torch.arange(n_compressed, device=x.device).unsqueeze(0).expand( + sq, -1 + ) + positions = torch.arange(1, sq + 1, device=x.device).unsqueeze(1) + causal_mask = torch.where( + causal_mask >= positions // self.compress_ratio, + float("-inf"), + 0.0, + ).unsqueeze(0).expand(b, -1, -1) # [b, sq, n_compressed] + + if self.training and torch.is_grad_enabled(): + q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk( + x_det, qr_det, packed_seq_params + ) + indexer_loss_coeff = getattr(self.config, 'dsa_indexer_loss_coeff', 0.0) + # compressed_kv is [n, b, hn]; expand to [n, b, np, hn] for loss + key_for_loss = compressed_kv.unsqueeze(2).expand(-1, -1, np, -1) + # ``FusedDSAIndexerLoss`` does not accept a separate + # indexer_softmax_scale; apply it here via the + # weights-scaling trick so the effective weights match + # the pre-scale-split behaviour. + weights_for_unfused = ( + weights_indexer * self.indexer.softmax_scale + ) + topk_indices_compressed, indexer_loss = FusedDSAIndexerLoss.apply( + q_indexer, + weights_for_unfused, + k_indexer, + query.detach(), + key_for_loss.detach(), + self.softmax_scale, + min(self.indexer.index_topk, n_compressed), + indexer_loss_coeff, + causal_mask, + getattr(self.config, "dsa_indexer_use_sparse_loss", True), + self.indexer.pg_collection, + ) + if indexer_loss_coeff > 0: + DSAIndexerLossLoggingHelper.save_loss_to_tracker( + loss=indexer_loss, + layer_number=self.layer_number, + num_layers=self.config.num_layers, + ) + else: + _, topk_indices_compressed = self.indexer( + x_det, qr_det, mask=causal_mask, packed_seq_params=packed_seq_params + ) + + n_valid_per_pos = positions // self.compress_ratio # [sq, 1] + valid = topk_indices_compressed < n_valid_per_pos + compress_topk_idxs = torch.where( + valid, topk_indices_compressed + offset, torch.tensor(-1, device=x.device) + ) + else: + compress_topk_idxs = get_compress_topk_idxs( + self.compress_ratio, b, sq, offset, query.device + ) + + topk_idxs = torch.cat([window_idxs, compress_topk_idxs], dim=-1) + nvtx_range_pop("compressed_indices") + else: + topk_idxs = window_idxs + + topk_idxs = topk_idxs.int() + + # --- Step 5: Sparse attention --- + nvtx_range_push("sparse_attn_kernel") + output = unfused_compressed_sparse_attn( + query, kv_full, self.attn_sink.float(), topk_idxs, self.softmax_scale + ) + nvtx_range_pop("sparse_attn_kernel") + + else: + raise ValueError("Fused path is not supported for CompressedSparseAttention") + + # --- Step 6: Attach indexer loss --- + if indexer_loss is not None and self.training and torch.is_grad_enabled(): + output = DSAIndexerLossAutoScaler.apply(output, indexer_loss) + + nvtx_range_pop("compressed_sparse_attn") + return output 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..7f4256ba2a4 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -0,0 +1,886 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + +import math +from dataclasses import dataclass +from typing import NoReturn, Optional, Union + +import torch + +from megatron.core import tensor_parallel +from megatron.core.dist_checkpointing.mapping import ShardedObject +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.tensor_parallel.layers import ColumnParallelLinear +from megatron.core.tensor_parallel.mappings import ( + gather_from_tensor_model_parallel_region, + scatter_to_sequence_parallel_region, +) +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_kv_split, + ) +except Exception: + fused_mla_rope_inplace = None + fused_mla_rope_kv_split = None + + +if HAVE_TE: + from megatron.core.extensions.transformer_engine import ( + TEColumnParallelLinear, + TELinear, + set_save_original_input, + ) + from megatron.core.post_training.modelopt.layers import Linear +else: + ( + TEColumnParallelLinear, + TELayerNormColumnParallelLinear, + TELinear, + Linear, + set_save_original_input, + split_te_layernorm_column_parallel_linear, + ) = (None, None, None, 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_up_proj: Union[ModuleSpec, type] = None + linear_qkv_down_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, + ) -> None: + + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + attention_type=attention_type, + attn_mask_type=attn_mask_type, + pg_collection=pg_collection, + ) + self.config: MLATransformerConfig + + 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 + + rope_base = self.config.rotary_base + compress_ratio = self.config.csa_compress_ratios[layer_number - 1] + if compress_ratio > 1: + rope_base = self.config.csa_compress_rotary_base + if self.config.rope_type == "rope": + 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, + ) + elif self.config.rope_type == "yarn": + 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, + ) + else: + raise ValueError( + f"Unsupported RoPE type: {self.config.rope_type}, supported types are " + "'rope' and 'yarn'" + ) + + core_attn_extra_kwargs = {"rotary_pos_emb": self.rotary_pos_emb} + 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. + tp_size = get_pg_size(self.pg_collection.tp) + assert self.config.o_groups % tp_size == 0, ( + "o_groups must be divisible by tp_size" + ) + self.o_local_groups = self.config.o_groups // tp_size + 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." + + # ===================== + # Query, Key, and Value + # ===================== + # Get the query, key and value tensors based on the type of attention - + # self or cross attn. + qkv_linear_manager = off_interface(self.offload_qkv_linear, hidden_states, "qkv_linear") + with qkv_linear_manager as hidden_states: + query, key, value, q_compressed, kv_compressed = self.get_query_key_value_tensors( + hidden_states, + key_value_states, + position_ids, + packed_seq_params, + inference_context=inference_context, + ) + query = qkv_linear_manager.group_offload(query, forced_released_tensors=[hidden_states]) + + # 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" + ) + if self.checkpoint_core_attention and self.training: + core_attn_out = self._checkpointed_attention_forward( + query, key, value, attention_mask, packed_seq_params=packed_seq_params + ) + else: + extra_kwargs = {} + if self.config.experimental_attention_variant in ("dsa", "dsv4_hybrid"): + # For dsa we need to pass in the original hidden states and the compressed + # query representation. + extra_kwargs["x"] = hidden_states + extra_kwargs["qr"] = q_compressed + with core_attn_manager as query: + core_attn_out = self.core_attention( + query, + key, + value, + attention_mask, + packed_seq_params=packed_seq_params, + **extra_kwargs, + ) + core_attn_out = core_attn_manager.group_offload( + core_attn_out, forced_released_tensors=[query, key, value] + ) + + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + # reshape to same output shape as unpacked case + # (t, np, hn) -> (t, b=1, h=np*hn) + # t is the pack size = sum (sq_i) + # note that batch is a dummy dimension in the packed case + core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1) + + 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) + packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + if packed_seq: + cu_seqlens_kv = ( + packed_seq_params.cu_seqlens_kv_padded + if packed_seq_params.cu_seqlens_kv_padded is not None + else packed_seq_params.cu_seqlens_kv + ) + rope_seqlen = cu_seqlens_kv[-1].item() + else: + cu_seqlens_kv = None + rope_seqlen = seq_len + mscale = 1.0 + rotary_pos_cos = None + rotary_pos_sin = None + if self.config.rope_type == "rope": + rotary_pos_emb = self.rotary_pos_emb(rope_seqlen, packed_seq=packed_seq) + else: + if self.config.apply_rope_fusion: + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + rope_seqlen, dtype=hidden_states.dtype, packed_seq=packed_seq + ) + 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" + else: + rotary_pos_emb, mscale = self.rotary_pos_emb(rope_seqlen, packed_seq=packed_seq) + if self.config.apply_rope_fusion: + core_attn_out = fused_mla_rope_inplace( + core_attn_out, + rotary_pos_cos, + rotary_pos_sin, + nope_dim, + pos_dim, + cu_seqlens_kv, + self.pg_collection.cp.rank(), + self.pg_collection.cp.size(), + inverse=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=cu_seqlens_kv, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=True, + inverse=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, + ): + if pg_collection is None: + 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, + ) + + q_down_proj_kwargs = {} + if submodules.linear_q_down_proj in [TELinear]: + q_down_proj_kwargs['parallel_mode'] = 'duplicated' + elif submodules.linear_q_down_proj in [ + Linear, + TEColumnParallelLinear, + ColumnParallelLinear, + ]: + q_down_proj_kwargs['gather_output'] = False + 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=( + pg_collection.tp + if q_down_proj_kwargs.get('parallel_mode') != 'duplicated' + 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, + ) + + self.linear_kv_down_proj = None + self.linear_kv_up_proj = build_module( + submodules.linear_kv_up_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, + ) + 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" + if packed_seq_params is not None: + assert ( + packed_seq_params.local_cp_size is None + ), "dynamic_context_parallel is not supported with MLA yet and is planned for future. \ + Please disable dynamic_context_parallel." + + 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, packed_seq_params + ) + + # rotary_pos_emb:[s, b, 1, 64] + mscale = 1.0 + rotary_pos_cos = None + rotary_pos_sin = None + packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == 'thd' + if self.config.rope_type == "rope": + rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) + else: + if self.config.apply_rope_fusion: + rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb.get_cached_cos_sin( + rotary_seq_len, dtype=hidden_states.dtype, packed_seq=packed_seq + ) + 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 + and fused_mla_rope_kv_split is not None + ), "Fused MLA RoPE apply is not imported successfully" + else: + rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) + + if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': + if packed_seq_params.cu_seqlens_q_padded is not None: + cu_seqlens_q = packed_seq_params.cu_seqlens_q_padded + else: + cu_seqlens_q = packed_seq_params.cu_seqlens_q + if packed_seq_params.cu_seqlens_kv_padded is not None: + cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded + else: + cu_seqlens_kv = packed_seq_params.cu_seqlens_kv + else: + cu_seqlens_q = cu_seqlens_kv = None + + # ========================================= + # QKV down projection and layernorm + # ========================================= + # if linear_q_down_proj is ColumnParallelLinear: + # q_compressed: [s, b, q_lora_rank / TP] + # elif linear_q_down_proj is Linear: + # q_compressed: [s / TP, b, q_lora_rank] + q_compressed, _ = self.linear_q_down_proj(hidden_states) + + # When output is sharded (ColumnParallelLinear), two things are needed to be + # identical to a normal Linear. + # 1. Manually gather output to restore output dim q_lora_rank; + # 2. Scatter sequence back to s / TP if sequence-parallel since it was + # gathered by ColumnParallelLinear. + if q_compressed.size(-1) != self.config.q_lora_rank: + q_compressed = gather_from_tensor_model_parallel_region(q_compressed) + if self.config.sequence_parallel: + q_compressed = scatter_to_sequence_parallel_region(q_compressed) + + kv_compressed = hidden_states + k_pos_emb = None + + if packed_seq_params is not None: + # If sequence packing, TE expect [t, h, d] shaped qkv input. + # In Megatron-Core, the qkv shape is [t, 1, h, d]. + # So we need to reshape qkv from [t, 1, h, d] to [t, h, d]. + q_compressed = q_compressed.squeeze(1) + + # ========================================= + # 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 query and key. + When sequence packing enabled, the input tensors adopt a packed shape of [t, ...]; + otherwise, they maintain the unpacked shape [s, b, ...]. In subsequent code comments, + we uniformly use [num_tokens, ...] to denote [s, b, ...] or [t, ...] for two cases. + """ + # q_compressed: [num_tokens, q_lora_rank] + # q: [num_tokens, 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_up_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) + + # todo add assert about fusions and caching + 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, + cu_seqlens_q, + cp_rank, + cp_size, + ) + 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, + cu_seqlens_q, + cp_rank, + cp_size, + ) + key = kv + value = kv + else: + q_len = q.size()[0] + if packed_seq_params is None or self.config.context_parallel_size == 1: + # Shorten rotary_pos_emb to the sequence length when inference_params + # is not provided. This makes sure we can run forward directly with + # any sequence length. During training, the sequence length is always + # the full rotary_pos_emb length, except for sequence packing + CP. + # When sequence packing and context parallel are both enabled, the + # position embedding will not split rotary_pos_emb, so it may exceed + # the sequence length on this CP rank, but we need the full rotary_pos_emb + # to cover the full sequence, so we do not shorten it here. + 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=cu_seqlens_q, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=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=cu_seqlens_kv, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=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_up_proj.backward_dw() + if self.linear_kv_down_proj is not None: + self.linear_kv_down_proj.backward_dw() + + def _backward_q_proj(self): + """Computes weight gradients of Q projection layers""" + if self.config.q_lora_rank is None: + self.linear_q_proj.backward_dw() + else: + 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.""" + if self.config.q_lora_rank is not None: + set_save_original_input(self.linear_q_down_proj) + if self.linear_kv_down_proj is not None: + set_save_original_input(self.linear_kv_down_proj) + elif not self.recompute_up_proj: + set_save_original_input(self.linear_kv_up_proj) + + def clip_qk(self): + """ + QK Clipping is a technique to clip the query and key attention logits to prevent the + attention logits from exploding. Per MuonClip usage, we update the weight by calling this + function after Muon optimizer step. + """ + + if not self.config.qk_clip: + raise ValueError("qk_clip option needs to be enabled") + + if self.core_attention.current_max_attn_logits is None: + raise ValueError("current_max_attn_logits is None") + + # Check if we're in absorption mode + if self.cache_mla_latents and not hasattr(self, 'linear_kv_up_proj'): + raise ValueError( + "qk_clip is not supported when cache_mla_latents is enabled and absorption is " + "active. The linear_kv_up_proj layer has been deleted during absorption " + "preparation." + ) + + assert self.core_attention.current_max_attn_logits.shape == ( + self.num_attention_heads_per_partition, + ), f"current_max_attn_logits shape is not ({self.num_attention_heads_per_partition}, ) \ + but {self.core_attention.current_max_attn_logits.shape}" + + # only update the weight if any head has + # current_max_attn_logits > qk_clip_threshold + if torch.any(self.core_attention.current_max_attn_logits > self.config.qk_clip_threshold): + # Use num_attention_heads_per_partition for tensor parallel scenarios + + # qk_clip_balancing_eta (n, 1, 1) + assert self.core_attention.current_max_attn_logits.shape == ( + self.num_attention_heads_per_partition, + ), f"current_max_attn_logits shape is not ({self.num_attention_heads_per_partition},) \ + but {self.core_attention.current_max_attn_logits.shape}" + self.qk_clip_balancing_eta = torch.clamp( + self.config.qk_clip_threshold / self.core_attention.current_max_attn_logits, max=1.0 + ).view(self.num_attention_heads_per_partition, 1, 1) + assert torch.all(self.qk_clip_balancing_eta <= 1.0) + + # Update q side weight, keep qk_pos_emb_head_dim side weight unchanged + if self.config.q_lora_rank is None: + q_proj_weight = self.linear_q_proj.weight + else: + q_proj_weight = self.linear_q_up_proj.weight + + # Handle different weight access patterns (main_param vs direct access) + if hasattr(q_proj_weight, 'main_param'): + q_proj_weight.main_param.data.copy_( + self._clip_q_proj_weight(q_proj_weight.main_param.data) + ) + q_proj_weight.data.copy_(self._clip_q_proj_weight(q_proj_weight.data)) + + # Update k side weight, keep v side weight unchanged + kv_proj_weight = self.linear_kv_up_proj.weight + + # Handle different weight access patterns + if hasattr(kv_proj_weight, 'main_param'): + kv_proj_weight.main_param.data.copy_( + self._clip_kv_proj_weight(kv_proj_weight.main_param.data) + ) + kv_proj_weight.data.copy_(self._clip_kv_proj_weight(kv_proj_weight.data)) + + # reset current_max_attn_logits + self.core_attention.current_max_attn_logits = None + + def _clip_q_proj_weight(self, weight): + """Clip q_proj_weight""" + # Reshape to (n, a + b, -1) + weight_reshaped = weight.view( + self.num_attention_heads_per_partition, + self.config.qk_head_dim + self.config.qk_pos_emb_head_dim, + -1, + ) + + # Split into qk_head_dim and qk_pos_emb_head_dim parts: (n, a, -1) and (n, b, -1) + weight_q_nope = weight_reshaped[:, : self.config.qk_head_dim, :] + weight_q_pe = weight_reshaped[:, self.config.qk_head_dim :, :] + + # Clipping + weight_q_nope.mul_(torch.pow(self.qk_clip_balancing_eta, self.config.qk_clip_alpha)) + weight_q_pe.mul_(self.qk_clip_balancing_eta) + + # Concatenate back and reshape to original shape + weight_q_updated = torch.cat([weight_q_nope, weight_q_pe], dim=1) + weight_q_updated = weight_q_updated.view( + self.num_attention_heads_per_partition + * (self.config.qk_head_dim + self.config.qk_pos_emb_head_dim), + -1, + ) + + return weight_q_updated + + def _clip_kv_proj_weight(self, weight): + """Clip kv_proj_weight""" + # shape: (n, qk_head_dim + v_head_dim, kv_lora_rank) + weight_reshaped = weight.view( + self.num_attention_heads_per_partition, + self.config.qk_head_dim + self.config.v_head_dim, + -1, + ) + + # Split into qk_head_dim and v_head_dim parts: (n, a, -1) and (n, b, -1) + weight_k = weight_reshaped[:, : self.config.qk_head_dim, :] + weight_v = weight_reshaped[:, self.config.qk_head_dim :, :] + + # Clipping + weight_k.mul_(torch.pow(self.qk_clip_balancing_eta, 1 - self.config.qk_clip_alpha)) + + # Concatenate back and reshape to original shape + weight_kv_updated = torch.cat([weight_k, weight_v], dim=1) + weight_kv_updated = weight_kv_updated.view( + self.num_attention_heads_per_partition + * (self.config.qk_head_dim + self.config.v_head_dim), + -1, + ) + + return weight_kv_updated diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index 5c5f77363dc..2ded8e3561e 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -167,6 +167,7 @@ def compute_dsa_indexer_loss( loss_coeff: float, sparse_loss: bool, pg_collection: ProcessGroupCollection, + causal_mask_override: Optional[torch.Tensor] = None, ) -> torch.Tensor: """ Compute KL divergence loss between index_scores and true attention_scores. @@ -203,29 +204,60 @@ def compute_dsa_indexer_loss( # Reshape to [b, np, sq, sk] attention_scores = attention_scores.reshape(b, np, sq, sk) - # causal_mask [sq, sk] - causal_mask = torch.triu( - torch.full((sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device), - diagonal=1, - ) + # causal_mask: use caller-provided mask when available (handles compressed KV), + # otherwise fall back to standard upper-triangular causal mask. + if causal_mask_override is not None: + causal_mask = causal_mask_override.to(dtype=torch.float32) # [b, sq, sk] + else: + causal_mask = torch.triu( + torch.full((sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device), + diagonal=1, + ) # index_mask [b, sq, sk] index_mask = torch.full( (b, sq, sk), float("-inf"), dtype=torch.float32, device=causal_mask.device ).scatter_(-1, topk_indices, 0) - # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv] - attention_scores += causal_mask.view(1, 1, sq, sk) + # Apply causal mask to attention_scores + # causal_mask: [b, sq, sk] (from causal_mask_override) or [sq, sk] (from triu) + if causal_mask.dim() == 3: + attention_scores = attention_scores + causal_mask.unsqueeze(1) # [b,1,sq,sk] + else: + attention_scores = attention_scores + causal_mask.view(1, 1, sq, sk) if sparse_loss: # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk] attention_scores += index_mask.view(b, 1, sq, sk) # [b, sq, sk] + [b, sq, sk] -> [b, sq, sk] index_scores += index_mask + # Identify rows where all KV positions are masked (e.g., early query positions with + # compress_ratio=4 have zero valid compressed KV entries). These rows would produce NaN + # from softmax(all -inf). We zero out their logits before softmax and mask out their + # contributions after, so NaN is never produced. + # row_valid: [b, sq] or [sq] — True if the row has at least one unmasked position. + row_valid = (causal_mask > float('-inf')).any(dim=-1) + if row_valid.dim() == 1: + # [sq] -> broadcast for attention_scores [b, np, sq, sk] and index_scores [b, sq, sk] + attn_row_mask = row_valid.view(1, 1, sq, 1) # [1, 1, sq, 1] + idx_row_mask = row_valid.view(1, sq, 1) # [1, sq, 1] + else: + # [b, sq] + attn_row_mask = row_valid.view(b, 1, sq, 1) # [b, 1, sq, 1] + idx_row_mask = row_valid.view(b, sq, 1) # [b, sq, 1] + + # Zero out fully-masked rows before softmax so it produces valid uniform distribution + attention_scores = attention_scores.masked_fill(~attn_row_mask, 0.0) + index_scores = index_scores.masked_fill(~idx_row_mask, 0.0) + # [b, np, sq, sk] -> [b, np, sq, sk] attention_scores = torch.nn.functional.softmax(attention_scores, dim=-1, dtype=torch.float32) # [b, sq, sk] -> [b, sq, sk] index_scores = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) + # Zero out invalid rows so they contribute nothing to loss/gradients + attention_scores = attention_scores * attn_row_mask.float() + index_scores = index_scores * idx_row_mask.float() + # Sum attention scores across heads. # [batch, heads, seqlen_q, seqlen_k] -> [batch, seqlen_q, seqlen_k] attention_scores = attention_scores.sum(dim=1) @@ -234,7 +266,9 @@ def compute_dsa_indexer_loss( torch.distributed.all_reduce(attention_scores.contiguous(), group=pg_collection.tp) # L1 normalize target on the last dimension. Doesn't use abs() because attention_scores are # obtained from softmax so they are already non-negative. - attention_scores = attention_scores / attention_scores.sum(dim=-1, keepdim=True) + attention_scores = attention_scores / ( + attention_scores.sum(dim=-1, keepdim=True).clamp(min=1e-10) + ) # Compute KL divergence: KL(target || index) = target(x) * log(target(x) / index(x)) # kl_per_element [b, sq, sk] @@ -338,6 +372,7 @@ def fwd_fused_indexer_loss_naive( loss_coeff, sparse_loss, pg_collection, + causal_mask_override=mask, ) return topk_indices, indexer_loss @@ -355,6 +390,7 @@ def bwd_fused_indexer_loss_naive( sparse_loss, grad_loss, pg_collection, + causal_mask_override=None, ): """Naive implementation of backward pass for indexer loss.""" index_scores = _compute_index_scores(q, weights, k) # [B, Sq, Sk] @@ -374,23 +410,30 @@ def bwd_fused_indexer_loss_naive( # Reshape to [b, np, sq, sk] attention_scores = attention_scores.reshape(b, np, sq, sk) - # causal_mask [sq, sk] - causal_mask = torch.triu( - torch.full((sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device), - diagonal=1, - ) + # causal_mask: use caller-provided mask when available (handles compressed KV), + # otherwise fall back to standard upper-triangular causal mask. + if causal_mask_override is not None: + causal_mask = causal_mask_override.to(dtype=torch.float32) # [b, sq, sk] + else: + causal_mask = torch.triu( + torch.full( + (sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device + ), + diagonal=1, + ) # index_mask [b, sq, sk] index_mask = torch.full( (b, sq, sk), float("-inf"), dtype=torch.float32, device=causal_mask.device ).scatter_(-1, topk_indices, 0) # Apply causal mask to both attention and index scores - # [b, np, sq, skv] + [1, 1, sq, skv] -> [b, np, sq, skv] - attention_scores = attention_scores + causal_mask.view(1, 1, sq, sk) - # [b, sq, sk] + [1, sq, sk] -> [b, sq, sk] - index_scores = index_scores + causal_mask.unsqueeze(0) - # Free causal_mask - no longer needed - del causal_mask + # attention_scores: [b, np, sq, sk], causal_mask: [b, sq, sk] or [sq, sk] + if causal_mask.dim() == 3: + attention_scores = attention_scores + causal_mask.unsqueeze(1) # [b,1,sq,sk] + index_scores = index_scores + causal_mask # [b,sq,sk] + else: + attention_scores = attention_scores + causal_mask.view(1, 1, sq, sk) + index_scores = index_scores + causal_mask.unsqueeze(0) if sparse_loss: # [b, np, sq, sk] + [b, 1, sq, sk] -> [b, np, sq, sk] @@ -398,17 +441,40 @@ def bwd_fused_indexer_loss_naive( # [b, sq, sk] + [b, sq, sk] -> [b, sq, sk] index_scores = index_scores + index_mask - # Compute softmax for both + # Identify rows where all KV positions are masked (e.g., early query positions with + # compress_ratio=4 have zero valid compressed KV entries). Zero out their logits before + # softmax and mask out contributions after, so NaN is never produced. + row_valid = (causal_mask > float('-inf')).any(dim=-1) + # Free causal_mask - no longer needed + del causal_mask + if row_valid.dim() == 1: + attn_row_mask = row_valid.view(1, 1, sq, 1) + idx_row_mask = row_valid.view(1, sq, 1) + else: + attn_row_mask = row_valid.view(b, 1, sq, 1) + idx_row_mask = row_valid.view(b, sq, 1) + + # Zero out fully-masked rows before softmax + attention_scores = attention_scores.masked_fill(~attn_row_mask, 0.0) + index_scores = index_scores.masked_fill(~idx_row_mask, 0.0) + + # Compute softmax attention_scores_softmax = torch.nn.functional.softmax( attention_scores, dim=-1, dtype=torch.float32 ) # Free attention_scores immediately del attention_scores - index_scores_softmax = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) + index_scores_softmax = torch.nn.functional.softmax( + index_scores, dim=-1, dtype=torch.float32 + ) # Free index_scores - no longer needed after softmax del index_scores + # Zero out invalid rows so they contribute nothing to gradients + attention_scores_softmax = attention_scores_softmax * attn_row_mask.float() + index_scores_softmax = index_scores_softmax * idx_row_mask.float() + # Sum attention scores across heads: [b, np, sq, sk] -> [b, sq, sk] attention_scores_sum = attention_scores_softmax.sum(dim=1) # Free attention_scores_softmax @@ -421,7 +487,7 @@ def bwd_fused_indexer_loss_naive( # L1 normalize attention_scores_normalized = attention_scores_sum / attention_scores_sum.sum( dim=-1, keepdim=True - ) + ).clamp(min=1e-10) # Free attention_scores_sum - no longer needed after normalization del attention_scores_sum @@ -452,10 +518,17 @@ def bwd_fused_indexer_loss_naive( # Zero out gradients for masked positions # Create a mask for valid (non-masked) positions - # Causal mask: position (i, j) is valid if j <= i - causal_valid_mask = torch.tril( - torch.ones((sq, sk), device=q.device, dtype=torch.bool) - ) # [sq, sk] + if causal_mask_override is not None: + # Derive valid mask from the causal_mask_override: valid where mask == 0 + _cm = causal_mask_override.to(dtype=torch.float32) + if _cm.dim() == 2: + _cm = _cm.unsqueeze(0) # [1, sq, sk] + causal_valid_mask = (_cm == 0).squeeze(0) if _cm.shape[0] == 1 else (_cm == 0) + else: + # Standard causal: position (i, j) is valid if j <= i + causal_valid_mask = torch.tril( + torch.ones((sq, sk), device=q.device, dtype=torch.bool) + ) # [sq, sk] if sparse_loss: # Also apply index mask - only topk positions are valid index_valid_mask = index_mask == 0 # [b, sq, sk] @@ -543,7 +616,7 @@ def forward( ) # Save for backward (recomputation strategy) - ctx.save_for_backward(q, weights, k, query, key, topk_indices) + ctx.save_for_backward(q, weights, k, query, key, topk_indices, mask) ctx.softmax_scale = softmax_scale ctx.loss_coeff = loss_coeff ctx.sparse_loss = sparse_loss @@ -556,7 +629,7 @@ def backward(ctx, grad_topk_indices, grad_loss): """ Backward: Recompute what we need. """ - q, weights, k, query, key, topk_indices = ctx.saved_tensors + q, weights, k, query, key, topk_indices, mask = ctx.saved_tensors grad_q, grad_weights, grad_k = bwd_fused_indexer_loss_naive( q, @@ -570,6 +643,7 @@ def backward(ctx, grad_topk_indices, grad_loss): ctx.sparse_loss, grad_loss, ctx.pg_collection, + causal_mask_override=mask, ) # query and key are detached in forward, so return None for their gradients diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 4251ee9d8a1..cb2b1be7e71 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -264,8 +264,8 @@ 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.""" #################### # DSA @@ -286,6 +286,18 @@ class TransformerConfig(ModelParallelConfig): """Whether to use sparse DSA indexer loss. If True, the indexer loss will be computed using the top-k indices.""" + #################### + # DeepSeek-v4 hybrid attention + #################### + csa_window_size: int = 128 + """Sliding window size for compressed sparse attention.""" + + csa_compress_ratios: Optional[List[int]] = None + """Per-layer compress ratios, e.g. [1, 1, 4, 128, 4, 128, ...].""" + + csa_compress_rotary_base: float = 40000.0 + """RoPE base for compressed KV positions in compressed sparse attention.""" + #################### # linear attention #################### @@ -1263,6 +1275,15 @@ def __post_init__(self): ) elif self.experimental_attention_variant == "dsa": pass + 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" + assert ( + len(self.csa_compress_ratios) == self.num_layers + ), "csa_compress_ratios must have the same length as num_layers" + assert ( + all(ratio in [1, 4, 128] for ratio in self.csa_compress_ratios) + ), "csa_compress_ratios must be 1, 4, or 128" if self.fp8: # cannot support first last layer bf16 with delayed scaling @@ -2737,6 +2758,12 @@ 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 @@ -2755,6 +2782,21 @@ def __post_init__(self): 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." + log_single_rank( + logger, + logging.WARNING, + f"DSv4 hybrid mode is enabled, deriving qk_head_dim and kv_lora_rank from " + f"v_head_dim and qk_pos_emb_head_dim" + ) + derived = self.v_head_dim - self.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/training/arguments.py b/megatron/training/arguments.py index 4b9f9790a0c..9bf0b451a8c 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -271,6 +271,21 @@ def no_rope_freq_type(x): return int(x) +def compress_ratios_type(x): + """Per-layer compress ratios for compressed sparse attention. + + Accepts a string containing a Python list expression, e.g.: + "[1,1,4,128,4,128]" + "([1]+[4,128]*2)*3" + The result must be a list of integers. Each value represents the + compression ratio for the corresponding transformer layer. + """ + 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. @@ -2426,6 +2441,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", @@ -4579,6 +4595,18 @@ def _add_mla_args(parser): 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 output (wo_a). 0 = single linear." + ) + 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', @@ -4612,6 +4640,17 @@ def _add_experimental_attention_variant_args(parser): '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 string containing a Python list expression, e.g.: ' + '"[1,1,4,128,4,128]" or "([1]+[4,128]*2)*3". ' + 'Each value is the compression ratio for the corresponding ' + 'transformer layer (valid values: 1, 4, 128). ' + 'The list length must equal num_layers.' + ) return parser diff --git a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py index 1a0c19d5222..712044f933d 100644 --- a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py +++ b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py @@ -16,12 +16,12 @@ try: from megatron.core.fusions.fused_mla_yarn_rope_apply import ( - fused_apply_mla_rope_for_kv, - fused_apply_mla_rope_for_q, + fused_mla_rope_inplace, + fused_mla_rope_kv_split, ) -except: - fused_apply_mla_rope_for_kv = None - fused_apply_mla_rope_for_q = None +except Exception: + fused_mla_rope_inplace = None + fused_mla_rope_kv_split = None def dtype_tols(dtype): @@ -43,8 +43,8 @@ def rank(self): return 0 -def _test_fused_apply_mla_rope_for_q(input_format): - assert fused_apply_mla_rope_for_q is not None +def _test_fused_mla_rope_inplace(input_format, inverse=False): + assert fused_mla_rope_inplace is not None num_heads = 32 q_dim = 128 emb_dim = 64 @@ -104,12 +104,14 @@ def _test_fused_apply_mla_rope_for_q(input_format): mscale=mscale, cp_group=FakeCPGroup(), mla_rotary_interleaved=True, + inverse=inverse, ) pytorch_output = torch.concat([no_pe, pe_output], dim=-1) pytorch_output.backward(pytorch_bwd_input, retain_graph=True) - fused_output = fused_apply_mla_rope_for_q( - fused_fwd_input, cos, sin, q_dim, emb_dim, cu_seqlens_q=cu_seqlens + fused_output = fused_mla_rope_inplace( + fused_fwd_input, cos, sin, q_dim, emb_dim, + cu_seqlens_q=cu_seqlens, inverse=inverse, ) fused_output.backward(fused_bwd_input, retain_graph=True) @@ -128,8 +130,8 @@ def _test_fused_apply_mla_rope_for_q(input_format): ) -def _test_fused_apply_mla_rope_for_kv(input_format): - assert fused_apply_mla_rope_for_kv is not None +def _test_fused_mla_rope_kv_split(input_format): + assert fused_mla_rope_kv_split is not None num_heads = 32 k_dim = 128 v_dim = 128 @@ -214,7 +216,7 @@ def _test_fused_apply_mla_rope_for_kv(input_format): (pytorch_k_output, pytorch_v_output), (pytorch_bwd_k_input, pytorch_bwd_v_input) ) - fused_k_output, fused_v_output = fused_apply_mla_rope_for_kv( + fused_k_output, fused_v_output = fused_mla_rope_kv_split( fused_fwd_kv_input, fused_fwd_emb_input, cos, @@ -260,12 +262,15 @@ def _test_fused_apply_mla_rope_for_kv(input_format): @pytest.mark.skipif(not is_torch_min_version("2.5.0"), reason="Requires PyTorch >= 2.5.0") @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.parametrize("input_format", ["sbhd", "thd"]) -class TestFusedApplyMLARope: - def test_forward_backward_for_q(self, input_format): - _test_fused_apply_mla_rope_for_q(input_format) +class TestFusedMLARope: + def test_inplace_forward_backward(self, input_format): + _test_fused_mla_rope_inplace(input_format, inverse=False) - def test_forward_backward_for_kv(self, input_format): - _test_fused_apply_mla_rope_for_kv(input_format) + def test_inplace_inverse_forward_backward(self, input_format): + _test_fused_mla_rope_inplace(input_format, inverse=True) + + def test_kv_split_forward_backward(self, input_format): + _test_fused_mla_rope_kv_split(input_format) class TestApplyRotaryPosEmbMlaFusionConflict: From caeff8c80d88e3b61e268282ad62b21b55836a42 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Tue, 28 Apr 2026 05:22:05 +0000 Subject: [PATCH 02/16] add CSA dense mode; add tests; minor fix --- .../experimental_attention_variant/csa.py | 6 +- .../deepseek_v4_hybrid_attention.py | 10 - .../core/transformer/transformer_config.py | 6 +- .../test_attention_variant_csa.py | 864 ++++++++++++++++++ .../test_dsv4_hybrid_attention.py | 407 +++++++++ 5 files changed, 1281 insertions(+), 12 deletions(-) create mode 100644 tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py create mode 100644 tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 9493f8348e4..3c5890fdf96 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -631,7 +631,11 @@ def __init__( self.compressor = None # Conditionally build Indexer (ratio == 4) - if self.compress_ratio == 4 and submodules.indexer is not None: + if ( + self.compress_ratio == 4 + and not config.csa_dense_mode + and submodules.indexer is not None + ): self.indexer = build_module( submodules.indexer, config=config, 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 index 7f4256ba2a4..53aabed1d08 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -1,14 +1,12 @@ # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -import math from dataclasses import dataclass from typing import NoReturn, Optional, Union import torch from megatron.core import tensor_parallel -from megatron.core.dist_checkpointing.mapping import ShardedObject from megatron.core.extensions.transformer_engine import HAVE_TE from megatron.core.models.common.embeddings import ( RotaryEmbedding, @@ -778,14 +776,6 @@ def clip_qk(self): if self.core_attention.current_max_attn_logits is None: raise ValueError("current_max_attn_logits is None") - # Check if we're in absorption mode - if self.cache_mla_latents and not hasattr(self, 'linear_kv_up_proj'): - raise ValueError( - "qk_clip is not supported when cache_mla_latents is enabled and absorption is " - "active. The linear_kv_up_proj layer has been deleted during absorption " - "preparation." - ) - assert self.core_attention.current_max_attn_logits.shape == ( self.num_attention_heads_per_partition, ), f"current_max_attn_logits shape is not ({self.num_attention_heads_per_partition}, ) \ diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index cb2b1be7e71..e084e76b3d6 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -11,7 +11,7 @@ from megatron.core.enums import Fp4Recipe, Fp8Recipe from megatron.core.quantization.quant_config import RecipeConfig -from megatron.core.transformer.enums import AttnBackend, CudaGraphScope, LayerType +from megatron.core.transformer.enums import AttnBackend, CudaGraphScope from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from megatron.core.utils import experimental_api @@ -298,6 +298,10 @@ class TransformerConfig(ModelParallelConfig): csa_compress_rotary_base: float = 40000.0 """RoPE base for compressed KV positions in compressed sparse attention.""" + csa_dense_mode: bool = False + """Whether to use dense mode for compressed sparse attention. If True, the CSA indexer will be + disabled.""" + #################### # linear attention #################### diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py new file mode 100644 index 00000000000..b3fb375cfff --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -0,0 +1,864 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from unittest.mock import patch + +import pytest +import torch + +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.experimental_attention_variant.csa import ( + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, + Compressor, + CompressorSubmodules, + CSAIndexer, + CSAIndexerSubmodules, + get_compress_topk_idxs, + get_window_topk_idxs, + unfused_compressed_sparse_attn, +) +from megatron.core.transformer.transformer_config import MLATransformerConfig +from tests.unit_tests.test_utilities import Utils + + +def mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: + """Mock implementation of hadamard_transform for testing without the library installed.""" + return x * scale + + +@pytest.fixture(autouse=True) +def patch_hadamard_if_needed(): + """Automatically patch hadamard_transform in both dsa and csa modules if 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 + + +# =========================================================================== +# Helper function tests +# =========================================================================== + + +class TestGetWindowTopkIdxs: + """Test get_window_topk_idxs helper.""" + + def test_basic_shape(self): + batch_size, seqlen, window_size = 2, 16, 4 + idxs = get_window_topk_idxs(window_size, batch_size, seqlen, torch.device("cpu")) + assert idxs.shape == (batch_size, seqlen, window_size) + + def test_causal_no_future(self): + """Indices should never exceed the query position.""" + seqlen, window_size = 32, 8 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + for i in range(seqlen): + valid = idxs[0, i][idxs[0, i] >= 0] + assert torch.all(valid <= i), f"Position {i} has future indices" + + def test_invalid_marked_minus_one(self): + """Early positions that cannot fill the window should use -1.""" + seqlen, window_size = 8, 4 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + assert idxs[0, 0, 0] == -1 or idxs[0, 0, 0] == 0 + for pos in range(window_size, seqlen): + assert torch.all(idxs[0, pos] >= 0), f"Position {pos} has invalid -1" + + def test_window_larger_than_seqlen(self): + """Window larger than sequence length should still work.""" + seqlen, window_size = 4, 16 + idxs = get_window_topk_idxs(window_size, 1, seqlen, torch.device("cpu")) + assert idxs.shape == (1, seqlen, window_size) + + +class TestGetCompressTopkIdxs: + """Test get_compress_topk_idxs helper.""" + + def test_basic_shape(self): + ratio, batch_size, seqlen, offset = 4, 2, 32, 32 + idxs = get_compress_topk_idxs(ratio, batch_size, seqlen, offset, torch.device("cpu")) + n_compressed = seqlen // ratio + assert idxs.shape == (batch_size, seqlen, n_compressed) + + def test_offset_applied(self): + """Valid indices should be >= offset.""" + ratio, seqlen, offset = 4, 32, 100 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + valid = idxs[idxs >= 0] + if valid.numel() > 0: + assert torch.all(valid >= offset), "Valid indices should be offset" + + def test_causal_no_future(self): + """Compressed indices should respect causality.""" + ratio, seqlen, offset = 4, 32, 32 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + for i in range(seqlen): + n_valid = (i + 1) // ratio + valid = idxs[0, i][idxs[0, i] >= 0] + assert valid.numel() <= n_valid, f"Position {i} has too many valid compressed indices" + + def test_ratio_128(self): + """Test with large compression ratio.""" + ratio, seqlen, offset = 128, 256, 256 + idxs = get_compress_topk_idxs(ratio, 1, seqlen, offset, torch.device("cpu")) + assert idxs.shape == (1, seqlen, seqlen // ratio) + + +# =========================================================================== +# unfused_compressed_sparse_attn tests +# =========================================================================== + + +class TestUnfusedCompressedSparseAttn: + """Test the unfused compressed sparse attention kernel.""" + + @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() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_output_shape(self): + """Test output shape of unfused compressed sparse attention.""" + sq, b, np_, hn = 16, 2, 4, 64 + n_kv = sq + sq // 4 + topk = 8 + + query = torch.randn(sq, b, np_, hn, dtype=torch.bfloat16).cuda() + kv_full = torch.randn(n_kv, b, hn, dtype=torch.bfloat16).cuda() + attn_sink = torch.zeros(np_, dtype=torch.float32).cuda() + topk_indices = torch.randint(0, n_kv, (b, sq, topk), dtype=torch.int32).cuda() + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn(query, kv_full, attn_sink, topk_indices, softmax_scale) + + assert output.shape == (sq, b, np_ * hn) + assert output.dtype == query.dtype + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_invalid_indices_masked(self): + """Test that -1 indices are properly masked.""" + sq, b, np_, hn = 8, 1, 2, 32 + n_kv = sq + topk = 4 + + query = torch.randn(sq, b, np_, hn, dtype=torch.bfloat16).cuda() + kv_full = torch.randn(n_kv, b, hn, dtype=torch.bfloat16).cuda() + attn_sink = torch.zeros(np_, dtype=torch.float32).cuda() + + topk_indices = torch.full((b, sq, topk), -1, dtype=torch.int32).cuda() + topk_indices[:, :, 0] = 0 + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn(query, kv_full, attn_sink, topk_indices, softmax_scale) + assert not torch.isnan(output).any(), "Output should not contain NaN" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_gradient_flow(self): + """Test that gradients flow through sparse attention.""" + sq, b, np_, hn = 8, 1, 2, 32 + n_kv = sq + topk = 4 + + query = torch.randn(sq, b, np_, hn, dtype=torch.float32).cuda().requires_grad_(True) + kv_full = torch.randn(n_kv, b, hn, dtype=torch.float32).cuda().requires_grad_(True) + attn_sink = torch.nn.Parameter(torch.zeros(np_, dtype=torch.float32).cuda()) + + topk_indices = torch.randint(0, n_kv, (b, sq, topk), dtype=torch.int32).cuda() + softmax_scale = hn**-0.5 + + output = unfused_compressed_sparse_attn(query, kv_full, attn_sink, topk_indices, softmax_scale) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert kv_full.grad is not None + assert attn_sink.grad is not None + + +# =========================================================================== +# Compressor tests +# =========================================================================== + + +def _make_mla_config( + num_layers=4, + hidden_size=256, + num_attention_heads=16, + v_head_dim=64, + qk_pos_emb_head_dim=32, + csa_compress_ratios=None, + csa_window_size=8, + csa_dense_mode=False, + 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, + dsa_indexer_use_sparse_loss=False, +): + """Helper to create MLATransformerConfig for CSA tests.""" + if csa_compress_ratios is None: + csa_compress_ratios = [1] * num_layers + 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, + tensor_model_parallel_size=tensor_model_parallel_size, + sequence_parallel=sequence_parallel, + q_lora_rank=64, + kv_lora_rank=64, + 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, + 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, + csa_dense_mode=csa_dense_mode, + 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, + dsa_indexer_use_sparse_loss=dsa_indexer_use_sparse_loss, + ) + + +def _make_compressor_submodules(): + """Create Compressor submodules spec.""" + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + return CompressorSubmodules( + linear_wkv=ModuleSpec(module=TELinear), + linear_wgate=ModuleSpec(module=TELinear), + norm=ModuleSpec(module=TENorm), + ) + + +def _make_csa_indexer_submodules(): + """Create CSAIndexer submodules spec.""" + from megatron.core.extensions.transformer_engine import TELinear, TENorm + from megatron.core.transformer.spec_utils import ModuleSpec + + return CSAIndexerSubmodules( + linear_wq_b=ModuleSpec(module=TELinear), + linear_weights_proj=ModuleSpec(module=TELinear), + compressor=ModuleSpec( + module=Compressor, + submodules=_make_compressor_submodules(), + ), + ) + + +def _make_csa_submodules(): + """Create CompressedSparseAttention submodules spec.""" + from megatron.core.transformer.spec_utils import ModuleSpec + + return CompressedSparseAttentionSubmodules( + compressor=ModuleSpec( + module=Compressor, + submodules=_make_compressor_submodules(), + ), + indexer=ModuleSpec( + module=CSAIndexer, + submodules=_make_csa_indexer_submodules(), + ), + ) + + +# =========================================================================== +# Compressor tests +# =========================================================================== + + +@pytest.mark.parametrize("compress_ratio", [4, 128]) +class TestCompressor: + """Test Compressor module.""" + + @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(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config( + csa_compress_ratios=[4, 128, 4, 128], + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp'] + ) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_output_shape(self, compress_ratio): + """Test that compressor produces correct output shape.""" + seq_len = 256 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + output = compressor(x) + + expected_len = seq_len // compress_ratio + assert output is not None + assert output.shape == (expected_len, batch_size, head_dim) + assert output.dtype == torch.bfloat16 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_too_short_input(self, compress_ratio): + """Test that compressor returns None when input is shorter than compress_ratio.""" + short_len = compress_ratio - 1 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = torch.randn( + short_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + output = compressor(x) + assert output is None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_compressor_gradient_flow(self, compress_ratio): + """Test that gradients flow through the compressor.""" + seq_len = 256 + batch_size = 2 + head_dim = self.config.v_head_dim + + compressor = Compressor( + config=self.config, + submodules=_make_compressor_submodules(), + compress_ratio=compress_ratio, + head_dim=head_dim, + rotate=False, + rotary_pos_emb=self.rotary_pos_emb, + pg_collection=self.pg_collection, + ).cuda() + + x = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda().requires_grad_(True) + output = compressor(x) + loss = output.sum() + loss.backward() + + assert x.grad is not None + for name, param in compressor.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"Parameter {name} has no gradient" + + +# =========================================================================== +# CSAIndexer tests +# =========================================================================== + + +@pytest.mark.parametrize("seqlen", [32, 128]) +class TestCSAIndexer: + """Test CSAIndexer module basic functionality.""" + + @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(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.compress_ratio = 4 + cls.config = _make_mla_config( + csa_compress_ratios=[4, 4, 4, 4], + dsa_indexer_topk=8, + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp'] + ) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + cls.indexer = CSAIndexer( + config=cls.config, + submodules=_make_csa_indexer_submodules(), + compress_ratio=cls.compress_ratio, + rotary_pos_emb=cls.rotary_pos_emb, + pg_collection=cls.pg_collection, + ) + + yield + Utils.destroy_model_parallel() + + def test_csa_indexer_constructor(self, seqlen): + """Test CSAIndexer initialization.""" + assert isinstance(self.indexer, CSAIndexer) + assert self.indexer.compress_ratio == self.compress_ratio + assert self.indexer.index_n_heads == self.config.dsa_indexer_n_heads + assert self.indexer.index_head_dim == self.config.dsa_indexer_head_dim + assert self.indexer.index_topk == self.config.dsa_indexer_topk + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_forward(self, seqlen): + """Test CSAIndexer forward pass.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + index_scores, topk_indices = self.indexer(x, qr) + n_compressed = seqlen // self.compress_ratio + effective_topk = min(self.config.dsa_indexer_topk, n_compressed) + + assert index_scores.shape == (batch_size, seqlen, n_compressed) + assert topk_indices.shape == (batch_size, seqlen, effective_topk) + assert index_scores.dtype == torch.float32 + assert topk_indices.dtype == torch.long + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_forward_before_topk(self, seqlen): + """Test CSAIndexer forward_before_topk.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + q, k, weights = self.indexer.forward_before_topk(x, qr) + + assert q.shape == (seqlen, batch_size, self.config.dsa_indexer_n_heads, self.config.dsa_indexer_head_dim) + n_compressed = seqlen // self.compress_ratio + assert k.shape == (n_compressed, batch_size, self.config.dsa_indexer_head_dim) + assert weights.shape == (seqlen, batch_size, self.config.dsa_indexer_n_heads) + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_csa_indexer_with_mask(self, seqlen): + """Test CSAIndexer with causal mask.""" + batch_size = 2 + self.indexer.cuda() + + x = torch.randn(seqlen, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seqlen, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + n_compressed = seqlen // self.compress_ratio + causal_mask = torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(seqlen, -1) + positions = torch.arange(1, seqlen + 1, device=x.device).unsqueeze(1) + causal_mask = torch.where( + causal_mask >= positions // self.compress_ratio, + float("-inf"), + 0.0, + ).unsqueeze(0).expand(batch_size, -1, -1) + + index_scores, topk_indices = self.indexer(x, qr, mask=causal_mask) + + effective_topk = min(self.config.dsa_indexer_topk, n_compressed) + assert index_scores.shape == (batch_size, seqlen, n_compressed) + assert topk_indices.shape == (batch_size, seqlen, effective_topk) + + +# =========================================================================== +# CompressedSparseAttention tests +# =========================================================================== + + +class TestCompressedSparseAttentionRatio1: + """Test CompressedSparseAttention with compress_ratio=1 (window-only).""" + + @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(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config( + csa_compress_ratios=[1, 1, 1, 1], + csa_window_size=8, + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp'] + ) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + cls.csa = CompressedSparseAttention( + config=cls.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=cls.pg_collection, + rotary_pos_emb=rotary_pos_emb, + ) + + yield + Utils.destroy_model_parallel() + + def test_ratio1_no_compressor(self): + """With ratio=1, compressor and indexer should not be built.""" + assert self.csa.compressor is None + assert self.csa.indexer is None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_ratio1_forward(self): + """Test forward pass with window-only attention.""" + seq_len = 32 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + self.csa.cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = self.csa( + query=query, key=key, value=value, attention_mask=None, x=x, qr=qr, + ) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert output.dtype == torch.bfloat16 + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_ratio1_backward(self): + """Test backward pass with window-only attention.""" + seq_len = 32 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + self.csa.train() + self.csa.cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.float32).cuda().requires_grad_(True) + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.float32).cuda().requires_grad_(True) + value = key.clone().detach().requires_grad_(True) + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = self.csa( + query=query, key=key, value=value, attention_mask=None, x=x, qr=qr, + ) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert key.grad is not None + + +@pytest.mark.parametrize("compress_ratio", [4, 128]) +class TestCompressedSparseAttentionCompressed: + """Test CompressedSparseAttention with compress_ratio > 1.""" + + @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(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config( + csa_compress_ratios=[4, 128, 4, 128], + csa_window_size=8, + dsa_indexer_topk=8, + dsa_indexer_loss_coeff=1.0, + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp'] + ) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + def _get_layer_number(self, compress_ratio): + """Return a layer_number (1-indexed) whose compress_ratio matches.""" + for i, r in enumerate(self.config.csa_compress_ratios): + if r == compress_ratio: + return i + 1 + raise ValueError(f"No layer with compress_ratio={compress_ratio}") + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_constructor(self, compress_ratio): + """Test that compressor/indexer are conditionally built.""" + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + ).cuda() + + assert csa.compressor is not None + if compress_ratio == 4: + assert csa.indexer is not None + elif compress_ratio == 128: + assert csa.indexer is None + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_forward(self, compress_ratio): + """Test forward pass with compressed attention.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + ).cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_backward(self, compress_ratio): + """Test backward pass with compressed attention.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + ).cuda() + csa.train() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.float32).cuda().requires_grad_(True) + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.float32).cuda().requires_grad_(True) + value = key.clone().detach().requires_grad_(True) + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + loss = output.sum() + loss.backward() + + assert query.grad is not None + assert key.grad is not None + + for name, param in csa.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"Parameter {name} has no gradient" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_eval_mode(self, compress_ratio): + """Test forward pass in eval mode.""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + layer_number = self._get_layer_number(compress_ratio) + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=layer_number, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + ).cuda() + csa.eval() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + with torch.no_grad(): + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any() + + +# =========================================================================== +# csa_dense_mode tests +# =========================================================================== + + +class TestCompressedSparseAttentionDenseMode: + """Test that csa_dense_mode=True disables the indexer for ratio=4 layers.""" + + @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(123) + model_parallel_cuda_manual_seed(123) + + cls = request.cls + cls.config = _make_mla_config( + csa_compress_ratios=[4, 128, 4, 128], + csa_window_size=8, + csa_dense_mode=True, + ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups( + required_pgs=['tp', 'cp'] + ) + + from megatron.core.models.common.embeddings import RotaryEmbedding + + cls.rotary_pos_emb = RotaryEmbedding( + cls.config.qk_pos_emb_head_dim, + rotary_percent=cls.config.rotary_percent, + rotary_base=cls.config.rotary_base, + cp_group=cls.pg_collection.cp, + ) + + yield + Utils.destroy_model_parallel() + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dense_mode_disables_indexer_for_ratio4(self): + """With csa_dense_mode=True, ratio=4 layers should NOT build an indexer.""" + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + ).cuda() + + assert csa.compress_ratio == 4 + assert csa.compressor is not None, "Compressor should still be built" + assert csa.indexer is None, "Indexer should be disabled in dense mode" + + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_dense_mode_forward_ratio4(self): + """Forward pass should work for ratio=4 in dense mode (uses all compressed positions).""" + seq_len = 256 + batch_size = 2 + np_ = self.config.num_attention_heads + hn = self.config.v_head_dim + + csa = CompressedSparseAttention( + config=self.config, + submodules=_make_csa_submodules(), + layer_number=1, + attn_mask_type=AttnMaskType.causal, + attention_type='self', + pg_collection=self.pg_collection, + rotary_pos_emb=self.rotary_pos_emb, + ).cuda() + + query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() + key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.bfloat16).cuda() + value = key.clone() + x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() + qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() + + output = csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) + + assert output.shape == (seq_len, batch_size, np_ * hn) + assert not torch.isnan(output).any() 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..87a1a61e652 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py @@ -0,0 +1,407 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. + +from unittest.mock import patch + +import pytest +import torch + +import megatron.core.parallel_state as parallel_state +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, +): + """Create an MLATransformerConfig for DSv4 hybrid attention tests.""" + if csa_compress_ratios is None: + csa_compress_ratios = [1, 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, + 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, + ) + + +def _make_attention_spec(config): + """Build the full DSv4HybridSelfAttention ModuleSpec using the canonical spec builder.""" + from megatron.core.models.gpt.experimental_attention_variant_module_specs import ( + get_dsv4_hybrid_module_spec_for_backend, + ) + + return get_dsv4_hybrid_module_spec_for_backend(config=config) + + +def _build_attention(config, layer_number, pg_collection): + """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) + + +# =========================================================================== +# 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 across TP sizes.""" + + @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_up_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') + assert attn.linear_kv_down_proj is None + + 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 + + @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 = [1, 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 From 0ff05b7623a634d44b4e544f003740bae264bf05 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Tue, 28 Apr 2026 05:39:28 +0000 Subject: [PATCH 03/16] update ratio args --- .../transformer/experimental_attention_variant/csa.py | 2 +- megatron/core/transformer/transformer_config.py | 6 +++--- megatron/training/arguments.py | 8 ++++---- .../test_attention_variant_csa.py | 4 ++-- .../test_dsv4_hybrid_attention.py | 2 +- 5 files changed, 11 insertions(+), 11 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 3c5890fdf96..b895a1421b2 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -570,7 +570,7 @@ class CompressedSparseAttention(MegatronModule): provides compressor and indexer submodule specs; this ``__init__`` inspects ``config.csa_compress_ratios[layer_idx]`` and conditionally builds them: - * ``ratio == 1``: window-only (compressor and indexer NOT built) + * ``ratio == 0``: window-only (compressor and indexer NOT built) * ``ratio == 4``: window + 4x compressed + learned Indexer (both built) * ``ratio == 128``: window + 128x compressed, attend to all (compressor built only) """ diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index e084e76b3d6..e15171f5bc5 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -293,7 +293,7 @@ class TransformerConfig(ModelParallelConfig): """Sliding window size for compressed sparse attention.""" csa_compress_ratios: Optional[List[int]] = None - """Per-layer compress ratios, e.g. [1, 1, 4, 128, 4, 128, ...].""" + """Per-layer compress ratios, e.g. [0, 0, 4, 128, 4, 128, ...].""" csa_compress_rotary_base: float = 40000.0 """RoPE base for compressed KV positions in compressed sparse attention.""" @@ -1286,8 +1286,8 @@ def __post_init__(self): len(self.csa_compress_ratios) == self.num_layers ), "csa_compress_ratios must have the same length as num_layers" assert ( - all(ratio in [1, 4, 128] for ratio in self.csa_compress_ratios) - ), "csa_compress_ratios must be 1, 4, or 128" + all(ratio in [0, 4, 128] for ratio in self.csa_compress_ratios) + ), "csa_compress_ratios must be 0, 4, or 128" if self.fp8: # cannot support first last layer bf16 with delayed scaling diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index 9bf0b451a8c..6a108a0d6d0 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -275,8 +275,8 @@ def compress_ratios_type(x): """Per-layer compress ratios for compressed sparse attention. Accepts a string containing a Python list expression, e.g.: - "[1,1,4,128,4,128]" - "([1]+[4,128]*2)*3" + "[0,0,4,128,4,128]" + "([0]+[4,128]*2)*3" The result must be a list of integers. Each value represents the compression ratio for the corresponding transformer layer. """ @@ -4646,9 +4646,9 @@ def _add_experimental_attention_variant_args(parser): default=None, help='Per-layer compress ratios for compressed sparse attention. ' 'Accepts a string containing a Python list expression, e.g.: ' - '"[1,1,4,128,4,128]" or "([1]+[4,128]*2)*3". ' + '"[0,0,4,128,4,128]" or "([0]+[4,128]*2)*3". ' 'Each value is the compression ratio for the corresponding ' - 'transformer layer (valid values: 1, 4, 128). ' + 'transformer layer (valid values: 0, 4, 128). ' 'The list length must equal num_layers.' ) return parser diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py index b3fb375cfff..9ba05bab3cf 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -212,7 +212,7 @@ def _make_mla_config( ): """Helper to create MLATransformerConfig for CSA tests.""" if csa_compress_ratios is None: - csa_compress_ratios = [1] * num_layers + csa_compress_ratios = [0] * num_layers return MLATransformerConfig( num_layers=num_layers, hidden_size=hidden_size, @@ -533,7 +533,7 @@ def setup_method(self, request): cls = request.cls cls.config = _make_mla_config( - csa_compress_ratios=[1, 1, 1, 1], + csa_compress_ratios=[0, 0, 0, 0], csa_window_size=8, ) cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups( 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 index 87a1a61e652..0d62f4b6a5a 100644 --- 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 @@ -68,7 +68,7 @@ def _make_config( ): """Create an MLATransformerConfig for DSv4 hybrid attention tests.""" if csa_compress_ratios is None: - csa_compress_ratios = [1, 4, 128, 4] + csa_compress_ratios = [0, 4, 128, 4] return MLATransformerConfig( num_layers=num_layers, hidden_size=hidden_size, From 0b7af8280446c8885c90c9a6439973bf4b495be6 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Wed, 29 Apr 2026 07:50:01 +0000 Subject: [PATCH 04/16] fix MTP support with dsv4_hybrid; fix test error --- megatron/core/transformer/attention.py | 6 ++++++ .../experimental_attention_variant/csa.py | 4 ++-- .../deepseek_v4_hybrid_attention.py | 15 +++++++++++++-- .../core/transformer/multi_latent_attention.py | 6 ++++++ megatron/core/transformer/transformer_config.py | 9 ++++++--- megatron/core/transformer/transformer_layer.py | 4 ++++ .../test_attention_variant_csa.py | 8 ++++++++ .../test_dsv4_hybrid_attention.py | 2 +- 8 files changed, 46 insertions(+), 8 deletions(-) diff --git a/megatron/core/transformer/attention.py b/megatron/core/transformer/attention.py index 2ff3eacc071..1a893096812 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -258,12 +258,14 @@ def __init__( cp_comm_type: str | None = None, pg_collection: ProcessGroupCollection | None = None, pp_layer_offset: Optional[int] = None, + is_mtp_layer: bool = False, ): super().__init__(config=config) 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 @@ -1357,6 +1359,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, ): super().__init__( config=config, @@ -1367,6 +1370,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, ) self.linear_qkv_out_dim = self.query_projection_size + 2 * self.kv_projection_size @@ -1766,6 +1770,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, ): super().__init__( config=config, @@ -1775,6 +1780,7 @@ def __init__( attention_type="cross", cp_comm_type=cp_comm_type, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, ) if self.config.num_query_groups != self.config.num_attention_heads: diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index b895a1421b2..381c4796720 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -589,6 +589,7 @@ def __init__( cp_comm_type: str = "p2p", pg_collection: Optional[ProcessGroupCollection] = None, rotary_pos_emb: nn.Module = None, + compress_ratio: int = 0, ): super().__init__(config=config) @@ -599,8 +600,7 @@ def __init__( self.pg_collection = pg_collection self.layer_number = layer_number - layer_idx = layer_number - 1 # layer_number is 1-indexed - self.compress_ratio = config.csa_compress_ratios[layer_idx] + self.compress_ratio = compress_ratio self.window_size = config.csa_window_size self.v_head_dim = config.v_head_dim 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 index 53aabed1d08..cb3ffc9582b 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -96,6 +96,7 @@ def __init__( attention_type: str, cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, + is_mtp_layer: bool = False, ) -> None: super().__init__( @@ -105,6 +106,7 @@ def __init__( attention_type=attention_type, attn_mask_type=attn_mask_type, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, ) self.config: MLATransformerConfig @@ -123,8 +125,12 @@ def __init__( 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] rope_base = self.config.rotary_base - compress_ratio = self.config.csa_compress_ratios[layer_number - 1] if compress_ratio > 1: rope_base = self.config.csa_compress_rotary_base if self.config.rope_type == "rope": @@ -152,7 +158,10 @@ def __init__( "'rope' and 'yarn'" ) - core_attn_extra_kwargs = {"rotary_pos_emb": self.rotary_pos_emb} + core_attn_extra_kwargs = { + "rotary_pos_emb": self.rotary_pos_emb, + "compress_ratio": compress_ratio, + } self.core_attention = build_module( submodules.core_attention, config=self.config, @@ -416,6 +425,7 @@ def __init__( attn_mask_type=AttnMaskType.padding, cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, + is_mtp_layer: bool = False, ): if pg_collection is None: pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -428,6 +438,7 @@ def __init__( attention_type="self", cp_comm_type=cp_comm_type, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, ) q_down_proj_kwargs = {} diff --git a/megatron/core/transformer/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 60841989994..c8f00084c5d 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -139,6 +139,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, ) -> None: super().__init__( @@ -149,6 +150,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, ) self.config: MLATransformerConfig @@ -472,6 +474,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, ): if pg_collection is None: pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -485,6 +488,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, ) if self.config.q_lora_rank is None: @@ -1210,6 +1214,7 @@ def __init__( attn_mask_type=AttnMaskType.padding, cp_comm_type: Optional[str] = None, pg_collection: Optional[ProcessGroupCollection] = None, + is_mtp_layer: bool = False, ): if pg_collection is None: pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -1223,6 +1228,7 @@ def __init__( attention_type="self", cp_comm_type=cp_comm_type, pg_collection=pg_collection, + is_mtp_layer=is_mtp_layer, ) assert self.config.q_lora_rank is not None, ( diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index e15171f5bc5..2e8af3489f5 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1282,9 +1282,12 @@ def __post_init__(self): 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" - assert ( - len(self.csa_compress_ratios) == self.num_layers - ), "csa_compress_ratios must have the same length as num_layers" + 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" diff --git a/megatron/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index 5392b87f7dd..271744b57a3 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -340,6 +340,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( @@ -1472,6 +1474,7 @@ def __init__( hidden_dropout: Optional[float] = None, pg_collection: Optional[ProcessGroupCollection] = None, vp_stage: Optional[int] = None, + is_mtp_layer: bool = False, ): super().__init__( config=config, @@ -1480,6 +1483,7 @@ def __init__( hidden_dropout=hidden_dropout, pg_collection=pg_collection, vp_stage=vp_stage, + is_mtp_layer=is_mtp_layer, ) if submodules.cross_attention_hyper_connection is not IdentityOp: diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py index 9ba05bab3cf..b540a621192 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -22,6 +22,14 @@ 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 + def mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor: """Mock implementation of hadamard_transform for testing without the library installed.""" 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 index 0d62f4b6a5a..98560c599d6 100644 --- 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 @@ -175,7 +175,7 @@ def test_rope_base_varies_with_compress_ratio(self, layer_number): torch.manual_seed(_SEED) model_parallel_cuda_manual_seed(_SEED) - ratios = [1, 4, 128, 4] + 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) From 0a74b71ee66cedaab1f3cd1d360dfee00a776680 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Wed, 29 Apr 2026 14:21:23 +0000 Subject: [PATCH 05/16] remove redundant codes and add some comments --- .../models/common/embeddings/rope_utils.py | 1 + ...rimental_attention_variant_module_specs.py | 2 +- .../experimental_attention_variant/csa.py | 16 +- .../deepseek_v4_hybrid_attention.py | 265 ++++-------------- .../core/transformer/transformer_config.py | 10 +- .../test_dsv4_hybrid_attention.py | 3 +- 6 files changed, 64 insertions(+), 233 deletions(-) diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index 0608abb040a..c558260f03f 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -147,6 +147,7 @@ def _apply_rotary_pos_emb_bshd( # Fallback to original permutation # DSv4 applies rope on V and O, so we need to uninterleave the tensor. + # The existing MLA code is safe because the dot product is permutation-invariant. if mla_rotary_interleaved: x1, x2 = torch.chunk(t, 2, dim=-1) t = torch.stack((x1, x2), dim=-1).flatten(start_dim=-2) diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py index 984a9d4ea7f..55aec0cd419 100644 --- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -188,7 +188,7 @@ def get_dsv4_hybrid_module_spec_for_backend( submodules=DSv4HybridSelfAttentionSubmodules( linear_q_down_proj=backend.linear(), linear_q_up_proj=backend.column_parallel_linear(), - linear_kv_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, diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 381c4796720..3e81d29dc13 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -15,7 +15,6 @@ ) from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region from megatron.core.transformer.enums import AttnMaskType from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexerLossAutoScaler, @@ -27,7 +26,7 @@ from megatron.core.transformer.module import MegatronModule from megatron.core.transformer.spec_utils import ModuleSpec, build_module from megatron.core.transformer.transformer_config import TransformerConfig -from megatron.core.utils import get_pg_size, nvtx_range_pop, nvtx_range_push +from megatron.core.utils import nvtx_range_pop, nvtx_range_push # --------------------------------------------------------------------------- @@ -348,9 +347,6 @@ def forward(self, x: torch.Tensor) -> Optional[torch.Tensor]: """ nvtx_range_push("compressor") - if self.config.sequence_parallel and get_pg_size(self.pg_collection.tp) > 1: - x = gather_from_sequence_parallel_region(x, group=self.pg_collection.tp) - sq, b, _ = x.size() ratio = self.compress_ratio @@ -500,11 +496,6 @@ def forward_before_topk( """Compute Q, compressed K, and weights before top-k selection.""" nvtx_range_push("indexer_before_topk") - # Gather for SP - if self.config.sequence_parallel and get_pg_size(self.pg_collection.tp) > 1: - x = gather_from_sequence_parallel_region(x, group=self.pg_collection.tp) - qr = gather_from_sequence_parallel_region(qr, group=self.pg_collection.tp) - sq, bsz, _ = x.size() # Q path @@ -604,8 +595,7 @@ def __init__( self.window_size = config.csa_window_size self.v_head_dim = config.v_head_dim - n_local_heads = config.num_attention_heads // get_pg_size(pg_collection.tp) - self.n_local_heads = n_local_heads + self.n_local_heads = config.num_attention_heads if softmax_scale is None: softmax_scale = config.v_head_dim ** -0.5 @@ -614,7 +604,7 @@ def __init__( self.force_unfused_dsa = getattr(config, 'force_unfused_dsa', True) # Learnable attention sink per head - self.attn_sink = nn.Parameter(torch.zeros(n_local_heads, dtype=torch.float32)) + self.attn_sink = nn.Parameter(torch.zeros(self.n_local_heads, dtype=torch.float32)) # Conditionally build Compressor (ratio > 1) if self.compress_ratio > 1 and submodules.compressor is not None: 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 index cb3ffc9582b..e8e004cb3f1 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -18,10 +18,6 @@ ) from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.tensor_parallel.layers import ColumnParallelLinear -from megatron.core.tensor_parallel.mappings import ( - gather_from_tensor_model_parallel_region, - scatter_to_sequence_parallel_region, -) from megatron.core.transformer.attention import Attention from megatron.core.transformer.enums import AttnMaskType from megatron.core.transformer.spec_utils import ModuleSpec, build_module @@ -34,18 +30,13 @@ ) try: - from megatron.core.fusions.fused_mla_yarn_rope_apply import ( - fused_mla_rope_inplace, - fused_mla_rope_kv_split, - ) + from megatron.core.fusions.fused_mla_yarn_rope_apply import fused_mla_rope_inplace except Exception: fused_mla_rope_inplace = None - fused_mla_rope_kv_split = None if HAVE_TE: from megatron.core.extensions.transformer_engine import ( - TEColumnParallelLinear, TELinear, set_save_original_input, ) @@ -53,12 +44,10 @@ else: ( TEColumnParallelLinear, - TELayerNormColumnParallelLinear, TELinear, Linear, set_save_original_input, - split_te_layernorm_column_parallel_linear, - ) = (None, None, None, None, None, None) + ) = (None, None, None, None) @torch.compile @@ -77,8 +66,7 @@ class DSv4HybridSelfAttentionSubmodules: linear_q_down_proj: Union[ModuleSpec, type] = None linear_q_up_proj: Union[ModuleSpec, type] = None - linear_kv_up_proj: Union[ModuleSpec, type] = None - linear_qkv_down_proj: Union[ModuleSpec, type] = None + linear_kv_proj: Union[ModuleSpec, type] = None core_attention: Union[ModuleSpec, type] = None linear_proj: Union[ModuleSpec, type] = None @@ -110,6 +98,17 @@ def __init__( ) 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 @@ -177,11 +176,7 @@ def __init__( ) # Output. - tp_size = get_pg_size(self.pg_collection.tp) - assert self.config.o_groups % tp_size == 0, ( - "o_groups must be divisible by tp_size" - ) - self.o_local_groups = self.config.o_groups // tp_size + 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" ) @@ -189,11 +184,11 @@ def __init__( 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, - ) + 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) @@ -265,16 +260,13 @@ def forward( # ===================== # Get the query, key and value tensors based on the type of attention - # self or cross attn. - qkv_linear_manager = off_interface(self.offload_qkv_linear, hidden_states, "qkv_linear") - with qkv_linear_manager as hidden_states: - query, key, value, q_compressed, kv_compressed = self.get_query_key_value_tensors( - hidden_states, - key_value_states, - position_ids, - packed_seq_params, - inference_context=inference_context, - ) - query = qkv_linear_manager.group_offload(query, forced_released_tensors=[hidden_states]) + query, key, value, q_compressed, kv_compressed = self.get_query_key_value_tensors( + hidden_states, + key_value_states, + position_ids, + packed_seq_params, + inference_context=inference_context, + ) # TODO: Currently, TE can only accept contiguous tensors for MLA query = query.contiguous() @@ -288,29 +280,19 @@ def forward( core_attn_manager = off_interface( self.offload_core_attention and self.training, query, "core_attn" ) - if self.checkpoint_core_attention and self.training: - core_attn_out = self._checkpointed_attention_forward( - query, key, value, attention_mask, packed_seq_params=packed_seq_params - ) - else: - extra_kwargs = {} - if self.config.experimental_attention_variant in ("dsa", "dsv4_hybrid"): - # For dsa we need to pass in the original hidden states and the compressed - # query representation. - extra_kwargs["x"] = hidden_states - extra_kwargs["qr"] = q_compressed - with core_attn_manager as query: - core_attn_out = self.core_attention( - query, - key, - value, - attention_mask, - packed_seq_params=packed_seq_params, - **extra_kwargs, - ) - core_attn_out = core_attn_manager.group_offload( - core_attn_out, forced_released_tensors=[query, key, value] + with core_attn_manager as query: + core_attn_out = self.core_attention( + query, + key, + value, + attention_mask, + packed_seq_params=packed_seq_params, + x=hidden_states, + qr=q_compressed, ) + core_attn_out = core_attn_manager.group_offload( + core_attn_out, forced_released_tensors=[query, key, value] + ) if packed_seq_params is not None and packed_seq_params.qkv_format == 'thd': # reshape to same output shape as unpacked case @@ -337,7 +319,7 @@ def forward( if packed_seq_params.cu_seqlens_kv_padded is not None else packed_seq_params.cu_seqlens_kv ) - rope_seqlen = cu_seqlens_kv[-1].item() + rope_seqlen = cu_seqlens_kv else: cu_seqlens_kv = None rope_seqlen = seq_len @@ -444,12 +426,6 @@ def __init__( q_down_proj_kwargs = {} if submodules.linear_q_down_proj in [TELinear]: q_down_proj_kwargs['parallel_mode'] = 'duplicated' - elif submodules.linear_q_down_proj in [ - Linear, - TEColumnParallelLinear, - ColumnParallelLinear, - ]: - q_down_proj_kwargs['gather_output'] = False else: raise ValueError(f"Unsupported linear_q_down_proj: {submodules.linear_q_down_proj}") @@ -464,11 +440,7 @@ def __init__( is_expert=False, tp_comm_buffer_name='q_down_proj', skip_weight_param_allocation=False, - tp_group=( - pg_collection.tp - if q_down_proj_kwargs.get('parallel_mode') != 'duplicated' - else None - ), + tp_group=None, **q_down_proj_kwargs, ) @@ -486,9 +458,8 @@ def __init__( tp_group=pg_collection.tp, ) - self.linear_kv_down_proj = None - self.linear_kv_up_proj = build_module( - submodules.linear_kv_up_proj, + self.linear_kv_proj = build_module( + submodules.linear_kv_proj, self.config.hidden_size, self.config.v_head_dim, config=self.config, @@ -562,7 +533,6 @@ def get_query_key_value_tensors( assert inference_context is None, "Inference with MLA RoPE fusion is not supported" assert ( fused_mla_rope_inplace is not None - and fused_mla_rope_kv_split is not None ), "Fused MLA RoPE apply is not imported successfully" else: rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) @@ -582,22 +552,9 @@ def get_query_key_value_tensors( # ========================================= # QKV down projection and layernorm # ========================================= - # if linear_q_down_proj is ColumnParallelLinear: - # q_compressed: [s, b, q_lora_rank / TP] - # elif linear_q_down_proj is Linear: - # q_compressed: [s / TP, b, q_lora_rank] + # q_compressed: [s, b, q_lora_rank] q_compressed, _ = self.linear_q_down_proj(hidden_states) - # When output is sharded (ColumnParallelLinear), two things are needed to be - # identical to a normal Linear. - # 1. Manually gather output to restore output dim q_lora_rank; - # 2. Scatter sequence back to s / TP if sequence-parallel since it was - # gathered by ColumnParallelLinear. - if q_compressed.size(-1) != self.config.q_lora_rank: - q_compressed = gather_from_tensor_model_parallel_region(q_compressed) - if self.config.sequence_parallel: - q_compressed = scatter_to_sequence_parallel_region(q_compressed) - kv_compressed = hidden_states k_pos_emb = None @@ -634,14 +591,13 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po 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_up_proj(kv_compressed) + 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) - # todo add assert about fusions and caching if self.config.apply_rope_fusion: cp_rank = self.pg_collection.cp.rank() cp_size = self.pg_collection.cp.size() @@ -749,17 +705,12 @@ def backward_dw(self) -> NoReturn: def _backward_kv_proj(self): """Computes weight gradients of KV projection layers""" - self.linear_kv_up_proj.backward_dw() - if self.linear_kv_down_proj is not None: - self.linear_kv_down_proj.backward_dw() + self.linear_kv_proj.backward_dw() def _backward_q_proj(self): """Computes weight gradients of Q projection layers""" - if self.config.q_lora_rank is None: - self.linear_q_proj.backward_dw() - else: - self.linear_q_down_proj.backward_dw() - self.linear_q_up_proj.backward_dw() + 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""" @@ -767,121 +718,5 @@ def _backward_output_proj(self): def set_for_recompute_input_layernorm(self): """Set the attention layer for recompute input_layernorm. Only needed for fp8/fp4.""" - if self.config.q_lora_rank is not None: - set_save_original_input(self.linear_q_down_proj) - if self.linear_kv_down_proj is not None: - set_save_original_input(self.linear_kv_down_proj) - elif not self.recompute_up_proj: - set_save_original_input(self.linear_kv_up_proj) - - def clip_qk(self): - """ - QK Clipping is a technique to clip the query and key attention logits to prevent the - attention logits from exploding. Per MuonClip usage, we update the weight by calling this - function after Muon optimizer step. - """ - - if not self.config.qk_clip: - raise ValueError("qk_clip option needs to be enabled") - - if self.core_attention.current_max_attn_logits is None: - raise ValueError("current_max_attn_logits is None") - - assert self.core_attention.current_max_attn_logits.shape == ( - self.num_attention_heads_per_partition, - ), f"current_max_attn_logits shape is not ({self.num_attention_heads_per_partition}, ) \ - but {self.core_attention.current_max_attn_logits.shape}" - - # only update the weight if any head has - # current_max_attn_logits > qk_clip_threshold - if torch.any(self.core_attention.current_max_attn_logits > self.config.qk_clip_threshold): - # Use num_attention_heads_per_partition for tensor parallel scenarios - - # qk_clip_balancing_eta (n, 1, 1) - assert self.core_attention.current_max_attn_logits.shape == ( - self.num_attention_heads_per_partition, - ), f"current_max_attn_logits shape is not ({self.num_attention_heads_per_partition},) \ - but {self.core_attention.current_max_attn_logits.shape}" - self.qk_clip_balancing_eta = torch.clamp( - self.config.qk_clip_threshold / self.core_attention.current_max_attn_logits, max=1.0 - ).view(self.num_attention_heads_per_partition, 1, 1) - assert torch.all(self.qk_clip_balancing_eta <= 1.0) - - # Update q side weight, keep qk_pos_emb_head_dim side weight unchanged - if self.config.q_lora_rank is None: - q_proj_weight = self.linear_q_proj.weight - else: - q_proj_weight = self.linear_q_up_proj.weight - - # Handle different weight access patterns (main_param vs direct access) - if hasattr(q_proj_weight, 'main_param'): - q_proj_weight.main_param.data.copy_( - self._clip_q_proj_weight(q_proj_weight.main_param.data) - ) - q_proj_weight.data.copy_(self._clip_q_proj_weight(q_proj_weight.data)) - - # Update k side weight, keep v side weight unchanged - kv_proj_weight = self.linear_kv_up_proj.weight - - # Handle different weight access patterns - if hasattr(kv_proj_weight, 'main_param'): - kv_proj_weight.main_param.data.copy_( - self._clip_kv_proj_weight(kv_proj_weight.main_param.data) - ) - kv_proj_weight.data.copy_(self._clip_kv_proj_weight(kv_proj_weight.data)) - - # reset current_max_attn_logits - self.core_attention.current_max_attn_logits = None - - def _clip_q_proj_weight(self, weight): - """Clip q_proj_weight""" - # Reshape to (n, a + b, -1) - weight_reshaped = weight.view( - self.num_attention_heads_per_partition, - self.config.qk_head_dim + self.config.qk_pos_emb_head_dim, - -1, - ) - - # Split into qk_head_dim and qk_pos_emb_head_dim parts: (n, a, -1) and (n, b, -1) - weight_q_nope = weight_reshaped[:, : self.config.qk_head_dim, :] - weight_q_pe = weight_reshaped[:, self.config.qk_head_dim :, :] - - # Clipping - weight_q_nope.mul_(torch.pow(self.qk_clip_balancing_eta, self.config.qk_clip_alpha)) - weight_q_pe.mul_(self.qk_clip_balancing_eta) - - # Concatenate back and reshape to original shape - weight_q_updated = torch.cat([weight_q_nope, weight_q_pe], dim=1) - weight_q_updated = weight_q_updated.view( - self.num_attention_heads_per_partition - * (self.config.qk_head_dim + self.config.qk_pos_emb_head_dim), - -1, - ) - - return weight_q_updated - - def _clip_kv_proj_weight(self, weight): - """Clip kv_proj_weight""" - # shape: (n, qk_head_dim + v_head_dim, kv_lora_rank) - weight_reshaped = weight.view( - self.num_attention_heads_per_partition, - self.config.qk_head_dim + self.config.v_head_dim, - -1, - ) - - # Split into qk_head_dim and v_head_dim parts: (n, a, -1) and (n, b, -1) - weight_k = weight_reshaped[:, : self.config.qk_head_dim, :] - weight_v = weight_reshaped[:, self.config.qk_head_dim :, :] - - # Clipping - weight_k.mul_(torch.pow(self.qk_clip_balancing_eta, 1 - self.config.qk_clip_alpha)) - - # Concatenate back and reshape to original shape - weight_kv_updated = torch.cat([weight_k, weight_v], dim=1) - weight_kv_updated = weight_kv_updated.view( - self.num_attention_heads_per_partition - * (self.config.qk_head_dim + self.config.v_head_dim), - -1, - ) - - return weight_kv_updated + set_save_original_input(self.linear_q_down_proj) + set_save_original_input(self.linear_kv_proj) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 2e8af3489f5..4b6dba27f20 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -1291,6 +1291,10 @@ def __post_init__(self): 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 not self.qk_clip, "QK clipping is not supported with DSv4 Hybrid Attention." if self.fp8: # cannot support first last layer bf16 with delayed scaling @@ -2724,10 +2728,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.""" 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 index 98560c599d6..2449c45328b 100644 --- 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 @@ -150,13 +150,12 @@ def test_basic_construction(self): assert isinstance(attn, DSv4HybridSelfAttention) assert hasattr(attn, 'linear_q_down_proj') assert hasattr(attn, 'linear_q_up_proj') - assert hasattr(attn, 'linear_kv_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') - assert attn.linear_kv_down_proj is None def test_q_head_dim_equals_v_head_dim(self): """q_head_dim must equal v_head_dim for DSv4 hybrid.""" From 4aee0cae3674c6bd293bb53927f41ee3812c7f3a Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Wed, 29 Apr 2026 15:19:47 +0000 Subject: [PATCH 06/16] fix lint --- .../core/fusions/fused_mla_yarn_rope_apply.py | 16 +- ...rimental_attention_variant_module_specs.py | 3 +- .../experimental_attention_variant/csa.py | 58 +++----- .../deepseek_v4_hybrid_attention.py | 67 ++++----- .../experimental_attention_variant/dsa.py | 8 +- .../core/transformer/transformer_config.py | 16 +- .../fusions/test_mla_yarn_rope_apply.py | 3 +- .../test_attention_variant_csa.py | 137 +++++++++--------- .../test_dsv4_hybrid_attention.py | 36 +++-- 9 files changed, 164 insertions(+), 180 deletions(-) diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py index 90acee5814b..9529cbc2a3f 100644 --- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py +++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py @@ -107,8 +107,12 @@ def _mla_rope_fwd_inplace_kernel( cos_left = tl.load(COS + token_idx * stride_cos_seq + tl.arange(0, emb_dim // 2)) sin_left = tl.load(SIN + token_idx * stride_sin_seq + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * stride_cos_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * stride_sin_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_right = tl.load( + COS + token_idx * stride_cos_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + sin_right = tl.load( + SIN + token_idx * stride_sin_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) if INVERSE: sin_left = -sin_left sin_right = -sin_right @@ -189,8 +193,12 @@ def _mla_rope_bwd_inplace_kernel( cos_left = tl.load(COS + token_idx * stride_cos_seq + tl.arange(0, emb_dim // 2)) sin_left = tl.load(SIN + token_idx * stride_sin_seq + tl.arange(0, emb_dim // 2)) - cos_right = tl.load(COS + token_idx * stride_cos_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2)) - sin_right = tl.load(SIN + token_idx * stride_sin_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2)) + cos_right = tl.load( + COS + token_idx * stride_cos_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) + sin_right = tl.load( + SIN + token_idx * stride_sin_seq + emb_dim // 2 + tl.arange(0, emb_dim // 2) + ) if INVERSE: sin_left = -sin_left sin_right = -sin_right diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py index 55aec0cd419..6f34d36a5b1 100644 --- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -177,8 +177,7 @@ def get_dsv4_hybrid_module_spec_for_backend( core_attention = ModuleSpec( module=CompressedSparseAttention, submodules=CompressedSparseAttentionSubmodules( - compressor=compressor_spec, - indexer=indexer_spec, + compressor=compressor_spec, indexer=indexer_spec ), ) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 3e81d29dc13..fe8fa1985d8 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -9,10 +9,7 @@ import torch.nn as nn from megatron.core.fusions.fused_mla_yarn_rope_apply import fused_mla_rope_inplace -from megatron.core.models.common.embeddings import ( - RotaryEmbedding, - apply_rotary_pos_emb, -) +from megatron.core.models.common.embeddings import RotaryEmbedding, apply_rotary_pos_emb from megatron.core.packed_seq_params import PackedSeqParams from megatron.core.process_groups_config import ProcessGroupCollection from megatron.core.transformer.enums import AttnMaskType @@ -267,9 +264,7 @@ def __init__( super().__init__(config=config) if pg_collection is None: - pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=['tp', 'cp'] - ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) self.pg_collection = pg_collection self.compress_ratio = compress_ratio @@ -309,10 +304,7 @@ def __init__( # keep to high precision _ape = torch.empty( - compress_ratio, - proj_out_dim, - device=torch.cuda.current_device(), - dtype=torch.float32 + compress_ratio, proj_out_dim, device=torch.cuda.current_device(), dtype=torch.float32 ) config.init_method(_ape) self.ape = nn.Parameter(_ape) @@ -430,9 +422,7 @@ def __init__( super().__init__(config=config) if pg_collection is None: - pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=['tp', 'cp'] - ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) self.pg_collection = pg_collection self.compress_ratio = compress_ratio @@ -446,7 +436,7 @@ def __init__( self.index_head_dim = config.dsa_indexer_head_dim self.index_topk = config.dsa_indexer_topk - self.softmax_scale: float = self.index_head_dim ** -0.5 + self.softmax_scale: float = self.index_head_dim**-0.5 self.rotary_pos_emb = rotary_pos_emb @@ -488,10 +478,7 @@ def __init__( ) def forward_before_topk( - self, - x: torch.Tensor, - qr: torch.Tensor, - packed_seq_params: Optional[PackedSeqParams] = None, + self, x: torch.Tensor, qr: torch.Tensor, packed_seq_params: Optional[PackedSeqParams] = None ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """Compute Q, compressed K, and weights before top-k selection.""" nvtx_range_push("indexer_before_topk") @@ -517,7 +504,7 @@ def forward_before_topk( k = self.compressor(x) # [sq//ratio, b, index_head_dim] weights, _ = self.linear_weights_proj(x) # [sq, b, n_heads] - weights = weights * (self.index_n_heads ** -0.5) + weights = weights * (self.index_n_heads**-0.5) nvtx_range_pop("indexer_before_topk") return q, k, weights @@ -585,9 +572,7 @@ def __init__( super().__init__(config=config) if pg_collection is None: - pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=['tp', 'cp'] - ) + pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) self.pg_collection = pg_collection self.layer_number = layer_number @@ -598,7 +583,7 @@ def __init__( self.n_local_heads = config.num_attention_heads if softmax_scale is None: - softmax_scale = config.v_head_dim ** -0.5 + softmax_scale = config.v_head_dim**-0.5 self.softmax_scale = softmax_scale self.force_unfused_dsa = getattr(config, 'force_unfused_dsa', True) @@ -662,8 +647,9 @@ def forward( output: [sq, b, np * v_head_dim] """ nvtx_range_push("compressed_sparse_attn") - assert packed_seq_params is None, \ - "Packed sequence not supported for CompressedSparseAttention" + assert ( + packed_seq_params is None + ), "Packed sequence not supported for CompressedSparseAttention" sq, b, np, hn = query.size() @@ -698,15 +684,17 @@ def forward( x_det = x.detach() qr_det = qr.detach() - causal_mask = torch.arange(n_compressed, device=x.device).unsqueeze(0).expand( - sq, -1 + causal_mask = ( + torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(sq, -1) ) positions = torch.arange(1, sq + 1, device=x.device).unsqueeze(1) - causal_mask = torch.where( - causal_mask >= positions // self.compress_ratio, - float("-inf"), - 0.0, - ).unsqueeze(0).expand(b, -1, -1) # [b, sq, n_compressed] + causal_mask = ( + torch.where( + causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0 + ) + .unsqueeze(0) + .expand(b, -1, -1) + ) # [b, sq, n_compressed] if self.training and torch.is_grad_enabled(): q_indexer, k_indexer, weights_indexer = self.indexer.forward_before_topk( @@ -719,9 +707,7 @@ def forward( # indexer_softmax_scale; apply it here via the # weights-scaling trick so the effective weights match # the pre-scale-split behaviour. - weights_for_unfused = ( - weights_indexer * self.indexer.softmax_scale - ) + weights_for_unfused = weights_indexer * self.indexer.softmax_scale topk_indices_compressed, indexer_loss = FusedDSAIndexerLoss.apply( q_indexer, weights_for_unfused, 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 index e8e004cb3f1..4eed51b07ae 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -24,10 +24,7 @@ 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, -) +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 @@ -36,18 +33,10 @@ if HAVE_TE: - from megatron.core.extensions.transformer_engine import ( - TELinear, - set_save_original_input, - ) + from megatron.core.extensions.transformer_engine import TELinear, set_save_original_input, from megatron.core.post_training.modelopt.layers import Linear else: - ( - TEColumnParallelLinear, - TELinear, - Linear, - set_save_original_input, - ) = (None, None, None, None) + (TEColumnParallelLinear, TELinear, Linear, set_save_original_input) = (None, None, None, None) @torch.compile @@ -58,8 +47,7 @@ def _q_rms_norm(q: torch.Tensor, eps: float) -> torch.Tensor: @dataclass class DSv4HybridSelfAttentionSubmodules: - """Submodules for the DSv4HybridAttention layer. - """ + """Submodules for the DSv4HybridAttention layer.""" q_layernorm: LayerNormBuilder kv_layernorm: LayerNormBuilder @@ -72,8 +60,7 @@ class DSv4HybridSelfAttentionSubmodules: class DSv4HybridAttention(Attention): - """DeepSeek-v4 Hybrid Attention layer. - """ + """DeepSeek-v4 Hybrid Attention layer.""" def __init__( self, @@ -177,9 +164,9 @@ def __init__( # 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" - ) + 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 @@ -244,16 +231,21 @@ def forward( 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 ( + 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." # ===================== # Query, Key, and Value @@ -334,9 +326,7 @@ def forward( rope_seqlen, dtype=hidden_states.dtype, packed_seq=packed_seq ) rotary_pos_emb = None - assert ( - inference_context is None - ), "Inference with MLA RoPE fusion is not supported" + 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" @@ -507,8 +497,9 @@ def get_query_key_value_tensors( ), "dynamic_context_parallel is not supported with MLA yet and is planned for future. \ Please disable dynamic_context_parallel." - assert inference_context is None and inference_params is None, \ - "Inference is not supported for DSv4HybridSelfAttention." + assert ( + inference_context is None and inference_params is None + ), "Inference is not supported for DSv4HybridSelfAttention." # ========================================= # Prepare RoPE and seqlen related params @@ -658,9 +649,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po 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 - ) + 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( diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index 2ded8e3561e..7ff336fe886 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -210,7 +210,9 @@ def compute_dsa_indexer_loss( causal_mask = causal_mask_override.to(dtype=torch.float32) # [b, sq, sk] else: causal_mask = torch.triu( - torch.full((sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device), + torch.full( + (sq, sk), float('-inf'), dtype=torch.float32, device=attention_scores.device + ), diagonal=1, ) # index_mask [b, sq, sk] @@ -465,9 +467,7 @@ def bwd_fused_indexer_loss_naive( # Free attention_scores immediately del attention_scores - index_scores_softmax = torch.nn.functional.softmax( - index_scores, dim=-1, dtype=torch.float32 - ) + index_scores_softmax = torch.nn.functional.softmax(index_scores, dim=-1, dtype=torch.float32) # Free index_scores - no longer needed after softmax del index_scores diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index 4b6dba27f20..dc2619280ee 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -264,7 +264,9 @@ class TransformerConfig(ModelParallelConfig): #################### # attention variant #################### - experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa', 'dsv4_hybrid']] = None + 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.""" #################### @@ -1288,12 +1290,12 @@ def __post_init__(self): 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) + 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.tensor_model_parallel_size == 1 + ), "DSv4 Hybrid Attention only supports TP size 1." assert not self.qk_clip, "QK clipping is not supported with DSv4 Hybrid Attention." if self.fp8: @@ -2804,7 +2806,7 @@ def __post_init__(self): logger, logging.WARNING, f"DSv4 hybrid mode is enabled, deriving qk_head_dim and kv_lora_rank from " - f"v_head_dim and qk_pos_emb_head_dim" + f"v_head_dim and qk_pos_emb_head_dim", ) derived = self.v_head_dim - self.qk_pos_emb_head_dim self.qk_head_dim = derived diff --git a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py index 712044f933d..089949e4201 100644 --- a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py +++ b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py @@ -110,8 +110,7 @@ def _test_fused_mla_rope_inplace(input_format, inverse=False): pytorch_output.backward(pytorch_bwd_input, retain_graph=True) fused_output = fused_mla_rope_inplace( - fused_fwd_input, cos, sin, q_dim, emb_dim, - cu_seqlens_q=cu_seqlens, inverse=inverse, + fused_fwd_input, cos, sin, q_dim, emb_dim, cu_seqlens_q=cu_seqlens, inverse=inverse, ) fused_output.backward(fused_bwd_input, retain_graph=True) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py index b540a621192..f6327b2ffb3 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -40,12 +40,15 @@ def mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tensor def patch_hadamard_if_needed(): """Automatically patch hadamard_transform in both dsa and csa modules if 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), + 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: @@ -150,7 +153,9 @@ def test_output_shape(self): topk_indices = torch.randint(0, n_kv, (b, sq, topk), dtype=torch.int32).cuda() softmax_scale = hn**-0.5 - output = unfused_compressed_sparse_attn(query, kv_full, attn_sink, topk_indices, softmax_scale) + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) assert output.shape == (sq, b, np_ * hn) assert output.dtype == query.dtype @@ -170,7 +175,9 @@ def test_invalid_indices_masked(self): topk_indices[:, :, 0] = 0 softmax_scale = hn**-0.5 - output = unfused_compressed_sparse_attn(query, kv_full, attn_sink, topk_indices, softmax_scale) + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) assert not torch.isnan(output).any(), "Output should not contain NaN" @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @@ -187,7 +194,9 @@ def test_gradient_flow(self): topk_indices = torch.randint(0, n_kv, (b, sq, topk), dtype=torch.int32).cuda() softmax_scale = hn**-0.5 - output = unfused_compressed_sparse_attn(query, kv_full, attn_sink, topk_indices, softmax_scale) + output = unfused_compressed_sparse_attn( + query, kv_full, attn_sink, topk_indices, softmax_scale + ) loss = output.sum() loss.backward() @@ -271,10 +280,7 @@ def _make_csa_indexer_submodules(): return CSAIndexerSubmodules( linear_wq_b=ModuleSpec(module=TELinear), linear_weights_proj=ModuleSpec(module=TELinear), - compressor=ModuleSpec( - module=Compressor, - submodules=_make_compressor_submodules(), - ), + compressor=ModuleSpec(module=Compressor, submodules=_make_compressor_submodules()), ) @@ -283,14 +289,8 @@ def _make_csa_submodules(): from megatron.core.transformer.spec_utils import ModuleSpec return CompressedSparseAttentionSubmodules( - compressor=ModuleSpec( - module=Compressor, - submodules=_make_compressor_submodules(), - ), - indexer=ModuleSpec( - module=CSAIndexer, - submodules=_make_csa_indexer_submodules(), - ), + compressor=ModuleSpec(module=Compressor, submodules=_make_compressor_submodules()), + indexer=ModuleSpec(module=CSAIndexer, submodules=_make_csa_indexer_submodules()), ) @@ -312,12 +312,8 @@ def setup_method(self, request): model_parallel_cuda_manual_seed(123) cls = request.cls - cls.config = _make_mla_config( - csa_compress_ratios=[4, 128, 4, 128], - ) - cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=['tp', 'cp'] - ) + cls.config = _make_mla_config(csa_compress_ratios=[4, 128, 4, 128],) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) from megatron.core.models.common.embeddings import RotaryEmbedding @@ -373,9 +369,7 @@ def test_compressor_too_short_input(self, compress_ratio): pg_collection=self.pg_collection, ).cuda() - x = torch.randn( - short_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 - ).cuda() + x = torch.randn(short_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() output = compressor(x) assert output is None @@ -396,9 +390,11 @@ def test_compressor_gradient_flow(self, compress_ratio): pg_collection=self.pg_collection, ).cuda() - x = torch.randn( - seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 - ).cuda().requires_grad_(True) + x = ( + torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) output = compressor(x) loss = output.sum() loss.backward() @@ -428,13 +424,8 @@ def setup_method(self, request): cls = request.cls cls.compress_ratio = 4 - cls.config = _make_mla_config( - csa_compress_ratios=[4, 4, 4, 4], - dsa_indexer_topk=8, - ) - cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=['tp', 'cp'] - ) + cls.config = _make_mla_config(csa_compress_ratios=[4, 4, 4, 4], dsa_indexer_topk=8) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) from megatron.core.models.common.embeddings import RotaryEmbedding @@ -493,7 +484,12 @@ def test_csa_indexer_forward_before_topk(self, seqlen): q, k, weights = self.indexer.forward_before_topk(x, qr) - assert q.shape == (seqlen, batch_size, self.config.dsa_indexer_n_heads, self.config.dsa_indexer_head_dim) + assert q.shape == ( + seqlen, + batch_size, + self.config.dsa_indexer_n_heads, + self.config.dsa_indexer_head_dim + ) n_compressed = seqlen // self.compress_ratio assert k.shape == (n_compressed, batch_size, self.config.dsa_indexer_head_dim) assert weights.shape == (seqlen, batch_size, self.config.dsa_indexer_n_heads) @@ -510,11 +506,11 @@ def test_csa_indexer_with_mask(self, seqlen): n_compressed = seqlen // self.compress_ratio causal_mask = torch.arange(n_compressed, device=x.device).unsqueeze(0).expand(seqlen, -1) positions = torch.arange(1, seqlen + 1, device=x.device).unsqueeze(1) - causal_mask = torch.where( - causal_mask >= positions // self.compress_ratio, - float("-inf"), - 0.0, - ).unsqueeze(0).expand(batch_size, -1, -1) + causal_mask = ( + torch.where(causal_mask >= positions // self.compress_ratio, float("-inf"), 0.0) + .unsqueeze(0) + .expand(batch_size, -1, -1) + ) index_scores, topk_indices = self.indexer(x, qr, mask=causal_mask) @@ -540,13 +536,8 @@ def setup_method(self, request): model_parallel_cuda_manual_seed(123) cls = request.cls - cls.config = _make_mla_config( - csa_compress_ratios=[0, 0, 0, 0], - csa_window_size=8, - ) - cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=['tp', 'cp'] - ) + cls.config = _make_mla_config(csa_compress_ratios=[0, 0, 0, 0], csa_window_size=8) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) from megatron.core.models.common.embeddings import RotaryEmbedding @@ -591,9 +582,7 @@ def test_ratio1_forward(self): x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() - output = self.csa( - query=query, key=key, value=value, attention_mask=None, x=x, qr=qr, - ) + output = self.csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) assert output.shape == (seq_len, batch_size, np_ * hn) assert output.dtype == torch.bfloat16 @@ -609,15 +598,19 @@ def test_ratio1_backward(self): self.csa.train() self.csa.cuda() - query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.float32).cuda().requires_grad_(True) - key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.float32).cuda().requires_grad_(True) + query = ( + torch.randn(seq_len, batch_size, np_, hn, dtype=torch.float32) + .cuda() + .requires_grad_(True) + ) + key = ( + torch.randn(seq_len, batch_size, 1, hn, dtype=torch.float32).cuda().requires_grad_(True) + ) value = key.clone().detach().requires_grad_(True) x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() - output = self.csa( - query=query, key=key, value=value, attention_mask=None, x=x, qr=qr, - ) + output = self.csa(query=query, key=key, value=value, attention_mask=None, x=x, qr=qr) loss = output.sum() loss.backward() @@ -644,9 +637,7 @@ def setup_method(self, request): dsa_indexer_topk=8, dsa_indexer_loss_coeff=1.0, ) - cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=['tp', 'cp'] - ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) from megatron.core.models.common.embeddings import RotaryEmbedding @@ -737,8 +728,14 @@ def test_backward(self, compress_ratio): ).cuda() csa.train() - query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.float32).cuda().requires_grad_(True) - key = torch.randn(seq_len, batch_size, 1, hn, dtype=torch.float32).cuda().requires_grad_(True) + query = ( + torch.randn(seq_len, batch_size, np_, hn, dtype=torch.float32) + .cuda() + .requires_grad_(True) + ) + key = ( + torch.randn(seq_len, batch_size, 1, hn, dtype=torch.float32).cuda().requires_grad_(True) + ) value = key.clone().detach().requires_grad_(True) x = torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16).cuda() qr = torch.randn(seq_len, batch_size, self.config.q_lora_rank, dtype=torch.bfloat16).cuda() @@ -805,13 +802,9 @@ def setup_method(self, request): cls = request.cls cls.config = _make_mla_config( - csa_compress_ratios=[4, 128, 4, 128], - csa_window_size=8, - csa_dense_mode=True, - ) - cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups( - required_pgs=['tp', 'cp'] + csa_compress_ratios=[4, 128, 4, 128], csa_window_size=8, csa_dense_mode=True ) + cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) from megatron.core.models.common.embeddings import RotaryEmbedding 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 index 2449c45328b..afdbef7e733 100644 --- 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 @@ -32,12 +32,15 @@ def _mock_hadamard_transform(x: torch.Tensor, scale: float = 1.0) -> torch.Tenso 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), + 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: @@ -48,6 +51,7 @@ def patch_hadamard_if_needed(): # Config / spec helpers # --------------------------------------------------------------------------- + def _make_config( num_layers=4, hidden_size=256, @@ -214,9 +218,7 @@ def setup_method(self, request): model_parallel_cuda_manual_seed(_SEED) cls = request.cls - cls.config = _make_config( - dsa_indexer_loss_coeff=1.0, - ) + cls.config = _make_config(dsa_indexer_loss_coeff=1.0) cls.pg = ProcessGroupCollection.use_mpu_process_groups() yield @@ -231,7 +233,9 @@ def test_forward_output_shape(self, layer_number): 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 = _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 @@ -252,12 +256,16 @@ def test_backward_gradient_flow(self, layer_number): 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 = _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) + 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() From 93189166c24610cb915b11dc0457c0adac78948b Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Thu, 30 Apr 2026 01:53:10 +0000 Subject: [PATCH 07/16] minor fix --- .../core/transformer/experimental_attention_variant/csa.py | 2 ++ .../deepseek_v4_hybrid_attention.py | 2 +- tests/unit_tests/fusions/test_mla_yarn_rope_apply.py | 2 +- .../test_attention_variant_csa.py | 4 ++-- 4 files changed, 6 insertions(+), 4 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index fe8fa1985d8..ab17ca9d5c1 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -31,6 +31,7 @@ # --------------------------------------------------------------------------- +# TODO: the lru_cache may not work well with packed sequence @lru_cache(maxsize=8) def _get_window_topk_idxs_cached(window_size: int, seqlen: int, device_str: str) -> torch.Tensor: """Compute sliding-window indices for a single sequence (cached). @@ -53,6 +54,7 @@ def get_window_topk_idxs( return matrix.unsqueeze(0).expand(batch_size, -1, -1) +# TODO: the lru_cache may not work well with packed sequence @lru_cache(maxsize=8) def _get_compress_topk_idxs_cached( ratio: int, seqlen: int, offset: int, device_str: str 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 index 4eed51b07ae..9a0396cf457 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -33,7 +33,7 @@ if HAVE_TE: - from megatron.core.extensions.transformer_engine import TELinear, set_save_original_input, + from megatron.core.extensions.transformer_engine import TELinear, set_save_original_input from megatron.core.post_training.modelopt.layers import Linear else: (TEColumnParallelLinear, TELinear, Linear, set_save_original_input) = (None, None, None, None) diff --git a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py index 089949e4201..8fd6acd1653 100644 --- a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py +++ b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py @@ -110,7 +110,7 @@ def _test_fused_mla_rope_inplace(input_format, inverse=False): pytorch_output.backward(pytorch_bwd_input, retain_graph=True) fused_output = fused_mla_rope_inplace( - fused_fwd_input, cos, sin, q_dim, emb_dim, cu_seqlens_q=cu_seqlens, inverse=inverse, + fused_fwd_input, cos, sin, q_dim, emb_dim, cu_seqlens_q=cu_seqlens, inverse=inverse ) fused_output.backward(fused_bwd_input, retain_graph=True) diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py index f6327b2ffb3..41edaf924c5 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -312,7 +312,7 @@ def setup_method(self, request): model_parallel_cuda_manual_seed(123) cls = request.cls - cls.config = _make_mla_config(csa_compress_ratios=[4, 128, 4, 128],) + cls.config = _make_mla_config(csa_compress_ratios=[4, 128, 4, 128]) cls.pg_collection = ProcessGroupCollection.use_mpu_process_groups(required_pgs=['tp', 'cp']) from megatron.core.models.common.embeddings import RotaryEmbedding @@ -488,7 +488,7 @@ def test_csa_indexer_forward_before_topk(self, seqlen): seqlen, batch_size, self.config.dsa_indexer_n_heads, - self.config.dsa_indexer_head_dim + self.config.dsa_indexer_head_dim, ) n_compressed = seqlen // self.compress_ratio assert k.shape == (n_compressed, batch_size, self.config.dsa_indexer_head_dim) From cf76ce5681c4b8508b08ef974ace2a1dc0371a7c Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Thu, 30 Apr 2026 02:43:00 +0000 Subject: [PATCH 08/16] fix format --- ...perimental_attention_variant_module_specs.py | 17 ++++++++--------- .../experimental_attention_variant/csa.py | 1 - 2 files changed, 8 insertions(+), 10 deletions(-) diff --git a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py index 6f34d36a5b1..329b8f259ea 100644 --- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -6,25 +6,24 @@ from megatron.core.models.backends import BackendSpecProvider from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules from megatron.core.transformer.enums import AttnMaskType, LayerType -from megatron.core.transformer.experimental_attention_variant.dsa import ( - DSAIndexer, - DSAIndexerSubmodules, - DSAttention, - DSAttentionSubmodules, -) from megatron.core.transformer.experimental_attention_variant.csa import ( + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, Compressor, CompressorSubmodules, CSAIndexer, CSAIndexerSubmodules, - CompressedSparseAttention, - CompressedSparseAttentionSubmodules, ) from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( DSv4HybridSelfAttention, DSv4HybridSelfAttentionSubmodules, ) - +from megatron.core.transformer.experimental_attention_variant.dsa import ( + DSAIndexer, + DSAIndexerSubmodules, + DSAttention, + DSAttentionSubmodules, +) from megatron.core.transformer.hyper_connection import HyperConnectionModule from megatron.core.transformer.identity_op import IdentityOp from megatron.core.transformer.multi_latent_attention import ( diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index ab17ca9d5c1..8256fc5af9e 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -25,7 +25,6 @@ from megatron.core.transformer.transformer_config import TransformerConfig from megatron.core.utils import nvtx_range_pop, nvtx_range_push - # --------------------------------------------------------------------------- # Helper functions for index computation # --------------------------------------------------------------------------- From 961eafe558adef6745bba4706094634ffadfa334 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Thu, 30 Apr 2026 02:49:00 +0000 Subject: [PATCH 09/16] fix format --- .../deepseek_v4_hybrid_attention.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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 index 9a0396cf457..f04840957a3 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -17,7 +17,6 @@ FineGrainedActivationOffloadingInterface as off_interface, ) from megatron.core.process_groups_config import ProcessGroupCollection -from megatron.core.tensor_parallel.layers import ColumnParallelLinear from megatron.core.transformer.attention import Attention from megatron.core.transformer.enums import AttnMaskType from megatron.core.transformer.spec_utils import ModuleSpec, build_module @@ -34,9 +33,8 @@ if HAVE_TE: from megatron.core.extensions.transformer_engine import TELinear, set_save_original_input - from megatron.core.post_training.modelopt.layers import Linear else: - (TEColumnParallelLinear, TELinear, Linear, set_save_original_input) = (None, None, None, None) + (TEColumnParallelLinear, TELinear, set_save_original_input) = (None, None, None) @torch.compile From 823959d347c81bb9951ddfef2de24d1b7130c824 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Thu, 30 Apr 2026 05:29:16 +0000 Subject: [PATCH 10/16] fix some tests --- .../transformer/experimental_attention_variant/dsa.py | 9 +++++++-- tests/unit_tests/models/test_hybrid_moe_model.py | 4 ++++ .../test_attention_variant_csa.py | 7 +++++++ 3 files changed, 18 insertions(+), 2 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index 7ff336fe886..5d7566b3926 100644 --- a/megatron/core/transformer/experimental_attention_variant/dsa.py +++ b/megatron/core/transformer/experimental_attention_variant/dsa.py @@ -529,15 +529,20 @@ def bwd_fused_indexer_loss_naive( causal_valid_mask = torch.tril( torch.ones((sq, sk), device=q.device, dtype=torch.bool) ) # [sq, sk] + + if causal_valid_mask.dim() == 2: + causal_valid_mask = causal_valid_mask.unsqueeze(0) + causal_valid_mask = causal_valid_mask.expand(b, sq, sk) + if sparse_loss: # Also apply index mask - only topk positions are valid index_valid_mask = index_mask == 0 # [b, sq, sk] del index_mask # Free index_mask immediately after use - valid_mask = causal_valid_mask.unsqueeze(0) & index_valid_mask # [b, sq, sk] + valid_mask = causal_valid_mask & index_valid_mask # [b, sq, sk] del index_valid_mask else: del index_mask # Free index_mask even if not used for sparse_loss - valid_mask = causal_valid_mask.unsqueeze(0).expand(b, sq, sk) # [b, sq, sk] + valid_mask = causal_valid_mask # [b, sq, sk] del causal_valid_mask grad_index_scores_logits = grad_index_scores_logits * valid_mask.float() diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index 1df280853fc..81aca3abea5 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -69,6 +69,10 @@ "cpu_offloading_weights": False, "cross_entropy_fusion_impl": "native", "cross_entropy_loss_fusion": True, + "csa_compress_ratios": None, + "csa_compress_rotary_base": 40000.0, + "csa_dense_mode": False, + "csa_window_size": 128, "cuda_graph_impl": "none", "cuda_graph_retain_backward_graph": False, "cuda_graph_scope": [], diff --git a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py index 41edaf924c5..83c153d698e 100644 --- a/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -556,6 +556,7 @@ def setup_method(self, request): attention_type='self', pg_collection=cls.pg_collection, rotary_pos_emb=rotary_pos_emb, + compress_ratio=0, ) yield @@ -670,6 +671,7 @@ def test_constructor(self, compress_ratio): attention_type='self', pg_collection=self.pg_collection, rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, ).cuda() assert csa.compressor is not None @@ -695,6 +697,7 @@ def test_forward(self, compress_ratio): attention_type='self', pg_collection=self.pg_collection, rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, ).cuda() query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() @@ -725,6 +728,7 @@ def test_backward(self, compress_ratio): attention_type='self', pg_collection=self.pg_collection, rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, ).cuda() csa.train() @@ -768,6 +772,7 @@ def test_eval_mode(self, compress_ratio): attention_type='self', pg_collection=self.pg_collection, rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=compress_ratio, ).cuda() csa.eval() @@ -829,6 +834,7 @@ def test_dense_mode_disables_indexer_for_ratio4(self): attention_type='self', pg_collection=self.pg_collection, rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=4, ).cuda() assert csa.compress_ratio == 4 @@ -851,6 +857,7 @@ def test_dense_mode_forward_ratio4(self): attention_type='self', pg_collection=self.pg_collection, rotary_pos_emb=self.rotary_pos_emb, + compress_ratio=4, ).cuda() query = torch.randn(seq_len, batch_size, np_, hn, dtype=torch.bfloat16).cuda() From dd6e5f02a2edd2e5fffdb1d188106bb8c09149d5 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Thu, 30 Apr 2026 08:23:29 +0000 Subject: [PATCH 11/16] fix test --- .../test_dsv4_hybrid_attention.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) 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 index afdbef7e733..5f9a3a74440 100644 --- 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 @@ -80,6 +80,7 @@ def _make_config( 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, @@ -105,11 +106,12 @@ def _make_config( 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.models.gpt.experimental_attention_variant_module_specs import ( get_dsv4_hybrid_module_spec_for_backend, ) - return get_dsv4_hybrid_module_spec_for_backend(config=config) + return get_dsv4_hybrid_module_spec_for_backend(config=config, backend=TESpecProvider()) def _build_attention(config, layer_number, pg_collection): From 80a2d9184232cb3baf1c4c98fbcf5aea582f7d53 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Thu, 30 Apr 2026 09:44:02 +0000 Subject: [PATCH 12/16] fix missing import in PR #4481 --- megatron/core/transformer/transformer_config.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/megatron/core/transformer/transformer_config.py b/megatron/core/transformer/transformer_config.py index dc2619280ee..11a3eb09014 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -11,7 +11,7 @@ from megatron.core.enums import Fp4Recipe, Fp8Recipe from megatron.core.quantization.quant_config import RecipeConfig -from megatron.core.transformer.enums import AttnBackend, CudaGraphScope +from megatron.core.transformer.enums import AttnBackend, CudaGraphScope, LayerType from megatron.core.transformer.pipeline_parallel_layer_layout import PipelineParallelLayerLayout from megatron.core.utils import experimental_api From b13ec2ab49659d4a979381fd4694bee7430b6ce4 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Thu, 30 Apr 2026 12:07:20 +0000 Subject: [PATCH 13/16] revert to dsv3 rope and add a new arg mla_output_remove_interleaving --- .../core/fusions/fused_mla_yarn_rope_apply.py | 99 +++++++++++++++---- .../models/common/embeddings/rope_utils.py | 17 +++- .../experimental_attention_variant/csa.py | 2 + .../deepseek_v4_hybrid_attention.py | 6 ++ .../fusions/test_mla_yarn_rope_apply.py | 32 ++++-- 5 files changed, 123 insertions(+), 33 deletions(-) diff --git a/megatron/core/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py index 9529cbc2a3f..6eed7581d03 100644 --- a/megatron/core/fusions/fused_mla_yarn_rope_apply.py +++ b/megatron/core/fusions/fused_mla_yarn_rope_apply.py @@ -82,6 +82,7 @@ def _mla_rope_fwd_inplace_kernel( cp_rank, cp_size, INVERSE: tl.constexpr, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ @@ -134,8 +135,14 @@ def _mla_rope_fwd_inplace_kernel( x_left = x_1 * cos_left - x_2 * sin_left x_right = x_2 * cos_right + x_1 * sin_right - tl.store(Q + x_1_off, x_left, mask=mask) - tl.store(Q + x_2_off, x_right, mask=mask) + if REMOVE_INTERLEAVING: + tl.store(Q + x_1_off, x_left, mask=mask) + tl.store(Q + x_2_off, x_right, mask=mask) + else: + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + tl.store(Q + x_left_off, x_left, mask=mask) + tl.store(Q + x_right_off, x_right, mask=mask) @triton.autotune( @@ -170,6 +177,7 @@ def _mla_rope_bwd_inplace_kernel( cp_rank, cp_size, INVERSE: tl.constexpr, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ @@ -211,10 +219,18 @@ def _mla_rope_bwd_inplace_kernel( x_off = tl.arange(0, BLOCK_H)[:, None] * stride_x_nheads + nope_dim mask = x_off < head_num * stride_x_nheads - x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 - x_2_off = x_1_off + 1 - x_left = tl.load(DO + x_1_off, mask=mask) - x_right = tl.load(DO + x_2_off, mask=mask) + if REMOVE_INTERLEAVING: + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_left = tl.load(DO + x_1_off, mask=mask) + x_right = tl.load(DO + x_2_off, mask=mask) + else: + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left = tl.load(DO + x_left_off, mask=mask) + x_right = tl.load(DO + x_right_off, mask=mask) + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 x_1 = x_left * cos_left + x_right * sin_right x_2 = -x_left * sin_left + x_right * cos_right @@ -242,6 +258,7 @@ def forward( cp_size, rotary_interleaved=False, inverse=False, + remove_interleaving=False, ): """ Forward function for _FusedMLARoPEInplace. @@ -291,6 +308,7 @@ def forward( cp_rank, cp_size, INVERSE=inverse, + REMOVE_INTERLEAVING=remove_interleaving, ) ctx.save_for_backward(cos, sin) ctx.nope_dim = nope_dim @@ -298,6 +316,7 @@ def forward( ctx.cu_seqlens_q = cu_seqlens_q ctx.rotary_interleaved = rotary_interleaved ctx.inverse = inverse + ctx.remove_interleaving = remove_interleaving ctx.cp_rank = cp_rank ctx.cp_size = cp_size if cu_seqlens_q is None: @@ -344,10 +363,11 @@ def backward(ctx, grad): ctx.cp_rank, ctx.cp_size, INVERSE=ctx.inverse, + REMOVE_INTERLEAVING=ctx.remove_interleaving, ) if ctx.cu_seqlens_q is None: grad = grad.view(max_seqlen, batch_size, nheads, headdim) - return grad, None, None, None, None, None, None, None, None, None + return grad, None, None, None, None, None, None, None, None, None, None def fused_mla_rope_inplace( @@ -361,6 +381,7 @@ def fused_mla_rope_inplace( cp_size: int = 1, rotary_interleaved: bool = False, inverse: bool = False, + remove_interleaving: bool = False, ): """ Fused RoPE applied inplace to the trailing emb_dim elements of a tensor, @@ -381,12 +402,23 @@ def fused_mla_rope_inplace( cu_seqlens_q: [seq_num + 1] accumulated sequence lengths for thd format rotary_interleaved: whether to apply RoPE interleaved, only supports False for now inverse: if True, apply the inverse rotation + remove_interleaving: if True, output RoPE dims in non-interleaved layout Returns: t: inplace modified input tensor """ return _FusedMLARoPEInplace.apply( - t, cos, sin, nope_dim, emb_dim, cu_seqlens_q, cp_rank, cp_size, rotary_interleaved, inverse + t, + cos, + sin, + nope_dim, + emb_dim, + cu_seqlens_q, + cp_rank, + cp_size, + rotary_interleaved, + inverse, + remove_interleaving, ) @@ -427,6 +459,7 @@ def _mla_rope_fwd_kv_split_kernel( stride_v_nheads, cp_rank, cp_size, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ @@ -487,14 +520,24 @@ def _mla_rope_fwd_kv_split_kernel( x_left = x_left.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) x_right = x_right.expand_dims(0).broadcast_to(BLOCK_H, emb_dim // 2) - x_1_off = ( - tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads - + k_dim - + tl.arange(0, emb_dim // 2)[None, :] * 2 - ) - x_2_off = x_1_off + 1 - tl.store(K_ptr + x_1_off, x_left, mask=mask) - tl.store(K_ptr + x_2_off, x_right, mask=mask) + if REMOVE_INTERLEAVING: + x_1_off = ( + tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] * 2 + ) + x_2_off = x_1_off + 1 + tl.store(K_ptr + x_1_off, x_left, mask=mask) + tl.store(K_ptr + x_2_off, x_right, mask=mask) + else: + x_left_off = ( + tl.arange(0, BLOCK_H)[:, None] * stride_k_nheads + + k_dim + + tl.arange(0, emb_dim // 2)[None, :] + ) + x_right_off = x_left_off + emb_dim // 2 + tl.store(K_ptr + x_left_off, x_left, mask=mask) + tl.store(K_ptr + x_right_off, x_right, mask=mask) @triton.autotune( @@ -534,6 +577,7 @@ def _mla_rope_bwd_kv_split_kernel( stride_demb_seq, cp_rank, cp_size, + REMOVE_INTERLEAVING: tl.constexpr, BLOCK_H: tl.constexpr, ): """ @@ -582,10 +626,16 @@ def _mla_rope_bwd_kv_split_kernel( dK_ptr = dK + pid_m * stride_dk_seq + i * BLOCK_H * stride_dk_nheads x_off = tl.arange(0, BLOCK_H)[:, None] * stride_dk_nheads + k_dim mask = x_off < head_num * stride_dk_nheads - x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 - x_2_off = x_1_off + 1 - x_left = tl.load(dK_ptr + x_1_off, mask=mask) - x_right = tl.load(dK_ptr + x_2_off, mask=mask) + if REMOVE_INTERLEAVING: + x_1_off = x_off + tl.arange(0, emb_dim // 2)[None, :] * 2 + x_2_off = x_1_off + 1 + x_left = tl.load(dK_ptr + x_1_off, mask=mask) + x_right = tl.load(dK_ptr + x_2_off, mask=mask) + else: + x_left_off = x_off + tl.arange(0, emb_dim // 2)[None, :] + x_right_off = x_left_off + emb_dim // 2 + x_left = tl.load(dK_ptr + x_left_off, mask=mask) + x_right = tl.load(dK_ptr + x_right_off, mask=mask) x_left_accum += x_left x_right_accum += x_right x_left_accum = tl.sum(x_left_accum, axis=0) @@ -625,6 +675,7 @@ def forward( cp_rank, cp_size, rotary_interleaved=False, + remove_interleaving=False, ): """ Forward function for _FusedMLARoPEKVSplit. @@ -685,8 +736,10 @@ def forward( o_value.stride(1), cp_rank, cp_size, + REMOVE_INTERLEAVING=remove_interleaving, ) ctx.save_for_backward(cos, sin) + ctx.remove_interleaving = remove_interleaving ctx.rotary_interleaved = rotary_interleaved ctx.emb_dim = emb_dim ctx.k_dim = k_dim @@ -753,11 +806,12 @@ def backward(ctx, dk, dv): d_emb.stride(0), ctx.cp_rank, ctx.cp_size, + REMOVE_INTERLEAVING=ctx.remove_interleaving, ) if ctx.cu_seqlens_kv is None: d_kv = d_kv.view(max_seqlen, batch_size, nheads, ctx.k_dim + ctx.v_dim) d_emb = d_emb.view(max_seqlen, batch_size, 1, ctx.emb_dim) - return d_kv, d_emb, None, None, None, None, None, None, None, None, None + return d_kv, d_emb, None, None, None, None, None, None, None, None, None, None def fused_mla_rope_kv_split( @@ -772,6 +826,7 @@ def fused_mla_rope_kv_split( cp_rank: int = 0, cp_size: int = 1, rotary_interleaved: bool = False, + remove_interleaving: bool = False, ): """ Fused function for applying RoPE to MLA's key and value. @@ -789,6 +844,7 @@ def fused_mla_rope_kv_split( cos/sin: [max_seq_len, 1, 1, emb_dim] cu_seqlens_kv: [seq_num + 1] accumulated sequence lengths for thd format rotary_interleaved: whether to apply RoPE interleaved, only supports False for now + remove_interleaving: if True, output RoPE dims in non-interleaved layout Returns: key: [seq_len, batch_size, head_num, emb_dim + k_dim] @@ -807,6 +863,7 @@ def fused_mla_rope_kv_split( cp_rank, cp_size, rotary_interleaved, + remove_interleaving, ) diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index c558260f03f..4442d8c9bed 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -96,6 +96,7 @@ def _apply_rotary_pos_emb_bshd( mla_rotary_interleaved: bool = False, mscale: float = 1.0, inverse: bool = False, + mla_output_remove_interleaving: bool = False, multi_latent_attention: Optional[bool] = None, ) -> Tensor: """Apply rotary positional embedding to input tensor T. @@ -148,7 +149,7 @@ def _apply_rotary_pos_emb_bshd( # Fallback to original permutation # DSv4 applies rope on V and O, so we need to uninterleave the tensor. # The existing MLA code is safe because the dot product is permutation-invariant. - if mla_rotary_interleaved: + if mla_rotary_interleaved and mla_output_remove_interleaving: x1, x2 = torch.chunk(t, 2, dim=-1) t = torch.stack((x1, x2), dim=-1).flatten(start_dim=-2) @@ -212,6 +213,7 @@ def _apply_rotary_pos_emb_thd( mla_rotary_interleaved: bool = False, mscale: float = 1.0, inverse: bool = False, + mla_output_remove_interleaving: bool = False, cp_group: torch.distributed.ProcessGroup = None, multi_latent_attention: Optional[bool] = None, ) -> Tensor: @@ -266,6 +268,7 @@ def _apply_rotary_pos_emb_thd( mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, ).squeeze(1) else: # CASE 2: Traditional mapping without offsets @@ -283,6 +286,7 @@ def _apply_rotary_pos_emb_thd( mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, ).squeeze(1) @@ -295,6 +299,7 @@ def apply_rotary_pos_emb( cp_group: torch.distributed.ProcessGroup = None, mla_rotary_interleaved: bool = False, inverse: bool = False, + mla_output_remove_interleaving: bool = False, ): """ Reroute to the appropriate apply_rotary_pos_emb function depending on @@ -337,7 +342,12 @@ def apply_rotary_pos_emb( use_unfused = True if not use_unfused: assert fused_apply_rotary_pos_emb is not None, "apply_rope_fusion is not available." - return fused_apply_rotary_pos_emb(t, freqs, interleaved=config.rotary_interleaved) + return fused_apply_rotary_pos_emb( + t, + freqs, + interleaved=config.rotary_interleaved, + mla_output_remove_interleaving=mla_output_remove_interleaving, + ) else: assert fused_apply_rotary_pos_emb_thd is not None, "apply_rope_fusion is not available." return fused_apply_rotary_pos_emb_thd( @@ -347,6 +357,7 @@ def apply_rotary_pos_emb( cp_size=cp_group.size(), cp_rank=cp_group.rank(), interleaved=config.rotary_interleaved, + mla_output_remove_interleaving=mla_output_remove_interleaving, ) # use unfused implementation if cu_seqlens is None: @@ -357,6 +368,7 @@ def apply_rotary_pos_emb( mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, ) else: return _apply_rotary_pos_emb_thd( @@ -368,6 +380,7 @@ def apply_rotary_pos_emb( mscale=mscale, cp_group=cp_group, inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, ) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index 8256fc5af9e..f27d022411a 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -140,6 +140,7 @@ def _apply_rope( None, cp_group.rank(), cp_group.size(), + mla_output_remove_interleaving=True, ) else: x_nope, x_pe = torch.split(x, [nope_dim, pos_dim], dim=-1) @@ -151,6 +152,7 @@ def _apply_rope( mscale=mscale, cp_group=cp_group, mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, ) out = torch.cat([x_nope, x_pe], dim=-1) if squeeze_head: 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 index f04840957a3..b7018e527bf 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -341,6 +341,7 @@ def forward( self.pg_collection.cp.rank(), self.pg_collection.cp.size(), inverse=True, + mla_output_remove_interleaving=True, ) else: content_part, rot_part = torch.split( @@ -355,6 +356,7 @@ def forward( 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) @@ -599,6 +601,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po cu_seqlens_q, cp_rank, cp_size, + mla_output_remove_interleaving=True, ) kv = kv.unsqueeze(-2) kv = fused_mla_rope_inplace( @@ -610,6 +613,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po cu_seqlens_q, cp_rank, cp_size, + mla_output_remove_interleaving=True, ) key = kv value = kv @@ -642,6 +646,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po 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) @@ -658,6 +663,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po 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] diff --git a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py index 8fd6acd1653..762195b5d7f 100644 --- a/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py +++ b/tests/unit_tests/fusions/test_mla_yarn_rope_apply.py @@ -43,7 +43,7 @@ def rank(self): return 0 -def _test_fused_mla_rope_inplace(input_format, inverse=False): +def _test_fused_mla_rope_inplace(input_format, inverse=False, remove_interleaving=False): assert fused_mla_rope_inplace is not None num_heads = 32 q_dim = 128 @@ -105,12 +105,20 @@ def _test_fused_mla_rope_inplace(input_format, inverse=False): cp_group=FakeCPGroup(), mla_rotary_interleaved=True, inverse=inverse, + mla_output_remove_interleaving=remove_interleaving, ) pytorch_output = torch.concat([no_pe, pe_output], dim=-1) pytorch_output.backward(pytorch_bwd_input, retain_graph=True) fused_output = fused_mla_rope_inplace( - fused_fwd_input, cos, sin, q_dim, emb_dim, cu_seqlens_q=cu_seqlens, inverse=inverse + fused_fwd_input, + cos, + sin, + q_dim, + emb_dim, + cu_seqlens_q=cu_seqlens, + inverse=inverse, + remove_interleaving=remove_interleaving, ) fused_output.backward(fused_bwd_input, retain_graph=True) @@ -129,7 +137,7 @@ def _test_fused_mla_rope_inplace(input_format, inverse=False): ) -def _test_fused_mla_rope_kv_split(input_format): +def _test_fused_mla_rope_kv_split(input_format, remove_interleaving=False): assert fused_mla_rope_kv_split is not None num_heads = 32 k_dim = 128 @@ -204,6 +212,7 @@ def _test_fused_mla_rope_kv_split(input_format): mscale=mscale, cp_group=FakeCPGroup(), mla_rotary_interleaved=True, + mla_output_remove_interleaving=remove_interleaving, ) if input_format == "sbhd": pe_output = pe_output.expand(-1, -1, num_heads, -1) @@ -224,6 +233,7 @@ def _test_fused_mla_rope_kv_split(input_format): k_dim, v_dim, cu_seqlens_kv=cu_seqlens, + remove_interleaving=remove_interleaving, ) torch.autograd.backward( (fused_k_output, fused_v_output), (fused_bwd_k_input, fused_bwd_v_input) @@ -262,14 +272,16 @@ def _test_fused_mla_rope_kv_split(input_format): @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") @pytest.mark.parametrize("input_format", ["sbhd", "thd"]) class TestFusedMLARope: - def test_inplace_forward_backward(self, input_format): - _test_fused_mla_rope_inplace(input_format, inverse=False) - - def test_inplace_inverse_forward_backward(self, input_format): - _test_fused_mla_rope_inplace(input_format, inverse=True) + @pytest.mark.parametrize("inverse", [False, True]) + @pytest.mark.parametrize("remove_interleaving", [False, True]) + def test_inplace_forward_backward(self, input_format, inverse, remove_interleaving): + _test_fused_mla_rope_inplace( + input_format, inverse=inverse, remove_interleaving=remove_interleaving + ) - def test_kv_split_forward_backward(self, input_format): - _test_fused_mla_rope_kv_split(input_format) + @pytest.mark.parametrize("remove_interleaving", [False, True]) + def test_kv_split_forward_backward(self, input_format, remove_interleaving): + _test_fused_mla_rope_kv_split(input_format, remove_interleaving=remove_interleaving) class TestApplyRotaryPosEmbMlaFusionConflict: From f0fdb2d6b818353d4c5dbc521be148f12518d1c8 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Thu, 30 Apr 2026 12:32:58 +0000 Subject: [PATCH 14/16] minor fix --- .../core/transformer/experimental_attention_variant/csa.py | 2 +- .../deepseek_v4_hybrid_attention.py | 6 +++--- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/megatron/core/transformer/experimental_attention_variant/csa.py b/megatron/core/transformer/experimental_attention_variant/csa.py index f27d022411a..1c24ecda5c7 100644 --- a/megatron/core/transformer/experimental_attention_variant/csa.py +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -140,7 +140,7 @@ def _apply_rope( None, cp_group.rank(), cp_group.size(), - mla_output_remove_interleaving=True, + remove_interleaving=True, ) else: x_nope, x_pe = torch.split(x, [nope_dim, pos_dim], dim=-1) 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 index b7018e527bf..7aa321a3cd1 100644 --- a/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -341,7 +341,7 @@ def forward( self.pg_collection.cp.rank(), self.pg_collection.cp.size(), inverse=True, - mla_output_remove_interleaving=True, + remove_interleaving=True, ) else: content_part, rot_part = torch.split( @@ -601,7 +601,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po cu_seqlens_q, cp_rank, cp_size, - mla_output_remove_interleaving=True, + remove_interleaving=True, ) kv = kv.unsqueeze(-2) kv = fused_mla_rope_inplace( @@ -613,7 +613,7 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po cu_seqlens_q, cp_rank, cp_size, - mla_output_remove_interleaving=True, + remove_interleaving=True, ) key = kv value = kv From 8aefd1043a8339048954e7046d1cc48164de4de5 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Thu, 30 Apr 2026 13:48:56 +0000 Subject: [PATCH 15/16] minor fix --- megatron/core/models/common/embeddings/rope_utils.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index 4442d8c9bed..b3c9344ac71 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -346,7 +346,6 @@ def apply_rotary_pos_emb( t, freqs, interleaved=config.rotary_interleaved, - mla_output_remove_interleaving=mla_output_remove_interleaving, ) else: assert fused_apply_rotary_pos_emb_thd is not None, "apply_rope_fusion is not available." @@ -357,7 +356,6 @@ def apply_rotary_pos_emb( cp_size=cp_group.size(), cp_rank=cp_group.rank(), interleaved=config.rotary_interleaved, - mla_output_remove_interleaving=mla_output_remove_interleaving, ) # use unfused implementation if cu_seqlens is None: From 8470a3b67fbae46e7e1040fa379ed74ef0634390 Mon Sep 17 00:00:00 2001 From: Hongxiao Bai Date: Thu, 30 Apr 2026 13:58:42 +0000 Subject: [PATCH 16/16] fix lint --- megatron/core/models/common/embeddings/rope_utils.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/megatron/core/models/common/embeddings/rope_utils.py b/megatron/core/models/common/embeddings/rope_utils.py index b3c9344ac71..c97f738771b 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -342,11 +342,7 @@ def apply_rotary_pos_emb( use_unfused = True if not use_unfused: assert fused_apply_rotary_pos_emb is not None, "apply_rope_fusion is not available." - return fused_apply_rotary_pos_emb( - t, - freqs, - interleaved=config.rotary_interleaved, - ) + return fused_apply_rotary_pos_emb(t, freqs, interleaved=config.rotary_interleaved) else: assert fused_apply_rotary_pos_emb_thd is not None, "apply_rope_fusion is not available." return fused_apply_rotary_pos_emb_thd(