llm_engine.py

# SPDX-License-Identifier: MIT
# Copyright (C) 2024-2025, Advanced Micro Devices, Inc. All rights reserved.

import itertools
import logging
import time
from collections import Counter
from dataclasses import fields
from typing import Any, Dict, List, Optional, Union

from atom.config import Config
from atom.model_engine.engine_core_mgr import CoreManager
from atom.model_engine.multimodal import get_mrope_input_positions
from atom.model_engine.sequence import Sequence
from atom.sampling_params import SamplingParams
from atom.utils import envs
from transformers import AutoTokenizer, PreTrainedTokenizerFast

logger = logging.getLogger("atom")


def _load_tokenizer(model: str, trust_remote_code: bool = False):
    tokenizer = AutoTokenizer.from_pretrained(
        model, use_fast=True, trust_remote_code=trust_remote_code
    )
    probe = "Hello world 你好"
    if tokenizer.decode(tokenizer.encode(probe), skip_special_tokens=True) != probe:
        logger.warning(
            "AutoTokenizer round-trip failed, falling back to PreTrainedTokenizerFast"
        )
        tokenizer = PreTrainedTokenizerFast.from_pretrained(model)
    return tokenizer


class LLMEngine:

    def __init__(self, model, tokenizer=None, **kwargs):
        config_fields = {field.name for field in fields(Config)}
        config_kwargs = {k: v for k, v in kwargs.items() if k in config_fields}
        data_parallel_size = kwargs.get("data_parallel_size", 1)
        data_parallel_master_port = kwargs.get("data_parallel_master_port", None)
        config = Config(model, **config_kwargs)
        self.config = config
        self.tokenizer = tokenizer or _load_tokenizer(
            config.model, config.trust_remote_code
        )
        config.bos_token_id = self.tokenizer.bos_token_id
        config.eos_token_id = self.tokenizer.eos_token_id
        stop_token_ids = set(config.stop_token_ids)
        # separate eos_token_id from stop_token_ids
        stop_token_ids.discard(config.eos_token_id)
        config.stop_token_ids = list(stop_token_ids)
        # Set data parallel size in config
        config.parallel_config.data_parallel_size = data_parallel_size
        if data_parallel_master_port is not None:
            config.parallel_config.data_parallel_master_port = data_parallel_master_port
        self.data_parallel_size = data_parallel_size
        # PCP and DP-attention are not yet compatible: PCP stripe-splits
        # input_ids to 1/pcp_size in ForCausalLM.forward, but DP-attention's
        # `_gather_ids_for_dp` all-gathers using dp_metadata sizes computed on
        # the FULL (un-split) token count, so all_gatherv asserts
        # `1/pcp_size != full`.
        if config.prefill_context_parallel_size > 1 and config.enable_dp_attention:
            raise ValueError(
                "prefill_context_parallel_size > 1 (-pcp) combined with "
                "--enable-dp-attention is not supported yet (may be supported "
                "in a future release): PCP splits tokens to 1/pcp_size while "
                "DP-attention's id-gather expects the full token count, "
                "causing an all_gatherv size mismatch. For now, disable one of "
                "them (use -tp N -pcp M without DP-attention, or -dp N "
                "--enable-dp-attention without -pcp)."
            )
        # PCP + TBO prefill: supported via coordinated splitting.
        # PCP + TBO decode: not yet supported (pcp_all_reduce semantics under
        # per-request ubatch split are unverified).
        if config.prefill_context_parallel_size > 1 and config.enable_tbo:
            # TBO overlaps compute with the attn<->MoE PCP collectives, which
            # only exist in MoE merge mode (ATOM_PCP_MOE_MERGE=1). With
            # ATOM_PCP_MOE_MERGE=0, MoE runs on each rank's 1/W token shard
            # with NO extra comm between attn and MoE, so TBO has nothing
            # to overlap and only adds ubatch-splitting overhead. Force it off.
            if not envs.ATOM_PCP_MOE_MERGE:
                logger.warning(
                    "Disabling TBO because ATOM_PCP_MOE_MERGE=0: in this "
                    "situation it runs MoE on each rank's 1/W token shard with "
                    "no extra attn<->MoE communication, so TBO has nothing to "
                    "overlap and only adds overhead."
                )
                config.enable_tbo = False
                config.enable_tbo_decode = False
            elif config.tensor_parallel_size > 1:
                # Cross-communicator (PCP x TP) RCCL deadlock:
                # under TBO the PCP collectives (comm_stream) run concurrently
                # and UNORDERED with the TP-group all_reduces (compute_stream,
                # from attention wo_b / MoE RowParallelLinear). On the TPxPCP
                # rank grid with large collectives this forms a cross-rank
                # circular wait -> hang (reproduced MI355 TP4PCP2/TP2PCP4 merge
                # +TBO, 64k/c32). Serialized (non-TBO single stream) is fine, and
                # TP=1 has no TP communicator so it cannot form the cycle. Only
                # TP=1 + PCP keeps TBO; TP>1 falls back to non-TBO.
                logger.warning(
                    "Disabling TBO: prefill_context_parallel_size > 1 (-pcp) "
                    "with tensor_parallel_size > 1 (-tp) hangs under TBO due to "
                    "a PCP<->TP cross-communicator RCCL deadlock (concurrent "
                    "unordered collectives on comm/compute streams; see "
                    "PCP_TBO.md 14.4). TBO with PCP is only supported at -tp 1 "
                    "(e.g. -tp 1 -pcp 8)."
                )
                config.enable_tbo = False
                config.enable_tbo_decode = False
            elif config.enable_tbo_decode:
                raise ValueError(
                    "prefill_context_parallel_size > 1 (-pcp) combined with "
                    "--enable-tbo all (decode TBO) is not supported yet. "
                    "Use --enable-tbo (prefill only) with -pcp."
                )
            # Under PCP, TBO prefill uses a request-boundary split (the
            # non-default TBO split mode; never token-midpoint split), so
            # ATOM_TBO_PREFILL_TOKEN_SPLIT is ignored.
            if config.enable_tbo and envs.ATOM_TBO_PREFILL_TOKEN_SPLIT:
                logger.warning(
                    "ATOM_TBO_PREFILL_TOKEN_SPLIT is ignored under PCP: TBO "
                    "prefill uses request-boundary balanced grouping."
                )
        self.rquest_ids = set()
        self.io_processor = InputOutputProcessor(
            config, self.tokenizer, config.kv_cache_block_size
        )
        self.core_mgr = CoreManager(config)
        self._step_lock = None
        self._pending_results = {}
        import json

        kv_config_str = kwargs.get("kv_transfer_config", "{}")
        try:
            config.kv_transfer_config = json.loads(kv_config_str)
            logger.info(f"KV transfer config loaded: {config.kv_transfer_config}")
        except json.JSONDecodeError:
            config.kv_transfer_config = {}
        logger.info(
            f"LLMEngine init with {self.data_parallel_size} data parallel ranks"
        )
        logger.info(
            f"LLMEngine init with {self.data_parallel_size} data parallel ranks"
        )

    def close(self):
        """Shut down engine and release all GPU resources."""
        if hasattr(self, "core_mgr"):
            self.core_mgr.close()

    def add_request(
        self,
        prompt_or_tokens_list: List[Union[str, List[int]]],
        sampling_params_list: SamplingParams | List[SamplingParams],
        stream_callback=None,
        multimodal_data_list: List[dict] | None = None,
        request_ids: Optional[list[str]] = None,
    ):
        # if sampling params is not list, use it for all prompts
        if not isinstance(sampling_params_list, list):
            sampling_params_iter = itertools.repeat(sampling_params_list)
        else:
            # otherwise check num elements first
            if len(prompt_or_tokens_list) != len(sampling_params_list):
                raise ValueError(
                    f"number of elements in prompt_or_tokens_list and sampling_params_list is different: "
                    f"{len(prompt_or_tokens_list)=} vs {len(sampling_params_list)=}"
                )
            sampling_params_iter = sampling_params_list

        # Handle stream_callback
        if stream_callback is not None and not isinstance(stream_callback, list):
            stream_callback_iter = itertools.repeat(stream_callback)
        elif isinstance(stream_callback, list):
            if len(stream_callback) != len(prompt_or_tokens_list):
                raise ValueError(
                    f"number of elements in prompt_or_tokens_list and stream_callback is different: "
                    f"{len(prompt_or_tokens_list)=} vs {len(stream_callback)=}"
                )
            stream_callback_iter = stream_callback
        else:
            stream_callback_iter = itertools.repeat(None)

        # Handle multimodal data
        if multimodal_data_list is not None:
            if len(prompt_or_tokens_list) != len(multimodal_data_list):
                raise ValueError(
                    f"number of elements in prompt_or_tokens_list and multimodal_data_list is different: "
                    f"{len(prompt_or_tokens_list)=} vs {len(multimodal_data_list)=}"
                )
            mm_data_iter = multimodal_data_list
        else:
            mm_data_iter = itertools.repeat(None)

        # Handle request_ids
        if request_ids is not None:
            if len(request_ids) != len(prompt_or_tokens_list):
                raise ValueError(
                    "number of elements in prompt_or_tokens_list and request_ids is different: "
                    f"{len(prompt_or_tokens_list)=} vs {len(request_ids)=}"
                )
            request_id_iter = iter(request_ids)
        else:
            request_id_iter = itertools.repeat(None)

        reqs = []
        for prompt, sampling_param, callback, mm_data, request_id in zip(
            prompt_or_tokens_list,
            sampling_params_iter,
            stream_callback_iter,
            mm_data_iter,
            request_id_iter,
        ):
            req = self.io_processor.preprocess(
                prompt,
                sampling_param,
                stream_callback=callback,
                multimodal_data=mm_data,
                request_id=request_id,
            )
            reqs.append(req)
        self.core_mgr.add_request(reqs)

    def step(self) -> list[Sequence]:
        seqs = self.core_mgr.get_output()
        return seqs

    def is_finished(self):
        return not self.io_processor.has_pending_requests()

    def generate(
        self,
        prompts: list[str],
        sampling_params: SamplingParams | list[SamplingParams],
        request_ids: Optional[list[str]] = None,
    ) -> list[str]:
        # Reset round-robin counter to ensure consistent DP not core dump
        self.core_mgr._rr_counter = 0

        self.add_request(prompts, sampling_params, request_ids=request_ids)
        outputs = {}
        while not self.is_finished() and (
            self.core_mgr.is_alive() or self.core_mgr.is_rest()
        ):
            seqs = self.step()
            outs = self.io_processor.postprocess(seqs)
            outputs.update(outs)

        outputs = [outputs[seq_id] for seq_id in sorted(outputs)]
        return outputs

    def generate_multimodal(
        self,
        token_ids_list: list[list[int]],
        sampling_params: SamplingParams | list[SamplingParams],
        multimodal_data_list: list[dict],
    ) -> list[dict]:
        """Generate completions for multimodal inputs (token IDs + vision data)."""
        self.core_mgr._rr_counter = 0
        self.add_request(
            token_ids_list,
            sampling_params,
            multimodal_data_list=multimodal_data_list,
        )
        outputs = {}
        while not self.is_finished() and (
            self.core_mgr.is_alive() or self.core_mgr.is_rest()
        ):
            seqs = self.step()
            outs = self.io_processor.postprocess(seqs)
            outputs.update(outs)

        outputs = [outputs[seq_id] for seq_id in sorted(outputs)]
        return outputs

    def start_profile(self):
        self.core_mgr.broadcast_utility_command_sync("start_profile")
        logger.info("Profiling started")

    def stop_profile(self) -> List[Dict[str, Any]]:
        responses = self.core_mgr.broadcast_utility_command_sync(
            "stop_profile", timeout=envs.ATOM_PROFILER_TIMEOUT
        )
        return [resp.get("result", {}) for resp in responses]

    def print_mtp_statistics(self):
        self.core_mgr.send_utility_command("get_mtp_stats")

    def get_mtp_statistics(self, timeout: float = 30.0) -> Dict[str, Any]:
        """Return aggregated speculative decoding statistics across DP ranks."""
        responses = self.core_mgr.broadcast_utility_command_sync(
            "get_mtp_statistics", timeout=timeout
        )
        rank_stats = [
            resp.get("result", resp)
            for resp in responses
            if resp.get("result", resp).get("enabled", False)
        ]

        distribution: Counter[int] = Counter()
        for stats in rank_stats:
            distribution.update(
                {
                    int(accepted): int(steps)
                    for accepted, steps in stats.get("distribution", {}).items()
                }
            )

        total_draft_tokens = sum(
            int(stats.get("total_draft_tokens", 0)) for stats in rank_stats
        )
        total_accepted_tokens = sum(
            int(stats.get("total_accepted_tokens", 0)) for stats in rank_stats
        )
        total_steps = sum(distribution.values())

        return {
            "enabled": bool(rank_stats),
            "total_draft_tokens": total_draft_tokens,
            "total_accepted_tokens": total_accepted_tokens,
            "acceptance_rate": (
                total_accepted_tokens / total_draft_tokens
                if total_draft_tokens
                else 0.0
            ),
            "average_tokens_per_forward": (
                1 + total_accepted_tokens / total_steps if total_steps else 0.0
            ),
            "distribution": dict(sorted(distribution.items())),
            "distribution_percent": {
                k: v / total_steps if total_steps else 0.0
                for k, v in sorted(distribution.items())
            },
        }


