概述

SGLang 的 Pipeline Parallelism (PP) 模式下，HiCache 负责 KV cache 的异步 prefetch（Host→GPU）和 backup（GPU→Host）。本文从 PP 调度器的外层事件循环出发，追踪 Load 和 Write 的完整时序，揭示 write_ack 比 load_ack 多延迟的根本原因。

一、PP 事件循环

每个 iteration 执行以下 4 步：

1
2
3
4
5

iter=N:
  ① check_hicache_events()      ← 查询 HiCache 异步事件（load_ack / write_ack）
  ② get_next_batch_to_run()     ← 选下一批 batch（prefix match → eviction → load）
  ③ _pp_launch_batch()          ← launch forward
  ④ _pp_process_batch_result()  ← 处理上一批 batch 的结果（insert → write）

核心设计：process_batch_result 处理的是 mbs[next_mb_id]——上一轮 launch 的 batch。同一个 iter 内，调度器同时在处理两个不同 batch 的生命周期阶段。

二、Load 时序（Host→GPU）

触发时机

iter=N 的 get_next_batch_to_run —— prefill 阶段 prefix match 发现 host_hit，发起 load_back。

完整时序

关键：load 是 prefetch——在 forward 之前触发，CUDA copy 和 forward 可以重叠（逐层等待、逐层执行）。

时序图

1
2
3
4
5
6
7
8
9
10
11
12

iter=N:
  ① check_hicache_events()
     └─ 消费更早的 load_ack
  ② get_next_batch_to_run()
     └─ match_prefix → host_hit → load_back() → start_loading()
     └─ CUDA copy 启动，event 入队
  ③ _pp_launch_batch()
     └─ forward 逐层等待 load 完成

iter=N+1:
  ① check_hicache_events()
     └─ loading_check() → event.query()=True → load_ack ✓

延迟：load_ack 比 load 发起晚 1 iter。

三、Write 时序（GPU→Host）

触发时机

iter=N 的 process_batch_result —— 处理上一轮 launch 的 batch 结果，insert 时触发 write_backup。

完整时序

关键：write 是 post-write——必须在 forward 算完、拿到完整 KV cache 后才能触发。

时序图

1
2
3
4
5
6
7
8
9
10
11
12
13

iter=N-1:
  ② get_next_batch_to_run() → 选出 Prefill A
  ③ _pp_launch_batch() → forward(Prefill A)

iter=N:
  ④ _pp_process_batch_result()
     └─ 处理 iter=N-1 的 Prefill A
     └─ insert() → write_backup() → start_writing()
     └─ CUDA copy 启动，event 入队

iter=N+1:
  ① check_hicache_events()
     └─ writing_check() → event.query()=True → write_ack ✓

延迟：write_ack 比 write 发起晚 1 iter，但 write 本身比 forward 晚 1 iter（因为 process_batch_result 处理上一批 batch）。

四、延迟对比

以同一个 Prefill batch 为基准

追踪一个 Prefill batch 从 launch 到 write_ack 的完整生命周期：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

iter=1:
  ② get_next_batch_to_run()
     └─ 选出 Prefill A
     └─ match_prefix → host_hit → load_back() → start_loading()
  ③ _pp_launch_batch() → forward(Prefill A)

iter=2:
  ① check_hicache_events()
     └─ loading_check() → load_ack ✓
        ↑ load_ack: 1 iter delay

iter=3:
  ④ process_batch_result()
     └─ 处理 iter=1 的 Prefill A
     └─ insert() → write_backup() → start_writing()
  ① check_hicache_events()
     └─ (write_ack 还没完成)

iter=4:
  ① check_hicache_events()
     └─ writing_check() → write_ack ✓
        ↑ write_ack: 3 iter delay from Prefill launch

对比表

为什么 write_ack 多延迟？

两个因素叠加：

1. process_batch_result 滞后一轮

它处理 mbs[next_mb_id]（上一轮 launch 的 batch），以 PP2 为例偏移 2 iter。所以 write 比 load 晚 2 iter 才触发。

2. CUDA event query 需要等下一轮 check

无论 load 还是 write，ack 入队后最早等下一轮 check_hicache_events 才能消费。两者各加 1 iter。

综合：

1
2

Load:  iter=1 发起 → iter=1 ack 入队 → iter=2 消费
Write: iter=1 forward → iter=3 process_result 发起 → iter=3 ack 入队 → iter=4 消费

Load 是 prefetch（forward 之前），Write 是 post-write（forward 之后）。这个架构差异决定了 Write 必然比 Load 多延迟。

五、PP 间同步问题

PP0 和 PP1 共享同一个 HiCache 实例（radix tree + host memory）。由于 output relay 延迟，PP0 和 PP1 的 iter 进度存在 1-2 iter 的偏移，需要三层同步机制保证一致性。

5.1 逻辑时钟保证重放一致性

PP1 的 writing_check/loading_check 不再直接消费 ack，而是将事件通过 Gloo P2P 通道 replay 给 PP0。PP0 作为唯一的事件消费者，按逻辑时钟顺序处理所有 ack，确保 PP0 和 PP1 的 radix tree 操作顺序一致。

问题：早期实现中 PP1 的 writing_check() 绕过 check_hicache_events guard，直接消费 write_ack（即 ack theft），导致 PP0 端 pending event 永远无法完成，radix tree 分叉。

修复：将 PP rank 分支移入 writing_check/loading_check 内部，用 PPHiCacheEventsReq 控制请求替代 dict wrapper，强制 PP1 replay 事件给 PP0。

5.2 Count Sync 保证 CP 一致性

PP0 比 PP1 早 1-2 iter 积累 write ack（output relay 延迟），导致 ack_write_queue 积累差异。PR #22878 通过 piggybacking write-ack consumption counts 在 PP ranks 间同步：

1
2
3

iter=N:
  PP0: writing_check() → 消费 3 个 write_ack → count=3
  PP1: 等待 PP0 的 count → count=3 → radix tree 操作对齐

Count sync 确保 PP0 和 PP1 对同一个 radix tree node 的 checkpoint（CP）操作一致，避免一个 stage 认为 node 已 backup、另一个 stage 还在等待的情况。

5.3 PP1 同步消费 PP0 的 ack

PP1 不再独立消费 ack，而是通过同步机制确保 PP0 消费 ack 后，PP1 的 radix tree 状态与 PP0 对齐：

1
2
3
4
5

PP0: writing_check() → event.query()=True → 消费 write_ack
     → radix tree: insert() → finalize() → node 状态更新
     ↓ TP all_reduce sync
PP1: 等待 PP0 完成 → radix tree: 同步执行 finalize()
     → node 状态与 PP0 一致

三层保障：

核心目标：无论 PP0 和 PP1 的 iter 偏移多少，radix tree 的结构和节点状态在两个 stage 上始终保持一致。

六、总结

核心结论：

1. PP 调度器每个 iter 同时做两件事：get_next_batch_to_run 选下一批 batch，process_batch_result 处理上一批 batch 的结果
2. Load 是 prefetch（forward 之前触发），Write 是 post-write（forward 之后触发）
3. process_batch_result 处理 mbs[next_mb_id] 的 iter 偏移是 Write 延迟的根本原因
4. PP 间需要 count sync 补偿 output relay 带来的时序差异

本文基于 SGLang 源码分析，涉及文件：hiradix_cache.py、cache_controller.py。

Pipeline Parallelism（PP）是多卡推理中的核心技术之一。SGLang 的 PP 实现有一套独立的事件循环和调度机制，和普通的 single-batch 路径完全不同。本文将从代码出发，深入分析 SGLang 的 PP 实现。

为什么需要 Pipeline Parallelism

单机推理的瓶颈

单卡推理时，GPU 利用率受限于模型大小和显存。一个 70B 模型单卡放不下，必须拆分到多卡。

模型并行方案对比

SGLang 通常 TP + PP 组合使用：节点内 TP，跨节点 PP。

SGLang PP 架构概览

PP 路径的独立事件循环

SGLang 的调度器有三条事件循环路径：

• event_loop_normal - 普通单 batch 调度
• event_loop_overlap - 计算和通信 overlap 的调度
• event_loop_pp - PP 路径的独立事件循环

PP 路径有自己独立的事件循环，定义在 scheduler_pp_mixin.py 的 SchedulerPPMixin 类中。它和 normal/overlap 路径完全隔离，因为 PP 的通信模式（send/recv proxy tensors 跨 stage）和单机调度逻辑差异太大，硬塞进同一个 event loop 会很脏。

event_loop_pp 核心流程

1
2
3
4
5
6
7
8
9
10
11

# scheduler_pp_mixin.py 第 72-145 行
while True:
    for mb_id in range(pp_loop_size):  # 遍历 microbatch
        1. recv_requests()                                  # 接收第 i 个请求（从上一个 PP stage 或 tokenizer）
        2. _pp_send_pyobj_to_next_stage()                   # 立即把第 i 个请求转发给下一个 PP rank（async，提前发送）
        3. get_next_batch_to_run()                          # 调度当前 microbatch 的 batch
        4. _pp_recv_proxy_tensors()                         # 接收第 i 个 proxy（hidden states，来自上一个 PP stage）
        5. _pp_launch_batch()                               # 执行第 i 个 batch 的 forward
        6. _pp_commit_send_output_work_and_preprocess...()  # 发送+接收 output + 预处理 next_mb_id 槽位的结果
        7. _pp_process_batch_result()                       # 处理 next_mb_id 槽位的 batch result
        8. _pp_send_dict_to_next_stage(msg_type="proxy")    # 发送第 i 个 proxy（hidden states）给下一个 PP stage

**⚠️ step 5 和 step 6 的顺序取决于 pp_async_batch_depth**：

• pp_async_batch_depth > 0：step 6（output 收发+预处理）在 step 5（forward）之前执行，与 forward 重叠，隐藏延迟
• pp_async_batch_depth == 0：step 6（output 收发+预处理）在 step 5（forward）之后执行，串行不重叠

上图展示的是 pp_async_batch_depth == 0 的顺序（step 6 在 step 5 之后）。

Microbatch 索引计算：

核心概念：commit = 等待上一轮的异步操作完成

所有 _pp_commit_comm_work(work) 做的事情就是：

1
2
3
4

def _pp_commit_comm_work(self: Scheduler, work: List[P2PWork]) -> None:
    for p2p_work in work:
        p2p_work.work.wait()  # 阻塞等待 NCCL/Gloo 异步通信完成
    work.clear()

即：commit 是一个同步屏障，确保上一轮发起的异步 send 已经完成，发送缓冲区可以安全复用。这是一种延迟等待模式–发送时不阻塞，到下一轮再等待完成，从而让发送和 GPU 计算重叠。

Microbatch 索引计算：

1
2

next_mb_id = (mb_id + 1) % self.pp_loop_size              # 当前 rank 下一个要处理的 microbatch 槽位
next_first_rank_mb_id = (mb_id + self.pp_size) % self.pp_loop_size  # first rank 对应的"下一轮"槽位

• next_mb_id：当前 rank 下一个要处理的 microbatch 槽位（偏移 1，因为每次迭代处理一个 mb）
• next_first_rank_mb_id：first rank 对应的”下一轮”microbatch 槽位（偏移 pp_size，因为流水线深度–first rank 产出的 batch 要经过 pp_size 个迭代才到 last rank 完成）

步骤与函数的对应关系：

关键设计：

• send req 提前：收到请求后立即转发给下一个 rank（step 2），不等 forward 完成，隐藏 P2P 通信延迟
• output 发送在 _pp_commit_send_output_work_and_preprocess_output_tensors 中处理：output 的发送和接收都在 step 6 中完成，last rank 把 output 发给 rank 0，中间 rank 转发
• **output 处理时机取决于 pp_async_batch_depth**：
  • > 0：output 收发+预处理在 forward 之前执行，与当前 forward 重叠，隐藏延迟
  • == 0：output 收发+预处理在 forward 之后执行，串行不重叠
• 延迟等待异步通信：_pp_commit_comm_work 在下一轮迭代中调用 .wait() 确保上一轮的 async send 已完成，让发送和 GPU 计算重叠
• 每个 microbatch 槽位独立管理状态
• async send + recv 实现计算和通信的 overlap
• 主循环轮转处理每个 microbatch，填充 pipeline bubble

完整流程示例（pp_size=2, pp_async_batch_depth=0）

以 pp_size=2, pp_async_batch_depth=0 为例，pp_loop_size=2，只有槽位 0 和 1：

1
2
3
4
5
6

参数：
  pp_size = 2（PP0 = first rank, PP1 = last rank）
  pp_async_batch_depth = 0
  pp_loop_size = 2（槽位：mb_id ∈ {0, 1}）
  next_mb_id = (mb_id + 1) % 2              # 另一个槽位
  next_first_rank_mb_id = (mb_id + 2) % 2 = mb_id  # 注意：next_first_rank_mb_id == mb_id

单次迭代（mb_id = N）的完整步骤：

步骤 1：恢复上下文

1
2
3
4

running_batch = running_mbs[N]
last_batch = last_mbs[N]
next_first_rank_mb_id = N          # 例如 N=0 时，next_first_rank_mb_id = 0
next_mb_id = (N + 1) % 2           # 例如 N=0 时，next_mb_id = 1（另一个槽位）

步骤 2：接收并处理请求

1
2

recv_reqs = recv_requests()
process_input_requests(recv_reqs)

• PP0：从 tokenizer_manager / detokenizer 收请求（zmq 通道）
• PP1：从 PP0 收转发过来的请求（P2P，含 PPHiCacheEventsReq 控制消息）

步骤 3：HiCache 事件同步 + 转发请求

1
2
3
4
5
6

if enable_hierarchical_cache:
    tree_cache.check_hicache_events()

if not is_last_rank:  # 即 PP0
    _pp_commit_comm_work(send_req_work)     # ← 等待上一轮的 send_req 异步操作完成
    send_req_work = async_send(pp_send_payload) → PP1

commit 含义：上一轮（槽位 (N-1)%2）发起的 send_req_work 可能还没完成，这里阻塞等它完成后，才能安全地复用发送缓冲区发送本轮的请求。

步骤 4：调度决策

1
2
3

mbs[N] = get_next_batch_to_run()
running_mbs[N] = running_batch
cur_batch = mbs[N]

步骤 5：接收 Proxy 张量

1
2
3
4

if cur_batch is not None:
    pp_proxy_tensors = _pp_recv_proxy_tensors()
    # PP0: 直接返回 None（第一个 stage 不需要接收）
    # PP1: 阻塞等待从 PP0 接收 hidden_states + residual (msg_type="proxy")

步骤 6：等待上一轮 Proxy 发送完成

1

_pp_commit_comm_work(send_proxy_work)  # ← 等待上一轮的 proxy 异步发送完成

commit 含义：上一轮（槽位 (N-1)%2）PP0 发起的 proxy 异步发送可能还没完成，这里确保它完成，释放发送缓冲区。

步骤 7：GPU Forward 计算

1
2
3
4
5

if cur_batch:
    result, launch_event = _pp_launch_batch(N, pp_proxy_tensors, ...)
    # 在 forward_stream 上执行 run_batch
    # 记录 CUDA event
    # PP1(last rank): 将 output 张量 push 到 last_rank_comm_queue

步骤 8：发送 Output + 接收 Output + 预处理

因为 pp_async_batch_depth == 0，这一步在 forward 之后执行（无法与 GPU 计算 overlap）。

1
2
3

_pp_commit_send_output_work_and_preprocess_output_tensors(
    next_first_rank_mb_id=N, next_mb_id=(N+1)%2,
)

内部逻辑（_pp_send_recv_and_preprocess_output_tensors）：

1
2
3
4
5
6
7
8
9
10
11
12
13

# 1. 等待上一轮的 output 异步发送完成
_pp_commit_comm_work(send_output_work)

# 2. 发送 output:
# PP1(last rank): 如果 mbs[next_first_rank_mb_id=N] 不为空，
#   从 last_rank_comm_queue 取出 output，异步发送给 PP0 (msg_type="output")
# PP0(非 last rank): 如果有上一轮暂存的 pp_outputs，转发给 PP1

# 3. 接收 output + 预处理:
# 如果 mbs[next_mb_id] 不为空:
#   阻塞接收 output 张量 (msg_type="output")
#   在 copy_stream 上执行 _pp_prep_batch_result()（组装 GenerationBatchResult + D2H 拷贝）
#   记录 d2h_event

commit 含义：上一轮（槽位 (N-1)%2）发起的 output 异步发送可能还没完成，这里确保完成后才能发起新的 output 发送。

步骤 9：后处理 next_mb_id 槽位的 Batch 结果

1
2
3
4
5

if mbs[next_mb_id] is not None:
    d2h_event.synchronize()  # 等待 copy_stream 上的 D2H 拷贝完成
    _pp_process_batch_result(mbs[next_mb_id], next_batch_result)
    # 更新请求状态、判断 EOS、发送给 detokenizer
    last_mbs[next_mb_id] = mbs[next_mb_id]

步骤 10：发送 Proxy 张量（仅 PP0）

1
2
3

if not is_last_rank and cur_batch:
    torch.cuda.current_stream().wait_event(launch_event)  # 等 forward 完成
    send_proxy_work = async_send(hidden_states + residual) → PP1  # msg_type="proxy"

步骤 11：保存状态

1

pp_outputs = next_pp_outputs  # 暂存本轮接收到的 output，供下一轮步骤 8 转发

完整时序（稳态，两个槽位交替）

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29

┌─────────────────────────────────────────────────────────────────┐
│ mb_id=0                                                          │
├─────────────────────────────────────────────────────────────────┤
│ 1. recv_requests + process                                       │
│ 2. commit(send_req_work[上轮mb=1的]) → async_send(reqs)         │
│ 3. get_next_batch_to_run → cur_batch                            │
│ 4. recv_proxy (PP1阻塞等PP0; PP0跳过)                            │
│ 5. GPU forward (forward_stream)                                  │
│ 6. commit(send_output_work[上轮mb=1的])                          │
│    send_output → recv_output → prep_batch_result (copy_stream)  │
│ 7. d2h_event.sync → process_batch_result(mbs[next_mb_id]的结果) │
│ 8. commit(send_proxy_work[上轮mb=1的])                           │
│ 9. PP0: wait(launch_event) → async_send(proxy) → PP1            │
│10. pp_outputs = next_pp_outputs                                  │
├─────────────────────────────────────────────────────────────────┤
│ mb_id=1                                                          │
├─────────────────────────────────────────────────────────────────┤
│ 1. recv_requests + process                                       │
│ 2. commit(send_req_work[上轮mb=0的]) → async_send(reqs)         │
│ 3. get_next_batch_to_run → cur_batch                            │
│ 4. recv_proxy (PP1阻塞等PP0; PP0跳过)                            │
│ 5. GPU forward (forward_stream)                                  │
│ 6. commit(send_output_work[上轮mb=0的])                          │
│    send_output → recv_output → prep_batch_result (copy_stream)  │
│ 7. d2h_event.sync → process_batch_result(mbs[next_mb_id]的结果) │
│ 8. commit(send_proxy_work[上轮mb=0的])                           │
│ 9. PP0: wait(launch_event) → async_send(proxy) → PP1            │
│10. pp_outputs = next_pp_outputs                                  │
└─────────────────────────────────────────────────────────────────┘

PP0 与 PP1 的交互时序（稳态）

1
2
3
4
5
6
7
8
9
10
11
12
13

时间轴 →

PP0 (mb_id=0):
recv_req → commit(send_req) → send_req→PP1 → schedule → [skip proxy]
→ FORWARD → commit(send_output)
→ send_output(转发) → recv_output(从PP1) → prep → process_result
→ commit(send_proxy) → send_proxy→PP1

PP1 (mb_id=0):
recv_req(从PP0) → process → schedule → recv_proxy(阻塞等PP0)
→ FORWARD → commit(send_output)
→ send_output→PP0 → recv_output(从PP0转发) → prep → process_result
→ commit(send_proxy)

关键依赖链：

• PP1 的 recv_proxy 必须等 PP0 的 send_proxy（上一轮的）完成
• PP0 的 recv_output 必须等 PP1 的 send_output 完成
• 因为 pp_async_batch_depth=0，output 的发送/接收在 forward 之后，无法与 GPU 计算重叠

所有 commit 汇总

模式统一：每个 async_send 在下一轮同一类型操作之前被 commit，形成”发起→（做其他事）→等待完成→再发起”的流水线模式。即使 pp_async_batch_depth=0 没有计算-通信 overlap，这种 commit 模式仍然保证了同类型通信操作之间不会冲突。

Microbatch：填充 Pipeline Bubble

问题：Pipeline Bubble

PP2（2 个 stage）如果只用一个 batch，会出现大量空闲：

1
2
3
4

时间 →
Stage 0: [ batch 0 ] [ batch 1 ]
Stage 1:              [ batch 0 ] [ batch 1 ]
              ↑ 空闲     ↑ 空闲

Stage 0 做完 batch 0 后必须等 Stage 1 处理完才能做 batch 1（因为只有一个 batch 在 pipeline 里流转），一半的时间都在空闲。这就是 pipeline bubble。

解法：Microbatch

把多个 batch 同时塞进 pipeline，这些 batch 就叫 microbatch：

1
2
3

时间 →
Stage 0: [ mb 0 ][ mb 1 ][ mb 2 ][ mb 3 ]
Stage 1:         [ mb 0 ][ mb 1 ][ mb 2 ][ mb 3 ]

Stage 0 做完 mb0 不用等，立刻做 mb1；Stage 1 也紧接着处理，bubble 大幅减少。

Microbatch 数量的确定

1
2

# scheduler_pp_mixin.py 第 514 行
self.pp_loop_size = self.pp_size + self.server_args.pp_async_batch_depth

• pp_size = 2（PP2）→ 至少 2 个 microbatch 槽位
• pp_async_batch_depth → 额外的 buffer 深度，进一步隐藏延迟

每个 microbatch 都有独立的状态：

1
2
3

self.mbs = [None] * self.pp_loop_size          # 当前 batch
self.last_mbs = [None] * self.pp_loop_size     # 最近一次处理的 batch
self.running_mbs = [ScheduleBatch(...) for _ in range(self.pp_loop_size)]  # running 状态

Microbatch vs 正常 Batch

Microbatch 和 batch 在数据结构上没有区别（都是 ScheduleBatch），区别只在于 PP 需要多个 batch 同时在不同 stage 中流转，所以给它们编号叫 microbatch。

Microbatch 数量不是越多越好

1

pp_loop_size = pp_size + pp_async_batch_depth

• 太少（= pp_size）：pipeline 容易断流，一个 stage 稍慢就导致下游空闲
• 太多：显存占用线性增长（每个 microbatch 都要有自己的 KV cache + activation），而且可能因为显存压力导致 batch size 被迫缩小，反而吞吐下降

调优经验：

• pp_async_batch_depth = 1~2 通常足够，再深收益递减
• 显存充裕时可以试更大的 depth，但要注意 activation memory 的开销
• 如果 single batch 的 token 数很大（长上下文），microbatch 多了容易 OOM

PP Rank 的请求接收路径

recv_requests() 根据 PP rank 的不同，有两条完全不同的接收路径。接收的请求包括两类：

• 外部 API 请求：tokenize 后的推理请求（TokenizedGenerateReqInput、TokenizedEmbeddingReqInput 等）
• 内部控制请求：PP 同步用的控制消息（如 PPHiCacheEventsReq 等）

PP Rank 0（第一个 stage）

从上游组件直接接收用户请求，走两条 zmq 通道：

1
2
3
4
5

# 1. 从 tokenizer 接收（zmq 通道）
recv_req = self.recv_from_tokenizer.recv_pyobj(zmq.NOBLOCK)

# 2. 从 rpc 接收（zmq 通道）
recv_rpc = self.recv_from_rpc.recv_pyobj(zmq.NOBLOCK)

接收的是 tokenize 后的请求对象，类型为：

• TokenizedGenerateReqInput - 生成请求
• TokenizedEmbeddingReqInput - embedding 请求
• 其他控制类消息（flush、abort 等）

PP Rank > 0（后续 stage）

不直接从 tokenizer 接收，而是从前一个 PP rank P2P 转发过来：

1
2
3
4
5
6
7

recv_reqs = point_to_point_pyobj(
    [],
    self.pp_rank * self.tp_size + dp_offset,  # 本 rank
    self.world_group.cpu_group,
    (self.pp_rank - 1) * self.tp_size + dp_offset,  # src: 前一个 PP rank
    self.pp_rank * self.tp_size + dp_offset,  # dst: 本 rank
)

对应 event_loop_pp 中的流转

1
2
3
4
5
6
7
8
9

# event_loop_pp 第 80-88 行
recv_reqs = self.recv_requests()  # PP0 从 zmq 收，PP1 从 PP0 的 P2P 收
self.process_input_requests(recv_reqs)  # 本地处理

if not self.pp_group.is_last_rank:
    # 步骤 2: send reqs to next PP stage
    self.send_req_work = self._pp_send_pyobj_to_next_stage(  # 转发给下一个 PP rank
        recv_reqs, async_send=True,
    )

流转全图

1
2
3
4
5
6
7

Tokenizer / RPC
     │ zmq
     ▼
PP Rank 0: recv_requests() ── point_to_point_pyobj (async) ──→ PP Rank 1: recv_requests()
     │                              │
     ▼                              ▼
process_input_requests()     process_input_requests()

关键设计：所有 PP rank 收到的内容是一样的（同样的请求列表），这样每个 stage 都能为同一批请求做 get_next_batch_to_run() 调度，保证各 stage 的 batch 一致。

process_input_requests：请求路由器

recv_requests() 收到请求后，process_input_requests() 负责按类型分发到对应的 handler：

1
2
3
4
5
6
7
8
9
10
11
12
13

# scheduler_pp_mixin.py 第 1545-1566 行
def process_input_requests(self, recv_reqs: List):
    now = time.monotonic()
    self.session_controller.maybe_reap(now)  # 1. 清理过期 session
    for recv_req in recv_reqs:
        # 2. 健康检查请求：服务器忙时跳过
        if is_health_check_generate_req(recv_req) and not self.is_fully_idle():
            continue
        # 3. 按类型分发
        output = self._request_dispatcher(recv_req)
        # 4. 有即时返回结果的，回送给 tokenizer/rpc
        if output is not None:
            self.send_to_tokenizer.send_output(output, recv_req)

_request_dispatcher 是一个 TypeBasedDispatcher，根据请求类型路由到不同 handler：

核心推理请求：

控制类请求：

权重管理请求：

一句话总结：process_input_requests 就是一个请求路由器–生成请求进入 waiting_queue 等待后续调度，控制请求立即处理并回送结果，所有 PP rank 都执行同样的分发逻辑保证状态一致。

_pp_recv_proxy_tensors：接收上游 Stage 的中间结果

_pp_recv_proxy_tensors（scheduler_pp_mixin.py 第 993-1004 行）负责从上一个 PP stage 接收 forward 的中间激活值：

1
2
3
4
5
6
7
8
9
10

