[Feat] use host-pinned memory with dual CPU/device addresses for transport buffers by yumingyue624 · Pull Request #1024 · ModelEngine-Group/unified-cache-management

yumingyue624 · 2026-06-12T03:57:02Z

Purpose

Switch ASU send/flag buffers from plain host memory to host-pinned
memory so that CPU code packs SQEs through the local mapping while
HCOMM/RDMA uses the device-visible mapping of the same allocation.

Modifications

BufferManager: allocate host-pinned memory via aclrtMallocHost +
aclrtHostRegisterV2, obtain device pointer via
aclrtHostGetDevicePointer. ScatterGatherEntry gains device_addr.
RegisterMemory uses device address for host-pinned regions.
AsuTransportImpl: send/flag buffers use HOST_PINNED instead of HOST.
asu_submit_flow: pass device_addr to SendIoBatch.
sqe_request: use flagBuffer.device_addr for response_buffer_addr.
Tests: added host-pinned dual-address and device_addr assertions.

Test

buffer_manager_test: HostPinnedRegistersDeviceAddress.
asu_submit_flow_test: BuildSubBatchSendBuffersUsesHostPinnedDeviceAddresses.
sqe_request_test: packed response address matches device_addr.

… transport buffers ## Purpose Switch ASU send/flag buffers from plain host memory to host-pinned memory so that CPU code packs SQEs through the local mapping while HCOMM/RDMA uses the device-visible mapping of the same allocation. ## Modifications 1. BufferManager: allocate host-pinned memory via aclrtMallocHost + aclrtHostRegisterV2, obtain device pointer via aclrtHostGetDevicePointer. ScatterGatherEntry gains device_addr. RegisterMemory uses device address for host-pinned regions. 2. AsuTransportImpl: send/flag buffers use HOST_PINNED instead of HOST. 3. asu_submit_flow: pass device_addr to SendIoBatch. 4. sqe_request: use flagBuffer.device_addr for response_buffer_addr. 5. Tests: added host-pinned dual-address and device_addr assertions. ## Test - buffer_manager_test: HostPinnedRegistersDeviceAddress. - asu_submit_flow_test: BuildSubBatchSendBuffersUsesHostPinnedDeviceAddresses. - sqe_request_test: packed response address matches device_addr.

Infinite666 · 2026-06-12T08:48:42Z

        }

