A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
disable graph partition in custom op (#26952)
Signed-off-by: Boyuan Feng <boyuan@meta.com> Signed-off-by: Boyuan Feng <fby.1994@gmail.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
Boyuan Feng committed
Oct 16, 2025 at 20:08 UTC
08405609cc14d81f20e9faf61b2cd87e7909b797
2 files changed
+15
-1
vllm/model_executor/layers/fused_moe/fused_moe.py
+6
-1
@@ -49,6 +49,7 @@ from vllm.model_executor.layers.fused_moe.utils import (
49
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4
50
from vllm.model_executor.layers.quantization.utils.mxfp6_utils import dequant_mxfp6
51
from vllm.model_executor.layers.quantization.utils.ocp_mx_utils import OCP_MX_Scheme
52
+from vllm.model_executor.utils import maybe_disable_graph_partition
53
from vllm.platforms import current_platform
54
from vllm.triton_utils import tl, triton
55
from vllm.utils import direct_register_custom_op, is_torch_equal_or_newer
@@ -1145,7 +1146,11 @@ def fused_topk_bias(
1146
1147
1148
# This is used by the Deepseek-V2 and Deepseek-V3 model
1148
-@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend)
1149
+@torch.compile(
1150
+ dynamic=True,
1151
+ backend=current_platform.simple_compile_backend,
1152
+ options=maybe_disable_graph_partition(current_platform.simple_compile_backend),
1153
+)
1154
def grouped_topk(
1155
hidden_states: torch.Tensor,
1156
gating_output: torch.Tensor,
vllm/model_executor/utils.py
+9
@@ -7,6 +7,8 @@ from typing import Any
7
8
import torch
9
10
+from vllm.utils import is_torch_equal_or_newer
11
+
12
13
def set_random_seed(seed: int) -> None:
14
from vllm.platforms import current_platform
@@ -83,3 +85,10 @@ def get_moe_expert_mapping(
85
if child_map is not None:
86
return child_map()
87
return []
88
+
89
+
90
+def maybe_disable_graph_partition(current_backend: str) -> dict[str, bool]:
91
+ if current_backend == "inductor" and is_torch_equal_or_newer("2.9.0.dev"):
92
+ return {"graph_partition": False}
93
+ else:
94
+ return {}