词表维上的并行账：嵌入查表、LM Head 切分与 logits 收集的跨框架拆解

为什么先 pad 再切

词表大小往往不是 TP 度的整数倍。129280 除以 8 恰好整除，但绝大多数模型没这么巧。vLLM 的做法是在构造时把词表补齐到 64 的倍数：

1
DEFAULT_VOCAB_PADDING_SIZE = 64
2

3
def pad_vocab_size(vocab_size: int, pad_to: int = DEFAULT_VOCAB_PADDING_SIZE) -> int:
4
    return ((vocab_size + pad_to - 1) // pad_to) * pad_to

补到 64 的倍数，再除以 TP 度，就能保证每个 shard 大小一致，GEMM 的 N 维也落在对齐边界上。补出来的那几行权重不参与真实 token，但在数学上必须有值——通常是零——否则采样时要额外做一次越界判断。

更复杂的是 LoRA。vLLM 支持运行时挂载 LoRA adapter，而 LoRA 可以带来额外词表（新增 token）。这些 token 的 embedding 不在基础权重里，加载顺序也不可控。vLLM 因此把词表划分成两段：基础词表与追加词表，两段各自 pad、各自切分，然后在同一个张量里首尾相接。源码里那张 ASCII 示意图讲得很细：

1
TP2, rank 0:
2
|< --------------------BASE--------------------- >|<-LORA-->|<-LORA PAD->|
3
|  0 |  1 | ... | 497 | 498 | ... | 511 | 1010 | ... | 1025 |   -1  | ... |

这样做的好处是：无论基础词表怎么 pad，LoRA 新增的 id 永远落在张量末尾，加载器只需要知道两段各自的起止下标，不需要重新排列已有权重。

掩码查表：非所有者的行必须变成零

VocabParallelEmbedding.forward 面对的第一件事，是把全局 id 翻译成”本 rank 的局部 id”，并把不属于本 rank 的 id 屏蔽掉：

1
org_vocab_mask = (input_ >= org_vocab_start_index) & (input_ < org_vocab_end_index)
2
added_vocab_mask = (input_ >= added_vocab_start_index) & (input_ < added_vocab_end_index)
3
added_offset = added_vocab_start_index - (org_vocab_end_index - org_vocab_start_index) - num_org_vocab_padding
4
valid_offset = org_vocab_start_index * org_vocab_mask + added_offset * added_vocab_mask
5
vocab_mask = org_vocab_mask | added_vocab_mask
6
input_ = vocab_mask * (input_ - valid_offset)
7
return input_, ~vocab_mask

这段代码值得逐行拆开。valid_offset 是两段各自的偏移：基础段的 id 减去本段起点，就变成从 0 开始的局部下标；追加段的 id 需要再跨过基础段的长度与 padding。vocab_mask 为真的位置保留，为假的位置直接把 id 置 0——置 0 而不是置 -1，因为后面要做 embedding 查表，负下标会越界。查出来的那一行随后用 masked_fill_ 清零：

1
output_parallel = self.quant_method.embedding(self, masked_input.long())
2
...
3
output_parallel.masked_fill_(input_mask.unsqueeze(-1), 0)

清零之后再 all-reduce，数学上就成立了：每个 id 恰好在一个 rank 上查出非零向量，其余 rank 贡献零向量，求和得到唯一正确的那一行。这就是前面说的所有权不变量的落地方式。

这里还有一个不太起眼但很讲究的优化。当权重是 FP8 时，代码先把输出按 int8 视图做 all-reduce：

1
if is_fp8(output_dtype):
2
    # Each vocab token has one owner, so FP8 bytes can use int8 SUM.
3
    output_parallel = output_parallel.view(torch.int8)

因为每个 token 只有一个所有者，其余 rank 全是零字节，所以”FP8 求和”与”INT8 求和”在结果上完全等价，而后者不涉及任何浮点运算，NCCL 可以直接按字节做归约。这是一个利用稀疏性把通信变成纯内存操作的小技巧。

用融合 kernel 把掩码、偏移、查表合并

上面那段 Python 掩码逻辑会在每次前向产生好几个逐元素 kernel。vLLM 在支持 CUDA 且权重未量化时，改走一个融合算子：

1
self.use_fused_embedding = (
2
    self.tp_size > 1
3
    and current_platform.is_cuda()
4
    and isinstance(quant_method, UnquantizedEmbeddingMethod)
5
)
6

7
output_parallel = ops.vocab_parallel_embedding(
8
    input_, self.weight,
9
    shard_indices.org_vocab_start_index, shard_indices.org_vocab_end_index,
10
    shard_indices.num_org_vocab_padding,
11
    shard_indices.added_vocab_start_index, shard_indices.added_vocab_end_index,
12
)

融合的收益不在算术，而在少了几次全量读写：掩码、偏移、gather 三步对一块 [T, H] 的数据各扫一遍，融合后只扫一遍。注释里还写了它只适用于”plain row-major shard”，量化权重的布局不满足这个前提，因此只能退回通用路径。

LM head 与 logits 收集

vLLM 的 ParallelLMHead 继承自 VocabParallelEmbedding，但它把 forward 直接写成异常：

1
class ParallelLMHead(VocabParallelEmbedding):
2
    def forward(self, input_):
3
        del input_
4
        raise RuntimeError("LMHead's weights should be used in the sampler.")

这个设计在 DeepInfer 里被原样继承。理由不是”懒”，而是职责边界：head 的权重要被采样器直接拿去做 GEMM，而不是作为一个 nn.Module 参与模型前向的调用栈。模型只负责输出 hidden state，head 与采样在同一个内核管线里完成，CUDA Graph 捕获时也更稳定。

采样前的收集发生在 LogitsProcessor._gather_logits：

1
def _gather_logits(self, logits):
2
    if self.use_all_gather:
3
        logits = tensor_model_parallel_all_gather(logits)
4
    else:
5
        logits = tensor_model_parallel_gather(logits)
6
    return logits

两条路径的差别在于谁最终持有完整 logits。all_gather 让每个 rank 都拿到整份 [B, V]，所有 rank 可以独立、确定地做同样的采样；gather 只把结果收到 rank 0，之后再把采样出的 token 广播出去。vLLM 的注释解释了为什么某些平台必须用 all-gather：XLA 要求严格 SPMD，每个设备都要执行相同的后续操作，不能出现”rank 0 有值、其他 rank 是 None”的分支。

收集之后还有一步容易忽略的裁剪：logits[..., :self.org_vocab_size]。前面补齐的 padding 列在这里被切掉，保证采样看到的列数与 tokenizer 的词表大小一致。

不收集整份 logits 的贪心

如果只是贪心解码，其实没必要把整份 [B, V] 收回来。vLLM 提供了 get_top_tokens：每个 rank 先在本地 shard 上做 argmax，然后只 all-gather 一对 (value, index)：

1
local_max_vals, local_max_indices = logits.max(dim=-1)
2
global_indices = local_max_indices + vocab_start
3
local_pair = torch.stack([local_max_vals.float(), global_indices.float()], dim=-1)
4
gathered = tensor_model_parallel_all_gather(local_pair, dim=-1)
5
gathered = gathered.view(hidden_states.shape[0], tp_size, 2)
6
max_rank_idx = gathered[:, :, 0].argmax(dim=-1, keepdim=True)
7
top_tokens = gathered[:, :, 1].gather(dim=-1, index=max_rank_idx)

通信量从 $O(B \times V)$ 降到 $O(B \times 2 \times n)$。$B = 256$、$n = 8$ 时，前者是 66 MB，后者是 16 KB，四个数量级的差距。注释里还特意强调索引用 float32 而不是 bf16：bf16 的尾数只有 7 位，在 12 万这个量级上无法精确表示每一个整数，转换成 bf16 会让相邻 token id 撞在一起。

代价是这条路径只适用于 argmax（以及带一定约束的 top-k），因为 top-p、min-p、典型采样都需要完整的分布做归一化与截断。工程上的取舍很清楚：能只做归约就不要做广播，但前提是采样算法允许。

下面这张图是 vLLM 官方文档里的 TP=4 进程拓扑，可以看到每个 worker 只持有一部分权重，词表矩阵也被切成四份：

vLLM 的张量并行切分与 worker 组织在模型并行与 Worker 层一篇里已经完整拆过；本文只关心它在词表这一维上多出来的那两段特殊逻辑。

路径一：DP attention 先收 hidden，再算 head，再散回 logits

DP attention 的核心是：每个 DP rank 只跑自己那一批请求的注意力，但 MoE 层需要全局 token 一起路由。于是 hidden state 的分布是”每 rank 一份本地 batch”，而 LM head 的权重可能是全局共享的。

SGLang 的选择是先把各 DP rank 的 hidden state 收成一份全局张量，算完全局 logits，再把属于各 rank 的行散回去：

1
hidden_states, local_hidden_states = self._gather_dp_attn_hidden_states(
2
    hidden_states, logits_metadata
3
)
4
...
5
logits = self._compute_lm_head(hidden_states, lm_head, embedding_bias, out=logits_buffer)
6
...
7
logits = self._scatter_dp_attn_logits(logits, local_hidden_states, logits_metadata)

_scatter_dp_attn_logits 里新建的缓冲形状是 [local_rows, vocab]，再用 dp_scatter 把全局 logits 的对应行搬回来。这条路径的账很清楚：多了一次 hidden 的 all-gather（小）和一次 logits 的 scatter（大），但 head 只算一遍，且每个 rank 的采样看到的都是完整词表。

源码里那句注释正是对这笔账的警告：[batch * num_dp_ranks, vocab] 的 all-gather 在显存紧张时会 OOM。这也是为什么 SGLang 在 autotune 的 dummy run 里允许跳过 LM head——只为测注意力 kernel，没必要把几百 MB 的 logits 也搬一遍。

路径二：attention TP 组内 all-gather

当 enable_dp_lm_head 打开时，logits 的收集范围从全局 TP 组缩小到 attention TP 组：

1
self.use_attn_tp_group = get_parallel().enable_dp_lm_head
2
...
3
if self.use_attn_tp_group:
4
    logits = self._gather_attn_tp_logits(logits)

这一步的意义是把”必须同步的 rank 数”从 $n_{\text{tp}}$ 降到 $n_{\text{attn\_tp}}$。DP attention 下，不同 DP rank 的 batch 是独立的，没必要为了一个 rank 的采样把其他 rank 的 logits 也拉进来。收集范围越小，单次通信的参与者越少，尾延迟越可控。

路径三：all-to-all 交换行块

第三条路径最有意思。它不收集 logits，而是交换 logits 的行块：

1
def _tp_lm_head_all_to_all(self, logits):
2
    """Exchange only the row block owned by each destination DP rank."""
3
    logits = logits.contiguous()
4
    all_to_all_output = torch.empty_like(logits)
5
    get_parallel().tp_group.all_to_all_single(
6
        all_to_all_output.view(-1), logits.view(-1)
7
    )
8
    return _reassemble_tp_lm_head_all_to_all_output(all_to_all_output)

前提条件写在 _can_use_tp_lm_head_all_to_all 里：每个 rank 的本地行数必须相同（或者处于 decode CUDA Graph 的等长 padding 布局），且 lm_head.tp_size 与全局 TP 度一致。满足条件时，all-to-all 的通信量与 all-gather 其实是同一个量级——每卡发出 $B \times V/n$、收回 $B \times V/n$——但它换来的是每个 rank 只保留自己那部分行对应的完整列，不需要在本地物化整份 [B, V]。对显存紧张的场景，这是关键差别。

对称内存与 multimem：把 all-gather 做成一次写

当收集不可避免时，SGLang 试图把这次 all-gather 本身做快。triton_symm_mem_ag.py 实现了一个基于对称内存与 multicast 的 all-gather：所有 rank 的缓冲区映射到同一段虚拟地址，发送方用 multimem.st 一次写入，所有 rank 的副本同时更新，不需要 ring 式的逐跳转发。

1
class MultimemAllGatherer:
2
    def __call__(self, x):
3
        if (state is not None and x.dtype == torch.bfloat16 and x.dim() == 2
4
                and x.is_contiguous() and 0 < x.shape[0] <= state.max_token_num
5
                and x.data_ptr() % 16 == 0
6
                and x.shape[-1] * state.world_size <= state.hidden_dim):
7
            return all_gather_inner(state, x, tp_hidden_dim=x.shape[-1] * state.world_size)
8
        return tensor_model_parallel_all_gather(x, dim=-1)

这段条件判断值得逐条看，因为它暴露了这类”零拷贝集合通信”的全部约束：

• x.dtype == torch.bfloat16：kernel 按 128 位打包写，只对 bf16 有实现；
• x.dim() == 2 且连续：kernel 假设输入是行主序的二维块；
• x.shape[0] <= state.max_token_num：对称内存缓冲在初始化时就按最大 token 数分配，超出只能退回 NCCL；
• x.data_ptr() % 16 == 0：multimem.st 要求 16 字节对齐；
• x.shape[-1] * world_size <= state.hidden_dim：目标缓冲的维度必须够宽。

任何一条不满足就 fall back 到 NCCL。更细的一条在 _build 里：如果当前正处于 CUDA Graph 捕获中，不能分配对称内存缓冲，于是返回 _UNINIT 等待下一次 eager 调用再建。这说明连”通信原语的初始化”都要服从 CUDA Graph 的约束。

all_gather_inner 内部还有一个按 token 分块的外层循环。注释写得很清楚：一次搬完整个 [B, V] 会超出对称内存缓冲，按 max_tokens 切块可以复用同一段缓冲，代价是多次 kernel 启动。这与 vLLM 的 LogitsGatherer 思路一致——大块通信要分块做，否则显存峰值不可控。

下面这张 SGLang 官方图展示了 DP attention 的结构：注意力在 DP 组内独立，MoE 走全局专家并行，LM head 与 logits 的归属需要在这两种拓扑之间做选择：

DP attention 的调度与专家并行细节在 SGLang DP Attention 与专家并行 一篇里已经展开；本文只取它对 logits 归属的影响。

LogitsProcessor 只算需要的行

TensorRT-LLM 的 LogitsProcessor.forward 很短：

1
if not return_context_logits:
2
    last_tokens = torch.cumsum(attn_metadata.seq_lens_cuda, dim=0, dtype=torch.long) - 1
3
    hidden_states = hidden_states[last_tokens]
4
logits = lm_head(hidden_states)
5
logits = logits.float()

先用 cumsum(seq_lens) - 1 算出每个请求最后一个 token 在扁平序列里的位置，只把这些行喂给 head。这一步把 $B$ 从”批内总 token 数”降到”请求数”，在 chunked prefill 场景下可能是几十倍的差距。

.float() 是第二处细节。head 的 GEMM 用 bf16 算，结果立刻升到 FP32 再做 softcap、缩放与采样。理由和 vLLM 里索引必须用 FP32 一样：bf16 只有 7 位尾数，12 万级别的 logits 在累加和比较时精度不够，尤其是 top-p 需要把很多小概率相加。

lm_head 的量化取舍

TensorRT-LLM 对 lm_head 是否量化有一条显式规则：只有当 checkpoint 里真的存了 lm_head.weight_scale 时才量化。注释里给了两个理由：

• 很多 bf16 checkpoint 根本不包含 lm_head 的量化参数，强行量化会引入额外误差；
• tie_word_embeddings 时 lm_head 与 embedding 共享同一块稠密 bf16 权重，量化会破坏共享关系；
• 在 ADP 里做 lm_head TP 时，权重是在 forward 时切片的，也不能预先量化。

这条规则看起来是工程细节，实际上决定了词表矩阵在显存预算里的分量：量化到 FP8 可以把 1.23 GiB 压到 0.62 GiB，量化到 NVFP4 可以压到 0.31 GiB。对一张 80 GB 的卡，这是”能不能多放几层 KV”的差别。TensorRT-LLM 的量化全链路在量化全链路一篇里已经拆过，这里只强调词表矩阵的特殊性：它是少数可以整块复制、也最值得量化的权重。

output 层是一块独立的放置单元

权重加载时，llama.cpp 把层编号分成两类：重复层（0 到 n_layer - 1）与输出层（编号 n_layer）。两者走同一个设备分配函数：

1
// assign the input layer
2
// there is very little benefit to offloading the input layer, so always keep it on the CPU
3
pimpl->dev_input = { cpu_dev, &pimpl->cpu_buft_list };
4

5
// assign the repeating layers to the devices according to the splits
6
for (int il = 0; il < n_layer_all; ++il) {
7
    pimpl->dev_layer[il] = get_layer_buft_list(il);
8
}
9

10
// assign the output layer
11
pimpl->dev_output = get_layer_buft_list(n_layer_all);

输入层强制放在 CPU，注释写的理由是”offload 输入层几乎没有收益”——embedding 查表只读几十行，把 1.23 GiB 的矩阵搬到显存是一次性成本，而每次前向读的量很小。输出层则跟着 -ngl 的层数走，可以被 offload 到 GPU。

这两条决策对应的是两个完全不同的瓶颈：

• embedding 是随机访问，读放大严重，但总量小；放 CPU 时通过统一内存或 pinned buffer 访问，延迟高但带宽需求低。
• lm_head 是稠密 GEMM，计算量是 $2 B V H$，放 CPU 会直接把 decode 速度打下来；放 GPU 后又要多传一次 hidden state，但 hidden 只有 $B \times H$，与 logits 的 $B \times V$ 相比微不足道。

源码的日志也把这件事说得很直白：offloading output layer to GPU 与 offloaded %d/%d layers to GPU。在显存不够的情况下，用户会看到”层被分成两段”——前几层在 GPU、后几层在 CPU，或者反过来只把输出层留在 CPU。

只算需要的行，在图的末端切

一次 prefill 的输入可能有几千个 token，但只有每个序列的最后一个 token 需要 logits。llama.cpp 在图的倒数第二层之后插入一次按行 gather：

1
if (il == n_layer - 1 && inp_out_ids) {
2
    cur   = ggml_get_rows(ctx0,   cur, inp_out_ids);
3
    inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
4
}

ggml_get_rows 只挑出需要的行，后面的输出归一化与 lm_head 只处理这几十行。这一步的收益是 $T / B$ 倍的节省——prefill 8192 token、batch 32 时，lm_head 的计算量从 8192 行降到 32 行，直接省掉 256 倍。

有意思的是源码里一段被注释掉的优化：

1
// note: when all tokens are output, we could skip this optimization to spare the ggml_get_rows() calls,
2
//       but this would make the graph topology depend on the number of output tokens, which can interfere with
3
//       features that require constant topology such as pipeline parallelism

如果所有 token 都要输出（比如训练或逐 token 打分），这次 gather 就是多余的。但跳过它会让图的结构依赖输出行数，而流水并行（PP）要求各 rank 在不同 micro-batch 上跑同一个拓扑的图。为了保住拓扑稳定，llama.cpp 选择永远保留这次 gather，牺牲一点点 eager 模式的性能。这和 vLLM 在 CUDA Graph 下要求形状稳定是同一类约束，只是发生在图构建阶段而不是捕获阶段。

量化表与词表

llama.cpp 的 lm_head 与 embedding 都按 GGUF 的块量化存储，常见的是 Q4_K、Q6_K 等 k-quants 格式。块量化把每 32 个或 256 个权重连同缩放因子打包成一个 block，embedding 查表时读的是一整块而不是连续行。这对随机访问更友好——块内连续、缓存命中率高；对 lm_head 的 GEMM 则需要 mmq 内核在计算时边解码边乘。GGUF 的块格式与 mmq 内核在 GGUF 与 k-quants 块格式和 llama.cpp 后端与 Kernel两篇里已经拆过。

llama.cpp 的结论对本文的意义在于：词表切分的必要性来自张量并行，而不是词表本身。当并行拓扑不切 hidden 维时，词表也可以完整保留；此时要优化的是”放哪、算什么精度、只算几行”，而不是”怎么收集 logits”。

30 行的 VocabParallelEmbedding

1
class VocabParallelEmbedding(torch.nn.Module):
2
    def __init__(self, num_embeddings: int, embedding_dim: int, weight: torch.nn.Parameter) -> None:
3
        super().__init__()
4
        assert weight.shape == (num_embeddings, embedding_dim)
5
        self.weight = reconstruct_param(weight)
6

7
    def forward(self, input_: torch.Tensor) -> torch.Tensor:
8
        output = gather(self.weight, input_)
9
        return output
10

11

12
class ParallelLMHead(VocabParallelEmbedding):
13
    def forward(self, input_: torch.Tensor) -> None:
14
        del input_
15
        raise RuntimeError("LMHead's weights should be used in the sampler.")

和 vLLM 的同名类相比，这里既没有 shard_indices，也没有掩码与 all-reduce：forward 就是一次按行 gather。名字里的 “Parallel” 是历史遗留，实际行为是”全复制 + 本地查表”。这篇源码在开头的版权头里写着 Derived from vLLM; modified by DeepSeek.，改动的核心就是删掉分片逻辑。

reconstruct_param 那一层包装的理由在模型装配与前向主循环一篇里解释过：装配阶段先物化 HF 参数壳再构造运行时模型，重新包一层 Parameter 可以共享 data 但隔离生命周期。

gather 内核：为什么不用 torch 的 index_select

查表走的是 tilelang_ops/gather.py：

1
def gather(x, index, out=None):
2
    """Row gather ``x[index]`` with out-of-range indices zero-filled."""
3
    ...
4
    if out is None:
5
        out = torch.empty(out_shape, dtype=x.dtype, device=x.device)
6
    x = _flat_dense_int8(x)
7
    y = _flat_dense_int8(out)
8
    kernel = ops.gather_tl(x.shape[1], x.dtype, index.dtype, x.stride(0))
9
    kernel(x, index, y)
10
    return out

三个细节值得讲。

第一，越界下标零填充。y[i] = x[index[i]] 若 index[i] 不在 [0, n)，输出整行为 0。这和 vLLM 的掩码逻辑目的相同，但位置不同：vLLM 是在查表之前把 id 置 0 并记住掩码，查完再清零；DeepInfer 则是让内核在越界时直接写零。之所以需要这个行为，是因为图像占位 token 用的是 vocab_size + offset 的哨兵值——这些 id 在真实词表之外，查 embedding 时必须安全地变成零向量，再由多模态路径回填视觉特征。

第二，按字节视图处理。_flat_dense_int8 把任意 dtype 的张量 view 成 int8，gather 内核只做字节搬运，不关心元素类型。这让同一个内核可以服务 bf16、fp8、fp4 各种权重，只要行是稠密的。对于 FP8 的 embedding，这意味着查表也是零拷贝语义。

第三，strided 输入不复制。注释里说明，非连续但”行内稠密”的张量可以直接用 strided 内核，只有任意内部布局才先 contiguous。参数化测试里专门有 strided_x 这一维，用 randn(n, d+17)[:, :d] 构造行间有间隙的输入，验证内核在 stride 不为 d 时仍然正确。

复制词表的代价与收益

在 DP attention + Wide-EP 的拓扑下，每个 rank 都要独立处理自己的一批请求。如果词表按 TP 切分，每次 embedding 都要一次跨 rank all-reduce——但在 DP 拓扑里没有”TP 组”这个概念，all-reduce 就要退化成 DP 组内的全收集，通信量会随 DP rank 数线性增长。反过来，复制一份词表只需要 1.23 GiB × 2 的显存，对一张 192 GB 的卡来说是 1.3% 左右。

更重要的是采样侧没有集合通信。每个 rank 本地持有完整词表，本地算出完整 logits，本地采样，输出 token 直接从 rank 返回。整个解码关键路径上少了一次 [B, V] 的 all-gather 或 all-to-all。在 EP 规模到几十、上百 rank 的部署里，这次通信的尾延迟比那 2.5 GiB 显存更贵。

源码里唯一的跨 rank 同步出现在 prompt logprobs：只有 CP rank 0 物化结果，因为 prompt logprobs 只用于调试与评测，不需要让每个 rank 都算一遍。

1
if PG.context_parallel_group().rank() == 0:
2
    result = torch.empty((full_batch_size,), dtype=torch.float, device=hidden_states.device)

注意注释里那句”每个 CP rank 必须收到相同的 hidden-state 行数，这样 chunk 边界才对齐”——即使只有 rank 0 需要结果，所有 rank 也要按相同分块走完循环，否则跨 rank 的同步会错位。

lm_head 的投影写在哪

GPTHead 里有一个普通的 ReplicatedLinear：

1
class GPTHead(nn.Module):
2
    def __init__(self, cfg, layer_id=None):
3
        super().__init__()
4
        self.norm = DropoutAddRMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps, ...)
