[Kernel] Decouple Tile Size from Block Size in Triton Unified Attention Kernel (#21197)

Signed-off-by: Jan van Lunteren <jvl@zurich.ibm.com>
2025-09-18 16:27:01 +02:00
parent bc19d75985
commit 01a583fea4
2 changed files with 70 additions and 55 deletions
--- a/tests/kernels/attention/test_triton_unified_attention.py
+++ b/tests/kernels/attention/test_triton_unified_attention.py
@@ -102,9 +102,6 @@ def test_triton_unified_attn(
 ) -> None:
    torch.set_default_device("cuda")

-    if q_dtype is not None and q_dtype.itemsize < 2 and block_size < 32:
-        pytest.skip("block size must be at least 32 for fp8")
-
    current_platform.seed_everything(0)
    num_seqs = len(seq_lens)
    query_lens = [x[0] for x in seq_lens]