[CI] Revert flaky `test_quark_int8_w8a8_moe` (#54991)
Signed-off-by: Felix Marty <Felix.Marty@amd.com>
fxmarty-amd committed
Sep 2, 2026 at 19:58 UTC
488e6fd53cb29b4255364b69a4d298b1bbc26062
1 file changed
+21
-21
tests/quantization/test_quark.py
+21
-21
@@ -1020,32 +1020,32 @@ def test_quark_int8_w_per_tensor_a_per_tensor(monkeypatch, dist_init, workspace_
1020
assert torch.isfinite(logits).all()
1021
1022
1023
-def test_quark_int8_w8a8_moe(monkeypatch, dist_init, workspace_init):
1023
+@pytest.mark.parametrize("tp", [1])
1024
+def test_quark_int8_w8a8_moe(vllm_runner, tp):
1025
"""Test W8A8 INT8 MoE quantization with a tiny Qwen3 MoE model."""
1026
model_path = "amd/tiny-qwen3-moe-w8a8-int8"
1026
- model, vllm_config = load_model_without_vllm_runner(
1027
+ with vllm_runner(
1028
model_path,
1028
- model_config_kwargs={"hf_overrides": {"num_hidden_layers": 3}},
1029
- )
1029
+ enforce_eager=True,
1030
+ tensor_parallel_size=tp,
1031
+ gpu_memory_utilization=0.1,
1032
+ ) as llm:
1033
+
1034
+ def check_model(model):
1035
+ layer = model.model.layers[0]
1036
+ # MoE experts should use QuarkW8A8Int8MoEMethod
1037
+ moe = layer.mlp.experts
1038
+ assert isinstance(moe._quant_method, QuarkW8A8Int8MoEMethod), (
1039
+ f"Expected QuarkW8A8Int8MoEMethod, got {type(moe._quant_method)}"
1040
+ )
1041
+ # Non-MoE linear layers should use QuarkW8A8Int8
1042
+ qkv_proj = layer.self_attn.qkv_proj
1043
+ assert isinstance(qkv_proj.scheme, QuarkW8A8Int8)
1044
1031
- layer = model.model.layers[0]
1032
- moe = layer.mlp.experts
1033
- assert isinstance(moe._quant_method, QuarkW8A8Int8MoEMethod), (
1034
- f"Expected QuarkW8A8Int8MoEMethod, got {type(moe._quant_method)}"
1035
- )
1036
- qkv_proj = layer.self_attn.qkv_proj
1037
- assert isinstance(qkv_proj.scheme, QuarkW8A8Int8)
1045
+ llm.apply_model(check_model)
1046
1039
- monkeypatch.setattr(Attention, "forward", lambda _, q, k, v: q.contiguous())
1040
- input_ids = torch.tensor([1, 2, 3, 4], device=DEVICE_TYPE)
1041
- positions = torch.arange(input_ids.numel(), device=DEVICE_TYPE)
1042
- with (
1043
- set_current_vllm_config(vllm_config),
1044
- set_forward_context(None, vllm_config, num_tokens=input_ids.numel()),
1045
- ):
1046
- hidden_states = model(input_ids, positions, None)
1047
- logits = model.compute_logits(hidden_states)
1048
- assert torch.isfinite(logits).all()
1047
+ output = llm.generate_greedy("Hello", max_tokens=4)
1048
+ assert output
1049
1050
1051
@pytest.mark.skipif(