[Bugfix][Attention] Avoid NaN in the Triton softcap for large attention logits (#56579)
Signed-off-by: Kushal Dabbe <72650064+kushaldabbe@users.noreply.github.com> Co-authored-by: opencode <noreply@opencode.ai> Co-authored-by: Misha Goin <mgoin64@gmail.com>
Kushal committed
Sep 23, 2026 at 01:11 UTC
066ee911975133dc148433ed9e65061b4587c969
2 files changed
+47
-5
tests/kernels/attention/test_triton_unified_attention.py
+38
@@ -10,6 +10,7 @@ from vllm.triton_utils import tl, triton
10
from vllm.utils.math_utils import next_power_of_2
11
from vllm.utils.torch_utils import set_random_seed
12
from vllm.v1.attention.ops.triton_attention_helpers import (
13
+ apply_softcap,
14
compute_tile_loop_bounds,
15
)
16
from vllm.v1.attention.ops.triton_unified_attention import unified_attention
@@ -899,3 +900,40 @@ def test_triton_unified_attn_use_td_tile_clamp(
900
soft_cap=None,
901
seq_threshold_3D=0,
902
)
903
+
904
+
905
+@triton.jit
906
+def _softcap_probe_kernel(
907
+ scores_ptr,
908
+ out_ptr,
909
+ numel,
910
+ soft_cap,
911
+ BLOCK: tl.constexpr,
912
+):
913
+ offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)
914
+ mask = offs < numel
915
+ scores = tl.load(scores_ptr + offs, mask=mask, other=0.0)
916
+ tl.store(out_ptr + offs, apply_softcap(scores, soft_cap), mask=mask)
917
+
918
+
919
+@torch.inference_mode()
920
+def test_softcap_does_not_overflow_on_large_scores() -> None:
921
+ """Scores above ~88 * soft_cap must not overflow to inf/NaN.
922
+
923
+ The exp-based softcap computes ``(exp(y) - exp(-y)) / (exp(y) + exp(-y))``
924
+ with ``y = S / soft_cap``. For ``|y| > ~88`` the exponentials overflow to
925
+ ``inf`` and the ratio becomes ``inf / inf = NaN``, poisoning the whole
926
+ attention row. Gemma-2 style models use ``attn_logit_softcapping = 50``,
927
+ so scores above 4400 are in range for their large attention logits.
928
+ """
929
+ soft_cap = 50.0
930
+ scores = torch.tensor(
931
+ [1.0, 100.0, 1000.0, 3000.0, 5000.0, 10000.0, -1.0, -10000.0, 4400.0],
932
+ device=DEVICE_TYPE,
933
+ dtype=torch.float32,
934
+ )
935
+ out = torch.empty_like(scores)
936
+ _softcap_probe_kernel[(1,)](scores, out, scores.numel(), soft_cap, BLOCK=16)
937
+ ref = soft_cap * torch.tanh(scores / soft_cap)
938
+ assert torch.isfinite(out).all(), out
939
+ torch.testing.assert_close(out, ref, atol=1e-3, rtol=1e-3)
vllm/v1/attention/ops/triton_attention_helpers.py
+9
-5
@@ -28,12 +28,16 @@ def cdiv_fn(x, y):
28
def apply_softcap(S, x):
29
"""Softcap (aka tanh-style clamp) used to bound attention scores.
30
31
- ``x * tanh(S / x)`` rewritten to avoid a direct ``tanh`` call.
31
+ ``x * tanh(S / x)`` computed without a direct ``tanh`` call.
32
+
33
+ Use ``tanh(y) = 1 - 2 / (exp(2y) + 1)`` rather than
34
+ ``(exp(y) - exp(-y)) / (exp(y) + exp(-y))``. The symmetric-exponential
35
+ form overflows to ``inf`` for ``|y| > ~88`` and then evaluates to
36
+ ``inf / inf = NaN``, poisoning the attention row. The single-exponential
37
+ form saturates cleanly at +/-1 for either sign.
38
"""
33
- Sdiv = S / x
34
- p1 = tl.exp(Sdiv)
35
- p2 = tl.exp(-Sdiv)
36
- return x * (p1 - p2) / (p1 + p2)
39
+ y = S / x
40
+ return x * (1.0 - 2.0 / (tl.exp(2.0 * y) + 1.0))
41
42
43
# ===========================================================================