[Pallas] Fold tritonbench-bridge results into the tpu dashboard column#3082
Open
norx1991 wants to merge 2 commits into
Open
[Pallas] Fold tritonbench-bridge results into the tpu dashboard column#3082norx1991 wants to merge 2 commits into
norx1991 wants to merge 2 commits into
Conversation
norx1991
force-pushed
the
yifeixu/tpu-bridge-dashboard
branch
from
July 17, 2026 18:17
71b01fc to
7da2241
Compare
norx1991
marked this pull request as ready for review
July 17, 2026 19:23
jansel
approved these changes
Jul 20, 2026
…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
force-pushed
the
yifeixu/tpu-bridge-dashboard
branch
from
July 21, 2026 20:01
5a375b9 to
86dc68e
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Folds the tritonbench-bridge TPU results (from #3028) into the existing
tpudashboard 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.py—get_device_name()returns"TPU v7"on TPU (wasNone→"unknown"), matchingrun_tpu.py, so bridge rows carry a real device label.build_dashboard_data.py—parse_runfolds thetpu_bridgeartifact into thetpuplatform_short, and treatspallasas a native backend (no platform suffix, liketriton) so pallas-on-tpu stays in thetpucolumn instead of a droppedtpu_pallas. Thetpuplatform's expected-kernels now unionskernels_tpu+kernels_tpu_bridge.benchmark_dispatch.yml— trimskernels_tputo the holdouts (bmm— no tritonbench operator;cross_entropy— Helion Mosaic-codegen holdout;flash_attention— the bridge's torch.compile comparator isflex_attention, which doesn't run on TPU, so the bridge would drop flash_attention'storch_compilecolumn, 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 transitionalrun_tpu_bridgeinput and gates both TPU jobs (run-tpuandrun-tpu-bridge) on the singlerun_tpuinput — so one switch runs all TPU benchmarks. No nightly change needed; the existingrun_tpu: 'true'now triggers both.Together with #3028, Helion's TPU nightly runs its 7 bridgeable kernels through tritonbench on the shared
tpudashboard column, withbmm/cross_entropy/flash_attentionremaining onrun_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 singlerun_tpu=truefired bothrun-tpu(run_tpu.py →bmm) andrun-tpu-bridge(the bridge →gemm), and the two resulting artifacts fold into onetpudashboard column — both kernels undertpu, labeledTPU v7, with notpu_bridge/tpu_pallassplit (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_compilecomparator (flex_attention) doesn't run on TPU, and itssdpabaseline hitsfused_sdp_choice(): no viable SDPBackend ... for TorchTPU(tritonbench forcesSDPBackend.FLASH_ATTENTION, which TorchTPU rejects; run_tpu.py's plainscaled_dot_product_attentionlets TorchTPU choose and works). Keeping it on run_tpu.py also preserves its existing shapes/numbers with no discontinuity.