Route the eager forward path through the CUDA graph input-buffer registry (#27407)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

python/sglang/srt/environ.py

@@ -763,6 +763,10 @@ class Envs:
     # CUDA graph
     SGLANG_PREP_IN_CUDA_GRAPH = EnvBool(True)
 
+    # Eager forward wraps the ForwardBatch's own tensors instead of copying them
+    # into the CUDA graph buffer registry (no per-iter device-to-device copy).
+    SGLANG_EAGER_INPUT_NO_COPY = EnvBool(False)
+
     # Distributed
     SGLANG_DSV4_FIX_TP_ATTN_A2A_SCATTER = EnvBool(True)
     SGLANG_SHARED_EXPERT_TP1 = EnvBool(False)

python/sglang/srt/model_executor/cuda_graph_buffer_registry.py

@@ -462,20 +462,11 @@ class CudaGraphBufferRegistry:
         padded_num_tokens: int,
         forward_batch_template: "ForwardBatch",
     ) -> "ForwardBatch":
-        """Return a FB view backed by registry slot buffers.
+        """Return a FB view (``dataclasses.replace`` of ``forward_batch_template``)
-
+        whose slot fields are buffer views and whose non-slot fields are carried
-        ``forward_batch_template`` provides the non-slot fields
+        from the template. A plain copy slot whose FB field is ``None`` this iter
-        (``forward_mode`` / ``spec_info`` / ``sampling_info`` /
+        is carried (not exposed as a stale buffer); computed slots are always
-        ``capture_hidden_mode`` / ``dp_*`` / ``lora_ids`` / ...). Slot
+        exposed.
-        fields are replaced with views into the registry buffers via
-        ``dataclasses.replace`` — the template itself is not mutated.
-
-        NOTE: currently parked / unused. It is NOT a drop-in for the decode
-        replay path's ``build_replay_fb_view``: it returns the *padded*
-        out_cache_loc slot slice (vs the raw ``fb.out_cache_loc`` that path
-        keeps), does not recompute ``seq_lens_sum`` for the padded tail, and
-        does not split ``forward_mode`` vs ``actual_forward_mode``. Reconcile
-        those before wiring it into any replay path.
         """
         import dataclasses
 
@@ -488,6 +479,14 @@ class CudaGraphBufferRegistry:
             # 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)
 
@@ -507,6 +506,7 @@ def build_decode_registry(
     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:
@@ -643,9 +643,12 @@ def build_decode_registry(
             )
         )
 
+    # 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):
-        # Filled with the padded token count on the gathered (DP) path; left
+            # Only the gathered (DP) path writes a value; otherwise left as init.
-        # untouched otherwise. Not an FB copy (copy_from_fb=False).
             if require_gathered_buffer:
                 buf.fill_(ctx.padded_num_tokens)
 
@@ -654,7 +657,10 @@ def build_decode_registry(
             if require_mlp_tp_gather
             else (lambda _bs, _mt: (1,))
         )
-    for _global_name in ("global_num_tokens_gpu", "global_num_tokens_for_logprob_gpu"):
+        for _global_name in (
+            "global_num_tokens_gpu",
+            "global_num_tokens_for_logprob_gpu",
+        ):
             slots.append(
                 GraphSlot(
                     _global_name,
@@ -760,12 +766,17 @@ def build_prefill_registry(
     hidden_size: int = 0,
     embed_dtype: Optional[torch.dtype] = None,
     enable_mamba_track: 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.replay_prepare``: ``input_ids`` / ``positions``
     / ``out_cache_loc`` / ``mrope_positions`` / ``input_embeds`` reset their
@@ -828,6 +839,7 @@ def build_prefill_registry(
                 slice_fn=lambda buf, n: buf[:, :n],
             )
         )
+        if register_input_embeds:
             slots.append(
                 GraphSlot(
                     "input_embeds",

python/sglang/srt/model_executor/model_runner.py

@@ -26,7 +26,7 @@ import socket
 import threading
 import time
 from collections import defaultdict
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import Any, Callable, List, Optional, Tuple, Union
 
@@ -140,6 +140,11 @@ from sglang.srt.model_executor.breakable_cuda_graph_runner import (
     BreakableCudaGraphRunner,
 )
 from sglang.srt.model_executor.cpu_graph_runner import CPUGraphRunner
+from sglang.srt.model_executor.cuda_graph_buffer_registry import (
+    CudaGraphBufferRegistry,
+    build_decode_registry,
+    build_prefill_registry,
+)
 from sglang.srt.model_executor.cuda_graph_runner import (
     CudaGraphRunner,
     _allocate_decode_buffers,
@@ -214,7 +219,7 @@ from sglang.srt.utils import (
     set_cuda_arch,
     slow_rank_detector,
 )
-from sglang.srt.utils.common import ceil_align, require_mlp_sync
+from sglang.srt.utils.common import ceil_align, next_power_of_2, require_mlp_sync
 from sglang.srt.utils.network import NetworkAddress, get_local_ip_auto
 from sglang.srt.utils.nvtx_pytorch_hooks import PytHooks
 from sglang.srt.utils.offloader import (
@@ -338,6 +343,14 @@ class ModelRunnerOutput:
     indexer_topk_output: Optional[TopkCaptureOutput] = None
 
 
+@dataclass
+class _EagerBufferRegistry:
+    # Lazily-built eager input-buffer registry plus the capacity it was sized to.
+    registry: Optional["CudaGraphBufferRegistry"] = None
+    max_bs: int = 0
+    max_num_tokens: int = 0
+
+
 class ModelRunner(ModelRunnerKVCacheMixin):
     """ModelRunner runs the forward passes of the models."""
 
@@ -414,6 +427,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         self.enable_elastic_ep = server_args.elastic_ep_backend is not None
         self.forward_pass_id = 0
         self.init_new_workspace = False
+        self._eager_decode_registry = _EagerBufferRegistry()
+        self._eager_prefill_registry = _EagerBufferRegistry()
         self.draft_model_idx = draft_model_idx
         self.enable_hisparse = server_args.enable_hisparse
 
@@ -3077,11 +3092,118 @@ class ModelRunner(ModelRunnerKVCacheMixin):
     def update_decode_attn_backend(self, stream_idx: int):
         self.decode_attn_backend = self.decode_attn_backend_group[stream_idx]
 
+    def _ensure_eager_registry(
+        self,
+        cache: _EagerBufferRegistry,
+        raw_bs: int,
+        raw_num_tokens: int,
+        build: Callable[[int, int], "CudaGraphBufferRegistry"],
+    ) -> "CudaGraphBufferRegistry":
+        # Built on first use and grown (next power of two) when a batch exceeds
+        # the current capacity.
+        if (
+            cache.registry is not None
+            and raw_bs <= cache.max_bs
+            and raw_num_tokens <= cache.max_num_tokens
+        ):
+            return cache.registry
+        cache.max_bs = next_power_of_2(max(raw_bs, cache.max_bs))
+        cache.max_num_tokens = next_power_of_2(
+            max(raw_num_tokens, cache.max_num_tokens)
+        )
+        cache.registry = build(cache.max_bs, cache.max_num_tokens)
+        return cache.registry
+
+    def _ensure_eager_decode_registry(
+        self, raw_bs: int, raw_num_tokens: int
+    ) -> "CudaGraphBufferRegistry":
+        is_encoder_decoder = self.model_config.is_encoder_decoder
+        return self._ensure_eager_registry(
+            self._eager_decode_registry,
+            raw_bs,
+            raw_num_tokens,
+            lambda bs, num_tokens: build_decode_registry(
+                device=self.device,
+                max_bs=bs,
+                max_num_token=num_tokens,
+                # Eager has no padding so this sentinel is never read; 0 avoids the
+                # cuda-graph-only fill-value method that some backends lack.
+                seq_len_fill_value=0,
+                cache_loc_dtype=torch.int64,
+                enable_mamba_track=(
+                    self.server_args.enable_mamba_extra_buffer()
+                    and self.spec_algorithm.is_none()
+                ),
+                is_encoder_decoder=is_encoder_decoder,
+                encoder_len_fill_value=(
+                    getattr(self.model_config.hf_config, "max_source_positions", 0)
+                    if is_encoder_decoder
+                    else 0
+                ),
+                enable_num_token_non_padded=False,
+                register_global_num_tokens=False,
+                require_gathered_buffer=False,
+                require_mlp_tp_gather=False,
+                dp_size=self.server_args.dp_size,
+                share_pool=False,
+                source=None,
+            ),
+        )
+
+    def _ensure_eager_prefill_registry(
+        self, raw_bs: int, raw_num_tokens: int
+    ) -> "CudaGraphBufferRegistry":
+        return self._ensure_eager_registry(
+            self._eager_prefill_registry,
+            raw_bs,
+            raw_num_tokens,
+            lambda bs, num_tokens: build_prefill_registry(
+                device=self.device,
+                max_bs=bs,
+                max_num_token=num_tokens,
+                cache_loc_dtype=torch.int64,
+                is_multimodal=self.is_multimodal,
+                enable_mamba_track=False,
+                register_input_embeds=False,
+                share_pool=False,
+                source=None,
+            ),
+        )
+
+    def _eager_fb_view(
+        self, forward_batch: ForwardBatch, pp_proxy_tensors=None
+    ) -> ForwardBatch:
+        if envs.SGLANG_EAGER_INPUT_NO_COPY.get():
+            return replace(forward_batch)
+        raw_bs = forward_batch.batch_size
+        raw_num_tokens = forward_batch.input_ids.shape[0]
+        ensure = (
+            self._ensure_eager_prefill_registry
+            if forward_batch.forward_mode.is_extend(include_draft_extend_v2=True)
+            else self._ensure_eager_decode_registry
+        )
+        registry = ensure(raw_bs, raw_num_tokens)
+        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 forward_decode(
         self,
         forward_batch: ForwardBatch,
         pp_proxy_tensors=None,
     ) -> Union[LogitsProcessorOutput, PPProxyTensors]:
+        if not self.server_args.enable_pdmux:
+            forward_batch = self._eager_fb_view(forward_batch, pp_proxy_tensors)
         # Set extra arguments
         pdmux_override = False
         if forward_batch.needs_forward_metadata_init():
@@ -3170,6 +3292,9 @@ class ModelRunner(ModelRunnerKVCacheMixin):
                 ret = self.piecewise_cuda_graph_runner.replay(forward_batch, **kwargs)
             return (ret, can_run_graph)
 
+        if not self.server_args.enable_pdmux:
+            forward_batch = self._eager_fb_view(forward_batch, pp_proxy_tensors)
+
         # Launch model forward
         if forward_batch.needs_forward_metadata_init():
             if hasattr(self.model, "prepare_forward_batch"):
@@ -3203,6 +3328,8 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         # called from the idle path can re-read a prior batch's req_pool
         # indices and trigger SWA mapping use-after-free.
         if forward_batch.batch_size > 0:
+            if not self.server_args.enable_pdmux:
+                forward_batch = self._eager_fb_view(forward_batch, pp_proxy_tensors)
             self.attn_backend.init_forward_metadata(forward_batch)
         else:
             self.attn_backend.forward_metadata = None

test/registered/unit/batch_overlap/test_tbo_filter_batch_marker.py

@@ -67,5 +67,74 @@ class TestTboFilterBatchMarker(CustomTestCase):
         self.assertFalse(child.forward_metadata_replan_equivalent)
 
 
+def _make_valued_batch(bs: int) -> ForwardBatch:
+    # Distinct per-position values so a filtered slice is unambiguous.
+    return ForwardBatch(
+        forward_mode=ForwardMode.TARGET_VERIFY,
+        batch_size=bs,
+        input_ids=torch.arange(bs, dtype=torch.long),
+        positions=torch.arange(bs, dtype=torch.long),
+        out_cache_loc=torch.arange(bs, dtype=torch.long) + 100,
+        req_pool_indices=torch.arange(bs, dtype=torch.long),
+        seq_lens=torch.arange(1, bs + 1, dtype=torch.int32),
+        seq_lens_cpu=torch.arange(1, bs + 1, dtype=torch.int32),
+        seq_lens_sum=int(torch.arange(1, bs + 1).sum()),
+        spec_info=None,
+    )
+
+
+class TestTboFilterBatchOnRegistryView(CustomTestCase):
+    """TBO runs one forward on the parent and splits via filter_batch, so an
+    eager registry-backed view must filter into children identically to the raw
+    batch."""
+
+    def _registry_view(self, batch: ForwardBatch) -> ForwardBatch:
+        from sglang.srt.model_executor.cuda_graph_buffer_registry import (
+            build_decode_registry,
+        )
+
+        bs = batch.batch_size
+        reg = build_decode_registry(
+            device=torch.device("cpu"),
+            max_bs=bs,
+            max_num_token=bs,
+            seq_len_fill_value=1,
+            cache_loc_dtype=torch.int64,
+            register_global_num_tokens=False,  # eager decode config
+            share_pool=False,
+            source=None,
+        )
+        reg.fill_from(
+            batch, raw_bs=bs, padded_bs=bs, raw_num_tokens=bs, padded_num_tokens=bs
+        )
+        view = reg.extract_buffer(
+            padded_bs=bs, padded_num_tokens=bs, forward_batch_template=batch
+        )
+        # The buffered fields are registry buffers (distinct storage), same values.
+        self.assertNotEqual(view.input_ids.data_ptr(), batch.input_ids.data_ptr())
+        self.assertTrue(torch.equal(view.input_ids, batch.input_ids))
+        return view
+
+    def test_filter_registry_view_matches_raw_batch(self):
+        raw_child = _filter(_make_valued_batch(8), lo=0, hi=4)
+        view_child = _filter(self._registry_view(_make_valued_batch(8)), lo=0, hi=4)
+
+        self.assertEqual(view_child.batch_size, raw_child.batch_size)
+        for f in (
+            "input_ids",
+            "positions",
+            "out_cache_loc",
+            "seq_lens",
+            "seq_lens_cpu",
+            "req_pool_indices",
+        ):
+            self.assertTrue(
+                torch.equal(getattr(view_child, f), getattr(raw_child, f)),
+                f"{f} differs between raw-batch and registry-view filtering",
+            )
+        # filter_batch still resets the plan marker on the child.
+        self.assertFalse(view_child.forward_metadata_ready)
+
+
 if __name__ == "__main__":
     unittest.main()

test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py

@@ -437,6 +437,65 @@ class TestMissingAndOptionalSlots(unittest.TestCase):
             )
         )
 
+    def test_extract_carries_none_for_absent_plain_slot(self):
+        # A plain copy slot absent this iter (mrope on a non-multimodal batch)
+        # must be carried as None, not exposed as the stale/zero buffer.
+        r = _make_registry(max_bs=4, max_num_tokens=8)
+        r.register_slot(
+            GraphSlot("input_ids", lambda bs, mt: (mt,), torch.int64, axis="tokens")
+        )
+        r.register_slot(
+            GraphSlot(
+                "mrope_positions",
+                lambda bs, mt: (3, mt),
+                torch.int64,
+                axis="tokens",
+                slice_fn=lambda buf, n: buf[:, :n],
+            )
+        )
+        fb = _MiniForwardBatch(
+            batch_size=2,
+            input_ids=torch.arange(2, dtype=torch.int64),
+            mrope_positions=None,  # non-multimodal: FB doesn't carry it
+        )
+        r.fill_from(fb, raw_bs=2, padded_bs=2, raw_num_tokens=2, padded_num_tokens=2)
+        fb_view = r.extract_buffer(
+            padded_bs=2, padded_num_tokens=2, forward_batch_template=fb
+        )
+        # input_ids was present -> buffer-backed; mrope absent -> carried None.
+        self.assertEqual(
+            fb_view.input_ids.data_ptr(), r.get_slot("input_ids").buffer.data_ptr()
+        )
+        self.assertIsNone(fb_view.mrope_positions)
+
+    def test_extract_exposes_computed_slot_even_when_fb_field_none(self):
+        # A computed slot (copy_from_fb=False) is always exposed, even when its
+        # FB field is None — the None-skip carry applies only to plain copies.
+        def _fill_two(buf, fb, ctx):
+            buf.fill_(2)
+
+        r = _make_registry(max_bs=4, max_num_tokens=8)
+        r.register_slot(
+            GraphSlot(
+                "global_num_tokens_gpu",
+                lambda bs, mt: (1,),
+                torch.int32,
+                axis="none",
+                copy_from_fb=False,
+                post_fill=_fill_two,
+            )
+        )
+        fb = _MiniForwardBatch(batch_size=2, global_num_tokens_gpu=None)
+        r.fill_from(fb, raw_bs=2, padded_bs=2, raw_num_tokens=2, padded_num_tokens=2)
+        fb_view = r.extract_buffer(
+            padded_bs=2, padded_num_tokens=2, forward_batch_template=fb
+        )
+        self.assertIsNotNone(fb_view.global_num_tokens_gpu)
+        self.assertEqual(
+            fb_view.global_num_tokens_gpu.data_ptr(),
+            r.get_slot("global_num_tokens_gpu").buffer.data_ptr(),
+        )
+
 
 class TestPostFillHook(unittest.TestCase):
     def test_post_fill_runs_after_copy(self):
@@ -888,6 +947,61 @@ class TestBuildDecodeRegistry(unittest.TestCase):
         # local = clamp(100 - rank*4, 0, 4) = 4  (NOT the raw FB copy of 100).
         self.assertEqual(int(src.num_token_non_padded.item()), 4)
 
+    def test_register_global_num_tokens_false_carries_fb_values(self):
+        # register_global_num_tokens=False (eager) excludes the computed
+        # global_num_tokens_* slots so the batch's DP values are carried, not
+        # clobbered by the zero buffer extract_buffer would otherwise expose.
+        from sglang.srt.model_executor.cuda_graph_buffer_registry import (
+            build_decode_registry,
+        )
+
+        reg = build_decode_registry(
+            device=torch.device("cpu"),
+            max_bs=4,
+            max_num_token=8,
+            seq_len_fill_value=5,
+            cache_loc_dtype=torch.int64,
+            register_global_num_tokens=False,
+            share_pool=False,
+            source=None,
+        )
+        self.assertFalse(reg.has_slot("global_num_tokens_gpu"))
+        self.assertFalse(reg.has_slot("global_num_tokens_for_logprob_gpu"))
+
+        gnt = torch.tensor([37], dtype=torch.int32)
+        gntlp = torch.tensor([41], dtype=torch.int32)
+        fb = _MiniForwardBatch(
+            batch_size=2,
+            input_ids=torch.arange(2, dtype=torch.int64),
+            positions=torch.arange(2, dtype=torch.int64),
+            out_cache_loc=torch.arange(2, dtype=torch.int64),
+            req_pool_indices=torch.zeros(2, dtype=torch.int64),
+            seq_lens=torch.full((2,), 5, dtype=torch.int32),
+            seq_lens_cpu=torch.full((2,), 5, dtype=torch.int32),
+            global_num_tokens_gpu=gnt,
+            global_num_tokens_for_logprob_gpu=gntlp,
+        )
+        reg.fill_from(fb, raw_bs=2, padded_bs=2, raw_num_tokens=2, padded_num_tokens=2)
+        fb_view = reg.extract_buffer(
+            padded_bs=2, padded_num_tokens=2, forward_batch_template=fb
+        )
+        # Carried from the batch (same tensors), not a zero registry buffer.
+        self.assertIs(fb_view.global_num_tokens_gpu, gnt)
+        self.assertIs(fb_view.global_num_tokens_for_logprob_gpu, gntlp)
+
+        # Default (graph path) still registers the computed slots.
+        reg2 = build_decode_registry(
+            device=torch.device("cpu"),
+            max_bs=4,
+            max_num_token=8,
+            seq_len_fill_value=5,
+            cache_loc_dtype=torch.int64,
+            share_pool=False,
+            source=None,
+        )
+        self.assertTrue(reg2.has_slot("global_num_tokens_gpu"))
+        self.assertTrue(reg2.has_slot("global_num_tokens_for_logprob_gpu"))
+
     def test_source_with_ngram_registers_structured_slots(self):
         from sglang.srt.model_executor.cuda_graph_buffer_registry import (
             build_decode_registry,
@@ -1188,6 +1302,42 @@ class TestBuildPrefillRegistry(unittest.TestCase):
         )
         self.assertTrue(torch.all(ids[3:8] == 0))
 
+    def test_register_input_embeds_false_keeps_mrope_carries_embeds(self):
+        # register_input_embeds=False (eager): mrope stays registered but
+        # input_embeds is carried from the FB (a read input), not a zero buffer.
+        from sglang.srt.model_executor.cuda_graph_buffer_registry import (
+            build_prefill_registry,
+        )
+
+        reg = build_prefill_registry(
+            device=torch.device("cpu"),
+            max_bs=2,
+            max_num_token=8,
+            cache_loc_dtype=torch.int64,
+            is_multimodal=True,
+            hidden_size=4,
+            embed_dtype=torch.float32,
+            register_input_embeds=False,
+            share_pool=False,
+            source=None,
+        )
+        self.assertTrue(reg.has_slot("mrope_positions"))
+        self.assertFalse(reg.has_slot("input_embeds"))
+        # extract_buffer carries the FB's real input_embeds (not a zero buffer).
+        embeds = torch.randn(3, 4)
+        fb = _MiniForwardBatch(
+            batch_size=1,
+            input_ids=torch.tensor([1, 2, 3], dtype=torch.int64),
+            positions=torch.tensor([0, 1, 2], dtype=torch.int64),
+            out_cache_loc=torch.tensor([7, 8, 9], dtype=torch.int64),
+            input_embeds=embeds,
+        )
+        reg.fill_from(fb, raw_bs=1, padded_bs=1, raw_num_tokens=3, padded_num_tokens=3)
+        fb_view = reg.extract_buffer(
+            padded_bs=1, padded_num_tokens=3, forward_batch_template=fb
+        )
+        self.assertIs(fb_view.input_embeds, embeds)
+
 
 class TestFillOncePolicy(unittest.TestCase):
     """FILL_ONCE initializes the whole buffer at alloc and never resets the