kvcache-ai/Mooncake

Fork 0

mirror of https://github.com/kvcache-ai/Mooncake.git synced 2026-11-05 01:17:55 +00:00

[GH-ISSUE #4874] [Bug]: [Ascend Direct] 注册未按 2 MiB 对齐的 DEVICE 地址后，HCCS 建链延迟失败 #1240

New issue

Closed

opened 2026-10-29 21:24:19 +00:00 by bsakaguchi · 2 comments

bsakaguchi commented

2026-10-29 21:24:19 +00:00

Owner

Originally created by @MarvekW on GitHub (Oct 23, 2026).
Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4874

Bug Report

问题描述

通过 Mooncake Ascend Direct 路径注册 NPU MEM_DEVICE 内存时，如果基址没有按 2 MiB 对齐，注册过程可能返回成功，但后续准备对端连接时会失败。失败发生在 HCCL 为已注册内存导出 IPC handle 的阶段，因此错误没有在注册时暴露，而是延迟表现为通用的连接失败。

本例中的 buffer 来自 PyTorch NPU caching allocator 的子分配。分配器的底层 segment 按 2 MiB 对齐，但返回的 tensor data_ptr() 偏移了 512 字节，即 data_ptr() % 2 MiB == 512。buffer 长度为 16,580,608 字节。目前确认的是基址对齐要求；尚未确认长度是否也必须按 2 MiB 对齐。

环境

Mooncake：部署版本 0.3.8；另检查了 0.3.13.post1 的同一注册路径。
传输后端：Ascend Direct / ADXL over HCCS。
HIXL/ADXL：8.5.0。
内存类型：NPU 设备内存（MEM_DEVICE）。
问题出现在 PyTorch NPU caching allocator 返回的子分配地址上。

复现步骤

分配一个 data_ptr() 未按 2 MiB 对齐的 NPU tensor。本例中 data_ptr() % (2 * 1024 * 1024) == 512。
使用 Ascend Direct 传输，将该地址和长度注册给 Mooncake client（例如通过 register_buffer 路径）。
建立一个需要为该新注册设备内存导出 IPC handle 的 HCCS 对端连接。

注册调用可能看起来成功。之后在准备连接、尝试为该内存导出 IPC handle 时才会触发失败。

实际结果

HCCL/驱动日志报告 IPC 导出失败，错误类似 halShmemCreateHandle: EINVAL，对应页大小为 2 MiB。错误沿调用链传播为 rtsIpcMemGetExportKey（507899）、HcclCommPrepare（ret 15）和 ADXL connect（503900）。Mooncake 调用方最终只看到通用连接失败，没有收到指出内存地址对齐问题的注册错误。

在对照测试中，所有新注册设备内存的基址都对齐时，连接准备成功；包含上述未对齐地址时，连接准备失败。一个无效地址就足以导致整个通信通道准备失败。初始的 HcclRegisterGlobalMem 调用仍可能报告成功，因此问题会延迟到后续建链阶段才暴露。

预期结果

Mooncake 应在注册阶段检测这类输入，并返回明确错误，指出所需对齐和不符合要求的 buffer，而不是先接受注册、等到对端建链时才失败。如果 Ascend Direct 需要支持任意调用方 buffer，则需要使用满足对齐要求的底层分配或 staging buffer。

请问能否在 Ascend Direct 注册路径中针对适用的 HCCS/IPC 路径校验基址，并在文档中说明该要求？增加覆盖对齐和未对齐 MEM_DEVICE 地址的回归测试，也可以避免这类错误延迟暴露。本例已确认基址对齐要求，但尚未确认长度对齐要求。

Mooncake 代码路径

当前观察到的调用路径到达 AscendDirectTransport::registerLocalMemory，并将调用方提供的设备地址传给 ADXL RegisterMem。在检查过的版本中，该路径没有在注册前拒绝未对齐的基址。失败随后出现在 HcclCommPrepare 阶段的 HCCL IPC handle 导出过程中。

问题代码

调用方传入的地址

问题 buffer 来自 RH2D null-slot scratch。原实现直接把 PyTorch tensor 的 data_ptr() 作为注册地址，没有先确保它按 2 MiB 对齐：

# 原逻辑的关键部分（省略缓存和锁等外围代码）
scratch = torch.empty(row_bytes, dtype=torch.uint8, device=device)
scratch_addr = scratch.data_ptr()
spans.add((scratch_addr, row_bytes))

# 后续经 Mooncake client 注册该地址和长度
client.register_buffer(scratch_addr, row_bytes)

本次失败样本中，scratch_addr % (2 * 1024 * 1024) == 512。这是实测到的分配结果，不代表每次分配都会固定偏移 512 字节。

Mooncake Ascend Direct 注册路径

Mooncake 0.3.8 的 AscendDirectTransport::registerLocalMemory 将地址和长度放入 MemDesc，识别内存类型后直接调用 ADXL RegisterMem。下面是相关代码摘录（mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp）：

adxl::MemDesc mem_desc{};
mem_desc.addr = static_cast<uint64_t>(reinterpret_cast<uintptr_t>(addr));
mem_desc.len = length;

// 上方识别 mem_type；这里仅跳过启用 buffer pool 时的 HOST 注册
if (mem_type == adxl::MEM_HOST && use_buffer_pool_) {
    return 0;
}

adxl::MemHandle mem_handle;
auto adxl_ret = adxl_->RegisterMem(mem_desc, mem_type, mem_handle);
if (adxl_ret != adxl::SUCCESS) {
    LOG(ERROR) << "Register mem ret:" << adxl_ret << ".";
    return -1;
}

在设置 mem_desc.addr 到调用 RegisterMem 之间，没有针对 MEM_DEVICE 基址的 2 MiB 对齐校验。0.3.13.post1 将注册调用移至 TransferExecutorBase::registerMem，但该路径同样把 addr、length 放入 MemDesc 后调用 ADXL RegisterMem，没有做此校验：

adxl::MemDesc mem_desc{};
mem_desc.addr = reinterpret_cast<uintptr_t>(addr);
mem_desc.len = length;
// ...切换到对应 engine 的 context...
auto adxl_ret = adxl_engines_[engine_idx]->RegisterMem(
    mem_desc, mem_type, mem_handle);

对应源码位置：Mooncake 0.3.8 的 ascend_direct_transport.cpp:395-450；0.3.13.post1 的 ascend_direct_transport.cpp:351-398 和 transfer_executor_base.cpp:469-526。注册阶段 ADXL 可以返回成功，直到后续 HCCL 导出 IPC handle 才因基址不满足对齐要求而失败。

Before submitting...

Ensure you searched for relevant issues and read the [documentation]

Originally created by @MarvekW on GitHub (Oct 23, 2026). Original GitHub issue: https://github.com/kvcache-ai/Mooncake/issues/4874 ### Bug Report ## 问题描述通过 Mooncake Ascend Direct 路径注册 NPU `MEM_DEVICE` 内存时，如果基址没有按 2 MiB 对齐，注册过程可能返回成功，但后续准备对端连接时会失败。失败发生在 HCCL 为已注册内存导出 IPC handle 的阶段，因此错误没有在注册时暴露，而是延迟表现为通用的连接失败。本例中的 buffer 来自 PyTorch NPU caching allocator 的子分配。分配器的底层 segment 按 2 MiB 对齐，但返回的 tensor `data_ptr()` 偏移了 512 字节，即 `data_ptr() % 2 MiB == 512`。buffer 长度为 16,580,608 字节。目前确认的是基址对齐要求；尚未确认长度是否也必须按 2 MiB 对齐。 ## 环境 - Mooncake：部署版本 0.3.8；另检查了 0.3.13.post1 的同一注册路径。 - 传输后端：Ascend Direct / ADXL over HCCS。 - HIXL/ADXL：8.5.0。 - 内存类型：NPU 设备内存（`MEM_DEVICE`）。 - 问题出现在 PyTorch NPU caching allocator 返回的子分配地址上。 ## 复现步骤 1. 分配一个 `data_ptr()` 未按 2 MiB 对齐的 NPU tensor。本例中 `data_ptr() % (2 * 1024 * 1024) == 512`。 2. 使用 Ascend Direct 传输，将该地址和长度注册给 Mooncake client（例如通过 `register_buffer` 路径）。 3. 建立一个需要为该新注册设备内存导出 IPC handle 的 HCCS 对端连接。注册调用可能看起来成功。之后在准备连接、尝试为该内存导出 IPC handle 时才会触发失败。 ## 实际结果 HCCL/驱动日志报告 IPC 导出失败，错误类似 `halShmemCreateHandle: EINVAL`，对应页大小为 2 MiB。错误沿调用链传播为 `rtsIpcMemGetExportKey`（507899）、`HcclCommPrepare`（ret 15）和 ADXL connect（503900）。Mooncake 调用方最终只看到通用连接失败，没有收到指出内存地址对齐问题的注册错误。在对照测试中，所有新注册设备内存的基址都对齐时，连接准备成功；包含上述未对齐地址时，连接准备失败。一个无效地址就足以导致整个通信通道准备失败。初始的 `HcclRegisterGlobalMem` 调用仍可能报告成功，因此问题会延迟到后续建链阶段才暴露。 ## 预期结果 Mooncake 应在注册阶段检测这类输入，并返回明确错误，指出所需对齐和不符合要求的 buffer，而不是先接受注册、等到对端建链时才失败。如果 Ascend Direct 需要支持任意调用方 buffer，则需要使用满足对齐要求的底层分配或 staging buffer。请问能否在 Ascend Direct 注册路径中针对适用的 HCCS/IPC 路径校验基址，并在文档中说明该要求？增加覆盖对齐和未对齐 `MEM_DEVICE` 地址的回归测试，也可以避免这类错误延迟暴露。本例已确认基址对齐要求，但尚未确认长度对齐要求。 ## 相关参考 - [cann/hixl issue #14](https://gitcode.com/cann/hixl/issues/14)：维护者确认了相同的 HCCS 注册地址 2 MiB 对齐要求，并记录了类似的 IPC 导出错误。 - [Ascend/pytorch issue #2077](https://gitcode.com/Ascend/pytorch/issues/2077)：为 PyTorch NPU 内存池增加 `sub_block_alignment` 参数的 RFC，可帮助调用方申请满足对齐要求的子分配。 ## Mooncake 代码路径当前观察到的调用路径到达 `AscendDirectTransport::registerLocalMemory`，并将调用方提供的设备地址传给 ADXL `RegisterMem`。在检查过的版本中，该路径没有在注册前拒绝未对齐的基址。失败随后出现在 `HcclCommPrepare` 阶段的 HCCL IPC handle 导出过程中。 ## 问题代码 ### 调用方传入的地址问题 buffer 来自 RH2D null-slot scratch。原实现直接把 PyTorch tensor 的 `data_ptr()` 作为注册地址，没有先确保它按 2 MiB 对齐： ```python # 原逻辑的关键部分（省略缓存和锁等外围代码） scratch = torch.empty(row_bytes, dtype=torch.uint8, device=device) scratch_addr = scratch.data_ptr() spans.add((scratch_addr, row_bytes)) # 后续经 Mooncake client 注册该地址和长度 client.register_buffer(scratch_addr, row_bytes) ``` 本次失败样本中，`scratch_addr % (2 * 1024 * 1024) == 512`。这是实测到的分配结果，不代表每次分配都会固定偏移 512 字节。 ### Mooncake Ascend Direct 注册路径 Mooncake 0.3.8 的 `AscendDirectTransport::registerLocalMemory` 将地址和长度放入 `MemDesc`，识别内存类型后直接调用 ADXL `RegisterMem`。下面是相关代码摘录（`mooncake-transfer-engine/src/transport/ascend_transport/ascend_direct_transport/ascend_direct_transport.cpp`）： ```cpp adxl::MemDesc mem_desc{}; mem_desc.addr = static_cast<uint64_t>(reinterpret_cast<uintptr_t>(addr)); mem_desc.len = length; // 上方识别 mem_type；这里仅跳过启用 buffer pool 时的 HOST 注册 if (mem_type == adxl::MEM_HOST && use_buffer_pool_) { return 0; } adxl::MemHandle mem_handle; auto adxl_ret = adxl_->RegisterMem(mem_desc, mem_type, mem_handle); if (adxl_ret != adxl::SUCCESS) { LOG(ERROR) << "Register mem ret:" << adxl_ret << "."; return -1; } ``` 在设置 `mem_desc.addr` 到调用 `RegisterMem` 之间，没有针对 `MEM_DEVICE` 基址的 2 MiB 对齐校验。0.3.13.post1 将注册调用移至 `TransferExecutorBase::registerMem`，但该路径同样把 `addr`、`length` 放入 `MemDesc` 后调用 ADXL `RegisterMem`，没有做此校验： ```cpp adxl::MemDesc mem_desc{}; mem_desc.addr = reinterpret_cast<uintptr_t>(addr); mem_desc.len = length; // ...切换到对应 engine 的 context... auto adxl_ret = adxl_engines_[engine_idx]->RegisterMem( mem_desc, mem_type, mem_handle); ``` 对应源码位置：Mooncake 0.3.8 的 `ascend_direct_transport.cpp:395-450`；0.3.13.post1 的 `ascend_direct_transport.cpp:351-398` 和 `transfer_executor_base.cpp:469-526`。注册阶段 ADXL 可以返回成功，直到后续 HCCL 导出 IPC handle 才因基址不满足对齐要求而失败。 ### Before submitting... - [x] Ensure you searched for relevant issues and read the [documentation]

bsakaguchi

2026-10-29 21:24:19 +00:00

closed this issue
added the
bug
label

bsakaguchi commented

2026-10-29 21:24:19 +00:00

Author

Owner

@github-actions[bot] commented on GitHub (Oct 23, 2026):

Thanks for opening this issue, @MarvekW!

Field	Value
Issue	#4874
GitHub user ID	`276451463`
Reporter	@MarvekW

A maintainer will triage this when possible. To help us respond faster, please include:

Mooncake version or commit SHA
Environment (OS, CUDA/driver, RDMA stack if relevant)
Steps to reproduce and expected vs. actual behavior

Useful links: Documentation · Contributing guide

This message was posted automatically by the issue bot.

@github-actions[bot] commented on GitHub (Oct 23, 2026): Thanks for opening this issue, @MarvekW! | Field | Value | |-------|-------| | **Issue** | #4874 | | **GitHub user ID** | `276451463` | | **Reporter** | @MarvekW | A maintainer will triage this when possible. To help us respond faster, please include: - Mooncake version or commit SHA - Environment (OS, CUDA/driver, RDMA stack if relevant) - Steps to reproduce and expected vs. actual behavior Useful links: [Documentation](https://kvcache-ai.github.io/Mooncake/) · [Contributing guide](https://github.com/kvcache-ai/Mooncake/blob/main/CONTRIBUTING.md) > This message was posted automatically by the issue bot.

bsakaguchi commented

2026-10-29 21:24:20 +00:00

Author

Owner

@huangxiaobai-ydmy commented on GitHub (Oct 23, 2026):

Summary: Registering NPU device memory (MEM_DEVICE) through the Ascend Direct path with a base address that is not 2 MiB-aligned succeeds at AscendDirectTransport::registerLocalMemory, but the failure is delayed to connection preparation, where HCCL fails to export the IPC handle (halShmemCreateHandle: EINVAL -> rtsIpcMemGetExportKey (507899) -> HcclCommPrepare (ret 15) -> ADXL connect 503900). The caller only sees a generic connection failure with no hint about alignment.

The misaligned address came from a PyTorch NPU caching-allocator sub-allocation (data_ptr() % 2 MiB == 512): the underlying segment is 2 MiB-aligned but sub-allocations are not. The 2 MiB base-address requirement for HCCS registration was confirmed by the hixl maintainers in the gitcode issue linked above, and the behavior was verified on both 0.3.8 and 0.3.13.post1; current main still has no alignment check in the registration path.

Proposed fix:

registerLocalMemory validates 2 MiB base alignment when the buffer is registered as MEM_DEVICE with remote_accessible=true (only those are exported for peer HCCS connections) and returns an explicit error naming the requirement and the offending address, instead of accepting it and failing later at connect time;
document the alignment requirement for callers that register torch_nPU allocator sub-allocations;
add mock-based unit tests in ascend_direct_transport_test.cpp covering aligned and unaligned MEM_DEVICE addresses (the harness is mock-based, so it runs in CI without NPU hardware).

I have an Ascend 910B machine available for end-to-end verification. Does this direction sound good? I can start on the PR if so.

cc @alogfans @ascend-direct-dev — this touches the Ascend Direct registration path; would one of you have a moment to confirm the direction?

@huangxiaobai-ydmy commented on GitHub (Oct 23, 2026): **Summary**: Registering NPU device memory (`MEM_DEVICE`) through the Ascend Direct path with a base address that is not 2 MiB-aligned *succeeds* at `AscendDirectTransport::registerLocalMemory`, but the failure is delayed to connection preparation, where HCCL fails to export the IPC handle (`halShmemCreateHandle: EINVAL` -> `rtsIpcMemGetExportKey` (507899) -> `HcclCommPrepare` (ret 15) -> ADXL connect 503900). The caller only sees a generic connection failure with no hint about alignment. The misaligned address came from a PyTorch NPU caching-allocator sub-allocation (`data_ptr() % 2 MiB == 512`): the underlying segment is 2 MiB-aligned but sub-allocations are not. The 2 MiB base-address requirement for HCCS registration was confirmed by the hixl maintainers in the gitcode issue linked above, and the behavior was verified on both 0.3.8 and 0.3.13.post1; current `main` still has no alignment check in the registration path. **Proposed fix**: - `registerLocalMemory` validates 2 MiB base alignment when the buffer is registered as `MEM_DEVICE` with `remote_accessible=true` (only those are exported for peer HCCS connections) and returns an explicit error naming the requirement and the offending address, instead of accepting it and failing later at connect time; - document the alignment requirement for callers that register torch_nPU allocator sub-allocations; - add mock-based unit tests in `ascend_direct_transport_test.cpp` covering aligned and unaligned `MEM_DEVICE` addresses (the harness is mock-based, so it runs in CI without NPU hardware). I have an Ascend 910B machine available for end-to-end verification. Does this direction sound good? I can start on the PR if so. cc @alogfans @ascend-direct-dev — this touches the Ascend Direct registration path; would one of you have a moment to confirm the direction?

bsakaguchi referenced this issue

2026-10-29 21:32:46 +00:00

[PR #1240] [MERGED] [Fix] Fix broken link in doc #2098