class InputOutputProcessor:

    def __init__(self, config, tokenizer, block_size):
        self.config = config
        self.tokenizer = tokenizer
        self.block_size = block_size
        self.requests = {}
        # `has_per_req_cache` flags model architectures that need a
        # per-request stateful buffer outside the paged KV pool. Sequences
        # constructed for these models trigger BlockManager to reserve a
        # per-req cache slot. Currently: GDN-based models (Qwen3-Next /
        # Qwen3.5). Future stateful models (DeepseekV4, etc.) extend the set.
        self._external_to_internal: dict[str, int] = {}
        self._internal_to_external: dict[int, str] = {}
        self.has_per_req_cache = False
        self.num_speculative_tokens = 0
        if (
            hasattr(self.config, "speculative_config")
            and self.config.speculative_config is not None
        ):
            self.num_speculative_tokens = (
                self.config.speculative_config.num_speculative_tokens
            )
        if self.config.hf_config.model_type in self._per_req_cache_model_types():
            self.has_per_req_cache = True

    @staticmethod
    def _per_req_cache_model_types() -> frozenset[str]:
        """Single source of truth for which model_types use per-req cache.

        Read by Sequence-construction (here) AND by ModelRunner's startup
        sanity check, which asserts that any model whose attention builder
        returns `compute_per_req_cache_bytes() > 0` has its model_type
        registered here. Adding a new stateful-attention model means
        adding its model_type to this set.
        """
        return frozenset(
            {
                "qwen3_next",
                "qwen3_5_text",
                "qwen3_5_moe_text",
                "deepseek_v4",
            }
        )

    def preprocess(
        self,
        prompt_or_tokens: str | list[int],
        sampling_params: SamplingParams,
        stream_callback=None,
        kv_transfer_params=None,
        multimodal_data=None,
        request_id: Optional[str] = None,
    ):
        """responsible for:
        1) Tokenize
        2) Create Sequence object

        Single-sequence entry point. Rejects ``sampling_params.n > 1`` so that
        callers which expect exactly one ``Sequence`` back cannot silently
        drop the other siblings. Use :meth:`preprocess_fanout` for n > 1.
        """
        if getattr(sampling_params, "n", 1) > 1:
            raise ValueError(
                "preprocess() returns a single Sequence; for SamplingParams.n > 1 "
                "call preprocess_fanout() and manage the returned list."
            )
        seqs = self.preprocess_fanout(
            prompt_or_tokens,
            sampling_params,
            stream_callback=stream_callback,
            kv_transfer_params=kv_transfer_params,
            multimodal_data=multimodal_data,
            parent_request_id=request_id,
        )
        return seqs[0]

    def preprocess_fanout(
        self,
        prompt_or_tokens: str | list[int],
        sampling_params: SamplingParams,
        stream_callback=None,
        stream_callbacks: Optional[List] = None,
        kv_transfer_params=None,
        multimodal_data=None,
        parent_request_id: Optional[str] = None,
    ) -> List[Sequence]:
        """Tokenize once and materialize ``sampling_params.n`` Sequences.

        Returns a list of length ``n``. For ``n == 1`` this is functionally
        equivalent to the legacy single-sequence path. For ``n > 1``:

        * The prompt is tokenized a single time and the token list is copied
          into each sibling (``Sequence`` copies internally, so mutations stay
          isolated).
        * Every sibling is marked ``needs_independent_noise=True`` so the
          sampler generates fresh per-row noise instead of reusing the cached
          shared exponential tensor. Without this, siblings with identical
          logits would emit identical tokens.
        * Per-sibling ``stream_callbacks`` can be supplied to route streaming
          deltas to independent queues (one per choice index). Falls back to
          the scalar ``stream_callback`` for every sibling.
        """
        n = max(1, int(getattr(sampling_params, "n", 1)))

        tokens = (
            self.tokenizer.encode(prompt_or_tokens)
            if isinstance(prompt_or_tokens, str)
            else prompt_or_tokens
        )
        mrope_positions = None
        mrope_position_delta = 0
        if multimodal_data is not None:
            mrope_positions, mrope_position_delta = get_mrope_input_positions(
                self.config,
                tokens,
                multimodal_data,
            )

        stop_token_sequences = []
        if sampling_params.stop_strings:
            stops = (
                [sampling_params.stop_strings]
                if isinstance(sampling_params.stop_strings, str)
                else sampling_params.stop_strings
            )
            for stop_str in stops:
                stop_tokens = self.tokenizer.encode(stop_str, add_special_tokens=False)
                if stop_tokens:
                    stop_token_sequences.append(stop_tokens)

        if stream_callbacks is not None and len(stream_callbacks) != n:
            raise ValueError(
                f"stream_callbacks length {len(stream_callbacks)} does not match n={n}"
            )

        seqs: List[Sequence] = []
        for i in range(n):
            cb = (
                stream_callbacks[i] if stream_callbacks is not None else stream_callback
            )
            seq = Sequence(
                tokens,
                self.block_size,
                sampling_params,
                stop_token_sequences,
                stream_callback=cb,
                num_draft_tokens=self.num_speculative_tokens,
                has_per_req_cache=self.has_per_req_cache,
                kv_transfer_params=kv_transfer_params,
                multimodal_data=multimodal_data,
                mrope_positions=mrope_positions,
                mrope_position_delta=mrope_position_delta,
                needs_independent_noise=(n > 1),
                parent_request_id=parent_request_id,
                sibling_index=i,
                request_id=parent_request_id if n == 1 else None,
            )
            seq.arrive_time = time.time()
            self.requests[seq.id] = seq
            if seq.external_request_id is not None:
                self._external_to_internal[seq.external_request_id] = seq.id
                self._internal_to_external[seq.id] = seq.external_request_id
            seqs.append(seq)

        if n == 1:
            logger.info(
                f"Request {seqs[0].id} arrived, input tokens: {len(tokens)}, "
                f"pending requests: {len(self.requests)}"
            )
        else:
            logger.info(
                f"Request {parent_request_id or seqs[0].id} fanned out into "
                f"{n} siblings ({seqs[0].id}..{seqs[-1].id}), "
                f"input tokens: {len(tokens)}, "
                f"pending requests: {len(self.requests)}"
            )
        return seqs

    def postprocess(self, reqs: List[Sequence]):
        """responsible for:
        1) Compute stats for logging
        2) Detokenize"""
        outputs = {}
        for req in reqs:
            self.requests.pop(req.id)
            external_request_id = self._internal_to_external.pop(req.id, None)
            if external_request_id is not None:
                self._external_to_internal.pop(external_request_id, None)
            output_str = self.tokenizer.decode(req.completion_token_ids)
            req.leave_time = time.time()

            # Calculate TTFT (Time To First Token) and TPOT (Time Per Output Token)
            ttft = 0.0
            tpot = 0.0
            if req.first_token_time > 0:
                ttft = req.first_token_time - req.arrive_time
                # Calculate TPOT only if there are multiple output tokens
                if req.num_completion_tokens > 1:
                    tpot = (req.leave_time - req.first_token_time) / (
                        req.num_completion_tokens - 1
                    )

            logger.info(
                f"Request {req.id} finished with reason {req.leave_reason}. "
                f"Input tokens: {req.num_prompt_tokens}, output tokens: {req.num_completion_tokens}, "
                f"latency: {req.leave_time - req.arrive_time:.2f}s, "
                f"TTFT: {ttft:.3f}s, TPOT: {tpot:.3f}s"
            )
            outputs[req.id] = {
                "text": output_str,
                "token_ids": req.completion_token_ids,
                "logprobs": req.logprobs if req.return_logprobs else None,
                "latency": req.leave_time - req.arrive_time,
                "finish_reason": req.leave_reason,
                "num_tokens_input": req.num_prompt_tokens,
                "num_tokens_output": req.num_completion_tokens,
                "ttft": ttft,  # Time to first token in seconds
                "tpot": tpot,  # Time per output token in seconds
            }
        return outputs

    def has_pending_requests(self):
        return len(self.requests) > 0