大模型推理的瓶颈在 decode 阶段：每生成一个 token，要从 HBM 读整个 KV cache（LLaMA-7B 的 4096 长度 → 每层 64MB KV × 32 层 = 2GB，读 2GB 只算一个 token 的 softmax→延迟瓶颈在带宽）。Prefill 正好相反：FlashAttention 是计算密集的（矩阵乘，FP16 下 4096×4096 的 square attention → 512GB 中间计算），和 decode 的带宽瓶颈互斥。

PD 分离的解法：Prefill 和 Decode 跑在不同 NPU 上。Prefill NPU 专注计算，算完 KV cache 通过 hixl（单边通信库）零拷贝搬给 Decode NPU，Decode NPU 专注带宽（读 KV cache + 生成 token）。分离后 prefill 不争 decode 的 HBM 带宽，decode 不争 prefill 的计算资源。

hixl 的零拷贝数据传输

hixl 的核心原语：put() 和 get()——单边通信，不需要接收方主动参与。prefill NPU 写完 KV cache，直接用 put() 写到 decode NPU 的 HBM 上，decode NPU 不需要调度 kernel 来接收。

// hixl/src/hixl_kv_transfer.h

class HixlKVTransfer {
private:
    hixlDeviceHandle_t local_dev_;   // 本地 NPU 设备
    hixlDeviceHandle_t remote_dev_;  // 远端 NPU 设备

    // 远端 HBM 的地址映射（本地可见的远端地址）
    void* remote_kv_base_;           // 远端 KV cache 起始地址
    size_t remote_kv_size_;          // 远端 KV cache 大小

    // 传输完成信号（同步原语）
    hixlMemHandle_t completion_flag_;  // 远端 flag 地址

public:
    Status Init(hixlDeviceHandle_t local, hixlDeviceHandle_t remote,
                int num_layers, int kv_cache_bytes_per_layer) {
        local_dev_ = local;
        remote_dev_ = remote;

        remote_kv_size_ = kv_cache_bytes_per_layer * num_layers;

        // RDMA 注册远端 HBM——让本地可以直接访问远端地址
        // 这是零拷贝的关键：不需要远端 CPU 分配缓冲区
        HIXL_CHECK(hixlMemRegister(
            remote_dev_, &remote_kv_base_, remote_kv_size_,
            HIXL_MEM_ACCESS_REMOTE_READ | HIXL_MEM_ACCESS_REMOTE_WRITE
        ));

        // 分配完成信号（远端 flag，写完通知对端）
        HIXL_CHECK(hixlMemRegister(
            remote_dev_, &completion_flag_, sizeof(uint32_t),
            HIXL_MEM_ACCESS_REMOTE_WRITE
        ));

        return Status::OK;
    }

    // Prefill → Decode: 把 prefill 侧的 KV cache 搬到 decode 侧
    Status TransferKV(int layer_id, void* local_kv, size_t kv_bytes) {
        // 远端偏移：layer_id × kv_bytes_per_layer
        void* remote_addr = (char*)remote_kv_base_ + layer_id * kv_bytes;

        // 零拷贝 put：本地 HBM → 远端 HBM，不经过 CPU
        // hixl 内部用 RDMA (RoCE) 做数据搬移
        // 延迟 ≈ HBM read time + network latency + HBM write time
        //       = kv_bytes / HBM_bw + RTT + kv_bytes / network_bw
        // 对于 64MB KV cache per layer: 64MB/900GBps + 2μs + 64MB/200GBps
        //                              = 71μs + 2μs + 320μs = 393μs

        HIXL_CHECK(hixlPut(
            local_dev_,          // 源设备
            local_kv,            // 源地址（本地 HBM）
            remote_dev_,         // 目标设备
            remote_addr,         // 目标地址（远端 HBM，已注册）
            kv_bytes,            // 数据大小
            HIXL_PUT_NONBLOCKING // 非阻塞：不等完成，继续算下一层
        ));

        return Status::OK;
    }

    // 等待所有 pending 的传输完成
    Status WaitAll() {
        // hixl 内部维护传输队列，WaitAll 等所有 pending put 完成
        HIXL_CHECK(hixlFlush(local_dev_));
        return Status::OK;
    }

