Repository navigation
Conversation
Codex Review Summary: ✅ CompletedThis comment shows the latest Codex review activity on this pull request.
ℹ️ About Codex in GitHubYour team has set up Codex to review pull requests in this repo. Reviews are triggered when you
Codex reacts with 👀 while any review is running, comments if it has suggestions, and reacts with 👍 once all reviews finish with no findings. |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7a5245b2dd
ℹ️ About Codex in GitHub
Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".
|
|
||
| batch_ptrs, batch_sizes, _ = self._preprocess_tensors_for_put(batch_tensors) | ||
| batch_ptrs, batch_sizes, batch_tensors = self._preprocess_tensors_for_put(batch_tensors) | ||
| if self._buffer_pool is not None and len(set(batch_keys)) == len(batch_keys): |
There was a problem hiding this comment.
When protocol="rdma" is used with an Ascend/NPU TensorDict, _preprocess_tensors_for_put preserves NPU tensors (it only calls .cpu() for CUDA), yet this unconditional pool branch sends them to _copy_between_buffer_and_tensors, whose torch.frombuffer destination is CPU memory. The previous path registered each NPU data_ptr() directly (and _bind_npu_device exists to make those registrations work); now every pool-sized NPU value pays a synchronous device-to-host copy and loses the direct-device transfer path, while only oversized or pool-exhausted groups retain the old behavior. Restrict pooling to CPU-backed tensors so NPU puts keep the registered-pointer path.
Useful? React with 👍 / 👎.
chenjr97
changed the title
chenjr97
changed the title
chenjr97
changed the title
overview
Extend the pre-registered buffer pool approach from #175, which optimized GET, to CPU RDMA PUT.
Implementation
BufferPool, backed by the local buffer registered during store setup.Evaluation
Setup
Benchmarks ran on a single node with the PUT client and storage segment colocated, using RDMA with either 1 or 8 HCAs. Each HCA port reported an active Ethernet link rate of 200 Gb/s.
0.3.12.post1uint8tensorsLatency measures
await client.async_put(metadata=None), including metadata, storage, and controller work. Payload preparation, initialization, subsequent GET verification, and partition clearing run outside the timed region. The figure below presents the complete benchmark results.PUT latency
The table below includes the single-key, 1 GiB total, and largest tested cases for each payload size, with duplicates removed. Speedup is baseline median latency divided by pooled median latency.