From 7196981189f697e321a1cd9607d4d8d4b78592fb Mon Sep 17 00:00:00 2001 From: Ard <88337265+aw471@users.noreply.github.com> Date: Fri, 13 Nov 2026 14:54:14 -0800 Subject: [PATCH] Slice attention batches above the kernel limit in an operator CUDA attention kernels accept at most 65,535 batch entries per call, so `SDPA` slices larger batches. A compiled graph held one kernel call per slice of the batch it was traced with: a dynamic batch compiled a new graph for every number of slices. The slicing now runs inside the `sdm::sdpa` operator whenever the batch does not provably fit a single call. The operator slices when it runs, so one graph serves any batch. It returns a contiguous tensor, since the memory layout of a single slice depends on the attention kernel. Batches that provably fit and calls with gradients keep the plain call. The operator declares an output of the dtype of its query, and autocast is not visible where it declares it. `SDPA` therefore converts the inputs of the operator to the dtype that autocast computes in, as the kernel call does itself. --- sdm/nn/attention.py | 127 +++++++++++++++++++++++++++++++--------- test/nn/test_compile.py | 35 +++++++++++ 2 files changed, 135 insertions(+), 27 deletions(-) diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py index 2a6dc0f3c..a7e53d84a 100644 --- a/sdm/nn/attention.py +++ b/sdm/nn/attention.py @@ -9,6 +9,7 @@ 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 @@ -16,6 +17,74 @@ 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] diff --git a/test/nn/test_compile.py b/test/nn/test_compile.py index 1612bacda..8f2f2b2ee 100644 --- a/test/nn/test_compile.py +++ b/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"),