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 5, 04:45 UTC
[MoE] Move GPT OSS Triton kernel experts into fused_moe/experts/ (#39007)
Signed-off-by: Robert Shaw <robertgshaw2@gmail.com> Signed-off-by: Jackmin801 <ongjackm@gmail.com> Co-authored-by: Robert Shaw <robertgshaw2@gmail.com>
Jackmin801 committed
Apr 14, 2026 at 12:27 UTC
1a9353bb02e69de4631de6f7c0424aae5aa32f1a
8 files changed
+16
-12
docs/design/moe_kernel_features.md
+1
-1
@@ -86,7 +86,7 @@ To be used with a particular `FusedMoEPrepareAndFinalizeModular` subclass, MoE k
86
| cutlass_fp4 | standard,</br>batched | nvfp4 | A,T | silu | Y | Y | [`CutlassExpertsFp4`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp4] |
87
| cutlass_fp8 | standard,</br>batched | fp8 | A,T | silu, gelu | Y | Y | [`CutlassExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassExpertsFp8],</br>[`CutlasBatchedExpertsFp8`][vllm.model_executor.layers.fused_moe.cutlass_moe.CutlassBatchedExpertsFp8] |
88
| flashinfer | standard | nvfp4,</br>fp8 | T | <sup>5</sup> | N | Y | [`FlashInferExperts`][vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe.FlashInferExperts] |
89
-| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe.OAITritonExperts] |
89
+| gpt oss triton | standard | N/A | N/A | <sup>5</sup> | Y | Y | [`triton_kernel_fused_experts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.triton_kernel_fused_experts],</br>[`OAITritonExperts`][vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe.OAITritonExperts] |
90
| marlin | standard,</br>batched | <sup>3</sup> / N/A | <sup>3</sup> / N/A | silu,</br>swigluoai | Y | Y | [`fused_marlin_moe`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.fused_marlin_moe],</br>[`MarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.MarlinExperts],</br>[`BatchedMarlinExperts`][vllm.model_executor.layers.fused_moe.fused_marlin_moe.BatchedMarlinExperts] |
91
| trtllm | standard | mxfp4,</br>nvfp4 | G(16),G(32) | <sup>5</sup> | N | Y | [`TrtLlmMxfp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsMonolithic],</br>[`TrtLlmMxfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_mxfp4_moe.TrtLlmMxfp4ExpertsModular],</br>[`TrtLlmNvFp4ExpertsMonolithic`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsMonolithic],</br>[`TrtLlmNvfp4ExpertsModular`][vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe.TrtLlmNvFp4ExpertsModular] |
92
| rocm aiter moe | standard | mxfp4,</br>fp8 | G(32),G(128),A,T | silu, gelu,</br>swigluoai | Y | N | `rocm_aiter_fused_experts`,</br>`AiterExperts` |
tests/kernels/moe/test_gpt_oss_triton_kernels.py
+1
-1
@@ -25,7 +25,7 @@ from triton_kernels.tensor_details import layout
25
from triton_kernels.testing import assert_close
26
27
from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config
28
-from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
28
+from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (
29
triton_kernel_moe_forward,
30
)
31
from vllm.utils.math_utils import round_up
tests/kernels/moe/test_modular_oai_triton_moe.py
+1
-1
@@ -29,7 +29,7 @@ from vllm.model_executor.layers.fused_moe.all2all_utils import (
29
maybe_make_prepare_finalize,
30
)
31
from vllm.model_executor.layers.fused_moe.config import mxfp4_w4a16_moe_quant_config
32
-from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
32
+from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (
33
OAITritonExperts,
34
UnfusedOAITritonExperts,
35
)
tests/kernels/quantization/test_mxfp4_triton_ep.py
+7
-3
@@ -29,18 +29,22 @@ class TestTritonMoeForwardExpertMap:
29
torch.tensor([0, -1, 1, -1], device=device) if expert_map_present else None
30
)
31
32
+ from vllm.utils.import_utils import import_triton_kernels
33
+
34
+ import_triton_kernels()
35
+
36
with (
37
patch("triton_kernels.topk.topk") as mock_topk,
38
patch(
35
- "vllm.model_executor.layers.fused_moe."
39
+ "vllm.model_executor.layers.fused_moe.experts."
40
"gpt_oss_triton_kernels_moe.make_routing_data"
41
) as mock_make_routing,
42
patch(
39
- "vllm.model_executor.layers.fused_moe."
43
+ "vllm.model_executor.layers.fused_moe.experts."
44
"gpt_oss_triton_kernels_moe.triton_kernel_fused_experts"
45
) as mock_fused_experts,
46
):
43
- from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( # noqa: E501
47
+ from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501
48
triton_kernel_moe_forward,
49
)
50
vllm/lora/layers/fused_moe.py
+3
-3
@@ -19,6 +19,9 @@ from vllm.model_executor.layers.fused_moe import FusedMoE
19
from vllm.model_executor.layers.fused_moe.config import (
20
_get_config_dtype_str,
21
)
22
+from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import (
23
+ UnfusedOAITritonExperts,
24
+)
25
from vllm.model_executor.layers.fused_moe.fused_marlin_moe import (
26
MarlinExperts,
27
)
@@ -28,9 +31,6 @@ from vllm.model_executor.layers.fused_moe.fused_moe import (
31
from vllm.model_executor.layers.fused_moe.fused_moe_modular_method import (
32
FusedMoEModularMethod,
33
)
31
-from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
32
- UnfusedOAITritonExperts,
33
-)
34
from vllm.model_executor.layers.fused_moe.modular_kernel import (
35
FusedMoEKernel,
36
)
vllm/model_executor/layers/fused_moe/experts/gpt_oss_triton_kernels_moe.py
renamed
vllm/model_executor/layers/fused_moe/oracle/mxfp4.py
+2
-2
@@ -101,7 +101,7 @@ def backend_to_kernel_cls(
101
return [FlashInferExperts]
102
103
elif backend == Mxfp4MoeBackend.TRITON:
104
- from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
104
+ from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501
105
OAITritonExperts,
106
OAITritonMxfp4ExpertsMonolithic,
107
)
@@ -110,7 +110,7 @@ def backend_to_kernel_cls(
110
return [OAITritonMxfp4ExpertsMonolithic, OAITritonExperts]
111
112
elif backend == Mxfp4MoeBackend.TRITON_UNFUSED:
113
- from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
113
+ from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501
114
UnfusedOAITritonExperts,
115
)
116
vllm/model_executor/layers/quantization/quark/quark_moe.py
+1
-1
@@ -1591,7 +1591,7 @@ class QuarkOCP_MX_MoEMethod_OSS(QuarkOCP_MX_MoEMethod):
1591
"EPLB not supported for `QuarkW4MXFp4MoEMethod_OSS` yet."
1592
)
1593
1594
- from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( # noqa: E501
1594
+ from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501
1595
triton_kernel_moe_forward,
1596
)
1597