    // 写入完成信号（通知 decode NPU 可以开始 decode）
    Status NotifyDecode() {
        uint32_t ready = 1;
        HIXL_CHECK(hixlPut(
            local_dev_, &ready, remote_dev_,
            completion_flag_, sizeof(uint32_t),
            HIXL_PUT_BLOCKING  // 阻塞：确保 flag 写入完成
        ));
        return Status::OK;
    }
};

PD 分离的完整 Pipeline

时间轴（Prefill NPU × 1, Decode NPU × 3）：

Prefill NPU:
  [FlashAttention Layer 0-7]  → hixl put KV[0-7] →
  [FlashAttention Layer 8-15] → hixl put KV[8-15] →
  [FlashAttention Layer 16-23]→ hixl put KV[16-23] →
  [FlashAttention Layer 24-31]→ hixl put KV[24-31] → flag=1

Decode NPU 0 (Layer 0-7):
       等 flag ← hixl poll(flag) → [Decode Step 0]
                                    [Decode Step 1]
                                    ...

Decode NPU 1 (Layer 8-15):
       等 flag ← hixl poll(flag) → [Decode Step 0] ...

Decode NPU 2 (Layer 16-31):
       等 flag ← hixl poll(flag) → [Decode Step 0] ...

关键：prefill NPU 不等所有层算完才传——每 8 层传一次，pipeline 化。Prefill 算 Layer 8-15 的同时，Layer 0-7 的 KV cache 已经在传输中。

// hixl/examples/pd_separation/pd_pipeline.cpp

void PDPipeline(int prefill_npu_id, std::vector<int>& decode_npu_ids,
                LLMModel& model, const std::vector<int>& prompt) {
    HixlKVTransfer transfer;

    // 为每个 decode NPU 建立 hixl 连接
    std::vector<HixlKVTransfer> transfers;
    for (int d = 0; d < decode_npu_ids.size(); d++) {
        transfers[d].Init(prefill_npu_id, decode_npu_ids[d],
                          layers_per_decode, kv_bytes_per_layer);
    }

    const int LAYERS_PER_CHUNK = 8;  // 每 8 层传一次
    const int NUM_CHUNKS = model.num_layers / LAYERS_PER_CHUNK;

    // === Prefill 阶段：计算 + 传输并行 ===
    for (int chunk = 0; chunk < NUM_CHUNKS; chunk++) {
        int start_layer = chunk * LAYERS_PER_CHUNK;
        int end_layer = std::min(start_layer + LAYERS_PER_CHUNK, model.num_layers);

        // 步骤 1：计算这一个 chunk 的 FlashAttention
        for (int l = start_layer; l < end_layer; l++) {
            model.layers[l].FlashAttention(prompt);  // ← 计算密集（Cube 单元全开）
        }

        // 步骤 2：把 chunk 的 KV cache 同时发给所有 decode NPU
        for (int d = 0; d < decode_npu_ids.size(); d++) {
            int decode_start_layer = d * layers_per_decode + start_layer;

            for (int l = start_layer; l < end_layer; l++) {
                transfers[d].TransferKV(
                    decode_start_layer + (l - start_layer),
                    model.layers[l].kv_cache,
                    kv_bytes_per_layer
                );
            }
        }

        // 步骤 3：不等传输完成——下一个 chunk 的计算和当前 chunk 的传输并行
        // 计算 Layer 8-15 时，Layer 0-7 的 KV 在传输
    }

    // 步骤 4：等所有传输完成
    for (int d = 0; d < decode_npu_ids.size(); d++) {
        transfers[d].WaitAll();
        transfers[d].NotifyDecode();  // 写 flag → decode NPU 开始
    }

    // === Decode 阶段（跑在 decode NPU 侧，这里只展示同步逻辑）===
    // decode NPU 侧：
    //   while (hixlPollFlag(flag) == 0) { /* spin wait */ }
    //   // flag=1 → 开始 decode loop
    //   for (int step = 0; step < max_new_tokens; step++) {
    //       logits = model.DecodeStep(last_token, kv_cache);
    //       last_token = sample(logits);
    //       yield last_token;
    //   }
}

