From 48423cb2d9cc4ecfc3a2ddf841203285723ee7f6 Mon Sep 17 00:00:00 2001 From: Lei Sun Date: Fri, 16 Oct 2026 17:11:21 +0000 Subject: [PATCH 1/6] Fuse inference query scaling while preserving eager rounding --- sdm/_kernels/log_scale.py | 53 ++++++++++++++++++++++ sdm/_kernels/triton/log_scale.py | 63 ++++++++++++++++++++++++++ sdm/nn/scaling.py | 22 ++++++--- test/nn/test_scaling.py | 76 ++++++++++++++++++++++++++++++++ 4 files changed, 209 insertions(+), 5 deletions(-) create mode 100644 sdm/_kernels/log_scale.py create mode 100644 sdm/_kernels/triton/log_scale.py diff --git a/sdm/_kernels/log_scale.py b/sdm/_kernels/log_scale.py new file mode 100644 index 0000000000..1d9afbd515 --- /dev/null +++ b/sdm/_kernels/log_scale.py @@ -0,0 +1,53 @@ +# 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, Tensor, Tensor, Tensor | None], Tensor] + +_triton_log_scale: _LogScale | None = 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, + log_key_len: Tensor, + head_scale: Tensor, + gate: Tensor | None = None, +) -> Tensor: + """Scale query heads, preserving each eager operation's rounding.""" + inputs = (query, log_key_len, 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 log_key_len.numel() == 1 + and log_key_len.dim() <= query.dim() - 2 + and all( + x.device == query.device + and x.dtype == query.dtype + and x.is_contiguous() + for x in inputs + ) + and (gate is None or gate.shape == query.shape) + and not (torch.is_grad_enabled() and any(x.requires_grad for x in inputs)) + ): + return _triton_log_scale(query, log_key_len, head_scale, gate) + + head_scale = head_scale.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..0db2a933cf --- /dev/null +++ b/sdm/_kernels/triton/log_scale.py @@ -0,0 +1,63 @@ +# 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"]) +def _log_scale_kernel( + query_ptr, + length_ptr, + 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 = tl.load(length_ptr).to(tl.float32) + head = tl.load(head_ptr + index // channels % heads).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, + log_key_len: Tensor, + head_scale: Tensor, + gate: Tensor | None, +) -> Tensor: + out = torch.empty_like(query) + grid = (triton.cdiv(query.numel(), 1024),) + with torch.cuda.device(query.device): + cast(Any, _log_scale_kernel)[grid]( + query, + log_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..28e79a7ea7 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,6 +171,15 @@ 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( @@ -180,10 +191,12 @@ def forward( 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 log_scale( + query, + log_key_len, + self.head_scale.to(query.dtype), + gate, ) - return query * log_key_len[..., None, None] * head_scale class GatedLogScale(LogScale): @@ -236,5 +249,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..f7cf59ed5b 100644 --- a/test/nn/test_scaling.py +++ b/test/nn/test_scaling.py @@ -94,3 +94,79 @@ 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) From d4b98b785d8d4dd83f03251fc653a810e116e079 Mon Sep 17 00:00:00 2001 From: Lei Sun Date: Fri, 16 Oct 2026 17:13:21 +0000 Subject: [PATCH 2/6] Apply style checks to fused query scaling --- sdm/_kernels/log_scale.py | 12 +++++++++--- sdm/_kernels/triton/log_scale.py | 2 +- test/nn/test_scaling.py | 20 +++++++++++++++----- 3 files changed, 25 insertions(+), 9 deletions(-) diff --git a/sdm/_kernels/log_scale.py b/sdm/_kernels/log_scale.py index 1d9afbd515..203b9ee794 100644 --- a/sdm/_kernels/log_scale.py +++ b/sdm/_kernels/log_scale.py @@ -7,11 +7,15 @@ import torch from torch import Tensor -_LogScale: TypeAlias = Callable[[Tensor, Tensor, Tensor, Tensor | None], Tensor] +_LogScale: TypeAlias = Callable[ + [Tensor, Tensor, Tensor, Tensor | None], Tensor +] _triton_log_scale: _LogScale | None = None try: - from sdm._kernels.triton.log_scale import log_scale as _triton_log_scale_impl + from sdm._kernels.triton.log_scale import ( + log_scale as _triton_log_scale_impl, + ) except ImportError: pass else: @@ -42,7 +46,9 @@ def log_scale( for x in inputs ) and (gate is None or gate.shape == query.shape) - and not (torch.is_grad_enabled() and any(x.requires_grad for x in inputs)) + and not ( + torch.is_grad_enabled() and any(x.requires_grad for x in inputs) + ) ): return _triton_log_scale(query, log_key_len, head_scale, gate) diff --git a/sdm/_kernels/triton/log_scale.py b/sdm/_kernels/triton/log_scale.py index 0db2a933cf..ed63fd07b9 100644 --- a/sdm/_kernels/triton/log_scale.py +++ b/sdm/_kernels/triton/log_scale.py @@ -46,7 +46,7 @@ def log_scale( gate: Tensor | None, ) -> Tensor: out = torch.empty_like(query) - grid = (triton.cdiv(query.numel(), 1024),) + grid = ((query.numel() + 1023) // 1024,) with torch.cuda.device(query.device): cast(Any, _log_scale_kernel)[grid]( query, diff --git a/test/nn/test_scaling.py b/test/nn/test_scaling.py index f7cf59ed5b..859b966d8e 100644 --- a/test/nn/test_scaling.py +++ b/test/nn/test_scaling.py @@ -113,7 +113,9 @@ def _log_scale_reference( @withCUDA @pytest.mark.parametrize("gated", [False, True]) -@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) +@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( @@ -124,7 +126,9 @@ def test_log_scale_exact_inference( length_kind: str, ) -> None: module = ( - GatedLogScale(channels=7, num_heads=3, hidden_channels=11, device=device) + GatedLogScale( + channels=7, num_heads=3, hidden_channels=11, device=device + ) if gated else LogScale(num_heads=3, device=device) ) @@ -143,18 +147,24 @@ def test_log_scale_exact_inference( 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), + 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) + 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) + GatedLogScale( + channels=7, num_heads=3, hidden_channels=11, device=device + ) if gated else LogScale(num_heads=3, device=device) ) From 50342fbb2684d536539377902b173eb0e73c9fd8 Mon Sep 17 00:00:00 2001 From: Lei Sun Date: Fri, 16 Oct 2026 17:18:07 +0000 Subject: [PATCH 3/6] Cover query scaling dtype and optional-kernel fallbacks --- test/nn/test_scaling.py | 38 +++++++++++++++++++++++++++++++++++++- 1 file changed, 37 insertions(+), 1 deletion(-) diff --git a/test/nn/test_scaling.py b/test/nn/test_scaling.py index 859b966d8e..7d920d4e33 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 @@ -180,3 +181,38 @@ def test_log_scale_gradients(device: torch.device, gated: bool) -> None: 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) From deb8e2ca647a39243a74495785b994da3ef1fc13 Mon Sep 17 00:00:00 2001 From: Lei Sun Date: Fri, 16 Oct 2026 17:42:23 +0000 Subject: [PATCH 4/6] Prepare shared query scales inside the fused inference kernel --- sdm/_kernels/log_scale.py | 35 ++++++++++++++++++-------------- sdm/_kernels/triton/log_scale.py | 14 +++++++------ sdm/nn/scaling.py | 12 +---------- test/nn/test_scaling.py | 22 ++++++++++++++++++++ 4 files changed, 51 insertions(+), 32 deletions(-) diff --git a/sdm/_kernels/log_scale.py b/sdm/_kernels/log_scale.py index 203b9ee794..8a0a8a21ab 100644 --- a/sdm/_kernels/log_scale.py +++ b/sdm/_kernels/log_scale.py @@ -7,9 +7,7 @@ import torch from torch import Tensor -_LogScale: TypeAlias = Callable[ - [Tensor, Tensor, Tensor, Tensor | None], Tensor -] +_LogScale: TypeAlias = Callable[[Tensor, int, Tensor, Tensor | None], Tensor] _triton_log_scale: _LogScale | None = None try: @@ -24,12 +22,12 @@ def log_scale( query: Tensor, - log_key_len: Tensor, + key_len: Tensor | int, head_scale: Tensor, gate: Tensor | None = None, ) -> Tensor: """Scale query heads, preserving each eager operation's rounding.""" - inputs = (query, log_key_len, head_scale) + inputs = (query, head_scale) if gate is not None: inputs += (gate,) if ( @@ -37,22 +35,29 @@ def log_scale( and query.is_cuda and query.dtype in {torch.float16, torch.bfloat16, torch.float32} and query.numel() > 0 - and log_key_len.numel() == 1 - and log_key_len.dim() <= query.dim() - 2 - and all( - x.device == query.device - and x.dtype == query.dtype - and x.is_contiguous() - for x in inputs + 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 (gate is None or gate.shape == query.shape) and not ( torch.is_grad_enabled() and any(x.requires_grad for x in inputs) ) ): - return _triton_log_scale(query, log_key_len, head_scale, gate) + return _triton_log_scale(query, key_len, head_scale, gate) - head_scale = head_scale.reshape(*((1,) * (query.dim() - 2)), -1, 1) + 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()) diff --git a/sdm/_kernels/triton/log_scale.py b/sdm/_kernels/triton/log_scale.py index ed63fd07b9..83f33b0b1e 100644 --- a/sdm/_kernels/triton/log_scale.py +++ b/sdm/_kernels/triton/log_scale.py @@ -10,10 +10,10 @@ from triton.language.extra.cuda import libdevice -@triton.jit(do_not_specialize=["numel"]) +@triton.jit(do_not_specialize=["numel", "key_len"]) def _log_scale_kernel( query_ptr, - length_ptr, + key_len, head_ptr, gate_ptr, out_ptr, @@ -26,8 +26,10 @@ def _log_scale_kernel( mask = index < numel query = tl.load(query_ptr + index, mask=mask, other=0) dtype = query.dtype - length = tl.load(length_ptr).to(tl.float32) - head = tl.load(head_ptr + index // channels % heads).to(tl.float32) + 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) @@ -41,7 +43,7 @@ def _log_scale_kernel( def log_scale( query: Tensor, - log_key_len: Tensor, + key_len: int, head_scale: Tensor, gate: Tensor | None, ) -> Tensor: @@ -50,7 +52,7 @@ def log_scale( with torch.cuda.device(query.device): cast(Any, _log_scale_kernel)[grid]( query, - log_key_len, + key_len, head_scale, gate, out, diff --git a/sdm/nn/scaling.py b/sdm/nn/scaling.py index 28e79a7ea7..211a86e62c 100644 --- a/sdm/nn/scaling.py +++ b/sdm/nn/scaling.py @@ -186,17 +186,7 @@ def _scale( "`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) - return log_scale( - query, - log_key_len, - self.head_scale.to(query.dtype), - gate, - ) + return log_scale(query, key_len, self.head_scale, gate) class GatedLogScale(LogScale): diff --git a/test/nn/test_scaling.py b/test/nn/test_scaling.py index 7d920d4e33..3b6bef85f1 100644 --- a/test/nn/test_scaling.py +++ b/test/nn/test_scaling.py @@ -216,3 +216,25 @@ def test_gated_log_scale_dtype_fallback(dtype: torch.dtype) -> None: 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) From 834cfc5ad0a0d0eb7ede99c24abaf676de4c5e8a Mon Sep 17 00:00:00 2001 From: Lei Sun Date: Fri, 16 Oct 2026 17:48:32 +0000 Subject: [PATCH 5/6] Check graph capture for fused and fallback query scaling --- test/nn/test_scaling.py | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/test/nn/test_scaling.py b/test/nn/test_scaling.py index 3b6bef85f1..1432fcc402 100644 --- a/test/nn/test_scaling.py +++ b/test/nn/test_scaling.py @@ -238,3 +238,19 @@ def test_log_scale_integer_lengths(dtype: torch.dtype, key_len: int) -> None: 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) From ba793705be3cbeb25a6600f4e4bacf4ea38d2c96 Mon Sep 17 00:00:00 2001 From: Lei Sun Date: Fri, 16 Oct 2026 18:19:02 +0000 Subject: [PATCH 6/6] Keep native scaling on non-CUDA backends --- sdm/_kernels/log_scale.py | 17 +++++++++-------- 1 file changed, 9 insertions(+), 8 deletions(-) diff --git a/sdm/_kernels/log_scale.py b/sdm/_kernels/log_scale.py index 8a0a8a21ab..8a942937ef 100644 --- a/sdm/_kernels/log_scale.py +++ b/sdm/_kernels/log_scale.py @@ -10,14 +10,15 @@ _LogScale: TypeAlias = Callable[[Tensor, int, Tensor, Tensor | None], Tensor] _triton_log_scale: _LogScale | None = 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 +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(