Add pre-submit gating test (H2D/D2H/H2H)

[haochengd-google]

Add presubmit gates for the D2H/H2D and H2H transfer paths

Overview

This PR adds two presubmit checks that guard Raiden's data-transfer paths on a single host, both running in the existing run_benchmarks workflow on every PR:

Test	Path	Gates on	Why
D2H/H2D	TPU HBM ↔ host pinned DRAM (PCIe/DMA), the KV-cache persist/restore copies	correctness + performance — a d2h→h2d round-trip byte-equality check, then both directions' bandwidth vs a recorded floor	D2H/H2D on one host really does traverse PCIe/DMA, so throughput is a meaningful hardware number
H2H	host-to-host KV-cache push	correctness — byte-integrity of the transferred data	on one host H2H runs over loopback, so throughput is not a real bandwidth number; what's meaningful is that the bytes arrive intact

Smoke Test (Deliberately Set to Fail)

Perf gating for H2H belongs on real multi-host hardware, not here — hence H2H is correctness-only on a single host.

CI runtime (measured on the test branch): excluding queue time, a run takes ~11 min with a warm build cache (29050209960, 29050156485, 29042855892) and ~19 min cold, with no build cache (29051584644);
queue wait is on top and variable.

---

D2H/H2D performance gate

Measures, over the node's 8 chips, the median d2h and h2d bandwidth of a small fixed config set and fails the PR if either drops below its recorded floor. A d2h→h2d round-trip byte-equality check runs first (correctness), then the floor
comparison (performance).

Why these configs / this setting

The two configs are large-block [8, 128, 1024, 128] (int32 + float32, 1 layer), deliberately chosen to be saturated and stable:

• Saturated — ~8 MB-per-chip contiguous blocks that sit on the PCIe/DMA bandwidth roofline. Throughput is set by the hardware ceiling, not by per-op software overhead, so the number is steady and reflects the real transfer path.
• Stable within a run — CV ≈ 1% across iterations.
• Stable across runs — run-to-run median drift < 1%.

v5e D2H/H2D roofline — the two configs sit on the saturated plateau

distribution of all 10 configs (200 iters each); the two selected sit near the roof with a tight spread

This stability is exactly what a static-floor presubmit gate needs: the floor can sit just below the normal spread without false-failing innocent PRs. (Sensitive, overhead-bound "canary" shapes drift ~3% run-to-run — too much for a static floor — and are handled by same-machine A/B in post-submit — see Future Work.)

Why the threshold is the normal-core floor

The per-config floor is the lower edge of its own normal run-to-run spread:

floor = median − 3.5 · MADσ        MADσ = 1.4826 · median(|x − median|)

MADσ is an outlier-resistant standard deviation, so a few slow samples don't drag the bound down. A PR fails when its measured median falls below that floor — i.e. only on a move that is large relative to the config's natural noise.

the MAD floor (green) sits just below each config's normal spread

Baselines can be regenerated on demand

Baselines and floors live in h2d_d2h_gating_baselines.json. They are not hand-edited — a manually-triggerable record workflow re-measures on the gate machine (with a large --iters for a stable MAD) and rewrites the file. Re-run it to refresh the floors after a legitimate hardware/toolchain change.

---

H2H correctness gate

A thin Python driver spawns the C++ H2H runner as two processes on one host over loopback (sender/receiver pinned to opposite NUMA nodes). The receiver byte-compares the whole buffer against a position-dependent fill pattern; any
misplacement, wrong offset, partial write, or dropped data fails the gate. No baseline file is needed — the check is self-contained and stateless.

Setup

Aspect	Setting
Interface	loopback (--data_interface=lo, 127.0.0.1)
Placement	sender on NUMA node 0, receiver on NUMA node 1 (forces the copy across the socket interconnect)
Transport op	op=6 — explicit destination blocks, the working H2H push path
Fixed shape	32 layers × 1 shard (runner constants)
Fill pattern	(layer + block + byte_index + interface) & 0xFF — every byte distinct, so any structural error is caught

Why these configs

Unlike D2H/H2D, these configs are not about bandwidth saturation — each (block_size, num_blocks, parallelism) reaches a code path / bug class the others cannot:

block_size	num_blocks	parallelism	Bug class it catches
1 MiB	64	1	baseline block mapping / offset / copy correctness
1 MiB	64	8	races, write interleaving, ordering, stream partitioning
1 MiB − 3 B	64	4	alignment / boundary / partial-write off-by-one that round sizes hide

• Parallelism (P1 → P8) is the primary axis — the push fans blocks across N streams, so concurrency bugs only appear with P > 1; P1 proves the serial path, P8 applies concurrency stress.
• num_blocks = 64 exercises the block-id → offset mapping with many distinct offsets; 1 MiB crosses page boundaries while keeping memory bounded (≈ 4 GiB total, avoiding the OOM the 16/128 MiB configs caused).
• The non-round size (1 MiB − 3 B) probes tail / partial-write handling that power-of-two sizes mask.
• Each config runs only a few iterations — correctness just needs to exercise the path, not collect perf samples.

---

Bypassing the gate

• Correctness is never bypassable — the D2H/H2D round-trip check and the entire H2H gate always block on failure.
• A perf-floor failure can be waived per-PR by adding [skip-perf-gate] (or [skip-h2d-d2h-gating]) to the CL description. The gate still measures and reports, but does not block.
• Maintainers can switch the whole perf gate to report-only by setting "enforce": false in h2d_d2h_gating_baselines.json.

---

Additional information

• The H2H gate drives the C++ runner rather than a Python implementation. A pure-Python H2H path would require pulling in an extra Python API; to avoid that dependency the gate uses the existing C++ runner. We validated the choice: on a single host the C++ and Python H2H bandwidths agree closely (ratios 1.05 and 1.09), so the C++ runner is representative of the Python path.
• D2H/H2D was validated against the existing microbenchmark. The gate's d2h/h2d throughput was compared against the existing microbenchmark and the numbers are very close (ratios 0.98–1.01), confirming the gate measures the same thing.

---

Known limitations

• The container image is a temporary arrangement. container_image refers to a custom ml-build container built for this workload as part of this change, which adds the packages the workload requires on top of the base image. Per the discussion with BAP, these packages will not be added to the official public image, so the image must be self-hosted. It currently lives in a borrowed GCP project (wyzhang-dev) with read access granted to the runner service accounts, and should be moved to a dedicated project. It should not be switched to the official ml-build container, which lacks these packages.
• Baselines and floors are hardware-specific. They are recorded on v5e single-node hardware and must be re-recorded if the hardware or the config set changes.
• The H2H gate uses the C++ runner, not a Python implementation. A pure-Python H2H path was avoided because it would require pulling in an extra Python API; single-host C++ vs Python bandwidths agree closely (ratios 1.05 / 1.09 — see Additional information).

Future work (post-submit)

Two post-submit tests are planned; each is waiting on BAP — one on a machine, one on a platform feature:

• Real (multi-host) H2H bandwidth — waiting on BAP to provision multi-host machines. Measures the actual NIC path instead of single-host loopback, so H2H throughput becomes a meaningful performance number.
• Unsaturated D2H/H2D — waiting on BAP's native same-machine A/B feature
  (measure the PR and its base back-to-back on one machine and compare the delta). Adds the sensitive, overhead-bound "canary" configs to catch the per-op software regressions the saturated presubmit configs are blind to.

the overhead-bound canary's h2d median drifts ~3% between two runs

Why same-machine is required for the unsaturated case: across two independent runs, the overhead-bound canary's median drifts ~3% — comparable to the floor margin — while the saturated sentinels barely move (< 1%). A static floor on the canary would therefore be flaky. Same-machine A/B measures the PR and its base in the same run on the same machine, so this cross-run drift is common-mode and cancels in the delta — which is why the unsaturated test must run same-machine.

---

Files

Gate binaries & measurement

• tpu_raiden/benchmarks/h2d_d2h_benchmark_gating.py — D2H/H2D perf gate (measure, round-trip check, MAD floor compare, skip-tag / enforce handling)
• tpu_raiden/benchmarks/h2h_cpp_gate.py — H2H byte-integrity gate
• tpu_raiden/benchmarks/perf_core.py — shared measurement core (measure(), round-trip verify)

