← 返回案例集合

Case 002 · GPU kernel bugfix

修复 Tesla T4 上的 80 KiB 共享内存溢出

vLLM issue #36802 的首请求崩溃并不是显存不够,而是一个 Triton block 临时申请了超过 T4 硬件上限的共享内存。我把报错还原成可计算的 Q / K / V tile 账单,再用真实 2×T4 验证修复。

上游问题vLLM #36802
硬件2 × Tesla T4 · SM75
失败值81,920 B > 65,536 B
修复提交5f00137664
当前状态Fork 已验证,上游 issue 讨论中

先用大白话说

GPU 里有一块很快但很小的“工作台”,叫 shared memory。这个 attention 内核一次把 Q、K、V 三组数据都摆上工作台,共要 80 KiB;Tesla T4 的工作台只有 64 KiB,所以还没开始算就被硬件拒绝。修复不是换模型,而是让 T4 每轮少搬一半 K / V token:同样的活分更多轮做,工作台占用降到约 48 KiB。

这类错误容易被误解成“GPU 显存不够”。但错误里的 shared memory 指每个 CUDA block 的片上临时空间,不是显卡的 16 GB 全局显存。加 CPU offload 可以帮助模型装入显存,却不能改变单个 block 的 64 KiB 上限。

先认识六个词

Triton

用 Python 风格代码编写 GPU kernel 的语言和编译器;这里负责生成 attention 内核。

Kernel

在 GPU 上并行执行的一段函数。launch 前,运行时会检查它需要的硬件资源。

Shared memory

同一个 thread block 共享的高速片上内存,容量远小于显存。

Tile

一次搬进工作台并处理的数据小块。tile 越大,循环可能越少,但临时空间更大。

SM75

Turing 架构的 compute capability 7.5;Tesla T4 属于这一代。

BF16 → FP32

T4 不原生支持 BF16,vLLM 把模型计算提升为 FP32;每个元素从 2 字节变成 4 字节。

报错是什么,怎样做到同条件复现?

上游 issue 报告在 T4 上运行 nvidia/Nemotron-Content-Safety-Reasoning-4B,第一个请求就失败:

triton.runtime.errors.OutOfResources:
out of resource: shared memory
Required: 81920, Hardware limit: 65536

原报告是 4×T4。Kaggle 只有 2×T4,因此我使用 tensor parallel 2、cpu_offload_gb=4,并关闭 hybrid KV cache manager,让 FP32 模型能够装入。它们只解决容量问题,不会改变 head size、元素宽度或单 block shared-memory 上限。

项目验证环境为什么重要
模型Nemotron Content Safety Reasoning 4B与 issue 相同,包含目标 Gemma 3 attention 路径
vLLM0.17.1 · commit 95c0f928…对齐报告对应的软件版本
PyTorch / Triton2.10.0+cu128 / 3.6.0避免编译器版本差异掩盖问题
GPU2× Tesla T4 · CC 7.5真实 SM75,不用 mock 假装硬件约束
实际 dtypeFP32直接决定每个 Q / K / V 元素占 4 字节

从模型加载到内核崩溃,原本发生了什么?

Gemma 3 的优化策略为 Unified Attention 选择 32-token K / V tile。在 FP16 / BF16 下未必有问题,但 T4 上的实际执行 dtype 是 FP32,每个元素翻倍为 4 字节,于是同一块 tile 的临时空间也翻倍。

81,920 字节是可以一项项算出来的

  • Q tile:16 行 × 256 head size × 4 bytes = 16 KiB
  • K tile:256 head size × 32 tokens × 4 bytes = 32 KiB
  • V tile:32 tokens × 256 head size × 4 bytes = 32 KiB
  • 合计:16 + 32 + 32 = 80 KiB = 81,920 bytes

报错里的两个数字与这份账单完全吻合。这把“可能是 Triton 的问题”收敛成一个具体结论:SM75 + FP32 + head 256 + tile 32 的静态内核形状超过硬件限制。

一个很合理、但实际无效的尝试

Triton 的报错通常会建议降低 num_stages,因为流水线 stage 增多可能复制临时 buffer。于是我先验证了 num_stages=1。

结果:完整模型仍然请求 81,920 字节,数字一字不变。说明这 80 KiB 来自同时存活的 Q / K / V tile,而不是软件流水线的额外复制。

保留这个失败实验很重要:它排除了一个常见解释,也防止后续维护者再次走同一条路。真正有效的控制量必须是 K / V tile size。

修复:只在精确命中问题形状时把 tile 32 改成 16

我没有按模型名硬编码,也没有让所有 T4 都降级。策略只在三个条件同时成立时触发:

  1. 设备是 SM75;
  2. attention head size = 256;
  3. Query 和 KV 的元素宽度都是 4 bytes(FP32)。
if capability == 75 and head_size == 256
   and q_element_size == 4 and kv_element_size == 4:
    return 16  # K/V tile: 32 → 16

return existing_tile_policy(...)

tile 变小以后,同一段序列会通过更多轮循环处理,attention 公式、覆盖范围和输出 dtype 都不变。Q 仍是 16 KiB,K / V 各降到 16 KiB,总计约 48 KiB,低于 T4 的 64 KiB 上限。

验证分三层:策略、内核、完整模型

Layer 1 · Passed

Tile 策略单测

确认只有 SM75 + FP32 + head 256 选择 16;其它 head、dtype 和架构保持原策略。

Layer 2 · Covered

T4 内核回归

为目标形状补充 prefill / decode 测试,输出对照 reference attention。

Layer 3 · Passed

真实完整模型

2×T4 加载同一 Nemotron 模型,修复后成功生成 32 个 token,进程 exit 0。

Control · Reproduced

未修复基线

同模型、prompt、引擎参数和 GPU,首请求稳定复现 81,920 > 65,536。

运行Tile 策略实际结果
基线K / V tile = 32首请求失败:需要 80 KiB,T4 上限 64 KiB
仅 num_stages=1tile 仍为 32仍需要 80 KiB,失败
修复版目标形状 K / V tile = 16成功生成 32 token,未再出现 shared-memory OOR

影响范围、风险和当前状态

设备与输入修复后行为
SM75 · FP32 · head 256使用 tile 16,避免 80 KiB 分配
SM75 · FP32 · head 128不变
SM75 · FP16 / BF16 · head 256不变
Ampere / Ada / Hopper不变
Blackwell head 256 优化原有 8 warps / 2 stages 路径不变

主要风险是 T4 上更小的 tile 可能带来延迟或吞吐回退。当前证据已经证明它把“完全无法运行”变成“可以完成生成”,但合入前仍值得补充稳定的性能对比,看看能否在 16 与其它候选 tile 之间找到更好平衡。

真实 T4 基线已复现

错误数字与上游报告一致:81,920 > 65,536。

根因已用静态内存账单解释

Q / K / V 三块 tile 正好合计 80 KiB。

Fork 修复与完整模型验证已完成

修复分支公开,Kaggle 2×T4 成功生成 32 token。

等待上游维护者反馈

已在 issue 留下方案与证据;目前没有把“可提交候选修复”写成“已被 vLLM 接受”。

我学到的核心:看到 GPU OOR,先分清是哪一种内存;看到编译器建议,也要做对照实验。一个能解释精确报错数字、能在真实硬件上复现、并且只改变目标形状的补丁,才具备进入上游评审的基础。

上一篇 · CASE 001DSH Agent Team 异构模型路由←