diff --git a/openwam/model/kernels/attention.py b/openwam/model/kernels/attention.py index a0f5b378..50b8498d 100644 --- a/openwam/model/kernels/attention.py +++ b/openwam/model/kernels/attention.py @@ -10,6 +10,11 @@ cuDNN attention has no deterministic implementation (PyTorch finds no kernel under ``torch.use_deterministic_algorithms(True)``), so in deterministic mode ``sdpa`` is the default call. Off by default: PyTorch's backend selection. + +When cuDNN is preferred, calls that autograd will differentiate pad K/V to a +multiple of the cuDNN key tile (see ``_pad_kv_for_cudnn``): the backward kernels +shipped with PyTorch 2.7 can otherwise return NaN query gradients when the key +length is not a multiple of 128. """ from __future__ import annotations @@ -71,6 +76,51 @@ def _zero_masked_rows(out: torch.Tensor, nonempty: torch.Tensor) -> torch.Tensor return out.permute(order).masked_fill(empty, 0.0).permute(inverse) +# The cuDNN SDPA backward bundled with PyTorch 2.7 (cudnn-frontend 1.10, cuDNN 9.7) computes the last key tile +# without a padding mask when s_kv is not a multiple of 128, so the out-of-range key columns score 0 there. The +# forward pass is unaffected, but for a query row whose visible logits are all very negative (logsumexp below about +# -88.7, e.g. keys sharing a large component) the backward evaluates exp(0 - lse) = inf and inf * 0 = NaN in that +# row of dQ, which then poisons every parameter upstream of Q (NVIDIA/cudnn-frontend#160, pytorch/pytorch#166211; +# cudnn-frontend >= 1.15 applies a padding mask itself). Padding K/V with zero keys that the mask hides removes the +# out-of-range columns; the hidden keys get exactly zero attention, so outputs and K/V/Q gradients are unchanged. +# Only calls that will be differentiated are padded: the forward kernel is not affected. +_CUDNN_KV_TILE = 128 + + +def _pad_seq(x: torch.Tensor, pad: int) -> torch.Tensor: + """Append ``pad`` zero rows to the sequence axis (-2) of ``x``, keeping its physical axis order.""" + # q/k/v are often BSHD buffers viewed as BHSD; keep that layout (cf. ``_zero_masked_rows``). Broadcast + # (stride-0) axes go outermost so they do not become the innermost axis of the new buffer. + order = sorted(range(x.ndim), key=lambda axis: (x.stride(axis) == 0, x.stride(axis)), reverse=True) + inverse = [order.index(axis) for axis in range(x.ndim)] + physical = x.permute(order) + axis = order.index(x.ndim - 2) + shape = list(physical.shape) + shape[axis] += pad + # Zero-filled buffer + copy_ (vectorized for bf16, unlike torch.cat in PyTorch 2.7); autograd routes the + # gradient of the copied slice back to ``x``. + padded = physical.new_zeros(shape) + padded.narrow(axis, 0, physical.shape[axis]).copy_(physical) + return padded.permute(inverse) + + +def _pad_kv_for_cudnn(k: torch.Tensor, v: torch.Tensor, attn_mask): + """Pad K/V to a multiple of the cuDNN key tile and hide the padded keys in the (bool or additive) mask.""" + s_kv = k.shape[-2] + pad = -s_kv % _CUDNN_KV_TILE + if pad == 0: + return k, v, attn_mask + k, v = _pad_seq(k, pad), _pad_seq(v, pad) + if attn_mask is None: # key-padding mask that broadcasts against any query shape of the same rank + attn_mask = torch.ones(*([1] * (k.ndim - 1)), s_kv + pad, dtype=torch.bool, device=k.device) + attn_mask[..., s_kv:] = False + return k, v, attn_mask + if attn_mask.shape[-1] != s_kv: # broadcast over keys: materialize the key axis before extending it + attn_mask = attn_mask.expand(*attn_mask.shape[:-1], s_kv) + hidden = attn_mask.new_full((*attn_mask.shape[:-1], pad), False if attn_mask.dtype == torch.bool else -torch.inf) + return k, v, torch.cat([attn_mask, hidden], dim=-1) + + def sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, attn_mask=None) -> torch.Tensor: """``F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask)``, cuDNN first when enabled.""" if not _PREFER_CUDNN or not q.is_cuda or torch.are_deterministic_algorithms_enabled(): @@ -80,6 +130,8 @@ def sdpa(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, attn_mask=None) -> t nonempty = None if attn_mask is not None: attn_mask, nonempty = _safe_cudnn_mask(attn_mask) + if torch.is_grad_enabled() and (q.requires_grad or k.requires_grad or v.requires_grad): + k, v, attn_mask = _pad_kv_for_cudnn(k, v, attn_mask) # backward-only issue: inference stays unpadded with sdpa_kernel(_order(), set_priority=True): out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask) return out if nonempty is None else _zero_masked_rows(out, nonempty) diff --git a/tests/test_sdpa_kv_padding.py b/tests/test_sdpa_kv_padding.py new file mode 100644 index 00000000..679562ed --- /dev/null +++ b/tests/test_sdpa_kv_padding.py @@ -0,0 +1,191 @@ +"""K/V padding that keeps the cuDNN SDPA backward finite when s_kv is not a multiple of the key tile.""" + +import pytest +import torch +import torch.nn.functional as F +from torch.nn.attention import SDPBackend, sdpa_kernel + +from openwam.model.kernels import attention + + +def _mask(kind, sq, skv, *, device="cpu", dtype=torch.float64): + if kind is None: + return None + keep = torch.ones(2, 1, sq, skv, dtype=torch.bool, device=device) + keep[:, :, : sq // 3, skv // 2 :] = False + if kind == "bool": + return keep + bias = torch.zeros(keep.shape, dtype=dtype, device=device).masked_fill(~keep, -torch.inf) + bias[0, :, :3, :5] = -0.25 # finite biases must survive the padding unchanged + return bias + + +def _grads(fn, qkv, mask, upstream): + q, k, v = [x.detach().clone().requires_grad_() for x in qkv] + out = fn(q, k, v, mask) + return out, torch.autograd.grad(out, (q, k, v), upstream) + + +def _math(q, k, v, mask): + with sdpa_kernel(SDPBackend.MATH): + return F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + + +@pytest.mark.parametrize("transposed", [False, True]) +def test_pad_seq_appends_zero_keys_and_keeps_the_physical_layout(transposed): + x = torch.randn(2, 3, 5, 4) + if transposed: # BSHD buffer viewed as BHSD + x = x.transpose(1, 2).contiguous().transpose(1, 2) + padded = attention._pad_seq(x, 3) + assert padded.shape == (2, 3, 8, 4) + torch.testing.assert_close(padded[:, :, :5], x, rtol=0, atol=0) + assert torch.count_nonzero(padded[:, :, 5:]) == 0 + order = sorted(range(4), key=x.stride().__getitem__, reverse=True) + assert sorted(range(4), key=padded.stride().__getitem__, reverse=True) == order + + +@pytest.mark.parametrize("skv", [128, 256]) +def test_tile_multiples_are_left_untouched(skv): + k = torch.randn(1, 2, skv, 8) + v = torch.randn(1, 2, skv, 8) + mask = torch.ones(1, 1, 3, skv, dtype=torch.bool) + pk, pv, pmask = attention._pad_kv_for_cudnn(k, v, mask) + assert pk is k and pv is v and pmask is mask + + +@pytest.mark.parametrize("kind", ["bool", "additive", None]) +@pytest.mark.parametrize("skv", [77, 392, 64]) +def test_padding_preserves_outputs_and_gradients_cpu(kind, skv): + torch.manual_seed(3) + sq = 50 + qkv = [torch.randn(2, 4, s, 16, dtype=torch.float64) for s in (sq, skv, skv)] + mask = _mask(kind, sq, skv) + upstream = torch.randn(2, 4, sq, 16, dtype=torch.float64) + + def padded(q, k, v, m): + k, v, m = attention._pad_kv_for_cudnn(k, v, m) + assert k.shape[-2] % attention._CUDNN_KV_TILE == 0 + return _math(q, k, v, m) + + ref, ref_grads = _grads(_math, qkv, mask, upstream) + out, grads = _grads(padded, qkv, mask, upstream) + torch.testing.assert_close(out, ref, rtol=1e-12, atol=1e-12) + for got, want in zip(grads, ref_grads): + torch.testing.assert_close(got, want, rtol=1e-12, atol=1e-12) + + +def test_three_dimensional_inputs_without_mask_keep_their_rank(): + torch.manual_seed(5) + q, k, v = (torch.randn(3, s, 8, dtype=torch.float64) for s in (6, 77, 77)) + pk, pv, pmask = attention._pad_kv_for_cudnn(k, v, None) + assert pmask.shape == (1, 1, 128) # same rank as the (B, L, S) logits of a 3-D call + assert pmask[..., :77].all() and not pmask[..., 77:].any() + out = _math(q, pk, pv, pmask) + assert out.shape == (3, 6, 8) + torch.testing.assert_close(out, _math(q, k, v, None), rtol=1e-12, atol=1e-12) + + +def test_key_broadcast_mask_is_materialized_before_padding(): + k, v = torch.randn(1, 2, 5, 4), torch.randn(1, 2, 5, 4) + mask = torch.tensor([True, False])[:, None, None, None].expand(2, 1, 3, 1) + _, _, pmask = attention._pad_kv_for_cudnn(k, v, mask) + assert pmask.shape == (2, 1, 3, 128) + assert pmask[0, ..., :5].all() and not pmask[0, ..., 5:].any() + assert not pmask[1].any() + + +def _low_lse_inputs(sq, skv, *, device="cuda", dtype=torch.bfloat16, bshd=False, head_dim=128, row=5): + """Keys sharing a large component and one query row anti-aligned with it: that row's logits are all about + -136, so its logsumexp is far below -88.7 - the condition that turns the out-of-range key columns of the last + cuDNN tile into exp(0 - lse) = inf in the backward pass.""" + torch.manual_seed(0) + common = torch.ones(head_dim, device=device, dtype=dtype) + q = torch.randn(2, 4, sq, head_dim, device=device, dtype=dtype) + q[:, :, row] = -(136.0 / head_dim**0.5) * common + k = torch.randn(2, 4, skv, head_dim, device=device, dtype=dtype) * 0.5 + common + v = torch.randn(2, 4, skv, head_dim, device=device, dtype=dtype) + if bshd: # the callers pass BSHD buffers viewed as BHSD + q, k, v = (x.transpose(1, 2).contiguous().transpose(1, 2) for x in (q, k, v)) + return q, k, v + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA/cuDNN") +@pytest.mark.parametrize("bshd", [False, True]) +@pytest.mark.parametrize("kind", ["bool", "additive", None]) +@pytest.mark.parametrize("sq,skv", [(392, 392), (360, 64), (100, 77), (256, 392), (128, 64)]) +def test_cudnn_low_logsumexp_rows_have_finite_gradients(monkeypatch, kind, sq, skv, bshd): + monkeypatch.setattr(attention, "_PREFER_CUDNN", True) + # Force the affected backend; a silent fallback to another kernel cannot pass this test. + monkeypatch.setattr(attention, "_order", lambda: [SDPBackend.CUDNN_ATTENTION]) + qkv = _low_lse_inputs(sq, skv, bshd=bshd) + mask = _mask(kind, sq, skv, device="cuda", dtype=torch.bfloat16) + upstream = torch.randn(2, 4, sq, 128, device="cuda", dtype=torch.bfloat16) + + def wrapper(q, k, v, m): + return attention.sdpa(q, k, v, attn_mask=m) + + def reference(q, k, v, m): + m32 = m if m is None or m.dtype == torch.bool else m.float() + return _math(q.float(), k.float(), v.float(), m32).to(q.dtype) + + out, grads = _grads(wrapper, qkv, mask, upstream) + ref, ref_grads = _grads(reference, qkv, mask, upstream) + assert torch.isfinite(out).all() + assert all(torch.isfinite(g).all() for g in grads) + torch.testing.assert_close(out, ref, rtol=0.02, atol=0.02) + for got, want in zip(grads, ref_grads): + torch.testing.assert_close(got.float(), want.float(), rtol=0.05, atol=0.05 * want.float().abs().max().item()) + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA/cuDNN") +@pytest.mark.parametrize("head_dim", [64, 128]) +@pytest.mark.parametrize("sq,skv", [(392, 392), (100, 77)]) +def test_cudnn_low_logsumexp_row_in_last_query_tile(monkeypatch, head_dim, sq, skv): + """Same regression with the low row in the last, partial query tile and 64-wide heads (Cosmos bridges).""" + monkeypatch.setattr(attention, "_PREFER_CUDNN", True) + monkeypatch.setattr(attention, "_order", lambda: [SDPBackend.CUDNN_ATTENTION]) + qkv = _low_lse_inputs(sq, skv, head_dim=head_dim, row=sq - 1) + mask = _mask("bool", sq, skv, device="cuda") + upstream = torch.randn(2, 4, sq, head_dim, device="cuda", dtype=torch.bfloat16) + + def reference(q, k, v, m): + return _math(q.float(), k.float(), v.float(), m).to(q.dtype) + + out, grads = _grads(lambda q, k, v, m: attention.sdpa(q, k, v, attn_mask=m), qkv, mask, upstream) + ref, ref_grads = _grads(reference, qkv, mask, upstream) + assert all(torch.isfinite(g).all() for g in grads) + torch.testing.assert_close(out, ref, rtol=0.02, atol=0.02) + for got, want in zip(grads, ref_grads): + torch.testing.assert_close(got.float(), want.float(), rtol=0.05, atol=0.05 * want.float().abs().max().item()) + + +@pytest.mark.gpu +@pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA/cuDNN") +@pytest.mark.parametrize( + "grad_enabled,requires_grad,padded", [(False, True, False), (True, False, False), (True, True, True)] +) +def test_only_differentiated_calls_are_padded(monkeypatch, grad_enabled, requires_grad, padded): + """The issue is backward-only: calls autograd will not differentiate keep the original cuDNN forward bit for bit.""" + monkeypatch.setattr(attention, "_PREFER_CUDNN", True) + monkeypatch.setattr(attention, "_order", lambda: [SDPBackend.CUDNN_ATTENTION]) + calls = [] + real = attention._pad_kv_for_cudnn + + def spy(*args): + calls.append(1) + return real(*args) + + monkeypatch.setattr(attention, "_pad_kv_for_cudnn", spy) + q, k, v = (x.requires_grad_(requires_grad) for x in _low_lse_inputs(392, 392)) + mask = _mask("bool", 392, 392, device="cuda") + with torch.set_grad_enabled(grad_enabled): + out = attention.sdpa(q, k, v, attn_mask=mask) + assert bool(calls) == padded + with torch.no_grad(), sdpa_kernel(SDPBackend.CUDNN_ATTENTION): + reference = F.scaled_dot_product_attention(q, k, v, attn_mask=mask) + if padded: + torch.testing.assert_close(out, reference, rtol=0.02, atol=0.02) + else: + torch.testing.assert_close(out, reference, rtol=0, atol=0)