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

[DSV4][BCG] Optimize the heavy memory use of C4 Indexer when BCG is enabled (#36534)

Browse Source 
Co-authored-by: Yuwei An <ayw.sirius19@gmail.com>
This commit is contained in:
Siyuan Chen
2026-09-15 17:23:11 -07:00
committed by GitHub
co-authored by Yuwei An
parent bbef93e22a
commit c9a8fba991
23 changed files with 344 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/arg_groups/cuda_graph_hook.py
+64 -1
View File
@@ -4,7 +4,7 @@
from __future__ import annotations
import logging
from typing import Any
from typing import Any, Optional
from sglang.srt.arg_groups.overrides import (
attention_backends_of,
@@ -87,6 +87,12 @@ def parse_cuda_graph_config(server_args: Any):
# decode is implemented; today decode ignores it.
_set(Phase.DECODE, "tc_compiler", cfg.cuda_graph_tc_compiler)
_set(Phase.PREFILL, "tc_compiler", cfg.cuda_graph_tc_compiler)
if cfg.cuda_graph_prefill_max_context is not None:
_set(
Phase.PREFILL,
"max_context_size",
cfg.cuda_graph_prefill_max_context,
)
# ---- Explicit JSON config (highest precedence) ----
for phase, phase_config in explicit_input.items():
@@ -554,6 +560,63 @@ def validate_cuda_graph_config(server_args: Any):
)
def _resolve_max_context_size(
*, requested_size: Any, page_size: int, model_context_len: Optional[int]
) -> int:
if requested_size <= 0:
raise ValueError("--cuda-graph-prefill-max-context must be a positive integer")
aligned_size = int((requested_size + page_size - 1) // page_size * page_size)
if (
model_context_len is not None
and model_context_len > 0
and aligned_size > model_context_len
):
raise ValueError(
"--cuda-graph-prefill-max-context exceeds the model context length: "
f"aligned size {aligned_size} > {model_context_len}"
)
if requested_size != aligned_size:
logger.info(
"Page-aligning prefill CUDA graph max context size %d -> %d "
"(page_size=%d).",
requested_size,
aligned_size,
page_size,
)
return aligned_size
def finalize_cuda_graph_prefill_max_context(server_args: Any) -> None:
cfg = resolving_view(server_args)
requested_size = cfg.cuda_graph_config.prefill.max_context_size
if requested_size is None:
return
page_size = cfg.page_size
assert page_size is not None and page_size > 0, (
"page_size must be resolved before prefill CUDA graph max context size"
)
max_context_size = _resolve_max_context_size(
requested_size=requested_size,
page_size=page_size,
model_context_len=model_config_of(server_args).context_len,
)
logger.info(
"Prefill CUDA graph max context size: %d; graph keys remain token-only.",
max_context_size,
)
declare_resolution(
server_args,
"_finalize_cuda_graph_prefill_max_context",
cuda_graph_config=with_phase(
cfg.cuda_graph_config,
Phase.PREFILL,
max_context_size=max_context_size,
),
)
def generate_prefill_cuda_graph_batch_sizes(max_bs: int):
"""
Generate the list of batch sizes for prefill CUDA graph capture
python/sglang/srt/arg_groups/field_order.py
+1
View File
@@ -230,6 +230,7 @@ POSITIONAL_FIELD_ORDER = (
"cuda_graph_max_bs_prefill",
"cuda_graph_bs_decode",
"cuda_graph_bs_prefill",
"cuda_graph_prefill_max_context",
"cuda_graph_tc_compiler",
"disable_prefill_cuda_graph",
"disable_decode_cuda_graph",
python/sglang/srt/arg_groups/fields/exec_.py
+15
View File
@@ -37,6 +37,7 @@ from sglang.srt.model_executor.cuda_graph_config import (
CudaGraphConfig,
parse_cuda_graph_config_arg,
)
from sglang.srt.utils.common import human_readable_int
class ExecFeatures(msgspec.Struct):
@@ -493,6 +494,20 @@ class ExecGraph(msgspec.Struct):
Optional[List[int]],
"Explicit list of batch sizes to capture for the prefill cuda graph.",
] = None
cuda_graph_prefill_max_context: A[
Optional[int],
Arg(
help=(
"Maximum context length supported by DeepSeek-V4 breakable/full "
"prefill CUDA graphs. Context-shaped attention metadata and "
"indexer logits are allocated at this fixed size instead of "
"the model maximum. Larger live contexts fall back to eager."
f"\n\n{human_readable_int.__doc__}"
),
type_parser=human_readable_int,
aliases=["--context-bucket"],
),
] = None
cuda_graph_tc_compiler: A[
Optional[Literal["eager", "inductor"]],
"Compiler used by the tc_piecewise backend (currently only the prefill phase consumes it).",
python/sglang/srt/arg_groups/pipeline.py
+3
View File
@@ -190,6 +190,7 @@ def run_resolution_pipeline(server_args: Any) -> None:
apply_inkling_prefill_cuda_graph_default,
apply_muse_glimmer_prefill_cuda_graph_max_bs_default,
disable_prefill_cuda_graph_for_deepseek_trtllm_mla,
finalize_cuda_graph_prefill_max_context,
handle_cuda_graph_config,
)
@@ -369,6 +370,8 @@ def run_resolution_pipeline(server_args: Any) -> None:
# time; last declarations of the resolution, mirroring that order.
run_hook(handle_model_capability_adjustments, server_args)
finalize_cuda_graph_prefill_max_context(server_args)
# Validate after all batch-size declarations are visible.
run_hook(validate_deepep_v2_speculative_draft, server_args)
run_hook(validate_deepep_v2_dispatch_token_budget, server_args)
python/sglang/srt/layers/attention/base_attn_backend.py
+3
View File
@@ -153,6 +153,9 @@ class AttentionBackend(ABC):
# object during capture, and refresh its dynamic fields before each replay.
use_captured_forward_metadata_for_breakable_cuda_graph: bool = False
# True when prefill graph metadata can use ForwardBatch.max_seq_len_override.
supports_prefill_cuda_graph_max_context_size: bool = False
def shared_read_ends(self, fm: ForwardMode) -> SharedReadEnds:
"""Declare where this backend's scheduler-shared reads end per mode.
Override only for audited deviations from this conservative default."""
python/sglang/srt/layers/attention/deepseek_v4_backend.py
+19 -4
View File
@@ -587,6 +587,7 @@ class DeepseekV4AttnBackend(
AttentionBackend, C4IndexerBackendMixin, CompressorBackendMixin
):
use_captured_forward_metadata_for_breakable_cuda_graph: bool = True
supports_prefill_cuda_graph_max_context_size: bool = True
supports_ragged_verify_graph: bool = True
needs_cpu_seq_lens: bool = False
trtllm_attn: bool = False
@@ -1514,9 +1515,16 @@ class DeepseekV4AttnBackend(
assert self.swa_page_size % SWA_WINDOW == 0 and self.page_size % 128 == 0
if max_seq_len_override is None:
max_seq_len_override = getattr(forward_batch, "max_seq_len_override", None)
max_seq_len_override = forward_batch.max_seq_len_override
if max_seq_len_override is not None:
max_seq_len = max_seq_len_override
if seq_lens_cpu is not None and len(seq_lens_cpu) > 0:
actual_max_seq_len = int(seq_lens_cpu.max().item())
if actual_max_seq_len > max_seq_len:
raise ValueError(
"Prefill CUDA graph max context size is smaller than the "
f"live context: {max_seq_len=} < {actual_max_seq_len=}"
)
elif seq_lens_cpu is not None:
max_seq_len = int(seq_lens_cpu.max().item())
else:
@@ -1607,9 +1615,10 @@ class DeepseekV4AttnBackend(
def init_forward_metadata_for_breakable_cuda_graph_capture(
self, forward_batch: ForwardBatch
):
max_seq_len = forward_batch.max_seq_len_override or self.MAX_SEQ_LEN_FOR_CAPTURE
self.forward_metadata = self._build_forward_metadata(
forward_batch,
max_seq_len_override=self.MAX_SEQ_LEN_FOR_CAPTURE,
max_seq_len_override=max_seq_len,
use_prefill_cuda_graph=True,
)
return self.forward_metadata
@@ -1624,9 +1633,15 @@ class DeepseekV4AttnBackend(
# Build graph-compatible metadata against the padded static batch. The
# batch still carries live seq/extend lens, so the online c128 prefill
# plan remains batch-specific without constructing a second metadata set.
metadata_batch = (
static_forward_batch if static_forward_batch is not None else forward_batch
)
max_seq_len = (
metadata_batch.max_seq_len_override or self.MAX_SEQ_LEN_FOR_CAPTURE
)
static_metadata = self._build_forward_metadata(
static_forward_batch if static_forward_batch is not None else forward_batch,
max_seq_len_override=self.MAX_SEQ_LEN_FOR_CAPTURE,
metadata_batch,
max_seq_len_override=max_seq_len,
use_prefill_cuda_graph=True,
)
assert isinstance(capture_metadata, DSV4Metadata)
python/sglang/srt/layers/attention/hybrid_attn_backend.py
+4
View File
@@ -93,6 +93,10 @@ class HybridAttnBackend(AttentionBackend):
def supports_full_cuda_graph_chunked_prefix(self) -> bool:
return self.prefill_backend.supports_full_cuda_graph_chunked_prefix
@property
def supports_prefill_cuda_graph_max_context_size(self) -> bool:
return self.prefill_backend.supports_prefill_cuda_graph_max_context_size
def prepare_full_cuda_graph_chunked_prefix(
self,
forward_batch: ForwardBatch,
python/sglang/srt/layers/attention/tbo_backend.py
+5
View File
@@ -28,6 +28,10 @@ class TboAttnBackend(AttentionBackend):
primary, "extend_dummy_seqs_capped_by_req_pool", False
)
@property
def supports_prefill_cuda_graph_max_context_size(self) -> bool:
return self.primary.supports_prefill_cuda_graph_max_context_size
@classmethod
def init_new(cls, creator: Callable[[], AttentionBackend]):
return cls(
@@ -259,4 +263,5 @@ def _build_tbo_child_replay_fb_view(
else None
),
spec_info=child_spec_info,
max_seq_len_override=fb_view.max_seq_len_override,
)
python/sglang/srt/managers/scheduler_components/dp_attn.py
+7
View File
@@ -350,6 +350,13 @@ def _local_prefill_cuda_graph_vote(
capture_hidden_mode=None,
return_logprob=return_logprob,
lora_ineligible=prefill_graph_runner.enable_lora,
batch_max_context_len=(
int(local_batch.seq_lens_cpu.max().item())
if prefill_graph_runner.max_context_size is not None
and local_batch.seq_lens_cpu is not None
and local_batch.seq_lens_cpu.numel() > 0
else None
),
)
python/sglang/srt/model_executor/cuda_graph_config.py
+7 -1
View File
@@ -72,7 +72,8 @@ ALLOWED_BACKENDS_PER_PHASE = {
# For prefill, bs carries aggregate-token capture buckets for every backend;
# full_prefill_max_req separately controls Full's fixed request-slot count.
# full_prefill_max_req and full_prefill_prefix_chunk_tokens are prefill-only and
# only meaningful when backend == full.
# only meaningful when backend == full. max_context_size is shared by the
# breakable and full prefill body-capture backends.
ALLOWED_KEYS_PER_PHASE = {
Phase.DECODE: ("backend", "max_bs", "bs", "tc_compiler"),
Phase.PREFILL: (
@@ -80,6 +81,7 @@ ALLOWED_KEYS_PER_PHASE = {
"max_bs",
"bs",
"tc_compiler",
"max_context_size",
"full_prefill_max_req",
"full_prefill_prefix_chunk_tokens",
),
@@ -95,6 +97,10 @@ class PhaseConfig:
bs: Optional[List[int]] = None
# Only meaningful when backend == tc_piecewise; ignored otherwise.
tc_compiler: str = "eager"
# Effective for both full and breakable backends and currently only DSV4:
# fixed maximum context length used by context-shaped prefill graph metadata.
# Every token bucket shares this size; larger live contexts run eagerly.
max_context_size: Optional[int] = None
# Only meaningful for the prefill phase with backend == full: max number of
# request slots baked into each captured graph. Real bs <= full_prefill_max_req
# reuses the graph (unused slots become zero-length sentinels); larger
python/sglang/srt/model_executor/forward_batch_info.py
+5
View File
@@ -573,6 +573,11 @@ class ForwardBatch(ForwardBatchDeepSeekMHAMixin):
# Preallocated piecewise-graph attention output, set by RadixAttention.
_attn_output: Optional[torch.Tensor] = None
# Prefill body-CUDA-graph context limit. Attention backends that allocate
# context-shaped metadata use this fixed maximum instead of deriving a
# shape from the live batch. None preserves eager/default graph behavior.
max_seq_len_override: Optional[int] = None
# For logits and logprobs post processing
next_token_logits_buffer: torch.Tensor = None
temperature: torch.Tensor = None
python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py
+1
View File
@@ -197,6 +197,7 @@ def build_replay_fb_view(
out_cache_loc=getattr(forward_batch, "out_cache_loc", None),
out_cache_loc_virtual=forward_batch.out_cache_loc_virtual,
out_cache_loc_dsv4=getattr(forward_batch, "out_cache_loc_dsv4", None),
max_seq_len_override=forward_batch.max_seq_len_override,
# The mamba-track registry slot (VIRTUAL ids) is the v2p translate SOURCE
# for the backend, which copies the result into its own static buffer and
# reads THAT in the decode track-save — this slot is never mutated. None
python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py
+120 -14
View File
@@ -312,6 +312,18 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# --- runner bounds --------------------------------------------
self.max_num_tokens = max(self.capture_num_tokens)
self.max_bs = model_runner.req_to_token_pool.size
self.max_context_size = prefill_config.max_context_size
self._validate_max_context_capacity(
max_context_size=self.max_context_size,
table_width=model_runner.req_to_token_pool.req_to_token.shape[1],
)
if (
self.prefill_backend_name == Backend.TC_PIECEWISE
and self.max_context_size is not None
):
# TODO(SYChen123): Plumb max_seq_len_override through TcPiecewise
# metadata preparation before enabling the fixed context limit here.
self._ignore_max_context_size("tc_piecewise prefill CUDA graph")
# --- capture modes --------------------------------------------
self.capture_forward_mode = ForwardMode.EXTEND
@@ -434,6 +446,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
max_req = prefill_config.full_prefill_max_req
assert max_req is not None, "full_prefill_max_req must be resolved"
self._capture_req_slots = max_req
# BCG/Full record LoRA kernels, so the metadata they read must live in
# static buffers refreshed in place per batch; unsupported LoRA
# configs were already routed to the eager runner.
@@ -510,10 +523,22 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self.backend, BreakableCudaGraphBackend
) and should_enable_cp_bcg_capture(server_args)
if self.enable_cp_bcg_capture:
if self.max_context_size is not None:
# TODO(SYChen123): Preserve max_seq_len_override through CP's padded
# metadata preparation before enabling the fixed context limit.
self._ignore_max_context_size("CP breakable prefill CUDA graph")
self.capture_num_tokens = filter_prefill_cp_bcg_capture_num_tokens(
self.capture_num_tokens, server_args
)
self.prefill_cp_bcg_input = PrefillCPBCGInput.create(self)
if self.max_context_size is not None and not (
model_runner.attn_backend.supports_prefill_cuda_graph_max_context_size
):
raise ValueError(
"--cuda-graph-prefill-max-context is only supported by attention "
"backends that implement fixed-context prefill graph metadata; "
f"got {type(model_runner.attn_backend).__name__}"
)
# Static hidden_states buffer giving the captured graph a stable
# address; load_batch refreshes it from live spec_info at replay.
@@ -541,7 +566,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self.use_captured_attn_metadata = model_runner.attn_backend.use_captured_forward_metadata_for_breakable_cuda_graph
else:
self.use_captured_attn_metadata = False
self.attn_metadata_buffers: Optional[Dict[int, object]] = (
self.attn_metadata_buffers: Optional[Dict[ShapeKey, object]] = (
{} if self.use_captured_attn_metadata else None
)
@@ -862,6 +887,41 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
else table_width
)
@staticmethod
def _validate_max_context_capacity(
*, max_context_size: Optional[int], table_width: int
) -> None:
if max_context_size is not None and max_context_size > table_width:
raise ValueError(
"--cuda-graph-prefill-max-context exceeds the request-to-token "
f"pool capacity: requested size {max_context_size} > {table_width}"
)
def _ignore_max_context_size(self, execution_path: str) -> None:
if self.max_context_size is None:
return
logger.warning(
"Ignoring prefill CUDA graph max context size %d for %s; this path "
"does not yet preserve the fixed metadata extent.",
self.max_context_size,
execution_path,
)
self.max_context_size = None
@staticmethod
def _batch_max_context_len(forward_batch: ForwardBatch) -> Optional[int]:
seq_lens_cpu = forward_batch.seq_lens_cpu
if seq_lens_cpu is not None:
if torch.is_tensor(seq_lens_cpu):
return int(seq_lens_cpu.max().item()) if seq_lens_cpu.numel() > 0 else 0
return max((int(value) for value in seq_lens_cpu), default=0)
seq_lens = forward_batch.seq_lens
if seq_lens is None or seq_lens.numel() == 0:
return None
# Prefill normally has seq_lens_cpu. This fallback is for hand-built
# batches and intentionally synchronizes rather than guessing a bucket.
return int(seq_lens.max().item())
@staticmethod
def _resolve_prefix_chunk_shape(
model_runner, capture_req_slots: int
@@ -1054,7 +1114,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
)
def _init_forward_metadata_for_capture(
self, forward_batch: ForwardBatch, num_tokens: int
self, forward_batch: ForwardBatch, shape_key: ShapeKey
) -> None:
"""Capture-time metadata init for the BCG-with-captured-metadata
contract. For opt-in backends (DSV4), call the BCG-specific entry
@@ -1071,13 +1131,13 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
)
)
assert self.attn_metadata_buffers is not None
self.attn_metadata_buffers[num_tokens] = metadata
self.attn_metadata_buffers[shape_key] = metadata
def _prepare_forward_metadata_for_replay(
self,
forward_batch: ForwardBatch,
static_forward_batch: ForwardBatch,
num_tokens: int,
shape_key: ShapeKey,
) -> None:
"""Replay-time metadata refresh for the BCG-with-captured-metadata
contract. For opt-in backends, refresh the stashed per-bucket
@@ -1103,16 +1163,17 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
padded_view.req_pool_indices = s["req_pool_indices"][:r]
padded_view.extend_seq_lens = s["extend_seq_lens"][:r]
padded_view.extend_prefix_lens = s["extend_prefix_lens"][:r]
padded_view.max_seq_len_override = static_forward_batch.max_seq_len_override
attn_backend.init_forward_metadata_out_graph(padded_view)
return
if not self.use_captured_attn_metadata:
attn_backend.init_forward_metadata(forward_batch)
attn_backend.prepare_prefill_shared_read_snapshot(
forward_batch, num_qo_tokens=num_tokens
forward_batch, num_qo_tokens=shape_key.size
)
return
assert self.attn_metadata_buffers is not None
metadata = self.attn_metadata_buffers[num_tokens]
metadata = self.attn_metadata_buffers[shape_key]
attn_backend.prepare_forward_metadata_for_breakable_cuda_graph_replay(
metadata,
forward_batch,
@@ -1138,6 +1199,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
capture_hidden_mode,
return_logprob: bool,
lora_ineligible: bool = False,
batch_max_context_len: Optional[int] = None,
) -> bool:
"""Rank-local replay eligibility: the single source of truth for
``can_run_graph`` (ForwardBatch, forward time) and the dp mlp-sync
@@ -1180,6 +1242,12 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
return False
if return_logprob and not self._uses_eager_prefill_tail():
return False
if self.max_context_size is not None:
if (
batch_max_context_len is None
or batch_max_context_len > self.max_context_size
):
return False
if num_tokens is None:
return True
if num_tokens > self.max_num_tokens:
@@ -1206,6 +1274,11 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
return False
# Non-DP local check (sole decision for tp-only).
batch_max_context_len = (
self._batch_max_context_len(forward_batch)
if self.max_context_size is not None
else None
)
if not self.can_replay_locally(
batch_size=forward_batch.batch_size,
num_tokens=len(forward_batch.input_ids),
@@ -1222,6 +1295,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
forward_batch
)
),
batch_max_context_len=batch_max_context_len,
):
return False
if getattr(self, "enable_cp_bcg_capture", False) and is_cp_active(
@@ -1263,7 +1337,10 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
Returns ``(forward_batch, attn_backend)`` to mirror decode's
capture_prepare signature.
"""
context_length = self.model_runner.model_config.context_len
model_context_length = self.model_runner.model_config.context_len
context_length = min(
self.max_context_size or model_context_length, model_context_length
)
# A prefill bucket is an aggregate token count. Capture it as the
# fewest synthetic requests, with every request containing no more
# than context_length tokens.
@@ -1398,6 +1475,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# refreshes the static batch info with live values.
lora_ids=([None] * bs if self._capture_lora else None),
return_pooled_hidden_states=self.capture_return_pooled_hidden_states,
max_seq_len_override=self.max_context_size,
)
self.tbo_plugin.capture_one_batch_size(forward_batch, num_tokens=num_tokens)
return forward_batch, self.model_runner.attn_backend
@@ -1430,10 +1508,16 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self.model_runner.gpu_id,
empty_cache=False,
)
capture_shapes = list(reversed(self.capture_num_tokens))
if self.max_context_size is not None:
logger.info(
"Capturing %d prefill CUDA graph token shapes with fixed "
"max_context_size=%d (before execution variants).",
len(capture_shapes),
self.max_context_size,
)
capture_range = (
tqdm.tqdm(list(reversed(self.capture_num_tokens)))
if get_parallel().tp_rank == 0
else reversed(self.capture_num_tokens)
tqdm.tqdm(capture_shapes) if get_parallel().tp_rank == 0 else capture_shapes
)
for num_tokens in capture_range:
if get_parallel().tp_rank == 0:
@@ -1443,12 +1527,15 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
empty_cache=False,
)
capture_range.set_description(
f"Capturing num tokens ({num_tokens=} {avail_mem=:.2f} GB)"
f"Capturing prefill shape ({num_tokens=} {avail_mem=:.2f} GB)"
)
self.capture_one_shape(num_tokens)
if self._capture_chunked_prefix:
for captured_n in self._prefix_capture_variants:
self.capture_one_shape(num_tokens, prefix_num_chunks=captured_n)
self.capture_one_shape(
num_tokens,
prefix_num_chunks=captured_n,
)
def capture_one_shape(self, size: int, *, prefix_num_chunks: int = 0) -> None:
"""Per-shape capture: build dummy ForwardBatch + run_once,
@@ -1496,7 +1583,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
# Reaching into a backend-specific metadata cache here would make
# this path incompatible with the OSS FlashAttention backend.
else:
self._init_forward_metadata_for_capture(forward_batch, num_tokens)
self._init_forward_metadata_for_capture(forward_batch, shape_key)
def run_once():
# Record LoRA kernels even when capture uses base-model requests.
@@ -1569,6 +1656,17 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
)
self.raw_num_tokens = num_tokens
if self.max_context_size is not None:
batch_max_context_len = self._batch_max_context_len(forward_batch)
if (
batch_max_context_len is None
or batch_max_context_len > self.max_context_size
):
raise RuntimeError(
"Prefill CUDA graph replay was admitted without a fitting "
"maximum context size"
)
bs = forward_batch.batch_size
self.raw_bs = bs
@@ -1707,6 +1805,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
self.capture_return_pooled_hidden_states
or forward_batch.return_pooled_hidden_states
),
max_seq_len_override=self.max_context_size,
)
if self._is_full_backend:
forward_batch.next_token_logits_buffer = (
@@ -1775,8 +1874,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
)
metadata_forward_batch = static_forward_batch
shape_key = self._shape_key(static_num_tokens, forward_batch)
self._prepare_forward_metadata_for_replay(
metadata_forward_batch, static_forward_batch, static_num_tokens
metadata_forward_batch, static_forward_batch, shape_key
)
return static_forward_batch
@@ -1860,6 +1960,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
raw_num_tokens: int,
**kwargs,
):
assert self.max_context_size is None, (
"tc_piecewise replay does not support a fixed prefill context size"
)
with self._prefill_forward_context(
static_forward_batch,
num_tokens=static_num_tokens,
@@ -1937,6 +2040,9 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
)
if self.enable_cp_bcg_capture:
assert self.max_context_size is None, (
"CP-v2 BCG replay does not support a fixed prefill context size"
)
output = execute_prefill_cp_bcg(
self,
forward_batch,