From 733c5506e3c31817c81f4cf641230883bd797d90 Mon Sep 17 00:00:00 2001 From: hx Date: Thu, 30 Apr 2026 23:22:26 +0800 Subject: [PATCH] [dev] [DeepSeek-v4] Part 1: Hybrid Attention with CSA and HCA (#4458) --- .../core/fusions/fused_mla_yarn_rope_apply.py | 270 ++++-- .../models/common/embeddings/rope_utils.py | 37 + ...rimental_attention_variant_module_specs.py | 71 ++ megatron/core/models/gpt/gpt_model.py | 4 +- megatron/core/transformer/attention.py | 6 + .../experimental_attention_variant/csa.py | 770 ++++++++++++++++ .../deepseek_v4_hybrid_attention.py | 715 ++++++++++++++ .../experimental_attention_variant/dsa.py | 137 ++- .../transformer/multi_latent_attention.py | 6 + .../core/transformer/transformer_config.py | 65 +- .../core/transformer/transformer_layer.py | 497 ++++++++++ megatron/training/arguments.py | 29 + .../fusions/test_mla_yarn_rope_apply.py | 47 +- .../models/test_hybrid_moe_model.py | 4 + .../test_attention_variant_csa.py | 872 ++++++++++++++++++ .../test_dsv4_hybrid_attention.py | 416 +++++++++ 16 files changed, 3808 insertions(+), 138 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 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/fusions/fused_mla_yarn_rope_apply.py b/megatron/core/fusions/fused_mla_yarn_rope_apply.py index 1fd5dcfae37..6eed7581d03 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,21 @@ 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, + REMOVE_INTERLEAVING: 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 +106,17 @@ 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 +124,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 +135,14 @@ 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) + 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( @@ -145,11 +160,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 +172,21 @@ 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, + REMOVE_INTERLEAVING: 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 +199,17 @@ 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 +217,32 @@ 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) + 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 - 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 +251,25 @@ def forward( q, cos, sin, - qk_head_dim, + nope_dim, emb_dim, cu_seqlens_q, cp_rank, cp_size, rotary_interleaved=False, + inverse=False, + remove_interleaving=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 +285,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 +303,20 @@ def forward( cu_seqlens_q, q.stride(0), q.stride(1), + cos.stride(0), + sin.stride(0), cp_rank, cp_size, + INVERSE=inverse, + REMOVE_INTERLEAVING=remove_interleaving, ) 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.remove_interleaving = remove_interleaving ctx.cp_rank = cp_rank ctx.cp_size = cp_size if cu_seqlens_q is None: @@ -284,11 +326,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 +346,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 +358,67 @@ 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, + 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 + return grad, None, 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, + remove_interleaving: 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 + remove_interleaving: if True, output RoPE dims in non-interleaved layout 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, + remove_interleaving, ) @@ -376,7 +436,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, @@ -399,12 +459,12 @@ def rotary_fwd_kv_kernel( stride_v_nheads, cp_rank, cp_size, + REMOVE_INTERLEAVING: tl.constexpr, 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 +520,24 @@ 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 = ( - 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) + 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( @@ -484,7 +554,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, @@ -507,10 +577,11 @@ def rotary_bwd_kv_kernel( stride_demb_seq, cp_rank, cp_size, + REMOVE_INTERLEAVING: tl.constexpr, 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 +626,16 @@ 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) + 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) @@ -578,9 +655,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 @@ -597,9 +675,10 @@ def forward( cp_rank, cp_size, rotary_interleaved=False, + remove_interleaving=False, ): """ - Forward function for ApplyMLARotaryEmbKV. + Forward function for _FusedMLARoPEKVSplit. Args: kv: [seq_len, batch_size, head_num, k_dim + v_dim] @@ -634,7 +713,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, @@ -657,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 @@ -674,7 +755,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 +783,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, @@ -725,14 +806,15 @@ 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_apply_mla_rope_for_kv( +def fused_mla_rope_kv_split( kv: torch.Tensor, k_pos_emb: torch.Tensor, cos: torch.Tensor, @@ -744,9 +826,10 @@ def fused_apply_mla_rope_for_kv( cp_rank: int = 0, cp_size: int = 1, rotary_interleaved: bool = False, + remove_interleaving: 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. @@ -761,13 +844,14 @@ def fused_apply_mla_rope_for_kv( 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] 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, @@ -779,4 +863,12 @@ def fused_apply_mla_rope_for_kv( cp_rank, cp_size, rotary_interleaved, + remove_interleaving, ) + + +# --------------------------------------------------------------------------- +# 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..c97f738771b 100644 --- a/megatron/core/models/common/embeddings/rope_utils.py +++ b/megatron/core/models/common/embeddings/rope_utils.py @@ -95,6 +95,8 @@ def _apply_rotary_pos_emb_bshd( rotary_interleaved: bool = False, 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. @@ -118,6 +120,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 +141,18 @@ 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. + # The existing MLA code is safe because the dot product is permutation-invariant. + 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) + return torch.cat((t, t_pass), dim=-1) @@ -193,6 +212,8 @@ def _apply_rotary_pos_emb_thd( rotary_interleaved: bool = False, 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: @@ -246,6 +267,8 @@ def _apply_rotary_pos_emb_thd( rotary_interleaved=rotary_interleaved, 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 @@ -262,6 +285,8 @@ def _apply_rotary_pos_emb_thd( rotary_interleaved=rotary_interleaved, mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, ).squeeze(1) @@ -273,6 +298,8 @@ def apply_rotary_pos_emb( mscale: float = 1.0, 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 @@ -307,6 +334,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 +361,8 @@ def apply_rotary_pos_emb( rotary_interleaved=config.rotary_interleaved, 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( @@ -338,6 +373,8 @@ def apply_rotary_pos_emb( mla_rotary_interleaved=mla_rotary_interleaved, mscale=mscale, cp_group=cp_group, + inverse=inverse, + mla_output_remove_interleaving=mla_output_remove_interleaving, ) 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 8231a2a3764..2f96fd9fba0 100644 --- a/megatron/core/models/gpt/experimental_attention_variant_module_specs.py +++ b/megatron/core/models/gpt/experimental_attention_variant_module_specs.py @@ -6,6 +6,18 @@ 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.csa import ( + CompressedSparseAttention, + CompressedSparseAttentionSubmodules, + Compressor, + CompressorSubmodules, + CSAIndexer, + CSAIndexerSubmodules, +) +from megatron.core.transformer.experimental_attention_variant.deepseek_v4_hybrid_attention import ( + DSv4HybridSelfAttention, + DSv4HybridSelfAttentionSubmodules, +) from megatron.core.transformer.experimental_attention_variant.dsa import ( DSAIndexer, DSAIndexerSubmodules, @@ -128,6 +140,63 @@ 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_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 +209,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 ac2e3f8bab1..2a932a26272 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, @@ -392,7 +392,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 not InferenceMode.is_active() 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/attention.py b/megatron/core/transformer/attention.py index b27f90c53d0..5019476ab8a 100644 --- a/megatron/core/transformer/attention.py +++ b/megatron/core/transformer/attention.py @@ -291,6 +291,7 @@ def __init__( pg_collection: ProcessGroupCollection | None = None, pp_layer_offset: Optional[int] = None, name: str | None = None, + is_mtp_layer: bool = False, ): """ Args: @@ -301,6 +302,7 @@ def __init__( self.config = config self.layer_number = layer_number self._pp_layer_offset = pp_layer_offset + self.is_mtp_layer = is_mtp_layer self.attn_mask_type = attn_mask_type self.attention_type = attention_type @@ -1386,6 +1388,7 @@ def __init__( pg_collection: ProcessGroupCollection | None = None, pp_layer_offset: Optional[int] = None, name: str | None = None, + is_mtp_layer: bool = False, ): """ Args: @@ -1401,6 +1404,7 @@ def __init__( pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, name=name, + is_mtp_layer=is_mtp_layer, ) self.linear_qkv_out_dim = self.query_projection_size + 2 * self.kv_projection_size @@ -1802,6 +1806,7 @@ def __init__( cp_comm_type: str | None = None, pg_collection: ProcessGroupCollection | None = None, name: str | None = None, + is_mtp_layer: bool = False, ): """ Args: @@ -1816,6 +1821,7 @@ def __init__( cp_comm_type=cp_comm_type, pg_collection=pg_collection, name=name, + 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 new file mode 100644 index 00000000000..1c24ecda5c7 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/csa.py @@ -0,0 +1,770 @@ +# 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.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 nvtx_range_pop, nvtx_range_push + +# --------------------------------------------------------------------------- +# Helper functions for index computation +# --------------------------------------------------------------------------- + + +# 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). + + 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) + + +# 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 +) -> 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(), + remove_interleaving=True, + ) + 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, + mla_output_remove_interleaving=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") + + 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") + + 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 == 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) + """ + + 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, + compress_ratio: int = 0, + ): + 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 + self.compress_ratio = compress_ratio + self.window_size = config.csa_window_size + self.v_head_dim = config.v_head_dim + + self.n_local_heads = config.num_attention_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(self.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 not config.csa_dense_mode + 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..7aa321a3cd1 --- /dev/null +++ b/megatron/core/transformer/experimental_attention_variant/deepseek_v4_hybrid_attention.py @@ -0,0 +1,715 @@ +# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. + + +from dataclasses import dataclass +from typing import NoReturn, Optional, Union + +import torch + +from megatron.core import tensor_parallel +from megatron.core.extensions.transformer_engine import HAVE_TE +from megatron.core.models.common.embeddings import ( + RotaryEmbedding, + YarnRotaryEmbedding, + apply_rotary_pos_emb, +) +from megatron.core.pipeline_parallel.fine_grained_activation_offload import ( + FineGrainedActivationOffloadingInterface as off_interface, +) +from megatron.core.process_groups_config import ProcessGroupCollection +from megatron.core.transformer.attention import Attention +from megatron.core.transformer.enums import AttnMaskType +from megatron.core.transformer.spec_utils import ModuleSpec, build_module +from megatron.core.transformer.torch_norm import LayerNormBuilder +from megatron.core.transformer.transformer_config import MLATransformerConfig +from megatron.core.typed_torch import apply_module +from megatron.core.utils import get_pg_size, is_te_min_version + +try: + from megatron.core.fusions.fused_mla_yarn_rope_apply import fused_mla_rope_inplace +except Exception: + fused_mla_rope_inplace = None + + +if HAVE_TE: + from megatron.core.extensions.transformer_engine import TELinear, set_save_original_input +else: + (TEColumnParallelLinear, TELinear, set_save_original_input) = (None, None, None) + + +@torch.compile +def _q_rms_norm(q: torch.Tensor, eps: float) -> torch.Tensor: + """Fused RMS normalization for query tensor (no learnable weight).""" + return q * torch.rsqrt(q.square().mean(-1, keepdim=True) + eps) + + +@dataclass +class DSv4HybridSelfAttentionSubmodules: + """Submodules for the DSv4HybridAttention layer.""" + + q_layernorm: LayerNormBuilder + kv_layernorm: LayerNormBuilder + + linear_q_down_proj: Union[ModuleSpec, type] = None + linear_q_up_proj: Union[ModuleSpec, type] = None + linear_kv_proj: Union[ModuleSpec, type] = None + core_attention: Union[ModuleSpec, type] = None + linear_proj: Union[ModuleSpec, type] = None + + +class DSv4HybridAttention(Attention): + """DeepSeek-v4 Hybrid Attention layer.""" + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSv4HybridSelfAttentionSubmodules, + layer_number: int, + attn_mask_type: AttnMaskType, + attention_type: str, + cp_comm_type: Optional[str] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + is_mtp_layer: bool = False, + ) -> 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, + is_mtp_layer=is_mtp_layer, + ) + self.config: MLATransformerConfig + + assert ( + get_pg_size(self.pg_collection.tp) == 1 + ), "DSv4 Hybrid Attention only supports TP size 1." + + assert ( + not self.checkpoint_core_attention + ), "Checkpoint core attention is not supported in DSv4 Hybrid Attention." + assert ( + not self.offload_qkv_linear + ), "Offload qkv linear is not supported in DSv4 Hybrid Attention." + + self.query_projection_size = self.config.v_head_dim * self.config.num_attention_heads + + self.q_head_dim = self.config.v_head_dim + + self.key_hidden_size = self.q_head_dim + self.val_hidden_size = self.config.v_head_dim + + self.recompute_up_proj = ( + self.config.recompute_granularity == 'selective' + and "mla_up_proj" in self.config.recompute_modules + ) + self.qkv_up_checkpoint = None + + self.softmax_scale = None + + if is_mtp_layer: + layer_idx = self.config.num_layers + layer_number - 1 + compress_ratio = self.config.csa_compress_ratios[layer_idx] + else: + compress_ratio = self.config.csa_compress_ratios[layer_number - 1] + rope_base = self.config.rotary_base + 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, + "compress_ratio": compress_ratio, + } + self.core_attention = build_module( + submodules.core_attention, + config=self.config, + layer_number=self.layer_number, + attn_mask_type=self.attn_mask_type, + attention_type=self.attention_type, + softmax_scale=self.softmax_scale, + k_channels=self.q_head_dim, + v_channels=self.config.v_head_dim, + cp_comm_type=cp_comm_type, + pg_collection=self.pg_collection, + **core_attn_extra_kwargs, + ) + + # Output. + self.o_local_groups = self.config.o_groups + assert ( + self.query_projection_size % self.config.o_groups == 0 + ), "num_attention_heads * v_head_dim must be divisible by o_groups" + group_proj_in_size = self.query_projection_size // self.config.o_groups + group_proj_out_size = self.config.o_groups * self.config.o_lora_rank + + _linear_o_group_proj = torch.empty( + group_proj_out_size, + group_proj_in_size, + device=torch.cuda.current_device(), + dtype=self.config.params_dtype, + ) + self.config.init_method(_linear_o_group_proj) + self.linear_o_group_proj = torch.nn.Parameter(_linear_o_group_proj) + + linear_proj_in_size = self.config.o_groups * self.config.o_lora_rank + + self.linear_proj = build_module( + submodules.linear_proj, + linear_proj_in_size, + self.config.hidden_size, + config=self.config, + init_method=self.config.output_layer_init_method, + bias=self.config.add_bias_linear, + input_is_parallel=True, + skip_bias_add=True, + is_expert=False, + tp_comm_buffer_name='proj', + tp_group=self.pg_collection.tp, + ) + + if ( + HAVE_TE + and isinstance(self.linear_proj, TELinear) + and ( + ( + self.config.fp8 + and self.config.fp8_recipe != 'delayed' + and is_te_min_version("2.6.0dev0") + ) + or (self.config.fp4 and is_te_min_version("2.7.0.dev0")) + ) + ): + # For fp8/fp4 training, the output of the fused core_attn is saved by itself, and + # linear_proj also saves the quantized tensor of this output. Here we set the + # linear_proj to save the original input tensors to avoid the extra memory usage of + # the quantized tensor. + set_save_original_input(self.linear_proj) + + def forward( + self, + hidden_states, + attention_mask, + key_value_states=None, + inference_context=None, + rotary_pos_emb=None, + rotary_pos_cos=None, + rotary_pos_sin=None, + rotary_pos_cos_sin=None, + attention_bias=None, + packed_seq_params=None, + position_ids=None, + sequence_len_offset=None, + *, + inference_params=None, + ): + """Forward pass for DeepSeek-v4 Hybrid Attention""" + assert ( + rotary_pos_emb is None + ), "Rotary position embeddings should not be passed into DSv4HybridAttention." + assert ( + attention_bias is None + ), "Attention bias should not be passed into DSv4HybridAttention." + assert ( + rotary_pos_cos is None and rotary_pos_sin is None + ), "DSv4HybridAttention does not support Flash Decoding" + assert ( + not rotary_pos_cos_sin + ), "Flash-infer rope has not been tested with DSv4HybridAttention." + assert ( + inference_context is None and inference_params is None + ), "Inference is not supported for DSv4HybridAttention." + + # ===================== + # Query, Key, and Value + # ===================== + # Get the query, key and value tensors based on the type of attention - + # self or cross attn. + query, key, value, q_compressed, kv_compressed = self.get_query_key_value_tensors( + hidden_states, + key_value_states, + position_ids, + packed_seq_params, + inference_context=inference_context, + ) + + # TODO: Currently, TE can only accept contiguous tensors for MLA + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + # ================================== + # core attention computation + # ================================== + # Need corresponding TE change + core_attn_manager = off_interface( + self.offload_core_attention and self.training, query, "core_attn" + ) + with core_attn_manager as query: + core_attn_out = self.core_attention( + query, + key, + value, + attention_mask, + packed_seq_params=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 + # (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 + 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, + remove_interleaving=True, + ) + else: + content_part, rot_part = torch.split( + core_attn_out, [core_attn_out.size(-1) - pos_dim, pos_dim], dim=-1 + ) + rot_part = apply_rotary_pos_emb( + rot_part, + rotary_pos_emb, + self.config, + cu_seqlens=cu_seqlens_kv, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=True, + inverse=True, + mla_output_remove_interleaving=True, + ) + core_attn_out = torch.cat([content_part, rot_part], dim=-1) + core_attn_out = core_attn_out.view(seq_len, core_attn_out.size(1), -1) + + # Grouped output + core_attn_out = core_attn_out.view( + core_attn_out.size(0), core_attn_out.size(1), self.o_local_groups, -1 + ) + wo_a_weight = self.linear_o_group_proj.view( + self.o_local_groups, self.config.o_lora_rank, -1 + ) + core_attn_out = torch.einsum("...gd,grd->...gr", core_attn_out, wo_a_weight) + core_attn_out = core_attn_out.reshape(*core_attn_out.shape[:-2], -1) + + # ================= + # Output. [sq, b, h] + # ================= + attn_proj_manager = off_interface(self.offload_attn_proj, core_attn_out, "attn_proj") + with attn_proj_manager as core_attn_out: + output, bias = self.linear_proj(core_attn_out) + output = attn_proj_manager.group_offload(output, forced_released_tensors=[core_attn_out]) + + return output, bias + + +class DSv4HybridSelfAttention(DSv4HybridAttention): + """DSv4Hybrid Self-attention layer class + + Self-attention layer takes input with size [s, b, h] + and returns output of the same size. + """ + + def __init__( + self, + config: MLATransformerConfig, + submodules: DSv4HybridSelfAttentionSubmodules, + layer_number: int, + attn_mask_type=AttnMaskType.padding, + cp_comm_type: Optional[str] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + is_mtp_layer: bool = False, + ): + 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, + is_mtp_layer=is_mtp_layer, + ) + + q_down_proj_kwargs = {} + if submodules.linear_q_down_proj in [TELinear]: + q_down_proj_kwargs['parallel_mode'] = 'duplicated' + else: + raise ValueError(f"Unsupported linear_q_down_proj: {submodules.linear_q_down_proj}") + + self.linear_q_down_proj = build_module( + submodules.linear_q_down_proj, + self.config.hidden_size, + self.config.q_lora_rank, + config=self.config, + init_method=self.config.init_method, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name='q_down_proj', + skip_weight_param_allocation=False, + tp_group=None, + **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_proj = build_module( + submodules.linear_kv_proj, + self.config.hidden_size, + self.config.v_head_dim, + config=self.config, + init_method=self.config.init_method, + gather_output=False, + bias=False, + skip_bias_add=False, + is_expert=False, + tp_comm_buffer_name='kv_up_proj', + tp_group=pg_collection.tp, + ) + 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 + ), "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 + # ========================================= + # q_compressed: [s, b, q_lora_rank] + q_compressed, _ = self.linear_q_down_proj(hidden_states) + + 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_proj(kv_compressed) + kv = self.kv_layernorm(kv) + + # [num_tokens, qk_pos_emb_head_dim] -> [num_tokens, 1, qk_pos_emb_head_dim] + if k_pos_emb is not None: + k_pos_emb = torch.unsqueeze(k_pos_emb, -2) + + if self.config.apply_rope_fusion: + cp_rank = self.pg_collection.cp.rank() + cp_size = self.pg_collection.cp.size() + query = fused_mla_rope_inplace( + q, + rotary_pos_cos, + rotary_pos_sin, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + cu_seqlens_q, + cp_rank, + cp_size, + remove_interleaving=True, + ) + kv = kv.unsqueeze(-2) + kv = fused_mla_rope_inplace( + kv, + rotary_pos_cos, + rotary_pos_sin, + self.config.qk_head_dim, + self.config.qk_pos_emb_head_dim, + cu_seqlens_q, + cp_rank, + cp_size, + remove_interleaving=True, + ) + 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, + mla_output_remove_interleaving=True, + ) + # query: [num_tokens, n, (qk_head_dim + v_head_dim)] + query = torch.cat([q_no_pe, q_pos_emb], dim=-1) + + pos_dim = self.config.qk_pos_emb_head_dim + kv_no_pe, k_pos_emb = torch.split(kv, [kv.size(-1) - pos_dim, pos_dim], dim=-1) + + # k_pos_emb:[num_tokens, 1, qk_pos_emb_head_dim] + k_pos_emb = apply_rotary_pos_emb( + k_pos_emb, + rotary_pos_emb, + config=self.config, + cu_seqlens=cu_seqlens_kv, + mscale=mscale, + cp_group=self.pg_collection.cp, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + ) + + # Single head: key = value = [num_tokens, 1, v_head_dim] + kv = torch.cat([kv_no_pe, k_pos_emb], dim=-1).unsqueeze(-2) + key = kv + value = kv + + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + return query, key, value + + if self.recompute_up_proj: + quantization = self.config.fp8 or self.config.fp4 + self.qkv_up_checkpoint = tensor_parallel.CheckpointWithoutOutput(fp8=quantization) + query, key, value = self.qkv_up_checkpoint.checkpoint( + qkv_up_proj_and_rope_apply, q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) + else: + query, key, value = qkv_up_proj_and_rope_apply( + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) + + return query, key, value, q_compressed, kv_compressed + + def backward_dw(self) -> NoReturn: + """Execute weight gradient computation""" + self._backward_kv_proj() + self._backward_q_proj() + self._backward_output_proj() + + def _backward_kv_proj(self): + """Computes weight gradients of KV projection layers""" + self.linear_kv_proj.backward_dw() + + def _backward_q_proj(self): + """Computes weight gradients of Q projection layers""" + self.linear_q_down_proj.backward_dw() + self.linear_q_up_proj.backward_dw() + + def _backward_output_proj(self): + """Computes weight gradients of output projection layer""" + self.linear_proj.backward_dw() + + def set_for_recompute_input_layernorm(self): + """Set the attention layer for recompute input_layernorm. Only needed for fp8/fp4.""" + set_save_original_input(self.linear_q_down_proj) + set_save_original_input(self.linear_kv_proj) diff --git a/megatron/core/transformer/experimental_attention_variant/dsa.py b/megatron/core/transformer/experimental_attention_variant/dsa.py index 5c5f77363dc..5d7566b3926 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,62 @@ 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 +268,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 +374,7 @@ def fwd_fused_indexer_loss_naive( loss_coeff, sparse_loss, pg_collection, + causal_mask_override=mask, ) return topk_indices, indexer_loss @@ -355,6 +392,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 +412,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,7 +443,24 @@ 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 ) @@ -409,6 +471,10 @@ def bwd_fused_indexer_loss_naive( # 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,19 +518,31 @@ 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 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() @@ -543,7 +621,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 +634,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 +648,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/multi_latent_attention.py b/megatron/core/transformer/multi_latent_attention.py index 8023f53056e..abd211c681c 100644 --- a/megatron/core/transformer/multi_latent_attention.py +++ b/megatron/core/transformer/multi_latent_attention.py @@ -145,6 +145,7 @@ def __init__( pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, name: str | None = None, + is_mtp_layer: bool = False, ) -> None: # TODO(nschank): Restructure so that the Attention initializer knows which specific # submodules it will construct, so that MLASelfAttentionSubmodules honors that interface. @@ -157,6 +158,7 @@ def __init__( pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, name=name, + is_mtp_layer=is_mtp_layer, ) self.config: MLATransformerConfig @@ -485,6 +487,7 @@ def __init__( pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, name: str | None = None, + is_mtp_layer: bool = False, ): if pg_collection is None: pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -499,6 +502,7 @@ def __init__( pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, name=name, + is_mtp_layer=is_mtp_layer, ) if self.config.q_lora_rank is None: @@ -1231,6 +1235,7 @@ def __init__( pg_collection: Optional[ProcessGroupCollection] = None, pp_layer_offset: Optional[int] = None, name: str | None = None, + is_mtp_layer: bool = False, ): if pg_collection is None: pg_collection = ProcessGroupCollection.use_mpu_process_groups() @@ -1246,6 +1251,7 @@ def __init__( pg_collection=pg_collection, pp_layer_offset=pp_layer_offset, name=name, + 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 ab84abd3a17..0860bf1b281 100644 --- a/megatron/core/transformer/transformer_config.py +++ b/megatron/core/transformer/transformer_config.py @@ -280,8 +280,10 @@ class TransformerConfig(ModelParallelConfig): #################### # attention variant #################### - experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa']] = None - """Type of attention variant to use. Currently support gated_delta_net and dsa.""" + experimental_attention_variant: Optional[Literal['gated_delta_net', 'dsa', 'dsv4_hybrid']] = ( + None + ) + """Type of attention variant to use. Currently support gated_delta_net, dsa, and dsv4_hybrid.""" #################### # DSA @@ -302,6 +304,22 @@ 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. [0, 0, 4, 128, 4, 128, ...].""" + + 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 #################### @@ -1261,6 +1279,22 @@ 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" + mtp_layers = self.mtp_num_layers or 0 + expected_len = self.num_layers + mtp_layers + assert len(self.csa_compress_ratios) == expected_len, ( + f"csa_compress_ratios length ({len(self.csa_compress_ratios)}) must equal " + f"num_layers + mtp_num_layers ({self.num_layers} + {mtp_layers} = {expected_len})" + ) + assert all( + ratio in [0, 4, 128] for ratio in self.csa_compress_ratios + ), "csa_compress_ratios must be 0, 4, or 128" + assert ( + self.tensor_model_parallel_size == 1 + ), "DSv4 Hybrid Attention only supports TP size 1." + assert 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 @@ -2598,10 +2632,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.""" @@ -2639,6 +2675,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 @@ -2657,6 +2699,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/core/transformer/transformer_layer.py b/megatron/core/transformer/transformer_layer.py index ddd1e7d34cd..a9a4d93941b 100644 --- a/megatron/core/transformer/transformer_layer.py +++ b/megatron/core/transformer/transformer_layer.py @@ -357,6 +357,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( @@ -1350,6 +1352,501 @@ def get_layer_norm_weights(self): return +class HyperConnectionTransformerLayer(TransformerLayer): + """A transformer layer with Manifold-Constrained Hyper-Connections (mHC). + + Extends TransformerLayer by adding hyper connection modules around self-attention + and MLP. The n-stream hidden states are aggregated before each sub-layer and + expanded back afterwards using learned mappings (H_pre, H_post, H_res). + + Cross-attention hyper connection is not supported. + """ + + def __init__( + self, + config: TransformerConfig, + submodules: TransformerLayerSubmodules, + layer_number: int = 1, + hidden_dropout: Optional[float] = None, + pg_collection: Optional[ProcessGroupCollection] = None, + vp_stage: Optional[int] = None, + is_mtp_layer: bool = False, + ): + super().__init__( + config=config, + submodules=submodules, + layer_number=layer_number, + 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: + raise ValueError( + "HyperConnectionTransformerLayer does not support cross-attention " + "hyper connections. Use IdentityOp for cross_attention_hyper_connection." + ) + + assert submodules.self_attention_hyper_connection is not IdentityOp, ( + "HyperConnectionTransformerLayer requires self_attention_hyper_connection. " + "Use TransformerLayer instead if hyper connections are not needed." + ) + assert submodules.mlp_hyper_connection is not IdentityOp, ( + "HyperConnectionTransformerLayer requires mlp_hyper_connection. " + "Use TransformerLayer instead if hyper connections are not needed." + ) + + self.self_attention_hyper_connection = build_module( + submodules.self_attention_hyper_connection, + config=self.config, + layer_number=self.layer_number, + ) + + self.mlp_hyper_connection = build_module( + submodules.mlp_hyper_connection, config=self.config, layer_number=self.layer_number + ) + + # When mHC recompute is active, skip checkpointing if the layernorm + # is IdentityOp (fused into TE linear) — there is nothing to recompute. + self.mhc_checkpoint_input_layernorm = not isinstance(self.input_layernorm, IdentityOp) + self.mhc_checkpoint_pre_mlp_layernorm = not isinstance(self.pre_mlp_layernorm, IdentityOp) + + def get_layer_static_inputs(self, seq_length, micro_batch_size): + """Override to produce n-stream hidden_states of shape [s, b, n*C]. + + CUDA graph capture creates static buffers whose shapes are determined by + this method. The base class returns [s, b, C], but mHC layers operate on + n-stream hidden states of shape [s, b, n*C]. + """ + static_inputs = super().get_layer_static_inputs(seq_length, micro_batch_size) + hs = static_inputs["hidden_states"] + n = self.config.num_residual_streams + static_inputs["hidden_states"] = torch.ones( + (hs.shape[0], hs.shape[1], n * self.config.hidden_size), + dtype=hs.dtype, + requires_grad=hs.requires_grad, + device=hs.device, + ) + + # Add input_ids for hash-based MoE routing under CUDA graphs. + # Only add for layers that actually use hash routing, + # since other layers (e.g. on later PP stages) receive input_ids=None. + if ( + self.is_moe_layer + and self.config.moe_n_hash_layers > 0 + and getattr(self.mlp.router, 'is_hash_layer', False) + ): + static_inputs["input_ids"] = torch.zeros( + (micro_batch_size, seq_length), dtype=torch.long, device=torch.cuda.current_device() + ) + + return static_inputs + + def _get_submodules_under_cudagraphs(self): + """Override to include hyper connection modules. + + The base TransformerLayer._get_submodules_under_cudagraphs does not include + self_attention_hyper_connection / mlp_hyper_connection. Their learnable + parameters (mapping_proj, alpha_*, bias) need manual pre-forward hooks + during CUDA graph replay so that parameter all-gathers are triggered. + """ + submodules = super()._get_submodules_under_cudagraphs() + + if not self.config.cuda_graph_scope: + return submodules + + if CudaGraphScope.attn in self.config.cuda_graph_scope: + submodules.append(self.self_attention_hyper_connection) + if (not self.is_moe_layer and CudaGraphScope.mlp in self.config.cuda_graph_scope) or ( + self.is_moe_layer + and ( + CudaGraphScope.moe in self.config.cuda_graph_scope + or CudaGraphScope.moe_router in self.config.cuda_graph_scope + ) + ): + submodules.append(self.mlp_hyper_connection) + return submodules + + def forward(self, *args, **kwargs): + """Forward pass with MHC recompute manager support.""" + kwargs.pop("dynamic_inference_decode_only", None) + + mhc_recompute_manager = getattr(self, '_mhc_recompute_manager', None) + + hidden_states, context = self._forward_attention( + *args, mhc_recompute_manager=mhc_recompute_manager, **kwargs + ) + + output = self._forward_mlp( + hidden_states, + kwargs.get("inference_context", None), + padding_mask=kwargs.get("padding_mask", None), + input_ids=kwargs.get("input_ids", None), + mhc_recompute_manager=mhc_recompute_manager, + ) + return output, context + + def _forward_attention( + self, + hidden_states: Tensor, + attention_mask: Optional[Tensor] = None, + context: Optional[Tensor] = None, + context_mask: Optional[Tensor] = None, + rotary_pos_emb: Optional[Tensor] = None, + rotary_pos_cos: Optional[Tensor] = None, + rotary_pos_sin: Optional[Tensor] = None, + rotary_pos_cos_sin: Optional[Tensor] = None, + attention_bias: Optional[Tensor] = None, + inference_context: Optional[Any] = None, + packed_seq_params: Optional[PackedSeqParams] = None, + sequence_len_offset: Optional[Tensor] = None, + padding_mask: Optional[Tensor] = None, + input_ids: Optional[Tensor] = None, + mhc_recompute_manager: Optional['CheckpointManager'] = None, + *, + inference_params: Optional[Any] = None, + ): + """Forward attention with hyper connection pre/post processing on self-attention.""" + inference_context = deprecate_inference_params(inference_context, inference_params) + + residual = hidden_states + + nvtx_range_push(suffix="self_attention_hyper_connection") + hidden_states, self_attn_h_res, self_attn_hc_h_post = self.self_attention_hyper_connection( + hidden_states, mhc_recompute_manager=mhc_recompute_manager + ) + nvtx_range_pop(suffix="self_attention_hyper_connection") + + # Optional Input Layer norm + checkpoint_input_layernorm = self.recompute_input_layernorm or ( + mhc_recompute_manager is not None and self.mhc_checkpoint_input_layernorm + ) + attn_norm_manager = self.off_interface(self.offload_attn_norm, hidden_states, "attn_norm") + if checkpoint_input_layernorm: + self.input_layernorm_checkpoint = tensor_parallel.CheckpointWithoutOutput( + ckpt_manager=mhc_recompute_manager + ) + with attn_norm_manager as hidden_states: + input_layernorm_output = self.input_layernorm_checkpoint.checkpoint( + self.input_layernorm, hidden_states + ) + else: + with attn_norm_manager as hidden_states: + input_layernorm_output = self.input_layernorm(hidden_states) + + # Self attention. + nvtx_range_push(suffix="self_attention") + attention_output_with_bias = self.self_attention( + input_layernorm_output, + attention_mask=attention_mask, + inference_context=inference_context, + rotary_pos_emb=rotary_pos_emb, + rotary_pos_cos=rotary_pos_cos, + rotary_pos_sin=rotary_pos_sin, + rotary_pos_cos_sin=rotary_pos_cos_sin, + attention_bias=attention_bias, + packed_seq_params=packed_seq_params, + sequence_len_offset=sequence_len_offset, + ) + nvtx_range_pop(suffix="self_attention") + + if checkpoint_input_layernorm: + self.input_layernorm_checkpoint.discard_output_and_register_recompute( + attention_output_with_bias[0] + ) + + nvtx_range_push(suffix="self_attention_fused_h_res_h_post_bda") + with self.bias_dropout_add_exec_handler(): + hidden_states = self.self_attention_hyper_connection.fused_h_res_h_post_bda( + self_attn_h_res, + residual, + self_attn_hc_h_post, + attention_output_with_bias, + self.hidden_dropout, + self.training, + self.config.bias_dropout_fusion, + mhc_recompute_manager, + ) + nvtx_range_pop(suffix="self_attention_fused_h_res_h_post_bda") + + hidden_states = attn_norm_manager.group_offload(hidden_states) + + # Cross-attention (no hyper connection support). + residual = hidden_states + pre_cross_attn_layernorm_output = self.pre_cross_attn_layernorm(hidden_states) + + attention_output_with_bias = self.cross_attention( + pre_cross_attn_layernorm_output, + attention_mask=context_mask, + key_value_states=context, + inference_context=inference_context, + ) + + if isinstance(attention_output_with_bias, dict) and "context" in attention_output_with_bias: + context = attention_output_with_bias["context"] + + with self.bias_dropout_add_exec_handler(): + hidden_states = self.cross_attn_bda(self.training, self.config.bias_dropout_fusion)( + attention_output_with_bias, residual, self.hidden_dropout + ) + + return hidden_states, context + + def _forward_mlp( + self, + hidden_states, + inference_context=None, + padding_mask=None, + input_ids=None, + mhc_recompute_manager: Optional['CheckpointManager'] = None, + ): + """Forward MLP with hyper connection pre/post processing.""" + is_last_in_recompute_block = bool( + mhc_recompute_manager is not None + and getattr(mhc_recompute_manager, "is_last_layer_in_recompute_block", False) + ) + mhc_mlp_bda_manager = None if is_last_in_recompute_block else mhc_recompute_manager + + residual = hidden_states + + nvtx_range_push(suffix="mlp_hyper_connection") + hidden_states, mlp_h_res, mlp_hc_h_post = self.mlp_hyper_connection( + hidden_states, mhc_recompute_manager=mhc_recompute_manager + ) + nvtx_range_pop(suffix="mlp_hyper_connection") + + # Optional Layer norm post the cross-attention. + checkpoint_pre_mlp_layernorm = self.recompute_pre_mlp_layernorm or ( + mhc_recompute_manager is not None and self.mhc_checkpoint_pre_mlp_layernorm + ) + self.mlp_norm_manager = self.off_interface(self.offload_mlp_norm, hidden_states, "mlp_norm") + if checkpoint_pre_mlp_layernorm: + self.pre_mlp_norm_checkpoint = tensor_parallel.CheckpointWithoutOutput( + ckpt_manager=mhc_recompute_manager + ) + with self.mlp_norm_manager as hidden_states: + pre_mlp_layernorm_output = self.pre_mlp_norm_checkpoint.checkpoint( + self.pre_mlp_layernorm, hidden_states + ) + else: + with self.mlp_norm_manager as hidden_states: + pre_mlp_layernorm_output = self.pre_mlp_layernorm(hidden_states) + + nvtx_range_push(suffix="mlp") + should_chunk_mlp_for_prefill = ( + self.config.mlp_chunks_for_prefill > 1 + and inference_context is not None + and not inference_context.is_decode_only() + and not isinstance(self.mlp, IdentityOp) + and not self.config.transformer_impl == "inference_optimized" + ) + + moe_kwargs = {} + if self.is_moe_layer and input_ids is not None: + moe_kwargs['input_ids'] = input_ids + + if self.recompute_mlp: + if self.config.fp8 or self.config.fp4: + from megatron.core.extensions.transformer_engine import te_checkpoint + + mlp_output_with_bias = te_checkpoint( + self.mlp, + False, + tensor_parallel.random.get_cuda_rng_tracker, + self.pg_collection.tp, + pre_mlp_layernorm_output, + padding_mask=padding_mask, + **moe_kwargs, + ) + else: + mlp_output_with_bias = tensor_parallel.checkpoint( + functools.partial(self.mlp, padding_mask=padding_mask, **moe_kwargs), + False, + pre_mlp_layernorm_output, + ) + elif should_chunk_mlp_for_prefill: + num_chunks = min(self.config.mlp_chunks_for_prefill, pre_mlp_layernorm_output.shape[0]) + chunks = pre_mlp_layernorm_output.chunk(num_chunks, dim=0) + outputs = [self.mlp(chunk) for chunk in chunks] + mlp_output = torch.cat([out for out, _ in outputs], dim=0) + bias_chunks = [bias for _, bias in outputs if bias is not None] + bias_output = torch.stack(bias_chunks, dim=0).sum(dim=0) if bias_chunks else None + mlp_output_with_bias = (mlp_output, bias_output) + else: + mlp_output_with_bias = self.mlp( + pre_mlp_layernorm_output, padding_mask=padding_mask, **moe_kwargs + ) + + nvtx_range_pop(suffix="mlp") + + # During TE CUDA graph partial MoE capture, skip HC post-processing and return + # intermediate outputs + HC state. The post-processing will be done during replay. + if ( + self.is_moe_layer + and self.config.cuda_graph_impl == "transformer_engine" + and self.training + and is_graph_capturing() + and CudaGraphScope.moe_router in self.config.cuda_graph_scope + ): + if self.recompute_pre_mlp_layernorm or ( + mhc_recompute_manager is not None and self.mhc_checkpoint_pre_mlp_layernorm + ): + for tensor in mlp_output_with_bias: + self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute(tensor) + # Append HC state (mlp_hc_h_post, mlp_h_res, residual) for replay. + return list(mlp_output_with_bias) + [mlp_hc_h_post, mlp_h_res, residual] + + return self._forward_post_mlp_with_fused_hyper_connection( + mlp_output_with_bias, mlp_h_res, residual, mlp_hc_h_post, mhc_mlp_bda_manager + ) + + def _forward_post_mlp_with_fused_hyper_connection( + self, + mlp_output_with_bias, + mlp_h_res, + residual, + mlp_hc_h_post, + mhc_mlp_bda_recompute_manager: Optional['CheckpointManager'] = None, + ): + """ + Perform operations after the MLP computation with fused hyper connection kernel. + + This method uses the fused kernel combining apply_h_res, apply_h_post and bias-dropout-add. + + Args: + mlp_output_with_bias (Tensor): Output tensor of the MLP layer with bias. + mlp_h_res (Tensor): [s, b, n, n] - residual mixing matrix from hyper connection. + residual (Tensor): [s, b, n*C] - original residual (n-stream hidden states). + mlp_hc_h_post (Tensor): [s, b, n] - expansion weights from hyper connection. + mhc_recompute_manager: Optional CheckpointManager for checkpoint management. + + Returns: + output (Tensor): Transformed hidden states of shape [s, b, h]. + """ + if self.recompute_pre_mlp_layernorm or ( + mhc_mlp_bda_recompute_manager is not None and self.mhc_checkpoint_pre_mlp_layernorm + ): + self.pre_mlp_norm_checkpoint.discard_output_and_register_recompute( + mlp_output_with_bias[0] + ) + + nvtx_range_push(suffix="mlp_fused_h_res_h_post_bda") + with self.bias_dropout_add_exec_handler(): + hidden_states = self.mlp_hyper_connection.fused_h_res_h_post_bda( + mlp_h_res, + residual, + mlp_hc_h_post, + mlp_output_with_bias, + self.hidden_dropout, + self.training, + self.config.bias_dropout_fusion, + mhc_mlp_bda_recompute_manager, + ) + nvtx_range_pop(suffix="mlp_fused_h_res_h_post_bda") + + hidden_states = self.mlp_norm_manager.group_offload(hidden_states) + + output = make_viewless_tensor( + inp=hidden_states, requires_grad=hidden_states.requires_grad, keep_graph=True + ) + return output + + def _te_cuda_graph_replay_impl(self, args, kwargs, context): + """Implementation of _te_cuda_graph_replay with hyper connection support. + + Overrides the parent's _te_cuda_graph_replay_impl so that the + delay_offload_until_cuda_graph lifecycle (enter_replay/exit_replay) in + the parent's _te_cuda_graph_replay is preserved. + + During MoE partial CUDA graph capture, the graph outputs include HC state + (mlp_hc_h_post, mlp_h_res) in addition to the base class outputs. This method + extracts the HC state and uses it for post-processing after resuming the MoE forward. + """ + cuda_graph_output = list( + GraphableMegatronModule._te_cuda_graph_replay(self, *args, **kwargs) + ) + + # Flush delayed offload groups from previous layers after graph replay. + if self.config.delay_offload_until_cuda_graph: + self.off_interface.flush_delayed_groups() + + if kwargs.get('context') is not None: + context = cuda_graph_output.pop() + + if ( + not self.config.cuda_graph_scope + or (not self.is_moe_layer and CudaGraphScope.mlp in self.config.cuda_graph_scope) + or (self.is_moe_layer and CudaGraphScope.moe in self.config.cuda_graph_scope) + ): + assert len(cuda_graph_output) == 1, "CUDA Graph output should be the layer output." + output = cuda_graph_output.pop() + assert ( + not self.config.overlap_moe_expert_parallel_comm + ), "EP overlap must be \ + disabled when CUDA graph captures the whole MLP/MoE part." + elif self.is_moe_layer and CudaGraphScope.moe_router in self.config.cuda_graph_scope: + # Pop HC state (appended during capture in _forward_mlp). + residual = cuda_graph_output.pop() + mlp_h_res = cuda_graph_output.pop() + mlp_hc_h_post = cuda_graph_output.pop() + + shared_expert_output, routing_map = None, None + if ( + self.config.moe_shared_expert_intermediate_size is not None + and not self.config.moe_shared_expert_overlap + ): + shared_expert_output = cuda_graph_output.pop() + + if CudaGraphScope.moe_preprocess in self.config.cuda_graph_scope: + (hidden_states, probs), attr_outputs = ( + cuda_graph_output[:2], + cuda_graph_output[2:], + ) + valid_cudagraph_attrs = self.mlp.token_dispatcher.valid_cudagraph_attrs + assert len(attr_outputs) == len( + valid_cudagraph_attrs + ), f"attr_outputs: {len(attr_outputs)} != {len(valid_cudagraph_attrs)}" + for i, attr_name in enumerate(valid_cudagraph_attrs): + self.mlp.token_dispatcher.set_cudagraph_attr(attr_name, attr_outputs[i]) + else: + assert len(cuda_graph_output) == 3, ( + "CUDA graph output should be [hidden_states, probs, routing_map], " + f"but got {len(cuda_graph_output)} elements" + ) + hidden_states, probs, routing_map = cuda_graph_output + + # Resume the MoELayer forward pass from the end of the CUDA graph scope. + nvtx_range_push(suffix="mlp") + self.mlp.cudagraph_tensor_store.set( + hidden_states=hidden_states, + probs=probs, + routing_map=routing_map, + shared_expert_output=shared_expert_output, + ) + # If EP overlap is enabled, remaining of mlp will be called as fine_grained_callables + # and should be skipped here. + if self.config.overlap_moe_expert_parallel_comm: + probs, routing_map = self.mlp.route(hidden_states) + hidden_states, probs = self.mlp.preprocess(hidden_states, probs, routing_map) + nvtx_range_pop(suffix="mlp") + return residual, hidden_states, probs, shared_expert_output + mlp_output_with_bias = self.mlp(hidden_states) + self.mlp.cudagraph_tensor_store.clear() + nvtx_range_pop(suffix="mlp") + + # HC post-processing with fused h_res, h_post and BDA. + recompute_pre_mlp_layernorm = self.recompute_pre_mlp_layernorm + self.recompute_pre_mlp_layernorm = False + output = self._forward_post_mlp_with_fused_hyper_connection( + mlp_output_with_bias, mlp_h_res, residual, mlp_hc_h_post + ) + self.recompute_pre_mlp_layernorm = recompute_pre_mlp_layernorm + else: + output = self._forward_mlp(*cuda_graph_output, input_ids=kwargs.get("input_ids", None)) + return output, context + + class MoETransformerLayer(TransformerLayer): """ A Transformer layer specialized for Mixture-of-Experts (MoE) architectures. diff --git a/megatron/training/arguments.py b/megatron/training/arguments.py index e7415dc3019..a0ad5e3eeda 100644 --- a/megatron/training/arguments.py +++ b/megatron/training/arguments.py @@ -334,6 +334,23 @@ def no_rope_freq_type(x): # it's a single int but in str 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.: + "[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. + """ + 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. @@ -2038,6 +2055,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", @@ -3181,6 +3199,10 @@ def _add_mla_args(parser): help="Mscale for YaRN RoPE in multi-latent attention.") group.add_argument('--mscale-all-dim', type=float, default=0.0, help="Mscale all dimensions for YaRN RoPE in multi-latent attention.") + group.add_argument('--o-groups', type=int, default=8, + help="Number of groups for grouped 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', default=False, help="If set caches the mla down projected latents with mla flash decode.") group.add_argument( @@ -3205,6 +3227,13 @@ def _add_experimental_attention_variant_args(parser): 'where 1 indicates an LA layer and 0 indicates a SDPA layer. ' 'Examples: "([0]+[1]*23)": 1 SDPA layer followed by 23 LA layers, ' '"([1]*3+[0]*2)*2": Three LA layers followed by two SDPA layers, repeated twice.') + group.add_argument('--csa-compress-ratios', type=compress_ratios_type, default=None, + help='Per-layer compress ratios for compressed sparse attention. ' + 'Accepts a string containing a Python list expression, e.g.: ' + '"[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: 0, 4, 128). ' + 'The list length must equal num_layers.') return parser def _add_heterogeneous_args(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 9059d0157aa..c80430105ca 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, remove_interleaving=False): + assert fused_mla_rope_inplace is not None num_heads = 32 q_dim = 128 emb_dim = 64 @@ -104,12 +104,21 @@ def _test_fused_apply_mla_rope_for_q(input_format): mscale=mscale, 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_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, + remove_interleaving=remove_interleaving, ) fused_output.backward(fused_bwd_input, retain_graph=True) @@ -128,8 +137,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, remove_interleaving=False): + assert fused_mla_rope_kv_split is not None num_heads = 32 k_dim = 128 v_dim = 128 @@ -203,6 +212,7 @@ def _test_fused_apply_mla_rope_for_kv(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) @@ -214,7 +224,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, @@ -223,6 +233,7 @@ def _test_fused_apply_mla_rope_for_kv(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) @@ -265,8 +276,16 @@ class TestFusedApplyMLARope: def test_forward_backward_for_q(self, input_format): _test_fused_apply_mla_rope_for_q(input_format) - def test_forward_backward_for_kv(self, input_format): - _test_fused_apply_mla_rope_for_kv(input_format) + @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 + ) + + @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: diff --git a/tests/unit_tests/models/test_hybrid_moe_model.py b/tests/unit_tests/models/test_hybrid_moe_model.py index c781dd11dd8..d2ad83f926f 100644 --- a/tests/unit_tests/models/test_hybrid_moe_model.py +++ b/tests/unit_tests/models/test_hybrid_moe_model.py @@ -68,6 +68,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_modules": [], 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..83c153d698e --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_attention_variant_csa.py @@ -0,0 +1,872 @@ +# 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 + +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.""" + 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 = [0] * 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=[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 + + 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, + compress_ratio=0, + ) + + 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, + compress_ratio=compress_ratio, + ).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, + compress_ratio=compress_ratio, + ).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, + compress_ratio=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) + ) + 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, + compress_ratio=compress_ratio, + ).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, + compress_ratio=4, + ).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, + compress_ratio=4, + ).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..5f9a3a74440 --- /dev/null +++ b/tests/unit_tests/transformer/experimental_attention_variant/test_dsv4_hybrid_attention.py @@ -0,0 +1,416 @@ +# 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 = [0, 4, 128, 4] + return MLATransformerConfig( + num_layers=num_layers, + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + use_cpu_initialization=True, + bf16=True, + params_dtype=torch.bfloat16, + add_bias_linear=False, + tensor_model_parallel_size=tensor_model_parallel_size, + sequence_parallel=sequence_parallel, + q_lora_rank=q_lora_rank, + kv_lora_rank=v_head_dim - qk_pos_emb_head_dim, + qk_head_dim=v_head_dim - qk_pos_emb_head_dim, + qk_pos_emb_head_dim=qk_pos_emb_head_dim, + v_head_dim=v_head_dim, + o_groups=o_groups, + o_lora_rank=o_lora_rank, + rope_type='rope', + rotary_base=10000, + rotary_percent=1.0, + multi_latent_attention=True, + experimental_attention_variant='dsv4_hybrid', + csa_compress_ratios=csa_compress_ratios, + csa_window_size=csa_window_size, + dsa_indexer_n_heads=dsa_indexer_n_heads, + dsa_indexer_head_dim=dsa_indexer_head_dim, + dsa_indexer_topk=dsa_indexer_topk, + dsa_indexer_loss_coeff=dsa_indexer_loss_coeff, + ) + + +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, backend=TESpecProvider()) + + +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_proj') + assert hasattr(attn, 'linear_proj') + assert hasattr(attn, 'linear_o_group_proj') + assert hasattr(attn, 'core_attention') + assert hasattr(attn, 'q_layernorm') + assert hasattr(attn, 'kv_layernorm') + + def test_q_head_dim_equals_v_head_dim(self): + """q_head_dim must equal v_head_dim for DSv4 hybrid.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + config = _make_config() + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=1, pg_collection=pg) + + assert attn.q_head_dim == config.v_head_dim + + @pytest.mark.parametrize("layer_number", [1, 2, 3, 4]) + def test_rope_base_varies_with_compress_ratio(self, layer_number): + """Layers with compress_ratio > 1 should use csa_compress_rotary_base.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + ratios = [0, 4, 128, 4] + config = _make_config(csa_compress_ratios=ratios) + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=layer_number, pg_collection=pg) + + ratio = ratios[layer_number - 1] + if ratio > 1: + expected_base = config.csa_compress_rotary_base + else: + expected_base = config.rotary_base + + # inv_freq is derived from rotary_base; verify the correct base was used + dim = config.qk_pos_emb_head_dim + recomputed_inv_freq = 1.0 / ( + expected_base ** (torch.arange(0, dim, 2, dtype=torch.float32) / dim) + ) + assert torch.allclose( + attn.rotary_pos_emb.inv_freq.cpu(), recomputed_inv_freq, rtol=1e-5, atol=1e-5 + ) + + +# =========================================================================== +# Forward / backward tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridAttentionForwardBackward: + """Test forward and backward passes of DSv4HybridSelfAttention.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.config = _make_config(dsa_indexer_loss_coeff=1.0) + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + @pytest.mark.parametrize("layer_number", [1, 2, 3, 4]) + def test_forward_output_shape(self, layer_number): + """Forward should produce [sq, b, hidden_size] output.""" + seq_len = 256 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention( + self.config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + output, bias = attn(hidden_states=hidden, attention_mask=None) + + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + assert output.dtype == torch.bfloat16 + assert not torch.isnan(output).any() + + @pytest.mark.parametrize("layer_number", [1, 2]) + def test_backward_gradient_flow(self, layer_number): + """Backward should produce gradients for all trainable parameters.""" + seq_len = 256 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention( + self.config, layer_number=layer_number, pg_collection=self.pg + ).cuda() + attn.train() + + hidden = ( + torch.randn(seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16) + .cuda() + .requires_grad_(True) + ) + + output, bias = attn(hidden_states=hidden, attention_mask=None) + loss = output.sum() + loss.backward() + + assert hidden.grad is not None, "No gradient on hidden_states" + for name, param in attn.named_parameters(): + if param.requires_grad: + assert param.grad is not None, f"No gradient for parameter {name}" + + def test_eval_mode(self): + """Forward should work in eval mode.""" + seq_len = 128 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + attn.eval() + + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + with torch.no_grad(): + output, bias = attn(hidden_states=hidden, attention_mask=None) + + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + assert not torch.isnan(output).any() + + def test_different_seq_lengths(self): + """Forward should handle various sequence lengths.""" + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=2, pg_collection=self.pg).cuda() + + for seq_len in [64, 128, 256]: + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + output, bias = attn(hidden_states=hidden, attention_mask=None) + assert output.shape == (seq_len, batch_size, self.config.hidden_size) + + +# =========================================================================== +# get_query_key_value_tensors tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridQKV: + """Test get_query_key_value_tensors internals.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self, request): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + cls = request.cls + cls.config = _make_config() + cls.pg = ProcessGroupCollection.use_mpu_process_groups() + + yield + Utils.destroy_model_parallel() + + def test_qkv_shapes(self): + """Query, key, value should have correct shapes.""" + seq_len = 64 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + q, k, v, q_compressed, kv_compressed = attn.get_query_key_value_tensors(hidden) + + n_heads = self.config.num_attention_heads + v_dim = self.config.v_head_dim + + assert q.shape == (seq_len, batch_size, n_heads, v_dim) + # key and value are single-head (MQA-style) with an extra head dim + assert k.shape[-1] == v_dim + assert v.shape[-1] == v_dim + + def test_key_equals_value(self): + """In the wkv path, key and value should be the same tensor.""" + seq_len = 64 + batch_size = 2 + + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + attn = _build_attention(self.config, layer_number=1, pg_collection=self.pg).cuda() + hidden = torch.randn( + seq_len, batch_size, self.config.hidden_size, dtype=torch.bfloat16 + ).cuda() + + q, k, v, _, _ = attn.get_query_key_value_tensors(hidden) + assert torch.equal(k, v), "key and value should be identical in wkv path" + + +# =========================================================================== +# Grouped output projection tests +# =========================================================================== + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +@pytest.mark.skipif(not HAVE_TE, reason="transformer_engine not available") +class TestDSv4HybridGroupedOutput: + """Test that grouped output projection (wo_a) parameters are created.""" + + @pytest.fixture(scope='class', autouse=True) + def setup_method(self): + Utils.initialize_model_parallel( + tensor_model_parallel_size=1, pipeline_model_parallel_size=1 + ) + yield + Utils.destroy_model_parallel() + + def test_o_group_proj_shape(self): + """linear_o_group_proj should have the correct shape.""" + torch.manual_seed(_SEED) + model_parallel_cuda_manual_seed(_SEED) + + o_groups = 8 + o_lora_rank = 64 + config = _make_config(o_groups=o_groups, o_lora_rank=o_lora_rank) + pg = ProcessGroupCollection.use_mpu_process_groups() + attn = _build_attention(config, layer_number=1, pg_collection=pg) + + expected_out = o_groups * o_lora_rank + expected_in = (config.v_head_dim * config.num_attention_heads) // o_groups + assert attn.linear_o_group_proj.shape == (expected_out, expected_in) + assert attn.linear_o_group_proj.requires_grad