sgl-project--sglang/python/sglang/srt/model_executor/cpu_graph_runner.py

# Copyright 2023-2024 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.
# ==============================================================================
"""Run the model with cpu torch compile."""

# The implementation of CPUGraphRunner follows the CudaGraphRunner

from __future__ import annotations

import bisect
import logging
from contextlib import contextmanager
from typing import TYPE_CHECKING, Callable, Optional, Union

import psutil
import torch
import tqdm

from sglang.srt.distributed.parallel_state import GroupCoordinator
from sglang.srt.layers.logits_processor import LogitsProcessorOutput
from sglang.srt.model_executor.forward_batch_info import (
    CaptureHiddenMode,
    ForwardBatch,
    ForwardMode,
    PPProxyTensors,
    enable_num_token_non_padded,
)
from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
from sglang.srt.model_executor.runner_utils.capture_mode import model_capture_mode
from sglang.srt.runtime_context import get_flags, get_parallel
from sglang.srt.utils import (
    empty_context,
    log_info_on_rank0,
    require_attn_tp_gather,
    require_gathered_buffer,
    require_mlp_sync,
    require_mlp_tp_gather,
)
from sglang.srt.utils.patch_torch import monkey_patch_torch_compile

logger = logging.getLogger(__name__)

# ---------------------------------------------------------------------------
# skip_cross_attention capture-mode helpers (CPU graph only)
# ---------------------------------------------------------------------------
# When CPUGraphRunner captures two graphs per batch size (one with cross-
# attention, one without), it uses this context variable so that
# encoder-decoder models (e.g. mllama) receive a compile-time-constant value
# for skip_cross_attention instead of a data-dependent branch to avoid recompiles.

_capture_skip_cross_attention: Optional[bool] = None


def get_capture_skip_cross_attention() -> Optional[bool]:
    """Return the active skip_cross_attention override, or None if not set."""
    return _capture_skip_cross_attention


@contextmanager
def capture_with_skip_cross_attention(skip: bool):
    """Pin skip_cross_attention to *skip* for the duration of the context."""
    global _capture_skip_cross_attention
    previous = _capture_skip_cross_attention
    _capture_skip_cross_attention = skip
    try:
        yield
    finally:
        _capture_skip_cross_attention = previous


if TYPE_CHECKING:
    from sglang.srt.model_executor.model_runner import ModelRunner


@contextmanager
def patch_model(
    model: torch.nn.Module,
    enable_compile: bool,
    num_tokens: int,
    tp_group: GroupCoordinator,
):
    """Patch the model to make it compatible with torch.compile"""
    backup_ca_comm = None

    try:
        if enable_compile:
            backup_ca_comm = tp_group.ca_comm
            # Use custom-allreduce here.
            # We found the custom allreduce is much faster than the built-in allreduce in torch,
            # even with ENABLE_INTRA_NODE_COMM=1.
            # tp_group.ca_comm = None
            yield torch.compile(
                torch.no_grad()(model.forward),
                dynamic=False,
            )
        else:
            yield model.forward
    finally:
        if enable_compile:
            tp_group.ca_comm = backup_ca_comm


def set_torch_compile_config():
    import torch._dynamo.config
    import torch._inductor.config

    torch._inductor.config.fx_graph_cache = True  # Experimental feature to reduce compilation times, will be on by default in future
    torch._inductor.config.freezing = True
    torch._dynamo.config.accumulated_cache_size_limit = 1024
    if hasattr(torch._dynamo.config, "cache_size_limit"):
        torch._dynamo.config.cache_size_limit = 1024
    register_inductor_fallback_ops()
    monkey_patch_torch_compile()


def get_batch_sizes_to_capture(model_runner: ModelRunner):
    # torch compile speeds up decoding by reducing python overhead on CPU
    server_args = model_runner.server_args
    # Reuse cuda_graph_config[decode].bs here.
    # Users can customize the batch sizes supported by cpu_graph, such as:
    # --cuda-graph-bs-decode 1 2 4 8 16
    capture_bs = server_args.cuda_graph_config.decode.bs
    assert (
        max(capture_bs) <= server_args.torch_compile_max_bs
    ), f"{capture_bs=}, {server_args.torch_compile_max_bs=}"
    capture_bs = [bs for bs in capture_bs if bs <= model_runner.req_to_token_pool.size]
    capture_bs = list(sorted(set(capture_bs)))
    assert len(capture_bs) > 0 and capture_bs[0] > 0, f"{capture_bs=}"
    return capture_bs


_CPU_COMPILE_FAKE_OPS: set[str] = set()


def register_cpu_compile_fake(op_name: str):
    _CPU_COMPILE_FAKE_OPS.add(op_name)
    return torch.library.register_fake(f"sgl_kernel::{op_name}")


def register_inductor_fallback_ops():
    from torch._inductor.lowering import lowerings, make_fallback

    sgl_kernel_ops = torch.ops.sgl_kernel
    for op_name in sorted(_CPU_COMPILE_FAKE_OPS):
        try:
            op = getattr(getattr(sgl_kernel_ops, op_name), "default")
        except AttributeError:
            continue
        if op not in lowerings:
            make_fallback(op, warn=False)


def register_fake_ops(tp_size: int):
    """
    Registers fake/meta implementations for all custom sgl_kernel CPU operators
    using torch.library.register_fake to support torch.compile
    """

    none_return_ops = [
        "shm_allreduce",
        "bmm_cpu",
        "fused_add_rmsnorm_cpu",
        "decode_attention_cpu",
        "extend_attention_cpu",
        "gemma_fused_add_rmsnorm_cpu",
        "layernorm_cpu",
        "fused_add_layernorm_cpu",
    ]
    for op in none_return_ops:

        @register_cpu_compile_fake(op)
        def _(*args, **kwargs):
            return

    for op in [
        "rmsnorm_cpu",
        "l2norm_cpu",
        "fused_experts_cpu",
        "fused_rmsnorm_gated_cpu",
        "shared_expert_cpu",
        "causal_conv1d_update_cpu",
        "causal_conv1d_fwd_cpu",
        "gemma_rmsnorm_cpu",
        "gemma3_rmsnorm_cpu",
        "gemma4_rmsnorm_cpu",
    ]:

        @register_cpu_compile_fake(op)
        def _(input, *args, **kwargs):
            return torch.empty_like(input)

    @register_cpu_compile_fake("shm_allgather")
    def _(data, dim):
        return torch.cat([data] * tp_size, dim=dim)

    @register_cpu_compile_fake("qkv_proj_with_rope")
    def _(
        hidden_states,
        q_a_proj_weight,
        q_b_proj_weight,
        kv_a_proj_weight,
        w_kc,
        q_a_layernorm_weight,
        kv_a_layernorm_weight,
        positions,
        cos_sin_cache,
        eps,
        use_int8_w8a8,
        use_fp8_w8a16,
        q_a_proj_scale,
        q_b_proj_scale,
        kv_a_proj_scale,
        is_vnni,
        block_size,
    ):
        num_seqs = hidden_states.shape[0]
        num_heads = w_kc.shape[0]
        kv_lora_rank = w_kc.shape[1]
        qk_rope_head_dim = kv_a_proj_weight.shape[0] - kv_lora_rank
        q_input = torch.empty(
            num_seqs,
            num_heads,
            kv_lora_rank + qk_rope_head_dim,
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )
        k_input = torch.empty(
            num_seqs,
            1,
            kv_lora_rank + qk_rope_head_dim,
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )
        v_input = k_input.narrow(-1, 0, kv_lora_rank)
        return q_input, k_input, v_input

    @register_cpu_compile_fake("rotary_embedding_cpu")
    def _(positions, query, key, head_size, cos_sin_cache, is_neox):
        if query.ndim == 2:
            return query, key
        else:
            return torch.empty_like(query), torch.empty_like(key)

    @register_cpu_compile_fake("apply_rotary_pos_emb_cpu")
    def _(query, key, cos, sin):
        return query, key

    @register_cpu_compile_fake("multimodal_rotary_embedding_cpu")
    def _(
        positions,
        query,
        key,
        head_size,
        cos_sin_cache,
        mrope_section,
        mrope_interleaved,
        is_neox,
    ):
        return query, key

    @register_cpu_compile_fake("qkv_proj_with_rope_fused_weight")
    def _(
        hidden_states,
        q_a_proj_weight,
        q_b_proj_weight,
        w_kc,
        q_a_layernorm_weight,
        kv_a_layernorm_weight,
        positions,
        cos_sin_cache,
        eps,
        use_int8_w8a8,
        use_fp8_w8a16,
        qkv_a_proj_scale,
        q_b_proj_scale,
        w_scale,
        is_vnni,
        block_size,
        q_lora_rank,
        kv_lora_rank,
        qk_rope_head_dim,
    ):
        num_seqs = hidden_states.shape[0]
        num_heads = w_kc.shape[0]
        kv_lora_rank = w_kc.shape[1]
        weight_chunks = torch.split(
            q_a_proj_weight, [q_lora_rank, kv_lora_rank + qk_rope_head_dim], dim=0
        )
        qk_rope_head_dim = weight_chunks[1].shape[0] - kv_lora_rank
        q_input = torch.empty(
            num_seqs,
            num_heads,
            kv_lora_rank + qk_rope_head_dim,
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )
        k_input = torch.empty(
            num_seqs,
            1,
            kv_lora_rank + qk_rope_head_dim,
            dtype=hidden_states.dtype,
            device=hidden_states.device,
        )
        v_input = k_input.narrow(-1, 0, kv_lora_rank)
        return q_input, k_input, v_input

    def get_n_size(mat2, is_vnni):
        tile_n = 16
        if mat2.dtype == torch.float32:
            return mat2.shape[1]
        if not is_vnni and mat2.dim() == 2 and mat2.shape[0] < tile_n:
            return mat2.shape[1]
        return mat2.shape[0]

    @register_cpu_compile_fake("weight_packed_linear")
    def _(mat1, mat2, bias, is_vnni):
        M = mat1.shape[0]
        N = get_n_size(mat2, is_vnni)
        return mat1.new_empty(M, N)

    @register_cpu_compile_fake("per_token_quant_int8_cpu")
    def _(input):
        M = input.shape[0]
        K = input.shape[1]
        Aq = input.new_empty(M, K, dtype=torch.int8)
        As = input.new_empty(M, dtype=torch.float32)
        return Aq, As

    @register_cpu_compile_fake("int8_scaled_mm_cpu")
    def _(mat1, mat2, scales1, scales2, bias, out_dtype, is_vnni):
        M = mat1.shape[0]
        N = mat2.shape[0]
        out = mat1.new_empty(M, N, dtype=out_dtype)
        return out

    @register_cpu_compile_fake("grouped_topk_cpu")
    def _(
        hidden_states,
        gating_output,
        topk,
        renormalize,
        num_expert_group,
        topk_group,
        num_fused_shared_experts,
        routed_scaling_factor,
        num_token_non_padded,
    ):
        num_tokens = hidden_states.shape[0]
        shape = (num_tokens, topk)
        device = hidden_states.device
        topk_weights = torch.empty(shape, device=device, dtype=torch.float32)
        topk_ids = torch.empty(shape, device=device, dtype=torch.int)
        return topk_weights, topk_ids

    @register_cpu_compile_fake("biased_grouped_topk_cpu")
    def _(
        hidden_states,
        gating_output,
        correction_bias,
        topk,
        renormalize,
        num_expert_group,
        topk_group,
        num_fused_shared_experts,
        routed_scaling_factor,
        num_token_non_padded,
    ):
        num_tokens = hidden_states.shape[0]
        shape = (num_tokens, topk)
        device = hidden_states.device
        topk_weights = torch.empty(shape, device=device, dtype=torch.float32)
        topk_ids = torch.empty(shape, device=device, dtype=torch.int)
        return topk_weights, topk_ids

    @register_cpu_compile_fake("topk_sigmoid_cpu")
    def _(hidden_states, gating_output, topk, renormalize):
        num_tokens = hidden_states.shape[0]
        shape = (num_tokens, topk)
        return (
            torch.empty(shape, device=hidden_states.device, dtype=torch.float),
            torch.empty(shape, device=hidden_states.device, dtype=torch.int),
        )

    @register_cpu_compile_fake("topk_softmax_cpu")
    def _(
        hidden_states,
        gating_output,
        topk,
        renormalize,
    ):
        num_tokens = hidden_states.shape[0]
        shape = (num_tokens, topk)
        return (
            torch.empty(shape, device=hidden_states.device, dtype=torch.float),
            torch.empty(shape, device=hidden_states.device, dtype=torch.int),
        )

    for act_op in [
        "silu_and_mul_cpu",
        "gelu_tanh_and_mul_cpu",
        "gelu_and_mul_cpu",
    ]:

        @register_cpu_compile_fake(act_op)
        def _(input):
            sizes = list(input.shape)
            last_dim = input.dim() - 1
            d = sizes[last_dim] // 2
            sizes[last_dim] = d
            return input.new_empty(sizes)

    @register_cpu_compile_fake("int8_scaled_mm_with_quant")
    def _(
        mat1,
        mat2,
        scales2,
        bias,
        out_dtype,
        is_vnni,
    ):
        M = mat1.shape[0]
        N = mat2.shape[0]
        return mat1.new_empty(M, N, dtype=out_dtype)

    @register_cpu_compile_fake("fp8_scaled_mm_cpu")
    def _(
        mat1,
        mat2,
        scales2,
        block_size,
        bias,
        out_dtype,
        is_vnni,
    ):
        M = mat1.shape[0]
        N = mat2.shape[0]
        return mat1.new_empty(M, N, dtype=out_dtype)

    @register_cpu_compile_fake("mxfp4_scaled_mm_cpu")
    def _(mat1, mat2, scales2, bias, is_vnni):
        sizes = list(mat1.shape)
        sizes[-1] = mat2.shape[0]
        return mat1.new_empty(sizes)

    @register_cpu_compile_fake("int4_scaled_mm_cpu")
    def _(x, w, w_zeros, w_scales, bias):
        sizes = list(x.shape)
        sizes[-1] = w_scales.shape[0] * w_scales.shape[-1]
        return x.new_empty(sizes)

    @register_cpu_compile_fake("fused_linear_sigmoid_mul")
    def _(
        mat1,
        mat2,
        bias,
        is_vnni,
        post_mul_mat,
    ):
        M = mat1.shape[0]
        N = post_mul_mat.shape[1]
        return mat1.new_empty(M, N)

    @register_cpu_compile_fake("fused_qkvzba_split_reshape_cat_cpu")
    def _(mixed_qkvz, mixed_ba, num_heads_qk, num_heads_v, head_qk, head_v):
        batch = mixed_qkvz.shape[0]
        qkv_dim = num_heads_qk * head_qk * 2 + num_heads_v * head_v
        mixed_qkv = mixed_qkvz.new_empty(batch, qkv_dim)
        z = mixed_qkvz.new_empty(batch, num_heads_v, head_v)
        b = mixed_ba.new_empty(batch, num_heads_v)
        a = mixed_ba.new_empty(batch, num_heads_v)
        return mixed_qkv, z, b, a

    @register_cpu_compile_fake("fused_qkvzba_split_reshape_cat_contiguous_cpu")
    def _(mixed_qkvz, mixed_ba, num_heads_qk, num_heads_v, head_qk, head_v):
        batch = mixed_qkvz.shape[0]
        qkv_dim = num_heads_qk * head_qk * 2 + num_heads_v * head_v
        mixed_qkv = mixed_qkvz.new_empty(batch, qkv_dim)
        z = mixed_qkvz.new_empty(batch, num_heads_v, head_v)
        b = mixed_ba.new_empty(batch, num_heads_v)
        a = mixed_ba.new_empty(batch, num_heads_v)
        return mixed_qkv, z, b, a

    @register_cpu_compile_fake("fused_sigmoid_gating_delta_rule_update_cpu")
    def _(
        A_log,
        dt_bias,
        q,
        k,
        v,
        a,
        b,
        initial_state_source,
        initial_state_indices,
        cu_seqlens,
        use_qk_l2norm_in_kernel,
        softplus_beta=1.0,
        softplus_threshold=20.0,
    ):
        assert q.dim() == 4
        assert v.dim() == 4
        batch_size = q.shape[1]
        seq_len = q.shape[0]
        v_num_heads = v.shape[2]
        v_head_dim = v.shape[3]
        return q.new_empty(batch_size, seq_len, v_num_heads, v_head_dim)

    @register_cpu_compile_fake("fused_gdn_gating_cpu")
    def _(A_log, a, b, dt_bias):
        batch = a.shape[0]
        num_heads = a.shape[1]
        out = a.new_empty(1, batch, num_heads, dtype=torch.float)
        beta = b.new_empty(1, batch, num_heads)
        return out, beta

    @register_cpu_compile_fake("chunk_gated_delta_rule_cpu")
    def _(
        query,
        key,
        value,
        g,
        beta,
        initial_state,
        output_final_state,
        cu_seqlens,
        head_first,
        use_qk_l2norm_in_kernel,
        initial_state_indices,
        eps=1e-6,
    ):
        output = torch.empty_like(value)
        assert initial_state is not None
        final_state = initial_state.to(torch.float32)

        return output, final_state


