decoding.py

from __future__ import annotations

from typing import Literal

import torch
import torch.nn.functional as F

from .kv_cache import KVCache, initialize_kv_caches

DecodingStrategy = Literal["greedy", "sample", "top_k", "top_p"]


def decode_greedy(logits: torch.Tensor) -> torch.Tensor:
    return torch.argmax(logits, dim=-1)


def decode_top_k(logits: torch.Tensor, k: int) -> torch.Tensor:
    if k <= 0:
        raise ValueError(f"k must be positive, got {k}.")
    k = min(k, logits.size(-1))
    topk_values, topk_indices = torch.topk(logits, k=k, dim=-1)
    probs = F.softmax(topk_values, dim=-1)
    samples = torch.multinomial(probs, num_samples=1)
    return torch.gather(topk_indices, -1, samples).squeeze(-1)


def decode_top_p(logits: torch.Tensor, p: float) -> torch.Tensor:
    if not 0 < p <= 1:
        raise ValueError(f"p must be in (0, 1], got {p}.")
    sorted_logits, sorted_indices = torch.sort(logits, dim=-1, descending=True)
    sorted_probs = F.softmax(sorted_logits, dim=-1)
    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

    sorted_indices_to_remove = cumulative_probs > p
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = False
    filtered_sorted_logits = sorted_logits.masked_fill(sorted_indices_to_remove, float("-inf"))

    filtered_logits = torch.full_like(logits, float("-inf"))
    filtered_logits.scatter_(dim=-1, index=sorted_indices, src=filtered_sorted_logits)
    probs = F.softmax(filtered_logits, dim=-1)
    return torch.multinomial(probs, num_samples=1).squeeze(-1)


def sample_next_token(
    logits: torch.Tensor,
    *,
    temperature: float = 1.0,
    strategy: DecodingStrategy = "greedy",
    top_k: int | None = None,
    top_p: float | None = None,
) -> torch.Tensor:
    if temperature <= 0:
        return decode_greedy(logits)

    scaled_logits = logits / temperature

    if strategy == "greedy":
        return decode_greedy(scaled_logits)
    if strategy == "sample":
        probs = F.softmax(scaled_logits, dim=-1)
        return torch.multinomial(probs, num_samples=1).squeeze(-1)
    if strategy == "top_k":
        if top_k is None:
            raise ValueError("top_k must be provided when strategy='top_k'.")
        return decode_top_k(scaled_logits, top_k)
    if strategy == "top_p":
        if top_p is None:
            raise ValueError("top_p must be provided when strategy='top_p'.")
        return decode_top_p(scaled_logits, top_p)
    raise ValueError(f"Unknown decoding strategy: {strategy}.")


def allocate_kv_caches(
    model,
    *,
    batch_size: int,
    total_seq_len: int,
    device: torch.device | None = None,
) -> list[KVCache]:
    model_device = device or next(model.parameters()).device
    model_dtype = next(model.parameters()).dtype
    return initialize_kv_caches(
        model.num_layers,
        batch_size=batch_size,
        num_heads=getattr(model, "num_kv_heads", model.num_heads),
        max_seq_len=total_seq_len,
        head_dim=model.head_dim,
        device=model_device,
        dtype=model_dtype,
    )


@torch.no_grad()
def prefill(
    model,
    input_ids: torch.Tensor,
    *,
    kv_caches: list[KVCache] | None = None,
    total_seq_len: int | None = None,
) -> tuple[torch.Tensor, list[KVCache]]:
    if input_ids.ndim == 1:
        input_ids = input_ids.unsqueeze(0)
    elif input_ids.ndim != 2:
        raise ValueError(f"input_ids must have shape (seq_len,) or (batch, seq_len), got {tuple(input_ids.shape)}.")
    if input_ids.size(1) == 0:
        raise ValueError("input_ids must contain at least one token for prefill().")

    model_device = next(model.parameters()).device
    prompt = input_ids.to(model_device)
    if kv_caches is None:
        max_seq_len = total_seq_len or getattr(model, "context_length", prompt.size(1))
        kv_caches = allocate_kv_caches(model, batch_size=prompt.size(0), total_seq_len=max_seq_len, device=model_device)
    logits, kv_caches = model(prompt, kv_caches=kv_caches, use_cache=True)
    return logits, kv_caches


