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

Support mnnvl all2allv from Flashinfer (#21003)

Signed-off-by: Shu Wang <shuw@nvidia.com> Signed-off-by: Shu Wang. <shuw@nvidia.com> Signed-off-by: Tyler Michael Smith <tyler@neuralmagic.com> Signed-off-by: Tyler Michael Smith <tlrmchlsmth@gmail.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com> Co-authored-by: Tyler Michael Smith <tlrmchlsmth@gmail.com>

Shu Wang committed Sep 24, 2025 at 13:38 UTC 54e42b72dbf7ec259de1c28e66ea9575fe594685
10 files changed +410 -40
tests/kernels/moe/modular_kernel_tools/mk_objects.py
+3 -2
@@ -222,7 +222,8 @@ if (has_flashinfer_cutlass_fused_moe()
222 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( # noqa: E501
223 FlashInferExperts)
224 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
225 - FlashInferCutlassMoEPrepareAndFinalize)
225 + FlashInferCutlassMoEPrepareAndFinalize,
226 + create_flashinfer_prepare_finalize)
227
228 register_prepare_and_finalize(
229 FlashInferCutlassMoEPrepareAndFinalize,
@@ -373,7 +374,7 @@ def make_prepare_finalize(
374 assert prepare_finalize is not None
375 return prepare_finalize
376 elif prepare_finalize_type == FlashInferCutlassMoEPrepareAndFinalize:
376 - return FlashInferCutlassMoEPrepareAndFinalize(
377 + return create_flashinfer_prepare_finalize(
378 use_dp=moe.moe_parallel_config.dp_size > 1)
379 else:
380 return MoEPrepareAndFinalizeNoEP()
vllm/distributed/device_communicators/all2all.py
+110 -15
@@ -10,9 +10,15 @@ from vllm.distributed import get_dp_group
10 from vllm.forward_context import get_forward_context
11 from vllm.logger import init_logger
12 from vllm.utils import has_deep_ep, has_pplx
13 +from vllm.utils.flashinfer import has_flashinfer_all2all
14
15 from .base_device_communicator import All2AllManagerBase, Cache
16
17 +if has_flashinfer_all2all():
18 + from flashinfer.comm import Mapping
19 + from flashinfer.comm.mnnvl import MnnvlConfig
20 + from flashinfer.comm.trtllm_alltoall import MnnvlMoe
21 +
22 logger = init_logger(__name__)
23
24
@@ -47,24 +53,22 @@ class NaiveAll2AllManager(All2AllManagerBase):
53
54 def dispatch(self, hidden_states: torch.Tensor,
55 router_logits: torch.Tensor):
50 - cu_tokens_across_dp_cpu = get_forward_context(
51 - ).dp_metadata.cu_tokens_across_dp_cpu
56 + sizes = get_forward_context(
57 + ).dp_metadata.get_chunk_sizes_across_dp_rank()
58 + hidden_states, router_logits = get_dp_group().all_gatherv(
59 + [hidden_states, router_logits],
60 + dim=0,
61 + sizes=sizes,
62 + )
63
53 - hidden_states = self.naive_multicast(hidden_states,
54 - cu_tokens_across_dp_cpu)
55 - router_logits = self.naive_multicast(router_logits,
56 - cu_tokens_across_dp_cpu)
64 return hidden_states, router_logits
65
66 def combine(self, hidden_states: torch.Tensor) -> torch.Tensor:
60 - cu_tokens_across_dp_cpu = get_forward_context(
61 - ).dp_metadata.cu_tokens_across_dp_cpu
62 - start = 0 if self.dp_rank == 0 else cu_tokens_across_dp_cpu[
63 - self.dp_rank - 1]
64 - end = cu_tokens_across_dp_cpu[self.dp_rank]
65 -
66 - all_hidden_states = self.dp_group.all_reduce(hidden_states)
67 - hidden_states = all_hidden_states[start:end, :]
67 + sizes = get_forward_context(
68 + ).dp_metadata.get_chunk_sizes_across_dp_rank()
69 + hidden_states = get_dp_group().reduce_scatterv(hidden_states,
70 + dim=0,
71 + sizes=sizes)
72 return hidden_states
73
74 def destroy(self):
@@ -300,4 +304,95 @@ class DeepEPLLAll2AllManager(DeepEPAll2AllManagerBase):
304
305 # DeepEP LL uses RDMA so no SMs are used for communication
306 def max_sms_used(self) -> Optional[int]:
303 - return 0
\ No newline at end of file
307 + return 0
308 +
309 +
310 +class FlashInferAllToAllManager(All2AllManagerBase):
311 + """
312 + All2All communication based on flashinfer kernels.
313 + """
314 +
315 + def __init__(self, cpu_group):
316 + assert has_flashinfer_all2all(
317 + ), "flashinfer all2all module not found. Please install/check flashinfer" # noqa
318 + super().__init__(cpu_group)
319 + logger.debug(
320 + "Initialize for flashinfer All2All "
321 + "rank=%d, world size=%d", self.rank, self.world_size)
322 + self.initialized = False
323 + self.alltoall_info = None
324 +
325 + def initialize(
326 + self,
327 + world_size: int,
328 + rank: int,
329 + gpus_per_node: int,
330 + ):
331 + """Initialize workspace"""
332 + if self.initialized:
333 + return
334 +
335 + self.cleanup()
336 + logger.debug("making map: "
337 + "rank=%d, world size=%d", rank, world_size)
338 + self.mapping = Mapping(
339 + world_size,
340 + rank,
341 + gpus_per_node,
342 + tp_size=world_size,
343 + )
344 +
345 + from vllm.distributed.device_communicators.mnnvl_compat import (
346 + CustomCommunicator)
347 + dp_config = MnnvlConfig(
348 + comm_backend=CustomCommunicator(get_dp_group().cpu_group),
349 + fabric_page_size=1 << 29, # 512MB
350 + allocation_granularity=0 # Auto-detect
351 + )
352 +
353 + self.workspace_tensor = MnnvlMoe.get_moe_workspaces(
354 + self.mapping, dp_config)
355 + self.prepare_workspace_tensor = MnnvlMoe.get_moe_prepare_workspace(
356 + self.mapping, dp_config)
357 +
358 + self.world_size = world_size
359 + self.rank = rank
360 + self.gpus_per_node = gpus_per_node
361 + self.initialized = True
362 +
363 + logger.info("FlashInfer All2All initialized for rank %s, size %s",
364 + rank, world_size)
365 +
366 + def ensure_alltoall_workspace_initialized(self):
367 + """Ensure workspace is initialized"""
368 + if not has_flashinfer_all2all():
369 + return False
370 +
371 + if self.world_size <= 1:
372 + return False
373 +
374 + if not self.initialized:
375 + self.initialize(
376 + world_size=self.world_size,
377 + rank=self.rank,
378 + gpus_per_node=torch.cuda.device_count,
379 + )
380 + return self.initialized
381 +
382 + def get_handle(self, kwargs):
383 + return self
384 +
385 + def cleanup(self):
386 + """Clean up workspace"""
387 + if self.initialized and self.workspace_tensor is not None \
388 + and self.prepare_workspace_tensor is not None:
389 + try:
390 + del self.workspace_tensor
391 + del self.prepare_workspace_tensor
392 + except Exception as e:
393 + logger.warning("Failed to cleanup FlashInfer workspace: %s", e)
394 + finally:
395 + self.workspace_tensor = None
396 + self.prepare_workspace_tensor = None
397 + self.mapping = None
398 + self.initialized = False
\ No newline at end of file
vllm/distributed/device_communicators/cuda_communicator.py
+5
@@ -114,6 +114,11 @@ class CudaCommunicator(DeviceCommunicatorBase):
114 from .all2all import DeepEPLLAll2AllManager
115 self.all2all_manager = DeepEPLLAll2AllManager(self.cpu_group)
116 logger.info("Using DeepEP Low-Latency all2all manager.")
117 + elif all2all_backend == "flashinfer_all2allv":
118 + from .all2all import FlashInferAllToAllManager
119 + self.all2all_manager = FlashInferAllToAllManager(
120 + self.cpu_group)
121 + logger.info("Using Flashinfer all2allv manager.")
122 else:
123 raise ValueError(f"Unknown all2all backend: {all2all_backend}")
124
vllm/distributed/device_communicators/mnnvl_compat.py new
+28
@@ -0,0 +1,28 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch.distributed as dist
4 +from flashinfer.comm.mnnvl import CommBackend as CommBackend
5 +
6 +from vllm.utils.flashinfer import has_flashinfer_all2all
7 +
8 +assert has_flashinfer_all2all(), "Flashinfer alltoallv module cannot be found"
9 +
10 +
11 +class CustomCommunicator(CommBackend):
12 +
13 + def __init__(self, group):
14 + self._group = group
15 +
16 + def Get_rank(self) -> int:
17 + return self._group.rank()
18 +
19 + def Get_size(self) -> int:
20 + return self._group.size()
21 +
22 + def allgather(self, data: int):
23 + gathered = [None] * self.Get_size()
24 + dist.all_gather_object(gathered, data, group=self._group)
25 + return gathered
26 +
27 + def Split(self, color: int, key: int) -> 'CustomCommunicator':
28 + return self
vllm/envs.py
+5 -2
@@ -156,7 +156,8 @@ if TYPE_CHECKING:
156 VLLM_ALL2ALL_BACKEND: Literal["naive", "pplx",
157 "deepep_high_throughput",
158 "deepep_low_latency",
159 - "allgather_reducescatter"] = \
159 + "allgather_reducescatter",
160 + "flashinfer_all2allv"] = \
161 "allgather_reducescatter"
162 VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE: int = 163840
163 VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1
@@ -1209,12 +1210,14 @@ environment_variables: dict[str, Callable[[], Any]] = {
1210 # - "pplx": use pplx kernels
1211 # - "deepep_high_throughput", use deepep high-throughput kernels
1212 # - "deepep_low_latency", use deepep low-latency kernels
1213 + # - "flashinfer_all2allv", use flashinfer alltoallv kernels for mnnvl
1214 "VLLM_ALL2ALL_BACKEND":
1215 env_with_choices("VLLM_ALL2ALL_BACKEND", "allgather_reducescatter",
1216 ["naive", "pplx",
1217 "deepep_high_throughput",
1218 "deepep_low_latency",
1217 - "allgather_reducescatter"]),
1219 + "allgather_reducescatter",
1220 + "flashinfer_all2allv"]),
1221
1222 # Flashinfer MoE backend for vLLM's fused Mixture-of-Experts support.
1223 # Both require compute capability 10.0 or above.
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py
+3 -4
@@ -8,7 +8,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)
11 + create_flashinfer_prepare_finalize)
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,
@@ -108,7 +108,7 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute):
108 of each tuple must be the number of tokens.
109 """
110 aq_m, aq_n = aq.shape
111 - workspace2 = ()
111 + workspace2 = (0, )
112 output_shape = (aq_m, aq_n * 2) if self.quant_dtype != \
113 torch.float8_e4m3fn else (aq_m, aq_n)
114 workspace_dtype = a.dtype
@@ -192,9 +192,8 @@ def flashinfer_cutlass_moe_fp4(
192 expert_map: Optional[torch.Tensor] = None,
193 apply_router_weight_on_input: bool = False,
194 ) -> torch.Tensor:
195 -
195 fused_experts = mk.FusedMoEModularKernel(
197 - FlashInferCutlassMoEPrepareAndFinalize(use_dp=False),
196 + create_flashinfer_prepare_finalize(use_dp=False),
197 FlashInferExperts(
198 out_dtype=hidden_states.dtype,
199 quant_config=quant_config,
vllm/model_executor/layers/fused_moe/flashinfer_cutlass_prepare_finalize.py
+220 -13
@@ -5,7 +5,9 @@ from typing import Optional
5 import torch
6
7 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
8 -from vllm.distributed import get_dp_group
8 +from vllm.distributed import get_dp_group, get_ep_group
9 +from vllm.distributed.device_communicators.base_device_communicator import (
10 + All2AllManagerBase)
11 from vllm.forward_context import get_forward_context
12 from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
13 from vllm.model_executor.layers.fused_moe.utils import (
@@ -18,6 +20,7 @@ def get_local_sizes():
20
21
22 class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
23 + """Base class for FlashInfer MoE prepare and finalize operations."""
24
25 def __init__(
26 self,
@@ -42,6 +45,39 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
45 def num_dispatchers(self) -> int:
46 return self.num_dispatchers_
47
48 + def _apply_router_weight_on_input(
49 + self,
50 + a1: torch.Tensor,
51 + topk_weights: torch.Tensor,
52 + topk_ids: torch.Tensor,
53 + apply_router_weight_on_input: bool,
54 + ) -> None:
55 + """Apply router weight on input if needed."""
56 + if apply_router_weight_on_input:
57 + topk = topk_ids.size(1)
58 + assert topk == 1, \
59 + "apply_router_weight_on_input is only implemented for topk=1"
60 + a1.mul_(topk_weights.to(a1.dtype))
61 +
62 +
63 +class FlashInferAllToAllMoEPrepareAndFinalize(
64 + FlashInferCutlassMoEPrepareAndFinalize):
65 + """FlashInfer implementation using AllToAll communication."""
66 +
67 + def __init__(
68 + self,
69 + use_dp: bool,
70 + num_dispatchers: int = 1,
71 + ):
72 + super().__init__(use_dp, num_dispatchers)
73 + self.alltoall_info = None
74 +
75 + # Initialize all2all_manager only for DP case
76 + self.all2all_manager = None
77 + if self.use_dp:
78 + self.all2all_manager = get_ep_group(
79 + ).device_communicator.all2all_manager
80 +
81 def prepare(
82 self,
83 a1: torch.Tensor,
@@ -53,12 +89,84 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
89 quant_config: FusedMoEQuantConfig,
90 ) -> mk.PrepareResultType:
91
56 - if apply_router_weight_on_input:
57 - topk = topk_ids.size(1)
58 - # TODO: this only works for topK=1, will need to update for topK>1
59 - assert topk == 1, \
60 - "apply_router_weight_on_input is only implemented for topk=1"
61 - a1.mul_(topk_weights.to(a1.dtype))
92 + self._apply_router_weight_on_input(a1, topk_weights, topk_ids,
93 + apply_router_weight_on_input)
94 +
95 + if not self.use_dp:
96 + # Non-DP case: standard quantization
97 + a1q, a1q_scale = moe_kernel_quantize_input(
98 + a1,
99 + quant_config.a1_gscale,
100 + quant_config.quant_dtype,
101 + quant_config.per_act_token_quant,
102 + quant_config.block_shape,
103 + is_fp4_scale_swizzled=not self.use_dp,
104 + )
105 + else:
106 + # DP case: use FlashInfer AllToAll
107 + global_num_tokens_cpu = get_local_sizes()
108 + top_k = topk_ids.size(1)
109 +
110 + (self.alltoall_info, topk_ids, topk_weights, a1q,
111 + a1q_scale) = flashinfer_alltoall_dispatch(
112 + self.all2all_manager,
113 + global_num_tokens_cpu,
114 + a1,
115 + quant_config.a1_gscale,
116 + topk_ids,
117 + topk_weights,
118 + top_k,
119 + num_experts,
120 + quant_config,
121 + )
122 +
123 + return a1q, a1q_scale, None, topk_ids, topk_weights
124 +
125 + def finalize(
126 + self,
127 + output: torch.Tensor,
128 + fused_expert_output: torch.Tensor,
129 + topk_weights: torch.Tensor,
130 + topk_ids: torch.Tensor,
131 + apply_router_weight_on_input: bool,
132 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
133 + ) -> None:
134 + if self.use_dp:
135 + top_k = topk_ids.size(1)
136 + token_count = output.shape[0]
137 + fused_expert_output = flashinfer_alltoall_combine(
138 + self.all2all_manager,
139 + fused_expert_output,
140 + top_k=top_k,
141 + token_count=token_count,
142 + alltoall_info=self.alltoall_info,
143 + )
144 + output.copy_(fused_expert_output)
145 +
146 +
147 +class FlashInferAllGatherMoEPrepareAndFinalize(
148 + FlashInferCutlassMoEPrepareAndFinalize):
149 +
150 + def __init__(
151 + self,
152 + use_dp: bool,
153 + num_dispatchers: int = 1,
154 + ):
155 + super().__init__(use_dp, num_dispatchers)
156 +
157 + def prepare(
158 + self,
159 + a1: torch.Tensor,
160 + topk_weights: torch.Tensor,
161 + topk_ids: torch.Tensor,
162 + num_experts: int,
163 + expert_map: Optional[torch.Tensor],
164 + apply_router_weight_on_input: bool,
165 + quant_config: FusedMoEQuantConfig,
166 + ) -> mk.PrepareResultType:
167 +
168 + self._apply_router_weight_on_input(a1, topk_weights, topk_ids,
169 + apply_router_weight_on_input)
170
171 a1q, a1q_scale = moe_kernel_quantize_input(
172 a1,
@@ -66,7 +174,6 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
174 quant_config.quant_dtype,
175 quant_config.per_act_token_quant,
176 quant_config.block_shape,
69 - # Swizzling after communication
177 is_fp4_scale_swizzled=not self.use_dp,
178 )
179 if self.use_dp:
@@ -76,17 +183,117 @@ class FlashInferCutlassMoEPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize):
183 dim=0,
184 sizes=get_local_sizes(),
185 )
79 - a1_m, a1_n = a1q.shape
186 a1q_scale = nvfp4_block_scale_interleave(a1q_scale)
187
188 return a1q, a1q_scale, None, topk_ids, topk_weights
189
84 - def finalize(self, output: torch.Tensor, fused_expert_output: torch.Tensor,
85 - topk_weights: torch.Tensor, topk_ids: torch.Tensor,
86 - apply_router_weight_on_input: bool,
87 - weight_and_reduce_impl: mk.TopKWeightAndReduce) -> None:
190 + def finalize(
191 + self,
192 + output: torch.Tensor,
193 + fused_expert_output: torch.Tensor,
194 + topk_weights: torch.Tensor,
195 + topk_ids: torch.Tensor,
196 + apply_router_weight_on_input: bool,
197 + weight_and_reduce_impl: mk.TopKWeightAndReduce,
198 + ) -> None:
199
200 if self.use_dp:
201 fused_expert_output = get_dp_group().reduce_scatterv(
202 fused_expert_output, dim=0, sizes=get_local_sizes())
203 output.copy_(fused_expert_output)
204 +
205 +
206 +def flashinfer_alltoall_dispatch(
207 + all2all_manager: All2AllManagerBase,
208 + global_num_tokens_cpu: list[int],
209 + x: torch.Tensor,
210 + gs: torch.Tensor,
211 + topk_ids: torch.Tensor,
212 + topk_weights: torch.Tensor,
213 + top_k: int,
214 + num_experts: int,
215 + quant_config: FusedMoEQuantConfig,
216 +):
217 + from flashinfer.comm.trtllm_alltoall import MnnvlMoe
218 + assert (all2all_manager.ensure_alltoall_workspace_initialized()
219 + ), "FlashInfer AllToAll workspace not available"
220 +
221 + ep_rank = all2all_manager.rank
222 + ep_size = all2all_manager.world_size
223 + max_num_token = max(global_num_tokens_cpu
224 + ) if global_num_tokens_cpu is not None else x.shape[0]
225 + alltoall_info, topk_ids, topk_weights, _ = (
226 + MnnvlMoe.mnnvl_moe_alltoallv_prepare_without_allgather(
227 + topk_ids,
228 + topk_weights,
229 + None,
230 + all2all_manager.prepare_workspace,
231 + max_num_token,
232 + ep_rank,
233 + ep_size,
234 + num_experts,
235 + num_experts,
236 + top_k,
237 + ))
238 +
239 + x, x_sf = moe_kernel_quantize_input(
240 + x,
241 + gs,
242 + quant_config.quant_dtype,
243 + quant_config.per_act_token_quant,
244 + quant_config.block_shape,
245 + is_fp4_scale_swizzled=False, # delay swizzle to after comm
246 + )
247 + x = MnnvlMoe.mnnvl_moe_alltoallv(
248 + x,
249 + alltoall_info,
250 + all2all_manager.workspace_tensor,
251 + ep_rank,
252 + ep_size,
253 + )
254 +
255 + x_sf = MnnvlMoe.mnnvl_moe_alltoallv(
256 + x_sf,
257 + alltoall_info,
258 + all2all_manager.workspace_tensor,
259 + ep_rank,
260 + ep_size,
261 + )
262 + x_sf = nvfp4_block_scale_interleave(x_sf)
263 + return alltoall_info, topk_ids, topk_weights, x, x_sf
264 +
265 +
266 +def flashinfer_alltoall_combine(
267 + all2all_manager: All2AllManagerBase,
268 + output: torch.Tensor,
269 + top_k: int,
270 + token_count: int,
271 + alltoall_info,
272 +):
273 + from flashinfer.comm.trtllm_alltoall import MnnvlMoe
274 + assert (all2all_manager.ensure_alltoall_workspace_initialized()
275 + ), "FlashInfer AllToAll workspace not available"
276 + return MnnvlMoe.mnnvl_moe_alltoallv_combine(
277 + output,
278 + alltoall_info,
279 + all2all_manager.workspace_tensor,
280 + ep_rank=all2all_manager.rank,
281 + ep_size=all2all_manager.world_size,
282 + top_k=top_k,
283 + token_count=token_count,
284 + )
285 +
286 +
287 +def create_flashinfer_prepare_finalize(
288 + use_dp: bool,
289 + use_nvfp4: bool = False,
290 + enable_alltoallv: bool = False,
291 +) -> FlashInferCutlassMoEPrepareAndFinalize:
292 + """Factory function to create the appropriate FlashInfer implementation."""
293 + if use_nvfp4:
294 + if enable_alltoallv:
295 + return FlashInferAllToAllMoEPrepareAndFinalize(use_dp)
296 + else:
297 + return FlashInferAllGatherMoEPrepareAndFinalize(use_dp)
298 + # Fp8 only supports AllGather
299 + return FlashInferAllGatherMoEPrepareAndFinalize(use_dp)
vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py
+4 -2
@@ -12,7 +12,7 @@ from vllm.model_executor.layers.fused_moe.config import (FusedMoEConfig,
12 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
13 FlashInferExperts)
14 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
15 - FlashInferCutlassMoEPrepareAndFinalize)
15 + create_flashinfer_prepare_finalize)
16 from vllm.platforms import current_platform
17 from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
18
@@ -51,7 +51,9 @@ def build_flashinfer_fp4_cutlass_moe_prepare_finalize(
51 moe: FusedMoEConfig) -> mk.FusedMoEPrepareAndFinalize:
52 """Create a FlashInfer CUTLASS fused-MoE prepare finalize kernel"""
53 use_dp = moe.moe_parallel_config.dp_size > 1
54 - return FlashInferCutlassMoEPrepareAndFinalize(use_dp)
54 + enable_alltoallv = envs.VLLM_ALL2ALL_BACKEND == "flashinfer_all2allv"
55 + return create_flashinfer_prepare_finalize(
56 + use_dp=use_dp, use_nvfp4=True, enable_alltoallv=enable_alltoallv)
57
58
59 def select_nvfp4_gemm_impl(
vllm/model_executor/layers/quantization/utils/flashinfer_utils.py
+2 -2
@@ -13,7 +13,7 @@ from vllm.model_executor.layers.fused_moe.config import (FusedMoEConfig,
13 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
14 FlashInferExperts)
15 from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_prepare_finalize import ( # noqa: E501
16 - FlashInferCutlassMoEPrepareAndFinalize)
16 + create_flashinfer_prepare_finalize)
17
18 logger = init_logger(__name__)
19
@@ -173,7 +173,7 @@ def build_flashinfer_fp8_cutlass_moe_prepare_finalize(
173 moe: Optional[FusedMoEConfig], ) -> mk.FusedMoEPrepareAndFinalize:
174 """Create a FlashInfer CUTLASS fused-MoE prepare finalize kernel"""
175 use_dp = moe.moe_parallel_config.dp_size > 1 if moe is not None else False
176 - return FlashInferCutlassMoEPrepareAndFinalize(use_dp)
176 + return create_flashinfer_prepare_finalize(use_dp)
177
178
179 def select_cutlass_fp8_gemm_impl(
vllm/utils/flashinfer.py
+30
@@ -97,6 +97,34 @@ autotune = _lazy_import_wrapper(
97 fallback_fn=lambda *args, **kwargs: contextlib.nullcontext())
98
99
100 +@functools.cache
101 +def has_flashinfer_comm() -> bool:
102 + """Return ``True`` if FlashInfer comm module is available."""
103 + return has_flashinfer() and importlib.util.find_spec(
104 + "flashinfer.comm") is not None
105 +
106 +
107 +@functools.cache
108 +def has_flashinfer_all2all() -> bool:
109 + """Return ``True`` if FlashInfer mnnvl all2all is available."""
110 + if not has_flashinfer_comm():
111 + return False
112 +
113 + # Check if all required functions are available
114 + required_functions = [
115 + ("flashinfer.comm", "Mapping"),
116 + ("flashinfer.comm.mnnvl", "MnnvlMemory"),
117 + ("flashinfer.comm.trtllm_alltoall", "MnnvlMoe"),
118 + ("flashinfer.comm.trtllm_alltoall", "MoEAlltoallInfo"),
119 + ]
120 +
121 + for module_name, attr_name in required_functions:
122 + mod = _get_submodule(module_name)
123 + if not mod or not hasattr(mod, attr_name):
124 + return False
125 + return True
126 +
127 +
128 @functools.cache
129 def has_flashinfer_moe() -> bool:
130 """Return ``True`` if FlashInfer MoE module is available."""
@@ -402,6 +430,8 @@ __all__ = [
430 "trtllm_fp4_block_scale_moe",
431 "autotune",
432 "has_flashinfer_moe",
433 + "has_flashinfer_comm",
434 + "has_flashinfer_all2all",
435 "has_flashinfer_cutlass_fused_moe",
436 "has_nvidia_artifactory",
437 "supports_trtllm_attention",