[ROCm][Test] Fix AITER MXFP4 oracle contract (#50728)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Aug 3, 2026 at 15:53 UTC
e279f7158322d81f7b10ec9733fd7f9420d2fdc6
1 file changed
+11
-4
tests/kernels/moe/test_ocp_mx_moe.py
+11
-4
@@ -1203,8 +1203,8 @@ ROCM_BACKEND_CONFIGS = {
1203
"requires_gfx950": False,
1204
},
1205
"AITER_MXFP4_BF16": {
1206
- "activation": "SILU",
1207
- "rtol": 1.0,
1206
+ "activation": "SWIGLUOAI",
1207
+ "rtol": 0.1,
1208
"percent": 0.7,
1209
"requires_aiter": True,
1210
"requires_gfx950": True,
@@ -1337,6 +1337,13 @@ def test_rocm_mxfp4_moe_oracle(
1337
w2_quant = w2_quant_2d.reshape(num_experts, hidden_size, -1)
1338
w2_scale = w2_scale_2d.reshape(num_experts, hidden_size, -1)
1339
1340
+ # AITER conversion mutates checkpoint-format weights and scales in place.
1341
+ # Preserve their original layout for the reference calculation below.
1342
+ w13_quant_ref = w13_quant.clone()
1343
+ w13_scale_ref = w13_scale.clone()
1344
+ w2_quant_ref = w2_quant.clone()
1345
+ w2_scale_ref = w2_scale.clone()
1346
+
1347
w13_bias = torch.randn(
1348
num_experts, 2 * intermediate_size, dtype=dtype, device=device
1349
)
@@ -1453,10 +1460,10 @@ def test_rocm_mxfp4_moe_oracle(
1460
1461
# Dequantize weights for reference computation
1462
w13_dq = upcast_from_mxfp(
1456
- w13_quant.view(torch.uint8), w13_scale, torch.bfloat16, axis=-1
1463
+ w13_quant_ref.view(torch.uint8), w13_scale_ref, torch.bfloat16, axis=-1
1464
)
1465
w2_dq = upcast_from_mxfp(
1459
- w2_quant.view(torch.uint8), w2_scale, torch.bfloat16, axis=-1
1466
+ w2_quant_ref.view(torch.uint8), w2_scale_ref, torch.bfloat16, axis=-1
1467
)
1468
1469
# Determine activation type and layout