diff --git a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py index 46daed9cc..2869c4bd1 100644 --- a/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py +++ b/python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py @@ -241,8 +241,8 @@ class NPUGraphRunner(DecodeCudaGraphRunner): output = self.backend.replay_with_input_update( graph_key, seq_lens=seq_lens, - attr_name=self.attr_name[AttentionArch.MLA], - attr_type=self.attr_type[AttentionArch.MLA], + attr_name=self._get_update_attr_name(), + attr_type=self._get_update_attr_type(), ) else: output = self.backend.replay(graph_key, forward_batch) diff --git a/python/sglang/srt/mem_cache/allocator/swa.py b/python/sglang/srt/mem_cache/allocator/swa.py index 4b44b39f1..66e49aadb 100644 --- a/python/sglang/srt/mem_cache/allocator/swa.py +++ b/python/sglang/srt/mem_cache/allocator/swa.py @@ -10,6 +10,8 @@ from sglang.srt.utils.common import get_num_new_pages _is_npu = is_npu() if _is_npu: + import torch_npu + from sglang.srt.hardware_backend.npu.allocator_npu import ( NPUPagedTokenToKVPoolAllocator, ) @@ -297,8 +299,11 @@ class SWATokenToKVPoolAllocator(BaseTokenToKVPoolAllocator): return None if _is_npu: - self.full_to_swa_index_mapping[alloc_full_indices.to(torch.int64)] = ( - alloc_swa_indices.to(torch.int64) + indices_2d = alloc_full_indices.to(torch.int64).unsqueeze(-1) + torch_npu.npu_scatter_nd_update_( + self.full_to_swa_index_mapping, + indices_2d, + alloc_swa_indices.to(torch.int64), ) else: self.full_to_swa_index_mapping[alloc_full_indices] = alloc_swa_indices diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index e54abb5a4..7b8a402d6 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/python/sglang/srt/speculative/eagle_info_v2.py @@ -370,7 +370,11 @@ class EagleDraftExtendInputV2Mixin: # (the `_batch_size == batch_size` assertion, see #27091); the # marked pre-pad metadata is used as-is, matching the proven # skip_attn_backend_init=True behavior. - forward_batch.mark_forward_metadata_ready() + # On NPU with --disable-cuda-graph, block_table shape won't match + # after prepare_mlp_sync_batch padding; defer re-init to + # forward_extend (post-pad) instead. + if not _is_npu or can_cuda_graph: + forward_batch.mark_forward_metadata_ready() return forward_batch diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 1f0e6d925..1f77b4269 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -562,7 +562,12 @@ class MultiLayerEagleDraftWorker(BaseDraftWorker): # pre-plan (see warning above). Mark the batch so the forward path # keeps skipping metadata init — preserves the pre-existing # behavior; the latent issue is tracked by the warning. - forward_batch.mark_forward_metadata_ready() + # On NPU with --disable-cuda-graph, leave each draft runner to init + # its own metadata in forward_extend (post-pad), otherwise + # per-runner attn_backend.forward_metadata is never initialized for + # draft_runner_list[1+]. + if not _is_npu or can_cuda_graph: + forward_batch.mark_forward_metadata_ready() for step in range(self.speculative_num_steps): # log_info_on_rank0(logger, f"step: {step}, forward_batch.input_ids: {forward_batch.input_ids}") @@ -821,7 +826,10 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): # prepare_for_v2_verify only plans when cuda-graph replay_prepare ran. # eagle_worker_v2 re-inits the non-graph path instead (post-pad); this # worker has not adopted that fix, so preserve its behavior verbatim. - verify_forward_batch.mark_forward_metadata_ready() + # On NPU with --disable-cuda-graph, non-graph verify needs metadata init + # in forward_extend (post-pad); only mark ready for the cuda-graph path. + if not _is_npu or can_run_cuda_graph: + verify_forward_batch.mark_forward_metadata_ready() # Run target verify batch in the main compute stream forward_batch_output = self.target_worker.forward_batch_generation( batch=None,