# TODO Remove unnecessary settings for CPUGraphRunner.
# Re-abstract the graph runner and restructure CPUGraphRunner to reuse the same logic.
class CPUGraphRunner:
    """A CPUGraphRunner runs the forward pass of a model with cpu torch.compile."""

    def __init__(self, model_runner: ModelRunner):
        # Parse args
        self.model_runner = model_runner
        self.device = model_runner.device
        self.enable_return_hidden_states = (
            model_runner.server_args.enable_return_hidden_states
        )
        # bs -> compiled fn (text-only / skip_cross_attention=True)
        self.graphs = {}
        # bs -> compiled fn (cross-attention / skip_cross_attention=False, enc-dec only)
        self.graphs_cross = {}
        self.output_buffers = {}
        self.enable_torch_compile = get_flags().capture.enable_torch_compile
        self.disable_padding = model_runner.server_args.disable_cuda_graph_padding
        self.is_encoder_decoder = model_runner.model_config.is_encoder_decoder
        self.require_gathered_buffer = require_gathered_buffer(model_runner.server_args)
        self.require_mlp_tp_gather = require_mlp_tp_gather(model_runner.server_args)
        self.require_mlp_sync = require_mlp_sync(model_runner.server_args)
        self.require_attn_tp_gather = require_attn_tp_gather(model_runner.server_args)
        self.enable_two_batch_overlap = (
            model_runner.server_args.enable_two_batch_overlap
        )
        self.speculative_algorithm = model_runner.server_args.speculative_algorithm
        self.enable_profile_cuda_graph = (
            model_runner.server_args.enable_profile_cuda_graph
        )
        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.capture_forward_mode = ForwardMode.DECODE
        self.capture_hidden_mode = CaptureHiddenMode.NULL
        self.num_tokens_per_bs = 1

        # If returning hidden states is enabled, set initial capture hidden mode to full to avoid double-capture on startup
        if self.enable_return_hidden_states:
            self.capture_hidden_mode = CaptureHiddenMode.FULL

        assert (
            not self.model_runner.server_args.enable_lora
        ), "CPUGraphRunner does not support LoRA yet."
        assert (
            not self.enable_two_batch_overlap
        ), "CPUGraphRunner does not support two batch overlap yet."
        assert (
            not self.require_mlp_tp_gather
        ), "CPUGraphRunner does not support MLP TP gather yet."
        assert (
            not self.require_mlp_sync
        ), "CPUGraphRunner does not support MLP sync yet."
        assert (
            not self.require_gathered_buffer
        ), "CPUGraphRunner does not support gathered buffer yet."
        assert (
            model_runner.spec_algorithm.is_none()
        ), "CPUGraphRunner does not support speculative inference yet."

        assert self.dp_size == 1, "CPUGraphRunner does not support DP yet."
        assert self.pp_size == 1, "CPUGraphRunner does not support PP yet."

        # Batch sizes to capture
        self.capture_bs = get_batch_sizes_to_capture(model_runner)
        log_info_on_rank0(logger, f"Capture cpu graph bs {self.capture_bs}")
        # bs -> ForwardBatch (text-only / skip_cross_attention=True)
        self.captured_forward_batches = {}
        # bs -> ForwardBatch (cross-attention / skip=False, enc-dec only)
        self.captured_forward_batches_cross = {}
        # Attention backend
        self.max_bs = max(self.capture_bs)
        self.max_num_token = self.max_bs * self.num_tokens_per_bs
        self.model_runner.attn_backend.init_cpu_graph_state(
            self.max_bs, self.max_num_token
        )

        self.encoder_len_fill_value = 0
        self.seq_len_fill_value = (
            self.model_runner.attn_backend.get_cpu_graph_seq_len_fill_value()
        )

        if self.enable_torch_compile:
            register_fake_ops(self.tp_size)
            set_torch_compile_config()

        # Graph inputs
        with torch.device(self.device):
            self.input_ids = torch.zeros((self.max_num_token,), dtype=torch.int64)
            self.req_pool_indices = torch.zeros((self.max_bs,), dtype=torch.int64)
            self.seq_lens = torch.full(
                (self.max_bs,), self.seq_len_fill_value, dtype=torch.int64
            )
            self.out_cache_loc = torch.zeros((self.max_num_token,), dtype=torch.int64)
            self.positions = torch.zeros((self.max_num_token,), dtype=torch.int64)
            self.mrope_positions = torch.zeros((3, self.max_bs), dtype=torch.int64)
            self.num_token_non_padded = torch.zeros((1,), dtype=torch.int64)
            self.custom_mask = torch.ones(
                (
                    (self.seq_lens.sum().item() + self.max_num_token)
                    * self.num_tokens_per_bs
                ),
                dtype=torch.bool,
                device=self.device,
            )
            if self.is_encoder_decoder:
                self.encoder_lens = torch.full(
                    (self.max_bs,), self.encoder_len_fill_value, dtype=torch.int64
                )
            else:
                self.encoder_lens = None

        # Capture
        try:
            # use model_capture_mode for encoder-decoder models to
            # set skip_cross_attention to avoid
            # "Graph Break Reason: Data-dependent branching" caused by
            # skip_cross_attention = forward_batch.encoder_lens.max() == 0
            capture_context = (
                model_capture_mode if self.is_encoder_decoder else empty_context
            )
            with capture_context():
                self.capture()
        except RuntimeError as e:
            raise Exception(
                f"Capture CPU graph failed: {e}\n{CPU_GRAPH_CAPTURE_FAILED_MSG}"
            )

    def _get_skip_cross_attention(self, forward_batch: ForwardBatch) -> bool:
        """Return True when cross-attention layers should be skipped.

        Non-encoder-decoder models have no cross-attention at all, so they
        always use self.graphs (the skip=True / text-only graph dict).
        For encoder-decoder models, skip when no request in the batch has
        encoder output (i.e. no images).
        """
        if not self.is_encoder_decoder:
            return True
        return bool(forward_batch.encoder_lens.max() == 0)

    def can_run_graph(self, forward_batch: ForwardBatch):
        is_bs_supported = (
            forward_batch.batch_size in self.graphs
            if self.disable_padding
            else forward_batch.batch_size <= self.max_bs
        )

        requested_capture_hidden_mode = max(
            forward_batch.capture_hidden_mode,
            (
                forward_batch.spec_info.capture_hidden_mode
                if getattr(forward_batch.spec_info, "capture_hidden_mode", None)
                is not None
                else CaptureHiddenMode.NULL
            ),
        )
        capture_hidden_mode_matches = (
            requested_capture_hidden_mode == CaptureHiddenMode.NULL
            or requested_capture_hidden_mode == self.capture_hidden_mode
        )

        return is_bs_supported and capture_hidden_mode_matches

    def capture(self) -> None:
        capture_range = (
            tqdm.tqdm(list(reversed(self.capture_bs)))
            if get_parallel().tp_rank == 0
            else reversed(self.capture_bs)
        )
        for bs in capture_range:
            if get_parallel().tp_rank == 0:
                avail_mem = psutil.virtual_memory().available / (1 << 30)
                capture_range.set_description(
                    f"Capturing batches ({bs=} {avail_mem=:.2f} GB)"
                )

            with patch_model(
                self.model_runner.model,
                bs in self.capture_bs,
                num_tokens=bs * self.num_tokens_per_bs,
                tp_group=self.model_runner.tp_group,
            ) as forward:
                graph, output_buffers = self.capture_one_batch_size(
                    bs, forward, skip_cross_attention=True
                )
                self.graphs[bs] = graph
                self.output_buffers[bs] = output_buffers
                if self.is_encoder_decoder:
                    # Capture a second graph with cross-attention enabled
                    # (used when the batch contains images).
                    graph_cross, _ = self.capture_one_batch_size(
                        bs, forward, skip_cross_attention=False
                    )
                    self.graphs_cross[bs] = graph_cross

        # Re-init states for qwen3-next as
        # torch.compile may change the states
        self._reset_mamba_cache_if_needed()

    def _reset_mamba_cache_if_needed(self) -> None:

        mamba_pool = getattr(self.model_runner.req_to_token_pool, "mamba_pool", None)
        if mamba_pool is None:
            return
        mamba_cache = getattr(mamba_pool, "mamba_cache", None)
        if mamba_cache is None:
            return

        def _zero_nested(obj):
            if isinstance(obj, torch.Tensor):
                obj.zero_()
            elif isinstance(obj, (list, tuple)):
                for it in obj:
                    _zero_nested(it)

        for v in vars(mamba_cache).values():
            _zero_nested(v)

    def capture_one_batch_size(
        self, bs: int, forward: Callable, skip_cross_attention: bool = False
    ):
        num_tokens = bs * self.num_tokens_per_bs

        # Graph inputs
        input_ids = self.input_ids[:num_tokens]
        req_pool_indices = self.req_pool_indices[:bs]
        seq_lens = self.seq_lens[:bs]
        out_cache_loc = self.out_cache_loc[:num_tokens]
        positions = self.positions[:num_tokens]
        mrope_positions = self.mrope_positions[:, :num_tokens]
        self.num_token_non_padded[...] = num_tokens
        if self.is_encoder_decoder:
            encoder_lens = self.encoder_lens[:bs]
        else:
            encoder_lens = None

        spec_info = self.get_spec_info(num_tokens)
        if self.capture_hidden_mode != CaptureHiddenMode.FULL:
            self.capture_hidden_mode = (
                spec_info.capture_hidden_mode if spec_info else CaptureHiddenMode.NULL
            )

        forward_batch = ForwardBatch(
            forward_mode=self.capture_forward_mode,
            batch_size=bs,
            input_ids=input_ids,
            req_pool_indices=req_pool_indices,
            seq_lens=seq_lens,
            out_cache_loc=out_cache_loc,
            seq_lens_sum=seq_lens.sum().item(),
            encoder_lens=encoder_lens,
            encoder_lens_cpu=encoder_lens,
            return_logprob=False,
            positions=positions,
            mrope_positions=mrope_positions,
            spec_algorithm=self.model_runner.spec_algorithm,
            spec_info=spec_info,
            capture_hidden_mode=self.capture_hidden_mode,
            num_token_non_padded=self.num_token_non_padded,
            global_forward_mode=self.capture_forward_mode,
        )
        # Wrap all forward calls with capture_with_skip_cross_attention so that
        # mllama (and any other encoder-decoder model) sees the correct compile-
        # time constant for skip_cross_attention during tracing.
        skip_ctx = (
            capture_with_skip_cross_attention(skip_cross_attention)
            if self.is_encoder_decoder
            else empty_context()
        )
        with skip_ctx:
            with forward_context(
                ForwardContext(attn_backend=self.model_runner.attn_backend)
            ):
                self.model_runner.attn_backend.init_forward_metadata_capture_cpu_graph(
                    bs,
                    num_tokens,
                    req_pool_indices,
                    seq_lens,
                    None,
                    forward_batch.forward_mode,
                    forward_batch.spec_info,
                )
                with torch.no_grad():
                    self.model_runner.tp_group.barrier()
                    self.model_runner.model.forward(
                        forward_batch.input_ids,
                        forward_batch.positions,
                        forward_batch,
                    )

                # Run and capture
                def run_once():
                    # Clean intermediate result cache for DP attention
                    forward_batch.dp_local_start_pos = (
                        forward_batch.dp_local_num_tokens
                    ) = None
                    logits_output_or_pp_proxy_tensors = forward(
                        forward_batch.input_ids,
                        forward_batch.positions,
                        forward_batch,
                    )
                    return logits_output_or_pp_proxy_tensors

                with torch.no_grad():
                    for _ in range(2):
                        self.model_runner.tp_group.barrier()
                        out = run_once()
                    # Save the captured forward_batch in the appropriate dict
                    if skip_cross_attention:
                        self.captured_forward_batches[bs] = forward_batch
                    else:
                        self.captured_forward_batches_cross[bs] = forward_batch
                    return forward, out

    def recapture_if_needed(self, forward_batch: ForwardBatch):

        # If the required capture_hidden_mode changes, we need to recapture the graph

        # These are the different factors that can influence the capture_hidden_mode
        capture_hidden_mode_required_by_forward_batch = (
            forward_batch.capture_hidden_mode
        )
        capture_hidden_mode_required_by_spec_info = getattr(
            forward_batch.spec_info, "capture_hidden_mode", CaptureHiddenMode.NULL
        )
        capture_hidden_mode_required_for_returning_hidden_states = (
            CaptureHiddenMode.FULL
            if self.enable_return_hidden_states
            else CaptureHiddenMode.NULL
        )

        # Determine the highest capture_hidden_mode required
        # (If we have FULL, we can emulate LAST or NULL)
        # (If we have LAST, we can emulate NULL)
        required_capture_hidden_mode = max(
            capture_hidden_mode_required_by_forward_batch,
            capture_hidden_mode_required_by_spec_info,
            capture_hidden_mode_required_for_returning_hidden_states,
        )

        # If the current hidden mode is no longer aligned with the required hidden mode, we need to set it to what is required and re-capture
        if self.capture_hidden_mode != required_capture_hidden_mode:
            self.capture_hidden_mode = required_capture_hidden_mode
            self.capture()

    def prepare_replay(
        self,
        forward_batch: ForwardBatch,
        skip: bool = False,
    ):
        self.recapture_if_needed(forward_batch)

        graphs = self.graphs_cross if not skip else self.graphs
        cfbs = (
            self.captured_forward_batches_cross
            if not skip
            else self.captured_forward_batches
        )

        raw_bs = forward_batch.batch_size
        if raw_bs in graphs:
            # Keep encoder_out_cache_loc consistent with the captured graph (None).
            if self.is_encoder_decoder:
                # encoder_out_cache_loc is never accessed during decode (k/v are
                # None so the KV-write path is skipped in the kernel).  Use None
                # consistently at both capture time and runtime.
                forward_batch.encoder_out_cache_loc = None
            self.model_runner.attn_backend.init_forward_metadata(forward_batch)
            return forward_batch

        raw_num_token = raw_bs * self.num_tokens_per_bs
        index = bisect.bisect_left(self.capture_bs, raw_bs)
        bs = self.capture_bs[index]
        assert bs > raw_bs
        self.raw_bs = raw_bs
        self.raw_num_token = raw_num_token
        self.bs = bs

        captured_forward_batch = cfbs[bs]
        assert captured_forward_batch is not None
        captured_forward_batch.seq_lens.fill_(self.seq_len_fill_value)
        captured_forward_batch.out_cache_loc.zero_()
        # Pair with seq_lens fill: padded rows must point at reserved
        # req_pool slot 0 (req_to_token[0, :] is all zeros from init).
        captured_forward_batch.req_pool_indices.zero_()
        captured_forward_batch.input_ids[:raw_num_token].copy_(forward_batch.input_ids)
        captured_forward_batch.req_pool_indices[:raw_bs].copy_(
            forward_batch.req_pool_indices
        )
        captured_forward_batch.seq_lens[:raw_bs].copy_(forward_batch.seq_lens)
        captured_forward_batch.out_cache_loc[:raw_num_token].copy_(
            forward_batch.out_cache_loc
        )
        captured_forward_batch.positions[:raw_num_token].copy_(forward_batch.positions)
        if forward_batch.mrope_positions is not None:
            self.mrope_positions[:, :raw_num_token].copy_(forward_batch.mrope_positions)

        if self.is_encoder_decoder:
            captured_forward_batch.encoder_lens[:raw_bs].copy_(
                forward_batch.encoder_lens
            )
            captured_forward_batch.encoder_out_cache_loc = None
        if enable_num_token_non_padded():
            captured_forward_batch.num_token_non_padded.copy_(
                forward_batch.num_token_non_padded
            )

        self.model_runner.attn_backend.init_forward_metadata(captured_forward_batch)
        return captured_forward_batch

    def execute(
        self,
        forward_batch: ForwardBatch,
        pp_proxy_tensors: Optional[PPProxyTensors] = None,
    ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
        assert (
            pp_proxy_tensors is None
        ), "PPProxyTensors is not supported in CPUGraphRunner yet."

        replay_context = (
            model_capture_mode if self.is_encoder_decoder else empty_context
        )
        # Determine which compiled graph to use and pin skip_cross_attention so
        # that any torch.compile re-tracing sees the same compile-time constant.
        skip = self._get_skip_cross_attention(forward_batch)
        graphs = self.graphs_cross if not skip else self.graphs
        skip_ctx = (
            capture_with_skip_cross_attention(skip)
            if self.is_encoder_decoder
            else empty_context()
        )
        with replay_context():
            with skip_ctx:
                prepared_forward_batch = self.prepare_replay(forward_batch, skip=skip)
                output = graphs[prepared_forward_batch.batch_size](
                    prepared_forward_batch.input_ids,
                    prepared_forward_batch.positions,
                    prepared_forward_batch,
                )
        if forward_batch.batch_size in graphs:
            return output

        assert isinstance(output, LogitsProcessorOutput)
        return LogitsProcessorOutput(
            next_token_logits=output.next_token_logits[: self.raw_num_token],
            hidden_states=(
                output.hidden_states[: self.raw_num_token]
                if output.hidden_states is not None
                else None
            ),
        )

    def get_spec_info(self, num_tokens: int):
        spec_info = None
        if (
            self.model_runner.spec_algorithm.is_eagle()
            or self.model_runner.spec_algorithm.is_standalone()
        ):
            from sglang.srt.speculative.eagle_info import EagleVerifyInput

            if self.model_runner.is_draft_worker:
                raise RuntimeError("This should not happen.")
            else:
                spec_info = EagleVerifyInput(
                    draft_token=None,
                    custom_mask=self.custom_mask,
                    positions=None,
                    retrieve_index=None,
                    retrieve_next_token=None,
                    retrieve_next_sibling=None,
                    retrieve_cum_len=None,
                    spec_steps=self.model_runner.server_args.speculative_num_steps,
                    topk=self.model_runner.server_args.speculative_eagle_topk,
                    draft_token_num=self.model_runner.server_args.speculative_num_draft_tokens,
                    capture_hidden_mode=CaptureHiddenMode.FULL,
                    seq_lens_sum=None,
                    seq_lens_cpu=None,
                )

        return spec_info


CPU_GRAPH_CAPTURE_FAILED_MSG = (
    "Possible solutions:\n"
    "1. set --mem-fraction-static to a smaller value (e.g., 0.8 or 0.7)\n"
    "2. set --torch-compile-max-bs to a smaller value (e.g., 8)\n"
    "3. disable torch compile by not using --enable-torch-compile\n"
    "Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
)