1. 背景

当前 torch_npu inductor 动态 shape 的 gear 分档主要在初始化阶段确定，后续运行过程中基本保持不变。这种静态分档策略在实际 workload 中可能暴露出两类问题：

1. 部分档位利用率低：初始化的档位分布可能与实际 shape 分布不匹配，导致某些档位长期闲置，浪费预留资源
2. padding 比例过高：档位设置不合理时，大量请求会被映射到偏大的档位，导致受控维度上持续的高比例 padding，降低计算效率
   为解决上述问题，需要引入自适应档位更新机制，根据运行时统计动态调整档位分布，在保持稳定性的同时提升资源利用率。

系统在运行期间持续回答三个问题：当前哪些 gear 值得保留、哪些可以安全删除、哪些区域值得新增。

2. 设计目标

1. 优先保留近期命中率高、padding/split 代价低的 gear
2. 在有限资源预算下，让 gear 集合朝更高收益方向演化
3. 支持单维 BATCHSIZE / SEQLEN 及双维组合场景
4. 更新不破坏请求一致性：通过锁保证原子性，通过快照保证请求线程不受影响

3. 总体运行时流程

3.1 端到端闭环

flowchart LR
    A[请求进入 shape handling] --> B[读取当前 gears 快照]
    B --> C[按 current gears 执行 transform]
    C --> D[记录 raw shape / mapped gear / pad / split]
    D --> E[执行 compiled graph]
    E --> F[执行 recover]
    F --> G[请求返回]

    H[Worker 线程定时醒来] --> I{检查触发条件}
    I -- 距上次更新 ≥ 300s --> J[执行更新]
    I -- 显存使用率 ≥ 阈值 --> J
    I -- 都不满足 --> H
    J --> K[加锁: 获取快照 + 构建统计 + 打分 + 淘汰 + 新增 + 预算判断]
    K --> L[更新 current gears 快照]
    L --> M[调用 pool.remove_by_keys 清理 graph]
    M --> H

3.2 职责划分

• 请求线程：get_snapshot() 读快照 → transform → src_fn(compile+execute) → consume_recent_keys() → record_event → recover
• Worker daemon 线程：每 60s 检查触发条件 → run_once（在 _commit_lock 下完成打分、淘汰、新增、提交）→ 锁外 pool.remove_by_keys

sequenceDiagram
    participant Req as 请求线程
    participant RT as AdaptiveGearRuntime
    participant Worker as Worker 线程 (daemon)
    participant WK as GearUpdateWorker

    Note over Worker: 每 60s 醒来一次

    Req->>RT: get_snapshot() (无锁读)
    RT-->>Req: gears 快照副本
    Req->>RT: record_event(...)
    RT-->>Req: 返回
    Note over Req,RT: 请求线程只做快照读取和事件记录

    Worker->>Worker: _should_trigger_update()
    alt 触发条件满足
        Worker->>WK: run_once(now_ts)
        Note over WK: 在 _commit_lock 保护下
        WK->>RT: get_snapshot() + build_stats_snapshot()
        WK->>WK: score / evict / add / budget
        WK->>RT: commit_update()
        Note over WK: 释放 _commit_lock
        WK->>WK: pool.remove_by_keys(keys)
    else 条件不满足
        Worker->>Worker: 继续等待
    end

3.3 并发控制

锁	保护对象	使用场景
_sample_lock	_states 读写	record_event、build_stats_snapshot、_collect_cleanup_keys
_commit_lock	更新计算串行化	run_once 全程
GearSnapshotStore._lock	快照写入	publish_snapshot

get_snapshot() 无锁（clone-on-read），请求线程不阻塞。graph 清理在所有锁外执行。

3.4 Graph 清理机制

图资源通过 GraphResourcePool（per-device 单例）管理，使用不透明整数句柄作为 key：

pool = GraphResourcePool.get_pool(device_index)

key = pool.register(resource)          # 下层 — graph capture 期间 → int
keys = pool.consume_recent_keys()      # 上层 — src_fn 返回后 → List[int]
pool.remove_by_keys([42, 43])          # 后台 — gear 淘汰后 (幂等)

只有下层构造 key，上层通过 consume_recent_keys() 读回，从根本上消除 key 一致性问题。

完整生命周期：

