Hierarchical Speculative Decoding with Dynamic Windows for Efficient Language Model Inference

import torch
import torch.nn.functional as F

def calculate_entropy(logits):
    """
    Calculate the Shannon entropy of the probability distribution.
    Formula: H(p) = -sum(p * log(p))
    """
    probs = F.softmax(logits, dim=-1)
    log_probs = F.log_softmax(logits, dim=-1)
    entropy = -(probs * log_probs).sum(dim=-1)
    return entropy

def sample_greedy(logits):
    """
    Extract the token with the highest logit value.
    """
    return torch.argmax(logits, dim=-1, keepdim=True)

def speculative_decoding_dynamic_window(
    input_ids,
    draft_model,  # M_q: Small draft model
    target_model, # M_p: Large target model
    max_new_tokens=100,
    gamma=5,      # Maximum lookahead window size
    initial_tau=0.0
):
    """
    Speculative Decoding with Dynamic Window (Self-Verify) using Greedy Sampling.
    """
    curr_ids = input_ids.clone()
    tau = initial_tau
    history_rejected_entropies = []
    generated_count = 0
    
    while generated_count < max_new_tokens:
        
        # --- 1. Draft Phase (with Self-Verify) ---
        draft_tokens = []
        draft_logits_list = []
        
        for k in range(gamma):
            with torch.no_grad():
                draft_outputs = draft_model(curr_ids)
                next_token_logits = draft_outputs.logits[:, -1, :]
            
            # [Self-Verify Mechanism]
            current_entropy = calculate_entropy(next_token_logits).item()
            
            # Stop drafting if draft model uncertainty exceeds threshold tau
            if len(history_rejected_entropies) > 0 and current_entropy > tau:
                break
                
            # Greedy sampling from draft model
            next_token = sample_greedy(next_token_logits)
            
            curr_ids = torch.cat([curr_ids, next_token], dim=-1)
            draft_tokens.append(next_token)
            draft_logits_list.append(next_token_logits)
        
        # Fallback if no tokens were drafted
        if len(draft_tokens) == 0:
            with torch.no_grad():
                outputs = target_model(curr_ids)
                next_token = sample_greedy(outputs.logits[:, -1, :])
                curr_ids = torch.cat([curr_ids, next_token], dim=-1)
                generated_count += 1
            continue

        # --- 2. Verification Phase ---
        with torch.no_grad():
            target_outputs = target_model(curr_ids)
        
        # Align target logits: target_outputs.logits[:, i, :] predicts token i+1
        start_idx = curr_ids.shape[1] - len(draft_tokens) - 1
        target_logits_seq = target_outputs.logits[:, start_idx:-1, :]
        
        accepted_count = 0
        rejected = False
        
        for i, draft_token in enumerate(draft_tokens):
            target_logits = target_logits_seq[:, i, :]
            
            # Greedy sampling from target model for verification
            target_token = sample_greedy(target_logits)
            
            if draft_token.item() == target_token.item():
                accepted_count += 1
            else:
                # --- Rejection & Threshold Update ---
                rejected = True
                rejected_entropy = calculate_entropy(draft_logits_list[i]).item()
                history_rejected_entropies.append(rejected_entropy)
                
                # Update tau: Mean of historical rejected draft entropies
                tau = sum(history_rejected_entropies) / len(history_rejected_entropies)
                
                # Rollback to the point of divergence
                keep_len = curr_ids.shape[1] - len(draft_tokens) + accepted_count
                curr_ids = curr_ids[:, :keep_len]
                
                # Append the correct token from the target model
                curr_ids = torch.cat([curr_ids, target_token], dim=-1)
                break 

        # Extra token if all drafts were correct
        if not rejected:
            last_logits = target_outputs.logits[:, -1, :]
            bonus_token = sample_greedy(last_logits)
            curr_ids = torch.cat([curr_ids, bonus_token], dim=-1)
            accepted_count += 1
            
        generated_count += accepted_count
        
    return curr_ids

import torch
import torch.nn.functional as F

def calculate_entropy(logits):
    """
    Calculate entropy H(p) = -sum(p * log(p))
    """
    probs = F.softmax(logits, dim=-1)
    log_probs = F.log_softmax(logits, dim=-1)
    entropy = -(probs * log_probs).sum(dim=-1)
    return entropy

def algorithm_2_hsddw(
    input_ids,
    model_m1,     # Fast Draft Model (e.g., 68M)
    model_m2,     # Strong Draft Model / Proxy Verifier (e.g., 7B)
    model_mp,     # Target Model (e.g., 70B)
    max_new_tokens=200,
    initial_tau=0.1
):
    """
    Algorithm 2: HSDDW - Hierarchical Speculative Decoding with Dynamic Window
    """
    
    curr_ids = input_ids.clone()
    
    # Initialize dynamic thresholds for M1 and M2
    tau_q1 = initial_tau  # Threshold for M1
    tau_q2 = initial_tau  # Threshold for M2
    
    # History of rejected entropies for updating thresholds
    history_h_q1 = []
    history_h_q2 = []
    
    # Initial State
    # mode: "Draft" means M1 generates, M2 verifies.
    # mode: "Verify" means we send accumulated tokens to Mp.
    mode = "Draft" 
    
    # 'pending_tokens' stores tokens generated by M1/M2 that haven't been seen by Mp yet.
    pending_tokens = [] 
    
    step = 0
    while step < max_new_tokens:
        
        if mode == "Draft":
            # --- PHASE 1: DRAFT LOOP (M1 generates, M2 verifies) ---
            
            # 1. M1 generates tokens (Self-verify with tau_q1)
            # -------------------------------------------------
            new_draft_tokens = []
            m1_logits_list = []
            
            # M1 tries to generate a small chunk (e.g., window size 5)
            # Stops early if M1 is uncertain (Entropy > tau_q1)
            for _ in range(5): 
                with torch.no_grad():
                    # Input is current confirmed context + any pending tokens
                    context = torch.cat([curr_ids, torch.tensor(pending_tokens)], dim=-1)
                    m1_out = model_m1(context)
                    next_logits = m1_out.logits[:, -1, :]
                
                # Self-verify: Check M1's confidence
                h_m1 = calculate_entropy(next_logits).item()
                
                if len(history_h_q1) > 0 and h_m1 > tau_q1:
                    break # Stop drafting
                
                token = torch.argmax(next_logits, dim=-1, keepdim=True) # Greedy for demo
                new_draft_tokens.append(token)
                m1_logits_list.append(next_logits)
                
                # Add to temporary context for next M1 step
                pending_tokens.append(token)

            if not new_draft_tokens:
                 # Force one step if M1 generates nothing, to avoid infinite loop
                 # (Implementation detail usually handled by forcing 1 token)
                 pass

            # 2. M2 Verifies M1 (Pre-verify)
            # -------------------------------------------------
            # M2 acts as the "Target" for M1 in this phase.
            
            # We verify the *newly* generated tokens from M1 against M2
            # Context includes confirmed curr_ids + pending_tokens (before this batch)
            # Note: For simplicity, assume we run M2 on the full sequence
            full_seq = torch.cat([curr_ids, torch.tensor(pending_tokens)], dim=-1)
            
            with torch.no_grad():
                m2_out = model_m2(full_seq)
            
            # Identify the logits corresponding to the new draft tokens
            # We only verify the tokens M1 just added
            start_idx = full_seq.shape[1] - len(new_draft_tokens) - 1
            m2_logits_seq = m2_out.logits[:, start_idx:-1, :]
            
            accepted_count = 0
            rejected = False
            
            for i, token in enumerate(new_draft_tokens):
                m2_logits = m2_logits_seq[:, i, :]
                m1_logits = m1_logits_list[i]
                
                # Verification (Standard SD logic / Rejection Sampling)
                target_token = torch.argmax(m2_logits, dim=-1)
                
                if token.item() == target_token.item():
                    accepted_count += 1
                else:
                    rejected = True
                    # [Update M1 Threshold]
                    # M1 was wrong, record its entropy to history
                    h_rejected = calculate_entropy(m1_logits).item()
                    history_h_q1.append(h_rejected)
                    tau_q1 = sum(history_h_q1) / len(history_h_q1)
                    
                    # Truncate pending_tokens to remove the bad token and subsequent ones
                    # pending_tokens currently has ALL pending (old + new). 
                    # We need to cut at the rejection point.
                    valid_len = len(pending_tokens) - len(new_draft_tokens) + accepted_count
                    pending_tokens = pending_tokens[:valid_len]
                    
                    # Append M2's correction
                    pending_tokens.append(target_token)
                    break
            
            # 3. M2 Self-Verify (Decide whether to loop back or verify)
            # -------------------------------------------------
            
            # Check M2's entropy on the LAST token it just verified/corrected
            last_m2_logits = m2_out.logits[:, -1, :]
            h_m2 = calculate_entropy(last_m2_logits).item()
            
            # LOGIC: 
            # If M2 accepted everything (or corrected) AND M2 is confident:
            # We stay in "Draft" mode. We trust M2 to guide M1 further.
            if not rejected and h_m2 <= tau_q2:
                mode = "Draft"
                # Loop continues... M1 will generate more on top of pending_tokens
            else:
                # If M2 rejected something OR M2 is uncertain:
                # We stop the draft loop. It's time to ask the big boss (Mp).
                mode = "Verify"

        elif mode == "Verify":
            # --- PHASE 2: VERIFY LOOP (Mp verifies everything) ---
            
            # Now we have a list of 'pending_tokens' accumulated from the Draft loop.
            # Mp needs to verify them all.
            
            full_seq = torch.cat([curr_ids, torch.tensor(pending_tokens)], dim=-1)
            
            with torch.no_grad():
                mp_out = model_mp(full_seq)
            
            # Verify pending_tokens against Mp
            start_idx = curr_ids.shape[1] - 1
            mp_logits_seq = mp_out.logits[:, start_idx:-1, :]
            
            # Note: We also need M2's logits here for rejection sampling math 
            # if we were doing strict sampling. For greedy, we compare IDs.
            
            final_accepted_count = 0
            
            for i, token in enumerate(pending_tokens):
                mp_logits = mp_logits_seq[:, i, :]
                target_token = torch.argmax(mp_logits, dim=-1)
                
                if token.item() == target_token.item():
                    final_accepted_count += 1
                else:
                    # [Update M2 Threshold]
                    # Mp rejected a token that M2 had approved/generated.
                    # We need to know M2's entropy at this position to update tau_q2.
                    # (Assuming we cached it or recompute it)
                    # h_m2_rejected = ... 
                    # history_h_q2.append(h_m2_rejected)
                    # tau_q2 = sum(history_h_q2) / len(history_h_q2)
                    
                    # Correct with Mp's token
                    final_accepted_count += 1 # The correction counts as 1 valid step
                    # Truncate and fix pending_tokens logic would happen here
                    break
            
            # Commit accepted tokens to curr_ids
            # (Simplification: assuming we take the verified prefix + 1 correction)
            curr_ids = torch.cat([curr_ids, torch.tensor(pending_tokens[:final_accepted_count])], dim=-1)
            
            step += final_accepted_count
            
            # Reset for next round
            pending_tokens = []
            mode = "Draft" # Go back to drafting with M1

    return curr_ids