chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

[NPU][1/N] NPU basic functions refactor and new modelslim quant type (#13359)

Browse Source
This commit is contained in:
Even Zhou
2025-12-04 16:15:31 +08:00
committed by GitHub
parent d6c490192d
commit 894c0dc57c
43 changed files with 2500 additions and 2058 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/speculative/draft_utils.py
+4 -2
View File
@@ -176,7 +176,7 @@ class DraftBackendFactory:
)
def _create_ascend_decode_backend(self):
from sglang.srt.layers.attention.ascend_backend import (
from sglang.srt.hardware_backend.npu.attention.ascend_backend import (
AscendAttnMultiStepDraftBackend,
)
@@ -231,7 +231,9 @@ class DraftBackendFactory:
return TRTLLMMLABackend(self.draft_model_runner, skip_prefill=False)
def _create_ascend_prefill_backend(self):
from sglang.srt.layers.attention.ascend_backend import AscendAttnBackend
from sglang.srt.hardware_backend.npu.attention.ascend_backend import (
AscendAttnBackend,
)
return AscendAttnBackend(self.draft_model_runner)
python/sglang/srt/speculative/eagle_draft_extend_npu_graph_runner.py
-68
View File
@@ -1,68 +0,0 @@
# Copyright 2024-2025 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 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_nsa
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
EAGLEDraftExtendCudaGraphRunner,
)
if TYPE_CHECKING:
from sglang.srt.speculative.eagle_worker import EAGLEWorker
class EAGLEDraftExtendNpuGraphRunner(EAGLEDraftExtendCudaGraphRunner):
def __init__(self, eagle_worker: EAGLEWorker):
super().__init__(eagle_worker)
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_update(self, seq_lens):
self.graphs[self.bs].update(
cpu_update_input=[{"actual_seq_lengths_kv": seq_lens}]
)
def _replay(self, forward_batch: ForwardBatch):
if not is_deepseek_nsa(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()
else:
self.graphs[self.bs].replay()
python/sglang/srt/speculative/eagle_draft_npu_graph_runner.py
-81
View File
@@ -1,81 +0,0 @@
# Copyright 2025 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 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
import torch
from sglang.srt.configs.model_config import is_deepseek_nsa
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
EAGLEDraftCudaGraphRunner,
)
if TYPE_CHECKING:
from sglang.srt.speculative.eagle_worker import EAGLEWorker
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: EAGLEWorker):
super().__init__(eagle_worker)
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_update(self, seq_lens):
self.graphs[self.bs].update(
cpu_update_input=[{"actual_seq_lengths_kv": seq_lens}]
)
def _replay(self, forward_batch: ForwardBatch):
if not is_deepseek_nsa(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()
else:
self.graphs[self.bs].replay()
python/sglang/srt/speculative/eagle_info.py
+8 -14
View File
@@ -38,9 +38,7 @@ from sglang.srt.speculative.spec_utils import (
get_src_tgt_cache_loc,
get_target_cache_loc,
)
from sglang.srt.utils import is_cuda, is_npu, next_power_of_2
_is_npu = is_npu()
from sglang.srt.utils import is_cuda, next_power_of_2
if is_cuda():
from sgl_kernel import (
@@ -77,22 +75,18 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
@classmethod
def create_idle_input(cls, topk: int, spec_steps: int, num_verify_tokens: int):
if not _is_npu:
device = "cuda"
else:
device = "npu"
return cls(
draft_token=torch.empty((0,), dtype=torch.long, device=device),
custom_mask=torch.full((0,), True, dtype=torch.bool, device=device),
positions=torch.empty((0,), dtype=torch.int64, device=device),
draft_token=torch.empty((0,), dtype=torch.long, device="cuda"),
custom_mask=torch.full((0,), True, dtype=torch.bool, device="cuda"),
positions=torch.empty((0,), dtype=torch.int64, device="cuda"),
retrive_index=torch.full(
(0, num_verify_tokens), -1, dtype=torch.long, device=device
(0, num_verify_tokens), -1, dtype=torch.long, device="cuda"
),
retrive_next_token=torch.full(
(0, num_verify_tokens), -1, dtype=torch.long, device=device
(0, num_verify_tokens), -1, dtype=torch.long, device="cuda"
),
retrive_next_sibling=torch.full(
(0, num_verify_tokens), -1, dtype=torch.long, device=device
(0, num_verify_tokens), -1, dtype=torch.long, device="cuda"
),
retrive_cum_len=None,
topk=topk,
@@ -282,7 +276,7 @@ class EagleVerifyInput(SpecInput, EagleVerifyInputV2Mixin):
"Falling back to greedy verification."
)
if is_all_greedy or not TREE_SPEC_KERNEL_AVAILABLE or _is_npu:
if is_all_greedy or not TREE_SPEC_KERNEL_AVAILABLE:
target_predict = torch.argmax(logits_output.next_token_logits, dim=-1)
target_predict = target_predict.reshape(bs, self.draft_token_num)
predict, accept_index, accept_length = verify_tree_greedy_func(
python/sglang/srt/speculative/eagle_worker.py
+3 -1
View File
@@ -5,6 +5,9 @@ from typing import List, Optional, Tuple
import torch
from sglang.srt.distributed import get_tp_group
from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner import (
EAGLEDraftNpuGraphRunner,
)
from sglang.srt.layers.dp_attention import get_attention_tp_group
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.layers.moe.utils import speculative_moe_backend_context
@@ -31,7 +34,6 @@ from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
EAGLEDraftExtendCudaGraphRunner,
)
from sglang.srt.speculative.eagle_draft_npu_graph_runner import EAGLEDraftNpuGraphRunner
from sglang.srt.speculative.eagle_info import (
EagleDraftInput,
EagleVerifyInput,
python/sglang/srt/speculative/eagle_worker_v2.py
+6 -4
View File
@@ -6,6 +6,12 @@ from typing import List, Optional, Tuple
import torch
from sglang.srt.environ import envs
from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_extend_npu_graph_runner import (
EAGLEDraftExtendNpuGraphRunner,
)
from sglang.srt.hardware_backend.npu.graph_runner.eagle_draft_npu_graph_runner import (
EAGLEDraftNpuGraphRunner,
)
from sglang.srt.layers.moe.utils import speculative_moe_backend_context
from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput
from sglang.srt.managers.schedule_batch import ModelWorkerBatch
@@ -21,10 +27,6 @@ from sglang.srt.speculative.eagle_draft_cuda_graph_runner import (
from sglang.srt.speculative.eagle_draft_extend_cuda_graph_runner import (
EAGLEDraftExtendCudaGraphRunner,
)
from sglang.srt.speculative.eagle_draft_extend_npu_graph_runner import (
EAGLEDraftExtendNpuGraphRunner,
)
from sglang.srt.speculative.eagle_draft_npu_graph_runner import EAGLEDraftNpuGraphRunner
from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput
from sglang.srt.speculative.eagle_info_v2 import (
assign_extend_cache_locs,