[Kernel] Raise an exception in MoE kernel if the batch size is larger…

… then 65k (#5939)
vllm-project · Jun 29, 2024 · f7dac83 · f7dac83
1 parent 7c01f70
commit f7dac83
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -423,6 +423,11 @@ def fused_experts(hidden_states: torch.Tensor,
     M, _ = hidden_states.shape
     E, N, _ = w1.shape
 
+    if M > 65536:
+        # https://github.com/vllm-project/vllm/issues/5938
+        raise ValueError("MoE kernel does not support more than 65536 tokens, "
+                         f"but got {M}")
+
     if override_config:
         config = override_config
     else: