Cascade Speculative Drafting for Even Faster LLM Inference

import random

class Model:
    def __init__(self, name):
        self.name = name
    
    def generate(self, prefix, max_len=1):
        """模拟生成：Base Case (如 MaG) 直接生成"""
        return [f"{self.name}_token"]
        
    def review(self, prefix, draft, lenience=1):
        """模拟审核：Target Review"""
        # 简单模拟：假设以一定概率接受
        # 论文中 lenience (l) 用于放宽中间层的审核标准
        accepted = []
        for token in draft:
            if random.random() < 0.8: # 模拟接受率
                accepted.append(token)
            else:
                break # 拒绝，截断
        return accepted

def cascade_speculative_drafting_step(draft_models, target_model, prefix, K_matrix, l_lenience):
    """
    对应论文 Algorithm 1: CascadeSpeculativeDraftingStep
    """
    
    # --- Base Case (递归终止条件) ---
    # 对应论文: if draftList is empty then ...
    # 当没有更小的草稿模型可用时，当前的最底层模型(target_model, 即 MaG)直接生成
    if not draft_models:
        # 这里模拟 MaG 的非自回归生成或简单生成
        return target_model.generate(prefix)

    # 初始化
    cur_gen = [] # 对应 curGen (相对于 prefix 的增量)
    full_prefix = prefix[:] # 保持当前完整的上下文
    
    # 获取当前层级的 horizontal 步长配置
    # 对应论文: [k1, ..., kn-1] <- first row of Knn
    current_level_ks = K_matrix[0] 
    
    # --- Horizontal Cascade Loop (水平级联循环) ---
    # 对应论文: for i <- 1 to n do ...
    # 这个循环体现了"水平级联"：随着 i 增加，负责审核的 target 变小(draft_models[i])
    for i, k_i in enumerate(current_level_ks):
        
        # 1. 确定当前的"临时目标"和"临时助手列表"
        # 对应论文: curTarget <- draftList[i]
        cur_target = draft_models[i]
        
        # 对应论文: curDraftList <- sublist of draftList starting from i+1
        cur_draft_list = draft_models[i+1:]
        
        # 对应论文: curK <- submatrix of Knn ...
        # (简化处理：切片取剩余的配置)
        cur_K = [row[i+1:] for row in K_matrix[1:]] 
        
        # 2. 生成当前片段 (Segment Generation)
        # 对应论文: while curGen.length ... < ki do
        # 只要当前片段还没凑够 k_i 个词，就持续递归调用
        segment_len = 0
        while segment_len < k_i:
            
            # --- Vertical Cascade (垂直级联递归) ---
            # 对应论文: Recursive Call -> CascadeSpeculativeDraftingStep
            # 这里的关键是：用更小的模型列表 (cur_draft_list) 为当前模型 (cur_target) 生成草稿
            sub_draft = cascade_speculative_drafting_step(
                cur_draft_list, 
                cur_target, 
                full_prefix + cur_gen, 
                cur_K, 
                l_lenience
            )
            
            cur_gen.extend(sub_draft)
            segment_len += len(sub_draft)
            
            # (为了防止无限循环的模拟保护)
            if len(sub_draft) == 0: break 
            
    # --- Final Review (最终审核) ---
    # 对应论文: review(Mt, curGen, ...)
    # 所有的片段拼凑好后，由本层原本的 target_model 进行统一度量
    # 注意：论文提到如果是最顶层调用，l=1；中间层调用 l>1 (Lenience)
    accepted_draft = target_model.review(prefix, cur_gen, lenience=l_lenience)
    
    return accepted_draft

# ==========================================
# 模拟运行
# ==========================================

# 定义模型层级: Target -> Draft1 -> Draft2 -> MaG
# 实际传入函数的 draft_models 列表是 [Draft1, Draft2, MaG]
target_mt = Model("Target(LLaMA-70B)")
draft_1 = Model("Draft1(LLaMA-7B)")
draft_2 = Model("Draft2(Tiny)")
mag = Model("MaG(N-gram)")

draft_list = [draft_1, draft_2, mag]

# 定义 K 矩阵 (简化版)
# K_matrix[0] 控制最顶层 Horizontal 分配：
#   - 位置 0 (交给 Draft1 负责): 目标生成 3 个词
#   - 位置 1 (交给 Draft2 负责): 目标生成 2 个词
#   - 位置 2 (交给 MaG 负责):    目标生成 1 个词
# Row 0 ([3, 2, 1]): 决定了最终输出序列是由谁拼凑的。（一段 Draft1 的，一段 Draft2 的，一段 MaG 的）。
# Row 1 ([2, 1]): 决定了 Draft 1 是如何被加速的。
# Row 2 ([1]): 决定了 Draft 2 是如何被加速的。
K_matrix = [
    [3, 2, 1], 
    [2, 1],    # 下一层的配置
    [1]        # 再下一层的配置
]

print("Running Algo 1 (Cascade Speculative Drafting)...")
final_result = cascade_speculative_drafting_step(
    draft_models=draft_list,
    target_model=target_mt,
    prefix=["The", "future"],
    K_matrix=K_matrix,
    l_lenience=1
)

print(f"Final Accepted Draft: {final_result}")