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

[NPU] perf update with kvcache nz & w4a8 quant (#14423)

Browse Source
This commit is contained in:
liupeng374
2025-12-13 17:39:55 +08:00
committed by GitHub
parent 0e7d7969d5
commit d36299ad77
6 changed files with 219 additions and 108 deletions
Download Patch File Download Diff File Expand all files Collapse all files

78
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
View File

@@ -8,6 +8,7 @@ import torch_npu
from sglang.srt.configs.model_config import AttentionArch
from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
is_fia_nz,
is_mla_preprocess_enabled,
)
from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
@@ -27,6 +28,14 @@ import logging
import numpy as np
def _reshape_kv_for_fia_nz(
tensor: torch.Tensor, num_heads: int, head_dim: int, page_size: int
) -> torch.Tensor:
"""Reshapes a tensor for FIA NZ format."""
return tensor.view(-1, 1, num_heads * head_dim // 16, page_size, 16)
logger = logging.getLogger(__name__)
@@ -881,12 +890,20 @@ class AscendAttnBackend(AttentionBackend):
)
c_kv, k_rope = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
k_rope_cache = k_rope.view(
-1, layer.tp_k_head_num, self.page_size, self.qk_rope_head_dim
)
c_kv_cache = c_kv.view(
-1, layer.tp_v_head_num, self.page_size, self.kv_lora_rank
)
if is_fia_nz():
k_rope_cache = _reshape_kv_for_fia_nz(
k_rope, layer.tp_k_head_num, self.qk_rope_head_dim, self.page_size
)
c_kv_cache = _reshape_kv_for_fia_nz(
c_kv, layer.tp_v_head_num, self.kv_lora_rank, self.page_size
)
else:
k_rope_cache = k_rope.view(
-1, layer.tp_k_head_num, self.page_size, self.qk_rope_head_dim
)
c_kv_cache = c_kv.view(
-1, layer.tp_v_head_num, self.page_size, self.kv_lora_rank
)
q_nope = q.view(-1, layer.tp_q_head_num, self.kv_lora_rank).contiguous()
q_rope = q_rope.view(-1, layer.tp_q_head_num, self.qk_rope_head_dim)
@@ -1082,15 +1099,24 @@ class AscendAttnBackend(AttentionBackend):
)
else:
c_kv, k_rope = forward_batch.token_to_kv_pool.get_kv_buffer(layer.layer_id)
k_rope_cache = k_rope.view(
-1, layer.tp_k_head_num, self.page_size, self.qk_rope_head_dim
)
c_kv_cache = c_kv.view(
-1, layer.tp_v_head_num, self.page_size, self.kv_lora_rank
)
if is_fia_nz():
k_rope_cache = _reshape_kv_for_fia_nz(
k_rope, layer.tp_k_head_num, self.qk_rope_head_dim, self.page_size
)
c_kv_cache = _reshape_kv_for_fia_nz(
c_kv, layer.tp_v_head_num, self.kv_lora_rank, self.page_size
)
else:
k_rope_cache = k_rope.view(
-1, self.page_size, layer.tp_k_head_num * self.qk_rope_head_dim
)
c_kv_cache = c_kv.view(
-1, self.page_size, layer.tp_k_head_num * self.kv_lora_rank
)
q_nope = q.view(-1, 1, layer.tp_q_head_num, self.kv_lora_rank).contiguous()
q_rope = q_rope.view(-1, 1, layer.tp_q_head_num, self.qk_rope_head_dim)
q_nope = q.view(-1, layer.tp_q_head_num, 1, self.kv_lora_rank).contiguous()
q_rope = q_rope.view(-1, layer.tp_q_head_num, 1, self.qk_rope_head_dim)
if self.forward_metadata.seq_lens_cpu_int is None:
actual_seq_len_kv = self.forward_metadata.seq_lens_cpu_list
else:
@@ -1108,7 +1134,7 @@ class AscendAttnBackend(AttentionBackend):
num_key_value_heads=layer.tp_k_head_num,
block_table=self.forward_metadata.block_tables,
block_size=self.page_size,
input_layout="BNSD",
input_layout="BSND",
scale=layer.scaling,
actual_seq_lengths_kv=actual_seq_len_kv,
antiquant_mode=0,
@@ -1128,7 +1154,7 @@ class AscendAttnBackend(AttentionBackend):
num_key_value_heads=layer.tp_k_head_num,
block_table=self.forward_metadata.block_tables,
block_size=self.page_size,
input_layout="BNSD",
input_layout="BSND",
scale=layer.scaling,
actual_seq_lengths_kv=actual_seq_len_kv,
antiquant_mode=0,
@@ -1243,12 +1269,20 @@ class AscendAttnBackend(AttentionBackend):
if self.use_fia and (layer.tp_q_head_num // layer.tp_k_head_num) >= 8:
"""layer.tp_q_head_num // layer.tp_k_head_num < 8 will support in the later version of CANN"""
kv_c = kv_c.view(
-1, self.page_size, layer.tp_k_head_num * self.kv_lora_rank
)
k_pe = k_pe.view(
-1, self.page_size, layer.tp_k_head_num * self.qk_rope_head_dim
)
if is_fia_nz():
kv_c = _reshape_kv_for_fia_nz(
kv_c, layer.tp_k_head_num, self.kv_lora_rank, self.page_size
)
k_pe = _reshape_kv_for_fia_nz(
k_pe, layer.tp_k_head_num, self.qk_rope_head_dim, self.page_size
)
else:
kv_c = kv_c.view(
-1, self.page_size, layer.tp_k_head_num * self.kv_lora_rank
)
k_pe = k_pe.view(
-1, self.page_size, layer.tp_k_head_num * self.qk_rope_head_dim
)
q = q.view(
forward_batch.batch_size, -1, layer.tp_q_head_num, self.kv_lora_rank
)

38
python/sglang/srt/hardware_backend/npu/attention/mla_preprocess.py
View File

@@ -16,6 +16,16 @@ def is_mla_preprocess_enabled() -> bool:
return get_bool_env_var("SGLANG_NPU_USE_MLAPO")
@lru_cache(maxsize=1)
def is_fia_nz() -> bool:
is_fia_nz_ = get_bool_env_var("SGLANG_USE_FIA_NZ")
if is_fia_nz_:
assert (
is_mla_preprocess_enabled()
), "SGLANG_USE_FIA_NZ must be enable with SGLANG_NPU_USE_MLAPO"
return is_fia_nz_
def round_up(val: int, align: int) -> int:
if align == 0:
return 0
@@ -84,7 +94,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
def preprocess_weights(self, hidden_states):
self.dummy = torch.empty(
self.dummy = torch.zeros(
(hidden_states.shape[-1]),
dtype=hidden_states.dtype,
device=hidden_states.device,
@@ -281,7 +291,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
-1, 1, 1, self.kv_lora_rank + self.qk_rope_head_dim
) # (B*S,N,1,D)
cache_mode = "PA_BNSD"
cache_mode = "PA_NZ" if is_fia_nz() else "PA_BNSD"
self.kvCache = self.kvCache.view(
-1,
forward_batch.attn_backend.page_size,
@@ -328,9 +338,22 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
dtype=input_dtype,
device=hidden_states.device,
)
if is_fia_nz():
kv_shape, kv_rope_shape = k_cache.shape, v_cache.shape
num_blocks, block_size, num_heads, _ = kv_shape
k_cache = k_cache.view(
num_blocks, num_heads * self.kv_lora_rank // 16, block_size, 16
)
v_cache = v_cache.view(
num_blocks, num_heads * self.qk_rope_head_dim // 16, block_size, 16
)
# TODO: dummy inputs to be removed
# https://github.com/sgl-project/sgl-kernel-npu/issues/78
if hasattr(self.q_a_layernorm, "bias"):
q_a_layernorm_bias = self.q_a_layernorm.bias
else:
q_a_layernorm_bias = self.dummy
torch.ops.npu.mla_preprocess(
hidden_states,
self.dummy,
@@ -338,7 +361,7 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
self.qkv_a_proj_weight_nz,
self.qkv_a_proj_deq_scale_kvq,
self.q_a_layernorm.weight,
self.q_a_layernorm.bias,
q_a_layernorm_bias,
self.q_b_proj_weight_nz,
self.q_b_proj_deq_scale,
self.kv_a_layernorm.weight,
@@ -354,13 +377,18 @@ class NPUFusedMLAPreprocess(torch.nn.Module):
quant_scale1=self.q_b_proj.input_scale,
quant_offset1=self.q_b_proj_input_offset,
bias1=self.q_b_proj_quant_bias,
cache_mode="krope_ctkv",
cache_mode="nzcache" if is_fia_nz() else "krope_ctkv",
quant_mode="per_tensor_quant_asymm",
q_out0=q_nope_out,
kv_cache_out0=k_cache,
q_out1=q_rope_out,
kv_cache_out1=v_cache,
)
if is_fia_nz():
k_cache = k_cache.view(kv_shape)
v_cache = v_cache.view(kv_rope_shape)
return (
q_rope_out,
v_cache,

61
python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py
View File

@@ -1,9 +1,11 @@
from typing import TYPE_CHECKING
import torch
import torch_npu
from sglang.srt.hardware_backend.npu.attention.mla_preprocess import (
NPUFusedMLAPreprocess,
is_fia_nz,
is_mla_preprocess_enabled,
)
from sglang.srt.layers.attention.nsa.utils import (
@@ -62,26 +64,39 @@ def forward_mha_prepare_npu(
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]
B, S = q.shape[0], 1
cos, sin = m.rotary_emb.get_cos_sin_cache(
positions, hidden_states.dtype, offsets=None
)
q_pe = torch_npu.npu_interleave_rope(
q_pe.reshape(B, -1, S, m.qk_rope_head_dim),
cos,
sin,
)
q_pe = q_pe.reshape(B, -1, m.qk_rope_head_dim)
ckv_cache, k_rope_cache = forward_batch.token_to_kv_pool.get_kv_buffer(m.layer_id)
_, _, k_pe, kv_a = torch_npu.npu_kv_rmsnorm_rope_cache(
latent_cache.view(-1, 1, 1, m.kv_lora_rank + m.qk_rope_head_dim), # bnsd
m.kv_a_layernorm.weight,
cos,
sin,
forward_batch.out_cache_loc.to(torch.int64),
k_rope_cache,
ckv_cache,
k_rope_scale=None,
c_kv_scale=None,
k_rope_offset=None,
c_kv_offset=None,
epsilon=m.kv_a_layernorm.variance_epsilon,
cache_mode="PA_NZ" if is_fia_nz() else "PA_BNSD",
is_output_kv=True,
) # adapter NZ
k_pe = k_pe.reshape(B, -1, m.qk_rope_head_dim)
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]
@@ -233,15 +248,15 @@ def forward_mla_core_npu(
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),
(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),
)
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

99
python/sglang/srt/hardware_backend/npu/quantization/fused_moe_method_npu.py
View File

@@ -324,7 +324,7 @@ class NPUW8A8Int8DynamicMoEMethod(FusedMoEMethodBase):
class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase):
def __init__(self) -> None:
self.group_size = 256
self.group_size = 0
self.tp_size = 1
def create_weights(
@@ -338,6 +338,7 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase):
) -> None:
from sglang.srt.layers.moe.fused_moe_triton import FusedMoeWeightScaleSupported
self.is_per_channel_weight = self.group_size == 0
self.num_experts = num_experts
extra_weight_attrs.update(
{"quant_method": FusedMoeWeightScaleSupported.CHANNEL.value}
@@ -399,52 +400,55 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase):
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)
if not self.is_per_channel_weight:
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_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)
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(
@@ -466,6 +470,11 @@ class NPUW4A8Int4DynamicMoEMethod(FusedMoEMethodBase):
def process_scale(self, weight: torch.Tensor, scale, per_group_scale):
scale = scale.transpose(1, 2).contiguous()
if self.is_per_channel_weight:
scale_np = scale.cpu().numpy()
scale_np.dtype = np.uint32
scale_uint64_tensor = torch.from_numpy(scale_np.astype(np.int64)).npu()
return scale_uint64_tensor, None
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

46
python/sglang/srt/layers/rotary_embedding.py
View File

@@ -806,6 +806,10 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
/ yarn_get_mscale(self.scaling_factor, float(mscale_all_dim))
* attn_factor
)
self.cos_cached_total = None
self.sin_cached_total = None
self.cos_cached = None
self.sin_cached = None
self.device = device
super().__init__(
head_size, rotary_dim, max_position_embeddings, base, is_neox_style, dtype
@@ -854,8 +858,41 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
cos = freqs.cos() * self.mscale
sin = freqs.sin() * self.mscale
cache = torch.cat((cos, sin), dim=-1)
emb = torch.cat((freqs, freqs), dim=-1)
self.cos_cached_total = torch.cos(emb) * self.mscale
self.sin_cached_total = torch.sin(emb) * self.mscale
return cache
def get_cos_cached_total(self):
return self.cos_cached_total
def get_sin_cached_total(self):
return self.sin_cached_total
def get_cos_sin_cache(
self, positions, dtype, offsets: Optional[torch.Tensor] = None
):
self.cos_cached = (
self.cos_cached_total[
torch.add(positions, offsets) if offsets is not None else positions
]
.unsqueeze(-2)
.unsqueeze(-2)
.to(dtype)
)
self.sin_cached = (
self.sin_cached_total[
torch.add(positions, offsets) if offsets is not None else positions
]
.unsqueeze(-2)
.unsqueeze(-2)
.to(dtype)
)
cos = self.cos_cached.to(positions.device)
sin = self.sin_cached.to(positions.device)
return cos, sin
def forward_native(
self,
positions: torch.Tensor,
@@ -906,14 +943,7 @@ class DeepseekScalingRotaryEmbedding(RotaryEmbedding):
num_tokens, num_q_heads, _ = query.shape
num_k_heads = key.shape[1]
self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(positions.device)
cos_sin = self.cos_sin_cache[
torch.add(positions, offsets) if offsets is not None else positions
]
cos, sin = cos_sin.chunk(2, dim=-1)
# Reshape to [batchsize, head_dim, seq, rotary_dim]
cos = cos.repeat(1, 2).unsqueeze(-2).unsqueeze(-2)
sin = sin.repeat(1, 2).unsqueeze(-2).unsqueeze(-2)
cos, sin = self.get_cos_sin_cache(positions, query.dtype, offsets)
query_rot = query[..., : self.rotary_dim]
key_rot = key[..., : self.rotary_dim]

5
python/sglang/srt/model_executor/forward_batch_info.py
View File

@@ -541,11 +541,6 @@ class ForwardBatch:
num_tokens_per_dp=num_tokens_per_dp,
)
self.num_token_non_padded_cpu = compute_local_num_token_non_padded(
global_num_token_non_padded=self.num_token_non_padded_cpu,
num_tokens_per_dp=num_tokens_per_dp,
)
def merge_mm_inputs(self) -> Optional[MultimodalInputs]:
"""
Merge all multimodal inputs in the batch into a single MultiModalInputs object.