如何用8台桌面机跑744B大模型百万token上下文?一个CUDA寄存器大小引发的血案

8台桌面级DGX Spark跑出744B模型的百万token上下文。卡死它的不是算力不是显存,而是一个99KB的CUDA共享内存限制。

如何用8台桌面机跑744B大模型百万token上下文?一个CUDA寄存器大小引发的血案 封面图

结论先行:744B参数、百万token上下文、8台桌面级DGX Spark。你以为瓶颈是算力?是显存?都不是。真正卡死这个系统的,是GPU调度算法中99KB共享内存导致worker group超过48个上限——加一行判断分支就解决了。

大模型领域有个潜规则:跑大模型=烧H100集群。百万token的上下文窗口?那得几十台数据中心级GPU才能撑起KV cache。

但有人不信邪。他用8台NVIDIA最小的桌面AI工作站——DGX Spark(每台128GB统一内存,GB10芯片),硬是把744B参数的GLM 5.2跑到了1,048,576 token的完整上下文窗口,还能生成连贯文本。

一路踩了两个"墙"。

第一堵墙:Tensor Parallelism把KV Cache复制了8份

大模型并行最常见的方案是Tensor Parallelism(张量并行)——每台机器算每一层的部分矩阵运算,互相实时通信。

但对GLM 5.2带MLA(多头潜在注意力)的模型,有个坑爹特性:压缩后的KV cache无法分片,必须完整拷贝到每台节点上。

实测数据:百万token时,每台DGX Spark要装 93GB的KV cache。

但每台只有128GB总内存,加载模型权重后,剩下的空间根本塞不下。加上Tensor Parallelism下,不管加多少台机器,KV cache都会被完整复制——上下文天花板固定在约20万token,加机器也没用。

解法:Pipeline Parallelism

不搞每台算全部层,改成每台负责不同层:节点1算前10层,节点2算接下来10层,以此类推。

这个切换有个意外的副作用:KV cache按层自动拆分。每个节点只存自己那几层的KV,93GB瞬间变成~10GB。

上下文墙消失——理论上可以直奔百万token了。

Tensor vs Pipeline并行策略对比:左侧是所有节点拷贝93GB KV cache,右侧是KV按层拆分

第二堵墙:99KB共享内存引发的血案

切换到Pipeline Parallelism后,作者信心满满设置百万token上下文——启动就崩。

报错信息看着像天书:

Persistent top K would oversubscribe, and the fallback requires at least 128 KB of shared memory per block, but you only have 99 KB.

翻译成人话:“GPU上有47个工位,但你叫来了100个人——多出来的人站着挡路,备用方案又需要更大的桌子才能坐下,但桌子只有这么大。”

来拆一下到底发生了什么。

罪魁祸首:Sparse Attention的Indexer

GLM 5.2的关键技术叫深度稀疏注意力(DSA)——不是每个词都看其他所有词,而是由一个叫Indexer的组件,为每个新token挑选最相关的2048个历史token。

这个Indexer用的是一个专门CUDA Kernel叫 persistent_top_k。

在长序列下,这个kernel把任务切分为小块,在GPU上并行运行大量worker group。问题来了:这种协同启动有硬性约束——所有worker group必须同时驻留在GPU上。GB10芯片的限制是 48个。

worker group的数量 = 序列长度 ÷ 每个组能处理的元素数。

共享内存是真正的瓶颈

每个组能处理多少元素,受限于它分配到多少共享内存(Shared Memory)——GPU上的超快速片上缓存。

代码里有一个对短序列友好的优化:当batch较小时,主动减少申请的共享内存。这个优化对短请求合情合理。

但到了百万token级别,每个组只能啃一丁点序列,所以需要的组数远超48——oversubscribe了。

更绝的是,当触发oversubscribe时,kernel试图回退到另一个更安全的路径。但那条路需要每block至少128KB共享内存。

GB10只允许你申请约 99KB(101,376字节)。数据中心的Blackwell芯片给227KB,但GB10没有。

128 > 99,所以回退路径在此硬件上直接不可能执行——必崩。

根因分析:跟模型大小无关,跟序列长度有关

作者花了很久才意识到一个关键洞察:这跟744B的大模型无关,跟并行策略无关,纯粹是序列长度太长,导致top K切分过细,worker group数超了48上限。

之前没人遇到这个Bug,是因为Tensor Parallelism在20万token时就因KV内存墙卡住了——20万token时top K还在48以内。

当Pipeline Parallelism解除了KV内存墙、得以推高上下文时,第二堵墙才暴露出来。

修复:一行branch的精准手术

一旦理清问题,修复极其简单:

对于长序列,不让kernel缩减共享内存,直接让它申请GB10允许的全部99KB。

一个worker group就能处理约24,000个元素(之前只处理几千个)。算算百万token的账:

1,000,000 ÷ 24,000 ≈ 43组

43 < 48。完美容纳,不需要128KB的回退路径。

在代码层面,就是在kernel启动逻辑里加一个分支:

  • 短序列:保持原有的batch-size优化逻辑
  • 长序列:不论batch大小,直接申请最大共享内存

修改完,重编译CUDA扩展,推送到所有节点——启动成功。KV cache容量报告显示超过200万token容量,启动自检通过,文本生成连贯。

CUDA共享内存修复示意图:99KB限制下43个worker组刚好塞进48上限

代价与取舍:Pipeline Bubble

Pipeline Parallelism有个固有缺陷——pipeline bubble(流水线气泡)。

单次请求时,只有当前节点在计算,其他节点干等着。单流生成约 3-4 token/s。

方案最大上下文推理速度适合场景
Tensor并行~200K tokens快实时对话、低延迟场景
Pipeline并行(修复版)1M tokens~3-4 tok/s大批量文档处理、长上下文推理

如果追求快速交互聊天,20万token的Tensor并行更合适。如果需要分析百万token级别文档——Pipeline并行才是正确选择。

你是选速度,还是选长度?按工作负载决定。

复现的三个坑

如果你有DGX Spark想复现:

  1. Python 2 vs 3:构建脚本调Python,容器里只有Python 3
  2. Git repo检查:源码目录不是git checkout的话,版本工具会失败,设一个假版本环境变量
  3. 依赖重拉取:构建会通过git重新拉取依赖然后失败,告诉它保持离线模式

作者已将修复提交到 Jazelle vLLM fork 上游(视频描述中有链接)。

总结

744B参数模型 → 百万token上下文 → 8台桌面级机器。

卡住它的不是GPU算力不足、不是显存不够——是构建kernel启动逻辑时,为了短查询优化而申请的共享内存太小,导致长序列下worker group超限。

一句话:一个CUDA kernel里99KB的共享内存限制,一条额外的判断分支,解锁了一个数量级。

作者最后说:“我在上游vLLM、公共GB10文档和GLM 5.2移植中都没找到这个修复或相关文档。所以我在上游提了bug。百万token是我的个人测试结果,尚未被独立复现——把它视为一个可信的数据点,而非教科书。”

如果你对在Blackwell上跑稀疏注意力模型感兴趣,希望这个实战故事能帮你省下一个抓狂的下午。

← 返回文章列表