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
This commit is contained in:
yy-fighting
parent
ffda226f5a
commit
a3b8f1b2ab
@ -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 双场景报告口径）
|
||||
|
||||
|
||||
66
experiments/pro6000/glm53_ppmtp_r37_verify_graph/README.md
Normal file
66
experiments/pro6000/glm53_ppmtp_r37_verify_graph/README.md
Normal file
@ -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 记录为准。
|
||||
File diff suppressed because it is too large
Load Diff
@ -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],
|
||||
)
|
||||
@ -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 留盘可随时切回。
|
||||
@ -0,0 +1,15 @@
|
||||
=== conc_test round 1 ===
|
||||
req6: ok=True 23.8s text="\n\t\t{{\n\t\t}}</arg_value></think>I notice you've pasted what ap"
|
||||
req7: ok=True 23.0s text='\n\t\t#-name\n\t\t}\n\t}\n}\n</arg_value>`\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</think># Java 8 Stream API: Complete Guide\n\n## Introductio'
|
||||
req7: ok=True 6.5s text='languagelanguage=’s\n\t\t);\n\t}\n}\n</arg_value>`\n\n</think># 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</arg_value>`\n\nThis is a corrupted/garbled te'
|
||||
req7: ok=True 5.5s text='\n</think>I need to stop here and clarify something important'
|
||||
stage 8x16384: OK total=5.7s
|
||||
FINAL: ALL-OK
|
||||
@ -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
|
||||
==================================================
|
||||
@ -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
|
||||
==================================================
|
||||
@ -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
|
||||
@ -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 =====
|
||||
@ -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 "<flags>" [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
|
||||
Loading…
x
Reference in New Issue
Block a user