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
[Bugfix] Disable moe inplace for torch >= 2.9 (#26497)
Signed-off-by: Bill Nell <bnell@redhat.com>
bnellnm committed
Oct 9, 2025 at 14:07 UTC
a462331e367191d61414878072dcee8f23bb9214
4 files changed
+22
-6
vllm/model_executor/layers/fused_moe/fused_marlin_moe.py
+6
-2
@@ -14,7 +14,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import moe_align_block_size
14
from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
15
TopKWeightAndReduceNoOP,
16
)
17
-from vllm.model_executor.layers.fused_moe.utils import _resize_cache
17
+from vllm.model_executor.layers.fused_moe.utils import _resize_cache, disable_inplace
18
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
19
marlin_make_workspace_new,
20
marlin_moe_intermediate_size,
@@ -235,7 +235,11 @@ def fused_marlin_moe(
235
).view(-1, topk, K)
236
237
if output is None:
238
- output = hidden_states if inplace else torch.empty_like(hidden_states)
238
+ if inplace and not disable_inplace():
239
+ output = hidden_states
240
+ else:
241
+ output = torch.empty_like(hidden_states)
242
+
243
return torch.sum(intermediate_cache3.view(-1, topk, K), dim=1, out=output)
244
245
vllm/model_executor/layers/fused_moe/fused_moe.py
+6
-2
@@ -39,6 +39,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
39
from vllm.model_executor.layers.fused_moe.utils import (
40
_resize_cache,
41
activation_without_mul,
42
+ disable_inplace,
43
moe_kernel_quantize_input,
44
)
45
from vllm.model_executor.layers.quantization.utils.mxfp4_utils import dequant_mxfp4
@@ -1516,7 +1517,7 @@ def torch_vllm_outplace_fused_experts(**kwargs) -> torch.Tensor:
1517
1518
1519
def dispatch_fused_experts_func(inplace: bool) -> Callable[..., torch.Tensor]:
1519
- if inplace:
1520
+ if inplace and not disable_inplace():
1521
return torch_vllm_inplace_fused_experts
1522
return torch_vllm_outplace_fused_experts
1523
@@ -1766,7 +1767,10 @@ def fused_experts_impl(
1767
else:
1768
raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")
1769
1769
- out_hidden_states = hidden_states if inplace else torch.empty_like(hidden_states)
1770
+ if inplace and not disable_inplace():
1771
+ out_hidden_states = hidden_states
1772
+ else:
1773
+ out_hidden_states = torch.empty_like(hidden_states)
1774
1775
if ocp_mx_scheme is not None:
1776
# TODO: On platforms for which `current_platform.supports_mx()` is True
vllm/model_executor/layers/fused_moe/modular_kernel.py
+2
-1
@@ -13,6 +13,7 @@ from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
13
from vllm.model_executor.layers.fused_moe.utils import (
14
_resize_cache,
15
count_expert_num_tokens,
16
+ disable_inplace,
17
)
18
from vllm.utils import cdiv
19
from vllm.v1.worker.ubatching import (
@@ -1139,7 +1140,7 @@ class FusedMoEModularKernel(torch.nn.Module):
1140
- torch.Tensor: The output tensor after applying the MoE layer.
1141
"""
1142
1142
- if inplace and self.shared_experts is None:
1143
+ if inplace and self.shared_experts is None and not disable_inplace():
1144
output = hidden_states
1145
else:
1146
output = torch.zeros_like(hidden_states)
vllm/model_executor/layers/fused_moe/utils.py
+8
-1
@@ -23,7 +23,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
23
mxfp8_e4m3_quantize,
24
)
25
from vllm.triton_utils import tl, triton
26
-from vllm.utils import cdiv
26
+from vllm.utils import cdiv, is_torch_equal_or_newer
27
from vllm.utils.flashinfer import flashinfer_fp4_quantize
28
29
@@ -321,3 +321,10 @@ def _validate_scale_shape(
321
322
def activation_without_mul(activation: str) -> str:
323
return activation + "_no_mul"
324
+
325
+
326
+# Torch custom ops can't deal with outputs aliasing inputs so we need to
327
+# disable inplace for torch >= 2.9.
328
+# See https://github.com/vllm-project/vllm/issues/26378
329
+def disable_inplace() -> bool:
330
+ return is_torch_equal_or_newer("2.9")