Skip to content

[BUG][mthreads]TTGIR miscompile scf.if causes double math.exp → NaN output #973

Open
Open
[BUG][mthreads]TTGIR miscompile scf.if causes double math.exp → NaN output#973

Description

@O1iveira7
O1iveira7
opened on Aug 12, 2026

Environment

  • Triton 3.6.0 (FlagTree mthreads3.6), MUSA backend, MTT S5000
  • torch_musa 2.7.1, MUSA 4.3.5, Python 3.10

Symptom

Semantically-identical transform (wrap a runtime-conditional if/else with three
identical branches around an assignment inside a loop) makes the kernel output
NaN where the original kernel outputs correct values.

The same kernel pair was also tested on NVIDIA (CUDA): original and transformed
kernels produce identical, correct outputs (no NaN, equality check PASS).

FlagTree mthreads3.6 output:

original: tensor([[2.6185, 1.7739, 3.2257, 2.3810],
[2.6185, 1.7739, 3.2257, 2.3810]])
variant: tensor([[nan, nan, nan, nan],
[nan, nan, nan, nan]])
FAIL

Triton3.7 RTX3070 output:

original: tensor([[2.6185, 1.7739, 3.2257, 2.3810],
[2.6185, 1.7739, 3.2257, 2.3810]])
variant: tensor([[2.6185, 1.7739, 3.2257, 2.3810],
[2.6185, 1.7739, 3.2257, 2.3810]])
PASS

Minimal reproducer

Difference between the two kernels

The two kernels differ in exactly one place — the computation of
attention_logits_shift inside the loop:

Original kernel — straight-line assignment:

  attention_logits_shift = attention_logits - max_value[:, None]

Variant kernel — the same assignment wrapped in nested runtime conditionals,
plus a new pointer argument _fz_switch_ptr (a 2-element int32 buffer):

  opaque_zero = tl.load(_fz_switch_ptr)         # runtime scalar (host value: 0)
  opaque_one  = tl.load(_fz_switch_ptr + 1)     # runtime scalar (host value: 1)
  opaque_true  = opaque_zero < opaque_one       # True at runtime
  opaque_false = opaque_zero > opaque_one       # False at runtime

  if opaque_true:
      if opaque_false:
          attention_logits_shift = attention_logits - max_value[:, None]
      else:
          attention_logits_shift = attention_logits - max_value[:, None]
  else:
      attention_logits_shift = attention_logits - max_value[:, None]

Why the semantics are identical

  1. All three branches assign the same expression
    (attention_logits - max_value[:, None]), so whichever branch executes,
    the value is identical.
  2. The switch buffer is only read (no side effects), and its values are
    used solely to select between identical branches.
  3. Therefore for any input (Q, K, V) both kernels produce bit-identical
    outputs — f(Q, K, V) is unchanged by the transform.

The conditions must be runtime (loaded from device memory) rather than
compile-time constants: a Python-level constant would be folded away by the
frontend and no scf.if would be emitted. The device-loaded scalars force
the compiler to emit real nested scf.if regions — and that is the only
difference the compiler sees. (The interpreter mode of Triton confirms the
equivalence: it executes the source step-by-step and both kernels give
bit-identical results; the NaN appears only on the compiled path.)

from __future__ import annotations

import sys

import torch
import triton
import triton.language as tl

M, N, d = 2, 3, 4
Q = torch.tensor([[1.0, 2.0, 3.0, 4.0], [0.5, 1.5, 2.5, 3.5]], dtype=torch.float32)
K = torch.tensor(
    [
        [1.0, 0.5, 0.0, -0.5],
        [-0.5, 1.0, 0.5, 0.0],
        [0.0, -0.5, 1.0, 0.5],
    ],
    dtype=torch.float32,
)
V = torch.tensor(
    [[1.0, 2.0, 3.0, 4.0], [4.0, 3.0, 2.0, 1.0], [2.0, 1.0, 4.0, 3.0]],
    dtype=torch.float32,
)


@triton.jit
def softmax_attention(
    Q_ptr,
    K_ptr,
    V_ptr,
    OUT_ptr,
    M,
    N,
    d,
    Q_stride_M,
    Q_stride_d,
    K_stride_N,
    K_stride_d,
    V_stride_N,
    V_stride_d,
    OUT_stride_M,
    OUT_stride_d,
    BLOCKSIZE_M: tl.constexpr,
    BLOCKSIZE_N: tl.constexpr,
    BLOCKSIZE_d: tl.constexpr,
):
    pid0 = tl.program_id(0)
    pid1 = tl.program_id(1)

    offset_M = pid0 * BLOCKSIZE_M + tl.arange(0, BLOCKSIZE_M)
    offset_d = pid1 * BLOCKSIZE_d + tl.arange(0, BLOCKSIZE_d)
    offset_N = tl.arange(0, BLOCKSIZE_N)

    Q_offset = offset_M[:, None] * Q_stride_M + offset_d[None, :] * Q_stride_d
    Q_mask = (offset_M[:, None] < M) & (offset_d[None, :] < d)
    Q_data = tl.load(Q_ptr + Q_offset, mask=Q_mask)

    accumulator = tl.zeros((BLOCKSIZE_M, BLOCKSIZE_d), dtype=tl.float32)
    softmax_running_sum = tl.zeros([BLOCKSIZE_M], dtype=tl.float32)
    softmax_current_max = tl.full([BLOCKSIZE_M], float("-inf"), dtype=tl.float32)
    attention_logits_scale = tl.sqrt(d + 0.0)

    for current_index in range(0, N, BLOCKSIZE_N):
        current_k_offset = current_index + offset_N
        current_v_offset = current_k_offset

        K_offset = (
            current_k_offset[:, None] * K_stride_N + offset_d[None, :] * K_stride_d
        )
        V_offset = (
            current_v_offset[:, None] * V_stride_N + offset_d[None, :] * V_stride_d
        )

        K_mask = (current_k_offset[:, None] < N) & (offset_d[None, :] < d)
        V_mask = (current_v_offset[:, None] < N) & (offset_d[None, :] < d)

        K_data = tl.load(K_ptr + K_offset, mask=K_mask)
        V_data = tl.load(V_ptr + V_offset, mask=V_mask)

        attention_logits = tl.dot(Q_data, tl.trans(K_data)) / attention_logits_scale

        attention_logits_mask = (offset_M[:, None] < M) & (
            current_k_offset[None, :] < N
        )
        attention_logits = tl.where(
            attention_logits_mask, attention_logits, float("-inf")
        )

        current_block_max = tl.max(attention_logits, axis=-1)
        max_value = tl.maximum(current_block_max, softmax_current_max)

        alpha = tl.exp(softmax_current_max - max_value)
        softmax_current_max = max_value

        attention_logits_shift = attention_logits - max_value[:, None]

        softmax_nom = tl.exp(attention_logits_shift)
        softmax_denom = tl.sum(softmax_nom, axis=1)

        softmax_running_sum = tl.fma(softmax_running_sum, alpha, softmax_denom)
        accumulator = tl.fma(accumulator, alpha[:, None], tl.dot(softmax_nom, V_data))

    accumulator /= softmax_running_sum[:, None]

    OUT_offset = offset_M[:, None] * OUT_stride_M + offset_d[None, :] * OUT_stride_d
    OUT_mask = (offset_M[:, None] < M) & (offset_d[None, :] < d)
    tl.store(OUT_ptr + OUT_offset, accumulator, mask=OUT_mask)


@triton.jit
def softmax_attention_v(
    _fz_switch_ptr,
    Q_ptr,
    K_ptr,
    V_ptr,
    OUT_ptr,
    M,
    N,
    d,
    Q_stride_M,
    Q_stride_d,
    K_stride_N,
    K_stride_d,
    V_stride_N,
    V_stride_d,
    OUT_stride_M,
    OUT_stride_d,
    BLOCKSIZE_M: tl.constexpr,
    BLOCKSIZE_N: tl.constexpr,
    BLOCKSIZE_d: tl.constexpr,
):
    opaque_zero = tl.load(_fz_switch_ptr)
    opaque_one = tl.load(_fz_switch_ptr + 1)
    opaque_true = opaque_zero < opaque_one
    opaque_false = opaque_zero > opaque_one

    pid0 = tl.program_id(0)
    pid1 = tl.program_id(1)

    offset_M = pid0 * BLOCKSIZE_M + tl.arange(0, BLOCKSIZE_M)
    offset_d = pid1 * BLOCKSIZE_d + tl.arange(0, BLOCKSIZE_d)
    offset_N = tl.arange(0, BLOCKSIZE_N)

    Q_offset = offset_M[:, None] * Q_stride_M + offset_d[None, :] * Q_stride_d
    Q_mask = (offset_M[:, None] < M) & (offset_d[None, :] < d)
    Q_data = tl.load(Q_ptr + Q_offset, mask=Q_mask)

    accumulator = tl.zeros((BLOCKSIZE_M, BLOCKSIZE_d), dtype=tl.float32)
    softmax_running_sum = tl.zeros([BLOCKSIZE_M], dtype=tl.float32)
    softmax_current_max = tl.full([BLOCKSIZE_M], float("-inf"), dtype=tl.float32)
    attention_logits_scale = tl.sqrt(d + 0.0)

    for current_index in range(0, N, BLOCKSIZE_N):
        current_k_offset = current_index + offset_N
        current_v_offset = current_k_offset

        K_offset = (
            current_k_offset[:, None] * K_stride_N + offset_d[None, :] * K_stride_d
        )
        V_offset = (
            current_v_offset[:, None] * V_stride_N + offset_d[None, :] * V_stride_d
        )

        K_mask = (current_k_offset[:, None] < N) & (offset_d[None, :] < d)
        V_mask = (current_v_offset[:, None] < N) & (offset_d[None, :] < d)

        K_data = tl.load(K_ptr + K_offset, mask=K_mask)
        V_data = tl.load(V_ptr + V_offset, mask=V_mask)

        attention_logits = tl.dot(Q_data, tl.trans(K_data)) / attention_logits_scale

        attention_logits_mask = (offset_M[:, None] < M) & (
            current_k_offset[None, :] < N
        )
        attention_logits = tl.where(
            attention_logits_mask, attention_logits, float("-inf")
        )

        current_block_max = tl.max(attention_logits, axis=-1)
        max_value = tl.maximum(current_block_max, softmax_current_max)

        alpha = tl.exp(softmax_current_max - max_value)
        softmax_current_max = max_value

        if opaque_true:
            if opaque_false:
                attention_logits_shift = attention_logits - max_value[:, None]
            else:
                attention_logits_shift = attention_logits - max_value[:, None]
        else:
            attention_logits_shift = attention_logits - max_value[:, None]

        softmax_nom = tl.exp(attention_logits_shift)
        softmax_denom = tl.sum(softmax_nom, axis=1)

        softmax_running_sum = tl.fma(softmax_running_sum, alpha, softmax_denom)
        accumulator = tl.fma(accumulator, alpha[:, None], tl.dot(softmax_nom, V_data))

    accumulator /= softmax_running_sum[:, None]

    OUT_offset = offset_M[:, None] * OUT_stride_M + offset_d[None, :] * OUT_stride_d
    OUT_mask = (offset_M[:, None] < M) & (offset_d[None, :] < d)
    tl.store(OUT_ptr + OUT_offset, accumulator, mask=OUT_mask)


def solve(kernel, q, k, v, out, switch=None):
    block_m, block_d, block_n = 32, 64, 64
    grid = (triton.cdiv(M, block_m), triton.cdiv(d, block_d))
    common = dict(
        Q_ptr=q,
        K_ptr=k,
        V_ptr=v,
        OUT_ptr=out,
        M=M,
        N=N,
        d=d,
        Q_stride_M=q.stride(0),
        Q_stride_d=q.stride(1),
        K_stride_N=k.stride(0),
        K_stride_d=k.stride(1),
        V_stride_N=v.stride(0),
        V_stride_d=v.stride(1),
        OUT_stride_M=out.stride(0),
        OUT_stride_d=out.stride(1),
        BLOCKSIZE_M=block_m,
        BLOCKSIZE_N=block_n,
        BLOCKSIZE_d=block_d,
        num_warps=4,
    )
    if switch is None:
        kernel[grid](**common)
    else:
        kernel[grid](switch, **common)


def main() -> int:
    if not torch.musa.is_available():
        print("MUSA is unavailable", file=sys.stderr)
        return 2

    q, k, v = Q.to("musa"), K.to("musa"), V.to("musa")
    switch = torch.tensor([0, 1], dtype=torch.int32, device="musa")
    original = torch.empty((M, d), device="musa")
    strict = torch.empty((M, d), device="musa")

    solve(softmax_attention, q, k, v, original)
    solve(softmax_attention_v, q, k, v, strict, switch)
    torch.musa.synchronize()

    passed = torch.equal(original, strict)
    print(f"original: {original.cpu()}")
    print(f"strict:   {strict.cpu()}")
    print("PASS" if passed else "FAIL")
    return 0 if passed else 1


if __name__ == "__main__":
    raise SystemExit(main())

IR evidence (TTGIR, dumped)

Correct:

    %softmax_nom = math.exp %shift                    # one exp
    %alloc = ttg.local_alloc %softmax_nom             # conversion AFTER exp

  Buggy (nested scf.if):
    %3:2 = scf.if %opaque_true -> (blocked, dot_op) { # double-yield
      %alloc = ttg.local_alloc %shift                  # ← conversion PROMOTED
      %load  = ttg.local_load %alloc                   #   into the branch,
      scf.yield %shift, %load                          #   object = shift not exp
    }
    %softmax_nom = math.exp %3#1                       # exp on smem read-back copy
    %softmax_nom_121 = math.exp %3#0                   # exp applied TWICE

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions