| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass |
| 4 | from types import SimpleNamespace |
| 5 | from typing import Any |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | import vllm._custom_ops as ops |
| 10 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 11 | from tests.kernels.moe.utils import make_test_weights, per_token_cast_to_fp8 |
| 12 | from tests.kernels.quantization.nvfp4_utils import ( |
| 13 | FLOAT4_E2M1_MAX, |
| 14 | FLOAT8_E4M3_MAX, |
| 15 | dequantize_nvfp4_to_dtype, |
| 16 | ) |
| 17 | from tests.kernels.utils import torch_experts |
| 18 | from vllm._aiter_ops import rocm_aiter_ops |
| 19 | from vllm.config import VllmConfig |
| 20 | from vllm.distributed import ( |
| 21 | get_dp_group, |
| 22 | get_pcp_group, |
| 23 | get_tensor_model_parallel_world_size, |
| 24 | ) |
| 25 | from vllm.forward_context import set_forward_context |
| 26 | from vllm.model_executor.layers.fused_moe import fused_topk |
| 27 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 28 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 29 | maybe_make_prepare_finalize, |
| 30 | ) |
| 31 | from vllm.model_executor.layers.fused_moe.config import ( |
| 32 | FusedMoEConfig, |
| 33 | FusedMoEParallelConfig, |
| 34 | FusedMoEQuantConfig, |
| 35 | RoutingMethodType, |
| 36 | ) |
| 37 | from vllm.model_executor.layers.quantization.utils.fp8_utils import is_fp8 |
| 38 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 39 | QuantKey, |
| 40 | kFp8Dynamic128Sym, |
| 41 | kFp8DynamicTensorSym, |
| 42 | kFp8DynamicTokenSym, |
| 43 | kFp8Static128BlockSym, |
| 44 | kFp8StaticChannelSym, |
| 45 | kFp8StaticTensorSym, |
| 46 | ) |
| 47 | from vllm.utils.import_utils import ( |
| 48 | has_aiter, |
| 49 | has_deep_ep, |
| 50 | has_deep_ep_v2, |
| 51 | has_deep_gemm, |
| 52 | has_mori, |
| 53 | ) |
| 54 | from vllm.utils.math_utils import next_power_of_2 |
| 55 | |
| 56 | from .mk_objects import ( |
| 57 | TestMoEQuantConfig, |
| 58 | expert_info, |
| 59 | make_fused_experts, |
| 60 | prepare_finalize_info, |
| 61 | ) |
| 62 | from .parallel_utils import ProcessGroupInfo |
| 63 | |
| 64 | |
| 65 | def _describe_tensor(t: torch.Tensor | None, name: str) -> str: |
| 66 | if t is None: |
| 67 | return f"{name} : None" |
| 68 | else: |
| 69 | return f"{name} : {t.shape} {t.dtype} {t.device}" |
| 70 | |
| 71 | |
| 72 | @dataclass |
| 73 | class Config: |
| 74 | Ms: list[int] | int |
| 75 | K: int |
| 76 | N: int |
| 77 | E: int |
| 78 | topks: list[int] | int |
| 79 | dtype: torch.dtype |
| 80 | quant_config: TestMoEQuantConfig | None |
| 81 | |
| 82 | prepare_finalize_type: mk.FusedMoEPrepareAndFinalize |
| 83 | fused_experts_type: mk.FusedMoEExperts |
| 84 | |
| 85 | world_size: int |
| 86 | |
| 87 | activation: MoEActivation = MoEActivation.SILU |
| 88 | |
| 89 | torch_trace_dir_path: str | None = None |
| 90 | |
| 91 | def __post_init__(self): |
| 92 | if self.quant_config is None: |
| 93 | self.quant_config = TestMoEQuantConfig(None, False, False, None) |
| 94 | |
| 95 | def describe(self) -> str: |
| 96 | s = "" |
| 97 | s += "== Config:\n" |
| 98 | s += f" world_size={self.world_size}\n" |
| 99 | s += f" PF={self.prepare_finalize_type.__name__}\n" |
| 100 | s += f" FE={self.fused_experts_type.__name__}\n" |
| 101 | s += f" E={self.E}\n" |
| 102 | s += f" Ms={self.Ms}\n" |
| 103 | s += f" N={self.N}\n" |
| 104 | s += f" K={self.K}\n" |
| 105 | s += f" topk={self.topks}\n" |
| 106 | s += f" dtype={self.dtype}\n" |
| 107 | s += " Quant:\n" |
| 108 | if self.quant_config is not None: |
| 109 | s += f" q_dtype={self.quant_dtype}\n" |
| 110 | s += f" q_block_shape={self.quant_block_shape}\n" |
| 111 | s += f" q_per_out_ch_quant={self.is_per_out_ch_quant}\n" |
| 112 | s += f" q_per_act_token={self.is_per_act_token_quant}\n" |
| 113 | else: |
| 114 | s += " quant=None\n" |
| 115 | return s |
| 116 | |
| 117 | @property |
| 118 | def M(self) -> int: |
| 119 | assert isinstance(self.Ms, int) |
| 120 | return self.Ms |
| 121 | |
| 122 | @property |
| 123 | def quant_dtype(self) -> torch.dtype | str | None: |
| 124 | assert self.quant_config is not None |
| 125 | return self.quant_config.quant_dtype |
| 126 | |
| 127 | @property |
| 128 | def is_per_act_token_quant(self) -> bool: |
| 129 | assert self.quant_config is not None |
| 130 | return self.quant_config.per_act_token_quant |
| 131 | |
| 132 | @property |
| 133 | def is_per_tensor_act_quant(self) -> bool: |
| 134 | return not self.is_per_act_token_quant and self.quant_block_shape is None |
| 135 | |
| 136 | @property |
| 137 | def is_per_out_ch_quant(self) -> bool: |
| 138 | assert self.quant_config is not None |
| 139 | return self.quant_config.per_out_ch_quant |
| 140 | |
| 141 | @property |
| 142 | def quant_block_shape(self) -> list[int] | None: |
| 143 | assert self.quant_config is not None |
| 144 | return self.quant_config.block_shape |
| 145 | |
| 146 | @property |
| 147 | def topk(self) -> int: |
| 148 | assert isinstance(self.topks, int) |
| 149 | return self.topks |
| 150 | |
| 151 | @property |
| 152 | def num_local_experts(self) -> int: |
| 153 | return self.E // self.world_size |
| 154 | |
| 155 | def make_env_data(self) -> tuple[VllmConfig, dict[Any, Any]]: |
| 156 | """Make env data for vllm launch.""" |
| 157 | vllm_config = VllmConfig() |
| 158 | vllm_config.model_config = SimpleNamespace( |
| 159 | enforce_eager=True, |
| 160 | is_moe=True, |
| 161 | ) |
| 162 | vllm_config.parallel_config.data_parallel_size = self.world_size |
| 163 | vllm_config.parallel_config.enable_expert_parallel = True |
| 164 | |
| 165 | env_dict = { |
| 166 | "VLLM_USE_DEEP_GEMM": str(int(self.needs_deep_gemm())), |
| 167 | } |
| 168 | |
| 169 | vllm_config.parallel_config.all2all_backend = self.all2all_backend() |
| 170 | |
| 171 | return vllm_config, env_dict |
| 172 | |
| 173 | def fp8_quant_key_pair(self) -> tuple[QuantKey, QuantKey]: |
| 174 | """Derive the (weight_quant_key, activation_quant_key) pair an FP8 |
| 175 | quant config of this shape corresponds to (either OCP or FNUZ FP8, |
| 176 | see ``current_platform.fp8_dtype()``).""" |
| 177 | if self.quant_block_shape is not None: |
| 178 | return kFp8Static128BlockSym, kFp8Dynamic128Sym |
| 179 | if self.is_per_out_ch_quant: |
| 180 | return ( |
| 181 | kFp8StaticChannelSym, |
| 182 | kFp8DynamicTokenSym |
| 183 | if self.is_per_act_token_quant |
| 184 | else kFp8StaticTensorSym, |
| 185 | ) |
| 186 | return ( |
| 187 | kFp8StaticTensorSym, |
| 188 | kFp8DynamicTensorSym |
| 189 | if self.is_per_act_token_quant |
| 190 | else kFp8StaticTensorSym, |
| 191 | ) |
| 192 | |
| 193 | def fe_supports_quant_scheme(self) -> bool: |
| 194 | """Check if the fused experts class supports this quant config. |
| 195 | See https://github.com/ROCm/aiter/issues/2419 for AITER gaps.""" |
| 196 | if self.quant_config is None or self.quant_dtype is None: |
| 197 | return True |
| 198 | if not is_fp8(self.quant_dtype): |
| 199 | return True |
| 200 | w_key, a_key = self.fp8_quant_key_pair() |
| 201 | fe_cls = self.fused_experts_type |
| 202 | if hasattr(fe_cls, "_supports_quant_scheme"): |
| 203 | try: |
| 204 | return fe_cls._supports_quant_scheme(w_key, a_key) |
| 205 | except NotImplementedError: |
| 206 | pass |
| 207 | return True |
| 208 | |
| 209 | def is_fp8_block_quantized(self): |
| 210 | return is_fp8(self.quant_dtype) and self.quant_block_shape is not None |
| 211 | |
| 212 | def is_batched_prepare_finalize(self): |
| 213 | info = prepare_finalize_info(self.prepare_finalize_type) |
| 214 | return mk.FusedMoEActivationFormat.BatchedExperts == info.activation_format |
| 215 | |
| 216 | def is_batched_fused_experts(self): |
| 217 | info = expert_info(self.fused_experts_type) |
| 218 | return mk.FusedMoEActivationFormat.BatchedExperts == info.activation_format |
| 219 | |
| 220 | def is_standard_fused_experts(self): |
| 221 | info = expert_info(self.fused_experts_type) |
| 222 | return mk.FusedMoEActivationFormat.Standard == info.activation_format |
| 223 | |
| 224 | def fe_supported_types(self): |
| 225 | info = expert_info(self.fused_experts_type) |
| 226 | return info.supported_dtypes |
| 227 | |
| 228 | def pf_supported_types(self): |
| 229 | info = prepare_finalize_info(self.prepare_finalize_type) |
| 230 | return info.supported_dtypes |
| 231 | |
| 232 | def is_block_quant_supported(self): |
| 233 | info = expert_info(self.fused_experts_type) |
| 234 | return info.blocked_quantization_support |
| 235 | |
| 236 | def supports_apply_weight_on_input(self): |
| 237 | info = prepare_finalize_info(self.prepare_finalize_type) |
| 238 | return info.supports_apply_weight_on_input |
| 239 | |
| 240 | def needs_deep_gemm(self): |
| 241 | info = expert_info(self.fused_experts_type) |
| 242 | return info.needs_deep_gemm |
| 243 | |
| 244 | def needs_deep_ep(self): |
| 245 | info = prepare_finalize_info(self.prepare_finalize_type) |
| 246 | return ( |
| 247 | info.backend == "deepep_high_throughput" |
| 248 | or info.backend == "deepep_low_latency" |
| 249 | ) |
| 250 | |
| 251 | def needs_deep_ep_v2(self): |
| 252 | info = prepare_finalize_info(self.prepare_finalize_type) |
| 253 | return info.backend == "deepep_v2" |
| 254 | |
| 255 | def needs_aiter(self): |
| 256 | info = expert_info(self.fused_experts_type) |
| 257 | return info.needs_aiter |
| 258 | |
| 259 | def needs_mori(self): |
| 260 | info = prepare_finalize_info(self.prepare_finalize_type) |
| 261 | return info.backend in ("mori_high_throughput", "mori_low_latency") |
| 262 | |
| 263 | def all2all_backend(self): |
| 264 | info = prepare_finalize_info(self.prepare_finalize_type) |
| 265 | return info.backend |
| 266 | |
| 267 | def is_valid(self) -> tuple[bool, str | None]: |
| 268 | # Check prepare-finalize and fused-experts compatibility |
| 269 | if self.is_batched_prepare_finalize(): |
| 270 | if not self.is_batched_fused_experts(): |
| 271 | return False, "Mismatched format." |
| 272 | else: |
| 273 | if not self.is_standard_fused_experts(): |
| 274 | return False, "Mismatched format." |
| 275 | |
| 276 | # Check quantization sanity |
| 277 | if ( |
| 278 | int(self.is_per_act_token_quant) |
| 279 | + int(self.is_per_tensor_act_quant) |
| 280 | + int(self.quant_block_shape is not None) |
| 281 | ) > 1: |
| 282 | # invalid quant config |
| 283 | return False, f"Bad quant_config {self.quant_config}." |
| 284 | |
| 285 | # check type support |
| 286 | if self.quant_dtype is None: |
| 287 | if ( |
| 288 | self.dtype not in self.pf_supported_types() |
| 289 | or self.dtype not in self.fe_supported_types() |
| 290 | ): |
| 291 | return False, ( |
| 292 | f"Unsupported type {self.dtype} not in " |
| 293 | f"{self.pf_supported_types()} and " |
| 294 | f"{self.fe_supported_types()}." |
| 295 | ) |
| 296 | else: |
| 297 | if ( |
| 298 | self.quant_dtype not in self.pf_supported_types() |
| 299 | or self.quant_dtype not in self.fe_supported_types() |
| 300 | ): |
| 301 | return False, ( |
| 302 | f"Unsupported quant type {self.quant_dtype} " |
| 303 | f"not in {self.pf_supported_types()} and " |
| 304 | f"{self.fe_supported_types()}." |
| 305 | ) |
| 306 | |
| 307 | # Check quant scheme compatibility with fused experts class |
| 308 | if not self.fe_supports_quant_scheme(): |
| 309 | return False, ( |
| 310 | f"FE {self.fused_experts_type.__name__} does not support " |
| 311 | f"quant scheme (per_out_ch={self.is_per_out_ch_quant}, " |
| 312 | f"per_act_token={self.is_per_act_token_quant}, " |
| 313 | f"block={self.quant_block_shape})" |
| 314 | ) |
| 315 | |
| 316 | # Check activation support; NotImplementedError means no opinion. |
| 317 | try: |
| 318 | if not self.fused_experts_type._supports_activation(self.activation): |
| 319 | return False, ( |
| 320 | f"FE {self.fused_experts_type.__name__} does not support " |
| 321 | f"activation {self.activation}" |
| 322 | ) |
| 323 | except NotImplementedError: |
| 324 | pass |
| 325 | |
| 326 | # Check block quantization support |
| 327 | is_block_quantized = self.quant_block_shape is not None |
| 328 | if is_block_quantized and self.quant_dtype is None: |
| 329 | return False, "No block quantization support." |
| 330 | |
| 331 | if is_block_quantized and not self.is_block_quant_supported(): |
| 332 | return False, "Mismatched block quantization support." |
| 333 | |
| 334 | # deep_gemm only works with block-quantized |
| 335 | if self.needs_deep_gemm() and not is_block_quantized: |
| 336 | return False, "Needs DeepGEMM but not block quantized." |
| 337 | |
| 338 | # Check dependencies (turn into asserts?) |
| 339 | if self.needs_deep_ep() and not has_deep_ep(): |
| 340 | return False, "Needs DeepEP, but DeepEP not available." |
| 341 | if self.needs_deep_ep_v2() and not has_deep_ep_v2(): |
| 342 | return False, "Needs DeepEP v2, but DeepEP v2 not available." |
| 343 | if self.needs_deep_gemm() and not has_deep_gemm(): |
| 344 | return ( |
| 345 | False, |
| 346 | "Needs DeepGEMM, but the current vLLM environment does not provide it.", |
| 347 | ) |
| 348 | if self.needs_aiter() and not has_aiter(): # noqa: SIM103 |
| 349 | return False, "Needs Aiter, but Aiter not available." |
| 350 | if self.needs_mori() and not has_mori(): # noqa: SIM103 |
| 351 | return False, "Needs MoRI, but MoRI not available." |
| 352 | if self.needs_mori() and not rocm_aiter_ops.is_fused_moe_enabled(): |
| 353 | return False, ( |
| 354 | "Mori requires AITER's fused-moe backend to be enabled " |
| 355 | "(VLLM_ROCM_USE_AITER=1 and VLLM_ROCM_USE_AITER_MOE=1)." |
| 356 | ) |
| 357 | |
| 358 | try: |
| 359 | if not self.fused_experts_type._supports_current_device(): |
| 360 | return ( |
| 361 | False, |
| 362 | f"{self.fused_experts_type} not supported on the current device.", |
| 363 | ) |
| 364 | except NotImplementedError: |
| 365 | pass |
| 366 | |
| 367 | return True, None |
| 368 | |
| 369 | |
| 370 | @dataclass |
| 371 | class WeightTensors: |
| 372 | w1: torch.Tensor |
| 373 | w2: torch.Tensor |
| 374 | w1_scale: torch.Tensor | None |
| 375 | w2_scale: torch.Tensor | None |
| 376 | w1_gs: torch.Tensor | None = None |
| 377 | w2_gs: torch.Tensor | None = None |
| 378 | |
| 379 | def describe(self): |
| 380 | s = "" |
| 381 | s += "== Weight Tensors: \n" |
| 382 | s += f" - {_describe_tensor(self.w1, 'w1')} \n" |
| 383 | s += f" - {_describe_tensor(self.w2, 'w2')} \n" |
| 384 | s += f" - {_describe_tensor(self.w1_scale, 'w1_scale')} \n" |
| 385 | s += f" - {_describe_tensor(self.w2_scale, 'w2_scale')} \n" |
| 386 | s += f" - {_describe_tensor(self.w1_gs, 'w1_gs')} \n" |
| 387 | s += f" - {_describe_tensor(self.w2_gs, 'w2_gs')} \n" |
| 388 | return s |
| 389 | |
| 390 | def is_quantized(self) -> bool: |
| 391 | # or w1_scale is not None? |
| 392 | return ( |
| 393 | is_fp8(self.w1.dtype) |
| 394 | or self.w1.dtype == torch.uint8 |
| 395 | or self.w1.dtype == torch.int8 |
| 396 | ) |
| 397 | |
| 398 | def to_current_device(self): |
| 399 | device = torch.accelerator.current_device_index() |
| 400 | self.w1 = self.w1.to(device=device) |
| 401 | self.w2 = self.w2.to(device=device) |
| 402 | |
| 403 | if self.w1_scale is not None: |
| 404 | self.w1_scale = self.w1_scale.to(device=device) |
| 405 | if self.w2_scale is not None: |
| 406 | self.w2_scale = self.w2_scale.to(device=device) |
| 407 | |
| 408 | if self.w1_gs is not None: |
| 409 | self.w1_gs = self.w1_gs.to(device=device) |
| 410 | if self.w2_gs is not None: |
| 411 | self.w2_gs = self.w2_gs.to(device=device) |
| 412 | |
| 413 | def slice_weights(self, rank: int, num_local_experts: int) -> "WeightTensors": |
| 414 | s = rank * num_local_experts |
| 415 | e = s + num_local_experts |
| 416 | w1 = self.w1[s:e, :, :] |
| 417 | w2 = self.w2[s:e, :, :] |
| 418 | w1_scale = self.w1_scale[s:e, :, :] if self.w1_scale is not None else None |
| 419 | w2_scale = self.w2_scale[s:e, :, :] if self.w2_scale is not None else None |
| 420 | w1_gs = self.w1_gs[s:e] if self.w1_gs is not None else None |
| 421 | w2_gs = self.w2_gs[s:e] if self.w2_gs is not None else None |
| 422 | |
| 423 | return WeightTensors(w1, w2, w1_scale, w2_scale, w1_gs, w2_gs) |
| 424 | |
| 425 | @staticmethod |
| 426 | def make(config: Config) -> "WeightTensors": |
| 427 | (_, w1, w1_scale, w1_gs), (_, w2, w2_scale, w2_gs) = make_test_weights( |
| 428 | e=config.E, |
| 429 | n=config.N, |
| 430 | k=config.K, |
| 431 | in_dtype=config.dtype, |
| 432 | quant_dtype=config.quant_dtype, |
| 433 | block_shape=config.quant_block_shape, |
| 434 | # or config.is_per_out_ch_quant |
| 435 | per_out_ch_quant=config.is_per_act_token_quant, |
| 436 | ) |
| 437 | return WeightTensors( |
| 438 | w1=w1, w2=w2, w1_scale=w1_scale, w2_scale=w2_scale, w1_gs=w1_gs, w2_gs=w2_gs |
| 439 | ) |
| 440 | |
| 441 | |
| 442 | @dataclass |
| 443 | class RankTensors: |
| 444 | hidden_states: torch.Tensor |
| 445 | hidden_states_scale: torch.Tensor | None |
| 446 | |
| 447 | topk_weights: torch.Tensor |
| 448 | topk_ids: torch.Tensor |
| 449 | expert_map: torch.Tensor | None |
| 450 | |
| 451 | def describe(self): |
| 452 | s = "" |
| 453 | s += "== Rank Tensors: \n" |
| 454 | s += f" - {_describe_tensor(self.hidden_states, 'HS')} \n" |
| 455 | s += f" - {_describe_tensor(self.hidden_states_scale, 'HS_scale')} \n" |
| 456 | s += f" - {_describe_tensor(self.topk_weights, 'topk_weights')} \n" |
| 457 | s += f" - {_describe_tensor(self.topk_ids, 'topk_ids')} \n" |
| 458 | s += f" - {_describe_tensor(self.expert_map, 'expert_map')} \n" |
| 459 | return s |
| 460 | |
| 461 | @staticmethod |
| 462 | def make_hidden_states( |
| 463 | config: Config, |
| 464 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 465 | """Return hidden_states.""" |
| 466 | m, k, dtype = (config.M, config.K, config.dtype) |
| 467 | device = torch.accelerator.current_device_index() |
| 468 | a = torch.randn((m, k), device=device, dtype=dtype) / 15.0 |
| 469 | |
| 470 | if config.quant_dtype is None: |
| 471 | return a, None |
| 472 | |
| 473 | # We dequant and use that as hidden_states so the tests are stable. |
| 474 | # quantizing and dequantizing yield slightly different results |
| 475 | # depending on the hardware. Here we, quantize and dequantize |
| 476 | # first - so further quantize and dequantize will yield the same |
| 477 | # values. |
| 478 | if config.is_per_tensor_act_quant: |
| 479 | a_q, a_scales = ops.scaled_fp8_quant(a, use_per_token_if_dynamic=False) |
| 480 | return a_q.float().mul(a_scales).to(dtype), a_scales |
| 481 | |
| 482 | if config.is_per_act_token_quant: |
| 483 | a_q, a_scales = ops.scaled_fp8_quant(a, use_per_token_if_dynamic=True) |
| 484 | return a_q.float().mul(a_scales).to(dtype), None |
| 485 | |
| 486 | assert config.quant_block_shape is not None |
| 487 | block_k = config.quant_block_shape[1] |
| 488 | a_q, a_scales = per_token_cast_to_fp8(a, block_size=block_k) |
| 489 | return a_q.float().view((-1, block_k)).mul(a_scales.view(-1, 1)).view(m, k).to( |
| 490 | dtype |
| 491 | ), None |
| 492 | |
| 493 | @staticmethod |
| 494 | def make(config: Config, pgi: ProcessGroupInfo): |
| 495 | dtype = config.dtype |
| 496 | topk, m, _ = (config.topk, config.M, config.K) |
| 497 | hidden_states, hidden_states_scale = RankTensors.make_hidden_states(config) |
| 498 | |
| 499 | num_local_experts, global_num_experts = (config.num_local_experts, config.E) |
| 500 | score = torch.randn((m, global_num_experts), device="cuda", dtype=dtype) |
| 501 | topk_weights, topk_ids, _ = fused_topk(hidden_states, score, topk, False) |
| 502 | |
| 503 | # distribute topk_ids evenly |
| 504 | device = torch.accelerator.current_device_index() |
| 505 | for mi in range(m): |
| 506 | topk_ids[mi] = torch.randperm(config.E)[:topk] |
| 507 | topk_ids = topk_ids.to(device=device) |
| 508 | |
| 509 | expert_map = None |
| 510 | if config.world_size > 1: |
| 511 | expert_map = torch.full( |
| 512 | (global_num_experts,), fill_value=-1, dtype=torch.int32 |
| 513 | ) |
| 514 | s = pgi.rank * num_local_experts |
| 515 | e = s + num_local_experts |
| 516 | expert_map[s:e] = torch.tensor(list(range(num_local_experts))) |
| 517 | expert_map = expert_map.to(device=device, dtype=torch.int32) |
| 518 | |
| 519 | return RankTensors( |
| 520 | hidden_states=hidden_states, |
| 521 | hidden_states_scale=hidden_states_scale, |
| 522 | topk_weights=topk_weights, |
| 523 | topk_ids=topk_ids, |
| 524 | expert_map=expert_map, |
| 525 | ) |
| 526 | |
| 527 | |
| 528 | def reference_moe_impl( |
| 529 | config: Config, weights: WeightTensors, rank_tensors: RankTensors |
| 530 | ) -> torch.Tensor: |
| 531 | if config.quant_dtype == "nvfp4": |
| 532 | quant_blocksize = 16 |
| 533 | dtype = config.dtype |
| 534 | |
| 535 | w1_q = weights.w1 |
| 536 | w1_blockscale = weights.w1_scale |
| 537 | w1_gs = weights.w1_gs |
| 538 | |
| 539 | w2_q = weights.w2 |
| 540 | w2_blockscale = weights.w2_scale |
| 541 | w2_gs = weights.w2_gs |
| 542 | |
| 543 | a_global_scale = ( |
| 544 | (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) |
| 545 | / torch.amax(rank_tensors.hidden_states.flatten(), dim=-1) |
| 546 | ).to(torch.float32) |
| 547 | |
| 548 | assert w1_gs is not None |
| 549 | assert w2_gs is not None |
| 550 | assert w1_blockscale is not None |
| 551 | assert w2_blockscale is not None |
| 552 | |
| 553 | assert w1_blockscale.shape[1] % 128 == 0 |
| 554 | assert w1_blockscale.shape[2] % 4 == 0 |
| 555 | assert w2_blockscale.shape[1] % 128 == 0 |
| 556 | assert w2_blockscale.shape[2] % 4 == 0 |
| 557 | |
| 558 | a_fp4, a_scale_interleaved = ops.scaled_fp4_quant( |
| 559 | rank_tensors.hidden_states, a_global_scale |
| 560 | ) |
| 561 | |
| 562 | a = dequantize_nvfp4_to_dtype( |
| 563 | a_fp4, |
| 564 | a_scale_interleaved, |
| 565 | a_global_scale, |
| 566 | dtype=dtype, |
| 567 | device=a_fp4.device, |
| 568 | block_size=quant_blocksize, |
| 569 | ) |
| 570 | |
| 571 | e = w1_q.shape[0] |
| 572 | n = w1_q.shape[1] // 2 |
| 573 | k = w2_q.shape[1] |
| 574 | |
| 575 | w1 = torch.zeros((e, 2 * n, k), device="cuda", dtype=dtype) |
| 576 | w2 = torch.zeros((e, k, n), device="cuda", dtype=dtype) |
| 577 | |
| 578 | for idx in range(0, e): |
| 579 | w1[idx] = dequantize_nvfp4_to_dtype( |
| 580 | w1_q[idx], |
| 581 | w1_blockscale[idx], |
| 582 | w1_gs[idx], |
| 583 | dtype=dtype, |
| 584 | device=w1_q.device, |
| 585 | block_size=quant_blocksize, |
| 586 | ) |
| 587 | w2[idx] = dequantize_nvfp4_to_dtype( |
| 588 | w2_q[idx], |
| 589 | w2_blockscale[idx], |
| 590 | w2_gs[idx], |
| 591 | dtype=dtype, |
| 592 | device=w2_q.device, |
| 593 | block_size=quant_blocksize, |
| 594 | ) |
| 595 | a_scale = None |
| 596 | w1_scale = None |
| 597 | w2_scale = None |
| 598 | quant_dtype = None |
| 599 | per_act_token_quant = False |
| 600 | block_shape = None |
| 601 | else: |
| 602 | a = rank_tensors.hidden_states |
| 603 | a_scale = rank_tensors.hidden_states_scale |
| 604 | w1 = weights.w1 |
| 605 | w1_scale = weights.w1_scale |
| 606 | w2 = weights.w2 |
| 607 | w2_scale = weights.w2_scale |
| 608 | quant_dtype = config.quant_dtype |
| 609 | per_act_token_quant = config.is_per_act_token_quant |
| 610 | block_shape = config.quant_block_shape |
| 611 | |
| 612 | return torch_experts( |
| 613 | a=a, |
| 614 | w1=w1, |
| 615 | w2=w2, |
| 616 | topk_weight=rank_tensors.topk_weights, |
| 617 | topk_ids=rank_tensors.topk_ids, |
| 618 | global_num_experts=config.E, |
| 619 | expert_map=None, |
| 620 | activation=config.activation, |
| 621 | w1_scale=w1_scale, |
| 622 | w2_scale=w2_scale, |
| 623 | a1_scale=a_scale, |
| 624 | quant_dtype=quant_dtype, |
| 625 | per_act_token_quant=per_act_token_quant, |
| 626 | block_shape=block_shape, |
| 627 | apply_router_weights_on_input=config.topk == 1 |
| 628 | and config.supports_apply_weight_on_input(), |
| 629 | ) |
| 630 | |
| 631 | |
| 632 | def _make_gscale(num_experts: int) -> torch.Tensor: |
| 633 | return torch.ones( |
| 634 | (num_experts,), |
| 635 | device=torch.accelerator.current_device_index(), |
| 636 | dtype=torch.float32, |
| 637 | ) |
| 638 | |
| 639 | |
| 640 | def make_modular_kernel( |
| 641 | config: Config, |
| 642 | vllm_config: VllmConfig, |
| 643 | quant_config: FusedMoEQuantConfig, |
| 644 | ) -> mk.FusedMoEKernel: |
| 645 | # make moe config |
| 646 | moe_parallel_config: FusedMoEParallelConfig = FusedMoEParallelConfig.make( |
| 647 | tp_size_=get_tensor_model_parallel_world_size(), |
| 648 | pcp_size_=get_pcp_group().world_size, |
| 649 | dp_size_=get_dp_group().world_size, |
| 650 | sp_size_=1, |
| 651 | vllm_parallel_config=vllm_config.parallel_config, |
| 652 | ) |
| 653 | |
| 654 | moe = FusedMoEConfig( |
| 655 | num_experts=config.E, |
| 656 | experts_per_token=config.topk, |
| 657 | hidden_dim=config.K, |
| 658 | intermediate_size=config.N, |
| 659 | num_local_experts=config.num_local_experts, |
| 660 | num_logical_experts=config.E, |
| 661 | moe_parallel_config=moe_parallel_config, |
| 662 | in_dtype=config.dtype, |
| 663 | max_num_tokens=next_power_of_2(config.M), |
| 664 | activation=config.activation, |
| 665 | device=vllm_config.device_config.device, |
| 666 | routing_method=RoutingMethodType.DeepSeekV3, |
| 667 | ) |
| 668 | |
| 669 | prepare_finalize = maybe_make_prepare_finalize( |
| 670 | moe=moe, |
| 671 | quant_config=quant_config, |
| 672 | allow_new_interface=True, |
| 673 | ) |
| 674 | assert prepare_finalize is not None |
| 675 | |
| 676 | fused_experts = make_fused_experts( |
| 677 | config.fused_experts_type, |
| 678 | moe, |
| 679 | quant_config, |
| 680 | prepare_finalize.num_dispatchers(), |
| 681 | config.N, |
| 682 | ) |
| 683 | |
| 684 | modular_kernel = mk.FusedMoEKernel( |
| 685 | prepare_finalize=prepare_finalize, |
| 686 | fused_experts=fused_experts, |
| 687 | ) |
| 688 | |
| 689 | return modular_kernel |
| 690 | |
| 691 | |
| 692 | def _maybe_convert_weights_for_experts( |
| 693 | config: Config, |
| 694 | rank_weights: WeightTensors, |
| 695 | ) -> WeightTensors: |
| 696 | """Convert weights to expert-specific format (e.g., TrtLLM BlockMajorK).""" |
| 697 | from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( |
| 698 | Fp8MoeBackend, |
| 699 | convert_to_fp8_moe_kernel_format, |
| 700 | ) |
| 701 | |
| 702 | fe_type = config.fused_experts_type |
| 703 | fe_name = getattr(fe_type, "__name__", "") |
| 704 | |
| 705 | backend: Fp8MoeBackend | None = None |
| 706 | if fe_name == "TrtLlmFp8ExpertsModular": |
| 707 | backend = Fp8MoeBackend.FLASHINFER_TRTLLM |
| 708 | elif fe_name == "FlashInferExperts": |
| 709 | backend = Fp8MoeBackend.FLASHINFER_CUTLASS |
| 710 | |
| 711 | if backend is None or not rank_weights.is_quantized(): |
| 712 | return rank_weights |
| 713 | |
| 714 | mock_layer = SimpleNamespace( |
| 715 | weight_block_size=config.quant_block_shape, |
| 716 | moe_config=SimpleNamespace( |
| 717 | is_act_and_mul=True, |
| 718 | intermediate_size_per_partition=config.N, |
| 719 | ), |
| 720 | activation=SimpleNamespace(is_gated=True), |
| 721 | ) |
| 722 | |
| 723 | w1, w2, w1_scale, w2_scale = convert_to_fp8_moe_kernel_format( |
| 724 | fp8_backend=backend, |
| 725 | layer=mock_layer, |
| 726 | w13=rank_weights.w1, |
| 727 | w2=rank_weights.w2, |
| 728 | w13_scale=rank_weights.w1_scale, |
| 729 | w2_scale=rank_weights.w2_scale, |
| 730 | w13_input_scale=None, |
| 731 | w2_input_scale=None, |
| 732 | ) |
| 733 | |
| 734 | return WeightTensors( |
| 735 | w1=w1, |
| 736 | w2=w2, |
| 737 | w1_scale=w1_scale, |
| 738 | w2_scale=w2_scale, |
| 739 | w1_gs=rank_weights.w1_gs, |
| 740 | w2_gs=rank_weights.w2_gs, |
| 741 | ) |
| 742 | |
| 743 | |
| 744 | def run_modular_kernel( |
| 745 | pgi: ProcessGroupInfo, |
| 746 | vllm_config: VllmConfig, |
| 747 | config: Config, |
| 748 | weights: WeightTensors, |
| 749 | rank_tensors: RankTensors, |
| 750 | ) -> torch.Tensor: |
| 751 | assert isinstance(config.Ms, int) |
| 752 | assert isinstance(config.topks, int) |
| 753 | |
| 754 | # weights for rank |
| 755 | rank_weights = weights.slice_weights(pgi.rank, config.num_local_experts) |
| 756 | rank_weights = _maybe_convert_weights_for_experts(config, rank_weights) |
| 757 | |
| 758 | if config.quant_dtype == "nvfp4": |
| 759 | gscale = _make_gscale(config.num_local_experts) |
| 760 | else: |
| 761 | gscale = None |
| 762 | |
| 763 | quant_config = FusedMoEQuantConfig.make( |
| 764 | config.quant_dtype, |
| 765 | w1_scale=rank_weights.w1_scale, |
| 766 | w2_scale=rank_weights.w2_scale, |
| 767 | a1_scale=rank_tensors.hidden_states_scale, |
| 768 | g1_alphas=(1 / rank_weights.w1_gs) if rank_weights.w1_gs is not None else None, |
| 769 | g2_alphas=(1 / rank_weights.w2_gs) if rank_weights.w2_gs is not None else None, |
| 770 | a1_gscale=gscale, |
| 771 | a2_gscale=gscale, |
| 772 | block_shape=config.quant_block_shape, |
| 773 | per_act_token_quant=config.is_per_act_token_quant, |
| 774 | per_out_ch_quant=config.is_per_out_ch_quant, |
| 775 | ) |
| 776 | |
| 777 | mk = make_modular_kernel(config, vllm_config, quant_config) |
| 778 | |
| 779 | # impls might update the tensor in place |
| 780 | hidden_states = rank_tensors.hidden_states.clone() |
| 781 | |
| 782 | topk_ids = rank_tensors.topk_ids.to(mk.prepare_finalize.topk_indices_dtype()) |
| 783 | |
| 784 | mk_kwargs = { |
| 785 | "hidden_states": hidden_states, |
| 786 | "w1": rank_weights.w1, |
| 787 | "w2": rank_weights.w2, |
| 788 | "topk_weights": rank_tensors.topk_weights, |
| 789 | "topk_ids": topk_ids, |
| 790 | "activation": config.activation, |
| 791 | "expert_map": rank_tensors.expert_map, |
| 792 | "global_num_experts": config.E, |
| 793 | "apply_router_weight_on_input": config.topk == 1 |
| 794 | and config.supports_apply_weight_on_input(), |
| 795 | } |
| 796 | |
| 797 | num_tokens = rank_tensors.hidden_states.shape[0] |
| 798 | num_tokens_across_dp = torch.tensor( |
| 799 | [num_tokens] * config.world_size, device="cuda", dtype=torch.int |
| 800 | ) |
| 801 | |
| 802 | torch.distributed.barrier() |
| 803 | |
| 804 | with set_forward_context( |
| 805 | None, |
| 806 | vllm_config, |
| 807 | num_tokens=num_tokens, |
| 808 | num_tokens_across_dp=num_tokens_across_dp, |
| 809 | ): |
| 810 | out = mk.apply(**mk_kwargs) |
| 811 | |
| 812 | return out |