AITER test_compile_constexpr_mutation.py Triton Test Fails with 99%+ Mismatched Elements

[jayfurmanek]

🐛 Describe the bug

Problem Description

At a fixed AITER commit (ROCm/aiter@de79999), Triton unit test op_tests/triton_tests/torch_compile/test_compile_constexpr_mutation.py fails with a big accuracy mismatch in all cases:

AssertionError: Tensor-likes are not close!
Mismatched elements 99%+

At first glance, it's a simple RMSNorm kernel:

@triton.jit
def _rmsnorm_constexpr_kernel(
    x_ptr,
    weight_ptr,
    out_ptr,
    eps,
    stride_x,
    stride_out,
    N: tl.constexpr,
    BLOCK_N: tl.constexpr,
):
    row = tl.program_id(0)
    cols = tl.arange(0, BLOCK_N)
    mask = cols < N

    x = tl.load(x_ptr + row * stride_x + cols, mask=mask, other=0.0).to(tl.float32)
    variance = tl.sum(x * x, axis=0) / N
    x_hat = x / tl.sqrt(variance + eps)

    w = tl.load(weight_ptr + cols, mask=mask, other=0.0).to(tl.float32)
    tl.store(out_ptr + row * stride_out + cols, (x_hat * w).to(tl.float16), mask=mask)

It's possible to get the test passing, keeping AITER code base the same and changing Triton compiler.

Software Commits

• AITER commit: ROCm/aiter@de79999
• Triton good commit - test pass: triton-lang/triton@8ea50fa
• Triton bad commit - test fail: triton-lang/triton@454ec9e

It was possible to determine these commits through bisection of Triton compiler:

Date	Status	Commit
Dec 02 2005	Good	triton-lang/triton@8ea50fa
Dec 02 2025	First bad	triton-lang/triton@454ec9e
May 07 2026	v3.7.0 tag, still bad	triton-lang/triton@5f3f125
Jun 24 2026	Current TOT, still bad	triton-lang/triton@7917e3f

Environment

• OS: Ubuntu 24.04.4 LTS (Noble Numbat)
• CPU: AMD EPYC 9575F
• GPU: AMD Instinct MI355X (gfx950) and AMD Instinct MI300X (gfx942)
• ROCm 7.2.4

Steps to Reproduce

1. Setup reproducer container

Run on host:

docker pull rocm/pytorch:latest && \
docker run -dt \
    --device /dev/kfd \
    --device /dev/dri \
    --shm-size 16g \
    --network host \
    --ipc host \
    --cap-add SYS_PTRACE \
    --security-opt seccomp=unconfined \
    --user root \
    --group-add $(getent group video | cut -d: -f3) \
    --workdir /workspace \
    --name triton_internal_issue_1883 \
    rocm/pytorch:latest && \
docker exec -it triton_internal_issue_1883 bash

2. Compile Triton (at good commit)

pushd /workspace && \
git clone https://github.com/triton-lang/triton.git && \
pushd /workspace/triton && \
git checkout 8ea50fa4098642ee0a6afc0b932f83a38e37719a && \
pip uninstall --yes triton && \
pip install --requirement python/requirements.txt && \
pip install --verbose --no-build-isolation --editable . && \
pip show triton && \
python python/tutorials/01-vector-add.py && \
popd && popd

3. Compile AITER

pushd /workspace && \
git clone --recursive https://github.com/ROCm/aiter.git && \
pushd /workspace/aiter && \
pip uninstall --yes amd-aiter && \
AITER_USE_SYSTEM_TRITON=1 python setup.py develop && \
pip show amd-aiter && \
python -c 'import aiter' && \
popd && popd

4. Verify good commit

pushd /workspace/aiter && \
rm -rf ~/.triton/cache && \
pytest op_tests/triton_tests/torch_compile/test_compile_constexpr_mutation.py && \
popd

Expected result: 6 passed, 14 warnings

5. Checkout bad commit and compile Triton again

pushd /workspace/triton && \
git checkout 454ec9eea6cf6deab971ed66dad45d957c561fe4 && \
pip uninstall --yes triton && \
git clean -xdf && \
pip install --requirement python/requirements.txt && \
pip install --verbose --no-build-isolation --editable . && \
pip show triton && \
python python/tutorials/01-vector-add.py && \
popd

6. Reproduce failure at bad commit

pushd /workspace/aiter && \
rm -rf ~/.triton/cache && \
pytest op_tests/triton_tests/torch_compile/test_compile_constexpr_mutation.py && \
popd

Expected result: 6 failed, 14 warnings

AssertionError: Tensor-likes are not close!
Mismatched elements 99%+

7. Export reproducer container

docker commit triton_internal_issue_1883 rocm/pytorch-private:triton_internal_issue_1883
docker login -u rocmshared
docker push rocm/pytorch-private:triton_internal_issue_1883
docker logout

You can go straight to step 6 after pulling the container exported in step 7.

Matching Triton issue here: https://github.com/AMD-Triton/triton-tickets/issues/1883

Error logs

No response

Versions

[jayfurmanek]

Triton team determined root cause:

Root cause

Triton PR #8846 rewrote ASTFunction.deserialize/build_value so it calls fn.set_arg_attr(cursor, …) for every kernel parameter that has an attribute.

It crashes when a parameter is simultaneously:

1. treated as a constant (no MLIR IR argument at ttir construction), and
2. carrying a tt.divisibility attribute (16-aligned scalar).

Native Triton never produces that combination (it keeps constants and attrs disjoint — jit.py::_pack_args force-blanks attrs for constexpr-typed params). But PyTorch's generate_ttir (used to analyze which tensors a kernel mutates) marks every non-tensor arg as a constant while still emitting divisibility attrs for aligned scalars (e.g. stride_x, stride_out = multiples of 16). So set_arg_attr is called on a non-existent IR arg → IndexError: Function argument index out of range.

Where the bug happens

ASTFunction.build_value (inside deserialize) in python/triton/compiler/code_generator.py:

# BUGGY (#8846)
def build_value(path, ty):
    nonlocal cursor, handles
    # > set attributes
    attr_specs = self.attrs.get(path, [])
    for attr_name, attr_val in attr_specs:
        fn.set_arg_attr(cursor, attr_name, attr_val)   # <-- for a constant param, cursor
                                                       #     points past the real IR args
    # > build frontend value
    val, cursor = ty._unflatten_ir(handles, cursor)
    set_iterable_path(vals, path, val)

Note: a genuine tl.constexpr does not trigger this — it has no attr entry, so attr_specs is [] and set_arg_attr is never called. The crash needs a param that is both IR-arg-less and has an attr (the constant-folded aligned scalar) triggered by pytorch's generate_ttir.

Throw site — the C++ binding set_arg_attr in python/src/ir.cc:

if (arg_no >= self.getNumArguments())
  throw py::index_error("Function argument index out of range");

PyTorch Eliminates Kernel

identify_mutated_tensors catches the exception and "fails safe" by returning all tensor inputs — BUT its fallback is [k for k,v in kwargs if isinstance(v, Tensor)]. This runs at 4 pipeline stages:

• 3 tracing stages use FakeTensor (a torch.Tensor) → fallback returns the inputs → safe.
• The Inductor lowering stage uses TensorBox (not a torch.Tensor) → fallback returns [] → "nothing is mutated."

With the output reported as unmutated, Inductor concludes the kernel is side-effect-free and dead-code-eliminates it, returning the uninitialized output buffer → garbage results.

Possible Triton fix

In ASTFunction.build_value: unflatten first, then apply argument attributes only if the param consumed a real IR argument (cursor > arg_start), guards that only if there exists an actual IR at the correct index, set an attr for it.

# FIXED
def build_value(path, ty):
    nonlocal cursor, handles
    # > build frontend value
    arg_start = cursor
    val, cursor = ty._unflatten_ir(handles, cursor)
    # > set attributes (only for args that map to real IR arguments;
    # > constexpr args consume no handles and have no IR argument to attach to)
    if cursor > arg_start:
        attr_specs = self.attrs.get(path, [])
        for attr_name, attr_val in attr_specs:
            fn.set_arg_attr(arg_start, attr_name, attr_val)
    set_iterable_path(vals, path, val)

Verified: aiter test 6/6 pass; frontend unit + gluon suites pass; native divisibility attrs unaffected.

aiter fix (stopgap)

Suppress the alignment attr on the 16-aligned scalar args so the constant+attr collision never forms:

@triton.jit(do_not_specialize_on_alignment=["stride_x", "stride_out"])
def _rmsnorm_constexpr_kernel(x_ptr, weight_ptr, out_ptr, eps,
                              stride_x, stride_out,
                              N: tl.constexpr, BLOCK_N: tl.constexpr):
    ...

Works (verified), but seems like a hack.

Reproducer triggered by torch.compile

Returns garbage (silently) on buggy Triton; correct with the fix. Run with rm -rf ~/.triton/cache && TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 python repro.py.

import torch
import triton
import triton.language as tl


@triton.jit
def scale_kernel(in_ptr, out_ptr, n, BLOCK: tl.constexpr):
    offs = tl.arange(0, BLOCK)
    mask = offs < n
    tl.store(out_ptr + offs, tl.load(in_ptr + offs, mask=mask) * 2.0, mask=mask)


def fn(x):
    out = torch.empty_like(x)
    scale_kernel[(1,)](x, out, x.numel(), BLOCK=64)   # n=64 is 16-aligned -> tt.divisibility
    return out


x = torch.arange(64, device="cuda", dtype=torch.float32)
out = torch.compile(fn, backend="inductor", fullgraph=True)(x)
print("max_abs_err =", (out - x * 2).abs().max().item())
assert torch.allclose(out, x * 2), "kernel was DCE'd -> garbage output"
print("PASS")

• Buggy Triton: logs IndexError('Function argument index out of range'), then fails the assert (out is uninitialized).
• Fixed Triton: max_abs_err = 0.0 → PASS.

Key commits / versions

• Bad: triton-lang/triton#8846 (454ec9e); last good parent 8ea50fa (same LLVM hash — no LLVM bump involved).
• Env: Triton 3.7.0+ (incl. current TOT), torch 2.10.0, ROCm 7.2.4, gfx950/gfx942.

[jayfurmanek]

For this issue, its a combination of pytorch inductor/triton. Inductor promotes a runtime arg to a constant and tries to assign an attribute to it. Prior to commit PR #8846, triton had a separate loop to iterate over constant args and it was ignoring the attribute assignment. PR #triton-lang/triton#8846 consolidated everything into a single loop, pulls the attribute from map and tries to assign it to IR for efficiency. Constants do not have an IR in triton and this results in index out of bounds error in triton ttir generation. The inductor catches the exception, usually this would make the analysis fold to conservative assumption but in this case it makes an aggressive decision that the kernel has no side-effects and dead code eliminates it. We can have a stop-gap fix in triton, but the fix should also go in pytorch. It should not dead code eliminate a kernel when it cannot prove that it is side-effect free. Above kernel without inductor would just be fine, and in this case it is triggered by combination of inductor/triton. Can you help me create a pytorch issue for this (who should we contact and ticket template)? I can then use this as a reference when making the triton modification.

[iupaikov-amd]

@jayfurmanek it looks like you used pytorch 2.10 judging by the env comment. Pytorch has pytorch#170808 addressing the issue. This kernel would be conservatively marked as mutating and not DCE'd, even with the buggy Triton. I think you should take the latest pytorch behavior as base for your fix.

There also is pytorch#170049 which addresses related things I think. It's also merged and completed.

We probably can't really change how upstream addresses such kernels, but I can look into that if you want (or conditionalize for ROCm).

[iupaikov-amd]

On the PyTorch-side ask: "don't DCE a kernel when it can't prove it's side-effect free"

Confirmed on MI300X (gfx942) that this is exactly what pytorch#170808 already does, and verified it resolves the silent-garbage failure.

When TTIR generation raises the IndexError, identify_mutated_tensors (renamed identify_accessed_tensors on newer branches) catches it and conservatively marks inputs as mutated so the kernel can't be DCE'd. pytorch#170808 fixed the fallback to keep both Tensor and TensorBox:

return [
    key
    for key, value in kwargs.items()
    if isinstance(value, (Tensor, torch._inductor.ir.TensorBox))  # #170808
]

Before pytorch#170808 the fallback kept only Tensor. During Inductor lowering the kernel args are TensorBox, so they were dropped -> the kernel had no mutation outputs -> Inductor treated it as side-effect-free and dead-code-eliminated it -> uninitialized output buffer -> the 99% mismatch.

A/B test — hold Triton buggy so the IndexError keeps firing, toggle only the fallback:

PyTorch fallback	compiled result
(Tensor, TensorBox) (pytorch#170808)	max_abs_err = 0.0 ✅ kernel kept
(Tensor,) (pre-pytorch#170808)	max_abs_err = 126.0 ❌ kernel DCE'd -> garbage

Env: torch 2.11.0+rocm7.13, Triton 3.6.0, MI300X. Triggered with the buggy ASTFunction.deserialize from triton@454ec9e; forcing the analysis to raise gives the identical result.

Minimal, version-independent repro (forces the analysis failure, so it does not depend on the Triton build):

import os, torch, triton, triton.language as tl
import torch._higher_order_ops.triton_kernel_wrap as tkw

# Force the compile-time mutation analysis to fail the same way buggy Triton does.
tkw.generate_ttir = lambda *a, **k: (_ for _ in ()).throw(
    IndexError("Function argument index out of range"))

@triton.jit
def scale_kernel(in_ptr, out_ptr, n, BLOCK: tl.constexpr):
    offs = tl.arange(0, BLOCK)
    mask = offs < n
    tl.store(out_ptr + offs, tl.load(in_ptr + offs, mask=mask) * 2.0, mask=mask)

def fn(x):
    out = torch.empty_like(x)
    scale_kernel[(1,)](x, out, x.numel(), BLOCK=64)  # n=64 -> 16-aligned
    return out

x = torch.arange(64, device="cuda", dtype=torch.float32)
out = torch.compile(fn, fullgraph=True)(x)
print("max_abs_err =", (out - x * 2).abs().max().item())
assert torch.allclose(out, x * 2), "kernel was DCE'd -> garbage output"
print("PASS")

Run: rm -rf ~/.triton/cache && TORCHINDUCTOR_FORCE_DISABLE_CACHES=1 python repro.py

• With [inductor][triton_kernel_wrap] Return TensorBox values when Triton compile failed in identify_mutated_tensors() pytorch/pytorch#170808: max_abs_err = 0.0 -> PASS
• Reverting [inductor][triton_kernel_wrap] Return TensorBox values when Triton compile failed in identify_mutated_tensors() pytorch/pytorch#170808 (fallback back to isinstance(value, Tensor)): max_abs_err = 126.0 -> FAIL (kernel DCE'd)

So pytorch#170808 satisfies the "don't DCE when the analysis is inconclusive" requirement.