| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from dataclasses import dataclass |
| 5 | from typing import Any |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.config import get_current_vllm_config |
| 10 | from vllm.distributed import ( |
| 11 | GroupCoordinator, |
| 12 | get_ep_group, |
| 13 | ) |
| 14 | from vllm.logger import init_logger |
| 15 | from vllm.model_executor.layers.fused_moe.config import ( |
| 16 | FusedMoEConfig, |
| 17 | FusedMoEParallelConfig, |
| 18 | FusedMoEQuantConfig, |
| 19 | ) |
| 20 | from vllm.model_executor.layers.fused_moe.modular_kernel import ( |
| 21 | FusedMoEPrepareAndFinalize, |
| 22 | ) |
| 23 | from vllm.model_executor.layers.fused_moe.prepare_finalize import ( |
| 24 | BatchedPrepareAndFinalize, |
| 25 | make_moe_prepare_and_finalize_naive_dp_ep, |
| 26 | make_moe_prepare_and_finalize_no_dp_ep, |
| 27 | ) |
| 28 | from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided import ( # noqa: E501 |
| 29 | FlashInferNVLinkOneSidedPrepareAndFinalize, |
| 30 | ) |
| 31 | from vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided import ( # noqa: E501 |
| 32 | FlashInferNVLinkTwoSidedPrepareAndFinalize, |
| 33 | ) |
| 34 | from vllm.platforms import current_platform |
| 35 | from vllm.utils.import_utils import ( |
| 36 | has_deep_ep, |
| 37 | has_deep_ep_v2, |
| 38 | has_moonep, |
| 39 | has_mori, |
| 40 | has_nixl_ep, |
| 41 | ) |
| 42 | |
| 43 | |
| 44 | @dataclass(frozen=True) |
| 45 | class FlashInferOneSidedDispatchLayout: |
| 46 | x_bytes_per_token: int |
| 47 | x_sf_bytes_per_token: int |
| 48 | |
| 49 | |
| 50 | def flashinfer_one_sided_dispatch_layout( |
| 51 | hidden_dim: int, |
| 52 | quant_config: FusedMoEQuantConfig, |
| 53 | input_dtype: torch.dtype | None = None, |
| 54 | ) -> FlashInferOneSidedDispatchLayout: |
| 55 | """Return the one-sided activation payload layout.""" |
| 56 | if input_dtype is not None: |
| 57 | if input_dtype not in (torch.float16, torch.bfloat16): |
| 58 | raise ValueError( |
| 59 | "flashinfer_nvlink_one_sided unpacked inputs must be float16 " |
| 60 | f"or bfloat16, got {input_dtype}" |
| 61 | ) |
| 62 | return FlashInferOneSidedDispatchLayout(hidden_dim * input_dtype.itemsize, 0) |
| 63 | if quant_config.quant_dtype is None: |
| 64 | return FlashInferOneSidedDispatchLayout(hidden_dim * 2, 0) |
| 65 | if quant_config.quant_dtype == "nvfp4": |
| 66 | scale_elems = hidden_dim // 16 |
| 67 | return FlashInferOneSidedDispatchLayout(hidden_dim // 2, scale_elems) |
| 68 | if quant_config.quant_dtype == "mxfp8": |
| 69 | align = quant_config.mx_alignment |
| 70 | padded_k = ( |
| 71 | ((hidden_dim + align - 1) // align) * align if align > 0 else hidden_dim |
| 72 | ) |
| 73 | scale_elems = padded_k // 32 |
| 74 | return FlashInferOneSidedDispatchLayout(hidden_dim, scale_elems) |
| 75 | if ( |
| 76 | quant_config.use_fp8_w8a8 |
| 77 | and quant_config.quant_dtype == current_platform.fp8_dtype() |
| 78 | and quant_config.block_shape == [128, 128] |
| 79 | ): |
| 80 | if hidden_dim % 128 != 0: |
| 81 | raise NotImplementedError( |
| 82 | "flashinfer_nvlink_one_sided DeepSeek Blockwise FP8 dispatch " |
| 83 | f"requires hidden_dim divisible by 128; got {hidden_dim}" |
| 84 | ) |
| 85 | scale_elems = hidden_dim // 128 |
| 86 | scale_bytes = scale_elems * torch.float32.itemsize |
| 87 | return FlashInferOneSidedDispatchLayout(hidden_dim, scale_bytes) |
| 88 | raise NotImplementedError( |
| 89 | "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, " |
| 90 | "DeepSeek Blockwise FP8 (E4M3 with FP32 1x128 scales), and bf16 " |
| 91 | "(quant_dtype=None) today; got " |
| 92 | f"quant_dtype={quant_config.quant_dtype!r}, " |
| 93 | f"use_fp8_w8a8={quant_config.use_fp8_w8a8!r}, " |
| 94 | f"block_shape={quant_config.block_shape!r}" |
| 95 | ) |
| 96 | |
| 97 | |
| 98 | logger = init_logger(__name__) |
| 99 | |
| 100 | if current_platform.is_cuda_alike(): |
| 101 | if has_deep_ep(): |
| 102 | from .prepare_finalize.deepep_ht import DeepEPHTPrepareAndFinalize |
| 103 | from .prepare_finalize.deepep_ll import ( |
| 104 | DEEPEP_QUANT_BLOCK_SHAPE, |
| 105 | DeepEPLLPrepareAndFinalize, |
| 106 | ) |
| 107 | if has_deep_ep_v2(): |
| 108 | from .prepare_finalize.deepep_v2 import DeepEPV2PrepareAndFinalize |
| 109 | if has_mori(): |
| 110 | from .prepare_finalize.mori import MoriPrepareAndFinalize |
| 111 | if has_nixl_ep(): |
| 112 | from .prepare_finalize.nixl_ep import ( |
| 113 | NIXL_EP_QUANT_BLOCK_SHAPE, |
| 114 | NixlEPPrepareAndFinalize, |
| 115 | ) |
| 116 | if has_moonep(): |
| 117 | from .prepare_finalize.moonep import ( |
| 118 | MOONEP_DEFAULT_NUM_PREFETCH_SLOTS, |
| 119 | MOONEP_DEFAULT_NUM_SMS, |
| 120 | MOONEP_DEFAULT_TOKEN_PADDING, |
| 121 | MoonEPPrepareAndFinalize, |
| 122 | ) |
| 123 | |
| 124 | |
| 125 | def get_ep_all2all_manager( |
| 126 | ep_group: GroupCoordinator | None = None, |
| 127 | ) -> Any: |
| 128 | if ep_group is None: |
| 129 | ep_group = get_ep_group() |
| 130 | device_communicator = ep_group.device_communicator |
| 131 | assert device_communicator is not None |
| 132 | all2all_manager = device_communicator.all2all_manager |
| 133 | assert all2all_manager is not None |
| 134 | return all2all_manager |
| 135 | |
| 136 | |
| 137 | def maybe_roundup_layer_hidden_size( |
| 138 | hidden_size: int, |
| 139 | act_dtype: torch.dtype, |
| 140 | moe_parallel_config: FusedMoEParallelConfig, |
| 141 | ) -> int: |
| 142 | """Given layer hidden size and MoE configurations, round up hidden_size |
| 143 | if necessary. |
| 144 | |
| 145 | Args: |
| 146 | hidden_size: Layer hidden-size |
| 147 | act_dtype: Data type of the layer activations. |
| 148 | moe_parallel_config: Fused MoE parallelization strategy configuration. |
| 149 | |
| 150 | Return: |
| 151 | Rounded up hidden_size if rounding up is required based on the configs |
| 152 | and all2all backend. |
| 153 | Original hidden size otherwise. |
| 154 | |
| 155 | """ |
| 156 | if moe_parallel_config.use_deepep_ht_kernels: |
| 157 | hidden_size = DeepEPHTPrepareAndFinalize.maybe_roundup_layer_hidden_size( |
| 158 | hidden_size, act_dtype |
| 159 | ) |
| 160 | |
| 161 | if moe_parallel_config.use_deepep_ll_kernels: |
| 162 | hidden_size = DeepEPLLPrepareAndFinalize.maybe_roundup_layer_hidden_size( |
| 163 | hidden_size |
| 164 | ) |
| 165 | |
| 166 | if moe_parallel_config.use_deepep_v2_kernels: |
| 167 | hidden_size = DeepEPV2PrepareAndFinalize.maybe_roundup_layer_hidden_size( |
| 168 | hidden_size, act_dtype |
| 169 | ) |
| 170 | |
| 171 | if moe_parallel_config.use_nixl_ep_kernels: |
| 172 | hidden_size = NixlEPPrepareAndFinalize.maybe_roundup_layer_hidden_size( |
| 173 | hidden_size |
| 174 | ) |
| 175 | |
| 176 | return hidden_size |
| 177 | |
| 178 | |
| 179 | def maybe_make_prepare_finalize( |
| 180 | moe: FusedMoEConfig, |
| 181 | quant_config: FusedMoEQuantConfig | None, |
| 182 | routing_tables: tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None = None, |
| 183 | allow_new_interface: bool = False, |
| 184 | use_monolithic: bool = False, |
| 185 | all2all_manager: Any | None = None, |
| 186 | input_dtype: torch.dtype | None = None, |
| 187 | ) -> FusedMoEPrepareAndFinalize | None: |
| 188 | if not moe.moe_parallel_config.use_all2all_kernels: |
| 189 | if not allow_new_interface: |
| 190 | return None |
| 191 | |
| 192 | # Opt-in XPU batched path: reorganize tokens into E x T x K locally |
| 193 | # (no all-to-all) so BatchedTritonExperts (moe_mmk TD) can run. |
| 194 | if current_platform.is_xpu() and moe.moe_backend == "batched_triton": |
| 195 | return BatchedPrepareAndFinalize( |
| 196 | max_num_tokens=moe.max_num_tokens, |
| 197 | num_local_experts=moe.num_local_experts, |
| 198 | num_dispatchers=1, |
| 199 | rank=moe.moe_parallel_config.ep_rank, |
| 200 | ) |
| 201 | |
| 202 | # For DP/TP case, fall back to naive P/F. |
| 203 | if moe.moe_parallel_config.dp_size > 1: |
| 204 | logger.info_once( |
| 205 | "Detected DP deployment with no --enable-expert-parallel. " |
| 206 | "Falling back to AllGather+ReduceScatter dispatch/combine." |
| 207 | ) |
| 208 | if all2all_manager is None: |
| 209 | all2all_manager = get_ep_all2all_manager() |
| 210 | return make_moe_prepare_and_finalize_naive_dp_ep( |
| 211 | is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel, |
| 212 | num_dispatchers=all2all_manager.world_size, |
| 213 | use_monolithic=use_monolithic, |
| 214 | ) |
| 215 | else: |
| 216 | return make_moe_prepare_and_finalize_no_dp_ep(use_monolithic) |
| 217 | |
| 218 | if all2all_manager is None: |
| 219 | all2all_manager = get_ep_all2all_manager() |
| 220 | |
| 221 | prepare_finalize: FusedMoEPrepareAndFinalize | None = None |
| 222 | |
| 223 | if moe.use_deepep_ht_kernels: |
| 224 | assert moe.dp_size == all2all_manager.dp_world_size |
| 225 | |
| 226 | all_to_all_args: dict[str, Any] = dict() |
| 227 | handle = all2all_manager.get_handle(all_to_all_args) |
| 228 | prepare_finalize = DeepEPHTPrepareAndFinalize( |
| 229 | handle, |
| 230 | num_dispatchers=all2all_manager.world_size, |
| 231 | dp_size=all2all_manager.dp_world_size, |
| 232 | rank_expert_offset=all2all_manager.rank * moe.num_local_experts, |
| 233 | ) |
| 234 | |
| 235 | elif moe.use_deepep_ll_kernels: |
| 236 | assert quant_config is not None |
| 237 | global_to_physical = physical_to_global = local_expert_global_ids = None |
| 238 | if routing_tables is not None: |
| 239 | ( |
| 240 | global_to_physical, |
| 241 | physical_to_global, |
| 242 | local_expert_global_ids, |
| 243 | ) = routing_tables |
| 244 | all_to_all_args = dict( |
| 245 | max_num_tokens_per_dp_rank=moe.max_num_tokens, |
| 246 | token_hidden_size=moe.hidden_dim, |
| 247 | num_ep_ranks=all2all_manager.world_size, |
| 248 | num_global_experts=moe.num_experts, |
| 249 | num_local_experts=moe.num_experts // all2all_manager.world_size, |
| 250 | ) |
| 251 | handle = all2all_manager.get_handle(all_to_all_args) |
| 252 | |
| 253 | # Note: We may want to use FP8 dispatch just to reduce |
| 254 | # data movement. |
| 255 | use_fp8_dispatch = ( |
| 256 | quant_config.quant_dtype == current_platform.fp8_dtype() |
| 257 | and quant_config.block_shape == DEEPEP_QUANT_BLOCK_SHAPE |
| 258 | ) |
| 259 | |
| 260 | prepare_finalize = DeepEPLLPrepareAndFinalize( |
| 261 | handle, |
| 262 | max_tokens_per_rank=moe.max_num_tokens, |
| 263 | num_dispatchers=all2all_manager.world_size, |
| 264 | use_fp8_dispatch=use_fp8_dispatch, |
| 265 | global_to_physical=global_to_physical, |
| 266 | physical_to_global=physical_to_global, |
| 267 | local_expert_global_ids=local_expert_global_ids, |
| 268 | ) |
| 269 | elif moe.use_deepep_v2_kernels: |
| 270 | assert moe.dp_size == all2all_manager.dp_world_size |
| 271 | |
| 272 | use_fp8_dispatch = ( |
| 273 | quant_config is not None |
| 274 | and quant_config.quant_dtype == current_platform.fp8_dtype() |
| 275 | and quant_config.is_block_quantized |
| 276 | ) |
| 277 | all_to_all_args = dict( |
| 278 | num_max_tokens_per_rank=moe.max_num_tokens, |
| 279 | hidden=moe.hidden_dim, |
| 280 | num_topk=moe.experts_per_token, |
| 281 | num_experts=moe.num_experts, |
| 282 | use_fp8_dispatch=use_fp8_dispatch, |
| 283 | ) |
| 284 | handle = all2all_manager.get_handle(all_to_all_args) |
| 285 | vllm_config = get_current_vllm_config() |
| 286 | use_cudagraph = not vllm_config.model_config.enforce_eager |
| 287 | |
| 288 | prepare_finalize = DeepEPV2PrepareAndFinalize( |
| 289 | buffer=handle, |
| 290 | num_dispatchers=all2all_manager.world_size, |
| 291 | dp_size=all2all_manager.dp_world_size, |
| 292 | rank_expert_offset=all2all_manager.rank * moe.num_local_experts, |
| 293 | num_experts=moe.num_experts, |
| 294 | num_topk=moe.experts_per_token, |
| 295 | use_fp8_dispatch=use_fp8_dispatch, |
| 296 | use_cudagraph=use_cudagraph, |
| 297 | sp_size=moe.moe_parallel_config.sp_size, |
| 298 | ) |
| 299 | |
| 300 | elif moe.use_moonep_kernels: |
| 301 | all_to_all_args = dict( |
| 302 | max_num_tokens_per_dp_rank=moe.max_num_tokens, |
| 303 | token_hidden_size=moe.hidden_dim, |
| 304 | num_topk=moe.experts_per_token, |
| 305 | num_global_experts=moe.num_experts, |
| 306 | num_prefetch_slots=MOONEP_DEFAULT_NUM_PREFETCH_SLOTS, |
| 307 | token_padding=MOONEP_DEFAULT_TOKEN_PADDING, |
| 308 | num_sms=MOONEP_DEFAULT_NUM_SMS, |
| 309 | ) |
| 310 | handle = all2all_manager.get_handle(all_to_all_args) |
| 311 | |
| 312 | # The [E+B] weight layout is picked up from the experts (their |
| 313 | # process_weights_after_loading hook) once the layer has loaded and |
| 314 | # converted its weights. |
| 315 | prepare_finalize = MoonEPPrepareAndFinalize( |
| 316 | buffer_pool=handle, |
| 317 | max_tokens_per_rank=moe.max_num_tokens, |
| 318 | num_dispatchers=all2all_manager.world_size, |
| 319 | num_global_experts=moe.num_experts, |
| 320 | ) |
| 321 | |
| 322 | elif moe.use_mori_kernels: |
| 323 | assert quant_config is not None |
| 324 | |
| 325 | # Note: We may want to use FP8 dispatch just to reduce |
| 326 | # data movement. |
| 327 | use_fp8_dispatch = ( |
| 328 | quant_config.is_per_act_token |
| 329 | or quant_config.is_block_quantized |
| 330 | or quant_config.is_per_tensor |
| 331 | ) |
| 332 | if use_fp8_dispatch: |
| 333 | # For PTPC (per token per channel) or per-tensor quant, |
| 334 | # scale dim is 1. For 1x128 quant, scale dim is |
| 335 | # hidden_dim // 128 |
| 336 | quant_dtype = quant_config.quant_dtype |
| 337 | scale_dim = ( |
| 338 | 1 |
| 339 | if (quant_config.is_per_act_token or quant_config.is_per_tensor) |
| 340 | else moe.hidden_dim // 128 |
| 341 | ) |
| 342 | else: |
| 343 | # Unquantized dispatch (e.g. AITER with defer_input_quant): |
| 344 | # dispatch raw BF16/FP16 data, no scales needed. |
| 345 | quant_dtype = moe.in_dtype |
| 346 | scale_dim = 0 |
| 347 | all_to_all_args = dict( |
| 348 | rank=all2all_manager.rank, |
| 349 | num_ep_ranks=all2all_manager.world_size, |
| 350 | quant_dtype=quant_dtype, |
| 351 | token_hidden_size=moe.hidden_dim, |
| 352 | scale_dim=scale_dim, |
| 353 | scale_type_size=0 if scale_dim == 0 else torch.float32.itemsize, |
| 354 | max_num_tokens_per_dp_rank=moe.max_num_tokens, |
| 355 | input_dtype=moe.in_dtype, |
| 356 | num_local_experts=moe.num_experts // all2all_manager.world_size, |
| 357 | num_experts_per_token=moe.experts_per_token, |
| 358 | ) |
| 359 | handle = all2all_manager.get_handle(all_to_all_args) |
| 360 | |
| 361 | prepare_finalize = MoriPrepareAndFinalize( |
| 362 | handle, |
| 363 | max_tokens_per_rank=moe.max_num_tokens, |
| 364 | num_dispatchers=all2all_manager.world_size, |
| 365 | use_fp8_dispatch=use_fp8_dispatch, |
| 366 | ) |
| 367 | |
| 368 | elif moe.use_fi_nvl_two_sided_kernels: |
| 369 | assert quant_config is not None |
| 370 | prepare_finalize = FlashInferNVLinkTwoSidedPrepareAndFinalize( |
| 371 | num_dispatchers=all2all_manager.world_size, |
| 372 | ) |
| 373 | |
| 374 | elif moe.use_fi_nvl_one_sided_kernels: |
| 375 | assert quant_config is not None |
| 376 | max_num_tokens = ( |
| 377 | get_current_vllm_config().scheduler_config.max_num_batched_tokens |
| 378 | ) |
| 379 | dispatch_layout = flashinfer_one_sided_dispatch_layout( |
| 380 | moe.hidden_dim, quant_config, input_dtype=input_dtype |
| 381 | ) |
| 382 | prepare_finalize = FlashInferNVLinkOneSidedPrepareAndFinalize( |
| 383 | max_num_tokens=max_num_tokens, |
| 384 | top_k=moe.experts_per_token, |
| 385 | num_experts=moe.num_experts, |
| 386 | hidden_size=moe.hidden_dim, |
| 387 | num_dispatchers=all2all_manager.world_size, |
| 388 | x_bytes_per_token=dispatch_layout.x_bytes_per_token, |
| 389 | x_sf_bytes_per_token=dispatch_layout.x_sf_bytes_per_token, |
| 390 | ) |
| 391 | |
| 392 | elif moe.use_ag_rs_all2all_kernels and allow_new_interface: |
| 393 | prepare_finalize = make_moe_prepare_and_finalize_naive_dp_ep( |
| 394 | use_monolithic=use_monolithic, |
| 395 | is_sequence_parallel=moe.moe_parallel_config.is_sequence_parallel, |
| 396 | num_dispatchers=all2all_manager.world_size, |
| 397 | ) |
| 398 | |
| 399 | elif moe.use_nixl_ep_kernels: |
| 400 | assert quant_config is not None |
| 401 | global_to_physical = physical_to_global = local_expert_global_ids = None |
| 402 | if routing_tables is not None: |
| 403 | ( |
| 404 | global_to_physical, |
| 405 | physical_to_global, |
| 406 | local_expert_global_ids, |
| 407 | ) = routing_tables |
| 408 | all_to_all_args = dict( |
| 409 | max_num_tokens_per_dp_rank=moe.max_num_tokens, |
| 410 | token_hidden_size=moe.hidden_dim, |
| 411 | num_local_experts=moe.num_local_experts, |
| 412 | ) |
| 413 | handle = all2all_manager.get_handle(all_to_all_args) |
| 414 | |
| 415 | # Note: We may want to use FP8 dispatch just to reduce |
| 416 | # data movement. |
| 417 | use_fp8_dispatch = ( |
| 418 | quant_config.quant_dtype == current_platform.fp8_dtype() |
| 419 | and quant_config.block_shape == NIXL_EP_QUANT_BLOCK_SHAPE |
| 420 | ) |
| 421 | |
| 422 | prepare_finalize = NixlEPPrepareAndFinalize( |
| 423 | handle, |
| 424 | max_tokens_per_rank=moe.max_num_tokens, |
| 425 | num_dispatchers=all2all_manager.max_num_ep_ranks, |
| 426 | expert_capacity=(moe.num_local_experts * all2all_manager.max_num_ep_ranks), |
| 427 | use_fp8_dispatch=use_fp8_dispatch, |
| 428 | global_to_physical=global_to_physical, |
| 429 | physical_to_global=physical_to_global, |
| 430 | local_expert_global_ids=local_expert_global_ids, |
| 431 | ) |
| 432 | |
| 433 | return prepare_finalize |