Skip to content

Work: 自动组装 Numba 融合内核并接入 G1WalkFlat #922

Description

@TATP-233

一句话问题

G1 walk 的 Numba 路径由手写固定 TERM_ORDER 和整段融合内核维护,配置变化不能自动形成一致的 fused kernel,且 observation noise 仍不受支持。

为什么现在做

最小交付结果

实现由 trusted resolved plan 自动组装的 prange fused kernel,并让 G1WalkFlatnumba_acceleration=true 时使用它;达到 parity/performance gate 后移除 G1 手写整段 aggregator 和固定名单维护。

In scope

  • 根据 canonical ordered term/preamble、workspace layout、dtype 和 implementation version 生成稳定 cache key。
  • 在冷路径组装并 JIT 一个融合 kernel;注册的 Numba item function 被直接调用,热路径无 Python per-term dispatch。
  • scale/parameters 作为运行时参数;零 scale 用显式、可测试的 runtime branch 跳过数学。
  • 同进程相同 plan 复用已编译 callable;记录首次 JIT 和 cache hit 时间,不建设长期 benchmark 服务。
  • 接入 G1 reward/obs/termination 与既有 noise 输入;noise 随机数仍由 task owner 产生,不新增 RNG contract。
  • parity/SPS 达标后删除 G1 task-owned 手写 aggregation、重复 term order 和只服务旧路径的代码。

Non-goals

Owner 与预计改动

  • Owner:term Numba assembler/cache、G1 registration/adapter 和局部 benchmark/tests。
  • 预计 9–15 个文件,≤800 行净手写改动,1 个 PR。
  • PR base:refactor/issue-902-mjwarp-squashed-delivery;不得向 main 提 PR。

Acceptance criteria

  • 相同 resolved plan 产生稳定 cache key;term 顺序、preamble、workspace、dtype 或实现版本变化会 miss。
  • G1 numba_acceleration=true 使用自动 fused kernel,不再调用手写整段 aggregator。
  • 未知/缺失 Numba implementation 在 materialization 边界 fail closed,或按配置显式选择 Tier 1;不得静默混跑。
  • reward/obs/termination 与 Tier 1 在冻结容差内一致;noise 开启时 shape、dtype 和 actor/critic clean/noisy 边界正确。
  • N=2048/8192 SPS ≥ 迁移前 G1 手写 Numba baseline 的 90%。
  • 首次 JIT 与同进程 cache-hit 数据记录在 PR;局部 benchmark 不进入常规 CI gate。
  • 局部测试通过;创建/更新 PR 前通过 make test-all

Stop conditions

  • 超过 15 个文件或 800 行净手写改动。
  • SPS 低于 baseline 90%,或生成方案需要新 runner/execution path。
  • 稳定 cache key 需要读取 hot-path asset/backend metadata。
  • 为支持 G1 必须引入通用 DAG、graph scheduler 或表达式语言。

Dependencies

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

    enhancementNew feature or request

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions