一条直线就能预测我的笔记本生成 LLM token 的速度

一条直线预测我的笔记本电脑生成 LLM token 的速度

我深入探究了一番，想看看内存带宽如何限制 LLM 生成 token 的速率。

我上一篇博文留了个悬念。对于 Qwen2.5-0.5B 模型，我发现不同提示大小下，每个输出 token 的时间（TPOT）大约都是 9 ms。但我当时没有解释为什么会这样。

后来我了解到，在大语言模型（LLM）推理的解码阶段，主要瓶颈是内存带宽，也就是 GPU 从 RAM 读取 LLM 模型权

对于每个模型，我通过两个步骤计算了所花时间的占比：

• 用模型大小除以有效带宽（163 GB/s）来估算读取权重所花的时间。

• 将这个读取时间除以预测的 TPOT。

可以看到，对于较小的模型，系统开销在生成一个新 token 所需时间中占据更大比例。对于最小的模型，我的笔记本电脑大约有 45% 的时间用于读取权重，而对于最大的模型，这一比例为 91%。

我是如何测量的

对于每个模型，我使用与上一篇博文中类似的 ttft_tpot 函数来测量 TPOT，并采用 PyTorch 的 MPS 后端，这让 PyTorch 可以在我的 MacBook 集成 GPU 上运行代码。

        import statistics
import time
import torch

def timed_generate(model, ids, n):
    torch.mps.synchronize()  # wait for any earlier GPU work to finish
    t0 = time.perf_counter()
    model.generate(ids, min_new_tokens=n, max_new_tokens=n, do_sample=False)
    torch.mps.synchronize()  # wait for this call's GPU work to finish
    return time.perf_counter() - t0

def ttft_tpot(model, ids, num_samples=5):
    timed_generate(model, ids, 1)  # warm-up
    ttft = statistics.median(timed_generate(model, ids, 1) for _ in range(num_samples))
    total = statistics.median(timed_generate(model, ids, 101) for _ in range(num_samples))
    return ttft, (total - ttft) / 100
      

然后我遍历这些模型：

        import gc
from transformers import AutoModelForCausalLM, AutoTokenizer

models = [
    ("google/gemma-3-270m-it", torch.bfloat16),
    ("google/gemma-3-270m-it", torch.float32),
    ("Qwen/Qwen2.5-0.5B-Instruct", torch.bfloat16),
    ("Qwen/Qwen2.5-0.5B-Instruct", torch.float32),
    ("google/gemma-3-1b-it", torch.bfloat16),
    ("google/gemma-3-1b-it", torch.float32),
    ("Qwen/Qwen2.5-1.5B-Instruct", torch.bfloat16),
    ("Qwen/Qwen2.5-1.5B-Instruct", torch.float32),
    ("Qwen/Qwen3-0.6B", torch.bfloat16),
    ("Qwen/Qwen3-0.6B", torch.float32),
    ("Qwen/Qwen3-1.7B", torch.bfloat16),
    ("Qwen/Qwen3-1.7B", torch.float32),
]

text = open("input.txt").read()
rows = []  # (size in GB, TPOT in seconds)

for name, dtype in models:
    model = AutoModelForCausalLM.from_pretrained(name, dtype=dtype).to("mps").eval()
    tok = AutoTokenizer.from_pretrained(name)
    input_ids = tok(text[:20_000], return_tensors="pt").input_ids.to("mps")

size_gb = sum(p.numel() * p.element_size() for p in model.parameters()) / 1e9
    _, tpot = ttft_tpot(model, input_ids[:, :128])
    rows.append((size_gb, tpot))

del model
    gc.collect()
    torch.mps.empty_cache()
      

对于每个模型，我使用 Tiny Shakespeare 文本的前 128 个 token 作为输入提示。

收集数据后，进行线性拟合就像这样简单：

        import numpy as np

gb = [r[0] for r in rows]
ms = [r[1] * 1000 for r in rows]  # seconds -> ms
slope, intercept = np.polyfit(gb, ms, 1)  # ms per GB, ms
      

这个测量的注意事项

这只是一个快速实验，用来感受一下这些数字，并不是严格的基准测试。话虽如此，我的方法有几个缺陷应该指出来。

一个特殊之处是，仅对 Gemma 3 270M 的数据点（bfloat16 和 float32）做线性拟合，会 implied 约 258 GB/s 的内存带宽，这远高于 200 GB/s 的规格上限。我能够在多次运行中复现类似结果，这让我认为这不是噪声导致的。

我认为简单的结论是，固定成本并不总是对每个模型都相同。即便考虑的是同一个模型的 bfloat16 或 float32 版本，这一点也成立。我还不完全清楚底层原因是什么，但我怀疑这可能与模型架构有关。

另一个需要注意的是，我没有使用优化过的推理服务器（例如 vLLM 或 llama.cpp）来运行这个实验。我预计，更优化的设置会降低这里测得的固定成本，但可能不会把带宽利用率提升那么多，因为进一步提升的空间很可能不大。用 mlx-lm 而不是 PyTorch 重新做这个实验可能会很有意思，mlx-lm 为 Apple silicon 提供了优化内核。

我的收获

现在，让我们回到我最初的问题：对于 Qwen2.5-0.5B，为什么 TPOT 约为 9 ms？线性拟合预测，对于这个模型，我的笔记本电脑读取权重大约需要 6.05 ms，再加上 4.03 ms 的固定成本，总计为 10.08 ms。这比我测得的 8.99 ms 高 12%，但它确实描绘出了生成每个 token 时时间是如何花掉的。

我承认，看到数据如此符合预期的线性趋势，我感到很惊讶。但现在我理解了这背后的机制，一切就说得通了。

标签