从大模型到 TPU 执行：框架、算子、内核、DLC、vLLM 与 Linux 内核

Posted on Edited on In , Views:

最近在梳理大模型基础设施的系统分层图：大模型、深度学习框架、Tensor、算子、算子内核、DLC/编译器、vLLM、TPU Runtime、Linux 内核和 TPU 硬件分别处在哪一层，各自解决什么问题。

这篇文章试图用更工程化的方式回答这个问题。最核心的关系是：大模型定义计算目标，框架表达计算图，算子定义数学语义，算子内核定义硬件实现，编译器/DLC 生成目标硬件可执行产物，vLLM 管理在线推理调度与 KV cache，Runtime 通过 Linux TPU 驱动把任务提交给 TPU 执行。

先分清两个 Kernel

“Kernel” 在这条栈里至少有两个完全不同的含义。第一个是 Linux Kernel，也就是操作系统内核；第二个是 Operator Kernel，也就是某个算子在 GPU/TPU/NPU 上执行的计算内核。

名称 专业含义 运行位置 主要职责
Linux Kernel 操作系统内核 CPU 特权态 进程调度、虚拟内存、文件系统、设备驱动、中断处理、权限隔离
TPU Operator Kernel TPU 算子内核 TPU/NPU/GPU 加速器 执行 MatMul、Softmax、Attention、LayerNorm 等算子的具体硬件实现

Linux Kernel 通常不理解 MatMul、Softmax、Attention 的数学语义；它负责让用户态程序安全、稳定地访问 TPU 设备。真正执行矩阵乘的是 TPU 上的 operator kernel。

整体分层关系

下面这张图把推理/编译/执行链路放在一张分层图中。图中有两条容易混淆的路径：DLC/Compiler 偏离线编译，vLLM 偏在线推理服务调度。二者最终都需要通过后端 Runtime、Linux Driver 和硬件执行能力完成计算。

Large Language Model / 大语言模型Transformer Blocks, Weights, Tokenizer-facing Semantics Deep Learning Framework / 深度学习框架PyTorch, TensorFlow, JAX, Transformers Computational Graph / 计算图Operators as Nodes, Tensors as Edges Compiler / DLC ToolchainGraph Optimization, Lowering, Code Generation vLLM Serving EngineScheduler, Continuous Batching, PagedAttention, KV Cache Runtime / SDK / 用户态运行时Load Artifacts, Allocate Buffers, Submit Command Buffers Linux Kernel + TPU Device Driverioctl, mmap, DMA Mapping, Command Queue, Interrupts TPU Hardware / TPU 硬件Global Memory, Local SRAM, DMA, Matrix Units, Vector Units

从模型到计算图

大模型：高层函数与参数集合

大模型可以看作一个由大量参数组成的函数。以 LLM 为例，模型通常由 Embedding、若干 Transformer Block、Normalization、MLP、Attention 和 LM Head 组成。大模型本身描述的是整体函数形式，而不是目标硬件上的具体执行策略。

框架：模型表达、张量管理与图导出

深度学习框架负责描述模型结构、管理 Tensor、加载权重、执行自动微分、导出计算图，并调用后端 Runtime 或编译器。常见框架包括 PyTorch、TensorFlow、JAX、MindSpore 和 PaddlePaddle。

模型代码通常是命令式或函数式写法；框架会把其中的张量运算记录成一张可分析、可优化的计算图。这张图不是图片，而是由节点和边组成的中间表示：节点通常是算子，边通常表示 Tensor 数据依赖。

例如 PyTorch 代码：

1
y = torch.softmax(x @ w, dim=-1)

从图语义上看，这段代码至少包含输入 Tensor x、权重 Tensor w、矩阵乘 MatMul、以及 Softmax。它可以抽象成：

1
2
3
x ─┐
   ├─ MatMul ── Softmax ── y
w ─┘

这就是计算图的核心作用：把“用户写的模型代码”转成“编译器和 Runtime 能理解的算子依赖关系”。框架层负责表达这个关系，但不一定决定它最终在 TPU 上如何分块、如何搬运数据、如何复用片上缓存。

Tensor：计算图边上的数据对象

Tensor 不只是多维数组，还包含一组影响编译和执行的重要元数据。

属性 含义 示例
shape 维度大小 $[B,S,H]$
dtype 元素类型 FP32、FP16、BF16、INT8
layout 内存布局 NCHW、NHWC、blocked layout
stride 各维度内存步长 contiguous / non-contiguous
device 所在设备 CPU、GPU、TPU、NPU

LLM 中常见 hidden states 可以表示为 [B, S, H]，例如 [1, 2048, 4096]：B 是 batch size，S 是 token 序列长度，H 是 hidden size。图像模型里常见的 NCHW / NHWC 则是在描述同一批图片数据的不同内存排列：N 是 batch，C 是 channel，H 是 height，W 是 width。

在大多数 row-major 存储里，最后一个维度最连续。因此 NCHW 更容易连续扫同一个 channel 的宽度方向，NHWC 更容易连续读取同一个像素位置的多个 channel。算子的访问模式如果和 layout 对齐，DMA 搬运、cache locality、向量化和 kernel selection 都会更友好；如果不对齐，就可能需要 stride access 或额外 layout conversion。

下面用一张纯文本图看同一张 N=1, C=3, H=2, W=3 图片在内存中的排列差异：

1
2
3
4
5
6
7
8
9
10
NCHW = [N, C, H, W]

内存顺序：先放完整 R 平面，再放完整 G 平面，再放完整 B 平面

R channel: R00 R01 R02 R10 R11 R12
G channel: G00 G01 G02 G10 G11 G12
B channel: B00 B01 B02 B10 B11 B12

连续方向：同一个 channel 内，H x W 空间位置连续
适合访问：同一 channel 的空间窗口，例如卷积在 feature map 上滑动
1
2
3
4
5
6
7
8
9
10
11
12
NHWC = [N, H, W, C]

内存顺序：每个像素位置的 R/G/B 连续放在一起

pixel(0,0): R00 G00 B00
pixel(0,1): R01 G01 B01
pixel(0,2): R02 G02 B02
pixel(1,0): R10 G10 B10
...

连续方向：同一个 pixel 的多个 channel 连续
适合访问：同一空间位置的 channel 向量，例如按 channel 向量化读取
Layout 最连续的维度 更友好的访问模式 不匹配时的代价
NCHW W，也就是同一 channel 内的宽度方向 同一 channel 的空间窗口扫描 读取同一 pixel 的多个 channel 可能是 stride access
NHWC C，也就是同一 pixel 的 channel 向量 同一 pixel 的多个 channel 一次性读取 扫同一 channel 的空间平面可能不连续

视觉上可以把它理解成下面这样：

NCHW / NHWC 的连续方向不同 NCHW 把 channel 平面分开存；NHWC 把同一个 pixel 的 channel 放在一起。 NCHW = [N, C, H, W] R planeH x W 连续 G planeH x W 连续 B planeH x W 连续 NHWC = [N, H, W, C] R G B R G B R G B pixel 0pixel 1pixel 2 NCHW适合扫同一 channel 的空间窗口 NHWC适合读同一 pixel 的 channel 向量 Layout 不匹配可能需要 stride access 或 transpose

Operator：计算图节点上的数学语义

Operator 定义“算什么”。例如 Elementwise Add、Reduction Sum、MatMul、Softmax、RMSNorm、Attention 都是算子。以矩阵乘为例：

1
C[i, j] = sum(k=1..K, A[i, k] * B[k, j])

如果 $A\in\mathbb{R}^{M\times K}$，$B\in\mathbb{R}^{K\times N}$，那么输出 $C\in\mathbb{R}^{M\times N}$。这个定义描述的是数学语义，不描述具体硬件上如何高效执行。

Operator Kernel：目标硬件上的执行实现

Operator Kernel 定义“怎么在硬件上算”。同一个 MatMul operator 可以对应 CPU kernel、CUDA kernel、Triton kernel、TPU kernel 或 NPU kernel。TPU MatMul kernel 通常会涉及 tile 划分、global memory 到 local SRAM 的搬运、DMA、matrix unit 调用、partial sum 累加和写回。

Operator / 算子数学语义Cᵢⱼ = Σₖ AᵢₖBₖⱼ Operator Kernel / 算子内核tiling, schedulingload → compute → store Hardware / TPU 硬件SRAM, DMAmatrix units, vector units loweringexecute

DLC / 编译器：从图语义到硬件执行计划

DLC 在不同厂商语境下可能表示 Deep Learning Compiler，也可能表示 Compiled Model Artifact。这两种含义都围绕同一件事：把高层模型图转换为目标硬件上可执行、正确且高效的执行计划。

生态 常见编译产物
NVIDIA TensorRT engine / plan
Qualcomm SNPE dlc
Ascend om
Sophon bmodel
其他 TPU/NPU 厂商 model binary / compiled artifact

典型编译链路如下：

Model ExportPyTorch / ONNX Graph IRoperators + tensors Graph Optimizefusion, folding LoweringGraph IR → Tensor IR Schedulingtiling, layout Memory Planningbuffers, reuse Code Generationkernel binary Compiled ArtifactDLC / model binary Runtimeload + execute

其中，Graph IR 保留较高层的算子语义；Tensor IR 更接近循环、索引和 buffer；Kernel IR 更接近目标硬件执行方式。编译器的核心任务不是“把 Python 翻译成二进制”这么简单，而是同时决定算子融合、layout、tiling、memory planning、kernel selection 和 code generation。

vLLM：在线推理服务引擎，不是离线编译器

vLLM 的定位是 LLM inference serving engine。它关注在线推理系统问题：如何接收多用户请求，如何调度 prefill 和 decode，如何做 continuous batching，如何管理 KV cache，如何通过后端 Runtime 调用硬件 kernel。

Client RequestsOpenAI API / Python API Tokenizertext ↔ token ids Request Scheduleradmission and priority Continuous Batchingmerge prefill and decode workloads PagedAttentionblock table maps tokens to KV blocks Model Executorprefill, decode, sampling Backend RuntimePyTorch / XLA / vendor SDK GPU / TPU / NPUoperator kernels execute

vLLM 与 DLC/Compiler 的边界可以这样理解：vLLM 管在线请求与 KV cache，编译器管图优化与 kernel 生成，Linux driver 管设备访问，TPU hardware 管实际计算。若要让 vLLM 使用 TPU/NPU，需要 XLA/torch_xla、厂商 runtime adapter 或专门的 backend 支持。

Prefill、Decode、KV Cache 与 PagedAttention

LLM 推理通常分成两个阶段：Prefill 和 Decode。Prefill 处理完整 prompt，计算量大但序列维度并行度高；Decode 每次生成一个 token，单步计算较小，但需要反复读取历史 KV cache。

KV cache 的近似内存占用可以写成：

1
KV_bytes = 2 × L × B × S × H_kv × D_head × bytes(dtype)

其中 2 表示 K 和 V，L 是层数，B 是 batch size，S 是序列长度，H_kv 是 KV head 数，D_head 是每个 head 的维度。

Prompt Tokensinput sequence Prefillbuild KV cache Decoderead cache, append KV PagedAttention: 用 block table 管理 KV cache 逻辑 token 顺序：模型看到 token 是连续的 tok 0-15 tok 16-31 tok 32-47 block table：记录逻辑块到物理块的映射 0→17 1→04 2→29 物理 KV blocks：真实内存里可分散、可复用 KV #17 KV #04 KV #29 好处：不要求 KV cache 连续分配，请求结束后释放物理块，其他请求可以复用。

PagedAttention 的核心思想类似虚拟内存分页：请求里的 token 逻辑上是连续的，但它们对应的 KV cache 不必放在一整块连续内存里。vLLM 会把 token 序列切成固定大小的逻辑块，再通过 block table 映射到真实设备内存中的 physical KV blocks。这样可以降低 KV cache 碎片；一个请求结束后，它占用的物理 KV blocks 可以被释放并复用给其他请求，从而提高长上下文和多请求并发时的设备内存利用率。

TPU 与 Linux 内核的关系

TPU 是加速器硬件，Linux Kernel 是操作系统内核。用户态 Runtime 通常通过 ioctl、mmap、poll 等接口与 TPU device driver 通信；驱动再负责 DMA 映射、command queue、interrupt handling 和错误状态返回。

User Space / 用户态 Application TPU Runtime / SDK DLC / Model Binary Kernel Space / Linux 内核态 ioctl / mmap / poll TPU Device Driver DMA / IOMMU Device / TPU 硬件 Command Queue TPU Cores Interrupt

因此，一个 MatMul 从 Python 到 TPU 的路径可以用专业语言概括为：框架识别 aten::matmul 或等价图节点，导出到 ONNX 或厂商 IR，编译器执行 lowering 和 kernel selection，Runtime 加载 compiled artifact 并提交 command buffer，Linux TPU driver 完成设备通信，TPU hardware 执行目标 MatMul kernel。

例子一：Linear + Bias + ReLU 融合

框架层表达式如下：

1
y = torch.relu(x @ w + b)

数学形式是：

1
Y = max(0, XW + b)

如果不融合，执行路径包含 MatMul、Add、ReLU 三个算子，可能产生两个中间张量；如果融合，编译器可以生成一个 Fused MatMul + Bias + ReLU Kernel，在累加完成后直接加 bias、应用激活函数并写回最终输出。

