# Copyright (c) 2026 LightSeek Foundation # # Permission is hereby granted, free of charge, to any person obtaining a copy # of this software and associated documentation files (the "Software"), to deal # in the Software without restriction, including without limitation the rights # to use, copy, modify, merge, publish, distribute, sublicense, and/or sell # copies of the Software, and to permit persons to whom the Software is # furnished to do so, subject to the following conditions: # # The above copyright notice and this permission notice shall be included in # all copies or substantial portions of the Software. # # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR # IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, # FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE # AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE # SOFTWARE. from __future__ import annotations import bisect import gc import queue from collections.abc import Callable from contextlib import contextmanager from typing import TYPE_CHECKING import torch import torch.distributed as dist import tqdm from tokenspeed.runtime.configs.paged_cache_spec import ( compute_max_logical_pages_for_capture, ) from tokenspeed.runtime.execution.context import ForwardContext from tokenspeed.runtime.execution.forward_batch_info import ( CaptureHiddenMode, ForwardMode, ) from tokenspeed.runtime.layers.attention.backends.base import ( init_backend_cuda_graph_state, ) from tokenspeed.runtime.sampling.backends.base import CUDA_GRAPH_VARIANT_DEFAULT from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo from tokenspeed.runtime.utils import ( get_available_gpu_memory, get_colorful_logger, ) from tokenspeed.runtime.utils.nvtx import nvtx_range if TYPE_CHECKING: from tokenspeed.runtime.execution.drafter.base import BaseDrafter from tokenspeed.runtime.execution.input_buffer import InputBuffers from tokenspeed.runtime.execution.model_executor import ModelExecutorConfig from tokenspeed.runtime.execution.runtime_states import RuntimeStates from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool from tokenspeed.runtime.sampling.backends.base import SamplingBackend logger = get_colorful_logger(__name__) _is_capture_mode = False def get_is_capture_mode() -> bool: return _is_capture_mode def _should_update_mamba_state_after_mtp_verify( drafter, attn_backend, forward_mode: ForwardMode ) -> bool: return ( drafter is not None and forward_mode.is_decode() and hasattr(attn_backend, "update_mamba_state_after_mtp_verify") ) @contextmanager def freeze_gc(enable_cudagraph_gc: bool): """ Optimize garbage collection during CUDA graph capture. Clean up, then freeze all remaining objects from being included in future collections if GC is disabled during capture. """ gc.collect() should_freeze = not enable_cudagraph_gc if should_freeze: gc.freeze() try: yield finally: if should_freeze: gc.unfreeze() gc.collect() def get_batch_sizes_to_capture(config: ModelExecutorConfig): capture_bs = config.cudagraph_capture_sizes max_bs = config.max_num_seqs // max(config.data_parallel_size, 1) if capture_bs is None: if config.disable_cuda_graph_padding: capture_bs = list(range(1, 33)) + [64, 96, 128, 160] else: capture_bs = [1, 2, 4] + [i * 8 for i in range(1, 21)] if max(capture_bs) > max_bs: capture_bs = list(sorted(set(capture_bs + [max_bs - 1] + [max_bs]))) effective_max = min(config.max_cudagraph_capture_size, max_bs) capture_bs = [bs for bs in capture_bs if 0 < bs <= effective_max] return capture_bs global_graph_memory_pool = None class DeepEPCudaGraphRunnerAdapter: """Manages DeepEP dispatch mode consistency across CUDA graph capture/replay. During capture the forward pass (including DeepEP low-latency RDMA dispatch/combine) is recorded. On replay the Python wrapper code that normally sets dispatch mode and manages the RDMA workspace never re-executes. This adapter restores both before each replay. Follows the same CUDA graph replay contract as the upstream DeepEP runner. """ def __init__(self): self._active = False @staticmethod def _get_buffer_cls(): try: from tokenspeed_kernel.ops.communication.deep_ep import ( DeepEPBuffer, ) return DeepEPBuffer except ImportError: return None def capture(self): """Call before ``torch.cuda.graph()`` capture.""" cls = self._get_buffer_cls() if cls is None or cls._buffer is None: return self._active = True cls.set_dispatch_mode_as_low_latency() def replay(self): """Call before every ``graph.replay()``; restores dispatch mode and resets RDMA workspace so stale sync state doesn't corrupt the combine kernel across replays.""" if not self._active: return cls = self._get_buffer_cls() if cls is None or cls._buffer is None: return cls.set_dispatch_mode_as_low_latency() cls.clean_buffer() class CudaGraphWrapper: """ Wraps a forward_func and transparently dispatches to either a captured CUDA graph (decode, supported batch size) or the eager path (prefill / unsupported batch size). Callers always use the same interface:: output_tokens, output_lengths, output_logprobs = runner( bs, ctx, sampling_info, req_to_page, extend_with_prefix=..., extend_prefix_lens=..., ) Internally the wrapper owns both paths and calls init_forward_metadata with use_cuda_graph=True/False to select the appropriate backend buffers. """ def __init__( self, forward_func: Callable, attn_backend: AttentionBackend, token_to_kv_pool: BaseTokenToKVPool, input_buffers: InputBuffers, config: ModelExecutorConfig, draft_attn_backend: AttentionBackend | None = None, draft_token_to_kv_pool: BaseTokenToKVPool | None = None, drafter: BaseDrafter | None = None, capturable_grammar=None, eager_grammar_buffers=None, sampling_backend: SamplingBackend | None = None, runtime_states: RuntimeStates | None = None, ): self.config = config self.attn_backend = attn_backend self.draft_attn_backend = draft_attn_backend self.draft_token_to_kv_pool = draft_token_to_kv_pool self.token_to_kv_pool = token_to_kv_pool self.drafter = drafter self.sampling_backend = sampling_backend self.input_buffers = input_buffers self.capturable_grammar = capturable_grammar self.eager_grammar_buffers = eager_grammar_buffers self.runtime_states = runtime_states self.enable_torch_compile = getattr(config, "enable_torch_compile", False) self.disable_padding = config.disable_cuda_graph_padding self.enable_cudagraph_gc = getattr(config, "enable_cudagraph_gc", True) self.device = config.device self.gpu_id = config.gpu_id self.global_rank = config.global_rank self.context_len = config.context_len self.vocab_size = config.vocab_size self.grammar_backend = config.grammar_backend self.capture_bs = get_batch_sizes_to_capture(config) self.max_bs = max(self.capture_bs) self.max_tokens_per_req = ( config.spec_num_tokens if config.spec_algo is not None else 1 ) self.overlap_schedule_depth = config.overlap_schedule_depth self.use_v4_mtp_paged_metadata = config.use_v4_mtp_paged_metadata self.dp_size = config.data_parallel_size self.world_size = config.world_size # Backends alias their cache_seqlens buffer. Draft backend aliases # the drafter-owned draft_seq_lens to keep InputBuffers read-only. init_backend_cuda_graph_state( attn_backend, self.max_bs, self.input_buffers.seq_lens_buf, paged_cache_group_specs=tuple(token_to_kv_pool.paged_cache_group_specs), max_tokens_per_req=self.max_tokens_per_req, overlap_schedule_depth=self.overlap_schedule_depth, ) if draft_attn_backend is not None: init_backend_cuda_graph_state( draft_attn_backend, self.max_bs, self.drafter.draft_seq_lens_buf, paged_cache_group_specs=tuple( draft_token_to_kv_pool.paged_cache_group_specs ), max_tokens_per_req=self.max_tokens_per_req, overlap_schedule_depth=self.overlap_schedule_depth, ) # Drafter (Eagle) is constructed with the target's req_to_page # (ModelExecutor passes the same self.req_to_page to both), and the # replay path hands both backends the same req_pool_indices. The # block-table gather is req_to_page[req_pool_indices] (see # _create_block_kv_indices; it does not depend on seq_lens), so both # backends would compute identical block_kv_indices. When the backing # buffer shapes/dtypes also line up, point the draft backend at the # target's buffer and skip its gather+copy in the replay path: the # target's metadata prep runs first and populates the shared buffer # (see init_forward_metadata_replay_cuda_graph). target_kv = getattr(attn_backend, "decode_cuda_graph_kv_indices", None) draft_kv = getattr(draft_attn_backend, "decode_cuda_graph_kv_indices", None) if ( target_kv is not None and draft_kv is not None and target_kv.shape == draft_kv.shape and target_kv.dtype == draft_kv.dtype ): draft_attn_backend.decode_cuda_graph_kv_indices = target_kv draft_attn_backend._block_table_aliased = True self.graph_variants = ( sampling_backend.cuda_graph_capture_variants(self.max_tokens_per_req) if sampling_backend is not None else (CUDA_GRAPH_VARIANT_DEFAULT,) ) self.graphs: dict[tuple[str, int], torch.cuda.CUDAGraph] = {} self.output_buffers: dict[tuple[str, int], tuple] = {} self._forward_func: Callable | None = forward_func self.disable = config.enforce_eager self.deepep_adapter = DeepEPCudaGraphRunnerAdapter() if not self.disable: self.capture() # ------------------------------------------------------------------ # Graph capture # ------------------------------------------------------------------ def capture(self): """ Capture CUDA graphs for all configured batch sizes. Args: forward_func: ModelExecutor.forward_step(bs, ctx, sampling_info). """ rank = self.global_rank with freeze_gc(self.enable_cudagraph_gc): self.stream = torch.cuda.Stream() # Capture backend-declared sampler variants explicitly. capture_items = [ (variant, bs) for variant in self._cuda_graph_capture_variants() for bs in self.capture_bs ] capture_range = tqdm.tqdm(capture_items) if rank == 0 else capture_items if rank == 0: logger.info("Capturing batches: %s", self.capture_bs) for variant, bs in capture_range: if rank == 0: avail_mem = get_available_gpu_memory( self.device, self.gpu_id, empty_cache=False ) variant_desc = ( "" if variant == CUDA_GRAPH_VARIANT_DEFAULT else f" variant={variant}" ) capture_range.set_description( f"Capturing batches ({bs=}{variant_desc} {avail_mem=:.2f} GB)" ) graph, output_buffers = self._capture_one(bs, variant=variant) self.graphs[(variant, bs)] = graph self.output_buffers[(variant, bs)] = output_buffers def _cuda_graph_capture_variants(self) -> tuple[str, ...]: if self.sampling_backend is None: return (CUDA_GRAPH_VARIANT_DEFAULT,) variants = self.sampling_backend.cuda_graph_capture_variants( self.max_tokens_per_req ) if not variants: return (CUDA_GRAPH_VARIANT_DEFAULT,) deduped = tuple(dict.fromkeys((CUDA_GRAPH_VARIANT_DEFAULT, *variants))) return deduped def _prepare_sampling_capture(self, bs: int, variant: str) -> None: if self.sampling_backend is None: return self.sampling_backend.prepare_capture_variant( bs=bs, num_tokens_per_req=self.max_tokens_per_req, variant=variant, ) def _cuda_graph_replay_variant(self) -> str: if self.sampling_backend is None: return CUDA_GRAPH_VARIANT_DEFAULT return self.sampling_backend.cuda_graph_replay_variant(self.max_tokens_per_req) def _cuda_graph_key(self, bs: int) -> tuple[str, int]: variant = self._cuda_graph_replay_variant() key = (variant, bs) if key in self.graphs: return key if variant != CUDA_GRAPH_VARIANT_DEFAULT: captured_variants = sorted( graph_variant for graph_variant, graph_bs in self.graphs if graph_bs == bs ) raise RuntimeError( "Sampling backend requested CUDA graph variant " f"{variant!r} for batch size {bs}, but it was not captured. " f"Captured variants for this batch size: {captured_variants}." ) return (CUDA_GRAPH_VARIANT_DEFAULT, bs) def _has_cuda_graph_for_bs(self, bs: int) -> bool: return (CUDA_GRAPH_VARIANT_DEFAULT, bs) in self.graphs def _capture_one(self, bs: int, variant: str = CUDA_GRAPH_VARIANT_DEFAULT): graph = torch.cuda.CUDAGraph() capture_forward_mode = ForwardMode.DECODE ctx = ForwardContext( attn_backend=self.attn_backend, token_to_kv_pool=self.token_to_kv_pool, bs=bs, num_extends=0, input_num_tokens=bs * self.max_tokens_per_req, forward_mode=capture_forward_mode, capture_hidden_mode=( CaptureHiddenMode.FULL if self.drafter is not None else CaptureHiddenMode.NULL ), ) # For DP mode, global_num_tokens must be set so that the MoE # all-gather comm layers know token counts for all DP ranks. # During capture, use uniform dummy counts across ranks. if self.dp_size > 1: ctx.global_num_tokens = [bs * self.max_tokens_per_req] * self.world_size # global_bs must ALSO be set at capture. The draft first step's # collective sizing (reported via report_collective_sizing) reads # global_bs; if left None at capture it records a single-rank # layout (fallback branch in comm_manager), but at replay global_bs # is the live per-rank batch list -> multi-rank layout. The mismatch # makes the captured (frozen-offset) gather read uninitialized # symm-mem -> NaN draft logits -> accept_rate 0. Set the matching # uniform dummy. ctx.global_bs = [bs] * self.world_size # Capture with is_all_greedy=False so the graph records the full # top_k_top_p_sampling path (greedy-only requests are served by the # same path with top_k=1 in the buffer, which effectively argmaxes). # is_all_greedy=True at capture would freeze the graph into # argmax and bypass per-request seeding at replay. ibd = self.input_buffers sampling_info = SamplingBatchInfo( req_pool_indices=ibd.req_pool_indices_buf[:bs], valid_cache_lengths=( self.runtime_states.valid_cache_lengths if self.runtime_states is not None else None ), is_all_greedy=False, vocab_size=self.vocab_size, device=self.device, ) from tokenspeed.runtime.grammar.capturable_grammar import ( bind_grammar_mask_buf, ) # Bind whichever grammar buffer is active so the captured sampler # records the apply_vocab_mask call. At replay, runtime fills the # bound buffer in place (hostfunc for capturable, sync H2D for # eager) — the captured graph reads from the same memory. bind_grammar_mask_buf( sampling_info, self.eager_grammar_buffers, bs, spec=self.drafter is not None, capturable=self.capturable_grammar, grammar_backend=self.grammar_backend, ) def run_once(): # Dummy add_batch keeps the grammar queue 1:1 with replays — # fetch_batch pops once per forward, so warmup + capture # would otherwise raise queue.Empty. if self.capturable_grammar is not None: self.capturable_grammar.add_batch( grammars=[None] * bs, bs=bs, has_candidates=False ) return self._forward_func(bs=bs, ctx=ctx, sampling_info=sampling_info) # Warm up before capture. for _ in range(4): torch.cuda.synchronize() dist.barrier() self._prepare_sampling_capture(bs=bs, variant=variant) # Keep warmup seq_lens >= q_len_per_req so no query row gets an # empty causal span; a stale seq_len of 1 overflows to non-finite KV. self.input_buffers.seq_lens_buf[:bs].fill_(self.max_tokens_per_req) self._init_capture_metadata(bs) run_once() # Clear any per-pool state that warm-up dirtied at pool row 0, # so the graph captures reads against a clean baseline. if self.sampling_backend is not None: self.sampling_backend.reset_capture_state() torch.cuda.synchronize() dist.barrier() # Warmups can switch a backend back to eager metadata objects. Restore # the graph-backed metadata immediately before capture so replay-time # metadata refreshes update the same tensors recorded by the graph. self._init_capture_metadata(bs) # Fill sampler buffers OUTSIDE the capture so RNG ops aren't recorded. self._prepare_sampling_capture(bs=bs, variant=variant) # Warmup forwards can mutate aliased metadata buffers, so refresh # them again immediately before graph capture records the final views. self._init_capture_metadata(bs) self.deepep_adapter.capture() global _is_capture_mode _is_capture_mode = True global global_graph_memory_pool with torch.cuda.graph(graph, pool=global_graph_memory_pool, stream=self.stream): out = run_once() torch.cuda.synchronize() dist.barrier() _is_capture_mode = False # Graph capture records the hostfunc launches without invoking # them, so the dummy run_once pushed stays queued — drain it, and # reset prev_batch/current_batch so the first real replay's build # doesn't advance the matcher from a stale warmup entry. if self.capturable_grammar is not None: while True: try: self.capturable_grammar.queue.get_nowait() except queue.Empty: break self.capturable_grammar.reset_state() global_graph_memory_pool = graph.pool() return graph, out def _capture_paged_cache_block_tables(self, bs: int, pool) -> dict | None: specs = tuple(pool.paged_cache_group_specs) if not specs: return None out = {} for spec in specs: max_pages = compute_max_logical_pages_for_capture( spec, max_context_len=( self.max_tokens_per_req * self.max_bs if self.context_len <= 0 else self.context_len ), max_tokens_per_req=self.max_tokens_per_req, overlap_schedule_depth=self.overlap_schedule_depth, ) out[str(spec.group_id)] = torch.zeros( (bs, max_pages), dtype=torch.int32, device=self.device, ) return out def _flat_cache_group_ids(self, pool) -> tuple[str, ...]: """Group ids for flat per-group CUDA-graph capture: real tables only arrive at replay, so capture needs just the ids to allocate its persistent per-group buffers.""" if not getattr(self.attn_backend, "uses_flat_cache_groups", False): return () return tuple(str(spec.group_id) for spec in pool.paged_cache_group_specs) def _draft_flat_group_ids(self) -> tuple[str, ...]: """The draft head shares the target full-attention group's page ids (EAGLE writes its own pool tensors at the same indices), so the drafter consumes exactly that group's table on the flat path.""" if self.draft_attn_backend is None or not getattr( self.draft_attn_backend, "uses_flat_cache_groups", False ): return () return tuple( str(spec.group_id) for spec in self.token_to_kv_pool.paged_cache_group_specs if spec.family != "state" and spec.retention == "full_history" ) def _draft_flat_tables(self, flat_block_tables): """Subset of the target's per-group tables the drafter consumes.""" gids = self._draft_flat_group_ids() if not gids or not flat_block_tables: return None subset = { gid: flat_block_tables[gid] for gid in gids if gid in flat_block_tables } return subset or None def _init_capture_metadata(self, bs: int): capture_kwargs = {} if self.input_buffers.has_mamba: capture_kwargs["mamba_pool_indices"] = ( self.input_buffers.mamba_pool_indices_buf[:bs] ) if self.attn_backend.uses_paged_cache_groups: paged_cache_block_tables = self._capture_paged_cache_block_tables( bs, self.token_to_kv_pool, ) if paged_cache_block_tables is not None: capture_kwargs["paged_cache_block_tables"] = paged_cache_block_tables if self.drafter is not None: capture_kwargs["num_tokens"] = bs * self.max_tokens_per_req flat_cache_group_ids = self._flat_cache_group_ids(self.token_to_kv_pool) if flat_cache_group_ids: capture_kwargs["flat_cache_group_ids"] = flat_cache_group_ids self.attn_backend.init_forward_metadata_capture_cuda_graph( bs, self.input_buffers.req_pool_indices_buf[:bs], self.input_buffers.seq_lens_buf[:bs], ForwardMode.DECODE, **capture_kwargs, ) if self.draft_attn_backend is not None: draft_kwargs = {} if ( self.draft_token_to_kv_pool is not None and self.draft_attn_backend.uses_paged_cache_groups ): draft_paged_cache_block_tables = self._capture_paged_cache_block_tables( bs, self.draft_token_to_kv_pool, ) if draft_paged_cache_block_tables is not None: draft_kwargs["paged_cache_block_tables"] = ( draft_paged_cache_block_tables ) draft_kwargs["num_tokens"] = bs * self.max_tokens_per_req draft_flat_ids = self._draft_flat_group_ids() if draft_flat_ids: draft_kwargs["flat_cache_group_ids"] = draft_flat_ids # Drafter mutates seq_lens_buf in place per step; backends alias. self.draft_attn_backend.init_forward_metadata_capture_cuda_graph( bs, self.input_buffers.req_pool_indices_buf[:bs], self.input_buffers.seq_lens_buf[:bs], ForwardMode.DECODE, **draft_kwargs, ) def _idle_flat_block_tables(self, padded_bs: int) -> dict | None: """Minimal per-group tables for the bs==0 idle replay: all rows are dummy rows, so one column of page-0 entries per group is valid. None when the pool publishes no groups.""" specs = tuple(self.token_to_kv_pool.paged_cache_group_specs) if not specs: return None table = torch.zeros((padded_bs, 1), dtype=torch.int32, device=self.device) return {str(spec.group_id): table for spec in specs} @staticmethod def _pad_block_tables_to_padded_bs( block_tables: dict, *, actual_bs: int, padded_bs: int, pad_value: int = -1, ) -> dict: """Pad each table with dummy ROWS up to padded_bs. Flat passes pad_value=0, radix/V4 keeps -1 — see the padding contract at the MHA backend's replay guard (backends/mha.py). """ if padded_bs <= actual_bs: return block_tables out = {} for key, table in block_tables.items(): if not isinstance(table, torch.Tensor): out[key] = table continue rows = int(table.shape[0]) if rows == padded_bs: out[key] = table continue out[key] = torch.nn.functional.pad( table, (0, 0, 0, padded_bs - rows), value=pad_value, ) return out @staticmethod def _pad_offsets_to_padded_bs( base_offsets: dict, *, actual_bs: int, padded_bs: int, ) -> dict: if padded_bs <= actual_bs: return base_offsets out = {} for key, off in base_offsets.items(): if not isinstance(off, torch.Tensor): out[key] = off continue rows = int(off.shape[0]) if rows == padded_bs: out[key] = off continue # Base 0: padded rows have no real request; the paired padded # table row is invalid (-1). out[key] = torch.nn.functional.pad( off, (0, padded_bs - rows), value=0, ) return out def _init_replay_metadata( self, padded_bs: int, actual_bs: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, req_to_page: torch.Tensor, forward_mode: ForwardMode, **kwargs, ): """Graph-replay path — update persistent cuda-graph buffers in place.""" paged_cache_block_tables = kwargs.pop("paged_cache_block_tables", None) paged_cache_block_table_base_offsets = kwargs.pop( "paged_cache_block_table_base_offsets", None ) flat_block_tables = kwargs.pop("flat_block_tables", None) target_uses_paged_groups = getattr( self.attn_backend, "uses_paged_cache_groups", False, ) draft_uses_paged_groups = self.draft_attn_backend is not None and getattr( self.draft_attn_backend, "uses_paged_cache_groups", False ) if paged_cache_block_tables is not None and ( target_uses_paged_groups or draft_uses_paged_groups ): table_bs = next( ( int(table.shape[0]) for table in paged_cache_block_tables.values() if isinstance(table, torch.Tensor) ), int(req_pool_indices.shape[0]), ) paged_cache_block_tables = self._pad_block_tables_to_padded_bs( paged_cache_block_tables, actual_bs=table_bs, padded_bs=padded_bs, ) if paged_cache_block_table_base_offsets is not None: paged_cache_block_table_base_offsets = self._pad_offsets_to_padded_bs( paged_cache_block_table_base_offsets, actual_bs=actual_bs, padded_bs=padded_bs, ) if target_uses_paged_groups: kwargs["paged_cache_block_tables"] = paged_cache_block_tables if paged_cache_block_table_base_offsets is not None: kwargs["paged_cache_block_table_base_offsets"] = ( paged_cache_block_table_base_offsets ) if flat_block_tables is not None and getattr( self.attn_backend, "uses_flat_cache_groups", False ): flat_table_bs = next( ( int(table.shape[0]) for table in flat_block_tables.values() if isinstance(table, torch.Tensor) ), int(req_pool_indices.shape[0]), ) kwargs["flat_block_tables"] = self._pad_block_tables_to_padded_bs( flat_block_tables, actual_bs=flat_table_bs, padded_bs=padded_bs, pad_value=0, ) if self.attn_backend.uses_padded_decode_token_mask: kwargs["actual_bs"] = actual_bs if target_uses_paged_groups and getattr(self, "drafter", None) is not None: kwargs["num_tokens"] = padded_bs * self.max_tokens_per_req self.attn_backend.init_forward_metadata_replay_cuda_graph( padded_bs, req_pool_indices, seq_lens, req_to_page=req_to_page, forward_mode=forward_mode, **kwargs, ) if self.draft_attn_backend is not None: draft_attn_kwargs = {} if draft_uses_paged_groups and paged_cache_block_tables is not None: draft_attn_kwargs["paged_cache_block_tables"] = paged_cache_block_tables if paged_cache_block_table_base_offsets is not None: draft_attn_kwargs["paged_cache_block_table_base_offsets"] = ( paged_cache_block_table_base_offsets ) if getattr(self.draft_attn_backend, "uses_padded_decode_token_mask", False): draft_attn_kwargs["actual_bs"] = actual_bs draft_flat = self._draft_flat_tables(kwargs.get("flat_block_tables")) if draft_flat is not None: draft_attn_kwargs["flat_block_tables"] = draft_flat draft_forward_mode = ForwardMode.DECODE if draft_uses_paged_groups: draft_attn_kwargs["num_tokens"] = padded_bs * self.max_tokens_per_req draft_seq_lens = self.drafter.draft_seq_lens_buf[:padded_bs] draft_seq_lens.copy_(seq_lens[:padded_bs]) self.draft_attn_backend.init_forward_metadata_replay_cuda_graph( padded_bs, req_pool_indices, draft_seq_lens, req_to_page=self.drafter.req_to_page, forward_mode=draft_forward_mode, **draft_attn_kwargs, ) @nvtx_range("attn_meta_prep", color="orange") def _init_forward_metadata( self, padded_bs: int, num_extends: int, req_pool_indices: torch.Tensor, seq_lens: torch.Tensor, req_to_page: torch.Tensor, forward_mode: ForwardMode, **kwargs, ): """Eager path — allocate/refresh metadata for the upcoming forward.""" if ( getattr(self.attn_backend, "uses_paged_cache_groups", False) and self.drafter is not None and forward_mode.is_decode() ): kwargs.setdefault("num_tokens", padded_bs * self.max_tokens_per_req) self.attn_backend.init_forward_metadata( bs=padded_bs, num_extends=num_extends, req_pool_indices=req_pool_indices, seq_lens=seq_lens, req_to_page=req_to_page, forward_mode=forward_mode, **kwargs, ) if self.draft_attn_backend is not None: draft_kwargs = {} if getattr(self.draft_attn_backend, "uses_paged_cache_groups", False): for key in ( "paged_cache_block_tables", "paged_cache_block_table_base_offsets", ): value = kwargs.get(key) if value is not None: draft_kwargs[key] = value draft_flat = self._draft_flat_tables(kwargs.get("flat_block_tables")) if draft_flat is not None: draft_kwargs["flat_block_tables"] = draft_flat # The drafter mutates draft_seq_lens_buf between MTP draft steps; # decode metadata must alias that buffer. draft_seq_lens = self.drafter.draft_seq_lens_buf[:padded_bs] draft_seq_lens.copy_(seq_lens[:padded_bs]) if forward_mode.is_extend_or_mixed(): # Non-V4 draft backends follow the legacy contract: a single # EXTEND/MIXED metadata init fills both first-step prefill # metadata and step 1+ decode metadata, with seq_lens aliased # to the drafter-owned mutable buffer. V4 additionally needs # the accepted-prefix view for first-step grouped-cache # metadata, then a separate decode init to prepare the draft # decode metadata from that first-step state. draft_prefill_seq_lens = ( seq_lens if self.use_v4_mtp_paged_metadata else draft_seq_lens ) # Drafter consumes only the full group's table (see _draft_flat_tables). draft_extend_kwargs = ( {**kwargs, "flat_block_tables": draft_flat} if kwargs.get("flat_block_tables") is not None else kwargs ) self.draft_attn_backend.init_forward_metadata( bs=padded_bs, num_extends=num_extends, req_pool_indices=req_pool_indices, seq_lens=draft_prefill_seq_lens, req_to_page=self.drafter.req_to_page, forward_mode=forward_mode, **draft_extend_kwargs, ) if self.use_v4_mtp_paged_metadata: self.draft_attn_backend.init_forward_metadata( bs=padded_bs, num_extends=0, req_pool_indices=req_pool_indices, seq_lens=draft_seq_lens, req_to_page=self.drafter.req_to_page, forward_mode=ForwardMode.DECODE, **draft_kwargs, ) else: draft_metadata_seq_lens = ( seq_lens if self.use_v4_mtp_paged_metadata else draft_seq_lens ) draft_forward_mode = ForwardMode.DECODE if getattr(self.draft_attn_backend, "uses_paged_cache_groups", False): draft_kwargs["num_tokens"] = padded_bs * self.max_tokens_per_req self.draft_attn_backend.init_forward_metadata( bs=padded_bs, num_extends=0, req_pool_indices=req_pool_indices, seq_lens=draft_metadata_seq_lens, req_to_page=self.drafter.req_to_page, forward_mode=draft_forward_mode, **draft_kwargs, ) def _global_graph_bs(self, ctx: ForwardContext) -> int | None: if self.dp_size <= 1 or ctx.global_num_tokens is None: return None max_num_tokens = max(ctx.global_num_tokens) return (max_num_tokens + self.max_tokens_per_req - 1) // self.max_tokens_per_req def _can_use_graph(self, bs: int, ctx: ForwardContext) -> bool: if self.disable: return False if not ctx.forward_mode.is_decode(): return False if self.dp_size > 1: if not ctx.all_decode_or_idle: return False global_bs = self._global_graph_bs(ctx) if global_bs is None or global_bs == 0: return False if self.disable_padding: return self._has_cuda_graph_for_bs(global_bs) return global_bs <= self.max_bs if self.disable_padding: return self._has_cuda_graph_for_bs(bs) return bs <= self.max_bs def can_run(self, bs: int, ctx: ForwardContext) -> bool: return self._can_use_graph(bs, ctx) def padded_bs(self, bs: int, ctx: ForwardContext) -> int: return self._padded_bs(bs, ctx) def _padded_bs(self, bs: int, ctx: ForwardContext) -> int: graph_bs = self._global_graph_bs(ctx) target_bs = graph_bs if graph_bs is not None else bs index = bisect.bisect_left(self.capture_bs, target_bs) return self.capture_bs[index] def _pad_graph_req_pool_indices( self, active_req_pool_indices: torch.Tensor, padded_bs: int ) -> torch.Tensor: pad = padded_bs - active_req_pool_indices.shape[0] if pad <= 0: return active_req_pool_indices if self.config.spec_algo == "DFLASH": # Route padding rows to the sentinel req-pool slot # (max_req_pool_size), not slot 0. The DFLASH draft derives each # row's block seq_len from valid_cache_lengths[req_pool], so # padding rows pointing at slot 0 would grow unbounded with # request 0's context and hang the draft block-decode kernel. # The sentinel row stays zero-init (length 0, dummy page 0). sentinel = int(self.config.max_req_pool_size) return torch.cat( [ active_req_pool_indices, active_req_pool_indices.new_full((pad,), sentinel), ] ) return torch.cat( [active_req_pool_indices, active_req_pool_indices.new_zeros(pad)] ) def _set_graph_state_write_indices( self, active_req_pool_indices: torch.Tensor, padded_bs: int ) -> None: state_indices = self.input_buffers.state_write_req_pool_indices_buf[:padded_bs] active_bs = active_req_pool_indices.shape[0] if active_bs > 0: state_indices[:active_bs].copy_(active_req_pool_indices) if active_bs < padded_bs: state_indices[active_bs:padded_bs].fill_(int(self.config.max_req_pool_size)) def __call__( self, bs: int, ctx: ForwardContext, sampling_info: SamplingBatchInfo, req_to_page: torch.Tensor, extend_with_prefix: bool = False, extend_prefix_lens: torch.Tensor | None = None, extend_prefix_lens_cpu: torch.Tensor | None = None, extend_seq_lens: torch.Tensor | None = None, extend_seq_lens_cpu: torch.Tensor | None = None, positions: torch.Tensor | None = None, out_cache_loc: torch.Tensor | None = None, mamba_pool_indices: torch.Tensor | None = None, mamba_cow_src_indices: torch.Tensor | None = None, mamba_branching_seqlens: torch.Tensor | None = None, mamba_track_pool_indices: torch.Tensor | None = None, spec_info=None, paged_cache_block_tables: dict | None = None, paged_cache_block_table_base_offsets: dict | None = None, flat_block_tables: dict | None = None, ): """ Unified forward entry point. Dispatches to the captured CUDA graph when possible; falls back to the eager forward_func otherwise. The caller does not need to know which path was taken. """ use_graph = self._can_use_graph(bs, ctx) padded_bs = self._padded_bs(bs, ctx) if use_graph else bs active_req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs] if use_graph and padded_bs != bs: ctx.bs = padded_bs pad = padded_bs - bs seq_lens = torch.nn.functional.pad( self.input_buffers.seq_lens_buf[:bs], (0, pad), value=1 ) req_pool_indices = self._pad_graph_req_pool_indices( active_req_pool_indices, padded_bs ) self.input_buffers.seq_lens_buf[:padded_bs].copy_(seq_lens) self.input_buffers.req_pool_indices_buf[:padded_bs].copy_(req_pool_indices) if mamba_pool_indices is not None: # Pad with -1 (PAD_SLOT_ID), NOT 0. Mamba slot 0 is a real # allocatable slot, so padding with 0 aliases a live request's # mamba state and corrupts it. -1 is the kernel-skipped pad id. mamba_pool_indices = torch.nn.functional.pad( mamba_pool_indices, (0, pad), value=-1 ) if mamba_cow_src_indices is not None: mamba_cow_src_indices = torch.nn.functional.pad( mamba_cow_src_indices, (0, pad), value=-1 ) if mamba_branching_seqlens is not None: mamba_branching_seqlens = torch.nn.functional.pad( mamba_branching_seqlens, (0, pad), value=-1 ) if mamba_track_pool_indices is not None: mamba_track_pool_indices = torch.nn.functional.pad( mamba_track_pool_indices, (0, pad), value=-1 ) else: seq_lens = self.input_buffers.seq_lens_buf[:padded_bs] req_pool_indices = self.input_buffers.req_pool_indices_buf[:padded_bs] if use_graph: self._set_graph_state_write_indices(active_req_pool_indices, padded_bs) mamba_kwargs = {} if mamba_pool_indices is not None: mamba_kwargs["mamba_pool_indices"] = mamba_pool_indices if mamba_cow_src_indices is not None: mamba_kwargs["mamba_cow_src_indices"] = mamba_cow_src_indices if mamba_branching_seqlens is not None: mamba_kwargs["mamba_branching_seqlens"] = mamba_branching_seqlens if mamba_track_pool_indices is not None: mamba_kwargs["mamba_track_pool_indices"] = mamba_track_pool_indices if use_graph: if ( bs == 0 and paged_cache_block_tables is None and self.attn_backend.uses_paged_cache_groups ): paged_cache_block_tables = self._capture_paged_cache_block_tables( padded_bs, self.token_to_kv_pool, ) # The backend's stale-table guard also covers the bs==0 idle # replay: synthesize minimal valid tables for it. if ( bs == 0 and not flat_block_tables and getattr(self.attn_backend, "uses_flat_cache_groups", False) ): flat_block_tables = self._idle_flat_block_tables(padded_bs) self._init_replay_metadata( padded_bs, bs, req_pool_indices, seq_lens, req_to_page=req_to_page, forward_mode=ctx.forward_mode, num_padding=padded_bs - bs if padded_bs != bs else 0, paged_cache_block_tables=paged_cache_block_tables, paged_cache_block_table_base_offsets=( paged_cache_block_table_base_offsets ), flat_block_tables=flat_block_tables, **mamba_kwargs, ) # Runtime prepare() is called by ModelExecutor with per-request rids # BEFORE self.forward_step — we don't refill here to avoid clobbering # the per-request generators with the capture-stub generator. self.deepep_adapter.replay() graph_key = self._cuda_graph_key(padded_bs) with nvtx_range("graph_replay", color="red"): self.graphs[graph_key].replay() ( output_tokens, output_lengths, output_logprobs, ) = self.output_buffers[graph_key] result = ( output_tokens[: bs * self.max_tokens_per_req], output_lengths[:bs], ( output_logprobs[: bs * self.max_tokens_per_req] if output_logprobs is not None else None ), ) else: # Eager parity with the replay stale-table guard: with >1 group # the single-table fallback would serve first-group pages to # every layer. Idle/bs==0 forwards carry no requests (exempt); # a single published group falls back to the single table. if ( bs > 0 and not ctx.forward_mode.is_idle() and not flat_block_tables and getattr(self.attn_backend, "uses_flat_cache_groups", False) and len(self.token_to_kv_pool.paged_cache_group_specs) > 1 ): raise RuntimeError( "CudaGraphWrapper eager forward: pool publishes " f"{len(self.token_to_kv_pool.paged_cache_group_specs)} " "flat cache groups and the backend consumes flat tables, " f"but flat_block_tables is missing/empty at bs={bs} " f"({ctx.forward_mode.name}); the single-table fallback " "would use one group's pages for all layers." ) metadata_num_tokens = ( {"num_tokens": ctx.input_num_tokens} if self.attn_backend.uses_paged_cache_groups else {} ) self._init_forward_metadata( padded_bs, ctx.num_extends, req_pool_indices, seq_lens, req_to_page=req_to_page, forward_mode=ctx.forward_mode, extend_with_prefix=extend_with_prefix, extend_prefix_lens=extend_prefix_lens, extend_prefix_lens_cpu=extend_prefix_lens_cpu, extend_seq_lens=extend_seq_lens, extend_seq_lens_cpu=extend_seq_lens_cpu, positions=positions, out_cache_loc=out_cache_loc, global_num_tokens=ctx.global_num_tokens, all_decode_or_idle=ctx.all_decode_or_idle, capture_hidden_mode=ctx.capture_hidden_mode, spec_info=spec_info, **metadata_num_tokens, paged_cache_block_tables=( paged_cache_block_tables if self.attn_backend.uses_paged_cache_groups else None ), paged_cache_block_table_base_offsets=( paged_cache_block_table_base_offsets if self.attn_backend.uses_paged_cache_groups else None ), flat_block_tables=( flat_block_tables if self.attn_backend.uses_flat_cache_groups else None ), **mamba_kwargs, ) result = self._forward_func(bs=bs, ctx=ctx, sampling_info=sampling_info) if use_graph and padded_bs != bs: ctx.bs = bs # Update mamba/GDN state after speculative verify if _should_update_mamba_state_after_mtp_verify( self.drafter, self.attn_backend, ctx.forward_mode ): accept_lengths = result[1] self.attn_backend.update_mamba_state_after_mtp_verify(accept_lengths, None) return result