Repository navigation
fix: break attention state ownership cycles - #1640
Merged
Merged
Conversation
Use an explicit weak reference and read-only infer_state property shared by attention states. Check initialized attention metadata for direct InferStateInfo references to prevent ownership cycles. Track decode graph capture and replay separately so eager overlap fallback always initializes FlashInfer wrappers. Add lifetime, graph-copy, wrapper initialization and scheduling regressions, and adapt existing tests to weak ownership.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
InferStateInfo 持有 attention state,而 attention state 再强引用 InferStateInfo,会形成循环引用。一次 forward 结束后,临时状态及其 CUDA Tensor 需要等待 Python 循环 GC 才能释放;这不一定是永久泄漏,但会延迟显存复用、抬高峰值,并可能触发运行时 OOM。
本 PR 同时修正 FlashInfer decode wrapper 的执行模式判断:decode overlap 中,只要任一 microbatch 超出 graph 上限,整体就会回退 eager。原先逐个 microbatch 判断 graph 是否可运行,会使较短 batch 跳过 wrapper 创建,随后在 eager 执行时访问空 wrapper。
实现
BaseAttState,各后端工厂显式设置_infer_state_weak_ref=weakref.ref(infer_state)。通过只读infer_stateproperty 获取原对象;未绑定时返回None,owner 已释放时抛出ReferenceError。不再接受原来的infer_state=构造参数或直接属性赋值,仓库内调用与测试已适配。InferStateInfo.init_att_state()完成主、辅 attention state 初始化后统一调用check_att_states(),拒绝 state 第一层实例属性直接保存任何InferStateInfo。检查不递归遍历容器和闭包。is_cuda_graph改为is_cuda_graph_capture,增加is_cuda_graph_replay。单 batch 与 overlap 均在初始化 attention state 前,根据实际执行分支设置标志。FlashInfer FP、FP8、MLA 共用该判断:两个标志同时为 True 时明确报错。现有 prefill/decode CUDA Graph 缓存继续强持有捕获用 InferStateInfo;decode 回放复制 Tensor metadata,并保留捕获 state 的弱引用绑定。
回归检查
git diff --check通过;GitHub Pre-commit 检查通过。真实模型验证
环境:NVIDIA H200、BF16、TP1。
Graph 精度与执行路径
矩阵使用 page_size=1、FlashInfer FP;hybrid 线性层使用 Triton GDN。Qwen3-4B、Qwen3.5-0.8B 分别测试 decode graph 开/关 × prefill graph 开/关四种组合,共记录 392 次真实 forward 输出。与同模型全 eager 路径比较,所有 logits 满足
rtol=0.01, atol=0.05,greedy top1 全部一致;最大绝对差为 0.0625。覆盖首次捕获、重复命中、batch 1→2 / 3→4 padding、batch 5 超限回退、KV 长度 160 超限回退、prefill token 数超限回退,以及分块前缀 KV 复用。测试配置为 decode 最大 batch 4、最大长度 128,prefill 最大处理 token 数 32。Hybrid 模型按生产流程,使用 HOLD request 预热并预捕获配置桶后验证真实请求。
显存问题复现
Qwen3.5-27B,4096-token prefill,开启 prompt logprobs,两个 graph 均关闭,默认 Python GC,不限制进程分配器预算。两组各启动新进程,并一致关闭新增的强引用检查以构造对照;强引用组仅额外恢复 attention → InferStateInfo 的强引用,数值计算路径相同。该检查在生产代码中保持开启。
强引用组结束后手动 GC 释放约 45.49 GiB,并清除滞留 owner;弱引用组每轮结束后的滞留 owner 为 0,结束后 GC 未额外释放 GPU Tensor 显存。正常迭代中不手动 GC。具体 OOM 轮数受默认 GC 触发节奏影响,不是服务请求阈值。
真实 HTTP
四个全新 Qwen3.5-27B 服务:强/弱引用 × decode graph 开/关;HTTP 验证中 prefill graph 关闭,其兼容性由上面的模型矩阵验证。统一配置为 256 请求槽位、65536 KV tokens、动态缓存和 chunked prefill 开启。强引用由实验探针注入,只豁免该实验属性的防护检查,其他初始化检查仍执行。
每组包含 32 个 decode-heavy 请求(并发 8,32 输入 / 256 输出)和 64 个长输入请求(串行,4096 输入 / 1 输出,返回 prompt logprobs)。共 384/384 个正式请求成功,无 HTTP OOM。
开启 graph 时,强/弱引用两组均保留正常的 8 个缓存状态,已从非缓存计数中排除。强弱引用对照及弱引用 graph 开关对照中,每对 96 个请求的生成文本均一致;decode 选中 token 的 logprob 最大差约 4.67e-5,长输入的生成及 prompt logprobs 对照差为 0。
活跃 Tensor 与 NVML 使用不同统计口径,NVML 包含 allocator 缓存和 CUDA context,差值不能直接视为泄漏量。探针不在 HTTP 请求路径手动 GC 或同步 CUDA,但 Python 记账可能影响 GC 时机;结果限定于上述配置。
验证边界