Spaces:
Running
Running
| import torch | |
| from collections import deque | |
| import torch.nn.functional as F | |
| from common import COMPUTE_DTYPE | |
| class KVCache: | |
| """ | |
| KV Cache designed for Flash Attention 3's flash_attn_with_kvcache API. | |
| Key differences from FA2-style cache: | |
| - Tensors are (B, T, H, D) not (B, H, T, D) | |
| - FA3 updates the cache in-place during flash_attn_with_kvcache | |
| - Position tracked per batch element via cache_seqlens tensor | |
| """ | |
| def __init__(self, batch_size, num_heads, seq_len, head_dim, num_layers, device, dtype): | |
| self.batch_size = batch_size | |
| self.max_seq_len = seq_len | |
| self.n_layers = num_layers | |
| self.n_heads = num_heads | |
| self.head_dim = head_dim | |
| # Pre-allocate cache tensors: (n_layers, B, T, H, D) | |
| self.k_cache = torch.zeros(num_layers, batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype) | |
| self.v_cache = torch.zeros(num_layers, batch_size, seq_len, num_heads, head_dim, device=device, dtype=dtype) | |
| # Current sequence length per batch element (FA3 needs int32) | |
| self.cache_seqlens = torch.zeros(batch_size, dtype=torch.int32, device=device) | |
| # Previous token's normalized embedding for smear (set by model forward pass) | |
| self.prev_embedding = None | |
| def reset(self): | |
| """Reset cache to empty state.""" | |
| self.cache_seqlens.zero_() | |
| self.prev_embedding = None | |
| def get_pos(self): | |
| """Get current position (assumes all batch elements at same position).""" | |
| return self.cache_seqlens[0].item() | |
| def get_layer_cache(self, layer_idx): | |
| """Return (k_cache, v_cache) views for a specific layer.""" | |
| return self.k_cache[layer_idx], self.v_cache[layer_idx] | |
| def advance(self, num_tokens): | |
| """Advance the cache position by num_tokens.""" | |
| self.cache_seqlens += num_tokens | |
| def prefill(self, other): | |
| """ | |
| Copy cached KV from another cache into this one. | |
| Used when we do batch=1 prefill and then want to generate multiple samples in parallel. | |
| """ | |
| assert self.get_pos() == 0, "Cannot prefill a non-empty KV cache" | |
| assert self.n_layers == other.n_layers and self.n_heads == other.n_heads and self.head_dim == other.head_dim | |
| assert self.max_seq_len >= other.max_seq_len | |
| other_pos = other.get_pos() | |
| self.k_cache[:, :, :other_pos, :, :] = other.k_cache[:, :, :other_pos, :, :] | |
| self.v_cache[:, :, :other_pos, :, :] = other.v_cache[:, :, :other_pos, :, :] | |
| self.cache_seqlens.fill_(other_pos) | |
| # Copy smear state: expand batch=1 prev_embedding to num_samples | |
| if other.prev_embedding is not None: | |
| self.prev_embedding = other.prev_embedding.expand(self.batch_size, -1, -1).clone() | |
| class RowState: | |
| # Per-row state tracking during generation | |
| def __init__(self, current_tokens=None): | |
| self.current_tokens = current_tokens or [] # Current token sequence for this row | |
| self.forced_tokens = deque() # Queue of tokens to force inject | |
| self.in_python_block = False # Whether we are inside a python block | |
| self.python_expr_tokens = [] # Tokens of the current python expression | |
| self.completed = False # Whether this row has completed generation | |
| def sample_next_token(logits, rng, temperature=1.0, top_k=None): | |
| """Sample a single next token from given logits of shape (B, vocab_size). Returns (B, 1).""" | |
| assert temperature >= 0.0, "temperature must be non-negative" | |
| if temperature == 0.0: | |
| return torch.argmax(logits, dim=-1, keepdim=True) | |
| if top_k is not None and top_k > 0: | |
| k = min(top_k, logits.size(-1)) | |
| vals, idx = torch.topk(logits, k, dim=-1) | |
| vals = vals / temperature | |
| probs = F.softmax(vals, dim=-1) | |
| choice = torch.multinomial(probs, num_samples=1, generator=rng) | |
| return idx.gather(1, choice) | |
| else: | |
| logits = logits / temperature | |
| probs = F.softmax(logits, dim=-1) | |
| return torch.multinomial(probs, num_samples=1, generator=rng) | |
| class Engine: | |
| def __init__(self, model, tokenizer): | |
| self.model = model | |
| self.tokenizer = tokenizer # needed for tool use | |
| def generate(self, tokens, negative_tokens=[], num_samples=1, max_tokens=None, temperature=1.0, top_k=None, seed=42): | |
| """Same as generate, but does single prefill and then clones the KV cache.""" | |
| assert isinstance(tokens, list) and isinstance(tokens[0], int), "expecting list of ints" | |
| device = self.model.get_device() | |
| # Allocate the KV cache in the compute dtype so it matches what the forward pass emits | |
| dtype = COMPUTE_DTYPE | |
| rng = torch.Generator(device=device) | |
| rng.manual_seed(seed) | |
| assistant_end = 1 | |
| # 1) Run a batch 1 prefill of the prompt tokens | |
| m = self.model.config | |
| kv_model_kwargs = {"num_heads": m.n_kv_head, "head_dim": m.n_embd // m.n_head, "num_layers": m.n_layer} | |
| kv_cache_prefill = KVCache( | |
| batch_size=1, | |
| seq_len=len(tokens), | |
| device=device, | |
| dtype=dtype, | |
| **kv_model_kwargs, | |
| ) | |
| ids = torch.tensor([tokens], dtype=torch.long, device=device) | |
| logits = self.model.forward(ids, kv_cache=kv_cache_prefill) | |
| logits = logits[:, -1, :].expand(num_samples, -1) # (num_samples, vocab_size) | |
| # 2) Replicate the KV cache for each sample/row | |
| kv_length_hint = (len(tokens) + max_tokens) if max_tokens is not None else self.model.config.sequence_len | |
| kv_cache_decode = KVCache( | |
| batch_size=num_samples, | |
| seq_len=kv_length_hint, | |
| device=device, | |
| dtype=dtype, | |
| **kv_model_kwargs, | |
| ) | |
| kv_cache_decode.prefill(kv_cache_prefill) | |
| del kv_cache_prefill # no need to keep this memory around | |
| # 3) Initialize states for each sample | |
| row_states = [RowState(tokens.copy()) for _ in range(num_samples)] | |
| # 4) Main generation loop | |
| num_generated = 0 | |
| while True: | |
| # Stop condition: we've reached max tokens | |
| if max_tokens is not None and num_generated >= max_tokens: | |
| break | |
| # Stop condition: all rows are completed | |
| if all(state.completed for state in row_states): | |
| break | |
| # Ban already-generated tags for each row | |
| for i, state in enumerate(row_states): | |
| banned = torch.tensor( | |
| list(set(state.current_tokens) | set(negative_tokens) - {0, 1}),#), # PAD=0, EOS=1 remain allowed | |
| dtype=torch.long, | |
| device=device, | |
| ) | |
| if len(banned) > 0: | |
| logits[i, banned] = -float("inf") | |
| # Sample the next token for each row | |
| next_ids = sample_next_token(logits, rng, temperature, top_k) # (B, 1) | |
| sampled_tokens = next_ids[:, 0].tolist() | |
| # Process each row: choose the next token, update state, optional tool use | |
| token_column = [] # contains the next token id along each row | |
| token_masks = [] # contains the mask (was it sampled (1) or forced (0)?) along each row | |
| for i, state in enumerate(row_states): | |
| # Select the next token in this row | |
| is_forced = len(state.forced_tokens) > 0 # are there tokens waiting to be forced in deque? | |
| token_masks.append(0 if is_forced else 1) # mask is 0 if forced, 1 if sampled | |
| next_token = state.forced_tokens.popleft() if is_forced else sampled_tokens[i] | |
| token_column.append(next_token) | |
| # Update the state of this row to include the next token | |
| state.current_tokens.append(next_token) | |
| # On <|assistant_end|> or <|bos|>, mark the row as completed | |
| if next_token == assistant_end: | |
| state.completed = True | |
| # Yield the token column | |
| yield token_column, token_masks | |
| num_generated += 1 | |
| # Prepare logits for next iteration | |
| ids = torch.tensor(token_column, dtype=torch.long, device=device).unsqueeze(1) | |
| logits = self.model.forward(ids, kv_cache=kv_cache_decode)[:, -1, :] # (B, vocab_size) |