@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 254a7f8fd6

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 11, 08:45 UTC

[Perf] Do FP4 quant before All gather on flashinfer trtllmgen MOE (#30014)

Signed-off-by: jiahanc <173873397+jiahanc@users.noreply.github.com>

jiahanc committed Dec 16, 2025 at 13:01 UTC 254a7f8fd613d6b6964abc277b73ca1f0b823cdb

9 files changed +165 -31

vllm/distributed/device_communicators/all2all.py

+24 -5

@@ -64,7 +64,12 @@ class NaiveAll2AllManager(All2AllManagerBase):

64 hidden_states: torch.Tensor,

65 router_logits: torch.Tensor,

66 is_sequence_parallel: bool = False,

67 + extra_tensors: list[torch.Tensor] | None = None,

68 ) -> tuple[torch.Tensor, torch.Tensor]:

69 + if extra_tensors is not None:

70 + raise NotImplementedError(

71 + "extra_tensors is not supported for NaiveAll2AllManager"

72 + )

73 sp_size = self.tp_group.world_size if is_sequence_parallel else 1

74 dp_metadata = get_forward_context().dp_metadata

75 assert dp_metadata is not None

@@ -76,6 +81,7 @@ class NaiveAll2AllManager(All2AllManagerBase):

81 router_logits = self.naive_multicast(

82 router_logits, cu_tokens_across_sp_cpu, is_sequence_parallel

83 )

84 +

85 return hidden_states, router_logits

87 def combine(

@@ -113,7 +119,11 @@ class AgRsAll2AllManager(All2AllManagerBase):

119 hidden_states: torch.Tensor,

120 router_logits: torch.Tensor,

121 is_sequence_parallel: bool = False,

116 - ) -> tuple[torch.Tensor, torch.Tensor]:

122 + extra_tensors: list[torch.Tensor] | None = None,

123 + ) -> (

124 + tuple[torch.Tensor, torch.Tensor]

125 + | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]

126 + ):

127 """

128 Gather hidden_states and router_logits from all dp ranks.

129 """

@@ -121,15 +131,22 @@ class AgRsAll2AllManager(All2AllManagerBase):

131 assert dp_metadata is not None

132 sizes = dp_metadata.get_chunk_sizes_across_dp_rank()

133 assert sizes is not None

124 -

134 dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()

135 assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]

127 - hidden_states, router_logits = dist_group.all_gatherv(

128 - [hidden_states, router_logits],

136 +

137 + tensors_to_gather = [hidden_states, router_logits]

138 + if extra_tensors is not None:

139 + tensors_to_gather.extend(extra_tensors)

140 +

141 + gathered_tensors = dist_group.all_gatherv(

142 + tensors_to_gather,

143 dim=0,

144 sizes=sizes,

145 )

132 - return hidden_states, router_logits

146 +

147 + if extra_tensors is not None:

148 + return (gathered_tensors[0], gathered_tensors[1], gathered_tensors[2:])

149 + return gathered_tensors[0], gathered_tensors[1]

150

151 def combine(

152 self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False

@@ -204,6 +221,7 @@ class PPLXAll2AllManager(All2AllManagerBase):

221 hidden_states: torch.Tensor,

222 router_logits: torch.Tensor,

223 is_sequence_parallel: bool = False,

224 + extra_tensors: list[torch.Tensor] | None = None,

225 ) -> tuple[torch.Tensor, torch.Tensor]:

226 raise NotImplementedError

227

@@ -251,6 +269,7 @@ class DeepEPAll2AllManagerBase(All2AllManagerBase):

269 hidden_states: torch.Tensor,

270 router_logits: torch.Tensor,

271 is_sequence_parallel: bool = False,

272 + extra_tensors: list[torch.Tensor] | None = None,

273 ) -> tuple[torch.Tensor, torch.Tensor]:

274 raise NotImplementedError

275

vllm/distributed/device_communicators/base_device_communicator.py

+6 -1

@@ -1,6 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import threading

4 +from typing import Any

5 from weakref import WeakValueDictionary

7 import torch

@@ -68,7 +69,11 @@ class All2AllManagerBase:

69 hidden_states: torch.Tensor,

70 router_logits: torch.Tensor,

71 is_sequence_parallel: bool = False,

71 - ):

72 + extra_tensors: list[torch.Tensor] | None = None,

73 + ) -> Any:

74 + # Subclasses should either:

75 + # - implement handling for extra_tensors, or

76 + # - raise a clear error if extra_tensors is not supported.

77 raise NotImplementedError

79 def set_num_sms(self, num_sms: int):

vllm/distributed/device_communicators/cuda_communicator.py

+11 -5

@@ -318,17 +318,23 @@ class CudaCommunicator(DeviceCommunicatorBase):

318

319 return output_list

320

321 - def dispatch(

321 + def dispatch( # type: ignore[override]

322 self,

323 hidden_states: torch.Tensor,

324 router_logits: torch.Tensor,

325 is_sequence_parallel: bool = False,

326 - ) -> tuple[torch.Tensor, torch.Tensor]:

326 + extra_tensors: list[torch.Tensor] | None = None,

327 + ) -> (

328 + tuple[torch.Tensor, torch.Tensor]

329 + | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]

330 + ):

331 assert self.all2all_manager is not None

328 - hidden_states, router_logits = self.all2all_manager.dispatch(

329 - hidden_states, router_logits, is_sequence_parallel

332 + return self.all2all_manager.dispatch(

333 + hidden_states,

334 + router_logits,

335 + is_sequence_parallel,

336 + extra_tensors, # type: ignore[call-arg]

337 )

331 - return hidden_states, router_logits

338

339 def combine(

340 self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False

vllm/distributed/parallel_state.py

+10 -3

@@ -1007,10 +1007,17 @@ class GroupCoordinator:

1007 hidden_states: torch.Tensor,

1008 router_logits: torch.Tensor,

1009 is_sequence_parallel: bool = False,

1010 - ) -> tuple[torch.Tensor, torch.Tensor]:

1010 + extra_tensors: list[torch.Tensor] | None = None,

1011 + ) -> (

1012 + tuple[torch.Tensor, torch.Tensor]

1013 + | tuple[torch.Tensor, torch.Tensor, list[torch.Tensor]]

1014 + ):

1015 if self.device_communicator is not None:

1012 - return self.device_communicator.dispatch(

1013 - hidden_states, router_logits, is_sequence_parallel

1016 + return self.device_communicator.dispatch( # type: ignore[call-arg]

1017 + hidden_states,

1018 + router_logits,

1019 + is_sequence_parallel,

1020 + extra_tensors,

1021 )

1022 else:

1023 return hidden_states, router_logits

vllm/model_executor/layers/fused_moe/fused_moe_method_base.py

+12

@@ -71,6 +71,18 @@ class FusedMoEMethodBase(QuantizeMethodBase):

71 "implementation based on the prepare_finalize"

72 )

74 + def prepare_dp_allgather_tensor(

75 + self,

76 + layer: "FusedMoE", # type: ignore[name-defined] # noqa: F821

77 + hidden_states: torch.Tensor,

78 + router_logits: torch.Tensor,

79 + ) -> tuple[torch.Tensor, list[torch.Tensor]]:

80 + """Hook to prepare tensors and extra tensors for DP allgather + EP dispatch."""

81 + raise NotImplementedError(

82 + "Method 'prepare_dp_allgather_tensor' is not implemented in "

83 + f"{self.__class__.__name__}."

84 + )

85 +

86 @abstractmethod

87 def get_fused_moe_quant_config(

88 self, layer: torch.nn.Module

vllm/model_executor/layers/fused_moe/layer.py

+39 -2

@@ -44,6 +44,7 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (

44 is_flashinfer_supporting_global_sf,

45 )

46 from vllm.platforms import current_platform

47 +from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe

48 from vllm.utils.math_utils import cdiv, round_up

49 from vllm.utils.torch_utils import (

50 aux_stream,

@@ -1933,10 +1934,46 @@ class FusedMoE(CustomOp):

1934 )

1935

1936 with sp_ctx:

1937 + extra_tensors = None

1938 if do_naive_dispatch_combine:

1937 - hidden_states_combined, router_logits = get_ep_group().dispatch(

1938 - hidden_states, router_logits, self.is_sequence_parallel

1939 + # Avoid circular import

1940 + from vllm.model_executor.layers.quantization.modelopt import (

1941 + ModelOptNvFp4FusedMoE,

1942 + )

1943 +

1944 + post_quant_allgather = (

1945 + has_flashinfer_trtllm_fused_moe()

1946 + and self.quant_method is not None

1947 + and self.dp_size > 1

1948 + and self.use_ep

1949 + and isinstance(self.quant_method, ModelOptNvFp4FusedMoE)

1950 + )

1951 + if post_quant_allgather:

1952 + hidden_states_to_dispatch, extra_tensors = (

1953 + self.quant_method.prepare_dp_allgather_tensor(

1954 + self, hidden_states, router_logits

1955 + )

1956 + )

1957 + else:

1958 + hidden_states_to_dispatch = hidden_states

1959 +

1960 + dispatch_res = get_ep_group().dispatch(

1961 + hidden_states_to_dispatch,

1962 + router_logits,

1963 + self.is_sequence_parallel,

1964 + extra_tensors=extra_tensors,

1965 )

1966 + if extra_tensors is not None:

1967 + hidden_states_combined, router_logits, extra_tensors_combined = (

1968 + dispatch_res

1969 + )

1970 + hidden_states_combined = (

1971 + hidden_states_combined,

1972 + extra_tensors_combined[0],

1973 + )

1974 + else:

1975 + hidden_states_combined, router_logits = dispatch_res

1976 +

1977 # Run shared experts before matrix multiply.

1978 # because matrix multiply maybe modify the hidden_states.

1979 if has_separate_shared_experts and not use_shared_experts_stream:

vllm/model_executor/layers/quantization/modelopt.py

+24 -1

@@ -1522,6 +1522,24 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1522 w2_blockscale_swizzled, requires_grad=False

1523 )

1524

1525 + def prepare_dp_allgather_tensor(

1526 + self,

1527 + layer: FusedMoE,

1528 + hidden_states: torch.Tensor,

1529 + router_logits: torch.Tensor,

1530 + ) -> tuple[torch.Tensor, list[torch.Tensor]]:

1531 + """Optionally prepare extra tensors to carry through DP allgather/EP."""

1532 + import flashinfer

1533 +

1534 + a1_gscale = layer.w13_input_scale_quant

1535 + hidden_states_fp4, hidden_states_sf = flashinfer.fp4_quantize(

1536 + hidden_states,

1537 + a1_gscale,

1538 + is_sf_swizzled_layout=False,

1539 + )

1540 + extra_tensors: list[torch.Tensor] = [hidden_states_sf]

1541 + return hidden_states_fp4, extra_tensors

1542 +

1543 def get_fused_moe_quant_config(

1544 self, layer: torch.nn.Module

1545 ) -> FusedMoEQuantConfig | None:

@@ -1576,8 +1594,13 @@ class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):

1594 e_score_correction_bias=layer.e_score_correction_bias,

1595 )

1596

1597 + # Hidden_states in select_experts is only used to extract metadata

1598 + if isinstance(x, tuple):

1599 + x_routing, _ = x

1600 + else:

1601 + x_routing = x

1602 topk_weights, topk_ids, _ = layer.select_experts(

1580 - hidden_states=x,

1603 + hidden_states=x_routing,

1604 router_logits=router_logits,

1605 )

1606

vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py

+22 -14

@@ -238,7 +238,7 @@ def prepare_static_weights_for_trtllm_fp4_moe(

238

239 def flashinfer_trtllm_fp4_moe(

240 layer: torch.nn.Module,

241 - x: torch.Tensor,

241 + x: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

242 router_logits: torch.Tensor,

243 top_k: int,

244 global_num_experts: int,

@@ -269,12 +269,16 @@ def flashinfer_trtllm_fp4_moe(

269 from vllm.model_executor.models.llama4 import Llama4MoE

270

271 # Quantize input to FP4

272 - a1_gscale = layer.w13_input_scale_quant

273 - (hidden_states_fp4, hidden_states_scale_linear_fp4) = flashinfer.fp4_quantize(

274 - x,

275 - a1_gscale,

276 - is_sf_swizzled_layout=False,

277 - )

272 + if isinstance(x, tuple):

273 + hidden_states_fp4, hidden_states_scale_linear_fp4 = x

274 + else:

275 + # hidden_states is the already quantized

276 + a1_gscale = layer.w13_input_scale_quant

277 + (hidden_states_fp4, hidden_states_scale_linear_fp4) = flashinfer.fp4_quantize(

278 + x,

279 + a1_gscale,

280 + is_sf_swizzled_layout=False,

281 + )

282

283 # Determine routing method type

284 use_llama4_routing = custom_routing_function is Llama4MoE.custom_routing_function

@@ -360,13 +364,17 @@ def flashinfer_trtllm_fp4_routed_moe(

364 torch.bfloat16

365 ).view(torch.int16)

366

363 - # Quantize input to FP4

364 - a1_gscale = layer.w13_input_scale_quant

365 - (hidden_states_fp4, hidden_states_scale_linear_fp4) = flashinfer.fp4_quantize(

366 - x,

367 - a1_gscale,

368 - is_sf_swizzled_layout=False,

369 - )

367 + if isinstance(x, tuple):

368 + # Hidden_states is the already quantized

369 + hidden_states_fp4, hidden_states_scale_linear_fp4 = x

370 + else:

371 + # Quantize input to FP4

372 + a1_gscale = layer.w13_input_scale_quant

373 + (hidden_states_fp4, hidden_states_scale_linear_fp4) = flashinfer.fp4_quantize(

374 + x,

375 + a1_gscale,

376 + is_sf_swizzled_layout=False,

377 + )

378

379 # Call TRT-LLM FP4 block-scale MoE kernel

380 out = flashinfer.fused_moe.trtllm_fp4_block_scale_routed_moe(

vllm/utils/flashinfer.py

+17

@@ -184,6 +184,23 @@ def has_flashinfer_cutedsl() -> bool:

184 )

185

186

187 +@functools.cache

188 +def has_flashinfer_trtllm_fused_moe() -> bool:

189 + """Return `True` if FlashInfer TRTLLM fused MoE is available."""

190 + if not has_flashinfer_moe():

191 + return False

192 + required_functions = [

193 + ("flashinfer.fused_moe", "trtllm_fp8_block_scale_moe"),

194 + ("flashinfer.fused_moe", "trtllm_fp8_per_tensor_scale_moe"),

195 + ("flashinfer.fused_moe", "trtllm_fp4_block_scale_moe"),

196 + ]

197 + for module_name, attr_name in required_functions:

198 + mod = _get_submodule(module_name)

199 + if not mod or not hasattr(mod, attr_name):

200 + return False

201 + return True

202 +

203 +

204 @functools.cache

205 def has_flashinfer_cutlass_fused_moe() -> bool:

206 """Return `True` if FlashInfer CUTLASS fused MoE is available."""