decoder.py

import torch
import time
from typing import Tuple
from transformers import AutoModelForCausalLM, AutoTokenizer


class SpeculativeDecoder:
    """
    Implements speculative sampling for accelerating LLM inference
    """

    def __init__(
        self,
        target_model_name: str,
        draft_model_name: str,
        device: str = "cuda",
        temperature: float = 1.0,
        top_p: float = 0.9,
    ):
        """
        Initialize speculative decoder with target and draft models.

        Args:
            target_model_name: HuggingFace model ID for target (e.g., "gpt2-large")
            draft_model_name: HuggingFace model ID for draft (e.g., "gpt2")
            device: Device to run models on
            temperature: Sampling temperature
            top_p: Nucleus sampling parameter
        """
        self.device = device
        self.temperature = temperature
        self.top_p = top_p

        # Load models and tokenizer
        print(f"Loading target model: {target_model_name}")
        self.target_model = AutoModelForCausalLM.from_pretrained(target_model_name).to(
            device
        )
        self.target_model.eval()

        print(f"Loading draft model: {draft_model_name}")
        self.draft_model = AutoModelForCausalLM.from_pretrained(draft_model_name).to(
            device
        )
        self.draft_model.eval()

        self.tokenizer = AutoTokenizer.from_pretrained(target_model_name)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token

        # Statistics tracking
        self.reset_stats()

    def reset_stats(self):
        """Reset statistics for tracking performance."""
        self.stats = {
            "total_tokens": 0,
            "total_draft_tokens": 0,
            "accepted_tokens": 0,
            "num_iterations": 0,
            "draft_time": 0.0,
            "target_time": 0.0,
        }

    def sample_from_logits(
        self, logits: torch.Tensor, temperature: float = 1.0, top_p: float = 0.9
    ) -> torch.Tensor:
        """
        Sample a token from logits using nucleus sampling.

        Args:
            logits: Raw logits tensor of shape [vocab_size]
            temperature: Temperature for scaling logits
            top_p: Nucleus sampling parameter (keep tokens with cumulative prob <= top_p)

        Returns:
            Sampled token ID (scalar tensor)
        """
        logits = logits / temperature
        probs = torch.softmax(logits, dim=-1)
        sorted_probs, sorted_indices = torch.sort(probs, descending=True)
        cumulative_probs = torch.cumsum(sorted_probs, dim=-1)
        cutoff_mask = cumulative_probs - sorted_probs > top_p
        sorted_probs[cutoff_mask] = 0.0
        probs_filtered = torch.zeros_like(probs)
        probs_filtered.scatter_(dim=-1, index=sorted_indices, src=sorted_probs)
        probs_filtered = probs_filtered / probs_filtered.sum(dim=-1, keepdim=True)
        token_id = torch.multinomial(probs_filtered, num_samples=1)
        return token_id

    @torch.no_grad()
    def generate_draft_tokens(
        self, input_ids: torch.Tensor, k: int
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Generate K draft tokens using the draft model auto-regressively.

        Args:
            input_ids: Current sequence [batch_size, seq_len]
            k: Number of draft tokens to generate

        Returns:
            draft_tokens: Generated tokens [batch_size, k]
            draft_probs: Probabilities of generated tokens [batch_size, k]
            draft_distributions: Full probability distributions [batch_size, k, vocab_size]
        """
        start_time = time.time()
        batch_size = input_ids.shape[0]

        draft_tokens = []
        draft_probs = []
        draft_distributions = []

        for t in range(k):
            outputs = self.draft_model(input_ids)
            logits = outputs.logits[:, -1, :]  # [batch_size, vocab_size]

            # Get full probability distribution
            probs = torch.softmax(logits / self.temperature, dim=-1)
            draft_distributions.append(probs)  # [batch_size, vocab_size]

            # Sample token
            if batch_size == 1:
                token_id = self.sample_from_logits(
                    logits[0], temperature=self.temperature, top_p=self.top_p
                )  # Returns shape: [1] or scalar

                # Ensure consistent shape [batch_size, 1]
                if token_id.dim() == 0:
                    token_id = token_id.unsqueeze(0).unsqueeze(0)
                elif token_id.dim() == 1:
                    token_id = token_id.unsqueeze(0)

                # Get probability of sampled token
                token_prob = probs[0, token_id[0, 0]]  # Scalar

            else:
                # Batched case
                token_id = []
                token_prob = []
                for b in range(batch_size):
                    tok = self.sample_from_logits(
                        logits[b], temperature=self.temperature, top_p=self.top_p
                    )
                    token_id.append(tok)
                    token_prob.append(probs[b, tok.item()])
                token_id = torch.stack(token_id).unsqueeze(-1)  # [batch_size, 1]
                token_prob = torch.stack(token_prob)  # [batch_size]

            draft_tokens.append(token_id.squeeze(-1))  # [batch_size]
            draft_probs.append(
                token_prob.unsqueeze(0) if token_prob.dim() == 0 else token_prob
            )  # [batch_size]

            # Append to sequence
            input_ids = torch.cat([input_ids, token_id], dim=-1)

        # Stack along the sequence dimension (dim=1)
        draft_tokens = torch.stack(draft_tokens, dim=1)  # [batch_size, k]
        draft_probs = torch.stack(draft_probs, dim=1)  # [batch_size, k]
        draft_distributions = torch.stack(
            draft_distributions, dim=1
        )  # [batch_size, k, vocab_size]

        self.stats["draft_time"] += time.time() - start_time
        return draft_tokens, draft_probs, draft_distributions

    @torch.no_grad()
    def score_draft_tokens(
        self, input_ids: torch.Tensor, draft_tokens: torch.Tensor
    ) -> torch.Tensor:
        """
        Score all K draft tokens in parallel using the target model.

        Args:
            input_ids: Original sequence [batch_size, seq_len]
            draft_tokens: Draft tokens to score [batch_size, k]

        Returns:
            target_probs: Probabilities from target model [batch_size, k+1, vocab_size]
        """
        start_time = time.time()

        seq_len = input_ids.shape[1]
        k = draft_tokens.shape[1]
        target_input_ids = torch.cat([input_ids, draft_tokens], dim=1)
        outputs = self.target_model(target_input_ids)
        target_logits = outputs.logits[:, seq_len - 1 : seq_len + k, :]
        target_probs = torch.softmax(target_logits / self.temperature, dim=-1)

        self.stats["target_time"] += time.time() - start_time
        return target_probs

    def modified_rejection_sampling(
        self,
        draft_tokens: torch.Tensor,
        draft_probs: torch.Tensor,
        draft_distributions: torch.Tensor,
        target_probs: torch.Tensor,
    ) -> Tuple[torch.Tensor, int]:
        """
        Accept/reject draft tokens using modified rejection sampling.

        Args:
            draft_tokens: Draft tokens [batch_size, k]
            draft_probs: p(x_t) from draft model [batch_size, k]
            draft_distributions: Full p(x) distributions [batch_size, k, vocab_size]
            target_probs: q(x) from target model [batch_size, k+1, vocab_size]

        Returns:
            accepted_tokens: Tokens that were accepted or resampled [batch_size, n_accepted]
            n_accepted: Number of tokens accepted (int, can be 1 to k+1)
        """
        batch_size, k = draft_tokens.shape
        accepted_tokens = []
        epsilon = 1e-10

        for t in range(k):
            q_t = target_probs[range(batch_size), t, draft_tokens[:, t]]
            p_t = draft_probs[:, t]
            ratio = q_t / (p_t + epsilon)
            r = torch.rand(batch_size, device=draft_tokens.device)
            accept = r < torch.min(torch.ones_like(ratio), ratio)

            if batch_size == 1:
                if accept.item():
                    # ACCEPTED: Add token and continue to next
                    accepted_tokens.append(draft_tokens[:, t])
                else:
                    # REJECTED: Resample from (q - p)+ and STOP
                    resampled_token = self.resample_from_adjusted_distribution(
                        target_probs[:, t, :],  # q(x) - full distribution
                        draft_distributions[:, t, :],  # p(x) - full distribution
                    )
                    accepted_tokens.append(resampled_token.squeeze(-1))
                    break  # CRITICAL: Stop processing remaining draft tokens
            else:
                # For batched case (more complex - handle per-batch acceptance)
                # Simple implementation: if any rejected, stop all
                if accept.all():
                    accepted_tokens.append(draft_tokens[:, t])
                else:
                    # At least one rejected - resample for rejected batches
                    resampled_token = self.resample_from_adjusted_distribution(
                        target_probs[:, t, :], draft_distributions[:, t, :]
                    )
                    accepted_tokens.append(resampled_token.squeeze(-1))
                    break

        # If all k tokens accepted, sample bonus token from target_probs[:, k, :]
        if len(accepted_tokens) == k:
            if batch_size == 1:
                bonus_token = self.sample_from_logits(
                    torch.log(
                        target_probs[0, k, :] + epsilon
                    ),  # Convert probs back to logits
                    temperature=1.0,  # Already applied temperature earlier
                    top_p=self.top_p,
                )
                # Ensure consistent shape [batch_size]
                if bonus_token.dim() == 1 and bonus_token.shape[0] == 1:
                    bonus_token = bonus_token.squeeze(0)  # [1] -> scalar
                if bonus_token.dim() == 0:
                    bonus_token = bonus_token.unsqueeze(0)  # scalar -> [1]
                accepted_tokens.append(bonus_token)
            else:
                # Batched bonus sampling
                bonus_tokens = []
                for b in range(batch_size):
                    bonus_token = self.sample_from_logits(
                        torch.log(target_probs[b, k, :] + epsilon),
                        temperature=1.0,
                        top_p=self.top_p,
                    )
                    if bonus_token.dim() == 1:
                        bonus_token = bonus_token.squeeze(0)
                    bonus_tokens.append(bonus_token)
                accepted_tokens.append(torch.stack(bonus_tokens))

        accepted_tokens = torch.stack(
            accepted_tokens, dim=1
        )  # [batch_size, n_accepted]
        n_accepted = accepted_tokens.shape[1]

        # Update statistics
        self.stats["accepted_tokens"] += n_accepted
        self.stats["total_draft_tokens"] += k

        return accepted_tokens, n_accepted

    def resample_from_adjusted_distribution(
        self, target_probs: torch.Tensor, draft_probs: torch.Tensor
    ) -> torch.Tensor:
        """
        Resample from the adjusted distribution (q - p)+

        Args:
            target_probs: q(x) from target model [batch_size, vocab_size]
            draft_probs: p(x) from draft model [batch_size, vocab_size]

        Returns:
            Resampled token ID [batch_size, 1]
        """
        diff = target_probs - draft_probs
        adjusted_probs = torch.clamp(diff, min=0.0)
        adjusted_probs = adjusted_probs / (
            adjusted_probs.sum(dim=-1, keepdim=True) + 1e-10
        )

        # Sample
        batch_size = target_probs.shape[0]
        resampled = []
        for b in range(batch_size):
            token = torch.multinomial(adjusted_probs[b], num_samples=1)
            resampled.append(token)

        return torch.stack(resampled, dim=0)

    @torch.no_grad()
    def generate(self, prompt: str, max_length: int = 128, k: int = 4) -> str:
        """
        Generate text using speculative decoding.

        Args:
            prompt: Input text prompt
            max_length: Maximum total sequence length
            k: Number of draft tokens per iteration (lookahead)

        Returns:
            Generated text
        """
        self.reset_stats()

        # Encode prompt
        input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
        n = input_ids.shape[1]  # Current sequence length

        print(f"Starting generation with prompt length: {n}")
        print(f"Target length: {max_length}, K={k}")
        print("-" * 60)

        # Main loop (Algorithm 2, line 4)
        while n < max_length:
            self.stats["num_iterations"] += 1

            # Step 1: Generate K draft tokens
            draft_tokens, draft_probs, draft_distribution = self.generate_draft_tokens(
                input_ids, k
            )

            # Step 2: Score draft tokens with target model in parallel
            target_probs = self.score_draft_tokens(input_ids, draft_tokens)

            # Step 3: Modified rejection sampling
            accepted_tokens, n_accepted = self.modified_rejection_sampling(
                draft_tokens, draft_probs, draft_distribution, target_probs
            )

            # Append accepted tokens to sequence
            input_ids = torch.cat([input_ids, accepted_tokens], dim=1)
            n = input_ids.shape[1]
            self.stats["total_tokens"] += n_accepted

            # Progress update
            if self.stats["num_iterations"] % 5 == 0:
                acceptance_rate = (
                    self.stats["accepted_tokens"] / self.stats["total_draft_tokens"]
                    if self.stats["total_draft_tokens"] > 0
                    else 0
                )
                print(
                    f"Iteration {self.stats['num_iterations']}: "
                    f"Length={n}, Accepted={n_accepted}/{k}, "
                    f"Avg acceptance rate={acceptance_rate:.2%}"
                )

        # Decode and return
        generated_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)
        self.print_statistics()
        return generated_text

    @torch.no_grad()
    def generate_baseline(
        self, prompt: str, max_length: int = 128
    ) -> tuple[str, float]:
        """
        Generate text using only the target model (baseline for comparison).

        Args:
            prompt: Input text prompt
            max_length: Maximum total sequence length

        Returns:
            Generated text and time taken
        """
        # Encode prompt
        input_ids = self.tokenizer.encode(prompt, return_tensors="pt").to(self.device)
        n = input_ids.shape[1]

        print(f"Baseline generation with prompt length: {n}")
        print(f"Target length: {max_length}")
        print("-" * 60)

        start_time = time.time()
        tokens_generated = 0

        # Standard autoregressive generation
        while n < max_length:
            outputs = self.target_model(input_ids)
            logits = outputs.logits[:, -1, :]

            # Sample next token
            token_id = self.sample_from_logits(
                logits[0], temperature=self.temperature, top_p=self.top_p
            )

            # Ensure proper shape
            if token_id.dim() == 0:
                token_id = token_id.unsqueeze(0).unsqueeze(0)
            elif token_id.dim() == 1:
                token_id = token_id.unsqueeze(0)

            input_ids = torch.cat([input_ids, token_id], dim=1)
            n = input_ids.shape[1]
            tokens_generated += 1

            if tokens_generated % 20 == 0:
                print(f"Generated {tokens_generated} tokens...")

        total_time = time.time() - start_time

        # Decode and return
        generated_text = self.tokenizer.decode(input_ids[0], skip_special_tokens=True)

        print("\n" + "=" * 60)
        print("BASELINE GENERATION STATISTICS")
        print("=" * 60)
        print(f"Total tokens generated: {tokens_generated}")
        print(f"Total time: {total_time:.3f}s")
        print(f"Tokens per second: {tokens_generated / total_time:.2f}")
        print("=" * 60)

        return generated_text, total_time

    def print_statistics(self):
        """Print performance statistics"""
        print("\n" + "=" * 60)
        print("SPECULATIVE DECODING STATISTICS")
        print("=" * 60)
        print(f"Total tokens generated: {self.stats['total_tokens']}")
        print(f"Number of iterations: {self.stats['num_iterations']}")
        print(
            f"Average tokens per iteration: "
            f"{self.stats['total_tokens'] / self.stats['num_iterations']:.2f}"
        )

        acceptance_rate = (
            self.stats["accepted_tokens"] / self.stats["total_draft_tokens"]
            if self.stats["total_draft_tokens"] > 0
            else 0
        )
        print(f"Overall acceptance rate: {acceptance_rate:.2%}")

        total_time = self.stats["draft_time"] + self.stats["target_time"]
        print(f"\nTiming:")
        print(f"  Draft model time: {self.stats['draft_time']:.3f}s")
        print(f"  Target model time: {self.stats['target_time']:.3f}s")
        print(f"  Total time: {total_time:.3f}s")
        print(f"  Tokens per second: {self.stats['total_tokens'] / total_time:.2f}")
        print("=" * 60)