diff --git a/python/sglang/srt/models/deepseek_v2.py b/python/sglang/srt/models/deepseek_v2.py index ed7022727..aa33a4273 100644 --- a/python/sglang/srt/models/deepseek_v2.py +++ b/python/sglang/srt/models/deepseek_v2.py @@ -490,7 +490,11 @@ class MoEGate(nn.Module): or mla_use_prefill_cp(forward_batch, self.mla_enable_prefill_cp) ) ): - logits = F.linear(hidden_states, self.weight, None) + if _is_cuda: + from sglang.jit_kernel.dsv4 import linear_bf16_fp32 + + return linear_bf16_fp32(hidden_states, self.weight) + return F.linear(hidden_states, self.weight, None) else: # NOTE(b8zhong): this threshold has been empirically verified max_router_gemm_tokens = 4 if _device_sm in (100, 103) else 16