KKKZOZ’s Blog

ML Compiler Overview ML Compiler 位于深度学习框架和底层硬件之间，作用是把 TensorFlow、PyTorch、JAX 等框架描述的高层张量计算转换成能在 CPU/GPU/TPU 等硬件上高效执行的程序。 它的典型架构是 Frontend → IR → Optimization → Backend: Frontend 接收框架已经构造或捕获出来的计算图/算子程序（例如 tf.function 得到的 TF Graph、PyTorch Dynamo 得到的 FX Graph），转换成统一的中间表示 IR 中间层进行算子融合、常量折叠、布局变换、内存规划、并行化等与机器学习计算相关的优化 Backend 再根据目标硬件进行 lowering、调度和代码生成 因而 ML compiler 输入通常是带有 tensor shape、dtype、算子及数据依赖关系的计算图或 IR，输出则是面向特定设备的低层 IR、kernel 或可执行程序。 需要特别区分的是：graph capture 不一定属于 compiler 本身——例如 TensorFlow 由 tf.function tracing 捕获图，PyTorch 由 TorchDynamo 捕获图，然后 ML compiler 才接手，对这个图做优化并最终生成高效机器码。 IR Pipeline 下面这套 IR 划分描述的是一个典型 ML Compiler 从高层模型计算逐渐 lowering 到硬件可执行 kernel 的过程。不同编译器实际使用的 IR 名称和数据结构可能不同，但抽象层次通常可以归纳为： ...

vLLM 中同时存在多条与 FX 和 torch.compile 相关的执行路径，它们虽然共享 PyTorch 的 graph infrastructure，但解决的问题并不相同。理解这些路径的边界，是区分“模型结构改写”“Inductor 编译优化”和“CUDA Graph replay”的前提。 进入正文前，可以先区分三类用法： HF Transformers modeling backend：直接使用 fx.Tracer 分析 Hugging Face forward，识别可融合结构并通过 AST 改写模型源码 vLLM 自定义 torch.compile 主路径：由 Dynamo 产生完整 FX GraphModule，再由 vLLM 控制 graph partition、shape specialization、custom passes、cache 和 CUDA Graph integration 普通 @torch.compile 小函数：FX 主要作为 Dynamo、AOTAutograd 和 Inductor 的内部 IR，vLLM 通常不直接处理这些 graph Piecewise CUDA Graph 不是第四条独立的 tracing 路径，而是建立在 vLLM Piecewise Compilation 之上的运行时优化：编译阶段先拆分并编译 graph regions，随后只对兼容区域执行 CUDA Graph capture/replay。 这篇文章的核心问题是：vLLM 如何围绕标准 torch.compile 增加一层 LLM-serving-aware 的 graph processing、compilation policy 和 runtime integration。 ...

核心结论 Piecewise CUDA Graph 的核心思想是：不要求整个 Transformer forward 都满足 CUDA Graph capture 条件，而是将 CUDA Graph 不兼容的算子作为边界，只 capture 其余 CUDA-Graph-safe 区域。 在 vLLM 中，最典型的边界是 Attention。因此，PIECEWISE 模式并不是“把一个 CUDA Graph 切开执行”，而是将模型划分为多个独立的 CUDA Graph，并在它们之间以 eager 模式执行 CUDA-Graph-unsafe 操作： CUDA Graph -> Attention (eager) -> CUDA Graph -> Attention (eager) -> CUDA Graph CUDA Graph 会提前 capture 一组 GPU operations，后续通过 graph replay 一次性提交，从而降低 CPU 逐个启动 kernel 的开销。Transformer 中的 RMSNorm、Linear/GEMM、activation、FFN 和 residual 等计算，在 capture size 确定时通常具有稳定的执行结构；Attention 则需要处理 KV Cache、动态 shape 和运行时 metadata，对 CUDA Graph 的兼容要求更高。 ...

