DeepSeek MLA 在 DCP 分布式环境中的 Prefill 阶段解析

前言 DeepSeek 提出的 MLA(Multi-head Latent Attention)通过将 KV 压缩到低维 latent 空间,大幅降低了推理时的 KV cache 开销。但在 DCP(Decode Context Parallel,即上下文并行)分布式环境下,MLA 的 prefill 阶段设计与 decode 阶段有显著差异。本文从实现角度展开分析。 什么是 DCP DCP(Decode Context Parallel)是一种将 KV cache 按序列维度切分到多个 GPU 的分布式策略。每个 rank 只持有完整 KV cache 的 1/dcp_world_size,从而减少单卡显存占用,支持更长的上下文。 与更常见的 DP(Data Parallel,扛并发)和 EP(Expert Parallel,分摊 MoE 参数显存)不同,DCP 解决的是单请求长上下文场景下 KV cache 放不下的问题。 MLA Prefill vs Decode:两条不同的路径 MLA 在 prefill 和 decode 阶段走了截然不同的计算路径: Prefill (forward_mha) Decode (forward_mqa) KV 形态 完整 MHA(N 头) Latent(1 头) Head dim P+R(~192) Lkv+R(~576) 计算特性 Sq ≈ Skv,计算密集 Sq ≪ Skv,避免显存搬运 Prefill 走 MHA 路径:kv_c 通过 W_UK/W_UV 解压成完整多头 K/V(N 个头),然后做标准的多头注意力。因为 prefill 时新 token 数和 context 长度在同一量级,展开 KV 做计算密集的 attention 是划算的。 ...

June 1, 2026 · 3 min · 557 words

vLLM GPU UBatchWrapper 中的 Barrier 机制解析

背景:什么是 UBatchWrapper 在 vLLM 的 GPU 推理引擎中,UBatchWrapper(位于 vllm/v1/worker/gpu_ubatch_wrapper.py)是一个模型包装器。它拦截对原始模型的调用,在内部将一次大 batch 的 forward 拆分成多个 micro-batch(ubatch),并利用多线程 + CUDA stream 实现计算与通信的重叠。 核心目标:解决 MoE 模型中 all2all 通信开销大的问题——让一个 ubatch 在做计算时,另一个 ubatch 的通信已经在背后并行执行。 线程模型与同步原语 threading.Barrier:集结号 UBatchWrapper 在初始化时创建一个 threading.Barrier: self.ready_barrier = threading.Barrier(num_ubatches + 1) threading.Barrier 不关心具体有哪些线程,只计数。内部维护一个计数器,每当有线程调用 barrier.wait(),计数器就 +1。当总数达到构造时指定的 parties 值(这里是 num_ubatches + 1),所有正在 wait() 的线程同时被释放,计数器归零。 这里的 +1 包含了 N 个 ubatch 线程 + 1 个主线程。 ready_barrier.wait() 在三处被调用: 文件 行号 调用方 gpu_ubatch_wrapper.py 261 主线程 在 _capture_ubatches 中 gpu_ubatch_wrapper.py 325 主线程 在 _run_ubatches 中 ubatching.py 56 每个 ubatch 线程 在 UBatchContext.__enter__ 中 少一个线程到达,所有人全部卡住。 ...

June 1, 2026 · 4 min · 776 words

mhc_pre_torch 的数学公式与 CUDA 代码解释

背景 mHC(multi-Head Combinatorial)是 DeepSeek V4 模型中引入的一种多头残差混合机制。它将传统 Transformer 中单一的残差向量扩展为 $M$ 个并行的残差副本(multi-head residual),并在每个 block 前后通过可学习的门控和组合矩阵对多头残差进行变换。 本文以 mhc_pre_torch 为核心,从数学公式出发逐行对照 PyTorch 代码,并延伸至 mhc_post_torch 和整个 mHC 流水线,帮助读者完整理解这一算子的设计思路与实现细节。 ...

June 1, 2026 · 5 min · 993 words

vLLM Cascade Attention:用法、场景与原理解析

概述 Cascade Attention(级联注意力)是 vLLM 推理引擎中针对多请求共享长前缀场景的一种注意力优化技术。它将标准 attention 拆解为 prefix(前缀)和 suffix(后缀)两个阶段,显著降低 KV cache 的全局内存读取量,在共享 system prompt 的批量推理场景中可实现最高数十倍的 attention 加速。 本文将深入剖析 cascade attention 的设计思想、使用方式、适用场景、实现原理,并与 FlashAttention、FlashDecoding 等其他优化技术进行对比。 一、Cascade Attention 的核心思想 标准 Attention 的冗余 在大语言模型的批量推理中,多个请求往往共享一个较长的 system prompt。例如: Chatbot 场景:所有请求共享 “You are a helpful assistant…” 等 system prompt Document QA:多用户对同一篇文档提问,文档内容为公共前缀 Self-Consistency:对同一 prompt 采样多条推理路径 标准 attention 的处理方式是每个 request 独立计算其完整的注意力——包括共享的 system prompt。这意味着同一份前缀 KV cache 被重复加载多次。 Cascade 的拆解思路 Cascade attention 将一次 attention 计算拆成三步: Prefix 阶段:将所有请求的 query 拼成一个"大序列",对共享前缀做一次非因果(bidirectional)attention Suffix 阶段:每个请求各自对其独有的后缀做因果(causal)attention Merge:通过 LSE(log-sum-exp)rescaling 将两阶段结果加权合并 数学上等价于标准 attention,但计算量和显存带宽需求大幅降低。 ...

June 1, 2026 · 6 min · 1143 words

vLLM TorchDynamo 编译:FX Graph 分割原理与实践

背景 vLLM 利用 PyTorch 2.x 的 torch.compile 路径,通过 TorchDynamo 捕获模型计算图,再经过图分割(graph splitting)和分段编译(piecewise compilation)来优化 GPU kernel 执行效率。本文将深入剖析 FX Graph 的分割原理及其背后的设计思想。 一、TorchDynamo 编译概览 TorchDynamo 是一个 Python 级别的 JIT 编译器,它通过 PEP 523 的 frame evaluation callback 在 Python 字节码执行之前捕获计算图。vLLM 利用这一机制,将模型 forward 函数中的计算捕获为 fx.GraphModule,然后送入自定义后端 VllmBackend 进行编译。 编译流程大致如下: model.forward() └── TorchDynamo 捕获计算图 └── fx.GraphModule (原始完整图) └── VllmBackend.__call__() ├── split_graph() → 图分割 ├── PiecewiseCompileInterpreter → 分段编译 └── codegen → 生成胶水函数 └── 返回可调用对象给 Dynamo 二、FX Graph 分割:split_graph 的工作原理 2.1 为什么需要分割 vLLM 的模型计算中包含两种性质不同的算子: ...

June 1, 2026 · 4 min · 754 words

DeepSeek Attention 中为什么除以 √d_k — 从代码到数学

引言 在 Transformer 的 Scaled Dot-Product Attention 中,有一个看似简单却至关重要的操作: scores = Q @ K^T / math.sqrt(d_k) 这个 ÷√d_k 几乎出现在每一个 Attention 实现中。DeepSeek 的 MLA(Multi-head Latent Attention)也不例外。本文从 DeepSeek V4 的 vLLM 实现代码出发,深入探讨这个缩放因子的数学动机、工程实现,以及 DeepSeek 特有的优化变体。 标准 Attention 中的缩放因子 Softmax 的"饱和"问题 原始 Attention 计算公式为: Attention(Q, K, V) = softmax(QK^T / √d_k) V 设 Q、K 的每个元素是独立同分布、均值为 0、方差为 1 的随机变量。对于向量 q 和 k(维度为 d_k),其点积的均值和方差为: E[q·k] = Σ E[q_i·k_i] = 0 Var(q·k) = Σ Var(q_i·k_i) = d_k 即点积的方差正比于 d_k。当 d_k 较大时(如 128、512),点积的绝对值会很大,导致 softmax 进入梯度极小的饱和区: 当某个 x_i 远大于其他值时,e^{x_i} 会主导分母,softmax 输出接近 one-hot,梯度趋近于 0,训练难以收敛。 ...

June 1, 2026 · 4 min · 721 words

DeepSeek MLA 中 QKV Head Dimension 的处理差异

