RFC: 为 torch_npu MemPool 增加 `sub_block_alignment` 参数

# RFC: 为 torch_npu MemPool 增加 `sub_block_alignment` 参数

- **状态**: Draft
- **作者**: A SGLang Contributor
- **创建日期**: 2026-05-23
- **相关项目**: torch_npu, CANN HCCL, Mooncake (AscendDirectTransport), SGLang
- **目标读者**: torch_npu 维护者、CANN HCCL 团队、SGLang / vLLM 等推理框架开发者

---

## 1. 背景 (Background)

在大模型 PD 分离（Prefill / Decode disaggregation）等场景中，SGLang 通过 Mooncake 的
`AscendDirectTransport` 把 KV Cache、辅助 buffer（output_ids、logprobs、hidden_states
等）注册到 **CANN HCCL IPC RMA**，实现跨节点零拷贝传输。

CANN HCCL IPC 注册路径（`rtsIpcMemGetExportKey` → `halShmemCreateHandle`）要求被注册的
每一段显存：

```
va  % 2 MiB (2 * 1024 * 1024) == 0
len % 2 MiB (2 * 1024 * 1024) == 0
```

否则在 `ADXL Connect` 阶段会直接报 **Error 503900**，驱动侧日志类似：

```
[DRV] Invalid para. va=0x12d6371b8000 page_size=2097152
      Create_para_check fail.
[DRV] Ipc node attr pack fail. len=2479489024
[GE]  Failed to connect, ErrorNo: 503900
```

底层分配链上，Mooncake 的 `AscendAllocator` 通过
`aclrtMalloc(..., ACL_MEM_MALLOC_HUGE_ONLY)` 拿到的**原始大块**本身已经是 2 MiB 对齐的；
但调用栈中间的 `torch_npu` CachingAllocator 会**切分（split）**这些大块以服务更小的
`torch.zeros` / `torch.empty` 请求，切出的子块（slab）起始地址往往**不再** 2 MiB 对齐：

```
torch.zeros(...)                       # 上层框架请求
    └─> torch_npu NPUCachingAllocator  # 切分 / 复用 slab，破坏对齐
            └─> Mooncake mc_ascend_malloc (2 MiB aligned raw block)
                    └─> aclrtMalloc(ACL_MEM_MALLOC_HUGE_ONLY)
```

由于 CachingAllocator 位于 PluggableAllocator 之上，**框架层目前没有公开 API 可以
影响子块的切分策略**，因此即便底层 allocator 完全合规，应用层拿到的 tensor 仍然
可能 `data_ptr() % 2 MiB != 0`，直接导致 IPC 注册失败。

---

## 2. 当前框架侧的临时方案 (Current Workaround in SGLang)

SGLang 在 `python/sglang/srt/hardware_backend/npu/alignment.py` 中实现了一个
**用户态自对齐**辅助函数 `zeros_2m_aligned`：

1. 按 `requested_bytes + 2 MiB` 多申请一个对齐块；
2. 计算第一个 `>= data_ptr()` 的 2 MiB 边界；
3. 返回从该边界开始、长度为 `requested_bytes` 的连续 view；
4. 通过 tensor 引用计数保持底层 storage 存活。

KV Cache（MHA / MLA）的每一层、MetadataBuffers 的每一个字段都改造为按上述方式
独立分配。同时在 `register_buffer_to_engine` 入口处加 `assert_2m_aligned_kv_args`
做 fail-fast 校验。

### 该方案的缺陷

- **显存开销**：每个独立 buffer 额外浪费最多 (2 MiB − 1) 字节。KV Cache 层数 ×
  K/V × （+ 可选 index_k）+ MetadataBuffers 多个字段时，单实例额外占用可达
  ~数百 MB ~ 1 GB 量级。
- **侵入式**：所有需要注册到 HCCL IPC 的 buffer 都要走专用 helper，无法直接复用
  `torch.zeros / torch.empty / nn.Parameter`，对上游代码侵入大。
- **生态不可扩展**：vLLM、TensorRT-LLM-Ascend、自研推理框架等都要重复实现一遍同样
  的对齐逻辑。
