sgl-project--sglang/python/sglang/srt/model_executor/runner/eager_runner.py

# Copyright 2023-2026 SGLang Team
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""No-cuda-graph phase runner; the eager dual of BaseCudaGraphRunner."""

from __future__ import annotations

import contextlib
import logging
from dataclasses import replace
from typing import TYPE_CHECKING, Any, Tuple, Union

import torch

from sglang.srt.dllm.config import DllmConfig
from sglang.srt.environ import envs
from sglang.srt.layers.cp.utils import (
    cp_gather_after_forward,
    cp_split_before_forward,
    is_cp_v2_active,
    prepare_cp_forward,
)
from sglang.srt.layers.pooler import EmbeddingPoolerOutput
from sglang.srt.model_executor.cuda_graph_buffer_registry import (
    build_eager_registry,
)
from sglang.srt.model_executor.forward_batch_deepseek_mha_mixin import (
    create_chunked_prefix_cache_kv_indices,
)
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors
from sglang.srt.model_executor.forward_context import (
    ForwardContext,
    forward_context,
    get_req_to_token_pool,
    get_token_to_kv_pool,
)
from sglang.srt.model_executor.runner.base_runner import BaseRunner
from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
    enable_tc_piecewise_cuda_graph,
    set_tc_piecewise_forward_context,
)
from sglang.srt.utils import is_hip
from sglang.srt.utils.common import ceil_align, require_mlp_sync

logger = logging.getLogger(__name__)

_is_hip = is_hip()

if TYPE_CHECKING:
    from sglang.srt.layers.logits_processor import LogitsProcessorOutput
    from sglang.srt.model_executor.model_runner import ModelRunner


class EagerRunner(BaseRunner):
    def __init__(self, model_runner: ModelRunner) -> None:
        super().__init__(model_runner)
        mr = model_runner
        sa = mr.server_args
        # Built first so the cg runners coalesce onto its buffers via the shared
        # input pool; size to the largest tokens/req across modes the worker hits.
        num_tokens_per_bs = 1
        if mr.spec_algorithm.is_speculative():
            # speculative_adaptive can grow draft tokens at runtime; size to the max.
            num_draft_tokens = sa.max_speculative_num_draft_tokens or 1
            if mr.is_draft_worker:
                num_tokens_per_bs = max(
                    sa.speculative_eagle_topk or 1,
                    num_draft_tokens,
                    (
                        2 * (sa.speculative_num_steps or 0)
                        if sa.enable_multi_layer_eagle
                        else 0
                    ),
                )
            else:
                num_tokens_per_bs = (
                    mr.spec_algorithm.get_num_tokens_per_bs_for_target_verify(
                        num_draft_tokens, mr.is_draft_worker
                    )
                )
        else:
            dllm_config = DllmConfig.from_server_args(sa)
            if dllm_config is not None:
                # dLLM runs block_size tokens/request (DLLM_EXTEND).
                num_tokens_per_bs = dllm_config.block_size
        max_bs = mr.max_running_requests
        if (
            mr.is_draft_worker
            and mr.spec_algorithm.is_frozen_kv_mtp()
            and sa.speculative_eagle_topk > 1
        ):
            # Frozen-KV MTP expands the draft batch by topk on the bs axis
            # (expand_for_topk_draft) before the eager fallback.
            max_bs *= sa.speculative_eagle_topk
        # Mirror prepare_mlp_sync_batch padding so the registry holds what load_batch copies.
        if require_mlp_sync(sa):
            from sglang.srt.layers.utils.cp_utils import get_cp_padding_align_size

            max_bs = ceil_align(max_bs, self.attn_tp_size)
            max_bs = ceil_align(max_bs, get_cp_padding_align_size())
        prefill_ceiling = max(mr.max_total_num_tokens, sa.max_prefill_buffer_tokens())
        max_num_token = max(prefill_ceiling, max_bs * num_tokens_per_bs)
        if require_mlp_sync(sa):
            max_num_token = ceil_align(max_num_token, self.attn_tp_size)
            max_num_token = ceil_align(max_num_token, get_cp_padding_align_size())
        self._eager_max_bs = max_bs
        self._eager_num_tokens_per_bs = num_tokens_per_bs
        is_encoder_decoder = mr.model_config.is_encoder_decoder
        self._eager_registry = build_eager_registry(
            device=mr.device,
            max_bs=max_bs,
            max_num_token=max_num_token,
            cache_loc_dtype=torch.int64,
            enable_mamba_track=(
                sa.enable_mamba_extra_buffer() and mr.spec_algorithm.is_none()
            ),
            is_encoder_decoder=is_encoder_decoder,
            encoder_len_fill_value=(
                getattr(mr.model_config.hf_config, "max_source_positions", 0)
                if is_encoder_decoder
                else 0
            ),
            encoder_lens_dtype=(
                torch.int64 if torch.device(mr.device).type == "cpu" else torch.int32
            ),
            dp_size=sa.dp_size,
        )
        # Eager has no capture step, so warm up here (run-once via mr._kernel_warmed_up).
        self.warmup()

    def _autotune_buffers(self) -> Tuple[Any, int]:
        """Decode-shaped dummy buffers (bs * num_tokens_per_bs) for the warmup
        flashinfer-autotune forward.

        flashinfer's MoE autotuner times candidate tactics against the buffer it
        is given, so it must match the live decode shape for the cached tactic to
        be optimal at decode. The eager input registry spans the prefill token
        ceiling; the dummy run only needs the decode-sized slice.
        """
        mr = self.model_runner
        num_tokens_per_bs = 1
        if mr.spec_algorithm.is_speculative():
            num_tokens_per_bs = (
                mr.spec_algorithm.get_num_tokens_per_bs_for_target_verify(
                    mr.server_args.speculative_num_draft_tokens, mr.is_draft_worker
                )
            )
        return (
            self._alloc_dummy_decode_buffers(
                self._eager_max_bs, num_tokens_per_bs=num_tokens_per_bs
            ),
            self._eager_max_bs,
        )

    def can_run_graph(self, forward_batch: ForwardBatch) -> bool:
        # Eager never runs a cuda graph; callers dispatch on isinstance(...,
        # EagerRunner) and must not route an eager batch into a replay branch.
        return False

    def load_batch(
        self, forward_batch: ForwardBatch, pp_proxy_tensors=None, **kwargs
    ) -> ForwardBatch:
        """Copy the live batch into the fixed-max eager static buffers (sliced to
        this batch's shape) — the eager counterpart of the cuda-graph runners'
        load_batch."""
        if envs.SGLANG_EAGER_INPUT_NO_COPY.get():
            return replace(forward_batch)
        raw_bs = forward_batch.batch_size
        if forward_batch.input_ids is not None:
            raw_num_tokens = forward_batch.input_ids.shape[0]
        elif forward_batch.input_embeds is not None:
            raw_num_tokens = forward_batch.input_embeds.shape[0]
        else:
            raw_num_tokens = 0
        registry = self._eager_registry
        registry.fill_from(
            forward_batch,
            raw_bs=raw_bs,
            padded_bs=raw_bs,
            raw_num_tokens=raw_num_tokens,
            padded_num_tokens=raw_num_tokens,
            pp_proxy_tensors=pp_proxy_tensors,
        )
        return registry.extract_buffer(
            padded_bs=raw_bs,
            padded_num_tokens=raw_num_tokens,
            forward_batch_template=forward_batch,
        )

    def execute(
        self, forward_batch: ForwardBatch, pp_proxy_tensors=None, **kwargs
    ) -> Any:
        mode = forward_batch.forward_mode
        if mode.is_decode():
            return self._execute_decode(forward_batch, pp_proxy_tensors)
        if mode.is_idle():
            return self._execute_idle(forward_batch, pp_proxy_tensors)
        if mode.is_extend(include_draft_extend_v2=True):
            return self._execute_extend(forward_batch, pp_proxy_tensors)
        raise ValueError(f"Invalid forward mode for eager runner: {mode}")

    def _resolve_decode_pdmux(
        self,
    ) -> Tuple[Any, contextlib.AbstractContextManager]:
        """Resolve the (attn_backend, forward_context) the eager decode forward
        runs under. PDmux selects a per-stream backend and publishes it via an
        active ForwardContext; non-pdmux uses attn_backend + the ambient ctx."""
        model_runner = self.model_runner
        if self.enable_pdmux:
            return model_runner.decode_attn_backend, forward_context(
                ForwardContext(attn_backend=model_runner.decode_attn_backend)
            )
        return model_runner.attn_backend, contextlib.nullcontext()

    def _execute_decode(
        self,
        forward_batch: ForwardBatch,
        pp_proxy_tensors=None,
    ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
        model_runner = self.model_runner
        enable_pdmux = self.enable_pdmux
        attn_backend, pdmux_ctx = self._resolve_decode_pdmux()
        if not enable_pdmux:
            forward_batch = self.load_batch(forward_batch, pp_proxy_tensors)
        if forward_batch.needs_forward_metadata_init():
            if hasattr(model_runner.model, "prepare_forward_batch"):
                # Prepare model-specific attention metadata before planning,
                # e.g. Moss-VL's prefill cross-attention custom mask.
                model_runner.model.prepare_forward_batch(forward_batch)
            attn_backend.init_forward_metadata(forward_batch)
        # FIXME: add pp_proxy_tensors arg to all models
        kwargs = model_runner._pp_kwargs(pp_proxy_tensors)

        ctx = (
            model_runner.device_timer.wrap(metadata={"category": "decode"})
            if model_runner.device_timer
            else contextlib.nullcontext()
        )

        with ctx, pdmux_ctx:
            return model_runner.model.forward(
                forward_batch.input_ids,
                forward_batch.positions,
                forward_batch,
                **kwargs,
            )

    def _execute_extend(
        self,
        forward_batch: ForwardBatch,
        pp_proxy_tensors=None,
    ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
        model_runner = self.model_runner
        kwargs = model_runner._extend_forward_kwargs(forward_batch, pp_proxy_tensors)

        if not self.enable_pdmux:
            forward_batch = self.load_batch(forward_batch, pp_proxy_tensors)

        if forward_batch.needs_forward_metadata_init():
            if hasattr(model_runner.model, "prepare_context_parallel_metadata_for_dcp"):
                # prepare kv cache buffer for dcp to gather kv cache
                forward_batch.attn_dcp_metadata = (
                    model_runner.model.prepare_context_parallel_metadata_for_dcp(
                        forward_batch.seq_lens,
                        forward_batch.extend_prefix_lens,
                        forward_batch.extend_prefix_lens_cpu,
                        forward_batch.extend_seq_lens,
                        forward_batch.req_pool_indices,
                        get_req_to_token_pool().req_to_token,
                        forward_batch.seq_lens_sum,
                        get_token_to_kv_pool().get_kv_buffer_shape()[0],
                        model_runner.kv_cache_dtype,
                        model_runner.device,
                        create_chunked_prefix_cache_kv_indices,
                    )
                )
            if hasattr(model_runner.model, "prepare_forward_batch"):
                # Prepare model-specific attention metadata before planning,
                # e.g. Moss-VL's prefill cross-attention custom mask.
                model_runner.model.prepare_forward_batch(forward_batch)
            model_runner.attn_backend.init_forward_metadata(forward_batch)

        cp_v2_active = is_cp_v2_active(forward_batch)
        forward_positions = forward_batch.positions
        if cp_v2_active:
            prepare_cp_forward(forward_batch)
            complete_hidden_states = kwargs.get("input_embeds")
            if complete_hidden_states is None:
                embed_layer = model_runner.model.get_input_embeddings()
                complete_hidden_states = embed_layer(forward_batch.input_ids)
            sharded_hidden_states, sharded_positions = cp_split_before_forward(
                complete_hidden_states,
                forward_batch.positions,
                forward_batch,
            )
            kwargs["input_embeds"] = sharded_hidden_states
            forward_positions = sharded_positions
        else:
            forward_batch.attn_cp_metadata = None

        category = (
            "target_verify"
            if forward_batch.forward_mode.is_target_verify()
            else "extend"
        )
        ctx = (
            model_runner.device_timer.wrap(metadata={"category": category})
            if model_runner.device_timer
            else contextlib.nullcontext()
        )
        with ctx:
            pcg_runner = model_runner.prefill_cuda_graph_runner
            if (
                _is_hip
                and pcg_runner is not None
                and not isinstance(pcg_runner, EagerRunner)
                and not cp_v2_active
            ):
                # HIP PCG eager fallback: enter the PCG context so Dynamo guards
                # and PCG-specific MoE/attention paths stay consistent.
                with (
                    enable_tc_piecewise_cuda_graph(),
                    set_tc_piecewise_forward_context(
                        forward_batch,
                        model_runner.attention_layers,
                        getattr(model_runner.model, "quant_config", None),
                        model_runner.moe_layers,
                        model_runner.moe_fusions,
                        dsa_indexers=model_runner.dsa_indexers,
                    ),
                ):
                    ret = model_runner.model.forward(
                        forward_batch.input_ids,
                        forward_positions,
                        forward_batch,
                        **kwargs,
                    )
            elif cp_v2_active:
                # CP-V2: drive .model directly to gather across CP ranks before logits.
                hidden_states = model_runner.model.model(
                    forward_batch.input_ids,
                    forward_positions,
                    forward_batch,
                    input_embeds=kwargs.get("input_embeds"),
                    pp_proxy_tensors=kwargs.get("pp_proxy_tensors"),
                )
                aux_hidden_states = None
                capture_aux_hidden_states = getattr(
                    model_runner.model, "capture_aux_hidden_states", False
                )
                if capture_aux_hidden_states:
                    hidden_states, aux_hidden_states = hidden_states
                if model_runner.model.pp_group.is_last_rank:
                    hidden_states = cp_gather_after_forward(
                        hidden_states,
                        forward_batch,
                        torch.cuda.current_stream(),
                    )
                    ret = model_runner.model.logits_processor(
                        forward_batch.input_ids,
                        hidden_states,
                        model_runner.model.lm_head,
                        forward_batch,
                        aux_hidden_states,
                    )
                elif capture_aux_hidden_states:
                    ret = hidden_states, aux_hidden_states
                else:
                    ret = hidden_states
            else:
                ret = model_runner.model.forward(
                    forward_batch.input_ids,
                    forward_positions,
                    forward_batch,
                    **kwargs,
                )
        return ret

    def _execute_idle(
        self, forward_batch: ForwardBatch, pp_proxy_tensors=None
    ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
        model_runner = self.model_runner
        # Padded idle (DP-attn MLP sync) needs metadata reinit; unpadded must
        # drop stale forward_metadata to avoid an SWA use-after-free on req_pool.
        if forward_batch.batch_size > 0:
            if not self.enable_pdmux:
                forward_batch = self.load_batch(forward_batch, pp_proxy_tensors)
            model_runner.attn_backend.init_forward_metadata(forward_batch)
        else:
            model_runner.attn_backend.forward_metadata = None

        kwargs = model_runner._pp_kwargs(pp_proxy_tensors)
        ctx = (
            model_runner.device_timer.wrap(metadata={"category": "idle"})
            if model_runner.device_timer
            else contextlib.nullcontext()
        )
        with ctx:
            return model_runner.model.forward(
                forward_batch.input_ids,
                forward_batch.positions,
                forward_batch,
                **kwargs,
            )