| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from typing import cast |
| 4 | |
| 5 | import pytest |
| 6 | import torch |
| 7 | from nvfp4_utils import ( |
| 8 | FLOAT4_E2M1_MAX, |
| 9 | FLOAT8_E4M3_MAX, |
| 10 | break_fp4_bytes, |
| 11 | kE2M1ToFloat, |
| 12 | ) |
| 13 | from safetensors import safe_open |
| 14 | |
| 15 | import vllm.model_executor.layers.fused_moe.modular_kernel as mk |
| 16 | from vllm.config import VllmConfig, set_current_vllm_config |
| 17 | from vllm.config.kernel import KernelConfig |
| 18 | from vllm.model_executor.kernels.linear import init_nvfp4_linear_kernel |
| 19 | from vllm.model_executor.kernels.linear.nvfp4.emulation import ( |
| 20 | EmulationNvFp4LinearKernel, |
| 21 | ) |
| 22 | from vllm.model_executor.layers.fused_moe.activation import MoEActivation |
| 23 | from vllm.model_executor.layers.fused_moe.config import ( |
| 24 | FusedMoEConfig, |
| 25 | FusedMoEParallelConfig, |
| 26 | FusedMoEQuantConfig, |
| 27 | RoutingMethodType, |
| 28 | nvfp4_moe_quant_config, |
| 29 | ) |
| 30 | from vllm.model_executor.layers.fused_moe.experts.nvfp4_emulation_moe import ( |
| 31 | Nvfp4QuantizationEmulationTritonExperts, |
| 32 | ) |
| 33 | from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts |
| 34 | from vllm.model_executor.layers.quantization.utils import ( |
| 35 | nvfp4_emulation_utils, |
| 36 | ) |
| 37 | from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import ( |
| 38 | dequantize_to_dtype, |
| 39 | ref_nvfp4_quant_dequant, |
| 40 | ) |
| 41 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 42 | QuantKey, |
| 43 | kNvfp4Dynamic, |
| 44 | kNvfp4Static, |
| 45 | ) |
| 46 | from vllm.platforms import current_platform |
| 47 | from vllm.transformers_utils.repo_utils import hf_api |
| 48 | from vllm.triton_utils import triton |
| 49 | from vllm.utils.torch_utils import set_random_seed |
| 50 | |
| 51 | if current_platform.is_rocm(): |
| 52 | from vllm.platforms.rocm import on_gfx950 |
| 53 | else: |
| 54 | |
| 55 | def on_gfx950() -> bool: |
| 56 | return False |
| 57 | |
| 58 | |
| 59 | MOE_MODEL_CONFIGS = { |
| 60 | "nvidia/Qwen3-30B-A3B-NVFP4": { |
| 61 | "shards": ["model-00001-of-00004.safetensors"], |
| 62 | "expert_prefix": "model.layers.9.mlp.experts.", |
| 63 | # Position of the expert index in the dot-split key. |
| 64 | "expert_idx_pos": 5, |
| 65 | } |
| 66 | } |
| 67 | |
| 68 | |
| 69 | @pytest.fixture(scope="module") |
| 70 | def loaded_model_files(): |
| 71 | return { |
| 72 | model_id: hf_api().snapshot_download( |
| 73 | repo_id=model_id, allow_patterns=config["shards"] |
| 74 | ) |
| 75 | for model_id, config in MOE_MODEL_CONFIGS.items() |
| 76 | } |
| 77 | |
| 78 | |
| 79 | class Nvfp4QuantizationEmulationTritonExpertsReference(TritonExperts): |
| 80 | """Extension of TritonExperts to support emulated NVFP4 MoE experts. |
| 81 | |
| 82 | It may be used for NVFP4 models when the device does not have |
| 83 | native support for this dtype. |
| 84 | """ |
| 85 | |
| 86 | def __init__( |
| 87 | self, |
| 88 | moe_config: FusedMoEConfig, |
| 89 | quant_config: FusedMoEQuantConfig, |
| 90 | ): |
| 91 | super().__init__(moe_config, quant_config) |
| 92 | |
| 93 | # `TritonExperts.apply` expects pre-dequantized weights, |
| 94 | # which we handle in `apply` below. |
| 95 | self.w1_scale_val = self.quant_config.w1_scale |
| 96 | self.w2_scale_val = self.quant_config.w2_scale |
| 97 | |
| 98 | self.quant_config._w1.scale = None |
| 99 | self.quant_config._w2.scale = None |
| 100 | |
| 101 | self.quantization_emulation = True |
| 102 | |
| 103 | @property |
| 104 | def quant_dtype(self) -> torch.dtype | str | None: |
| 105 | return "nvfp4" |
| 106 | |
| 107 | @property |
| 108 | def a1_scale(self) -> torch.Tensor | None: |
| 109 | return self.quant_config.a1_gscale |
| 110 | |
| 111 | @property |
| 112 | def expects_unquantized_inputs(self) -> bool: |
| 113 | return True |
| 114 | |
| 115 | @staticmethod |
| 116 | def _supports_quant_scheme( |
| 117 | weight_key: QuantKey | None, |
| 118 | activation_key: QuantKey | None, |
| 119 | ) -> bool: |
| 120 | return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic) |
| 121 | |
| 122 | def apply( |
| 123 | self, |
| 124 | output: torch.Tensor, |
| 125 | hidden_states: torch.Tensor, |
| 126 | w1: torch.Tensor, |
| 127 | w2: torch.Tensor, |
| 128 | topk_weights: torch.Tensor, |
| 129 | topk_ids: torch.Tensor, |
| 130 | activation: MoEActivation, |
| 131 | global_num_experts: int, |
| 132 | expert_map: torch.Tensor | None, |
| 133 | a1q_scale: torch.Tensor | None, |
| 134 | a2_scale: torch.Tensor | None, |
| 135 | workspace13: torch.Tensor, |
| 136 | workspace2: torch.Tensor, |
| 137 | expert_tokens_meta: mk.ExpertTokensMetadata | None, |
| 138 | apply_router_weight_on_input: bool, |
| 139 | ): |
| 140 | assert w1.dtype == torch.uint8 |
| 141 | assert w2.dtype == torch.uint8 |
| 142 | |
| 143 | # Dequantize w1 from packed NVFP4 to fp16/bf16 |
| 144 | w13_global_scale = self.quant_config.g1_alphas |
| 145 | |
| 146 | w1_dequant = dequantize_to_dtype( |
| 147 | tensor_fp4=w1, |
| 148 | tensor_sf=self.w1_scale_val, |
| 149 | global_scale=w13_global_scale, |
| 150 | dtype=hidden_states.dtype, |
| 151 | block_size=16, |
| 152 | swizzle=False, |
| 153 | ) |
| 154 | |
| 155 | # Dequantize w2 from packed NVFP4 to fp16/bf16 |
| 156 | w2_global_scale = self.quant_config.g2_alphas |
| 157 | |
| 158 | w2_dequant = dequantize_to_dtype( |
| 159 | tensor_fp4=w2, |
| 160 | tensor_sf=self.w2_scale_val, |
| 161 | global_scale=w2_global_scale, |
| 162 | dtype=hidden_states.dtype, |
| 163 | block_size=16, |
| 164 | swizzle=False, |
| 165 | ) |
| 166 | |
| 167 | super().apply( |
| 168 | output=output, |
| 169 | hidden_states=hidden_states, |
| 170 | w1=w1_dequant, |
| 171 | w2=w2_dequant, |
| 172 | topk_weights=topk_weights, |
| 173 | topk_ids=topk_ids, |
| 174 | activation=activation, |
| 175 | global_num_experts=global_num_experts, |
| 176 | expert_map=expert_map, |
| 177 | a1q_scale=None, |
| 178 | a2_scale=self.quant_config.a2_gscale, |
| 179 | workspace13=workspace13, |
| 180 | workspace2=workspace2, |
| 181 | expert_tokens_meta=expert_tokens_meta, |
| 182 | apply_router_weight_on_input=apply_router_weight_on_input, |
| 183 | ) |
| 184 | |
| 185 | |
| 186 | @pytest.mark.parametrize( |
| 187 | ("config_kwargs", "expected_reason"), |
| 188 | [ |
| 189 | ({"has_bias": True}, "kernel does not support bias"), |
| 190 | ({"is_lora_enabled": True}, "kernel does not support LoRA"), |
| 191 | ], |
| 192 | ) |
| 193 | def test_nvfp4_emulation_support_check_rejects_bias_and_lora( |
| 194 | config_kwargs: dict[str, bool], |
| 195 | expected_reason: str, |
| 196 | ) -> None: |
| 197 | moe_config = FusedMoEConfig( |
| 198 | num_experts=2, |
| 199 | experts_per_token=1, |
| 200 | hidden_dim=16, |
| 201 | intermediate_size=16, |
| 202 | num_local_experts=2, |
| 203 | num_logical_experts=2, |
| 204 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 205 | activation=MoEActivation.SILU, |
| 206 | in_dtype=torch.bfloat16, |
| 207 | device="cuda", |
| 208 | routing_method=RoutingMethodType.TopK, |
| 209 | **config_kwargs, |
| 210 | ) |
| 211 | |
| 212 | supported, reason = Nvfp4QuantizationEmulationTritonExperts.is_supported_config( |
| 213 | Nvfp4QuantizationEmulationTritonExperts, |
| 214 | moe_config, |
| 215 | kNvfp4Static, |
| 216 | kNvfp4Dynamic, |
| 217 | mk.FusedMoEActivationFormat.Standard, |
| 218 | ) |
| 219 | |
| 220 | assert not supported |
| 221 | assert reason == expected_reason |
| 222 | |
| 223 | |
| 224 | @pytest.mark.skipif( |
| 225 | not current_platform.is_cuda_alike(), |
| 226 | reason="Triton NVFP4 kernel requires CUDA.", |
| 227 | ) |
| 228 | def test_triton_dequantize_nvfp4(monkeypatch, loaded_model_files) -> None: |
| 229 | """Test the Triton dequantization kernel against the CPU reference |
| 230 | using real NVFP4 weights from a checkpoint. |
| 231 | |
| 232 | Tests both 2D (attention projection) and 3D (stacked MoE experts). |
| 233 | """ |
| 234 | checkpoint_path = loaded_model_files["nvidia/Qwen3-30B-A3B-NVFP4"] |
| 235 | shards = cast(list[str], MOE_MODEL_CONFIGS["nvidia/Qwen3-30B-A3B-NVFP4"]["shards"]) |
| 236 | shard_path = f"{checkpoint_path}/{shards[0]}" |
| 237 | block_size = 16 |
| 238 | |
| 239 | with safe_open(shard_path, framework="pt", device="cpu") as f: |
| 240 | all_keys = list(f.keys()) |
| 241 | |
| 242 | # 2D case: attention projection |
| 243 | tensor_fp4_2d = f.get_tensor("model.layers.9.self_attn.k_proj.weight") |
| 244 | tensor_sf_2d = f.get_tensor("model.layers.9.self_attn.k_proj.weight_scale") |
| 245 | global_scale_2d = f.get_tensor("model.layers.9.self_attn.k_proj.weight_scale_2") |
| 246 | |
| 247 | # 3D case: stack ALL experts for layer 9 up_proj |
| 248 | expert_prefix = "model.layers.9.mlp.experts." |
| 249 | expert_indices = sorted( |
| 250 | int(key.split(".")[5]) |
| 251 | for key in all_keys |
| 252 | if key.startswith(expert_prefix) and key.endswith(".up_proj.weight") |
| 253 | ) |
| 254 | assert len(expert_indices) > 0 |
| 255 | |
| 256 | all_fp4 = [] |
| 257 | all_sf = [] |
| 258 | all_global_scale = [] |
| 259 | for index in expert_indices: |
| 260 | name = f"{expert_prefix}{index}.up_proj" |
| 261 | all_fp4.append(f.get_tensor(f"{name}.weight")) |
| 262 | all_sf.append(f.get_tensor(f"{name}.weight_scale")) |
| 263 | all_global_scale.append(f.get_tensor(f"{name}.weight_scale_2")) |
| 264 | |
| 265 | tensor_fp4_3d = torch.stack(all_fp4) |
| 266 | tensor_sf_3d = torch.stack(all_sf) |
| 267 | global_scale_3d = torch.stack(all_global_scale) |
| 268 | |
| 269 | test_cases = [ |
| 270 | ("2D base", tensor_fp4_2d, tensor_sf_2d, global_scale_2d), |
| 271 | ( |
| 272 | "2D 2x rows", |
| 273 | tensor_fp4_2d.repeat(2, 1), |
| 274 | tensor_sf_2d.repeat(2, 1), |
| 275 | global_scale_2d, |
| 276 | ), |
| 277 | ( |
| 278 | "2D 4x rows", |
| 279 | tensor_fp4_2d.repeat(4, 1), |
| 280 | tensor_sf_2d.repeat(4, 1), |
| 281 | global_scale_2d, |
| 282 | ), |
| 283 | ( |
| 284 | "2D 2x cols", |
| 285 | tensor_fp4_2d.repeat(1, 2), |
| 286 | tensor_sf_2d.repeat(1, 2), |
| 287 | global_scale_2d, |
| 288 | ), |
| 289 | ("3D base", tensor_fp4_3d, tensor_sf_3d, global_scale_3d), |
| 290 | ( |
| 291 | "3D 2x experts", |
| 292 | tensor_fp4_3d.repeat(2, 1, 1), |
| 293 | tensor_sf_3d.repeat(2, 1, 1), |
| 294 | global_scale_3d.repeat(2), |
| 295 | ), |
| 296 | ( |
| 297 | "3D 2x rows", |
| 298 | tensor_fp4_3d.repeat(1, 2, 1), |
| 299 | tensor_sf_3d.repeat(1, 2, 1), |
| 300 | global_scale_3d, |
| 301 | ), |
| 302 | ( |
| 303 | "3D 2x cols", |
| 304 | tensor_fp4_3d.repeat(1, 1, 2), |
| 305 | tensor_sf_3d.repeat(1, 1, 2), |
| 306 | global_scale_3d, |
| 307 | ), |
| 308 | ] |
| 309 | |
| 310 | quantiles = [0.5, 0.001, 0.999] |
| 311 | |
| 312 | # Move the E2M1 lookup table to CUDA ahead of time, as would normally |
| 313 | # happen during model loading (process_weights_after_loading). Both the |
| 314 | # Triton and PyTorch reference paths run on CUDA. |
| 315 | nvfp4_emulation_utils.kE2M1ToFloat_handle.val = ( |
| 316 | nvfp4_emulation_utils.kE2M1ToFloat_handle.val.cuda() |
| 317 | ) |
| 318 | |
| 319 | for label, tensor_fp4, tensor_sf, global_scale in test_cases: |
| 320 | fp4_cuda = tensor_fp4.cuda() |
| 321 | sf_cuda = tensor_sf.cuda() |
| 322 | gs_cuda = global_scale.cuda() |
| 323 | |
| 324 | # Triton path |
| 325 | triton_result = dequantize_to_dtype( |
| 326 | fp4_cuda, |
| 327 | sf_cuda, |
| 328 | gs_cuda, |
| 329 | torch.bfloat16, |
| 330 | block_size, |
| 331 | swizzle=False, |
| 332 | ) |
| 333 | |
| 334 | # Reference path (PyTorch ops on CUDA, Triton dispatch disabled) |
| 335 | with monkeypatch.context() as m: |
| 336 | m.setattr( |
| 337 | nvfp4_emulation_utils.current_platform, |
| 338 | "is_cuda_alike", |
| 339 | lambda: False, |
| 340 | ) |
| 341 | reference = dequantize_to_dtype( |
| 342 | fp4_cuda, |
| 343 | sf_cuda, |
| 344 | gs_cuda, |
| 345 | torch.bfloat16, |
| 346 | block_size, |
| 347 | swizzle=False, |
| 348 | ) |
| 349 | |
| 350 | torch.testing.assert_close(triton_result, reference, atol=0, rtol=0) |
| 351 | |
| 352 | # Benchmark |
| 353 | shape = list(tensor_fp4.shape) |
| 354 | |
| 355 | def _triton_bench( |
| 356 | fp4_cuda=fp4_cuda, |
| 357 | scale_cuda=sf_cuda, |
| 358 | global_scale_cuda=gs_cuda, |
| 359 | block_size=block_size, |
| 360 | ): |
| 361 | return dequantize_to_dtype( |
| 362 | fp4_cuda, |
| 363 | scale_cuda, |
| 364 | global_scale_cuda, |
| 365 | torch.bfloat16, |
| 366 | block_size, |
| 367 | swizzle=False, |
| 368 | ) |
| 369 | |
| 370 | triton_ms, triton_min, triton_max = triton.testing.do_bench( |
| 371 | _triton_bench, quantiles=quantiles |
| 372 | ) |
| 373 | |
| 374 | def _reference_bench( |
| 375 | fp4_cuda=fp4_cuda, |
| 376 | scale_cuda=sf_cuda, |
| 377 | global_scale_cuda=gs_cuda, |
| 378 | block_size=block_size, |
| 379 | ): |
| 380 | with monkeypatch.context() as m2: |
| 381 | m2.setattr( |
| 382 | nvfp4_emulation_utils.current_platform, |
| 383 | "is_cuda_alike", |
| 384 | lambda: False, |
| 385 | ) |
| 386 | dequantize_to_dtype( |
| 387 | fp4_cuda, |
| 388 | scale_cuda, |
| 389 | global_scale_cuda, |
| 390 | torch.bfloat16, |
| 391 | block_size, |
| 392 | swizzle=False, |
| 393 | ) |
| 394 | |
| 395 | ref_ms, ref_min, ref_max = triton.testing.do_bench( |
| 396 | _reference_bench, quantiles=quantiles |
| 397 | ) |
| 398 | |
| 399 | speedup = ref_ms / triton_ms if triton_ms > 0 else float("inf") |
| 400 | print(f" dequantize {label} {shape}:") |
| 401 | print( |
| 402 | f" triton: median={triton_ms:.3f}ms, " |
| 403 | f"min={triton_min:.3f}ms, max={triton_max:.3f}ms" |
| 404 | ) |
| 405 | print( |
| 406 | f" reference: median={ref_ms:.3f}ms, " |
| 407 | f"min={ref_min:.3f}ms, max={ref_max:.3f}ms" |
| 408 | ) |
| 409 | print(f" speedup: {speedup:.2f}x") |
| 410 | |
| 411 | |
| 412 | @pytest.mark.skipif( |
| 413 | not current_platform.is_cuda_alike(), |
| 414 | reason="Triton NVFP4 kernel requires CUDA.", |
| 415 | ) |
| 416 | @pytest.mark.parametrize( |
| 417 | "m, k", |
| 418 | [ |
| 419 | (1, 16), |
| 420 | (1, 4096), |
| 421 | (2, 4096), |
| 422 | (4, 4096), |
| 423 | (8, 4096), |
| 424 | (16, 4096), |
| 425 | (24, 4096), |
| 426 | (32, 4096), |
| 427 | (1, 8192), |
| 428 | (2, 8192), |
| 429 | (4, 8192), |
| 430 | (8, 8192), |
| 431 | (16, 8192), |
| 432 | (24, 8192), |
| 433 | (32, 8192), |
| 434 | (1, 32), |
| 435 | (2, 48), |
| 436 | (7, 64), |
| 437 | (16, 128), |
| 438 | (33, 160), |
| 439 | (128, 256), |
| 440 | (256, 512), |
| 441 | (1024, 1024), |
| 442 | (5120, 2048), |
| 443 | (2048, 4096), |
| 444 | (4096, 7168), |
| 445 | (8192, 8192), |
| 446 | (128, 16384), |
| 447 | ], |
| 448 | ) |
| 449 | @pytest.mark.parametrize("global_scale_value", [0.5, 1.0, 0.001]) |
| 450 | def test_triton_nvfp4_quant_dequant( |
| 451 | monkeypatch, m: int, k: int, global_scale_value: float |
| 452 | ) -> None: |
| 453 | """Test the Triton quant-dequant kernel against the CPU reference.""" |
| 454 | block_size = 16 |
| 455 | x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda") |
| 456 | global_scale = torch.tensor(global_scale_value, dtype=torch.float32, device="cuda") |
| 457 | |
| 458 | # Triton path |
| 459 | triton_result = ref_nvfp4_quant_dequant(x, global_scale, block_size) |
| 460 | |
| 461 | # CPU reference path |
| 462 | with monkeypatch.context() as mp: |
| 463 | mp.setattr( |
| 464 | nvfp4_emulation_utils.current_platform, |
| 465 | "is_cuda_alike", |
| 466 | lambda: False, |
| 467 | ) |
| 468 | reference = ref_nvfp4_quant_dequant(x.cpu(), global_scale.cpu(), block_size) |
| 469 | |
| 470 | torch.testing.assert_close(triton_result.cpu(), reference, atol=0, rtol=0) |
| 471 | |
| 472 | # Benchmark (both paths on CUDA tensors for fair comparison) |
| 473 | quantiles = [0.5, 0.001, 0.999] |
| 474 | |
| 475 | def _triton_bench( |
| 476 | input_tensor=x, input_global_scale=global_scale, input_block_size=block_size |
| 477 | ): |
| 478 | return ref_nvfp4_quant_dequant( |
| 479 | input_tensor, input_global_scale, input_block_size |
| 480 | ) |
| 481 | |
| 482 | triton_ms, triton_min, triton_max = triton.testing.do_bench( |
| 483 | _triton_bench, quantiles=quantiles |
| 484 | ) |
| 485 | |
| 486 | def _reference_bench( |
| 487 | input_tensor=x, input_global_scale=global_scale, input_block_size=block_size |
| 488 | ): |
| 489 | with monkeypatch.context() as mp2: |
| 490 | mp2.setattr( |
| 491 | nvfp4_emulation_utils.current_platform, |
| 492 | "is_cuda_alike", |
| 493 | lambda: False, |
| 494 | ) |
| 495 | ref_nvfp4_quant_dequant(input_tensor, input_global_scale, input_block_size) |
| 496 | |
| 497 | ref_ms, ref_min, ref_max = triton.testing.do_bench( |
| 498 | _reference_bench, quantiles=quantiles |
| 499 | ) |
| 500 | |
| 501 | speedup = ref_ms / triton_ms if triton_ms > 0 else float("inf") |
| 502 | print(f" quant_dequant [{m}x{k}] gs={global_scale_value}:") |
| 503 | print( |
| 504 | f" triton: median={triton_ms:.3f}ms, " |
| 505 | f"min={triton_min:.3f}ms, max={triton_max:.3f}ms" |
| 506 | ) |
| 507 | print( |
| 508 | f" reference: median={ref_ms:.3f}ms, " |
| 509 | f"min={ref_min:.3f}ms, max={ref_max:.3f}ms" |
| 510 | ) |
| 511 | print(f" speedup: {speedup:.2f}x") |
| 512 | |
| 513 | |
| 514 | def _load_nvfp4_moe_weights( |
| 515 | model_files: dict[str, str], |
| 516 | model_id: str, |
| 517 | tensor_parallel_size: int, |
| 518 | max_experts: int | None = None, |
| 519 | ): |
| 520 | """Load and stack NVFP4 MoE weights from checkpoint shards. |
| 521 | |
| 522 | Returns (w1, w1_scale, w1_gscale, w2, w2_scale, w2_gscale, |
| 523 | a1_gscale, a2_gscale, num_experts, hidden_dim, |
| 524 | intermediate_size). |
| 525 | |
| 526 | When max_experts is set, only the first max_experts experts are loaded. |
| 527 | |
| 528 | When tensor_parallel_size > 1, the N dimension of w1 and the K |
| 529 | dimension of w2 are narrowed to the first TP shard (simulating |
| 530 | column-parallel on w1 / row-parallel on w2). |
| 531 | """ |
| 532 | cfg = MOE_MODEL_CONFIGS[model_id] |
| 533 | shards = cast(list[str], cfg["shards"]) |
| 534 | checkpoint_path = model_files[model_id] |
| 535 | |
| 536 | expert_prefix = cfg["expert_prefix"] |
| 537 | idx_pos = cast(int, cfg["expert_idx_pos"]) |
| 538 | |
| 539 | # Collect all tensors across shards into a flat dict — an expert's |
| 540 | # tensors may be split across multiple shard files. |
| 541 | all_tensors: dict[str, torch.Tensor] = {} |
| 542 | for shard_name in shards: |
| 543 | shard_path = f"{checkpoint_path}/{shard_name}" |
| 544 | with safe_open(shard_path, framework="pt", device="cpu") as f: |
| 545 | for key in f.keys(): # noqa: SIM118 |
| 546 | if key.startswith(expert_prefix): |
| 547 | all_tensors[key] = f.get_tensor(key) |
| 548 | |
| 549 | expert_indices = sorted( |
| 550 | { |
| 551 | int(key.split(".")[idx_pos]) |
| 552 | for key in all_tensors |
| 553 | if key.endswith(".gate_proj.weight") |
| 554 | } |
| 555 | ) |
| 556 | if max_experts is not None: |
| 557 | expert_indices = expert_indices[:max_experts] |
| 558 | num_experts = len(expert_indices) |
| 559 | |
| 560 | gate_weights, up_weights, down_weights = [], [], [] |
| 561 | gate_scales, up_scales, down_scales = [], [], [] |
| 562 | gate_gscales, up_gscales, down_gscales = [], [], [] |
| 563 | a1_scales, a2_scales = [], [] |
| 564 | |
| 565 | for idx in expert_indices: |
| 566 | prefix = f"{expert_prefix}{idx}" |
| 567 | gate_weights.append(all_tensors[f"{prefix}.gate_proj.weight"]) |
| 568 | gate_scales.append(all_tensors[f"{prefix}.gate_proj.weight_scale"]) |
| 569 | gate_gscales.append(all_tensors[f"{prefix}.gate_proj.weight_scale_2"]) |
| 570 | up_weights.append(all_tensors[f"{prefix}.up_proj.weight"]) |
| 571 | up_scales.append(all_tensors[f"{prefix}.up_proj.weight_scale"]) |
| 572 | up_gscales.append(all_tensors[f"{prefix}.up_proj.weight_scale_2"]) |
| 573 | down_weights.append(all_tensors[f"{prefix}.down_proj.weight"]) |
| 574 | down_scales.append(all_tensors[f"{prefix}.down_proj.weight_scale"]) |
| 575 | down_gscales.append(all_tensors[f"{prefix}.down_proj.weight_scale_2"]) |
| 576 | a1_scales.append(all_tensors[f"{prefix}.gate_proj.input_scale"]) |
| 577 | a2_scales.append(all_tensors[f"{prefix}.down_proj.input_scale"]) |
| 578 | |
| 579 | # Stack into MoE format. |
| 580 | # w1 = [E, 2*intermediate, hidden//2] (gate + up concatenated) |
| 581 | w1 = torch.stack( |
| 582 | [torch.cat([g, u], dim=0) for g, u in zip(gate_weights, up_weights)] |
| 583 | ).cuda() |
| 584 | w1_scale = torch.stack( |
| 585 | [torch.cat([g, u], dim=0) for g, u in zip(gate_scales, up_scales)] |
| 586 | ).cuda() |
| 587 | w1_gscale = torch.stack(gate_gscales).cuda() |
| 588 | |
| 589 | # w2 = [E, hidden, intermediate//2] |
| 590 | w2 = torch.stack(down_weights).cuda() |
| 591 | w2_scale = torch.stack(down_scales).cuda() |
| 592 | w2_gscale = torch.stack(down_gscales).cuda() |
| 593 | |
| 594 | a13_scale_raw = torch.stack(a1_scales).cuda() |
| 595 | a2_scale_raw = torch.stack(a2_scales).cuda() |
| 596 | |
| 597 | # Apply EMULATION transforms (matches oracle/nvfp4.py). |
| 598 | nvfp4_emulation_utils.kE2M1ToFloat_handle.val = ( |
| 599 | nvfp4_emulation_utils.kE2M1ToFloat_handle.val.cuda() |
| 600 | ) |
| 601 | a1_gscale = 1.0 / a13_scale_raw.max().to(torch.float32) |
| 602 | a2_gscale = 1.0 / a2_scale_raw.max().to(torch.float32) |
| 603 | |
| 604 | # ── Simulate TP sharding ── |
| 605 | # w1 (gate_up): column-parallel → shard the N dimension (dim 1). |
| 606 | # w2 (down): row-parallel → shard the K dimension (dim 2, |
| 607 | # which is the packed K//2 dim). |
| 608 | # Scales follow the same sharding on the corresponding dimension. |
| 609 | tp = tensor_parallel_size |
| 610 | if tp > 1: |
| 611 | n1 = w1.size(1) // tp |
| 612 | w1 = w1[:, :n1, :].contiguous() |
| 613 | w1_scale = w1_scale[:, :n1, :].contiguous() |
| 614 | |
| 615 | k2_packed = w2.size(2) // tp |
| 616 | k2_scale = w2_scale.size(2) // tp |
| 617 | w2 = w2[:, :, :k2_packed].contiguous() |
| 618 | w2_scale = w2_scale[:, :, :k2_scale].contiguous() |
| 619 | |
| 620 | hidden_dim = w1.size(2) * 2 |
| 621 | intermediate_size = w1.size(1) // 2 |
| 622 | |
| 623 | return ( |
| 624 | w1, |
| 625 | w1_scale, |
| 626 | w1_gscale, |
| 627 | w2, |
| 628 | w2_scale, |
| 629 | w2_gscale, |
| 630 | a1_gscale, |
| 631 | a2_gscale, |
| 632 | num_experts, |
| 633 | hidden_dim, |
| 634 | intermediate_size, |
| 635 | ) |
| 636 | |
| 637 | |
| 638 | @pytest.mark.skipif( |
| 639 | not current_platform.is_cuda_alike(), |
| 640 | reason="Triton NVFP4 kernel requires CUDA.", |
| 641 | ) |
| 642 | @pytest.mark.parametrize("num_tokens", [1, 2, 4, 1024]) |
| 643 | @pytest.mark.parametrize("top_k", [4]) |
| 644 | @pytest.mark.parametrize("model_id", list(MOE_MODEL_CONFIGS.keys())) |
| 645 | @pytest.mark.parametrize( |
| 646 | "tensor_parallel_size", |
| 647 | [pytest.param(val, id=f"tensor_parallel_size:{val}") for val in [1, 2, 4, 8]], |
| 648 | ) |
| 649 | def test_nvfp4_moe_correctness( |
| 650 | loaded_model_files, |
| 651 | num_tokens: int, |
| 652 | top_k: int, |
| 653 | model_id: str, |
| 654 | tensor_parallel_size: int, |
| 655 | ) -> None: |
| 656 | """Compare Nvfp4QuantizationEmulationTritonExperts (fused weight dequant + compute) |
| 657 | against the unfused reference Nvfp4QuantizationEmulationTritonExpertsReference. |
| 658 | |
| 659 | Both must produce bit-identical results. |
| 660 | """ |
| 661 | num_test_experts = max(8, top_k) |
| 662 | ( |
| 663 | w1, |
| 664 | w1_scale, |
| 665 | w1_gscale, |
| 666 | w2, |
| 667 | w2_scale, |
| 668 | w2_gscale, |
| 669 | a1_gscale, |
| 670 | a2_gscale, |
| 671 | num_experts, |
| 672 | hidden_dim, |
| 673 | intermediate_size, |
| 674 | ) = _load_nvfp4_moe_weights( |
| 675 | loaded_model_files, |
| 676 | model_id, |
| 677 | tensor_parallel_size, |
| 678 | max_experts=num_test_experts, |
| 679 | ) |
| 680 | |
| 681 | moe_config = FusedMoEConfig( |
| 682 | num_experts=num_experts, |
| 683 | experts_per_token=top_k, |
| 684 | hidden_dim=hidden_dim, |
| 685 | intermediate_size=intermediate_size, |
| 686 | num_local_experts=num_experts, |
| 687 | num_logical_experts=num_experts, |
| 688 | moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(), |
| 689 | activation=MoEActivation.SILU, |
| 690 | in_dtype=torch.bfloat16, |
| 691 | device="cuda", |
| 692 | routing_method=RoutingMethodType.TopK, |
| 693 | max_num_tokens=512, |
| 694 | ) |
| 695 | |
| 696 | def _make_quant_config(): |
| 697 | return nvfp4_moe_quant_config( |
| 698 | g1_alphas=w1_gscale.clone(), |
| 699 | g2_alphas=w2_gscale.clone(), |
| 700 | a1_gscale=a1_gscale.clone(), |
| 701 | a2_gscale=a2_gscale.clone(), |
| 702 | w1_scale=w1_scale.clone(), |
| 703 | w2_scale=w2_scale.clone(), |
| 704 | ) |
| 705 | |
| 706 | ref_experts = Nvfp4QuantizationEmulationTritonExpertsReference( |
| 707 | moe_config=moe_config, |
| 708 | quant_config=_make_quant_config(), |
| 709 | ) |
| 710 | fused_experts = Nvfp4QuantizationEmulationTritonExperts( |
| 711 | moe_config=moe_config, |
| 712 | quant_config=_make_quant_config(), |
| 713 | ) |
| 714 | |
| 715 | torch.manual_seed(42) |
| 716 | hidden_states = torch.randn( |
| 717 | num_tokens, hidden_dim, dtype=torch.bfloat16, device="cuda" |
| 718 | ) |
| 719 | |
| 720 | topk_weights = torch.randn( |
| 721 | num_tokens, top_k, dtype=torch.float32, device="cuda" |
| 722 | ).softmax(dim=-1) |
| 723 | topk_ids = torch.stack( |
| 724 | [torch.randperm(num_experts, device="cuda")[:top_k] for _ in range(num_tokens)] |
| 725 | ).to(torch.int32) |
| 726 | |
| 727 | N = w1.size(1) # 2 * intermediate |
| 728 | K = hidden_dim |
| 729 | |
| 730 | ws13_size = num_tokens * top_k * max(intermediate_size, K) |
| 731 | ws2_size = num_tokens * top_k * max(N, K) |
| 732 | |
| 733 | workspace13_ref = torch.zeros(ws13_size, dtype=torch.bfloat16, device="cuda") |
| 734 | workspace2_ref = torch.zeros(ws2_size, dtype=torch.bfloat16, device="cuda") |
| 735 | output_ref = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device="cuda") |
| 736 | |
| 737 | workspace13_fused = torch.zeros_like(workspace13_ref) |
| 738 | workspace2_fused = torch.zeros_like(workspace2_ref) |
| 739 | output_fused = torch.zeros_like(output_ref) |
| 740 | |
| 741 | apply_kwargs = dict( |
| 742 | hidden_states=hidden_states, |
| 743 | w1=w1, |
| 744 | w2=w2, |
| 745 | topk_weights=topk_weights, |
| 746 | topk_ids=topk_ids, |
| 747 | activation=MoEActivation.SILU, |
| 748 | global_num_experts=num_experts, |
| 749 | expert_map=None, |
| 750 | a1q_scale=None, |
| 751 | a2_scale=None, |
| 752 | expert_tokens_meta=None, |
| 753 | apply_router_weight_on_input=False, |
| 754 | ) |
| 755 | |
| 756 | # Unfused reference. |
| 757 | ref_experts.apply( |
| 758 | output=output_ref, |
| 759 | workspace13=workspace13_ref, |
| 760 | workspace2=workspace2_ref, |
| 761 | **apply_kwargs, |
| 762 | ) |
| 763 | |
| 764 | # Fused implementation. |
| 765 | fused_experts.apply( |
| 766 | output=output_fused, |
| 767 | workspace13=workspace13_fused, |
| 768 | workspace2=workspace2_fused, |
| 769 | **apply_kwargs, |
| 770 | ) |
| 771 | |
| 772 | # Not strict equality on H100, MI325, MI300 (< 0.1% elements). |
| 773 | # The fused on-the-fly dequant path can lower to a slightly |
| 774 | # different Triton/MMA tiling than the pre-dequantized |
| 775 | # reference; experiments with reference-like tiling/masking |
| 776 | # reduced some diffs were not kept because they regress |
| 777 | # the fused kernel speed. |
| 778 | # Strict equality validated on MI355. |
| 779 | torch.testing.assert_close( |
| 780 | output_fused, |
| 781 | output_ref, |
| 782 | atol=0.0 if on_gfx950() else 0.02, |
| 783 | rtol=0, |
| 784 | ) |
| 785 | |
| 786 | |
| 787 | def _dense_nvfp4_reference( |
| 788 | layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None |
| 789 | ) -> torch.Tensor: |
| 790 | """Decode E2M1 nibbles and both scale levels without production helpers.""" |
| 791 | levels = kE2M1ToFloat.cpu() |
| 792 | weight = break_fp4_bytes(layer.weight.cpu(), torch.float32) |
| 793 | weight *= ( |
| 794 | layer.weight_scale.cpu().float() * layer.weight_global_scale.cpu() |
| 795 | ).repeat_interleave(16, dim=-1) |
| 796 | blocks = x.cpu().float().reshape(x.shape[0], -1, 16) |
| 797 | global_scale = layer.input_global_scale_inv.cpu() |
| 798 | scales = ( |
| 799 | (blocks.abs().amax(-1, keepdim=True) * (1.0 / FLOAT4_E2M1_MAX) * global_scale) |
| 800 | .clamp(max=FLOAT8_E4M3_MAX) |
| 801 | .to(torch.float8_e4m3fn) |
| 802 | .float() |
| 803 | ) |
| 804 | inverse = torch.where(scales == 0, 0, global_scale / scales) |
| 805 | scaled = (blocks * inverse).abs().clamp(max=FLOAT4_E2M1_MAX) |
| 806 | distances = (scaled.unsqueeze(-1) - levels).abs() |
| 807 | # E2M1 ties round to the even significand (even encoding). |
| 808 | nearest_mask = distances == distances.amin(-1, keepdim=True) |
| 809 | even_mask = nearest_mask & (torch.arange(8, device="cpu") % 2 == 0) |
| 810 | candidates = torch.where(even_mask.any(-1, keepdim=True), even_mask, nearest_mask) |
| 811 | nearest = candidates.to(torch.int8).argmax(-1) |
| 812 | activation = levels[nearest] * blocks.sign() * (scales / global_scale) |
| 813 | expected = activation.reshape_as(x).to(x.dtype) @ weight.to(x.dtype).T |
| 814 | if bias is not None: |
| 815 | expected += bias.cpu() |
| 816 | return expected |
| 817 | |
| 818 | |
| 819 | @pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm NVFP4 fallback") |
| 820 | @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) |
| 821 | @pytest.mark.parametrize("m,n,k", [(1, 32, 64), (17, 128, 256)]) |
| 822 | @pytest.mark.parametrize("with_bias", [False, True]) |
| 823 | def test_nvfp4_dense_emulation( |
| 824 | dtype: torch.dtype, m: int, n: int, k: int, with_bias: bool |
| 825 | ) -> None: |
| 826 | """W4A4 fallback preserves packed scales, bias and outputs across replay.""" |
| 827 | set_random_seed(47) |
| 828 | layer = torch.nn.Module() |
| 829 | # Every signed E2M1 code appears, including both zero encodings. |
| 830 | codes = (torch.randperm(n * k, device="cuda") % 16).reshape(n, k) |
| 831 | layer.weight = (codes[:, ::2] | (codes[:, 1::2] << 4)).to(torch.uint8) |
| 832 | layer.weight_scale = ( |
| 833 | torch.arange(n * k // 16, device="cuda").reshape(n, k // 16) % 7 / 8 |
| 834 | ).to(torch.float8_e4m3fn) |
| 835 | layer.weight_global_scale = torch.tensor(0.25, device="cuda") |
| 836 | layer.input_global_scale_inv = torch.tensor(3.0, device="cuda") |
| 837 | x = torch.randn(m, k, device="cuda", dtype=dtype) |
| 838 | x[:, :16] = 0 |
| 839 | bias = torch.randn(n, device="cuda", dtype=dtype) if with_bias else None |
| 840 | with set_current_vllm_config( |
| 841 | VllmConfig(kernel_config=KernelConfig(linear_backend="emulation")) |
| 842 | ): |
| 843 | kernel = init_nvfp4_linear_kernel() |
| 844 | assert isinstance(kernel, EmulationNvFp4LinearKernel) |
| 845 | kernel.process_weights_after_loading(layer) |
| 846 | original_weight = layer.weight.clone() |
| 847 | original_scale = layer.weight_scale.view(torch.uint8).clone() |
| 848 | actual = kernel.apply_weights(layer, x, bias) |
| 849 | torch.testing.assert_close( |
| 850 | actual.cpu(), _dense_nvfp4_reference(layer, x, bias), atol=0.03, rtol=0.01 |
| 851 | ) |
| 852 | assert actual.abs().max() > 0.1 |
| 853 | saved = actual.clone() |
| 854 | graph = torch.cuda.CUDAGraph() |
| 855 | with torch.cuda.graph(graph): |
| 856 | captured = kernel.apply_weights(layer, x, bias) |
| 857 | for sign in (-1, 1): |
| 858 | x.mul_(sign * 0.75) |
| 859 | graph.replay() |
| 860 | torch.accelerator.synchronize() |
| 861 | torch.testing.assert_close( |
| 862 | captured.cpu(), _dense_nvfp4_reference(layer, x, bias), atol=0.03, rtol=0.01 |
| 863 | ) |
| 864 | torch.testing.assert_close(actual, saved, atol=0, rtol=0) |
| 865 | assert torch.equal(layer.weight, original_weight) |
| 866 | assert torch.equal(layer.weight_scale.view(torch.uint8), original_scale) |