def _pp_recv_proxy_tensors(self: Scheduler) -> Optional[PPProxyTensors]:
    pp_proxy_tensors = None
    if not self.pp_group.is_first_rank:
        pp_proxy_tensors = PPProxyTensors(
            self._pp_recv_typed_dict(
                expected_kind="proxy",
                all_gather_group=self.attn_tp_group if self.require_attn_tp_allgather else None,
            )
        )
    return pp_proxy_tensors

行为：

• PP Rank 0（first rank）：返回 None，因为它是第一个 stage，没有上一级给它传 hidden states，直接用 embedding 层的输出开始 forward。代码里有 is_first_rank 判断，PP0 不会做任何接收操作。
• PP Rank > 0：从前一个 PP rank 接收 tensor dict（msg_type="proxy"），包装成 PPProxyTensors

Proxy tensors 是什么：

看 profile_and_init_predictor 里的构造（第 609-622 行）就清楚了：

1
2
3
4
5

proxy_tensors = {
    "hidden_states": torch.zeros((seq_len, hidden_size), ...),
    "residual": torch.zeros((seq_len, hidden_size), ...),
}
pp_proxy = PPProxyTensors(proxy_tensors)

就是上一个 PP stage forward 输出的中间激活值：

• hidden_states - 当前层的隐藏状态
• residual - 残差连接

在 event_loop_pp 中的位置：

1
2
3
4
5
6
7

if self.cur_batch:
    pp_proxy_tensors = self._pp_recv_proxy_tensors()  # rank>0 在这里阻塞等 recv
    ...
    # 步骤 5: run_batch → 实际调用 _pp_launch_batch()
    result, self.launch_event = self._pp_launch_batch(
        mb_id, pp_proxy_tensors, ...  # 传入 forward，接着上一个 stage 继续算
    )

对应的发送端：在 event_loop_pp 末尾（第 129-139 行），非 last rank 发送 proxy：

1
2
3
4
5
6
7
8

if not self.pp_group.is_last_rank:
    if self.cur_batch:
        torch.cuda.current_stream().wait_event(self.launch_event)  # 等 forward 完成
        self.send_proxy_work = self._pp_send_dict_to_next_stage(
            result.pp_hidden_states_proxy_tensors.tensors,  # forward 输出的 hidden_states
            async_send=True,
            msg_type="proxy",
        )

一句话总结：_pp_recv_proxy_tensors 就是非首 rank 从前一个 PP stage 接收 forward 中间结果（hidden_states + residual），这样本 stage 的模型层才能接着算。这是 PP 的核心数据流–每个 stage 只有部分层，需要前一个 stage 的输出作为输入。

_pp_recv_dict_from_prev_stage（接收 output）与 _pp_recv_proxy_tensors（接收 proxy）的区别

这两个接收操作虽然底层都走 _pp_recv_typed_dict，但它们在 语义、时机、数据内容、通信方向 上完全不同。

一、核心区别对比

二、在流水线中的位置

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

PP Rank 0              PP Rank 1              PP Rank 2 (Last)
─────────              ──────────              ──────────────
                       _pp_recv_proxy_tensors
                       ← proxy (hidden_states + residual)
                       _pp_launch_batch (forward)
                       _pp_send_dict_to_next_stage
                       → proxy (hidden_states + residual)
                                                ← proxy (hidden_states + residual)
                                                _pp_launch_batch (forward + sample)
                                                _pp_send_dict_to_next_stage
                                                → output (next_token_ids + logprobs)
                       ← output (next_token_ids + logprobs)
                       _pp_recv_dict_from_prev_stage (msg_type="output")
                       _pp_commit_send_output_work_and_preprocess_output_tensors
                       process_batch_result (更新请求状态、流式输出)

三、数据内容详解

Proxy 张量（中间隐藏状态）：

1
2
3
4
5

# 发送端（_pp_launch_batch 之后）
result.pp_hidden_states_proxy_tensors.tensors = {
    "hidden_states": tensor([num_tokens, hidden_size]),  # 模型中间层输出
    "residual": tensor([num_tokens, hidden_size]),       # 残差连接
}

这是模型被纵向切分后，前半部分层的输出。下一个 PP stage 需要它作为输入继续跑后半部分层。没有它，下一个 stage 无法执行 forward。

Output 张量（最终输出）：

1
2
3
4
5
6
7
8
9

# 发送端（_pp_prepare_tensor_dict）
tensor_dict = {
    "next_token_ids": tensor([batch_size]),     # 采样得到的下一个 token
    # 可选：
    "input_token_logprobs": ...,
    "normalized_prompt_logprobs": ...,
    "prefill_top_logprobs": ...,
    ...
}

这是 last rank 完成整个模型 forward + 采样后的最终结果。PP0 需要它来执行后处理（更新请求状态、判断是否结束、发送给 tokenizer 等）。

四、为什么需要分开？

1. 时序不同：proxy 必须在 forward 前到位（是 forward 的输入）；output 可以在 forward 后异步接收（是上一轮的结果，用于 CPU 后处理）。

2. 方向不同：proxy 沿流水线正向流动（rank 0→1→2→…→last）；output 从 last rank 反向回到 rank 0（形成环路）。

3. 同一条 P2P 链路上交错到达：由于 PP 使用环形拓扑（last rank 的 next 就是 rank 0），同一对 rank 之间可能同时有 proxy 和 output 在传输。_pp_recv_typed_dict 的 demux 机制（按 msg_type 分拣）就是为了解决这个问题–收到不匹配类型的消息时暂存到 inbox，等需要时再取出。

4. 大小差异巨大：proxy 张量很大（num_tokens × hidden_size，可能几十 MB）；output 张量很小（batch_size 个 int，几 KB）。这影响通信调度策略。

五、一句话总结

• _pp_recv_proxy_tensors：接收的是 “半成品” –上一个 stage 跑完前半部分模型层后的隐藏状态，当前 stage 要拿它继续跑后半部分。
• _pp_recv_dict_from_prev_stage（output）：接收的是 “成品” –last rank 跑完整个模型后采样出的 token，PP0 拿它做后处理（更新状态、返回用户）。

_pp_launch_batch：执行当前 Microbatch 的 Forward

_pp_launch_batch（scheduler_pp_mixin.py）对应 event_loop_pp 流程中的 step 5: run_batch，负责启动当前 microbatch 的 forward 计算：

1
2
3
4
5
6
7
8
9

# 伪代码示意
def _pp_launch_batch(self, mb_id, pp_proxy_tensors, ...):
    batch = self.mbs[mb_id]
    # 如果有 proxy tensors（非首 rank），用 proxy 的 hidden_states 作为输入
    # 否则用 embedding 输出作为输入
    result = self.worker.forward_batch(batch, pp_proxy_tensors)
    launch_event = torch.cuda.Event()
    launch_event.record()
    return result, launch_event

关键行为：

• PP Rank 0（first rank）：不需要 proxy tensors，直接用 token embeddings 作为 forward 输入
• PP Rank > 0：用 _pp_recv_proxy_tensors() 收到的 hidden_states + residual 作为 forward 输入，接着上一个 stage 继续算
• 返回 result（包含 pp_hidden_states_proxy_tensors，供后续发送给下一个 stage）和 launch_event（用于同步）

在 event_loop_pp 中的调用位置：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17

for mb_id in range(self.pp_loop_size):
    ...
    # 步骤 4: 接收 proxy tensors
    pp_proxy_tensors = self._pp_recv_proxy_tensors()

    # 步骤 5: run_batch → _pp_launch_batch
    result, self.launch_event = self._pp_launch_batch(
        mb_id, pp_proxy_tensors, ...
    )

    # 步骤 7: 发送 proxy tensors 给下一个 stage
    if not self.pp_group.is_last_rank:
        torch.cuda.current_stream().wait_event(self.launch_event)
        self.send_proxy_work = self._pp_send_dict_to_next_stage(
            result.pp_hidden_states_proxy_tensors.tensors,
            async_send=True, msg_type="proxy",
        )

一句话总结：_pp_launch_batch 就是 PP 路径下的 forward 执行器–首 rank 用 embedding 输入，非首 rank 用上游传来的 hidden_states 输入，执行当前 stage 的模型层计算后输出给下一个 stage。

_pp_commit_send_output_work_and_preprocess_output_tensors：Output 的收发与预处理（对应 step 5: recv next_mb_id outputs + preprocess）

这个函数（scheduler_pp_mixin.py 第 854-873 行）对应 event_loop_pp 流程中的 step 5: recv next_mb_id outputs + preprocess，负责处理 next_mb_id 槽位对应的 output tensors 的发送、接收和预处理：

1
2
3
4
5
6
7
8
9
10
11
12
13
14

def _pp_commit_send_output_work_and_preprocess_output_tensors(
    self: Scheduler,
    next_first_rank_mb_id: int,
    next_mb_id: int,
) -> Tuple[PPProxyTensors, GenerationBatchResult, torch.cuda.Event]:
    # 1. 等待上一轮的 output 发送完成
    self._pp_commit_comm_work(work=self.send_output_work)
    # 2. 发送 + 接收 + 预处理
    (next_pp_outputs, next_batch_result, d2h_event, self.send_output_work) = \
        self._pp_send_recv_and_preprocess_output_tensors(
            next_first_rank_mb_id, next_mb_id,
            self.mbs, self.mb_metadata, self.last_rank_comm_queue, self.pp_outputs,
        )
    return next_pp_outputs, next_batch_result, d2h_event

核心逻辑在 _pp_send_recv_and_preprocess_output_tensors（第 1081-1116 行），做了三件事：

1. 发送 output（last rank → rank 0）

1
2
3
4
5
6
7
8
9
10
11
12
13
14

# Last rank: 从 last_rank_comm_queue 取出 forward 结果，发给 rank 0
if self.pp_group.is_last_rank:
    q_event, pp_outputs_to_send = last_rank_comm_queue.popleft()
    torch.cuda.current_stream().wait_event(q_event)
    send_output_work = self._pp_send_dict_to_next_stage(
        pp_outputs_to_send.tensors, msg_type="output"
    )

# 中间 rank: 把从前一个 rank 收到的 output 继续转发给下一个 rank
if not self.pp_group.is_last_rank:
    if pp_outputs:
        send_output_work = self._pp_send_dict_to_next_stage(
            pp_outputs.tensors, msg_type="output"
        )

2. 接收 output（rank 0 从 last rank 收）

1
2

if mbs[next_mb_id] is not None:
    next_pp_outputs = PPProxyTensors(self._pp_recv_dict_from_prev_stage())

3. 预处理：GPU→CPU 拷贝（copy stream 上）

1
2
3
4
5
6
7

with self.copy_stream_ctx:
    self.copy_stream.wait_stream(self.schedule_stream)
    batch_result = self._pp_prep_batch_result(  # 解包 next_token_ids, logprobs
        mbs[next_mb_id], mb_metadata[next_mb_id], next_pp_outputs
    )
    d2h_event = torch.cuda.Event()
    d2h_event.record(torch.cuda.current_stream())  # 记录 D2H 完成事件

Output tensors 里有什么：

看 _pp_prepare_tensor_dict（第 920-933 行）：

1
2
3
4
5

tensor_dict = {
    "next_token_ids": result.next_token_ids,  # 采样出的 token id
}
if batch.return_logprob:
    tensor_dict.update(logprob_dict)  # logprob 相关张量

这是 last rank 采样后的最终结果，不是中间的 hidden states。

在 event_loop_pp 中的调用位置：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

for mb_id in range(self.pp_loop_size):
    ...
    # ① 处理【next_mb_id 槽位】的 output 收发 + 预处理
    if self.server_args.pp_async_batch_depth > 0:
        next_pp_outputs, next_batch_result, d2h_event = \
            self._pp_commit_send_output_work_and_preprocess_output_tensors(
                next_first_rank_mb_id, next_mb_id,
            )

    # ② 执行【当前 microbatch】的 forward
    result, self.launch_event = self._pp_launch_batch(mb_id, ...)

    # ③ async_depth==0 时，forward 后再处理 output
    if self.server_args.pp_async_batch_depth == 0:
        next_pp_outputs, next_batch_result, d2h_event = \
            self._pp_commit_send_output_work_and_preprocess_output_tensors(...)

    # ④ 等 D2H 完成，处理 next_mb_id 槽位的结果
    if self.mbs[next_mb_id] is not None:
        d2h_event.synchronize()
        self._pp_process_batch_result(self.mbs[next_mb_id], next_batch_result)

pp_async_batch_depth 的影响：

一句话总结：_pp_commit_send_output_work_and_preprocess_output_tensors 负责把 last rank 的采样结果（next_token_ids、logprobs）从 GPU 传回 rank 0，并在 rank 0 上做 GPU→CPU 拷贝和结果解包。这是 pipeline 的”最后一公里”–把最终输出送回给 tokenizer/rpc 层。

next_batch_result 是什么：

next_batch_result 是 next_mb_id 槽位对应的 microbatch 的最终输出结果（GenerationBatchResult），包含 next_token_ids 等。这里的逻辑是：当前迭代处理的是 mb_id，但同时要处理 next_mb_id 槽位的结果（因为流水线的延迟，next_mb_id 的 batch 在之前的迭代中已经完成了 GPU forward，output 在这一轮才到达）。

⚠️ 注意：next_mb_id 不是”上一个 microbatch”

next_mb_id = (mb_id + 1) % pp_loop_size，从当前 stage 的视角看，它是下一个要处理结果的槽位。但从前一个 PP stage 的视角看，这个槽位对应的是它们刚刚处理完、准备发给我们的下一个 microbatch。

