Recent 📎Posts
IndexCache:DeepSeek 稀疏注意力中的跨层索引复用

技术背景#
MLA 已经解决了 KV Cache 存储问题#
DeepSeek-V2 提出的 MLA(Multi-head Latent Attention) 将 KV Cache 压缩到了极低的水平:
- 每 token 仅缓存 576 个元素(512 维压缩隐向量 + 64 维 RoPE 分量)
- 相当于 MHA 的 1.76%
- 128K 上下文下,FP8 KV Cache 仅需 8.72 GB(对比 Llama-70B 的 343 GB)
MLA 的压缩技巧在于:K 和 V 共享一个低秩隐向量 $c_t^{KV} \in \mathbb{R}^{512}$,推理时通过矩阵吸收技巧避免显式恢复高维 K/V。但是,MLA 仍然是全注意力——每个 query token 仍然要计算与所有历史 token 的注意力分数。当上下文长度增长到 128K 以上时,注意力计算量本身成为新瓶颈。
从全注意力到稀疏注意力#
DeepSeek-V3.2 引入的 DSA(DeepSeek Sparse Attention) 直接解决计算瓶颈:不再计算所有 token 的注意力,而是只选择 Top-K 个最重要的历史 token。闪电索引器 (Lightning Indexer) 是一个轻量级的注意力模块,为每个 query 快速计算与所有 key 的粗粒度相关性分数,然后通过 Top-K Selector 选出 top-k 个最相关的 token 索引。
DSA(Deepseek Sparse Attention)
稀疏注意力#
复杂度瓶颈#
Transformer 架构的核心是缩放点积注意力 (Scaled Dot-Product Attention):
$$o_i = \sum_{j=1}^{n} \alpha_{ij} \cdot v_j, \quad \alpha_{ij} = \operatorname{softmax}\left(\frac{q_i^T k_j}{\sqrt{d_k}}\right)_j$$其中 $Q, K, V \in \mathbb{R}^{n \times d}$ 分别是查询、键、值矩阵。
$QK^T$ 矩阵的形状为 $n \times n$,计算复杂度 $O(n^2 d)$,空间复杂度 $O(n^2)$。对于长序列($n=128\text{K}$,$d=4096$,FP16),仅注意力矩阵就需要 $128\text{K} \times 128\text{K} \times 2\text{B} \approx 32\text{ GB}$ 显存——单张 GPU 无法承受。
注意力矩阵的稀疏性来源于一个直觉观察:并非所有 token 对都同等重要。大多数 token 对之间的注意力权重接近零,仅有少数 token 对承载了主要的信息交互。
这种稀疏性有两层含义:
- 自然的局部性:相邻 token 之间的依赖远强于远距离 token
- 内容相关性:语义相关的 token 即使距离远也需要交互
稀疏注意力方法试图在信息完整性和计算效率之间寻找最优平衡。

MLA(Multi-Head Latent Attention)

MHA(多头注意力)#
$$Attention(Q, K, V) = softmax(Q·K^T / \sqrt{d_h}) · V$$$$q_t = W^Q h_t$$
$$k_t = W^K h_t$$
$$v_t = W^V h_t$$
- $d$: Embedding维度;
- $d_h$:每一个头的维度;
- $n_h$: Attention头数;
- $h_t ∈ ℝ^{d}$: Attention单层中第$t$ 个token的输入;
- $W^Q, W^K, W^V ∈ ℝ^{d_hn_h × d}$ ;
- $q_t,k_t,v_t ∈ ℝ^{d_hn_h}$ ; $$[q_{𝑡,1};q_{𝑡,2}; ...; q_{𝑡,𝑛_ℎ} ] = q_t$$ $$[k_{𝑡,1};k_{𝑡,2}; ...; k_{𝑡,𝑛_ℎ} ] = k_t$$ $$[v_{𝑡,1};v_{𝑡,2}; ...; v_{𝑡,𝑛_ℎ} ] = v_t$$ $$o_{t,j} = \sum_{j=1}^t{softmax_j(q_{t,i}^T·k_{j,i} / \sqrt{d_h})} · v_{j,i}$$ $$u_𝑡 = 𝑊^𝑂 [o_{𝑡,1}; o_{𝑡,2}; ...; o_{𝑡,𝑛_ℎ} ]$$
- $q_{t,i},k_{t,i},v_{t,i} ∈ ℝ^{d_h}$ : 第$i$个Attention头;
- $W^O ∈ ℝ^{d × d_hn_h}$ :输出投影矩阵;
- 每个 token 单层需要缓存的 cache 的是 $n_h$ 套 ($k,v$) — 共 $2n_hd_h$
MQA(多查询注意力)#
所有 $n_h$ 个 Q head 共享同一套 K、V,只剩 1 套,更快的推理速度:由于内存和缓存需求大幅下降,内存带宽开销也基本消除。虽然训练效率基本保持不变,但与 MHA 相比,共享 K/V 对可能会导致模型质量和输出精度略有下降
Split QKV + RMSNorm + RoPE 融合算子
源代码:
vllm-ascend/vllm_ascend/ops/triton/linearnorm/split_qkv_rmsnorm_rope.py
背景#
问题:内存墙#
LLM 推理(尤其是 decode 阶段)是典型的 memory-bound 场景。每一次算子调用都是一次「从 Global Memory 搬数据到片上 → 计算 → 搬回 Global Memory」的循环。如果不融合,Split QKV → RMSNorm → RoPE 这三个步骤各自独立执行:
非融合流程:
hidden_states ──[load]──> Split Q,K,V ──[store]──> q_in, k_in, v_in
q_in ──[load]──> RMSNorm ──[store]──> q_normed
k_in ──[load]──> RMSNorm ──[store]──> k_normed
q_normed ──[load]──> RoPE ──[store]──> q_out
k_normed ──[load]──> RoPE ──[store]──> k_out
每一次 load/store 都是一次 GM ↔ UB 的数据搬运。数据在总线上反复进出,但实际的计算量(几个乘加 + 一个开方)极小——这就是 memory-bandwidth bottleneck。
vLLM 分布式通信
class WorkerProc:
"""Wrapper that runs one Worker in a separate process."""
READY_STR = "READY"
rpc_broadcast_mq: MessageQueue | None
worker_response_mq: MessageQueue | None
@instrument(span_name="Worker init")
def __init__(...):
self.rank = rank
wrapper = WorkerWrapperBase(rpc_rank=local_rank, global_rank=rank)
...
wrapper.init_worker(all_kwargs)
self.worker = wrapper
...
self.worker.init_device()
if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH:
self.worker.elastic_ep_execute("load_model")
else:
self.worker.load_model()
。。
Worker init_device#
worker的init_device函数负责初始化设备的相关信息
- 根据当前worker的rank找到它所属的device,将它绑到指定的卡上,以及清空该device的显存,获取该device的显存大小等
- 对当前的worker做分布式环境初始化,也就是初始化当前worker的各种进程组(如模型并行、流水线并行、数据并行等)
- 构造当前worker的GPUModelRunner对象。维护着模型权重分片,还维护模型运行过程中所需要的一些数据结构,比如kv cache等,负责模型权重的加载(load_model),以及实际的推理执行过程等逻辑。
@instrument(span_name="Init device")
def init_device(self):
if self.device_config.device_type == "cuda":
...
# Ray 会设置 NCCL_ASYNC_ERROR_HANDLING,但这个环境变量会导致 CUDA graph
# 构建时出现异常。CUDA graph 需要同步执行,而该变量会引入异步错误处理,
# 两者可能引起冲突。
os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None)
...
# - Ray/external_launcher 场景:这些分布式执行器自己管理GPU映射
# - 多节点场景(nnodes_within_dp > 1):每个节点有独立的GPU集合,映射逻辑不同
# - Ray作为DP后端:Ray的resource pool处理GPU分配
if (
parallel_config.distributed_executor_backend
not in ("ray", "external_launcher")
and parallel_config.data_parallel_backend != "ray"
and parallel_config.nnodes_within_dp == 1 # 单节点场景
):
# local DP rank 表示在当前节点内的数据并行编号,而 global rank 可能
# 跨节点。在单节点场景下,GPU映射只看节点内的local rank。
dp_local_rank = self.parallel_config.data_parallel_rank_local
if dp_local_rank is None:
dp_local_rank = self.parallel_config.data_parallel_index
# DP副本 0 (dp_local_rank=0)
# original_local_rank=0 → GPU 0 + 0×2 = GPU 0
# original_local_rank=0 → GPU 0 + 1×2 = GPU 2 ← 偏移
# 偏移是为了计算出实际的rank信息,后续用来初始化device
self.local_rank += dp_local_rank * tp_pp_world_size
。。
self.device = torch.device(f"cuda:{self.local_rank}")
# PyTorch的设备API演进:从 torch.cuda.set_device() 到
# torch.accelerator.set_device_index()。 后续可以不用再手动指定,
torch.accelerator.set_device_index(self.device)
...
# 初始化分布式推理所需的所有环境,包括通信组信息
# 优先于内存快照处理逻辑,NCCL在初始化会分配内部缓存区,提前初始化用于保证显存计算的准确性
init_worker_distributed_environment(
self.vllm_config,
self.rank,
self.distributed_init_method,
self.local_rank,
current_platform.dist_backend,
)
# 1. gc.collect() 回收Python层的垃圾对象,释放可能的GPU引用
# 2. empty_cache() 释放PyTorch缓存的显存(包括NCCL缓冲区)
# 得到一个"干净"的显存状态,作为基准线
gc.collect()
torch.accelerator.empty_cache()
# 用于后续计算 KV cache 可用显存。
# init_snapshot 记录当前可用显存,request_memory 根据模型配置
# 计算需要预留的 KV cache 大小。
self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device)
self.requested_memory = request_memory(init_snapshot, self.cache_config)
# 最后初始化modelrunner,需要依赖设备、分布式通信环境,同时通过快照可以计算出可以分配的kv cache显存大小
if self.use_v2_model_runner:
...
self.model_runner: GPUModelRunner = GPUModelRunnerV2( # type: ignore
self.vllm_config, self.device
)
else:
...
self.model_runner = GPUModelRunnerV1(self.vllm_config, self.device)
...
init_worker_distributed_environment#
负责初始化分布式推理所需的所有环境组件
TorchInductor Pattern Matcher
PyTorch FX 图#
PyTorch FX 是用于捕获、分析和转换 PyTorch 计算图。FX 图是一种静态表示,它记录了 PyTorch 代码的执行流程。用户通过将模型表示为FX图,可以更轻松地进行各种转换,例如图优化,量化,算子融合等。
FX 图的核心组件包括:
torch.fx.Graph:计算图的容器torch.fx.Node:图中的节点,表示计算操作,如函数调用、方法调用等torch.fx.GraphModule:由图构建的可执行模块
graph TD
subgraph FX_Graph
A["Placeholder Node"] --> B["CallFunction Node"]
B --> C["CallMethod Node"]
C --> D["Output Node"]
E["Module Node"] --> B
end
subgraph Components
F["torch.fx.Graph"] --> FX_Graph
G["torch.fx.Node"] --> A
G --> B
G --> C
G --> D
H["torch.fx.GraphModule"] --> F
end
style FX_Graph stroke:#333,stroke-width:2px
style Components stroke:#333,stroke-width:2px
FX Symbolic Tracing#
FX 图的生成过程称为"符号追踪"(Symbolic Tracing),主要步骤包括:
- 追踪:使用
torch.fx.symbolic_trace()对 PyTorch 函数或模块进行追踪 - 捕获:捕获函数执行过程中的所有操作,构建计算图
- 表示:将计算图表示为
Graph对象,其中包含一系列Node对象 - 转换:对捕获的图进行分析和转换
- 执行:将转换后的图包装为
GraphModule,可像普通 PyTorch 模块一样执行
import torch
# Simple module for demonstration
class MyModule(torch.nn.Module):
def __init__(self) -> None:
super().__init__()
self.param = torch.nn.Parameter(torch.rand(3, 4))
self.linear = torch.nn.Linear(4, 5)
def forward(self, x):
return self.linear(x + self.param).clamp(min=0.0, max=1.0)
module = MyModule()
from torch.fx import symbolic_trace
# Symbolic tracing frontend - captures the semantics of the module
symbolic_traced: torch.fx.GraphModule = symbolic_trace(module)
# High-level intermediate representation (IR) - Graph representation
# 由一个列表组成 代表函数输入、调用点(函数、方法、 或 torch.nn.Module 实例),以及返回值。
print(symbolic_traced.graph)
"""
graph():
%x : [num_users=1] = placeholder[target=x]
%param : [num_users=1] = get_attr[target=param]
%add : [num_users=1] = call_function[target=operator.add](args = (%x, %param), kwargs = {})
%linear : [num_users=1] = call_module[target=linear](args = (%add,), kwargs = {})
%clamp : [num_users=1] = call_method[target=clamp](args = (%linear,), kwargs = {min: 0.0, max: 1.0})
return clamp
"""
# Code generation - valid Python code
# 使 FX 成为 Python 到 Python(或 模块到模块)转换工具包。对于每个 Graph IR,我们可以 创建与图语义匹配的有效 Python 代码。
print(symbolic_traced.code)
"""
def forward(self, x):
param = self.param
add = x + param; x = param = None
linear = self.linear(add); add = None
clamp = linear.clamp(min = 0.0, max = 1.0); linear = None
return clamp
"""
FX 图的特点
vllm DP (Data Parallel)
DP基本概念#
DP在推理场景下的核心思想,在多个 GPU/节点上完整复制同一个模型权重,每个副本独立处理不同的请求或批次,从而近似线性提升吞吐。与训练中的 DP 需要梯度聚合不同,推理 DP 没有参数同步,通信负担主要来自调度、路由、指标与可选的缓存协同。
具体来说
- 每个 GPU/设备都拥有模型的完整副本
- 输入数据被分割成多个批次,通过负载均衡分配给不同设备
- 各设备独立进行前向推理
- 每个设备产生各自批次的输出结果
在DP部署方式下由于单卡的计算效率基本保持不变,因此吞吐提升近似是线性:理论上2 张卡就是 2 倍吞吐,4 张卡就是 4 倍,以此类推。
在大规模部署DP的时候,由于整体可支持的吞吐翻倍,API服务器需要面临成倍的压力,因此API服务器(Tokenize等预处理)可能会成为系统瓶颈。vllm可以使用--api-server-count命令行选项来扩展,最终暴露给用户的是一个Endpoint,在内部实现API服务器的扩展。
vllm 异步调度解析
在vllm初始版本中只有一个同步调度策略,在该策略下GPU资源会在调度过程中形成空泡,造成GPU资源的浪费。vllm在v0.10.0版本后提供异步调度策略,并且在后续的迭代中不断加入对于其他特性(例如异步场景下的投机解码)的支持。原始PR内容可查看#19970 Implement Async Scheduling ,当前代码分析基于main branch(735284ed)。
EngineCore处理处理Step逻辑:
def _process_engine_step(self) -> bool:
"""Called only when there are unfinished local requests."""
# Step the engine core.
outputs, model_executed = self.step_fn()
# Put EngineCoreOutputs into the output queue.
for output in outputs.items() if outputs else ():
self.output_queue.put_nowait(output)
# Post-step hook.
self.post_step(model_executed)
return model_executed
同步调度策略#
def step(self) -> tuple[dict[int, EngineCoreOutputs], bool]:
if not self.scheduler.has_requests():
return {}, False
scheduler_output = self.scheduler.schedule()
# 通过FutureWrapper进行异步包装(复用异步调度的部分逻辑, 在同步调度逻辑里面会等待结果返回)
#
future = self.model_executor.execute_model(scheduler_output, non_block=True)
# 用于支持结构化输出等
grammar_output = self.scheduler.get_grammar_bitmask(scheduler_output)
with self.log_error_detail(scheduler_output):
# 同步
model_output = future.result()
if model_output is None:
model_output = self.model_executor.sample_tokens(grammar_output)
# 处理整个过程中abort的请求
self._process_aborts_queue()
engine_core_outputs = self.scheduler.update_from_output(
scheduler_output, model_output
)
return engine_core_outputs, scheduler_output.total_num_scheduled_tokens > 0
同步调度步骤: