Support DeepSeek V3.2 Exp (#11061)

Co-authored-by: Stefan He <11166516+hebiao064@users.noreply.github.com>
Co-authored-by: Liangsheng Yin <95566987+hnyls2002@users.noreply.github.com>
Co-authored-by: Baizhou Zhang <56809903+fridge003@users.noreply.github.com>
Co-authored-by: DarkSharpness <76582120+darksharpness@users.noreply.github.com>
Co-authored-by: ZhengdQin <46387172+zhengdqin@users.noreply.github.com>
Co-authored-by: DarkSharpness <2040703891@qq.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>
Co-authored-by: Zhengda Qin <zhengdqin@gmail.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
Co-authored-by: HAI <hixiao@gmail.com>
Co-authored-by: Baizhou Zhang <sobereddiezhang@gmail.com>

python/sglang/test/get_logits_ut.py

@@ -0,0 +1,57 @@
+import torch
+import torch.nn as nn
+
+
+class DummyModel(nn.Module):
+    def __init__(self, d_in=2048, n_heads=128, softmax_scale=0.5):
+        super().__init__()
+        self.weights_proj = nn.Linear(d_in, 1024)
+        self.n_heads = n_heads
+        self.softmax_scale = softmax_scale
+
+    def _get_logits_head_gate_orig(self, x: torch.Tensor, q_scale: torch.Tensor):
+        weights = self.weights_proj(x)
+        weights = weights * self.n_heads**-0.5
+        q_scale = q_scale.unsqueeze(1)  # (B,1,1)
+        weights = weights.unsqueeze(-1) * q_scale * self.softmax_scale
+        return weights
+
+    def _get_logits_head_gate_opt(self, x: torch.Tensor, q_scale: torch.Tensor):
+        weights = self.weights_proj(x)
+        q_scale = q_scale.unsqueeze(1)  # (B,1,1)
+        scale_const = self.n_heads**-0.5 * q_scale * self.softmax_scale  # (B,1,1)
+        weights = weights.unsqueeze(-1) * scale_const  # (B,1024,1)
+        return weights
+
+
+def main():
+    torch.manual_seed(0)
+    model = DummyModel(d_in=2048, n_heads=128, softmax_scale=0.5)
+    x = torch.randn(128, 2048)  # batch=128, d_in=2048
+    q_scale = torch.randn(128, 1)
+
+    import time
+
+    start = time.time()
+    for _ in range(1000):
+        out_orig = model._get_logits_head_gate_orig(x, q_scale)
+    print("Original version time:", time.time() - start)
+
+    start = time.time()
+    for _ in range(1000):
+        out_opt = model._get_logits_head_gate_opt(x, q_scale)
+    print("Optimized version time:", time.time() - start)
+
+    print("Difference:", (out_orig - out_opt).abs().max().item())
+    assert torch.allclose(out_orig, out_opt), "Mismatch between original and optimized"
+
+
+if __name__ == "__main__":
+    main()
+
+
+"""
+Original version time: 0.49235057830810547
+Optimized version time: 0.4087331295013428
+Difference: 1.4901161193847656e-08
+"""