引言
在 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，训练难以收敛。
缩放因子的推导
为了控制方差，在点积后除以 √d_k：
Var(q·k / √d_k) = (1/√d_k)² · Var(q·k) = d_k / d_k = 1
缩放后点积的方差稳定在 1，与 d_k 无关。softmax 的输入分布不再随维度增长而发散，梯度流动得以保持。
√d_k（而非 d_k）的选择是因为方差缩放是平方关系：Var(aX) = a²·Var(X)，要消除 d_k 的因子，令 a = 1/√d_k。
DeepSeek 中的 Attention 实现
MLA（Multi-head Latent Attention）架构
DeepSeek 使用 MLA 压缩 KV cache。核心思路是先将 KV 投影到低维 latent 空间再存，需要时再展开。代码中 nope_head_dim 是 NoPE（无位置编码）部分的维度：
# deepseek_v4_attention.py:L220
self.nope_head_dim = 448 # 7B scale factors
结合 RoPE 维度（64），总 head_dim 为 512：
constexpr int kHeadDim = 512;
constexpr int kRopeDim = 64;
constexpr int kNopeDim = kHeadDim - kRopeDim; // 448
Per-Head RMSNorm — 另一种"缩放"
DeepSeek 还在 Q 上做了 per-head RMSNorm（无 weight），作为一个隐式的缩放层：
# fused CUDA kernel: per-head RMSNorm
# 对每个 head 的向量做：
# sum(x²) → rsqrt → 乘回去
Triton 实现中 block_size 取 next_power_of_2(max(q_size, kv_size))：
block_size = triton.next_power_of_2(max(q_size, kv_size))
# q_size=1536, kv_size=128 → block_size=2048
这确保一个 Triton program 能完整处理 RMSNorm 的一行，通过 mask = block < SIZE 安全处理 padding。
Softmax Scale 在 Indexer 中的使用
DeepSeek 的 Lightning Indexer 中也显式使用了缩放因子：
fused_indexer_q_rope_quant(
indexer_weights, self.softmax_scale, self.n_head ** -0.5, ...
)
这里的 n_head ** -0.5 是另一种缩放形式，用于平衡 Indexer 的 per-head learned weights。
FP8 量化中的缩放——除以 64
这个问题的直接来源其实不是 softmax 缩放，而是 FP8 sparse KV cache 中 scale factor 的存储计算：
# deepseek_v4_attention.py:L222
self.nope_head_dim // 64 # 7B scale factors
nope_head_dim = 448，448 // 64 = 7。这意味着每 64 个 FP8 元素共享一个 scale factor（1 byte FP8 格式），448 个元素需要 7 个 scale factor。
CUDA kernel 中的量化逻辑：
constexpr int kQuantBlock = 64;
constexpr int kNumQuantBlocks = kNopeDim / kQuantBlock; // 7
constexpr int kScaleBytesPerToken = kNumQuantBlocks + 1; // 8 (7 real + 1 pad)
每个 quant block 内，4 个连续 lane 为一组，shuffle 计算 absmax：
absmax → exponent = ceil(log2(absmax / 448))
inv_scale = 2^(-exponent)
clamp 到 [-448, 448] → __nv_cvt_float_to_fp8 (E4M3)
这里的 ÷448 是量化到 UE8M0 格式所需的归一化因子，不是一个可学习的参数，而是固定的数值映射。
为什么 block_size = 64？
对于 448 维的 NoPE 向量，64 的 block_size 给出了整数个 block（7 个）。这与标准的 128 block size 不同（给出 3.5 个 block）。64 提供了更细粒度的量化，量化误差更小，代价是每个 token 多几个字节的 scale factor 存储。
注释中提到了另一种方案：
# 7B scale factors // block_size 128 * 4
即如果用 block_size=128 且每个 scale 占 4 bytes（FP32），需要 ceil(448/128) * 4 = 16 bytes。当前方案 block_size=64、scale 仅 1 byte，只需 7 bytes。
vLLM 中的实现细节
Paged Cache 布局
DeepSeek V4 的 SWA（Sliding Window Attention）cache 以 uint8 tensor 存储，形状为 (num_blocks, block_size, 584)：
每个 token 584 bytes：
- 448 bytes FP8 NoPE 数据
- 128 bytes bf16 RoPE 数据（64 × 2）
- 8 bytes scale factors（7 real + 1 pad）
每个 block 有 64 个 token，总计 64 × 584 = 37376 bytes。
Fused Kernel 架构
vLLM 使用水平融合的 CUDA kernel 一次性完成 Q 的 per-head RMSNorm、RoPE，以及 KV 的 RoPE、FP8 量化、cache 写入：
Grid: 1D，每个 block 8 个 warp（256 threads）
全局 warp idx = blockIdx.x * 8 + warpId
每个 token 有 num_heads_q + 1 个 slot：
< num_heads_q → Q 分支
== num_heads_q → KV 分支
tokenIdx = globalWarpIdx / (num_heads_q + 1)
slotIdx = globalWarpIdx % (num_heads_q + 1)
每个 lane 处理 16 个元素（512 bits / 32 lanes = 16），通过 uint4 向量化加载/存储。
压缩 KV Cache 的 Slot Mapping
DeepSeek 的 MLA 有时域压缩：每 compress_ratio 个 token 共享一个 cache slot。对应的 Triton kernel 计算 slot mapping：
is_valid = (pos + 1) % COMPRESS_RATIO == 0 # 每组最后一位才有效
pos_after_compress = pos // COMPRESS_RATIO # 压缩后的位置
slot_id = block_number * block_size + pos_after_compress % block_size
只有压缩组边界的 token（(pos+1) % ratio == 0）才写入 paged cache，中间的 token 只更新 compressor 的 running state。
Top-K 索引的全局映射
稀疏 attention 选出 top-k 局部索引后，通过 block table 转换为全局 slot ID：
block_indices = local_idx // block_size
block_numbers = block_table[req_idx, block_indices]
block_offsets = local_idx % block_size
slot_ids = block_numbers * block_size + block_offsets
无效条目（pad）标记为 -1，下游 kernel 通过 slot_id < 0 跳过。
DeepSeek 特有的缩放变体
总结 DeepSeek V4 中出现的所有"缩放"操作：
| 缩放操作 | 位置 | 目的 |
|---|---|---|
QK^T / √d_k | 标准 Attention | 控制点积方差 |
nope_head_dim // 64 | Cache 分配 | FP8 量化 block 数 |
absmax / 448 | FP8 量化核 | 量化归一化 |
per-head RMSNorm | Q 预处理 | 隐式特征缩放 |
n_head ** -0.5 | Indexer | 平衡 learned weights |
compress_ratio | MLA Cache | 时域压缩比率 |
每种缩放的动机不同：有的为了数值稳定性（√d_k），有的为了存储效率（// 64、compress_ratio），有的为了量化精度（absmax / 448）。
总结
÷√d_k 是 Attention 中最基础也最重要的缩放操作，它保证了点积方差与维度无关，使 softmax 训练稳定。在 DeepSeek 的实践中，缩放以多种形式出现——从标准的 softmax 缩放，到 FP8 量化中的每 64 元素分组缩放，再到 MLA 的时域压缩缩放。
理解每种缩放背后的数学和工程动机，才能真正掌握现代 Attention 机制的实现精髓。DeepSeek 的代码告诉我们：缩放不仅是一个公式，更是一个贯穿数值计算、存储布局和硬件利用的系统性工程考量。