[CI] Speed up quantization test group (#53291)
Signed-off-by: Felix Marty <Felix.Marty@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
fxmarty-amd committed
Sep 1, 2026 at 20:15 UTC
d1c15e589d4c92136eea18773570a927d8795718
20 files changed
+838
-739
tests/evals/gsm8k/configs/Qwen-1.5-MOE-W-MXFP4-A-MXFP6.yaml
new
+5
@@ -0,0 +1,5 @@
1
+model_name: "fxmarty/qwen1.5_moe_a2.7b_chat_w_fp4_a_fp6_e2m3"
2
+accuracy_threshold: 0.4405
3
+num_questions: 50
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 4096"
tests/evals/gsm8k/configs/Qwen-1.5-MOE-W-MXFP6-A-MXFP6.yaml
new
+5
@@ -0,0 +1,5 @@
1
+model_name: "fxmarty/qwen1.5_moe_a2.7b_chat_w_fp6_e3m2_a_fp6_e3m2"
2
+accuracy_threshold: 0.5049
3
+num_questions: 50
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 4096"
tests/evals/gsm8k/configs/Qwen3-30B-A3B-NVFP4-quark.yaml
new
+5
@@ -0,0 +1,5 @@
1
+model_name: "amd-quark/Qwen3-30B-A3B-nvfp4-quark"
2
+accuracy_threshold: 0.89
3
+num_questions: 1319
4
+num_fewshot: 5
5
+server_args: "--enforce-eager --max-model-len 4096"
tests/evals/gsm8k/configs/models-gfx950-large.txt
+3
@@ -6,3 +6,6 @@
6
# startup timeout fix has landed and is verified green.
7
# DeepSeek-V4-Pro-NVFP4.yaml
8
DeepSeek-V4-Flash-NVFP4.yaml
9
+Qwen3-30B-A3B-NVFP4-quark.yaml
10
+Qwen-1.5-MOE-W-MXFP4-A-MXFP6.yaml
11
+Qwen-1.5-MOE-W-MXFP6-A-MXFP6.yaml
\ No newline at end of file
tests/evals/gsm8k/configs/models-mi3xx.txt
+4
-1
@@ -3,5 +3,8 @@ DeepSeek-R1-DP_MI325.yaml
3
DeepSeek-V3.2-TP_MI325.yaml
4
DeepSeek-V3.2-DP_MI325.yaml
5
Qwen3-30B-A3B-NVFP4.yaml
6
+Qwen3-30B-A3B-NVFP4-quark.yaml
7
Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml
7
-Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml
\ No newline at end of file
8
+Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml
9
+Qwen-1.5-MOE-W-MXFP4-A-MXFP6.yaml
10
+Qwen-1.5-MOE-W-MXFP6-A-MXFP6.yaml
\ No newline at end of file
tests/quantization/test_auto_awq.py
+40
-32
@@ -18,7 +18,14 @@ from __future__ import annotations
18
import pytest
19
import torch
20
21
-from tests.quantization.utils import is_quant_method_supported
21
+from tests.quantization.utils import (
22
+ is_quant_method_supported,
23
+ load_model_without_vllm_runner,
24
+)
25
+from vllm.config import set_current_vllm_config
26
+from vllm.forward_context import set_forward_context
27
+from vllm.model_executor.layers.attention import Attention
28
+from vllm.platforms import current_platform
29
30
31
def _get_auto_awq_config_source() -> str:
@@ -186,42 +193,43 @@ AWQ_MODELS = [
193
reason="auto_awq is not supported on this GPU type.",
194
)
195
@pytest.mark.parametrize("model_id", AWQ_MODELS)
189
-def test_auto_awq_quantization_method(vllm_runner, model_id: str, monkeypatch):
196
+def test_auto_awq_quantization_method(
197
+ model_id: str, monkeypatch, dist_init, workspace_init
198
+):
199
"""Test that quantization='auto_awq' loads and runs correctly."""
200
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
192
-
193
- with vllm_runner(
201
+ model, vllm_config = load_model_without_vllm_runner(
202
model_id,
203
dtype=torch.float16,
204
quantization="auto_awq",
197
- max_model_len=2048,
198
- enforce_eager=True,
199
- ) as llm:
200
-
201
- def check_model(model):
202
- from vllm.model_executor.layers.quantization.auto_awq import (
203
- AutoAWQLinearMethod,
204
- AutoAWQMarlinLinearMethod,
205
- )
206
-
207
- for name, submodule in model.named_modules():
208
- if name == "model.layers.0.self_attn.qkv_proj":
209
- # Should use either AutoAWQLinearMethod (Triton) or
210
- # AutoAWQMarlinLinearMethod (Marlin) depending on hardware
211
- assert isinstance(
212
- submodule.quant_method,
213
- (AutoAWQLinearMethod, AutoAWQMarlinLinearMethod),
214
- ), (
215
- f"Expected AutoAWQLinearMethod or AutoAWQMarlinLinearMethod "
216
- f"for {name}, got {type(submodule.quant_method)}"
217
- )
218
- break
219
-
220
- llm.apply_model(check_model)
221
-
222
- outputs = llm.generate_greedy([PROMPT], max_tokens=8)
223
- assert outputs
224
- assert len(outputs[0][1]) > 0
205
+ model_config_kwargs={"max_model_len": 2048},
206
+ )
207
+ target_device = torch.device(current_platform.device_type)
208
+
209
+ from vllm.model_executor.layers.quantization.auto_awq import (
210
+ AutoAWQLinearMethod,
211
+ AutoAWQMarlinLinearMethod,
212
+ )
213
+
214
+ qkv_proj = model.model.layers[0].self_attn.qkv_proj
215
+ assert isinstance(
216
+ qkv_proj.quant_method,
217
+ (AutoAWQLinearMethod, AutoAWQMarlinLinearMethod),
218
+ ), (
219
+ "Expected AutoAWQLinearMethod or AutoAWQMarlinLinearMethod, "
220
+ f"got {type(qkv_proj.quant_method)}"
221
+ )
222
+
223
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
224
+ input_ids = torch.tensor([1, 2, 3, 4], device=target_device)
225
+ positions = torch.arange(input_ids.numel(), device=target_device)
226
+ with (
227
+ set_current_vllm_config(vllm_config),
228
+ set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
229
+ ):
230
+ hidden_states = model(input_ids, positions, None)
231
+ logits = model.compute_logits(hidden_states)
232
+ assert torch.isfinite(logits).all()
233
234
235
def test_auto_awq_config_get_name():
tests/quantization/test_auto_gptq.py
+27
-17
@@ -10,13 +10,20 @@ from types import SimpleNamespace
10
import pytest
11
import torch
12
13
-from tests.quantization.utils import is_quant_method_supported
13
+from tests.quantization.utils import (
14
+ is_quant_method_supported,
15
+ load_model_without_vllm_runner,
16
+)
17
+from vllm.config import set_current_vllm_config
18
+from vllm.forward_context import set_forward_context
19
+from vllm.model_executor.layers.attention import Attention
20
from vllm.model_executor.layers.fused_moe import RoutedExperts
21
from vllm.model_executor.layers.quantization.auto_gptq import (
22
AutoGPTQConfig,
23
AutoGPTQLinearMethod,
24
AutoGPTQMoEMethod,
25
)
26
+from vllm.platforms import current_platform
27
28
PROMPT = "On the surface of Mars, we found"
29
@@ -30,29 +37,32 @@ MODELS = [
37
reason="auto_gptq is not supported on this GPU type.",
38
)
39
@pytest.mark.parametrize("model_id", MODELS)
33
-def test_auto_gptq_quantization_method(vllm_runner, model_id: str, monkeypatch):
40
+def test_auto_gptq_quantization_method(
41
+ model_id: str, monkeypatch, dist_init, workspace_init
42
+):
43
"""Test that quantization='auto_gptq' loads and runs correctly."""
44
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
36
-
37
- with vllm_runner(
45
+ model, vllm_config = load_model_without_vllm_runner(
46
model_id,
47
dtype=torch.float16,
48
quantization="auto_gptq",
41
- max_model_len=2048,
42
- enforce_eager=True,
43
- ) as llm:
44
-
45
- def check_model(model):
46
- for name, submodule in model.named_modules():
47
- if name == "model.layers.0.self_attn.qkv_proj":
48
- assert isinstance(submodule.quant_method, AutoGPTQLinearMethod)
49
- break
49
+ model_config_kwargs={"max_model_len": 2048},
50
+ )
51
51
- llm.apply_model(check_model)
52
+ qkv_proj = model.model.layers[0].self_attn.qkv_proj
53
+ assert isinstance(qkv_proj.quant_method, AutoGPTQLinearMethod)
54
53
- outputs = llm.generate_greedy([PROMPT], max_tokens=8)
54
- assert outputs
55
- assert len(outputs[0][1]) > 0
55
+ target_device = torch.device(current_platform.device_type)
56
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
57
+ input_ids = torch.tensor([1, 2, 3, 4], device=target_device)
58
+ positions = torch.arange(input_ids.numel(), device=target_device)
59
+ with (
60
+ set_current_vllm_config(vllm_config),
61
+ set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
62
+ ):
63
+ hidden_states = model(input_ids, positions, None)
64
+ logits = model.compute_logits(hidden_states)
65
+ assert torch.isfinite(logits).all()
66
67
68
def test_auto_gptq_config_get_name():
tests/quantization/test_compressed_tensors.py
+128
-167
@@ -5,7 +5,6 @@
5
Run `pytest tests/quantization/test_compressed_tensors.py`.
6
"""
7
8
-from contextlib import contextmanager
8
from unittest.mock import Mock
9
10
import pytest
@@ -18,12 +17,17 @@ from compressed_tensors.quantization import (
17
)
18
19
from tests.models.utils import check_logprobs_close
20
+from tests.quantization.utils import load_model_without_vllm_runner
21
+from vllm.config import set_current_vllm_config
22
+from vllm.config.kernel import KernelConfig
23
+from vllm.forward_context import set_forward_context
24
from vllm.model_executor.kernels.linear import (
25
Fp8BlockScaledMMLinearKernel,
26
)
27
from vllm.model_executor.kernels.linear.scaled_mm import (
28
MarlinFP8ScaledMMLinearKernel,
29
)
30
+from vllm.model_executor.layers.attention import Attention
31
from vllm.model_executor.layers.fused_moe import UnquantizedFusedMoEMethod
32
from vllm.model_executor.layers.fused_moe.oracle.fp8 import Fp8MoeBackend
33
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors import ( # noqa: E501
@@ -99,7 +103,7 @@ def enable_pickle(monkeypatch):
103
),
104
],
105
)
102
-def test_compressed_tensors_w8a8_static_setup(vllm_runner, model_args):
106
+def test_compressed_tensors_w8a8_static_setup(model_args, dist_init, workspace_init):
107
model_path, strategy, quant_type, shape_0, is_symmetric = model_args
108
109
if (
@@ -108,55 +112,45 @@ def test_compressed_tensors_w8a8_static_setup(vllm_runner, model_args):
112
):
113
pytest.skip(f"Skip model {model_path} as it is not supported on ROCm.")
114
111
- with vllm_runner(model_path, enforce_eager=True) as llm:
112
-
113
- def check_model(model):
114
- layer = model.model.layers[0]
115
+ model, _ = load_model_without_vllm_runner(model_path)
116
+ layer = model.model.layers[0]
117
116
- qkv_proj = layer.self_attn.qkv_proj
117
- o_proj = layer.self_attn.o_proj
118
- gate_up_proj = layer.mlp.gate_up_proj
119
- down_proj = layer.mlp.down_proj
118
+ qkv_proj = layer.self_attn.qkv_proj
119
+ o_proj = layer.self_attn.o_proj
120
+ gate_up_proj = layer.mlp.gate_up_proj
121
+ down_proj = layer.mlp.down_proj
122
121
- # assert zp for symmetric and asymmetric cases
122
- def zp_valid(zp: torch.Tensor | None):
123
- if is_symmetric:
124
- return zp is None
123
+ def zp_valid(zp: torch.Tensor | None):
124
+ if is_symmetric:
125
+ return zp is None
126
126
- return zp is not None and zp.dtype is torch.int32
127
+ return zp is not None and zp.dtype is torch.int32
128
128
- assert zp_valid(qkv_proj.input_zero_point)
129
- assert zp_valid(o_proj.input_zero_point)
130
- assert zp_valid(gate_up_proj.input_zero_point)
131
- assert zp_valid(down_proj.input_zero_point)
129
+ assert zp_valid(qkv_proj.input_zero_point)
130
+ assert zp_valid(o_proj.input_zero_point)
131
+ assert zp_valid(gate_up_proj.input_zero_point)
132
+ assert zp_valid(down_proj.input_zero_point)
133
133
- assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
134
- assert isinstance(o_proj.quant_method, CompressedTensorsLinearMethod)
135
- assert isinstance(gate_up_proj.quant_method, CompressedTensorsLinearMethod)
136
- assert isinstance(down_proj.quant_method, CompressedTensorsLinearMethod)
137
- assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8)
134
+ assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
135
+ assert isinstance(o_proj.quant_method, CompressedTensorsLinearMethod)
136
+ assert isinstance(gate_up_proj.quant_method, CompressedTensorsLinearMethod)
137
+ assert isinstance(down_proj.quant_method, CompressedTensorsLinearMethod)
138
+ assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8)
139
139
- assert qkv_proj.scheme.strategy == strategy
140
- assert qkv_proj.scheme.is_static_input_scheme
141
- expected_type = torch.int8
142
-
143
- assert qkv_proj.weight.dtype is expected_type
144
- assert o_proj.weight.dtype is expected_type
145
- assert gate_up_proj.weight.dtype is expected_type
146
-
147
- if qkv_proj.scheme.strategy == "tensor":
148
- # Make sure it is a channelwise buffer
149
- # After running process_weights_after_loading
150
- assert len(qkv_proj.weight_scale.shape) == 2
151
- assert qkv_proj.weight_scale.shape[0] == shape_0
152
- assert qkv_proj.weight_scale.shape[1] == 1
153
- assert qkv_proj.weight_scale.dtype is torch.float32
154
- assert qkv_proj.input_scale.dtype is torch.float32
140
+ assert qkv_proj.scheme.strategy == strategy
141
+ assert qkv_proj.scheme.is_static_input_scheme
142
+ expected_type = torch.int8
143
156
- llm.apply_model(check_model)
144
+ assert qkv_proj.weight.dtype is expected_type
145
+ assert o_proj.weight.dtype is expected_type
146
+ assert gate_up_proj.weight.dtype is expected_type
147
158
- output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
159
- assert output
148
+ if qkv_proj.scheme.strategy == "tensor":
149
+ assert len(qkv_proj.weight_scale.shape) == 2
150
+ assert qkv_proj.weight_scale.shape[0] == shape_0
151
+ assert qkv_proj.weight_scale.shape[1] == 1
152
+ assert qkv_proj.weight_scale.dtype is torch.float32
153
+ assert qkv_proj.input_scale.dtype is torch.float32
154
155
156
@pytest.mark.parametrize(
@@ -243,10 +237,7 @@ def test_compressed_tensors_no_enforce_eager(vllm_runner):
237
"use_aiter", [True, False] if current_platform.is_rocm() else [False]
238
)
239
def test_compressed_tensors_w8a8_dynamic_per_token(
246
- vllm_runner,
247
- model_args,
248
- use_aiter,
249
- monkeypatch,
240
+ model_args, use_aiter, monkeypatch, dist_init, workspace_init
241
):
242
model_path, strategy = model_args
243
@@ -262,23 +253,13 @@ def test_compressed_tensors_w8a8_dynamic_per_token(
253
# this will enable VLLM_ROCM_USE_AITER_LINEAR
254
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
255
265
- with vllm_runner(model_path, enforce_eager=True, dtype=torch.float16) as llm:
266
-
267
- def check_model(model):
268
- layer = model.model.layers[0]
269
-
270
- qkv_proj = layer.self_attn.qkv_proj
271
-
272
- assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
273
- assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8)
274
- assert not qkv_proj.scheme.is_static_input_scheme
275
- assert qkv_proj.scheme.strategy == strategy
276
- assert qkv_proj.weight.dtype is torch.int8
277
-
278
- llm.apply_model(check_model)
279
-
280
- output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
281
- assert output
256
+ model, _ = load_model_without_vllm_runner(model_path, dtype=torch.float16)
257
+ qkv_proj = model.model.layers[0].self_attn.qkv_proj
258
+ assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
259
+ assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Int8)
260
+ assert not qkv_proj.scheme.is_static_input_scheme
261
+ assert qkv_proj.scheme.strategy == strategy
262
+ assert qkv_proj.weight.dtype is torch.int8
263
264
265
@pytest.mark.parametrize(
@@ -329,33 +310,36 @@ def test_compressed_tensors_wNa16(vllm_runner, wNa16_args):
310
assert output
311
312
332
-def test_compressed_tensors_fp8(vllm_runner):
313
+def test_compressed_tensors_fp8(monkeypatch, dist_init, workspace_init):
314
model_path = "nm-testing/Meta-Llama-3-8B-FP8-compressed-tensors-test"
334
- with vllm_runner(model_path, enforce_eager=True) as llm:
335
-
336
- def check_model(model):
337
- layer = model.model.layers[0]
338
-
339
- qkv_proj = layer.self_attn.qkv_proj
340
-
341
- assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
342
- assert isinstance(
343
- qkv_proj.scheme,
344
- (CompressedTensorsW8A8Fp8, CompressedTensorsW8A16Fp8),
345
- )
346
-
347
- assert qkv_proj.input_scale.dtype is torch.float32
315
+ model, vllm_config = load_model_without_vllm_runner(model_path)
316
+ target_device = torch.device(current_platform.device_type)
317
+
318
+ layer = model.model.layers[0]
319
+ qkv_proj = layer.self_attn.qkv_proj
320
+ assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
321
+ assert isinstance(
322
+ qkv_proj.scheme,
323
+ (CompressedTensorsW8A8Fp8, CompressedTensorsW8A16Fp8),
324
+ )
325
+ assert qkv_proj.input_scale.dtype is torch.float32
326
349
- if isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8):
350
- assert len(qkv_proj.input_scale.shape) == 0
351
- assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
352
- assert qkv_proj.weight_scale.dtype is torch.float32
353
- assert len(qkv_proj.weight_scale.shape) == 0
327
+ if isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8):
328
+ assert len(qkv_proj.input_scale.shape) == 0
329
+ assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
330
+ assert qkv_proj.weight_scale.dtype is torch.float32
331
+ assert len(qkv_proj.weight_scale.shape) == 0
332
355
- llm.apply_model(check_model)
356
-
357
- output = llm.generate_greedy("Hello my name is", max_tokens=4)
358
- assert output
333
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
334
+ input_ids = torch.tensor([1, 2, 3, 4], device=target_device)
335
+ positions = torch.arange(input_ids.numel(), device=target_device)
336
+ with (
337
+ set_current_vllm_config(vllm_config),
338
+ set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
339
+ ):
340
+ hidden_states = model(input_ids, positions, None)
341
+ logits = model.compute_logits(hidden_states)
342
+ assert torch.isfinite(logits).all()
343
344
345
def test_compressed_tensors_w8a8_fp8_moe_forwards_swiglu_params():
@@ -413,27 +397,6 @@ def test_compressed_tensors_kv_cache_fp8_per_attn_head(vllm_runner):
397
assert output
398
399
416
-@contextmanager
417
-def _nvfp4_marlin_error_context(model, capfd):
418
- is_rocm_and_unsupported = (
419
- model == "nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4A16"
420
- and current_platform.is_rocm()
421
- )
422
-
423
- if is_rocm_and_unsupported:
424
- expected_error = (
425
- "ValueError: Forced NVFP4 kernel MarlinNvFp4LinearKernel is not "
426
- "supported: Marlin FP4 not available"
427
- )
428
- with pytest.raises(RuntimeError, match="Engine core initialization failed"):
429
- yield
430
-
431
- captured = capfd.readouterr()
432
- assert expected_error in captured.out + captured.err
433
- else:
434
- yield
435
-
436
-
400
@pytest.mark.parametrize(
401
"args",
402
[
@@ -441,26 +404,30 @@ def _nvfp4_marlin_error_context(model, capfd):
404
("nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4", False),
405
],
406
)
444
-def test_compressed_tensors_nvfp4(vllm_runner, args, capfd):
445
- model, use_a16 = args
446
- with (
447
- _nvfp4_marlin_error_context(model, capfd),
448
- vllm_runner(model, enforce_eager=True) as llm,
449
- ):
450
-
451
- def check_model(model):
452
- layer = model.model.layers[0]
407
+def test_compressed_tensors_nvfp4(args, dist_init, workspace_init):
408
+ model_path, use_a16 = args
409
+ is_rocm_and_unsupported = (
410
+ model_path == "nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4A16"
411
+ and current_platform.is_rocm()
412
+ )
413
+ if is_rocm_and_unsupported:
414
+ with pytest.raises(
415
+ ValueError,
416
+ match="Forced NVFP4 kernel MarlinNvFp4LinearKernel is not supported",
417
+ ):
418
+ load_model_without_vllm_runner(model_path)
419
+ return
420
454
- qkv_proj = layer.self_attn.qkv_proj
455
- assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
456
- assert isinstance(qkv_proj.scheme, CompressedTensorsW4A4Fp4)
457
- assert qkv_proj.scheme.use_a16 == use_a16
458
- assert qkv_proj.scheme.group_size == 16
421
+ try:
422
+ model, _ = load_model_without_vllm_runner(model_path)
423
+ except ValueError as exc:
424
+ pytest.fail(f"Unexpected NVFP4 loading failure: {exc}")
425
460
- llm.apply_model(check_model)
461
- output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
462
- print(output)
463
- assert output
426
+ qkv_proj = model.model.layers[0].self_attn.qkv_proj
427
+ assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
428
+ assert isinstance(qkv_proj.scheme, CompressedTensorsW4A4Fp4)
429
+ assert qkv_proj.scheme.use_a16 == use_a16
430
+ assert qkv_proj.scheme.group_size == 16
431
432
433
@pytest.mark.skipif(
@@ -533,47 +500,41 @@ def test_compressed_tensors_transforms_perplexity(
500
@pytest.mark.parametrize(
501
"linear_backend", ["auto", "marlin"] if current_platform.is_cuda() else ["auto"]
502
)
536
-def test_compressed_tensors_fp8_block_enabled(vllm_runner, linear_backend):
503
+def test_compressed_tensors_fp8_block_enabled(
504
+ linear_backend, dist_init, workspace_init
505
+):
506
model_path = "RedHatAI/Qwen3-0.6B-FP8-BLOCK"
538
- with vllm_runner(
539
- model_path, enforce_eager=True, linear_backend=linear_backend
540
- ) as llm:
541
- fp8_dtype = current_platform.fp8_dtype()
542
-
543
- def check_model(model):
544
- layer = model.model.layers[0]
545
-
546
- qkv_proj = layer.self_attn.qkv_proj
547
- assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
548
- assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8)
549
- if linear_backend == "marlin":
550
- assert isinstance(
551
- qkv_proj.scheme.fp8_linear, MarlinFP8ScaledMMLinearKernel
552
- )
553
- assert qkv_proj.weight.dtype is torch.int32
554
- assert qkv_proj.weight_scale.dtype is qkv_proj.orig_dtype
555
- else:
556
- assert isinstance(
557
- qkv_proj.scheme.fp8_linear, Fp8BlockScaledMMLinearKernel
558
- )
559
- assert qkv_proj.weight.dtype is fp8_dtype
560
- assert qkv_proj.weight_scale.dtype is torch.float32
561
- assert len(qkv_proj.weight.shape) == 2
562
- assert len(qkv_proj.weight_scale.shape) == 2
563
-
564
- if linear_backend == "auto":
565
- input_quant_op = qkv_proj.scheme.fp8_linear.quant_fp8
566
- assert isinstance(input_quant_op, QuantFP8)
567
- assert input_quant_op._forward_method in (
568
- input_quant_op.forward_cuda,
569
- input_quant_op.forward_hip,
570
- input_quant_op.forward_xpu,
571
- )
572
-
573
- llm.apply_model(check_model)
574
-
575
- output = llm.generate_greedy("Hello my name is", max_tokens=4)
576
- assert output
507
+ model, _ = load_model_without_vllm_runner(
508
+ model_path,
509
+ vllm_config_kwargs={
510
+ "kernel_config": KernelConfig(linear_backend=linear_backend)
511
+ },
512
+ )
513
+ fp8_dtype = current_platform.fp8_dtype()
514
+ layer = model.model.layers[0]
515
+
516
+ qkv_proj = layer.self_attn.qkv_proj
517
+ assert isinstance(qkv_proj.quant_method, CompressedTensorsLinearMethod)
518
+ assert isinstance(qkv_proj.scheme, CompressedTensorsW8A8Fp8)
519
+ if linear_backend == "marlin":
520
+ assert isinstance(qkv_proj.scheme.fp8_linear, MarlinFP8ScaledMMLinearKernel)
521
+ assert qkv_proj.weight.dtype is torch.int32
522
+ assert qkv_proj.weight_scale.dtype is qkv_proj.orig_dtype
523
+ else:
524
+ assert isinstance(qkv_proj.scheme.fp8_linear, Fp8BlockScaledMMLinearKernel)
525
+ assert qkv_proj.weight.dtype is fp8_dtype
526
+ assert qkv_proj.weight_scale.dtype is torch.float32
527
+ assert len(qkv_proj.weight.shape) == 2
528
+ assert len(qkv_proj.weight_scale.shape) == 2
529
+
530
+ if linear_backend == "auto":
531
+ input_quant_op = qkv_proj.scheme.fp8_linear.quant_fp8
532
+ assert isinstance(input_quant_op, QuantFP8)
533
+ assert input_quant_op._forward_method in (
534
+ input_quant_op.forward_cuda,
535
+ input_quant_op.forward_hip,
536
+ input_quant_op.forward_xpu,
537
+ )
538
539
540
@pytest.mark.skipif(
tests/quantization/test_fp8.py
+43
-25
@@ -12,12 +12,20 @@ import pytest
12
import regex as re
13
import torch
14
15
-from tests.quantization.utils import is_quant_method_supported
15
+from tests.quantization.utils import (
16
+ is_quant_method_supported,
17
+ load_model_without_vllm_runner,
18
+)
19
from vllm import _custom_ops as ops
20
+from vllm.config import set_current_vllm_config
21
+from vllm.config.cache import CacheConfig
22
+from vllm.config.kernel import KernelConfig
23
from vllm.config.model import ModelConfig
24
+from vllm.forward_context import set_forward_context
25
from vllm.model_executor.kernels.linear.scaled_mm import (
26
MarlinFP8ScaledMMLinearKernel,
27
)
28
+from vllm.model_executor.layers.attention import Attention
29
from vllm.model_executor.layers.attention.attention import (
30
set_default_quant_scales,
31
)
@@ -116,21 +124,33 @@ def test_static_fp8_moe_input_scales_remain_scalar() -> None:
124
"use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
125
)
126
def test_model_load_and_run(
119
- vllm_runner, model_id: str, force_marlin: bool, use_rocm_aiter: bool, monkeypatch
127
+ model_id: str,
128
+ force_marlin: bool,
129
+ use_rocm_aiter: bool,
130
+ monkeypatch,
131
+ dist_init,
132
+ workspace_init,
133
) -> None:
134
if use_rocm_aiter:
135
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
136
124
- kwargs = {}
125
- if force_marlin:
126
- kwargs["linear_backend"] = "marlin"
127
- kwargs["moe_backend"] = "marlin"
128
-
129
- with vllm_runner(model_id, enforce_eager=True, **kwargs) as llm:
130
- # note: this does not test accuracy, just that we can run through
131
- # see lm-eval tests for accuracy
132
- outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
133
- print(outputs[0][1])
137
+ kernel_config = KernelConfig(
138
+ linear_backend="marlin" if force_marlin else "auto",
139
+ moe_backend="marlin" if force_marlin else "auto",
140
+ )
141
+ model, vllm_config = load_model_without_vllm_runner(
142
+ model_id,
143
+ model_config_kwargs={"hf_overrides": {"num_hidden_layers": 3}},
144
+ vllm_config_kwargs={"kernel_config": kernel_config},
145
+ )
146
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
147
+ input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
148
+ positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
149
+ with (
150
+ set_current_vllm_config(vllm_config),
151
+ set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
152
+ ):
153
+ model(input_ids, positions, None)
154
155
156
@pytest.mark.skipif(
@@ -525,20 +545,18 @@ def test_kv_cache_scale_sync_to_host_copies():
545
not is_quant_method_supported("fp8"),
546
reason="FP8 is not supported on this GPU type.",
547
)
528
-def test_kv_cache_dtype_skip_layers(vllm_runner, monkeypatch):
548
+def test_kv_cache_dtype_skip_layers(monkeypatch, dist_init, workspace_init):
549
"""Test that kv_cache_dtype_skip_layers skips quantization for specified layers."""
550
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
551
532
- with vllm_runner(
552
+ model, _ = load_model_without_vllm_runner(
553
"facebook/opt-125m",
534
- kv_cache_dtype="fp8",
535
- kv_cache_dtype_skip_layers=["0", "2"],
536
- enforce_eager=True,
537
- ) as llm:
538
-
539
- def check_layers(model):
540
- for i, layer in enumerate(model.model.decoder.layers):
541
- expected = "auto" if str(i) in ["0", "2"] else "fp8"
542
- assert layer.self_attn.attn.kv_cache_dtype == expected
543
-
544
- llm.apply_model(check_layers)
554
+ vllm_config_kwargs={
555
+ "cache_config": CacheConfig(
556
+ cache_dtype="fp8", kv_cache_dtype_skip_layers=["0", "2"]
557
+ )
558
+ },
559
+ )
560
+ for i, layer in enumerate(model.model.decoder.layers):
561
+ expected = "auto" if str(i) in ["0", "2"] else "fp8"
562
+ assert layer.self_attn.attn.kv_cache_dtype == expected
tests/quantization/test_fp8_per_channel.py
-34
@@ -67,37 +67,3 @@ def test_scaled_fp8_quant_per_channel_shape() -> None:
67
assert y.dtype == current_platform.fp8_dtype()
68
assert s.shape == (96, 1)
69
assert s.dtype == torch.float32
70
-
71
-
72
-@pytest.mark.skipif(
73
- not is_quant_method_supported("fp8"),
74
- reason="FP8 is not supported on this GPU type.",
75
-)
76
-def test_fp8_per_channel_online_quantization(
77
- vllm_runner,
78
- monkeypatch,
79
-) -> None:
80
- """End-to-end smoke: load `facebook/opt-125m` bf16 with
81
- `quantization='fp8_per_channel'`, check a dense Linear is wrapped by
82
- `Fp8PtpcOnlineLinearMethod`, its weights are fp8 with per-channel
83
- scales (shape `[N, 1]`), and a short greedy generation works.
84
- """
85
- monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
86
-
87
- with vllm_runner(
88
- "facebook/opt-125m",
89
- quantization="fp8_per_channel",
90
- enforce_eager=True,
91
- ) as llm:
92
-
93
- def check_model(model):
94
- fc1 = model.model.decoder.layers[0].fc1
95
- assert isinstance(fc1.quant_method, Fp8PtpcOnlineLinearMethod)
96
- assert fc1.weight.dtype == current_platform.fp8_dtype()
97
- assert fc1.weight_scale.ndim == 2
98
- assert fc1.weight_scale.shape[-1] == 1
99
- assert fc1.input_scale is None
100
-
101
- llm.apply_model(check_model)
102
- outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
103
- print(outputs[0][1])
tests/quantization/test_gptq_dynamic.py
+30
-39
@@ -10,6 +10,7 @@ Note: Only symmetric GPTQ models are supported after consolidation to Marlin.
10
import pytest
11
import torch
12
13
+from tests.quantization.utils import load_model_without_vllm_runner
14
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
15
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQLinearMethod
16
from vllm.model_executor.layers.quantization.utils.gptq_utils import (
@@ -28,48 +29,38 @@ MODELS = [
29
30
31
@pytest.mark.parametrize("model_id", MODELS)
31
-def test_gptq_with_dynamic(vllm_runner, model_id: str, monkeypatch):
32
+def test_gptq_with_dynamic(model_id: str, monkeypatch, dist_init, workspace_init):
33
# `LLM.apply_model` requires pickling a function.
34
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
35
36
linear_method_cls = AutoGPTQLinearMethod
37
37
- with vllm_runner(
38
- model_id, dtype=torch.float16, max_model_len=2048, enforce_eager=True
39
- ) as llm:
38
+ model, _ = load_model_without_vllm_runner(
39
+ model_id,
40
+ dtype=torch.float16,
41
+ model_config_kwargs={
42
+ "max_model_len": 2048,
43
+ "hf_overrides": {"num_hidden_layers": 3},
44
+ },
45
+ )
46
41
- def check_model(model):
42
- for name, submodule in model.named_modules():
43
- if name == "lm_head":
44
- assert isinstance(submodule.quant_method, linear_method_cls)
45
- elif name == "model.layers.0.self_attn.qkv_proj":
46
- # The first layer is quantized using bits=4, group_size=128
47
- # desc_act=True
48
- assert isinstance(submodule.quant_method, linear_method_cls)
49
- config = submodule.quant_method.quant_config
50
- assert config.weight_bits == 4
51
- assert config.group_size == 128
52
- assert config.desc_act
53
- elif name == "model.layers.1.self_attn.qkv_proj":
54
- # The second layer is quantized using bits=8, group_size=32
55
- # desc_act=False
56
- assert isinstance(submodule.quant_method, linear_method_cls)
57
- config = submodule.quant_method.quant_config
58
- assert (
59
- get_dynamic_override(config, layer_name=name, key="bits") == 8
60
- )
61
- assert (
62
- get_dynamic_override(config, layer_name=name, key="group_size")
63
- == 32
64
- )
65
- assert not get_dynamic_override(
66
- config, layer_name=name, key="desc_act"
67
- )
68
- elif (
69
- name == "model.layers.2.self_attn.qkv_proj"
70
- or name == "model.layers.2.mlp.gate_up_proj"
71
- ):
72
- # All other layers (layer index >= 2) are not quantized
73
- assert isinstance(submodule.quant_method, UnquantizedLinearMethod)
74
-
75
- llm.apply_model(check_model)
47
+ for name, submodule in model.named_modules():
48
+ if name == "lm_head":
49
+ assert isinstance(submodule.quant_method, linear_method_cls)
50
+ elif name == "model.layers.0.self_attn.qkv_proj":
51
+ assert isinstance(submodule.quant_method, linear_method_cls)
52
+ config = submodule.quant_method.quant_config
53
+ assert config.weight_bits == 4
54
+ assert config.group_size == 128
55
+ assert config.desc_act
56
+ elif name == "model.layers.1.self_attn.qkv_proj":
57
+ assert isinstance(submodule.quant_method, linear_method_cls)
58
+ config = submodule.quant_method.quant_config
59
+ assert get_dynamic_override(config, layer_name=name, key="bits") == 8
60
+ assert get_dynamic_override(config, layer_name=name, key="group_size") == 32
61
+ assert not get_dynamic_override(config, layer_name=name, key="desc_act")
62
+ elif name in (
63
+ "model.layers.2.self_attn.qkv_proj",
64
+ "model.layers.2.mlp.gate_up_proj",
65
+ ):
66
+ assert isinstance(submodule.quant_method, UnquantizedLinearMethod)
tests/quantization/test_lm_head.py
+16
-20
@@ -8,6 +8,7 @@ Run `pytest tests/quantization/test_quant_lm_head_true.py --forked`.
8
import pytest
9
import torch
10
11
+from tests.quantization.utils import load_model_without_vllm_runner
12
from vllm.model_executor.layers.quantization.auto_gptq import AutoGPTQLinearMethod
13
from vllm.model_executor.layers.vocab_parallel_embedding import (
14
UnquantizedEmbeddingMethod,
@@ -23,29 +24,24 @@ MODELS_QUANT = [
24
25
@pytest.mark.parametrize("model_id, lm_head_quantized", MODELS_QUANT)
26
def test_lm_head(
26
- vllm_runner,
27
model_id: str,
28
lm_head_quantized: bool,
29
monkeypatch,
30
+ dist_init,
31
+ workspace_init,
32
) -> None:
33
# `LLM.apply_model` requires pickling a function.
34
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
33
- with vllm_runner(
34
- model_id, dtype=torch.float16, max_model_len=2048, enforce_eager=True
35
- ) as vllm_model:
36
-
37
- def check_model(model):
38
- lm_head_layer = model.lm_head
39
- if lm_head_quantized:
40
- assert isinstance(
41
- lm_head_layer.quant_method,
42
- AutoGPTQLinearMethod,
43
- )
44
- else:
45
- assert isinstance(
46
- lm_head_layer.quant_method, UnquantizedEmbeddingMethod
47
- )
48
-
49
- vllm_model.apply_model(check_model)
50
-
51
- print(vllm_model.generate_greedy(["Hello my name is"], max_tokens=4)[0][1])
35
+ model, _ = load_model_without_vllm_runner(
36
+ model_id,
37
+ dtype=torch.float16,
38
+ model_config_kwargs={
39
+ "max_model_len": 2048,
40
+ "hf_overrides": {"num_hidden_layers": 3},
41
+ },
42
+ )
43
+ lm_head_layer = model.lm_head
44
+ if lm_head_quantized:
45
+ assert isinstance(lm_head_layer.quant_method, AutoGPTQLinearMethod)
46
+ else:
47
+ assert isinstance(lm_head_layer.quant_method, UnquantizedEmbeddingMethod)
tests/quantization/test_modelopt.py
+87
-99
@@ -12,14 +12,19 @@ from unittest.mock import MagicMock, Mock, patch
12
import pytest
13
import torch
14
15
-from tests.quantization.utils import is_quant_method_supported
15
+from tests.quantization.utils import (
16
+ is_quant_method_supported,
17
+ load_model_without_vllm_runner,
18
+)
19
from vllm.config import VllmConfig, set_current_vllm_config
20
from vllm.config.model import ModelConfig
21
+from vllm.forward_context import set_forward_context
22
from vllm.model_executor.kernels.linear import (
23
FlashInferCuteDslNvFp4W4A16LinearKernel,
24
HummingNvFp4LinearKernel,
25
MarlinNvFp4LinearKernel,
26
)
27
+from vllm.model_executor.layers.attention import Attention
28
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
29
from vllm.model_executor.layers.quantization.modelopt import (
30
ModelOptFp8Config,
@@ -419,124 +424,107 @@ def test_modelopt_fp8_checkpoint_setup(default_vllm_config, vllm_runner):
424
not is_quant_method_supported("modelopt"),
425
reason="ModelOpt FP8 is not supported on this GPU type.",
426
)
422
-def test_modelopt_fp8_pc_pt_checkpoint_setup(default_vllm_config, vllm_runner):
427
+def test_modelopt_fp8_pc_pt_checkpoint_setup(monkeypatch, dist_init, workspace_init):
428
"""Test ModelOpt FP8_PER_CHANNEL_PER_TOKEN checkpoint setup."""
429
model_id = "CedricHwang/qwen2.5-0.5b-modelopt-fp8-pc-pt"
430
model_path = _snapshot_download_or_skip(model_id)
431
427
- # Set model config as model_config.dtype is required in ModelOptFp8LinearMethod.
428
- default_vllm_config.model_config = ModelConfig()
429
- with vllm_runner(model_path, quantization="modelopt", enforce_eager=True) as llm:
430
-
431
- def check_model(model):
432
- layer = model.model.layers[0]
433
-
434
- qkv_proj = layer.self_attn.qkv_proj
435
- o_proj = layer.self_attn.o_proj
436
- gate_up_proj = layer.mlp.gate_up_proj
437
- down_proj = layer.mlp.down_proj
438
-
439
- from vllm.model_executor.layers.quantization.modelopt import (
440
- ModelOptFp8PcPtLinearMethod,
441
- )
442
-
443
- assert isinstance(qkv_proj.quant_method, ModelOptFp8PcPtLinearMethod)
444
- assert isinstance(o_proj.quant_method, ModelOptFp8PcPtLinearMethod)
445
- assert isinstance(gate_up_proj.quant_method, ModelOptFp8PcPtLinearMethod)
446
- assert isinstance(down_proj.quant_method, ModelOptFp8PcPtLinearMethod)
447
-
448
- fp8_dtype = current_platform.fp8_dtype()
449
- assert qkv_proj.weight.dtype == fp8_dtype
450
- assert o_proj.weight.dtype == fp8_dtype
451
- assert gate_up_proj.weight.dtype == fp8_dtype
452
- assert down_proj.weight.dtype == fp8_dtype
453
-
454
- # Per-channel scales; activations are dynamically scaled per token.
455
- assert hasattr(qkv_proj, "weight_scale")
456
- assert qkv_proj.weight_scale.dtype == torch.float32
457
- assert qkv_proj.weight_scale.dim() == 1
458
- assert not hasattr(qkv_proj, "input_scale")
459
-
460
- assert hasattr(o_proj, "weight_scale")
461
- assert o_proj.weight_scale.dtype == torch.float32
462
- assert o_proj.weight_scale.dim() == 1
463
- assert not hasattr(o_proj, "input_scale")
464
-
465
- assert hasattr(gate_up_proj, "weight_scale")
466
- assert gate_up_proj.weight_scale.dtype == torch.float32
467
- assert gate_up_proj.weight_scale.dim() == 1
468
- assert not hasattr(gate_up_proj, "input_scale")
432
+ model, vllm_config = load_model_without_vllm_runner(
433
+ model_path,
434
+ quantization="modelopt",
435
+ )
436
+ layer = model.model.layers[0]
437
470
- assert hasattr(down_proj, "weight_scale")
471
- assert down_proj.weight_scale.dtype == torch.float32
472
- assert down_proj.weight_scale.dim() == 1
473
- assert not hasattr(down_proj, "input_scale")
438
+ qkv_proj = layer.self_attn.qkv_proj
439
+ o_proj = layer.self_attn.o_proj
440
+ gate_up_proj = layer.mlp.gate_up_proj
441
+ down_proj = layer.mlp.down_proj
442
475
- llm.apply_model(check_model)
443
+ from vllm.model_executor.layers.quantization.modelopt import (
444
+ ModelOptFp8PcPtLinearMethod,
445
+ )
446
477
- output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
478
- assert output
479
- print(f"ModelOpt FP8_PER_CHANNEL_PER_TOKEN output: {output}")
447
+ assert isinstance(qkv_proj.quant_method, ModelOptFp8PcPtLinearMethod)
448
+ assert isinstance(o_proj.quant_method, ModelOptFp8PcPtLinearMethod)
449
+ assert isinstance(gate_up_proj.quant_method, ModelOptFp8PcPtLinearMethod)
450
+ assert isinstance(down_proj.quant_method, ModelOptFp8PcPtLinearMethod)
451
+
452
+ fp8_dtype = current_platform.fp8_dtype()
453
+ assert qkv_proj.weight.dtype == fp8_dtype
454
+ assert o_proj.weight.dtype == fp8_dtype
455
+ assert gate_up_proj.weight.dtype == fp8_dtype
456
+ assert down_proj.weight.dtype == fp8_dtype
457
+
458
+ # Per-channel scales; activations are dynamically scaled per token.
459
+ for projection in (qkv_proj, o_proj, gate_up_proj, down_proj):
460
+ assert hasattr(projection, "weight_scale")
461
+ assert projection.weight_scale.dtype == torch.float32
462
+ assert projection.weight_scale.dim() == 1
463
+ assert not hasattr(projection, "input_scale")
464
+
465
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
466
+ input_ids = torch.tensor([1, 2, 3, 4], device=current_platform.device_type)
467
+ positions = torch.arange(input_ids.numel(), device=current_platform.device_type)
468
+ with (
469
+ set_current_vllm_config(vllm_config),
470
+ set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
471
+ ):
472
+ hidden_states = model(input_ids, positions, None)
473
+ logits = model.compute_logits(hidden_states)
474
+ assert torch.isfinite(logits).all()
475
476
477
@pytest.mark.skipif(
478
not is_quant_method_supported("modelopt"),
479
reason="ModelOpt FP8 is not supported on this GPU type.",
480
)
486
-def test_modelopt_fp8_pb_wo_checkpoint_setup(default_vllm_config, vllm_runner):
481
+def test_modelopt_fp8_pb_wo_checkpoint_setup(monkeypatch, dist_init, workspace_init):
482
"""Test ModelOpt FP8_PB_WO checkpoint setup."""
483
model_id = "CedricHwang/qwen2.5-0.5b-modelopt-fp8-pb-wo"
484
model_path = _snapshot_download_or_skip(model_id)
485
491
- # Set model config as model_config.dtype is required in ModelOptFp8LinearMethod.
492
- default_vllm_config.model_config = ModelConfig()
493
- with vllm_runner(model_path, quantization="modelopt", enforce_eager=True) as llm:
494
-
495
- def check_model(model):
496
- layer = model.model.layers[0]
497
-
498
- qkv_proj = layer.self_attn.qkv_proj
499
- o_proj = layer.self_attn.o_proj
500
- gate_up_proj = layer.mlp.gate_up_proj
501
- down_proj = layer.mlp.down_proj
502
-
503
- from vllm.model_executor.layers.quantization.modelopt import (
504
- ModelOptFp8PbWoLinearMethod,
505
- )
506
-
507
- assert isinstance(qkv_proj.quant_method, ModelOptFp8PbWoLinearMethod)
508
- assert isinstance(o_proj.quant_method, ModelOptFp8PbWoLinearMethod)
509
- assert isinstance(gate_up_proj.quant_method, ModelOptFp8PbWoLinearMethod)
510
- assert isinstance(down_proj.quant_method, ModelOptFp8PbWoLinearMethod)
511
-
512
- fp8_dtype = current_platform.fp8_dtype()
513
- assert qkv_proj.weight.dtype == fp8_dtype
514
- assert o_proj.weight.dtype == fp8_dtype
515
- assert gate_up_proj.weight.dtype == fp8_dtype
516
- assert down_proj.weight.dtype == fp8_dtype
517
-
518
- # Block scales; should be materialized as a 2D [out_blk, in_blk] tensor.
519
- assert hasattr(qkv_proj, "weight_scale")
520
- assert qkv_proj.weight_scale.dtype == torch.float32
521
- assert qkv_proj.weight_scale.dim() == 2
486
+ model, vllm_config = load_model_without_vllm_runner(
487
+ model_path,
488
+ quantization="modelopt",
489
+ )
490
+ layer = model.model.layers[0]
491
523
- assert hasattr(o_proj, "weight_scale")
524
- assert o_proj.weight_scale.dtype == torch.float32
525
- assert o_proj.weight_scale.dim() == 2
492
+ qkv_proj = layer.self_attn.qkv_proj
493
+ o_proj = layer.self_attn.o_proj
494
+ gate_up_proj = layer.mlp.gate_up_proj
495
+ down_proj = layer.mlp.down_proj
496
527
- assert hasattr(gate_up_proj, "weight_scale")
528
- assert gate_up_proj.weight_scale.dtype == torch.float32
529
- assert gate_up_proj.weight_scale.dim() == 2
530
-
531
- assert hasattr(down_proj, "weight_scale")
532
- assert down_proj.weight_scale.dtype == torch.float32
533
- assert down_proj.weight_scale.dim() == 2
534
-
535
- llm.apply_model(check_model)
497
+ from vllm.model_executor.layers.quantization.modelopt import (
498
+ ModelOptFp8PbWoLinearMethod,
499
+ )
500
537
- output = llm.generate_greedy(["Hello my name is"], max_tokens=4)
538
- assert output
539
- print(f"ModelOpt FP8_PB_WO output: {output}")
501
+ assert isinstance(qkv_proj.quant_method, ModelOptFp8PbWoLinearMethod)
502
+ assert isinstance(o_proj.quant_method, ModelOptFp8PbWoLinearMethod)
503
+ assert isinstance(gate_up_proj.quant_method, ModelOptFp8PbWoLinearMethod)
504
+ assert isinstance(down_proj.quant_method, ModelOptFp8PbWoLinearMethod)
505
+
506
+ fp8_dtype = current_platform.fp8_dtype()
507
+ assert qkv_proj.weight.dtype == fp8_dtype
508
+ assert o_proj.weight.dtype == fp8_dtype
509
+ assert gate_up_proj.weight.dtype == fp8_dtype
510
+ assert down_proj.weight.dtype == fp8_dtype
511
+
512
+ # Block scales are materialized as a 2D [out_blk, in_blk] tensor.
513
+ for projection in (qkv_proj, o_proj, gate_up_proj, down_proj):
514
+ assert hasattr(projection, "weight_scale")
515
+ assert projection.weight_scale.dtype == torch.float32
516
+ assert projection.weight_scale.dim() == 2
517
+
518
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
519
+ input_ids = torch.tensor([1, 2, 3, 4], device=current_platform.device_type)
520
+ positions = torch.arange(input_ids.numel(), device=current_platform.device_type)
521
+ with (
522
+ set_current_vllm_config(vllm_config),
523
+ set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
524
+ ):
525
+ hidden_states = model(input_ids, positions, None)
526
+ logits = model.compute_logits(hidden_states)
527
+ assert torch.isfinite(logits).all()
528
529
530
def test_modelopt_nvfp4_config_dispatches_w4a4_method():
tests/quantization/test_online.py
+119
-52
@@ -13,16 +13,22 @@ from torch.distributed import ProcessGroup
13
from tests.quantization.utils import (
14
_test_online_quant_peak_mem_impl,
15
is_quant_method_supported,
16
+ load_model_without_vllm_runner,
17
)
18
from vllm import _custom_ops as ops
19
from vllm._aiter_ops import rocm_aiter_ops
20
from vllm._custom_ops import scaled_fp4_quant
21
+from vllm.config.quantization import resolve_quantization_config
22
+from vllm.forward_context import set_forward_context
23
+from vllm.model_executor.layers.attention import Attention
24
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
25
from vllm.model_executor.layers.quantization.online.fp8 import (
26
Fp8PerBlockOnlineLinearMethod,
27
Fp8PerBlockOnlineMoEMethod,
28
Fp8PerTensorOnlineLinearMethod,
29
Fp8PerTensorOnlineMoEMethod,
30
+ Fp8PtpcOnlineLinearMethod,
31
+ Fp8PtpcOnlineMoEMethod,
32
_fp8_channel_scale,
33
_fp8_quant_per_channel,
34
_fp8_scale,
@@ -43,6 +49,10 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
49
amax_for_tp_weight_quant,
50
weight_amax,
51
)
52
+from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader
53
+from vllm.model_executor.models.granitemoe import (
54
+ GraniteMoeModel,
55
+)
56
from vllm.platforms import current_platform
57
from vllm.utils.flashinfer import has_flashinfer_trtllm_fused_moe
58
@@ -128,6 +138,12 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
138
Fp8PerBlockOnlineLinearMethod,
139
Fp8PerBlockOnlineMoEMethod,
140
),
141
+ (
142
+ "fp8_per_channel",
143
+ None,
144
+ Fp8PtpcOnlineLinearMethod,
145
+ Fp8PtpcOnlineMoEMethod,
146
+ ),
147
# quantization='online' with per-layer-kind overrides
148
(
149
"online",
@@ -160,13 +176,14 @@ def test_online_nvfp4_reuses_kernel_when_weights_are_reprocessed(
176
"use_rocm_aiter", [True, False] if current_platform.is_rocm() else [False]
177
)
178
def test_online_quantization(
163
- vllm_runner,
179
quant_scheme: str,
180
online_quant_args: dict | None,
181
expected_linear_cls,
182
expected_moe_cls,
183
use_rocm_aiter: bool,
184
monkeypatch,
185
+ dist_init,
186
+ workspace_init,
187
) -> None:
188
"""
189
Tests that online quantization frontend configuration works -
@@ -188,67 +205,117 @@ def test_online_quantization(
205
if current_platform.is_xpu() and quant_scheme == "fp8_per_block":
206
pytest.skip("Skip test for online fp8_per_block on XPU platform.")
207
191
- # `LLM.apply_model` requires pickling a function.
192
- monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
193
-
194
- # a tiny model with both dense and MoE layers
208
model_name = "ibm-granite/granite-3.0-1b-a400m-base"
196
-
197
- runner_kwargs = dict(
209
+ model, vllm_config = load_model_without_vllm_runner(
210
+ model_name,
211
+ dtype="bfloat16",
212
quantization=quant_scheme,
199
- enforce_eager=True,
213
+ model_config_kwargs={
214
+ "quantization_config": resolve_quantization_config(
215
+ quant_scheme, online_quant_args
216
+ ),
217
+ "hf_overrides": {
218
+ "num_hidden_layers": 3,
219
+ "vocab_size": 256,
220
+ "hidden_size": 256,
221
+ "intermediate_size": 512,
222
+ "num_attention_heads": 4,
223
+ "num_key_value_heads": 4,
224
+ "max_position_embeddings": 64,
225
+ "num_local_experts": 4,
226
+ "num_experts_per_tok": 2,
227
+ },
228
+ },
229
+ model_loader_cls=DummyModelLoader,
230
)
201
- if online_quant_args is not None:
202
- runner_kwargs["quantization_config"] = online_quant_args
231
+
232
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
233
+
234
+ o_proj = model.model.layers[0].self_attn.o_proj
235
+ moe = model.model.layers[0].block_sparse_moe.experts
236
+ assert isinstance(o_proj.quant_method, expected_linear_cls)
237
+ assert isinstance(moe._quant_method, expected_moe_cls)
238
+
239
+ if quant_scheme == "mxfp4":
240
+ assert o_proj.weight.dtype == torch.uint8
241
+ elif current_platform.is_cuda() or current_platform.is_xpu():
242
+ assert o_proj.weight.dtype == torch.float8_e4m3fn
243
+ elif current_platform.is_rocm():
244
+ assert o_proj.weight.dtype == current_platform.fp8_dtype()
245
+ else:
246
+ pytest.skip("Only runs on CUDA and ROCm.")
247
+
248
+ if quant_scheme == "fp8_per_channel":
249
+ assert o_proj.weight_scale.ndim == 2
250
+ assert o_proj.weight_scale.shape[-1] == 1
251
+ assert o_proj.input_scale is None
252
+
253
+ if isinstance(online_quant_args, dict) and "ignore" in online_quant_args:
254
+ for layer_idx in range(len(model.model.layers)):
255
+ o_proj = model.model.layers[layer_idx].self_attn.o_proj
256
+ if layer_idx == 1:
257
+ assert isinstance(o_proj.quant_method, UnquantizedLinearMethod)
258
+ else:
259
+ assert isinstance(o_proj.quant_method, expected_linear_cls)
260
+
261
+ for layer in model.model.layers:
262
+ assert isinstance(
263
+ layer.self_attn.qkv_proj.quant_method, UnquantizedLinearMethod
264
+ )
265
+
266
+ input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE)
267
+ positions = torch.arange(input_ids.numel(), device=DEVICE)
268
+ with set_forward_context(None, vllm_config, num_tokens=input_ids.numel()):
269
+ hidden_states = model(input_ids, positions, None)
270
+ logits = model.compute_logits(hidden_states)
271
+ assert torch.isfinite(logits).all()
272
+
273
+
274
+@pytest.mark.skipif(
275
+ not is_quant_method_supported("fp8"),
276
+ reason="FP8 is not supported on this GPU type.",
277
+)
278
+def test_online_quantization_loads_real_weights(vllm_runner, monkeypatch) -> None:
279
+ """Verify online quantization loads a Granite-MoE checkpoint end to end."""
280
+ monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
281
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
282
+
283
+ original_load_weights = GraniteMoeModel.load_weights
284
+
285
+ def load_weights(self, weights):
286
+ weights = (
287
+ (name, weight)
288
+ for name, weight in weights
289
+ if not name.startswith("layers.") or int(name.split(".")[1]) < 3
290
+ )
291
+ return original_load_weights(self, weights)
292
+
293
+ monkeypatch.setattr(GraniteMoeModel, "load_weights", load_weights)
294
295
with vllm_runner(
205
- model_name,
206
- **runner_kwargs,
296
+ "ibm-granite/granite-3.0-1b-a400m-base",
297
+ quantization="fp8_per_tensor",
298
+ dtype="bfloat16",
299
+ enforce_eager=True,
300
+ hf_overrides={"num_hidden_layers": 3},
301
+ max_model_len=16,
302
+ max_num_seqs=1,
303
) as llm:
304
305
def check_model(model):
210
- # checks further down in the test case are hardcoded for this
211
- # model
212
- assert model_name == "ibm-granite/granite-3.0-1b-a400m-base"
213
-
214
- o_proj = model.model.layers[0].self_attn.o_proj
215
- moe = model.model.layers[0].block_sparse_moe.experts
216
-
217
- # o_proj and moe in layer 0 are always quantized (never ignored)
218
- # because of how we craft the test case inputs
219
- assert isinstance(o_proj.quant_method, expected_linear_cls)
220
- if moe is not None:
221
- assert isinstance(moe._quant_method, expected_moe_cls)
222
-
223
- if quant_scheme == "mxfp4":
224
- # Packed e2m1 values, two per byte.
225
- assert o_proj.weight.dtype == torch.uint8
226
- elif current_platform.is_cuda() or current_platform.is_xpu():
227
- assert o_proj.weight.dtype == torch.float8_e4m3fn
228
- elif current_platform.is_rocm():
229
- assert o_proj.weight.dtype == current_platform.fp8_dtype()
230
- else:
231
- pytest.skip("Only runs on CUDA and ROCm.")
232
-
233
- # Verify ignored layers are unquantized.
234
- if isinstance(online_quant_args, dict) and "ignore" in online_quant_args:
235
- # only .*1.self_attn_o_proj is skipped
236
- for layer_idx in range(len(model.model.layers)):
237
- o_proj = model.model.layers[layer_idx].self_attn.o_proj
238
- if layer_idx == 1:
239
- assert isinstance(o_proj.quant_method, UnquantizedLinearMethod)
240
- else:
241
- assert isinstance(o_proj.quant_method, expected_linear_cls)
242
-
243
- # every .*self_attn.qkv_proj is skipped
244
- for layer_idx in range(len(model.model.layers)):
245
- qkv_proj = model.model.layers[layer_idx].self_attn.qkv_proj
246
- assert isinstance(qkv_proj.quant_method, UnquantizedLinearMethod)
306
+ layer = model.model.layers[0]
307
+ assert isinstance(
308
+ layer.self_attn.o_proj.quant_method,
309
+ Fp8PerTensorOnlineLinearMethod,
310
+ )
311
+ assert isinstance(
312
+ layer.block_sparse_moe.experts._quant_method,
313
+ Fp8PerTensorOnlineMoEMethod,
314
+ )
315
316
llm.apply_model(check_model)
249
-
250
- outputs = llm.generate_greedy(["Hello my name is"], max_tokens=4)
251
- print(outputs[0][1])
317
+ outputs = llm.generate_greedy(["Hello my name is"], max_tokens=1)
318
+ assert outputs
319
320
321
@pytest.mark.skipif(
tests/quantization/test_quark.py
+142
-207
@@ -18,7 +18,12 @@ import pytest
18
import torch
19
from packaging import version
20
21
-from vllm._aiter_ops import is_aiter_found_and_supported
21
+from tests.quantization.utils import load_model_without_vllm_runner
22
+from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
23
+from vllm.config import set_current_vllm_config
24
+from vllm.config.cache import CacheConfig
25
+from vllm.forward_context import set_forward_context
26
+from vllm.model_executor.layers.attention import Attention
27
from vllm.model_executor.layers.fused_moe import (
28
RoutedExperts,
29
UnquantizedFusedMoEMethod,
@@ -76,6 +81,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
81
kNvfp4Dynamic,
82
kNvfp4Static,
83
)
84
+from vllm.model_executor.models.llama import LlamaForCausalLM
85
from vllm.platforms import current_platform
86
from vllm.transformers_utils.repo_utils import hf_api
87
@@ -903,157 +909,184 @@ def test_quant_method_dispatch_instantiation(case, monkeypatch, default_vllm_con
909
910
911
@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
906
-@pytest.mark.parametrize("tp", [1])
907
-def test_quark_fp8_w_per_tensor_a_per_tensor(vllm_runner, kv_cache_dtype, tp):
912
+def test_quark_fp8_w_per_tensor_a_per_tensor(
913
+ kv_cache_dtype: str, monkeypatch, dist_init, workspace_init
914
+):
915
model_path = "amd/Llama-3.1-8B-Instruct-FP8-KV-Quark-test"
909
- with vllm_runner(
910
- model_path,
911
- enforce_eager=True,
912
- kv_cache_dtype=kv_cache_dtype,
913
- tensor_parallel_size=tp,
914
- ) as llm:
915
-
916
- def check_model(model):
917
- layer = model.model.layers[0]
916
+ checkpoint_scales = {}
917
+ scale_names = {
918
+ "model.layers.0.self_attn.k_proj.output_scale",
919
+ "model.layers.0.self_attn.v_proj.output_scale",
920
+ }
921
+ original_load_weights = LlamaForCausalLM.load_weights
922
919
- qkv_proj = layer.self_attn.qkv_proj
923
+ def load_weights(self, weights):
924
+ def capture_scales():
925
+ for name, weight in weights:
926
+ if name in scale_names:
927
+ checkpoint_scales[name] = weight.detach().cpu()
928
+ yield name, weight
929
921
- assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
922
- assert isinstance(qkv_proj.scheme, QuarkW8A8Fp8)
930
+ return original_load_weights(self, capture_scales())
931
924
- if isinstance(qkv_proj.scheme, QuarkW8A8Fp8):
925
- assert len(qkv_proj.input_scale.shape) == 0
926
- assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
927
- assert len(qkv_proj.weight_scale.shape) == 0
932
+ monkeypatch.setattr(LlamaForCausalLM, "load_weights", load_weights)
933
+ model, vllm_config = load_model_without_vllm_runner(
934
+ model_path,
935
+ model_config_kwargs={"hf_overrides": {"num_hidden_layers": 3}},
936
+ vllm_config_kwargs={"cache_config": CacheConfig(cache_dtype=kv_cache_dtype)},
937
+ )
938
929
- llm.apply_model(check_model)
939
+ qkv_proj = model.model.layers[0].self_attn.qkv_proj
940
+ assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
941
+ assert isinstance(qkv_proj.scheme, QuarkW8A8Fp8)
942
+ assert len(qkv_proj.input_scale.shape) == 0
943
+ assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
944
+ assert len(qkv_proj.weight_scale.shape) == 0
945
+
946
+ attn = model.model.layers[0].self_attn.attn
947
+ if kv_cache_dtype == "fp8":
948
+ assert checkpoint_scales.keys() == scale_names
949
+ scale_multiplier = 2 if current_platform.is_fp8_fnuz() else 1
950
+ assert attn._k_scale_float == (
951
+ checkpoint_scales["model.layers.0.self_attn.k_proj.output_scale"].item()
952
+ * scale_multiplier
953
+ )
954
+ assert attn._v_scale_float == (
955
+ checkpoint_scales["model.layers.0.self_attn.v_proj.output_scale"].item()
956
+ * scale_multiplier
957
+ )
958
+ else:
959
+ assert attn._k_scale_float == 1.0
960
+ assert attn._v_scale_float == 1.0
961
931
- output = llm.generate_greedy("Hello my name is", max_tokens=4)
932
- assert output
962
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
963
+ input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
964
+ positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
965
+ with (
966
+ set_current_vllm_config(vllm_config),
967
+ set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
968
+ ):
969
+ hidden_states = model(input_ids, positions, None)
970
+ logits = model.compute_logits(hidden_states)
971
+ assert torch.isfinite(logits).all()
972
973
935
-@pytest.mark.parametrize("tp", [1])
936
-def test_quark_fp8_w_per_channel_a_per_token(vllm_runner, tp):
974
+def test_quark_fp8_w_per_channel_a_per_token(monkeypatch, dist_init, workspace_init):
975
model_path = "amd/Qwen2.5-1.5B-Instruct-ptpc-Quark-ts"
938
- with vllm_runner(model_path, enforce_eager=True, tensor_parallel_size=tp) as llm:
939
-
940
- def check_model(model):
941
- layer = model.model.layers[0]
942
-
943
- qkv_proj = layer.self_attn.qkv_proj
944
-
945
- assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
946
- assert isinstance(qkv_proj.scheme, QuarkW8A8Fp8)
947
-
948
- if isinstance(qkv_proj.scheme, QuarkW8A8Fp8):
949
- assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
950
- assert qkv_proj.weight_scale.shape[0] == qkv_proj.weight.shape[1]
951
- assert qkv_proj.weight_scale.shape[1] == 1
976
+ model, vllm_config = load_model_without_vllm_runner(
977
+ model_path,
978
+ model_config_kwargs={"hf_overrides": {"num_hidden_layers": 3}},
979
+ )
980
953
- llm.apply_model(check_model)
981
+ qkv_proj = model.model.layers[0].self_attn.qkv_proj
982
+ assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
983
+ assert isinstance(qkv_proj.scheme, QuarkW8A8Fp8)
984
+ assert qkv_proj.weight.dtype is current_platform.fp8_dtype()
985
+ assert qkv_proj.weight_scale.shape[0] == qkv_proj.weight.shape[1]
986
+ assert qkv_proj.weight_scale.shape[1] == 1
987
955
- output = llm.generate_greedy("Hello my name is", max_tokens=4)
956
- assert output
988
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
989
+ input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
990
+ positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
991
+ with (
992
+ set_current_vllm_config(vllm_config),
993
+ set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
994
+ ):
995
+ hidden_states = model(input_ids, positions, None)
996
+ logits = model.compute_logits(hidden_states)
997
+ assert torch.isfinite(logits).all()
998
999
959
-@pytest.mark.parametrize("tp", [1])
960
-def test_quark_int8_w_per_tensor_a_per_tensor(vllm_runner, tp):
1000
+def test_quark_int8_w_per_tensor_a_per_tensor(monkeypatch, dist_init, workspace_init):
1001
model_path = "amd/Llama-3.1-8B-Instruct-w-int8-a-int8-sym-test"
962
- with vllm_runner(model_path, enforce_eager=True, tensor_parallel_size=tp) as llm:
963
-
964
- def check_model(model):
965
- layer = model.model.layers[0]
966
-
967
- qkv_proj = layer.self_attn.qkv_proj
968
-
969
- assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
970
- assert isinstance(qkv_proj.scheme, QuarkW8A8Int8)
971
-
972
- llm.apply_model(check_model)
1002
+ model, vllm_config = load_model_without_vllm_runner(
1003
+ model_path,
1004
+ model_config_kwargs={"hf_overrides": {"num_hidden_layers": 3}},
1005
+ )
1006
+ with set_current_vllm_config(vllm_config):
1007
+ qkv_proj = model.model.layers[0].self_attn.qkv_proj
1008
+ assert isinstance(qkv_proj.quant_method, QuarkLinearMethod)
1009
+ assert isinstance(qkv_proj.scheme, QuarkW8A8Int8)
1010
974
- output = llm.generate_greedy("Hello my name is", max_tokens=4)
975
- assert output
1011
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
1012
+ input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
1013
+ positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
1014
+ with set_forward_context(None, vllm_config, num_tokens=input_ids.numel()):
1015
+ hidden_states = model(input_ids, positions, None)
1016
+ logits = model.compute_logits(hidden_states)
1017
+ assert torch.isfinite(logits).all()
1018
1019
978
-@pytest.mark.parametrize("tp", [1])
979
-def test_quark_int8_w8a8_moe(vllm_runner, tp):
1020
+def test_quark_int8_w8a8_moe(monkeypatch, dist_init, workspace_init):
1021
"""Test W8A8 INT8 MoE quantization with a tiny Qwen3 MoE model."""
1022
model_path = "amd/tiny-qwen3-moe-w8a8-int8"
982
- with vllm_runner(
1023
+ model, vllm_config = load_model_without_vllm_runner(
1024
model_path,
984
- enforce_eager=True,
985
- tensor_parallel_size=tp,
986
- gpu_memory_utilization=0.1,
987
- ) as llm:
988
-
989
- def check_model(model):
990
- layer = model.model.layers[0]
991
- # MoE experts should use QuarkW8A8Int8MoEMethod
992
- moe = layer.mlp.experts
993
- assert isinstance(moe._quant_method, QuarkW8A8Int8MoEMethod), (
994
- f"Expected QuarkW8A8Int8MoEMethod, got {type(moe._quant_method)}"
995
- )
996
- # Non-MoE linear layers should use QuarkW8A8Int8
997
- qkv_proj = layer.self_attn.qkv_proj
998
- assert isinstance(qkv_proj.scheme, QuarkW8A8Int8)
1025
+ model_config_kwargs={"hf_overrides": {"num_hidden_layers": 3}},
1026
+ )
1027
1000
- llm.apply_model(check_model)
1028
+ layer = model.model.layers[0]
1029
+ moe = layer.mlp.experts
1030
+ assert isinstance(moe._quant_method, QuarkW8A8Int8MoEMethod), (
1031
+ f"Expected QuarkW8A8Int8MoEMethod, got {type(moe._quant_method)}"
1032
+ )
1033
+ qkv_proj = layer.self_attn.qkv_proj
1034
+ assert isinstance(qkv_proj.scheme, QuarkW8A8Int8)
1035
1002
- output = llm.generate_greedy("Hello", max_tokens=4)
1003
- assert output
1036
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
1037
+ input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
1038
+ positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
1039
+ with (
1040
+ set_current_vllm_config(vllm_config),
1041
+ set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
1042
+ ):
1043
+ hidden_states = model(input_ids, positions, None)
1044
+ logits = model.compute_logits(hidden_states)
1045
+ assert torch.isfinite(logits).all()
1046
1047
1048
@pytest.mark.skipif(
1049
not (on_gfx950() or on_gfx942()),
1050
reason="Quark W4A8 (INT4-FP8) MoE requires the AITER kernel on gfx942/gfx950",
1051
)
1010
-@pytest.mark.parametrize("tp", [1])
1011
-def test_quark_w4a8_fp8_moe(vllm_runner, monkeypatch, tp):
1052
+def test_quark_w4a8_fp8_moe(monkeypatch, dist_init, workspace_init):
1053
"""Test W4A8 (INT4 weight + FP8 activation) MoE with a tiny Qwen3 MoE model.
1054
1055
W4A8 dispatches through the AITER fused MoE kernel, so AITER must be on.
1056
"""
1057
monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1")
1058
monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1")
1059
+ rocm_aiter_ops.refresh_env_variables()
1060
+
1061
model_path = "amd/tiny-qwen3-moe-w4a8"
1019
- with vllm_runner(
1062
+ model, vllm_config = load_model_without_vllm_runner(
1063
model_path,
1021
- enforce_eager=True,
1022
- tensor_parallel_size=tp,
1023
- gpu_memory_utilization=0.1,
1024
- ) as llm:
1025
-
1026
- def check_model(model):
1027
- moe = model.model.layers[0].mlp.experts
1028
- assert isinstance(moe._quant_method, QuarkW4A8Fp8MoEMethod), (
1029
- f"Expected QuarkW4A8Fp8MoEMethod, got {type(moe._quant_method)}"
1030
- )
1031
-
1032
- llm.apply_model(check_model)
1064
+ )
1065
+ with set_current_vllm_config(vllm_config):
1066
+ moe = model.model.layers[0].mlp.experts
1067
+ assert isinstance(moe._quant_method, QuarkW4A8Fp8MoEMethod), (
1068
+ f"Expected QuarkW4A8Fp8MoEMethod, got {type(moe._quant_method)}"
1069
+ )
1070
1034
- output = llm.generate_greedy("Hello", max_tokens=4)
1035
- assert output
1071
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
1072
+ input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
1073
+ positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
1074
+ with set_forward_context(None, vllm_config, num_tokens=input_ids.numel()):
1075
+ hidden_states = model(input_ids, positions, None)
1076
+ logits = model.compute_logits(hidden_states)
1077
+ assert torch.isfinite(logits).all()
1078
1079
1038
-def test_quark_fp8_parity(vllm_runner):
1080
+def test_quark_fp8_parity(dist_init, workspace_init):
1081
quark_model_id = "amd-quark/llama-tiny-fp8-quark-quant-method"
1082
fp8_model_id = "amd-quark/llama-tiny-fp8-quant-method"
1083
1042
- llm_kwargs = {
1043
- "tensor_parallel_size": 1,
1044
- "enforce_eager": True,
1045
- "gpu_memory_utilization": 0.1,
1046
- }
1047
- with (
1048
- vllm_runner(quark_model_id, **llm_kwargs) as quark_handle,
1049
- vllm_runner(fp8_model_id, **llm_kwargs) as fp8_handle,
1050
- ):
1051
-
1052
- def get_state_dict(model):
1053
- return {k: v.cpu() for k, v in model.state_dict().items()}
1084
+ def load_state_dict(model_id: str) -> dict[str, torch.Tensor]:
1085
+ model, _ = load_model_without_vllm_runner(model_id)
1086
+ return {k: v.cpu() for k, v in model.state_dict().items()}
1087
1055
- (quark_state_dict,) = quark_handle.apply_model(get_state_dict)
1056
- (fp8_state_dict,) = fp8_handle.apply_model(get_state_dict)
1088
+ quark_state_dict = load_state_dict(quark_model_id)
1089
+ fp8_state_dict = load_state_dict(fp8_model_id)
1090
1091
assert fp8_state_dict.keys() == quark_state_dict.keys()
1092
@@ -1097,104 +1130,6 @@ GSM8K_ACCURACY_CONFIGS = [
1130
),
1131
]
1132
1100
-WIKITEXT_ACCURACY_CONFIGS = [
1101
- AccuracyTestConfig(
1102
- model_name="fxmarty/qwen1.5_moe_a2.7b_chat_w_fp4_a_fp6_e2m3",
1103
- excepted_value=11.3,
1104
- ),
1105
- AccuracyTestConfig(
1106
- model_name="fxmarty/qwen1.5_moe_a2.7b_chat_w_fp6_e3m2_a_fp6_e3m2",
1107
- excepted_value=10.6,
1108
- ),
1109
- AccuracyTestConfig(
1110
- model_name="fxmarty/qwen_1.5-moe-a2.7b-mxfp4", excepted_value=12.45
1111
- ),
1112
-]
1113
-
1114
-
1115
-@pytest.mark.skipif(
1116
- not QUARK_MXFP4_AVAILABLE,
1117
- reason=f"amd-quark>={QUARK_MXFP4_MIN_VERSION} is not available",
1118
-)
1119
-@pytest.mark.parametrize(
1120
- "config",
1121
- [pytest.param(val, id=f"config:{val}") for val in WIKITEXT_ACCURACY_CONFIGS],
1122
-)
1123
-@pytest.mark.parametrize(
1124
- "tp_size", [pytest.param(val, id=f"tp_size:{val}") for val in [1, 2]]
1125
-)
1126
-def test_ocp_mx_wikitext_correctness(config: AccuracyTestConfig, tp_size: int):
1127
- device_count = torch.accelerator.device_count()
1128
- if device_count < tp_size:
1129
- pytest.skip(f"This test requires >={tp_size} gpus, got only {device_count}")
1130
-
1131
- task = "wikitext"
1132
- rtol = 0.1
1133
-
1134
- # Smaller cudagraph_capture_sizes to speed up the test.
1135
- results = lm_eval.simple_evaluate(
1136
- model="vllm",
1137
- model_args=config.get_model_args(
1138
- tp_size=tp_size, kwargs={"cudagraph_capture_sizes": [16]}
1139
- ),
1140
- tasks=task,
1141
- batch_size=64,
1142
- )
1143
-
1144
- EXPECTED_VALUE = config.excepted_value
1145
- measured_value = results["results"][task]["word_perplexity,none"]
1146
- assert (
1147
- measured_value < EXPECTED_VALUE + rtol
1148
- and measured_value > EXPECTED_VALUE - rtol
1149
- ), f"Expected: {EXPECTED_VALUE} | Measured: {measured_value}"
1150
-
1151
-
1152
-@pytest.mark.skipif(
1153
- not QUARK_MXFP4_AVAILABLE,
1154
- reason=f"amd-quark>={QUARK_MXFP4_MIN_VERSION} is not available",
1155
-)
1156
-@pytest.mark.parametrize("tp_size", [1, 2])
1157
-def test_nvfp4_wikitext_correctness(tp_size: int):
1158
- device_count = torch.accelerator.device_count()
1159
- if device_count < tp_size:
1160
- pytest.skip(f"This test requires >={tp_size} gpus, got only {device_count}")
1161
-
1162
- # NOTE: expected_value from nvidia/Qwen3-30B-A3B-NVFP4
1163
- expected_value = 11.2391
1164
-
1165
- model_name = "amd-quark/Qwen3-30B-A3B-nvfp4-quark"
1166
- task = "wikitext"
1167
-
1168
- rtol = 0.25
1169
-
1170
- config = AccuracyTestConfig(
1171
- model_name=model_name,
1172
- excepted_value=expected_value,
1173
- )
1174
-
1175
- model_args = config.get_model_args(
1176
- tp_size=tp_size,
1177
- kwargs={
1178
- "cudagraph_capture_sizes": [16],
1179
- },
1180
- )
1181
- model_args.pop("add_bos_token")
1182
-
1183
- # Smaller cudagraph_capture_sizes to speed up the test.
1184
- results = lm_eval.simple_evaluate(
1185
- model="vllm",
1186
- model_args=model_args,
1187
- tasks=task,
1188
- batch_size=64,
1189
- )
1190
-
1191
- EXPECTED_VALUE = config.excepted_value
1192
- measured_value = results["results"][task]["word_perplexity,none"]
1193
- assert (
1194
- measured_value < EXPECTED_VALUE + rtol
1195
- and measured_value > EXPECTED_VALUE - rtol
1196
- ), f"Expected: {EXPECTED_VALUE} | Measured: {measured_value}"
1197
-
1133
1134
@pytest.mark.parametrize("config", GSM8K_ACCURACY_CONFIGS)
1135
@pytest.mark.skipif(
tests/quantization/test_register_quantization_config.py
+27
-19
@@ -14,6 +14,10 @@ import pytest
14
import torch
15
import torch.nn.functional as F
16
17
+from tests.quantization.utils import load_model_without_vllm_runner
18
+from vllm.config import set_current_vllm_config
19
+from vllm.forward_context import set_forward_context
20
+from vllm.model_executor.layers.attention import Attention
21
from vllm.model_executor.layers.linear import (
22
LinearBase, # noqa: E501
23
UnquantizedLinearMethod,
@@ -26,6 +30,7 @@ from vllm.model_executor.layers.quantization import (
30
from vllm.model_executor.layers.quantization.base_config import (
31
QuantizationConfig, # noqa: E501
32
)
33
+from vllm.platforms import current_platform
34
35
36
class FakeQuantLinearMethod(UnquantizedLinearMethod):
@@ -124,23 +129,26 @@ def test_register_quantization_config(caplog_vllm):
129
"meta-llama/Llama-3.2-1B-Instruct",
130
],
131
)
127
-def test_custom_quant(vllm_runner, model, monkeypatch):
132
+def test_custom_quant(model, monkeypatch, dist_init, workspace_init):
133
"""Test infer with the custom quantization method."""
129
- # `LLM.apply_model` requires pickling a function.
130
- monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
131
-
132
- with vllm_runner(
133
- model_name=model, quantization="custom_quant", enforce_eager=True
134
- ) as llm:
135
-
136
- def check_model(model):
137
- layer = model.model.layers[0]
138
- qkv_proj = layer.self_attn.qkv_proj
139
-
140
- # Check the quantization method is FakeQuantLinearMethod
141
- assert isinstance(qkv_proj.quant_method, FakeQuantLinearMethod)
142
-
143
- llm.apply_model(check_model)
144
-
145
- output = llm.generate_greedy("Hello my name is", max_tokens=1)
146
- assert output
134
+ vllm_model, vllm_config = load_model_without_vllm_runner(
135
+ model,
136
+ quantization="custom_quant",
137
+ )
138
+ layer = vllm_model.model.layers[0]
139
+ qkv_proj = layer.self_attn.qkv_proj
140
+
141
+ # Check the quantization method is FakeQuantLinearMethod.
142
+ assert isinstance(qkv_proj.quant_method, FakeQuantLinearMethod)
143
+
144
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
145
+ device_type = current_platform.device_type
146
+ input_ids = torch.tensor([1, 2, 3, 4], device=device_type)
147
+ positions = torch.arange(input_ids.numel(), device=device_type)
148
+ with (
149
+ set_current_vllm_config(vllm_config),
150
+ set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
151
+ ):
152
+ hidden_states = vllm_model(input_ids, positions, None)
153
+ logits = vllm_model.compute_logits(hidden_states)
154
+ assert torch.isfinite(logits).all()
tests/quantization/test_torchao.py
+68
-24
@@ -7,6 +7,11 @@ import pytest
7
import torch
8
from packaging import version
9
10
+from tests.quantization.utils import load_model_without_vllm_runner
11
+from vllm.config import set_current_vllm_config
12
+from vllm.config.load import LoadConfig
13
+from vllm.forward_context import set_forward_context
14
+from vllm.model_executor.layers.attention import Attention
15
from vllm.model_executor.model_loader import get_model_loader
16
from vllm.platforms import current_platform
17
@@ -27,20 +32,39 @@ TORCHAO_VERSION_0_18_AVAILABLE = TORCHAO_AVAILABLE and version.parse(
32
) >= version.parse("0.18.0")
33
34
35
+# Fix for leaking `torch.set_float32_matmul_precision("high")`
36
+# https://github.com/pytorch/ao/blob/v0.17.0/torchao/quantization/utils.py#L696
37
+@pytest.fixture(autouse=True)
38
+def restore_float32_matmul_precision():
39
+ original_precision = torch.get_float32_matmul_precision()
40
+ try:
41
+ yield
42
+ finally:
43
+ torch.set_float32_matmul_precision(original_precision)
44
+
45
+
46
@pytest.mark.skipif(
47
current_platform.is_rocm() and current_platform.is_fp8_fnuz(),
48
reason="Only fp8_fnuz supported on CDNA3 architecture",
49
)
50
@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
35
-def test_pre_quantized_model(vllm_runner):
36
- with vllm_runner(
51
+def test_pre_quantized_model(monkeypatch, dist_init, workspace_init):
52
+ model, vllm_config = load_model_without_vllm_runner(
53
"torchao-testing/opt-125m-Float8WeightOnlyConfig-v2-0.15.0",
54
quantization="torchao",
55
dtype="bfloat16",
40
- enforce_eager=True,
41
- ) as llm:
42
- output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
43
- assert output
56
+ )
57
+
58
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
59
+ input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
60
+ positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
61
+ with (
62
+ set_current_vllm_config(vllm_config),
63
+ set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
64
+ ):
65
+ hidden_states = model(input_ids, positions, None)
66
+ logits = model.compute_logits(hidden_states)
67
+ assert torch.isfinite(logits).all()
68
69
70
@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
@@ -51,35 +75,55 @@ def test_pre_quantized_model(vllm_runner):
75
# {"": "cuda"},
76
],
77
)
54
-def test_opt_125m_int8wo_model_loading_with_params(vllm_runner, pt_load_map_location):
78
+def test_opt_125m_int8wo_model_loading_with_params(
79
+ pt_load_map_location, monkeypatch, dist_init, workspace_init
80
+):
81
torch._dynamo.reset()
82
model_name = "jerryzh168/opt-125m-int8wo-partial-quant"
57
- with vllm_runner(
58
- model_name=model_name,
83
+ model, vllm_config = load_model_without_vllm_runner(
84
+ model_name,
85
quantization="torchao",
60
- dtype="bfloat16",
61
- pt_load_map_location=pt_load_map_location,
62
- enforce_eager=True,
63
- ) as llm:
64
- output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
86
+ vllm_config_kwargs={
87
+ "load_config": LoadConfig(pt_load_map_location=pt_load_map_location)
88
+ },
89
+ )
90
66
- assert output
91
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
92
+ input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
93
+ positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
94
+ with (
95
+ set_current_vllm_config(vllm_config),
96
+ set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
97
+ ):
98
+ hidden_states = model(input_ids, positions, None)
99
+ logits = model.compute_logits(hidden_states)
100
+ assert torch.isfinite(logits).all()
101
102
103
@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
70
-def test_qwenvl_int8wo_model_loading_with_params(vllm_runner):
104
+def test_qwenvl_int8wo_model_loading_with_params(
105
+ monkeypatch, dist_init, workspace_init
106
+):
107
torch._dynamo.reset()
108
model_name = "mobicham/Qwen2.5-VL-3B-Instruct_int8wo_ao"
73
- with vllm_runner(
74
- model_name=model_name,
109
+ model, vllm_config = load_model_without_vllm_runner(
110
+ model_name,
111
quantization="torchao",
76
- dtype="bfloat16",
77
- pt_load_map_location=f"{DEVICE_TYPE}:0",
78
- enforce_eager=True,
79
- ) as llm:
80
- output = llm.generate_greedy(["The capital of France is"], max_tokens=4)
112
+ vllm_config_kwargs={
113
+ "load_config": LoadConfig(pt_load_map_location=f"{DEVICE_TYPE}:0")
114
+ },
115
+ )
116
82
- assert output
117
+ monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
118
+ input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
119
+ positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE).expand(3, -1)
120
+ with (
121
+ set_current_vllm_config(vllm_config),
122
+ set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
123
+ ):
124
+ hidden_states = model(input_ids, positions, None)
125
+ logits = model.compute_logits(hidden_states)
126
+ assert torch.isfinite(logits).all()
127
128
129
@pytest.mark.skipif(not TORCHAO_AVAILABLE, reason="torchao is not available")
tests/quantization/utils.py
+71
-1
@@ -2,13 +2,83 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
import logging
5
+from typing import Any
6
7
import regex as re
7
-
8
+import torch
9
+
10
+from vllm.config import (
11
+ CompilationConfig,
12
+ ModelConfig,
13
+ VllmConfig,
14
+ set_current_vllm_config,
15
+)
16
from vllm.model_executor.layers.quantization import get_quantization_config
17
+from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
18
from vllm.platforms import current_platform
19
20
21
+def _limit_num_hidden_layers(
22
+ model: torch.nn.Module, num_hidden_layers: int | None
23
+) -> None:
24
+ if num_hidden_layers is None:
25
+ return
26
+
27
+ original_load_weights = model.load_weights
28
+
29
+ def should_load_weight(name: str) -> bool:
30
+ for prefix in ("model.layers.", "layers."):
31
+ if name.startswith(prefix):
32
+ layer_idx = int(name.removeprefix(prefix).split(".", 1)[0])
33
+ return layer_idx < num_hidden_layers
34
+ return True
35
+
36
+ def load_weights(weights):
37
+ weights = (
38
+ (name, weight) for name, weight in weights if should_load_weight(name)
39
+ )
40
+ return original_load_weights(weights)
41
+
42
+ model.load_weights = load_weights
43
+
44
+
45
+def load_model_without_vllm_runner(
46
+ model_path: str,
47
+ *,
48
+ dtype: str | torch.dtype = "bfloat16",
49
+ quantization: str | None = None,
50
+ model_config_kwargs: dict[str, Any] | None = None,
51
+ vllm_config_kwargs: dict[str, Any] | None = None,
52
+ model_loader_cls: type = DefaultModelLoader,
53
+) -> tuple[torch.nn.Module, VllmConfig]:
54
+ """Instantiate a model, load weights, and process them for inference."""
55
+ model_config = ModelConfig(
56
+ model=model_path,
57
+ dtype=dtype,
58
+ quantization=quantization,
59
+ **(model_config_kwargs or {}),
60
+ )
61
+ vllm_config_args = dict(vllm_config_kwargs or {})
62
+ vllm_config_args.setdefault("compilation_config", CompilationConfig(mode=0))
63
+ vllm_config = VllmConfig(model_config=model_config, **vllm_config_args)
64
+ hf_overrides = (model_config_kwargs or {}).get("hf_overrides") or {}
65
+ num_hidden_layers = hf_overrides.get("num_hidden_layers")
66
+
67
+ with set_current_vllm_config(vllm_config):
68
+ model_loader = model_loader_cls(vllm_config.load_config)
69
+ if num_hidden_layers is not None:
70
+ original_load_weights = model_loader.load_weights
71
+
72
+ def load_weights(model, model_config):
73
+ _limit_num_hidden_layers(model, num_hidden_layers)
74
+ original_load_weights(model, model_config)
75
+
76
+ model_loader.load_weights = load_weights
77
+ model = model_loader.load_model(vllm_config, model_config)
78
+
79
+ return model, vllm_config
80
+
81
+
82
def is_quant_method_supported(quant_method: str) -> bool:
83
# Currently, quantization tests only run GPUs
84
if current_platform.is_cpu():
tests/utils_/test_torch_utils.py
+14
@@ -11,6 +11,7 @@ from vllm.utils.torch_utils import (
11
get_kv_cache_torch_dtype,
12
is_lossless_cast,
13
is_quantized_kv_cache,
14
+ set_default_torch_dtype,
15
set_torch_threads_for_runtime,
16
startup_omp_num_threads,
17
)
@@ -28,6 +29,19 @@ def test_nvfp4_4over6_cache_dtype() -> None:
29
assert get_kv_quant_mode(cache_config.cache_dtype) == KVQuantMode.NVFP4
30
31
32
+def test_set_default_torch_dtype_restores_dtype_after_exception() -> None:
33
+ original_dtype = torch.get_default_dtype()
34
+
35
+ with (
36
+ pytest.raises(RuntimeError, match="expected failure"),
37
+ set_default_torch_dtype(torch.bfloat16),
38
+ ):
39
+ assert torch.get_default_dtype() == torch.bfloat16
40
+ raise RuntimeError("expected failure")
41
+
42
+ assert torch.get_default_dtype() == original_dtype
43
+
44
+
45
@pytest.mark.parametrize(
46
("src_dtype", "tgt_dtype", "expected_result"),
47
[
vllm/utils/torch_utils.py
+4
-2
@@ -164,8 +164,10 @@ def set_default_torch_dtype(dtype: torch.dtype):
164
"""Sets the default torch dtype to the given dtype."""
165
old_dtype = torch.get_default_dtype()
166
torch.set_default_dtype(dtype)
167
- yield
168
- torch.set_default_dtype(old_dtype)
167
+ try:
168
+ yield
169
+ finally:
170
+ torch.set_default_dtype(old_dtype)
171
172
173
def _cgroup_cpu_limit() -> float | None: