Fuse shared-length query scaling during inference

sdm/_kernels/log_scale.py

@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from collections.abc import Callable
+from typing import TypeAlias
+
+import torch
+from torch import Tensor
+
+_LogScale: TypeAlias = Callable[[Tensor, int, Tensor, Tensor | None], Tensor]
+
+_triton_log_scale: _LogScale | None = None
+if torch.version.cuda is not None:
+    try:
+        from sdm._kernels.triton.log_scale import (
+            log_scale as _triton_log_scale_impl,
+        )
+    except ImportError:
+        pass
+    else:
+        _triton_log_scale = _triton_log_scale_impl
+
+
+def log_scale(
+    query: Tensor,
+    key_len: Tensor | int,
+    head_scale: Tensor,
+    gate: Tensor | None = None,
+) -> Tensor:
+    """Scale query heads, preserving each eager operation's rounding."""
+    inputs = (query, head_scale)
+    if gate is not None:
+        inputs += (gate,)
+    if (
+        _triton_log_scale is not None
+        and query.is_cuda
+        and query.dtype in {torch.float16, torch.bfloat16, torch.float32}
+        and query.numel() > 0
+        and isinstance(key_len, int)
+        and -(2**63) <= key_len < 2**63
+        and head_scale.dtype in {torch.float16, torch.bfloat16, torch.float32}
+        and all(x.device == query.device and x.is_contiguous() for x in inputs)
+        and (
+            gate is None
+            or (gate.shape == query.shape and gate.dtype == query.dtype)
+        )
+        and not (
+            torch.is_grad_enabled() and any(x.requires_grad for x in inputs)
+        )
+    ):
+        return _triton_log_scale(query, key_len, head_scale, gate)
+
+    length = (
+        key_len.float()
+        if isinstance(key_len, Tensor)
+        else query.new_full((), key_len, dtype=torch.float32)
+    )
+    log_key_len = length.clamp(min=1.0).log().to(query.dtype)
+    head_scale = head_scale.to(query.dtype).reshape(
+        *((1,) * (query.dim() - 2)), -1, 1
+    )
+    out = query * log_key_len[..., None, None] * head_scale
+    if gate is not None:
+        out = out * (1 + gate.tanh())
+    return out

sdm/_kernels/triton/log_scale.py