5
        self.head = ReplicatedLinear(cfg.hidden_size, cfg.vocab_size, bias=False, use_fp8_gemm=False)

但这个 head 并不在模型前向里被调用。真正算 logits 的是 sampler_utils.get_logits：

1
def get_logits(hidden_states, embedding, vocab_size):
2
    logits = torch.empty(hidden_states.size(0), embedding.size(0),
3
                         device=hidden_states.device, dtype=torch.float)
4
    bf16_gemm_nt(hidden_states, embedding, logits)
5
    assert logits.shape[-1] == vocab_size
6
    return logits

embedding 就是 VocabParallelEmbedding.weight，也就是那份复制的词表矩阵；bf16_gemm_nt 是”NT 布局”的 GEMM——A 是行主序的 hidden，B 是行主序的权重（即 $V \times H$），正好对应 $X W^\top$。输出缓冲直接分配为 FP32，与 TensorRT-LLM 的 logits.float() 是同样的理由。

ReplicatedLinear 里还有一套输出对齐逻辑：

1
def _allocate_output(self, shape, dtype, device):
2
    outputs_per_alignment = _OUTPUT_ALIGNMENT_BYTES // itemsize
3
    physical_output_size = (self.output_size + outputs_per_alignment - 1) // outputs_per_alignment * outputs_per_alignment
