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

Posted on In , Views:

最近在梳理“agent + TPU 算子生成”的学习路径时，我发现真正需要先建立的不是某个 SDK API 的记忆，而是一张稳定的系统分层图：大模型、深度学习框架、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。

例如 PyTorch 代码：

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

从图语义上看，这段代码至少包含矩阵乘和 Softmax 两类算子。框架层负责表达这个关系，但不一定决定它最终在 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]$。

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

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

如果 $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 PagedAttentionlogical tokens → physical 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 的近似内存占用可以写成：

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

Prompt Tokensinput sequence Prefillbuild KV cache Decoderead cache, append KV PagedAttention: logical token blocks → physical KV blocks logical 0 logical 1 logical 2 page 17 page 04 page 29

PagedAttention 的核心思想类似虚拟内存分页：逻辑上连续的 token block 不要求物理上连续存放。这样可以降低 KV cache 碎片，提高长上下文和多请求并发时的设备内存利用率。

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)

数学形式是：

如果不融合，执行路径包含 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

标准注意力公式是：

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

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

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 可以帮助判断算子更可能受计算限制还是带宽限制：

如果 $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。