a57c8228ff · tduncan/Nvidia-Optimization-Stack-VLLM

tests/kernels/moe/modular_kernel_tools/common.py

+1

@@ -620,6 +620,7 @@ def make_modular_kernel(

620 modular_kernel = mk.FusedMoEModularKernel(

621 prepare_finalize=prepare_finalize,

622 fused_experts=fused_experts,

623 + inplace=False,

624 )

625

626 return modular_kernel

tests/kernels/moe/test_batched_deepgemm.py

+10 -4

@@ -74,7 +74,11 @@ def test_batched_deepgemm_vs_triton(

74 quant_config=quant_config,

75 moe_config=make_dummy_moe_config(),

76 )

77 - mk_triton = FusedMoEModularKernel(prep_finalize, triton_experts)

77 + mk_triton = FusedMoEModularKernel(

78 + prep_finalize,

79 + triton_experts,

80 + inplace=False,

81 + )

82

83 out_triton = mk_triton(

84 hidden_states=a,

@@ -82,7 +86,6 @@ def test_batched_deepgemm_vs_triton(

86 w2=w2,

87 topk_weights=topk_weights,

88 topk_ids=topk_ids,

85 - inplace=False,

89 global_num_experts=E,

90 )

91

@@ -93,7 +96,11 @@ def test_batched_deepgemm_vs_triton(

96 quant_config=quant_config,

97 moe_config=make_dummy_moe_config(),

98 )

96 - mk_deepgemm = FusedMoEModularKernel(prep_finalize, deepgemm_experts)

99 + mk_deepgemm = FusedMoEModularKernel(

100 + prep_finalize,

101 + deepgemm_experts,

102 + inplace=False,

103 + )

104

105 out_deepgemm = mk_deepgemm(

106 hidden_states=a,

@@ -101,7 +108,6 @@ def test_batched_deepgemm_vs_triton(

108 w2=w2,

109 topk_weights=topk_weights,

110 topk_ids=topk_ids,

104 - inplace=False,

111 global_num_experts=E,

112 )

113

tests/kernels/moe/test_block_fp8.py

+2 -3

@@ -9,6 +9,7 @@ from tests.kernels.moe.utils import (

9 make_dummy_moe_config,

10 make_test_quant_config,

11 make_test_weights,

12 + modular_triton_fused_moe,

13 )

14 from tests.kernels.quant_utils import (

15 native_per_token_group_quant_fp8,

@@ -26,9 +27,6 @@ from vllm.model_executor.layers.fused_moe.config import (

27 from vllm.model_executor.layers.fused_moe.deep_gemm_moe import (

28 _valid_deep_gemm_shape,

29 )

29 -from vllm.model_executor.layers.fused_moe.fused_moe import (

30 - modular_triton_fused_moe,

31 -)

30 from vllm.model_executor.layers.fused_moe.prepare_finalize import (

31 MoEPrepareAndFinalizeNoEP,

32 )

@@ -261,6 +259,7 @@ def test_w8a8_block_fp8_deep_gemm_fused_moe(M, N, K, E, topk, seed, monkeypatch)

259 moe_config=make_dummy_moe_config(),

260 quant_config=quant_config,

261 ),

262 + inplace=False,

263 )

264

265 def deep_gemm_moe_fp8(a, w1, w2, w1_s, w2_s, topk_weights, topk_ids):

tests/kernels/moe/test_cutlass_moe.py

+2

@@ -207,6 +207,7 @@ def run_with_expert_maps(

207 ),

208 quant_config=new_quant_config,

209 ),

210 + inplace=False,

211 )

212 out_tensor = out_tensor + kernel(**kwargs)

213

@@ -266,6 +267,7 @@ def run_8_bit(

267 ),

268 quant_config=quant_config,

269 ),

270 + inplace=False,

271 )

272 return kernel(**kwargs)

273

tests/kernels/moe/test_deepep_deepgemm_moe.py

+10 -5

@@ -194,8 +194,11 @@ def make_ll_modular_kernel(

194 quant_config=quant_config,

195 moe_config=make_dummy_moe_config(),

196 )

197 - mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts)

198 - return mk

197 + return FusedMoEModularKernel(

198 + prepare_finalize=a2a,

199 + fused_experts=fused_experts,

200 + inplace=False,

201 + )

202

203

204 def make_ht_modular_kernel(

@@ -224,8 +227,11 @@ def make_ht_modular_kernel(

227 moe_config=make_dummy_moe_config(),

228 quant_config=quant_config,

229 )

227 - mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts)

228 - return mk

230 + return FusedMoEModularKernel(

231 + prepare_finalize=a2a,

232 + fused_experts=fused_experts,

233 + inplace=False,

234 + )

235

236

237 def make_modular_kernel(

@@ -318,7 +324,6 @@ def deepep_deepgemm_moe_impl(

324 w2=w2,

325 topk_weights=test_tensors.topk_weights,

326 topk_ids=test_tensors.topk,

321 - inplace=False,

327 activation="silu",

328 global_num_experts=num_experts,

329 expert_map=build_expert_map(),

tests/kernels/moe/test_deepep_moe.py

+5 -2

@@ -179,7 +179,11 @@ def make_modular_kernel(

179 quant_config=quant_config,

180 )

181

182 - mk = FusedMoEModularKernel(prepare_finalize=a2a, fused_experts=fused_experts)

182 + mk = FusedMoEModularKernel(

183 + prepare_finalize=a2a,

184 + fused_experts=fused_experts,

185 + inplace=False,

186 + )

187 return mk

188

189

@@ -256,7 +260,6 @@ def deep_ep_moe_impl(

260 w2=w2,

261 topk_weights=topk_weights_chunk,

262 topk_ids=topk_chunk,

259 - inplace=False,

263 activation="silu",

264 global_num_experts=num_experts,

265 expert_map=build_expert_map(),

tests/kernels/moe/test_deepgemm.py

+1 -1

@@ -115,6 +115,7 @@ def run_single_case(m, n, k, topk, num_experts, block_size):

115 moe_config=make_dummy_moe_config(),

116 quant_config=quant_config,

117 ),

118 + inplace=False,

119 )

120

121 # triton reference

@@ -135,7 +136,6 @@ def run_single_case(m, n, k, topk, num_experts, block_size):

136 w2=w2,

137 topk_weights=topk_weights,

138 topk_ids=topk_ids,

138 - inplace=False,

139 )

140 diff = calc_diff(out_deepgemm, out_triton)

141 assert diff < 0.001, f"Diff exceeded 1%: {diff}"

tests/kernels/moe/test_flashinfer.py

+1 -1

@@ -301,6 +301,7 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

301 moe_config=moe_config,

302 quant_config=quant_config,

303 ),

304 + inplace=False,

305 )

306

307 flashinfer_cutlass_output = kernel(

@@ -309,7 +310,6 @@ def test_flashinfer_cutlass_moe_fp8_no_graph(

310 td.layer.w2_weight,

311 topk_weights,

312 topk_ids,

312 - inplace=False,

313 activation=activation,

314 global_num_experts=e,

315 expert_map=None,

tests/kernels/moe/test_flashinfer_moe.py

+1

@@ -108,6 +108,7 @@ def test_flashinfer_fp4_moe_no_graph(

108 flashinfer_experts = FusedMoEModularKernel(

109 MoEPrepareAndFinalizeNoEP(),

110 FlashInferExperts(moe_config=moe_config, quant_config=quant_config),

111 + inplace=False,

112 )

113

114 fi_activation = {"silu_and_mul": "silu", "relu2": "relu2_no_mul"}[activation]

tests/kernels/moe/test_modular_oai_triton_moe.py

+5 -2

@@ -180,7 +180,11 @@ def oai_triton_moe_impl(

180 else:

181 fused_experts = OAITritonExperts(make_dummy_moe_config(), quant_config)

182

183 - mk = FusedMoEModularKernel(MoEPrepareAndFinalizeNoEP(), fused_experts)

183 + mk = FusedMoEModularKernel(

184 + MoEPrepareAndFinalizeNoEP(),

185 + fused_experts,

186 + inplace=False,

187 + )

188

189 return mk.forward(

190 hidden_states=x,

@@ -188,7 +192,6 @@ def oai_triton_moe_impl(

192 w2=w2,

193 topk_weights=topk_weights,

194 topk_ids=topk_ids,

191 - inplace=True,

195 activation="swigluoai",

196 global_num_experts=num_experts,

197 expert_map=None,

tests/kernels/moe/test_moe.py

+5 -4

@@ -18,7 +18,11 @@ from transformers import MixtralConfig

18 from transformers.models.mixtral.modeling_mixtral import MixtralSparseMoeBlock

19

20 import vllm.model_executor.layers.fused_moe # noqa

21 -from tests.kernels.moe.utils import fused_moe, make_dummy_moe_config

21 +from tests.kernels.moe.utils import (

22 + fused_moe,

23 + make_dummy_moe_config,

24 + modular_triton_fused_moe,

25 +)

26 from tests.kernels.utils import opcheck, stack_and_dev, torch_experts, torch_moe

27 from vllm._aiter_ops import rocm_aiter_ops

28 from vllm.config import VllmConfig, set_current_vllm_config

@@ -36,9 +40,6 @@ from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (

40 batched_fused_marlin_moe,

41 fused_marlin_moe,

42 )

39 -from vllm.model_executor.layers.fused_moe.fused_moe import (

40 - modular_triton_fused_moe,

41 -)

43 from vllm.model_executor.layers.quantization.utils.marlin_utils import (

44 marlin_permute_bias,

45 )

tests/kernels/moe/test_nvfp4_moe.py

+1

@@ -95,6 +95,7 @@ def test_cutlass_fp4_moe_no_graph(

95 moe_config=make_dummy_moe_config(),

96 quant_config=quant_config,

97 ),

98 + inplace=False,

99 )

100

101 cutlass_output = kernel(

tests/kernels/moe/test_pplx_cutlass_moe.py

+1

@@ -172,6 +172,7 @@ def pplx_cutlass_moe(

172 fused_cutlass_experts = FusedMoEModularKernel(

173 prepare_finalize,

174 experts,

175 + inplace=False,

176 )

177

178 a_chunk = chunk_by_rank(a, rank, world_size).to(device)

tests/kernels/moe/test_pplx_moe.py

+1

@@ -592,6 +592,7 @@ def pplx_moe(

592 prepare_finalize,

593 experts,

594 shared_experts,

595 + inplace=False,

596 )

597

598 # Note: for now use_compile will error out if the problem size is

tests/kernels/moe/utils.py

+23 -1

@@ -7,7 +7,11 @@ import vllm._custom_ops as ops

7 from tests.kernels.quant_utils import per_block_cast_to_int8

8 from tests.kernels.quantization.nvfp4_utils import FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX

9 from vllm.model_executor.layers.activation import SiluAndMul

10 -from vllm.model_executor.layers.fused_moe import fused_experts, fused_topk

10 +from vllm.model_executor.layers.fused_moe import (

11 + TritonExperts,

12 + fused_experts,

13 + fused_topk,

14 +)

15 from vllm.model_executor.layers.fused_moe.config import (

16 FusedMoEConfig,

17 FusedMoEParallelConfig,

@@ -20,6 +24,9 @@ from vllm.model_executor.layers.fused_moe.fused_batched_moe import (

24 NaiveBatchedExperts,

25 )

26 from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel

27 +from vllm.model_executor.layers.fused_moe.prepare_finalize import (

28 + MoEPrepareAndFinalizeNoEP,

29 +)

30 from vllm.model_executor.layers.fused_moe.utils import moe_kernel_quantize_input

31 from vllm.utils.deep_gemm import per_block_cast_to_fp8

32 from vllm.utils.math_utils import round_up

@@ -116,6 +123,7 @@ def batched_moe(

123 quant_config=quant_config,

124 moe_config=make_dummy_moe_config(),

125 ),

126 + inplace=False,

127 )

128

129 return fused_experts(a, w1, w2, topk_weight, topk_ids)

@@ -157,6 +165,7 @@ def naive_batched_moe(

165 quant_config=quant_config,

166 moe_config=make_dummy_moe_config(),

167 ),

168 + inplace=False,

169 )

170

171 return fused_experts(a, w1, w2, topk_weight, topk_ids)

@@ -554,3 +563,16 @@ def make_shared_experts(

563 return RealMLP(K, N, w1, w2, "silu", quant_config, w1_s=w1_s, w2_s=w2_s)

564 finally:

565 torch.set_default_dtype(old_dtype)

566 +

567 +

568 +def modular_triton_fused_moe(

569 + moe_config: FusedMoEConfig,

570 + quant_config: FusedMoEQuantConfig,

571 + shared_experts: torch.nn.Module | None = None,

572 +) -> FusedMoEModularKernel:

573 + return FusedMoEModularKernel(

574 + MoEPrepareAndFinalizeNoEP(),

575 + TritonExperts(moe_config, quant_config),

576 + shared_experts,

577 + inplace=False,

578 + )

vllm/model_executor/layers/fused_moe/config.py

+6 -3

@@ -1083,13 +1083,16 @@ class FusedMoEConfig:

1083 router_logits_dtype: torch.dtype | None = None

1084

1085 max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE

1086 -

1086 has_bias: bool = False

1088 -

1087 is_act_and_mul: bool = True

1090 -

1088 is_lora_enabled: bool = False

1089

1090 + # This flag is used to disable the inplace optimization

1091 + # in MoE kernels. If this flag is True then the kernel

1092 + # should not be using inplace. If the flag is false, the

1093 + # kernel is free to use inplace or not.

1094 + disable_inplace: bool = True

1095 +

1096 def __post_init__(self):

1097 if self.dp_size > 1:

1098 logger.debug_once(

vllm/model_executor/layers/fused_moe/cutlass_moe.py

+1

@@ -1165,6 +1165,7 @@ def cutlass_moe_w4a8_fp8(

1165 quant_config=quant_config,

1166 group_size=group_size,

1167 ),

1168 + inplace=False,

1169 )

1170

1171 return fn(

vllm/model_executor/layers/fused_moe/fused_marlin_moe.py

+2 -4

@@ -267,6 +267,7 @@ def fused_marlin_moe(

267

268 if inplace:

269 assert output is None, "Conflicting request"

270 + assert not disable_inplace()

271

272 quant_type = ScalarType.from_id(quant_type_id)

273 assert quant_type in [

@@ -356,10 +357,7 @@ def fused_marlin_moe(

357 ).view(-1, topk, K)

358

359 if output is None:

359 - if inplace and not disable_inplace():

360 - output = hidden_states

361 - else:

362 - output = torch.empty_like(hidden_states)

360 + output = hidden_states if inplace else torch.empty_like(hidden_states)

361

362 if moe_sum is None:

363 return torch.sum(moe_output.view(-1, topk, K), dim=1, out=output)

vllm/model_executor/layers/fused_moe/fused_moe.py

+5 -21

@@ -27,9 +27,6 @@ from vllm.model_executor.layers.fused_moe.config import (

27 from vllm.model_executor.layers.fused_moe.moe_align_block_size import (

28 moe_align_block_size,

29 )

30 -from vllm.model_executor.layers.fused_moe.prepare_finalize import (

31 - MoEPrepareAndFinalizeNoEP,

32 -)

30 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (

31 TopKWeightAndReduceNoOP,

32 )

@@ -1511,7 +1508,7 @@ def torch_vllm_outplace_fused_experts(**kwargs) -> torch.Tensor:

1508

1509

1510 def dispatch_fused_experts_func(inplace: bool) -> Callable[..., torch.Tensor]:

1514 - if inplace and not disable_inplace():

1511 + if inplace:

1512 return torch_vllm_inplace_fused_experts

1513 return torch_vllm_outplace_fused_experts

1514

@@ -1534,6 +1531,8 @@ def fused_experts(

1531 if quant_config is None:

1532 quant_config = FUSED_MOE_UNQUANTIZED_CONFIG

1533

1534 + assert not inplace or not disable_inplace()

1535 +

1536 return dispatch_fused_experts_func(inplace)(

1537 hidden_states=hidden_states,

1538 w1=w1,

@@ -1593,7 +1592,7 @@ def fused_experts_impl(

1592 w2: torch.Tensor,

1593 topk_weights: torch.Tensor,

1594 topk_ids: torch.Tensor,

1596 - inplace: bool = False,

1595 + inplace: bool,

1596 activation: str = "silu",

1597 apply_router_weight_on_input: bool = False,

1598 use_fp8_w8a8: bool = False,

@@ -1712,10 +1711,7 @@ def fused_experts_impl(

1711 else:

1712 raise ValueError(f"Unsupported compute_type: {hidden_states.dtype}")

1713

1715 - if inplace and not disable_inplace():

1716 - out_hidden_states = hidden_states

1717 - else:

1718 - out_hidden_states = torch.empty_like(hidden_states)

1714 + out_hidden_states = hidden_states if inplace else torch.empty_like(hidden_states)

1715

1716 if ocp_mx_scheme is not None:

1717 # TODO: On platforms for which `current_platform.supports_mx()` is True

@@ -2291,15 +2287,3 @@ class TritonWNA16Experts(TritonExperts):

2287

2288 # separate function is required for MoE + LoRA

2289 self.moe_sum(intermediate_cache3, output)

2294 -

2295 -

2296 -def modular_triton_fused_moe(

2297 - moe_config: FusedMoEConfig,

2298 - quant_config: FusedMoEQuantConfig,

2299 - shared_experts: torch.nn.Module | None = None,

2300 -) -> mk.FusedMoEModularKernel:

2301 - return mk.FusedMoEModularKernel(

2302 - MoEPrepareAndFinalizeNoEP(),

2303 - TritonExperts(moe_config, quant_config),

2304 - shared_experts,

2305 - )

vllm/model_executor/layers/fused_moe/fused_moe_method_base.py

-4

@@ -113,10 +113,6 @@ class FusedMoEMethodBase(QuantizeMethodBase):

113 def supports_eplb(self) -> bool:

114 return False

115

116 - @property

117 - def allow_inplace(self) -> bool:

118 - return False

119 -

116 @property

117 def method_name(self) -> str:

118 return self.__class__.__name__

vllm/model_executor/layers/fused_moe/fused_moe_modular_method.py

+2 -5

@@ -46,6 +46,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

46 old_quant_method: FusedMoEMethodBase,

47 prepare_finalize: FusedMoEPrepareAndFinalize,

48 shared_experts: torch.nn.Module | None,

49 + inplace: bool = False,

50 ) -> "FusedMoEModularMethod":

51 return FusedMoEModularMethod(

52 old_quant_method,

@@ -54,6 +55,7 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

55 old_quant_method.select_gemm_impl(prepare_finalize, moe_layer),

56 shared_experts,

57 moe_parallel_config=moe_layer.moe_parallel_config,

58 + inplace=inplace,

59 ),

60 )

61

@@ -61,10 +63,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

63 def supports_eplb(self) -> bool:

64 return self.old_quant_method.supports_eplb

65

64 - @property

65 - def allow_inplace(self) -> bool:

66 - return self.old_quant_method.allow_inplace

67 -

66 @property

67 def method_name(self) -> str:

68 return self.old_quant_method.method_name

@@ -99,7 +97,6 @@ class FusedMoEModularMethod(FusedMoEMethodBase, CustomOp):

97 w2=layer.w2_weight,

98 topk_weights=topk_weights,

99 topk_ids=topk_ids,

102 - inplace=self.allow_inplace,

100 activation=layer.activation,

101 global_num_experts=layer.global_num_experts,

102 apply_router_weight_on_input=layer.apply_router_weight_on_input,

vllm/model_executor/layers/fused_moe/layer.py

+10 -1

@@ -50,6 +50,9 @@ from vllm.model_executor.layers.fused_moe.router.router_factory import (

50 from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (

51 UnquantizedFusedMoEMethod,

52 )

53 +from vllm.model_executor.layers.fused_moe.utils import (

54 + disable_inplace,

55 +)

56 from vllm.model_executor.layers.quantization.base_config import (

57 QuantizationConfig,

58 )

@@ -560,6 +563,8 @@ class FusedMoE(CustomOp):

563 activation=activation,

564 device=vllm_config.device_config.device,

565 routing_method=self.routing_method_type,

566 + # TODO: in_dtype == out_dtype?

567 + disable_inplace=disable_inplace() or self.shared_experts is not None,

568 )

569 if self.use_mori_kernels:

570 assert self.rocm_aiter_fmoe_enabled, (

@@ -650,7 +655,11 @@ class FusedMoE(CustomOp):

655 "%s for %s(%s)", prepare_finalize.__class__.__name__, self, id(self)

656 )

657 self.quant_method = FusedMoEModularMethod.make(

653 - self, self.quant_method, prepare_finalize, self.shared_experts

658 + self,

659 + self.quant_method,

660 + prepare_finalize,

661 + self.shared_experts,

662 + inplace=not self.moe_config.disable_inplace,

663 )

664

665 @property

vllm/model_executor/layers/fused_moe/modular_kernel.py

+5 -4

@@ -811,11 +811,13 @@ class FusedMoEModularKernel(torch.nn.Module):

811 fused_experts: FusedMoEPermuteExpertsUnpermute,

812 shared_experts: torch.nn.Module | None = None,

813 moe_parallel_config: FusedMoEParallelConfig | None = None,

814 + inplace: bool = False,

815 ):

816 super().__init__()

817 self.prepare_finalize = prepare_finalize

818 self.fused_experts = fused_experts

819 self.shared_experts = shared_experts

820 + self.inplace = inplace

821

822 # prefer an explicit FusedMoEParallelConfig when available (from

823 # FusedMoE layers / tests).

@@ -1292,7 +1294,6 @@ class FusedMoEModularKernel(torch.nn.Module):

1294 w2: torch.Tensor,

1295 topk_weights: torch.Tensor,

1296 topk_ids: torch.Tensor,

1295 - inplace: bool = False,

1297 activation: str = "silu",

1298 global_num_experts: int = -1,

1299 expert_map: torch.Tensor | None = None,

@@ -1309,8 +1310,6 @@ class FusedMoEModularKernel(torch.nn.Module):

1310 - topk_weights (torch.Tensor): The topk weights applied at the end of

1311 the layer.

1312 - topk_ids (torch.Tensor): A map of row to expert id.

1312 - - inplace (bool): If True, perform the operation in-place.

1313 - Defaults to False.

1313 - activation (str): The activation function to apply after the first

1314 MoE layer.

1315 - global_num_experts (int): The total number of experts in the global

@@ -1326,7 +1325,9 @@ class FusedMoEModularKernel(torch.nn.Module):

1325 - torch.Tensor: The output tensor after applying the MoE layer.

1326 """

1327

1329 - if inplace and self.shared_experts is None and not disable_inplace():

1328 + if self.inplace:

1329 + assert self.shared_experts is None

1330 + assert not disable_inplace()

1331 output = hidden_states

1332 else:

1333 output = torch.zeros_like(hidden_states)

vllm/model_executor/layers/fused_moe/oracle/fp8.py

+6 -4

@@ -472,7 +472,7 @@ def make_fp8_moe_kernel(

472 fp8_backend: Fp8MoeBackend,

473 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

474 shared_experts: torch.nn.Module | None = None,

475 -) -> tuple[mk.FusedMoEModularKernel, bool]:

475 +) -> mk.FusedMoEModularKernel:

476 # Create Prepare/Finalize.

477 prepare_finalize = maybe_make_prepare_finalize(

478 moe=moe_config,

@@ -512,8 +512,10 @@ def make_fp8_moe_kernel(

512 else None

513 ),

514 moe_parallel_config=moe_config.moe_parallel_config,

515 + inplace=(

516 + not moe_config.disable_inplace

517 + and fp8_backend != Fp8MoeBackend.FLASHINFER_CUTLASS

518 + ),

519 )

520

517 - # TODO(rob): update inplace logic to be part of the kernel.

518 - inplace = fp8_backend != Fp8MoeBackend.FLASHINFER_CUTLASS

519 - return kernel, inplace

521 + return kernel

vllm/model_executor/layers/fused_moe/oracle/nvfp4.py

+1

@@ -437,6 +437,7 @@ def make_nvfp4_moe_kernel(

437 else None

438 ),

439 moe_parallel_config=moe_config.moe_parallel_config,

440 + inplace=False,

441 )

442

443 # TODO(rob): update inplace logic to be part of the kernel.

vllm/model_executor/layers/fused_moe/oracle/unquantized.py

+8 -6

@@ -154,11 +154,9 @@ def make_unquantized_moe_kernel(

154 backend: UnquantizedMoeBackend,

155 quant_config: FusedMoEQuantConfig,

156 moe_config: FusedMoEConfig,

157 -) -> tuple[mk.FusedMoEModularKernel | None, bool]:

158 - use_inplace = True

159 -

157 +) -> mk.FusedMoEModularKernel | None:

158 if backend in UNSUPPORTED_BACKEND:

161 - return None, use_inplace

159 + return None

160

161 if backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS:

162 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (

@@ -171,8 +169,9 @@ def make_unquantized_moe_kernel(

169 moe_config=moe_config,

170 quant_config=quant_config,

171 ),

172 + inplace=False,

173 )

175 - use_inplace = False

174 +

175 elif backend == UnquantizedMoeBackend.AITER:

176 from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import (

177 AiterExperts,

@@ -184,6 +183,7 @@ def make_unquantized_moe_kernel(

183 moe_config=moe_config,

184 quant_config=quant_config,

185 ),

186 + inplace=not moe_config.disable_inplace,

187 )

188 elif backend == UnquantizedMoeBackend.TRITON:

189 from vllm.model_executor.layers.fused_moe import TritonExperts

@@ -194,6 +194,7 @@ def make_unquantized_moe_kernel(

194 moe_config=moe_config,

195 quant_config=quant_config,

196 ),

197 + inplace=not moe_config.disable_inplace,

198 )

199 elif backend == UnquantizedMoeBackend.XPU:

200 from vllm.model_executor.layers.fused_moe import XPUExperts

@@ -204,5 +205,6 @@ def make_unquantized_moe_kernel(

205 moe_config=moe_config,

206 quant_config=quant_config,

207 ),

208 + inplace=not moe_config.disable_inplace,

209 )

208 - return kernel, use_inplace

210 + return kernel

vllm/model_executor/layers/fused_moe/unquantized_fused_moe_method.py

+1 -6

@@ -101,10 +101,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

101 def supports_eplb(self) -> bool:

102 return True

103

104 - @property

105 - def allow_inplace(self) -> bool:

106 - return True

107 -

104 def maybe_make_prepare_finalize(

105 self,

106 routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None,

@@ -225,7 +221,7 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

221 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

222 assert self.moe_quant_config is not None

223

228 - self.kernel, self.use_inplace = make_unquantized_moe_kernel(

224 + self.kernel = make_unquantized_moe_kernel(

225 backend=self.unquantized_backend,

226 quant_config=self.moe_quant_config,

227 moe_config=self.moe,

@@ -329,7 +325,6 @@ class UnquantizedFusedMoEMethod(FusedMoEMethodBase, CustomOp):

325 w2=layer.w2_weight,

326 topk_weights=topk_weights,

327 topk_ids=topk_ids,

332 - inplace=self.use_inplace,

328 activation=layer.activation,

329 apply_router_weight_on_input=layer.apply_router_weight_on_input,

330 global_num_experts=layer.global_num_experts,

vllm/model_executor/layers/quantization/awq_marlin.py

+1

@@ -785,4 +785,5 @@ class AWQMarlinMoEMethod(FusedMoEMethodBase):

785 w2_zeros=layer.w2_qzeros,

786 workspace=layer.workspace,

787 input_dtype=self.input_dtype,

788 + inplace=not self.moe.disable_inplace,

789 )

vllm/model_executor/layers/quantization/bitsandbytes.py

+1 -1

@@ -515,7 +515,7 @@ class BitsAndBytesMoEMethod(FusedMoEMethodBase):

515 w2=w2,

516 topk_weights=topk_weights,

517 topk_ids=topk_ids,

518 - inplace=True,

518 + inplace=not self.moe.disable_inplace,

519 activation=layer.activation,

520 apply_router_weight_on_input=layer.apply_router_weight_on_input,

521 global_num_experts=layer.global_num_experts,

vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py

+5 -6

@@ -357,7 +357,6 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):

357 layer.w2_weight,

358 topk_weights,

359 topk_ids,

360 - inplace=False,

360 activation=layer.activation,

361 global_num_experts=layer.global_num_experts,

362 expert_map=layer.expert_map,

@@ -669,7 +668,6 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod):

668 layer.w2_weight,

669 topk_weights,

670 topk_ids,

672 - inplace=False,

671 activation=layer.activation,

672 global_num_experts=layer.global_num_experts,

673 expert_map=layer.expert_map,

@@ -960,7 +958,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

958 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

959 if self.moe_quant_config:

960 assert self.experts_cls is not None

963 - self.moe_mk, self.use_inplace = make_fp8_moe_kernel(

961 + self.moe_mk = make_fp8_moe_kernel(

962 moe_quant_config=self.moe_quant_config,

963 moe_config=self.moe,

964 fp8_backend=self.fp8_backend,

@@ -1073,7 +1071,6 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

1071 layer.w2_weight,

1072 topk_weights,

1073 topk_ids,

1076 - inplace=self.use_inplace,

1074 activation=layer.activation,

1075 global_num_experts=layer.global_num_experts,

1076 # TODO(rob): investigate the disable_expert_map introduced by:

@@ -1212,7 +1209,7 @@ class CompressedTensorsW8A8Int8MoEMethod(CompressedTensorsMoEMethod):

1209 w2=layer.w2_weight,

1210 topk_weights=topk_weights,

1211 topk_ids=topk_ids,

1215 - inplace=True,

1212 + inplace=not self.moe.disable_inplace,

1213 activation=layer.activation,

1214 apply_router_weight_on_input=layer.apply_router_weight_on_input,

1215 global_num_experts=layer.global_num_experts,

@@ -1739,6 +1736,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1736 workspace=layer.workspace,

1737 input_dtype=self.marlin_input_dtype,

1738 is_k_full=self.is_k_full,

1739 + inplace=not self.moe.disable_inplace,

1740 )

1741

1742

@@ -1969,7 +1967,7 @@ class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):

1967 layer.w2_weight_packed,

1968 topk_weights=topk_weights,

1969 topk_ids=topk_ids,

1972 - inplace=True,

1970 + inplace=not self.moe.disable_inplace,

1971 activation=layer.activation,

1972 apply_router_weight_on_input=layer.apply_router_weight_on_input,

1973 global_num_experts=layer.global_num_experts,

@@ -2605,6 +2603,7 @@ class CompressedTensorsW4A8Fp8MoEMethod(CompressedTensorsMoEMethod):

2603 s_strides1=self.s_strides1,

2604 s_strides2=self.s_strides2,

2605 group_size=self.group_size,

2606 + apply_router_weight_on_input=layer.apply_router_weight_on_input,

2607 )

2608

2609 @property

vllm/model_executor/layers/quantization/experts_int8.py

+1 -1

@@ -149,7 +149,7 @@ class ExpertsInt8MoEMethod(FusedMoEMethodBase):

149 layer.w2_weight,

150 topk_weights=topk_weights,

151 topk_ids=topk_ids,

152 - inplace=True,

152 + inplace=not self.moe.disable_inplace,

153 activation=layer.activation,

154 apply_router_weight_on_input=layer.apply_router_weight_on_input,

155 global_num_experts=layer.global_num_experts,

vllm/model_executor/layers/quantization/fp8.py

+1 -6

@@ -854,7 +854,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

854 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

855 if self.moe_quant_config:

856 assert self.experts_cls is not None

857 - self.moe_mk, self.use_inplace = make_fp8_moe_kernel(

857 + self.moe_mk = make_fp8_moe_kernel(

858 moe_quant_config=self.moe_quant_config,

859 moe_config=self.moe,

860 fp8_backend=self.fp8_backend,

@@ -958,10 +958,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

958 def supports_eplb(self) -> bool:

959 return True

960

961 - @property

962 - def allow_inplace(self) -> bool:

963 - return True

964 -

961 @property

962 def is_monolithic(self) -> bool:

963 return self.fp8_backend == Fp8MoeBackend.FLASHINFER_TRTLLM

@@ -1032,7 +1028,6 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1028 layer.w2_weight,

1029 topk_weights,

1030 topk_ids,

1035 - inplace=self.use_inplace,

1031 activation=layer.activation,

1032 global_num_experts=layer.global_num_experts,

1033 expert_map=layer.expert_map,

vllm/model_executor/layers/quantization/gptq_marlin.py

+1

@@ -924,4 +924,5 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

924 workspace=layer.workspace,

925 is_k_full=self.is_k_full,

926 input_dtype=self.input_dtype,

927 + inplace=not self.moe.disable_inplace,

928 )

vllm/model_executor/layers/quantization/modelopt.py

+1 -3

@@ -853,7 +853,7 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

853 self.moe_quant_config = self.get_fused_moe_quant_config(layer)

854 if self.moe_quant_config:

855 assert self.experts_cls is not None

856 - self.moe_mk, self.use_inplace = make_fp8_moe_kernel(

856 + self.moe_mk = make_fp8_moe_kernel(

857 moe_quant_config=self.moe_quant_config,

858 moe_config=self.moe,

859 fp8_backend=self.fp8_backend,

@@ -967,7 +967,6 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase):

967 layer.w2_weight,

968 topk_weights,

969 topk_ids,

970 - inplace=self.use_inplace,

970 activation=layer.activation,

971 global_num_experts=layer.global_num_experts,

972 expert_map=layer.expert_map,

@@ -1538,7 +1537,6 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1537 layer.w2_weight,

1538 topk_weights,

1539 topk_ids,

1541 - inplace=False,

1540 activation=layer.activation,

1541 global_num_experts=layer.global_num_experts,

1542 expert_map=layer.expert_map,

vllm/model_executor/layers/quantization/moe_wna16.py

+1 -1

@@ -378,7 +378,7 @@ class MoeWNA16Method(FusedMoEMethodBase):

378 layer.w2_qweight,

379 topk_weights=topk_weights,

380 topk_ids=topk_ids,

381 - inplace=True,

381 + inplace=not self.moe.disable_inplace,

382 apply_router_weight_on_input=layer.apply_router_weight_on_input,

383 global_num_experts=layer.global_num_experts,

384 expert_map=layer.expert_map,

vllm/model_executor/layers/quantization/mxfp4.py

+1 -4

@@ -881,10 +881,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

881 f"Incompatible Mxfp4 backend ({self.mxfp4_backend}) for EP"

882 )

883

884 - @property

885 - def allow_inplace(self) -> bool:

886 - return True

887 -

884 @property

885 def is_monolithic(self) -> bool:

886 return (

@@ -923,6 +919,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

919 activation=layer.activation,

920 expert_map=layer.expert_map,

921 input_dtype=self.marlin_input_dtype,

922 + inplace=not self.moe.disable_inplace,

923 )

924

925 assert _can_support_mxfp4(

vllm/model_executor/layers/quantization/quark/quark_moe.py

+3 -6

@@ -388,6 +388,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

388 apply_router_weight_on_input=layer.apply_router_weight_on_input,

389 global_num_experts=layer.global_num_experts,

390 expert_map=layer.expert_map,

391 + inplace=not self.moe.disable_inplace,

392 )

393 else:

394 from vllm.model_executor.layers.fused_moe import fused_experts

@@ -398,7 +399,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

399 w2=layer.w2_weight,

400 topk_weights=topk_weights,

401 topk_ids=topk_ids,

401 - inplace=True,

402 + inplace=not self.moe.disable_inplace,

403 activation=layer.activation,

404 apply_router_weight_on_input=layer.apply_router_weight_on_input,

405 global_num_experts=layer.global_num_experts,

@@ -734,10 +735,6 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

735 block_shape=None,

736 )

737

737 - @property

738 - def allow_inplace(self) -> bool:

739 - return True

740 -

738 def apply(

739 self,

740 layer: FusedMoE,

@@ -769,7 +766,7 @@ class QuarkOCP_MX_MoEMethod(QuarkMoEMethod):

766 layer.w2_weight,

767 topk_weights=topk_weights,

768 topk_ids=topk_ids,

772 - inplace=True,

769 + inplace=not self.moe.disable_inplace,

770 activation=layer.activation,

771 global_num_experts=layer.global_num_experts,

772 apply_router_weight_on_input=layer.apply_router_weight_on_input,

[Moe Refactor] Make Inplace Flag for FusedMoEModularKernel part of the constructor (#33375)