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

[Bugfix] Fix accuracy issue when using flashinfer cutlass moe, TP=1 and modelopt. (#23125)

Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>

bnellnm committed Aug 19, 2025 at 14:00 UTC b94faf9d50360217659f5605fce45a562dde6834
4 files changed +135 -35
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+49
@@ -7,6 +7,8 @@ import torch
7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 from vllm.logger import init_logger
9 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
10 +from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
11 + FlashInferCutlassMoEPrepareAndFinalize)
12 from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
13 TopKWeightAndReduceNoOP)
14 from vllm.utils.flashinfer import (flashinfer_cutlass_fused_moe,
@@ -181,3 +183,50 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
183 ep_rank=self.ep_rank,
184 output=output,
185 )
186 +
187 +
188 +def flashinfer_cutlass_moe_fp4(
189 + hidden_states: torch.Tensor,
190 + w1: torch.Tensor,
191 + w2: torch.Tensor,
192 + topk_weights: torch.Tensor,
193 + topk_ids: torch.Tensor,
194 + w1_scale: torch.Tensor,
195 + w2_scale: torch.Tensor,
196 + g1_alphas: torch.Tensor,
197 + g2_alphas: torch.Tensor,
198 + a1_gscale: torch.Tensor,
199 + a2_gscale: torch.Tensor,
200 + inplace: bool = False,
201 + activation: str = "silu",
202 + global_num_experts: int = -1,
203 + expert_map: Optional[torch.Tensor] = None,
204 + apply_router_weight_on_input: bool = False,
205 +) -> torch.Tensor:
206 +
207 + fused_experts = mk.FusedMoEModularKernel(
208 + FlashInferCutlassMoEPrepareAndFinalize(use_dp=False,
209 + a1_gscale=a1_gscale),
210 + FlashInferExperts(
211 + g1_alphas=g1_alphas,
212 + g2_alphas=g2_alphas,
213 + a1_gscale=a1_gscale,
214 + a2_gscale=a2_gscale,
215 + out_dtype=hidden_states.dtype,
216 + quant_dtype="nvfp4",
217 + ))
218 +
219 + return fused_experts(
220 + hidden_states=hidden_states,
221 + w1=w1,
222 + w2=w2,
223 + topk_weights=topk_weights,
224 + topk_ids=topk_ids,
225 + inplace=inplace,
226 + activation=activation,
227 + global_num_experts=global_num_experts,
228 + expert_map=expert_map,
229 + w1_scale=w1_scale,
230 + w2_scale=w2_scale,
231 + apply_router_weight_on_input=apply_router_weight_on_input,
232 + )
vllm/model_executor/layers/fused_moe/layer.py
+2
@@ -198,6 +198,8 @@ class FusedMoEMethodBase(QuantizeMethodBase):
198 else:
199 return None
200
201 + # Note: init_prepare_finalize should only be called by
202 + # prepare_communication_buffer_for_model.
203 def init_prepare_finalize(self):
204 assert self.moe is not None
205 prepare_finalize = self.maybe_make_prepare_finalize(self.moe)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+27
@@ -388,6 +388,33 @@ class CompressedTensorsW4A4MoeMethod(CompressedTensorsMoEMethod):
388 apply_router_weight_on_input=apply_router_weight_on_input,
389 )
390
391 + elif self.allow_flashinfer:
392 + from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
393 + flashinfer_cutlass_moe_fp4)
394 +
395 + assert is_valid_flashinfer_cutlass_fused_moe(
396 + x, layer.w13_weight, layer.w2_weight), (
397 + "Flashinfer CUTLASS Fused MoE not applicable!")
398 +
399 + return flashinfer_cutlass_moe_fp4(
400 + hidden_states=x,
401 + w1=layer.w13_weight,
402 + w2=layer.w2_weight,
403 + topk_weights=topk_weights,
404 + topk_ids=topk_ids,
405 + inplace=False, # TODO(shuw): fix later, now output is high prec
406 + activation=activation,
407 + global_num_experts=global_num_experts,
408 + expert_map=expert_map,
409 + w1_scale=layer.w13_blockscale_swizzled,
410 + w2_scale=layer.w2_blockscale_swizzled,
411 + g1_alphas=layer.g1_alphas,
412 + g2_alphas=layer.g2_alphas,
413 + a1_gscale=layer.w13_input_scale_quant,
414 + a2_gscale=layer.w2_input_scale_quant,
415 + apply_router_weight_on_input=apply_router_weight_on_input,
416 + )
417 +
418 assert expert_map is None, ("Expert Parallelism / expert_map "
419 "is currently not supported for "
420 "CompressedTensorsW4A4MoeMethod.")
vllm/model_executor/layers/quantization/modelopt.py
+57 -35
@@ -966,22 +966,21 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
966 f"Unknown flashinfer moe backend: {flashinfer_moe_backend}"
967 f" expected one of {allowed_backends}")
968
969 - self.fused_experts: Optional[
970 - mk.FusedMoEModularKernel] = None # type: ignore[assignment]
971 -
969 def maybe_make_prepare_finalize(
970 self,
971 moe: FusedMoEConfig,
972 ) -> Optional[mk.FusedMoEPrepareAndFinalize]:
976 - if not self.allow_flashinfer:
977 - return super().maybe_make_prepare_finalize(moe)
978 -
979 - prepare_finalize = build_flashinfer_fp4_cutlass_moe_prepare_finalize(
980 - moe,
981 - a1_gscale=self.layer.w13_input_scale_quant,
982 - )
983 - logger.debug_once("%s", prepare_finalize.__class__.__name__)
984 - return prepare_finalize
973 + if (self.allow_flashinfer and self.flashinfer_moe_backend
974 + == FlashinferMoeBackend.CUTLASS):
975 + prepare_finalize = (
976 + build_flashinfer_fp4_cutlass_moe_prepare_finalize(
977 + moe,
978 + a1_gscale=self.layer.w13_input_scale_quant,
979 + ))
980 + logger.debug_once("%s", prepare_finalize.__class__.__name__)
981 + return prepare_finalize
982 +
983 + return super().maybe_make_prepare_finalize(moe)
984
985 def select_gemm_impl(
986 self,
@@ -1409,7 +1408,52 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1408 global_num_experts=global_num_experts,
1409 expert_map=expert_map)
1410
1412 - if self.fused_experts is None:
1411 + if self.fused_experts is not None:
1412 + assert self.allow_flashinfer and \
1413 + self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS
1414 +
1415 + assert is_valid_flashinfer_cutlass_fused_moe(
1416 + x, layer.w13_weight, layer.w2_weight), (
1417 + "Flashinfer CUTLASS Fused MoE not applicable!")
1418 +
1419 + out = self.fused_experts(
1420 + hidden_states=x,
1421 + w1=layer.w13_weight,
1422 + w2=layer.w2_weight,
1423 + topk_weights=topk_weights,
1424 + topk_ids=topk_ids,
1425 + inplace=False, # TODO(shuw): fix later, now output is high prec
1426 + activation=activation,
1427 + global_num_experts=global_num_experts,
1428 + expert_map=expert_map,
1429 + w1_scale=layer.w13_blockscale_swizzled,
1430 + w2_scale=layer.w2_blockscale_swizzled,
1431 + apply_router_weight_on_input=apply_router_weight_on_input,
1432 + )
1433 + elif (self.allow_flashinfer
1434 + and self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS):
1435 + from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
1436 + flashinfer_cutlass_moe_fp4)
1437 +
1438 + out = flashinfer_cutlass_moe_fp4(
1439 + hidden_states=x,
1440 + w1=layer.w13_weight,
1441 + w2=layer.w2_weight,
1442 + topk_weights=topk_weights,
1443 + topk_ids=topk_ids,
1444 + w1_scale=layer.w13_blockscale_swizzled,
1445 + w2_scale=layer.w2_blockscale_swizzled,
1446 + g1_alphas=layer.g1_alphas,
1447 + g2_alphas=layer.g2_alphas,
1448 + a1_gscale=layer.w13_input_scale_quant,
1449 + a2_gscale=layer.w2_input_scale_quant,
1450 + inplace=False, # TODO(shuw): fix later, now output is high prec
1451 + activation=activation,
1452 + global_num_experts=global_num_experts,
1453 + expert_map=expert_map,
1454 + apply_router_weight_on_input=apply_router_weight_on_input,
1455 + )
1456 + else:
1457 # If no modular kernel is provided, use cutlass_moe_fp4 for TP case
1458 # only (no EP).
1459 from vllm.model_executor.layers.fused_moe.cutlass_moe import (
@@ -1432,27 +1476,5 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
1476 e=layer.w13_weight.shape[0],
1477 expert_map=expert_map,
1478 apply_router_weight_on_input=apply_router_weight_on_input)
1435 - else:
1436 - assert self.allow_flashinfer and \
1437 - self.flashinfer_moe_backend == FlashinferMoeBackend.CUTLASS
1438 -
1439 - assert is_valid_flashinfer_cutlass_fused_moe(
1440 - x, layer.w13_weight, layer.w2_weight), (
1441 - "Flashinfer CUTLASS Fused MoE not applicable!")
1442 -
1443 - out = self.fused_experts(
1444 - hidden_states=x,
1445 - w1=layer.w13_weight,
1446 - w2=layer.w2_weight,
1447 - topk_weights=topk_weights,
1448 - topk_ids=topk_ids,
1449 - inplace=False, # TODO(shuw): fix later, now output is high prec
1450 - activation=activation,
1451 - global_num_experts=global_num_experts,
1452 - expert_map=expert_map,
1453 - w1_scale=layer.w13_blockscale_swizzled,
1454 - w2_scale=layer.w2_blockscale_swizzled,
1455 - apply_router_weight_on_input=apply_router_weight_on_input,
1456 - )
1479
1480 return out