minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

bugfix for npu mtp graph runner (#27808)

Browse Source
This commit is contained in:
Hexq0210
2026-06-11 10:09:50 +08:00
committed by GitHub
parent fc1fee528c
commit 475e9d25bf
8 changed files with 100 additions and 163 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
+39 -20
View File
@@ -1770,26 +1770,45 @@ class AscendAttnBackend(AttentionBackend):
mask = self.mtp_mask
sparse_mode = 4 if is_swa_layer else 3
attn_output, _ = torch_npu.npu_fused_infer_attention_score_v2(
query,
k_cache,
v_cache,
block_table=block_table,
block_size=self.page_size,
num_query_heads=layer.tp_q_head_num,
num_key_value_heads=layer.tp_k_head_num,
input_layout="TND",
atten_mask=mask,
softmax_scale=layer.scaling,
actual_seq_qlen=actual_seq_lengths,
actual_seq_kvlen=actual_seq_lengths_kv,
sparse_mode=sparse_mode,
pre_tokens=(
layer.sliding_window_size if is_swa_layer else FULL_ATTENTION_WINDOW
),
next_tokens=0 if is_swa_layer else FULL_ATTENTION_WINDOW,
learnable_sink=sinks,
)
if self.is_hybrid_swa:
attn_output, _ = torch_npu.npu_fused_infer_attention_score_v2(
query,
k_cache,
v_cache,
block_table=block_table,
block_size=self.page_size,
num_query_heads=layer.tp_q_head_num,
num_key_value_heads=layer.tp_k_head_num,
input_layout="TND",
atten_mask=mask,
softmax_scale=layer.scaling,
actual_seq_qlen=actual_seq_lengths,
actual_seq_kvlen=actual_seq_lengths_kv,
sparse_mode=sparse_mode,
pre_tokens=(
layer.sliding_window_size
if is_swa_layer
else FULL_ATTENTION_WINDOW
),
next_tokens=0 if is_swa_layer else FULL_ATTENTION_WINDOW,
learnable_sink=sinks,
)
else:
attn_output, _ = torch.ops.npu.npu_fused_infer_attention_score(
query,
k_cache,
v_cache,
block_table=self.forward_metadata.block_tables,
block_size=self.page_size,
num_heads=layer.tp_q_head_num,
num_key_value_heads=layer.tp_k_head_num,
input_layout="TND",
atten_mask=mask,
scale=layer.scaling,
actual_seq_lengths=actual_seq_lengths,
actual_seq_lengths_kv=actual_seq_lengths_kv,
sparse_mode=sparse_mode,
)
attn_output = attn_output.view(-1, layer.tp_q_head_num * layer.v_head_dim)
if (
not self.graph_mode
python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py
+8 -30
View File
@@ -11,17 +11,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Run the model with npu graph and torch.compile."""
from __future__ import annotations
import threading
from typing import TYPE_CHECKING
import torch
from sglang.srt.configs.model_config import is_deepseek_dsa
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
EAGLEDraftExtendCudaGraphRunner,
)
@@ -34,38 +31,19 @@ class EAGLEDraftExtendNpuGraphRunner(EAGLEDraftExtendCudaGraphRunner):
def __init__(self, eagle_worker: EagleDraftWorker):
super().__init__(eagle_worker)
def _create_graph(self):
return torch.npu.NPUGraph()
def _cache_loc_dtype(self):
return torch.int32
def _capture_init(self, run_once_fn):
for _ in range(2):
torch.npu.synchronize()
self.model_runner.tp_group.barrier()
run_once_fn()
def _capture_graph(self, graph, pool, stream, run_once_fn):
with torch.npu.graph(
graph, pool=pool, stream=stream, auto_dispatch_capture=True
):
out = run_once_fn()
return out
def _replay_update(self, seq_lens):
self.graphs[self.bs].update(
cpu_update_input=[{"actual_seq_lengths_kv": seq_lens}]
)
def _replay(self, forward_batch: ForwardBatch):
def _replay_graph(self, shape_key, forward_batch):
if not is_deepseek_dsa(self.model_runner.model_config.hf_config):
seq_lens = forward_batch.seq_lens_cpu.tolist() + [0] * (
self.bs - self.raw_bs
)
thread = threading.Thread(target=self._replay_update, args=(seq_lens,))
thread.start()
self.graphs[self.bs].replay()
thread.join()
return self.backend.replay_with_input_update(
shape_key,
seq_lens=seq_lens,
attr_name="actual_seq_lengths_kv",
attr_type=[],
)
else:
self.graphs[self.bs].replay()
return self.backend.replay(shape_key, forward_batch)
python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py
+12 -58
View File
@@ -11,19 +11,14 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Run the model with npu graph and torch.compile"""
from __future__ import annotations
import logging
import threading
from typing import TYPE_CHECKING, Dict, Union
import numpy as np
import torch
from sglang.srt.configs.model_config import AttentionArch, is_deepseek_dsa
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
EAGLEDraftCudaGraphRunner,
)
@@ -31,25 +26,11 @@ from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
if TYPE_CHECKING:
from sglang.srt.speculative.eagle_worker_v2 import EagleDraftWorker
from sglang.srt.utils import is_npu
logger = logging.getLogger(__name__)
if is_npu():
torch.cuda.CUDAGraph = torch.npu.NPUGraph
torch.cuda.synchronize = torch.npu.synchronize
torch.cuda.graph = torch.npu.graph
torch.cuda.stream = torch.npu.stream
torch.cuda.Stream = torch.npu.Stream
torch.cuda.current_stream = torch.npu.current_stream
class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
def __init__(self, eagle_worker: EagleDraftWorker):
super().__init__(eagle_worker)
self.update_attr_name = None
self.update_attr_type = None
self._init_arch_map()
super().__init__(eagle_worker)
def _init_arch_map(self):
self.attr_name: Dict[str, str] = {
@@ -61,21 +42,8 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
AttentionArch.MHA: torch.Tensor(),
}
def _create_graph(self):
return torch.npu.NPUGraph()
def _capture_init(self, run_once_fn):
for _ in range(2):
torch.npu.synchronize()
self.model_runner.tp_group.barrier()
run_once_fn()
def _capture_graph(self, graph, pool, stream, run_once_fn):
with torch.npu.graph(
graph, pool=pool, stream=stream, auto_dispatch_capture=True
):
out = run_once_fn()
return out
def _cache_loc_dtype(self):
return torch.int32
def _get_update_attr_name(self):
return self.attr_name[AttentionArch.MLA]
@@ -83,33 +51,19 @@ class EAGLEDraftNpuGraphRunner(EAGLEDraftCudaGraphRunner):
def _get_update_attr_type(self):
return self.attr_type[AttentionArch.MLA]
def _replay_update(self, seq_lens_list):
if isinstance(self.update_attr_type, torch.Tensor):
seq_lens = torch.from_numpy(np.array(seq_lens_list).astype(np.int32))
self.graphs[self.bs].update(
cpu_update_input=[
{self.update_attr_name: seq_lens} for seq_lens in seq_lens_list
]
)
def _replay(self, forward_batch: ForwardBatch):
self.update_attr_name = self._get_update_attr_name()
self.update_attr_type = self._get_update_attr_type()
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 = []
for speculative_step_id in range(self.speculative_num_steps - 1):
seq_lens_cpu = forward_batch.seq_lens_cpu + speculative_step_id + 1
seq_lens_cpu = (
forward_batch.seq_lens_cpu[: self.raw_bs] + speculative_step_id + 1
)
seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs)
seq_lens_for_each_draft_step.append(seq_lens)
thread = threading.Thread(
target=self._replay_update, args=(seq_lens_for_each_draft_step,)
attr_name = self._get_update_attr_name()
cpu_update_input = [{attr_name: sl} for sl in seq_lens_for_each_draft_step]
return self.backend.replay_with_input_update(
shape_key, seq_lens=None, cpu_update_input=cpu_update_input
)
thread.start()
self.graphs[self.bs].replay()
thread.join()
else:
self.graphs[self.bs].replay()
def _cache_loc_dtype(self):
return torch.int32
return self.backend.replay(shape_key, forward_batch)
python/sglang/srt/hardware_backend/npu/graph_runner/multi_layer_eagle_draft_extend_npu_graph_runner.py
+6 -29
View File
@@ -11,19 +11,16 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Run the multi-layer eagle draft extend model with npu graph."""
from __future__ import annotations
import logging
import threading
import time
from typing import TYPE_CHECKING, List, Optional
import torch
from sglang.srt.model_executor.cuda_graph_config import cuda_graph_fully_disabled
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.speculative.multi_layer_eagle_draft_extend_cuda_graph_runner import (
MultiLayerEagleDraftExtendCudaGraphRunner,
MultiLayerEagleMultiStepDraftExtendCudaGraphRunner,
@@ -44,36 +41,16 @@ class MultiLayerEagleDraftExtendNpuGraphRunner(
def __init__(self, eagle_worker: MultiLayerEagleDraftWorker, step: int):
super().__init__(eagle_worker, step)
def _create_graph(self):
return torch.npu.NPUGraph()
def _capture_init(self, run_once_fn):
for _ in range(2):
torch.npu.synchronize()
self.model_runner.tp_group.barrier()
run_once_fn()
def _capture_graph(self, graph, pool, stream, run_once_fn):
with torch.npu.graph(
graph,
pool=pool,
stream=stream,
auto_dispatch_capture=True,
):
out = run_once_fn()
return out
def _replay(self, forward_batch: ForwardBatch):
def _replay_graph(self, shape_key, forward_batch):
seq_lens = self.buffers.seq_lens_cpu[: self.raw_bs].tolist() + [0] * (
self.bs - self.raw_bs
)
thread = threading.Thread(
target=self.graphs[self.bs].update,
kwargs={"cpu_update_input": [{"actual_seq_kvlen": seq_lens}]},
return self.backend.replay_with_input_update(
shape_key,
seq_lens=seq_lens,
attr_name="actual_seq_kvlen",
attr_type=[],
)
thread.start()
self.graphs[self.bs].replay()
thread.join()
class MultiLayerEagleMultiStepDraftExtendNpuGraphRunner(
python/sglang/srt/hardware_backend/npu/graph_runner/npu_cudagraph_backend.py
+17 -7
View File
@@ -140,19 +140,29 @@ class NPUCudaGraphBackend(BaseCudaGraphBackend):
def replay_with_input_update(
self,
shape_key: Any,
seq_lens: list,
attr_name: str,
attr_type: Any,
seq_lens: Any,
attr_name: str = None,
attr_type: Any = None,
cpu_update_input: list = None,
) -> Any:
"""Rebind seq_lens on the recorded NPU graph in a background
thread, then replay. Used when the model is not deepseek-nsa."""
if isinstance(attr_type, torch.Tensor):
seq_lens = torch.from_numpy(np.array(seq_lens).astype(np.int32))
thread, then replay. Used when the model is not deepseek-nsa.
Two calling conventions:
1. (legacy) seq_lens + attr_name + attr_type:
Constructs cpu_update_input=[{attr_name: seq_lens}] internally.
2. cpu_update_input: A list of {attr_name: seq_lens} dicts,
one per speculative step. Used by EAGLE draft runners.
"""
if cpu_update_input is None:
if isinstance(attr_type, torch.Tensor):
seq_lens = torch.from_numpy(np.array(seq_lens).astype(np.int32))
cpu_update_input = [{attr_name: seq_lens}]
graph = self._graphs[shape_key]
def _update():
graph.update(cpu_update_input=[{attr_name: seq_lens}])
graph.update(cpu_update_input=cpu_update_input)
thread = threading.Thread(target=_update)
thread.start()
python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py
+6 -7
View File
@@ -26,7 +26,7 @@ from sglang.srt.model_executor.runner import (
get_batch_sizes_to_capture,
model_capture_mode,
)
from sglang.srt.model_executor.runner_backend import FullCudaGraphBackend
from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
from sglang.srt.model_executor.runner_backend_utils import (
CUDA_GRAPH_CAPTURE_FAILED_MSG,
)
@@ -226,11 +226,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
)
self.buffers.share_buffers()
# Backend (Full CUDA graph capture)
self.backend = FullCudaGraphBackend(
self,
enable_memory_saver=model_runner.server_args.enable_memory_saver,
)
self.backend = resolve_decode_backend(self)
# Capture
try:
@@ -241,6 +237,9 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
)
def _replay_graph(self, shape_key, forward_batch):
return self.backend.replay(shape_key, forward_batch)
# -----------------------------------------------------------------
# Helpers
# -----------------------------------------------------------------
@@ -533,7 +532,7 @@ class EAGLEDraftCudaGraphRunner(DecodeCudaGraphRunner):
else contextlib.nullcontext()
)
with timer_ctx:
out = self.backend.replay(shape_key, forward_batch)
out = self._replay_graph(shape_key, forward_batch)
if bs != raw_bs:
out = self._postprocess_output_to_raw_bs(out, raw_bs)
python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py
+6 -6
View File
@@ -26,7 +26,7 @@ from sglang.srt.model_executor.runner import (
get_batch_sizes_to_capture,
model_capture_mode,
)
from sglang.srt.model_executor.runner_backend import FullCudaGraphBackend
from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
from sglang.srt.model_executor.runner_backend_utils import (
CUDA_GRAPH_CAPTURE_FAILED_MSG,
)
@@ -242,10 +242,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
)
self.buffers.share_buffers()
self.backend = FullCudaGraphBackend(
self,
enable_memory_saver=model_runner.server_args.enable_memory_saver,
)
self.backend = resolve_decode_backend(self)
try:
with model_capture_mode():
@@ -255,6 +252,9 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
)
def _replay_graph(self, shape_key, forward_batch):
return self.backend.replay(shape_key, forward_batch)
def _cache_loc_dtype(self):
return torch.int64
@@ -556,7 +556,7 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
else contextlib.nullcontext()
)
with timer_ctx:
out = self.backend.replay(shape_key, forward_batch)
out = self._replay_graph(shape_key, forward_batch)
if self.forward_mode == ForwardMode.DRAFT_EXTEND_V2:
unpadding_bs = num_tokens
python/sglang/srt/speculative/multi_layer_eagle_draft_extend_cuda_graph_runner.py
+6 -6
View File
@@ -50,7 +50,7 @@ from sglang.srt.model_executor.runner import (
get_batch_sizes_to_capture,
model_capture_mode,
)
from sglang.srt.model_executor.runner_backend import FullCudaGraphBackend
from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend
from sglang.srt.model_executor.runner_backend_utils import (
CUDA_GRAPH_CAPTURE_FAILED_MSG,
)
@@ -272,10 +272,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
)
self.backend = FullCudaGraphBackend(
self,
enable_memory_saver=self.model_runner.server_args.enable_memory_saver,
)
self.backend = resolve_decode_backend(self)
try:
with model_capture_mode():
@@ -285,6 +282,9 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
f"Capture cuda graph failed: {e}\n{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
)
def _replay_graph(self, shape_key, forward_batch):
return self.backend.replay(shape_key, forward_batch)
def _make_graph_key(self, bs, stream_idx=None, variant_label=None):
return bs
@@ -588,7 +588,7 @@ class MultiLayerEagleDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
self.raw_bs = raw_bs
self.bs = bs
shape_key = self._make_graph_key(bs)
out = self.backend.replay(shape_key, forward_batch)
out = self._replay_graph(shape_key, forward_batch)
if self.forward_mode == ForwardMode.DRAFT_EXTEND_V2:
unpadding_bs = num_tokens