Apply torch.compile to fused_moe/grouped_topk (#12637)

2025-02-01 11:16:19 -05:00
parent 4f4d427ac2
commit 3194039c0e
2 changed files with 3 additions and 2 deletions
--- a/vllm/model_executor/layers/fused_moe/fused_moe.py
+++ b/vllm/model_executor/layers/fused_moe/fused_moe.py
@@ -759,6 +759,7 @@ def fused_topk(


 # This is used by the Deepseek-V2 and Deepseek-V3 model
+@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
 def grouped_topk(hidden_states: torch.Tensor,
                 gating_output: torch.Tensor,
                 topk: int,