4
    storage = torch.empty((*shape[:-1], physical_output_size), dtype=dtype, device=device)
5
    return storage[..., : self.output_size]

按对齐字节数分配物理宽度，再切出逻辑宽度。源码里 _OUTPUT_ALIGNMENT_BYTES = 32，也就是说每一行的物理跨度必须是 32 字节的整数倍。FP32 下是 8 个元素一组，129280 恰好能被 8 整除；换成 BF16 则是 16 个元素一组，同样整除。这样做让下游的 tile kernel 拿到的 stride 始终对齐，避免每次 GEMM 都做一次拷贝。

CUDA Graph 里的约束

全复制还有一个附带好处：在 CUDA Graph 下，词表矩阵是每 rank 一份的静态权重，head 的 GEMM 可以在图里被捕获并重放，没有任何集合通信需要在图外处理。相比 vLLM/SGLang 需要为 all-gather 做显存预算与分块，DeepInfer 的图捕获简单得多。代价是显存，收益是解码路径的确定性——这与 DeepInfer 在异步调度与 ZOS里”把 CPU 后处理移出关键路径”的整体思路是一致的。

下图是 DeepInfer 官方仓库给出的系统架构总览，可以看到它不做 TP、而是把并行度放在 DP Attention 与 Wide-EP 上：