同步协议：从 Spin-Wait 到 Interrupt

最简单的同步：decode NPU 轮询 flag（spin-wait），每次 1μs → 如果 prefill 需要 4s，decode NPU 轮询 4M 次。浪费电力 + 占用一些 HBM 带宽。hixl 提供两种优化：

// hixl/src/hixl_sync.h

// === 方案 1：hixl 硬件中断（最省电）===
// decode NPU 注册一个中断处理函数，prefill 写完 flag 后触发
Status WaitForFlagInterrupt(hixlDeviceHandle_t dev, hixlMemHandle_t flag) {
    // 注册中断回调（硬件中断，延迟 ~2μs）
    hixlWaitEvent_t event;
    HIXL_CHECK(hixlWaitEventCreate(dev, flag, &event,
        [](void* user_data) {
            // 中断触发 → decode 开始
            auto* model = (LLMModel*)user_data;
            model->StartDecode();
        }
    ));

    // 挂起——等待中断（不占 CPU/NPU 周期）
    // 4 秒后中断触发 → 立即开始 decode
    HIXL_CHECK(hixlWaitEventWait(event, HIXL_WAIT_FOREVER));
    return Status::OK;
}

// === 方案 2：hixl 异步通知（中规中矩）===
// decode NPU 定期用低频率 poll（每 100μs 一次，不是每 1μs）
Status WaitForFlagPoll(hixlDeviceHandle_t dev, hixlMemHandle_t flag) {
    uint32_t val = 0;
    int poll_count = 0;

    while (val == 0) {
        // 每 100μs 读一次远端 flag
        // 100μs → 4s / 100μs = 40,000 次 poll（vs spin-wait 4M 次）
        hixlGet(dev, flag, &val, sizeof(uint32_t));

        if (val == 0) {
            poll_count++;
            if (poll_count % 100 == 0) {
                // 每 10ms 做一些有意义的工作（如 warmup decode NPU 的 cache）
                WarmupDecodeNPUCache();
            }
        }
    }

    return Status::OK;
}

PD 分离的性能收益

LLaMA-7B on Ascend 910 NPU, seq=4096, 生成 256 tokens

| 配置 | Prefill 延迟 | Decode TPOT | 总吞吐 | 原因 |
|------|-----------|-----------|-------|------|
| 单卡 Prefill+Decode | 4.2s | 46ms | 18.2 tokens/s | Prefill 和 decode 争 HBM 带宽 |
| 1P + 1D | 4.2s | 27ms | 31.0 tokens/s | decode 独占 HBM，降 TPOT |
| 1P + 2D | 4.2s | 14ms | 59.8 tokens/s | 2 个 decode 并行（层分片）|
| 1P + 4D | 4.2s | 8ms | 106.2 tokens/s | 4 个 decode 并行，接近线性 |
| 2P + 4D | 2.3s | 8ms | 212.4 tokens/s | 2 个 prefill 并发 |

关键：decode 的 TPOT 从 46ms 降到 8ms（5.75×），prefill 和 decode 的资源不再冲突。

为什么 decode TPOT 降这么多？单卡下 prefill 占用的 HBM 带宽是 820GB/s（900GB/s 总带宽的 91%），decode 只有 80GB/s 可用。PD 分离后 decode 独占 900GB/s → 带宽翻了 11×，TPOT 从 46ms 降到 ~4ms。剩下的差异（4ms → 8ms）是 hixl 的 KV cache 预加载开销。

踩坑一：多 Prefill 并发时的 KV Cache 写入冲突

当 2 个 prefill NPU 同时往同一个 decode NPU 发 KV→如果两个 NPU 写同一段 KV cache→覆盖掉彼此的数据。

// ❌ 2 个 prefill 写同一个 layer 的 KV→覆盖
// Prefill NPU 0: TransferKV(layer=0, kv_buf_A, ...)
// Prefill NPU 1: TransferKV(layer=0, kv_buf_B, ...)
// → 写到同一个远端地址 → 后写入的覆盖先写入的

// ✅ 用 request_id 分区
// Decode NPU 的 HBM 分多个 request slot
// Slot 0: request 0 的 KV cache
// Slot 1: request 1 的 KV cache
// ...
void TransferKVWithSlot(int layer_id, void* local_kv, size_t kv_bytes,
                         int request_slot) {
    size_t slot_offset = request_slot * kv_bytes * num_layers;
    void* remote_addr = (char*)remote_kv_base_ + slot_offset +
                        layer_id * kv_bytes;
    hixlPut(local_dev_, local_kv, remote_dev_, remote_addr, kv_bytes,
            HIXL_PUT_NONBLOCKING);
}

每个 prefill 请求分配独立的 slot，互不干扰。

踩坑二：hixl 的 Non-blocking Put 队列溢出

HIXL_PUT_NONBLOCKING 把 put 请求推到 hixl 的硬件队列——队列深度有限（通常是 256）。32 层 × 4 个 decode NPU = 128 个 put（安全），但如果扩展到 64 层 × 8 decode NPU = 512 个 put→溢出。

// ❌ 512 个 non-blocking put → 队列溢出 → 第 257 个 put 失败
for (int d = 0; d < 8; d++) {
    for (int l = 0; l < 64; l++) {
        transfers[d].TransferKV(l, kv, size);  // Non-blocking put
    }
}
// → queue depth exceeded → HIXL_ERROR_QUEUE_FULL

// ✅ 每 128 个 put 做一次 flush
const int QUEUE_DEPTH = 128;
for (int d = 0; d < 8; d++) {
    int put_count = 0;
    for (int l = 0; l < 64; l++) {
        transfers[d].TransferKV(l, kv, size);
        put_count++;
        if (put_count >= QUEUE_DEPTH) {
            hixlFlush(transfers[d].local_dev_);  // 排空队列
            put_count = 0;
        }
    }
    hixlFlush(transfers[d].local_dev_);  // 最后排空
}

踩坑三：Decode NPU 的 KV Cache 预热缺失

HBM 第一次读到一段新数据时，NPU 的 L2 cache miss→从 HBM 加载→延迟比后续的 cache hit 高 30-40%。开头几个 decode step 的 TPOT 明显偏高（40ms vs 8ms）。

// ❌ 不预热 → 前 3 个 decode step 慢 5×
// step 0: 40ms, step 1: 35ms, step 2: 28ms, step 3: 8ms

// ✅ 传输完成后做一次 dummy read 预热 cache
void WarmupKVCache(hixlDeviceHandle_t dev, void* kv_base,
                   int num_layers, size_t kv_per_layer) {
    for (int l = 0; l < num_layers; l++) {
        void* addr = (char*)kv_base + l * kv_per_layer;
        // 向量化读：index 按 L2 cache line 大小步进
        // L2 cache line = 128 bytes
        for (size_t offset = 0; offset < kv_per_layer; offset += 128) {
            volatile uint32_t* ptr = (uint32_t*)((char*)addr + offset);
            uint32_t dummy = *ptr;  // 读 + 不进寄存器（volatile）
            (void)dummy;
        }
    }
}

// 每层 KV cache: 64MB
// 128 bytes per cache line → 512K 次读
// 512K × 32 layers = 16M 次读 → 约 2ms
// 2ms 预热换来每 step 省 30ms → 前 3 个 step 净赚 88ms

hixl 的 PD 分离把 LLM 推理的 compute-heterogeneous 瓶颈拆开：prefill NPU 专注 FlashAttention 矩阵乘（计算密集）、decode NPU 专注 KV cache 带宽（内存密集）。hixl 的单边通信原语（put/get）实现零拷贝数据传输——prefill 一侧写完 HBM，直接 RDMA 写到 decode 一侧的 HBM，延迟 393μs/层 × 32 层 = 12.6ms 传输开销。PD 分离后 decode TPOT 从 46ms 降到 8ms（5.75×），1P+4D 配置达到 106 tokens/s。三个核心点：多 prefill 并发时需要 request slot 分区防 KV 写入冲突、non-blocking put 队列有深度限制需定期 flush、KV cache 传输后必须 dummy read 预热 L2 cache（省前 3 step 的 30ms each）。