@@ -0,0 +1,65 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+from typing import Any, cast
+
+import torch
+import triton
+import triton.language as tl
+from torch import Tensor
+from triton.language.extra.cuda import libdevice
+
+
+@triton.jit(do_not_specialize=["numel", "key_len"])
+def _log_scale_kernel(
+    query_ptr,
+    key_len,
+    head_ptr,
+    gate_ptr,
+    out_ptr,
+    numel,
+    heads: tl.constexpr,
+    channels: tl.constexpr,
+    block: tl.constexpr,
+) -> None:
+    index = tl.program_id(0).to(tl.int64) * block + tl.arange(0, block)
+    mask = index < numel
+    query = tl.load(query_ptr + index, mask=mask, other=0)
+    dtype = query.dtype
+    length = libdevice.log(tl.maximum(key_len.to(tl.float32), 1.0))
+    length = length.to(dtype).to(tl.float32)
+    head = tl.load(head_ptr + index // channels % heads)
+    head = head.to(dtype).to(tl.float32)
+    # Eager execution rounds each product and gate operation separately.
+    out = (query.to(tl.float32) * length).to(dtype).to(tl.float32)
+    out = (out * head).to(dtype).to(tl.float32)
+    if gate_ptr is not None:
+        gate = tl.load(gate_ptr + index, mask=mask, other=0).to(tl.float32)
+        gate = libdevice.tanh(gate).to(dtype).to(tl.float32)
+        gate = (1.0 + gate).to(dtype).to(tl.float32)
+        out = out * gate
+    tl.store(out_ptr + index, out, mask=mask)
+
+
+def log_scale(
+    query: Tensor,
+    key_len: int,
+    head_scale: Tensor,
+    gate: Tensor | None,
+) -> Tensor:
+    out = torch.empty_like(query)
+    grid = ((query.numel() + 1023) // 1024,)
+    with torch.cuda.device(query.device):
+        cast(Any, _log_scale_kernel)[grid](
+            query,
+            key_len,
+            head_scale,
+            gate,
+            out,
+            query.numel(),
+            query.size(-2),
+            query.size(-1),
+            1024,
+            enable_fp_fusion=False,
+        )
+    return out

sdm/nn/scaling.py

@@ -12,6 +12,8 @@
 from torch import Tensor
 from torch.nn import GELU, Linear, Parameter, Sequential
 
+from sdm._kernels.log_scale import log_scale
+
 
 class QueryScaling(torch.nn.Module, abc.ABC):
     r"""Base class for query scaling modules in :class:`SDPA`.
@@ -169,21 +171,22 @@ def forward(
         key_len: Tensor | int,  # [..., 1] or [..., S] or scalar
     ) -> Tensor:  # [..., S, H, C]
         r""":meta private:"""  # noqa: D415
+        return self._scale(query, key_len=key_len)
+
+    def _scale(
+        self,
+        query: Tensor,
+        *,
+        key_len: Tensor | int,
+        gate: Tensor | None = None,
+    ) -> Tensor:
         # Avoid a silently wrong result from broadcasting the head dimension.
         if query.size(-2) != self.head_scale.numel():
             raise ValueError(
                 "`query.size(-2)` must equal the number of per-head scales "
                 f"(got {query.size(-2)} and {self.head_scale.numel()})"
             )
-        if isinstance(key_len, Tensor):
-            key_len_tensor = key_len.float()
-        else:
-            key_len_tensor = query.new_full((), key_len, dtype=torch.float32)
-        log_key_len = key_len_tensor.clamp(min=1.0).log().to(query.dtype)
-        head_scale = self.head_scale.to(query.dtype).reshape(
-            *((1,) * (query.dim() - 2)), -1, 1
-        )
-        return query * log_key_len[..., None, None] * head_scale
+        return log_scale(query, key_len, self.head_scale, gate)
 
 
 class GatedLogScale(LogScale):
@@ -236,5 +239,4 @@ def forward(
         key_len: Tensor | int,  # [..., 1] or [..., S] or scalar
     ) -> Tensor:  # [..., S, H, C]
         r""":meta private:"""  # noqa: D415
-        query_scaled = super().forward(query=query, key_len=key_len)
-        return query_scaled * (1 + self.gate(query).tanh())
+        return self._scale(query, key_len=key_len, gate=self.gate(query))

test/nn/test_scaling.py

@@ -1,14 +1,15 @@
 # SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
+import importlib
 from collections.abc import Callable
 
 import pytest
 import torch
 from torch import Tensor
 
 from sdm.nn import GatedLogScale, LogScale, QASSMax
-from sdm.testing import withCUDA
+from sdm.testing import onlyCUDA, withCUDA
 
 
 @withCUDA
@@ -94,3 +95,162 @@ def test_gated_log_scale_forward(device: torch.device) -> None:
     assert out.shape == query.shape
     assert out.dtype == query.dtype
     assert out.device == query.device
+
+
+def _log_scale_reference(
+    module: LogScale,
+    query: Tensor,
+    key_len: Tensor | int,
+) -> Tensor:
+    length = torch.as_tensor(key_len, device=query.device).float()
+    length = length.clamp(min=1).log().to(query.dtype)
+    heads = module.head_scale.to(query.dtype)
+    heads = heads.reshape(*((1,) * (query.dim() - 2)), -1, 1)
+    out = query * length[..., None, None] * heads
+    if isinstance(module, GatedLogScale):
+        out = out * (1 + module.gate(query).tanh())
+    return out
+
+
+@withCUDA
+@pytest.mark.parametrize("gated", [False, True])
+@pytest.mark.parametrize(
+    "dtype", [torch.float16, torch.bfloat16, torch.float32]
+)
+@pytest.mark.parametrize("layout", ["contiguous", "strided", "empty"])
+@pytest.mark.parametrize("length_kind", ["scalar", "batch", "query"])
+def test_log_scale_exact_inference(
+    device: torch.device,
+    gated: bool,
+    dtype: torch.dtype,
+    layout: str,
+    length_kind: str,
+) -> None:
+    module = (
+        GatedLogScale(
+            channels=7, num_heads=3, hidden_channels=11, device=device
+        )
+        if gated
+        else LogScale(num_heads=3, device=device)
+    )
+    with torch.no_grad():
+        for parameter in module.parameters():
+            parameter.uniform_(-0.5, 0.5)
+    rows = 0 if layout == "empty" else 17
+    query = torch.randn(2, rows, 3, 14, dtype=dtype, device=device)
+    query = query[..., ::2]
+    if layout != "strided":
+        query = query.contiguous()
+    key_len: Tensor | int = 311
+    if length_kind == "batch":
+        key_len = torch.tensor([[0], [311]], device=device)
+    elif length_kind == "query":
+        key_len = torch.arange(2 * rows, device=device).reshape(2, rows)
+    with (
+        torch.inference_mode(),
+        torch.autocast(
+            device.type, dtype=dtype, enabled=dtype != torch.float32
+        ),
+    ):
+        expected = _log_scale_reference(module, query, key_len)
+        actual = module(query, key_len=key_len)
+    torch.testing.assert_close(
+        actual, expected, rtol=0, atol=0, equal_nan=True
+    )
+
+
+@withCUDA
+@pytest.mark.parametrize("gated", [False, True])
+def test_log_scale_gradients(device: torch.device, gated: bool) -> None:
+    module = (
+        GatedLogScale(
+            channels=7, num_heads=3, hidden_channels=11, device=device
+        )
+        if gated
+        else LogScale(num_heads=3, device=device)
+    )
+    with torch.no_grad():
+        for parameter in module.parameters():
+            parameter.uniform_(-0.5, 0.5)
+    query = torch.randn(2, 17, 3, 7, device=device, requires_grad=True)
+    key_len = torch.tensor(11.0, device=device, requires_grad=True)
+    expected = _log_scale_reference(module, query, key_len)
+    actual = module(query, key_len=key_len)
+    inputs = (query, key_len, *module.parameters())
+    expected_grad = torch.autograd.grad(expected.sum(), inputs)
+    actual_grad = torch.autograd.grad(actual.sum(), inputs)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+    torch.testing.assert_close(actual_grad, expected_grad, rtol=0, atol=0)
+
+
+@onlyCUDA
+@pytest.mark.parametrize("gated", [False, True])
+def test_log_scale_without_triton(
+    monkeypatch: pytest.MonkeyPatch, gated: bool
+) -> None:
+    kernel = importlib.import_module("sdm._kernels.log_scale")
+    monkeypatch.setattr(kernel, "_triton_log_scale", None)
+    module = (
+        GatedLogScale(7, 3, 11, device="cuda")
+        if gated
+        else LogScale(3, device="cuda")
+    )
+    query = torch.randn(17, 3, 7, device="cuda", dtype=torch.float16)
+    with torch.inference_mode(), torch.autocast("cuda", dtype=query.dtype):
+        expected = _log_scale_reference(module, query, 31)
+        actual = module(query, key_len=31)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@onlyCUDA
+@pytest.mark.parametrize("dtype", [torch.float32, torch.float64])
+def test_gated_log_scale_dtype_fallback(dtype: torch.dtype) -> None:
+    module = GatedLogScale(7, 3, 11, device="cuda", dtype=dtype)
+    query = torch.randn(17, 3, 7, device="cuda", dtype=dtype)
+    with (
+        torch.inference_mode(),
+        torch.autocast(
+            "cuda", dtype=torch.float16, enabled=dtype != torch.float64
+        ),
+    ):
+        expected = _log_scale_reference(module, query, 31)
+        actual = module(query, key_len=31)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@onlyCUDA
+@pytest.mark.parametrize(
+    "dtype", [torch.float16, torch.bfloat16, torch.float32]
+)
+@pytest.mark.parametrize(
+    "key_len", [-17, 0, 1, 311, 2**24 + 1, 2**31 + 1, 2**63 - 1]
+)
+def test_log_scale_integer_lengths(dtype: torch.dtype, key_len: int) -> None:
+    module = GatedLogScale(7, 3, 11, device="cuda")
+    with torch.no_grad():
+        for parameter in module.parameters():
+            parameter.uniform_(-0.5, 0.5)
+    query = torch.randn(17, 3, 7, device="cuda", dtype=dtype)
+    with (
+        torch.inference_mode(),
+        torch.autocast("cuda", dtype=dtype, enabled=dtype != torch.float32),
+    ):
+        expected = _log_scale_reference(module, query, key_len)
+        actual = module(query, key_len=key_len)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+@onlyCUDA
+@pytest.mark.parametrize("tensor_length", [False, True])
+def test_log_scale_compilation(compilation: None, tensor_length: bool) -> None:
+    module = GatedLogScale(7, 3, 11, device="cuda")
+    with torch.no_grad():
+        for parameter in module.parameters():
+            parameter.uniform_(-0.5, 0.5)
+    query = torch.randn(17, 3, 7, device="cuda", dtype=torch.float16)
+    key_len = torch.tensor(31, device="cuda") if tensor_length else 31
+    compiled = torch.compile(module, backend="eager", fullgraph=True)
+    with torch.inference_mode(), torch.autocast("cuda", dtype=query.dtype):
+        expected = _log_scale_reference(module, query, key_len)
+        actual = compiled(query, key_len=key_len)
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)