chore: import upstream snapshot with attribution

python/sglang/srt/model_executor/cuda_graph_buffer_registry.py

@@ -0,0 +1,947 @@
+# 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.
+# ==============================================================================
+"""FB-shared slot registry for the CUDA graph forward paths.
+
+``CudaGraphBufferRegistry`` is the ForwardBatch → graph-resident buffer mirror
+used by capture / replay. It replaces the per-runner ``DecodeInputBuffers`` /
+``PrefillInputBuffers`` dataclasses and their hand-written
+``populate_from_forward_batch`` methods with a single ``GraphSlot``-driven
+registry.
+
+Backend-private buffers (kernel workspaces, derived page tables, etc.) stay
+on ``AttentionBackend.cuda_graph_*`` — the registry only owns FB-shared
+slots (FB attribute name maps 1:1 to slot name).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Tuple
+
+import torch
+
+from sglang.srt.model_executor.input_buffers import share_input_buffer
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+_has_foreach_copy = hasattr(torch, "_foreach_copy_")
+
+
+def _grouped_foreach_copy_(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:
+    """Call torch._foreach_copy_ grouped by (dst_dtype, src_dtype) pairs
+    (a single foreach call requires a uniform dtype pair)."""
+
+    def _foreach_copy(
+        group_dsts: List[torch.Tensor], group_srcs: List[torch.Tensor]
+    ) -> None:
+        if _has_foreach_copy:
+            torch._foreach_copy_(group_dsts, group_srcs)
+        else:
+            for dst, src in zip(group_dsts, group_srcs):
+                dst.copy_(src)
+
+    groups: Dict[Tuple[torch.dtype, torch.dtype], Tuple[List, List]] = {}
+    for dst, src in zip(dsts, srcs):
+        key = (dst.dtype, src.dtype)
+        if key not in groups:
+            groups[key] = ([], [])
+        groups[key][0].append(dst)
+        groups[key][1].append(src)
+    for group_dsts, group_srcs in groups.values():
+        _foreach_copy(group_dsts, group_srcs)
+
+
+class PaddingPolicy(Enum):
+    """How to handle ``raw_n < padded_n`` for a slot.
+
+    KEEP_PAD      — Leave the padded region as-is (caller proves the
+                    padded tail will not be read).
+    FILL_SENTINEL — Reset the padded region to ``slot.pad_value`` before
+                    copy (e.g. ``seq_lens`` filled with
+                    ``seq_len_fill_value``).
+    ZERO          — Reset the padded region to ``0`` (e.g.
+                    ``out_cache_loc`` / ``req_pool_indices`` — padded
+                    rows must point at slot 0 so dummy attention reads
+                    land harmlessly).
+    FOREACH_COPY  — Always copy ``raw_n`` from src; padded region is
+                    left as whatever the previous replay (or the init
+                    zeros) wrote. Caller is responsible for proving
+                    safety.
+    FILL_ONCE     — Fill the whole buffer to ``pad_value`` once at alloc;
+                    never reset per iter (e.g. ``encoder_lens`` init to
+                    ``encoder_len_fill_value``, copied head-only with the
+                    tail kept).
+    """
+
+    KEEP_PAD = "keep_pad"
+    FILL_SENTINEL = "fill_sentinel"
+    ZERO = "zero"
+    FOREACH_COPY = "foreach_copy"
+    FILL_ONCE = "fill_once"
+
+
+@dataclass
+class FillContext:
+    """Per-iteration shape context passed to ``GraphSlot.post_fill``.
+
+    Carries both the bs-axis and tokens-axis raw/padded counts so a hook can
+    derive values regardless of its own slot's axis — e.g. the padded token
+    count (``padded_num_tokens`` == padded_bs * num_tokens_per_bs), which the
+    global-num-tokens fill and the local-num-token-non-padded transform need.
+    """
+
+    raw_bs: int
+    padded_bs: int
+    raw_num_tokens: int
+    padded_num_tokens: int
+    # Side inputs that are not ForwardBatch attributes but are needed by a
+    # slot's source_fn — e.g. the pipeline-parallel proxy tensors, which the
+    # replay path receives as a separate argument rather than off the FB.
+    pp_proxy_tensors: Optional[Any] = None
+
+
+@dataclass
+class GraphSlot:
+    """A single FB-mirrored buffer.
+
+    Each slot mirrors one ``ForwardBatch`` attribute. ``name`` MUST match
+    the FB attribute name so ``fill_from`` can ``getattr(fb, name)`` and
+    ``extract_buffer`` can ``setattr`` the view back into a FB replace.
+
+    Fields:
+        name           — the FB attribute name mirrored by this slot.
+        shape_fn       — ``(max_bs, max_num_tokens) -> shape`` callable
+                         used at ``register_slot`` time to allocate the
+                         physical buffer.
+        dtype          — buffer dtype.
+        axis           — ``"bs"`` (slot is sliced ``[:bs]``) or
+                         ``"tokens"`` (sliced ``[:num_tokens]``) or
+                         ``"none"`` (no slicing — full buffer always
+                         exposed; used for scalar buffers and global
+                         counters).
+        device         — buffer device. ``None`` means use registry
+                         default; can be ``"cpu"`` for slots like
+                         ``seq_lens_cpu`` that must live on host.
+        padding_policy — see ``PaddingPolicy``.
+        pad_value      — sentinel for ``FILL_SENTINEL``.
+        enabled        — runtime gate; disabled slots are not allocated
+                         and skipped during fill / extract.
+        copy_from_fb   — when ``True`` (default), ``fill_from`` copies the
+                         same-named FB tensor into the buffer head. Set
+                         ``False`` for computed slots whose value is not a
+                         straight FB copy (e.g. ``global_num_tokens_*``,
+                         filled by a ``post_fill`` instead).
+        post_fill      — optional ``(buffer, forward_batch, FillContext)
+                         -> None`` hook run after the grouped copy. Used for
+                         compute-then-write slots (local-num-token-non-padded
+                         transform, global-num-tokens fill).
+        slice_fn       — optional ``(buffer, padded_n) -> Tensor``
+                         override for slots with non-trivial slicing
+                         (e.g. ``mrope_positions`` shape ``[3, T]`` is
+                         sliced on axis 1 not 0).
+        source_fn      — optional ``(forward_batch, FillContext) -> Tensor |
+                         None`` override for the copy *source*. When set,
+                         ``fill_from`` copies ``source_fn(fb, ctx)`` (instead of
+                         the same-named FB attribute) into
+                         ``buffer[:src.shape[0]]`` — a source-length slice for
+                         structured / side-sourced fields whose data lives on a
+                         nested FB dataclass (``ngram_embedding_info.*``) or an
+                         out-of-band argument (``pp_proxy_tensors``, carried on
+                         ``FillContext``). Returning ``None`` skips the copy for
+                         that iteration. Such slots use dotted names and are
+                         skipped by ``extract_buffer``.
+    """
+
+    name: str
+    shape_fn: Callable[[int, int], Tuple[int, ...]]
+    dtype: torch.dtype
+    axis: str = "tokens"
+    device: Optional[torch.device] = None
+    padding_policy: PaddingPolicy = PaddingPolicy.FOREACH_COPY
+    pad_value: Optional[Any] = None
+    enabled: bool = True
+    copy_from_fb: bool = True
+    post_fill: Optional[Callable[[torch.Tensor, ForwardBatch, FillContext], None]] = (
+        None
+    )
+    slice_fn: Optional[Callable[[torch.Tensor, int], torch.Tensor]] = None
+    source_fn: Optional[
+        Callable[[ForwardBatch, FillContext], Optional[torch.Tensor]]
+    ] = None
+
+    # runtime
+    buffer: Optional[torch.Tensor] = field(default=None, repr=False)
+
+    def __post_init__(self) -> None:
+        if self.axis not in ("bs", "tokens", "none"):
+            raise ValueError(
+                f"GraphSlot {self.name!r}: axis must be one of "
+                f"'bs'/'tokens'/'none', got {self.axis!r}"
+            )
+
+    def _padded_n(self, padded_bs: int, padded_num_tokens: int) -> int:
+        if self.axis == "bs":
+            return padded_bs
+        if self.axis == "tokens":
+            return padded_num_tokens
+        # axis == "none": no slicing
+        return self.buffer.shape[0] if self.buffer is not None else 0
+
+    def _raw_n(self, raw_bs: int, raw_num_tokens: int) -> int:
+        if self.axis == "bs":
+            return raw_bs
+        if self.axis == "tokens":
+            return raw_num_tokens
+        return self.buffer.shape[0] if self.buffer is not None else 0
+
+    def slice_for(self, padded_bs: int, padded_num_tokens: int) -> torch.Tensor:
+        """Return the ``[:padded_n]`` slice of the buffer consumed by callers.
+
+        This truncates the (full-length) buffer to the active region for the
+        current iteration — it is a slice, not a tensor reshape.
+        """
+        if self.buffer is None:
+            raise RuntimeError(f"GraphSlot {self.name!r}: buffer not allocated")
+        if self.slice_fn is not None:
+            return self.slice_fn(
+                self.buffer, self._padded_n(padded_bs, padded_num_tokens)
+            )
+        if self.axis == "none":
+            return self.buffer
+        return self.buffer[: self._padded_n(padded_bs, padded_num_tokens)]
+
+    def reset_padding(self, raw_n: int, padded_n: int) -> None:
+        """Reset the padded tail according to ``padding_policy``."""
+        if self.buffer is None or raw_n >= padded_n:
+            return
+        if self.padding_policy in (
+            PaddingPolicy.KEEP_PAD,
+            PaddingPolicy.FOREACH_COPY,
+            PaddingPolicy.FILL_ONCE,
+        ):
+            return
+        # slice_fn governs non-trivial layouts (e.g. mrope_positions [3, T]);
+        # the pad region is the same axis the slot exposes via slice_for().
+        if self.slice_fn is not None:
+            # slice_fn returns the [:padded_n] portion already; we need the
+            # tail [raw_n:padded_n]. We rely on slice_fn slicing the same
+            # axis used by slice_for(): take the padded slice first, then index
+            # the tail with the standard slice on axis 0 of the result.
+            padded_slice = self.slice_fn(self.buffer, padded_n)
+            tail = (
+                padded_slice[..., raw_n:padded_n]
+                if padded_slice.dim() > 1
+                else padded_slice[raw_n:padded_n]
+            )
+        else:
+            tail = self.buffer[raw_n:padded_n]
+        if self.padding_policy == PaddingPolicy.FILL_SENTINEL:
+            if self.pad_value is None:
+                raise RuntimeError(
+                    f"GraphSlot {self.name!r}: FILL_SENTINEL requires pad_value"
+                )
+            tail.fill_(self.pad_value)
+        elif self.padding_policy == PaddingPolicy.ZERO:
+            tail.zero_()
+
+
+class CudaGraphBufferRegistry:
+    """FB → graph-resident buffer mirror, shared across eager / capture / replay.
+
+    The registry holds a dict of ``GraphSlot`` instances, each mirroring
+    one ``ForwardBatch`` attribute. Slots are registered up-front (during
+    runner init), allocated at ``register_slot``, then filled per-iter via
+    ``fill_from(fb, ...)`` and consumed via ``extract_buffer(template) ->
+    ForwardBatch``. ``fill_from`` issues plain D2D copies on the caller's
+    current stream; cross-stream correctness (stream handoff) is handled by
+    the runners, not here.
+
+    Backend-private buffers (kernel workspace, derived page tables) are
+    NOT managed here — backends keep them on ``self.cuda_graph_*`` and
+    allocate via ``AttentionBackend.init_cuda_graph_state(...)``.
+
+    Usage::
+
+        registry = CudaGraphBufferRegistry(device=..., max_bs=..., max_num_tokens=...)
+        registry.register_slot(GraphSlot(name="input_ids", ...))
+        registry.register_slot(GraphSlot(name="seq_lens",
+                                         padding_policy=PaddingPolicy.FILL_SENTINEL,
+                                         pad_value=seq_len_fill_value, ...))
+        # per-iter:
+        registry.fill_from(fb, raw_bs=..., padded_bs=..., raw_num_tokens=...,
+                           padded_num_tokens=...)
+        fb_view = registry.extract_buffer(padded_bs=..., padded_num_tokens=...,
+                                          forward_batch_template=fb)
+        attn_backend.init_forward_metadata(fb_view)
+        model.forward(fb_view.input_ids, fb_view.positions, fb_view)
+    """
+
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        max_bs: int,
+        max_num_tokens: int,
+        share_pool: bool = False,
+    ) -> None:
+        self.device = device
+        self.max_bs = max_bs
+        self.max_num_tokens = max_num_tokens
+        # Coalesce allocated slot buffers through the global pool; only applies
+        # when allocating (bind/source bypasses the pool).
+        self.share_pool = share_pool
+        self._slots: Dict[str, GraphSlot] = {}
+
+    # ---- registration ------------------------------------------------------
+
+    def register_slot(
+        self, slot: GraphSlot, bind: Optional[torch.Tensor] = None
+    ) -> GraphSlot:
+        """Register a slot and allocate (or adopt) its physical buffer.
+
+        If ``bind`` is given, the slot adopts that existing tensor instead of
+        allocating a fresh one (and skips the pool / sentinel init — the bound
+        tensor is assumed already initialized). This lets a registry share
+        storage with the legacy ``DecodeInputBuffers`` by adopting its fields,
+        guaranteeing a stable, identical ``data_ptr`` for capture vs replay.
+
+        Returns the slot for caller convenience. Re-registering an existing
+        name raises.
+        """
+        if slot.name in self._slots:
+            raise ValueError(
+                f"GraphSlot {slot.name!r} already registered; "
+                "use enable()/disable() to gate per-iter."
+            )
+        if not slot.enabled:
+            # Even when disabled, keep the spec so callers can introspect
+            # by name; just don't allocate.
+            self._slots[slot.name] = slot
+            return slot
+        shape = slot.shape_fn(self.max_bs, self.max_num_tokens)
+        device = slot.device if slot.device is not None else self.device
+        if bind is not None:
+            expected = tuple(shape)
+            if tuple(bind.shape) != expected:
+                raise ValueError(
+                    f"bind tensor for slot {slot.name!r} has shape "
+                    f"{tuple(bind.shape)}, expected {expected}."
+                )
+            if bind.dtype != slot.dtype:
+                raise ValueError(
+                    f"bind tensor for slot {slot.name!r} has dtype {bind.dtype}, "
+                    f"expected {slot.dtype}."
+                )
+            slot.buffer = bind
+            self._slots[slot.name] = slot
+            return slot
+        buffer = torch.zeros(shape, dtype=slot.dtype, device=device)
+        if self.share_pool:
+            # Coalesce with any same-named buffer (e.g. the legacy
+            # DecodeInputBuffers field) so capture and replay see one
+            # physical allocation with a stable data_ptr.
+            buffer = share_input_buffer(slot.name, buffer)
+        if (
+            slot.padding_policy
+            in (PaddingPolicy.FILL_SENTINEL, PaddingPolicy.FILL_ONCE)
+            and slot.pad_value is not None
+        ):
+            buffer.fill_(slot.pad_value)
+        slot.buffer = buffer
+        self._slots[slot.name] = slot
+        return slot
+
+    def has_slot(self, name: str) -> bool:
+        return name in self._slots and self._slots[name].enabled
+
+    def get_slot(self, name: str) -> GraphSlot:
+        return self._slots[name]
+
+    def slot_names(self) -> List[str]:
+        return [name for name, s in self._slots.items() if s.enabled]
+
+    # ---- per-iter ----------------------------------------------------------
+
+    def fill_from(
+        self,
+        forward_batch: ForwardBatch,
+        *,
+        raw_bs: int,
+        padded_bs: int,
+        raw_num_tokens: int,
+        padded_num_tokens: int,
+        pp_proxy_tensors: Optional[Any] = None,
+    ) -> None:
+        """Copy FB → registry buffers.
+
+        Phase 1 — reset the padded tail per slot ``padding_policy``.
+        Phase 2 — grouped D2D copy of all enabled slots from FB (or from a
+                  slot's ``source_fn`` for structured / side-sourced fields).
+        Phase 3 — run ``post_fill`` hooks for slots that need
+                  post-copy transforms.
+
+        ``pp_proxy_tensors`` is the out-of-band pipeline-parallel input; it is
+        not an FB attribute, so it reaches ``source_fn`` slots via
+        ``FillContext.pp_proxy_tensors``.
+
+        Slots whose FB attribute (or ``source_fn`` result) is ``None`` are
+        silently skipped (the FB doesn't carry that field for the current
+        request).
+        """
+        ctx = FillContext(
+            raw_bs=raw_bs,
+            padded_bs=padded_bs,
+            raw_num_tokens=raw_num_tokens,
+            padded_num_tokens=padded_num_tokens,
+            pp_proxy_tensors=pp_proxy_tensors,
+        )
+
+        # Phase 1: reset padded regions where it matters.
+        for slot in self._slots.values():
+            if not slot.enabled or slot.buffer is None:
+                continue
+            raw_n = slot._raw_n(raw_bs, raw_num_tokens)
+            padded_n = slot._padded_n(padded_bs, padded_num_tokens)
+            slot.reset_padding(raw_n, padded_n)
+
+        # Phase 2: collect (dst, src) pairs and dispatch a grouped copy.
+        gpu_dsts: List[torch.Tensor] = []
+        gpu_srcs: List[torch.Tensor] = []
+        cpu_dsts: List[torch.Tensor] = []
+        cpu_srcs: List[torch.Tensor] = []
+        for slot in self._slots.values():
+            if not slot.enabled or slot.buffer is None or not slot.copy_from_fb:
+                continue
+            if slot.source_fn is not None:
+                # Structured / side-sourced slot: source comes from a nested FB
+                # dataclass or an out-of-band input, and the copy is sliced to
+                # the source's own length rather than a bs/tokens axis.
+                src = slot.source_fn(forward_batch, ctx)
+                if src is None:
+                    continue
+                dst = slot.buffer[: src.shape[0]]
+            else:
+                src = getattr(forward_batch, slot.name, None)
+                if src is None:
+                    continue
+                if not isinstance(src, torch.Tensor):
+                    # Non-tensor FB fields (e.g. dicts, dataclasses) are not
+                    # auto-copied — caller handles via source_fn or post_fill.
+                    continue
+                raw_n = slot._raw_n(raw_bs, raw_num_tokens)
+                if slot.slice_fn is not None:
+                    dst = slot.slice_fn(slot.buffer, raw_n)
+                elif slot.axis == "none":
+                    dst = slot.buffer
+                else:
+                    dst = slot.buffer[:raw_n]
+            # foreach_copy_ requires same-device tensors per call — bucket
+            # by device.
+            if dst.device.type == "cpu":
+                cpu_dsts.append(dst)
+                cpu_srcs.append(src)
+            else:
+                gpu_dsts.append(dst)
+                gpu_srcs.append(src)
+        if gpu_dsts:
+            _grouped_foreach_copy_(gpu_dsts, gpu_srcs)
+        for dst, src in zip(cpu_dsts, cpu_srcs):
+            dst.copy_(src)
+
+        # Phase 3: post-fill hooks (compute-then-write slots).
+        for slot in self._slots.values():
+            if not slot.enabled or slot.buffer is None or slot.post_fill is None:
+                continue
+            slot.post_fill(slot.buffer, forward_batch, ctx)
+
+    def extract_buffer(
+        self,
+        *,
+        padded_bs: int,
+        padded_num_tokens: int,
+        forward_batch_template: ForwardBatch,
+    ) -> ForwardBatch:
+        """Return a FB view (``dataclasses.replace`` of ``forward_batch_template``)
+        whose slot fields are buffer views and whose non-slot fields are carried
+        from the template. A plain copy slot whose FB field is ``None`` this iter
+        is carried (not exposed as a stale buffer); computed slots are always
+        exposed.
+        """
+        import dataclasses
+
+        replace_kwargs: Dict[str, Any] = {"batch_size": padded_bs}
+        for slot in self._slots.values():
+            if not slot.enabled or slot.buffer is None:
+                continue
+            # Structured slots use dotted names ("<field>.<sub>") and are not
+            # top-level FB attributes — their data is consumed in place off the
+            # adopted backing object, not re-attached to the FB view here.
+            if "." in slot.name:
+                continue
+            is_computed = slot.post_fill is not None or not slot.copy_from_fb
+            if (
+                not is_computed
+                and slot.source_fn is None
+                and getattr(forward_batch_template, slot.name, None) is None
+            ):
+                # Absent this iter (fill_from skipped it): carry the template.
+                continue
+            replace_kwargs[slot.name] = slot.slice_for(padded_bs, padded_num_tokens)
+        return dataclasses.replace(forward_batch_template, **replace_kwargs)
+
+
+def build_decode_registry(
+    *,
+    device: torch.device,
+    max_bs: int,
+    max_num_token: int,
+    seq_len_fill_value: int,
+    cache_loc_dtype: torch.dtype,
+    enable_mamba_track: bool = False,
+    is_encoder_decoder: bool = False,
+    encoder_len_fill_value: int = 0,
+    encoder_lens_dtype: torch.dtype = torch.int32,
+    enable_num_token_non_padded: bool = False,
+    require_gathered_buffer: bool = False,
+    enable_prefill_cp: bool = False,
+    require_mlp_tp_gather: bool = False,
+    dp_size: int = 1,
+    register_global_num_tokens: bool = True,
+    share_pool: bool = True,
+    source: Optional[Any] = None,
+) -> CudaGraphBufferRegistry:
+    """Registry mirroring the always-on (+ mamba / mrope) FB-shared decode
+    buffers, with the per-slot padding policy that resets the padded tail on
+    each replay:
+
+      - ``seq_lens`` / ``seq_lens_cpu`` -> FILL_SENTINEL(seq_len_fill_value)
+      - ``req_pool_indices`` / ``out_cache_loc`` / ``mamba_track_*`` -> ZERO
+      - ``positions`` / ``mrope_positions`` -> ZERO: the flashinfer verify-path
+        plan reads the padded tail, so leaving stale out-of-range values there
+        triggers an illegal memory access (issue #24361).
+      - ``input_ids`` -> FOREACH_COPY: head ``[:raw_n]`` is overwritten by the
+        copy and the padded tail is not read.
+
+    ``custom_mask`` / ``next_token_logits_buffer`` / ``input_embeds`` are not
+    registered here — they are not per-replay FB copies (allocated and written
+    elsewhere), so the runner keeps owning them.
+
+    When ``source`` is given (the decode buffer namespace from
+    ``_allocate_decode_buffers``), each slot adopts the same-named tensor off
+    it instead of allocating, so the registry shares one physical allocation
+    with that object. With ``source=None`` the registry allocates its own.
+    """
+    reg = CudaGraphBufferRegistry(
+        device=device,
+        max_bs=max_bs,
+        max_num_tokens=max_num_token,
+        share_pool=share_pool,
+    )
+
+    def _tokens(_bs: int, mt: int) -> Tuple[int, ...]:
+        return (mt,)
+
+    def _bs(bs: int, _mt: int) -> Tuple[int, ...]:
+        return (bs,)
+
+    slots = [
+        GraphSlot("input_ids", _tokens, torch.int64, axis="tokens"),
+        GraphSlot(
+            "positions",
+            _tokens,
+            torch.int64,
+            axis="tokens",
+            padding_policy=PaddingPolicy.ZERO,
+        ),
+        GraphSlot(
+            "out_cache_loc",
+            _tokens,
+            cache_loc_dtype,
+            axis="tokens",
+            padding_policy=PaddingPolicy.ZERO,
+        ),
+        GraphSlot(
+            "req_pool_indices",
+            _bs,
+            torch.int64,
+            axis="bs",
+            padding_policy=PaddingPolicy.ZERO,
+        ),
+        GraphSlot(
+            "seq_lens",
+            _bs,
+            torch.int64,
+            axis="bs",
+            padding_policy=PaddingPolicy.FILL_SENTINEL,
+            pad_value=seq_len_fill_value,
+        ),
+        GraphSlot(
+            "seq_lens_cpu",
+            _bs,
+            torch.int64,
+            axis="bs",
+            device=torch.device("cpu"),
+            padding_policy=PaddingPolicy.FILL_SENTINEL,
+            pad_value=seq_len_fill_value,
+        ),
+        GraphSlot(
+            "mrope_positions",
+            lambda _bs2, mt: (3, mt),
+            torch.int64,
+            axis="tokens",
+            slice_fn=lambda buf, n: buf[:, :n],
+            padding_policy=PaddingPolicy.ZERO,
+        ),
+    ]
+    if enable_mamba_track:
+        slots.append(
+            GraphSlot(
+                "mamba_track_indices",
+                _bs,
+                torch.int64,
+                axis="bs",
+                padding_policy=PaddingPolicy.ZERO,
+            )
+        )
+        slots.append(
+            GraphSlot(
+                "mamba_track_mask",
+                _bs,
+                torch.bool,
+                axis="bs",
+                padding_policy=PaddingPolicy.ZERO,
+            )
+        )
+    if is_encoder_decoder:
+        # Initialized once to encoder_len_fill_value, copied head-only, never
+        # reset per iter — matching the legacy DecodeInputBuffers behavior.
+        slots.append(
+            GraphSlot(
+                "encoder_lens",
+                _bs,
+                encoder_lens_dtype,
+                axis="bs",
+                padding_policy=PaddingPolicy.FILL_ONCE,
+                pad_value=encoder_len_fill_value,
+            )
+        )
+    if enable_num_token_non_padded:
+        from sglang.srt.model_executor.forward_batch_info import (
+            compute_local_num_token_non_padded,
+        )
+
+        def _num_token_non_padded_post_fill(buf, fb, ctx):
+            # Gathered (DP) path overwrites the plain FB copy with this rank's
+            # local count; the non-gathered path keeps the copied value.
+            if require_gathered_buffer and not enable_prefill_cp:
+                buf.copy_(
+                    compute_local_num_token_non_padded(
+                        global_num_token_non_padded=fb.num_token_non_padded,
+                        num_tokens_per_dp=ctx.padded_num_tokens,
+                    )
+                )
+
+        slots.append(
+            GraphSlot(
+                "num_token_non_padded",
+                lambda _bs, _mt: (1,),
+                torch.int32,
+                axis="none",
+                post_fill=_num_token_non_padded_post_fill,
+            )
+        )
+
+    # Computed slots, always exposed by extract_buffer; callers that already set
+    # global_num_tokens_* on the batch pass register_global_num_tokens=False.
+    if register_global_num_tokens:
+
+        def _global_num_tokens_post_fill(buf, fb, ctx):
+            # Only the gathered (DP) path writes a value; otherwise left as init.
+            if require_gathered_buffer:
+                buf.fill_(ctx.padded_num_tokens)
+
+        _global_shape = (
+            (lambda _bs, _mt: (dp_size,))
+            if require_mlp_tp_gather
+            else (lambda _bs, _mt: (1,))
+        )
+        for _global_name in (
+            "global_num_tokens_gpu",
+            "global_num_tokens_for_logprob_gpu",
+        ):
+            slots.append(
+                GraphSlot(
+                    _global_name,
+                    _global_shape,
+                    torch.int32,
+                    axis="none",
+                    copy_from_fb=False,
+                    post_fill=_global_num_tokens_post_fill,
+                )
+            )
+
+    for slot in slots:
+        bind = None
+        if source is not None:
+            bind = getattr(source, slot.name, None)
+            if bind is None:
+                raise ValueError(
+                    f"source is missing buffer {slot.name!r} required by the "
+                    "decode registry; cannot adopt."
+                )
+        reg.register_slot(slot, bind=bind)
+
+    # Structured slots whose backing storage still lives on the source object
+    # (adopt-only during migration): registered only when the source actually
+    # carries them. The per-replay copy source is a nested FB dataclass field,
+    # supplied via source_fn; head is copied (source-length slice), tail kept.
+    if source is not None:
+        ngram = getattr(source, "ngram_embedding_info", None)
+        if ngram is not None:
+
+            def _ngram_source(attr):
+                def _fn(fb, _ctx):
+                    info = getattr(fb, "ngram_embedding_info", None)
+                    return None if info is None else getattr(info, attr)
+
+                return _fn
+
+            for _attr in ("column_starts", "req_lens"):
+                backing = getattr(ngram, _attr)
+                reg.register_slot(
+                    GraphSlot(
+                        name=f"ngram_embedding_info.{_attr}",
+                        shape_fn=lambda _bs, _mt, _s=tuple(backing.shape): _s,
+                        dtype=backing.dtype,
+                        axis="none",
+                        padding_policy=PaddingPolicy.KEEP_PAD,
+                        source_fn=_ngram_source(_attr),
+                    ),
+                    bind=backing,
+                )
+
+        # Pipeline-parallel proxy tensors: a dict of per-key buffers, sourced
+        # from the out-of-band pp input on FillContext rather than the FB.
+        pp = getattr(source, "pp_proxy_tensors", None)
+        if pp is not None:
+
+            def _pp_source(key):
+                def _fn(_fb, ctx):
+                    ppx = ctx.pp_proxy_tensors
+                    return None if ppx is None else ppx.tensors[key]
+
+                return _fn
+
+            for _key, _backing in pp.items():
+                reg.register_slot(
+                    GraphSlot(
+                        name=f"pp_proxy_tensors.{_key}",
+                        shape_fn=lambda _bs, _mt, _s=tuple(_backing.shape): _s,
+                        dtype=_backing.dtype,
+                        axis="none",
+                        padding_policy=PaddingPolicy.KEEP_PAD,
+                        source_fn=_pp_source(_key),
+                    ),
+                    bind=_backing,
+                )
+
+        # KV-canary id buffers (off by default): plain bs-axis FB copies,
+        # adopt-only when the source carries them. Head [:raw_bs] is copied;
+        # the tail keeps its init (rids_int 0, bootstrap_room_ids_int -1).
+        for _cname in ("rids_int", "bootstrap_room_ids_int"):
+            canary = getattr(source, _cname, None)
+            if canary is not None:
+                reg.register_slot(
+                    GraphSlot(
+                        name=_cname,
+                        shape_fn=lambda _bs, _mt, _s=tuple(canary.shape): _s,
+                        dtype=canary.dtype,
+                        axis="bs",
+                    ),
+                    bind=canary,
+                )
+
+    return reg
+
+
+def build_prefill_registry(
+    *,
+    device: torch.device,
+    max_bs: int,
+    max_num_token: int,
+    cache_loc_dtype: torch.dtype,
+    is_multimodal: bool = False,
+    hidden_size: int = 0,
+    embed_dtype: Optional[torch.dtype] = None,
+    enable_mamba_track: bool = False,
+    enable_num_token_non_padded: bool = False,
+    register_input_embeds: bool = True,
+    share_pool: bool = True,
+    source: Optional[Any] = None,
+) -> CudaGraphBufferRegistry:
+    """Registry mirroring the **token-axis** FB-shared buffers for the
+    piecewise / breakable (prefill) cuda-graph runners.
+
+    ``register_input_embeds`` (default ``True``) registers the multimodal
+    ``input_embeds`` slot; the eager extend path passes ``False`` so it is
+    carried from the batch (a read input) rather than written in-graph.
+
+    Padding policies match the inline copy/zero in
+    ``PiecewiseCudaGraphRunner.load_batch``: ``input_ids`` / ``positions``
+    / ``out_cache_loc`` / ``mrope_positions`` / ``input_embeds`` reset their
+    padded tail ``[raw_num_tokens:padded_num_tokens]`` to ``0`` (the padded
+    tokens *are* processed by the graph, so they must be benign), then the head
+    ``[:raw_num_tokens]`` is copied from the FB. ``input_embeds`` is not an FB
+    copy — the model writes the embeds into it inside the graph — so it is
+    reset-only (``copy_from_fb=False``). ``mamba_track_*`` are bs-axis copies
+    with no padding reset (bs is not padded on this path).
+
+    The piecewise / breakable runners pass ``source=None``, so the registry
+    allocates (and owns) these buffers directly; ``share_pool`` then coalesces
+    them through the process-wide pool. (A ``source`` object, if given, would be
+    adopted instead — one shared allocation with stable ``data_ptr``.)
+    """
+    reg = CudaGraphBufferRegistry(
+        device=device,
+        max_bs=max_bs,
+        max_num_tokens=max_num_token,
+        share_pool=share_pool,
+    )
+
+    def _tokens(_bs: int, mt: int) -> Tuple[int, ...]:
+        return (mt,)
+
+    def _bs(bs: int, _mt: int) -> Tuple[int, ...]:
+        return (bs,)
+
+    slots = [
+        GraphSlot(
+            "input_ids",
+            _tokens,
+            torch.int64,
+            axis="tokens",
+            padding_policy=PaddingPolicy.ZERO,
+        ),
+        GraphSlot(
+            "positions",
+            _tokens,
+            torch.int64,
+            axis="tokens",
+            padding_policy=PaddingPolicy.ZERO,
+        ),
+        GraphSlot(
+            "out_cache_loc",
+            _tokens,
+            cache_loc_dtype,
+            axis="tokens",
+            padding_policy=PaddingPolicy.ZERO,
+        ),
+    ]
+    if is_multimodal:
+        slots.append(
+            GraphSlot(
+                "mrope_positions",
+                lambda _bs2, mt: (3, mt),
+                torch.int64,
+                axis="tokens",
+                padding_policy=PaddingPolicy.ZERO,
+                slice_fn=lambda buf, n: buf[:, :n],
+            )
+        )
+        if register_input_embeds:
+            slots.append(
+                GraphSlot(
+                    "input_embeds",
+                    lambda _bs2, mt: (mt, hidden_size),
+                    embed_dtype,
+                    axis="tokens",
+                    padding_policy=PaddingPolicy.ZERO,
+                    copy_from_fb=False,
+                )
+            )
+    if enable_mamba_track:
+        slots.append(GraphSlot("mamba_track_indices", _bs, torch.int64, axis="bs"))
+        slots.append(GraphSlot("mamba_track_mask", _bs, torch.bool, axis="bs"))
+        slots.append(GraphSlot("mamba_track_seqlens", _bs, torch.int32, axis="bs"))
+    if enable_num_token_non_padded:
+        slots.append(
+            GraphSlot(
+                "num_token_non_padded",
+                lambda _bs2, _mt: (1,),
+                torch.int32,
+                axis="none",
+            )
+        )
+
+    for slot in slots:
+        bind = None
+        if source is not None:
+            bind = getattr(source, slot.name, None)
+            if bind is None:
+                raise ValueError(
+                    f"source is missing buffer {slot.name!r} required by the "
+                    "prefill registry; cannot adopt."
+                )
+        reg.register_slot(slot, bind=bind)
+    return reg
+
+
+def build_eager_registry(
+    *,
+    device: torch.device,
+    max_bs: int,
+    max_num_token: int,
+    cache_loc_dtype: torch.dtype,
+    enable_mamba_track: bool = False,
+    is_encoder_decoder: bool = False,
+    encoder_len_fill_value: int = 0,
+    encoder_lens_dtype: torch.dtype = torch.int32,
+    dp_size: int = 1,
+) -> CudaGraphBufferRegistry:
+    """One fixed-max input registry for the ``EagerRunner``, serving BOTH eager
+    decode and eager prefill.
+
+    The decode slot set is a superset of eager prefill's needs (eager prefill
+    carries ``input_embeds`` from the batch and reads the bs-axis fields live),
+    so we reuse it, sized at ``(max_bs, max_num_token)`` where ``max_num_token``
+    is the prefill token ceiling. ``seq_len_fill_value=0`` because eager never
+    pads, so the sentinel tail is never read.
+
+    ``share_pool=True`` so same-named / same-size slots coalesce through the
+    process-wide pool. The ``EagerRunner`` is built before the cuda-graph runners
+    (see ``ModelRunner.init_backends``), so its (largest) allocations are
+    canonical and the cg runners' matching slots (prefill's token-axis at
+    ``max_num_token``, decode's bs-axis at ``max_bs``) adopt them.
+    """
+    return build_decode_registry(
+        device=device,
+        max_bs=max_bs,
+        max_num_token=max_num_token,
+        seq_len_fill_value=0,
+        cache_loc_dtype=cache_loc_dtype,
+        enable_mamba_track=enable_mamba_track,
+        is_encoder_decoder=is_encoder_decoder,
+        encoder_len_fill_value=encoder_len_fill_value,
+        encoder_lens_dtype=encoder_lens_dtype,
+        enable_num_token_non_padded=False,
+        register_global_num_tokens=False,
+        require_gathered_buffer=False,
+        require_mlp_tp_gather=False,
+        dp_size=dp_size,
+        share_pool=True,
+        source=None,
+    )

python/sglang/srt/model_executor/cuda_graph_config.py

@@ -0,0 +1,245 @@
+# 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.
+# ==============================================================================
+"""Phase / backend identifiers, the canonical default for
+cuda_graph_config, and the --cuda-graph-config JSON CLI parser.
+
+Module-level imports are pure stdlib — no torch / sglang.srt deps — so
+ServerArgs can import everything here without pulling in backend
+classes. check_cuda_graph_backend lazy-imports get_server_args
+inside the function body to preserve that invariant.
+"""
+
+import argparse
+import dataclasses
+import json
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional
+
+
+class Phase:
+    """The two phases of model forward."""
+
+    DECODE = "decode"
+    PREFILL = "prefill"
+    ALL = (DECODE, PREFILL)
+
+
+class Backend:
+    """CUDA graph capture backends a phase can use."""
+
+    FULL = "full"
+    BREAKABLE = "breakable"
+    TC_PIECEWISE = "tc_piecewise"
+    DISABLED = "disabled"
+    ALL = (FULL, BREAKABLE, TC_PIECEWISE, DISABLED)
+
+
+ALLOWED_BACKENDS_PER_PHASE = {
+    Phase.DECODE: (
+        Backend.FULL,
+        Backend.BREAKABLE,
+        Backend.TC_PIECEWISE,
+        Backend.DISABLED,
+    ),
+    # full for prefill captures one whole-forward graph per num_tokens
+    # bucket (bs=1 only); replay pads num_tokens up to the nearest
+    # captured bucket. Opt-in: the padding waste is the operator's call.
+    Phase.PREFILL: (
+        Backend.FULL,
+        Backend.BREAKABLE,
+        Backend.TC_PIECEWISE,
+        Backend.DISABLED,
+    ),
+}
+
+# Per-phase settings schema. Keys other than backend are runner-level
+# (read by any backend in that phase); tc_compiler is the lone
+# backend-specific knob (only meaningful when backend == tc_piecewise).
+# For prefill, bs carries the captured shape size (token count for
+# tc_piecewise, request count for breakable) — one shape knob per phase.
+# full_prefill_max_req is prefill-only and only meaningful when backend == full.
+ALLOWED_KEYS_PER_PHASE = {
+    Phase.DECODE: ("backend", "max_bs", "bs", "tc_compiler"),
+    Phase.PREFILL: ("backend", "max_bs", "bs", "tc_compiler", "full_prefill_max_req"),
+}
+
+
+@dataclass
+class PhaseConfig:
+    """Per-phase CUDA graph settings."""
+
+    backend: str = Backend.DISABLED
+    max_bs: Optional[int] = None
+    bs: Optional[List[int]] = None
+    # Only meaningful when backend == tc_piecewise; ignored otherwise.
+    tc_compiler: str = "eager"
+    # Only meaningful for the prefill phase with backend == full: max number of
+    # request slots baked into each captured graph. Real bs <= full_prefill_max_req
+    # reuses the graph (unused slots become zero-length sentinels); larger
+    # batches fall back to eager. Ignored by BCG (bs=1 only) and TC_PIECEWISE
+    # (bs-invariant via torch.compile). None auto-derives chunked_prefill_size // 512.
+    full_prefill_max_req: Optional[int] = None
+
+
+def default_prefill_backend() -> str:
+    """BCG (breakable) is the prefill default on CUDA only; other platforms
+    (HIP/NPU/...) keep tc_piecewise until BCG is validated there. Lazy import
+    keeps this module's stdlib-only import invariant (see module docstring)."""
+    from sglang.srt.utils import is_cuda
+
+    return Backend.BREAKABLE if is_cuda() else Backend.TC_PIECEWISE
+
+
+@dataclass
+class CudaGraphConfig:
+    """Top-level CUDA graph config: one PhaseConfig per phase."""
+
+    decode: PhaseConfig = field(
+        default_factory=lambda: PhaseConfig(backend=Backend.FULL)
+    )
+    prefill: PhaseConfig = field(
+        default_factory=lambda: PhaseConfig(backend=default_prefill_backend())
+    )
+
+    def __getitem__(self, phase: str) -> PhaseConfig:
+        """Phase-string lookup; kept for migration ergonomics."""
+        if phase not in Phase.ALL:
+            raise KeyError(phase)
+        return getattr(self, phase)
+
+    def to_dict(self) -> Dict[str, Dict[str, Any]]:
+        # Diff-only, not asdict: the parser locks every (phase, key) it sees,
+        # so emitting defaults would lock fields the caller never set.
+        baseline = default_cuda_graph_config()
+        return {
+            Phase.DECODE: _diff_phase(self.decode, baseline.decode),
+            Phase.PREFILL: _diff_phase(self.prefill, baseline.prefill),
+        }
+
+    @classmethod
+    def from_dict(cls, raw: Optional[Dict[str, Dict[str, Any]]]) -> "CudaGraphConfig":
+        """Build from a (partial) dict of overrides, defaults fill the rest.
+        Unknown phases / keys are silently dropped — the JSON-input
+        validator (parse_cuda_graph_config_arg) rejects them upstream."""
+        cfg = cls()
+        if not raw:
+            return cfg
+        for phase, phase_settings in raw.items():
+            if phase not in Phase.ALL or not isinstance(phase_settings, dict):
+                continue
+            phase_cfg = getattr(cfg, phase)
+            allowed = ALLOWED_KEYS_PER_PHASE[phase]
+            for key, value in phase_settings.items():
+                if key in allowed:
+                    setattr(phase_cfg, key, value)
+        return cfg
+
+
+def default_cuda_graph_config() -> CudaGraphConfig:
+    """Fresh CudaGraphConfig populated with canonical defaults."""
+    return CudaGraphConfig()
+
+
+def _diff_phase(actual: PhaseConfig, baseline: PhaseConfig) -> Dict[str, Any]:
+    """Return only fields whose value differs from the per-phase default."""
+    return {
+        f.name: getattr(actual, f.name)
+        for f in dataclasses.fields(actual)
+        if getattr(actual, f.name) != getattr(baseline, f.name)
+    }
+
+
+def check_cuda_graph_backend(phase: str, backend: str) -> bool:
+    """True if cuda_graph_config[phase].backend == backend on the
+    global server args. Returns False if the global server args have not
+    been initialized yet (e.g. unit tests, early startup)."""
+    from sglang.srt.runtime_context import get_server_args
+
+    try:
+        server_args = get_server_args()
+    except ValueError:
+        return False
+    cfg = server_args.cuda_graph_config
+    if cfg is None or phase not in Phase.ALL:
+        return False
+    return getattr(cfg, phase).backend == backend
+
+
+def cuda_graph_fully_disabled() -> bool:
+    """True iff cuda_graph_config has Backend.DISABLED on every phase.
+
+    Use at sites that ask the legacy server_args.disable_cuda_graph
+    question ("no CG anywhere globally") — e.g., preallocating buffers
+    that any captured graph would otherwise reuse, or one-shot init
+    that's a no-op when CG is completely off.
+    """
+    return check_cuda_graph_backend(
+        Phase.DECODE, Backend.DISABLED
+    ) and check_cuda_graph_backend(Phase.PREFILL, Backend.DISABLED)
+
+
+def parse_cuda_graph_config_arg(raw: str) -> Dict[str, Dict[str, Any]]:
+    """argparse type for --cuda-graph-config: parse JSON dict of
+    phase → settings dict. Each phase's settings dict is itself validated
+    against ALLOWED_KEYS_PER_PHASE. Returns a plain dict — the
+    precedence pipeline in ServerArgs converts to CudaGraphConfig
+    after merging."""
+    try:
+        parsed = json.loads(raw)
+    except json.JSONDecodeError as e:
+        raise argparse.ArgumentTypeError(f"--cuda-graph-config must be JSON: {e}")
+    if not isinstance(parsed, dict):
+        raise argparse.ArgumentTypeError(
+            f"--cuda-graph-config must be a JSON object, got {type(parsed).__name__}"
+        )
+
+    result: Dict[str, Dict[str, Any]] = {}
+    for phase, phase_settings in parsed.items():
+        phase = str(phase)
+        if phase not in Phase.ALL:
+            raise argparse.ArgumentTypeError(
+                f"--cuda-graph-config: unknown phase '{phase}', expected one of {Phase.ALL}"
+            )
+        if not isinstance(phase_settings, dict):
+            raise argparse.ArgumentTypeError(
+                f"--cuda-graph-config['{phase}'] must be a JSON object, got "
+                f"{type(phase_settings).__name__}"
+            )
+        allowed = ALLOWED_KEYS_PER_PHASE[phase]
+        result[phase] = {}
+        for key, value in phase_settings.items():
+            if key not in allowed:
+                raise argparse.ArgumentTypeError(
+                    f"--cuda-graph-config['{phase}']: unknown key '{key}', expected one of {allowed}"
+                )
+            result[phase][key] = value
+    return result
+
+
+def explicit_keys_in(
+    settings: Optional[Dict[str, Dict[str, Any]]],
+) -> set:
+    """Return the set of (phase, key) tuples present in settings
+    (the raw dict form, as it arrives from CLI/SDK). Used by ServerArgs
+    to track keys the user explicitly set so the auto-disable cascade can
+    skip them."""
+    out: set = set()
+    if not settings:
+        return out
+    for phase, phase_settings in settings.items():
+        if not isinstance(phase_settings, dict):
+            continue
+        for key in phase_settings.keys():
+            out.add((phase, key))
+    return out

python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py

@@ -0,0 +1,225 @@
+# Mixin class for metadata management of Deepseek MHA forward (chunked prefix cache)
+# More details can be found in python/sglang/srt/models/deepseek_common/attention_forward_methods/forward_mha.py
+
+from typing import List, Optional
+
+import torch
+
+from sglang.kernels.ops.kvcache.kv_indices import (
+    create_chunked_prefix_cache_kv_indices,
+    create_flashinfer_kv_indices_triton,
+)
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.forward_context import (
+    get_req_to_token_pool,
+    get_token_to_kv_pool,
+)
+
+
+class ForwardBatchDeepSeekMHAMixin:
+    # For MLA chunked prefix cache used in chunked prefill
+    # Tell attention backend whether the kv cache needs to be attended in current pass
+    attn_attend_prefix_cache: Optional[bool] = None
+    # Number of prefix cache chunks
+    num_prefix_chunks: Optional[int] = None
+    # Index of current chunk, used by attention backend
+    prefix_chunk_idx: Optional[int] = None
+    # Maximum number of tokens in each chunk per sequence. Computed from maximum chunk capacity
+    prefix_chunk_len: Optional[int] = None
+    # Start positions of prefix cache for each chunk, (num_prefix_chunks, batch_size)
+    prefix_chunk_starts: Optional[torch.Tensor] = None
+    # Start positions of prefix cache for each chunk, (num_prefix_chunks, batch_size), need prefix_chunk_starts_cpu for dcp all gather kv cache
+    prefix_chunk_starts_cpu: Optional[torch.Tensor] = None
+    # length of prefix cache for each chunk, (num_prefix_chunks, batch_size)
+    prefix_chunk_seq_lens_cpu: Optional[torch.Tensor] = None
+    # Lengths of prefix cache for each chunk, (num_prefix_chunks, batch_size)
+    prefix_chunk_seq_lens: Optional[torch.Tensor] = None
+    # Accumulated lengths of prefix cache for each chunk, (num_prefix_chunks, batch_size + 1)
+    prefix_chunk_cu_seq_lens: Optional[torch.Tensor] = None
+    # Max lengths of prefix cache for each chunk, (num_prefix_chunks,)
+    prefix_chunk_max_seq_lens: Optional[List[int]] = None
+    # Per-chunk flag: True if any sequence has kv_len==0 in that chunk.
+    # Precomputed on CPU to avoid GPU-CPU sync in the hot path.
+    prefix_chunk_has_zero_kv: Optional[List[bool]] = None
+    # Number of tokens in each prefix cache chunk, (num_prefix_chunks,)
+    prefix_chunk_num_tokens: Optional[List[int]] = None
+    # KV Indices for each chunk
+    prefix_chunk_kv_indices: Optional[List[torch.Tensor]] = None
+    # For MLA chunked prefix cache used in chunked prefill
+    # Tell attention backend whether lse needs to be returned
+    mha_return_lse: Optional[bool] = None
+    # Whether to apply MHA_ONE_SHOT forward method
+    mha_one_shot: Optional[bool] = None
+    # KV Indices for MHA_ONE_SHOT forward method
+    mha_one_shot_kv_indices: Optional[torch.Tensor] = None
+
+    def get_max_chunk_capacity(self):
+        return envs.SGLANG_MAX_KV_CHUNK_CAPACITY.get()
+
+    def set_prefix_chunk_idx(self, idx: int):
+        self.prefix_chunk_idx = idx
+
+    def set_attn_attend_prefix_cache(self, attn_attend_prefix_cache: bool):
+        self.attn_attend_prefix_cache = attn_attend_prefix_cache
+
+    def prepare_chunked_kv_indices(self, device: torch.device):
+        self.prefix_chunk_kv_indices = []
+        req_to_token = get_req_to_token_pool().req_to_token
+        for idx in range(self.num_prefix_chunks):
+            chunk_starts = self.prefix_chunk_starts[idx]
+            chunk_seq_lens = self.prefix_chunk_seq_lens[idx]
+            chunk_cu_seq_lens = self.prefix_chunk_cu_seq_lens[idx]
+            num_chunk_tokens = self.prefix_chunk_num_tokens[idx]
+
+            chunk_kv_indices = torch.empty(
+                num_chunk_tokens, dtype=torch.int32, device=device
+            )
+
+            create_chunked_prefix_cache_kv_indices[(self.batch_size,)](
+                req_to_token,
+                self.req_pool_indices,
+                chunk_starts,
+                chunk_seq_lens,
+                chunk_cu_seq_lens,
+                chunk_kv_indices,
+                req_to_token.shape[1],
+            )
+            self.prefix_chunk_kv_indices.append(chunk_kv_indices)
+
+    # Here we suppose the length of each chunk is equal
+    # For example, if we have 4 sequences with prefix length [256, 512, 768, 1024], prefix_chunk_len = 256
+    # num_prefix_chunks = cdiv(1024, 256) = 4
+    # prefix_chunk_starts = [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512], [768, 768, 768, 768]]
+    # prefix_chunk_ends = [[256, 256, 256, 256], [256, 512, 512, 512], [256, 512, 768, 768], [256, 512, 768, 1024]]
+    # prefix_chunk_seq_lens = [[256, 256, 256, 256], [0, 256, 256, 256], [0, 0, 256, 256], [0, 0, 0, 256]]
+    # TODO: Implement a better way to allocate chunk lengths that uses memory spaces more efficiently.
+    def get_prefix_chunk_seq_lens(
+        self, prefix_lens: torch.Tensor, num_prefix_chunks: int, prefix_chunk_len: int
+    ):
+        device = prefix_lens.device
+        prefix_chunk_starts = (
+            torch.arange(num_prefix_chunks, device=device, dtype=torch.int32)
+            .unsqueeze(1)
+            .expand(-1, self.batch_size)
+            * prefix_chunk_len
+        )
+        prefix_chunk_ends = torch.min(
+            prefix_lens.unsqueeze(0),
+            prefix_chunk_starts + prefix_chunk_len,
+        ).to(torch.int32)
+
+        prefix_chunk_seq_lens = (
+            (prefix_chunk_ends - prefix_chunk_starts).clamp(min=0).to(torch.int32)
+        )
+
+        return prefix_chunk_starts, prefix_chunk_seq_lens
+
+    # Called before each attention module if using chunked kv cache for prefill
+    # Some of the codes are adapted from https://github.com/vllm-project/vllm/blob/main/vllm/v1/attention/backends/mla/common.py
+    def prepare_chunked_prefix_cache_info(self, device: torch.device):
+
+        from sglang.srt.mem_cache.memory_pool import (
+            HybridLinearKVPool,
+            MLATokenToKVPool,
+        )
+
+        token_to_kv_pool = get_token_to_kv_pool()
+        assert isinstance(token_to_kv_pool, MLATokenToKVPool) or (
+            isinstance(token_to_kv_pool, HybridLinearKVPool)
+            and isinstance(token_to_kv_pool.full_kv_pool, MLATokenToKVPool)
+        ), "Currently chunked prefix cache can only be used by Deepseek models"
+
+        if not any(self.extend_prefix_lens_cpu):
+            self.num_prefix_chunks = 0
+            return
+
+        if self.prefix_chunk_len is not None:
+            # Chunked kv cache info already prepared by prior modules
+            return
+
+        self.prefix_chunk_idx = -1
+
+        # chunk_capacity is the maximum number of tokens in each chunk
+        chunk_capacity = self.get_max_chunk_capacity()
+        self.prefix_chunk_len = chunk_capacity // self.batch_size
+
+        self.num_prefix_chunks = (
+            max(self.extend_prefix_lens_cpu) + self.prefix_chunk_len - 1
+        ) // self.prefix_chunk_len
+
+        # Here we compute chunk lens twice to avoid stream sync, once on gpu and once on cpu.
+        prefix_chunk_starts_cuda, prefix_chunk_seq_lens_cuda = (
+            self.get_prefix_chunk_seq_lens(
+                self.extend_prefix_lens,
+                self.num_prefix_chunks,
+                self.prefix_chunk_len,
+            )
+        )
+        prefix_chunk_starts_cpu, prefix_chunk_seq_lens_cpu = (
+            self.get_prefix_chunk_seq_lens(
+                torch.tensor(self.extend_prefix_lens_cpu),
+                self.num_prefix_chunks,
+                self.prefix_chunk_len,
+            )
+        )
+        self.prefix_chunk_starts = prefix_chunk_starts_cuda
+        self.prefix_chunk_seq_lens = prefix_chunk_seq_lens_cuda
+
+        # set prefix_chunk_starts_cpu and prefix_chunk_seq_lens_cpu for dcp to gather chunk kv cache with arbitrary lens
+        self.prefix_chunk_starts_cpu = prefix_chunk_starts_cpu
+        self.prefix_chunk_seq_lens_cpu = prefix_chunk_seq_lens_cpu
+        # Metadata for attention backend
+        self.prefix_chunk_cu_seq_lens = torch.zeros(
+            self.num_prefix_chunks,
+            self.batch_size + 1,
+            device=device,
+            dtype=torch.int32,
+        )
+        self.prefix_chunk_cu_seq_lens[:, 1:] = prefix_chunk_seq_lens_cuda.cumsum(
+            dim=1
+        ).to(torch.int32)
+        self.prefix_chunk_max_seq_lens = prefix_chunk_seq_lens_cpu.max(
+            dim=1
+        ).values.tolist()
+
+        self.prefix_chunk_num_tokens = prefix_chunk_seq_lens_cpu.sum(dim=1).tolist()
+        assert max(self.prefix_chunk_num_tokens) <= self.get_max_chunk_capacity()
+
+        # Per-chunk flag: does any sequence have kv_len == 0?
+        # Pure CPU check (prefix_chunk_seq_lens_cpu is on CPU), no GPU sync.
+        self.prefix_chunk_has_zero_kv = [
+            bool((prefix_chunk_seq_lens_cpu[i] == 0).any())
+            for i in range(self.num_prefix_chunks)
+        ]
+
+        # Precompute the kv indices for each chunk
+        self.prepare_chunked_kv_indices(device)
+
+    def fetch_mha_one_shot_kv_indices(self):
+        if self.mha_one_shot_kv_indices is not None:
+            return self.mha_one_shot_kv_indices
+        batch_size = self.batch_size
+        paged_kernel_lens_sum = sum(self.seq_lens_cpu)
+        kv_indices = torch.empty(
+            paged_kernel_lens_sum,
+            dtype=torch.int32,
+            device=self.req_pool_indices.device,
+        )
+        kv_indptr = torch.zeros(
+            batch_size + 1,
+            dtype=torch.int32,
+            device=self.req_pool_indices.device,
+        )
+        kv_indptr[1:] = torch.cumsum(self.seq_lens, dim=0)
+        req_to_token = get_req_to_token_pool().req_to_token
+        create_flashinfer_kv_indices_triton[(self.batch_size,)](
+            req_to_token,
+            self.req_pool_indices,
+            self.seq_lens,
+            kv_indptr,
+            None,
+            kv_indices,
+            req_to_token.shape[1],
+        )
+        self.mha_one_shot_kv_indices = kv_indices
+        return kv_indices

python/sglang/srt/model_executor/forward_context.py

@@ -0,0 +1,84 @@
+"""Per-forward-call control context.
+
+Owns ForwardContext — a frozen dataclass holding control configs the model
+layer reads at depth via get_forward_context(). The only mandatory field
+today is attn_backend; pool refs are derived from attn_backend.*
+(every backend caches them at __init__), so a published ForwardContext
+is enough to resolve the active pools without a separate global.
+
+ModelRunner._forward_raw publishes a fresh ForwardContext for the
+duration of each forward; callers that need a per-call override (PDmux
+per-stream backend, frozen-KV MTP draft loop, TBO per-child dispatch) use
+dataclasses.replace and wrap the override scope with forward_context().
+
+Distinct from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph.TcPiecewiseForwardContext,
+which collects compilation-time refs for the piecewise CUDA graph backend.
+
+Concurrency: _current is a plain module-level global, not thread-local.
+This matches the global_server_args precedent and is safe because each
+forward runs synchronously on a single Python thread per worker process. If
+worker threads ever share a process, migrate to contextvars.ContextVar.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+    from sglang.srt.mem_cache.memory_pool import KVCache, ReqToTokenPool
+
+
+@dataclass(frozen=True, slots=True)
+class ForwardContext:
+    """Per-forward-call control configs. Read via get_forward_context();
+    extend by adding fields here. Frozen so accidental mutation raises at
+    write time — use dataclasses.replace for per-call overrides."""
+
+    attn_backend: AttentionBackend
+
+
+_current: Optional[ForwardContext] = None
+
+
+def set_forward_context(ctx: Optional[ForwardContext]) -> Optional[ForwardContext]:
+    """Set the active context; return the previous one for explicit
+    save/restore. Prefer the forward_context() context manager."""
+    global _current
+    prev, _current = _current, ctx
+    return prev
+
+
+def has_forward_context() -> bool:
+    return _current is not None
+
+
+def get_forward_context() -> ForwardContext:
+    assert _current is not None, (
+        "no forward context active — call forward_context(...) or set_forward_context(...) "
+        "before reading get_forward_context()."
+    )
+    return _current
+
+
+def get_attn_backend() -> AttentionBackend:
+    return get_forward_context().attn_backend
+
+
+def get_token_to_kv_pool() -> KVCache:
+    return get_attn_backend().token_to_kv_pool
+
+
+def get_req_to_token_pool() -> ReqToTokenPool:
+    return get_attn_backend().req_to_token_pool
+
+
+@contextmanager
+def forward_context(ctx: ForwardContext):
+    prev = set_forward_context(ctx)
+    try:
+        yield
+    finally:
+        set_forward_context(prev)

python/sglang/srt/model_executor/graph_shared_output.py

@@ -0,0 +1,66 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Dict, Optional
+
+import torch
+
+from sglang.srt.model_executor.cuda_graph_config import Backend
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+
+class GraphSharedOutput:
+    """``(max_rows, vocab)`` logits buffer, shared by every cuda-graph runner."""
+
+    _process_shared: Optional[GraphSharedOutput] = None
+
+    def __init__(
+        self,
+        *,
+        device: torch.device,
+        max_rows: int,
+    ) -> None:
+        self.device = torch.device(device)
+        self.max_rows = max_rows
+        self._logits_buffers: Dict[int, torch.Tensor] = {}
+
+    @classmethod
+    def create_for_model_runner(
+        cls, model_runner: ModelRunner
+    ) -> Optional[GraphSharedOutput]:
+        cuda_graph_config = model_runner.server_args.cuda_graph_config
+        if cuda_graph_config is None:
+            return None
+
+        max_rows = 0
+        decode = cuda_graph_config.decode
+        if decode.backend != Backend.DISABLED and decode.bs:
+            max_rows = max(max_rows, model_runner.max_decode_logits_rows())
+
+        if max_rows <= 0:
+            return None
+
+        device = torch.device(model_runner.device)
+        shared = cls._process_shared
+        if (
+            shared is not None
+            and shared.device == device
+            and shared.max_rows >= max_rows
+        ):
+            return shared
+        cls._process_shared = cls(device=device, max_rows=max_rows)
+        return cls._process_shared
+
+    def get_logits_buffer(self, vocab_size: int, *, rows: int) -> torch.Tensor:
+        assert rows <= self.max_rows, (
+            f"shared logits buffer holds {self.max_rows} rows but caller "
+            f"needs {rows} (vocab_size={vocab_size})"
+        )
+        buffer = self._logits_buffers.get(vocab_size)
+        if buffer is None:
+            buffer = torch.zeros(
+                (self.max_rows, vocab_size), dtype=torch.float, device=self.device
+            )
+            self._logits_buffers[vocab_size] = buffer
+        return buffer[:rows]

python/sglang/srt/model_executor/hook_manager.py

@@ -0,0 +1,82 @@
+import fnmatch
+import importlib
+import logging
+from typing import Any, Callable, List, Optional
+
+import torch.nn as nn
+
+logger = logging.getLogger(__name__)
+
+
+def register_forward_hooks(model: nn.Module, hook_specs: List[dict[str, Any]]) -> None:
+    """
+    hook_specs is a list of dicts from server_args.forward_hooks.
+    Attaches forward hooks to the matching modules.
+    """
+    name_to_module = dict(model.named_modules())
+
+    for spec in hook_specs:
+        spec_name = spec.get("name", "")
+        target_patterns = spec.get("target_modules", [])
+        if not target_patterns:
+            logger.warning(f"Hook spec '{spec_name}' has no 'target_modules', skipping")
+            continue
+
+        hook_factory_path = spec.get("hook_factory")
+        if not hook_factory_path:
+            logger.warning(f"Hook spec '{spec_name}' has no 'hook_factory', skipping")
+            continue
+
+        config = spec.get("config") or {}
+        hook_factory = resolve_callable(hook_factory_path)
+
+        hook = hook_factory(config) if hook_factory else None
+        if hook is None:
+            logger.warning(
+                f"Hook factory '{hook_factory_path}' for spec '{spec_name}' "
+                "returned None, not registering any hook"
+            )
+            continue
+
+        # Resolve patterns like "model.layers.*.mlp"
+        matched = []
+        for name, module in name_to_module.items():
+            if any(fnmatch.fnmatch(name, pattern) for pattern in target_patterns):
+                matched.append((name, module))
+
+        if not matched:
+            logger.warning(
+                f"No modules matched hook spec '{spec_name}' "
+                f"patterns={target_patterns}"
+            )
+            continue
+
+        for module_name, module in matched:
+            _ = module.register_forward_hook(hook)
+            logger.info(f"Registered forward hook '{spec_name}' " f"on {module_name}")
+
+
+def resolve_callable(path: Optional[str]) -> Optional[Callable]:
+    if path is None:
+        return None
+
+    if ":" in path:
+        module_name, fn_name = path.split(":", 1)
+    else:
+        parts = path.split(".")
+        if len(parts) < 2:
+            raise ValueError(
+                f"Invalid hook callable path '{path}'. "
+                "Expected 'module.submodule:factory' or 'module.submodule.factory'."
+            )
+        *mod_parts, fn_name = parts
+        module_name = ".".join(mod_parts)
+
+    module = importlib.import_module(module_name)
+    try:
+        return getattr(module, fn_name)
+    except AttributeError as e:
+        raise AttributeError(
+            f"Module '{module_name}' has no attribute '{fn_name}' "
+            f"(from hook path '{path}')"
+        ) from e

python/sglang/srt/model_executor/input_buffers.py

@@ -0,0 +1,104 @@
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass, fields
+from typing import Dict, Tuple
+
+import torch
+
+from sglang.srt.utils import is_npu
+
+# Process-wide pool keyed by (name, numel, dtype, device); see share_input_buffer.
+_PoolKey = Tuple[str, int, torch.dtype, torch.device]
+_forward_input_buffer_pool: Dict[_PoolKey, torch.Tensor] = {}
+
+
+def share_input_buffer(name: str, new_buffer: torch.Tensor) -> torch.Tensor:
+    """Coalesce a buffer by ``(name, size, dtype, device)`` into the
+    process-wide input-buffer pool.
+
+    Distinct callers that request the same field ``name`` with the same
+    size/dtype/device share one physical allocation (and therefore one
+    ``data_ptr``): the first registrant's buffer becomes canonical and every
+    later identical request is returned as a view aliased onto it. Requests
+    that differ in size get their own allocation — they never reuse or displace
+    an existing entry — so the sharing *structure* is independent of
+    registration order and no already-captured buffer is ever repointed.
+
+    This pool is process-wide and governs *every* ``share_buffers()`` caller —
+    including graph runners not yet on the registry (the speculative draft /
+    draft-extend / frozen-kv-mtp / multi-layer-eagle runners), which register
+    identically-named ``input_ids`` / ``positions`` / ``out_cache_loc`` /
+    ``mrope_positions``. Cross-runner sharing is safe because those buffers are
+    filled immediately before each replay and the forwards that use them are
+    sequential / mutually exclusive.
+    """
+    key: _PoolKey = (name, new_buffer.numel(), new_buffer.dtype, new_buffer.device)
+    canonical = _forward_input_buffer_pool.get(key, None)
+    if canonical is None:
+        _forward_input_buffer_pool[key] = new_buffer
+        canonical = new_buffer
+    return canonical.as_strided(new_buffer.size(), new_buffer.stride())
+
+
+def share_input_buffers_in(obj) -> None:
+    """Pool every tensor buffer on ``obj`` (dataclass / ``SimpleNamespace``)
+    through the process-wide pool, in place. No-op on NPU; recurses into dict /
+    dataclass buffer fields (``pp_proxy_tensors`` / ``ngram_embedding_info``)."""
+    if is_npu():
+        return
+
+    for name, buffer in list(vars(obj).items()):
+        if buffer is None:
+            continue
+        if dataclasses.is_dataclass(buffer):
+            buffer = vars(buffer)
+        if isinstance(buffer, dict):
+            for sub_name, sub_buffer in buffer.items():
+                assert isinstance(
+                    sub_buffer, torch.Tensor
+                ), f"Field {name}.{sub_name} is expected to be a torch.Tensor, but got {type(sub_buffer)}."
+                buffer[sub_name] = share_input_buffer(f"{name}.{sub_name}", sub_buffer)
+        else:
+            assert isinstance(
+                buffer, torch.Tensor
+            ), f"Field {name} is expected to be a torch.Tensor, a dict of torch.Tensor, or a dataclass of torch.Tensor, but got {type(buffer)}."
+            setattr(obj, name, share_input_buffer(name, buffer))
+
+
+@dataclass
+class ForwardInputBuffers:
+
+    def _share_one_buffer(self, name: str, new_buffer: torch.Tensor) -> torch.Tensor:
+        return share_input_buffer(name, new_buffer)
+
+    def share_buffers(self):
+        # disable share input buffer on npu due to accuracy issue
+        if is_npu():
+            return
+
+        for f in fields(self):
+            name = f.name
+            buffer = getattr(self, name)
+
+            if buffer is None:
+                continue
+
+            if dataclasses.is_dataclass(buffer):
+                buffer = vars(buffer)
+
+            if isinstance(buffer, dict):
+                for sub_name, sub_buffer in buffer.items():
+                    assert isinstance(
+                        sub_buffer, torch.Tensor
+                    ), f"Field {name}.{sub_name} is expected to be a torch.Tensor, but got {type(sub_buffer)}."
+                    new_buffer = self._share_one_buffer(
+                        f"{name}.{sub_name}", sub_buffer
+                    )
+                    buffer[sub_name] = new_buffer
+            else:
+                assert isinstance(
+                    buffer, torch.Tensor
+                ), f"Field {name} is expected to be a torch.Tensor, a dict of torch.Tensor, or a dataclass of torch.Tensor, but got {type(buffer)}."
+                new_buffer = self._share_one_buffer(name, buffer)
+                setattr(self, name, new_buffer)

python/sglang/srt/model_executor/mindspore_runner.py

@@ -0,0 +1,120 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the SGLang project
+"""ms_runner launch MindSpore distributed modules."""
+
+import logging
+import multiprocessing as mp
+import os
+import sys
+from pathlib import Path
+
+import mindspore as ms
+import torch
+from mindspore._c_expression import GroupOptions
+from mindspore.communication import create_group
+
+from sglang.srt.distributed.parallel_state import _groups
+
+logger = logging.getLogger(__name__)
+
+
+class _Tmp:
+    def __init__(self):
+        self.sched_p = None
+
+    def set_sched_process(self, p):
+        self.sched_p = p
+
+    def __del__(self):
+        if self.sched_p:
+            self.sched_p.kill()
+
+
+_tmp = _Tmp()
+
+
+def _get_host_and_ip(distributed_init_method):
+    try:
+        _, ip_str, port_str = distributed_init_method.split(":")
+        ip = ip_str.split("/")[-1]
+        port = int(port_str)
+    except Exception as e:
+        raise RuntimeError(
+            "Cannot get host and port information from %s, error: %s!"
+            % (distributed_init_method, str(e))
+        )
+
+    return ip, port
+
+
+def run_scheduler_init(rank, local_rank, world_size, master_addr, master_port):
+    with open(str(Path() / "schedule.log"), "w") as scheduler_f:
+        # For Python outputs.
+        sys.stdout = scheduler_f
+        sys.stderr = scheduler_f
+        # For C++ outputs.
+        os.dup2(scheduler_f.fileno(), 1)
+        os.dup2(scheduler_f.fileno(), 2)
+        os.environ["DEVICE_ID"] = str(local_rank)
+        os.environ["MS_WORKER_NUM"] = str(world_size)
+        os.environ["MS_ROLE"] = "MS_SCHED"
+        os.environ["MS_NODE_ID"] = str(rank)
+        os.environ["MS_SCHED_HOST"] = str(master_addr)
+        os.environ["MS_SCHED_PORT"] = str(master_port)
+        # This function is blocked until the whole cluster exits.
+        ms.communication.init()
+
+
+def set_ms_parallel_env(rank, local_rank, world_size, init_method):
+    master_addr, master_port = _get_host_and_ip(init_method)
+    # change port avoiding port conflicts with torch
+    master_port = master_port + 35 if master_port < 65500 else master_port - 35
+    if not os.getenv("MS_ROLE"):
+        if rank == 0:
+            # Create a subprocess for scheduler of MindSpore, just for internal collaboration, not for collective communication
+            sched_p = mp.Process(
+                target=run_scheduler_init,
+                args=(rank, local_rank, world_size, master_addr, master_port),
+            )
+            sched_p.start()
+            global _tmp
+            _tmp.set_sched_process(sched_p)
+
+        os.environ["DEVICE_ID"] = str(local_rank)
+        os.environ["MS_WORKER_NUM"] = str(world_size)
+        os.environ["MS_ROLE"] = "MS_WORKER"
+        os.environ["MS_NODE_ID"] = str(rank)
+        os.environ["MS_SCHED_HOST"] = str(master_addr)
+        os.environ["MS_SCHED_PORT"] = str(master_port)
+
+
+def reuse_hccl_comm():
+    for group_name, group in _groups.items():
+        # Torch ProcessGroupHccl
+        device_group = group().device_group
+        hccl_comm_handle = device_group._get_backend(torch.device("npu")).get_hccl_comm(
+            group().local_rank
+        )
+        logger.info(
+            f"MindSpore reuse torch group: {device_group}, group_name: {group_name}, local rank: {group().local_rank},"
+            f"hccl communicator handle: {hex(hccl_comm_handle)}",
+        )
+        # Create MS communication group by hccl comm handle to reuse Torch group.
+        group_options = GroupOptions()
+        group_options.hccl_config = {"hccl_comm": hccl_comm_handle}
+        create_group(group_name, group().ranks, group_options)
+
+
+def init_ms_distributed(world_size, rank, local_rank, server_args, port):
+    if server_args.dist_init_addr:
+        dist_init_method = f"tcp://{server_args.dist_init_addr}"
+    else:
+        dist_init_method = f"tcp://{server_args.host}:{port}"
+    set_ms_parallel_env(rank, local_rank, world_size, dist_init_method)
+
+    ms.set_context(infer_boost="on", jit_level="O0")
+    ms.set_context(mode=ms.context.PYNATIVE_MODE)
+    ms.set_device("Ascend", local_rank)
+    ms.communication.init("hccl")
+    # After distributed job is initialized, reuse hccl comms for MindSpore.
+    reuse_hccl_comm()

python/sglang/srt/model_executor/ngram_token_table.py

@@ -0,0 +1,51 @@
+"""Utilities for updating LongCat ngram embedding token tables."""
+
+from __future__ import annotations
+
+import torch
+
+from sglang.jit_kernel.ngram_embedding import update_token_table
+
+
+def update_ngram_token_table_after_sampling(
+    *,
+    ngram_embedding_info,
+    next_token_ids: torch.Tensor,
+    req_pool_indices: torch.Tensor,
+    seq_lens: torch.Tensor,
+    batch_size: int,
+) -> bool:
+    """Update the ngram token table with sampled tokens.
+
+    Returns whether the token table was updated.
+    """
+    skip_token_table_update = ngram_embedding_info.skip_token_table_update
+    if skip_token_table_update is not None:
+        # Skip chunked (not-yet-finished) prefill requests: their sampled token
+        # is a pseudo prediction and must not pollute the token table.
+        indices = (~skip_token_table_update).nonzero(as_tuple=True)[0]
+        if indices.numel() == 0:
+            return False
+        update_token_table(
+            ne_token_table=ngram_embedding_info.token_table,
+            tokens=next_token_ids[indices].to(torch.int32),
+            row_indices=req_pool_indices[indices],
+            column_starts=seq_lens[indices].to(torch.int32),
+            req_lens=torch.ones(
+                indices.numel(), dtype=torch.int32, device=next_token_ids.device
+            ),
+            ignore_tokens=None,
+        )
+        return True
+
+    ngram_embedding_info.out_column_starts[:batch_size] = seq_lens
+    ngram_embedding_info.out_req_lens[:batch_size] = 1
+    update_token_table(
+        ne_token_table=ngram_embedding_info.token_table,
+        tokens=next_token_ids.to(torch.int32),
+        row_indices=req_pool_indices,
+        column_starts=ngram_embedding_info.out_column_starts,
+        req_lens=ngram_embedding_info.out_req_lens,
+        ignore_tokens=None,
+    )
+    return True

python/sglang/srt/model_executor/pool_configurator.py

@@ -0,0 +1,794 @@
+"""Memory pool configurators for profiling and sizing KV cache pools.
+
+Each model architecture has its own configurator that computes pool sizes
+from available GPU memory using a unified coeff+bias model:
+
+    available_bytes = max_tokens * coeff + bias
+    max_tokens = (available_bytes - bias) / coeff
+
+Two entry points, same core computation:
+- calculate_pool_sizes(available_bytes, page_size): profiling path
+- calculate_pool_sizes_from_max_tokens(max_tokens, page_size): constraint path
+"""
+
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.configs.model_config import (
+    get_dsa_index_head_dim,
+    get_minimax_sparse_attention_config,
+    get_minimax_sparse_disable_value_layer_ids,
+    get_minimax_sparse_layer_ids,
+    is_deepseek_dsa,
+    is_deepseek_v4,
+    is_minimax_sparse,
+)
+from sglang.srt.environ import envs
+from sglang.srt.mem_cache.common import get_alloc_len_per_decode
+from sglang.srt.mem_cache.deepseek_v4_memory_pool import get_compress_state_ring_size
+from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.utils.common import (
+    ceil_align,
+    ceil_div,
+    is_float4_e2m1fn_x2,
+    spec_decode_alloc_len_per_request,
+)
+
+
+@dataclass
+class MemoryPoolConfig:
+    """Resolved memory pool config, shared between target and draft workers."""
+
+    max_total_num_tokens: int
+    max_running_requests: Optional[int] = None
+    full_max_total_num_tokens: Optional[int] = None
+    swa_max_total_num_tokens: Optional[int] = None
+
+    # DSV4 compressed-attention pool sizes (target only; draft workers leave at 0).
+    c4_max_total_num_tokens: int = 0
+    c128_max_total_num_tokens: int = 0
+    c4_state_pool_size: int = 0
+    c128_state_pool_size: int = 0
+
+    mem_fraction_static: Optional[float] = None
+
+    def __post_init__(self):
+        if self.max_total_num_tokens <= 0:
+            msg = "Not enough memory. Please try to increase --mem-fraction-static."
+            if self.mem_fraction_static is not None:
+                msg += f" Current value: mem_fraction_static={self.mem_fraction_static}"
+            raise RuntimeError(msg)
+
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.model_runner import ModelRunner
+
+logger = logging.getLogger(__name__)
+
+
+def _get_dsv4_compress_state_dtype_sizes() -> tuple[int, int]:
+    dtype_name = envs.SGLANG_DSV4_COMPRESS_STATE_DTYPE.get().strip().lower()
+    if dtype_name in ("float32", "fp32"):
+        return 4, 4
+    if dtype_name in ("bfloat16", "bf16"):
+        if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
+            raise ValueError(
+                "SGLANG_DSV4_COMPRESS_STATE_DTYPE=bf16 is not supported when "
+                "SGLANG_OPT_USE_ONLINE_COMPRESS=1; online c128 state must stay float32."
+            )
+        return 2, 2
+    raise ValueError(
+        "Unsupported SGLANG_DSV4_COMPRESS_STATE_DTYPE="
+        f"{dtype_name!r}. Expected one of: float32, fp32, bfloat16, bf16."
+    )
+
+
+class MemoryPoolConfigurator:
+    """Base class for memory pool configurators.
+
+    Subclasses compute pool sizes for their architecture via coeff+bias model.
+    Both entry points return MemoryPoolConfig (with max_running_requests=None,
+    to be filled by the consumer).
+    """
+
+    def calculate_pool_sizes(
+        self, available_bytes: int, page_size: int
+    ) -> MemoryPoolConfig:
+        """Profiling path: compute pool sizes from available bytes."""
+        raise NotImplementedError
+
+    def calculate_pool_sizes_from_max_tokens(
+        self, max_total_num_tokens: int, page_size: int
+    ) -> MemoryPoolConfig:
+        """Constraint path: recalculate pool sizes from a constrained max_tokens."""
+        raise NotImplementedError
+
+    def finalize_with_max_running_requests(
+        self, config: MemoryPoolConfig
+    ) -> MemoryPoolConfig:
+        return config
+
+
+class DefaultPoolConfigurator(MemoryPoolConfigurator):
+    """Configurator for standard models: MHA, MLA, DSA, FP4.
+
+    coeff = cell_size (bytes per token across all layers)
+    bias = 0
+    """
+
+    def __init__(self, mr: ModelRunner):
+        # Determine effective number of layers for KV cache
+        if mambaish := mr.mambaish_config:
+            effective_layer_ids = [
+                i
+                for i in mambaish.full_attention_layer_ids
+                if mr.start_layer <= i < mr.end_layer
+            ]
+            num_layers = len(effective_layer_ids)
+        else:
+            num_layers = mr.num_effective_layers
+
+        self._cell_size = self._compute_cell_size(mr, num_layers)
+
+        # EAGLE/STANDALONE: scale cell_size to account for draft model KV cache.
+        # Assumes draft and target share the same per-layer KV size (head_dim,
+        # num_kv_heads, dtype), which holds for EAGLE/MTP draft models that
+        # reuse the target architecture's attention config.
+        if (
+            mr.spec_algorithm.is_eagle() or mr.spec_algorithm.is_standalone()
+        ) and not mr.is_draft_worker:
+            eagle_draft_num_layers = getattr(mr, "eagle_draft_num_layers", None)
+            if (
+                eagle_draft_num_layers is not None
+                and int(eagle_draft_num_layers) > 0
+                and int(num_layers) > 0
+            ):
+                self._cell_size = int(
+                    self._cell_size
+                    * (1 + int(eagle_draft_num_layers) / int(num_layers))
+                )
+
+        # DFLASH/DSPARK: scale cell_size to account for draft model KV cache
+        if mr.spec_algorithm.is_dflash_family() and not mr.is_draft_worker:
+            from sglang.srt.speculative.dflash_utils import (
+                scale_kv_cell_size_per_token_for_dflash,
+            )
+
+            draft_num_layers = mr.dflash_family_draft_num_layers
+            if (
+                draft_num_layers is not None
+                and int(draft_num_layers) > 0
+                and int(num_layers) > 0
+            ):
+                self._cell_size = scale_kv_cell_size_per_token_for_dflash(
+                    target_cell_size_per_token=self._cell_size,
+                    target_num_layers=int(num_layers),
+                    draft_num_layers=int(draft_num_layers),
+                )
+
+    def _compute_cell_size(self, mr: ModelRunner, num_layers: int) -> int:
+        """Compute per-token KV cache cost in bytes. Subclasses can override."""
+        # args to config cell size
+        model_config = mr.model_config
+        kv_cache_dtype = mr.kv_cache_dtype
+        from sglang.srt.layers.cp.utils import (
+            get_glm_dsa_layer_split_effective_num_layers,
+        )
+
+        effective_num_layers = get_glm_dsa_layer_split_effective_num_layers(
+            mr, num_layers
+        )
+
+        kv_size = torch._utils._element_size(kv_cache_dtype)
+        tp_size = get_parallel().attn_tp_size
+
+        if mr.use_mla_backend:
+            cell_size = (
+                (model_config.kv_lora_rank + model_config.qk_rope_head_dim)
+                * effective_num_layers
+                * kv_size
+            )
+            if is_float4_e2m1fn_x2(kv_cache_dtype):
+                # kv_scale_buffer
+                scale_block_size = 16
+                cell_size = (cell_size // 2) + (
+                    (
+                        (model_config.kv_lora_rank + model_config.qk_rope_head_dim)
+                        // scale_block_size
+                    )
+                    * effective_num_layers
+                    * kv_size
+                )
+
+            # Add indexer KV cache overhead for DSA models (DeepSeek V3.2)
+            if is_deepseek_dsa(model_config.hf_config):
+                index_head_dim = get_dsa_index_head_dim(model_config.hf_config)
+                indexer_size_per_token = (
+                    index_head_dim
+                    + index_head_dim // DSATokenToKVPool.quant_block_size * 4
+                )
+                element_size = torch._utils._element_size(
+                    DSATokenToKVPool.index_k_with_scale_buffer_dtype
+                )
+                cell_size += (
+                    indexer_size_per_token * effective_num_layers * element_size
+                )
+        elif is_minimax_sparse(model_config.hf_config):
+            # Mirrors MiniMaxSparseKVPool: main pool (K+V all layers) + indexer pool
+            # (sparse-only, single-head; kv layers store K+V, k-only layers store K).
+            sparse_cfg = get_minimax_sparse_attention_config(model_config.hf_config)
+            dense_layer_ids, sparse_layer_ids = get_minimax_sparse_layer_ids(sparse_cfg)
+            indexer_k_only_layer_ids = set(
+                get_minimax_sparse_disable_value_layer_ids(sparse_cfg)
+            )
+
+            local_dense_layer_ids = [
+                l for l in dense_layer_ids if mr.start_layer <= l < mr.end_layer
+            ]
+            local_sparse_layer_ids = [
+                l for l in sparse_layer_ids if mr.start_layer <= l < mr.end_layer
+            ]
+            num_dense = len(local_dense_layer_ids)
+            num_sparse = len(local_sparse_layer_ids)
+            num_indexer_k_only = sum(
+                1 for l in local_sparse_layer_ids if l in indexer_k_only_layer_ids
+            )
+            num_indexer_kv = num_sparse - num_indexer_k_only
+
+            kv_heads = model_config.get_num_kv_heads(get_parallel().attn_tp_size)
+            head_dim = model_config.head_dim
+            indexer_head_dim = sparse_cfg["sparse_index_dim"]
+            indexer_dtype_size = torch._utils._element_size(mr.dtype)
+
+            main_pool_bytes = (
+                (num_dense + num_sparse) * 2 * kv_heads * head_dim * kv_size
+            )
+            indexer_bytes = (
+                (num_indexer_kv * 2 + num_indexer_k_only)
+                * indexer_head_dim
+                * indexer_dtype_size
+            )
+            # FP4 scale buffer adjustment doesn't apply to MiniMax sparse:
+            # cell_size is already a sum over heterogeneous sub-pools.
+            return main_pool_bytes + indexer_bytes
+        else:
+            cell_size = (
+                model_config.get_num_kv_heads(tp_size)
+                * (model_config.head_dim + model_config.v_head_dim)
+                * effective_num_layers
+                * kv_size
+            )
+
+            if is_float4_e2m1fn_x2(kv_cache_dtype):
+                # kv_scale_buffer
+                scale_block_size = 16
+                n = model_config.get_num_kv_heads(tp_size)
+                k = model_config.head_dim
+                cell_size = (cell_size // 2) + (
+                    (n * k * effective_num_layers * 2 * kv_size) // scale_block_size
+                )
+
+        return cell_size
+
+    def calculate_pool_sizes(
+        self, available_bytes: int, page_size: int
+    ) -> MemoryPoolConfig:
+        max_total_num_tokens = available_bytes // self._cell_size
+        max_total_num_tokens = max_total_num_tokens // page_size * page_size
+        return MemoryPoolConfig(max_total_num_tokens=max_total_num_tokens)
+
+    def calculate_pool_sizes_from_max_tokens(
+        self, max_total_num_tokens: int, page_size: int
+    ) -> MemoryPoolConfig:
+        max_total_num_tokens = max_total_num_tokens // page_size * page_size
+        return MemoryPoolConfig(max_total_num_tokens=max_total_num_tokens)
+
+
+class HybridSWAPoolConfigurator(MemoryPoolConfigurator):
+    """Configurator for hybrid sliding window attention models (Gemma2, Command-R, MiMo).
+
+    Splits available memory between full attention and SWA pools.
+    Does NOT inherit DefaultPoolConfigurator — different coeff model.
+    """
+
+    def __init__(self, mr: ModelRunner):
+        model_config = mr.model_config
+        kv_cache_dtype = mr.kv_cache_dtype
+        kv_size = torch._utils._element_size(kv_cache_dtype)
+        tp_size = get_parallel().attn_tp_size
+
+        self._full_layers_num = len(model_config.full_attention_layer_ids)
+        self._swa_layers_num = len(model_config.swa_attention_layer_ids)
+        assert (
+            self._swa_layers_num > 0
+        ), "Hybrid SWA model must have at least one SWA layer"
+
+        self._swa_full_tokens_ratio = mr.server_args.swa_full_tokens_ratio
+
+        # Full layer per-token memory (bytes)
+        self._full_per_token = (
+            model_config.get_num_kv_heads(tp_size)
+            * (model_config.head_dim + model_config.v_head_dim)
+            * kv_size
+        )
+
+        # SWA layer per-token memory (bytes)
+        self._swa_per_token = (
+            model_config.get_swa_num_kv_heads(tp_size)
+            * (model_config.swa_head_dim + model_config.swa_v_head_dim)
+            * kv_size
+        )
+
+        # EAGLE/STANDALONE draft KV pool inherits max_total tokens with its
+        # full-attn layers; budget into the full term.
+        self._draft_full_layers_num = 0
+        if (
+            mr.spec_algorithm.is_eagle() or mr.spec_algorithm.is_standalone()
+        ) and not mr.is_draft_worker:
+            draft_layers = getattr(mr, "eagle_draft_num_layers", None)
+            if draft_layers is not None and int(draft_layers) > 0:
+                self._draft_full_layers_num = int(draft_layers)
+
+        # Bytes per token of max_total_num_tokens.
+        #
+        # Hybrid (full_layers > 0): max_total = full_tokens, so cell_size accounts
+        # for both pools: F*nf + r*S*ns (where swa_tokens = full_tokens * r).
+        #
+        # All-SWA (full_layers == 0): max_total = swa_tokens directly. The ratio
+        # is meaningless here -- there is no full pool to relate to, and every
+        # token beyond the sliding window can be evicted. So cell_size = S*ns,
+        # with no ratio factor applied.
+        if self._full_layers_num == 0:
+            self._cell_size = (
+                self._swa_per_token * self._swa_layers_num
+                + self._full_per_token * self._draft_full_layers_num
+            )
+        else:
+            self._cell_size = (
+                self._full_per_token
+                * (self._full_layers_num + self._draft_full_layers_num)
+                + self._swa_full_tokens_ratio
+                * self._swa_per_token
+                * self._swa_layers_num
+            )
+
+    def _solve_pool_sizes(
+        self, max_total_num_tokens: int, page_size: int
+    ) -> MemoryPoolConfig:
+        """Core computation: split max_total_num_tokens into full/swa pool sizes."""
+
+        def align_page_size(x: int) -> int:
+            return (x // page_size) * page_size
+
+        if self._full_layers_num == 0:
+            # All-SWA: no full pool, max_total = actual SWA pool size.
+            # Ratio is not applied -- see __init__ comment.
+            swa_tokens = align_page_size(max_total_num_tokens)
+            logger.info(
+                f"Use sliding window memory pool (all SWA). "
+                f"swa_layer_tokens={swa_tokens}"
+            )
+            return MemoryPoolConfig(
+                max_total_num_tokens=swa_tokens,
+                full_max_total_num_tokens=0,
+                swa_max_total_num_tokens=swa_tokens,
+            )
+
+        # Hybrid: full_tokens = max_total_num_tokens, swa_tokens = full_tokens * ratio
+        full_tokens = align_page_size(max_total_num_tokens)
+        swa_tokens = align_page_size(int(full_tokens * self._swa_full_tokens_ratio))
+
+        logger.info(
+            f"Use sliding window memory pool. "
+            f"full_layer_tokens={full_tokens}, swa_layer_tokens={swa_tokens}"
+        )
+
+        return MemoryPoolConfig(
+            max_total_num_tokens=full_tokens,
+            full_max_total_num_tokens=full_tokens,
+            swa_max_total_num_tokens=swa_tokens,
+        )
+
+    def calculate_pool_sizes(
+        self, available_bytes: int, page_size: int
+    ) -> MemoryPoolConfig:
+        max_total_num_tokens = int(available_bytes // self._cell_size)
+        return self._solve_pool_sizes(max_total_num_tokens, page_size)
+
+    def calculate_pool_sizes_from_max_tokens(
+        self, max_total_num_tokens: int, page_size: int
+    ) -> MemoryPoolConfig:
+        return self._solve_pool_sizes(max_total_num_tokens, page_size)
+
+
+class SWAChunkCapPoolConfigurator(HybridSWAPoolConfigurator):
+    """Hybrid SWA configurator with the SWA pool sized from a fixed token cap.
+
+    When max_running_requests is explicit, the SWA pool's worst-case
+    footprint is bounded per request. The SWA pool is sized tightly from that
+    cap and the freed memory is redirected to the full pool, instead of sizing
+    both pools by swa_full_tokens_ratio.
+    """
+
+    def __init__(self, mr: ModelRunner):
+        super().__init__(mr)
+        assert self._full_layers_num > 0
+
+        sa = mr.server_args
+        page_size = mr.page_size
+        window = mr.sliding_window_size
+        draft_tokens = sa.speculative_num_draft_tokens or 1
+        eviction_interval = max(1, envs.SGLANG_SWA_EVICTION_INTERVAL.get())
+
+        """
+        __________[padding][eviction_interval][window]
+        Padding to make sure eviction point is page-aligned.
+        """
+        trailing_tokens = window + eviction_interval * draft_tokens + page_size
+        if sa.speculative_algorithm is None:
+            decode_alloc = page_size
+        elif sa.disable_overlap_schedule:
+            # spec-v1: new_tokens_required_next_decode per request.
+            decode_alloc = spec_decode_alloc_len_per_request(sa)
+        else:
+            # spec-v2: the overlap allocator keeps 2 * alloc_len outstanding
+            # (eagle_utils.eagle_prepare_for_decode: kv_committed_len + 2 * alloc_len).
+            decode_alloc = 2 * get_alloc_len_per_decode(sa)
+        per_request = trailing_tokens + decode_alloc
+
+        num_reqs = sa.max_running_requests // mr.dp_size
+        if sa.disaggregation_mode == "decode":
+            self._swa_cap = (
+                per_request * num_reqs
+                + (window + page_size) * sa.disaggregation_decode_extra_slots
+            )
+        else:
+            chunks_in_flight = 1 if sa.disable_overlap_schedule else 2
+            self._swa_cap = (
+                per_request * num_reqs
+                + chunks_in_flight * sa.chunked_prefill_size
+                + page_size
+            )
+
+    @staticmethod
+    def is_applicable(mr: ModelRunner) -> bool:
+        """True when SWAChunkCache can be sized from explicit max requests."""
+        sa = mr.server_args
+        if sa.max_running_requests is None:
+            return False
+        if not sa.disable_radix_cache:
+            return False
+        if sa.chunked_prefill_size is None:
+            return False
+        if mr.sliding_window_size is None:
+            return False
+        return len(mr.model_config.full_attention_layer_ids) > 0
+
+    def calculate_pool_sizes(
+        self, available_bytes: int, page_size: int
+    ) -> MemoryPoolConfig:
+        # SWA pool sized tightly from the cap; the rest of the budget goes to full.
+        swa_tokens = ceil_align(self._swa_cap, page_size)
+        fixed_swa_bytes = swa_tokens * self._swa_per_token * self._swa_layers_num
+        full_cell_size = self._full_per_token * (
+            self._full_layers_num + self._draft_full_layers_num
+        )
+        full_tokens = (
+            int((available_bytes - fixed_swa_bytes) // full_cell_size) // page_size
+        ) * page_size
+        if full_tokens <= 0:
+            raise RuntimeError(
+                f"SWA pool cap ({swa_tokens} tokens, "
+                f"{fixed_swa_bytes / (1 << 30):.2f} GiB) leaves no room for the full "
+                f"KV pool within the available {available_bytes / (1 << 30):.2f} GiB. "
+                f"Reduce --max-running-requests, lower SGLANG_SWA_EVICTION_INTERVAL, "
+                f"or increase --mem-fraction-static."
+            )
+        return MemoryPoolConfig(
+            max_total_num_tokens=full_tokens,
+            full_max_total_num_tokens=full_tokens,
+            swa_max_total_num_tokens=swa_tokens,
+        )
+
+    def calculate_pool_sizes_from_max_tokens(
+        self, max_total_num_tokens: int, page_size: int
+    ) -> MemoryPoolConfig:
+        # Constrained max_total goes to the full pool; SWA stays at its cap.
+        swa_tokens = ceil_align(self._swa_cap, page_size)
+        full_tokens = (max_total_num_tokens // page_size) * page_size
+        return MemoryPoolConfig(
+            max_total_num_tokens=full_tokens,
+            full_max_total_num_tokens=full_tokens,
+            swa_max_total_num_tokens=min(swa_tokens, max_total_num_tokens),
+        )
+
+
+@dataclass
+class _DSV4PoolSizes:
+    full_max_total_num_tokens: int
+    swa_max_total_num_tokens: int
+    c4_max_total_num_tokens: int
+    c128_max_total_num_tokens: int
+    c4_state_pool_size: int
+    c128_state_pool_size: int
+
+
+class DSV4PoolConfigurator(MemoryPoolConfigurator):
+    """Configurator for DSV4 compressed-attention models.
+
+    Splits available memory across full / swa / c4 / c128 + c4_state / c128_state
+    pools. coeff is bytes_per_full_token (inflated by (T+D)/T when speculative
+    decode reserves a draft worker, mirroring dflash's cell_size scaling); bias = 0.
+    """
+
+    def __init__(self, mr: ModelRunner):
+        cfg = mr.model_config
+        self.qk_nope_head_dim = cfg.qk_nope_head_dim
+        self.qk_rope_head_dim = cfg.qk_rope_head_dim
+        self.indexer_head_dim = cfg.index_head_dim
+        self.context_len = mr.model_config.context_len
+        # PP-local slice; matches DeepSeekV4TokenToKVPool's stage_ratios.
+        self.compression_ratios = cfg.compress_ratios[mr.start_layer : mr.end_layer]
+        if mr.pp_size > 1:
+            logger.info(
+                f"DSV4 pool PP slice: rank={mr.pp_group.rank_in_group} "
+                f"layers=[{mr.start_layer},{mr.end_layer}) "
+                f"local={len(self.compression_ratios)}/{len(cfg.compress_ratios)}"
+            )
+        self.swa_page_size = cfg.window_size
+        self.swa_ratio = mr.server_args.swa_full_tokens_ratio
+        self.is_speculative = mr.server_args.speculative_algorithm is not None
+        self.online_c128_mtp_max_draft_tokens = (
+            mr.server_args.max_speculative_num_draft_tokens or 0
+        )
+        self.requested_max_running_requests_per_worker = (
+            mr.server_args.max_running_requests // mr.dp_size
+            if mr.server_args.max_running_requests is not None
+            else None
+        )
+        self.disaggregation_mode = mr.server_args.disaggregation_mode
+        self.disaggregation_decode_extra_slots = (
+            mr.server_args.disaggregation_decode_extra_slots or 0
+        )
+        if mr.enable_hisparse:
+            from sglang.srt.mem_cache.sparsity import parse_hisparse_config
+
+            self.c4_shrink_factor = parse_hisparse_config(
+                mr.server_args
+            ).host_to_device_ratio
+        else:
+            self.c4_shrink_factor = 1
+        assert self.c4_shrink_factor >= 1
+        if self.c4_shrink_factor > 1:
+            logger.info(f"HiSparse c4 host-to-device ratio = {self.c4_shrink_factor}")
+
+        self.c4_ring_size = get_compress_state_ring_size(4, self.is_speculative)
+        self.c128_ring_size = get_compress_state_ring_size(128, self.is_speculative)
+
+        self.num_layers_total = len(self.compression_ratios)
+        self.num_layers_ca4 = sum(1 for r in self.compression_ratios if r == 4)
+        self.num_layers_ca128 = sum(1 for r in self.compression_ratios if r == 128)
+
+        self.bytes_per_full_token = self._get_bytes_per_full_token()
+        if self.is_speculative:
+            # Reserve memory for the speculative draft worker by inflating
+            # per-token bytes by (target+draft)/target. Equivalent to dflash's
+            # scale_kv_cell_size_per_token_for_dflash but applied to
+            # bytes_per_full_token: tokens = avail / (bpft * (T+D)/T).
+            draft_layers = 1
+            target_layers = self.num_layers_total
+            self.bytes_per_full_token *= (target_layers + draft_layers) / target_layers
+
+        # Online c128 keeps a single in-progress (max, sum, kv) state per index
+        # and assumes a strict forward-only schedule. Speculative decode (MTP)
+        # would need rollback / replay across draft and verify, which the
+        # online path doesn't support yet.
+        if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
+            allow_experimental_online_c128_mtp = (
+                envs.SGLANG_EXPERIMENTAL_ONLINE_C128_MTP.get()
+                and mr.spec_algorithm.is_eagle()
+            )
+            assert mr.spec_algorithm.is_none() or allow_experimental_online_c128_mtp, (
+                "SGLANG_OPT_USE_ONLINE_COMPRESS does not support speculative decode "
+                "(MTP) yet, except the experimental EAGLE topk=1 path gated by "
+                "SGLANG_EXPERIMENTAL_ONLINE_C128_MTP=1"
+            )
+            if allow_experimental_online_c128_mtp:
+                assert self.online_c128_mtp_max_draft_tokens > 0, (
+                    "SGLANG_EXPERIMENTAL_ONLINE_C128_MTP requires "
+                    "speculative_num_draft_tokens to be set."
+                )
+                logger.warning(
+                    "DSV4 compressed attention: experimental online c128 + MTP enabled "
+                    f"(EAGLE topk=1 only, "
+                    f"draft_banks={self.online_c128_mtp_max_draft_tokens}). "
+                    "Validate correctness carefully."
+                )
+            else:
+                logger.info(
+                    "DSV4 compressed attention: online c128 enabled (ring_size=1)"
+                )
+
+    def _get_bytes_per_full_token(self) -> float:
+        kv_bytes = self.qk_nope_head_dim + self.qk_rope_head_dim * 2 + 8
+
+        quant_block_size = 128
+        indexer_bytes = (
+            self.indexer_head_dim + self.indexer_head_dim // quant_block_size * 4
+        )
+
+        attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        c4_state_dtype_size, c128_state_dtype_size = (
+            _get_dsv4_compress_state_dtype_sizes()
+        )
+        c4_state_bytes = 2 * 2 * attn_head_dim * c4_state_dtype_size
+        # Online c128 stores (max, sum, kv) per slot (3*head_dim) instead of
+        # raw (kv, score) (2*head_dim). Combined with ring_size=1 this still
+        # nets a large reduction (~3/256x) but the per-slot bytes go up.
+        c128_online = envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get()
+        c128_state_bytes = (
+            (3 if c128_online else 2 * 1) * attn_head_dim * c128_state_dtype_size
+        )
+        c4_indexer_state_bytes = 2 * 2 * self.indexer_head_dim * c4_state_dtype_size
+
+        c4_state_ratio = self.c4_ring_size / self.swa_page_size
+        # C128 state is request-scoped and is finalized after
+        # max_running_requests is known, so it should not scale with
+        # full-token capacity here.
+        c128_state_ratio = 0
+
+        c4_frac = 1 / (4 * self.c4_shrink_factor)
+        return (
+            self.swa_ratio * kv_bytes * self.num_layers_total
+            + c4_frac * kv_bytes * self.num_layers_ca4
+            + 1 / 128 * kv_bytes * self.num_layers_ca128
+            + 1 / 4 * indexer_bytes * self.num_layers_ca4
+            + self.swa_ratio * c4_state_ratio * c4_state_bytes * self.num_layers_ca4
+            + c128_state_ratio * c128_state_bytes * self.num_layers_ca128
+            + self.swa_ratio
+            * c4_state_ratio
+            * c4_indexer_state_bytes
+            * self.num_layers_ca4
+        )
+
+    def _compute_dsv4_sizes(self, full_token: int, page_size: int) -> _DSV4PoolSizes:
+        full_token = full_token // page_size * page_size
+        swa_tokens = int(full_token * self.swa_ratio) // page_size * page_size
+        return _DSV4PoolSizes(
+            full_max_total_num_tokens=full_token,
+            swa_max_total_num_tokens=swa_tokens,
+            c4_max_total_num_tokens=full_token // (4 * self.c4_shrink_factor),
+            c128_max_total_num_tokens=full_token // 128,
+            c4_state_pool_size=swa_tokens // self.swa_page_size * self.c4_ring_size,
+            c128_state_pool_size=0,
+        )
+
+    def _get_num_req_slots(self, max_running_requests: int) -> int:
+        if self.disaggregation_mode == "decode":
+            return max_running_requests + self.disaggregation_decode_extra_slots + 1
+        return max_running_requests + 1
+
+    def _get_c128_state_fixed_bytes(self, max_running_requests: int) -> int:
+        if self.num_layers_ca128 == 0:
+            return 0
+
+        _, c128_state_dtype_size = _get_dsv4_compress_state_dtype_sizes()
+        attn_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim
+        num_req_slots = self._get_num_req_slots(max_running_requests)
+
+        if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
+            state_rows = num_req_slots + self.c128_ring_size + 1
+            state_rows *= 1 + self.online_c128_mtp_max_draft_tokens
+            state_last_dim = 3 * attn_head_dim
+        else:
+            state_pool_size = num_req_slots * self.c128_ring_size
+            state_rows = state_pool_size + self.c128_ring_size + 1
+            state_rows = ceil_div(state_rows, 128) * 128
+            state_last_dim = 2 * attn_head_dim
+
+        return (
+            state_rows * state_last_dim * c128_state_dtype_size * self.num_layers_ca128
+        )
+
+    def _get_c128_state_fixed_bytes_for_token_capacity(
+        self, token_capacity: int
+    ) -> int:
+        if self.requested_max_running_requests_per_worker is not None:
+            return self._get_c128_state_fixed_bytes(
+                self.requested_max_running_requests_per_worker
+            )
+
+        estimated = int(token_capacity / self.context_len * 512)
+        estimated = max(min(estimated, 4096), 2048)
+        max_running_requests = min(estimated, token_capacity // 2)
+        return self._get_c128_state_fixed_bytes(max_running_requests)
+
+    def _to_config(self, sizes: _DSV4PoolSizes) -> MemoryPoolConfig:
+        full = sizes.full_max_total_num_tokens
+        swa = sizes.swa_max_total_num_tokens
+        logger.info(
+            f"DSV4 pool sizes: full={full}, swa={swa}, "
+            f"c4={sizes.c4_max_total_num_tokens}, "
+            f"c128={sizes.c128_max_total_num_tokens}, "
+            f"c4_state={sizes.c4_state_pool_size}, "
+            f"c128_state={sizes.c128_state_pool_size}"
+        )
+        return MemoryPoolConfig(
+            max_total_num_tokens=full,
+            full_max_total_num_tokens=full,
+            swa_max_total_num_tokens=swa,
+            c4_max_total_num_tokens=sizes.c4_max_total_num_tokens,
+            c128_max_total_num_tokens=sizes.c128_max_total_num_tokens,
+            c4_state_pool_size=sizes.c4_state_pool_size,
+            c128_state_pool_size=sizes.c128_state_pool_size,
+        )
+
+    def finalize_with_max_running_requests(
+        self, config: MemoryPoolConfig
+    ) -> MemoryPoolConfig:
+        assert config.max_running_requests is not None
+        num_req_slots = self._get_num_req_slots(config.max_running_requests)
+        if envs.SGLANG_OPT_USE_ONLINE_COMPRESS.get():
+            config.c128_state_pool_size = num_req_slots
+        else:
+            config.c128_state_pool_size = num_req_slots * self.c128_ring_size
+        return config
+
+    def calculate_pool_sizes(
+        self, available_bytes: int, page_size: int
+    ) -> MemoryPoolConfig:
+        assert (
+            page_size % 128 == 0
+        ), "page_size must be multiple of 128 for compressed attention"
+
+        if self.requested_max_running_requests_per_worker is not None:
+            c128_state_fixed_bytes = self._get_c128_state_fixed_bytes(
+                self.requested_max_running_requests_per_worker
+            )
+        else:
+            full_token = int(available_bytes / self.bytes_per_full_token)
+            c128_state_fixed_bytes = (
+                self._get_c128_state_fixed_bytes_for_token_capacity(full_token)
+            )
+
+        available_bytes_for_tokens = max(available_bytes - c128_state_fixed_bytes, 0)
+        full_token = int(available_bytes_for_tokens / self.bytes_per_full_token)
+
+        sizes = self._compute_dsv4_sizes(full_token, page_size)
+        logger.info(
+            f"DSV4 memory calculation: "
+            f"bytes_per_full_token={self.bytes_per_full_token:.2f}, "
+            f"available_bytes={available_bytes / (1 << 30):.2f} GB, "
+            f"c128_state_fixed={c128_state_fixed_bytes / (1 << 30):.2f} GB, "
+            f"full_token={sizes.full_max_total_num_tokens}"
+        )
+        return self._to_config(sizes)
+
+    def calculate_pool_sizes_from_max_tokens(
+        self, max_total_num_tokens: int, page_size: int
+    ) -> MemoryPoolConfig:
+        assert (
+            page_size % 128 == 0
+        ), "page_size must be multiple of 128 for compressed attention"
+        sizes = self._compute_dsv4_sizes(max_total_num_tokens, page_size)
+        return self._to_config(sizes)
+
+
+def create_memory_pool_configurator(
+    mr: ModelRunner,
+) -> MemoryPoolConfigurator:
+    """Factory: select the right configurator for the model architecture."""
+    if is_deepseek_v4(mr.model_config.hf_config) and mr.is_hybrid_swa:
+        return DSV4PoolConfigurator(mr)
+    if mr.is_hybrid_swa:
+        if SWAChunkCapPoolConfigurator.is_applicable(mr):
+            return SWAChunkCapPoolConfigurator(mr)
+        return HybridSWAPoolConfigurator(mr)
+    # Future: MambaPoolConfigurator
+    return DefaultPoolConfigurator(mr)

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

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

@@ -0,0 +1,31 @@
+"""Capture-mechanism backends for CUDA graphs.
+
+A backend owns *how* a captured artifact is produced and replayed for
+one shape; it is phase-agnostic. Runners (cuda_graph_runner/) own
+*what* data flows in and out.
+
+Public API:
+  - BaseCudaGraphBackend — abstract interface.
+  - FullCudaGraphBackend — single torch.cuda.CUDAGraph per shape.
+  - BreakableCudaGraphBackend — segmented capture with eager break
+    markers; no torch.compile.
+  - TcPiecewiseCudaGraphBackend — torch.compile-driven piecewise
+    capture; FX-splits the model at attention layers.
+"""
+
+from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (  # noqa: F401
+    BaseCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import (  # noqa: F401
+    BreakableCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import (  # noqa: F401
+    FullCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.tc_piecewise_cuda_graph_backend import (  # noqa: F401
+    TcPiecewiseCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.utils import (  # noqa: F401
+    resolve_decode_backend,
+    resolve_prefill_backend,
+)

python/sglang/srt/model_executor/runner_backend/base_cuda_graph_backend.py

@@ -0,0 +1,81 @@
+# 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.
+# ==============================================================================
+"""Backend interface for CUDA graph capture/replay."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Callable, Iterator, Optional
+
+import torch
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.runner.shape_key import ShapeKey
+
+
+class BaseCudaGraphBackend(ABC):
+    """Pure ABC: no state, no defaults. Each implementation owns its
+    per-backend state and binds the handles it needs from the
+    cuda_graph_runner passed to its __init__.
+
+    Methods:
+      - capture_session(stream) — context wrapping the runner's outer
+        capture loop; backends bind stream / pool and open per-backend
+        capture flags here.
+      - capture_one(shape_key, forward_fn, dummies, post_warmup_hook)
+        — record the replayable artifact for shape_key; one call per
+        shape inside capture_session.
+      - can_run(forward_batch, shape_key) — can this backend replay
+        for the given batch at the given shape.
+      - replay_session() — context wrapping replay-time model code;
+        backends open the "we are replaying" flag here when they have
+        one.
+      - replay(shape_key, static_forward_batch, **kwargs) — invoke
+        the captured artifact.
+      - cleanup() — release pool and drop captured artifacts.
+
+    Notes:
+      - The outer capture loop is runner-specific; it lives on the
+        runner, not here.
+    """
+
+    @abstractmethod
+    def capture_session(self, stream: torch.cuda.Stream) -> Iterator[None]: ...
+
+    @abstractmethod
+    def capture_one(
+        self,
+        shape_key: ShapeKey,
+        forward_fn,
+        dummies: Optional[Any] = None,
+        post_warmup_hook: Optional[Callable[[], None]] = None,
+    ) -> None: ...
+
+    @abstractmethod
+    def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool: ...
+
+    @abstractmethod
+    def replay_session(self) -> Iterator[None]: ...
+
+    @abstractmethod
+    def replay(
+        self,
+        shape_key: ShapeKey,
+        static_forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any: ...
+
+    @abstractmethod
+    def cleanup(self) -> None: ...

python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py

@@ -0,0 +1,252 @@
+# 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.
+# ==============================================================================
+"""BreakableCudaGraphBackend — segment-captured graphs with eager break
+markers (eager_on_graph decorators on attention / mamba layers).
+No torch.compile.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+
+import torch
+
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    set_graph_pool_id,
+)
+from sglang.srt.model_executor.forward_batch_info import PPProxyTensors
+from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
+    BaseCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.cuda_graph_dedup_mixin import (
+    DedupedCudaGraphMixin,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
+    BreakableCUDAGraph,
+    BreakableCUDAGraphCapture,
+    eager_on_graph,
+    enable_breakable_cuda_graph,
+)
+from sglang.srt.model_executor.runner_utils.pool import (
+    get_or_create_global_graph_memory_pool,
+)
+from sglang.srt.utils import get_bool_env_var
+from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
+        BaseCudaGraphRunner,
+    )
+    from sglang.srt.model_executor.runner.shape_key import ShapeKey
+
+
+class BreakableCudaGraphBackend(DedupedCudaGraphMixin, BaseCudaGraphBackend):
+    """Segmented capture: graphs break at attention / mamba boundaries;
+    attention metadata is recomputed at replay outside captured segments.
+    """
+
+    def __init__(
+        self,
+        cuda_graph_runner: BaseCudaGraphRunner,
+        *,
+        enable_memory_saver: bool = False,
+        debug_eager: bool = False,
+    ) -> None:
+        self._model_runner = cuda_graph_runner.model_runner
+        self._graphs: Dict[Any, BreakableCUDAGraph] = {}
+        self._outputs: Dict[Any, Any] = {}
+        self._pool = None
+        self._device_module = cuda_graph_runner.device_module
+        self._tp_group = cuda_graph_runner.model_runner.tp_group
+        self._capture_stream: Optional[torch.cuda.Stream] = None
+        self._debug_eager = debug_eager
+        self._shared_output_buffer: Optional[Any] = None
+        self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
+            enable=enable_memory_saver
+            and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
+        )
+        if (
+            self._memory_saver_adapter is not None
+            and self._memory_saver_adapter.enabled
+        ):
+            raise NotImplementedError(
+                "Breakable CUDA graph is not compatible with memory saver mode"
+            )
+
+    @contextmanager
+    def capture_session(self, stream: torch.cuda.Stream):
+        if self._pool is None:
+            self._pool = get_or_create_global_graph_memory_pool(self._device_module)
+        set_graph_pool_id(self._pool)
+        self._capture_stream = stream
+        self._shared_output_buffer = None
+        self.begin_cuda_graph_capture()
+        try:
+            with self.replay_session():
+                yield
+        finally:
+            try:
+                self.end_cuda_graph_capture()
+            finally:
+                self._capture_stream = None
+
+    def capture_one(
+        self,
+        shape_key: ShapeKey,
+        forward_fn: Callable[[], Any],
+        dummies: Optional[Any] = None,
+        post_warmup_hook: Optional[Callable[[], None]] = None,
+    ) -> None:
+        warmup_out = None
+        for _ in range(2):
+            self._device_module.synchronize()
+            self._tp_group.barrier()
+            warmup_out = forward_fn()
+            if post_warmup_hook is not None:
+                post_warmup_hook()
+
+        graph = BreakableCUDAGraph()
+        captured_fn = (
+            eager_on_graph(True)(forward_fn) if self._debug_eager else forward_fn
+        )
+        size = shape_key.size
+        if self._shared_output_buffer is None:
+            self._shared_output_buffer = self._alloc_full_buffer(warmup_out, size)
+        with BreakableCUDAGraphCapture(
+            cuda_graph=graph,
+            pool=self._pool,
+            stream=self._capture_stream,
+        ):
+            out = captured_fn()
+            out_rows = self._output_rows(out, size)
+            self._copy_output_to_buffer(out, self._shared_output_buffer, out_rows)
+
+        stored = self._slice_output(self._shared_output_buffer, out_rows)
+        self._graphs[shape_key] = graph
+        self._outputs[shape_key] = stored
+
+    def _output_rows(self, output: Any, cap: int) -> int:
+        """Leading-dim row count actually produced by the body, clamped to ``cap``.
+
+        A body that shards or prunes its output along dim 0 returns fewer than
+        ``cap`` rows; everything else returns exactly ``cap``.
+        """
+        if torch.is_tensor(output):
+            return min(cap, output.shape[0])
+        if isinstance(output, PPProxyTensors):
+            rows = [t.shape[0] for t in output.tensors.values()]
+            return min([cap, *rows])
+        if isinstance(output, (list, tuple)) and output:
+            return min(self._output_rows(o, cap) for o in output if o is not None)
+        return cap
+
+    def _alloc_full_buffer(self, output: Any, size: int) -> Any:
+        """A same-structure buffer as ``output`` but with ``size`` leading rows."""
+        if output is None:
+            return None
+        if torch.is_tensor(output):
+            return output.new_empty((size, *output.shape[1:]))
+        if isinstance(output, PPProxyTensors):
+            return PPProxyTensors(
+                {
+                    key: t.new_empty((size, *t.shape[1:]))
+                    for key, t in output.tensors.items()
+                }
+            )
+        if isinstance(output, tuple):
+            return tuple(self._alloc_full_buffer(o, size) for o in output)
+        if isinstance(output, list):
+            return [self._alloc_full_buffer(o, size) for o in output]
+        raise TypeError(f"Unsupported BCG output type: {type(output)}")
+
+    def _slice_output(self, output: Any, num_tokens: int) -> Any:
+        if output is None:
+            return None
+        if torch.is_tensor(output):
+            return output[:num_tokens]
+        if isinstance(output, PPProxyTensors):
+            return output[:num_tokens]
+        if isinstance(output, tuple):
+            return tuple(self._slice_output(item, num_tokens) for item in output)
+        if isinstance(output, list):
+            return [self._slice_output(item, num_tokens) for item in output]
+        raise TypeError(f"Unsupported BCG output type: {type(output)}")
+
+    def _copy_output_to_buffer(
+        self, output: Any, output_buffer: Any, num_tokens: int
+    ) -> None:
+        if output is None or output_buffer is None:
+            if output is None and output_buffer is None:
+                return
+            raise ValueError(
+                "BCG output structure changed between capture sizes: "
+                f"{type(output)} vs {type(output_buffer)}"
+            )
+        if torch.is_tensor(output) and torch.is_tensor(output_buffer):
+            output_buffer[:num_tokens].copy_(output[:num_tokens])
+            return
+        if isinstance(output, PPProxyTensors) and isinstance(
+            output_buffer, PPProxyTensors
+        ):
+            if output.tensors.keys() != output_buffer.tensors.keys():
+                raise ValueError(
+                    "BCG output proxy structure changed between capture sizes: "
+                    f"{output.tensors.keys()} != {output_buffer.tensors.keys()}"
+                )
+            for key, tensor in output.tensors.items():
+                self._copy_output_to_buffer(
+                    tensor, output_buffer.tensors[key], num_tokens
+                )
+            return
+        if isinstance(output, (list, tuple)) and isinstance(
+            output_buffer, type(output)
+        ):
+            if len(output) != len(output_buffer):
+                raise ValueError(
+                    "BCG output sequence structure changed between capture sizes: "
+                    f"{len(output)} != {len(output_buffer)}"
+                )
+            for item, buffer in zip(output, output_buffer):
+                self._copy_output_to_buffer(item, buffer, num_tokens)
+            return
+        raise TypeError(
+            "Unsupported BCG output buffer pair: "
+            f"{type(output)} vs {type(output_buffer)}"
+        )
+
+    def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool:
+        return shape_key in self._graphs
+
+    @contextmanager
+    def replay_session(self):
+        with enable_breakable_cuda_graph():
+            yield
+
+    def replay(
+        self,
+        shape_key: ShapeKey,
+        static_forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any:
+        self._graphs[shape_key].replay()
+        return self._outputs[shape_key]
+
+    def cleanup(self) -> None:
+        self.close()
+        self._graphs.clear()
+        self._outputs.clear()
+        self._pool = None
+        self._shared_output_buffer = None

python/sglang/srt/model_executor/runner_backend/cuda_graph_dedup_mixin.py

@@ -0,0 +1,375 @@
+"""Shared CUDA graph executable-dedup plumbing for CUDA graph backends."""
+
+from __future__ import annotations
+
+import heapq
+import logging
+from dataclasses import dataclass, field
+
+import torch
+
+try:
+    from cuda.bindings import driver as cuda_drv
+    from cuda.bindings import runtime as cuda_rt
+except ImportError:
+    cuda_drv = None
+    cuda_rt = None
+
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.cuda_utils import (
+    checkCudaErrors,
+)
+from sglang.srt.utils import get_bool_env_var
+
+logger = logging.getLogger(__name__)
+
+
+def dedup_update(graph_exec: int, raw_graph: int) -> tuple[bool, str]:
+    assert cuda_rt is not None
+    err, info = cuda_rt.cudaGraphExecUpdate(graph_exec, raw_graph)
+    if info is None:
+        return False, f"err={int(err)}"
+    result = info.result
+    ok = (
+        err == cuda_rt.cudaError_t.cudaSuccess
+        and result == cuda_rt.cudaGraphExecUpdateResult.cudaGraphExecUpdateSuccess
+    )
+    return ok, "" if ok else f"err={int(err)} result={result}"
+
+
+def maybe_cuda_result(result):
+    return None if int(result[0]) != 0 else checkCudaErrors(result)
+
+
+def kernel_name(params) -> str:
+    assert cuda_drv is not None
+    for handle, getter in (
+        (getattr(params, "kern", None), cuda_drv.cuKernelGetName),
+        (getattr(params, "func", None), cuda_drv.cuFuncGetName),
+    ):
+        if handle is None or int(handle) == 0:
+            continue
+        name = maybe_cuda_result(getter(handle))
+        if name is not None:
+            return name.decode("utf-8", "replace")
+    return f"func:{int(getattr(params, 'func', 0))}"
+
+
+def kernel_attrs(node) -> tuple[tuple[str, object], ...]:
+    assert cuda_drv is not None
+    attrs = []
+    for name, attr_name, get_value in (
+        (
+            "cooperative",
+            "CU_LAUNCH_ATTRIBUTE_COOPERATIVE",
+            lambda v: int(v.cooperative),
+        ),
+        (
+            "clusterDim",
+            "CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION",
+            lambda v: (
+                int(v.clusterDim.x),
+                int(v.clusterDim.y),
+                int(v.clusterDim.z),
+            ),
+        ),
+        (
+            "clusterSchedulingPolicyPreference",
+            "CU_LAUNCH_ATTRIBUTE_CLUSTER_SCHEDULING_POLICY_PREFERENCE",
+            lambda v: int(v.clusterSchedulingPolicyPreference),
+        ),
+        (
+            "preferredClusterDim",
+            "CU_LAUNCH_ATTRIBUTE_PREFERRED_CLUSTER_DIMENSION",
+            lambda v: (
+                int(v.preferredClusterDim.x),
+                int(v.preferredClusterDim.y),
+                int(v.preferredClusterDim.z),
+            ),
+        ),
+        (
+            "sharedMemCarveout",
+            "CU_LAUNCH_ATTRIBUTE_PREFERRED_SHARED_MEMORY_CARVEOUT",
+            lambda v: int(v.sharedMemCarveout),
+        ),
+    ):
+        attr = getattr(cuda_drv.CUkernelNodeAttrID, attr_name, None)
+        if attr is None:
+            continue
+        value = maybe_cuda_result(cuda_drv.cuGraphKernelNodeGetAttribute(node, attr))
+        if value is not None:
+            attrs.append((name, get_value(value)))
+    return tuple(attrs)
+
+
+def kernel_node_payload(node):
+    assert cuda_drv is not None
+    params = checkCudaErrors(cuda_drv.cuGraphKernelNodeGetParams(node))
+    return (
+        kernel_name(params),
+        (int(params.gridDimX), int(params.gridDimY), int(params.gridDimZ)),
+        (int(params.blockDimX), int(params.blockDimY), int(params.blockDimZ)),
+        int(params.sharedMemBytes),
+        kernel_attrs(node),
+    )
+
+
+def graph_node_payload(node):
+    assert cuda_drv is not None
+    node_type = checkCudaErrors(cuda_drv.cuGraphNodeGetType(node))
+    match node_type:
+        case cuda_drv.CUgraphNodeType.CU_GRAPH_NODE_TYPE_KERNEL:
+            payload = kernel_node_payload(node)
+        case cuda_drv.CUgraphNodeType.CU_GRAPH_NODE_TYPE_MEMCPY:
+            params = checkCudaErrors(cuda_drv.cuGraphMemcpyNodeGetParams(node))
+            payload = (int(params.srcMemoryType), int(params.dstMemoryType))
+        case cuda_drv.CUgraphNodeType.CU_GRAPH_NODE_TYPE_MEMSET:
+            params = checkCudaErrors(cuda_drv.cuGraphMemsetNodeGetParams(node))
+            payload = (int(params.elementSize),)
+        case cuda_drv.CUgraphNodeType.CU_GRAPH_NODE_TYPE_GRAPH:
+            child_graph = checkCudaErrors(cuda_drv.cuGraphChildGraphNodeGetGraph(node))
+            payload = graph_signature(child_graph)
+        case cuda_drv.CUgraphNodeType.CU_GRAPH_NODE_TYPE_EMPTY:
+            payload = ()
+        case _:
+            payload = ()
+    return (node_type.name, payload)
+
+
+def graph_signature(raw_graph: int):
+    assert cuda_drv is not None
+    _, num_nodes = checkCudaErrors(cuda_drv.cuGraphGetNodes(raw_graph, 0))
+    nodes, _ = checkCudaErrors(cuda_drv.cuGraphGetNodes(raw_graph, num_nodes))
+    node_indices = {int(node): i for i, node in enumerate(nodes)}
+
+    _, _, _, num_edges = checkCudaErrors(cuda_drv.cuGraphGetEdges(raw_graph, 0))
+    from_nodes, to_nodes, _, _ = checkCudaErrors(
+        cuda_drv.cuGraphGetEdges(raw_graph, num_edges)
+    )
+    edges = [
+        (node_indices[int(src)], node_indices[int(dst)])
+        for src, dst in zip(from_nodes, to_nodes)
+    ]
+
+    children = [[] for _ in nodes]
+    indegree = [0] * len(nodes)
+    for src, dst in edges:
+        children[src].append(dst)
+        indegree[dst] += 1
+
+    ready = [i for i, degree in enumerate(indegree) if degree == 0]
+    heapq.heapify(ready)
+    order = []
+    while ready:
+        node_idx = heapq.heappop(ready)
+        order.append(node_idx)
+        for child_idx in sorted(children[node_idx]):
+            indegree[child_idx] -= 1
+            if indegree[child_idx] == 0:
+                heapq.heappush(ready, child_idx)
+    assert len(order) == len(nodes), "CUDA graph contains a dependency cycle"
+
+    topo_indices = {node_idx: i for i, node_idx in enumerate(order)}
+    topo_edges = tuple(
+        sorted((topo_indices[src], topo_indices[dst]) for src, dst in edges)
+    )
+    return (
+        tuple(graph_node_payload(nodes[node_idx]) for node_idx in order),
+        topo_edges,
+    )
+
+
+@dataclass(slots=True)
+class GraphExecGroup:
+    graph_exec: int
+    current_raw_graph: int
+    compat_exec: int | None
+    graphs: list[DedupedCudaGraph] = field(default_factory=list)
+
+
+@dataclass(eq=False, slots=True)
+class DedupedCudaGraph:
+    raw_graph: int
+    original_graph: object | None
+    registry: DedupedCudaGraphRegistry
+    group: GraphExecGroup | None = None
+
+    def replay(self, stream: int | None = None) -> None:
+        if stream is None:
+            stream = torch.cuda.current_stream().cuda_stream
+        self.registry.replay(self, stream)
+
+
+class DedupedCudaGraphRegistry:
+    def __init__(self):
+        self.groups: dict[tuple, GraphExecGroup] = {}
+        self.sealed = False
+
+    def instantiate(self, raw_graph: int) -> int:
+        assert cuda_rt is not None
+        graph_exec = checkCudaErrors(
+            cuda_rt.cudaGraphInstantiateWithFlags(raw_graph, 0)
+        )
+        return graph_exec
+
+    def destroy_exec(self, graph_exec: int) -> None:
+        assert cuda_rt is not None
+        checkCudaErrors(cuda_rt.cudaGraphExecDestroy(graph_exec))
+
+    def register(self, captured_graph) -> DedupedCudaGraph:
+        assert not self.sealed
+        raw_graph = captured_graph.raw_cuda_graph()
+        signature = graph_signature(raw_graph)
+        graph = DedupedCudaGraph(raw_graph, captured_graph, self)
+
+        group = self.groups.get(signature)
+        if group is not None:
+            assert group.compat_exec is not None
+            ok, detail = dedup_update(group.compat_exec, graph.raw_graph)
+            assert ok, f"CUDA graph dedup register update failed ({detail})"
+            graph.group = group
+            group.graphs.append(graph)
+            return graph
+
+        group = GraphExecGroup(
+            graph_exec=self.instantiate(graph.raw_graph),
+            current_raw_graph=graph.raw_graph,
+            compat_exec=self.instantiate(graph.raw_graph),
+            graphs=[graph],
+        )
+        graph.group = group
+        self.groups[signature] = group
+        return graph
+
+    def seal(self) -> None:
+        if self.sealed:
+            return
+        self.sealed = True
+        for group in self.groups.values():
+            if group.compat_exec is not None:
+                self.destroy_exec(group.compat_exec)
+                group.compat_exec = None
+
+    def stats(self) -> tuple[int, int]:
+        return sum(len(group.graphs) for group in self.groups.values()), len(
+            self.groups
+        )
+
+    def replay(self, graph: DedupedCudaGraph, stream: int) -> None:
+        assert cuda_rt is not None
+        group = graph.group
+        assert (
+            group is not None
+        ), "captured CUDA graph does not belong to this dedup state"
+
+        raw_graph = graph.raw_graph
+        graph_exec = group.graph_exec
+        if group.current_raw_graph != raw_graph:
+            ok, detail = dedup_update(graph_exec, raw_graph)
+            assert ok, (
+                "CUDA graph dedup replay update failed "
+                f"({detail}); captured graph is not compatible with its dedup group"
+            )
+            group.current_raw_graph = raw_graph
+
+        checkCudaErrors(cuda_rt.cudaGraphLaunch(graph_exec, stream))
+
+    def close(self) -> None:
+        self.sealed = True
+
+        for group in self.groups.values():
+            if group.compat_exec is not None:
+                self.destroy_exec(group.compat_exec)
+                group.compat_exec = None
+            self.destroy_exec(group.graph_exec)
+            for graph in group.graphs:
+                if graph.original_graph is not None:
+                    graph.original_graph.reset()
+                graph.original_graph = None
+                graph.group = None
+            group.graphs.clear()
+
+        self.groups.clear()
+
+
+class DedupedCudaGraphMixin:
+    deduped_cuda_graph: DedupedCudaGraphRegistry | None = None
+
+    def _dedup_registries(self) -> list[DedupedCudaGraphRegistry]:
+        registries = getattr(self, "_deduped_cuda_graph_registries", None)
+        if registries is None:
+            registries = []
+            self._deduped_cuda_graph_registries = registries
+        return registries
+
+    def _memory_saver_cuda_graph_enabled(self) -> bool:
+        adapter = getattr(self, "_memory_saver_adapter", None)
+        if adapter is not None and getattr(adapter, "enabled", False):
+            return True
+
+        model_runner = getattr(self, "model_runner", None)
+        if model_runner is None:
+            model_runner = getattr(self, "_model_runner", None)
+        server_args = getattr(model_runner, "server_args", None)
+        return bool(
+            server_args is not None
+            and getattr(server_args, "enable_memory_saver", False)
+            and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
+        )
+
+    def build_deduped_cuda_graph(self):
+        if not envs.SGLANG_ENABLE_CUDA_GRAPH_DEDUP.get():
+            return None
+        if cuda_drv is None or cuda_rt is None:
+            return None
+        try:
+            graph = torch.cuda.CUDAGraph(keep_graph=True)
+            if not hasattr(graph, "raw_cuda_graph"):
+                return None
+            return DedupedCudaGraphRegistry()
+        except TypeError:
+            return None
+        except Exception as e:
+            logger.warning(
+                "[CudaGraph][dedup] %s init failed (%s); using plain executables.",
+                type(self).__name__,
+                e,
+            )
+            return None
+
+    def begin_cuda_graph_capture(self) -> None:
+        if self.deduped_cuda_graph is not None:
+            self.end_cuda_graph_capture()
+
+        if self._memory_saver_cuda_graph_enabled():
+            self.deduped_cuda_graph = None
+            return
+
+        self.deduped_cuda_graph = self.build_deduped_cuda_graph()
+        if self.deduped_cuda_graph is not None:
+            self._dedup_registries().append(self.deduped_cuda_graph)
+
+    def end_cuda_graph_capture(self) -> None:
+        dedup = self.deduped_cuda_graph
+        self.deduped_cuda_graph = None
+        if dedup is not None:
+            captured, execs = dedup.stats()
+            dedup.seal()
+            logger.info("captured %d CUDA graphs, deduped to %d execs", captured, execs)
+
+    def close(self) -> None:
+        registries = self._dedup_registries()
+        seen: set[int] = set()
+        for registry in [self.deduped_cuda_graph, *registries]:
+            if registry is None or id(registry) in seen:
+                continue
+            seen.add(id(registry))
+            registry.close()
+        registries.clear()
+        self.deduped_cuda_graph = None
+
+    def __del__(self):
+        try:
+            self.close()
+        except Exception:
+            pass

python/sglang/srt/model_executor/runner_backend/full_cuda_graph_backend.py

@@ -0,0 +1,135 @@
+# 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.
+# ==============================================================================
+"""FullCudaGraphBackend — captures the entire model forward as one
+torch.cuda.CUDAGraph per shape.
+"""
+
+from __future__ import annotations
+
+from contextlib import AbstractContextManager, contextmanager
+from functools import partial
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+
+import torch
+
+from sglang.srt.constants import GPU_MEMORY_TYPE_CUDA_GRAPH
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    set_graph_pool_id,
+)
+from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
+    BaseCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_utils.pool import (
+    get_or_create_global_graph_memory_pool,
+)
+from sglang.srt.utils import get_bool_env_var
+from sglang.srt.utils.torch_memory_saver_adapter import TorchMemorySaverAdapter
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
+        BaseCudaGraphRunner,
+    )
+    from sglang.srt.model_executor.runner.shape_key import ShapeKey
+
+
+class FullCudaGraphBackend(BaseCudaGraphBackend):
+    """One torch.cuda.CUDAGraph per shape; attention metadata is
+    captured inside the graph. Memory-saver-aware.
+    """
+
+    def __init__(
+        self,
+        cuda_graph_runner: BaseCudaGraphRunner,
+        *,
+        enable_memory_saver: bool = False,
+    ) -> None:
+        self._graphs: Dict[Any, torch.cuda.CUDAGraph] = {}
+        self._outputs: Dict[Any, Any] = {}
+        self._pool = None
+        self._device_module = cuda_graph_runner.device_module
+        self._tp_group = cuda_graph_runner.model_runner.tp_group
+        self._capture_stream: Optional[torch.cuda.Stream] = None
+        self._memory_saver_adapter: Optional[Any] = TorchMemorySaverAdapter.create(
+            enable=enable_memory_saver
+            and get_bool_env_var("SGLANG_MEMORY_SAVER_CUDA_GRAPH")
+        )
+
+    @contextmanager
+    def capture_session(self, stream: torch.cuda.Stream):
+        if self._pool is None:
+            self._pool = get_or_create_global_graph_memory_pool(self._device_module)
+        set_graph_pool_id(self._pool)
+        self._capture_stream = stream
+        try:
+            yield
+        finally:
+            self._capture_stream = None
+
+    def capture_one(
+        self,
+        shape_key: ShapeKey,
+        forward_fn: Callable[[], Any],
+        dummies: Optional[Any] = None,
+        post_warmup_hook: Optional[Callable[[], None]] = None,
+    ) -> None:
+        # Two warmups so kernels are loaded and one-time setup is paid before capture.
+        # post_warmup_hook lets the attention backend reset state that warmup mutated.
+        for _ in range(2):
+            self._device_module.synchronize()
+            self._tp_group.barrier()
+            forward_fn()
+            if post_warmup_hook is not None:
+                post_warmup_hook()
+
+        graph = torch.cuda.CUDAGraph()
+
+        graph_ctx: Callable[..., AbstractContextManager]
+        if (
+            self._memory_saver_adapter is not None
+            and self._memory_saver_adapter.enabled
+        ):
+            graph_ctx = partial(
+                self._memory_saver_adapter.cuda_graph,
+                tag=GPU_MEMORY_TYPE_CUDA_GRAPH,
+            )
+        else:
+            graph_ctx = self._device_module.graph
+
+        with graph_ctx(cuda_graph=graph, pool=self._pool, stream=self._capture_stream):
+            out = forward_fn()
+
+        self._graphs[shape_key] = graph
+        self._outputs[shape_key] = out
+
+    def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool:
+        return shape_key in self._graphs
+
+    @contextmanager
+    def replay_session(self):
+        yield
+
+    def replay(
+        self,
+        shape_key: ShapeKey,
+        static_forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any:
+        self._graphs[shape_key].replay()
+        return self._outputs[shape_key]
+
+    def cleanup(self) -> None:
+        self._graphs.clear()
+        self._outputs.clear()
+        self._pool = None

python/sglang/srt/model_executor/runner_backend/tc_piecewise_cuda_graph_backend.py

@@ -0,0 +1,259 @@
+# 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.
+# ==============================================================================
+"""TcPiecewiseCudaGraphBackend — torch.compile-driven piecewise CUDA graph.
+
+FX-splits the model forward at attention layers; per-shape compiled
+callables internally capture sub-graphs via
+compilation/cuda_piecewise_backend. torch.compile owns the per-shape
+cache so this backend has no _graphs table — only a single
+_compiled_fn reused for every shape.
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any, Callable, Optional
+
+import torch
+import tqdm
+
+from sglang.srt.compilation.compilation_config import CompilationConfig
+from sglang.srt.compilation.compile import install_torch_compiled
+from sglang.srt.compilation.compile_phase import (
+    enable_torch_compile_warmup,
+    set_pcg_capture_stream,
+)
+from sglang.srt.distributed.device_communicators.pynccl_allocator import (
+    set_graph_pool_id,
+)
+from sglang.srt.layers.moe.utils import get_moe_a2a_backend
+from sglang.srt.layers.utils import MultiPlatformOp
+from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
+    BaseCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph import (
+    enable_tc_piecewise_cuda_graph,
+)
+from sglang.srt.model_executor.runner_utils.pool import (
+    get_or_create_global_graph_memory_pool,
+)
+from sglang.srt.runtime_context import get_parallel
+from sglang.srt.utils import is_hip
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
+        BaseCudaGraphRunner,
+    )
+    from sglang.srt.model_executor.runner.shape_key import ShapeKey
+    from sglang.srt.server_args import ServerArgs
+
+
+_VALID_COMPILERS = ("eager", "inductor")
+
+
+def _toggle_multi_platform_ops(
+    model: torch.nn.Module, *, reverse: bool, num_tokens: int
+) -> None:
+    """Recursively flip MultiPlatformOp submodules into / out of
+    torch.compile mode."""
+    for sub in model._modules.values():
+        if isinstance(sub, MultiPlatformOp):
+            if reverse:
+                sub.leave_torch_compile()
+            else:
+                sub.enter_torch_compile(num_tokens=num_tokens)
+        if isinstance(sub, torch.nn.Module):
+            _toggle_multi_platform_ops(sub, reverse=reverse, num_tokens=num_tokens)
+
+
+class TcPiecewiseCudaGraphBackend(BaseCudaGraphBackend):
+    """torch.compile-driven piecewise capture; attention metadata
+    recomputed at replay outside the compiled callable's sub-graphs.
+    """
+
+    def __init__(self, cuda_graph_runner: BaseCudaGraphRunner) -> None:
+        model_runner = cuda_graph_runner.model_runner
+        self._pool = None
+        self._device_module = cuda_graph_runner.device_module
+        self._tp_group = model_runner.tp_group
+        self._capture_stream: Optional[torch.cuda.Stream] = None
+        self._compile_config: CompilationConfig = self.build_compilation_config(
+            model_runner.server_args
+        )
+        self._language_model: torch.nn.Module = getattr(
+            model_runner.model, "language_model", model_runner.model
+        )
+        self._run_compile_pass(cuda_graph_runner)
+        # model_runner.model.forward is the wrapper that builds LogitsProcessorOutput.
+        # The compiled trampoline is dispatched internally by it.
+        self._compiled_fn: Callable = model_runner.model.forward
+
+    @staticmethod
+    def build_compilation_config(server_args: ServerArgs) -> CompilationConfig:
+        """Construct a CompilationConfig from ServerArgs and
+        register the MoE A2A split-op when DeepEP / Mooncake is in use."""
+        prefill = server_args.cuda_graph_config.prefill
+        num_tokens = prefill.bs
+        compiler = prefill.tc_compiler
+        assert num_tokens is not None, "cuda_graph_config[prefill].bs is not set"
+        assert compiler in _VALID_COMPILERS, (
+            f"By now, only {_VALID_COMPILERS} are supported for the "
+            "tc_piecewise prefill compiler."
+        )
+
+        config = CompilationConfig(
+            num_tokens,
+            compiler,
+            server_args.enable_torch_compile_debug_mode,
+        )
+
+        if get_moe_a2a_backend().is_deepep() or get_moe_a2a_backend().is_mooncake():
+            config.add_split_op("sglang.moe_forward_piecewise_cuda_graph_impl")
+
+        return config
+
+    @staticmethod
+    def install_compile(
+        language_model: Any,
+        *,
+        compile_config: CompilationConfig,
+        graph_pool: Any,
+        fullgraph: bool = True,
+        dynamic_arg_dims: Optional[Any] = None,
+    ) -> None:
+        """Wrap language_model.model.forward with torch.compile."""
+        install_torch_compiled(
+            language_model,
+            fullgraph=fullgraph,
+            dynamic_arg_dims=dynamic_arg_dims,
+            compile_config=compile_config,
+            graph_pool=graph_pool,
+        )
+
+    def _run_compile_pass(self, cuda_graph_runner: BaseCudaGraphRunner) -> None:
+        """JIT-activate kernels at the smallest shape, install
+        torch.compile, then run one forward per shape inside
+        enable_torch_compile_warmup to drive FX / inductor through
+        every shape without capturing cuda graphs yet."""
+        language_model = self._language_model
+        # Some multimodal models (e.g. Gemma4) store the inner transformer
+        # directly as `language_model` rather than wrapping it in a
+        # ForCausalLM that has a `.model` child.  Fall back to the module
+        # itself when `.model` is absent.
+        inner_model = getattr(language_model, "model", language_model)
+        compiler = self._compile_config.compiler
+        with enable_tc_piecewise_cuda_graph():
+            try:
+                if compiler != "eager":
+                    _toggle_multi_platform_ops(
+                        inner_model, reverse=False, num_tokens=16
+                    )
+
+                cuda_graph_runner._run_dummy_forward(
+                    num_tokens=cuda_graph_runner.capture_num_tokens[0]
+                )
+
+                if self._pool is None:
+                    self._pool = get_or_create_global_graph_memory_pool(
+                        self._device_module
+                    )
+                set_graph_pool_id(self._pool)
+
+                self.install_compile(
+                    inner_model,
+                    compile_config=self._compile_config,
+                    graph_pool=self._pool,
+                )
+
+                with enable_torch_compile_warmup():
+                    if is_hip():
+                        # AMD: single Dynamo trace is sufficient; the capture
+                        # phase does per-shape JIT kernel warmup before each
+                        # CUDA graph recording.  The N-iteration loop is
+                        # redundant and extremely slow on ROCm (~30 min).
+                        cuda_graph_runner._run_dummy_forward(
+                            num_tokens=cuda_graph_runner.capture_num_tokens[-1]
+                        )
+                    else:
+                        compile_range = (
+                            tqdm.tqdm(
+                                list(reversed(cuda_graph_runner.capture_num_tokens))
+                            )
+                            if get_parallel().tp_rank == 0
+                            else reversed(cuda_graph_runner.capture_num_tokens)
+                        )
+                        for num_tokens in compile_range:
+                            if get_parallel().tp_rank == 0:
+                                compile_range.set_description(
+                                    f"Compiling num tokens ({num_tokens=})"
+                                )
+                            cuda_graph_runner._run_dummy_forward(num_tokens=num_tokens)
+            finally:
+                _toggle_multi_platform_ops(inner_model, reverse=True, num_tokens=16)
+
+    @contextmanager
+    def capture_session(self, stream: torch.cuda.Stream):
+        self._capture_stream = stream
+        try:
+            with self.replay_session():
+                with set_pcg_capture_stream(stream):
+                    yield
+        finally:
+            self._capture_stream = None
+
+    def capture_one(
+        self,
+        shape_key: ShapeKey,
+        forward_fn: Callable[[], Any],
+        dummies: Optional[Any] = None,
+        post_warmup_hook: Optional[Callable[[], None]] = None,
+    ) -> None:
+        # Call 1 warms FX state; call 2 captures the cuda graph inside capture_session.
+        # See cuda_piecewise_backend.py for the FX backend that drives the capture.
+        for _ in range(2):
+            self._device_module.synchronize()
+            self._tp_group.barrier()
+            forward_fn()
+            if post_warmup_hook is not None:
+                post_warmup_hook()
+
+    def can_run(self, forward_batch: ForwardBatch, shape_key: ShapeKey) -> bool:
+        # torch.compile manages its per-shape cache internally.
+        # _run_compile_pass warms every shape in capture_num_tokens at __init__.
+        return True
+
+    @contextmanager
+    def replay_session(self):
+        with enable_tc_piecewise_cuda_graph():
+            yield
+
+    def replay(
+        self,
+        shape_key: ShapeKey,
+        static_forward_batch: ForwardBatch,
+        **kwargs,
+    ) -> Any:
+        return self._compiled_fn(
+            static_forward_batch.input_ids,
+            static_forward_batch.positions,
+            static_forward_batch,
+            **kwargs,
+        )
+
+    def cleanup(self) -> None:
+        self._compiled_fn = None
+        self._compile_config = None
+        self._language_model = None
+        self._pool = None

python/sglang/srt/model_executor/runner_backend/utils.py

@@ -0,0 +1,124 @@
+# 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.
+# ==============================================================================
+"""runner_backend utilities — phase → BaseCudaGraphBackend resolution.
+
+Centralizes per-phase backend resolution so platform overrides (NPU,
+out-of-tree) and future backend additions can plug in without
+modifying the runner files. Phase / backend identifiers used here
+live in :mod:`.cuda_graph_config`.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from sglang.srt.model_executor.cuda_graph_config import Backend
+from sglang.srt.model_executor.runner_backend.base_cuda_graph_backend import (
+    BaseCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import (
+    BreakableCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.full_cuda_graph_backend import (
+    FullCudaGraphBackend,
+)
+from sglang.srt.model_executor.runner_backend.tc_piecewise_cuda_graph_backend import (
+    TcPiecewiseCudaGraphBackend,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
+        BaseCudaGraphRunner,
+    )
+
+logger = logging.getLogger(__name__)
+
+# Track first occurrence of each fallback warning to avoid log spam.
+_TC_PIECEWISE_DECODE_FALLBACK_LOGGED = False
+
+
+def resolve_decode_backend(
+    cuda_graph_runner: BaseCudaGraphRunner,
+) -> BaseCudaGraphBackend:
+    """Pick a backend instance from cuda_graph_config['decode']['backend'].
+
+    NPU device returns NPUCudaGraphBackend regardless of mode (only
+    the Full-style backend is wired for NPU today).
+    """
+    model_runner = cuda_graph_runner.model_runner
+    cfg = model_runner.server_args.cuda_graph_config
+    backend_name = cfg.decode.backend if cfg is not None else Backend.FULL
+
+    enable_memory_saver = model_runner.server_args.enable_memory_saver
+
+    if model_runner.device == "npu":
+        from sglang.srt.hardware_backend.npu.graph_runner.npu_cudagraph_backend import (
+            NPUCudaGraphBackend,
+        )
+
+        return NPUCudaGraphBackend(
+            cuda_graph_runner, enable_memory_saver=enable_memory_saver
+        )
+    elif model_runner.device == "xpu":
+        if backend_name not in (Backend.FULL, Backend.DISABLED):
+            raise ValueError(
+                f"XPU only supports cuda_graph_config decode backend 'full', got '{backend_name}'"
+            )
+        from sglang.srt.hardware_backend.xpu.graph_runner.xpu_full_graph_backend import (
+            FullXPUGraphBackend,
+        )
+
+        return FullXPUGraphBackend(cuda_graph_runner)
+
+    if backend_name == Backend.BREAKABLE:
+        return BreakableCudaGraphBackend(
+            cuda_graph_runner,
+            enable_memory_saver=enable_memory_saver,
+            debug_eager=model_runner.server_args.debug_cuda_graph,
+        )
+    if backend_name == Backend.TC_PIECEWISE:
+        global _TC_PIECEWISE_DECODE_FALLBACK_LOGGED
+        if not _TC_PIECEWISE_DECODE_FALLBACK_LOGGED:
+            logger.warning(
+                "cuda_graph_config decode='tc_piecewise' is not yet implemented; "
+                "falling back to 'full'."
+            )
+            _TC_PIECEWISE_DECODE_FALLBACK_LOGGED = True
+    return FullCudaGraphBackend(
+        cuda_graph_runner, enable_memory_saver=enable_memory_saver
+    )
+
+
+def resolve_prefill_backend(
+    cuda_graph_runner: BaseCudaGraphRunner,
+) -> BaseCudaGraphBackend:
+    """Pick a backend instance from cuda_graph_config['prefill']['backend']."""
+    model_runner = cuda_graph_runner.model_runner
+    cfg = model_runner.server_args.cuda_graph_config
+    backend_name = cfg.prefill.backend if cfg is not None else Backend.TC_PIECEWISE
+
+    if backend_name == Backend.BREAKABLE:
+        return BreakableCudaGraphBackend(
+            cuda_graph_runner,
+            enable_memory_saver=model_runner.server_args.enable_memory_saver,
+            debug_eager=model_runner.server_args.debug_cuda_graph,
+        )
+    if backend_name == Backend.FULL:
+        return FullCudaGraphBackend(
+            cuda_graph_runner,
+            enable_memory_saver=model_runner.server_args.enable_memory_saver,
+        )
+    # Default: tc_piecewise.
+    return TcPiecewiseCudaGraphBackend(cuda_graph_runner)

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

@@ -0,0 +1,28 @@
+"""Low-level primitives used by the CUDA graph backends.
+
+Subpackages:
+  - breakable_cuda_graph: BreakableCUDAGraph + capture context,
+    eager_on_graph decorator, is_in_breakable_cuda_graph flag.
+  - piecewise_cuda_graph: shared piecewise context manager
+    (set_tc_piecewise_forward_context, is_in_tc_piecewise_cuda_graph).
+
+Backends in cuda_graph_backend/ import from here. Runners do not.
+"""
+
+# Generic failure-message hint for decode-style CUDA graph capture paths
+# (Full backend used by decode + EAGLE draft runners).
+CUDA_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 --cuda-graph-max-bs-decode to a smaller value (e.g., 16)\n"
+    "3. disable decode CUDA graph by --cuda-graph-backend-decode=disabled. "
+    "(Not recommended. Huge performance loss)\n"
+    "Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
+)
+
+PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG = (
+    "Fail when using backend: {backend} for prefill runner.\n"
+    "Possible suggestions:\n"
+    "{suggestions}"
+    "Open an issue on GitHub https://github.com/sgl-project/sglang/issues/new/choose \n"
+)

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/__init__.py

@@ -0,0 +1,21 @@
+"""Breakable primitives — segmented CUDA graph capture with eager break points.
+
+Public API (also reachable via the deeper module paths):
+  - BreakableCUDAGraph, BreakableCUDAGraphCapture — capture/replay
+  - eager_on_graph — decorator that marks a callable as a graph break
+  - break_graph — helper that inserts a bare graph break
+  - enable_breakable_cuda_graph — context that flips the Breakable runtime flag
+  - is_in_breakable_cuda_graph — runtime flag getter
+
+"""
+
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (  # noqa: F401
+    BreakableCUDAGraph,
+    BreakableCUDAGraphCapture,
+    break_graph,
+    eager_on_graph,
+)
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (  # noqa: F401
+    enable_breakable_cuda_graph,
+    is_in_breakable_cuda_graph,
+)

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py

@@ -0,0 +1,391 @@
+# 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.
+# ==============================================================================
+"""Breakable CUDA Graph: capture a region as a sequence of
+torch.cuda.CUDAGraph segments separated by eager break points.
+
+Each segment is a real torch.cuda.CUDAGraph. Its destructor calls
+releasePool on the shared mempool, so the pool's use_count tracks how
+many segments are alive; the pool stays pinned as long as any segment graph
+is alive. This lets weak_ref_tensor views of intermediate pool-allocated
+tensors remain valid across replays — we don't need Python-managed bridge
+buffers to keep break-point tensors at stable addresses.
+"""
+
+import logging
+import threading
+from contextvars import ContextVar
+from typing import Any, Callable
+
+import torch
+
+try:
+    from cuda.bindings import runtime as rt
+except ImportError:
+    rt = None
+
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.cuda_utils import (
+    checkCudaErrors,
+)
+from sglang.srt.utils import is_hip
+
+logger = logging.getLogger(__name__)
+
+__all__ = [
+    "eager_on_graph",
+    "BreakableCUDAGraph",
+    "BreakableCUDAGraphCapture",
+    "break_graph",
+]
+
+
+def _check_cuda_bindings():
+    if rt is None:
+        raise ImportError(
+            "Breakable CUDA graph on NVIDIA requires the 'cuda-python' package. "
+            "Install it with: pip install cuda-python"
+        )
+
+
+# Active BreakableCUDAGraphCapture context for the currently-capturing thread.
+# eager_on_graph's wrapper uses this to split the current torch.cuda.CUDAGraph
+# at break points.
+_current_capture_var: ContextVar["BreakableCUDAGraphCapture | None"] = ContextVar(
+    "current_capture", default=None
+)
+_current_stream_var: ContextVar[torch.cuda.Stream | None] = ContextVar(
+    "current_stream", default=None
+)
+_forked_streams_var: ContextVar[set[torch.cuda.Stream] | None] = ContextVar(
+    "forked_streams", default=None
+)
+
+
+def get_current_stream(device: torch.device | None = None) -> torch.cuda.Stream:
+    stream = _current_stream_var.get()
+    if stream is None:
+        return torch.cuda.current_stream(device)
+    return stream
+
+
+def _capture_status(stream_ptr: int) -> "rt.cudaStreamCaptureStatus":
+    _check_cuda_bindings()
+    status, *_ = checkCudaErrors(rt.cudaStreamGetCaptureInfo(stream_ptr))
+    return status
+
+
+def _is_stream_capturing(stream: torch.cuda.Stream) -> bool:
+    # On ROCm/HIP, cuda-python is unavailable, so use the portable torch API
+    # (which maps to the HIP runtime). On NVIDIA, keep querying the CUDA runtime
+    # directly via cuda-python: torch.cuda.is_current_stream_capturing() has
+    # proven unreliable there, so we preserve the original behavior.
+    if is_hip():
+        with torch.cuda.stream(stream):
+            return torch.cuda.is_current_stream_capturing()
+    return (
+        _capture_status(stream.cuda_stream)
+        == rt.cudaStreamCaptureStatus.cudaStreamCaptureStatusActive
+    )
+
+
+# Hook torch.cuda.Stream.wait_stream to track side-stream forks/joins that happen
+# during breakable capture. We need this because capture_end() on a torch
+# CUDAGraph fails if there are still side streams participating in the capture
+# — so before ending each segment we auto-join any forked-but-not-rejoined streams.
+_original_wait_stream: Callable | None = None
+_hook_lock = threading.Lock()
+_hook_refcount = 0
+
+
+def _hooked_wait_stream(self: torch.cuda.Stream, other: torch.cuda.Stream):
+    assert _original_wait_stream is not None
+    forked = _forked_streams_var.get()
+    if forked is None:
+        _original_wait_stream(self, other)
+        return
+    capturing = _current_stream_var.get()
+    if capturing is None:
+        _original_wait_stream(self, other)
+        return
+
+    cap_ptr = capturing.cuda_stream
+    is_self_cap = self is capturing or self.cuda_stream == cap_ptr
+    is_other_cap = other is capturing or other.cuda_stream == cap_ptr
+
+    if is_self_cap and not is_other_cap:
+        if not _is_stream_capturing(other):
+            return
+        _original_wait_stream(self, other)
+        forked.discard(other)
+    elif is_other_cap and not is_self_cap:
+        _original_wait_stream(self, other)
+        forked.add(self)
+    else:
+        _original_wait_stream(self, other)
+
+
+def _install_wait_stream_hook():
+    global _original_wait_stream, _hook_refcount
+    with _hook_lock:
+        if _hook_refcount == 0:
+            _original_wait_stream = torch.cuda.Stream.wait_stream
+            torch.cuda.Stream.wait_stream = _hooked_wait_stream  # type: ignore[assignment]
+        _hook_refcount += 1
+
+
+def _uninstall_wait_stream_hook():
+    global _original_wait_stream, _hook_refcount
+    with _hook_lock:
+        _hook_refcount -= 1
+        if _hook_refcount == 0:
+            assert _original_wait_stream is not None, "wait_stream hook not installed"
+            torch.cuda.Stream.wait_stream = _original_wait_stream  # type: ignore[assignment]
+            _original_wait_stream = None
+
+
+def _weak_ref_if_tensor(x):
+    """Return a weak-ref tensor view (shared storage, no refcount) for tensors;
+    recurse into tuples/lists; pass-through for non-tensors. Weak-ref'ing
+    captured args lets the shared mempool reclaim per-layer intermediates
+    between segments — storage stays alive for each segment CUDAGraph's
+    lifetime via its pool use_count.
+
+    weak_ref_tensors is imported lazily because it hard-raises on
+    platforms without a CUDA/HIP/NPU backend; we only reach this code during
+    an active Breakable capture, which runs only on those backends."""
+    if torch.is_tensor(x):
+        from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors
+
+        return weak_ref_tensors(x)
+    if isinstance(x, tuple):
+        return tuple(_weak_ref_if_tensor(e) for e in x)
+    if isinstance(x, list):
+        return [_weak_ref_if_tensor(e) for e in x]
+    return x
+
+
+def _copy_output(dst: Any, src: Any) -> Any:
+    """Copy src output into dst in-place where possible.
+
+    Handles plain tensors, tuples/lists of tensors, dataclass/object with
+    tensor attributes, and dicts of tensors. Returns dst if in-place copy
+    succeeded, otherwise returns src.
+    """
+    if torch.is_tensor(dst) and torch.is_tensor(src):
+        dst.copy_(src)
+        return dst
+
+    if (
+        isinstance(dst, (tuple, list))
+        and isinstance(src, (tuple, list))
+        and len(dst) == len(src)
+    ):
+        copied = [_copy_output(d, s) for d, s in zip(dst, src)]
+        return tuple(copied) if isinstance(dst, tuple) else copied
+
+    if hasattr(dst, "__dict__") and hasattr(src, "__dict__"):
+        for key, src_val in src.__dict__.items():
+            dst_val = getattr(dst, key, None)
+            if torch.is_tensor(dst_val) and torch.is_tensor(src_val):
+                dst_val.copy_(src_val)
+            else:
+                setattr(dst, key, src_val)
+        return dst
+
+    if isinstance(dst, dict) and isinstance(src, dict):
+        for key, src_val in src.items():
+            dst_val = dst.get(key)
+            if torch.is_tensor(dst_val) and torch.is_tensor(src_val):
+                dst_val.copy_(src_val)
+            else:
+                dst[key] = src_val
+        return dst
+
+    return src
+
+
+def eager_on_graph(enable: bool):
+    def decorator(inner: Callable):
+        if not enable:
+            return inner
+
+        def wrapper(*args, **kwargs):
+            capture = _current_capture_var.get()
+            if capture is None:
+                return inner(*args, **kwargs)
+
+            logger.debug("Break graph due to function: %s", inner.__name__)
+
+            # End the segment that captured up to this break point.
+            capture._end_current_segment()
+
+            # Run the eager function once so it allocates its outputs and
+            # writes real data into them.
+            output = inner(*args, **kwargs)
+
+            # Weak-ref captured inputs produced by graph segments. Their storage
+            # is pinned by the segment CUDAGraphs' mempool use-count, so Python
+            # refs do not need to keep every intermediate alive.
+            captured_inner = inner
+            captured_args = tuple(_weak_ref_if_tensor(a) for a in args)
+            captured_kwargs = {k: _weak_ref_if_tensor(v) for k, v in kwargs.items()}
+            # The eager break output is different: it is allocated between graph
+            # captures and is the static input address consumed by the next
+            # captured segment. Keep a strong reference so replay can safely
+            # copy fresh eager output into that bridge buffer.
+            captured_output = output
+
+            def replay_fn():
+                new_out = captured_inner(*captured_args, **captured_kwargs)
+                return _copy_output(captured_output, new_out)
+
+            capture.cuda_graph._break_fns.append(replay_fn)
+
+            # Start a fresh CUDAGraph segment for the remainder of the forward.
+            capture._begin_new_segment()
+            return output
+
+        return wrapper
+
+    return decorator
+
+
+class BreakableCUDAGraph:
+    """Container holding one torch.cuda.CUDAGraph per segment plus an
+    eager break function between consecutive segments."""
+
+    def __init__(self, deduped_cuda_graph=None) -> None:
+        self._segments: list[Any] = []
+        self._break_fns: list[Callable[[], Any]] = []
+        self._deduped_cuda_graph = deduped_cuda_graph
+
+    def replay(self) -> None:
+        stream = torch.cuda.current_stream()
+        token = _current_stream_var.set(stream)
+        try:
+            for i, seg in enumerate(self._segments):
+                seg.replay()
+                if i < len(self._break_fns):
+                    self._break_fns[i]()
+        finally:
+            _current_stream_var.reset(token)
+
+    def _append_segment(
+        self, graph: torch.cuda.CUDAGraph, needs_instantiate: bool
+    ) -> None:
+        if self._deduped_cuda_graph is not None:
+            self._segments.append(self._deduped_cuda_graph.register(graph))
+            return
+        if needs_instantiate:
+            graph.instantiate()
+        self._segments.append(graph)
+
+
+class BreakableCUDAGraphCapture:
+    """Context manager that captures the enclosed code as one or more
+    torch.cuda.CUDAGraph segments separated by eager break points.
+
+    Each segment shares the supplied pool (MempoolId_t tuple) so
+    pool-allocated intermediates can be reused across segments. While any
+    segment is alive, its beginAllocateToPool call keeps the mempool's
+    use_count > 0, which makes weak_ref_tensor of segment-allocated
+    tensors safe across subsequent replays.
+    """
+
+    def __init__(
+        self,
+        cuda_graph: BreakableCUDAGraph,
+        pool=None,
+        stream: torch.cuda.Stream | None = None,
+        capture_error_mode: str = "global",
+    ):
+        assert isinstance(
+            cuda_graph, BreakableCUDAGraph
+        ), "cuda_graph must be a BreakableCUDAGraph"
+        self.cuda_graph = cuda_graph
+        self._pool = pool if pool is not None else (0, 0)
+        self._stream = stream
+        self._capture_error_mode = capture_error_mode
+        self._stream_ctx = None
+        self._capture_token = None
+        self._stream_token = None
+        self._forked_token = None
+        self._current_graph: torch.cuda.CUDAGraph | None = None
+        self._current_graph_needs_instantiate = False
+
+    def __enter__(self):
+        _install_wait_stream_hook()
+        if self._stream is not None:
+            self._stream_ctx = torch.cuda.stream(self._stream)
+            self._stream_ctx.__enter__()
+        self._capture_token = _current_capture_var.set(self)
+        self._stream_token = _current_stream_var.set(
+            self._stream or torch.cuda.current_stream()
+        )
+        self._forked_token = _forked_streams_var.set(set())
+        self._begin_new_segment()
+        return self
+
+    def __exit__(self, *args: object):
+        try:
+            self._end_current_segment()
+        finally:
+            _forked_streams_var.reset(self._forked_token)
+            _current_stream_var.reset(self._stream_token)
+            _current_capture_var.reset(self._capture_token)
+            if self._stream_ctx is not None:
+                self._stream_ctx.__exit__(*args)
+                self._stream_ctx = None
+            _uninstall_wait_stream_hook()
+        return False
+
+    def _begin_new_segment(self) -> None:
+        # keep_graph retains the raw graph for dedup; skip it on the plain path.
+        if self.cuda_graph._deduped_cuda_graph is not None:
+            try:
+                graph = torch.cuda.CUDAGraph(keep_graph=True)
+                self._current_graph_needs_instantiate = True
+            except TypeError:
+                graph = torch.cuda.CUDAGraph()
+                self._current_graph_needs_instantiate = False
+        else:
+            graph = torch.cuda.CUDAGraph()
+            self._current_graph_needs_instantiate = False
+        graph.capture_begin(
+            pool=self._pool, capture_error_mode=self._capture_error_mode
+        )
+        self._current_graph = graph
+
+    def _end_current_segment(self) -> None:
+        # Auto-join any side streams forked during this segment but not joined.
+        main_stream = get_current_stream()
+        forked = _forked_streams_var.get()
+        if forked:
+            assert _original_wait_stream is not None
+            for side in list(forked):
+                if _is_stream_capturing(side):
+                    _original_wait_stream(main_stream, side)
+            forked.clear()
+        graph = self._current_graph
+        assert graph is not None
+        graph.capture_end()
+        self.cuda_graph._append_segment(graph, self._current_graph_needs_instantiate)
+        self._current_graph = None
+        self._current_graph_needs_instantiate = False
+
+
+@eager_on_graph(True)
+def break_graph() -> None:
+    """Insert a graph break. The @eager_on_graph decorator does the actual
+    segment split; this function body intentionally does nothing."""
+    pass

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/context.py

@@ -0,0 +1,60 @@
+# 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.
+# ==============================================================================
+"""Runtime state for the breakable CUDA graph runner."""
+
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+
+from sglang.srt.model_executor.cuda_graph_config import Backend
+from sglang.srt.model_executor.runner_backend_utils import (
+    PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG,
+)
+
+logger = logging.getLogger(__name__)
+
+_in_breakable_cuda_graph = False
+
+
+def is_in_breakable_cuda_graph() -> bool:
+    return _in_breakable_cuda_graph
+
+
+@contextmanager
+def enable_breakable_cuda_graph():
+    """Mark the enclosed scope as inside a BCG capture/replay. Any exception
+    raised inside is logged with the BCG-specific failure hint, then re-raised
+    for the caller to handle."""
+    global _in_breakable_cuda_graph
+    _in_breakable_cuda_graph = True
+    try:
+        yield
+    except Exception as exc:
+        msg = PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG.format(
+            backend=Backend.BREAKABLE, suggestions=BCG_FAILURE_HINT
+        )
+        logger.error(f"{type(exc).__name__}: {exc}\n{msg}")
+        raise
+    finally:
+        _in_breakable_cuda_graph = False
+
+
+BCG_FAILURE_HINT = (
+    "1. change to tc_piecewise by --cuda-graph-backend-prefill=tc_piecewise\n"
+    "2. disable the prefill CUDA graph by --cuda-graph-backend-prefill=disabled\n"
+    "3. if it is an OOM problem, set --mem-fraction-static to a smaller value "
+    "(e.g., 0.8 or 0.7) or set --cuda-graph-max-bs-prefill to a smaller value "
+    "(e.g., 2048)\n"
+)

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/cuda_utils.py

@@ -0,0 +1,48 @@
+# 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.
+# ==============================================================================
+"""CUDA runtime binding utilities."""
+
+try:
+    from cuda.bindings import runtime as rt
+except ImportError:
+    rt = None
+
+
+def _cudaGetErrorString(error):
+    if rt is None:
+        return "<cuda.bindings not available>"
+    err, msg = rt.cudaGetErrorString(error)
+    if err != rt.cudaError_t.cudaSuccess:
+        return "<unknown>"
+    if isinstance(msg, bytes):
+        return msg.decode("utf-8", "replace")
+    return str(msg)
+
+
+def checkCudaErrors(result):
+    if rt is None:
+        raise RuntimeError(
+            "cuda.bindings is not available. "
+            "Install it with: pip install cuda-python"
+        )
+    if result[0] != rt.cudaError_t.cudaSuccess:
+        raise RuntimeError(
+            f"CUDA error {int(result[0])}({_cudaGetErrorString(result[0])})"
+        )
+    if len(result) == 1:
+        return None
+    elif len(result) == 2:
+        return result[1]
+    else:
+        return result[1:]

python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/__init__.py

@@ -0,0 +1,22 @@
+"""Piecewise CUDA graph utilities — shared between Breakable and tc_piecewise backends.
+
+Public API:
+  - is_in_tc_piecewise_cuda_graph() — true while inside any piecewise capture.
+  - enable_tc_piecewise_cuda_graph() — context manager that toggles the flag.
+  - TcPiecewiseForwardContext + set_tc_piecewise_forward_context + get_tc_piecewise_forward_context.
+  - TCPCG_FAILURE_HINT — backend-switch suggestion plugged into
+    PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG by the prefill runner.
+
+The torch.compile-warmup flag (is_in_torch_compile_warmup) lives in
+sglang.srt.compilation.compile_phase — it is torch.compile-internal,
+not piecewise-shared.
+"""
+
+from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph.context_manager import (  # noqa: F401
+    TCPCG_FAILURE_HINT,
+    TcPiecewiseForwardContext,
+    enable_tc_piecewise_cuda_graph,
+    get_tc_piecewise_forward_context,
+    is_in_tc_piecewise_cuda_graph,
+    set_tc_piecewise_forward_context,
+)

python/sglang/srt/model_executor/runner_backend_utils/tc_piecewise_cuda_graph/context_manager.py

@@ -0,0 +1,121 @@
+"""CUDA graph capture context manager + forward-context propagation.
+
+Owns two pieces of cross-cutting state used by *every* piecewise-style
+backend (currently breakable + tc_piecewise):
+
+* _in_tc_piecewise_cuda_graph — a process-global flag set true while we
+  are inside the capture or replay window of a piecewise CUDA graph.
+  Read by model code that needs to take the static-buffer / fixed-shape
+  branch. See refactor/plan.md §6.5 for the full semantics.
+* TcPiecewiseForwardContext — a dataclass propagated across attention/MoE
+  layers during capture and replay so that submodules can reach the
+  current ForwardBatch and per-layer metadata without threading
+  arguments through every call site. Named TcPiecewise… (matches
+  Backend.TC_PIECEWISE + enable_tc_piecewise_cuda_graph) to
+  disambiguate from the per-forward-call
+  sglang.srt.model_executor.forward_context.ForwardContext.
+
+This module deliberately does **not** own torch.compile-specific state
+(warmup flag, capture stream); those live in compilation/compile_phase.py.
+"""
+
+from __future__ import annotations
+
+import logging
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, List, Optional
+
+from sglang.srt.model_executor.cuda_graph_config import Backend
+from sglang.srt.model_executor.runner_backend_utils import (
+    PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG,
+)
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+
+logger = logging.getLogger(__name__)
+
+_in_tc_piecewise_cuda_graph = False
+
+
+def is_in_tc_piecewise_cuda_graph() -> bool:
+    """True while inside tc_piecewise CUDA graph capture/replay."""
+    return _in_tc_piecewise_cuda_graph
+
+
+@contextmanager
+def enable_tc_piecewise_cuda_graph():
+    """Mark the enclosed scope as inside a tc_piecewise CUDA graph
+    capture/replay. Any exception raised inside is logged with the
+    PCG-specific failure hint, then re-raised for the caller to handle.
+    """
+    global _in_tc_piecewise_cuda_graph
+    _in_tc_piecewise_cuda_graph = True
+    try:
+        yield
+    except Exception as exc:
+        msg = PREFILL_CUDA_GRAPH_CAPTURE_FAILED_MSG.format(
+            backend=Backend.TC_PIECEWISE, suggestions=TCPCG_FAILURE_HINT
+        )
+        logger.error(f"{type(exc).__name__}: {exc}\n{msg}")
+        raise
+    finally:
+        _in_tc_piecewise_cuda_graph = False
+
+
+@dataclass
+class TcPiecewiseForwardContext:
+    forward_batch: Optional[ForwardBatch] = None
+    attention_layers: Optional[List[Any]] = field(default=None)
+    quant_config: Any = None
+    moe_layers: Optional[List[Any]] = field(default=None)
+    moe_fusions: Optional[List[Any]] = field(default=None)
+    dsa_indexers: Optional[List[Any]] = field(default=None)
+    num_tokens: Optional[int] = None
+    raw_num_tokens: Optional[int] = None
+
+
+_tc_piecewise_forward_context: Optional[TcPiecewiseForwardContext] = None
+
+
+def get_tc_piecewise_forward_context() -> Optional[TcPiecewiseForwardContext]:
+    return _tc_piecewise_forward_context
+
+
+@contextmanager
+def set_tc_piecewise_forward_context(
+    forward_batch: ForwardBatch,
+    attention_layers: List[Any],
+    quant_config: Any,
+    moe_layers: List[Any],
+    moe_fusions: List[Any],
+    dsa_indexers: Optional[List[Any]] = None,
+    num_tokens: Optional[int] = None,
+    raw_num_tokens: Optional[int] = None,
+):
+    global _tc_piecewise_forward_context
+    _tc_piecewise_forward_context = TcPiecewiseForwardContext(
+        forward_batch=forward_batch,
+        attention_layers=attention_layers,
+        quant_config=quant_config,
+        moe_layers=moe_layers,
+        moe_fusions=moe_fusions,
+        dsa_indexers=dsa_indexers,
+        num_tokens=num_tokens,
+        raw_num_tokens=raw_num_tokens,
+    )
+    try:
+        yield
+    finally:
+        _tc_piecewise_forward_context = None
+
+
+TCPCG_FAILURE_HINT = (
+    "1. change to breakable by --cuda-graph-backend-prefill=breakable\n"
+    "2. disable the prefill CUDA graph by --cuda-graph-backend-prefill=disabled\n"
+    "3. if it is an OOM problem, set --mem-fraction-static to a smaller value "
+    "(e.g., 0.8 or 0.7) or set --cuda-graph-max-bs-prefill to a smaller value "
+    "(e.g., 2048)\n"
+)

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

@@ -0,0 +1,28 @@
+"""Low-level utilities used by the CUDA graph runners.
+
+Mirror of cuda_graph_backend_utils/ for runner-side state — buffer
+dataclasses, process-global capture flags, the speculative-shared
+graph memory pool, and the DeepEP capture/replay adapter. Runners in
+cuda_graph_runner/ import from here; nothing here should import
+back into cuda_graph_runner/.
+"""
+
+from sglang.srt.model_executor.runner_utils.buffers import (  # noqa: F401
+    DecodeInputBuffers,
+    PrefillInputBuffers,
+    _grouped_foreach_copy_,
+)
+from sglang.srt.model_executor.runner_utils.capture_mode import (  # noqa: F401
+    _set_capture_lora_variant,
+    compile_in_capture_mode,
+    get_capture_lora_variant,
+    get_is_capture_mode,
+    model_capture_mode,
+)
+from sglang.srt.model_executor.runner_utils.deepep_adapter import (  # noqa: F401
+    DeepEPCudaGraphRunnerAdapter,
+)
+from sglang.srt.model_executor.runner_utils.pool import (  # noqa: F401
+    get_global_graph_memory_pool,
+    set_global_graph_memory_pool,
+)

python/sglang/srt/model_executor/runner_utils/buffers.py

@@ -0,0 +1,435 @@
+# 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.
+# ==============================================================================
+"""Static-buffer dataclasses used by the CUDA graph runners.
+
+DecodeInputBuffers backs the decode-phase capture/replay path.
+PrefillInputBuffers backs the prefill-phase capture/replay path.
+
+Both subclass ForwardInputBuffers so that buffer-pool sharing works
+the same way as for non-cuda-graph forward paths.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+from sglang.srt.environ import envs
+from sglang.srt.model_executor.forward_batch_info import (
+    ForwardBatch,
+    NgramEmbeddingInfo,
+    PPProxyTensors,
+    compute_local_num_token_non_padded,
+)
+from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
+
+_has_foreach_copy = hasattr(torch, "_foreach_copy_")
+
+
+def _grouped_foreach_copy_(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:
+    """Call torch._foreach_copy_ grouped by (dst_dtype, src_dtype) pairs."""
+
+    def foreach_copy(dsts: List[torch.Tensor], srcs: List[torch.Tensor]) -> None:
+        if _has_foreach_copy:
+            torch._foreach_copy_(dsts, srcs)
+        else:
+            for dst, src in zip(dsts, srcs):
+                dst.copy_(src)
+
+    groups: Dict[Tuple[torch.dtype, torch.dtype], Tuple[List, List]] = {}
+    for dst, src in zip(dsts, srcs):
+        key = (dst.dtype, src.dtype)
+        if key not in groups:
+            groups[key] = ([], [])
+        groups[key][0].append(dst)
+        groups[key][1].append(src)
+    for group_dsts, group_srcs in groups.values():
+        foreach_copy(group_dsts, group_srcs)
+
+
+@dataclass
+class DecodeInputBuffers(ForwardInputBuffers):
+    input_ids: torch.Tensor
+    input_embeds: torch.Tensor
+    req_pool_indices: torch.Tensor
+    seq_lens: torch.Tensor
+    seq_lens_cpu: torch.Tensor
+    out_cache_loc: torch.Tensor
+    positions: torch.Tensor
+    mrope_positions: torch.Tensor
+    num_token_non_padded: torch.Tensor
+    custom_mask: torch.Tensor
+    next_token_logits_buffer: torch.Tensor
+    mamba_track_indices: Optional[torch.Tensor]
+    mamba_track_mask: Optional[torch.Tensor]
+    global_num_tokens_gpu: torch.Tensor
+    global_num_tokens_for_logprob_gpu: torch.Tensor
+    encoder_lens: Optional[torch.Tensor]
+    pp_proxy_tensors: Optional[Dict[str, torch.Tensor]]
+    ngram_embedding_info: Optional[NgramEmbeddingInfo]
+    rids_int: Optional[torch.Tensor]
+    bootstrap_room_ids_int: Optional[torch.Tensor]
+
+    @classmethod
+    def create(
+        cls,
+        *,
+        device: torch.device,
+        max_bs: int,
+        max_num_token: int,
+        hidden_size: int,
+        next_token_logits_buffer: torch.Tensor,
+        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,
+    ) -> DecodeInputBuffers:
+        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,
+            )
+            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:
+                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 cls(
+            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,
+        )
+
+    def populate_from_forward_batch(
+        self,
+        *,
+        forward_batch: ForwardBatch,
+        raw_bs: int,
+        raw_num_token: int,
+        bs: int,
+        seq_len_fill_value: int,
+        require_gathered_buffer: bool,
+        num_tokens_per_bs: int,
+        dsa_enable_prefill_cp: bool,
+        enable_num_token_non_padded_flag: bool,
+        pp_proxy_tensors: Optional[PPProxyTensors] = None,
+    ):
+        if bs != raw_bs:
+            self.seq_lens.fill_(seq_len_fill_value)
+            self.out_cache_loc.zero_()
+            if self.mamba_track_indices is not None:
+                self.mamba_track_indices.zero_()
+            if self.mamba_track_mask is not None:
+                self.mamba_track_mask.fill_(False)
+
+        # Build batched copy lists for all GPU tensors.
+        dsts = [
+            self.input_ids[:raw_num_token],
+            self.req_pool_indices[:raw_bs],
+            self.seq_lens[:raw_bs],
+            self.out_cache_loc[:raw_num_token],
+            self.positions[:raw_num_token],
+        ]
+        srcs = [
+            forward_batch.input_ids,
+            forward_batch.req_pool_indices,
+            forward_batch.seq_lens,
+            forward_batch.out_cache_loc,
+            forward_batch.positions,
+        ]
+
+        if self.ngram_embedding_info is not None:
+            ngram_embedding_info = forward_batch.ngram_embedding_info
+            self.ngram_embedding_info.column_starts[:raw_bs].copy_(
+                ngram_embedding_info.column_starts
+            )
+            self.ngram_embedding_info.req_lens[:raw_bs].copy_(
+                ngram_embedding_info.req_lens
+            )
+
+        if (
+            self.mamba_track_indices is not None
+            and forward_batch.mamba_track_indices is not None
+        ):
+            dsts.append(self.mamba_track_indices[:raw_bs])
+            srcs.append(forward_batch.mamba_track_indices)
+        if (
+            self.mamba_track_mask is not None
+            and forward_batch.mamba_track_mask is not None
+        ):
+            dsts.append(self.mamba_track_mask[:raw_bs])
+            srcs.append(forward_batch.mamba_track_mask)
+
+        if self.encoder_lens is not None and forward_batch.encoder_lens is not None:
+            dsts.append(self.encoder_lens[:raw_bs])
+            srcs.append(forward_batch.encoder_lens)
+
+        if forward_batch.mrope_positions is not None:
+            dsts.append(self.mrope_positions[:, :raw_num_token])
+            srcs.append(forward_batch.mrope_positions)
+
+        if self.rids_int is not None and forward_batch.rids_int is not None:
+            dsts.append(self.rids_int[:raw_bs])
+            srcs.append(forward_batch.rids_int)
+        if (
+            self.bootstrap_room_ids_int is not None
+            and forward_batch.bootstrap_room_ids_int is not None
+        ):
+            dsts.append(self.bootstrap_room_ids_int[:raw_bs])
+            srcs.append(forward_batch.bootstrap_room_ids_int)
+
+        if require_gathered_buffer:
+            self.global_num_tokens_gpu.fill_(bs * num_tokens_per_bs)
+            self.global_num_tokens_for_logprob_gpu.fill_(bs * num_tokens_per_bs)
+
+        if enable_num_token_non_padded_flag:
+            if require_gathered_buffer and not dsa_enable_prefill_cp:
+                num_tokens_per_dp = bs * num_tokens_per_bs
+                local = compute_local_num_token_non_padded(
+                    global_num_token_non_padded=forward_batch.num_token_non_padded,
+                    num_tokens_per_dp=num_tokens_per_dp,
+                )
+                dsts.append(self.num_token_non_padded)
+                srcs.append(local)
+            else:
+                dsts.append(self.num_token_non_padded)
+                srcs.append(forward_batch.num_token_non_padded)
+
+        # Pipeline-parallel proxy tensors.
+        if pp_proxy_tensors is not None and self.pp_proxy_tensors is not None:
+            for key, buf in self.pp_proxy_tensors.items():
+                src = pp_proxy_tensors.tensors[key]
+                dim = src.shape[0]
+                dsts.append(buf[:dim])
+                srcs.append(src)
+
+        # Batch all GPU copies, grouped by dtype pair.
+        _grouped_foreach_copy_(dsts, srcs)
+
+        if forward_batch.seq_lens_cpu is not None:
+            if bs != raw_bs:
+                self.seq_lens_cpu.fill_(seq_len_fill_value)
+            self.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
+
+
+@dataclass
+class PrefillInputBuffers(ForwardInputBuffers):
+    input_ids: torch.Tensor
+    out_cache_loc: torch.Tensor
+    num_token_non_padded: torch.Tensor
+    mamba_track_indices: Optional[torch.Tensor]
+    mamba_track_mask: Optional[torch.Tensor]
+    mamba_track_seqlens: Optional[torch.Tensor]
+    positions: torch.Tensor
+    input_embeds: Optional[torch.Tensor]
+    mrope_positions: Optional[torch.Tensor]
+
+    @classmethod
+    def create(
+        cls,
+        *,
+        device: torch.device,
+        max_bs: int,
+        max_num_tokens: int,
+        cache_loc_dtype: torch.dtype,
+        is_multimodal: bool,
+        hidden_size: int,
+        dtype: torch.dtype,
+        enable_mamba_track: bool,
+    ) -> PrefillInputBuffers:
+        with torch.device(device):
+            input_ids = torch.zeros((max_num_tokens,), dtype=torch.int64)
+            out_cache_loc = torch.zeros((max_num_tokens,), dtype=cache_loc_dtype)
+            num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
+            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
+            )
+            mamba_track_seqlens = (
+                torch.zeros((max_bs,), dtype=torch.int32)
+                if enable_mamba_track
+                else None
+            )
+            positions = torch.zeros((max_num_tokens,), dtype=torch.int64)
+
+            if is_multimodal:
+                input_embeds = torch.zeros((max_num_tokens, hidden_size), dtype=dtype)
+                mrope_positions = torch.zeros((3, max_num_tokens), dtype=torch.int64)
+            else:
+                input_embeds = None
+                mrope_positions = None
+
+        return cls(
+            input_ids=input_ids,
+            out_cache_loc=out_cache_loc,
+            num_token_non_padded=num_token_non_padded,
+            mamba_track_indices=mamba_track_indices,
+            mamba_track_mask=mamba_track_mask,
+            mamba_track_seqlens=mamba_track_seqlens,
+            positions=positions,
+            input_embeds=input_embeds,
+            mrope_positions=mrope_positions,
+        )
+
+    def populate_from_forward_batch(
+        self,
+        *,
+        forward_batch: ForwardBatch,
+        raw_num_tokens: int,
+        static_num_tokens: int,
+        is_multimodal: bool,
+    ) -> None:
+        """Copy serving-batch values into static buffers and zero out
+        the padding region between raw_num_tokens and
+        static_num_tokens.
+        """
+        if static_num_tokens != raw_num_tokens:
+            self.out_cache_loc.zero_()
+            self.input_ids[raw_num_tokens:static_num_tokens].zero_()
+            self.positions[raw_num_tokens:static_num_tokens].zero_()
+            if is_multimodal:
+                self.input_embeds[raw_num_tokens:static_num_tokens].zero_()
+            if forward_batch.mrope_positions is not None:
+                self.mrope_positions[:, raw_num_tokens:static_num_tokens].zero_()
+
+        bs = forward_batch.batch_size
+
+        self.input_ids[:raw_num_tokens].copy_(forward_batch.input_ids)
+        self.positions[:raw_num_tokens].copy_(forward_batch.positions)
+        self.out_cache_loc[:raw_num_tokens].copy_(forward_batch.out_cache_loc)
+
+        if (
+            self.mamba_track_indices is not None
+            and forward_batch.mamba_track_indices is not None
+        ):
+            self.mamba_track_indices[:bs].copy_(forward_batch.mamba_track_indices)
+        if (
+            self.mamba_track_mask is not None
+            and forward_batch.mamba_track_mask is not None
+        ):
+            self.mamba_track_mask[:bs].copy_(forward_batch.mamba_track_mask)
+        if (
+            self.mamba_track_seqlens is not None
+            and forward_batch.mamba_track_seqlens is not None
+        ):
+            self.mamba_track_seqlens[:bs].copy_(forward_batch.mamba_track_seqlens)
+
+        if forward_batch.mrope_positions is not None:
+            self.mrope_positions[:, :raw_num_tokens].copy_(
+                forward_batch.mrope_positions
+            )

python/sglang/srt/model_executor/runner_utils/capture_mode.py

@@ -0,0 +1,78 @@
+# 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.
+# ==============================================================================
+"""Process-global capture-mode flags shared by the decode runner and the
+speculative-draft runners. Read by model code that needs to take a
+capture-time branch (e.g. lora dual-graph capture decides per-batch
+which variant to use).
+"""
+
+from __future__ import annotations
+
+from contextlib import contextmanager
+from typing import Optional
+
+import torch
+
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.context import (
+    is_in_breakable_cuda_graph,
+)
+
+# Detect whether the current forward pass is in capture mode.
+is_capture_mode = False
+
+# When capturing dual MoE backends, tracks which variant is being captured.
+# None = not dual, "lora" = capturing lora variant, "nolora" = capturing nolora variant.
+_capture_lora_variant: Optional[str] = None
+
+
+def get_is_capture_mode() -> bool:
+    return is_capture_mode or is_in_breakable_cuda_graph()
+
+
+def compile_in_capture_mode(func):
+    """Decorator: wrap func with torch.compile only when defined
+    inside model capture mode; passthrough otherwise.
+
+    Used by model code (e.g. DeepSeek-V4) to opt nested helpers into
+    torch.compile during cuda-graph capture without paying the
+    compilation cost in the eager forward path.
+    """
+    if is_capture_mode:
+        return torch.compile(func)
+    return func
+
+
+def get_capture_lora_variant() -> Optional[str]:
+    """Return the lora variant being captured, or None if not in dual capture."""
+    return _capture_lora_variant
+
+
+def _set_capture_lora_variant(variant: Optional[str]) -> None:
+    global _capture_lora_variant
+    _capture_lora_variant = variant
+
+
+@contextmanager
+def model_capture_mode():
+    global is_capture_mode
+    from sglang.srt.runtime_context import get_flags
+
+    # Disable dispose_tensor() during capture: freeing mid-capture records data_ptr()==0 into the graph.
+    is_capture_mode = True
+    get_flags().capture.disable_dispose_tensor = True
+    try:
+        yield
+    finally:
+        is_capture_mode = False
+        get_flags().capture.disable_dispose_tensor = False

python/sglang/srt/model_executor/runner_utils/deepep_adapter.py

@@ -0,0 +1,42 @@
+# 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.
+# ==============================================================================
+"""DeepEP capture/replay adapter — records the dispatch mode used during
+capture and re-applies it during replay so DeepEP all-to-all has
+consistent expert routing across the captured graph.
+"""
+
+from __future__ import annotations
+
+from sglang.srt.layers.moe.token_dispatcher.deepep import DeepEPBuffer
+from sglang.srt.layers.moe.utils import get_deepep_mode, get_moe_a2a_backend
+
+
+class DeepEPCudaGraphRunnerAdapter:
+    def __init__(self) -> None:
+        # Record DeepEP mode used during capture to ensure replay consistency.
+        self._captured_deepep_mode = None
+
+    def capture(self, is_extend_in_batch: bool) -> None:
+        if not get_moe_a2a_backend().is_deepep():
+            return
+        self._captured_deepep_mode = get_deepep_mode().resolve(
+            is_extend_in_batch=is_extend_in_batch
+        )
+        DeepEPBuffer.set_dispatch_mode(self._captured_deepep_mode)
+
+    def replay(self) -> None:
+        if not get_moe_a2a_backend().is_deepep():
+            return
+        assert self._captured_deepep_mode is not None
+        DeepEPBuffer.set_dispatch_mode(self._captured_deepep_mode)

python/sglang/srt/model_executor/runner_utils/pool.py

@@ -0,0 +1,40 @@
+# 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.
+# ==============================================================================
+"""Process-wide CUDA graph memory pool shared across the prefill and
+decode graph backends. The two phases never replay concurrently, so
+sharing one pool reserves only the larger phase's capture footprint.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Optional
+
+from sglang.srt.runtime_context import get_resources
+
+
+def get_global_graph_memory_pool() -> Optional[Any]:
+    return get_resources().graph_memory_pool
+
+
+def set_global_graph_memory_pool(val: Any) -> None:
+    get_resources().graph_memory_pool = val
+
+
+def get_or_create_global_graph_memory_pool(device_module: Any) -> Any:
+    """Return the shared graph memory pool, creating it on first use so
+    later backends reuse the same handle."""
+    resources = get_resources()
+    if resources.graph_memory_pool is None:
+        resources.graph_memory_pool = device_module.graph_pool_handle()
+    return resources.graph_memory_pool