| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """End-to-end accuracy test for GPT-OSS model quantization. |
| 4 | |
| 5 | Config: |
| 6 | Task: gsm8k_platinum |
| 7 | Filter: flexible-extract |
| 8 | n-shot: 5 |
| 9 | Metric: exact_match |
| 10 | |
| 11 | Run: pytest tests/models/quantization/test_gpt_oss.py |
| 12 | """ |
| 13 | |
| 14 | import importlib.metadata |
| 15 | import importlib.util |
| 16 | from dataclasses import dataclass |
| 17 | |
| 18 | import huggingface_hub |
| 19 | import lm_eval |
| 20 | import pytest |
| 21 | import torch |
| 22 | from packaging import version |
| 23 | |
| 24 | from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops |
| 25 | from vllm.config import get_current_vllm_config |
| 26 | from vllm.forward_context import set_forward_context |
| 27 | from vllm.model_executor.layers.fused_moe.config import ( |
| 28 | mxfp4_w4a8_moe_quant_config, |
| 29 | mxfp4_w4a16_moe_quant_config, |
| 30 | ) |
| 31 | from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a8_moe import ( |
| 32 | aiter_triton_kernel_w4a8_moe_forward, |
| 33 | ) |
| 34 | from vllm.model_executor.layers.fused_moe.experts.aiter_mxfp4_w4a16_moe import ( |
| 35 | aiter_triton_kernel_w4a16_moe_forward, |
| 36 | ) |
| 37 | from vllm.model_executor.layers.fused_moe.oracle.mxfp4 import ( |
| 38 | Mxfp4MoeBackend, |
| 39 | convert_gpt_oss_weight_to_mxfp4_moe_kernel_format, |
| 40 | ) |
| 41 | from vllm.platforms import current_platform |
| 42 | from vllm.transformers_utils.repo_utils import hf_api |
| 43 | from vllm.utils.torch_utils import set_random_seed |
| 44 | |
| 45 | if current_platform.is_rocm(): |
| 46 | from vllm.platforms.rocm import on_gfx950, on_gfx1250 |
| 47 | else: |
| 48 | |
| 49 | def on_gfx950() -> bool: |
| 50 | return False |
| 51 | |
| 52 | def on_gfx1250() -> bool: |
| 53 | return False |
| 54 | |
| 55 | |
| 56 | MODEL_ACCURACIES = { |
| 57 | # Full quantization: attention linears and MoE linears |
| 58 | "amd/gpt-oss-20b-WFP8-AFP8-KVFP8": 0.89, |
| 59 | # MoE linears only quantization |
| 60 | "amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-FP8-KV-FP8": 0.89, |
| 61 | # MoE linears only quantization |
| 62 | # "amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-MXFP4-KV-FP8": 0.90, |
| 63 | } |
| 64 | |
| 65 | QUARK_MXFP4_AVAILABLE = importlib.util.find_spec("quark") is not None and version.parse( |
| 66 | importlib.metadata.version("amd-quark") |
| 67 | ) >= version.parse("0.9.0") |
| 68 | |
| 69 | |
| 70 | def has_huggingface_access(repo): |
| 71 | try: |
| 72 | hf_api().list_repo_refs(repo) |
| 73 | return True |
| 74 | except huggingface_hub.errors.RepositoryNotFoundError: |
| 75 | return False |
| 76 | |
| 77 | |
| 78 | HF_HUB_AMD_ORG_ACCESS = all( |
| 79 | [has_huggingface_access(model_name) for model_name in MODEL_ACCURACIES] |
| 80 | ) |
| 81 | |
| 82 | |
| 83 | @dataclass |
| 84 | class ModelCase: |
| 85 | model_id: str |
| 86 | tp: int |
| 87 | |
| 88 | |
| 89 | @dataclass |
| 90 | class EvaluationConfig: |
| 91 | model_name: str |
| 92 | |
| 93 | def get_model_args(self, tp_size: int): |
| 94 | return { |
| 95 | "pretrained": self.model_name, |
| 96 | "chat_template_args": {"reasoning_effort": "low"}, |
| 97 | "enable_thinking": True, |
| 98 | "think_end_token": "200008", |
| 99 | "tensor_parallel_size": tp_size, |
| 100 | "dtype": "auto", |
| 101 | "gpu_memory_utilization": 0.95, |
| 102 | "trust_remote_code": False, |
| 103 | "enable_prefix_caching": False, |
| 104 | "enforce_eager": False, |
| 105 | } |
| 106 | |
| 107 | |
| 108 | @pytest.mark.skipif(not QUARK_MXFP4_AVAILABLE, reason="amd-quark>=0.9 is not available") |
| 109 | @pytest.mark.skipif( |
| 110 | not HF_HUB_AMD_ORG_ACCESS, |
| 111 | reason="Read access to huggingface.co/amd is required for this test.", |
| 112 | ) |
| 113 | @pytest.mark.parametrize("tp_size", [1, 2, 4, 8]) |
| 114 | @pytest.mark.parametrize("model_name, expected_accuracy", MODEL_ACCURACIES.items()) |
| 115 | def test_gpt_oss_attention_quantization( |
| 116 | model_name: str, |
| 117 | tp_size: int, |
| 118 | expected_accuracy: float, |
| 119 | monkeypatch: pytest.MonkeyPatch, |
| 120 | ): |
| 121 | if tp_size > current_platform.device_count(): |
| 122 | pytest.skip("Not enough GPUs to run this test case") |
| 123 | |
| 124 | if "amd/gpt-oss-20b-MoE-Quant-W-MXFP4-A-FP8-KV-FP8" in model_name and on_gfx950(): |
| 125 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 126 | |
| 127 | model_args = EvaluationConfig(model_name).get_model_args(tp_size) |
| 128 | |
| 129 | extra_run_kwargs = { |
| 130 | "gen_kwargs": {"max_gen_toks": 8000}, |
| 131 | "apply_chat_template": True, |
| 132 | "fewshot_as_multiturn": True, |
| 133 | "num_fewshot": 5, |
| 134 | } |
| 135 | |
| 136 | lm_eval_out = lm_eval.simple_evaluate( |
| 137 | model="vllm", |
| 138 | model_args=model_args, |
| 139 | tasks="gsm8k_platinum", |
| 140 | batch_size="auto", |
| 141 | **extra_run_kwargs, |
| 142 | ) |
| 143 | measured_accuracy = float( |
| 144 | lm_eval_out["results"]["gsm8k_platinum"]["exact_match,flexible-extract"] |
| 145 | ) |
| 146 | |
| 147 | rtol = 0.02 |
| 148 | assert measured_accuracy >= expected_accuracy - rtol, ( |
| 149 | f"Accuracy {measured_accuracy:.4f} is below threshold " |
| 150 | f"{expected_accuracy - rtol:.4f} (expected >= {expected_accuracy} - {rtol})" |
| 151 | ) |
| 152 | |
| 153 | |
| 154 | @pytest.mark.skipif( |
| 155 | not (on_gfx950() or on_gfx1250()), |
| 156 | reason="AITER MXFP4 MoE requires gfx950 or gfx1250", |
| 157 | ) |
| 158 | @pytest.mark.skipif(not is_aiter_found_and_supported(), reason="aiter is not installed") |
| 159 | @pytest.mark.skipif(torch.accelerator.device_count() == 0, reason="no gpu available") |
| 160 | @pytest.mark.parametrize( |
| 161 | "mxfp4_backend", |
| 162 | [Mxfp4MoeBackend.AITER_MXFP4_FP8, Mxfp4MoeBackend.AITER_MXFP4_BF16], |
| 163 | ids=["w4a8", "w4a16"], |
| 164 | ) |
| 165 | @pytest.mark.parametrize("num_experts", [32, 64, 128]) |
| 166 | def test_aiter_mxfp4_moe_ignores_padded_rows( |
| 167 | mxfp4_backend: Mxfp4MoeBackend, |
| 168 | num_experts: int, |
| 169 | monkeypatch: pytest.MonkeyPatch, |
| 170 | dist_init, |
| 171 | ) -> None: |
| 172 | """Garbage in cudagraph padding rows must not reach the unpadded outputs. |
| 173 | |
| 174 | A cudagraph replay of a size-`TOKENS_PADDED` graph driven by |
| 175 | `TOKENS_UNPADDED` real tokens leaves the trailing padding row containing |
| 176 | -inf/inf/nan. |
| 177 | """ |
| 178 | TOKENS_PADDED = 8 |
| 179 | TOKENS_UNPADDED = 7 |
| 180 | |
| 181 | # `rocm_aiter_ops` snapshots the environment at import time, so setting the |
| 182 | # env var alone is not enough. |
| 183 | monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") |
| 184 | rocm_aiter_ops.refresh_env_variables() |
| 185 | |
| 186 | set_random_seed(0) |
| 187 | device = torch.device("cuda") |
| 188 | hidden_size = 2048 |
| 189 | intermediate_size = 512 |
| 190 | |
| 191 | num_experts = 32 |
| 192 | topk = 4 |
| 193 | is_w4a8 = mxfp4_backend is Mxfp4MoeBackend.AITER_MXFP4_FP8 |
| 194 | |
| 195 | # TODO: clean this in oracle/mxfp4. |
| 196 | # `convert_gpt_oss_weight_to_mxfp4_moe_kernel_format` maps |
| 197 | # AITER_MXFP4_BF16 onto the CK layout (`shuffle_weight_a16w4`), which the |
| 198 | # triton monolithic cannot consume -- it wants the `_swizzle_mxfp4` |
| 199 | # PrecisionConfig that the TRITON branch of the same converter produces. |
| 200 | weight_backend = ( |
| 201 | Mxfp4MoeBackend.AITER_MXFP4_FP8 if is_w4a8 else Mxfp4MoeBackend.TRITON |
| 202 | ) |
| 203 | |
| 204 | layer = torch.nn.Module() |
| 205 | layer.w13_weight = torch.randint( |
| 206 | 0, |
| 207 | 256, |
| 208 | (num_experts, 2 * intermediate_size, hidden_size // 2), |
| 209 | dtype=torch.uint8, |
| 210 | device=device, |
| 211 | ) |
| 212 | layer.w2_weight = torch.randint( |
| 213 | 0, |
| 214 | 256, |
| 215 | (num_experts, hidden_size, intermediate_size // 2), |
| 216 | dtype=torch.uint8, |
| 217 | device=device, |
| 218 | ) |
| 219 | |
| 220 | # Keep the exponents near 127 (2**0) so the dequantized weights stay in a sane |
| 221 | # range and 255 (e8m0 NaN) is never hit. |
| 222 | layer.w13_weight_scale = torch.randint( |
| 223 | 124, |
| 224 | 131, |
| 225 | (num_experts, 2 * intermediate_size, hidden_size // 32), |
| 226 | dtype=torch.uint8, |
| 227 | device=device, |
| 228 | ) |
| 229 | layer.w2_weight_scale = torch.randint( |
| 230 | 124, |
| 231 | 131, |
| 232 | (num_experts, hidden_size, intermediate_size // 32), |
| 233 | dtype=torch.uint8, |
| 234 | device=device, |
| 235 | ) |
| 236 | layer.w13_bias = torch.randn( |
| 237 | (num_experts, 2 * intermediate_size), dtype=torch.float32, device=device |
| 238 | ) |
| 239 | layer.w2_bias = torch.randn( |
| 240 | (num_experts, hidden_size), dtype=torch.float32, device=device |
| 241 | ) |
| 242 | # Read off `layer`, not off the arguments, by the AITER_MXFP4_FP8 branch. |
| 243 | layer.w13_input_scale = torch.full( |
| 244 | (num_experts,), 0.1, dtype=torch.float32, device=device |
| 245 | ) |
| 246 | layer.w2_input_scale = torch.full( |
| 247 | (num_experts,), 0.1, dtype=torch.float32, device=device |
| 248 | ) |
| 249 | |
| 250 | ( |
| 251 | w13_weight, |
| 252 | w2_weight, |
| 253 | w13_scale, |
| 254 | w2_scale, |
| 255 | w13_bias, |
| 256 | w2_bias, |
| 257 | ) = convert_gpt_oss_weight_to_mxfp4_moe_kernel_format( |
| 258 | weight_backend, |
| 259 | layer, |
| 260 | layer.w13_weight, |
| 261 | layer.w2_weight, |
| 262 | layer.w13_weight_scale, |
| 263 | layer.w2_weight_scale, |
| 264 | w13_bias=layer.w13_bias, |
| 265 | w2_bias=layer.w2_bias, |
| 266 | ) |
| 267 | |
| 268 | quant_config_factory = ( |
| 269 | mxfp4_w4a8_moe_quant_config if is_w4a8 else mxfp4_w4a16_moe_quant_config |
| 270 | ) |
| 271 | quant_config = quant_config_factory( |
| 272 | w1_scale=w13_scale, |
| 273 | w2_scale=w2_scale, |
| 274 | w1_bias=w13_bias, |
| 275 | w2_bias=w2_bias, |
| 276 | ) |
| 277 | |
| 278 | hidden_states = torch.randn( |
| 279 | (TOKENS_PADDED, hidden_size), dtype=torch.bfloat16, device=device |
| 280 | ) |
| 281 | gating_output = torch.randn( |
| 282 | (TOKENS_PADDED, num_experts), dtype=torch.bfloat16, device=device |
| 283 | ) |
| 284 | |
| 285 | hidden_states[TOKENS_UNPADDED:, 0] = float("inf") |
| 286 | hidden_states[TOKENS_UNPADDED:, 1] = float("-inf") |
| 287 | hidden_states[TOKENS_UNPADDED:, 2] = float("nan") |
| 288 | |
| 289 | gating_output[TOKENS_UNPADDED:, :] = float("-inf") |
| 290 | gating_output[TOKENS_UNPADDED:, 0] = float("inf") |
| 291 | gating_output[TOKENS_UNPADDED:, 1] = float("nan") |
| 292 | gating_output[TOKENS_UNPADDED:, 2] = -float("nan") |
| 293 | |
| 294 | is_padding = torch.zeros(TOKENS_PADDED, dtype=torch.bool, device=device) |
| 295 | is_padding[TOKENS_UNPADDED:] = True |
| 296 | |
| 297 | # Poison memory: after `del blocks`, further `torch.empty` in routing with wrongful |
| 298 | # expert ids may read into these memory sections resulting in |
| 299 | # potential memory access fault. |
| 300 | blocks = [] |
| 301 | for numel in (16, 32, 64, 128, 256, 512, 1024, 4096, 1 << 14, 1 << 16, 1 << 20): |
| 302 | blocks.append(torch.full((numel,), 100e7, dtype=torch.int32, device=device)) |
| 303 | # return blocks |
| 304 | del blocks |
| 305 | torch.accelerator.synchronize() |
| 306 | |
| 307 | with set_forward_context(None, get_current_vllm_config(), is_padding=is_padding): |
| 308 | if is_w4a8: |
| 309 | # AITER_MXFP4_FP8 routes through the triton monolithic expert, which |
| 310 | # does its own top-k inside `aiter.ops.triton.moe.moe_routing`. |
| 311 | output = aiter_triton_kernel_w4a8_moe_forward( |
| 312 | hidden_states=hidden_states, |
| 313 | w1=w13_weight, |
| 314 | w2=w2_weight, |
| 315 | gating_output=gating_output, |
| 316 | topk=topk, |
| 317 | renormalize=True, |
| 318 | quant_config=quant_config, |
| 319 | global_num_experts=num_experts, |
| 320 | unpadded_N_w1=2 * intermediate_size, |
| 321 | unpadded_K_w1=hidden_size, |
| 322 | unpadded_N_w2=hidden_size, |
| 323 | unpadded_K_w2=intermediate_size, |
| 324 | ) |
| 325 | else: |
| 326 | output = aiter_triton_kernel_w4a16_moe_forward( |
| 327 | hidden_states=hidden_states, |
| 328 | w1=w13_weight, |
| 329 | w2=w2_weight, |
| 330 | gating_output=gating_output, |
| 331 | topk=topk, |
| 332 | renormalize=True, |
| 333 | quant_config=quant_config, |
| 334 | global_num_experts=num_experts, |
| 335 | unpadded_N_w1=2 * intermediate_size, |
| 336 | unpadded_K_w1=hidden_size, |
| 337 | unpadded_N_w2=hidden_size, |
| 338 | unpadded_K_w2=intermediate_size, |
| 339 | ) |
| 340 | |
| 341 | assert torch.isfinite(output[:TOKENS_UNPADDED]).all() |