在和 GGML 打交道时不知道已经写了几篇文档了，这篇一定是最后一个 这篇文档主要记录我用 GGML 实现一个简单的 LLM Inference Engine 时遇到的问题 C++ Initialization 在 C++ 中， class 中的成员在构造函数体开始之前就被默认构造了，要控制其行为，需要使用初始化列表 class Member { public: Member() { std::cout << "Member 默认构造\n"; } Member(int x) { std::cout << "Member 带参构造: " << x << "\n"; } }; class MyClass { Member m1; // 成员对象 Member m2; int value; public: // 情况1: 不使用初始化列表 MyClass() { std::cout << "构造函数体开始\n"; value = 10; // 这是赋值，不是初始化！ } // 情况2: 使用初始化列表 MyClass(int v) : m1(1), m2(2), value(v) { std::cout << "构造函数体开始\n"; } }; 情况 1 输出 Member 默认构造 // m1 在构造函数体前被默认构造 Member 默认构造 // m2 在构造函数体前被默认构造 构造函数体开始 情况 2 输出 Member 带参构造: 1 // m1 在构造函数体前初始化 Member 带参构造: 2 // m2 在构造函数体前初始化 构造函数体开始 对于指针来说，初始化就是将其设置为 nullptr, 对于 STL 容器，就是初始化空容器 ...

向量内积 物理意义 内积本质上是一个向量在另一个向量方向上的投影长度，与基准向量长度的乘积。 如果提取特征或求取分量，令 $\mathbf{u}$ 为单位向量（$\|\mathbf{u}\| = 1$），则 $\mathbf{v} \cdot \mathbf{u}$ 直接输出 $\mathbf{v}$ 在 $\mathbf{u}$ 方向上的标量投影。 工程应用： 在信号处理（如傅里叶变换）中，信号与正交基函数的内积，就是在提取该信号在特定频率上的能量分量。在经典力学中，功的计算 $W = \mathbf{F} \cdot \mathbf{d}$ 就是提取力在位移方向上的有效分量并相乘。 由于内积公式中包含 $\cos(\theta)$，它是衡量高维空间中两个向量方向“一致性”或“对齐程度”的线性算子。 当 $\mathbf{a} \cdot \mathbf{b} > 0$ 时，夹角为锐角，两者存在正相关性。 当 $\mathbf{a} \cdot \mathbf{b} = 0$ 时，$\cos(\theta) = 0$，两向量正交（垂直）。在工程上，这意味着两个系统、信号或特征完全独立，互不干涉（即协方差为零）。 当 $\mathbf{a} \cdot \mathbf{b} < 0$ 时，夹角为钝角，存在负相关性。 工程应用： 在机器学习和数据挖掘中，将向量归一化后求内积，即为余弦相似度（Cosine Similarity），常用于衡量文本词向量的语义相似性或推荐系统中用户偏好的匹配度。 Note 向量内积天然满足交换律，即 $\mathbf{a} \cdot \mathbf{b} = \mathbf{b} \cdot \mathbf{a}$，即 $\mathbf{a}$ 投影在 $\mathbf{b}$ 上和 $\mathbf{b}$ 投影在 $\mathbf{a}$ 上的数值是相等的 ...

Notes on L1 and L2 norms, KL divergence, and their applications in machine learning.

记录一些我在使用 cli 时简单整理的一些东西 Setup Docker curl -fsSL https://get.docker.com -o get-docker.sh sudo DOWNLOAD_URL=https://mirrors.ustc.edu.cn/docker-ce sh get-docker.sh sudo groupadd docker # optional sudo usermod -aG docker $USER newgrp docker sudo systemctl start docker apt 替换镜像源： For 24.04: sed -i "s@http://.*archive.ubuntu.com@https://mirrors.aliyun.com/@g" /etc/apt/sources.list.d/ubuntu.sources sed -i "s@http://.*security.ubuntu.com@https://mirrors.aliyun.com/@g" /etc/apt/sources.list.d/ubuntu.sources sed -i "s@http://ports.ubuntu.com@https://mirrors.aliyun.com@g" /etc/apt/sources.list.d/ubuntu.sources For 22.04: sed -i "s@http://.*archive.ubuntu.com@https://mirrors.aliyun.com/@g" /etc/apt/sources.list sed -i "s@http://.*security.ubuntu.com@https://mirrors.aliyun.com/@g" /etc/apt/sources.list Network # 查看网关 ip route # 查看 ip 地址 ip addr # 查看每个进程的网速 sudo apt install nethogs nethogs apt sudo 会为了安全重置环境变量, 加上 -E 参数来保留当前的环境变量 ...

Transformer 是基于全局视角处理序列的，它通过 Self-Attention 一次性让所有 token 互相交互。而 RNN（Recurrent Neural Network，循环神经网络）的本质是基于时间步的顺序迭代。 如果把 Transformer 比作同时看到整段句子的“上帝视角”，RNN 则像是按照从左到右的顺序逐字阅读，并在脑海中维护一个不断更新的“记忆向量”。 RNN 的核心结构是一个循环单元。它在处理序列时，并不是一次性输入整个序列（如 [batch, seq_len, dim]），而是将序列沿着 seq_len 维度切开，在每一个时间步 $t$ 输入一个 token。 为了保留历史上下文，RNN 引入了隐藏状态 $h_t$。 $h_t$ 是一个固定维度的向量，它包含了从时间步 $0$ 到 $t$ 的所有历史信息压缩。 在每一个时间步，RNN 接收两个输入：当前时刻的输入 $x_t$ 和 上一时刻的隐藏状态 $h_{t-1}$。 RNN 会使用同一套权重矩阵（参数共享）来融合这两个输入，生成新的 $h_t$。 对于序列中的第 $t$ 个 token，RNN 内部只做两步基本的线性变换和一次非线性激活： 更新隐藏状态： $$h_t = \tanh(W_{xh} x_t + W_{hh} h_{t-1} + b_h)$$这里 $W_{xh}$ 负责投影当前输入，$W_{hh}$ 负责投影历史记忆。两者的结果相加后，通过 $\tanh$ 激活函数将数值压缩到 $[-1, 1]$ 之间，防止在长序列迭代中数值爆炸。 计算当前输出（可选）： $$y_t = W_{hy} h_t + b_y$$如果需要每个时间步都输出（比如序列标注），就用当前的 $h_t$ 映射出 $y_t$。如果只需要句向量（比如文本分类），则直接取最后一个时间步的 $h_n$ 即可。 ...

Vanilla RNN 的缺陷：为什么我们需要 LSTM？ 在上一篇中我们提到，RNN 的核心公式是 $h_t = \tanh(W_{xh} x_t + W_{hh} h_{t-1} + b_h)$。这种设计在处理长序列时会遇到两个工程上的致命问题： 梯度消失与梯度爆炸（Vanishing/Exploding Gradients） 在反向传播计算梯度时，误差需要沿着时间步反向传递。这会导致权重矩阵 $W_{hh}$ 被连乘 $t$ 次。根据线性代数原理，如果 $W_{hh}$ 的最大特征值小于 1，连乘后梯度会呈指数级衰减趋近于 0（梯度消失）；如果大于 1，则会指数级放大（梯度爆炸）。梯度消失意味着网络根本无法学习到长距离的依赖关系。 信息覆盖（Information Overwrite） RNN 只有一个隐藏状态 $h_t$。在每一个时间步，新的输入 $x_t$ 都会强制与历史信息 $h_{t-1}$ 混合。没有任何机制能够保护早期非常重要但最近没有出现的信息。这就好比一个容量有限的栈，新数据不断涌入，旧数据很快就被冲刷掉了。 LSTM 的核心思想：分离状态与引入门控机制 为了解决上述问题，LSTM 对架构进行了大改，其核心创新在于：将内部状态拆分为两个，并引入了“门（Gates）”来进行精确的信息路由。 细胞状态 $c_t$ (Cell State)： 这是 LSTM 的“主干道”或“长期记忆”。它在整个链条上贯穿运行，只有一些少量的线性交互。这种设计使得梯度可以通过 $c_t$ 顺畅地无损反向传播，直接解决了梯度消失问题。 隐藏状态 $h_t$ (Hidden State)： 类似于 Vanilla RNN 的 $h_t$，作为“短期记忆”或当前时间步的输出。 门控机制 (Gating Mechanism)： 门本质上是经过 Sigmoid 激活的全连接层。Sigmoid 的输出在 $[0, 1]$ 之间，用于控制信息的保留比例（0 代表完全丢弃，1 代表完全保留）。 数学工作流：LSTM 的四个核心步骤 对于第 $t$ 个 token，LSTM 内部执行以下运算。为了方便，我们通常将前一个隐藏状态 $h_{t-1}$ 和当前输入 $x_t$ 拼接在一起计算。 ...

要理解 Mamba，我们需要将思维从 Transformer 的“全局注意力”切回到 RNN 的“时序状态机”，但这次，我们要解决 Vanilla RNN 和 LSTM 共同的工程死穴：训练阶段无法并行。 Mamba 的核心目标非常明确：既要拥有 Transformer 级别的高效并行训练能力，又要保持 RNN $O(1)$ 的推理复杂度和无限上下文潜力。 它主要通过整合状态空间模型（State Space Model, SSM）、**选择性机制（Selective Mechanism）和底层硬件优化（Hardware-aware Algorithm）**来实现这一目标。 以下从数学和工程的视角详细拆解 Mamba 架构。 数学基础：从连续到离散的状态空间模型 (SSM) Mamba 的前身是 S4 等状态空间模型。SSM 的思想源于控制论，它假设系统的状态演化可以用一组连续的微分方程来描述： $$h'(t) = A h(t) + B x(t)$$$$y(t) = C h(t)$$ $x(t)$ 是输入。 $h(t)$ 是隐藏状态（类比 RNN 的 $h_t$ 或 LSTM 的 $c_t$）。 $A$ 是状态转移矩阵，描述系统本身的演化规律。 $B$ 和 $C$ 是输入输出投影矩阵。 为了在深度学习中使用（处理离散的 token），必须对连续系统进行离散化（Discretization）。通过引入一个步长参数 $\Delta$（Delta），使用零阶保持（Zero-Order Hold）等数学技巧，上述方程可以转换为离散的递归形式： $$h_t = \bar{A} h_{t-1} + \bar{B} x_t$$$$y_t = C h_t$$其中 $\bar{A} = \exp(\Delta A)$。 你可以看到，离散化后的公式与 Vanilla RNN 非常相似：当前状态 $h_t$ 是前一个状态 $h_{t-1}$ 的线性变换加上当前输入 $x_t$ 的线性变换。因为它是纯线性的（没有 LSTM 那样的 $\tanh$ 或 Sigmoid 阻断），这为后续的数学化简和并行计算奠定了基础。 ...

linear_attention 是把 Mamba2 的 gating 机制 和 DeltaNet 的 delta rule 结合起来形成的新结构 先简单分析一下 线性注意力是什么 Overview 从 Transformer 视角 标准 causal self-attention 的一层，大致是： $$ Q = XW_Q,\quad K = XW_K,\quad V = XW_V $$$$ A = \text{softmax}(QK^\top + \text{mask}),\quad O = AV $$对第 $t$ 个 token 来说，本质上是在做： $$ o_t = \sum_{i \le t} \alpha_{t,i} v_i $$ O(n, t) 中的每一行是把 attn_map 中的对应行作为权重，对 value 矩阵的行进行线性组合(加权) 也就是：当前 token 用 query 去和历史所有 key 打分，再把历史 value 加权求和。这给了 Transformer 很强的“按内容检索历史”的能力，但代价是训练时通常要处理完整的 token-token 交互矩阵。Mamba2、DeltaNet、Gated DeltaNet 这一类工作，核心目标就是：尽量保留这种“从历史取信息”的能力，但不要每次都真的显式看全部历史 token。 ...

Kernel Paradigm triton 的 kernel 有两种写法: 传统 kernel Persistent Kernel 传统 Kernel grid = (num_tiles,) 每个 block 处理一个 tile → 结束 每个 SM 从 grid 里领一个 work item，做完就退出, CTA 调度由硬件自动完成 Persistent Kernel @triton.jit def kernel(..., n_tiles: tl.constexpr, ...): pid = tl.program_id(0) n_progs = tl.num_programs(0) # grid(0) 启动的 program 数 tile = pid while tile < n_tiles: # 处理 tile 对应的那一块工作 # ... tile += n_progs # 跳到下一个属于自己的 tile 每个 CTA 长期驻留在 SM 上，自己软件调度 work persistent kernel 常见动机： 任务总 tile 数不够多（比如小 batch、小矩阵、很多小 GEMM / MoE / grouped GEMM）:普通写法 programs 数 < SM，GPU 吃不满；persistent 让每个 SM 都有活干，并通过循环把剩余工作吃完。 ...