[Cherry-Pick][XPU] Enable CudaGraph capture for MTP draft model(#7864)#7941
[Cherry-Pick][XPU] Enable CudaGraph capture for MTP draft model(#7864)#7941Clarity256 wants to merge 2 commits into
Conversation
|
Thanks for your contribution! |
|
|
Based on PR PaddlePaddle#7864, this adds draft model CudaGraph support: - Enable step_use_cudagraph for draft model with proper gating logic - Pass forward_meta and use_cudagraph to xpu_pre_process in draft path - Add padding_cudagraph_inputs() for draft model buffer management - Slice model output by real_token_num when graph is active - Include target model cudagraph changes (xpu_model_runner, xpu_pre_and_post_process) - Add build_sampling_params XPU kernel for MTP Verified stable with benchmark_serving (100 requests, concurrency=16, 0 failures). Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Clarity256
force-pushed
the
feature/mtp-cudagraph-draft-model-xpu
branch
from
5a6882b to
bf150a4
Compare
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
CI报告基于以下代码生成（30分钟更新一次）: 1 任务总览当前暂无 required 失败任务，但 7 个 Workflow 处于
2 任务状态汇总日志列说明：失败任务直接使用工具生成的日志链接；运行中任务使用 Job 页面链接。 2.1 Required任务 : 0/0 通过
2.2 可选任务 — 1/2 通过
3 失败详情（仅 required）无 required 失败任务。本轮不对 optional 失败任务做深度分析； |
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## develop #7941 +/- ##
==========================================
Coverage ? 64.06%
==========================================
Files ? 467
Lines ? 65065
Branches ? 9977
==========================================
Hits ? 41687
Misses ? 20556
Partials ? 2822
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
This comment was marked as outdated.
This comment was marked as outdated.
Sorry, something went wrong.
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
Clarity256
force-pushed
the
feature/mtp-cudagraph-draft-model-xpu
branch
from
4ef7a1f to
00299e4
Compare
PaddlePaddle-bot
left a comment
PaddlePaddle-bot
left a comment
There was a problem hiding this comment.
🤖 Paddle-CI-Agent | pr_review |
2026-05-28 14:46:48
📋 Review 摘要
PR 概述：为 XPU 平台 MTP 投机解码的 draft model 补齐 CUDAGraph 支持，新增 build_sampling_params XPU 算子替代 Python 侧实现
变更范围：custom_ops/xpu_ops（新增算子）、fastdeploy/spec_decode/mtp_xpu.py、fastdeploy/worker/xpu_model_runner.py、sampler.py、xpu_pre_and_post_process.py
影响面 Tag：[XPU] [Speculative Decoding] [OP]
问题
| 级别 | 文件 | 概述 |
|---|---|---|
| 🟡 建议 | xpu_model_runner.py:868 |
重构移除了 proposer 的 MoE phase 同步逻辑，MoE+EP+投机解码场景可能回归 |
历史 Findings 修复情况
| Finding | 问题 | 状态 |
|---|---|---|
| F1 | substep 参数未在方法签名中声明 |
|
| F2 | infer_seed 在非投机解码路径不再更新 |
|
| F3 | 同 F1，substep 参数问题 |
|
| F4 | capture_size 整除截断 |
|
| F5 | kernel 中 pad_start O(bs²) 性能问题 | |
| F6 | sampling_metadata.topp_seed 语义不一致 |
|
| F7 | mtp_xpu.py 遗留大量注释代码 | |
| F8 | 同 F4，capture_size 整除截断 |
📝 PR 规范检查
标题格式符合 Cherry-Pick 规范 [Cherry-Pick][XPU] ... ✅。但 ## Usage or Command 段为空（仅模板注释），Checklist 全部未勾选（实际已新增单测 test_build_sampling_params.py 且提供了精度对比截图）。建议补全如下：
标题建议（可直接复制）：
[Cherry-Pick][XPU][Speculative Decoding] Enable CUDAGraph capture for MTP draft model (#7864)
PR 描述建议（点击展开，可直接复制）
## Motivation
在 PR #7864 中，XPU 平台 MTP 投机解码的 target model 已支持 CUDAGraph，但 draft model 侧仍未启用 CUDAGraph capture。本 PR 在 #7864 基础上，为 MTP draft model 补齐 CUDAGraph 支持，主要包括：
1. draft model 前向推理启用 `step_use_cudagraph` 门控逻辑，并在 multi-step 执行中仅对首步进行 capture。
2. draft model 推理路径中传递 `forward_meta` 和 `use_cudagraph` 到 `xpu_pre_process`，确保 `cu_seqlens_q_output` / `batch_id_per_token_output` 在 cudagraph 模式下使用 `copy_` 原地更新，保证 tensor 地址稳定性。
3. 新增 `padding_cudagraph_inputs()` 方法处理 draft model 的 buffer padding，并在 graph replay 时按 `real_token_num` 切片 model output。
4. target model 侧投机解码 warmup 流程适配（capture size 计算、`accept_all_drafts` 参数传递）。
5. 将 `padding_sampling_params`（Python 侧 CPU 实现）替换为 `build_sampling_params` XPU 自定义算子，在算子内部完成 `infer_seed` 的原地更新，避免在 cudagraph 外额外操作。
6. `increment_value` 改为与投机解码 token 数联动（`(num_speculative_tokens + 1) * 4`）。
## Modifications
- **`custom_ops/xpu_ops/src/ops/mtp/build_sampling_params.cc`**：新增 Paddle 自定义算子入口，注册 `build_sampling_params` op。
- **`custom_ops/xpu_ops/src/plugin/include/xpu/plugin.h`**：声明 `build_sampling_params` C 接口。
- **`custom_ops/xpu_ops/src/plugin/src/kernel/kunlun3cpp/mtp_kernel/build_sampling_params.xpu`**：新增 XPU3 kernel 实现。
- **`custom_ops/xpu_ops/src/plugin/src/wrapper/mtp_wrapper/build_sampling_params.cpp`**：新增 CPU wrapper 和 XPU3 wrapper。
- **`custom_ops/xpu_ops/test/test_build_sampling_params.py`**：新增单元测试，覆盖纯 decoder、纯 encoder、混合、单条、seed wrap-around 等场景。
- **`fastdeploy/model_executor/layers/sample/sampler.py`**：`forward_xpu` 改用 `build_sampling_params` XPU 算子替代 `padding_sampling_params`；新增 `increment_value` 参数。
- **`fastdeploy/model_executor/xpu_pre_and_post_process.py`**：cudagraph 模式下改用 `copy_` 原地更新 `cu_seqlens_q_output` 和 `batch_id_per_token_output`，保证 graph 捕获的 tensor 地址稳定。
- **`fastdeploy/spec_decode/mtp_xpu.py`**：draft model 启用 `step_use_cudagraph` 门控；`_propose` 新增 cudagraph padding 逻辑与 output slicing；`_initialize_forward_meta` 传递 cudagraph 参数。
- **`fastdeploy/worker/xpu_model_runner.py`**：`increment_value` 与投机解码 token 数联动；warmup capture 流程适配 speculative decoding；`infer_seed` 更新移入 `build_sampling_params` 算子内部；draft model propose 传递 `step_use_cudagraph`。
- **`tests/xpu_ci/4cards_cases/run_mtp_cudagraph.py → test_mtp_cudagraph.py`**：重命名测试脚本以符合 CI 命名规范。
## Usage or Command
N/A
## Accuracy Tests
- MTP with CUDAGraph：输出与参考结果一致（见 PR 截图）
- MTP without CUDAGraph：输出与参考结果一致（见 PR 截图）
## Checklist
- [x] Add at least a tag in the PR title.
- Tag list: [`[FDConfig]`,`[APIServer]`,`[Engine]`, `[Scheduler]`, `[PD Disaggregation]`, `[Executor]`, `[Graph Optimization]`, `[Speculative Decoding]`, `[RL]`, `[Models]`, `[Quantization]`, `[Loader]`, `[OP]`, `[KVCache]`, `[DataProcessor]`, `[BugFix]`, `[Docs]`, `[CI]`, `[Optimization]`, `[Feature]`, `[Benchmark]`, `[Others]`, `[XPU]`, `[HPU]`, `[GCU]`, `[DCU]`, `[Iluvatar]`, `[Metax]`]
- You can add new tags based on the PR content, but the semantics must be clear.
- [ ] Format your code, run `pre-commit` before commit.
- [x] Add unit tests. Please write the reason in this PR if no unit tests.
- [x] Provide accuracy results.
- [x] If the current PR is submitting to the `release` branch, make sure the PR has been submitted to the `develop` branch, then cherry-pick it to the `release` branch with the `[Cherry-Pick]` PR tag.总体评价
整体实现思路清晰，CUDAGraph 地址稳定性处理（copy_ 原地更新、output slicing）和新增 XPU 算子替代 Python 侧实现的方案合理。建议关注历史 Findings 中多个未修复问题（尤其 F2 infer_seed 更新缺失和 F4 整除截断），以及本次新增的 proposer MoE phase 同步遗漏。
| if ( | ||
| self.fd_config.scheduler_config.splitwise_role == "mixed" and envs.FD_XPU_ENABLE_MIXED_EP_MODE | ||
| ): # Centralized scenario: the phase is initialized as "prefill" by default. During inference runtime, different types of batches can achieve phase switching at this point. | ||
| if self.fd_config.parallel_config.use_ep and self.fd_config.scheduler_config.splitwise_role == "mixed": |
There was a problem hiding this comment.
🟡 建议 移除了 self.proposer.fd_config.parallel_config.moe_phase.phase 的运行时更新
原代码在 use_cudagraph 分支内对 proposer 的 MoE phase 进行了同步更新：
if self.speculative_decoding:
self.proposer.fd_config.parallel_config.moe_phase.phase = "decode" if if_only_decode else "prefill"本次重构后该逻辑被删除。若 draft model 使用 MoE + EP，其 phase 将始终保持初始值，可能导致 EP dispatch/combine 路径不正确。
建议确认 draft model 是否涉及 MoE；若涉及，需在此处补回 proposer 的 phase 同步逻辑。
4 participants
Motivation
在 PR #7864 中，XPU 平台 MTP 投机解码的 target model 已支持 CUDAGraph，但 draft model 侧仍未启用 CUDAGraph capture。本 PR 在 #7864 基础上，为 MTP draft model 补齐 CUDAGraph 支持，主要包括：
step_use_cudagraph门控逻辑，并在 multi-step 执行中仅对首步进行 capture。forward_meta和use_cudagraph到xpu_pre_process，确保cu_seqlens_q_output/batch_id_per_token_output在 cudagraph 模式下使用copy_原地更新，保证 tensor 地址稳定性。padding_cudagraph_inputs()方法处理 draft model 的 buffer padding，并在 graph replay 时按real_token_num切片 model output。accept_all_drafts参数传递）。padding_sampling_params（Python 侧 CPU 实现）替换为build_sampling_paramsXPU 自定义算子，在算子内部完成infer_seed的原地更新，避免在 cudagraph 外额外操作。increment_value改为与投机解码 token 数联动（(num_speculative_tokens + 1) * 4）。Modifications
custom_ops/xpu_ops/src/ops/mtp/build_sampling_params.cc：新增 Paddle 自定义算子入口，注册build_sampling_paramsop。custom_ops/xpu_ops/src/plugin/include/xpu/plugin.h：声明build_sampling_paramsC 接口。custom_ops/xpu_ops/src/plugin/src/kernel/kunlun3cpp/mtp_kernel/build_sampling_params.xpu：新增 XPU3 kernel 实现。custom_ops/xpu_ops/src/plugin/src/wrapper/mtp_wrapper/build_sampling_params.cpp：新增 CPU wrapper 和 XPU3 wrapper。custom_ops/xpu_ops/test/test_build_sampling_params.py：新增单元测试，覆盖纯 decoder、纯 encoder、混合、单条、seed wrap-around 等场景。fastdeploy/model_executor/layers/sample/sampler.py：forward_xpu改用build_sampling_paramsXPU 算子替代padding_sampling_params；新增increment_value参数。fastdeploy/model_executor/xpu_pre_and_post_process.py：cudagraph 模式下改用copy_原地更新cu_seqlens_q_output和batch_id_per_token_output，保证 graph 捕获的 tensor 地址稳定。fastdeploy/spec_decode/mtp_xpu.py：draft model 启用step_use_cudagraph门控；_propose新增 cudagraph padding 逻辑与 output slicing；_initialize_forward_meta传递 cudagraph 参数。fastdeploy/worker/xpu_model_runner.py：increment_value与投机解码 token 数联动；warmup capture 流程适配 speculative decoding；infer_seed更新移入build_sampling_params算子内部；draft model propose 传递step_use_cudagraph。tests/xpu_ci/4cards_cases/run_mtp_cudagraph.py → test_mtp_cudagraph.py：重命名测试脚本以符合 CI 命名规范。Usage or Command
Accuracy Tests
Checklist
[FDConfig],[APIServer],[Engine],[Scheduler],[PD Disaggregation],[Executor],[Graph Optimization],[Speculative Decoding],[RL],[Models],[Quantization],[Loader],[OP],[KVCache],[DataProcessor],[BugFix],[Docs],[CI],[Optimization],[Feature],[Benchmark],[Others],[XPU],[HPU],[GCU],[DCU],[Iluvatar],[Metax]]pre-commitbefore commit.releasebranch, make sure the PR has been submitted to thedevelopbranch, then cherry-pick it to thereleasebranch with the[Cherry-Pick]PR tag.