[pallas] Support dependent Tile.end bounds in Pallas - #3201
Merged
Conversation
ethche
requested review from
AmesingFlank,
norx1991 and
thcmbs
and removed request for
AmesingFlank,
norx1991 and
thcmbs
July 29, 2026 18:50
ethche
force-pushed
the
pallas-causal-tile-range-poc
branch
from
July 30, 2026 04:31
240acb0 to
9c46bc7
Compare
Contributor
|
For my own understanding, in our existing jagged hstu kernel, the end of the kv loop is fetched from memory ( |
Contributor
Author
|
@AmesingFlank For streaming fori_loop and emit_pipeline this is does not require much additional machinery. However, the new features are really
|
AmesingFlank
approved these changes
Jul 30, 2026
Follow-up cleanups to the dependent tile-end support: remove the ways the three inner-loop lowerings could drift, and drop a config knob that the resident lowering silently ignored. - Build LoopDimInfo through one _loop_dim_infos helper. The three lowerings each constructed it separately and the resident loop had fallen out of step, omitting end_expr. Setting it changes no generated code. - Allocate the resident loop's carry tuple through new_var. Two sibling bodies in one kernel both emitted a bare _carry; a nested one would have shadowed the enclosing carry and silently rebound reads of it. - Drop pallas_pre_broadcast when the loop type cannot apply it. The transform widens loop-carried VMEM scratch, which the resident lowering does not allocate, so both settings were autotuning as distinct configs that generate one kernel. - Spell out which bound forms _ordered_source_end accepts, and name the two other places that recognize the same shape and have to move with it: tracing_ops._dependent_tile_end_expr for kernels with no worklist plan, and tile_strategy._fold_tile_end_op for every backend. - Delete the worklist search gating tests: they assert behavior keyed on ConfigSpec.pallas_worklist_search_enabled, which does not exist. Replace the recognizer test with test_dependent_bound_grammar, covering accepted and rejected forms and both begin-matching rules. Generated code is unchanged across 42 kernel/config combinations spanning every loop type, grouping, and both dense and worklist paths, except for the intended carry rename.
thcmbs
approved these changes
Jul 31, 2026
norx1991
approved these changes
Jul 31, 2026
added 3 commits
July 31, 2026 11:43
_lift_sympy_arg returned the loop offset for the whole GridOrigin family,
but only a tile's begin is the offset. TileEndOrigin, TileCountOrigin and
TileIdOrigin each render as a compound expression, so a symbol carrying one
of those origins was silently replaced by the tile's begin.
This is reachable whenever such a symbol survives into a sympy expression
instead of its own op. min()/max() do exactly that: they are device
function replacements that fold their operands into a single sympy Min/Max
at trace time, dropping the tile_end node. So
for tile_k in hl.tile(0, min(n, tile_q.end)):
lowered its trip count from tile_q.begin -- zero iterations on the first
tile, and short by one block on every later one. tile.count collapsed the
same way (begin instead of cdiv), as did tile.id (begin instead of
offset // block). Plain arithmetic on a tile edge (tile.end + 8) keeps the
op and was never affected, nor was min() over tile.begin, which is what the
offset already meant.
Dispatch on the concrete origin instead. The block id is resolved before
rendering and carried into a replaced copy, because host_str() looks up
offset_var and active_device_loops by block id and an aliased symbol's own
block may have no active loop. The rendered edge is parenthesized: it
becomes a sympy Symbol name inside the enclosing expression, so an
unbracketed tile.id would turn 2 * (offset // block) into
(2 * offset) // block.
Covered for all three edges under both ops, on an input whose final tile is
partial. The test is Triton-only: Pallas lowers a dependent tile bound
through its own loop codegen and never reaches this path, so its generated
code is unchanged across 42 kernel/config combinations.
test_min_max_over_derived_tile_edge_keeps_its_own_formula only inspects generated code, so under HELION_INTERPRET=1 it fails twice over: ref eager exposes no tunable block sizes, making the declared block_sizes=[128, 128] an InvalidConfig, and RefEagerTestBase asserts every test either calls run_ref or skips. Mark it skipIfRefEager like the other codegen-only tests in this file.
Both tile a symbolic V, so under Pallas interpret they used to hit the backend-wide gate that rejected pallas_loop_type="unroll" whenever a kernel had symbolic or data-dependent bounds. "Support dependent Pallas tile ends" replaced that gate with a per-loop check, so these kernels now compile and pass in interpret mode and the xfail reports an unexpected success. Bisected across the branch: main and the mask-elision commit still xfail; the dependent-tile-end commit is where they start passing.
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
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.
Support Pallas inner tile loops whose end depends on an enclosing tile:
This PR supports both tile streaming loop types (foriloop, emit_pipeline) but also loop type = unroll. This also supports flatten_worklist (in the case where q_begin, q_end, kv_begin, kv_end) are also runtime dynamic.
Note that similar to the static range case, pallas_loop_type="unroll" will load the inner loop operands in VMEM once. But instead of a static python for loop, it will use a jax.lax.fori_loop.
Bound recognition
The dense path recognizes a direct enclosing
Tile.end, or:Recognition uses
TileEndOriginprovenance rather than attention-specific graph matching. Unsupported dynamic bounds continue to raiseInvalidConfigforpallas_loop_type="unroll".The previous backend-wide dynamic-bound rejection was moved to per-loop dispatch. This allows the supported dependent end while continuing to reject unrelated jagged or symbolic unroll loops when worklist grouping is disabled.
Unroll lowering
Static loops retain the existing Python-unrolled lowering. A supported dependent end instead emits a resident, value-carried loop:
This reuses the existing unroll tensor placement and offset/index/mask helpers. Dense operands remain resident through their outer Pallas BlockSpecs. Flattened worklists retain their existing range-keyed resident windows.
Unlike streaming
fori_loop, this path does not allocate per-tile DMA resources, semaphores, or VMEM scratch for loop-carried values. There is no resident-to-streaming fallback: configurations that cannot fit the resident working set fail explicitly.Streaming lowering
pallas_loop_type="fori_loop"keeps its existing explicit per-tile DMA lowering. Only its dynamic trip count and offsets use the dependent end.pallas_loop_type="emit_pipeline"keeps one buffered pipeline. Its dynamic grid and BlockSpec index maps use the dependent end, preserving the existing DMA/compute overlap.The change also records host padding for partial outer tiles used by dense streaming DMA and BlockSpecs.
Flattened worklists
Flattened worklists distinguish the resident source range from the active compute range:
compute_endcontrols the loop trip count. The fullrange_lencontinues to control:This allows successive Q work items to reuse the same resident K/V range while executing progressively different causal compute extents.
Grouping 2 generates its base-block and double-block body variants as before. Resident prepared-cache allocation and refill are shared between those variants and emitted only once.
Scalar minimum
Builtin
min()now accepts mixed integer, SymInt, and 0-D tensor arguments. This supports bounds such asmin(kv_end, tile_q.end)without adding tensor-awaremax; dependent lower bounds remain deferred to sliding-window support.Results: jagged causal self-attn — B=1152, H=4, D=128, blocks 512×512, grouping=1, unroll
(TF/s = causal useful; gdpa verified correct rel 2.87e-3. Pruning wins under the default unroll here because the jagged masked baseline is inherently a dynamic fori — no static-unroll advantage to offset it.)
Results: Dense causal flash — B=8, H=32, D=256, S=4096, emit_pipeline