A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[CI] Retry flaky fp8 cutlass mla tests (#24536)
Signed-off-by: Nick Hill <nhill@redhat.com>
Nick Hill committed
Sep 9, 2025 at 20:33 UTC
7e7db04310922f45524d4bfbe72f1dc7205ae7d6
1 file changed
+7
-1
tests/kernels/test_cutlass_mla_decode.py
+7
-1
@@ -49,7 +49,13 @@ CUTLASS_MLA_UNSUPPORTED_REASON = \
49
@pytest.mark.parametrize("block_size", [64])
50
@pytest.mark.parametrize("causal", [True])
51
@pytest.mark.parametrize("varlen", [False, True])
52
-@pytest.mark.parametrize("torch_dtype", [torch.bfloat16, torch.float8_e4m3fn])
52
+@pytest.mark.parametrize(
53
+ "torch_dtype",
54
+ [
55
+ torch.bfloat16,
56
+ # fp8 can have occasional precision-related failures.
57
+ pytest.param(torch.float8_e4m3fn, marks=pytest.mark.flaky(reruns=2))
58
+ ])
59
@torch.inference_mode()
60
def test_cutlass_mla_decode(b, s_q, mean_sk, h_q, h_kv, d, dv, block_size,
61
causal, varlen, torch_dtype):