之所以叫 next_mb_id，是因为从前一个 stage 的角度，它是”下一个要发给我们的 microbatch”。不是”上一个”。

简单记：next_mb_id = 从前一个 PP stage 传过来的下一个 microbatch 槽位。

Output 与 Proxy 的发送完全分开：

_pp_commit_send_output_work_and_preprocess_output_tensors 只发 output，不发 proxy。

• output（next_token_ids）：通过 _pp_send_output_to_next_stage 发送（msg_type="output"），方向是 last rank → rank 0（环形）
• proxy（hidden_states + residual）：在 event_loop_pp 循环末尾单独发送（msg_type="proxy"），方向是 rank k → rank k+1（正向）

两者完全独立，不会混在一起。

_pp_commit_comm_work：延迟等待异步通信完成

1
2
3
4

def _pp_commit_comm_work(self: Scheduler, work: List[P2PWork]) -> None:
    for p2p_work in work:
        p2p_work.work.wait()
    work.clear()

作用：等待上一轮的异步通信完成。

因为 _pp_send_pyobj_to_next_stage、_pp_send_dict_to_next_stage 都用了 async_send=True，返回的是一个 P2PWork 列表（底层是 dist.isend 的 future）。_pp_commit_comm_work 就是在下一轮迭代中调用 .wait() 来确保上一轮的发送已经完成，然后清空 work 列表。

这是一种延迟等待的模式：发送时不阻塞，到下一轮再等待完成，从而让发送和 GPU 计算重叠。

_pp_commit_comm_work 的调用位置：

在 event_loop_pp 中，_pp_commit_comm_work 被调用于两个地方：

1
2
3
4
5
6
7
8
9
10

for mb_id in range(self.pp_loop_size):
    ...
    # ① 等待上一轮的 output 发送完成
    self._pp_commit_comm_work(work=self.send_output_work)

    # ② 执行 forward
    result, self.launch_event = self._pp_launch_batch(mb_id, ...)

    # ③ 等待上一轮的 proxy 发送完成
    self._pp_commit_comm_work(work=self.send_proxy_work)

• self._pp_commit_comm_work(work=self.send_output_work) - 等待上一轮 output tensor 的 async send 完成
• self._pp_commit_comm_work(work=self.send_proxy_work) - 等待上一轮 proxy tensor 的 async send 完成

两者都是延迟等待，让上一轮的通信和当前轮的 GPU 计算重叠。

数据流全图：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

Last Rank          中间 Rank              Rank 0
─────────          ──────────             ──────
forward 完成
  │
  ├─ next_token_ids 入 comm_queue
  │
  └─ send "output" ──────────→ recv & 转发 "output" ──────────→ recv "output"
                                    │
                              copy_stream:
                              prep_batch_result (D2H)
                                    │
                              d2h_event.sync()
                                    │
                              process_batch_result()
                              (更新请求状态、流式输出)

PP 路径中的 HiCache（分层 KV Cache）

HiCache（Hierarchical Cache）是 SGLang 的多级 KV Cache 管理机制，把 KV cache 存在 L3（CPU 内存或磁盘），需要时再加载到 GPU。在 PP scheduler 中，HiCache 不在 event_loop_pp 本身，而是在 event loop 调用的子流程中发挥作用。

HiCache 在 PP 调度中的四个阶段

1. 请求入队时 - 预取（Prefetch）

1
2
3
4
5
6
7
8
9

# process_input_requests → handle_generate_request → _add_request_to_queue → _prefetch_kvcache
# 第 1868-1888 行
def _prefetch_kvcache(self, req: Req):
    if self.enable_hicache_storage:
        req.init_next_round_input(self.tree_cache, cow_mamba=False)
        # 从外部存储（L3）异步预取 KV cache 到 GPU
        self.tree_cache.prefetch_from_storage(
            req.rid, last_host_node, new_input_tokens, last_hash, prefix_keys
        )

时机：请求刚进入 waiting_queue 时，立即发起异步预取，利用等待调度的时间把 KV cache 从外部存储搬到 GPU。

2. 调度组 batch 时 - 检查预取进度 + 加载

get_next_batch_to_run → _get_new_batch_prefill_raw 中有三处：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

# 2a. 检查异步事件完成（第 2282-2283 行）
if self.enable_hierarchical_cache:
    self.tree_cache.check_hicache_events()  # 推进异步 D2H/H2D 事件

# 2b. 逐请求检查预取是否完成（第 2390-2398 行）
for req in self.waiting_queue:
    ...
    if self.enable_hicache_storage:
        prefetch_done = self.tree_cache.check_prefetch_progress(req.rid)
        if not prefetch_done:
            continue  # 预取未完成，跳过这个请求，不放进本轮 batch
        # 预取完成，拿到从 L3 加载了多少 token 的 KV
        req.storage_hit_length = self.tree_cache.pop_prefetch_loaded_tokens(req.rid)

# 2c. batch 创建后，准备从 host 加载到 GPU（第 2459-2463 行）
if self.enable_hierarchical_cache:
    new_batch.hicache_consumer_index = (
        self.tree_cache.ready_to_load_host_cache()  # 标记哪些 KV 需要从 host→GPU
    )

3. 请求被 abort 时 - 释放预取资源

1
2
3

# 在多个 abort 路径中（第 1965、1997、3146 行）
if self.enable_hicache_storage:
    self.tree_cache.release_aborted_request(req.rid)  # 取消预取，释放资源

4. flush cache 时 - 等待异步操作完成

1
2

# 第 2877 行
# HiCache: in-flight async ops (GPU↔️Host↔️L3) must drain before flush

check_hicache_events()：异步事件推进器

check_hicache_events() 是 HiCache 的核心–每轮调度前轮询一次，把已完成的 GPU↔️Host↔️L3 异步拷贝操作确认掉（解锁节点、释放资源、推进状态）。

1
2
3
4
5
6

# tree_cache.py 第 1136-1144 行
def check_hicache_events(self):
    self.writing_check()        # ① 检查 GPU→Host 写入完成
    self.loading_check()        # ② 检查 Host→GPU 加载完成
    if self.enable_storage:
        self.drain_storage_control_queues()  # ③ 处理 L3 存储控制队列

① writing_check()（第 676-717 行）- GPU→Host 写入

KV cache 从 GPU 写到 Host 内存（备份），是异步的。这个函数检查哪些写入完成了：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15

# 遍历 ack_write_queue 里的 CUDA event
for _, finish_event, ack_list in self.cache_controller.ack_write_queue:
    if not finish_event.query():  # CUDA event 未完成，停止检查
        break
    finish_count += 1

# TP 多卡时取 min，确保所有 TP rank 都完成
torch.distributed.all_reduce(queue_size, op=ReduceOp.MIN, group=self.tp_group)

# 对已完成的写入：
for ack_id in ack_list:
    backuped_node = self.ongoing_write_through.pop(ack_id)  # 从进行中移除
    self.dec_lock_ref(backuped_node)  # 解锁，允许被驱逐
    if self.enable_storage:
        self.write_backup_storage(backuped_node)  # 继续写到 L3 存储

作用：GPU 上的 KV cache 备份到 Host 完成后，解锁节点（可以被驱逐释放 GPU 显存），如果开了 L3 存储还会继续往 L3 写。

② loading_check()（第 719-732 行）- Host→GPU 加载

KV cache 从 Host 内存加载回 GPU，也是异步的：

1
2
3
4
5
6
7
8
9

for _, finish_event, ack_list in self.cache_controller.ack_load_queue:
    if not finish_event.query():  # CUDA event 未完成
        break
    finish_count += 1
    for ack_id in ack_list:
        end_node = self.ongoing_load_back.pop(ack_id)  # 从进行中移除
        self.dec_lock_ref(end_node)  # 解锁节点

del self.cache_controller.ack_load_queue[:finish_count]

作用：之前从 Host 加载回 GPU 的 KV cache 完成后，解锁节点，这些 KV cache 就可以被 forward 使用了。

③ drain_storage_control_queues()（第 1146-1172 行）- L3 存储控制

处理三个队列：

1
2
3
4
5
6
7
8

cc = self.cache_controller
qsizes = torch.tensor([
    cc.prefetch_revoke_queue.qsize(),   # 取消预取
    cc.ack_backup_queue.qsize(),        # L3 备份确认
    cc.host_mem_release_queue.qsize(),  # Host 内存释放
])
# TP 同步后统一处理
torch.distributed.all_reduce(qsizes, op=ReduceOp.MIN, group=self.tp_group)

HiCache 三级缓存的数据流

1
2
3
4
5
6
7
8

check_hicache_events() 推进的异步操作
════════════════════════════════════

 GPU (L1) ◄────── loading_check() ────── Host (L2) ◄──── drain (prefetch from L3)
   │  ▲
   └──── writing_check() ────────────────────┘
   │
   └──── drain (backup to L3) ────→ L3 Storage

HiCache 数据流

1
2
3
4
5
6
7
8
9
10

请求进入 waiting_queue
  │
  ├─ _prefetch_kvcache() → 异步预取 KV cache 从 L3 到 GPU
  │
  ├─ get_next_batch_to_run()
  │    ├─ check_hicache_events() → 推进 D2H/H2D 事件
  │    ├─ check_prefetch_progress() → 检查预取是否完成
  │    └─ ready_to_load_host_cache() → 标记 host→GPU 加载
  │
  └─ 预取完成 → 请求进入 batch → forward

一句话总结：HiCache 在 PP 调度中贯穿请求的完整生命周期–入队时异步预取、调度时检查进度、abort 时释放资源、flush 时等待完成。核心思想是利用请求在 waiting_queue 中的等待时间，提前把 KV cache 从 L3 搬到 GPU，隐藏加载延迟。

Host 内存驱逐：LRU 最小堆

当 host 内存不足时，evict_host() 使用 LRU 最小堆来决定驱逐哪些节点。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

# 1. 收集候选节点
leaves = list(self.evictable_host_leaves)

# evictable_host_leaves 中的节点满足：
# - node.evicted = True（GPU 上的 value 已被驱逐，只剩 host_value）
# - node.lock_ref == 0（没有被锁定）
# - 没有 evicted 的子节点（是叶子）

# 2. 构建三元组堆
eviction_heap = [
    (
        self.eviction_strategy.get_priority(node),  # 第一排序：LRU 优先级
        self._evict_tie_breaker(node),              # 第二排序：平局打破
        node,                                        # 节点本身
    )
    for node in leaves
]

# 3. 建堆 O(n)
heapq.heapify(eviction_heap)

# 4. 循环弹出堆顶，逐个驱逐
while num_evicted < num_tokens and len(eviction_heap):
    _priority, _tb, x = heapq.heappop(eviction_heap)
    num_evicted += self.cache_controller.evict_host(x.host_value)
    x.parent.children.pop(key)  # 从 radix tree 彻底删除

三元组比较：

PP 场景下的 tie-breaker 设计：

• PP=1：用 node.id 作为 tie-breaker，简单高效
• PP>1：用内容 hash（node.hash_value）作为 tie-breaker，保证不同 rank 对相同节点计算出相同的优先级顺序，确保驱逐一致性

一句话总结：把所有可驱逐的 host 叶子节点放进 LRU 最小堆，循环弹出堆顶（最久没访问的）逐个驱逐，直到释放够 num_tokens 个 token 的 host 内存。tie-breaker 保证 PP 多 rank 场景下驱逐顺序一致。

PP + HiCache 的 Tree Cache 同步问题（PR #22878）

PP + HiCache 组合使用时存在一个经典问题：各 PP rank 的 writing_check() 独立消费 GPU→Host 的 write-ack，由于各 rank 的 batch 内容和异步进度不同，消费数量可能不一致，导致各 rank 的 radix tree 状态出现分歧（哪些节点已备份、哪些可以驱逐不一致）。

解决方案：上游 rank 的消费数作为下游的预算上限，通过 piggyback 在已有的 PP 请求转发中传递。

1. writing_check() 末尾，PP rank 0 记录本轮消费数

1
2
3

# hiradix_cache.py
if self.pp_size > 1 and self.pp_rank < self.pp_size - 1:
    self._pp_last_write_ack_consumed += consumed_count

2. 通过 recv_reqs 捎带传给下游（scheduler_pp_mixin.py）

1
2
3
4
5
6
7
8
9

# 非 last rank 发送请求时，把 ack_count 捎带上
if self.enable_hicache_storage:
    ack_count = self.tree_cache.get_pp_last_write_ack_consumed()
    if ack_count > 0:
        pp_send_payload = {
            "recv_reqs": recv_reqs,          # 原始请求
            "pp_write_ack_count": ack_count, # 捎带的 ack 消费数
        }
        self.send_req_work = self._pp_send_pyobj_to_next_stage(pp_send_payload)

3. 下游 rank 接收并设置预算（scheduler.py）

1
2
3
4
5

# recv_requests() 中，PP rank > 0 解包捎带的 ack_count
if self.pp_rank > 0 and isinstance(recv_reqs, dict) and "pp_write_ack_count" in recv_reqs:
    pp_write_ack_count = recv_reqs["pp_write_ack_count"]
    recv_reqs = recv_reqs["recv_reqs"]
    self.tree_cache.set_pp_upstream_write_ack_count(pp_write_ack_count)

4. 下游 rank 的 writing_check 受限于上游预算

1
2
3
4

# writing_check() 中，下游 rank 不能消费超过上游的数量
if self.pp_rank > 0 and self._pp_write_ack_budget_from_upstream is not None:
    finish_count = min(finish_count, self._pp_write_ack_budget_from_upstream)
    self._pp_write_ack_budget_from_upstream -= finish_count

数据流：

