先用大白话说
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 上限。
先认识六个词
用 Python 风格代码编写 GPU kernel 的语言和编译器;这里负责生成 attention 内核。
在 GPU 上并行执行的一段函数。launch 前,运行时会检查它需要的硬件资源。
同一个 thread block 共享的高速片上内存,容量远小于显存。
一次搬进工作台并处理的数据小块。tile 越大,循环可能越少,但临时空间更大。
Turing 架构的 compute capability 7.5;Tesla T4 属于这一代。
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 路径 |
| vLLM | 0.17.1 · commit 95c0f928… | 对齐报告对应的软件版本 |
| PyTorch / Triton | 2.10.0+cu128 / 3.6.0 | 避免编译器版本差异掩盖问题 |
| GPU | 2× Tesla T4 · CC 7.5 | 真实 SM75,不用 mock 假装硬件约束 |
| 实际 dtype | FP32 | 直接决定每个 Q / K / V 元素占 4 字节 |
从模型加载到内核崩溃,原本发生了什么?
Before · the first request never reaches generation
4 bytes / element
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 都降级。策略只在三个条件同时成立时触发:
- 设备是 SM75;
- attention head size = 256;
- 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 上限。
Narrow guard · no model-name special case
验证分三层:策略、内核、完整模型
Tile 策略单测
确认只有 SM75 + FP32 + head 256 选择 16;其它 head、dtype 和架构保持原策略。
T4 内核回归
为目标形状补充 prefill / decode 测试,输出对照 reference attention。
真实完整模型
2×T4 加载同一 Nemotron 模型,修复后成功生成 32 个 token,进程 exit 0。
未修复基线
同模型、prompt、引擎参数和 GPU,首请求稳定复现 81,920 > 65,536。
| 运行 | Tile 策略 | 实际结果 |
|---|---|---|
| 基线 | K / V tile = 32 | 首请求失败:需要 80 KiB,T4 上限 64 KiB |
| 仅 num_stages=1 | tile 仍为 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 之间找到更好平衡。
错误数字与上游报告一致:81,920 > 65,536。
Q / K / V 三块 tile 正好合计 80 KiB。
修复分支公开,Kaggle 2×T4 成功生成 32 token。
已在 issue 留下方案与证据;目前没有把“可提交候选修复”写成“已被 vLLM 接受”。
我学到的核心:看到 GPU OOR,先分清是哪一种内存;看到编译器建议,也要做对照实验。一个能解释精确报错数字、能在真实硬件上复现、并且只改变目标形状的补丁,才具备进入上游评审的基础。