[Bugfix][Tests] Stabilize B12X linear kernel checks (#54996)
Signed-off-by: Luke Alonso <lalonso@gmail.com>
Luke Alonso committed
Sep 2, 2026 at 11:46 UTC
1945a9457563e28a32baa03dca620ed52bd18d10
2 files changed
+9
-2
tests/kernels/quantization/test_block_fp8.py
+4
-2
@@ -406,5 +406,7 @@ def test_w8a8_block_fp8_b12x_matmul(M, N, K):
406
ref_out.float().flatten(),
407
dim=0,
408
)
409
- assert rel_diff < 0.002
410
- assert cosine >= 0.9999
409
+ # Four-way split-K uses atomic BF16 reductions, so nondeterministic atomic
410
+ # ordering can flap an output between adjacent BF16 values one ULP apart.
411
+ assert rel_diff < 0.003
412
+ assert cosine >= 0.99999
tests/model_executor/kernels/test_b12x_linear.py
+5
@@ -785,6 +785,11 @@ def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None:
785
786
monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
787
monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
788
+ monkeypatch.setitem(
789
+ linear_mod._POSSIBLE_NVFP4_KERNELS,
790
+ PlatformEnum.CUDA,
791
+ [B12xNvFp4LinearKernel, MarlinNvFp4LinearKernel],
792
+ )
793
monkeypatch.setattr(
794
MarlinNvFp4LinearKernel,
795
"is_supported",