[Pallas] Sample TPU bridge shapes equally-spaced (cap kl_div / rms_norm to feasible shapes)

[norx1991]

The TPU bridge sampled only --num-inputs 2 with the default first-k mode, so the dashboard showed just the two smallest shapes per kernel. This samples 6 shapes equally-spaced across each operator's range (matching the GPU nightly's --input-sample-mode equally-spaced-k), so small/mid/large are all represented.

Two kernels take the feasible low end via first-k instead of spanning to the top, because their largest sampled shape isn't runnable on TPU:

• kl_div: at fp32 its reduction dim reaches V=131072, which OOMs TPU vmem and autotunes for hours → first-k 4 (V≤32768).
• rms_norm / rms_norm-bwd: their largest shape (H=32768) sits ~66 KB over the 32 MB scoped-vmem ceiling, and helion's autotuner bails with "No working config found" → first-k 5 (H≤16384).

Validated via manual TPU-bridge dispatch: with these caps all bridged kernels run under the 6h job cap.

The rms_norm H=32768 cap is temporary. Root cause is a helion autotuner limitation (the autotune search abandons when its initial configs all miss the scoped-vmem ceiling by a hair, rather than finding a config that fits); a separate PR adds an autotune_baseline_fn for rms_norm (removes the earlier "Default config failed while computing baseline" crash, a prerequisite), and the autotuner search-robustness fix is tracked separately. Once that lands, the H≤16384 cap can be lifted.

[norx1991]

Welford is having some timeout issue in a trial dispatch. Need to hold on this before figuring it out.