Inside PR #10699: How a Missing elif Let NaN Win in Triton's Interpreter

A line-by-line walkthrough of the fix; from IEEE 754 semantics to why == was quietly the wrong tool for the job.

TL;DR: tl.argmin / tl.argmax with tie_break_left=False gave different answers in Triton's interpreter than in compiled (JIT) mode whenever the input contained NaN. The interpreter's fast dispatch table had a branch for tie_break_left=True but not for the tie_break_left=False variant, so that variant silently fell back to a generic Python reduction where value1 > value2 — which is False whenever either operand is NaN - let NaN win instead of being ignored. The fix is four lines of routing logic, a doc comment about a divergence it doesn't fix, a small correctness cleanup in how functions are compared, and 59 lines of regression tests. PR #10699, fixing issue #10697, merged into triton-lang/triton on Jun 23, 2026.

If you read the first writeup on this bug, that one covered the shape of the problem. This one is the actual autopsy, the real diff, why each piece of it exists, and the one architectural decision in the fix that’s easy to miss but is the part worth remembering.

1. The bug, reproduced

Triton kernels can run two ways: compiled, where your Python gets lowered through LLVM to real GPU/CPU instructions, and interpreted (TRITON_INTERPRET=1), where the same kernel runs as pure Python/NumPy on CPU so you can print() and debug it like normal code. The entire premise of interpreter mode is that it's a faithful stand-in for the compiled path — same inputs, same outputs, easier to poke at.

Here’s the repro from the original issue:

@triton.jit
def reduction_nan_kernel(x_ptr, out_ptr, N: tl.constexpr, BLOCK: tl.constexpr):
    offsets = tl.arange(0, BLOCK)
    mask = offsets < N
    x = tl.load(x_ptr + offsets, mask=mask, other=float("inf"))
    result = tl.argmin(x, axis=0)
    tl.store(out_ptr, result)

x = torch.tensor([float("nan"), 2.0, 1.0], device="cuda", dtype=torch.float32)
# compiled:    tl.argmin result: 2   (correct - ignores NaN, finds the min at index 2)
# interpreter: tl.argmin result: 0   (wrong - NaN "wins")

Two different answers for the same tensor. And silently; no warning, no crash. That’s the dangerous part: it’s a debugging tool that lies to you exactly when you’re using it to debug something else.

2. Why NaN “wins” - the one-sentence version of IEEE 754

If you already know this, skip to §3. If you don’t, this is the whole reason the bug exists:

Source: Image by the author.

Every ordered comparison against NaN ; >, <, ==, >=, <= ;evaluates to False. Not "undefined," not an error. False. It's part of the IEEE 754 spec, and it's consistent: NaN means "not comparable," so no comparison can succeed.

That single fact is enough to break a huge class of naive max/min-tracking code. If your loop looks like if candidate > best: best = candidate, and best starts out as (or becomes) NaN, that condition is False forever, best is stuck. This is exactly what was happening inside Triton's interpreter.

Hardware doesn’t have this problem, because GPU reduction instructions don’t use >/< at all, they use dedicated fmaxf/fminf instructions that implement IEEE 754's minNum/maxNum operations, which are defined to treat NaN as "not a real candidate" and return the other operand. That's the ground truth the interpreter is supposed to match.

3. Anatomy of the divergence

Triton’s interpreter doesn’t literally simulate hardware instruction-by-instruction. For performance, common reduction patterns get fast-dispatched straight to NumPy: ReduceOps.apply_impl in python/triton/runtime/interpreter.py looks at which specific combine function a reduction is using and, if it recognizes it, routes straight to something like np.nanargmin instead of running a slow generic Python loop.

The bug was a gap in that recognition:

Source: Image by the author.

tl.argmin/tl.argmax take a tie_break_left flag that controls which index wins when two elements are exactly tied. That flag selects between two different combine functions internally, _argmin_combine_tie_break_left and _argmin_combine_tie_break_fast (same for argmax). Before this PR, the dispatch table only had a case for the _tie_break_left variant. The _tie_break_fast variant, which is what you get with tie_break_left=False, arguably the more common default path, matched nothing in the table and fell through to generic_reduce, the slow, fully-Python fallback that applies the combine function element-by-element using plain > comparisons. Which, per §2, is exactly where NaN quietly wins.

So the two tie_break_left variants of the same operation were on completely different code paths: one NaN-safe (NumPy), one not (Python >). That's the whole bug in one sentence.

4. The fix

The final diff is small and touches exactly two files : python/triton/runtime/interpreter.py (+11/−5) and a new regression test in python/test/unit/language/test_core.py (+59). Here's the dispatch table, before and after:

Before:

def apply_impl(self, input):
    if self.combine_fn == tl.standard._argmin_combine_tie_break_left:
        return self.min_max(input[0], val_reduce_op=np.nanmin, idx_reduce_op=np.nanargmin)
    elif self.combine_fn == tl.standard._argmax_combine_tie_break_left:
        return self.min_max(input[0], val_reduce_op=np.nanmax, idx_reduce_op=np.nanargmax)
    elif self.combine_fn == tl.standard._elementwise_max:
        return self.min_max(input[0], val_reduce_op=np.nanmax, idx_reduce_op=None)
    elif self.combine_fn == tl.standard._elementwise_min:
        return self.min_max(input[0], val_reduce_op=np.nanmin, idx_reduce_op=None)
    elif self.combine_fn == tl.standard._sum_combine:
        return self.sum(input[0])
    else:
        # Fall back to the slow mode

After:

def apply_impl(self, input):
    # Note: np.nanargmin/np.nanargmax always return the leftmost index
    # for equal values, whereas tie_break_fast on hardware returns an
    # arbitrary index. This is a known remaining divergence between the
    # interpreter and JIT for inputs with equal non-NaN elements.
    if (self.combine_fn is tl.standard._argmin_combine_tie_break_left
            or self.combine_fn is tl.standard._argmin_combine_tie_break_fast):
        return self.min_max(input[0], val_reduce_op=np.nanmin, idx_reduce_op=np.nanargmin)
    elif (self.combine_fn is tl.standard._argmax_combine_tie_break_left
            or self.combine_fn is tl.standard._argmax_combine_tie_break_fast):
        return self.min_max(input[0], val_reduce_op=np.nanmax, idx_reduce_op=np.nanargmax)
    elif self.combine_fn is tl.standard._elementwise_max:
        return self.min_max(input[0], val_reduce_op=np.nanmax, idx_reduce_op=None)
    elif self.combine_fn is tl.standard._elementwise_min:
        return self.min_max(input[0], val_reduce_op=np.nanmin, idx_reduce_op=None)
    elif self.combine_fn is tl.standard._sum_combine:
        return self.sum(input[0])
    else:
        # Fall back to the slow mode

Three things happened here, and they shipped as three of the PR’s five separate commits rather than one blob:

• **fix(interpreter): handle NaN in argmin/argmax with tie_break_left=False** - the actual fix. Both fast-path variants now route to the same NaN-aware np.nanmin/np.nanargmin path as their tie_break_left=True siblings, or-ed into the same condition.
• **fix(interpreter): consolidate argmin/argmax NaN branches** :cleanup pass merging what could have been four separate elif branches into two, once it was clear the NaN-handling behavior is identical for both tie-break modes.
• **fix(interpreter): document tie-breaking divergence in argmin/argmax** - the comment block. More on this in §5.

5. The detail worth slowing down on: == became is

Notice every == in the "before" block became is in the "after" block, including on the branches that weren't buggy (_elementwise_max, _elementwise_min, _sum_combine). This was its own commit: **fix(interpreter): use identity checks in apply_impl for jit functions**.

This is the kind of change a junior engineer might reasonably skip as a style nit, and it’s exactly the kind of change a senior reviewer notices and appreciates. self.combine_fn here is a reference to a function wrapped by @triton.jit. Comparing JIT-decorated function objects with == invokes Python's normal equality machinery — which, depending on how __eq__ is or isn't implemented on the wrapper, can be slower than necessary, or in principle behave in a way you didn't intend, for a check that only ever means one thing: "is this the literal same function object I'm expecting?" is says exactly that and nothing else ; identity, not equality; and it's cheaper. Once one branch needed to check "is this function A or function B," using is consistently across the whole dispatch table isn't just a fix, it's a small piece of defensive consistency: the kind of change that costs nothing and removes a whole category of "wait, why does this branch behave differently" questions for the next person reading the table.

6. What the fix deliberately did not fix

Read the comment again:

np.nanargmin/np.nanargmax always return the leftmost index for equal values, whereas tie_break_fast on hardware returns an arbitrary index. This is a known remaining divergence between the interpreter and JIT for inputs with equal non-NaN elements.

This is arguably the most senior-engineer-flavored line in the whole PR, and it’s not code, it’s a comment. tie_break_left=False is supposed to mean "don't guarantee which index wins on an exact tie, hardware picks whatever's fastest." NumPy's nanargmin/nanargmax, on the other hand, always deterministically pick the leftmost index. So even after this fix, if you feed the interpreter [3.0, 5.0, 5.0] and ask for argmax with tie_break_left=False, it will always say index 1, while compiled mode might legitimately say index 1 or 2 depending on hardware. That's still a divergence.

The PR doesn’t fix that, and it says so, explicitly, right next to the code. That’s the difference between “I made the symptom in front of me go away” and “I understand the exact shape of what I fixed and what I didn’t, and I’ve left a note so nobody has to rediscover this the hard way.” Fixing the NaN bug closed issue #10697. Silently also “fixing” the tie-break-order behavior would have been scope creep into a different semantic contract, one where “arbitrary” doesn’t mean “wrong,” so there’d be nothing to fix, just a decision to make about whether the interpreter should bother matching an intentionally-unspecified hardware behavior at all.

7. The regression test

The new test, test_argmax_argmin_tie_break_fast_with_nan, is marked @pytest.mark.interpreter - Triton's test suite can select interpreter-only tests, which matters because this is purely an interpreter-mode correctness bug; it has nothing to verify on the compiled path. The test builds two tiny @triton.jit kernels ; one for argmax, one for argmin; and runs each against three tensors:

# argmax: [nan, 6, 8] -> max=8.0, argmax=2        (NaN at the start)
# argmax: [3, 5, nan] -> max=5.0, argmax=1        (NaN at the end)
# argmin: [3, nan, 1] -> min=1.0, argmin=2        (NaN in the middle)

That’s not padding , NaN-at-start, NaN-at-end, and NaN-in-the-middle are three genuinely different code paths through a reduction loop (first element becomes the initial “best” vs. NaN has to displace an existing “best” vs. NaN sits somewhere it never gets compared as a boundary case), and a fix that only handled one position wouldn’t actually prove the general case. Covering all three in ~15 lines is the kind of test-writing instinct that’s cheap to do and expensive to skip.

8. What happened after it merged

Source: Image by the author.

A fix that only exists on main is a fix nobody's using yet. What actually happened here: the commit got cherry-picked into Meta's internal fb-experimental-triton fork multiple times as that fork periodically re-synced against upstream, meaning someone judged it worth carrying forward rather than waiting for a routine sync. A week later, two follow-up PRs (#10760, #10761) landed, specifically closing a gap where test_reduce1d didn't cover the tie_break_left=False path at all. That's a second, independent signal: someone else looked at the area this PR touched and found — and closed — a testing hole next to it.

Neither of those things were guaranteed to happen. A correct fix can still merge and sit unused. This one didn’t.

9. The takeaway

If you’re newer to this: the lesson isn’t “NaN is scary,” it’s that any codebase with two independent implementations of the same logic (an interpreter and a compiler, a mock and the real client, a fast path and a slow path) has a standing liability that they’ll quietly disagree, and the only defense is tests that specifically probe the disagreement — not just the happy path both implementations already agree on.

If you’re further along: the parts of this PR worth stealing for your own reviews aren’t the four lines that fixed the bug. They’re the three things wrapped around those four lines, consolidating duplicate branches once the fix made them identical, converting == to is once you're comparing function identity, and writing down in the code the exact boundary of what you did not fix. None of that was required to close the issue. All of it is why the fix was trusted enough to get cherry-picked into a production-adjacent fork within a week.

PR: triton-lang/triton#10699 · Issue: #10697 · Part 1: Debugging NaN Semantics Between Triton’s Interpreter and JIT