[Dynamic shapes 4/n] Slice attention batches above the kernel limit in an operator

sdm/nn/attention.py

@@ -9,13 +9,82 @@
 import torch
 import torch.nn.functional as F
 from torch import Tensor
+from torch.fx.experimental.symbolic_shapes import statically_known_true
 from torch.nn import Linear
 
 from sdm._memory import _split_tensor, chunk_memory_limit
 from sdm.cache import KVCacheEntry
 from sdm.nn import QueryScaling
 
 
+def _sdpa(
+    query: Tensor,  # [B, Hq, Q, C]
+    key: Tensor,  # [B, Hkv, KV, C]
+    value: Tensor,  # [B, Hkv, KV, C]
+    attn_mask: Tensor | None,  # [B, Hq, Q, KV]
+    dropout_p: float,
+    enable_gqa: bool,
+    is_causal: bool,
+    scale: float | None,
+) -> Tensor:  # [B, Q, Hq, C]
+    # CUDA attention kernels accept at most 65,535 batch entries per call, so
+    # a larger batch is sliced. A batch that provably fits is not split:
+    # splitting a dynamic size would add a guard to a compiled graph.
+    queries = _split_tensor(query, 65_535, dim=0)
+    masks = (
+        (None,) * len(queries)
+        if attn_mask is None
+        else _split_tensor(attn_mask, 65_535, dim=0)
+    )
+    outs = []
+    for chunk_query, chunk_key, chunk_value, chunk_mask in zip(
+        queries,
+        _split_tensor(key, 65_535, dim=0),
+        _split_tensor(value, 65_535, dim=0),
+        masks,
+        strict=True,
+    ):
+        out = F.scaled_dot_product_attention(
+            query=chunk_query,
+            key=chunk_key,
+            value=chunk_value,
+            attn_mask=chunk_mask,
+            dropout_p=dropout_p,
+            enable_gqa=enable_gqa,
+            is_causal=is_causal,
+            scale=scale,
+        )
+        outs.append(out.transpose(-3, -2))
+    return outs[0] if len(outs) == 1 else torch.cat(outs)
+
+
+@torch.library.custom_op("sdm::sdpa", mutates_args=())
+def _sdpa_operator(
+    query: Tensor,
+    key: Tensor,
+    value: Tensor,
+    attn_mask: Tensor | None,
+    dropout_p: float,
+    enable_gqa: bool,
+    is_causal: bool,
+    scale: float | None,
+) -> Tensor:
+    """Slice the batch when the operator runs, not when a graph is traced."""
+    # The memory layout of a single slice depends on the attention kernel.
+    return _sdpa(
+        query, key, value, attn_mask, dropout_p, enable_gqa, is_causal, scale
+    ).contiguous()
+
+
+@_sdpa_operator.register_fake
+def _sdpa_operator_fake(
+    query: Tensor, key: Tensor, value: Tensor, *args: Any, **kwargs: Any
+) -> Tensor:
+    return query.new_empty(
+        (query.size(0), query.size(2), query.size(1), value.size(-1))
+    )
+
+
 class SDPA(torch.nn.Module):
     r"""Scaled Dot-Product Attention (SDPA).
 
@@ -189,34 +258,38 @@ def forward(
         key = key.transpose(-3, -2)  # [B, Hkv, KV, C]
         value = value.transpose(-3, -2)  # [B, Hkv, KV, C]
 
-        # CUDA attention kernels accept at most 65,535 batch entries per call.
-        # Slicing a dynamic batch that provably fits would add a size guard.
-        queries = _split_tensor(query, 65_535, dim=0)
-        masks = (
-            (None,) * len(queries)
-            if attn_mask is None
-            else _split_tensor(attn_mask, 65_535, dim=0)
-        )
-        outs = []
-        for chunk_query, chunk_key, chunk_value, chunk_mask in zip(
-            queries,
-            _split_tensor(key, 65_535, dim=0),
-            _split_tensor(value, 65_535, dim=0),
-            masks,
-            strict=True,
+        # Under `torch.compile`, the slicing loop of `_sdpa` is unrolled into
+        # the graph, which then serves only batches with that many slices.
+        # `_sdpa_operator` runs the same loop but is traced as one call, so
+        # one graph serves any batch size. It is skipped with gradients, which
+        # it lacks, and when the batch provably fits one kernel call.
+        sdpa = _sdpa
+        if not torch.is_grad_enabled() and not statically_known_true(
+            query.size(0) <= 65_535
         ):
-            out = F.scaled_dot_product_attention(
-                query=chunk_query,
-                key=chunk_key,
-                value=chunk_value,
-                attn_mask=chunk_mask,  # [B, Hq, Q, KV]
-                dropout_p=self.dropout if self.training else 0.0,
-                enable_gqa=enable_gqa,
-                is_causal=is_causal,
-                scale=self.scale,
-            )
-            outs.append(out.transpose(-3, -2))  # [B, Q, Hq, C]
-        out = outs[0] if len(outs) == 1 else torch.cat(outs)
+            sdpa = _sdpa_operator
+            device_type = query.device.type
+            if torch.is_autocast_enabled(device_type):
+                # Autocast converts the inputs of the attention kernel to its
+                # 16-bit dtype. `torch.compile` applies autocast while it
+                # traces and does not trace into the operator, so the kernel
+                # inside it would get unconverted inputs. They are converted
+                # here instead. As in autocast, 64-bit inputs stay as they are.
+                dtype = torch.get_autocast_dtype(device_type)
+                query, key, value = (
+                    x if x.dtype == torch.float64 else x.to(dtype)
+                    for x in (query, key, value)
+                )
+        out = sdpa(
+            query,
+            key,
+            value,
+            attn_mask,  # [B, Hq, Q, KV]
+            self.dropout if self.training else 0.0,
+            enable_gqa,
+            is_causal,
+            self.scale,
+        )  # [B, Q, Hq, C]
 
         return out.view(batch_shape + out.size()[-3:])  # [..., Q, Hq, C]
 

test/nn/test_compile.py

@@ -7,6 +7,7 @@
 import pytest
 import torch
 from torch import Tensor
+from torch._dynamo.testing import CompileCounter
 
 from sdm.nn import (
     SDPA,
@@ -184,6 +185,40 @@ def test_sdpa_compile_dynamic_batch() -> None:
     torch.testing.assert_close(out, expected)
 
 
+def test_sdpa_compile_batch_above_kernel_limit() -> None:
+    module = SDPA(num_query_heads=2).eval()
+    counter = CompileCounter()
+    compiled = torch.compile(module, fullgraph=True, backend=counter)
+    with torch.inference_mode():
+        for batch_size in [3, 65_536]:
+            query = torch.randn(batch_size, 4, 2, 8)
+            key = torch.randn(batch_size, 5, 2, 8)
+            value = torch.randn_like(key)
+            for tensor in (query, key, value):
+                torch._dynamo.mark_dynamic(tensor, 0)
+            expected = module(query=query, key=key, value=value)
+            out = compiled(query=query, key=key, value=value)
+            torch.testing.assert_close(out, expected)
+    assert counter.frame_count == 1
+
+
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_sdpa_compile_dynamic_batch_autocast(dtype: torch.dtype) -> None:
+    module = SDPA(num_query_heads=2).eval()
+    query = torch.randn(3, 4, 2, 8, dtype=dtype)
+    key = torch.randn(3, 5, 2, 8, dtype=dtype)
+    value = torch.randn_like(key)
+    for tensor in (query, key, value):
+        torch._dynamo.mark_dynamic(tensor, 0)
+    with torch.inference_mode(), torch.autocast("cpu", torch.bfloat16):
+        expected = module(query=query, key=key, value=value)
+        out = torch.compile(module, fullgraph=True)(
+            query=query, key=key, value=value
+        )
+    assert out.dtype == expected.dtype
+    torch.testing.assert_close(out, expected)
+
+
 @withCUDA
 @pytest.mark.parametrize(
     ("num_key_value_heads", "qassmax", "self_attn"),