Fold gather(load(t, [..., :, ...]), dim, idx) into direct indirect load#2684
Open
AmesingFlank wants to merge 1 commit into
Open
Fold gather(load(t, [..., :, ...]), dim, idx) into direct indirect load#2684AmesingFlank wants to merge 1 commit into
AmesingFlank wants to merge 1 commit into
Conversation
AmesingFlank
force-pushed
the
AmesingFlank/stack/63
branch
from
1b36076 to
b70396c
Compare
This was referenced Jun 4, 2026
[Pallas] Lower aten gather using one_hot + sum for TPU compatibility, unblocking cross_entropy
#2060
Open
meta-cla
Bot
added
the
CLA Signed
AmesingFlank
added a commit
that referenced
this pull request
Jun 4, 2026
The cross_entropy pattern (logits[tile_n, :].gather(1, idx[tile_n].unsqueeze(1))) was producing invalid Triton (NameError on the load) when the reduction roller tried to roll the surrounding amax/sum: a _for_loop output can't carry the rdim-shaped logits_rows out to feed the gather sitting outside the loop. Rewrite gather(load(t, [..., :, ...]), dim, idx) at the FX layer to a direct indirect load(t, [..., idx, ...]). The two forms compute the same values, but the direct form skips the wide load entirely — so the rdim-shaped intermediate never exists and the roller's existing logic handles the surrounding reductions naturally. The CuTe backend already does this fold at codegen time (aten_lowering.codegen_gather_cute); lifting it to FX surfaces the same simplification to the Triton backend and the rolling analysis. The fold is gated to the cross_entropy-style pattern: load's dim axis is a full slice, gather index has a singleton at dim and the same rank as the load's subscript, no extra_mask. Other gather shapes go through the existing aten.gather path. After this, examples/cross_entropy.py runs end-to-end: autotuning finds rolled configs (block_sizes=[1], reduction_loops=[16384]) and the kernel is ~3x faster than torch eager. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> stack-info: PR: #2684, branch: AmesingFlank/stack/63
AmesingFlank
force-pushed
the
AmesingFlank/stack/63
branch
from
b70396c to
674c822
Compare
AmesingFlank
added a commit
that referenced
this pull request
Jun 4, 2026
The cross_entropy pattern (logits[tile_n, :].gather(1, idx[tile_n].unsqueeze(1))) was producing invalid Triton (NameError on the load) when the reduction roller tried to roll the surrounding amax/sum: a _for_loop output can't carry the rdim-shaped logits_rows out to feed the gather sitting outside the loop. Rewrite gather(load(t, [..., :, ...]), dim, idx) at the FX layer to a direct indirect load(t, [..., idx, ...]). The two forms compute the same values, but the direct form skips the wide load entirely — so the rdim-shaped intermediate never exists and the roller's existing logic handles the surrounding reductions naturally. The CuTe backend already does this fold at codegen time (aten_lowering.codegen_gather_cute); lifting it to FX surfaces the same simplification to the Triton backend and the rolling analysis. The fold is gated to the cross_entropy-style pattern: load's dim axis is a full slice, gather index has a singleton at dim and the same rank as the load's subscript, no extra_mask. Other gather shapes go through the existing aten.gather path. After this, examples/cross_entropy.py runs end-to-end: autotuning finds rolled configs (block_sizes=[1], reduction_loops=[16384]) and the kernel is ~3x faster than torch eager. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> stack-info: PR: #2684, branch: AmesingFlank/stack/63
AmesingFlank
force-pushed
the
AmesingFlank/stack/63
branch
from
674c822 to
190a1a3
Compare
jansel
approved these changes
Jun 4, 2026
choijon5
approved these changes
Jun 4, 2026
The cross_entropy pattern (logits[tile_n, :].gather(1, idx[tile_n].unsqueeze(1))) was producing invalid Triton (NameError on the load) when the reduction roller tried to roll the surrounding amax/sum: a _for_loop output can't carry the rdim-shaped logits_rows out to feed the gather sitting outside the loop. Rewrite gather(load(t, [..., :, ...]), dim, idx) at the FX layer to a direct indirect load(t, [..., idx, ...]). The two forms compute the same values, but the direct form skips the wide load entirely — so the rdim-shaped intermediate never exists and the roller's existing logic handles the surrounding reductions naturally. The CuTe backend already does this fold at codegen time (aten_lowering.codegen_gather_cute); lifting it to FX surfaces the same simplification to the Triton backend and the rolling analysis. The fold is gated to the cross_entropy-style pattern: load's dim axis is a full slice, gather index has a singleton at dim and the same rank as the load's subscript, no extra_mask. Other gather shapes go through the existing aten.gather path. After this, examples/cross_entropy.py runs end-to-end: autotuning finds rolled configs (block_sizes=[1], reduction_loops=[16384]) and the kernel is ~3x faster than torch eager. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com> stack-info: PR: #2684, branch: AmesingFlank/stack/63
AmesingFlank
force-pushed
the
AmesingFlank/stack/63
branch
from
190a1a3 to
13ba6b2
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
3 participants
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.
Stacked PRs:
Fold gather(load(t, [..., :, ...]), dim, idx) into direct indirect load
The cross_entropy pattern (logits[tile_n, :].gather(1, idx[tile_n].unsqueeze(1)))
was producing invalid Triton (NameError on the load) when the reduction roller
tried to roll the surrounding amax/sum: a _for_loop output can't carry the
rdim-shaped logits_rows out to feed the gather sitting outside the loop.
Rewrite gather(load(t, [..., :, ...]), dim, idx) at the FX layer to a direct
indirect load(t, [..., idx, ...]). The two forms compute the same values, but
the direct form skips the wide load entirely — so the rdim-shaped intermediate
never exists and the roller's existing logic handles the surrounding reductions
naturally. The CuTe backend already does this fold at codegen time
(aten_lowering.codegen_gather_cute); lifting it to FX surfaces the same
simplification to the Triton backend and the rolling analysis.
The fold is gated to the cross_entropy-style pattern: load's dim axis is a
full slice, gather index has a singleton at dim and the same rank as the
load's subscript, no extra_mask. Other gather shapes go through the existing
aten.gather path.
After this, examples/cross_entropy.py runs end-to-end: autotuning finds rolled
configs (block_sizes=[1], reduction_loops=[16384]) and the kernel is ~3x
faster than torch eager.
Co-Authored-By: Claude Opus 4.7 (1M context) noreply@anthropic.com