1
2
3
4
5
6
7
8
9
10
11
12
13
14

PP Rank 0                      PP Rank 1
─────────                      ─────────
writing_check()
consumed 3 acks ──→ _pp_last_write_ack_consumed = 3
  │
  ├─ send_pyobj({recv_reqs: [...], pp_write_ack_count: 3})
  │
  └──────────→ recv_requests() 解包
               _pp_write_ack_budget = 3
               │
               writing_check()
               本地完成了 5 个
               但 min(5, 3) = 3 ← 只消费 3 个
               保持和 Rank 0 一致

设计亮点：不新增通信通道，而是 piggyback 在已有的 _pp_send_pyobj_to_next_stage 上，仅 +61 行代码解决问题。

一句话总结：把上游 PP rank 的 write-ack 消费数量捎带在 recv_reqs 的 pyobj 通信上传给下游，下游以此为上限限制自己的消费，确保所有 PP rank 的 radix tree 状态变更保持一致，防止缓存驱逐决策不同步。

PP + HiCache 一致性修复进度（Issue #22607）：

Channel A：Host Tree Event Replay

Channel B（write-back count sync）只限制 PP1 不超前于 PP0，但不能保证两个 rank 在每一轮都消费相同数量。如果 PP1 的 CUDA event 本地没完成，min(local_ready=0, budget=3) = 0，PP1 这轮仍然消费 0。这意味着 host tree 的状态变更时机在两个 rank 之间可以有偏差，一旦偏差积累，就会导致 host tree 结构不同。

Channel A 的思路：不指望两个 rank 独立做出相同决策，而是让 PP0 做决策，PP1 回放 PP0 的结果。

具体场景：PP2 下的 host tree 分歧

假设 PP2 配置：PP0 管 layer 0-29，PP1 管 layer 30-59，两个 rank 各自独立维护一棵 radix tree。

第 1 步：正常运行，两个 rank 一致

请求 A、B、C 进来，两个 rank 收到相同请求、做相同调度。节点 X（系统提示词的 KV cache）在两个 rank 的 GPU 上都存在：

1
2
3

PP0 radix tree: root → X(value=GPU, host_value=None, hit_count=5)
PP1 radix tree: root → X(value=GPU, host_value=None, hit_count=5)
✅ 一致

第 2 步：hit_count 达到阈值，触发 write_backup

1
2
3

PP0: write_backup(X) → CUDA 异步拷贝启动 → X 加入 ongoing_write_through
PP1: write_backup(X) → CUDA 异步拷贝启动 → X 加入 ongoing_write_through
✅ 一致

第 3 步：writing_check 完成时间分歧（关键！）

1
2
3
4
5
6
7

PP0: CUDA event 已完成 → finish_count=1 → X.lock_ref=0
     X.backuped=True, X.host_value 有值
     X 进入 evictable_leaves（可被驱逐）

PP1: CUDA event 未完成 → finish_count=0 → X.lock_ref=1 仍然锁定
     X 不在 evictable_leaves 中
⚠️ 开始分歧

Channel B 此时 PP0 把 consumed_count=1 传给 PP1，PP1 的 budget 变成 1。但 PP1 本地 finish_count=0，min(0, 1)=0，这轮 PP1 仍然消费 0。Channel B 只保证 PP1 不会超过 PP0，不能强迫 PP1 追上 PP0。

第 4 步：GPU 内存压力 → PP0 驱逐 X

1
2
3
4
5
6
7
8

PP0: X 在 evictable_leaves 中（lock_ref=0, backuped=True）
     → _evict_backuped(X) → X.value=None → X.evicted=True
     X 现在只存在于 Host 内存

PP1: X 不在 evictable_leaves（lock_ref=1, 还在 ongoing_write_through）
     → 驱逐了另一个节点 Y
     X 仍在 GPU 上
⚠️ 分歧加大

第 5 步：PP0 的 Host 内存也满了 → evict_host 删除 X

1
2
3
4
5
6
7
8

PP0: write_backup(Z) → host 内存不足 → evict_host()
     → X 是 evictable_host_leaves 中优先级最低的
     → parent.children.pop(X) → X 从 host tree 彻底删除！

PP1: host 内存没有压力（X 没被驱逐到 host）
     → 不触发 evict_host
     → X 仍在 radix tree 中
❌ 两个 rank 的 tree 结构不同了

第 6 步：新请求匹配前缀 → crash

1
2
3
4
5
6
7

PP0: _match_prefix_helper → 找不到 X → device_indices=[], host_hit_length=0
PP1: _match_prefix_helper → 匹配到 X → device_indices=X.value, host_hit_length=0

PP0: prefix_indices 长度 = 0 → extend_input_len = 完整长度
PP1: prefix_indices 长度 = len(X) → extend_input_len = 更短

两个 rank 对同一个请求算出不同的 extend_input_len → batch 形状不一致 → crash

Channel A 如何解决

1
2
3
4
5
6
7
8
9
10
11
12
13

PP0 (权威源)              PP1 (回放)
───────────────           ─────────
writing_check: X 完成
emit FINALIZE(hash=X_hash) ─────→ replay FINALIZE: 标记 X 为 backuped
（不管本地 CUDA event 是否完成）

evict: 驱逐 X
emit EVICT(hash=X_hash) ─────→ replay EVICT: 也驱逐 X

evict_host: 删除 X
emit REVOKE(hash=X_hash) ─────→ replay REVOKE: 也从 host tree 删除 X

PP1 的 host tree 始终和 PP0 一致 ✅

三个 Channel 的分工

1
2
3
4
5
6
7
8
9
10
11
12
13
14

问题链条：
writing_check 时机不同 → lock_ref 不同 → evictable 集合不同
→ 驱逐不同 → host tree 结构不同 → match_prefix 不同 → crash

Channel B (count sync): ────────┐
 限制 PP1 不超前于 PP0          │ 减小分歧窗口，
 减少 lock_ref 偏差             │ 但不能消除
                                │
逻辑时钟 + tie-breaker: ────────┤ 保证驱逐顺序在"相同输入"下一致
LRU 顺序确定性                  │ 但前提是输入（evictable 集合）相同
                                │
Channel A (event replay): ────────┘ 根本解决：PP1 不独立决策
PP0 的 host tree 变更直接回放到 PP1
host tree 结构保证一致

一句话总结：Channel B + 逻辑时钟是”努力让两个 rank 独立做出相同决策”，但由于异步操作的固有非确定性，无法 100% 保证。Channel A 是”放弃独立决策，直接同步结果”，从根本上消除分歧。

PP 路径的调度细节

get_next_batch_to_run

这个调度器在 PP 路径下要考虑的东西比 normal 多得多：

1. 当前 stage 的显存 - 不能 OOM
2. Microbatch 的依赖关系 - 上下游 stage 的进度
3. Pipeline 的吞吐 - 尽量让每个 stage 都有活干

Async Send/Recv 的 Overlap

第 2 步 send reqs 和 recv proxy tensors 是 async 的，理想情况下计算和通信能 overlap。但如果 recv 阻塞了，整个 loop 就卡住了。

1
2

Stage 0: [compute mb0] → [send hidden to S1] → [recv hidden from S1] → [compute mb1]
Stage 1: [recv hidden from S0] → [compute mb0] → [send hidden to S0] → [recv hidden from S0]

通信和计算的重叠程度决定了 pipeline 的效率。

PP 路径 vs 非 PP 路径对比

非 PP 路径（event_loop_normal / event_loop_overlap）

1
2

时间 →
GPU: [ batch 0 forward ] → [ batch 1 forward ] → [ batch 2 forward ]

调度一个 batch，forward 一个 batch，处理结果，再调度下一个。同一时刻只有一个 batch。

PP 路径（event_loop_pp）

1
2
3

时间 →
Stage 0: [ mb 0 ][ mb 1 ][ mb 2 ][ mb 3 ]
Stage 1:         [ mb 0 ][ mb 1 ][ mb 2 ][ mb 3 ]

多个 microbatch 同时在不同 stage 中流转，主循环轮转处理每个 microbatch。

总结

SGLang 的 PP 实现有几个关键设计：

1. 独立事件循环 - PP 路径和 normal/overlap 路径完全隔离，避免复杂度交叉
2. Microbatch 填充 bubble - 通过 pp_loop_size = pp_size + async_depth 控制 pipeline 中的并发度
3. Async 通信 - send/recv 异步化，尽量 overlap 计算和通信
4. 独立状态管理 - 每个 microbatch 槽位有独立的 running_batch、last_batch 等状态

PP 的核心挑战始终是 如何减少 bubble、提升 GPU 利用率，同时控制显存开销。Microbatch 数量的调优是 PP 性能调优的关键。

参考

• SGLang 源码：scheduler_pp_mixin.py
• Pipeline Parallelism 论文：GPipe: Efficient Training of Giant Neural Networks using Pipeline Parallelism

FlashMLA Sparse Decode 完整计算过程详解

本文以一个最小化的具体数值例子，逐步展示 FlashMLA Sparse Decode 的每一个计算步骤，并标注每一步与 Hopper (SM90) / Blackwell (SM100/GB200) 硬件特性的关系。

0. 问题设定

参考配置

本文的模型参数来自 DeepSeek-V3.2 的官方 Hugging Face 配置:

• 配置文件: https://huggingface.co/deepseek-ai/DeepSeek-V3.2/blob/main/config.json

以下是从 config.json 中提取的注意力相关参数，以及它们在 FlashMLA 内核中的对应关系:

MLA (Multi-head Latent Attention) 的核心思想：不存储 128 个 head × 128 维的完整 KV cache (= 16384 维)，而是存储一个 512 维的 LoRA 压缩向量 + 64 维的 RoPE 向量 (= 576 维)。这将 KV cache 压缩了 28 倍 (16384 → 576)。FlashMLA 的所有计算都在这个压缩空间中进行。

本文使用的示例参数

我们使用与 config.json 一致的模型配置，但将 batch size 和 topk 缩小以便手算:

模型配置 (来自 config.json):

1
2
3
4
5
6

d_qk = 576 = kv_lora_rank(512) + qk_rope_head_dim(64)
d_v = 512 = kv_lora_rank
d_nope = 512 = kv_lora_rank (NoPE 部分，做 FP8 量化)
d_rope = 64 = qk_rope_head_dim (RoPE 部分，保持 BF16 不量化)
h_q = 128 = num_attention_heads
h_kv = 1 (MLA 压缩后：所有 head 共享 1 份 KV cache)

示例 Batch 配置 (缩小规模以便展示):

1
2
3
4

b = 2 (batch size)
s_q = 1 (每个请求只 decode 1 个 token)
topk = 128 (原始 config 中 index_topk=2048, 此处缩小便于计算)
page_block_size = 64

硬件配置 (H800 / SM90):

1

num_sms = 132 (SM 数量)

注意: 实际部署中 topk=2048，h_q=128，这意味着每个 query token 需要从 KV cache 中 gather 2048 个 token 并做注意力计算。本文为了展示清晰，将 topk 缩小为 128。

为了简化展示，我们只跟踪 2 个 token（token 0 和 token 1）在一个 head 上的完整计算，其余的计算过程完全相同。

第 1 步：Python 层入口

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

sched_meta, _ = flash_mla.get_mla_metadata() # 创建空的调度元数据

out, lse = flash_mla.flash_mla_with_kvcache(
    q,
    # shape: [batch_size, seq_len_q, num_heads_q, head_dim]
    # = [2, 1, 128, 576] dtype=BF16
    # - batch_size=2: 同时处理 2 个请求
    # - seq_len_q=1: decode 阶段每次只生成 1 个新 token
    # - num_heads_q=128: 来自 config.json 的 num_attention_heads=128
    # - head_dim=576: = kv_lora_rank(512) + qk_rope_head_dim(64)
    #   即每个 head 的 Q 向量是 576 维 (512 维 NoPE + 64 维 RoPE)
    
    k_cache,
    # shape: [num_blocks, page_block_size, num_heads_k, bytes_per_token]
    # = [num_blocks, 64, 1, 656] dtype=FP8
    # - num_blocks: KV cache 的总页数 (由推理框架的 paged attention 管理)
    # - page_block_size=64: 每个页面存放 64 个 token 的 KV 向量
    # - num_heads_k=1: MLA 的核心 — 所有 128 个 head 共享 1 份压缩后的 KV cache
    # - bytes_per_token=656: 每个 token 的 FP8 KV cache 大小:
    #   512 字节 (NoPE, FP8_E4M3) + 16 字节 (4×float32 缩放因子) + 128 字节 (RoPE, BF16)
    
    block_table=None,        # 稀疏模式不需要 (dense 模式才需要页表做地址翻译)
    cache_seqlens=None,      # 稀疏模式不需要 (dense 模式才需要知道每个序列的实际长度)
    head_dim_v=512,          # V 的 head 维度 = kv_lora_rank = 512
    
    tile_scheduler_metadata=sched_meta,
    is_fp8_kvcache=True,
    
    indices=indices,
    # shape: [batch_size, seq_len_q, topk]
    # = [2, 1, 128] dtype=int32
    # - batch_size=2: 对应 2 个请求
    # - seq_len_q=1: 每个新 token 各有自己的 top-k 索引列表
    # - topk=128: 每个 query token 关注的 KV token 数量 (实际部署中为 index_topk=2048)
    # 值的含义：indices[i][j][k] = global_token_index
    # 例：indices[0][0][3] = 47 表示 batch 0 的 query 的第 4 个关注 token
    #     位于第 0 个 page block (47÷64=0) 的第 47 个位置 (47%64=47)
)

# 返回值:
# out: [batch_size, seq_len_q, num_heads_q, head_dim_v] = [2, 1, 128, 512] BF16
# lse: [batch_size, num_heads_q, seq_len_q] = [2, 128, 1] FP32

关于 dense vs sparse 的选择: FlashMLA 本身不决定使用哪条路径。当调用方（vLLM/SGLang 等推理框架）传入 indices 参数时走 sparse 路径；不传 indices 而传入 block_table + cache_seqlens 时走 dense 路径。当上下文长度 ≤ index_topk(2048) 时，top-k 等于全部 token 数，稀疏无意义，框架层面应直接调用 dense 路径。

关键路由逻辑 (flash_mla_interface.py 第 151 行):

1
2
3

if topk is not None:
    # topk 不为 None → 走稀疏路径
    out, lse, ... = flash_mla_cuda.sparse_decode_fwd(...)

→ 无硬件特性依赖，纯 Python 路由

第 2 步：C++ 接口层 — 输入验证与参数准备

文件: csrc/api/sparse_decode.h

2.1 架构检测与实现选择

1
2
3
4
5
6
7
8
9
10
11

Arch arch = Arch();
// arch.major=9, arch.minor=0 → SM90a (Hopper)
// 或 arch.major=10 → SM100f (Blackwell)

if (arch.is_sm100f()) {
    if (h_q == 64) impl = new Decode_Sm100_Head64_Impl();
    // SM100 使用 TMEM + UTCMMA
} else if (arch.is_sm90a()) {
    impl = new Decode_Sm90_Impl();
    // SM90 使用 TMA + WGMMA
}

🔧 硬件特性: cudaGetDeviceProperties 获取 compute capability。SM90 = Hopper, SM100 = Blackwell。

2.2 FP8 KV Cache 形状验证

1
2
3
4
5

// V3.2: 每个 token 占 656 字节
// 656 = 512 (NoPE FP8) + 16 (scales FP32) + 128 (RoPE BF16)
int bytes_per_token = 512 + 4*sizeof(float) + 64*sizeof(nv_bfloat16);
KU_CHECK_SHAPE(kv, num_blocks, page_block_size, h_kv, bytes_per_token);
// 即 kv shape = [num_blocks, 64, 1, 656]

656 字节的内存布局：

1
2
3
4
5
6

┌─────────────────────────────────────────────────────────┐
│ 偏移 0-511:   512 × FP8_E4M3 NoPE 部分 (量化)           │
│ 偏移 512-527: 4 × float32 缩放因子 (每 128 个一组)        │
│ 偏移 528-655: 64 × BF16 RoPE 部分 (不量化)               │
│ 总计 656 字节                                           │
└─────────────────────────────────────────────────────────┘

🔧 硬件特性: FP8 (E4M3) 是 Hopper/Blackwell 原生支持的数据类型。Hopper 的 Tensor Core 可以直接做 FP8 矩阵乘法，但 FlashMLA 选择先反量化为 BF16 再做 WGMMA，以获得更高精度。

2.3 填充参数结构体

这一步是”翻译层”— 把 PyTorch tensor 的元信息（shape、stride、data_ptr）和标量参数，打包成 CUDA kernel 能直接使用的 C struct。kernel 内部只看这个 struct，不再接触 PyTorch 对象。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42

SparseAttnDecodeParams params = {
    // ===== 基础维度 (直接从输入 tensor 的 shape 读取) =====
    .b = 2,                    // q.size(0) = batch_size，同时处理 2 个请求
    .s_q = 1,                  // q.size(1) = seq_len_q，decode 阶段每次只生成 1 个新 token
    .h_q = 128,                // q.size(2) = num_attention_heads，来自 config.json
    .h_kv = 1,                 // kv.size(2)，MLA 压缩后 128 个 head 共享 1 份 KV cache
                               // 注意：计算时会被 broadcast 到 h_q=128
    .d_qk = 576,               // q.size(3) = head_dim = kv_lora_rank(512) + qk_rope_head_dim(64)
    .d_v = 512,                // head_dim_v 参数 = kv_lora_rank
    
    // ===== Softmax 缩放因子 =====
    // 标准 Scaled Dot-Product Attention: softmax(Q·K^T / sqrt(d)) · V
    .sm_scale = 1.0 / sqrt(576) = 0.04167,
    // 即公式中的 1/sqrt(d)，防止点积值过大导致 softmax 饱和
    
    .sm_scale_div_log2 = 0.04167 * log2(e) = 0.04167 * 1.4427 = 0.06010,
    // 性能优化的预计算。kernel 中用 exp2f() 代替 expf() 计算 softmax:
    // expf(x * sm_scale) = exp2f(x * sm_scale * log2(e))
    //                    = exp2f(x * sm_scale_div_log2)
    // GPU 上 exp2f() 比 expf() 快约 2 倍 (SFU 原生支持 base-2 指数运算)，
    // 所以预乘好 log2(e)，避免 kernel 里每个元素都重复乘一次
    
    // ===== 稀疏注意力参数 =====
    .topk = 128,               // indices.size(2)，每个 query 关注的 KV token 数量
    .model_type = ModelType::V32,
    // 由 d_qk 决定：576 → V32 (DeepSeek-V3/V3.2), 512 → MODEL1
    // 不同 model_type 的 FP8 KV cache 布局和缩放因子格式不同
    
    // ===== 指针 (从 tensor.data_ptr() 获取) =====
    .q = (bf16*)q.data_ptr(),           // Q 矩阵的 GPU 内存地址
    .kv = (bf16*)kv.data_ptr(),         // FP8 KV cache 的 GPU 内存地址
    .indices = (int*)indices.data_ptr(), // top-k 索引数组的 GPU 内存地址
    .out = (bf16*)out.data_ptr(),       // 输出矩阵的 GPU 内存地址
    .lse = (float*)lse.data_ptr(),      // log-sum-exp 的 GPU 内存地址
    
    // ===== Stride (从 tensor.stride() 获取，单位是元素数) =====
    // Stride 描述 tensor 在内存中的布局，让 kernel 知道如何从一个元素跳到下一个元素
    .stride_q_b = q.stride(0),     // 跳到下一个 batch: 1×128×576 = 73728
    .stride_q_s_q = q.stride(1),   // 跳到下一个 query token: 128×576 = 73728
    .stride_q_h_q = q.stride(2),   // 跳到下一个 head: 576
    // ... kv, indices, out, lse 的 stride 类似
};

→ 无硬件特性依赖，纯数据结构准备

第 3 步：Tile Scheduler — 工作分配

文件: csrc/smxx/decode/get_decoding_sched_meta/get_decoding_sched_meta.cu

这一步决定了 哪个 SM 处理哪些 batch 的哪些 KV block。

3.1 计算每个 batch 的处理块数

⚠️ “block” 概念澄清: 这里的 “block” 不是 KV cache page block（64 个连续 token 组成的页面），而是对 top-k indices 列表的分块。

top-k 选择是 token 级别的 — indices 数组中存储的是 2048 个（本例中 128 个）散落在 KV cache 各处的 token 的绝对位置，这些 token 之间通常不连续。

由于 kernel 无法一次性处理所有 top-k token，所以按 TOPK_BLOCK_SIZE=64 将 indices 列表切块处理：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19

Launch: <<<1, 32>>> (单 CUDA block, 32 线程，即 1 个 warp)

indices[batch 0] = [token_47, token_2091, token_583, ..., token_8821] ← 共 128 个 token index (散落各处)

将 128 个 token index 按 TOPK_BLOCK_SIZE=64 分块处理:
  处理块 0: indices[0:63] → 去 KV cache 中 gather 这 64 个 token (位于不同的 page)
  处理块 1: indices[64:127] → gather 另外 64 个 token

对于 batch 0:
  topk = 128, TOPK_BLOCK_SIZE = 64
  num_blocks[0] = ceil(128/64) = 2 ← 需要 2 个处理块来遍历全部 128 个 top-k token

对于 batch 1:
  topk = 128
  num_blocks[1] = ceil(128/64) = 2

fixed_overhead_num_blocks = 5 (每个 batch 的固定 CUDA CTA 开销)

total_num_blocks = (2+5) + (2+5) = 14 ← 需要调度的 CUDA CTA 总数

🔧 硬件特性: 使用 __shfl_xor_sync 做 warp 内规约求和 — 这是 CUDA warp shuffle 指令，所有 NVIDIA GPU 支持。

注意: FP8 量化将每个 token 从 1152 字节 (576×BF16) 压缩到 656 字节，减少 ~43% 的内存带宽消耗。Decode 阶段是 memory-bound（瓶颈在显存带宽而非计算），因此减少每 token 的字节数可以直接提升吞吐量。

3.2 背景：SM (Streaming Multiprocessor) 是什么

SM 是 NVIDIA GPU 的基本计算单元。可以把 GPU 理解为一个拥有很多独立计算核心的处理器，每个 SM 就是其中一个核心。

• 每个 SM 有自己的寄存器文件（65536 个 32-bit 寄存器）、共享内存（228KB，SM90）、Tensor Core（矩阵乘法加速器）、SFU（特殊函数单元，做 exp2f 等）
• H800 有 132 个 SM，它们并行工作
• 一个 CUDA kernel 启动时，会生成很多 **CTA (Cooperative Thread Array)**，也叫 thread block。GPU 的硬件调度器把这些 CTA 分配到各个 SM 上执行
• 一个 SM 可以同时执行多个 CTA（如果资源允许），但 FlashMLA 的 kernel 每个 CTA 用了接近满载的寄存器和共享内存，所以通常一个 SM 只跑 1 个 CTA

Cluster 是 Hopper 新增的概念：多个 SM 组成一个 cluster，cluster 内的 SM 可以直接读写彼此的共享内存。FlashMLA 中 h_q=128 时 CLUSTER_SIZE=2，即 2 个 SM 组成一个 cluster，协作处理 128 个 head。

3.3 计算 SM Partition 数

1
2
3
4
5
6
7
8
9
10
11
12

h_q = 128 个 head
BLOCK_M = 64 个 head / CTA
→ 处理全部 head 需要 128/64 = 2 个 CTA (组成 1 个 cluster)

s_q = 1 (只有 1 个 query token)

num_sm_parts = max(132 / s_q / (h_q/64), 1)
             = max(132 / 1 / 2, 1)
             = 66

含义：132 个 SM，每 2 个 SM 组成一个 cluster 处理一个"任务单元"，
     所以有 66 个可用的"工位" (SM partition)

3.4 贪心调度算法

调度器的目标是：把所有 batch 的工作均匀分配到 66 个 partition 上。

1
2
3
4
5
6
7
8
9

第一步：计算每个 partition 的负载上限 (payload)

payload = ceil(total_num_blocks / num_sm_parts) + fixed_overhead_num_blocks
        = ceil(14 / 66) + 5
        = 1 + 5
        = 6

每个 partition 最多处理 6 个"单位"的工作。
(+5 是因为每接手一个新 batch 有固定成本：加载 Q、初始化 softmax 状态、写回结果等)

然后，单线程从左到右逐个 partition 分配工作 (代码第 66-91 行):

状态变量:

1
2
3

now_req_idx = 0    // 当前正在分配第几个 batch
now_block = 0      // 当前 batch 已经分配了多少 block
remain_payload = 6 // 当前 partition 剩余容量

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

═══════════ Partition 0, remain_payload=6 ═══════════

看 batch 0: 共 2 个 block
能否把整个 batch 0 放进来？需要 2(block) + 5(开销) = 7 ... 7 > 6, 放不下!

那放部分：remain_payload - fixed_overhead = 6 - 5 = 1 个 block
→ 只分配 batch 0 的 block 0
→ now_block = 1 (batch 0 还剩 1 个 block)
→ 容量用完，跳出

结果:
  begin_req=0, end_req=0, begin_block=0, end_block=1
  is_first_req_splitted=false ← batch 0 从这里开始
  is_last_req_splitted=true   ← batch 0 没处理完 (被 split 了!)

═══════════ Partition 1, remain_payload=6 ═══════════

继续 batch 0: 还剩 1 个 block
能否整个放进来？需要 1 + 5 = 6 ... 6 <= 6, 可以!
→ remain_payload = 6 - 6 = 0
→ batch 0 全部完成，now_req_idx=1, now_block=0
→ 容量用完，跳出

结果:
  begin_req=0, end_req=0, begin_block=1, end_block=2
  is_first_req_splitted=true  ← batch 0 从中间开始 (上个 partition 没处理完)
  is_last_req_splitted=false  ← batch 0 在这里结束

═══════════ Partition 2, remain_payload=6 ═══════════

看 batch 1: 共 2 个 block, 同理 7 > 6 放不下整个
→ 只放 1 个 block

结果:
  begin_req=1, end_req=1, begin_block=0, end_block=1
  is_first_req_splitted=false, is_last_req_splitted=true

═══════════ Partition 3, remain_payload=6 ═══════════

继续 batch 1: 剩 1 个 block, 1+5=6 <= 6, 放进来
→ batch 1 完成

结果:
  begin_req=1, end_req=1, begin_block=1, end_block=2
  is_first_req_splitted=true, is_last_req_splitted=false

═══════════ Partition 4~65 ═══════════
所有 batch 已分配完，begin_req >= batch_size → kernel 直接 return

3.5 Split 与 Combine

当一个 batch 被分到多个 partition 时，就产生了 split — 每个 partition 只计算部分 KV block 的注意力，得到局部结果 (partial O, partial LSE)。最后需要 combine kernel 把各个 split 的结果合并。

1
2
3
4
5
6
7
8
9
10

