[NPU][Bugfix] Fix accuracy issue in no-graph with MTP (#27752)

This commit is contained in:
iridiumine
2026-06-12 16:05:06 +08:00
committed by GitHub
parent a52ccd2179
commit 60e4f14953
4 changed files with 24 additions and 7 deletions
@@ -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)
+7 -2
View File
@@ -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
@@ -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
@@ -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,