minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 13 Packages Projects Releases Wiki Activity

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

Browse Source 
Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This commit is contained in:
Cheng Wan
2026-06-03 20:52:56 -07:00
committed by GitHub
co-authored by Claude Opus 4.8
parent 3b7a258f63
commit 10ab7c919f
7 changed files with 277 additions and 318 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/breakable_cuda_graph_runner.py
+5 -39
View File
@@ -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()
python/sglang/srt/model_executor/cuda_graph_buffer_registry.py
+17 -9
View File
@@ -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,
python/sglang/srt/model_executor/cuda_graph_runner.py
+133 -199
View File
@@ -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,
)
python/sglang/srt/model_executor/input_buffers.py
+33
View File
@@ -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
+2 -2
View File
@@ -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
+4 -69
View File
@@ -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