| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from types import SimpleNamespace |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | |
| 8 | import vllm.distributed as distributed |
| 9 | from tests.kernels.moe.utils import make_dummy_moe_config |
| 10 | from vllm.model_executor.layers.fused_moe import ( |
| 11 | unquantized_fused_moe_method as unquantized_module, |
| 12 | ) |
| 13 | from vllm.model_executor.layers.fused_moe.config import RoutingMethodType |
| 14 | from vllm.model_executor.layers.fused_moe.routed_experts import RoutedExperts |
| 15 | from vllm.model_executor.layers.quantization import fp8 as fp8_module |
| 16 | from vllm.model_executor.layers.quantization.fp8 import Fp8Config |
| 17 | |
| 18 | |
| 19 | def _make_fp8_tp_experts( |
| 20 | monkeypatch, |
| 21 | tp_size, |
| 22 | tp_rank, |
| 23 | backend="flashinfer_trtllm", |
| 24 | *, |
| 25 | mock_backend=True, |
| 26 | num_experts=2, |
| 27 | hidden_dim=256, |
| 28 | intermediate_size=640, |
| 29 | quant_config=None, |
| 30 | config_overrides=None, |
| 31 | ): |
| 32 | # Exercise allocation and the public weight loader without selecting a GPU kernel. |
| 33 | if mock_backend: |
| 34 | monkeypatch.setattr( |
| 35 | fp8_module, "select_fp8_moe_backend", lambda **kwargs: (None, None) |
| 36 | ) |
| 37 | monkeypatch.setattr( |
| 38 | unquantized_module, |
| 39 | "select_unquantized_moe_backend", |
| 40 | lambda **kwargs: (None, None), |
| 41 | ) |
| 42 | monkeypatch.setattr( |
| 43 | fp8_module, "get_tensor_model_parallel_world_size", lambda: tp_size |
| 44 | ) |
| 45 | # The block-shape validator imports this function from vllm.distributed. |
| 46 | monkeypatch.setattr( |
| 47 | distributed, "get_tensor_model_parallel_world_size", lambda: tp_size |
| 48 | ) |
| 49 | config = make_dummy_moe_config( |
| 50 | num_experts=num_experts, |
| 51 | hidden_dim=hidden_dim, |
| 52 | intermediate_size=intermediate_size, |
| 53 | experts_per_token=min(num_experts, 10), |
| 54 | ) |
| 55 | config.routing_method = RoutingMethodType.RenormalizeNaive |
| 56 | config.moe_parallel_config.tp_size = tp_size |
| 57 | config.moe_parallel_config.tp_rank = tp_rank |
| 58 | config.intermediate_size_per_partition = intermediate_size // tp_size |
| 59 | config.intermediate_size_per_partition_unpadded = intermediate_size // tp_size |
| 60 | config.moe_backend = backend |
| 61 | for field, value in (config_overrides or {}).items(): |
| 62 | setattr(config, field, value) |
| 63 | return RoutedExperts( |
| 64 | "model.layers.0.mlp.experts", |
| 65 | torch.bfloat16, |
| 66 | config, |
| 67 | quant_config |
| 68 | or Fp8Config(is_checkpoint_fp8_serialized=True, weight_block_size=[128, 128]), |
| 69 | expert_map_manager=SimpleNamespace( |
| 70 | local_num_experts=num_experts, |
| 71 | placement_strategy="linear", |
| 72 | expert_map=None, |
| 73 | expert_mask=None, |
| 74 | routing_tables=None, |
| 75 | map_global_to_local=lambda index: index, |
| 76 | ), |
| 77 | ) |
| 78 | |
| 79 | |
| 80 | @pytest.mark.parametrize("tp_size", [2, 4, 8]) |
| 81 | @pytest.mark.parametrize("batched", [False, True]) |
| 82 | @pytest.mark.parametrize("intermediate_size", [640, 896]) |
| 83 | def test_fp8_block_aligned_tp_preserves_checkpoint( |
| 84 | monkeypatch, tp_size, batched, intermediate_size |
| 85 | ): |
| 86 | """All ranks reconstruct the original dequantized projections, including |
| 87 | padding-only ranks. Reloading must clear stale weights and scales. |
| 88 | """ |
| 89 | generator = torch.Generator().manual_seed(42) |
| 90 | weights = { |
| 91 | "w1": torch.randn(2, intermediate_size, 256, generator=generator).to( |
| 92 | torch.float8_e4m3fn |
| 93 | ), |
| 94 | "w3": torch.randn(2, intermediate_size, 256, generator=generator).to( |
| 95 | torch.float8_e4m3fn |
| 96 | ), |
| 97 | "w2": torch.randn(2, 256, intermediate_size, generator=generator).to( |
| 98 | torch.float8_e4m3fn |
| 99 | ), |
| 100 | } |
| 101 | scales = { |
| 102 | name: torch.rand(2, w.shape[1] // 128, w.shape[2] // 128, generator=generator) |
| 103 | + 0.1 |
| 104 | for name, w in weights.items() |
| 105 | } |
| 106 | reconstructed: dict[str, list[torch.Tensor]] = {name: [] for name in weights} |
| 107 | num_blocks = intermediate_size // 128 |
| 108 | for rank in range(tp_size): |
| 109 | layer = _make_fp8_tp_experts( |
| 110 | monkeypatch, tp_size, rank, intermediate_size=intermediate_size |
| 111 | ) |
| 112 | assert layer.quant_method.weight_scale_refine is None |
| 113 | assert layer.quant_method.moe_block_shape == [128, 128] |
| 114 | width = layer.moe_config.intermediate_size_per_partition |
| 115 | assert width == ((num_blocks + tp_size - 1) // tp_size) * 128 |
| 116 | for name in weights: |
| 117 | prefix = "w2" if name == "w2" else "w13" |
| 118 | for suffix, checkpoint in ( |
| 119 | ("weight", weights[name]), |
| 120 | ("weight_scale_inv", scales[name]), |
| 121 | ): |
| 122 | param_name = f"{prefix}_{suffix}" |
| 123 | param = getattr(layer, param_name) |
| 124 | # Also reload after dirtying the destination, to catch padding leaks. |
| 125 | for _ in range(2): |
| 126 | if name != "w3": |
| 127 | param.data.fill_(7) |
| 128 | if batched: |
| 129 | assert param.weight_loader( |
| 130 | param, checkpoint, param_name, name, 0, return_success=True |
| 131 | ) |
| 132 | else: |
| 133 | for expert in range(2): |
| 134 | assert param.weight_loader( |
| 135 | param, |
| 136 | checkpoint[expert], |
| 137 | param_name, |
| 138 | name, |
| 139 | expert, |
| 140 | return_success=True, |
| 141 | ) |
| 142 | w = getattr(layer, f"{prefix}_weight").float() |
| 143 | s = getattr(layer, f"{prefix}_weight_scale_inv") |
| 144 | if name in ("w1", "w3"): |
| 145 | half = 0 if name == "w1" else 1 |
| 146 | w = w.chunk(2, dim=1)[half] |
| 147 | s = s.chunk(2, dim=1)[half] |
| 148 | dim = 2 if name == "w2" else 1 |
| 149 | valid = max(0, min(width, intermediate_size - rank * width)) |
| 150 | assert torch.count_nonzero(w.narrow(dim, valid, width - valid)) == 0 |
| 151 | assert torch.all(s.narrow(dim, valid // 128, (width - valid) // 128) == 1) |
| 152 | dequant = w * s.repeat_interleave(128, dim=1).repeat_interleave(128, dim=2) |
| 153 | reconstructed[name].append(dequant.narrow(dim, 0, valid)) |
| 154 | for name, parts in reconstructed.items(): |
| 155 | expected = weights[name].float() * scales[name].repeat_interleave( |
| 156 | 128, dim=1 |
| 157 | ).repeat_interleave(128, dim=2) |
| 158 | assert torch.equal(torch.cat(parts, dim=2 if name == "w2" else 1), expected) |
| 159 | |
| 160 | |
| 161 | @pytest.mark.parametrize("backend", ["auto", "triton"]) |
| 162 | @pytest.mark.parametrize("tp_size,block", [(2, 64), (4, 32)]) |
| 163 | def test_fp8_tp_default_keeps_refined_layout(monkeypatch, backend, tp_size, block): |
| 164 | layer = _make_fp8_tp_experts(monkeypatch, tp_size, 0, backend) |
| 165 | assert layer.moe_config.intermediate_size_per_partition == 640 // tp_size |
| 166 | assert layer.quant_method.moe_block_shape == [block, block] |
| 167 | |
| 168 | |
| 169 | @pytest.mark.parametrize( |
| 170 | "field,value", |
| 171 | [ |
| 172 | ("hidden_dim", 192), |
| 173 | ("intermediate_size", 672), |
| 174 | ("is_lora_enabled", True), |
| 175 | ("has_bias", True), |
| 176 | ], |
| 177 | ) |
| 178 | def test_fp8_block_aligned_tp_rejects_unsupported_layout(monkeypatch, field, value): |
| 179 | with pytest.raises(ValueError, match="Block-aligned FP8 TP sharding requires"): |
| 180 | _make_fp8_tp_experts(monkeypatch, 4, 0, config_overrides={field: value}) |
| 181 | |
| 182 | |
| 183 | @pytest.mark.parametrize("tp_size", [1, 5]) |
| 184 | def test_fp8_aligned_tp_keeps_original_layout(monkeypatch, tp_size): |
| 185 | layer = _make_fp8_tp_experts(monkeypatch, tp_size, 0) |
| 186 | assert layer.moe_config.intermediate_size_per_partition == 640 // tp_size |
| 187 | assert not layer.moe_config.tp_shard_with_padding |
| 188 | assert layer.quant_method.weight_scale_refine is None |
| 189 | |
| 190 | |
| 191 | def test_fp8_skipped_layer_keeps_original_tp_layout(monkeypatch, default_vllm_config): |
| 192 | layer = _make_fp8_tp_experts( |
| 193 | monkeypatch, |
| 194 | 4, |
| 195 | 0, |
| 196 | quant_config=Fp8Config( |
| 197 | is_checkpoint_fp8_serialized=True, |
| 198 | weight_block_size=[128, 128], |
| 199 | ignored_layers=["model.layers.0.mlp.experts"], |
| 200 | ), |
| 201 | ) |
| 202 | assert isinstance(layer.quant_method, unquantized_module.UnquantizedFusedMoEMethod) |
| 203 | assert not layer.moe_config.tp_shard_with_padding |
| 204 | assert layer.moe_config.intermediate_size_per_partition == 160 |
| 205 | assert layer.w13_weight.shape == (2, 320, 256) |
| 206 | assert layer.w2_weight.shape == (2, 256, 160) |
| 207 | |
| 208 | |
| 209 | def test_fp8_block_aligned_tp_rejects_presharded_weights(monkeypatch): |
| 210 | layer = _make_fp8_tp_experts(monkeypatch, 4, 0) |
| 211 | with pytest.raises(ValueError, match="unsharded checkpoint"): |
| 212 | layer.w13_weight.weight_loader( |
| 213 | layer.w13_weight, |
| 214 | torch.zeros(160, 256, dtype=torch.float8_e4m3fn), |
| 215 | "w13_weight", |
| 216 | "w1", |
| 217 | 0, |
| 218 | ) |
| 219 | |
| 220 | |
| 221 | @pytest.mark.parametrize("tp_size", [2, 4, 8]) |
| 222 | @pytest.mark.parametrize("num_tokens", [1, 17]) |
| 223 | @torch.inference_mode() |
| 224 | def test_fp8_block_aligned_tp_flashinfer_matches_unsharded( |
| 225 | monkeypatch, |
| 226 | tp_size, |
| 227 | num_tokens, |
| 228 | dist_init, |
| 229 | workspace_init, |
| 230 | ): |
| 231 | """Compare sharded FI with unsharded FI, reporting the cross-backend |
| 232 | difference separately because activation quantization can differ. |
| 233 | """ |
| 234 | from vllm.config import VllmConfig, set_current_vllm_config |
| 235 | from vllm.model_executor.layers.fused_moe.config import fp8_w8a8_moe_quant_config |
| 236 | from vllm.model_executor.layers.fused_moe.fused_moe import fused_experts |
| 237 | from vllm.platforms import current_platform |
| 238 | |
| 239 | if not current_platform.is_device_capability_family(100): |
| 240 | pytest.skip("Requires datacenter Blackwell and FlashInfer TRTLLM") |
| 241 | torch.manual_seed(42) |
| 242 | e, h, n, topk = 16, 2560, 640, 10 |
| 243 | with torch.device("cuda"), set_current_vllm_config(VllmConfig()): |
| 244 | x = torch.randn(num_tokens, h, dtype=torch.bfloat16) / 10 |
| 245 | logits = torch.randn(num_tokens, e, dtype=torch.float32) |
| 246 | probabilities = logits.softmax(dim=-1) |
| 247 | topk_weights, topk_ids = probabilities.topk(topk, dim=-1) |
| 248 | topk_weights /= topk_weights.sum(dim=-1, keepdim=True) |
| 249 | weights = { |
| 250 | "w1": torch.randn(e, n, h).to(torch.float8_e4m3fn), |
| 251 | "w3": torch.randn(e, n, h).to(torch.float8_e4m3fn), |
| 252 | "w2": torch.randn(e, h, n).to(torch.float8_e4m3fn), |
| 253 | } |
| 254 | scales = { |
| 255 | name: (torch.rand(e, w.shape[1] // 128, w.shape[2] // 128) + 0.5) / 32 |
| 256 | for name, w in weights.items() |
| 257 | } |
| 258 | ref = fused_experts( |
| 259 | x, |
| 260 | torch.cat([weights["w1"], weights["w3"]], dim=1), |
| 261 | weights["w2"], |
| 262 | topk_weights, |
| 263 | topk_ids, |
| 264 | quant_config=fp8_w8a8_moe_quant_config( |
| 265 | w1_scale=torch.cat([scales["w1"], scales["w3"]], dim=1), |
| 266 | w2_scale=scales["w2"], |
| 267 | block_shape=[128, 128], |
| 268 | ), |
| 269 | ) |
| 270 | total = torch.zeros_like(ref, dtype=torch.float32) |
| 271 | fi_reference = None |
| 272 | for rank in range(-1, tp_size): |
| 273 | layer = _make_fp8_tp_experts( |
| 274 | monkeypatch, |
| 275 | 1 if rank == -1 else tp_size, |
| 276 | max(rank, 0), |
| 277 | mock_backend=False, |
| 278 | num_experts=e, |
| 279 | hidden_dim=h, |
| 280 | ) |
| 281 | for name in weights: |
| 282 | prefix = "w2" if name == "w2" else "w13" |
| 283 | for suffix, checkpoint in ( |
| 284 | ("weight", weights[name]), |
| 285 | ("weight_scale_inv", scales[name]), |
| 286 | ): |
| 287 | param_name = f"{prefix}_{suffix}" |
| 288 | param = getattr(layer, param_name) |
| 289 | param.weight_loader(param, checkpoint, param_name, name, 0) |
| 290 | method = layer.quant_method |
| 291 | method.process_weights_after_loading(layer) |
| 292 | assert method.is_monolithic |
| 293 | out = method.apply_monolithic(layer, x, logits) |
| 294 | assert torch.isfinite(out).all() |
| 295 | if rank == -1: |
| 296 | fi_reference = out.float().clone() |
| 297 | else: |
| 298 | total += out.float() |
| 299 | # Different FP8 GEMM reduction orders need not be bitwise identical. |
| 300 | assert fi_reference is not None |
| 301 | relative_l2 = (total - fi_reference).norm() / fi_reference.norm().clamp_min( |
| 302 | 1e-8 |
| 303 | ) |
| 304 | cross_backend_l2 = ( |
| 305 | fi_reference - ref.float() |
| 306 | ).norm() / ref.float().norm().clamp_min(1e-8) |
| 307 | print( |
| 308 | f"TP{tp_size} M={num_tokens}: shard_l2={relative_l2.item():.6f}, " |
| 309 | f"unsharded_fi_vs_triton_l2={cross_backend_l2.item():.6f}" |
| 310 | ) |
| 311 | assert relative_l2 < 0.01, relative_l2.item() |