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

[NPU] optimization for dsv3.2 (#14572)

Browse Source
This commit is contained in:
ZhengdQin
2025-12-12 14:52:16 +08:00
committed by GitHub
parent edb172e927
commit c05d3afb5d
11 changed files with 140 additions and 67 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -573,8 +573,12 @@ class AscendAttnBackend(AttentionBackend):
key_rope=k_pe,
sparse_indices=topk_indices,
scale_value=layer.scaling,
actual_seq_lengths_query=actual_seq_qlen,
actual_seq_lengths_kv=actual_seq_lengths_kv.to(q.device),
actual_seq_lengths_query=actual_seq_qlen.to(
device=q_nope.device, dtype=torch.int32
),
actual_seq_lengths_kv=actual_seq_lengths_kv.to(
device=q_nope.device, dtype=torch.int32
),
block_table=self.forward_metadata.block_tables,
sparse_block_size=1,
layout_query="TND",

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

@@ -273,39 +273,51 @@ def forward_dsa_prepare_npu(
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
)
mla_event = torch.npu.Event()
mla_event.record()
with torch.npu.stream(m.alt_stream):
torch.npu.current_stream().wait_event(mla_event)
(
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)
torch.npu.current_stream().wait_stream(m.alt_stream)
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]
# overlap qk norm
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)
m.alt_stream.wait_stream(torch.npu.current_stream())
with torch.npu.stream(m.alt_stream):
q = m.q_b_proj(q_lora)[0].view(-1, m.num_local_heads, m.qk_head_dim)
q.record_stream(m.alt_stream)
q_event = m.alt_stream.record_event()
k_nope, k_pe = latent_cache.unsqueeze(1).split(
[m.kv_lora_rank, m.qk_rope_head_dim], dim=-1
)
k_nope = m.kv_a_layernorm(k_nope).unsqueeze(1)
torch.npu.current_stream().wait_event(q_event)
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)
@@ -367,7 +379,11 @@ def forward_dsa_core_npu(
device=attn_output.device,
)
if not forward_batch.forward_mode.is_decode():
if (
forward_batch.forward_mode.is_extend()
and not forward_batch.forward_mode.is_draft_extend(include_v2=True)
and not forward_batch.forward_mode.is_target_verify()
):
attn_output = attn_output.transpose(0, 1)
torch.bmm(
attn_output,

15
python/sglang/srt/hardware_backend/npu/moe/topk.py
View File

@@ -39,8 +39,7 @@ def fused_topk_npu(
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
# Force set routed_scaling_factor = 1 to optimize renormalize
topk_weights, topk_ids, _ = torch.ops.npu.npu_moe_gating_top_k(
router_logits.to(torch.float32),
k=topk_config.top_k,
@@ -50,18 +49,12 @@ def fused_topk_npu(
group_select_mode=1,
renorm=0,
norm_type=1,
routed_scaling_factor=routed_scaling_factor,
routed_scaling_factor=(
1 if renormalize else topk_config.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(

8
python/sglang/srt/hardware_backend/npu/quantization/modelslim.py
View File

@@ -159,8 +159,10 @@ class ModelSlimConfig(QuantizationConfig):
proj_name, packed_modules_mapping_subset[proj_name][0]
)
self.is_dynamic = (
self.quant_description[prefix_in_quant_config + ".weight"]
self.quant_description.get(prefix_in_quant_config + ".weight", "")
== "W8A8_DYNAMIC"
or self.quant_description.get("quant_method", "")
== "modelslim" # TODO: This path is for compress-tensor config，needs refactor @zhengdqin
)
if self.is_layer_skipped(prefix, packed_modules_mapping_subset):
return UnquantizedLinearMethod()
@@ -199,7 +201,7 @@ class ModelSlimConfig(QuantizationConfig):
is_skipped = None
for shard_prefix in shard_prefixes:
is_shard_skipped = (
self.quant_description[shard_prefix + ".weight"] == "FLOAT"
self.quant_description.get(shard_prefix + ".weight", "") == "FLOAT"
)
if is_skipped is None:
@@ -211,7 +213,7 @@ class ModelSlimConfig(QuantizationConfig):
"to have the same precision."
)
else:
is_skipped = self.quant_description[prefix + ".weight"] == "FLOAT"
is_skipped = self.quant_description.get(prefix + ".weight", "") == "FLOAT"
assert is_skipped is not None
return is_skipped

8
python/sglang/srt/hardware_backend/npu/utils.py
View File

@@ -13,6 +13,7 @@ if TYPE_CHECKING:
logger = logging.getLogger(__name__)
_is_npu = is_npu()
indexer_weight_stream = None
class NPUACLFormat(IntEnum):
@@ -110,3 +111,10 @@ def npu_format_cast(
import torch_npu
return torch_npu.npu_format_cast(tensor, acl_format.value)
def get_indexer_weight_stream():
global indexer_weight_stream
if indexer_weight_stream is None:
indexer_weight_stream = torch.npu.Stream()
return indexer_weight_stream

64
python/sglang/srt/layers/attention/nsa/nsa_indexer.py
View File

@@ -10,12 +10,18 @@ from sglang.srt.custom_op import CustomOp
from sglang.srt.layers.layernorm import LayerNorm
from sglang.srt.utils import add_prefix, ceil_align, is_cuda, is_hip, is_npu
global _use_multi_stream
if is_cuda():
try:
import deep_gemm
except ImportError as e:
deep_gemm = e
if is_npu():
import custom_ops # noqa: F401
import torch_npu
from sglang.srt.hardware_backend.npu.utils import get_indexer_weight_stream
from sglang.srt.layers import deep_gemm_wrapper
from sglang.srt.layers.attention.nsa.utils import (
@@ -980,19 +986,47 @@ class Indexer(CustomOp):
sin = sin.repeat(1, 2).view(-1, 1, 1, self.rope_head_dim)
bs = x.shape[0]
q = self.wq_b(q_lora)[0] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128]
q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128]
q_pe, q_nope = torch.split(
q,
[self.rope_head_dim, self.head_dim - self.rope_head_dim],
dim=-1,
) # [bs, 64, 64 + 64]
if self.alt_stream is not None:
self.alt_stream.wait_stream(torch.npu.current_stream())
with torch.npu.stream(self.alt_stream):
q = self.wq_b(q_lora)[
0
] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128]
wq_b_event = self.alt_stream.record_event()
q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128]
q_pe, q_nope = torch.split(
q,
[self.rope_head_dim, self.head_dim - self.rope_head_dim],
dim=-1,
) # [bs, 64, 64 + 64]
q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim)
q_pe = torch_npu.npu_rotary_mul(q_pe, cos, sin).view(
bs, self.n_heads, self.rope_head_dim
) # [bs, n, d]
q = torch.cat([q_pe, q_nope], dim=-1)
q.record_stream(self.alt_stream)
q_rope_event = self.alt_stream.record_event()
else:
q = self.wq_b(q_lora)[0] # [bs, 1536] @ [1536, 64 * 128] = [bs, 64 * 128]
q = q.view(bs, self.n_heads, self.head_dim) # [bs, 64, 128]
q_pe, q_nope = torch.split(
q,
[self.rope_head_dim, self.head_dim - self.rope_head_dim],
dim=-1,
) # [bs, 64, 64 + 64]
q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim)
q_pe = torch_npu.npu_rotary_mul(q_pe, cos, sin).view(
bs, self.n_heads, self.rope_head_dim
) # [bs, n, d]
q = torch.cat([q_pe, q_nope], dim=-1)
q_pe = q_pe.view(bs, self.n_heads, 1, self.rope_head_dim)
q_pe = torch.ops.npu.npu_rotary_mul(q_pe, cos, sin).view(
bs, self.n_heads, self.rope_head_dim
) # [bs, n, d]
q = torch.cat([q_pe, q_nope], dim=-1)
indexer_weight_stream = get_indexer_weight_stream()
indexer_weight_stream.wait_stream(torch.npu.current_stream())
with torch.npu.stream(indexer_weight_stream):
x = x.view(-1, self.hidden_size)
weights = self.weights_proj(x.float())[0].to(torch.bfloat16)
weights.record_stream(indexer_weight_stream)
weights_event = indexer_weight_stream.record_event()
k_proj = self.wk(x)[0] # [b, s, 7168] @ [7168, 128] = [b, s, 128]
k = self.k_norm(k_proj)
@@ -1072,8 +1106,10 @@ class Indexer(CustomOp):
past_key_states = forward_batch.token_to_kv_pool.get_index_k_buffer(layer_id)
x = x.view(-1, self.hidden_size)
weights = self.weights_proj(x.float())[0].to(torch.bfloat16)
if self.alt_stream is not None:
torch.npu.current_stream().wait_event(q_rope_event)
torch.npu.current_stream().wait_event(weights_event)
block_table = forward_batch.attn_backend.forward_metadata.block_tables
if (
is_prefill

14
python/sglang/srt/layers/layernorm.py
View File

@@ -343,19 +343,7 @@ class LayerNorm(CustomOp):
self,
x: torch.Tensor,
) -> torch.Tensor:
orig_dtype = x.dtype
x = x.to(self.dtype)
mean = x.mean(dim=-1, keepdim=True)
variance = (x - mean).pow(2).mean(dim=-1, keepdim=True)
x = (x - mean) * torch.rsqrt(variance + self.variance_epsilon)
if self.elementwise_affine:
x = x * self.weight.to(self.dtype)
if self.use_bias:
x = x + self.bias.to(self.dtype)
return x.to(orig_dtype)
return self.forward_native(x)
def forward_cpu(
self,

5
python/sglang/srt/layers/linear.py
View File

@@ -245,6 +245,11 @@ class ReplicatedLinear(LinearBase):
else:
raise ValueError(f"{loaded_weight} are not all equal")
if param.dtype == torch.int8 or loaded_weight.dtype == torch.int8:
assert (
param.dtype == loaded_weight.dtype
), "init para dtype and loaded weight dtype should be the same"
assert param.size() == loaded_weight.size()
param.data.copy_(loaded_weight)

2
python/sglang/srt/layers/quantization/compressed_tensors/utils.py
View File

@@ -79,7 +79,7 @@ def check_equal_or_regex_match(layer_name: str, targets: Iterable[str]) -> bool:
if target starts with 're:' to any target in list.
"""
for target in targets:
if _is_equal_or_regex_match(layer_name, target):
if _is_equal_or_regex_match(layer_name, target, check_contains=True):
return True
return False

2
python/sglang/srt/models/deepseek_nextn.py
View File

@@ -97,7 +97,7 @@ class DeepseekModelNextN(nn.Module):
self.eh_proj = nn.Linear(2 * config.hidden_size, config.hidden_size, bias=False)
self.alt_stream = torch.cuda.Stream() if _is_cuda else None
self.alt_stream = torch.cuda.Stream() if _is_cuda or _is_npu else None
layer_name = "decoder"
if _is_npu and (

29
python/sglang/srt/models/deepseek_v2.py
View File

@@ -700,6 +700,7 @@ class DeepseekV2MoE(nn.Module):
dict(tp_rank=0, tp_size=1)
if get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_ascend_fuseep()
or should_use_flashinfer_cutlass_moe_fp4_allgather()
else {}
),
@@ -738,7 +739,11 @@ class DeepseekV2MoE(nn.Module):
self.top_k = config.num_experts_per_tok
if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
if (
get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_ascend_fuseep()
):
# TODO: we will support tp < ep in the future
self.ep_size = get_moe_expert_parallel_world_size()
self.num_experts = (
@@ -755,7 +760,9 @@ class DeepseekV2MoE(nn.Module):
)
self._enable_a2a_moe = (
get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake()
get_moe_a2a_backend().is_deepep()
or get_moe_a2a_backend().is_mooncake()
or get_moe_a2a_backend().is_ascend_fuseep()
)
self._fuse_shared_experts_inside_sbo = SboFlags.fuse_shared_experts_inside_sbo()
@@ -986,7 +993,14 @@ class DeepseekV2MoE(nn.Module):
# router_logits: (num_tokens, n_experts)
router_logits = self.gate(hidden_states, forward_batch=forward_batch)
if not sbo_enabled_flag:
shared_output = self._forward_shared_experts(hidden_states)
if self.alt_stream is not None:
self.alt_stream.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(self.alt_stream):
shared_output = self._forward_shared_experts(hidden_states)
shared_output.record_stream(self.alt_stream)
shared_event = self.alt_stream.record_event()
else:
shared_output = self._forward_shared_experts(hidden_states)
topk_output = self.topk(
hidden_states,
router_logits,
@@ -1105,6 +1119,12 @@ class DeepseekV2MoE(nn.Module):
topk_output=topk_output,
)
if (
hidden_states.shape[0] > 0
and not sbo_enabled_flag
and self.alt_stream is not None
):
torch.cuda.current_stream().wait_event(shared_event)
if shared_output is not None:
x = shared_output
if self.experts.should_fuse_routed_scaling_factor_in_topk:
@@ -2991,7 +3011,8 @@ class DeepseekV2Model(nn.Module):
else:
self.embed_tokens = PPMissingLayer()
self.alt_stream = torch.cuda.Stream() if _is_cuda else None
self.alt_stream = torch.cuda.Stream() if _is_cuda or _is_npu else None
self.layers, self.start_layer, self.end_layer = make_layers(
config.num_hidden_layers,
lambda idx, prefix: DeepseekV2DecoderLayer(