AstrAI/astrai/inference/core/executor.py

import logging
from dataclasses import dataclass
from typing import List, Optional

import torch
from torch import Tensor

from astrai.extension.attention_backend import (
    ATTN_BACKEND,
    CudaBackend,
    attn_backend,
    get_backend,
)
from astrai.inference.core.cache import PagePool
from astrai.inference.core.graph import CudaGraphContext
from astrai.inference.core.task import Task
from astrai.inference.core.workspace import InferenceWorkspace
from astrai.inference.sample import sample
from astrai.model.automodel import AutoModel
from astrai.tokenize.tokenizer import AutoTokenizer

logger = logging.getLogger(__name__)


@dataclass
class SamplingBatchInfo:
    """Per-batch sampling parameters, cached across decode steps.

    Sampling params are constant for a given ordered task set, so they are
    built once (pinned-memory async H2D) and reused until the task set
    changes.  ``top_ks`` is int32 to match the native consumers.
    """

    temperatures: Tensor  # float32 [B]
    top_ks: Tensor  # int32  [B]
    top_ps: Tensor  # float32 [B]
    freq_penalties: Tensor  # float32 [B]
    has_freq: bool  # any frequency_penalty != 0 (avoids per-step GPU .any())


def _build_sampling_batch_info(tasks: List[Task], device) -> SamplingBatchInfo:
    pin = str(device).startswith("cuda")
    freq_penalties = torch.tensor(
        [t.frequency_penalty for t in tasks], dtype=torch.float32, pin_memory=pin
    ).to(device, non_blocking=True)
    return SamplingBatchInfo(
        temperatures=torch.tensor(
            [t.temperature for t in tasks], dtype=torch.float32, pin_memory=pin
        ).to(device, non_blocking=True),
        top_ks=torch.tensor(
            [t.top_k for t in tasks], dtype=torch.int32, pin_memory=pin
        ).to(device, non_blocking=True),
        top_ps=torch.tensor(
            [t.top_p for t in tasks], dtype=torch.float32, pin_memory=pin
        ).to(device, non_blocking=True),
        freq_penalties=freq_penalties,
        has_freq=bool((freq_penalties != 0).any()),
    )


def _warmup_cuda_graphs(
    model: AutoModel,
    pool: PagePool,
    ws: InferenceWorkspace,
    gctx: "CudaGraphContext",
    max_batch_size: int,
    prompt_len: int = 32,
    device: Optional[str] = None,
):
    batch_sizes = [1]
    n = 2
    while n <= max_batch_size:
        batch_sizes.append(n)
        n *= 2
    if max_batch_size not in batch_sizes:
        batch_sizes.append(max_batch_size)

    dev = device or next(model.parameters()).device

    for b in batch_sizes:
        task_ids = [f"_gr_{b}_{i}" for i in range(b)]
        prompt_tokens = [list(range(prompt_len)) for _ in range(b)]
        alloc_ok = True
        for tid, pt in zip(task_ids, prompt_tokens):
            if not pool.task_alloc(tid, pt):
                alloc_ok = False
                break
        if not alloc_ok:
            for tid in task_ids:
                pool.task_free(tid)
            continue

        with torch.inference_mode(), attn_backend(ATTN_BACKEND.CUDA):
            kv_cache = pool.bind_tasks(task_ids, ws, start_pos=0)
            ids_in = torch.tensor(prompt_tokens, dtype=torch.long, device=dev)
            pos_in = torch.arange(prompt_len, device=dev).unsqueeze(0).expand(b, -1)
            model(
                ids_in,
                input_mask=pos_in.unsqueeze(-1) >= torch.arange(prompt_len, device=dev),
                kv_cache=kv_cache,
                position_ids=pos_in,
            )

            for step in range(2):
                seq_pos = prompt_len + step
                ws.position_ids[:b] = seq_pos
                for tid in task_ids:
                    pool.task_extend(tid, seq_pos)
                kv = pool.bind_tasks(task_ids, ws)
                input_mask = ws.decode_mask(ws.position_ids[:b], ws.max_seq_len)
                ids_buf = ws.fill_input_ids([step] * b)
                gctx.forward(
                    model,
                    key=(b,),
                    input_ids=ids_buf.unsqueeze(1),
                    input_mask=input_mask,
                    kv_cache=kv,
                    position_ids=ws.position_ids[:b].unsqueeze(1),
                )

        for tid in task_ids:
            pool.task_free(tid)
        torch.cuda.synchronize()


class Executor:
    """Model forward passes for prefill and decode phases."""

    def __init__(
        self,
        model: AutoModel,
        tokenizer: AutoTokenizer,
        kv_cache: PagePool,
        device: Optional[str] = None,
        dtype: Optional[torch.dtype] = None,
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.kv_cache = kv_cache
        self.device = device or next(model.parameters()).device
        self.dtype = dtype or next(model.parameters()).dtype

        # Per-step decode cache for the steady-state case where the same
        # ordered task set decodes one token per step.  Sampling params are
        # constant across steps; position_ids grows by exactly 1.  Single-slot:
        # any task-set change is a cache miss.
        self._decode_cache: Optional[tuple] = None

        # Pre-allocated fixed-shape buffers for the decode hot path
        # (input_ids, decode mask, KV bind metadata).  Eagerly sized at init
        # so the workspace is CUDA-graph-capture friendly — no allocation
        # during capture.
        config = model.config
        max_q_heads = config.num_attention_heads
        head_dim = config.hidden_size // config.num_attention_heads
        self._workspace = InferenceWorkspace(
            max_batch_size=kv_cache.max_batch_size,
            max_seq_len=kv_cache.max_seq_len,
            max_q_heads=max_q_heads,
            head_dim=head_dim,
            device=self.device,
            dtype=self.dtype,
        )

        # CUDA-graph capture: one graph per (batch_size,) key.
        # Enabled at init-time via _warmup_cuda_graphs for CudaBackend
        # on supported head_dims; left disabled otherwise.
        self._graph_ctx = CudaGraphContext()
        self._try_enable_cuda_graph()

    def _try_enable_cuda_graph(self):
        on_cuda = "cuda" in str(self.device)
        head_dim = (
            self.model.config.hidden_size // self.model.config.num_attention_heads
        )
        if not on_cuda or head_dim not in (32, 64, 128, 256):
            return

        self._graph_ctx.set_enabled(True)
        _warmup_cuda_graphs(
            self.model,
            self.kv_cache,
            self._workspace,
            self._graph_ctx,
            max_batch_size=self.kv_cache.max_batch_size,
            prompt_len=32,
            device=self.device,
        )

    def _sample_logits(
        self,
        logits: Tensor,
        tasks: List[Task],
        return_logprobs: bool = False,
        info: Optional[SamplingBatchInfo] = None,
    ):
        info = info or _build_sampling_batch_info(tasks, self.device)
        if info.has_freq:
            history_lists = [
                t.prompt_ids[-t.rep_window :] + t.output_ids for t in tasks
            ]
            history_lens = [len(ids) for ids in history_lists]
            max_len = max(history_lens, default=0)
            padded_ids = torch.zeros(
                len(tasks), max_len, dtype=torch.long, device=self.device
            )
            padded_mask = torch.zeros(
                len(tasks), max_len, dtype=torch.bool, device=self.device
            )
            for i, ids in enumerate(history_lists):
                length = len(ids)
                padded_ids[i, :length] = torch.as_tensor(
                    ids, dtype=torch.long, device=self.device
                )
                padded_mask[i, :length] = True
        else:
            padded_ids = None
            padded_mask = None

        result = sample(
            logits,
            temperature=info.temperatures,
            top_k=info.top_ks,
            top_p=info.top_ps,
            frequency_penalty=info.freq_penalties,
            input_ids=padded_ids,
            input_mask=padded_mask,
            return_logprobs=return_logprobs,
        )
        if not return_logprobs:
            return result.tolist()

        tokens, logprobs = result
        tokens_list = tokens.tolist()
        logprobs_list = logprobs.tolist()
        for task, logprob in zip(tasks, logprobs_list):
            task.output_logprobs.append(float(logprob))
        return list(zip(tokens_list, logprobs_list))

    def execute_prefill(
        self,
        tasks: List[Task],
        prompt_len: int,
        start_pos: int = 0,
        return_logprobs: bool = False,
    ):
        if start_pos >= prompt_len:
            return []

        tasks = sorted(tasks, key=lambda t: t.task_id)
        batch_sz = len(tasks)

        input_ids = torch.tensor(
            [t.prompt_ids[start_pos:prompt_len] for t in tasks],
            dtype=torch.long,
            device=self.device,
        )

        task_ids = [t.task_id for t in tasks]
        position_ids = (
            torch.arange(start_pos, prompt_len, dtype=torch.long, device=self.device)
            .unsqueeze(0)
            .expand(batch_sz, -1)
        )
        input_mask = position_ids.unsqueeze(-1) >= torch.arange(
            prompt_len, device=self.device
        )

        with torch.inference_mode():
            outputs = self.model(
                input_ids,
                input_mask=input_mask,
                position_ids=position_ids,
                kv_cache=self.kv_cache.bind_tasks(
                    task_ids,
                    self._workspace,
                    start_pos=start_pos,
                ),
            )
            logits = outputs["logits"][:, -1, :]

        return tasks, self._sample_logits(logits, tasks, return_logprobs)

    def execute_decode(
        self, tasks: List[Task], return_logprobs: bool = False
    ) -> List[int]:
        """Decode next token for each task.

        Args:
            return_logprobs: When ``True``, also record (and return)
                the log-probability of each sampled token under the
                post-strategy sampling distribution.  The logprob is
                appended to ``task.output_logprobs`` and the return
                list becomes ``List[Tuple[int, float]]``.

        Returns:
            ``List[int]`` of sampled token IDs, or
            ``List[Tuple[int, float]]`` of ``(token_id, logprob)`` when
            ``return_logprobs`` is ``True``.
        """
        if not tasks:
            return []

        b = len(tasks)
        ws = self._workspace

        # ---- pre-replay: update input buffers in-place ----

        input_ids = ws.fill_input_ids(
            [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks]
        )

        task_ids = [t.task_id for t in tasks]
        cur_positions = [t.next_pos for t in tasks]

        sig = tuple(task_ids)
        cached = self._decode_cache
        if (
            cached is not None
            and cached[0] == sig
            and cur_positions == [p + 1 for p in cached[1]]
        ):
            info = cached[2]
            ws.position_ids[:b] += 1
            self._decode_cache = (sig, cur_positions, info)
        else:
            info = _build_sampling_batch_info(tasks, self.device)
            ws.position_ids[:b].copy_(
                torch.tensor(cur_positions, dtype=torch.long, device=self.device)
            )
            self._decode_cache = (sig, cur_positions, info)

        total_len = max(cur_positions) + 1
        input_mask = ws.decode_mask(ws.position_ids[:b], total_len)

        kv_cache = self.kv_cache.bind_tasks(task_ids, ws)

        # ---- forward (graph replay or live run + capture) ----

        use_graph = (
            self._graph_ctx.enabled
            and "cuda" in str(self.device)
            and isinstance(get_backend(), CudaBackend)
        )
        if use_graph:
            head_dim = (
                self.model.config.hidden_size // self.model.config.num_attention_heads
            )
            use_graph = use_graph and head_dim in (32, 64, 128, 256)
        key = (b,)
        if use_graph:
            input_mask = ws.decode_mask(ws.position_ids[:b], ws.max_seq_len)

        with torch.inference_mode():
            if use_graph:
                outputs = self._graph_ctx.forward(
                    self.model,
                    key=key,
                    input_ids=input_ids.unsqueeze(1),
                    input_mask=input_mask,
                    kv_cache=kv_cache,
                    position_ids=ws.position_ids[:b].unsqueeze(1),
                )
            else:
                outputs = self.model(
                    input_ids.unsqueeze(1),
                    input_mask=input_mask,
                    kv_cache=kv_cache,
                    position_ids=ws.position_ids[:b].unsqueeze(1),
                )
            logits = outputs["logits"][:, -1, :]

        return self._sample_logits(logits, tasks, return_logprobs, info=info)