Key Takeaways
• Marlin 是一种混合精度矩阵乘法内核，实现了 LLM（大型语言模型）矩阵乘法性能的重大突破，利用 FP16xINT4 运算在批处理大小达到 32 时实现了 4 倍加速。
• Marlin 采用了三项高级技术来优化 GPU 资源使用，如异步全局权重加载、循环共享内存队列以及任务调度和同步（asynchronous global weight loads, a circular shared memory queue, and task scheduling and synchronization）。这些方法提升了整体计算效率，充分发挥了 GPU 的性能。
• Marlin 实现了端到端的加速，通过加速生成首个 token 的时间（TTFT）和每个输出 token 的时间（TPOT）来提升常见推理 API 场景的性能。
在快速发展的 LLM 推理领域中，提高现代 GPU 的速度和效率已成为关键挑战。此时出现了 Marlin，一个突破性的 Mixed Auto-Regressive Linear 内核，可以在 FP16xINT4 矩阵乘法中实现前所未有的性能提升。Marlin 由 IST-DASLab 的 Elias Frantar 开发，并以地球上最快的鱼类之一命名。这种混合精度 LLM 推理内核在批量大小达到 32 个 token 时实现了理想的 4 倍加速，远超以往仅能在批处理大小为 1 时实现的有限加速。这使得 Marlin 成为大型服务、推测解码和高级多推理方案的出色选择。Neural Magic 已将 Marlin 集成并在 nm-vllm 中进一步扩展，以实现最佳的 GPU 上 LLM 服务。
在本文中，我们将深入探讨 Marlin 卓越的 LLM 推理性能背后的技术创新，探索优化现代 GPU 推理内核的挑战，并剖析 Marlin 为实现其无与伦比的效率而采用的关键技术和优化。
混合精度 LLM 推理背景
混合精度 LLM 推理计算对于机器学习而言是革命性的。通过结合不同的数值精度，混合精度推理既能减少内存占用，还能在不牺牲精度的情况下显著提高速度。
量化权重比量化激活更容易，因为权重在推理期间保持不变，而激活会根据输入动态变化。权重的静态特性使量化过程更简单，因为最佳量化参数可以离线确定并在推理中一致应用。而量化激活则需要在推理期间进行实时量化，由于激活值的动态范围和分布，这更具挑战性。此外，激活量化误差会在网络中传播，可能导致比权重量化更显著的精度下降。因此，权重量化通常被认为是简化 LLM 模型的首选方法。
GPTQ 是一种基于近似二阶信息的领先单步权重量化方法，既高效又准确。像 GPTQ 这样的技术可以将权重从 16 位压缩到 4 位，且对精度的影响很小。然而，混合精度 LLM 推理的全部潜力往往受到性能瓶颈和 GPU 资源利用不足的限制。尽管当计算资源充足时可以通过 4 位权重最大化内存带宽的有效利用，实现 4 倍于 FP16 权重的理想加速，但低效的计算利用率使得混合精度对多用户 LLM 服务不够理想。这时 Marlin 的出现扩展了混合精度 LLM 推理的有效范围。
为什么 Marlin 是混合精度 LLM 推理的最新成果
Marlin 的架构基于一个关键观察：现代 GPU 通常具有 100 到 200 的浮点运算（FLOPs）与字节比。因此，只要每 4 位量化权重执行少于 25-50 次张量核乘累加运算，就有可能实现接近于 FP16 权重的理想 4 倍加速。这意味着仅权重量化的性能优势理论上应适用于比现有方法大 4-8 倍的批量大小。
下图来自 Marlin 的原始概述，展示了 Marlin 与其他常用 4 位推理内核在 NVIDIA A10 GPU 上的性能对比。所有内核都以组大小为 128 执行，但需注意比例格式并不完全一致。
现有内核在批处理大小为 1 时接近 3.87 倍的最佳加速（考虑到组尺度的 0.125 位存储开销），但随着 token 数量增加，它们的性能迅速下降。相反，Marlin 在所有批量大小下均能实现接近理想的加速，在批处理大小达到约 16-32 时可实现最大 3.87 倍加速。
混合精度 LLM 推理的关键技术和优化
Marlin 同时优化 GPU 资源如全局内存、L2 缓存、共享内存、向量核心和张量核心，主要采用以下几种技术：
- 异步全局权重加载
- 循环共享内存队列
- 任务调度与同步
以下是这些技术的详细说明。
异步全局权重加载
Marlin 利用 NVIDIA Ampere 架构中引入的 异步数据传输 功能来优化从全局内存加载权重。使用 cuda::memcpy_async API，Marlin 可以将神经网络权重非阻塞地直接复制到共享内存中。此技术无需线程或寄存器来管理数据传输，使这些资源可完全专注于计算任务。该并行执行模式使 Marlin 能够将权重传输与进行中的计算重叠，从而有效地掩盖传统上访问全局内存的延迟。
流水线处理策略进一步优化了此方法。Marlin 使用 cuda::pipeline 预取下一部分 N+1 的权重，同时对当前部分 N 进行计算。这种流水线预取对保持计算的连续性非常重要，有效地隐藏了全局内存访问的延迟。此外，这种异步权重加载方法旨在降低对 L2 缓存的影响。通过将权重直接传输到共享内存，绕过 L2 缓存，Marlin 保留了缓存容量用于存储输入和输出激活，从而提高了缓存利用率和性能效率。
// Asynchronously fetch the next A, B and s tile from global to the next shared memory pipeline location.
auto fetch_to_shared = [&] (int pipe, int a_off, bool pred = true) {
if (pred) {
int4* sh_a_stage = sh_a + a_sh_stage * pipe;
#pragma unroll
for (int i = 0; i < a_sh_wr_iters; i++) {
cp_async4_pred(
&sh_a_stage[a_sh_wr_trans[i]],
&A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off],
a_sh_wr_pred[i]
);
}
int4* sh_b_stage = sh_b + b_sh_stage * pipe;
#pragma unroll
for (int i = 0; i < b_sh_wr_iters; i++) {
cp_async4_stream(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr], B_ptr[i]);
B_ptr[i] += b_gl_rd_delta_o;
}
// Only fetch scales if this tile starts a new group
if (group_blocks != -1 && pipe % (group_blocks / thread_k_blocks) == 0) {
int4* sh_s_stage = sh_s + s_sh_stage * pipe;
if (s_sh_wr_pred)
cp_async4_stream(&sh_s_stage[s_sh_wr], &s[s_gl_rd]);
s_gl_rd += s_gl_rd_delta;
}
}
// Insert a fence even when we are winding down the pipeline to ensure that waiting is also correct at this point.
cp_async_fence();
};
循环共享内存队列
Marlin 包含多个共享内存缓冲区以构建循环缓冲系统。该系统旨在高效管理共享内存中的数据加载和使用，允许数据的连续加载、处理和卸载，而无需可能导致计算停顿的同步点。
循环队列架构包含多个共享内存缓冲区，使 Marlin 能够在当前缓冲区执行计算时将数据预加载到后续缓冲区中。此策略显著提高了数据处理的吞吐量，确保计算和数据加载可同时进行而互不干扰。循环缓冲机制实现了缓冲区之间的转换，每个缓冲区根据循环的位置动态分配为数据加载的下一个目标或计算的当前来源。
// Register storage for double buffer of shared memory reads.
FragA frag_a[2][thread_m_blocks];
I4 frag_b_quant[2];
FragC frag_c[thread_m_blocks][4][2];
FragS frag_s[2][4];
…
// Wait until the next thread tile has been loaded to shared memory.
auto wait_for_stage = [&] () {
// We only have `stages - 2` active fetches since we are double buffering and can only issue the next fetch when
// it is guaranteed that the previous shared memory load is fully complete (as it may otherwise be overwritten).
cp_async_wait<stages - 2>();
__syncthreads();
};
任务调度与同步
Marlin 采用了高效的任务调度和同步策略，最大限度地利用了矩阵乘法期间 GPU 资源，特别是流多处理器（SM）。该策略的核心是 Stream-K 并行化 方法，其中问题的 M 和 N 维度均匀划分，K 维度则不均匀划分。
例如，考虑一个问题形状为 32x512x4096（MxNxK）的矩阵乘法（C = AxB）。通过将 M 和 N 分别按 1 和 4 均匀划分，我们得到 4 个大小相等的（32x128）输出瓦片标记为 0-3。
然而，对 K 维度的均匀划分会导致 SM 的次优利用。假设有 6 个 SM 的 GPU，仅对 K 均匀划分将导致 4 个可用的 SM 分配了输出瓦片，剩余 2 个 SM 空闲。Marlin 通过在 K 维度上采用非均匀划分方案解决了此问题。
为了尽量减少同步开销，Marlin 按顺序分配每个输出瓦片的工作。即，所有与输出瓦片 0 相关的工作在继续分配给输出瓦片 1 之前全部完成。这样的顺序分配确保了同步和缩减操作的高效执行。
Marlin 的一个关键优势在于它能够利用 L2 缓存来高效进行部分结果缩减。在 LLM 推理的解码阶段（这是 Marlin 的主要目标），输出大小相对较小。通过确保输出适合 L2 缓存，Marlin 显著加快了部分结果缩减操作，减少了同步对性能的影响。
通过将 Stream-K 并行化方法与智能工作分配、高效同步和有效利用 L2 缓存相结合，Marlin 在广泛的 GPU 架构和问题规模上实现了最佳性能和可扩展性。这种复杂的任务调度和同步策略是 Marlin 能够在混合精度 LLM 推理中实现接近理想加速的关键。
对实际应用的影响
Marlin 在多种推理方案中的出色表现使其在混合精度 LLM 推理领域成为最佳解决方案。相关图表数据显示，无论是 TTFT（生成首个 token 的时间）还是 TPOT（每个输出 token 的时间）延迟测量，Marlin 在更高的每秒查询数（QPS）下均优于其他方法。随着 QPS 增加，即代表实际应用中更高的需求和吞吐量，Marlin 相较于包括 GPTQ、AWQ 和 FP16 在内的其他方法始终保持较低的延迟。
在 10 QPS 下，TTFT 延迟约为 217 毫秒，而最接近的竞争方法延迟在 300 毫秒左右。对于 TPOT 延迟，Marlin 在 10 QPS 时保持 86 毫秒的延迟，远超其他方法（它们接近或超过 200 毫秒）。这些结果对于性能工程师至关重要，因为它们表明 Marlin 即使在需求增加时也能保持近乎最佳的加速和 GPU 利用率。
结论
Marlin 在混合精度 LLM 推理领域具有突破性的进展，推动了现代 GPU 性能和效率的极限。Marlin 的性能提升在实际应用中有广泛影响，特别是在高吞吐量和低延迟至关重要的场景中。