MTP推理加速实战：多Token预测如何提升大模型解码速度

1次阅读

共计 1494 个字符，预计需要花费 4 分钟才能阅读完成。

传统语言模型使用逐 Token 生成方式（autoregressive decoding），每个时间步只能产生一个输出 Token。实际测试表明，在 A100 显卡上运行 LLaMA-7B 模型时：

当 batch_size= 1 时，生成速度约 18 Token/s
当 batch_size= 8 时，QPS 仅提升到 35 Token/s

这种线性解码方式导致 GPU 利用率长期低于 30%，成为推理性能的主要瓶颈。

Multi-Token Prediction 的核心思想是让模型同时预测多个未来 Token。结构改进主要体现在：

graph TD
    A[输入序列] --> B[传统 Transformer]
    B --> C[单输出头]
    A --> D[MTP 改造]
    D --> E[多输出头并行]
    E --> F[Token1]
    E --> G[Token2]
    E --> H[Token3]

关键修改点：

输出层扩展为 N 个独立 head（N= 预测 Token 数）
损失函数计算改为 $\mathcal{L} = \sum_{i=1}^N \lambda_i \cdot CE(y_i, \hat{y}_i)$
解码时并行执行多个 Top- K 采样

class MTPGenerationMixin:
    def generate(self, input_ids, **kwargs):
        # 新增 mtp_num 参数控制预测 Token 数
        mtp_num = kwargs.pop('mtp_num', 3)  

        # 修改 logits 处理逻辑
        logits = model(input_ids).logits  # [batch, seq_len, vocab_size * mtp_num]
        logits = logits.view(-1, mtp_num, vocab_size)  # 重塑维度

        # 并行采样
        samples = []
        for i in range(mtp_num):
            probs = F.softmax(logits[:,i,:], dim=-1)
            samples.append(torch.multinomial(probs, 1))
        return torch.cat(samples, dim=1)  # [batch, mtp_num]

import triton

@triton.jit
def parallel_top_k(logits, output, k: int):
    # 每个 block 处理一个预测头
    bid = tl.program_id(0)
    offset = bid * vocab_size

    # 使用共享内存加速排序
    shmem = tl.zeros((BLOCK_SIZE,), dtype=tl.float32)
    for i in range(0, vocab_size, BLOCK_SIZE):
        shmem = logits[offset + i : offset + i + BLOCK_SIZE]
        # 并行 Top- K 算法实现...

测试环境：NVIDIA A100 80GB, LLaMA-7B 模型