fix: correct MTP speculative sampling semantics - #1295
Conversation
LLLLKKKK
left a comment
There was a problem hiding this comment.
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-92的rejection_sampling_test_utils导出、仅被models_py/bindings/cuda/ops/tests:rejection_sampling_op_test消费,且无 ROCm 变体,更新不会波及其他 target。同时补两条守护用例:do_sample=true且same_token成立但u*p>=q必须走残差采样;do_sample=true全接受但存在same_token=false时 bonus 必须等于target_token_ids。本地跑通该 target 三个用例再提交,并在 PR 描述中说明旧黄金期望为何失效。
- 建议:在本 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_tokensbonus 门控到 ROCm 内核;更优做法是把这段串行接受循环抽成两侧共享的__device__inline 辅助函数(放入共享头文件),消除「CUDA 内核 / ROCm 内核 / CPU 参考」三份并存的判据副本,并给 ROCm 侧补一条复用同一黄金参考的 op 测试。若本 PR 暂不覆盖 ROCm,请在 PR 描述与sampling.cu注释中显式声明「ROCm 语义未同步」及跟进计划,避免后续误认为已修复。
- 建议:在本 PR 同步移植 CUDA 侧的接受判据、
- 随机采样路径的三处新语义没有任何判别性测试覆盖 @
rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:275- 建议:补三类确定性 kernel 用例:(1)
do_sample=true、draft_probs 为 point mass、q(draft)=0.3、u=0.5,断言即使same_token成立也被拒绝且发射 token 落在relu(q-p)支撑集;(2)do_sample=true在中间位置被拒,按固定u断言accepted=pos+1、output[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)重采样,实际走贪心直写分支)。
- 建议:补三类确定性 kernel 用例:(1)
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。
- 建议:在 PR 描述与发布说明中明确「prefill/decode 必须同版本升级」并给出灰度顺序建议;更稳妥的做法是让 decode 侧不依赖对端概率语义——既然
- 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 分支。
- 建议:明确并固化索引口径:建议在 d2t 映射之后再构造 one-hot(或对
- 随机请求接受率将系统性下降,缺少实测数据与运维回滚手段 @
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_INFOfail-fast,并在注释中声明「不得用于 rejection sampling」。同时补一条top_k=2用例断言all_probs等于torch::softmax(logits, -1),防止 point-mass 逻辑误扩散到多候选路径。
- 建议:明确该分支契约:要么统一按 point-mass 语义处理(对 top-k 候选各自置 1 或按需归一化,并把 d2t 映射的
P3
- top_k==1 路径 softmax 成为无用计算并新增每步全词表分配 @
rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:12- 建议:
top_k==1分支直接对logits取torch::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_testdeps 中显式加入exec_ops_test_lib并补上与兄弟 target 一致的注释。同时参照兄弟用例先断言token_ids的dim()与各维大小、增加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 边界,混版滚动升级期静默不一致
draftall_probs不只在进程内流转:MtpBatchStreamProcessor.cc:1001/1039把draft_sampler_output.all_probs存入流的sp_output_buffer->all_probs,PrefillRpcServer.cc:400-413将其序列化进 proto 字段propose_probs(model_rpc/proto/model_rpc_service.proto:311),decode 侧经DecodeRpcServer.cc:293-297或StreamCacheResource.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)、新值为q,q=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)、新值为q,q=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=1、token_offset=0调用execMappingDraft2Target,其实现按index = i*token_stride + j(CudaSampleOp.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)、新值为q,q=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.cc(FastTopKSamplerTest套件)及目标speculative_sampler_test,其 BUILD 注释明确写了FastTopKSampler::forward calls execMappingDraft2Target并显式声明exec_ops_test_lib+cuda_impl(speculative/test/BUILD:29-37)。新用例却放进MtpExecutorTest.cc:该文件其余用例一律用FakeFastTopKSampler(:194)打桩、从不执行真实forward;mtp_executor_testdeps 未声明exec_ops_test_lib,仅靠device_test_utils间接带入(链接可用)。另:550的item<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、:349的if (pos != num_speculative_tokens)恒真、:339的min(pos + 1, num_speculative_tokens)恒等于pos + 1,而「there is no draft_probs for the bonus token」注释描述的分支已永不进入(旧内核在all_same_token=false且pos=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_token、same_token || (do_sample && u*p<q)(:336)、if (all_same_token)(:349)与if (s_all_same_token) return;(:364),本 PR 未改。该文件在rocm/BUILD:83-97的speculative_sampling目标中实际参与构建,并由core/CudaSampleOp.cc:819-836的 ROCmrejectionSampling调用。而FastTopKSampler是平台无关代码,one-hot 会同时喂给 ROCm:p由 softmax 变为 1.0 使判据收紧,但same_token仍短路接受、拒绝时仍直吐 verifier token——结果既不同于修复前也不同于修复后的 CUDA,接受率、拒绝后 t - [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue
随机采样路径的三处新语义没有任何判别性测试覆盖
本次改变三种行为:(1)do_sample=true且same_token但u*p>=q由接受变拒绝;(2)do_sample=true拒绝后走残差采样而非直吐 verifier token;(3) 全接受时 bonus 固定取target_token_ids且不再消费uniform_samples末列。但cuda/test/CudaSpeculativeSamplingTest.cc四个既有用例均为退化输入:AllAccept(:36) 虽do_sample=ones但uniform_samples全 0 且 probs 已是 one-hot,必然全接受;ImmediateReject(:90)/PartialAccept(:143) 的do_sample均为torch::zeros;BatchSizeZero(:203) 为空批。逐个复算后新旧输出完全相同,无判别力。本 PR 唯一新增用例只覆盖FastTopKSampler的 one-hot 输出,不触及 kernel;比 - [6.1] Tests — 边界 case 覆盖(空、单元素、最大值) → issue
新增用例落位偏离既有专用测试文件,断言未校验形状且含冗余拷贝
仓内已有FastTopKSampler的专属单测SpeculativeSamplerTest.cc(FastTopKSamplerTest套件)及目标speculative_sampler_test,其 BUILD 注释明确写了FastTopKSampler::forward calls execMappingDraft2Target并显式声明exec_ops_test_lib+cuda_impl(speculative/test/BUILD:29-37)。新用例却放进MtpExecutorTest.cc:该文件其余用例一律用FakeFastTopKSampler(:194)打桩、从不执行真实forward;mtp_executor_testdeps 未声明exec_ops_test_lib,仅靠device_test_utils间接带入(链接可用)。另:550的item<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_idsbonus 并提前 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。复算runReferenceCasesbatch 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.cc(FastTopKSamplerTest套件)及目标speculative_sampler_test,其 BUILD 注释明确写了FastTopKSampler::forward calls execMappingDraft2Target并显式声明exec_ops_test_lib+cuda_impl(speculative/test/BUILD:29-37)。新用例却放进MtpExecutorTest.cc:该文件其余用例一律用FakeFastTopKSampler(:194)打桩、从不执行真实forward;mtp_executor_testdeps 未声明exec_ops_test_lib,仅靠device_test_utils间接带入(链接可用)。另:550的item<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、其后补-1、accepted=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二次重采样更贴合请求的采样配置,也与forceSpAccept(SpeculativeSampler.cc:174)提取 bonus 的位置一致。 - 关键分支都有解释「为什么」的注释(如「残差采样只对随机投机解码成立」);新增用例用
checkTensorEqual(内部torch::equal精确比较)断言精确 one-hot 值{0,0,1,0}而非仅 argmax,固定 logits 与期望值均可精确表示,无 flaky 风险。 - 依赖链已核实:
mtp_executor_test经device_test_utils(cpp/testing/BUILD:30-39默认 select 分支)间接获得exec_ops_test_lib与cuda_impl,新用例可正常 link,排除链接失败误报。
21dbb33 to
5cbe31d
Compare
LLLLKKKK
left a comment
There was a problem hiding this comment.
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 阶段的显存/带宽实测对照。
- 建议:把「计算/搬运」与「传输兼容」拆开:1)用运行期开关(或按是否启用 PD 分离判定)控制是否物化 draft 概率,decode 单机场景直接跳过 softmax 与 stack/clone;2)把 MtpBatchStreamProcessor.cc:903 的硬 CHECK 降级为「仅在确实需要构造
- 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.h的invokeRejectionSampling声明上方补契约注释,明确deterministic_draft=true时draft_probs允许为 nullptr、提议分布由draft_token_ids隐式给出,ROCm 同名声明保持同一措辞。
- 建议:1)让该标志由提议来源推导而非常量:由
- deterministic_draft=true 与贪心 / 混合 batch 的生产主组合零覆盖 @
rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:273- 建议:在
runDeterministicDraftCases补两组字面量用例:1)deterministic_draft=true且do_sample全 false,分别覆盖「verifier 全命中、bonus 取target_token_ids[num_spec]」与「首位失配,输出 target token、其后补-1、accepted==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),drafttoken_ids用 int64 CPU 张量以覆盖 dtype/device 转换,喂入固定的 targetall_probs/token_ids,断言accept_tokens/accept_len与手工推演一致、未定义draft_probs不触发校验异常、-1 -> 0归零生效;至少覆盖「全接受」与「首位拒绝」两种配置,形成对装配契约的回归防线。
- 建议:在
- deterministic_draft 新增的参数校验分支缺少负例测试 @
rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:400- 建议:在
runRejectsInvalidTensorMetadata增补三条负例:deterministic_draft=false且draft_probs_d未定义应抛异常;deterministic_draft=false且draft_probs_d的 batch/num_spec/vocab 维与draft_token_ids_d、target_probs_d不一致应抛异常;deterministic_draft=true且draft_token_ids_d与target_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与>1、temperature=0与>0四种组合,量化预期回退幅度;3)在 PR description 中写明「相同 seed 的随机采样请求输出会变化、accept_len 预期下降」这一兼容性影响,并明确回滚手段(版本回滚,或可运行期切换的判据开关),避免被误判为性能回归。
- 建议:这是必要的分布正确性修复,不建议改回旧逻辑,但请补齐运维面:1)在 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,把转换责任上移到生产者。
- 建议:调整顺序,使被持有的张量与拷贝源一致:先在 host 侧完成
- 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 侧同步处理,避免后续两平台改漏。
- 建议:在两个 vocab 循环之前取到寄存器:
- 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_vec、relu_q_minus_p类型统一;在两处 kernel 模板参数列表上给这两个 bool 加简短注释,降低写反风险。
- 建议:为 ROCm 侧补一个与
- 异常断言过宽且未遵循仓库关闭 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=false且draft_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_probs与uniform_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 已无 draftall_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_probs与uniform_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_rocmselect 与alwayslink=1闭合到 HIP 实现,torch_deps()/tags/exec_properties与同目录既有目标一致,ROCm 首次获得该 kernel 的算子级回归。
4532394 to
b5cd0ee
Compare
b5cd0ee to
764cdbf
Compare
Vinkle-hzt
left a comment
There was a problem hiding this comment.
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 描述和代码边界注释明确了同版本部署要求、回滚方式与线上观察指标。
|
已按最终 review 收口到 commit |
LLLLKKKK
left a comment
There was a problem hiding this comment.
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透传到SpeculativeSampler并RTP_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快路径保留为显式可选项);同时建议新增/复用接受长度指标,使灰度期具备对照与回滚判断依据。
- 建议:合入前给出 MTP
- propose_probs 取值语义与 dtype 静默变更,三处反序列化零校验且 same-build 条款不覆盖生成路径 @
rtp_llm/cpp/model_rpc/proto/model_rpc_service.proto:312- 建议:在
PrefillLoadCaller.cc与DecodeRpcServer.cc的还原处对propose_probs校验 dtype(FP32)与末维长度,与相邻propose_tokens的 INT32 校验保持一致,不满足时按请求粒度返回 gRPC 错误并打印 request_id 与 peer,让契约违约在传输边界 fail-fast 而非下沉为共享采样器异常。同时把 same-build 契约范围显式扩写到生成路径 RPC(RemoteGenerate/StartLoad的propose_*载荷),或增加提议分布语义标识(枚举/最小协议版本号)供接收方 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=1且do_sample默认 true 与显式 false)、小 vocab 的 one-hot targetall_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.cc与RocmRejectionSamplingOpTest.cc同时注册,使两平台的贪心契约都有不依赖参考实现的锚点。
- 建议:在
- mapDraftProbsToTarget 的直通分支与新增十余处校验分支缺少用例 @
rtp_llm/cpp/normal_engine/speculative/test/SpeculativeSamplerTest.cc:45- 建议:复用已正确保存/恢复
user_ft_core_dump_on_exception的FastTopKSamplerValidationTestfixture 补齐:d2t_map未定义且 vocab 相等时断言返回值与入参同data_ptr(零拷贝直通)、vocab 不等时EXPECT_THROW(..., RTPException);构造函数侧覆盖 d2t_map 含负值、最大值等于/超过target_vocab_size、dtype 非 int64;mapDraftProbsToTarget侧覆盖非 3-D、非 float32、d2t_mapnumel 不匹配。
- 建议:复用已正确保存/恢复
- 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,并配一条可观测的计数/日志,避免后续扩展提议分布类型时出现无声错误。
- 建议:在 op 层
- 同文件遗留 ChainSpeculativeSampling 保留旧语义,形成语义分叉的死代码 @
rtp_llm/models_py/bindings/rocm/speculative_sampling/sampling.cu:458- 建议:在本 PR 或紧随的清理 PR 中删除
ChainSpeculativeSampling/execChainSpeculativeSampling链路(删除前按 R.I.2 再做一次全仓消费者搜索确认);若因外部依赖需保留,请在其定义处加一行注释标明「legacy,已废弃,权威实现为rejection_sampling_kernel,语义不同,勿作为基准」。
- 建议:在本 PR 或紧随的清理 PR 中删除
- 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复用。
- 建议:把校验抽成共享 helper(如
- 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/形状要求补进权威处一次。
- 建议:沿用 proto:610-611 已采用的引用式写法,统一为「以
- 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_map下atomicAdd的累加顺序会破坏可复现性,需改用确定性归约。
- 建议:用 profiler 或
- dtype 转换使 hold_host 的 H2D 生命周期契约被绕过 @
rtp_llm/cpp/normal_engine/speculative/SpeculativeSampler.cc:131- 建议:先显式做 dtype 转换得到 host 张量、
hold_host该结果、再做纯设备迁移(to(device, /*non_blocking=*/true)),使被持有对象与实际参与异步拷贝的 buffer 一致;或在该处加一行注释说明 dtype 不一致时拷贝为同步、无需持有,使该保护的适用条件显式化。
- 建议:先显式做 dtype 转换得到 host 张量、
- 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 延迟的前后对比。
- 建议:非阻塞优化建议:契约层面为 top-1 场景约定稀疏表示(仅传 token id,接收方本地展开为点质量),字段保留以支持未来 top-p 分布;或明确允许多种
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 idscatter_(draft 空间),:53-54随后对token_ids调execMappingDraft2Target就地改写为 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 idscatter_(draft 空间),:53-54随后对token_ids调execMappingDraft2Target就地改写为 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:373持side_channel_mutex时按值深拷贝整个SideChannelData(含propose_probs),:397再CopyFrom一次,等于每请求在锁内外各复制一份全词表张量,直接落在首 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:211裸CopyFrom(同函数propose_tokens(:202)、position_ids(:222)都校验 INT32),StreamCacheResource.cc只判存在性,DecodeRpcServer.cc:293-297直接transTensor(...).to(kCUDA)。混版本 PD 下:单步stack(MtpBatchStreamProcessor.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()>=0、max()<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 作「不校验」哨兵::33的target_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:373持side_channel_mutex时按值深拷贝整个SideChannelData(含propose_probs),:397再CopyFrom一次,等于每请求在锁内外各复制一份全词表张量,直接落在首 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_samplebatch 只出现在runReferenceCases(do_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 把 draftall_probs重定义为「产生该 token 的精确分布」并启用了checkInputs,但两处夹具未同步:(1):714的stream1_draft_token_probs={{0,0,1,0}}点质量在 index 2,而:716的StreamSpecUpdateInfo第三字段draft_token=3(字段序见GenerateStream.h:51-57),该夹具又被:828拼进期望张量(配:827的token_ids={3,2,1,3});(2):1160的token_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()>=0、max()<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-oned2t_map下重复下标互相覆盖导致质量丢失;旧的「仅宽度不等才映射」条件使等宽非恒等映射被整段跳过,kernel 侧用 target id 索引 draft 空间的 p。MapsEqualWidthNonIdentityVocabMap与MapsProposalDistributionWithManyToOneVocabMap(含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()>=0、max()<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为准」的单向引用而非平行复制,本部分回滚成本为零。
| @@ -27,4 +27,4 @@ cudaError_t invokeRejectionSampling(DType* draft_probs, | |||
| int num_speculative_tokens, | |||
There was a problem hiding this comment.
📍 实际位置 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 本次只改了文件末尾换行,签名与注释未变;validateRejectionSamplingParams(CudaSampleOp.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)) { |
There was a problem hiding this comment.
[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 说明动机与设计
| 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); |
There was a problem hiding this comment.
[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_map 下 atomicAdd 的累加顺序会破坏可复现性,需改用确定性归约。
| 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); |
There was a problem hiding this comment.
[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, |
There was a problem hiding this comment.
[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)); |
There was a problem hiding this comment.
[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:373 持 side_channel_mutex 时按值深拷贝整个 SideChannelData(含 propose_probs),:397 再 CopyFrom 一次,等于每请求在锁内外各复制一份全词表张量,直接落在首 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
left a comment
There was a problem hiding this comment.
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.h的invokeRejectionSampling上方补齐三条契约的注释(target 空间索引、p(draft_id)>0、do_sample双分支语义),并在validateRejectionSamplingParams增加可选的调试期校验(例如draft_probs_d.gather(-1, draft_token_ids_d.to(kLong))全大于 0),使违约在 op 边界以可读信息失败,而不是在内核里静默改变分布。同时为OpData.h的do_sample_d补一行语义注释。
- 建议:在
- 随机采样行接受率系统性下降,缺少实测数据与引擎级运维回滚手段 @
rtp_llm/models_py/bindings/cuda/kernels/speculative_sampling/sampling.cu:268- 建议:在 PR description 补充随机采样场景下接受长度与端到端 TPS 的前后对比数据;同时提供一个进程级配置开关(例如允许把随机行强制按
!do_sample的精确匹配路径验证,或把提案分布退化为 top-p 过滤后的真实分布),使线上出现吞吐劣化时可在不回滚二进制的前提下降级。建议在既有接受长度指标上按do_sample拆分维度,便于灰度对比。
- 建议:在 PR description 补充随机采样场景下接受长度与端到端 TPS 的前后对比数据;同时提供一个进程级配置开关(例如允许把随机行强制按
- 契约注释把 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_CHECK。P2PConnectorResourceStore.h、PrefillLoadCaller.h、OpData.h、GenerateStream.h的同类注释建议统一引用该处权威描述,避免六处复述措辞漂移。
- 建议:拆开两字段的空间描述:
- batchSample 的三处语义变更零覆盖且恰被 fake 屏蔽,mapDraftProbsToTarget 用例全在 CPU @
rtp_llm/cpp/normal_engine/speculative/test/SpeculativeSamplerTest.cc:45- 建议:补
batchSample层测试:构造top_k=1/do_sample=true、top_k>1/do_sample=true、do_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」这一新契约具备真实回归能力。
- 建议:在字面用例中构造至少一行「全接受 + bonus 位
- 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>=0、max<target_vocab_size、numel==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,也为后续平滑提案分布预先划清确定性边界。
- 建议:在构造期一次性判定 d2t_map 是否单射/恒等(
- 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 约定,便于容量评估。
- 建议:对 top-1 提议考虑不再传输稠密分布:或仅传被提议 token 的概率标量,或直接省略
- 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- 建议:若确认该路径已废弃,连同
SpeculativeSamplingParams、chainSpeculativeSampling、execChainSpeculativeSampling与对应 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:397的propose_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()==3与kFloat32(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:结果仍等价(q被top_k_renorm_probs归一为 argmax 点质量,u*1<q退化为精确匹配,失配时relu(q-p)亦为 argmax 处的 one-hot),但内核不再走direct_target_fallback的提前 return,每次失配都要额外执行两遍全词表的 block reduce/scan 残差采样。反向地,do_sample=false且top_k=0的请求由旧的随机路径改为精确匹配,属于改善。 - [6.1] Architecture — 回滚路径:风险行为存在运维回滚手段 → issue
随机采样行接受率系统性下降,缺少实测数据与引擎级运维回滚手段
对do_sample=true且top_k>1/top_p<1的行,旧判据为same_token || u*p_softmax(draft_id) < q(p_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-1590的total_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:结果仍等价(q被top_k_renorm_probs归一为 argmax 点质量,u*1<q退化为精确匹配,失配时relu(q-p)亦为 argmax 处的 one-hot),但内核不再走direct_target_fallback的提前 return,每次失配都要额外执行两遍全词表的 block reduce/scan 残差采样。反向地,do_sample=false且top_k=0的请求由旧的随机路径改为精确匹配,属于改善。 - [6.1] Architecture — 错误语义:fail-fast/retry/fallback/silent 行为显式 → issue
d2t_map 校验在三处不统一,作为写索引的 mapDraftProbsToTarget 反而不校验取值范围
同一份d2t_map在三处被校验且强度不同:FastTopKSamplerctor 校验 1-D、int64 与[0, target_vocab_size)取值域(:16-21);mapDraftProbsToTarget只校验 1-D、int64、device、numel(:78-83),不校验取值范围;SpeculativeSamplerctor 完全不校验(.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=true且top_k>1/top_p<1的行,旧判据为same_token || u*p_softmax(draft_id) < q(p_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-1590的total_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_list与catVectors(...)的调用,改用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 == 0时u = uniform * 0 = 0,aggregate > 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在三处被校验且强度不同:FastTopKSamplerctor 校验 1-D、int64 与[0, target_vocab_size)取值域(:16-21);mapDraftProbsToTarget只校验 1-D、int64、device、numel(:78-83),不校验取值范围;SpeculativeSamplerctor 完全不校验(.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结果始终连续,与checkRejectionSamplingTensor的is_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_utils的select()链接 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 层二进制兼容。
| @@ -27,4 +27,4 @@ cudaError_t invokeRejectionSampling(DType* draft_probs, | |||
| int num_speculative_tokens, | |||
There was a problem hiding this comment.
📍 实际位置 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_d,do_sample_d 无任何语义说明。
建议: 在 sampling.h 的 invokeRejectionSampling 上方补齐三条契约的注释(target 空间索引、p(draft_id)>0、do_sample 双分支语义),并在 validateRejectionSamplingParams 增加可选的调试期校验(例如 draft_probs_d.gather(-1, draft_token_ids_d.to(kLong)) 全大于 0),使违约在 op 边界以可读信息失败,而不是在内核里静默改变分布。同时为 OpData.h 的 do_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)) { |
There was a problem hiding this comment.
[P2] 随机采样行接受率系统性下降,缺少实测数据与引擎级运维回滚手段
对 do_sample=true 且 top_k>1/top_p<1 的行,旧判据为 same_token || u*p_softmax(draft_id) < q(p_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-1590 的 total_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 |
There was a problem hiding this comment.
[P2] 契约注释把 draft 空间的 probs 与已映射为 target 空间的 propose_token_ids 绑定描述
新注释写作「用于产生 propose_token_ids 的精确归一化分布,位于 draft 词表空间」,但两字段不在同一索引空间:PrefillRpcServer.cc:397 的 propose_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()==3 与 kFloat32(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_CHECK。P2PConnectorResourceStore.h、PrefillLoadCaller.h、OpData.h、GenerateStream.h 的同类注释建议统一引用该处权威描述,避免六处复述措辞漂移。
Checklist: [6.1] 兼容性:外部 HTTP/RPC API、持久数据、配置、环境迁移安全
| EXPECT_TRUE(torch::allclose(out.all_probs, expected)); | ||
| } | ||
|
|
||
| TEST(SpeculativeSamplerTest, MapsProposalDistributionWithManyToOneVocabMap) { |
There was a problem hiding this comment.
[P2] batchSample 的三处语义变更零覆盖且恰被 fake 屏蔽,mapDraftProbsToTarget 用例全在 CPU
新增用例只覆盖 FastTopKSampler::forward 与静态 mapDraftProbsToTarget(SpeculativeSamplerTest.cc:27-108),batchSample 无任何测试入口;而 MtpExecutorTest 用 FakeSpeculativeSampler(:222-233)与 FakeFastTopKSampler(:194-201)替换了真实实现。因此本次行为切换面最大的三行——do_sample_d = generateConfig()->do_sample(:145)、draft_token_ids.to(target_device, kInt32, true)(:133)、mapDraftProbsToTarget 与 draft_probs_target_buffer_ 的生产接线(:169)——零覆盖。此外 7 个映射用例全部由 torch::tensor(...) 构造、运行在 CPU,scatter_add_ 的 CUDA 原子路径与 :80-81 的 devi...
建议: 补 batchSample 层测试:构造 top_k=1/do_sample=true、top_k>1/do_sample=true、do_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)); |
There was a problem hiding this comment.
[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 约定,便于容量评估。
| spec::SpeculativeSamplerOutput forward(const std::list<GenerateStreamPtr>& streams, | ||
| SamplerOutput& draft_sampler_output, | ||
| SamplerOutput& target_sampler_output) override { | ||
| if (!input_holder.test_data.empty()) { |
There was a problem hiding this comment.
[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 { | |||
|
|
|||
There was a problem hiding this comment.
📍 实际位置 rtp_llm/cpp/normal_engine/speculative/test/MtpExecutorTest.cc:116(不在 diff 展示范围内,就近挂载)
[P3] MtpExecutorTest 中 catVectors 辅助函数变为死代码
本 PR 在 testSingleBatchDecode / testMultiBatchDecode 中删除了 draft_all_probs_list 与 catVectors(...) 的调用,改用 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; | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
📍 实际位置 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 消费者。未来若有人接线,会拿到一个不支持贪心验证的内核。
建议: 若确认该路径已废弃,连同 SpeculativeSamplingParams、chainSpeculativeSampling、execChainSpeculativeSampling 与对应 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
left a comment
There was a problem hiding this comment.
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-94、PrefillLoadCaller.h:29、P2PConnectorResourceStore.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)供FastTopKSampler与SpeculativeSampler两个构造函数复用(在构造期做一次性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_ids对do_sample=true行必须来自真随机采样;uniform_samples[i]用于位置 i 的接受判定而[pos+1]用于残差采样。ROCm 侧同步。(2)新增一组do_sample=true、全部接受、bonus 位target_probs[N]为非退化多峰分布且target_token_ids[N]明确不等于残差采样命中值的字面用例(可照搬runReferenceCasesbatch 3 的{0.10,0.20,0.30,0.25,0.15}+uniform=0.55),独立锁定「bonus 来自 target sampler」契约。
- 建议:(1)在
- 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。
- 建议:(1)为 4 个映射用例增加 CUDA 变体(对输入
- do_sample=false 贪心直出分支无字面期望值,ROCm 侧只能与同源参考实现互证 @
rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:165- 建议:新增
runGreedySemanticsCases(或在现有函数内)补 1~2 条do_sample=false的runLiteralCase,硬编码期望覆盖:(1)逐位与 verifier 一致直至 bonus;(2)首个不一致位直接输出 verifier token、accepted_token_num = pos + 1、其后补-1;(3)混合do_sample的 batch。在RocmRejectionSamplingOpTest.cc与CudaRejectionSamplingOpTest.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_id的atomicMin归约。CUDA 与 ROCm 共用同一 helper,一次补齐覆盖两个平台。
- 建议:补一条大 vocab 用例(例如
- 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 刚强化的点质量契约固化下来,而不是继续以裸字面量表达、任其漂移。
- 建议:把 :714 改为 index 3 的点质量
- 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_ImmediateMismatchEmitsVerifierToken、GreedyVerify_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 级测试负责。
- 建议:(1)把两个贪心用例重命名为体现
- runReferenceCases 中为旧 bonus 位残差采样准备的 fixture 已成死数据 @
rtp_llm/cpp/testing/RejectionSamplingOpTest.hpp:130- 建议:删除 :130 与 :133 这两处已失效的 fixture 设置,或把 batch 3 改造成在某一位真正被拒绝的场景(令
u * p >= q),使该 batch 自身重新覆盖随机残差采样路径;同时在函数头补一行注释说明各 batch 分别对应哪条语义分支,避免后续维护者误以为 :130/:133 仍在覆盖 bonus 位残差采样。
- 建议:删除 :130 与 :133 这两处已失效的 fixture 设置,或把 batch 3 改造成在某一位真正被拒绝的场景(令
- 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 哨兵。
- 建议:(1)后续可在 payload 中增加「退化分布」标记(只传 token id + 标志位),decode 侧按标记直接构造 point mass,既保留未来 top-p 提案传稠密分布的能力,又省掉跨机带宽与显存清零开销;建议先在注释中记录该演进方向,避免实现被 one-hot 语义锁死。(2)把
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=1024下smem_size可能超限)。CUDA 侧走cudaLaunchKernel并抛出 launch 错误。调用方rejectionSampling的RTP_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」路径准备的 :130setProbRow(target_probs, ..., 3, 3, {0.10,...})与 :133uniform_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 把testSingleBatchDecode与testMultiBatchDecode中构造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)却没有虚析构函数。MtpExecutor以std::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=1024下smem_size可能超限)。CUDA 侧走cudaLaunchKernel并抛出 launch 错误。调用方rejectionSampling的RTP_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_data是queue,get()会 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=1024、VEC_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 自身写出的-1padding。- 用例选点专业:
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 段调整,
P2PConnectorStartLoadResponsePB的reserved 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)) { |
There was a problem hiding this comment.
[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 |
There was a problem hiding this comment.
[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-94、PrefillLoadCaller.h:29、P2PConnectorResourceStore.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. |
There was a problem hiding this comment.
[P2] 新增契约遗漏 FP32 与 rank 硬约束,跨节点接收边界零校验
消费侧 mapDraftProbsToTarget 在 d2t_map 分支之前就硬断言 dim()==3 与 scalar_type()==kFloat32(SpeculativeSampler.cc:64-66),不满足即 myAssert 终止。但 TensorPB schema 允许 FP16/BF16(proto:5-18),P2PConnector.cc:392-397 对 fp16/bf16 payload 原样 CopyFrom 转发,接收路径 StreamCacheResource.cc 与 DecodeRpcServer 只做 .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. |
There was a problem hiding this comment.
[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"); |
There was a problem hiding this comment.
[P2] d2t_map 取值域校验放在了不使用它作索引的类里,mapDraftProbsToTarget 越界退化为 device assert
FastTopKSampler 构造函数校验了 d2t_map.min() >= 0 与 max() < 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) 供 FastTopKSampler 与 SpeculativeSampler 两个构造函数复用(在构造期做一次性 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()); | |||
There was a problem hiding this comment.
📍 实际位置 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_ImmediateMismatchEmitsVerifierToken、GreedyVerify_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); | |||
There was a problem hiding this comment.
📍 实际位置 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()) { |
There was a problem hiding this comment.
[P3] Fake 采样器在期望队列耗尽后静默跳过输入校验
if (!input_holder.test_data.empty()) { checkInputs(...); }(:229-231):test_data 是 queue,get() 会 pop。因此当 forward 被调用次数超过 setInputs 推入的期望数时,超出部分不再校验且不报错。这一守卫是为兼容未设置期望的旧用例,但也意味着「期望数少于实际调用数」的用例会静默通过——本次新接入的 draft 概率契约校验强度因此取决于测试作者是否恰好补齐了每一步期望。
建议: 改为「一旦某用例设置过期望,就要求调用次数与期望数严格一致」:例如新增显式的 expect_inputs_ 开关,或在 TearDown 中断言 input_holder.test_data.empty(),确保推入的期望全部被消费。
Checklist: [6.1] 新逻辑有聚焦单测 + 相关集成/smoke 测试
| @@ -193,7 +193,7 @@ class FakeModel: public ModelBase { | |||
|
|
|||
There was a problem hiding this comment.
📍 实际位置 rtp_llm/cpp/normal_engine/speculative/test/MtpExecutorTest.cc:116(不在 diff 展示范围内,就近挂载)
[P3] MtpExecutorTest 中 catVectors 在删除调用点后成为死代码
本 PR 把 testSingleBatchDecode 与 testMultiBatchDecode 中构造 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. |
There was a problem hiding this comment.
[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] 同一功能用统一工具函数
Summary
a887c9848c465df6e9768941960f497b1014033donto the latestmain.all_probs/propose_probspath.Design
u * p < q; matching verifier/draft IDs do not bypass this test.relu(q - p).scatter_add, so many draft IDs mapping to one target ID preserve total probability mass.top_k > 1remains rejected inFastTopKSampler; 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_probsnow explicitly means the actual normalized proposal distribution used to sample the transmitted draft token, not raw model softmax scores.Review checklist
deterministic_draftruntime/template branch and kept one explicit-distribution kernel contract.N=0, andN=1.relu(q - p)sampling with a point-mass token in a non-leading vector lane.batch > 1with a changed draft-step capacity and asserted contiguous mapped probabilities.RTPExceptioncontract with core dumps disabled for the fixture.Validation
//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_testclang-format --dry-run --Werrorgit diff --checkrules_pythonfetch. CUDA/ROCm compilation and ROCm hardware execution remain delegated to CI.