diff --git a/src/mcore_bridge/model/gpts/deepseek_v4.py b/src/mcore_bridge/model/gpts/deepseek_v4.py index 03d0641..2d8e625 100644 --- a/src/mcore_bridge/model/gpts/deepseek_v4.py +++ b/src/mcore_bridge/model/gpts/deepseek_v4.py @@ -60,6 +60,38 @@ def _patch_YarnRotaryEmbedding(config): delattr(config, attr) +def _apply_mla_rope(t, freqs, *, config, cu_seqlens, cp_group, inverse=False): + """Apply DSv4's MLA RoPE to a tensor whose frequencies are already expanded per token. + + `GPTModel` pre-indexes the rotary table by `position_ids`, so `freqs` is row-aligned with + `t`: row i holds the frequency of token i. That holds for every layout DSv4 supports -- + unpacked, packed (thd), and packed CP: swift CP-splits `position_ids` with the same + partition mode as the hidden states, so the pre-indexed frequencies come out rank-local + while still carrying absolute positions. The multiply is therefore purely elementwise and + needs no `cu_seqlens`-based segment alignment, under any `cp_partition_mode`. + + Enforcing that invariant here matters: when it does not hold, the generic + `apply_rotary_pos_emb` thd path re-derives positions from `cu_seqlens` assuming a *zigzag* + CP split, which is wrong for DSv4's contiguous split and would corrupt positions silently. + Asserting row alignment turns any future layout change into an immediate, explicit failure. + """ + assert freqs.shape[0] == t.shape[0], ( + f'DSv4 MLA RoPE expects per-token frequencies row-aligned with the input, got ' + f'freqs.shape[0]={freqs.shape[0]} vs tokens={t.shape[0]}. `GPTModel` must pre-index the ' + 'rotary table by `position_ids` (requires `apply_rope_fusion=False`), and under CP the ' + '`position_ids` must be split with the same partition mode as the hidden states.') + return apply_rotary_pos_emb( + t, + freqs, + config=config, + cu_seqlens=cu_seqlens, + cp_group=cp_group, + mla_rotary_interleaved=True, + mla_output_remove_interleaving=True, + inverse=inverse, + ) + + class DSv4HybridSelfAttention(McoreDSv4HybridSelfAttention): def __init__(self, config, *args, **kwargs): @@ -153,6 +185,15 @@ def qkv_up_proj_and_rope_apply(q_compressed, 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. + + RoPE frequency layout: `GPTModel` pre-indexes the rotary table by `position_ids` + (see gpt_model.py, the `not apply_rope_fusion` branch), so `rotary_pos_emb` here is + already expanded per token -- row i belongs to token i -- rather than being a + position->frequency lookup table. Every RoPE call below is therefore an elementwise + multiply; see `_apply_mla_rope` for why that invariant is asserted. Under CP the + frequencies arrive rank-local because `position_ids` is split alongside the hidden + states, which is also why the boundary rows carry their own frequencies + (`boundary_rotary_pos_emb`) instead of being re-derived from positions. """ # q_compressed: [num_tokens, q_lora_rank] # q: [num_tokens, n * (qk_head_dim + qk_pos_emb_head_dim)] @@ -166,6 +207,8 @@ def qkv_up_proj_and_rope_apply(q_compressed, if boundary_kv_compressed is not None: boundary_rows = boundary_kv_compressed.shape[0] kv_projection_input = torch.cat([boundary_kv_compressed, kv_compressed], dim=0) + # The boundary rows precede this rank's block, so their frequencies must precede + # too -- keeping kv_rotary_pos_emb row-aligned with kv_projection_input. kv_rotary_pos_emb = torch.cat([boundary_rotary_pos_emb, rotary_pos_emb], dim=0) else: kv_projection_input = kv_compressed @@ -182,14 +225,12 @@ def qkv_up_proj_and_rope_apply(q_compressed, # 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 = _apply_mla_rope( q_pos_emb, rotary_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q, 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) @@ -197,14 +238,12 @@ def qkv_up_proj_and_rope_apply(q_compressed, 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 = _apply_mla_rope( k_pos_emb, kv_rotary_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv, 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] @@ -384,15 +423,13 @@ def forward( rot_part_in = rot_part.squeeze(1) else: rot_part_in = rot_part - rot_part_out = apply_rotary_pos_emb( + rot_part_out = _apply_mla_rope( rot_part_in, rotary_pos_emb, - self.config, + config=self.config, cu_seqlens=cu_seqlens_kv, cp_group=self.pg_collection.cp, - mla_rotary_interleaved=True, inverse=True, - mla_output_remove_interleaving=True, ) if packed_seq: rot_part = rot_part_out.unsqueeze(1) diff --git a/src/mcore_bridge/patcher.py b/src/mcore_bridge/patcher.py index abf8ad0..df5347e 100644 --- a/src/mcore_bridge/patcher.py +++ b/src/mcore_bridge/patcher.py @@ -197,16 +197,17 @@ def forward(self, position_ids, mrope_section: List[int], mrope_interleaved: boo def _apply_rotary_pos_emb_thd(t: torch.Tensor, cu_seqlens: torch.Tensor, freqs: torch.Tensor, *args, **kwargs) -> torch.Tensor: cp_group = kwargs.pop('cp_group', None) - if cp_group is not None: - cp_size = cp_group.size() - else: - cp_size = mpu.get_context_parallel_world_size() + if cp_group is None: cp_group = mpu.get_context_parallel_group() - cu_seqlens_for_batched = cu_seqlens // cp_size - use_batched_rope = (freqs.dim() >= 1 and freqs.shape[0] == cu_seqlens_for_batched[-1]).item() - # The determination of mla_output_remove_interleaving: a quick solution for identifying deepseek_v4 - # (TODO: refactor) - if not use_batched_rope and not kwargs.get('mla_output_remove_interleaving', False): + # The fast path below reinterprets the thd tensor as bshd and multiplies it by `freqs` + # directly, which is only valid when `freqs` is already expanded per token. Compare against + # the token count of `t` itself rather than deriving one from `cu_seqlens`: callers that + # change the sequence length -- e.g. deepseek_v4's CSA compressor, which shortens kv and + # passes the compressed cu_seqlens -- must fall back to the upstream thd kernel, which + # aligns freqs per segment. For an ordinary packed batch both are equal, so the fast path + # is preserved. + use_batched_rope = freqs.dim() >= 1 and freqs.shape[0] == t.shape[0] + if not use_batched_rope: logger.warning_once('Using non-batched RoPE, which may affect performance.') return _origin_apply_rotary_pos_emb_thd(t, cu_seqlens, freqs, *args, cp_group=cp_group, **kwargs)