Skip to content

fix: correct MTP speculative sampling semantics - #1295

Open
Vinkle-hzt wants to merge 4 commits into
alibaba:mainfrom
Vinkle-hzt:codex/fix-mtp-speculative-sampling-semantics
Open

fix: correct MTP speculative sampling semantics#1295
Vinkle-hzt wants to merge 4 commits into
alibaba:mainfrom
Vinkle-hzt:codex/fix-mtp-speculative-sampling-semantics

Conversation

@Vinkle-hzt

@Vinkle-hzt Vinkle-hzt commented Aug 13, 2026

Copy link
Copy Markdown
Collaborator

Summary

  • Cherry-pick a887c9848c465df6e9768941960f497b1014033d onto the latest main.
  • Make greedy and stochastic rejection sampling follow their respective contracts.
  • Pass the exact normalized distribution used to generate each draft token through the existing all_probs / propose_probs path.
  • Map proposal probabilities from draft-vocabulary space to target-vocabulary space before verification, including non-identity and many-to-one maps.
  • Keep one generic CUDA/ROCm rejection-sampling path. There is no deterministic-draft flag, template specialization, or implicit proposal distribution.

Design

  • Stochastic acceptance is u * p < q; matching verifier/draft IDs do not bypass this test.
  • Stochastic rejection samples from relu(q - p).
  • Greedy verification accepts equal IDs and otherwise emits the verifier token directly.
  • Full acceptance emits the final token already produced by the target sampler as the bonus and consumes no extra rejection-sampling uniform value.
  • The current top-1 draft sampler publishes its actual proposal law as a normalized point mass in draft-vocabulary space. A future top-p draft sampler can publish its filtered and renormalized distribution through the same interface without changing the rejection kernel.
  • Draft-to-target probability mapping uses scatter_add, so many draft IDs mapping to one target ID preserve total probability mass.
  • The reusable mapped-probability buffer only reuses an exact step dimension, preserving contiguous layout when a future sampler changes its draft step count.
  • top_k > 1 remains rejected in FastTopKSampler; a future non-top-1 sampler must implement and return its actual proposal distribution rather than reuse top-k candidate scores.

Compatibility and rollout

  • propose_probs now explicitly means the actual normalized proposal distribution used to sample the transmitted draft token, not raw model softmax scores.
  • Prefill, decode, P2P, remote-cache, and memory-cache peers must run the same build, as required by the existing model RPC/KV-cache protocol contract. Mixed-version rolling deployment is unsupported.
  • Roll out and roll back all peers together. For stochastic top-1 draft requests, corrected semantics can reduce acceptance versus the previous biased implementation and can change same-seed outputs.
  • Monitor the existing proposed-token, accepted-token, average-accept-length, and end-to-end TPS metrics during rollout.

Review checklist

  • Removed every deterministic_draft runtime/template branch and kept one explicit-distribution kernel contract.
  • Synchronized CUDA, ROCm, and the shared CPU golden reference.
  • Covered same-ID stochastic rejection, middle rejection, full-accept bonus, greedy fallback, N=0, and N=1.
  • Covered non-degenerate proposal distributions and explicit top-1 point masses through the same kernel interface.
  • Covered vectorized relu(q - p) sampling with a point-mass token in a non-leading vector lane.
  • Covered non-identity and many-to-one draft-to-target vocabulary mappings and probability-mass conservation.
  • Covered batch > 1 with a changed draft-step capacity and asserted contiguous mapped probabilities.
  • Added fail-fast metadata checks and negative tests for undefined or mismatched proposal tensors.
  • Made validation tests assert the precise RTPException contract with core dumps disabled for the fixture.
  • Registered a ROCm op test that reuses the same cross-backend fixture.
  • Documented the wire-field semantics and same-build deployment requirement at in-memory, P2P, and RPC boundaries.

Validation

  • CUDA 12.9 before the final buffer-only amendment: 4/4 passed
    • //rtp_llm/cpp/normal_engine/speculative/test:mtp_executor_test
    • //rtp_llm/cpp/normal_engine/speculative/test:speculative_sampler_test
    • //rtp_llm/models_py/bindings/cuda/ops/tests:rejection_sampling_op_test
    • //rtp_llm/models_py/bindings/cuda/test:cuda_speculative_sampling_test
  • Final amendment checks:
    • focused PyTorch stride reproduction: old view non-contiguous, exact-step buffer contiguous
    • independent final-diff review: no P0/P1/P2 findings
    • clang-format --dry-run --Werror
    • git diff --check
  • The final incremental Bazel rerun was blocked before analysis by the container's internal SSH-only rules_python fetch. CUDA/ROCm compilation and ROCm hardware execution remain delegated to CI.

@LLLLKKKK LLLLKKKK left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AI Code Review - PR #1295

Status: BLOCKING

Summary: P0/1 · P1/2 · P2/5 · P3/3

Reviewed: commit a3d9e9da1230 · 2026-08-13 13:03 UTC+8

Blocking Issues

P0

  • kernel 接受/bonus 语义变更未同步共享黄金参考实现,既有 rejection sampling op 测试确定性失败 @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:288
    • 建议:在本 PR 内同步改写 referenceRejectionSampling:接受条件改为 (do_sample && u*p<q) || (!do_sample && same_token),bonus 门控改为 pos == num_speculative_tokens,并补上 !do_sample 拒绝时写 target_id 后跳过残差采样的分支,使其与 kernel 逐位对齐。该 header 由 cpp/testing/BUILD:79-92rejection_sampling_test_utils 导出、仅被 models_py/bindings/cuda/ops/tests:rejection_sampling_op_test 消费,且无 ROCm 变体,更新不会波及其他 target。同时补两条守护用例:do_sample=truesame_token 成立但 u*p>=q 必须走残差采样;do_sample=true 全接受但存在 same_token=false 时 bonus 必须等于 target_token_ids。本地跑通该 target 三个用例再提交,并在 PR 描述中说明旧黄金期望为何失效。

P1

  • ROCm 孪生内核未同步,同一请求在 CUDA/ROCm 上产出不同 token @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:268
    • 建议:在本 PR 同步移植 CUDA 侧的接受判据、direct_target_fallback 分支与 pos == num_speculative_tokens bonus 门控到 ROCm 内核;更优做法是把这段串行接受循环抽成两侧共享的 __device__ inline 辅助函数(放入共享头文件),消除「CUDA 内核 / ROCm 内核 / CPU 参考」三份并存的判据副本,并给 ROCm 侧补一条复用同一黄金参考的 op 测试。若本 PR 暂不覆盖 ROCm,请在 PR 描述与 sampling.cu 注释中显式声明「ROCm 语义未同步」及跟进计划,避免后续误认为已修复。
  • 随机采样路径的三处新语义没有任何判别性测试覆盖 @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:275
    • 建议:补三类确定性 kernel 用例:(1) do_sample=true、draft_probs 为 point mass、q(draft)=0.3u=0.5,断言即使 same_token 成立也被拒绝且发射 token 落在 relu(q-p) 支撑集;(2) do_sample=true 在中间位置被拒,按固定 u 断言 accepted=pos+1output[pos] 为残差采样结果、其后全为 -1;(3) do_sample=true 全接受但存在 same_token=false 的行,断言 bonus 等于 target_token_ids 末列。边界上再补 num_speculative_tokens=1 且首位即被拒一例。同时修正 CudaSpeculativeSamplingTest.cc:134-135 已过时的注释(其称 pos 0 由 relu(q-p) 重采样,实际走贪心直写分支)。

Non-blocking Suggestions

P2

  • draft all_probs 语义变更跨越 PD 分离的 RPC 边界,混版滚动升级期静默不一致 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:25
    • 建议:在 PR 描述与发布说明中明确「prefill/decode 必须同版本升级」并给出灰度顺序建议;更稳妥的做法是让 decode 侧不依赖对端概率语义——既然 top_k==1 的提议是确定性的,可只传 draft token id 并由 decode 侧本地构造 point mass(同时省掉每步一个全词表张量的跨进程传输),或在 payload 中加一个提议分布类型标识(softmax/point_mass),由 kernel 按标识选择判据并对未知标识 fail-fast。
  • one-hot 写在 d2t 映射之前,词表等宽且映射非恒等时 p 恒为 0 导致无条件接受 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:33
    • 建议:明确并固化索引口径:建议在 d2t 映射之后再构造 one-hot(或对 draft_probs 同步做一次 d2t 重排),使概率下标空间与 draft_token_ids 始终一致;同时在 validateRejectionSamplingParams 增加快速失败校验,当 d2t_map 已定义且非恒等时要求 draft_probs 已按 target 词表对齐,避免错位以「接受率 100%」的形式静默通过。测试侧补一条非恒等 d2t_map(如 {3,2,1,0}、等宽)用例,断言 all_probs 中值为 1.0 的下标与 kernel 读取 p 的下标一致;再补一条 draft 词表窄于 target 的用例覆盖 padding 分支。
  • 随机请求接受率将系统性下降,缺少实测数据与运维回滚手段 @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:286
    • 建议:补充 do_sample=true 场景下修复前后的 accept_len、每步 token 数与端到端 TPS 对比,以及分布正确性数据(例如与纯 target 采样做频率/KL 对比)并写入 PR 描述,同时说明同 seed 输出变化对 golden/复现类用例的影响。建议在 metrics 中按贪心/随机区分 accept rate 以便线上观测,并评估是否需要一个默认开启新语义、可临时关闭的配置项,使线上加速比回退时有运维回滚手段。若接受率损失不可接受,更标准的做法是让 do_sample 的流按 draft 自身分布采样,使 p 为真实分布,兼顾正确性与接受率。
  • 全接受时 bonus 直取 target_token_ids 的隐式契约无注释无测试 @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:292
    • 建议:在 pos == num_speculative_tokens 分支上方补注释,写明「bonus token 复用目标采样器的采样结果,要求 do_sample 行的 target_token_ids 必须来自随机采样而非 argmax」,并在 op 测试中加一条 do_sample=true 全接受用例固化该期望,使上游改动目标采样快路径条件时立刻在测试层暴露。顺带说明 uniform_samples 末列在新语义下不再被该路径消费,避免上游误以为需继续为它保留随机数预算。
  • top_k>1 分支保留被本 PR 修复的错误语义且无契约约束与断言 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:26
    • 建议:明确该分支契约:要么统一按 point-mass 语义处理(对 top-k 候选各自置 1 或按需归一化,并把 d2t 映射的 token_stride 改为实际列数),要么在 top_k>1 时用 RTP_LLM_CHECK_WITH_INFO fail-fast,并在注释中声明「不得用于 rejection sampling」。同时补一条 top_k=2 用例断言 all_probs 等于 torch::softmax(logits, -1),防止 point-mass 逻辑误扩散到多候选路径。

P3

  • top_k==1 路径 softmax 成为无用计算并新增每步全词表分配 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:12
    • 建议:top_k==1 分支直接对 logitstorch::argmax/torch::max,跳过 softmax;one-hot 写入一个与 draft_probs_padding_buffer_ 同风格的 grow-only 复用成员(narrow + zero_() + scatter_),避免每步新分配。若后续要彻底省掉整块全词表缓冲,可在 kernel 侧为 point-mass proposal 增加专用分支(p 恒为 1、残差为置零 q[draft_id]),但那属 kernel 接口变更,建议单独提并附 accept rate / TPS 数据。
  • 残差采样段遗留变量遮蔽、新引入的死分支与过时注释 @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:277
    • 建议:把 :277/:293 的循环下标改名(如 pad_idx,与文件末尾 :379-1 的循环保持同一命名)以消除遮蔽;删除或改写残差段中已不可达的 pos != num_speculative_tokens 判断、min(pos + 1, ...) 夹逼与相关注释,改为直接使用 pos + 1,并在残差段入口补一行说明「进入此处时 pos 必然属于 [0, num_speculative_tokens)」,避免陈旧注释与死分支误导后续维护者。
  • 新增用例落位偏离既有专用测试文件,断言未校验形状且含冗余拷贝 @ rtp_llm/cpp/normal_engine/speculative/test/MtpExecutorTest.cc:542
    • 建议:建议将该用例迁到 SpeculativeSamplerTest.cc 与既有 FastTopKSamplerTest 并列,那里依赖声明显式且正确,也能让 MtpExecutorTest.cc 继续专注 MtpExecutor 编排;若确需留在当前文件,请在 mtp_executor_test deps 中显式加入 exec_ops_test_lib 并补上与兄弟 target 一致的注释。同时参照兄弟用例先断言 token_idsdim() 与各维大小、增加 ASSERT_EQ(output.all_probs.sizes(), logits.sizes()) 前置断言,使失败信息直接指向形状契约变化而非张量降维异常;去掉多余的 .to(torch::kCUDA);再补一条 batch=2 且两行 argmax 不同的用例覆盖多行 scatter_ 边界。

Checklist Findings (13 fail / 26 total)

General Principles Checklist

  • [6.1] Architecture — 兼容性:外部 HTTP/RPC API、持久数据、配置、环境迁移安全 → issue draft all_probs 语义变更跨越 PD 分离的 RPC 边界,混版滚动升级期静默不一致
    draft all_probs 不只在进程内流转:MtpBatchStreamProcessor.cc:1001/1039draft_sampler_output.all_probs 存入流的 sp_output_buffer->all_probsPrefillRpcServer.cc:400-413 将其序列化进 proto 字段 propose_probsmodel_rpc/proto/model_rpc_service.proto:311),decode 侧经 DecodeRpcServer.cc:293-297StreamCacheResource.cc:224-225 还原后直接喂给 rejection sampling kernel。本 PR 把该张量语义由 softmax 改为 one-hot,而接受判据在对端 kernel 内,两侧必须同版本才成立:旧 prefill + 新 decode 时 p 仍是 softmax(p≪1),u*p<q 近似恒真,分布修复静默失效;新 prefill + 旧 decode
  • [6.1] Architecture — 可观测性:日志/指标/超时可操作、非噪声 → issue 随机请求接受率将系统性下降,缺少实测数据与运维回滚手段
    do_sample=true 的行(do_sample = !stream->generateConfig()->top1()SpeculativeSampler.cc:78),新接受集合 {u*p<q} 是旧集合 {same_token} ∪ {u*p<q} 的严格子集,output_accepted_token_num:286)只降不升:point-mass 下单位置接受概率由 min(1, q/p_softmax)p_softmax 通常远小于 1,接受概率常被抬到接近 1)降为 q(draft);把 same_token 粗估为概率约 q 的事件,旧值约 q*(2-q)、新值为 qq=0.5 时约 0.75 降到 0.5。另 SpeculativeSampler.cc:87-88 注释说明 uniform_samples 用 seeded generator 保证「deterministic acceptance for reproducible iter_count」,故同 seed 的输出 toke
  • [6.1] Architecture — 回滚路径:风险行为存在运维回滚手段 → issue 随机请求接受率将系统性下降,缺少实测数据与运维回滚手段
    do_sample=true 的行(do_sample = !stream->generateConfig()->top1()SpeculativeSampler.cc:78),新接受集合 {u*p<q} 是旧集合 {same_token} ∪ {u*p<q} 的严格子集,output_accepted_token_num:286)只降不升:point-mass 下单位置接受概率由 min(1, q/p_softmax)p_softmax 通常远小于 1,接受概率常被抬到接近 1)降为 q(draft);把 same_token 粗估为概率约 q 的事件,旧值约 q*(2-q)、新值为 qq=0.5 时约 0.75 降到 0.5。另 SpeculativeSampler.cc:87-88 注释说明 uniform_samples 用 seeded generator 保证「deterministic acceptance for reproducible iter_count」,故同 seed 的输出 toke
  • [6.1] Architecture — 状态不变量:创建/更新/失败/重试/回滚路径有效 → issue 全接受时 bonus 直取 target_token_ids 的隐式契约无注释无测试
    :288-296 只要 pos == num_speculative_tokens 就把 target_token_ids 末列当作 bonus 并在 :303 提前 return,不再像旧路径那样在 bonus 行按 relu(q-0)=q 重采样,也不再消费 uniform_samples[row*(n+1)+num_speculative_tokens]。这隐式要求 do_sample=true 的行其 target_token_ids 必须是目标分布的随机采样结果(当前取自 SpeculativeSampler.cc:68)。core/CudaSampleOp.cc:432-446 的 argmax 快路径要求整批 top_k 全为 1 且未请求 output_all_probs,投机路径需要 all_probs 故当前不会命中;但一旦该触发条件放宽,bonus 会静默退化成 greedy,kernel 侧没有注释或测试固化该期望。
  • [6.1] Architecture — 错误语义:fail-fast/retry/fallback/silent 行为显式 → issue top_k>1 分支保留被本 PR 修复的错误语义且无契约约束与断言
    注释称「保留既有多候选行为」,但 top_k>1 的 proposal 同样是确定性 top-k 选择,all_probs 仍为完整 softmax,一旦接入 rejection sampling 会重现本 PR 修复的过度接受偏差。更严重的是此时 token_ids 形状为 [batch, top_k],而 :33 固定以 token_stride=1token_offset=0 调用 execMappingDraft2Target,其实现按 index = i*token_stride + jCudaSampleOp.cc:495-501)遍历,只会映射扁平张量的前 batch_size 个元素,多候选下 d2t 映射本身即是错的。仓内生产调用点(MtpExecutor.cc:799/1166/1854)全为 top_k=1(默认值见 SpeculativeSampler.h:44),该分支既无使用者也无 all_probs 断言(SpeculativeSamplerTest.cc:21-32 只断言 `token_i
  • [6.1] Quality — PR description 说明动机与设计 → issue 随机请求接受率将系统性下降,缺少实测数据与运维回滚手段
    do_sample=true 的行(do_sample = !stream->generateConfig()->top1()SpeculativeSampler.cc:78),新接受集合 {u*p<q} 是旧集合 {same_token} ∪ {u*p<q} 的严格子集,output_accepted_token_num:286)只降不升:point-mass 下单位置接受概率由 min(1, q/p_softmax)p_softmax 通常远小于 1,接受概率常被抬到接近 1)降为 q(draft);把 same_token 粗估为概率约 q 的事件,旧值约 q*(2-q)、新值为 qq=0.5 时约 0.75 降到 0.5。另 SpeculativeSampler.cc:87-88 注释说明 uniform_samples 用 seeded generator 保证「deterministic acceptance for reproducible iter_count」,故同 seed 的输出 toke
  • [6.1] Software Engineering — DRY:重复非平凡逻辑被抽取或显式复用 → issue 新增用例落位偏离既有专用测试文件,断言未校验形状且含冗余拷贝
    仓内已有 FastTopKSampler 的专属单测 SpeculativeSamplerTest.ccFastTopKSamplerTest 套件)及目标 speculative_sampler_test,其 BUILD 注释明确写了 FastTopKSampler::forward calls execMappingDraft2Target 并显式声明 exec_ops_test_lib + cuda_implspeculative/test/BUILD:29-37)。新用例却放进 MtpExecutorTest.cc:该文件其余用例一律用 FakeFastTopKSampler(:194)打桩、从不执行真实 forwardmtp_executor_test deps 未声明 exec_ops_test_lib,仅靠 device_test_utils 间接带入(链接可用)。另 :550item<int64_t>() 依赖单元素隐式降维,未如兄弟用例(`SpeculativeSamplerTest.cc:1
  • [6.1] Software Engineering — KISS/YAGNI:无投机性抽象 → issue 残差采样段遗留变量遮蔽、新引入的死分支与过时注释
    新增的 padding 循环 for (int p = i + 1; ...)(:277)位于外层 float ... p = draft_probs[...](:264)作用域内,int p 遮蔽了表示 draft 概率的 p(旧代码同名循环变量在 p 作用域之外,不存在遮蔽),在概率语义的 kernel 里易误读,也可能在开启 -Wshadow 时告警。另外 :303 提前返回后 pos 必然小于 num_speculative_tokens,因此 :318:349if (pos != num_speculative_tokens) 恒真、:339min(pos + 1, num_speculative_tokens) 恒等于 pos + 1,而「there is no draft_probs for the bonus token」注释描述的分支已永不进入(旧内核在 all_same_token=falsepos=num_spec 时确实会走到),死分支由本 PR 新引入。
  • [6.1] Tests — 分布式/跨平台变更有对应覆盖 → issue ROCm 孪生内核未同步,同一请求在 CUDA/ROCm 上产出不同 token
    rtp_llm/models_py/bindings/rocm/speculative_sampling/sampling.cu:320-366 是本内核的逐行复制版,仍保留 s_all_same_tokensame_token || (do_sample && u*p<q)(:336)、if (all_same_token)(:349)与 if (s_all_same_token) return;(:364),本 PR 未改。该文件在 rocm/BUILD:83-97speculative_sampling 目标中实际参与构建,并由 core/CudaSampleOp.cc:819-836 的 ROCm rejectionSampling 调用。而 FastTopKSampler 是平台无关代码,one-hot 会同时喂给 ROCm:p 由 softmax 变为 1.0 使判据收紧,但 same_token 仍短路接受、拒绝时仍直吐 verifier token——结果既不同于修复前也不同于修复后的 CUDA,接受率、拒绝后 t
  • [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue 随机采样路径的三处新语义没有任何判别性测试覆盖
    本次改变三种行为:(1) do_sample=truesame_tokenu*p>=q 由接受变拒绝;(2) do_sample=true 拒绝后走残差采样而非直吐 verifier token;(3) 全接受时 bonus 固定取 target_token_ids 且不再消费 uniform_samples 末列。但 cuda/test/CudaSpeculativeSamplingTest.cc 四个既有用例均为退化输入:AllAccept(:36) 虽 do_sample=onesuniform_samples 全 0 且 probs 已是 one-hot,必然全接受;ImmediateReject(:90)/PartialAccept(:143) 的 do_sample 均为 torch::zerosBatchSizeZero(:203) 为空批。逐个复算后新旧输出完全相同,无判别力。本 PR 唯一新增用例只覆盖 FastTopKSampler 的 one-hot 输出,不触及 kernel;比
  • [6.1] Tests — 边界 case 覆盖(空、单元素、最大值) → issue 新增用例落位偏离既有专用测试文件,断言未校验形状且含冗余拷贝
    仓内已有 FastTopKSampler 的专属单测 SpeculativeSamplerTest.ccFastTopKSamplerTest 套件)及目标 speculative_sampler_test,其 BUILD 注释明确写了 FastTopKSampler::forward calls execMappingDraft2Target 并显式声明 exec_ops_test_lib + cuda_implspeculative/test/BUILD:29-37)。新用例却放进 MtpExecutorTest.cc:该文件其余用例一律用 FakeFastTopKSampler(:194)打桩、从不执行真实 forwardmtp_executor_test deps 未声明 exec_ops_test_lib,仅靠 device_test_utils 间接带入(链接可用)。另 :550item<int64_t>() 依赖单元素隐式降维,未如兄弟用例(`SpeculativeSamplerTest.cc:1

RTP-LLM Checklist

  • [I] 代码质量 — 删除或重命名内部 file、registry entry、model name、metric enum、op binding、plugin symbol 时,必须全仓搜索消费者,并提供替代实现、迁移说明或 smoke 覆盖;只有暴露到 HTTP/RPC/config/persisted format 时才按外部兼容性处理 → issue kernel 接受/bonus 语义变更未同步共享黄金参考实现,既有 rejection sampling op 测试确定性失败
    新 kernel 在 pos == num_speculative_tokens 时一律写 target_token_ids bonus 并提前 return(:288-303),接受判据按 do_sample 分流(:268)。共享参考 rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:352/363 仍是旧语义(same_token || (do_sample && u*p<q) + all_same_token),且不在本 PR。复算 runReferenceCases batch 3(do_sample=true,draft={2,4,1},target={0,4,3,4}):i=0 时 u*p=0.2*0.5=0.10<q=0.80 接受但 same_token=false,三个 draft 全接受后参考走 relu(q-0)(阈值 0.55,累加 0.10/0.30/0.60)返回 2,新 kernel 发 bonus=4,assertVectorEqual 在 flat
  • [I] 代码质量 — 同一功能用统一工具函数 → issue 新增用例落位偏离既有专用测试文件,断言未校验形状且含冗余拷贝
    仓内已有 FastTopKSampler 的专属单测 SpeculativeSamplerTest.ccFastTopKSamplerTest 套件)及目标 speculative_sampler_test,其 BUILD 注释明确写了 FastTopKSampler::forward calls execMappingDraft2Target 并显式声明 exec_ops_test_lib + cuda_implspeculative/test/BUILD:29-37)。新用例却放进 MtpExecutorTest.cc:该文件其余用例一律用 FakeFastTopKSampler(:194)打桩、从不执行真实 forwardmtp_executor_test deps 未声明 exec_ops_test_lib,仅靠 device_test_utils 间接带入(链接可用)。另 :550item<int64_t>() 依赖单元素隐式降维,未如兄弟用例(`SpeculativeSamplerTest.cc:1

Strengths

  • 修复与标准 speculative sampling 闭合:point-mass 提议 + relu(q-p) 残差,残差和为 1-q(draft),发射任一 token 的概率恰为 q(y),target 分布被精确保持。
  • 同时修掉旧实现两处真实偏差:随机请求下 same_token 无条件短路接受(接受率被高估、与残差分布不自洽),以及位置 0 被拒时跳过残差采样直吐 verifier token。
  • 贪心路径(do_sample=false)逐位推演与旧实现完全等价(失配写 target_id、其后补 -1accepted=pos+1),我复算 AllAccept/ImmediateReject/PartialAccept 三个既有用例断言全部仍成立,向后兼容明确。
  • 提前返回条件 s_direct_target_fallback || s_pos == num_speculative_tokens:303)只依赖 thread 0 写入、__syncthreads() 之后读取的两个共享标量,全 block 取值一致,不引入 __syncthreads() 分歧挂死;贪心失配与全接受两条路径各省掉一整轮 O(vocab) block reduce,是热路径净收益。
  • bonus token 改取已经过 top_k/top_p renorm 的 target_token_ids,比旧实现从 target_probs 二次重采样更贴合请求的采样配置,也与 forceSpAcceptSpeculativeSampler.cc:174)提取 bonus 的位置一致。
  • 关键分支都有解释「为什么」的注释(如「残差采样只对随机投机解码成立」);新增用例用 checkTensorEqual(内部 torch::equal 精确比较)断言精确 one-hot 值 {0,0,1,0} 而非仅 argmax,固定 logits 与期望值均可精确表示,无 flaky 风险。
  • 依赖链已核实:mtp_executor_testdevice_test_utilscpp/testing/BUILD:30-39 默认 select 分支)间接获得 exec_ops_test_libcuda_impl,新用例可正常 link,排除链接失败误报。

