diff --git a/docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md b/docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md index 0032c0f14..567b6049f 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_phase7_materialize_triton_plan.md @@ -518,6 +518,44 @@ forward_mla.py 5. 该路径不替代后续 attention 前的 shared KV materialize,只减少 persistent write 阶段的 remap/quant/store kernel 碎片。 +### P0 fused NSA index MQA prepare 接入 + +在 shared NSA index K/scale materialize 之后、`deep_gemm.fp8_mqa_logits(...)` +之前,CP in-seq-split pair path 原来还会执行: + +```text +GetK(index_buffer, block_tables[0], kv_len) +GetS(index_buffer, block_tables[0], kv_len) +torch.arange(ke_start, ke_start + valid_q_count) +torch.zeros_like(ke) +``` + +对应的 tai-kernel fast path 已通过实验开关接入: + +```text +SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE=0/1 +``` + +默认关闭。开启后只影响 `_get_topk_ragged_with_cp(...)` 的单序列 +`cp_index is None` 且 `current_index_kv is None` 路径: + +```text +nsa_indexer.py + _get_topk_ragged_with_cp(...) + -> try_tai_prepare_cp_mqa_index(...) + -> tai_kernel.nsa_prefill.prepare_cp_mqa_kv_and_range( + index_buffer, block_tables[0], kv_len, + valid_q_count, ke_start, page_size=64, index_head_dim=128) + -> fallback GetK + GetS + torch arange/zeros +``` + +接入约束: + +1. 仅支持 CUDA、`page_size=64`、`index_head_dim=128`、`index_buffer.dtype=uint8`。 +2. 不覆盖 multi-batch `cp_index is not None` path。 +3. 不覆盖 `current_index_kv` reuse path。 +4. 不融合 `fp8_mqa_logits` 本身,也不改变 distributed top-k 语义。 + ## Benchmark plan 新增 benchmark 只测本地 materialize,不包含 CP all-reduce: diff --git a/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md b/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md index d9c393187..70218741f 100644 --- a/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md +++ b/docs/advanced_features/nsa_prefill_cp_phase8_mla_prefix_prefetch_plan.md @@ -31,6 +31,9 @@ prefetch 时构造 `slot_logical_pages`、`page_inverse`、paged `dense_pages` ```text SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0/1 # 生产开关,默认关闭;同时控制 MLA KV 与 index K/scale prefetch SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH=0/1 # 调试日志,只打印 probe layer,默认关闭 +SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES=-1 + # 小 prefix skip threshold;默认 -1 表示使用 cp_size; + # 0 表示不跳过;正数表示绝对 page 数阈值 ``` 已移除的策略开关: @@ -411,6 +414,7 @@ Phase 8 不能改变 dense page id,否则现有 `page_table_1` / `topk_indices ```text SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0/1 +SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES=-1 ``` 默认关闭。调试时可额外开启: @@ -419,6 +423,8 @@ SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH=0/1 SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH=0/1 ``` +`SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES` 同时作用于 MLA KV prefix prefetch 和 NSA index K/scale prefix prefetch。默认 `-1` 使用 `cp_size` 作为阈值,跳过无法覆盖全部 CP lane 的 tiny prefix,避免为收益很小的前缀付出 dense buffer 分配、zero-fill、异步 all-reduce setup 和 event 管理开销。设置为 `0` 可恢复旧行为;设置为正数时按绝对 prefix page 数判断。 + 不新增/不保留: ```text @@ -817,6 +823,7 @@ Phase 8 的 local copy/remap helper 应尽量复用 Phase 7 的 tai materialize ```text SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False) +SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES = EnvInt(-1) forward_batch.cp_shared_kv_mla_prefetcher forward_batch.cp_shared_kv_index_prefetcher ``` diff --git a/docs/advanced_features/nsa_prefill_cp_phase9_two_request_overlap_plan.md b/docs/advanced_features/nsa_prefill_cp_phase9_two_request_overlap_plan.md new file mode 100644 index 000000000..7e1d681ca --- /dev/null +++ b/docs/advanced_features/nsa_prefill_cp_phase9_two_request_overlap_plan.md @@ -0,0 +1,654 @@ +# NSA Prefill CP Phase9:多请求重叠执行计划 + +## 目标 + +Phase9 为 NSA prefill CP + shared KV 增加一条 CP 专用的多请求 overlap 路径。 +目标是在同一个 CP group 内同时保留多个独立 prefill 请求,用一个请求的本地计算、attention 或 MLP/MoE 工作,隐藏另一个请求在 attention 前发生的 shared-KV 通信与 materialize 等待。 + +Phase9 主要隐藏两类 attention 前窗口: + +1. **NSA index K/scale materialize**:top-k 选择前需要用到的 index cache。 +2. **MLA KV materialize**:最终 MQA / NSA attention 前需要用到的 MLA KV cache。 + +这个阶段不减少总通信字节数,也不改变 shared KV layout。Phase9 是 latency hiding:让一个请求的 CP 通信在另一个请求的有效工作期间发生。 + +初版可以从 2 个 slot 起步,但设计不能写死为 pair-only。`SGLANG_CP_TBO_MAX_SLOTS > 2` 表示多个独立 request slot,而不是一个 true `ForwardBatch(batch_size=N)`。 + +## 非目标 + +- 初版不把多个请求合并成一个 true `ForwardBatch(batch_size=N)`。 + 当前 shared-KV 路径在 `nsa_cp_metadata`、page-aligned split、local out-cache loc ownership、direct-write、materialize remap、top-k 等位置仍有 batch-size-1 假设。 +- 不改变 Phase4/Phase5 已定义的 shared-KV ownership/layout。 +- 不支持 decode CP。decode 当前不开 CP。 +- 初版不支持 `round-robin` CP mode。 +- 不直接复用现有 EP two-batch-overlap 作为 CP 调度器。现有 EP TBO 是拆一个 parent batch;Phase9 调度多个独立 request slot。 +- 初版不支持 CUDA graph capture/replay。 +- 初版不要求 slot 动态 refill。第一版可以 admitted slots 全部跑完后再 admission 下一组;后续再做 finished-slot refill。 +- 初版不重新启用 Phase8 prefetch。Phase9 初版必须先稳定 collective ordering 与 async handle 边界;Phase8 prefetch 的重新接入放到后置子阶段。 + +## 术语:slot 与 batch size + +Phase9 的并发单位是 **slot**。 + +一个 slot 拥有一个独立的: + +- `ScheduleBatch` +- `ModelWorkerBatch` +- `ForwardBatch` +- hidden/residual/positions 状态 +- attention backend metadata +- CP materialize handles + +初版每个 slot 内部仍保持 `batch_size == 1`。 + +区别: + +- `SGLANG_CP_TBO_MAX_SLOTS=4`:同时交错 4 个独立请求流。这是 Phase9 初版支持 `>2` 并发的方式。 +- true `ForwardBatch.batch_size > 1`:一个 slot 内有多个 sequence。这需要 batch-aware NSA CP metadata,是后续阶段。 + +## 当前单请求数据流 + +单个 prefill 请求在每层内大致如下: + +```text +hidden local slice + -> indexer local q/k/proj/rope/rotate + -> CP gather / rerange / materialize index K/scale + -> top-k transform + -> shared index direct write + -> MLA q/k/v prepare + shared MLA direct write + -> CP gather / materialize MLA KV + -> MQA / NSA attention + -> residual / MLP / MoE +``` + +当前 profile 中较明显的 bubble 在: + +1. index K/scale materialize / all-reduce / all-gather 之后、top-k/MQA 之前; +2. MLA KV materialize / all-reduce 之后、attention 之前。 + +Phase8 prefetch 能隐藏单请求 pipeline 内的部分 prefix materialize,但它仍是单请求内部的一层提前预取。Phase9 要利用另一个请求的计算窗口来隐藏这些等待。 + +## Phase9 目标数据流 + +Phase9 维护一个固定大小 slot table。每个 slot 是一个独立请求状态机。不同 slot 可以处在不同 layer,但所有 CP rank 的 collective enqueue 顺序必须一致。 + +Phase9 的调度应当是 **attention-window driven**,不是“slot0 发起 CP comm,然后随便拿 slot1 的 local prepare 来 overlap”。真正需要隐藏的是 attention 前的 CP materialize/gather 等待,所以更合理的窗口是: + +```text +一个 slot 正在跑 MQA / NSA attention +另一个 slot 在 CP comm stream 上发起 index 或 MLA materialize/gather +``` + +也就是说,CP comm 应优先和另一个请求的 attention overlap。local prepare 可以作为铺垫阶段,但不应该作为主要 overlap 目标。原因是 local prepare 通常更碎、更短,并且可能和后续 top-k/vector kernels 抢 compute stream;attention 窗口更稳定,且正好位于下一个请求需要消费 KV 前。 + +每个 slot 的 layer-local stage: + +```text +LOCAL_INDEX_PREP +INDEX_COMM_READY_TO_LAUNCH +INDEX_COMM_LAUNCHED +INDEX_READY +TOPK_DONE +MLA_COMM_READY_TO_LAUNCH +MLA_COMM_LAUNCHED +MLA_READY +ATTENTION_READY +ATTENTION_RUNNING +ATTENTION_DONE +MLP_DONE +NEXT_LAYER 或 DONE +``` + +两个 slot 的目标执行形态应类似下面这样。重点是:slot1 的 CP comm 挂在 slot0 的 attention 窗口下,而不是只和 slot0/slot1 的 local prepare overlap。 + +```text +# warmup:先让 slot0 准备到可以跑 attention +slot0 layer L local index prepare -> index comm -> top-k -> MLA comm -> MLA ready + +# 为 slot1 做不含重 CP wait 的本地铺垫 +slot1 layer M local index prepare +slot1 layer M index/top-k 或 MLA launch 条件准备好 + +# 主 overlap 窗口 +slot1 layer M launch index/MLA CP comm on CP comm stream +slot0 layer L run MQA / NSA attention on compute stream overlaps slot1 CP comm + +slot1 layer M wait CP handle only when data is consumed +slot1 layer M run MQA / NSA attention overlaps another slot's CP comm + +slot0 layer L MLP/MoE -> advance L+1 +slot1 layer M MLP/MoE -> advance M+1 +``` + +如果 slot1 的 MLA materialize 依赖 top-k,那么调度器应该尽量把 slot1 的 index/top-k 提前到 peer attention 前完成,使 peer attention 窗口主要隐藏 slot1 的 MLA materialize。若 index 仍未 ready,则可以在 peer attention 中先 launch index comm,但这只是次优路径,因为 index 完成后还要 top-k 才能 launch MLA comm。 + +因此 Phase9 的优先级应是: + +1. **最优**:peer attention 窗口隐藏另一个 slot 的 MLA materialize/gather。 +2. **次优**:peer attention 窗口隐藏另一个 slot 的 index materialize/gather。 +3. **保底**:CP comm 与另一个 slot 的 MLP/dense compute overlap。 +4. **不作为目标**:CP comm 仅与另一个 slot 的 local prepare overlap。 + +关键不变量: + +```text +所有 CP rank 必须以完全一致的顺序 enqueue collective。 +``` + +例如所有 rank 都必须看到同样的序列: + +```text +slot0.reqA.index_comm(layer=10) +slot1.reqB.index_comm(layer=8) +slot0.reqA.mla_comm(layer=10) +slot2.reqC.index_comm(layer=3) +slot1.reqB.mla_comm(layer=8) +``` + +不能出现 rank0 enqueue `slot0.index, slot1.index`,rank1 enqueue `slot1.index, slot0.index`。调度决策不能依赖 rank-local CUDA event readiness、rank-local fallback、local token ownership 或本地时间。 + +## 为什么不能直接复用现有 TBO + +现有 SGLang TBO 相关文件: + +- `python/sglang/srt/batch_overlap/two_batch_overlap.py` +- `python/sglang/srt/batch_overlap/operations.py` +- `python/sglang/srt/layers/attention/tbo_backend.py` + +现有 TBO 的语义是:把一个 parent `ForwardBatch` 拆成两个 child batch,用 `execute_overlapped_operations(...)` overlap EP/MoE stages。 + +Phase9 需要的是: + +- 多个独立 `ScheduleBatch` / `ForwardBatch`; +- 每个请求独立 radix/cache/PD-transfer 状态; +- 每个请求独立 mutable attention metadata; +- 跨 rank 一致的 CP collective order; +- slot 之间允许不同 layer; +- 不依赖 parent batch token/sequence split。 + +因此现有 TBO 的 child backend 思路可以借鉴,但 batch splitting 机制不能直接复用。 + +## 当前代码地图与涉及模块 + +### 1. Admission 与结果回传 + +- `python/sglang/srt/disaggregation/prefill.py` + - `event_loop_overlap_disagg_prefill()` 是当前 PD prefill overlap loop。 + - 当前流程:pop bootstrapped requests -> 构造一个 `ScheduleBatch` -> `run_batch(batch)` -> `result_queue.append((batch.copy(), result))` -> 处理上一轮结果。 + - Phase9 在这里加 gated slot branch。 + - 不能绕过 PD bootstrap。admitted request 必须已经从 `disagg_prefill_bootstrap_queue` 出来。 + +- `python/sglang/srt/managers/scheduler.py` + - `get_new_batch_prefill()` / `_get_new_batch_prefill_raw()` 管 waiting queue mutation、`PrefillAdder`、chunked-prefill、radix cache、KV allocation。 + - `run_batch()` 管 model-worker conversion、overlap stream、result copy、scheduler-side cache update。 + - Phase9 不应手写 request construction。应该加窄的 slot-admission wrapper,复用当前 allocation 路径,并在初版强制每个 slot `max_prefill_bs=1`。 + +- `python/sglang/srt/managers/schedule_batch.py` + - `ScheduleBatch.copy()` 只复制 `process_batch_result()` 需要的字段。 + - Phase9 每个 slot 完成后都要独立 append `(batch.copy(), result)`。 + - 不要构造 synthetic combined batch,否则 PD transfer、radix cache 与 request bookkeeping 容易错。 + +### 2. Worker 与 ForwardBatch 构建 + +- `python/sglang/srt/managers/tp_worker.py` + - `forward_batch_generation()` 当前只处理一个 `ModelWorkerBatch`。 + - Phase9 需要 sibling entrypoint:`forward_batch_generation_slots(...)`。 + +- `python/sglang/srt/model_executor/forward_batch_info.py` + - `ForwardBatch` 已有 CP shared-KV 字段: + - `nsa_cp_metadata` + - `uses_cp_shared_kv` + - `cp_shared_kv_layout` + - `cp_local_out_cache_loc` + - `cp_local_physical_out_cache_loc` + - `cp_shared_kv_mla_prefetcher` + - `cp_shared_kv_index_prefetcher` + - Phase9 初版不要让一个 `ForwardBatch` 表示多个请求。最多加 slot identity/state 字段。 + +- `python/sglang/srt/model_executor/model_runner.py` + - `init_attention_backend()` 当前在 `--enable-two-batch-overlap` 时选择现有 EP `TboAttnBackend`。 + - `forward_extend()` 当前初始化一次 attention metadata,然后调用 monolithic model forward。 + - Phase9 需要 CP 专用 slot forward path,不要复用 EP TBO flag。 + +### 3. DeepSeek / GLM MLA 层执行 + +- `python/sglang/srt/models/deepseek_v2.py` + - `DeepseekV2ForCausalLM.forward()` 在 `can_cp_split(...)` 成功时创建 `nsa_cp_metadata`。 + - `DeepseekV2Model.forward()` 做 embedding、CP token split、逐层 forward,并在 MoE 部分可进入现有 EP TBO。 + - `DeepseekV2DecoderLayer.forward()` 是 layer 级边界,包含 pre-attn comm、attention、pre-MLP comm、MLP/MoE、postprocess。 + - Phase9 需要 model-specific resumable slot runner,保存:`hidden_states`、`residual`、`positions`、`zero_allocator`、`layer_idx`。 + +- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py` + - `forward_absorb_prepare()` 当前包含 local MLA projection、RoPE、indexer/top-k、shared MLA direct write、current-rank CP KV rebuild。 + - `forward_absorb_core()` 调用 `attn_mqa(...)`,进入 NSA backend,在那里 materialize shared MLA KV。 + - Phase9 需要在这里拆 stage hook: + 1. local index/key prepare; + 2. index materialize launch/wait; + 3. top-k; + 4. local MLA KV write; + 5. MLA materialize launch/wait; + 6. MQA/attention。 + +### 4. NSA / shared KV / CP primitives + +- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` + - `_maybe_materialize_shared_index_buffer()` 同步调用 `materialize_shared_paged_buffer(...)`。 + - `_get_q_k_bf16()` 同步调用 `cp_all_gather_rerange_output(...)`。 + - `_get_topk_in_seq_cp_pair()` 仍有 single-batch 假设,所以 Phase9 初版必须保持每 slot batch-size-1。 + +- `python/sglang/srt/layers/attention/nsa_backend.py` + - `forward_extend()` 将 top-k 转成 `page_table_1`,然后在 shared KV path 调用 `materialize_shared_token_kv_buffer(...)`。 + - 这是 Phase9 第二个主要 overlap 目标。 + +- `python/sglang/srt/layers/attention/nsa/utils.py` + - `_cp_attn_tp_all_gather_padded_tensor()` 分配 padded gather output 并 enqueue `cp_all_gather_into_tensor_async(...)`。 + - `cp_all_gather_rerange_output()` 当前立即 rerange 并返回最终 tensor。 + - Phase9 需要拆成 launch/finish。 + +- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py` + - `materialize_shared_token_kv_buffer()` 构造 dense MLA KV 后 all-reduce。 + - `materialize_shared_paged_buffer()` 构造 dense index K/scale pages 后 all-reduce。 + - Phase9 需要给两者增加 handle API。 + +- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py` + - Phase8 prefetch 已经有 stream/event handle 思路。 + - Phase9 初版不能直接让它自主 enqueue CP comm。后置阶段要把 Phase8 prefetch 接入 Phase9 scheduler,由 Phase9 统一管理 collective order。 + +## 模块边界设计 + +### 新增模块 + +#### `python/sglang/srt/batch_overlap/cp_multi_request_overlap.py` + +职责: + +- CP-TBO env/config snapshot; +- slot eligibility 与 fallback reason; +- `CpTboStage`; +- `CpTboStageClass`; +- `CpTboRequestState`; +- `CpTboSlotResult`; +- rank-consistent deterministic scheduler; +- EP overlap policy。 + +建议核心数据结构: + +```python +@dataclass +class CpTboRequestState: + slot_id: int + forward_batch: ForwardBatch + hidden_states: torch.Tensor + residual: torch.Tensor | None + positions: torch.Tensor + layer_idx: int + layer_end: int + stage: CpTboStage + pending_index_handle: object | None = None + pending_mla_handle: object | None = None + done: bool = False +``` + +#### `python/sglang/srt/layers/attention/cp_tbo_backend.py` + +职责: + +- `CpTboAttnBackend`; +- 每个 slot 一个 child backend; +- slot stage 执行前绑定 `forward_batch.attn_backend`; +- 隔离 `NsaBackend.forward_metadata`。 + +### 修改模块 + +- `python/sglang/srt/disaggregation/prefill.py` + - 加 gated slot admission。 +- `python/sglang/srt/managers/scheduler.py` + - 加 slot-aware admission 与 `run_batch_slots(...)`。 +- `python/sglang/srt/managers/tp_worker.py` + - 加 `forward_batch_generation_slots(...)`。 +- `python/sglang/srt/model_executor/model_runner.py` + - 加 CP slot forward entrypoint。 +- `python/sglang/srt/models/deepseek_v2.py` + - 加 DeepSeek/GLM MLA slot runner。 +- `python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mla.py` + - 拆 MLA prepare/core 内部阶段。 +- `python/sglang/srt/layers/attention/nsa/utils.py` + - 加 CP all-gather rerange launch/finish。 +- `python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py` + - 加 shared index / shared MLA materialize launch/finish。 +- `python/sglang/srt/layers/attention/nsa/nsa_indexer.py` + - 支持外部传入 materialized index 或 handle result。 +- `python/sglang/srt/layers/attention/nsa_backend.py` + - 支持外部传入 materialized MLA KV 或 handle result。 +- `python/sglang/srt/environ.py` + - 加私有 env flag,初期不加 CLI args。 + +## 启用条件与环境变量 + +初版只在全部条件满足时启用: + +- `SGLANG_CP_TWO_BATCH_OVERLAP=1`; +- disaggregation mode 是 prefill; +- `--enable-nsa-prefill-context-parallel`; +- `--nsa-prefill-cp-mode in-seq-split`; +- shared KV active; +- 每个 slot `batch_size == 1`; +- extend/prefill forward mode,不是 decode/idle/spec; +- 不走 return-logprob; +- 不走 hidden-state capture; +- 不走 CUDA graph; +- admitted slots 在所有 rank 上 eligibility 一致; +- token length ratio 不超过阈值; +- workspace 足够容纳多个 live hidden/residual/materialize buffer。 + +初始 env: + +```text +SGLANG_CP_TWO_BATCH_OVERLAP=1 +SGLANG_CP_TBO_MAX_SLOTS=2 +SGLANG_CP_TBO_MAX_INFLIGHT_COMM=1 +SGLANG_CP_TBO_DEBUG=1 +SGLANG_CP_TBO_TOKEN_RATIO=2.0 +SGLANG_CP_TBO_EP_OVERLAP_POLICY=avoid_ep_comm +``` + +后续可以增加 alias: + +```text +SGLANG_CP_MULTI_REQUEST_OVERLAP=1 +``` + +## EP overlap 策略 + +Phase9 可以和 EP work overlap,但不能初版就盲目叠加 CP all-gather 与 DeepEP all-to-all。 + +分三层: + +1. **自然 slot overlap**:slot0 等 CP,slot1 跑 attention/MLP。 +2. **带宽感知 CP/EP scheduling**:stage 分类为 `CP_COMM`、`ATTN_COMPUTE`、`MLP_GEMM`、`EP_COMM`。默认 `avoid_ep_comm`。 +3. **与现有 EP TBO 深度融合**:等 CP-TBO 稳定后再做。 + +推荐 Phase9 默认实现第 1、2 层,不做第 3 层。 + +## 子阶段计划 + +### P9.1:slot framework,顺序执行 + +目的:验证调度、admission、result plumbing,不改变计算顺序。 + +任务: + +1. 增加 CP-TBO env gates,包括 `SGLANG_CP_TBO_MAX_SLOTS`。 +2. 增加 slot eligibility helper。 +3. 增加 `get_next_disagg_prefill_slots_to_run(max_slots)`。 +4. 增加 `run_batch_slots()`,初版内部顺序调用 `run_batch(slot_i)`。 +5. 每个 slot 独立 `batch.copy()` 并进入 result queue。 +6. 增加 slot admission/fallback debug log。 +7. 增加 eligibility、FIFO/fairness、`max_slots > 2` admission 单测。 + +预期:无加速;功能上与单请求路径一致。 + +### P9.2:slot state machine 与 metadata 隔离 + +目的:让多个 request state 可以安全交错,包括不同 layer 的 slot。 + +任务: + +1. 增加 `CpTboRequestState`、`CpTboStage`、`CpTboStageClass`。 +2. 增加 deterministic slot scheduler。 +3. 增加 CP slot backend wrapper。 +4. 每个 `ForwardBatch` 使用自己的 NSA metadata/backend。 +5. fake backend 单测验证 slot1 不覆盖 slot0 metadata。 +6. backend 隔离不可用时 fallback 到 P9.1 顺序执行。 + +预期:可以表达 layer-skew,不发生 metadata corruption。 + +### P9.3:index gather 与 MLA materialize 的 launch/wait handle + +目的:把同步通信/materialize 拆成可调度接口。 + +任务: + +1. `cp_all_gather_rerange_output(...)` 拆出 launch/finish。 +2. `materialize_shared_paged_buffer(...)` 拆出 launch/finish。 +3. `materialize_shared_token_kv_buffer(...)` 拆出 launch/finish。 +4. 保留原同步 wrapper:`launch` 后立即 `finish`。 +5. 增加 launch/wait NVTX。 +6. 增加 fake handle 单测验证 collective order。 + +预期:调度器可以控制通信何时发起、何时等待。 + +### P9.4:async multi-slot layer executor + +目的:实现真正的跨请求 overlap。 + +任务: + +1. 增加 `List[CpTboRequestState]` layer loop。 +2. 支持至少两个 slot,且 slot 可处于不同 `layer_idx`。 +3. 用同一 slot table 支持 `SGLANG_CP_TBO_MAX_SLOTS > 2`。 +4. 每 rank 使用一个 ordered CP comm stream。 +5. 只在真实 consumer 前 wait CUDA event。 +6. Phase9 active 时禁用 Phase8 prefetch。 +7. unsupported condition fallback 到 P9.1。 + +预期:Nsight 可看到 slotA CP materialize/gather 主要被 slotB MQA/NSA attention 覆盖;attention 不可用时才退到 dense compute/MLP overlap。 + +### P9.5:EP-aware scheduling + +目的:避免 CP comm 与 DeepEP comm 争抢带宽导致反向变慢。 + +任务: + +1. 标记 stage class:`CP_COMM`、`ATTN_COMPUTE`、`MLP_GEMM`、`EP_COMM`。 +2. 默认策略 `SGLANG_CP_TBO_EP_OVERLAP_POLICY=avoid_ep_comm`。 +3. `avoid_ep_comm` 下,优先在另一个 slot 处于 attention/dense compute 时发 CP comm。 +4. 避免在 DeepEP dispatch/combine 期间启动新的 heavy CP comm。 +5. 增加 debug override policy 供 profile。 +6. 增加 CP-TBO stage decision / DeepEP overlap NVTX。 + +预期:减少 pre-attn bubble,同时不引入更大的 fabric contention bubble。 + +### P9.6:profile、调参、运行稳定化 + +目的:让功能具备实际部署价值。 + +任务: + +1. 增加 `SGLANG_CP_TBO_DEBUG=1` 下简洁日志。 +2. 增加 NVTX: + - `cp_tbo.slotN.index.launch/wait` + - `cp_tbo.slotN.mla.launch/wait` + - `cp_tbo.slotN.attn` + - `cp_tbo.slotN.mlp` + - `cp_tbo.scheduler.decision` +3. 对比 CP-TBO off/on。 +4. 分别测试 2、3、4 个 concurrent long prompts。 +5. 调整 token-ratio、slot count、in-flight comm、EP policy。 + +预期:端到端 concurrent long-prompt throughput 比 Phase8 单请求 pipeline 更好。 + +### P9.7:重新接入 Phase8 prefetch + +目的:在 Phase9 初版稳定后,让单请求 layer-prefetch 和多请求 overlap 可以共存。 + +这个阶段不在初版实现中展开具体代码位置。原因是 Phase9 初版完成前,最终的 stage boundary、handle ownership、CP comm stream、collective ordering API 仍可能调整。现在只定义接入原则和验收条件。 + +接入原则: + +1. Phase8 prefetch 不能自主 enqueue 可能改变 collective order 的 CP 通信。 +2. Phase8 prefetch 必须作为 Phase9 scheduler 管理下的一个 stage/handle。 +3. prefetch 的 launch order 必须进入全 rank 一致的 CP collective schedule。 +4. prefetch 只允许在 Phase9 明确判定安全的 bandwidth window 发起。 +5. prefetch hit/miss 不允许造成不同 rank 走不同 collective 序列。 +6. Phase8 prefetch 与 CP-TBO 共同启用时,fallback 必须清晰打日志。 + +后续设计方向: + +```text +Phase8 current: + layer L attention 后,自主 prefetch layer L+1 prefix + +Phase9-compatible: + CP-TBO scheduler 决定是否为 slotN layer L+1 发起 prefix prefetch + prefetch handle 进入 slot state + slotN 到达 layer L+1 时 consume scheduler-owned handle +``` + +新增 env 建议后续再定名,初步倾向: + +```text +SGLANG_CP_TBO_ENABLE_LAYER_PREFETCH=1 +``` + +验收条件: + +- CP-TBO off + Phase8 prefetch 行为不变; +- CP-TBO on + prefetch off 行为不变; +- CP-TBO on + prefetch on 时 collective order 全 rank 一致; +- Nsight 显示 prefetch 没有和 CP-TBO 主路径重复 materialize 同一窗口; +- 没有 fallback spam、没有 hang、输出不劣化。 + +### P9.8:true `ForwardBatch.batch_size > 1` 后续扩展 + +目的:在 slot scheduler 稳定后,再支持一个 slot 内多个 sequence。 + +前置条件: + +- batch-aware `NSAContextParallelMetadata`; +- batch-aware page-aligned split; +- batch-aware local out-cache loc ownership; +- top-k CP pair 逻辑支持 multi-batch; +- shared KV materialize remap 支持 per-request mapping。 + +这不是 Phase9 初版目标。 + +## 正确性风险 + +### CP collective order mismatch + +如果不同 rank enqueue 顺序不同,可能 hang 或产生错误结果。所有 admission、stage scheduling、fallback 都必须 rank-consistent。 + +### mutable backend metadata + +`NsaBackend.forward_metadata` 是 mutable state。多个 slot 共享一个 backend 会互相覆盖 metadata。必须先做 backend slot isolation。 + +### hidden synchronous fallback + +如果 `cp_all_gather_into_tensor_async` 在某些 backend 下退化为同步路径,CP-TBO 的 overlap 会消失。必须检测并日志提示,必要时 fallback。 + +### 额外显存压力 + +多个 slot 会同时持有 hidden/residual、index dense buffer、MLA dense buffer、pending handle。admission 不只检查 KV pool,也要检查 workspace budget。 + +### CP/EP 带宽争用 + +CP materialize/gather、DeepEP A2A、attention 可能争抢 NVLink/IB。默认策略不能盲目叠加通信。 + +### Phase8 prefetch 重新接入风险 + +Phase8 prefetch 原本是单请求局部优化。重新接入 Phase9 后,如果不纳入全局 CP collective order,可能导致 rank 间 collective 序列不一致。因此 P9.7 必须在 P9.1-P9.6 初版稳定后再做。 + +## 测试计划 + +### 单元测试 + +新增/扩展测试文件: + +- 新增 `test/registered/unit/batch_overlap/test_cp_multi_request_overlap.py` + - env/config parsing; + - slot eligibility; + - deterministic schedule generation; + - EP-overlap policy。 + +- 新增 `test/registered/unit/layers/test_cp_tbo_backend.py` + - child backend creation; + - per-slot metadata isolation; + - unsupported backend fallback。 + +- 扩展 `test/registered/unit/layers/test_nsa_cp_utils.py` + - CP all-gather rerange launch/finish wrapper 与旧同步 wrapper 等价; + - fake handle order。 + +- 扩展 `test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py` + - shared paged materialize launch/finish wrapper 等价; + - shared token MLA materialize launch/finish wrapper 等价; + - async all-reduce unavailable fallback reason。 + +远端容器中运行,不在本地跑: + +```bash +cd /sgl-workspace/sglang-tai +PYTHONPATH=/sgl-workspace/sglang-tai/python:/mnt/beegfs/cjy/tai-kernel/python \ +pytest -q \ + test/registered/unit/batch_overlap/test_cp_multi_request_overlap.py \ + test/registered/unit/layers/test_cp_tbo_backend.py \ + test/registered/unit/layers/test_nsa_cp_utils.py \ + test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +``` + +必须覆盖: + +- 只接受 supported CP shared-KV in-seq-split prefill batch; +- 拒绝 return-logprob、hidden capture、CUDA graph、speculative、rank-local fallback 不一致; +- fewer than 2 eligible requests 时 fallback single-batch; +- `max_slots > 2` admission 不构造 true multi-sequence `ForwardBatch`; +- fake handle 验证不同 layer slot 的 collective order; +- fake backend 验证 slot0/slot1/slot2 metadata 隔离; +- `avoid_ep_comm` 不在 marked DeepEP comm window 启动新 CP comm; +- P9.7 时验证 Phase8 prefetch 的 handle 由 CP-TBO scheduler 管理。 + +### 远端功能测试 + +GLM5 CP prefill 环境: + +1. CP-TBO disabled,发送两个 concurrent long-prompt 请求。 +2. CP-TBO enabled: + +```text +SGLANG_CP_TWO_BATCH_OVERLAP=1 +SGLANG_CP_TBO_MAX_SLOTS=2 +``` + +发送相同请求。 + +3. 测试: + +```text +SGLANG_CP_TBO_MAX_SLOTS=3 +SGLANG_CP_TBO_MAX_SLOTS=4 +``` + +并发送对应数量 concurrent long prompts。 + +4. deterministic generation 设置下对比输出。 +5. 检查: + - 无 CP shared KV fallback spam; + - 无 health-check hang; + - 无 request loss; + - 输出正常。 + +### Profiling 检查 + +Nsight 应看到: + +- MQA 前 large idle bubble 减少; +- CP index/MLA communication ranges 优先被其他 slot 的 MQA/NSA attention 覆盖;attention 窗口不足时才被 dense compute/允许的 MLP work 覆盖; +- slot 可处在不同 layer; +- CP collective order 没有乱序; +- 初版没有 Phase8 prefetch 的额外自主通信; +- P9.7 后 Phase8 prefetch range 纳入 CP-TBO scheduler; +- 默认 EP policy 下没有 CP/DeepEP 带宽 pile-up。 + +## 成功标准 + +Phase9 成功条件: + +1. 两个及以上 concurrent long-prompt prefill 请求能在 CP shared KV 下正确完成。 +2. 至少两个 slot 可处于不同 layer id,且不破坏 metadata 与 collective order。 +3. `SGLANG_CP_TBO_MAX_SLOTS > 2` 表示多个独立 request slot,并可正常运行。 +4. fallback path 安全,保持单请求行为不变。 +5. Nsight 显示两个 attention 前 communication windows 主要被其他 slot 的 MQA/NSA attention 窗口隐藏。 +6. concurrent long-prompt 端到端吞吐高于 Phase8 单请求 pipeline。 +7. P9.7 后 Phase8 prefetch 可在 CP-TBO scheduler 管理下重新启用,不引入 collective order 风险。 diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 0edd9de4b..929ddf3d5 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -208,8 +208,10 @@ class Envs: SGLANG_CP_SHARED_KV_CURRENT_REUSE = EnvBool(False) SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False) SGLANG_CP_SHARED_KV_FUSED_MLA_STORE = EnvBool(False) + SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE = EnvBool(False) SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False) SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH = EnvBool(False) + SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES = EnvInt(-1) SGLANG_TEST_REQUEST_TIME_STATS = EnvBool(False) SGLANG_DISABLE_TP_MEMORY_INBALANCE_CHECK = EnvBool(False) SGLANG_SIMULATE_ACC_LEN = EnvFloat(-1) diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py index 8abd6e190..c945d7f0c 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_prefetch.py @@ -12,6 +12,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( cp_shared_kv_debug_enabled, cp_shared_kv_mla_prefetch_enabled, cp_shared_kv_mla_prefetch_log, + cp_shared_kv_mla_prefetch_min_prefix_pages, cp_shared_kv_mla_prefetch_should_log_layer, filter_locs_mappable_to_physical_pool, filter_pages_mappable_to_physical_pool, @@ -181,6 +182,9 @@ class CpSharedKVMlaPrefetcher: int(real_page_table.numel()), ) return None + min_prefix_pages = cp_shared_kv_mla_prefetch_min_prefix_pages(layout.cp_size) + if prefix_pages < min_prefix_pages: + return None cp_group = get_attention_cp_group() if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1: @@ -595,6 +599,9 @@ class CpSharedKVIndexPrefetcher: int(real_page_table.numel()), ) return None + min_prefix_pages = cp_shared_kv_mla_prefetch_min_prefix_pages(layout.cp_size) + if prefix_pages < min_prefix_pages: + return None cp_group = get_attention_cp_group() if getattr(cp_group, "pynccl_comm", None) is None and layout.cp_size > 1: diff --git a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py index 6bd3644a4..150f4fbb9 100644 --- a/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py +++ b/python/sglang/srt/layers/attention/nsa/cp_shared_kv_runtime.py @@ -41,6 +41,10 @@ def cp_shared_kv_tai_fused_mla_store_enabled() -> bool: return envs.SGLANG_CP_SHARED_KV_FUSED_MLA_STORE.get() +def cp_shared_kv_tai_index_mqa_prepare_enabled() -> bool: + return envs.SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE.get() + + def cp_shared_kv_mla_prefetch_enabled() -> bool: return envs.SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH.get() @@ -49,6 +53,20 @@ def cp_shared_kv_mla_prefetch_log_enabled() -> bool: return envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.get() +def cp_shared_kv_mla_prefetch_min_prefix_pages(cp_size: int) -> int: + """Minimum prefix pages required to enable Phase8 prefetch. + + Negative env values mean "use cp_size" so the default skips tiny prefixes + that cannot cover all CP lanes. Set the env to 0 to disable the gate, or to + a positive absolute page count for workload-specific tuning. + """ + + configured = envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.get() + if configured < 0: + return max(int(cp_size), 0) + return max(int(configured), 0) + + def cp_shared_kv_mla_prefetch_log(message: str, *args) -> None: if cp_shared_kv_mla_prefetch_log_enabled(): logger.info("[CP_SHARED_KV_MLA_PREFETCH] " + message, *args) @@ -219,6 +237,16 @@ def _load_tai_fused_mla_store_kernel(): return None +@lru_cache(maxsize=1) +def _load_tai_index_mqa_prepare_kernel(): + try: + from tai_kernel.nsa_prefill import prepare_cp_mqa_kv_and_range + + return prepare_cp_mqa_kv_and_range + except Exception: + return None + + def _log_tai_materialize_fallback( key: str, message: str, @@ -249,6 +277,51 @@ def _contiguous_for_tai(tensor: torch.Tensor) -> torch.Tensor: return tensor if tensor.is_contiguous() else tensor.contiguous() +def try_tai_prepare_cp_mqa_index( + *, + index_buffer: torch.Tensor, + page_indices: torch.Tensor, + kv_len: int, + valid_q_count: int, + ke_start: int, + page_size: int, + index_head_dim: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | None: + """Try the TAI fused GetK/GetS + MQA range prepare path. + + The fallback path in SGLang launches separate GetK/GetS kernels and then + creates `ks`/`ke` tensors with torch elementwise ops. The TAI kernel fuses + those preparation steps into one launch for the single-sequence CP pair path. + """ + + if not cp_shared_kv_tai_index_mqa_prepare_enabled(): + return None + + kernel = _load_tai_index_mqa_prepare_kernel() + if kernel is None: + return None + + if index_buffer.dtype != torch.uint8: + return None + if not index_buffer.is_contiguous(): + return None + if index_head_dim != 128 or page_size != 64: + return None + + try: + return kernel( + index_buffer, + _contiguous_for_tai(page_indices), + kv_len=int(kv_len), + valid_q_count=int(valid_q_count), + ke_start=int(ke_start), + page_size=int(page_size), + index_head_dim=int(index_head_dim), + ) + except Exception: + return None + + def try_tai_fused_mla_store( *, token_to_kv_pool, diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 18f75eb6e..659e8c02c 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -24,6 +24,7 @@ from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( materialize_shared_paged_buffer, tensor_debug_checksum, tensor_debug_summary, + try_tai_prepare_cp_mqa_index, ) from sglang.srt.layers.dp_attention import attn_tp_all_gather_into_tensor from sglang.srt.layers.layernorm import LayerNorm @@ -1097,39 +1098,59 @@ class Indexer(MultiPlatformOp): ) ke_start = cp_kv_end - actual_seq_q + 1 - ke_offset = torch.arange( - ke_start, - ke_start + valid_q_count, - dtype=torch.int32, - device=q_fp8.device, - ) q_fp8 = q_fp8[:valid_q_count] weights = weights[:valid_q_count] kv_len = min(cp_kv_end, logical_kv_limit) if current_index_kv is None: assert index_buffer is not None - k_fp8 = index_buf_accessor.GetK.execute( - forward_batch.token_to_kv_pool, - index_buffer, - seq_len=kv_len, - page_indices=block_tables[0], - ) - k_scale = index_buf_accessor.GetS.execute( - forward_batch.token_to_kv_pool, - index_buffer, - seq_len=kv_len, + tai_prepared = try_tai_prepare_cp_mqa_index( + index_buffer=index_buffer, page_indices=block_tables[0], + kv_len=kv_len, + valid_q_count=valid_q_count, + ke_start=ke_start, + page_size=page_size, + index_head_dim=forward_batch.token_to_kv_pool.index_head_dim, ) + if tai_prepared is not None: + k_fp8_u8, k_scale, ks, ke_offset = tai_prepared + k_fp8 = k_fp8_u8.view(torch.float8_e4m3fn) + else: + ke_offset = torch.arange( + ke_start, + ke_start + valid_q_count, + dtype=torch.int32, + device=q_fp8.device, + ) + k_fp8 = index_buf_accessor.GetK.execute( + forward_batch.token_to_kv_pool, + index_buffer, + seq_len=kv_len, + page_indices=block_tables[0], + ) + k_scale = index_buf_accessor.GetS.execute( + forward_batch.token_to_kv_pool, + index_buffer, + seq_len=kv_len, + page_indices=block_tables[0], + ) - k_fp8 = k_fp8.view(torch.float8_e4m3fn) - k_scale = k_scale.view(torch.float32).squeeze(-1) + k_fp8 = k_fp8.view(torch.float8_e4m3fn) + k_scale = k_scale.view(torch.float32).squeeze(-1) + ks = torch.zeros_like(ke_offset) else: + ke_offset = torch.arange( + ke_start, + ke_start + valid_q_count, + dtype=torch.int32, + device=q_fp8.device, + ) k_fp8, k_scale = current_index_kv k_fp8 = k_fp8[:kv_len].contiguous() k_scale = k_scale[:kv_len].view(torch.float32).squeeze(-1).contiguous() + ks = torch.zeros_like(ke_offset) kv_fp8 = (k_fp8, k_scale) - ks = torch.zeros_like(ke_offset) ke = ke_offset with self._with_real_sm_count(): diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index 89d3c7588..a6ec657d3 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -458,6 +458,24 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase): with envs.SGLANG_CP_SHARED_KV_LOG_MLA_PREFETCH.override(True): self.assertTrue(cp_shared_kv_mla_prefetch_log_enabled()) + def test_mla_prefetch_min_prefix_pages_defaults_to_cp_size_and_can_override(self): + from sglang.srt.environ import envs + from sglang.srt.layers.attention.nsa.cp_shared_kv_runtime import ( + cp_shared_kv_mla_prefetch_min_prefix_pages, + ) + + envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.clear() + self.assertEqual(cp_shared_kv_mla_prefetch_min_prefix_pages(8), 8) + + with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.override(0): + self.assertEqual(cp_shared_kv_mla_prefetch_min_prefix_pages(8), 0) + + with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.override(16): + self.assertEqual(cp_shared_kv_mla_prefetch_min_prefix_pages(8), 16) + + with envs.SGLANG_CP_SHARED_KV_MLA_PREFETCH_MIN_PREFIX_PAGES.override(-2): + self.assertEqual(cp_shared_kv_mla_prefetch_min_prefix_pages(4), 4) + def test_fused_mla_store_uses_tai_kernel_when_enabled(self): from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout