[SparseMLA] Propagate benchmark failures

Let compile and runtime exceptions propagate from table and single benchmarks in both forward and decode paths. Reserve skip/NaN behavior for optional providers that are genuinely unavailable, and reject non-finite measured timings.

Author: sunnycase

python/test/regression/test_sparse_mla_autotune_configs.py

@@ -1,7 +1,10 @@
 import inspect
 import importlib.util
+import math
 from pathlib import Path
 
+import pytest
+
 
 def _load_sparse_mla_module():
     path = (Path(__file__).resolve().parents[2] / "tutorials" / "tle" / "deepseek_v32" / "02-sparse-mla.py")
@@ -74,3 +77,78 @@ def test_sparse_mla_bench_seed_is_explicit_and_shared():
     assert inspect.signature(module.run_bench_table).parameters["seed"].default == module.BENCH_DEFAULT_SEED
     assert inspect.signature(module.bench_sparse_mla_fwd).parameters["seed"].default == module.BENCH_DEFAULT_SEED
     assert "seed" in inspect.signature(module.benchmark_sparse_mla_fwd.fn).parameters
+
+
+def _run_forward_benchmark_callback(module, provider):
+    return module.benchmark_sparse_mla_fwd.fn(
+        B=1,
+        S=1,
+        SKV=1,
+        H=1,
+        HKV=1,
+        DQK=1,
+        DV=1,
+        topk=1,
+        provider=provider,
+        warmup=1,
+        rep=1,
+        tilelang_block_I=1,
+        tilelang_num_stages=1,
+        tilelang_threads=32,
+        input_mode="flashmla",
+        seed=1,
+    )
+
+
+def _stub_forward_benchmark_inputs(monkeypatch, module):
+    monkeypatch.setattr(module, "_get_bench_sparse_mla_inputs", lambda *args, **kwargs: (None, None, None, None))
+
+
+def _benchmark_failure(message):
+    def fail(*args, **kwargs):
+        raise RuntimeError(message)
+
+    return fail
+
+
+@pytest.mark.parametrize("provider", ["tle-flashmla-prefill", "tilelang"])
+def test_sparse_mla_benchmark_propagates_supported_provider_failure(monkeypatch, provider):
+    module = _load_sparse_mla_module()
+    _stub_forward_benchmark_inputs(monkeypatch, module)
+    monkeypatch.setattr(module, "_HAVE_TILELANG", True)
+    monkeypatch.setattr(module.triton.testing, "do_bench", _benchmark_failure("compile failed"))
+
+    with pytest.raises(RuntimeError, match="compile failed"):
+        _run_forward_benchmark_callback(module, provider)
+
+
+def test_sparse_mla_benchmark_skips_unavailable_optional_provider(monkeypatch):
+    module = _load_sparse_mla_module()
+    _stub_forward_benchmark_inputs(monkeypatch, module)
+    monkeypatch.setattr(module, "_HAVE_TILELANG", False)
+    monkeypatch.setattr(module.triton.testing, "do_bench", _benchmark_failure("unavailable provider was executed"))
+
+    result = _run_forward_benchmark_callback(module, "tilelang")
+
+    assert all(math.isnan(value) for value in result)
+
+
+def test_sparse_mla_single_benchmark_propagates_provider_failure(monkeypatch):
+    module = _load_sparse_mla_module()
+    _stub_forward_benchmark_inputs(monkeypatch, module)
+    monkeypatch.setattr(module, "triton_sparse_mla_fwd_interface", lambda *args, **kwargs: (None, None))
+    monkeypatch.setattr(module, "tle_sparse_mla_fwd_interface", _benchmark_failure("TLE compile failed"))
+    monkeypatch.setattr(module, "_bench_ms", lambda *args, **kwargs: 1.0)
+    monkeypatch.setattr(module, "_sparse_mla_tflops_from_topk_length", lambda *args, **kwargs: 1.0)
+
+    with pytest.raises(RuntimeError, match="TLE compile failed"):
+        module.bench_sparse_mla_fwd(B=1, S=1, SKV=1, H=1, HKV=1, DQK=1, DV=1, topk=1, check_outputs=False)
+
+
+def test_sparse_mla_benchmark_rejects_non_finite_timings(monkeypatch):
+    module = _load_sparse_mla_module()
+    _stub_forward_benchmark_inputs(monkeypatch, module)
+    monkeypatch.setattr(module.triton.testing, "do_bench", lambda *args, **kwargs: (float("nan"), 1.0, 2.0))
+
+    with pytest.raises(RuntimeError, match="non-finite timings"):
+        _run_forward_benchmark_callback(module, "triton")

python/tutorials/tle/deepseek_v32/02-sparse-mla.py

@@ -26,7 +26,9 @@
     from tilelang import language as T
 
     _HAVE_TILELANG = True
-except Exception:  # pragma: no cover - optional dependency
+except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
+    if exc.name != "tilelang":
+        raise
     tilelang = None
     T = None
     _HAVE_TILELANG = False
@@ -35,7 +37,9 @@
     import flash_mla
 
     _HAVE_FLASHMLA = True
-except Exception:  # pragma: no cover - optional dependency
+except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
+    if exc.name != "flash_mla":
+        raise
     flash_mla = None
     _HAVE_FLASHMLA = False
 
@@ -3103,7 +3107,10 @@ def _sparse_mla_tflops_from_topk_length(topk_length, H, DQK, DV, ms):
 
 def _bench_ms(fn, warmup=BENCH_DEFAULT_WARMUP_MS, rep=BENCH_DEFAULT_REP_MS):
     ms = triton.testing.do_bench(fn, warmup=warmup, rep=rep)
-    return float(ms if not isinstance(ms, tuple) else ms[0])
+    ms = float(ms if not isinstance(ms, tuple) else ms[0])
+    if not math.isfinite(ms):
+        raise RuntimeError(f"Benchmark returned a non-finite timing: {ms}")
+    return ms
 
 
 _DECODE_BENCH_PROVIDERS = (["triton", "tle", "tle-pipe-pipelined"] +
@@ -3356,9 +3363,13 @@ def run():
         )
     except Exception as exc:  # pragma: no cover - depends on runtime/resource limits
         print(f"[bench:{provider}] failed for "
-              f"(B={B}, S={S}, SKV={SKV}, H={H}, HKV={HKV}, DQK={DQK}, DV={DV}, topk={topk}): {exc}")
-        return float("nan"), float("nan"), float("nan")
-    return ms, max_ms, min_ms
+              f"(B={B}, S={S}, SKV={SKV}, H={H}, HKV={HKV}, DQK={DQK}, DV={DV}, topk={topk}): {exc}",
+              flush=True)
+        raise
+    result = (ms, max_ms, min_ms)
+    if not all(math.isfinite(float(value)) for value in result):
+        raise RuntimeError(f"Benchmark provider {provider!r} returned non-finite timings: {result}")
+    return result
 
 
 @triton.testing.perf_report(
@@ -3530,9 +3541,13 @@ def run():
         )
     except Exception as exc:  # pragma: no cover - depends on runtime/resource limits
         print(f"[decode-bench:{provider}] failed for "
-              f"(B={B}, S={S}, SKV={SKV}, H={H}, HKV={HKV}, DQK={DQK}, DV={DV}, topk={topk}): {exc}")
-        return float("nan"), float("nan"), float("nan")
-    return ms, max_ms, min_ms
+              f"(B={B}, S={S}, SKV={SKV}, H={H}, HKV={HKV}, DQK={DQK}, DV={DV}, topk={topk}): {exc}",
+              flush=True)
+        raise
+    result = (ms, max_ms, min_ms)
+    if not all(math.isfinite(float(value)) for value in result):
+        raise RuntimeError(f"Decode benchmark provider {provider!r} returned non-finite timings: {result}")
+    return result
 
 
 def run_bench_table(warmup=BENCH_DEFAULT_WARMUP_MS, rep=BENCH_DEFAULT_REP_MS, show_plots=False, tilelang_block_I=64,
@@ -3987,7 +4002,8 @@ def run_tle():
         tle_tflops = _sparse_mla_tflops_from_topk_length(topk_length, H, DQK, DV, tle_ms)
         results.append(("tle", tle_ms, tle_tflops))
     except Exception as exc:  # pragma: no cover - depends on tle/runtime constraints
-        print(f"TLE bench skipped due to compile/runtime error: {exc}")
+        print(f"TLE bench failed due to compile/runtime error: {exc}", flush=True)
+        raise
 
     def run_tle_pipe():
         return tle_pipe_sparse_mla_fwd_interface(q, kv, indices, topk_length=topk_length, sm_scale=sm_scale, d_v=DV,
@@ -3999,7 +4015,8 @@ def run_tle_pipe():
         tle_pipe_tflops = _sparse_mla_tflops_from_topk_length(topk_length, H, DQK, DV, tle_pipe_ms)
         results.append(("tle-pipe-pipelined", tle_pipe_ms, tle_pipe_tflops))
     except Exception as exc:  # pragma: no cover - depends on tle/runtime constraints
-        print(f"TLE pipe-pipelined bench skipped due to compile/runtime error: {exc}")
+        print(f"TLE pipe-pipelined bench failed due to compile/runtime error: {exc}", flush=True)
+        raise
 
     def run_tle_flashmla_prefill():
         return tle_flashmla_prefill_interface(q, kv, indices, topk_length=topk_length, sm_scale=sm_scale, d_v=DV,
@@ -4012,7 +4029,8 @@ def run_tle_flashmla_prefill():
                                                                           tle_flashmla_prefill_ms)
         results.append(("tle-flashmla-prefill", tle_flashmla_prefill_ms, tle_flashmla_prefill_tflops))
     except Exception as exc:  # pragma: no cover - depends on tle/runtime constraints
-        print(f"TLE FlashMLA-prefill bench skipped due to compile/runtime error: {exc}")
+        print(f"TLE FlashMLA-prefill bench failed due to compile/runtime error: {exc}", flush=True)
+        raise
 
     if _HAVE_TILELANG:
         resolved_block_i = _resolve_tilelang_block_i(topk, tilelang_block_I)
@@ -4040,7 +4058,8 @@ def run_tilelang():
             tilelang_tflops = _sparse_mla_tflops_from_topk_length(topk_length, H, DQK, DV, tilelang_ms)
             results.append(("tilelang", tilelang_ms, tilelang_tflops))
         except Exception as exc:  # pragma: no cover - depends on tilelang/runtime constraints
-            print(f"TileLang bench skipped due to compile/runtime error: {exc}")
+            print(f"TileLang bench failed due to compile/runtime error: {exc}", flush=True)
+            raise
     else:
         print("TileLang is not installed, skip TileLang bench.")
 
@@ -4067,7 +4086,8 @@ def run_tilelang_pipelined():
                                                                             tilelang_pipelined_ms)
             results.append(("tilelang-pipelined", tilelang_pipelined_ms, tilelang_pipelined_tflops))
         except Exception as exc:  # pragma: no cover - depends on tilelang/runtime constraints
-            print(f"Pipelined TileLang bench skipped due to compile/runtime error: {exc}")
+            print(f"Pipelined TileLang bench failed due to compile/runtime error: {exc}", flush=True)
+            raise
 
         def run_tilelang_seesaw():
             return tilelang_sparse_mla_fwd_seesaw_interface(
@@ -4089,7 +4109,8 @@ def run_tilelang_seesaw():
             tilelang_seesaw_tflops = _sparse_mla_tflops_from_topk_length(topk_length, H, DQK, DV, tilelang_seesaw_ms)
             results.append(("tilelang-seesaw", tilelang_seesaw_ms, tilelang_seesaw_tflops))
         except Exception as exc:  # pragma: no cover - depends on tilelang/runtime constraints
-            print(f"Seesaw TileLang bench skipped due to compile/runtime error: {exc}")
+            print(f"Seesaw TileLang bench failed due to compile/runtime error: {exc}", flush=True)
+            raise
 
     if _HAVE_FLASHMLA:
         try:
@@ -4111,7 +4132,8 @@ def run_flashmla():
             flashmla_tflops = _sparse_mla_tflops_from_topk_length(topk_length, H, DQK, DV, flashmla_ms)
             results.append(("flashmla", flashmla_ms, flashmla_tflops))
         except Exception as exc:  # pragma: no cover - depends on flashmla/runtime constraints
-            print(f"FlashMLA bench skipped due to compile/runtime error: {exc}")
+            print(f"FlashMLA bench failed due to compile/runtime error: {exc}", flush=True)
+            raise
     else:
         print("FlashMLA is not installed, skip FlashMLA bench.")
 
@@ -4254,7 +4276,8 @@ def run_tle():
         tle_tflops = _sparse_mla_tflops_from_topk_length(inputs["topk_length_flat"], H, DQK, DV, tle_ms)
         results.append(("tle", tle_ms, tle_tflops))
     except Exception as exc:  # pragma: no cover - depends on tle/runtime constraints
-        print(f"TLE decode bench skipped due to compile/runtime error: {exc}")
+        print(f"TLE decode bench failed due to compile/runtime error: {exc}", flush=True)
+        raise
 
     def run_tle_pipe():
         return tle_pipe_sparse_mla_fwd_interface(
@@ -4274,7 +4297,8 @@ def run_tle_pipe():
         tle_pipe_tflops = _sparse_mla_tflops_from_topk_length(inputs["topk_length_flat"], H, DQK, DV, tle_pipe_ms)
         results.append(("tle-pipe-pipelined", tle_pipe_ms, tle_pipe_tflops))
     except Exception as exc:  # pragma: no cover - depends on tle/runtime constraints
-        print(f"TLE pipe-pipelined decode bench skipped due to compile/runtime error: {exc}")
+        print(f"TLE pipe-pipelined decode bench failed due to compile/runtime error: {exc}", flush=True)
+        raise
 
     if _HAVE_TILELANG:
         resolved_block_i = _resolve_tilelang_block_i(topk, tilelang_block_I)
@@ -4300,7 +4324,8 @@ def run_tilelang():
             tilelang_tflops = _sparse_mla_tflops_from_topk_length(inputs["topk_length_flat"], H, DQK, DV, tilelang_ms)
             results.append(("tilelang", tilelang_ms, tilelang_tflops))
         except Exception as exc:  # pragma: no cover - depends on tilelang/runtime constraints
-            print(f"TileLang decode bench skipped due to compile/runtime error: {exc}")
+            print(f"TileLang decode bench failed due to compile/runtime error: {exc}", flush=True)
+            raise
 
         def run_tilelang_pipelined():
             return tilelang_sparse_mla_fwd_pipelined_interface(
@@ -4324,7 +4349,8 @@ def run_tilelang_pipelined():
                                                                             tilelang_pipelined_ms)
             results.append(("tilelang-pipelined", tilelang_pipelined_ms, tilelang_pipelined_tflops))
         except Exception as exc:  # pragma: no cover - depends on tilelang/runtime constraints
-            print(f"Pipelined TileLang decode bench skipped due to compile/runtime error: {exc}")
+            print(f"Pipelined TileLang decode bench failed due to compile/runtime error: {exc}", flush=True)
+            raise
 
         def run_tilelang_seesaw():
             return tilelang_sparse_mla_fwd_seesaw_interface(
@@ -4348,7 +4374,8 @@ def run_tilelang_seesaw():
                                                                          tilelang_seesaw_ms)
             results.append(("tilelang-seesaw", tilelang_seesaw_ms, tilelang_seesaw_tflops))
         except Exception as exc:  # pragma: no cover - depends on tilelang/runtime constraints
-            print(f"Seesaw TileLang decode bench skipped due to compile/runtime error: {exc}")
+            print(f"Seesaw TileLang decode bench failed due to compile/runtime error: {exc}", flush=True)
+            raise
     else:
         print("TileLang is not installed, skip TileLang decode bench.")
 
@@ -4372,7 +4399,8 @@ def run_flashmla():
             flashmla_tflops = _sparse_mla_tflops_from_topk_length(inputs["topk_length_flat"], H, DQK, DV, flashmla_ms)
             results.append(("flashmla", flashmla_ms, flashmla_tflops))
         except Exception as exc:  # pragma: no cover - depends on flashmla/runtime constraints
-            print(f"FlashMLA decode bench skipped due to compile/runtime error: {exc}")
+            print(f"FlashMLA decode bench failed due to compile/runtime error: {exc}", flush=True)
+            raise
     else:
         print("FlashMLA is not installed, skip FlashMLA decode bench.")