@LLLLKKKK LLLLKKKK left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AI Code Review - PR #1295 (non-blocking suggestions)

8 条 P2/P3 建议,不阻塞合并。阻塞判定与完整摘要见上一条 review。

Comment thread rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc Outdated
Comment thread rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc
Comment thread rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc
Comment thread rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc Outdated
Comment thread rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu Outdated
Comment thread rtp_llm/cpp/normal_engine/speculative/test/MtpExecutorTest.cc Outdated
@Vinkle-hzt
Vinkle-hzt force-pushed the codex/fix-mtp-speculative-sampling-semantics branch 2 times, most recently from 21dbb33 to 5cbe31d Compare August 13, 2026 09:36
@Vinkle-hzt
Vinkle-hzt requested a review from LLLLKKKK August 13, 2026 09:43

@LLLLKKKK LLLLKKKK left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AI Code Review - PR #1295

Status: LGTM

Summary: P0/0 · P1/0 · P2/8 · P3/7

Reviewed: commit 5cbe31d71858 · 2026-08-13 18:27 UTC+8

lgtm ready to ci

Non-blocking Suggestions

P2

  • draft all_probs 已成热路径死负载,但仍被逐步物化且硬校验 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:15
    • 建议:把「计算/搬运」与「传输兼容」拆开:1)用运行期开关(或按是否启用 PD 分离判定)控制是否物化 draft 概率,decode 单机场景直接跳过 softmax 与 stack/clone;2)把 MtpBatchStreamProcessor.cc:903 的硬 CHECK 降级为「仅在确实需要构造 propose_probs 时才校验」,避免为无验证消费者的数据 fail-fast;3)在 SpeculativeSampler.cc:13 的注释与 PR description 中写明灰度窗口与移除条件(目标版本或 issue 编号),灰度结束后连同 MtpBatchStreamProcessor 侧的 stack/clone 一并删除;4)若要以性能为由推进,请附上 decode 阶段的显存/带宽实测对照。
  • deterministic_draft 在唯一生产调用点硬编码 true,提议律契约无断言且 draft_probs 被静默忽略 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:96
    • 建议:1)让该标志由提议来源推导而非常量:由 MtpExecutor 构造 SpeculativeSampler 时传入提议律标识(如 ProposalLaw::Deterministic),或由 draft 采样器上报 is_deterministic_proposal,未来接入带 temperature 的 draft model 时在构造期即暴露;2)在 validateRejectionSamplingParams 增加 deterministic_draft => !draft_probs_d.defined() 的显式 RTP_LLM_CHECK,让「传了却被忽略」立即 fail-fast;3)在 sampling.hinvokeRejectionSampling 声明上方补契约注释,明确 deterministic_draft=truedraft_probs 允许为 nullptr、提议分布由 draft_token_ids 隐式给出,ROCm 同名声明保持同一措辞。
  • deterministic_draft=true 与贪心 / 混合 batch 的生产主组合零覆盖 @ rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:273
    • 建议:在 runDeterministicDraftCases 补两组字面量用例:1)deterministic_draft=truedo_sample 全 false,分别覆盖「verifier 全命中、bonus 取 target_token_ids[num_spec]」与「首位失配,输出 target token、其后补 -1accepted==pos+1」,顺带确认空 draft_probs 在贪心分支下不被解引用;2)同 batch 混合 do_sample={true,false},断言 false 行直发 verifier token、true 行按 relu(q-p) 重采样、两行互不污染且 accepted_token_num 逐行正确。再给 runZeroAndOneSpeculativeTokenCases(N=0/1)加 deterministic_draft=true 变体。harness 为 CUDA/ROCm 共用,补齐后两平台自动同步覆盖。
  • 点质量的向量化 lane 构造只在 VEC_SIZE=1 下被验证 @ rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:277
    • 建议:在 runDeterministicDraftCases 追加一组 vocab_size 为 4 的倍数(如 8 或 12)的字面量用例,使 vec_size 落到 2 与 4,并让被拒绝位置的 draft_token_ids 处于非 0 lane(draft_id % VEC_SIZE != 0,例如 1、6、11),使 p_vec[j] 的 lane 偏移与 q_vec.load 偏移的一致性得到实际验证。该 helper 位于共享 harness,补完后 CUDA 与 ROCm 同时受益。
  • SpeculativeSampler::batchSample 的参数装配路径完全没有测试覆盖 @ rtp_llm/cpp/normal_engine/speculative/test/SpeculativeSamplerTest.cc:10
    • 建议:在 SpeculativeSamplerTest.cc 增加针对 SpeculativeSampler::forward 的用例:构造 2~3 个 stream(含 top1 与非 top1、含 forceSpAccept),draft token_ids 用 int64 CPU 张量以覆盖 dtype/device 转换,喂入固定的 target all_probs/token_ids,断言 accept_tokens/accept_len 与手工推演一致、未定义 draft_probs 不触发校验异常、-1 -> 0 归零生效;至少覆盖「全接受」与「首位拒绝」两种配置,形成对装配契约的回归防线。
  • deterministic_draft 新增的参数校验分支缺少负例测试 @ rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:400
    • 建议:在 runRejectsInvalidTensorMetadata 增补三条负例:deterministic_draft=falsedraft_probs_d 未定义应抛异常;deterministic_draft=falsedraft_probs_d 的 batch/num_spec/vocab 维与 draft_token_ids_dtarget_probs_d 不一致应抛异常;deterministic_draft=truedraft_token_ids_dtarget_probs_d 的 batch 维不一致应抛异常。同时把 makeParams() 改为具名字段初始化,避免后续新增字段时位置错位。
  • 随机请求接受率将系统性下降,缺少实测数据与运维回滚手段 @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:274
    • 建议:这是必要的分布正确性修复,不建议改回旧逻辑,但请补齐运维面:1)在 accept_len 指标上报处按 do_sample(或 top1())维度拆分接受率/接受长度,使灰度期间可分别观察两类请求;2)补充非贪心场景 accept_len 均值与 decode 吞吐的前后对比,覆盖 propose_step=1>1temperature=0>0 四种组合,量化预期回退幅度;3)在 PR description 中写明「相同 seed 的随机采样请求输出会变化、accept_len 预期下降」这一兼容性影响,并明确回滚手段(版本回滚,或可运行期切换的判据开关),避免被误判为性能回归。
  • FakeSpeculativeSampler::checkInputs 是永不执行的死断言 @ rtp_llm/cpp/normal_engine/speculative/test/MtpExecutorTest.cc:226
    • 建议:在 FakeSpeculativeSampler::forward 返回前调用 checkInputs(streams, draft_sampler_output, target_sampler_output)(与 FakeFastTopKSampler 一致),并核对 :833/:1198 已推入的期望在新语义下是否仍成立(argmax(logits) 与旧 max(softmax(logits)) 索引等价,预计无需改数值);若暂不打开断言,应删除对应 setInputs 调用以免留下误导性的「已覆盖」表象。

P3

  • dtype 转换使 hold_host 的 H2D 生命周期契约被绕过 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:49
    • 建议:调整顺序,使被持有的张量与拷贝源一致:先在 host 侧完成 to(torch::kInt32)(必要时 pin_memory()),再 hold_host 该结果,最后只做纯设备搬运 to(target_device, /*non_blocking=*/true);或在入口断言 draft token dtype 必为 int32,把转换责任上移到生产者。
  • draft_token_ids 形状校验在 shape 来源切换后成为恒真断言 @ rtp_llm/models_py/bindings/core/CudaSampleOp.cc:93
    • 建议:删除这两条恒真断言(真正的交叉约束已由 target_probs_d.size(0/1)uniform_samples_d.size(0/1)target_token_ids_d.size(0)==target_token_rows 覆盖),或改为显式校验 shape 推导源与其他张量的一致性,避免留下无效校验噪声。
  • FastTopKSampler 保留了只允许取值 1 的 top_k 参数 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:10
    • 建议:移除 top_k 形参并同步更新 MtpExecutor.cc:1854 的显式传参,必要时把类/方法重命名为体现 argmax 语义的名字;若需保留虚函数签名以兼容测试替身(MtpExecutorTest.cc:198 的 override),请在头文件注释中写明该参数已废弃且仅接受 1。
  • kernel 内点质量构造重复 4 份且在 vocab 循环内反复读取循环不变量 @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:326
    • 建议:在两个 vocab 循环之前取到寄存器:IdType det_draft_id = 0; if constexpr (DETERMINISTIC_DRAFT) { det_draft_id = draft_token_ids[row_idx * num_speculative_tokens + pos]; },两个循环共用该值以消除重复取址;接受循环中把 q/p 的读取移入 if (do_sample[row_idx]) 内。建议把点质量填充抽成 __device__ __forceinline__ 小函数(不影响内联与向量化),ROCm 侧同步处理,避免后续两平台改漏。
  • ROCm 侧无共享头声明、p_vec 元素类型不一致、相邻 bool 模板参数易写反 @ rtp_llm/models_py/bindings/rocm/speculative_sampling/sampling.cu:400
    • 建议:为 ROCm 侧补一个与 cuda/kernels/speculative_sampling/sampling.h 对称的头文件(或复用同一头文件并按 USING_ROCM 切换返回类型),让定义与调用方共享同一声明;p_vec[j] 改为 1.0f/0.0f 或把 p_vec 声明为 vec_t<DType, VEC_SIZE>,与 q_vecrelu_q_minus_p 类型统一;在两处 kernel 模板参数列表上给这两个 bool 加简短注释,降低写反风险。
  • 异常断言过宽且未遵循仓库关闭 core dump 开关的惯例 @ rtp_llm/cpp/normal_engine/speculative/test/SpeculativeSamplerTest.cc:32
    • 建议:参照 BlockPoolTest.cc 为该测试加 fixture,在 SetUp 中保存并将 StaticConfig::user_ft_core_dump_on_exception 置为 false、TearDown 中恢复;同时把断言收紧为 EXPECT_THROW(sampler.forward(logits, 2), rtp_llm::RTPException)
  • runLiteralCase 的 13 个位置参数与裸字面量墙影响可维护性,且断言不具备 fatal 语义 @ rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:45
    • 建议:将入参收敛为具名结构体(如 LiteralCase{batch_size; num_spec; vocab_size; ...})并在调用点使用指定初始化器书写,使布局意图自解释;同时在 :70 之后加入 ASSERT_NO_FATAL_FAILURE 包裹或 if (::testing::Test::HasFatalFailure()) return;,让参考实现比对失败即终止该用例。

Checklist Findings (14 fail / 26 total)

