From a3b8f1b2ab46fa605107d098f9c168755130c7d9 Mon Sep 17 00:00:00 2001 From: yy-fighting <2351884576@qq.com> Date: Wed, 9 Sep 2026 01:24:29 +0800 Subject: [PATCH] r37: fix PP+MTP verify CUDA graph (pre-planned path missing pp_proxy fill); beats A16 by 40% on killer, now serving on 60.8 --- deploy/CURRENT.md | 2 +- .../glm53_ppmtp_r37_verify_graph/README.md | 66 + .../patches/decode_cuda_graph_runner_fix.py | 1770 +++++++++++++++++ .../patches/eagle_worker_common_bisect.py | 796 ++++++++ .../results/bench_results.md | 42 + .../results/conc_r37fix.txt | 15 + .../results/killer_s6402_a16.txt | 12 + .../results/killer_s6403_a16.txt | 12 + .../results/md5s.txt | 3 + .../results/qg_final.txt | 12 + .../scripts/deploy_ppmtp_r37.sh | 81 + 11 files changed, 2810 insertions(+), 1 deletion(-) create mode 100644 experiments/pro6000/glm53_ppmtp_r37_verify_graph/README.md create mode 100644 experiments/pro6000/glm53_ppmtp_r37_verify_graph/patches/decode_cuda_graph_runner_fix.py create mode 100644 experiments/pro6000/glm53_ppmtp_r37_verify_graph/patches/eagle_worker_common_bisect.py create mode 100644 experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/bench_results.md create mode 100644 experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/conc_r37fix.txt create mode 100644 experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/killer_s6402_a16.txt create mode 100644 experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/killer_s6403_a16.txt create mode 100644 experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/md5s.txt create mode 100644 experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/qg_final.txt create mode 100644 experiments/pro6000/glm53_ppmtp_r37_verify_graph/scripts/deploy_ppmtp_r37.sh diff --git a/deploy/CURRENT.md b/deploy/CURRENT.md index 9a553ed..00ee5b8 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 进程。**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/` | +| 60.8 | `glm53-nvfp4`（Up，:30000，restart=unless-stopped） | **r37 PP+MTP verify-graph 口径在役**（09-08 深夜取代 A16：TP4PP2 + EAGLE 3/1/4 + verify CUDA graph + draft eager + de-GLOO + SYNC_MASK=127，9 挂载文件，启动 `bash /root/deploy_ppmtp_r37.sh '--tp 4 --pp-size 2 --disable-overlap-schedule --max-prefill-tokens 16384' mtp 8192 0.88 1 1 0 0`）。质量门 7/7 ×2、conc_test×3、killer×6 零崩溃；killer cc16 双种子 **63.8/64.7s vs A16 105.6/96.0s（-40%）**、TPOT 41.5/39.3ms、TTFT 33s vs 58-61s（近乎减半）。修复根因=pre-planned 早退路径缺 pp_proxy 补拷（详见 `experiments/pro6000/glm53_ppmtp_r37_verify_graph/`）。A16 切回脚本 deploy_glm53_605.sh 留盘；sglang_patch2/eagle_worker_common.py 已升级 bisect 版（"1"/"0" 语义兼容） | `experiments/pro6000/glm53_ppmtp_r37_verify_graph/`（patches+scripts+results 全量）；前史 `experiments/pro6000/glm53_ppmtp_r36_degloo/` | ## 方案 A-F 一览（GLM-5.3-NVFP4 @ pro6000，2026-09-08 双场景报告口径） diff --git a/experiments/pro6000/glm53_ppmtp_r37_verify_graph/README.md b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/README.md new file mode 100644 index 0000000..4e9db33 --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/README.md @@ -0,0 +1,66 @@ +# GLM-5.3-NVFP4 PP+MTP verify CUDA graph 修复（r37）— 2026-09-08 @ 174.1.60.8 + +上一个实验（glm53_ppmtp_r36_degloo）判定：PP+MTP eager 模式正确但 MTP 机器开销吃光收益（step 239ms@cc16 vs +净胜拐点 148ms），图模式（r36g）则正确性崩溃（accept 2.85→1.05 + 数字乱码）。本实验继续修复图捕获并兑现 +MTP+PP 双优势。**结果：修复成功，killer e2e -40%、TPOT -40%、TTFT 近乎减半，全面胜过在役 A16 配置并留役生产。** + +## 环境 +- 镜像 lmsysorg/sglang:nightly-dev-cu13-20260901-07c8f729，GLM-5.3-NVFP4，8×RTX 6000D（SM120） +- TP4PP2 + EAGLE（3 steps / topk1 / 4 draft tokens），DSA，hicache 3，memfrac 0.88，chunk 8192， + --disable-overlap-schedule，SYNC_MASK=127，de-GLOO 中继（r36 遗产） +- r36 的 8 个挂载文件全部沿用 + 本实验新增第 9 挂载（decode_cuda_graph_runner） + +## 排查过程（按时间序） +1. **:1515 切片修复**（r33 注释的根因猜测）：decode runner `execute()` 对非末 stage 的 PPProxyTensors + 按 `self.bs`（请求数）切片，截断 TARGET_VERIFY 树行（bs×4）→ 改按 `raw_num_token`。**必要但不充分**： + r37 实测 accept 仍 1.02-1.27 + 乱码（bs=1 零 padding 下确定性复现 → 排除竞态/padding）。 +2. **按 stage 二分**：把 /root/sglang_patch2/eagle_worker_common.py 的 `_PP_FORCE_EAGER_VERIFY` 门扩展为 + 1/2/3 模式（全 eager / 仅首 stage eager / 仅末 stage eager）。判决： + - 模式 2（PP0 eager + PP1 图）→ 仍乱码；模式 3（PP0 图 + PP1 eager）→ 完全正确。 + - **缺陷唯一在末 stage（PP1）的 verify 图**；PP0 图（含 :1515 修复）被实证正确。 +3. **遥测**（runner 内置 VRFY-IN/BUF/G/SLOT + xcheck，SGLANG_PP_SPEC_DEBUG=2/3）： + - 入口 relayed hidden 每轮新鲜正常；返回 hidden_states 新鲜；logits argmax 全是高频英文 token（' of'/' the'/':'， + tokenizer 解码实证）→ "模型没看到上下文"signature。 + - VRFY-SLOT：pp 槽位注册健康（enabled=True copy_from_fb=True buffer=静态张量 data_ptr 一致）。 + - **VRFY-BUF：fill_from 之后静态 pp 缓冲仍全零；VRFY-IN needs_init=False** → 真凶。 + +## 根因（真·根因，取代 r33 的 :1515 猜测为唯一根因） +spec worker 的 plan-stream 在 prepare 阶段预跑 `load_batch`（eagle_prepare_for_verify 的 overlap 设计）——彼时 +**PP 中继尚未到达**，`pp_proxy_tensors=None`，`fill_from` 跳过全部 pp 槽位，仅完成常规槽位 + 注意力元数据， +batch 被标记 pre-planned（needs_forward_metadata_init()=False）。前向流 `execute()` 于是走 pre-planned 早退路径， +**该路径只补拷 input_ids/positions 两个 spec 必需字段** → pp_proxy（hidden_states/residual/topk_indices） +静态缓冲永远是分配时的全零。捕获图内的 clone kernel（run_once 体内，:1212）从静态缓冲拷零 → 末 stage 在零激活上 +计算 → logits 为"无上下文"垃圾 → 全拒 + 乱码。 + +一切旁证自洽： +- PP0（首 stage）不吃 pp 输入（input_ids/positions 恰好是早退路径补拷的两个）→ 模式 3 正确； +- B'（nomtp+图+PP）无 spec worker plan-stream → 走完整 fill_from 路径 → 正确； +- TP8 A16 无 PP → 正确； +- 图 logits ≈ eager 参考滞后一轮（xcheck 实测）：零 token 嵌入 + 上一轮 KV 历史的注意力产出。 + +## 修复（patches/decode_cuda_graph_runner_fix.py，第 9 挂载文件） +两处改动： +1. `execute()` :1515：PPProxyTensors 抽取改按 `self.raw_num_token` 切（原 :bs 截断树行；decode width=1 时二者等价于 + raw/padded 行数之差，树行场景必须 raw_num_token）。 +2. `load_batch()` pre-planned 早退路径：新增 pp_proxy 补拷——`pp_proxy_tensors` 各 key 拷入静态缓冲 + `[:src.shape[0]]`，pad 尾（`self.bs*captured_req_width` 之外）清零保确定性。这是本修复的核心。 + +生产运行时遥测代码（VRFY-* / xcheck）随文件留存，由 SGLANG_PP_SPEC_DEBUG 门控（0=关闭，无开销）。 + +## 验证阶梯（全部通过） +- 质量门 7/7 ×2 boot（GSM8K×5/中文推理/tool call）； +- accept 恢复 3.59-3.75（eager 档）； +- conc_test ×3 ALL-OK、killer ×6 零崩溃（三个 boot 累计）； +- 性能见 results/bench_results.md：vs eager -40% e2e；vs A16 在役配置 -40% e2e、TTFT 33s vs 58-61s。 + +## 生产终态 +60.8 glm53-nvfp4:30000 = r37 胜者配置（verify 图 + draft eager + de-GLOO + mask127）， +`--restart unless-stopped`，已取代 A16 留役（A16 脚本 deploy_glm53_605.sh 留盘可切回）。 +draft 图（DRAFTEAGER=0）实测无增益（噪声带内互有胜负），保持 eager 以少一个变量。 + +## 遗产与提醒 +- /root/sglang_patch2/eagle_worker_common.py 已被本实验覆盖为 bisect 版（md5 babb6461…）：env "1"/"0" 语义与 + r36 版完全兼容，新增 "2"/"3" 单侧模式；r36 档案中的旧 md5 以本目录为准。 +- 语料池 61,772 token 剩余（< 一次 i8k 262,144），i8k 永久无重跑余量，killer random-ids 为标准口径。 +- killer 的 raw bench JSON（/data/hf_models/bs_results/san_cc16_s*.json）被 A16 复测覆盖， + 胜者数字以 results/bench_results.md 记录为准。 diff --git a/experiments/pro6000/glm53_ppmtp_r37_verify_graph/patches/decode_cuda_graph_runner_fix.py b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/patches/decode_cuda_graph_runner_fix.py new file mode 100644 index 0000000..dce3fdf --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/patches/decode_cuda_graph_runner_fix.py @@ -0,0 +1,1770 @@ +# Copyright 2023-2026 SGLang Team +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""DecodeCudaGraphRunner — runs DECODE / TARGET_VERIFY / DLLM_EXTEND under +a pluggable backend. + +Backend selection comes from cuda_graph_config.decode: + - "full" — default, FullCudaGraphBackend: one + torch.cuda.CUDAGraph per shape. + - "breakable" — experimental, BreakableCudaGraphBackend: + segmented capture (no torch.compile). + - "tc_piecewise" — not implemented for decode; logs a one-shot warning + and falls back to "full". +""" + +from __future__ import annotations + +import contextlib +import inspect +import logging +import os +from types import SimpleNamespace +from typing import TYPE_CHECKING, Callable, Optional, Union + +import torch +import tqdm +from torch.profiler import ProfilerActivity, profile + +from sglang.srt.compilation import torch_compile_decoration +from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config +from sglang.srt.distributed.parallel_state import ( + graph_capture, + set_pdmux_status, +) +from sglang.srt.dllm.config import DllmConfig +from sglang.srt.environ import envs +from sglang.srt.layers.attention.base_attn_backend import ( + AttentionBackend, + SharedReadEnds, +) +from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp +from sglang.srt.layers.dp_attention import ( + DpPaddingMode, + set_dp_buffer_len, + set_is_extend_in_batch, +) +from sglang.srt.layers.logits_processor import LogitsProcessorOutput +from sglang.srt.layers.utils.cp_utils import is_mla_prefill_cp_enabled +from sglang.srt.model_executor.cuda_graph_buffer_registry import ( + CudaGraphBufferRegistry, + build_decode_registry, +) +from sglang.srt.model_executor.forward_batch_info import ( + CaptureHiddenMode, + ForwardBatch, + ForwardMode, + PPProxyTensors, + compute_local_num_token_non_padded, + enable_num_token_non_padded, + get_required_capture_hidden_mode, +) +from sglang.srt.model_executor.forward_context import ForwardContext, forward_context +from sglang.srt.model_executor.runner.base_cuda_graph_runner import ( + BaseCudaGraphRunner, + freeze_gc, + get_batch_sizes_to_capture, +) +from sglang.srt.model_executor.runner.flashinfer_autotune import ( + maybe_flashinfer_autotune_speculative_draft, +) +from sglang.srt.model_executor.runner.metadata_glue_graph import MetadataGlueGraph +from sglang.srt.model_executor.runner.shape_key import ShapeKey +from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import ( + BreakableCudaGraphBackend, +) +from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backend +from sglang.srt.model_executor.runner_backend_utils import ( + CUDA_GRAPH_CAPTURE_FAILED_MSG, +) +from sglang.srt.model_executor.runner_utils.buffers import ( + DecodeInputBuffers, +) +from sglang.srt.model_executor.runner_utils.capture_mode import ( + _set_capture_dsa_variant, + _set_capture_lora_variant, + model_capture_mode, +) +from sglang.srt.model_executor.runner_utils.deepep_adapter import ( + DeepEPCudaGraphRunnerAdapter, +) +from sglang.srt.model_executor.runner_utils.pool import ( + get_or_create_global_graph_capture_stream, +) +from sglang.srt.model_executor.runner_utils.shared_read_event import make_external_event +from sglang.srt.multiplex.pdmux_context import get_current_stream_idx, get_stream_groups +from sglang.srt.runtime_context import ( + get_exec, + get_flags, + get_parallel, + get_spec, +) +from sglang.srt.speculative.ragged_verify import resolve_ragged_verify_layout +from sglang.srt.utils import ( + empty_context, + get_available_gpu_memory, + is_hip, + require_attn_tp_gather, + require_mlp_tp_gather, +) +from sglang.srt.utils.device_timer import device_timer_ctx +from sglang.srt.utils.profile_utils import ( + export_cuda_graph_capture_trace, + graph_capture_profile_dir, +) + +try: + from kt_kernel import KTMoEWrapper + + KTRANSFORMERS_AVAILABLE = True +except ImportError: + KTRANSFORMERS_AVAILABLE = False + +logger = logging.getLogger(__name__) + +# SGLANG_PP_SPEC_DEBUG=2: per-replay telemetry for TARGET_VERIFY graphs — +# relayed proxy inputs at entry (non-first stages) and extracted logits at +# exit (last stage). First 24 verify rounds only. +# =3: additionally re-run the same batch through the eager forward on the +# last stage and log its logits as a reference (KV rewrite is idempotent: +# same tokens, same slots, same values). +_SPEC_LOGITS = os.getenv("SGLANG_PP_SPEC_DEBUG", "0") in ("2", "3") +_SPEC_XCHECK = os.getenv("SGLANG_PP_SPEC_DEBUG", "0") == "3" + +if TYPE_CHECKING: + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.spec_info import SpeculativeAlgorithm + + +def ragged_verify_compact_graphs_enabled(spec_algorithm: SpeculativeAlgorithm) -> bool: + if not spec_algorithm.supports_ragged_verify(): + return False + from sglang.srt.speculative.ragged_verify import ragged_verify_compact_enabled + + return ragged_verify_compact_enabled() + + +def build_replay_fb_view( + forward_batch: ForwardBatch, + buffers: DecodeInputBuffers, + bs: int, + raw_bs: int, + num_tokens: int, + seq_len_fill_value: int, + capture_forward_mode: ForwardMode, + is_encoder_decoder: bool, +) -> SimpleNamespace: + """Construct a ForwardBatch-like view for backend replay-side init. + + Combines the original forward_batch (for unpadded / per-iter + fields like spec_info, out_cache_loc, and the runtime + actual_forward_mode) with the padded capture-time buffers from + buffers (for req_pool_indices, seq_lens, seq_lens_cpu, + positions, encoder_lens). + + forward_mode is the capture-time mode (used by backends for + bucket / dispatch decisions); actual_forward_mode is the + runtime mode (may be IDLE while the captured graph targets DECODE + — DSV4's replay metadata prep uses this for IDLE substitution). + + Subsumes the _replay_forward_batch side channel that DSV4 used to + read out-of-band before the init_forward_metadata 3-method ABC. + """ + return SimpleNamespace( + batch_size=bs, + forward_mode=capture_forward_mode, + actual_forward_mode=forward_batch.forward_mode, + input_ids=buffers.input_ids[:num_tokens], + positions=buffers.positions[:num_tokens], + req_pool_indices=buffers.req_pool_indices[:bs], + seq_lens=buffers.seq_lens[:bs], + seq_lens_sum=( + None + if forward_batch.seq_lens_sum is None + else forward_batch.seq_lens_sum + (bs - raw_bs) * seq_len_fill_value + ), + # Propagate mirror absence: the pinned buffer is not refreshed when the + # batch has no CPU mirror; a stale non-None tensor defeats None-guards. + seq_lens_cpu=( + None if forward_batch.seq_lens_cpu is None else buffers.seq_lens_cpu[:bs] + ), + num_padding=bs - raw_bs, + encoder_lens=buffers.encoder_lens[:bs] if is_encoder_decoder else None, + out_cache_loc=getattr(forward_batch, "out_cache_loc", None), + out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None), + # The mamba-track registry slot (VIRTUAL ids) is the v2p translate SOURCE + # for the backend, which copies the result into its own static buffer and + # reads THAT in the decode track-save — this slot is never mutated. None + # when mamba-track is disabled, slice to [:bs] like every other buffer + mamba_track_indices=( + None + if buffers.mamba_track_indices is None + else buffers.mamba_track_indices[:bs] + ), + spec_info=forward_batch.spec_info, + ) + + +class DecodeCudaGraphRunner(BaseCudaGraphRunner): + """Decode-phase CUDA graph runner. + + Owns: static input buffers (DecodeInputBuffers), capture-bs list, + attention backend, two-batch-overlap plugin, DeepEP adapter, and the + pluggable self.backend that handles the actual capture/replay. + """ + + def __init__( + self, + model_runner: ModelRunner, + *, + attn_backend=None, + speculative_num_steps: Optional[int] = None, + speculative_num_draft_tokens: Optional[int] = None, + ): + super().__init__(model_runner) + + # In-graph metadata prep: shared buffers -> in-graph private data + self.in_graph_metadata_prep_done: Optional[torch.cuda.Event] = None + + # --- core state ------------------------------------------------ + self.enable_torch_compile = get_flags().capture.enable_torch_compile + self.disable_padding = model_runner.server_args.disable_cuda_graph_padding + self.is_encoder_decoder = model_runner.model_config.is_encoder_decoder + self.require_mlp_tp_gather = ( + require_mlp_tp_gather() and not self._forward_is_dp_local(model_runner) + ) + self.require_attn_tp_gather = require_attn_tp_gather() + # Composite predicates derive from the instance values so the dp-local + # draft exemption above stays consistent (require_gathered_buffer == + # mlp_tp_gather or attn_tp_gather; require_mlp_sync adds dp attention). + self.require_gathered_buffer = ( + self.require_mlp_tp_gather or self.require_attn_tp_gather + ) + self.require_mlp_sync = ( + get_parallel().enable_dp_attention or self.require_gathered_buffer + ) + self.enable_two_batch_overlap = ( + model_runner.server_args.enable_two_batch_overlap + ) + self.use_ngram_embedding = model_runner.ngram_embedding_manager.enabled + if self.use_ngram_embedding: + hf_config = model_runner.model_config.hf_config + self.ngram_embedding_n = hf_config.ngram_embedding_n + self.ngram_embedding_k = hf_config.ngram_embedding_k + self.speculative_algorithm = get_spec().speculative_algorithm + self.enable_profile_cuda_graph = ( + model_runner.server_args.enable_profile_cuda_graph + ) + + # --- DSA dense-decode dual-graph ------------------------------- + # Capture a "dense" (k-only, skip-indexer) and a "sparse" (full indexer) + # decode graph per bs bucket, and dispatch on max_kv_len vs index_topk at + # replay. Auto-enabled for DSA models (index_topk present in the HF + # config) — correct for mixed lengths since any request with + # kv_len > index_topk falls back to the sparse graph. Adds ~52 graphs and + # ~2x capture time. + # + # Scoped to HIP (AMD): the k-only dense-decode fast path has only been + # validated on MI355X. This is common (non-hardware-gated) code, so on + # CUDA we deliberately keep the original behavior (no dual-graph) to + # avoid silently changing the CUDA decode path for DSA models (e.g. + # DeepSeek-V3.2). CUDA can opt in later once validated there. + self.dsa_dual_graph = False + self.dsa_index_topk: Optional[int] = None + from sglang.srt.configs.model_config import ( + get_dsa_index_topk, + is_deepseek_dsa, + ) + + hf_config = model_runner.model_config.hf_config + if is_hip() and is_deepseek_dsa(hf_config): + self.dsa_index_topk = get_dsa_index_topk(hf_config) + self.dsa_dual_graph = True + logger.info( + "[dense-decode] DSA dual-graph enabled: capturing " + "dense (k-only) + sparse (full indexer) decode graphs; " + "dispatch on max_kv_len vs index_topk=%d.", + self.dsa_index_topk, + ) + + self.attn_tp_size = get_parallel().attn_tp_size + self.attn_tp_rank = get_parallel().attn_tp_rank + # True if a DSACPLayerCommunicator-style prefill-CP flavor is active + # (DSA or MLA). These flavors feed a zigzag-split rank-local layout + # into the runner; MHA-arch prefill CP (Qwen3/Qwen2 MoE via PR + # #18233) uses the plain LayerCommunicator with an attn_tp-replicated + # layout and is intentionally excluded so the attn_tp-local + # num_token_non_padded adjustment still runs for it. + self.enable_prefill_cp = ( + is_dsa_enable_prefill_cp() or is_mla_prefill_cp_enabled() + ) + + self.deepep_adapter = DeepEPCudaGraphRunnerAdapter() + + self.dllm_config = DllmConfig.from_server_args(model_runner.server_args) + self.is_dllm = self.dllm_config is not None + self.attn_backend = attn_backend or model_runner.attn_backend + self.speculative_num_steps = ( + get_spec().speculative_num_steps + if speculative_num_steps is None + else speculative_num_steps + ) + self.speculative_num_draft_tokens = ( + get_spec().speculative_num_draft_tokens + if speculative_num_draft_tokens is None + else speculative_num_draft_tokens + ) + + # --- capture mode + tokens-per-bs ------------------------------ + self.capture_forward_mode = ForwardMode.DECODE + self.capture_hidden_mode = self.return_hidden_states_mode + # Static capture width. + self.captured_req_width = model_runner.decode_num_tokens_per_req( + num_draft_tokens=self.speculative_num_draft_tokens + ) + if model_runner.spec_algorithm.is_speculative(): + if self.model_runner.is_draft_worker: + # Draft workers can use TARGET_VERIFY mode. + if ( + not self.model_runner.spec_algorithm.supports_target_verify_for_draft() + ): + raise RuntimeError("This should not happen") + self.capture_forward_mode = ForwardMode.TARGET_VERIFY + elif self.is_dllm: + self.capture_forward_mode = ForwardMode.DLLM_EXTEND + + # --- bucket sizes --------------------------------------------- + self.capture_bs, self.compile_bs = get_batch_sizes_to_capture( + model_runner, self.captured_req_width + ) + if KTRANSFORMERS_AVAILABLE: + KTMoEWrapper.set_capture_batch_sizes(self.capture_bs) + + self.ragged_verify_mode = ( + ragged_verify_compact_graphs_enabled(self.model_runner.spec_algorithm) + and (self.capture_forward_mode == ForwardMode.TARGET_VERIFY) + and not self.model_runner.is_draft_worker + ) + self.capture_num_tokens: Optional[list[int]] = ( + self._build_ragged_verify_token_buckets() + if self.ragged_verify_mode + else None + ) + self._ragged_graph_size = 0 + # Per-tier capture layouts; their verify_lens / qo_indptr tensors are + # baked into the captured graphs and refreshed in place each replay. + self._captured_ragged_layouts: dict[int, object] = {} + if self.ragged_verify_mode and ( + self.enable_two_batch_overlap + or model_runner.lora_manager is not None + or self.disable_padding + ): + raise ValueError( + "Compact ragged verify does not support two-batch-overlap, " + "LoRA, or disable-cuda-graph-padding (bs pads to the captured " + "tier); disable SGLANG_RAGGED_VERIFY_MODE or the conflicting " + "feature." + ) + + # Attention backend + self.max_bs = max(self.capture_bs) + self.max_num_token = self.max_bs * self.captured_req_width + self.attn_backend.init_cuda_graph_state(self.max_bs, self.max_num_token) + + # Init PDMux if needed + self.maybe_init_pdmux() + self.seq_len_fill_value = ( + self.attn_backend.get_cuda_graph_seq_len_fill_value() + if self.dllm_config is None + else self.dllm_config.block_size + ) + + # Non-zero encoder length ensures cross-attention kernels are captured in the graph. + self.encoder_len_fill_value = ( + getattr(model_runner.model_config.hf_config, "max_source_positions", 0) + if self.is_encoder_decoder + else 0 + ) + + if self.enable_torch_compile: + set_torch_compile_config() + + if self.model_runner.lora_manager is not None: + # Phase 2 of LoRA CUDA graph init: dense LoRA batch metadata. + # Phase 1 (MoE buffers) was handled earlier in ModelRunner via + # lora_manager.init_cuda_graph_moe_buffers(). + self.model_runner.lora_manager.init_cuda_graph_batch_info( + max_bs_in_cuda_graph=self.max_bs, + num_tokens_per_req=self.captured_req_width, + ) + + enable_mamba_track = ( + self.model_runner.server_args.enable_mamba_extra_buffer() + and self.model_runner.spec_algorithm.is_none() + ) + + if self.require_gathered_buffer: + assert self.require_mlp_tp_gather or self.require_attn_tp_gather + + # --- buffers --------------------------------------------------- + self.buffers: DecodeInputBuffers = DecodeInputBuffers.create( + device=self.device, + max_bs=self.max_bs, + max_num_token=self.max_num_token, + hidden_size=self.model_runner.model_config.hidden_size, + next_token_logits_buffer=self.model_runner.graph_shared_output.get_logits_buffer( + self.model_runner.model_config.vocab_size, rows=self.max_num_token + ), + dtype=self.model_runner.model_config.dtype, + dp_size=self.dp_size, + pp_size=self.pp_size, + is_encoder_decoder=self.is_encoder_decoder, + require_mlp_tp_gather=self.require_mlp_tp_gather, + seq_len_fill_value=self.seq_len_fill_value, + encoder_len_fill_value=self.encoder_len_fill_value, + num_tokens_per_req=self.captured_req_width, + cache_loc_dtype=self._cache_loc_dtype(), + enable_mamba_track=enable_mamba_track, + ne_token_table=( + model_runner.ngram_embedding_manager.table + if self.use_ngram_embedding + else None + ), + hc_hidden_size=getattr( + self.model_runner.model_config, "hc_hidden_size", None + ), + pp_proxy_topk_size=self.model_runner.get_pp_proxy_topk_size(), + pp_proxy_residual_num_blocks=( + self.model_runner.get_pp_proxy_residual_num_blocks() + ), + ) + self.buffers.share_buffers() + # FB-shared slot registry adopting DecodeInputBuffers storage (same + # physical tensors, stable data_ptr for capture vs replay). Provides + # the unified fill_from / slot access surface, replacing + # populate_from_forward_batch on capture/replay paths. + self.buffer_registry: CudaGraphBufferRegistry = build_decode_registry( + device=self.device, + max_bs=self.max_bs, + max_num_token=self.max_num_token, + seq_len_fill_value=self.seq_len_fill_value, + cache_loc_dtype=self._cache_loc_dtype(), + enable_mamba_track=enable_mamba_track, + is_encoder_decoder=self.is_encoder_decoder, + encoder_len_fill_value=self.encoder_len_fill_value, + enable_num_token_non_padded=enable_num_token_non_padded(), + require_gathered_buffer=self.require_gathered_buffer, + enable_prefill_cp=self.enable_prefill_cp, + require_mlp_tp_gather=self.require_mlp_tp_gather, + dp_size=self.dp_size, + source=self.buffers, + ) + + # Captures the per-replay attention-metadata prep into a small CUDA + # graph; see metadata_glue_graph.py for the correctness contract. + # Force-off for DFlash-family spec: verify installs host-fed fast + # plans (sync-free begin_forward that recomputes plan inputs on the + # host every replay), and capturing one freezes the capture-time + # plan — drafts go stale and accept length collapses to ~1. + enable_metadata_glue = envs.SGLANG_ENABLE_METADATA_GLUE_GRAPH.get() + if enable_metadata_glue and model_runner.spec_algorithm.is_dflash_family(): + logger.warning( + "SGLANG_ENABLE_METADATA_GLUE_GRAPH is incompatible with " + "DFlash-family speculative decoding (host-fed fast verify " + "plans must re-run on the host every replay); disabling the " + "metadata glue graph." + ) + enable_metadata_glue = False + self._metadata_glue = ( + MetadataGlueGraph(self.device) if enable_metadata_glue else None + ) + + # --- backend --------------------------------------------------- + self.backend = resolve_decode_backend(self) + + # --- capture -------------------------------------------------- + try: + with model_capture_mode(): + self.capture() + except RuntimeError as e: + raise Exception( + f"Capture cuda graph failed: {e}\n" f"{CUDA_GRAPH_CAPTURE_FAILED_MSG}" + ) + + def _record_in_graph_metadata_prep_done(self): + # Purely a marker at this point in the graph; where the shared reads + # actually end is the attn backend's call. + if not self.device_module.is_current_stream_capturing(): + # Warmup shares this body. Breakable capture still plants: it opens + # segment 1 on context entry and every segment re-arms the node. + # Routed through device_module so XPU (torch.xpu) is picked up + # instead of hitting torch.cuda dummy stubs on non-CUDA builds. + return + if self.in_graph_metadata_prep_done is None: + self.in_graph_metadata_prep_done = make_external_event(self.device_module) + event = self.in_graph_metadata_prep_done + if event is not None: + # Stays None without external-event support, so the read-end + # resolution below never hands out an unrecorded event. + event.record() + + def _replay_attn_backend(self) -> AttentionBackend: + # Under pdmux each stream replays on its own group member. + if self.enable_pdmux: + return self.model_runner.decode_attn_backend_group[get_current_stream_idx()] + return self.attn_backend + + def _resolve_shared_read_ends(self, attn_backend, forward_mode) -> SharedReadEnds: + declared = attn_backend.shared_read_ends(forward_mode) + if ( + declared is SharedReadEnds.IN_REPLAY + and self.in_graph_metadata_prep_done is None + ): + # TODO: this lands EARLIER than declared; POST_REPLAY is the sound one. + return SharedReadEnds.PRE_REPLAY + return declared + + def _publish_read_done(self, in_graph: bool): + """Hand the scheduler's WAR barrier the event marking this phase's + shared-buffer reads as done.""" + if in_graph: + # Reads end at the in-graph marker: wire it through, don't re-record. + self.model_runner.shared_read_done_event = self.in_graph_metadata_prep_done + else: + read_done = self.device_module.Event() + read_done.record() + self.model_runner.shared_read_done_event = read_done + + def _build_ragged_verify_token_buckets(self) -> list[int]: + buckets = sorted({bs * self.captured_req_width for bs in self.capture_bs}) + assert buckets and buckets[0] > 0, f"{buckets=}" + return buckets + + def _autotune_buffers(self): + """Reuse these static decode buffers (sized to max_bs) for the warmup + flashinfer-autotune dummy forward instead of allocating a throwaway set + — see BaseRunner._autotune_buffers / BaseRunner._dummy_run. + + The dummy forward derives its shape from max_bs and must match these + buffers exactly; _dummy_run asserts that. Every autotune-reachable + decode shape (plain decode, spec target-verify) matches. DLLM would not + (its buffers hold block_size tokens/bs while the dummy run derives 1), + but DLLM does not use a flashinfer MoE backend, so autotune never runs + for it and this is never reached there. + """ + return self.buffers, self.max_bs + + def maybe_init_pdmux(self): + if self.enable_pdmux: + self.stream_groups = get_stream_groups() + for attn_backend in self.model_runner.decode_attn_backend_group: + attn_backend.init_cuda_graph_state(self.max_bs, self.max_num_token) + + def _cache_loc_dtype(self): + return torch.int64 + + def _make_graph_key( + self, size, stream_idx=None, variant_label=None, dsa_variant=None + ): + return ShapeKey( + size=size, + stream_idx=stream_idx, + variant_label=variant_label, + dsa_variant=dsa_variant, + ) + + def _capture_graph_size(self, *, bs: int, num_tokens: int) -> int: + return num_tokens if self.ragged_verify_mode else bs + + def _resolve_dsa_variant(self, forward_batch: ForwardBatch) -> Optional[str]: + """Host dispatch: pick which pre-captured DSA decode graph to replay + from the batch-max kv_len. If any request has kv_len > index_topk + the dense (k-only) graph would be wrong for it, so the whole batch uses + the sparse (full indexer) graph. Returns None when dual-graph is off.""" + if not getattr(self, "dsa_dual_graph", False): + return None + seq_lens_cpu = getattr(forward_batch, "seq_lens_cpu", None) + if seq_lens_cpu is not None and seq_lens_cpu.numel() > 0: + # Host-side mirror (maintained incrementally for plain decode) — no + # d2h sync needed. + max_kv_len = int(seq_lens_cpu.max().item()) + elif forward_batch.seq_lens is not None and forward_batch.seq_lens.numel() > 0: + # Fallback: a single scalar reduction d2h (cheap, per-step). + max_kv_len = int(forward_batch.seq_lens.max().item()) + else: + # No length info: be safe and use the correct-for-all sparse graph. + return "sparse" + return "dense" if max_kv_len <= self.dsa_index_topk else "sparse" + + def _resolve_lora_variant(self, forward_batch: ForwardBatch): + if not getattr(self, "record_nolora_graph", False): + return None + if forward_batch.lora_ids is not None and any( + uid is not None for uid in forward_batch.lora_ids + ): + return "lora" + return "nolora" + + @staticmethod + def _forward_is_dp_local(model_runner) -> bool: + """The DSpark dense draft runs attn-TP-local (draft_tp_context): each + DP rank drafts independently with no cross-DP collective, so its + hand-built batches carry no dp-global metadata and must key graphs by + local batch size. Everything else keeps the dp-global padding path.""" + if not model_runner.is_draft_worker: + return False + if not model_runner.spec_algorithm.is_dspark(): + return False + from sglang.srt.speculative.dspark_components.dspark_config import ( + draft_is_deepseek_v4, + ) + + return not draft_is_deepseek_v4() + + def _ragged_capture_slots(self, num_tokens: int) -> int: + if envs.SGLANG_TEST_RAGGED_VERIFY_FORCE_UNIFORM_CAPTURE.get(): + return num_tokens // self.captured_req_width + return min(num_tokens, self.max_bs) + + def _capture_ragged_verify_layout(self, num_tokens: int): + if not self.ragged_verify_mode: + return None + if envs.SGLANG_TEST_RAGGED_VERIFY_FORCE_UNIFORM_CAPTURE.get(): + return None + from sglang.srt.speculative.ragged_verify import ( + RaggedVerifyLayout, + build_capture_verify_lens, + ) + + verify_lens_cpu = build_capture_verify_lens( + num_tokens=num_tokens, + num_slots=self._ragged_capture_slots(num_tokens), + num_draft_tokens=self.captured_req_width, + ) + layout = RaggedVerifyLayout.from_verify_lens( + verify_lens_cpu=verify_lens_cpu, + device=self.device, + grid=self.capture_num_tokens, + ) + self._captured_ragged_layouts[num_tokens] = layout + return layout + + def _stage_ragged_verify_layout(self, ragged_layout, graph_size_key: int) -> None: + # Without this refresh every replay reuses the capture-time synthetic + # verify_lens / qo_indptr and mis-slices the packed q rows. + cap_layout = self._captured_ragged_layouts.get(graph_size_key) + if cap_layout is None: + return + live = ragged_layout + if live.bs != cap_layout.bs or live.cap is None: + live = live.padded_to_bucket( + padded_bs=cap_layout.bs, cap=self.captured_req_width + ) + cap_layout.verify_lens.copy_(live.verify_lens) + cap_layout.qo_indptr_device.copy_(live.qo_indptr_device) + + @staticmethod + def _max_dp_batch_size(forward_batch: ForwardBatch) -> int: + request_counts = forward_batch.original_global_num_tokens_cpu + if request_counts is None: + raise RuntimeError( + "DP CUDA graph replay requires raw per-rank request counts" + ) + return max(request_counts) + + def can_run_graph(self, forward_batch: ForwardBatch): + # Disable for token embedding overrides (dynamic per-request) + if forward_batch.replace_embeds is not None: + return False + + ragged_layout = ( + resolve_ragged_verify_layout(forward_batch) + if self.ragged_verify_mode + else None + ) + if ragged_layout is not None: + return self._can_run_ragged_verify_graph(forward_batch, ragged_layout) + if self.ragged_verify_mode and forward_batch.forward_mode.is_target_verify(): + return False + + # Uniform-width replay invariant: the batch's actual per-request width + # must match this runner's capture width; anything else falls back to + # eager. (Unset widths pass: not every path fills the field yet.) + spec_info = forward_batch.spec_info + if ( + spec_info is not None + and spec_info.num_tokens_per_req > 0 + and spec_info.num_tokens_per_req != self.captured_req_width + ): + return False + + if self.require_mlp_tp_gather: + cuda_graph_bs = self._max_dp_batch_size(forward_batch) + else: + cuda_graph_bs = forward_batch.batch_size + + graph_key = self._make_graph_key( + cuda_graph_bs, + stream_idx=get_current_stream_idx() if self.enable_pdmux else None, + variant_label=self._resolve_lora_variant(forward_batch), + ) + + is_bs_supported = ( + self.backend.can_run(forward_batch, graph_key) + if self.disable_padding + else cuda_graph_bs <= self.max_bs + ) + + if self.require_mlp_sync: + is_bs_supported = ( + is_bs_supported and forward_batch.can_run_decode_cuda_graph + ) + + # NOTE: cuda graph cannot handle mixed batch (encoder_len = 0) + # If mixed batch cannot be supported, then encoder_lens can be removed in cuda graph + # because the full_text_row_masked_out_mask tensor will always be ones + is_encoder_lens_supported = ( + torch.all(forward_batch.encoder_lens > 0) + if self.is_encoder_decoder + else True + ) + + is_tbo_supported = ( + forward_batch.can_run_tbo if self.enable_two_batch_overlap else True + ) + + is_ngram_supported = ( + ( + forward_batch.batch_size * self.captured_req_width + == forward_batch.input_ids.numel() + ) + if self.model_runner.spec_algorithm.is_ngram() + else True + ) + + return ( + is_bs_supported + and is_encoder_lens_supported + and is_tbo_supported + and is_ngram_supported + ) + + def _can_run_ragged_verify_graph(self, forward_batch: ForwardBatch, ragged_layout): + if not self.attn_backend.supports_ragged_verify_graph: + return False + + admission_tokens = ragged_layout.graph_num_tokens + is_tokens_supported = admission_tokens <= self.capture_num_tokens[ + -1 + ] and forward_batch.batch_size <= self._ragged_capture_slots(admission_tokens) + + is_dp_supported = ( + forward_batch.can_run_decode_cuda_graph if self.require_mlp_sync else True + ) + + is_encoder_lens_supported = ( + torch.all(forward_batch.encoder_lens > 0) + if self.is_encoder_decoder + else True + ) + + capture_hidden_mode_matches = ( + forward_batch.capture_hidden_mode <= self.capture_hidden_mode + ) + + return ( + is_tokens_supported + and is_dp_supported + and is_encoder_lens_supported + and capture_hidden_mode_matches + ) + + def _graph_batch_capture_active(self) -> bool: + """Whether the per-batch-size capture-trace feature is active. + + Gated by SGLANG_GRAPH_BATCH_CAPTURE. The original single-trace export + (SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE) takes precedence: when both are + set we fall back to the original behavior. + """ + return ( + envs.SGLANG_GRAPH_BATCH_CAPTURE.get() + and not envs.SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE.get() + ) + + def _init_profile_context_and_memory_record(self): + if self._graph_batch_capture_active(): + # Per-batch-size capture traces (SGLANG_GRAPH_BATCH_CAPTURE): a + # scheduled profiler is stepped once per batch size (see + # FullCudaGraphBackend.capture_one) and on_trace_ready writes one + # chrome trace per bs. + rank = get_parallel().tp_rank + runner_name = type(self).__name__ + trace_dir = graph_capture_profile_dir() + os.makedirs(trace_dir, exist_ok=True) + + # Track which BS is currently being captured for trace file naming + self._profile_bs_list = list(reversed(self.capture_bs)) + self._profile_bs_idx = 0 + + def on_trace_ready(prof): + bs = self._profile_bs_list[self._profile_bs_idx] + trace_file = os.path.join( + trace_dir, f"{runner_name}_bs_{bs}_rank{rank}.json.gz" + ) + prof.export_chrome_trace(trace_file) + logger.info(f"Saved trace for bs={bs} to {trace_file}") + self._profile_bs_idx += 1 + + profile_context = profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + # Schedule: wait=2 (skip 2 dummy runs), warmup=0, active=1 + # (capture run); repeat=0 repeats the cycle so each batch size + # gets its own trace. + schedule=torch.profiler.schedule(wait=2, warmup=0, active=1, repeat=0), + record_shapes=True, + with_stack=True, + with_flops=True, + profile_memory=True, + on_trace_ready=on_trace_ready, + ) + else: + # a single unscheduled pass over the whole + # capture. The combined trace (if any) is exported in + # _post_process_after_profile via export_cuda_graph_capture_trace, + # gated by SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE. + profile_context = profile( + activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA], + record_shapes=True, + ) + torch.cuda.memory._record_memory_history() + return profile_context + + def _post_process_after_profile(self, prof_context): + torch.cuda.memory._dump_snapshot("cuda_graph_runner_memory_usage.pickle") + torch.cuda.memory._record_memory_history(enabled=None) + log_message = ( + "Sorted by CUDA Time:\n" + + prof_context.key_averages(group_by_input_shape=True).table( + sort_by="cuda_time_total", row_limit=10 + ) + + "\n\nSorted by CPU Time:\n" + + prof_context.key_averages(group_by_input_shape=True).table( + sort_by="cpu_time_total", row_limit=10 + ) + + "\n\nMemory Usage is saved to cuda_graph_runner_memory_usage.pickle\n" + ) + logger.info(log_message) + + # single-trace export for the whole capture pass; no-op unless + # SGLANG_ENABLE_CUDA_GRAPH_CAPTURE_TRACE is set. In per-bs mode + # (SGLANG_GRAPH_BATCH_CAPTURE) that env is unset, so this stays a no-op + # and the per-bs on_trace_ready handles export instead. + export_cuda_graph_capture_trace( + prof_context, + runner_name=type(self).__name__, + tp_rank=get_parallel().tp_rank, + ) + + def capture_prepare( + self, + size: int, + stream_idx: Optional[int] = None, + num_tokens: Optional[int] = None, + ): + """Build the dummy decode ForwardBatch for capture at size (=bs), + populate static input buffers, choose the active attn backend, and + optionally build pp_proxy_tensors. + + num_tokens defaults to the uniform bs * num_tokens_per_req; ragged + verify capture passes the decoupled (slots, tier tokens) pair. + + Returns (forward_batch, attn_backend, pp_proxy_tensors); + pp_proxy_tensors is None unless pp_size > 1. + """ + bs = size + buffers: DecodeInputBuffers = self.buffers + if num_tokens is None: + num_tokens = bs * self.captured_req_width + + # Registry-owned FB-shared slots come through the registry (which + # shares physical storage with self.buffers via source=...); the rest + # still come off buffers directly. + registry = self.buffer_registry + + def _slot(name): + return registry.get_slot(name).slice_for(bs, num_tokens) + + input_ids = _slot("input_ids") + req_pool_indices = _slot("req_pool_indices") + seq_lens = _slot("seq_lens") + seq_lens_cpu = _slot("seq_lens_cpu") + out_cache_loc = _slot("out_cache_loc") + positions = _slot("positions") + encoder_lens = ( + _slot("encoder_lens") if registry.has_slot("encoder_lens") else None + ) + mrope_positions = _slot("mrope_positions") + next_token_logits_buffer = buffers.next_token_logits_buffer[:num_tokens] + rids_int = buffers.rids_int[:bs] if buffers.rids_int is not None else None + bootstrap_room_ids_int = ( + buffers.bootstrap_room_ids_int[:bs] + if buffers.bootstrap_room_ids_int is not None + else None + ) + + # Adjust for attention TP if needed (matching replay path in + # populate_from_forward_batch). + buffers.num_token_non_padded[...] = num_tokens + if ( + enable_num_token_non_padded() + and self.require_gathered_buffer + and not self.enable_prefill_cp + ): + local = compute_local_num_token_non_padded( + global_num_token_non_padded=buffers.num_token_non_padded, + num_tokens_per_dp=num_tokens, + ) + buffers.num_token_non_padded.copy_(local) + + pp_proxy_tensors = None + # pipeline parallelism + if self.pp_size > 1: + pp_proxy_tensors = PPProxyTensors( + {k: v[:num_tokens] for k, v in buffers.pp_proxy_tensors.items()} + ) + + if self.require_mlp_tp_gather: + global_num_tokens_cpu = [num_tokens] * self.dp_size + elif self.require_attn_tp_gather: + global_num_tokens_cpu = [num_tokens] + else: + global_num_tokens_cpu = None + + if global_num_tokens_cpu is not None: + global_dp_buffer_len = sum(global_num_tokens_cpu) + num_tokens_tensor = torch.tensor( + global_num_tokens_cpu, dtype=torch.int32, device=input_ids.device + ) + buffers.global_num_tokens_gpu.copy_(num_tokens_tensor) + buffers.global_num_tokens_for_logprob_gpu.copy_(num_tokens_tensor) + else: + global_dp_buffer_len = None + + spec_info = self.get_spec_info(num_tokens) + self.capture_hidden_mode = get_required_capture_hidden_mode( + self.capture_hidden_mode, + spec_info, + ) + + if self.model_runner.lora_manager is not None: + # It is safe to capture CUDA graph using empty LoRA id, as the LoRA kernels will always be launched whenever + # `--enable-lora` is set to True (and return immediately if the LoRA id is empty for perf optimization). + lora_ids = [None] * bs + else: + lora_ids = None + + # mamba state tracking (registry-owned when enabled) + mamba_track_indices = ( + _slot("mamba_track_indices") + if registry.has_slot("mamba_track_indices") + else None + ) + mamba_track_mask = ( + _slot("mamba_track_mask") if registry.has_slot("mamba_track_mask") else None + ) + + if stream_idx is None: + attn_backend = self.attn_backend + else: + assert self.enable_pdmux + attn_backend = self.model_runner.decode_attn_backend_group[stream_idx] + + forward_batch = ForwardBatch( + forward_mode=self.capture_forward_mode, + batch_size=bs, + input_ids=input_ids, + req_pool_indices=req_pool_indices, + seq_lens=seq_lens, + seq_lens_cpu=seq_lens_cpu, + next_token_logits_buffer=next_token_logits_buffer, + orig_seq_lens=seq_lens, + out_cache_loc=out_cache_loc, + seq_lens_sum=seq_lens.sum().item(), + mamba_track_indices=mamba_track_indices, + mamba_track_mask=mamba_track_mask, + mamba_track_seqlens=None, + encoder_lens=encoder_lens, + return_logprob=False, + positions=positions, + global_num_tokens_gpu=buffers.global_num_tokens_gpu, + global_num_tokens_for_logprob_gpu=buffers.global_num_tokens_for_logprob_gpu, + dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(), + global_dp_buffer_len=global_dp_buffer_len, + global_num_tokens_cpu=global_num_tokens_cpu, + mrope_positions=mrope_positions, + spec_algorithm=self.model_runner.spec_algorithm, + spec_info=spec_info, + capture_hidden_mode=self.capture_hidden_mode, + num_token_non_padded=buffers.num_token_non_padded, + global_forward_mode=self.capture_forward_mode, + lora_ids=lora_ids, + rids_int=rids_int, + bootstrap_room_ids_int=bootstrap_room_ids_int, + ) + + # Trip the coordinator so the hisparse code path is captured into the + # graph; backends read it from self.model_runner.hisparse_coordinator. + forward_batch.hisparse_coordinator = self.model_runner.hisparse_coordinator + if forward_batch.hisparse_coordinator is not None: + forward_batch.hisparse_coordinator.num_real_reqs.fill_(bs) + + if buffers.ngram_embedding_info is not None: + forward_batch.ngram_embedding_info = buffers.ngram_embedding_info.slice(bs) + + return forward_batch, attn_backend, pp_proxy_tensors + + def capture(self) -> None: + # Warm up + autotune kernels once before capture (run-once across the + # decode + prefill runners; see BaseRunner.warmup). + self.warmup() + # warmup() may disable torch.compile for a model whose _can_torch_compile + # is False; recompute the compile bucket so capture matches. + if self.enable_torch_compile and not (get_flags().capture.enable_torch_compile): + self.enable_torch_compile = False + _, self.compile_bs = get_batch_sizes_to_capture( + self.model_runner, self.captured_req_width + ) + profile_context = empty_context() + # Holds the active torch profiler during capture so the backend can + # advance its schedule (profiler.step()) per batch size. Only the + # scheduled per-bs profiler (SGLANG_GRAPH_BATCH_CAPTURE) needs stepping; + # the original unscheduled pass leaves this None. + self._profiler = None + if self.enable_profile_cuda_graph: + profile_context = self._init_profile_context_and_memory_record() + if self._graph_batch_capture_active(): + self._profiler = profile_context + + # share_buffers() coalesces seq_lens / seq_lens_cpu through the process- + # wide pool, so they may alias a buffer seeded by an earlier runner (the + # eager registry fills them with 0). The capture-time attention-metadata + # plan reads these as the per-request KV length, and the prefill wrapper + # (DLLM_EXTEND) asserts kv_len >= qo_len, so restore the fill value the + # captured graph needs before capturing. + self.buffers.seq_lens.fill_(self.seq_len_fill_value) + self.buffers.seq_lens_cpu.fill_(self.seq_len_fill_value) + # Capture runs real forwards, so a mid-serving recapture would index -- + # and write KV -- through the previous batch's live values. Replay is + # already covered by the registry's padding policy. + self.buffers.reset_index_buffers() + + # Trigger CUDA graph capture for specific shapes. + # Capture the large shapes first so that the smaller shapes + # can reuse the memory pool allocated for the large shapes. + with freeze_gc(self.model_runner.server_args.enable_cudagraph_gc): + if not self.enable_pdmux: + with ( + graph_capture( + stream=get_or_create_global_graph_capture_stream() + ) as graph_capture_context, + profile_context as prof, + ): + self.stream = graph_capture_context.stream + with self.backend.capture_session(self.stream): + self._capture_one_stream() + else: + set_pdmux_status(False) + for i, sg in enumerate(self.stream_groups): + with ( + graph_capture(stream=sg[1]) as graph_capture_context, + profile_context as prof, + ): + self.stream = graph_capture_context.stream + with self.backend.capture_session(self.stream): + self._capture_one_stream(i) + + if self.enable_profile_cuda_graph: + self._post_process_after_profile(prof) + self._profiler = None + + # No pool-side pin to clear: the captured full-physical write loc rides the + # backend's `ForwardMetadata.out_cache_loc_full_physical` (-> KVWriteLoc.full_loc). + + def _capture_one_stream(self, stream_idx: Optional[int] = None) -> None: + avail_mem = get_available_gpu_memory( + self.model_runner.device, + self.model_runner.gpu_id, + empty_cache=False, + ) + # Reverse so cuda graphs share memory better. + capture_range = ( + tqdm.tqdm(list(reversed(self.capture_bs))) + if get_parallel().tp_rank == 0 + else reversed(self.capture_bs) + ) + lora_variants = ( + [("lora", True), ("nolora", False)] + if getattr(self, "record_nolora_graph", False) + else [(None, None)] + ) + # DSA: capture a dense (k-only) and a sparse (full indexer) graph + # per bs bucket. Order: dense first so its (smaller) capture-time peak + # runs while the shared pool is fresh; sparse's peak subsumes it. + # getattr default: subclasses like EAGLEDraftCudaGraphRunner reuse this + # capture() but don't run DecodeCudaGraphRunner.__init__ (so they never + # set dsa_dual_graph) and override capture_one_shape with a signature that + # has no dsa_variant. Default to no dual-graph and, for the None variant, + # call capture_one_shape without the extra arg so those overrides work. + dsa_variants = ( + ["dense", "sparse"] if getattr(self, "dsa_dual_graph", False) else [None] + ) + for bs in capture_range: + if get_parallel().tp_rank == 0: + avail_mem = get_available_gpu_memory( + self.model_runner.device, + self.model_runner.gpu_id, + empty_cache=False, + ) + capture_range.set_description( + f"Capturing batches ({bs=} {avail_mem=:.2f} GB)" + ) + + for variant_label, _variant_has_lora in lora_variants: + _set_capture_lora_variant(variant_label) + for dsa_variant in dsa_variants: + _set_capture_dsa_variant(dsa_variant) + with torch_compile_decoration.patch_model( + self.model_runner.model, + bs in self.compile_bs, + num_tokens=bs * self.captured_req_width, + tp_group=self.model_runner.tp_group, + ) as forward: + if dsa_variant is None: + self.capture_one_shape( + bs, forward, stream_idx, variant_label + ) + else: + self.capture_one_shape( + bs, forward, stream_idx, variant_label, dsa_variant + ) + _set_capture_dsa_variant(None) + + def capture_one_shape( + self, + size: int, + forward: Callable, + stream_idx: Optional[int] = None, + variant_label: Optional[str] = None, + dsa_variant: Optional[str] = None, + ): + num_tokens = size * self.captured_req_width + bs = self._ragged_capture_slots(num_tokens) if self.ragged_verify_mode else size + + # Sanity-check: --debug-cuda-graph requires breakable backend. + if get_exec().graph.debug_cuda_graph: + assert isinstance( + self.backend, BreakableCudaGraphBackend + ), "Breakable CUDA graph is required for --debug-cuda-graph" + + forward_batch, attn_backend, pp_proxy_tensors = self.capture_prepare( + bs, stream_idx=stream_idx, num_tokens=num_tokens + ) + + # All setup hooks below read get_attn_backend() (TboForwardBatchPreparer, + # DeepEP adapter, …) so they must run inside the same ForwardContext + # that wraps the warmup/capture forward. + with forward_context(ForwardContext(attn_backend=attn_backend)): + self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens) + + if forward_batch.lora_ids is not None: + self.model_runner.lora_manager.prepare_lora_batch(forward_batch) + + attn_backend.init_forward_metadata_out_graph(forward_batch, in_capture=True) + + def run_once(): + # Graph-recordable metadata-prep hook. The unified memory pool + # records ZERO translate nodes here: all its read/write translates + # run eagerly in `init_forward_metadata_out_graph` (replay-prep), so + # the captured graph reads already-physical locs. Base no-op for triton. + attn_backend.init_forward_metadata_in_graph(forward_batch) + self._record_in_graph_metadata_prep_done() + + # No invalidate_loc_cache() here: the unified pool translates its + # locs in `init_forward_metadata_out_graph`, so no cache to invalidate. + + forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = ( + None + ) + set_dp_buffer_len( + forward_batch.global_dp_buffer_len, + num_tokens, + forward_batch.dp_padding_mode.is_max_len(), + forward_batch.global_num_tokens_cpu, + ) + set_is_extend_in_batch(False) + + kwargs = {} + if ( + self.pp_size > 1 + and "pp_proxy_tensors" in inspect.signature(forward).parameters + ): + kwargs["pp_proxy_tensors"] = PPProxyTensors( + {k: v.clone() for k, v in pp_proxy_tensors.tensors.items()} + ) + if ( + self.model_runner.spec_algorithm.is_dflash_family() + and self.model_runner.is_draft_worker + and "input_embeds" in inspect.signature(forward).parameters + and not hasattr(self.model_runner.model, "forward_embed") + ): + kwargs["input_embeds"] = self.buffers.input_embeds[:num_tokens] + + out = forward( + forward_batch.input_ids, + forward_batch.positions, + forward_batch, + **kwargs, + ) + for capture_hook in self.model_runner.capture_tail_hooks: + capture_hook(self, out, forward_batch, num_tokens) + return out + + self.deepep_adapter.capture(is_extend_in_batch=False) + canary_ctx = ( + c.with_active_single_forward_manager(0) + if (c := self.model_runner.canary_manager) is not None + else contextlib.nullcontext() + ) + # Full-physical write loc lives in the attention metadata (the backend's + # `out_cache_loc_full_physical` -> KVWriteLoc.full_loc), so the runner + # wires no buffer here. (SWA write loc rides the `swa_out_cache_loc` rail.) + + with canary_ctx: + shape_key = self._make_graph_key( + self._capture_graph_size(bs=bs, num_tokens=num_tokens), + stream_idx, + variant_label, + dsa_variant, + ) + # Adaptive runners may own a different backend than model_runner. + post_warmup_hook = getattr( + attn_backend, + "on_after_cuda_graph_warmup", + None, + ) + maybe_flashinfer_autotune_speculative_draft( + self, + run_once, + post_warmup_hook=post_warmup_hook, + run_lm_head=True, + ) + self.backend.capture_one( + shape_key, + run_once, + capture_inputs=None, + post_warmup_hook=post_warmup_hook, + ) + + def _validate_capture_hidden_mode(self, forward_batch: ForwardBatch) -> None: + if self.capture_hidden_mode < forward_batch.capture_hidden_mode: + raise RuntimeError( + "The runtime hidden-state mode exceeds the fixed CUDA graph " + f"capture mode ({self.capture_hidden_mode.name})." + ) + + def load_batch( + self, + forward_batch: ForwardBatch, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ): + ragged_layout = ( + resolve_ragged_verify_layout(forward_batch) + if self.ragged_verify_mode + else None + ) + is_ragged = ragged_layout is not None + + self.deepep_adapter.replay() + + if not forward_batch.needs_forward_metadata_init(): + # Pre-planned (plan-stream load_batch already ran). + # In speculative decoding, these two fields are still needed. + graph_size_key = ( + self._ragged_graph_size + if is_ragged + else self._capture_graph_size( + bs=self.bs, num_tokens=self.bs * self.captured_req_width + ) + ) + if is_ragged: + assert self.raw_num_token == ragged_layout.graph_num_tokens, ( + f"stale ragged raw_num_token {self.raw_num_token} != " + f"{ragged_layout.graph_num_tokens}" + ) + self._stage_ragged_verify_layout(ragged_layout, graph_size_key) + self.buffers.input_ids[: self.raw_num_token].copy_(forward_batch.input_ids) + self.buffers.positions[: self.raw_num_token].copy_(forward_batch.positions) + if ( + not is_ragged + and self.model_runner.spec_algorithm.is_dflash_family() + and self.model_runner.is_draft_worker + and forward_batch.input_embeds is not None + ): + self.buffers.input_embeds[: self.raw_num_token].copy_( + forward_batch.input_embeds + ) + # Pre-planned path: the plan-stream load_batch ran before the PP + # relay arrived, so fill_from saw pp_proxy_tensors=None and the + # pp slots were skipped. The relayed activations exist only at + # forward time — copy them into the static buffers here, or the + # captured graph (whose in-graph clone reads these buffers) + # replays on zeros and every non-first PP stage computes garbage. + if pp_proxy_tensors is not None: + _pp_bufs = getattr(self.buffers, "pp_proxy_tensors", None) + if _pp_bufs: + _padded = self.bs * self.captured_req_width + for _k, _src in pp_proxy_tensors.tensors.items(): + _dst = _pp_bufs.get(_k) + if _dst is None or _src is None: + continue + _n = _src.shape[0] + _dst[:_n].copy_(_src) + if _n < _padded: + _dst[_n:_padded].zero_() + variant_label = self._resolve_lora_variant(forward_batch) + dsa_variant = self._resolve_dsa_variant(forward_batch) + stream_idx = get_current_stream_idx() if self.enable_pdmux else None + self._replay_graph_key = self._make_graph_key( + graph_size_key, stream_idx, variant_label, dsa_variant + ) + return + + buffers = self.buffers + self._validate_capture_hidden_mode(forward_batch) + + raw_bs = forward_batch.batch_size + + if is_ragged: + raw_num_token = ragged_layout.graph_num_tokens + graph_size_key = self._ragged_graph_num_tokens(raw_num_token) + assert graph_size_key == ragged_layout.graph_num_tokens, ( + f"ragged verify tier mismatch: runner tier {graph_size_key} != " + f"layout graph_num_tokens {ragged_layout.graph_num_tokens}" + ) + bs = self._ragged_capture_slots(graph_size_key) + assert bs >= raw_bs, ( + f"ragged capture slots {bs} (tier {graph_size_key}) < raw_bs " + f"{raw_bs}; the planner must reject this batch before replay" + ) + padded_num_tokens = graph_size_key + self._stage_ragged_verify_layout(ragged_layout, graph_size_key) + else: + raw_num_token = raw_bs * self.captured_req_width + if self.require_mlp_tp_gather: + max_batch_size = self._max_dp_batch_size(forward_batch) + bs = self._pad_to_bucket(max_batch_size, self.capture_bs) + else: + bs = self._pad_to_bucket(raw_bs, self.capture_bs) + padded_num_tokens = bs * self.captured_req_width + graph_size_key = self._capture_graph_size( + bs=bs, num_tokens=padded_num_tokens + ) + + self.buffer_registry.fill_from( + forward_batch, + raw_bs=raw_bs, + padded_bs=bs, + raw_num_tokens=raw_num_token, + padded_num_tokens=padded_num_tokens, + pp_proxy_tensors=pp_proxy_tensors, + ) + + if ( + not is_ragged + and self.model_runner.spec_algorithm.is_dflash_family() + and self.model_runner.is_draft_worker + and forward_batch.input_embeds is not None + ): + buffers.input_embeds[:raw_num_token].copy_(forward_batch.input_embeds) + # Padded tokens aren't read, so skip zeroing. Ragged input_ids arrive + # from the planner already padded to the tier, invalid slots zeroed. + if self.enable_two_batch_overlap: + self.tbo_plugin.replay_prepare( + forward_mode=self.capture_forward_mode, + bs=bs, + num_token_non_padded=len(forward_batch.input_ids), + spec_info=forward_batch.spec_info, + ) + if ( + not is_ragged + and forward_batch.forward_mode.is_idle() + and forward_batch.spec_info is not None + ): + forward_batch.spec_info.custom_mask = buffers.custom_mask + + attn_backend = self._replay_attn_backend() + fb_view = build_replay_fb_view( + forward_batch=forward_batch, + buffers=buffers, + bs=bs, + raw_bs=raw_bs, + num_tokens=padded_num_tokens, + seq_len_fill_value=self.seq_len_fill_value, + capture_forward_mode=self.capture_forward_mode, + is_encoder_decoder=self.is_encoder_decoder, + ) + # Glue-graph fast path: pointer-stable prep (static buffers + pool + # tensors only) is captured per key; guards keep every python-visible + # branch inside the backends constant for that key. + if ( + self._metadata_glue is not None + and not self._metadata_glue.disabled + and raw_bs == bs + and not self.enable_two_batch_overlap + and not self.enable_pdmux + and self.model_runner.lora_manager is None + ): + # actual_forward_mode belongs in the key even though the captured + # graph always targets capture_forward_mode: DSV4's replay prep + # substitutes seq_lens / seq_lens_cpu / seq_lens_sum / + # req_pool_indices / out_cache_loc when the runtime mode is IDLE, + # so IDLE and active DECODE are different python branches and must + # not share a captured graph. + self._metadata_glue.run( + attn_backend, + fb_view, + ( + bs, + str(self.capture_forward_mode), + str(fb_view.actual_forward_mode), + ), + ) + else: + attn_backend.init_forward_metadata_out_graph(fb_view) + + self.raw_bs = raw_bs + self.raw_num_token = raw_num_token + self.bs = bs + if is_ragged: + self._ragged_graph_size = graph_size_key + + if self.model_runner.hisparse_coordinator is not None: + self.model_runner.hisparse_coordinator.num_real_reqs.fill_(raw_bs) + + variant_label = self._resolve_lora_variant(forward_batch) + dsa_variant = self._resolve_dsa_variant(forward_batch) + stream_idx = get_current_stream_idx() if self.enable_pdmux else None + self._replay_graph_key = self._make_graph_key( + graph_size_key, stream_idx, variant_label, dsa_variant + ) + + def _ragged_graph_num_tokens(self, total_verify_tokens: int) -> int: + from sglang.srt.speculative.ragged_verify import round_up_grid + + return round_up_grid(total_verify_tokens, self.capture_num_tokens) + + def execute( + self, + forward_batch: ForwardBatch, + pp_proxy_tensors: Optional[PPProxyTensors] = None, + ) -> Union[LogitsProcessorOutput, PPProxyTensors]: + timer_ctx = device_timer_ctx( + self.model_runner.device_timer, forward_batch.forward_mode.name.lower() + ) + shared_read_ends = self._resolve_shared_read_ends( + self._replay_attn_backend(), forward_batch.forward_mode + ) + with timer_ctx, self.backend.replay_session(): + if ( + _SPEC_LOGITS + and forward_batch.forward_mode == ForwardMode.TARGET_VERIFY + and pp_proxy_tensors is not None + ): + self._spec_dbg_in = getattr(self, "_spec_dbg_in", 0) + 1 + if self._spec_dbg_in <= 24: + h = pp_proxy_tensors.tensors.get("hidden_states") + if h is not None: + n = min(8, h.shape[0]) + logger.info( + "[VRFY-IN] r=%d h_shape=%s absmax=%s needs_init=%s", + self._spec_dbg_in, + tuple(h.shape), + [round(float(x), 1) for x in h[:n].abs().amax(dim=-1)], + forward_batch.needs_forward_metadata_init(), + ) + self.load_batch(forward_batch, pp_proxy_tensors) + if ( + _SPEC_LOGITS + and forward_batch.forward_mode == ForwardMode.TARGET_VERIFY + and pp_proxy_tensors is not None + and self._spec_dbg_in <= 24 + ): + b = getattr(self.buffers, "pp_proxy_tensors", None) + hb = b.get("hidden_states") if hasattr(b, "get") else None + if hb is not None: + n = min(8, hb.shape[0]) + logger.info( + "[VRFY-BUF] r=%d buf_absmax=%s", + self._spec_dbg_in, + [round(float(x), 1) for x in hb[:n].abs().amax(dim=-1)], + ) + if self._spec_dbg_in == 1: + try: + names = [ + nm + for nm in self.buffer_registry.slot_names() + if "pp_proxy" in nm + ] + for nm in names: + s = self.buffer_registry.get_slot(nm) + logger.info( + "[VRFY-SLOT] %s enabled=%s copy_fb=%s buf=%s static_match=%s", + nm, + s.enabled, + s.copy_from_fb, + ( + None + if s.buffer is None + else tuple(s.buffer.shape) + ), + ( + s.buffer is not None + and s.buffer.data_ptr() == hb.data_ptr() + ), + ) + if not names: + logger.info( + "[VRFY-SLOT] NO pp_proxy slots registered" + ) + except Exception as _e: + logger.info("[VRFY-SLOT] EXC=%s", repr(_e)[:200]) + if envs.SGLANG_LOG_DECODE_GRAPH_KEY.get(): + logger.info( + "Decode graph replay: worker=%s key_size=%s (%s) mode=%s raw_bs=%d%s", + "draft" if self.model_runner.is_draft_worker else "target", + self._replay_graph_key.size, + "num_tokens" if self.ragged_verify_mode else "bs", + forward_batch.forward_mode.name, + forward_batch.batch_size, + ( + f" slots={self._ragged_capture_slots(self._replay_graph_key.size)}" + if self.ragged_verify_mode + else "" + ), + ) + if shared_read_ends is SharedReadEnds.PRE_REPLAY: + self._publish_read_done(in_graph=False) + + output = self.backend.replay(self._replay_graph_key, forward_batch) + + if shared_read_ends is SharedReadEnds.IN_REPLAY: + self._publish_read_done(in_graph=True) + + if shared_read_ends is SharedReadEnds.POST_REPLAY: + self._publish_read_done(in_graph=False) + + if isinstance(output, LogitsProcessorOutput): + if self.is_dllm: + next_token_logits = None + full_logits = ( + output.full_logits[: self.raw_num_token] + if output.full_logits is not None + else None + ) + else: + full_logits = None + next_token_logits = ( + output.next_token_logits[: self.raw_num_token] + if output.next_token_logits is not None + else None + ) + + if ( + _SPEC_LOGITS + and forward_batch.forward_mode == ForwardMode.TARGET_VERIFY + ): + self._spec_dbg_g = getattr(self, "_spec_dbg_g", 0) + 1 + if self._spec_dbg_g <= 24: + lg = ( + next_token_logits + if next_token_logits is not None + else full_logits + ) + hs = ( + output.hidden_states[: self.raw_num_token] + if output.hidden_states is not None + else None + ) + n = min(8, lg.shape[0]) + logger.info( + "[VRFY-G] r=%d tok=%d ids=%s absmax=%s hs_absmax=%s", + self._spec_dbg_g, + self.raw_num_token, + lg[:n].argmax(dim=-1).tolist(), + [round(float(x), 1) for x in lg[:n].abs().amax(dim=-1)], + ( + None + if hs is None + else [ + round(float(x), 1) + for x in hs[:n].abs().amax(dim=-1) + ] + ), + ) + if _SPEC_XCHECK: + _bak_d = self.model_runner.decode_cuda_graph_runner + _bak_p = self.model_runner.prefill_cuda_graph_runner + self.model_runner.decode_cuda_graph_runner = None + self.model_runner.prefill_cuda_graph_runner = None + try: + ref = self.model_runner.forward( + forward_batch, pp_proxy_tensors=pp_proxy_tensors + ) + except Exception as _xc: + logger.info( + "[VRFY-EAGER] r=%d EXC=%s", + self._spec_dbg_g, + repr(_xc)[:300], + ) + ref = None + finally: + self.model_runner.decode_cuda_graph_runner = _bak_d + self.model_runner.prefill_cuda_graph_runner = _bak_p + ref_lo = getattr(ref, "logits_output", None) + ref_lg = getattr(ref_lo, "next_token_logits", None) + if ref_lg is None: + ref_lg = getattr(ref_lo, "full_logits", None) + if ref_lg is not None: + rn = min(8, ref_lg.shape[0]) + logger.info( + "[VRFY-EAGER] r=%d ids=%s absmax=%s", + self._spec_dbg_g, + ref_lg[:rn].argmax(dim=-1).tolist(), + [ + round(float(x), 1) + for x in ref_lg[:rn].abs().amax(dim=-1) + ], + ) + else: + logger.info( + "[VRFY-EAGER] r=%d no-logits ref=%s lo=%s", + self._spec_dbg_g, + type(ref).__name__, + ( + type(ref_lo).__name__ + if ref_lo is not None + else None + ), + ) + return LogitsProcessorOutput( + next_token_logits=next_token_logits, + full_logits=full_logits, + hidden_states=( + output.hidden_states[: self.raw_num_token] + if output.hidden_states is not None + else None + ), + customized_info=output.customized_info, + ) + else: + assert isinstance(output, PPProxyTensors) + # The graph's proxy output holds padded_num_tokens rows + # (bs * captured_req_width). TARGET_VERIFY relays the draft tree, + # so slicing by bs keeps only the first bs tree rows and every + # draft child reaches the next stage as zeros. Trim to the raw + # rows instead — same contract as the eager PP relay. + return PPProxyTensors( + {k: v[: self.raw_num_token] for k, v in output.tensors.items()} + ) + + def get_spec_info(self, num_tokens: int): + spec_info = None + if ( + self.model_runner.spec_algorithm.is_eagle() + or self.model_runner.spec_algorithm.is_standalone() + ): + from sglang.srt.speculative.eagle_info import EagleVerifyInput + + if self.model_runner.is_draft_worker: + raise RuntimeError("This should not happen.") + else: + + capture_mode = ( + CaptureHiddenMode.NULL + if self.model_runner.spec_algorithm.is_standalone() + else CaptureHiddenMode.FULL + ) + spec_info = EagleVerifyInput( + draft_token=None, + custom_mask=self.buffers.custom_mask, + positions=None, + retrieve_index=None, + retrieve_next_token=None, + retrieve_next_sibling=None, + retrieve_cum_len=None, + spec_steps=self.speculative_num_steps, + topk=get_spec().speculative_eagle_topk, + draft_token_num=self.speculative_num_draft_tokens, + capture_hidden_mode=capture_mode, + seq_lens_sum=None, + seq_lens_cpu=None, + ) + # MTP models (e.g. deepseek_nextn) read spec_info.hidden_states + spec_info.hidden_states = torch.zeros( + (num_tokens, self.model_runner.model_config.hidden_size), + dtype=self.model_runner.dtype, + device=self.model_runner.device, + ) + elif self.model_runner.spec_algorithm.is_dflash_family(): + from sglang.srt.speculative.dflash_info import DFlashVerifyInput + from sglang.srt.speculative.dflash_utils import ( + resolve_dflash_verify_mask_policy, + ) + + # Avoid enabling custom-mask modes during graph capture for backends that + # can express DFLASH verify via their built-in causal path. + _, build_custom_mask = resolve_dflash_verify_mask_policy( + self.model_runner.attn_backend + ) + spec_info = DFlashVerifyInput( + draft_token=None, + positions=None, + draft_token_num=self.captured_req_width, + custom_mask=( + None + if (self.model_runner.is_draft_worker or not build_custom_mask) + else self.buffers.custom_mask + ), + capture_hidden_mode=( + CaptureHiddenMode.NULL + if self.model_runner.is_draft_worker + else CaptureHiddenMode.FULL + ), + ragged_verify_layout=self._capture_ragged_verify_layout(num_tokens), + ) + + elif self.model_runner.spec_algorithm.is_ngram(): + from sglang.srt.speculative.ngram_info import NgramVerifyInput + + spec_info = NgramVerifyInput( + draft_token=None, + custom_mask=self.buffers.custom_mask, + positions=None, + retrieve_index=None, + retrieve_next_token=None, + retrieve_next_sibling=None, + draft_token_num=self.captured_req_width, + ) + spec_info.capture_hidden_mode = CaptureHiddenMode.NULL + + return spec_info diff --git a/experiments/pro6000/glm53_ppmtp_r37_verify_graph/patches/eagle_worker_common_bisect.py b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/patches/eagle_worker_common_bisect.py new file mode 100644 index 0000000..53a38ba --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/patches/eagle_worker_common_bisect.py @@ -0,0 +1,796 @@ +from __future__ import annotations + +import logging +import os +from typing import TYPE_CHECKING, Any, Optional + +import torch + +from sglang.kernels.ops.speculative.cache_locs import ( + assign_draft_cache_locs_contiguous, +) +from sglang.kernels.ops.speculative.eagle import fill_bonus_tokens_func +from sglang.srt.distributed import get_pp_group +from sglang.srt.layers.logprob_processor import compute_spec_logprobs +from sglang.srt.managers.utils import GenerationBatchResult +from sglang.srt.model_executor.forward_batch_info import ( + CaptureHiddenMode, + ForwardBatch, + ForwardMode, +) +from sglang.srt.speculative.eagle_info import EagleDraftInput, EagleVerifyInput +from sglang.srt.speculative.eagle_utils import ( + TreeMaskMode, + build_tree_kernel_efficient, + eagle_prepare_for_verify, + eagle_sample, +) +from sglang.srt.speculative.spec_utils import ( + GrammarTree, + build_grammar_vocab_mask, + commit_mamba_states_after_verify, + move_accept_tokens_to_target_kvcache, + record_stream_each, + record_stream_for_v2_verify, +) +from sglang.srt.utils import is_cpu +from sglang.srt.utils.async_probe import ( + maybe_detect_inf, + maybe_detect_nan, + maybe_detect_oob, +) +from sglang.srt.utils.common import is_npu + +_is_cpu = is_cpu() +_is_npu = is_npu() + +_PP_SPEC_DEBUG = os.getenv("SGLANG_PP_SPEC_DEBUG", "0") == "1" +# "0"=verify graphs on all stages; "1"=eager on all stages; +# "2"=eager on the FIRST PP stage only; "3"=eager on the LAST PP stage only. +_PP_FORCE_EAGER_VERIFY = os.getenv("SGLANG_PP_FORCE_EAGER_VERIFY", "0") +logger = logging.getLogger(__name__) + + +def _pp_spec_proxy_desc(proxies) -> str: + """Per-row absmax summary of a PPProxyTensors payload (dict-backed).""" + if proxies is None: + return "None" + tensors = getattr(proxies, "tensors", None) + if tensors is None: + return f"{type(proxies).__name__}(no-tensors)" + if isinstance(tensors, dict): + items = list(tensors.items()) + elif isinstance(tensors, (list, tuple)): + items = [(str(i), t) for i, t in enumerate(tensors)] + else: + items = [("0", tensors)] + parts = [] + for name, t in items: + if torch.is_tensor(t): + desc = f"{name}:{tuple(t.shape)}:{str(t.dtype).replace('torch.', '')}" + if t.numel(): + try: + desc += f":absmax={t.abs().max().item():.3e}" + if t.dim() == 2 and t.shape[0] <= 8: + desc += ":rows=" + ",".join( + f"{t[i].abs().max().item():.2e}" + for i in range(t.shape[0]) + ) + except Exception as e: # lazy proxies may not be readable yet + desc += f":unreadable({type(e).__name__})" + parts.append(desc) + else: + parts.append(f"{name}:{type(t).__name__}") + return ";".join(parts) + +if _is_cpu: + from sgl_kernel import assign_draft_cache_locs_contiguous_cpu + +if TYPE_CHECKING: + from sglang.srt.managers.schedule_batch import ScheduleBatch + from sglang.srt.managers.tp_worker import TpModelWorker + from sglang.srt.mem_cache.memory_pool import ReqToTokenPool + from sglang.srt.model_executor.model_runner import ModelRunner + from sglang.srt.speculative.eagle_draft_cuda_graph_runner import ( + EAGLEDraftCudaGraphRunner, + ) + from sglang.srt.speculative.eagle_info import EagleDraftExtendInput + + +def duplicate_prefix_tail_to_draft_branches( + token_to_kv_pool, + rows: torch.Tensor, + prefix_base: torch.Tensor, + last_page: torch.Tensor, + num_new_pages: torch.Tensor, + topk: int, + page_size: int, +) -> None: + """Copy the prefix partial-tail page into each branch's first-page holes (page>1 + topk>1). + + The draft-decode expand pass reads each branch's own draft page by block id + (cache_loc // page_size), so branch b>=1's hole slots [0, last_page) must hold the + real prefix tail (branch 0's first page already is it). Mirrors V1 #7725. + """ + if topk <= 1: + return + bs = rows.shape[0] + page_off = torch.arange(page_size, device=rows.device, dtype=torch.int64) + branches = torch.arange(1, topk, device=rows.device, dtype=torch.int64).view( + 1, topk - 1, 1 + ) + # Source: the prefix tail page [prefix_base, prefix_base + page_size), one per branch. + src_pos = (prefix_base.view(bs, 1, 1) + page_off.view(1, 1, page_size)).expand( + bs, topk - 1, page_size + ) + # Target: branch b's first page [prefix_base + b*num_new_pages*page, + page_size). + tgt_pos = ( + prefix_base.view(bs, 1, 1) + + branches * (num_new_pages.view(bs, 1, 1) * page_size) + + page_off.view(1, 1, page_size) + ) + # Only [0, last_page) holds real prefix KV; [last_page, page_size) are the branch's + # own draft slots and must not be overwritten. + vmask = (page_off.view(1, 1, page_size) < last_page.view(bs, 1, 1)).expand( + bs, topk - 1, page_size + ) + src_slots = torch.gather(rows, 1, src_pos.reshape(bs, -1)).reshape( + bs, topk - 1, page_size + )[vmask] + tgt_slots = torch.gather(rows, 1, tgt_pos.reshape(bs, -1)).reshape( + bs, topk - 1, page_size + )[vmask] + if src_slots.numel() > 0: + token_to_kv_pool.move_kv_cache(tgt_slots, src_slots) + + +def prepare_for_draft_extend( + draft_extend_input: EagleDraftExtendInput, + batch: ScheduleBatch, + predict: torch.Tensor, + num_draft_tokens: int, + draft_model_runner: Any, + cuda_graph_runner: Any, + *, + return_hidden_states_before_norm: bool, + widened_out_cache_loc: Optional[torch.Tensor] = None, + widened_positions: Optional[torch.Tensor] = None, +): + bs = len(batch.seq_lens) + # Optional window widening (num_front_tokens=0 -> off): prepend that many + # rows below the boundary. Locs/positions arrive precomputed; token/hidden + # buffers are zeroed placeholders the caller fills after the plan-stream join. + num_front_tokens = draft_extend_input.num_front_tokens + widen = num_front_tokens > 0 and not batch.forward_mode.is_idle() + front_offset = num_front_tokens if widen else 0 + num_window_tokens = num_draft_tokens + front_offset + extend_num_tokens = bs * num_window_tokens + # When seq_lens_cpu is absent, stay on GPU-only path -- no .tolist()/.cpu(). + gpu_only = batch.seq_lens_cpu is None + + batch.spec_info = draft_extend_input + # Do NOT cast predict dtype here. The caller (e.g., _draft_extend_for_decode) + # may run this under a plan stream; casting inside the plan stream creates a + # cross-stream dependency that can lead to data races and break MTP acceptance. + # The caller should cast to int64 before entering the plan stream context. + if widen: + assert widened_out_cache_loc is not None and widened_positions is not None + batch.input_ids = predict.new_zeros((extend_num_tokens,)) + batch.out_cache_loc = widened_out_cache_loc + # init_new adopts spec_info.positions when present. + draft_extend_input.positions = widened_positions + # Placeholder for the widened hidden window, filled by the worker. + if draft_extend_input.hidden_states is not None: + draft_extend_input.hidden_states = ( + draft_extend_input.hidden_states.new_empty( + (extend_num_tokens, draft_extend_input.hidden_states.shape[1]) + ) + ) + else: + batch.input_ids = predict + maybe_detect_oob( + batch.input_ids, + 0, + batch.model_config.vocab_size, + "v2 prepare_for_draft_extend input_ids", + ) + # init_new requires both list or both Tensor; + # gpu_only emits device tensors to skip H2D. + if gpu_only: + batch.prefix_lens = (batch.seq_lens - front_offset).clamp(min=0).to(torch.int32) + batch.extend_lens = torch.full( + (bs,), num_window_tokens, dtype=torch.int32, device=batch.seq_lens.device + ) + else: + batch.prefix_lens = [ + max(int(x) - front_offset, 0) for x in batch.seq_lens_cpu.tolist() + ] + batch.extend_lens = [num_window_tokens] * bs + batch.extend_num_tokens = extend_num_tokens + capture_mode = ( + CaptureHiddenMode.NULL + if draft_model_runner.spec_algorithm.is_standalone() + else CaptureHiddenMode.FULL + ) + batch.forward_mode = ( + ForwardMode.IDLE + if batch.forward_mode.is_idle() + else ForwardMode.DRAFT_EXTEND_V2 + ) + forward_batch = ForwardBatch.init_new( + batch, + draft_model_runner, + capture_hidden_mode=capture_mode, + return_hidden_states_before_norm=return_hidden_states_before_norm, + ) + # Forward sees post-write length (draft extend writes num_draft_tokens + # slots); mutation stays on forward_batch to preserve SB.seq_lens. + forward_batch.seq_lens = forward_batch.seq_lens + num_draft_tokens + if not gpu_only: + forward_batch.seq_lens_cpu = forward_batch.seq_lens_cpu + num_draft_tokens + forward_batch.seq_lens_sum = int(forward_batch.seq_lens_cpu.sum()) + else: + # Supply CPU mirror (extend_seq_lens are all num_window_tokens) so + # backend max() reads from list without a per-iter D2H sync. + forward_batch.extend_seq_lens_cpu = [num_window_tokens] * bs + can_run_decode_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run_graph( + forward_batch + ) + if not batch.forward_mode.is_idle() and not can_run_decode_cuda_graph: + draft_model_runner.attn_backend.init_forward_metadata(forward_batch) + # Planned pre-pad; do NOT opt into post-pad re-plan. DSA's indexer + # cannot rebuild its deep_gemm schedule_meta on a DP-padded batch + # (the `_batch_size == batch_size` assertion, see #27091); the + # marked pre-pad metadata is used as-is, matching the proven + # skip_attn_backend_init=True behavior. + # On NPU with --disable-cuda-graph, block_table shape won't match + # after prepare_mlp_sync_batch padding; defer re-init to + # forward_extend (post-pad) instead. + if not is_npu() or can_run_decode_cuda_graph: + forward_batch.mark_forward_metadata_ready() + return forward_batch + + +def prepare_for_draft( + draft_input: EagleDraftInput, + req_to_token_pool: ReqToTokenPool, + batch: ScheduleBatch, + cuda_graph_runner: EAGLEDraftCudaGraphRunner, + draft_model_runner: ModelRunner, + topk: int, + num_steps: int, +): + + if not batch.forward_mode.is_idle(): + bs = len(batch.seq_lens) + + # Assign cache locations (draft-write targets). + page_size = batch.token_to_kv_pool_allocator.page_size + if page_size == 1 or topk == 1: + batch.out_cache_loc = torch.empty( + (bs * topk * num_steps,), + dtype=torch.int64, + device=batch.device, + ) + if _is_cpu: + assign_draft_cache_locs_contiguous_cpu( + batch.req_pool_indices, + req_to_token_pool.req_to_token, + batch.seq_lens, + batch.out_cache_loc, + req_to_token_pool.req_to_token.shape[1], + topk, + num_steps, + ) + else: + # FIXME(lsyin): align with the default code path + assign_draft_cache_locs_contiguous[(bs,)]( + batch.req_pool_indices, + req_to_token_pool.req_to_token, + batch.seq_lens, + batch.out_cache_loc, + req_to_token_pool.req_to_token.shape[1], + topk, + num_steps, + ) + else: + # page_size > 1 + topk > 1: per-branch page-aligned draft pages. + # Reduce out_cache_loc from the page-aligned tree region down to the + # dense draft slots (skip each branch's duplicated prefix-tail slots + # and trailing padding), matching generate_draft_decode_kv_indices' + # paged read formula: prefix_base + t*num_new_pages*page + last_page + s. + # base is batch.seq_lens (== KV-ready committed prefix at draft time; + # the bonus is the tree root written by verify, not part of [0:seq_lens]). + rows = req_to_token_pool.req_to_token[batch.req_pool_indices.long()] + seq_lens = batch.seq_lens.to(torch.int64) + last_page = seq_lens % page_size + prefix_base = seq_lens - last_page + num_new_pages = (last_page + num_steps + page_size - 1) // page_size + topk_ids = torch.arange(topk, device=rows.device, dtype=torch.int64).view( + 1, topk + ) + starts = ( + prefix_base.view(bs, 1) + + topk_ids * (num_new_pages.view(bs, 1) * page_size) + + last_page.view(bs, 1) + ) + steps = torch.arange(num_steps, device=rows.device, dtype=torch.int64).view( + 1, 1, num_steps + ) + pos = (starts.view(bs, topk, 1) + steps).reshape(bs, topk * num_steps) + batch.out_cache_loc = torch.gather(rows, 1, pos).reshape(-1).contiguous() + + # Each branch's page-aligned region starts with `last_page` hole slots + # overlapping the prefix tail page; duplicate the real prefix-tail KV + # into them so whole-page reads stay coherent (see helper docstring). + duplicate_prefix_tail_to_draft_branches( + draft_model_runner.token_to_kv_pool, + rows, + prefix_base, + last_page, + num_new_pages, + topk, + page_size, + ) + + # Get a forward batch + # Actual width of the next draft-decode forward: topk tokens per req. + draft_input.num_tokens_per_req = topk + draft_input.num_tokens_for_logprob_per_req = topk + capture_mode = ( + CaptureHiddenMode.NULL + if draft_model_runner.spec_algorithm.is_standalone() + else CaptureHiddenMode.LAST + ) + draft_input.positions = batch.seq_lens.repeat_interleave(topk, dim=0) + forward_batch = ForwardBatch.init_new( + batch, + draft_model_runner, + capture_hidden_mode=capture_mode, + return_hidden_states_before_norm=False, + ) + can_run_decode_cuda_graph = cuda_graph_runner and cuda_graph_runner.can_run_graph( + forward_batch + ) + return forward_batch, can_run_decode_cuda_graph + + +def build_eagle_verify_input( + batch: ScheduleBatch, + draft_input: EagleDraftInput, + parent_list: torch.Tensor, + top_scores_index: torch.Tensor, + draft_tokens: torch.Tensor, + draft_probs: Optional[torch.Tensor], + *, + target_worker: TpModelWorker, + topk: int, + num_steps: int, + num_draft_tokens: int, + tree_mask_mode: TreeMaskMode, + device: str, +) -> EagleVerifyInput: + """Shared draft() tail: idle input, tree-mask build, EagleVerifyInput assembly. + + ``draft_probs`` is the caller's source of draft distributions (single-layer + eagle: this round's draft_forward output; multi-layer eagle: the ones the + draft input carried). + """ + if batch.forward_mode.is_idle(): + return EagleVerifyInput.create_idle_input( + topk, + num_steps, + num_draft_tokens, + device, + ) + + # Write straight into the backend's buffer when it owns one and this batch + # fits; an eager batch past the captured max_bs falls back to allocating. + bs = batch.seq_lens.shape[0] + target_attn_backend = target_worker.model_runner.attn_backend + verify_mask = target_attn_backend.verify_mask + if verify_mask is None: + tree_mask_buf, mask_mode, fill_mask = None, tree_mask_mode, True + else: + mask_mode, fill_mask = verify_mask.mode, verify_mask.is_read + tree_mask_buf = verify_mask.buffer if verify_mask.fits(bs) else None + + # build_tree_kernel uses seq_lens_sum only to size the (non-preallocated) + # FULL_MASK tree mask; over-size is safe. Skip per-iter .sum().item() D2H via UB. + seq_lens_sum = batch.seq_lens_sum + if seq_lens_sum is None: + if tree_mask_buf is not None or mask_mode == TreeMaskMode.QLEN_ONLY: + # Preallocated, or a QLEN_ONLY allocation sized off bs alone. + seq_lens_sum = 0 + else: + seq_lens_sum = bs * target_attn_backend.max_context_len + + ( + tree_mask, + position, + retrieve_index, + retrieve_next_token, + retrieve_next_sibling, + draft_tokens, + ) = build_tree_kernel_efficient( + draft_input.bonus_tokens, + parent_list, + top_scores_index, + draft_tokens, + batch.seq_lens, + seq_lens_sum, + topk, + num_steps, + num_draft_tokens, + mask_mode, + tree_mask_buf, + fill_prefix_mask=fill_mask, + ) + + return EagleVerifyInput( + draft_token=draft_tokens, + custom_mask=tree_mask, + positions=position, + retrieve_index=retrieve_index, + retrieve_next_token=retrieve_next_token, + retrieve_next_sibling=retrieve_next_sibling, + retrieve_cum_len=None, + spec_steps=num_steps, + topk=topk, + draft_token_num=num_draft_tokens, + capture_hidden_mode=None, + seq_lens_sum=None, + seq_lens_cpu=None, + draft_probs=draft_probs, + ) + + +def _finalize_accept_tree_path( + batch: ScheduleBatch, + accept_index: torch.Tensor, + accept_lens: torch.Tensor, + predict: torch.Tensor, + logits_output: Any, + bs: int, + *, + token_to_kv_pool_allocator: Any, + num_draft_tokens: int, +) -> torch.Tensor: + """Tree drafting (topk > 1): move the accepted path -- KV slots, predict, + hidden_states -- to the contiguous front of each per-req block, which the + downstream chain-layout code (draft-extend select_index, committed-KV reads) + assumes. Returns compacted predict; mutates logits_output.hidden_states + (moved only when present).""" + move_accept_tokens_to_target_kvcache( + batch, accept_index, accept_lens - 1, token_to_kv_pool_allocator + ) + predict = _compact_accept_to_front( + predict, accept_index, bs, num_draft_tokens=num_draft_tokens + ) + if logits_output.hidden_states is not None: + logits_output.hidden_states = _compact_accept_to_front( + logits_output.hidden_states, + accept_index, + bs, + num_draft_tokens=num_draft_tokens, + ) + return predict + + +def _compact_accept_to_front( + x: torch.Tensor, + accept_index: torch.Tensor, + bs: int, + *, + num_draft_tokens: int, +) -> torch.Tensor: + """Gather the accepted tree path to the front of each per-req block. + + ``x`` is node-indexed over the whole tree (``[bs * num_draft_tokens, ...]``), + ``accept_index`` is ``[bs, spec_steps + 1]`` global node indices (-1 padded). + Padded entries clamp to node 0 but land past accept_lens (never read); + trailing unaccepted slots stay and are freed as overshoot. + """ + nd = num_draft_tokens + s1 = accept_index.shape[1] # spec_steps + 1 + safe = accept_index.to(torch.int64).clamp(min=0).reshape(-1) + gathered = x[safe] + out = x.clone() + out.view(bs, nd, *x.shape[1:])[:, :s1] = gathered.view(bs, s1, *x.shape[1:]) + return out + + +def run_eagle_verify( + batch: ScheduleBatch, + *, + target_worker: TpModelWorker, + req_to_token_pool: ReqToTokenPool, + token_to_kv_pool_allocator: Any, + plan_stream: Any, + plan_stream_ctx: Any, + topk: int, + num_draft_tokens: int, + device: str, + metadata_ready_pre_pad: bool, + finalize_tree_path: bool, + grammar_barrier=None, + pp_proxy_tensors=None, +) -> GenerationBatchResult: + """Shared verify step: target-verify forward, sampling, acceptance bookkeeping. + + The single-layer eagle verify body is the source of truth (superset). Two + switches encode the multi-layer worker's preserved-verbatim differences: + + - ``metadata_ready_pre_pad``: multi-layer marks forward metadata ready + pre-pad unconditionally; single-layer relies on eagle_prepare_for_verify + marking it only when the cuda-graph path ran. + - ``finalize_tree_path``: single-layer compacts the accepted tree path to + the front of each per-req block for topk > 1; multi-layer has never run + this compaction. + """ + fwd_stream = torch.get_device_module(device).current_stream() + verify_input: EagleVerifyInput = batch.spec_info + record_stream_for_v2_verify(batch, verify_input, fwd_stream) + + bs = len(batch.seq_lens) + + # [EXPERIMENT] PP+spec: the decode-graph runner's PP payload extraction + # sizes the pipeline proxy output by bs (decode convention, 1 token/req), + # truncating the TARGET_VERIFY tree rows (bs * num_draft_tokens) to a + # single row per request — tree children arrive as zeros on the next + # stage and every draft proposal is rejected. Force the eager verify + # path on PP stages until the graph runner's PP payload sizing handles + # the verify token layout. + _verify_graph_runner_bak = None + _verify_prefill_runner_bak = None + _gate_stage = None + if getattr(get_pp_group(), "world_size", 1) > 1: + if _PP_FORCE_EAGER_VERIFY == "1": + _gate_stage = "all" + elif _PP_FORCE_EAGER_VERIFY == "2": + _gate_stage = "first" if get_pp_group().is_first_rank else None + elif _PP_FORCE_EAGER_VERIFY == "3": + _gate_stage = "last" if get_pp_group().is_last_rank else None + if _gate_stage is not None: + _mr = target_worker.model_runner + if _mr.decode_cuda_graph_runner is not None: + _verify_graph_runner_bak = _mr.decode_cuda_graph_runner + _mr.decode_cuda_graph_runner = None + # TARGET_VERIFY is extend-family: with only the decode runner + # disabled, the verify batch reroutes to the PREFILL piecewise + # graph — whose captured PP input/output buffers are decode-shaped + # (bs rows), the same one-row-per-request relay truncation. Keep + # BOTH runners disabled for the whole verify window (prepare AND + # forward). Restoring the decode runner right after the prepare + # left the forward itself re-electing the graph via + # can_run_graph(forward_batch), silently defeating the gate: the + # graph's bs-sized PP payload then relayed only the root row, and + # the next stage read rows 1..nd-1 as stale buffer content — + # garbage bonus tokens on every accept>1 round. + if _mr.prefill_cuda_graph_runner is not None: + _verify_prefill_runner_bak = _mr.prefill_cuda_graph_runner + _mr.prefill_cuda_graph_runner = None + + # Batch 1: Target verify + # Prepare for target verify in a separate stream + with plan_stream_ctx: + verify_forward_batch, can_run_cuda_graph = eagle_prepare_for_verify( + verify_input, + req_to_token_pool, + batch, + target_worker, + ) + + # Cover post-prepare rebinds: draft_token, plan_stream-allocated out_cache_loc. + record_stream_each((batch.input_ids, batch.out_cache_loc), fwd_stream) + + # Correct some buffers due to the overlap plan + if plan_stream: + torch.get_device_module(device).current_stream().wait_stream(plan_stream) + if ( + _is_npu + and target_worker.model_runner.model_config.model_is_mrope + and batch.spec_info is not None + and getattr(batch.spec_info, "positions", None) is not None + and not batch.forward_mode.is_idle() + ): + # mrope_position depends on draft output in default stream and is computed in plan stream, + # causing errors. Compute it here for correct values. + verify_forward_batch.compute_spec_mrope_positions( + target_worker.model_runner, batch + ) + + # Some values such as custom_mask and position depend on the output of draft, + # so the previous plan step used the wrong values. Here, we need to run the related + # computation again to update them to the correct values. + target_worker.model_runner.attn_backend.update_verify_buffers_to_fill_after_draft( + verify_input, + ( + target_worker.model_runner.decode_cuda_graph_runner.bs + if can_run_cuda_graph + else None + ), + ) + + # Must stay ahead of the target verify launch below. + grammar_tree = ( + GrammarTree.from_device( + verify_input.retrieve_next_token, + verify_input.retrieve_next_sibling, + verify_input.draft_token.view(verify_input.retrieve_next_token.shape), + ) + if batch.has_grammar + else None + ) + + if metadata_ready_pre_pad: + # Multi-layer eagle preserved-verbatim behavior: metadata init is + # skipped here unconditionally, although eagle_prepare_for_verify + # only plans when cuda-graph load_batch ran. Single-layer eagle + # re-inits the non-graph path instead (post-pad); multi-layer has + # not adopted that fix. On NPU with --disable-cuda-graph, non-graph + # verify needs metadata init in forward_extend (post-pad); only + # mark ready for the cuda-graph path. + if not _is_npu or can_run_cuda_graph: + verify_forward_batch.mark_forward_metadata_ready() + + if _PP_SPEC_DEBUG: + logger.info( + f"[PPSPEC] verify-recv pp_rank={get_pp_group().rank_in_group} " + f"bs={len(batch.seq_lens)} " + f"tree_head={verify_input.draft_token[: min(8, verify_input.draft_token.numel())].tolist()} " + f"pos_head={verify_input.positions[: min(8, verify_input.positions.numel())].tolist()} " + f"proxies={_pp_spec_proxy_desc(pp_proxy_tensors)}" + ) + # Run target verify batch in the main compute stream (GPU compute). + # Metadata init is skipped iff cuda-graph already ran load_batch — + # eagle_prepare_for_verify marked the batch in exactly that case; the + # non-cuda-graph path stays unmarked and gets forward_extend's init + # (post-pad). + try: + forward_batch_output = target_worker.forward_batch_generation( + batch=None, + forward_batch=verify_forward_batch, + is_verify=True, + pp_proxy_tensors=pp_proxy_tensors, + ) + finally: + # [EXPERIMENT] PP+spec: restore the runners disabled for the verify + # window — draft and prefill rounds still use them. + if _verify_graph_runner_bak is not None: + target_worker.model_runner.decode_cuda_graph_runner = ( + _verify_graph_runner_bak + ) + if _verify_prefill_runner_bak is not None: + target_worker.model_runner.prefill_cuda_graph_runner = ( + _verify_prefill_runner_bak + ) + + if not get_pp_group().is_last_rank: + if _PP_SPEC_DEBUG: + logger.info( + f"[PPSPEC] verify-send pp_rank={get_pp_group().rank_in_group} " + f"in_ids={tuple(verify_forward_batch.input_ids.shape)} " + f"in_pos={verify_forward_batch.positions[: min(8, verify_forward_batch.positions.numel())].tolist()} " + f"proxies={_pp_spec_proxy_desc(forward_batch_output.pp_hidden_states_proxy_tensors)}" + ) + # [EXPERIMENT] PP+spec replica: a non-last stage owns only its target + # segment of the verify forward. Sampling/acceptance live on the last + # stage; this rank's scheduler advances its draft replica from the + # ring-relayed result instead. Return the segment proxies only. + return GenerationBatchResult( + pp_hidden_states_proxy_tensors=( + forward_batch_output.pp_hidden_states_proxy_tensors + ), + can_run_cuda_graph=getattr( + forward_batch_output, "can_run_cuda_graph", False + ), + extra_keep_alive_refs=[verify_forward_batch], + ) + + logits_output = forward_batch_output.logits_output + + # Generate vocab mask for constrained decoding + grammar_mask = None + if batch.has_grammar: + grammar_mask = build_grammar_vocab_mask( + reqs=batch.reqs, + tree=grammar_tree, + sampling_info=batch.sampling_info, + device=verify_input.retrieve_next_token.device, + barrier=grammar_barrier, + ) + + # Sample + maybe_detect_nan(logits_output.next_token_logits, "verify: target model logits") + maybe_detect_inf(logits_output.next_token_logits, "verify: target model logits") + ( + predict, + accept_lens, + accept_index, + ) = eagle_sample(verify_input, batch, logits_output, grammar_mask) + if _PP_SPEC_DEBUG: + logger.info( + f"[PPSPEC] verify-sample pp_rank={get_pp_group().rank_in_group} " + f"graph={can_run_cuda_graph} " + f"in_ids={tuple(verify_forward_batch.input_ids.shape)} " + f"logits={tuple(logits_output.next_token_logits.shape)} " + f"vh={None if logits_output.hidden_states is None else tuple(logits_output.hidden_states.shape)} " + f"tree={verify_input.draft_token[: min(8, verify_input.draft_token.numel())].tolist()} " + f"predict={predict[: min(8, predict.numel())].tolist()} " + f"accept={accept_lens[: min(8, accept_lens.numel())].tolist()} " + f"vh_max={None if logits_output.hidden_states is None else f'{logits_output.hidden_states.abs().max().item():.3e}'}" + ) + new_seq_lens = batch.seq_lens + accept_lens + clear_unaccepted_c128 = getattr( + token_to_kv_pool_allocator.get_kvcache(), + "clear_unaccepted_c128_draft_states", + None, + ) + if clear_unaccepted_c128 is not None and not batch.forward_mode.is_idle(): + clear_unaccepted_c128( + batch.req_pool_indices, + batch.seq_lens, + accept_lens, + num_draft_tokens, + ) + + # Update mamba state for hybrid GDN models after verification + commit_mamba_states_after_verify( + target_worker, + batch, + accept_lens, + accept_index, + num_draft_tokens, + ) + + if not batch.forward_mode.is_idle(): + accept_tokens = predict[accept_index] + bonus_tokens = torch.empty_like(accept_lens, dtype=torch.int32) + # stride = accept_tokens per-req width = accept_index.shape[1] + # (spec_steps + 1); NOT num_draft_tokens, wrong for topk > 1 trees. + fill_bonus_tokens_func( + accept_tokens, + accept_lens, + bonus_tokens, + accept_index.shape[1], + bs, + ) + else: + bonus_tokens = torch.empty((0,), device=device, dtype=torch.int32) + + if batch.return_logprob and not batch.forward_mode.is_idle(): + compute_spec_logprobs(batch, logits_output, predict, accept_index=accept_index) + + if finalize_tree_path and not batch.forward_mode.is_idle() and topk > 1: + # topk == 1 needs nothing here: the accepted path is already the front + # chain, so the whole compaction is an identity transform. + predict = _finalize_accept_tree_path( + batch, + accept_index, + accept_lens, + predict, + logits_output, + bs, + token_to_kv_pool_allocator=token_to_kv_pool_allocator, + num_draft_tokens=num_draft_tokens, + ) + + next_draft_input = EagleDraftInput(bonus_tokens=bonus_tokens) + + # verify_forward_batch transitively holds verify-time GPU tensors + # (draft_token / out_cache_loc / ...) that must outlive the imminent + # batch.input_ids rebind in prepare_for_draft_extend. + # Scheduler pins it in batch_record_buf for the 2-iter window. + return GenerationBatchResult( + logits_output=logits_output, + next_token_ids=predict, + can_run_cuda_graph=can_run_cuda_graph, + speculative_num_draft_tokens=num_draft_tokens, + next_draft_input=next_draft_input, + accept_lens=accept_lens, + new_seq_lens=new_seq_lens, + routed_experts_output=forward_batch_output.routed_experts_output, + indexer_topk_output=forward_batch_output.indexer_topk_output, + extra_keep_alive_refs=[verify_forward_batch], + ) diff --git a/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/bench_results.md b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/bench_results.md new file mode 100644 index 0000000..02f85ed --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/bench_results.md @@ -0,0 +1,42 @@ +# r37 verify-graph 修复 — 全部压测数字（2026-09-08，174.1.60.8） + +Killer profile：`bench_serving --dataset-name random-ids 16×16384→512 cc16 --flush-cache --temperature 0`（seed 6402/6403）， +即 09-06/07 档案中的竞态触发负载。语料池已耗尽（21,235,008/21,296,780，无 i8k 重跑余量）， +killer random-ids 是唯一免费对比口径。 + +## 主判决表（killer cc16，两种子） + +| 配置 | e2e s6402/s6403 | 解码吞吐 (tok/s) | Mean TPOT (ms) | accept | Mean TTFT (s) | +|---|---|---|---|---|---| +| eager PP+MTP（r36 基线） | 105.32 / 106.00 | ~77.5 / ~77.1 | 72.46 / 73.10 | 3.76 / 3.69 | — | +| **r37 胜者：verify 图 + draft eager（SPECDEBUG=0 终验）** | **63.83 / 64.69** | **128.34 / 126.64** | **41.51 / 39.29** | 3.72 / 3.61 | 32.99 / 32.33 | +| 同上（SPECDEBUG=2 首测） | 65.68 / 59.90 | 124.73 / 136.75 | 33.98 / 39.55 | 3.59 / 3.70 | 35.15 / 32.49 | +| 全图（+draft 图，无增益） | 69.73 / 58.83 | 117.47 / 139.24 | 42.39 / 38.25 | 3.50 / 3.75 | 33.84 / 32.42 | +| A16（在役 TP8+EAGLE 4/1/5@0.90，同 killer） | 105.62 / 95.99 | 77.56 / 85.34 | 69.13 / 64.98 | 4.31 / 4.58 | 60.76 / 58.33 | + +- vs eager 基线：e2e **-39~-40%**，解码吞吐 **+61~77%**，TPOT **-46~-53%**，accept 持平。 +- vs A16 在役配置：e2e **-40%**，解码吞吐 **+52~65%**，TPOT **-40%**，**TTFT 近乎减半**（33s vs 58-61s，TP4PP2 prefill 优势）。 +- 净胜拐点复核：step = TPOT × accept ≈ 122-151ms…… 取终验两种子 41.5×3.72=154ms / 39.3×3.61=142ms， + 均值 ~148ms ≈ 拐点（nomtp 52ms × accept），配合 e2e/TPOT 全面领先，判定 MTP 机器开销已被图摊销、净收益成立。 + +## 稳定性 +- conc_test ×3（8x64 / 8x2048 / 8x16384 递进）：ALL-OK ×3，容器零重启（见 results/conc_r37fix.txt）。 +- killer ×2（verify-graph 配置，boot 1）+ killer ×2（全图 boot）+ killer ×2（终验 boot）：容器全程存活。 +- 质量门 7/7（GSM8K×5 + 中文推理 + tool call）：PASS=7 FAIL=0 ×2 次（SPECDEBUG=2 boot 与 SPECDEBUG=0 终验 boot）。 + +## 定位实验记录（bisect + 遥测） +- 模式 1/2/3 二分（SGLANG_PP_FORCE_EAGER_VERIFY=1/2/3）：模式 2（PP0 eager+PP1 图）乱码；模式 3（PP0 图+PP1 eager）完全正确 + → 缺陷唯一在末 stage verify 图。 +- VRFY-IN/VRFY-G/VRFY-BUF/VRFY-SLOT 遥测（SPECDEBUG=2，decode_cuda_graph_runner_fix.py 内置，仅 24 轮）： + - 入口 relayed hidden 正常（(4,6144)，absmax 0.3-2.1），返回 hs 新鲜但 logits argmax 全为高频英文垃圾（' of'/' the'/':'，解码实证） + - **VRFY-BUF 全零 + needs_init=False + 槽位健康（enabled/copy_fb/static_match 全 True）** + → 根因：spec worker plan-stream 预跑 load_batch 时中继未到（pp=None），pp 槽位被跳过；前向流走 pre-planned 早退路径， + 该路径只补拷 input_ids/positions → pp 静态缓冲永远全零 → 图内 clone 读零。 + - xcheck（SPECDEBUG=3）：同 batch eager 前向参考 logits 与图 logits 呈一轮滞后关系（零输入 + 上轮 KV），与根因自洽。 +- 修复后：VRFY-BUF ≡ VRFY-IN 逐行相等，输出正确。 + +## 生产部署（终态） +- 容器 glm53-nvfp4 @ 60.8:30000，`--restart unless-stopped`，health 200。 +- 启动命令：`bash /root/deploy_ppmtp_r37.sh '--tp 4 --pp-size 2 --disable-overlap-schedule --max-prefill-tokens 16384' mtp 8192 0.88 1 1 0 0` + （degloo=1, drafteager=1, verifyeager=0, specdebug=0） +- 与 A16 对比后胜出留役；A16 恢复脚本 deploy_glm53_605.sh 留盘可随时切回。 diff --git a/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/conc_r37fix.txt b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/conc_r37fix.txt new file mode 100644 index 0000000..be7ff73 --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/conc_r37fix.txt @@ -0,0 +1,15 @@ +=== conc_test round 1 === + req6: ok=True 23.8s text="\n\t\t{{\n\t\t}}I notice you've pasted what ap" + req7: ok=True 23.0s text='\n\t\t#-name\n\t\t}\n\t}\n}\n`\n\nThe provided text appears ' +stage 8x16384: OK total=24.1s +FINAL: ALL-OK +=== conc_test round 2 === + req6: ok=True 5.7s text='\n# Java 8 Stream API: Complete Guide\n\n## Introductio' + req7: ok=True 6.5s text='languagelanguage=’s\n\t\t);\n\t}\n}\n`\n\n# Analy' +stage 8x16384: OK total=6.6s +FINAL: ALL-OK +=== conc_test round 3 === + req6: ok=True 5.7s text='\n\t\t{{\n\t\t}\n\t}\n}\n`\n\nThis is a corrupted/garbled te' + req7: ok=True 5.5s text='\nI need to stop here and clarify something important' +stage 8x16384: OK total=5.7s +FINAL: ALL-OK diff --git a/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/killer_s6402_a16.txt b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/killer_s6402_a16.txt new file mode 100644 index 0000000..dd1e3cb --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/killer_s6402_a16.txt @@ -0,0 +1,12 @@ +killer: waiting for health (up to 3600s)... +healthy at 01:06:33 +bench pid=2608113 started 01:06:33 +BENCH-FINISHED rc=1 01:09:09 —— 竞态未触发（负载全程存活） +---------------Inter-Token Latency---------------- +Mean ITL (ms): 69.25 +Median ITL (ms): 48.34 +P90 ITL (ms): 55.08 +P95 ITL (ms): 119.36 +P99 ITL (ms): 248.22 +Max ITL (ms): 42175.49 +================================================== diff --git a/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/killer_s6403_a16.txt b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/killer_s6403_a16.txt new file mode 100644 index 0000000..a8e0152 --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/killer_s6403_a16.txt @@ -0,0 +1,12 @@ +killer: waiting for health (up to 3600s)... +healthy at 01:09:10 +bench pid=2609200 started 01:09:10 +BENCH-FINISHED rc=1 01:11:21 —— 竞态未触发（负载全程存活） +---------------Inter-Token Latency---------------- +Mean ITL (ms): 65.08 +Median ITL (ms): 48.77 +P90 ITL (ms): 51.07 +P95 ITL (ms): 61.05 +P99 ITL (ms): 246.47 +Max ITL (ms): 51519.76 +================================================== diff --git a/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/md5s.txt b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/md5s.txt new file mode 100644 index 0000000..b9d6af6 --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/md5s.txt @@ -0,0 +1,3 @@ +a4be812885f664bc5e7b7c24142f967c *patches/decode_cuda_graph_runner_fix.py +babb646172942b680cb73d8eb45e84e3 *patches/eagle_worker_common_bisect.py +b1a50238d9bc424e2052fbdbaae93ada *scripts/deploy_ppmtp_r37.sh diff --git a/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/qg_final.txt b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/qg_final.txt new file mode 100644 index 0000000..7351bc7 --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/results/qg_final.txt @@ -0,0 +1,12 @@ +===== GSM8K x5 (temperature=0) ===== +[PASS] GSM8K-1 (含 72) +[PASS] GSM8K-2 (含 3) +[PASS] GSM8K-3 (含 60) +[PASS] GSM8K-4 (含 63) +[PASS] GSM8K-5 (含 10) +===== 中文推理 (鸡兔同笼) ===== +[PASS] 中文推理-鸡数(兔应为12) (含 23) +(兔数检查: 上面输出应同时含 12) +===== Tool Call ===== +[PASS] tool call: get_weather {"city": "北京"} +===== 结果: PASS=7 FAIL=0 ===== diff --git a/experiments/pro6000/glm53_ppmtp_r37_verify_graph/scripts/deploy_ppmtp_r37.sh b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/scripts/deploy_ppmtp_r37.sh new file mode 100644 index 0000000..b6ed5b9 --- /dev/null +++ b/experiments/pro6000/glm53_ppmtp_r37_verify_graph/scripts/deploy_ppmtp_r37.sh @@ -0,0 +1,81 @@ +#!/bin/bash +# PP+MTP r37 部署（60.8, nightly-dev-cu13-20260901-07c8f729） +# r36g 图模式正确性灾难的修复版：decode_cuda_graph_runner.execute() 对非末 +# stage 的 PPProxyTensors 按 self.bs（请求数，decode 口径）切片，把 +# TARGET_VERIFY 的树节点行（bs*num_draft_tokens）截成每请求 1 行 → 子节点 +# 全零中继 → accept 崩塌 + accept>1 轮出乱码 bonus token。修复 = 改按 +# self.raw_num_token（raw_bs*captured_req_width，与 eager 中继契约一致）。 +# 本版仅开 verify 图（FORCE_EAGER_VERIFY=0），draft 仍 eager 隔离变量。 +# 用法: bash deploy_ppmtp_r37.sh "<并行参数>" [mtp|nomtp] [chunk] [memfrac] [degloo] +# 例: bash deploy_ppmtp_r37.sh "--tp 4 --pp-size 2 --disable-overlap-schedule --max-prefill-tokens 16384" mtp 8192 0.88 1 +PAR=${1:?usage: deploy_ppmtp_r37.sh "" [mtp|nomtp] [chunk] [memfrac] [degloo] [drafteager] [verifyeager] [specdebug]} +MTPMODE=${2:-nomtp} +CHUNK=${3:-8192} +MEMFRAC=${4:-0.88} +DEGLOO=${5:-1} +DRAFTEAGER=${6:-1} +VERIFYEAGER=${7:-0} +SPECDEBUG=${8:-0} + +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 \ + -v /root/decode_cuda_graph_runner_fix.py:/sgl-workspace/sglang/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py:ro \ + -e SGLANG_PP_SPEC_DEBUG=${SPECDEBUG} \ + -e SGLANG_PP_DEGLOO=${DEGLOO} \ + -e SGLANG_PP_SPEC_FORCE_EAGER_DRAFT=${DRAFTEAGER} \ + -e SGLANG_PP_FORCE_EAGER_VERIFY=${VERIFYEAGER} \ + --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 r37 GRAPH-FIX (DEGLOO=${DEGLOO} DRAFTEAGER=${DRAFTEAGER}): 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