[fix] Fallback DeepGEMM activation for unsupported shapes (#26025)

2026-05-22 09:01:40 +00:00
parent 4486a339a3
commit bd6c7e713c
1 changed files with 5 additions and 2 deletions
@@ -863,8 +863,11 @@ def _varlen_deep_gemm_silu_mul_quant(
        dtype=torch.float8_e4m3fn,
    )

-    if envs.SGLANG_OPT_USE_JIT_EP_ACTIVATION.get():
-        assert N % 4 == 0 and G % 4 == 0
+    use_jit_ep_activation = envs.SGLANG_OPT_USE_JIT_EP_ACTIVATION.get()
+    if N % 4 != 0 or G % 4 != 0:
+        use_jit_ep_activation = False
+
+    if use_jit_ep_activation:
        packed_ue8m0 = deep_gemm_wrapper.DEEPGEMM_SCALE_UE8M0
        down_input_scale = torch.empty(
            (E, G // 4, N) if packed_ue8m0 else (E, N, G),