三种收集策略的适用边界

all-gather 是最”通用”的做法：所有 rank 最终持有完整 logits，采样逻辑不需要知道并行拓扑，任意采样算法都能跑。代价是显存峰值最高，因为每个 rank 都物化一份 [B, V]。vLLM 在支持 P2P 的 GPU 平台上默认走这条，SGLang 的 multimem 路径则是它的加速版。

gather 到单 rank 把显存压力集中到一个 rank，其他 rank 只保留分片。它适合”采样逻辑复杂、但只需要一份结果”的场景；代价是结果要再广播回所有 rank，且非目标 rank 在这段时间空转。

不收集、只归约 是通信量最小的路径，但只对”逐元素可分解”的采样算法成立。argmax 可以；top-k 可以（局部 top-k 再合并）；top-p、min-p 不行，因为阈值依赖整份分布的累积和。

DeepInfer 的”零收集”其实是第四类：不是把归约做得更聪明，而是从拓扑上消除了收集的必要性。这也解释了为什么它的 VocabParallelEmbedding 只有 30 行——当并行度不切词表时，这段逻辑确实不需要存在。

显存与通信的换算

把前面的数字代进来，可以做一个粗略的换算。以 TP=8、$B=256$、BF16 logits 为例：

• all-gather 的每卡接收量约 $7/8 \times 66.2$ MB $= 57.9$ MB；
• 若有效带宽 300 GB/s，时间约 193 μs；
• 若改用局部 argmax，通信量降到 16 KB，时间可以忽略。

再看显存。全复制词表在 8 卡上多花 $8 \times 2.47 - 2.47 = 17.3$ GiB 的总显存，但每卡只多 2.16 GiB。对 80 GB 卡来说，这 2.16 GiB 大概相当于 3 到 5 万 token 的 KV 容量（按 V4.1-Flash 的 890 字节/token 估算，约 2.4 万 token）。也就是说，DeepInfer 的选择是用每卡约 2.4 万 token 的 KV 容量，换掉解码路径上的一次 193 μs 通信。在 EP 规模很大、decode 延迟以毫秒计的场景里，这笔交易是划算的；在小规模单机部署里，可能就不划算。

这也是为什么同一个问题在不同框架里有不同答案：答案取决于拓扑，而不是取决于数学。

softmax 的分母是全局量

给定一行 logits $z \in \mathbb{R}^V$，softmax 是：

分母要对整个词表求和。词表切成 $n$ 份后，每个 rank 只有 $z^{(r)}$，它算出的局部和 $\sum_{j \in \text{shard}_r} \exp(z_j - m_r)$ 用的还是局部最大值 $m_r$。要合并成全局分布，必须做两步：

1. 先 all-reduce 出全局最大值 $m = \max_r m_r$；
2. 每个 rank 用 $m$ 重新缩放局部和，再 all-reduce 求和。

这是标准的 log-sum-exp 合并，通信量是每行两个标量，很小。但它是两次串行归约，中间还夹着一次全量的 exp 重算。如果每步都要为采样做这两次归约，延迟就会叠加在关键路径上。

vLLM 与 SGLang 选择”先把 logits 收齐，再在一个 rank 上做一次 softmax”，本质上是用带宽换延迟确定性：一次大通信，换掉两次小但串行的归约。DeepInfer 因为本地就有完整词表，两次归约都退化成一次本地扫描。

top-k、top-p、min-p 各自需要什么

表里最容易被忽略的是最后一行。结构化输出（JSON Schema、正则约束）在采样前要把不允许的 token 的 logits 置为 $-\infty$。这份 bitmask 是按全局 token id 生成的，长度等于词表大小；词表一旦分片，每个 rank 只能取自己那一段 bitmask：

偏移算错一位，整个约束就会错位到别的 token 上。vLLM 与 TensorRT-LLM 都要求 bitmask 按 TP 度补齐并对齐，原因就在这里。SGLang 的语法编译与掩码路径在 SGLang 采样与语法一篇里拆过，DeepInfer 的对应实现在结构化输出与确定性一篇里。

logprobs 的语义

返回给用户的 logprobs 是 $\log p_i$，同样依赖全局分母。如果只在分片上算，得到的是”局部 softmax”，数值会比真实值大——因为分母少了几万项。这类错误不会让程序崩溃，只会让 API 返回的概率之和不等于 1，很难在测试里发现。

