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..e6e530552f2 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,6 +368,19 @@ 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_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, + ) + + for size in (192, 256, 512, 1024): + with self.subTest(head_size=size): + node = self._sdpa_node((1, 2, 64, size)) + self.assertTrue( + ReplaceEdgeOpWithTritonOpPass._sdpa_kernel_supports(node) + ) + 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..3dc2b45695e 100644 --- a/backends/cuda/triton/kernels/sdpa.py +++ b/backends/cuda/triton/kernels/sdpa.py @@ -270,6 +270,13 @@ 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, 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( + "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: raise RuntimeError( @@ -1273,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. 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)