sequenceDiagram
    participant NF as new_fn (上层)
    participant SF as src_fn
    participant RF as record_function (下层)
    participant Pool as GraphResourcePool
    participant MGR as AdaptiveGearRuntime
    participant WK as GearUpdateWorker (后台)

    Note over NF: variant 0: BATCHSIZE→32, SEQLEN→128

    NF->>SF: src_fn(call_args)
    activate SF

    loop 每个子图 (graph break)
        SF->>RF: record_function(sub_inputs, func_id)
        RF->>RF: node = NPUGraphNode(...)
        RF->>Pool: key = pool.register(node)
        Note over Pool: _entries[key] = node<br/>_pending_by_thread[tid].append(key)
        Pool-->>RF: key (int, e.g. 100)
    end

    SF-->>NF: result
    deactivate SF

    NF->>Pool: keys = pool.consume_recent_keys()
    Pool-->>NF: [100, 101, 102]
    Note over Pool: _pending_by_thread[tid] 已清空

    loop 每个 key
        NF->>MGR: record_event(..., cleanup_key=key)
        MGR->>MGR: state["BATCHSIZE:32"].cleanup_keys.add(key)
        MGR->>MGR: state["SEQLEN:128"].cleanup_keys.add(key)
    end

    Note over WK: --- 后台：gear 淘汰 (每 60s) ---

    WK->>MGR: commit_update(removed_gears=["BATCHSIZE:32"])
    MGR->>MGR: 收集 cleanup_keys = {100, 101, 102}
    MGR-->>WK: removed_keys = [100, 101, 102]

    WK->>Pool: remove_by_keys([100, 101, 102])

    loop 每个 key
        Pool->>Pool: resource = _entries.pop(key, None)
        alt resource 存在
            Pool->>Pool: torch.npu.synchronize()
            Pool->>Pool: resource.release() / reset()
        else resource 不存在 (已删)
            Note over Pool: pop 返回 None，跳过
        end
    end

线程隔离：register 和 consume_recent_keys 构成 per-thread 生产者-消费者对。register 在 _lock 下自增分配 key，追加到 _pending_by_thread[tid]；consume 排空当前线程的 pending 列表。不同线程互不干扰，同一函数断图产生的 N 个 key 在同一线程上一次性全部消费。

Key-Gear 关联：一个 variant 内产生的所有子图 key 挂到该 variant 命中的全部 gear 上。任意 gear 淘汰后该 shape 组合不再生成，所有依赖子图都应清理。收集到的 key 传给 remove_by_keys，使用 _entries.pop(key, None) 保证幂等——同一 key 被多个 gear 共享时，首次淘汰 pop 拿到资源并释放，后续 pop 返回 None 跳过。

4. 时间窗口与统计规则

• 窗口长度 300s（window_seconds），半衰期 150s，衰减系数 $\lambda = \frac{\ln 2}{150}$
• 事件权重：$w_t(e) = \exp(-\lambda \cdot \Delta t_e)$

对 gear $g$：

$$hit\_rate_g = \frac{\sum_{e \in g} w_t(e)}{\sum_{e \in all} w_t(e)} \qquad avg\_pad_g = \frac{\sum_{e \in g} w_t(e) \cdot pad\_ratio(e)}{\sum_{e \in g} w_t(e)} \qquad avg\_split_g = \frac{\sum_{e \in g} w_t(e) \cdot split\_ratio(e)}{\sum_{e \in g} w_t(e)}$$

双维场景按维度独立统计。

5. 基础打分

单一公式：$Score_g = w_h H_g - w_p P_g - w_{sp} SP_g$

默认权重：$w_h = 0.60$，$w_p = 0.20$，$w_{sp} = 0.20$。命中率越高分越高，padding/split 越高分越低。双维按两个独立一维问题处理。

6. 淘汰机制

触发：定时（距上次 ≥ update_interval_seconds，默认 300s）或显存压力（≥ device_memory_usage_threshold_ratio）。

多层过滤（每维度至多淘汰 1 个）：

1. 跳过最大 gear（max_gear_by_type）
2. 使用保护：now - max(created_ts, last_hit_ts) < recent_use_protect_seconds（默认 300s）→ 跳过
3. 数量下限：len(gears) ≤ min_gear_count_per_type → 跳过
4. 替代损失：计算删除后 raw_samples 重映射到最佳替代 gear 的平均损失，超过 replace_loss_threshold（0.60）→ 跳过并延长保护

