diff --git a/src/mcore_bridge/model/gpt_model.py b/src/mcore_bridge/model/gpt_model.py index 367a91f..5294610 100644 --- a/src/mcore_bridge/model/gpt_model.py +++ b/src/mcore_bridge/model/gpt_model.py @@ -311,6 +311,8 @@ def forward( if self.config.moe_n_hash_layers > 0 or getattr(self.config, 'ple_layer_ids', None): extra_block_kwargs['input_ids'] = input_ids + if getattr(self.config, 'indexer_n_heads', None) is not None: + extra_block_kwargs['position_ids'] = position_ids # Run decoder. decoder_output = self.decoder( @@ -334,6 +336,9 @@ def forward( # MTP: https://github.com/NVIDIA/Megatron-LM/issues/1661 extra_block_kwargs.pop('input_ids', None) + # self.mtp below takes position_ids explicitly; leaving it here would + # collide with the explicit kwarg. + extra_block_kwargs.pop('position_ids', None) return self._postprocess( hidden_states=hidden_states, input_ids=input_ids, diff --git a/src/mcore_bridge/model/gpts/qwen4_exp.py b/src/mcore_bridge/model/gpts/qwen4_exp.py index 4852f92..86649e6 100644 --- a/src/mcore_bridge/model/gpts/qwen4_exp.py +++ b/src/mcore_bridge/model/gpts/qwen4_exp.py @@ -204,14 +204,17 @@ def _qsa_select_indices_thd(self, hidden_states, rotary_pos_emb, packed_seq_para psp_for_cp.cu_seqlens_q_padded = packed_seq_params.cu_seqlens_q_padded.to(torch.int32) else: psp_for_cp = packed_seq_params + local_len = hidden_states.shape[0] hidden_states = reconstruct_tensor_cp(hidden_states, psp_for_cp, dim=0) # Per-token rotary angles. Without rope fusion gpt_model already indexes # the freq table by position_ids, so what arrives is per-token (zigzag- # sharded under CP -- undo it like hidden). With fusion the raw table # arrives and must be indexed by the (CP-reconstructed) per-doc ids. freqs = rotary_pos_emb - fused_table = freqs.shape[0] != hidden_states.shape[0] if self.config.context_parallel_size > 1: + fused_table = ( + self.config.position_embedding_type != 'mrope' + and (self.config.apply_rope_fusion or freqs.shape[0] != local_len)) if fused_table: if position_ids is None: raise RuntimeError('QSA thd selection under CP needs position_ids to index the fused rotary ' @@ -221,15 +224,12 @@ def _qsa_select_indices_thd(self, hidden_states, rotary_pos_emb, packed_seq_para freqs = freqs[pos.reshape(-1)] else: freqs = reconstruct_tensor_cp(freqs, psp_for_cp, dim=0) - elif fused_table: - # Same problem without CP, and here there is no reconstruct step to hide - # behind: the indexer would slice the raw table's first T rows, treating - # row i as token i's angle. In a packed batch token i sits at in-document - # position i - cu[doc], so those angles belong to the wrong positions -- - # silently degrading the selection instead of failing. - raise RuntimeError(f'QSA thd selection got a fused rotary table ({freqs.shape[0]} rows for ' - f'{hidden_states.shape[0]} tokens): apply_rope_fusion=true hands over the raw ' - 'table rather than per-token freqs. Set --apply_rope_fusion false.') + else: + fused_table = freqs.shape[0] != hidden_states.shape[0] + if fused_table: + raise RuntimeError(f'QSA thd selection got a fused rotary table ({freqs.shape[0]} rows for ' + f'{hidden_states.shape[0]} tokens): apply_rope_fusion=true hands over the raw ' + 'table rather than per-token freqs. Set --apply_rope_fusion false.') # the CP reconstruct (like TE's thd kernels) works in the padded pack # space, so align against the padded cu when present cu = packed_seq_params.cu_seqlens_q_padded @@ -240,7 +240,8 @@ def _qsa_select_indices_thd(self, hidden_states, rotary_pos_emb, packed_seq_para 'boundaries, but it is missing.') cu = Qwen4ExpTextPLELayer._normalize_cu_seqlens(cu, hidden_states.shape[0]) hidden_tok = hidden_states.reshape(hidden_states.shape[0], -1) - return self.self_attention.indexer.select_token_indices_thd(hidden_tok, freqs, cu) + return self.self_attention.indexer.select_token_indices_thd( + hidden_tok, freqs, cu, force_materialize=self.config.context_parallel_size > 1) def _qsa_select_mask(self, hidden_states, attn_kwargs): # Bool-mask QSA on TE's `arbitrary` mask. Only reached for sbhd with CP==1 -- diff --git a/src/mcore_bridge/model/modules/qsa_indexer.py b/src/mcore_bridge/model/modules/qsa_indexer.py index d81af5d..f739daa 100644 --- a/src/mcore_bridge/model/modules/qsa_indexer.py +++ b/src/mcore_bridge/model/modules/qsa_indexer.py @@ -267,8 +267,11 @@ def selection_as_token_indices(self, hidden_states: torch.Tensor, freqs: torch.T return torch.cat([top_idx, tail_idx], dim=-1).to(torch.int64) @torch.no_grad() - def select_token_indices_thd(self, hidden_tok: torch.Tensor, freqs: torch.Tensor, - cu_seqlens: torch.Tensor) -> torch.Tensor: + def select_token_indices_thd(self, + hidden_tok: torch.Tensor, + freqs: torch.Tensor, + cu_seqlens: torch.Tensor, + force_materialize: bool = False) -> torch.Tensor: """QSA selection for packed (thd) inputs, indices in pack space. A standalone implementation: it does *not* call @@ -288,7 +291,13 @@ def select_token_indices_thd(self, hidden_tok: torch.Tensor, freqs: torch.Tensor ``[T, K]`` int64 pack-space indices (``K = block_topk*R + R``), ``-1`` unused; every index stays inside its query's document and causal prefix. ``None`` when selection is a no-op for every document, in which - case TE's packed causal kernel reproduces the selection exactly. + case TE's packed causal kernel reproduces the selection exactly -- + unless ``force_materialize`` is set, which builds the (select-all) + indices anyway. Callers pass that when the TE packed fallback is not + usable, i.e. thd under context parallelism: TE restricts thd+all_gather + to FusedAttention/FlashAttention-v3, and on sm100 neither is available + (FA3 is sm90-only), so the no-op must stay on the CP-aware sparse + kernel instead of degrading to a TE dense path that has no backend. """ T, _ = hidden_tok.shape R = self.compress_ratio @@ -296,8 +305,7 @@ def select_token_indices_thd(self, hidden_tok: torch.Tensor, freqs: torch.Tensor doc_lens = (cu_seqlens[1:] - cu_seqlens[:-1]).long() # [D] D = doc_lens.numel() full_blocks = doc_lens // R # complete R-blocks per document - # No-op when no document's causal prefix can exceed the budget. - if int(full_blocks.max().item()) <= self.block_topk: + if not force_materialize and int(full_blocks.max().item()) <= self.block_topk: return None # ---- per-token document id / in-doc position ----