diff --git a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md index c2037cfc8..236bf1dbb 100644 --- a/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md +++ b/docs/advanced_features/nsa_prefill_cp_shared_kv_bs_gt1_l1_prefetch_zero_sm_plan_zh.md @@ -876,3 +876,122 @@ SGLang 接入位置: 1. partial-current 的 current rows 仍是本 rank 当前 forward 产生的临时 buffer,source layout 不是长期 L1 page buffer;不能直接复用 peer page IPC。 2. 该 current rows all_reduce 需要单独设计 owner-aware current-source IPC/fused compose kernel。 3. 0SM CE path 仍未实现,本阶段只是先消除 bs>1 prefix/suffix 上不必要的 collective。 + +--- + +## 9. 0SM/CE 与 symmetric memory 路线更新(2026-06-12) + +### 9.1 结论 + +可以基于 symmetric memory 做一版更高性能的 0SM transport,但需要把它定义为 **CE/NCCL/NVSHMEM transport primitive**,不是普通 CUDA copy kernel。 + +关键区别: + +1. 现有 TAI IPC materialize 是 GPU kernel peer-read,copy 逻辑运行在 SM 上,会和 attention/GEMM/MoE 抢 SM。 +2. `cudaMemcpyBatchAsync` CE baseline 是 0SM,但每个 descriptor 仍要由 CPU 提交给 driver/CE,短 prefix / bs1-2 时 CPU submit 开销很明显。 +3. symmetric memory 本身只解决跨 GPU 地址/窗口注册与 peer 访问合同;真正的 0SM 需要使用 Copy Engine/NCCL CE collective/one-sided RMA 这类 transport。用 symmetric pointer 写一个 CUDA kernel 仍然会占 SM。 + +### 9.2 当前已验证 baseline + +在 tai-kernel 增加了 CE baseline 方向的单测设计: + +```text +materialize_cuda_ipc_peer_pages_slot_indices_ce( + peer_ptrs: CPU int64 pointer table, + dst: CUDA dense buffer, + owner_ranks/src_page_indices/dst_page_indices: CPU int64 descriptors, + page_nbytes: int, +) +``` + +合同: + +- CPU descriptor 是显式要求,避免隐藏 CUDA descriptor -> CPU 的同步拷贝。 +- 使用 `cudaMemcpyBatchAsync` 提交 D2D/IPC page copies。 +- 不支持 invalid/sentinel zero-fill;prefix spans 进入 CE path 前必须已经 prune 干净。需要 sentinel 语义时使用 SM path 或提前清零。 +- 该 baseline 主要用于验证 0SM 语义与 benchmark,不应直接宣称高性能。 + +远端 quick benchmark: + +```text +# 未合并连续 descriptor 前,小规模 sanity: +cached=4096/req extend=1024/req bs=2 page=64 kv_dim=64 fp16, 2GPU +all_reduce_full: gpu_p50=0.079ms cpu_p50=0.155ms +SM IPC prefix+current: gpu_p50=0.101ms cpu_p50=0.134ms +CE batch prefix + SM current: gpu_p50=0.384ms cpu_p50=0.406ms + +# 合并同 owner/src/dst 连续 descriptor 后: +cached=4096/req extend=1024/req bs=2 page=64 kv_dim=64 fp16, 2GPU +all_reduce_full: gpu_p50=0.076ms cpu_p50=0.094ms +SM IPC prefix+current: gpu_p50=0.094ms cpu_p50=0.113ms +CE batch prefix + SM current: gpu_p50=0.120ms cpu_p50=0.138ms + +cached=102400/req extend=10240/req bs=2 page=64 kv_dim=656 fp8, 8GPU +all_reduce_full: gpu_p50=0.856ms cpu_p50=0.882ms +SM IPC prefix+current: gpu_p50=0.701ms cpu_p50=0.718ms +CE batch prefix + SM current: gpu_p50=0.447ms cpu_p50=0.467ms +``` + +结论:`cudaMemcpyBatchAsync` baseline 语义正确;未合并 descriptor 时短 prefix 明显慢,合并连续 span 后生产尺度 prefix 已经优于 SM IPC/all-reduce。下一步重点不是直接接 SGLang,而是把 span coalescing 与 symmetric-memory/window 管理做成稳定 transport 合同。 + +### 9.3 symmetric memory 高性能方案 + +下一版应优先做 symmetric-memory transport,而不是继续优化普通 CUDA IPC batch submit: + +1. 为每个 CP rank 的 owner-local staging/prefix buffer 建立 symmetric memory/window。 +2. 让所有 rank 用一致的 per-rank window/base offset 访问 owner pages,避免每层反复 open IPC handle 或构造完整 peer pointer table。 +3. 0SM path 只处理 prefix/history pages: + - prefix 已经 backed、只读、ready ordering 简单; + - current suffix 仍先保留现有 ready-counter + SM IPC path,后续再设计 CE signal/ready 合同。 +4. descriptor 粒度从 page list 提升为 coalesced spans: + - bs>1 下先用 request-aware spans; + - 对同 owner、src 连续、dst 连续的条目合并,减少 CE command 数。 +5. benchmark 必须覆盖: + - bs=1/2/5/10; + - cached 100k-300k,extend 10k-65k; + - bf16/fp8; + - descriptor 构造在 timed region 与跨 layer 复用两种模式。 + +### 9.4 风险 + +1. PyTorch symmetric memory 可用性需要按部署版本确认。远端 `cjy-glm5-new` 中 `torch.distributed._symmetric_memory` 可 import,具备 `empty/rendezvous` 接口;`pynvshmem` 不可用。 +2. 如果使用 PyTorch private `_symmetric_memory`,接口稳定性弱,生产风险高;需要封装在 tai-kernel/sglang 边界后面,不要让业务路径直接依赖 private API。 +3. 如果 fallback 到 symmetric pointer + CUDA kernel,则不是 0SM;文档和日志必须明确标记为 SM path。 +4. current suffix 的 ready/wait 如果改到 CE,需要新的信号协议;不能用普通 CE copy 直接替代当前 wait-ready SM kernel。 + +### 9.5 symmetric window prototype 结果(2026-06-12) + +已在 tai-kernel 增加 PyTorch symmetric-memory window helper: + +```text +allocate_cuda_ipc_symmetric_window(num_bytes, device, group) + -> local_buffer: uint8 CUDA symmetric buffer + -> peer_ptrs: CPU int64 remote pointer table + -> peer_buffers/handle: 保持 symmetric rendezvous 生命周期 +``` + +验证: + +```text +torchrun --nproc_per_node=2 -m pytest tests/nsa_prefill/test_cuda_ipc_symmetric_ce.py::test_symmetric_window_ce_materializes_remote_pages_without_sm_ipc_handles +结果:pass +``` + +benchmark 新增 `cache_hit_symm_ce_prefix_current_compose`,其 prefix 使用 symmetric window + CE submit,current 仍沿用 ready-counter + SM IPC。 + +远端 8GPU sanity: + +```text +cached=102400/req extend=10240/req bs=2 page=64 kv_dim=656 fp8, 8GPU +all_reduce_full: gpu_p50=0.859ms cpu_p50=0.880ms +SM IPC prefix+current: gpu_p50=0.703ms cpu_p50=0.722ms +IPC CE prefix + SM current: gpu_p50=0.447ms cpu_p50=0.471ms +SYMM CE prefix + SM current: gpu_p50=0.521ms cpu_p50=0.546ms +``` + +当前结论: + +1. symmetric window 指针合同可用,CE 可以直接从 symmetric peer buffer copy 到本地 dense buffer。 +2. 当前 PyTorch symmetric-memory prototype 没有优于 CUDA IPC CE;推测原因包括 symmetric memory backend/registration overhead、统一 window size 带来的空间浪费、以及仍使用 cudaMemcpyBatchAsync per-span submit。 +3. 短期最有价值路径是继续用 IPC CE 做 prefix prefetch baseline,同时保留 symmetric window helper 作为后续 NCCL CE/NVSHMEM transport 的地址合同实验入口。 +4. 生产接入前必须继续做 span coalescing、跨 layer descriptor 复用,以及更大 bs/不同 cached/extend sweep;不能只凭 symmetric window 概念直接替换现有路径。