[refactor] Retire DecodeInputBuffers / PrefillInputBuffers in favor of CudaGraphBufferRegistry (#27192)

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

python/sglang/srt/model_executor/breakable_cuda_graph_runner.py

@@ -172,30 +172,9 @@ class BreakableCudaGraphRunner:
         from sglang.srt.model_executor.cuda_graph_buffer_registry import (
             build_prefill_registry,
         )
-        from sglang.srt.model_executor.piecewise_cuda_graph_runner import (
-            PrefillInputBuffers,
-        )
         from sglang.srt.utils import is_npu
 
-        with torch.device(self.device):
+        cache_loc_dtype = torch.int64 if not is_npu() else torch.int32
-            input_ids = torch.zeros((self.max_num_tokens,), dtype=torch.int64)
-            out_cache_loc = torch.zeros(
-                (self.max_num_tokens,),
-                dtype=torch.int64 if not is_npu() else torch.int32,
-            )
-            positions = torch.zeros((self.max_num_tokens,), dtype=torch.int64)
-            if self.is_multimodal:
-                input_embeds = torch.zeros(
-                    (self.max_num_tokens, model_runner.model_config.hidden_size),
-                    dtype=model_runner.dtype,
-                )
-                mrope_positions = torch.zeros(
-                    (3, self.max_num_tokens), dtype=torch.int64
-                )
-            else:
-                input_embeds = None
-                mrope_positions = None
-
         if model_runner.is_draft_worker:
             from sglang.srt.speculative.eagle_utils import get_draft_hidden_dim
 
@@ -206,31 +185,18 @@ class BreakableCudaGraphRunner:
                 device=self.device,
             )
 
-        self.buffers = PrefillInputBuffers(
+        # Registry owns (allocates + pools) the token-axis input buffers.
-            input_ids=input_ids,
-            out_cache_loc=out_cache_loc,
-            mamba_track_indices=None,
-            mamba_track_mask=None,
-            mamba_track_seqlens=None,
-            positions=positions,
-            input_embeds=input_embeds,
-            mrope_positions=mrope_positions,
-        )
-        self.buffers.share_buffers()
-
-        # Token-axis FB-shared slot registry adopting the PrefillInputBuffers
-        # storage. Breakable has no mamba track and bs is not padded here, so
-        # there are no bs-axis slots (max_bs is unused).
         self.buffer_registry = build_prefill_registry(
             device=self.device,
             max_bs=1,
             max_num_token=self.max_num_tokens,
-            cache_loc_dtype=torch.int64 if not is_npu() else torch.int32,
+            cache_loc_dtype=cache_loc_dtype,
             is_multimodal=self.is_multimodal,
             hidden_size=model_runner.model_config.hidden_size,
             embed_dtype=model_runner.dtype,
             enable_mamba_track=False,
-            source=self.buffers,
+            share_pool=not is_npu(),
+            source=None,
         )
 
     @torch.no_grad()

python/sglang/srt/model_executor/cuda_graph_buffer_registry.py

@@ -288,9 +288,8 @@ class CudaGraphBufferRegistry:
         self.device = device
         self.max_bs = max_bs
         self.max_num_tokens = max_num_tokens
-        # When True, slot buffers are coalesced by name through the global
+        # Coalesce allocated slot buffers through the global pool; only applies
-        # ForwardInputBuffers pool, so a registry can share physical storage
+        # when allocating (bind/source bypasses the pool).
-        # (and data_ptr) with the legacy DecodeInputBuffers during migration.
         self.share_pool = share_pool
         self._slots: Dict[str, GraphSlot] = {}
 
@@ -470,6 +469,13 @@ class CudaGraphBufferRegistry:
         ``capture_hidden_mode`` / ``dp_*`` / ``lora_ids`` / ...). Slot
         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
 
@@ -519,9 +525,10 @@ def build_decode_registry(
     registered here — they are not per-replay FB copies (allocated and written
     elsewhere), so the runner keeps owning them.
 
-    When ``source`` is given, each slot adopts the same-named tensor off
+    When ``source`` is given (the decode buffer namespace from
-    ``source`` (e.g. a ``DecodeInputBuffers``) instead of allocating, so the
+    ``_allocate_decode_buffers``), each slot adopts the same-named tensor off
-    registry shares one physical allocation with that object.
+    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,
@@ -769,9 +776,10 @@ def build_prefill_registry(
     reset-only (``copy_from_fb=False``). ``mamba_track_*`` are bs-axis copies
     with no padding reset (bs is not padded on this path).
 
-    When ``source`` is given, each slot adopts the same-named tensor off
+    The piecewise / breakable runners pass ``source=None``, so the registry
-    ``source`` (the ``PrefillInputBuffers``) instead of allocating, so the
+    allocates (and owns) these buffers directly; ``share_pool`` then coalesces
-    registry shares one physical allocation (and ``data_ptr``) with it.
+    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,

python/sglang/srt/model_executor/cuda_graph_runner.py

@@ -22,10 +22,9 @@ import inspect
 import logging
 import os
 from contextlib import contextmanager
-from dataclasses import dataclass
 from functools import partial
 from types import SimpleNamespace
-from typing import TYPE_CHECKING, Callable, Dict, Optional, Union
+from typing import TYPE_CHECKING, Callable, Optional, Union
 
 import torch
 import tqdm
@@ -69,7 +68,7 @@ from sglang.srt.model_executor.forward_batch_info import (
     enable_num_token_non_padded,
 )
 from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
-from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
+from sglang.srt.model_executor.input_buffers import share_input_buffers_in
 from sglang.srt.multiplex.pdmux_context import get_current_stream_idx, get_stream_groups
 from sglang.srt.utils import (
     empty_context,
@@ -109,7 +108,7 @@ if TYPE_CHECKING:
 
 def build_replay_fb_view(
     forward_batch: "ForwardBatch",
-    buffers: "DecodeInputBuffers",
+    buffers,
     bs: int,
     raw_bs: int,
     num_tokens: int,
@@ -158,183 +157,130 @@ def build_replay_fb_view(
     )
 
 
-@dataclass
+def _allocate_decode_buffers(
-class DecodeInputBuffers(ForwardInputBuffers):
+    *,
-
+    device: torch.device,
-    input_ids: torch.Tensor
+    max_bs: int,
-    input_embeds: torch.Tensor
+    max_num_token: int,
-    req_pool_indices: torch.Tensor
+    hidden_size: int,
-    seq_lens: torch.Tensor
+    vocab_size: int,
-    seq_lens_cpu: torch.Tensor
+    dtype: torch.dtype,
-    out_cache_loc: torch.Tensor
+    dp_size: int,
-    positions: torch.Tensor
+    pp_size: int,
-    mrope_positions: torch.Tensor
+    is_encoder_decoder: bool,
-    num_token_non_padded: torch.Tensor
+    require_mlp_tp_gather: bool,
-    custom_mask: torch.Tensor
+    seq_len_fill_value: int,
-    next_token_logits_buffer: torch.Tensor
+    encoder_len_fill_value: int,
-    mamba_track_indices: Optional[torch.Tensor]
+    num_tokens_per_bs: int,
-    mamba_track_mask: Optional[torch.Tensor]
+    cache_loc_dtype: torch.dtype,
-    global_num_tokens_gpu: torch.Tensor
+    enable_mamba_track: bool,
-    global_num_tokens_for_logprob_gpu: torch.Tensor
+    ne_token_table: Optional[torch.Tensor] = None,
-    encoder_lens: Optional[torch.Tensor]
+    hc_hidden_size: Optional[int] = None,
-    pp_proxy_tensors: Optional[Dict[str, torch.Tensor]]
+) -> SimpleNamespace:
-    ngram_embedding_info: Optional["NgramEmbeddingInfo"]
+    """Allocate the FB-shared decode buffers as a namespace adopted by
-    rids_int: Optional[torch.Tensor]
+    ``build_decode_registry(source=...)``."""
-    bootstrap_room_ids_int: Optional[torch.Tensor]
+    with torch.device(device):
-
+        input_ids = torch.zeros((max_num_token,), dtype=torch.int64)
-    @classmethod
+        input_embeds = torch.zeros((max_num_token, hidden_size), dtype=dtype)
-    def create(
+        req_pool_indices = torch.zeros((max_bs,), dtype=torch.int64)
-        cls,
+        seq_lens = torch.full((max_bs,), seq_len_fill_value, dtype=torch.int32)
-        *,
+        out_cache_loc = torch.zeros((max_num_token,), dtype=cache_loc_dtype)
-        device: torch.device,
+        positions = torch.zeros((max_num_token,), dtype=torch.int64)
-        max_bs: int,
+        mrope_positions = torch.zeros((3, max_num_token), dtype=torch.int64)
-        max_num_token: int,
+        num_token_non_padded = torch.zeros((1,), dtype=torch.int32)
-        hidden_size: int,
+        custom_mask = torch.ones(
-        vocab_size: int,
+            (max_bs * seq_len_fill_value + max_num_token) * num_tokens_per_bs,
-        dtype: torch.dtype,
+            dtype=torch.bool,
-        dp_size: int,
+        )
-        pp_size: int,
+        next_token_logits_buffer = torch.zeros(
-        is_encoder_decoder: bool,
+            (max_num_token, vocab_size),
-        require_mlp_tp_gather: bool,
+            dtype=torch.float,
-        seq_len_fill_value: int,
+        )
-        encoder_len_fill_value: int,
+        mamba_track_indices = (
-        num_tokens_per_bs: int,
+            torch.zeros((max_bs,), dtype=torch.int64) if enable_mamba_track else None
-        cache_loc_dtype: torch.dtype,
+        )
-        enable_mamba_track: bool,
+        mamba_track_mask = (
-        ne_token_table: Optional[torch.Tensor] = None,
+            torch.zeros((max_bs,), dtype=torch.bool) if enable_mamba_track else None
-        hc_hidden_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.int32)
-            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
-                    )
-            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),
-                )
-                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
-
-        # Keep seq_lens_cpu as a true CPU tensor, like the old implementation.
-        seq_lens_cpu = torch.full(
-            (max_bs,),
-            seq_len_fill_value,
-            dtype=torch.int32,
-            device="cpu",
         )
 
-        return cls(
+        if pp_size > 1:
-            input_ids=input_ids,
+            # mHC (e.g. DSV4) flattens residual into hidden_states (size = hc_hidden_size).
-            input_embeds=input_embeds,
+            is_mhc = hc_hidden_size is not None
-            req_pool_indices=req_pool_indices,
+            hs = hc_hidden_size if is_mhc else hidden_size
-            seq_lens=seq_lens,
+            pp_proxy_tensors = {
-            seq_lens_cpu=seq_lens_cpu,
+                "hidden_states": torch.zeros((max_bs, hs), dtype=dtype),
-            out_cache_loc=out_cache_loc,
+            }
-            positions=positions,
+            if not is_mhc:
-            mrope_positions=mrope_positions,
+                pp_proxy_tensors["residual"] = torch.zeros(
-            num_token_non_padded=num_token_non_padded,
+                    (max_bs, hidden_size), dtype=dtype
-            custom_mask=custom_mask,
+                )
-            next_token_logits_buffer=next_token_logits_buffer,
+        else:
-            mamba_track_indices=mamba_track_indices,
+            pp_proxy_tensors = None
-            mamba_track_mask=mamba_track_mask,
+
-            encoder_lens=encoder_lens,
+        if is_encoder_decoder:
-            global_num_tokens_gpu=global_num_tokens_gpu,
+            encoder_lens = torch.full(
-            global_num_tokens_for_logprob_gpu=global_num_tokens_for_logprob_gpu,
+                (max_bs,), encoder_len_fill_value, dtype=torch.int32
-            pp_proxy_tensors=pp_proxy_tensors,
+            )
-            ngram_embedding_info=ngram_embedding_info,
+        else:
-            rids_int=rids_int,
+            encoder_lens = None
-            bootstrap_room_ids_int=bootstrap_room_ids_int,
+
+        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),
+            )
+            if ne_token_table is not None
+            else None
         )
 
