从请求到 TPU 执行:vLLM、框架、算子、LLVM、Runtime 与 Linux 驱动

当用户向大模型服务发送一条请求时,这条请求不会直接变成一条“TPU 指令”。它要先经过在线服务引擎的调度,被框架表示成 Tensor 和算子,再由编译器将高层计算逐步 lowering 成目标硬件可以执行的程序。运行时负责加载程序、管理设备内存和提交任务,Linux 驱动负责安全地访问设备,最后才由 TPU 的矩阵单元、向量单元和存储系统完成计算。

把整条链路压缩成一行,就是:

1
2
3
4
5
6
7
8
请求
→ vLLM / API Server
→ Model Runner / 深度学习框架
→ Tensor、Operator 与计算图
→ 编译器、Kernel 与目标可执行程序
→ Runtime
→ Linux Device Driver
→ TPU Hardware

这篇文章沿着一次 LLM 请求向下追踪,解释每一层接收什么、产生什么,以及 LLVM、MLIR、XLA、vLLM 和 Linux 内核分别处在什么位置。

先建立正确的整体关系

这条栈里同时存在三类问题:

  1. 服务问题:请求排队、continuous batching、KV cache、并行与吞吐,由 vLLM 一类推理引擎负责。
  2. 计算问题:模型如何表示成 Tensor 和算子,算子如何优化并变成目标程序,由框架和编译器负责。
  3. 设备问题:程序如何被加载、提交并在硬件上执行,由 Runtime、驱动和硬件共同负责。

它们不是几条互相替代的路径,而是一次请求经过的不同层次。

Request / 用户请求prompt, sampling parameters, SLA vLLM / Serving EngineAPI, scheduler, batching, KV cache, distributed coordination Model Runner / Deep Learning Frameworkmodel forward through framework operators Tensor + Operator + Graph IRMatMul, RMSNorm, Attention; shape, dtype, dependency Compiler / 编译器fusion, layout, tiling, scheduling, lowering, code generationMLIR / XLA / LLVM may participate here Kernel / Target Executablecompiled program, kernel binary, command metadata Runtime / 用户态运行时load, allocate, bind buffers, launch, synchronize Linux Kernel + Device Driverdevice access, queues, DMA/IOMMU, interrupts, isolation TPU Hardwarememory, DMA, matrix units, vector units 服务与控制模型语义编译与生成提交与执行

这张图需要注意两个边界:

  • vLLM 不取代框架和编译器。它决定“这一轮运行哪些请求”,Model Runner 和编译后端决定“这一轮计算怎样执行”。
  • Runtime 不等于驱动。Runtime 通常位于用户态,驱动位于 Linux 内核态,二者通过系统调用和设备接口通信。

每一层究竟传递什么

层 接收的主要输入 产生的主要输出 主要职责
vLLM / Serving 文本请求、采样参数、请求状态 本轮需要执行的 token/sequence batch 调度、批处理、KV cache、并行协调
Framework / Model Runner token ids、position、KV block table、模型权重 Tensor 运算与算子调用 执行模型 forward,维护张量语义
Graph / Operator IR 算子、Tensor、shape、dtype、依赖关系 可优化的中间表示 表达“算什么”
Compiler Graph IR、硬件信息、编译配置 kernel 或目标 executable 融合、lowering、布局、切块、代码生成
Runtime executable、输入输出 buffer、stream/queue 设备命令及同步事件 加载程序、分配内存、绑定参数、提交执行
Linux Driver 用户态命令、内存映射请求 设备队列项、DMA 映射、完成事件 设备访问、隔离、资源管理、中断处理
TPU Hardware 命令、数据、目标程序 输出 Tensor 真正执行矩阵、向量和数据搬运

“请求 → 框架 → 算子 → Runtime → 驱动 → 硬件”是正确的大方向,但在算子和 Runtime 之间还必须补上 编译器与可执行程序。Runtime 通常不会现场理解 Attention 的全部数学含义;它主要加载编译结果并把本轮输入交给设备执行。

请求如何进入框架:vLLM 负责调度,不负责替代计算栈

以在线生成请求为例,服务端先完成 tokenization,然后把请求交给推理引擎。vLLM 的核心任务包括:

  • 根据请求状态决定本轮调度哪些 sequence;
  • 把不同请求动态组成执行批次;
  • 区分和安排 prefill、decode 工作;
  • 分配、回收并映射 KV cache block;
  • 协调 tensor parallel、pipeline parallel 或 data parallel worker;
  • 在模型执行结束后完成采样和流式返回。

Prefill 一次处理一段 prompt,序列维度并行度较高;Decode 通常每个 sequence 每轮只生成一个新 token,却需要持续读取历史 KV cache。两者使用相同模型,但 shape、访存特征和性能瓶颈不同。

vLLM 的 scheduler 产生的不是硬件指令,而是一份“本轮执行计划”:哪些请求继续运行、每个请求处理哪些 token、对应的 KV block 在哪里。Worker/Model Runner 再把这些信息组织成 Tensor,调用框架里的模型 forward。

在当前 vLLM 架构中,可以把主路径简化成:

1
2
3
4
5
API Server
→ Engine Core:scheduler + KV cache manager
→ Worker / Model Runner:执行模型 forward
→ logits / sampling
→ 输出 token

因此,vLLM 位于框架执行之上,而不是和编译器并列的另一条硬件路径。

框架如何把模型变成 Tensor 和算子

模型库负责提供 Transformer、Attention、MLP 等模型实现;PyTorch、JAX、TensorFlow 这类深度学习框架负责 Tensor、算子、设备和执行语义。二者不应混为同一层。本文从 vLLM 请求出发时,Model Runner 对应的是 vLLM 所使用的 PyTorch 执行路径;JAX 和 TensorFlow 是其他 TPU 软件栈可能采用的框架,并不是说一次 vLLM 请求会依次经过三种框架。

以一个高度简化的 Transformer 片段为例:

1
2
3
4
5
x = rms_norm(hidden_states)
qkv = linear(x, qkv_weight)
q, k, v = split_qkv(qkv)
o = attention(q, k, v, kv_cache)
hidden_states = linear(o, out_weight)

从框架视角看,它会形成一系列带有 shape、dtype、layout 和依赖关系的 Tensor 运算:

1
2
3
4
5
6
7
hidden_states
→ RMSNorm
→ MatMul / Linear
→ Split / View / RoPE
→ Attention
→ MatMul / Linear
→ output

Tensor 是带元数据的数据对象

Tensor 不只是一个多维数组。影响编译和执行的关键属性至少包括:

属性 示例 为什么重要
shape [B, S, H]、[B, H_kv, S, D] 决定循环范围、并行度和 buffer 大小
dtype BF16、FP16、FP8、INT8 决定精度、带宽和硬件指令
layout / stride contiguous、blocked、QKV packed 决定数据怎样排列和访问
device CPU、GPU、TPU 决定使用哪个执行后端
sharding TP/DP 分片规则 决定单设备 local shape 和通信

对于 LLM,更值得关心的是 QKV、KV cache 和并行分片后的 layout,而不是泛泛地讨论 NCHW/NHWC。

Operator 定义“算什么”

Operator 描述数学或张量语义。例如:

1
2
3
MatMul(A, B) -> C
RMSNorm(X, weight, eps) -> Y
Attention(Q, K, V, mask) -> O

如果 $A\in\mathbb{R}^{M\times K}$、$B\in\mathbb{R}^{K\times N}$,那么 MatMul 的语义是:

这个定义不包含 tile 大小、DMA 次序、片上 SRAM 分配或矩阵单元指令。

Kernel 或 executable 定义“在目标硬件上怎样算”

同一个 MatMul operator 可以对应 CPU、CUDA、Triton、TPU 或其他 NPU 实现。目标实现会进一步确定:

  • M/N/K 如何分块;
  • 数据怎样从全局内存搬到片上存储;
  • 使用哪些矩阵或向量指令;
  • 是否与 Bias、Activation、Residual 等算子融合;
  • 中间结果何时保留、何时写回;
  • 搬运和计算是否使用 double buffering 重叠。

Operator 和 Kernel 并不是一一对应关系:一个 operator 可能根据 shape 选择多个 kernel;多个 operator 也可能被融合成一个 kernel 或一个更大的目标 executable。

编译器:从算子语义到目标可执行程序

框架可以 eager 地逐算子执行,也可以通过 graph capture、export 或 tracing 获得适合编译的图。具体机制因框架而异,不能笼统地说“所有框架总会先生成一张静态图”。

一个典型的机器学习编译流程是:

1
2
3
4
5
6
7
8
9
10
Framework Program
→ Graph Capture / Export
→ Graph IR
→ Canonicalization 与 Shape 分析
→ Operator Fusion
→ Lowering
→ Layout / Tiling / Scheduling
→ Bufferization 与 Memory Planning
→ Target Code Generation
→ Kernel / Executable

其中:

  • Graph IR 保留 MatMul、Attention、Reduce 等较高层语义;
  • Tensor/Loop IR 更接近索引、循环、tile 和并行映射;
  • 低层 IR 更接近指针、控制流、向量指令和目标 ABI;
  • 目标 executable 是 Runtime 能够加载或调用的最终产物。

编译既可能是 AOT,也可能是 JIT。在线推理引擎可以在模型加载或首次遇到某类 shape 时编译,并在后续请求中复用缓存的 executable。因此,“编译器只负责离线、vLLM 只负责在线”并不准确。

LLVM、MLIR 与 XLA 在哪里

LLVM 最容易被误画成“框架和 Runtime 之间必经的一层”。更准确的说法是:LLVM 是许多编译后端会采用的基础设施和低层 IR,但不同硬件后端可以选择不同 lowering 路径;TPU 并不天然要求先变成 LLVM IR。

Framework / Captured GraphFX, exported program, model graph High-level ML IRStableHLO / HLO / MLIR tensor dialects / vendor IR CPU / GPU or LLVM-based backend TPU / Dedicated accelerator backend MLIR progressive loweringtensor → linalg/loop/vector → LLVM dialect XLA / Vendor target loweringfusion, layout, sharding, target scheduling LLVM IR + LLVM Backendoptimization and target code generation TPU Target Executablenot required to pass through LLVM IR CPU object / GPU kernel binaryloaded or launched by its runtime PJRT / Vendor Runtimeload and execute on TPU 后端可复用部分 MLIR/LLVM 基础设施,但不是固定必经关系

LLVM IR 是低层、硬件相关代码生成之前的表示

LLVM IR 是 SSA 形式的低层中间表示,擅长表达控制流、标量/向量运算、指针和内存访问。LLVM 的优化器和目标后端可以继续把 LLVM IR 生成 x86、AArch64、NVPTX、AMDGPU 等目标代码。

在机器学习栈里,LLVM 通常位于编译流程的后半段。MatMul、Attention 这类高层语义往往已经被分解、融合或映射成循环、向量和目标 intrinsic,才进入 LLVM IR。

MLIR 不是 LLVM IR 的新名字

MLIR 是一个多层 IR 基础设施,允许不同抽象层使用不同 dialect:Tensor、Linalg、Affine、SCF、Vector、GPU、LLVM dialect 等。一个典型 progressive lowering 可以写成:

1
2
3
4
5
6
7
高层 Tensor / Linalg
→ 循环、索引、并行与 Vector dialect
→ Bufferization:Tensor 语义变成显式内存
→ LLVM dialect
→ LLVM IR
→ LLVM target backend
→ 机器码或目标 kernel

