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

[CPU] [BF16] Call fused_experts_cpu, weight_packed_linear and bmm_cpu kernel in DeepSeek model (#6641)

Browse Source 
Co-authored-by: Thien Tran <gau.nernst@yahoo.com.sg>
This commit is contained in:
Chunyuan WU
2025-06-25 16:43:33 +08:00
committed by GitHub
parent bc2e5645c4
commit 7eb47b0f3d
9 changed files with 343 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
python/sglang/srt/layers/logits_processor.py
View File

@@ -442,11 +442,20 @@ class LogitsProcessor(nn.Module):
dp_gather_replicate(hidden_states, local_hidden_states, logits_metadata)
if hasattr(lm_head, "weight"):
logits = torch.matmul(
hidden_states.to(lm_head.weight.dtype), lm_head.weight.T
)
if getattr(lm_head, "use_intel_amx_backend", False):
logits = torch.ops.sgl_kernel.weight_packed_linear(
hidden_states.to(lm_head.weight.dtype),
lm_head.weight,
None, # bias
True, # is_vnni
)
else:
logits = torch.matmul(
hidden_states.to(lm_head.weight.dtype), lm_head.weight.T
)
else:
# GGUF models
# TODO: use weight_packed_linear for GGUF models
logits = lm_head.quant_method.apply(lm_head, hidden_states, embedding_bias)
if self.logit_scale is not None: