@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 8ae169286f

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 14, 05:30 UTC

[torch.compile] Unwrap fused_marlin_moe custom op (#26739)

Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com>

Varun Sundar Rabindranath committed Oct 13, 2025 at 22:22 UTC 8ae169286f93dfdf1c84f0a31eb2ee6a8debdbce

10 files changed +22 -52

tests/kernels/moe/test_moe.py

+3 -2

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

26 int4_w4a16_moe_quant_config,

27 int8_w8a16_moe_quant_config,

28 )

29 +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe

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

31 fused_topk,

32 modular_triton_fused_moe,

@@ -724,7 +725,7 @@ def test_fused_marlin_moe(

725 with set_current_vllm_config(vllm_config):

726 torch_output = torch_moe(a, w_ref1, w_ref2, score, topk, expert_map=e_map)

727

727 - marlin_output = torch.ops.vllm.fused_marlin_moe(

728 + marlin_output = fused_marlin_moe(

729 a,

730 qweight1,

731 qweight2,

@@ -837,7 +838,7 @@ def test_fused_marlin_moe_with_bias(m):

838 with set_current_vllm_config(vllm_config):

839 torch_output = torch_moe(a, w_ref1, w_ref2, score, topk, b_bias1, b_bias2)

840

840 - marlin_output = torch.ops.vllm.fused_marlin_moe(

841 + marlin_output = fused_marlin_moe(

842 a,

843 qweight1,

844 qweight2,

vllm/model_executor/layers/fused_moe/__init__.py

-1

@@ -51,7 +51,6 @@ __all__ = [

52 if HAS_TRITON:

53 # import to register the custom ops

54 - import vllm.model_executor.layers.fused_moe.fused_marlin_moe # noqa

54 from vllm.model_executor.layers.fused_moe.batched_deep_gemm_moe import (

55 BatchedDeepGemmExperts,

56 )

vllm/model_executor/layers/fused_moe/fused_marlin_moe.py

-39

@@ -19,7 +19,6 @@ from vllm.model_executor.layers.quantization.utils.marlin_utils import (

19 maybe_warn_marlin_atomic_add,

20 )

21 from vllm.scalar_type import ScalarType, scalar_types

22 -from vllm.utils import direct_register_custom_op

24 def fused_marlin_moe(

@@ -241,44 +240,6 @@ def fused_marlin_moe(

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

241

242

244 -def fused_marlin_moe_fake(

245 - hidden_states: torch.Tensor,

246 - w1: torch.Tensor,

247 - w2: torch.Tensor,

248 - w1_scale: torch.Tensor,

249 - w2_scale: torch.Tensor,

250 - gating_output: torch.Tensor | None,

251 - topk_weights: torch.Tensor,

252 - topk_ids: torch.Tensor,

253 - quant_type_id: int,

254 - apply_router_weight_on_input: bool = False,

255 - global_num_experts: int = -1,

256 - global_scale1: torch.Tensor | None = None,

257 - global_scale2: torch.Tensor | None = None,

258 - expert_map: torch.Tensor | None = None,

259 - g_idx1: torch.Tensor | None = None,

260 - g_idx2: torch.Tensor | None = None,

261 - sort_indices1: torch.Tensor | None = None,

262 - sort_indices2: torch.Tensor | None = None,

263 - w1_zeros: torch.Tensor | None = None,

264 - w2_zeros: torch.Tensor | None = None,

265 - workspace: torch.Tensor | None = None,

266 - intermediate_cache13: torch.Tensor | None = None,

267 - intermediate_cache2: torch.Tensor | None = None,

268 - is_k_full: bool = True,

269 - output: torch.Tensor | None = None,

270 - inplace: bool = False,

271 -) -> torch.Tensor:

272 - return torch.empty_like(hidden_states)

273 -

274 -

275 -direct_register_custom_op(

276 - op_name="fused_marlin_moe",

277 - op_func=fused_marlin_moe,

278 - fake_impl=fused_marlin_moe_fake,

279 -)

280 -

281 -

243 class MarlinExperts(mk.FusedMoEPermuteExpertsUnpermute):

244 def __init__(self, quant_config: FusedMoEQuantConfig):

245 # TODO (varun) : Enable activation quantization

vllm/model_executor/layers/quantization/awq_marlin.py

+2 -1

@@ -14,6 +14,7 @@ from vllm.model_executor.layers.fused_moe.config import (

14 FusedMoEConfig,

15 FusedMoEQuantConfig,

16 )

17 +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe

18 from vllm.model_executor.layers.fused_moe.layer import (

19 FusedMoE,

20 FusedMoEMethodBase,

@@ -604,7 +605,7 @@ class AWQMoEMethod(FusedMoEMethodBase):

605 indices_type=self.topk_indices_dtype,

606 )

607

607 - return torch.ops.vllm.fused_marlin_moe(

608 + return fused_marlin_moe(

609 x,

610 layer.w13_qweight,

611 layer.w2_qweight,

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

+4 -3

@@ -34,6 +34,7 @@ from vllm.model_executor.layers.fused_moe.cpu_fused_moe import select_experts

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

35 is_valid_flashinfer_cutlass_fused_moe,

36 )

37 +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe

38 from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compressed_tensors_wNa16 import ( # noqa

39 WNA16_SUPPORTED_BITS,

40 WNA16_SUPPORTED_TYPES_MAP,

@@ -462,7 +463,7 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):

463 #

464 if self.use_marlin:

465 assert self.fused_experts is None

465 - return torch.ops.vllm.fused_marlin_moe(

466 + return fused_marlin_moe(

467 x,

468 layer.w13_weight,

469 layer.w2_weight,

@@ -1067,7 +1068,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):

1068 if self.use_marlin:

1069 assert activation == "silu", f"{activation} not supported for Marlin MoE."

1070 assert self.fused_experts is None

1070 - return torch.ops.vllm.fused_marlin_moe(

1071 + return fused_marlin_moe(

1072 x,

1073 layer.w13_weight,

1074 layer.w2_weight,

@@ -1654,7 +1655,7 @@ class CompressedTensorsWNA16MarlinMoEMethod(CompressedTensorsMoEMethod):

1655 indices_type=self.topk_indices_dtype,

1656 )

1657

1657 - return torch.ops.vllm.fused_marlin_moe(

1658 + return fused_marlin_moe(

1659 x,

1660 layer.w13_weight_packed,

1661 layer.w2_weight_packed,

vllm/model_executor/layers/quantization/fp8.py

+2 -1

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

26 FusedMoEQuantConfig,

27 fp8_w8a8_moe_quant_config,

28 )

29 +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe

30 from vllm.model_executor.layers.fused_moe.layer import UnquantizedFusedMoEMethod

31 from vllm.model_executor.layers.linear import (

32 LinearBase,

@@ -1196,7 +1197,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):

1197 elif self.use_marlin:

1198 assert activation == "silu", f"{activation} not supported for Marlin MoE."

1199 assert self.fused_experts is None

1199 - result = torch.ops.vllm.fused_marlin_moe(

1200 + result = fused_marlin_moe(

1201 x,

1202 layer.w13_weight,

1203 layer.w2_weight,

vllm/model_executor/layers/quantization/gptq_marlin.py

+2 -1

@@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe.config import (

15 FusedMoEConfig,

16 FusedMoEQuantConfig,

17 )

18 +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe

19 from vllm.model_executor.layers.fused_moe.layer import (

20 FusedMoE,

21 FusedMoEMethodBase,

@@ -765,7 +766,7 @@ class GPTQMarlinMoEMethod(FusedMoEMethodBase):

766 indices_type=self.topk_indices_dtype,

767 )

768

768 - return torch.ops.vllm.fused_marlin_moe(

769 + return fused_marlin_moe(

770 x,

771 layer.w13_qweight,

772 layer.w2_qweight,

vllm/model_executor/layers/quantization/modelopt.py

+2 -1

@@ -21,6 +21,7 @@ from vllm.model_executor.layers.fused_moe.config import (

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

22 is_valid_flashinfer_cutlass_fused_moe,

23 )

24 +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe

25 from vllm.model_executor.layers.fused_moe.layer import (

26 FusedMoE,

27 FusedMoEMethodBase,

@@ -1701,7 +1702,7 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1702 #

1703 if self.use_marlin:

1704 assert self.fused_experts is None

1704 - return torch.ops.vllm.fused_marlin_moe(

1705 + return fused_marlin_moe(

1706 x,

1707 layer.w13_weight,

1708 layer.w2_weight,

vllm/model_executor/layers/quantization/mxfp4.py

+5 -2

@@ -21,7 +21,10 @@ from vllm.model_executor.layers.fused_moe.config import (

21 mxfp4_w4a16_moe_quant_config,

22 ocp_mx_moe_quant_config,

23 )

24 -from vllm.model_executor.layers.fused_moe.fused_marlin_moe import MarlinExperts

24 +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (

25 + MarlinExperts,

26 + fused_marlin_moe,

27 +)

28 from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (

29 OAITritonExperts,

30 )

@@ -947,7 +950,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):

950 e_score_correction_bias=e_score_correction_bias,

951 )

952

950 - return torch.ops.vllm.fused_marlin_moe(

953 + return fused_marlin_moe(

954 x,

955 layer.w13_weight,

956 layer.w2_weight,

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

+2 -1

@@ -20,6 +20,7 @@ from vllm.model_executor.layers.fused_moe.config import (

20 fp8_w8a8_moe_quant_config,

21 ocp_mx_moe_quant_config,

22 )

23 +from vllm.model_executor.layers.fused_moe.fused_marlin_moe import fused_marlin_moe

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

25 is_rocm_aiter_moe_enabled,

26 )

@@ -402,7 +403,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):

403 )

404 if self.use_marlin:

405 assert activation == "silu", f"{activation} not supported for Marlin MoE."

405 - return torch.ops.vllm.fused_marlin_moe(

406 + return fused_marlin_moe(

407 x,

408 layer.w13_weight,

409 layer.w2_weight,