引言 Multi-head Latent Attention(MLA)是 DeepSeek-V2/V3 系列模型中最核心的架构创新之一。它在标准 Multi-Head Attention(MHA)的基础上引入了低秩压缩,大幅降低了 KV cache 的显存占用,同时保持了与 MHA 相当的模型质量。 本文从 QKV head dimension 处理差异 这一视角切入,深入分析 MLA 的设计原理、vLLM 中的具体实现,以及这种设计带来的性能收益。 1. 标准 MHA 回顾 在标准 MHA 中,给定输入序列 X ∈ ℝ^{S×D},Q、K、V 通过三个独立的线性变换得到: Q = X @ W_Q → [S, H, d_head] K = X @ W_K → [S, H, d_head] V = X @ W_V → [S, H, d_head] 其中 d_head = D / H,三者完全相等。注意力计算为: Attention(Q, K, V) = softmax(Q @ K^T / √d_head) @ V → [S, H, d_head] 输出再经 W_O 映射回 D 维。这种 Q/K/V 共享同一 head dim 的设计深入人心,以至于很多人默认这是注意力机制的"必须要求"。 ...

June 1, 2026 · 5 min · 909 words

DeepSeek V3 MoE 模块计算与通信逻辑详解

概述 DeepSeek V3 的 MoE(Mixture of Experts)模块是其核心组成部分,采用 Shared + Routed Expert 架构。本文基于 vLLM 代码库,深入分析其计算流程、通信模式、量化方案以及性能优化策略。 ...

June 1, 2026 · 6 min · 1086 words

vLLM TurboQuant 介绍:极低比特 KV Cache 量化

引言 大语言模型(LLM)推理时,KV Cache 是显存消耗的主要来源之一。以 LLaMA-70B 为例,在 4096 上下文长度下,KV Cache 占用超过 40 GB 显存。传统的 FP16 KV Cache 每个 token 需要 2 × head_dim × num_heads × 2 bytes = 大量显存。如果能将 KV Cache 压缩到 3-4 bit,就可以将显存占用降低到原来的 1/4 到 1/5,从而支持更长的上下文、更大的 batch size,或在更少的 GPU 上运行同一模型。 TurboQuant(Zandieh et al., ICLR 2026)正是为此而生——它是 vLLM 中实现的一种极低比特 KV Cache 量化方法,通过 Hadamard 旋转 + Lloyd-Max 最优标量量化 + Norm Correction 的组合,将 KV Cache 压缩至 3-4 bit,最高实现 4.9× 压缩比。 为什么需要 TurboQuant? 在 TurboQuant 之前,vLLM 已经支持了 FP8 KV Cache(2× 压缩,几乎无损)。但 2× 压缩在某些场景下仍然不够: ...

June 1, 2026 · 5 min · 884 words

vLLM 量化方法详解:原理、算法与实现

本文档聚焦于 vLLM 中与 NVIDIA GPU 兼容的通用量化方法,排除了框架特定方案(GGUF、Compressed-Tensors、Quark、Hummingbird)和硬件厂商专有方案(ModelOpt、INC、NVFP4)。涵盖:数学原理、量化/反量化公式、位宽与分组结构、核心算法步骤、Kernel 级实现细节及硬件要求。 目录 FP8 W8A8(动态/静态量化) AWQ(激活感知权重量化) GPTQ / Marlin AWQ Marlin BitsAndBytes(NF4 / QLoRA) MXFP4(OCP MX 规范) Online Quantization(在线量化) TurboQuant FP8 / INT8 KV Cache 对比总结 1. FP8 W8A8(动态/静态量化) 论文 FP8 格式标准: FP8 Formats for Deep Learning (Micikevicius et al., 2022, arXiv:2209.05433) vLLM 实现: 主要面向 H100/H200 GPU,利用 torch._scaled_mm 调用 Tensor Core。 概述 FP8 W8A8 将权重和激活均量化为 FP8 格式。支持两种调度模式: 静态量化:校准集上预计算 scale,推理时固定。 动态量化:推理时实时统计最大值计算 scale。 量化公式 FP8 有两种指数位宽子格式: 格式 符号位 指数位 尾数位 最大值 最小值 E4M3 1 4 3 448 1/16 E5M2 1 5 2 57344 1/64 量化核心操作: ...

June 1, 2026 · 11 min · 2171 words