[TRITON] Reduce MHA UTs

op_tests/triton_tests/attention/test_chunked_pa_prefill.py

@@ -9,12 +9,11 @@
 from aiter.ops.triton.utils.types import str_to_torch_dtype
 
 NUM_HEADS = [64]
-NUM_QUERIES_PER_KV = [1, 8, 64]
-HEAD_SIZES = [128, 96, 24]
+NUM_QUERIES_PER_KV = [1, 8]
+HEAD_SIZES = [128]
 DTYPES = [torch.float16]
-CUDA_DEVICES = [f"cuda:{i}" for i in range(1)]
-SLIDING_WINDOW = [0, 16, 64, 128, 256, 512, 2048]
-KV_CACHE_DTYPES = ["auto", "fp8e4m3", "fp8e5m2"]
+SLIDING_WINDOW = [0, 256, 1024]
+KV_CACHE_DTYPES = ["auto", "fp8e4m3"]
 
 
 def context_attention_fwd_torch(
@@ -316,7 +315,6 @@ def input_helper(
 @pytest.mark.parametrize("head_size", HEAD_SIZES)
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
-@pytest.mark.parametrize("device", CUDA_DEVICES)
 @pytest.mark.parametrize("sliding_window", SLIDING_WINDOW)
 @torch.inference_mode()
 def test_contexted_kv_attention(
@@ -326,8 +324,8 @@ def test_contexted_kv_attention(
     sliding_window: int,
     dtype: torch.dtype,
     kv_cache_dtype: str,
-    device: str,
 ) -> None:
+    device = "cuda:0"
     (
         query,
         k,
@@ -399,16 +397,15 @@ def test_contexted_kv_attention(
 @pytest.mark.parametrize("head_size", HEAD_SIZES)
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("kv_cache_dtype", KV_CACHE_DTYPES)
-@pytest.mark.parametrize("device", CUDA_DEVICES)
 @torch.inference_mode()
 def test_contexted_kv_attention_alibi(
     num_heads: int,
     num_queries_per_kv: int,
     head_size: int,
     dtype: torch.dtype,
     kv_cache_dtype: str,
-    device: str,
 ) -> None:
+    device = "cuda:0"
     (
         query,
         k,

op_tests/triton_tests/attention/test_fav3_sage.py

@@ -203,18 +203,17 @@ def input_helper(
 @pytest.mark.parametrize(
     "NUM_Q_HEADS, NUM_K_HEADS", [(1, 1), (16, 16), (2, 1), (48, 8)]
 )
-@pytest.mark.parametrize("HEAD_SZ", [128])
 @pytest.mark.parametrize("layout", ["bhsd", "bshd"])
 def test_sage(
     BATCH: int,
     SEQLEN_Q: int,
     SEQLEN_K: int,
     NUM_Q_HEADS: int,
     NUM_K_HEADS: int,
-    HEAD_SZ: int,
     layout: str,
     dtype=torch.bfloat16,
 ):
+    HEAD_SZ = 128
     torch.cuda.empty_cache()
 
     softmax_scale = 1.0 / math.sqrt(HEAD_SZ)
@@ -461,24 +460,21 @@ def test_sage_block_sparse_empty_kv_blocks(layout: str, dtype=torch.bfloat16):
 @pytest.mark.parametrize(
     "NUM_Q_HEADS, NUM_K_HEADS", [(1, 1), (16, 16), (2, 1), (48, 8)]
 )
-@pytest.mark.parametrize("HEAD_SZ", [128])
-@pytest.mark.parametrize("layout", ["bhsd"])
 @pytest.mark.parametrize("causal", [True, False])
 @pytest.mark.parametrize("qsmooth", [True, False])
-@pytest.mark.parametrize("hadamard_rotate", [True])  # TODO: hadamard expected to be on
 def test_sage_mxfp4(
     BATCH: int,
     SEQLEN_Q: int,
     SEQLEN_K: int,
     NUM_Q_HEADS: int,
     NUM_K_HEADS: int,
-    HEAD_SZ: int,
-    layout: str,
     causal: bool,
     qsmooth: bool,
-    hadamard_rotate: bool,
     dtype=torch.bfloat16,
 ):
+    HEAD_SZ = 128
+    layout = "bhsd"
+    hadamard_rotate = True  # hadamard expected to be on
 
     if not (arch_info.is_fp4_avail()):
         pytest.skip("MXFP4 not supported on this architecture")

op_tests/triton_tests/attention/test_flash_attn_kvcache.py

@@ -157,30 +157,24 @@ def _generate_block_kvcache(
     return k_cache, v_cache, block_table, k_cache_paged, v_cache_paged, num_blocks
 
 
-@pytest.mark.parametrize("dtype", [torch.bfloat16])
-@pytest.mark.parametrize("mha_type", ["mha", "mqa", "gqa"])
+@pytest.mark.parametrize("mha_type", ["mha", "gqa"])
 @pytest.mark.parametrize("new_kv", [False, True])
 @pytest.mark.parametrize("causal", [False, True])
 @pytest.mark.parametrize("seqlen_new_eq_seqlen_q", [True, False])
 @pytest.mark.parametrize("paged_kv_block_size", [None, 256])
 @pytest.mark.parametrize(
     "seqlen_q,seqlen_k",
     [
-        (1, 128),
         (1, 339),
         (3, 1024),
-        (64, 800),
         (3, 799),
         (64, 2048),
         (128, 128),
-        (1, 3131),
         (8, 3131),
         (1, 1024),
-        (3, 799),
-        (1, 339),
     ],
 )
-@pytest.mark.parametrize("d", [64, 128, 256])
+@pytest.mark.parametrize("d", [64, 128])
 def test_flash_attn_kvcache(
     seqlen_q,
     seqlen_k,
@@ -190,8 +184,8 @@ def test_flash_attn_kvcache(
     causal,
     new_kv,
     mha_type,
-    dtype,
 ):
+    dtype = torch.bfloat16
     if seqlen_q > seqlen_k and new_kv:
         pytest.skip()
 
@@ -340,13 +334,12 @@ def test_flash_attn_kvcache(
 @pytest.mark.parametrize("new_kv", [False, True])
 @pytest.mark.parametrize("causal", [True, False])
 @pytest.mark.parametrize("mha_type", ["mha", "gqa"])
-@pytest.mark.parametrize("d", [128])
 def test_flash_attn_kvcache_torch_compile(
-    d,
     mha_type,
     causal,
     new_kv,
 ):
+    d = 128
     device = "cuda"
     torch.random.manual_seed(SEED)
     torch.cuda.manual_seed(SEED)
@@ -412,8 +405,8 @@ def fn(q, k_cache, v_cache, k_new, v_new, cache_seqlens):
 
 @pytest.mark.parametrize("new_kv", [False, True])
 @pytest.mark.parametrize("mha_type", ["mha", "gqa"])
-@pytest.mark.parametrize("d", [128])
-def test_flash_attn_kvcache_hipgraph_capture(d, mha_type, new_kv):
+def test_flash_attn_kvcache_hipgraph_capture(mha_type, new_kv):
+    d = 128
     device = "cuda"
     torch.random.manual_seed(SEED)
     torch.cuda.manual_seed(SEED)

op_tests/triton_tests/attention/test_la.py

@@ -356,7 +356,7 @@ def test_persistent_lean_attention(
 
     assert batch == len(n_ctx)
     try:
-        sum_n_ctx = sum(int(n) for n in n_ctx)
+        sum(int(n) for n in n_ctx)
     except ValueError:
         print(f"N_CTX contains non-numeric values: {n_ctx}")
 
@@ -562,7 +562,7 @@ def main():
     assert batch == len(n_ctx)
 
     try:
-        sum_n_ctx = sum(int(n) for n in n_ctx)
+        sum(int(n) for n in n_ctx)
     except ValueError:
         print(f"N_CTX contains non-numeric values: {n_ctx}")
     print(f"causal={causal}, batch={batch}")