【Bug已解决】[WebGPU EP] Leak Detector failure when running GenAI LLM models 解决方案

【Bug已解决】[WebGPU EP] Leak Detector failure when running GenAI LLM models 解决方案一、现象长什么样开启 ONNX Runtime 内置的内存泄漏检测器ORT_LEAK_DETECTOR1或 debug 构建后用 WebGPU EP 跑 GenAI LLM 模型如 Llama/Qwen 的对话生成带 KV cache 反复复用泄漏检测器在推理结束时误报泄漏或直接崩溃而实际并没有内存泄漏——那些被报“泄漏”的其实是 LLM 在生成过程中故意常驻的 KV cache / 权重 buffer。现象# 现象 A推理结束报“泄漏”但实际没漏 # LeakDetector: 123 allocations not freed at session cleanup # 但这些是 KV cache 的持久 buffer本就应跨 step 存活 # 现象 B泄漏检测器自身崩溃 # double free / heap corruption at LeakDetector::Check() # 因为 WebGPU 的分配是异步的GPU 侧 free 延后 # 检测器在 CPU 侧计数已减、GPU 侧还没真正释放时就报不一致 # 现象 C只在 WebGPU EP GenAI LLM 触发 # CPU EP 跑同样模型泄漏检测正常WebGPU EP 跑非 LLM无 KV cache也正常最坑的是现象 ACI 里开了泄漏检测当门禁结果 WebGPU EP LLM 永远误报要么门禁被无奈关掉失去意义要么大家绕着走不测 WebGPU LLM掩盖了真实泄漏。二、背景ONNX Runtime 的泄漏检测器LeakDetector在 debug 构建下对所有Allocator::Alloc/Free做配对计数每次 Alloc 记一笔Free 减一笔session 销毁时若还有未 Free 的就报泄漏。它假设分配/释放是同步且立即配对的。WebGPU EP 跑 LLM 时有两个特点打破这个假设KV cache 是常驻的生成第 1 个 token 时分配 KV cache buffer之后每个 token 复用session 结束时才统一释放。泄漏检测器若在每个 step 结束时检查会看到“未释放的 KV cache”误报为泄漏。WebGPU 分配是异步的WebGPU 的destroy()是把释放命令排队到 GPU 队列CPU 侧Free调用返回了GPU 侧真正释放延后。泄漏检测器若在Free返回瞬间就减计数并立即 Check会和 GPU 实际状态不一致甚至 double-free。这是泄漏检测/异步资源审查里典型的坑同步假设的泄漏检测器遇到“常驻 buffer 异步释放”的 WebGPU LLM误报或崩溃。三、根因泄漏检测器把常驻 KV cache 当成泄漏它在每个 step/子图结束就 Check而 KV cache 本应跨 step 常驻到 session 结束被误报现象 A。WebGPU 异步释放导致计数不一致Free在 CPU 侧返回但 GPU 侧延后检测器在返回瞬间 Check 计数状态错位甚至 double-free现象 B。缺少“WebGPU LLM”的泄漏检测白名单/延迟 CheckCI 没覆盖这个组合误报长期存在。本质是同步泄漏检测器假设分配/释放立即配对遇到 WebGPU 的常驻 KV cache 与异步释放误报或崩溃且缺组合测试。四、最小可运行复现下面用 Python 模拟“同步泄漏检测器把常驻 KV cache 误报为泄漏 异步释放计数错位”class LeakDetectorBuggy: def __init__(self): self._live 0 def alloc(self): self._live 1 def free_sync(self): self._live - 1 def check_after_step(self): # buggy: 每步结束就 check if self._live ! 0: raise AssertionError(fLEAK: {self._live} allocs unfreed) det LeakDetectorBuggy() det.alloc() # 分配 KV cache常驻 det.check_after_step() # 误报KV cache 本应跨 step 存活 class LeakDetectorFixed: def __init__(self): self._live 0 self._persistent set() # 标记为常驻、不在 step check 计入泄漏 def alloc(self, persistentFalse): self._live 1 if persistent: self._persistent.add(id(self)) def free_async(self): # 异步释放CPU 侧先减GPU 侧延后这里简化为延迟减 self._live - 1 def check_after_step(self): # 常驻 buffer 不计入“泄漏”只有真正游离的才算 pass # KV cache 常驻不误报 def check_at_session_end(self): if self._live ! 0: raise AssertionError(freal leak: {self._live}) df LeakDetectorFixed() df.alloc(persistentTrue) # KV cache 常驻 df.check_after_step() # 不误报 df.free_async() df.check_at_session_end() # 会话结束时才真正校验buggy每步 check 误报 KV cachefixed区分常驻与会话结束才校验。五、解决方案第一层最小直接修复最小修复泄漏检测器的 Check 分两级——step 级只查“游离”分配排除标记为常驻的 KV cachesession 级才查全部且 WebGPU 的异步释放用延迟计数避免错位// 修正常驻 buffer 标记为 persistentstep check 排除session end 才全查 void Allocator::AllocTracked(size_t size, bool persistent) { leak_detector_.OnAlloc(size, persistent); } void LeakDetector::CheckAfterStep() { // 只检查非 persistent 的游离分配 if (live_non_persistent_ ! 0) ORT_THROW(leak: , live_non_persistent_); } void LeakDetector::CheckAtSessionEnd() { if (live_total_ ! 0) // 会话结束才查全部含 KV cache 是否已释放 ORT_THROW(real leak: , live_total_); }这一层改动最小加 persistent 标记 两级 Check误报消失。但依赖“每处分配都标对 persistent”下看第二层。六、解决方案第二层结构性改进把“泄漏检测的常驻豁免 异步释放延迟计数”固化成单一事实来源。下面这个 dataclass 集中管理检测契约from dataclasses import dataclass, field from typing import Dict, Set dataclass class WebGpuLeakDetectorPolicy: 单一事实来源WebGPU LLM 场景的泄漏检测契约。 _live_total: int 0 _persistent: Set[int] field(default_factoryset) _pending_free: int 0 # 异步释放待确认的计数 def alloc(self, buf_id: int, persistent: bool False) - None: self._live_total 1 if persistent: self._persistent.add(buf_id) def free_async(self, buf_id: int) - None: # WebGPU 异步释放先记 pendingGPU 真正释放后再确认 if buf_id in self._persistent: self._persistent.discard(buf_id) self._pending_free 1 def confirm_gpu_free(self) - None: # GPU 队列排空后确认 self._live_total - self._pending_free self._pending_free 0 def check_after_step(self) - None: # step 级只查游离非 persistent且已确认的未释放 freeable self._live_total - len(self._persistent) - self._pending_free if freeable 0: raise AssertionError(fleak detector state corrupt: {freeable}) def check_at_session_end(self) - None: self.confirm_gpu_free() if self._live_total ! 0: raise AssertionError(freal leak at session end: {self._live_total})这一层的关键收益常驻豁免persistent集合让 KV cache 不被 step check 误报异步对齐_pending_free等 GPU 确认后才减live_total避免计数错位/double-free两级 Checkstep 级查游离、session 级查全部语义正确单一事实来源所有泄漏检测约定收口在WebGpuLeakDetectorPolicy。七、解决方案第三层断言 / CI 守护把第二层钉成 pytest挂进 CI覆盖 WebGPU LLM 场景import pytest from your_package.webgpu_leak import WebGpuLeakDetectorPolicy def test_kv_cache_not_false_leak(): # 断言 1常驻 KV cache 不被 step check 误报 p WebGpuLeakDetectorPolicy() p.alloc(1, persistentTrue) # KV cache p.check_after_step() # 不抛异常 def test_async_free_no_corruption(): # 断言 2异步释放不会计数为负/错乱 p WebGpuLeakDetectorPolicy() p.alloc(1, persistentTrue) p.free_async(1) p.check_after_step() # 不报 corrupt p.confirm_gpu_free() p.check_at_session_end() # 全部释放无泄漏 def test_real_leak_detected(): # 断言 3真泄漏游离 buffer 没释放在 session end 被抓 p WebGpuLeakDetectorPolicy() p.alloc(99, persistentFalse) # 游离 buffer 忘了释放 with pytest.raises(AssertionError): p.check_at_session_end() def test_pending_free_deferred(): # 断言 4异步释放在 GPU 确认前不减少 live_total p WebGpuLeakDetectorPolicy() p.alloc(5, persistentTrue) p.free_async(5) assert p._pending_free 1 and p._live_total 1 p.confirm_gpu_free() assert p._live_total 0四条断言从“KV cache 不误报”“异步不混乱”“真泄漏被抓”“pending 延迟”四面把误报/崩溃钉死在 CI。八、排查清单WebGPU EP 开泄漏检测跑 LLM 误报/崩溃时报的“泄漏”是不是 KV cache 这类常驻 buffer是就确认检测器是否把常驻当泄漏现象 A。崩溃在LeakDetector::Check的 double free查 WebGPU 异步释放是否让计数错位现象 B。是否只在 WebGPU LLM 触发是就确认检测器假设了同步释放。用第二层WebGpuLeakDetectorPolicy常驻豁免 异步延迟计数 两级 Check。加第三层 pytest断言“KV cache 不误报、异步不混乱、真泄漏被抓、pending 延迟”。泄漏检测门禁不能因误报就关掉应在检测器侧修对异步/常驻语义。九、小结WebGPU EP 的泄漏检测器在 GenAI LLM 下误报/崩溃本质是同步假设的泄漏检测器把 LLM 常驻的 KV cache 误报为泄漏且 WebGPU 的异步释放让 CPU 侧计数与 GPU 实际状态错位甚至 double-freeCI 误报导致门禁被关或绕开掩盖真实泄漏。修复分三层——第一层加 persistent 标记 step/session 两级 Check第二层用WebGpuLeakDetectorPolicy这个 dataclass 把“常驻豁免 异步延迟计数 两级校验”收口成单一事实来源第三层用四条 pytest 把“KV cache 不误报、异步不混乱、真泄漏被抓、pending 延迟”钉死在 CI。核心心法泄漏检测器必须区分常驻 buffer 与游离分配、并适配异步释放的延迟语义否则对 WebGPU LLM 必然误报或崩溃。