@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e8db44f883
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

[DP/EP][GPTOSS] Use triton matmul-ogs kernels for GPTOSS DP/EP (#24588)

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

Varun Sundar Rabindranath committed Sep 23, 2025 at 00:01 UTC e8db44f8834161259b25ebb13300bd502d37af9f
6 files changed +274 -75
vllm/model_executor/layers/fused_moe/config.py
+21 -1
@@ -288,7 +288,11 @@ class FusedMoEQuantConfig:
288
289 @property
290 def use_mxfp4_w4a4(self) -> bool:
291 - return self.quant_dtype == "mxfp4"
291 + return (self._a1.dtype == "mxfp4" and self._w1.dtype == "mxfp4")
292 +
293 + @property
294 + def use_mxfp4_w4a16(self) -> bool:
295 + return (self._a1.dtype is None and self._w1.dtype == "mxfp4")
296
297 @property
298 def use_nvfp4_w4a4(self) -> bool:
@@ -453,6 +457,22 @@ def int8_w8a8_moe_quant_config(
457 )
458
459
460 +def mxfp4_w4a16_moe_quant_config(
461 + w1_scale: Union[torch.Tensor, "PrecisionConfig"],
462 + w2_scale: Union[torch.Tensor, "PrecisionConfig"],
463 + w1_bias: Optional[torch.Tensor] = None,
464 + w2_bias: Optional[torch.Tensor] = None) -> FusedMoEQuantConfig:
465 + """
466 + Construct a quant config for unquantized activations and mxfp4 weights.
467 + """
468 + return FusedMoEQuantConfig(
469 + _a1=FusedMoEQuantDesc(),
470 + _a2=FusedMoEQuantDesc(),
471 + _w1=FusedMoEQuantDesc("mxfp4", None, w1_scale, None, None, w1_bias),
472 + _w2=FusedMoEQuantDesc("mxfp4", None, w2_scale, None, None, w2_bias),
473 + )
474 +
475 +
476 def mxfp4_w4a4_moe_quant_config(
477 w1_scale: Union[torch.Tensor, "PrecisionConfig"],
478 w2_scale: Union[torch.Tensor, "PrecisionConfig"],
vllm/model_executor/layers/fused_moe/deepep_ht_prepare_finalize.py
+18
@@ -11,6 +11,7 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
11 TopKWeightAndReduceContiguous, TopKWeightAndReduceDelegate)
12 from vllm.model_executor.layers.fused_moe.utils import (
13 moe_kernel_quantize_input)
14 +from vllm.utils import round_up
15
16
17 class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
@@ -18,6 +19,23 @@ class DeepEPHTPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
19 Prepare/Finalize using DeepEP High-Throughput kernels.
20 """
21
22 + @staticmethod
23 + def maybe_roundup_layer_hidden_size(hidden_size: int,
24 + dtype: torch.dtype) -> int:
25 + # Round up hidden size so it is compatible with DeepEP High Throughput
26 + # kernels.
27 + # DeepEP intranode kernels make copies in units of,
28 + # 32(warp-size) int4 elements. Round up hidden size to respect this.
29 + # For example, an input hidden size of 2880 with dtype torch.bfloat16
30 + # will be rounded up to 3072.
31 + hidden_size_bytes = hidden_size * dtype.itemsize
32 + xfer_atom_size = 512 # 32 * 16 (size(int4))
33 + if hidden_size_bytes % xfer_atom_size == 0:
34 + return hidden_size
35 +
36 + hidden_size_bytes = round_up(hidden_size_bytes, xfer_atom_size)
37 + return hidden_size_bytes // dtype.itemsize
38 +
39 def __init__(self, buffer: deep_ep.Buffer, num_dispatchers: int,
40 dp_size: int, rank_expert_offset: int):
41 super().__init__()
vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py
+140 -34
@@ -9,7 +9,8 @@ from vllm.logger import init_logger
9 from vllm.model_executor.layers.fused_moe.config import (
10 FUSED_MOE_UNQUANTIZED_CONFIG, FusedMoEQuantConfig)
11 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
12 - TopKWeightAndReduceDelegate)
12 + TopKWeightAndReduceNoOP)
13 +from vllm.triton_utils import tl, triton
14 from vllm.utils import has_triton_kernels
15
16 logger = init_logger(__name__)
@@ -19,13 +20,55 @@ if has_triton_kernels():
20 import triton_kernels.swiglu
21 from triton_kernels.matmul_ogs import (FnSpecs, FusedActivation,
22 matmul_ogs)
22 - from triton_kernels.routing import routing
23 + from triton_kernels.routing import (RoutingData, routing,
24 + routing_from_bitmatrix)
25 + from triton_kernels.tensor import Bitmatrix
26 except (ModuleNotFoundError, AttributeError) as e:
27 logger.error(
28 "Failed to import Triton kernels. Please make sure your triton "
29 "version is compatible. Error: %s", e)
30
31
32 +@triton.jit
33 +def pack_bitmatrix(
34 + bitmatrix,
35 + topk_ids,
36 + n_rows, # n_rows in bitmatrix / topk_ids
37 + bm_cols: tl.constexpr, # n int32_t bitpacks in bitmatrix
38 + n_expts_act, # num_topk
39 + BLOCK_SIZE_M: tl.constexpr,
40 + BLOCK_SIZE_K: tl.constexpr,
41 +):
42 + """
43 + Packs topk_ids into a bitmatrix.
44 + code reference:
45 + https://github.com/triton-lang/triton/blob/dd1bbc52b34d202dfe5ffea1e04fb16166c5c04e/python/triton_kernels/bench/distributed.py#L264
46 + """
47 + pid_m = tl.program_id(0)
48 + offsets_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
49 + offsets_k = tl.arange(0, BLOCK_SIZE_K)
50 + offsets = offsets_m[:, None] * n_expts_act + offsets_k[None, :]
51 + mask = (offsets_m < n_rows)[:, None] & (offsets_k < n_expts_act)[None, :]
52 + indices = tl.load(topk_ids + offsets, mask=mask, other=-1)
53 + div = indices // 32
54 + rem = indices % 32
55 + one = tl.cast(1, tl.uint32)
56 +
57 + # Iterate through all the relevant bitmatrix columns.
58 + for i in range(bm_cols):
59 + # When BLOCK_SIZE_K=32, offs is just the column index.
60 + offs = tl.arange(0, BLOCK_SIZE_K // 32) + i * (BLOCK_SIZE_K // 32)
61 + # All topks that need to go into this column has the correct bit set.
62 + # Other bits are 0. x is a 2D tensor.
63 + x = tl.where(div[:, :, None] == offs[None, None, :],
64 + (one << rem)[:, :, None], 0)
65 + # Reduce x to get a single int32_t bitpack.
66 + y = tl.reduce_or(x, axis=1)
67 + bitmatrix_ptrs = bitmatrix + offsets_m[:,
68 + None] * bm_cols + offs[None, :]
69 + tl.store(bitmatrix_ptrs, y, mask=offsets_m[:, None] < n_rows)
70 +
71 +
72 def triton_kernel_moe_forward(
73 hidden_states: torch.Tensor,
74 w1, # Tensor or triton_kernels.Tensor
@@ -124,34 +167,88 @@ def triton_kernel_fused_experts(
167 return intermediate_cache3
168
169
127 -class BatchedOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
170 +def make_routing_data(
171 + topk_ids: torch.Tensor,
172 + topk_weights: torch.Tensor,
173 + num_local_experts: int,
174 +) -> tuple["RoutingData", torch.Tensor, torch.Tensor]:
175 +
176 + topk_ids = topk_ids.to(torch.int16)
177 + topk_weights = topk_weights.to(torch.bfloat16)
178 +
179 + n_rows, num_topk = topk_ids.size()
180 +
181 + BLOCK_SIZE_M = 512
182 + BLOCK_SIZE_K = 32
183 +
184 + bm_cols = triton.cdiv(num_local_experts, BLOCK_SIZE_K) # n_bitpacks
185 + bitmatrix = torch.zeros((n_rows, bm_cols),
186 + dtype=torch.uint32,
187 + device=topk_ids.device)
188 +
189 + grid = (triton.cdiv(n_rows, BLOCK_SIZE_M), )
190 + pack_bitmatrix[grid](
191 + bitmatrix,
192 + topk_ids,
193 + n_rows,
194 + bm_cols,
195 + num_topk,
196 + BLOCK_SIZE_M=BLOCK_SIZE_M,
197 + BLOCK_SIZE_K=BLOCK_SIZE_K,
198 + )
199 +
200 + bitmatrix_shape = [n_rows, bm_cols * 32]
201 + bitmatrix_shape_max = [n_rows, None]
202 + bitmatrix = Bitmatrix(bitmatrix,
203 + shape=bitmatrix_shape,
204 + shape_max=bitmatrix_shape_max,
205 + scratchpad=None)
206 +
207 + # matmul_ogs expects invalid topk_weights to be -1s
208 + topk_weights = torch.where(topk_ids == -1, -1.0, topk_weights)
209 + routing_data, gather_indx, scatter_indx = routing_from_bitmatrix(
210 + bitmatrix, topk_weights, topk_ids, num_local_experts, num_topk)
211 +
212 + return routing_data, gather_indx, scatter_indx
213 +
214 +
215 +class BaseOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
216 +
217 + def __init__(self, quant_config: FusedMoEQuantConfig):
218 + super().__init__(quant_config)
219 +
220 + def supports_expert_map(self) -> bool:
221 + return True
222 +
223 + def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
224 + # Weight application and reduction happens in the fused_experts kernel.
225 + return TopKWeightAndReduceNoOP()
226
129 - def __init__(
227 + def _make_routing_data(
228 self,
131 - max_num_tokens: int,
132 - num_dispatchers: int,
133 - quant_config: FusedMoEQuantConfig,
134 - ):
229 + topk_ids: torch.Tensor,
230 + topk_weights: torch.Tensor,
231 + num_local_experts: int,
232 + ) -> tuple["RoutingData", torch.Tensor, torch.Tensor]:
233 + return make_routing_data(topk_ids, topk_weights, num_local_experts)
234 +
235 +
236 +class OAITritonExperts(BaseOAITritonExperts):
237 +
238 + def __init__(self, quant_config: FusedMoEQuantConfig):
239 + # TODO (varun) : Enable activation quantization
240 + assert quant_config.use_mxfp4_w4a16, "Supports only mxfp4_w4a16"
241 super().__init__(quant_config)
136 - self.max_num_tokens = max_num_tokens
137 - self.num_dispatchers = num_dispatchers
242
243 @property
244 def activation_formats(
245 self
246 ) -> tuple[mk.FusedMoEActivationFormat, mk.FusedMoEActivationFormat]:
143 - return (mk.FusedMoEActivationFormat.BatchedExperts,
144 - mk.FusedMoEActivationFormat.BatchedExperts)
247 + return (mk.FusedMoEActivationFormat.Standard,
248 + mk.FusedMoEActivationFormat.Standard)
249
250 def supports_chunking(self) -> bool:
147 - return False
148 -
149 - def supports_expert_map(self) -> bool:
150 - return False
151 -
152 - def finalize_weight_and_reduce_impl(self) -> mk.TopKWeightAndReduce:
153 - # Let PrepareAndFinalize::finalize() decide the impl.
154 - return TopKWeightAndReduceDelegate()
251 + return True
252
253 def workspace_shapes(
254 self, a: torch.Tensor, aq: torch.Tensor, M: int, N: int, K: int,
@@ -159,13 +256,10 @@ class BatchedOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
256 expert_tokens_meta: Optional[mk.ExpertTokensMetadata]
257 ) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], torch.dtype]:
258 # workspace are allocated inside the kernel
162 - assert a.dim() == 2
163 - num_dp = self.num_dispatchers
164 - num_experts = local_num_experts
165 - max_num_tokens = self.max_num_tokens
166 - workspace2 = (0, 0, 0)
167 - output = (num_experts, max_num_tokens * num_dp, N)
168 - return (output, workspace2, output, a.dtype)
259 + workspace1 = (M, K)
260 + workspace2 = (0, 0)
261 + output = (M, K)
262 + return (workspace1, workspace2, output, a.dtype)
263
264 def apply(
265 self,
@@ -185,17 +279,29 @@ class BatchedOAITritonExperts(mk.FusedMoEPermuteExpertsUnpermute):
279 expert_tokens_meta: Optional[mk.ExpertTokensMetadata],
280 apply_router_weight_on_input: bool,
281 ):
188 - return triton_kernel_fused_experts(
189 - output,
282 + if expert_map is not None:
283 + topk_ids = expert_map[topk_ids]
284 +
285 + local_num_experts = w1.size(0)
286 + if global_num_experts == -1:
287 + global_num_experts = local_num_experts
288 +
289 + routing_data, gather_indx, scatter_indx = self._make_routing_data(
290 + topk_ids, topk_weights, local_num_experts)
291 +
292 + experts_output = triton_kernel_fused_experts(
293 + None,
294 hidden_states,
295 w1,
296 w2,
193 - routing_data=None,
194 - gather_indx=None,
195 - scatter_indx=None,
297 + routing_data,
298 + gather_indx,
299 + scatter_indx,
300 activation=activation,
301 quant_config=self.quant_config,
302 apply_router_weight_on_input=False,
199 - global_num_experts=global_num_experts,
200 - expert_map=expert_map,
303 + global_num_experts=local_num_experts,
304 + expert_map=None, # applied already
305 a1q_scale=a1q_scale)
306 +
307 + output.copy_(experts_output, non_blocking=True)
vllm/model_executor/layers/fused_moe/layer.py
+60 -22
@@ -800,6 +800,49 @@ def get_compressed_expert_map(expert_map: torch.Tensor) -> str:
800 for local_index, global_index in zip(local_indices, global_indices))
801
802
803 +def maybe_roundup_hidden_size(
804 + hidden_size: int, act_dtype: torch.dtype,
805 + quant_config: Optional[QuantizationConfig],
806 + moe_parallel_config: FusedMoEParallelConfig) -> int:
807 + """
808 + Given layer hidden size and MoE configurations, round up hidden_size
809 + if necessary.
810 +
811 + Args:
812 + hidden_size(int): Layer hidden-size
813 + act_dtype: Data type of the layer activations.
814 + quant_config(FusedMoEQuantConfig): Fused MoE quantization configuration.
815 + moe_parallel_config(FusedMoEParallelConfig): Fused MoE parallelization
816 + strategy configuration.
817 +
818 + Return:
819 + Rounded up hidden_size if rounding up is required based on the configs.
820 + Original hidden size otherwise.
821 + """
822 +
823 + if (moe_parallel_config.use_deepep_ht_kernels):
824 + hidden_size = (
825 + DeepEPHTPrepareAndFinalize.maybe_roundup_layer_hidden_size(
826 + hidden_size, act_dtype))
827 +
828 + # we are padding globally so EP buffer allocation works
829 + if quant_config and quant_config.get_name() == "mxfp4":
830 +
831 + from vllm.model_executor.layers.quantization.mxfp4 import (
832 + Mxfp4Backend, get_mxfp4_backend)
833 + current_mxfp4_backend = get_mxfp4_backend()
834 + if (current_mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16
835 + or current_mxfp4_backend
836 + == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS):
837 + hidden_size = round_up(hidden_size, 128)
838 + elif (current_platform.is_rocm() or current_mxfp4_backend
839 + == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
840 + or current_mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16):
841 + hidden_size = round_up(hidden_size, 256)
842 +
843 + return hidden_size
844 +
845 +
846 @CustomOp.register("fused_moe")
847 class FusedMoE(CustomOp):
848 """FusedMoE layer for MoE models.
@@ -856,6 +899,18 @@ class FusedMoE(CustomOp):
899 params_dtype = torch.get_default_dtype()
900 self.params_dtype = params_dtype
901
902 + vllm_config = get_current_vllm_config()
903 +
904 + # FIXME (varun): We should have a better way of inferring the activation
905 + # datatype. This works for now as the tensor datatype entering the MoE
906 + # operation is typically unquantized (i.e. float16/bfloat16).
907 + if vllm_config.model_config is not None:
908 + moe_in_dtype = vllm_config.model_config.dtype
909 + else:
910 + # TODO (bnell): This is a hack to get test_mixtral_moe to work
911 + # since model_config is not set in the pytest test.
912 + moe_in_dtype = params_dtype
913 +
914 tp_size_ = (tp_size if tp_size is not None else
915 get_tensor_model_parallel_world_size())
916 dp_size_ = (dp_size
@@ -865,7 +920,6 @@ class FusedMoE(CustomOp):
920 if self.is_sequence_parallel:
921 self.sp_size = tp_size_
922
868 - vllm_config = get_current_vllm_config()
923 self.moe_parallel_config: FusedMoEParallelConfig = (
924 FusedMoEParallelConfig.make(
925 tp_size_=tp_size_,
@@ -874,19 +928,10 @@ class FusedMoE(CustomOp):
928
929 self.global_num_experts = num_experts + num_redundant_experts
930
877 - # we are padding globally so EP buffer allocation works
878 - if quant_config and quant_config.get_name() == "mxfp4":
879 - from vllm.model_executor.layers.quantization.mxfp4 import (
880 - Mxfp4Backend, get_mxfp4_backend)
881 - current_mxfp4_backend = get_mxfp4_backend()
882 - if (current_mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16
883 - or current_mxfp4_backend
884 - == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS):
885 - hidden_size = round_up(hidden_size, 128)
886 - elif (current_platform.is_rocm() or current_mxfp4_backend
887 - == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM or
888 - current_mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16):
889 - hidden_size = round_up(hidden_size, 256)
931 + # Round up hidden size if needed.
932 + hidden_size = maybe_roundup_hidden_size(hidden_size, moe_in_dtype,
933 + quant_config,
934 + self.moe_parallel_config)
935
936 # For smuggling this layer into the fused moe custom op
937 compilation_config = vllm_config.compilation_config
@@ -967,20 +1012,13 @@ class FusedMoE(CustomOp):
1012 raise ValueError("Only softmax scoring function is supported for "
1013 "non-grouped topk.")
1014
970 - if vllm_config.model_config is not None:
971 - model_dtype = vllm_config.model_config.dtype
972 - else:
973 - # TODO (bnell): This is a hack to get test_mixtral_moe to work
974 - # since model_config is not set in the pytest test.
975 - model_dtype = params_dtype
976 -
1015 moe = FusedMoEConfig(
1016 num_experts=self.global_num_experts,
1017 experts_per_token=top_k,
1018 hidden_dim=hidden_size,
1019 num_local_experts=self.local_num_experts,
1020 moe_parallel_config=self.moe_parallel_config,
983 - in_dtype=model_dtype,
1021 + in_dtype=moe_in_dtype,
1022 max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE,
1023 has_bias=has_bias,
1024 )
vllm/model_executor/layers/fused_moe/modular_kernel.py
+1 -1
@@ -76,7 +76,7 @@ def _moe_problem_size(
76 """
77 assert w1.dim() == 3 and w2.dim() == 3
78 E, N, _ = w1.size()
79 - K = w2.size(1)
79 + K = a1.size(-1)
80
81 if a1.dim() == 2:
82 # Make sure we are using the correct a1 (pre-permute).
vllm/model_executor/layers/quantization/mxfp4.py
+34 -17
@@ -13,7 +13,10 @@ from vllm.model_executor.layers.fused_moe import (FusedMoE, FusedMoEConfig,
13 FusedMoEMethodBase)
14 from vllm.model_executor.layers.fused_moe import modular_kernel as mk
15 from vllm.model_executor.layers.fused_moe.config import (
16 - FusedMoEQuantConfig, mxfp4_w4a4_moe_quant_config)
16 + FusedMoEQuantConfig, mxfp4_w4a4_moe_quant_config,
17 + mxfp4_w4a16_moe_quant_config)
18 +from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import (
19 + OAITritonExperts)
20 from vllm.model_executor.layers.fused_moe.trtllm_moe import TrtLlmGenExperts
21 from vllm.model_executor.layers.linear import (LinearBase,
22 UnquantizedLinearMethod)
@@ -578,9 +581,14 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
581 layer.w13_bias = Parameter(w13_bias, requires_grad=False)
582 layer.w2_bias = Parameter(w2_bias, requires_grad=False)
583
581 - # FIXME warp need to be adjusted based on batch size
582 - # only apply to batched mode
583 - if self.moe.use_ep:
584 + # Ideally we'd use FusedMoEModularKernel.prepare_finalize object
585 + # (stored in self.fused_experts) to determine if the MoE has a
586 + # batched activation format. As self.fused_experts is not
587 + # initialized at this point, we resort to checking the MoE config
588 + # directly.
589 + is_batched_moe = (self.moe.use_pplx_kernels
590 + or self.moe.use_deepep_ll_kernels)
591 + if is_batched_moe:
592 num_warps = 4 if envs.VLLM_MOE_DP_CHUNK_SIZE <= 512 else 8
593 else:
594 num_warps = 8
@@ -640,16 +648,21 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
648 if self.mxfp4_backend == Mxfp4Backend.TRITON:
649 w1_scale = self.w13_precision_config
650 w2_scale = self.w2_precision_config
651 + return mxfp4_w4a16_moe_quant_config(
652 + w1_bias=layer.w13_bias,
653 + w2_bias=layer.w2_bias,
654 + w1_scale=w1_scale,
655 + w2_scale=w2_scale,
656 + )
657 else:
658 w1_scale = layer.w13_weight_scale
659 w2_scale = layer.w2_weight_scale
646 -
647 - return mxfp4_w4a4_moe_quant_config(
648 - w1_bias=layer.w13_bias,
649 - w2_bias=layer.w2_bias,
650 - w1_scale=w1_scale,
651 - w2_scale=w2_scale,
652 - )
660 + return mxfp4_w4a4_moe_quant_config(
661 + w1_bias=layer.w13_bias,
662 + w2_bias=layer.w2_bias,
663 + w1_scale=w1_scale,
664 + w2_scale=w2_scale,
665 + )
666
667 def select_gemm_impl(
668 self,
@@ -661,6 +674,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
674 raise NotImplementedError(
675 "Mxfp4 does not support batched experts format for EP")
676 else:
677 + assert self.moe_quant_config is not None
678 if (self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM
679 or self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_BF16):
680 # B200 code-path
@@ -671,13 +685,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
685 # TODO(bnell): part of quant_config
686 "max_capture_size": self.max_capture_size,
687 }
674 - assert self.moe_quant_config is not None
688 return TrtLlmGenExperts(self.moe, self.moe_quant_config,
689 **kwargs)
690 else:
678 - # Use matmul_ogs from triton_kernels here!
679 - raise NotImplementedError(
680 - "Mxfp4 does not support non-batched experts format for EP")
691 + return OAITritonExperts(self.moe_quant_config)
692
693 def _route_and_experts(
694 self,
@@ -722,10 +733,16 @@ class Mxfp4MoEMethod(FusedMoEMethodBase):
733 logical_to_physical_map=logical_to_physical_map,
734 logical_replica_count=logical_replica_count)
735
736 + w13_weight = (self.w13_weight_triton_tensor
737 + if layer.w13_weight is None else layer.w13_weight)
738 + w2_weight = (self.w2_weight_triton_tensor
739 + if layer.w2_weight is None else layer.w2_weight)
740 + assert all([w is not None for w in [w13_weight, w2_weight]])
741 +
742 return self.fused_experts(
743 hidden_states=x,
727 - w1=layer.w13_weight,
728 - w2=layer.w2_weight,
744 + w1=w13_weight,
745 + w2=w2_weight,
746 topk_weights=topk_weights,
747 topk_ids=topk_ids,
748 inplace=True,