batch 0 被 split 成 2 份:
  Partition 0 → block 0 的注意力结果 (partial_O_0, partial_LSE_0)
  Partition 1 → block 1 的注意力结果 (partial_O_1, partial_LSE_1)
  → combine kernel 根据 LSE 加权合并：O_final = w0 × partial_O_0 + w1 × partial_O_1

batch 1 同理。

num_splits 前缀和：[0, 2, 4]
  batch 0 的 splits 在 index [0, 2) → 2 个 splits
  batch 1 的 splits 在 index [2, 4) → 2 个 splits

如果一个 batch 没有被 split (is_no_split=true)，则 decode kernel 直接输出最终结果，不需要 combine kernel，省去一次 kernel launch。在实际部署中 (topk=2048, b=128)，大部分 batch 都会被 split 到多个 partition 以充分利用 132 个 SM 的并行度。

🔧 硬件特性: 调度器本身是架构无关的 (smxx/ 目录)，但 num_sm_parts 根据 GPU 的 SM 数量和 kernel 的 cluster size 计算，与具体硬件相关。

第 4 步：FP8 量化回顾 — KV Cache 是如何存储的

在 decode 之前，预填充阶段已经将 KV cache 量化为 FP8。我们用具体数字展示：

4.1 原始 BF16 KV 向量 (一个 token)

1
2
3

原始 K 向量 (d=576):
  K_nope[0:511] = [0.25, -0.5, 0.125, ..., 0.75]  (512 个 BF16 值)
  K_rope[512:575] = [0.1, -0.2, 0.3, ..., 0.05]   (64 个 BF16 值)

4.2 分组量化

NoPE 部分按 128 个一组量化（V3.2 有 4 组）：

1
2
3
4
5
6
7
8
9
10
11
12

Group 0: K_nope[0:127]
  max_abs = max(|K_nope[0]|, ..., |K_nope[127]|) = 0.75
  scale_inv = 0.75 / 448.0 = 0.001674
  scale_inv_ue8m0 = 2^(ceil(log2(0.001674))) = 2^(-9) = 0.001953
  scale = 1 / 0.001953 = 512.0
  
  量化：K_nope_fp8[i] = round_to_fp8(K_nope[i] / scale_inv_ue8m0)
  例：K_nope_fp8[0] = round_to_fp8(0.25 / 0.001953) = round_to_fp8(128.0) = 128.0 (FP8)

Group 1: K_nope[128:255] → scale_1
Group 2: K_nope[256:383] → scale_2
Group 3: K_nope[384:511] → scale_3

4.3 存储布局 (656 字节)

🔧 硬件特性:

• FP8 (E4M3): Hopper/Blackwell 原生数据类型，1 符号位 + 4 指数位 + 3 尾数位，范围 ±448，精度约 3 位有效数字
• UE8M0 缩放因子: 纯 2 的幂，无尾数位，确保反量化时乘法精确无舍入误差
• Blackwell (SM100) 额外支持: FP8_E8M0FNU 格式的缩放因子，MODEL1 使用此格式

第 5 步：CUDA Kernel 启动

5.1 SM90 (Hopper) 启动配置

文件: csrc/sm90/decode/sparse_fp8/splitkv_mla.cuh 底部

1
2
3
4
5
6
7
8
9

// SM90 启动配置
cutlass::ClusterLaunchParams launch_params = {
    dim3(2, 1, 66),              // grid: (NUM_M_BLOCKS=2, s_q=1, num_sm_parts=66)
    dim3(384, 1, 1),             // block: 384 线程 = 3 warpgroups
    dim3(2, 1, 1),               // cluster size = 2 (h_q=128 → CLUSTER_SIZE=NUM_M_BLOCKS=2)
    sizeof(SharedMemoryPlan),    // 动态共享内存
    params.stream
};
cutlass::launch_kernel_on_cluster(launch_params, kernel, params, tma_params);

🔧 Hopper 硬件特性:

• Cluster Launch: launch_kernel_on_cluster 使用 Hopper 的 Thread Block Cluster 功能。h_q=128 时 CLUSTER_SIZE=2，两个 CTA 组成一个 cluster，可以直接访问对方的共享内存 (通过 XOR 寻址)
• 384 线程: 3 个 warpgroup（每个 128 线程 = 4 个 warp），这是 Hopper WGMMA 指令的基本执行单位
• NUM_M_BLOCKS=2: h_q=128 / BLOCK_M=64 = 2，每个 CTA 处理 64 个 head，2 个 CTA 覆盖全部 128 个 head

5.2 SM100 (Blackwell / GB200) 启动配置

文件: csrc/sm100/decode/head64/kernel.cuh 底部

1
2
3
4
5
6
7
8

// SM100 启动配置 — 注意与 SM90 的关键区别
mla_kernel<<<
    dim3(params.s_q, params.num_sm_parts, 1),  // grid: (s_q=1, num_sm_parts=132, 1)
    dim3(384, 1, 1),                           // block: 384 线程 = 3 warpgroups (相同)
    smem_size,                                 // 动态共享内存
    params.stream
>>>(params, tma_params);
// 注意：没有使用 cluster launch, 也没有使用 PDL

关键区别:

1
2

// SM90 的 KU_ASSERT: h_q == 64 || h_q == 128 等
// SM100 head64 的 KU_ASSERT: h_q == 64 (即 B_H == 64)

SM100 为什么 h_q == 64 而不是 128?
SM100 head64 实现 (csrc/sm100/decode/head64/) 要求 h_q == B_H == 64。当推理框架处理 DeepSeek-V3.2 (h_q=128) 时，需要通过其他方式拆分 head（如 head64x2 或 head128 实现），或者多次调用。这与 SM90 通过 cluster (2 个 CTA 各处理 64 head) 的方案不同。

5.3 SM100 的三大硬件创新及其在 Kernel 中的应用

5.3.1 TMEM (Tensor Memory) — 新增的片上存储层

SM100 在寄存器和共享内存之外，新增了一层 TMEM (Tensor Memory):

1
2
3
4
5

SM100 存储层次:
  寄存器 (Register File)     — 每线程私有，最快
  TMEM (Tensor Memory)       — 512KB, 每 SM 共享，Tensor Core 可直接读写 ← 新增!
  共享内存 (Shared Memory)   — 228KB, CTA 内共享
  L2 Cache → HBM

在 FlashMLA 中的具体应用 (csrc/sm100/decode/head64/kernel.cuh):

1
2
3
4
5
6
7
8
9
10
11
12

// TMEM 列布局 (config.h 第 72-80 行)
struct tmem_cols {
    static constexpr int O = 0;     // 列 0~255: 输出 O 累加器 (FP32)
    static constexpr int Q = 256;   // 列 256~399: Q 矩阵 (BF16)
    static constexpr int P = 400;   // 列 400~463: P = Q·K^T 的结果 (FP32)
};

// kernel 初始化时分配 512 列 TMEM (kernel.cuh 第 63 行)
cute::TMEM::Allocator1Sm().allocate(512, plan.tmem_start_addr.data());

// kernel 结束时释放 (kernel.cuh 第 425 行)
cute::TMEM::Allocator1Sm().free(0, 512);

SM90 的对比: SM90 没有 TMEM，Q 和 O 都在共享内存或寄存器中:

• Q 通过 TMA → 共享内存 → WGMMA 直接从共享内存读取
• O 累加器在寄存器中
• P 在寄存器中

SM100 的优势:

• TMEM 带宽远高于共享内存，Q 常驻 TMEM 后每个 block 的 QK^T 都不需要重新加载 Q
• O 累加器放在 TMEM 中，减轻寄存器压力 (Warpgroup 0 从 SM90 的 192 regs 变为 224 regs，但 regs 不用存 O 了)
• P 的 softmax 结果直接在 TMEM 中读取，无需通过共享内存传递

5.3.2 UTCMMA (Unified Tensor Core MMA) — 替代 WGMMA

SM100 使用 UTCMMA (tcgen05.mma) 指令替代 SM90 的 WGMMA (wgmma.mma_async):

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16

// QK^T 矩阵乘法:
// SM90: WGMMA SS 模式 — Q 和 K 都在共享内存
gemm<true, -1>(tiled_mma_QK, sQ, sK, rP);
// PTX: wgmma.mma_async.sync.aligned.m64n64k16.f32.bf16.bf16 ...

// SM100: UTCMMA TS 模式 — Q 在 TMEM, K 在共享内存
ku::utcmma_ts(tiled_mma_P, tQ_in_tmem, sK_in_smem, tP_in_tmem, true);
// PTX: tcgen05.mma.ws.cta_group::1.kind::f16 [tmem_c], [tmem_a], smem_desc_b, ...

// SV 矩阵乘法:
// SM90: WGMMA RS 模式 — S 在寄存器，V 在共享内存
gemm<false, -1>(tiled_mma_PV, rS, sV, rO);

// SM100: UTCMMA SS 模式 — S 和 V 都在共享内存，O 在 TMEM
ku::utcmma_ss(tiled_mma_O, sS, sV, tO_in_tmem, false);
// PTX: tcgen05.mma.ws.cta_group::1.kind::f16 [tmem_c], smem_desc_a, smem_desc_b, ...

区别总结:

5.3.3 TMA Gather4 — 硬件稀疏 Gather

这是对 sparse decode 最重要的 SM100 创新。SM90 通过线程协作做 __ldg 逐 token gather，SM100 直接用硬件 TMA 做 2D 稀疏 gather:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18

// SM90: 每个线程用 __ldg 加载 1 个 token 的一部分 (splitkv_mla.cuh 第 509-530 行)
int token_index = __ldg(gIndices + ...);  // 读 index
int block_id = token_index / page_block_size;
int offset = token_index % page_block_size;
fp8x16 data = load_128b_from_gmem(k_ptr + block_id*stride + offset*row_stride + dim_offset);
// → 128 个线程各自发起 1 个 128-bit 的全局内存读取

// SM100: 1 个线程用 TMA Gather4 一次加载 4 个 token (kernel.cuh 第 596-611 行)
ku::tma_gather4(
    &tma_params.tensor_map_kv_nope,    // TMA 描述符 (2D tensor layout)
    plan.bar_raw_ready[buf_idx],       // 完成后通知这个 barrier
    plan.u.kv.raw_nope[buf_idx].data(), // 目标：shared memory
    0,                                 // column index (NoPE 起始列)
    cur_indices,                       // int4: 4 个 row index (token 位置)
    TMA::CacheHintSm90::EVICT_LAST
);
// PTX: cp.async.bulk.tensor.2d.shared::cta.global.tile::gather4.mbarrier::complete_tx::bytes
//      .cta_group::1.L2::cache_hint [smem], [desc, {col, row0, row1, row2, row3}], [mbar], hint;

SM100 的 gather4 工作流程:

1
2
3

64 个 top-k token 需要 gather → 每次 gather4 加载 4 个 token → 需要 16 次 TMA gather4
每次 gather4: 从 2D tensor [num_tokens × D_NOPE] 中按 row index 取 4 行
             硬件自动处理地址计算、跨 page block 的非连续访问

对比:
| 架构 | Gather 方式 |
|—|—|
| SM90 | 128 线程 × 多轮 __ldg → ~128 个散射的全局内存请求 |
| SM100 | 1 线程 × 16 次 TMA gather4 → 16 个硬件管理的 DMA 请求 |

优势:

• 线程无需参与地址计算和数据搬运，可以做其他工作
• TMA 引擎内部有合并和调度优化，比线程发起的 __ldg 更高效
• 数据直接落入 shared memory，无需经过寄存器

5.3.4 SM100 UTCCP — Q 从共享内存搬到 TMEM

SM100 加载 Q 的过程分两步:

1
2
3
4
5
6
7
8
9

// Step 1: TMA 把 Q 从 global memory → shared memory (和 SM90 相同)
ku::launch_tma_copy(tma_params.tma_Q_SW128, gQ, sQ, plan.bar_q_tma, ...);

// Step 2: UTCCP 把 Q 从 shared memory → TMEM (SM100 独有)
SM100_UTCCP_128dp256bit_1cta::copy(sQ_desc, tmem_cols::Q + offset);
// PTX: tcgen05.cp ... (shared → TMEM 的异步拷贝)

// 之后 Q 常驻 TMEM，每个 KV block 的 QK^T 直接从 TMEM 读取 Q
// 不像 SM90 每次都从共享内存读

5.4 SM100 的三 Warpgroup 分工

5.4.1 什么是 Warpgroup？

Warpgroup 是 Hopper (SM90) 引入的线程组织层次，介于 warp 和 CTA (thread block) 之间：

1
2
3
4
5
6

CTA (Thread Block) = 384 线程 (FlashMLA 的配置)
  ├── Warpgroup 0 = 128 线程 = 4 个 warp
  ├── Warpgroup 1 = 128 线程 = 4 个 warp
  └── Warpgroup 2 = 128 线程 = 4 个 warp

其中每个 warp = 32 线程 (NVIDIA GPU 的最小调度单位)

为什么需要 Warpgroup？因为 WGMMA 指令要求它。

在 Hopper 之前，Tensor Core 的 MMA 指令由单个 warp（32 线程）发起。Hopper 引入了 **WGMMA (Warpgroup MMA)**，需要 128 线程（4 个 warp）协作发起一条矩阵乘法指令，矩阵规模更大（如 64×64×16），吞吐更高。所以 warpgroup = 4 个 warp = 128 线程，是 WGMMA 的基本执行单位。

5.4.2 SM90 vs SM100 的 Warpgroup 分工对比

SM100 的三个 Warpgroup 分工与 SM90 完全不同:

1
2
3
4
5
6

SM90 (Hopper):
┌──────────────────────────────────────────────────────────────────┐
│ WG0 (128 线程，192 regs) — Consumer A: QK^T + softmax + S·V_left │
│ WG1 (128 线程，160 regs) — Consumer B: S·V_right                 │
│ WG2 (128 线程，152 regs) — Producer: FP8 gather + 反量化          │
└──────────────────────────────────────────────────────────────────┘

在 SM90 中，warpgroup 的概念很”实”——WGMMA 确实需要 128 线程一起发起。WG0 的 128 线程共同发起 WGMMA (QK^T)，然后做 softmax，再共同发起 WGMMA (SV)。WG2 的 128 线程各自 __ldg 加载 FP8 数据并反量化。每个线程都在忙。

1
2
3
4
5
6
7
8
9
10
11
12
13
14

SM100 (Blackwell):
┌──────────────────────────────────────────────────────────────────┐
│ WG0 (128 线程，224 regs) — Softmax: 从 TMEM 读 P → exp2 → 写 S   │
│                            (纯标量计算，不参与任何 MMA)           │
│                                                                  │
│ WG1 (128 线程，72 regs) — MMA + Produce:                         │
│   Warp 4 (1 线程): 发起所有 UTCMMA (QK^T, SV) + TMA 加载 Q        │
│   Warp 5 (1 线程): TMA Gather4 加载 FP8 NoPE 到 shared memory    │
│   Warp 6 (1 线程): TMA Gather4 加载 BF16 RoPE 到 shared memory   │
│   Warp 7 (32 线程): 读 indices → 计算 TMA 坐标 + 加载 scale      │
│                                                                  │
│ WG2 (128 线程，208 regs) — Dequant: FP8→BF16 反量化              │
│                            从 raw_nope (FP8) → dequant.nope (BF16)│
└──────────────────────────────────────────────────────────────────┘

在 SM100 中，warpgroup 的概念变”虚”了——UTCMMA 只需 1 个线程发起，所以同一个 warpgroup 内的线程可以干完全不同的事。SM100 的 WG1 中，大量线程其实是空闲的（Warp 4/5/6 各只需 1 个线程），但仍然被组织为一个 warpgroup，主要是为了共享寄存器配额和 barrier 同步的方便。

这体现了 Blackwell 的设计哲学：硬件加速器（TMA、UTCMMA）接管数据搬运和矩阵计算，线程只处理无法硬件加速的”缝隙”工作。

5.4.3 关键设计差异

1. SM90 的 MMA 由 warpgroup (128 线程) 发起; SM100 的 UTCMMA 由 1 个线程发起

   • UTCMMA 是 “widthless” 指令 (tcgen05.mma.ws)，只需 1 个线程 issue，Tensor Core 自动完成矩阵计算并写入 TMEM
   • 这解放了大量线程去做其他工作

2. Softmax 独立成 WG0 (128 线程)

   • SM90: softmax 由 Consumer A (WG0) 在 WGMMA 之间插入
   • SM100: softmax 是 WG0 的唯一任务 — 从 TMEM 读 P、做 exp2、写 S 到 shared memory
   • 好处：softmax 是标量密集计算，给它 224 个寄存器足够存所有中间状态

3. TMA Gather4 替代线程协作 __ldg

   • SM90: WG2 的 128 个线程各自 __ldg 加载 FP8 数据
   • SM100: WG1 中仅 Warp 5 的 1 个线程发起所有 TMA Gather4，硬件 DMA 完成实际搬运

4. 四重缓冲索引 (NUM_INDEX_BUFS=4)

   • SM90 的 indices 无显式缓冲
   • SM100 用 4 个 buffer 轮转：indices → TMA 坐标 → scale → valid mask
   • 因为 SM100 的 TMA Gather4 需要预计算 TMA 坐标 (tma_coord)，这个过程是异步的

🔧 SM100 的 NUM_INDEX_BUFS=4 vs SM90 的 NUM_K_BUFS=2:
SM100 需要更多 buffer 因为 TMA Gather4 的 pipeline 更深:

• Warp 7 计算 TMA 坐标 (写 buf N)
• Warp 5/6 用 TMA Gather4 加载 FP8/RoPE (读 buf N, 写 KV buf M)
• WG2 反量化 (读 KV buf M)
• WG1-Warp4 发起 UTCMMA (读反量化后的 KV)
• WG0 做 softmax (读 TMEM 中的 P)
  每一步都需要独立的 buffer 来实现流水线

第 6 步：TMA 加载 Q 矩阵

Producer (Warpgroup 2 中线程 0) 发起 TMA 拷贝：

1
2
3

// 加载 Q: [64, 576] 的一个 tile 从 global memory → shared memory
launch_tma_copy(tma_Q, gQ, sQ, plan.bar_q, TMA::CacheHintSm90::EVICT_FIRST);
plan.bar_q.arrive_and_expect_tx(64 * 576 * 2); // 期望 64×576×2 = 73728 字节

具体数据:

1
2
3
4

Q tile shape: [64 heads, 576 dims] × BF16 = 73,728 字节
(每个 CTA 加载 64 个 head, 另外 64 个 head 由 cluster 中的 peer CTA 处理)
从 global memory 地址 q_ptr + batch_0 * stride 开始
拷贝到 shared memory sQ (SwizzledLayout_SW128)

🔧 Hopper (SM90) TMA 特性:

• TMA (Tensor Memory Accelerator): 硬件异步拷贝引擎，无需线程参与
• 只需 1 个线程 发起，硬件自动完成整个 73KB 拷贝
• EVICT_FIRST: L2 cache 提示，Q 只用一次所以优先驱逐
• ClusterTransactionBarrier: TMA 完成后自动减少 barrier 计数
• Swizzle 布局 (SW128): 128 字节粒度的地址映射，消除 shared memory bank conflict

🔧 Blackwell (SM100) TMA 改进:

• SM100 新增 TMA Gather4: 可以从 2D tensor 中非连续地 gather 4 行，专为稀疏访问设计
• SM100 用 TMA 加载 Q 到 shared memory，然后用 UTCCP 将 Q 从 shared memory 搬到 TMEM (Tensor Memory)
• TMEM 是 Blackwell 新增的 512KB 片上存储，比 shared memory 带宽更高

第 7 步：Producer Warpgroup — FP8 反量化

这是稀疏 decode 最核心的步骤之一。Warpgroup 2 (线程 256-383) 负责：

7.1 加载 top-k 索引

1
2
3
4
5

int* gIndices = params.indices + batch_0 * stride;
// indices = [47, 2091, 583, 12, ...] 共 128 个 int32

// 线程 256 负责 token index 0:
int token_index = __ldg(gIndices + 0); // = 47

7.2 计算物理地址

1
2
3
4
5
6

int block_index = 47 / 64 = 0;           // 第 0 个 page block
int rel_idx_in_block = 47 % 64 = 47;     // block 内第 47 个 token

fp8* gK_base = kv_ptr                    // KV cache 起始地址
             + block_index * stride_kv_block  // 跳到第 0 个 block
             + rel_idx_in_block * 656;        // 跳到第 47 个 token (每个 656 字节)

7.3 加载缩放因子

1
2
3
4
5
6
7
8
9
10

// 从偏移 512 处加载 4 个 float32 缩放因子 (128 bits)
float scales_float[4];
*(float4*)(scales_float) = load_128b_from_gmem<float4,
    L1CacheHint::EVICT_LAST,    // 缩放因子可能被复用
    L2PrefetchHint::B128        // 预取 128 字节
>(gK_base + 512);

// scales_float = [0.001953, 0.003906, 0.001953, 0.007812]
// 转为 BF16
bf16 scales[4] = {bf16(0.001953), bf16(0.003906), bf16(0.001953), bf16(0.007812)};

🔧 Hopper 硬件特性:

• load_128b_from_gmem: 通过内联 PTX 汇编实现 128-bit 全局内存加载，附带精确的 L1/L2 cache 控制
• PTX: ld.global.nc.L1::evict_last.L2::128B.v4.s32 {%0,%1,%2,%3}, [%4];
• .nc = non-coherent (只读)，.L1::evict_last = 在 L1 中最后被驱逐（优先保留）
• 这些 cache hint 对 Hopper 和 Blackwell 都有效

7.4 FP8 → BF16 反量化（核心步骤）

每个线程处理 16 个 FP8 值：

1
2
3
4
5

// 加载 16 个 FP8 值 (128 bits)
fp8x16 cur_fp8x16 = load_128b_from_gmem<fp8x16,
    L1CacheHint::EVICT_LAST,
    L2PrefetchHint::B256    // 预取 256 字节 (一次性加载整个 NoPE)
>(gK_nope + dim_idx * 64);

反量化公式:

1

BF16_value = float(FP8_value) × scale

具体例子 (dim_idx=0, Group 0 的前 8 个值):

1
2
3
4

FP8 原始值：  [128.0, -256.0, 64.0, -32.0, 192.0, -16.0, 384.0, -448.0]
scale_0   =   0.001953

BF16 结果：   [0.2500, -0.5000, 0.1250, -0.0625, 0.3750, -0.0312, 0.7500, -0.8750]

C++ 实现 (components/dequant.h):

1
2
3
4
5
6
7

bf16x8 cvt_fp8x8_bf16x8(const fp8x8 &inputs, const __nv_bfloat162 &scale_bf162) {
    // 每次处理 4 个 FP8
    float4 fp32x4 = (float4)(inputs.lo);      // FP8 → FP32 (硬件隐式转换)
    bf16x2 out_lo = __float22bfloat162_rn({fp32x4.x, fp32x4.y}) * scale_bf162;  // FP32 → BF16 × scale
    bf16x2 out_hi = __float22bfloat162_rn({fp32x4.z, fp32x4.w}) * scale_bf162;
    // ... 共 8 个值
}

🔧 Hopper 硬件特性:

• (float4)(fp8x4): Hopper Tensor Core 原生支持 FP8→FP32 的类型转换
• __float22bfloat162_rn(): 硬件 FP32→BF16 转换，round-to-nearest
• 每个线程处理 16 个值，128 线程共处理 2048 个值 → 4 组 × 512/4 = 32 个线程覆盖一个 token 的 NoPE

🔧 Blackwell (SM100) 改进:

• SM100 新增 FP8_E8M0FNU 格式缩放因子（MODEL1 使用），用 __nv_cvt_e8m0x2_to_bf162raw 硬件指令转换
• SM100 的反量化在 Warpgroup 2 中完成后，结果直接写入 shared memory，随后 UTCMMA 从 shared memory 读取

7.5 写入共享内存

1
2

// 写入 sK (shared memory), 使用 interleaved 布局
*(__int128_t*)(sK_nope_base + smem_offset) = *(__int128_t*)&cur_bf16x8;

对于 CLUSTER_SIZE=2 的情况（h_q=128），还需写入 peer CTA 的共享内存：

1
2
3

if constexpr (CLUSTER_SIZE == 2) {
    st_async_128b(sK_nope_peer_base + smem_offset, cur_bf16x8, peer_bar_k_remote_ready);
}

🔧 Hopper Cluster 特性:

• get_peer_addr(): 通过 XOR 地址高位（16MB 偏移）直接访问 cluster 中邻居 SM 的 shared memory
• st_async_128b: 异步写入 + mbarrier 完成通知
• PTX: st.async.weak.shared::cluster.mbarrier::complete_tx::bytes.v2.s64 [dst], {data}, [mbar]
• 这允许两个 CTA 共享反量化后的 K 矩阵，无需经过 global memory

7.6 加载 RoPE 部分（不量化）

1
2
3
4
5
6
7
8

// RoPE 部分直接从 global memory 加载 BF16，不需反量化
bf16x8 cur_bf16x8 = load_128b_from_gmem<bf16x8,
    L1CacheHint::EVICT_LAST,
    L2PrefetchHint::B128
>(gK_rope + dim_idx * 32);

// 直接写入 shared memory
*(__int128_t*)(sK_rope_base + smem_offset) = *(__int128_t*)&cur_bf16x8;

7.7 设置有效性标记

1
2
3
4
5
6

// 线程 0-31 检查索引有效性
if (idx_in_warpgroup < 32) {
    int2 indices = __ldg((int2*)(indices_base + lane_idx*2));
    plan.is_kv_valid[buf_idx][lane_idx*2]   = (indices.x != -1);
    plan.is_kv_valid[buf_idx][lane_idx*2+1] = (indices.y != -1);
}

7.8 通知 Consumer

1
2
3

fence_view_async_shared();                      // 确保所有共享内存写入对 async proxy 可见
plan.bar_k_local_ready[buf_idx].arrive();       // 通知 Consumer: K 数据已就绪
bar_phase_k ^= 1 << buf_idx;                    // 翻转 phase (双缓冲)

🔧 Hopper 硬件特性:

• fence_view_async_shared(): 确保 shared memory 写入对 TMA/async proxy 可见的 fence 指令
• Transaction Barrier: bar_k_local_ready 是 ClusterTransactionBarrier，128 线程共同 arrive，Consumer 的 wait 才会通过
• 双缓冲 (NUM_K_BUFS=2): Producer 写 buf 0 的同时 Consumer 读 buf 1，实现流水线

第 8 步：QK^T 矩阵乘法

8.1 SM90 (Hopper) — WGMMA

Consumer A (Warpgroup 0) 等待 K 就绪后执行：

1
2
3
4
5
6
7

// 等待 Producer 完成反量化
plan.bar_k_local_ready[buf_idx].wait(bar_phase_k >> buf_idx & 1);

// Cluster 模式下还需等待 peer CTA
if constexpr (CLUSTER_SIZE == 2) {
    plan.bar_k_remote_ready[buf_idx].wait(...);
}

WGMMA 计算 P = Q · K^T：

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18