fix(pt): reduce border op CUDA blocking

[njzjz-bot]

Summary

• add C++ Meta kernels for deepmd_export::border_op and border_op_backward so output metadata is available without running the real implementation
• expose gpuMemcpyAsync in DeepMD's CUDA/HIP device abstraction and use it for local self-send device-to-device copies on PyTorch's current stream
• document why the Meta kernels are safe (border_op returns g1 metadata; backward returns grad_g1 metadata) and why the copy still uses PyTorch's current stream

Motivation / notes

The pt_expt border communication wrappers already have Python fake impls, but the C++ op registrations did not provide a Meta dispatch implementation. PyTorch documents that a FakeTensor kernel (also called a meta kernel) is required for custom operators to work with compile/export/FX, and opcheck verifies that the fake/meta kernel returns the same metadata (sizes/strides/dtype/device/etc.) as the real op:

• https://docs.pytorch.org/docs/stable/library.html#torch.library.opcheck
• https://docs.pytorch.org/tutorials/advanced/cpp_custom_ops.html

For these wrappers, metadata is input-derived:

• border_op(..., g1, ...) -> Tensor returns a cloned export output with the same metadata as g1
• border_op_backward(..., grad_g1, ...) -> Tensor returns a cloned export output with the same metadata as grad_g1

So the Meta kernels can use torch::empty_like(...) and avoid touching data_ptr(), MPI, or CUDA/HIP runtime state during shape-only dispatch.

For CUDA/HIP local self-send copies, this PR now exposes gpuMemcpyAsync in gpu_cuda.h / gpu_rocm.h and uses that abstraction in comm.cc. The stream is still obtained from PyTorch's current CUDA/HIP stream because this is a PyTorch custom op; enqueueing work on the default stream would risk incorrect ordering or unintended synchronization with neighbouring PyTorch ops.

This does not try to make the MPI send/recv/barrier path asynchronous; those remain host-side synchronization points by nature.

Tests

• git diff --check HEAD~1..HEAD

A local build/test was not run in this environment because PyTorch is not installed for the active Python interpreter.

Authored by OpenClaw (model: custom-chat-jinzhezeng-group/gpt-5.5)

Summary by CodeRabbit

• Bug Fixes

  • Improved GPU device-to-device transfers during MPI border self-sends by performing async copies on the active CUDA/HIP stream for GPU-resident tensors (CPU paths still use host copies).

• New Features

  • Added Meta (FakeTensor) dispatch implementations for deepmd_export border forward/backward operations to support export/AOT workflows without real communication.

• Chores

  • Standardized async GPU memcpy usage by introducing gpuMemcpyAsync aliases for both CUDA and HIP backends.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

source/op/pt/comm.cc gains GPU memcpy convenience macros, a gpu_memcpy_d2d_current_stream helper dispatching to the current CUDA/HIP stream, and switches MPI self-send paths from synchronous to async copies. Meta-dispatch stubs for border_op and border_op_backward are added and registered under TORCH_LIBRARY_IMPL(deepmd_export, Meta, ...).

Changes

Stream-aware GPU D2D Memcpy with Meta Dispatch

Layer / File(s)	Summary
GPU memcpy async convenience macros source/lib/include/gpu_cuda.h, source/lib/include/gpu_rocm.h	Introduces gpuMemcpyAsync macro aliases mapping to cudaMemcpyAsync and hipMemcpyAsync respectively.
gpu_memcpy_d2d_current_stream helper and build configuration source/op/pt/comm.cc	Adds <cstddef>, conditionally includes CUDA/HIP context headers for stream access, and defines gpu_memcpy_d2d_current_stream which dispatches async copies on the current framework stream with early exit when element count is zero.
Forward and backward MPI self-send path migrations source/op/pt/comm.cc	Replaces synchronous device-to-device copies with gpu_memcpy_d2d_current_stream in both forward and backward border exchange self-send branches when recv_g1_tensor.is_cuda() is true.
Meta-dispatch stubs and registration source/op/pt/comm.cc	border_op_meta returns torch::empty_like(g1_tensor), border_op_backward_meta returns torch::empty_like(grad_g1), and both are bound in a new TORCH_LIBRARY_IMPL(deepmd_export, Meta, ...) block to enable FakeTensor tracing.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 8.33% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.
Title check	✅ Passed	The title 'fix(pt): reduce border op CUDA blocking' directly and clearly summarizes the main objective: addressing CUDA blocking issues in the PyTorch border operation by converting synchronous copies to asynchronous ones.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

source/op/pt/comm.cc (1)

167-181: ⚠️ Potential issue | 🟠 Major

Guard self-send copy with nsend > 0 to avoid UB from uninitialized send_g1.

send_g1 is only initialized when nsend != 0 (Line 150), but the self-send copy branch (Lines 176-181) executes unconditionally. When nsend == 0, reading send_g1 as a function argument is undefined behavior.

Proposed fix

-        if (recv_g1_tensor.is_cuda()) {
-          gpu_memcpy_d2d_current_stream(recv_g1, send_g1,
-                                        (std::size_t)nsend * tensor_size);
-        } else {
-          memcpy(recv_g1, send_g1,
-                 (unsigned long)nsend * tensor_size * sizeof(FPTYPE));
-        }
+        if (nsend > 0) {
+          if (recv_g1_tensor.is_cuda()) {
+            gpu_memcpy_d2d_current_stream(
+                recv_g1, send_g1, (std::size_t)nsend * tensor_size);
+          } else {
+            memcpy(recv_g1, send_g1,
+                   (unsigned long)nsend * tensor_size * sizeof(FPTYPE));
+          }
+        }

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@source/op/pt/comm.cc` around lines 167 - 181, The self-send copy branch
containing the gpu_memcpy_d2d_current_stream and memcpy calls for recv_g1_tensor
execution is not guarded against the case where nsend equals zero, but send_g1
is only initialized when nsend is non-zero. Add a condition checking nsend > 0
to wrap the entire self-send copy branch (the conditional block checking
recv_g1_tensor.is_cuda()) to prevent reading the uninitialized send_g1 pointer
when nsend is zero, which would result in undefined behavior.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@source/op/pt/comm.cc`:
- Around line 167-181: The self-send copy branch containing the
gpu_memcpy_d2d_current_stream and memcpy calls for recv_g1_tensor execution is
not guarded against the case where nsend equals zero, but send_g1 is only
initialized when nsend is non-zero. Add a condition checking nsend > 0 to wrap
the entire self-send copy branch (the conditional block checking
recv_g1_tensor.is_cuda()) to prevent reading the uninitialized send_g1 pointer
when nsend is zero, which would result in undefined behavior.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Repository UI

Review profile: CHILL

Plan: Pro

Run ID: 6b2beea4-4486-457f-8282-d1775d0a55f5

📥 Commits

Reviewing files that changed from the base of the PR and between 5a0d505 and ff9fea3.

📒 Files selected for processing (1)

• source/op/pt/comm.cc

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 82.18%. Comparing base (5a0d505) to head (248c36e).
⚠️ Report is 223 commits behind head on master.

Additional details and impacted files

@@            Coverage Diff             @@
##           master    #5563      +/-   ##
==========================================
+ Coverage   82.15%   82.18%   +0.03%     
==========================================
  Files         898      898              
  Lines      103306   103584     +278     
  Branches     4410     4432      +22     
==========================================
+ Hits        84867    85130     +263     
+ Misses      17065    17059       -6     
- Partials     1374     1395      +21     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.