chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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

46
python/sglang/srt/model_executor/model_runner.py
View File

@@ -78,6 +78,7 @@ from sglang.srt.eplb.expert_location import (
set_global_expert_location_metadata,
)
from sglang.srt.eplb.expert_location_updater import ExpertLocationUpdater
from sglang.srt.hardware_backend.npu.graph_runner.npu_graph_runner import NPUGraphRunner
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.attention.attention_registry import (
ATTENTION_BACKENDS,
@@ -104,10 +105,7 @@ from sglang.srt.mem_cache.allocator import (
SWATokenToKVPoolAllocator,
TokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.allocator_ascend import AscendPagedTokenToKVPoolAllocator
from sglang.srt.mem_cache.memory_pool import (
AscendMLAPagedTokenToKVPool,
AscendTokenToKVPool,
DoubleSparseTokenToKVPool,
HybridLinearKVPool,
HybridReqToTokenPool,
@@ -132,7 +130,6 @@ from sglang.srt.model_executor.forward_batch_info import (
)
from sglang.srt.model_executor.hook_manager import register_forward_hooks
from sglang.srt.model_executor.input_buffers import GraphInputBuffers
from sglang.srt.model_executor.npu_graph_runner import NPUGraphRunner
from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
PiecewiseCudaGraphRunner,
)
@@ -187,6 +184,17 @@ from sglang.srt.weight_sync.tensor_bucket import (
FlattenedTensorMetadata,
)
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_npu = is_npu()
_is_cpu_amx_available = cpu_has_amx_support()
_is_xpu_xmx_available = xpu_has_xmx_support()
if _is_npu:
from sglang.srt.hardware_backend.npu.utils import init_npu_backend
init_npu_backend()
MLA_ATTENTION_BACKENDS = [
"aiter",
"flashinfer",
@@ -224,12 +232,6 @@ def add_chunked_prefix_cache_attention_backend(backend_name):
)
_is_cuda = is_cuda()
_is_hip = is_hip()
_is_npu = is_npu()
_is_cpu_amx_available = cpu_has_amx_support()
_is_xpu_xmx_available = xpu_has_xmx_support()
# Use a small KV cache pool size for tests in CI
SGLANG_CI_SMALL_KV_SIZE = os.getenv("SGLANG_CI_SMALL_KV_SIZE", None)
@@ -241,12 +243,6 @@ MAMBA_CACHE_SIZE_MAX_RUNNING_REQUESTS_RATIO = 3
logger = logging.getLogger(__name__)
if _is_npu:
import torch_npu
torch.npu.config.allow_internal_format = True
torch_npu.npu.set_compile_mode(jit_compile=False)
def resolve_language_model(model: nn.Module) -> nn.Module:
model_cls_name = model.__class__.__name__
@@ -1814,7 +1810,11 @@ class ModelRunner:
is_nsa_model = is_deepseek_nsa(self.model_config.hf_config)
if self.server_args.attention_backend == "ascend":
if self.use_mla_backend:
self.token_to_kv_pool = AscendMLAPagedTokenToKVPool(
from sglang.srt.hardware_backend.npu.memory_pool_npu import (
NPUMLATokenToKVPool,
)
self.token_to_kv_pool = NPUMLATokenToKVPool(
self.max_total_num_tokens,
page_size=self.page_size,
dtype=self.kv_cache_dtype,
@@ -1828,7 +1828,11 @@ class ModelRunner:
end_layer=self.end_layer,
)
else:
self.token_to_kv_pool = AscendTokenToKVPool(
from sglang.srt.hardware_backend.npu.memory_pool_npu import (
NPUMHATokenToKVPool,
)
self.token_to_kv_pool = NPUMHATokenToKVPool(
self.max_total_num_tokens,
page_size=self.page_size,
dtype=self.kv_cache_dtype,
@@ -1986,7 +1990,11 @@ class ModelRunner:
self.server_args.attention_backend == "ascend"
or self.hybrid_gdn_config is not None
):
self.token_to_kv_pool_allocator = AscendPagedTokenToKVPoolAllocator(
from sglang.srt.hardware_backend.npu.allocator_npu import (
NPUPagedTokenToKVPoolAllocator,
)
self.token_to_kv_pool_allocator = NPUPagedTokenToKVPoolAllocator(
self.max_total_num_tokens,
page_size=self.page_size,
dtype=self.kv_cache_dtype,

172
python/sglang/srt/model_executor/npu_graph_runner.py
View File

@@ -1,172 +0,0 @@
# Copyright 2023-2024 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 os
import threading
from pathlib import Path
from typing import TYPE_CHECKING, Dict, Optional, Union
import numpy as np
import torch
from sglang.srt.configs.model_config import AttentionArch, is_deepseek_nsa
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.model_executor.cuda_graph_runner import CudaGraphRunner
from sglang.srt.utils import is_npu
is_npu = is_npu()
if is_npu:
import torch_npu
from torch_npu.profiler import ProfilerActivity, profile
logger = logging.getLogger(__name__)
if TYPE_CHECKING:
from sglang.srt.model_executor.model_runner import ModelRunner
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
class NPUGraphRunner(CudaGraphRunner):
"""A NPUGraphRunner runs the forward pass of a model with npu graph and torch.compile."""
def __init__(self, model_runner: ModelRunner):
super().__init__(model_runner)
self.update_attr_name = None
self.update_attr_type = None
self.model_runner = model_runner
self._init_arch_map()
def _init_arch_map(self):
self.attr_name: Dict[str, str] = {
AttentionArch.MLA: "actual_seq_lengths_kv",
AttentionArch.MHA: "context_lens",
}
self.attr_type: Dict[str, Union[list, torch.Tensor]] = {
AttentionArch.MLA: [],
AttentionArch.MHA: torch.Tensor(),
}
def _create_device_graph(self):
return torch.npu.NPUGraph()
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 _get_update_attr_name(self, model_runner):
if self.bs < get_attention_tp_size():
return self.attr_name[AttentionArch.MLA]
return self.attr_name[model_runner.model_config.attention_arch]
def _get_update_attr_type(self, model_runner):
if self.bs < get_attention_tp_size():
return self.attr_type[AttentionArch.MLA]
return self.attr_type[model_runner.model_config.attention_arch]
def _update_inputs(self, seq_lens):
if isinstance(self.update_attr_type, torch.Tensor):
seq_lens = torch.from_numpy(np.array(seq_lens).astype(np.int32))
self.graphs[self.bs].update(
cpu_update_input=[{self.update_attr_name: seq_lens}]
)
def _cache_loc_dtype(self):
return torch.int32
def _init_profile_context_and_memory_record(self):
output_dir = os.path.join(
os.getenv("SGLANG_TORCH_PROFILER_DIR", "/tmp"), "graph_capture_profile"
)
if not Path(output_dir).exists():
Path(output_dir).mkdir(parents=True, exist_ok=True)
logger.info(
f"Profiling starts for graph capture for NPU. Traces will be saved to: {output_dir}"
)
experimental_config = torch_npu.profiler._ExperimentalConfig(
export_type=[torch_npu.profiler.ExportType.Text],
profiler_level=torch_npu.profiler.ProfilerLevel.Level1,
)
profile_context = profile(
activities=[ProfilerActivity.CPU, ProfilerActivity.NPU],
record_shapes=True,
profile_memory=True,
on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(
output_dir, async_mode=True
),
experimental_config=experimental_config,
)
return profile_context
def _post_process_after_profile(self, prof_context):
# for NPU, profile data will be saved to disk for further analysis.
pass
def replay(
self,
forward_batch: ForwardBatch,
skip_attn_backend_init: bool = False,
pp_proxy_tensors: Optional[PPProxyTensors] = None,
) -> Union[LogitsProcessorOutput, PPProxyTensors]:
if not skip_attn_backend_init:
self.replay_prepare(forward_batch, pp_proxy_tensors)
else:
# In speculative decoding, these two fields are still needed.
self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids)
self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions)
self.update_attr_name = self._get_update_attr_name(self.model_runner)
self.update_attr_type = self._get_update_attr_type(self.model_runner)
# Replay
if not is_deepseek_nsa(self.model_runner.model_config.hf_config):
if forward_batch.forward_mode.is_target_verify():
seq_lens_cpu = forward_batch.seq_lens.cpu() + self.num_tokens_per_bs
seq_lens = seq_lens_cpu.tolist() + [0] * (self.bs - self.raw_bs)
else:
seq_lens = forward_batch.seq_lens.cpu().tolist() + [0] * (
self.bs - self.raw_bs
)
thread = threading.Thread(target=self._update_inputs, args=(seq_lens,))
thread.start()
self.graphs[self.bs].replay()
thread.join()
else:
self.graphs[self.bs].replay()
output = self.output_buffers[self.bs]
if isinstance(output, LogitsProcessorOutput):
return LogitsProcessorOutput(
next_token_logits=output.next_token_logits[: self.raw_num_token],
hidden_states=(
output.hidden_states[: self.raw_num_token]
if output.hidden_states is not None
else None
),
)
else:
assert isinstance(output, PPProxyTensors)
return PPProxyTensors({k: v[: self.bs] for k, v in output.tensors.items()})