CUDA out of memory

快速结论：该报错发生在 vLLM 使用 FlashMLA 稀疏 FP8 混合批次解码路径时，稀疏解码内核会分配一个随 --max-num-batched-tokens 线性增长的工作区，且该工作区不受启动时内存预检约束，批次增大后触发运行时 CUDA OOM。优先排查方向是调低 --max-num-batched-tokens，而不是调整 --gpu-memory-utilization。

适用环境：vLLM + GLM-5.2 FP8 模型，8×H200（GPU 单卡 139.80 GiB 显存），TP=8 + EP 启用，--kv-cache-dtype fp8，--max-num-batched-tokens 32768，FlashMLA 稀疏解码路径，混合 prefill+decode 批次场景。

最快修复方案：将 --max-num-batched-tokens 从 32768 调低到 8192，并将 --gpu-memory-utilization 提高到 0.945；该组合在 4×H200 生产环境连续运行约两个月未再触发 OOM。Issue 中提到的 PR #49357 与 #53755 属于上游修复补丁，尚未合并到当前版本。

注意事项：调低 token 预算会限制单步最大吞吐，需要根据实际并发请求数和 prefill 长度重新评估性能。Issue 作者同时维护了一个本地补丁，用于限制 prefill 工作区分配并修改混合批次路径的 tiling，但该补丁未合入上游，使用时需注意随版本升级维护成本。当前修复方案是在 flashmla_sparse.py 的混合批次路径上验证的，不适用于标准 BF16 prefill 路径。

问题场景

在 vLLM 服务 GLM-5.2 FP8 模型时，使用 8×H200 配置，开启 --tensor-parallel-size 8、--enable-expert-parallel、FP8 KV cache 以及 MTP 投机解码，服务启动正常，但在持续运行约 61 小时后，当某一步的 prefill+decode 混合批次填满 token 预算时，所有 8 个 worker 在同一个 step 内触发 torch.OutOfMemoryError，随后 EngineCore 进程因致命错误退出，所有在途请求以 EngineDeadError 失败。

报错原文

torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 8.00 GiB.
GPU 0 has a total capacity of 139.80 GiB of which 6.40 GiB is free.
Including non-PyTorch memory, this process has 133.39 GiB memory in use.
Of the allocated memory 124.39 GiB is allocated by PyTorch, with 3.23 GiB allocated in
private pools (e.g., CUDA Graphs), and 5.01 GiB is reserved by PyTorch but unallocated.

(EngineCore pid=614) ERROR [core.py:1233] EngineCore encountered a fatal error.
(APIServer pid=1) ERROR [serving.py:812] vllm.v1.engine.exceptions.EngineDeadError:
    EngineCore encountered an issue. See stack trace (above) for the root cause.

原因分析

FlashMLA 稀疏 FP8 混合批次路径（当每 rank 头数 < 32 时自动选择）中的 sparse_decode_fwd 内核会分配一个工作区，其大小与 --max-num-batched-tokens 线性相关，计算公式为 max_num_batched_tokens × index_topk × index_head_dim。在默认配置下（32768 × 2048 × 128 × 1 byte = 8 GiB），这个工作区不受启动时内存 profiling 约束，也不会被现有 workspace limit 限制。

关键点是：这个 OOM 不是 KV cache 压力导致的——崩溃前 1 分钟 GPU KV cache 使用率仅 2.0–5.4%，并发请求最多 8 个。vLLM 在启动时通过 memory profiling 预留了 KV cache 和模型权重空间，但稀疏解码内核的工作区是在运行时按需分配的，当混合批次中首次出现足够大的 batch 时，这个一次性分配直接击穿剩余显存。

可能原因还包括：混合批次路径绕过了独立路径（BF16 prefill 路径）使用的 workspace 边界检查，导致 get_prefill_workspace_size 中“5 × max_model_len × 576 × 2 bytes”的内存评估没有覆盖到稀疏解码的 8 GiB 工作区。

环境排查

确认 vLLM 版本是否包含 flashmla_sparse.py 中的 MIN_HEADS_FOR_BF16_PREFILL = 32 逻辑和 fp8_use_mixed_batch 分支
确认启动参数中 --max-num-batched-tokens 的值（32768 高风险，8192 验证安全）
确认 GPU 显存总量与 --gpu-memory-utilization 的乘积是否留有足够余量给稀疏解码工作区
检查单卡显存占用分布：PyTorch 分配、CUDA Graphs 私有池、预留未分配空间
确认每 rank 头数是否小于 32（高 TP 配置下常见），据此判断是否必然走混合批次路径
记录崩溃前 KV cache 使用率，排除 KV 压力误判

解决步骤

第一步：将 --max-num-batched-tokens 从 32768 调低到 8192。这是验证过的首选改动，4×H200 生产环境约两个月未触发 OOM。调低后单步工作区从 8 GiB 降至 2 GiB。
第二步：将 --gpu-memory-utilization 从默认 0.90 提高到 0.945。注意：仅提高利用率收益有限，0.945 → 0.942 仅多出约 0.43 GiB，不能单独解决问题，必须与第一步配合。
第三步：如果上游已合入修复 PR #49357（workspace 限制）或 #53755（FlashMLA 修复），优先升级到包含这些修复的版本。当前 Issue 关闭时这些 PR 尚未合并。
第四步（可选）：如果 GPU 预算允许，考虑 prefill-decode disaggregation（PD 分离）作为规避方案，将长 prefill 路由到独立实例。
第五步（可选，需自行维护）：参考 Issue 作者的本地补丁思路，修改 vllm/v1/attention/backends/mla/flashmla_sparse.py，对混合批次路径的 prefill 工作区分配增加上限检查，并调整 tiling 策略。此补丁未合入上游，升级 vLLM 后需重新应用。

验证方法

用调整后的参数重启服务，持续驱动混合 prefill+decode 流量直到填满 token 预算，观察是否还出现 torch.OutOfMemoryError。同时监控两点：单卡显存峰值是否始终低于 gpu-memory-utilization × 总显存，以及日志中 sparse_decode_fwd 的分配是否被 workspace 限制拦截。生产环境建议至少持续运行一周以上，覆盖高峰期并发样本。如果使用本地补丁，需在 patch 后重新运行显存 profiling 确认工作区不再超限。