[Frontend] Unblock buffer reuse and triton_helpers on the Triton route - #308
Open
YWHyuk wants to merge 10 commits into
Open
[Frontend] Unblock buffer reuse and triton_helpers on the Triton route#308YWHyuk wants to merge 10 commits into
YWHyuk wants to merge 10 commits into
Conversation
…ton 3.6 triton-npu is built against triton 3.6. 2.10 is the first torch release whose Inductor targets that version, so anything older emits kernels its passes were not written for. The pins live here rather than in a requirements file because the base image tag is derived from this file's hash: editing it moves the tag, and nothing else has to be kept in step by hand.
The toolchain layer is about 1.8 GiB -- LLVM, Spike, the triton runtime, the triton-shared binary -- which no other job needs, so it gets its own image and its own workflow instead of growing the one every test pulls. torchsim_tnpu_base is pinned by thirdparty/triton-npu.json plus Dockerfile.tnpu, the same derived-tag scheme the base image uses. Building it needs secrets.TNPU_TOKEN and a toolchain release on PSAL-POSTECH/triton-npu, both private; the token wants Contents: Read-only and must not be a push token, since this tree ships inside the image. The pin is 8ebe408, which drops llvm.intr.assume in normalize_upstream and moves triton_shared to 9017bd4 for the llvm dialect registration. Both halves are needed: without registration the op stops stage 2 at the parse. Inductor's mm template emits tl.assume(pid >= 0) unconditionally, so every matmul from that template carried it.
A second path from an fx graph to a simulated kernel: Inductor's own Triton backend emits the kernel, and triton-npu lowers it, instead of the MLIR templates doing both. It is selected by TORCHSIM_TRITON_CODEGEN, read once at device registration, so no test knows which route it is on. kernel_spec captures what tnpu needs while V.graph is still live -- there is no second chance once codegen returns -- and turns it into a tnpu KernelSpec. Several things have to be carried that Inductor never writes down for a device that benchmarks its own kernels: The grid. kernel.numels is keyed by iteration-space prefix, and the axes go outermost-first for tile decisions but x-first for the pid order tnpu reads. Getting either wrong silently produces a launch that matches nothing. The strides. Inductor allocates outputs with empty_strided and indexes them by stride, so a launch assuming contiguous writes elements to the wrong places -- and the wrongness is quiet. The scalars, which reach the tnpu wrapper as typed C arguments. helpers_shim copies torch's triton_helpers into the tnpu venv, since the kernel Inductor writes imports it. _triton_compat strips what only a GPU can honour: num_warps, num_stages, tl.debug_barrier, set_driver_to_gpu. The strip has to keep libdevice and tl_math, which used to go out with the torch import that named them. functional writes and reads tensors in storage order via as_strided, so a non-contiguous input round-trips through a .raw file unchanged. A kernel that cannot be described is rejected here with the missing field named, rather than written out to fail somewhere the cause is unrecoverable.
The lowered kernel still has to become a tile-operation graph with a cycle cost per tile, the same as the MLIR route's output. timing drives gem5 over a body reduced to one work-item, then hands TOGSim the trace. The grid is taken at run time rather than baked in, so one trace serves every shape the kernel is launched with. Two fixes the shared passes needed. generate() went through memory_plan_reuse directly, which skips the wrapper IR passes estimate_peak depends on; it now enters through run_wrapper_ir_passes, the entry point upstream provides. And both wrap_kernel_call sites stringify their arguments, which were being formatted as objects.
…rnel mm and conv went to extern_kernels, where on npu they either raise convolution_overrideable not implemented or fall back to eager and simulate nothing. Passing was not the same as working. The templates themselves are not GPU-specific: torch ships one triton_mm.py.jinja for cuda, xpu, mtia and cpu, and the per-backend difference is a config table. What gates them is use_triton_template asking is_gpu, and GPU_TYPES is a hardcoded list with no registration hook, so npu is appended to it here. Selection is a fixed choice, not a benchmark: there is no device to time on. The TODO is to rank by simulated cycles, which timing.run_togsim already returns per kernel. Until then the offered order wins, which is deterministic and not a claim about speed. max_autotune_gemm rather than max_autotune, and epilogue-fusion benchmarking off: both render a benchmark-flavoured kernel whose harness imports arrive indented in the real module, which broke kernels that already worked. addmm and baddbmm carry a bias as input_nodes[0] and need their own heuristic; without one the mm entry is used with prefix_args=0 and def_kernel asserts. device_guard returned "pass", which the caller writes as `with pass:`.
TORCHSIM_TRITON_CODEGEN is read at device registration, so the suite can be run on the other route without any test knowing. The sweep does that and produces three things: a gate (triton_route_passing.txt), a report bucketed by cause and by the tnpu stage a kernel last cleared, and per-failure artifacts -- the generated kernel, the stage IRs, the error -- so a problem can be reported upstream with the code rather than a description of it. Emitting a kernel is not enough to call the route exercised. A test can generate a kernel for part of its graph and still send the op it is named for to aten: test_matmul_scalar generated the mul and called extern_kernels.mm. Those are reported with the op that took the exit and kept out of the gate. Tests run in parallel, each with its own dump and TORCHINDUCTOR_CACHE_DIR; sharing the cache made results depend on what ran first.
…irst op test_triton_codegen exercises the route directly, so a break in it is not only visible as a change in the sweep. test_sparse_core passed device positionally into a helper that takes it by keyword. test_mlp_cpu imported torchsummary, matplotlib and numpy, none of which it uses: the one summary() call is commented out. Importing a package the repo does not depend on was enough to stop it. It then loaded ./128_784_32_10_cpu/best_model.pth, which is neither the name computed ten lines above nor anything in the repo -- a checkpoint from whoever ran it last. The test trains and saves its own further down, so the load is a warm start, not a requirement.
YWHyuk
force-pushed
the
feature/triton-helpers
branch
from
August 5, 2026 08:51
3bcf2aa to
aebeb9c
Compare
README.md is the setup, the seven tnpu stages, and what each module here owes. It is referenced from the workflow, extension_config, the device registration and the route test. CLAUDE.md gains the triton_npu workflow and the torch 2.10 pin.
Two ways the choice list came back empty, both raising NoValidChoicesError before any kernel was emitted. A zero-length axis has no tile, so the heuristics offer no config at all. A MoE expert that routes no tokens produces exactly that: [0, K] @ [K, N]. Those are short-circuited to zeros without going near the selection machinery -- correct whether the output is empty (M or N zero) or a sum over nothing (K zero). addmm and baddbmm with K == 0 are beta * bias rather than zeros, so they keep the original lowering unless the output is empty. The second was ours. max_autotune_gemm_backends is global but the heuristics are registered for npu only, so an ordinary cpu gemm in the same graph -- test_moe_cpu runs [32, 1000] @ [1000, 10] -- found no backend it could use. ATEN goes back in the list and pick_config ranks it last, so npu still takes a template wherever one exists. Measured over the whole suite afterwards: no test reaches an extern kernel. tnpu delivery goes from 60 to 61 of 69.
…e now The prose had grown past what it was explaining. Removed: the account of a torch 2.8 API skew that no longer exists, a comparison against the MLIR route's line count, section banners in capitals, and the sentences arguing that a decision was the right one rather than saying what it does. Two comments were also stale. __init__ called the route scaffolding and said to expect failures, not results; fourteen tests now pass through it end to end. And inductor_templates pointed at tl.assume as the thing tests stop on, which was fixed upstream in triton-npu 8ebe408. Kept the pipeline diagram, the file and line references, and the TODOs. No code changed.
YWHyuk
force-pushed
the
feature/triton-helpers
branch
from
August 5, 2026 10:48
ee1026c to
fc192b3
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 ondevelopyet, so this targetsfeature/triton-codegen.1. Enter memory planning through the pass upstream provides
generate()calledmemory_plan_reuse()directly, bypassingrun_wrapper_ir_passes()— the method that chooses a planner and sets up thestate that planner needs. Three things followed:
self.estimate_peak, created inrun_wrapper_ir_passesandread by
should_reuse_bufferduring planning. Every graph that reached bufferreuse died with
AttributeError. In CI that is 16 of 69 tests on theTriton route — Llama, CLIP, ConvNeXtV2, MobileNet, Mixtral, and every
attention test.
memory_plan()override was unreachable:run_wrapper_ir_passesisits only caller.
generate()tookis_inferenceand never used it.Calling the upstream entry point fixes all three. Behaviour is identical today:
config.memory_planningdefaults toFalse, so upstream takes the same branchwe 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 itInductor'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.pyimports nothing fromtorch — only
.triton_compat, and that module touches torch in three places(
torch.version.hip×2,torch.autograd.profiler×1), none of whichtriton_helpersneeds. So the installed torch's file is copied verbatim besidethe kernel and paired with a small
triton_compatthat resolves the same sevennames 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 bypath, so a sibling package would not otherwise be importable from it.
Verification
tests/ops/elementwise/test_activation.py: the three relu kernels usetriton_helpers.maximumand now reach an ELF and a trace producer, where beforethey 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 notrun either, because the test aborts at its first failure and relu is first.
tests/ops/attention/test_gqa.pygets pastestimate_peakto its next blocker.MLIR route unaffected:
test_add,test_matmul,test_prologue_fusion,test_softmaxall still pass.Gate (
scripts/ci/triton_route_passing.txt) still 11/11.Not fixed here
Both fixes now converge on the same next blocker:
libdeviceexternintrinsics (
exp,tanh,rsqrt,erf,isnan), which have notriton_shared implementation. Worth noting that
tl.sigmoidlowers tomath.expand works, so a substitution may be cheaper than a lowering pass —that is the next thing to measure, not part of this PR.