Skip to content

Perf: mjwarp 后端引入 CUDA graph capture #1271

Description

@TATP-233

一句话问题

当前 mjwarp 后端的稳态物理热路径逐次提交 CUDA kernel,尚未复用 CUDA graph,固定形状的大批量仿真存在可避免的 launch 开销。

为什么现在做

Roadmap #1252 对照 mjlab v1.6.0 后确认,UniLab 的 src/unilab/base/backend/mjwarp/backend.py 目前没有 graph capture/replay;本任务只补齐这一后端内部性能机制,并用同机同配置基准验证实际收益。

最小交付结果

符合运行条件时,mjwarp 后端对固定形状稳态热路径使用 CUDA graph replay;条件不满足或 capture 失败时透明回退到现有执行路径,物理与后端契约不变。

In scope

  • 在 mjwarp backend 冷路径检测 CUDA driver ≥ 12.4 与 mempool/capture 条件。
  • 在后端内部建立固定形状热路径的 capture、replay 与必要失效处理。
  • 保留普通执行路径,并为不支持条件与 capture 失败提供可诊断的安全降级。
  • 增加靠近 backend 风险边界的单元/集成测试。
  • 用同一硬件、task、world 数与测量方法报告修改前后吞吐。

Non-goals

Owner 与预计改动

Owner 为 src/unilab/base/backend/mjwarp/;预计 2–6 个文件、≤500 行净手写改动、1 个 PR。合并后的永久维护项仅为 mjwarp 内部 graph eligibility、capture/replay 与 fallback 路径。

Acceptance criteria

  • 支持条件下,稳态固定形状路径在 warm-up/capture 后通过 graph replay 执行。
  • driver/mempool 条件不满足或 capture 失败时,后端继续走现有路径且行为正确。
  • capture 生命周期不会绕开现有 model mutation/reset/step 契约。
  • 局部 correctness/fallback 测试通过,make test-all 通过。
  • PR 中记录可复现的修改前后吞吐、硬件/软件条件与测量参数。

Stop conditions

若需要新增公共 contract、改动 env/runner/learner、引入新的执行生命周期、修改常规 CI,或超过 6 文件/500 行预算,则暂停并回到 maintainer 决策。

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions