【Bug已解决】[Bug]: MTP speculative decoding crash with illegal memory access on long sequences (Qwen3.6-

📅 2026/7/30 21:23:34
【Bug已解决】[Bug]: MTP speculative decoding crash with illegal memory access on long sequences (Qwen3.6-
【Bug已解决】[Bug] MTP speculative decoding crash with illegal memory access on long sequences (Qwen3.6-27B-FP8, v0.19.1) 解决方案一、现象长什么样在Qwen3.6-27B-FP8 vLLM 0.19.1 上开MTPmulti-token prediction多 token 预测投机解码短序列一切正常但一旦序列接近 KV 缓存上限长文档、长对话forward 中途进程崩溃报非法内存访问torch.OutOfMemoryError: CUDA out of memory (有时) RuntimeError: CUDA error: an illegal memory access was encountered或者更明确指向投机解码层illegal memory access at kernel mtp_draft_forward: slot index 16384 num_kv_slots 16384几个特征短序列比如 2k token 以内完全正常长序列接近--max-model-len必崩。只在开 MTP 时崩关掉 MTPspeculative_configNone长序列也能跑。崩的位置是 MTP 的 draft 前向不是主模型前向。报错有时是illegal memory access有时是out of memory本质是同一件事MTP 往「超出已分配 KV 槽位」的地方写了数据。本质MTP 一次草拟k个 token会把当前位置往后推k格去写 KV 缓存当序列已经很长、剩余 KV 槽位不足k个时草拟位置越界kernel 写到了未分配显存 → 非法内存访问。二、背景MTP 的做法是主模型算出下一个 token 后MTP 头基于「已生成的序列 刚算出的 token」再一次性草拟出接下来 k 个 token比如 k3然后主模型并行验证这 k 个。为了草拟第k个 tokenMTP 前向需要把位置pos, pos1, ..., posk-1的 KV 都写进缓存、并读出来算下一层。问题在于「写 KV 缓存」这一步和「KV 缓存容量」的耦合KV 缓存是按max_model_len预分配好固定槽位的比如 16384 个 slot。普通解码每次只推进 1 个位置永远不会越界因为调度器保证序列长度 ≤ max_model_len。但 MTP 一次要推进k个位置。调度器在计算「这个序列还能不能接着生成」时往往只按「主模型 1」来算剩余槽位没把 MTP 要额外占的k-1个槽位算进去。于是当序列长度 max_model_len - 2时调度器认为「还能生成」MTP 却要写pos到pos2共 3 个槽位最后一个槽位pos2 max_model_len已经越界 → kernel 写未分配显存 → 非法内存访问。这和普通「序列超长」不同普通情况调度器会拦下但 MTP 把「一次占用的槽位数」从 1 变成了 k调度器的边界判断没跟着改漏洞就出现了。三、根因根因是MTP 草拟长度k没有被纳入 KV 槽位的边界核算导致长序列末尾草拟位置越界三层第一层主因调度器的「剩余槽位」判断没加 MTP 的k余量。调度器决定「这个序列还能不能生成下一个 token」时检查的是seq_len 1 max_model_len。但 MTP 实际上需要seq_len k max_model_len。差了k-1个槽位序列在max_model_len - k seq_len max_model_len - 1这段区间里调度器放行、MTP 越界。第二层MTP draft 前向没有对 slot 做边界断言。draft kernel 拿到pos和k后直接kv_cache[pos i] ...没有任何if pos i num_slots: 截断/报错的防护。它假设「调用方保证槽位够」但调用方调度器的保证是错的于是越界写直接发生。第三层错误表现不稳定IMA vs OOM。越界写的后果取决于「越界到哪」若越界到同一块已分配显存的邻近区域可能只是静默污染偶尔还能跑完但结果错若越界到未映射显存就是illegal memory access若越界触发了一次额外的显存分配就是out of memory。同一个根因三种表象增加排查难度。一句话MTP 的草拟长度没被调度器算进 KV 边界长序列末尾草拟越界写未分配显存表现为非法内存访问或偶发 OOM。四、最小可运行复现下面用纯 Python 模拟「MTP 草拟 k 个 token但调度器只按 1 判断边界长序列末尾越界」的控制流不需要 GPUclass KVCache: def __init__(self, num_slots): self.slots [None] * num_slots self.num num_slots def write(self, pos, k, value): # MTP draft 前向写 pos .. posk-1 for i in range(k): idx pos i if idx self.num: # 原版没有这个检查直接越界 raise IndexError(fslot {idx} num_slots {self.num}) self.slots[idx] value def can_generate(seq_len, max_len, k, speculative): # 调度器的边界判断 needed seq_len (k if speculative else 1) return needed max_len def main(): max_len, k 16, 3 cache KVCache(max_len) # 序列化到 seq_len 14max_len - 2 seq_len max_len - 2 speculative True # 调度器认为14 1 15 16放行 print(调度器放行:, can_generate(seq_len, max_len, k, speculative)) # 但 MTP 要写 14,15,16 - 16 越界 try: cache.write(seq_len, k, draft_token) print(写成功实际会越界) except IndexError as e: print(复现成功:, e) if __name__ __main__: main()跑出来会打印调度器放行: True然后复现成功: slot 16 num_slots 16——调度器以为能生成、MTP 却越界和线上「长序列末尾崩溃」完全一致。五、解决方案第一层最小直接修复最省事的救火关掉 MTP退回普通解码长序列立刻能跑。代价是吞吐下降失去投机加速llm LLM( modelQwen3.6-27B-FP8, # speculative_configNone # 不启用 MTP )或者把--max-model-len调大一点给 MTP 的k余量留出空间代价是 KV 缓存显存变大vllm serve Qwen3.6-27B-FP8 \ --speculative-config {method:mtp,num_speculative_tokens:3} \ --max-model-len 16384 \ --gpu-memory-utilization 0.8 # 留出 KV 余量更精准的临时规避限制 MTP 只在「剩余槽位充足」时启用剩余不足k就退回单 token 解码。这是第一层的「带保护」版本def safe_num_draft(seq_len, max_len, k): # 剩余槽位不足以支撑 k 个草拟时自动缩减到 1普通解码 remaining max_len - seq_len return min(k, max(1, remaining))六、解决方案第二层结构性改进第一层是「避开/手动留余量」第二层是「让调度器和 MTP 用同一套边界规则」——核心是把 MTP 的k纳入「可生成判定」和「KV 槽位核算」的单一事实来源from dataclasses import dataclass dataclass class SeqBounds: max_len: int num_speculative: int 1 def can_generate(self, seq_len: int) - bool: # 单一边界规则主模型 全部草拟 token 都必须落在 max_len 内 needed seq_len self.num_speculative return needed self.max_len def draft_slots_ok(self, pos: int, k: int) - bool: # MTP draft 前向的边界检查pos .. posk-1 必须全部合法 return pos k self.max_len def check_position(self, pos: int, k: int) - None: assert self.draft_slots_ok(pos, k), ( fMTP draft 越界: pos{pos} k{k} 需要槽位 {posk} f但 max_len{self.max_len} ) def check_expert(self, pos: int, k: int) - None: # 专家路由侧的同样检查MoE 下 token 也要落到合法 slot self.check_position(pos, k)调度器在决定是否继续生成时统一调用can_generate把num_speculative算进去MTP draft kernel 入口先check_position(pos, k)再写 KVdef mtp_draft_forward(kv_cache, pos, k, bounds: SeqBounds): bounds.check_position(pos, k) # 越界立刻报错绝不写未分配显存 for i in range(k): kv_cache.write(pos i, compute_token(pos i))这样「边界规则」只有一份调度器和 MTP 不可能再各算各的。七、解决方案第三层断言 / CI 守护把「MTP 不越界」「调度器按 k 判断」「长序列末尾安全降级」固化成测试import pytest def test_draft_within_bounds_ok(): b SeqBounds(max_len16, num_speculative3) b.check_position(10, 3) # 10..12 16应通过 def test_draft_at_boundary_raises(): b SeqBounds(max_len16, num_speculative3) with pytest.raises(AssertionError): b.check_position(14, 3) # 14..16 越界 def test_scheduler_accounts_for_k(): b SeqBounds(max_len16, num_speculative3) # seq_len14 时1431716调度器应拒绝继续生成 assert b.can_generate(14) is False assert b.can_generate(13) is True # 13316 16 def test_long_seq_tail_safe_degrade(): # 长序列末尾MTP 自动退化成单 token不越界 b SeqBounds(max_len16, num_speculative3) seq_len 15 k min(b.num_speculative, 16 - seq_len) # k1 b.check_position(seq_len, k) # 15..15 合法 def test_no_ima_on_max_len(): # 端到端在 max_len 处停止草拟不应触发越界 b SeqBounds(max_len16, num_speculative3) for seq_len in range(0, 16): k 3 if b.can_generate(seq_len) else 0 if k: b.check_position(seq_len, k) assert True再加一个端到端回归长序列 MTP 跑到max_model_len不崩def test_mtp_long_sequence_no_ima(): engine make_engine(modelQwen3.6-27B-FP8, speculative{method: mtp, num_speculative_tokens: 3}, max_model_len16384) out engine.generate(超长文档... * 500, max_tokens16384) assert out is not None # 不应 illegal memory access八、排查清单看报错是否illegal memory access/out of memory且栈指向 MTP draft 前向 → 坐实本问题。短序列能跑、长序列崩且只在开 MTP 时崩 → 基本是 MTP 越界。临时救火关 MTP或调大--max-model-len或在长序列末尾手动降num_speculative。检查调度器「剩余槽位」判断是否包含 MTP 的k余量最常见疏漏。长期修复边界规则单一化调度器与 MTP 共用SeqBoundsdraft 前向前做check_position。升级 vLLM 到合了 MTP 边界修复的版本并跑上面的长序列回归。若用CUDA_LAUNCH_BLOCKING1TORCH_USE_CUDA_DSA1复现能让越界错误定位到精确 kernel 行。九、小结MTP 长序列非法内存访问不是 FP8 或 Qwen 的锅而是MTP 一次草拟 k 个 token但调度器的 KV 边界判断只按 1 算长序列末尾草拟位置越界写未分配显存。最小修复是关 MTP / 调大 max_model_len / 末尾降 k结构性修复是把边界规则收敛成单一SeqBounds、draft 前向前做check_position最后用 pytest 把「不越界」「调度器按 k 判断」「长序列安全降级」锁死。配合CUDA_LAUNCH_BLOCKING1能快速定位越界 kernel。抓住「投机解码一次性占用的槽位数 ≠ 1」这条所有 spec decode 的边界坑都能照此排查。