Tencent AngelSpec Speculative Decoding Framework Deep Dive: MTP + Block Diffusion Dual-Draft Strategy and D-cut High-Concurrency Throughput Optimization

import numpy as np

class WorkloadAnalyzer:
    """Quantifies why different workloads need different draft strategies"""
    
    def __init__(self):
        self.profiles = {
            'conversation': {'entropy': 0.85, 
                'rates': [0.81, 0.66, 0.52, 0.39, 0.29, 0.22, 0.17]},
            'code': {'entropy': 0.35,
                'rates': [0.92, 0.85, 0.78, 0.72, 0.66, 0.61, 0.57]},
            'math': {'entropy': 0.40,
                'rates': [0.90, 0.82, 0.74, 0.67, 0.60, 0.54, 0.49]}
        }
    
    def analyze(self):
        print("Workload Analysis for Draft Strategy Selection:")
        for name, prof in self.profiles.items():
            mtp_speedup = self._mtp_speedup(prof['rates'])
            block_speedup = self._block_speedup(prof['rates'])
            print(f"  {name}: entropy={prof['entropy']:.2f}")
            print(f"    MTP optimal speedup: {max(mtp_speedup):.2f}x")
            print(f"    Block optimal speedup: {max(block_speedup):.2f}x")
            print(f"    Best strategy: {'MTP' if max(mtp_speedup) > max(block_speedup) else 'Block'}")
    
    def _mtp_speedup(self, rates):
        return [sum(rates[:L]) / (L * 0.1 + 1.0) for L in range(1, len(rates) + 1)]
    
    def _block_speedup(self, rates):
        return [sum(rates[:L]) / (0.1 + 1.0) for L in range(1, len(rates) + 1)]

WorkloadAnalyzer().analyze()

import torch
import torch.nn as nn
import torch.nn.functional as F

class MTPWithTTT(nn.Module):
    """MTP with Training-Time Test — eliminating train-inference mismatch"""
    
    def __init__(self, hidden_dim=4096, vocab_size=128000, num_depths=3):
        super().__init__()
        self.num_depths = num_depths
        self.mtp_block = nn.Sequential(
            nn.Linear(hidden_dim * 2, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, vocab_size)
        )
    
    def forward(self, hidden_states: torch.Tensor, embedding_matrix: torch.Tensor):
        """TTT forward: autoregressive unrolling during training"""
        batch_size = hidden_states.shape[0]
        current_hidden = hidden_states[:, -1:, :].detach()
        draft_logits = []
        
        for depth in range(self.num_depths):
            if depth == 0:
                combined = torch.cat([current_hidden, 
                    torch.zeros_like(current_hidden)], dim=-1)
            else:
                prev_token = draft_logits[-1].argmax(dim=-1)
                token_embed = F.embedding(prev_token, embedding_matrix)
                combined = torch.cat([current_hidden, token_embed], dim=-1)
            
            logits = self.mtp_block(combined)
            draft_logits.append(logits)
        
        return torch.cat(draft_logits, dim=1)

    def compute_loss(self, hidden_states, target_tokens, embedding_matrix):
        draft_logits = self.forward(hidden_states, embedding_matrix)
        targets = target_tokens[:, :self.num_depths]
        loss = sum(F.cross_entropy(draft_logits[:, d, :], targets[:, d]) 
                  for d in range(self.num_depths))
        return loss / self.num_depths

class DFlyDrafter(nn.Module):
    """Block diffusion drafter with parallel generation + correction head"""
    
    def __init__(self, hidden_dim=4096, vocab_size=128000, block_size=8):
        super().__init__()
        self.block_size = block_size
        
        self.diffusion = nn.TransformerEncoder(
            nn.TransformerEncoderLayer(hidden_dim, 32, batch_first=True), 6)
        self.pos_embed = nn.Embedding(block_size, hidden_dim)
        self.hidden_proj = nn.Linear(hidden_dim, hidden_dim)
        self.output_proj = nn.Linear(hidden_dim, vocab_size)
        
        self.correction = nn.Sequential(
            nn.Linear(hidden_dim + vocab_size, hidden_dim),
            nn.LayerNorm(hidden_dim), nn.GELU(),
            nn.Linear(hidden_dim, vocab_size))
    
    def forward(self, target_hidden: torch.Tensor):
        B, D = target_hidden.shape
        state = self.hidden_proj(target_hidden).unsqueeze(1)
        state = state.expand(-1, self.block_size, -1)
        state = state + self.pos_embed.weight.unsqueeze(0)
        state = self.diffusion(state)
        block_logits = self.output_proj(state)  # [B, block_size, V]
        
        corrected = []
        for pos in range(self.block_size):
            if pos == 0:
                corrected.append(block_logits[:, 0, :])
            else:
                prev = F.one_hot(corrected[-1].argmax(dim=-1), 
                                num_classes=block_logits.size(-1)).float()
                corr = self.correction(torch.cat([state[:, pos, :], prev], dim=-1))
                corrected.append(block_logits[:, pos, :] + corr)
        
        return torch.stack(corrected, dim=1)

class DCutScheduler:
    """Dynamic verification pruning for high-concurrency throughput"""
    
    def __init__(self, max_verify=4096):
        self.max_verify = max_verify
        self.latency_table = {1: 1.0, 2: 1.2, 4: 1.5, 8: 2.0, 16: 3.0}
    
    def schedule(self, requests: list) -> list:
        candidates = []
        for req in requests:
            conf = req['confidence']
            best = max(
                ((d, np.prod(conf[:d]) * d / self._cost(d)) 
                 for d in range(1, len(conf)+1)),
                key=lambda x: x[1])
            candidates.append({'id': req['id'], 'depth': best[0], 
                              'benefit': best[1], 'cost': self._cost(best[0])})
        
        candidates.sort(key=lambda x: x['benefit'], reverse=True)
        
        scheduled, total = [], 0
        for c in candidates:
            if total + c['cost'] <= self.max_verify:
                scheduled.append(c)
                total += c['cost']
        
        return scheduled
    
    def _cost(self, depth):
        depths = sorted(self.latency_table.keys())
        if depth <= depths[0]: return self.latency_table[depths[0]]
        if depth >= depths[-1]: return self.latency_table[depths[-1]]
        for i in range(len(depths)-1):
            if depths[i] <= depth <= depths[i+1]:
                r = (depth - depths[i]) / (depths[i+1] - depths[i])
                return self.latency_table[depths[i]] + r * (
                    self.latency_table[depths[i+1]] - self.latency_table[depths[i]])
        return self.latency_table[depths[-1]]