-        if (subBatchContext.flagBuffer.addr == 0 || subBatchContext.flagBuffer.length == 0) {
+        if (subBatchContext.sendSge.device_addr == 0 || subBatchContext.flagBuffer.addr == 0 ||


What about sendSge.addr and sendSge.length?

Infinite666 · 2026-06-12T08:52:19Z

+    TransProvider::MemType providerMemType{TransProvider::MemType::MEM_HOST};
+};
+
+class BufferRegionCreator : public Trans::AscendBuffer {


Naming problem, the parent class is a Buffer and the child class is a Creator?

Infinite666 · 2026-06-12T08:56:46Z

-    }
+    BufferRegionCreator regionCreator;
+    BufferRegion region;
+    auto allocStatus = regionCreator.MakeRegion(memory_type_, total, region);


Violating the inherit spirit, the child class has its own public function and will be called directly.

ygwpz · 2026-06-13T06:14:38Z

+        // The owner keeps the ACL registration alive until after HCOMM has
+        // unregistered the region in BufferManager's destructor.
+        auto owner = std::shared_ptr<void>(hostAddr, [](void* addr) {
+            aclrtHostUnregister(addr);


[Warning] Custom deleter lacks error handling

The custom deleter calls aclrtHostUnregister(addr) and aclrtFreeHost(addr) without checking return values. If aclrtHostUnregister fails, calling aclrtFreeHost on still-registered memory may cause undefined behavior or resource leaks.

Consider:

Logging errors in the deleter (even if silent)

Ensuring the cleanup order is safe per ACL documentation

Using aclrtFreeHost only if unregister succeeded

ygwpz · 2026-06-13T06:14:41Z

+            case MemoryType::HOST: {
+                auto owner = MakeHostBuffer(size);
+                if (!owner) { return AllocationFailed("host"); }
+                region = {owner, owner.get(), owner.get(), TransProvider::MemType::MEM_HOST};


[Suggestion] device_addr semantic inconsistency for HOST memory type

For MemoryType::HOST, deviceAddr is set equal to localAddr (owner.get()). This is semantically ambiguous because plain host memory does not have a device-visible mapping. The comment in buffer_manager.h states device_addr is for device-visible address used by HCOMM/RDMA, but setting it equal to addr for HOST type contradicts this.

Consider:

Setting device_addr = 0 for HOST type (forcing callers to handle the case)

Or explicitly documenting that addr == device_addr for HOST memory means the same physical address is used

ygwpz · 2026-06-13T06:14:43Z

        }

-        if (subBatchContext.flagBuffer.addr == 0 || subBatchContext.flagBuffer.length == 0) {
+        if (subBatchContext.sendSge.device_addr == 0 || subBatchContext.flagBuffer.addr == 0 ||


[Critical] Missing validation for sendSge.length

The validation now checks sendSge.device_addr, flagBuffer.addr, flagBuffer.device_addr, and flagBuffer.length, but sendSge.length is not validated. If sendSge.length is 0 or invalid, the SendIoBatch will have an incorrect length parameter, potentially causing buffer overflow or underflow in RDMA operations.

Add: sendSge.length == 0 to the validation condition.

ygwpz · 2026-06-13T06:14:46Z

+    const auto offset = idx * slot_size_;
+    sge.addr = reinterpret_cast<std::uint64_t>(static_cast<char*>(memory_.get()) + offset);
+    sge.device_addr =
+        reinterpret_cast<std::uint64_t>(static_cast<char*>(device_memory_) + offset);


[Warning] Potential null pointer arithmetic if device_memory_ is null

The device_addr calculation uses static_cast<char*>(device_memory_) + offset. If device_memory_ is nullptr (e.g., due to a failed Init that did not properly set it), this pointer arithmetic produces a garbage address value.

While Allocate() checks memory_ is valid at line 201, it does not verify device_memory_ is non-null. Consider adding: if (!device_memory_) return Status::Error(...)

ygwpz · 2026-06-13T06:14:50Z

+    // Local address used by CPU code for SQE packing and completion polling.
    std::uint64_t addr{0};
+    // Device-visible address used by HCOMM/HIXL and remote RDMA operations.
+    std::uint64_t device_addr{0};


[Performance] Struct size increase may impact cache efficiency

Adding device_addr increases ScatterGatherEntry from ~24 bytes to ~32 bytes (assuming 64-bit system). In high-throughput scenarios with many pending SQEs, this 33% size increase could:

Increase memory footprint for sub-batch contexts

Reduce cache efficiency when processing batches

Impact memory bandwidth in tight loops

This is acceptable given the feature requirement, but worth noting for performance-sensitive deployments.

ygwpz · 2026-06-13T06:15:08Z

    request.dtype = GetTransportConfigAttr<std::uint8_t>(attrs, "dtype");
    request.dspec = GetTransportConfigAttr<std::uint8_t>(attrs, "dspec");
-    request.response_buffer_addr = flagBuffer.addr;
+    request.response_buffer_addr = flagBuffer.device_addr;


[Suggestion] Architectural consistency for user-provided buffer addresses

The response_buffer_addr now correctly uses flagBuffer.device_addr for RDMA operations. However, the entry.buffer_addr at line 212 uses entries[index].buffer.region.addr (CPU address). If user-provided KVBuffer regions are also host-pinned memory, they should similarly use device addresses for RDMA operations.

This creates a semantic inconsistency: internal buffers use device_addr for RDMA, but user buffers use addr. Consider:

Adding a device_addr field to MemoryRegion struct

Or documenting that user-provided buffers must use device-visible addresses if they are host-pinned

Or providing a registration API that captures both addresses

yumingyue624 requested review from FangRun2, Infinite666, Tarrei, Wwwzff, flesher0813, mag1c-h and ygwpz as code owners June 12, 2026 03:57

yumingyue624 changed the base branch from develop to feature_26h1 June 12, 2026 03:57

yumingyue624 requested a review from nrj868 as a code owner June 12, 2026 03:57

Infinite666 reviewed Jun 12, 2026

View reviewed changes

ygwpz reviewed Jun 13, 2026

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Feat] use host-pinned memory with dual CPU/device addresses for transport buffers#1024

[Feat] use host-pinned memory with dual CPU/device addresses for transport buffers#1024
yumingyue624 wants to merge 1 commit into
ModelEngine-Group:feature_26h1from
yumingyue624:adapt_connection

yumingyue624 commented Jun 12, 2026

Uh oh!

Infinite666 Jun 12, 2026

Uh oh!

Infinite666 Jun 12, 2026

Uh oh!

Infinite666 Jun 12, 2026

Uh oh!

ygwpz Jun 13, 2026

Uh oh!

ygwpz Jun 13, 2026

Uh oh!

ygwpz Jun 13, 2026

Uh oh!

ygwpz Jun 13, 2026

Uh oh!

ygwpz Jun 13, 2026

Uh oh!

ygwpz Jun 13, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

yumingyue624 commented Jun 12, 2026

Purpose

Modifications

Test

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants