[NPU] Fix MTP IndexShare warm-up for attention DP and prefill CP (#32210)

python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py

@@ -849,7 +849,10 @@ class AscendAttnBackend(AttentionBackend):
         q_nope_next = q_nope_next.contiguous()
         q_rope_prev = q_rope_prev.contiguous()
         q_rope_next = q_rope_next.contiguous()
-        topk_indices_prev, topk_indices_next = topk_indices
+        topk_indices = _expand_dsa_sparse_indices(topk_indices)
+        topk_indices_prev, topk_indices_next = torch.split(
+            topk_indices, split_len, dim=0
+        )
 
         actual_seq_qlen_prev, actual_seq_qlen_next = actual_seq_qlen
         actual_seq_lengths_kv_prev, actual_seq_lengths_kv_next = actual_seq_lengths_kv

python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py

@@ -24,6 +24,7 @@ from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
 )
 
 if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
     from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
 
 
@@ -51,6 +52,36 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
     def _get_update_attr_type(self):
         return self.attr_type[AttentionArch.MLA]
 
+    def can_run_graph(self, forward_batch: ForwardBatch):
+        can_run_graph = super().can_run_graph(forward_batch)
+        if (
+            not self.eagle_worker.seed_dsa_topk_from_draft_extend
+            or self.attn_dp_size <= 1
+        ):
+            return can_run_graph
+
+        # PR #30839 falls back to eager when an IndexShare seed is unavailable.
+        # Under attention DP, seed availability is request-local: a real rank can
+        # miss the PD seed while idle ranks have no request at all. All ranks must
+        # nevertheless choose the same graph/eager path because the draft forward
+        # contains TP/EP collectives. Reduce the final decision across the model TP
+        # group, which contains all attention-DP ranks for this pipeline stage.
+        spec_info = forward_batch.spec_info
+        seed_ready = forward_batch.forward_mode.is_idle() or (
+            spec_info is not None and spec_info.dsa_topk_indices is not None
+        )
+        decision = torch.tensor(
+            int(can_run_graph and seed_ready),
+            dtype=torch.int32,
+            device=self.device,
+        )
+        torch.distributed.all_reduce(
+            decision,
+            op=torch.distributed.ReduceOp.MIN,
+            group=self.model_runner.tp_group.device_group,
+        )
+        return bool(decision.item())
+
     def _replay_graph(self, shape_key, forward_batch):
         if not is_deepseek_dsa(self.model_runner.model_config.hf_config):
             seq_lens_for_each_draft_step = []

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -2398,7 +2398,7 @@ class Indexer(MultiPlatformOp):
             sparse_count=self.index_topk,
             sparse_mode=3,
         )
-        return topk_indices_prev[0], topk_indices_next[0]
+        return torch.cat([topk_indices_prev[0], topk_indices_next[0]], dim=0).squeeze(1)
 
 
 @register_custom_op(mutates_args=["topk_result"])