Skip to content

Conversation

@linamy85
Copy link
Collaborator

  • 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
Collaborator Author

Thanks Chi, the PR has been updated

@linamy85 linamy85 force-pushed the feature/pcie-transfer-benchmark branch 21 times, most recently from abd5271 to dcf0697 Compare January 20, 2026 02:14
@linamy85 linamy85 force-pushed the feature/pcie-transfer-benchmark branch 4 times, most recently from b7ce14b to 6137df9 Compare January 20, 2026 04:52
@linamy85 linamy85 force-pushed the feature/pcie-transfer-benchmark branch 5 times, most recently from 54230c0 to d280255 Compare January 20, 2026 06:12
- 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 4 times, most recently from 81e951a to 8803f87 Compare January 20, 2026 08:16
@linamy85 linamy85 force-pushed the feature/pcie-transfer-benchmark branch from 8803f87 to d777e69 Compare January 20, 2026 08:34
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