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 12, 12:15 UTC
[Bugfix] Make compressed-tensors MoEs respect ignored layers (#28878)
Signed-off-by: HDCharles <charlesdavidhernandez@gmail.com>
HDCharles committed
Nov 26, 2025 at 21:35 UTC
df01eda4dc570dbf9aa45dd196e288d13f427fab
5 files changed
+133
-52
.buildkite/test-pipeline.yaml
+1
@@ -632,6 +632,7 @@ steps:
632
# we can only upgrade after this is resolved
633
# TODO(jerryzh168): resolve the above comment
634
- uv pip install --system torchao==0.13.0 --index-url https://download.pytorch.org/whl/cu129
635
+ - uv pip install --system conch-triton-kernels
636
- VLLM_TEST_FORCE_LOAD_FORMAT=auto pytest -v -s quantization/ --ignore quantization/test_blackwell_moe.py
637
638
- label: LM Eval Small Models # 53min
tests/quantization/test_compressed_tensors.py
+48
@@ -10,6 +10,7 @@ import torch
10
from compressed_tensors.quantization import QuantizationType
11
12
from tests.models.utils import check_logprobs_close
13
+from vllm.model_executor.layers.fused_moe import UnquantizedFusedMoEMethod
14
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
15
CompressedTensors24,
16
CompressedTensorsLinearMethod,
@@ -767,3 +768,50 @@ def test_compressed_tensors_fp8_block_enabled(vllm_runner):
768
769
output = llm.generate_greedy("Hello my name is", max_tokens=4)
770
assert output
771
+
772
+
773
+@pytest.mark.skipif(
774
+ not current_platform.is_cuda(),
775
+ reason="This test is not for non-CUDA platforms",
776
+)
777
+def test_compressed_tensors_moe_ignore_with_model(vllm_runner):
778
+ """
779
+ Integration test for MoE layer ignore functionality with a real model.
780
+
781
+ This test would verify that when loading a compressed-tensors quantized
782
+ MoE model where some MoE layers are in the ignore list, those layers
783
+ use UnquantizedFusedMoEMethod while non-ignored layers use the
784
+ quantized method.
785
+
786
+ Expected model structure:
787
+ - Compressed-tensors quantized MoE model (e.g., Mixtral-based)
788
+ - Config with ignore list containing specific MoE layers
789
+ - Multiple MoE layers where some are quantized and some are not
790
+ """
791
+
792
+ # model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only" # CT 12.3
793
+ model_path = "nm-testing/tinysmokeqwen3moe-W4A16-first-only-CTstable" # CT 12.2
794
+
795
+ with vllm_runner(model_path, enforce_eager=True) as llm:
796
+
797
+ def check_model(model):
798
+ from vllm.model_executor.layers.fused_moe import FusedMoE
799
+ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
800
+ CompressedTensorsMoEMethod,
801
+ )
802
+
803
+ # Check layer 0 MoE (should be quantized)
804
+ layer_quantized = model.model.layers[0].mlp.experts
805
+ assert isinstance(layer_quantized, FusedMoE)
806
+ assert isinstance(layer_quantized.quant_method, CompressedTensorsMoEMethod)
807
+
808
+ # Check layer 10 MoE (should be unquantized + ignored)
809
+ layer_unquantized = model.model.layers[3].mlp.experts
810
+ assert isinstance(layer_unquantized, FusedMoE)
811
+ assert isinstance(layer_unquantized.quant_method, UnquantizedFusedMoEMethod)
812
+
813
+ llm.apply_model(check_model)
814
+
815
+ # Verify the model can generate output
816
+ output = llm.generate_greedy("Hello, my name is", max_tokens=4)
817
+ assert output
vllm/model_executor/layers/fused_moe/__init__.py
+4
@@ -18,6 +18,9 @@ from vllm.model_executor.layers.fused_moe.modular_kernel import (
18
FusedMoEPrepareAndFinalize,
19
)
20
from vllm.model_executor.layers.fused_moe.shared_fused_moe import SharedFusedMoE
21
+from vllm.model_executor.layers.fused_moe.unquantized_fused_moe_method import (
22
+ UnquantizedFusedMoEMethod,
23
+)
24
from vllm.model_executor.layers.fused_moe.utils import activation_without_mul
25
from vllm.triton_utils import HAS_TRITON
26
@@ -41,6 +44,7 @@ __all__ = [
44
"FusedMoE",
45
"FusedMoEConfig",
46
"FusedMoEMethodBase",
47
+ "UnquantizedFusedMoEMethod",
48
"FusedMoeWeightScaleSupported",
49
"FusedMoEPermuteExpertsUnpermute",
50
"FusedMoEActivationFormat",
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+53
-19
@@ -158,9 +158,23 @@ class CompressedTensorsConfig(QuantizationConfig):
158
if isinstance(layer, Attention):
159
return CompressedTensorsKVCacheMethod(self)
160
if isinstance(layer, FusedMoE):
161
- return CompressedTensorsMoEMethod.get_moe_method(self, layer)
161
+ return CompressedTensorsMoEMethod.get_moe_method(self, layer, prefix)
162
return None
163
164
+ def _add_fused_moe_to_target_scheme_map(self):
165
+ """
166
+ Helper function to update target_scheme_map
167
+ since linear layers get fused into FusedMoE
168
+ targetting 'Linear' needs to also match
169
+ FusedMoE modules.
170
+ """
171
+ if (
172
+ "Linear" not in self.target_scheme_map
173
+ or "FusedMoE" in self.target_scheme_map
174
+ ):
175
+ return
176
+ self.target_scheme_map["FusedMoE"] = self.target_scheme_map["Linear"]
177
+
178
@classmethod
179
def from_config(cls, config: dict[str, Any]) -> "CompressedTensorsConfig":
180
ignore: list[str] = cast(list[str], config.get("ignore", []))
@@ -655,25 +669,13 @@ class CompressedTensorsConfig(QuantizationConfig):
669
to select the CompressedTensorsScheme used for inference.
670
"""
671
658
- # Find the "target" in the compressed-tensors config
659
- # that our layer conforms to.
660
- # TODO (@kylesayrs): support ignore module names with ct matching utils
661
- if should_ignore_layer(
662
- layer_name, ignore=self.ignore, fused_mapping=self.packed_modules_mapping
663
- ):
664
- return None
665
-
666
- # Will be empty for models with only sparsity
667
- weight_quant = input_quant = None
668
- if self.target_scheme_map:
669
- matched_target = find_matched_target(
670
- layer_name=layer_name,
671
- module=layer,
672
- targets=self.target_scheme_map.keys(),
673
- fused_mapping=self.packed_modules_mapping,
674
- )
672
+ # Use the new get_quant_args method to extract QuantizationArgs
673
+ scheme_dict = self.get_scheme_dict(layer, layer_name)
674
676
- scheme_dict = self.target_scheme_map[matched_target]
675
+ weight_quant = None
676
+ input_quant = None
677
+ format = None
678
+ if scheme_dict:
679
weight_quant = scheme_dict.get("weights")
680
input_quant = scheme_dict.get("input_activations")
681
format = scheme_dict.get("format")
@@ -732,6 +734,38 @@ class CompressedTensorsConfig(QuantizationConfig):
734
logger.debug("Using scheme: %s for %s", scheme.__class__.__name__, layer_name)
735
return scheme
736
737
+ def get_scheme_dict(
738
+ self, layer: torch.nn.Module, layer_name: str | None = None
739
+ ) -> dict[str, QuantizationArgs | str | None] | None:
740
+ """
741
+ Extract the QuantizationArgs for a given layer.
742
+
743
+ Returns:
744
+ dict with {
745
+ "weights": QuantizationArgs,
746
+ "input_activations": QuantizationArgs | None,
747
+ "format": str | None
748
+ } | None
749
+ """
750
+ # TODO (@kylesayrs): support ignore module names with ct matching utils
751
+ if should_ignore_layer(
752
+ layer_name, ignore=self.ignore, fused_mapping=self.packed_modules_mapping
753
+ ):
754
+ return None
755
+
756
+ # Will be empty for models with only sparsity
757
+ if self.target_scheme_map:
758
+ matched_target = find_matched_target(
759
+ layer_name=layer_name,
760
+ module=layer,
761
+ targets=self.target_scheme_map.keys(),
762
+ fused_mapping=self.packed_modules_mapping,
763
+ )
764
+
765
+ return self.target_scheme_map[matched_target]
766
+
767
+ return None
768
+
769
def get_cache_scale(self, name: str) -> str | None:
770
"""
771
Check whether the param name matches the format for k/v cache scales
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+27
-33
@@ -23,6 +23,7 @@ from vllm.model_executor.layers.fused_moe import (
23
FusedMoEMethodBase,
24
FusedMoEPermuteExpertsUnpermute,
25
FusedMoeWeightScaleSupported,
26
+ UnquantizedFusedMoEMethod,
27
)
28
from vllm.model_executor.layers.fused_moe.config import (
29
FusedMoEQuantConfig,
@@ -45,9 +46,6 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes.compress
46
WNA16_SUPPORTED_BITS,
47
WNA16_SUPPORTED_TYPES_MAP,
48
)
48
-from vllm.model_executor.layers.quantization.compressed_tensors.utils import (
49
- find_matched_target,
50
-)
49
from vllm.model_executor.layers.quantization.utils import replace_parameter
50
from vllm.model_executor.layers.quantization.utils.flashinfer_fp4_moe import (
51
build_flashinfer_fp4_cutlass_moe_prepare_finalize,
@@ -113,39 +111,35 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
111
def get_moe_method(
112
quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
113
layer: torch.nn.Module,
114
+ prefix: str,
115
) -> "CompressedTensorsMoEMethod":
117
- # TODO: @dsikka: refactor this to use schemes as other kernels
118
- # are supported + check if the layer is being ignored.
119
- # Check if a using "Linear" to select schemes
120
- if "Linear" in quant_config.target_scheme_map:
121
- matched_target = "Linear"
122
- else:
123
- # May have instead defined the linear layers in the fused model
124
-
125
- fused_layers = ["re:.*down_proj.*", "re:.*gate_proj.*", "re:.*up_proj.*"]
126
- current_scheme = None
127
- for fused_layer in fused_layers:
128
- # Check if one of the fused layers are defined in quant_config
129
- matched_target = find_matched_target(
130
- layer_name=fused_layer,
131
- module=layer,
132
- targets=quant_config.target_scheme_map.keys(),
133
- fused_mapping=quant_config.packed_modules_mapping,
134
- )
116
+ # FusedMoE was made by combining multiple Linears so need to
117
+ # make sure quantization config for Linear can target it
118
+ quant_config._add_fused_moe_to_target_scheme_map()
119
+ unfused_names = [
120
+ prefix + proj_name
121
+ for proj_name in [".0.gate_proj", ".0.up_proj", ".0.down_proj"]
122
+ ]
123
+ # TODO: refactor this to use expert_mapping and check all layer numbers
124
+ all_scheme_dicts = [
125
+ quant_config.get_scheme_dict(layer, name) for name in unfused_names
126
+ ]
127
+ scheme_dict = all_scheme_dicts.pop()
128
+
129
+ # multiple schemes found
130
+ if not all([cur_dict == scheme_dict for cur_dict in all_scheme_dicts]):
131
+ raise ValueError(
132
+ "All MoE projections need to have same "
133
+ "quantization scheme but found multiple"
134
+ )
135
136
- # Only valid if down_proj, gate_proj, and up_proj
137
- # are mapped to the same quant scheme in the quant_config
138
- if current_scheme is None:
139
- current_scheme = quant_config.target_scheme_map.get(matched_target)
140
- else:
141
- assert current_scheme == quant_config.target_scheme_map.get(
142
- matched_target
143
- )
136
+ if scheme_dict is None: # ignored layer
137
+ return UnquantizedFusedMoEMethod(layer.moe_config)
138
145
- weight_quant = quant_config.target_scheme_map[matched_target].get("weights")
146
- input_quant = quant_config.target_scheme_map[matched_target].get(
147
- "input_activations"
148
- )
139
+ # TODO: @dsikka: refactor this to use schemes as other kernels
140
+ # are supported + check if the layer is being ignored.
141
+ weight_quant = scheme_dict.get("weights")
142
+ input_quant = scheme_dict.get("input_activations")
143
144
if quant_config._is_wNa16_group_channel(weight_quant, input_quant):
145
# group_size=None means channelwise