[DO NOT REVIEW] [Pallas] Add TPU nightly benchmark workflow and runner

[norx1991]

Caution

The push trigger in benchmark_tpu_nightly.yml is temporary for CI testing — remove before merging.

Summary

Add a nightly CI workflow that runs Helion examples with autotuning on TPU, with results published to pytorch benchmark hub. This mirrors the GPU benchmark infrastructure (benchmark_nightly.yml + benchmarks/run.py) but uses a standalone runner since TritonBench has hard dependencies on triton and CUDA at import/install/runtime levels.

New files

• benchmarks/run_tpu.py — TPU benchmark runner

  • CLI: --kernel/--op (comma-separated), --output (JSON), --list-kernels
  • Multi-shape benchmarking: each kernel tested at multiple input sizes (like TritonBench does for GPU)
  • Per-shape accuracy check + timing vs torch baseline with speedup display
  • JSON output in pytorch benchmark hub format
  • 5 reliable kernels: exp, add, softmax_two_pass, welford, layer_norm

• .github/workflows/benchmark_tpu.yml — Reusable benchmark workflow

  • Runner: linux.google.tpuv7x.1
  • Setup: PyTorch CPU nightly, JAX/Pallas, builds torch_tpu from pinned commit
  • Two-pass pattern: autotune → sleep 1min → HELION_ASSERT_CACHE_HIT=1 verify + record
  • Upload to pytorch benchmark hub

• .github/workflows/benchmark_tpu_nightly.yml — Nightly trigger

  • Cron: daily at 2 AM PST (10 AM UTC)
  • workflow_dispatch with kernels input for manual runs

Modified files

• helion/_testing.py — run_example() return type changed from None to dict[str, float] (mapping implementation names to their benchmark times in ms), enabling run_tpu.py to capture and display per-shape speedup without duplicating timing logic

Excluded kernels (with reasons)

• rms_norm: InductorLoweringError in torch.mean reduction codegen for fori_loop/emit_pipeline
• geglu/swiglu: autotuning takes >15min per kernel (large shape 8x2048x4096), many configs fail to compile
• low_mem_dropout: ~37% element accuracy mismatch on all configs except block_sizes=[128]

Example output

===========================================================================
Summary
===========================================================================
Kernel                 Shape            Status   Helion (ms)    Torch (ms)     Speedup
---------------------------------------------------------------------------
exp                    [1024]           PASS     0.3410         0.2676         0.78x
exp                    [4096]           PASS     0.2829         0.2320         0.82x
exp                    [16384]          PASS     0.2924         0.2276         0.78x
exp                    [65536]          PASS     0.2744         0.2173         0.79x
exp                    [262144]         PASS     0.3246         0.2355         0.73x
exp                    [1048576]        PASS     0.3076         0.2134         0.69x
add                    [128,128]        PASS     0.2688         0.2087         0.78x
add                    [256,256]        PASS     0.2727         0.2149         0.79x
add                    [512,512]        PASS     0.3202         0.2259         0.71x
add                    [1024,1024]      PASS     0.3165         0.2225         0.70x
add                    [2048,2048]      PASS     0.3604         0.2440         0.68x
softmax_two_pass       [1024,256]       PASS     0.3279         0.2375         0.72x
softmax_two_pass       [1024,512]       PASS     0.3701         0.2371         0.64x
softmax_two_pass       [1024,1024]      PASS     0.3610         0.2391         0.66x
softmax_two_pass       [1024,2048]      PASS     0.3112         0.2195         0.71x
softmax_two_pass       [1024,4096]      PASS     0.3373         0.2317         0.69x
---------------------------------------------------------------------------
Total: 3/3 passed
===========================================================================

Test plan

• Tested run_tpu.py on TPU pod with all 5 default kernels
• Verified per-shape accuracy and speedup display
• Verified run_example() return value is correctly consumed
• CI run triggered via temporary push trigger on this branch
• Remove push trigger before merging