From 277c443279ef9524bff17086a429b7e7433d4d31 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Sun, 1 Nov 2026 22:26:32 -0800 Subject: [PATCH 1/3] Keep wide attention heads that are not a power of 2 off the Triton SDPA kernel The Triton SDPA kernel for head sizes that are not a power of 2 loads the whole head in one tile, and that tile is at most 256 wide. For a wider head, such as 640 in the mid-block attention of a video VAE encoder, it reads only the first 256 values of each query and key, and it never writes the rest of the output, which then holds whatever was in memory. The program runs and returns wrong numbers. The shape check that the kernel and the replacement pass share now rejects such heads, so the pass leaves the call to the regular CUDA lowering. Heads that are a power of 2 use the other kernel, which covers the whole head, so they still use Triton. test_sdpa_head_sizes lowers attention over head sizes on both sides of the limit and compares each result with eager. Two smaller tests check that the pass still gives heads of 512 and 1024 to Triton, and that a direct call with a wide head raises. --- backends/cuda/tests/test_cuda_export.py | 37 +++++++++++++++++++ .../tests/test_sdpa_splitk_replacement.py | 15 ++++++++ backends/cuda/tests/test_triton_sdpa.py | 6 +++ backends/cuda/triton/kernels/sdpa.py | 6 +++ 4 files changed, 64 insertions(+) diff --git a/backends/cuda/tests/test_cuda_export.py b/backends/cuda/tests/test_cuda_export.py index 637ff4036b1..52bf2f7677d 100644 --- a/backends/cuda/tests/test_cuda_export.py +++ b/backends/cuda/tests/test_cuda_export.py @@ -683,6 +683,43 @@ def test_sdpa_single_kernel(self): "SDPA single kernel operation export failed", ) + @unittest.skipIf( + torch.version.hip is not None, + "loads through the Python runtime, which the ROCm build does not include", + ) + def test_sdpa_head_sizes(self): + from executorch.runtime import Runtime + + class Attention(torch.nn.Module): + def forward(self, *qkv): + return tuple( + torch.nn.functional.scaled_dot_product_attention(*qkv[i : i + 3]) + for i in range(0, len(qkv), 3) + ) + + # Heads that are not a power of 2 and are above 256 are wider than one + # tile of the Triton kernel for such heads. + head_sizes = (96, 192, 256, 512, 288, 320, 384, 640, 768) + inputs = tuple( + torch.randn(1, 2, 64, size, dtype=torch.bfloat16) + for size in head_sizes + for _ in range(3) + ) + model = Attention() + edge = self._export_to_cuda_with_lower(model, inputs) + with tempfile.TemporaryDirectory() as directory: + path = os.path.join(directory, "model.pte") + with open(path, "wb") as f: + edge.to_executorch().write_to_file(f) + method = Runtime.get().load_program(path).load_method("forward") + outputs = method.execute(list(inputs)) + expected = model(*(x.cuda() for x in inputs)) + for size, output, reference in zip(head_sizes, outputs, expected, strict=True): + with self.subTest(head_size=size): + torch.testing.assert_close( + output.cpu(), reference.cpu(), atol=2e-2, rtol=2e-2 + ) + def test_triton_kernel_mode_off(self): """ Test CUDA export with triton_kernel_mode set to OFF for SDPA kernel. diff --git a/backends/cuda/tests/test_sdpa_splitk_replacement.py b/backends/cuda/tests/test_sdpa_splitk_replacement.py index e41ffd09201..96f081378bc 100644 --- a/backends/cuda/tests/test_sdpa_splitk_replacement.py +++ b/backends/cuda/tests/test_sdpa_splitk_replacement.py @@ -364,6 +364,21 @@ def test_named_inputs_are_not_supported(self): node.kwargs = {"query": query, "key": key, "value": value} self.assertFalse(ReplaceEdgeOpWithTritonOpPass._sdpa_kernel_supports(node)) + def test_wide_heads_are_supported_only_as_powers_of_2(self): + # A hand-built graph shows which calls the pass gives to the kernel + # without compiling the kernel. + from executorch.backends.cuda.triton.replacement_pass import ( + ReplaceEdgeOpWithTritonOpPass, + ) + + sizes = ((192, True), (264, False), (640, False), (512, True), (1024, True)) + for size, supported in sizes: + with self.subTest(head_size=size): + node = self._sdpa_node((1, 2, 64, size)) + self.assertEqual( + ReplaceEdgeOpWithTritonOpPass._sdpa_kernel_supports(node), supported + ) + if __name__ == "__main__": unittest.main() diff --git a/backends/cuda/tests/test_triton_sdpa.py b/backends/cuda/tests/test_triton_sdpa.py index 1fdb1dedef9..8b878af8a4e 100644 --- a/backends/cuda/tests/test_triton_sdpa.py +++ b/backends/cuda/tests/test_triton_sdpa.py @@ -573,6 +573,12 @@ def test_causal_lq_ne_lkv_rejected(self): with self.assertRaises(RuntimeError): self.sdpa(q, k, v, is_causal=True) + def test_wide_non_pow2_head_rejected(self): + """A non-pow2 head wider than one tile should raise, not be cut short.""" + q = torch.randn(1, 2, 64, 640, dtype=torch.bfloat16, device="cuda") + with self.assertRaises(RuntimeError): + self.sdpa(q, q, q) + def test_non_pow2_no_mask(self): """Non-pow2 head dim without mask should work (mask_ptr=0 path).""" B, H, Lq, Lk, D = 1, 4, 4, 64, 40 # D=40 is not pow2 diff --git a/backends/cuda/triton/kernels/sdpa.py b/backends/cuda/triton/kernels/sdpa.py index e01c0922cef..5edcd349a96 100644 --- a/backends/cuda/triton/kernels/sdpa.py +++ b/backends/cuda/triton/kernels/sdpa.py @@ -270,6 +270,12 @@ def _validate_qkv_shapes( raise RuntimeError( f"Head dimension must match across Q, K, V; got D_q={D_q}, D_k={D_k}, D_v={D_v}." ) + # Heads that are not a power of 2 must fit in one tile, which + # _next_power_of_2 caps at 256. + if not _is_power_of_2(D_q) and D_q > 256: + raise RuntimeError( + f"Head dimensions that are not a power of 2 must be at most 256; got D={D_q}." + ) # Key and Value sequence lengths must match if L_kv_k != L_kv_v: raise RuntimeError( From 47603ba92cd737fd74df4891c52cb7cad8696659 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 05:55:02 -0800 Subject: [PATCH 2/3] Tie the wide-head check to the tile size, and test wide heads end to end The check now compares the head size with the tile that _next_power_of_2 picks, so it stays right if that cap changes. The docstring of triton::sdpa states the limit for direct callers. A head of 640 is now exported through the regular lowering in the tests, which also checks the logged reason. The hand-built support test keeps only the heads that the kernel still takes, since some CI GPUs lack the shared memory to compile a head of 1024. --- .../cuda/tests/test_sdpa_splitk_replacement.py | 18 ++++++++++-------- backends/cuda/triton/kernels/sdpa.py | 12 +++++++----- 2 files changed, 17 insertions(+), 13 deletions(-) diff --git a/backends/cuda/tests/test_sdpa_splitk_replacement.py b/backends/cuda/tests/test_sdpa_splitk_replacement.py index 96f081378bc..dd2a08bf1d0 100644 --- a/backends/cuda/tests/test_sdpa_splitk_replacement.py +++ b/backends/cuda/tests/test_sdpa_splitk_replacement.py @@ -275,6 +275,10 @@ def test_value_head_size_other_than_query_is_not_replaced(self): q, v = _bf16(1, 2, 1, 16), _bf16(1, 2, 256, 32) self.assertStays(self._logs(q, _bf16(1, 2, 256, 16), v), "Head dimension") + def test_wide_head_that_is_not_a_power_of_2_is_not_replaced(self): + q = _bf16(1, 2, 64, 640) + self.assertStays(self._logs(q, q, q), "not a power of 2") + def test_key_value_batch_1_under_batch_2_is_not_replaced(self): q, kv = _bf16(2, 2, 1, 16), _bf16(1, 2, 256, 16) self.assertStays(self._logs(q, kv, kv), "Batch dimension must match") @@ -320,7 +324,7 @@ def test_decode_with_a_shared_key_value_head_still_uses_splitk(self): class TestSDPAKernelSupportCheck(unittest.TestCase): - """Calls that cannot be exported end to end, on hand-built graphs.""" + """Support checks on hand-built graphs, which compile nothing.""" def setUp(self): if importlib.util.find_spec("triton") is None: @@ -364,19 +368,17 @@ def test_named_inputs_are_not_supported(self): node.kwargs = {"query": query, "key": key, "value": value} self.assertFalse(ReplaceEdgeOpWithTritonOpPass._sdpa_kernel_supports(node)) - def test_wide_heads_are_supported_only_as_powers_of_2(self): - # A hand-built graph shows which calls the pass gives to the kernel - # without compiling the kernel. + def test_heads_that_fit_the_kernel_are_supported(self): + # Some CI GPUs lack the shared memory to compile heads of 1024. from executorch.backends.cuda.triton.replacement_pass import ( ReplaceEdgeOpWithTritonOpPass, ) - sizes = ((192, True), (264, False), (640, False), (512, True), (1024, True)) - for size, supported in sizes: + for size in (192, 256, 512, 1024): with self.subTest(head_size=size): node = self._sdpa_node((1, 2, 64, size)) - self.assertEqual( - ReplaceEdgeOpWithTritonOpPass._sdpa_kernel_supports(node), supported + self.assertTrue( + ReplaceEdgeOpWithTritonOpPass._sdpa_kernel_supports(node) ) diff --git a/backends/cuda/triton/kernels/sdpa.py b/backends/cuda/triton/kernels/sdpa.py index 5edcd349a96..f77a8cd67b3 100644 --- a/backends/cuda/triton/kernels/sdpa.py +++ b/backends/cuda/triton/kernels/sdpa.py @@ -270,11 +270,12 @@ def _validate_qkv_shapes( raise RuntimeError( f"Head dimension must match across Q, K, V; got D_q={D_q}, D_k={D_k}, D_v={D_v}." ) - # Heads that are not a power of 2 must fit in one tile, which - # _next_power_of_2 caps at 256. - if not _is_power_of_2(D_q) and D_q > 256: + # Heads that are not a power of 2 must fit in one tile, and + # _next_power_of_2 caps the tile at 256. + if not _is_power_of_2(D_q) and _next_power_of_2(D_q) < D_q: raise RuntimeError( - f"Head dimensions that are not a power of 2 must be at most 256; got D={D_q}." + "Head dimensions that are not a power of 2 must be at most " + f"{_next_power_of_2(D_q)}; got D={D_q}." ) # Key and Value sequence lengths must match if L_kv_k != L_kv_v: @@ -1279,7 +1280,8 @@ def sdpa( share K/V loads) based on a tile-utilization heuristic from FlashAttention. Args: - query: Query tensor [B, H_q, L_q, D], dtype torch.bfloat16 + query: Query tensor [B, H_q, L_q, D], dtype torch.bfloat16. D must be + a power of 2 or at most 256. key: Key tensor [B, H_kv, L_kv, D], dtype torch.bfloat16 value: Value tensor [B, H_kv, L_kv, D], dtype torch.bfloat16 attn_mask: Optional bool mask [B, 1, L_q, L_kv] (broadcast over heads) From 82b00ec0729a14730bb89dbd41c2bfa4bda0c4d5 Mon Sep 17 00:00:00 2001 From: Anthony Shoumikhin Date: Mon, 2 Nov 2026 07:14:10 -0800 Subject: [PATCH 3/3] Say only the head-size limit this change adds in the triton::sdpa docstring The docstring said any power-of-2 head works, but the power-of-2 kernel has its own shared-memory limits for very wide heads. It now states only the limit for heads that are not a power of 2, and the support test is named for what it checks: which heads the pass gives to the kernel. --- backends/cuda/tests/test_sdpa_splitk_replacement.py | 2 +- backends/cuda/triton/kernels/sdpa.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/backends/cuda/tests/test_sdpa_splitk_replacement.py b/backends/cuda/tests/test_sdpa_splitk_replacement.py index dd2a08bf1d0..e6e530552f2 100644 --- a/backends/cuda/tests/test_sdpa_splitk_replacement.py +++ b/backends/cuda/tests/test_sdpa_splitk_replacement.py @@ -368,7 +368,7 @@ def test_named_inputs_are_not_supported(self): node.kwargs = {"query": query, "key": key, "value": value} self.assertFalse(ReplaceEdgeOpWithTritonOpPass._sdpa_kernel_supports(node)) - def test_heads_that_fit_the_kernel_are_supported(self): + def test_heads_the_kernel_takes_are_supported(self): # Some CI GPUs lack the shared memory to compile heads of 1024. from executorch.backends.cuda.triton.replacement_pass import ( ReplaceEdgeOpWithTritonOpPass, diff --git a/backends/cuda/triton/kernels/sdpa.py b/backends/cuda/triton/kernels/sdpa.py index f77a8cd67b3..3dc2b45695e 100644 --- a/backends/cuda/triton/kernels/sdpa.py +++ b/backends/cuda/triton/kernels/sdpa.py @@ -1280,8 +1280,8 @@ def sdpa( share K/V loads) based on a tile-utilization heuristic from FlashAttention. Args: - query: Query tensor [B, H_q, L_q, D], dtype torch.bfloat16. D must be - a power of 2 or at most 256. + query: Query tensor [B, H_q, L_q, D], dtype torch.bfloat16. A D that + is not a power of 2 must be at most 256. key: Key tensor [B, H_kv, L_kv, D], dtype torch.bfloat16 value: Value tensor [B, H_kv, L_kv, D], dtype torch.bfloat16 attn_mask: Optional bool mask [B, 1, L_q, L_kv] (broadcast over heads)