[pallas] Support dependent Tile.end bounds in Pallas

[ethche]

Support Pallas inner tile loops whose end depends on an enclosing tile:

for tile_q in hl.tile(q_begin,q_end):                                                                                                                                                                                         
  for tile_k in hl.tile(kv_begin, min(kv_end, tile_q.end)):                                                                                                                                      

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:

min(host_evaluable_source_end, enclosing_tile.end)

Recognition uses TileEndOrigin provenance rather than attention-specific graph matching. Unsupported dynamic bounds continue to raise InvalidConfig for pallas_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:

trip_count = cdiv(active_end - begin, block_k)

def dynamic_unroll_body(tile_id, carry):
    offset_k = begin + tile_id * block_k
    indices_k = offset_k + jnp.arange(block_k)
    mask_k = indices_k < active_end
    return original_body(carry)

result = jax.lax.fori_loop(
    0,
    trip_count,
    dynamic_unroll_body,
    initial_carry,
)

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:

ordered_begin = range_start_ref[wid]
source_end = ordered_begin + range_len_ref[wid]
compact_end = tile_start_ref[wid] + tile_extent_ref[wid]
compute_end = jnp.minimum(source_end, compact_end)

compute_end controls the loop trip count. The full range_len continues to control:

• Resident-window sizing and overflow validation
• Resident cache keys and refills
• Prepared-cache full and tail refills

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 as min(kv_end, tile_q.end) without adding tensor-aware max; 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

kernel	before (masked)	after (pruned)	speedup
gdpa	76.0	118.5	1.56×
hstu	81.9	126.1	1.54×
flash	63.5	97.7	1.54×

(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

block	before (masked)	after (pruned)	speedup
512×512 (bb2)	85.4	131.9	1.54×
1024×1024 (bb2)	158.8	225.8	1.42×
2048×2048 (bb1)	144.2	174.9	1.21×

[AmesingFlank]

For my own understanding, in our existing jagged hstu kernel, the end of the kv loop is fetched from memory (seq_offsets[seq_idx + 1]), which is already a dynamic value. What makes min(seq_offsets[seq_idx + 1], tile_q.end) "more dynamic" than just seq_offsets[seq_idx + 1] that causes existing lowering to fail for it?

[ethche]

@AmesingFlank For streaming fori_loop and emit_pipeline this is does not require much additional machinery. However, the new features are really

• Support for loop unrolling. Previously pallas_loop_type = unroll would just raise an invalid config error because the loop extent is dynamic. This PR now supports it.
• Support for flatten worklist. We have to adjust the flatten worklist lowering as the num blocks arithmetic is different when we have a custom end point. We are tracing the end point from the offset, so we need extra logic to encode additional operations (e.g. taking the min with a runtime jagged extent).