[Pallas] Defer tile load masks past transposes onto the sublane axis

[ethche]

A Pallas tiled load masks out-of-bounds lanes multiplicatively in the load's own layout (ref[idx] * mask). When the value is only relayouted by an axis permutation (transpose / .T) before a dot or reduction, that consumer already inserts _mask_to(x, 0), which re-materializes the mask in the consumer layout.

In general, we see that masking on the sublane axis is much cheaper, while masking on a major axis adds a lot of vector load overhead. This informs our optimization: we defer loads only when the masked axis is a major/outer dim at the load (eager mask is expensive there) and the permutation moves it into the last-two (sublane/lane) dims at the consumer (mask is cheap there). We don't defer for other types of ops (view reshape etc). If the permutation moves the mask from sublane/lane to major axis, deferring can make things slower, so this gate is important.

Before:

  mask_1 = offset < tile_start + tile_extent
  load_4 = q[pl.ds(0, BLOCK), :, :] * mask_1.astype(jnp.bfloat16)[:, None, None]
  qbk    = jnp.transpose(load_4, [1, 0, 2])
  scores = lax.dot_general(qbk, k_t, ...)

After:

  mask_1  = offset < tile_start + tile_extent
  load_4  = q[pl.ds(0, BLOCK), :, :]
  qbk     = jnp.transpose(load_4, [1, 0, 2])                      
  _mask_to = jnp.where(mask_1.astype(jnp.float32)[None, :, None], qbk, 0)   # mask in consumer layout
  scores  = lax.dot_general(_mask_to, k_t, ...)

Performance on jagged gdpa dense kv (bf16). same shapes as in:#2782

shape	dist	block	baseline TFLOP/s	with fix TFLOP/s	speedup
A (B=256)	uniform	q4096	200	257	1.29×
A	ramp	q1024	142	158	1.11×
A	random	q4096	150	189	1.27×
B (B=768)	uniform	q4096	213	281	1.32×
B	ramp	q1024	156	176	1.13×
B	random	q4096	159	210	1.32×

New FX graph pass defer_pallas_load_masks (Pallas-only, run just before remove_unnecessary_masking). For each load whose masked tile dim is provably re-masked downstream, it records the deferred block ids and marks the load's masked value unknown so the _mask_to survives; load codegen then skips the eager mask for those dims and the mask is materialized after the transpose.

Correctness rests on a dataflow proof: every use of the load reaches a mask_to(, 0) crossing only pure axis-permutation ops (aten.permute.default), with the masked dim surviving as a unique tile dim at each step. Only gated to pallas + major to sublane/lane permutes.

[meta-codesync]

@ethche has imported this pull request. If you are a Meta employee, you can view this in D109034842.

[AmesingFlank]

lgtm. Altho, do you expect this rewrite to always be profitable? Or do you think there is value is making this tunable?

[ethche]

lgtm. Altho, do you expect this rewrite to always be profitable? Or do you think there is value is making this tunable?

The gating is pretty conservative so I think it would only fire when it's reasonable. But I think it would be good for us to keep an eye out on whether there's could be part of a more general optimization pass (e.g. re-ordering operations).