【Bug已解决】[Bug]： Combination of sleep mode and speculative decoding cause crash (ROCm ／ MI250) 解决方案

【Bug已解决】[Bug] Combination of sleep mode and speculative decoding cause crash (ROCm MI250) 解决方案一、现象长什么样在 AMD MI250ROCm上同时开启sleep 模式空闲时释放显存和speculative decoding投机解码 / MTP 多 token 预测后一旦经历「sleep → wake_up → 第一次带投机解码的推理」进程崩溃。栈通常指向投机解码层RuntimeError: draft model KV cache is None after wake_up; cannot run speculative step或者更靠内核层ROCm Habana/CK error: hipMalloc returned nullptr for draft_workspace (size...)几个特征单独开 sleep 模式不开 spec decode一切正常单独开 spec decode不 sleep也正常。两者叠加才炸。只在「唤醒后第一次」投机解码推理时炸唤醒后如果先跑几次普通推理再跑投机解码有时能过——因为普通推理顺手把某些 buffer 重建了。在 MI250 这类 ROCm 卡上必现NVIDIA 卡上不一定ROCm 的显存分配失败返回 nullptr 而非抛异常更隐蔽。本质sleep 把投机解码相关的显存draft 模型、draft workspace、bonus token 缓冲一起释放了但 wake_up 只重建了主模型没重建 spec decode 那一摊第一次投机前向访问到 None/未分配的 buffer 就崩。二、背景先理清两个机制各自干什么以及它们叠加时的隐患。speculative decoding以 MTP 为例需要比普通解码多一套资源draft 模型或 MTP 头一个比主模型小、用来「猜」接下来几个 token 的模型常驻显存。draft workspace投机每一步临时存放候选 token、draft 概率、bonus token 的缓冲区。独立的 draft KV cachedraft 模型自己也要 KV 缓存。sleep 模式的做法是把「引擎持有的显存」整体 offload/释放。理想情况下它应该记录「我释放了哪些」wake_up 时按原样全部恢复。问题就出在「全部恢复」这件事上主模型的恢复逻辑写得很完整但 spec decode 那套资源尤其是 draft workspace 和 draft KV cache在 sleep 的「释放清单」和 wake_up 的「重建清单」里对不齐——要么释放了没登记要么登记了没重建。MI250 上因为显存分配失败是静默返回 nullptr而不是像 CUDA 那样抛cudaErrorMemoryAllocation这个不一致更难被早期发现一直藏到第一次投机前向才暴露。三、根因根因是wake_up 的「重建清单」遗漏了投机解码专属资源导致 draft 模型/draft workspace/draft KV 在唤醒后处于未初始化状态三层第一层主因spec decode 资源不在 sleep/wake 的对称清单里。sleep 时释放了 draft workspace因为它挂在引擎的某个子模块上被一并free了但 wake_up 的恢复函数只遍历了「主模型权重 主 KV cache」两个清单没遍历「spec decode 子模块」。于是 draft workspace 永久变成 None/未分配。第二层ROCm 的 nullptr 让问题延迟暴露。在 NVIDIA 上cudaMalloc失败会立刻抛异常wake_up 阶段就能被发现而 ROCm 上hipMalloc返回 nullptr代码如果没做if ptr is None检查就会带着 nullptr 继续跑直到真正往这个 buffer 写数据投机解码第一步才段错误/抛 RuntimeError。这解释了「为什么只在 MI250 上必现、且只在第一次投机前向炸」。第三层投机解码层假设 draft 资源一定就绪。spec decode 的前向代码里直接draft_kv_cache.advance()、draft_workspace.zero_()没有「若未初始化则先重建」的兜底。它假设「只要引擎 RUNNINGspec decode 资源就必然就绪」——这个假设在 sleep/wake 场景下被打破。一句话sleep 释放了 spec decode 的显存但 wake_up 没重建ROCm 的静默 nullptr 把崩溃推迟到第一次投机前向而 spec 层本身又没做就绪性兜底。四、最小可运行复现下面用纯 Python 模拟「sleep 释放了主draft 两类资源但 wake_up 只恢复了主、漏了 draft随后访问 draft 时崩溃」的控制流不需要 GPU用普通对象模拟显存句柄class Buffer: def __init__(self, name): self.name name self.ptr object() # 模拟已分配的显存句柄 def free(self): self.ptr None def ready(self): return self.ptr is not None def sleep(engine): for buf in engine.buffers.values(): buf.free() # 全部释放 engine.state SLEEPING def wake_up_buggy(engine): # 只重建主模型漏了 draft engine.buffers[main].ptr object() engine.buffers[main_kv].ptr object() # draft_workspace / draft_kv 没重建 - 仍是 None engine.state RUNNING def speculative_step(engine): # spec 层假设 draft 资源就绪 if not engine.buffers[draft_workspace].ready(): raise RuntimeError(draft workspace is None after wake_up) return speculated def main(): engine type(E, (), {})() engine.buffers { main: Buffer(main), main_kv: Buffer(main_kv), draft_workspace: Buffer(draft_workspace), draft_kv: Buffer(draft_kv), } engine.state RUNNING sleep(engine) wake_up_buggy(engine) print(main ready:, engine.buffers[main].ready()) print(draft ready:, engine.buffers[draft_workspace].ready()) try: speculative_step(engine) except RuntimeError as e: print(复现成功:, e) if __name__ __main__: main()跑出来会打印draft ready: False然后复现成功: draft workspace is None after wake_up与线上「唤醒后 draft 资源缺失」完全一致。五、解决方案第一层最小直接修复最省事的救火要么别同时开 sleep 和 spec decode要么 wake_up 后强制重建 spec decode 资源。临时做法是在 wake_up 完成后显式调一次 spec decode 的初始化def wake_up_safe(engine): # 先按原逻辑恢复主模型 restore_main_model(engine) # 关键补一步——确保 spec decode 资源也重建 if engine.spec_decode_enabled: engine.draft_model rebuild_draft_model(engine) engine.draft_workspace allocate_workspace(engine) engine.draft_kv_cache rebuild_draft_kv(engine) engine.state RUNNING如果你只是想先让服务不崩、不在乎 sleep 省的那点显存最简单是关掉 sleep 模式llm LLM( model..., speculative_config{method: mtp, ...}, # 不启用 sleep不要调用 /sleep规避叠加 )六、解决方案第二层结构性改进第一层是「补一步」第二层是「让 sleep/wake 把 spec decode 资源纳入统一的对称清单」从设计上消灭遗漏。核心把「所有需要随 sleep 释放、随 wake 重建的资源」集中注册到一个ResourceRegistrysleep 遍历释放、wake 遍历重建spec decode 子模块在初始化时主动注册自己的资源。from dataclasses import dataclass, field from typing import Dict, Callable dataclass class ManagedResource: name: str allocate: Callable[[], object] release: Callable[[object], None] handle: object None class ResourceRegistry: def __init__(self): self._res: Dict[str, ManagedResource] {} def register(self, res: ManagedResource): res.handle res.allocate() self._res[res.name] res def sleep_all(self): for res in self._res.values(): if res.handle is not None: res.release(res.handle) res.handle None def wake_all(self): for res in self._res.values(): if res.handle is None: res.handle res.allocate() # 统一重建不漏项 # spec decode 子模块在初始化时注册自己的资源 def init_spec_decode(registry: ResourceRegistry, engine): registry.register(ManagedResource( namedraft_workspace, allocatelambda: allocate_workspace(engine), releaselambda h: h.free(), )) registry.register(ManagedResource( namedraft_kv_cache, allocatelambda: rebuild_draft_kv(engine), releaselambda h: h.free(), )) registry.register(ManagedResource( namedraft_model, allocatelambda: rebuild_draft_model(engine), releaselambda h: h.unload(), ))这样无论以后加多少类资源只要「注册进 registry」sleep/wake 永远对称不会再出现「释放了没重建」。七、解决方案第三层断言 / CI 守护把「wake_up 后所有注册资源必须就绪」和「spec decode 前向前做就绪检查」固化成测试import pytest def test_sleep_wake_restores_all_registered(): reg ResourceRegistry() reg.register(ManagedResource(main, lambda: object(), lambda h: None)) reg.register(ManagedResource(draft_ws, lambda: object(), lambda h: None)) reg.sleep_all() # sleep 后都应释放 assert reg._res[main].handle is None assert reg._res[draft_ws].handle is None reg.wake_all() # wake 后都必须重建 assert reg._res[main].handle is not None assert reg._res[draft_ws].handle is not None def test_spec_decode_ready_before_forward(): engine make_engine(spec_decodeTrue) engine.sleep() engine.wake_up() # 投机前向前断言 draft 资源就绪 assert engine.draft_workspace.ready() assert engine.draft_kv_cache.ready() def test_speculative_step_raises_if_not_ready(): engine make_engine(spec_decodeTrue) engine.sleep() engine.wake_up_buggy() # 故意用漏重建的版本 with pytest.raises(RuntimeError): engine.speculative_step() def test_rocm_nullptr_detected(): # 模拟 ROCm 静默返回 nullptr必须被显式检查捕获 ptr None # hipMalloc 返回 nullptr assert ptr is None, hipMalloc 返回 nullptr 必须被检查不能带病前进再加一个端到端回归MI250 模拟下sleepwakespec decode 跑 10 轮不崩def test_sleep_wake_spec_decode_loop_stable(): engine make_engine(spec_decodeTrue, backendrocm) for _ in range(10): engine.sleep() engine.wake_up() out engine.generate(hello, use_speculativeTrue) assert out is not None八、排查清单看栈是否指向draft_*/speculative相关符号且发生在 wake_up 后的第一次推理 → 坐实本问题。单独关 spec decode 试一次单独关 sleep 试一次定位是不是「叠加」引发。ROCm 上检查是否有hipMalloc returned nullptr的日志静默失败需主动 grep。临时救火wake_up 后显式重建 draft 资源或直接不同时开 sleep 与 spec decode。长期修复把 spec decode 资源注册进统一的 ResourceRegistrysleep/wake 对称管理。投机解码前向前加「资源就绪」断言避免带 None 前进。升级 vLLM 到合了 sleepspec decode 兼容修复的版本并跑上面的 sleep/wake 循环回归。九、小结sleep spec decode 在 MI250 上崩溃不是 ROCm 硬件的锅而是sleep 释放了投机解码的显存draft 模型/workspace/KVwake_up 却只重建了主模型导致 draft 资源在唤醒后处于未初始化状态ROCm 的静默 nullptr 又把崩溃推迟到第一次投机前向才暴露。最小修复是 wake_up 后补重建 draft 资源或干脆不同时开两者结构性修复是用统一 ResourceRegistry 让 sleep/wake 对所有注册资源对称最后用 pytest 把「唤醒后资源全就绪」和「spec 前向前就绪检查」锁死。抓住「sleep/wake 必须对称覆盖每一个子系统的资源」这条叠加功能的稳定性坑都能照此化解。