TensorRT-LLM 与 vLLM 都在返回前把 logits 升到 FP32 并做全局归一化；DeepInfer 因为本地有完整词表，process_sampling 可以直接对 [total_q, vocab_size] 的 FP32 张量做 log-softmax。词表是否分片，决定了这一步是本地循环还是跨 rank 归约。

确定性与 batch 不变性

还有一个更隐蔽的维度：同一请求在不同 batch 组成下，logits 的归约顺序会变，浮点结果会有 1-2 ULP 的差异。词表分片把这个问题放大了——all-gather 的拼接顺序、all-reduce 的树形结构，都会随 TP 度与拓扑变化。DeepInfer 的全复制反而让每个 rank 的 head GEMM 完全一致，只有 batch 内行数变化带来的 tile 划分差异。

batch-invariant 的完整讨论在 Batch-invariant 推理一篇里；这里只指出一个结论：词表并行是 batch 不变性的一道额外裂缝，因为它的归约维度不在 hidden 上，而在一个通常没人检查的词表维上。

padding 列必须在采样前屏蔽

分片要求词表能被 TP 度整除，于是补零。补出来的列在 lm_head 里也会产生 logits，且值可能不小——因为权重是零，logits 是零，但在 softmax 之后仍有非零概率。若不屏蔽，模型有一定概率采样到一个”不存在的 token id”。

vLLM 的做法是在局部 argmax 前把 padding 列置 $-\infty$：

1
num_pad = lm_head.shard_indices.num_org_vocab_padding
2
if num_pad > 0:
3
    logits[..., -num_pad:] = -float("inf")

然后在收集完成后切掉尾部：logits[..., :self.org_vocab_size]。两步缺一不可——只切不屏蔽，argmax 可能先选中 padding 列；只屏蔽不切，列数与 tokenizer 对不上。

LoRA 追加词表会打乱”下标等于 id”

基础词表与 LoRA 词表分开切分后，物理列的顺序是”基础段所有 rank → 追加段所有 rank”。收集回来的 logits 列序与 token id 不再一一对应。vLLM 为此提供了一个重排映射：

1
def get_sharded_to_full_mapping(self) -> list[int] | None:
2
    """Get a mapping that can be used to reindex the gathered logits for sampling."""

它按 rank 顺序枚举基础段、追加段与 padding，生成一个长度等于 num_embeddings_padded 的索引表，采样前用 index_select 把列重排成 id 顺序。这个细节在多 LoRA 服务里尤其重要：不同请求可能挂不同 adapter，词表大小不同，重排映射也要按请求切换。

CUDA Graph 与集合通信的形状

图捕获要求每次重放的 kernel 序列与形状固定。词表相关的三个量都会破坏这一点：

• batch 行数：vLLM 用 logits_indices 的 padding 把行数对齐到桶；
• logits 宽度：分片宽度固定，但 padding 数随模型变；
• all-gather 的输出缓冲：SGLang 的对称内存缓冲按 max_tokens 预分配，捕获期间不能新建。

SGLang 的 _build 里专门判断了 torch.cuda.is_current_stream_capturing()，捕获中返回 _UNINIT，等 eager 调用时再建。TensorRT-LLM 则要求 gather_output 的形状在各 rank 一致，否则 all-gather 不能走等长路径。这些约束在 KV Cache 与注意力上已经被讨论过很多次，在词表上同样成立。

FP8 与 BF16 的精度分工

词表矩阵通常是 BF16，但 embedding 的输出与 logits 的计算可以走 FP8：

• vLLM 的 FP8 embedding 用 int8 视图做 all-reduce，因为每个 token 只有一个所有者；
• DeepInfer 的 lm_head 显式关闭 FP8（use_fp8_gemm=False），因为词表矩阵不在 prepare_meta_fp8() 的量化范围内；
• TensorRT-LLM 只在 checkpoint 带 scale 时量化 lm_head。

三种做法的共同点是：词表矩阵是精度最敏感的一块。它参与的是”最后一层到 token”的映射，误差会直接改变采样结果；而中间层的 MoE 权重误差会被后续层平滑掉。这也是为什么量化方案通常把 lm_head 留在 BF16 或 FP8，而不敢轻易上 FP4。

通信与计算的流水

当收集不可避免时，最后的优化手段是把 head 的 GEMM 按列分块，算完一块就发一块，与下一块的 GEMM 重叠。vLLM 的 LogitsGatherer 与 SGLang 的 MultimemAllGatherer 都支持按 token 分块；DeepInfer 则因为没有收集而省掉了这一整层调度。

这类流水的收益取决于 head 的计算时间与通信时间的比例。$B$ 很小时，head 的 GEMM 只有几十微秒，通信占主导，流水的收益有限；$B$ 很大时，GEMM 变成毫秒级，分块重叠能把通信藏进去大半。