Skip to content

Add torch_dma_kv_cache_benchmark.py under examples/microbenchmarks/ providing PyTorch parity with jax_dma_kv_cache_benchmark.py: - #871

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

Add torch_dma_kv_cache_benchmark.py under examples/microbenchmarks/ providing PyTorch parity with jax_dma_kv_cache_benchmark.py:#871
copybara-service[bot] wants to merge 1 commit into
mainfrom
test_976038290

Conversation

@copybara-service

Copy link
Copy Markdown

Add torch_dma_kv_cache_benchmark.py under examples/microbenchmarks/ providing PyTorch parity with jax_dma_kv_cache_benchmark.py:

  • Single-host raw PCIe DMA microbenchmark testing KVCacheManager D2H and H2D throughput and latency across variable layers (1 to 128) and dtypes (bf16, fp32, int32, fp8).
  • Implements comparative baselines against native PyTorch pinned host memory (pin_memory=True, non_blocking=True) and standard unpinned host memory.
  • Logs statistical metrics (mean, median, sample stddev, 95% confidence intervals, throughput in GB/s) to JSONL telemetry log for golden result comparisons in CI/cron.
  • Registers :torch_dma_kv_cache_benchmark_test target in examples/microbenchmarks/BUILD.
  • Updates examples/microbenchmarks/README.md with execution instructions.

@google-cla

google-cla Bot commented Sep 6, 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.

…/` providing PyTorch parity with `jax_dma_kv_cache_benchmark.py`:

- Single-host raw PCIe DMA microbenchmark testing KVCacheManager D2H and H2D throughput and latency across variable layers (1 to 128) and dtypes (bf16, fp32, int32, fp8).
- Implements comparative baselines against native PyTorch pinned host memory (`pin_memory=True`, `non_blocking=True`) and standard unpinned host memory.
- Logs statistical metrics (mean, median, sample stddev, 95% confidence intervals, throughput in GB/s) to JSONL telemetry log for golden result comparisons in CI/cron.
- Registers `:torch_dma_kv_cache_benchmark_test` target in `examples/microbenchmarks/BUILD`.
- Updates `examples/microbenchmarks/README.md` with execution instructions.

PiperOrigin-RevId: 976038290
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