Blame for src/speculative_decoding.py

"""
src/speculative_decoding.py

Zero-SVD Self-Speculative Decoding for Q-TensorFormer.
Leverages nested Tensor-Train core slicing so that a SINGLE trained model
serves as both its own ultra-fast low-rank draft model and high-capacity target model.
No secondary draft model weights are trained or stored in memory (0 MB parameter overhead).
"""

import time
from typing import Dict, Any, List, Tuple, Optional
import torch
import torch.nn as nn


class ZeroSVDSpeculativeDecoder:
    """
    Self-Speculative Decoding engine for Q-TensorFormer.
    
    Operates by slicing nested TT cores on the fly:
    - Drafting: Slices weights to draft_rank (e.g. r=1 or r=2) to generate K draft tokens.
    - Verification: Slices weights to target_rank (e.g. r=8 or r=16) to verify all K tokens
      in a single parallel forward pass.
    """

    def __init__(
        self,
        model: nn.Module,
        draft_rank: int = 2,
        target_rank: int = 8,
        lookahead_k: int = 3,
        temperature: float = 0.0,
    ):
        self.model = model
        self.draft_rank = draft_rank
        self.target_rank = target_rank
        self.lookahead_k = lookahead_k
        self.temperature = max(0.0, temperature)

    def _set_model_rank(self, rank: int):
        """Sets active rank across all TT layers in the model."""
        for module in self.model.modules():
            if hasattr(module, "set_rank") and callable(module.set_rank):
                module.set_rank(rank)

    @torch.no_grad()
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 32,
    ) -> Tuple[torch.Tensor, Dict[str, Any]]:
        """
        Runs self-speculative decoding.
        
        Returns:
            generated_ids: Tensor of shape (B, S + generated_tokens)
            telemetry: Dict with acceptance_rate, draft_tokens_total, accepted_tokens_total, speedup
        """
        device = input_ids.device
        curr_ids = input_ids.clone()
        total_accepted = 0
        total_drafted = 0
        speculation_steps = 0
        tokens_generated = 0

        t_start = time.perf_counter()

        while tokens_generated < max_new_tokens:
            step_k = min(self.lookahead_k, max_new_tokens - tokens_generated)
            if step_k <= 0:
                break

            # 1. Drafting Phase: Set model to ultra-low draft_rank
            self._set_model_rank(self.draft_rank)
            draft_tokens = []
            draft_probs = []
            draft_prefix = curr_ids.clone()

            for _ in range(step_k):
                outputs = self.model(draft_prefix)
                if isinstance(outputs, torch.Tensor):
                    logits = outputs
                elif isinstance(outputs, (tuple, list)):
                    logits = outputs[0]
                elif hasattr(outputs, "logits"):
                    logits = outputs.logits
                else:
                    logits = outputs

                next_logits = logits[:, -1, :]

                if self.temperature > 0.0:
                    probs = torch.softmax(next_logits / self.temperature, dim=-1)
                    next_tok = torch.multinomial(probs, num_samples=1)
                else:
                    probs = torch.softmax(next_logits, dim=-1)
                    next_tok = next_logits.argmax(dim=-1, keepdim=True)

                tok_id = next_tok[0, 0].item()
                draft_tokens.append(tok_id)
                draft_probs.append(probs[0, tok_id].item())
                draft_prefix = torch.cat([draft_prefix, next_tok], dim=1)

            total_drafted += step_k
            speculation_steps += 1

            # 2. Verification Phase: Set model to target_rank and evaluate in parallel
            self._set_model_rank(self.target_rank)
            
            # Form parallel candidate sequence: curr_ids + draft_tokens
            candidate_seq = torch.cat([
                curr_ids,
                torch.tensor([draft_tokens], device=device, dtype=curr_ids.dtype)
            ], dim=1)

            verify_outputs = self.model(candidate_seq)
            if isinstance(verify_outputs, torch.Tensor):
                verify_logits = verify_outputs
            elif isinstance(verify_outputs, (tuple, list)):
                verify_logits = verify_outputs[0]
            elif hasattr(verify_outputs, "logits"):
                verify_logits = verify_outputs.logits
            else:
                verify_logits = verify_outputs

            # Verification indices: positions evaluating the draft tokens
            start_pos = curr_ids.shape[1] - 1
            n_accepted = 0

            for i in range(step_k):
                pos = start_pos + i
                target_logits_pos = verify_logits[:, pos, :]
                target_prob = torch.softmax(target_logits_pos, dim=-1)[0, draft_tokens[i]].item()
                draft_prob = max(1e-7, draft_probs[i])

                if self.temperature == 0.0:
                    # Greedy acceptance: exact match on argmax
                    target_arg = target_logits_pos.argmax(dim=-1).item()
                    if target_arg == draft_tokens[i]:
                        n_accepted += 1
                    else:
                        # Reject draft and append target's true argmax
                        break
                else:
                    ratio = target_prob / draft_prob
                    if ratio >= 1.0 or torch.rand(1).item() < ratio:
                        n_accepted += 1
                    else:
                        break

            total_accepted += n_accepted

            # Append accepted draft tokens
            if n_accepted > 0:
                accepted_tensor = torch.tensor([draft_tokens[:n_accepted]], device=device, dtype=curr_ids.dtype)
                curr_ids = torch.cat([curr_ids, accepted_tensor], dim=1)
                tokens_generated += n_accepted

            # Append one correction token from the target model if space remains
            if tokens_generated < max_new_tokens:
                correction_pos = start_pos + n_accepted
                correction_logits = verify_logits[:, correction_pos, :]
                if self.temperature > 0.0:
                    corr_probs = torch.softmax(correction_logits / self.temperature, dim=-1)
                    corr_tok = torch.multinomial(corr_probs, num_samples=1)
                else:
                    corr_tok = correction_logits.argmax(dim=-1, keepdim=True)
                curr_ids = torch.cat([curr_ids, corr_tok], dim=1)
                tokens_generated += 1

        t_elapsed = time.perf_counter() - t_start
        acceptance_rate = total_accepted / max(1, total_drafted)

        telemetry = {
            "elapsed_seconds": t_elapsed,
            "tokens_generated": tokens_generated,
            "tokens_per_sec": tokens_generated / max(1e-6, t_elapsed),
            "total_drafted": total_drafted,
            "total_accepted": total_accepted,
            "acceptance_rate": acceptance_rate,
            "speculation_steps": speculation_steps,
            "draft_rank": self.draft_rank,
            "target_rank": self.target_rank,
            "lookahead_k": self.lookahead_k,
            "zero_extra_vram_overhead": True,
            "scientific_classification": "MEASURED",
        }

        # Restore model to target rank
        self._set_model_rank(self.target_rank)
        return curr_ids, telemetry