feat: enable Qwen35-MoE MTP && sp prefill cuda graph

[amd-yilizhao]

Summary

Enable ROCm paged prefill CUDA graph capture/replay for MTP (speculative decoding) scenarios on Qwen3.5-MoE.

Key changes:

• C++ cuda_graph layer: Distinguish embedding-style vs MTP draft prefill graph modes (isEmbeddingStylePrefillCudaGraph / isMtpDraftPrefillCudaGraph); skip invalid batch-1 post-check for MTP draft; add finish_capture_session() to defer trt_allreduce IPC handle exchange and avoid ProcessGroupNCCL watchdog races
• FusedRopeKVCacheOp (ROCm): Add updatePrefillRuntime() pybind method to pass runtime max_seq_len/max_prefix_len from Python, eliminating .item() host-device sync during graph capture; conditionally skip k/v output buffer allocation in paged mode
• aiter.py (ROCm): Add HIP graph capture/replay support for AiterPrefillAttnOpPaged and AiterPrefillImplPaged — pre-allocate metadata buffers, replace .to(device) with pre-moved tensors, replace .item() with pre-computed values, disable triton short-seq path during graph capture
• rocm_rccl.py: Add RCCL collective warmup after comm bootstrap (fixes Mi355 tp=8 ncclAllGather failure during graph capture); defer consume_capture() to finish_hipgraph_capture_session(); enable trt_allreduce priority path during graph capture; add lazy comm init from ProcessGroup
• trt_allreduce.py: Add has_pending_capture() and _capture_handles_pending flag for deferred IPC handle finalization

[LLLLKKKK]

AI Code Review - PR #969

Status: BLOCKING

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

Blocking Issues

P1

• HIPGraph 收尾异常被静默吞掉 @ rtp_llm/models_py/distributed/rocm_rccl.py:459
  • 建议：至少记录异常并重新抛出，或将 pending 状态显式清理/降级，避免 graph capture 成功表象下进入不可恢复的 collective 状态。
• 缺少 ROCm MTP prefill graph 回归覆盖 @ rtp_llm/models_py/modules/factory/attention/rocm_impl/aiter.py:1164
  • 建议：补充至少一个 ROCm MTP draft prefill CUDA graph 用例，覆盖 prefix、all_gather/allreduce capture、finish_capture_session 成功/失败路径。

Checklist Violations (12 fail / 81 total)

General Principles Checklist

• [6.1] Architecture — 状态不变量：创建/更新/失败/重试/回滚路径有效 → issue HIPGraph 收尾异常被静默吞掉
  finish_hipgraph_capture_session 静默吞掉 consume_capture 异常，pending capture handle 状态失败后没有可见回滚。
• [6.1] Architecture — 错误语义：fail-fast/retry/fallback/silent 行为显式 → issue HIPGraph 收尾异常被静默吞掉
  except Exception: pass 让 capture 收尾失败变成 silent 行为，调用侧无法区分成功和失败。
• [6.1] Architecture — 可观测性：日志/指标/超时可操作、非噪声 → issue HIPGraph 收尾异常被静默吞掉
  finish_hipgraph_capture_session 捕获所有异常后不记录日志，线上无法定位 handle exchange 或 barrier 失败。
• [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue 缺少 ROCm MTP prefill graph 回归覆盖
  diff_paths 没有测试文件，但新增了 MTP prefill graph、ROCm collectives capture 和 session finish 行为。
• [6.1] Tests — 边界 case 覆盖（空、单元素、最大值） → issue 缺少 ROCm MTP prefill graph 回归覆盖
  MTP graph 分支涉及 active_bs、inactive batch、prefix_len/max_bs 边界，但未见新增边界测试。
• [6.1] Tests — 分布式/跨平台变更有对应覆盖 → issue 缺少 ROCm MTP prefill graph 回归覆盖
  ROCm RCCL/trt_allreduce capture 生命周期改动没有对应 ROCm/TP 测试变更。

RTP-LLM Checklist

• [E] 分布式 — 跨 rank 数据一致性 → issue HIPGraph 收尾异常被静默吞掉
  consume_capture 失败被吞掉后，各 rank 可能对 capture handle 是否已交换产生不一致状态。
• [E] 分布式 — collective op 匹配 → issue HIPGraph 收尾异常被静默吞掉
  finish session 的 barrier/all_gather_object 异常被静默吞掉，会掩盖 collective 不匹配或某 rank 未进入的问题。
• [F] 跨平台 — ROCm 路径错误处理非静默 → issue HIPGraph 收尾异常被静默吞掉
  finish_hipgraph_capture_session 对所有异常 pass，ROCm capture 收尾失败不可见。
• [H] 测试与 CI — 测试覆盖充分：大重构等价覆盖，新功能端到端测试 → issue 缺少 ROCm MTP prefill graph 回归覆盖
  高风险 ROCm graph/collective/MTP prefill 新逻辑未见新增测试或 smoke 覆盖。

Python Static-First Checklist

• [P.A] 静态结构与类型纪律 — 禁止 getattr/setattr literal 访问 → checklist-only
  aiter.py 新增 getattr(attn_inputs, ...) 和 getattr(self.rope_params, ...)；这是 pybind 兼容写法，当前未形成具体崩溃链路，暂不单独升级。
• [P.B] 错误处理 — 禁止 bare except 或静默吞异常 → issue HIPGraph 收尾异常被静默吞掉
  rocm_rccl.py 的 finish_hipgraph_capture_session 使用 except Exception 后 pass。

Strengths

• 新增 is_draft_prefill_cuda_graph_capture 字段带 C++ 默认值，pybind 和 pyi 同步暴露，减少旧路径未设置字段时的崩溃风险。
• ROCm paged prefill 路径把 graph replay 需要的 seqlen、block table、descale buffer 预分配到对象内，方向上符合 CUDA graph 避免 replay 期分配的要求。

[LLLLKKKK]

AI Code Review - PR #969

Status: BLOCKING

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

Blocking Issues

P1

• C++ 收尾吞掉 HIPGraph finalize 异常 @ rtp_llm/cpp/cuda_graph/cuda_graph_device_shims.cc:207
  • 建议：记录日志后重新抛出，或显式清理 pending capture 状态并让本次 graph capture 失败退出。
• MTP graph replay 准备阶段引入 GPU 同步 @ rtp_llm/models_py/modules/factory/attention/rocm_impl/aiter.py:1138
  • 建议：优先使用 CPU 侧 input_lengths/cu_seqlens_host 计算这些标量，或在 C++ prepareInputs 中传入已知值，避免每次 graph replay 前同步 GPU。

Non-blocking Suggestions

P2

• MTP prefill graph replay 缺少端到端覆盖 @ rtp_llm/models_py/modules/factory/attention/rocm_impl/test/test_aiter_prefill_op.py:297
  • 建议：补一个 ROCm MTP draft prefill graph 用例，至少覆盖 prefix=0/非0、replay batch 小于 capture batch、finish_capture_session 成功和失败路径。

Checklist Violations (10 fail / 81 total)

General Principles Checklist

• [6.1] Architecture — 状态不变量：创建/更新/失败/重试/回滚路径有效 → issue C++ 收尾吞掉 HIPGraph finalize 异常
  C++ finish_capture_session 捕获 Python finalize 异常后继续，pending capture handle 失败时没有失败退出或回滚。
• [6.1] Architecture — 错误语义：fail-fast/retry/fallback/silent 行为显式 → issue C++ 收尾吞掉 HIPGraph finalize 异常
  Python 单测要求 consume_capture 异常传播，但 C++ shim 将其降级为 warning，调用侧无法区分 finalize 成功和失败。
• [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue MTP prefill graph replay 缺少端到端覆盖
  新增单测覆盖 rocm_rccl finish 和 support()，但没有覆盖 MTP prefill graph replay 的核心 prepare/update/replay 行为。
• [6.1] Tests — 分布式/跨平台变更有对应覆盖 → issue MTP prefill graph replay 缺少端到端覆盖
  ROCm graph collective 有 mock 单测，但 MTP prefill graph 的分布式 replay/TP collective 组合未被端到端覆盖。

RTP-LLM Checklist

• [D] 性能 — 设备同步 .item()/.cpu()/.tolist() 在关键路径 → issue MTP graph replay 准备阶段引入 GPU 同步
  prepare_cuda_graph/replay 前新增 q_lengths.max().item()、kv_lengths.max().item()、q_lengths.sum().item()。
• [E] 分布式 — 跨 rank 数据一致性 → issue C++ 收尾吞掉 HIPGraph finalize 异常
  consume_capture/barrier 失败被 C++ catch 后继续，可能导致各 rank 对 captured handle 是否完成交换认知不一致。
• [E] 分布式 — collective op 匹配 → issue C++ 收尾吞掉 HIPGraph finalize 异常
  finalize 内 barrier/all_gather_object 异常被 C++ 降级为 warning，会掩盖 rank 未进入或 collective 不匹配。
• [F] 跨平台 — ROCm 路径错误处理非静默 → issue C++ 收尾吞掉 HIPGraph finalize 异常
  ROCm C++ shim 对 Python finish 异常只 warning 后继续，错误没有传播给 capture 调用链。
• [H] 测试与 CI — 测试覆盖充分：大重构等价覆盖，新功能端到端测试 → issue MTP prefill graph replay 缺少端到端覆盖
  新增测试未覆盖 MTP draft prefill graph 的 prepare_cuda_graph/update_prefill_runtime/replay 端到端行为。

Python Static-First Checklist

• [P.A] 静态结构与类型纪律 — 禁止 getattr/setattr literal 访问 → checklist-only
  aiter.py 新增 getattr(attn_inputs, ...) 和 getattr(self.rope_params, ...)；这是兼容 pybind/旧对象字段的局部写法，未形成独立运行时缺陷。

Strengths

• Python 侧 finish_hipgraph_capture_session 增加 has_pending_capture 判断，并用单测覆盖 pending、非 pending、异常传播和 ImportError 场景。
• 新增 PyAttentionInputs.is_draft_prefill_cuda_graph_capture 带 C++ 默认值，并同步 pybind 和 pyi，降低旧路径未设置字段的崩溃风险。

[LLLLKKKK]

AI Code Review - PR #969

Status: BLOCKING

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

Blocking Issues

P1

• FP8 paged prefill 返回未定义 Q buffer @ rtp_llm/models_py/bindings/rocm/FusedRopeKVCacheOp.cc:264
  • 建议：FP8 paged 路径返回 q_fp8_buf，非 paged FP8 才返回 qkv_buf_fp8，并补充 FP8 paged prefill 覆盖。
• MTP prefill replay 填错 padded cu_kv_seqlens @ rtp_llm/cpp/cuda_graph/cuda_graph_runner.cc:262
  • 建议：分别维护 query_total 和 kv_total；cu_kv_seqlens padding tail 用最后有效 KV 累积值填充，并增加 prefix+padding replay 测试。
• max_prefix_len 计算会低估真实 prefix @ rtp_llm/models_py/modules/factory/attention/rocm_impl/aiter.py:1182
  • 建议：直接使用 prefix_lengths.max() 计算 max_prefix_len，或从 live cu_seqlens 得到每请求 prefix 后取最大值。

Checklist Violations (9 fail / 97 total)

General Principles Checklist

• [6.1] Architecture — 状态不变量：创建/更新/失败/重试/回滚路径有效 → issue MTP prefill replay 填错 padded cu_kv_seqlens
  MTP prefill replay padding 使用 query total 填 cu_kv_seqlens，破坏 KV 累积长度不变量。
• [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue FP8 paged prefill 返回未定义 Q buffer
  新增测试覆盖 capture finalize，但未覆盖 FP8 paged prefill forward 返回 buffer。
• [6.1] Tests — 边界 case 覆盖（空、单元素、最大值） → issue max_prefix_len 计算会低估真实 prefix
  未覆盖最大 prefix 与最大 query 分属不同 batch 的边界。

RTP-LLM Checklist

• [B] 正确性与逻辑 — 逻辑错误、off-by-one、null/zero 检查 → issue FP8 paged prefill 返回未定义 Q buffer
  FP8 paged 分支写 q_fp8_buf，却返回未分配的 qkv_buf_fp8。
• [B] 正确性与逻辑 — 边界 case（空输入、单元素、最大值） → issue max_prefix_len 计算会低估真实 prefix
  max_prefix_len 使用 max_k - max_q，最大 prefix 与最大 query 分离时会低估。
• [D] 性能 — CUDA Graph replay 参数过期 → issue MTP prefill replay 填错 padded cu_kv_seqlens
  replay 前更新的 cu_kv_seqlens padding tail 使用 query 累计值，prefix+padding 时 graph 参数失效。
• [D] 性能 — Model input 修改处理 CUDA Graph padding 填充 → issue MTP prefill replay 填错 padded cu_kv_seqlens
  cu_kv_seqlens padding tail 未按 KV prefix 后的累计长度填充。
• [H] 测试与 CI — 测试覆盖充分：大重构等价覆盖，新功能端到端测试 → issue FP8 paged prefill 返回未定义 Q buffer
  测试未覆盖 FP8 paged prefill graph forward，无法捕获未定义 buffer 返回。

Python Static-First Checklist

• [P.A] 静态结构与类型纪律 — 禁止 getattr/setattr literal 访问 → checklist-only
  aiter.py 使用 getattr 兼容 pybind 新旧字段；这是局部兼容写法，未形成独立缺陷链。

Strengths

• finish_capture_session 将 ROCm consume_capture 延后到 capture loop 之后，并在单测里覆盖 pending、非 pending、异常传播和 ImportError 分支。
• PyAttentionInputs 新增 is_draft_prefill_cuda_graph_capture 带 C++ 默认值，并同步 pybind 与 pyi，兼容旧路径未设置字段的场景。

[LLLLKKKK]

AI Code Review - PR #969

Status: LGTM

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

lgtm ready to ci

Checklist Violations (1 fail / 56 total)

Python Static-First Checklist

• [P.A] 静态结构与类型纪律 — 禁止 getattr/setattr literal 访问 → checklist-only
  aiter.py 使用 getattr(attn_inputs, ...) 和 getattr(self.rope_params, ...) 兼容 pybind 新旧字段；当前有默认值和 callable 检查，未形成独立缺陷。

Strengths

• PyAttentionInputs 新增字段带 C++ 默认值，并同步 pybind 与 pyi，降低旧对象未设置字段时的兼容风险。
• ROCm HIPGraph capture 收尾改为显式 finish session，consume_capture 异常会向 C++ 调用链传播。
• MTP prefill graph 的 prefix、cu_seqlens、FP8 paged 返回路径补了针对性单测。