[Hardware][AMD][CI] Fix Kernels Quantization test timeout (#46859)
Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com>
Matt committed
Jun 26, 2026 at 15:19 UTC
95e6442a6b6973f827783d162709627304cd13f2
1 file changed
+29
-28
tests/kernels/quantization/test_nvfp4_emulation.py
+29
-28
@@ -43,6 +43,26 @@ else:
43
return False
44
45
46
+MOE_MODEL_CONFIGS = {
47
+ "nvidia/Qwen3-30B-A3B-NVFP4": {
48
+ "shards": ["model-00001-of-00004.safetensors"],
49
+ "expert_prefix": "model.layers.9.mlp.experts.",
50
+ # Position of the expert index in the dot-split key.
51
+ "expert_idx_pos": 5,
52
+ }
53
+}
54
+
55
+
56
+@pytest.fixture(scope="module")
57
+def loaded_model_files():
58
+ return {
59
+ model_id: huggingface_hub.snapshot_download(
60
+ repo_id=model_id, allow_patterns=config["shards"]
61
+ )
62
+ for model_id, config in MOE_MODEL_CONFIGS.items()
63
+ }
64
+
65
+
66
class Nvfp4QuantizationEmulationTritonExpertsReference(TritonExperts):
67
"""
68
Extension of TritonExperts to support emulated NVFP4 MoE experts.
@@ -193,17 +213,15 @@ def test_nvfp4_emulation_support_check_rejects_bias_and_lora(
213
not current_platform.is_cuda_alike(),
214
reason="Triton NVFP4 kernel requires CUDA.",
215
)
196
-def test_triton_dequantize_nvfp4(monkeypatch) -> None:
216
+def test_triton_dequantize_nvfp4(monkeypatch, loaded_model_files) -> None:
217
"""Test the Triton dequantization kernel against the CPU reference
218
using real NVFP4 weights from a checkpoint.
219
220
Tests both 2D (attention projection) and 3D (stacked MoE experts).
221
"""
202
- checkpoint_path = huggingface_hub.snapshot_download(
203
- "nvidia/Qwen3-30B-A3B-NVFP4",
204
- allow_patterns=["model-00001-of-00004.safetensors"],
205
- )
206
- shard_path = f"{checkpoint_path}/model-00001-of-00004.safetensors"
222
+ checkpoint_path = loaded_model_files["nvidia/Qwen3-30B-A3B-NVFP4"]
223
+ shards = cast(list[str], MOE_MODEL_CONFIGS["nvidia/Qwen3-30B-A3B-NVFP4"]["shards"])
224
+ shard_path = f"{checkpoint_path}/{shards[0]}"
225
block_size = 16
226
227
with safe_open(shard_path, framework="pt", device="cpu") as f:
@@ -481,25 +499,8 @@ def test_triton_nvfp4_quant_dequant(
499
print(f" speedup: {speedup:.2f}x")
500
501
484
-MOE_MODEL_CONFIGS = {
485
- "nvidia/Qwen3-30B-A3B-NVFP4": {
486
- "shards": ["model-00001-of-00004.safetensors"],
487
- "expert_prefix": "model.layers.9.mlp.experts.",
488
- # Position of the expert index in the dot-split key.
489
- "expert_idx_pos": 5,
490
- },
491
- "nvidia/Kimi-K2.6-NVFP4": {
492
- "shards": [
493
- "model-00001-of-00060.safetensors",
494
- "model-00002-of-00060.safetensors",
495
- ],
496
- "expert_prefix": "language_model.model.layers.1.mlp.experts.",
497
- "expert_idx_pos": 6,
498
- },
499
-}
500
-
501
-
502
def _load_nvfp4_moe_weights(
503
+ model_files: dict[str, str],
504
model_id: str,
505
tensor_parallel_size: int,
506
max_experts: int | None = None,
@@ -518,10 +519,8 @@ def _load_nvfp4_moe_weights(
519
"""
520
cfg = MOE_MODEL_CONFIGS[model_id]
521
shards = cast(list[str], cfg["shards"])
521
- checkpoint_path = huggingface_hub.snapshot_download(
522
- model_id,
523
- allow_patterns=shards,
524
- )
522
+ checkpoint_path = model_files[model_id]
523
+
524
expert_prefix = cfg["expert_prefix"]
525
idx_pos = cast(int, cfg["expert_idx_pos"])
526
@@ -636,6 +635,7 @@ def _load_nvfp4_moe_weights(
635
[pytest.param(val, id=f"tensor_parallel_size:{val}") for val in [1, 2, 4, 8]],
636
)
637
def test_nvfp4_moe_correctness(
638
+ loaded_model_files,
639
num_tokens: int,
640
top_k: int,
641
model_id: str,
@@ -660,6 +660,7 @@ def test_nvfp4_moe_correctness(
660
hidden_dim,
661
intermediate_size,
662
) = _load_nvfp4_moe_weights(
663
+ loaded_model_files,
664
model_id,
665
tensor_parallel_size,
666
max_experts=num_test_experts,