From 10ab7c919f7b74223f39ac288cf7b6ac8c2ba3ed Mon Sep 17 00:00:00 2001 From: Cheng Wan <54331508+ch-wan@users.noreply.github.com> Date: Wed, 3 Jun 2026 20:52:56 -0700 Subject: [PATCH] [refactor] Retire DecodeInputBuffers / PrefillInputBuffers in favor of CudaGraphBufferRegistry (#27192) Co-authored-by: Claude Opus 4.8 (1M context) --- .../breakable_cuda_graph_runner.py | 44 +-- .../cuda_graph_buffer_registry.py | 26 +- .../srt/model_executor/cuda_graph_runner.py | 332 +++++++----------- .../srt/model_executor/input_buffers.py | 33 ++ .../sglang/srt/model_executor/model_runner.py | 4 +- .../piecewise_cuda_graph_runner.py | 73 +--- .../test_cuda_graph_buffer_registry.py | 83 +++++ 7 files changed, 277 insertions(+), 318 deletions(-) diff --git a/python/sglang/srt/model_executor/breakable_cuda_graph_runner.py b/python/sglang/srt/model_executor/breakable_cuda_graph_runner.py index 15681ea13..e78746a5f 100644 --- a/python/sglang/srt/model_executor/breakable_cuda_graph_runner.py +++ b/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): - 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 - + cache_loc_dtype = torch.int64 if not is_npu() else torch.int32 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( - 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). + # Registry owns (allocates + pools) the token-axis input buffers. 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() diff --git a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py b/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py index 90e9fadc4..6c189d0ca 100644 --- a/python/sglang/srt/model_executor/cuda_graph_buffer_registry.py +++ b/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 - # ForwardInputBuffers pool, so a registry can share physical storage - # (and data_ptr) with the legacy DecodeInputBuffers during migration. + # 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] = {} @@ -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 - ``source`` (e.g. a ``DecodeInputBuffers``) instead of allocating, so the - registry shares one physical allocation with that object. + 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, @@ -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 - ``source`` (the ``PrefillInputBuffers``) instead of allocating, so the - registry shares one physical allocation (and ``data_ptr``) with it. + 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, diff --git a/python/sglang/srt/model_executor/cuda_graph_runner.py b/python/sglang/srt/model_executor/cuda_graph_runner.py index cf40e18a2..3c90c034e 100644 --- a/python/sglang/srt/model_executor/cuda_graph_runner.py +++ b/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 -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, - 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, - ) -> "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", +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, +) -> SimpleNamespace: + """Allocate the FB-shared decode buffers as a namespace adopted by + ``build_decode_registry(source=...)``.""" + 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 ) - 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, + 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 ) - 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, - registry, - pp_proxy_tensors: Optional[PPProxyTensors] = None, - ): - # Reset padded tails + copy FB into the registry-adopted graph buffers - # (same storage the old per-field populate wrote). - registry.fill_from( - forward_batch, - raw_bs=raw_bs, - padded_bs=bs, - raw_num_tokens=raw_num_token, - padded_num_tokens=bs * num_tokens_per_bs, - pp_proxy_tensors=pp_proxy_tensors, - ) + 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.int32, + 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, + ) # 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() - # FB-shared slot registry, adopting the DecodeInputBuffers storage so - # 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. + share_input_buffers_in(self.buffers) + # The registry adopts these buffers (one data_ptr for capture + replay). 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 - # registry (it adopted the DecodeInputBuffers storage, so these are the - # same physical tensors); the rest still come off `buffers` directly. + # Graph inputs: owned slots come from the registry; the rest off `buffers`. 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( - forward_batch=forward_batch, + self.buffer_registry.fill_from( + forward_batch, raw_bs=raw_bs, - raw_num_token=raw_num_token, - bs=bs, - seq_len_fill_value=self.seq_len_fill_value, - 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, + padded_bs=bs, + raw_num_tokens=raw_num_token, + padded_num_tokens=bs * self.num_tokens_per_bs, pp_proxy_tensors=pp_proxy_tensors, ) diff --git a/python/sglang/srt/model_executor/input_buffers.py b/python/sglang/srt/model_executor/input_buffers.py index efb1f4255..444522f21 100644 --- a/python/sglang/srt/model_executor/input_buffers.py +++ b/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: diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index ce056571c..ba7ec25d2 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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, diff --git a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py b/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py index e455efaf2..ec2f5b2c2 100644 --- a/python/sglang/srt/model_executor/piecewise_cuda_graph_runner.py +++ b/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, Optional, Union +from typing import TYPE_CHECKING, 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 ( - 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). + # Registry owns (allocates + pools) the token-axis input buffers. 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 diff --git a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py b/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py index 602d3f9c3..0b1322589 100644 --- a/test/registered/unit/model_executor/test_cuda_graph_buffer_registry.py +++ b/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