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
python/sglang/srt/environ.py
+3
View File
@@ -232,6 +232,9 @@ class Envs:
SGLANG_ROCM_FUSED_DECODE_MLA = EnvBool(False)
SGLANG_ROCM_DISABLE_LINEARQUANT = EnvBool(False)
# NPU
SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT = EnvBool(False)
# Quantization
SGLANG_INT4_WEIGHT = EnvBool(False)
SGLANG_CPU_QUANTIZATION = EnvBool(False)
python/sglang/srt/eplb/expert_distribution.py
+2 -8
View File
@@ -31,9 +31,7 @@ from sglang.srt.environ import envs
from sglang.srt.metrics.collector import ExpertDispatchCollector
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import Withable, get_int_env_var, is_npu
_is_npu = is_npu()
from sglang.srt.utils import Withable, get_int_env_var
if TYPE_CHECKING:
from sglang.srt.eplb.expert_location import ExpertLocationMetadata
@@ -465,10 +463,6 @@ def _list_sum(a: List, b: List) -> List:
class _LayerBasedGpuSinglePassGatherer(_SinglePassGatherer):
def __init__(self, *args, enable_global_physical_experts: bool, **kwargs):
super().__init__(*args, **kwargs)
if not _is_npu:
device = "cuda"
else:
device = "npu"
self._enable_global_physical_experts = enable_global_physical_experts
self._data = torch.zeros(
(
@@ -480,7 +474,7 @@ class _LayerBasedGpuSinglePassGatherer(_SinglePassGatherer):
),
),
dtype=torch.int,
device=device,
device="cuda",
)
def reset(self):
python/sglang/srt/mem_cache/allocator_ascend.py → python/sglang/srt/hardware_backend/npu/allocator_npu.py
+7 -9
View File
@@ -1,17 +1,15 @@
from __future__ import annotations
from typing import TYPE_CHECKING
import torch
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import KVCache
from sglang.srt.mem_cache.allocator import PagedTokenToKVPoolAllocator
from sglang.srt.utils import get_num_new_pages, next_power_of_2
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import KVCache
def alloc_extend_kernel_ascend(
def _alloc_extend_naive(
prefix_lens,
seq_lens,
last_loc,
@@ -65,14 +63,14 @@ def alloc_extend_kernel_ascend(
).view(-1)
class AscendPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
class NPUPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
def __init__(
self,
size: int,
page_size: int,
dtype: torch.dtype,
device: str,
kvcache: KVCache,
kvcache: "KVCache",
need_sort: bool,
):
super().__init__(size, page_size, dtype, device, kvcache, need_sort)
@@ -130,7 +128,7 @@ class AscendPagedTokenToKVPoolAllocator(PagedTokenToKVPoolAllocator):
dtype=torch.int32,
device=self.device,
)
alloc_extend_kernel_ascend(
_alloc_extend_naive(
prefix_lens,
seq_lens,
last_loc,
python/sglang/srt/layers/attention/ascend_backend.py → python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
+3 -1
View File
@@ -7,8 +7,10 @@ import torch
import torch_npu
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
is_mla_preprocess_enabled,
)
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
from sglang.srt.layers.attention.npu_ops.mla_preprocess import is_mla_preprocess_enabled
from sglang.srt.layers.attention.torch_native_backend import TorchNativeAttnBackend
from sglang.srt.layers.dp_attention import get_attention_tp_size
from sglang.srt.layers.radix_attention import AttentionType
python/sglang/srt/layers/attention/npu_ops/mla_preprocess.py → python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py
+19 -24
View File
@@ -1,23 +1,19 @@
from functools import lru_cache
from typing import TYPE_CHECKING, Optional
import torch
import torch.nn.functional as F
from sglang.srt.utils import get_bool_env_var, is_npu
from sglang.srt.hardware_backend.npu.utils import npu_format_cast
from sglang.srt.utils import get_bool_env_var
_is_npu = is_npu()
_ENABLE_MLA_PREPROCESS_FLAG = get_bool_env_var("SGLANG_NPU_USE_MLAPO")
_NPU_FORMAT_NZ = 29
if TYPE_CHECKING:
from sglang.srt.layers.quantization.base_config import QuantizationConfig
@lru_cache(maxsize=1)
def is_mla_preprocess_enabled() -> bool:
return _is_npu and _ENABLE_MLA_PREPROCESS_FLAG
if is_mla_preprocess_enabled():
import sgl_kernel_npu # noqa: F401
import torch_npu
torch.npu.config.allow_internal_format = True
torch.npu.set_compile_mode(jit_compile=False)
return get_bool_env_var("SGLANG_NPU_USE_MLAPO")
def round_up(val: int, align: int) -> int:
@@ -66,6 +62,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
num_local_heads,
qk_nope_head_dim,
qk_rope_head_dim,
quant_config: Optional["QuantizationConfig"] = None,
):
super().__init__()
self.qkv_a_proj = fused_qkv_a_proj_with_mqa
@@ -75,6 +72,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
self.w_kc = w_kc.contiguous()
self.rotary_emb = rotary_emb
self.layer_id = layer_id
self.quant_config = quant_config
self.has_preprocess_weights = False
self.dtype = None
@@ -124,9 +122,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
.unsqueeze(0)
.contiguous()
)
self.qkv_a_proj_weight_nz = torch_npu.npu_format_cast(
fused_qkv_a_proj_with_mqa_weight_nz, _NPU_FORMAT_NZ
)
self.qkv_a_proj_weight_nz = npu_format_cast(fused_qkv_a_proj_with_mqa_weight_nz)
# matmul_0 deq_scale [2112]
fused_qkv_a_proj_with_mqa_deq_scale_q = self.qkv_a_proj.deq_scale.data[
@@ -198,9 +194,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
q_b_proj_weight_nz = (
transdata(q_b_proj_weight, block_size=(16, 32)).unsqueeze(0).contiguous()
)
self.q_b_proj_weight_nz = torch_npu.npu_format_cast(
q_b_proj_weight_nz, _NPU_FORMAT_NZ
)
self.q_b_proj_weight_nz = npu_format_cast(q_b_proj_weight_nz)
# matmul_1 deq_scale [num_head * 192]
q_b_proj_deq_scale = self.q_b_proj.deq_scale.data.clone()
@@ -280,7 +274,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
q_pe = q_pe.view(-1, self.num_local_heads, 1, self.qk_rope_head_dim)
cos = cos.view(-1, 1, 1, self.qk_rope_head_dim)
sin = sin.view(-1, 1, 1, self.qk_rope_head_dim)
q_pe = torch_npu.npu_interleave_rope(q_pe, cos, sin) # (B,N,S,D)
q_pe = torch.ops.npu.npu_interleave_rope(q_pe, cos, sin) # (B,N,S,D)
q_pe = q_pe.view(cos.shape[0], self.num_local_heads, self.qk_rope_head_dim)
latent_cache = latent_cache.view(
@@ -300,7 +294,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
1,
forward_batch.attn_backend.qk_rope_head_dim,
)
k_rope, k_nope, _, _ = torch_npu.npu_kv_rmsnorm_rope_cache(
k_rope, k_nope, _, _ = torch.ops.npu.npu_kv_rmsnorm_rope_cache(
latent_cache,
self.kv_a_layernorm.weight,
cos,
@@ -378,10 +372,11 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
)
def forward(self, positions, hidden_states, forward_batch, zero_allocator):
assert self.quant_config and self.quant_config.get_name() == "modelslim"
# route by `qkv_a_proj` quant type as MTP layers can be unquantized
_is_w8a8 = (
hasattr(self.qkv_a_proj.quant_method, "quantization_config")
and self.qkv_a_proj.quant_method.quantization_config.get_name()
== "w8a8_int8"
hasattr(self.qkv_a_proj.quant_method, "quant_config")
and self.qkv_a_proj.quant_method.quant_config.get_name() == "modelslim"
)
if _is_w8a8:
return self.forward_mlapo(
python/sglang/srt/hardware_backend/npu/cmo.py
+54
View File
@@ -0,0 +1,54 @@
import torch
cmo_stream = None
def get_cmo_stream():
"""
Cache Management Operation(CMO).
Launch a new stream to prefetch the weight of matmul when running other
AIV or communication kernels, aiming to overlap the memory access time.
"""
global cmo_stream
return cmo_stream
def set_cmo_stream(stream):
global cmo_stream
cmo_stream = stream
def prepare_weight_cache(handle, cache, PREFETCH_MAX_SIZE=1000000000):
"""
PREFETCH_MAX_SIZE: maximum size (bytes) for each prefetch operation.
This affects the time spent in prefetch:
time ≈ PREFETCH_MAX_SIZE / system_bandwidth
"""
import torch_npu
stream = get_cmo_stream()
if stream is None:
stream = torch.npu.Stream()
set_cmo_stream(stream)
stream.wait_stream(torch.npu.current_stream())
with torch.npu.stream(stream):
if isinstance(cache, list):
for weight in cache:
torch_npu.npu_prefetch(
weight,
handle,
PREFETCH_MAX_SIZE,
)
else:
torch_npu.npu_prefetch(
cache,
handle,
PREFETCH_MAX_SIZE,
)
def wait_cmo_stream():
stream = get_cmo_stream()
if stream is not None:
cur_stream = torch.npu.current_stream()
cur_stream.wait_stream(stream)
python/sglang/srt/speculative/eagle_draft_extend_npu_graph_runner.py → python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_extend_npu_graph_runner.py
View File
python/sglang/srt/speculative/eagle_draft_npu_graph_runner.py → python/sglang/srt/hardware_backend/npu/graph_runner/eagle_draft_npu_graph_runner.py
View File
python/sglang/srt/model_executor/npu_graph_runner.py → python/sglang/srt/hardware_backend/npu/graph_runner/npu_graph_runner.py
View File
python/sglang/srt/hardware_backend/npu/memory_pool_npu.py
+300
View File
@@ -0,0 +1,300 @@
from typing import TYPE_CHECKING, Optional
import torch
import torch_npu
from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE
from sglang.srt.mem_cache.memory_pool import (
MHATokenToKVPool,
MLATokenToKVPool,
get_tensor_size_bytes,
)
if TYPE_CHECKING:
from sglang.srt.layers.radix_attention import RadixAttention
class NPUMHATokenToKVPool(MHATokenToKVPool):
def _create_buffers(self):
with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
# [size, head_num, head_dim] for each layer
# The padded slot 0 is used for writing dummy outputs from padded tokens.
# Continuous memory improves the efficiency of Ascend`s transmission backend,
# while other backends remain unchanged.
self.kv_buffer = torch.zeros(
(
2,
self.layer_num,
self.size // self.page_size + 1,
self.page_size,
self.head_num,
self.head_dim,
),
dtype=self.store_dtype,
device=self.device,
)
self.k_buffer = self.kv_buffer[0]
self.v_buffer = self.kv_buffer[1]
# for disagg
def get_contiguous_buf_infos(self):
# layer_num x [seq_len, head_num, head_dim]
# layer_num x [page_num, page_size, head_num, head_dim]
kv_data_ptrs = [
self.get_key_buffer(i).data_ptr()
for i in range(self.start_layer, self.start_layer + self.layer_num)
] + [
self.get_value_buffer(i).data_ptr()
for i in range(self.start_layer, self.start_layer + self.layer_num)
]
kv_data_lens = [
self.get_key_buffer(i).nbytes
for i in range(self.start_layer, self.start_layer + self.layer_num)
] + [
self.get_value_buffer(i).nbytes
for i in range(self.start_layer, self.start_layer + self.layer_num)
]
kv_item_lens = [
self.get_key_buffer(i)[0].nbytes
for i in range(self.start_layer, self.start_layer + self.layer_num)
] + [
self.get_value_buffer(i)[0].nbytes
for i in range(self.start_layer, self.start_layer + self.layer_num)
]
return kv_data_ptrs, kv_data_lens, kv_item_lens
def set_kv_buffer(
self,
layer: "RadixAttention",
loc: torch.Tensor,
cache_k: torch.Tensor,
cache_v: torch.Tensor,
k_scale: Optional[float] = None,
v_scale: Optional[float] = None,
layer_id_override: Optional[int] = None,
):
if layer_id_override is not None:
layer_id = layer_id_override
else:
layer_id = layer.layer_id
if cache_k.dtype != self.dtype:
if k_scale is not None:
cache_k.div_(k_scale)
if v_scale is not None:
cache_v.div_(v_scale)
cache_k = cache_k.to(self.dtype)
cache_v = cache_v.to(self.dtype)
if self.store_dtype != self.dtype:
cache_k = cache_k.view(self.store_dtype)
cache_v = cache_v.view(self.store_dtype)
torch_npu._npu_reshape_and_cache(
key=cache_k,
value=cache_v,
key_cache=self.k_buffer[layer_id - self.start_layer].view(
-1, self.page_size, self.head_num, self.head_dim
),
value_cache=self.v_buffer[layer_id - self.start_layer].view(
-1, self.page_size, self.head_num, self.head_dim
),
slot_indices=loc,
)
class NPUMLATokenToKVPool(MLATokenToKVPool):
def __init__(
self,
size: int,
page_size: int,
dtype: torch.dtype,
kv_lora_rank: int,
qk_rope_head_dim: int,
index_head_dim: Optional[int],
layer_num: int,
device: str,
enable_memory_saver: bool,
start_layer: Optional[int] = None,
end_layer: Optional[int] = None,
):
super(MLATokenToKVPool, self).__init__(
size,
page_size,
dtype,
layer_num,
device,
enable_memory_saver,
start_layer,
end_layer,
)
self.kv_lora_rank = kv_lora_rank
self.qk_rope_head_dim = qk_rope_head_dim
self.index_head_dim = index_head_dim
self.custom_mem_pool = None
with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
# The padded slot 0 is used for writing dummy outputs from padded tokens.
self.k_buffer = torch.zeros(
(
layer_num,
self.size // self.page_size + 1,
self.page_size,
1,
self.kv_lora_rank,
),
dtype=self.store_dtype,
device=self.device,
)
self.v_buffer = torch.zeros(
(
layer_num,
self.size // self.page_size + 1,
self.page_size,
1,
self.qk_rope_head_dim,
),
dtype=self.store_dtype,
device=self.device,
)
if self.index_head_dim is not None:
self.index_k_buffer = torch.zeros(
(
layer_num,
self.size // self.page_size + 1,
self.page_size,
1,
self.index_head_dim,
),
dtype=self.store_dtype,
device=self.device,
)
self._finalize_allocation_log(size)
def get_kv_size_bytes(self):
assert hasattr(self, "k_buffer")
assert hasattr(self, "v_buffer")
kv_size_bytes = 0
for k_cache in self.k_buffer:
kv_size_bytes += get_tensor_size_bytes(k_cache)
for v_cache in self.v_buffer:
kv_size_bytes += get_tensor_size_bytes(v_cache)
if self.index_head_dim is not None:
assert hasattr(self, "index_k_buffer")
for index_k_cache in self.index_k_buffer:
kv_size_bytes += get_tensor_size_bytes(index_k_cache)
return kv_size_bytes
def get_kv_buffer(self, layer_id: int):
if self.layer_transfer_counter is not None:
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
return (
self.k_buffer[layer_id - self.start_layer],
self.v_buffer[layer_id - self.start_layer],
)
def get_key_buffer(self, layer_id: int):
if self.layer_transfer_counter is not None:
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
if self.store_dtype != self.dtype:
return self.k_buffer[layer_id - self.start_layer].view(self.dtype)
return self.k_buffer[layer_id - self.start_layer]
def get_value_buffer(self, layer_id: int):
if self.layer_transfer_counter is not None:
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
if self.store_dtype != self.dtype:
return self.v_buffer[layer_id - self.start_layer].view(self.dtype)
return self.v_buffer[layer_id - self.start_layer]
def get_index_k_buffer(self, layer_id: int):
if self.layer_transfer_counter is not None:
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
if self.store_dtype != self.dtype:
return self.index_k_buffer[layer_id - self.start_layer].view(self.dtype)
return self.index_k_buffer[layer_id - self.start_layer]
# for disagg
def get_contiguous_buf_infos(self):
# MLA has only one kv_buffer, so only the information of this buffer needs to be returned.
kv_data_ptrs = [self.k_buffer[i].data_ptr() for i in range(self.layer_num)] + [
self.v_buffer[i].data_ptr() for i in range(self.layer_num)
]
kv_data_lens = [self.k_buffer[i].nbytes for i in range(self.layer_num)] + [
self.v_buffer[i].nbytes for i in range(self.layer_num)
]
kv_item_lens = [self.k_buffer[i][0].nbytes for i in range(self.layer_num)] + [
self.v_buffer[i][0].nbytes for i in range(self.layer_num)
]
if self.index_head_dim is not None:
kv_data_ptrs += [
self.index_k_buffer[i].data_ptr() for i in range(self.layer_num)
]
kv_data_lens += [
self.index_k_buffer[i].nbytes for i in range(self.layer_num)
]
kv_item_lens += [
self.index_k_buffer[i][0].nbytes for i in range(self.layer_num)
]
return kv_data_ptrs, kv_data_lens, kv_item_lens
def set_kv_buffer(
self,
layer: "RadixAttention",
loc: torch.Tensor,
cache_k: torch.Tensor,
cache_v: torch.Tensor,
):
layer_id = layer.layer_id
if cache_k.dtype != self.dtype:
cache_k = cache_k.to(self.dtype)
cache_v = cache_v.to(self.dtype)
if self.store_dtype != self.dtype:
cache_k = cache_k.view(self.store_dtype)
cache_v = cache_v.view(self.store_dtype)
if cache_v is None:
cache_k, cache_v = cache_k.split(
[self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
)
torch_npu.npu_scatter_nd_update_(
self.k_buffer[layer_id - self.start_layer].view(-1, 1, self.kv_lora_rank),
loc.view(-1, 1),
cache_k.view(-1, 1, self.kv_lora_rank),
)
torch_npu.npu_scatter_nd_update_(
self.v_buffer[layer_id - self.start_layer].view(
-1, 1, self.qk_rope_head_dim
),
loc.view(-1, 1),
cache_v.view(-1, 1, self.qk_rope_head_dim),
)
def set_index_k_buffer(
self,
layer_id: int,
loc: torch.Tensor,
index_k: torch.Tensor,
):
if index_k.dtype != self.dtype:
index_k = index_k.to(self.dtype)
if self.store_dtype != self.dtype:
index_k = index_k.view(self.store_dtype)
torch_npu.npu_scatter_nd_update_(
self.index_k_buffer[layer_id - self.start_layer].view(
-1, 1, self.index_head_dim
),
loc.view(-1, 1),
index_k.view(-1, 1, self.index_head_dim),
)
python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py
+380
View File
@@ -0,0 +1,380 @@
from typing import TYPE_CHECKING
import torch
from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
NPUFusedMLAPreprocess,
is_mla_preprocess_enabled,
)
from sglang.srt.layers.attention.nsa.utils import (
cp_split_and_rebuild_position,
enable_prefill_cp,
)
from sglang.srt.layers.communicator import get_attn_tp_context
if TYPE_CHECKING:
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.models.deepseek_v2 import DeepseekV2AttentionMLA
from sglang.srt.utils import BumpAllocator
# region MHA
def forward_mha_prepare_npu(
m: "DeepseekV2AttentionMLA",
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: "ForwardBatch",
zero_allocator: "BumpAllocator",
):
if m.q_lora_rank is not None:
q, latent_cache = (
get_attn_tp_context()
.fetch_qkv_latent()
.split(
[m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim],
dim=-1,
)
)
# NSA Indexer: cache quantized keys, auto-skip topk for sequences <= nsa_index_topk
if m.use_nsa:
q_lora = m.q_a_layernorm(q)
q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
_ = m.indexer(
x=hidden_states,
q_lora=q_lora,
positions=positions,
forward_batch=forward_batch,
layer_id=m.layer_id,
return_indices=False,
)
else:
q = m.q_a_layernorm(q)
q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim)
else:
q = m.q_proj(hidden_states)[0].view(-1, m.num_local_heads, m.qk_head_dim)
latent_cache = m.kv_a_proj_with_mqa(hidden_states)[0]
_, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1)
kv_a, _ = latent_cache.split([m.kv_lora_rank, m.qk_rope_head_dim], dim=-1)
latent_cache = latent_cache.unsqueeze(1)
kv_a = m.kv_a_layernorm(kv_a)
kv = m.kv_b_proj(kv_a)[0]
k_pe = latent_cache[:, :, m.kv_lora_rank :]
if m.rotary_emb is not None:
q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
q[..., m.qk_nope_head_dim :] = q_pe
m._set_mla_kv_buffer(latent_cache, kv_a, k_pe, forward_batch)
if forward_batch.mha_one_shot and sum(forward_batch.extend_prefix_lens_cpu) != 0:
if m.use_nsa and m.kv_cache_dtype == "fp8_e4m3":
# FP8 path: dequantize NSA-specific FP8 format to BF16
kv_a, k_pe = m._get_mla_kv_buffer_from_fp8(forward_batch)
else:
# BF16/FP16 path: directly fetch from cache
kv_a, k_pe = m._get_mla_kv_buffer(
forward_batch.fetch_mha_one_shot_kv_indices(),
q.dtype,
forward_batch,
)
kv = m.kv_b_proj(kv_a)[0]
kv = kv.view(-1, m.num_local_heads, m.qk_nope_head_dim + m.v_head_dim)
k_nope = kv[..., : m.qk_nope_head_dim]
v = kv[..., m.qk_nope_head_dim :]
k = m._concat_and_cast_mha_k(k_nope, k_pe, forward_batch)
return q, k, v, forward_batch
def forward_mha_core_npu(
m: "DeepseekV2AttentionMLA",
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
forward_batch: "ForwardBatch",
) -> torch.Tensor:
attn_output = m.attn_mha(q, k, v, forward_batch, save_kv_cache=False)
attn_output = attn_output.reshape(-1, m.num_local_heads * m.v_head_dim)
output, _ = m.o_proj(attn_output)
return output
# endregion
# region MLA
def forward_mla_prepare_npu(
m: "DeepseekV2AttentionMLA",
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: "ForwardBatch",
zero_allocator: "BumpAllocator",
):
if is_mla_preprocess_enabled():
if not hasattr(m, "mla_preprocess"):
m.mla_preprocess = NPUFusedMLAPreprocess(
m.fused_qkv_a_proj_with_mqa,
m.q_a_layernorm,
m.kv_a_layernorm,
m.q_b_proj,
m.w_kc,
m.rotary_emb,
m.layer_id,
m.num_local_heads,
m.qk_nope_head_dim,
m.qk_rope_head_dim,
m.quant_config,
)
(
q_pe,
k_pe,
q_nope_out,
k_nope,
forward_batch,
zero_allocator,
positions,
) = m.mla_preprocess.forward(
positions, hidden_states, forward_batch, zero_allocator
)
topk_indices = None
else:
q_lora = None
if m.q_lora_rank is not None:
q, latent_cache = (
get_attn_tp_context()
.fetch_qkv_latent()
.split(
[m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim],
dim=-1,
)
)
k_nope = latent_cache[..., : m.kv_lora_rank]
q = m.q_a_layernorm(q)
k_nope = m.kv_a_layernorm(k_nope)
# q_lora needed by indexer
if m.use_nsa:
q_lora = q
k_nope = k_nope.unsqueeze(1)
q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim)
else:
q = m.q_proj(hidden_states)[0].view(-1, m.num_local_heads, m.qk_head_dim)
latent_cache = m.kv_a_proj_with_mqa(hidden_states)[0]
k_nope = latent_cache[..., : m.kv_lora_rank]
k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1)
q_nope, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1)
k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1)
q_nope_out = torch.bmm(q_nope.transpose(0, 1), m.w_kc)
q_nope_out = q_nope_out.transpose(0, 1)
if enable_prefill_cp(forward_batch, m.nsa_enable_prefill_cp):
positions = cp_split_and_rebuild_position(forward_batch, positions)
q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
if enable_prefill_cp(forward_batch, m.nsa_enable_prefill_cp):
# support allgather+rerrange
k_nope, k_pe = m.rebuild_cp_kv_cache(
latent_cache, forward_batch, k_nope, k_pe
)
topk_indices = None
if q_lora is not None:
topk_indices = m.indexer(
x=hidden_states,
q_lora=q_lora,
positions=positions,
forward_batch=forward_batch,
layer_id=m.layer_id,
)
return (
q_pe,
k_pe,
q_nope_out,
k_nope,
forward_batch,
zero_allocator,
positions,
topk_indices,
)
def forward_mla_core_npu(
m: "DeepseekV2AttentionMLA",
q_pe: torch.Tensor,
k_pe: torch.Tensor,
q_nope_out: torch.Tensor,
k_nope: torch.Tensor,
forward_batch: "ForwardBatch",
zero_allocator: "BumpAllocator",
positions: torch.Tensor,
topk_indices: torch.Tensor,
) -> torch.Tensor:
attn_output = m.attn_mqa(
q_nope_out,
k_nope,
k_nope,
forward_batch,
q_rope=q_pe,
k_rope=k_pe,
**(dict(topk_indices=topk_indices) if topk_indices is not None else {}),
)
attn_output = attn_output.view(-1, m.num_local_heads, m.kv_lora_rank)
attn_bmm_output = torch.empty(
(attn_output.shape[0], m.num_local_heads * m.v_head_dim),
dtype=attn_output.dtype,
device=attn_output.device,
)
torch.bmm(
attn_output.transpose(0, 1),
m.w_vc,
out=attn_bmm_output.view(-1, m.num_local_heads, m.v_head_dim).transpose(0, 1),
)
output, _ = m.o_proj(attn_bmm_output)
return output
# endregion
# region DSA
def forward_dsa_prepare_npu(
m: "DeepseekV2AttentionMLA",
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: "ForwardBatch",
zero_allocator: "BumpAllocator",
):
if is_mla_preprocess_enabled() and forward_batch.forward_mode.is_decode():
if not hasattr(m, "mla_preprocess"):
m.mla_preprocess = NPUFusedMLAPreprocess(
m.fused_qkv_a_proj_with_mqa,
m.q_a_layernorm,
m.kv_a_layernorm,
m.q_b_proj,
m.w_kc,
m.rotary_emb,
m.layer_id,
m.num_local_heads,
m.qk_nope_head_dim,
m.qk_rope_head_dim,
m.quant_config,
)
(
q_pe,
k_pe,
q_nope_out,
k_nope,
forward_batch,
zero_allocator,
positions,
) = m.mla_preprocess.forward(
positions, hidden_states, forward_batch, zero_allocator
)
fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0]
q, _ = fused_qkv_a_proj_out.split(
[m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
)
q_lora = m.q_a_layernorm(q)
else:
fused_qkv_a_proj_out = m.fused_qkv_a_proj_with_mqa(hidden_states)[0]
q, latent_cache = fused_qkv_a_proj_out.split(
[m.q_lora_rank, m.kv_lora_rank + m.qk_rope_head_dim], dim=-1
)
k_nope = latent_cache[..., : m.kv_lora_rank]
q = m.q_a_layernorm(q)
k_nope = m.kv_a_layernorm(k_nope)
q_lora = q.clone() # required for topk_indices
k_nope = k_nope.unsqueeze(1)
q = m.q_b_proj(q)[0].view(-1, m.num_local_heads, m.qk_head_dim)
q_nope, q_pe = q.split([m.qk_nope_head_dim, m.qk_rope_head_dim], dim=-1)
k_pe = latent_cache[..., m.kv_lora_rank :].unsqueeze(1)
q_nope_out = torch.bmm(q_nope.transpose(0, 1), m.w_kc)
q_nope_out = q_nope_out.transpose(0, 1)
q_pe, k_pe = m.rotary_emb(positions, q_pe, k_pe)
topk_indices = m.indexer(
hidden_states, q_lora, positions, forward_batch, m.layer_id
)
return (
q_pe,
k_pe,
q_nope_out,
k_nope,
topk_indices,
forward_batch,
zero_allocator,
positions,
)
def forward_dsa_core_npu(
m: "DeepseekV2AttentionMLA",
q_pe: torch.Tensor,
k_pe: torch.Tensor,
q_nope_out: torch.Tensor,
k_nope: torch.Tensor,
topk_indices: torch.Tensor,
forward_batch: "ForwardBatch",
zero_allocator: "BumpAllocator",
positions: torch.Tensor,
) -> torch.Tensor:
attn_output = m.attn_mqa(
q_nope_out.contiguous(),
k_nope.contiguous(),
k_nope.contiguous(),
forward_batch,
save_kv_cache=True, # False if forward_batch.forward_mode.is_extend() else True,
q_rope=q_pe.contiguous(),
k_rope=k_pe.contiguous(),
topk_indices=topk_indices,
)
attn_output = attn_output.view(-1, m.num_local_heads, m.kv_lora_rank)
attn_bmm_output = torch.empty(
(attn_output.shape[0], m.num_local_heads, m.v_head_dim),
dtype=attn_output.dtype,
device=attn_output.device,
)
if not forward_batch.forward_mode.is_decode():
attn_output = attn_output.transpose(0, 1)
torch.bmm(
attn_output,
m.w_vc,
out=attn_bmm_output.view(-1, m.num_local_heads, m.v_head_dim).transpose(
0, 1
),
)
else:
attn_output = attn_output.contiguous()
torch.ops.npu.batch_matmul_transpose(attn_output, m.w_vc, attn_bmm_output)
attn_bmm_output = attn_bmm_output.reshape(-1, m.num_local_heads * m.v_head_dim)
output, _ = m.o_proj(attn_bmm_output)
return output
# endregion
python/sglang/srt/hardware_backend/npu/moe/topk.py
+79
View File
@@ -0,0 +1,79 @@
from typing import TYPE_CHECKING, Optional
import torch
from sgl_kernel_npu.norm.l1_norm import l1_norm
from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
from sglang.srt.eplb.expert_location_dispatch import topk_ids_logical_to_physical
from sglang.srt.layers.moe.topk import StandardTopKOutput, select_experts
if TYPE_CHECKING:
from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
from sglang.srt.layers.moe.topk import TopKConfig, TopKOutput
def fused_topk_npu(
hidden_states: torch.Tensor,
router_logits: torch.Tensor,
topk_config: "TopKConfig",
num_token_non_padded: Optional[torch.Tensor] = None,
expert_location_dispatch_info: Optional["ExpertLocationDispatchInfo"] = None,
) -> "TopKOutput":
use_grouped_topk = topk_config.use_grouped_topk
renormalize = topk_config.renormalize
correction_bias = topk_config.correction_bias
if not use_grouped_topk:
topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k_softmax(
router_logits,
k=topk_config.top_k,
)
if renormalize:
topk_weights = l1_norm(
topk_weights
if topk_config.num_fused_shared_experts == 0
else topk_weights[:, :-1]
)
topk_weights = topk_weights.to(torch.float32)
elif use_grouped_topk and correction_bias is not None:
routed_scaling_factor = topk_config.routed_scaling_factor or 1
topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k(
router_logits.to(torch.float32),
k=topk_config.top_k,
bias=correction_bias.to(torch.float32),
k_group=topk_config.topk_group,
group_count=topk_config.num_expert_group,
group_select_mode=1,
renorm=0,
norm_type=1,
routed_scaling_factor=routed_scaling_factor,
eps=float(1e-20),
)
if renormalize:
topk_weights_sum = (
topk_weights.sum(dim=-1, keepdim=True)
if topk_config.num_fused_shared_experts == 0
else topk_weights[:, :-1].sum(dim=-1, keepdim=True)
)
topk_weights = topk_weights / topk_weights_sum
else:
topk_config.torch_native = True
return select_experts(
hidden_states=hidden_states,
router_logits=router_logits,
topk_config=topk_config,
num_token_non_padded=num_token_non_padded,
expert_location_dispatch_info=expert_location_dispatch_info,
)
if expert_location_dispatch_info is not None:
topk_ids = topk_ids_logical_to_physical(topk_ids, expert_location_dispatch_info)
get_global_expert_distribution_recorder().on_select_experts(topk_ids=topk_ids)
return StandardTopKOutput(topk_weights, topk_ids, router_logits)
python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py
+916
View File
@@ -0,0 +1,916 @@
from typing import TYPE_CHECKING
import numpy as np
import torch
from sglang.srt.hardware_backend.npu.utils import npu_format_cast
from sglang.srt.layers.quantization.base_config import FusedMoEMethodBase
from sglang.srt.utils import set_weight_attrs
if TYPE_CHECKING:
from sglang.srt.layers.moe import MoeRunnerConfig
from sglang.srt.layers.moe.token_dispatcher import (
CombineInput,
StandardDispatchOutput,
)
def npu_fused_experts(
hidden_states: torch.Tensor,
w13: torch.Tensor,
w13_scale: torch.Tensor,
w2: torch.Tensor,
w2_scale: torch.Tensor,
topk_weights: torch.Tensor,
topk_ids: torch.Tensor,
top_k: int,
**kwargs,
):
w13_offset = kwargs.get("w13_offset", None)
w2_offset = kwargs.get("w2_offset", None)
use_wna16 = kwargs.get("use_wna16", False)
original_shape = hidden_states.shape
original_dtype = hidden_states.dtype
scale_dtype = original_dtype if original_dtype == torch.bfloat16 else torch.float32
if len(original_shape) == 3:
hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
num_tokens = hidden_states.shape[0]
num_experts = w13.shape[0]
row_idx_len = num_tokens * top_k
row_idx = (
torch.arange(0, row_idx_len, dtype=torch.int32, device=topk_weights.device)
.view(top_k, -1)
.permute(1, 0)
.contiguous()
)
hidden_states, expanded_row_idx, expanded_expert_idx = (
torch.ops.npu.npu_moe_init_routing(
hidden_states, row_idx=row_idx, expert_idx=topk_ids, active_num=num_tokens
)
)
expert_tokens = torch.ops.npu.npu_moe_compute_expert_tokens(
expanded_expert_idx, num_experts
)
expert_tokens = expert_tokens.to(torch.int64)
# gmm1: gate_up_proj
if not use_wna16:
hidden_states, pertoken_scale = torch.ops.npu.npu_dynamic_quant(hidden_states)
scale_args13 = {
"scale": [w13_scale.to(scale_dtype)],
"per_token_scale": [pertoken_scale],
}
else:
scale_args13 = {
"antiquant_scale": [w13_scale],
"antiquant_offset": [w13_offset],
}
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[w13],
**scale_args13,
split_item=2,
group_list_type=0,
group_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)[0]
# act_fn: swiglu
hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
if not use_wna16:
hidden_states, pertoken_scale = torch.ops.npu.npu_dynamic_quant(hidden_states)
scale_args2 = {
"scale": [w2_scale.to(scale_dtype)],
"per_token_scale": [pertoken_scale],
}
else:
scale_args2 = {"antiquant_scale": [w2_scale], "antiquant_offset": [w2_offset]}
# gmm2: down_proj
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[w2],
**scale_args2,
split_item=2,
group_list_type=0,
group_type=0,
group_list=expert_tokens,
output_dtype=original_dtype,
)[0]
final_hidden_states = torch.ops.npu.npu_moe_finalize_routing(
hidden_states,
skip1=None,
skip2=None,
bias=None,
scales=topk_weights,
expanded_src_to_dst_row=expanded_row_idx,
export_for_source_row=topk_ids,
)
if len(original_shape) == 3:
final_hidden_states = final_hidden_states.view(original_shape)
return final_hidden_states
def npu_fused_moe_without_routing_weights_bf16(
layer, hidden_states, group_list_type, group_list, output_dtype
):
# gmm1: gate_up_proj
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w13_weight.permute(0, 2, 1)],
split_item=2,
group_list_type=group_list_type,
group_type=0,
group_list=group_list,
output_dtype=output_dtype,
)[0]
hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
# gmm2: down_proj
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w2_weight.permute(0, 2, 1)],
split_item=2,
group_list_type=group_list_type,
group_type=0,
group_list=group_list,
output_dtype=output_dtype,
)[0]
return hidden_states
class NPUW8A8Int8DynamicMoEMethod(FusedMoEMethodBase):
def create_weights(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
) -> None:
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
self.num_experts = num_experts
extra_weight_attrs.update(
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
)
# weight
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size,
dtype=torch.int8,
),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition,
dtype=torch.int8,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# scale
w13_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
w2_weight_scale = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
# offset
w13_weight_offset = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_offset", w13_weight_offset)
set_weight_attrs(w13_weight_offset, extra_weight_attrs)
w2_weight_offset = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset", w2_weight_offset)
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
def release_weight_cache(self, weight: torch.Tensor):
# .contiguous() introduces additional memory overhead and needs to be released using resize_(0)
origin_weight = weight.data.transpose(1, 2)
new_weight = origin_weight.contiguous()
origin_weight.untyped_storage().resize_(0)
return new_weight
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
weight_data = self.release_weight_cache(layer.w13_weight.data)
layer.w13_weight = torch.nn.Parameter(weight_data, requires_grad=False)
weight_data = self.release_weight_cache(layer.w2_weight.data)
layer.w2_weight = torch.nn.Parameter(weight_data, requires_grad=False)
layer.w13_weight_scale = torch.nn.Parameter(
layer.w13_weight_scale.data.squeeze(-1).contiguous().to(torch.float32),
requires_grad=False,
)
layer.w2_weight_scale = torch.nn.Parameter(
layer.w2_weight_scale.data.squeeze(-1).contiguous(), requires_grad=False
)
layer.w13_weight_offset = torch.nn.Parameter(
layer.w13_weight_offset.data.squeeze(-1).contiguous(), requires_grad=False
)
layer.w2_weight_offset = torch.nn.Parameter(
layer.w2_weight_offset.data.squeeze(-1).contiguous(), requires_grad=False
)
layer.w13_weight.data = npu_format_cast(layer.w13_weight.data)
layer.w2_weight.data = npu_format_cast(layer.w2_weight.data)
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig"
):
self.moe_runner_config = moe_runner_config
def apply(
self,
layer,
dispatch_output: "StandardDispatchOutput",
) -> "CombineInput":
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
x = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
topk_weights, topk_ids, _ = topk_output
topk_ids = topk_ids.to(torch.int32)
topk_weights = topk_weights.to(x.dtype)
output = npu_fused_experts(
hidden_states=x,
w13=layer.w13_weight,
w13_scale=layer.w13_weight_scale,
w2=layer.w2_weight,
w2_scale=layer.w2_weight_scale,
topk_weights=topk_weights,
topk_ids=topk_ids,
top_k=topk_ids.shape[1],
)
return StandardCombineInput(hidden_states=output)
def apply_without_routing_weights(
self,
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
):
# gmm1: gate_up_proj
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w13_weight],
split_item=2,
group_list_type=group_list_type,
group_type=0,
group_list=group_list,
output_dtype=torch.int32,
)[0]
# act_fn: swiglu
hidden_states, swiglu_out_scale = torch.ops.npu.npu_dequant_swiglu_quant(
x=hidden_states,
weight_scale=layer.w13_weight_scale,
activation_scale=hidden_states_scale,
bias=None,
quant_scale=None,
quant_offset=None,
group_index=group_list,
activate_left=True,
quant_mode=1,
)
# gmm2: down_proj
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w2_weight],
scale=[layer.w2_weight_scale.to(output_dtype)],
per_token_scale=[swiglu_out_scale],
split_item=2,
group_list_type=group_list_type,
group_type=0,
group_list=group_list,
output_dtype=output_dtype,
)[0]
return hidden_states
class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase):
def __init__(self) -> None:
self.group_size = 256
self.tp_size = 1
def create_weights(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
) -> None:
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
self.num_experts = num_experts
extra_weight_attrs.update(
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
)
# >> weight
w13_output_size = intermediate_size_per_partition
w2_output_size = hidden_size // 2
w13_weight = torch.nn.Parameter(
torch.empty(num_experts, w13_output_size, hidden_size, dtype=torch.int8),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
w2_output_size,
intermediate_size_per_partition,
dtype=torch.int8,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# >> scale
w13_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
w2_weight_scale = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
# >> offset
w13_weight_offset = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_weight_offset", w13_weight_offset)
set_weight_attrs(w13_weight_offset, extra_weight_attrs)
w2_weight_offset = torch.nn.Parameter(
torch.empty(num_experts, hidden_size, 1, dtype=torch.float32),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset", w2_weight_offset)
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
# >>> special param for w4a8
w13_weight_scale_second = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale_second", w13_weight_scale_second)
set_weight_attrs(w13_weight_scale_second, extra_weight_attrs)
w13_weight_offset_second = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_offset_second", w13_weight_offset_second)
set_weight_attrs(w13_weight_offset_second, extra_weight_attrs)
w2_weight_scale_second = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale_second", w2_weight_scale_second)
set_weight_attrs(w2_weight_scale_second, extra_weight_attrs)
w2_weight_offset_second = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=torch.float32,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset_second", w2_weight_offset_second)
set_weight_attrs(w2_weight_offset_second, extra_weight_attrs)
w13_scale_bias = torch.nn.Parameter(
torch.empty(
num_experts, 2 * intermediate_size_per_partition, 1, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w13_scale_bias", w13_scale_bias)
set_weight_attrs(w13_scale_bias, extra_weight_attrs)
w2_scale_bias = torch.nn.Parameter(
torch.empty(
num_experts, hidden_size, 16 // self.tp_size, dtype=torch.float32
),
requires_grad=False,
)
layer.register_parameter("w2_scale_bias", w2_scale_bias)
set_weight_attrs(w2_scale_bias, extra_weight_attrs)
def process_scale(self, weight: torch.Tensor, scale, per_group_scale):
scale = scale.transpose(1, 2).contiguous()
per_group_scale = per_group_scale.transpose(1, 2).contiguous()
group_num, k, n = weight.shape
# the weight of the new version is reduced by half by pack n, so it needs to be restored
n = n * 2
per_group_scale = per_group_scale.reshape(group_num, -1, n)
group_num, quantgroup_num, n = per_group_scale.shape
bias = None
scale_fp32 = (scale * per_group_scale).to(torch.float16).to(torch.float32)
scale_fp32_np = scale_fp32.cpu().numpy()
scale_fp32_np.dtype = np.uint32
sscale_uint64 = np.zeros((group_num, quantgroup_num, n * 2), dtype=np.uint32)
sscale_uint64[..., ::2] = scale_fp32_np
sscale_uint64_buffer = np.frombuffer(
sscale_uint64.tobytes(), dtype=np.int64
).copy()
sscale_uint64_tensor = torch.from_numpy(sscale_uint64_buffer).reshape(
group_num, quantgroup_num, n
)
sscale_uint64_tensor = sscale_uint64_tensor.npu()
return sscale_uint64_tensor, bias
def update_bias(self, layer, w13_bias, w2_bias):
layer.w13_scale_bias.data = (
layer.w13_scale_bias.data.transpose(1, 2).contiguous().sum(axis=1)
)
layer.w2_scale_bias.data = (
layer.w2_scale_bias.data.transpose(1, 2).contiguous().sum(axis=1)
)
def pack_to_int32(self, weight: torch.Tensor):
# pack 4 int8(int4*2) to int32, because in pytorch, we need to use int32 to represent int4
assert (
weight.shape[-1] % 4 == 0
), "the last dim of weight needs to be divided by 4"
return weight.view(torch.int32).contiguous()
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
layer.w13_weight = torch.nn.Parameter(
layer.w13_weight.data.transpose(1, 2).contiguous(), requires_grad=False
)
layer.w2_weight = torch.nn.Parameter(
layer.w2_weight.data.transpose(1, 2).contiguous(), requires_grad=False
)
w13_weight_scale_second = (
layer.w13_weight_scale_second.data
if hasattr(layer, "w13_weight_scale_second")
else None
)
w2_weight_scale_second = (
layer.w2_weight_scale_second.data
if hasattr(layer, "w2_weight_scale_second")
else None
)
layer.w13_weight_scale.data, w13_bias = self.process_scale(
layer.w13_weight, layer.w13_weight_scale.data, w13_weight_scale_second
)
layer.w2_weight_scale.data, w2_bias = self.process_scale(
layer.w2_weight, layer.w2_weight_scale.data, w2_weight_scale_second
)
if hasattr(layer, "w13_weight_scale_second"):
# scale_second is no longer used, release this part of the memory
del layer.w13_weight_scale_second
del layer.w2_weight_scale_second
del layer.w13_weight_offset_second
del layer.w2_weight_offset_second
self.update_bias(layer, w13_bias, w2_bias)
layer.w13_weight.data = npu_format_cast(layer.w13_weight.data)
layer.w2_weight.data = npu_format_cast(layer.w2_weight.data)
layer.w13_weight.data = self.pack_to_int32(layer.w13_weight.data)
layer.w2_weight.data = self.pack_to_int32(layer.w2_weight.data)
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig"
):
self.moe_runner_config = moe_runner_config
def apply(
self,
layer,
dispatch_output: "StandardDispatchOutput",
) -> "CombineInput":
# FIXME W4A8 only support with deepep
raise NotImplementedError(
f"W4A8 only support with deepep for now, please enable --moe-a2a-backend deepep"
)
def apply_without_routing_weights(
self,
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
):
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[self.w13_weight],
scale=[self.w13_weight_scale],
bias=[self.w13_scale_bias],
per_token_scale=[hidden_states_scale],
group_list=group_list,
split_item=2,
group_type=0,
group_list_type=group_list_type,
output_dtype=output_dtype,
)[0]
# act_fn: swiglu
hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
hidden_states, swiglu_out_scale = torch.ops.npu.npu_dynamic_quant(hidden_states)
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[self.w2_weight],
scale=[self.w2_weight_scale],
bias=[self.w2_scale_bias],
per_token_scale=[swiglu_out_scale],
group_list=group_list,
split_item=2,
group_type=0,
group_list_type=group_list_type,
output_dtype=output_dtype,
)[0]
return hidden_states
class NPUW4A16Int4DynamicMoEMethod(FusedMoEMethodBase):
def __init__(self, quantization_config) -> None:
self.pack_factor = 8 # weight dtype is int4, but use int32 to create
target = (
"MoEGMM" if "MoEGMM" in quantization_config.target_scheme_map else "Linear"
)
if target in quantization_config.target_scheme_map:
self.group_size = quantization_config.target_scheme_map[target][
"weights"
].group_size
else:
self.group_size = 128
def create_weights(
self,
layer: torch.nn.Module,
num_experts: int,
hidden_size: int,
intermediate_size_per_partition: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
) -> None:
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
self.num_experts = num_experts
if (
extra_weight_attrs.get(
"intermediate_size_full", intermediate_size_per_partition
)
// intermediate_size_per_partition
> 1
):
quant_method = FusedMoeWeightScaleSupported.GROUP.value
else:
quant_method = FusedMoeWeightScaleSupported.CHANNEL.value
extra_weight_attrs.update({"quant_method": quant_method})
# weight
w13_weight = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.pack_factor,
dtype=torch.int32,
),
requires_grad=False,
)
layer.register_parameter("w13_weight", w13_weight)
set_weight_attrs(w13_weight, extra_weight_attrs)
w2_weight = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.pack_factor,
dtype=torch.int32,
),
requires_grad=False,
)
layer.register_parameter("w2_weight", w2_weight)
set_weight_attrs(w2_weight, extra_weight_attrs)
# scale
weight_scale_dtype = torch.bfloat16
w13_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=weight_scale_dtype,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_scale", w13_weight_scale)
set_weight_attrs(w13_weight_scale, extra_weight_attrs)
w2_weight_scale = torch.nn.Parameter(
torch.empty(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=weight_scale_dtype,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_scale", w2_weight_scale)
set_weight_attrs(w2_weight_scale, extra_weight_attrs)
# offset
w13_weight_offset = torch.nn.Parameter(
torch.zeros(
num_experts,
2 * intermediate_size_per_partition,
hidden_size // self.group_size,
dtype=weight_scale_dtype,
),
requires_grad=False,
)
layer.register_parameter("w13_weight_offset", w13_weight_offset)
set_weight_attrs(w13_weight_offset, extra_weight_attrs)
w2_weight_offset = torch.nn.Parameter(
torch.zeros(
num_experts,
hidden_size,
intermediate_size_per_partition // self.group_size,
dtype=weight_scale_dtype,
),
requires_grad=False,
)
layer.register_parameter("w2_weight_offset", w2_weight_offset)
set_weight_attrs(w2_weight_offset, extra_weight_attrs)
def pack_to_int32(self, weight: torch.Tensor):
assert weight.dim() == 3
if weight.dtype == torch.int32:
# pack 8 int4 to int32, we use a int32 to represent a int4
assert (
weight.shape[-1] % 8 == 0
), "the last dim of weight needs to be divided by 8"
new_weight = torch.ops.npu.npu_convert_weight_to_int4pack(
weight.flatten(0, 1)
)
new_weight = new_weight.view(weight.shape[0], weight.shape[1], -1)
elif weight.dtype == torch.int8:
# pack 4 int8(int4*2) to int32, because in pytorch, we need to use int32 to represent int4
assert (
weight.shape[-1] % 4 == 0
), "the last dim of weight needs to be divided by 4"
new_weight = weight.view(torch.int32).contiguous()
else:
raise ValueError(f"{weight.dtype=} is not supported !")
return new_weight
def unpack_from_int32(
self,
value: torch.Tensor,
num_bits: int,
shape: torch.Size = None,
packed_dim=1,
) -> torch.Tensor:
"""
Unpacks a tensor of packed int32 weights into individual int8s, maintaining the
original bit range.
Return tensors in int8
:param value: tensor to unpack
:param num_bits: number of bits to unpack each data point into
:param shape: shape to unpack into, used to remove padding
:returns: unpacked int8 tensor
"""
if value.dtype is not torch.int32:
raise ValueError(
f"Expected {torch.int32} but got {value.dtype}, Aborting unpack."
)
if num_bits > 8:
raise ValueError("Unpacking is only supported for less than 8 bits")
pack_factor = 32 // num_bits
# unpack
mask = (1 << num_bits) - 1
if packed_dim == 1:
unpacked = torch.zeros(
(value.shape[0], value.shape[1] * pack_factor),
device=value.device,
dtype=torch.int32,
)
for i in range(pack_factor):
unpacked[:, i::pack_factor] = (value >> (num_bits * i)) & mask
# remove padding
if shape is not None:
original_row_size = int(shape[1])
unpacked = unpacked[:, :original_row_size]
else:
unpacked = torch.zeros(
(value.shape[0] * pack_factor, value.shape[1]),
device=value.device,
dtype=torch.int32,
)
for i in range(pack_factor):
unpacked[i::pack_factor, :] = (value >> (num_bits * i)) & mask
# remove padding
original_row_size = int(shape[0])
unpacked = unpacked[:original_row_size, :]
# bits are packed in unsigned format, reformat to signed
# update the value range from unsigned to signed
offset = pow(2, num_bits) // 2
unpacked = (unpacked - offset).to(torch.int8)
return unpacked
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
w13_weight_scale = layer.w13_weight_scale.data.transpose(-1, -2).contiguous()
w2_weight_scale = layer.w2_weight_scale.data.transpose(-1, -2).contiguous()
layer.w13_weight_scale = torch.nn.Parameter(
w13_weight_scale, requires_grad=False
)
layer.w2_weight_scale = torch.nn.Parameter(w2_weight_scale, requires_grad=False)
layer.w13_weight_offset = torch.nn.Parameter(
layer.w13_weight_offset.data.transpose(-1, -2).contiguous(),
requires_grad=False,
)
layer.w2_weight_offset = torch.nn.Parameter(
layer.w2_weight_offset.data.transpose(-1, -2).contiguous(),
requires_grad=False,
)
# w = [n, k // 8] --> [k, n // 8]
# w13_weight = layer.w13_weight.data.transpose(1, 2).contiguous()
# w2_weight = layer.w2_weight.data.transpose(1, 2).contiguous()
unpacked_w13_weight = (
self.unpack_from_int32(layer.w13_weight.data.flatten(0, 1), 4)
.view(layer.w13_weight.data.shape[0], layer.w13_weight.data.shape[1], -1)
.transpose(1, 2)
.contiguous()
.int()
)
unpacked_w2_weight = (
self.unpack_from_int32(layer.w2_weight.data.flatten(0, 1), 4)
.view(layer.w2_weight.data.shape[0], layer.w2_weight.data.shape[1], -1)
.transpose(1, 2)
.contiguous()
.int()
)
w13_weight = self.pack_to_int32(unpacked_w13_weight)
w2_weight = self.pack_to_int32(unpacked_w2_weight)
layer.w13_weight = torch.nn.Parameter(w13_weight, requires_grad=False)
layer.w2_weight = torch.nn.Parameter(w2_weight, requires_grad=False)
def create_moe_runner(
self, layer: torch.nn.Module, moe_runner_config: "MoeRunnerConfig"
):
self.moe_runner_config = moe_runner_config
def apply(
self,
layer,
dispatch_output: "StandardDispatchOutput",
) -> "CombineInput":
from sglang.srt.layers.moe.token_dispatcher import StandardCombineInput
x = dispatch_output.hidden_states
topk_output = dispatch_output.topk_output
topk_weights, topk_ids, _ = topk_output
topk_ids = topk_ids.to(torch.int32)
topk_weights = topk_weights.to(x.dtype)
output = npu_fused_experts(
hidden_states=x,
w13=layer.w13_weight,
w13_scale=layer.w13_weight_scale,
w13_offset=layer.w13_weight_offset,
w2=layer.w2_weight,
w2_scale=layer.w2_weight_scale,
w2_offset=layer.w2_weight_offset,
topk_weights=topk_weights,
topk_ids=topk_ids,
top_k=topk_ids.shape[1],
use_wna16=True,
)
return StandardCombineInput(hidden_states=output)
def apply_without_routing_weights(
self,
layer,
hidden_states,
hidden_states_scale,
group_list_type,
group_list,
output_dtype,
):
if hidden_states_scale is None:
# gmm1: gate_up_proj
hidden_states = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w13_weight],
antiquant_scale=[layer.w13_weight_scale],
antiquant_offset=[layer.w13_weight_offset],
split_item=2,
group_list_type=group_list_type,
group_type=0,
group_list=group_list,
output_dtype=output_dtype,
)[0]
# act_fn: swiglu
hidden_states = torch.ops.npu.npu_swiglu(hidden_states)
# gmm2: down_proj
out_hidden = torch.ops.npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w2_weight],
antiquant_scale=[layer.w2_weight_scale],
antiquant_offset=[layer.w2_weight_offset],
split_item=2,
group_list_type=group_list_type,
group_type=0,
group_list=group_list,
output_dtype=output_dtype,
)[0]
else:
raise ValueError(
"when weight is int4, hidden_states only supports non-quant dtype!"
)
return out_hidden
python/sglang/srt/hardware_backend/npu/quantization/linear_method_npu.py
+215
View File
@@ -0,0 +1,215 @@
from typing import TYPE_CHECKING, List, Optional
import torch
from sglang.srt.hardware_backend.npu.utils import npu_format_cast
from sglang.srt.layers.parameter import (
ChannelQuantScaleParameter,
ModelWeightParameter,
PerTensorScaleParameter,
)
from sglang.srt.layers.quantization.base_config import LinearMethodBase
if TYPE_CHECKING:
from sglang.srt.layers.quantization.base_config import QuantizationConfig
class _NPULinearMethodBase(LinearMethodBase):
def __init__(
self,
quant_config: Optional["QuantizationConfig"] = None,
):
super().__init__()
self.quant_config = quant_config
class NPUW8A8Int8LinearMethod(_NPULinearMethodBase):
def create_weights(
self,
layer: torch.nn.Module,
input_size_per_partition: int,
output_partition_sizes: List[int],
input_size: int,
output_size: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
):
weight_loader = extra_weight_attrs.get("weight_loader")
output_size_per_partition = sum(output_partition_sizes)
weight = ModelWeightParameter(
data=torch.empty(
(output_size_per_partition, input_size_per_partition), dtype=torch.int8
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight", weight)
weight_scale = ChannelQuantScaleParameter(
data=torch.empty((output_size_per_partition, 1), dtype=params_dtype),
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight_scale", weight_scale)
weight_offset = ChannelQuantScaleParameter(
data=torch.empty((output_size_per_partition, 1), dtype=params_dtype),
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight_offset", weight_offset)
input_scale = PerTensorScaleParameter(
data=torch.empty(1, dtype=params_dtype),
weight_loader=weight_loader,
)
input_scale.ignore_warning = True
layer.register_parameter("input_scale", input_scale)
input_offset = PerTensorScaleParameter(
data=torch.empty(1, dtype=params_dtype),
weight_loader=weight_loader,
)
input_offset.ignore_warning = True
layer.register_parameter("input_offset", input_offset)
quant_bias = ChannelQuantScaleParameter(
data=torch.empty(output_size_per_partition, dtype=torch.int32),
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("quant_bias", quant_bias)
if params_dtype == torch.bfloat16:
deq_scale_dtype = torch.float32
elif params_dtype == torch.float16:
deq_scale_dtype = torch.int64
else:
raise ValueError(f"Unsupported params_dtype: {params_dtype}")
deq_scale = ChannelQuantScaleParameter(
data=torch.empty(output_size_per_partition, dtype=deq_scale_dtype),
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("deq_scale", deq_scale)
def apply(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
from sglang.srt.layers.linear import RowParallelLinear
original_dtype = x.dtype
if original_dtype != torch.int8:
x = torch.ops.npu.npu_quantize(
x,
layer.aclnn_input_scale_reciprocal,
layer.aclnn_input_offset,
torch.qint8,
-1,
False,
)
# Only fuse bias add into GEMM for rank 0 (this ensures that
# bias will not get added more than once in Attention TP>1 case)
if isinstance(layer, RowParallelLinear) and layer.tp_rank > 0:
quant_bias = None
else:
quant_bias = layer.quant_bias
return torch.ops.npu.npu_quant_matmul(
x,
layer.weight,
layer.deq_scale,
bias=quant_bias,
output_dtype=original_dtype,
)
def process_weights_after_loading(self, layer: torch.nn.Module):
layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
layer.weight.data = npu_format_cast(layer.weight.data)
layer.weight_scale.data = torch.flatten(layer.weight_scale.data)
layer.weight_offset.data = torch.flatten(layer.weight_offset.data)
expanding_factor = layer.weight.data.shape[0]
layer.aclnn_input_scale = torch.nn.Parameter(
layer.input_scale.data.repeat(expanding_factor).to(device="npu"),
requires_grad=False,
)
layer.aclnn_input_scale_reciprocal = 1 / torch.nn.Parameter(
layer.input_scale.data.repeat(expanding_factor).to(device="npu"),
requires_grad=False,
)
layer.aclnn_input_offset = torch.nn.Parameter(
layer.input_offset.data.repeat(expanding_factor).to(device="npu"),
requires_grad=False,
)
class NPUW8A8Int8DynamicLinearMethod(_NPULinearMethodBase):
def create_weights(
self,
layer: torch.nn.Module,
input_size_per_partition: int,
output_partition_sizes: List[int],
input_size: int,
output_size: int,
params_dtype: torch.dtype,
**extra_weight_attrs,
):
weight_loader = extra_weight_attrs.get("weight_loader")
output_size_per_partition = sum(output_partition_sizes)
weight = ModelWeightParameter(
data=torch.empty(
(output_size_per_partition, input_size_per_partition), dtype=torch.int8
),
input_dim=1,
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight", weight)
weight_scale = ChannelQuantScaleParameter(
data=torch.empty((output_size_per_partition, 1), dtype=params_dtype),
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight_scale", weight_scale)
weight_offset = ChannelQuantScaleParameter(
data=torch.empty((output_size_per_partition, 1), dtype=params_dtype),
output_dim=0,
weight_loader=weight_loader,
)
layer.register_parameter("weight_offset", weight_offset)
def apply(
self,
layer: torch.nn.Module,
x: torch.Tensor,
bias: Optional[torch.Tensor] = None,
) -> torch.Tensor:
original_dtype = x.dtype
quant_out, dynamic_scale = torch.ops.npu.npu_dynamic_quant(x)
return torch.ops.npu.npu_quant_matmul(
quant_out,
layer.weight,
layer.weight_scale,
pertoken_scale=dynamic_scale,
bias=bias,
output_dtype=original_dtype,
)
def process_weights_after_loading(self, layer: torch.nn.Module):
layer.weight.data = layer.weight.data.transpose(0, 1).contiguous()
layer.weight.data = npu_format_cast(layer.weight.data)
layer.weight_scale.data = layer.weight_scale.data.flatten()
layer.weight_offset.data = layer.weight_offset.data.flatten()
python/sglang/srt/hardware_backend/npu/quantization/modelslim.py
+241
View File
@@ -0,0 +1,241 @@
from __future__ import annotations
from types import MappingProxyType
from typing import Any, Dict, List, Mapping, Optional, Tuple, Union, cast
import torch
from compressed_tensors.quantization import QuantizationStrategy
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
NPUW4A8Int4DynamicMoEMethod,
NPUW4A16Int4DynamicMoEMethod,
NPUW8A8Int8DynamicMoEMethod,
)
from sglang.srt.hardware_backend.npu.quantization.linear_method_npu import (
NPUW8A8Int8DynamicLinearMethod,
NPUW8A8Int8LinearMethod,
)
from sglang.srt.layers.quantization.base_config import (
QuantizationConfig,
QuantizeMethodBase,
)
from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import (
CompressedTensorsConfig,
)
from sglang.srt.layers.quantization.compressed_tensors.utils import should_ignore_layer
from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
from sglang.srt.utils import apply_module_patch
# func refers to RMSNorm.__init__
def npu_wrapper_rmsnorm_init(func):
def init(self, hidden_size: int, **extra_args) -> None:
func(self, hidden_size, **extra_args)
self.ignore_anti = True
# The Ascend w8a8_int8 quantization requires adding a bias in rmsnorm
self.bias = torch.nn.Parameter(torch.zeros(hidden_size), requires_grad=False)
return init
# func refers to RMSNorm.forward_oot
def npu_wrapper_rmsnorm_forward(func):
def _rmsnorm_forward_oot(
self,
x: torch.Tensor,
residual: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
from sgl_kernel_npu.norm.add_rmsnorm_bias import add_rmsnorm_bias
if not x.is_contiguous():
x = x.contiguous()
if residual is not None:
out, residual_out = add_rmsnorm_bias(
x,
residual,
self.weight.data,
self.bias,
self.variance_epsilon,
)
return out.to(x.dtype), residual_out
out = torch.ops.npu.npu_rms_norm(x, self.weight.data, self.variance_epsilon)[0]
out = out + self.bias
return out.to(x.dtype)
return _rmsnorm_forward_oot
class ModelSlimConfig(QuantizationConfig):
"""
Config class for ModelSlim Quantization, a NPU-specific quantization type.
"""
def __init__(self, quant_config: Dict[str, Any] = {}):
super().__init__()
self.quant_description = quant_config
self.is_dynamic = quant_config.get("is_dynamic", False)
self.is_moe_w4_dynamic = False
ignore = cast(List[str], quant_config.get("ignore", []))
self.ignore = ignore if ignore is not None else []
packed_modules_mapping = quant_config.get("packed_modules_mapping", {})
self.packed_modules_mapping = (
packed_modules_mapping if packed_modules_mapping is not None else {}
)
self.target_scheme_map = (
CompressedTensorsConfig._quantization_scheme_map_from_config(
config=quant_config
)
)
target = "MoEGMM" if "MoEGMM" in self.target_scheme_map else "Linear"
target_scheme = self.target_scheme_map.get(target, None)
if target_scheme is None:
self.is_moe_w4_dynamic = False
else:
weight_quant = target_scheme.get("weights")
input_quant = target_scheme.get("input_activations")
self.is_moe_w4_dynamic = self.is_dynamic_token_w4(weight_quant, input_quant)
self.is_moe_input_quant = input_quant
for name in self.quant_description.keys():
if "norm.bias" in name:
apply_module_patch(
"sglang.srt.layers.layernorm.RMSNorm",
"__init__",
[npu_wrapper_rmsnorm_init],
)
apply_module_patch(
"sglang.srt.layers.layernorm.RMSNorm",
"forward_npu",
[npu_wrapper_rmsnorm_forward],
)
@classmethod
def get_supported_act_dtypes(cls) -> List[torch.dtype]:
return [torch.int8, torch.float16, torch.bfloat16]
@classmethod
def get_min_capability(cls) -> int:
return 0
@classmethod
def get_name(self) -> str:
return "modelslim"
@classmethod
def get_config_filenames(cls) -> List[str]:
filenames = ["quant_model_description.json"]
return filenames
@classmethod
def from_config(cls, config: Dict[str, Any]) -> ModelSlimConfig:
return cls(config)
def get_quant_method(
self,
layer: torch.nn.Module,
prefix: str,
) -> Optional[QuantizeMethodBase]:
from sglang.srt.layers.linear import LinearBase
from sglang.srt.layers.moe.fused_moe_triton import FusedMoE
if isinstance(layer, LinearBase):
if should_ignore_layer(
prefix,
ignore=self.ignore,
fused_mapping=self.packed_modules_mapping,
):
return UnquantizedLinearMethod()
key = "model"
if "vision_model" in prefix:
key = "vision_model"
elif "visual" in prefix:
key = "visual"
packed_modules_mapping_subset = self.packed_modules_mapping.get(key, {})
prefix_in_quant_config = prefix
proj_name = prefix.split(".")[-1]
if proj_name in packed_modules_mapping_subset:
prefix_in_quant_config = prefix.replace(
proj_name, packed_modules_mapping_subset[proj_name][0]
)
self.is_dynamic = (
self.quant_description[prefix_in_quant_config + ".weight"]
== "W8A8_DYNAMIC"
)
if self.is_layer_skipped(prefix, packed_modules_mapping_subset):
return UnquantizedLinearMethod()
return (
NPUW8A8Int8DynamicLinearMethod(self)
if self.is_dynamic
else NPUW8A8Int8LinearMethod(self)
)
elif isinstance(layer, FusedMoE):
prefix_in_quant_config = prefix + ".0.down_proj.weight"
is_moe_w4a8_dynamic = (
self.quant_description.get(prefix_in_quant_config, "STATIC")
== "W4A8_DYNAMIC"
)
if (
self.is_moe_w4_dynamic and self.is_moe_input_quant is not None
) or is_moe_w4a8_dynamic:
return NPUW4A8Int4DynamicMoEMethod()
elif self.is_moe_w4_dynamic and self.is_moe_input_quant is None:
return NPUW4A16Int4DynamicMoEMethod(self)
else:
return NPUW8A8Int8DynamicMoEMethod()
return None
def is_layer_skipped(
self, prefix: str, fused_mapping: Mapping[str, List[str]] = MappingProxyType({})
):
# adapted from vllm.model_executor.layers.quantization.utils.quant_utils.is_layer_skipped
proj_name = prefix.split(".")[-1]
if proj_name in fused_mapping:
shard_prefixes = [
prefix.replace(proj_name, shard_proj_name)
for shard_proj_name in fused_mapping[proj_name]
]
is_skipped = None
for shard_prefix in shard_prefixes:
is_shard_skipped = (
self.quant_description[shard_prefix + ".weight"] == "FLOAT"
)
if is_skipped is None:
is_skipped = is_shard_skipped
elif is_shard_skipped != is_skipped:
raise ValueError(
f"Detected some but not all shards of {prefix} "
"are quantized. All shards of fused layers "
"to have the same precision."
)
else:
is_skipped = self.quant_description[prefix + ".weight"] == "FLOAT"
assert is_skipped is not None
return is_skipped
def get_scaled_act_names(self) -> List[str]:
return []
def is_dynamic_token_w4(self, weight_quant, input_quant) -> bool:
is_w4 = weight_quant.num_bits == 4
weight_strategy = (
weight_quant.strategy == QuantizationStrategy.TENSOR.value
or weight_quant.strategy == QuantizationStrategy.CHANNEL.value
or weight_quant.strategy == QuantizationStrategy.GROUP.value
)
if input_quant is not None:
is_token = (
weight_strategy
and input_quant.strategy == QuantizationStrategy.TOKEN.value
)
is_dynamic = not weight_quant.dynamic and input_quant.dynamic
else:
is_token = weight_strategy
is_dynamic = not weight_quant.dynamic
# Both symmetric and asymmetric input quantization supported.
# Only symmetric weight quantization supported.
return is_w4 and weight_quant.symmetric and is_token and is_dynamic
python/sglang/srt/hardware_backend/npu/utils.py
+104
View File
@@ -0,0 +1,104 @@
import functools
import logging
from enum import IntEnum
from typing import TYPE_CHECKING, Callable
import torch
from sglang.srt.environ import envs
from sglang.srt.utils import is_npu
if TYPE_CHECKING:
from sglang.srt.server_args import ServerArgs
logger = logging.getLogger(__name__)
_is_npu = is_npu()
class NPUACLFormat(IntEnum):
ACL_FORMAT_UNDEFINED = -1
ACL_FORMAT_ND = 2
ACL_FORMAT_FRACTAL_NZ = 29
def _call_once(fn: Callable):
@functools.wraps(fn)
def wrapper(*args, **kwargs):
if getattr(fn, "_has_been_called", False):
logger.debug("Function {} has already been called.", fn.__name__)
return
fn._has_been_called = True
return fn(*args, **kwargs)
return wrapper
def set_default_server_args(args: "ServerArgs"):
"""
Set default server arguments for NPU backend.
"""
# NPU only works with "ascend" attention backend for now
args.attention_backend = "ascend"
args.prefill_attention_backend = "ascend"
args.decode_attention_backend = "ascend"
if args.page_size is None:
args.page_size = 128
# NPU does not support CustomAllReduce
args.disable_custom_all_reduce = True
# handles hierarchical cache configs
if args.enable_hierarchical_cache:
args.hicache_io_backend = "kernel_ascend"
if args.use_mla_backend():
args.hicache_mem_layout = "page_first_kv_split"
else:
args.hicache_mem_layout = "page_first_direct"
@_call_once
def init_npu_backend():
"""
Initialize NPU backend. This function should be called only once.
"""
assert _is_npu, "NPU backend initialization called on non-NPU device."
import sgl_kernel_npu # noqa: F401
import torch_npu
from torch_npu.contrib import transfer_to_npu # noqa: F401
# Re-mock torch.cuda.is_available cuz transfer_to_npu mocks it True
torch.cuda.is_available = lambda: False
torch_npu.npu.config.allow_internal_format = True
torch_npu.npu.set_compile_mode(jit_compile=False)
def npu_format_cast(
tensor: torch.Tensor,
acl_format: NPUACLFormat = NPUACLFormat.ACL_FORMAT_FRACTAL_NZ,
) -> torch.Tensor:
"""
Cast a tensor to a specific NPU ACL format.
Args:
tensor (torch.Tensor): The input tensor.
acl_format (NPUACLFormat): The target NPU ACL format.
Returns:
torch.Tensor: The tensor cast to the specified NPU ACL format.
"""
if not _is_npu:
return tensor
if envs.SGLANG_NPU_DISABLE_ACL_FORMAT_WEIGHT.get():
return tensor
import torch_npu
return torch_npu.npu_format_cast(tensor, acl_format.value)
python/sglang/srt/layers/attention/attention_registry.py
+3 -1
View File
@@ -70,7 +70,9 @@ def create_wave_backend(runner):
@register_attention_backend("ascend")
def create_ascend_backend(runner):
from sglang.srt.layers.attention.ascend_backend import AscendAttnBackend
from sglang.srt.hardware_backend.npu.attention.ascend_backend import (
AscendAttnBackend,
)
return AscendAttnBackend(runner)
python/sglang/srt/layers/communicator.py
+5 -2
View File
@@ -56,9 +56,9 @@ from sglang.srt.utils import (
is_flashinfer_available,
is_gfx95_supported,
is_hip,
is_npu,
is_sm90_supported,
is_sm100_supported,
prepare_weight_cache,
)
_is_cuda = is_cuda()
@@ -67,11 +67,14 @@ _is_sm90_supported = _is_cuda and is_sm90_supported()
_is_sm100_supported = _is_cuda and is_sm100_supported()
_use_aiter = get_bool_env_var("SGLANG_USE_AITER") and is_hip()
_is_gfx95_supported = is_gfx95_supported()
_is_npu = is_npu()
if _use_aiter and _is_gfx95_supported:
from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_group_quant
from sglang.srt.layers.quantization.rocm_mxfp4_utils import fused_rms_mxfp4_quant
elif _is_npu:
from sglang.srt.hardware_backend.npu.cmo import prepare_weight_cache
FUSE_ALLREDUCE_MAX_BATCH_SIZE = 2048
@@ -776,7 +779,7 @@ class CommunicateWithAllReduceAndLayerNormFn:
)
else:
hidden_states = tensor_model_parallel_all_reduce(hidden_states)
if context.cache is not None:
if _is_npu and context.cache is not None:
_ = prepare_weight_cache(hidden_states, context.cache)
hidden_states, residual = layernorm(hidden_states, residual)
return hidden_states, residual
python/sglang/srt/layers/moe/ep_moe/layer.py
+3 -27
View File
@@ -318,6 +318,9 @@ class DeepEPMoE(FusedMoE):
assert self.quant_method is not None
assert self.moe_runner_config.activation == "silu"
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
npu_fused_moe_without_routing_weights_bf16,
)
from sglang.srt.layers.moe.token_dispatcher import DispatchOutputChecker
# NOTE: Ascend's Dispatch & Combine does not support FP16
@@ -388,33 +391,6 @@ class DeepEPMoE(FusedMoE):
return hidden_states
def npu_fused_moe_without_routing_weights_bf16(
layer, hidden_states, group_list_type, group_list, output_dtype
):
# gmm1: gate_up_proj
hidden_states = torch_npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w13_weight.permute(0, 2, 1)],
split_item=2,
group_list_type=group_list_type,
group_type=0,
group_list=group_list,
output_dtype=output_dtype,
)[0]
hidden_states = torch_npu.npu_swiglu(hidden_states)
# gmm2: down_proj
hidden_states = torch_npu.npu_grouped_matmul(
x=[hidden_states],
weight=[layer.w2_weight.permute(0, 2, 1)],
split_item=2,
group_list_type=group_list_type,
group_type=0,
group_list=group_list,
output_dtype=output_dtype,
)[0]
return hidden_states
class NpuFuseEPMoE(DeepEPMoE):
def __init__(
self,
python/sglang/srt/layers/moe/topk.py
+9 -78
View File
@@ -91,9 +91,6 @@ if _use_aiter:
from aiter import biased_grouped_topk as aiter_biased_grouped_topk
except ImportError:
raise ImportError("aiter is required when SGLANG_USE_AITER is set to True")
if _is_npu:
import torch_npu
from sgl_kernel_npu.norm.l1_norm import l1_norm
# -------------------------------- TopKConfig ---------------------------------------
@@ -344,76 +341,15 @@ class TopK(CustomOp):
expert_location_dispatch_info: Optional[ExpertLocationDispatchInfo] = None,
) -> TopKOutput:
use_grouped_topk = self.topk_config.use_grouped_topk
torch_native = self.topk_config.torch_native
renormalize = self.topk_config.renormalize
from sglang.srt.hardware_backend.npu.moe.topk import fused_topk_npu
if not use_grouped_topk and not torch_native:
topk_weights, topk_ids, _ = torch_npu.npu_moe_gating_top_k_softmax(
router_logits,
k=self.topk_config.top_k,
)
if renormalize:
topk_weights = l1_norm(
topk_weights
if self.topk_config.num_fused_shared_experts == 0
else topk_weights[:, :-1]
)
topk_weights = topk_weights.to(torch.float32)
if expert_location_dispatch_info is not None:
topk_ids = topk_ids_logical_to_physical(
topk_ids, expert_location_dispatch_info
)
get_global_expert_distribution_recorder().on_select_experts(
topk_ids=topk_ids
)
return StandardTopKOutput(topk_weights, topk_ids, _)
if use_grouped_topk and not torch_native and router_logits.shape[-1] == 256:
# NOTE: now npu_moe_gating_top_k can only support `group_count=256` pattern
routed_scaling_factor = self.topk_config.routed_scaling_factor or 1
topk_weights, topk_ids, _ = torch_npu.npu_moe_gating_top_k(
router_logits.to(torch.float32),
k=self.topk_config.top_k,
bias=self.topk_config.correction_bias.to(torch.float32),
k_group=self.topk_config.topk_group,
group_count=self.topk_config.num_expert_group,
group_select_mode=1,
renorm=0,
norm_type=1,
routed_scaling_factor=routed_scaling_factor,
eps=float(1e-20),
)
if renormalize:
topk_weights_sum = (
topk_weights.sum(dim=-1, keepdim=True)
if self.topk_config.num_fused_shared_experts == 0
else topk_weights[:, :-1].sum(dim=-1, keepdim=True)
)
topk_weights = topk_weights / topk_weights_sum
if expert_location_dispatch_info is not None:
topk_ids = topk_ids_logical_to_physical(
topk_ids, expert_location_dispatch_info
)
get_global_expert_distribution_recorder().on_select_experts(
topk_ids=topk_ids
)
return StandardTopKOutput(topk_weights, topk_ids, _)
else:
self.topk_config.torch_native = True
return select_experts(
hidden_states=hidden_states,
router_logits=router_logits,
topk_config=self.topk_config,
num_token_non_padded=num_token_non_padded,
expert_location_dispatch_info=expert_location_dispatch_info,
)
return fused_topk_npu(
hidden_states=hidden_states,
router_logits=router_logits,
topk_config=self.topk_config,
num_token_non_padded=num_token_non_padded,
expert_location_dispatch_info=expert_location_dispatch_info,
)
def empty_topk_output(self, device: torch.device) -> TopKOutput:
topk = self.topk_config.top_k - self.topk_config.num_fused_shared_experts
@@ -548,7 +484,7 @@ def fused_topk(
# This is used by the Deepseek V2/V3/R1 series models
@torch.compile(dynamic=True, backend=get_compiler_backend())
@torch.compile(dynamic=True, backend=get_compiler_backend(), disable=_is_npu)
def grouped_topk_gpu(
hidden_states: torch.Tensor,
gating_output: torch.Tensor,
@@ -565,9 +501,6 @@ def grouped_topk_gpu(
assert hidden_states.shape[0] == gating_output.shape[0], "Number of tokens mismatch"
scores = torch.softmax(gating_output, dim=-1)
# NPU compiler limitation
if _is_npu and scores.dtype == torch.bfloat16:
scores = scores.to(torch.float16)
num_token = scores.shape[0]
num_experts = scores.shape[1]
group_scores = (
@@ -584,7 +517,6 @@ def grouped_topk_gpu(
.reshape(num_token, -1)
) # [n, e]
tmp_scores = scores.masked_fill(~score_mask.bool(), 0.0) # [n, e]
# TODO: NPU can't support directly evaluating a comparison for now
topk_weights, topk_ids = torch.topk(
tmp_scores,
k=topk,
@@ -729,7 +661,6 @@ def biased_grouped_topk_impl(
tmp_scores = scores_for_choice.masked_fill(
~score_mask.bool(), float("-inf")
) # [n, e]
# TODO: NPU can't support directly evaluating a comparison for now
_, topk_ids = torch.topk(
tmp_scores,
k=topk,
python/sglang/srt/layers/quantization/__init__.py
+10 -1
View File
@@ -39,7 +39,7 @@ from sglang.srt.layers.quantization.quark.quark import QuarkConfig
from sglang.srt.layers.quantization.w4afp8 import W4AFp8Config
from sglang.srt.layers.quantization.w8a8_fp8 import W8A8Fp8Config
from sglang.srt.layers.quantization.w8a8_int8 import W8A8Int8Config
from sglang.srt.utils import is_cuda, is_hip, mxfp_supported
from sglang.srt.utils import is_cuda, is_hip, is_npu, mxfp_supported
_is_mxfp_supported = mxfp_supported()
@@ -78,6 +78,15 @@ if is_cuda() or (_is_mxfp_supported and is_hip()):
}
)
if is_npu():
from sglang.srt.hardware_backend.npu.quantization.modelslim import ModelSlimConfig
BASE_QUANTIZATION_METHODS.update(
{
"modelslim": ModelSlimConfig,
}
)
QUANTIZATION_METHODS = {**BASE_QUANTIZATION_METHODS}
python/sglang/srt/layers/quantization/awq.py
+3 -1
View File
@@ -7,6 +7,9 @@ from typing import TYPE_CHECKING, Any, Dict, List, Optional
import torch
from sglang.srt.hardware_backend.npu.quantization.fused_moe_method_npu import (
npu_fused_experts,
)
from sglang.srt.layers.linear import LinearBase, set_weight_attrs
from sglang.srt.layers.parameter import GroupQuantScaleParameter, PackedvLLMParameter
from sglang.srt.layers.quantization.base_config import (
@@ -31,7 +34,6 @@ from sglang.srt.layers.quantization.marlin_utils import (
)
from sglang.srt.layers.quantization.unquant import UnquantizedLinearMethod
from sglang.srt.layers.quantization.utils import get_scalar_types, replace_parameter
from sglang.srt.layers.quantization.w8a8_int8 import npu_fused_experts
from sglang.srt.utils.patch_torch import register_fake_if_exists
if TYPE_CHECKING:
python/sglang/srt/layers/quantization/w8a8_int8.py
+13 -1163
View File
File diff suppressed because it is too large Load Diff
python/sglang/srt/lora/lora_manager.py
+1 -8
View File
@@ -37,16 +37,9 @@ from sglang.srt.lora.utils import (
from sglang.srt.managers.io_struct import LoRAUpdateOutput
from sglang.srt.model_executor.forward_batch_info import ForwardBatch
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import is_npu, replace_submodule
from sglang.srt.utils import replace_submodule
from sglang.srt.utils.hf_transformers_utils import AutoConfig
if is_npu():
from torch_npu.contrib import transfer_to_npu # noqa: F401
# Re-mock torch.cuda.is_available cuz transfer_to_npu mocks it to True
torch.cuda.is_available = lambda: False
logger = logging.getLogger(__name__)
python/sglang/srt/managers/cache_controller.py
+1 -1
View File
@@ -384,7 +384,7 @@ class HiCacheController:
self.tp_size = get_tensor_model_parallel_world_size()
self.dp_rank = 0
# Currently, AscendMLAPagedTokenToKVPool is the subclass of MLATokenToKVPool.
# Currently, NPUMLATokenToKVPool is the subclass of MLATokenToKVPool.
is_mla_backend = isinstance(self.mem_pool_device, MLATokenToKVPool)
return HiCacheStorageConfig(
python/sglang/srt/managers/schedule_batch.py
+1 -7
View File
@@ -83,11 +83,8 @@ from sglang.srt.sampling.sampling_batch_info import SamplingBatchInfo
from sglang.srt.sampling.sampling_params import SamplingParams
from sglang.srt.server_args import ServerArgs, get_global_server_args
from sglang.srt.utils import flatten_nested_list
from sglang.srt.utils.common import is_npu
from sglang.srt.utils.cuda_ipc_transport_utils import CudaIpcTensorTransportProxy
_is_npu = is_npu()
if TYPE_CHECKING:
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.speculative.eagle_info import EagleDraftInput
@@ -1152,10 +1149,7 @@ class ScheduleBatch(ScheduleBatchDisaggregationDecodeMixin):
has_grammar: bool = False
# Device
if not _is_npu:
device: str = "cuda"
else:
device: str = "npu"
device: str = "cuda"
# Speculative decoding
spec_algorithm: SpeculativeAlgorithm = None
python/sglang/srt/mem_cache/memory_pool.py
+18 -290
View File
@@ -62,8 +62,6 @@ logger = logging.getLogger(__name__)
GB = 1024 * 1024 * 1024
_is_cuda = is_cuda()
_is_npu = is_npu()
if _is_npu:
import torch_npu
def get_tensor_size_bytes(t: Union[torch.Tensor, List[torch.Tensor]]):
@@ -1001,10 +999,16 @@ class HybridLinearKVPool(KVCache):
assert not enable_kvcache_transpose
self.use_mla = use_mla
if not use_mla:
TokenToKVPoolClass = MHATokenToKVPool
if _is_npu:
TokenToKVPoolClass = AscendTokenToKVPool
else:
TokenToKVPoolClass = MHATokenToKVPool
from sglang.srt.hardware_backend.npu.memory_pool_npu import (
NPUMHATokenToKVPool,
)
TokenToKVPoolClass = NPUMHATokenToKVPool
self.full_kv_pool = TokenToKVPoolClass(
size=size,
page_size=self.page_size,
@@ -1016,7 +1020,16 @@ class HybridLinearKVPool(KVCache):
enable_memory_saver=enable_memory_saver,
)
else:
TokenToKVPoolClass = MLATokenToKVPool
if _is_npu:
from sglang.srt.hardware_backend.npu.memory_pool_npu import (
NPUMLATokenToKVPool,
)
TokenToKVPoolClass = NPUMLATokenToKVPool
self.full_kv_pool = TokenToKVPoolClass(
size=size,
page_size=self.page_size,
@@ -1290,95 +1303,6 @@ class SWAKVPool(KVCache):
)
class AscendTokenToKVPool(MHATokenToKVPool):
def _create_buffers(self):
with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
# [size, head_num, head_dim] for each layer
# The padded slot 0 is used for writing dummy outputs from padded tokens.
# Continuous memory improves the efficiency of Ascend`s transmission backend,
# while other backends remain unchanged.
self.kv_buffer = torch.zeros(
(
2,
self.layer_num,
self.size // self.page_size + 1,
self.page_size,
self.head_num,
self.head_dim,
),
dtype=self.store_dtype,
device=self.device,
)
self.k_buffer = self.kv_buffer[0]
self.v_buffer = self.kv_buffer[1]
# for disagg
def get_contiguous_buf_infos(self):
# layer_num x [seq_len, head_num, head_dim]
# layer_num x [page_num, page_size, head_num, head_dim]
kv_data_ptrs = [
self.get_key_buffer(i).data_ptr()
for i in range(self.start_layer, self.start_layer + self.layer_num)
] + [
self.get_value_buffer(i).data_ptr()
for i in range(self.start_layer, self.start_layer + self.layer_num)
]
kv_data_lens = [
self.get_key_buffer(i).nbytes
for i in range(self.start_layer, self.start_layer + self.layer_num)
] + [
self.get_value_buffer(i).nbytes
for i in range(self.start_layer, self.start_layer + self.layer_num)
]
kv_item_lens = [
self.get_key_buffer(i)[0].nbytes
for i in range(self.start_layer, self.start_layer + self.layer_num)
] + [
self.get_value_buffer(i)[0].nbytes
for i in range(self.start_layer, self.start_layer + self.layer_num)
]
return kv_data_ptrs, kv_data_lens, kv_item_lens
def set_kv_buffer(
self,
layer: RadixAttention,
loc: torch.Tensor,
cache_k: torch.Tensor,
cache_v: torch.Tensor,
k_scale: Optional[float] = None,
v_scale: Optional[float] = None,
layer_id_override: Optional[int] = None,
):
if layer_id_override is not None:
layer_id = layer_id_override
else:
layer_id = layer.layer_id
if cache_k.dtype != self.dtype:
if k_scale is not None:
cache_k.div_(k_scale)
if v_scale is not None:
cache_v.div_(v_scale)
cache_k = cache_k.to(self.dtype)
cache_v = cache_v.to(self.dtype)
if self.store_dtype != self.dtype:
cache_k = cache_k.view(self.store_dtype)
cache_v = cache_v.view(self.store_dtype)
torch_npu._npu_reshape_and_cache(
key=cache_k,
value=cache_v,
key_cache=self.k_buffer[layer_id - self.start_layer].view(
-1, self.page_size, self.head_num, self.head_dim
),
value_cache=self.v_buffer[layer_id - self.start_layer].view(
-1, self.page_size, self.head_num, self.head_dim
),
slot_indices=loc,
)
class MLATokenToKVPool(KVCache):
def __init__(
self,
@@ -1879,202 +1803,6 @@ class NSATokenToKVPool(MLATokenToKVPool):
return kv_size_bytes
class AscendMLAPagedTokenToKVPool(MLATokenToKVPool):
def __init__(
self,
size: int,
page_size: int,
dtype: torch.dtype,
kv_lora_rank: int,
qk_rope_head_dim: int,
index_head_dim: Optional[int],
layer_num: int,
device: str,
enable_memory_saver: bool,
start_layer: Optional[int] = None,
end_layer: Optional[int] = None,
):
super(MLATokenToKVPool, self).__init__(
size,
page_size,
dtype,
layer_num,
device,
enable_memory_saver,
start_layer,
end_layer,
)
self.kv_lora_rank = kv_lora_rank
self.qk_rope_head_dim = qk_rope_head_dim
self.index_head_dim = index_head_dim
self.custom_mem_pool = None
with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
# The padded slot 0 is used for writing dummy outputs from padded tokens.
self.k_buffer = torch.zeros(
(
layer_num,
self.size // self.page_size + 1,
self.page_size,
1,
self.kv_lora_rank,
),
dtype=self.store_dtype,
device=self.device,
)
self.v_buffer = torch.zeros(
(
layer_num,
self.size // self.page_size + 1,
self.page_size,
1,
self.qk_rope_head_dim,
),
dtype=self.store_dtype,
device=self.device,
)
if self.index_head_dim is not None:
self.index_k_buffer = torch.zeros(
(
layer_num,
self.size // self.page_size + 1,
self.page_size,
1,
self.index_head_dim,
),
dtype=self.store_dtype,
device=self.device,
)
self._finalize_allocation_log(size)
def get_kv_size_bytes(self):
assert hasattr(self, "k_buffer")
assert hasattr(self, "v_buffer")
kv_size_bytes = 0
for k_cache in self.k_buffer:
kv_size_bytes += get_tensor_size_bytes(k_cache)
for v_cache in self.v_buffer:
kv_size_bytes += get_tensor_size_bytes(v_cache)
if self.index_head_dim is not None:
assert hasattr(self, "index_k_buffer")
for index_k_cache in self.index_k_buffer:
kv_size_bytes += get_tensor_size_bytes(index_k_cache)
return kv_size_bytes
def get_kv_buffer(self, layer_id: int):
if self.layer_transfer_counter is not None:
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
return (
self.k_buffer[layer_id - self.start_layer],
self.v_buffer[layer_id - self.start_layer],
)
def get_key_buffer(self, layer_id: int):
if self.layer_transfer_counter is not None:
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
if self.store_dtype != self.dtype:
return self.k_buffer[layer_id - self.start_layer].view(self.dtype)
return self.k_buffer[layer_id - self.start_layer]
def get_value_buffer(self, layer_id: int):
if self.layer_transfer_counter is not None:
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
if self.store_dtype != self.dtype:
return self.v_buffer[layer_id - self.start_layer].view(self.dtype)
return self.v_buffer[layer_id - self.start_layer]
def get_index_k_buffer(self, layer_id: int):
if self.layer_transfer_counter is not None:
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
if self.store_dtype != self.dtype:
return self.index_k_buffer[layer_id - self.start_layer].view(self.dtype)
return self.index_k_buffer[layer_id - self.start_layer]
# for disagg
def get_contiguous_buf_infos(self):
# MLA has only one kv_buffer, so only the information of this buffer needs to be returned.
kv_data_ptrs = [self.k_buffer[i].data_ptr() for i in range(self.layer_num)] + [
self.v_buffer[i].data_ptr() for i in range(self.layer_num)
]
kv_data_lens = [self.k_buffer[i].nbytes for i in range(self.layer_num)] + [
self.v_buffer[i].nbytes for i in range(self.layer_num)
]
kv_item_lens = [self.k_buffer[i][0].nbytes for i in range(self.layer_num)] + [
self.v_buffer[i][0].nbytes for i in range(self.layer_num)
]
if self.index_head_dim is not None:
kv_data_ptrs += [
self.index_k_buffer[i].data_ptr() for i in range(self.layer_num)
]
kv_data_lens += [
self.index_k_buffer[i].nbytes for i in range(self.layer_num)
]
kv_item_lens += [
self.index_k_buffer[i][0].nbytes for i in range(self.layer_num)
]
return kv_data_ptrs, kv_data_lens, kv_item_lens
def set_kv_buffer(
self,
layer: RadixAttention,
loc: torch.Tensor,
cache_k: torch.Tensor,
cache_v: torch.Tensor,
):
layer_id = layer.layer_id
if cache_k.dtype != self.dtype:
cache_k = cache_k.to(self.dtype)
cache_v = cache_v.to(self.dtype)
if self.store_dtype != self.dtype:
cache_k = cache_k.view(self.store_dtype)
cache_v = cache_v.view(self.store_dtype)
if cache_v is None:
cache_k, cache_v = cache_k.split(
[self.kv_lora_rank, self.qk_rope_head_dim], dim=-1
)
torch_npu.npu_scatter_nd_update_(
self.k_buffer[layer_id - self.start_layer].view(-1, 1, self.kv_lora_rank),
loc.view(-1, 1),
cache_k.view(-1, 1, self.kv_lora_rank),
)
torch_npu.npu_scatter_nd_update_(
self.v_buffer[layer_id - self.start_layer].view(
-1, 1, self.qk_rope_head_dim
),
loc.view(-1, 1),
cache_v.view(-1, 1, self.qk_rope_head_dim),
)
def set_index_k_buffer(
self,
layer_id: int,
loc: torch.Tensor,
index_k: torch.Tensor,
):
if index_k.dtype != self.dtype:
index_k = index_k.to(self.dtype)
if self.store_dtype != self.dtype:
index_k = index_k.view(self.store_dtype)
torch_npu.npu_scatter_nd_update_(
self.index_k_buffer[layer_id - self.start_layer].view(
-1, 1, self.index_head_dim
),
loc.view(-1, 1),
index_k.view(-1, 1, self.index_head_dim),
)
class DoubleSparseTokenToKVPool(KVCache):
def __init__(
self,
python/sglang/srt/mem_cache/memory_pool_host.py
+1 -1
View File
@@ -666,7 +666,7 @@ class MLATokenToKVPoolHost(HostKVCache):
1,
self.kv_lora_rank + self.qk_rope_head_dim,
)
# Ascend-specific: Aligns with AscendMLAPagedTokenToKVPool layout
# Ascend-specific: Aligns with NPUMLATokenToKVPool layout
# Separately allocate k_buffer and v_buffer for easier data transfer.
elif self.layout == "page_first_kv_split":
base_dims = (
python/sglang/srt/model_executor/model_runner.py
+27 -19
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,
python/sglang/srt/models/deepseek_v2.py
+43 -301
View File
@@ -54,10 +54,6 @@ from sglang.srt.eplb.expert_location_dispatch import ExpertLocationDispatchInfo
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.activation import SiluAndMul
from sglang.srt.layers.amx_utils import PackWeightMethod
from sglang.srt.layers.attention.npu_ops.mla_preprocess import (
NPUFusedMLAPreprocess,
is_mla_preprocess_enabled,
)
from sglang.srt.layers.attention.nsa.dequant_k_cache import dequantize_k_cache_paged
from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer
from sglang.srt.layers.attention.nsa.utils import (
@@ -214,10 +210,14 @@ elif _is_hip:
awq_dequantize_triton as awq_dequantize,
)
elif _is_npu:
import custom_ops # noqa: F401
import sgl_kernel_npu # noqa: F401
import torch_npu # noqa: F401
from sglang.srt.hardware_backend.npu.modules.deepseek_v2_attention_mla_npu import (
forward_dsa_core_npu,
forward_dsa_prepare_npu,
forward_mha_core_npu,
forward_mha_prepare_npu,
forward_mla_core_npu,
forward_mla_prepare_npu,
)
from sglang.srt.layers.quantization.awq_triton import (
awq_dequantize_decomposition as awq_dequantize,
)
@@ -260,9 +260,6 @@ class AttnForwardMethod(IntEnum):
# Use absorbed multi-latent attention
MLA = auto()
# Use Deepseek V3.2 sparse multi-latent attention
NPU_MLA_SPARSE = auto()
# Use multi-head attention, but with KV cache chunked.
# This method can avoid OOM when prefix lengths are long.
MHA_CHUNKED_KV = auto()
@@ -277,6 +274,15 @@ class AttnForwardMethod(IntEnum):
# Use MLA with fused RoPE kernel for CPU
MLA_FUSED_ROPE_CPU = auto()
# Use multi-head attention for NPU
MHA_NPU = auto()
# Use absorbed multi-latent attention for NPU
MLA_NPU = auto()
# Use Deepseek V3.2 sparse multi-latent attention for NPU
DSA_NPU = auto()
def _dispatch_mla_subtype(attn, forward_batch):
if _is_hip:
@@ -311,14 +317,14 @@ def handle_attention_ascend(attn, forward_batch):
and not forward_batch.forward_mode.is_draft_extend_v2()
):
if hasattr(attn, "indexer"):
return AttnForwardMethod.NPU_MLA_SPARSE
return AttnForwardMethod.DSA_NPU
else:
return AttnForwardMethod.MHA
return AttnForwardMethod.MHA_NPU
else:
if hasattr(attn, "indexer"):
return AttnForwardMethod.NPU_MLA_SPARSE
return AttnForwardMethod.DSA_NPU
else:
return AttnForwardMethod.MLA
return AttnForwardMethod.MLA_NPU
def _get_sum_extend_prefix_lens(forward_batch):
@@ -1243,6 +1249,7 @@ class DeepseekV2AttentionMLA(nn.Module):
self.v_head_dim = v_head_dim
self.q_lora_rank = q_lora_rank
self.kv_lora_rank = kv_lora_rank
self.quant_config = quant_config
attn_tp_rank = get_attention_tp_rank()
attn_tp_size = get_attention_tp_size()
self.use_nsa = is_deepseek_nsa(config)
@@ -1470,12 +1477,6 @@ class DeepseekV2AttentionMLA(nn.Module):
self.weight_block_size = (
self.fused_qkv_a_proj_with_mqa.quant_method.quant_config.weight_block_size
)
self.is_mla_preprocess_enabled = is_mla_preprocess_enabled()
if self.is_mla_preprocess_enabled:
assert (
quant_config is None or quant_config.get_name() == "w8a8_int8"
), "MLA Preprocess only works with Unquant or W8A8Int8"
self.mla_preprocess = None
def dispatch_attn_forward_method(
self, forward_batch: ForwardBatch
@@ -1571,31 +1572,7 @@ class DeepseekV2AttentionMLA(nn.Module):
positions, hidden_states, forward_batch, zero_allocator
)
elif attn_forward_method == AttnForwardMethod.MLA:
if not self.is_mla_preprocess_enabled:
inner_state = self.forward_absorb_prepare(
positions, hidden_states, forward_batch, zero_allocator
)
else:
# TODO(iforgetmyname): to be separated as a standalone func
if self.mla_preprocess is None:
self.mla_preprocess = NPUFusedMLAPreprocess(
self.fused_qkv_a_proj_with_mqa,
self.q_a_layernorm,
self.kv_a_layernorm,
self.q_b_proj,
self.w_kc,
self.rotary_emb,
self.layer_id,
self.num_local_heads,
self.qk_nope_head_dim,
self.qk_rope_head_dim,
)
inner_state = self.mla_preprocess.forward(
positions, hidden_states, forward_batch, zero_allocator
)
inner_state = (*inner_state, None) # add a position for topk_indices
elif attn_forward_method == AttnForwardMethod.NPU_MLA_SPARSE:
inner_state = self.forward_npu_sparse_prepare(
inner_state = self.forward_absorb_prepare(
positions, hidden_states, forward_batch, zero_allocator
)
elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE:
@@ -1606,6 +1583,18 @@ class DeepseekV2AttentionMLA(nn.Module):
inner_state = self.forward_absorb_fused_mla_rope_cpu_prepare(
positions, hidden_states, forward_batch, zero_allocator
)
elif attn_forward_method == AttnForwardMethod.MHA_NPU:
inner_state = forward_mha_prepare_npu(
self, positions, hidden_states, forward_batch, zero_allocator
)
elif attn_forward_method == AttnForwardMethod.MLA_NPU:
inner_state = forward_mla_prepare_npu(
self, positions, hidden_states, forward_batch, zero_allocator
)
elif attn_forward_method == AttnForwardMethod.DSA_NPU:
inner_state = forward_dsa_prepare_npu(
self, positions, hidden_states, forward_batch, zero_allocator
)
else:
raise NotImplementedError
return None, attn_forward_method, forward_batch, inner_state
@@ -1625,12 +1614,16 @@ class DeepseekV2AttentionMLA(nn.Module):
return self.forward_normal_one_shot_core(*inner_state)
elif attn_forward_method == AttnForwardMethod.MLA:
return self.forward_absorb_core(*inner_state)
elif attn_forward_method == AttnForwardMethod.NPU_MLA_SPARSE:
return self.forward_npu_sparse_core(*inner_state)
elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE:
return self.forward_absorb_fused_mla_rope_core(*inner_state)
elif attn_forward_method == AttnForwardMethod.MLA_FUSED_ROPE_CPU:
return self.forward_absorb_fused_mla_rope_cpu_core(*inner_state)
elif attn_forward_method == AttnForwardMethod.MHA_NPU:
return forward_mha_core_npu(self, *inner_state)
elif attn_forward_method == AttnForwardMethod.MLA_NPU:
return forward_mla_core_npu(self, *inner_state)
elif attn_forward_method == AttnForwardMethod.DSA_NPU:
return forward_dsa_core_npu(self, *inner_state)
else:
raise NotImplementedError
@@ -2147,253 +2140,6 @@ class DeepseekV2AttentionMLA(nn.Module):
return output
def forward_npu_sparse_prepare(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
forward_batch: ForwardBatch,
zero_allocator: BumpAllocator,
):
"""
Reuse `self.q_lora_rank is not None` branch from forward_absorb_prepare
"""
if self.is_mla_preprocess_enabled and forward_batch.forward_mode.is_decode():
if self.mla_preprocess is None:
self.mla_preprocess = NPUFusedMLAPreprocess(
self.fused_qkv_a_proj_with_mqa,
self.q_a_layernorm,
self.kv_a_layernorm,
self.q_b_proj,
self.w_kc,
self.rotary_emb,
self.layer_id,
self.num_local_heads,
self.qk_nope_head_dim,
self.qk_rope_head_dim,
)
(
q_pe,
k_pe,
q_nope_out,
k_nope,
forward_batch,
zero_allocator,
positions,
) = self.mla_preprocess.forward(
positions, hidden_states, forward_batch, zero_allocator
)
fused_qkv_a_proj_out = self.fused_qkv_a_proj_with_mqa(hidden_states)[0]
q, _ = fused_qkv_a_proj_out.split(
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
)
q_lora = self.q_a_layernorm(q)
else:
from sglang.srt.model_executor.cuda_graph_runner import get_is_capture_mode
if (
(not isinstance(hidden_states, tuple))
and hidden_states.shape[0] <= 16
and self.use_min_latency_fused_a_gemm
):
fused_qkv_a_proj_out = dsv3_fused_a_gemm(
hidden_states, self.fused_qkv_a_proj_with_mqa.weight.T
)
else:
fused_qkv_a_proj_out = self.fused_qkv_a_proj_with_mqa(hidden_states)[0]
q, latent_cache = fused_qkv_a_proj_out.split(
[self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim], dim=-1
)
k_nope = latent_cache[..., : self.kv_lora_rank]
# overlap qk norm
if self.alt_stream is not None and get_is_capture_mode():
current_stream = torch.cuda.current_stream()
self.alt_stream.wait_stream(current_stream)
q = self.q_a_layernorm(q)
with torch.cuda.stream(self.alt_stream):
k_nope = self.kv_a_layernorm(k_nope)
current_stream.wait_stream(self.alt_stream)
else:
if _use_aiter_gfx95 and self.q_b_proj.weight.dtype == torch.uint8:
q, k_nope, *_ = fused_rms_mxfp4_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
k_nope,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
)
else:
if (
_use_aiter_gfx95
and self.q_b_proj.weight.dtype == torch.float8_e4m3fn
):
q, _, k_nope, _ = fused_rms_fp8_group_quant(
q,
self.q_a_layernorm.weight,
self.q_a_layernorm.variance_epsilon,
k_nope,
self.kv_a_layernorm.weight,
self.kv_a_layernorm.variance_epsilon,
group_size=128,
dtype_quant=torch.float8_e4m3fn,
res1=None,
output_unquantized_inp1=False,
)
else:
q = self.q_a_layernorm(q)
k_nope = self.kv_a_layernorm(k_nope)
q_lora = q.clone() # required for topk_indices
k_nope = k_nope.unsqueeze(1)
q = self.q_b_proj(q)[0].view(-1, self.num_local_heads, self.qk_head_dim)
q_nope, q_pe = q.split(
[self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
)
k_pe = latent_cache[..., self.kv_lora_rank :].unsqueeze(1)
if self.use_deep_gemm_bmm:
q_nope_val, q_nope_scale, masked_m, expected_m, aligned_m = (
per_token_group_quant_mla_deep_gemm_masked_fp8(
q_nope.transpose(0, 1)
)
)
q_nope_out = q_nope.new_empty(
(self.num_local_heads, aligned_m, self.kv_lora_rank)
)
deep_gemm_wrapper.grouped_gemm_nt_f8f8bf16_masked(
(q_nope_val, q_nope_scale),
(self.w_kc, self.w_scale_k),
q_nope_out,
masked_m,
expected_m,
)
q_nope_out = q_nope_out[:, :expected_m, :]
elif _is_hip:
# TODO(haishaw): add bmm_fp8 to ROCm
if _use_aiter_gfx95 and self.w_kc.dtype == torch.uint8:
x = q_nope.transpose(0, 1)
q_nope_out = torch.empty(
x.shape[0],
x.shape[1],
self.w_kc.shape[2],
device=x.device,
dtype=torch.bfloat16,
)
batched_gemm_afp4wfp4_pre_quant(
x,
self.w_kc.transpose(-2, -1),
self.w_scale_k.transpose(-2, -1),
torch.bfloat16,
q_nope_out,
)
else:
if _use_aiter_gfx95 and self.w_kc.dtype == torch.float8_e4m3fn:
q_nope_out = batched_gemm_a8w8_a_per_token_group_prequant_w_per_batched_tensor_quant(
X=q_nope,
WQ=self.w_kc.transpose(-1, -2),
w_scale=self.w_scale, #
group_size=128,
YQ=None, # allocate (B, M, N)
transpose_bm=False, # (B, M, N)
transpose_bm_in=True, # (M, B, K)
dtype=torch.bfloat16,
)
else:
q_nope_out = torch.bmm(
q_nope.to(torch.bfloat16).transpose(0, 1),
self.w_kc.to(torch.bfloat16) * self.w_scale,
)
elif self.w_kc.dtype == torch.float8_e4m3fn:
q_nope_val, q_nope_scale = per_tensor_quant_mla_fp8(
q_nope.transpose(0, 1),
zero_allocator.allocate(1),
)
q_nope_out = bmm_fp8(
q_nope_val, self.w_kc, q_nope_scale, self.w_scale, torch.bfloat16
)
else:
q_nope_out = torch.bmm(q_nope.transpose(0, 1), self.w_kc)
q_nope_out = q_nope_out.transpose(0, 1)
if not self._fuse_rope_for_trtllm_mla(forward_batch) and (
not _use_aiter or not _is_gfx95_supported
):
q_pe, k_pe = self.rotary_emb(positions, q_pe, k_pe)
# TODO: multi-stream indexer
topk_indices = self.indexer(
hidden_states, q_lora, positions, forward_batch, self.layer_id
)
return (
q_pe,
k_pe,
q_nope_out,
k_nope,
topk_indices,
forward_batch,
zero_allocator,
positions,
)
def forward_npu_sparse_core(
self,
q_pe,
k_pe,
q_nope_out,
k_nope,
topk_indices,
forward_batch,
zero_allocator,
positions,
):
attn_output = self.attn_mqa(
q_nope_out.contiguous(),
k_nope.contiguous(),
k_nope.contiguous(),
forward_batch,
save_kv_cache=True, # False if forward_batch.forward_mode.is_extend() else True,
q_rope=q_pe.contiguous(),
k_rope=k_pe.contiguous(),
topk_indices=topk_indices,
)
attn_output = attn_output.view(-1, self.num_local_heads, self.kv_lora_rank)
attn_bmm_output = torch.empty(
(attn_output.shape[0], self.num_local_heads, self.v_head_dim),
dtype=attn_output.dtype,
device=attn_output.device,
)
if not forward_batch.forward_mode.is_decode():
attn_output = attn_output.transpose(0, 1)
torch.bmm(
attn_output,
self.w_vc,
out=attn_bmm_output.view(
-1, self.num_local_heads, self.v_head_dim
).transpose(0, 1),
)
else:
attn_output = attn_output.contiguous()
torch.ops.npu.batch_matmul_transpose(
attn_output, self.w_vc, attn_bmm_output
)
attn_bmm_output = attn_bmm_output.reshape(
-1, self.num_local_heads * self.v_head_dim
)
output, _ = self.o_proj(attn_bmm_output)
return output
def forward_absorb_fused_mla_rope_prepare(
self,
positions: torch.Tensor,
@@ -4126,12 +3872,8 @@ class DeepseekV2ForCausalLM(nn.Module):
del self.lm_head.weight
self.model.embed_tokens.weight = embed
self.lm_head.weight = head
if not _is_npu:
torch.cuda.empty_cache()
torch.cuda.synchronize()
else:
torch.npu.empty_cache()
torch.npu.synchronize()
torch.cuda.empty_cache()
torch.cuda.synchronize()
@classmethod
def get_model_config_for_expert_location(cls, config):
python/sglang/srt/models/qwen3.py
+3 -7
View File
@@ -30,13 +30,7 @@ from sglang.srt.model_loader.weight_utils import (
from sglang.srt.models.qwen2 import Qwen2MLP as Qwen3MLP
from sglang.srt.models.qwen2 import Qwen2Model
from sglang.srt.server_args import get_global_server_args
from sglang.srt.utils import (
add_prefix,
get_cmo_stream,
is_cuda,
is_npu,
wait_cmo_stream,
)
from sglang.srt.utils import add_prefix, is_cuda, is_npu
Qwen3Config = None
@@ -47,6 +41,8 @@ _is_npu = is_npu()
if _is_npu:
from sgl_kernel_npu.norm.split_qkv_rmsnorm_rope import split_qkv_rmsnorm_rope
from sglang.srt.hardware_backend.npu.cmo import get_cmo_stream, wait_cmo_stream
class Qwen3Attention(nn.Module):
def __init__(
python/sglang/srt/server_args.py
+9 -26
View File
@@ -103,6 +103,7 @@ QUANTIZATION_CHOICES = [
"mxfp4",
"auto-round",
"compressed-tensors", # for Ktransformers
"modelslim", # for NPU
]
ATTENTION_BACKEND_CHOICES = [
@@ -615,6 +616,7 @@ class ServerArgs:
# Handle device-specific backends.
self._handle_hpu_backends()
self._handle_cpu_backends()
self._handle_npu_backends()
# Apply model-specific adjustments.
self._handle_model_specific_adjustments()
@@ -928,6 +930,12 @@ class ServerArgs:
self.attention_backend = "intel_amx"
self.sampling_backend = "pytorch"
def _handle_npu_backends(self):
if self.device == "npu":
from sglang.srt.hardware_backend.npu.utils import set_default_server_args
set_default_server_args(self)
def _handle_model_specific_adjustments(self):
from sglang.srt.configs.model_config import is_deepseek_nsa
@@ -1300,8 +1308,6 @@ class ServerArgs:
self.attention_backend = "fa3"
elif is_hip():
self.attention_backend = "aiter"
elif is_npu():
self.attention_backend = "ascend"
else:
self.attention_backend = (
"flashinfer" if is_flashinfer_available() else "triton"
@@ -1319,8 +1325,6 @@ class ServerArgs:
self.attention_backend = "aiter"
else:
self.attention_backend = "triton"
elif is_npu():
self.attention_backend = "ascend"
else:
self.attention_backend = "triton"
@@ -1421,13 +1425,6 @@ class ServerArgs:
if model_config.context_len > 8192:
self.mem_fraction_static *= 0.85
# NPU platforms backends
if is_npu() and self.attention_backend in ["ascend"]:
logger.warning(
"At this moment Ascend attention backend only supports a page_size of 128, change page_size to 128."
)
self.page_size = 128
# Other platforms backends
if (
self.attention_backend == "intel_amx"
@@ -1620,7 +1617,7 @@ class ServerArgs:
)
if self.hicache_mem_layout == "page_first_direct":
if self.hicache_io_backend != "direct":
if self.hicache_io_backend not in ["direct", "kernel_ascend"]:
self.hicache_io_backend = "direct"
logger.warning(
"Page first direct layout only support direct io backend"
@@ -1644,20 +1641,6 @@ class ServerArgs:
"Setting hicache_io_backend to vanilla I/O, which may lead to suboptimal performance with small page sizes."
)
# Below are the only parameters currently supported on Ascend
if self.enable_hierarchical_cache and is_npu():
# FIXME(iforgetmyname) fix decode_attention_backend on ascend
self.decode_attention_backend = "ascend"
self.hicache_io_backend = "kernel_ascend"
if self.use_mla_backend():
self.hicache_mem_layout = "page_first_kv_split"
else:
self.hicache_mem_layout = "page_first_direct"
logger.warning(
f"Ascend NPU Platform detected, change `hicache_io_backend` to `kernel_ascend` and "
f"`hicache_mem_layout` to `{self.hicache_mem_layout}`"
)
def _handle_speculative_decoding(self):
if (
self.speculative_draft_model_path is not None
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_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,
python/sglang/srt/utils/common.py
-54
View File
@@ -645,60 +645,6 @@ def make_layers_non_pp(
return layers
cmo_stream = None
def get_cmo_stream():
"""
Cache Management Operation(CMO).
Launch a new stream to prefetch the weight of matmul when running other
AIV or communication kernels, aiming to overlap the memory access time.
"""
global cmo_stream
return cmo_stream
def set_cmo_stream(stream):
global cmo_stream
cmo_stream = stream
def prepare_weight_cache(handle, cache, PREFETCH_MAX_SIZE=1000000000):
"""
PREFETCH_MAX_SIZE: maximum size (bytes) for each prefetch operation.
This affects the time spent in prefetch:
time ≈ PREFETCH_MAX_SIZE / system_bandwidth
"""
import torch_npu
stream = get_cmo_stream()
if stream is None:
stream = torch.get_device_module().Stream()
set_cmo_stream(stream)
stream.wait_stream(torch.get_device_module().current_stream())
with torch.get_device_module().stream(stream):
if isinstance(cache, list):
for weight in cache:
torch_npu.npu_prefetch(
weight,
handle,
PREFETCH_MAX_SIZE,
)
else:
torch_npu.npu_prefetch(
cache,
handle,
PREFETCH_MAX_SIZE,
)
def wait_cmo_stream():
stream = get_cmo_stream()
if stream is not None:
cur_stream = torch.get_device_module().current_stream()
cur_stream.wait_stream(stream)
@lru_cache(maxsize=1)
def get_device_module():
return torch.get_device_module()
python/sglang/srt/utils/numa_utils.py
+1 -1
View File
@@ -6,8 +6,8 @@ import time
from contextlib import contextmanager
from pathlib import Path
from sglang import ServerArgs
from sglang.srt.environ import envs
from sglang.srt.server_args import ServerArgs
logger = logging.getLogger(__name__)