diff --git a/tests/quantization/test_quark.py b/tests/quantization/test_quark.py index bbb1702502c7..49af26fd885c 100644 --- a/tests/quantization/test_quark.py +++ b/tests/quantization/test_quark.py @@ -1020,32 +1020,32 @@ def test_quark_int8_w_per_tensor_a_per_tensor(monkeypatch, dist_init, workspace_ assert torch.isfinite(logits).all() -def test_quark_int8_w8a8_moe(monkeypatch, dist_init, workspace_init): +@pytest.mark.parametrize("tp", [1]) +def test_quark_int8_w8a8_moe(vllm_runner, tp): """Test W8A8 INT8 MoE quantization with a tiny Qwen3 MoE model.""" model_path = "amd/tiny-qwen3-moe-w8a8-int8" - model, vllm_config = load_model_without_vllm_runner( + with vllm_runner( model_path, - model_config_kwargs={"hf_overrides": {"num_hidden_layers": 3}}, - ) + enforce_eager=True, + tensor_parallel_size=tp, + gpu_memory_utilization=0.1, + ) as llm: + + def check_model(model): + layer = model.model.layers[0] + # MoE experts should use QuarkW8A8Int8MoEMethod + moe = layer.mlp.experts + assert isinstance(moe._quant_method, QuarkW8A8Int8MoEMethod), ( + f"Expected QuarkW8A8Int8MoEMethod, got {type(moe._quant_method)}" + ) + # Non-MoE linear layers should use QuarkW8A8Int8 + qkv_proj = layer.self_attn.qkv_proj + assert isinstance(qkv_proj.scheme, QuarkW8A8Int8) - layer = model.model.layers[0] - moe = layer.mlp.experts - assert isinstance(moe._quant_method, QuarkW8A8Int8MoEMethod), ( - f"Expected QuarkW8A8Int8MoEMethod, got {type(moe._quant_method)}" - ) - qkv_proj = layer.self_attn.qkv_proj - assert isinstance(qkv_proj.scheme, QuarkW8A8Int8) + llm.apply_model(check_model) - monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous()) - input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE) - positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE) - with ( - set_current_vllm_config(vllm_config), - set_forward_context(None, vllm_config, num_tokens=input_ids.numel()), - ): - hidden_states = model(input_ids, positions, None) - logits = model.compute_logits(hidden_states) - assert torch.isfinite(logits).all() + output = llm.generate_greedy("Hello", max_tokens=4) + assert output @pytest.mark.skipif(