多Token预测：大模型推理加速的新范式

生成速度 ≈ GPU内存带宽 / (参数量 × 每参数字节数)

class MultiTokenPredictionHeads(nn.Module):
    def __init__(self, d_model, vocab_size, n_predict=3):
        super().__init__()
        self.n_predict = n_predict
        # 每个预测偏移量一个独立头
        self.heads = nn.ModuleList([
            nn.Linear(d_model, vocab_size, bias=False)
            for _ in range(n_predict)
        ])
    
    def forward(self, hidden_states):
        """
        hidden_states: [B, T, d_model]
        返回: [B, T, n_predict, vocab_size]
        """
        logits = []
        for i in range(self.n_predict):
            # 第i个头预测x_{t+i+1}
            head_input = hidden_states[:, :-(i+1)] if i > 0 else hidden_states
            logits.append(self.heads[i](head_input))
        return torch.stack(logits, dim=2)

def mtp_greedy_decode(model, input_ids, max_new_tokens, n_predict=3):
    generated = input_ids.clone()
    for step in range(max_new_tokens):
        with torch.no_grad():
            outputs = model(generated)
            # MTP heads输出
            logits = outputs['mtp_logits']  # [B, T, n_predict, V]
        
        # 贪婪选择：取argmax
        next_tokens = []
        all_correct = True
        for k in range(n_predict):
            next_token = logits[0, -1, k].argmax(dim=-1)
            # 验证：用前k个token作为输入，检查第k+1个预测
            if k > 0:
                verify_input = torch.cat([generated, torch.stack(next_tokens)], dim=-1)
                verify_logits = model(verify_input)['logits'][0, -1]
                if verify_logits.argmax() != next_token:
                    all_correct = False
                    break
            next_tokens.append(next_token)
        
        if not all_correct:
            # 只接受已验证的token
            break
        
        generated = torch.cat([generated, torch.stack(next_tokens).unsqueeze(0)], dim=-1)
    return generated

def mtp_speculative_decode(model, input_ids, n_predict=3):
    """
    用MTP头推测n_predict个token，主模型验证
    """
    # MTP头生成候选token
    with torch.no_grad():
        mtp_logits = model.mtp_heads(input_ids)  # [B, T, n_predict, V]
    candidates = mtp_logits.argmax(dim=-1)  # [B, T, n_predict]
    
    # 主模型验证
    verify_logits = model.main_head(
        torch.cat([input_ids, candidates], dim=-1)
    )['logits']
    
    # 接受所有与主模型预测一致的token
    accepted = 0
    for k in range(n_predict):
        if (verify_logits[..., k, :].argmax(dim=-1) == candidates[..., k]).all():
            accepted += 1
        else:
            break
    
    return input_ids, accepted

class ProgressiveMTPTrainer:
    def __init__(self, model, schedule_epochs):
        self.model = model
        self.schedule = schedule_epochs  # e.g., {0: 1, 10: 2, 30: 3}
    
    def get_n_predict(self, epoch):
        for e, n in sorted(self.schedule.items(), reverse=True):
            if epoch >= e:
                return n
        return 1
    
    def loss(self, batch, epoch):
        n_predict = self.get_n_predict(epoch)
        outputs = self.model(batch['input_ids'], n_predict=n_predict)
        
        loss = 0
        for k in range(n_predict):
            # 第k个MTP头的loss，权重递减
            weight = 1.0 / (k + 1)
            head_loss = F.cross_entropy(
                outputs['mtp_logits'][:, :-k-1, k],
                batch['labels'][:, k+1:]
            )
            loss += weight * head_loss
        
        return loss

def blockwise_mtp(model, input_ids, block_size=4):
    """
    将序列分块，每块内并行预测
    """
    B, T = input_ids.shape
    n_blocks = (T + block_size - 1) // block_size
    
    outputs = []
    for b in range(n_blocks):
        start = b * block_size
        end = min((b + 1) * block_size, T)
        block_input = input_ids[:, start:end]
        
        # 块内并行预测
        block_output = model.mtp_block_forward(
            block_input,
            n_predict=block_size
        )
        outputs.append(block_output)
    
    return torch.cat(outputs, dim=1)

def update_kv_cache(kv_cache, new_keys, new_values, accepted_count):
    """
    根据接受的token数更新KV Cache
    """
    for layer in range(len(kv_cache)):
        k_cache, v_cache = kv_cache[layer]
        # 只保留已接受token的KV
        kv_cache[layer] = (
            torch.cat([k_cache, new_keys[layer][:accepted_count]], dim=1),
            torch.cat([v_cache, new_values[layer][:accepted_count]], dim=1)
        )
    return kv_cache

def dynamic_batch_mtp(generated_sequences, accepted_counts):
    """
    动态批处理：按接受的token数分组
    """
    # 按accepted_count分组
    groups = {}
    for i, n_accepted in enumerate(accepted_counts):
        groups.setdefault(n_accepted, []).append(i)
    
    # 每组分别处理
    results = []
    for n_accepted, indices in groups.items():
        group_inputs = generated_sequences[indices]
        group_outputs = model.forward(group_inputs)
        results.append((indices, group_outputs))
    
    return results