karylab_agents/vllm
1
0
Fork 0
forked from Karylab-cklius/vllm
Code Pull Requests 2 Actions 1 Packages Activity

[MoE] Remove inplace fused experts mechanism (#43727)

Browse Source 
Signed-off-by: Yongye Zhu <zyy1102000@gmail.com>
This commit is contained in:
Yongye Zhu
2026-05-27 20:00:19 -07:00
committed by GitHub
parent 05eec7120e
commit 2d2c660104
44 changed files with 11 additions and 237 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/kernels/moe/modular_kernel_tools/common.py
-1
View File
@@ -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
tests/kernels/moe/test_batched_deepgemm.py
-2
View File
@@ -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(
tests/kernels/moe/test_block_fp8.py
-1
View File
@@ -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):
tests/kernels/moe/test_cutlass_moe.py
-2
View File
@@ -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)
tests/kernels/moe/test_deepep_deepgemm_moe.py
-3
View File
@@ -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,
)
tests/kernels/moe/test_deepep_moe.py
-1
View File
@@ -185,7 +185,6 @@ def make_modular_kernel(
mk = FusedMoEKernel(
prepare_finalize=a2a,
fused_experts=fused_experts,
inplace=False,
)
return mk
tests/kernels/moe/test_deepgemm.py
-3
View File
@@ -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
tests/kernels/moe/test_flashinfer.py
-3
View File
@@ -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(
tests/kernels/moe/test_flashinfer_b12x_moe.py
-1
View File
@@ -200,7 +200,6 @@ def test_flashinfer_b12x_moe(
use_monolithic=False,
),
experts,
inplace=False,
)
score = torch.randn((m, e), device="cuda", dtype=dtype)
tests/kernels/moe/test_flashinfer_moe.py
-1
View File
@@ -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(
tests/kernels/moe/test_modular_oai_triton_moe.py
-1
View File
@@ -190,7 +190,6 @@ def oai_triton_moe_impl(
use_monolithic=False,
),
fused_experts,
inplace=False,
)
return mk.apply(
tests/kernels/moe/test_moe_layer.py
-1
View File
@@ -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,
tests/kernels/moe/test_nvfp4_moe.py
-2
View File
@@ -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(
tests/kernels/moe/test_trtllm_nvfp4_moe.py
-1
View File
@@ -201,7 +201,6 @@ def test_trtllm_fp4_moe_no_graph(
use_monolithic=False,
),
trtllm_inner,
inplace=False,
)
trtllm_output = trtllm_experts.apply(
tests/kernels/moe/utils.py
-3
View File
@@ -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,
)
vllm/model_executor/layers/fused_moe/config.py
-6
View File
@@ -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(
vllm/model_executor/layers/fused_moe/eep_reconfigure.py
-1
View File
@@ -118,6 +118,5 @@ def make_eep_staged_quant_method(
mk.FusedMoEKernel(
prepare_finalize,
experts,
inplace=moe_kernel.inplace,
),
)
vllm/model_executor/layers/fused_moe/experts/marlin_moe.py
+1 -13
View File
@@ -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(
vllm/model_executor/layers/fused_moe/experts/rocm_aiter_moe.py
-2
View File
@@ -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:
vllm/model_executor/layers/fused_moe/fused_moe.py
+7 -124
View File
@@ -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,
vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py
-2
View File
@@ -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,
),
)
vllm/model_executor/layers/fused_moe/layer.py
-5
View File
@@ -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,
)
)
vllm/model_executor/layers/fused_moe/modular_kernel.py
+1 -19
View File
@@ -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)
vllm/model_executor/layers/fused_moe/oracle/fp8.py
-4
View File
@@ -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
vllm/model_executor/layers/fused_moe/oracle/int8.py
-1
View File
@@ -213,7 +213,6 @@ def make_int8_moe_kernel(
kernel = mk.FusedMoEKernel(
prepare_finalize,
experts,
inplace=not moe_config.disable_inplace,
)
return kernel
vllm/model_executor/layers/fused_moe/oracle/int_wna16.py
-1
View File
@@ -202,7 +202,6 @@ def make_wna16_moe_kernel(
return mk.FusedMoEKernel(
prepare_finalize,
experts,
inplace=not moe_config.disable_inplace,
)
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
-3
View File
@@ -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
vllm/model_executor/layers/fused_moe/oracle/nvfp4.py
-2
View File
@@ -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
vllm/model_executor/layers/fused_moe/oracle/unquantized.py
-1
View File
@@ -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
vllm/model_executor/layers/fused_moe/oracle/w4a8.py
-1
View File
@@ -192,5 +192,4 @@ def make_w4a8_moe_kernel(
return mk.FusedMoEKernel(
prepare_finalize,
experts,
inplace=not moe_config.disable_inplace,
)
vllm/model_executor/layers/fused_moe/oracle/w4a8_int8.py
-1
View File
@@ -355,7 +355,6 @@ def make_w4a8_int8_moe_kernel(
kernel = mk.FusedMoEKernel(
prepare_finalize,
experts,
inplace=not moe_config.disable_inplace,
)
return kernel
vllm/model_executor/layers/fused_moe/prepare_finalize/deepep_ht.py
-2
View File
@@ -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
vllm/model_executor/layers/fused_moe/prepare_finalize/naive_dp_ep.py
-1
View File
@@ -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)
vllm/model_executor/layers/fused_moe/prepare_finalize/no_dp_ep.py
-1
View File
@@ -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)
vllm/model_executor/layers/fused_moe/runner/shared_experts.py
-1
View File
@@ -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
vllm/model_executor/layers/fused_moe/utils.py
-10
View File
@@ -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,
vllm/model_executor/layers/quantization/bitsandbytes.py
-1
View File
@@ -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,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16.py
-1
View File
@@ -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,
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_wna16_marlin.py
-1
View File
@@ -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,
)
vllm/model_executor/layers/quantization/moe_wna16.py
-1
View File
@@ -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,
vllm/model_executor/layers/quantization/quark/quark_moe.py
-3
View File
@@ -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,
vllm/model_executor/models/arctic.py
-1
View File
@@ -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)
vllm/model_executor/models/bert_with_rope.py
+1 -1
View File
@@ -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,
vllm/model_executor/models/minicpm.py
+1 -1
View File
@@ -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: