[Pallas] Fix pre-broadcasting transformation bug when non-broadcast dims exceed PRE_BROADCAST_SIZE

[AmesingFlank]

Stacked PRs:

• A jagged_hstu_attention example that works on Pallas TPU #2218
• ->[Pallas] Fix pre-broadcasting transformation bug when non-broadcast dims exceed PRE_BROADCAST_SIZE #2223

---

[Pallas] Fix pre-broadcasting transformation bug when non-broadcast dims exceed PRE_BROADCAST_SIZE

The pre-broadcast optimization pass (added in #2103) appends a trailing
PRE_BROADCAST_SIZE=128 dimension to loop-carried scratch buffers like
m_i and l_i to avoid costly implicit broadcasts on TPU. When
head_dim equals 128, this works correctly because the pre-broadcast
dimension matches head_dim. However, when head_dim > 128 (e.g. 256),
the outer graph produced after the inner loop has a shape mismatch.

Repro: Run the attention kernel from dunfanlu_notes/tpu/attn/attn_helion.py
with D=256 and pallas_pre_broadcast=True.

Error: helion.exc.ShapeMismatch: Shape mismatch between [u0, u1, 256] and [u0, u1, 128] at acc = acc / l_i[:, :, None] in the outer graph.

Root cause: _rewrite_outer_subscripts_for_pre_broadcast correctly rewrites
subscript[:, :, None] to identity slicing [:, :] for pre-broadcast
results exiting the inner loop, and updates their meta shapes to include
the trailing 128 dim. But it never inserted _pre_broadcast_tile nodes to
expand the 128-wide values to match wider-dim consumers (e.g. acc with
head_dim=256). The inner graph handled this via Step 3 of
_annotate_pre_broadcast, but the outer graph lacked the equivalent logic.

Fix: Add a tile-insertion pass to _rewrite_outer_subscripts_for_pre_broadcast
that mirrors the inner graph's Step 3. It scans non-pre-broadcast nodes in
the outer graph whose last dim exceeds 128, checks if any of their args are
pre-broadcast nodes with trailing dim 128, and inserts _pre_broadcast_tile
to expand them (e.g. jnp.tile(subscript_2, 2) for 128→256).

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

[norx1991]

Do we cover indirect consumer? I found this test case with help of claude:

@helion.kernel(backend="pallas", static_shapes=True)
  def outer_chain_scale(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
      batch, m, k = a.size()
      head_dim = hl.specialize(b.size(-1))
      out = torch.empty([batch, m, head_dim], device=a.device, dtype=a.dtype)
      for tile_b, tile_m in hl.tile([batch, m]):
          running = hl.zeros([tile_b, tile_m], dtype=torch.float32)
          acc = hl.zeros([tile_b, tile_m, head_dim], dtype=torch.float32)
          for tile_k in hl.tile(k):
              chunk = a[tile_b, tile_m, tile_k]
              running = running + torch.sum(chunk, -1)
              acc = acc + running[:, :, None]
          scale = torch.rsqrt(running[:, :, None] + 1.0)
          out[tile_b, tile_m, :] = (acc * scale).to(out.dtype)
      return out

[AmesingFlank]

good catch! Update the PR to fix this and added a test for this case