Skip to content

[Pallas] Fold tritonbench-bridge results into the tpu dashboard column#3082

Open
norx1991 wants to merge 2 commits into
mainfrom
yifeixu/tpu-bridge-dashboard
Open

[Pallas] Fold tritonbench-bridge results into the tpu dashboard column#3082
norx1991 wants to merge 2 commits into
mainfrom
yifeixu/tpu-bridge-dashboard

Conversation

@norx1991

@norx1991 norx1991 commented Jul 16, 2026

Copy link
Copy Markdown
Contributor

Folds the tritonbench-bridge TPU results (from #3028) into the existing tpu dashboard column, and turns the bridge on in the nightly cron — so the TPU dashboard's bridgeable kernels are measured through tritonbench, in one column alongside the run_tpu.py holdouts.

  • helion/_compat.pyget_device_name() returns "TPU v7" on TPU (was None"unknown"), matching run_tpu.py, so bridge rows carry a real device label.
  • build_dashboard_data.pyparse_run folds the tpu_bridge artifact into the tpu platform_short, and treats pallas as a native backend (no platform suffix, like triton) so pallas-on-tpu stays in the tpu column instead of a dropped tpu_pallas. The tpu platform's expected-kernels now unions kernels_tpu + kernels_tpu_bridge.
  • benchmark_dispatch.yml — trims kernels_tpu to the holdouts (bmm — no tritonbench operator; cross_entropy — Helion Mosaic-codegen holdout; flash_attention — the bridge's torch.compile comparator is flex_attention, which doesn't run on TPU, so the bridge would drop flash_attention's torch_compile column, whereas run_tpu.py compiles the SDPA baseline and keeps it); the other 7 kernels come from the bridge, so nothing is double-reported (the two sets are disjoint — 3 + 7). Consolidates the toggles: removes the transitional run_tpu_bridge input and gates both TPU jobs (run-tpu and run-tpu-bridge) on the single run_tpu input — so one switch runs all TPU benchmarks. No nightly change needed; the existing run_tpu: 'true' now triggers both.

Together with #3028, Helion's TPU nightly runs its 7 bridgeable kernels through tritonbench on the shared tpu dashboard column, with bmm/cross_entropy/flash_attention remaining on run_tpu.py. The bridge install recipe was validated end-to-end on the real TPU runner.

Verified end-to-end on the TPU runner via manual workflow_dispatch (PRs don't trigger the benchmark workflows, so this isn't visible in PR CI): a single run_tpu=true fired both run-tpu (run_tpu.py → bmm) and run-tpu-bridge (the bridge → gemm), and the two resulting artifacts fold into one tpu dashboard column — both kernels under tpu, labeled TPU v7, with no tpu_bridge/tpu_pallas split (run 29604064037). The dashboard parse/union logic is also covered by a local test against synthetic + real artifacts.

Measurement / dashboard-continuity notes

The 7 migrated kernels re-baseline. run_tpu.py measures each kernel on its own hand-written shapes (benchmarks/run_tpu.py::*_shapes); the bridge measures on tritonbench's operator shapes instead. These differ, so when the bridge goes live the migrated kernels' dashboard numbers shift in a one-time discontinuity — they are not comparable to the prior run_tpu.py values (different shapes → different FLOPs). This is intended: it aligns the TPU columns with the shapes the GPU dashboard already uses via the same tritonbench operators. We have not (and cannot meaningfully) verify run_tpu.py-vs-bridge numeric equivalence, since the shapes are different by construction.

Why flash_attention stays on run_tpu.py (beyond the compile-column reason): the bridge cannot measure it on TPU at all — its torch_compile comparator (flex_attention) doesn't run on TPU, and its sdpa baseline hits fused_sdp_choice(): no viable SDPBackend ... for TorchTPU (tritonbench forces SDPBackend.FLASH_ATTENTION, which TorchTPU rejects; run_tpu.py's plain scaled_dot_product_attention lets TorchTPU choose and works). Keeping it on run_tpu.py also preserves its existing shapes/numbers with no discontinuity.

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Jul 16, 2026
@norx1991
norx1991 force-pushed the yifeixu/tpu-bridge-dashboard branch from 71b01fc to 7da2241 Compare July 17, 2026 18:17
@norx1991
norx1991 marked this pull request as ready for review July 17, 2026 19:23
@norx1991 norx1991 changed the title [TPU] Fold tritonbench-bridge results into the tpu dashboard column [Pallas] Fold tritonbench-bridge results into the tpu dashboard column Jul 17, 2026
norx1991 added 2 commits July 21, 2026 13:00
…lumn

Makes the bridge nightly (benchmark_tpu_bridge.yml, #3028) feed the existing
`tpu` dashboard platform, and turns it on in the nightly cron:

- helion/_compat.py: get_device_name() returns "TPU v7" on TPU (was None ->
  "unknown"), matching run_tpu.py, so bridge rows carry a real device label.
- build_dashboard_data.py: fold the `tpu_bridge` artifact into the `tpu`
  platform_short, and treat `pallas` as a native backend (no platform suffix,
  like `triton`) so pallas-on-tpu stays in the `tpu` column instead of a dropped
  `tpu_pallas`. Expected-kernels for `tpu` now unions kernels_tpu +
  kernels_tpu_bridge.
- benchmark_dispatch.yml: trim kernels_tpu to the bridge's holdouts (bmm: no
  tritonbench operator; cross_entropy: Helion Mosaic-codegen holdout); the other
  8 kernels now come from the bridge, so nothing is double-reported.
- benchmark_nightly.yml: set run_tpu_bridge: 'true' so the bridge runs nightly.

Stacked on #3028. Net: the TPU dashboard's 8 bridgeable kernels are measured
through tritonbench, bmm + cross_entropy stay on run_tpu.py, all in one `tpu`
column.
…lumn)

The bridge's torch_compile comparator for flash_attention is flex_attention,
which does not run on TPU (its torch.compile HOP path doesn't lower on
torch_tpu) — verified: a bridge flash_attention run emits only helion_speedup,
no torch_compile_speedup, while run_tpu.py (which compiles the SDPA baseline)
emits torch_compile_speedup on TPU. Migrating flash_attention to the bridge
would silently drop its compile column, so keep it a run_tpu.py holdout
(kernels_tpu) alongside bmm/cross_entropy. The two kernel sets stay disjoint
(3 + 7), so the tpu column still covers the same 10 kernels with no dup.
@norx1991
norx1991 force-pushed the yifeixu/tpu-bridge-dashboard branch from 5a375b9 to 86dc68e Compare July 21, 2026 20:01
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot.

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants