突破内存墙：在标准 GPU 上实现 3,000 tokens/s 的推理速度

当前 LLM 推理的状态通常通过聚合吞吐量来衡量——即一个服务器在数百个用户中可以生成的总 token 数量。虽然这对云提供商来说是一个关键指标，但对于下一代 AI：自主智能体（autonomous agents）来说，这是一个错误的指标。智能体工作流本质上是顺序执行的；智能体必须在一个循环中进行检查、规划、编辑和测试，其中每一步都依赖于前一步。在这种情况下，唯一重要的指标是单请求解码速度。

Kog AI 发布了 Kog Inference Engine (KIE) 的技术预览版，证明了标准数据中心 GPU 在处理 2B 模型时，每个请求可以实现 3,000 tokens/s 的速度（在 8× AMD MI300X 上）和 2,100 tokens/s 的速度（在 8× NVIDIA H200 上）。通过将模型架构、运行时（runtime）和底层 GPU 内核（kernels）视为一个协同设计的单一系统，Kog 正试图缩小通用 GPU 与专用推理芯片之间的差距。

瓶颈：是内存带宽，而非计算

对于单请求解码（batch size 1），主要的约束不是每秒浮点运算次数 (FLOPS)，而是内存带宽。在自回归解码中，GPU 必须为生成的每一个 token 移动每一个活跃的模型权重，从高带宽内存 (HBM) 移动到计算处理器。

这创造了一个低算术强度（arithmetic intensity）的场景，其中系统几乎完全受限于权重流式传输的速度。虽然像 H200 这样的现代 GPU 拥有巨大的理论带宽，但大多数推理栈（inference stacks）都无法有效地利用它。对于一个 2B 模型，在 FP16 精度下，8× H200 节点的理论“光速”上限大约是 7,700 tokens/s。大多数框架运行速度远低于此，这一事实表明瓶颈不在于硬件，而在于软件栈。

微秒级的损失在哪里

在 3,000 tokens/s 的目标速度下，每个 token 的总时间预算仅为 333 微秒。在标准的推理栈（例如 vLLM 或 TensorRT-LLM）中，这一预算会被开销（overheads）消耗掉，这些开销在较高的 batch size 时通常可以忽略不计，但对于单请求延迟来说却是灾难性的：

  • Kernel Launch Overhead（内核启动开销）： 启动和清理一个 kernel 可以花费约 4.5µs。如果一个 Transformer 层需要 10 个 kernel，那么 25 层就会在进行任何实际计算之前产生超过 1,100µs 的开销。
  • CPU Scheduling（CPU 调度）： 主机端逻辑和 GPU-CPU 通信会在采样和执行中引入延迟。
  • Synchronization（同步）： 全局 GPU 同步和 GPU 间集体通信（如张量并行中的 AllReduce）会为每次操作增加数微秒。
  • Topology Ignorance（拓扑感知缺失）： GPU chiplet 架构内的非均匀内存访问 (NUMA) 效应可能会引入延迟偏差。

Kog 的解决方案：协同设计与 Monokernel

为了夺回这些流失的微秒，Kog 在关键路径中避免使用 PyTorch 或 Triton 等高级框架，而是选择手工编写的 GPU 代码 (CUDA/PTX 和 HIP/CDNA ISA)。他们的方案依赖于三个核心创新：

1. The Monokernel Runtime

而不是执行一系列 kernel，Kog 使用的是持久化 monokernel。这是一个驻留在 GPU 上的单一程序，它处理整个解码路径——MatMul、attention、normalization 和 sampling——而无需将控制权返回给 CPU。这完全消除了 kernel launch 开销和主机端调度。

2. KCCL (Kog Collective Communication Layer)

标准通信库是为吞吐量优化的，而不是为微秒级的延迟。KCCL 是一个在汇编级别进行调优的自定义层，旨在将 AllReduce 延迟保持在 3µs 以下，确保 GPU 间通信不会阻塞内存流式传输过程。

3. Laneformer 架构

Kog 引入了 Laneformer 架构，该架构利用了 Delayed Tensor Parallelism (DTP)。DTP 改变了解码过程的依赖结构，允许跨设备通信与计算重叠，而不是阻塞关键路径。

扩展到前沿 MoE 模型

虽然 3,000 tokens/s 的基准测试使用的是 2B 模型，但 Kog argue 认为同样的原理也适用于大型混合专家模型 (MoE)。因为 MoE 模型在每个 token 时仅激活其总参数量的一小部分，所以带宽瓶颈取决于活跃参数，而非总参数。

Kog 估计，通过将他们的 monokernel 和 KCCL 优化应用于前沿模型（如 DeepSeek 或 Qwen），他们可以在标准数据中心 GPU 上将输出速度推向 1,000–5,000 tokens/s 的范围，特别是随着下一代硬件（如 NVIDIA Rubin 或 AMD MI450）增加 HBM 带宽。

批判性视角与权衡

尽管基准测试结果令人印象深刻，但社区就该技术的实际应用提出了几个重要观点：

  • 模型效用 vs. 速度： 一些用户指出，如果模型缺乏前沿模型的智能水平，那么高 token 吞吐量是毫无意义的。正如一位评论者所说，“3.4k tok/s 的纯粹废话”并不是一个可行的产品。

  • 硬件定义： 对于什么是“标准 GPU”，存在争议。批评者指出 8× H200 节点是顶级的、精英级的数据中心硬件，而非易于获取的“标准”GPU。

  • Scaling 扩展性问题： 怀疑论者质疑，在 2B 模型上看到的线性增益是否能转化到 30B+ 模型上，因为在这些模型中，内存压力和通信开销的扩展方式不同。

最终，Kog 的工作凸显了 AI 基础设施领域的根本性转变：随着我们从聊天机器人转向自主智能体，行业必须从追求最大化聚合吞吐量转向最小化单个、顺序思维链的延迟。

Sources