diff --git a/tests/distributed/test_eplb_quant_scale_consistency.py b/tests/distributed/test_eplb_quant_scale_consistency.py new file mode 100644 index 0000000000..3b50b315f0 --- /dev/null +++ b/tests/distributed/test_eplb_quant_scale_consistency.py @@ -0,0 +1,299 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""EPLB rearrangement consistency tests for quant-method derived state. + +EPLB rearranges every registered Parameter of a RoutedExperts in place, +sliced along dim 0 (see RoutedExperts.get_expert_weights). Quant methods +must therefore register all derived per-expert tensors as Parameters and +alias the same storage in their FusedMoEQuantConfig, so the kernels observe +rearranged values with no extra bookkeeping. These tests verify that +contract: simulating a rearrangement on the registered Parameters must be +indistinguishable from loading a checkpoint with permuted experts. +""" + +import pytest +import torch + +from vllm.config import VllmConfig, set_current_vllm_config +from vllm.model_executor.layers.fused_moe.activation import MoEActivation +from vllm.model_executor.layers.fused_moe.config import ( + FusedMoEConfig, + FusedMoEParallelConfig, + RoutingMethodType, +) +from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4 import ( # noqa: E501 + CompressedTensorsW4A4Nvfp4MoEMethod, +) +from vllm.platforms import current_platform + +EPLB_NVFP4_BACKENDS = ["flashinfer_cutedsl", "flashinfer_trtllm"] + +NUM_EXPERTS = 8 +HIDDEN_SIZE = 128 +INTERMEDIATE_SIZE = 256 +EXPERT_PERMUTATION = [3, 0, 5, 1, 7, 2, 6, 4] + +QUANT_CONFIG_TENSORS = ( + "w1_scale", + "w2_scale", + "g1_alphas", + "g2_alphas", + "a1_gscale", + "a2_gscale", +) + + +class _RoutedExpertsStub(torch.nn.Module): + """Minimal ``RoutedExperts`` stand-in for NVFP4 weight processing.""" + + def __init__(self, moe_config: FusedMoEConfig) -> None: + super().__init__() + self.moe_config = moe_config + self.activation = moe_config.activation + + def _expert_routing_tables( + self, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: + return None + + +def _make_moe_config(backend: str) -> FusedMoEConfig: + parallel_config = FusedMoEParallelConfig( + tp_size=1, + pcp_size=1, + dp_size=1, + ep_size=1, + tp_rank=0, + pcp_rank=0, + dp_rank=0, + ep_rank=0, + sp_size=1, + use_ep=True, + all2all_backend="allgather_reducescatter", + enable_eplb=True, + ) + return FusedMoEConfig( + num_experts=NUM_EXPERTS, + experts_per_token=2, + hidden_dim=HIDDEN_SIZE, + intermediate_size=INTERMEDIATE_SIZE, + num_local_experts=NUM_EXPERTS, + num_logical_experts=NUM_EXPERTS, + activation=MoEActivation.SILU, + device="cuda", + routing_method=RoutingMethodType.TopK, + moe_parallel_config=parallel_config, + in_dtype=torch.bfloat16, + moe_backend=backend, + ) + + +def _make_raw_weights( + device: torch.device, generator: torch.Generator +) -> dict[str, torch.Tensor]: + """Random raw NVFP4 checkpoint tensors, indexable per-expert on dim 0.""" + e, h, i = NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE + + def packed(*shape: int) -> torch.Tensor: + return torch.randint( + 0, 256, shape, dtype=torch.uint8, device=device, generator=generator + ) + + def block_scale(*shape: int) -> torch.Tensor: + return torch.randint( + 0, 128, shape, dtype=torch.uint8, device=device, generator=generator + ).view(torch.float8_e4m3fn) + + def global_scale(*shape: int) -> torch.Tensor: + return torch.rand(shape, device=device, generator=generator) + 0.5 + + return { + "w13_weight_packed": packed(e, 2 * i, h // 2), + "w2_weight_packed": packed(e, h, i // 2), + "w13_weight_scale": block_scale(e, 2 * i, h // 16), + "w2_weight_scale": block_scale(e, h, i // 16), + "w13_weight_global_scale": global_scale(e, 2), + "w2_weight_global_scale": global_scale(e), + "w13_input_global_scale": global_scale(e, 2), + "w2_input_global_scale": global_scale(e), + } + + +def _build_processed_layer( + backend: str, raw: dict[str, torch.Tensor], device: torch.device +) -> tuple[CompressedTensorsW4A4Nvfp4MoEMethod, _RoutedExpertsStub]: + """Create the method + layer, load ``raw`` and run the real + process_weights_after_loading (which builds the flashinfer kernel).""" + moe_config = _make_moe_config(backend) + method = CompressedTensorsW4A4Nvfp4MoEMethod(moe_config, "layer.0", use_a16=False) + layer = _RoutedExpertsStub(moe_config).to(device) + method.create_weights( + layer, + num_experts=NUM_EXPERTS, + hidden_size=HIDDEN_SIZE, + intermediate_size_per_partition=INTERMEDIATE_SIZE, + params_dtype=torch.bfloat16, + weight_loader=lambda *args, **kwargs: None, + ) + layer = layer.to(device) + with torch.no_grad(): + for name, value in raw.items(): + getattr(layer, name).copy_(value) + method.process_weights_after_loading(layer) + return method, layer + + +def _assert_tensors_equal( + actual: torch.Tensor, expected: torch.Tensor, name: str +) -> None: + if actual.dtype == torch.float8_e4m3fn: + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)), name + elif actual.dtype == torch.uint8: + assert torch.equal(actual, expected), name + else: + torch.testing.assert_close(actual, expected, msg=lambda m: f"{name}: {m}") + + +def _simulate_eplb_rearrangement(layer: torch.nn.Module, perm: torch.Tensor) -> None: + """Permute experts of every registered Parameter in place, the way + EPLB's in-place rearrangement moves expert slices along dim 0.""" + with torch.no_grad(): + for _, param in layer.named_parameters(): + param.copy_(param[perm]) + + +def _ensure_world1_distributed() -> None: + """The enable_eplb config makes weight processing all-reduce activation + scale amaxes over the EP group, so a (single-rank) EP group must exist.""" + from vllm.distributed.parallel_state import ( + ensure_model_parallel_initialized, + init_distributed_environment, + ) + + if not torch.distributed.is_initialized(): + init_distributed_environment( + world_size=1, + rank=0, + distributed_init_method="tcp://127.0.0.1:0", + local_rank=0, + ) + ensure_model_parallel_initialized(1, 1) + + +@pytest.mark.parametrize("backend", EPLB_NVFP4_BACKENDS) +def test_nvfp4_eplb_rearrangement_matches_reload(backend: str) -> None: + if not ( + current_platform.is_cuda() and current_platform.is_device_capability_family(100) + ): + pytest.skip("NVFP4 CuteDSL/TRTLLM MoE backends require Blackwell (SM100).") + + device = torch.device("cuda:0") + torch.accelerator.set_device_index(device) + perm = torch.tensor(EXPERT_PERMUTATION, device=device) + + with set_current_vllm_config(VllmConfig()): + _ensure_world1_distributed() + generator = torch.Generator(device=device).manual_seed(1234) + raw = _make_raw_weights(device, generator) + raw_permuted = {name: value[perm].contiguous() for name, value in raw.items()} + + try: + method, layer = _build_processed_layer(backend, raw, device) + except ValueError as exc: + pytest.skip(f"{backend} NVFP4 MoE backend unavailable: {exc}") + ref_method, ref_layer = _build_processed_layer(backend, raw_permuted, device) + + # The EPLB contract: every registered Parameter must be an + # expert-major contiguous tensor so get_expert_weights can view it + # as (E, -1) and rearrange expert slices in place. + for name, param in layer.named_parameters(): + assert param.is_contiguous(), f"{name} is not contiguous" + assert param.shape[0] == NUM_EXPERTS, ( + f"{name} is not expert-major: {tuple(param.shape)}" + ) + + # Derived per-expert scales must live in registered Parameters that + # the quant config aliases; anything else goes stale on rearrangement. + quant_config = method.moe_quant_config + assert quant_config is not None + params = dict(layer.named_parameters()) + assert quant_config.g1_alphas.data_ptr() == ( + params["w13_weight_scale_2"].data_ptr() + ) + assert quant_config.g2_alphas.data_ptr() == ( + params["w2_weight_scale_2"].data_ptr() + ) + assert quant_config.w1_scale.data_ptr() == ( + params["w13_weight_scale"].data_ptr() + ) + assert quant_config.w2_scale.data_ptr() == params["w2_weight_scale"].data_ptr() + + _simulate_eplb_rearrangement(layer, perm) + + # After rearrangement, all registered state must equal what a fresh + # load of the permuted experts produces. + ref_params = dict(ref_layer.named_parameters()) + assert params.keys() == ref_params.keys() + for name in params: + _assert_tensors_equal(params[name], ref_params[name], name) + + # And so must the kernel-visible quant config tensors (the CuteDSL + # MMA scale views alias the registered Parameters' storage). + ref_quant_config = ref_method.moe_quant_config + assert ref_quant_config is not None + for name in QUANT_CONFIG_TENSORS: + actual = getattr(quant_config, name) + expected = getattr(ref_quant_config, name) + assert (actual is None) == (expected is None), name + if actual is not None: + _assert_tensors_equal(actual, expected, f"quant_config.{name}") + + +def test_fp8_per_tensor_alphas_registered_and_aliased() -> None: + """CPU-only: the fp8 per-tensor oracle must register the fused + (w_scale * a_scale) products as layer Parameters aliased by the quant + config, so EPLB rearrangement keeps kernels consistent.""" + from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( + Fp8MoeBackend, + make_fp8_moe_quant_config, + ) + + e = NUM_EXPERTS + generator = torch.Generator().manual_seed(1234) + + def scales() -> tuple[torch.Tensor, torch.Tensor]: + w_scale = torch.rand((e,), generator=generator) + 0.5 + a_scale = torch.rand((), generator=generator) + 0.5 + return w_scale, a_scale.expand(e).contiguous() + + w1_scale, a1_scale = scales() + w2_scale, a2_scale = scales() + + layer = torch.nn.Module() + quant_config = make_fp8_moe_quant_config( + fp8_backend=Fp8MoeBackend.FLASHINFER_CUTLASS, + w1_scale=w1_scale, + w2_scale=w2_scale, + a1_scale=a1_scale, + a2_scale=a2_scale, + layer=layer, + ) + + params = dict(layer.named_parameters()) + assert quant_config.g1_alphas.data_ptr() == params["g1_alphas"].data_ptr() + assert quant_config.g2_alphas.data_ptr() == params["g2_alphas"].data_ptr() + + perm = torch.tensor(EXPERT_PERMUTATION) + ref_config = make_fp8_moe_quant_config( + fp8_backend=Fp8MoeBackend.FLASHINFER_CUTLASS, + w1_scale=w1_scale[perm].contiguous(), + w2_scale=w2_scale[perm].contiguous(), + a1_scale=a1_scale, + a2_scale=a2_scale, + layer=None, + ) + _simulate_eplb_rearrangement(layer, perm) + + torch.testing.assert_close(quant_config.g1_alphas, ref_config.g1_alphas) + torch.testing.assert_close(quant_config.g2_alphas, ref_config.g2_alphas) diff --git a/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py index 1688c4c648..d8e449ec0a 100644 --- a/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py +++ b/tests/kernels/moe/test_flashinfer_cutedsl_nvfp4_moe.py @@ -193,7 +193,22 @@ def test_flashinfer_cutedsl_fp4_moe( hidden_states, score, topk, renormalize=False ) - fake_layer = SimpleNamespace(activation=activation) + activation = MoEActivation.RELU2_NO_MUL + moe_config = FusedMoEConfig( + num_experts=e, + experts_per_token=topk, + hidden_dim=k, + intermediate_size=n, + num_local_experts=e, + num_logical_experts=e, + activation=activation, + device="cuda", + moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), + in_dtype=dtype, + routing_method=RoutingMethodType.TopK, + max_num_tokens=next_power_of_2(m), + ) + fake_layer = SimpleNamespace(activation=activation, moe_config=moe_config) a1_scale = torch.ones(1, device="cuda", dtype=torch.float32) a2_scale = torch.ones(1, device="cuda", dtype=torch.float32) ( @@ -230,20 +245,6 @@ def test_flashinfer_cutedsl_fp4_moe( gemm1_beta=beta, gemm1_clamp_limit=limit, ) - moe_config = FusedMoEConfig( - num_experts=e, - experts_per_token=topk, - hidden_dim=k, - intermediate_size=n, - num_local_experts=e, - num_logical_experts=e, - activation=activation, - device="cuda", - moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), - in_dtype=dtype, - routing_method=RoutingMethodType.TopK, - max_num_tokens=next_power_of_2(m), - ) cutedsl_experts = mk.FusedMoEKernel( maybe_make_prepare_finalize( diff --git a/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py b/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py index 5a73774396..9bec23b1b4 100644 --- a/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py +++ b/tests/quantization/test_trtllm_nvfp4_hidden_dim_padding.py @@ -19,7 +19,12 @@ def test_shared_nvfp4_input_scales_have_writable_storage(monkeypatch): monkeypatch.setattr(flashinfer_fp4_moe, "swizzle_blockscale", lambda x: x) num_experts = 3 - layer = SimpleNamespace(activation=SimpleNamespace(is_gated=False)) + layer = SimpleNamespace( + activation=SimpleNamespace(is_gated=False), + moe_config=SimpleNamespace( + moe_parallel_config=SimpleNamespace(enable_eplb=False) + ), + ) w13 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8) w2 = torch.zeros((num_experts, 2, 1), dtype=torch.uint8) w13_scale = torch.zeros((num_experts, 2, 1), dtype=torch.float8_e4m3fn) diff --git a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py index c7813b8548..fbeaa198d0 100644 --- a/tests/v1/worker/test_gpu_model_runner_v2_eplb.py +++ b/tests/v1/worker/test_gpu_model_runner_v2_eplb.py @@ -108,8 +108,10 @@ def test_v2_load_model_registers_moe_with_eplb(monkeypatch): ) monkeypatch.setattr( eplb, - "is_mixture_of_experts", - lambda loaded_model: getattr(loaded_model, "is_moe", False), + "get_mixture_of_experts_model", + lambda loaded_model: ( + loaded_model if getattr(loaded_model, "is_moe", False) else None + ), ) runner = _make_runner(is_last_pp_rank=False) @@ -138,7 +140,7 @@ def test_v2_load_model_with_dummy_weights_skips_eplb_registration(monkeypatch): "init_model_state", lambda *args: SimpleNamespace(num_new_sampled_tokens_per_step=1), ) - monkeypatch.setattr(eplb, "is_mixture_of_experts", lambda *_: True) + monkeypatch.setattr(eplb, "get_mixture_of_experts_model", lambda model: model) runner = _make_runner(is_last_pp_rank=False) mrv2.GPUModelRunner.load_model(runner, load_dummy_weights=True) @@ -153,7 +155,7 @@ def test_v2_setup_eplb_from_mapping_rebuilds_state(monkeypatch): FakeEplbState.instances.clear() FakeEplbState.from_mapping_kwargs = None monkeypatch.setattr(eplb, "EplbState", FakeEplbState) - monkeypatch.setattr(eplb, "is_mixture_of_experts", lambda *_: True) + monkeypatch.setattr(eplb, "get_mixture_of_experts_model", lambda model: model) runner = _make_runner(model=SimpleNamespace(is_moe=True)) mapping = torch.tensor([[0, 1, 2, 3]], dtype=torch.int64) diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 7cb57a44eb..7501943668 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -628,6 +628,17 @@ def make_fp8_moe_quant_config( and block_shape is None ): assert a1_scale is not None and a2_scale is not None + g1_alphas = w1_scale * a1_scale + g2_alphas = w2_scale * a2_scale + if layer is not None: + layer.register_parameter( + "g1_alphas", torch.nn.Parameter(g1_alphas, requires_grad=False) + ) + layer.register_parameter( + "g2_alphas", torch.nn.Parameter(g2_alphas, requires_grad=False) + ) + g1_alphas = layer.g1_alphas + g2_alphas = layer.g2_alphas return fp8_w8a8_moe_quant_config( w1_scale=w1_scale, w2_scale=w2_scale, @@ -637,8 +648,8 @@ def make_fp8_moe_quant_config( a2_scale=a2_scale, a1_gscale=(1.0 / a1_scale), a2_gscale=(1.0 / a2_scale), - g1_alphas=(w1_scale * a1_scale).squeeze(), - g2_alphas=(w2_scale * a2_scale).squeeze(), + g1_alphas=g1_alphas, + g2_alphas=g2_alphas, gemm1_clamp_limit=swiglu_limit, ) # MXFP8 (block [1, 32]) dispatches to the mxfp8 activation quant. Scales are diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index d540e3ae0c..de3fb6fa58 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -19,6 +19,7 @@ from vllm.model_executor.layers.fused_moe.config import ( ) from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import ( + nvfp4_swizzled_scale_to_cutedsl_mma_view, prepare_nvfp4_moe_layer_for_fi_or_cutlass, prepare_nvfp4_moe_layer_for_flashinfer_cutedsl, ) @@ -506,6 +507,10 @@ def make_nvfp4_moe_quant_config( gemm1_clamp_limit=swiglu_limit, ) + if backend == NvFp4MoeBackend.FLASHINFER_CUTEDSL: + w13_scale = nvfp4_swizzled_scale_to_cutedsl_mma_view(w13_scale) + w2_scale = nvfp4_swizzled_scale_to_cutedsl_mma_view(w2_scale) + # Pass w13_scale_2 / w2_scale_2 directly as g1/g2_alphas. # The expert's process_weights_after_loading will fuse activation # scales in-place. Since the quant config references the same tensor diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py index 9ae807fcb6..4980a1a511 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_nvfp4.py @@ -15,6 +15,7 @@ from vllm.model_executor.layers.fused_moe.config import ( FusedMoEQuantConfig, ) from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import ( + NvFp4MoeBackend, convert_to_nvfp4_moe_kernel_format, is_global_sf_supported_for_nvfp4_backend, make_nvfp4_moe_kernel, @@ -30,6 +31,15 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.model_executor.utils import replace_parameter, set_weight_attrs +# NVFP4 backends that have been verified to support EPLB for this quantization recipe. +_EPLB_SUPPORTED_NVFP4_BACKENDS = frozenset( + { + NvFp4MoeBackend.FLASHINFER_CUTEDSL, + NvFp4MoeBackend.FLASHINFER_CUTEDSL_BATCHED, + NvFp4MoeBackend.FLASHINFER_TRTLLM, + } +) + logger = init_logger(__name__) @@ -54,6 +64,10 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): self.nvfp4_backend ) + @property + def supports_eplb(self) -> bool: + return self.nvfp4_backend in _EPLB_SUPPORTED_NVFP4_BACKENDS + def create_weights( self, layer: torch.nn.Module, @@ -223,8 +237,8 @@ class CompressedTensorsW4A4Nvfp4MoEMethod(CompressedTensorsMoEMethod): replace_parameter(layer, "w13_weight_scale", w13_scale) replace_parameter(layer, "w2_weight", w2) replace_parameter(layer, "w2_weight_scale", w2_scale) - layer.w13_weight_scale_2 = w13_scale_2 - layer.w2_weight_scale_2 = w2_scale_2 + replace_parameter(layer, "w13_weight_scale_2", w13_scale_2) + replace_parameter(layer, "w2_weight_scale_2", w2_scale_2) layer.w13_input_scale = a13_scale layer.w2_input_scale = a2_scale diff --git a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py index d4f71bd76f..8cdc928b5e 100644 --- a/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py +++ b/vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w8a8_fp8.py @@ -290,7 +290,9 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod): assert self.input_quant.strategy == QuantizationStrategy.TENSOR assert w13_input_scale is not None and w2_input_scale is not None w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe( - w13_input_scale, w2_input_scale + w13_input_scale, + w2_input_scale, + layer.moe_config.moe_parallel_config.enable_eplb, ) replace_parameter(layer, "w13_input_scale", w13_input_scale) replace_parameter(layer, "w2_input_scale", w2_input_scale) diff --git a/vllm/model_executor/layers/quantization/fp8.py b/vllm/model_executor/layers/quantization/fp8.py index b377a16eba..ea77d8c0e8 100644 --- a/vllm/model_executor/layers/quantization/fp8.py +++ b/vllm/model_executor/layers/quantization/fp8.py @@ -709,7 +709,9 @@ class Fp8MoEMethod(FusedMoEMethodBase): assert not self.block_quant assert w13_input_scale is not None and w2_input_scale is not None w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe( - w13_input_scale, w2_input_scale + w13_input_scale, + w2_input_scale, + layer.moe_config.moe_parallel_config.enable_eplb, ) replace_parameter(layer, "w13_input_scale", w13_input_scale) replace_parameter(layer, "w2_input_scale", w2_input_scale) diff --git a/vllm/model_executor/layers/quantization/modelopt.py b/vllm/model_executor/layers/quantization/modelopt.py index 8e22a6b73e..7bf77cf069 100644 --- a/vllm/model_executor/layers/quantization/modelopt.py +++ b/vllm/model_executor/layers/quantization/modelopt.py @@ -899,7 +899,9 @@ class ModelOptFp8MoEMethod(FusedMoEMethodBase): # Per tensor kernels require single activation scale. Use the max. w13_input_scale, w2_input_scale = process_fp8_input_tensor_strategy_moe( - w13_input_scale, w2_input_scale + w13_input_scale, + w2_input_scale, + layer.moe_config.moe_parallel_config.enable_eplb, ) replace_parameter(layer, "w13_input_scale", w13_input_scale) replace_parameter(layer, "w2_input_scale", w2_input_scale) diff --git a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py index 8fe1a22371..e0a8c1a310 100644 --- a/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py +++ b/vllm/model_executor/layers/quantization/utils/flashinfer_fp4_moe.py @@ -15,6 +15,9 @@ from vllm.model_executor.layers.quantization.utils.flashinfer_utils import ( from vllm.model_executor.layers.quantization.utils.nvfp4_utils import ( swizzle_blockscale, ) +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + amax_for_moe_activation_quant, +) if TYPE_CHECKING: from vllm.model_executor.layers.fused_moe import RoutedExperts @@ -129,12 +132,13 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( and interleaves w13 gate/linear rows for gated activations. Non-gated activations use a single w13 projection and keep its row order unchanged. """ - from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout - # Global scaling factors (same as other FlashInfer backends). num_experts = w13.shape[0] - a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts) - a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts) + enable_eplb = layer.moe_config.moe_parallel_config.enable_eplb + a13_scale = amax_for_moe_activation_quant(a13_scale, enable_eplb).repeat( + num_experts + ) + a2_scale = amax_for_moe_activation_quant(a2_scale, enable_eplb).repeat(num_experts) if layer.activation.is_gated: w13, w13_scale = reorder_w13_to_w31_for_flashinfer_cutedsl( @@ -145,29 +149,8 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( w13 = interleave_linear_and_gate(w13, group_size=64, dim=1) w13_scale = interleave_linear_and_gate(w13_scale, group_size=64, dim=1) - # Convert w13 scale factors: linear → swizzled → MMA layout. w13_scale = swizzle_blockscale(w13_scale) - E, M_padded, K_sf_padded = w13_scale.shape - w13_scale_flat = w13_scale.reshape(E * M_padded, K_sf_padded) - w13_scale = convert_sf_to_mma_layout( - w13_scale_flat, - m=M_padded, - k=K_sf_padded * 16, - num_groups=E, - sf_vec_size=16, - ) - - # Convert w2 scale factors: linear → swizzled → MMA layout. w2_scale = swizzle_blockscale(w2_scale) - E, M_padded, K_sf_padded = w2_scale.shape - w2_scale_flat = w2_scale.reshape(E * M_padded, K_sf_padded) - w2_scale = convert_sf_to_mma_layout( - w2_scale_flat, - m=M_padded, - k=K_sf_padded * 16, - num_groups=E, - sf_vec_size=16, - ) return ( w13, @@ -181,6 +164,31 @@ def prepare_nvfp4_moe_layer_for_flashinfer_cutedsl( ) +def nvfp4_swizzled_scale_to_cutedsl_mma_view(scale: torch.Tensor) -> torch.Tensor: + """View a swizzled (E, M_padded, K_sf_padded) block-scale tensor in the + MMA layout expected by the CuteDSL MoE kernel. + + The returned tensor aliases `scale`'s storage, so in-place updates of the + registered Parameter (weight reloads, EPLB rearrangement) are visible to + the kernel with no extra bookkeeping. + """ + from flashinfer.cute_dsl.utils import convert_sf_to_mma_layout + + num_experts, m_padded, k_sf_padded = scale.shape + mma_view = convert_sf_to_mma_layout( + scale.reshape(num_experts * m_padded, k_sf_padded), + m=m_padded, + k=k_sf_padded * 16, + num_groups=num_experts, + sf_vec_size=16, + ) + assert mma_view.data_ptr() == scale.data_ptr(), ( + "convert_sf_to_mma_layout no longer returns a view of its input; " + "the quant config would go stale after weight updates." + ) + return mma_view + + def prepare_static_weights_for_trtllm_fp4_moe( # args_dequant, # args, @@ -363,8 +371,13 @@ def prepare_nvfp4_moe_layer_for_fi_or_cutlass( # For some FI kernels, the input scales are shared by all experts. if is_global_sf_supported_for_nvfp4_backend(backend): num_experts = w13.shape[0] - a13_scale = a13_scale.max().to(torch.float32).repeat(num_experts) - a2_scale = a2_scale.max().to(torch.float32).repeat(num_experts) + enable_eplb = layer.moe_config.moe_parallel_config.enable_eplb + a13_scale = amax_for_moe_activation_quant(a13_scale, enable_eplb).repeat( + num_experts + ) + a2_scale = amax_for_moe_activation_quant(a2_scale, enable_eplb).repeat( + num_experts + ) else: a13_scale = a13_scale.max(dim=1).values.to(torch.float32) diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 1fce368247..7910db1925 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -14,6 +14,7 @@ import vllm.envs as envs from vllm import _custom_ops as ops from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.quant_utils import ( + amax_for_moe_activation_quant, get_fp8_min_max, ) from vllm.model_executor.layers.quantization.utils.w8a8_utils import ( @@ -1437,6 +1438,7 @@ def process_fp8_weight_tensor_strategy_moe( def process_fp8_input_tensor_strategy_moe( w13_input_scale: torch.Tensor, w2_input_scale: torch.Tensor, + enable_eplb: bool, ) -> tuple[torch.Tensor, torch.Tensor]: """Process moe input scales for tensor-wise quantization strategy.""" @@ -1447,4 +1449,7 @@ def process_fp8_input_tensor_strategy_moe( "for each layer." ) - return w13_input_scale.max(), w2_input_scale.max() + return ( + amax_for_moe_activation_quant(w13_input_scale, enable_eplb), + amax_for_moe_activation_quant(w2_input_scale, enable_eplb), + ) diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index a63db87fe6..b127f2dd61 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -62,6 +62,28 @@ def amax_for_moe_weight_quant(amax: torch.Tensor, moe_tp_size: int) -> torch.Ten return amax +def amax_for_moe_activation_quant( + a_scale: torch.Tensor, enable_eplb: bool +) -> torch.Tensor: + """Reduce a per-expert activation scale to one value shared by all experts. + + Note: when EPLB is enabled and since this quantization scales get + folded into the per-expert dequantization alphas, we can only + ensure that the quant/dequant scales match by having a single + quantization scale shared across all ranks. + """ + a_max = a_scale.max().to(torch.float32) + if enable_eplb: + from vllm.distributed.parallel_state import get_ep_group + + torch.distributed.all_reduce( + a_max, + op=torch.distributed.ReduceOp.MAX, + group=get_ep_group().device_group, + ) + return a_max + + def get_fp8_min_max() -> tuple[float, float]: """Get the min and max values for FP8 quantization.""" # Using the default value (240.0) from pytorch will cause accuracy diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 8f365359ee..46d0ed30df 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -5,7 +5,6 @@ import asyncio from collections.abc import ( AsyncGenerator, Callable, - Iterable, Mapping, MutableSequence, Sequence, @@ -949,7 +948,7 @@ class MixtureOfExperts(Protocol): num_redundant_experts: int """Number of redundant experts in this model.""" - moe_layers: Iterable["MoERunner"] + moe_layers: Sequence["MoERunner"] """List of MoE layers in this model.""" def set_eplb_state( @@ -992,6 +991,37 @@ class MixtureOfExperts(Protocol): ) -> None: ... +def get_mixture_of_experts_model(model: object) -> MixtureOfExperts | None: + """Return the MixtureOfExperts contained within an arbitrary model. + + - If the model itself is a MixtureOfExperts, return the model directly. + - If the model is a multi-modal model, and its `language_model` is a + MixtureOfExperts, return the `language_model`. + - If neither, return None. + + Args: + model: Model being served. + + Returns: + The MixtureOfExperts instance contained within the model, or None. + """ + + if is_mixture_of_experts(model): + return model + + if isinstance(model, SupportsMultiModal): + try: + mm_language_model = model.get_language_model() + return ( + mm_language_model if is_mixture_of_experts(mm_language_model) else None + ) + except NotImplementedError: + logger.info_once("Cannot fetch language_model from MultiModal model") + return None + + return None + + def is_mixture_of_experts(model: object) -> TypeIs[MixtureOfExperts]: return ( isinstance(model, MixtureOfExperts) and getattr(model, "num_moe_layers", 0) > 0 diff --git a/vllm/v1/worker/gpu/eplb_utils.py b/vllm/v1/worker/gpu/eplb_utils.py index aea6fdeff8..57a4e6db2c 100644 --- a/vllm/v1/worker/gpu/eplb_utils.py +++ b/vllm/v1/worker/gpu/eplb_utils.py @@ -12,23 +12,12 @@ from vllm.config import ModelConfig from vllm.distributed.eplb.eplb_state import EplbState from vllm.logger import init_logger from vllm.model_executor.models.interfaces import ( - SupportsMultiModal, - is_mixture_of_experts, + get_mixture_of_experts_model, ) logger = init_logger(__name__) -def _unwrap_moe(model: nn.Module) -> nn.Module: - # VLM wrappers (e.g. KimiK25ForConditionalGeneration) hold the MoE - # language model under `.language_model` but don't implement - # MixtureOfExperts themselves. Mirror the V1 path - # (see vllm/v1/worker/gpu_model_runner.py, PR #39805). - if not is_mixture_of_experts(model) and isinstance(model, SupportsMultiModal): - return model.get_language_model() - return model - - def step_eplb_after(*, is_dummy: bool = False) -> Callable: """Step EPLB after a model runner method completes successfully.""" @@ -78,7 +67,8 @@ class EPLBController: return False draft_model = speculator.model - if not is_mixture_of_experts(draft_model): + draft_moe_model = get_mixture_of_experts_model(draft_model) + if draft_moe_model is None: return False assert not self.parallel_config.enable_elastic_ep, ( @@ -88,7 +78,7 @@ class EPLBController: assert speculative_config.draft_model_config is not None assert self.state is not None self.state.add_model( - draft_model, + draft_moe_model, speculative_config.draft_model_config, ) speculator.set_eplb_state(self.state) @@ -104,13 +94,15 @@ class EPLBController: if not self.parallel_config.enable_eplb or load_dummy_weights: return False - model = _unwrap_moe(model) - if not is_mixture_of_experts(model): + moe_model = get_mixture_of_experts_model(model) + if moe_model is None: return False - logger.info_once("EPLB is enabled for model %s.", model_config.model) + logger.info_once( + "EPLB is enabled for MoE part of model %s.", model_config.model + ) assert self.state is not None - self.state.add_model(model, model_config) + self.state.add_model(moe_model, model_config) self._has_registered_models = True return True @@ -154,11 +146,11 @@ class EPLBController: expanded_physical_to_logical: torch.Tensor, old_num_physical_experts: int, ) -> None: - model = _unwrap_moe(model) - assert is_mixture_of_experts(model) + moe_model = get_mixture_of_experts_model(model) + assert moe_model is not None self.state = EplbState.from_mapping( - model=model, + model=moe_model, model_config=model_config, device=self.device, parallel_config=self.parallel_config, diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index c42d02cb22..b56e3cc5c1 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -82,7 +82,7 @@ from vllm.model_executor.models.interfaces import ( SupportsMRoPE, SupportsMultiModal, SupportsXDRoPE, - is_mixture_of_experts, + get_mixture_of_experts_model, supports_eagle3, supports_mrope, supports_multimodal_pruning, @@ -5364,9 +5364,14 @@ class GPUModelRunner( if hasattr(self.drafter, "load_model"): self.drafter.load_model(self.model) if ( - hasattr(self.drafter, "model") - and is_mixture_of_experts(self.drafter.model) - and self.parallel_config.enable_eplb + self.parallel_config.enable_eplb + and hasattr(self.drafter, "model") + and ( + drafter_moe_model := get_mixture_of_experts_model( + self.drafter.model + ) + ) + is not None ): assert not self.parallel_config.enable_elastic_ep, ( "Elastic EP is not supported with drafter model." @@ -5375,7 +5380,7 @@ class GPUModelRunner( assert spec_config is not None assert spec_config.draft_model_config is not None logger.info_once( - "EPLB is enabled for drafter model %s.", + "EPLB is enabled for MoE part of drafter model %s.", spec_config.draft_model_config.model, ) if self.eplb_state is None: @@ -5383,7 +5388,7 @@ class GPUModelRunner( self.parallel_config, self.device ) self.eplb_state.add_model( - self.drafter.model, + drafter_moe_model, spec_config.draft_model_config, ) assert hasattr(self.drafter, "set_eplb_state") @@ -5396,21 +5401,15 @@ class GPUModelRunner( # VLM models (e.g. KimiK25ForConditionalGeneration) wrap the # actual MoE language model but don't implement # MixtureOfExperts themselves. - moe_candidate = self.model - if not is_mixture_of_experts(moe_candidate) and isinstance( - moe_candidate, SupportsMultiModal - ): - moe_candidate = moe_candidate.get_language_model() - if is_mixture_of_experts(moe_candidate): - self._moe_model = moe_candidate + self._moe_model = get_mixture_of_experts_model(self.model) if ( - self._moe_model is not None - and self.parallel_config.enable_eplb + self.parallel_config.enable_eplb and not load_dummy_weights + and self._moe_model is not None ): logger.info_once( - "EPLB is enabled for model %s.", + "EPLB is enabled for MoE part of model %s.", self.model_config.model, ) assert self.eplb_state is not None @@ -5450,8 +5449,8 @@ class GPUModelRunner( ) # Temporary hack for dynamic res video w/o support for bs>1 yet if ( - self._moe_model is not None - and self.parallel_config.enable_eplb + self.parallel_config.enable_eplb + and self._moe_model is not None and not load_dummy_weights and self.eplb_state is not None and self.eplb_state.is_async