| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Tests online quantization.""" |
| 4 | |
| 5 | import json |
| 6 | from pathlib import Path |
| 7 | from types import SimpleNamespace |
| 8 | from typing import Any, cast |
| 9 | from unittest.mock import Mock |
| 10 | |
| 11 | import pytest |
| 12 | import torch |
| 13 | from torch.distributed import ProcessGroup |
| 14 | |
| 15 | from tests.quantization.utils import ( |
| 16 | _test_online_quant_peak_mem_impl, |
| 17 | is_quant_method_supported, |
| 18 | load_model_without_vllm_runner, |
| 19 | quant_config_args, |
| 20 | ) |
| 21 | from vllm import _custom_ops as ops |
| 22 | from vllm._aiter_ops import rocm_aiter_ops |
| 23 | from vllm._custom_ops import scaled_fp4_quant |
| 24 | from vllm.config.cache import CacheConfig |
| 25 | from vllm.config.kernel import KernelConfig |
| 26 | from vllm.config.load import LoadConfig |
| 27 | from vllm.config.model import ModelConfig |
| 28 | from vllm.config.quantization import ( |
| 29 | QuantizationConfigArgs, |
| 30 | resolve_quantization_config, |
| 31 | ) |
| 32 | from vllm.config.vllm import VllmConfig |
| 33 | from vllm.forward_context import set_forward_context |
| 34 | from vllm.model_executor.kernels.linear.mxfp8.emulation import ( |
| 35 | EmulationMxfp8LinearKernel, |
| 36 | ) |
| 37 | from vllm.model_executor.kernels.linear.mxfp8.marlin import ( |
| 38 | MarlinMxfp8LinearKernel, |
| 39 | ) |
| 40 | from vllm.model_executor.kernels.linear.scaled_mm import ( |
| 41 | MarlinFP8ScaledMMLinearKernel, |
| 42 | ) |
| 43 | from vllm.model_executor.layers.attention import Attention |
| 44 | from vllm.model_executor.layers.fused_moe import FusedMoEFactory |
| 45 | from vllm.model_executor.layers.fused_moe.oracle import nvfp4 as nvfp4_oracle |
| 46 | from vllm.model_executor.layers.linear import ( |
| 47 | ColumnParallelLinear, |
| 48 | LinearBase, |
| 49 | UnquantizedLinearMethod, |
| 50 | ) |
| 51 | from vllm.model_executor.layers.quantization.auto_awq import AutoAWQConfig |
| 52 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 |
| 53 | CompressedTensorsConfig, |
| 54 | CompressedTensorsLinearMethod, |
| 55 | ) |
| 56 | from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 |
| 57 | CompressedTensorsMoEMethod, |
| 58 | ) |
| 59 | from vllm.model_executor.layers.quantization.modelopt import ( |
| 60 | ModelOptFp8Config, |
| 61 | ModelOptMxFp8Config, |
| 62 | ) |
| 63 | from vllm.model_executor.layers.quantization.online.base import ( |
| 64 | OnlineQuantizationConfig, |
| 65 | ) |
| 66 | from vllm.model_executor.layers.quantization.online.fp8 import ( |
| 67 | Fp8PerBlockOnlineLinearMethod, |
| 68 | Fp8PerBlockOnlineMoEMethod, |
| 69 | Fp8PerTensorOnlineLinearMethod, |
| 70 | Fp8PerTensorOnlineMoEMethod, |
| 71 | Fp8PtpcOnlineLinearMethod, |
| 72 | Fp8PtpcOnlineMoEMethod, |
| 73 | _fp8_channel_scale, |
| 74 | _fp8_quant_per_channel, |
| 75 | _fp8_scale, |
| 76 | _is_tp_sharded, |
| 77 | ) |
| 78 | from vllm.model_executor.layers.quantization.online.int8 import Int8OnlineMoEMethod |
| 79 | from vllm.model_executor.layers.quantization.online.mxfp4 import ( |
| 80 | Mxfp4OnlineLinearMethod, |
| 81 | Mxfp4OnlineMoEMethod, |
| 82 | ) |
| 83 | from vllm.model_executor.layers.quantization.online.mxfp8 import ( |
| 84 | Mxfp8OnlineLinearMethod, |
| 85 | ) |
| 86 | from vllm.model_executor.layers.quantization.online.nvfp4 import ( |
| 87 | Nvfp4OnlineMoEMethod, |
| 88 | _quantize_moe_weight_to_nvfp4, |
| 89 | ) |
| 90 | from vllm.model_executor.layers.quantization.quark.quark import ( |
| 91 | QuarkConfig, |
| 92 | QuarkLinearMethod, |
| 93 | ) |
| 94 | from vllm.model_executor.layers.quantization.utils import quant_utils |
| 95 | from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( |
| 96 | MXFP8_VALUE_DTYPE, |
| 97 | ) |
| 98 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 99 | amax_for_moe_weight_quant, |
| 100 | amax_for_tp_weight_quant, |
| 101 | kMxfp8Dynamic, |
| 102 | weight_amax, |
| 103 | ) |
| 104 | from vllm.model_executor.layers.vocab_parallel_embedding import ( |
| 105 | VocabParallelEmbedding, |
| 106 | ) |
| 107 | from vllm.model_executor.model_loader import weight_utils |
| 108 | from vllm.model_executor.model_loader.base_loader import log_online_quantization |
| 109 | from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader |
| 110 | from vllm.model_executor.models.granitemoe import ( |
| 111 | GraniteMoeModel, |
| 112 | ) |
| 113 | from vllm.platforms import current_platform |
| 114 | from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe |
| 115 | |
| 116 | if current_platform.is_rocm(): |
| 117 | from vllm.platforms.rocm import on_gfx942, on_gfx950 |
| 118 | else: |
| 119 | |
| 120 | def on_gfx950() -> bool: |
| 121 | return False |
| 122 | |
| 123 | def on_gfx942() -> bool: |
| 124 | return False |
| 125 | |
| 126 | |
| 127 | DEVICE = current_platform.device_type |
| 128 | GRANITE_MODEL_NAME = "ibm-granite/granite-3.0-1b-a400m-base" |
| 129 | PARTIALLY_PREQUANTIZED_MODEL_NAME = ( |
| 130 | "nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable" |
| 131 | ) |
| 132 | |
| 133 | |
| 134 | @pytest.mark.parametrize( |
| 135 | "overrides", |
| 136 | [ |
| 137 | None, |
| 138 | {"ignore": ["model.layers.0.self_attn.o_proj"]}, |
| 139 | { |
| 140 | "linear": "fp8_per_block", |
| 141 | "ignore": ["model.layers.0.self_attn.o_proj"], |
| 142 | }, |
| 143 | ], |
| 144 | ids=["defaults", "ignore", "linear-and-ignore"], |
| 145 | ) |
| 146 | def test_legacy_fp8_online_quantization_uses_per_tensor_shorthand( |
| 147 | tmp_path, caplog, disable_log_dedup, overrides |
| 148 | ) -> None: |
| 149 | """The legacy alias preserves overrides and fills in per-tensor defaults.""" |
| 150 | _write_minimal_llama_config(tmp_path) |
| 151 | model_config = ModelConfig( |
| 152 | model=str(tmp_path), quantization="fp8", quantization_config=overrides |
| 153 | ) |
| 154 | |
| 155 | result = weight_utils.get_quant_config(model_config, LoadConfig()) |
| 156 | |
| 157 | assert isinstance(result, OnlineQuantizationConfig) |
| 158 | assert result.args == resolve_quantization_config("fp8_per_tensor", overrides) |
| 159 | assert "--quantization fp8 is deprecated for online quantization" in caplog.text |
| 160 | assert "--quantization fp8_per_tensor instead" in caplog.text |
| 161 | |
| 162 | |
| 163 | @pytest.mark.parametrize( |
| 164 | "shorthand,method_cls", |
| 165 | [ |
| 166 | ("fp8_per_tensor", Fp8PerTensorOnlineLinearMethod), |
| 167 | ("fp8_per_block", Fp8PerBlockOnlineLinearMethod), |
| 168 | ("fp8_per_channel", Fp8PtpcOnlineLinearMethod), |
| 169 | ], |
| 170 | ) |
| 171 | def test_qianfan_online_fp8_keeps_vision_layers_unquantized( |
| 172 | shorthand, method_cls |
| 173 | ) -> None: |
| 174 | """Model exclusions preserve the language backbone and shared defaults.""" |
| 175 | from vllm.model_executor.models.qianfan_ocr import ( |
| 176 | QianfanOCRForConditionalGeneration, |
| 177 | ) |
| 178 | |
| 179 | args = resolve_quantization_config(shorthand, None) |
| 180 | assert args is not None |
| 181 | config = OnlineQuantizationConfig(args) |
| 182 | other_config = OnlineQuantizationConfig(args) |
| 183 | model = QianfanOCRForConditionalGeneration.__new__( |
| 184 | QianfanOCRForConditionalGeneration # type: ignore[type-abstract] |
| 185 | ) |
| 186 | model._patch_quant_config( |
| 187 | SimpleNamespace(vision_config=SimpleNamespace(num_hidden_layers=2)), |
| 188 | config, |
| 189 | ) |
| 190 | |
| 191 | layer = Mock(spec=LinearBase) |
| 192 | for prefix in ( |
| 193 | "vision_model.encoder.layers.0.attn.qkv", |
| 194 | "vision_model.encoder.layers.1.mlp.fc2", |
| 195 | "language_model.lm_head", |
| 196 | "mlp1.1", |
| 197 | "mlp1.3", |
| 198 | ): |
| 199 | assert config.resolve_quant_method_cls(layer, prefix) is None |
| 200 | |
| 201 | resolved = config.resolve_quant_method_cls( |
| 202 | layer, "language_model.model.layers.0.self_attn.qkv_proj" |
| 203 | ) |
| 204 | assert resolved is not None and resolved[-1] is method_cls |
| 205 | assert other_config.ignored_layers == [] |
| 206 | assert args.ignore == [] |
| 207 | |
| 208 | |
| 209 | def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed( |
| 210 | monkeypatch, |
| 211 | ) -> None: |
| 212 | method = object.__new__(Nvfp4OnlineMoEMethod) |
| 213 | method.moe = SimpleNamespace(is_act_and_mul=True) |
| 214 | method.nvfp4_backend = object() |
| 215 | method.experts_cls = object |
| 216 | method.moe_quant_config = None |
| 217 | method.moe_kernel = None |
| 218 | |
| 219 | layer = Mock() |
| 220 | converted_weights = [ |
| 221 | tuple(torch.full((2,), value + index) for index in range(8)) |
| 222 | for value in (1.0, 2.0, 0.5) |
| 223 | ] |
| 224 | convert_weights = Mock(side_effect=converted_weights) |
| 225 | quant_config = SimpleNamespace( |
| 226 | g1_alphas=converted_weights[0][2].clone(), |
| 227 | g2_alphas=converted_weights[0][6].clone(), |
| 228 | ) |
| 229 | alpha_ptrs = (quant_config.g1_alphas.data_ptr(), quant_config.g2_alphas.data_ptr()) |
| 230 | observed_scales = [] |
| 231 | process_weights = Mock( |
| 232 | side_effect=lambda _: observed_scales.append( |
| 233 | (quant_config.g1_alphas.clone(), quant_config.g2_alphas.clone()) |
| 234 | ) |
| 235 | ) |
| 236 | kernel = SimpleNamespace( |
| 237 | fused_experts=SimpleNamespace( |
| 238 | process_weights_after_loading=process_weights, |
| 239 | ) |
| 240 | ) |
| 241 | make_kernel = Mock(return_value=kernel) |
| 242 | get_quant_config = Mock(return_value=quant_config) |
| 243 | monkeypatch.setattr(method, "get_fused_moe_quant_config", get_quant_config) |
| 244 | |
| 245 | monkeypatch.setattr( |
| 246 | "vllm.model_executor.layers.quantization.online.nvfp4." |
| 247 | "convert_to_nvfp4_moe_kernel_format", |
| 248 | convert_weights, |
| 249 | ) |
| 250 | monkeypatch.setattr( |
| 251 | "vllm.model_executor.layers.quantization.online.nvfp4.replace_parameter", |
| 252 | Mock(), |
| 253 | ) |
| 254 | monkeypatch.setattr( |
| 255 | "vllm.model_executor.layers.quantization.online.nvfp4.make_nvfp4_moe_kernel", |
| 256 | make_kernel, |
| 257 | ) |
| 258 | |
| 259 | for weights in converted_weights: |
| 260 | method._setup_kernel(layer) |
| 261 | # Expert post-processing must observe fresh scales, with stable storage. |
| 262 | torch.testing.assert_close(observed_scales[-1][0], weights[2]) |
| 263 | torch.testing.assert_close(observed_scales[-1][1], weights[6]) |
| 264 | assert method.moe_quant_config is quant_config |
| 265 | assert ( |
| 266 | quant_config.g1_alphas.data_ptr(), |
| 267 | quant_config.g2_alphas.data_ptr(), |
| 268 | ) == alpha_ptrs |
| 269 | |
| 270 | assert method.moe_kernel is kernel |
| 271 | assert convert_weights.call_count == 3 |
| 272 | make_kernel.assert_called_once() |
| 273 | get_quant_config.assert_called_once() |
| 274 | assert process_weights.call_count == 3 |
| 275 | |
| 276 | |
| 277 | def _fully_quantized_quark_config() -> QuarkConfig: |
| 278 | return QuarkConfig( |
| 279 | { |
| 280 | "exclude": [], |
| 281 | "global_quant_config": { |
| 282 | "weight": { |
| 283 | "dtype": "int8", |
| 284 | "qscheme": "per_tensor", |
| 285 | "is_dynamic": False, |
| 286 | "symmetric": True, |
| 287 | }, |
| 288 | "input_tensors": { |
| 289 | "dtype": "int8", |
| 290 | "qscheme": "per_tensor", |
| 291 | "is_dynamic": False, |
| 292 | "symmetric": True, |
| 293 | }, |
| 294 | }, |
| 295 | "layer_quant_config": {}, |
| 296 | "layer_type_quant_config": {}, |
| 297 | } |
| 298 | ) |
| 299 | |
| 300 | |
| 301 | def _fully_quantized_modelopt_config() -> ModelOptFp8Config: |
| 302 | return ModelOptFp8Config( |
| 303 | quant_method="FP8", |
| 304 | is_checkpoint_fp8_serialized=True, |
| 305 | kv_cache_quant_method=None, |
| 306 | exclude_modules=[], |
| 307 | ) |
| 308 | |
| 309 | |
| 310 | def _moe_only_compressed_tensors_config() -> CompressedTensorsConfig: |
| 311 | return CompressedTensorsConfig( |
| 312 | target_scheme_map={"RoutedExperts": {}}, |
| 313 | ignore=[], |
| 314 | quant_format="pack-quantized", |
| 315 | ) |
| 316 | |
| 317 | |
| 318 | def _write_minimal_llama_config( |
| 319 | model_path: Path, quantization_config: dict[str, Any] | None = None |
| 320 | ) -> None: |
| 321 | config: dict[str, Any] = { |
| 322 | "architectures": ["LlamaForCausalLM"], |
| 323 | "model_type": "llama", |
| 324 | "hidden_size": 32, |
| 325 | "intermediate_size": 64, |
| 326 | "num_attention_heads": 4, |
| 327 | "num_hidden_layers": 1, |
| 328 | "vocab_size": 32, |
| 329 | "max_position_embeddings": 32, |
| 330 | } |
| 331 | if quantization_config is not None: |
| 332 | config["quantization_config"] = quantization_config |
| 333 | (model_path / "config.json").write_text(json.dumps(config)) |
| 334 | |
| 335 | |
| 336 | @pytest.mark.parametrize( |
| 337 | "checkpoint_config_factory,raises_conflict", |
| 338 | [ |
| 339 | pytest.param(_fully_quantized_quark_config, True, id="quark"), |
| 340 | pytest.param(_fully_quantized_modelopt_config, True, id="modelopt"), |
| 341 | pytest.param( |
| 342 | _moe_only_compressed_tensors_config, |
| 343 | False, |
| 344 | id="compressed_tensors", |
| 345 | ), |
| 346 | ], |
| 347 | ) |
| 348 | def test_online_prequantized_compatibility( |
| 349 | checkpoint_config_factory, |
| 350 | raises_conflict: bool, |
| 351 | default_vllm_config, |
| 352 | dist_init, |
| 353 | ) -> None: |
| 354 | """Online weights replace only layers left unquantized by a checkpoint.""" |
| 355 | default_vllm_config.model_config = ModelConfig() |
| 356 | checkpoint_config = checkpoint_config_factory() |
| 357 | |
| 358 | checkpoint_config.online_quantization_config = OnlineQuantizationConfig( |
| 359 | quant_config_args(linear="mxfp8") |
| 360 | ) |
| 361 | config = checkpoint_config |
| 362 | |
| 363 | layer_kwargs = { |
| 364 | "input_size": 32, |
| 365 | "output_size": 32, |
| 366 | "bias": False, |
| 367 | "params_dtype": torch.bfloat16, |
| 368 | "quant_config": config, |
| 369 | "prefix": "model.layers.0.self_attn.o_proj", |
| 370 | "disable_tp": True, |
| 371 | } |
| 372 | |
| 373 | if raises_conflict: |
| 374 | with pytest.raises(ValueError, match="pre-quantized layer"): |
| 375 | ColumnParallelLinear(**layer_kwargs) |
| 376 | else: |
| 377 | layer = ColumnParallelLinear(**layer_kwargs) |
| 378 | assert isinstance(layer.quant_method, Mxfp8OnlineLinearMethod) |
| 379 | |
| 380 | |
| 381 | def test_online_target_rejects_prequantized_layer( |
| 382 | default_vllm_config, dist_init |
| 383 | ) -> None: |
| 384 | """A targets match participates in checkpoint compatibility checks.""" |
| 385 | default_vllm_config.model_config = ModelConfig() |
| 386 | prefix = "model.layers.0.self_attn.o_proj" |
| 387 | quant_config = _fully_quantized_quark_config() |
| 388 | quant_config.online_quantization_config = OnlineQuantizationConfig( |
| 389 | QuantizationConfigArgs(targets={prefix: "mxfp8"}) |
| 390 | ) |
| 391 | |
| 392 | with pytest.raises(ValueError, match="pre-quantized layer"): |
| 393 | ColumnParallelLinear( |
| 394 | input_size=32, |
| 395 | output_size=32, |
| 396 | bias=False, |
| 397 | params_dtype=torch.bfloat16, |
| 398 | quant_config=quant_config, |
| 399 | prefix=prefix, |
| 400 | disable_tp=True, |
| 401 | ) |
| 402 | |
| 403 | |
| 404 | def test_online_ignore_keeps_checkpoint_quantization_linear( |
| 405 | default_vllm_config, dist_init, monkeypatch |
| 406 | ): |
| 407 | """Ignoring online quantization does not replace a checkpoint method (linear).""" |
| 408 | default_vllm_config.model_config = ModelConfig() |
| 409 | quant_config = _fully_quantized_quark_config() |
| 410 | prefix = "model.layers.0.self_attn.o_proj" |
| 411 | quant_config.online_quantization_config = OnlineQuantizationConfig( |
| 412 | quant_config_args(linear="mxfp8", ignore=[prefix]) |
| 413 | ) |
| 414 | monkeypatch.setattr( |
| 415 | quant_config.online_quantization_config, |
| 416 | "get_quant_method", |
| 417 | lambda *args: pytest.fail("online method should not be constructed"), |
| 418 | ) |
| 419 | |
| 420 | layer = ColumnParallelLinear( |
| 421 | input_size=32, |
| 422 | output_size=32, |
| 423 | bias=False, |
| 424 | params_dtype=torch.bfloat16, |
| 425 | quant_config=quant_config, |
| 426 | prefix=prefix, |
| 427 | disable_tp=True, |
| 428 | ) |
| 429 | |
| 430 | assert isinstance(layer.quant_method, QuarkLinearMethod) |
| 431 | |
| 432 | |
| 433 | def test_online_quantization_rejects_prequantized_moe( |
| 434 | default_vllm_config, dist_init |
| 435 | ) -> None: |
| 436 | """Online linear and MoE quantization reject a pre-quantized MoE layer.""" |
| 437 | default_vllm_config.model_config = ModelConfig() |
| 438 | prefix = "model.layers.0.mlp.experts" |
| 439 | quant_config = _fully_quantized_quark_config() |
| 440 | quant_config.online_quantization_config = OnlineQuantizationConfig( |
| 441 | quant_config_args(linear="mxfp4", moe="mxfp4") |
| 442 | ) |
| 443 | |
| 444 | with pytest.raises(ValueError, match="pre-quantized layer"): |
| 445 | FusedMoEFactory( |
| 446 | num_experts=4, |
| 447 | top_k=2, |
| 448 | hidden_size=32, |
| 449 | intermediate_size=64, |
| 450 | params_dtype=torch.bfloat16, |
| 451 | quant_config=quant_config, |
| 452 | prefix=prefix, |
| 453 | ) |
| 454 | |
| 455 | |
| 456 | def test_activation_only_override_applies_to_checkpoint_method( |
| 457 | default_vllm_config, tmp_path |
| 458 | ) -> None: |
| 459 | """Activation-only overrides remain available to the checkpoint method.""" |
| 460 | _write_minimal_llama_config(tmp_path) |
| 461 | model_config = ModelConfig( |
| 462 | model=str(tmp_path), |
| 463 | quantization="compressed-tensors", |
| 464 | quantization_config=quant_config_args(moe={"activation": "mxfp8"}), |
| 465 | hf_overrides={ |
| 466 | "quantization_config": { |
| 467 | "quant_method": "compressed-tensors", |
| 468 | "format": "pack-quantized", |
| 469 | "config_groups": {}, |
| 470 | } |
| 471 | }, |
| 472 | ) |
| 473 | |
| 474 | result = weight_utils.get_quant_config(model_config, LoadConfig()) |
| 475 | |
| 476 | assert isinstance(result, CompressedTensorsConfig) |
| 477 | assert isinstance(result.online_quantization_config, OnlineQuantizationConfig) |
| 478 | |
| 479 | default_vllm_config.model_config = model_config |
| 480 | from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( |
| 481 | _resolve_activation_key, |
| 482 | ) |
| 483 | |
| 484 | assert _resolve_activation_key(None) == kMxfp8Dynamic |
| 485 | |
| 486 | |
| 487 | def test_online_overlay_requires_normalized_quantization_config(tmp_path) -> None: |
| 488 | """The loader rejects a raw config dict instead of silently ignoring it.""" |
| 489 | _write_minimal_llama_config(tmp_path) |
| 490 | model_config = ModelConfig( |
| 491 | model=str(tmp_path), |
| 492 | quantization="compressed-tensors", |
| 493 | quantization_config={"linear": "mxfp8"}, |
| 494 | hf_overrides={ |
| 495 | "quantization_config": { |
| 496 | "quant_method": "compressed-tensors", |
| 497 | "format": "pack-quantized", |
| 498 | "config_groups": {}, |
| 499 | } |
| 500 | }, |
| 501 | ) |
| 502 | |
| 503 | # assert isinstance(online_args, QuantizationConfigArgs) |
| 504 | # in maybe_compose_online_quantization |
| 505 | with pytest.raises(AssertionError): |
| 506 | weight_utils.get_quant_config(model_config, LoadConfig()) |
| 507 | |
| 508 | |
| 509 | def test_online_overlay_loads_checkpoint_config_file(tmp_path) -> None: |
| 510 | """An online overlay must not bypass checkpoint config file loading.""" |
| 511 | _write_minimal_llama_config(tmp_path) |
| 512 | (tmp_path / "quantize_config.json").write_text( |
| 513 | '{"w_bit": 4, "q_group_size": 128, "zero_point": true}' |
| 514 | ) |
| 515 | model_config = ModelConfig( |
| 516 | model=str(tmp_path), |
| 517 | quantization="awq", |
| 518 | quantization_config=quant_config_args(linear="mxfp8"), |
| 519 | ) |
| 520 | |
| 521 | result = weight_utils.get_quant_config(model_config, LoadConfig()) |
| 522 | |
| 523 | assert isinstance(result, AutoAWQConfig) |
| 524 | assert result.weight_bits == 4 |
| 525 | assert result.group_size == 128 |
| 526 | assert result.zero_point |
| 527 | assert isinstance(result.online_quantization_config, OnlineQuantizationConfig) |
| 528 | assert result.online_quantization_config.args is model_config.quantization_config |
| 529 | |
| 530 | |
| 531 | @pytest.mark.parametrize( |
| 532 | "quantization,checkpoint_quantization_config,is_config_file", |
| 533 | [ |
| 534 | pytest.param("fp8_per_block", None, False, id="fp8-per-block-online"), |
| 535 | pytest.param("mxfp4", None, False, id="mxfp4-online"), |
| 536 | pytest.param("mxfp8", None, False, id="mxfp8-online"), |
| 537 | pytest.param( |
| 538 | "mxfp8", |
| 539 | {"quant_method": "mxfp8", "ignored_layers": []}, |
| 540 | False, |
| 541 | id="mxfp8-inline-checkpoint-config", |
| 542 | ), |
| 543 | pytest.param( |
| 544 | "mxfp8", |
| 545 | {"quant_method": "mxfp8", "ignored_layers": []}, |
| 546 | True, |
| 547 | id="mxfp8-checkpoint-config-file", |
| 548 | ), |
| 549 | ], |
| 550 | ) |
| 551 | def test_online_shorthand_selects_checkpoint_or_online_config( |
| 552 | tmp_path, quantization, checkpoint_quantization_config, is_config_file |
| 553 | ) -> None: |
| 554 | """Online shorthands defer to checkpoint metadata when available.""" |
| 555 | if is_config_file: |
| 556 | (tmp_path / "hf_quant_config.json").write_text( |
| 557 | '{"quant_method": "mxfp8", "ignored_layers": []}' |
| 558 | ) |
| 559 | |
| 560 | _write_minimal_llama_config( |
| 561 | tmp_path, |
| 562 | None if is_config_file else checkpoint_quantization_config, |
| 563 | ) |
| 564 | model_config = ModelConfig(model=str(tmp_path), quantization=quantization) |
| 565 | |
| 566 | result = weight_utils.get_quant_config(model_config, LoadConfig()) |
| 567 | |
| 568 | if checkpoint_quantization_config is not None: |
| 569 | assert isinstance(result, ModelOptMxFp8Config) |
| 570 | assert result.online_quantization_config is None |
| 571 | else: |
| 572 | assert isinstance(result, OnlineQuantizationConfig) |
| 573 | assert result.online_quantization_config is None |
| 574 | |
| 575 | |
| 576 | def test_log_online_quantization_for_composable_config(monkeypatch) -> None: |
| 577 | """Composable configs log their nested online quantization results.""" |
| 578 | online_config = OnlineQuantizationConfig(quant_config_args(linear="mxfp8")) |
| 579 | online_config.quantized_layers = { |
| 580 | "model.layers.0.self_attn.o_proj": ("linear", "mxfp8", None), |
| 581 | "model.layers.1.self_attn.o_proj": ("linear", "mxfp8", None), |
| 582 | } |
| 583 | vllm_config = VllmConfig(quant_config=online_config) |
| 584 | log_args = [] |
| 585 | monkeypatch.setattr( |
| 586 | "vllm.model_executor.model_loader.base_loader.logger.info", |
| 587 | lambda *args: log_args.append(args), |
| 588 | ) |
| 589 | |
| 590 | log_online_quantization(vllm_config) |
| 591 | |
| 592 | assert log_args == [ |
| 593 | ( |
| 594 | "Quantized %d layers of types: %s", |
| 595 | 2, |
| 596 | "; ".join(online_config.quantized_layer_summaries), |
| 597 | ) |
| 598 | ] |
| 599 | |
| 600 | |
| 601 | def test_checkpoint_quantization_rejects_online_shorthand(tmp_path) -> None: |
| 602 | """A checkpoint quant method cannot be replaced by an online shorthand.""" |
| 603 | _write_minimal_llama_config(tmp_path, {"quant_method": "mxfp4"}) |
| 604 | |
| 605 | with pytest.raises(ValueError, match="does not match the quantization"): |
| 606 | ModelConfig(model=str(tmp_path), quantization="fp8_per_channel") |
| 607 | |
| 608 | |
| 609 | def test_nvfp4_per_token_backend_contract() -> None: |
| 610 | from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutedsl_moe import ( |
| 611 | FlashInferCuteDSLExperts, |
| 612 | ) |
| 613 | from vllm.model_executor.layers.fused_moe.experts.flashinfer_cutlass_moe import ( |
| 614 | FlashInferExperts, |
| 615 | ) |
| 616 | from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( |
| 617 | TrtLlmNvFp4ExpertsModular, |
| 618 | TrtLlmNvFp4ExpertsMonolithic, |
| 619 | ) |
| 620 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 621 | kNvfp4DynamicToken, |
| 622 | kNvfp4Static, |
| 623 | ) |
| 624 | |
| 625 | scheme = (kNvfp4Static, kNvfp4DynamicToken) |
| 626 | assert TrtLlmNvFp4ExpertsMonolithic._supports_quant_scheme(*scheme) |
| 627 | assert TrtLlmNvFp4ExpertsModular._supports_quant_scheme(*scheme) |
| 628 | assert FlashInferCuteDSLExperts._supports_quant_scheme(*scheme) |
| 629 | assert FlashInferCuteDSLExperts._supports_no_act_and_mul() |
| 630 | assert not FlashInferExperts._supports_quant_scheme(*scheme) |
| 631 | |
| 632 | |
| 633 | @pytest.mark.parametrize("per_token_activation", [False, True]) |
| 634 | @pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16]) |
| 635 | def test_nvfp4_one_sided_sizes_dispatched_activations( |
| 636 | monkeypatch, per_token_activation: bool, input_dtype: torch.dtype |
| 637 | ) -> None: |
| 638 | from tests.kernels.moe.utils import make_dummy_moe_config |
| 639 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 640 | flashinfer_one_sided_dispatch_layout, |
| 641 | ) |
| 642 | from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig |
| 643 | from vllm.model_executor.layers.fused_moe.experts.trtllm_nvfp4_moe import ( |
| 644 | TrtLlmNvFp4ExpertsModular, |
| 645 | ) |
| 646 | |
| 647 | config = make_dummy_moe_config(hidden_dim=6144, in_dtype=input_dtype) |
| 648 | config.moe_parallel_config.dp_size = 2 |
| 649 | config.moe_parallel_config.ep_size = 2 |
| 650 | config.moe_parallel_config.use_ep = True |
| 651 | config.moe_parallel_config.all2all_backend = "flashinfer_nvlink_one_sided" |
| 652 | expert_quant_config = FusedMoEQuantConfig.make("nvfp4") |
| 653 | |
| 654 | def check_dispatch_layout(*, quant_config, input_dtype, **kwargs): |
| 655 | assert quant_config is expert_quant_config |
| 656 | layout = flashinfer_one_sided_dispatch_layout( |
| 657 | config.hidden_dim, quant_config, input_dtype=input_dtype |
| 658 | ) |
| 659 | if per_token_activation: |
| 660 | assert input_dtype == config.in_dtype |
| 661 | assert ( |
| 662 | layout.x_bytes_per_token == config.hidden_dim * config.in_dtype.itemsize |
| 663 | ) |
| 664 | assert layout.x_sf_bytes_per_token == 0 |
| 665 | else: |
| 666 | assert input_dtype is None |
| 667 | assert layout.x_bytes_per_token == config.hidden_dim // 2 |
| 668 | assert layout.x_sf_bytes_per_token == config.hidden_dim // 16 |
| 669 | assert expert_quant_config.quant_dtype == "nvfp4" |
| 670 | raise RuntimeError("dispatch layout verified before allocation") |
| 671 | |
| 672 | monkeypatch.setattr( |
| 673 | nvfp4_oracle, "maybe_make_prepare_finalize", check_dispatch_layout |
| 674 | ) |
| 675 | with pytest.raises( |
| 676 | RuntimeError, match="dispatch layout verified before allocation" |
| 677 | ): |
| 678 | nvfp4_oracle.make_nvfp4_moe_kernel( |
| 679 | moe_quant_config=expert_quant_config, |
| 680 | moe_config=config, |
| 681 | experts_cls=TrtLlmNvFp4ExpertsModular, |
| 682 | backend=nvfp4_oracle.NvFp4MoeBackend.FLASHINFER_TRTLLM, |
| 683 | per_token_activation=per_token_activation, |
| 684 | ) |
| 685 | |
| 686 | |
| 687 | @pytest.mark.parametrize( |
| 688 | "input_dtype", |
| 689 | [torch.float32, torch.float8_e4m3fn, torch.float8_e5m2, torch.int8], |
| 690 | ) |
| 691 | def test_nvfp4_one_sided_rejects_unsupported_input_dtype(input_dtype: torch.dtype): |
| 692 | from vllm.model_executor.layers.fused_moe.all2all_utils import ( |
| 693 | flashinfer_one_sided_dispatch_layout, |
| 694 | ) |
| 695 | from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig |
| 696 | |
| 697 | with pytest.raises(ValueError, match="unpacked inputs must be float16 or bfloat16"): |
| 698 | flashinfer_one_sided_dispatch_layout( |
| 699 | 6144, FusedMoEQuantConfig.make("nvfp4"), input_dtype=input_dtype |
| 700 | ) |
| 701 | |
| 702 | |
| 703 | @pytest.mark.skipif( |
| 704 | not is_quant_method_supported("fp8"), |
| 705 | reason="FP8 is not supported on this GPU type.", |
| 706 | ) |
| 707 | @pytest.mark.parametrize( |
| 708 | ( |
| 709 | "model_name,quant_scheme,online_quant_args,expected_linear_cls,expected_moe_cls,extra_runner_kwargs" |
| 710 | ), |
| 711 | [ |
| 712 | # simple case - quantization='fp8_per_tensor' |
| 713 | pytest.param( |
| 714 | GRANITE_MODEL_NAME, |
| 715 | "fp8_per_tensor", |
| 716 | None, |
| 717 | Fp8PerTensorOnlineLinearMethod, |
| 718 | Fp8PerTensorOnlineMoEMethod, |
| 719 | {}, |
| 720 | id="fp8_per_tensor", |
| 721 | ), |
| 722 | # FP8 KV cache with the fp8_per_tensor shorthand. |
| 723 | pytest.param( |
| 724 | GRANITE_MODEL_NAME, |
| 725 | "fp8_per_tensor", |
| 726 | None, |
| 727 | Fp8PerTensorOnlineLinearMethod, |
| 728 | Fp8PerTensorOnlineMoEMethod, |
| 729 | {"kv_cache_dtype": "fp8"}, |
| 730 | id="fp8_per_tensor-fp8-kv-cache", |
| 731 | ), |
| 732 | pytest.param( |
| 733 | GRANITE_MODEL_NAME, |
| 734 | "fp8_per_tensor", |
| 735 | None, |
| 736 | Fp8PerTensorOnlineLinearMethod, |
| 737 | Fp8PerTensorOnlineMoEMethod, |
| 738 | {"linear_backend": "marlin", "moe_backend": "marlin"}, |
| 739 | marks=pytest.mark.skipif( |
| 740 | not current_platform.is_cuda(), |
| 741 | reason="Marlin online FP8 coverage is CUDA-only.", |
| 742 | ), |
| 743 | id="fp8_per_tensor-marlin", |
| 744 | ), |
| 745 | pytest.param( |
| 746 | GRANITE_MODEL_NAME, |
| 747 | "fp8_per_tensor", |
| 748 | None, |
| 749 | Fp8PerTensorOnlineLinearMethod, |
| 750 | Fp8PerTensorOnlineMoEMethod, |
| 751 | { |
| 752 | "kv_cache_dtype": "fp8", |
| 753 | "linear_backend": "marlin", |
| 754 | "moe_backend": "marlin", |
| 755 | }, |
| 756 | marks=pytest.mark.skipif( |
| 757 | not current_platform.is_cuda(), |
| 758 | reason="Marlin online FP8 coverage is CUDA-only.", |
| 759 | ), |
| 760 | id="fp8_per_tensor-fp8-kv-cache-marlin", |
| 761 | ), |
| 762 | # simple case - quantization='fp8_per_block' |
| 763 | pytest.param( |
| 764 | GRANITE_MODEL_NAME, |
| 765 | "fp8_per_block", |
| 766 | None, |
| 767 | Fp8PerBlockOnlineLinearMethod, |
| 768 | Fp8PerBlockOnlineMoEMethod, |
| 769 | {}, |
| 770 | id="fp8_per_block", |
| 771 | ), |
| 772 | pytest.param( |
| 773 | GRANITE_MODEL_NAME, |
| 774 | "fp8_per_channel", |
| 775 | None, |
| 776 | Fp8PtpcOnlineLinearMethod, |
| 777 | Fp8PtpcOnlineMoEMethod, |
| 778 | {}, |
| 779 | id="fp8_per_channel", |
| 780 | ), |
| 781 | # quantization='online' with per-layer-kind overrides |
| 782 | pytest.param( |
| 783 | GRANITE_MODEL_NAME, |
| 784 | "online", |
| 785 | { |
| 786 | "linear": "fp8_per_block", |
| 787 | "moe": "fp8_per_tensor", |
| 788 | }, |
| 789 | Fp8PerBlockOnlineLinearMethod, |
| 790 | Fp8PerTensorOnlineMoEMethod, |
| 791 | {}, |
| 792 | id="per_layer_kind_overrides", |
| 793 | ), |
| 794 | # quantization='online' with per-layer target patterns |
| 795 | pytest.param( |
| 796 | GRANITE_MODEL_NAME, |
| 797 | "online", |
| 798 | { |
| 799 | "targets": { |
| 800 | r"re:.*self_attn\.o_proj": "fp8_per_block", |
| 801 | r"re:.*block_sparse_moe\.experts": "fp8_per_tensor", |
| 802 | } |
| 803 | }, |
| 804 | Fp8PerBlockOnlineLinearMethod, |
| 805 | Fp8PerTensorOnlineMoEMethod, |
| 806 | {}, |
| 807 | id="targets", |
| 808 | ), |
| 809 | # ignore with direct layer name |
| 810 | pytest.param( |
| 811 | GRANITE_MODEL_NAME, |
| 812 | "fp8_per_tensor", |
| 813 | # qkv_proj is fused from q_proj/k_proj/v_proj. The shard regex |
| 814 | # remains supported alongside direct fused-name regexes. |
| 815 | {"ignore": ["model.layers.1.self_attn.o_proj", "re:.*[qkv]_proj"]}, |
| 816 | Fp8PerTensorOnlineLinearMethod, |
| 817 | Fp8PerTensorOnlineMoEMethod, |
| 818 | {}, |
| 819 | id="ignore", |
| 820 | ), |
| 821 | pytest.param( |
| 822 | GRANITE_MODEL_NAME, |
| 823 | "mxfp4", |
| 824 | None, |
| 825 | Mxfp4OnlineLinearMethod, |
| 826 | Mxfp4OnlineMoEMethod, |
| 827 | {}, |
| 828 | id="mxfp4", |
| 829 | ), |
| 830 | pytest.param( |
| 831 | PARTIALLY_PREQUANTIZED_MODEL_NAME, |
| 832 | None, |
| 833 | {"targets": {"model.layers.1.self_attn.o_proj": "mxfp8"}}, |
| 834 | Mxfp8OnlineLinearMethod, |
| 835 | CompressedTensorsMoEMethod, |
| 836 | {}, |
| 837 | id="partially_prequantized_checkpoint", |
| 838 | ), |
| 839 | ], |
| 840 | ) |
| 841 | @pytest.mark.parametrize( |
| 842 | "use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False] |
| 843 | ) |
| 844 | def test_online_quantization( |
| 845 | model_name: str, |
| 846 | quant_scheme: str | None, |
| 847 | online_quant_args: dict | None, |
| 848 | expected_linear_cls, |
| 849 | expected_moe_cls, |
| 850 | extra_runner_kwargs: dict, |
| 851 | use_rocm_aiter: bool, |
| 852 | monkeypatch, |
| 853 | dist_init, |
| 854 | workspace_init, |
| 855 | ) -> None: |
| 856 | """Tests that online quantization frontend configuration works - |
| 857 | selecting quant schemes, overriding quant schemes by type, ignoring |
| 858 | layers. |
| 859 | |
| 860 | Does not test performance, peak memory usage, etc. |
| 861 | """ |
| 862 | # TODO: Relax this condition once there is a native MXFP4_MXFP4 |
| 863 | # linear/moe backend supported on cuda. |
| 864 | if quant_scheme == "mxfp4" and not (on_gfx950() or on_gfx942()): |
| 865 | pytest.skip("mxfp4 online quantization is only tested on AMD gfx942, gfx950.") |
| 866 | |
| 867 | if current_platform.is_rocm(): |
| 868 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1" if use_rocm_aiter else "0") |
| 869 | rocm_aiter_ops.refresh_env_variables() |
| 870 | |
| 871 | if current_platform.is_xpu() and quant_scheme in ( |
| 872 | "fp8_per_block", |
| 873 | "fp8_per_channel", |
| 874 | ): |
| 875 | pytest.skip(f"Skip test for online {quant_scheme} on XPU platform.") |
| 876 | |
| 877 | runner_kwargs = dict(extra_runner_kwargs) |
| 878 | kv_cache_dtype = runner_kwargs.get("kv_cache_dtype", "auto") |
| 879 | force_marlin = extra_runner_kwargs.get("linear_backend") == "marlin" |
| 880 | if kv_cache_dtype == "fp8" and current_platform.is_device_capability_family(90): |
| 881 | # FA3 requires BF16 output when the query input is FP8. |
| 882 | runner_kwargs["dtype"] = "bfloat16" |
| 883 | |
| 884 | vllm_config_kwargs = { |
| 885 | "cache_config": CacheConfig(cache_dtype=kv_cache_dtype), |
| 886 | "kernel_config": KernelConfig( |
| 887 | linear_backend=runner_kwargs.get("linear_backend", "auto"), |
| 888 | moe_backend=runner_kwargs.get("moe_backend", "auto"), |
| 889 | ), |
| 890 | } |
| 891 | |
| 892 | if model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME: |
| 893 | model, vllm_config = load_model_without_vllm_runner( |
| 894 | model_name, |
| 895 | dtype=runner_kwargs.get("dtype", "bfloat16"), |
| 896 | model_config_kwargs={ |
| 897 | "quantization_config": resolve_quantization_config( |
| 898 | None, online_quant_args |
| 899 | ) |
| 900 | }, |
| 901 | vllm_config_kwargs=vllm_config_kwargs, |
| 902 | model_loader_cls=DummyModelLoader, |
| 903 | ) |
| 904 | assert isinstance( |
| 905 | model.model.layers[1].self_attn.o_proj.quant_method, |
| 906 | expected_linear_cls, |
| 907 | ) |
| 908 | assert isinstance( |
| 909 | model.model.layers[0].mlp.experts._quant_method, |
| 910 | expected_moe_cls, |
| 911 | ) |
| 912 | assert isinstance( |
| 913 | model.model.layers[0].self_attn.o_proj.quant_method, |
| 914 | CompressedTensorsLinearMethod, |
| 915 | ) |
| 916 | else: |
| 917 | model, vllm_config = load_model_without_vllm_runner( |
| 918 | model_name, |
| 919 | dtype=runner_kwargs.get("dtype", "bfloat16"), |
| 920 | quantization=quant_scheme, |
| 921 | model_config_kwargs={ |
| 922 | "quantization_config": resolve_quantization_config( |
| 923 | quant_scheme, online_quant_args |
| 924 | ), |
| 925 | "hf_overrides": { |
| 926 | "num_hidden_layers": 3, |
| 927 | "vocab_size": 256, |
| 928 | "hidden_size": 256, |
| 929 | "intermediate_size": 512, |
| 930 | "num_attention_heads": 4, |
| 931 | "num_key_value_heads": 4, |
| 932 | "max_position_embeddings": 64, |
| 933 | "num_local_experts": 4, |
| 934 | "num_experts_per_tok": 2, |
| 935 | }, |
| 936 | }, |
| 937 | vllm_config_kwargs=vllm_config_kwargs, |
| 938 | model_loader_cls=DummyModelLoader, |
| 939 | ) |
| 940 | |
| 941 | monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous()) |
| 942 | |
| 943 | if model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME: |
| 944 | o_proj = model.model.layers[1].self_attn.o_proj |
| 945 | moe = model.model.layers[0].mlp.experts |
| 946 | else: |
| 947 | o_proj = model.model.layers[0].self_attn.o_proj |
| 948 | moe = model.model.layers[0].block_sparse_moe.experts |
| 949 | assert isinstance(o_proj.quant_method, expected_linear_cls) |
| 950 | assert isinstance(moe._quant_method, expected_moe_cls) |
| 951 | |
| 952 | if model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME and isinstance( |
| 953 | o_proj.quant_method.kernel, MarlinMxfp8LinearKernel |
| 954 | ): |
| 955 | assert o_proj.weight.dtype == torch.int32 |
| 956 | elif model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME and isinstance( |
| 957 | o_proj.quant_method.kernel, EmulationMxfp8LinearKernel |
| 958 | ): |
| 959 | assert o_proj.weight.dtype == torch.bfloat16 |
| 960 | elif model_name == PARTIALLY_PREQUANTIZED_MODEL_NAME: |
| 961 | assert o_proj.weight.dtype == MXFP8_VALUE_DTYPE |
| 962 | elif quant_scheme == "mxfp4": |
| 963 | assert o_proj.weight.dtype == torch.uint8 |
| 964 | elif current_platform.is_cuda() or current_platform.is_xpu(): |
| 965 | if current_platform.supports_fp8() and not force_marlin: |
| 966 | assert o_proj.weight.dtype == torch.float8_e4m3fn |
| 967 | assert not isinstance( |
| 968 | o_proj.quant_method.fp8_linear, MarlinFP8ScaledMMLinearKernel |
| 969 | ) |
| 970 | else: |
| 971 | assert o_proj.weight.dtype == torch.int32 |
| 972 | assert isinstance( |
| 973 | o_proj.quant_method.fp8_linear, MarlinFP8ScaledMMLinearKernel |
| 974 | ) |
| 975 | elif current_platform.is_rocm(): |
| 976 | assert o_proj.weight.dtype == current_platform.fp8_dtype() |
| 977 | else: |
| 978 | pytest.skip("Only runs on CUDA and ROCm.") |
| 979 | |
| 980 | if quant_scheme == "fp8_per_channel": |
| 981 | assert o_proj.weight_scale.ndim == 2 |
| 982 | assert o_proj.weight_scale.shape[-1] == 1 |
| 983 | assert o_proj.input_scale is None |
| 984 | |
| 985 | if isinstance(online_quant_args, dict) and "ignore" in online_quant_args: |
| 986 | for layer_idx in range(len(model.model.layers)): |
| 987 | o_proj = model.model.layers[layer_idx].self_attn.o_proj |
| 988 | if layer_idx == 1: |
| 989 | assert isinstance(o_proj.quant_method, UnquantizedLinearMethod) |
| 990 | else: |
| 991 | assert isinstance(o_proj.quant_method, expected_linear_cls) |
| 992 | |
| 993 | for layer in model.model.layers: |
| 994 | assert isinstance( |
| 995 | layer.self_attn.qkv_proj.quant_method, UnquantizedLinearMethod |
| 996 | ) |
| 997 | |
| 998 | if kv_cache_dtype == "fp8": |
| 999 | attn = model.model.layers[0].self_attn.attn |
| 1000 | assert attn._k_scale == 1.0 |
| 1001 | assert attn._v_scale == 1.0 |
| 1002 | |
| 1003 | if ( |
| 1004 | model_name == GRANITE_MODEL_NAME |
| 1005 | and isinstance(online_quant_args, dict) |
| 1006 | and "targets" in online_quant_args |
| 1007 | ): |
| 1008 | for layer in model.model.layers: |
| 1009 | assert isinstance( |
| 1010 | layer.self_attn.qkv_proj.quant_method, UnquantizedLinearMethod |
| 1011 | ) |
| 1012 | |
| 1013 | input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE) |
| 1014 | positions = torch.arange(input_ids.numel(), device=DEVICE) |
| 1015 | with set_forward_context(None, vllm_config, num_tokens=input_ids.numel()): |
| 1016 | hidden_states = model(input_ids, positions, None) |
| 1017 | logits = model.compute_logits(hidden_states) |
| 1018 | assert torch.isfinite(logits).all() |
| 1019 | |
| 1020 | |
| 1021 | @pytest.mark.skipif( |
| 1022 | not is_quant_method_supported("fp8"), |
| 1023 | reason="FP8 is not supported on this GPU type.", |
| 1024 | ) |
| 1025 | @pytest.mark.parametrize( |
| 1026 | "model_name", [GRANITE_MODEL_NAME, PARTIALLY_PREQUANTIZED_MODEL_NAME] |
| 1027 | ) |
| 1028 | def test_online_quantization_loads_real_weights( |
| 1029 | model_name, vllm_runner, monkeypatch |
| 1030 | ) -> None: |
| 1031 | """Verify online quantization loads real checkpoints end to end.""" |
| 1032 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 1033 | monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") |
| 1034 | |
| 1035 | original_load_weights = GraniteMoeModel.load_weights |
| 1036 | |
| 1037 | def load_weights(self, weights): |
| 1038 | weights = ( |
| 1039 | (name, weight) |
| 1040 | for name, weight in weights |
| 1041 | if not name.startswith("layers.") or int(name.split(".")[1]) < 3 |
| 1042 | ) |
| 1043 | return original_load_weights(self, weights) |
| 1044 | |
| 1045 | monkeypatch.setattr(GraniteMoeModel, "load_weights", load_weights) |
| 1046 | |
| 1047 | if model_name == GRANITE_MODEL_NAME: |
| 1048 | runner_kwargs = { |
| 1049 | "quantization": "fp8_per_tensor", |
| 1050 | "dtype": "bfloat16", |
| 1051 | "enforce_eager": True, |
| 1052 | "hf_overrides": {"num_hidden_layers": 3}, |
| 1053 | "max_model_len": 16, |
| 1054 | "max_num_seqs": 1, |
| 1055 | } |
| 1056 | else: |
| 1057 | runner_kwargs = { |
| 1058 | "quantization_config": { |
| 1059 | "targets": {"model.layers.1.self_attn.o_proj": "mxfp8"} |
| 1060 | }, |
| 1061 | "enforce_eager": True, |
| 1062 | } |
| 1063 | |
| 1064 | with vllm_runner(model_name, **runner_kwargs) as llm: |
| 1065 | if model_name == GRANITE_MODEL_NAME: |
| 1066 | |
| 1067 | def check_model(model): |
| 1068 | layer = model.model.layers[0] |
| 1069 | assert isinstance( |
| 1070 | layer.self_attn.o_proj.quant_method, |
| 1071 | Fp8PerTensorOnlineLinearMethod, |
| 1072 | ) |
| 1073 | assert isinstance( |
| 1074 | layer.block_sparse_moe.experts._quant_method, |
| 1075 | Fp8PerTensorOnlineMoEMethod, |
| 1076 | ) |
| 1077 | |
| 1078 | else: |
| 1079 | |
| 1080 | def check_model(model): |
| 1081 | assert isinstance( |
| 1082 | model.model.layers[1].self_attn.o_proj.quant_method, |
| 1083 | Mxfp8OnlineLinearMethod, |
| 1084 | ) |
| 1085 | assert isinstance( |
| 1086 | model.model.layers[0].self_attn.o_proj.quant_method, |
| 1087 | CompressedTensorsLinearMethod, |
| 1088 | ) |
| 1089 | assert isinstance( |
| 1090 | model.model.layers[0].mlp.experts._quant_method, |
| 1091 | CompressedTensorsMoEMethod, |
| 1092 | ) |
| 1093 | |
| 1094 | llm.apply_model(check_model) |
| 1095 | outputs = llm.generate_greedy(["Hello my name is"], max_tokens=1) |
| 1096 | assert outputs |
| 1097 | |
| 1098 | |
| 1099 | @pytest.mark.skipif( |
| 1100 | not is_quant_method_supported("fp8"), |
| 1101 | reason="FP8 is not supported on this GPU type.", |
| 1102 | ) |
| 1103 | @pytest.mark.parametrize( |
| 1104 | "targets,prefix,expected_method_cls,unmatched_prefix,expected_metadata", |
| 1105 | [ |
| 1106 | ( |
| 1107 | {r"re:.*self_attn\.o_proj": "fp8_per_block"}, |
| 1108 | "model.layers.0.self_attn.o_proj", |
| 1109 | Fp8PerBlockOnlineLinearMethod, |
| 1110 | "model.layers.0.self_attn.qkv_proj", |
| 1111 | ("targets", "fp8_per_block", r"re:.*self_attn\.o_proj"), |
| 1112 | ), |
| 1113 | ( |
| 1114 | {r"re:.*qkv_proj.*": "fp8_per_tensor"}, |
| 1115 | "model.layers.0.self_attn.qkv_proj", |
| 1116 | Fp8PerTensorOnlineLinearMethod, |
| 1117 | "model.layers.0.self_attn.o_proj", |
| 1118 | ("targets", "fp8_per_tensor", r"re:.*qkv_proj.*"), |
| 1119 | ), |
| 1120 | ( |
| 1121 | {r"re:.*[qkv]_proj": "fp8_per_tensor"}, |
| 1122 | "model.layers.0.self_attn.qkv_proj", |
| 1123 | Fp8PerTensorOnlineLinearMethod, |
| 1124 | "model.layers.0.self_attn.o_proj", |
| 1125 | ("targets", "fp8_per_tensor", r"re:.*[qkv]_proj"), |
| 1126 | ), |
| 1127 | ], |
| 1128 | ids=["linear_regex", "direct_fused_regex", "legacy_fused_regex"], |
| 1129 | ) |
| 1130 | def test_online_quantization_targets( |
| 1131 | default_vllm_config, |
| 1132 | dist_init, |
| 1133 | targets: dict[str, str], |
| 1134 | prefix: str, |
| 1135 | expected_method_cls, |
| 1136 | unmatched_prefix: str, |
| 1137 | expected_metadata: tuple[str, str, str], |
| 1138 | ) -> None: |
| 1139 | """Target patterns select the real online linear methods.""" |
| 1140 | default_vllm_config.model_config = ModelConfig() |
| 1141 | config = OnlineQuantizationConfig(QuantizationConfigArgs(targets=targets)) |
| 1142 | config.packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]} |
| 1143 | |
| 1144 | layer = ColumnParallelLinear( |
| 1145 | input_size=1, |
| 1146 | output_size=1, |
| 1147 | bias=False, |
| 1148 | disable_tp=True, |
| 1149 | ) |
| 1150 | |
| 1151 | method = config.get_quant_method(layer, prefix) |
| 1152 | assert isinstance(method, expected_method_cls) |
| 1153 | assert config.quantized_layers == {prefix: expected_metadata} |
| 1154 | |
| 1155 | unmatched_method = config.get_quant_method(layer, unmatched_prefix) |
| 1156 | assert isinstance(unmatched_method, UnquantizedLinearMethod) |
| 1157 | assert config.quantized_layers == {prefix: expected_metadata} |
| 1158 | |
| 1159 | |
| 1160 | @pytest.mark.skipif( |
| 1161 | not is_quant_method_supported("fp8"), |
| 1162 | reason="FP8 is not supported on this GPU type.", |
| 1163 | ) |
| 1164 | def test_online_quantization_records_global_config( |
| 1165 | default_vllm_config, dist_init |
| 1166 | ) -> None: |
| 1167 | default_vllm_config.model_config = ModelConfig() |
| 1168 | config = OnlineQuantizationConfig(quant_config_args(linear="fp8_per_block")) |
| 1169 | prefix = "model.layers.0.self_attn.o_proj" |
| 1170 | layer = ColumnParallelLinear( |
| 1171 | input_size=1, |
| 1172 | output_size=1, |
| 1173 | bias=False, |
| 1174 | disable_tp=True, |
| 1175 | ) |
| 1176 | |
| 1177 | method = config.get_quant_method(layer, prefix) |
| 1178 | |
| 1179 | assert isinstance(method, Fp8PerBlockOnlineLinearMethod) |
| 1180 | assert config.quantized_layers == { |
| 1181 | prefix: ("linear", str(config.args.linear), None) |
| 1182 | } |
| 1183 | |
| 1184 | |
| 1185 | def test_online_quantization_targets_ignore_collision() -> None: |
| 1186 | """A targets/ignore collision is reported when the layer is dispatched.""" |
| 1187 | config = OnlineQuantizationConfig( |
| 1188 | QuantizationConfigArgs( |
| 1189 | targets={"model.layers.0.self_attn.o_proj": "fp8_per_tensor"}, |
| 1190 | ignore=["model.layers.0.self_attn.o_proj"], |
| 1191 | ) |
| 1192 | ) |
| 1193 | with pytest.raises(ValueError, match="matches both quantization_config.ignore"): |
| 1194 | config.resolve_quant_method_cls( |
| 1195 | Mock(spec=LinearBase), "model.layers.0.self_attn.o_proj" |
| 1196 | ) |
| 1197 | |
| 1198 | |
| 1199 | def test_online_quantization_targets_reject_unsupported_layer() -> None: |
| 1200 | """A targets match on a non-linear, non-MoE layer is rejected.""" |
| 1201 | config = OnlineQuantizationConfig( |
| 1202 | QuantizationConfigArgs(targets={"lm_head": "fp8_per_tensor"}) |
| 1203 | ) |
| 1204 | lm_head = VocabParallelEmbedding( |
| 1205 | num_embeddings=1, |
| 1206 | embedding_dim=1, |
| 1207 | disable_tp=True, |
| 1208 | ) |
| 1209 | |
| 1210 | with pytest.raises( |
| 1211 | ValueError, |
| 1212 | match=( |
| 1213 | "Layer lm_head was matched by quantization_config.targets " |
| 1214 | r"\(lm_head\), but online quantization is not supported for " |
| 1215 | "VocabParallelEmbedding." |
| 1216 | ), |
| 1217 | ): |
| 1218 | config.get_quant_method(lm_head, "lm_head") |
| 1219 | |
| 1220 | |
| 1221 | def test_log_online_quantization(default_vllm_config, monkeypatch) -> None: |
| 1222 | config = OnlineQuantizationConfig(quant_config_args(linear="fp8_per_tensor")) |
| 1223 | config.quantized_layers = { |
| 1224 | "model.layers.0.mlp.down_proj": ("linear", "fp8_per_tensor", None), |
| 1225 | "model.layers.1.mlp.down_proj": ("linear", "fp8_per_tensor", None), |
| 1226 | "model.layers.0.self_attn.qkv_proj": ( |
| 1227 | "targets", |
| 1228 | "mxfp4", |
| 1229 | r"re:.*qkv_proj.*", |
| 1230 | ), |
| 1231 | } |
| 1232 | default_vllm_config.quant_config = config |
| 1233 | |
| 1234 | logged_messages: list[str] = [] |
| 1235 | |
| 1236 | def record_info(message: str, *args: object) -> None: |
| 1237 | logged_messages.append(message % args) |
| 1238 | |
| 1239 | monkeypatch.setattr( |
| 1240 | "vllm.model_executor.model_loader.base_loader.logger.info", record_info |
| 1241 | ) |
| 1242 | log_online_quantization(default_vllm_config) |
| 1243 | |
| 1244 | assert logged_messages == [ |
| 1245 | "Quantized 3 layers of types: mlp.down_proj: 2 (from linear: " |
| 1246 | "fp8_per_tensor); self_attn.qkv_proj: 1 (from targets: " |
| 1247 | "re:.*qkv_proj.*, mxfp4)" |
| 1248 | ] |
| 1249 | |
| 1250 | |
| 1251 | @pytest.mark.skipif( |
| 1252 | not ( |
| 1253 | current_platform.is_cuda() |
| 1254 | and current_platform.is_device_capability_family(100) |
| 1255 | and has_flashinfer_trtllm_fused_moe() |
| 1256 | ), |
| 1257 | reason="nvfp4_per_token needs a Blackwell (SM100) GPU + FlashInfer TRTLLM MoE.", |
| 1258 | ) |
| 1259 | def test_online_nvfp4_per_token_moe(vllm_runner, monkeypatch) -> None: |
| 1260 | """Online NVFP4 quantizes the MoE and leaves dense layers unquantized.""" |
| 1261 | monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") |
| 1262 | |
| 1263 | with vllm_runner( |
| 1264 | "ibm-granite/granite-3.0-1b-a400m-base", |
| 1265 | quantization="nvfp4_per_token", |
| 1266 | enforce_eager=True, |
| 1267 | ) as llm: |
| 1268 | |
| 1269 | def check_model(model): |
| 1270 | layer = model.model.layers[0] |
| 1271 | assert isinstance( |
| 1272 | layer.block_sparse_moe.experts._quant_method, Nvfp4OnlineMoEMethod |
| 1273 | ) |
| 1274 | assert isinstance( |
| 1275 | layer.self_attn.o_proj.quant_method, UnquantizedLinearMethod |
| 1276 | ) |
| 1277 | |
| 1278 | llm.apply_model(check_model) |
| 1279 | outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4) |
| 1280 | print(outputs[0][1]) |
| 1281 | |
| 1282 | |
| 1283 | def _patch_max_reduce(monkeypatch, full_amax) -> None: |
| 1284 | """Stand in for the TP/EP MAX all-reduce, returning the unsharded amax.""" |
| 1285 | expected = cast(ProcessGroup, object()) |
| 1286 | stub = SimpleNamespace(device_group=expected) |
| 1287 | monkeypatch.setattr(quant_utils, "get_tp_group", lambda: stub) |
| 1288 | monkeypatch.setattr(quant_utils, "get_ep_group", lambda: stub) |
| 1289 | |
| 1290 | def fake_all_reduce(tensor, op, group): |
| 1291 | assert op == torch.distributed.ReduceOp.MAX |
| 1292 | assert group is expected |
| 1293 | tensor.copy_(full_amax) |
| 1294 | |
| 1295 | monkeypatch.setattr(torch.distributed, "all_reduce", fake_all_reduce) |
| 1296 | |
| 1297 | |
| 1298 | def test_is_tp_sharded_false_when_scale_is_already_global() -> None: |
| 1299 | """Replicated and column-parallel-with-channel-scales need no collective.""" |
| 1300 | replicated = SimpleNamespace( |
| 1301 | tp_size=4, |
| 1302 | input_size=64, |
| 1303 | output_size=32, |
| 1304 | input_size_per_partition=64, |
| 1305 | output_size_per_partition=32, |
| 1306 | ) |
| 1307 | assert not _is_tp_sharded(replicated) |
| 1308 | |
| 1309 | column = SimpleNamespace( |
| 1310 | tp_size=4, |
| 1311 | input_size=64, |
| 1312 | output_size=32, |
| 1313 | input_size_per_partition=64, |
| 1314 | output_size_per_partition=8, |
| 1315 | ) |
| 1316 | assert not _is_tp_sharded(column, reduces_output_dim=False) |
| 1317 | assert _is_tp_sharded(column) |
| 1318 | |
| 1319 | |
| 1320 | def _quantize_linear(weight, scheme, is_sharded): |
| 1321 | if scheme == "per_tensor": |
| 1322 | amax = weight_amax(weight).reshape(1) |
| 1323 | scale = _fp8_scale(amax_for_tp_weight_quant(amax, is_sharded)) |
| 1324 | return ops.scaled_fp8_quant(weight, scale=scale)[0], scale |
| 1325 | amax = weight_amax(weight, dim=-1, keepdim=True) |
| 1326 | scale = _fp8_channel_scale(amax_for_tp_weight_quant(amax, is_sharded)) |
| 1327 | return _fp8_quant_per_channel(weight, scale), scale |
| 1328 | |
| 1329 | |
| 1330 | @pytest.mark.skipif( |
| 1331 | not is_quant_method_supported("fp8"), |
| 1332 | reason="FP8 is not supported on this GPU type.", |
| 1333 | ) |
| 1334 | @pytest.mark.parametrize("scheme", ["per_tensor", "per_channel"]) |
| 1335 | @pytest.mark.parametrize("shard_dim", [0, 1]) |
| 1336 | def test_online_linear_tp_weight_quant_matches_unsharded( |
| 1337 | monkeypatch, scheme: str, shard_dim: int |
| 1338 | ) -> None: |
| 1339 | """TP shards pack the same FP8 values and scales as the unsharded weight.""" |
| 1340 | torch.manual_seed(0) |
| 1341 | weight = torch.randn(32, 64, device=DEVICE, dtype=torch.bfloat16) |
| 1342 | weight[-1, -1] = 64.0 |
| 1343 | |
| 1344 | full_weight, full_scale = _quantize_linear(weight, scheme, False) |
| 1345 | |
| 1346 | # Per-channel scales reduce only the input dim, so a column (dim 0) shard |
| 1347 | # already matches without a collective. |
| 1348 | is_sharded = scheme == "per_tensor" or shard_dim == 1 |
| 1349 | if is_sharded: |
| 1350 | full_amax = ( |
| 1351 | weight_amax(weight).reshape(1) |
| 1352 | if scheme == "per_tensor" |
| 1353 | else weight_amax(weight, dim=-1, keepdim=True) |
| 1354 | ) |
| 1355 | _patch_max_reduce(monkeypatch, full_amax) |
| 1356 | |
| 1357 | shard_size = weight.shape[shard_dim] // 2 |
| 1358 | shard = weight.narrow(shard_dim, 0, shard_size).contiguous() |
| 1359 | tp_weight, tp_scale = _quantize_linear(shard, scheme, is_sharded) |
| 1360 | |
| 1361 | assert torch.equal(tp_weight, full_weight.narrow(shard_dim, 0, shard_size)) |
| 1362 | if scheme == "per_channel" and shard_dim == 0: |
| 1363 | assert torch.equal(tp_scale, full_scale.narrow(0, 0, shard_size)) |
| 1364 | else: |
| 1365 | assert torch.equal(tp_scale, full_scale) |
| 1366 | |
| 1367 | |
| 1368 | def _quantize_moe(weight, scheme, moe_tp_size): |
| 1369 | if scheme == "nvfp4": |
| 1370 | return _quantize_moe_weight_to_nvfp4(weight, moe_tp_size) |
| 1371 | if scheme == "per_tensor": |
| 1372 | amax = weight_amax(weight.flatten(1), dim=-1) |
| 1373 | scale = _fp8_scale(amax_for_moe_weight_quant(amax, moe_tp_size)) |
| 1374 | quant = lambda w, s: ops.scaled_fp8_quant(w, scale=s)[0] # noqa: E731 |
| 1375 | else: |
| 1376 | amax = weight_amax(weight, dim=-1, keepdim=True) |
| 1377 | scale = _fp8_channel_scale(amax_for_moe_weight_quant(amax, moe_tp_size)) |
| 1378 | quant = _fp8_quant_per_channel |
| 1379 | qweight = torch.stack([quant(w, s) for w, s in zip(weight, scale)]) |
| 1380 | return qweight, scale |
| 1381 | |
| 1382 | |
| 1383 | @pytest.mark.parametrize("scheme", ["per_tensor", "per_channel", "nvfp4"]) |
| 1384 | def test_online_moe_tp_weight_quant_matches_ep(monkeypatch, scheme: str) -> None: |
| 1385 | """TP shards of w2 pack the same values and scales as full experts.""" |
| 1386 | if scheme == "nvfp4": |
| 1387 | if ( |
| 1388 | not ( |
| 1389 | current_platform.is_cuda() |
| 1390 | and current_platform.is_device_capability_family(100) |
| 1391 | ) |
| 1392 | or current_platform.is_xpu() |
| 1393 | ): |
| 1394 | pytest.skip("NVFP4 weight quantization needs a Blackwell (SM100) GPU.") |
| 1395 | elif not is_quant_method_supported("fp8"): |
| 1396 | pytest.skip("FP8 is not supported on this GPU type.") |
| 1397 | |
| 1398 | torch.manual_seed(0) |
| 1399 | weight = torch.randn(2, 32, 32, device=DEVICE, dtype=torch.bfloat16) |
| 1400 | weight[:, -1, -1] = torch.tensor([32.0, 64.0], device=DEVICE) |
| 1401 | |
| 1402 | ep_out = _quantize_moe(weight, scheme, 1) |
| 1403 | |
| 1404 | full_amax = ( |
| 1405 | weight_amax(weight, dim=-1, keepdim=True) |
| 1406 | if scheme == "per_channel" |
| 1407 | else weight_amax(weight.flatten(1), dim=-1).to(torch.float32) |
| 1408 | ) |
| 1409 | _patch_max_reduce(monkeypatch, full_amax) |
| 1410 | |
| 1411 | # w2 is sharded along its last (intermediate) dim. |
| 1412 | shard_size = weight.shape[2] // 2 |
| 1413 | tp_out = _quantize_moe(weight[:, :, :shard_size], scheme, 2) |
| 1414 | |
| 1415 | packing = 2 if scheme == "nvfp4" else 1 |
| 1416 | assert torch.equal(tp_out[0], ep_out[0][:, :, : shard_size // packing]) |
| 1417 | if scheme == "nvfp4": |
| 1418 | assert torch.equal(tp_out[1], ep_out[1][:, :, : shard_size // 16]) |
| 1419 | assert torch.equal(tp_out[-1], ep_out[-1]) |
| 1420 | |
| 1421 | |
| 1422 | def test_online_int8_moe_w2_scale_matches_unsharded(monkeypatch) -> None: |
| 1423 | """Int8 MoE w2 reduces over the sharded intermediate dim.""" |
| 1424 | torch.manual_seed(0) |
| 1425 | w13 = torch.randn(2, 16, 8, dtype=torch.bfloat16) |
| 1426 | w2 = torch.randn(2, 8, 16, dtype=torch.bfloat16) |
| 1427 | w2[:, -1, -1] = 64.0 |
| 1428 | |
| 1429 | def quantize(w2_in, moe_tp_size): |
| 1430 | layer = torch.nn.Module() |
| 1431 | layer.w13_weight = torch.nn.Parameter(w13, requires_grad=False) |
| 1432 | layer.w2_weight = torch.nn.Parameter(w2_in, requires_grad=False) |
| 1433 | layer.num_experts = layer.local_num_experts = w13.shape[0] |
| 1434 | method = SimpleNamespace(moe=SimpleNamespace(tp_size=moe_tp_size)) |
| 1435 | Int8OnlineMoEMethod._quantize_weights(method, layer) |
| 1436 | return layer.w2_weight, layer.w2_scale |
| 1437 | |
| 1438 | full_weight, full_scale = quantize(w2, 1) |
| 1439 | |
| 1440 | _patch_max_reduce(monkeypatch, weight_amax(w2, dim=-1)) |
| 1441 | |
| 1442 | shard_size = w2.shape[2] // 2 |
| 1443 | tp_weight, tp_scale = quantize(w2[:, :, :shard_size].contiguous(), 2) |
| 1444 | |
| 1445 | assert torch.equal(tp_weight, full_weight[:, :, :shard_size]) |
| 1446 | assert torch.equal(tp_scale, full_scale) |
| 1447 | |
| 1448 | |
| 1449 | @pytest.mark.skipif( |
| 1450 | not ( |
| 1451 | current_platform.is_cuda() and current_platform.is_device_capability_family(100) |
| 1452 | ), |
| 1453 | reason="NVFP4 weight quantization needs a Blackwell (SM100) GPU.", |
| 1454 | ) |
| 1455 | def test_online_nvfp4_quantizes_original_expert_weights() -> None: |
| 1456 | torch.manual_seed(0) |
| 1457 | weight = torch.randn(2, 32, 32, device="cuda", dtype=torch.bfloat16) |
| 1458 | |
| 1459 | quantized, block_scale, global_decode_scale = _quantize_moe_weight_to_nvfp4(weight) |
| 1460 | global_encode_scale = 1.0 / global_decode_scale |
| 1461 | expected = [ |
| 1462 | scaled_fp4_quant( |
| 1463 | expert_weight, |
| 1464 | expert_scale, |
| 1465 | is_sf_swizzled_layout=False, |
| 1466 | ) |
| 1467 | for expert_weight, expert_scale in zip( |
| 1468 | weight, |
| 1469 | global_encode_scale, |
| 1470 | strict=True, |
| 1471 | ) |
| 1472 | ] |
| 1473 | |
| 1474 | assert torch.equal( |
| 1475 | quantized, |
| 1476 | torch.stack([expert_weight for expert_weight, _ in expected]), |
| 1477 | ) |
| 1478 | assert torch.equal( |
| 1479 | block_scale, |
| 1480 | torch.stack([expert_scale for _, expert_scale in expected]), |
| 1481 | ) |
| 1482 | |
| 1483 | |
| 1484 | @pytest.mark.skipif( |
| 1485 | not is_quant_method_supported("fp8"), |
| 1486 | reason="FP8 is not supported on this GPU type.", |
| 1487 | ) |
| 1488 | def test_online_quant_peak_mem( |
| 1489 | vllm_runner, |
| 1490 | caplog_mp_spawn, |
| 1491 | monkeypatch, |
| 1492 | ) -> None: |
| 1493 | _test_online_quant_peak_mem_impl( |
| 1494 | "fp8_per_tensor", vllm_runner, caplog_mp_spawn, monkeypatch |
| 1495 | ) |
| 1496 | |
| 1497 | |
| 1498 | @pytest.mark.skipif( |
| 1499 | not is_quant_method_supported("fp8"), |
| 1500 | reason="FP8 is not supported on this GPU type.", |
| 1501 | ) |
| 1502 | def test_online_quant_load_format_dummy( |
| 1503 | vllm_runner, |
| 1504 | monkeypatch, |
| 1505 | caplog, |
| 1506 | ) -> None: |
| 1507 | with vllm_runner( |
| 1508 | "ibm-granite/granite-3.0-1b-a400m-base", |
| 1509 | quantization="fp8_per_tensor", |
| 1510 | enforce_eager=True, |
| 1511 | load_format="dummy", |
| 1512 | ) as llm: |
| 1513 | outputs = llm.generate_greedy(["The future of AI is"], max_tokens=4) |
| 1514 | print(outputs[0][1]) |