Config & baselines

• tpu_raiden/benchmarks/h2d_d2h_gating_baselines.json — the 2 configs + recorded baselines/floors (+ sigma_k, iters, warmup, enforce)
• tpu_raiden/benchmarks/benchmark_registry.pbtxt — registers both benchmarks (h2d_d2h_gating, h2h_cpp_gate): workload target, v5e env, metric names
• tpu_raiden/benchmarks/benchmark_registry_record.pbtxt — registry for the baseline-record run
• tpu_raiden/benchmarks/h2d_d2h_record_baselines.sh — records/refreshes baselines

CI workflows

• .github/workflows/run_benchmarks.yml — presubmit gate workflow (runs both benchmarks via BAP on PRs)
• .github/workflows/update_d2h_h2d_baseline.yml — manually-triggerable record workflow that rewrites the baselines file

Build infrastructure (needed to build the workload in CI)

• .bazelrc — adds the oss and ci build configs (toolchain/flags, RBE remote cache, and the torch_tpu module override)
• third_party/torch_tpu_stub/BUILD.bazel, third_party/torch_tpu_stub/MODULE.bazel — a stub module so Bazel resolves the torch_tpu dependency without pulling in real torch (target of the .bazelrc --override_module)
• MODULE.bazel — dependency declarations for the above
• tpu_raiden/benchmarks/BUILD — build targets for the gate binaries
• Custom ml-build container image (see Known limitations) — self-hosted,
  carries the extra packages the workload needs

Docs

• tpu_raiden/benchmarks/H2D_D2H_SINGLE_HOST_TEST.md
• tpu_raiden/benchmarks/H2H_SINGLE_HOST_CORRECTNESS_TEST.md

[linamy85]

Thank you so much for putting together this benchmarking framework! While it LGTM, we need a small favor to get this merged. tpu-raiden relies on Google's Copybara system to keep this repository bidirectionally synced with our internal Source of Truth (Google3/Piper). Because of this setup, our automation is strictly configured to ignore GitHub-specific workflow files (like .github/**) and OSS-only configuration files (like .bazelrc) when importing a PR for internal review.

Since this PR currently mixes both managed core source files and OSS-only infrastructure files, our import_pr automation tool will either fail or silently drop the GitHub Action updates during the import process.

To ensure nothing gets lost and everything routes through the correct review pipeline, could you please split this into two separate PRs?

1. PR 1: The CI / OSS Infrastructure (To be merged directly on GitHub) Please include only these files:

• .bazelrc
• .github/workflows/run_benchmarks.yml
• .github/workflows/update_d2h_h2d_baseline.yml

2. PR 2: The Core Benchmarking Code (To be synced to Piper) Please include the remaining codebase modifications. We can get this on Critique and submit first before the PR 1.

[haochengd-google]

Thank you so much for putting together this benchmarking framework! While it LGTM, we need a small favor to get this merged. tpu-raiden relies on Google's Copybara system to keep this repository bidirectionally synced with our internal Source of Truth (Google3/Piper). Because of this setup, our automation is strictly configured to ignore GitHub-specific workflow files (like .github/**) and OSS-only configuration files (like .bazelrc) when importing a PR for internal review.

Since this PR currently mixes both managed core source files and OSS-only infrastructure files, our import_pr automation tool will either fail or silently drop the GitHub Action updates during the import process.

To ensure nothing gets lost and everything routes through the correct review pipeline, could you please split this into two separate PRs?

1. PR 1: The CI / OSS Infrastructure (To be merged directly on GitHub) Please include only these files:

• .bazelrc
• .github/workflows/run_benchmarks.yml
• .github/workflows/update_d2h_h2d_baseline.yml

2. PR 2: The Core Benchmarking Code (To be synced to Piper) Please include the remaining codebase modifications. We can get this on Critique and submit first before the PR 1.

Hi Amy, thanks for your comment! This completely makes sense, and I have submitted a CL as PR 2.