chore: import upstream snapshot with attribution

python/sglang/srt/model_executor/runner/__init__.py

@@ -0,0 +1,55 @@
+"""Phase-aware CUDA graph runners.
+
+One concrete runner per phase. Each runner owns its phase-specific
+shape semantics (decode → batch size; prefill → token count) and
+delegates capture/replay mechanics to a pluggable
+BaseCudaGraphBackend chosen via cuda_graph_config.
+
+Public API:
+  - BaseRunner — minimal abstract base shared by the cuda-graph runners
+    and the eager runner (shared __init__ + warmup + abstract
+    can_run_graph/load_batch/execute).
+  - BaseCudaGraphRunner — abstract cuda-graph base; bucket padding +
+    capture-loop scaffolding on top of BaseRunner.
+  - DecodeCudaGraphRunner — concrete decode-phase runner.
+  - PrefillCudaGraphRunner — concrete prefill-phase runner.
+  - EagerRunner — no-cuda-graph runner; runs model.forward live (the
+    eager dual of the cuda-graph runners), mode-dispatched over decode +
+    extend + idle.
+  - Buffer dataclasses, capture-mode flags, the global memory pool,
+    and the DeepEP adapter live in
+    sglang.srt.model_executor.runner_utils; they are
+    re-exported here for the EAGLE / multi-step draft cuda graph
+    runners that were authored against the legacy public surface.
+"""
+
+from sglang.srt.model_executor.runner.base_cuda_graph_runner import (  # noqa: F401
+    BaseCudaGraphRunner,
+    freeze_gc,
+    get_batch_sizes_to_capture,
+)
+from sglang.srt.model_executor.runner.base_runner import BaseRunner  # noqa: F401
+from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
+    DecodeCudaGraphRunner,
+)
+from sglang.srt.model_executor.runner.eager_runner import EagerRunner  # noqa: F401
+from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (  # noqa: F401
+    PrefillCudaGraphRunner,
+)
+from sglang.srt.model_executor.runner.shape_key import ShapeKey  # noqa: F401
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (  # noqa: F401
+    TCPCG_FAILURE_HINT,
+)
+from sglang.srt.model_executor.runner_utils import (  # noqa: F401
+    DecodeInputBuffers,
+    DeepEPCudaGraphRunnerAdapter,
+    PrefillInputBuffers,
+    _grouped_foreach_copy_,
+    _set_capture_lora_variant,
+    compile_in_capture_mode,
+    get_capture_lora_variant,
+    get_global_graph_memory_pool,
+    get_is_capture_mode,
+    model_capture_mode,
+    set_global_graph_memory_pool,
+)

python/sglang/srt/model_executor/runner/base_cuda_graph_runner.py

@@ -0,0 +1,158 @@
+# 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.
+# ==============================================================================
+"""Shared scaffolding for the prefill and decode CUDA graph runners."""
+
+from __future__ import annotations
+
+import bisect
+import gc
+import logging
+from abc import abstractmethod
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, List, Sequence, Tuple
+
+from sglang.srt.model_executor.runner.base_runner import BaseRunner
+from sglang.srt.runtime_context import get_flags, get_parallel
+from sglang.srt.utils import require_gathered_buffer
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
+        BaseCudaGraphBackend,
+    )
+
+logger = logging.getLogger(__name__)
+
+
+@contextmanager
+def freeze_gc(enable_cudagraph_gc: bool):
+    """Optimize garbage collection during CUDA graph capture.
+
+    Clean up first, then freeze remaining objects from being included in
+    future collections if GC is disabled during capture.
+    """
+    gc.collect()
+    should_freeze = not enable_cudagraph_gc
+    if should_freeze:
+        gc.freeze()
+    try:
+        yield
+    finally:
+        if should_freeze:
+            gc.unfreeze()
+            gc.collect()
+
+
+def get_batch_sizes_to_capture(
+    model_runner: ModelRunner, num_tokens_per_bs: int = 1
+) -> Tuple[List[int], List[int]]:
+    """Build the (capture_bs, compile_bs) lists for the decode runner.
+
+    Filters cuda_graph_config[decode].bs by attention-tp/cp alignment
+    constraints and clamps to req_to_token_pool.size.
+    """
+
+    server_args = model_runner.server_args
+    capture_bs = list(server_args.cuda_graph_config.decode.bs)
+    num_max_requests = model_runner.req_to_token_pool.size
+
+    mul_base = 1
+    if server_args.enable_two_batch_overlap:
+        mul_base *= 2
+        num_tokens_per_bs = 1
+
+    if require_gathered_buffer(server_args):
+        mul_base *= get_parallel().attn_tp_size
+
+    if mul_base % get_parallel().attn_cp_size != 0:
+        mul_base *= get_parallel().attn_cp_size
+
+    # pad `num_max_requests` to avoid being filtered out
+    num_max_requests = (num_max_requests + mul_base - 1) // mul_base * mul_base
+    if max(capture_bs) > num_max_requests:
+        # In some cases (e.g., with a small GPU or --max-running-requests), the #max-running-requests
+        # is very small. We add more values here to make sure we capture the maximum bs.
+        capture_bs += [num_max_requests]
+
+    # Model input token count = bs * num_tokens_per_bs; must be a multiple of attn_tp_size.
+    capture_bs = [bs for bs in capture_bs if bs * num_tokens_per_bs % mul_base == 0]
+    capture_bs = [bs for bs in capture_bs if bs <= num_max_requests]
+    capture_bs = list(sorted(set(capture_bs)))
+
+    assert len(capture_bs) > 0 and capture_bs[0] > 0, f"{capture_bs=}"
+    compile_bs = (
+        [bs for bs in capture_bs if bs <= server_args.torch_compile_max_bs]
+        if get_flags().capture.enable_torch_compile
+        else []
+    )
+    return capture_bs, compile_bs
+
+
+class BaseCudaGraphRunner(BaseRunner):
+    """Abstract base for phase-specific cuda-graph runners.
+
+    A subclass (DecodeCudaGraphRunner / PrefillCudaGraphRunner) owns one
+    phase and plugs in a BaseCudaGraphBackend that handles the
+    capture / replay mechanics. The runner orchestrates bucket
+    selection, static buffer population, attention metadata init,
+    replay dispatch, and output slicing.
+
+    Adds the capture/shape machinery on top of BaseRunner:
+      - capture_prepare(size, ...) — build the dummy ForwardBatch and
+        per-shape local state needed by capture_one_shape.
+      - capture() — one-time setup; iterates over shapes and calls
+        capture_one_shape for each.
+      - capture_one_shape(size, ...) — drive one model forward at this
+        shape into the backend's captured artifact.
+      - _pad_to_bucket(...) — round a raw shape up to the nearest captured
+        bucket.
+
+    Inherits from BaseRunner: __init__ and the abstract
+    can_run_graph / load_batch / execute.
+
+    Notes:
+      - buffers and backend are populated by the subclass before
+        capture(); the base only declares them.
+    """
+
+    # Subclasses populate before calling capture().
+    buffers: ForwardInputBuffers
+    backend: BaseCudaGraphBackend
+
+    @staticmethod
+    def _pad_to_bucket(raw_size: int, buckets: Sequence[int]) -> int:
+        """Return the smallest buckets[i] >= raw_size.
+
+        Caller's can_run_graph must reject raw_size > max(buckets) before
+        reaching load_batch; this assertion makes the contract
+        explicit (bisect_left returns len(buckets) when the value
+        exceeds all buckets, which would otherwise IndexError below
+        with no diagnostic).
+        """
+        assert raw_size <= buckets[-1], (
+            f"size {raw_size} exceeds max captured bucket {buckets[-1]}; "
+            f"can_run_graph should have rejected this batch"
+        )
+        index = bisect.bisect_left(buckets, raw_size)
+        return buckets[index]
+
+    @abstractmethod
+    def capture_prepare(self, size: int, *args, **kwargs) -> Any: ...
+
+    @abstractmethod
+    def capture(self) -> None: ...
+
+    @abstractmethod
+    def capture_one_shape(self, size: int, *args, **kwargs) -> Any: ...

python/sglang/srt/model_executor/runner/base_runner.py

