""" 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