fix qwen4-exp thd cp

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,

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 --

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,16 +291,21 @@ 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
         device = hidden_tok.device
         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 ----