@@ -0,0 +1,608 @@
+# 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.
+# ==============================================================================
+"""Base class shared by EagerRunner and BaseCudaGraphRunner."""
+
+from __future__ import annotations
+
+import inspect
+import logging
+from abc import ABC, abstractmethod
+from types import SimpleNamespace
+from typing import TYPE_CHECKING, Any, Optional, Tuple
+
+import torch
+
+from sglang.srt.batch_overlap.two_batch_overlap import TboCudaGraphRunnerPlugin
+from sglang.srt.compilation.torch_compile_decoration import set_torch_compile_config
+from sglang.srt.environ import envs
+from sglang.srt.layers import deep_gemm_wrapper
+from sglang.srt.layers.dp_attention import (
+    DpPaddingMode,
+    set_dp_buffer_len,
+    set_is_extend_in_batch,
+)
+from sglang.srt.model_executor.forward_batch_info import (
+    CaptureHiddenMode,
+    ForwardBatch,
+    ForwardMode,
+    NgramEmbeddingInfo,
+    PPProxyTensors,
+)
+from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
+from sglang.srt.model_executor.runner.flashinfer_autotune import (
+    run_flashinfer_autotune_forward,
+    should_run_flashinfer_autotune,
+)
+from sglang.srt.runtime_context import get_flags, get_parallel
+from sglang.srt.speculative.spec_info import create_dummy_verify_input
+from sglang.srt.utils import (
+    empty_context,
+    log_info_on_rank0,
+    require_attn_tp_gather,
+    require_gathered_buffer,
+    require_mlp_tp_gather,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+logger = logging.getLogger(__name__)
+
+
+def _allocate_decode_buffers(
+    *,
+    device: torch.device,
+    max_bs: int,
+    max_num_token: int,
+    hidden_size: int,
+    vocab_size: int,
+    dtype: torch.dtype,
+    dp_size: int,
+    pp_size: int,
+    is_encoder_decoder: bool,
+    require_mlp_tp_gather: bool,
+    seq_len_fill_value: int,
+    encoder_len_fill_value: int,
+    num_tokens_per_bs: int,
+    cache_loc_dtype: torch.dtype,
+    enable_mamba_track: bool,
+    ne_token_table: Optional[torch.Tensor] = None,
+    hc_hidden_size: Optional[int] = None,
+    pp_proxy_topk_size: Optional[int] = None,
+) -> SimpleNamespace:
+    """Allocate the FB-shared decode buffers."""
+    with torch.device(device):
+        input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
+        input_embeds = torch.zeros((max_num_token, hidden_size), dtype=dtype)
+        req_pool_indices = torch.zeros((max_bs,), dtype=torch.int64)
+        seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int64)
+        out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype)
+        positions = torch.zeros((max_num_token,), dtype=torch.int64)
+        mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64)
+        num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
+        custom_mask = torch.ones(
+            (max_bs * seq_len_fill_value + max_num_token) * num_tokens_per_bs,
+            dtype=torch.bool,
+        )
+        next_token_logits_buffer = torch.zeros(
+            (max_num_token, vocab_size),
+            dtype=torch.float,
+        )
+        mamba_track_indices = (
+            torch.zeros((max_bs,), dtype=torch.int64) if enable_mamba_track else None
+        )
+        mamba_track_mask = (
+            torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None
+        )
+
+        if pp_size > 1:
+            # mHC (e.g. DSV4) flattens residual into hidden_states (size = hc_hidden_size).
+            is_mhc = hc_hidden_size is not None
+            hs = hc_hidden_size if is_mhc else hidden_size
+            pp_proxy_tensors = {
+                "hidden_states": torch.zeros((max_bs, hs), dtype=dtype),
+            }
+            if not is_mhc:
+                pp_proxy_tensors["residual"] = torch.zeros(
+                    (max_bs, hidden_size), dtype=dtype
+                )
+            if pp_proxy_topk_size is not None:
+                pp_proxy_tensors["topk_indices"] = torch.zeros(
+                    (max_num_token, pp_proxy_topk_size), dtype=torch.int32
+                )
+        else:
+            pp_proxy_tensors = None
+
+        if is_encoder_decoder:
+            encoder_lens = torch.full(
+                (max_bs,), encoder_len_fill_value, dtype=torch.int32
+            )
+        else:
+            encoder_lens = None
+
+        if require_mlp_tp_gather:
+            global_num_tokens_gpu = torch.zeros((dp_size,), dtype=torch.int32)
+            global_num_tokens_for_logprob_gpu = torch.zeros(
+                (dp_size,), dtype=torch.int32
+            )
+        else:
+            global_num_tokens_gpu = torch.zeros((1,), dtype=torch.int32)
+            global_num_tokens_for_logprob_gpu = torch.zeros((1,), dtype=torch.int32)
+
+        ngram_embedding_info = (
+            NgramEmbeddingInfo(
+                token_table=ne_token_table,
+                column_starts=torch.zeros([max_bs], dtype=torch.int32),
+                req_lens=torch.ones([max_bs], dtype=torch.int32),
+                out_column_starts=torch.zeros([max_bs], dtype=torch.int32),
+                out_req_lens=torch.ones([max_bs], dtype=torch.int32),
+                skip_token_table_update=torch.zeros([max_bs], dtype=torch.bool),
+            )
+            if ne_token_table is not None
+            else None
+        )
+
+        if envs.SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE.get():
+            rids_int = torch.zeros((max_bs,), dtype=torch.int64)
+            bootstrap_room_ids_int = torch.full((max_bs,), -1, dtype=torch.int64)
+        else:
+            rids_int = None
+            bootstrap_room_ids_int = None
+
+    seq_lens_cpu = torch.full(
+        (max_bs,),
+        seq_len_fill_value,
+        dtype=torch.int64,
+        device="cpu",
+    )
+
+    return SimpleNamespace(
+        input_ids=input_ids,
+        input_embeds=input_embeds,
+        req_pool_indices=req_pool_indices,
+        seq_lens=seq_lens,
+        seq_lens_cpu=seq_lens_cpu,
+        out_cache_loc=out_cache_loc,
+        positions=positions,
+        mrope_positions=mrope_positions,
+        num_token_non_padded=num_token_non_padded,
+        custom_mask=custom_mask,
+        next_token_logits_buffer=next_token_logits_buffer,
+        mamba_track_indices=mamba_track_indices,
+        mamba_track_mask=mamba_track_mask,
+        encoder_lens=encoder_lens,
+        global_num_tokens_gpu=global_num_tokens_gpu,
+        global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
+        pp_proxy_tensors=pp_proxy_tensors,
+        ngram_embedding_info=ngram_embedding_info,
+        rids_int=rids_int,
+        bootstrap_room_ids_int=bootstrap_room_ids_int,
+    )
+
+
+class BaseRunner(ABC):
+    def __init__(self, model_runner: ModelRunner) -> None:
+        self.model_runner = model_runner
+        self.device = model_runner.device
+        self.device_module = torch.get_device_module(self.device)
+        self.tp_size = model_runner.server_args.tp_size
+        self.dp_size = model_runner.server_args.dp_size
+        self.pp_size = model_runner.server_args.pp_size
+        self.enable_pdmux = model_runner.server_args.enable_pdmux
+        self.enable_return_hidden_states = (
+            model_runner.server_args.enable_return_hidden_states
+        )
+        self.attn_tp_size = get_parallel().attn_tp_size
+        self.attn_tp_rank = get_parallel().attn_tp_rank
+        self.tbo_plugin = TboCudaGraphRunnerPlugin()
+
+    def warmup(self) -> None:
+        """Run kernel warmup + autotune once, gated by mr._kernel_warmed_up."""
+        mr = self.model_runner
+        if getattr(mr, "_kernel_warmed_up", False):
+            return
+        mr._kernel_warmed_up = True
+
+        if mr.device != "cuda":
+            return
+
+        self._pre_initialize_flashinfer_allreduce_workspace()
+
+        if should_run_flashinfer_autotune(self.model_runner):
+            buffers, batch_size = self._autotune_buffers()
+            assert (
+                buffers is not None
+            ), "_autotune_buffers() must return a reusable buffer set for autotune"
+            self._flashinfer_autotune(buffers=buffers, batch_size=batch_size)
+
+        if (
+            envs.SGLANG_PP_PARALLEL_DEEPGEMM_WARMUP.get()
+            and deep_gemm_wrapper.ENABLE_JIT_DEEPGEMM
+            and mr.pp_size > 1
+            and not mr.spec_algorithm.is_speculative()
+        ):
+            from sglang.srt.layers.deep_gemm_wrapper.compile_utils import (
+                pp_parallel_deep_gemm_warmup,
+            )
+
+            pp_parallel_deep_gemm_warmup(self)
+
+    def _pre_initialize_flashinfer_allreduce_workspace(self):
+        """Allocate flashinfer allreduce workspaces; must run before CG capture
+        to keep broadcasts/barriers outside the capture context (else deadlock
+        with custom_all_reduce.register_graph_buffers).
+        """
+        mr = self.model_runner
+        if mr.server_args.flashinfer_allreduce_fusion_backend is None:
+            return
+
+        from sglang.srt.layers.communicator import FUSE_ALLREDUCE_MAX_BATCH_SIZE
+        from sglang.srt.layers.flashinfer_comm_fusion import pre_initialize_workspaces
+
+        pre_initialize_workspaces(
+            max_token_num=FUSE_ALLREDUCE_MAX_BATCH_SIZE,
+            hidden_dim=mr.model_config.hidden_size,
+            dtype=mr.dtype,
+        )
+
+    def _flashinfer_autotune(self, *, buffers, batch_size):
+        """Run flashinfer autotune.
+
+        buffers / batch_size: a prepared static decode-buffer set and its bs,
+        reused for the dummy forward instead of allocating a throwaway set.
+        Supplied by warmup() (the decode runner's captured buffers when a graph
+        runner exists; a freshly-allocated dummy set in the eager path).
+        """
+        mr = self.model_runner
+        canary_run_ctx = (
+            c.with_active_single_forward_manager(0)
+            if (c := mr.canary_manager) is not None
+            else empty_context()
+        )
+
+        def forward_fn():
+            self._dummy_run(
+                batch_size=batch_size,
+                buffers=buffers,
+                run_ctx=canary_run_ctx,
+            )
+
+        run_flashinfer_autotune_forward(self.model_runner, forward_fn, skip_logits=True)
+
+    def _alloc_dummy_decode_buffers(self, max_bs: int, *, num_tokens_per_bs: int = 1):
+        """Allocate one static decode-buffer set for a dummy forward, sized to
+        (max_bs, max_bs * num_tokens_per_bs).
+
+        The PP-parallel DeepGEMM warmup sweeps batch sizes far larger than any
+        runner's max_bs (up to ~n_sms*block_m), so no pre-allocated runner buffer
+        set fits; it builds one here and hands it to _dummy_run (reused across the
+        sweep; _dummy_run slices it per shape). Eager FlashInfer autotune also
+        allocates decode-shaped scratch buffers here. Decode cuda-graph autotune
+        reuses the captured runner buffers instead.
+        """
+        mr = self.model_runner
+        return _allocate_decode_buffers(
+            device=mr.device,
+            max_bs=max_bs,
+            max_num_token=max_bs * num_tokens_per_bs,
+            hidden_size=mr.model_config.hidden_size,
+            vocab_size=mr.model_config.vocab_size,
+            dtype=mr.model_config.dtype,
+            dp_size=mr.server_args.dp_size,
+            pp_size=mr.server_args.pp_size,
+            is_encoder_decoder=mr.model_config.is_encoder_decoder,
+            require_mlp_tp_gather=require_mlp_tp_gather(mr.server_args),
+            seq_len_fill_value=mr.attn_backend.get_cuda_graph_seq_len_fill_value(),
+            encoder_len_fill_value=(
+                getattr(mr.model_config.hf_config, "max_source_positions", 0)
+                if mr.model_config.is_encoder_decoder
+                else 0
+            ),
+            num_tokens_per_bs=num_tokens_per_bs,
+            cache_loc_dtype=torch.int64,
+            enable_mamba_track=False,
+            ne_token_table=mr.token_table if mr.use_ngram_embedding else None,
+            hc_hidden_size=getattr(mr.model_config, "hc_hidden_size", None),
+            pp_proxy_topk_size=mr.get_pp_proxy_topk_size(),
+        )
+
+    def _dummy_run(
+        self,
+        batch_size: int,
+        run_ctx=None,
+        forward_mode_override: Optional[ForwardMode] = None,
+        *,
+        buffers,
+    ):
+        """Run a dummy forward pass for warmup/profiling.
+
+        forward_mode_override forces EXTEND/DECODE regardless of
+        is_generation (used by the PP-parallel DeepGEMM warmup).
+
+        buffers: a prepared static buffer set (or lightweight adapter exposing
+        the same fields), sized >= this dummy shape, which _dummy_run slices to
+        (batch_size, num_tokens). The caller owns the shape and the allocation --
+        the flashinfer autotune reuses an existing runner's buffers via
+        _autotune_buffers (the eager input registry, or the decode cuda-graph
+        runner's captured buffers); the PP-DeepGEMM warmup builds one via
+        _alloc_dummy_decode_buffers. _dummy_run never allocates and never re-pads
+        (autotune must run at the reused shape; the PP warmup pre-pads and sizes
+        its buffer to match). next_token_logits_buffer is optional -- a live
+        autotune forward returns logits fresh, so the eager-reuse path passes
+        None (only the PP warmup set still carries one).
+        """
+        mr = self.model_runner
+        if forward_mode_override is not None:
+            capture_forward_mode = forward_mode_override
+        elif mr.is_generation:
+            capture_forward_mode = ForwardMode.DECODE
+        else:
+            capture_forward_mode = ForwardMode.EXTEND
+        capture_hidden_mode = CaptureHiddenMode.NULL
+        num_tokens_per_bs = 1
+        if mr.spec_algorithm.is_speculative():
+            if mr.is_draft_worker:
+                if not mr.spec_algorithm.supports_target_verify_for_draft():
+                    raise RuntimeError("This should not happen")
+            capture_forward_mode = ForwardMode.TARGET_VERIFY
+            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
+                )
+            )
+
+        if mr.server_args.enable_return_hidden_states:
+            capture_hidden_mode = CaptureHiddenMode.FULL
+
+        num_tokens = batch_size * num_tokens_per_bs
+
+        # Caller owns the shape: passes a static buffer >= the dummy shape; no
+        # allocation, no re-padding (would overflow the reused buffers).
+        assert (
+            buffers is not None
+            and num_tokens <= buffers.input_ids.shape[0]
+            and batch_size <= buffers.seq_lens.shape[0]
+        ), (
+            f"_dummy_run needs a static buffer >= (num_tokens={num_tokens}, "
+            f"batch_size={batch_size}); got "
+            + (
+                "None"
+                if buffers is None
+                else f"(input_ids={buffers.input_ids.shape[0]}, "
+                f"seq_lens={buffers.seq_lens.shape[0]})"
+            )
+        )
+
+        seq_len_fill_value = mr.attn_backend.get_cuda_graph_seq_len_fill_value()
+
+        if get_flags().capture.enable_torch_compile:
+            set_torch_compile_config()
+            should_disable_torch_compile = not getattr(
+                mr.model, "_can_torch_compile", True
+            )
+            if should_disable_torch_compile:
+                log_info_on_rank0(
+                    logger,
+                    "Transformers backend model reports it is not torch.compile "
+                    "compatible (e.g. dynamic rope scaling). Disabling torch.compile.",
+                )
+                get_flags().capture.enable_torch_compile = False
+
+        # NOTE: aux hidden state capture (eagle3/dflash) is already
+        # configured by init_aux_hidden_state_capture() in initialize().
+
+        require_mlp_tp_gather_ = require_mlp_tp_gather(mr.server_args)
+        if require_gathered_buffer(mr.server_args):
+            assert require_mlp_tp_gather_ or require_attn_tp_gather(mr.server_args)
+
+        input_ids = buffers.input_ids[:num_tokens]
+        positions = buffers.positions[:num_tokens]
+        out_cache_loc = buffers.out_cache_loc[:num_tokens]
+        # Eager-reuse drops the logits buffer; only buffer sets that carry one slice it.
+        next_token_logits_buffer = (
+            buffers.next_token_logits_buffer[:num_tokens]
+            if buffers.next_token_logits_buffer is not None
+            else None
+        )
+        mrope_positions = buffers.mrope_positions[:, :num_tokens]
+        req_pool_indices = buffers.req_pool_indices[:batch_size]
+        seq_lens = buffers.seq_lens[:batch_size]
+        seq_lens_cpu = buffers.seq_lens_cpu[:batch_size]
+        encoder_lens = (
+            buffers.encoder_lens[:batch_size]
+            if buffers.encoder_lens is not None
+            else None
+        )
+
+        buffers.num_token_non_padded[...] = num_tokens
+
+        # For extend mode
+        if capture_forward_mode == ForwardMode.EXTEND:
+            extend_prefix_lens_cpu = [0] * batch_size
+            extend_seq_lens_cpu = [seq_len_fill_value] * batch_size
+            extend_num_tokens = num_tokens
+            extend_seq_lens = torch.full(
+                (batch_size,), seq_len_fill_value, dtype=torch.int32, device=mr.device
+            )
+            extend_prefix_lens = torch.zeros(
+                (batch_size,), dtype=torch.int32, device=mr.device
+            )
+            extend_start_loc = torch.arange(
+                0, num_tokens, num_tokens_per_bs, dtype=torch.int32, device=mr.device
+            )
+        else:
+            extend_prefix_lens_cpu = None
+            extend_seq_lens_cpu = None
+            extend_num_tokens = None
+            extend_seq_lens = None
+            extend_prefix_lens = None
+            extend_start_loc = None
+
+        if mr.server_args.pp_size > 1:
+            # PP0 already cp-split hidden_states before send.
+            pp_hidden_tokens = num_tokens
+            if (
+                capture_forward_mode == ForwardMode.EXTEND
+                and mr.pp_rank != 0
+                and mr.attn_cp_size > 1
+            ):
+                pp_hidden_tokens = num_tokens // mr.attn_cp_size
+            pp_proxy_tensors = PPProxyTensors(
+                {k: v[:pp_hidden_tokens] for k, v in buffers.pp_proxy_tensors.items()}
+            )
+
+        if require_mlp_tp_gather_:
+            global_num_tokens_cpu = [num_tokens] * mr.server_args.dp_size
+        elif require_attn_tp_gather(mr.server_args):
+            global_num_tokens_cpu = [num_tokens]
+        else:
+            global_num_tokens_cpu = None
+
+        if global_num_tokens_cpu is not None:
+            global_dp_buffer_len = sum(global_num_tokens_cpu)
+            num_tokens_tensor = torch.tensor(
+                global_num_tokens_cpu, dtype=torch.int32, device=mr.device
+            )
+            buffers.global_num_tokens_gpu.copy_(num_tokens_tensor)
+            buffers.global_num_tokens_for_logprob_gpu.copy_(num_tokens_tensor)
+        else:
+            global_dp_buffer_len = None
+            global_num_tokens_cpu = None
+
+        spec_info = create_dummy_verify_input(
+            mr.spec_algorithm,
+            mr.server_args,
+            buffers.custom_mask,
+            num_tokens_per_bs,
+            mr.is_draft_worker,
+        )
+        if spec_info is not None and (
+            mr.spec_algorithm.is_eagle() or mr.spec_algorithm.is_standalone()
+        ):
+            # MTP models (e.g. deepseek_nextn) read spec_info.hidden_states
+            # during forward; provide a dummy so warmup doesn't crash.
+            spec_info.hidden_states = torch.zeros(
+                (num_tokens, mr.model_config.hidden_size),
+                dtype=mr.dtype,
+                device=mr.device,
+            )
+        if capture_hidden_mode != CaptureHiddenMode.FULL:
+            capture_hidden_mode = (
+                spec_info.capture_hidden_mode if spec_info else CaptureHiddenMode.NULL
+            )
+
+        if mr.server_args.enable_lora:
+            lora_ids = [None] * batch_size
+        else:
+            lora_ids = None
+
+        forward_batch = ForwardBatch(
+            forward_mode=capture_forward_mode,
+            batch_size=batch_size,
+            input_ids=input_ids,
+            req_pool_indices=req_pool_indices,
+            seq_lens=seq_lens,
+            seq_lens_cpu=seq_lens_cpu,
+            next_token_logits_buffer=next_token_logits_buffer,
+            orig_seq_lens=seq_lens,
+            out_cache_loc=out_cache_loc,
+            seq_lens_sum=seq_lens.sum().item(),
+            encoder_lens=encoder_lens,
+            return_logprob=False,
+            positions=positions,
+            extend_num_tokens=extend_num_tokens,
+            extend_seq_lens=extend_seq_lens,
+            extend_prefix_lens=extend_prefix_lens,
+            extend_start_loc=extend_start_loc,
+            extend_prefix_lens_cpu=extend_prefix_lens_cpu,
+            extend_seq_lens_cpu=extend_seq_lens_cpu,
+            global_num_tokens_gpu=buffers.global_num_tokens_gpu,
+            global_num_tokens_cpu=global_num_tokens_cpu,
+            global_num_tokens_for_logprob_gpu=buffers.global_num_tokens_for_logprob_gpu,
+            dp_padding_mode=DpPaddingMode.get_default_mode_in_cuda_graph(),
+            global_dp_buffer_len=global_dp_buffer_len,
+            mrope_positions=mrope_positions,
+            spec_algorithm=mr.spec_algorithm,
+            spec_info=spec_info,
+            capture_hidden_mode=capture_hidden_mode,
+            num_token_non_padded=buffers.num_token_non_padded,
+            global_forward_mode=capture_forward_mode,
+            lora_ids=lora_ids,
+        )
+        if buffers.ngram_embedding_info is not None:
+            forward_batch.ngram_embedding_info = buffers.ngram_embedding_info.slice(
+                batch_size
+            )
+
+        if lora_ids is not None:
+            mr.lora_manager.prepare_lora_batch(forward_batch)
+
+        mr.attn_backend.init_forward_metadata(forward_batch)
+
+        def run_once():
+            forward_batch.dp_local_start_pos = forward_batch.dp_local_num_tokens = None
+            set_dp_buffer_len(
+                global_dp_buffer_len,
+                num_tokens,
+                forward_batch.dp_padding_mode.is_max_len(),
+                global_num_tokens_cpu,
+            )
+            set_is_extend_in_batch(False)
+
+            kwargs = {}
+            if (
+                mr.server_args.pp_size > 1
+                and "pp_proxy_tensors" in inspect.signature(mr.model.forward).parameters
+            ):
+                kwargs["pp_proxy_tensors"] = PPProxyTensors(
+                    {k: v.clone() for k, v in pp_proxy_tensors.tensors.items()}
+                )
+            if not mr.is_generation:
+                kwargs["get_embedding"] = True
+
+            logits_output_or_pp_proxy_tensors = mr.model.forward(
+                input_ids,
+                forward_batch.positions,
+                forward_batch,
+                **kwargs,
+            )
+            return logits_output_or_pp_proxy_tensors
+
+        torch.get_device_module(mr.device).synchronize()
+        mr.tp_group.barrier()
+        with forward_context(ForwardContext(attn_backend=mr.attn_backend)):
+            with torch.inference_mode(), run_ctx or empty_context():
+                run_once()
+
+    def _autotune_buffers(self) -> Tuple[Optional[Any], Optional[int]]:
+        """Return (buffers, bs) for the autotune dummy forward to reuse; the
+        EagerRunner and DecodeCudaGraphRunner override this."""
+        return None, None
+
+    @abstractmethod
+    def can_run_graph(self, forward_batch: ForwardBatch) -> bool: ...
+
+    @abstractmethod
+    def load_batch(
+        self,
+        forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any: ...
+
+    @abstractmethod
+    def execute(
+        self,
+        forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any: ...

python/sglang/srt/model_executor/runner/eager_runner.py

@@ -0,0 +1,414 @@
+# 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,
+            )

python/sglang/srt/model_executor/runner/flashinfer_autotune.py

@@ -0,0 +1,224 @@
+# 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.
+# ==============================================================================
+from __future__ import annotations
+
+import contextlib
+import datetime
+import hashlib
+import logging
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable, Optional
+
+import torch
+
+from sglang.srt.environ import envs
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+    from sglang.srt.model_executor.runner.base_runner import BaseRunner
+
+logger = logging.getLogger(__name__)
+
+
+def should_run_flashinfer_autotune(
+    model_runner: ModelRunner, *, for_speculative_draft: bool = False
+) -> bool:
+    """Check if flashinfer autotune should be run."""
+    mr = model_runner
+    if mr.device != "cuda":
+        return False
+    if mr.server_args.disable_flashinfer_autotune:
+        return False
+
+    # CuteDSL v1 (cutedsl runner + deepep a2a) bypasses MoeRunner and must not
+    # be autotuned -- its _dummy_run would dispatch more tokens per rank than
+    # SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK, tripping a DeepEP assert.
+    # Read server_args directly to avoid depending on initialize_moe_config()
+    # having already populated the MoE backend globals.
+    if (
+        mr.server_args.moe_runner_backend == "flashinfer_cutedsl"
+        and mr.server_args.moe_a2a_backend == "deepep"
+    ):
+        return False
+
+    backend_str = mr.server_args.moe_runner_backend
+
+    # TODO smor- support other cases for flashinfer autotune, such as, mamba backend
+
+    moe_needs_autotune = backend_str in [
+        "flashinfer_trtllm",
+        "flashinfer_trtllm_routed",
+        "flashinfer_mxfp4",
+        "flashinfer_cutedsl",
+        "flashinfer_cutlass",
+    ]
+
+    from sglang.srt.layers.quantization.fp4_utils import (
+        get_fp4_gemm_runner_backend,
+    )
+
+    model_quantization = mr.model_config.quantization
+    model_uses_fp4 = model_quantization in (
+        "modelopt_fp4",
+        "modelopt_mixed",
+    )
+    fp4_gemm_needs_autotune = model_uses_fp4 and (
+        get_fp4_gemm_runner_backend().is_flashinfer_cutlass()
+        or get_fp4_gemm_runner_backend().is_flashinfer_cutedsl()
+    )
+
+    from sglang.srt.layers.quantization.fp8_utils import (
+        get_fp8_gemm_runner_backend,
+    )
+    from sglang.srt.utils import is_sm100_supported
+
+    model_uses_modelopt_fp8 = model_quantization in (
+        "modelopt",
+        "modelopt_fp8",
+        "modelopt_mixed",
+    )
+    # Online MXFP8 (microscaling) linears dispatch to flashinfer's
+    # ``mm_mxfp8``, which the flashinfer fp8 autotune dummy run does not
+    # exercise correctly -- it triggers an illegal memory access inside the
+    # mxfp8 cutlass cubin. The mxfp8 gemm is fixed-config and needs no
+    # tuning, so skip autotune for these models.
+    model_uses_mxfp8 = "mxfp8" in (model_quantization or "")
+    fp8_gemm_needs_autotune = not model_uses_mxfp8 and (
+        get_fp8_gemm_runner_backend().is_flashinfer_cutlass()
+        or (model_uses_modelopt_fp8 and is_sm100_supported())
+    )
+
+    if not (moe_needs_autotune or fp4_gemm_needs_autotune or fp8_gemm_needs_autotune):
+        return False
+
+    if torch.cuda.get_device_capability()[0] < 9:
+        return False
+
+    if mr.spec_algorithm.is_speculative():
+        return mr.is_draft_worker if for_speculative_draft else not mr.is_draft_worker
+
+    return True
+
+
+def flashinfer_autotune_cache_path(model_runner: ModelRunner) -> Path:
+    import flashinfer
+
+    mr = model_runner
+    major, minor = torch.cuda.get_device_capability(mr.device)
+    arch = f"sm{major}{minor}"
+    flashinfer_version = getattr(flashinfer, "__version__", "unknown")
+
+    server_args = mr.server_args
+    model_key_parts = [
+        str(server_args.model_path),
+        str(mr.dtype),
+        str(server_args.quantization),
+        str(server_args.moe_runner_backend),
+        str(mr.tp_size),
+        str(mr.pp_size),
+        str(mr.dp_size),
+        str(mr.moe_ep_size),
+        str(mr.model_config.hf_config.__class__.__name__),
+    ]
+    if mr.is_draft_worker:
+        model_key_parts.append(f"draft_quant={mr.model_config.quantization}")
+    model_key = "|".join(model_key_parts)
+    cache_key = hashlib.sha256(model_key.encode()).hexdigest()[:16]
+    cache_dir = (
+        Path(envs.SGLANG_CACHE_DIR.get())
+        / "flashinfer"
+        / "autotune"
+        / flashinfer_version
+        / arch
+        / cache_key
+    )
+    cache_dir.mkdir(parents=True, exist_ok=True)
+    return cache_dir / f"rank_tp{mr.tp_rank}_pp{mr.pp_rank}_dp{mr.dp_rank or 0}.json"
+
+
+@contextlib.contextmanager
+def flashinfer_autotune_context(model_runner: ModelRunner, *, skip_logits: bool):
+    from flashinfer.autotuner import autotune
+
+    mr = model_runner
+    cache_path = flashinfer_autotune_cache_path(mr)
+    if envs.SGLANG_FLASHINFER_AUTOTUNE_CACHE.get():
+        autotune_cache = cache_path
+        logger.info("Running FlashInfer autotune with cache: %s", autotune_cache)
+    else:
+        timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
+        runs_dir = cache_path.parent / "runs"
+        runs_dir.mkdir(parents=True, exist_ok=True)
+        autotune_cache = runs_dir / f"{cache_path.stem}.{timestamp}{cache_path.suffix}"
+        logger.info(
+            "Running FlashInfer autotune (cache reuse DISABLED via "
+            "SGLANG_FLASHINFER_AUTOTUNE_CACHE=0); writing fresh result to: %s",
+            autotune_cache,
+        )
+
+    # Run warmup on the non-default stream to avoid NCCL 2.29+ cudaMemcpyBatchAsync
+    # calls on default stream (unsupported by CUDA) when --enable-symm-mem is used.
+    mr.forward_stream.wait_stream(torch.cuda.current_stream())
+    with torch.get_device_module(mr.device).stream(mr.forward_stream):
+        maybe_skip_logits = contextlib.nullcontext()
+        if skip_logits:
+            from sglang.srt.layers.logits_processor import autotune_dummy_run_mode
+
+            maybe_skip_logits = autotune_dummy_run_mode()
+        with torch.inference_mode(), autotune(
+            True, cache=str(autotune_cache)
+        ), maybe_skip_logits:
+            yield
+    torch.cuda.current_stream().wait_stream(mr.forward_stream)
+    logger.info("FlashInfer autotune completed.")
+
+
+def run_flashinfer_autotune_forward(
+    model_runner: ModelRunner, forward_fn: Callable[[], None], *, skip_logits: bool
+) -> None:
+    """Run flashinfer autotune forward."""
+    with flashinfer_autotune_context(model_runner, skip_logits=skip_logits):
+        forward_fn()
+
+
+def maybe_flashinfer_autotune_speculative_draft(
+    runner: BaseRunner,
+    forward_fn: Callable[[], None],
+    *,
+    post_warmup_hook: Optional[Callable[[], None]] = None,
+    skip_logits: bool = False,
+) -> None:
+    """Run speculative draft flashinfer autotune."""
+    mr = runner.model_runner
+    phase_key = f"{runner.__class__.__module__}.{runner.__class__.__qualname__}"
+    tuned_phases = getattr(mr, "_flashinfer_spec_draft_autotuned_phases", None)
+    if tuned_phases is None:
+        tuned_phases = set()
+        mr._flashinfer_spec_draft_autotuned_phases = tuned_phases
+    if phase_key in tuned_phases:
+        return
+    if (
+        not mr.spec_algorithm.is_speculative()
+        or not mr.is_draft_worker
+        or not should_run_flashinfer_autotune(mr, for_speculative_draft=True)
+    ):
+        return
+
+    def run_and_reset():
+        forward_fn()
+        if post_warmup_hook is not None:
+            post_warmup_hook()
+
+    run_flashinfer_autotune_forward(mr, run_and_reset, skip_logits=skip_logits)
+    tuned_phases.add(phase_key)

python/sglang/srt/model_executor/runner/shape_key.py

@@ -0,0 +1,36 @@
+# 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.
+# ==============================================================================
+"""ShapeKey — typed identifier for one captured CUDA-graph shape."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Optional
+
+
+@dataclass(frozen=True)
+class ShapeKey:
+    """Identifies one captured CUDA-graph shape across all runners.
+
+    size: the per-phase capture size — what the runner iterates over.
+        - prefill: num_tokens
+        - decode:  bs
+    stream_idx:   pdmux stream index, or None for single-stream runners.
+    variant_label: LoRA-variant label ("lora" / "nolora"), or None
+        for runners that don't record per-variant graphs.
+    """
+
+    size: int
+    stream_idx: Optional[int] = None
+    variant_label: Optional[str] = None