# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Test model set-up and weight loading for llmcompressor-quantized models. Run `pytest tests/quantization/test_compressed_tensors.py`. """ from contextlib import contextmanager from unittest.mock import Mock import pytest import torch from compressed_tensors.quantization import ( ActivationOrdering, QuantizationArgs, QuantizationStrategy, QuantizationType, ) from tests.models.utils import check_logprobs_close from vllm.model_executor.kernels.linear import ( Fp8BlockScaledMMLinearKernel, ) from vllm.model_executor.layers.fused_moe import UnquantizedFusedMoEMethod from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501 CompressedTensorsConfig, CompressedTensorsLinearMethod, CompressedTensorsW4A4Fp4, CompressedTensorsW4A4Mxfp4, CompressedTensorsW4A8Fp8, CompressedTensorsW8A8Fp8, CompressedTensorsW8A8Int8, CompressedTensorsW8A8Mxfp8, CompressedTensorsW8A16Fp8, CompressedTensorsWNA8O8Int, CompressedTensorsWNA16, ) from vllm.model_executor.layers.quantization.compressed_tensors.utils import ( find_matched_target, ) from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8 from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.platforms import current_platform from vllm.v1.attention.backends.fa_utils import get_flash_attn_version # AITER only supports per-channel-per-channel INT8 gemm # and per-tensor-per-tensor INT8 GEMM. # It does not support mix precision MM and mix quantization scheme. ROCM_AITER_SUPPORTED_INT8_MODEL = [ "neuralmagic/Llama-3.2-1B-quantized.w8a8", "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2", ] # TritonInt8ScaledMMLinearKernel only supports symmetric quantization. ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL = [ "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change", "nm-testing/tinyllama-oneshot-w8-channel-a8-tensor", "neuralmagic/Llama-3.2-1B-quantized.w8a8", "nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2", "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2", ] @pytest.fixture(scope="function", autouse=True) def enable_pickle(monkeypatch): """`LLM.apply_model` requires pickling a function.""" monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1") @pytest.mark.parametrize( "model_args", [ ( "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change", "tensor", QuantizationType.INT, 2560, True, ), ( "nm-testing/asym-w8w8-int8-static-per-tensor-tiny-llama", "tensor", QuantizationType.INT, 2560, False, ), ], ) def test_compressed_tensors_w8a8_static_setup(vllm_runner, model_args): model_path, strategy, quant_type, shape_0, is_symmetric = model_args if ( current_platform.is_rocm() and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL ): pytest.skip(f"Skip model {model_path} as it is not supported on ROCm.") with vllm_runner(model_path, enforce_eager=True) as llm: def check_model(model): layer = model.model.layers[0] qkv_proj = layer.self_attn.qkv_proj o_proj = layer.self_attn.o_proj gate_up_proj = layer.mlp.gate_up_proj down_proj = layer.mlp.down_proj # assert zp for symmetric and asymmetric cases def zp_valid(zp: torch.Tensor | None): if is_symmetric: return zp is None return zp is not None and zp.dtype is torch.int32 assert zp_valid(qkv_proj.input_zero_point) assert zp_valid(o_proj.input_zero_point) assert zp_valid(gate_up_proj.input_zero_point) assert zp_valid(down_proj.input_zero_point) assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) assert isinstance(o_proj.quant_method, CompressedTensorsLinearMethod) assert isinstance(gate_up_proj.quant_method, CompressedTensorsLinearMethod) assert isinstance(down_proj.quant_method, CompressedTensorsLinearMethod) assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8) assert qkv_proj.scheme.strategy == strategy assert qkv_proj.scheme.is_static_input_scheme expected_type = torch.int8 assert qkv_proj.weight.dtype is expected_type assert o_proj.weight.dtype is expected_type assert gate_up_proj.weight.dtype is expected_type if qkv_proj.scheme.strategy == "tensor": # Make sure it is a channelwise buffer # After running process_weights_after_loading assert len(qkv_proj.weight_scale.shape) == 2 assert qkv_proj.weight_scale.shape[0] == shape_0 assert qkv_proj.weight_scale.shape[1] == 1 assert qkv_proj.weight_scale.dtype is torch.float32 assert qkv_proj.input_scale.dtype is torch.float32 llm.apply_model(check_model) output = llm.generate_greedy(["Hello my name is"], max_tokens=4) assert output @pytest.mark.parametrize( "model_path", [ "neuralmagic/Llama-3.2-1B-quantized.w8a8", ], ) @pytest.mark.parametrize("max_tokens", [4]) @pytest.mark.parametrize("num_logprobs", [10]) @pytest.mark.parametrize( "use_aiter", [True, False] if current_platform.is_rocm() else [False] ) def test_compressed_tensors_w8a8_logprobs( hf_runner, vllm_runner, example_prompts, model_path, max_tokens, num_logprobs, use_aiter, monkeypatch, ): if ( current_platform.is_rocm() and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL ): pytest.skip(f"Skip model {model_path} as it is not supported on ROCm.") if use_aiter: if model_path not in ROCM_AITER_SUPPORTED_INT8_MODEL: pytest.skip(f"Skip model {model_path} as it is not support by aiter.") # this will enable VLLM_ROCM_USE_AITER_LINEAR monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") dtype = "bfloat16" # skip language translation prompt for the static per tensor models if model_path in ( "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Sym", "nm-testing/Meta-Llama-3-8B-Instruct-W8A8-Static-Per-Tensor-Asym", ): example_prompts = example_prompts[0:-1] with hf_runner(model_path, dtype=dtype) as hf_model: hf_outputs = hf_model.generate_greedy_logprobs_limit( example_prompts, max_tokens, num_logprobs ) with vllm_runner(model_path, dtype=dtype, enforce_eager=True) as vllm_model: vllm_outputs = vllm_model.generate_greedy_logprobs( example_prompts, max_tokens, num_logprobs ) check_logprobs_close( outputs_0_lst=hf_outputs, outputs_1_lst=vllm_outputs, name_0="hf", name_1="vllm", ) if current_platform.is_rocm(): torch.accelerator.synchronize() def test_compressed_tensors_no_enforce_eager(vllm_runner): model_path = "nm-testing/tinyllama-oneshot-w8w8-test-static-shape-change" with vllm_runner(model_path) as llm: output = llm.generate_greedy("Hello my name is", max_tokens=4) assert output @pytest.mark.parametrize( "model_args", [ ("nm-testing/tinyllama-oneshot-w8a8-dynamic-token-v2", "tensor"), ( "nm-testing/tinyllama-oneshot-w8a8-channel-dynamic-token-v2", "channel", ), ], ) @pytest.mark.parametrize( "use_aiter", [True, False] if current_platform.is_rocm() else [False] ) def test_compressed_tensors_w8a8_dynamic_per_token( vllm_runner, model_args, use_aiter, monkeypatch, ): model_path, strategy = model_args if ( current_platform.is_rocm() and model_path not in ROCM_TRITON_SCALED_MM_SUPPORTED_INT8_MODEL ): pytest.skip(f"Skip model {model_path} as it is not supported on ROCm.") if use_aiter: if model_path not in ROCM_AITER_SUPPORTED_INT8_MODEL: pytest.skip(f"Skip model {model_path} as it is not support by aiter.") # this will enable VLLM_ROCM_USE_AITER_LINEAR monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") with vllm_runner(model_path, enforce_eager=True, dtype=torch.float16) as llm: def check_model(model): layer = model.model.layers[0] qkv_proj = layer.self_attn.qkv_proj assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8) assert not qkv_proj.scheme.is_static_input_scheme assert qkv_proj.scheme.strategy == strategy assert qkv_proj.weight.dtype is torch.int8 llm.apply_model(check_model) output = llm.generate_greedy(["Hello my name is"], max_tokens=4) assert output @pytest.mark.parametrize( "wNa16_args", [ ( "nm-testing/tinyllama-oneshot-w4a16-channel-v2", "channel", None, 8, True, False, ), ( "nm-testing/TinyLlama-1.1B-Chat-v1.0-W4A16-G128-Asym-Updated-ActOrder", "group", 128, 8, False, True, ), ], ) @pytest.mark.skipif( not current_platform.is_cuda(), reason="The tests are skipped on non-CUDA platform." ) def test_compressed_tensors_wNa16(vllm_runner, wNa16_args): model, strategy, group, pack_factor, symmetric, has_g_idx = wNa16_args with vllm_runner(model, enforce_eager=True) as llm: def check_model(model): layer = model.model.layers[0] qkv_proj = layer.self_attn.qkv_proj assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) assert isinstance(qkv_proj.scheme, CompressedTensorsWNA16) assert qkv_proj.scheme.strategy == strategy assert qkv_proj.scheme.group_size == (-1 if group is None else group) assert qkv_proj.scheme.pack_factor == pack_factor assert qkv_proj.scheme.symmetric == symmetric assert qkv_proj.scheme.has_g_idx == has_g_idx llm.apply_model(check_model) output = llm.generate_greedy("Hello my name is", max_tokens=4) assert output def test_compressed_tensors_fp8(vllm_runner): model_path = "nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test" with vllm_runner(model_path, enforce_eager=True) as llm: def check_model(model): layer = model.model.layers[0] qkv_proj = layer.self_attn.qkv_proj assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) assert isinstance( qkv_proj.scheme, (CompressedTensorsW8A8Fp8, CompressedTensorsW8A16Fp8), ) assert qkv_proj.input_scale.dtype is torch.float32 if isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8): assert len(qkv_proj.input_scale.shape) == 0 assert qkv_proj.weight.dtype is current_platform.fp8_dtype() assert qkv_proj.weight_scale.dtype is torch.float32 assert len(qkv_proj.weight_scale.shape) == 0 llm.apply_model(check_model) output = llm.generate_greedy("Hello my name is", max_tokens=4) assert output @pytest.mark.skipif( not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." ) def test_compressed_tensors_kv_cache_fp8_per_tensor(vllm_runner): model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-kvcache-fp8-tensor" with vllm_runner(model_path) as llm: output = llm.generate_greedy("Hello world!", max_tokens=4) assert output @pytest.mark.skipif( not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." ) def test_compressed_tensors_kv_cache_fp8_per_attn_head(vllm_runner): model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-kvcache-fp8-attn_head" try: fa_version = get_flash_attn_version() except Exception: pytest.skip("This test requires FlashAttention backend.") if fa_version is None or fa_version < 3: pytest.skip("This test requires FlashAttention version >= 3.") with vllm_runner(model_path, attention_config={"backend": "FLASH_ATTN"}) as llm: output = llm.generate_greedy("Hello world!", max_tokens=4) assert output @contextmanager def _nvfp4_marlin_error_context(model, capfd): is_rocm_and_unsupported = ( model == "nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4A16" and current_platform.is_rocm() ) if is_rocm_and_unsupported: expected_error = ( "ValueError: Forced NVFP4 kernel MarlinNvFp4LinearKernel is not " "supported: Marlin FP4 not available" ) with pytest.raises(RuntimeError, match="Engine core initialization failed"): yield captured = capfd.readouterr() assert expected_error in captured.out + captured.err else: yield @pytest.mark.parametrize( "args", [ ("nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4A16", True), ("nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4", False), ], ) def test_compressed_tensors_nvfp4(vllm_runner, args, capfd): model, use_a16 = args with ( _nvfp4_marlin_error_context(model, capfd), vllm_runner(model, enforce_eager=True) as llm, ): def check_model(model): layer = model.model.layers[0] qkv_proj = layer.self_attn.qkv_proj assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) assert isinstance(qkv_proj.scheme, CompressedTensorsW4A4Fp4) assert qkv_proj.scheme.use_a16 == use_a16 assert qkv_proj.scheme.group_size == 16 llm.apply_model(check_model) output = llm.generate_greedy(["Hello my name is"], max_tokens=4) print(output) assert output @pytest.mark.skipif( not current_platform.is_cuda() or not current_platform.has_device_capability(90), reason="W4A8 FP8 is not yet supported on this GPU type.", ) @pytest.mark.parametrize( "args", [("czhu-cohere/TinyLlama-1.1B-Chat-v1.0-W4A8-e2e", CompressedTensorsW4A8Fp8)], ) def test_compressed_tensors_w4a8_fp8(vllm_runner, args): model, scheme = args with vllm_runner(model, enforce_eager=True) as llm: def check_model(model): layer = model.model.layers[0] qkv_proj = layer.self_attn.qkv_proj o_proj = layer.self_attn.o_proj gate_up_proj = layer.mlp.gate_up_proj down_proj = layer.mlp.down_proj for proj in (qkv_proj, o_proj, gate_up_proj, down_proj): assert isinstance(proj.quant_method, CompressedTensorsLinearMethod) assert isinstance(proj.scheme, scheme) assert proj.weight_packed.dtype is torch.int32 assert proj.weight_scale.dtype is torch.float8_e4m3fn assert proj.weight_chan_scale.dtype is torch.float32 assert proj.scheme.group_size == 128 llm.apply_model(check_model) output = llm.generate_greedy("Hello my name is", max_tokens=4) print(output) assert output @pytest.mark.skipif( not current_platform.is_cuda(), reason="This test is skipped on non-CUDA platform." ) @pytest.mark.parametrize( "model,prompt,exp_perplexity", [ ( "nm-testing/Llama-3.2-1B-Instruct-spinquantR1R2R4-w4a16", "Flat is better than nested.\nSparse is better than dense.", 150.0, ), ( "nm-testing/Llama-3.2-1B-Instruct-quip-w4a16", "Flat is better than nested.\nSparse is better than dense.", 150.0, ), ], ) def test_compressed_tensors_transforms_perplexity( vllm_runner, model, prompt, exp_perplexity ): with vllm_runner(model, enforce_eager=True) as llm: perplexity = llm.generate_prompt_perplexity([prompt])[0] print(perplexity) assert perplexity <= exp_perplexity def test_compressed_tensors_fp8_block_enabled(vllm_runner): model_path = "RedHatAI/Qwen3-0.6B-FP8-BLOCK" with vllm_runner(model_path, enforce_eager=True) as llm: fp8_dtype = current_platform.fp8_dtype() def check_model(model): layer = model.model.layers[0] qkv_proj = layer.self_attn.qkv_proj assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod) assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8) assert isinstance(qkv_proj.scheme.fp8_linear, Fp8BlockScaledMMLinearKernel) assert qkv_proj.weight.dtype is fp8_dtype assert qkv_proj.weight_scale.dtype is torch.float32 assert len(qkv_proj.weight.shape) == 2 assert len(qkv_proj.weight_scale.shape) == 2 input_quant_op = qkv_proj.scheme.fp8_linear.quant_fp8 assert isinstance(input_quant_op, QuantFP8) assert input_quant_op._forward_method in ( input_quant_op.forward_cuda, input_quant_op.forward_hip, input_quant_op.forward_xpu, ) llm.apply_model(check_model) output = llm.generate_greedy("Hello my name is", max_tokens=4) assert output @pytest.mark.skipif( not current_platform.is_cuda(), reason="This test is not for non-CUDA platforms", ) def test_compressed_tensors_moe_ignore_with_model(vllm_runner): """ Integration test for MoE layer ignore functionality with a real model. This test would verify that when loading a compressed-tensors quantized MoE model where some MoE layers are in the ignore list, those layers use UnquantizedFusedMoEMethod while non-ignored layers use the quantized method. Expected model structure: - Compressed-tensors quantized MoE model (e.g., Mixtral-based) - Config with ignore list containing specific MoE layers - Multiple MoE layers where some are quantized and some are not """ # model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only" # CT 12.3 model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable" # CT 12.2 with vllm_runner(model_path, enforce_eager=True) as llm: def check_model(model): from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501 CompressedTensorsMoEMethod, ) # Check layer 0 MoE (should be quantized) layer_quantized = model.model.layers[0].mlp.experts assert isinstance(layer_quantized, MoERunner) assert isinstance(layer_quantized._quant_method, CompressedTensorsMoEMethod) # Check layer 10 MoE (should be unquantized + ignored) layer_unquantized = model.model.layers[3].mlp.experts assert isinstance(layer_unquantized, MoERunner) assert isinstance( layer_unquantized._quant_method, UnquantizedFusedMoEMethod ) llm.apply_model(check_model) # Verify the model can generate output output = llm.generate_greedy("Hello, my name is", max_tokens=4) assert output def test_w4a16_moe_torch_compile(vllm_runner): """Regression test: MoE quant_config must be initialized inside the moe_forward custom op, not just in forward_native which is compiled by Dynamo (attribute mutations are not replayed at runtime). Without the fix in _moe_forward/_moe_forward_shared, this hits: AssertionError: Hidden size mismatch 2048 != 1024 because use_int4_w4a16 is False (moe_quant_config stays None). """ model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable" with vllm_runner( model_path, enforce_eager=False, max_model_len=256, compilation_config={ "cudagraph_mode": "NONE", }, ) as llm: output = llm.generate_greedy("Hi", max_tokens=1) assert output def _make_ct_config(*, target: str = "Linear") -> CompressedTensorsConfig: """Build a minimal CompressedTensorsConfig with INT8 channel quant.""" weight_quant = QuantizationArgs( num_bits=8, type=QuantizationType.INT, strategy=QuantizationStrategy.CHANNEL, symmetric=True, dynamic=False, ) return CompressedTensorsConfig( target_scheme_map={ target: { "weights": weight_quant, "input_activations": None, "format": "pack-quantized", } }, ignore=[], quant_format="pack-quantized", ) def test_get_quant_method_returns_linear_method_for_parallel_lm_head(): """ParallelLMHead whose name matches a target must get a quantised method.""" config = _make_ct_config(target="re:.*lm_head") mock_lm_head = Mock(spec=ParallelLMHead) mock_lm_head.__class__ = ParallelLMHead method = config.get_quant_method(mock_lm_head, prefix="model.lm_head") assert isinstance(method, CompressedTensorsLinearMethod), ( f"Expected CompressedTensorsLinearMethod, got {type(method).__name__}" ) def test_get_quant_method_returns_none_for_ignored_parallel_lm_head(): """ParallelLMHead on the ignore list should be left unquantized (None).""" config = _make_ct_config(target="re:.*lm_head") config.ignore = ["re:.*lm_head"] mock_lm_head = Mock(spec=ParallelLMHead) mock_lm_head.__class__ = ParallelLMHead method = config.get_quant_method(mock_lm_head, prefix="model.lm_head") assert method is None, ( f"Expected None for ignored ParallelLMHead, got {type(method).__name__}" ) def test_get_quant_method_returns_none_for_unmatched_parallel_lm_head(): """ParallelLMHead with target='Linear' (typical real model) must not crash. Most compressed-tensors models only target 'Linear'. ParallelLMHead does not match that target, so get_quant_method should return None (unquantized) instead of raising ValueError. """ config = _make_ct_config(target="Linear") mock_lm_head = Mock(spec=ParallelLMHead) mock_lm_head.__class__ = ParallelLMHead method = config.get_quant_method(mock_lm_head, prefix="model.lm_head") assert method is None, ( f"Expected None for unmatched ParallelLMHead, got {type(method).__name__}" ) def test_find_matched_target_returns_none_on_no_match(): result = find_matched_target( layer_name="model.layers.0.self_attn.qkv_proj", module=Mock(spec=torch.nn.Linear), targets=["no_match_target"], ) assert result is None def test_get_scheme_dict_returns_none_on_no_match(): config = _make_ct_config(target="matched_layer") result = config.get_scheme_dict( layer=Mock(spec=torch.nn.Linear), layer_name="model.layers.0.unmatched_layer", ) assert result is None # Test constants for activation quantization _STATIC_SYM_INT8_ACT = QuantizationArgs( num_bits=8, type=QuantizationType.INT, strategy=QuantizationStrategy.TENSOR.value, symmetric=True, dynamic=False, ) _STATIC_ASYM_INT8_ACT = QuantizationArgs( num_bits=8, type=QuantizationType.INT, strategy=QuantizationStrategy.TENSOR.value, symmetric=False, dynamic=False, ) _DYNAMIC_INT8_ACT = QuantizationArgs( num_bits=8, type=QuantizationType.INT, strategy=QuantizationStrategy.TOKEN.value, symmetric=True, dynamic=True, ) @pytest.mark.parametrize( "weight_bits,weight_strategy,input_act,output_act,format,expected_scheme", [ # W8A8 int-quantized -> W8A8Int8 (regression test for #46389) pytest.param( 8, QuantizationStrategy.CHANNEL.value, _STATIC_SYM_INT8_ACT, None, "int-quantized", CompressedTensorsW8A8Int8, id="w8a8_channel_static_sym", ), pytest.param( 8, QuantizationStrategy.CHANNEL.value, _STATIC_ASYM_INT8_ACT, None, "int-quantized", CompressedTensorsW8A8Int8, id="w8a8_channel_static_asym", ), pytest.param( 8, QuantizationStrategy.TENSOR.value, _STATIC_SYM_INT8_ACT, None, "int-quantized", CompressedTensorsW8A8Int8, id="w8a8_tensor_static", ), pytest.param( 8, QuantizationStrategy.CHANNEL.value, _DYNAMIC_INT8_ACT, None, "int-quantized", CompressedTensorsW8A8Int8, id="w8a8_channel_dynamic", ), # W8A8O8 int-quantized -> WNA8O8Int (both input and output) pytest.param( 8, QuantizationStrategy.CHANNEL.value, _STATIC_SYM_INT8_ACT, _STATIC_SYM_INT8_ACT, "int-quantized", CompressedTensorsWNA8O8Int, id="w8a8o8_channel", ), pytest.param( 4, QuantizationStrategy.GROUP.value, _STATIC_SYM_INT8_ACT, _STATIC_SYM_INT8_ACT, "int-quantized", CompressedTensorsWNA8O8Int, id="w4a8o8_group", ), # Weight-only pack-quantized -> WNA16 pytest.param( 8, QuantizationStrategy.CHANNEL.value, None, None, "pack-quantized", CompressedTensorsWNA16, id="w8_pack", ), pytest.param( 4, QuantizationStrategy.GROUP.value, None, None, "pack-quantized", CompressedTensorsWNA16, id="w4_pack", ), pytest.param( 2, QuantizationStrategy.GROUP.value, None, None, "pack-quantized", CompressedTensorsWNA16, id="w2_pack", ), pytest.param( 3, QuantizationStrategy.GROUP.value, None, None, "pack-quantized", CompressedTensorsWNA16, id="w3_pack", ), pytest.param( 5, QuantizationStrategy.GROUP.value, None, None, "pack-quantized", CompressedTensorsWNA16, id="w5_pack", ), pytest.param( 6, QuantizationStrategy.GROUP.value, None, None, "pack-quantized", CompressedTensorsWNA16, id="w6_pack", ), pytest.param( 7, QuantizationStrategy.GROUP.value, None, None, "pack-quantized", CompressedTensorsWNA16, id="w7_pack", ), ], ) def test_scheme_selection( weight_bits, weight_strategy, input_act, output_act, format, expected_scheme ): """Test that _get_scheme_from_parts selects the correct scheme. This parametrized test verifies scheme selection for various combinations of weight bits, quantization strategies, input/output activations, and compression formats. Key regression test: W8A8 int-quantized models with channel-wise weights should use W8A8Int8 (true int8 gemm), not WNA8O8Int (fake-quant). WNA8O8Int should only match when BOTH input and output activations are present. """ weight_quant = QuantizationArgs( num_bits=weight_bits, type=QuantizationType.INT, strategy=weight_strategy, symmetric=True, dynamic=False, group_size=128 if weight_strategy == QuantizationStrategy.GROUP.value else None, ) config = CompressedTensorsConfig( target_scheme_map={}, ignore=[], quant_format=format, ) scheme = config._get_scheme_from_parts( weight_quant=weight_quant, input_quant=input_act, output_quant=output_act, format=format, ) assert isinstance(scheme, expected_scheme), ( f"Expected {expected_scheme.__name__} for " f"W{weight_bits} {weight_strategy} + " f"input_act={input_act} + output_act={output_act} + " f"format={format}, got {type(scheme).__name__}" ) @pytest.mark.skipif( not current_platform.is_cuda() or not current_platform.has_device_capability(75), reason="MXFP8 requires Turing (sm_75+) or newer.", ) def test_compressed_tensors_mxfp8_moe_setup(vllm_runner): """Verify MXFP8 scheme, dtypes, and generation for a MoE model.""" model_path = "AliEdalati97/Qwen3-30B-A3B-MXFP8" with vllm_runner( model_path, enforce_eager=True, load_format="dummy", hf_overrides={"num_hidden_layers": 4}, ) as llm: def check_model(model): from vllm.model_executor.layers.fused_moe import MoERunner from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_w8a8_mxfp8 import ( # noqa: E501 CompressedTensorsW8A8Mxfp8MoEMethod, ) layer = model.model.layers[0] qkv = layer.self_attn.qkv_proj assert isinstance(qkv.quant_method, CompressedTensorsLinearMethod) assert isinstance(qkv.scheme, CompressedTensorsW8A8Mxfp8) experts = layer.mlp.experts assert isinstance(experts, MoERunner) assert isinstance( experts._quant_method, CompressedTensorsW8A8Mxfp8MoEMethod ) llm.apply_model(check_model) output = llm.generate_greedy("Hello my name is", max_tokens=4) assert output @pytest.mark.parametrize( "actorder,group_size,part,full,expected", [ # actorder="group" with real grouping: must load full-K w2 scales and, # when sharded (part != full), report is_k_full=False. (ActivationOrdering.GROUP, 32, 64, 128, (True, 128, False)), # actorder="group" but unsharded (part == full): full scales, k_full. (ActivationOrdering.GROUP, 32, 128, 128, (True, 128, True)), # actorder="group" with channel-wise (group_size == -1): no full load. (ActivationOrdering.GROUP, -1, 64, 128, (False, 64, False)), # "static"/"weight" reorder at quant time -> shard normally + k_full. # Regression: static actorder under TP must keep is_k_full=True so the # Marlin kernel never gets the invalid (group_size=16, is_k_full=0). ("static", 32, 64, 128, (False, 64, True)), ("weight", 32, 64, 128, (False, 64, True)), (None, 32, 64, 128, (False, 64, True)), ], ) def test_wna16_marlin_moe_w2_scale_sharding(actorder, group_size, part, full, expected): from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe.compressed_tensors_moe_wna16_marlin import ( # noqa: E501 CompressedTensorsWNA16MarlinMoEMethod, ) result = CompressedTensorsWNA16MarlinMoEMethod._w2_scale_sharding( actorder, group_size, part, full ) assert result == expected @pytest.mark.skipif( not current_platform.is_cuda() or not current_platform.has_device_capability(80), reason="MXFP4 requires ampere or newer", ) def test_compressed_tensors_mxfp4(vllm_runner): model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-MXFP4" with vllm_runner(model_path, enforce_eager=True) as llm: def check_model(model): layer = model.model.layers[0] qkv_proj = layer.self_attn.qkv_proj o_proj = layer.self_attn.o_proj gate_up_proj = layer.mlp.gate_up_proj down_proj = layer.mlp.down_proj for proj in (qkv_proj, o_proj, gate_up_proj, down_proj): assert isinstance(proj.quant_method, CompressedTensorsLinearMethod) assert isinstance(proj.scheme, CompressedTensorsW4A4Mxfp4) # Verify group size assert proj.scheme.group_size == 32 llm.apply_model(check_model) output = llm.generate_greedy("Hello my name is", max_tokens=4) assert output