-    def populate_from_forward_batch(
+        if envs.SGLANG_KV_CANARY_ENABLE_TOKEN_ORACLE.get():
-        self,
+            rids_int = torch.zeros((max_bs,), dtype=torch.int64)
-        *,
+            bootstrap_room_ids_int = torch.full((max_bs,), -1, dtype=torch.int64)
-        forward_batch: ForwardBatch,
+        else:
-        raw_bs: int,
+            rids_int = None
-        raw_num_token: int,
+            bootstrap_room_ids_int = None
-        bs: int,
+
-        seq_len_fill_value: int,
+    seq_lens_cpu = torch.full(
-        require_gathered_buffer: bool,
+        (max_bs,),
-        num_tokens_per_bs: int,
+        seq_len_fill_value,
-        dsa_enable_prefill_cp: bool,
+        dtype=torch.int32,
-        enable_num_token_non_padded_flag: bool,
+        device="cpu",
-        registry,
+    )
-        pp_proxy_tensors: Optional[PPProxyTensors] = None,
+
-    ):
+    return SimpleNamespace(
-        # Reset padded tails + copy FB into the registry-adopted graph buffers
+        input_ids=input_ids,
-        # (same storage the old per-field populate wrote).
+        input_embeds=input_embeds,
-        registry.fill_from(
+        req_pool_indices=req_pool_indices,
-            forward_batch,
+        seq_lens=seq_lens,
-            raw_bs=raw_bs,
+        seq_lens_cpu=seq_lens_cpu,
-            padded_bs=bs,
+        out_cache_loc=out_cache_loc,
-            raw_num_tokens=raw_num_token,
+        positions=positions,
-            padded_num_tokens=bs * num_tokens_per_bs,
+        mrope_positions=mrope_positions,
-            pp_proxy_tensors=pp_proxy_tensors,
+        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,
+    )
 
 
 # Detect whether the current forward pass is in capture mode
@@ -630,7 +576,7 @@ class CudaGraphRunner:
 
         if self.require_gathered_buffer:
             assert self.require_mlp_tp_gather or self.require_attn_tp_gather
-        self.buffers: DecodeInputBuffers = DecodeInputBuffers.create(
+        self.buffers = _allocate_decode_buffers(
             device=self.device,
             max_bs=self.max_bs,
             max_num_token=self.max_num_token,
@@ -653,11 +599,8 @@ class CudaGraphRunner:
                 self.model_runner.model_config, "hc_hidden_size", None
             ),
         )
-        self.buffers.share_buffers()
+        share_input_buffers_in(self.buffers)
-        # FB-shared slot registry, adopting the DecodeInputBuffers storage so
+        # The registry adopts these buffers (one data_ptr for capture + replay).
-        # it mirrors the same physical buffers (stable data_ptr for capture vs
-        # replay). This is the unified fill/extract surface that eager /
-        # capture / replay migrate onto, replacing populate_from_forward_batch.
         self.buffer_registry = build_decode_registry(
             device=self.device,
             max_bs=self.max_bs,
@@ -897,14 +840,12 @@ class CudaGraphRunner:
     def capture_one_batch_size(
         self, bs: int, forward: Callable, stream_idx: Optional[int] = None
     ):
-        buffers: DecodeInputBuffers = self.buffers
+        buffers = self.buffers
         graph = self._create_device_graph()
         stream = self.stream
         num_tokens = bs * self.num_tokens_per_bs
 
-        # Graph inputs. The registry-owned FB-shared slots come from the
+        # Graph inputs: owned slots come from the registry; the rest off `buffers`.
-        # registry (it adopted the DecodeInputBuffers storage, so these are the
-        # same physical tensors); the rest still come off `buffers` directly.
         registry = self.buffer_registry
 
         def _slot(name):
@@ -1168,19 +1109,12 @@ class CudaGraphRunner:
             index = bisect.bisect_left(self.capture_bs, raw_bs)
         bs = self.capture_bs[index]
 
-        buffers.populate_from_forward_batch(
+        self.buffer_registry.fill_from(
-            forward_batch=forward_batch,
+            forward_batch,
             raw_bs=raw_bs,
-            raw_num_token=raw_num_token,
+            padded_bs=bs,
-            bs=bs,
+            raw_num_tokens=raw_num_token,
-            seq_len_fill_value=self.seq_len_fill_value,
+            padded_num_tokens=bs * self.num_tokens_per_bs,
-            require_gathered_buffer=self.require_gathered_buffer,
-            num_tokens_per_bs=self.num_tokens_per_bs,
-            # Parameter name retained for API stability; semantically this is
-            # "any prefill-CP flavor enabled" (DSA CP or MLA CP).
-            dsa_enable_prefill_cp=self.enable_prefill_cp,
-            enable_num_token_non_padded_flag=enable_num_token_non_padded(),
-            registry=self.buffer_registry,
             pp_proxy_tensors=pp_proxy_tensors,
         )
 

python/sglang/srt/model_executor/input_buffers.py

@@ -24,6 +24,14 @@ def share_input_buffer(name: str, new_buffer: torch.Tensor) -> torch.Tensor:
     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)
@@ -33,6 +41,31 @@ def share_input_buffer(name: str, new_buffer: torch.Tensor) -> torch.Tensor:
     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:
 

python/sglang/srt/model_executor/model_runner.py

@@ -142,7 +142,7 @@ from sglang.srt.model_executor.breakable_cuda_graph_runner import (
 from sglang.srt.model_executor.cpu_graph_runner import CPUGraphRunner
 from sglang.srt.model_executor.cuda_graph_runner import (
     CudaGraphRunner,
-    DecodeInputBuffers,
+    _allocate_decode_buffers,
     set_torch_compile_config,
 )
 from sglang.srt.model_executor.forward_batch_info import (
@@ -2566,7 +2566,7 @@ class ModelRunner(ModelRunnerKVCacheMixin):
         if require_gathered_buffer(self.server_args):
             assert require_mlp_tp_gather_ or require_attn_tp_gather(self.server_args)
 
-        buffers: DecodeInputBuffers = DecodeInputBuffers.create(
+        buffers = _allocate_decode_buffers(
             device=self.device,
             max_bs=batch_size,
             max_num_token=num_tokens,

python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py

@@ -20,8 +20,7 @@ import gc
 import logging
 import warnings
 from contextlib import contextmanager
-from dataclasses import dataclass
+from typing import TYPE_CHECKING, Union
-from typing import TYPE_CHECKING, Optional, Union
 
 import torch
 import tqdm
@@ -60,7 +59,6 @@ from sglang.srt.model_executor.forward_batch_info import (
     PPProxyTensors,
 )
 from sglang.srt.model_executor.forward_context import ForwardContext, forward_context
-from sglang.srt.model_executor.input_buffers import ForwardInputBuffers
 from sglang.srt.utils import (
     get_available_gpu_memory,
     is_musa,
@@ -79,18 +77,6 @@ if TYPE_CHECKING:
 _is_musa = is_musa()
 
 
-@dataclass
-class PrefillInputBuffers(ForwardInputBuffers):
-    input_ids: torch.Tensor
-    out_cache_loc: 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]
-
-
 @contextmanager
 def freeze_gc(enable_cudagraph_gc: bool):
     """
@@ -242,62 +228,10 @@ class PiecewiseCudaGraphRunner:
         # CUDA graph capture must use the same flag value as replay for those models.
         self.capture_return_pooled_hidden_states = not model_runner.is_generation
 
-        # Graph inputs
         with torch.device(self.device):
-            input_ids = torch.zeros((self.max_num_tokens,), dtype=torch.int64)
-            out_cache_loc = torch.zeros(
-                (self.max_num_tokens,), dtype=self._cache_loc_dtype()
-            )
-            mamba_track_indices = (
-                torch.zeros((self.max_bs,), dtype=torch.int64)
-                if self.mamba_track_enabled
-                else None
-            )
-            mamba_track_mask = (
-                torch.zeros((self.max_bs,), dtype=torch.bool)
-                if self.mamba_track_enabled
-                else None
-            )
-            mamba_track_seqlens = (
-                torch.zeros((self.max_bs,), dtype=torch.int32)
-                if self.mamba_track_enabled
-                else None
-            )
-            positions = torch.zeros((self.max_num_tokens,), dtype=torch.int64)
-
             self.tbo_plugin = TboCudaGraphRunnerPlugin()
 
-            if (
+        # Registry owns (allocates + pools) the token-axis input buffers.
-                self.is_multimodal
-            ):  # Only create input_embeds and mrope_positions for multimodal model to save memory
-                # 1. In multimodal, we only compile and capture the language model part.
-                # 2. The embedder is outside of the graph, but cuda graph requires the input embeds to have a fixed memory address.
-                # 3. Input embeds is a pre-allocated buffer. In model.forward, we copy the embed output to this buffer.
-                input_embeds = torch.zeros(
-                    (self.max_num_tokens, self.model_runner.model_config.hidden_size),
-                    dtype=self.model_runner.dtype,
-                )
-                mrope_positions = torch.zeros(
-                    (3, self.max_num_tokens), dtype=torch.int64
-                )
-            else:
-                input_embeds = None
-                mrope_positions = None
-
-        self.buffers = PrefillInputBuffers(
-            input_ids=input_ids,
-            out_cache_loc=out_cache_loc,
-            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,
-        )
-        self.buffers.share_buffers()
-
-        # Token-axis FB-shared slot registry, adopting the PrefillInputBuffers
-        # storage (one data_ptr shared with capture + replay).
         self.buffer_registry = build_prefill_registry(
             device=self.device,
             max_bs=self.max_bs,
@@ -307,7 +241,8 @@ class PiecewiseCudaGraphRunner:
             hidden_size=self.model_runner.model_config.hidden_size,
             embed_dtype=self.model_runner.dtype,
             enable_mamba_track=self.mamba_track_enabled,
-            source=self.buffers,
+            share_pool=not is_npu(),
+            source=None,
         )
 
         self.attention_layers = self.model_runner.attention_layers

test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py

@@ -842,6 +842,52 @@ class TestBuildDecodeRegistry(unittest.TestCase):
                 name,
             )
 
+    def test_num_token_non_padded_gathered_dp_branch(self):
+        import unittest.mock as mock
+
+        from sglang.srt.model_executor import forward_batch_info as fbi
+        from sglang.srt.model_executor.cuda_graph_buffer_registry import (
+            build_decode_registry,
+        )
+
+        ntnp = torch.zeros(1, dtype=torch.int32)
+        src = SimpleNamespace(
+            input_ids=torch.zeros(8, dtype=torch.int64),
+            positions=torch.zeros(8, dtype=torch.int64),
+            out_cache_loc=torch.zeros(8, dtype=torch.int64),
+            req_pool_indices=torch.zeros(4, dtype=torch.int64),
+            seq_lens=torch.full((4,), 5, dtype=torch.int32),
+            seq_lens_cpu=torch.full((4,), 5, dtype=torch.int32),
+            mrope_positions=torch.zeros((3, 8), dtype=torch.int64),
+            num_token_non_padded=ntnp,
+            global_num_tokens_gpu=torch.zeros(1, dtype=torch.int32),
+            global_num_tokens_for_logprob_gpu=torch.zeros(1, dtype=torch.int32),
+        )
+        # Gathered (DP) path: post_fill overwrites the FB copy with the local
+        # count. Pin attn-TP (size=2, rank=0) so the result is deterministic.
+        with mock.patch.object(
+            fbi, "get_attention_tp_size", return_value=2
+        ), mock.patch.object(fbi, "get_attention_tp_rank", return_value=0):
+            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,
+                enable_num_token_non_padded=True,
+                require_gathered_buffer=True,
+                source=src,
+            )
+            fb = _MiniForwardBatch(
+                num_token_non_padded=torch.tensor([100], dtype=torch.int32),
+            )
+            reg.fill_from(
+                fb, raw_bs=4, padded_bs=4, raw_num_tokens=4, padded_num_tokens=8
+            )
+        # tokens_per_rank = padded_num_tokens(8) // attn_tp_size(2) = 4;
+        # 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_source_with_ngram_registers_structured_slots(self):
         from sglang.srt.model_executor.cuda_graph_buffer_registry import (
             build_decode_registry,
@@ -1105,6 +1151,43 @@ class TestBuildPrefillRegistry(unittest.TestCase):
         reg.fill_from(fb, raw_bs=2, padded_bs=2, raw_num_tokens=3, padded_num_tokens=8)
         self.assertTrue(torch.equal(idx, torch.tensor([3, 4], dtype=torch.int64)))
 
+    def test_source_none_owns_allocated_buffers(self):
+        # source=None -> the registry allocates (owns) every slot.
+        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=16,
+            cache_loc_dtype=torch.int64,
+            is_multimodal=True,
+            hidden_size=4,
+            embed_dtype=torch.float32,
+            enable_mamba_track=True,
+            share_pool=False,
+            source=None,
+        )
+        self.assertEqual(tuple(reg.get_slot("input_ids").buffer.shape), (16,))
+        self.assertEqual(tuple(reg.get_slot("positions").buffer.shape), (16,))
+        self.assertEqual(tuple(reg.get_slot("out_cache_loc").buffer.shape), (16,))
+        self.assertEqual(tuple(reg.get_slot("mrope_positions").buffer.shape), (3, 16))
+        self.assertEqual(tuple(reg.get_slot("input_embeds").buffer.shape), (16, 4))
+        self.assertEqual(tuple(reg.get_slot("mamba_track_indices").buffer.shape), (2,))
+        # Fills + ZERO-tails the pad with no backing source.
+        fb = _MiniForwardBatch(
+            input_ids=torch.tensor([1, 2, 3], dtype=torch.int64),
+            positions=torch.tensor([4, 5, 6], dtype=torch.int64),
+            out_cache_loc=torch.tensor([7, 8, 9], dtype=torch.int64),
+        )
+        reg.fill_from(fb, raw_bs=1, padded_bs=1, raw_num_tokens=3, padded_num_tokens=8)
+        ids = reg.get_slot("input_ids").buffer
+        self.assertTrue(
+            torch.equal(ids[:3], torch.tensor([1, 2, 3], dtype=torch.int64))
+        )
+        self.assertTrue(torch.all(ids[3:8] == 0))
+
 
 class TestFillOncePolicy(unittest.TestCase):
     """FILL_ONCE initializes the whole buffer at alloc and never resets the