flowchart TD
    A[按维度分组构建候选队列] --> B{候选为空?}
    B -- 是 --> Z[结束]
    B -- 否 --> C[按维度遍历]
    C --> D{该维度 gear 数 > 下限?}
    D -- 否 --> C
    D -- 是 --> E[过滤: 跳过最大 gear / 使用保护窗口过滤]
    E --> F{过滤后仍有候选?}
    F -- 否 --> C
    F -- 是 --> G[取最低分候选]
    G --> H{替代损失超过阈值?}
    H -- 是 --> I[延长保护窗口]
    I --> C
    H -- 否 --> J[标记 gear 为待淘汰]
    J --> K[继续下一个维度或结束]
    K --> Z

提交阶段：加锁更新 current_gears，收集 cleanup_keys，删除 GearRuntimeState。清理阶段：锁外 pool.remove_by_keys(keys)。

7. 新增机制

两类独立触发（共用同一份窗口统计）：

• padding 驱动：pad_sample_count ≥ add_min_samples 且 avg_pad ≥ pad_add_threshold（默认 0.35）
• split 驱动：split_sample_count ≥ add_min_samples 且 avg_split ≥ split_add_threshold（默认 0.20），仅对最大 gear 生效

候选值 = median(raw_samples)，去重后按 pressure 排序。显存超阈值时 budget = len(removed_gears)（替换式新增），否则不限制。每个 shape_type 上限 max_gears_per_type（默认 64）。

flowchart TD
    A["遍历所有 gear 的窗口统计"] --> B{"当前 gear 是最大档位?"}
    B -- 是 --> C["检查 split:<br>avg_split >= theta 且样本数 >= N_min?"]
    B -- 是 --> D["检查 pad:<br>avg_pad >= theta 且样本数 >= N_min?"]
    B -- 否 --> D
    C -- 是 --> E["g_new = median(raw_samples)<br>候选类型: 新最大档位"]
    D -- 是 --> E
    C -- 否 --> F["不触发 split 新增"]
    D -- 否 --> G["不触发 pad 新增"]
    E --> H["加入候选集合"]
    H --> I["所有候选按 pressure 排序"]
    I --> J{"显存压力高?"}
    J -- 是 --> K["budget = 已淘汰 gear 数<br>(替换式新增)"]
    J -- 否 --> L["budget = 无限制<br>(直接新增)"]
    K --> M
    L --> M["按 pressure 顺序遍历候选"]
    M --> N{"budget 耗尽?"}
    N -- 是 --> O["结束"]
    N -- 否 --> P["加入 active_gears<br>budget -= 1"]
    P --> M

8. 资源预算

• device_memory_usage_ratio = 1 - free / total，通过 torch.npu.mem_get_info() 获取
• 达到 device_memory_usage_threshold_ratio（默认 0.90）时立即触发更新，新增转为替换式（addition_budget = len(removed_gears)）
• 配置为 None 时不启用显存检查

9. 测试覆盖

测试文件：test/_inductor/test_shape_handling.py（63 个用例）、test/npu/test_graph_tree.py（15 个用例）。

元数据采集 ─── transform metadata → raw/mapped shapes → pad/split ratios

统计采集 ─── snapshot 创建 → record_event → build_stats
                                     ├─ pad/split 分离统计
                                     ├─ 双维独立统计
                                     └─ cleanup_keys 追踪 (Set[int])

淘汰决策 ─── 基础打分 → 保护过滤 → 替代损失 → 提交
                 ├─ 最近使用/新建保护
                 ├─ 零使用优先淘汰
                 └─ 最大 gear 跳过

新增决策 ─── pad 驱动 / split 驱动 → 候选生成 → 重复过滤 → 提交

资源守卫 ─── 显存压力 ─┬─ 触发更新
                        └─ 阻止直接新增

并发安全 ─── daemon 线程 ── _commit_lock 串行化 ── 快照隔离 ── 高并发稳定性

GraphResourcePool ─── register → consume_recent_keys → remove_by_keys
                          ├─ per-thread 线程隔离
                          └─ 幂等删除

10. 实现边界限制

1. 双维场景分维独立决策，不构造联合候选
2. 显存使用率基于当前观测值，不预测新增后的增量
3. 新增候选仅做精确重复检查，不做距离阈值过滤
4. 每轮每维度至多淘汰 1 个 gear，数量远超下限时需多轮收敛

11. 总结

本方案引入运行时自适应齿轮调整机制：通过时间衰减窗口统计量化 gear 价值，经多层过滤实现可解释的淘汰决策，由 padding/split 信号驱动新增，在显存预算约束下控制 gear 数量。请求线程无锁读取快照，后台 daemon 线程自主触发更新，graph 清理在锁外幂等执行。