General Principles Checklist

  • [6.1] Architecture — 兼容性:外部 HTTP/RPC API、持久数据、配置、环境迁移安全 → issue 随机请求接受率将系统性下降,缺少实测数据与运维回滚手段
    旧判据 same_token || (do_sample && u*p_softmax < q) 既有 same_token 短路接受,p_softmax <= 1 又使接受概率达 min(1, q/p_softmax) >= q。新判据在 deterministic_draft 下对 do_sample=true 行退化为 u < q,接受概率严格降为 q(draft_id),接受率单调不增;同一 random_seed 的请求也会产出与上一版本不同的 token 序列(uniform 下标 min(pos+1,N) -> pos+1 已核实恒等,差异仅来自判据)。贪心行逐位等价,故影响面限于非贪心请求。但 PR 未给出 accept_len/吞吐对照,且 MtpExecutor::stageAcceptLenMetrics(MtpExecutor.cc:334-359)只上报跨全部流求和的 accept_len 与 propose_token_num,无法区分贪心与随机请求。
  • [6.1] Architecture — 可观测性:日志/指标/超时可操作、非噪声 → issue 随机请求接受率将系统性下降,缺少实测数据与运维回滚手段
    旧判据 same_token || (do_sample && u*p_softmax < q) 既有 same_token 短路接受,p_softmax <= 1 又使接受概率达 min(1, q/p_softmax) >= q。新判据在 deterministic_draft 下对 do_sample=true 行退化为 u < q,接受概率严格降为 q(draft_id),接受率单调不增;同一 random_seed 的请求也会产出与上一版本不同的 token 序列(uniform 下标 min(pos+1,N) -> pos+1 已核实恒等,差异仅来自判据)。贪心行逐位等价,故影响面限于非贪心请求。但 PR 未给出 accept_len/吞吐对照,且 MtpExecutor::stageAcceptLenMetrics(MtpExecutor.cc:334-359)只上报跨全部流求和的 accept_len 与 propose_token_num,无法区分贪心与随机请求。
  • [6.1] Architecture — 回滚路径:风险行为存在运维回滚手段 → issue 随机请求接受率将系统性下降,缺少实测数据与运维回滚手段
    旧判据 same_token || (do_sample && u*p_softmax < q) 既有 same_token 短路接受,p_softmax <= 1 又使接受概率达 min(1, q/p_softmax) >= q。新判据在 deterministic_draft 下对 do_sample=true 行退化为 u < q,接受概率严格降为 q(draft_id),接受率单调不增;同一 random_seed 的请求也会产出与上一版本不同的 token 序列(uniform 下标 min(pos+1,N) -> pos+1 已核实恒等,差异仅来自判据)。贪心行逐位等价,故影响面限于非贪心请求。但 PR 未给出 accept_len/吞吐对照,且 MtpExecutor::stageAcceptLenMetrics(MtpExecutor.cc:334-359)只上报跨全部流求和的 accept_len 与 propose_token_num,无法区分贪心与随机请求。
  • [6.1] Architecture — 状态不变量:创建/更新/失败/重试/回滚路径有效 → issue dtype 转换使 hold_host 的 H2D 生命周期契约被绕过
    代码先 buffer_holder_.hold_host(draft_token_ids)(:50),再 draft_token_ids.to(target_device, torch::kInt32, true)(:52)。TensorHolder::hold_host(ModelTypes.h:123)只持有传入的那个 CPU 张量;当来源为 host 且 dtype 非 int32 时,ATen 会在 host 侧先生成一个未被 holder 持有、且非 pinned 的 int32 临时张量作为拷贝源,holder 持有的却是转换前的缓冲,从而失去 :27-29 刻意维护的 pinned + holder 快路径与 release 契约。当前生产链路已在 MtpBatchStreamProcessor.cc:911 统一 .to(torch::kInt32),故属契约表达不一致的潜在风险,而非现网故障。
  • [6.1] Architecture — 错误语义:fail-fast/retry/fallback/silent 行为显式 → issue deterministic_draft 新增的参数校验分支缺少负例测试
    CudaSampleOp.cc:53 把 draft_probs_d 的 dtype/dim/contiguous/device 校验整体收进 if (!params.deterministic_draft),:86-91 新增 shape[0..2] 三条一致性检查,batch_size/num_speculative_tokens 的取值来源也改为 draft_token_ids_d(:72-73)。但 runRejectsInvalidTensorMetadata(:400)本次未做任何扩展,其 makeParams()(:422-431)仍是 8 字段位置聚合初始化、deterministic_draft 取默认 false。结果是:新增三条 shape 检查无测试;「deterministic_draft=falsedraft_probs_d 未定义」这一组合无测试,而 SpeculativeSampler.cc:96 恰恰传入 torch::Tensor(),一旦该 flag 被误传 false 即命中该组合。
  • [6.1] Quality — PR description 说明动机与设计 → issue 随机请求接受率将系统性下降,缺少实测数据与运维回滚手段
    旧判据 same_token || (do_sample && u*p_softmax < q) 既有 same_token 短路接受,p_softmax <= 1 又使接受概率达 min(1, q/p_softmax) >= q。新判据在 deterministic_draft 下对 do_sample=true 行退化为 u < q,接受概率严格降为 q(draft_id),接受率单调不增;同一 random_seed 的请求也会产出与上一版本不同的 token 序列(uniform 下标 min(pos+1,N) -> pos+1 已核实恒等,差异仅来自判据)。贪心行逐位等价,故影响面限于非贪心请求。但 PR 未给出 accept_len/吞吐对照,且 MtpExecutor::stageAcceptLenMetrics(MtpExecutor.cc:334-359)只上报跨全部流求和的 accept_len 与 propose_token_num,无法区分贪心与随机请求。
  • [6.1] Software Engineering — DRY:重复非平凡逻辑被抽取或显式复用 → issue runLiteralCase 的 13 个位置参数与裸字面量墙影响可维护性,且断言不具备 fatal 语义
    runLiteralCase 有 13 个位置参数,其中 4 个 int 与 4 个同类型 vector 相邻;调用点(:172、:239、:277、:295)是没有任何字段名的 float/int 字面量数组,batch × (num_spec+1) × vocab 的三维布局只能靠人工数元素个数校验,后续新增用例极易把 target_probsuniform_samples 或行序写错而仍然编译通过。此外 assertVectorEqual(:480)内的 ASSERT_EQ 只能从该 helper 返回,:69-70 参考实现比对失败后仍会继续构造张量并执行 rejectionSampling(:87),失败输出会叠加参考实现与 kernel 两侧错误,难以第一时间区分是 oracle 写错还是 kernel 回退。
  • [6.1] Software Engineering — ISP:调用方不依赖无关大接口 → issue FastTopKSampler 保留了只允许取值 1 的 top_k 参数
    forward(const torch::Tensor&, int top_k = 1)(SpeculativeSampler.h:44)签名不变,但函数体第一行即 RTP_LLM_CHECK_WITH_INFO(top_k == 1, ...),即接口暴露一个只有唯一合法取值的参数,把编译期可消除的错误推迟为运行期抛异常;类名 FastTopKSampler 与「仅支持 argmax」的新语义也不再匹配。全仓三处生产调用点(MtpExecutor.cc:799/1166/1854)均用默认值或显式 1,该检查线上不会触发。
  • [6.1] Software Engineering — KISS/YAGNI:无投机性抽象 → issue FastTopKSampler 保留了只允许取值 1 的 top_k 参数
    forward(const torch::Tensor&, int top_k = 1)(SpeculativeSampler.h:44)签名不变,但函数体第一行即 RTP_LLM_CHECK_WITH_INFO(top_k == 1, ...),即接口暴露一个只有唯一合法取值的参数,把编译期可消除的错误推迟为运行期抛异常;类名 FastTopKSampler 与「仅支持 argmax」的新语义也不再匹配。全仓三处生产调用点(MtpExecutor.cc:799/1166/1854)均用默认值或显式 1,该检查线上不会触发。
  • [6.1] Tests — 分布式/跨平台变更有对应覆盖 → issue 点质量的向量化 lane 构造只在 VEC_SIZE=1 下被验证
    kernel 用 vec_size = std::gcd(16 / sizeof(DType), target_vocab_size)(sampling.cu:424)经 DISPATCH_ALIGNED_VEC_SIZE 选模板参数,float 下取 1/2/4。runDeterministicDraftCases 两组 vocab 均为 5、runStochasticSemanticsCases 为 3 → gcd(4,V)=1,只实例化 VEC_SIZE=1,内层 j 恒为 0;vocab 为 4 与 16 的用例(harness:314/358、CudaSpeculativeSamplingTest.cc:39/94/146)又全部 deterministic_draft=false,走 p_vec.load 而非 one-hot 分支。生产词表为 4 的倍数,走的正是未覆盖的 token_id=(i*BLOCK_THREADS+tx)*VEC_SIZE+j(sampling.cu:329/365,ROCm:398/435);lane 偏
  • [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue 异常断言过宽且未遵循仓库关闭 core dump 开关的惯例
    TopKGreaterThanOneIsRejected 使用 EXPECT_ANY_THROW(sampler.forward(logits, 2))。而 RTP_LLM_CHECK_WITH_INFO 最终走 myAssert(rtp_llm/cpp/utils/AssertUtils.cc:13-22),当 StaticConfig::user_ft_core_dump_on_exception 为真时是 abort() 而非抛异常。仓库既有惯例明确:rtp_llm/cpp/cache/test/BlockPoolTest.cc:24-33 在 SetUp/TearDown 中显式置 false 并还原。当前用例依赖静态初值,若执行环境开启该开关会以进程 abort 形式失败;且 EXPECT_ANY_THROW 会被 torch 侧任意无关异常满足,无法证明失败原因就是 top_k != 1 的契约检查。
  • [6.1] Tests — 边界 case 覆盖(空、单元素、最大值) → issue 点质量的向量化 lane 构造只在 VEC_SIZE=1 下被验证
    kernel 用 vec_size = std::gcd(16 / sizeof(DType), target_vocab_size)(sampling.cu:424)经 DISPATCH_ALIGNED_VEC_SIZE 选模板参数,float 下取 1/2/4。runDeterministicDraftCases 两组 vocab 均为 5、runStochasticSemanticsCases 为 3 → gcd(4,V)=1,只实例化 VEC_SIZE=1,内层 j 恒为 0;vocab 为 4 与 16 的用例(harness:314/358、CudaSpeculativeSamplingTest.cc:39/94/146)又全部 deterministic_draft=false,走 p_vec.load 而非 one-hot 分支。生产词表为 4 的倍数,走的正是未覆盖的 token_id=(i*BLOCK_THREADS+tx)*VEC_SIZE+j(sampling.cu:329/365,ROCm:398/435);lane 偏

RTP-LLM Checklist

  • [I] 代码质量 — 删除或重命名内部 file、registry entry、model name、metric enum、op binding、plugin symbol 时,必须全仓搜索消费者,并提供替代实现、迁移说明或 smoke 覆盖;只有暴露到 HTTP/RPC/config/persisted format 时才按外部兼容性处理 → issue draft all_probs 已成热路径死负载,但仍被逐步物化且硬校验
    SpeculativeSampler.cc:96/104 固定传未定义 torch::Tensor()deterministic_draft=true,kernel 在 DETERMINISTIC_DRAFT 分支从不读 draft 概率(sampling.cu:266-270),rejection sampling 已无 draft all_probs 数值消费者。但生产侧仍全量物化::15 每个 draft step 做一次全词表 torch::softmax;MtpBatchStreamProcessor.cc:919 仍 torch::stack(...).contiguous(),:1002/:1040 再 narrow(...).to(kCUDA).clone()。该数据现仅服务 PD 分离的 propose_probs 载荷(DecodeRpcServer.cc:297、StreamCacheResource.cc:225 回写 sp_output_buffer->all_probs),却仍由 :903 的硬 `RTP_LLM_C
  • [I] 代码质量 — 同一功能用统一工具函数 → issue runLiteralCase 的 13 个位置参数与裸字面量墙影响可维护性,且断言不具备 fatal 语义
    runLiteralCase 有 13 个位置参数,其中 4 个 int 与 4 个同类型 vector 相邻;调用点(:172、:239、:277、:295)是没有任何字段名的 float/int 字面量数组,batch × (num_spec+1) × vocab 的三维布局只能靠人工数元素个数校验,后续新增用例极易把 target_probsuniform_samples 或行序写错而仍然编译通过。此外 assertVectorEqual(:480)内的 ASSERT_EQ 只能从该 helper 返回,:69-70 参考实现比对失败后仍会继续构造张量并执行 rejectionSampling(:87),失败输出会叠加参考实现与 kernel 两侧错误,难以第一时间区分是 oracle 写错还是 kernel 回退。

Strengths

  • 语义修正方向正确且有理论支撑:提议由 torch::argmax(SpeculativeSampler.cc:16)产生,提议律确为点质量,p=1 下接受概率 min(1,q/1)=q、残差 relu(q-onehot) 归一化为 1-q(draft_id),均为标准投机解码结论;旧代码用 draft softmax 当提议律又对 same_token 无条件短路,会同时高估接受率并偏向 draft token。
  • 顺带消除一处潜在偏差:旧 padding 仅在 draft_token_probs.size(2) != target_token_probs.size(2) 时触发,若两侧词表等宽而 d2t 映射非恒等,p 取自 draft 空间而 q 取自 target 空间;改为完全不依赖 draft 概率后,该不一致在结构上不再可能。
  • FastTopKSampler::forward 收敛到 top_k==1 修掉另一处隐患:execMappingDraft2Target 一直按 token_stride=1 调用,旧 top_k>1 分支会产出 [batch, top_k] 而只映射前 batch 个元素。
  • 删除 draft_probs_padding_buffer_ 及其 zero_() + index_put_({..., d2t_map_}),去掉热路径上一次 [batch, num_spec, target_vocab] 规模的 scatter 与一处依赖 shape 的 host 分支,对 CUDA graph 捕获稳定性亦为正向收益;if constexpr 编译期分派使热路径无新增逐 token 运行时分支。
  • 输出契约收紧为「N+1 个位置必被写入」,三条出口(全接受发 bonus、贪心直发、随机残差)各自负责 -1 填充;测试用 torch::full(..., -7) 哨兵预填(harness:79)即可检出漏写,逐路径核对未发现未写入槽位。
  • 全接受时 bonus 复用 target sampler 的最终 token(sampling.cu:299-302)而非 argmax 快路径,且不额外消耗 uniform 采样,符合无偏性要求;新增注释明确了这一契约。
  • 测试是真正的双 oracle:runLiteralCase(harness:45)先断言 CPU 参考实现等于硬编码期望(:69),再断言 kernel 等于同一期望(:89),避免「参考实现与 kernel 同步写错仍通过」。
  • validateRejectionSamplingParams 把 shape 基准从 draft_probs_d 迁到 draft_token_ids_d/target_probs_d(CudaSampleOp.cc:72-75),并在非确定性分支补齐 draft_probs_d 三维一致性校验,比原先只信任 draft_probs_d 更健壮;两平台 rejectionSampling 均对未定义 draft probs 显式取 nullptr(:476/:834)。
  • draft_token_ids.to(target_device, torch::kInt32, true) 显式统一 dtype,使 forceSpAccept 路径上的 torch::cat/torch::where(:133-139)不再因上游 int64 而污染 accept_tokens 类型。
  • 跨平台闭环完整:新增 RocmRejectionSamplingOpTest.cc 复用平台无关 harness,5 个 TEST_F 与 CUDA 侧一一对应;依赖经 //rtp_llm/cpp/testing:rejection_sampling_test_utils(testing/BUILD:80-92)→ rejection_sampling_ops_test_lib(core/BUILD:302-331)的 using_rocm select 与 alwayslink=1 闭合到 HIP 实现,torch_deps()/tags/exec_properties 与同目录既有目标一致,ROCm 首次获得该 kernel 的算子级回归。

Comment thread rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc Outdated
Comment thread rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc Outdated
Comment thread rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp Outdated
Comment thread rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp
Comment thread rtp_llm/cpp/normal_engine/speculative/test/SpeculativeSamplerTest.cc Outdated
Comment thread rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc Outdated
Comment thread rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu Outdated
Comment thread rtp_llm/models_py/bindings/rocm/speculative_sampling/sampling.cu Outdated
Comment thread rtp_llm/cpp/normal_engine/speculative/test/SpeculativeSamplerTest.cc Outdated
Comment thread rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp
@LLLLKKKK
LLLLKKKK dismissed their stale review August 13, 2026 10:27

LGTM:阻断已解除(rtpcli 自动清除旧红标)

@Vinkle-hzt
Vinkle-hzt force-pushed the codex/fix-mtp-speculative-sampling-semantics branch 3 times, most recently from 4532394 to b5cd0ee Compare August 13, 2026 11:51
@Vinkle-hzt
Vinkle-hzt force-pushed the codex/fix-mtp-speculative-sampling-semantics branch from b5cd0ee to 764cdbf Compare August 13, 2026 12:34

@Vinkle-hzt Vinkle-hzt left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AI Code Review - PR #1295

Status: LGTM

Summary: P0/0 · P1/0 · P2/0 · P3/0

Reviewed: commit 764cdbf8ad33 · 2026-08-13 20:36 UTC+8

lgtm ready to ci

Checklist ✅ (102 items passed)

Strengths

  • proposal probability 被明确为“生成该 draft token 时实际使用的归一化提议分布”,避免将原始 softmax 或候选分数误当提议律,并为未来 top-p 迁移保留统一接口。
  • draft-to-target 映射使用 scatter_add,正确覆盖非恒等与多对一映射并保持概率质量。
  • CUDA、ROCm 与共享 CPU golden 使用同一接受、拒绝残差和 bonus 语义,覆盖随机与贪心混合路径。
  • 输入张量的 dtype、shape、contiguous 与定义状态均 fail-fast,负向测试断言精确异常契约。
  • grow-only 映射缓冲仅复用 batch 容量,step 数变化时重新分配精确布局,避免 future sampler 动态 step 产生非连续 view。
  • PR 描述和代码边界注释明确了同版本部署要求、回滚方式与线上观察指标。

@Vinkle-hzt
Vinkle-hzt requested a review from LLLLKKKK August 13, 2026 12:40
@Vinkle-hzt

Copy link
Copy Markdown
Collaborator Author

已按最终 review 收口到 commit 764cdbf8ad33:没有引入 determined/deterministic draft probability 模式;rejection sampler 始终消费实际、归一化的 proposal distribution。当前 top-1 producer 写入真实 point-mass proposal,未来 top-p producer 可通过同一接口写入 filtered + renormalized proposal,无需修改 kernel。另补了动态 step buffer 的 batch>1 contiguous 回归。所有旧 blocking threads 已 resolve;请基于最终 head 做 fresh review。

@LLLLKKKK LLLLKKKK left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AI Code Review - PR #1295

Status: LGTM

Summary: P0/0 · P1/0 · P2/10 · P3/10

Reviewed: commit 764cdbf8ad33 · 2026-08-13 21:34 UTC+8

lgtm ready to ci

Non-blocking Suggestions

P2

  • do_sample 判据改用 generateConfig 后,greedy 正确性依赖跨三文件的未声明重归一化不变量 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:145
    • 建议:建议改为 do_sample && !top1():q 为 one-hot 时两条路径输出位相同,既恢复 greedy 早退快路径,又消除对重归一化不变量的隐式依赖,并保留 top1() 的消费者。若坚持当前写法,请在取 target_token_probs 处显式注释该不变量(q 必须是 post-top-k/top-p 重归一化分布),并二选一加固:(1) 把 return_original_all_probs 透传到 SpeculativeSamplerRTP_LLM_CHECK;(2) 加锚定测试断言 MTP 路径下该标志恒为 false,使后续改动该模式时 CI 立即失败。同时清理已无消费者的 top1() 或注明保留原因。
  • rejection kernel 的三条入参语义契约既未在 sampling.h 声明也无法 fail-fast @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.h:17
    • 建议:在 invokeRejectionSampling 声明处写清上述三条契约,并注明 do_sample 选择的是两套不同算法而非同一算法的开关。同时建议在 op 层或 debug 构建下增加一条不变量校验(按行检查 draft_probs[draft_id] > 0 或行和接近 1),让契约被破坏时 fail-fast,而不是静默输出错误分布。
  • 随机采样行接受率系统性下降,缺少实测数据、接受率指标与运维回滚手段 @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:268
    • 建议:合入前给出 MTP do_sample=true 场景的 accept_len 与 TPS 前后对比数据,并在 PR description 中说明该吞吐变化及原因。若下降幅度超出预期,提供运维可控的降级开关(例如把旧 same_token 快路径保留为显式可选项);同时建议新增/复用接受长度指标,使灰度期具备对照与回滚判断依据。
  • propose_probs 取值语义与 dtype 静默变更,三处反序列化零校验且 same-build 条款不覆盖生成路径 @ rtp_llm/cpp/model_rpc/proto/model_rpc_service.proto:312
    • 建议:在 PrefillLoadCaller.ccDecodeRpcServer.cc 的还原处对 propose_probs 校验 dtype(FP32)与末维长度,与相邻 propose_tokens 的 INT32 校验保持一致,不满足时按请求粒度返回 gRPC 错误并打印 request_id 与 peer,让契约违约在传输边界 fail-fast 而非下沉为共享采样器异常。同时把 same-build 契约范围显式扩写到生成路径 RPC(RemoteGenerate/StartLoadpropose_* 载荷),或增加提议分布语义标识(枚举/最小协议版本号)供接收方 fail-fast 或降级为 greedy 校验并打一次 WARNING;PR description 需明确「PD 对必须同版本发布与同版本回滚」。
  • 契约注释把 draft 空间的 probs 与已映射为 target 空间的 propose_token_ids 绑定描述 @ rtp_llm/cpp/model_rpc/proto/model_rpc_service.proto:311
    • 建议:拆成两句并点明索引空间差异:propose_token_ids 已由 d2t_map 映射为 target 词表 id;propose_probs 保持 draft 词表列索引,消费方必须先用同源 d2t_map 映射(参考 SpeculativeSampler::mapDraftProbsToTarget)后再与 target 概率对齐。同时补齐 dtype=FP32、rank/shape 约定与字段缺省语义,并在 field 8 上补一行 target-space 说明,使 proto 注释成为可据以实现的完整契约。
  • batchSample 的三处语义变更零覆盖且恰被 fake 屏蔽,mapDraftProbsToTarget 用例全在 CPU @ rtp_llm/cpp/normal_engine/speculative/test/SpeculativeSamplerTest.cc:96
    • 建议:补一条针对真实 SpeculativeSampler::forward 的窄测试:构造 2 条最小 stream(top_k=1do_sample 默认 true 与显式 false)、小 vocab 的 one-hot target all_probs 与可控 uniform_samples,断言 accept_tokens/accept_len 落在各自分支的期望值上,把 do_sample 映射与 int32 归一化钉住;至少要有一个用例锁定「top_k=1 + 默认 do_sample」这一存量流量组合。若 GenerateStream 装配成本过高,可把该派生抽为可测纯函数(如 resolveDoSample(const GenerateConfig&))做 top_k×do_sample 表驱动断言。同时补 CUDA 上的 mapDraftProbsToTarget 用例(many-to-one 与等宽非恒等各一),断言结果对 draft_probs.sum(-1) 守恒。
  • bonus token 来源变更在全部新旧用例上均不可判别 @ rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:168
    • 建议:在 runStochasticSemanticsCases 补一个字面量用例:do_sample=true、全部提议被随机接受(draft 与 verifier token 不同),且 target_token_ids 末列刻意不等于「按给定 uniform_samples[N]relu(q-p) 采样」的结果,断言输出为末列 token。放在共享 harness 即可同时在 CUDA 与 ROCm 两个 target 生效,锁定新 bonus 语义。
  • greedy 直出分支与混合 do_sample batch 缺字面期望值,ROCm 仅与同源参考实现互证 @ rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:113
    • 建议:在 runStochasticSemanticsCases 同级补 runGreedySemanticsCases(),用 runLiteralCase 给出手工期望值,至少覆盖首位即不匹配、中途不匹配、全部匹配三种形态;并增加一条同批 do_sample={false,true} 的混合用例,让贪心直出与残差重采样在同一次 launch 内同时发生。在 CudaRejectionSamplingOpTest.ccRocmRejectionSamplingOpTest.cc 同时注册,使两平台的贪心契约都有不依赖参考实现的锚点。
  • mapDraftProbsToTarget 的直通分支与新增十余处校验分支缺少用例 @ rtp_llm/cpp/normal_engine/speculative/test/SpeculativeSamplerTest.cc:45
    • 建议:复用已正确保存/恢复 user_ft_core_dump_on_exceptionFastTopKSamplerValidationTest fixture 补齐:d2t_map 未定义且 vocab 相等时断言返回值与入参同 data_ptr(零拷贝直通)、vocab 不等时 EXPECT_THROW(..., RTPException);构造函数侧覆盖 d2t_map 含负值、最大值等于/超过 target_vocab_size、dtype 非 int64;mapDraftProbsToTarget 侧覆盖非 3-D、非 float32、d2t_map numel 不匹配。
  • MtpExecutorTest 两处夹具 all_probs 与 token_ids 自相矛盾,违反本 PR 建立的点质量契约 @ rtp_llm/cpp/normal_engine/speculative/test/MtpExecutorTest.cc:714
    • 建议:把 :714 改为 {{0.0f, 0.0f, 0.0f, 1.0f}}(对齐 token 3,参照 :807-:808 的兄弟写法),:1162 改为 {0.0f,1.0f,0.0f,0.0f, 0.0f,0.0f,1.0f,0.0f}(对齐 token_ids={1,2}),并同步核对由其拼出的期望张量与相关 stream 状态断言。可进一步在 FakeSpeculativeSampler::checkInputs 中加一条 argmax(all_probs, -1)token_ids 的一致性断言,让这类夹具错配在编写期即暴露。

P3

  • 残差分布和为 0 时静默回退到词表末位 token @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:335
    • 建议:在 op 层 validateRejectionSamplingParams 增加一条 draft_probs 行和校验,或在 kernel 内对 sum_relu_q_minus_p <= 0 显式回退到 target_token_ids 的 verifier token,并配一条可观测的计数/日志,避免后续扩展提议分布类型时出现无声错误。
  • 同文件遗留 ChainSpeculativeSampling 保留旧语义,形成语义分叉的死代码 @ rtp_llm/models_py/bindings/rocm/speculative_sampling/sampling.cu:458
    • 建议:在本 PR 或紧随的清理 PR 中删除 ChainSpeculativeSampling / execChainSpeculativeSampling 链路(删除前按 R.I.2 再做一次全仓消费者搜索确认);若因外部依赖需保留,请在其定义处加一行注释标明「legacy,已废弃,权威实现为 rejection_sampling_kernel,语义不同,勿作为基准」。
  • d2t_map 校验在三处不统一,作为写索引的 mapDraftProbsToTarget 反而不校验取值范围 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:78
    • 建议:把校验抽成共享 helper(如 validateD2TMap(d2t_map, target_vocab_size)),在 FastTopKSampler 构造、SpeculativeSampler 构造与 mapDraftProbsToTarget 三处统一调用,避免同一功能两套校验。若担心热路径 .item() 引入 D2H 同步,可在 SpeculativeSampler 构造期做一次性校验并缓存结果供 mapDraftProbsToTarget 复用。
  • FastTopKSampler 保留了只允许取值 1 的 top_k 参数,且 target_vocab_size=0 静默关闭全部校验 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.h:42
    • 建议:删除 top_k 形参(内部固定 top-1),或在声明处注明当前唯一合法值与预期演进方向,并把类名/注释语义收敛为「Top-1 提议采样器」。target_vocab_size 建议改为必填且构造期无条件校验 > 0(测试桩显式传真实 draft vocab),或改用 std::optional 明确表达「无映射」,避免哨兵 0 同时承担「无映射」与「关闭校验」两种含义。
  • CudaSpeculativeSamplingTest 用 -1 初始化输出,padding 断言无法区分写入与未写入 @ rtp_llm/models_py/bindings/cuda/test/CudaSpeculativeSamplingTest.cc:169
    • 建议:把该文件中 output_token_ids 的初值改为 kernel 不可能写出的哨兵(如 -7,与共享 harness 保持一致),使 padding 断言成为真实的写入校验。
  • 同一提议分布契约在五处复述且措辞不一,缺少单一权威来源 @ rtp_llm/cpp/cache/connector/p2p/P2PConnectorResourceStore.h:38
    • 建议:沿用 proto:610-611 已采用的引用式写法,统一为「以 GenerateRequestPB.propose_probs 为唯一权威契约来源」,其余各处只保留一行指向该处的引用而不复述内容,并把 dtype/形状要求补进权威处一次。
  • scatter_add_ 的广播索引开销未验证,原子累加使未来平滑提议分布失去可复现性 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:100
    • 建议:用 profiler 或 max_memory_allocated 确认广播索引未被 materialize,并把结论写进函数注释;若确实 materialize,改用支持 stride-0 索引的路径或按 target-index 预分组。同时在 :47-49 的扩展点注释旁补一句:一旦提议分布不再是 one-hot,多对一 d2t_mapatomicAdd 的累加顺序会破坏可复现性,需改用确定性归约。
  • dtype 转换使 hold_host 的 H2D 生命周期契约被绕过 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:131
    • 建议:先显式做 dtype 转换得到 host 张量、hold_host 该结果、再做纯设备迁移(to(device, /*non_blocking=*/true)),使被持有对象与实际参与异步拷贝的 buffer 一致;或在该处加一行注释说明 dtype 不一致时拷贝为同步、无需持有,使该保护的适用条件显式化。
  • runLiteralCase 的 12 个位置参数与裸字面量墙影响可维护性,kernel 输出断言不具备 fatal 语义 @ rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:45
    • 建议:把入参聚合为一个具名结构体(LiteralCase{batch_size, num_spec, vocab_size, stride, draft_probs, ...})并在调用点使用指定初始化器,使字段名随实参出现;同时给 :85-86 也套上 ASSERT_NO_FATAL_FAILURE,让首个失败即终止该测试,避免级联噪声掩盖首因。
  • top-1 点质量仍按全词表稠密 float32 上线传输,PD 字节翻倍且在锁内深拷贝 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:50
    • 建议:非阻塞优化建议:契约层面为 top-1 场景约定稀疏表示(仅传 token id,接收方本地展开为点质量),字段保留以支持未来 top-p 分布;或明确允许多种 TensorPB.data_type 由接收方转换(相应把 mapDraftProbsToTarget 的 float32 硬断言放宽为可恢复错误)。若暂不改协议,至少把 P2PConnector.cc:373 的深拷贝改为锁内 move/shared_ptr 交接以缩短锁持有时间,并给出该字段实测字节数与首 token 延迟的前后对比。

Checklist Findings (15 fail / 26 total)

General Principles Checklist

  • [6.1] Architecture — 兼容性:外部 HTTP/RPC API、持久数据、配置、环境迁移安全 → issue 契约注释把 draft 空间的 probs 与已映射为 target 空间的 propose_token_ids 绑定描述
    注释「Exact normalized distribution used to produce propose_token_ids, in draft-vocabulary space」把紧邻的 field 8 与 field 10 置于同一索引空间。实际二者不同:SpeculativeSampler.cc:50-51{batch, draft_vocab} 分配 all_probs 并按 draft id scatter_(draft 空间),:53-54 随后对 token_idsexecMappingDraft2Target 就地改写为 target id(mapping.cu:18 为绝对 id 替换),PrefillRpcServer.cc:397-398 发出的即映射后结果。field 8 上无任何 target-space 说明。按注释字面实现的消费方若用 propose_probs[propose_token_ids[i]] 取值,在 d2t_map 非恒等时取到错误概率、在 `draft_vocab < tar
  • [6.1] Architecture — 分层边界:新概念在正确层级,不泄漏内部 → issue 契约注释把 draft 空间的 probs 与已映射为 target 空间的 propose_token_ids 绑定描述
    注释「Exact normalized distribution used to produce propose_token_ids, in draft-vocabulary space」把紧邻的 field 8 与 field 10 置于同一索引空间。实际二者不同:SpeculativeSampler.cc:50-51{batch, draft_vocab} 分配 all_probs 并按 draft id scatter_(draft 空间),:53-54 随后对 token_idsexecMappingDraft2Target 就地改写为 target id(mapping.cu:18 为绝对 id 替换),PrefillRpcServer.cc:397-398 发出的即映射后结果。field 8 上无任何 target-space 说明。按注释字面实现的消费方若用 propose_probs[propose_token_ids[i]] 取值,在 d2t_map 非恒等时取到错误概率、在 `draft_vocab < tar
  • [6.1] Architecture — 可观测性:日志/指标/超时可操作、非噪声 → issue top-1 点质量仍按全词表稠密 float32 上线传输,PD 字节翻倍且在锁内深拷贝
    torch::zeros({batch_size, draft_vocab_size}, logits.options().dtype(torch::kFloat32)) 把提议分布固定为 fp32,而改动前 torch::softmax(logits, -1) 继承 logits dtype(本仓常见 bf16/fp16),同一字段在 PD 线上的字节数相比改前翻倍。当前唯一生产者恒为 one-hot(forward 强制 top_k==1),信息量可由 propose_token_ids 完全还原,却仍按 draft_vocab×4B/步 稠密序列化(PrefillRpcServer.cc:403-413)。P2P 侧信道进一步放大:P2PConnector.cc:373side_channel_mutex 时按值深拷贝整个 SideChannelData(含 propose_probs),:397CopyFrom 一次,等于每请求在锁内外各复制一份全词表张量,直接落在首 token 关键路径上。
  • [6.1] Architecture — 回滚路径:风险行为存在运维回滚手段 → issue propose_probs 取值语义与 dtype 静默变更,三处反序列化零校验且 same-build 条款不覆盖生成路径
    载荷由 softmax(logits)(继承 fp16/bf16)改为 zeros(..., kFloat32).scatter_(1, argmax, 1.0)SpeculativeSampler.cc:50-51),字段号/类型/形状不变,注释「Exact normalized distribution」对旧 softmax 同样成立,接收端无判别信号。三个入参点均不校验:PrefillLoadCaller.cc:211CopyFrom(同函数 propose_tokens(:202)、position_ids(:222)都校验 INT32),StreamCacheResource.cc 只判存在性,DecodeRpcServer.cc:293-297 直接 transTensor(...).to(kCUDA)。混版本 PD 下:单步 stackMtpBatchStreamProcessor.cc:919)保留 fp16 → 命中 mapDraftProbsToTarget 的 float32 硬断言,异常落在整批 `b
  • [6.1] Architecture — 状态不变量:创建/更新/失败/重试/回滚路径有效 → issue dtype 转换使 hold_host 的 H2D 生命周期契约被绕过
    :130-132 仅在非 CUDA 时 hold_host(draft_token_ids),随后 :133 执行 draft_token_ids.to(target_device, torch::kInt32, true)。当源为 CPU 且 dtype 不是 int32 时,ATen 会先在 CPU 上生成一个未被 hold 的临时张量再异步拷贝,buffer_holder_ 持有的是转换前的张量,真正参与 H2D 的 buffer 不在持有范围内,hold_host 的语义在该路径上失效。当前无实际风险:临时张量为 pageable 内存、cudaMemcpyAsync 对主机同步;且 updateOneStepDraftSamplerOutput 已提前 .to(kInt32),同 dtype 路径不产生临时张量。但「hold 住即安全」的约定因此变成依赖实现细节的巧合。
  • [6.1] Architecture — 错误语义:fail-fast/retry/fallback/silent 行为显式 → issue d2t_map 校验在三处不统一,作为写索引的 mapDraftProbsToTarget 反而不校验取值范围
    FastTopKSampler 构造期对 d2t_map 做完整校验(:16-21,含 min()>=0max()<target_vocab_size_),但真正把 d2t_map 当作写索引散射进 target 宽缓冲的 mapDraftProbsToTarget 只校验 dim/dtype/device/numel(:78-83),无值域校验;SpeculativeSampler 构造函数(SpeculativeSampler.h:54-55)对 d2t_map_ 完全不校验。两者的 target_vocab_size 还来自两条独立来源(前者构造期 vocab_size_,后者运行时 target_token_probs_d_t.size(2))。当前引擎路径安全(MtpExecutor.cc:641-642 把同一个 map 交给两者,ctor 会在初始化期 fail-fast),但 mapDraftProbsToTarget 已是 public static API 且被测试直接调用,越界索引在 CUDA `s
  • [6.1] Quality — PR description 说明动机与设计 → issue 随机采样行接受率系统性下降,缺少实测数据、接受率指标与运维回滚手段
    判据由 same_token || (do_sample && u*p<q) 改为 (do_sample && u*p<q) || (!do_sample && same_token)。对随机行,MTP top-1 提议的 p 恒为 1.0,接受概率退化为 q[draft_id],原先「verifier 与 draft 相同即无条件接受」的快路径消失;温度高、词表大的请求 accept_len 会明显缩短,残差采样分支(两遍 full-vocab block 归约)触发频率同步上升。配套把 do_sample 来源改为 generateConfig()->do_sample(默认 true)后走随机路径的行数进一步增加。diff 内既无 accept_len/TPS 前后对比,也无 accept-rate 指标,也没有可退回旧行为的 env/config 开关。
  • [6.1] Software Engineering — DRY:重复非平凡逻辑被抽取或显式复用 → issue runLiteralCase 的 12 个位置参数与裸字面量墙影响可维护性,kernel 输出断言不具备 fatal 语义
    runLiteralCase 接受 12 个位置参数,其中 3 个是 vector<float>draft_probs/uniform_samples/target_probs)、4 个是 vector<int32_t>draft_token_ids/target_token_ids/expected_token_ids/expected_accepted_token_num),调用点是数十行无标注的裸字面量(如 :168-:231),相邻同类型实参写反不会触发编译错误。另外 :67-68 对参考实现用了 ASSERT_NO_FATAL_FAILURE,但 :85-86 对 kernel 输出是裸 assertVectorEqual——该 helper 内部的 ASSERT_EQ(:527-:529)只从 helper 返回,因此首个 case 的 kernel 输出不匹配后,同一测试内后续 runLiteralCase 仍会继续执行。
  • [6.1] Software Engineering — ISP:调用方不依赖无关大接口 → issue FastTopKSampler 保留了只允许取值 1 的 top_k 参数,且 target_vocab_size=0 静默关闭全部校验
    forward(logits, top_k=1) 第一行即 RTP_LLM_CHECK_WITH_INFO(top_k == 1, ...)SpeculativeSampler.cc:25),参数只剩「必须传 1」一种合法取值;生产调用点 MtpExecutor.cc:799/:1166 不传、:1854 显式传 1,测试中唯一非 1 用法就是验证抛异常。该形参仍在虚函数签名上,所有子类(含测试 fake)必须重复它。同时新增的 target_vocab_size 用 0 作「不校验」哨兵::33target_vocab_size_ == 0 ? draft_vocab_size : target_vocab_size_ 使默认构造与 FastTopKSampler(t, 0)MtpExecutorTest.cc:196)都静默跳过 vocab 一致性校验,构造函数 5 处检查也因 d2t_map 未定义而全部 return,调用方漏传或误传 0 无任何提示。
  • [6.1] Software Engineering — KISS/YAGNI:无投机性抽象 → issue top-1 点质量仍按全词表稠密 float32 上线传输,PD 字节翻倍且在锁内深拷贝
    torch::zeros({batch_size, draft_vocab_size}, logits.options().dtype(torch::kFloat32)) 把提议分布固定为 fp32,而改动前 torch::softmax(logits, -1) 继承 logits dtype(本仓常见 bf16/fp16),同一字段在 PD 线上的字节数相比改前翻倍。当前唯一生产者恒为 one-hot(forward 强制 top_k==1),信息量可由 propose_token_ids 完全还原,却仍按 draft_vocab×4B/步 稠密序列化(PrefillRpcServer.cc:403-413)。P2P 侧信道进一步放大:P2PConnector.cc:373side_channel_mutex 时按值深拷贝整个 SideChannelData(含 propose_probs),:397CopyFrom 一次,等于每请求在锁内外各复制一份全词表张量,直接落在首 token 关键路径上。
  • [6.1] Tests — 分布式/跨平台变更有对应覆盖 → issue greedy 直出分支与混合 do_sample batch 缺字面期望值,ROCm 仅与同源参考实现互证
    runLiteralCase 的 5 处调用(:168、:235、:273、:293、:309)do_sample 全为 true。新增的贪心 direct_target_fallback 分支(含内联 -1 padding 循环)与混合 do_sample batch 只出现在 runReferenceCasesdo_sample{true,false,true,true},:113)与 runZeroAndOneSpeculativeTokenCases,而这两组期望值取自 referenceRejectionSampling 自身,该参考实现在同一提交内被同步改写(:560-615,与 kernel 逐分支同源),两侧同时写错时测试仍会通过。CUDA 侧有 CudaSpeculativeSamplingTest.cc:134-138/:194-198 的贪心字面断言兜底,但该文件依赖 CUDA runtime、不参与 ROCm 构建。
  • [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue MtpExecutorTest 两处夹具 all_probs 与 token_ids 自相矛盾,违反本 PR 建立的点质量契约
    本 PR 把 draft all_probs 重定义为「产生该 token 的精确分布」并启用了 checkInputs,但两处夹具未同步:(1) :714stream1_draft_token_probs={{0,0,1,0}} 点质量在 index 2,而 :716StreamSpecUpdateInfo 第三字段 draft_token=3(字段序见 GenerateStream.h:51-57),该夹具又被 :828 拼进期望张量(配 :827token_ids={3,2,1,3});(2) :1160token_ids={1,2}:1162{0,1,0,0, 0,1,0,0},batch 1 点质量落在 index 1 而非 2。同一次修复已改正 :803-:808 等多行,这两处被漏掉。测试仍通过(两侧读同一份夹具),但因此丧失对该契约的校验能力;真实链路中此类 p=0 会使 u*p=0<q 恒成立而无条件接受,正是本 PR 要修的缺陷。
  • [6.1] Tests — 边界 case 覆盖(空、单元素、最大值) → issue CudaSpeculativeSamplingTest 用 -1 初始化输出,padding 断言无法区分写入与未写入
    output_token_ids 在 :55、:108、:169 均以 torch::full(..., -1, intCuda()) 初始化,随后 :136-138、:198 断言尾部为 -1。初值与期望值相同,断言无法区分「kernel 显式写入 -1」与「kernel 根本没写」,而本 PR 恰好迁移了 -1 padding 的写入位置(全接受分支移除该循环、贪心 fallback 分支新增内联循环),弱化断言正落在被改动的行为上。共享 harness 已用 -7 哨兵(RejectionSamplingOpTest.hpp:77)解决同类问题。

RTP-LLM Checklist

  • [I] 代码质量 — 删除或重命名内部 file、registry entry、model name、metric enum、op binding、plugin symbol 时,必须全仓搜索消费者,并提供替代实现、迁移说明或 smoke 覆盖;只有暴露到 HTTP/RPC/config/persisted format 时才按外部兼容性处理 → issue 同文件遗留 ChainSpeculativeSampling 保留旧语义,形成语义分叉的死代码
    rejection_sampling_kernel 已修正接受判定与 bonus 语义,但同文件的 ChainSpeculativeSampling(:458,CUDA 侧对应 cuda/kernels/sampling/flashinfer/sampling.cuh:2040)仍是旧实现:无 do_sample 概念、pos == num_speculative_tokens 时从 target_probs 采样 bonus,并保留 min(pos+1, N)pos != num_speculative_tokens 等本次被删除的旧写法。全仓搜索 execChainSpeculativeSampling 仅有 ExecOps.h:97 声明与 ExecOps.cc:599 定义,无任何调用方,也未导出到 Python。两套语义分叉的 rejection sampling 长期共存,后续维护者无法判断哪份权威,也存在照搬旧写法的风险。
  • [I] 代码质量 — 同一功能用统一工具函数 → issue d2t_map 校验在三处不统一,作为写索引的 mapDraftProbsToTarget 反而不校验取值范围
    FastTopKSampler 构造期对 d2t_map 做完整校验(:16-21,含 min()>=0max()<target_vocab_size_),但真正把 d2t_map 当作写索引散射进 target 宽缓冲的 mapDraftProbsToTarget 只校验 dim/dtype/device/numel(:78-83),无值域校验;SpeculativeSampler 构造函数(SpeculativeSampler.h:54-55)对 d2t_map_ 完全不校验。两者的 target_vocab_size 还来自两条独立来源(前者构造期 vocab_size_,后者运行时 target_token_probs_d_t.size(2))。当前引擎路径安全(MtpExecutor.cc:641-642 把同一个 map 交给两者,ctor 会在初始化期 fail-fast),但 mapDraftProbsToTarget 已是 public static API 且被测试直接调用,越界索引在 CUDA `s

Strengths

  • 拒绝采样数学首次成立:点质量提议下 u*p<q 退化为「以 q[draft_id] 概率接受」,relu(q-p) 恰为正确残差分布;去掉 same_token || 短路消除了原先偏向 draft 一致性的分布偏差(旧实现把全 softmax 当作 argmax 提议的 p,属实质偏差)。贪心路径经新旧逐行比对与改前输出、accepted_token_num 完全一致,不引入 greedy 回归。
  • mapDraftProbsToTarget 修复两个真实缺陷:index_put_ 在 many-to-one d2t_map 下重复下标互相覆盖导致质量丢失;旧的「仅宽度不等才映射」条件使等宽非恒等映射被整段跳过,kernel 侧用 target id 索引 draft 空间的 p。MapsEqualWidthNonIdentityVocabMapMapsProposalDistributionWithManyToOneVocabMap(含 mapped.sum(-1)==draft_probs.sum(-1) 守恒断言)精准覆盖两者。
  • top_k != 1 改为 fail-fast 是实质修复:旧路径 torch::topk 返回 k 列,但 execMappingDraft2Target({..., stride=1}) 只映射首元素,k>1 是静默错误;TopKGreaterThanOneIsRejected 为被删用例提供等价替代。
  • 校验前移与归因改善:FastTopKSampler 构造期一次性校验 d2t_map 的 dim/dtype/值域(含 min()>=0max()<target_vocab_size),把配置类错误从每步热路径挪到引擎初始化,.item() 的 D2H 同步只发生一次;validateRejectionSamplingParams 改以 target_probs_d 为设备与词表基准,补齐 draft_probs_d 三维形状校验,并消除了原先「shape 来源即被校验方」的恒真断言。
  • 全接受行直接提前返回,省掉一整轮 vocab 级两遍 block scan(旧实现在「随机接受但 token 不同」时仍要跑残差采样求 bonus),且该 bonus 本身就是 target 分布的抽样,比额外残差采样更自洽。
  • runLiteralCase 先断言 CPU 参考实现等于手工字面期望、再断言 kernel 输出等于同一期望,打破了「kernel 与参考实现互证」的自证闭环;输出缓冲改用 -7 哨兵(:77)使「剩余位填 -1」成为可验证断言,正对应本 PR 迁移 -1 padding 写入位置的改动。已复算 vocab=3/5/8 全部 6 组用例(覆盖 VEC_SIZE=gcd(4,V) 的标量与向量化两条 dispatch),与新语义完全一致。
  • 跨平台一致性到位:ROCm kernel 与 CUDA 逐行等价(仅 load/cast_load 差异),新增 rocm/ops/tests:rejection_sampling_op_test 与 CUDA target 同构、复用同一 harness,tags/exec_properties 与同目录既有 target 约定一致;validateRejectionSamplingParams 位于 #if USING_CUDA 之前的匿名命名空间,两平台共用,故 4 个新增 draft_probs_d 负例跨平台生效。
  • FakeSpeculativeSampler::forward 现在真正调用此前是死代码的 checkInputs,期望张量由 2-D logits 拼接改为与生产一致的 3-D point-mass,多处夹具同步修正;FastTopKSamplerValidationTest 正确保存/恢复 user_ft_core_dump_on_exception,使 EXPECT_THROW(..., RTPException) 成立。
  • proto/头文件首次把此前散落多层的 propose_probs 语义写成显式契约,且严格保持 wire 与 ABI 兼容:字段号 10、TensorPB 枚举、reserved 名单、消息布局零变更,SideChannelPayloadPB 采用「以 GenerateRequestPB 为准」的单向引用而非平行复制,本部分回滚成本为零。

Comment thread rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc Outdated
@@ -27,4 +27,4 @@ cudaError_t invokeRejectionSampling(DType* draft_probs,
int num_speculative_tokens,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📍 实际位置 rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.h:17(不在 diff 展示范围内,就近挂载)

[P2] rejection kernel 的三条入参语义契约既未在 sampling.h 声明也无法 fail-fast

新语义下随机行的接受判定与残差分布完全由 q=target_probs[...]p=draft_probs[...+draft_id] 决定(sampling.cu:263-264),隐含三条硬前提:target_probs 必须是与实际抽样一致的重归一化分布;draft_probs 必须已完成 d2t 映射且在 draft_token_ids 处为正;target_token_ids 末列必须是该分布的真实抽样。三者均无声明——sampling.h 本次只改了文件末尾换行,签名与注释未变;validateRejectionSamplingParamsCudaSampleOp.cc:50-108)只校验 shape/dtype/device/contiguous。任一侧漏做映射时 p 变 0,判定退化为「q>0 即接受」,输出分布静默失真且无任何报错。

建议:invokeRejectionSampling 声明处写清上述三条契约,并注明 do_sample 选择的是两套不同算法而非同一算法的开关。同时建议在 op 层或 debug 构建下增加一条不变量校验(按行检查 draft_probs[draft_id] > 0 或行和接近 1),让契约被破坏时 fail-fast,而不是静默输出错误分布。


bool same_token = target_id == draft_id;
if (same_token || (do_sample[row_idx] && u * p < q)) {
if ((do_sample[row_idx] && u * p < q) || (!do_sample[row_idx] && same_token)) {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] 随机采样行接受率系统性下降,缺少实测数据、接受率指标与运维回滚手段

判据由 same_token || (do_sample && u*p<q) 改为 (do_sample && u*p<q) || (!do_sample && same_token)。对随机行,MTP top-1 提议的 p 恒为 1.0,接受概率退化为 q[draft_id],原先「verifier 与 draft 相同即无条件接受」的快路径消失;温度高、词表大的请求 accept_len 会明显缩短,残差采样分支(两遍 full-vocab block 归约)触发频率同步上升。配套把 do_sample 来源改为 generateConfig()->do_sample(默认 true)后走随机路径的行数进一步增加。diff 内既无 accept_len/TPS 前后对比,也无 accept-rate 指标,也没有可退回旧行为的 env/config 开关。

建议: 合入前给出 MTP do_sample=true 场景的 accept_len 与 TPS 前后对比数据,并在 PR description 中说明该吞吐变化及原因。若下降幅度超出预期,提供运维可控的降级开关(例如把旧 same_token 快路径保留为显式可选项);同时建议新增/复用接受长度指标,使灰度期具备对照与回滚判断依据。

Checklist: [6.1] PR description 说明动机与设计

Comment thread rtp_llm/cpp/model_rpc/proto/model_rpc_service.proto
Comment thread rtp_llm/cpp/model_rpc/proto/model_rpc_service.proto
Comment thread rtp_llm/cpp/cache/connector/p2p/P2PConnectorResourceStore.h
auto target_probs = buffer.narrow(0, 0, batch_size).narrow(1, 0, num_steps);
target_probs.zero_();
auto target_indices = d2t_map.view({1, 1, draft_vocab_size}).expand_as(draft_probs);
target_probs.scatter_add_(-1, target_indices, draft_probs);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P3] scatter_add_ 的广播索引开销未验证,原子累加使未来平滑提议分布失去可复现性

d2t_map.view({1,1,V}).expand_as(draft_probs) 产生 dim0/dim1 stride 为 0 的广播索引,其在 CUDA scatter_add_ 下是否被 materialize 成 [batch, num_spec, draft_vocab] 的 int64 临时张量未经验证——按大词表典型形状该临时张量可达数百 MB,与输出 buffer 同量级,属热路径上未量化的内存/带宽风险。另外 CUDA 上重复下标的 scatter_add_ 退化为 atomicAdd,浮点累加顺序不确定,与 :154-155「ensuring deterministic acceptance for reproducible iter_count」的目标冲突。当前无害(draft_probs 是精确 one-hot,1.0+0.0 位精确),但 :47-49 已明确预留 top-p sampler 复用该接口的扩展点,届时会影响 u*p<q 的边界决定。

建议: 用 profiler 或 max_memory_allocated 确认广播索引未被 materialize,并把结论写进函数注释;若确实 materialize,改用支持 stride-0 索引的路径或按 target-index 预分组。同时在 :47-49 的扩展点注释旁补一句:一旦提议分布不再是 one-hot,多对一 d2t_mapatomicAdd 的累加顺序会破坏可复现性,需改用确定性归约。

buffer_holder_.hold_host(draft_token_ids);
auto draft_token_ids_d_t = draft_token_ids.to(target_device, true);
if (!draft_token_ids.is_cuda()) {
buffer_holder_.hold_host(draft_token_ids);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P3] dtype 转换使 hold_host 的 H2D 生命周期契约被绕过

:130-132 仅在非 CUDA 时 hold_host(draft_token_ids),随后 :133 执行 draft_token_ids.to(target_device, torch::kInt32, true)。当源为 CPU 且 dtype 不是 int32 时,ATen 会先在 CPU 上生成一个未被 hold 的临时张量再异步拷贝,buffer_holder_ 持有的是转换前的张量,真正参与 H2D 的 buffer 不在持有范围内,hold_host 的语义在该路径上失效。当前无实际风险:临时张量为 pageable 内存、cudaMemcpyAsync 对主机同步;且 updateOneStepDraftSamplerOutput 已提前 .to(kInt32),同 dtype 路径不产生临时张量。但「hold 住即安全」的约定因此变成依赖实现细节的巧合。

建议: 先显式做 dtype 转换得到 host 张量、hold_host 该结果、再做纯设备迁移(to(device, /*non_blocking=*/true)),使被持有对象与实际参与异步拷贝的 buffer 一致;或在该处加一行注释说明 dtype 不一致时拷贝为同步、无需持有,使该保护的适用条件显式化。

Checklist: [6.1] 状态不变量:创建/更新/失败/重试/回滚路径有效

}
return cpu.to(torch::kCUDA);
}
void runLiteralCase(int batch_size,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P3] runLiteralCase 的 12 个位置参数与裸字面量墙影响可维护性,kernel 输出断言不具备 fatal 语义

runLiteralCase 接受 12 个位置参数,其中 3 个是 vector<float>draft_probs/uniform_samples/target_probs)、4 个是 vector<int32_t>draft_token_ids/target_token_ids/expected_token_ids/expected_accepted_token_num),调用点是数十行无标注的裸字面量(如 :168-:231),相邻同类型实参写反不会触发编译错误。另外 :67-68 对参考实现用了 ASSERT_NO_FATAL_FAILURE,但 :85-86 对 kernel 输出是裸 assertVectorEqual——该 helper 内部的 ASSERT_EQ(:527-:529)只从 helper 返回,因此首个 case 的 kernel 输出不匹配后,同一测试内后续 runLiteralCase 仍会继续执行。

建议: 把入参聚合为一个具名结构体(LiteralCase{batch_size, num_spec, vocab_size, stride, draft_probs, ...})并在调用点使用指定初始化器,使字段名随实参出现;同时给 :85-86 也套上 ASSERT_NO_FATAL_FAILURE,让首个失败即终止该测试,避免级联噪声掩盖首因。

Checklist: [6.1] DRY:重复非平凡逻辑被抽取或显式复用

// The probability tensor always describes the distribution that produced
// token_ids. Top-1 is a point mass; a future top-p sampler can return its
// filtered and normalized distribution through the same contract.
output.all_probs = torch::zeros({batch_size, draft_vocab_size}, logits.options().dtype(torch::kFloat32));

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P3] top-1 点质量仍按全词表稠密 float32 上线传输,PD 字节翻倍且在锁内深拷贝

torch::zeros({batch_size, draft_vocab_size}, logits.options().dtype(torch::kFloat32)) 把提议分布固定为 fp32,而改动前 torch::softmax(logits, -1) 继承 logits dtype(本仓常见 bf16/fp16),同一字段在 PD 线上的字节数相比改前翻倍。当前唯一生产者恒为 one-hot(forward 强制 top_k==1),信息量可由 propose_token_ids 完全还原,却仍按 draft_vocab×4B/步 稠密序列化(PrefillRpcServer.cc:403-413)。P2P 侧信道进一步放大:P2PConnector.cc:373side_channel_mutex 时按值深拷贝整个 SideChannelData(含 propose_probs),:397CopyFrom 一次,等于每请求在锁内外各复制一份全词表张量,直接落在首 token 关键路径上。

建议: 非阻塞优化建议:契约层面为 top-1 场景约定稀疏表示(仅传 token id,接收方本地展开为点质量),字段保留以支持未来 top-p 分布;或明确允许多种 TensorPB.data_type 由接收方转换(相应把 mapDraftProbsToTarget 的 float32 硬断言放宽为可恢复错误)。若暂不改协议,至少把 P2PConnector.cc:373 的深拷贝改为锁内 move/shared_ptr 交接以缩短锁持有时间,并给出该字段实测字节数与首 token 延迟的前后对比。

Checklist: [6.1] 可观测性:日志/指标/超时可操作、非噪声;[6.1] KISS/YAGNI:无投机性抽象

@LLLLKKKK LLLLKKKK left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AI Code Review - PR #1295

Status: LGTM

Summary: P0/0 · P1/0 · P2/8 · P3/10

Reviewed: commit d860b41dbdf9 · 2026-08-14 16:23 UTC+8

lgtm ready to ci

Non-blocking Suggestions

P2

  • do_sample 判据改用 generateConfig 后,正确性依赖跨四处文件的未声明重归一化不变量 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:145
    • 建议:在 do_sample_d 赋值处与 MtpExecutor.cc:1663 两侧互相加注释,说明「MTP 强制 DEFAULT 是随机接受判据成立的前提,改动此行必须同步复核 batchSample」。更稳妥的做法是由 SamplerOutput 透传 probs_are_renormalized 标记,batchSample 在不满足时 RTP_LLM_CHECK 报错,或把该 stream 的 do_sample_d 降级为 false 退回精确匹配验证,作为运维回滚路径。
  • rejection kernel 的入参语义契约既未在 sampling.h 声明也无法 fail-fast @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.h:17
    • 建议:在 sampling.hinvokeRejectionSampling 上方补齐三条契约的注释(target 空间索引、p(draft_id)>0do_sample 双分支语义),并在 validateRejectionSamplingParams 增加可选的调试期校验(例如 draft_probs_d.gather(-1, draft_token_ids_d.to(kLong)) 全大于 0),使违约在 op 边界以可读信息失败,而不是在内核里静默改变分布。同时为 OpData.hdo_sample_d 补一行语义注释。
  • 随机采样行接受率系统性下降,缺少实测数据与引擎级运维回滚手段 @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:268
    • 建议:在 PR description 补充随机采样场景下接受长度与端到端 TPS 的前后对比数据;同时提供一个进程级配置开关(例如允许把随机行强制按 !do_sample 的精确匹配路径验证,或把提案分布退化为 top-p 过滤后的真实分布),使线上出现吞吐劣化时可在不回滚二进制的前提下降级。建议在既有接受长度指标上按 do_sample 拆分维度,便于灰度对比。
  • 契约注释把 draft 空间的 probs 与已映射为 target 空间的 propose_token_ids 绑定描述 @ rtp_llm/cpp/model_rpc/proto/model_rpc_service.proto:311
    • 建议:拆开两字段的空间描述:propose_token_ids 已是经 d2t_map(绝对映射,非偏移映射)转换后的 target 词表 id;propose_probs 仍以 draft 词表 id 为索引,需经 d2t_map scatter 后才能与 target 空间对齐;并写明 data_type 必须为 FP32、单流 shape 为 [1, draft_vocab_size](decode 侧按步维 stack 为 3-D)。建议在 DecodeRpcServer / StreamCacheResource 填充 all_probs 处补 dtype 与 rank 校验并返回明确协议错误码,使 fail-fast 发生在协议层而非下游 RTP_LLM_CHECKP2PConnectorResourceStore.hPrefillLoadCaller.hOpData.hGenerateStream.h 的同类注释建议统一引用该处权威描述,避免六处复述措辞漂移。
  • batchSample 的三处语义变更零覆盖且恰被 fake 屏蔽,mapDraftProbsToTarget 用例全在 CPU @ rtp_llm/cpp/normal_engine/speculative/test/SpeculativeSamplerTest.cc:45
    • 建议:补 batchSample 层测试:构造 top_k=1/do_sample=truetop_k>1/do_sample=truedo_sample=false 三类 stream,断言 do_sample_d 取值与 int32 转换结果;若难以直接驱动,可先把 do_sample_d 的推导抽成输入 GenerateConfig、输出 bool 的纯函数再做参数化测试。同时为 mapDraftProbsToTarget 增加一组 CUDA 用例(多对一映射 + buffer 复用)与一条无 d2t_map 的直通用例,使 GPU 原子累加、device 校验与直通分支进入回归网。
  • greedy 直出分支与混合 do_sample batch 缺字面期望值,ROCm 仅与同源参考实现互证 @ rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:113
    • 建议:在 RejectionSamplingOpTest.hpp 增加至少一条 do_sample=false 且含中途 mismatch 的 runLiteralCase(硬编码 expected_token_ids 与 accepted_num),使 greedy 直出语义在 CUDA 与 ROCm 双平台都具备独立 oracle;同时为混合 do_sample 的 batch 补一组固定断言,覆盖同批内贪心行与随机行互不干扰。
  • 全接受路径的 bonus token 来源变更在全部新旧用例上均不可判别 @ rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:168
    • 建议:在字面用例中构造至少一行「全接受 + bonus 位 target_token_ids 值 ≠ argmax(target_probs[bonus_pos])」的数据(例如该位 target_probs[0.1,0.1,0.8] 而 verifier 采样结果给 1),使「bonus 必须取自 target sampler」这一新契约具备真实回归能力。
  • MtpExecutorTest 单 batch 夹具的 all_probs 与 draft_token 自相矛盾,未编码本 PR 建立的点质量契约 @ rtp_llm/cpp/normal_engine/speculative/test/MtpExecutorTest.cc:714
    • 建议:把 stream1_draft_token_probs 改为 {{0.0f,0.0f,0.0f,1.0f}},与 draft_token = 3 / combo_tokens = {3} 对齐,使单 batch 夹具与多 batch 夹具在同一不变量下保持一致,从而对未来破坏点质量语义的改动具备真实回归能力;若该槽位的 draft 分布确属 don't care,请在夹具处加注释说明,避免后续维护者误以为它编码了契约。

P3

  • 残差分布和为 0 时静默回退到词表末位 token @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:335
    • 建议:在残差分支对 sum_relu_q_minus_p 增加调试期 assert 或一次性告警,并在 sampling.h 的函数注释中显式声明所依赖的输入不变量;同时补一条 sum==0 的边界用例固化预期行为。属非阻塞健壮性观察,正常路径不受影响。
  • FastTopKSampler 保留只允许取值 1 的 top_k 形参,且 target_vocab_size=0 静默关闭全部校验 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.h:42
    • 建议:删除 top_k 形参(连同 MtpExecutor.cc:1854 的显式实参),让「只支持 top-1 提案」成为编译期契约;若确要扩展到 top-p 提案(:47-49 注释提到该意图),改为传入提案策略枚举而非恒为 1 的整型。target_vocab_size == 0 建议改为显式的 std::optional 或独立的「无映射」构造语义,避免用魔法值表达「跳过校验」。
  • d2t_map 校验在三处不统一,作为写索引的 mapDraftProbsToTarget 反而不校验取值范围 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:78
    • 建议:把取值域校验收敛到唯一权威位置:抽一个共享的 validateD2tMap(d2t_map, target_vocab_size)(1-D、int64、min>=0max<target_vocab_sizenumel==draft_vocab_size)供三处复用;GPU 张量上的 min/max 开销可通过只在构造期做一次范围检查、映射期仅做形状/设备检查并在注释中声明该前置条件来规避。
  • scatter_add_ 的原子累加使多对一映射在未来平滑分布下失去可复现性,且未判定单射/恒等映射 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:100
    • 建议:在构造期一次性判定 d2t_map 是否单射/恒等(unique().numel()arange 比较)并缓存:恒等时直通、单射时走确定性的 index_copy_/gather、非单射时在启动日志或注释中显式声明该场景放弃逐位可复现性。这样既省掉热路径上的一次全词表清零与 scatter,也为后续平滑提案分布预先划清确定性边界。
  • mapDraftProbsToTarget 以 public static 加原始指针出参暴露内部缓冲管理细节 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.h:57
    • 建议:把映射逻辑收为 speculative 命名空间内的自由函数或私有实例方法(直接使用成员缓冲),公开接口只保留 forward;若为可测性必须保留独立入口,则在复用判定前补一条 RTP_LLM_CHECK_WITH_INFO(!buffer.defined() || buffer.dim() == 3, ...),与该函数其余入参校验风格保持一致。
  • top-1 点质量仍按全词表稠密张量上线传输,与 propose_token_ids 完全信息冗余 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:50
    • 建议:对 top-1 提议考虑不再传输稠密分布:或仅传被提议 token 的概率标量,或直接省略 propose_probs,由 decode 侧依据 propose_token_ids 就地重建 one-hot(顺带消除协议层的 dtype/词表空间歧义)。若因后续 top-p 提案规划必须保留稠密张量,请在 proto 注释中记录该带宽代价与 FP32 约定,便于容量评估。
  • do_sample 默认为 true 使 top_k=1 请求由精确匹配切到随机验证路径 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:145
    • 建议:在 do_sample_d 赋值处补注释说明该行为矩阵变化,并考虑在 top_k==1 时把 do_sample_d 收敛为 false——此时随机判据与精确匹配数学等价,却能省掉失配路径上的全词表残差采样;如保留现状,请在接受率/耗时指标中区分这两类流量以便观测。
  • FakeSpeculativeSampler 的输入断言变为机会式,期望未被消费也不会失败 @ rtp_llm/cpp/normal_engine/speculative/test/MtpExecutorTest.cc:229
    • 建议:去掉该守卫恢复无条件校验(与 FakeFastTopKSampler 一致),或在测试 TearDown 中断言 input_holder.test_data.empty(),把「期望被全部消费」显式检查出来,而不是依赖调用次数恰好匹配。
  • MtpExecutorTest 中 catVectors 辅助函数变为死代码 @ rtp_llm/cpp/normal_engine/speculative/test/MtpExecutorTest.cc:116
    • 建议:删除 catVectors 模板定义以保持测试文件整洁;后续若再需要拼接工具,在使用点重新引入即可。
  • 同文件遗留的 ChainSpeculativeSampling 是不支持 do_sample 的重复实现且无在库调用方 @ rtp_llm/models_py/bindings/rocm/speculative_sampling/sampling.cu:458
    • 建议:若确认该路径已废弃,连同 SpeculativeSamplingParamschainSpeculativeSamplingexecChainSpeculativeSampling 与对应 pybind 导出一并删除(CUDA/ROCm 双侧),避免两份能力不同的拒绝采样实现长期共存;若需保留作为参考,请在函数上方注明「deprecated,不支持 do_sample 贪心验证,勿新增调用方」。

Checklist Findings (16 fail / 26 total)

