chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 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
python/sglang/srt/hardware_backend/npu/attention/ascend_backend.py
+6 -2
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",
python/sglang/srt/hardware_backend/npu/modules/deepseek_v2_attention_mla_npu.py
+34 -18
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,
python/sglang/srt/hardware_backend/npu/moe/topk.py
+4 -11
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(
python/sglang/srt/hardware_backend/npu/quantization/modelslim.py
+5 -3
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
python/sglang/srt/hardware_backend/npu/utils.py
+8
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