fix: size flashinfer prefill workspace dynamically

[Vinkle-hzt]

Summary

• Dynamically size PyFlashInfer fa2 prefill workspace and round the allocation up to a power of two.
• Share the PyFlashInfer workspace across CUDA graph captured prefill instances to avoid per-graph workspace retention.
• Resolve the actual FlashInfer prefill backend before wrapper construction and add a post-plan plan_info check to catch workspace underestimation.

[LLLLKKKK]

AI Code Review - PR #1099

Status: BLOCKING

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

Blocking Issues

P1

• CUDA graph replay 前新增 GPU 到 CPU 同步 @ rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:255
  • 建议：replay 路径避免读取 CUDA tensor；改用已有 host mirror/input_lengths 计算，或在 capture/plan 阶段缓存 workspace 上界。

Non-blocking Suggestions

P2

• 运行逻辑依赖 FlashInfer 私有属性 @ rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:425
  • 建议：封装 FlashInfer 兼容层并显式校验版本/属性，优先使用 public API；无法避免时给出清晰 fallback 或启动期 fail-fast。
• 动态 workspace 分支缺少大尺寸回归覆盖 @ rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:596
  • 建议：补充 mocked plan_info 或大序列 GPU 用例，覆盖普通扩容、capture 预扩容、replay forbid_realloc 报错和 plan_info 校验。

Checklist Violations (11 fail / 67 total)

General Principles Checklist

• [6.1] Software Engineering — DIP：高层策略不依赖非必要具体细节 → issue 运行逻辑依赖 FlashInfer 私有属性
  运行逻辑直接读取 FlashInfer wrapper 私有字段，具体实现细节依赖形成兼容风险。
• [6.1] Architecture — 分层边界：新概念在正确层级，不泄漏内部 → issue 运行逻辑依赖 FlashInfer 私有属性
  _FlashInfer 私有 _backend/_plan_info/int_workspace_buffer 泄漏到 RTP attention 实现逻辑。
• [6.1] Architecture — 兼容性：公开 API/持久数据/配置/环境迁移安全 → issue 运行逻辑依赖 FlashInfer 私有属性
  新增逻辑依赖 FlashInfer 私有属性，依赖版本升级时兼容性不稳。
• [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue 动态 workspace 分支缺少大尺寸回归覆盖
  动态 workspace 扩容、plan_info 校验和 forbid_realloc 分支没有新增覆盖。
• [6.1] Tests — 边界 case 覆盖（空、单元素、最大值） → issue 动态 workspace 分支缺少大尺寸回归覆盖
  required_bytes 超过默认 128MB、估算不足和 replay 禁止扩容等边界未见覆盖。

RTP-LLM Checklist

• [D] 性能 — per-forward C++→Python 调用引入 GIL + GPU sync → issue CUDA graph replay 前新增 GPU 到 CPU 同步
  CudaGraph replay 前会进入 Python prepare_cuda_graph，新增 estimator 在该链路读取 CUDA tensor 到 CPU。
• [D] 性能 — 设备同步 .item()/.cpu()/.tolist() 在关键路径 → issue CUDA graph replay 前新增 GPU 到 CPU 同步
  qo_indptr.detach().to("cpu").tolist() 位于 prepare/prepare_cuda_graph 链路。
• [D] 性能 — GIL 在 hot path 上被持有 → issue CUDA graph replay 前新增 GPU 到 CPU 同步
  prepare_cuda_graph 是 replay 前 Python 调用，新增 CPU 回读延长了该 GIL 持有窗口。

Python Static-First Checklist

• [P.A] 静态结构与类型纪律 — 禁止 getattr/setattr literal 访问 → issue 运行逻辑依赖 FlashInfer 私有属性
  _新增 getattr(..., "_backend"/"max_total_num_rows") 访问 FlashInfer 私有属性。
• [P.G] 测试规范 — 数据驱动测试用 pytest.mark.parametrize → checklist-only
  diff 未新增测试；大尺寸 workspace 覆盖缺口已由独立 issue 记录，此处不按 pytest 风格单独升级。
• [P.G] 测试规范 — pytest.raises 带 match 参数 → checklist-only
  新增 RuntimeError 分支未配套 pytest.raises 覆盖；主测试缺口已由 issue 记录。

Strengths

• workspace 池按 device 和容量筛选 buffer，降低跨设备或尺寸不匹配复用风险。
• 新增 plan 后校验和 forbid_realloc 失败路径，错误信息包含 estimated/actual/current 等定位信息。

[LLLLKKKK]

AI Code Review - PR #1099

Status: BLOCKING

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

Blocking Issues

P1

• prepare 热路径新增 CUDA 到 CPU 同步 @ rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:255
  • 建议：改用已有 host metadata 计算长度，或只在确需扩容时缓存/同步估算，避免每次 prepare 回读 CUDA tensor。
• _plan_info 为 list 时会在 workspace 校验崩溃 @ rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:333
  • 建议：兼容 list/tuple/Tensor：仅 Tensor 调用 detach()，其他序列转 list[int]；补充旧 FlashInfer/fa2 prefill 覆盖。

Non-blocking Suggestions

P2

• workspace 扩容逻辑缺少回归测试 @ rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:596
  • 建议：补充 PyFlashInfer paged prefill 单测，覆盖扩容成功、replay 禁止扩容报错、plan 后实际 workspace 校验和旧 FlashInfer plan_info 类型。

Checklist Violations (6 fail / 97 total)

General Principles Checklist

• [6.1] Architecture — 兼容性：公开 API/持久数据/配置/环境迁移安全 → issue ``_plan_info 为 list 时会在 workspace 校验崩溃
  _line 333 假定 FlashInfer `plan_info` 一定是 Tensor；0.2.5/fa2 路径可返回 List[int]，缺少兼容处理。
• [6.1] Tests — 新逻辑有聚焦单测 + 相关集成/smoke 测试 → issue workspace 扩容逻辑缺少回归测试
  新增 workspace sizing/realloc/replay/plan_info 校验，但 diff 未增加测试文件覆盖关键分支。
• [6.1] Tests — 边界 case 覆盖（空、单元素、最大值） → issue workspace 扩容逻辑缺少回归测试
  缺少 >128MB、估算低估、replay forbid_realloc、plan_info=list 等边界路径覆盖。

RTP-LLM Checklist

• [B] 正确性与逻辑 — 接口返回类型变更有兼容处理 → issue ``_plan_info 为 list 时会在 workspace 校验崩溃
  _新增校验把第三方 `plan_info` 当 Tensor 处理，未兼容 FlashInfer 0.2.5 返回 List[int] 的路径。
• [B] 正确性与逻辑 — 边界 case（空输入、单元素、最大值） → issue workspace 扩容逻辑缺少回归测试
  最大 workspace、估算低估、replay 禁止重分配和旧 plan_info 类型缺少验证。
• [D] 性能 — 设备同步 .item()/.cpu()/.tolist() 在关键路径 → issue prepare 热路径新增 CUDA 到 CPU 同步
prepare() 新增对 CUDA qo_indptr 的 .to("cpu").tolist()，attention plan 前会强制 host/device 同步。

Strengths

• workspace pool 按 device 复用，降低跨 GPU 误用 workspace buffer 的风险。
• plan 后校验实际 workspace 大小，有助于更早暴露 FlashInfer workspace 估算不足。

[LLLLKKKK]

AI Code Review - PR #1099

Status: BLOCKING

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

Blocking Issues

P1

• FA2 workspace estimator 的 CTA tile 选择与实际 plan 不一致 @ rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:275
  • 建议：按实际 FlashInfer 版本/plan 规则选择 CTA tile，或 plan 失败后扩容重试，并用真实 plan_info 覆盖旧版与新版规则。

Non-blocking Suggestions

P2

• prefill workspace 估算会在热路径同步 CUDA tensor @ rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:549
  • 建议：优先使用 CPU 侧 indptr/host mirror 参与 workspace 估算，仅在没有 host 数据时显式回读，并避免默认允许 CUDA 同步。

Checklist Violations (7 fail / 92 total)

General Principles Checklist

• [6.1] Architecture — 状态不变量：创建/更新/失败/重试/回滚路径有效 → issue FA2 workspace estimator 的 CTA tile 选择与实际 plan 不一致
  estimator 的 CTA tile 规则与实际 FlashInfer plan 不一致，capture 阶段记录的 workspace upper bound 可能小于 plan 实际需求。
• [6.1] Architecture — 兼容性：公开 API/持久数据/配置/环境迁移安全 → issue FA2 workspace estimator 的 CTA tile 选择与实际 plan 不一致
  新增 attention 路径依赖 FlashInfer FA2 planner 常量，但未对仓库仍支持的旧 FlashInfer 规则做兼容分支。
• [6.1] Tests — 分布式/跨平台变更有对应覆盖 → issue FA2 workspace estimator 的 CTA tile 选择与实际 plan 不一致
  改动影响 CUDA 12.9、cuda_pre_12_9 与 ARM 不同 FlashInfer 版本，测试未覆盖旧版 planner 常量差异。

RTP-LLM Checklist

• [B] 正确性与逻辑 — 逻辑错误、off-by-one、null/zero 检查 → issue FA2 workspace estimator 的 CTA tile 选择与实际 plan 不一致
  CTA tile 分支使用 64/8 会偏离 FA2 plan 的 workspace 行数计算，触发 actual_bytes > estimated_bytes。
• [D] 性能 — 设备同步 .item()/.cpu()/.tolist() 在关键路径 → issue prefill workspace 估算会在热路径同步 CUDA tensor
  _workspace_tensor_to_host_list 对 qo_indptr 执行 .to("cpu").tolist()，普通 prefill 的 qo_indptr 来自 CUDA cu_seqlens，会在 prepare/plan 热路径同步。
• [D] 性能 — 硬编码容量值 → issue FA2 workspace estimator 的 CTA tile 选择与实际 plan 不一致
  workspace 估算硬编码 CTA tile 和 cta_tile_q==8 的 grid 放大策略；这些值需要与 FlashInfer plan 完全一致。
• [H] 测试与 CI — 测试覆盖充分：大重构等价覆盖，新功能端到端测试 → issue FA2 workspace estimator 的 CTA tile 选择与实际 plan 不一致
  新增测试覆盖 fake wrapper/单套估算规则，但没有覆盖 cuda_pre_12_9/ARM 的 FlashInfer 0.2.5 planner 常量差异。

Strengths

• 为 PyFlashInfer workspace 扩容、CUDA graph replay 上界和 wrapper layout 校验补充了聚焦单测。

[LLLLKKKK]

AI Code Review - PR #1099

Status: BLOCKING

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

Blocking Issues

P1

• workspace retry 会把 int workspace 失败误判为 float workspace 扩容 @ rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:315
  • 建议：缩窄 retry 匹配，只对 tmp_v/tmp_s/float_workspace 扩容；或同步扩容 int workspace，并补充 int-workspace failure 测试。
• workspace 超限时 normal prefill 没有重分配恢复 @ rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:499
  • 建议：forbid_realloc=False 时按 actual_bytes 扩容并重新 plan；仅在 CUDA graph replay 或扩容后仍不足时抛错。

Non-blocking Suggestions

P2

• fa2 workspace 估算在 prefill prepare 热路径做 CPU 回读 @ rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:283
  • 建议：避免每次 prepare 回读 GPU tensor；只在 capture/扩容诊断路径读取，或使用 FlashInfer 直接返回的 CPU workspace 元数据。

Checklist Violations (8 fail / 92 total)

General Principles Checklist

• [6.1] Architecture — 状态不变量：创建/更新/失败/重试/回滚路径有效 → issue workspace 超限时 normal prefill 没有重分配恢复
  normal prefill 允许 realloc，但 actual_bytes > current_bytes 后直接抛错，没有扩容并重试 plan。
• [6.1] Architecture — 错误语义：fail-fast/retry/fallback/silent 行为显式 → issue workspace retry 会把 int workspace 失败误判为 float workspace 扩容
  通用 "Failed to allocate memory" 会进入 float workspace retry；int workspace 或通用 OOM 失败不会被正确区分。
• [6.1] Tests — 边界 case 覆盖（空、单元素、最大值） → issue workspace 超限时 normal prefill 没有重分配恢复
  新增测试覆盖 overflow 抛错，但缺少超过默认 workspace 后 normal prefill 扩容成功的最大值场景。

RTP-LLM Checklist

• [B] 正确性与逻辑 — 逻辑错误、off-by-one、null/zero 检查 → issue workspace 超限时 normal prefill 没有重分配恢复
  actual_bytes > current_bytes 的 normal prefill 分支未走扩容恢复，直接抛错。
• [B] 正确性与逻辑 — 边界 case（空输入、单元素、最大值） → issue workspace 超限时 normal prefill 没有重分配恢复
  缺少真实 FlashInfer 下超过默认 workspace 后应扩容继续执行的最大值路径。
• [D] 性能 — 设备同步 .item()/.cpu()/.tolist() 在关键路径 → issue fa2 workspace 估算在 prefill prepare 热路径做 CPU 回读
  fa2 workspace sizing 在每次 prefill prepare 后对 plan_info 执行 to("cpu").tolist()。
• [D] 性能 — 硬编码容量值 → issue workspace retry 会把 int workspace 失败误判为 float workspace 扩容
  _retry 只增长 float workspace，继续复用固定 int_workspace_buffer，无法处理 int workspace 容量不足。
• [H] 测试与 CI — 测试覆盖充分：大重构等价覆盖，新功能端到端测试 → issue workspace 超限时 normal prefill 没有重分配恢复
  新增测试覆盖 helper/fake wrapper，但没有覆盖真实 normal prefill workspace 扩容后的成功路径。

Strengths

• 新增 workspace sizing 单测覆盖 plan_info 解析、CUDA graph replay 上界、retry 路径和 FlashInfer wrapper 兼容性失败路径。

[LLLLKKKK]

AI Code Review - PR #1099

Status: BLOCKING

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

Blocking Issues

P1

• FlashInfer 0.2.5 环境会在导入/初始化时失效 @ rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:10
  • 建议：为旧 FlashInfer 做 lazy import/版本门禁和旧 wrapper fallback，或同步移除 cuda_pre_12_9 支持并调整依赖/BUILD。
• 非 replay 路径的 workspace 溢出只报错不扩容 @ rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:499
  • 建议：forbid_realloc=False 时按 actual_upper_bound_bytes 扩容 workspace 并重新 plan，仅 replay 禁止重分配时抛错。

Non-blocking Suggestions

P2

• workspace retry 会误处理 int workspace 容量错误 @ rtp_llm/models_py/modules/factory/attention/cuda_impl/py_flashinfer_mha.py:315
  • 建议：只匹配 batch_prefill_tmp_v/tmp_s/float_workspace，或同时扩容 int workspace，并补充 int workspace 失败用例。

Checklist Violations (9 fail / 99 total)

General Principles Checklist

• [6.1] Architecture — 状态不变量：创建/更新/失败/重试/回滚路径有效 → issue 非 replay 路径的 workspace 溢出只报错不扩容
  py_flashinfer_mha.py:499 对非 replay 的 actual_bytes > current_bytes 直接抛错，缺少扩容和重新 plan，状态恢复路径不完整。
• [6.1] Architecture — 错误语义：fail-fast/retry/fallback/silent 行为显式 → issue workspace retry 会误处理 int workspace 容量错误
  py_flashinfer_mha.py:310 的错误匹配过宽，plan retry 不能区分 float workspace 与 int workspace 容量错误。
• [6.1] Architecture — 兼容性：公开 API/持久数据/配置/环境迁移安全 → issue FlashInfer 0.2.5 环境会在导入/初始化时失效
  cuda_pre_12_9 仍使用 flashinfer-python==0.2.5，但新增 import 和 wrapper 私有字段校验按新版 FlashInfer 布局执行。
• [6.1] Tests — 边界 case 覆盖（空、单元素、最大值） → issue 非 replay 路径的 workspace 溢出只报错不扩容
  新增测试覆盖 overflow 报错，但未覆盖非 replay 场景 actual workspace 超过当前 buffer 后可恢复扩容继续运行。

RTP-LLM Checklist

• [A] 兼容性与配置 — 可选依赖 lazy import，pybind 新字段有 C++ 默认值 → issue FlashInfer 0.2.5 环境会在导入/初始化时失效
  PyFlashInfer 模块顶层导入 flashinfer.utils，并在初始化强依赖新版 wrapper 私有成员，旧 FlashInfer 缺少门禁。
• [B] 正确性与逻辑 — 逻辑错误、off-by-one、null/zero 检查 → issue 非 replay 路径的 workspace 溢出只报错不扩容
  actual workspace 大于 current 时普通路径缺少 resize/replan 分支，只直接抛 RuntimeError。
• [B] 正确性与逻辑 — 边界 case（空输入、单元素、最大值） → issue 非 replay 路径的 workspace 溢出只报错不扩容
  大 batch/长序列导致 plan_info 上界超过默认 workspace 时，非 replay 路径仍不可恢复失败。
• [D] 性能 — 硬编码容量值 → issue 非 replay 路径的 workspace 溢出只报错不扩容
  超过默认 workspace 后没有按 plan_info 上界扩容，实际容量仍受当前 buffer 硬限制。
• [H] 测试与 CI — 测试覆盖充分：大重构等价覆盖，新功能端到端测试 → issue FlashInfer 0.2.5 环境会在导入/初始化时失效
  新增 fake wrapper 测试未覆盖 cuda_pre_12_9 + flashinfer-python==0.2.5 的真实导入/初始化兼容路径。

Strengths

• 新增 workspace sizing 单测覆盖 plan_info 估算、扩容重试和 CUDA graph replay 禁止扩容路径。