diff --git a/deploy/CURRENT.md b/deploy/CURRENT.md index c64e732..9a553ed 100644 --- a/deploy/CURRENT.md +++ b/deploy/CURRENT.md @@ -14,7 +14,7 @@ | 60.5 | `glm53-nvfp4`（Up 29h） | **NVFP4 团队生产**（deploy_glm53_605.sh，md5 fcd9109b）。生产机铁律：不实验、不重启、不覆盖脚本 | `profiles/pro6000/glm53_nvfp4_pro6000_sglang_tp8eagle.env`（方案 A 口径） | | 60.6 | 无容器，但 8 卡被外部裸金属实验占用（`/data/hzy/sparse-opd-*`，09-08 晚实测） | 外部任务，勿动（此前台账漏记） | — | | 60.7 | 基本空（4 卡仍有 `/home/user/dirA_exp` 外部小任务，09-08 晚实测） | 09-08 已拆除清空（方案 F 前身单机实验 + 场景一深优资产留盘），不再恢复 | — | -| 60.8 | `glm53-nvfp4`（Up，:30000，restart=unless-stopped） | **A16 口径在役**（09-08 i8k/o1k/c16 压测收官保留：tp8eagle.env 场景二变体 = A + MRR32 + graph bs 4/8/12/16，池 276,864，质量门 7/7）。同场景实测最优为 B（TP4PP2，out 265 vs A16 219 tok/s），如转纯吞吐用途可切 B。压测前经授权清理了 GPU3/6 的 dirA_ext 外部 eval 进程 | `profiles/pro6000/glm53_nvfp4_pro6000_sglang_tp8eagle.env`（+注释中场景二高并发变体）；压测记录 `experiments/pro6000/glm53_nvfp4_i8k_o1k_c16_bench/` | +| 60.8 | `glm53-nvfp4`（Up，:30000，restart=unless-stopped） | **A16 口径在役**（09-08 i8k/o1k/c16 压测收官保留：tp8eagle.env 场景二变体 = A + MRR32 + graph bs 4/8/12/16，池 276,864，质量门 7/7）。同场景实测最优为 B（TP4PP2，out 265 vs A16 219 tok/s），如转纯吞吐用途可切 B。压测前经授权清理了 GPU3/6 的 dirA_ext 外部 eval 进程。**09-08 晚 PP+MTP r36 实验后已恢复本口径**（实验判决：去 GLOO 修复成立且稳定，但 i8k 场景 PP+MTP 111-116s 差 B' 61.5s 1.81×，图模式 r36g 为正确性灾难——详见 `experiments/pro6000/glm53_ppmtp_r36_degloo/`；语料消耗至 21,235,008/21,296,780 无复测余量） | `profiles/pro6000/glm53_nvfp4_pro6000_sglang_tp8eagle.env`（+注释中场景二高并发变体）；压测记录 `experiments/pro6000/glm53_nvfp4_i8k_o1k_c16_bench/`、`experiments/pro6000/glm53_ppmtp_r36_degloo/` | ## 方案 A-F 一览（GLM-5.3-NVFP4 @ pro6000，2026-09-08 双场景报告口径） diff --git a/experiments/pro6000/glm53_ppmtp_r36_degloo/README.md b/experiments/pro6000/glm53_ppmtp_r36_degloo/README.md new file mode 100644 index 0000000..2b7e425 --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r36_degloo/README.md @@ -0,0 +1,95 @@ +# GLM-5.3-NVFP4 PP+MTP r36（去边界 GLOO）修复与判决 — 2026-09-08 @174.1.60.8 + +接续 `glm53_nvfp4_i8k_o1k_c16_bench`（A16/B 基线）与 r33 PP+MTP 深挖（ppmtp_deepdive）。 +目标：修复自研 TP4PP2+MTP 补丁并让 MTP 与 PP 双优势同时兑现；判负则 60.8 恢复 A16 在役。 + +## 结论速览 + +| 判决项 | 结果 | +|---|---| +| 去 GLOO（r36） | **成立**：稳定（conc_test×3 + cc16 杀手×2 零崩溃、可复现 ±1%），prefill 密集档 8×16384 **-22%**（46.2→36.0s） | +| i8k/o1024/cc16/nreq32 判决（PG19 语料） | **PP+MTP 判负**：e2e P50 **111.2/116.3s** vs B' 同栈对照 **61.5/61.6s**（≡B 基线 61.8s）vs A16 66.2s，**差 1.81×**，双门未过 | +| 缺口归属 | **100% 属 MTP 机制**（B' 在同一 r36 栈上完美复现 B 基线，栈本身零回退）；verify-step ~235ms@cc16 / ~148ms@cc8，break-even 需 <148ms（=52ms×accept2.85） | +| Phase4 图模式（r36g） | **正确性灾难，判死**：`cuda graph: True` 生效但 accept **2.85→1.05**（接受率 0.02）、GSM8K 输出数字噪声（draft+verify 双坏）。**FORCE_EAGER 双开关是承重墙**，非性能旋钮 | +| sanitizer 根治路线 | **不可行**：320 错误全为 NCCL 探测噪声（error 209）；真死因=memcheck 开销 ~10GB/卡 吃掉 0.88 的 KV 余量（最低可行 memfrac 0.959，复现条件被破坏）。沿用 mask=127（4.9ms/轮）+ 实证稳定性，风险显式记录 | +| 在役状态 | **A16 已恢复**（TP8 EAGLE 4/1/5@0.90 + MRR32 + graph bs 4/8/12/16，restart=unless-stopped） | + +## r35→r36：一次死锁与一次修复（本次核心工程产出） + +**r35**（scheduler_pp_mixin_r35.py，r33 fb8c96f7 基础上去 GLOO）： +- dict 通道元数据：`send_object`（2 条 GLOO/字典 × 3 字典/轮）→ **8KB 固定尺寸 NCCL uint8 缓冲**（pickle 零填充，`pickle.loads` 天然忽略 STOP 后字节——已验证）；CPU 张量随车 GPU 化 + `__pp_cpu_keys__` 回填。 +- pyobj 通道（请求中继/控制面）：GLOO 阻塞 `dist.send` → **两段式 NCCL**（int64 size + uint8 payload，空载优化 size=0），走 `world_group.device_group`（与 dict 通道的 pp_group 通信器隔离，FIFO 不串扰）。 +- `_pp_nccl_stage_guard`：fill-event 门 + 双流 record_stream（r33 实证：isend 内核跑在调度器 ambient 流而非 send_stream，guard 必须在 isend 入队前）。 +- 修掉两处自捕 bug：guard 原在 isend 之后（无效门控）→ 重排；`P2PWork` 须钉住实际发送对象（contiguous() 副本）。 + +**r35 首部署死锁**（py-spy 双端定格）： +- PP0_TP0 卡在 `_pp_send_pyobj_nccl` isend 内（懒加载 2-rank NCCL communicator 创建自旋，102% CPU） +- PP1_TP0 停在 `request_receiver.py:142 _pull_raw_reqs` 的 GLOO `point_to_point_pyobj` irecv +- **根因：通道断裂**——请求中继的发送端在 mixin 里（已切 NCCL），接收端却在 `scheduler_components/request_receiver.py`（补丁树之外，仍走 GLOO），两端永不相遇。全库仅 2 个 `point_to_point_pyobj` 使用者，无第三者顺序风险。 + +**r36 修复**（request_receiver_degloo.py，第 8 个挂载文件）： +- `_pull_raw_reqs` p2p 分支（pp_rank≠0 且 attn_tp/cp_rank==0）同步切两段式 NCCL 接收，与发送端同通道同 FIFO 语义（含空载 size=0 锁步）。 +- `SGLANG_PP_DEGLOO` 门控，=0 整体回退 r33 GLOO 行为。 +- 部署即刻通过：健康 380s（与 r33 完全一致）、池 384,960、冒烟 ALL-OK、accept 2.09-3.54。 + +## 判决数据（i8k/o1024/cc16/nreq32，PG19 真实语料，温度0，warm+2） + +| 配置 | run | e2e P50 | out tok/s | TPOT P50 | TTFT P50 | accept | retract | +|---|---|---|---|---|---|---|---| +| PP+MTP r36 | 9402 | 111.2s | 136.0 | 105.2ms | 4.09s | 2.873 | 0 | +| PP+MTP r36 | 9403 | 116.3s | 133.6 | 106.5ms | 4.07s | 2.842 | 0 | +| B'（同栈 nomtp） | 9405 | **61.5s** | **266.7** | 51.4ms | 9.24s | — | 0 | +| B'（同栈 nomtp） | 9406 | **61.6s** | **265.7** | 51.8ms | 9.29s | — | 0 | +| （参照）B 基线 | 09-08 上午 | 61.8s | 264.9 | 50.3ms | — | — | 0 | +| （参照）A16 基线 | 09-08 上午 | 66.2s | 218.9 | 55ms | 4.0s | 2.50 | 7/7 | + +语料窗口：PP+MTP [19,662,144, 20,448,576)、B' [20,448,576, 21,235,008)，各 3×262,144（=warm+2，i8k 全唯一 prompt 窗口=32×8192）。 +消耗至 **21,235,008 / 21,296,780**（余 61,772，无复测余量，复测须换语料）。run-id 9401-9406（9401/9404 为 warm，不入判决）。 +缓存纪律验证：两配置 hit_rate=0.0（窗口互不重叠、radix 干净）。 + +**MTP 增量拆解**：稳态 16 并发 decode 190 tok/s → per-req 11.9 tok/s → step≈239ms；8 并发段 step≈148ms（≈A16 的 137ms）。 +步长随批超线性 + spec 每轮 4 次 forward（3 draft + 1 verify，全 eager）+ 边界中继 → 2.85× 放大率不敌 4.5× 步长。MTP 在本栈 cc16 为**净负收益 1.59×**。 + +## cc16 杀手 profile（16×16384/512 random-ids，历史竞态触发器，语料免费） + +| 配置 | 种子 | wall | out tok/s | E2E 中位 | TPOT 中位 | accept | +|---|---|---|---|---|---|---| +| r36 eager | 6402 | 105.3s | 77.8 | 72.2s | 72.5ms | 3.76 | +| r36 eager | 6403 | 106.0s | 77.3 | 72.6s | 73.1ms | 3.69 | + +零崩溃、±1% 可复现（SYNC_MASK=127 掩码下；历史 SYNC_MASK=0 会触发 async-IMA 的同一 profile）。 +注意：random-ids 的 accept 3.7+ 相对真实语料 2.85 虚高，与既有结论一致。 + +## Phase4 图模式判决（r36g = r36 + FORCE_EAGER_DRAFT/VERIFY=0） + +- `cuda graph: True` 确认生效；conc_test ALL-OK 但 **accept len 1.05-1.11 / 接受率 0.02-0.04**（draft 提议几乎全灭） +- 质量门：GSM8K-1 输出即数字噪声 → **verify 图同样损坏**（非仅 draft） +- r33 注释"draft-decode cuda-graph capture issues on PP stages"实证为 draft+verify 双重捕获缺陷 +- **修复方向**（未来工程）：按 cudagraph×spec 冲突清单（桶形状/地址预分配/host 控制流等 8 类）做预分配补丁 + 图外 canary，非配置级可解 + +## 文件清单（md5） + +| 文件 | md5 | +|---|---| +| patches/scheduler_pp_mixin_r35.py | efce72af34b2e59f20d2eb0874f9f2b2 | +| patches/request_receiver_degloo.py | ad6c978a3e1efea395ec3d537a4fedd8 | +| scripts/deploy_ppmtp_r35.sh | 86fb4b576fb04eef7b71628324796ab7 | +| scripts/deploy_ppmtp_r36.sh | 7b80f4af0debccef8e87df4a2532d8dc | +| scripts/deploy_ppmtp_r36g.sh | dca9dd387a5e980dc81e25590b895646 | +| scripts/deploy_ppmtp_sanitize.sh | 5bb5a7dfaf00569d5408e487267129a6 | +| scripts/killer_cc16.sh | ac3dda7e05b416aecf1da978325e2584 | +| patches/request_receiver_orig_0901.py | vanilla 参考（镜像原版） | + +results/bench_logs/：9401-9406 六轮语料 SUMMARY、killer s6402/6403、质量门 r36（7/7）/r36g（灾难证据）、 +四个部署日志、san_crash_prehealth_s6401.log（290K，sanitizer 不可行的完整证据链：320×error209 噪声 + KV OOM 死因）。 + +60.8 服务器侧资产：/root/scheduler_pp_mixin_r35.py、request_receiver_degloo.py、deploy_ppmtp_r35/36/36g.sh、 +sglang_patch2/（r33 七件套不动）、corpus_ids.json（消耗至 21,235,008）。 + +## 风险与遗留显式记录 + +1. **async-IMA 竞态未根治**：sanitizer 路线不可行（见上），r36 全部验证在 SYNC_MASK=127 掩码下通过。 + 动掩码/改流拓扑/图化任一变动都可能重新暴露，动前必须重跑全套稳定性门。 +2. **PP+MTP 的兑现条件**（后续若再战）：verify-step 需压到 <148ms（现 235@cc16/148@cc8）——图捕获修复（Phase4 深水区） + 是最大单点杠杆；否则 MTP 在 PP 内为净负收益，不如 B（TP4PP2 nomtp）。 +3. A16 三参数口径（MRR32 + cuda-graph-max-bs-decode 16 + bs 4/8/12/16）已从档案找回并用于恢复部署。 diff --git a/experiments/pro6000/glm53_ppmtp_r36_degloo/patches/request_receiver_degloo.py b/experiments/pro6000/glm53_ppmtp_r36_degloo/patches/request_receiver_degloo.py new file mode 100644 index 0000000..03664b3 --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r36_degloo/patches/request_receiver_degloo.py @@ -0,0 +1,336 @@ +from __future__ import annotations + +from dataclasses import dataclass +from http import HTTPStatus +from typing import ( + TYPE_CHECKING, + Any, + Callable, + List, + Optional, + Union, +) + +import os +import pickle + +import torch +import torch.distributed +import zmq +from torch.distributed import barrier + +from sglang.srt.disaggregation.utils import prepare_abort +from sglang.srt.environ import envs +from sglang.srt.managers.io_struct import ( + BatchTokenizedEmbeddingReqInput, + BatchTokenizedGenerateReqInput, + TokenizedEmbeddingReqInput, + TokenizedGenerateReqInput, + sock_recv, +) +from sglang.srt.managers.mm_utils import ( + has_shm_features, + unwrap_shm_features, +) +from sglang.srt.runtime_context import get_disagg, get_parallel, is_ep_scale_joiner +from sglang.srt.utils import ( + broadcast_pyobj, + point_to_point_pyobj, +) +from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method + +if TYPE_CHECKING: + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.distributed.parallel_state_wrapper import ParallelState + from sglang.srt.rust_server.server import RustServer + from sglang.srt.server_args import ServerArgs + from sglang.test.scripted_runtime.scheduler_hook import ScriptedSchedulerHook + from sglang.test.scripted_runtime.tokenizer_recv_proxy import ( + ScriptedTokenizerRecvProxy, + ) + +# [r36 de-GLOO] Request-relay receiver side. The sender moved to a two-phase +# NCCL exchange on world_group.device_group (scheduler_pp_mixin +# _pp_send_pyobj_nccl); this end MUST use the same channel, or the pair never +# meets — a GLOO irecv here while the peer isends on NCCL wedges the first +# request (observed as the 09-08 r35 startup deadlock: PP0 stuck inside +# isend's lazy 2-rank comm init, PP1 parked in the GLOO irecv). +_PP_DEGLOO = os.getenv("SGLANG_PP_DEGLOO", "1") == "1" + + +def _pp_recv_pyobj_nccl(world_group, global_src: int): + """[r36 de-GLOO] Two-phase NCCL object recv: int64 size, then uint8 + payload. Mirrors scheduler_pp_mixin._pp_send_pyobj_nccl, including the + empty optimization (size=0 means no payload follows), so the relay stays + in FIFO lockstep on the same pair communicator the sender uses.""" + device = torch.device(torch.cuda.current_device()) + group = world_group.device_group + size_tensor = torch.empty(1, dtype=torch.int64, device=device) + work = torch.distributed.irecv(size_tensor, global_src, group=group) + work.wait() + size = int(size_tensor.item()) + if size == 0: + return [] + payload_tensor = torch.empty(size, dtype=torch.uint8, device=device) + work = torch.distributed.irecv(payload_tensor, global_src, group=group) + work.wait() + return pickle.loads(bytes(payload_tensor.cpu().numpy().tobytes())) + + +@dataclass(kw_only=True, slots=True, frozen=True) +class SchedulerRequestReceiver: + recv_from_tokenizer: Union[zmq.Socket, ScriptedTokenizerRecvProxy, RustServer] + recv_from_rpc: Optional[zmq.Socket] + recv_skipper: Any + input_blocker: Any + mm_receiver: Any + ps: ParallelState + tp_group: Any + tp_cpu_group: Any + attn_tp_group: Any + attn_tp_cpu_group: Any + attn_cp_group: Any + attn_cp_cpu_group: Any + world_group: Any + server_args: ServerArgs + model_config: ModelConfig + max_recv_per_poll: int + stream_output: Callable[..., None] + get_last_batch: Callable[[], Any] + scripted_scheduler_hook: Optional[ScriptedSchedulerHook] = None + + def recv_limit_reached(self, num_recv_reqs: int) -> bool: + if self.max_recv_per_poll < 0: + return False + return num_recv_reqs >= self.max_recv_per_poll + + @scheduler_nvtx_method("scheduler.recv_requests") + def recv_requests( + self, + ) -> List[Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput, Any]]: + """Receive results at tp_rank = 0 and broadcast it to all other TP ranks.""" + + if self.scripted_scheduler_hook is not None: + self.scripted_scheduler_hook.step() + + if self.recv_skipper is not None: + if not self.recv_skipper.handle(self.get_last_batch()): + return [] + + recv_reqs = self._pull_raw_reqs() + + if self.input_blocker is not None: + recv_reqs = self.input_blocker.handle(recv_reqs) + + recv_reqs = self._broadcast_reqs_across_ranks(recv_reqs) + + if self.ps.pp_rank == 0: + self.unwrap_pickle_wrapper(recv_reqs) + + recv_reqs = self._apply_mm_receiver(recv_reqs) + + self._finalize_shm_features(recv_reqs) + + return recv_reqs + + def _pull_raw_reqs(self) -> Optional[List]: + if self.ps.pp_rank == 0: + if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: + recv_reqs = [] + + # Rust ringbuffer backend: drain the in-process ring fed by the + # embedded Rust TokenizerManager instead of a zmq socket. Same + # non-blocking, msgpack-decoded contract as the zmq path below. + if envs.SGLANG_RUST_SERVER.get(): + recv_reqs.extend( + self.recv_from_tokenizer.drain(self.max_recv_per_poll) + ) + return recv_reqs + + while True: + try: + if self.recv_limit_reached(len(recv_reqs)): + break + recv_req = sock_recv(self.recv_from_tokenizer, zmq.NOBLOCK) + except zmq.ZMQError: + break + recv_reqs.append(recv_req) + + while True: + try: + if self.recv_limit_reached(len(recv_reqs)): + break + recv_rpc = sock_recv(self.recv_from_rpc, zmq.NOBLOCK) + except zmq.ZMQError: + break + recv_reqs.append(recv_rpc) + else: + recv_reqs = None + else: + if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: + dp_offset = ( + self.ps.attn_dp_rank * self.ps.attn_cp_size * self.ps.attn_tp_size + ) + if _PP_DEGLOO: + recv_reqs = _pp_recv_pyobj_nccl( + self.world_group, + (self.ps.pp_rank - 1) * self.ps.tp_size + dp_offset, + ) + else: + recv_reqs = point_to_point_pyobj( + [], + self.ps.pp_rank * self.ps.tp_size + dp_offset, + self.world_group.cpu_group, + (self.ps.pp_rank - 1) * self.ps.tp_size + dp_offset, + self.ps.pp_rank * self.ps.tp_size + dp_offset, + ) + else: + recv_reqs = None + return recv_reqs + + def _broadcast_reqs_across_ranks(self, recv_reqs: Optional[List]) -> List: + if get_parallel().enable_dp_attention: + if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: + work_reqs, control_reqs = self._split_work_and_control_reqs(recv_reqs) + else: + work_reqs = None + control_reqs = None + + if self.ps.attn_tp_size != 1: + work_reqs = broadcast_pyobj( + work_reqs, + self.attn_tp_group.rank, + self.attn_tp_cpu_group, + src=self.attn_tp_group.ranks[0], + ) + + if self.ps.attn_cp_size != 1: + work_reqs = broadcast_pyobj( + work_reqs, + self.attn_cp_group.rank, + self.attn_cp_cpu_group, + src=self.attn_cp_group.ranks[0], + ) + + # When dp_attention_local_control_broadcast is enabled, each DP + # group leader already receives control messages from the DP + # controller, so we broadcast within attn_tp_group + attn_cp_group + # instead of the full tp_group. This avoids an expensive + # all-ranks gloo sync. + _local_ctrl = ( + get_parallel().enable_dp_attention_local_control_broadcast + or is_ep_scale_joiner() + ) + if _local_ctrl: + if self.ps.attn_tp_size != 1: + control_reqs = broadcast_pyobj( + control_reqs, + self.attn_tp_group.rank, + self.attn_tp_cpu_group, + src=self.attn_tp_group.ranks[0], + ) + if self.ps.attn_cp_size != 1: + control_reqs = broadcast_pyobj( + control_reqs, + self.attn_cp_group.rank, + self.attn_cp_cpu_group, + src=self.attn_cp_group.ranks[0], + ) + elif self.ps.tp_size != 1: + control_reqs = broadcast_pyobj( + control_reqs, + self.tp_group.rank, + self.tp_cpu_group, + src=self.tp_group.ranks[0], + ) + recv_reqs = work_reqs + control_reqs + elif self.ps.tp_size != 1: + recv_reqs = broadcast_pyobj( + recv_reqs, + self.tp_group.rank, + self.tp_cpu_group, + src=self.tp_group.ranks[0], + ) + return recv_reqs + + def unwrap_pickle_wrapper(self, recv_reqs: Optional[List]) -> None: + if not recv_reqs: + return + + for req in recv_reqs: + if isinstance(req, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput)): + req.unwrap_pickle_fields() + elif isinstance( + req, (BatchTokenizedGenerateReqInput, BatchTokenizedEmbeddingReqInput) + ): + for sub_req in req: + sub_req.unwrap_pickle_fields() + + def _apply_mm_receiver(self, recv_reqs: List) -> List: + # Process MM requests under EPD-disaggregation mode + if ( + self.ps.pp_rank == 0 + and get_disagg().language_only + and get_disagg().encoder_transfer_backend + in ["zmq_to_scheduler", "mooncake"] + ): + recv_reqs, abort_reqs = self.mm_receiver.process_waiting_requests(recv_reqs) + for req, error_msg, error_code in abort_reqs: + if error_code is None: + status_code = HTTPStatus.INTERNAL_SERVER_ERROR + elif isinstance(error_code, HTTPStatus): + status_code = error_code + else: + status_code = HTTPStatus(int(error_code)) + prepare_abort(req, error_msg, status_code=status_code) + self.stream_output([req], req.return_logprob) + return recv_reqs + + def _finalize_shm_features(self, recv_reqs: Optional[List]) -> None: + # Unwrap shared memory features AFTER all broadcasts complete, + # so that ShmPointerMMData metadata (not full tensor data) is what + # gets serialized during broadcast_pyobj. + if recv_reqs: + if self.model_config.is_multimodal and has_shm_features(recv_reqs): + # The broadcast source returns with its original objects while + # peer ranks may still be unpickling ShmPointerMMData + # (-> shm_open). Synchronize the same CPU groups that carried + # SHM-backed work requests before materialize() unlinks them. + if get_parallel().enable_dp_attention: + if self.ps.attn_tp_size > 1: + barrier(group=self.attn_tp_cpu_group) + if self.ps.attn_cp_size > 1: + barrier(group=self.attn_cp_cpu_group) + elif self.ps.tp_size > 1: + barrier(group=self.tp_cpu_group) + for req in recv_reqs: + unwrap_shm_features(req) + + def _split_work_and_control_reqs(self, recv_reqs: List): + work_reqs = [ + req + for req in recv_reqs + if isinstance( + req, + ( + TokenizedGenerateReqInput, + TokenizedEmbeddingReqInput, + BatchTokenizedGenerateReqInput, + BatchTokenizedEmbeddingReqInput, + ), + ) + ] + control_reqs = [ + req + for req in recv_reqs + if not isinstance( + req, + ( + TokenizedGenerateReqInput, + TokenizedEmbeddingReqInput, + BatchTokenizedGenerateReqInput, + BatchTokenizedEmbeddingReqInput, + ), + ) + ] + return work_reqs, control_reqs diff --git a/experiments/pro6000/glm53_ppmtp_r36_degloo/patches/request_receiver_orig.py b/experiments/pro6000/glm53_ppmtp_r36_degloo/patches/request_receiver_orig.py new file mode 100644 index 0000000..cbb6ebb --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r36_degloo/patches/request_receiver_orig.py @@ -0,0 +1,298 @@ +from __future__ import annotations + +from dataclasses import dataclass +from http import HTTPStatus +from typing import ( + TYPE_CHECKING, + Any, + Callable, + List, + Optional, + Union, +) + +import zmq +from torch.distributed import barrier + +from sglang.srt.disaggregation.utils import prepare_abort +from sglang.srt.environ import envs +from sglang.srt.managers.io_struct import ( + BatchTokenizedEmbeddingReqInput, + BatchTokenizedGenerateReqInput, + TokenizedEmbeddingReqInput, + TokenizedGenerateReqInput, + sock_recv, +) +from sglang.srt.managers.mm_utils import ( + has_shm_features, + unwrap_shm_features, +) +from sglang.srt.runtime_context import get_disagg, get_parallel, is_ep_scale_joiner +from sglang.srt.utils import ( + broadcast_pyobj, + point_to_point_pyobj, +) +from sglang.srt.utils.nvtx_utils import scheduler_nvtx_method + +if TYPE_CHECKING: + from sglang.srt.configs.model_config import ModelConfig + from sglang.srt.distributed.parallel_state_wrapper import ParallelState + from sglang.srt.rust_server.server import RustServer + from sglang.srt.server_args import ServerArgs + from sglang.test.scripted_runtime.scheduler_hook import ScriptedSchedulerHook + from sglang.test.scripted_runtime.tokenizer_recv_proxy import ( + ScriptedTokenizerRecvProxy, + ) + + +@dataclass(kw_only=True, slots=True, frozen=True) +class SchedulerRequestReceiver: + recv_from_tokenizer: Union[zmq.Socket, ScriptedTokenizerRecvProxy, RustServer] + recv_from_rpc: Optional[zmq.Socket] + recv_skipper: Any + input_blocker: Any + mm_receiver: Any + ps: ParallelState + tp_group: Any + tp_cpu_group: Any + attn_tp_group: Any + attn_tp_cpu_group: Any + attn_cp_group: Any + attn_cp_cpu_group: Any + world_group: Any + server_args: ServerArgs + model_config: ModelConfig + max_recv_per_poll: int + stream_output: Callable[..., None] + get_last_batch: Callable[[], Any] + scripted_scheduler_hook: Optional[ScriptedSchedulerHook] = None + + def recv_limit_reached(self, num_recv_reqs: int) -> bool: + if self.max_recv_per_poll < 0: + return False + return num_recv_reqs >= self.max_recv_per_poll + + @scheduler_nvtx_method("scheduler.recv_requests") + def recv_requests( + self, + ) -> List[Union[TokenizedGenerateReqInput, TokenizedEmbeddingReqInput, Any]]: + """Receive results at tp_rank = 0 and broadcast it to all other TP ranks.""" + + if self.scripted_scheduler_hook is not None: + self.scripted_scheduler_hook.step() + + if self.recv_skipper is not None: + if not self.recv_skipper.handle(self.get_last_batch()): + return [] + + recv_reqs = self._pull_raw_reqs() + + if self.input_blocker is not None: + recv_reqs = self.input_blocker.handle(recv_reqs) + + recv_reqs = self._broadcast_reqs_across_ranks(recv_reqs) + + if self.ps.pp_rank == 0: + self.unwrap_pickle_wrapper(recv_reqs) + + recv_reqs = self._apply_mm_receiver(recv_reqs) + + self._finalize_shm_features(recv_reqs) + + return recv_reqs + + def _pull_raw_reqs(self) -> Optional[List]: + if self.ps.pp_rank == 0: + if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: + recv_reqs = [] + + # Rust ringbuffer backend: drain the in-process ring fed by the + # embedded Rust TokenizerManager instead of a zmq socket. Same + # non-blocking, msgpack-decoded contract as the zmq path below. + if envs.SGLANG_RUST_SERVER.get(): + recv_reqs.extend( + self.recv_from_tokenizer.drain(self.max_recv_per_poll) + ) + return recv_reqs + + while True: + try: + if self.recv_limit_reached(len(recv_reqs)): + break + recv_req = sock_recv(self.recv_from_tokenizer, zmq.NOBLOCK) + except zmq.ZMQError: + break + recv_reqs.append(recv_req) + + while True: + try: + if self.recv_limit_reached(len(recv_reqs)): + break + recv_rpc = sock_recv(self.recv_from_rpc, zmq.NOBLOCK) + except zmq.ZMQError: + break + recv_reqs.append(recv_rpc) + else: + recv_reqs = None + else: + if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: + dp_offset = ( + self.ps.attn_dp_rank * self.ps.attn_cp_size * self.ps.attn_tp_size + ) + recv_reqs = point_to_point_pyobj( + [], + self.ps.pp_rank * self.ps.tp_size + dp_offset, + self.world_group.cpu_group, + (self.ps.pp_rank - 1) * self.ps.tp_size + dp_offset, + self.ps.pp_rank * self.ps.tp_size + dp_offset, + ) + else: + recv_reqs = None + return recv_reqs + + def _broadcast_reqs_across_ranks(self, recv_reqs: Optional[List]) -> List: + if get_parallel().enable_dp_attention: + if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: + work_reqs, control_reqs = self._split_work_and_control_reqs(recv_reqs) + else: + work_reqs = None + control_reqs = None + + if self.ps.attn_tp_size != 1: + work_reqs = broadcast_pyobj( + work_reqs, + self.attn_tp_group.rank, + self.attn_tp_cpu_group, + src=self.attn_tp_group.ranks[0], + ) + + if self.ps.attn_cp_size != 1: + work_reqs = broadcast_pyobj( + work_reqs, + self.attn_cp_group.rank, + self.attn_cp_cpu_group, + src=self.attn_cp_group.ranks[0], + ) + + # When dp_attention_local_control_broadcast is enabled, each DP + # group leader already receives control messages from the DP + # controller, so we broadcast within attn_tp_group + attn_cp_group + # instead of the full tp_group. This avoids an expensive + # all-ranks gloo sync. + _local_ctrl = ( + get_parallel().enable_dp_attention_local_control_broadcast + or is_ep_scale_joiner() + ) + if _local_ctrl: + if self.ps.attn_tp_size != 1: + control_reqs = broadcast_pyobj( + control_reqs, + self.attn_tp_group.rank, + self.attn_tp_cpu_group, + src=self.attn_tp_group.ranks[0], + ) + if self.ps.attn_cp_size != 1: + control_reqs = broadcast_pyobj( + control_reqs, + self.attn_cp_group.rank, + self.attn_cp_cpu_group, + src=self.attn_cp_group.ranks[0], + ) + elif self.ps.tp_size != 1: + control_reqs = broadcast_pyobj( + control_reqs, + self.tp_group.rank, + self.tp_cpu_group, + src=self.tp_group.ranks[0], + ) + recv_reqs = work_reqs + control_reqs + elif self.ps.tp_size != 1: + recv_reqs = broadcast_pyobj( + recv_reqs, + self.tp_group.rank, + self.tp_cpu_group, + src=self.tp_group.ranks[0], + ) + return recv_reqs + + def unwrap_pickle_wrapper(self, recv_reqs: Optional[List]) -> None: + if not recv_reqs: + return + + for req in recv_reqs: + if isinstance(req, (TokenizedGenerateReqInput, TokenizedEmbeddingReqInput)): + req.unwrap_pickle_fields() + elif isinstance( + req, (BatchTokenizedGenerateReqInput, BatchTokenizedEmbeddingReqInput) + ): + for sub_req in req: + sub_req.unwrap_pickle_fields() + + def _apply_mm_receiver(self, recv_reqs: List) -> List: + # Process MM requests under EPD-disaggregation mode + if ( + self.ps.pp_rank == 0 + and get_disagg().language_only + and get_disagg().encoder_transfer_backend + in ["zmq_to_scheduler", "mooncake"] + ): + recv_reqs, abort_reqs = self.mm_receiver.process_waiting_requests(recv_reqs) + for req, error_msg, error_code in abort_reqs: + if error_code is None: + status_code = HTTPStatus.INTERNAL_SERVER_ERROR + elif isinstance(error_code, HTTPStatus): + status_code = error_code + else: + status_code = HTTPStatus(int(error_code)) + prepare_abort(req, error_msg, status_code=status_code) + self.stream_output([req], req.return_logprob) + return recv_reqs + + def _finalize_shm_features(self, recv_reqs: Optional[List]) -> None: + # Unwrap shared memory features AFTER all broadcasts complete, + # so that ShmPointerMMData metadata (not full tensor data) is what + # gets serialized during broadcast_pyobj. + if recv_reqs: + if self.model_config.is_multimodal and has_shm_features(recv_reqs): + # The broadcast source returns with its original objects while + # peer ranks may still be unpickling ShmPointerMMData + # (-> shm_open). Synchronize the same CPU groups that carried + # SHM-backed work requests before materialize() unlinks them. + if get_parallel().enable_dp_attention: + if self.ps.attn_tp_size > 1: + barrier(group=self.attn_tp_cpu_group) + if self.ps.attn_cp_size > 1: + barrier(group=self.attn_cp_cpu_group) + elif self.ps.tp_size > 1: + barrier(group=self.tp_cpu_group) + for req in recv_reqs: + unwrap_shm_features(req) + + def _split_work_and_control_reqs(self, recv_reqs: List): + work_reqs = [ + req + for req in recv_reqs + if isinstance( + req, + ( + TokenizedGenerateReqInput, + TokenizedEmbeddingReqInput, + BatchTokenizedGenerateReqInput, + BatchTokenizedEmbeddingReqInput, + ), + ) + ] + control_reqs = [ + req + for req in recv_reqs + if not isinstance( + req, + ( + TokenizedGenerateReqInput, + TokenizedEmbeddingReqInput, + BatchTokenizedGenerateReqInput, + BatchTokenizedEmbeddingReqInput, + ), + ) + ] + return work_reqs, control_reqs diff --git a/experiments/pro6000/glm53_ppmtp_r36_degloo/patches/scheduler_pp_mixin_r35.py b/experiments/pro6000/glm53_ppmtp_r36_degloo/patches/scheduler_pp_mixin_r35.py new file mode 100644 index 0000000..08df3a6 --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r36_degloo/patches/scheduler_pp_mixin_r35.py @@ -0,0 +1,2389 @@ +from __future__ import annotations + +import logging +import math +import time +from array import array +from collections import defaultdict, deque +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple + +import numpy as np +import torch +import torch.distributed +from tqdm import tqdm + +import os +import pickle +# [EXPERIMENT] PP+spec debug probes (see _pp_replay_spec_round). +_PP_SPEC_DEBUG = os.getenv("SGLANG_PP_SPEC_DEBUG", "0") == "1" + +# [r35 de-GLOO] Boundary control channel off GLOO (see _pp_send_tensor_dict_nccl). +# On by default; SGLANG_PP_DEGLOO=0 reverts to the r33 GLOO metadata path. +_PP_DEGLOO = os.getenv("SGLANG_PP_DEGLOO", "1") == "1" +# Fixed-size metadata buffer. The pickle stream is zero-padded into it; +# pickle.loads stops at the STOP opcode and ignores the padding (verified), +# so no length prefix and no .item() host sync is needed on the receiver. +_PP_META_BUF_BYTES = 8192 + +from sglang.srt.disaggregation.base.conn import KVPoll +from sglang.srt.disaggregation.utils import poll_and_all_reduce_attn_cp_tp_group +from sglang.srt.distributed.parallel_state import ( + P2PWork, + TensorMetadata, + _split_tensor_dict, +) +from sglang.srt.environ import envs +from sglang.srt.layers.dp_attention import ( + get_attention_dp_rank, + get_attention_dp_size, + is_dp_attention_enabled, + set_is_extend_in_batch, +) +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.managers.overlap_utils import RelayPayload +from sglang.srt.managers.schedule_batch import FINISH_ABORT, Req, ScheduleBatch +from sglang.srt.managers.utils import ( + GenerationBatchResult, + get_logprob_dict_from_result, + get_logprob_from_pp_outputs, +) +from sglang.srt.mem_cache.common import release_kv_cache +from sglang.srt.model_executor.forward_batch_info import ( + ForwardBatch, + ForwardMode, + PPProxyTensors, +) +from sglang.srt.observability.req_time_stats import set_time_batch +from sglang.srt.runtime_context import get_disagg, get_parallel +from sglang.srt.sampling.sampling_observer_pp import ( + add_auxiliary_output_to_pp_tensors, + pop_auxiliary_output_from_pp_tensors, +) +from sglang.srt.sampling.sampling_params import SamplingParams +from sglang.srt.utils import DynamicGradMode, broadcast_pyobj, point_to_point_pyobj +from sglang.srt.utils.common import get_device_module, is_xpu + +logger = logging.getLogger(__name__) + +if TYPE_CHECKING: + from sglang.srt.managers.scheduler import Scheduler + + +def _pp_can_skip_output_comm(batch: ScheduleBatch) -> bool: + """Check if output send/recv can be skipped for this batch.""" + # [EXPERIMENT] PP+spec: every round's relayed state (hidden/accept) feeds + # the non-last stages' draft replicas. Never skip output comm under + # speculation — a skipped round desynchronizes the replicas. + try: + from sglang.srt.runtime_context import get_spec + + spec = get_spec() + if spec is not None and getattr(spec, "speculative_algorithm", None): + return False + except Exception: + return False + return ( + envs.SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM.get() + and batch is not None + and batch.forward_mode == ForwardMode.EXTEND + and len(batch.reqs) == 1 + and not batch.contains_last_prefill_chunk + and not batch.return_logprob + ) + + +@dataclass +class PPBatchMetadata: + can_run_cuda_graph: bool + + +class SchedulerPPMixin: + @DynamicGradMode() + def event_loop_pp(self: Scheduler): + """ + A scheduler loop for pipeline parallelism. + Notes: + 1. Each stage runs in the same order and is notified by the previous stage. + 2. We use async send but sync recv to avoid desynchronization while minimizing the communication overhead. + 3. We can use async batch depth to buffer the outputs in the last stage for to allow overlapping the GPU computation and CPU processing and avoid last PP rank staggler. + + Unified Schedule: + ==================================================================== + Stage P + recv ith req from previous stage + recv ith proxy from previous stage + run ith batch + recv prev (i+1)% mb_size th outputs + process batch result of prev (i+1)% mb_size th batch (can be run in parallel with the curr batch GPU computation) + send ith req to next stage + send ith proxy to next stage + send current stage's outputs to next stage(can be stashed and delayed to send later) + + the above order can be optimized and reordered to minimize communication-related CPU stall and overhead bubbles. + + ==================================================================== + """ + self.init_pp_loop_state() + while True: + server_is_idle = True + for mb_id in range(self.pp_loop_size): + self.running_batch = self.running_mbs[mb_id] + self.last_batch = self.last_mbs[mb_id] + next_first_rank_mb_id = (mb_id + self.ps.pp_size) % self.pp_loop_size + next_mb_id = (mb_id + 1) % self.pp_loop_size + with torch.profiler.record_function("recv_requests"): + recv_reqs = self.request_receiver.recv_requests() + self.process_input_requests(recv_reqs) + if not self.pp_group.is_last_rank: + self._pp_commit_comm_work(self.send_req_work) + with torch.profiler.record_function("send_reqs_to_next_stage"): + self.send_req_work = self._pp_send_pyobj_to_next_stage( + recv_reqs, + async_send=True, + ) + with torch.profiler.record_function("get_next_batch_to_run"): + plan = self.get_next_batch_to_run( + running_batch=self.running_batch, last_batch=self.last_batch + ) + self.running_batch = plan.running_batch + self.mbs[mb_id] = plan.batch_to_run + self.running_mbs[mb_id] = self.running_batch + cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id] + self.cur_batch_for_debug = cur_batch + if _PP_SPEC_DEBUG: + _nt = "" + if cur_batch is not None: + _nt = ( + f" tokens={getattr(cur_batch, 'extend_num_tokens', None)}" + if cur_batch.forward_mode.is_extend() + else f" bs={len(cur_batch.reqs)}" + ) + logger.info( + f"[PPSPEC] LOOP pp_rank={self.pp_group.rank_in_group} " + f"mb={mb_id} cur=" + f"{None if cur_batch is None else cur_batch.forward_mode.name}" + f"{_nt} " + f"sched_done={self.schedule_stream.query()} " + f"fwd_done={self.forward_stream.query()}" + ) + if cur_batch: + server_is_idle = False + pp_proxy_tensors = self._pp_recv_proxy_tensors() + next_pp_outputs = None + next_batch_result = None + d2h_event = None + if get_parallel().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, + ) + ) + self._pp_commit_comm_work(self.send_proxy_work) + if cur_batch: + result, self.launch_event = self._pp_launch_batch( + mb_id, + cur_batch, + pp_proxy_tensors, + self.mb_metadata, + self.last_rank_comm_queue, + ) + if get_parallel().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, + ) + ) + if self.mbs[next_mb_id] is not None: + d2h_event.synchronize() + with torch.profiler.record_function("process_batch_result"): + self._pp_process_batch_result( + self.mbs[next_mb_id], + next_batch_result, + ) + self.last_mbs[next_mb_id] = self.mbs[next_mb_id] + if not self.pp_group.is_last_rank: + if cur_batch: + self.device_module.current_stream().wait_event( + self.launch_event + ) + with torch.profiler.record_function( + "send_proxy_dict_to_next_stage" + ): + self.send_proxy_work = self._pp_send_dict_to_next_stage( + result.pp_hidden_states_proxy_tensors.tensors, + async_send=True, + msg_type="proxy", + ) + + self.pp_outputs = next_pp_outputs + + # When the server is idle, self-check and re-init some states + if server_is_idle: + self.on_idle() + + @DynamicGradMode() + def event_loop_pp_disagg_prefill(self: Scheduler): + """ + This is the prefill server event loop for pipeline parallelism. + + Notes: + 1. Following the same rules as the event_loop_pp. + 2. Adds extra steps for KV transfer process: bootstrap + release. + + Prefill Server Schedule: + ==================================================================== + Stage P + recv ith req from previous stage + recv ith bootstrap req from previous stage + recv ith transferred req from previous stage + recv ith proxy from previous stage + run ith batch + recv prev (i+1) % mb_size th consensus bootstrapped req from previous stage + local consensus on bootstrapped req + recv prev (i+1) % mb_size th release req from previous stage + local consensus on release req + recv prev (i+1) % mb_size th outputs + process batch result of prev (i+1)% mb_size th batch (can be run in parallel with the curr batch GPU computation) + send ith req to next stage + send ith bootstrap req to next stage + send ith transferred req to next stage + send ith proxy to next stage + send current stage's outputs to next stage (can be stashed and delayed to send later) + + the above order can be optimized and reordered to minimize communication-related CPU stall and overhead bubbles. + ==================================================================== + + There are two additional elements compared to the regular schedule: + + Bootstrap Requests + Release Requests: + - Both can have local failure and need to be consensus on. PP needs to guarantee eventual consistency of local failure and flush malfunc requests out as soft error. + + """ + self.init_pp_loop_state() + + # PD additional state initialization + bmbs = [None] * self.pp_loop_size + tmbs = [None] * self.pp_loop_size + consensus_bootstrapped_rids: Optional[List[str]] = None + transferred_rids: List[str] = [] + release_rids: Optional[List[str]] = None + send_bootstrapped_work = [] + send_transfer_work = [] + send_consensus_bootstrapped_work = [] + send_release_work = [] + + while True: + server_is_idle = True + for mb_id in range(self.pp_loop_size): + self.running_batch = self.running_mbs[mb_id] + self.last_batch = self.last_mbs[mb_id] + next_first_rank_mb_id = (mb_id + self.ps.pp_size) % self.pp_loop_size + next_mb_id = (mb_id + 1) % self.pp_loop_size + + next_pp_outputs = None + next_release_rids = None + next_consensus_bootstrapped_rids = None + d2h_event = None + next_batch_result = None + + recv_reqs = self.request_receiver.recv_requests() + self.process_input_requests(recv_reqs) + + if not self.pp_group.is_last_rank: + self._pp_commit_comm_work(self.send_req_work) + + bootstrapped_rids = self._pp_pd_get_bootstrapped_ids() + bmbs[mb_id] = bootstrapped_rids + self._pp_commit_comm_work(send_bootstrapped_work) + + transferred_rids = self._pp_pd_get_prefill_transferred_ids() + self._pp_commit_comm_work(send_transfer_work) + tmbs[mb_id] = transferred_rids + + self.process_prefill_chunk( + last_batch=self.last_batch, running_batch=self.running_batch + ) + prefill_plan = self.get_new_batch_prefill(self.running_batch) + batch = prefill_plan.batch_to_run + self.running_batch = prefill_plan.running_batch + batch = self.dp_attn_adapter.maybe_prepare_mlp_sync_batch(batch) + self.mbs[mb_id] = batch + self.running_mbs[mb_id] = self.running_batch + + cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id] + self.cur_batch_for_debug = cur_batch + if cur_batch: + server_is_idle = False + pp_proxy_tensors = self._pp_recv_proxy_tensors() + + if get_parallel().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, + ) + ) + self._pp_commit_comm_work(self.send_proxy_work) + if cur_batch: + if self.enable_staging: + self.maybe_prefetch_staging_for_batch(cur_batch) + result, self.launch_event = self._pp_launch_batch( + mb_id, + cur_batch, + pp_proxy_tensors, + self.mb_metadata, + self.last_rank_comm_queue, + ) + if get_parallel().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, + ) + ) + send_consensus_bootstrapped_work, consensus_bootstrapped_rids = ( + self._pp_pd_send_consensus_bootstrapped_ids( + bmbs, + next_first_rank_mb_id, + consensus_bootstrapped_rids, + bootstrapped_rids, + ) + ) + send_release_work, release_rids = ( + self._pp_pd_send_consensus_release_ids( + tmbs, next_first_rank_mb_id, release_rids, transferred_rids + ) + ) + + if bmbs[next_mb_id] is not None: + next_consensus_bootstrapped_rids = ( + self._pp_recv_pyobj_from_prev_stage() + ) + next_consensus_bootstrapped_rids = self.process_bootstrapped_queue( + next_consensus_bootstrapped_rids + ) + self._pp_commit_comm_work(send_consensus_bootstrapped_work) + if tmbs[next_mb_id] is not None: + next_release_rids = self._pp_recv_pyobj_from_prev_stage() + self._pp_commit_comm_work(send_release_work) + # post-process the coming microbatch + 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, + ) + self.last_mbs[next_mb_id] = self.mbs[next_mb_id] + + if tmbs[next_mb_id] is not None: + self.process_disagg_prefill_inflight_queue(next_release_rids) + if not self.pp_group.is_last_rank: + self.send_req_work = self._pp_send_pyobj_to_next_stage( + recv_reqs, async_send=True + ) + send_bootstrapped_work = self._pp_send_pyobj_to_next_stage( + bootstrapped_rids, async_send=True + ) + send_transfer_work = self._pp_send_pyobj_to_next_stage( + transferred_rids, async_send=True + ) + if cur_batch: + self.device_module.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", + ) + + self.pp_outputs = next_pp_outputs + release_rids = next_release_rids + consensus_bootstrapped_rids = next_consensus_bootstrapped_rids + + self.running_batch.batch_is_full = False + + # When the server is idle, self-check and re-init some states + if server_is_idle and len(self.disagg_prefill_inflight_queue) == 0: + self.on_idle() + + @DynamicGradMode() + def event_loop_pp_disagg_decode(self: Scheduler): + self.init_pp_loop_state() + + # PD additional state initialization + rmbs = [None] * self.pp_loop_size + pmbs = [None] * self.pp_loop_size + tmbs = [None] * self.pp_loop_size + consensus_retract_rids: Optional[List[str]] = None + consensus_prealloc_rids: Optional[List[str]] = None + release_rids: Optional[List[str]] = None # consensus transferred rids + send_retract_work = [] + send_prealloc_work = [] + send_transfer_work = [] + send_consensus_retract_work = [] + send_consensus_prealloc_work = [] + send_release_work = [] + + while True: + server_is_idle = True + for mb_id in range(self.pp_loop_size): + self.running_batch = self.running_mbs[mb_id] + self.last_batch = self.last_mbs[mb_id] + next_first_rank_mb_id = (mb_id + self.ps.pp_size) % self.pp_loop_size + next_mb_id = (mb_id + 1) % self.pp_loop_size + + next_pp_outputs = None + next_consensus_retract_rids = None + next_consensus_prealloc_rids = None + next_release_rids = None + d2h_event = None + next_batch_result = None + + recv_reqs = self.request_receiver.recv_requests() + self.process_input_requests(recv_reqs) + + if not self.pp_group.is_last_rank: + self._pp_commit_comm_work(self.send_req_work) + + # reaching consensus through PP ranks + retract_rids = self._pp_pd_get_retract_ids(mb_id) + rmbs[mb_id] = retract_rids + self._pp_commit_comm_work(send_retract_work) + + prealloc_rids = self._pp_pd_get_prealloc_ids() + pmbs[mb_id] = prealloc_rids + self._pp_commit_comm_work(send_prealloc_work) + + transferred_rids = self._pp_pd_get_decode_transferred_ids() + tmbs[mb_id] = transferred_rids + self._pp_commit_comm_work(send_transfer_work) + + # get batch to run and proxy tensors if needed + plan = self.get_next_disagg_decode_batch_to_run( + running_batch=self.running_batch + ) + self.running_batch = plan.running_batch + batch = plan.batch_to_run + self.mbs[mb_id] = batch + self.running_mbs[mb_id] = self.running_batch + + cur_batch: Optional[ScheduleBatch] = self.mbs[mb_id] + self.cur_batch_for_debug = cur_batch + if cur_batch: + server_is_idle = False + pp_proxy_tensors = None + if not cur_batch.forward_mode.is_prebuilt(): + pp_proxy_tensors = self._pp_recv_proxy_tensors() + + # early send output if possible + if get_parallel().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, + ) + ) + self._pp_commit_comm_work(self.send_proxy_work) + + if cur_batch: + result, self.launch_event = self._pp_launch_batch( + mb_id, + cur_batch, + pp_proxy_tensors, + self.mb_metadata, + self.last_rank_comm_queue, + ) + + if get_parallel().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, + ) + ) + + # reach consensus on last rank and send to PP=0 + # otherwise, just pass along previous consensus + send_consensus_retract_work, consensus_retract_rids = ( + self._pp_pd_send_consensus_bootstrapped_ids( + rmbs, + next_first_rank_mb_id, + consensus_retract_rids, + retract_rids, + ) + ) + + send_consensus_prealloc_work, consensus_prealloc_rids = ( + self._pp_pd_send_consensus_bootstrapped_ids( + pmbs, + next_first_rank_mb_id, + consensus_prealloc_rids, + prealloc_rids, + ) + ) + + send_release_work, release_rids = ( + self._pp_pd_send_consensus_release_ids( + tmbs, next_first_rank_mb_id, release_rids, transferred_rids + ) + ) + + if get_disagg().disaggregation_decode_enable_offload_kvcache: + self.decode_offload_manager.check_offload_progress() + + if rmbs[next_mb_id] is not None: + next_consensus_retract_rids = self._pp_recv_pyobj_from_prev_stage() + next_consensus_retract_rids = self.process_retract_queue( + next_consensus_retract_rids + ) + self._pp_commit_comm_work(send_consensus_retract_work) + + if pmbs[next_mb_id] is not None: + next_consensus_prealloc_rids = self._pp_recv_pyobj_from_prev_stage() + next_consensus_prealloc_rids = self.process_prealloc_queue( + next_consensus_prealloc_rids + ) + self._pp_commit_comm_work(send_consensus_prealloc_work) + + if tmbs[next_mb_id] is not None: + next_release_rids = self._pp_recv_pyobj_from_prev_stage() + next_release_rids = self.process_decode_transfer_queue( + next_release_rids + ) + self._pp_commit_comm_work(send_release_work) + + # post-process the coming microbatch + if self.mbs[next_mb_id] is not None: + if not self.mbs[next_mb_id].forward_mode.is_prebuilt(): + d2h_event.synchronize() + self._pp_process_batch_result( + self.mbs[next_mb_id], + next_batch_result, + ) + self.last_mbs[next_mb_id] = self.mbs[next_mb_id] + + if not self.pp_group.is_last_rank: + self.send_req_work = self._pp_send_pyobj_to_next_stage( + recv_reqs, async_send=True + ) + send_retract_work = self._pp_send_pyobj_to_next_stage( + retract_rids, async_send=True + ) + send_prealloc_work = self._pp_send_pyobj_to_next_stage( + prealloc_rids, async_send=True + ) + send_transfer_work = self._pp_send_pyobj_to_next_stage( + transferred_rids, async_send=True + ) + if cur_batch and not cur_batch.forward_mode.is_prebuilt(): + self.device_module.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", + ) + + self.pp_outputs = next_pp_outputs + release_rids = next_release_rids + consensus_retract_rids = next_consensus_retract_rids + consensus_prealloc_rids = next_consensus_prealloc_rids + + self.running_batch.batch_is_full = False + + # When the server is idle, self-check and re-init some states + queue_size = ( + len(self.waiting_queue) + + len(self.disagg_decode_transfer_queue.queue) + + len(self.disagg_decode_prealloc_queue.queue) + ) + if get_disagg().disaggregation_decode_enable_offload_kvcache: + queue_size += len(self.decode_offload_manager.ongoing_offload) + + if server_is_idle and queue_size == 0: + self.on_idle() + + def init_pp_loop_state(self: Scheduler): + self.pp_loop_size: int = self.ps.pp_size + get_parallel().pp_async_batch_depth + # In CP mode, attention weights are duplicated, eliminating the need for the attention TP all-gather operation. + self.require_attn_tp_allgather = ( + not get_parallel().enable_dsa_prefill_context_parallel + ) + self.mbs = [None] * self.pp_loop_size + self.last_mbs = [None] * self.pp_loop_size + self.running_mbs = [ + ScheduleBatch(reqs=[], batch_is_full=False) + for _ in range(self.pp_loop_size) + ] + self.mb_metadata: List[Optional[PPBatchMetadata]] = [None] * self.pp_loop_size + self.pp_outputs: Optional[PPProxyTensors] = None + self.last_rank_comm_queue: deque[Tuple[torch.Event, PPProxyTensors]] = deque() + + self.send_req_work = [] + self.send_proxy_work = [] + self.send_output_work = [] + self.launch_event = None + self._pp_tensor_dict_inbox: Dict[str, deque[Dict[str, torch.Tensor]]] = ( + defaultdict(deque) + ) + + def profile_and_init_predictor(self: Scheduler): + """ + Profile prefill latency for dynamic chunk sizing. + + Only runs on PP0 (first rank), then broadcasts data to all ranks. + All ranks fit coefficients using the same data. + """ + seq_lens: List[int] = [] + latencies: List[float] = [] + + if self.pp_group.is_first_rank: + model_runner = self.tp_worker.model_runner + model_config = model_runner.model_config + input_ids_list: List[array[int]] = [] + for i in range(128): + chunk_size = int( + self.chunked_prefill_size * 1.25 + - i * (self.chunked_prefill_size * 1.25 // 128) + ) + if chunk_size <= 0: + break + input_ids = array( + "q", + np.random.randint( + 0, 10000, size=chunk_size, dtype=np.int64 + ).tobytes(), + ) + input_ids_list.append(input_ids) + + sampling_params = SamplingParams( + temperature=0, + max_new_tokens=1, + ) + # Create and profile requests + for i, input_ids in enumerate( + tqdm( + input_ids_list, + desc="Profiling prefill latency for dynamic chunking", + ) + ): + req = Req( + rid=str(i), + origin_input_text="", + origin_input_ids=input_ids, + sampling_params=sampling_params, + ) + req.full_untruncated_fill_ids = req.origin_input_ids + req.logprob_start_len = -1 + req.set_extend_range( + len(req.prefix_indices), len(req.full_untruncated_fill_ids) + ) + + # Prepare batch + batch = ScheduleBatch.init_new( + [req], + self.req_to_token_pool, + self.token_to_kv_pool_allocator, + self.tree_cache, + self.model_config, + False, + self.spec_algorithm, + ) + + current_seq_len = req.extend_range.end + + if is_dp_attention_enabled(): + # For profiling, we only have one request on PP0 + # Set global_num_tokens to indicate this rank has tokens, others have 0 + dp_size = get_attention_dp_size() + global_num_tokens = [0] * dp_size + dp_rank = get_attention_dp_rank() + global_num_tokens[dp_rank] = current_seq_len + batch.global_num_tokens = global_num_tokens + batch.global_num_tokens_for_logprob = global_num_tokens + + hs = ( + getattr(model_config, "hc_hidden_size", None) + or model_config.hidden_size + ) + proxy_tensors = { + "hidden_states": torch.zeros( + (current_seq_len, hs), + dtype=model_config.dtype, + device=self.device, + ), + "residual": torch.zeros( + (current_seq_len, model_config.hidden_size), + dtype=model_config.dtype, + device=self.device, + ), + } + pp_proxy_topk_size = model_runner.get_pp_proxy_topk_size() + if pp_proxy_topk_size is not None: + proxy_tensors["topk_indices"] = torch.zeros( + (current_seq_len, pp_proxy_topk_size), + dtype=torch.int32, + device=self.device, + ) + + pp_proxy = PPProxyTensors(proxy_tensors) + + # Measure latency with device synchronization for accurate timing + device_module = get_device_module() + # Synchronize before starting timing to ensure clean measurement + device_module.synchronize() + + start = time.perf_counter() + batch.prepare_for_extend() + + # Resolve deferred H2D: prepare_for_extend now leaves input_ids=None + if batch.input_ids is None and batch.prefill_input_ids_cpu is not None: + batch.input_ids = batch.prefill_input_ids_cpu.to( + self.device, non_blocking=True + ) + batch.prefill_input_ids_cpu = None + + forward_batch = ForwardBatch.init_new( + batch, + model_runner, + return_hidden_states_before_norm=False, + ) + set_is_extend_in_batch(batch.forward_mode.is_extend()) + + _ = model_runner.forward( + forward_batch=forward_batch, pp_proxy_tensors=pp_proxy + ) + + # Synchronize after forward to ensure GPU operations complete + device_module.synchronize() + + latency_seconds = time.perf_counter() - start + latency_ms = latency_seconds * 1e3 # Convert to milliseconds + seq_lens.append(len(input_ids)) + latencies.append(latency_ms) + + # Release KV and Mamba cache + if req.kv.holds_kv: + release_kv_cache(req, self.tree_cache, is_insert=False) + + logger.info( + f"[PP Dynamic Chunk] [PP0] Profiled {len(seq_lens)} samples: " + f"seq_lens={seq_lens}, latencies_ms={latencies}" + ) + + if self.ps.attn_tp_size > 1: + data_to_sync_tp = [seq_lens, latencies] + data_to_sync_tp = broadcast_pyobj( + data_to_sync_tp, + self.attn_tp_group.rank, + self.attn_tp_cpu_group, + src=self.attn_tp_group.ranks[0], + ) + seq_lens, latencies = data_to_sync_tp + + if self.ps.attn_cp_size > 1: + data_to_sync_tp = [seq_lens, latencies] + data_to_sync_tp = broadcast_pyobj( + data_to_sync_tp, + self.attn_cp_group.rank, + self.attn_cp_cpu_group, + src=self.attn_cp_group.ranks[0], + ) + + # Broadcast data to all ranks + if torch.distributed.is_available() and torch.distributed.is_initialized(): + data_to_sync = [seq_lens, latencies] + self.pp_group.broadcast_object_list(data_to_sync, src=0) + seq_lens, latencies = data_to_sync + + # Quadratic model: f(l) = al^2 + bl + c + self.length_predictor = ChunkSizePredictor() + self.length_predictor.fit(seq_lens, latencies) + self.length_predictor.set_target_latency(self.chunked_prefill_size) + self.length_predictor.is_ready = True + logger.info( + f"[PP Dynamic Chunk] [PP{self.ps.pp_rank}] Predictor ready (quadratic). " + f"Target latency: {self.length_predictor.target_latency:.2f}ms" + ) + + def predict_next_chunk_size(self: Scheduler, history_len: int) -> Optional[int]: + """ + Predict next chunk size dynamically based on current history length. + + Args: + history_len: Current sequence length + + Returns: + Predicted chunk size, or None to use default chunked_prefill_size + """ + if ( + not self.enable_dynamic_chunking + or self.length_predictor is None + or not self.length_predictor.is_ready + ): + return None + + max_chunk_size = self.max_prefill_tokens + predicted_size = self.length_predictor.predict_next_chunk_size( + history_len=history_len, + base_chunk_size=self.chunked_prefill_size, + page_size=self.page_size, + context_len=self.model_config.context_len, + max_chunk_size=max_chunk_size, + ) + + if predicted_size is not None: + logger.debug( + f"[PP Dynamic Chunk] [PP{self.ps.pp_rank}] Predicted chunk size: " + f"{predicted_size} (history_len={history_len})" + ) + + return predicted_size + + def process_bootstrapped_queue( + self: Scheduler, bootstrapped_rids: Optional[List[str]] + ): + # finished consensus bootstrapped reqs and prepare the waiting queue + if bootstrapped_rids is not None: + ( + good_consensus_bootstrapped_rids, + bad_consensus_bootstrapped_rids, + ) = bootstrapped_rids + good_reqs, failed_reqs = ( + self.disagg_prefill_bootstrap_queue.pop_bootstrapped( + return_failed_reqs=True, + pp_good_rids=good_consensus_bootstrapped_rids, + pp_bad_rids=bad_consensus_bootstrapped_rids, + ) + ) + self.waiting_queue.extend(good_reqs) + return [[req.rid for req in good_reqs], [req.rid for req in failed_reqs]] + return None + + def _pp_pd_get_bootstrapped_ids(self: Scheduler): + # communicate pre-consensus bootstrapp reqs + if self.pp_group.is_first_rank: + # First rank, pop the bootstrap reqs from the bootstrap queue + good_bootstrapped_rids, bad_bootstrapped_rids = self.get_rids( + self.disagg_prefill_bootstrap_queue.queue, + True, + [KVPoll.WaitingForInput], + [KVPoll.Failed], + ) + else: + # Other ranks, receive the bootstrap reqs info from the previous rank and ensure the consensus + prev_bootstrapped_rids = self._pp_recv_pyobj_from_prev_stage() + prev_good_bootstrapped_rids, prev_bad_bootstrapped_rids = ( + prev_bootstrapped_rids + ) + curr_good_bootstrapped_rids, curr_bad_bootstrapped_rids = self.get_rids( + self.disagg_prefill_bootstrap_queue.queue, + True, + [KVPoll.WaitingForInput], + [KVPoll.Failed], + ) + good_bootstrapped_rids = list( + set(prev_good_bootstrapped_rids) & set(curr_good_bootstrapped_rids) + ) + bad_bootstrapped_rids = list( + set(prev_bad_bootstrapped_rids) | set(curr_bad_bootstrapped_rids) + ) + # Route locally-aborted reqs through the bad-union consensus so every PP + # rank flushes them in the same consensus round, regardless of when the + # AbortReq reaches each rank and regardless of whether + # disagg_kv_sender.abort() drives the poll to Failed (it is optional). + aborted_rids = { + req.rid + for req in self.disagg_prefill_bootstrap_queue.queue + if isinstance(req.finished_reason, FINISH_ABORT) + } + good_bootstrapped_rids, bad_bootstrapped_rids = self._route_aborts_to_bad( + good_bootstrapped_rids, bad_bootstrapped_rids, aborted_rids + ) + return [good_bootstrapped_rids, bad_bootstrapped_rids] + + def _pp_pd_get_prefill_transferred_ids(self: Scheduler): + # get the current stage transfer success + if self.pp_group.is_first_rank: + transferred_rids = self.get_rids( + self.disagg_prefill_inflight_queue, + True, + [KVPoll.Success, KVPoll.Failed], + ) + # if other ranks, do intersection with the previous rank's transferred rids + else: + # 2 (Release): Receive the transferred rids from the previous rank + # 1. recv previous stage's transferred reqs info + prev_transferred_rids = self._pp_recv_pyobj_from_prev_stage() + # 2. get the current stage's transferred reqs info + curr_transferred_rids = self.get_rids( + self.disagg_prefill_inflight_queue, + True, + [KVPoll.Success, KVPoll.Failed], + ) + # 3. new consensus rids = intersection(previous consensus rids, transfer finished rids) + transferred_rids = list( + set(prev_transferred_rids) & set(curr_transferred_rids) + ) + return transferred_rids + + def _pp_pd_send_consensus_bootstrapped_ids( + self: Scheduler, + bmbs: List[List[str]], + next_first_rank_mb_id: int, + consensus_bootstrapped_rids: List[str], + bootstrapped_rids: List[str], + ): + # 3 (Release): send the release rids from last stage to the first stage + send_consensus_bootstrapped_work = [] + if self.pp_group.is_last_rank: + if bmbs[next_first_rank_mb_id] is not None: + consensus_bootstrapped_rids = bootstrapped_rids + send_consensus_bootstrapped_work = self._pp_send_pyobj_to_next_stage( + consensus_bootstrapped_rids, async_send=True + ) + # 4 (Release): send the release rids from non last rank to the next rank + else: + if consensus_bootstrapped_rids is not None: + send_consensus_bootstrapped_work = self._pp_send_pyobj_to_next_stage( + consensus_bootstrapped_rids, async_send=True + ) + return send_consensus_bootstrapped_work, consensus_bootstrapped_rids + + def _pp_pd_send_consensus_release_ids( + self: Scheduler, + tmbs: List[List[str]], + next_first_rank_mb_id: int, + release_rids: List[str], + transferred_rids: List[str], + ): + send_release_work = [] + if self.pp_group.is_last_rank: + if tmbs[next_first_rank_mb_id] is not None: + release_rids = transferred_rids + send_release_work = self._pp_send_pyobj_to_next_stage( + release_rids, async_send=True + ) + # 4 (Release): send the release rids from non last rank to the next rank + else: + if release_rids is not None: + send_release_work = self._pp_send_pyobj_to_next_stage( + release_rids, async_send=True + ) + return send_release_work, release_rids + + def _pp_commit_comm_work(self: Scheduler, work: List[P2PWork]) -> None: + for p2p_work in work: + p2p_work.work.wait() + work.clear() + + def _pp_commit_send_output_work_and_preprocess_output_tensors( + self: Scheduler, + next_first_rank_mb_id: int, + next_mb_id: int, + ) -> Tuple[ + Optional[PPProxyTensors], + Optional[GenerationBatchResult], + Optional[torch.Event], + ]: + self._pp_commit_comm_work(work=self.send_output_work) + ( + 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 + + def _pp_send_pyobj_to_next_stage(self: Scheduler, data, async_send: bool = False): + p2p_work = [] + if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: + dp_offset = ( + self.ps.attn_dp_rank * self.ps.attn_cp_size * self.ps.attn_tp_size + ) + if _PP_DEGLOO: + # [r35 de-GLOO] Request relay over NCCL (world device group, + # a separate communicator from the pp dict channel, so its + # FIFO ordering can never interleave with proxy/output dict + # traffic). The r33 path used point_to_point_pyobj on + # world_group.cpu_group: a BLOCKING gloo dist.send every + # iteration (even for empty req lists) that rendezvous-stalled + # the sender for hundreds of ms under stage phase drift, + # skewing its TP partners' allreduces. NCCL isend never + # blocks the host; the receiver's waits are bounded by one + # small transfer. + p2p_work = self._pp_send_pyobj_nccl( + data, + ((self.ps.pp_rank + 1) % self.ps.pp_size) * self.ps.tp_size + + dp_offset, + ) + else: + p2p_work = point_to_point_pyobj( + data, + self.ps.pp_rank * self.ps.tp_size + dp_offset, + self.world_group.cpu_group, + self.ps.pp_rank * self.ps.tp_size + dp_offset, + ((self.ps.pp_rank + 1) % self.ps.pp_size) * self.ps.tp_size + dp_offset, + async_send=async_send, + ) + return p2p_work + + def _pp_nccl_stage_guard(self: Scheduler, tensors: List[torch.Tensor]) -> None: + """[r35 de-GLOO] Cross-stream staging guard for freshly built NCCL + send buffers. + + The boundary isend empirically executes on the scheduler's ambient + stream rather than the calling context's current stream (see the + q_event note in _pp_send_output_to_next_stage). Buffers built on the + current stream must therefore (a) gate the ambient stream on the + fill BEFORE the isend is enqueued, and (b) record_stream on both + candidates so the allocator delays reuse either way — the same + hedge the caller applies to dict tensors in + _pp_send_dict_to_next_stage. Call this right before the isends. + """ + _fill_ev = torch.cuda.Event() + _fill_ev.record() + self.schedule_stream.wait_event(_fill_ev) + _cur_stream = self.device_module.current_stream() + for _t in tensors: + _t.record_stream(self.schedule_stream) + if _cur_stream != self.schedule_stream: + _t.record_stream(_cur_stream) + + def _pp_send_pyobj_nccl( + self: Scheduler, data, global_dst: int + ) -> List[P2PWork]: + """[r35 de-GLOO] Two-phase NCCL object send: int64 size, then payload. + + Mirrors point_to_point_pyobj's empty optimization (size=0 means no + payload follows) so the receiver stays in FIFO lockstep. Buffers are + held by the returned P2PWorks until commit, exactly like the GLOO + path, so the caching allocator sees the same lifetime. + """ + device = self.device_module.current_device() + group = self.world_group.device_group + p2p_works: List[P2PWork] = [] + if len(data) == 0: + size_tensor = torch.zeros(1, dtype=torch.int64, device=device) + self._pp_nccl_stage_guard([size_tensor]) + work = torch.distributed.isend(size_tensor, global_dst, group=group) + p2p_works.append(P2PWork(work, size_tensor)) + else: + payload = pickle.dumps(data) + payload_tensor = ( + torch.frombuffer(bytearray(payload), dtype=torch.uint8) + .to(device, non_blocking=False) + .contiguous() + ) + size_tensor = torch.tensor( + [payload_tensor.numel()], dtype=torch.int64, device=device + ) + self._pp_nccl_stage_guard([size_tensor, payload_tensor]) + work = torch.distributed.isend(size_tensor, global_dst, group=group) + p2p_works.append(P2PWork(work, size_tensor)) + work = torch.distributed.isend(payload_tensor, global_dst, group=group) + p2p_works.append(P2PWork(work, payload_tensor)) + return p2p_works + + def _pp_recv_pyobj_nccl(self: Scheduler, global_src: int): + """[r35 de-GLOO] Receiving side of _pp_send_pyobj_nccl.""" + device = self.device_module.current_device() + group = self.world_group.device_group + size_tensor = torch.empty(1, dtype=torch.int64, device=device) + work = torch.distributed.irecv(size_tensor, global_src, group=group) + work.wait() + size = int(size_tensor.item()) + if size == 0: + return [] + payload_tensor = torch.empty(size, dtype=torch.uint8, device=device) + work = torch.distributed.irecv(payload_tensor, global_src, group=group) + work.wait() + return pickle.loads(bytes(payload_tensor.cpu().numpy().tobytes())) + + def _pp_recv_pyobj_from_prev_stage(self: Scheduler): + if self.ps.attn_tp_rank == 0 and self.ps.attn_cp_rank == 0: + dp_offset = ( + self.ps.attn_dp_rank * self.ps.attn_cp_size * self.ps.attn_tp_size + ) + if _PP_DEGLOO: + data = self._pp_recv_pyobj_nccl( + ((self.ps.pp_rank - 1) % self.ps.pp_size) * self.ps.tp_size + + dp_offset, + ) + else: + data = point_to_point_pyobj( + [], + self.ps.pp_rank * self.ps.tp_size + dp_offset, + self.world_group.cpu_group, + ((self.ps.pp_rank - 1) % self.ps.pp_size) * self.ps.tp_size + dp_offset, + self.ps.pp_rank * self.ps.tp_size + dp_offset, + ) + else: + data = None + + if self.ps.attn_tp_size > 1: + data = broadcast_pyobj( + data, + self.attn_tp_group.rank, + self.attn_tp_cpu_group, + src=self.attn_tp_group.ranks[0], + ) + + if self.ps.attn_cp_size > 1: + data = broadcast_pyobj( + data, + self.attn_cp_group.rank, + self.attn_cp_cpu_group, + src=self.attn_cp_group.ranks[0], + ) + + return data + + def _pp_prepare_tensor_dict( + self: Scheduler, result: GenerationBatchResult, batch: ScheduleBatch + ) -> Dict[str, torch.Tensor]: + tensor_dict = { + "next_token_ids": result.next_token_ids, + } + + # Draft extend runs only on the last stage, but every rank needs its relayed + # output to fill PD auxiliary buffers. + draft_input = result.next_draft_input + if draft_input is not None and draft_input.topk_p is not None: + tensor_dict["draft_topk_p"] = draft_input.topk_p.contiguous() + tensor_dict["draft_topk_index"] = draft_input.topk_index.contiguous() + tensor_dict["draft_hidden_states"] = draft_input.hidden_states.contiguous() + + # [EXPERIMENT] PP+spec replica: relay the verify/extend hidden + # states and acceptance lens so non-last stages can replay the + # draft-extend and keep their draft replicas in lockstep. Gated on + # the same condition as the draft keys above so the receiving side + # sees rebind + replay inputs together or not at all. + if ( + result.logits_output is not None + and result.logits_output.hidden_states is not None + and result.logits_output.hidden_states.numel() > 0 + ): + tensor_dict["spec_hidden_states"] = ( + result.logits_output.hidden_states.contiguous() + ) + if result.accept_lens is not None: + tensor_dict["spec_accept_lens"] = result.accept_lens.contiguous() + + # [EXPERIMENT] PP+spec: the verify-assigned draft-KV slots + # (stashed on the result inside the forward window, before + # _forward_isolation reverts the batch edits). The non-last + # stages' decode replay must commit accepted rows into these + # exact slots. + verify_slots = getattr(result, "pp_verify_slots", None) + if verify_slots is not None: + tensor_dict["spec_verify_slots"] = verify_slots.contiguous() + + # [EXPERIMENT] PP+spec: the last stage's shifted draft-extend + # input ids (stashed on next_draft_input by the extend branch) — + # non-last stages' prefill replay consumes them verbatim because + # their batch.input_ids is cleared before the pack runs. + spec_prefill_ids = getattr(draft_input, "spec_prefill_input_ids", None) + if spec_prefill_ids is not None: + tensor_dict["spec_prefill_input_ids"] = spec_prefill_ids.contiguous() + + if batch.return_logprob: + logprob_dict = get_logprob_dict_from_result(result) + tensor_dict = { + **tensor_dict, + **logprob_dict, + } + auxiliary_output = ( + result.logits_output.auxiliary_device_output + if result.logits_output is not None + else None + ) + add_auxiliary_output_to_pp_tensors(tensor_dict, auxiliary_output) + return tensor_dict + + def _pp_send_dict_to_next_stage( + self: Scheduler, + tensor_dict: Dict[str, torch.Tensor], + async_send: bool = True, + msg_type: str = "default", + ): + # Warn once if using default untyped messages + if msg_type == "default": + logger.warning_once( + "PP send: using default untyped message. " + "Consider adding msg_type='proxy' or 'output' to avoid recv conflicts." + ) + tensor_dict["__msg_type__"] = msg_type + # [PP+spec FIX] Lifetime of async-sent tensors: the P2P isend runs + # on the scheduler's ambient stream while several of these tensors + # were produced on the forward stream (last-rank results, proxy + # tensors). ProcessGroupNCCL's async work does not record its + # stream on the tensors, so once the P2PWork list is released + # (the next commit's work.wait() is a host no-op in the default + # non-blocking mode) the caching allocator can reuse or free the + # block while the isend kernel still reads it. Under concurrent + # load this surfaced as an async CUDA illegal memory access on + # PP0 (cc8 run, ~54s in; masked in debug builds by the post-recv + # schedule_stream drain). record_stream makes the allocator delay + # reuse until the send stream is done with the block. + _send_cur_stream = self.device_module.current_stream() + for _t in tensor_dict.values(): + if torch.is_tensor(_t) and _t.is_cuda: + _t.record_stream(self.schedule_stream) + if _send_cur_stream != self.schedule_stream: + _t.record_stream(_send_cur_stream) + p2p_work = [] + if _PP_DEGLOO: + # [r35 de-GLOO] Metadata over a fixed-size NCCL buffer instead + # of send_object's 2 GLOO messages. The GLOO isend rendezvous + # (completed at the next commit's work.wait()) host-blocked the + # sender until the peer posted its irecv — the measured + # 300-540ms stalls that skewed TP allreduce arrival by up to + # 132ms/round. NCCL isend is fire-and-forget on the host. + # Tensors and metadata share the pp device group's per-pair + # FIFO, so the per-direction message order (and therefore the + # proxy/output demux) is preserved exactly. + p2p_work = self._pp_send_tensor_dict_nccl( + tensor_dict, async_send=async_send + ) + else: + p2p_work.extend( + self.pp_group.send_tensor_dict( + tensor_dict=tensor_dict, + # [PP+spec FIX] Never slice+allgather stage-to-stage + # messages. recv_tensor_dict leaves the reassembly + # allgathers in flight on the schedule stream via pynccl, + # sharing one NCCL comm (and op sequence) with the model's + # lm_head allgathers. The proxy/output demux recv order is + # per-TP-pair timing-dependent (a stray proxy gets recv'd + # and stashed inside a result recv, issuing its + # allgathers), so the per-rank pynccl issue sequence can + # diverge across TP ranks; NCCL then pairs mismatched + # collectives which spin forever (wedged 6/6 at + # chunked-prefill chunk-2). Full-tensor pairwise P2P + # issues no collectives and has no cross-rank sequence + # invariant. Boundary tensors are replicated across TP + # ranks (pure TP), so this is semantically identical; + # only per-pair wire bytes grow ~4x. + all_gather_group=None, + async_send=async_send, + ) + ) + return p2p_work + + def _pp_send_tensor_dict_nccl( + self: Scheduler, + tensor_dict: Dict[str, torch.Tensor], + async_send: bool = True, + ) -> List[P2PWork]: + """[r35 de-GLOO] send_tensor_dict with GLOO metadata replaced by a + fixed-size NCCL buffer. + + Wire format per message (per direction, per pair, FIFO on the pp + device group): [meta_buf(8192B uint8), tensor0, tensor1, ...]. + CPU tensors are moved to the device for the NCCL ride and restored + on the receiver; their original keys are listed in the metadata + under __pp_cpu_keys__ so both sides stay byte-compatible with the + GLOO path's semantics. + """ + group = self.pp_group + dst = (group.rank_in_group + 1) % group.world_size + device = self.device_module.current_device() + + cpu_keys: List[str] = [] + moved: Dict[str, torch.Tensor] = {} + for key, value in tensor_dict.items(): + if isinstance(value, torch.Tensor) and value.is_cpu: + cpu_keys.append(key) + moved[key] = value + if cpu_keys: + for key in cpu_keys: + tensor_dict[key] = tensor_dict[key].to( + device, non_blocking=False + ).contiguous() + + metadata_list, tensor_list = _split_tensor_dict(tensor_dict) + if cpu_keys: + metadata_list.append(("__pp_cpu_keys__", cpu_keys)) + + meta_bytes = pickle.dumps(metadata_list) + assert len(meta_bytes) <= _PP_META_BUF_BYTES, ( + f"[r35] boundary metadata pickle too big for the fixed buffer: " + f"{len(meta_bytes)} > {_PP_META_BUF_BYTES}" + ) + meta_buf = torch.empty(_PP_META_BUF_BYTES, dtype=torch.uint8, device=device) + meta_buf[: len(meta_bytes)].copy_( + torch.frombuffer(bytearray(meta_bytes), dtype=torch.uint8), + non_blocking=False, + ) + # [r35] Gate the freshly built buffers onto the ambient stream and + # record both candidate use streams BEFORE any isend is enqueued + # (see _pp_nccl_stage_guard for the empirical isend-stream note). + _guard_tensors = [meta_buf] + if cpu_keys: + # Moved CPU->GPU copies were allocated after the caller's + # record_stream pass, so they need the same two-stream hedge. + _guard_tensors.extend(tensor_dict[key] for key in cpu_keys) + self._pp_nccl_stage_guard(_guard_tensors) + + p2p_works: List[P2PWork] = [] + send_func = torch.distributed.isend if async_send else torch.distributed.send + work = send_func(meta_buf, group.ranks[dst], group=group.device_group) + if async_send: + p2p_works.append(P2PWork(work, meta_buf)) + for tensor in tensor_list: + if tensor.numel() == 0: + # Skip sending empty tensors (vanilla semantics). + continue + # Hold the tensor that is actually on the wire: if contiguous() + # had to copy, the P2PWork must pin the copy, not the original. + send_tensor = ( + tensor if tensor.is_contiguous() else tensor.contiguous() + ) + if send_tensor is not tensor: + # The fresh copy is filled on the current stream while the + # isend runs on the ambient stream — gate it first. + self._pp_nccl_stage_guard([send_tensor]) + work = send_func( + send_tensor, group.ranks[dst], group=group.device_group + ) + if async_send: + p2p_works.append(P2PWork(work, send_tensor)) + if cpu_keys: + # Restore the caller's view: the moved copies must not leak back + # into the caller's dict (it may reuse the tensors after return). + for key in cpu_keys: + tensor_dict[key] = moved[key] + return p2p_works + + def _pp_recv_typed_dict( + self: Scheduler, + expected_kind: str = "default", + all_gather_group: Optional = None, + ) -> Dict[str, torch.Tensor]: + """Receive a typed tensor dict, demultiplexing by msg_type. + + If a message of the wrong kind is received, it's stashed in the queue + and we continue receiving until we get the expected kind. + """ + if expected_kind in self._pp_tensor_dict_inbox: + inbox_queue = self._pp_tensor_dict_inbox[expected_kind] + if inbox_queue: + return inbox_queue.popleft() + + while True: + if _PP_DEGLOO: + tensor_dict = self._pp_recv_tensor_dict_nccl() + else: + tensor_dict = self.pp_group.recv_tensor_dict( + all_gather_group=all_gather_group + ) + # [PP+spec FIX] The recv buffers are allocated on the ambient + # (schedule) stream but read cross-stream: by the copy stream + # (prep / d2h / draft replay) and by the forward stream (the + # stage forward consuming proxy tensors). Those reads are + # unordered against the tensors' eventual free, so the + # allocator could reuse the block while a reader still runs. + # record_stream closes the window (see the matching send-side + # note in _pp_send_dict_to_next_stage). + for _t in tensor_dict.values(): + if torch.is_tensor(_t) and _t.is_cuda: + _t.record_stream(self.forward_stream) + _t.record_stream(self.copy_stream) + received_kind = tensor_dict.get("__msg_type__", "default") + if received_kind == expected_kind: + if received_kind == "default": + logger.warning_once( + f"PP recv: got default untyped message. Content keys: {tensor_dict.keys()}" + "Consider adding msg_type='proxy' or 'output' to avoid recv conflicts." + ) + return tensor_dict + else: + # [PP+spec] visible at INFO: a wrong-kind message was + # consumed mid-recv (channel timing); this is exactly the + # event whose per-rank asymmetry used to poison the pynccl + # sequence before the all_gather_group=None fix. + logger.info( + f"[PPSPEC] DEMUX-STASH pp_rank={self.pp_group.rank_in_group} " + f"expected={expected_kind} got={received_kind} " + f"keys={ [k for k in tensor_dict.keys() if k != '__msg_type__'] }" + ) + self._pp_tensor_dict_inbox[received_kind].append(tensor_dict) + + def _pp_recv_tensor_dict_nccl(self: Scheduler) -> Dict[str, torch.Tensor]: + """[r35 de-GLOO] Receiving side of _pp_send_tensor_dict_nccl. + + One fixed-size metadata irecv replaces recv_object's two blocking + GLOO irecv waits; the only host sync left is the 8KB D2H copy of + the metadata buffer, which is bounded by a single small NCCL + transfer. Data tensors are irecv'd one by one exactly like the + vanilla path. + """ + group = self.pp_group + src = (group.rank_in_group - 1) % group.world_size + device = self.device_module.current_device() + + meta_buf = torch.empty(_PP_META_BUF_BYTES, dtype=torch.uint8, device=device) + work = torch.distributed.irecv(meta_buf, group.ranks[src], group=group.device_group) + work.wait() + metadata_list = pickle.loads(bytes(meta_buf.cpu().numpy().tobytes())) + + tensor_dict: Dict[str, Any] = {} + for key, value in metadata_list: + if isinstance(value, TensorMetadata): + tensor = torch.empty(value.size, dtype=value.dtype, device=value.device) + if tensor.numel() == 0: + tensor_dict[key] = tensor + continue + work = torch.distributed.irecv( + tensor, group.ranks[src], group=group.device_group + ) + work.wait() + tensor_dict[key] = tensor + else: + tensor_dict[key] = value + cpu_keys = tensor_dict.pop("__pp_cpu_keys__", None) + if cpu_keys: + for key in cpu_keys: + tensor_dict[key] = tensor_dict[key].cpu() + return tensor_dict + + 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( + # [PP+spec FIX] all_gather_group=None — see the long note in + # _pp_send_dict_to_next_stage. Collectives in the recv path + # made the pynccl op sequence depend on per-pair demux + # timing, deadlocking NCCL at chunked-prefill chunk-2. + self._pp_recv_typed_dict( + expected_kind="proxy", + all_gather_group=None, + ) + ) + return pp_proxy_tensors + + def _pp_recv_dict_from_prev_stage( + self: Scheduler, + ) -> Dict[str, torch.Tensor]: + return self._pp_recv_typed_dict( + expected_kind="output", + # [PP+spec FIX] all_gather_group=None — see + # _pp_send_dict_to_next_stage. + all_gather_group=None, + ) + + def _pp_make_skip_output_result( + self: Scheduler, + batch: ScheduleBatch, + mb_metadata: Optional[PPBatchMetadata], + ): + bs = len(batch.reqs) + placeholder = torch.zeros(bs, dtype=torch.int64, device=self.device) + # next_pp_outputs = None so non-last ranks skip forwarding + # (pp_outputs is None gate). Placeholder carried in + # batch_result.next_token_ids for process_batch_result_prefill. + batch_result = GenerationBatchResult( + logits_output=None, + pp_hidden_states_proxy_tensors=None, + next_token_ids=placeholder, + can_run_cuda_graph=( + mb_metadata.can_run_cuda_graph if mb_metadata else False + ), + skipped_output_comm=True, + ) + d2h_event = self.device_module.Event() + d2h_event.record(self.device_module.current_stream()) + return None, batch_result, d2h_event + + def _pp_prep_batch_result( + self: Scheduler, + batch: ScheduleBatch, + mb_metadata: PPBatchMetadata, + pp_outputs: PPProxyTensors, + ): + from sglang.srt.managers.scheduler import GenerationBatchResult + + logits_output = None + extend_input_len_per_req = None + extend_logprob_start_len_per_req = None + + if batch.return_logprob: + ( + logits_output, + extend_input_len_per_req, + extend_logprob_start_len_per_req, + ) = get_logprob_from_pp_outputs(pp_outputs) + if self.pp_group.is_first_rank: + observer = self.tp_worker.model_runner.sampling_observer + auxiliary_output = pop_auxiliary_output_from_pp_tensors( + pp_outputs.tensors, + observer, + ) + if auxiliary_output is not None: + if logits_output is None: + logits_output = LogitsProcessorOutput(next_token_logits=None) + logits_output.auxiliary_device_output = auxiliary_output + next_token_ids = pp_outputs["next_token_ids"].to(torch.int64) + + # Rebind the last stage's ring proposal as batch.spec_info so the PD result + # processor sees the same object on every rank. + next_draft_input = None + if "draft_topk_p" in pp_outputs.tensors: + from sglang.srt.speculative.eagle_info import EagleDraftInput + + # [PPSPEC] decode rounds: the ring's next_token_ids is the padded + # predict matrix [bs*stride]; the last stage's + # EagleDraftInput.bonus_tokens is the per-req last-committed + # (bonus) token [bs] — the next tree's root (build_tree_kernel). + # Prefill rounds already carry per-req [bs] tokens. + if "spec_accept_lens" in pp_outputs.tensors: + accept = pp_outputs.tensors["spec_accept_lens"].to(torch.int64) + bs = len(batch.seq_lens) + flat = next_token_ids.view(-1) + if flat.numel() != bs: + stride = flat.numel() // bs + idx = ( + torch.arange(bs, device=flat.device) * stride + accept - 1 + ) + per_req_bonus = flat[idx] + else: + per_req_bonus = next_token_ids + else: + per_req_bonus = next_token_ids + next_draft_input = EagleDraftInput( + topk_p=pp_outputs["draft_topk_p"], + topk_index=pp_outputs["draft_topk_index"], + hidden_states=pp_outputs["draft_hidden_states"], + bonus_tokens=per_req_bonus, + num_tokens_per_req=1, + num_tokens_for_logprob_per_req=1, + ) + batch.spec_info = next_draft_input + + # [EXPERIMENT] PP+spec replica: replay the last stage's in-forward + # draft-extend on this rank NOW — same pre-filter batch state the + # last stage saw (input_ids intact, nothing freed yet), strictly + # before process_batch_result filters/frees and before this + # batch's next draft(). Keeps this stage's draft KV + DSA seeds + # in lockstep with the last stage's replica. + if ( + not self.pp_group.is_last_rank + and "spec_hidden_states" in pp_outputs.tensors + ): + self._pp_replay_spec_round( + batch, + next_draft_input, + pp_outputs, + next_token_ids, + ) + + # PP rank 0 also relays into output_tokens_buf so the next iter's + # resolve_forward_inputs finds these tokens for the decode portion + # of mixed-chunk batches (which gather via mix_running_indices). + # [EXPERIMENT] PP+spec: for spec decode rounds the ring's + # next_token_ids is the padded predict matrix [bs*stride]; the + # future_map stash is per-req — stash each req's last accepted token + # (the sequence-continuation token) instead. + if next_draft_input is not None: + stash_bonus_tokens = next_draft_input.bonus_tokens + else: + stash_bonus_tokens = next_token_ids + self.future_map.stash( + batch.req_pool_indices, + RelayPayload( + bonus_tokens=stash_bonus_tokens, + topk_p=None if next_draft_input is None else next_draft_input.topk_p, + topk_index=( + None if next_draft_input is None else next_draft_input.topk_index + ), + hidden_states=( + None if next_draft_input is None else next_draft_input.hidden_states + ), + ), + ) + batch.input_ids = None + output_result = GenerationBatchResult( + logits_output=logits_output, + pp_hidden_states_proxy_tensors=None, + next_token_ids=pp_outputs["next_token_ids"], + next_draft_input=next_draft_input, + extend_input_len_per_req=extend_input_len_per_req, + extend_logprob_start_len_per_req=extend_logprob_start_len_per_req, + can_run_cuda_graph=mb_metadata.can_run_cuda_graph, + ) + output_result.copy_auxiliary_output_to_cpu() + # [EXPERIMENT] PP+spec: the spec decode processor + # (_resolve_spec_v2_tokens) asserts CPU next_token_ids + accept_lens + # and needs the draft stride — all local-result fields on the last + # stage, ring-only here. Blocking copies: the caller records the + # d2h_event right after, gating the processor's read. + if "spec_accept_lens" in pp_outputs.tensors: + accept_lens_t = pp_outputs.tensors["spec_accept_lens"] + new_seq_lens = batch.seq_lens + accept_lens_t + # [EXPERIMENT] PP+spec: run_batch's spec branch advances + # batch.seq_lens from result.new_seq_lens, but that result is the + # LOCAL worker's — only the last rank's carries new_seq_lens; + # non-last stages' local results are proxies-only, so without + # this their batch.seq_lens never advances, the draft replica's + # positions freeze at the prefill length, and its tree diverges + # from the last stage's within a few rounds (mismatched verify + # context poisons every round after). Advance the relayed batch + # here — the same pre-filter object the ring rebind targets and + # the next round's draft() reads. The replay above already ran + # with the pre-advance value, mirroring the last stage's + # in-forward draft-extend (which also sees pre-advance seq_lens, + # since run_batch advances only after forward_batch_generation). + if not self.pp_group.is_last_rank: + batch.seq_lens = new_seq_lens + if batch.seq_lens_cpu is not None: + batch.seq_lens_cpu = new_seq_lens.to("cpu") + batch.seq_lens_sum = int(batch.seq_lens_cpu.sum()) + output_result.accept_lens = accept_lens_t + output_result.new_seq_lens = new_seq_lens + stride = getattr( + self.model_worker, "speculative_num_draft_tokens", None + ) + if stride is None: + stride = ( + output_result.next_token_ids.numel() + // len(batch.seq_lens) + ) + output_result.speculative_num_draft_tokens = stride + output_result.next_token_ids = output_result.next_token_ids.cpu() + output_result.accept_lens = output_result.accept_lens.cpu() + return output_result + + def _pp_replay_spec_round( + self: Scheduler, + batch: ScheduleBatch, + next_draft_input, + pp_outputs: PPProxyTensors, + next_token_ids: torch.Tensor, + ): + """[EXPERIMENT] PP+spec: advance this (non-last) stage's draft replica + from the ring-relayed round state, mirroring the last stage's + in-forward draft-extend. Called from _pp_prep_batch_result on the same + pre-filter batch object the ring rebind targets. + + Classification mirrors the last stage's in-forward branch choice: + accept_lens present -> decode verify round (draft-extend for decode); + otherwise -> extend/mixed round (draft-extend for prefill). + """ + worker = self.model_worker + replay_prefill = getattr(worker, "replay_draft_extend_for_prefill", None) + if replay_prefill is None: + logger.warning_once( + "[PPSPEC] model worker lacks draft-replay methods; non-last " + "stage draft replica will desynchronize" + ) + return + spec_hidden = pp_outputs.tensors["spec_hidden_states"] + # [EXPERIMENT] PP+spec replica: run the replay on the FORWARD stream, + # not the ambient copy stream. The replay enqueues draft-model TP + # collectives; keeping every stage-local TP collective on one stream + # removes any cross-stream NCCL issue-order ambiguity, and it also + # takes the replay out of the copy stream's wait chain entirely (the + # copy-stream path wedged reproducibly at chunked-prefill chunk-2 — + # all four stage0 ranks stuck at the first host sync inside the + # replay — across three different stream-ordering patch states, so + # the trigger is the copy-stream placement itself, not a specific + # wait edge). Order the forward stream behind the streams that + # produced the replay's inputs: the schedule stream (relay tensors + # from _do_recv) and the copy stream (any prep-side tensor ops + # enqueued earlier in this same _pp_prep_batch_result call). + if _PP_SPEC_DEBUG: + logger.info( + f"[PPSPEC] replay-pre pp_rank={self.pp_group.rank_in_group} " + f"round={'decode' if 'spec_accept_lens' in pp_outputs.tensors else 'prefill'} " + f"sched_done={self.schedule_stream.query()} " + f"fwd_done={self.forward_stream.query()} " + f"copy_done={self.device_module.current_stream().query()}" + ) + with self.forward_stream_ctx: + self.forward_stream.wait_stream(self.schedule_stream) + self.forward_stream.wait_stream(self.device_module.current_stream()) + if "spec_accept_lens" in pp_outputs.tensors: + verify_slots = pp_outputs.tensors.get("spec_verify_slots") + if verify_slots is None: + logger.warning_once( + "[PPSPEC] decode replay round missing spec_verify_slots; " + "skipping replica maintenance (replica will desynchronize)" + ) + return + worker.replay_draft_extend_for_decode( + batch, + verify_hidden_states=spec_hidden, + accept_lens=pp_outputs.tensors["spec_accept_lens"], + next_token_ids=next_token_ids, + rebound_spec_info=next_draft_input, + verify_slots=verify_slots, + ) + else: + prefill_ids = pp_outputs.tensors.get("spec_prefill_input_ids") + if prefill_ids is None: + logger.warning_once( + "[PPSPEC] prefill replay round missing spec_prefill_input_ids; " + "skipping replica maintenance (replica will desynchronize)" + ) + return + worker.replay_draft_extend_for_prefill( + batch, + target_hidden_states=spec_hidden, + next_token_ids=next_token_ids, + rebound_spec_info=next_draft_input, + prebuilt_input_ids=prefill_ids, + ) + if _PP_SPEC_DEBUG: + # Non-blocking snapshot printed after the replay's python path + # (its internal .item() probes have already synced the forward + # stream when debug is on). + logger.info( + f"[PPSPEC] replay-post pp_rank={self.pp_group.rank_in_group} " + f"fwd_done={self.forward_stream.query()}" + ) + + def _pp_process_batch_result( + self: Scheduler, batch: ScheduleBatch, output_result: GenerationBatchResult + ): + # [EXPERIMENT] PP+spec: run_batch's copy_to_cpu is gated on + # pp_size == 1, so the last rank's local spec decode result still + # holds GPU tensors here while the processor asserts CPU. The ring + # pack already captured device tensors at launch, so CPU-ifying the + # local result now is safe. Stage0's prepped result is already CPU. + if ( + output_result is not None + and output_result.accept_lens is not None + and output_result.next_token_ids is not None + and not output_result.next_token_ids.is_cpu + ): + output_result.next_token_ids = output_result.next_token_ids.cpu() + output_result.accept_lens = output_result.accept_lens.cpu() + self.process_batch_result(batch, output_result) + + def _pp_send_output_to_next_stage( + self: Scheduler, + next_first_rank_mb_id: int, + mbs: List[ScheduleBatch], + last_rank_comm_queue: deque, + pp_outputs: PPProxyTensors | None, + ) -> List[P2PWork]: + send_output_work = [] + # [EXPERIMENT] PP+spec: output-dict sends must NOT share the schedule + # stream. The prep's copy stream waits the schedule stream with a + # point-in-time snapshot (see _do_recv), so an output isend still + # sitting there drags the prep — and the host's d2h sync, and the + # proxy send at the bottom of the event loop — behind the PEER's + # matching recv; meanwhile the peer's own d2h waits on this rank's + # proxy send. That cycle wedges the scheduler pair (observed: + # chunked-prefill chunk-2 replay, both stages py-spied inside it). + # Route every output send through a dedicated stream so no + # copy-stream wait or d2h event ever covers a cross-stage send. + # Send inputs still order correctly: the non-last branch forwards + # tensors received on the schedule stream (wait_stream below), and + # the last-rank branch sends tensors produced on the forward stream + # (wait_event on the round's launch event below). + send_stream = getattr(self, "_pp_output_send_stream", None) + if send_stream is None: + send_stream = self.device_module.Stream() + self._pp_output_send_stream = send_stream + # [EXPERIMENT] PP+spec: order the send's INPUTS via the recv-done + # event, never a full wait_stream(schedule) — a point-in-time full + # stream wait would also cover this iteration's pending cross-stage + # recv, and the peer's send is ordered the same way: the two recvs + # each wait the other stage's send, a self-deadlock (observed as the + # chunked-prefill chunk-2 wedge). The last-rank branch below adds + # wait_event(q_event) for its forward-stream-produced tensors. + _recv_ev = getattr(self, "_pp_recv_done_event", None) + if _recv_ev is not None: + send_stream.wait_event(_recv_ev) + if _PP_SPEC_DEBUG: + _n_keys = ( + 0 + if pp_outputs is None + else len(pp_outputs.tensors) + ) + logger.info( + f"[PPSPEC] OUTSEND pp_rank={self.pp_group.rank_in_group} " + f"last={self.pp_group.is_last_rank} fwd_keys={_n_keys} " + f"target_slot={next_first_rank_mb_id} " + f"target=" + f"{None if mbs[next_first_rank_mb_id] is None else mbs[next_first_rank_mb_id].forward_mode.name}" + ) + with torch.cuda.stream(send_stream): + if self.pp_group.is_last_rank: + # send ready PP output to rank 0 + target = mbs[next_first_rank_mb_id] + if target is not None: + q_event, pp_outputs_to_send = last_rank_comm_queue.popleft() + if ( + not target.forward_mode.is_prebuilt() + and not _pp_can_skip_output_comm(target) + ): + send_stream.wait_event(q_event) + # [PP+spec FIX] The isend kernels empirically run + # on the scheduler's ambient stream, not on + # send_stream (see the send_first parity note), + # so the wait above does not actually gate them. + # Gate the data on the stream that executes the + # op, or the send races the forward that produces + # it once the debug probes' host syncs are off. + self.schedule_stream.wait_event(q_event) + with torch.profiler.record_function( + "send_res_dict_to_next_stage" + ): + send_output_work = self._pp_send_dict_to_next_stage( + pp_outputs_to_send.tensors, + async_send=True, + msg_type="output", + ) + # send the outputs from the last round to let the next stage worker run post processing + if not self.pp_group.is_last_rank: + if pp_outputs: + with torch.profiler.record_function("send_res_dict_to_next_stage"): + send_output_work = self._pp_send_dict_to_next_stage( + pp_outputs.tensors, + async_send=True, + msg_type="output", + ) + return send_output_work + + def _pp_send_recv_and_preprocess_output_tensors( + self: Scheduler, + next_first_rank_mb_id: int, + next_mb_id: int, + mbs: List[ScheduleBatch], + mb_metadata: List[PPBatchMetadata], + last_rank_comm_queue: deque[Tuple[torch.Event, PPProxyTensors]], + pp_outputs: PPProxyTensors | None, + ) -> Tuple[ + Optional[PPProxyTensors], + Optional[GenerationBatchResult], + Optional[torch.Event], + List[P2PWork], + ]: + next_pp_outputs = None + d2h_event = None + batch_result = None + send_output_work = [] + + # On some backends isend is effectively blocking and does not + # return until the peer posts a matching recv; if every PP rank + # sends first, all ranks block waiting for a receiver and the + # ring deadlocks. Order send/recv by pp_rank parity (even: + # send->recv, odd: recv->send) so each adjacent pair has one + # sender and one receiver posted at the same time. + + # [PP+spec FIX] Apply the parity ordering on CUDA too. The + # boundary P2P ops go through torch.distributed.isend/irecv + # (send_tensor_dict/recv_tensor_dict), and NCCL COLL logs of the + # wedged run show every Send/Recv of a rank executing on ONE + # stream (the scheduler's ambient stream) regardless of the + # `with torch.cuda.stream(send_stream)` context — so a rank's + # recv is stream-ordered BEHIND its own pending send. With + # send-first on both stages, each side's big rendezvous send + # (spec ring-back carries chunk-sized hidden states, ~100MB for + # an 8192-token chunk) spins waiting for the peer's recv, which + # is queued behind the peer's own spinning send — a FIFO cycle. + # Observed at chunked-prefill chunk-2: all 8 ranks inside the + # post-recv sync, both NCCL send kernels at 100% GPU forever + # (7/7 repros). Small ring-backs (non-spec: token ids only) take + # NCCL's buffered path and complete without the peer, which is + # why plain PP never hits this. Parity makes each stream's head + # op pair with the peer's head op, unwinding any such round. + send_first = (self.ps.pp_rank % 2) == 0 + + def _do_send(): + return self._pp_send_output_to_next_stage( + next_first_rank_mb_id, + mbs, + last_rank_comm_queue, + pp_outputs, + ) + + def _do_recv(): + nonlocal next_pp_outputs, batch_result, d2h_event + target = mbs[next_mb_id] + if target is None or target.forward_mode.is_prebuilt(): + if _PP_SPEC_DEBUG: + logger.info( + f"[PPSPEC] RECV-SKIP pp_rank={self.pp_group.rank_in_group} " + f"mb={next_mb_id}" + ) + return + if _pp_can_skip_output_comm(target): + next_pp_outputs, batch_result, d2h_event = ( + self._pp_make_skip_output_result(target, mb_metadata[next_mb_id]) + ) + return + if _PP_SPEC_DEBUG: + logger.info( + f"[PPSPEC] RECV-IN pp_rank={self.pp_group.rank_in_group} " + f"mb={next_mb_id} mode={target.forward_mode.name}" + ) + with torch.profiler.record_function("recv_res_dict_from_prev_stage"): + next_pp_outputs = PPProxyTensors(self._pp_recv_dict_from_prev_stage()) + # [EXPERIMENT] PP+spec: record the recv's completion on the + # schedule stream. The non-last rank forwards these tensors + # next iteration from the dedicated send stream; that send + # must order behind THIS event only — never behind a full + # wait_stream(schedule): the point-in-time snapshot can + # include this iteration's pending cross-stage recv, which + # waits the peer's send, which (symmetrically) waits its own + # pending recv — the mutual dependency wedged the ring at + # chunked-prefill chunk-2 across every patch variant tried. + _recv_ev = self.device_module.Event() + _recv_ev.record(self.schedule_stream) + self._pp_recv_done_event = _recv_ev + # [EXPERIMENT] PP+spec: recv_tensor_dict leaves the attn-TP + # allgathers that reassemble each sliced tensor in flight on + # the schedule stream (they are enqueued but never waited). + # R24 trace: at chunked-prefill chunk-2 these never complete + # on stage0 (sched_done stays False forever) while chunk-1's + # identical set completes moments earlier. Drain them here + # before any copy/fwd point-in-time wait is armed: either the + # sync returns (stream-progress issue fixed, downstream waits + # see an empty stream) or the host blocks (proves the + # allgathers themselves are stuck; NCCL_DEBUG opCounts then + # pinpoint the divergence). + # [EXPERIMENT] PP+spec: unconditional drain of the schedule + # stream after the recv's irecvs. CUDA_LAUNCH_BLOCKING and + # debug builds both pass 8x16384 while the no-drain build + # hits an async CUDA illegal access — the drain serializes + # the in-flight P2P sends/recvs against everything the host + # does next, closing a race window that is still being + # bisected. Keep it unconditional (logs stay debug-gated) + # until the racing pair is identified. + if _PP_SPEC_DEBUG: + logger.info( + f"[PPSPEC] RECV-SYNC-ENTER pp_rank={self.pp_group.rank_in_group} " + f"mb={next_mb_id}" + ) + self.schedule_stream.synchronize() + if _PP_SPEC_DEBUG: + logger.info( + f"[PPSPEC] RECV-SYNC-DONE pp_rank={self.pp_group.rank_in_group} " + f"mb={next_mb_id}" + ) + if _PP_SPEC_DEBUG: + logger.info( + f"[PPSPEC] RECV-OUT pp_rank={self.pp_group.rank_in_group} " + f"mb={next_mb_id} " + f"sched_done={self.schedule_stream.query()} " + f"fwd_done={self.forward_stream.query()} " + f"copy_done={self.copy_stream.query()} " + f"shapes={ {k: tuple(v.shape) for k, v in next_pp_outputs.tensors.items() if hasattr(v, 'shape')} }" + ) + with self.copy_stream_ctx: + self.copy_stream.wait_stream(self.schedule_stream) + batch_result = self._pp_prep_batch_result( + target, mb_metadata[next_mb_id], next_pp_outputs + ) + d2h_event = self.device_module.Event() + d2h_event.record(self.device_module.current_stream()) + + if send_first: + send_output_work = _do_send() + _do_recv() + else: + _do_recv() + send_output_work = _do_send() + + return next_pp_outputs, batch_result, d2h_event, send_output_work + + def _pp_launch_batch( + self: Scheduler, + mb_id: int, + cur_batch: ScheduleBatch, + pp_proxy_tensors: PPProxyTensors, + mb_metadata: List[Optional[PPBatchMetadata]], + last_rank_comm_queue: deque, + ): + with torch.profiler.record_function("run_batch"): + with self.forward_stream_ctx: + self.forward_stream.wait_stream(self.schedule_stream) + # [EXPERIMENT] PP+spec replica: the previous iteration's prep + # wrote batch.seq_lens (accept-lens advance) and other relayed + # state on the copy stream; this round's draft()/verify() on + # the forward stream reads them. Point-in-time wait, same + # semantics as the schedule-stream wait above. + self.forward_stream.wait_stream(self.copy_stream) + set_time_batch( + cur_batch.reqs, + "set_run_batch_cpu_start_time", + trace_only=True, + ) + result = self.run_batch(cur_batch, pp_proxy_tensors) + set_time_batch( + cur_batch.reqs, + "set_run_batch_cpu_end_time", + trace_only=True, + attrs={"pp_mb_id": mb_id}, + ) + mb_metadata[mb_id] = PPBatchMetadata( + can_run_cuda_graph=result.can_run_cuda_graph, + ) + event = self.device_module.Event() + event.record(self.device_module.current_stream()) + if self.pp_group.is_last_rank: + # (last rank) buffer the outputs for async batch depth + last_rank_comm_queue.append( + ( + event, + PPProxyTensors( + self._pp_prepare_tensor_dict(result, cur_batch) + ), + ) + ) + return result, event + + def get_rids( + self: Scheduler, req_queue: List[Req], is_send: bool, *poll_statuses_group + ): + """ + Used by PP, get the required rids with the given poll statuses. + """ + polls = poll_and_all_reduce_attn_cp_tp_group( + [req.disagg_kv_sender if is_send else req.kv_receiver for req in req_queue], + self.attn_cp_cpu_group, + self.attn_tp_cpu_group, + ) + rids: List = [] + for poll_statuses in poll_statuses_group: + rids.append( + [ + req.rid if is_send else req.req.rid + for req, poll in zip(req_queue, polls) + if poll in poll_statuses + ] + ) + return tuple(rids) if len(rids) > 1 else rids[0] + + def _pp_pd_get_retract_ids(self: Scheduler, mb_id: int): + # communicate pre-consensus retracted reqs + for req in self.disagg_decode_prealloc_queue.retracted_queue: + # assign retracted reqs to the current microbatch + if req.retraction_mb_id is None: + req.retraction_mb_id = mb_id + curr_retract_rids = [ + req.rid + for req in self.disagg_decode_prealloc_queue.retracted_queue + if req.retraction_mb_id == mb_id + ] + if self.pp_group.is_first_rank: + # First rank, get all retracted req ids for the microbatch + return curr_retract_rids + else: + # Other ranks, receive the retracted reqs info from the previous rank and ensure the consensus + prev_retract_rids = self._pp_recv_pyobj_from_prev_stage() + return list(set(prev_retract_rids) & set(curr_retract_rids)) + + def _pp_pd_get_prealloc_ids(self: Scheduler): + # communicate pre-consensus prealloc reqs + if self.pp_group.is_first_rank: + # First rank, pop the preallocated reqs from the prealloc queue + good_prealloc_rids, bad_prealloc_rids = self.get_rids( + self.disagg_decode_prealloc_queue.queue, + False, + [KVPoll.WaitingForInput], + [KVPoll.Failed], + ) + else: + # Other ranks, receive the preallocated reqs info from the previous rank and ensure the consensus + prev_prealloc_rids = self._pp_recv_pyobj_from_prev_stage() + prev_good_prealloc_rids, prev_bad_prealloc_rids = prev_prealloc_rids + curr_good_prealloc_rids, curr_bad_prealloc_rids = self.get_rids( + self.disagg_decode_prealloc_queue.queue, + False, + [KVPoll.WaitingForInput], + [KVPoll.Failed], + ) + good_prealloc_rids = list( + set(prev_good_prealloc_rids) & set(curr_good_prealloc_rids) + ) + bad_prealloc_rids = list( + set(prev_bad_prealloc_rids) | set(curr_bad_prealloc_rids) + ) + # Same abort routing as the prefill bootstrap consensus above. + aborted_rids = { + decode_req.req.rid + for decode_req in self.disagg_decode_prealloc_queue.queue + if isinstance(decode_req.req.finished_reason, FINISH_ABORT) + } + good_prealloc_rids, bad_prealloc_rids = self._route_aborts_to_bad( + good_prealloc_rids, bad_prealloc_rids, aborted_rids + ) + return [good_prealloc_rids, bad_prealloc_rids] + + @staticmethod + def _route_aborts_to_bad(good_rids, bad_rids, aborted_rids): + """Move aborted rids out of the good (intersection) set and into the + bad (union) set, so PP consensus fails them uniformly on every rank. + + This also flushes aborted reqs that never reached good/bad consensus + (e.g. stuck in Bootstrapping with a sender that has no working abort()). + """ + if not aborted_rids: + return good_rids, bad_rids + good_rids = [rid for rid in good_rids if rid not in aborted_rids] + bad_rids = list(set(bad_rids) | set(aborted_rids)) + return good_rids, bad_rids + + def _pp_pd_get_decode_transferred_ids(self: Scheduler): + # get the current stage transfer success + if self.pp_group.is_first_rank: + transferred_rids = self.get_rids( + self.disagg_decode_transfer_queue.queue, + False, + [KVPoll.Success, KVPoll.Failed], + ) + # if other ranks, do intersection with the previous rank's transferred rids + else: + # 2 (Release): Receive the transferred rids from the previous rank + # 1. recv previous stage's transferred reqs info + prev_transferred_rids = self._pp_recv_pyobj_from_prev_stage() + # 2. get the current stage's transferred reqs info + curr_transferred_rids = self.get_rids( + self.disagg_decode_transfer_queue.queue, + False, + [KVPoll.Success, KVPoll.Failed], + ) + # 3. new consensus rids = intersection(previous consensus rids, transfer finished rids) + transferred_rids = list( + set(prev_transferred_rids) & set(curr_transferred_rids) + ) + return transferred_rids + + def process_retract_queue(self: Scheduler, retract_rids: Optional[List[str]]): + if retract_rids is not None: + # try to resume retracted requests if there are enough space for another `num_reserved_decode_tokens` decode steps + resumed_reqs = self.disagg_decode_prealloc_queue.resume_retracted_reqs( + retract_rids + ) + self.waiting_queue.extend(resumed_reqs) + return [req.rid for req in resumed_reqs] + return None + + def process_prealloc_queue(self: Scheduler, prealloc_rids: Optional[List[str]]): + if len(self.disagg_decode_prealloc_queue.retracted_queue) > 0: + # if there are still retracted requests, we do not allocate new requests + return [[], []] + + if prealloc_rids is not None: + ( + good_consensus_prealloc_rids, + bad_consensus_prealloc_rids, + ) = prealloc_rids + good_reqs, failed_reqs = self.disagg_decode_prealloc_queue.pop_preallocated( + pp_good_rids=good_consensus_prealloc_rids, + pp_bad_rids=bad_consensus_prealloc_rids, + ) + self.disagg_decode_transfer_queue.extend(good_reqs) + return [ + [req.req.rid for req in good_reqs], + [req.req.rid for req in failed_reqs], + ] + return None + + def process_decode_transfer_queue( + self: Scheduler, release_rids: Optional[List[str]] + ): + # Resolve held deferred releases every call, independent of release_rids, + # so ack/timeout-driven releases still fire when no rids are being polled. + self.disagg_decode_transfer_queue.resolve_deferred_releases() + if release_rids is not None: + released_reqs = self.disagg_decode_transfer_queue.pop_transferred( + release_rids + ) + if self.enable_hisparse: + for req in released_reqs: + self.hisparse_coordinator.admit_request_direct(req) + self.waiting_queue.extend(released_reqs) + return [req.rid for req in released_reqs] + return None + + +class ChunkSizePredictor: + """ + Predictor for dynamic chunk size based on quadratic latency model. + + Models latency as: f(l) = a*l^2 + b*l + c + Predicts next chunk size x such that: f(L+x) - f(L) = target_latency + """ + + def __init__(self): + self.quadratic_coeff_a = 0.0 + self.linear_coeff_b = 0.0 + self.constant_coeff_c = 0.0 + self.target_latency: Optional[float] = None + self.is_ready = False + + def fit(self, seq_lens: List[int], latencies: List[float]): + """Fit quadratic coefficients f(l) = al^2 + bl + c from data points.""" + # Skip the first data point to reduce fitting bias, as the first run is slower without warmup + L = np.array(seq_lens[1:], dtype=np.float64) + T = np.array(latencies[1:], dtype=np.float64) + + if len(L) < 8: + raise ValueError( + f"Not enough data points for quadratic fitting ({len(L)} < 8). " + "Need at least 8 samples with different sequence lengths." + ) + + # Build design matrix for f(l) = al^2 + bl + c + X = np.column_stack([L * L, L, np.ones_like(L)]) # [l^2, l, 1] + + try: + coeffs, residuals, rank, s = np.linalg.lstsq(X, T, rcond=None) + if len(coeffs) >= 3: + fitted_a = float(coeffs[0]) # quadratic coefficient + fitted_b = float(coeffs[1]) # linear coefficient + fitted_c = float(coeffs[2]) # constant coefficient + else: + raise ValueError("Failed to fit coefficients: insufficient rank") + except np.linalg.LinAlgError as e: + raise ValueError(f"Failed to fit f(l) = al^2 + bl + c: {e}") + + # Validate coefficients + if fitted_a <= 0: + raise ValueError( + f"Fitted quadratic coefficient a={fitted_a:.2e} is not positive. " + "Attention has O(n^2) complexity, so a must be positive. " + "Check warmup data quality." + ) + + if fitted_b < 0: + logger.warning( + f"Fitted linear coefficient b={fitted_b:.2e} is negative. Setting b=0." + ) + fitted_b = 0.0 + + self.quadratic_coeff_a = fitted_a + self.linear_coeff_b = fitted_b + self.constant_coeff_c = fitted_c + + logger.info( + f"[ChunkSizePredictor] Fitted coefficients: a={fitted_a:.2e}, " + f"b={fitted_b:.2e}, c={fitted_c:.2e}" + ) + + def set_target_latency(self, base_chunk_size: int): + """Set target latency based on base chunk size: target = f(base_chunk_size) - f(0).""" + + def f(length: float) -> float: + """Total latency function: f(length) = a*length^2 + b*length + c.""" + return ( + self.quadratic_coeff_a * length * length + + self.linear_coeff_b * length + + self.constant_coeff_c + ) + + self.target_latency = f(float(base_chunk_size)) - f(0.0) + + if self.target_latency <= 0: + raise ValueError( + f"Calculated target_latency={self.target_latency:.2f}ms is not positive. " + "Check warmup data quality." + ) + + logger.info( + f"[ChunkSizePredictor] Target latency: {self.target_latency:.2f}ms " + f"(base_chunk_size={base_chunk_size})" + ) + + def predict_next_chunk_size( + self, + history_len: int, + base_chunk_size: int, + page_size: int, + context_len: int, + max_chunk_size: Optional[int] = None, + ) -> Optional[int]: + """ + Predict next chunk size x such that f(history_len + x) - f(history_len) = target_latency. + + Args: + history_len: Current sequence length (L) + base_chunk_size: Base chunk size + page_size: Page size for alignment + context_len: Maximum context length + max_chunk_size: Maximum allowed chunk size (optional) + + Returns: + Predicted chunk size, or None if prediction fails + """ + if not self.is_ready or self.target_latency is None: + return None + + # Handle quadratic model: f(l) = al^2 + bl + c + if self.quadratic_coeff_a <= 0: + return None + + # Solve f(L+x) - f(L) = T + # where f(L) = a*L^2 + b*L + c + # This expands to: ax^2 + (2aL+b)x - T = 0 + # A = a, B = 2aL + b, C = -T + A = self.quadratic_coeff_a + B = 2 * self.quadratic_coeff_a * history_len + self.linear_coeff_b + C = -self.target_latency + + discriminant = B * B - 4 * A * C + + if discriminant < 0: + logger.warning( + f"Discriminant is negative ({discriminant:.2e}). " + f"No real solution for chunk size. L={history_len}, T={self.target_latency:.2f}ms." + ) + return None + + sqrt_discriminant = math.sqrt(discriminant) + calculated_chunk_size_float = (-B + sqrt_discriminant) / (2 * A) + + if calculated_chunk_size_float <= 0: + logger.warning( + f"Calculated chunk size is non-positive ({calculated_chunk_size_float:.2f}). " + f"L={history_len}, T={self.target_latency:.2f}ms." + ) + return None + + # Use a smooth coefficient to reduce the abrupt decrease in chunk size + smooth_coeff = envs.SGLANG_DYNAMIC_CHUNKING_SMOOTH_FACTOR.get() + smoothed_chunk_size = base_chunk_size + smooth_coeff * ( + calculated_chunk_size_float - base_chunk_size + ) + # Make sure the dynamic chunk size is at least 1/4 of the base chunk size + calculated_chunk_size = max(int(smoothed_chunk_size), base_chunk_size // 4) + + # Align to page_size (minimum alignment size is 64) + alignment_size = max(page_size, 64) + dynamic_chunk_size = (calculated_chunk_size // alignment_size) * alignment_size + + # Ensure aligned size is at least alignment_size + if dynamic_chunk_size < alignment_size: + dynamic_chunk_size = alignment_size + + # Apply constraints + max_allowed = context_len - history_len - 100 # Leave 100 tokens margin + if max_chunk_size is not None: + max_allowed = min(max_allowed, max_chunk_size) + dynamic_chunk_size = min(dynamic_chunk_size, max_allowed) + + # Align again after min operation + dynamic_chunk_size = (dynamic_chunk_size // alignment_size) * alignment_size + + if dynamic_chunk_size < alignment_size: + return None + + return dynamic_chunk_size diff --git a/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/deploy_ppmtp_r35.sh b/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/deploy_ppmtp_r35.sh new file mode 100644 index 0000000..2a80484 --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/deploy_ppmtp_r35.sh @@ -0,0 +1,72 @@ +#!/bin/bash +# PP+MTP r35 部署（60.8, nightly-dev-cu13-20260901-07c8f729） +# 与 deploy_ppmtp_mask.sh 差异：scheduler 挂载切到 r35（r33 + 去边界 GLOO， +# SGLANG_PP_DEGLOO 默认开，=0 回退 r33 行为）。eagle worker 仍挂 mask 版保留热开关。 +# 用法: bash deploy_ppmtp_r35.sh "<并行参数>" [mtp|nomtp] [chunk] [memfrac] [degloo] +# 例: bash deploy_ppmtp_r35.sh "--tp 4 --pp-size 2 --disable-overlap-schedule --max-prefill-tokens 16384" mtp 8192 0.88 1 +PAR=${1:?usage: deploy_ppmtp_r35.sh "" [mtp|nomtp] [chunk] [memfrac] [degloo]} +MTPMODE=${2:-nomtp} +CHUNK=${3:-8192} +MEMFRAC=${4:-0.88} +DEGLOO=${5:-1} + +MTPARGS="" +if [ "$MTPMODE" = "mtp" ]; then + MTPARGS="--speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4" +fi + +docker update --restart=no glm53-nvfp4 >/dev/null 2>&1 +for i in 1 2 3 4 5; do + docker rm -f glm53-nvfp4 >/dev/null 2>&1 + sleep 2 + docker ps -a --format '{{.Names}}' 2>/dev/null | grep -q '^glm53-nvfp4$' || break +done +if docker ps -a --format '{{.Names}}' 2>/dev/null | grep -q '^glm53-nvfp4$'; then + echo "ERROR: old container cannot be removed"; exit 1 +fi +for i in $(seq 1 15); do ss -ltn 2>/dev/null | grep -q ":30000 " || break; sleep 2; done + +docker run -d --name glm53-nvfp4 --gpus all --shm-size 64g --ipc=host --cap-add SYS_PTRACE \ + -v /root/sglang_patch2/layer_setup.py:/sgl-workspace/sglang/python/sglang/srt/model_executor/model_runner_components/layer_setup.py:ro \ + -v /root/sglang_patch2/validation_hook.py:/sgl-workspace/sglang/python/sglang/srt/arg_groups/validation_hook.py:ro \ + -v /root/eagle_worker_v2_mask.py:/sgl-workspace/sglang/python/sglang/srt/speculative/eagle_worker_v2.py:ro \ + -v /root/sglang_patch2/eagle_worker_common.py:/sgl-workspace/sglang/python/sglang/srt/speculative/eagle_worker_common.py:ro \ + -v /root/sglang_patch2/deepseek_nextn.py:/sgl-workspace/sglang/python/sglang/srt/models/deepseek_nextn.py:ro \ + -v /root/scheduler_pp_mixin_r35.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler_pp_mixin.py:ro \ + -v /root/sglang_patch2/deepseek_v2.py:/sgl-workspace/sglang/python/sglang/srt/models/deepseek_v2.py:ro \ + -e SGLANG_PP_SPEC_DEBUG=0 \ + -e SGLANG_PP_DEGLOO=${DEGLOO} \ + -e SGLANG_PP_SPEC_FORCE_EAGER_DRAFT=1 \ + -e SGLANG_PP_FORCE_EAGER_VERIFY=1 \ + --restart no -p 30000:30000 \ + -v /data/hf_models:/data/hf_models \ + lmsysorg/sglang:nightly-dev-cu13-20260901-07c8f729 \ + python3 -m sglang.launch_server \ + --model-path /data/hf_models/GLM-5.3-NVFP4 \ + --tp 8 \ + --mem-fraction-static ${MEMFRAC} \ + --max-running-requests 16 \ + --chunked-prefill-size ${CHUNK} \ + --disable-shared-experts-fusion \ + --moe-runner-backend flashinfer_cutlass \ + --disable-flashinfer-autotune \ + --reasoning-parser glm45 --tool-call-parser glm47 \ + --enable-hierarchical-cache --hicache-ratio 3 \ + ${MTPARGS} \ + ${PAR} \ + --host 0.0.0.0 --port 30000 + +echo "deployed r35 (DEGLOO=${DEGLOO}): par=[${PAR}] mtp=${MTPMODE}; waiting for health..." +for i in $(seq 10 10 1800); do + code=$(curl -s -o /dev/null -m3 -w '%{http_code}' http://127.0.0.1:30000/health 2>/dev/null) + if [ "$code" = "200" ]; then + echo "healthy after ${i}s" + docker logs glm53-nvfp4 2>&1 | grep -oE "max_total_num_tokens = [0-9]+" | head -1 + exit 0 + fi + if ! docker ps --format '{{.Names}}' | grep -q '^glm53-nvfp4$'; then + echo "CONTAINER DIED after ${i}s"; docker logs --tail 60 glm53-nvfp4 2>&1 | grep -iE "error|assert|not support|incompatible" | tail -8; exit 1 + fi + sleep 10 +done +echo "TIMEOUT waiting for health"; exit 1 diff --git a/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/deploy_ppmtp_r36.sh b/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/deploy_ppmtp_r36.sh new file mode 100644 index 0000000..818642d --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/deploy_ppmtp_r36.sh @@ -0,0 +1,74 @@ +#!/bin/bash +# PP+MTP r36 部署（60.8, nightly-dev-cu13-20260901-07c8f729） +# r35 首部署死锁修复版：r35 只把请求中继发送端切到 NCCL，接收端在 +# request_receiver.py:_pull_raw_reqs 仍走 GLOO，通道断裂 → 首请求 wedge。 +# r36 = r35 scheduler + request_receiver_degloo.py（接收端同步切两段式 NCCL）。 +# 用法: bash deploy_ppmtp_r36.sh "<并行参数>" [mtp|nomtp] [chunk] [memfrac] [degloo] +# 例: bash deploy_ppmtp_r36.sh "--tp 4 --pp-size 2 --disable-overlap-schedule --max-prefill-tokens 16384" mtp 8192 0.88 1 +PAR=${1:?usage: deploy_ppmtp_r36.sh "" [mtp|nomtp] [chunk] [memfrac] [degloo]} +MTPMODE=${2:-nomtp} +CHUNK=${3:-8192} +MEMFRAC=${4:-0.88} +DEGLOO=${5:-1} + +MTPARGS="" +if [ "$MTPMODE" = "mtp" ]; then + MTPARGS="--speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4" +fi + +docker update --restart=no glm53-nvfp4 >/dev/null 2>&1 +for i in 1 2 3 4 5; do + docker rm -f glm53-nvfp4 >/dev/null 2>&1 + sleep 2 + docker ps -a --format '{{.Names}}' 2>/dev/null | grep -q '^glm53-nvfp4$' || break +done +if docker ps -a --format '{{.Names}}' 2>/dev/null | grep -q '^glm53-nvfp4$'; then + echo "ERROR: old container cannot be removed"; exit 1 +fi +for i in $(seq 1 15); do ss -ltn 2>/dev/null | grep -q ":30000 " || break; sleep 2; done + +docker run -d --name glm53-nvfp4 --gpus all --shm-size 64g --ipc=host --cap-add SYS_PTRACE \ + -v /root/sglang_patch2/layer_setup.py:/sgl-workspace/sglang/python/sglang/srt/model_executor/model_runner_components/layer_setup.py:ro \ + -v /root/sglang_patch2/validation_hook.py:/sgl-workspace/sglang/python/sglang/srt/arg_groups/validation_hook.py:ro \ + -v /root/eagle_worker_v2_mask.py:/sgl-workspace/sglang/python/sglang/srt/speculative/eagle_worker_v2.py:ro \ + -v /root/sglang_patch2/eagle_worker_common.py:/sgl-workspace/sglang/python/sglang/srt/speculative/eagle_worker_common.py:ro \ + -v /root/sglang_patch2/deepseek_nextn.py:/sgl-workspace/sglang/python/sglang/srt/models/deepseek_nextn.py:ro \ + -v /root/scheduler_pp_mixin_r35.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler_pp_mixin.py:ro \ + -v /root/request_receiver_degloo.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler_components/request_receiver.py:ro \ + -v /root/sglang_patch2/deepseek_v2.py:/sgl-workspace/sglang/python/sglang/srt/models/deepseek_v2.py:ro \ + -e SGLANG_PP_SPEC_DEBUG=0 \ + -e SGLANG_PP_DEGLOO=${DEGLOO} \ + -e SGLANG_PP_SPEC_FORCE_EAGER_DRAFT=1 \ + -e SGLANG_PP_FORCE_EAGER_VERIFY=1 \ + --restart no -p 30000:30000 \ + -v /data/hf_models:/data/hf_models \ + lmsysorg/sglang:nightly-dev-cu13-20260901-07c8f729 \ + python3 -m sglang.launch_server \ + --model-path /data/hf_models/GLM-5.3-NVFP4 \ + --tp 8 \ + --mem-fraction-static ${MEMFRAC} \ + --max-running-requests 16 \ + --chunked-prefill-size ${CHUNK} \ + --disable-shared-experts-fusion \ + --moe-runner-backend flashinfer_cutlass \ + --disable-flashinfer-autotune \ + --reasoning-parser glm45 --tool-call-parser glm47 \ + --enable-hierarchical-cache --hicache-ratio 3 \ + ${MTPARGS} \ + ${PAR} \ + --host 0.0.0.0 --port 30000 + +echo "deployed r36 (DEGLOO=${DEGLOO}): par=[${PAR}] mtp=${MTPMODE}; waiting for health..." +for i in $(seq 10 10 1800); do + code=$(curl -s -o /dev/null -m3 -w '%{http_code}' http://127.0.0.1:30000/health 2>/dev/null) + if [ "$code" = "200" ]; then + echo "healthy after ${i}s" + docker logs glm53-nvfp4 2>&1 | grep -oE "max_total_num_tokens = [0-9]+" | head -1 + exit 0 + fi + if ! docker ps --format '{{.Names}}' | grep -q '^glm53-nvfp4$'; then + echo "CONTAINER DIED after ${i}s"; docker logs --tail 60 glm53-nvfp4 2>&1 | grep -iE "error|assert|not support|incompatible" | tail -8; exit 1 + fi + sleep 10 +done +echo "TIMEOUT waiting for health"; exit 1 diff --git a/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/deploy_ppmtp_r36g.sh b/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/deploy_ppmtp_r36g.sh new file mode 100644 index 0000000..3bd34a0 --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/deploy_ppmtp_r36g.sh @@ -0,0 +1,74 @@ +#!/bin/bash +# PP+MTP r36 部署（60.8, nightly-dev-cu13-20260901-07c8f729） +# r35 首部署死锁修复版：r35 只把请求中继发送端切到 NCCL，接收端在 +# request_receiver.py:_pull_raw_reqs 仍走 GLOO，通道断裂 → 首请求 wedge。 +# r36 = r35 scheduler + request_receiver_degloo.py（接收端同步切两段式 NCCL）。 +# 用法: bash deploy_ppmtp_r36.sh "<并行参数>" [mtp|nomtp] [chunk] [memfrac] [degloo] +# 例: bash deploy_ppmtp_r36.sh "--tp 4 --pp-size 2 --disable-overlap-schedule --max-prefill-tokens 16384" mtp 8192 0.88 1 +PAR=${1:?usage: deploy_ppmtp_r36.sh "" [mtp|nomtp] [chunk] [memfrac] [degloo]} +MTPMODE=${2:-nomtp} +CHUNK=${3:-8192} +MEMFRAC=${4:-0.88} +DEGLOO=${5:-1} + +MTPARGS="" +if [ "$MTPMODE" = "mtp" ]; then + MTPARGS="--speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4" +fi + +docker update --restart=no glm53-nvfp4 >/dev/null 2>&1 +for i in 1 2 3 4 5; do + docker rm -f glm53-nvfp4 >/dev/null 2>&1 + sleep 2 + docker ps -a --format '{{.Names}}' 2>/dev/null | grep -q '^glm53-nvfp4$' || break +done +if docker ps -a --format '{{.Names}}' 2>/dev/null | grep -q '^glm53-nvfp4$'; then + echo "ERROR: old container cannot be removed"; exit 1 +fi +for i in $(seq 1 15); do ss -ltn 2>/dev/null | grep -q ":30000 " || break; sleep 2; done + +docker run -d --name glm53-nvfp4 --gpus all --shm-size 64g --ipc=host --cap-add SYS_PTRACE \ + -v /root/sglang_patch2/layer_setup.py:/sgl-workspace/sglang/python/sglang/srt/model_executor/model_runner_components/layer_setup.py:ro \ + -v /root/sglang_patch2/validation_hook.py:/sgl-workspace/sglang/python/sglang/srt/arg_groups/validation_hook.py:ro \ + -v /root/eagle_worker_v2_mask.py:/sgl-workspace/sglang/python/sglang/srt/speculative/eagle_worker_v2.py:ro \ + -v /root/sglang_patch2/eagle_worker_common.py:/sgl-workspace/sglang/python/sglang/srt/speculative/eagle_worker_common.py:ro \ + -v /root/sglang_patch2/deepseek_nextn.py:/sgl-workspace/sglang/python/sglang/srt/models/deepseek_nextn.py:ro \ + -v /root/scheduler_pp_mixin_r35.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler_pp_mixin.py:ro \ + -v /root/request_receiver_degloo.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler_components/request_receiver.py:ro \ + -v /root/sglang_patch2/deepseek_v2.py:/sgl-workspace/sglang/python/sglang/srt/models/deepseek_v2.py:ro \ + -e SGLANG_PP_SPEC_DEBUG=0 \ + -e SGLANG_PP_DEGLOO=${DEGLOO} \ + -e SGLANG_PP_SPEC_FORCE_EAGER_DRAFT=0 \ + -e SGLANG_PP_FORCE_EAGER_VERIFY=0 \ + --restart no -p 30000:30000 \ + -v /data/hf_models:/data/hf_models \ + lmsysorg/sglang:nightly-dev-cu13-20260901-07c8f729 \ + python3 -m sglang.launch_server \ + --model-path /data/hf_models/GLM-5.3-NVFP4 \ + --tp 8 \ + --mem-fraction-static ${MEMFRAC} \ + --max-running-requests 16 \ + --chunked-prefill-size ${CHUNK} \ + --disable-shared-experts-fusion \ + --moe-runner-backend flashinfer_cutlass \ + --disable-flashinfer-autotune \ + --reasoning-parser glm45 --tool-call-parser glm47 \ + --enable-hierarchical-cache --hicache-ratio 3 \ + ${MTPARGS} \ + ${PAR} \ + --host 0.0.0.0 --port 30000 + +echo "deployed r36g GRAPH (DEGLOO=${DEGLOO}): par=[${PAR}] mtp=${MTPMODE}; waiting for health..." +for i in $(seq 10 10 1800); do + code=$(curl -s -o /dev/null -m3 -w '%{http_code}' http://127.0.0.1:30000/health 2>/dev/null) + if [ "$code" = "200" ]; then + echo "healthy after ${i}s" + docker logs glm53-nvfp4 2>&1 | grep -oE "max_total_num_tokens = [0-9]+" | head -1 + exit 0 + fi + if ! docker ps --format '{{.Names}}' | grep -q '^glm53-nvfp4$'; then + echo "CONTAINER DIED after ${i}s"; docker logs --tail 60 glm53-nvfp4 2>&1 | grep -iE "error|assert|not support|incompatible" | tail -8; exit 1 + fi + sleep 10 +done +echo "TIMEOUT waiting for health"; exit 1 diff --git a/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/deploy_ppmtp_sanitize.sh b/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/deploy_ppmtp_sanitize.sh new file mode 100644 index 0000000..e04193d --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/deploy_ppmtp_sanitize.sh @@ -0,0 +1,75 @@ +#!/bin/bash +# PP+MTP 竞态定位部署（60.8, nightly-dev-cu13-20260901-07c8f729） +# 与 deploy_ppmtp_mask.sh 差异： +# 1) 入口用 compute-sanitizer --tool memcheck 包裹（LD_PRELOAD 随 mp.spawn 传满 8 rank） +# 2) SGLANG_PP_SPEC_SYNC_MASK=0 —— 拆掉 7 同步遮罩暴露竞态 +# 3) 健康等待窗口放宽到 3600s（sanitizer 下启动 ~20min+） +# 用法: bash deploy_ppmtp_sanitize.sh "<并行参数>" [mtp|nomtp] [chunk] [memfrac] +# 例: bash deploy_ppmtp_sanitize.sh "--tp 4 --pp-size 2 --disable-overlap-schedule --max-prefill-tokens 16384" mtp +PAR=${1:?usage: deploy_ppmtp_sanitize.sh "" [mtp|nomtp] [chunk] [memfrac]} +MTPMODE=${2:-nomtp} +CHUNK=${3:-8192} +MEMFRAC=${4:-0.88} + +MTPARGS="" +if [ "$MTPMODE" = "mtp" ]; then + MTPARGS="--speculative-algorithm EAGLE --speculative-num-steps 3 --speculative-eagle-topk 1 --speculative-num-draft-tokens 4" +fi + +# 顽固容器清理（同 bench 版） +docker update --restart=no glm53-nvfp4 >/dev/null 2>&1 +for i in 1 2 3 4 5; do + docker rm -f glm53-nvfp4 >/dev/null 2>&1 + sleep 2 + docker ps -a --format '{{.Names}}' 2>/dev/null | grep -q '^glm53-nvfp4$' || break +done +if docker ps -a --format '{{.Names}}' 2>/dev/null | grep -q '^glm53-nvfp4$'; then + echo "ERROR: old container cannot be removed"; exit 1 +fi +for i in $(seq 1 15); do ss -ltn 2>/dev/null | grep -q ":30000 " || break; sleep 2; done + +docker run -d --name glm53-nvfp4 --gpus all --shm-size 64g --ipc=host --cap-add SYS_PTRACE \ + -v /root/sglang_patch2/layer_setup.py:/sgl-workspace/sglang/python/sglang/srt/model_executor/model_runner_components/layer_setup.py:ro \ + -v /root/sglang_patch2/validation_hook.py:/sgl-workspace/sglang/python/sglang/srt/arg_groups/validation_hook.py:ro \ + -v /root/eagle_worker_v2_mask.py:/sgl-workspace/sglang/python/sglang/srt/speculative/eagle_worker_v2.py:ro \ + -v /root/sglang_patch2/eagle_worker_common.py:/sgl-workspace/sglang/python/sglang/srt/speculative/eagle_worker_common.py:ro \ + -v /root/sglang_patch2/deepseek_nextn.py:/sgl-workspace/sglang/python/sglang/srt/models/deepseek_nextn.py:ro \ + -v /root/sglang_patch2/scheduler_pp_mixin.py:/sgl-workspace/sglang/python/sglang/srt/managers/scheduler_pp_mixin.py:ro \ + -v /root/sglang_patch2/deepseek_v2.py:/sgl-workspace/sglang/python/sglang/srt/models/deepseek_v2.py:ro \ + -e SGLANG_PP_SPEC_DEBUG=0 \ + -e SGLANG_PP_SPEC_SYNC_MASK=0 \ + -e SGLANG_PP_SPEC_FORCE_EAGER_DRAFT=1 \ + -e SGLANG_PP_FORCE_EAGER_VERIFY=1 \ + --restart no -p 30000:30000 \ + -v /data/hf_models:/data/hf_models \ + lmsysorg/sglang:nightly-dev-cu13-20260901-07c8f729 \ + compute-sanitizer --tool memcheck --launch-timeout 0 --target-processes all \ + python3 -m sglang.launch_server \ + --model-path /data/hf_models/GLM-5.3-NVFP4 \ + --tp 8 \ + --mem-fraction-static ${MEMFRAC} \ + --max-running-requests 16 \ + --chunked-prefill-size ${CHUNK} \ + --disable-shared-experts-fusion \ + --moe-runner-backend flashinfer_cutlass \ + --disable-flashinfer-autotune \ + --reasoning-parser glm45 --tool-call-parser glm47 \ + --enable-hierarchical-cache --hicache-ratio 3 \ + ${MTPARGS} \ + ${PAR} \ + --host 0.0.0.0 --port 30000 + +echo "deployed(sanitizer memcheck, SYNC_MASK=0): par=[${PAR}] mtp=${MTPMODE}; waiting for health (up to 3600s)..." +for i in $(seq 10 10 3600); do + code=$(curl -s -o /dev/null -m3 -w '%{http_code}' http://127.0.0.1:30000/health 2>/dev/null) + if [ "$code" = "200" ]; then + echo "healthy after ${i}s" + docker logs glm53-nvfp4 2>&1 | grep -oE "max_total_num_tokens = [0-9]+" | head -1 + exit 0 + fi + if ! docker ps --format '{{.Names}}' | grep -q '^glm53-nvfp4$'; then + echo "CONTAINER DIED after ${i}s"; docker logs --tail 80 glm53-nvfp4 2>&1 | tail -40; exit 1 + fi + sleep 10 +done +echo "TIMEOUT waiting for health"; exit 1 diff --git a/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/killer_cc16.sh b/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/killer_cc16.sh new file mode 100644 index 0000000..3815bd3 --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r36_degloo/scripts/killer_cc16.sh @@ -0,0 +1,64 @@ +#!/bin/bash +# Sanitizer 竞态触发 + 崩溃收割（60.8, deploy_ppmtp_sanitize.sh 配套） +# 自带健康等待；健康后打 bench cc16 剖面杀手负载（16x16384->512 random-ids, +# flush-cache 队列churn——09-06/07 档案中唯一可靠的竞态触发器）； +# 容器死亡时收割全量 docker logs 到 /root/san_crash_s.log。 +# 用法: nohup bash killer_cc16.sh [seed] > /root/killer_s.log 2>&1 & +SEED=${1:-6401} +URL=http://127.0.0.1:30000 + +echo "killer: waiting for health (up to 3600s)..." +HEALTHY=0 +for i in $(seq 1 360); do + code=$(curl -s -o /dev/null -m5 -w '%{http_code}' $URL/health 2>/dev/null) + if [ "$code" = "200" ]; then HEALTHY=1; break; fi + if ! docker ps --format '{{.Names}}' 2>/dev/null | grep -q '^glm53-nvfp4$'; then + echo "CONTAINER-DIED-BEFORE-HEALTH at $(date +%T)" + docker logs glm53-nvfp4 > /root/san_crash_prehealth_s${SEED}.log 2>&1 + echo "log saved: /root/san_crash_prehealth_s${SEED}.log" + exit 2 + fi + sleep 10 +done +[ $HEALTHY -eq 1 ] || { echo "HEALTH-TIMEOUT"; exit 4; } +echo "healthy at $(date +%T)" + +docker exec glm53-nvfp4 python3 -m sglang.bench_serving \ + --backend sglang --host 127.0.0.1 --port 30000 \ + --dataset-name random-ids --tokenizer /data/hf_models/GLM-5.3-NVFP4 \ + --num-prompts 16 --random-input-len 16384 --random-output-len 512 \ + --random-range-ratio 1.0 --max-concurrency 16 \ + --temperature 0.0 --flush-cache --warmup-requests 1 --seed ${SEED} \ + --output-file /data/hf_models/bs_results/san_cc16_s${SEED}.json \ + > /root/san_bench_s${SEED}.log 2>&1 & +BPID=$! +echo "bench pid=$BPID started $(date +%T)" + +BADH=0 +while kill -0 $BPID 2>/dev/null; do + sleep 20 + if ! docker ps --format '{{.Names}}' 2>/dev/null | grep -q '^glm53-nvfp4$'; then + echo "CONTAINER-DIED at $(date +%T)" + docker logs glm53-nvfp4 > /root/san_crash_s${SEED}.log 2>&1 + echo "full log saved: /root/san_crash_s${SEED}.log ($(wc -l < /root/san_crash_s${SEED}.log) lines)" + grep -nE "=========|Invalid|invalid|illegal|CUDA_ERROR|out of bounds|Error Summ" /root/san_crash_s${SEED}.log | head -40 + exit 2 + fi + code=$(curl -s -o /dev/null -m5 -w '%{http_code}' $URL/health 2>/dev/null) + if [ "$code" != "200" ]; then + BADH=$((BADH+1)) + if [ $BADH -ge 3 ]; then + echo "HEALTH-DEAD at $(date +%T)" + docker logs glm53-nvfp4 > /root/san_crash_s${SEED}.log 2>&1 + echo "full log saved: /root/san_crash_s${SEED}.log ($(wc -l < /root/san_crash_s${SEED}.log) lines)" + grep -nE "=========|Invalid|invalid|illegal|CUDA_ERROR|out of bounds|Error Summ" /root/san_crash_s${SEED}.log | head -40 + exit 3 + fi + else + BADH=0 + fi +done +wait $BPID; RC=$? +echo "BENCH-FINISHED rc=$RC $(date +%T) —— 竞态未触发（负载全程存活）" +tail -8 /root/san_bench_s${SEED}.log +exit $RC