@torch.no_grad()
def decode_step(
    model,
    input_ids: torch.Tensor,
    kv_caches: list[KVCache],
) -> tuple[torch.Tensor, list[KVCache]]:
    if input_ids.ndim == 1:
        input_ids = input_ids.unsqueeze(-1)
    elif input_ids.ndim != 2:
        raise ValueError(f"input_ids must have shape (batch,) or (batch, 1), got {tuple(input_ids.shape)}.")
    logits, kv_caches = model(input_ids, kv_caches=kv_caches, use_cache=True)
    return logits, kv_caches


def branch_kv_caches(kv_caches: list[KVCache], batch_size: int) -> list[KVCache]:
    return [cache.branch(batch_size=batch_size) for cache in kv_caches]


@torch.no_grad()
def generate(
    model,
    input_ids: torch.Tensor,
    *,
    max_new_tokens: int,
    temperature: float = 1.0,
    strategy: DecodingStrategy = "greedy",
    top_k: int | None = None,
    top_p: float | None = None,
    bos_token_id: int | None = None,
    eos_token_id: int | None = None,
    use_kv_cache: bool = True,
) -> torch.Tensor:
    if max_new_tokens < 0:
        raise ValueError(f"max_new_tokens must be non-negative, got {max_new_tokens}.")
    if input_ids.ndim == 1:
        input_ids = input_ids.unsqueeze(0)
        squeeze_output = True
    elif input_ids.ndim == 2:
        squeeze_output = False
    else:
        raise ValueError(f"input_ids must have shape (seq_len,) or (batch, seq_len), got {tuple(input_ids.shape)}.")

    model_device = next(model.parameters()).device
    generated = input_ids.to(model_device)
    if max_new_tokens == 0:
        return generated.squeeze(0) if squeeze_output else generated

    strip_prefix_tokens = 0
    if generated.size(1) == 0:
        if bos_token_id is None:
            raise ValueError("Empty prompts require bos_token_id to seed generation.")
        generated = torch.full(
            (generated.size(0), 1),
            bos_token_id,
            dtype=input_ids.dtype,
            device=model_device,
        )
        strip_prefix_tokens = 1

    context_length = getattr(model, "context_length", None)
    if context_length is not None and generated.size(1) + max_new_tokens > context_length:
        raise ValueError(
            f"Requested total sequence length {generated.size(1) + max_new_tokens} exceeds context length {context_length}."
        )
    target_seq_len = generated.size(1) + max_new_tokens

    was_training = model.training
    model.eval()
    finished = torch.zeros(generated.size(0), dtype=torch.bool, device=generated.device)
    kv_caches: list[KVCache] | None = None

    try:
        if use_kv_cache:
            total_seq_len = generated.size(1) + max_new_tokens
            kv_caches = allocate_kv_caches(
                model,
                batch_size=generated.size(0),
                total_seq_len=total_seq_len,
                device=model_device,
            )
            logits, kv_caches = prefill(model, generated, kv_caches=kv_caches)
        else:
            logits = model(generated)

        for _ in range(max_new_tokens):
            next_token = sample_next_token(
                logits[:, -1, :],
                temperature=temperature,
                strategy=strategy,
                top_k=top_k,
                top_p=top_p,
            )

            if eos_token_id is not None:
                eos_tensor = torch.full_like(next_token, eos_token_id)
                next_token = torch.where(finished, eos_tensor, next_token)

            generated = torch.cat((generated, next_token.unsqueeze(-1)), dim=-1)

            if eos_token_id is not None:
                finished = finished | (next_token == eos_token_id)
                if torch.all(finished):
                    break

            if generated.size(1) >= target_seq_len:
                break

            if use_kv_cache:
                logits, kv_caches = decode_step(model, next_token.unsqueeze(-1), kv_caches)
            else:
                logits = model(generated)
    finally:
        model.train(was_training)

    output = generated[:, strip_prefix_tokens:]
    return output.squeeze(0) if squeeze_output else output