@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / f01482408c
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 5, 12:15 UTC

[MoE Refactor][Test] FusedMoE layer test (#24675)

Signed-off-by: Bill Nell <bnell@redhat.com> Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com>

bnellnm committed Apr 6, 2026 at 13:17 UTC f01482408c9e1f8a7e1647aab96d339ba3234cca
6 files changed +1858 -55
.buildkite/test_areas/kernels.yaml
+18
@@ -180,3 +180,21 @@ steps:
180 - pytest -v -s kernels/moe/test_flashinfer_moe.py
181 - pytest -v -s kernels/moe/test_nvfp4_moe.py
182 - pytest -v -s kernels/moe/test_ocp_mx_moe.py
183 +
184 +
185 +- label: Kernels FusedMoE Layer Test (2 H100s)
186 + timeout_in_minutes: 90
187 + device: h100
188 + num_devices: 2
189 + optional: true
190 + commands:
191 + - pytest -v -s kernels/moe/test_moe_layer.py
192 +
193 +
194 +- label: Kernels FusedMoE Layer Test (2 B200s)
195 + timeout_in_minutes: 90
196 + device: b200
197 + num_devices: 2
198 + optional: true
199 + commands:
200 + - pytest -v -s kernels/moe/test_moe_layer.py
tests/kernels/moe/conftest.py new
+14
@@ -0,0 +1,14 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import pytest
4 +
5 +
6 +def pytest_addoption(parser):
7 + parser.addoption(
8 + "--subtests", action="store", type=str, default=None, help="subtest ids"
9 + )
10 +
11 +
12 +@pytest.fixture
13 +def subtests(request):
14 + return request.config.getoption("--subtests")
tests/kernels/moe/modular_kernel_tools/parallel_utils.py
+23 -9
@@ -11,7 +11,11 @@ from torch.multiprocessing import spawn # pyright: ignore[reportPrivateImportUs
11 from typing_extensions import ParamSpec
12
13 from vllm.config import VllmConfig, set_current_vllm_config
14 -from vllm.distributed import init_distributed_environment, initialize_model_parallel
14 +from vllm.distributed import (
15 + cleanup_dist_env_and_memory,
16 + init_distributed_environment,
17 + initialize_model_parallel,
18 +)
19 from vllm.utils.network_utils import get_open_port
20
21 ## Parallel Processes Utils
@@ -36,10 +40,17 @@ def _set_vllm_config(
40
41 temp_file = tempfile.mkstemp()[1]
42
43 + # When DP is enabled, processes are organized as:
44 + # rank = dp_rank * tp_pp_world_size + tp_pp_rank
45 + tp_pp_world_size = vllm_config.parallel_config.world_size
46 + vllm_config.parallel_config.data_parallel_rank = rank // tp_pp_world_size
47 + tp_pp_rank = rank % tp_pp_world_size
48 + vllm_config.parallel_config.rank = tp_pp_rank
49 +
50 with set_current_vllm_config(vllm_config):
51 init_distributed_environment(
41 - world_size=world_size,
42 - rank=rank,
52 + world_size=tp_pp_world_size,
53 + rank=tp_pp_rank,
54 distributed_init_method=f"file://{temp_file}",
55 local_rank=local_rank,
56 backend="nccl",
@@ -59,11 +70,11 @@ def _worker_parallel_launch(
70 world_local_size: int,
71 node_rank: int,
72 init_method: str,
62 - worker: Callable[Concatenate[ProcessGroupInfo, VllmConfig | None, Any, P], None],
73 + worker: Callable[..., None],
74 vllm_config: VllmConfig | None,
75 env_dict: dict | None,
65 - *args: P.args,
66 - **kwargs: P.kwargs,
76 + worker_kwargs: dict[str, Any],
77 + *args: Any,
78 ) -> None:
79 rank = node_rank * world_local_size + local_rank
80 torch.accelerator.set_device_index(local_rank)
@@ -98,14 +109,17 @@ def _worker_parallel_launch(
109 vllm_config,
110 cpu_group,
111 *args,
101 - **kwargs,
112 + **worker_kwargs,
113 )
114 except Exception as ex:
115 print(ex)
116 traceback.print_exc()
117 raise
118 finally:
108 - torch.distributed.destroy_process_group()
119 + if vllm_config is not None:
120 + cleanup_dist_env_and_memory()
121 + else:
122 + torch.distributed.destroy_process_group()
123
124
125 def parallel_launch_with_config(
@@ -116,7 +130,6 @@ def parallel_launch_with_config(
130 *args: P.args,
131 **kwargs: P.kwargs,
132 ) -> None:
119 - assert not kwargs
133 spawn(
134 _worker_parallel_launch,
135 args=(
@@ -127,6 +140,7 @@ def parallel_launch_with_config(
140 worker,
141 vllm_config,
142 env_dict,
143 + kwargs,
144 )
145 + args,
146 nprocs=world_size,
tests/kernels/moe/test_moe_layer.py new
+1727
@@ -0,0 +1,1727 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for the MOE layer.
4 +
5 +Run `pytest tests/kernels/test_moe_layer.py`.
6 +"""
7 +
8 +import functools
9 +import os
10 +import traceback
11 +import types
12 +from collections.abc import Callable
13 +from dataclasses import astuple, dataclass, fields
14 +from itertools import product
15 +from typing import get_args
16 +
17 +import pytest
18 +import torch
19 +
20 +from tests.kernels.moe.modular_kernel_tools.parallel_utils import (
21 + ProcessGroupInfo,
22 + _set_vllm_config,
23 + parallel_launch_with_config,
24 +)
25 +from tests.kernels.moe.utils import TestMLP, make_test_weights, moe_quantize_weights
26 +from vllm.config import (
27 + CompilationConfig,
28 + ParallelConfig,
29 + VllmConfig,
30 + set_current_vllm_config,
31 +)
32 +from vllm.distributed.eplb.eplb_communicator import create_eplb_communicator
33 +from vllm.distributed.eplb.rebalance_execute import rearrange_expert_weights_inplace
34 +from vllm.distributed.parallel_state import (
35 + get_ep_group,
36 + get_eplb_group,
37 +)
38 +from vllm.forward_context import set_forward_context
39 +from vllm.model_executor.layers.fused_moe import FusedMoE, SharedFusedMoE, fused_experts
40 +from vllm.model_executor.layers.fused_moe.activation import MoEActivation
41 +from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
42 +from vllm.model_executor.layers.fused_moe.router.router_factory import (
43 + create_fused_moe_router,
44 +)
45 +from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
46 +from vllm.model_executor.layers.quantization.modelopt import (
47 + ModelOptFp8Config,
48 + ModelOptNvFp4Config,
49 +)
50 +from vllm.platforms import current_platform
51 +from vllm.utils.flashinfer import (
52 + has_flashinfer_nvlink_one_sided,
53 + has_flashinfer_nvlink_two_sided,
54 +)
55 +from vllm.utils.import_utils import has_deep_ep, has_mori, has_nixl_ep
56 +from vllm.utils.math_utils import cdiv
57 +from vllm.utils.torch_utils import set_random_seed
58 +from vllm.v1.worker.workspace import (
59 + init_workspace_manager,
60 + is_workspace_manager_initialized,
61 +)
62 +
63 +fp8_dtype = torch.float8_e4m3fn # current_platform.fp8_dtype
64 +
65 +SHAPE_COMBOS = [
66 + (1, 128, 256),
67 + (32, 1024, 512),
68 + (222, 2048, 2048), # should be big enough to exercise DP chunking
69 +]
70 +
71 +NUM_EXPERTS = [8, 64]
72 +TOP_KS = [2, 6]
73 +
74 +# dp_size, tp_size, use_ep
75 +# Note: DP+TP is not yet supported in the FusedMoE layer.
76 +PARALLEL_COMBOS = [
77 + [1, 2, False],
78 + [1, 4, False],
79 + [2, 1, True],
80 + [4, 1, True],
81 +]
82 +
83 +# TODO: should this even be set manually? let oracles handle this
84 +BACKENDS = ["allgather_reducescatter"]
85 +
86 +if has_mori():
87 + BACKENDS += ["mori"]
88 +
89 +if has_flashinfer_nvlink_two_sided():
90 + BACKENDS += ["flashinfer_nvlink_two_sided"]
91 +
92 +if has_flashinfer_nvlink_one_sided():
93 + BACKENDS += ["flashinfer_nvlink_one_sided"]
94 +
95 +if has_deep_ep():
96 + BACKENDS += ["deepep_low_latency", "deepep_high_throughput"]
97 +
98 +if has_nixl_ep():
99 + BACKENDS += ["nixl_ep"]
100 +
101 +QUANT_METHODS = [
102 + None,
103 + "fp8",
104 + "modelopt_fp8",
105 + "modelopt_fp4",
106 +]
107 +
108 +# Which quantization methods each backend supports.
109 +# fmt: off
110 +BACKEND_SUPPORTED_QUANTS: dict[str, set[str | None]] = {
111 + "allgather_reducescatter": {None, "fp8", "modelopt_fp8", "modelopt_fp4"},
112 + "mori": {None, "fp8", "modelopt_fp8"},
113 + "flashinfer_nvlink_two_sided": {None, "modelopt_fp8", "modelopt_fp4"},
114 + "flashinfer_nvlink_one_sided": {None, "modelopt_fp8", "modelopt_fp4"},
115 + "deepep_low_latency": {None, "fp8", "modelopt_fp8", "modelopt_fp4"},
116 + "deepep_high_throughput": {None, "fp8", "modelopt_fp8", "modelopt_fp4"},
117 + "nixl_ep": {None, "fp8", "modelopt_fp8"},
118 +}
119 +# fmt: on
120 +
121 +# Which quantization methods support EPLB.
122 +# ModelOptFp8MoEMethod inherits supports_eplb=False from FusedMoEMethodBase.
123 +# TODO: double check modelopt fp8
124 +# modelopt_fp4 excluded: get_expert_weights() can't handle NvFP4 packed format.
125 +EPLB_SUPPORTED_QUANTS: list[str | None] = [None, "fp8"]
126 +
127 +# Which backends support EPLB.
128 +# deepep backends fail in get_expert_weights / rearrange_expert_weights_inplace.
129 +# TODO(bnell): check this
130 +EPLB_SUPPORTED_BACKENDS: list[str] = ["allgather_reducescatter"]
131 +
132 +
133 +def maybe_roundup_layer_hidden_size(
134 + hidden_size: int,
135 + act_dtype: torch.dtype,
136 + backend: str | None,
137 +) -> int:
138 + """
139 + Given layer hidden size and MoE configurations, round up hidden_size
140 + if necessary.
141 +
142 + Args:
143 + hidden_size: Layer hidden-size
144 + act_dtype: Data type of the layer activations.
145 + moe_parallel_config: Fused MoE parallelization strategy configuration.
146 +
147 + Return:
148 + Rounded up hidden_size if rounding up is required based on the configs
149 + and all2all backend.
150 + Original hidden size otherwise.
151 + """
152 + if backend == "deepep_high_throughput":
153 + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht import (
154 + DeepEPHTPrepareAndFinalize,
155 + )
156 +
157 + hidden_size = DeepEPHTPrepareAndFinalize.maybe_roundup_layer_hidden_size(
158 + hidden_size, act_dtype
159 + )
160 +
161 + if backend == "deepep_low_latency":
162 + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import (
163 + DeepEPLLPrepareAndFinalize,
164 + )
165 +
166 + hidden_size = DeepEPLLPrepareAndFinalize.maybe_roundup_layer_hidden_size(
167 + hidden_size
168 + )
169 +
170 + return hidden_size
171 +
172 +
173 +def rank_chunk(num: int, r: int, w: int) -> int:
174 + rem = num % w
175 + return (num // w) + (1 if r < rem else 0)
176 +
177 +
178 +def chunk_by_rank(
179 + t: torch.Tensor,
180 + r: int,
181 + w: int,
182 + dim: int = 0,
183 + device: torch.device | None = None,
184 +) -> torch.Tensor:
185 + chunk = cdiv(t.shape[dim], w)
186 + t = t.narrow(dim, r * chunk, chunk)
187 + if device is not None:
188 + t = t.to(device)
189 + return t
190 +
191 +
192 +def maybe_chunk_by_rank(
193 + t: torch.Tensor | None,
194 + r: int,
195 + w: int,
196 + dim: int = 0,
197 + device: torch.device | None = None,
198 +) -> torch.Tensor | None:
199 + if t is not None:
200 + return chunk_by_rank(t, r, w, dim, device)
201 + else:
202 + return t
203 +
204 +
205 +def tp_chunk_gate_up(
206 + w: torch.Tensor,
207 + tp_rank: int,
208 + tp_size: int,
209 + dim: int,
210 + device: torch.device | int | None = None,
211 +) -> torch.Tensor:
212 + """TP-chunk a combined [gate; up] weight, splitting each half separately
213 + so every rank gets a portion of both gate and up."""
214 + half = w.shape[dim] // 2
215 + gate = chunk_by_rank(
216 + w.narrow(dim, 0, half), tp_rank, tp_size, dim=dim, device=device
217 + )
218 + up = chunk_by_rank(
219 + w.narrow(dim, half, half), tp_rank, tp_size, dim=dim, device=device
220 + )
221 + return torch.cat([gate, up], dim=dim)
222 +
223 +
224 +@dataclass
225 +class MoETestConfig:
226 + m: int
227 + n: int
228 + k: int
229 + num_experts: int
230 + top_k: int
231 + in_dtype: torch.dtype
232 + quantization: str | None
233 + use_shared_experts: bool
234 + use_gate: bool
235 + use_routed_input_transform: bool
236 + enable_eplb: bool = False
237 + reduce_results: bool = False
238 + backend: str | None = None
239 + ep_size: int = 1
240 + dp_size: int = 1
241 + tp_size: int = 1
242 +
243 + # TODO: add more error messages
244 + def id(self) -> str:
245 + def proc(s: str) -> str:
246 + return s.removeprefix("torch.")
247 +
248 + id_str = "-".join([proc(str(f)) for f in astuple(self)])
249 + return f"[{id_str}]"
250 +
251 + # TODO: add more error messages
252 + @staticmethod
253 + def from_id(id: str) -> "MoETestConfig":
254 + id = id[1:-1]
255 + str_values = id.split("-")
256 +
257 + def convert(v: str, ty):
258 + if isinstance(ty, types.UnionType):
259 + sub_ty = list(get_args(ty))
260 + assert len(sub_ty) == 2 and types.NoneType in sub_ty
261 + sub_ty.remove(types.NoneType)
262 + return sub_ty[0](v) if v != "None" else None
263 + elif ty is torch.dtype:
264 + ty_val = getattr(torch, v, None)
265 + assert isinstance(ty_val, torch.dtype)
266 + return ty_val
267 + elif ty is bool:
268 + return v == "True"
269 + else:
270 + return ty(v)
271 +
272 + values = tuple(
273 + [convert(v, f.type) for v, f in zip(str_values, fields(MoETestConfig))]
274 + )
275 + return MoETestConfig(*values)
276 +
277 +
278 +def generate_valid_test_configs(
279 + backend: str,
280 + ep_size: int,
281 + dp_size: int,
282 + tp_size: int,
283 + enable_eplb: bool,
284 + verbosity: int = 0,
285 +) -> list[MoETestConfig]:
286 + configs: list[MoETestConfig] = []
287 +
288 + for (
289 + shape,
290 + num_experts,
291 + top_k,
292 + quantization,
293 + use_shared_experts,
294 + use_gate,
295 + use_routed_input_transform,
296 + reduce_results,
297 + ) in product(
298 + SHAPE_COMBOS,
299 + NUM_EXPERTS,
300 + TOP_KS,
301 + QUANT_METHODS,
302 + [False, True], # shared
303 + [False, True], # gate
304 + [False, True], # routed input exform
305 + [False, True], # reduce results
306 + ):
307 + config = MoETestConfig(
308 + shape[0], # m
309 + shape[1], # n
310 + shape[2], # k
311 + num_experts,
312 + top_k,
313 + torch.bfloat16,
314 + quantization,
315 + use_shared_experts,
316 + use_gate,
317 + use_routed_input_transform,
318 + enable_eplb,
319 + reduce_results,
320 + backend,
321 + ep_size,
322 + dp_size,
323 + tp_size,
324 + )
325 +
326 + valid, reason = is_valid_config(config)
327 + if valid:
328 + configs.append(config)
329 + elif verbosity > 1:
330 + print(f"Skipping invalid config {config} - {reason}")
331 +
332 + return configs
333 +
334 +
335 +# TODO: break this up into sections
336 +def is_valid_config(config: MoETestConfig) -> tuple[bool, str | None]:
337 + # routed_input_transform only makes sense with shared_experts (latent MoE)
338 + # TODO: not sure this is true
339 + if config.use_routed_input_transform and not config.use_shared_experts:
340 + return False, "routed_input_transform requires shared_experts"
341 +
342 + # TODO: disable for now
343 + if config.use_routed_input_transform and config.enable_eplb:
344 + return False, "routed_input_transform not supported with EPLB."
345 +
346 + # TODO: disable for now
347 + if config.use_routed_input_transform and config.use_gate:
348 + return (
349 + False,
350 + "routed_input_transform not supported with gate because of "
351 + "padding problems",
352 + )
353 +
354 + # TODO: disable for now
355 + if config.use_routed_input_transform and config.backend in [
356 + "deepep_low_latency",
357 + "deepep_high_throughput",
358 + ]:
359 + return (
360 + False,
361 + "routed_input_transform not supported with DeepEP backends because "
362 + "of padding problems",
363 + )
364 +
365 + # routed_input_transform + quantization + high hidden dimensions
366 + # TODO: Disable >= 2048 w/fp8 + deepep LL for now due to insane errors.
367 + if (
368 + (config.use_routed_input_transform or config.backend == "deepep_low_latency")
369 + and config.quantization is not None
370 + and config.k >= 2048
371 + ):
372 + return (
373 + False,
374 + "routed_input_transform + quantization + higher hidden dimensions "
375 + "leads to large differences.",
376 + )
377 +
378 + # gate requires shared_experts (use_overlapped mode)
379 + # TODO: also not sure this is true
380 + if config.use_gate and not config.use_shared_experts:
381 + return False, "gate requires shared_experts (use_overlapped mode)"
382 +
383 + # Skip modelopt_fp4 if not on B100+ (compute capability 10.0+)
384 + if (
385 + config.quantization == "modelopt_fp4"
386 + and not current_platform.has_device_capability(100)
387 + ):
388 + return False, "modelopt_fp4 not supported on H100+ GPUs"
389 +
390 + # Skip flashinfer_nvlink if not on H100+ (compute capability 10.0+)
391 + if (
392 + config.backend is not None
393 + and config.backend.startswith("flashinfer_nvlink")
394 + and not current_platform.has_device_capability(90)
395 + ):
396 + return False, "flashinfer_nvlink needs an H100+ GPUs"
397 +
398 + # reduce_results incompatibilities
399 + if config.reduce_results and config.use_shared_experts:
400 + return False, "reduce_results=True is not compatible with shared_experts=True"
401 +
402 + if config.reduce_results and config.quantization is not None:
403 + return (
404 + False,
405 + "reduce_results=True only tested with unquantized data types in "
406 + "order to limit number of tests run",
407 + )
408 +
409 + # Backend-specific checks
410 + if config.backend is not None:
411 + supported_quants = BACKEND_SUPPORTED_QUANTS.get(config.backend)
412 + if supported_quants is not None and config.quantization not in supported_quants:
413 + return (
414 + False,
415 + f"{config.backend} does not support quantization={config.quantization}",
416 + )
417 +
418 + if config.backend == "deepep_low_latency":
419 + from vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll import ( # noqa: E501
420 + DeepEPLLPrepareAndFinalize,
421 + )
422 +
423 + if config.k not in DeepEPLLPrepareAndFinalize.SUPPORTED_HIDDEN_SIZES:
424 + return (
425 + False,
426 + f"Skipping unsupported K {config.k} in {config.backend} w/o EP.",
427 + )
428 +
429 + if config.backend == "nixl_ep":
430 + from vllm.model_executor.layers.fused_moe.nixl_ep_prepare_finalize import ( # noqa: E501
431 + NixlEPPrepareAndFinalize,
432 + )
433 +
434 + if config.k not in NixlEPPrepareAndFinalize.SUPPORTED_HIDDEN_SIZES:
435 + return (
436 + False,
437 + f"Skipping unsupported K {config.k} in {config.backend} w/o EP.",
438 + )
439 +
440 + if config.enable_eplb and config.ep_size == 1:
441 + return False, "EPLB requires EP."
442 +
443 + if config.enable_eplb and config.quantization not in EPLB_SUPPORTED_QUANTS:
444 + return False, f"EPLB not supported with {config.quantization} quantization."
445 +
446 + if config.enable_eplb and config.backend not in EPLB_SUPPORTED_BACKENDS:
447 + return False, f"EPLB not supported with {config.backend}."
448 +
449 + world_size = config.tp_size * config.dp_size
450 + if config.reduce_results and world_size == 1:
451 + return False, "reduce_results=True only makes sense for multi-GPU tests"
452 +
453 + if (
454 + config.backend is not None
455 + and config.backend.startswith("flashinfer_nvlink")
456 + and config.ep_size > 1
457 + ):
458 + return False, "flashinfer_nvlink EP not yet supported."
459 +
460 + if config.enable_eplb and config.num_experts % config.dp_size != 0:
461 + return False, "EPLB requires num_experts divisible by ep_size"
462 +
463 + if config.enable_eplb and config.ep_size == 1:
464 + return False, "EPLB only works with EP+DP"
465 +
466 + return True, None
467 +
468 +
469 +def chunk_scales_by_rank(
470 + t: torch.Tensor | None,
471 + r: int,
472 + w: int,
473 + device: torch.device | None = None,
474 +) -> torch.Tensor | None:
475 + if t is not None and t.numel() > 1:
476 + # Calculate start index by summing chunk sizes for all previous ranks
477 + # start = sum(rank_chunk(t.shape[0], i, w) for i in range(r))
478 + # chunk = rank_chunk(t.shape[0], r, w)
479 + # t = t[start:(start + chunk)]
480 + chunk = rank_chunk(t.shape[0], r, w)
481 + t = t[(r * chunk) : max(t.shape[0], (r + 1) * chunk)]
482 +
483 + if t is not None and device is not None:
484 + t = t.to(device)
485 +
486 + return t
487 +
488 +
489 +def chunk_scales(
490 + t: torch.Tensor | None,
491 + start: int,
492 + end: int,
493 + device: torch.device | None = None,
494 +) -> torch.Tensor | None:
495 + if t is not None and t.numel() > 1:
496 + t = t[start:end]
497 +
498 + if t is not None and device is not None:
499 + t = t.to(device)
500 +
501 + return t
502 +
503 +
504 +@dataclass
505 +class QuantizedWeights:
506 + w13_weight: torch.Tensor
507 + w2_weight: torch.Tensor
508 + w13_weight_scale: torch.Tensor | None = None
509 + w2_weight_scale: torch.Tensor | None = None
510 + w13_weight_scale_2: torch.Tensor | None = None
511 + w2_weight_scale_2: torch.Tensor | None = None
512 + w13_input_scale: torch.Tensor | None = None
513 + w2_input_scale: torch.Tensor | None = None
514 +
515 +
516 +def _quantize_fp8_halves(
517 + w1: torch.Tensor,
518 + w2: torch.Tensor,
519 +) -> QuantizedWeights:
520 + """Quantize w13 gate/up halves separately to FP8, producing per-shard scales."""
521 + half = w1.shape[1] // 2
522 + w1q_a, w1s_a, _ = moe_quantize_weights(
523 + w1[:, :half, :], None, fp8_dtype, False, None
524 + )
525 + w1q_b, w1s_b, _ = moe_quantize_weights(
526 + w1[:, half:, :], None, fp8_dtype, False, None
527 + )
528 + assert w1s_a is not None and w1s_b is not None
529 +
530 + w2q, w2s, _ = moe_quantize_weights(w2, None, fp8_dtype, False, None)
531 + assert w2s is not None
532 +
533 + return QuantizedWeights(
534 + w13_weight=torch.cat([w1q_a, w1q_b], dim=1),
535 + w2_weight=w2q,
536 + # Each w1s_x is (E, 1, 1) -> reshape to (E, 1), cat to (E, 2)
537 + w13_weight_scale=torch.cat([w1s_a.view(-1, 1), w1s_b.view(-1, 1)], dim=1),
538 + # w2s is (E, 1, 1) -> reshape to (E,)
539 + w2_weight_scale=w2s.view(-1),
540 + )
541 +
542 +
543 +def quantization_to_quant_dtype(
544 + quantization: str | None,
545 +) -> torch.dtype | str | None:
546 + if quantization is None:
547 + return None
548 + elif quantization in ["fp8", "modelopt_fp8"]:
549 + return fp8_dtype
550 + elif quantization in ["modelopt_fp4"]:
551 + return "nvfp4"
552 + else:
553 + raise NotImplementedError(f"Unsupported quantization: {quantization}")
554 +
555 +
556 +def make_quant_config(
557 + quantization: str | None,
558 + w1: torch.Tensor,
559 + w2: torch.Tensor,
560 + num_experts: int,
561 +) -> tuple[QuantizationConfig | None, QuantizedWeights]:
562 + from vllm.model_executor.layers.quantization.fp8 import Fp8Config
563 +
564 + if quantization is None:
565 + return None, QuantizedWeights(w13_weight=w1, w2_weight=w2)
566 +
567 + if quantization == "fp8":
568 + return Fp8Config(True), _quantize_fp8_halves(w1, w2)
569 +
570 + if quantization == "modelopt_fp8":
571 + qw = _quantize_fp8_halves(w1, w2)
572 + # why?
573 + qw.w13_input_scale = torch.ones(
574 + num_experts, dtype=torch.float32, device=w1.device
575 + )
576 + # why?
577 + qw.w2_input_scale = torch.ones(
578 + num_experts, dtype=torch.float32, device=w2.device
579 + )
580 + quant_config = ModelOptFp8Config(
581 + quant_method="FP8",
582 + is_checkpoint_fp8_serialized=True,
583 + kv_cache_quant_method=None,
584 + exclude_modules=[],
585 + )
586 + return quant_config, qw
587 +
588 + if quantization == "modelopt_fp4":
589 + # Quantize full w13 at once so both gate/up halves share the same
590 + # global scale per expert. process_weights_after_loading uses
591 + # w13_weight_scale_2[:, 0] for the entire tensor, so the two shard
592 + # scales must match.
593 + w1q, w1s, w1gs = moe_quantize_weights(w1, None, "nvfp4", False, None)
594 + assert w1s is not None and w1gs is not None
595 +
596 + w2q, w2s, w2gs = moe_quantize_weights(w2, None, "nvfp4", False, None)
597 + assert w2s is not None and w2gs is not None
598 +
599 + qw = QuantizedWeights(
600 + w13_weight=w1q,
601 + w2_weight=w2q,
602 + w13_weight_scale=w1s,
603 + w2_weight_scale=w2s,
604 + # weight_scale_2 = 1/w_gs: the kernel computes
605 + # g_alphas = a_scale * w_scale_2, and correct dequant needs 1/w_gs.
606 + # Expand per-expert scalar to (E, 2) for the two shards.
607 + w13_weight_scale_2=(1.0 / w1gs).unsqueeze(1).expand(-1, 2).contiguous(),
608 + w2_weight_scale_2=1.0 / w2gs,
609 + w13_input_scale=torch.ones(
610 + (num_experts, 2), dtype=torch.float32, device=w1.device
611 + ),
612 + w2_input_scale=torch.ones(
613 + num_experts, dtype=torch.float32, device=w2.device
614 + ),
615 + )
616 + quant_config = ModelOptNvFp4Config(
617 + is_checkpoint_nvfp4_serialized=True,
618 + kv_cache_quant_algo=None,
619 + exclude_modules=[],
620 + )
621 + return quant_config, qw
622 +
623 + raise NotImplementedError(f"Unsupported quantization: {quantization}")
624 +
625 +
626 +@dataclass
627 +class SharedExpertsConfig:
628 + w1: torch.Tensor
629 + w2: torch.Tensor
630 + w1_s: torch.Tensor | None = None
631 + w2_s: torch.Tensor | None = None
632 + quant_dtype: torch.dtype | str | None = None
633 +
634 +
635 +@dataclass
636 +class MoETestData:
637 + """Container for MOE test data and transforms."""
638 +
639 + w1: torch.Tensor
640 + w2: torch.Tensor
641 + hidden_states: torch.Tensor
642 + router_logits: torch.Tensor
643 + shared_experts_config: SharedExpertsConfig | None
644 + gate: torch.nn.Module | None
645 + routed_input_transform: torch.nn.Module | None
646 + routed_output_transform: torch.nn.Module | None
647 + routed_expert_hidden_size: int
648 +
649 +
650 +class SimpleGate(torch.nn.Module):
651 + """Simple gate module for testing: computes router logits from hidden states."""
652 +
653 + def __init__(
654 + self,
655 + hidden_size: int,
656 + num_experts: int,
657 + dtype: torch.dtype,
658 + device: str = "cuda",
659 + ):
660 + super().__init__()
661 + self.weight = torch.nn.Parameter(
662 + torch.randn(num_experts, hidden_size, device=device, dtype=dtype) / 10
663 + )
664 +
665 + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, None]:
666 + """Returns (router_logits, None) to match expected signature."""
667 + router_logits = torch.nn.functional.linear(hidden_states, self.weight)
668 + return router_logits, None
669 +
670 +
671 +class SimpleRoutedInputTransform(torch.nn.Module):
672 + """Simple linear transform for testing routed input transform
673 + (e.g., latent projection).
674 + """
675 +
676 + def __init__(
677 + self,
678 + in_features: int,
679 + out_features: int,
680 + dtype: torch.dtype,
681 + device: str = "cuda",
682 + ):
683 + super().__init__()
684 + self.weight = torch.nn.Parameter(
685 + torch.randn(out_features, in_features, device=device, dtype=dtype) / 10
686 + )
687 +
688 + def forward(self, x: torch.Tensor) -> torch.Tensor:
689 + return torch.nn.functional.linear(x, self.weight)
690 +
691 +
692 +def create_shared_experts_from_config(
693 + shared_experts_config: SharedExpertsConfig | None,
694 + in_dtype: torch.dtype,
695 + tp_size: int = 1,
696 + tp_rank: int = 0,
697 + device: torch.device | str | None = None,
698 +) -> TestMLP | None:
699 + """Create TestMLP for shared experts from config.
700 +
701 + Args:
702 + shared_experts_config: Configuration for shared experts
703 + in_dtype: Output data type
704 + tp_size: Tensor parallel size (for weight chunking)
705 + tp_rank: Tensor parallel rank (for weight chunking)
706 + device: Device to move weights to (optional)
707 +
708 + Returns:
709 + TestMLP instance or None if config is None
710 + """
711 + if shared_experts_config is None:
712 + return None
713 +
714 + s_w1 = shared_experts_config.w1
715 + s_w2 = shared_experts_config.w2
716 +
717 + # Apply TP chunking if needed
718 + if tp_size > 1:
719 + s_w1 = tp_chunk_gate_up(s_w1, tp_rank, tp_size, dim=1, device=device)
720 + s_w2 = chunk_by_rank(s_w2, tp_rank, tp_size, dim=0, device=device)
721 + else:
722 + s_w1 = s_w1.to(device)
723 + s_w2 = s_w2.to(device)
724 +
725 + return TestMLP(w1=s_w1, w2=s_w2, out_dtype=in_dtype)
726 +
727 +
728 +# Make version that takes a MoETestConfig?
729 +def setup_moe_test_data(
730 + m: int,
731 + k: int,
732 + n: int,
733 + num_experts: int,
734 + in_dtype: torch.dtype,
735 + use_shared_experts: bool,
736 + use_gate: bool,
737 + use_routed_input_transform: bool,
738 + backend: str | None,
739 + device: str = "cuda",
740 +) -> MoETestData:
741 + """Setup test data and transforms for MOE tests.
742 +
743 + Args:
744 + m: Number of tokens
745 + k: Hidden size
746 + n: Intermediate size
747 + num_experts: Number of experts
748 + in_dtype: Data type for tensors
749 + use_shared_experts: Whether to create shared experts config
750 + use_gate: Whether to create gate module
751 + use_routed_input_transform: Whether to create routed input/output transforms
752 + device: Device to create tensors on ("cuda" or "cpu")
753 +
754 + Returns:
755 + MoETestData containing all test data and transforms
756 + """
757 + # For latent MoE: latent_size = k // 2
758 + latent_size = k // 2
759 +
760 + # k = maybe_roundup_layer_hidden_size(k, in_dtype, backend)
761 + # latent_size = maybe_roundup_layer_hidden_size(latent_size, in_dtype, backend)
762 +
763 + # Determine dimensions for routed experts (may be transformed)
764 + # For latent MoE, routed experts operate entirely in latent space
765 + # (k//2). The routed_output_transform then projects back to k before
766 + # adding with shared experts.
767 + # w1: (E, 2*N, latent_size) - input latent_size
768 + # w2: (E, latent_size, N) - output latent_size (fused_experts returns
769 + # same shape as input)
770 + routed_expert_hidden_size = latent_size if use_routed_input_transform else k
771 +
772 + # Create expert weights
773 + (w1, _, _, _), (w2, _, _, _) = make_test_weights(
774 + num_experts,
775 + n,
776 + routed_expert_hidden_size, # Both w1 input and w2 output use latent_size
777 + in_dtype=in_dtype,
778 + )
779 +
780 + # Create shared experts config if needed
781 + if use_shared_experts:
782 + shared_experts_config = SharedExpertsConfig(
783 + w1=torch.randn((k, n * 2), device=device, dtype=in_dtype) / 15,
784 + w2=torch.randn((n, k), device=device, dtype=in_dtype) / 15,
785 + )
786 + else:
787 + shared_experts_config = None
788 +
789 + # Create routed input transform if needed
790 + routed_input_transform = (
791 + SimpleRoutedInputTransform(k, latent_size, in_dtype, device=device)
792 + if use_routed_input_transform
793 + else None
794 + )
795 +
796 + # Create gate if needed
797 + # Note: gate is called AFTER routed_input_transform, so it should expect
798 + # the transformed dimension (latent_size) when routed_input_transform is used
799 + gate_input_dim = latent_size if use_routed_input_transform else k
800 + gate = (
801 + SimpleGate(gate_input_dim, num_experts, in_dtype, device=device)
802 + if use_gate
803 + else None
804 + )
805 +
806 + # Create routed output transform if needed (projects latent space back to original)
807 + routed_output_transform = (
808 + SimpleRoutedInputTransform(latent_size, k, in_dtype, device=device)
809 + if use_routed_input_transform
810 + else None
811 + )
812 +
813 + # Create test inputs
814 + hidden_states = torch.randn((m, k), device=device, dtype=in_dtype) / 10
815 + router_logits = torch.randn((m, num_experts), device=device, dtype=in_dtype)
816 +
817 + return MoETestData(
818 + w1=w1,
819 + w2=w2,
820 + hidden_states=hidden_states,
821 + router_logits=router_logits,
822 + shared_experts_config=shared_experts_config,
823 + gate=gate,
824 + routed_input_transform=routed_input_transform,
825 + routed_output_transform=routed_output_transform,
826 + routed_expert_hidden_size=routed_expert_hidden_size,
827 + )
828 +
829 +
830 +def make_fused_moe_layer(
831 + quantization: str | None,
832 + use_ep: bool,
833 + hidden_size: int,
834 + intermediate_size: int,
835 + in_dtype: torch.dtype,
836 + tp_size: int,
837 + ep_size: int,
838 + dp_size: int,
839 + reduce_results: bool,
840 + w1: torch.Tensor,
841 + w2: torch.Tensor,
842 + top_k: int,
843 + global_num_experts: int,
844 + renormalize: bool = False,
845 + shared_experts: torch.nn.Module | None = None,
846 + use_grouped_topk: bool = False,
847 + topk_group: int | None = None,
848 + num_expert_group: int | None = None,
849 + custom_routing_function: Callable | None = None,
850 + scoring_func: str = "softmax",
851 + routed_scaling_factor: float = 1.0,
852 + e_score_correction_bias: torch.Tensor | None = None,
853 + apply_router_weight_on_input: bool = False,
854 + activation: str = "silu",
855 + indices_type: torch.dtype | None = None,
856 + expert_map: torch.Tensor | None = None,
857 + enable_eplb: bool = False,
858 + expert_load_view: torch.Tensor | None = None,
859 + logical_to_physical_map: torch.Tensor | None = None,
860 + logical_replica_count: torch.Tensor | None = None,
861 + num_redundant_experts: int = 0,
862 + has_bias: bool = False,
863 + gate: torch.nn.Module | None = None,
864 + routed_input_transform: torch.nn.Module | None = None,
865 + routed_output_transform: torch.nn.Module | None = None,
866 + pcp_size: int | None = 1,
867 +) -> tuple[Callable, FusedMoE]:
868 + quant_config, qw = make_quant_config(quantization, w1, w2, global_num_experts)
869 +
870 + kwargs = dict()
871 + if shared_experts is None:
872 + builder = FusedMoE
873 + else:
874 + builder = SharedFusedMoE
875 + kwargs["shared_experts"] = shared_experts
876 +
877 + # Add gate and routed_input_transform if provided
878 + if gate is not None:
879 + kwargs["gate"] = gate
880 + if routed_input_transform is not None:
881 + kwargs["routed_input_transform"] = routed_input_transform
882 +
883 + layer = builder(
884 + num_experts=global_num_experts,
885 + top_k=top_k,
886 + hidden_size=hidden_size,
887 + intermediate_size=intermediate_size,
888 + params_dtype=in_dtype,
889 + reduce_results=reduce_results,
890 + renormalize=renormalize,
891 + use_grouped_topk=use_grouped_topk,
892 + num_expert_group=num_expert_group,
893 + topk_group=topk_group,
894 + quant_config=quant_config,
895 + tp_size=tp_size,
896 + ep_size=ep_size,
897 + dp_size=dp_size,
898 + pcp_size=pcp_size,
899 + prefix="from_forward_context",
900 + custom_routing_function=custom_routing_function,
901 + scoring_func=scoring_func,
902 + routed_scaling_factor=routed_scaling_factor,
903 + e_score_correction_bias=e_score_correction_bias,
904 + apply_router_weight_on_input=apply_router_weight_on_input,
905 + activation=activation,
906 + enable_eplb=enable_eplb,
907 + num_redundant_experts=num_redundant_experts,
908 + has_bias=has_bias,
909 + **kwargs,
910 + )
911 +
912 + for name, value in [
913 + ("w13_weight", qw.w13_weight),
914 + ("w2_weight", qw.w2_weight),
915 + ("w13_weight_scale", qw.w13_weight_scale),
916 + ("w2_weight_scale", qw.w2_weight_scale),
917 + ("w13_weight_scale_2", qw.w13_weight_scale_2),
918 + ("w2_weight_scale_2", qw.w2_weight_scale_2),
919 + ("w13_input_scale", qw.w13_input_scale),
920 + ("w2_input_scale", qw.w2_input_scale),
921 + ]:
922 + if value is not None:
923 + layer.register_parameter(
924 + name, torch.nn.Parameter(value, requires_grad=False)
925 + )
926 +
927 + layer.quant_method.process_weights_after_loading(layer)
928 +
929 + def _moe(
930 + hidden_states: torch.Tensor,
931 + router_logits: torch.Tensor,
932 + ) -> torch.Tensor:
933 + if shared_experts is None:
934 + final_shared_states = None
935 + final_hidden_states = layer(hidden_states, router_logits)
936 + else:
937 + final_shared_states, final_hidden_states = layer(
938 + hidden_states, router_logits
939 + )
940 +
941 + # Apply routed output transform if provided
942 + # (e.g., latent space -> original space)
943 + if routed_output_transform is not None:
944 + final_hidden_states = routed_output_transform(final_hidden_states)
945 +
946 + if shared_experts is not None:
947 + assert not reduce_results
948 + assert final_shared_states is not None
949 + final_hidden_states += final_shared_states
950 +
951 + if not reduce_results and layer.tp_size > 1:
952 + final_hidden_states = layer.maybe_all_reduce_tensor_model_parallel(
953 + final_hidden_states
954 + )
955 +
956 + return final_hidden_states
957 +
958 + return _moe, layer
959 +
960 +
961 +def make_fake_moe_layer(
962 + w1: torch.Tensor,
963 + w2: torch.Tensor,
964 + top_k: int,
965 + global_num_experts: int,
966 + in_dtype: torch.dtype,
967 + quant_dtype: torch.dtype | None,
968 + renormalize: bool = False,
969 + shared_experts_config: SharedExpertsConfig | None = None,
970 + use_grouped_topk: bool = False,
971 + topk_group: int | None = None,
972 + num_expert_group: int | None = None,
973 + custom_routing_function: Callable | None = None,
974 + scoring_func: str = "softmax",
975 + routed_scaling_factor: float = 1.0,
976 + e_score_correction_bias: torch.Tensor | None = None,
977 + apply_router_weight_on_input: bool = False,
978 + activation: str = "silu",
979 + indices_type: torch.dtype | None = None,
980 + expert_map: torch.Tensor | None = None,
981 + enable_eplb: bool = False,
982 + expert_load_view: torch.Tensor | None = None,
983 + logical_to_physical_map: torch.Tensor | None = None,
984 + logical_replica_count: torch.Tensor | None = None,
985 + gate: torch.nn.Module | None = None,
986 + routed_input_transform: torch.nn.Module | None = None,
987 + routed_output_transform: torch.nn.Module | None = None,
988 + use_ep: bool = False,
989 + tp_size: int = 1,
990 + dp_size: int = 1,
991 + ep_size: int = 1,
992 + reduce_results: bool = False,
993 +) -> Callable:
994 + activation = MoEActivation.from_str(activation)
995 +
996 + router = create_fused_moe_router(
997 + top_k=top_k,
998 + global_num_experts=global_num_experts,
999 + # eplb_state=None, # TODO
1000 + renormalize=renormalize,
1001 + use_grouped_topk=use_grouped_topk,
1002 + num_expert_group=num_expert_group,
1003 + topk_group=topk_group,
1004 + custom_routing_function=custom_routing_function,
1005 + scoring_func=scoring_func,
1006 + routed_scaling_factor=routed_scaling_factor,
1007 + e_score_correction_bias=e_score_correction_bias,
1008 + num_fused_shared_experts=0, # TODO
1009 + enable_eplb=enable_eplb,
1010 + # TODO(bnell): once we can construct the MK at init time, we
1011 + # can make this a value.
1012 + indices_type_getter=lambda: indices_type,
1013 + )
1014 +
1015 + if quant_dtype is not None:
1016 + w1, w1_s, _ = moe_quantize_weights(w1, None, quant_dtype, False, None)
1017 + w2, w2_s, _ = moe_quantize_weights(w2, None, quant_dtype, False, None)
1018 + else:
1019 + w1_s = None
1020 + w2_s = None
1021 +
1022 + shared_experts = create_shared_experts_from_config(
1023 + shared_experts_config, in_dtype, 1, 0, "cuda"
1024 + )
1025 +
1026 + quant_config = FusedMoEQuantConfig.make(
1027 + quant_dtype,
1028 + w1_scale=w1_s,
1029 + w2_scale=w2_s,
1030 + )
1031 +
1032 + def _moe(
1033 + hidden_states: torch.Tensor,
1034 + router_logits: torch.Tensor,
1035 + ) -> torch.Tensor:
1036 + # Save original hidden_states for shared experts (before transform)
1037 + original_hidden_states = hidden_states
1038 +
1039 + # Apply routed input transform if provided
1040 + if routed_input_transform is not None:
1041 + hidden_states = routed_input_transform(hidden_states)
1042 +
1043 + # If gate provided, compute router_logits from hidden_states
1044 + # Note: gate operates on transformed hidden_states (after
1045 + # routed_input_transform)
1046 + if gate is not None:
1047 + router_logits, _ = gate(hidden_states)
1048 +
1049 + topk_weights, topk_ids = router.select_experts(
1050 + hidden_states=hidden_states,
1051 + router_logits=router_logits,
1052 + )
1053 +
1054 + # Shared experts use original (untransformed) hidden_states
1055 + if shared_experts is not None:
1056 + shared_output = shared_experts(original_hidden_states)
1057 + else:
1058 + shared_output = None
1059 +
1060 + # Routed experts use transformed hidden_states
1061 + output = fused_experts(
1062 + hidden_states=hidden_states,
1063 + w1=w1,
1064 + w2=w2,
1065 + quant_config=quant_config,
1066 + topk_weights=topk_weights,
1067 + topk_ids=topk_ids,
1068 + inplace=False,
1069 + activation=activation,
1070 + apply_router_weight_on_input=apply_router_weight_on_input,
1071 + global_num_experts=global_num_experts,
1072 + expert_map=expert_map,
1073 + )
1074 +
1075 + # Apply routed output transform if provided
1076 + # (e.g., latent space -> original space)
1077 + if routed_output_transform is not None:
1078 + output = routed_output_transform(output)
1079 +
1080 + if shared_experts is not None:
1081 + assert shared_output is not None
1082 + output += shared_output
1083 +
1084 + # Apply TP/DP reduction if not already reduced
1085 + # if (tp_size > 1 or dp_size > 1):
1086 + # output = tensor_model_parallel_all_reduce(output)
1087 +
1088 + return output
1089 +
1090 + return _moe
1091 +
1092 +
1093 +def _test_body_regular(
1094 + moe_fn: Callable,
1095 + hidden_states: torch.Tensor,
1096 + router_logits: torch.Tensor,
1097 + vllm_config: VllmConfig,
1098 + num_tokens: int,
1099 + num_tokens_across_dp: torch.Tensor,
1100 + **kwargs,
1101 +) -> tuple[torch.Tensor, torch.Tensor]:
1102 + """Regular MoE test body: compare layer output to baseline."""
1103 + baseline_output = kwargs["baseline_output"]
1104 +
1105 + with set_forward_context(
1106 + None,
1107 + vllm_config,
1108 + num_tokens=num_tokens,
1109 + num_tokens_across_dp=num_tokens_across_dp,
1110 + ):
1111 + output = moe_fn(hidden_states, router_logits)
1112 +
1113 + return baseline_output, output
1114 +
1115 +
1116 +def _test_body_eplb(
1117 + moe_fn: Callable,
1118 + moe_layer: FusedMoE,
1119 + hidden_states: torch.Tensor,
1120 + router_logits: torch.Tensor,
1121 + vllm_config: VllmConfig,
1122 + num_tokens: int,
1123 + num_tokens_across_dp: torch.Tensor,
1124 + cpu_group,
1125 + in_dtype: torch.dtype,
1126 + quantization: str | None,
1127 + use_ep: bool,
1128 + tp_size: int,
1129 + ep_size: int,
1130 + dp_size: int,
1131 + w1: torch.Tensor,
1132 + w2: torch.Tensor,
1133 + num_experts: int,
1134 + k: int,
1135 + n: int,
1136 + top_k: int,
1137 + shared_experts,
1138 + reduce_results: bool,
1139 + gate: torch.nn.Module | None,
1140 + routed_input_transform: torch.nn.Module | None,
1141 + routed_output_transform: torch.nn.Module | None,
1142 + **kwargs,
1143 +) -> tuple[torch.Tensor, torch.Tensor]:
1144 + device = torch.accelerator.current_accelerator()
1145 +
1146 + """EPLB test body: compare output before and after expert weight rearrangement."""
1147 + # Get "before" output with original weight arrangement
1148 + with set_forward_context(
1149 + None,
1150 + vllm_config,
1151 + num_tokens=num_tokens,
1152 + num_tokens_across_dp=num_tokens_across_dp,
1153 + ):
1154 + output_before = moe_fn(hidden_states, router_logits)
1155 +
1156 + # Create a fresh FusedMoE layer with enable_eplb=True
1157 + # Delete the original layer's registration so the constructor can
1158 + # re-use the same "from_forward_context" prefix
1159 + cc = vllm_config.compilation_config
1160 + del cc.static_forward_context["from_forward_context"]
1161 + cc.static_all_moe_layers.remove("from_forward_context")
1162 +
1163 + # Determine hidden size for MoE layer
1164 + # When using routed_input_transform, experts operate in latent space
1165 + hidden_size_for_layer = k // 2 if routed_input_transform is not None else k
1166 +
1167 + moe_fn, moe_layer = make_fused_moe_layer(
1168 + quantization=quantization,
1169 + use_ep=use_ep,
1170 + hidden_size=hidden_size_for_layer,
1171 + intermediate_size=n,
1172 + in_dtype=in_dtype,
1173 + tp_size=tp_size,
1174 + ep_size=ep_size,
1175 + dp_size=dp_size,
1176 + reduce_results=reduce_results,
1177 + w1=w1,
1178 + w2=w2,
1179 + top_k=top_k,
1180 + global_num_experts=num_experts,
1181 + shared_experts=shared_experts,
1182 + enable_eplb=True,
1183 + gate=gate,
1184 + routed_input_transform=routed_input_transform,
1185 + routed_output_transform=routed_output_transform,
1186 + )
1187 +
1188 + # Necessary?
1189 + if moe_layer._expert_map is not None:
1190 + moe_layer._expert_map = moe_layer._expert_map.to(device)
1191 +
1192 + # All ranks must generate the same permutation
1193 + initial_indices = torch.arange(num_experts, dtype=torch.long)
1194 + shuffled_indices = initial_indices[torch.randperm(num_experts)]
1195 +
1196 + expert_weights = [list(moe_layer.get_expert_weights())]
1197 +
1198 + communicator = create_eplb_communicator(
1199 + group_coordinator=get_eplb_group(),
1200 + backend=vllm_config.parallel_config.eplb_config.communicator,
1201 + expert_weights=expert_weights[0],
1202 + )
1203 +
1204 + # Rearrange expert weights across EP ranks
1205 + rearrange_expert_weights_inplace(
1206 + old_global_expert_indices=initial_indices.unsqueeze(0),
1207 + new_global_expert_indices=shuffled_indices.unsqueeze(0),
1208 + expert_weights=expert_weights,
1209 + ep_group=cpu_group,
1210 + communicator=communicator,
1211 + )
1212 +
1213 + # Build logical_to_physical_map from shuffled_indices
1214 + # shuffled_indices[physical] = logical, we need the inverse
1215 + logical_to_physical = torch.empty(num_experts, dtype=torch.int32, device=device)
1216 + logical_to_physical[shuffled_indices.to(device)] = torch.arange(
1217 + num_experts, dtype=torch.int32, device=device
1218 + )
1219 +
1220 + moe_layer.set_eplb_state(
1221 + moe_layer_idx=0,
1222 + expert_load_view=torch.zeros(
1223 + (1, num_experts),
1224 + dtype=torch.int32,
1225 + device=device,
1226 + ),
1227 + logical_to_physical_map=logical_to_physical.reshape(num_experts, 1).unsqueeze(
1228 + 0
1229 + ),
1230 + logical_replica_count=torch.ones(
1231 + (1, num_experts),
1232 + dtype=torch.int32,
1233 + device=device,
1234 + ),
1235 + )
1236 +
1237 + moe_layer.eplb_state.should_record_tensor = torch.ones(
1238 + (), dtype=torch.bool, device=device
1239 + )
1240 +
1241 + # Get "after" output with rearranged weights and EPLB routing
1242 + with set_forward_context(
1243 + None,
1244 + vllm_config,
1245 + num_tokens=num_tokens,
1246 + num_tokens_across_dp=num_tokens_across_dp,
1247 + ):
1248 + output_after = moe_fn(hidden_states, router_logits)
1249 +
1250 + return output_before, output_after
1251 +
1252 +
1253 +# TODO: make this take a MoETestConfig
1254 +def _run_one_config(
1255 + vllm_config: VllmConfig,
1256 + ep_size: int,
1257 + dp_size: int,
1258 + tp_size: int,
1259 + dp_rank: int,
1260 + tp_rank: int,
1261 + m: int,
1262 + n: int,
1263 + k: int,
1264 + num_experts: int,
1265 + top_k: int,
1266 + quantization: str | None,
1267 + reduce_results: bool,
1268 + backend: str | None,
1269 + test_body_fn: Callable,
1270 + use_shared_experts: bool,
1271 + use_gate: bool,
1272 + use_routed_input_transform: bool,
1273 + **kwargs,
1274 +) -> None:
1275 + set_random_seed(7)
1276 +
1277 + """Generic test loop that sets up environment and delegates to test_body_fn.
1278 +
1279 + This function is called directly by test_moe_layer and test_moe_layer_eplb
1280 + via parallel_launch_with_config, passing either _test_body_regular or
1281 + _test_body_eplb as the test_body_fn parameter.
1282 + """
1283 + world_size = tp_size * dp_size
1284 + use_ep = ep_size > 1
1285 +
1286 + assert vllm_config.parallel_config.enable_expert_parallel == use_ep
1287 +
1288 + in_dtype = torch.bfloat16
1289 + device = torch.accelerator.current_accelerator()
1290 +
1291 + if not is_workspace_manager_initialized():
1292 + init_workspace_manager(device)
1293 +
1294 + # Create test data and transforms
1295 + test_data = setup_moe_test_data(
1296 + m=m,
1297 + k=k,
1298 + n=n,
1299 + num_experts=num_experts,
1300 + in_dtype=in_dtype,
1301 + use_shared_experts=use_shared_experts,
1302 + use_gate=use_gate,
1303 + use_routed_input_transform=use_routed_input_transform,
1304 + backend=backend,
1305 + device=device,
1306 + )
1307 +
1308 + # Extract data from test_data
1309 + hidden_states = test_data.hidden_states
1310 + router_logits = test_data.router_logits
1311 + w1 = test_data.w1
1312 + w2 = test_data.w2
1313 + shared_experts_config = test_data.shared_experts_config
1314 + gate = test_data.gate
1315 + routed_input_transform = test_data.routed_input_transform
1316 + routed_output_transform = test_data.routed_output_transform
1317 +
1318 + baseline_layer = make_fake_moe_layer(
1319 + w1=w1,
1320 + w2=w2,
1321 + top_k=top_k,
1322 + global_num_experts=num_experts,
1323 + in_dtype=in_dtype,
1324 + quant_dtype=None, # quantization_to_quant_dtype(quantization),
1325 + renormalize=False,
1326 + shared_experts_config=shared_experts_config,
1327 + gate=gate,
1328 + routed_input_transform=routed_input_transform,
1329 + routed_output_transform=routed_output_transform,
1330 + use_ep=use_ep,
1331 + tp_size=tp_size,
1332 + ep_size=ep_size,
1333 + dp_size=dp_size,
1334 + reduce_results=reduce_results,
1335 + )
1336 +
1337 + baseline_output = baseline_layer(hidden_states, router_logits)
1338 +
1339 + del baseline_layer
1340 + torch.accelerator.empty_cache()
1341 +
1342 + with set_current_vllm_config(vllm_config):
1343 + # Chunk weights for EP/TP (after baseline is created)
1344 + if ep_size > 1:
1345 + w1 = chunk_by_rank(w1, dp_rank, dp_size, dim=0, device=device)
1346 + w2 = chunk_by_rank(w2, dp_rank, dp_size, dim=0, device=device)
1347 +
1348 + if tp_size > 1:
1349 + w1 = tp_chunk_gate_up(w1, tp_rank, tp_size, dim=1, device=device)
1350 + w2 = chunk_by_rank(w2, tp_rank, tp_size, dim=2, device=device)
1351 +
1352 + # Setup shared experts if needed
1353 + shared_experts = create_shared_experts_from_config(
1354 + shared_experts_config, in_dtype, tp_size, tp_rank, device
1355 + )
1356 +
1357 + # Determine hidden size for MoE layer
1358 + # When using routed_input_transform, experts operate in latent space
1359 + hidden_size_for_layer = k // 2 if routed_input_transform is not None else k
1360 +
1361 + # Create initial MoE layer
1362 + moe_fn, moe_layer = make_fused_moe_layer(
1363 + quantization=quantization,
1364 + use_ep=use_ep,
1365 + hidden_size=hidden_size_for_layer,
1366 + intermediate_size=n,
1367 + in_dtype=in_dtype,
1368 + tp_size=tp_size,
1369 + ep_size=ep_size,
1370 + dp_size=dp_size,
1371 + reduce_results=reduce_results,
1372 + w1=w1,
1373 + w2=w2,
1374 + top_k=top_k,
1375 + global_num_experts=num_experts,
1376 + shared_experts=shared_experts,
1377 + gate=gate,
1378 + routed_input_transform=routed_input_transform,
1379 + routed_output_transform=routed_output_transform,
1380 + )
1381 +
1382 + # Necessary?
1383 + if moe_layer._expert_map is not None:
1384 + moe_layer._expert_map = moe_layer._expert_map.to(device)
1385 +
1386 + num_tokens = m
1387 + num_tokens_across_dp = torch.tensor(
1388 + [num_tokens] * world_size,
1389 + device=device,
1390 + dtype=torch.int,
1391 + )
1392 +
1393 + # Call the test body function with all necessary context
1394 + expected, actual = test_body_fn(
1395 + moe_fn=moe_fn,
1396 + moe_layer=moe_layer,
1397 + hidden_states=hidden_states,
1398 + router_logits=router_logits,
1399 + vllm_config=vllm_config,
1400 + num_tokens=num_tokens,
1401 + num_tokens_across_dp=num_tokens_across_dp,
1402 + in_dtype=in_dtype,
1403 + quantization=quantization,
1404 + use_ep=use_ep,
1405 + tp_size=tp_size,
1406 + ep_size=ep_size,
1407 + dp_size=dp_size,
1408 + w1=w1,
1409 + w2=w2,
1410 + num_experts=num_experts,
1411 + k=k,
1412 + n=n,
1413 + m=m,
1414 + top_k=top_k,
1415 + shared_experts=shared_experts,
1416 + reduce_results=reduce_results,
1417 + gate=gate,
1418 + routed_input_transform=routed_input_transform,
1419 + routed_output_transform=routed_output_transform,
1420 + baseline_output=baseline_output,
1421 + **kwargs,
1422 + )
1423 +
1424 + # Common tolerance logic
1425 + # TODO: consider associating tolerances with quant methods.
1426 + if quantization is None:
1427 + if k >= 2048:
1428 + atol, rtol = 7.6e-2, 7.6e-2
1429 + else:
1430 + atol, rtol = 3.5e-2, 3.5e-2
1431 + elif quantization in ("fp8", "modelopt_fp8"):
1432 + if k >= 2048:
1433 + atol, rtol = 7.6e-2, 7.6e-2
1434 + else:
1435 + atol, rtol = 6e-2, 6e-2
1436 + elif quantization == "modelopt_fp4":
1437 + atol = rtol = 1e-1 + k * 5e-4
1438 + else:
1439 + atol, rtol = 6e-2, 6e-2
1440 +
1441 + torch.accelerator.synchronize() # TODO: Is this needed?
1442 + torch.testing.assert_close(expected, actual, atol=atol, rtol=rtol)
1443 +
1444 +
1445 +# Test for non-parallel cases (world_size == 1) - backend doesn't matter
1446 +@pytest.mark.parametrize("m, n, k", SHAPE_COMBOS)
1447 +@pytest.mark.parametrize("num_experts", NUM_EXPERTS)
1448 +@pytest.mark.parametrize("top_k", TOP_KS)
1449 +@pytest.mark.parametrize("quantization", QUANT_METHODS)
1450 +@pytest.mark.parametrize("use_shared_experts", [False, True])
1451 +@pytest.mark.parametrize("use_gate", [False, True])
1452 +@pytest.mark.parametrize("use_routed_input_transform", [False, True])
1453 +def test_moe_layer_no_parallel(
1454 + m: int,
1455 + n: int,
1456 + k: int,
1457 + num_experts: int,
1458 + top_k: int,
1459 + quantization: str | None,
1460 + use_shared_experts: bool,
1461 + use_gate: bool,
1462 + use_routed_input_transform: bool,
1463 + monkeypatch,
1464 +):
1465 + """Test MoE layer without parallelism (dp_size=1, tp_size=1, use_ep=False)."""
1466 +
1467 + if os.environ.get("VLLM_LOGGING_LEVEL") is None:
1468 + monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR")
1469 +
1470 + test_config = MoETestConfig(
1471 + m,
1472 + n,
1473 + k,
1474 + num_experts,
1475 + top_k,
1476 + torch.bfloat16,
1477 + quantization,
1478 + use_shared_experts,
1479 + use_gate,
1480 + use_routed_input_transform,
1481 + )
1482 +
1483 + valid, reason = is_valid_config(test_config)
1484 + if not valid:
1485 + pytest.skip(reason)
1486 +
1487 + set_random_seed(7)
1488 +
1489 + parallel_config = ParallelConfig()
1490 + compilation_config = CompilationConfig()
1491 + compilation_config.pass_config.fuse_allreduce_rms = False
1492 +
1493 + vllm_config = VllmConfig(
1494 + parallel_config=parallel_config, compilation_config=compilation_config
1495 + )
1496 +
1497 + # Initialize distributed environment for single GPU
1498 + _set_vllm_config(vllm_config, 1, rank=0, local_rank=0)
1499 +
1500 + _run_one_config(
1501 + vllm_config,
1502 + test_config.ep_size,
1503 + test_config.dp_size,
1504 + test_config.tp_size,
1505 + 0,
1506 + 0,
1507 + test_config.m,
1508 + test_config.n,
1509 + test_config.k,
1510 + test_config.num_experts,
1511 + test_config.top_k,
1512 + test_config.quantization,
1513 + test_config.reduce_results,
1514 + test_config.backend,
1515 + _test_body_regular,
1516 + use_shared_experts=test_config.use_shared_experts,
1517 + use_gate=test_config.use_gate,
1518 + use_routed_input_transform=test_config.use_routed_input_transform,
1519 + )
1520 +
1521 +
1522 +def _test_body_config(test_config: MoETestConfig, cpu_group, **kwargs):
1523 + if not test_config.enable_eplb:
1524 + return _test_body_regular(**kwargs)
1525 + else:
1526 + return _test_body_eplb(**kwargs, cpu_group=cpu_group)
1527 +
1528 +
1529 +def _parallel_worker(
1530 + pgi: ProcessGroupInfo,
1531 + vllm_config: VllmConfig,
1532 + cpu_group,
1533 + test_configs: list[MoETestConfig],
1534 + verbosity: int,
1535 + **kwargs,
1536 +) -> None:
1537 + set_random_seed(7)
1538 +
1539 + total = 0
1540 + passed = 0
1541 + failed = 0
1542 + fail_ids = []
1543 +
1544 + dp_rank = vllm_config.parallel_config.data_parallel_rank
1545 +
1546 + for test_config in test_configs:
1547 + cc = vllm_config.compilation_config
1548 + if "from_forward_context" in cc.static_forward_context:
1549 + del cc.static_forward_context["from_forward_context"]
1550 + cc.static_all_moe_layers.remove("from_forward_context")
1551 +
1552 + tp_rank = pgi.rank % test_config.tp_size
1553 +
1554 + if verbosity > 0:
1555 + print(f"subtest: {test_config.id()}", end="")
1556 +
1557 + try:
1558 + _run_one_config(
1559 + vllm_config,
1560 + test_config.ep_size,
1561 + test_config.dp_size,
1562 + test_config.tp_size,
1563 + dp_rank,
1564 + tp_rank,
1565 + test_config.m,
1566 + test_config.n,
1567 + test_config.k,
1568 + test_config.num_experts,
1569 + test_config.top_k,
1570 + test_config.quantization,
1571 + test_config.reduce_results,
1572 + test_config.backend,
1573 + functools.partial(
1574 + _test_body_config, test_config=test_config, cpu_group=cpu_group
1575 + ),
1576 + use_shared_experts=test_config.use_shared_experts,
1577 + use_gate=test_config.use_gate,
1578 + use_routed_input_transform=test_config.use_routed_input_transform,
1579 + )
1580 + if verbosity > 0:
1581 + print(" PASSED")
1582 + else:
1583 + print(".", end="")
1584 + passed = passed + 1
1585 + except Exception as ex:
1586 + fail_ids.append(test_config.id())
1587 + failed = failed + 1
1588 + if verbosity > 0:
1589 + traceback.print_exc()
1590 + print(f"\n{str(ex)}\nFAILED {ex.__class__}")
1591 + else:
1592 + print("F", end="")
1593 + finally:
1594 + # Note: for some reason DeepEP buffers don't seem to be
1595 + # entirely reusable on B200. In order to work around this
1596 + # we clear the all2all manager's cache after each testpoint.
1597 + cap = current_platform.get_device_capability()
1598 + if (
1599 + cap is not None
1600 + and cap.major == 10
1601 + and (
1602 + test_config.backend == "deepep_low_latency"
1603 + or test_config.backend == "deepep_high_throughput"
1604 + )
1605 + ):
1606 + torch.accelerator.synchronize()
1607 + all2all_manager = get_ep_group().device_communicator.all2all_manager
1608 + if all2all_manager is not None:
1609 + all2all_manager.destroy()
1610 + total = total + 1
1611 +
1612 + skipped = total - (passed + failed)
1613 +
1614 + fails = f"{failed} failed" if failed > 0 else ""
1615 + sep = ", " if fails != "" else ""
1616 + skips = f"{sep}{skipped} skipped" if skipped > 0 else ""
1617 + sep = ", " if skips != "" or fails != "" else ""
1618 + passes = f"{sep}{passed} passed" if passed > 0 else ""
1619 +
1620 + report = (
1621 + f"============= {fails}{skips}{passes} of {total} total tests ============="
1622 + )
1623 +
1624 + sep = "\n" if verbosity == 0 else ""
1625 + print(f"{sep}{report}")
1626 +
1627 + if failed > 0:
1628 + fail_ids_str = "\n".join(fail_ids)
1629 + raise RuntimeError(
1630 + f"\n============= Failed subtests =============\n{fail_ids_str}\n{report}"
1631 + )
1632 +
1633 +
1634 +# TODO: add cudagraphs/torch.compile tests
1635 +@pytest.mark.parametrize("dp_size, tp_size, use_ep", PARALLEL_COMBOS)
1636 +@pytest.mark.parametrize("backend", BACKENDS)
1637 +@pytest.mark.parametrize("enable_eplb", [False, True])
1638 +def test_moe_layer(
1639 + dp_size: int,
1640 + tp_size: int,
1641 + use_ep: bool,
1642 + backend: str,
1643 + enable_eplb: bool,
1644 + monkeypatch,
1645 + pytestconfig,
1646 + subtests,
1647 +):
1648 + """Test MoE layer with parallelism (multi-GPU or TP/EP enabled).
1649 +
1650 + For non-parallel cases (world_size == 1), use test_moe_layer_no_parallel instead.
1651 + """
1652 + num_gpus = current_platform.device_count()
1653 + world_size = tp_size * dp_size
1654 + ep_size = 1 if not use_ep else world_size # or dp_size?
1655 + assert world_size > 1
1656 +
1657 + # Check if enough GPUs available
1658 + if world_size is not None and num_gpus is not None and world_size > num_gpus:
1659 + pytest.skip(f"Not enough GPUs got {num_gpus}, expected {world_size}.")
1660 +
1661 + if enable_eplb and not use_ep:
1662 + pytest.skip("EPLB requires EP.")
1663 +
1664 + verbosity = pytestconfig.getoption("verbose")
1665 +
1666 + test_env = dict()
1667 + test_env["VLLM_MOE_DP_CHUNK_SIZE"] = "128"
1668 + monkeypatch.setenv("VLLM_MOE_DP_CHUNK_SIZE", "128")
1669 + if os.environ.get("VLLM_LOGGING_LEVEL") is None:
1670 + monkeypatch.setenv("VLLM_LOGGING_LEVEL", "ERROR")
1671 +
1672 + # TODO
1673 + # VLLM_FLASHINFER_MOE_BACKEND=latency
1674 + # VLLM_USE_FLASHINFER_MOE_FP16=1
1675 + # VLLM_USE_FLASHINFER_MOE_FP8
1676 + # VLLM_USE_FLASHINFER_MOE_FP4
1677 + # VLLM_USE_FLASHINFER_MOE_INT4
1678 +
1679 + parallel_config = ParallelConfig(
1680 + pipeline_parallel_size=1,
1681 + data_parallel_size=dp_size,
1682 + tensor_parallel_size=tp_size,
1683 + enable_expert_parallel=use_ep,
1684 + all2all_backend=backend,
1685 + enable_eplb=enable_eplb,
1686 + )
1687 +
1688 + compilation_config = CompilationConfig()
1689 + # compilation_config.mode = CompilationMode.NONE # for now
1690 + compilation_config.pass_config.fuse_allreduce_rms = False # for now
1691 +
1692 + vllm_config = VllmConfig(
1693 + parallel_config=parallel_config, compilation_config=compilation_config
1694 + )
1695 +
1696 + test_configs = generate_valid_test_configs(
1697 + backend, ep_size, dp_size, tp_size, enable_eplb, verbosity
1698 + )
1699 +
1700 + if subtests is not None:
1701 + new_test_configs = []
1702 + for subtest in subtests.split(","):
1703 + sub_test_config = MoETestConfig.from_id(subtest)
1704 + if sub_test_config in test_configs:
1705 + new_test_configs.append(sub_test_config)
1706 + else:
1707 + pytest.skip(
1708 + f"subtest config {subtest} does not match any valid test "
1709 + "configuration"
1710 + )
1711 + test_configs = new_test_configs
1712 +
1713 + if len(test_configs) == 0:
1714 + pytest.skip("No supported configs found for this testpoint.")
1715 +
1716 + try:
1717 + parallel_launch_with_config(
1718 + world_size,
1719 + _parallel_worker,
1720 + vllm_config,
1721 + test_env,
1722 + test_configs,
1723 + verbosity,
1724 + )
1725 + finally:
1726 + torch.accelerator.synchronize() # TODO: Is this needed?
1727 + torch.accelerator.empty_cache()
tests/kernels/moe/utils.py
+76 -45
@@ -248,7 +248,7 @@ def make_quantized_test_activations(
248 return a, a_q, a_scale
249
250
251 -def moe_quantize_weights(
251 +def moe_quantize_weights_2d(
252 w: torch.Tensor,
253 w_s: torch.Tensor | None,
254 quant_dtype: torch.dtype | str | None,
@@ -293,6 +293,40 @@ def moe_quantize_weights(
293 return w, w_s, w_gs
294
295
296 +def moe_quantize_weights(
297 + w: torch.Tensor,
298 + w_s: torch.Tensor | None,
299 + quant_dtype: torch.dtype | str | None,
300 + per_token_quant: bool,
301 + block_shape: list[int] | None,
302 +) -> tuple[torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
303 + assert w.dim() == 3
304 + e, rows, cols = w.shape
305 + w_l = [None] * e
306 + w_s_l = [None] * e
307 + w_gs_l = [None] * e
308 + for idx in range(e):
309 + w_l[idx], w_s_l[idx], w_gs_l[idx] = moe_quantize_weights_2d(
310 + w[idx], None, quant_dtype, per_token_quant, block_shape
311 + )
312 +
313 + w = torch.stack(w_l)
314 + w_s = torch.stack(w_s_l)
315 + w_gs = torch.stack(w_gs_l) if e > 0 and w_gs_l[0] is not None else None
316 +
317 + if w_s.ndim == 2:
318 + assert w_s.shape[-1] == 1
319 + w_s = w_s.view(-1, 1, 1)
320 +
321 + if block_shape is not None:
322 + block_n, block_k = block_shape
323 + n_tiles = (rows + block_n - 1) // block_n
324 + k_tiles = (cols + block_k - 1) // block_k
325 + assert w_s.shape == (e, n_tiles, k_tiles)
326 +
327 + return w, w_s, w_gs
328 +
329 +
330 def make_test_weight(
331 e: int,
332 rows: int,
@@ -303,30 +337,11 @@ def make_test_weight(
337 per_out_ch_quant: bool = False,
338 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
339 w_16 = torch.randn((e, rows, cols), device="cuda", dtype=in_dtype) / 15
306 - w_gs = None
340
341 if quant_dtype is not None:
309 - w_l = [None] * e
310 - w_s_l = [None] * e
311 - w_gs_l = [None] * e
312 - for idx in range(e):
313 - w_l[idx], w_s_l[idx], w_gs_l[idx] = moe_quantize_weights(
314 - w_16[idx], None, quant_dtype, per_out_ch_quant, block_shape
315 - )
316 -
317 - w = torch.stack(w_l)
318 - w_s = torch.stack(w_s_l)
319 - if e > 0 and w_gs_l[0] is not None:
320 - w_gs = torch.stack(w_gs_l)
321 - if w_s.ndim == 2:
322 - assert w_s.shape[-1] == 1
323 - w_s = w_s.view(-1, 1, 1)
324 -
325 - if block_shape is not None:
326 - block_n, block_k = block_shape
327 - n_tiles = (rows + block_n - 1) // block_n
328 - k_tiles = (cols + block_k - 1) // block_k
329 - assert w_s.shape == (e, n_tiles, k_tiles)
342 + w, w_s, w_gs = moe_quantize_weights(
343 + w_16, None, quant_dtype, per_out_ch_quant, block_shape
344 + )
345 else:
346 w = w_16
347 w_s = None
@@ -454,7 +469,6 @@ def fused_moe(
469 )
470
471
457 -# CustomOp?
472 class BaselineMM(torch.nn.Module):
473 def __init__(
474 self,
@@ -462,13 +476,22 @@ class BaselineMM(torch.nn.Module):
476 out_dtype: torch.dtype,
477 ):
478 super().__init__()
465 - self.b = b.to(dtype=torch.float32)
479 + self.b = torch.nn.Parameter(b.to(dtype=torch.float32))
480 self.out_dtype = out_dtype
481
482 def forward(self, a: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor | None]:
483 return torch.mm(a.to(dtype=torch.float32), self.b).to(self.out_dtype), None
484
485
486 +class BaselineSiluAndMul(torch.nn.Module):
487 + def __init__(self):
488 + super().__init__()
489 +
490 + def forward(self, x: torch.Tensor) -> torch.Tensor:
491 + d = x.shape[-1] // 2
492 + return torch.nn.functional.silu(x[..., :d]) * x[..., d:]
493 +
494 +
495 class TestMLP(torch.nn.Module):
496 def __init__(
497 self,
@@ -479,7 +502,7 @@ class TestMLP(torch.nn.Module):
502 super().__init__()
503 self.gate_up_proj = BaselineMM(w1, out_dtype)
504 self.down_proj = BaselineMM(w2, out_dtype)
482 - self.act_fn = SiluAndMul()
505 + self.act_fn = BaselineSiluAndMul()
506
507 def forward(self, x):
508 x, _ = self.gate_up_proj(x)
@@ -564,35 +587,24 @@ class RealMLP(torch.nn.Module):
587 return x
588
589
567 -def make_shared_experts(
590 +def make_shared_experts_with_weights(
591 N: int,
592 K: int,
570 - in_dtype: torch.dtype = torch.bfloat16,
593 + in_dtype: torch.dtype,
594 + w1: torch.Tensor,
595 + w2: torch.Tensor,
596 + w1_s: torch.Tensor | None = None,
597 + w2_s: torch.Tensor | None = None,
598 quant_dtype: torch.dtype | str | None = None,
599 ) -> torch.nn.Module:
573 - from vllm.model_executor.layers.quantization.fp8 import Fp8Config
574 -
575 - (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights(
576 - 1,
577 - N,
578 - K,
579 - in_dtype=in_dtype,
580 - quant_dtype=quant_dtype,
581 - )
600 old_dtype = torch.get_default_dtype()
601 try:
602 torch.set_default_dtype(in_dtype)
603 if quant_dtype == torch.float8_e4m3fn:
586 - w1 = w1[0].transpose(0, 1)
587 - w2 = w2[0].transpose(0, 1)
588 - w1_s = w1_s[0].transpose(0, 1) if w1_s is not None else None
589 - w2_s = w2_s[0].transpose(0, 1) if w2_s is not None else None
604 + from vllm.model_executor.layers.quantization.fp8 import Fp8Config
605 +
606 quant_config = Fp8Config(True)
607 else:
592 - w1 = w1[0]
593 - w2 = w2[0]
594 - w1_s = None
595 - w2_s = None
608 quant_config = None
609
610 return RealMLP(K, N, w1, w2, "silu", quant_config, w1_s=w1_s, w2_s=w2_s)
@@ -614,3 +626,22 @@ def modular_triton_fused_moe(
626 TritonExperts(moe_config, quant_config),
627 inplace=False,
628 )
629 +
630 +
631 +def make_shared_experts(
632 + N: int,
633 + K: int,
634 + in_dtype: torch.dtype = torch.bfloat16,
635 + quant_dtype: torch.dtype | str | None = None,
636 +) -> torch.nn.Module:
637 + (_, w1, w1_s, _), (_, w2, w2_s, _) = make_test_weights(
638 + 1,
639 + N,
640 + K,
641 + in_dtype=in_dtype,
642 + quant_dtype=quant_dtype,
643 + )
644 +
645 + return make_shared_experts_with_weights(
646 + N, K, in_dtype, w1, w2, w1_s=w1_s, w2_s=w2_s, quant_dtype=quant_dtype
647 + )
vllm/model_executor/layers/fused_moe/experts/trtllm_bf16_moe.py
-1
@@ -76,7 +76,6 @@ class TrtLlmBf16Experts(mk.FusedMoEExpertsMonolithic):
76 activation_key: QuantKey | None,
77 ) -> bool:
78 return routing_method in [
79 - RoutingMethodType.Default,
79 RoutingMethodType.DeepSeekV3,
80 RoutingMethodType.Llama4,
81 RoutingMethodType.Renormalize,