# 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 time from dataclasses import dataclass from typing import TYPE_CHECKING import torch from tokenspeed_kernel.ops.tuning import freeze_autotuning from tokenspeed_kernel.platform import current_platform from tokenspeed.runtime.configs.model_config import ModelConfig from tokenspeed.runtime.configs.paged_cache_spec import ( scheduler_ext_flat_kvcache, validate_flat_scheduler_config, ) from tokenspeed.runtime.configs.utils import get_rope_parameters from tokenspeed.runtime.engine.scheduler_utils import ( flat_block_tables_from_forward_op, paged_cache_block_table_base_offsets_from_forward_op, paged_cache_block_tables_from_forward_op, ) from tokenspeed.runtime.execution.cache_loc_kernel import update_block_table from tokenspeed.runtime.execution.context import ForwardContext from tokenspeed.runtime.execution.cuda_graph_wrapper import CudaGraphWrapper from tokenspeed.runtime.execution.drafter.dflash import DFlash from tokenspeed.runtime.execution.drafter.eagle import Eagle from tokenspeed.runtime.execution.forward_batch_info import ( CaptureHiddenMode, ForwardMode, ) from tokenspeed.runtime.execution.input_buffer import InputBuffers from tokenspeed.runtime.execution.model_runner import ModelRunner from tokenspeed.runtime.execution.nan_guard import NanGuard from tokenspeed.runtime.execution.prefill_graph import PrefillGraph from tokenspeed.runtime.execution.runtime_states import RuntimeStates from tokenspeed.runtime.execution.types import ModelExecutionResult from tokenspeed.runtime.grammar.capturable_grammar import setup_grammar_step from tokenspeed.runtime.layers.logits_processor import LogitsProcessorOutput from tokenspeed.runtime.layers.paged_attention import ( validate_paged_cache_group_ids, ) from tokenspeed.runtime.sampling.backends.base import SamplingBackend from tokenspeed.runtime.sampling.dp_sampling_config import ( DpSamplingRuntimeLimits, DpSamplingTopology, resolve_dp_sampling_runtime, resolve_dp_sampling_support, ) from tokenspeed.runtime.sampling.sampling_batch_info import SamplingBatchInfo from tokenspeed.runtime.utils import get_colorful_logger, set_random_seed from tokenspeed.runtime.utils.common import maybe_inference_mode from tokenspeed.runtime.utils.env import envs from tokenspeed.runtime.utils.nvtx import nvtx_range from tokenspeed.runtime.utils.server_args import ServerArgs if TYPE_CHECKING: from tokenspeed.runtime.layers.attention.backends.base import AttentionBackend from tokenspeed.runtime.layers.attention.kv_cache.base import BaseTokenToKVPool from tokenspeed.runtime.sampling.sampling_params import SamplingParams logger = get_colorful_logger(__name__) _DRAFTER_MAPPING = {"EAGLE3": Eagle, "MTP": Eagle, "DFLASH": DFlash} LOG_MM_TIMING = envs.TOKENSPEED_LOG_MM_TIMING.get() def _eagle_aux_layer_ids(hf_config) -> list[int] | None: """Draft's eagle_aux_hidden_state_layer_ids (nested or top-level), or None.""" eagle_config = getattr(hf_config, "eagle_config", None) if isinstance(eagle_config, dict): ids = eagle_config.get("eagle_aux_hidden_state_layer_ids") elif eagle_config is not None: ids = getattr(eagle_config, "eagle_aux_hidden_state_layer_ids", None) else: ids = getattr(hf_config, "eagle_aux_hidden_state_layer_ids", None) return list(ids) if ids else None def _draft_idle_global_num_tokens_for_step( step_idx: int, global_num_tokens: list[int], global_bs: list[int] | None, ) -> list[int]: if step_idx == 0 or global_bs is None: return global_num_tokens return global_bs PREFILL_GRAPH_DEFAULT_MAX_TOKENS = 2048 def _resolve_prefill_graph_max_tokens(server_args) -> int: """Largest prefill-graph bucket: explicit value, or min(2048, chunk, kv budget).""" if server_args.prefill_graph_max_tokens is not None: return int(server_args.prefill_graph_max_tokens) cap = PREFILL_GRAPH_DEFAULT_MAX_TOKENS if server_args.chunked_prefill_size: cap = min(cap, int(server_args.chunked_prefill_size)) if server_args.max_total_tokens: cap = min(cap, int(server_args.max_total_tokens)) return cap @dataclass class ModelExecutorConfig: """ Scalar configuration for ModelExecutor. Contains only primitive values — no heavy objects. Created once via from_server_args() and injected into ModelExecutor. """ # Rank-local graph-padding req-pool index. The C++ scheduler owns real rows # 1..max_batch_size and row 0 is reserved, so this must sit after the # scheduler-owned range. max_req_pool_size: int output_length: int enforce_eager: bool block_size: int max_num_seqs: int chunked_prefill_size: int vocab_size: int context_len: int device: str gpu_id: int global_rank: int num_total_pages: int decode_log_interval: int cudagraph_capture_sizes: list[int] | None disable_cuda_graph_padding: bool max_cudagraph_capture_size: int model_is_mrope: bool enable_nan_detection: bool = False # ====== DP ========= data_parallel_size: int = 1 world_size: int = 1 world_group: list[int] | None = None # ====== SPEC ========= spec_algo: str | None = None spec_num_steps: int | None = None # spec_num_tokens == spec_num_steps + 1 for now (without Tree Attention) spec_num_tokens: int | None = None overlap_schedule_depth: int = 0 dp_sampling: bool = False dp_sampling_min_bs: int | None = None use_v4_mtp_paged_metadata: bool = False # ====== GRAMMAR ========= # "none" disables all grammar handling; otherwise the backend name # (currently only "xgrammar" is implemented). grammar_backend: str = "xgrammar" # Force the synchronous eager grammar fallback even on CUDA. For # parity-testing the captured-grammar path. disable_capturable_grammar: bool = False # ====== PREFILL CUDA GRAPH (breakable) ========= disable_prefill_graph: bool = False # Opt-in: > 0 enables the prefill graph and caps the largest token bucket. prefill_graph_max_tokens: int = 0 # Explicit bucket list overriding the ladder (see get_prefill_token_buckets). prefill_graph_capture_sizes: list[int] | None = None @staticmethod def from_server_args( server_args: ServerArgs, model_config: ModelConfig, max_req_pool_size: int, gpu_id: int, global_rank: int, num_total_pages: int, overlap_schedule_depth: int = 0, ) -> ModelExecutorConfig: output_length = ( server_args.speculative_num_draft_tokens if server_args.speculative_algorithm else 1 ) rope_parameters = get_rope_parameters(model_config.hf_text_config) model_is_mrope = bool(rope_parameters and "mrope_section" in rope_parameters) return ModelExecutorConfig( max_req_pool_size=max_req_pool_size, output_length=output_length, enforce_eager=server_args.enforce_eager, block_size=server_args.block_size, max_num_seqs=server_args.max_num_seqs, chunked_prefill_size=server_args.chunked_prefill_size, vocab_size=model_config.vocab_size, context_len=model_config.context_len, device=server_args.device, gpu_id=gpu_id, global_rank=global_rank, num_total_pages=num_total_pages, decode_log_interval=server_args.decode_log_interval, cudagraph_capture_sizes=server_args.cudagraph_capture_sizes, disable_cuda_graph_padding=server_args.disable_cuda_graph_padding, max_cudagraph_capture_size=server_args.max_cudagraph_capture_size, disable_prefill_graph=bool(server_args.disable_prefill_graph), prefill_graph_max_tokens=_resolve_prefill_graph_max_tokens(server_args), prefill_graph_capture_sizes=server_args.prefill_graph_capture_sizes, model_is_mrope=model_is_mrope, data_parallel_size=server_args.mapping.attn.dp_size, world_size=server_args.mapping.world_size, world_group=server_args.mapping.world_group, spec_algo=server_args.speculative_algorithm, spec_num_steps=server_args.speculative_num_steps, spec_num_tokens=server_args.speculative_num_draft_tokens, overlap_schedule_depth=overlap_schedule_depth, dp_sampling=server_args.dp_sampling, dp_sampling_min_bs=server_args.dp_sampling_min_bs, enable_nan_detection=server_args.enable_nan_detection, use_v4_mtp_paged_metadata=model_config.use_v4_mtp_paged_metadata, grammar_backend=server_args.grammar_backend, disable_capturable_grammar=server_args.disable_capturable_grammar, ) class ModelExecutor: """ Orchestrates model forward execution. """ def __init__( self, config: ModelExecutorConfig, model_runner: ModelRunner, attn_backend: AttentionBackend, token_to_kv_pool: BaseTokenToKVPool, sampling_backend: SamplingBackend, draft_model_runner: ModelRunner | None = None, draft_attn_backend: AttentionBackend | None = None, draft_token_to_kv_pool: BaseTokenToKVPool | None = None, mamba_pool: object | None = None, ): self.device = config.device self.config = config self.model_runner = model_runner self.sampling_backend = sampling_backend self.attn_backend = attn_backend self.token_to_kv_pool = token_to_kv_pool # Full-attention group mirrored into req_to_page each step (flat+spec). _group_specs = getattr(token_to_kv_pool, "paged_cache_group_specs", ()) or () self._flat_full_group_id = next( ( str(spec.group_id) for spec in _group_specs if getattr(spec, "family", "history") != "state" and getattr(spec, "retention", None) == "full_history" ), None, ) self._mirror_idx_cpu: torch.Tensor | None = None self._mirror_idx_dev: torch.Tensor | None = None self._mirror_row_buf: torch.Tensor | None = None self.draft_attn_backend = draft_attn_backend self.draft_token_to_kv_pool = draft_token_to_kv_pool self._layerwise_mamba_cow_done = None # Must precede CUDA-graph capture: unsupported flat combinations # (e.g. spec on a backend without flat_spec_capable) would otherwise # die on a capture-path assert instead of this actionable error. validate_flat_scheduler_config( flat_kvcache_ext=scheduler_ext_flat_kvcache(), paged_cache_groups=_group_specs, attn_backend=attn_backend, kv_pool=token_to_kv_pool, speculative_algorithm=config.spec_algo, ) if config.spec_algo is not None: # The DFLASH overlap scheduler reserves a fresh draft block per # decode step a request stays scheduled, including the few steps it # lingers between finishing and eviction, so peak page count runs # ~1 page past context_len + spec_num_tokens. Without headroom # req_to_page overflows and the next draft block's page write goes # out of bounds, hanging the attention kernel. Pad generously; a few # int32 columns per request. Non-DFLASH algorithms do not need this. draft_block_reservation_slack = ( config.spec_num_tokens * 64 if config.spec_algo == "DFLASH" else 0 ) max_num_pages_per_req = ( config.context_len + config.spec_num_tokens + draft_block_reservation_slack + config.block_size - 1 ) // config.block_size else: max_num_pages_per_req = ( config.context_len + config.block_size ) // config.block_size max_bs = config.max_num_seqs // max(config.data_parallel_size, 1) self.req_to_page = torch.zeros( (config.max_req_pool_size + 1, max_num_pages_per_req), dtype=torch.int32, device=self.device, ) spec_num_tokens = config.spec_num_tokens if config.spec_algo is not None else 1 self.input_buffers = InputBuffers( max_bs=max_bs, max_num_tokens=config.chunked_prefill_size, page_size=config.block_size, # token_to_kv_pool allocates size+page_size slots; index `size` is # the reserved dummy slot (see MHATokenToKVPool._create_buffers). dummy_kv_slot=0, state_write_padding_pool_index=config.max_req_pool_size, device=self.device, has_mamba=(mamba_pool is not None), ) self.runtime_states = RuntimeStates( req_pool_size=config.max_req_pool_size, context_len=config.context_len, vocab_size=config.vocab_size, device=self.device, output_length=config.output_length, mamba_pool=mamba_pool, ) # Sized like InputBuffers.max_bs so the padded graph-bucket bs fits. self.nan_guard = NanGuard.create( config.enable_nan_detection, max_bs, self.device, ) if self.config.spec_algo is not None: DrafterImpl = _DRAFTER_MAPPING[config.spec_algo] self.drafter = DrafterImpl( spec_num_tokens=config.spec_num_tokens, spec_num_steps=config.spec_num_steps, draft_model_runner=draft_model_runner, page_size=config.block_size, runtime_states=self.runtime_states, input_buffers=self.input_buffers, req_to_page=self.req_to_page, attn_backend=draft_attn_backend, token_to_kv_pool=draft_token_to_kv_pool, vocab_size=config.vocab_size, ) if hasattr(self.drafter, "bind_target_model"): self.drafter.bind_target_model(self.model_runner.model) # EAGLE3/MTP share the target's embed + lm_head; DFLASH ships its # own draft weights, so it must NOT inherit the target's. if config.spec_algo in ("EAGLE3", "MTP"): embed, head = self.model_runner.model.get_embed_and_head() draft_model_runner.model.set_embed_and_head(embed, head) target_hf = self.model_runner.model_config.hf_config mm_pad_substitute_id = getattr( target_hf, "image_token_id", None ) or getattr(target_hf, "media_placeholder_token_id", None) if mm_pad_substitute_id is not None and hasattr( self.drafter, "set_mm_pad_substitute_id" ): self.drafter.set_mm_pad_substitute_id(mm_pad_substitute_id) if config.spec_algo in ("EAGLE3",) and hasattr( self.model_runner.model, "set_eagle3_layers_to_capture" ): # capture the layers the draft was trained on, not the default aux_layer_ids = _eagle_aux_layer_ids( draft_model_runner.model_config.hf_config ) self.model_runner.model.set_eagle3_layers_to_capture(aux_layer_ids) if config.spec_algo == "DFLASH": if not hasattr(self.model_runner.model, "set_dflash_layers_to_capture"): raise ValueError( "DFLASH requires the target model to support " "set_dflash_layers_to_capture." ) self.model_runner.model.set_dflash_layers_to_capture( self.drafter.target_layer_ids ) else: self.drafter = None # Single grammar handle: CapturableGrammarExecutor on CUDA (uses # cudaLaunchHostFunc on a side stream so the xgrammar fill + # H2D overlap with the forward, and is also CUDA-graph-capturable), # EagerGrammarBuffers on non-CUDA (synchronous fallback). # ``disable_capturable_grammar`` forces the eager path on CUDA too # for parity-testing. self.grammar_runtime = None if config.grammar_backend != "none": from tokenspeed.runtime.grammar.capturable_grammar import ( CapturableGrammarExecutor, EagerGrammarBuffers, ) use_captured = ( current_platform().is_nvidia and not config.disable_capturable_grammar ) if use_captured: self.grammar_runtime = CapturableGrammarExecutor( max_bs=max_bs, vocab_size=config.vocab_size, max_tokens_per_req=spec_num_tokens, device=self.device, ) else: self.grammar_runtime = EagerGrammarBuffers( max_bs=max_bs, vocab_size=config.vocab_size, max_tokens_per_req=spec_num_tokens, device=self.device, ) attn_backend.configure_runtime( sliding_window_size=model_runner.sliding_window_size, req_to_page=self.req_to_page, ) if draft_attn_backend is not None: draft_attn_backend.configure_runtime( sliding_window_size=model_runner.sliding_window_size, req_to_page=self.req_to_page, ) validate_paged_cache_group_ids( model_runner.model, token_to_kv_pool.paged_cache_group_specs, ) if draft_model_runner is not None and draft_token_to_kv_pool is not None: validate_paged_cache_group_ids( draft_model_runner.model, draft_token_to_kv_pool.paged_cache_group_specs, ) processor = self.model_runner.model.logits_processor dp_topology = DpSamplingTopology( tp_rank=processor.tp_rank, tp_size=processor.tp_size, tp_group=processor.tp_group, skip_all_gather=processor.skip_all_gather, tie_word_embeddings=bool( getattr(processor.config, "tie_word_embeddings", False) ), ) dp_support = resolve_dp_sampling_support( requested=self.config.dp_sampling, drafter_available=self.drafter is not None, backend_supports_verify=bool( getattr(self.sampling_backend, "_SUPPORTS_DP_VERIFY", False) ), topology=dp_topology, ) lm_head_rows = 0 if dp_support.enabled: lm_head_weight = self.model_runner.model.lm_head.weight if lm_head_weight.ndim < 1: raise RuntimeError( "dp_sampling LM head weight must be at least 1D, got " f"{lm_head_weight.ndim}D" ) lm_head_rows = int(lm_head_weight.shape[0]) dp_runtime_config = resolve_dp_sampling_runtime( support=dp_support, lm_head_rows=lm_head_rows, topology=dp_topology, limits=DpSamplingRuntimeLimits( runtime_vocab_size=self.config.vocab_size, max_num_seqs=config.max_num_seqs, data_parallel_size=config.data_parallel_size, num_tokens_per_req=spec_num_tokens, configured_min_bs=self.config.dp_sampling_min_bs, device=self.device, ), ) self.dp_sampling_runtime_config = dp_runtime_config self._last_dp_sampling_route_log: ( tuple[str, int, bool, int, int, bool, int] | None ) = None if dp_runtime_config.enabled: self.sampling_backend.configure_dp_sampling(dp_runtime_config) processor.configure_dp_logits_layout(dp_runtime_config) logger.info( "Batch-DP spec-verify: requested=%s, infra_supports=%s, enabled=%s " "min_bs=%s (drafter=%s, backend_supports_dp=%s, " "tp_size=%s, tp_group=%s)", dp_support.requested, dp_support.infra_supports, dp_support.enabled, dp_runtime_config.min_bs, dp_support.drafter_available, dp_support.backend_supports_verify, dp_support.tp_size, dp_support.tp_group_set, ) self._active_multimodal_context = None self._active_positions_override = None self.forward_step = CudaGraphWrapper( forward_func=self._forward_step, attn_backend=attn_backend, token_to_kv_pool=token_to_kv_pool, input_buffers=self.input_buffers, config=config, drafter=self.drafter, draft_attn_backend=draft_attn_backend, draft_token_to_kv_pool=draft_token_to_kv_pool, capturable_grammar=self.capturable_grammar, eager_grammar_buffers=self.eager_grammar_buffers, sampling_backend=self.sampling_backend, runtime_states=self.runtime_states, ) # Breakable prefill (extend) CUDA graphs, the extend-mode analogue of # the decode wrapper above; captures in __init__, borrowing the decode # capture stream so all graphs share one mempool-reuse domain. self.prefill_graph = PrefillGraph( model_runner=self.model_runner, attn_backend=attn_backend, token_to_kv_pool=token_to_kv_pool, input_buffers=self.input_buffers, config=config, req_to_page=self.req_to_page, drafter=self.drafter, decode_wrapper=self.forward_step, ) # Encoder CUDA graph: install model-built wrappers by overriding # modality encoder callables (e.g. ``image_encoder``, ``video_encoder``). # Multimodal-encoder analogue of ``forward_step``'s ``CudaGraphWrapper``. self.encoder_graph_wrappers = {} _mm_model = self.model_runner.model if ( hasattr(_mm_model, "make_encoder_cudagraph_wrappers") and getattr(_mm_model, "is_multimodal_active", True) and envs.TOKENSPEED_MM_ENABLE_ENCODER_CUDA_GRAPH.get() and self.model_runner.server_args.mm_attention_backend != "flashinfer_cudnn" ): self.encoder_graph_wrappers = _mm_model.make_encoder_cudagraph_wrappers( _mm_model.mapping ) active_encoder_graph_wrappers = {} for encoder_attr, wrapper in self.encoder_graph_wrappers.items(): if not hasattr(_mm_model, encoder_attr): logger.warning( "Skipping encoder CUDA graph wrapper for missing attribute %s", encoder_attr, ) continue setattr(_mm_model, encoder_attr, wrapper) active_encoder_graph_wrappers[encoder_attr] = wrapper self.encoder_graph_wrappers = active_encoder_graph_wrappers self.execution_stream = torch.cuda.Stream() self.log_step = 0 self._seen_prefill_ids: set[str] = set() self._prev_decode_bs: int = 0 self._sentinel_neg1 = torch.tensor(-1, device=self.device, dtype=torch.int64) if config.model_is_mrope: mrope_decode_capacity = self.input_buffers.max_num_tokens # Double-buffered pinned host staging for the decode delta copy. # Under overlap scheduling the next decode forward is dispatched # before the previous result is synchronized, so a single reused # pinned buffer could be refilled by the next step while the prior # step's ``non_blocking=True`` H2D copy is still reading it (a race # that corrupts M-RoPE deltas). Ping-pong two buffers so a buffer is # never overwritten while its copy is in flight (overlap depth 1). self._mrope_decode_deltas_cpu = [ self._make_mrope_decode_deltas_cpu(mrope_decode_capacity), self._make_mrope_decode_deltas_cpu(mrope_decode_capacity), ] self._mrope_decode_deltas_cpu_idx = 0 self._mrope_decode_deltas_buf = torch.zeros( mrope_decode_capacity, device=self.device, dtype=torch.int64 ) else: self._mrope_decode_deltas_cpu = None self._mrope_decode_deltas_cpu_idx = 0 self._mrope_decode_deltas_buf = None # Decode stats — accumulated from synced results (no GPU sync needed) self.num_generated_tokens = 0 self.num_decode_steps = 0 self.last_decode_stats_tic = time.time() set_random_seed(48) # Startup has tuned every size class it serves; serving must never autotune. freeze_autotuning() logger.info("ModelExecutor initialized") @staticmethod def _make_mrope_decode_deltas_cpu(size: int) -> torch.Tensor: try: return torch.zeros(size, dtype=torch.int64, pin_memory=True) except RuntimeError: return torch.zeros(size, dtype=torch.int64) @property def capturable_grammar(self): """Captured-graph grammar handle, or None on the eager-fallback path. Used by ``_forward_step`` to fence the side-stream grammar fill against the captured forward — those calls only make sense for the captured flavor of grammar runtime. """ from tokenspeed.runtime.grammar.capturable_grammar import ( CapturableGrammarExecutor, ) return ( self.grammar_runtime if isinstance(self.grammar_runtime, CapturableGrammarExecutor) else None ) @property def eager_grammar_buffers(self): """Eager-fallback grammar buffer handle, or None on the captured path.""" from tokenspeed.runtime.grammar.capturable_grammar import ( EagerGrammarBuffers, ) return ( self.grammar_runtime if isinstance(self.grammar_runtime, EagerGrammarBuffers) else None ) def _mirror_flat_full_table_into_req_to_page( self, forward_op, flat_block_tables ) -> None: """Flat + spec: scatter the full-attention group's per-batch table into req_to_page rows, restoring the radix contract for every legacy consumer (input prep's out_cache_loc kernels, the drafter's per-step location chains). The flat scheduler never populates req_to_page itself; column tails zero-fill to the dummy page so stale longer rows can't leak.""" if ( self.drafter is None or not flat_block_tables or self._flat_full_group_id is None ): return table = flat_block_tables.get(self._flat_full_group_id) if table is None: return bs = len(forward_op.request_pool_indices) if self._mirror_idx_cpu is None or self._mirror_idx_cpu.shape[0] < bs: cap = max(bs, self.input_buffers.max_bs) self._mirror_idx_cpu = torch.empty(cap, dtype=torch.long, pin_memory=True) self._mirror_idx_dev = torch.empty( cap, dtype=torch.long, device=self.device ) self._mirror_idx_cpu[:bs] = torch.tensor( forward_op.request_pool_indices, dtype=torch.long ) self._mirror_idx_dev[:bs].copy_(self._mirror_idx_cpu[:bs], non_blocking=True) idx = self._mirror_idx_dev[:bs] width = table.shape[1] max_width = self.req_to_page.shape[1] if self._mirror_row_buf is None: self._mirror_row_buf = torch.zeros( (self.input_buffers.max_bs, max_width), dtype=self.req_to_page.dtype, device=self.device, ) # Staged rows + index_copy_ (advanced-index setitem costs ~150us dispatch). rows = self._mirror_row_buf[:bs] rows[:, :width].copy_(table) rows[:, :width].clamp_min_( 0 ) # -1 column pads -> dummy page 0 (negative locs otherwise) if width < max_width: rows[:, width:].zero_() self.req_to_page.index_copy_(0, idx, rows) @nvtx_range("target_forward", color="red") def _run_target_forward(self, bs: int, ctx: ForwardContext, req_pool_indices): positions = self._active_positions_override if positions is None: if self.config.model_is_mrope: positions = self.input_buffers.mrope_positions_buf[ :, : ctx.input_num_tokens ] else: positions = self.input_buffers.positions_buf[: ctx.input_num_tokens] # Prefill-graph replay when captured for this forward (the decode graph # replays one level up: it captures the whole _forward_step). The mode # check is LOAD-BEARING, not an optimization: the decode capture runs # this dispatch before prefill_graph exists (it is constructed after # the decode wrapper, whose capture stream it borrows), and decode-mode # forwards must short-circuit before touching it. mode = ctx.forward_mode if ( mode is not None and (mode.is_extend() or mode.is_mixed()) and self.prefill_graph.can_run(ctx, self._active_multimodal_context) ): return self.prefill_graph.replay( ctx, self.input_buffers.input_ids_buf[: ctx.input_num_tokens], self._active_multimodal_context, ) return self.model_runner.forward( ctx, self.input_buffers.input_ids_buf[: ctx.input_num_tokens], positions, self.input_buffers.out_cache_loc_buf[: ctx.input_num_tokens], req_pool_indices=req_pool_indices, seq_lens=self.input_buffers.seq_lens_buf[:bs], extend_prefix_lens=self.input_buffers.extend_prefix_lens_buf[ : ctx.num_extends ], multimodal_context=self._active_multimodal_context, ) def _apply_force_single_token_verify( self, accept_lengths: torch.Tensor, row_offset: int, row_count: int, decode_input_ids: list[int] | None, ) -> torch.Tensor: if decode_input_ids is None or row_count <= 0: return accept_lengths force_mask = self.input_buffers.force_single_token_verify_buf[ row_offset : row_offset + row_count ] return torch.where(force_mask, torch.ones_like(accept_lengths), accept_lengths) def _cap_accept_to_context_len( self, accept_lengths: torch.Tensor, decode_req_pool_indices: torch.Tensor, ) -> torch.Tensor: """Clamp spec-verify accept so committed length never exceeds ``context_len``. ``req_to_page`` is sized for ``context_len + spec_num_tokens`` pages. A request at the context limit whose ``max_new_tokens`` termination lags a step can accept past ``context_len``, so its next draft block needs a page beyond ``req_to_page``'s width — an out-of-bounds access that hangs the attention kernel. Clamping to the remaining budget keeps the table in range; the request is still removed a step later. Deterministic in ``valid_cache_lengths`` / ``accept_lengths``, so no cross-rank divergence. """ if accept_lengths.numel() == 0: return accept_lengths committed = self.runtime_states.valid_cache_lengths.index_select( 0, decode_req_pool_indices ).to(accept_lengths.dtype) remaining = (self.config.context_len - committed).clamp_(min=0) # In-place: the drafter reads this same buffer to size its next block. accept_lengths.copy_(torch.minimum(accept_lengths, remaining)) return accept_lengths def _clamp_committed_to_context_len( self, output_lengths: torch.Tensor, num_extends: int, bs: int, ) -> torch.Tensor: """Return ``output_lengths`` with decode rows clamped so committed KV length never exceeds ``context_len`` (post-forward, outside the CUDA graph). The clamp must reach BOTH ``_update_runtime_state`` (``valid_cache_lengths``) and the ``ModelExecutionResult`` that drives scheduler page reservation, so they stay in lock-step. Hence a FRESH tensor, not the persistent ``_accept_length_buf``: the verify path also mirrors accept counts into ``_output_pack_buf``, and an in-place clamp would leave that mirror (read by the packed-D2H fast path) uncapped, reserving a draft block past ``req_to_page``'s width and hanging the kernel. A fresh tensor forces the safe two-D2H fallback. Only decode rows ``[num_extends:bs]`` carry an accept delta; prefill rows pass through. Deterministic, so no cross-rank divergence. """ if bs <= num_extends: return output_lengths decode_rpi = self.input_buffers.req_pool_indices_buf[num_extends:bs] committed = self.runtime_states.valid_cache_lengths.index_select( 0, decode_rpi ).to(output_lengths.dtype) remaining = (self.config.context_len - committed).clamp_(min=0) capped_decode = torch.minimum(output_lengths[num_extends:bs], remaining) if num_extends == 0: return capped_decode return torch.cat([output_lengths[:num_extends], capped_decode]) @nvtx_range("sampling", color="yellow") def _run_sampling( self, logits_output: LogitsProcessorOutput, sampling_info: SamplingBatchInfo, ctx: ForwardContext, candidates: torch.Tensor | None = None, ): if self.drafter is None: return self.sampling_backend.sample(logits_output, sampling_info) num_extends = ctx.num_extends num_decodes = ctx.bs - num_extends if num_decodes == 0: return self.sampling_backend.sample(logits_output, sampling_info) if num_extends == 0: output_tokens, accept_lengths = self.sampling_backend.verify( logits_output, sampling_info, candidates ) accept_lengths = self._apply_force_single_token_verify( accept_lengths, 0, num_decodes, ctx.decode_input_ids ) if self.config.spec_algo == "DFLASH": accept_lengths = self._cap_accept_to_context_len( accept_lengths, sampling_info.req_pool_indices[:num_decodes] ) return output_tokens, accept_lengths logits = logits_output.next_token_logits prefill_out = LogitsProcessorOutput(next_token_logits=logits[:num_extends]) prefill_tokens, prefill_accept = self.sampling_backend.sample( prefill_out, sampling_info[:num_extends] ) decode_out = LogitsProcessorOutput(next_token_logits=logits[num_extends:]) decode_tokens, decode_accept = self.sampling_backend.verify( decode_out, sampling_info[num_extends:], candidates ) decode_accept = self._apply_force_single_token_verify( decode_accept, num_extends, num_decodes, ctx.decode_input_ids ) if self.config.spec_algo == "DFLASH": decode_accept = self._cap_accept_to_context_len( decode_accept, sampling_info.req_pool_indices[num_extends:] ) if ( prefill_out.next_token_logprobs is not None and decode_out.next_token_logprobs is not None ): logits_output.next_token_logprobs = torch.cat( [prefill_out.next_token_logprobs, decode_out.next_token_logprobs] ) return ( torch.cat([prefill_tokens, decode_tokens]), torch.cat([prefill_accept, decode_accept]), ) def _log_dp_sampling_route(self, bs: int, ctx: ForwardContext) -> None: runtime = self.dp_sampling_runtime_config if ( self.config.global_rank != 0 or not runtime.enabled or runtime.min_bs is None or runtime.topology is None or ctx.forward_mode is None or not ctx.forward_mode.is_decode() ): return use_graph = self.forward_step.can_run(bs=bs, ctx=ctx) effective_bs = self.forward_step.padded_bs(bs=bs, ctx=ctx) if use_graph else bs tp_size = runtime.topology.tp_size bucket_bs = ((effective_bs + tp_size - 1) // tp_size) * tp_size dp_sampling = effective_bs >= runtime.min_bs route_key = ( ctx.forward_mode.name, bs, use_graph, effective_bs, bucket_bs, dp_sampling, runtime.min_bs, ) if route_key == self._last_dp_sampling_route_log: return self._last_dp_sampling_route_log = route_key logger.debug( "Batch-DP route: forward_mode=%s bs=%d effective_bs=%d " "use_graph=%s bucket_bs=%d dp_sampling=%s min_bs=%d", ctx.forward_mode.name.lower(), bs, effective_bs, use_graph, bucket_bs, dp_sampling, runtime.min_bs, ) @maybe_inference_mode() def _forward_step( self, bs: int, ctx: ForwardContext, sampling_info: SamplingBatchInfo, ): req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs] # Fork grammar onto its side stream so fill + H2D overlap with # attention/MoE. Rejoined at wait_bitmask() before apply_mask. if self.capturable_grammar is not None: n = self.capturable_grammar.max_tokens_per_req is_spec_verify = n > 1 and ctx.forward_mode.is_decode() slice_ = ( self.input_buffers.input_ids_buf[: bs * n] if is_spec_verify else None ) self.capturable_grammar.schedule_fill(input_ids_buf_slice=slice_) logits_output = self._run_target_forward(bs, ctx, req_pool_indices) # Flag NaN per request and sanitize in place, before any sampling kernel. self.nan_guard.audit_logits(logits_output, ctx) candidates = ( self.drafter.get_candidates(ctx) if self.config.spec_algo is not None else None ) if self.capturable_grammar is not None: self.capturable_grammar.wait_bitmask() output_tokens, accept_lengths = self._run_sampling( logits_output, sampling_info, ctx, candidates ) # Backstop: flag any request whose sampled id falls outside [0, vocab) # so the output processor can terminate it. Covers sampler/verify kernel # corruption and DP-sharded steps that audit_logits cannot attribute. self.nan_guard.merge_oov(output_tokens, ctx, self.runtime_states.vocab_size) # Fork sampler-output D2H onto the grammar side stream so the # next step's build hostfunc can advance the matcher. if self.capturable_grammar is not None: self.capturable_grammar.schedule_post_sampler(output_tokens, accept_lengths) if self.drafter is not None: next_round_input_ids = self.drafter.run( base_ctx=ctx, logits_output=logits_output, output_tokens=output_tokens, accept_lengths=accept_lengths, ) # _update_runtime_state skips future_input_map when drafter is # active — drafter writes the next-round inputs directly. self.runtime_states.future_input_map[ self.input_buffers.state_write_req_pool_indices_buf[: ctx.bs] ] = next_round_input_ids.to(torch.int32) output_logprobs = logits_output.next_token_logprobs return output_tokens, accept_lengths, output_logprobs @nvtx_range("update_runtime_state", color="orange") def _update_runtime_state( self, req_pool_indices: torch.Tensor, output_tokens: torch.Tensor, accept_lengths: torch.Tensor, input_lengths: torch.Tensor, num_extends: int, ): """Write output tokens to future_input_map and update cache lengths. Must NOT be captured in CUDA graph — these writes are read by the next iteration's batch prep on the default stream, so they need explicit stream synchronization (see execute_forward_op). """ if self.drafter is None: # Without drafter, store output tokens for next round. # With drafter, _forward_step already wrote the drafter's # next-round input (verified + draft tokens) to future_input_map. tokens_per_req = self.config.output_length if num_extends == 0 else 1 next_round_input_ids = output_tokens.to(torch.int32).reshape( -1, tokens_per_req ) self.runtime_states.future_input_map[req_pool_indices, :tokens_per_req] = ( next_round_input_ids ) bs = req_pool_indices.shape[0] if num_extends == 0: deltas = accept_lengths elif num_extends == bs: deltas = input_lengths else: deltas = torch.cat( [input_lengths[:num_extends], accept_lengths[num_extends:]] ) self.runtime_states.update_valid_cache_length(req_pool_indices, deltas) def _build_sampling_info( self, bs: int, sampling_params_list: list[SamplingParams], ) -> SamplingBatchInfo: return SamplingBatchInfo( req_pool_indices=self.input_buffers.req_pool_indices_buf[:bs], valid_cache_lengths=self.runtime_states.valid_cache_lengths, is_all_greedy=all(p.top_k <= 1 for p in sampling_params_list), vocab_size=self.runtime_states.vocab_size, device=self.device, ) def accumulate_decode_stats(self, results: ModelExecutionResult, bs: int): """Accumulate decode stats from already-synced results. No GPU sync.""" self.num_generated_tokens += int(results.output_lengths.sum().item()) self.num_decode_steps += bs @staticmethod @torch.compile(dynamic=True) def _compute_mtp_snapshot_indices( valid_cache_lengths: torch.Tensor, req_pool_indices: torch.Tensor, accept_lengths: torch.Tensor, output_indices: torch.Tensor, track_indices: torch.Tensor, sentinel: torch.Tensor, page_size: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Fused elementwise pipeline computing snapshot src/dst for MTP. All operations are batched and fused by torch.compile into a single Triton kernel (plus the two gathers), eliminating the ~14 individual elementwise kernel launches of the eager implementation. """ new_cl = valid_cache_lengths[req_pool_indices] old_cl = new_cl - accept_lengths.to(new_cl.dtype) first_boundary = ((old_cl // page_size) + 1) * page_size step_raw = first_boundary - old_cl - 1 max_col = output_indices.shape[1] - 1 step = step_raw.clamp(min=0, max=max_col).to(torch.int64) bs = req_pool_indices.shape[0] req_range = torch.arange(bs, device=req_pool_indices.device) src_raw = output_indices[req_range, step].to(torch.int64) dst_raw = track_indices.to(torch.int64) invalid = ( (first_boundary > new_cl) | (dst_raw < 0) | (src_raw < 0) | (src_raw == dst_raw) | (step_raw < 0) ) src = torch.where(invalid, sentinel, src_raw) dst = torch.where(invalid, sentinel, dst_raw) return src, dst def _snapshot_mamba_checkpoints( self, accept_lengths: torch.Tensor, bs: int, num_extends: int, ) -> None: """Snapshot mamba states to checkpoint slots at page boundaries. Called after ``_update_runtime_state`` on the execution stream so ``valid_cache_lengths`` already reflects the accepted tokens. Non-MTP (accept_length == 1): The working slot holds the up-to-date state for the new cache_length. Pass the kernel page_size so it copies only when the new length is page-aligned. MTP (accept_length > 1): cache_length may jump over a page boundary. The intermediate state lives in ``mamba_output_indices[req, step]``. Boundary detection and source-slot selection are done entirely on GPU with -1 sentinels so the snapshot kernel skips invalid entries via its bounds check — no GPU-to-CPU sync, preserving overlap-schedule pipelining. """ if self.runtime_states.mamba_pool is None or num_extends > 0: return if not self.input_buffers.has_mamba: return req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs] track_indices = self.input_buffers.mamba_track_pool_indices_buf[:bs] page_size = self.config.block_size dev = req_pool_indices.device sentinel = self._sentinel_neg1 if self.drafter is not None: # -- MTP path: find the output slot at the crossed boundary -- backend = getattr( self.attn_backend, "linear_attn_backend", self.attn_backend ) fm = getattr(backend, "forward_metadata", None) if fm is None: return output_indices = fm.mamba_output_indices if output_indices is None: return src, dst = self._compute_mtp_snapshot_indices( self.runtime_states.valid_cache_lengths, req_pool_indices, accept_lengths[:bs].to(device=dev), output_indices, track_indices, sentinel, page_size, ) self.runtime_states.snapshot_mamba_checkpoints( src, dst, cache_lengths=None, page_size=0, num_valid=bs, ) else: # -- Non-MTP path: working slot IS the up-to-date state -- src_raw = self.input_buffers.mamba_pool_indices_buf[:bs].to( device=dev, dtype=torch.int64 ) dst_raw = track_indices.to(device=dev, dtype=torch.int64) invalid = (src_raw < 0) | (dst_raw < 0) | (src_raw == dst_raw) src = torch.where(invalid, sentinel, src_raw) dst = torch.where(invalid, sentinel, dst_raw) cache_lengths = self.runtime_states.valid_cache_lengths[req_pool_indices] self.runtime_states.snapshot_mamba_checkpoints( src, dst, cache_lengths=cache_lengths, page_size=page_size, num_valid=bs, ) def flush_mamba_draft_to_working_on_retract(self) -> None: """Copy accepted draft mamba state -> working slot for all previous-batch requests. Called from event_loop when retract WriteBackOps are detected. Uses the previous decode iteration's input_buffers (still valid since no new forward has overwritten them). Runs on execution_stream to respect ordering with previous forward writes. """ bs = self._prev_decode_bs if bs <= 0: return backend = getattr(self.attn_backend, "linear_attn_backend", self.attn_backend) pool = getattr(backend, "pool", None) if pool is None: return sentinel = self._sentinel_neg1 with torch.cuda.stream(self.execution_stream): req_pool_indices = self.input_buffers.req_pool_indices_buf[:bs] working = self.input_buffers.mamba_pool_indices_buf[:bs] req = req_pool_indices.to(dtype=torch.int64) current_input_size = int(pool.current_input_indices.shape[0]) in_bounds = (req >= 0) & (req < current_input_size) safe_req = req.clamp(0, current_input_size - 1) src_raw = pool.current_input_indices[safe_req].to(dtype=torch.int64) src_raw = torch.where(in_bounds, src_raw, sentinel) dst_raw = working.to(dtype=torch.int64) invalid = ( (~in_bounds) | (src_raw < 0) | (dst_raw < 0) | (src_raw == dst_raw) ) src = torch.where(invalid, sentinel, src_raw) dst = torch.where(invalid, sentinel, dst_raw) self.runtime_states.snapshot_mamba_checkpoints( src, dst, cache_lengths=None, page_size=0, num_valid=bs, ) def execute_forward_op_with_log( self, forward_op, sampling_params_list: list[SamplingParams], num_active_pages: int = 0, num_cached_pages: int = 0, num_queue_reqs: int = 0, dp_global_num_tokens=None, dp_global_bs=None, dp_all_decode_or_idle: bool = False, dp_all_extend: bool = False, grammar_inputs=None, multimodal_context=None, capture_next_input_ids: bool = False, ) -> ModelExecutionResult: self.log_step += 1 num_extends = forward_op.num_extends() bs = len(forward_op.request_ids) is_decode = num_extends <= 0 if not is_decode and self.config.global_rank == 0: mode = "Prefill" if num_extends == bs else "Mix" total_tokens = sum(forward_op.input_lengths) cached_tokens = sum( pl for rid, pl in zip( forward_op.request_ids[:num_extends], forward_op.extend_prefix_lens, ) if rid not in self._seen_prefill_ids ) if len(self._seen_prefill_ids) > 100_000: self._seen_prefill_ids.clear() # log-dedup only; bound the growth self._seen_prefill_ids.update(forward_op.request_ids[:num_extends]) logger.info( "%s batch. #new-seq: %s, #new-token: %s, #cached-token: %s, " "#running-req: %s, #queue-req: %s", mode, num_extends, total_tokens, cached_tokens, bs, num_queue_reqs, ) result = self.execute_forward_op( forward_op, sampling_params_list, dp_global_num_tokens, dp_global_bs, dp_all_decode_or_idle, dp_all_extend, grammar_inputs=grammar_inputs, multimodal_context=multimodal_context, capture_next_input_ids=capture_next_input_ids, ) if is_decode and ( self.config.global_rank == 0 and self.log_step % self.config.decode_log_interval == 0 ): now = time.time() gap = now - self.last_decode_stats_tic gen_throughput = self.num_generated_tokens / gap if gap > 0 else 0 avg_accept = ( self.num_generated_tokens / self.num_decode_steps if self.num_decode_steps > 0 else 0 ) accept_rate = ( (avg_accept - 1) / self.config.spec_num_steps if self.config.spec_num_steps else 0 ) num_total_pages = self.config.num_total_pages page_ratio = ( num_active_pages / num_total_pages if num_total_pages > 0 else 0 ) if self.config.spec_num_steps: logger.info( "Decode batch. #running-req: %s, " "#pages(active/cached/total): %s/%s/%s, " "page ratio: %.2f, gen throughput (token/s): %.2f, " "avg_accept_len: %.2f, accept_rate: %.2f, #queue-req: %s", bs, num_active_pages, num_cached_pages, num_total_pages, page_ratio, gen_throughput, avg_accept, accept_rate, num_queue_reqs, ) else: logger.info( "Decode batch. #running-req: %s, " "#pages(active/cached/total): %s/%s/%s, " "page ratio: %.2f, gen throughput (token/s): %.2f, " "#queue-req: %s", bs, num_active_pages, num_cached_pages, num_total_pages, page_ratio, gen_throughput, num_queue_reqs, ) self.token_to_kv_pool.maybe_log_paged_cache_group_pages() self.num_generated_tokens = 0 self.num_decode_steps = 0 self.last_decode_stats_tic = now return result def execute_idle_forward( self, global_num_tokens: list[int], global_bs: list[int], all_decode_or_idle: bool, ): """Run a zero-token forward so this rank participates in NCCL collectives. Called by the EventLoop when this DP rank has no work but other ranks do. The MoE all-to-all is a collective that requires ALL ranks to participate. """ graph_forward_mode = ForwardMode.DECODE ctx = ForwardContext( attn_backend=self.attn_backend, token_to_kv_pool=self.token_to_kv_pool, req_to_page=self.req_to_page, bs=0, num_extends=0, input_num_tokens=0, forward_mode=graph_forward_mode, global_num_tokens=global_num_tokens, global_bs=global_bs, all_decode_or_idle=all_decode_or_idle, ) sampling_info = SamplingBatchInfo( req_pool_indices=self.input_buffers.req_pool_indices_buf[:0], valid_cache_lengths=self.runtime_states.valid_cache_lengths, is_all_greedy=True, vocab_size=self.runtime_states.vocab_size, device=self.device, ) if self.forward_step.can_run(bs=0, ctx=ctx): padded_bs = self.forward_step.padded_bs(bs=0, ctx=ctx) self.input_buffers.fill_dummy_decode_buffers( batch_size=padded_bs, total_tokens=padded_bs * self.config.output_length, ) # Captured hostfunc pops one entry per replay; push a dummy # for this idle replay, same as run_once. if self.capturable_grammar is not None: self.capturable_grammar.add_batch( grammars=[None] * padded_bs, bs=padded_bs, has_candidates=False ) # IDLE doesn't produce tokens, so no sampler/drafter call here — # only the model forward, which still participates in collectives. with nvtx_range("forward_step idle", color="blue"): self.forward_step( bs=0, ctx=ctx, sampling_info=sampling_info, req_to_page=self.req_to_page, ) return # Run model forward with IDLE mode — skips attention but still # participates in MLP NCCL collectives (dense all-gather, MoE). ctx.forward_mode = ForwardMode.IDLE empty = torch.zeros(0, dtype=torch.int32, device=self.device) self.model_runner.forward( ctx, input_ids=empty, positions=empty, out_cache_loc=empty, ) # If a drafter is active, its model also has MoE layers that issue # NCCL collectives. Idle ranks must match those collectives: # 1 first-step forward + (spec_num_steps - 1) multi-step decode forwards. if self.drafter is not None: # DFLASH is a block drafter (idle_forward_steps=1); EAGLE3/MTP # default to spec_num_steps. Mirror the active rank's per-step # collective sizing either way. idle_forward_steps = getattr( self.drafter, "idle_forward_steps", self.drafter.spec_num_steps ) for step_idx in range(idle_forward_steps or 0): # Mirror active rank's catch-up step: when all non-idle ranks # are decoding, step 0 sizes collectives from bs/global_bs. draft_global_num_tokens = _draft_idle_global_num_tokens_for_step( step_idx, global_num_tokens, global_bs, ) draft_ctx = ForwardContext( attn_backend=self.drafter.attn_backend, token_to_kv_pool=self.drafter.token_to_kv_pool, req_to_page=self.drafter.req_to_page, bs=0, num_extends=0, input_num_tokens=0, forward_mode=ForwardMode.IDLE, global_num_tokens=draft_global_num_tokens, global_bs=global_bs, all_decode_or_idle=all_decode_or_idle, ) self.drafter.draft_model_runner.forward( draft_ctx, input_ids=empty, positions=empty, out_cache_loc=empty, spec_step_idx=step_idx, ) def update_block_table(self, forward_op) -> ModelExecutionResult: # Update page tables on the default stream before switching to execution stream. # HostTodevice segment begins with nvtx_range("update_block_table", color="cyan"): update_block_table( forward_op=forward_op, device=self.device, req_to_page=self.req_to_page, ) def reset_remote_prefill_mamba_inputs(self, forward_op) -> None: if self.runtime_states.mamba_pool is None: return if not hasattr(self.attn_backend, "reset_current_inputs"): return num_extends = forward_op.num_extends() if num_extends <= 0: return mamba_indices = list(getattr(forward_op, "mamba_pool_indices", [])) if not mamba_indices: return req_pool_indices = list(forward_op.request_pool_indices[:num_extends]) pairs = [ (int(req_pool_idx), int(mamba_idx)) for req_pool_idx, mamba_idx in zip( req_pool_indices, mamba_indices[:num_extends] ) if int(mamba_idx) >= 0 ] if not pairs: return req_pool_tensor = torch.tensor( [req_pool_idx for req_pool_idx, _ in pairs], dtype=torch.int64, device="cpu", pin_memory=True, ).to(self.device, non_blocking=True) mamba_tensor = torch.tensor( [mamba_idx for _, mamba_idx in pairs], dtype=torch.int64, device="cpu", pin_memory=True, ).to(self.device, non_blocking=True) self.attn_backend.reset_current_inputs(req_pool_tensor, mamba_tensor) @staticmethod def _contains_retracted_decode(forward_op) -> bool: # FlatForwardOperation stores hist_token_lens for decode rows only; # non-recovery rows use -1. return any( hist_token_len != -1 for hist_token_len in getattr(forward_op, "hist_token_lens", ()) ) @nvtx_range("reset_valid_cache_length", color="orange") def reset_valid_cache_length(self, forward_op) -> None: num_extends = forward_op.num_extends() is_prefill = num_extends > 0 # Retraction recovery: scheduler pushes -1 per decode op, overriding to # a real length only on ScheduleDecodeFromRetractedEvent. Decode rows # may follow prefill rows in a mixed batch, so this cannot be gated on # pure-decode mode. has_retract = self._contains_retracted_decode(forward_op) # Pure decode without retraction has nothing to do — skip the # cross-stream wait + stream-context entry entirely. if not is_prefill and not has_retract: return if has_retract: hist_token_lens_tensor = torch.tensor( forward_op.hist_token_lens, dtype=torch.int32, device="cpu", pin_memory=True, ) decode_pool_indices = torch.tensor( forward_op.request_pool_indices[num_extends:], dtype=torch.int64, device="cpu", pin_memory=True, ) else: hist_token_lens_tensor = None decode_pool_indices = None self.execution_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(self.execution_stream): if is_prefill: extend_request_pool_indices = torch.tensor( forward_op.request_pool_indices[:num_extends], dtype=torch.int64, device="cpu", pin_memory=True, ).to(self.device, non_blocking=True) extend_prefix_lens = torch.tensor( forward_op.extend_prefix_lens, dtype=torch.int32, device="cpu", pin_memory=True, ).to(self.device, non_blocking=True) self.runtime_states.reset_states( extend_request_pool_indices, extend_prefix_lens ) if hist_token_lens_tensor is not None: # Apply retraction recovery: override valid_cache_lengths with hist_token_lens # where the scheduler has specified a non-(-1) value, so that out_cache_loc # and position IDs are computed against the retracted KV length. pool_idx_dev = decode_pool_indices.to(self.device, non_blocking=True) hist_dev = hist_token_lens_tensor.to(self.device, non_blocking=True) mask_1d = hist_dev != -1 vcl = self.runtime_states.valid_cache_lengths[pool_idx_dev] self.runtime_states.valid_cache_lengths[pool_idx_dev] = torch.where( mask_1d, hist_dev, vcl ) def set_layerwise_mamba_cow_done( self, cow_by_src: dict[int, list[int]] | None ) -> None: self._layerwise_mamba_cow_done = ( { int(src): {int(dst) for dst in dsts} for src, dsts in cow_by_src.items() if dsts } if cow_by_src else None ) def _skip_completed_layerwise_mamba_cow( self, forward_op, bs: int ) -> torch.Tensor | None: cow_done = self._layerwise_mamba_cow_done self._layerwise_mamba_cow_done = None if not cow_done or not getattr(self.input_buffers, "has_mamba", False): return None cow_src_indices = getattr(forward_op, "mamba_cow_src_indices", None) working_indices = getattr(forward_op, "mamba_pool_indices", None) if cow_src_indices is None or working_indices is None: return None cow_src_indices = list(cow_src_indices)[:bs] working_indices = list(working_indices)[:bs] skipped_mask = [False] * bs changed = False for i, (cow_src, working) in enumerate(zip(cow_src_indices, working_indices)): cow_src = int(cow_src) working = int(working) if working in cow_done.get(cow_src, set()): cow_src_indices[i] = -1 skipped_mask[i] = True changed = True if not changed: return None self.input_buffers._mamba_cow_src_indices_cpu[:bs].copy_( torch.as_tensor(cow_src_indices, dtype=torch.int32) ) cow_src_buf = self.input_buffers.mamba_cow_src_indices_buf cow_src_buf[:bs].copy_( self.input_buffers._mamba_cow_src_indices_cpu[:bs], non_blocking=True ) return torch.tensor(skipped_mask, dtype=torch.bool, device=cow_src_buf.device) @staticmethod def _mamba_retract_reset_mask( mamba_cow_src: torch.Tensor, bs: int, skipped_layerwise_cow_mask: torch.Tensor | None, ) -> torch.Tensor: reset_mask = mamba_cow_src[:bs] >= 0 if skipped_layerwise_cow_mask is not None: reset_mask = reset_mask | skipped_layerwise_cow_mask[:bs].to( device=reset_mask.device, dtype=torch.bool ) return reset_mask def _reset_mamba_current_inputs( self, *, num_extends: int, bs: int, has_retract: bool, mamba_pool_indices: torch.Tensor, mamba_cow_src: torch.Tensor, skipped_layerwise_cow_mask: torch.Tensor | None, ) -> None: if not hasattr(self.attn_backend, "reset_current_inputs"): return if num_extends > 0: self.attn_backend.reset_current_inputs( self.input_buffers.req_pool_indices_buf[:num_extends], mamba_pool_indices[:num_extends], ) if not has_retract: return # FlatForwardOperation places prefill rows first. Restrict recovery to # the decode suffix so a prefix-cache COW on a prefill row cannot be # mistaken for a retracted decode. decode_begin = num_extends decode_count = bs - decode_begin if decode_count <= 0: return skipped_decode_mask = ( skipped_layerwise_cow_mask[decode_begin:bs] if skipped_layerwise_cow_mask is not None else None ) retract_mask = self._mamba_retract_reset_mask( mamba_cow_src[decode_begin:bs], decode_count, skipped_decode_mask, ) self.attn_backend.reset_current_inputs( self.input_buffers.req_pool_indices_buf[decode_begin:bs][retract_mask], mamba_pool_indices[decode_begin:bs][retract_mask], ) def execute_forward_op( self, forward_op, sampling_params_list: list[SamplingParams], dp_global_num_tokens=None, dp_global_bs=None, dp_all_decode_or_idle: bool = False, dp_all_extend: bool = False, grammar_inputs=None, multimodal_context=None, capture_next_input_ids: bool = False, ) -> ModelExecutionResult: num_extends = forward_op.num_extends() total_tokens = sum(forward_op.input_lengths) self._active_multimodal_context = multimodal_context self._active_positions_override = None timing_enabled = LOG_MM_TIMING timing_start = time.perf_counter() if timing_enabled else 0.0 input_fill_ms = 0.0 mrope_ms = 0.0 sampling_prep_ms = 0.0 forward_step_ms = 0.0 output_d2h_ms = 0.0 graph_capable = False graph_padded_bs = 0 with nvtx_range("pre_fill_setup", color="orange"): has_retract = self._contains_retracted_decode(forward_op) # Wait for previous iteration's runtime state updates # (future_input_map, valid_cache_lengths) on execution_stream to # complete before reading them. torch.cuda.current_stream().wait_stream(self.execution_stream) self.execution_stream.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(self.execution_stream): bs = len(forward_op.request_ids) # Outside the graph: in-graph sites only OR into the flag buffer. self.nan_guard.reset(bs) # Mirror the full group's flat table into req_to_page (flat+spec). flat_block_tables = flat_block_tables_from_forward_op( forward_op, device=self.device, num_reqs=bs, ) self._mirror_flat_full_table_into_req_to_page(forward_op, flat_block_tables) decode_input_ids = self.input_buffers.fill_input_buffers( forward_op=forward_op, runtime_states=self.runtime_states, req_to_page=self.req_to_page, total_tokens=total_tokens, ) if timing_enabled: input_fill_done = time.perf_counter() input_fill_ms = (input_fill_done - timing_start) * 1000.0 skipped_layerwise_cow_mask = self._skip_completed_layerwise_mamba_cow( forward_op, bs ) mrope_start = time.perf_counter() if timing_enabled else 0.0 self._active_positions_override = self._build_mrope_positions_override( forward_op=forward_op, multimodal_context=multimodal_context, total_tokens=total_tokens, ) if timing_enabled: mrope_ms = (time.perf_counter() - mrope_start) * 1000.0 forward_mode = ForwardMode.from_num_extends(num_extends, bs) if num_extends <= 0: self._prev_decode_bs = bs if self.runtime_states.mamba_pool is not None and ( num_extends > 0 or has_retract ): mamba_pool_indices = self.input_buffers.mamba_pool_indices_buf[:bs] mamba_cow_src = self.input_buffers.mamba_cow_src_indices_buf[:bs] self.runtime_states.copy_mamba_states( mamba_pool_indices, mamba_cow_src, bs ) if num_extends > 0: self.runtime_states.zero_mamba_states( mamba_pool_indices, mamba_cow_src, self.input_buffers.extend_prefix_lens_buf[:num_extends], num_extends, ) self._reset_mamba_current_inputs( num_extends=num_extends, bs=bs, has_retract=has_retract, mamba_pool_indices=mamba_pool_indices, mamba_cow_src=mamba_cow_src, skipped_layerwise_cow_mask=skipped_layerwise_cow_mask, ) grammar_completion = None if total_tokens == 0: # Fully prefix-cached prefill: no tokens to process. output_tokens = torch.zeros(0, dtype=torch.int32, device=self.device) output_lengths = torch.zeros(bs, dtype=torch.int32, device=self.device) output_logprobs = None else: gather_ids = None if num_extends > 0: num_decodes = bs - num_extends if self.drafter is not None and num_decodes > 0: # MIXED + spec: prefill rows pruned to last token, # decode block kept full at verify width. num_decode_tokens = num_decodes * self.config.spec_num_tokens num_prefill_tokens = total_tokens - num_decode_tokens gather_ids = torch.empty( num_extends + num_decode_tokens, dtype=torch.int64, device=self.device, ) gather_ids[:num_extends] = ( torch.cumsum( self.input_buffers.input_lengths_buf[:num_extends], dim=0, ) - 1 ) gather_ids[num_extends:] = torch.arange( num_prefill_tokens, total_tokens, device=self.device, dtype=torch.int64, ) else: # EXTEND, MIXED non-spec, or EXTEND + spec: last token # per request via cumsum. gather_ids = ( torch.cumsum( self.input_buffers.input_lengths_buf[:bs], dim=0 ) - 1 ) ctx = ForwardContext( attn_backend=self.attn_backend, token_to_kv_pool=self.token_to_kv_pool, req_to_page=self.req_to_page, bs=bs, num_extends=num_extends, input_num_tokens=total_tokens, forward_mode=forward_mode, capture_hidden_mode=( CaptureHiddenMode.FULL if self.drafter is not None else CaptureHiddenMode.NULL ), gather_ids=gather_ids, decode_input_ids=decode_input_ids, ) if self.config.data_parallel_size > 1: if dp_global_num_tokens is None: raise RuntimeError( "DP forward metadata must be gathered on CPU by " "the event loop before model execution." ) ctx.global_num_tokens = dp_global_num_tokens ctx.global_bs = dp_global_bs ctx.all_decode_or_idle = dp_all_decode_or_idle ctx.all_extend = dp_all_extend with nvtx_range("sampling_prep", color="yellow"): sampling_start = time.perf_counter() if timing_enabled else 0.0 sampling_info = self._build_sampling_info(bs, sampling_params_list) grammar_completion = setup_grammar_step( sampling_info=sampling_info, bs=bs, is_spec_decode=self.drafter is not None and num_extends < bs, spec_num_tokens=self.config.spec_num_tokens or 1, grammar_inputs=grammar_inputs, grammar_runtime=self.grammar_runtime, input_ids_buf=self.input_buffers.input_ids_buf, grammar_backend=self.config.grammar_backend, ) extend_with_prefix = num_extends > 0 and any( forward_op.extend_prefix_lens ) # Flip detection + per-slot scalar scatter + backend-owned # RNG state refill. Runs OUTSIDE the CUDA graph. Generators # are now backend-internal (pool-indexed, seeded on flip # from sp.seed), so the event loop no longer threads them # through. self.sampling_backend.prepare_step( request_ids=forward_op.request_ids, request_pool_indices=forward_op.request_pool_indices, sampling_params_list=sampling_params_list, num_tokens_per_req=self.config.output_length, ) if timing_enabled: sampling_prep_ms = ( time.perf_counter() - sampling_start ) * 1000.0 with nvtx_range( f"forward_step ext={num_extends} dec={bs - num_extends}", color="blue", ): mamba_kwargs = ( { "mamba_pool_indices": self.input_buffers.mamba_pool_indices_buf[ :bs ], "mamba_cow_src_indices": self.input_buffers.mamba_cow_src_indices_buf[ :bs ], "mamba_branching_seqlens": self.input_buffers.mamba_branching_seqlens_buf[ :bs ], "mamba_track_pool_indices": self.input_buffers.mamba_track_pool_indices_buf[ :bs ], } if self.input_buffers.has_mamba else {} ) paged_cache_block_tables = paged_cache_block_tables_from_forward_op( forward_op, device=self.device, num_reqs=bs, ) # flat_block_tables computed once in pre_fill above. ( paged_cache_block_table_base_offsets, _paged_cache_block_table_base_offset_max, ) = paged_cache_block_table_base_offsets_from_forward_op( forward_op, device=self.device, num_reqs=bs, ) self._log_dp_sampling_route(bs, ctx) forward_step_start = 0.0 if timing_enabled: graph_capable = self.forward_step.can_run(bs, ctx) graph_padded_bs = ( self.forward_step.padded_bs(bs, ctx) if graph_capable else bs ) forward_step_start = time.perf_counter() output_tokens, output_lengths, output_logprobs = self.forward_step( bs=bs, ctx=ctx, sampling_info=sampling_info, req_to_page=self.req_to_page, extend_with_prefix=extend_with_prefix, extend_prefix_lens=self.input_buffers.extend_prefix_lens_buf[ :num_extends ], extend_prefix_lens_cpu=self.input_buffers.extend_prefix_lens_cpu[ :num_extends ], extend_seq_lens=self.input_buffers.extend_seq_lens_buf[ :num_extends ], extend_seq_lens_cpu=self.input_buffers.extend_seq_lens_cpu[ :num_extends ], 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, ) if timing_enabled: forward_step_ms = ( time.perf_counter() - forward_step_start ) * 1000.0 if self.config.spec_algo == "DFLASH": # Clamp the committed-length delta so no request grows past # context_len. Done here (outside the graph) so it reaches # both _update_runtime_state and the scheduler page # reservation; see _clamp_committed_to_context_len. output_lengths = self._clamp_committed_to_context_len( output_lengths, num_extends, bs ) # Update runtime state on execution_stream (NOT in the CUDA graph). self._update_runtime_state( req_pool_indices=self.input_buffers.req_pool_indices_buf[:bs], output_tokens=output_tokens, accept_lengths=output_lengths, input_lengths=self.input_buffers.input_lengths_buf[:bs], num_extends=num_extends, ) self._snapshot_mamba_checkpoints( output_lengths, bs, num_extends, ) with nvtx_range("output_d2h", color="green"): output_d2h_start = time.perf_counter() if timing_enabled else 0.0 next_input_ids = None if ( capture_next_input_ids and self.drafter is not None and num_extends > 0 ): next_input_ids = self.runtime_states.future_input_map.index_select( 0, self.input_buffers.req_pool_indices_buf[:num_extends] ).to("cpu", non_blocking=True) # Defensive clamp into the valid vocab range (kept from the # pre-pack path). An out-of-range token id -- e.g. a stale/corrupt # value surfaced by the intermittent spec-decode decode-state race # -- would otherwise reach the detokenizer, whose HF # tokenizer.decode raises a fatal OverflowError on ids outside # [0, vocab) and tears down the whole server process tree. # It must run on-GPU *before* the non_blocking D2H: clamping the # CPU result afterwards would race the in-flight copy. In-place # (clamp_) so output_tokens keeps aliasing _output_pack_buf and # the get_packed_output_d2h data_ptr fast-path still fires -- and # in-place on the forward's inference tensors is only legal inside # inference mode, so re-enter it (maybe_inference_mode mirrors the # forward and reduces to no_grad when inference mode is disabled, # where output_tokens isn't an inference tensor anyway). vocab_size = self.runtime_states.vocab_size with maybe_inference_mode(): output_tokens.clamp_(0, vocab_size - 1) packed = self.sampling_backend.get_packed_output_d2h( output_tokens, output_lengths ) if packed is not None: output_tokens, output_lengths = packed else: output_tokens = output_tokens.to("cpu", non_blocking=True) output_lengths = output_lengths.to("cpu", non_blocking=True) if output_logprobs is not None: output_logprobs = output_logprobs.to("cpu", non_blocking=True) output_nan_flags = self.nan_guard.flags_cpu copy_event = torch.cuda.Event() copy_event.record() if timing_enabled: output_d2h_ms = (time.perf_counter() - output_d2h_start) * 1000.0 if timing_enabled and ( num_extends > 0 or self.log_step < 64 or self.log_step % 100 == 0 ): has_mm = ( multimodal_context is not None and multimodal_context.has_inputs() ) mm_count = 0 mm_delta_count = 0 if has_mm: for mm_input in multimodal_context.mm_inputs: if mm_input is None: continue mm_count += 1 if mm_input.mrope_position_delta is not None: mm_delta_count += 1 logger.info( "mm_timing forward_execute_ms total=%.3f input_fill=%.3f " "mrope=%.3f sampling=%.3f forward_step=%.3f output_d2h=%.3f " "mode=%s bs=%s total_tokens=%s graph=%s padded_bs=%s " "has_mm=%s mm_count=%s mm_delta_count=%s", (time.perf_counter() - timing_start) * 1000.0, input_fill_ms, mrope_ms, sampling_prep_ms, forward_step_ms, output_d2h_ms, forward_mode.name, bs, total_tokens, graph_capable, graph_padded_bs, has_mm, mm_count, mm_delta_count, ) return ModelExecutionResult( output_tokens=output_tokens, output_lengths=output_lengths, output_logprobs=output_logprobs, copy_event=copy_event, grammar_completion=grammar_completion, next_input_ids=next_input_ids, output_nan_flags=output_nan_flags, ) def write_remote_spec_candidate_ids( self, req_pool_idx: int, candidate_ids: list[int] ) -> None: # Remote spec candidates are CPU materialized; enqueue the H2D copy and # future_input_map update on execution_stream. The next forward's input # prep already waits on execution_stream before reading runtime state. with torch.cuda.stream(self.execution_stream): self.runtime_states.write_remote_spec_candidate_ids( req_pool_idx, candidate_ids ) def _expand_mrope_from_input(self, mm_input, seq_len: int) -> torch.Tensor: # Cache delta expansion for retracted/chunked requests. if mm_input.mrope_position_delta_repeated_cache is None: mm_input.mrope_position_delta_repeated_cache = ( (mm_input.mrope_position_delta - 1).flatten().unsqueeze(0).repeat(3, 1) ) return mm_input.mrope_position_delta_repeated_cache + seq_len @staticmethod def _mrope_delta_scalar(mm_input) -> int: delta = getattr(mm_input, "mrope_position_delta_scalar", None) if delta is not None: return int(delta) tensor = getattr(mm_input, "mrope_position_delta", None) if tensor is None: return 0 delta = int(tensor.flatten()[0].item()) mm_input.mrope_position_delta_scalar = delta return delta def _build_decode_mrope_positions_override( self, forward_op, mm_inputs, total_tokens: int, ) -> torch.Tensor: if ( self._mrope_decode_deltas_cpu is None or self._mrope_decode_deltas_buf is None ): raise RuntimeError( "M-RoPE decode buffers were not initialized for this model" ) base_positions = self.input_buffers.positions_buf[:total_tokens] # Ping-pong the pinned host staging buffer (see __init__): the previous # step's non_blocking H2D copy may still be reading the other buffer. cpu_staging = self._mrope_decode_deltas_cpu[self._mrope_decode_deltas_cpu_idx] self._mrope_decode_deltas_cpu_idx ^= 1 token_deltas_cpu = cpu_staging[:total_tokens] offset = 0 has_nonzero_delta = False for batch_idx, input_len in enumerate(forward_op.input_lengths): input_len = int(input_len) if input_len <= 0: continue delta = 0 mm_input = mm_inputs[batch_idx] if batch_idx < len(mm_inputs) else None # Honor scalar-only deltas: an upstream payload may set # mrope_position_delta_scalar while leaving the tensor field # mrope_position_delta as None (positions precomputed upstream). # _mrope_delta_scalar handles scalar, tensor, and the absent case # (returns 0), so call it whenever an mm_input is present. if mm_input is not None: delta = self._mrope_delta_scalar(mm_input) has_nonzero_delta = has_nonzero_delta or delta != 0 token_deltas_cpu[offset : offset + input_len].fill_(delta) offset += input_len if offset != total_tokens: token_deltas_cpu[offset:total_tokens].zero_() if has_nonzero_delta: token_deltas = self._mrope_decode_deltas_buf[:total_tokens] token_deltas.copy_(token_deltas_cpu, non_blocking=True) mrope_base = base_positions + token_deltas else: mrope_base = base_positions self.input_buffers.mrope_positions_buf[:, :total_tokens].copy_( mrope_base.unsqueeze(0).expand(3, -1) ) return self.input_buffers.mrope_positions_buf[:, :total_tokens] def _build_mrope_positions_override( self, forward_op, multimodal_context, total_tokens: int, ) -> torch.Tensor | None: if not self.config.model_is_mrope or total_tokens == 0: return None is_prefill = forward_op.num_extends() > 0 base_positions = self.input_buffers.positions_buf[:total_tokens] mm_inputs = ( multimodal_context.mm_inputs if multimodal_context is not None and multimodal_context.has_inputs() else [] ) if not is_prefill: return self._build_decode_mrope_positions_override( forward_op=forward_op, mm_inputs=mm_inputs, total_tokens=total_tokens, ) pos_chunks = torch.split(base_positions, list(forward_op.input_lengths), dim=0) mrope_chunks = [] for batch_idx, base_chunk in enumerate(pos_chunks): mm_input = mm_inputs[batch_idx] if batch_idx < len(mm_inputs) else None # Fall back to linear only when there is neither a per-token mrope table # nor a transferred scalar delta. A decode-only mm_input may carry just # the delta (post-image decode positions = base+delta); it must skip the # fallback and take the base+delta branch below. if mm_input is None or ( mm_input.mrope_positions is None and mm_input.mrope_position_delta is None ): mrope_chunks.append(base_chunk.unsqueeze(0).expand(3, -1)) continue if ( is_prefill and mm_input.mrope_positions is not None and batch_idx < len(forward_op.extend_prefix_lens) ): start = int(forward_op.extend_prefix_lens[batch_idx]) end = start + int(forward_op.input_lengths[batch_idx]) positions = mm_input.mrope_positions[:, start:end] if positions.numel() != 0: mrope_chunks.append( positions.to(device=self.device, dtype=torch.int64) ) continue if base_chunk.numel() == 1: seq_len = int(base_chunk[-1].item()) + 1 mrope_chunks.append( self._expand_mrope_from_input(mm_input, seq_len).to( device=self.device, dtype=torch.int64 ) ) continue delta = mm_input.mrope_position_delta if delta is None: delta = torch.zeros(1, dtype=torch.int64) delta = delta.flatten()[0].to(device=self.device, dtype=torch.int64) # Decode positions need (mrope_delta - 1) + seq_len. positions_buf # already stores the per-token zero-based position (seq_len - 1 for # decode), so this is the same value without a GPU-to-CPU sync. mrope_chunks.append((base_chunk + delta).unsqueeze(0).expand(3, -1)) mrope_positions = torch.cat(mrope_chunks, dim=1).contiguous() self.input_buffers.mrope_positions_buf[:, :total_tokens].copy_(mrope_positions) return self.input_buffers.mrope_positions_buf[:, :total_tokens]