Spaces:
Sleeping
Sleeping
| """ | |
| 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) | |
| 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 | |