MLIR 官方文档把生成 LLVM IR 描述为两步:先把其他 dialect 转换到可翻译的 LLVM dialect 或硬件相关 dialect,再翻译成 LLVM IR。并且 Tensor 类型在进入 LLVM dialect前需要先 bufferize,因为 LLVM IR 不直接保留高层 Tensor 抽象。

XLA 与 LLVM 的关系取决于目标后端

XLA 接收 StableHLO/HLO 一类高层机器学习表示,执行 target-independent 和 target-specific 优化,然后交给目标后端生成 executable。OpenXLA 官方架构明确说明:XLA 的 CPU 和 GPU 后端使用 LLVM 做低层 IR、优化与代码生成。

但是,不能由此推出“所有 TPU 程序都必须经过 LLVM IR”。TPU 有专门的目标后端和 executable 格式,其 lowering 路径可以绕开通用 LLVM IR。最稳妥的表达是:

1
2
3
4
5
6
7
PyTorch/XLA、JAX 或 TensorFlow
→ StableHLO / HLO
→ XLA 高层与目标相关优化
→ TPU backend
→ TPU executable
→ PJRT / TPU runtime
→ TPU

因此,在总分层图里应把 LLVM 放进“编译器内部”,而不是放在 Runtime、驱动或 TPU 之下。

Runtime:把编译结果变成一次真实执行

编译器负责生成 executable,Runtime 负责在一次具体请求中使用它。Runtime 通常工作在用户态,主要职责包括:

  • 初始化设备和 execution context;
  • 加载、缓存或选择目标 executable;
  • 分配、复用和释放设备 buffer;
  • 把输入 Tensor 绑定到 executable 参数;
  • 管理 stream、command queue、event 和同步;
  • 提交执行并等待或异步查询完成状态;
  • 处理多设备通信和错误状态。

不同生态中的 Runtime 名字不同,例如 CUDA Runtime/Driver API、PJRT、ONNX Runtime Execution Provider 或厂商 SDK。它们提供的抽象并不完全一致,但处在相似的系统位置。

编译器和 Runtime 的边界也不是绝对固定:JIT 系统可能由 Runtime 触发编译,Runtime 也可能根据动态 shape 选择不同 executable。判断它们的关键不是“发生在离线还是在线”,而是:

  • 编译器负责产生和优化程序;
  • Runtime 负责用具体输入组织并执行程序。

Linux 驱动:让用户态 Runtime 安全访问设备

Linux Kernel 和 Operator Kernel 是完全不同的两个概念:

名称 运行位置 主要职责
Linux Kernel / Device Driver CPU 内核态 设备访问、进程隔离、内存映射、队列、中断、错误处理
Operator Kernel / Device Program TPU/GPU/NPU 执行 MatMul、Attention、Norm 等计算的目标实现

以本地 PCIe 加速器为例,用户态 Runtime 可能通过 ioctl、mmap、poll 等接口与设备驱动通信。驱动可能负责:

  • 建立设备内存和主机内存映射;
  • 配置 DMA/IOMMU;
  • 管理 command queue 和 doorbell;
  • 校验并提交命令;
  • 处理中断、fault、timeout 和 reset;
  • 隔离不同进程或虚拟机的设备资源。

Linux 驱动通常不理解 Attention 的数学公式。到驱动这一层,高层算子大多已经变成 executable、command buffer、buffer address 和同步依赖。

需要注意的是,云 TPU、远程加速器和不同厂商设备可能把驱动细节封装在 libtpu、虚拟化层或远端服务中,不能把某一种 PCIe 驱动接口当成所有 TPU 的唯一实现。

硬件:矩阵、向量与数据搬运最终发生的地方

TPU 或其他 AI 加速器通常包含:

  • 全局设备内存或 HBM;
  • 片上 SRAM / scratchpad;
  • DMA 或其他数据搬运引擎;
  • 矩阵乘单元;
  • 向量/标量单元;
  • 片上互连和多设备互连;
  • command processor 和同步机制。

一个高效 MatMul 的目标程序不会只说“计算 $A\times B$”,而是会描述或隐含:如何分 tile、何时加载 A/B、partial sum 放在哪里、如何流水化搬运和计算,以及何时写回结果。

硬件只执行最终命令和目标程序,不认识 HTTP 请求,也不认识 vLLM request object。

完整追踪一次 Decode Step

现在把前面的层连接起来。假设多个请求都已经完成 prefill,vLLM 正准备生成下一个 token。

1. Scheduler 选择本轮请求

vLLM 根据 token budget、优先级、KV cache 可用空间和各 sequence 状态,决定本轮运行哪些请求,并生成调度结果。

2. Model Runner 组织输入 Tensor

Worker/Model Runner 根据调度结果准备:

  • 当前 token ids;
  • position ids;
  • sequence metadata;
  • KV cache block table;
  • attention slot mapping;
  • sampling 相关元数据。

3. 框架执行模型 forward

PyTorch、JAX 或 TensorFlow 模型依次表达 Embedding、RMSNorm、QKV Projection、Attention、MLP、LM Head 等算子。它们可能以 eager 方式调用,也可能进入已捕获和编译的图。

4. 编译器选择或生成 executable

如果对应 shape 和配置已有缓存,系统直接复用已编译程序;否则可能发生 graph capture、fusion、lowering 和 JIT compilation。

以 Attention 为例,高层表达式是:

编译后不一定真的生成 MatMul → Scale → Mask → Softmax → MatMul 五次独立 launch。它可以被融合或映射到专门的 attention kernel,并通过 block 算法避免完整写出 $[S,S]$ 中间矩阵。

5. Runtime 绑定 buffer 并提交执行

Runtime 把输入 Tensor、权重、KV cache 地址和输出 buffer 绑定到 executable,建立依赖和同步事件,然后向设备队列提交任务。

6. Linux 驱动把命令交给设备

驱动验证设备访问、建立 DMA 映射、更新 command queue,并通知硬件开始执行。设备完成后通过中断、fence 或轮询状态通知 Runtime。

7. TPU 执行目标程序

TPU 搬运权重和 activation,调用矩阵与向量单元,读取并追加 KV cache,最终生成 logits。

8. 结果回到 vLLM

Runtime 报告执行完成,Model Runner 返回 logits 或采样结果。vLLM 更新 sequence 状态,把新 token 流式返回客户端,并决定下一轮调度。

整条动态路径可以写成:

1
2
3
4
5
6
7
8
9
10
请求状态
→ vLLM Scheduler
→ Model Runner
→ Framework Operators
→ Compiled Executable
→ Runtime Launch
→ Driver Queue
→ TPU Execution
→ Logits / Token
→ 更新请求状态并进入下一轮

KV Cache 在各层分别是什么

KV cache 是理解“同一个对象在不同层呈现不同形态”的好例子。

层 看到的 KV cache
vLLM block allocator、block table、请求到物理块的映射
Framework 带 shape/dtype/device 的 K/V Tensor 或 cache view
Compiler buffer、索引计算、load/store、layout 和融合机会
Runtime 一组设备内存分配和 executable 参数
Driver 映射后的地址区间、DMA 和访问权限
Hardware 实际读写的设备内存与片上 tile

KV cache 的逻辑总量可以近似写成:

其中 2 表示 K 和 V,$L$ 是层数,$B$ 是 batch size,$S$ 是序列长度,$H{kv}$ 是 KV head 数,$D{head}$ 是 head dimension。这个公式是逻辑总量近似值,实际单设备占用还会受到 tensor parallel 分片、block 对齐、allocator metadata 和 padding 的影响。

PagedAttention 主要解决服务层和 kernel 层共同面对的 KV cache 管理问题:逻辑 token 顺序连续,但物理 KV block 可以分散分配。vLLM 维护逻辑块到物理块的映射,attention kernel 根据 block table 找到真实地址。这里可以看到:vLLM 负责“块属于谁”,kernel 负责“怎样按映射高效读取”。

性能优化分别发生在哪一层

层 常见优化 主要指标
vLLM / Scheduler continuous batching、chunked prefill、prefix cache、抢占策略 吞吐、TTFT、TPOT、并发数
Framework / Graph 去除冗余算子、静态化 shape、减少 graph break 可编译范围、Python overhead
Compiler fusion、layout、tiling、memory planning、quantization executable 性能、峰值内存
Kernel FlashAttention、fused RMSNorm、fused MoE、流水化 单 kernel latency、带宽利用率
Runtime executable cache、异步队列、buffer pool、通信重叠 launch overhead、设备利用率
Driver queue、内存映射、DMA、中断与隔离优化 提交开销、稳定性、尾延迟
Hardware 矩阵单元、HBM/SRAM、互连与 DMA 峰值算力、带宽、能效

Arithmetic intensity 可以帮助判断目标程序更偏计算受限还是带宽受限:

Prefill 中的大矩阵乘更容易利用矩阵单元吞吐;Decode 的小 batch 计算和 KV cache 读取则更容易受到内存带宽、调度与 launch overhead 影响。优化时不能只看某一个 kernel,也要看 vLLM 的 batch 组成和端到端请求指标。

五个最容易混淆的关系

1. vLLM 不是深度学习编译器

vLLM 首先是 LLM inference and serving engine。它会使用框架、编译器和 kernel,但核心价值是请求调度、KV cache 和高吞吐执行组织。

2. Operator 不等于 Kernel

Operator 是数学/张量语义,Kernel 是某种目标硬件实现。两者可能是一对多,也可能经过 fusion 变成多对一。

3. Runtime 不等于 Linux Driver

Runtime 位于用户态,面向 executable、buffer 和 queue;驱动位于内核态,面向设备资源、地址映射、命令提交与隔离。

4. Linux Kernel 不执行 MatMul

Linux 内核负责让进程安全访问设备。真正的矩阵计算由设备上的目标 kernel/program 和硬件单元执行。

5. LLVM 不是所有 TPU 流程的必经层

LLVM 常用于 CPU/GPU 等后端的低层优化与代码生成;MLIR 可以 lowering 到 LLVM IR,也可以走硬件专用 dialect 和后端。TPU 编译链更适合表述为 StableHLO/HLO → XLA/TPU backend → TPU executable,而不是强行插入 LLVM IR。

最后总结

一次 LLM 请求最终落到 TPU,可以分成连续的七层:

  1. vLLM 把用户请求转成一轮轮可执行的 batch,并管理 KV cache;
  2. Model Runner 与框架 把调度结果转成 Tensor 和模型 forward;
  3. Operator / Graph IR 表达 MatMul、Attention、RMSNorm 等计算语义;
  4. 编译器 完成 fusion、lowering、layout、tiling 和代码生成,产生 kernel 或 executable;
  5. Runtime 加载编译结果、管理 buffer、绑定参数并提交执行;
  6. Linux 驱动 完成设备访问、内存映射、队列和中断处理;
  7. TPU 硬件 执行数据搬运、矩阵与向量计算,产生最终输出。

LLVM 和 MLIR 位于第 4 层内部:MLIR 提供多层中间表示与 progressive lowering,LLVM IR 和 LLVM backend 为适合的目标提供低层优化和代码生成。XLA 则是一条面向机器学习图和多种硬件后端的编译体系;它的 CPU/GPU 后端可以使用 LLVM,而 TPU 后端不应被简单描述成必须经过 LLVM IR。

最短的一句话是:vLLM 决定这一轮算哪些请求,框架和算子描述要算什么,编译器决定怎样在目标硬件上算,Runtime 和驱动把程序送到设备,TPU 完成真正的计算。

参考资料