General Principles Checklist

  • [6.1] Architecture — 兼容性:外部 HTTP/RPC API、持久数据、配置、环境迁移安全 → issue 契约注释把 draft 空间的 probs 与已映射为 target 空间的 propose_token_ids 绑定描述
    新注释写作「用于产生 propose_token_ids 的精确归一化分布,位于 draft 词表空间」,但两字段不在同一索引空间:PrefillRpcServer.cc:397propose_token_ids 取自 getProposeToken(),而 FastTopKSampler::forward 返回前已调用 execMappingDraft2Target(SpeculativeSampler.cc:54)把 token 映射为 target id(CudaSampleOp.cc:500 为绝对映射 tokens[i]=d2t_map[token_id]);propose_probs(:413 取自 all_probs)仍以 draft id 为索引。对端若按注释字面用前者索引后者,在 d2t_map 非恒等时取到错误分量且不报错。注释也未声明代码已硬断言的 dim()==3kFloat32(SpeculativeSampler.cc:64-66),而 DecodeRpcServer.cc:297、`StreamC
  • [6.1] Architecture — 分层边界:新概念在正确层级,不泄漏内部 → issue mapDraftProbsToTarget 以 public static 加原始指针出参暴露内部缓冲管理细节
    mapDraftProbsToTarget 声明为 public static,并以 torch::Tensor* target_probs_buffer = nullptr 接收外部缓冲,唯一的生产调用方是同类的 batchSample(SpeculativeSampler.cc:169,传入自身成员 draft_probs_target_buffer_),其余调用方全部是测试。这把「缓冲区复用策略」这一实现细节提升到了类的公开接口。同时 :89-92 的复用判定依次访问 buffer.size(0)/size(1)/size(2),只有 defined() 短路保护而无维度校验,外部传入一个已定义的 1-D 或 2-D 张量时会抛 ATen 下标越界错误,而非该函数其他检查那样的可读信息。
  • [6.1] Architecture — 可观测性:日志/指标/超时可操作、非噪声 → issue do_sample 默认为 true 使 top_k=1 请求由精确匹配切到随机验证路径
    GenerateConfig::do_sample 默认 true(GenerateConfig.h:92)而 top1() 只看 top_k == 1(:128-129)。对显式设置 top_k=1 但保留 do_sample 默认值的请求,do_sample_d 由旧的 false 变为 true:结果仍等价(qtop_k_renorm_probs 归一为 argmax 点质量,u*1<q 退化为精确匹配,失配时 relu(q-p) 亦为 argmax 处的 one-hot),但内核不再走 direct_target_fallback 的提前 return,每次失配都要额外执行两遍全词表的 block reduce/scan 残差采样。反向地,do_sample=falsetop_k=0 的请求由旧的随机路径改为精确匹配,属于改善。
  • [6.1] Architecture — 回滚路径:风险行为存在运维回滚手段 → issue 随机采样行接受率系统性下降,缺少实测数据与引擎级运维回滚手段
    do_sample=truetop_k>1/top_p<1 的行,旧判据为 same_token || u*p_softmax(draft_id) < qp_softmax < 1,等价 u < q/p,且精确匹配免检);新判据为 u*1 < q,即 u < q(draft_id)。因 p_new = 1 ≥ p_old 且免检通道被移除,接受概率严格下降,MTP 的 token/step 与吞吐随之回落。这是换取分布无偏所必须的代价,但 PR 未给出任一模型/温度档位的接受率与 TPS 实测对比。接受长度指标本身已存在(MtpExecutor.cc:1577-1590total_accept_len/spec_steps),故缺的是灰度数据与开关:目前唯一降级手段是逐请求设 do_sample=false,没有进程级或配置级回滚。
  • [6.1] Architecture — 状态不变量:创建/更新/失败/重试/回滚路径有效 → issue do_sample 默认为 true 使 top_k=1 请求由精确匹配切到随机验证路径
    GenerateConfig::do_sample 默认 true(GenerateConfig.h:92)而 top1() 只看 top_k == 1(:128-129)。对显式设置 top_k=1 但保留 do_sample 默认值的请求,do_sample_d 由旧的 false 变为 true:结果仍等价(qtop_k_renorm_probs 归一为 argmax 点质量,u*1<q 退化为精确匹配,失配时 relu(q-p) 亦为 argmax 处的 one-hot),但内核不再走 direct_target_fallback 的提前 return,每次失配都要额外执行两遍全词表的 block reduce/scan 残差采样。反向地,do_sample=falsetop_k=0 的请求由旧的随机路径改为精确匹配,属于改善。
  • [6.1] Architecture — 错误语义:fail-fast/retry/fallback/silent 行为显式 → issue d2t_map 校验在三处不统一,作为写索引的 mapDraftProbsToTarget 反而不校验取值范围
    同一份 d2t_map 在三处被校验且强度不同:FastTopKSampler ctor 校验 1-D、int64 与 [0, target_vocab_size) 取值域(:16-21);mapDraftProbsToTarget 只校验 1-D、int64、device、numel(:78-83),不校验取值范围SpeculativeSampler ctor 完全不校验(.h:54-55)。而恰恰是 mapDraftProbsToTarget 把 d2t_map 当作 scatter_add_写索引(:100),越界值会触发 ATen/CUDA 的 device-side assert(黏性 CUDA 错误、进程级崩溃),而非该函数其他检查那样的可读 RTP_LLM_CHECK 信息。当前只因 MtpExecutor.cc:641-642 用同一张量先后构造两个对象、FastTopKSampler ctor 会先抛错,才间接获得保护。
  • [6.1] Quality — PR description 说明动机与设计 → issue 随机采样行接受率系统性下降,缺少实测数据与引擎级运维回滚手段
    do_sample=truetop_k>1/top_p<1 的行,旧判据为 same_token || u*p_softmax(draft_id) < qp_softmax < 1,等价 u < q/p,且精确匹配免检);新判据为 u*1 < q,即 u < q(draft_id)。因 p_new = 1 ≥ p_old 且免检通道被移除,接受概率严格下降,MTP 的 token/step 与吞吐随之回落。这是换取分布无偏所必须的代价,但 PR 未给出任一模型/温度档位的接受率与 TPS 实测对比。接受长度指标本身已存在(MtpExecutor.cc:1577-1590total_accept_len/spec_steps),故缺的是灰度数据与开关:目前唯一降级手段是逐请求设 do_sample=false,没有进程级或配置级回滚。
  • [6.1] Software Engineering — DRY:重复非平凡逻辑被抽取或显式复用 → issue 同文件遗留的 ChainSpeculativeSampling 是不支持 do_sample 的重复实现且无在库调用方
    ChainSpeculativeSampling(:458-580)与本次修改的 rejection_sampling_kernel 在同一文件内并存,重复实现了 relu(q-p) 残差采样与 block 级归约。需修正一处常见误读:它的判据本就是纯 u*p<q(:481),并非旧的 same_token || 有偏判据;真正的差距是它没有 do_sample 形参、无 greedy 分支、无 direct_target_fallback(:490-496),也未随本 PR 更新。其调用链 chainSpeculativeSampling → execChainSpeculativeSampling(ExecOps.cc:599-601)在全仓无任何 C++ 调用方,sampling_ops.cu:4 的 pybind 导出亦无 Python 消费者。未来若有人接线,会拿到一个不支持贪心验证的内核。
  • [6.1] Software Engineering — ISP:调用方不依赖无关大接口 → issue mapDraftProbsToTarget 以 public static 加原始指针出参暴露内部缓冲管理细节
    mapDraftProbsToTarget 声明为 public static,并以 torch::Tensor* target_probs_buffer = nullptr 接收外部缓冲,唯一的生产调用方是同类的 batchSample(SpeculativeSampler.cc:169,传入自身成员 draft_probs_target_buffer_),其余调用方全部是测试。这把「缓冲区复用策略」这一实现细节提升到了类的公开接口。同时 :89-92 的复用判定依次访问 buffer.size(0)/size(1)/size(2),只有 defined() 短路保护而无维度校验,外部传入一个已定义的 1-D 或 2-D 张量时会抛 ATen 下标越界错误,而非该函数其他检查那样的可读信息。
  • [6.1] Software Engineering — KISS/YAGNI:无投机性抽象 → issue MtpExecutorTest 中 catVectors 辅助函数变为死代码
    本 PR 在 testSingleBatchDecode / testMultiBatchDecode 中删除了 draft_all_probs_listcatVectors(...) 的调用,改用 torch::cat 直接构造 3-D 概率。全仓检索 catVectors 仅剩定义处(MtpExecutorTest.cc:116),无任何调用点;同文件的 toVec / checkTensorEqual 仍在使用。因其为模板函数,未实例化不会触发 -Wunused-function,不影响编译,但已成为无用残留。
  • [6.1] Software Engineering — SRP:模块/类职责单一 → issue mapDraftProbsToTarget 以 public static 加原始指针出参暴露内部缓冲管理细节
    mapDraftProbsToTarget 声明为 public static,并以 torch::Tensor* target_probs_buffer = nullptr 接收外部缓冲,唯一的生产调用方是同类的 batchSample(SpeculativeSampler.cc:169,传入自身成员 draft_probs_target_buffer_),其余调用方全部是测试。这把「缓冲区复用策略」这一实现细节提升到了类的公开接口。同时 :89-92 的复用判定依次访问 buffer.size(0)/size(1)/size(2),只有 defined() 短路保护而无维度校验,外部传入一个已定义的 1-D 或 2-D 张量时会抛 ATen 下标越界错误,而非该函数其他检查那样的可读信息。
  • [6.1] Tests — 分布式/跨平台变更有对应覆盖 → issue greedy 直出分支与混合 do_sample batch 缺字面期望值,ROCm 仅与同源参考实现互证
    新增字面用例 runStochasticSemanticsCases(.hpp:165-267)与 runPointMassDraftCases(:269-338)的 do_sample 全为 true。含 do_sample=false 的只有 runReferenceCases(:113 的 {true,false,true,true})与 runZeroAndOneSpeculativeTokenCases(:352),二者仅与 referenceRejectionSampling 对比,而该参考实现本次已随内核同步改写(:586 新判据、:590-593 的 direct_target_fallback、:600-608 的分流),同源逻辑错误无法暴露。因此本 PR 引入的「greedy 首个失配位直出 verifier token 并停止残差采样」在共享 harness 中没有独立手算 oracle;独立硬编码期望仅存在于 CUDA 专属的 CudaSpeculativeSamplingTest.cc(本次仅改注释、期望值未变),ROC
  • [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue FakeSpeculativeSampler 的输入断言变为机会式,期望未被消费也不会失败
    if (!input_holder.test_data.empty()) { checkInputs(...); }TestDataHolder::get() 会 pop 队列,因此该守卫的实际含义是「还有剩余期望时才校验」。对比同文件 FakeFastTopKSampler::forward(:198-201)无条件调用 checkInputs,期望耗尽时由 get() 抛异常暴露调用次数不符;当前实现下若 executor 调用 speculative sampler 的次数少于 setInputs 排入的期望数,剩余期望永远不被校验、测试静默通过,且 TearDown 也未断言队列已清空。
  • [6.1] Tests — 边界 case 覆盖(空、单元素、最大值) → issue 残差分布和为 0 时静默回退到词表末位 token
    残差采样在 sum_relu_q_minus_p == 0u = uniform * 0 = 0aggregate > u 恒不成立、循环不会 break,temp_storage.sampled_id 保持初值 target_vocab_size - 1(:335),即静默输出词表最后一个 token。我验证了正常输入下不可达:p 为 one-hot(p(draft_id)=1),拒绝要求 u >= q(draft_id),而 relu(q-p) 之和为 1 - q(draft_id)torch::rand 取值在 [0,1),故 q(draft_id)=1 时不可能被拒绝,拒绝时该和必为正。但不可达性完全依赖「q、p 均在 target_vocab_size 上归一化且 p(draft_id)>0」这一未声明前提。ROCm 侧 sampling.cu:404 逻辑相同。

RTP-LLM Checklist

  • [I] 代码质量 — 删除或重命名内部 file、registry entry、model name、metric enum、op binding、plugin symbol 时,必须全仓搜索消费者,并提供替代实现、迁移说明或 smoke 覆盖;只有暴露到 HTTP/RPC/config/persisted format 时才按外部兼容性处理 → issue 同文件遗留的 ChainSpeculativeSampling 是不支持 do_sample 的重复实现且无在库调用方
    ChainSpeculativeSampling(:458-580)与本次修改的 rejection_sampling_kernel 在同一文件内并存,重复实现了 relu(q-p) 残差采样与 block 级归约。需修正一处常见误读:它的判据本就是纯 u*p<q(:481),并非旧的 same_token || 有偏判据;真正的差距是它没有 do_sample 形参、无 greedy 分支、无 direct_target_fallback(:490-496),也未随本 PR 更新。其调用链 chainSpeculativeSampling → execChainSpeculativeSampling(ExecOps.cc:599-601)在全仓无任何 C++ 调用方,sampling_ops.cu:4 的 pybind 导出亦无 Python 消费者。未来若有人接线,会拿到一个不支持贪心验证的内核。
  • [I] 代码质量 — 同一功能用统一工具函数 → issue d2t_map 校验在三处不统一,作为写索引的 mapDraftProbsToTarget 反而不校验取值范围
    同一份 d2t_map 在三处被校验且强度不同:FastTopKSampler ctor 校验 1-D、int64 与 [0, target_vocab_size) 取值域(:16-21);mapDraftProbsToTarget 只校验 1-D、int64、device、numel(:78-83),不校验取值范围SpeculativeSampler ctor 完全不校验(.h:54-55)。而恰恰是 mapDraftProbsToTarget 把 d2t_map 当作 scatter_add_写索引(:100),越界值会触发 ATen/CUDA 的 device-side assert(黏性 CUDA 错误、进程级崩溃),而非该函数其他检查那样的可读 RTP_LLM_CHECK 信息。当前只因 MtpExecutor.cc:641-642 用同一张量先后构造两个对象、FastTopKSampler ctor 会先抛错,才间接获得保护。

Strengths

  • 数学上正确:top-1 提案的真实分布是点质量而非 softmax,接受概率 min(1,q/p) 退化为 q(x)、拒绝后从 relu(q-p) 重采样,目标分布无偏;改用 argmax + one-hot 同时省掉一次全词表 softmax kernel。
  • mapDraftProbsToTarget 改为「有 d2t_map 就一定映射」,修掉旧实现仅在宽度不等时才 padding 的静默错误;MapsEqualWidthNonIdentityVocabMap(SpeculativeSamplerTest.cc:56)精确覆盖此前出错的等宽非恒等场景。
  • scatter_add_ 而非 index_put_,使多对一 d2t_map 下同一 target token 的提案概率被正确求和,并以 mapped.sum(-1) == draft_probs.sum(-1) 断言质量守恒。
  • 校验前置且分层完整:ctor 校验 d2t_map 的 1-D/int64/取值域(:16-21),映射期校验 dim/dtype/device/长度(:63-83),op 层把锚点从 draft_probs_d 改到 target_probs_d 并新增三个维度一致性检查(CudaSampleOp.cc:51-86),把旧代码的静默假设换成 fail-fast。
  • 缓冲复用要求 size(1) == num_steps 精确匹配而非 grow-only,保证 narrow 结果始终连续,与 checkRejectionSamplingTensoris_contiguous() 要求对齐,并由 ReallocatesForDifferentStepCountToKeepOutputContiguous 专门覆盖。
  • greedy 路径改用显式 direct_target_fallback(sampling.cu:275-281),在首个失配位输出 verifier token 并 pad,与旧 all_same_token 隐式路径逐位等价(accept_len = pos+1 口径不变,CudaSpeculativeSamplingTest 期望值未变);三类路径均写满 num_spec+1 个输出位,无未初始化残留。
  • CUDA 与 ROCm 内核改动逐点一致(判据/fallback/bonus/残差下标),并新增 ROCm rejection_sampling_op_test,经 rejection_sampling_test_utilsselect() 链接 HIP 实现,与 CUDA 侧达成用例 parity。
  • runLiteralCase 先断言参考实现等于手算期望、再断言内核等于同一期望,提供独立 oracle,规避「参考模型随内核同源改动导致漏检」;runStochasticSemanticsCases 精准覆盖「verifier ID 相等但仍被拒绝」这一旧 bug 的精确触发点。
  • buffer_holder_.hold_host(draft_token_ids) 增加 !is_cuda() 前置判断,修掉旧代码把设备张量塞进 host holder 的问题。
  • proto 改动严格限于注释:字段号 10/11、SideChannelPayloadPB 的 map 结构与 reserved 1 to 8(:624-625)均未触碰,wire 层二进制兼容。

Comment thread rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc Outdated
@@ -27,4 +27,4 @@ cudaError_t invokeRejectionSampling(DType* draft_probs,
int num_speculative_tokens,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📍 实际位置 rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.h:17(不在 diff 展示范围内,就近挂载)

[P2] rejection kernel 的入参语义契约既未在 sampling.h 声明也无法 fail-fast

invokeRejectionSampling(sampling.h:17-29)为裸指针签名,本次仅补了文件末尾换行、无任何注释。本 PR 为该内核建立三条隐式契约:draft_probs 必须是经 d2t 映射到 target 空间、shape [batch, num_spec, target_vocab] 的精确归一化提案分布;draft_token_ids 同样已映射到 target 空间;且提案 token 处 p>0。三条均无法在内核内检测:sampling.cu:264 取 p = draft_probs[...+draft_id],若 p == 0(提案 token 与提案分布来自不一致的上游)则 u*p = 0 < q 恒成立 → 无条件接受,静默产生分布偏置。OpData.h:352 新增的一行注释只覆盖 draft_probs_ddo_sample_d 无任何语义说明。

建议:sampling.hinvokeRejectionSampling 上方补齐三条契约的注释(target 空间索引、p(draft_id)>0do_sample 双分支语义),并在 validateRejectionSamplingParams 增加可选的调试期校验(例如 draft_probs_d.gather(-1, draft_token_ids_d.to(kLong)) 全大于 0),使违约在 op 边界以可读信息失败,而不是在内核里静默改变分布。同时为 OpData.hdo_sample_d 补一行语义注释。


bool same_token = target_id == draft_id;
if (same_token || (do_sample[row_idx] && u * p < q)) {
if ((do_sample[row_idx] && u * p < q) || (!do_sample[row_idx] && same_token)) {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] 随机采样行接受率系统性下降,缺少实测数据与引擎级运维回滚手段

do_sample=truetop_k>1/top_p<1 的行,旧判据为 same_token || u*p_softmax(draft_id) < qp_softmax < 1,等价 u < q/p,且精确匹配免检);新判据为 u*1 < q,即 u < q(draft_id)。因 p_new = 1 ≥ p_old 且免检通道被移除,接受概率严格下降,MTP 的 token/step 与吞吐随之回落。这是换取分布无偏所必须的代价,但 PR 未给出任一模型/温度档位的接受率与 TPS 实测对比。接受长度指标本身已存在(MtpExecutor.cc:1577-1590total_accept_len/spec_steps),故缺的是灰度数据与开关:目前唯一降级手段是逐请求设 do_sample=false,没有进程级或配置级回滚。

建议: 在 PR description 补充随机采样场景下接受长度与端到端 TPS 的前后对比数据;同时提供一个进程级配置开关(例如允许把随机行强制按 !do_sample 的精确匹配路径验证,或把提案分布退化为 top-p 过滤后的真实分布),使线上出现吞吐劣化时可在不回滚二进制的前提下降级。建议在既有接受长度指标上按 do_sample 拆分维度,便于灰度对比。

Checklist: [6.1] 回滚路径:风险行为存在运维回滚手段;[6.1] PR description 说明动机与设计

repeated string peer_addrs = 7;
repeated int32 propose_token_ids = 8;
repeated int32 position_ids = 9;
// Exact normalized distribution used to produce propose_token_ids, in

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] 契约注释把 draft 空间的 probs 与已映射为 target 空间的 propose_token_ids 绑定描述

新注释写作「用于产生 propose_token_ids 的精确归一化分布,位于 draft 词表空间」,但两字段不在同一索引空间:PrefillRpcServer.cc:397propose_token_ids 取自 getProposeToken(),而 FastTopKSampler::forward 返回前已调用 execMappingDraft2Target(SpeculativeSampler.cc:54)把 token 映射为 target id(CudaSampleOp.cc:500 为绝对映射 tokens[i]=d2t_map[token_id]);propose_probs(:413 取自 all_probs)仍以 draft id 为索引。对端若按注释字面用前者索引后者,在 d2t_map 非恒等时取到错误分量且不报错。注释也未声明代码已硬断言的 dim()==3kFloat32(SpeculativeSampler.cc:64-66),而 DecodeRpcServer.cc:297、`Stre...

建议: 拆开两字段的空间描述:propose_token_ids 已是经 d2t_map(绝对映射,非偏移映射)转换后的 target 词表 id;propose_probs 仍以 draft 词表 id 为索引,需经 d2t_map scatter 后才能与 target 空间对齐;并写明 data_type 必须为 FP32、单流 shape 为 [1, draft_vocab_size](decode 侧按步维 stack 为 3-D)。建议在 DecodeRpcServer / StreamCacheResource 填充 all_probs 处补 dtype 与 rank 校验并返回明确协议错误码,使 fail-fast 发生在协议层而非下游 RTP_LLM_CHECKP2PConnectorResourceStore.hPrefillLoadCaller.hOpData.hGenerateStream.h 的同类注释建议统一引用该处权威描述,避免六处复述措辞漂移。

Checklist: [6.1] 兼容性:外部 HTTP/RPC API、持久数据、配置、环境迁移安全

EXPECT_TRUE(torch::allclose(out.all_probs, expected));
}

TEST(SpeculativeSamplerTest, MapsProposalDistributionWithManyToOneVocabMap) {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] batchSample 的三处语义变更零覆盖且恰被 fake 屏蔽,mapDraftProbsToTarget 用例全在 CPU

新增用例只覆盖 FastTopKSampler::forward 与静态 mapDraftProbsToTarget(SpeculativeSamplerTest.cc:27-108),batchSample 无任何测试入口;而 MtpExecutorTestFakeSpeculativeSampler(:222-233)与 FakeFastTopKSampler(:194-201)替换了真实实现。因此本次行为切换面最大的三行——do_sample_d = generateConfig()->do_sample(:145)、draft_token_ids.to(target_device, kInt32, true)(:133)、mapDraftProbsToTargetdraft_probs_target_buffer_ 的生产接线(:169)——零覆盖。此外 7 个映射用例全部由 torch::tensor(...) 构造、运行在 CPU,scatter_add_ 的 CUDA 原子路径与 :80-81 的 devi...

建议:batchSample 层测试:构造 top_k=1/do_sample=truetop_k>1/do_sample=truedo_sample=false 三类 stream,断言 do_sample_d 取值与 int32 转换结果;若难以直接驱动,可先把 do_sample_d 的推导抽成输入 GenerateConfig、输出 bool 的纯函数再做参数化测试。同时为 mapDraftProbsToTarget 增加一组 CUDA 用例(多对一映射 + buffer 复用)与一条无 d2t_map 的直通用例,使 GPU 原子累加、device 校验与直通分支进入回归网。

// The probability tensor always describes the distribution that produced
// token_ids. Top-1 is a point mass; a future top-p sampler can return its
// filtered and normalized distribution through the same contract.
output.all_probs = torch::zeros({batch_size, draft_vocab_size}, logits.options().dtype(torch::kFloat32));

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P3] top-1 点质量仍按全词表稠密张量上线传输,与 propose_token_ids 完全信息冗余

torch::zeros({batch_size, draft_vocab_size}, ...kFloat32)scatter_(1, ..., 1.0f) 产出的张量只有一个非零值,信息与同一 GenerateRequestPB 中的 propose_token_ids(field 8)完全等价,却在 PD 分离路径按 draft_vocab * 4 字节整张传输(PrefillRpcServer.cc:403-413 的 D2H + transTensorPB),并在 P2P 侧信道中做 protobuf 深拷贝。需澄清:dtype 发生变化——PyWrappedModel.cc:976/980/1016/1020 两条分支均产出 FP32 logits,旧的 torch::softmax(logits, -1) 本就是 FP32,且 CudaSampleOp 一直要求 draft_probs_d 为 FP32,故不存在「BF16/FP16 变 FP32、字节翻倍」的兼容性回归,本条仅为既有带宽开销的澄清。

建议: 对 top-1 提议考虑不再传输稠密分布:或仅传被提议 token 的概率标量,或直接省略 propose_probs,由 decode 侧依据 propose_token_ids 就地重建 one-hot(顺带消除协议层的 dtype/词表空间歧义)。若因后续 top-p 提案规划必须保留稠密张量,请在 proto 注释中记录该带宽代价与 FP32 约定,便于容量评估。

Comment thread rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc Outdated
spec::SpeculativeSamplerOutput forward(const std::list<GenerateStreamPtr>& streams,
SamplerOutput& draft_sampler_output,
SamplerOutput& target_sampler_output) override {
if (!input_holder.test_data.empty()) {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P3] FakeSpeculativeSampler 的输入断言变为机会式,期望未被消费也不会失败

if (!input_holder.test_data.empty()) { checkInputs(...); }TestDataHolder::get() 会 pop 队列,因此该守卫的实际含义是「还有剩余期望时才校验」。对比同文件 FakeFastTopKSampler::forward(:198-201)无条件调用 checkInputs,期望耗尽时由 get() 抛异常暴露调用次数不符;当前实现下若 executor 调用 speculative sampler 的次数少于 setInputs 排入的期望数,剩余期望永远不被校验、测试静默通过,且 TearDown 也未断言队列已清空。

建议: 去掉该守卫恢复无条件校验(与 FakeFastTopKSampler 一致),或在测试 TearDown 中断言 input_holder.test_data.empty(),把「期望被全部消费」显式检查出来,而不是依赖调用次数恰好匹配。

Checklist: [6.1] 新逻辑有聚焦单测 + 相关集成/smoke 测试

@@ -193,7 +193,7 @@ class FakeModel: public ModelBase {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📍 实际位置 rtp_llm/cpp/normal_engine/speculative/test/MtpExecutorTest.cc:116(不在 diff 展示范围内,就近挂载)

[P3] MtpExecutorTest 中 catVectors 辅助函数变为死代码

本 PR 在 testSingleBatchDecode / testMultiBatchDecode 中删除了 draft_all_probs_listcatVectors(...) 的调用,改用 torch::cat 直接构造 3-D 概率。全仓检索 catVectors 仅剩定义处(MtpExecutorTest.cc:116),无任何调用点;同文件的 toVec / checkTensorEqual 仍在使用。因其为模板函数,未实例化不会触发 -Wunused-function,不影响编译,但已成为无用残留。

建议: 删除 catVectors 模板定义以保持测试文件整洁;后续若再需要拼接工具,在使用点重新引入即可。

Checklist: [6.1] KISS/YAGNI:无投机性抽象

output_token_ids[row_idx * (num_speculative_tokens + 1) + pad_idx] = -1;
}
}
}

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📍 实际位置 rtp_llm/models_py/bindings/rocm/speculative_sampling/sampling.cu:458(不在 diff 展示范围内,就近挂载)

[P3] 同文件遗留的 ChainSpeculativeSampling 是不支持 do_sample 的重复实现且无在库调用方

ChainSpeculativeSampling(:458-580)与本次修改的 rejection_sampling_kernel 在同一文件内并存,重复实现了 relu(q-p) 残差采样与 block 级归约。需修正一处常见误读:它的判据本就是纯 u*p<q(:481),并非旧的 same_token || 有偏判据;真正的差距是它没有 do_sample 形参、无 greedy 分支、无 direct_target_fallback(:490-496),也未随本 PR 更新。其调用链 chainSpeculativeSampling → execChainSpeculativeSampling(ExecOps.cc:599-601)在全仓无任何 C++ 调用方,sampling_ops.cu:4 的 pybind 导出亦无 Python 消费者。未来若有人接线,会拿到一个不支持贪心验证的内核。

建议: 若确认该路径已废弃,连同 SpeculativeSamplingParamschainSpeculativeSamplingexecChainSpeculativeSampling 与对应 pybind 导出一并删除(CUDA/ROCm 双侧),避免两份能力不同的拒绝采样实现长期共存;若需保留作为参考,请在函数上方注明「deprecated,不支持 do_sample 贪心验证,勿新增调用方」。

Checklist: [6.1] DRY:重复非平凡逻辑被抽取或显式复用;[I] 删除或重命名内部 file、registry entry、model name、metric enum、op binding、plugin symbol 时,必须全仓搜索消费者,并提供替代实现、迁移说明或 smoke 覆盖;只有暴露到 HTTP/RPC/config/persisted format 时才按外部兼容性处理

@LLLLKKKK LLLLKKKK left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

AI Code Review - PR #1295

Status: LGTM

Summary: P0/0 · P1/0 · P2/12 · P3/9

Reviewed: commit 43e153b2a4bf · 2026-08-18 17:19 UTC+8

lgtm ready to ci

Non-blocking Suggestions

P2

  • 随机流接受率因判据收紧与点质量提案系统性下降,缺少实测数据与运维回滚手段 @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:268
    • 建议:补充 top1()==false 场景的 accept_len / TPS / 输出质量前后对比数据(至少一个 MTP 模型),并在 PR description 中写明这是有偏采样修复及其预期吞吐代价;同时确认现有接受长度指标能按 do_sample 维度区分,便于上线对比。建议评估更优的等价修法:让随机流的 draft token 从 softmax(draft_logits) 采样而非 argmax,同时保留 p=softmax,期望接受率为最大耦合 Σ min(p,q),通常远高于 q(argmax_p)。不建议长期保留双套采样语义的 feature flag。
  • 契约注释把 draft 空间的 probs 与已映射为 target 空间的 propose_token_ids 绑定描述 @ rtp_llm/cpp/model_rpc/proto/model_rpc_service.proto:319
    • 建议:注释中明确区分两个空间:propose_token_ids 已经过 d2t 映射、位于 target 词表空间;propose_probs 位于 draft 词表空间,消费方必须先用本地 d2t_map 映射(参见 SpeculativeSampler::mapDraftProbsToTarget)后才能与 propose_token_ids 配对,禁止直接用前者索引后者。同一契约目前以四种措辞散落于 proto:319-320、GenerateStream.h:93-94PrefillLoadCaller.h:29P2PConnectorResourceStore.h:38,建议统一以 GenerateRequestPB.propose_probs 为唯一权威定义、其余三处改为引用(与 SideChannelPayloadPB(proto:619)已采用的写法一致),避免后续语义调整产生文档漂移。
  • 新增契约遗漏 FP32 与 rank 硬约束,跨节点接收边界零校验 @ rtp_llm/cpp/cache/connector/p2p/PrefillLoadCaller.h:29
    • 建议:在三处注释中补齐硬约束:data_type 必须为 FP32、shape 为 [batch, propose_step, draft_vocab_size](单请求发送为 [1, draft_vocab_size])。并在接收边界(PrefillLoadCaller 解析处与 StreamCacheResource::applyP2PSideChannel 落盘前)补齐与 propose_tokens 一致的校验:dtype 非 FP32 或最后一维不等于本地 draft_vocab_size 时,带 request_id 与实际 dtype/shape 返回明确的 ErrorCodePB,而不是把非法 payload 带到采样阶段 abort。
  • 所援引的 same-build 契约作用域未覆盖 generate RPC 的采样语义变更 @ rtp_llm/cpp/model_rpc/proto/model_rpc_service.proto:320
    • 建议:二选一:(1) 把 proto:28-33 的契约段落显式扩写覆盖「proposal 分布语义」,并在 :319-320 按名引用该契约块而非泛指,同时说明混版表现是接受率静默偏移而非报错,便于运维判断回滚条件;(2) 在 payload 中增加 proposal 分布语义版本标量(或复用已有 build/protocol version),decode 侧检测不匹配时 fail-fast 或退回非投机路径。无论哪种,都请在 PR description / 发布说明中写明 prefill 与 decode 必须同批升级(或先 drain 再切换);也可在 decode 侧上报「收到的 propose_probs 最大值是否为 1」的指标使混版可被监控发现。
  • d2t_map 取值域校验放在了不使用它作索引的类里,mapDraftProbsToTarget 越界退化为 device assert @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:78
    • 建议:把取值域校验下沉到真正消费它的位置:抽出统一的 validateD2TMap(d2t_map, target_vocab_size)FastTopKSamplerSpeculativeSampler 两个构造函数复用(在构造期做一次性 min/max,避免每 forward .item() 同步),失败时抛出带 d2t_map min/max 与 target_vocab_size 的可读错误,使「init 期 fail-fast」成为类自身的不变量而非调用顺序的巧合。
  • kernel 入参语义契约未在 sampling.h 声明,bonus 来源变更也无独立字面用例 @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.h:17
    • 建议:(1)在 sampling.h 的声明上方补参数契约注释,明确:draft_probs 必须是 target 词表空间的精确归一化提案分布;target_probs 必须与产生 target_token_ids 的分布一致(若上游做过 top-k/top-p 过滤,此处必须是重归一化后的分布);target_token_idsdo_sample=true 行必须来自真随机采样;uniform_samples[i] 用于位置 i 的接受判定而 [pos+1] 用于残差采样。ROCm 侧同步。(2)新增一组 do_sample=true、全部接受、bonus 位 target_probs[N] 为非退化多峰分布且 target_token_ids[N] 明确不等于残差采样命中值的字面用例(可照搬 runReferenceCases batch 3 的 {0.10,0.20,0.30,0.25,0.15} + uniform=0.55),独立锁定「bonus 来自 target sampler」契约。
  • mapDraftProbsToTarget 单测全在 CPU,且无 d2t_map 的成功路径与直通分支完全无覆盖 @ rtp_llm/cpp/normal_engine/speculative/test/SpeculativeSamplerTest.cc:45
    • 建议:(1)为 4 个映射用例增加 CUDA 变体(对输入 .to(torch::kCUDA) 即可复用同一份期望),mapped.sum(-1) 归一性断言在 atomics 累加下改用带 tolerance 的比较;补一个「映射结果直接喂给 execRejectionSampling」的用例,确保返回视图同时满足 device/contiguous/shape 三项 op 契约。(2)补回无 map 成功路径:FastTopKSampler sampler(torch::Tensor(), 4)[1,4] logits,断言 token_ids 未被重映射且 all_probs 为点质量;mapDraftProbsToTarget(draft_probs, torch::Tensor(), draft_vocab) 断言返回张量与输入同 data_ptr。(3)补构造期负例:d2t_map 二维、非 int64、含负值、max >= target_vocab_size
  • do_sample=false 贪心直出分支无字面期望值,ROCm 侧只能与同源参考实现互证 @ rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:165
    • 建议:新增 runGreedySemanticsCases(或在现有函数内)补 1~2 条 do_sample=falserunLiteralCase,硬编码期望覆盖:(1)逐位与 verifier 一致直至 bonus;(2)首个不一致位直接输出 verifier token、accepted_token_num = pos + 1、其后补 -1;(3)混合 do_sample 的 batch。在 RocmRejectionSamplingOpTest.ccCudaRejectionSamplingOpTest.cc 同步注册,让贪心语义在两个平台都有不依赖参考模型的字面锚点。
  • 新增字面用例的 vocab 全部落在单轮 chunk 路径,改动过的多轮 chunk 逻辑无覆盖 @ rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:269
    • 建议:补一条大 vocab 用例(例如 vocab_size = 8192,使 VEC_SIZE=4 且至少 2 轮 chunk):把 target 概率质量放在最后一个 chunk 内的非首 lane,让提案 token 落在第一个 chunk 并被拒绝,同时验证跨 chunk 的 aggregate 累积与 sampled_idatomicMin 归约。CUDA 与 ROCm 共用同一 helper,一次补齐覆盖两个平台。
  • MtpExecutorTest 夹具的 all_probs 与 draft token 自相矛盾,且被 checkOutput 固化为 golden @ rtp_llm/cpp/normal_engine/speculative/test/MtpExecutorTest.cc:714
    • 建议:把 :714 改为 index 3 的点质量 {0,0,0,1}(或把 draft_token 改为 2);把 :1161-1162 第二行改为 {0.0f, 0.0f, 1.0f, 0.0f} 并同步 :1217 的第三个参数。同时建议在 checkOutput(:381-405)中补一条不变量断言(argmax(all_probs) == 该 stream 的下一个 propose token),把本 PR 刚强化的点质量契约固化下来,而不是继续以裸字面量表达、任其漂移。
  • runLiteralCase 的 12 个位置参数使相邻同类型向量可被静默错位 @ rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:45
    • 建议:引入具名字段的入参结构体(如 struct LiteralCase { int batch_size; ...; std::vector<float> draft_probs; std::vector<float> uniform_samples; std::vector<float> target_probs; ... };),调用方用指定初始化器 .uniform_samples = {...} 传入,使参数含义在调用点自解释、误换变成编译期错误;同时把扁平数组按 vocab_size 分行排版,让切分关系在源码上可见。
  • ROCm 侧 invokeRejectionSampling 无条件返回 hipSuccess,吞掉 kernel launch 失败 @ rtp_llm/models_py/bindings/rocm/speculative_sampling/sampling.cu:654
    • 建议:顺手在 DISPATCH_ALIGNED_VEC_SIZE 之后补 return hipGetLastError();(或改用与 CUDA 侧对齐的 launch 宏),使 rejectionSampling 现有的 RTP_LLM_CHECK_WITH_INFO 真正生效,让 CUDA/ROCm 的错误语义一致为 fail-fast。

P3

  • 多对一映射下 scatter_add_ 原子累加削弱 random_seed 的可复现性 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:100
    • 建议:明确记录该限制(例如注释说明「确定性复现仅在 d2t_map 为单射时成立」),或在需要严格复现的场景采用确定性累加(按 target id 预分组后 index_add / segment reduce),避免同一 seed 下 accept_len 抖动。
  • 插入 dtype 转换后 hold_host 不再保护真实的 H2D 暂存缓冲 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:133
    • 建议:要么在上游保证进入 batchSample 的 draft token 恒为 int32、此处仅做 RTP_LLM_CHECK 断言 dtype 而非静默转换(并补一行注释说明该假设);要么显式先在 CPU 上 to(kInt32) 并 pin、对转换结果调用 hold_host 后再做纯 device 搬运,使 TensorHolder 的保护对象与异步拷贝源一致。
  • SpeculativeSampler 有虚函数但缺少虚析构函数 @ rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.h:52
    • 建议:补充 virtual ~SpeculativeSampler() = default;,与同文件 FastTopKSampler 保持一致。
  • 删除 pos != num_speculative_tokens 守卫后 draft_probs 越界安全完全依赖上游 early-return @ rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:320
    • 建议:在 int pos = s_pos; 之后补一行 assert(pos < num_speculative_tokens)(debug 构建生效)或显式注释,把「进入残差块即 pos < num_speculative_tokens」这个不变量写死;ROCm 侧同步。
  • 内核级单测缺少 do_sample=true 拒绝分支且用例名与新语义不符 @ rtp_llm/models_py/bindings/cuda/test/CudaSpeculativeSamplingTest.cc:90
    • 建议:(1)把两个贪心用例重命名为体现 do_sample=false 语义的名字(如 GreedyVerify_ImmediateMismatchEmitsVerifierTokenGreedyVerify_PartialAccept)。(2)补一条 do_sample=true 的 kernel 直调用例:令 target_token_ids[0]==draft_token_ids[0]target_probs[0][draft_id] 足够小以触发拒绝,断言输出不等于 draft token 且等于 relu(q-p) 的唯一非零位置。若认为分层职责上该文件只做 launch 冒烟,也请在文件头显式注明随机语义由 op 级测试负责。
  • runReferenceCases 中为旧 bonus 位残差采样准备的 fixture 已成死数据 @ rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:130
    • 建议:删除 :130 与 :133 这两处已失效的 fixture 设置,或把 batch 3 改造成在某一位真正被拒绝的场景(令 u * p >= q),使该 batch 自身重新覆盖随机残差采样路径;同时在函数头补一行注释说明各 batch 分别对应哪条语义分支,避免后续维护者误以为 :130/:133 仍在覆盖 bonus 位残差采样。
  • Fake 采样器在期望队列耗尽后静默跳过输入校验 @ rtp_llm/cpp/normal_engine/speculative/test/MtpExecutorTest.cc:229
    • 建议:改为「一旦某用例设置过期望,就要求调用次数与期望数严格一致」:例如新增显式的 expect_inputs_ 开关,或在 TearDown 中断言 input_holder.test_data.empty(),确保推入的期望全部被消费。
  • MtpExecutorTest 中 catVectors 在删除调用点后成为死代码 @ rtp_llm/cpp/normal_engine/speculative/test/MtpExecutorTest.cc:116
    • 建议:删除 catVectors 模板定义。若后续仍需拼接 vector<vector<T>>,届时再引入更贴合当时需求的写法,避免留存无调用方的辅助函数。
  • 点质量提案仍按全词表稠密张量上线传输,且 P2PConnector 的 FP32 空值哨兵恒失效 @ rtp_llm/cpp/cache/connector/p2p/P2PConnectorResourceStore.h:38
    • 建议:(1)后续可在 payload 中增加「退化分布」标记(只传 token id + 标志位),decode 侧按标记直接构造 point mass,既保留未来 top-p 提案传稠密分布的能力,又省掉跨机带宽与显存清零开销;建议先在注释中记录该演进方向,避免实现被 one-hot 语义锁死。(2)把 P2PConnector.cc:392 的空值判断改为统一复用 tensorPbHasPayload 语义,去掉不成立的 dtype 哨兵。

Checklist Findings (15 fail / 26 total)

