tests: bound default suite runtime to five minutes

README.md

@@ -182,6 +182,15 @@ python benchmarks/bench_kda_bwd_wy_dqkg_sm90.py --mode both --heads 32
 ## Tests
 
 ```bash
+# Time-bounded default suite (representative CUDA shapes, about five minutes on GB200)
+python -m pytest -q
+
+# Extended CUDA parameter sweeps only
+python -m pytest -m cula_slow -q
+
+# Default tests plus every extended CUDA case
+python -m pytest -m cula_full -q
+
 # Tests for modular KDA forward against FLA Triton implementation
 python -m pytest tests/test_kda_sm100_chunk_vs_fla.py -v
 # Tests for modular KDA forward against naive KDA reference
@@ -195,7 +204,7 @@ python -m pytest tests/test_lightning_attn_prefill_dispatch.py tests/test_lightn
 # Tests for Lightning Attention decode
 python -m pytest tests/test_lightning_decode.py -v
 
-# test_kda_sm100_chunk_vs_naive.py and test_kda_sm100_chunk_vs_fla.py support a fast/slow split.
+# The two chunk-vs-reference files also retain their narrower KDA fast/slow split.
 # Fast (default) — representative correctness paths for default CI and local iteration
 python -m pytest tests/test_kda_sm100_chunk_vs_naive.py tests/test_kda_sm100_chunk_vs_fla.py -v
 # Slow — broader stress coverage for nightly or manual runs

tests/__init__.py

@@ -0,0 +1,8 @@
+# Copyright 2025-2026 Ant Group Co., Ltd.
+# SPDX-License-Identifier: Apache-2.0
+
+"""cuLA test package.
+
+Keeping the tests importable as a package prevents unrelated site-packages
+named ``tests`` from shadowing shared test helpers.
+"""

tests/conftest.py

@@ -12,11 +12,30 @@ def _is_sm90() -> bool:
     return torch.cuda.is_available() and torch.cuda.get_device_capability()[0] == 9
 
 
+def _prepare_markexpr(markexpr: str) -> tuple[str, bool, bool]:
+    """Expand suite aliases and report whether extended CUDA cases are requested."""
+    include_cula_slow = bool(re.search(r"\b(cula_slow|cula_full)\b", markexpr))
+    include_kda_slow = bool(re.search(r"\b(kda_slow|kda_full|cula_slow|cula_full)\b", markexpr))
+    markexpr = re.sub(r"\bcula_full\b", "(cula_slow or not cula_slow)", markexpr)
+    markexpr = re.sub(r"\bcula_slow\b", "(cula_slow or kda_slow)", markexpr)
+    markexpr = re.sub(r"\bkda_full\b", "(kda_fast or kda_slow)", markexpr)
+    return markexpr, include_cula_slow, include_kda_slow
+
+
 def pytest_configure(config):
     config.addinivalue_line("markers", "sm100_only: only run on SM100 devices")
     config.addinivalue_line("markers", "sm90_only: only run on SM90 devices")
     config.addinivalue_line("markers", "benchmark: long-running benchmark-shaped coverage")
     config.addinivalue_line("markers", "sanitizer: tests intended to run under compute-sanitizer")
+    config.addinivalue_line(
+        "markers",
+        "cula_fast: representative case kept in the default, time-bounded test suite",
+    )
+    config.addinivalue_line(
+        "markers",
+        "cula_slow: extended CUDA coverage excluded from the default suite; "
+        "run with '-m cula_slow' or include everything with '-m cula_full'",
+    )
     config.addinivalue_line(
         "markers",
         "kda_fast: KDA test case included in fast (default) mode",
@@ -32,9 +51,8 @@ def pytest_configure(config):
         "variant in fast mode (other fast configs run disable_recompute=False only)",
     )
 
-    markexpr = config.option.markexpr
-    if markexpr and "kda_full" in markexpr:
-        config.option.markexpr = re.sub(r"\bkda_full\b", "(kda_fast or kda_slow)", markexpr)
+    markexpr, config._cula_include_slow, config._kda_include_slow = _prepare_markexpr(config.option.markexpr or "")
+    config.option.markexpr = markexpr
 
 
 def pytest_collection_modifyitems(config, items):
@@ -43,21 +61,26 @@ def pytest_collection_modifyitems(config, items):
     skip_non_sm100 = pytest.mark.skip(reason="SM100-only test: skip on non-SM100 devices")
     skip_non_sm90 = pytest.mark.skip(reason="SM90-only test: skip on non-SM90 devices")
 
-    marker_expr = config.option.markexpr or ""
-    include_slow = "kda_slow" in marker_expr
+    include_slow = config._kda_include_slow
+    include_cula_slow = config._cula_include_slow
     skip_slow = pytest.mark.skip(
         reason="kda_slow case: run 'pytest -m kda_slow' or the full sweep with '-m kda_full' to include"
     )
     skip_fast_norecomp = pytest.mark.skip(
         reason="disable_recompute=True runs in fast mode only for kda_fast_norecomp configs; "
         "include the rest via '-m kda_slow' or '-m kda_full'"
     )
+    skip_cula_slow = pytest.mark.skip(
+        reason="extended CUDA case: run 'pytest -m cula_slow' or the full suite with '-m cula_full'"
+    )
 
     for item in items:
         if "sm100_only" in item.keywords and not is_sm100:
             item.add_marker(skip_non_sm100)
         if "sm90_only" in item.keywords and not is_sm90:
             item.add_marker(skip_non_sm90)
+        if "cula_slow" in item.keywords and "cula_fast" not in item.keywords and not include_cula_slow:
+            item.add_marker(skip_cula_slow)
         if include_slow:
             continue
         if "kda_slow" in item.keywords:

tests/test_kda_conv_mtp.py

@@ -34,6 +34,8 @@
 
 from cula.ops.kda.decode.mtp_conv import NWARP, kda_conv_decode_mtp_verify
 
+pytestmark = pytest.mark.cula_slow
+
 W = 4  # KDA conv width (short_conv_kernel_size)
 
 
@@ -233,6 +235,7 @@ def test_conv_mtp_hv_eq_h_softplus(N, T, H, HV, variant):
 
 
 @pytest.mark.parametrize("variant", ["small_batch", "large_batch"])
+@pytest.mark.cula_fast
 def test_conv_mtp_softplus_large_input_is_finite(variant):
     N, T, H, HV, K, V = 1, 4, 8, 8, 128, 128
     scale = K**-0.5

tests/test_kda_conv_mtp_kvbuffer.py

@@ -12,6 +12,8 @@
     kda_flush_kvbuffer,
 )
 
+pytestmark = pytest.mark.cula_slow
+
 W = 4
 K = 128
 V = 128
@@ -209,6 +211,7 @@ def test_fused_without_conv_bias(gate):
         torch.testing.assert_close(actual, expected, atol=4e-3, rtol=6e-3)
 
 
+@pytest.mark.cula_fast
 def test_fused_tensor_core_single_qk_cta():
     result = _run_pair(2, 4, 8, 8, num_v_tiles=1, seed=64)
     _, out_base, out_fused, base_bufs, fused_bufs, conv_state, windows, rolled, windows_ref = result

tests/test_kda_decode.py

@@ -35,6 +35,8 @@
 
 from cula.kda import fused_sigmoid_gating_delta_rule_update, kda_decode
 
+pytestmark = pytest.mark.cula_slow
+
 
 # ---------------------------------------------------------------------------
 # PyTorch reference
@@ -223,8 +225,7 @@ def run_kda_decode_triton_compatible(q, k, v, a, b, A_log, dt_bias, state_kv, sc
 # ---------------------------------------------------------------------------
 # Tests: Dense layout
 # ---------------------------------------------------------------------------
-@pytest.mark.parametrize("N", [1, 4, 16, 32, 64, 128])
-@pytest.mark.parametrize("H,HV", [(8, 16), (16, 32)])
+@pytest.mark.parametrize("N,H,HV", [(1, 8, 16), (16, 16, 32), (64, 8, 16), (128, 16, 32)])
 def test_kda_decode_dense(N, H, HV):
     K, V = 128, 128
     scale = K**-0.5
@@ -263,8 +264,7 @@ def test_kda_decode_dense(N, H, HV):
 # ---------------------------------------------------------------------------
 # Tests: Varlen layout
 # ---------------------------------------------------------------------------
-@pytest.mark.parametrize("N", [4, 16, 32, 64, 128])
-@pytest.mark.parametrize("H,HV", [(8, 16), (16, 32)])
+@pytest.mark.parametrize("N,H,HV", [(4, 8, 16), (32, 16, 32), (128, 8, 16)])
 def test_kda_decode_varlen(N, H, HV):
     K, V = 128, 128
     scale = K**-0.5
@@ -340,6 +340,7 @@ def test_kda_decode_large_v(N):
 # ---------------------------------------------------------------------------
 # Tests: Zero initial state
 # ---------------------------------------------------------------------------
+@pytest.mark.cula_fast
 def test_kda_decode_zero_state():
     N, H, HV, K, V = 4, 8, 16, 128, 128
     scale = K**-0.5

tests/test_kda_decode_mtp.py

@@ -43,6 +43,8 @@
     kda_flush_kvbuffer,
 )
 
+pytestmark = pytest.mark.cula_slow
+
 
 def torch_kda_mtp_ref(
     q, k, v, a, b, A_log, dt_bias, state, scale, use_l2norm=True, softplus_beta=1.0, softplus_threshold=20.0, lower_bound=None
@@ -379,6 +381,7 @@ def launch():
         assert torch.equal(st_i, st_ref), f"{kernel} state non-deterministic at iter {i}"
 
 
+@pytest.mark.cula_fast
 def test_intermediate_disable_state_update():
     """disable_state_update leaves the pool untouched; snapshots still fire and match the oracle."""
     N, T, H, HV, K, V = 4, 4, 8, 16, 128, 128

tests/test_kda_packed_decode.py

@@ -35,6 +35,8 @@
 
 K = 128
 
+pytestmark = pytest.mark.cula_slow
+
 
 def _pack_mixed_qkv(q, k, v, N):
     """q,k: (N,H,K) bf16 ; v: (N,HV,V) bf16 -> mixed_qkv (N, qkv_dim) bf16."""
@@ -81,9 +83,10 @@ def _run_packed_dense(q, k, v, a, b, A_log, dt_bias, state, scale):
     return o.squeeze(1), state_p  # (N,HV,V), (N,HV,V,K)
 
 
-@pytest.mark.parametrize("N", [1, 2, 8, 16, 32, 64, 128])
-@pytest.mark.parametrize("H,HV", [(8, 16), (16, 32)])
-@pytest.mark.parametrize("V", [128, 256])
+@pytest.mark.parametrize(
+    "N,H,HV,V",
+    [(1, 8, 16, 128), (8, 16, 32, 128), (32, 8, 16, 256), (128, 16, 32, 256)],
+)
 def test_packed_dense(N, H, HV, V):
     scale = K**-0.5
     q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V)
@@ -138,9 +141,10 @@ def _run_packed_varlen(q, k, v, a, b, A_log, dt_bias, state, scale):
     return o.squeeze(0), state_p  # (N,HV,V), (N,HV,V,K)
 
 
-@pytest.mark.parametrize("N", [2, 8, 16, 32, 64, 128])
-@pytest.mark.parametrize("H,HV", [(8, 16), (16, 32)])
-@pytest.mark.parametrize("V", [128, 256])
+@pytest.mark.parametrize(
+    "N,H,HV,V",
+    [(2, 8, 16, 128), (8, 16, 32, 128), (32, 8, 16, 256), (128, 16, 32, 256)],
+)
 def test_packed_varlen(N, H, HV, V):
     scale = K**-0.5
     q, k, v, a, b, A_log, dt_bias, state = make_inputs(N, H, HV, K, V)
@@ -155,8 +159,7 @@ def test_packed_varlen(N, H, HV, V):
 # ---------------------------------------------------------------------------
 # Equal-head coverage: H == HV (group ratio = 1)
 # ---------------------------------------------------------------------------
-@pytest.mark.parametrize("N", [1, 8, 64])
-@pytest.mark.parametrize("H", [8, 16, 32, 64])
+@pytest.mark.parametrize("N,H", [(1, 8), (8, 32), (64, 64)])
 def test_packed_dense_equal_heads(N, H):
     HV, V = H, 128
     scale = K**-0.5
@@ -169,8 +172,7 @@ def test_packed_dense_equal_heads(N, H):
     _assert_close("state", state_ref, state_p)
 
 
-@pytest.mark.parametrize("N", [2, 8, 64])
-@pytest.mark.parametrize("H", [8, 16, 32, 64])
+@pytest.mark.parametrize("N,H", [(2, 8), (8, 32), (64, 64)])
 def test_packed_varlen_equal_heads(N, H):
     HV, V = H, 128
     scale = K**-0.5
@@ -328,6 +330,7 @@ def test_packed_kv_state_layout(is_varlen):
 # ---------------------------------------------------------------------------
 # No L2 norm path
 # ---------------------------------------------------------------------------
+@pytest.mark.cula_fast
 def test_packed_no_l2norm():
     N, H, HV, V = 8, 8, 16, 128
     scale = K**-0.5

tests/test_kda_sm100_chunk_vs_fla.py

@@ -35,6 +35,7 @@
 _SLOW = [pytest.mark.kda_slow]
 
 # (B, T, H, HV, D, gln, mask_p, l2norm, gate, safe_gate, dtype), marks
+# FLA chunk KDA requires H == HV; dedicated decode/conv tests cover GVA.
 _FIXED_CONFIGS = [
     ((1, 63, 1, 1, 128, 1, 0, False, False, True, torch.bfloat16), _FAST_NR),  # small fixed (+ no_recomp)
     ((2, 500, 3, 3, 128, 1, 0, False, False, True, torch.bfloat16), _SLOW),
@@ -44,18 +45,14 @@
     ((4, 1024, 4, 4, 128, 1, 0, True, False, True, torch.bfloat16), _FAST),  # l2norm medium
     ((2, 1500, 4, 4, 128, 10, 0, False, True, True, torch.bfloat16), _FAST_NR),  # gated (+ no_recomp)
     ((4, 2048, 8, 8, 128, 1, 0, False, True, True, torch.bfloat16), _SLOW),
-    # GVA cases: HV > H
-    ((2, 1024, 4, 8, 128, 1, 0, True, False, True, torch.bfloat16), _FAST),  # GVA medium
-    ((2, 1500, 2, 4, 128, 10, 0, False, True, True, torch.bfloat16), _SLOW),
-    ((2, 2048, 4, 8, 128, 1, 0, False, True, True, torch.bfloat16), _SLOW),
 ]
 
 # (H, HV, D, mask_p, cu_seqlens, dtype, safe_gate), marks
 _VARLEN_CONFIGS = [
-    ((4, 4, 128, 0.1, [0, 15], torch.bfloat16, True), _FAST),  # short varlen smoke
+    ((4, 4, 128, 0.1, [0, 15], torch.bfloat16, True), _SLOW),
     ((4, 4, 128, 0.9, [0, 256, 500, 1000], torch.bfloat16, True), _SLOW),
     ((4, 4, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), _SLOW),
-    ((4, 4, 128, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16, True), _FAST),  # multi-batch varlen
+    ((4, 4, 128, 0, [0, 15, 100, 300, 1200, 2000], torch.bfloat16, True), _SLOW),
     ((4, 4, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), _SLOW),
     # ======Varlen test with simulated trace=======
     (
@@ -106,22 +103,6 @@
         ),
         _SLOW,
     ),
-    # ======GVA varlen cases: HV > H=======
-    ((2, 4, 128, 0.1, [0, 15], torch.bfloat16, True), _FAST),  # GVA varlen smoke
-    ((4, 8, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), _SLOW),
-    ((4, 8, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), _SLOW),
-    (
-        (
-            8,
-            32,
-            128,
-            0,
-            [0, 494, 1004, 1561, 1908, 2240, 2849, 3116, 4096, 4986, 5626, 6090, 6718, 7244, 7870, 8192],
-            torch.bfloat16,
-            True,
-        ),
-        _SLOW,
-    ),
 ]
 
 

tests/test_kda_sm100_chunk_vs_naive.py

@@ -36,6 +36,7 @@
 _SLOW = [pytest.mark.kda_slow]
 
 # (B, T, H, HV, D, gln, mask_p, l2norm, gate, safe_gate, dtype), marks
+# The FLA naive reference requires H == HV; dedicated decode/conv tests cover GVA.
 _FIXED_CONFIGS = [
     ((1, 63, 1, 1, 128, 1, 0, False, False, True, torch.bfloat16), _FAST_NR),  # small fixed (+ no_recomp)
     ((2, 500, 3, 3, 128, 1, 0, False, False, True, torch.bfloat16), _SLOW),
@@ -45,10 +46,6 @@
     ((4, 1024, 4, 4, 128, 1, 0, True, False, True, torch.bfloat16), _FAST),  # l2norm medium
     ((2, 1500, 4, 4, 128, 10, 0, False, True, True, torch.bfloat16), _FAST_NR),  # gated (+ no_recomp)
     ((4, 2048, 8, 8, 128, 1, 0, False, True, True, torch.bfloat16), _SLOW),
-    # GVA cases: HV > H
-    ((2, 1024, 4, 8, 128, 1, 0, True, False, True, torch.bfloat16), _FAST),  # GVA medium
-    ((2, 1500, 2, 4, 128, 10, 0, False, True, True, torch.bfloat16), _SLOW),
-    ((2, 2048, 4, 8, 128, 1, 0, False, True, True, torch.bfloat16), _SLOW),
 ]
 
 # (H, HV, D, mask_p, cu_seqlens, dtype, safe_gate), marks
@@ -107,22 +104,6 @@
         ),
         _SLOW,
     ),
-    # ======GVA varlen cases: HV > H=======
-    ((2, 4, 128, 0.1, [0, 15], torch.bfloat16, True), _FAST),  # GVA varlen smoke
-    ((4, 8, 128, 0.5, [0, 256, 500, 1000], torch.bfloat16, True), _SLOW),
-    ((4, 8, 128, 0, [0, 100, 300, 1200, 3000, 4096], torch.bfloat16, True), _SLOW),
-    (
-        (
-            8,
-            32,
-            128,
-            0,
-            [0, 494, 1004, 1561, 1908, 2240, 2849, 3116, 4096, 4986, 5626, 6090, 6718, 7244, 7870, 8192],
-            torch.bfloat16,
-            True,
-        ),
-        _SLOW,
-    ),
 ]
 
 

tests/test_kda_sm100_intra_fused_cutedsl.py

@@ -20,6 +20,8 @@
 )
 from cula.ops.kda.sm100.recompute_wu import recompute_w_u_fwd
 
+pytestmark = pytest.mark.cula_slow
+
 
 def _requires_sm100():
     if not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 0):