feat(attn): dynamic decode attention backend selection at CUDA graph capture

[SAzwj]

Summary

Add runtime selection of the fastest decode attention backend per batch-size
bucket during CUDA graph capture, replacing the fixed first-match priority.

Controlled by --enable_dynamic_decode_backend (env ENABLE_DYNAMIC_DECODE_BACKEND),
default OFF — zero behavior change when the flag is off.

How it works

At CUDA graph capture time, for each decode batch-size bucket:

1. Benchmark: run real on-device micro-benchmarks across a KV-length grid
   for each candidate backend (XQA TRT, FlashInfer XQA, PyFlashinfer, etc.)
2. Select: pick the stably fastest backend via select_stable (requires
   consistent wins, falls back to priority on tie)
3. Dispatch: prepare_fmha_impl looks up backend_plan[bs] at decode time
   (cache miss → fixed priority fallback)

Key changes

1. Deferred CG capture — move CUDA graph capture out of PyWrappedModel
   constructor to engine-driven triggerInitCapture(), opening a window for
   bench/dispatch before graphs are recorded.

2. dispatch/ package — new selector, backend_bench, backend_selector
   modules + C++ wiring (enable_dynamic_decode_backend config flag,
   CudaGraphDecode plan storage, module_base backend lookup).

3. PyFlashinfer CG replay fix — call plan() before every CG replay for
   all decode paths (FA2 and non-FA2) with CG-safe tensor placement, fixing a
   hang on replay.

4. Smoke tests — dyn_decode cases for dense (qwen3 fp8pb, fp8pb tp2) and
   moe (qwen3_moe fp8pt deepep). Also adds kernel_seq_size_per_block to
   next/moe_vl cases and updates affected goldens.

Testing

• CPU unit tests: test_selector, test_backend_bench, test_backend_selector
• GPU integration: decode CG capture/replay including non-FA2 hang detection
• Smoke: dense + moe dyn_decode cases with dedicated golden data

[LLLLKKKK]

AI Code Review - PR #1173

Status: LGTM

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

lgtm ready to ci

Checklist ✅ (56 items passed)

Strengths

• 动态选择路径默认关闭，并在失败时回退到固定优先级 backend，风险边界相对清晰。
• FlashInfer decode replay 补充了 plan 刷新、device indices、padding workspace 清零等保护，覆盖了 CUDA graph replay 的关键状态更新。
• 新增 selector/backend_selector/backend_bench 单测和 H20 smoke，覆盖了 disable flags、clone 隔离、graph replay、TP/FP8/MoE 等主要路径。

[LLLLKKKK]

AI Code Review - PR #1173

Status: BLOCKING

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

Blocking Issues

P1

• 默认 kernel page 推导在 Python/C++ 两侧不一致 @ rtp_llm/cpp/cache/CacheConfigCreator.cc:13
  • 建议：将 Python 与 C++ 的默认候选集合和 fallback 规则统一到同一策略：要么两侧都只允许 128/64/32/16 并在 Python 侧提前拒绝不支持的配置，要么同步扩展 C++ cache 侧对 8/4/2 的支持并补齐对应测试。建议增加覆盖非 16 倍数、8 倍数等边界 block size 的配置测试。

Checklist Violations (1 fail / 56 total)

General Principles Checklist

• [6.1] Architecture — 兼容性：公开 API/持久数据/配置/环境迁移安全 → issue 默认 kernel page 推导在 Python/C++ 两侧不一致
  当用户未显式设置 kernel_seq_size_per_block 时，Python 侧候选包含 8/4/2 并会给 seq_size_per_block=136 推导出 8；C++ cache 创建只接受 128/64/32/16，同一配置会在 seq_size_per_block % 16 != 0 处 fail-fast。该值来自启动配置，默认推导两侧不一致会导致可配置组合启动失败。

Strengths

• 新功能通过 enable_dynamic_decode_backend 显式开关控制且默认关闭，降低了默认路径回归面。
• Python/C++ 配置绑定保留了旧 pickle 字段数 fallback，兼顾了已有序列化兼容性。
• 动态 backend 选择失败时保留固定优先级 fallback，避免选择逻辑异常直接阻断推理。

[LLLLKKKK]

AI Code Review - PR #1173

Status: LGTM

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

lgtm ready to ci

Checklist ✅ (55 items passed)

Strengths

• 动态后端选择默认关闭，并在 Python/C++ 两侧保留 fixed-priority fallback，失败时不会阻塞启动。
• CUDA graph 捕获前的 benchmark 使用独立 graph pool、clone 输入和显式 cleanup，降低了对真实捕获状态的污染风险。
• 新增配置经过 server args、pybind、pickle 兼容路径，并补充了 dispatch 单测和动态 decode smoke 覆盖。

[LLLLKKKK]

AI Code Review - PR #1173

Status: BLOCKING

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

Blocking Issues

P1

• FlashInfer 探测实例释放前未清理 workspace 状态 @ rtp_llm/models_py/modules/factory/attention/dispatch/backend_selector.py:76
  • 建议：在 _eligible() 和 _fallback() 的真实后端实例生命周期中复用 FlashInfer cleanup 逻辑，并用 try/finally 保证清理后再释放实例；或者把 support/fallback 判定改为不触发 plan() 的静态能力判断。建议补一个单测覆盖 probe/fallback 后再进入 prepare_fmha_impl() 时 FlashInfer wrapper/workspace 状态已被清理。

Checklist Violations (2 fail / 55 total)

General Principles Checklist

• [6.1] Architecture — 状态不变量：创建/更新/失败/重试/回滚路径有效 → issue FlashInfer 探测实例释放前未清理 workspace 状态
_eligible() 为 support probe 直接构造真实 impl 后仅 del inst；_fallback() 也通过 get_fmha_impl() 构造后端后直接返回类型名。FlashInfer 构造会执行 prepare()/plan() 并复用全局 workspace，而同 PR 的 bench cleanup 明确需要清理旧 plan 状态，否则会污染后续 capture。
• [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue FlashInfer 探测实例释放前未清理 workspace 状态
_eligible() 为 support probe 直接构造真实 impl 后仅 del inst；_fallback() 也通过 get_fmha_impl() 构造后端后直接返回类型名。FlashInfer 构造会执行 prepare()/plan() 并复用全局 workspace，而同 PR 的 bench cleanup 明确需要清理旧 plan 状态，否则会污染后续 capture。

Strengths

• 新配置默认关闭，并在 server args、pybind、C++ config 和 pickle 兼容路径中做了串联。
• 动态选择逻辑对 CUDA graph、CUDA device、hybrid attention、MLA 等不适用场景有显式降级。
• 新增 selector/bench 单测、FlashInfer CUDA graph 测试和 H20 smoke case，覆盖了主要启用路径。

[LLLLKKKK]

AI Code Review - PR #1173

Status: BLOCKING

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

Blocking Issues

P1

• 新增 HWKernelConfig 字段未同步到 pybind 类型 stub @ rtp_llm/cpp/pybind/ConfigInit.cc:593
  • 建议：同步更新或重新生成 libth_transformer_config.pyi，把 enable_dynamic_decode_backend: bool 加到 HWKernelConfig，并确保 pybind 配置字段变更时 stub 一并更新。

Checklist Violations (1 fail / 55 total)

General Principles Checklist

• [6.1] Architecture — 兼容性：外部 HTTP/RPC API、持久数据、配置、环境迁移安全 → issue 新增 HWKernelConfig 字段未同步到 pybind 类型 stub
HWKernelConfig 已通过 pybind 暴露 enable_dynamic_decode_backend，server args 也会写入该字段；但 rtp_llm/ops/libth_transformer_config.pyi 的 HWKernelConfig 仍没有该属性，typed Python 调用和静态检查会看到过期配置合约。

Strengths

• 动态 backend 选择默认关闭，并在 CUDA-only、MLA/hybrid skip、TP broadcast、失败 fallback 等路径上做了保护。
• CUDA graph replay 前的 FlashInfer plan 刷新和 padding workspace 处理有明确注释，并补了相关单测与 smoke 配置。

[LLLLKKKK]

AI Code Review - PR #1173

Status: LGTM

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

lgtm ready to ci

Checklist ✅ (55 items passed)

Strengths

• 动态选择默认关闭，并且在非 CUDA、MLA、hybrid KV cache、target verify 等不适用场景有显式 guard，失败时回退固定优先级。
• C++ deferred capture 增加了 triggerInitCapture() 的 no-op/防重复语义，NormalEngine 在服务启动前触发，避免 graph 未 capture 就进入请求路径。
• dispatch 层把纯 CPU selector、GPU benchmark、TP 广播 orchestration 分层，并补充了单测与 H20 smoke 覆盖。

[wht21]

internal source has been updated, please review the changes!

[LLLLKKKK]

AI Code Review - PR #1173

Status: LGTM

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

lgtm ready to ci

Checklist ✅ (55 items passed)

Strengths

• 动态选择默认关闭，并且在非 CUDA、MLA、hybrid KV cache、target verify 等不适用场景有显式 guard，失败时回退固定优先级。
• C++ deferred capture 增加了 triggerInitCapture() 的 no-op/防重复语义，NormalEngine 在服务启动前触发，避免 graph 未 capture 就进入请求路径。
• dispatch 层把纯 CPU selector、GPU benchmark、TP 广播 orchestration 分层，并补充了单测与 H20 smoke 覆盖。

[wht21]

internal source has been updated, please review the changes!