Debugging NaN Semantics Between Triton’s Interpreter and JIT

A small but important fix for consistent reduction behavior when NaN values are present

When you’re debugging a Triton kernel, the interpreter is incredibly useful. It runs your code in pure Python/NumPy so you can step through logic, inspect values, and avoid spinning up a GPU. But the interpreter and the JIT-compiled path don’t always agree on every edge case.

Recently I ran into one of those disagreements involving NaN values in reductions — specifically with tl.argmin and tl.argmax when using tie_break_left=False.

The symptom

In the interpreter, when a tensor contained NaN and we used the fast tie-break path, NaN could end up being selected as the minimum or maximum. The JIT path (which uses hardware fmin/fmax instructions) correctly ignored the NaN, as IEEE 754 semantics require.

The result? The same kernel could return different answers depending on whether you were testing with the interpreter or running on GPU. That’s exactly the kind of divergence that makes interpreter-based testing less reliable.

Why NaN behavior is tricky in reductions

In IEEE 754, any comparison involving NaN returns False. So NaN > 5.0 is False, and NaN < 5.0 is also False. This means a naïve reduction loop using regular Python comparisons can accidentally let NaN “win” depending on how the logic is written.

Triton provides two tie-breaking modes for argmin/argmax:

• tie_break_left=True (the default in many cases)
• tie_break_left=False (the “fast” path)

Both should ignore NaN values, just like the hardware instructions do. The interpreter was already handling this correctly for the tie_break_left=True path by using NumPy’s nan* functions. The fast path had fallen back to regular min/max + argmin/argmax, which don’t ignore NaN in the same way.

Root cause in the interpreter

Triton’s interpreter dispatches reduction operations through a ReduceOps class. For the fast tie-break variants of argmin and argmax, it was routing to a generic reduction implementation that ultimately used regular NumPy comparison-based functions.

Those functions treat NaN according to Python’s comparison rules rather than the “treat NaN as missing” rule that np.nanmin, np.nanmax, np.nanargmin, and np.nanargmax follow.

Meanwhile, the JIT path lowers to LLVM intrinsics that match the hardware fminf/fmaxf behavior — NaN is never chosen as the result of a min or max.

The mismatch only appeared in one specific combination of operation + tie-break mode, which is why it slipped through for a while.

The fix

The change was small and localized. In python/triton/runtime/interpreter.py, the dispatch logic for the fast tie-break variants was updated to explicitly use NumPy’s NaN-aware reducers:

• argmin fast path → np.nanmin + np.nanargmin
• argmax fast path → np.nanmax + np.nanargmax

This makes the interpreter’s behavior match the JIT path for NaN inputs.

We also switched a couple of membership checks from in / equality to identity checks (is) to avoid triggering eq overrides on JIT-decorated functions (a separate but related robustness improvement).

A regression test was added that explicitly checks argmin/argmax with tie_break_left=False on tensors containing NaN. The test runs only under the interpreter marker so it catches exactly this class of divergence.

What stayed different (and documented)

Even after the fix, there is still one documented difference between interpreter and JIT for equal non-NaN values under certain tie-break settings. The interpreter tends to return the leftmost index in some cases, while the JIT can return any valid index. This is now called out in a comment so future readers aren’t surprised.

Takeaways

A few things stood out from this:

• Interpreter paths are incredibly valuable for debugging and testing, but they require the same level of attention to floating-point edge cases as the fast path.
• NumPy’s nanmin/nanmax family is usually the right tool when you want “ignore missing values” semantics. Rolling your own comparison logic is easy to get subtly wrong with NaN.
• Small semantic differences between interpreter and JIT can hide for a long time if your test data doesn’t include NaN or specific tie-break combinations. Adding targeted regression tests for these cases pays off.
• When fixing interpreter behavior, it’s worth checking whether the same logic exists in multiple code paths (the left-tie and fast-tie variants in this case).

If you rely on the Triton interpreter for local development or CI testing of kernels that do reductions, this particular divergence is now resolved for argmin/argmax.

The change landed in https://github.com/triton-lang/triton/pull/10699 and the original issue is at https://github.com/triton-lang/triton/issues/10697.