disable graph partition in custom op

vllm/model_executor/layers/fused_moe/layer.py

@@ -46,7 +46,7 @@
     QuantizationConfig,
     QuantizeMethodBase,
 )
-from vllm.model_executor.utils import set_weight_attrs
+from vllm.model_executor.utils import disable_graph_partition, set_weight_attrs
 from vllm.platforms import current_platform
 from vllm.platforms.interface import CpuArchEnum
 from vllm.utils import cdiv, direct_register_custom_op, has_deep_ep, has_pplx, round_up
@@ -1900,17 +1900,19 @@ def select_experts(
         if use_grouped_topk:
             assert topk_group is not None
             assert num_expert_group is not None
-            topk_weights, topk_ids = grouped_topk(
-                hidden_states=hidden_states,
-                gating_output=router_logits,
-                topk=top_k,
-                renormalize=renormalize,
-                num_expert_group=num_expert_group,
-                topk_group=topk_group,
-                scoring_func=scoring_func,
-                routed_scaling_factor=routed_scaling_factor,
-                e_score_correction_bias=e_score_correction_bias,
-            )
+
+            with disable_graph_partition():
+                topk_weights, topk_ids = grouped_topk(
+                    hidden_states=hidden_states,
+                    gating_output=router_logits,
+                    topk=top_k,
+                    renormalize=renormalize,
+                    num_expert_group=num_expert_group,
+                    topk_group=topk_group,
+                    scoring_func=scoring_func,
+                    routed_scaling_factor=routed_scaling_factor,
+                    e_score_correction_bias=e_score_correction_bias,
+                )
             if indices_type is not None:
                 topk_ids = topk_ids.to(dtype=indices_type)
         elif e_score_correction_bias is not None:

vllm/model_executor/utils.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 """Utils for model executor."""
 
+import contextlib
 import copy
 from typing import Any
 
@@ -83,3 +84,26 @@ def get_moe_expert_mapping(
             if child_map is not None:
                 return child_map()
         return []
+
+
+@contextlib.contextmanager
+def disable_graph_partition():
+    """Context manager to disable inductor graph partition.
+    This is used to avoid nested cudagraph capture.
+
+    Example:
+    1. We apply torch.compile directly on some ops (e.g., grouped_topk) wrapped
+    in custom ops. Inductor graph partition applies cudagraph within the custom op.
+    2. At the same time, we compile the model which uses these custom ops. Inductor
+    graph partition also wraps each graph partition with CUDAGraph. Some partitions
+    may include custom ops, which has already been applied cudagraph. This leads to
+    nested cudagraph which is not supported.
+
+    This context manager should be wrapped around torch.compile calls within custom ops
+    to avoid the nested cudagraph capture."""
+    old_val = torch._inductor.config.graph_partition
+    try:
+        torch._inductor.config.graph_partition = False
+        yield
+    finally:
+        torch._inductor.config.graph_partition = old_val