From d763448228cd962c9221574fd1e8a44f3bc1bdf4 Mon Sep 17 00:00:00 2001 From: Yifei Xu Date: Wed, 29 Jul 2026 17:58:27 -0700 Subject: [PATCH] [Pallas] Run kl_div and welford at fp32 on the TPU bridge to match the examples The bridge forced --precision bf16 for every kernel, but the canonical helion examples (which run_tpu.py mirrors) run kl_div and welford at fp32: kl_div keeps its loss reduction in fp32, and welford's bf16 layer-norm output drifts on near-zero values and trips the 1e-2 accuracy tolerance for every correct impl. Run those two at fp32 so each is compared against a same-precision baseline; the rest stay bf16 (matching their examples). welford's fp32 only takes effect once the tritonbench pin includes the operator change that makes its input loader honor --precision. --- .github/workflows/benchmark_tpu_bridge.yml | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/.github/workflows/benchmark_tpu_bridge.yml b/.github/workflows/benchmark_tpu_bridge.yml index 5383a11c3..e9953f49e 100644 --- a/.github/workflows/benchmark_tpu_bridge.yml +++ b/.github/workflows/benchmark_tpu_bridge.yml @@ -111,18 +111,31 @@ jobs: continue fi - # TPU differences vs GPU (benchmark.yml): --device tpu, --precision - # bf16 (match run_tpu), drop --cudagraph and + # TPU differences vs GPU (benchmark.yml): --device tpu, per-kernel + # --precision (see below), drop --cudagraph and # --latency-measure-mode triton_do_bench (CUDA/Triton-only; TPU uses # the wall-clock _do_bench_tpu path). Shapes are tritonbench's own, # count-capped by --num-inputs; per-kernel shape args (e.g. embedding # --B/--T/--D/--v-range) are supplied via inputs.env-vars / custom # args where a default shape exceeds TPU vmem. + # + # Precision follows the canonical helion examples per kernel: bf16 + # for gemm/softmax/layer_norm/rms_norm (examples/rms_norm.py runs + # both fwd and bwd at bf16), and fp32 for kl_div and welford + # (examples/kl_div.py and examples/welford.py use fp32 inputs). + # kl_div keeps its loss reduction in fp32; welford's bf16 layer-norm + # output drifts ~1 ULP on near-zero values and trips the 1e-2 + # accuracy tolerance for every correct impl. Force fp32 for those two + # so each impl isn't compared against a downcast bf16 baseline. + PRECISION=bf16 + case "$kernel" in + kl_div|welford) PRECISION=fp32 ;; + esac ${{ inputs.env-vars }} HELION_PRINT_OUTPUT_CODE=1 HELION_MEASURE_COMPILE_TIME=1 \ python benchmarks/run.py \ --op $kernel \ --device tpu \ - --precision bf16 \ + --precision $PRECISION \ --helion-backend pallas \ --metrics speedup,accuracy,latency \ --measure-compile-time \