General Principles Checklist

  • [6.1] Architecture — 兼容性:外部 HTTP/RPC API、持久数据、配置、环境迁移安全 → issue 点质量提案仍按全词表稠密张量上线传输,且 P2PConnector 的 FP32 空值哨兵恒失效
    注释把 propose_probs 定义为「Exact normalized proposal distribution in draft-vocabulary space」。在当前唯一的 top-1 提案实现下该张量恒为 one-hot:prefill 侧每请求需序列化并跨机传输 float32 × draft_vocab(十万级词表数百 KB),decode 侧再解析、H2D、参与 zero_ + scatter_add_,实际信息量只有一个 token id。wire format 未变故非本次引入的回归,但注释把这一冗余表示固化成了协议契约。同时 P2PConnector.cc:392 的空值判断以 data_type() != TensorPB::FP32 为首个条件,而 proto3 默认值 FP32 = 0(proto:7),未设置的 TensorPB 该条件恒为 false,空 payload 也会被塞入 tensors map。
  • [6.1] Architecture — 分层边界:新概念在正确层级,不泄漏内部 → issue 契约注释把 draft 空间的 probs 与已映射为 target 空间的 propose_token_ids 绑定描述
    新注释写作「Exact normalized distribution used to produce propose_token_ids, in draft-vocabulary space」,把二者绑定为同一索引空间。但实现中空间不同:FastTopKSampler::forward 在 draft 空间构造 all_probs(宽度 logits.size(1)),随后对 token_ids 调用 execMappingDraft2Target 映射为 target id(SpeculativeSampler.cc:50-54)。新增单测正好证实差异:d2t_map={0,1,3,2},argmax=2 → token_ids=3,而点质量仍在 index 2(SpeculativeSamplerTest.cc:28-42)。d2t_map 非恒等时 propose_probs[propose_token_ids[i]] 取错值甚至越界。
  • [6.1] Architecture — 可观测性:日志/指标/超时可操作、非噪声 → issue 随机流接受率因判据收紧与点质量提案系统性下降,缺少实测数据与运维回滚手段
    判据由 same_token || (do_sample && u*p<q) 变为按 do_sample 二分。对 do_sample=true!generateConfig()->top1(),SpeculativeSampler.cc:146)的行,旧行为是「verifier argmax 命中即无条件接受」或 min(1, q/p_softmax);叠加 all_probs 改为 one-hot(p 由 softmax 峰值 0.8~0.99 升为 1.0),接受概率退化为 q(draft_id)。两项均单调降低接受率。ROCm sampling.cu:336 同步变更。代码与测试中无接受率/TPS 数据,kernel 亦无 env/config 开关。
  • [6.1] Architecture — 回滚路径:风险行为存在运维回滚手段 → issue 所援引的 same-build 契约作用域未覆盖 generate RPC 的采样语义变更
    注释称「PD peers follow the same-build protocol contract」,但本文件唯一的契约段落(proto:28-33)标题为「KV-cache protocol compatibility contract」,正文明确限定适用于「peers that exchange KV-cache RPCs」,动机是分组拓扑与独立 block pool 改变了 block 寻址与 payload 结构,并未覆盖 GenerateRequestPB 的采样分布语义。本 PR 把 propose_probs 由全 softmax 分布改为 one-hot,而字段号 10、TensorPB 类型、张量形状均未变,接收端无从区分。滚动升级期(旧 prefill → 新 decode)不报错,但 p 仍为 0.x 且 same_token 兜底已删除,接受判据被系统性放大、输出分布有偏,无日志或指标能暴露。
  • [6.1] Architecture — 状态不变量:创建/更新/失败/重试/回滚路径有效 → issue 删除 pos != num_speculative_tokens 守卫后 draft_probs 越界安全完全依赖上游 early-return
    draft_probs 形状为 (batch, N, vocab),改动前两处 p_vec.load 都被 if (pos != num_speculative_tokens) 包裹以防 pos==N 时读到不存在的第 N 个 step(旧流程中随机全接受行确实会以 pos==N 进入该块)。现在守卫被移除,唯一保护是 :304 的 if (s_direct_target_fallback || s_pos == num_speculative_tokens) return;。当前逻辑成立,但该不变量与越界读之间隔了十余行且无断言或注释固化;后续若在 early-return 条件上再加分支,很容易在最后一行 batch 上产生越界读且不易复现。ROCm sampling.cu:389 同样处理。
  • [6.1] Architecture — 错误语义:fail-fast/retry/fallback/silent 行为显式 → issue ROCm 侧 invokeRejectionSampling 无条件返回 hipSuccess,吞掉 kernel launch 失败
    ROCm 的 invokeRejectionSampling<<<...>>> 启动 kernel 后直接 return hipSuccess;(sampling.cu:639-655),既未调用 hipGetLastError(),也未对动态共享内存做能力检查(BLOCK_THREADS=1024smem_size 可能超限)。CUDA 侧走 cudaLaunchKernel 并抛出 launch 错误。调用方 rejectionSamplingRTP_LLM_CHECK_WITH_INFO(err == hipSuccess, ...) 因此永不触发。这是既有问题,但本 PR 首次为该 ROCm kernel 引入 op 级测试:若 launch 静默失败,测试只会因输出仍为 -7 预填充值而报「结果不符」,错误指向语义而非 launch 失败。
  • [6.1] Quality — PR description 说明动机与设计 → issue 随机流接受率因判据收紧与点质量提案系统性下降,缺少实测数据与运维回滚手段
    判据由 same_token || (do_sample && u*p<q) 变为按 do_sample 二分。对 do_sample=true!generateConfig()->top1(),SpeculativeSampler.cc:146)的行,旧行为是「verifier argmax 命中即无条件接受」或 min(1, q/p_softmax);叠加 all_probs 改为 one-hot(p 由 softmax 峰值 0.8~0.99 升为 1.0),接受概率退化为 q(draft_id)。两项均单调降低接受率。ROCm sampling.cu:336 同步变更。代码与测试中无接受率/TPS 数据,kernel 亦无 env/config 开关。
  • [6.1] Quality — 逻辑变更未混入无关格式化 → issue runReferenceCases 中为旧 bonus 位残差采样准备的 fixture 已成死数据
    逐 batch 手算:batch 3(do_sample=true)三步全部接受(i=0: u*p=0.20*0.50=0.10 < q=0.80;i=1: 0.5*0.2 < 0.2;i=2: 0.10*0.50=0.05 < 0.90),因此 pos==N,bonus 直接取 target_token_ids[3]=4,不再进入残差采样。专为旧「pos==N 时从 relu(q-0) 采样 bonus」路径准备的 :130 setProbRow(target_probs, ..., 3, 3, {0.10,...}) 与 :133 uniform_samples[3*4+3]=0.55f 已不影响任何输出。(batch 2 在 i=1 处 u*p=0.35 >= q=0.10 被拒,仍会走残差采样,:124 的 uniform_samples[2*4+2]=0.90f 依然有效。)
  • [6.1] Software Engineering — KISS/YAGNI:无投机性抽象 → issue MtpExecutorTest 中 catVectors 在删除调用点后成为死代码
    本 PR 把 testSingleBatchDecodetestMultiBatchDecode 中构造 draft_all_probs_list 的代码改为 torch::cat(:828-832、:1190-1195),catVectors 的两处(也是全部)调用点随之消失。全仓检索确认 catVectors 仅剩 :116 一处定义、零引用。因其为函数模板,未实例化不会触发 -Wunused-function,所以不会打断构建,只是留下无人使用的测试辅助模板。
  • [6.1] Software Engineering — LSP:子类/重写保持基类契约 → issue SpeculativeSampler 有虚函数但缺少虚析构函数
    同文件的 FastTopKSampler 声明了 virtual ~FastTopKSampler() {}(:43),而 SpeculativeSampler(:52-82)有 virtual forward(:62)却没有虚析构函数。MtpExecutorstd::unique_ptr<speculative::SpeculativeSampler> 持有对象(MtpExecutor.h:207),测试中注入的 FakeSpeculativeSampler(含两个 TestDataHolder 成员,MtpExecutorTest.cc:255-257)经基类指针销毁属 UB,派生类成员的析构不会执行。本次修改了该类的构造函数并新增 static 方法但未一并补齐,两个相邻类的对称性缺口更明显。
  • [6.1] Tests — 分布式/跨平台变更有对应覆盖 → issue ROCm 侧 invokeRejectionSampling 无条件返回 hipSuccess,吞掉 kernel launch 失败
    ROCm 的 invokeRejectionSampling<<<...>>> 启动 kernel 后直接 return hipSuccess;(sampling.cu:639-655),既未调用 hipGetLastError(),也未对动态共享内存做能力检查(BLOCK_THREADS=1024smem_size 可能超限)。CUDA 侧走 cudaLaunchKernel 并抛出 launch 错误。调用方 rejectionSamplingRTP_LLM_CHECK_WITH_INFO(err == hipSuccess, ...) 因此永不触发。这是既有问题,但本 PR 首次为该 ROCm kernel 引入 op 级测试:若 launch 静默失败,测试只会因输出仍为 -7 预填充值而报「结果不符」,错误指向语义而非 launch 失败。
  • [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue Fake 采样器在期望队列耗尽后静默跳过输入校验
    if (!input_holder.test_data.empty()) { checkInputs(...); }(:229-231):test_dataqueueget() 会 pop。因此当 forward 被调用次数超过 setInputs 推入的期望数时,超出部分不再校验且不报错。这一守卫是为兼容未设置期望的旧用例,但也意味着「期望数少于实际调用数」的用例会静默通过——本次新接入的 draft 概率契约校验强度因此取决于测试作者是否恰好补齐了每一步期望。
  • [6.1] Tests — 被删除测试有等价替代覆盖 → issue mapDraftProbsToTarget 单测全在 CPU,且无 d2t_map 的成功路径与直通分支完全无覆盖
    4 个映射用例(:45/:56/:66/:81)全部使用 CPU 张量且全部传入 d2t_map,而生产路径中 draft_probs / d2t_map 均在 CUDA 上:CPU 的 scatter_add_ 串行累加、CUDA 走 atomics;d2t_map.device() 校验与返回视图的 is_cuda()/is_contiguous() op 约束(CudaSampleOp.cc:36/43)都是设备相关的。BUILD 已设 TEST_USING_DEVICE=CUDA 并链接 cuda_impl,却无 CUDA 用例。同时原 TopKOneReturnsArgmaxIndex(默认构造、无 map)被替换为必带 map 的 :27,无 map 的两个新用例(:96/:103)只断言抛异常——而「d2t_map 未定义且 draft vocab == target vocab」正是不做词表重映射的默认部署形态,mapDraftProbsToTarget 的零拷贝直通分支(:70-76)亦无覆盖。
  • [6.1] Tests — 边界 case 覆盖(空、单元素、最大值) → issue 新增字面用例的 vocab 全部落在单轮 chunk 路径,改动过的多轮 chunk 逻辑无覆盖
    残差采样按 ceil_div(target_vocab_size, BLOCK_THREADS * VEC_SIZE) 分 chunk 迭代(sampling.cu:314、341),BLOCK_THREADS=1024VEC_SIZE=gcd(16/sizeof(float), vocab)=gcd(4, vocab)≤4(:399),只有 vocab > 4096 才出现多轮迭代与跨 chunk 的 aggregate_relu_q_minus_p 累积。新增用例 vocab 为 3/3/5/5/8(:168、:235、:273、:293、:309),runReferenceCases 为 5,runZeroAndOne... 为 4,全部只跑 1 轮。而本 PR 恰好改动了两个 chunk 循环体内的 p_vec.load(:320、:347)与 u 的下标(:338),生产 vocab 约 10^5 走的正是未覆盖的多轮路径。

RTP-LLM Checklist

  • [I] 代码质量 — 同一功能用统一工具函数 → issue 点质量提案仍按全词表稠密张量上线传输,且 P2PConnector 的 FP32 空值哨兵恒失效
    注释把 propose_probs 定义为「Exact normalized proposal distribution in draft-vocabulary space」。在当前唯一的 top-1 提案实现下该张量恒为 one-hot:prefill 侧每请求需序列化并跨机传输 float32 × draft_vocab(十万级词表数百 KB),decode 侧再解析、H2D、参与 zero_ + scatter_add_,实际信息量只有一个 token id。wire format 未变故非本次引入的回归,但注释把这一冗余表示固化成了协议契约。同时 P2PConnector.cc:392 的空值判断以 data_type() != TensorPB::FP32 为首个条件,而 proto3 默认值 FP32 = 0(proto:7),未设置的 TensorPB 该条件恒为 false,空 payload 也会被塞入 tensors map。

Strengths

  • 修掉一个真实隐藏 bug:旧实现仅在 draft.size(2) != target.size(2) 时才重映射概率,而 token_ids 始终经 execMappingDraft2Target 映射;等宽非恒等 d2t_map 下 kernel 以 target id 索引 draft 空间概率,p 完全错位。新增 MapsEqualWidthNonIdentityVocabMap(SpeculativeSamplerTest.cc:56-64)正好锁住该场景。
  • 接受判据回归教科书式 rejection sampling:旧判据在 do_sample=true 时用 same_token 强制接受,使输出分布偏向 draft;sampling.cu:268 与 CPU 参考 RejectionSamplingOpTest.hpp:586 完全同构。经逐行推演确认贪心(do_sample=false)路径输出与改动前严格等价,语义收紧只作用于随机行。
  • 多对一 d2t_map 由 index_put_(后写覆盖、丢失质量、重复索引未定义)改为 scatter_add_(质量累加),并用 mapped.sum(-1) == draft_probs.sum(-1) 断言归一性(SpeculativeSamplerTest.cc:53)。
  • buffer 复用条件由 grow-only 改为精确匹配 size(1) == num_steps(SpeculativeSampler.cc:89),消除 narrow(1, ...) 产生非连续视图触发 op 层 is_contiguous 校验失败的隐患,并有 ReallocatesForDifferentStepCountToKeepOutputContiguous(:81-94)直接覆盖。
  • 重量级校验(min()/max().item() 需设备同步)放在构造函数而非 forward(SpeculativeSampler.cc:19-21),热路径与 CUDA graph 捕获未引入新的 host 同步点。
  • CUDA(sampling.cu:227-378)与 ROCm(sampling.cu:296-449)两份 kernel 改动逐字对应,无单平台漏改;新增 RocmRejectionSamplingOpTest.cc 复用同一份骨架与同组字面量期望,把「双写 kernel 必须语义一致」变成可执行断言,Bazel target 与 CUDA 侧完全对称。
  • runLiteralCase 先断言 CPU 参考等于硬编码期望、再断言 GPU kernel 等于同一期望(RejectionSamplingOpTest.hpp:57-86),打破旧的「参考实现与 kernel 同源互证」闭环;输出张量用 -7 而非 -1 预填充,可真正验证 kernel 自身写出的 -1 padding。
  • 用例选点专业:runStochasticSemanticsCases 刻意构造「verifier token 与 draft token 相同但仍被随机判据拒绝」的行,这正是新旧语义唯一能区分的场景;vocab=8 使 gcd(4,8)=4 覆盖向量化残差路径。
  • FakeSpeculativeSampler::checkInputs 此前定义但从未被调用(旧期望甚至把 logits 当概率、shape 写成 2-D),本次真正接入调用并改为 3-D [batch, step, vocab] 正确期望(MtpExecutorTest.cc:229-244),端到端锁住 draft 概率契约。
  • 越界与死代码清理正确:min(pos+1, N) 钳位与两处 if (pos != num_speculative_tokens) 守卫在新控制流下均不可达(sampling.cu:304 提前返回保证 pos < N),删除后 uniform_samples[pos+1](最大索引 N,宽度 N+1)与 draft_probs[pos] 均不越界;padding 循环变量改名 pad_idx,消除与概率变量 p 的同名遮蔽。
  • 语义变更未伴随任何字段号、类型或 reserved 段调整,P2PConnectorStartLoadResponsePBreserved 1 to 8 与 reserved 名单原样保留(proto:632-633),protobuf 线格式前后向兼容。


bool same_token = target_id == draft_id;
if (same_token || (do_sample[row_idx] && u * p < q)) {
if ((do_sample[row_idx] && u * p < q) || (!do_sample[row_idx] && same_token)) {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] 随机流接受率因判据收紧与点质量提案系统性下降,缺少实测数据与运维回滚手段

判据由 same_token || (do_sample && u*p<q) 变为按 do_sample 二分。对 do_sample=true!generateConfig()->top1(),SpeculativeSampler.cc:146)的行,旧行为是「verifier argmax 命中即无条件接受」或 min(1, q/p_softmax);叠加 all_probs 改为 one-hot(p 由 softmax 峰值 0.8~0.99 升为 1.0),接受概率退化为 q(draft_id)。两项均单调降低接受率。ROCm sampling.cu:336 同步变更。代码与测试中无接受率/TPS 数据,kernel 亦无 env/config 开关。

建议: 补充 top1()==false 场景的 accept_len / TPS / 输出质量前后对比数据(至少一个 MTP 模型),并在 PR description 中写明这是有偏采样修复及其预期吞吐代价;同时确认现有接受长度指标能按 do_sample 维度区分,便于上线对比。建议评估更优的等价修法:让随机流的 draft token 从 softmax(draft_logits) 采样而非 argmax,同时保留 p=softmax,期望接受率为最大耦合 Σ min(p,q),通常远高于 q(argmax_p)。不建议长期保留双套采样语义的 feature flag。

Checklist: [6.1] 可观测性:日志/指标/超时可操作、非噪声;[6.1] PR description 说明动机与设计

repeated string peer_addrs = 7;
repeated int32 propose_token_ids = 8;
repeated int32 position_ids = 9;
// Exact normalized distribution used to produce propose_token_ids, in

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] 契约注释把 draft 空间的 probs 与已映射为 target 空间的 propose_token_ids 绑定描述

新注释写作「Exact normalized distribution used to produce propose_token_ids, in draft-vocabulary space」,把二者绑定为同一索引空间。但实现中空间不同:FastTopKSampler::forward 在 draft 空间构造 all_probs(宽度 logits.size(1)),随后对 token_ids 调用 execMappingDraft2Target 映射为 target id(SpeculativeSampler.cc:50-54)。新增单测正好证实差异:d2t_map={0,1,3,2},argmax=2 → token_ids=3,而点质量仍在 index 2(SpeculativeSamplerTest.cc:28-42)。d2t_map 非恒等时 propose_probs[propose_token_ids[i]] 取错值甚至越界。

建议: 注释中明确区分两个空间:propose_token_ids 已经过 d2t 映射、位于 target 词表空间;propose_probs 位于 draft 词表空间,消费方必须先用本地 d2t_map 映射(参见 SpeculativeSampler::mapDraftProbsToTarget)后才能与 propose_token_ids 配对,禁止直接用前者索引后者。同一契约目前以四种措辞散落于 proto:319-320、GenerateStream.h:93-94PrefillLoadCaller.h:29P2PConnectorResourceStore.h:38,建议统一以 GenerateRequestPB.propose_probs 为唯一权威定义、其余三处改为引用(与 SideChannelPayloadPB(proto:619)已采用的写法一致),避免后续语义调整产生文档漂移。

Checklist: [6.1] 分层边界:新概念在正确层级,不泄漏内部

int32_t remote_reuse_len = 0;
int32_t memory_reuse_len = 0;
std::vector<int> propose_tokens;
// Exact normalized proposal distribution in draft-vocabulary space.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] 新增契约遗漏 FP32 与 rank 硬约束,跨节点接收边界零校验

消费侧 mapDraftProbsToTarget 在 d2t_map 分支之前就硬断言 dim()==3scalar_type()==kFloat32(SpeculativeSampler.cc:64-66),不满足即 myAssert 终止。但 TensorPB schema 允许 FP16/BF16(proto:5-18),P2PConnector.cc:392-397 对 fp16/bf16 payload 原样 CopyFrom 转发,接收路径 StreamCacheResource.ccDecodeRpcServer 只做 .to(kCUDA)(不改 dtype),兄弟字段 propose_tokens 却校验了 INT32(P2PConnector.cc:387)。本次新增的三处契约注释均未写明 dtype 与 shape,违约只会在采样器深处抛出不带 request_id 的通用断言。

建议: 在三处注释中补齐硬约束:data_type 必须为 FP32、shape 为 [batch, propose_step, draft_vocab_size](单请求发送为 [1, draft_vocab_size])。并在接收边界(PrefillLoadCaller 解析处与 StreamCacheResource::applyP2PSideChannel 落盘前)补齐与 propose_tokens 一致的校验:dtype 非 FP32 或最后一维不等于本地 draft_vocab_size 时,带 request_id 与实际 dtype/shape 返回明确的 ErrorCodePB,而不是把非法 payload 带到采样阶段 abort。

repeated int32 propose_token_ids = 8;
repeated int32 position_ids = 9;
// Exact normalized distribution used to produce propose_token_ids, in
// draft-vocabulary space. PD peers follow the same-build protocol contract.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] 所援引的 same-build 契约作用域未覆盖 generate RPC 的采样语义变更

注释称「PD peers follow the same-build protocol contract」,但本文件唯一的契约段落(proto:28-33)标题为「KV-cache protocol compatibility contract」,正文明确限定适用于「peers that exchange KV-cache RPCs」,动机是分组拓扑与独立 block pool 改变了 block 寻址与 payload 结构,并未覆盖 GenerateRequestPB 的采样分布语义。本 PR 把 propose_probs 由全 softmax 分布改为 one-hot,而字段号 10、TensorPB 类型、张量形状均未变,接收端无从区分。滚动升级期(旧 prefill → 新 decode)不报错,但 p 仍为 0.x 且 same_token 兜底已删除,接受判据被系统性放大、输出分布有偏,无日志或指标能暴露。

建议: 二选一:(1) 把 proto:28-33 的契约段落显式扩写覆盖「proposal 分布语义」,并在 :319-320 按名引用该契约块而非泛指,同时说明混版表现是接受率静默偏移而非报错,便于运维判断回滚条件;(2) 在 payload 中增加 proposal 分布语义版本标量(或复用已有 build/protocol version),decode 侧检测不匹配时 fail-fast 或退回非投机路径。无论哪种,都请在 PR description / 发布说明中写明 prefill 与 decode 必须同批升级(或先 drain 再切换);也可在 decode 侧上报「收到的 propose_probs 最大值是否为 1」的指标使混版可被监控发现。

Checklist: [6.1] 回滚路径:风险行为存在运维回滚手段

return draft_probs;
}

RTP_LLM_CHECK_WITH_INFO(d2t_map.dim() == 1, "d2t_map must be 1-D");

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P2] d2t_map 取值域校验放在了不使用它作索引的类里,mapDraftProbsToTarget 越界退化为 device assert

FastTopKSampler 构造函数校验了 d2t_map.min() >= 0max() < target_vocab_size(SpeculativeSampler.cc:19-21),但真正把 d2t_map 当作 scatter_add_ 写索引使用的 mapDraftProbsToTarget 只校验 dim/dtype/device/numel(:78-83),SpeculativeSampler 构造函数(SpeculativeSampler.h:54)也不校验。目前仅因 MtpExecutor.cc:640-642 用同一个 d2t_map 先后构造两个对象,才被 FastTopKSampler 侧「顺带」在 init 期挡住;该 public static 的其他调用方无任何保护。越界 id 在 CUDA 上触发 device-side assert(污染 context、错误信息无法指向 d2t_map),CPU 上抛 index out of bounds。

建议: 把取值域校验下沉到真正消费它的位置:抽出统一的 validateD2TMap(d2t_map, target_vocab_size)FastTopKSamplerSpeculativeSampler 两个构造函数复用(在构造期做一次性 min/max,避免每 forward .item() 同步),失败时抛出带 d2t_map min/max 与 target_vocab_size 的可读错误,使「init 期 fail-fast」成为类自身的不变量而非调用顺序的巧合。

@@ -103,8 +103,7 @@ TEST_F(SpeculativeSamplingKernelTest, RejectionSampling_ImmediateReject) {
auto draft_token_ids = torch::full({batch_size, num_spec}, 3, intCuda());

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📍 实际位置 rtp_llm/models_py/bindings/cuda/test/CudaSpeculativeSamplingTest.cc:90(不在 diff 展示范围内,就近挂载)

[P3] 内核级单测缺少 do_sample=true 拒绝分支且用例名与新语义不符

该文件是唯一直接调用 invokeRejectionSampling 的单测。RejectionSampling_ImmediateReject(:90)与 RejectionSampling_PartialAccept(:141)的 do_sample 都是 torch::zeros(:110、:171,贪心),走新增的贪心直出分支,完全不触发 relu(q-p) 残差采样;RejectionSampling_AllAccept(:36)虽用 torch::ones(:57)但 draft 与 target 押在同一 token 的全接受场景,无法区分 bonus 来自 target_token_ids 还是残差重采样。本次改动的核心分支——随机行被拒后进入残差采样——在 kernel 级单测中完全未被执行;两个贪心用例名仍叫 rejection sampling。

建议: (1)把两个贪心用例重命名为体现 do_sample=false 语义的名字(如 GreedyVerify_ImmediateMismatchEmitsVerifierTokenGreedyVerify_PartialAccept)。(2)补一条 do_sample=true 的 kernel 直调用例:令 target_token_ids[0]==draft_token_ids[0]target_probs[0][draft_id] 足够小以触发拒绝,断言输出不等于 draft token 且等于 relu(q-p) 的唯一非零位置。若认为分层职责上该文件只做 launch 冒烟,也请在文件头显式注明随机语义由 op 级测试负责。

@@ -119,6 +162,181 @@ class RejectionSamplingOpTest: public ::testing::Test {
assertVectorEqual(getTensorValues<int32_t>(params.output_accepted_token_num_d), expected.accepted_token_num);

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📍 实际位置 rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:130(不在 diff 展示范围内,就近挂载)

[P3] runReferenceCases 中为旧 bonus 位残差采样准备的 fixture 已成死数据

逐 batch 手算:batch 3(do_sample=true)三步全部接受(i=0: u*p=0.20*0.50=0.10 < q=0.80;i=1: 0.5*0.2 < 0.2;i=2: 0.10*0.50=0.05 < 0.90),因此 pos==N,bonus 直接取 target_token_ids[3]=4,不再进入残差采样。专为旧「pos==N 时从 relu(q-0) 采样 bonus」路径准备的 :130 setProbRow(target_probs, ..., 3, 3, {0.10,...}) 与 :133 uniform_samples[3*4+3]=0.55f 已不影响任何输出。(batch 2 在 i=1 处 u*p=0.35 >= q=0.10 被拒,仍会走残差采样,:124 的 uniform_samples[2*4+2]=0.90f 依然有效。)

建议: 删除 :130 与 :133 这两处已失效的 fixture 设置,或把 batch 3 改造成在某一位真正被拒绝的场景(令 u * p >= q),使该 batch 自身重新覆盖随机残差采样路径;同时在函数头补一行注释说明各 batch 分别对应哪条语义分支,避免后续维护者误以为 :130/:133 仍在覆盖 bonus 位残差采样。

Checklist: [6.1] 逻辑变更未混入无关格式化

spec::SpeculativeSamplerOutput forward(const std::list<GenerateStreamPtr>& streams,
SamplerOutput& draft_sampler_output,
SamplerOutput& target_sampler_output) override {
if (!input_holder.test_data.empty()) {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P3] Fake 采样器在期望队列耗尽后静默跳过输入校验

if (!input_holder.test_data.empty()) { checkInputs(...); }(:229-231):test_dataqueueget() 会 pop。因此当 forward 被调用次数超过 setInputs 推入的期望数时,超出部分不再校验且不报错。这一守卫是为兼容未设置期望的旧用例,但也意味着「期望数少于实际调用数」的用例会静默通过——本次新接入的 draft 概率契约校验强度因此取决于测试作者是否恰好补齐了每一步期望。

建议: 改为「一旦某用例设置过期望,就要求调用次数与期望数严格一致」:例如新增显式的 expect_inputs_ 开关,或在 TearDown 中断言 input_holder.test_data.empty(),确保推入的期望全部被消费。

Checklist: [6.1] 新逻辑有聚焦单测 + 相关集成/smoke 测试

@@ -193,7 +193,7 @@ class FakeModel: public ModelBase {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

📍 实际位置 rtp_llm/cpp/normal_engine/speculative/test/MtpExecutorTest.cc:116(不在 diff 展示范围内,就近挂载)

[P3] MtpExecutorTest 中 catVectors 在删除调用点后成为死代码

本 PR 把 testSingleBatchDecodetestMultiBatchDecode 中构造 draft_all_probs_list 的代码改为 torch::cat(:828-832、:1190-1195),catVectors 的两处(也是全部)调用点随之消失。全仓检索确认 catVectors 仅剩 :116 一处定义、零引用。因其为函数模板,未实例化不会触发 -Wunused-function,所以不会打断构建,只是留下无人使用的测试辅助模板。

建议: 删除 catVectors 模板定义。若后续仍需拼接 vector<vector<T>>,届时再引入更贴合当时需求的写法,避免留存无调用方的辅助函数。

Checklist: [6.1] KISS/YAGNI:无投机性抽象

int32_t remote_reuse_len = 0;
int32_t memory_reuse_len = 0;
std::vector<int> propose_tokens;
// Exact normalized proposal distribution in draft-vocabulary space.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[P3] 点质量提案仍按全词表稠密张量上线传输,且 P2PConnector 的 FP32 空值哨兵恒失效

注释把 propose_probs 定义为「Exact normalized proposal distribution in draft-vocabulary space」。在当前唯一的 top-1 提案实现下该张量恒为 one-hot:prefill 侧每请求需序列化并跨机传输 float32 × draft_vocab(十万级词表数百 KB),decode 侧再解析、H2D、参与 zero_ + scatter_add_,实际信息量只有一个 token id。wire format 未变故非本次引入的回归,但注释把这一冗余表示固化成了协议契约。同时 P2PConnector.cc:392 的空值判断以 data_type() != TensorPB::FP32 为首个条件,而 proto3 默认值 FP32 = 0(proto:7),未设置的 TensorPB 该条件恒为 false,空 payload 也会被塞入 tensors map。

建议: (1)后续可在 payload 中增加「退化分布」标记(只传 token id + 标志位),decode 侧按标记直接构造 point mass,既保留未来 top-p 提案传稠密分布的能力,又省掉跨机带宽与显存清零开销;建议先在注释中记录该演进方向,避免实现被 one-hot 语义锁死。(2)把 P2PConnector.cc:392 的空值判断改为统一复用 tensorPbHasPayload 语义,去掉不成立的 dtype 哨兵。

Checklist: [6.1] 兼容性:外部 HTTP/RPC API、持久数据、配置、环境迁移安全;[I] 同一功能用统一工具函数

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants