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 13, 12:00 UTC
[Quantization] fix attention quantization of gpt_oss model (#27334)
Signed-off-by: xuebwang-amd <xuebwang@amd.com>
xuebwang-amd committed
Nov 12, 2025 at 01:06 UTC
5a1271d83a65be5ed8dc3e4c990ed42074197db3
3 files changed
+101
-4
tests/models/quantization/test_gpt_oss_attn_quantization.py
new
+80
@@ -0,0 +1,80 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Test attention quantization of gpt-oss model.
4
+The qkv_proj and o_proj in self_attention can be either quantized or excluded.
5
+
6
+Run `pytest tests/models/quantization/test_gpt_oss_attn_quantization.py`.
7
+
8
+"""
9
+
10
+import importlib
11
+import importlib.metadata
12
+from dataclasses import dataclass
13
+
14
+import huggingface_hub
15
+import lm_eval
16
+import pytest
17
+from packaging import version
18
+
19
+MODEL_NAMES = ["amd/gpt-oss-20b-customized-attention-quantization"]
20
+
21
+QUARK_MXFP4_AVAILABLE = importlib.util.find_spec("quark") is not None and version.parse(
22
+ importlib.metadata.version("amd-quark")
23
+) >= version.parse("0.8.99")
24
+
25
+
26
+def has_huggingface_access(repo):
27
+ try:
28
+ huggingface_hub.list_repo_refs(repo)
29
+ return True
30
+ except huggingface_hub.errors.RepositoryNotFoundError:
31
+ return False
32
+
33
+
34
+HF_HUB_AMD_ORG_ACCESS = all(
35
+ [has_huggingface_access(model_name) for model_name in MODEL_NAMES]
36
+)
37
+
38
+
39
+@dataclass
40
+class ModelCase:
41
+ model_id: str
42
+ tp: int
43
+
44
+
45
+@dataclass
46
+class EvaluationConfig:
47
+ model_name: str
48
+
49
+ def get_model_args(self) -> str:
50
+ return (
51
+ f"pretrained={self.model_name},"
52
+ "tensor_parallel_size=4,dtype=auto,gpu_memory_utilization=0.9,trust_remote_code=False"
53
+ )
54
+
55
+
56
+EXPECTED_ACCURACIES = {"arc_challenge": 0.20}
57
+
58
+
59
+@pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available")
60
+@pytest.mark.skipif(
61
+ not HF_HUB_AMD_ORG_ACCESS,
62
+ reason="Read access to huggingface.co/amd is required for this test.",
63
+)
64
+@pytest.mark.parametrize("model_name", MODEL_NAMES)
65
+@pytest.mark.parametrize("task_name, expected_accuracy", EXPECTED_ACCURACIES.items())
66
+def test_gpt_oss_attention_quantization(
67
+ model_name: str, task_name: str, expected_accuracy: float
68
+):
69
+ measured_accuracy = lm_eval.simple_evaluate(
70
+ model="vllm",
71
+ model_args=EvaluationConfig(model_name).get_model_args(),
72
+ tasks=task_name,
73
+ batch_size="auto",
74
+ )["results"][task_name]["acc,none"]
75
+
76
+ rtol = 0.05
77
+ assert (
78
+ measured_accuracy - rtol < expected_accuracy
79
+ and measured_accuracy + rtol > expected_accuracy
80
+ ), f"Expected: {expected_accuracy} | Measured: {measured_accuracy}"
vllm/model_executor/layers/quantization/mxfp4.py
+13
-2
@@ -190,14 +190,25 @@ class Mxfp4Config(QuantizationConfig):
190
fused_mapping=self.packed_modules_mapping,
191
):
192
return UnquantizedLinearMethod()
193
- raise NotImplementedError("Mxfp4 linear layer is not implemented")
193
+ # TODO: Add support for MXFP4 Linear Method.
194
+ # MXFP4 LinearMethod is available in AMD-Quark, refer to that implementation
195
+ # if you are interested in enabling MXFP4 here.
196
+ logger.warning_once(
197
+ "MXFP4 linear layer is not implemented - falling back to "
198
+ "UnquantizedLinearMethod."
199
+ )
200
+ return UnquantizedLinearMethod()
201
elif isinstance(layer, FusedMoE):
202
if current_platform.is_xpu():
203
return IpexMxfp4MoEMethod(layer.moe_config)
204
else:
205
return Mxfp4MoEMethod(layer.moe_config)
206
elif isinstance(layer, Attention):
200
- raise NotImplementedError("Mxfp4 attention layer is not implemented")
207
+ # TODO: Add support for MXFP4 Attention.
208
+ logger.warning_once(
209
+ "MXFP4 attention layer is not implemented. "
210
+ "Skipping quantization for this layer."
211
+ )
212
return None
213
214
vllm/model_executor/models/gpt_oss.py
+8
-2
@@ -198,6 +198,7 @@ class TransformerBlock(torch.nn.Module):
198
def __init__(
199
self,
200
vllm_config: VllmConfig,
201
+ quant_config: QuantizationConfig,
202
prefix: str = "",
203
):
204
super().__init__()
@@ -207,7 +208,10 @@ class TransformerBlock(torch.nn.Module):
208
209
self.layer_idx = extract_layer_index(prefix)
210
self.attn = OAIAttention(
210
- config, prefix=f"{prefix}.attn", cache_config=cache_config
211
+ config,
212
+ prefix=f"{prefix}.attn",
213
+ quant_config=quant_config,
214
+ cache_config=cache_config,
215
)
216
self.mlp = MLPBlock(vllm_config, self.layer_idx, prefix=f"{prefix}.mlp")
217
self.input_layernorm = RMSNorm(config.hidden_size, eps=1e-5)
@@ -243,6 +247,7 @@ class GptOssModel(nn.Module):
247
):
248
super().__init__()
249
self.config = vllm_config.model_config.hf_config
250
+ self.quant_config = vllm_config.quant_config
251
self.parallel_config = vllm_config.parallel_config
252
self.config.hidden_size = self.config.hidden_size
253
self.embedding = VocabParallelEmbedding(
@@ -254,6 +259,7 @@ class GptOssModel(nn.Module):
259
lambda prefix: TransformerBlock(
260
vllm_config,
261
prefix=prefix,
262
+ quant_config=self.quant_config,
263
),
264
prefix=f"{prefix}.layers",
265
)
@@ -645,7 +651,7 @@ class GptOssModel(nn.Module):
651
652
653
class GptOssForCausalLM(nn.Module, SupportsPP, SupportsEagle3, SupportsLoRA):
648
- packed_modules_mapping = {"qkv": ["q_proj", "k_proj", "v_proj"]}
654
+ packed_modules_mapping = {"qkv_proj": ["q_proj", "k_proj", "v_proj"]}
655
656
hf_to_vllm_mapper = WeightsMapper(
657
orig_to_new_substr={