| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | """EPLB rearrangement consistency tests for quant-method derived state. |
| 5 | |
| 6 | EPLB rearranges every registered Parameter of a RoutedExperts in place, |
| 7 | sliced along dim 0 (see RoutedExperts.get_expert_weights). Quant methods |
| 8 | must therefore register all derived per-expert tensors as Parameters and |
| 9 | alias the same storage in their FusedMoEQuantConfig, so the kernels observe |
| 10 | rearranged values with no extra bookkeeping. These tests verify that |
| 11 | contract: simulating a rearrangement on the registered Parameters must be |
| 12 | indistinguishable from loading a checkpoint with permuted experts. |
| 13 | """ |
| 14 | |
| 15 | import pytest |
| 16 | import torch |
| 17 | |
| 18 | from vllm.config import VllmConfig, set_current_vllm_config |
| 19 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 20 | from vllm.model_executor.layers.fused_moe.config import ( |
| 21 | FusedMoEConfig, |
| 22 | FusedMoEParallelConfig, |
| 23 | RoutingMethodType, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w4a4_nvfp4 import ( # noqa: E501 |
| 26 | CompressedTensorsW4A4Nvfp4MoEMethod, |
| 27 | ) |
| 28 | from vllm.platforms import current_platform |
| 29 | |
| 30 | EPLB_NVFP4_BACKENDS = ["flashinfer_cutedsl", "flashinfer_trtllm"] |
| 31 | |
| 32 | NUM_EXPERTS = 8 |
| 33 | HIDDEN_SIZE = 128 |
| 34 | INTERMEDIATE_SIZE = 256 |
| 35 | EXPERT_PERMUTATION = [3, 0, 5, 1, 7, 2, 6, 4] |
| 36 | |
| 37 | QUANT_CONFIG_TENSORS = ( |
| 38 | "w1_scale", |
| 39 | "w2_scale", |
| 40 | "g1_alphas", |
| 41 | "g2_alphas", |
| 42 | "a1_gscale", |
| 43 | "a2_gscale", |
| 44 | ) |
| 45 | |
| 46 | |
| 47 | class _RoutedExpertsStub(torch.nn.Module): |
| 48 | """Minimal ``RoutedExperts`` stand-in for NVFP4 weight processing.""" |
| 49 | |
| 50 | def __init__(self, moe_config: FusedMoEConfig) -> None: |
| 51 | super().__init__() |
| 52 | self.moe_config = moe_config |
| 53 | self.activation = moe_config.activation |
| 54 | |
| 55 | def _expert_routing_tables( |
| 56 | self, |
| 57 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor] | None: |
| 58 | return None |
| 59 | |
| 60 | |
| 61 | def _make_moe_config(backend: str) -> FusedMoEConfig: |
| 62 | parallel_config = FusedMoEParallelConfig( |
| 63 | tp_size=1, |
| 64 | pcp_size=1, |
| 65 | dp_size=1, |
| 66 | ep_size=1, |
| 67 | tp_rank=0, |
| 68 | pcp_rank=0, |
| 69 | dp_rank=0, |
| 70 | ep_rank=0, |
| 71 | sp_size=1, |
| 72 | use_ep=True, |
| 73 | all2all_backend="allgather_reducescatter", |
| 74 | enable_eplb=True, |
| 75 | ) |
| 76 | return FusedMoEConfig( |
| 77 | num_experts=NUM_EXPERTS, |
| 78 | experts_per_token=2, |
| 79 | hidden_dim=HIDDEN_SIZE, |
| 80 | intermediate_size=INTERMEDIATE_SIZE, |
| 81 | num_local_experts=NUM_EXPERTS, |
| 82 | num_logical_experts=NUM_EXPERTS, |
| 83 | activation=MoEActivation.SILU, |
| 84 | device="cuda", |
| 85 | routing_method=RoutingMethodType.TopK, |
| 86 | moe_parallel_config=parallel_config, |
| 87 | in_dtype=torch.bfloat16, |
| 88 | moe_backend=backend, |
| 89 | ) |
| 90 | |
| 91 | |
| 92 | def _make_raw_weights( |
| 93 | device: torch.device, generator: torch.Generator |
| 94 | ) -> dict[str, torch.Tensor]: |
| 95 | """Random raw NVFP4 checkpoint tensors, indexable per-expert on dim 0.""" |
| 96 | e, h, i = NUM_EXPERTS, HIDDEN_SIZE, INTERMEDIATE_SIZE |
| 97 | |
| 98 | def packed(*shape: int) -> torch.Tensor: |
| 99 | return torch.randint( |
| 100 | 0, 256, shape, dtype=torch.uint8, device=device, generator=generator |
| 101 | ) |
| 102 | |
| 103 | def block_scale(*shape: int) -> torch.Tensor: |
| 104 | return torch.randint( |
| 105 | 0, 128, shape, dtype=torch.uint8, device=device, generator=generator |
| 106 | ).view(torch.float8_e4m3fn) |
| 107 | |
| 108 | def global_scale(*shape: int) -> torch.Tensor: |
| 109 | return torch.rand(shape, device=device, generator=generator) + 0.5 |
| 110 | |
| 111 | return { |
| 112 | "w13_weight_packed": packed(e, 2 * i, h // 2), |
| 113 | "w2_weight_packed": packed(e, h, i // 2), |
| 114 | "w13_weight_scale": block_scale(e, 2 * i, h // 16), |
| 115 | "w2_weight_scale": block_scale(e, h, i // 16), |
| 116 | "w13_weight_global_scale": global_scale(e, 2), |
| 117 | "w2_weight_global_scale": global_scale(e), |
| 118 | "w13_input_global_scale": global_scale(e, 2), |
| 119 | "w2_input_global_scale": global_scale(e), |
| 120 | } |
| 121 | |
| 122 | |
| 123 | def _build_processed_layer( |
| 124 | backend: str, raw: dict[str, torch.Tensor], device: torch.device |
| 125 | ) -> tuple[CompressedTensorsW4A4Nvfp4MoEMethod, _RoutedExpertsStub]: |
| 126 | """Create the method + layer, load ``raw`` and run the real |
| 127 | process_weights_after_loading (which builds the flashinfer kernel).""" |
| 128 | moe_config = _make_moe_config(backend) |
| 129 | method = CompressedTensorsW4A4Nvfp4MoEMethod(moe_config, "layer.0", use_a16=False) |
| 130 | layer = _RoutedExpertsStub(moe_config).to(device) |
| 131 | method.create_weights( |
| 132 | layer, |
| 133 | num_experts=NUM_EXPERTS, |
| 134 | hidden_size=HIDDEN_SIZE, |
| 135 | intermediate_size_per_partition=INTERMEDIATE_SIZE, |
| 136 | params_dtype=torch.bfloat16, |
| 137 | weight_loader=lambda *args, **kwargs: None, |
| 138 | ) |
| 139 | layer = layer.to(device) |
| 140 | with torch.no_grad(): |
| 141 | for name, value in raw.items(): |
| 142 | getattr(layer, name).copy_(value) |
| 143 | method.process_weights_after_loading(layer) |
| 144 | return method, layer |
| 145 | |
| 146 | |
| 147 | def _assert_tensors_equal( |
| 148 | actual: torch.Tensor, expected: torch.Tensor, name: str |
| 149 | ) -> None: |
| 150 | if actual.dtype == torch.float8_e4m3fn: |
| 151 | assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)), name |
| 152 | elif actual.dtype == torch.uint8: |
| 153 | assert torch.equal(actual, expected), name |
| 154 | else: |
| 155 | torch.testing.assert_close(actual, expected, msg=lambda m: f"{name}: {m}") |
| 156 | |
| 157 | |
| 158 | def _simulate_eplb_rearrangement(layer: torch.nn.Module, perm: torch.Tensor) -> None: |
| 159 | """Permute experts of every registered Parameter in place, the way |
| 160 | EPLB's in-place rearrangement moves expert slices along dim 0.""" |
| 161 | with torch.no_grad(): |
| 162 | for _, param in layer.named_parameters(): |
| 163 | param.copy_(param[perm]) |
| 164 | |
| 165 | |
| 166 | def _ensure_world1_distributed() -> None: |
| 167 | """The enable_eplb config makes weight processing all-reduce activation |
| 168 | scale amaxes over the EP group, so a (single-rank) EP group must exist.""" |
| 169 | from vllm.distributed.parallel_state import ( |
| 170 | ensure_model_parallel_initialized, |
| 171 | init_distributed_environment, |
| 172 | ) |
| 173 | |
| 174 | if not torch.distributed.is_initialized(): |
| 175 | init_distributed_environment( |
| 176 | world_size=1, |
| 177 | rank=0, |
| 178 | distributed_init_method="tcp://127.0.0.1:0", |
| 179 | local_rank=0, |
| 180 | ) |
| 181 | ensure_model_parallel_initialized(1, 1) |
| 182 | |
| 183 | |
| 184 | @pytest.mark.parametrize("backend", EPLB_NVFP4_BACKENDS) |
| 185 | def test_nvfp4_eplb_rearrangement_matches_reload(backend: str) -> None: |
| 186 | if not ( |
| 187 | current_platform.is_cuda() and current_platform.is_device_capability_family(100) |
| 188 | ): |
| 189 | pytest.skip("NVFP4 CuteDSL/TRTLLM MoE backends require Blackwell (SM100).") |
| 190 | |
| 191 | device = torch.device("cuda:0") |
| 192 | torch.accelerator.set_device_index(device) |
| 193 | perm = torch.tensor(EXPERT_PERMUTATION, device=device) |
| 194 | |
| 195 | with set_current_vllm_config(VllmConfig()): |
| 196 | _ensure_world1_distributed() |
| 197 | generator = torch.Generator(device=device).manual_seed(1234) |
| 198 | raw = _make_raw_weights(device, generator) |
| 199 | raw_permuted = {name: value[perm].contiguous() for name, value in raw.items()} |
| 200 | |
| 201 | try: |
| 202 | method, layer = _build_processed_layer(backend, raw, device) |
| 203 | except ValueError as exc: |
| 204 | pytest.skip(f"{backend} NVFP4 MoE backend unavailable: {exc}") |
| 205 | ref_method, ref_layer = _build_processed_layer(backend, raw_permuted, device) |
| 206 | |
| 207 | # The EPLB contract: every registered Parameter must be an |
| 208 | # expert-major contiguous tensor so get_expert_weights can view it |
| 209 | # as (E, -1) and rearrange expert slices in place. |
| 210 | for name, param in layer.named_parameters(): |
| 211 | assert param.is_contiguous(), f"{name} is not contiguous" |
| 212 | assert param.shape[0] == NUM_EXPERTS, ( |
| 213 | f"{name} is not expert-major: {tuple(param.shape)}" |
| 214 | ) |
| 215 | |
| 216 | # Derived per-expert scales must live in registered Parameters that |
| 217 | # the quant config aliases; anything else goes stale on rearrangement. |
| 218 | quant_config = method.moe_quant_config |
| 219 | assert quant_config is not None |
| 220 | params = dict(layer.named_parameters()) |
| 221 | assert quant_config.g1_alphas.data_ptr() == ( |
| 222 | params["w13_weight_scale_2"].data_ptr() |
| 223 | ) |
| 224 | assert quant_config.g2_alphas.data_ptr() == ( |
| 225 | params["w2_weight_scale_2"].data_ptr() |
| 226 | ) |
| 227 | assert quant_config.w1_scale.data_ptr() == ( |
| 228 | params["w13_weight_scale"].data_ptr() |
| 229 | ) |
| 230 | assert quant_config.w2_scale.data_ptr() == params["w2_weight_scale"].data_ptr() |
| 231 | |
| 232 | _simulate_eplb_rearrangement(layer, perm) |
| 233 | |
| 234 | # After rearrangement, all registered state must equal what a fresh |
| 235 | # load of the permuted experts produces. |
| 236 | ref_params = dict(ref_layer.named_parameters()) |
| 237 | assert params.keys() == ref_params.keys() |
| 238 | for name in params: |
| 239 | _assert_tensors_equal(params[name], ref_params[name], name) |
| 240 | |
| 241 | # And so must the kernel-visible quant config tensors (the CuteDSL |
| 242 | # MMA scale views alias the registered Parameters' storage). |
| 243 | ref_quant_config = ref_method.moe_quant_config |
| 244 | assert ref_quant_config is not None |
| 245 | for name in QUANT_CONFIG_TENSORS: |
| 246 | actual = getattr(quant_config, name) |
| 247 | expected = getattr(ref_quant_config, name) |
| 248 | assert (actual is None) == (expected is None), name |
| 249 | if actual is not None: |
| 250 | _assert_tensors_equal(actual, expected, f"quant_config.{name}") |
| 251 | |
| 252 | |
| 253 | def test_fp8_per_tensor_alphas_registered_and_aliased() -> None: |
| 254 | """CPU-only: the fp8 per-tensor oracle must register the fused |
| 255 | (w_scale * a_scale) products as layer Parameters aliased by the quant |
| 256 | config, so EPLB rearrangement keeps kernels consistent.""" |
| 257 | from vllm.model_executor.layers.fused_moe.oracle.fp8 import ( |
| 258 | Fp8MoeBackend, |
| 259 | make_fp8_moe_quant_config, |
| 260 | ) |
| 261 | |
| 262 | e = NUM_EXPERTS |
| 263 | generator = torch.Generator().manual_seed(1234) |
| 264 | |
| 265 | def scales() -> tuple[torch.Tensor, torch.Tensor]: |
| 266 | w_scale = torch.rand((e,), generator=generator) + 0.5 |
| 267 | a_scale = torch.rand((), generator=generator) + 0.5 |
| 268 | return w_scale, a_scale.expand(e).contiguous() |
| 269 | |
| 270 | w1_scale, a1_scale = scales() |
| 271 | w2_scale, a2_scale = scales() |
| 272 | |
| 273 | layer = torch.nn.Module() |
| 274 | quant_config = make_fp8_moe_quant_config( |
| 275 | fp8_backend=Fp8MoeBackend.FLASHINFER_CUTLASS, |
| 276 | w1_scale=w1_scale, |
| 277 | w2_scale=w2_scale, |
| 278 | a1_scale=a1_scale, |
| 279 | a2_scale=a2_scale, |
| 280 | layer=layer, |
| 281 | ) |
| 282 | |
| 283 | params = dict(layer.named_parameters()) |
| 284 | assert quant_config.g1_alphas.data_ptr() == params["g1_alphas"].data_ptr() |
| 285 | assert quant_config.g2_alphas.data_ptr() == params["g2_alphas"].data_ptr() |
| 286 | |
| 287 | perm = torch.tensor(EXPERT_PERMUTATION) |
| 288 | ref_config = make_fp8_moe_quant_config( |
| 289 | fp8_backend=Fp8MoeBackend.FLASHINFER_CUTLASS, |
| 290 | w1_scale=w1_scale[perm].contiguous(), |
| 291 | w2_scale=w2_scale[perm].contiguous(), |
| 292 | a1_scale=a1_scale, |
| 293 | a2_scale=a2_scale, |
| 294 | layer=None, |
| 295 | ) |
| 296 | _simulate_eplb_rearrangement(layer, perm) |
| 297 | |
| 298 | torch.testing.assert_close(quant_config.g1_alphas, ref_config.g1_alphas) |
| 299 | torch.testing.assert_close(quant_config.g2_alphas, ref_config.g2_alphas) |