diff --git a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py b/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py index e95072db6..8b403c38b 100644 --- a/python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py +++ b/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 diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py index 450a81f3a..78c82d864 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py +++ b/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 = [] diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer.py index 036db24d6..e5268aeab 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer.py +++ b/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"])