A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 5, 19:45 UTC
[Bugfix] Handle ParallelLMHead in compressed-tensors get_quant_method (#37291)
Signed-off-by: Matthias Gehre <matthias.gehre@amd.com> Co-authored-by: Michael Goin <mgoin64@gmail.com>
Matthias Gehre committed
Mar 30, 2026 at 16:30 UTC
e8b055a5ac1d16285545087d60f898085a485959
2 files changed
+88
-1
tests/quantization/test_compressed_tensors.py
+78
-1
@@ -5,13 +5,20 @@
5
Run `pytest tests/quantization/test_compressed_tensors.py`.
6
"""
7
8
+from unittest.mock import Mock
9
+
10
import pytest
11
import torch
10
-from compressed_tensors.quantization import QuantizationType
12
+from compressed_tensors.quantization import (
13
+ QuantizationArgs,
14
+ QuantizationStrategy,
15
+ QuantizationType,
16
+)
17
18
from tests.models.utils import check_logprobs_close
19
from vllm.model_executor.layers.fused_moe import UnquantizedFusedMoEMethod
20
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
21
+ CompressedTensorsConfig,
22
CompressedTensorsLinearMethod,
23
CompressedTensorsW4A4Fp4,
24
CompressedTensorsW4A8Fp8,
@@ -26,6 +33,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import W8A8BlockFp8
33
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
34
cutlass_fp4_supported,
35
)
36
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
37
from vllm.platforms import current_platform
38
from vllm.v1.attention.backends.fa_utils import get_flash_attn_version
39
@@ -558,3 +566,72 @@ def test_w4a16_moe_torch_compile(vllm_runner):
566
) as llm:
567
output = llm.generate_greedy("Hi", max_tokens=1)
568
assert output
569
+
570
+
571
+def _make_ct_config(*, target: str = "Linear") -> CompressedTensorsConfig:
572
+ """Build a minimal CompressedTensorsConfig with INT8 channel quant."""
573
+ weight_quant = QuantizationArgs(
574
+ num_bits=8,
575
+ type=QuantizationType.INT,
576
+ strategy=QuantizationStrategy.CHANNEL,
577
+ symmetric=True,
578
+ dynamic=False,
579
+ )
580
+ return CompressedTensorsConfig(
581
+ target_scheme_map={
582
+ target: {
583
+ "weights": weight_quant,
584
+ "input_activations": None,
585
+ "format": "pack-quantized",
586
+ }
587
+ },
588
+ ignore=[],
589
+ quant_format="pack-quantized",
590
+ sparsity_scheme_map={},
591
+ sparsity_ignore_list=[],
592
+ )
593
+
594
+
595
+def test_get_quant_method_returns_linear_method_for_parallel_lm_head():
596
+ """ParallelLMHead whose name matches a target must get a quantised method."""
597
+ config = _make_ct_config(target="re:.*lm_head")
598
+ mock_lm_head = Mock(spec=ParallelLMHead)
599
+ mock_lm_head.__class__ = ParallelLMHead
600
+
601
+ method = config.get_quant_method(mock_lm_head, prefix="model.lm_head")
602
+
603
+ assert isinstance(method, CompressedTensorsLinearMethod), (
604
+ f"Expected CompressedTensorsLinearMethod, got {type(method).__name__}"
605
+ )
606
+
607
+
608
+def test_get_quant_method_returns_none_for_ignored_parallel_lm_head():
609
+ """ParallelLMHead on the ignore list should be left unquantized (None)."""
610
+ config = _make_ct_config(target="re:.*lm_head")
611
+ config.ignore = ["re:.*lm_head"]
612
+ mock_lm_head = Mock(spec=ParallelLMHead)
613
+ mock_lm_head.__class__ = ParallelLMHead
614
+
615
+ method = config.get_quant_method(mock_lm_head, prefix="model.lm_head")
616
+
617
+ assert method is None, (
618
+ f"Expected None for ignored ParallelLMHead, got {type(method).__name__}"
619
+ )
620
+
621
+
622
+def test_get_quant_method_returns_none_for_unmatched_parallel_lm_head():
623
+ """ParallelLMHead with target='Linear' (typical real model) must not crash.
624
+
625
+ Most compressed-tensors models only target 'Linear'. ParallelLMHead does
626
+ not match that target, so get_quant_method should return None (unquantized)
627
+ instead of raising ValueError.
628
+ """
629
+ config = _make_ct_config(target="Linear")
630
+ mock_lm_head = Mock(spec=ParallelLMHead)
631
+ mock_lm_head.__class__ = ParallelLMHead
632
+
633
+ method = config.get_quant_method(mock_lm_head, prefix="model.lm_head")
634
+
635
+ assert method is None, (
636
+ f"Expected None for unmatched ParallelLMHead, got {type(method).__name__}"
637
+ )
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+10
@@ -62,6 +62,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
62
should_ignore_layer,
63
)
64
from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod
65
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
66
from vllm.platforms import current_platform
67
68
if TYPE_CHECKING:
@@ -179,6 +180,15 @@ class CompressedTensorsConfig(QuantizationConfig):
180
else:
181
return quant_method
182
183
+ if isinstance(layer, ParallelLMHead):
184
+ try:
185
+ quant_scheme = self.get_scheme(layer=layer, layer_name=prefix)
186
+ except ValueError:
187
+ quant_scheme = None
188
+ if quant_scheme is not None:
189
+ layer.scheme = quant_scheme
190
+ return CompressedTensorsLinearMethod(self)
191
+
192
if isinstance(layer, Attention):
193
return CompressedTensorsKVCacheMethod(self)
194
if isinstance(layer, FusedMoE):