[NPU][Bugfix] Fix accuracy issue in no-graph with MTP (#27752)
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user