diff --git a/sdm/_kernels/log_scale.py b/sdm/_kernels/log_scale.py new file mode 100644 index 0000000000..8a942937ef --- /dev/null +++ b/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 diff --git a/sdm/_kernels/triton/log_scale.py b/sdm/_kernels/triton/log_scale.py new file mode 100644 index 0000000000..83f33b0b1e --- /dev/null +++ b/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 diff --git a/sdm/nn/scaling.py b/sdm/nn/scaling.py index 4084b5a74f..211a86e62c 100644 --- a/sdm/nn/scaling.py +++ b/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)) diff --git a/test/nn/test_scaling.py b/test/nn/test_scaling.py index eeb84fb76c..1432fcc402 100644 --- a/test/nn/test_scaling.py +++ b/test/nn/test_scaling.py @@ -1,6 +1,7 @@ # 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 @@ -8,7 +9,7 @@ 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)