[Frontend] Unblock buffer reuse and triton_helpers on the Triton route by YWHyuk · Pull Request #308 · PSAL-POSTECH/PyTorchSim

YWHyuk · 2026-08-03T15:35:01Z

Two fixes that came out of running the whole test suite on the Triton route
(scripts/ci/triton_route_sweep.py). Independent of each other; both measured.

Stacked on #305 — PyTorchSimFrontend/triton_backend/ does not exist on
develop yet, so this targets feature/triton-codegen.

1. Enter memory planning through the pass upstream provides

generate() called memory_plan_reuse() directly, bypassing
run_wrapper_ir_passes() — the method that chooses a planner and sets up the
state that planner needs. Three things followed:

torch 2.10 added self.estimate_peak, created in run_wrapper_ir_passes and
read by should_reuse_buffer during planning. Every graph that reached buffer
reuse died with AttributeError. In CI that is 16 of 69 tests on the
Triton route — Llama, CLIP, ConvNeXtV2, MobileNet, Mixtral, and every
attention test.
our own memory_plan() override was unreachable: run_wrapper_ir_passes is
its only caller.
generate() took is_inference and never used it.

Calling the upstream entry point fixes all three. Behaviour is identical today:
config.memory_planning defaults to False, so upstream takes the same branch
we were taking by hand.

This file is shared with the MLIR route, so the fix applies there too.

2. Give the tnpu venv torch's `triton_helpers`, by copying it

Inductor's kernels call triton_helpers.maximum, .max2, .sort_with_index.
The module is inside torch and the tnpu venv deliberately has none, so those
kernels were rejected outright.

Rewriting them was never necessary. triton_helpers.py imports nothing from
torch — only .triton_compat, and that module touches torch in three places
(torch.version.hip ×2, torch.autograd.profiler ×1), none of which
triton_helpers needs. So the installed torch's file is copied verbatim beside
the kernel and paired with a small triton_compat that resolves the same seven
names straight from triton, mirroring upstream's fallbacks.

Copying rather than vendoring a snapshot keeps the helpers matched to the torch
that generated the kernel, and means a helper we have not seen yet needs no work.

The spec also puts its own directory on sys.path: the kernel is loaded by
path, so a sibling package would not otherwise be importable from it.

Verification

tests/ops/elementwise/test_activation.py: the three relu kernels use
triton_helpers.maximum and now reach an ELF and a trace producer, where before
they were rejected. The sigmoid, silu and swiglu kernels in the same file were
never blocked by this — they use tl.sigmoid, a triton builtin — but could not
run either, because the test aborts at its first failure and relu is first.

tests/ops/attention/test_gqa.py gets past estimate_peak to its next blocker.

MLIR route unaffected: test_add, test_matmul, test_prologue_fusion,
test_softmax all still pass.

Gate (scripts/ci/triton_route_passing.txt) still 11/11.

Not fixed here

Both fixes now converge on the same next blocker: libdevice extern
intrinsics (exp, tanh, rsqrt, erf, isnan), which have no
triton_shared implementation. Worth noting that tl.sigmoid lowers to
math.exp and works, so a substitution may be cheaper than a lowering pass —
that is the next thing to measure, not part of this PR.

generate() called memory_plan_reuse() directly, bypassing run_wrapper_ir_passes(), which is what chooses a planner and sets up the state that planner needs. Three things followed from that. torch 2.10 added self.estimate_peak, created in run_wrapper_ir_passes and read by should_reuse_buffer during planning, so every graph that reached buffer reuse died with AttributeError. In CI that is 16 of 69 tests on the Triton route -- Llama, CLIP, ConvNeXtV2, MobileNet, Mixtral and every attention test. Our own memory_plan() override was unreachable: run_wrapper_ir_passes is its only caller. generate() took is_inference and never used it. Calling the upstream entry point fixes all three and keeps behaviour identical today, since config.memory_planning defaults to False and that branch is the one we were already taking by hand. Verified: Triton route test_gqa gets past this to its next blocker; MLIR route test_add, test_matmul, test_prologue_fusion and test_softmax still pass.

Inductor's kernels call triton_helpers.maximum, .max2, .sort_with_index; the module is inside torch and the tnpu venv deliberately has none, so those kernels were rejected outright. Rewriting them was never necessary. triton_helpers.py imports nothing from torch -- only .triton_compat, and that module touches torch in three places (torch.version.hip twice, torch.autograd.profiler once), none of which triton_helpers needs. So the installed torch's file is copied verbatim beside the kernel and paired with a small triton_compat that resolves the same seven names straight from triton, mirroring upstream's fallbacks. Copying rather than vendoring a snapshot keeps the helpers matched to the torch that generated the kernel, and means a helper we have not seen yet needs no work. The spec now also puts its own directory on sys.path: the kernel is loaded by path, so a sibling package would not otherwise be importable from it. Verified on tests/ops/elementwise/test_activation.py: the three relu kernels use triton_helpers.maximum and now reach an ELF and a trace producer, where before they were rejected. The sigmoid, silu and swiglu kernels in the same file were never blocked by this -- they use tl.sigmoid, a triton builtin -- but could not run either, because the test aborts at its first failure and relu is first. Gate still 11/11.

triton/language/extra/libdevice.py is a file of empty stubs, and a backend binds them through get_module_map. triton_shared returned {}, so libdevice.tanh(x) evaluated to None and the kernel died on "cannot convert None of type NoneType to tensor". strip_for_tnpu then made it worse by rejecting the kernel outright with a claim that turned out to be false: that these intrinsics have no lowering here. They do, on both sides of the seam. triton-shared maps 22 of the 49 __nv_* symbols Inductor can emit to math.*, triton-npu's normalize_upstream derives 25 more, and lower_to_vcix turns math.exp/tanh/erf/sin/cos/log/atan into the VCIX custom instructions. The one symbol neither knows, __nv_llrintf, is refused by name. Nothing was missing but the binding, which is df7452a in the fork and d46995f in the harness. So the rejection goes, and libdevice is re-bound the way tl_math already was -- from triton, since torch only re-exports it and strip_for_tnpu drops the torch import. The whole prefix is unconditional now: an unused import costs nothing, triton's jit skips libdevice stubs when hashing, and the conditionals were three regexes deciding something they did not need to decide. Measured on this NPU target: libdevice.tanh and .exp reach the VCIX instructions with err 5.96e-08 and 2.38e-07, .rsqrt goes through math.rsqrt to LLVM with err 0, .isnan becomes arith.cmpf uno and is correct on NaN input. tests/ops/elementwise/test_exponent.py and test_transcendental.py now pass; the latter puts exp, tanh, sin, cos and erf through the VCIX path in one kernel. preflight grows a check for the manifest's new `also_reads`, because restore.sh clones triton_shared as of d46995f and a token without read on it would fail inside the image build instead of here.

fixed_config_for passed kernel.numels straight through, but that dict is keyed by iteration-space prefix ('x', 'r0') while parallel_axes looks for collect_meta's '<prefix>numel' form. Nothing matched, so every kernel came back with no parallel axes and fell through to the XBLOCK default.

Inductor emits tl.debug_barrier() before a store when the buffer is in-place and the index is broadcast. It is not part of the computation: the comment on codegen/triton.py:3554 says it works around triton-lang/triton#1615, where a load broadcast over several warps can observe a later store, and the barrier forces every warp to have read first. There are no warps here. tnpu registers work-items and replays them one at a time, so the race the barrier orders cannot occur -- the same reason set_driver_to_gpu is dropped, and it goes in the same place. Leaving it in was not free: it reaches ttir as `ttg.barrier`, an op of triton's GPU dialect, and triton-shared-opt stops at the parse -- "Dialect `ttg' not found for custom op 'ttg.barrier'" -- before any lowering runs. tests/ops/conv/test_pool.py now has no ttg.barrier in its ttir and gets through stage 2, where it stops on the same linalg.index rank error as tests/ops/reduce/test_reduce.py. That one is a tnpu pass bug, reported upstream.

The pattern was the bare word TOGSim, which appears in the INFO lines the simulator writes on a normal run. So any test that got far enough to simulate was filed under togsim whatever actually went wrong, and the further the route progressed the more it collected: 7 tests in the last sweep, which turn out to be four different things. It now matches only a TOGSim failure, and wrong_values is tested first -- a kernel that simulated cleanly and then compared wrong is a numerics failure, and by that point the log is full of TOGSim lines either way. FileNotFoundError gets its own bucket rather than landing in `other`. Re-running the classifier over the same sweep: togsim 7 -> wrong_values 3 (softmax, layernorm, indirect_access) missing_artifact 2 (transpose3D, vectorops) togsim 1 (test_hetro, KeyError 'vpu_num_lanes') other 1 (floormod) The three numerics failures are the ones worth having: they run end to end -- Spike, gem5 and TOGSim all succeed -- and produce wrong numbers. softmax is out by 0.204.

Three separate bugs wore one label. None of them was the block-size policy the README blames. parallel_axes matches "<prefix>numel" keys, and fixed_config_for was handing it kernel.numels, which is keyed by bare prefix. It matched nothing, so the axis list was always empty, no YBLOCK was ever written, and grid_of failed on YBLOCK=None. The policy the docstring describes had never run once -- its "inner blocks pinned to 1" warning had never fired either. With blocks set, the grid came out reversed. Inductor reads program_id(0) as x and tnpu's wrapper takes spec.grid positionally as gridX/Y/Z, but grid_of emitted outermost-first, so a 2-D kernel got its two extents swapped. pid_axes now names that order and the three places that build a grid tuple share it; parallel_axes keeps outermost-first for tile-shape decisions, which is what it was for. That left transpose+add wrong while plain transpose passed. Inductor allocates with empty_strided -- (62,34,44) at stride (1496,1,34) here -- and indexes by that stride, while write_inputs called .contiguous() and read_outputs did view_as. Both wrote the right values to the wrong places. collect_meta now records size and stride, and the launch moves through as_strided so element k in the file is element k in the kernel's index space. Probes isolated it: 2-D transpose with no mask, with a mask, with YBLOCK past the axis, with two operands whose lane axes conflict, and with runtime numels shadowed by constants all came back err 0 -- so the tile layout, the masking and select_lane_axis were never the problem. Also: the sweep gave each test its own dump dir but not its own TORCHINDUCTOR_CACHE_DIR. extension_config re-points that at codegen time, by which point Inductor has already put the first graph in the shared /tmp cache, so ten concurrent tests collided there and four failed on artifacts another process had removed -- one of them tripping the gate as a regression that did not reproduce alone. 18 of 69 now pass through the route, from 11. test_conv_fusion and test_conv_view_input turn out not to have been multi-axis failures at all: with the grid fixed they reach convolution_overrideable, which is the device layer.

A kernel like index_put writes some elements of a tensor and leaves the rest alone, so Inductor passes the tensor it mutates straight through: triton_npu_fused_index_put_0(arg2_1, arg1_1, arg0_1, 16384) out_ptr0 = arg0_1, a graph input It arrives in kernel.args.output_buffers, so we called it "out", and write_inputs fills an "out" with zeros. Every element the kernel did not write came back zero. Measured on a[idx, :] = b, 128x128 with 128 random indices covering 77 distinct rows: the 51 rows the indices never name were exactly the 51 wrong rows, and all of them were zero. Rows the indices did name were correct throughout, which is why this looked like a scatter bug rather than a seeding one. A graph input in output_buffers is being mutated, not produced, so it is "inout". tests/ops/misc/test_indirect_access.py now passes both of its scatter cases; it goes on to fail in index_add, which is a tnpu pipeline error and a separate thing. Gate unaffected. Note the label: test_scatter_full and test_indirect_vectoradd both report themselves as "Indirect VectorAdd", and it is the first that was failing.

A functional launch that failed said only RuntimeError: [Spike] triton_npu_fused_cat_0 failed: tnpu.pipeline.StageError: command failed with exit code 255 because tnpu.spike raises on the exit code and spike's own output goes to stage.log, which nothing read. The two kernels that fail this way are the only ones in the suite that compile to a working binary and then go wrong at run time, and neither could be diagnosed. functional.run now raises TnpuError with stage.log appended, the same way run_pipeline already did, and the signal pattern learns the lines spike ends on -- traps, illegal instructions, a rejected --isa. Both now name the fault: [Spike] triton_npu_fused_cat_0 failed User load segfault @ 0x00000004d0801220 [Spike] triton_npu_fused_constant_pad_nd_0 failed User load segfault @ 0x00000004d0800fe0 Both addresses are in the scratchpad window (0xd0000000 + 8 MiB), so the kernel reads past the end of a banked tile rather than corrupting a value -- which is a different bug from the one the numbers suggested. Gate still 18/18.

…be used Two failures in the sweep that were not about the Triton route. test_sparse_core built its model as SparseMLP(input, hidden, output, device), but the fourth parameter is sparsity_fc1 -- so a torch.device reached _apply_pruning and compared against a tensor: TypeError: '>' not supported between instances of 'Tensor' and 'torch.device' The MLIR route fails the same way, since this is in the test's __init__ before anything compiles. Passing device by keyword fixes it and the test passes on the Triton route, so it joins the allowlist. test_hetro runs on stonne_big_c1_simple_noc.yml, which has no vpu_num_lanes -- it describes a sparse accelerator, not a VPU. fixed_config_for read the key unconditionally and raised a bare KeyError from inside extension_config. It now names the config and what the route needs it for. The test still fails, because this route genuinely cannot pin block sizes without a lane count, but it says so. Gate 19/19.

…rnel aten.convolution, mm, bmm, sort and cat currently leave this route through extern_kernels. conv raises convolution_overrideable not implemented; the rest fall back to eager, which returns the right numbers and simulates nothing -- test_matmul and test_bmm "pass" today without emitting a kernel. Inductor's templates are the alternative and they are not GPU code: torch ships one triton_mm.py.jinja for cuda, xpu, mtia and cpu, and the body is the shape tnpu already compiles -- tl.zeros accumulator, masked tl.load, acc += tl.dot, tl.store. What keeps us out is use_triton_template gating on is_gpu, and GPU_TYPES being a hardcoded list; mtia is in it because Meta edited torch. So inductor_templates.py puts npu in that list and takes the consequences. Two of them are "there is no GPU here" and upstream cannot help: GPU_TYPES itself, and has_triton, which asks whether a supported device is available rather than whether triton is installed -- the same reasoning as the hash shim in _triton_compat. The rest goes through the extension points torch provides: register_template_heuristic for the config table. Selection is ours rather than autotuned. The autotuner compiles each candidate and times it on the device, and there is no device; pick_config takes the offered order and carries a TODO to rank by TOGSim cycles instead, which timing.run_togsim already produces per kernel. Precompilation is skipped for the same reason -- we need the chosen kernel's source, not a GPU binary. Two fixes fell out of turning the gate on. device_guard returned "pass", which the caller writes as `with pass:`; cpu gets away with that because device_need_guard is is_gpu. And wrap_kernel_call joins its arguments as strings, which a template kernel breaks by passing sympy constants. Both templates now render and reach tnpu: mm stops at stage 2 on tl.assume -> llvm.intr.assume, conv at stage 3 in select_lane_axis. Both are reported upstream with reproducers (PSAL-POSTECH/triton-npu#2), so this opens the path and the remaining blockers are on that side. No regression: add, transpose2D and test_triton_codegen still pass on the Triton route, add still passes on the MLIR route.

…he parts that break working kernels Enabling it by default cost six allowlisted tests. Two separate reasons, both from reaching for a bigger switch than the job needed. max_autotune also turns on pointwise autotuning, and benchmark_epilogue_fusion defaults to on: between them every fused kernel got a benchmark-flavoured rendering, whose harness imports arrive indented and made the module a SyntaxError at line 10. max_autotune_gemm is the flag that actually gates use_triton_template, and autotune_at_compile_time and benchmark_epilogue_fusion are now off explicitly. That fixed the syntax error and exposed the real cost: with templates on, addmm leaves a path that works today for one that stops at tl.assume in tnpu. Nothing to do here until that lands (PSAL-POSTECH/triton-npu#2 carries a reproducer), so the whole thing is opt-in through TORCHSIM_TRITON_TEMPLATES=1 -- on to work on it, and the default to make it the default. Allowlist back to 19/19.

…ay so Sending mm to aten is not a working state: the op runs in eager and is never simulated, so a test built on it was measuring nothing. The flag defaulted off to keep those tests green, which bought a number rather than a result -- TORCHSIM_TRITON_TEMPLATES=0 still opts out. With it on, extern_kernels disappears from all six and four of them stop at tl.assume instead. They come off the allowlist with that reason recorded in the file, so the gate keeps catching real breakage and this stays visible until PSAL-POSTECH/triton-npu#2 lands. Gate is 13/13, and now every one of them exercises the route -- the six that were counted before reached aten for the op they are named for.

Registering only mm_template left addmm and baddbmm resolving to it, and their bias is input_nodes[0], so def_kernel got three inputs for a two-argument template and asserted (2, 3, 0, 3). Eleven tests landed in the sweep's "other" bucket that way -- conv2d, mlp, transformer, the attention and fusion ones. Same shape as XPU's: an AddMMConfigMixin subclass registered under op_name addmm/baddbmm so prefix_args accounts for the bias. They now reach tl.assume like the rest, which is the real blocker (PSAL-POSTECH/triton-npu#2). Gate 13/13.

The first pass read 11/69 and buried the shape of the problem under a flat list of buckets. Two things changed underneath it. Coverage got stricter: a test that emits a kernel and still hands its own op to extern_kernels is no longer counted, which moved six out. And the route stopped taking that exit -- mm and addmm go through Inductor's template now, so sixteen tests newly reach tl.assume. They were not passing before; they were not simulating. What the numbers say now is narrower than a list of thirty-seven: tl.assume accounts for 16 and the select_lane_axis demand on a matmul result for 10, both upstream with reproducers. Three tests compile, run and return the wrong answer, which is the smallest bucket and the one worth reading first. Our own list is four tests long.

fixed_config_for left the reduction block unset and raised, on the grounds that the reduced axis has to stay inside a lane and the lane count would be the wrong value. That is the lowering pass's judgement, not codegen's. The only constraint here is that the block has to be a constant, because tnpu compiles one binary per kernel and there is no autotuner later. It now takes the whole reduced extent, so the kernel's r0 loop runs once -- the same shape Inductor bakes in for a persistent reduction. Whether that tile fits the lanes is for tnpu to answer, and it does: the five tests that stopped at SpecIncomplete now reach llvm.intr.assume or NoValidChoicesError, each of which names a real layer. Gate 13/13.

YWHyuk added 16 commits August 3, 2026 23:38

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Frontend] Unblock buffer reuse and triton_helpers on the Triton route - #308

[Frontend] Unblock buffer reuse and triton_helpers on the Triton route#308
YWHyuk wants to merge 16 commits into
feature/triton-codegenfrom
feature/triton-helpers

YWHyuk commented Aug 3, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

YWHyuk commented Aug 3, 2026

1. Enter memory planning through the pass upstream provides

2. Give the tnpu venv torch's triton_helpers, by copying it

Verification

Not fixed here

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

2. Give the tnpu venv torch's `triton_helpers`, by copying it