Skip to content

Conversation

@linamy85
Copy link

  • Implement benchmark_pcie_transfer.py to measure H2D and D2H transfer performance using JAX, supporting various transfer modes (Standard, Parallel, Threaded, Chunked).
  • Integrate the new benchmark into run_benchmark.py.
  • Add configuration files for single device, single chip, and single VM topologies in configs/pcie_transfer/.
  • Add scripts/run_pcie_transfer_benchmark.sh for bulk execution with numactl interleaving option.
  • Add guides/pcie_transfer to run on Kubernetes Ironwood 2x2x1

@chishuen chishuen self-requested a review January 16, 2026 09:44
@chishuen
Copy link
Collaborator

Thank you for the PR. Shall we rename pcie_transfer to something else everywhere, because in a TPU tray or VM (consists of 4 chips), there are a few places where PCIe is used if I understand it correctly from https://docs.cloud.google.com/tpu/docs/tpu7x. Replace the name with something like host_device? Wdyt?

@linamy85 linamy85 force-pushed the feature/pcie-transfer-benchmark branch from ed45549 to 252715f Compare January 16, 2026 14:51
@linamy85
Copy link
Author

Thanks Chi, the PR has been updated

- Implement benchmark_pcie_transfer.py to measure H2D and D2H transfer performance using JAX, supporting various transfer modes (Standard, Parallel, Threaded, Chunked).
- Integrate the new benchmark into run_benchmark.py.
- Add configuration files for single device, single chip, and single VM topologies in configs/pcie_transfer/.
- Add scripts/run_pcie_transfer_benchmark.sh for bulk execution with numactl interleaving option.
@linamy85 linamy85 force-pushed the feature/pcie-transfer-benchmark branch 5 times, most recently from 673d7ab to 64f9902 Compare January 19, 2026 05:29
@linamy85 linamy85 force-pushed the feature/pcie-transfer-benchmark branch 13 times, most recently from ef0e7a2 to 323ec75 Compare January 19, 2026 16:03
@linamy85 linamy85 force-pushed the feature/pcie-transfer-benchmark branch from 323ec75 to 2cf3f44 Compare January 19, 2026 16:11
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.

2 participants