Skip to content

Add torch_d2d_read_benchmark_runner.py under examples/microbenchmarks/ providing PyTorch parity with jax_d2d_read_benchmark_runner.py: - #924

Open
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_981398158
Open

copybara-service[bot] wants to merge 1 commit into
mainfrom
test_981398158

Conversation

@copybara-service

Copy link
Copy Markdown

Add torch_d2d_read_benchmark_runner.py under examples/microbenchmarks/ providing PyTorch parity with jax_d2d_read_benchmark_runner.py:

  • Implements symmetrical sender and receiver workflows: sender populates deterministic KV-cache tensors on TPU HBM, arms KVCacheManager for read, and coordinates via gRPC; receiver coordinates metadata, initiates start_read, polls transfer completion, computes bandwidth (Gbps and GB/s), and verifies tensor consistency byte-for-byte.
  • Adds target :torch_d2d_read_benchmark_runner in examples/microbenchmarks/BUILD.
  • Updates examples/microbenchmarks/README.md to document the PyTorch D2D read microbenchmark.

@google-cla

google-cla Bot commented Sep 14, 2026

Copy link
Copy Markdown

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

…rks/` providing PyTorch parity with `jax_d2d_read_benchmark_runner.py`:

- Implements symmetrical sender and receiver workflows: sender populates deterministic KV-cache tensors on TPU HBM, arms `KVCacheManager` for read, and coordinates via gRPC; receiver coordinates metadata, initiates `start_read`, polls transfer completion, computes bandwidth (Gbps and GB/s), and verifies tensor consistency byte-for-byte.
- Adds target `:torch_d2d_read_benchmark_runner` in `examples/microbenchmarks/BUILD`.
- Updates `examples/microbenchmarks/README.md` to document the PyTorch D2D read microbenchmark.

PiperOrigin-RevId: 981398158
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

0 participants