- **无法覆盖隐式分配**：模型权重 reshard、临时激活、PluggableAllocator 之外的 NPU
  分配路径无法自对齐，未来一旦扩展到更多 buffer 类型，方案会快速失控。

---

## 3. 提案 (Proposal)

我们建议在 `torch_npu` 中为 **MemPool**（以及/或 `NPUPluggableAllocator`）增加一个
新的可选参数 `sub_block_alignment`，让 CachingAllocator 在切分子块时保证返回地址
满足该对齐要求。

### 3.1 API 设计

**方案 A（推荐）：MemPool 构造参数**

```python
import torch
import torch_npu
from torch_npu.npu import NPUPluggableAllocator

allocator = NPUPluggableAllocator("libmc_ascend_alloc.so", "mc_ascend_malloc", "mc_ascend_free")

pool = torch_npu.npu.MemPool(
    allocator.allocator(),
    sub_block_alignment=2 * 1024 * 1024,   # 新增参数
)

with torch_npu.npu.use_mem_pool(pool):
    x = torch.zeros((1024, 128), dtype=torch.bfloat16, device="npu:0")

assert x.data_ptr() % (2 * 1024 * 1024) == 0
```

**方案 B：PluggableAllocator 构造参数**

```python
allocator = NPUPluggableAllocator(
    "libmc_ascend_alloc.so",
    "mc_ascend_malloc",
    "mc_ascend_free",
    sub_block_alignment=2 * 1024 * 1024,
)
```

两种方案不互斥，方案 A 粒度更细（同一进程可以同时存在对齐 / 非对齐多个 pool），
**优先推荐 A**；B 可作为补充，覆盖不显式使用 MemPool 的旧代码。

### 3.2 语义约定 (Semantics)

设 `A = sub_block_alignment`，要求满足 `A` 为 2 的幂且不小于 device page size：

1. **分配对齐**：通过此 pool 分配出的任意子块 `p` 满足 `p % A == 0`。
2. **长度对齐**：被记账（charge）的子块长度向上取整到 `A` 的倍数；
   即从 CachingAllocator 视角看，子块占用 `ceil(requested_bytes / A) * A` 字节。
3. **切分约束**：只有当一个空闲块 `[base, base + size)` 本身满足
   `base % A == 0 && size >= ceil(req, A)` 时才允许从中切出新的子块；不满足时
   要么向上级 allocator 重新申请，要么使用其他兼容的空闲块。
4. **释放与复用**：释放的子块仍保持其原始对齐特性，进入对齐感知的 free list，
   供后续相同 pool 的请求复用。
5. **向后兼容**：未指定 `sub_block_alignment`（或显式置 `None / 0`）时，行为与
   当前完全一致，**不引入任何额外开销**。
6. **校验**：传入非 2 的幂或小于 1 的值时直接抛 `ValueError`。

### 3.3 内存开销与碎片

- 对于绝大多数 KV / metadata buffer，单笔请求字节数远大于 2 MiB，向上取整带来的
  内部碎片占比 < 1%。
- 对于小 tensor（< 2 MiB），建议**不要**放入 `sub_block_alignment > 0` 的 pool；
  这是上层框架的责任，文档中需要给出明确指引。
- 长期累积碎片由 CachingAllocator 的常规回收策略管理，与现有路径一致。

### 3.4 与底层 allocator 的关系

- 该参数**只影响 CachingAllocator 的切分策略**，不改变底层 `aclrtMalloc` /
  PluggableAllocator 的 ABI。
- 底层 allocator 已返回的大块若本身不满足 `A` 对齐，CachingAllocator 内部应自行
  跳过对齐前缀（类似当前 SGLang 自对齐方案），并把不可用的前缀记入 overhead
  统计指标，便于排查。

---

## 4. 受益场景 (Use Cases)

1. **PD 分离 / KV 跨节点传输**：CANN HCCL IPC RMA、RDMA `ibv_reg_mr`、Mooncake
   `AscendDirectTransport` 等都对注册段有对齐要求，原生支持后整条链路免去手工自对齐。
2. **共享内存 / IPC**：跨进程共享 NPU 显存（如 vLLM disaggregated serving、
   Ray actor 间共享 KV）天然需要 page 粒度对齐。
3. **大页 / Huge Page 一致性**：与 `ACL_MEM_MALLOC_HUGE_ONLY` 的语义对齐，避免
   "底层大页、上层散页" 的语义割裂。
4. **生态统一**：vLLM-Ascend、SGLang、MindIE 等可以共用同一套对齐保证，去掉各自
   的 hack 层。

---

## 5. 兼容性 (Compatibility)

| 维度 | 影响 |
|------|------|
| 不使用新参数的现有代码 | **零影响**，默认行为不变 |
| ABI / C++ 接口 | MemPool 构造增加可选字段，旧调用兼容 |
| PluggableAllocator 协议 | 不变 |
| 已有 checkpoint / 权重加载 | 不涉及 |
| Profiling 工具（msprof 等） | 不涉及 |

---

## 6. 同行实现 (Prior Art) — vllm-ascend 也在做完全一样的 hack

为了证明这**不是 SGLang 一家的特殊需求**，我们调研了
[vllm-project/vllm-ascend](https://github.com/vllm-project/vllm-ascend) 主分支的实现。
结论非常直接：**vllm-ascend 在解决 503900 时独立收敛到了与 SGLang 几乎一字不差的临时
方案**——`size + 2 MiB` 超额分配 + 前向取齐 + 注册入口断言。

### 6.1 vllm-ascend 的对齐辅助函数

`vllm_ascend/worker/v2/attn_utils.py`：

```python
def _align_memory(tensor: torch.Tensor, alignment: int) -> torch.Tensor:
    data_ptr = tensor.data_ptr()
    aligned_addr = (data_ptr + alignment - 1) // alignment * alignment
    offset = (aligned_addr - data_ptr) // tensor.element_size()
    return tensor[int(offset) :]
```

与 SGLang `zeros_2m_aligned` 公式完全一致。

### 6.2 vllm-ascend 的 KV Cache 分配（v2 路径）

```python
# vllm_ascend/worker/v2/attn_utils.py::_allocate_kv_cache
alignment = 2 * 1024 * 1024
...
if vllm_config.kv_transfer_config is None:
    k_tensor = torch.zeros(k_tensor_size, dtype=torch.int8, device=device)
    v_tensor = torch.zeros(v_tensor_size, dtype=torch.int8, device=device)
else:
    # PD 分离场景：被迫超额分配 2 MiB 再切片
    k_tensor = torch.zeros(k_tensor_size + alignment, dtype=torch.int8, device=device)
    v_tensor = torch.zeros(v_tensor_size + alignment, dtype=torch.int8, device=device)
    k_tensor = _align_memory(k_tensor, alignment)[:k_tensor_size]
    v_tensor = _align_memory(v_tensor, alignment)[:v_tensor_size]
```

### 6.3 vllm-ascend 的 v1 路径——同一份代码被 copy-paste 三次

`vllm_ascend/worker/model_runner_v1.py::_allocate_kv_cache_tensors` 里
mamba / hybrid attn-mamba / 普通 attn 三个分支**几乎逐字重复**同一段对齐模板：

```python
# 分支 1: mamba / hybrid / cache_only_layers
if self.vllm_config.kv_transfer_config is None:
    tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=self.device)
else:
    cache_size_aligned = kv_cache_tensor.size + alignment
    tensor = torch.zeros(cache_size_aligned, dtype=torch.int8, device=self.device)
    tensor = self._align_memory(tensor, alignment)[: kv_cache_tensor.size]

# 分支 2: attn + use_compress（再来一遍）
if self.vllm_config.kv_transfer_config is None:
    tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=self.device)
else:
    cache_size_aligned = kv_cache_tensor.size + alignment
    tensor = torch.zeros(cache_size_aligned, dtype=torch.int8, device=self.device)
    tensor = self._align_memory(tensor, alignment)[: kv_cache_tensor.size]

# 分支 3: 普通 attn（K/V 拆分，又重复一遍）
...
```

### 6.4 vllm-ascend 的注册前断言（与 SGLang 的 fail-fast 思路相同）

`vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_layerwise_connector.py`：

```python
def create_kv_buffer(self, first_kv_cache_tuple):
    alignment = 2 * 1024 * 1024
    ...
    self.k_buffer = torch.zeros(first_k_cache.numel() + alignment, dtype=k_dtype, device=...)
    self.k_buffer = align_memory(self.k_buffer, alignment)[: first_k_cache.numel()].view(...)
    self.v_buffer = torch.zeros(first_v_cache.numel() + alignment, dtype=v_dtype, device=...)
    self.v_buffer = align_memory(self.v_buffer, alignment)[: first_v_cache.numel()].view(...)

    for tensor in (self.k_buffer, self.v_buffer):
        assert tensor.data_ptr() % alignment == 0, \
            "The address of the registered kv cache should be aligned to 2M"
        ret_value = self.engine.register_memory(tensor.data_ptr(),
                                                tensor.numel() * tensor.element_size())
        ...

# 在 register_kv_caches 注册前再断言一次
assert min(set(tensor_addrs)) % (2 * 1024 * 1024) == 0, \
    "Tensor start addr is not align with 2M."
```

### 6.5 这种现状非常恶心 (This Status Quo Is Frankly Awful)

把 SGLang 和 vllm-ascend 的实现放在一起看，问题已经不再是"某个框架的特例"，而是
**整个 Ascend 推理生态都在重复同一份丑陋的 workaround**：

1. **底层 ABI 的脏细节泄漏到了上层框架**。
   "底层 PluggableAllocator 给的是 2 MiB 对齐大块、但 torch_npu CachingAllocator 在
   切分时把对齐丢了"——这是一个**纯粹的 torch_npu 内部实现细节**，本不应该穿透到
   推理框架的代码里。可现在，每一个面向 Ascend 的推理框架都被迫去理解、复现、维护
   这同一段补丁。

2. **生态在重复造同一个轮子，而且越造越脏**。
   - SGLang 把它抽到 `hardware_backend/npu/alignment.py`，至少做了模块化。
   - vllm-ascend 直接把同一段 4 行公式 **copy-paste 了 4 处**（v1 三个分支 + v2 +
     connector 的 `create_kv_buffer`），任何一处忘了改都是显存泄漏或注册失败。
   - 未来 MindIE、自研推理框架、训练框架（一旦也跨节点共享 KV）都要再写一遍。
   - 没有任何一个上层框架有能力**真正修复**这个问题，因为根因在 CachingAllocator 内部。

3. **每个上层框架都要自己处理"应不应该启用对齐"的开关**。
   vllm-ascend 用 `kv_transfer_config is not None` 判断、SGLang 用 `is_npu()` 判断，
   语义完全不一致；同样的 buffer 在不同框架下要不要对齐取决于框架作者怎么想，
   而**对齐本身是硬件 + 驱动 + HCCL 的硬性要求**，跟上层业务逻辑没有任何关系。
   这种"业务层替硬件层做决策"的现状本身就是设计上的倒错。

4. **`torch.zeros` 的语义被悄悄改变**。
   上层框架不得不区分"普通 `torch.zeros`"和"对齐 `torch.zeros`"，前者完全可能在
   IPC 注册时炸掉，后者要付出最多 2 MiB / buffer 的显存代价。**用户/开发者每写一行
   显存分配代码都要担心"这块会不会被注册"**——这是非常糟糕的认知负担。

5. **显存浪费在生态尺度上被放大**。
   单个 KV layer 浪费最多 2 MiB，乘以 (层数 × K/V × 框架数 × 实例数)，整个 Ascend
   推理集群每天都在为 torch_npu 的一个内部实现细节多消耗可观的 HBM。

### 6.6 我们的立场

**底层硬件 / 驱动 / 通信库的差异，理应在统一的 torch_npu 层规避，而不是让上层框架
按底层差异各自适配。** `torch_npu` 作为 Ascend 在 PyTorch 生态的官方桥梁，是
唯一一个**有资格、有上下文、也有义务**消化这种差异的层级：

- 它直接持有 CachingAllocator 的内部状态，可以零侵入地修改切分策略；
- 它对所有上层框架是单一入口，一次修复全生态受益；
- 它本来就承诺"让用户像在 CUDA 上一样使用 PyTorch"，今天用户却被迫学习
  `aclrtMalloc` 大页、HCCL IPC 2 MiB 对齐这些 CUDA 上根本不存在的硬件细节。

因此，**`sub_block_alignment` 不是一个 nice-to-have 的优化，而是一个"把本应属于
torch_npu 的责任收回到 torch_npu 自己手里"的修复**。本 RFC 提出的 API 就是这件
事情的最小可行表达。

---

## 7. 替代方案对比 (Alternatives Considered)

1. **保持框架层自对齐（现状）**
   - 优点：无需修改 torch_npu。
   - 缺点：见 §2，显存浪费、侵入大、生态分裂。

2. **在 PluggableAllocator 内部按 2 MiB 切片**
   - 缺点：CachingAllocator 仍然可能在其之上再做不对齐切分，**不能根治**。

3. **HCCL 团队放松 2 MiB 限制**
   - 受限于硬件 / 驱动 page 表实现，短期不可行；即便放松也需要长周期发版。

4. **暴露 CachingAllocator 切分策略的 hook**
   - 表达力强但接口复杂，且容易被滥用造成碎片。`sub_block_alignment` 是一个
     语义明确、最小可用（minimal viable）的子集。

---

## 8. 落地计划 (Rollout Plan)

1. **Phase 1 — 设计评审**：torch_npu 维护者评审 API、语义与碎片影响。
2. **Phase 2 — 原型实现**：在 `torch_npu` 的 NPUCachingAllocator 中接入
   `sub_block_alignment`；在 MemPool / PluggableAllocator 暴露参数；新增单测覆盖
   对齐、复用、碎片、并发四类用例。
3. **Phase 3 — SGLang 适配**：将
   `python/sglang/srt/hardware_backend/npu/alignment.py` 中的 `zeros_2m_aligned`
   退化为对 MemPool 的薄封装；保留 `assert_2m_aligned_kv_args` 作为 fail-fast。
4. **Phase 4 — 文档与示例**：在 torch_npu 文档中补充 "与 HCCL IPC / Mooncake
   集成" 章节，给出端到端示例。
5. **Phase 5 — 下线 workaround**：在两个稳定版本后移除 SGLang 侧自对齐路径。

---

## 9. 验收标准 (Acceptance Criteria)

- [ ] 指定 `sub_block_alignment=2*1024*1024` 后，任意 `torch.zeros / empty /
      ones / Parameter` 在该 pool 上的 `data_ptr() % 2 MiB == 0`。
- [ ] 未指定该参数时，分配性能与显存占用与当前版本一致（±1%）。
- [ ] CANN HCCL IPC RMA 注册在 SGLang PD 全链路上不再出现 503900。
- [ ] 单测覆盖：对齐正确性、释放后复用对齐、并发分配、跨 stream、显存压力下的
      fallback 路径。
- [ ] 文档：API reference + 一个端到端的 IPC 注册示例。

---

## 10. 开放问题 (Open Questions)

1. `sub_block_alignment` 是否需要支持**运行时调整**（例如在已分配 pool 上动态修改）？
   倾向于：**不支持**，仅构造期设置，避免实现复杂度。
2. 是否需要同时暴露 `min_block_size`，让用户显式限制最小切分粒度？可作为后续 RFC。
3. 是否需要在 PyTorch 上游（CUDA CachingAllocator）做对应能力对齐？建议先在
   `torch_npu` 验证后再向上游讨论。
4. 与 `expandable_segments` / `garbage_collection_threshold` 等已有
   CachingAllocator 选项的交互语义，需要在实现阶段明确并补充矩阵测试。

---

## 11. 参考资料 (References)

- CANN HCCL IPC: `rtsIpcMemGetExportKey`、`halShmemCreateHandle` 文档
- Ascend Runtime: `aclrtMalloc` / `ACL_MEM_MALLOC_HUGE_ONLY`
- Mooncake `AscendDirectTransport` 源码
- SGLang 临时方案: https://github.com/whybeyoung/sglang/blob/fix_npu_2m_align/python/sglang/srt/hardware_backend/npu/alignment.py
- 相关错误码: ADXL Connect Error **503900**
- https://gitcode.com/cann/hixl/wiki/HIXL%E5%B8%B8%E8%A7%81%E9%97%AE%E9%A2%98%E5%AE%9A%E4%BD%8D%E6%89%8B%E5%86%8C.md

<img width="1583" height="1010" alt="Image" src="https://github.com/user-attachments/assets/507f9c8d-61e8-4613-876e-688857404f3f" />
- vllm-ascend 同类 workaround:
  - `vllm_ascend/worker/v2/attn_utils.py::_align_memory` / `_allocate_kv_cache`
  - `vllm_ascend/worker/model_runner_v1.py::_allocate_kv_cache_tensors`
  - `vllm_ascend/distributed/kv_transfer/kv_p2p/mooncake_layerwise_connector.py::create_kv_buffer`
- 相关错误码: ADXL Connect Error **503900**



Navigation Menu

Uh oh!

RFC: 为 torch_npu MemPool 增加 `sub_block_alignment` 参数 #139

RFC: 为 torch_npu MemPool 增加 `sub_block_alignment` 参数

1. 背景 (Background)

2. 当前框架侧的临时方案 (Current Workaround in SGLang)

该方案的缺陷

3. 提案 (Proposal)

3.1 API 设计

3.2 语义约定 (Semantics)

3.3 内存开销与碎片

3.4 与底层 allocator 的关系

4. 受益场景 (Use Cases)

5. 兼容性 (Compatibility)

6. 同行实现 (Prior Art) — vllm-ascend 也在做完全一样的 hack

6.1 vllm-ascend 的对齐辅助函数

6.2 vllm-ascend 的 KV Cache 分配（v2 路径）

6.3 vllm-ascend 的 v1 路径——同一份代码被 copy-paste 三次

6.4 vllm-ascend 的注册前断言（与 SGLang 的 fail-fast 思路相同）

6.5 这种现状非常恶心 (This Status Quo Is Frankly Awful)

6.6 我们的立场

7. 替代方案对比 (Alternatives Considered)

8. 落地计划 (Rollout Plan)

9. 验收标准 (Acceptance Criteria)

10. 开放问题 (Open Questions)

11. 参考资料 (References)

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

维度	影响
不使用新参数的现有代码	零影响，默认行为不变
ABI / C++ 接口	MemPool 构造增加可选字段，旧调用兼容
PluggableAllocator 协议	不变
已有 checkpoint / 权重加载	不涉及
Profiling 工具（msprof 等）	不涉及

Uh oh!

RFC: 为 torch_npu MemPool 增加 sub_block_alignment 参数 #139

Description

RFC: 为 torch_npu MemPool 增加 sub_block_alignment 参数

1. 背景 (Background)

2. 当前框架侧的临时方案 (Current Workaround in SGLang)

该方案的缺陷

3. 提案 (Proposal)

3.1 API 设计

3.2 语义约定 (Semantics)

3.3 内存开销与碎片

3.4 与底层 allocator 的关系

4. 受益场景 (Use Cases)

5. 兼容性 (Compatibility)

6. 同行实现 (Prior Art) — vllm-ascend 也在做完全一样的 hack

6.1 vllm-ascend 的对齐辅助函数

6.2 vllm-ascend 的 KV Cache 分配（v2 路径）

6.3 vllm-ascend 的 v1 路径——同一份代码被 copy-paste 三次

6.4 vllm-ascend 的注册前断言（与 SGLang 的 fail-fast 思路相同）

6.5 这种现状非常恶心 (This Status Quo Is Frankly Awful)

6.6 我们的立场

7. 替代方案对比 (Alternatives Considered)

8. 落地计划 (Rollout Plan)

9. 验收标准 (Acceptance Criteria)

10. 开放问题 (Open Questions)

11. 参考资料 (References)

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

RFC: 为 torch_npu MemPool 增加 `sub_block_alignment` 参数 #139

RFC: 为 torch_npu MemPool 增加 `sub_block_alignment` 参数