From 2d2c660104ee17178951ac68d668e540d690c39c Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Wed, 27 May 2026 23:00:19 -0400 Subject: [PATCH] [MoE] Remove inplace fused experts mechanism (#43727) Signed-off-by: Yongye Zhu --- .../moe/modular_kernel_tools/common.py | 1 - tests/kernels/moe/test_batched_deepgemm.py | 2 - tests/kernels/moe/test_block_fp8.py | 1 - tests/kernels/moe/test_cutlass_moe.py | 2 - tests/kernels/moe/test_deepep_deepgemm_moe.py | 3 - tests/kernels/moe/test_deepep_moe.py | 1 - tests/kernels/moe/test_deepgemm.py | 3 - tests/kernels/moe/test_flashinfer.py | 3 - tests/kernels/moe/test_flashinfer_b12x_moe.py | 1 - tests/kernels/moe/test_flashinfer_moe.py | 1 - .../moe/test_modular_oai_triton_moe.py | 1 - tests/kernels/moe/test_moe_layer.py | 1 - tests/kernels/moe/test_nvfp4_moe.py | 2 - tests/kernels/moe/test_trtllm_nvfp4_moe.py | 1 - tests/kernels/moe/utils.py | 3 - .../model_executor/layers/fused_moe/config.py | 6 - .../layers/fused_moe/eep_reconfigure.py | 1 - .../layers/fused_moe/experts/marlin_moe.py | 14 +- .../fused_moe/experts/rocm_aiter_moe.py | 2 - .../layers/fused_moe/fused_moe.py | 131 +----------------- .../fused_moe/fused_moe_modular_method.py | 2 - vllm/model_executor/layers/fused_moe/layer.py | 5 - .../layers/fused_moe/modular_kernel.py | 20 +-- .../layers/fused_moe/oracle/fp8.py | 4 - .../layers/fused_moe/oracle/int8.py | 1 - .../layers/fused_moe/oracle/int_wna16.py | 1 - .../layers/fused_moe/oracle/mxfp4.py | 3 - .../layers/fused_moe/oracle/nvfp4.py | 2 - .../layers/fused_moe/oracle/unquantized.py | 1 - .../layers/fused_moe/oracle/w4a8.py | 1 - .../layers/fused_moe/oracle/w4a8_int8.py | 1 - .../fused_moe/prepare_finalize/deepep_ht.py | 2 - .../fused_moe/prepare_finalize/naive_dp_ep.py | 1 - .../fused_moe/prepare_finalize/no_dp_ep.py | 1 - .../layers/fused_moe/runner/shared_experts.py | 1 - vllm/model_executor/layers/fused_moe/utils.py | 10 -- .../layers/quantization/bitsandbytes.py | 1 - .../compressed_tensors_moe_wna16.py | 1 - .../compressed_tensors_moe_wna16_marlin.py | 1 - .../layers/quantization/moe_wna16.py | 1 - .../layers/quantization/quark/quark_moe.py | 3 - vllm/model_executor/models/arctic.py | 1 - vllm/model_executor/models/bert_with_rope.py | 2 +- vllm/model_executor/models/minicpm.py | 2 +- 44 files changed, 11 insertions(+), 237 deletions(-) diff --git a/tests/kernels/moe/modular_kernel_tools/common.py b/tests/kernels/moe/modular_kernel_tools/common.py index 88aa9f2dc19..ea52a2d3398 100644 --- a/tests/kernels/moe/modular_kernel_tools/common.py +++ b/tests/kernels/moe/modular_kernel_tools/common.py @@ -647,7 +647,6 @@ def make_modular_kernel( modular_kernel = mk.FusedMoEKernel( prepare_finalize=prepare_finalize, fused_experts=fused_experts, - inplace=False, ) return modular_kernel diff --git a/tests/kernels/moe/test_batched_deepgemm.py b/tests/kernels/moe/test_batched_deepgemm.py index ab5672f8320..3fca9190dc3 100644 --- a/tests/kernels/moe/test_batched_deepgemm.py +++ b/tests/kernels/moe/test_batched_deepgemm.py @@ -80,7 +80,6 @@ def test_batched_deepgemm_vs_triton( mk_triton = FusedMoEKernel( prep_finalize, triton_experts, - inplace=False, ) out_triton = mk_triton.apply( @@ -105,7 +104,6 @@ def test_batched_deepgemm_vs_triton( mk_deepgemm = FusedMoEKernel( prep_finalize, deepgemm_experts, - inplace=False, ) out_deepgemm = mk_deepgemm.apply( diff --git a/tests/kernels/moe/test_block_fp8.py b/tests/kernels/moe/test_block_fp8.py index d11573bd31d..9a47c6e1d20 100644 --- a/tests/kernels/moe/test_block_fp8.py +++ b/tests/kernels/moe/test_block_fp8.py @@ -269,7 +269,6 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch) moe_config=moe_config, quant_config=quant_config, ), - inplace=False, ) def deep_gemm_moe_fp8(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids): diff --git a/tests/kernels/moe/test_cutlass_moe.py b/tests/kernels/moe/test_cutlass_moe.py index a613e7d2e29..32336f37ac6 100644 --- a/tests/kernels/moe/test_cutlass_moe.py +++ b/tests/kernels/moe/test_cutlass_moe.py @@ -214,7 +214,6 @@ def run_with_expert_maps( moe_config=moe_config, quant_config=new_quant_config, ), - inplace=False, ) out_tensor = out_tensor + kernel.apply(**kwargs) @@ -284,7 +283,6 @@ def run_8_bit( moe_config=moe_config, quant_config=quant_config, ), - inplace=False, ) return kernel.apply(**kwargs) diff --git a/tests/kernels/moe/test_deepep_deepgemm_moe.py b/tests/kernels/moe/test_deepep_deepgemm_moe.py index 6caa9d8c068..452bf64ed98 100644 --- a/tests/kernels/moe/test_deepep_deepgemm_moe.py +++ b/tests/kernels/moe/test_deepep_deepgemm_moe.py @@ -193,7 +193,6 @@ def make_ll_modular_kernel( return FusedMoEKernel( prepare_finalize=a2a, fused_experts=fused_experts, - inplace=False, ) @@ -226,7 +225,6 @@ def make_ht_modular_kernel( return FusedMoEKernel( prepare_finalize=a2a, fused_experts=fused_experts, - inplace=False, ) @@ -353,7 +351,6 @@ def triton_impl( w2=w2, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=False, quant_config=quant_config, ) diff --git a/tests/kernels/moe/test_deepep_moe.py b/tests/kernels/moe/test_deepep_moe.py index fd6dda384c1..5e0303c3df7 100644 --- a/tests/kernels/moe/test_deepep_moe.py +++ b/tests/kernels/moe/test_deepep_moe.py @@ -185,7 +185,6 @@ def make_modular_kernel( mk = FusedMoEKernel( prepare_finalize=a2a, fused_experts=fused_experts, - inplace=False, ) return mk diff --git a/tests/kernels/moe/test_deepgemm.py b/tests/kernels/moe/test_deepgemm.py index 9095bbc2c76..a9c6e9f5b19 100644 --- a/tests/kernels/moe/test_deepgemm.py +++ b/tests/kernels/moe/test_deepgemm.py @@ -126,7 +126,6 @@ def run_single_case(m, n, k, topk, num_experts, block_size): moe_config=moe_config, quant_config=quant_config, ), - inplace=False, ) # triton reference @@ -136,7 +135,6 @@ def run_single_case(m, n, k, topk, num_experts, block_size): w2=w2, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=False, quant_config=quant_config, ) @@ -323,7 +321,6 @@ def run_single_fp4_case(m, n, k, topk, num_experts): moe_config=moe_config, quant_config=quant_config, ), - inplace=False, ) # DeepGEMM FP4 path diff --git a/tests/kernels/moe/test_flashinfer.py b/tests/kernels/moe/test_flashinfer.py index 9e35be0db80..c8dc02927fa 100644 --- a/tests/kernels/moe/test_flashinfer.py +++ b/tests/kernels/moe/test_flashinfer.py @@ -233,7 +233,6 @@ def test_flashinfer_per_tensor_moe_fp8_no_graph( td.w2_quantized, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=False, activation=activation, global_num_experts=e, expert_map=None, @@ -321,7 +320,6 @@ def test_flashinfer_cutlass_moe_fp8_no_graph( td.w2_quantized, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=False, activation=activation, global_num_experts=e, expert_map=None, @@ -364,7 +362,6 @@ def test_flashinfer_cutlass_moe_fp8_no_graph( moe_config=moe_config, quant_config=quant_config, ), - inplace=False, ) flashinfer_cutlass_output = kernel.apply( diff --git a/tests/kernels/moe/test_flashinfer_b12x_moe.py b/tests/kernels/moe/test_flashinfer_b12x_moe.py index ec0a9594fe1..85d0bbe06d7 100644 --- a/tests/kernels/moe/test_flashinfer_b12x_moe.py +++ b/tests/kernels/moe/test_flashinfer_b12x_moe.py @@ -200,7 +200,6 @@ def test_flashinfer_b12x_moe( use_monolithic=False, ), experts, - inplace=False, ) score = torch.randn((m, e), device="cuda", dtype=dtype) diff --git a/tests/kernels/moe/test_flashinfer_moe.py b/tests/kernels/moe/test_flashinfer_moe.py index 6ae67fa0e98..2cec0bad1cb 100644 --- a/tests/kernels/moe/test_flashinfer_moe.py +++ b/tests/kernels/moe/test_flashinfer_moe.py @@ -117,7 +117,6 @@ def test_flashinfer_fp4_moe_no_graph( use_monolithic=False, ), FlashInferExperts(moe_config=moe_config, quant_config=quant_config), - inplace=False, ) flashinfer_output = flashinfer_experts.apply( diff --git a/tests/kernels/moe/test_modular_oai_triton_moe.py b/tests/kernels/moe/test_modular_oai_triton_moe.py index 589d90d1eca..0315d8d89e5 100644 --- a/tests/kernels/moe/test_modular_oai_triton_moe.py +++ b/tests/kernels/moe/test_modular_oai_triton_moe.py @@ -190,7 +190,6 @@ def oai_triton_moe_impl( use_monolithic=False, ), fused_experts, - inplace=False, ) return mk.apply( diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index 1bb899c1c89..e0f73cd657e 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -1159,7 +1159,6 @@ def make_fake_moe_layer( quant_config=quant_config, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=False, activation=activation, apply_router_weight_on_input=apply_router_weight_on_input, global_num_experts=global_num_experts, diff --git a/tests/kernels/moe/test_nvfp4_moe.py b/tests/kernels/moe/test_nvfp4_moe.py index e2a6cd1a7dc..dd5614ce119 100644 --- a/tests/kernels/moe/test_nvfp4_moe.py +++ b/tests/kernels/moe/test_nvfp4_moe.py @@ -105,7 +105,6 @@ def test_cutlass_fp4_moe_no_graph( moe_config=moe_config, quant_config=quant_config, ), - inplace=False, ) cutlass_output = kernel.apply( @@ -226,7 +225,6 @@ def test_cutlass_fp4_moe_swiglustep( moe_config=make_dummy_moe_config(), quant_config=quant_config, ), - inplace=False, ) cutlass_output = kernel.apply( diff --git a/tests/kernels/moe/test_trtllm_nvfp4_moe.py b/tests/kernels/moe/test_trtllm_nvfp4_moe.py index ca10ebaddce..4b4c3e712be 100644 --- a/tests/kernels/moe/test_trtllm_nvfp4_moe.py +++ b/tests/kernels/moe/test_trtllm_nvfp4_moe.py @@ -201,7 +201,6 @@ def test_trtllm_fp4_moe_no_graph( use_monolithic=False, ), trtllm_inner, - inplace=False, ) trtllm_output = trtllm_experts.apply( diff --git a/tests/kernels/moe/utils.py b/tests/kernels/moe/utils.py index acb2c21b389..3503ce4cdeb 100644 --- a/tests/kernels/moe/utils.py +++ b/tests/kernels/moe/utils.py @@ -142,7 +142,6 @@ def batched_moe( quant_config=quant_config, moe_config=moe_config, ), - inplace=False, ) return fused_experts.apply( @@ -195,7 +194,6 @@ def naive_batched_moe( quant_config=quant_config, moe_config=moe_config, ), - inplace=False, ) return fused_experts.apply( @@ -631,7 +629,6 @@ def modular_triton_fused_moe( use_monolithic=False, ), TritonExperts(moe_config, quant_config), - inplace=False, ) diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 8e92d0cfb67..da9d5d3e234 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1296,12 +1296,6 @@ class FusedMoEConfig: # cannot silently select one and drop the clamp. swiglu_limit: float | None = None - # This flag is used to disable the inplace optimization - # in MoE kernels. If this flag is True then the kernel - # should not be using inplace. If the flag is false, the - # kernel is free to use inplace or not. - disable_inplace: bool = True - def __post_init__(self): if self.dp_size > 1: logger.debug_once( diff --git a/vllm/model_executor/layers/fused_moe/eep_reconfigure.py b/vllm/model_executor/layers/fused_moe/eep_reconfigure.py index 6d40c6749f4..5ae055b6221 100644 --- a/vllm/model_executor/layers/fused_moe/eep_reconfigure.py +++ b/vllm/model_executor/layers/fused_moe/eep_reconfigure.py @@ -118,6 +118,5 @@ def make_eep_staged_quant_method( mk.FusedMoEKernel( prepare_finalize, experts, - inplace=moe_kernel.inplace, ), ) diff --git a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py index d2d4444d8ca..a358796a7d0 100644 --- a/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/marlin_moe.py @@ -30,7 +30,6 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( ) from vllm.model_executor.layers.fused_moe.utils import ( _resize_cache, - disable_inplace, swiglu_limit_func, ) from vllm.model_executor.layers.quantization.utils.marlin_utils import ( @@ -257,7 +256,6 @@ def fused_marlin_moe( is_k_full: bool = True, output: torch.Tensor | None = None, input_dtype: torch.dtype | None = None, - inplace: bool = False, clamp_limit: float | None = None, ) -> torch.Tensor: """ @@ -286,10 +284,6 @@ def fused_marlin_moe( - torch.Tensor: The output tensor after applying the MoE layer. """ - if inplace: - assert output is None, "Conflicting request" - assert not disable_inplace() - quant_type = ScalarType.from_id(quant_type_id) assert quant_type in [ scalar_types.uint4, @@ -379,7 +373,7 @@ def fused_marlin_moe( ).view(-1, topk, K) if output is None: - output = hidden_states if inplace else torch.empty_like(hidden_states) + output = torch.empty_like(hidden_states) if moe_sum is None: return torch.sum(moe_output.view(-1, topk, K), dim=1, out=output) @@ -417,7 +411,6 @@ def batched_fused_marlin_moe( is_k_full: bool = True, output: torch.Tensor | None = None, input_dtype: torch.dtype | None = None, - inplace: bool = False, clamp_limit: float | None = None, ) -> torch.Tensor: """ @@ -449,8 +442,6 @@ def batched_fused_marlin_moe( f"hidden states must be batched. e.g. [B, MAX_TOKENS, K]." f"But got {hidden_states.size()}" ) - if inplace: - assert output is None, "Conflicting request." quant_type = ScalarType.from_id(quant_type_id) assert quant_type in [ @@ -508,9 +499,6 @@ def batched_fused_marlin_moe( expert_num_tokens=expert_num_tokens, ) - if output is None and inplace: - output = hidden_states - # TODO (varun): This can be avoided by plumbing the marlin kernel to # ignore topk_weights when topk_weights_ptr is a nullptr. topk_weights = torch.ones( diff --git a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py index 64c8604acd6..b272a458b17 100644 --- a/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py @@ -17,7 +17,6 @@ from vllm.model_executor.layers.fused_moe.config import ( from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import ( TopKWeightAndReduceNoOP, ) -from vllm.model_executor.layers.fused_moe.utils import disable_inplace from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, kFp8Dynamic128Sym, @@ -516,7 +515,6 @@ class AiterExperts(mk.FusedMoEExpertsModular): and output.is_contiguous() and result.is_contiguous() and output._base is None - and disable_inplace() ): output.set_(result) else: diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index 8d974ea5671..e13ef52c01d 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -5,7 +5,6 @@ import functools import json import os -from collections.abc import Callable from typing import Any import torch @@ -27,7 +26,6 @@ from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( moe_align_block_size, ) from vllm.model_executor.layers.fused_moe.utils import ( - disable_inplace, moe_kernel_quantize_input, ) from vllm.platforms import current_platform @@ -1334,99 +1332,7 @@ def try_get_optimal_moe_config( return config -def inplace_fused_experts( - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - activation: str = "silu", - apply_router_weight_on_input: bool = False, - use_fp8_w8a8: bool = False, - use_int8_w8a8: bool = False, - use_int8_w8a16: bool = False, - use_int4_w4a16: bool = False, - ocp_mx_scheme: str | None = None, - per_channel_quant: bool = False, - global_num_experts: int = -1, - expert_map: torch.Tensor | None = None, - w1_scale: torch.Tensor | None = None, - w2_scale: torch.Tensor | None = None, - w1_zp: torch.Tensor | None = None, - w2_zp: torch.Tensor | None = None, - a1_scale: torch.Tensor | None = None, - a2_scale: torch.Tensor | None = None, - block_shape: list[int] | None = None, - w1_bias: torch.Tensor | None = None, - w2_bias: torch.Tensor | None = None, -) -> None: - fused_experts_impl( - hidden_states, - w1, - w2, - topk_weights, - topk_ids, - True, - activation, - apply_router_weight_on_input, - use_fp8_w8a8, - use_int8_w8a8, - use_int8_w8a16, - use_int4_w4a16, - ocp_mx_scheme, - per_channel_quant, - global_num_experts, - expert_map, - w1_scale, - w2_scale, - w1_zp, - w2_zp, - a1_scale, - a2_scale, - block_shape, - w1_bias, - w2_bias, - ) - - -def inplace_fused_experts_fake( - hidden_states: torch.Tensor, - w1: torch.Tensor, - w2: torch.Tensor, - topk_weights: torch.Tensor, - topk_ids: torch.Tensor, - activation: str = "silu", - apply_router_weight_on_input: bool = False, - use_fp8_w8a8: bool = False, - use_int8_w8a8: bool = False, - use_int8_w8a16: bool = False, - use_int4_w4a16: bool = False, - ocp_mx_scheme: str | None = None, - per_channel_quant: bool = False, - global_num_experts: int = -1, - expert_map: torch.Tensor | None = None, - w1_scale: torch.Tensor | None = None, - w2_scale: torch.Tensor | None = None, - w1_zp: torch.Tensor | None = None, - w2_zp: torch.Tensor | None = None, - a1_scale: torch.Tensor | None = None, - a2_scale: torch.Tensor | None = None, - block_shape: list[int] | None = None, - w1_bias: torch.Tensor | None = None, - w2_bias: torch.Tensor | None = None, -) -> None: - pass - - -direct_register_custom_op( - op_name="inplace_fused_experts", - op_func=inplace_fused_experts, - mutates_args=["hidden_states"], - fake_impl=inplace_fused_experts_fake, -) - - -def outplace_fused_experts( +def fused_experts_op( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, @@ -1458,7 +1364,6 @@ def outplace_fused_experts( w2, topk_weights, topk_ids, - False, activation, apply_router_weight_on_input, use_fp8_w8a8, @@ -1481,7 +1386,7 @@ def outplace_fused_experts( ) -def outplace_fused_experts_fake( +def fused_experts_op_fake( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, @@ -1511,28 +1416,12 @@ def outplace_fused_experts_fake( direct_register_custom_op( - op_name="outplace_fused_experts", - op_func=outplace_fused_experts, - fake_impl=outplace_fused_experts_fake, + op_name="fused_experts", + op_func=fused_experts_op, + fake_impl=fused_experts_op_fake, ) -def torch_vllm_inplace_fused_experts(**kwargs) -> torch.Tensor: - torch.ops.vllm.inplace_fused_experts(**kwargs) - hidden_states = kwargs["hidden_states"] - return hidden_states - - -def torch_vllm_outplace_fused_experts(**kwargs) -> torch.Tensor: - return torch.ops.vllm.outplace_fused_experts(**kwargs) - - -def dispatch_fused_experts_func(inplace: bool) -> Callable[..., torch.Tensor]: - if inplace: - return torch_vllm_inplace_fused_experts - return torch_vllm_outplace_fused_experts - - def _prepare_expert_assignment( topk_ids: torch.Tensor, config: dict[str, Any], @@ -1582,15 +1471,12 @@ def _prepare_expert_assignment( ) -# TODO (bnell): replace this with modular op. Can get rid of inplace/outplace -# torch ops. def fused_experts( hidden_states: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, - inplace: bool = False, activation: MoEActivation = MoEActivation.SILU, apply_router_weight_on_input: bool = False, global_num_experts: int = -1, @@ -1601,9 +1487,7 @@ def fused_experts( if quant_config is None: quant_config = FUSED_MOE_UNQUANTIZED_CONFIG - assert not inplace or not disable_inplace() - - return dispatch_fused_experts_func(inplace)( + return torch.ops.vllm.fused_experts( hidden_states=hidden_states, w1=w1, w2=w2, @@ -1667,7 +1551,6 @@ def fused_experts_impl( w2: torch.Tensor, topk_weights: torch.Tensor, topk_ids: torch.Tensor, - inplace: bool, activation: str = "silu", apply_router_weight_on_input: bool = False, use_fp8_w8a8: bool = False, @@ -1775,7 +1658,7 @@ def fused_experts_impl( else: raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}") - out_hidden_states = hidden_states if inplace else torch.empty_like(hidden_states) + out_hidden_states = torch.empty_like(hidden_states) qhidden_states, a1q_scale = moe_kernel_quantize_input( A=hidden_states, diff --git a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py index bbf06df8e47..3ebb63b0057 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py @@ -51,14 +51,12 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp): moe_layer: torch.nn.Module, old_quant_method: FusedMoEMethodBase, prepare_finalize: FusedMoEPrepareAndFinalizeModular, - inplace: bool = False, ) -> "FusedMoEModularMethod": return FusedMoEModularMethod( old_quant_method, FusedMoEKernel( prepare_finalize, old_quant_method.select_gemm_impl(prepare_finalize, moe_layer), - inplace=inplace, ), ) diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 776e86770ef..b5d8f5fd016 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -50,9 +50,6 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import ( UnquantizedFusedMoEMethod, ) -from vllm.model_executor.layers.fused_moe.utils import ( - disable_inplace, -) from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, ) @@ -337,7 +334,6 @@ class FusedMoE(PluggableLayer): routing_method=self.routing_method_type, swiglu_limit=swiglu_limit, # TODO: in_dtype == out_dtype? - disable_inplace=disable_inplace() or shared_experts is not None, ) if self.moe_config.use_mori_kernels: assert self.rocm_aiter_fmoe_enabled, ( @@ -473,7 +469,6 @@ class FusedMoE(PluggableLayer): self, self.base_quant_method, prepare_finalize, - inplace=not self.moe_config.disable_inplace, ) ) diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index 6d9e63278ab..6fbc1bffaac 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -27,7 +27,6 @@ from vllm.model_executor.layers.fused_moe.runner.shared_experts import ( ) from vllm.model_executor.layers.fused_moe.utils import ( _resize_cache, - disable_inplace, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, @@ -1025,11 +1024,9 @@ class FusedMoEKernelModularImpl: self, prepare_finalize: FusedMoEPrepareAndFinalizeModular, fused_experts: FusedMoEExpertsModular, - inplace: bool = False, ): self.prepare_finalize = prepare_finalize self.fused_experts = fused_experts - self.inplace = inplace moe_parallel_config = fused_experts.moe_config.moe_parallel_config self.moe_parallel_config = moe_parallel_config self.is_dp_ep = ( @@ -1392,12 +1389,7 @@ class FusedMoEKernelModularImpl: Returns: - torch.Tensor: The output tensor after applying the MoE layer. """ - if self.inplace: - assert shared_experts is None - assert not disable_inplace() - output = hidden_states - else: - output = torch.empty_like(hidden_states) + output = torch.empty_like(hidden_states) local_num_experts = w1.shape[0] if global_num_experts == -1: @@ -1473,7 +1465,6 @@ class FusedMoEKernelMonolithicImpl: that have fused router + experts (e.g. FLASHINFER_TRTLLM). """ - # TODO(rob): add inplace support. a1q, a1q_scale, router_logits = self.prepare_finalize.prepare( hidden_states, router_logits=router_logits, @@ -1509,7 +1500,6 @@ class FusedMoEKernel: self, prepare_finalize: FusedMoEPrepareAndFinalize, fused_experts: FusedMoEExperts, - inplace: bool = False, ): super().__init__() @@ -1521,13 +1511,11 @@ class FusedMoEKernel: self.impl = FusedMoEKernelModularImpl( prepare_finalize, fused_experts, - inplace, ) elif isinstance( prepare_finalize, FusedMoEPrepareAndFinalizeMonolithic ) and isinstance(fused_experts, FusedMoEExpertsMonolithic): - assert not inplace self.impl = FusedMoEKernelMonolithicImpl( prepare_finalize, fused_experts, @@ -1549,12 +1537,6 @@ class FusedMoEKernel: else: return False - @property - def inplace(self) -> bool: - if isinstance(self.impl, FusedMoEKernelModularImpl): - return self.impl.inplace - return False - @property def is_monolithic(self) -> bool: return isinstance(self.impl, FusedMoEKernelMonolithicImpl) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 562128d66d4..bbbad3e8271 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -627,10 +627,6 @@ def make_fp8_moe_kernel( kernel = mk.FusedMoEKernel( prepare_finalize, experts, - inplace=( - not moe_config.disable_inplace - and fp8_backend != Fp8MoeBackend.FLASHINFER_CUTLASS - ), ) return kernel diff --git a/vllm/model_executor/layers/fused_moe/oracle/int8.py b/vllm/model_executor/layers/fused_moe/oracle/int8.py index 36c4ccdd00b..6d50a3ba0ee 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int8.py @@ -213,7 +213,6 @@ def make_int8_moe_kernel( kernel = mk.FusedMoEKernel( prepare_finalize, experts, - inplace=not moe_config.disable_inplace, ) return kernel diff --git a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py index c7edd9a500d..64583540495 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py +++ b/vllm/model_executor/layers/fused_moe/oracle/int_wna16.py @@ -202,7 +202,6 @@ def make_wna16_moe_kernel( return mk.FusedMoEKernel( prepare_finalize, experts, - inplace=not moe_config.disable_inplace, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py index 3ecc4f018f1..a01718c5879 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/mxfp4.py @@ -1720,9 +1720,6 @@ def make_mxfp4_moe_kernel( kernel = mk.FusedMoEKernel( prepare_finalize, experts, - inplace=( - not moe_config.disable_inplace and mxfp4_backend not in TRTLLM_BACKENDS - ), ) return kernel diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index 5a2c64cecc0..21c44aad685 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -555,8 +555,6 @@ def make_nvfp4_moe_kernel( kernel = mk.FusedMoEKernel( prepare_finalize, experts, - inplace=False, ) - # TODO(rob): update inplace logic to be part of the kernel. return kernel diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 9c0df9d153b..4a9b190335d 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -359,7 +359,6 @@ def make_unquantized_moe_kernel( kernel = mk.FusedMoEKernel( prepare_finalize, experts, - inplace=(not moe_config.disable_inplace and not is_monolithic), ) return kernel diff --git a/vllm/model_executor/layers/fused_moe/oracle/w4a8.py b/vllm/model_executor/layers/fused_moe/oracle/w4a8.py index e886d9d5da5..d8e026ba319 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/w4a8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/w4a8.py @@ -192,5 +192,4 @@ def make_w4a8_moe_kernel( return mk.FusedMoEKernel( prepare_finalize, experts, - inplace=not moe_config.disable_inplace, ) diff --git a/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py b/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py index 7afa5ac2231..e4f4d497568 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py @@ -355,7 +355,6 @@ def make_w4a8_int8_moe_kernel( kernel = mk.FusedMoEKernel( prepare_finalize, experts, - inplace=not moe_config.disable_inplace, ) return kernel diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py index 9ca94521363..45f9e815ac8 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py @@ -383,7 +383,6 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): if event.event is not None: event.current_stream_wait() dbo_switch_to_comm() - # Respect inplace outputs. output.copy_(combined_x, non_blocking=True) # TODO(lucas): refactor the modular kernel so this will be @@ -394,7 +393,6 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalizeModular): else: # TODO(lucas): support this case with the refactored modular kernel assert not dbo_enabled() - # Respect inplace outputs. output.copy_(combined_x, non_blocking=True) return None diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py index b8633726c72..ffbb4c4a7d3 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py @@ -127,7 +127,6 @@ class MoEPrepareAndFinalizeNaiveDPEPModular(mk.FusedMoEPrepareAndFinalizeModular assert topk == 1, ( "apply_router_weight_on_input is only implemented for topk=1" ) - # Note: do not use inplace for shared experts overlap a1 = a1 * topk_weights.to(a1.dtype) a1q, scales = _quantize_and_setup_dispatch(a1, quant_config, defer_input_quant) diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py index 82079127251..69587770c40 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py @@ -71,7 +71,6 @@ class MoEPrepareAndFinalizeNoDPEPModular(mk.FusedMoEPrepareAndFinalizeModular): assert topk == 1, ( "apply_router_weight_on_input is only implemented for topk=1" ) - # Note: do not use inplace for shared experts overlap a1 = a1 * topk_weights.to(a1.dtype) a1q, a1q_scale = _quantize_input(a1, quant_config, defer_input_quant) diff --git a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py index e6492cef61c..f017829417f 100644 --- a/vllm/model_executor/layers/fused_moe/runner/shared_experts.py +++ b/vllm/model_executor/layers/fused_moe/runner/shared_experts.py @@ -119,7 +119,6 @@ class SharedExperts: if experts_order == SharedExpertsOrder.MULTI_STREAM_OVERLAPPED: assert self._stream is not None - assert self._moe_config.disable_inplace # Record that the clone will be used by shared_experts_stream # to avoid gc issue from deallocation of hidden_states_clone diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index b254dc7b7a7..cb2cd5e94a5 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -1,6 +1,5 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import functools from math import prod import torch @@ -32,7 +31,6 @@ from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv -from vllm.utils.torch_utils import is_torch_equal_or_newer @triton.jit @@ -368,14 +366,6 @@ def normalize_batched_scales_shape( return scales -# Torch custom ops can't deal with outputs aliasing inputs so we need to -# disable inplace for torch >= 2.9. -# See https://github.com/vllm-project/vllm/issues/26378 -@functools.cache -def disable_inplace() -> bool: - return is_torch_equal_or_newer("2.9") - - @triton.jit def _pack_topk_ids_weights_kernel( topk_ids_ptr, diff --git a/vllm/model_executor/layers/quantization/bitsandbytes.py b/vllm/model_executor/layers/quantization/bitsandbytes.py index a32ed3196c1..02267b8f682 100644 --- a/vllm/model_executor/layers/quantization/bitsandbytes.py +++ b/vllm/model_executor/layers/quantization/bitsandbytes.py @@ -497,7 +497,6 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase): w2=w2, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=not self.moe.disable_inplace, activation=layer.activation, apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py index f6418194fd7..88303f189f5 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py @@ -256,7 +256,6 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod): layer.w2_weight_packed, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=not self.moe.disable_inplace, activation=layer.activation, apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py index c60691df583..c56962422bc 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py @@ -573,5 +573,4 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod): workspace=layer.workspace, input_dtype=self.marlin_input_dtype, is_k_full=self.is_k_full, - inplace=not self.moe.disable_inplace, ) diff --git a/vllm/model_executor/layers/quantization/moe_wna16.py b/vllm/model_executor/layers/quantization/moe_wna16.py index d6729bbbcb7..471febab044 100644 --- a/vllm/model_executor/layers/quantization/moe_wna16.py +++ b/vllm/model_executor/layers/quantization/moe_wna16.py @@ -377,7 +377,6 @@ class MoeWNA16Method(FusedMoEMethodBase): layer.w2_qweight, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=not self.moe.disable_inplace, apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, diff --git a/vllm/model_executor/layers/quantization/quark/quark_moe.py b/vllm/model_executor/layers/quantization/quark/quark_moe.py index 627a44a77b5..9ee901a9910 100644 --- a/vllm/model_executor/layers/quantization/quark/quark_moe.py +++ b/vllm/model_executor/layers/quantization/quark/quark_moe.py @@ -497,7 +497,6 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod): apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, expert_map=layer.expert_map, - inplace=not self.moe.disable_inplace, ) else: from vllm.model_executor.layers.fused_moe import fused_experts @@ -508,7 +507,6 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod): w2=layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=not self.moe.disable_inplace, activation=layer.activation, apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, @@ -809,7 +807,6 @@ class QuarkW8A8Int8MoEMethod(QuarkMoEMethod): w2=layer.w2_weight, topk_weights=topk_weights, topk_ids=topk_ids, - inplace=not self.moe.disable_inplace, activation=layer.activation, apply_router_weight_on_input=layer.apply_router_weight_on_input, global_num_experts=layer.global_num_experts, diff --git a/vllm/model_executor/models/arctic.py b/vllm/model_executor/models/arctic.py index 6ab55a4b1bf..fee9eaaac8d 100644 --- a/vllm/model_executor/models/arctic.py +++ b/vllm/model_executor/models/arctic.py @@ -213,7 +213,6 @@ class ArcticMoE(nn.Module): self.w2s, topk_weights, topk_ids, - inplace=True, ) if self.reduce_results and self.tp_size > 1: final_hidden_states = tensor_model_parallel_all_reduce(final_hidden_states) diff --git a/vllm/model_executor/models/bert_with_rope.py b/vllm/model_executor/models/bert_with_rope.py index 22bcdeb453c..020ef993474 100644 --- a/vllm/model_executor/models/bert_with_rope.py +++ b/vllm/model_executor/models/bert_with_rope.py @@ -326,7 +326,7 @@ class NomicMoE(nn.Module): hidden_states, router_logits, self.top_k, renormalize=False ) - final_hidden_states = torch.ops.vllm.outplace_fused_experts( + final_hidden_states = torch.ops.vllm.fused_experts( hidden_states=hidden_states, w1=self.w1, w2=self.w2, diff --git a/vllm/model_executor/models/minicpm.py b/vllm/model_executor/models/minicpm.py index 54870eb2ede..ee9b31f140b 100644 --- a/vllm/model_executor/models/minicpm.py +++ b/vllm/model_executor/models/minicpm.py @@ -181,7 +181,7 @@ class MiniCPMMoE(nn.Module): ) final_hidden_states = fused_experts( - hidden_states, self.ws, self.w2s, topk_weights, topk_ids, inplace=False + hidden_states, self.ws, self.w2s, topk_weights, topk_ids ) if self.tp_size > 1: