[Pallas] Fix out-of-bound DMA caused by tiles from non-zero begins

[AmesingFlank]

Stacked PRs:

• [Pallas] Fix pre-broadcasting transformation bug when non-broadcast dims exceed PRE_BROADCAST_SIZE #2223
• A jagged_hstu_attention example that works on Pallas TPU #2218
• [Pallas] Fix DMA scratch buffer offset bug in nested fori_loop codegen #2217
• [Pallas] Cast bool masks to float before expanding in _mask_to codegen #2216
• [Pallas] Fix scratch Ref scoping bug in fori_loop/emit_pipeline codegen #2215
• [Pallas] Apply tile masks at load time to zero out-of-bounds data #2214
• ->[Pallas] Fix out-of-bound DMA caused by tiles from non-zero begins #2213

---

[Pallas] Fix out-of-bound DMA caused by tiles from non-zero begins

When pl.ds(offset, block_size) reads into a tensor, the last block
can overshoot the tensor boundary. The previous padding formula
(-shape) % block_size only accounted for rounding up to a block
boundary, but tiles from hl.tile(start, end) with non-zero start
can begin at arbitrary offsets, requiring additional headroom.

Pass an extra_pad value through _record_pad_info →
_compute_pad_info → _ds_pad_dims so the launcher pads by
(-shape) % block_size + extra_pad. The extra_pad is:

• 0 when the loop starts at offset 0
• begin % block_size for a provably constant begin
• block_size - 1 for a data-dependent begin

Co-Authored-By: Claude Opus 4.6 noreply@anthropic.com

[norx1991]

I am a bit surprised by this test case: does this create a pl.ds expr?
Also, if we are explicitly reading out of bound data in the helion kernel, shouldn't we throw an error? This test can pass because 3 is small.

[AmesingFlank]

Ah, I agree with you that hl.tile(3, 128 + 3) is a confusing test, because it seems to be explicitly requested OOB access. I updated the PR to use hl.tile(3, 128, block_size=16) instead. Here, the final tile still goes out-of-bound, but it is up to helion compiler to apply the correct padding/masking, which is what this PR is trying to fix.

does this create a pl.ds expr?

Yes, the tile is across (3, 128), so our compiler recognizes that this is tiling across the full dim of the tensor, so that dimension is not tiled via BlockSpec, and we are using pl.ds to do sliced access of each tile. This gist contains the full generated Pallas code, in case you are curious

[norx1991]

It seems to me that we can put the logic of extra_pad into _record_pad_info so we only need a single call here.

[AmesingFlank]

I just gave this a try,but it's not straightforward because _record_pad_info is called from three different contexts with different amounts of loop state available:

1. _ds_expr in codegen.py — the DeviceLoopState is registered in active_device_loops and LoopDimInfo.begin_expr is set, so
   _loop_begin_extra_pad works correctly here.
2. _make_block_spec in _codegen_emit_pipeline — called before the EmitPipelineLoopState is added to active_device_loops, and its
   LoopDimInfo doesn't set begin_expr.
3. _build_hbm_dma_slice in _codegen_fori_loop — same issue as (2).

For (2) and (3), the begin info only exists as codegen-level string expressions (begin_exprs) in the enclosing scope, not in LoopDimInfo. To make _record_pad_info self-contained, we'd need to either propagate begin_expr into the LoopDimInfo for emit_pipeline/fori_loop AND register the loop state earlier, or pass the begin info through a different channel — both add more complexity than the current approach. So I'd prefer to leave this as is

[norx1991]

I see. Thanks for giving it a try!