Before Fusion / 融合前 MatMul Add ReLU Output tmp1tmp2 After Fusion / 融合后 Fused MatMul + Bias + ReLU Kernelaccumulate → add bias → activation → store once Output

融合的直接收益是减少中间 tensor、减少 global memory 读写、减少 kernel launch overhead，并提升端到端吞吐。

例子二：Scaled Dot-Product Attention

Scaled Dot-Product Attention 的经典公式来自 Transformer 论文 Attention Is All You Need 的 3.2.1 节：

1
Attention(Q, K, V) = softmax((QK^T) / sqrt(d_k)) V

其中 d_k 是 key/query 向量维度，除以 sqrt(d_k) 是为了避免点积值随维度增大而过大，导致 softmax 进入梯度很小的区域。在 LLM 的 causal self-attention 中，softmax 前通常还会加入 causal mask，阻止当前位置看到未来 token。

这里最容易抽象的是 scores 为什么会变成 [B, H, S, S]。先看 Q、K、V 的常见 shape：

1
2
3
Q: [B, H, S, D]
K: [B, H, S, D]
V: [B, H, S, D]

其中 B 是 batch size，H 是 attention head 数，S 是序列长度，D 是每个 head 的维度。对每个 batch、每个 head 来说，注意力会让“每个 query token”去和“每个 key token”算相似度：

1
2
3
Q[b,h]      : [S_query, D]
K[b,h]^T    : [D, S_key]
Q[b,h]K[b,h]^T -> [S_query, S_key]

自注意力里 S_query = S_key = S，所以每个 head 会得到一个 [S, S] 的 scores 矩阵；再乘上 batch 和 head 两个维度，整体就是 [B, H, S, S]。这两个 S 不是重复写错了：前一个 S 表示 query token 位置，后一个 S 表示 key token 位置。

普通实现可能 materialize scores 和 probs 两个 [B, H, S, S] 规模的中间矩阵。当序列长度 S 很大时，内存读写会成为主要瓶颈。

Q[B,H,S,D] K[B,H,S,D] V[B,H,S,D] QKᵀ scores[B,H,S_query,S_key] Scale1 / sqrt(dₖ) Softmaxprobabilities MatMul Vweighted sum Output O[B,H,S,D] Optimizationblockingonline softmaxavoid full[B,H,S,S]

FlashAttention 类优化的核心不是改变数学公式，而是改变执行方式：按 block 计算 $QK^T$，在线维护 softmax 的 max 和 sum，不完整写出 scores/probs，从而减少 HBM/global memory traffic。

优化到底在优化什么

优化项 专业含义 主要目的 例子
Operator Fusion 多个相邻算子合成 fused kernel 减少中间张量和 kernel launch MatMul + Bias + ReLU
Layout Optimization 选择硬件友好的内存布局 提高连续访存和对齐效率 NCHW、NHWC、blocked layout
Memory Planning 规划 buffer 生命周期和复用 降低峰值设备内存 activation buffer reuse
Kernel Selection 根据 shape/dtype/layout 选实现 提高单算子性能 small GEMM vs large GEMM
Tiling / Blocking 把大计算切成 tile 提高片上缓存复用 MatMul 的 M/N/K tile
Quantization 降低数值精度 降低带宽和存储，提高吞吐 FP16、BF16、INT8、INT4
DMA/Compute Overlap 搬运和计算流水化 隐藏 DMA latency double buffering
Continuous Batching 动态合并在线请求 提高 LLM serving 吞吐 vLLM scheduler
KV Cache Paging 按 block 管理 KV cache 降低碎片，提高并发 vLLM PagedAttention

Arithmetic intensity 可以帮助判断算子更可能受计算限制还是带宽限制：

1
AI = FLOPs / Bytes moved

如果 $AI$ 很低，算子通常更容易受内存带宽限制；如果 $AI$ 很高，才更可能受计算单元吞吐限制。

最后总结

整条链路可以概括为：Large Model 定义任务语义，Framework 或 vLLM 负责模型表达与在线服务，Computational Graph 将模型分解为 Operators，Compiler 或 Backend Runtime 选择和生成 Operator Kernels，Runtime 通过 Linux Driver 提交硬件任务，TPU Hardware 最终执行这些算子内核。

更短地说：模型由算子组成，算子由硬件内核执行，Runtime 通过 Linux 驱动把硬件任务提交给 TPU。Linux kernel 让 TPU 能被系统安全访问；TPU operator kernel 让算子能在 TPU 上高效计算；vLLM 让大模型在线推理能高吞吐地调度请求和管理 KV cache。