import logging import math from typing import List, Optional import torch from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func from sglang.kernels.ops.speculative.dflash import ( _compute_dflash_accept_bonus_triton_unchecked, _prepare_dflash_draft_block_unchecked, ) from sglang.srt.distributed import get_tp_group from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.scheduler import GenerationBatchResult from sglang.srt.managers.tp_worker import TpModelWorker from sglang.srt.model_executor.cuda_graph_config import Backend from sglang.srt.model_executor.forward_batch_info import ( CaptureHiddenMode, ForwardBatch, ForwardMode, compute_position, ) from sglang.srt.server_args import ServerArgs from sglang.srt.speculative.base_spec_worker import BaseSpecWorker from sglang.srt.speculative.dflash_info import DFlashVerifyInput from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 from sglang.srt.speculative.dflash_utils import ( apply_dflash_verify_logits_adjustments, can_dflash_use_fused_qkv_proj, compute_dflash_correct_drafts_and_bonus, compute_dflash_sampling_correct_drafts_and_bonus, is_dflash_sampling_verify_available, parse_dflash_draft_config, ) from sglang.srt.speculative.draft_worker_common import ( build_block_pos_offsets, build_draft_tp_worker, make_draft_block_spec_info, make_draft_input_v2, make_draft_sampler_capture_hook, ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func from sglang.srt.utils import get_available_gpu_memory, is_cuda, is_hip, is_npu _is_npu = is_npu() logger = logging.getLogger(__name__) _FusedKVMaterializeHelper = None def _get_fused_kv_materialize_helper(): global _FusedKVMaterializeHelper if _FusedKVMaterializeHelper is None: from sglang.kernels.ops.speculative.fused_kv_materialize import ( FusedKVMaterializeHelper, ) _FusedKVMaterializeHelper = FusedKVMaterializeHelper return _FusedKVMaterializeHelper class _DflashDraftSampler: """Capture-safe greedy argmax over the target LM head, run inside the draft cuda graph so the draft sampling is captured and counted in fwd_occupancy. DFLASH's draft has no head of its own; it borrows the target `lm_head`. tp=1 / no-added-vocab only; TP>1 stays eager in the worker. """ def __init__(self, *, weight, block_size, num_org, org_vocab_start, max_bs): self.weight = weight self.block_size = int(block_size) self.num_org = int(num_org) self.org_vocab_start = int(org_vocab_start) # Proposed draft tokens: written in-graph, read by the worker after replay. self.out = torch.empty( (int(max_bs) * (self.block_size - 1),), dtype=torch.int64, device=weight.device, ) def __call__(self, hidden_states, input_ids=None): # draft tokens are block positions 1: (pos 0 is the seeded bonus token) bs = hidden_states.shape[0] // self.block_size hs = hidden_states.view(bs, self.block_size, -1)[:, 1:, :].reshape( -1, hidden_states.shape[-1] ) if hs.dtype != self.weight.dtype: hs = hs.to(self.weight.dtype) logits = torch.matmul(hs, self.weight[: self.num_org].T) tokens = torch.argmax(logits, dim=-1).to(torch.long) if self.org_vocab_start: tokens += self.org_vocab_start self.out[: tokens.shape[0]].copy_(tokens) class DFlashWorkerV2(BaseSpecWorker): """DFLASH speculative decoding worker (spec-v2). Drives both overlap and non-overlap scheduling, same as EAGLE: the scheduler runs it synchronously when overlap is disabled. """ def __init__( self, server_args: ServerArgs, gpu_id: int, tp_rank: int, dp_rank: Optional[int], moe_ep_rank: int, attn_cp_rank: int, moe_dp_rank: int, nccl_port: int, target_worker: TpModelWorker, ): self.server_args = server_args self.gpu_id = gpu_id self.tp_rank = tp_rank self.dp_rank = dp_rank self.moe_ep_rank = moe_ep_rank self.attn_cp_rank = attn_cp_rank self.moe_dp_rank = moe_dp_rank self.nccl_port = nccl_port self._target_worker = target_worker self.model_runner = target_worker.model_runner self._need_mamba_verify_commit = False self.page_size = server_args.page_size # Normalized in arg_groups.speculative_hook.handle_speculative_decoding. self.draft_window_size: Optional[int] = ( server_args.speculative_draft_window_size ) self.use_compact_draft_cache = self.draft_window_size is not None self.device = target_worker.device self._warned_sampling_fallback = False self._logged_first_verify = False bundle = build_draft_tp_worker( server_args=server_args, gpu_id=gpu_id, tp_rank=tp_rank, dp_rank=dp_rank, moe_ep_rank=moe_ep_rank, attn_cp_rank=attn_cp_rank, moe_dp_rank=moe_dp_rank, nccl_port=nccl_port, target_model_config=target_worker.model_runner.model_config, algo_label="DFLASH", ) self._draft_worker = bundle.draft_worker self.draft_model_runner = bundle.draft_model_runner self._draft_sampler = None self.draft_model = bundle.draft_model draft_config = parse_dflash_draft_config( draft_hf_config=self.draft_model_runner.model_config.hf_config ) if server_args.speculative_num_draft_tokens is None: # Should not happen (ServerArgs should have inferred it), but keep a fallback. self.block_size = int(draft_config.resolve_block_size(default=16)) else: self.block_size = int(server_args.speculative_num_draft_tokens) model_block_size = draft_config.block_size if model_block_size is None: model_block_size = getattr(self.draft_model, "block_size", None) if model_block_size is not None and int(model_block_size) != int( self.block_size ): logger.warning( "DFLASH block size mismatch: using speculative_num_draft_tokens=%s but draft config block_size=%s.", self.block_size, model_block_size, ) self.speculative_num_draft_tokens = int(self.block_size) self._mask_token = draft_config.mask_token self._mask_token_id_override = draft_config.mask_token_id self._mask_token_id = self._resolve_mask_token_id( mask_token=self._mask_token, mask_token_id=self._mask_token_id_override, ) if self.tp_rank == 0: logger.info( "Initialized DFLASH draft runner. attention_backend=%s, model=%s, block_size=%s, draft_window_size=%s, compact_cache=%s", bundle.resolved_attention_backend, self.draft_model.__class__.__name__, self.block_size, self.draft_window_size, self.use_compact_draft_cache, ) logger.info( "DFLASH draft runner ready. mask_token=%s, mask_token_id=%s, mask_token_id_override=%s", self._mask_token, self._mask_token_id, self._mask_token_id_override, ) self._block_pos_offsets = build_block_pos_offsets( length=self.block_size, device=self.device ) self._draft_block_ids_buf: Optional[torch.Tensor] = None # [cap_bs, block_size] self._draft_block_positions_buf: Optional[torch.Tensor] = ( None # [cap_bs, block_size] ) self._draft_block_tokens_buf: Optional[torch.Tensor] = ( None # [cap_bs, block_size] ) self._draft_verify_out_cache_loc_buf: Optional[torch.Tensor] = ( None # [cap_bs, block_size] ) self._draft_block_end_buf: Optional[torch.Tensor] = None # [cap_bs] self._draft_seq_lens_cpu_buf: Optional[torch.Tensor] = None # [cap_bs] on CPU self._draft_block_spec_info = make_draft_block_spec_info( draft_token_num=int(self.block_size), device=self.device ) self._draft_greedy_gathered_max_buf: Optional[torch.Tensor] = None self._draft_greedy_gathered_ids_buf: Optional[torch.Tensor] = None self._draft_greedy_gather_cap: int = 0 self._draft_greedy_local_max_buf: Optional[torch.Tensor] = None self._draft_greedy_local_arg_buf: Optional[torch.Tensor] = None self._draft_greedy_local_cap: int = 0 self._draft_greedy_best_rank_buf: Optional[torch.Tensor] = None self._draft_greedy_rank_index_buf: Optional[torch.Tensor] = None self._draft_greedy_selected_ids_buf: Optional[torch.Tensor] = None self._draft_greedy_index_cap: int = 0 self._use_fused_kv_materialize = is_cuda() or is_hip() self._fused_kv_helper: Optional[object] = None if self._use_fused_kv_materialize: self._init_fused_kv_helper() supports_gpu_triton = is_cuda() or is_hip() self._use_triton_prepare_block = supports_gpu_triton self._use_triton_accept_bonus = supports_gpu_triton self._accept_bonus_buffer_cap: int = 0 self._accept_bonus_buffer_slot: int = 0 self._accept_len_buf: Optional[torch.Tensor] = None self._commit_lens_bufs: List[torch.Tensor] = [] self._bonus_id_bufs: List[torch.Tensor] = [] self._out_tokens_bufs: List[torch.Tensor] = [] self._new_seq_lens_bufs: List[torch.Tensor] = [] @property def target_worker(self) -> TpModelWorker: return self._target_worker @property def draft_worker(self): # DFLASH drives the draft model through a plain TpModelWorker: the # draft KV is materialized from target hidden states, so there is no # EagleDraftWorkerBase draft/draft_extend split to wrap it in. return self._draft_worker @property def spec_v2_attn_backends(self) -> tuple: # Every attn backend a spec_v2 forward touches; consumed by # decide_needs_cpu_seq_lens to gate the seq_lens_cpu D2H. return ( self._target_worker.model_runner.attn_backend, self.draft_model_runner.attn_backend, ) def alloc_memory_pool( self, memory_pool_config=None, req_to_token_pool=None, token_to_kv_pool_allocator=None, ): # Without draft windowing, the draft worker aliases the target # request->token mapping and allocation state. With draft windowing # enabled, the draft worker keeps a private compact req->token table # over the same global KV index space, so radix-cache/prefix-hit KV # remains reusable while draft attention sees only the recent window. self._draft_worker.alloc_memory_pool( memory_pool_config=memory_pool_config, req_to_token_pool=( None if self.use_compact_draft_cache else req_to_token_pool ), token_to_kv_pool_allocator=token_to_kv_pool_allocator, ) def init_attention_backends(self): self._draft_worker.init_attention_backends() self._need_mamba_verify_commit = ( self.model_runner.mambaish_config is not None and hasattr( self.model_runner.attn_backend, "update_mamba_state_after_mtp_verify", ) ) def init_cuda_graphs(self): capture_decode_cuda_graph = ( self.server_args.cuda_graph_config.decode.backend != Backend.DISABLED ) if is_cuda() and capture_decode_cuda_graph: available_mem = get_available_gpu_memory(self.device, self.gpu_id) if available_mem < 1.0: capture_decode_cuda_graph = False logger.warning( "Disable DFLASH draft cuda graph because only %.2f GB GPU " "memory is available after target backend initialization.", available_mem, ) if capture_decode_cuda_graph: # Must run before capture so the draft graph folds the head in. self._draft_sampler = self._maybe_build_draft_sampler() if self._draft_sampler is not None: self.draft_model_runner.capture_tail_hooks.append( make_draft_sampler_capture_hook(self._draft_sampler) ) self._draft_worker.init_cuda_graphs( capture_decode_cuda_graph=capture_decode_cuda_graph ) def _maybe_build_draft_sampler(self): def _eager(reason): if self.tp_rank == 0: logger.info("DFLASH draft greedy head kept eager (reason=%s).", reason) return None if get_tp_group().world_size != 1: return _eager("tp>1") if self.block_size <= 1: return _eager("block_size<=1") target_model = self._target_worker.model_runner.model lm_head = getattr(target_model, "lm_head", None) if lm_head is None or not hasattr(lm_head, "weight"): return _eager("no target lm_head") if not torch.is_floating_point(lm_head.weight): # Quantized lm_head (FP8/INT) would break the static matmul. return _eager("quantized lm_head") if not hasattr(lm_head, "shard_indices"): num_org = int(lm_head.weight.shape[0]) org_vocab_start = 0 else: shard = lm_head.shard_indices if int(shard.num_added_elements) != 0: return _eager("added vocab") num_org = int(shard.num_org_elements) org_vocab_start = int(shard.org_vocab_start_index) if self.tp_rank == 0: logger.info("DFLASH draft greedy head folded into the draft cuda graph.") return _DflashDraftSampler( weight=lm_head.weight, block_size=self.block_size, num_org=num_org, org_vocab_start=org_vocab_start, max_bs=max(self.server_args.cuda_graph_config.decode.bs), ) def _init_fused_kv_helper(self) -> None: """Initialize the fused KV materialization helper with pre-stacked weights.""" try: layers = self.draft_model.layers fused_disable_reason: Optional[str] = None if len(layers) == 0: fused_disable_reason = "no layers found" elif not getattr(self.draft_model, "supports_fused_context_kv", True): fused_disable_reason = "draft model does not support fused context KV" if fused_disable_reason is not None: if self.tp_rank == 0: logger.info( "DFLASH fused KV materialization disabled: %s", fused_disable_reason, ) self._use_fused_kv_materialize = False self._fused_kv_helper = None return for layer_idx, layer in enumerate(layers): attn = layer.self_attn eligible, reason = can_dflash_use_fused_qkv_proj(attn.qkv_proj) if not eligible: fused_disable_reason = f"{reason}: layer={layer_idx}" break # Keep semantics aligned with set_kv_buffer scaling behavior. k_scale = getattr(attn.attn, "k_scale", None) v_scale = getattr(attn.attn, "v_scale", None) if k_scale is not None and not math.isclose(float(k_scale), 1.0): fused_disable_reason = ( "non-unit k_scale is not supported for fused KV path: " f"layer={layer_idx}, k_scale={k_scale}" ) break if v_scale is not None and not math.isclose(float(v_scale), 1.0): fused_disable_reason = ( "non-unit v_scale is not supported for fused KV path: " f"layer={layer_idx}, v_scale={v_scale}" ) break rope_is_neox_style = bool( getattr(attn.rotary_emb, "is_neox_style", True) ) if not rope_is_neox_style: fused_disable_reason = ( "non-neox RoPE is not supported for fused KV path: " f"layer={layer_idx}, rope_is_neox_style={rope_is_neox_style}" ) break if fused_disable_reason is not None: if self.tp_rank == 0: logger.info( "DFLASH fused KV materialization disabled: %s", fused_disable_reason, ) self._use_fused_kv_materialize = False self._fused_kv_helper = None return FusedKVMaterializeHelper = _get_fused_kv_materialize_helper() first_attn = layers[0].self_attn rotary_emb = first_attn.rotary_emb self._fused_kv_helper = FusedKVMaterializeHelper( layers=layers, rotary_emb=rotary_emb, num_kv_heads=first_attn.num_kv_heads, head_dim=first_attn.head_dim, device=self.device, max_position_hint=self.target_worker.model_runner.model_config.context_len + int(self.block_size), ) if self.tp_rank == 0: logger.info( "DFLASH fused KV materialization enabled. " "n_layers=%d, num_kv_heads=%d, head_dim=%d", len(layers), first_attn.num_kv_heads, first_attn.head_dim, ) except Exception as e: logger.warning( "DFLASH fused KV initialization failed, falling back to sequential path: %s", e, ) self._use_fused_kv_materialize = False self._fused_kv_helper = None def _ensure_draft_block_buffers(self, bs: int) -> None: cap = ( 0 if self._draft_block_ids_buf is None else int(self._draft_block_ids_buf.shape[0]) ) if cap >= int(bs): return new_cap = max(int(bs), cap * 2 if cap > 0 else int(bs)) device = self.device block_size = int(self.block_size) self._draft_block_ids_buf = torch.empty( (new_cap, block_size), dtype=torch.long, device=device ) self._draft_block_positions_buf = torch.empty( (new_cap, block_size), dtype=torch.int64, device=device ) self._draft_block_tokens_buf = torch.empty( (new_cap, block_size), dtype=torch.long, device=device ) self._draft_verify_out_cache_loc_buf = torch.empty( (new_cap, block_size), dtype=torch.int64, device=device ) self._draft_block_end_buf = torch.empty( (new_cap,), dtype=torch.int32, device=device ) self._draft_seq_lens_cpu_buf = torch.empty( (new_cap,), dtype=torch.int32, device="cpu" ) def __getattr__(self, name): # Delegate anything not implemented yet to the target worker. Guard # the backing field so a lookup before __init__ sets it raises # AttributeError instead of recursing through the property. if name == "_target_worker": raise AttributeError(name) return getattr(self.target_worker, name) def clear_cache_pool(self): # The target worker owns the shared KV allocator/cache. For the compact # sliding-window path, the draft req->token view is rebuilt from committed # target state before each draft forward, so there is nothing persistent # to flush here. pass def _gather_req_to_token_masked( self, *, req_to_token: torch.Tensor, req_pool_indices: torch.Tensor, pos2d: torch.Tensor, mask: torch.Tensor, context: str, ) -> torch.Tensor: if pos2d.ndim != 2: raise RuntimeError( f"{context} expected 2D positions, got shape={tuple(pos2d.shape)}." ) if mask.shape != pos2d.shape: raise RuntimeError( f"{context} mask/position shape mismatch: {tuple(mask.shape)} vs {tuple(pos2d.shape)}." ) if req_pool_indices.dtype != torch.int64: req_pool_indices = req_pool_indices.to(torch.int64) if mask.dtype != torch.bool: mask = mask.to(torch.bool) table_width = int(req_to_token.shape[1]) if table_width <= 0: if bool(mask.any().item()): raise RuntimeError( f"{context} req_to_token table is empty but gather mask is non-empty." ) return torch.empty((0,), dtype=torch.int64, device=self.device) # Only the masked-off rectangular padding can be out of range in the normal # ragged-batch case. Replace those don't-care columns with a valid in-range # position before the gather so the kernel only sees real positions. safe_pos2d = pos2d.masked_fill(~mask, 0) return req_to_token[req_pool_indices[:, None], safe_pos2d][mask].to(torch.int64) def _gather_req_to_token_segments( self, *, req_to_token: torch.Tensor, req_pool_indices: torch.Tensor, start: torch.Tensor | None, lengths: torch.Tensor, ) -> torch.Tensor: lengths = lengths.to(torch.int64) if lengths.numel() == 0: return torch.empty((0,), dtype=torch.int64, device=self.device) max_len = int(lengths.max().item()) if max_len <= 0: return torch.empty((0,), dtype=torch.int64, device=self.device) if req_pool_indices.dtype != torch.int64: req_pool_indices = req_pool_indices.to(torch.int64) offsets = torch.arange( max_len, device=self.device, dtype=torch.int64 ).unsqueeze(0) if start is None: pos2d = offsets.expand(req_pool_indices.shape[0], -1) else: pos2d = start.to(torch.int64).unsqueeze(1) + offsets mask = offsets < lengths.unsqueeze(1) return self._gather_req_to_token_masked( req_to_token=req_to_token, req_pool_indices=req_pool_indices, pos2d=pos2d, mask=mask, context="DFLASH req_to_token segment gather", ) def _compute_compact_draft_seq_lens(self, seq_lens: torch.Tensor) -> torch.Tensor: assert self.draft_window_size is not None visible_lens = torch.clamp( seq_lens.to(dtype=torch.int32, device=self.device), max=int(self.draft_window_size), ) if self.page_size <= 1: return visible_lens # Paged FA backends derive the page table from local token positions, so the # compact suffix must start on a page boundary. Keep up to page_size - 1 extra # tokens on the left to preserve valid local page structure. seq_lens_i64 = seq_lens.to(torch.int64) visible_lens_i64 = visible_lens.to(torch.int64) visible_start = seq_lens_i64 - visible_lens_i64 aligned_start = visible_start - torch.remainder(visible_start, self.page_size) return (seq_lens_i64 - aligned_start).to(torch.int32) def _resolve_mask_token_id( self, *, mask_token: str, mask_token_id: Optional[int] = None ) -> int: if not isinstance(mask_token, str) or not mask_token: raise ValueError( f"DFLASH mask_token must be a non-empty string, got {mask_token!r}." ) vocab_size = int(self.target_worker.model_runner.model_config.vocab_size) if mask_token_id is not None: resolved_id = int(mask_token_id) if resolved_id >= vocab_size: raise ValueError( "DFLASH mask_token_id is outside the target vocab size. " f"mask_token_id={resolved_id}, vocab_size={vocab_size}. " f"This likely means mask_token={mask_token!r} requires vocab expansion beyond the model's embedding size. " "SGLang does not support resizing target embeddings for DFLASH yet." ) tokenizer = getattr(self.target_worker, "tokenizer", None) if tokenizer is not None: token_id_from_vocab = tokenizer.get_vocab().get(mask_token, None) if ( token_id_from_vocab is not None and int(token_id_from_vocab) != resolved_id ): raise ValueError( "DFLASH config mismatch: dflash_config.mask_token_id conflicts with tokenizer vocab id " f"for dflash_config.mask_token. mask_token={mask_token!r}, " f"mask_token_id={resolved_id}, tokenizer_vocab_id={int(token_id_from_vocab)}." ) return resolved_id tokenizer = getattr(self.target_worker, "tokenizer", None) if tokenizer is None: raise RuntimeError( "DFLASH requires tokenizer initialization when dflash_config.mask_token_id is not set " "(skip_tokenizer_init is not supported in this mode)." ) resolved_id = None if getattr(tokenizer, "mask_token", None) == mask_token: resolved_id = getattr(tokenizer, "mask_token_id", None) if resolved_id is None: # Prefer checking the explicit vocab mapping first. vocab = tokenizer.get_vocab() resolved_id = vocab.get(mask_token, None) if resolved_id is None: # Mirror the reference DFlash HF demo by adding the mask token to the tokenizer. # This is safe only when the resulting id stays within the target model vocab size. added = tokenizer.add_special_tokens({"mask_token": mask_token}) resolved_id = getattr(tokenizer, "mask_token_id", None) if resolved_id is None: resolved_id = tokenizer.convert_tokens_to_ids(mask_token) if added and self.tp_rank == 0: logger.info( "Added DFLASH mask token to tokenizer. token=%s, mask_token_id=%s, tokenizer_len=%s, model_vocab_size=%s", mask_token, resolved_id, len(tokenizer), vocab_size, ) if resolved_id is None or int(resolved_id) < 0: raise ValueError( "DFLASH requires resolving a mask token id, but it could not be resolved. " f"mask_token={mask_token!r}." ) if resolved_id >= vocab_size: raise ValueError( "DFLASH mask_token_id is outside the target vocab size. " f"mask_token_id={resolved_id}, vocab_size={vocab_size}. " f"This likely means mask_token={mask_token!r} requires vocab expansion beyond the model's embedding size. " "SGLang does not support resizing target embeddings for DFLASH yet." ) return int(resolved_id) def _greedy_sample_from_vocab_parallel_head( self, *, hidden_states: torch.Tensor, lm_head, chunk_size: int = 256, ) -> torch.Tensor: """Greedy argmax over the target LM head in a TP-safe way. We cannot materialize full logits for large vocabularies efficiently, and with TP>1 each rank only owns a shard of the LM head weight. This computes the per-rank max, gathers candidates across TP ranks, and selects the global max. """ if hidden_states.numel() == 0: return torch.empty((0,), dtype=torch.long, device=hidden_states.device) weight = lm_head.weight # [local_vocab_padded, hidden] weight_dtype = weight.dtype num_tokens = int(hidden_states.shape[0]) out_tokens = torch.empty( (num_tokens,), dtype=torch.long, device=hidden_states.device ) def _cast_hs(x: torch.Tensor) -> torch.Tensor: return x if x.dtype == weight_dtype else x.to(weight_dtype) if not hasattr(lm_head, "shard_indices"): for start in range(0, num_tokens, int(chunk_size)): end = min(num_tokens, start + int(chunk_size)) hs = _cast_hs(hidden_states[start:end]) logits = torch.matmul(hs, weight.T) out_tokens[start:end] = torch.argmax(logits, dim=-1).to(torch.long) return out_tokens shard = lm_head.shard_indices tp_group = get_tp_group() tp_size = int(tp_group.world_size) # Valid ranges in the local shard (excluding padding): # base vocab: [0, num_org) # added vocab: [num_org_padded, num_org_padded + num_added) num_org = int(shard.num_org_elements) num_org_padded = int(shard.num_org_elements_padded) num_added = int(shard.num_added_elements) org_vocab_start = int(shard.org_vocab_start_index) added_vocab_start = int(shard.added_vocab_start_index) def _ensure_local_reduce_buffers( chunk_len: int, value_dtype: torch.dtype, device: torch.device, ) -> tuple[torch.Tensor, torch.Tensor]: if ( self._draft_greedy_local_cap < chunk_len or self._draft_greedy_local_max_buf is None or self._draft_greedy_local_arg_buf is None or self._draft_greedy_local_max_buf.dtype != value_dtype or self._draft_greedy_local_max_buf.device != device or self._draft_greedy_local_arg_buf.device != device ): cap = max(int(chunk_size), chunk_len) self._draft_greedy_local_max_buf = torch.empty( (cap,), dtype=value_dtype, device=device ) self._draft_greedy_local_arg_buf = torch.empty( (cap,), dtype=torch.int64, device=device ) self._draft_greedy_local_cap = cap return ( self._draft_greedy_local_max_buf[:chunk_len], self._draft_greedy_local_arg_buf[:chunk_len], ) # Fast path (common): single-rank greedy sampling over the base vocab shard. # Avoids extra max/id bookkeeping that is only needed for TP sync or added vocab. # # DFLASH draft sampling only materializes a small fixed block of hidden states # each step. On tp=1, splitting those states into many 256-token chunks adds # extra matmul/argmax launches without reducing peak memory meaningfully. if tp_size == 1 and num_added == 0: fast_chunk_size = max(int(chunk_size), 1024) for start in range(0, num_tokens, fast_chunk_size): end = min(num_tokens, start + fast_chunk_size) hs = _cast_hs(hidden_states[start:end]) if num_org > 0: base_logits = torch.matmul(hs, weight[:num_org].T) local_max, local_arg = _ensure_local_reduce_buffers( end - start, base_logits.dtype, hs.device ) torch.max(base_logits, dim=-1, out=(local_max, local_arg)) out_tokens[start:end].copy_(local_arg) out_tokens[start:end].add_(org_vocab_start) else: out_tokens[start:end] = 0 return out_tokens for start in range(0, num_tokens, int(chunk_size)): end = min(num_tokens, start + int(chunk_size)) hs = _cast_hs(hidden_states[start:end]) chunk_len = int(hs.shape[0]) # Base vocab logits. if num_org > 0: base_logits = torch.matmul(hs, weight[:num_org].T) local_max, local_arg = _ensure_local_reduce_buffers( chunk_len, base_logits.dtype, hs.device ) torch.max(base_logits, dim=-1, out=(local_max, local_arg)) else: local_max = torch.full( (chunk_len,), torch.finfo(weight_dtype).min, dtype=weight_dtype, device=hs.device, ) local_arg = torch.zeros( (chunk_len,), dtype=torch.int64, device=hs.device ) # Added vocab logits (e.g., LoRA-added embeddings), if present. if num_added > 0: added_slice_start = num_org_padded added_slice_end = num_org_padded + num_added added_logits = torch.matmul( hs, weight[added_slice_start:added_slice_end].T ) added_max, added_arg = torch.max(added_logits, dim=-1) use_added = added_max > local_max local_max = torch.where(use_added, added_max, local_max) # For base/added conversion below, keep local_arg expressed in the full local # weight index space (base + padding + added), matching `lm_head.weight`. local_arg = torch.where( use_added, added_arg.to(local_arg.dtype) + num_org_padded, local_arg ) # Convert local argmax indices to global token ids. if num_added == 0: local_arg.add_(org_vocab_start) global_ids = local_arg else: global_ids = torch.empty( (chunk_len,), dtype=torch.int64, device=hs.device ) is_base = local_arg < num_org global_ids[is_base] = org_vocab_start + local_arg[is_base] global_ids[~is_base] = added_vocab_start + ( local_arg[~is_base] - num_org_padded ) if tp_size == 1: out_tokens[start:end] = global_ids.to(torch.long) continue # Gather per-rank maxima and associated global ids, then select the global max. needed = tp_size * chunk_len chunk_cap = int(chunk_size) if ( self._draft_greedy_gather_cap < needed or self._draft_greedy_gathered_max_buf is None or self._draft_greedy_gathered_ids_buf is None or self._draft_greedy_gathered_max_buf.dtype != local_max.dtype or self._draft_greedy_gathered_max_buf.device != hs.device ): # Allocate enough space for the max chunk size to avoid reallocations. cap = tp_size * chunk_cap self._draft_greedy_gathered_max_buf = torch.empty( (cap,), dtype=local_max.dtype, device=hs.device ) self._draft_greedy_gathered_ids_buf = torch.empty( (cap,), dtype=global_ids.dtype, device=hs.device ) self._draft_greedy_gather_cap = cap if ( self._draft_greedy_index_cap < chunk_len or self._draft_greedy_best_rank_buf is None or self._draft_greedy_rank_index_buf is None or self._draft_greedy_selected_ids_buf is None or self._draft_greedy_best_rank_buf.device != hs.device or self._draft_greedy_selected_ids_buf.device != hs.device ): self._draft_greedy_best_rank_buf = torch.empty( (chunk_cap,), dtype=torch.int64, device=hs.device ) self._draft_greedy_rank_index_buf = torch.empty( (1, chunk_cap), dtype=torch.int64, device=hs.device ) self._draft_greedy_selected_ids_buf = torch.empty( (1, chunk_cap), dtype=torch.int64, device=hs.device ) self._draft_greedy_index_cap = chunk_cap gathered_max = self._draft_greedy_gathered_max_buf[:needed] gathered_ids = self._draft_greedy_gathered_ids_buf[:needed] tp_group.all_gather_into_tensor(gathered_max, local_max.contiguous()) tp_group.all_gather_into_tensor(gathered_ids, global_ids.contiguous()) gathered_max = gathered_max.view(tp_size, chunk_len) gathered_ids = gathered_ids.view(tp_size, chunk_len) best_rank = self._draft_greedy_best_rank_buf[:chunk_len] torch.argmax(gathered_max, dim=0, out=best_rank) rank_index = self._draft_greedy_rank_index_buf[:, :chunk_len] rank_index[0].copy_(best_rank) selected_ids = self._draft_greedy_selected_ids_buf[:, :chunk_len] torch.gather(gathered_ids, 0, rank_index, out=selected_ids) out_tokens[start:end].copy_(selected_ids.view(-1)) return out_tokens def _append_target_hidden_to_draft_kv_by_loc( self, *, target_hidden: torch.Tensor, cache_loc: torch.Tensor, positions: torch.Tensor, cache_loc_2d: Optional[torch.Tensor] = None, commit_lens: Optional[torch.Tensor] = None, ) -> None: """Materialize target context features into the draft KV cache at explicit slots. For the spec-v2 overlap path, callers can pass dense `[bs, block_size]` `cache_loc_2d` plus `commit_lens`; the prefix-valid writer then commits only the live prefix rows without constructing masked/packed index tensors. """ if target_hidden is None: raise RuntimeError("DFLASH missing target hidden context features.") if target_hidden.numel() == 0: return if target_hidden.ndim != 2: raise ValueError( "DFLASH target_hidden must be 2D, " f"got shape={tuple(target_hidden.shape)}." ) if cache_loc.ndim != 1: raise ValueError( f"DFLASH cache_loc must be 1D, got shape={tuple(cache_loc.shape)}." ) if positions.ndim != 1: raise ValueError( f"DFLASH positions must be 1D, got shape={tuple(positions.shape)}." ) num_tokens = int(target_hidden.shape[0]) if int(cache_loc.numel()) != num_tokens: raise ValueError( "DFLASH cache_loc length mismatch: " f"cache_loc={int(cache_loc.numel())}, target_hidden={num_tokens}." ) if int(positions.numel()) != num_tokens: raise ValueError( "DFLASH positions length mismatch: " f"positions={int(positions.numel())}, target_hidden={num_tokens}." ) if cache_loc_2d is not None: if cache_loc_2d.ndim != 2: raise ValueError( "DFLASH cache_loc_2d must be 2D, " f"got shape={tuple(cache_loc_2d.shape)}." ) if int(cache_loc_2d.numel()) != num_tokens: raise ValueError( "DFLASH cache_loc_2d size mismatch: " f"cache_loc_2d={int(cache_loc_2d.numel())}, target_hidden={num_tokens}." ) if commit_lens is None: raise ValueError( "DFLASH cache_loc_2d requires commit_lens for prefix-valid writes." ) device = self.model_runner.device if cache_loc.device != device: cache_loc = cache_loc.to(device, non_blocking=True) if positions.device != device: positions = positions.to(device, non_blocking=True) if target_hidden.device != device: target_hidden = target_hidden.to(device, non_blocking=True) if cache_loc.dtype != torch.int64: cache_loc = cache_loc.to(torch.int64) if positions.dtype != torch.int64: positions = positions.to(torch.int64) if cache_loc_2d is not None: if cache_loc_2d.device != device: cache_loc_2d = cache_loc_2d.to(device, non_blocking=True) if cache_loc_2d.dtype != torch.int64: cache_loc_2d = cache_loc_2d.to(torch.int64) if commit_lens is not None: if commit_lens.device != device: commit_lens = commit_lens.to(device, non_blocking=True) if commit_lens.dtype != torch.int32: commit_lens = commit_lens.to(torch.int32) with torch.inference_mode(): ctx_hidden = self.draft_model.project_target_hidden(target_hidden) if cache_loc_2d is not None: bs = int(commit_lens.shape[0]) if int(cache_loc_2d.shape[0]) != bs: raise ValueError( "DFLASH cache_loc_2d batch size mismatch: " f"cache_loc_2d={tuple(cache_loc_2d.shape)}, commit_lens={tuple(commit_lens.shape)}." ) if bs == 0: return if self._use_fused_kv_materialize and self._fused_kv_helper is not None: try: self._append_target_hidden_fused( ctx_hidden=ctx_hidden, ctx_positions=positions, ctx_cache_loc=cache_loc, ctx_cache_loc_2d=cache_loc_2d, commit_lens=commit_lens, ) return except Exception as e: logger.warning( "DFLASH fused prefix-direct KV append failed; falling back to the per-layer prefix-direct path: %s", e, ) self._use_fused_kv_materialize = False self._fused_kv_helper = None for layer in self.draft_model.layers: attn = layer.self_attn layer_ctx_hidden = self.draft_model.prepare_context_hidden_for_kv( layer, ctx_hidden ) k, v = attn.kv_proj_only(layer_ctx_hidden) k = attn.apply_k_norm(k) k = attn.apply_k_rope(positions, k) k = k.view(-1, attn.num_kv_heads, attn.head_dim) v = v.view(-1, attn.num_kv_heads, attn.head_dim) self.draft_model_runner.token_to_kv_pool.set_kv_buffer_prefix_valid( attn.attn, cache_loc_2d, commit_lens, k, v, attn.attn.k_scale, attn.attn.v_scale, ) return if self._use_fused_kv_materialize and self._fused_kv_helper is not None: try: self._append_target_hidden_fused( ctx_hidden=ctx_hidden, ctx_positions=positions, ctx_cache_loc=cache_loc, ) return except Exception as e: logger.warning( "DFLASH fused KV append-by-loc failed; falling back to sequential path: %s", e, ) self._use_fused_kv_materialize = False self._fused_kv_helper = None self._append_target_hidden_sequential( ctx_hidden=ctx_hidden, ctx_positions=positions, ctx_cache_loc=cache_loc, ) def _append_target_hidden_sequential( self, ctx_hidden: torch.Tensor, ctx_positions: torch.Tensor, ctx_cache_loc: torch.Tensor, ) -> None: for layer in self.draft_model.layers: attn = layer.self_attn layer_ctx_hidden = self.draft_model.prepare_context_hidden_for_kv( layer, ctx_hidden ) if _is_npu: _, k, v = attn.forward_prepare_npu(ctx_positions, layer_ctx_hidden) else: k, v = attn.kv_proj_only(layer_ctx_hidden) k = attn.apply_k_norm(k) k = attn.apply_k_rope(ctx_positions, k) k = k.view(-1, attn.num_kv_heads, attn.head_dim) v = v.view(-1, attn.num_kv_heads, attn.head_dim) self.draft_model_runner.token_to_kv_pool.set_kv_buffer( attn.attn, ctx_cache_loc, k, v, attn.attn.k_scale, attn.attn.v_scale, ) def _append_target_hidden_fused( self, ctx_hidden: torch.Tensor, ctx_positions: torch.Tensor, ctx_cache_loc: torch.Tensor, ctx_cache_loc_2d: Optional[torch.Tensor] = None, commit_lens: Optional[torch.Tensor] = None, ) -> None: """Fused KV materialization using batched projection + Triton kernel.""" token_to_kv_pool = self.draft_model_runner.token_to_kv_pool if self._fused_kv_helper is None: raise RuntimeError("DFLASH fused KV helper is not initialized.") def _write_layer_kv( layer_idx: int, cache_k: torch.Tensor, cache_v: torch.Tensor, ) -> None: attn = self.draft_model.layers[layer_idx].self_attn.attn if ctx_cache_loc_2d is not None and commit_lens is not None: token_to_kv_pool.set_kv_buffer_prefix_valid( attn, ctx_cache_loc_2d, commit_lens, cache_k, cache_v, attn.k_scale, attn.v_scale, ) else: token_to_kv_pool.set_kv_buffer( attn, ctx_cache_loc, cache_k, cache_v, attn.k_scale, attn.v_scale, ) self._fused_kv_helper.materialize( ctx_hidden=ctx_hidden, positions=ctx_positions, write_layer_kv=_write_layer_kv, ) def _update_target_mamba_state_after_verify( self, *, batch: ScheduleBatch, seq_lens_pre_verify: torch.Tensor, commit_lens: torch.Tensor, ) -> None: """Commit Mamba intermediate states for accepted verify steps. During TARGET_VERIFY, Mamba kernels run with `disable_state_update=True` and cache per-step intermediate states. After acceptance, we need to commit the state corresponding to each request's last accepted step. """ if not self._need_mamba_verify_commit: return attn_backend = self.target_worker.model_runner.attn_backend last_correct_step_indices = commit_lens.to(torch.int64) - 1 mamba_steps_to_track = None if batch.mamba_track_indices is not None: mamba_track_interval = self.server_args.mamba_track_interval to_track_mask = ( seq_lens_pre_verify // mamba_track_interval != batch.seq_lens // mamba_track_interval ) tracking_point = ( batch.seq_lens // mamba_track_interval * mamba_track_interval ) to_track_ith = torch.clamp(tracking_point - seq_lens_pre_verify - 1, min=0) can_track_mask = to_track_mask & ( to_track_ith < commit_lens.to(to_track_ith.dtype) ) mamba_steps_to_track = torch.where( can_track_mask, to_track_ith.to(torch.int64), torch.full_like(to_track_ith, -1, dtype=torch.int64), ) attn_backend.update_mamba_state_after_mtp_verify( last_correct_step_indices=last_correct_step_indices, mamba_track_indices=batch.mamba_track_indices, mamba_steps_to_track=mamba_steps_to_track, model=self.target_worker.model_runner.model, ) def _ensure_accept_bonus_buffers(self, bs: int) -> None: if self._accept_bonus_buffer_cap >= int(bs): return new_cap = max( int(bs), ( self._accept_bonus_buffer_cap * 2 if self._accept_bonus_buffer_cap > 0 else int(bs) ), ) device = self.device block_size = int(self.block_size) self._accept_len_buf = torch.empty((new_cap,), dtype=torch.int32, device=device) self._commit_lens_bufs = [ torch.empty((new_cap,), dtype=torch.int32, device=device) for _ in range(2) ] self._bonus_id_bufs = [ torch.empty((new_cap,), dtype=torch.int32, device=device) for _ in range(2) ] self._out_tokens_bufs = [ torch.empty((new_cap, block_size), dtype=torch.int64, device=device) for _ in range(2) ] self._new_seq_lens_bufs = [ torch.empty((new_cap,), dtype=torch.int64, device=device) for _ in range(2) ] self._accept_bonus_buffer_cap = new_cap def _next_accept_bonus_buffers(self, bs: int) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, ]: self._ensure_accept_bonus_buffers(bs) assert self._accept_len_buf is not None slot = self._accept_bonus_buffer_slot self._accept_bonus_buffer_slot = (slot + 1) % 2 return ( self._accept_len_buf[:bs], self._commit_lens_bufs[slot][:bs], self._bonus_id_bufs[slot][:bs], self._out_tokens_bufs[slot][:bs], self._new_seq_lens_bufs[slot][:bs], ) def _validate_phase1_sampling_support(self, batch: ScheduleBatch) -> None: sampling_info = batch.sampling_info if sampling_info is None or sampling_info.is_all_greedy: return if ( not is_dflash_sampling_verify_available() and not self._warned_sampling_fallback and self.tp_rank == 0 ): logger.warning( "DFLASH non-greedy verification is unavailable on this build/device; " "falling back to greedy argmax verification." ) self._warned_sampling_fallback = True def _make_next_draft_input_prefill( self, *, bonus_tokens: torch.Tensor, seq_lens: torch.Tensor, ) -> DFlashDraftInputV2: return make_draft_input_v2(bonus_tokens=bonus_tokens, new_seq_lens=seq_lens) def _make_next_draft_input_decode( self, *, bonus_tokens: torch.Tensor, new_seq_lens: torch.Tensor, ) -> DFlashDraftInputV2: return make_draft_input_v2(bonus_tokens=bonus_tokens, new_seq_lens=new_seq_lens) def forward_batch_generation( self, batch: ScheduleBatch, on_publish=None, ) -> GenerationBatchResult: if getattr(batch, "return_logprob", False): raise ValueError( "DFLASH speculative decoding does not support return_logprob yet." ) self._validate_phase1_sampling_support(batch) if batch.forward_mode.is_extend() or batch.is_extend_in_batch: # Target prefill: capture DFlash aux hidden states for prompt tokens. batch.capture_hidden_mode = CaptureHiddenMode.FULL batch_output = self.target_worker.forward_batch_generation(batch) logits_output, next_token_ids = ( batch_output.logits_output, batch_output.next_token_ids, ) batch_output.new_seq_lens = batch.seq_lens if on_publish is not None: on_publish(batch_output.new_seq_lens) if logits_output.hidden_states is None: raise RuntimeError( "DFLASH requires target aux hidden capture for prefill, but got None. " "Make sure the target model has DFlash layers-to-capture configured." ) if batch.extend_lens is None or batch.prefix_lens is None: raise RuntimeError( "DFLASH expected extend_lens / prefix_lens to be populated in extend mode, " "but got None." ) # Materialize prompt tokens into the draft KV cache immediately. This is required # for radix cache safety (the scheduler may update radix after prefill returns). device = next_token_ids.device ctx_lens = torch.tensor(batch.extend_lens, dtype=torch.int32, device=device) draft_seq_lens = torch.tensor( batch.prefix_lens, dtype=torch.int32, device=device ) if batch.out_cache_loc is None: raise RuntimeError( "DFLASH prefill expected out_cache_loc, but got None." ) positions, _ = compute_position( self.model_runner.server_args.attention_backend, draft_seq_lens, ctx_lens, int(sum(batch.extend_lens)), ) self._append_target_hidden_to_draft_kv_by_loc( target_hidden=logits_output.hidden_states, cache_loc=batch.out_cache_loc, positions=positions, ) # Avoid copying large hidden-state buffers to CPU in overlap scheduling. logits_output.hidden_states = None batch_output.next_draft_input = self._make_next_draft_input_prefill( bonus_tokens=next_token_ids, seq_lens=batch.seq_lens, ) return batch_output # Decode / target-verify stage. if batch.spec_info is None: batch.spec_info = DFlashDraftInputV2.create_idle_input(device=self.device) draft_input = batch.spec_info if not isinstance(draft_input, DFlashDraftInputV2): raise RuntimeError( "DFLASH spec-v2 expected DFlashDraftInputV2 state on the running batch." ) if batch.forward_mode.is_idle(): empty_ids = torch.empty((0,), dtype=torch.int64, device=self.device) empty_lens = torch.empty((0,), dtype=torch.int32, device=self.device) next_draft_input = self._make_next_draft_input_decode( bonus_tokens=torch.empty((0,), device=self.device, dtype=torch.int64), new_seq_lens=torch.empty((0,), device=self.device, dtype=torch.int64), ) if on_publish is not None: on_publish(next_draft_input.new_seq_lens) return GenerationBatchResult( logits_output=None, next_token_ids=empty_ids, accept_lens=empty_lens, next_draft_input=next_draft_input, can_run_cuda_graph=False, speculative_num_draft_tokens=int(self.block_size), new_seq_lens=next_draft_input.new_seq_lens, ) # `seq_lens` is carried over from the previous overlap iteration and may have been # produced on another stream. batch.seq_lens.record_stream( torch.get_device_module(self.device).current_stream() ) bs = len(batch.seq_lens) device = self.device # --- 1) Draft a fixed block with the draft model. target_model = self.target_worker.model_runner.model embed_module = target_model.get_input_embeddings() lm_head = getattr(target_model, "lm_head", None) if lm_head is None or not hasattr(lm_head, "weight"): raise RuntimeError( "DFLASH requires the target model to expose `lm_head` with `weight`." ) block_size = int(self.block_size) self._ensure_draft_block_buffers(bs) assert self._draft_block_ids_buf is not None assert self._draft_block_positions_buf is not None assert self._draft_block_tokens_buf is not None assert self._draft_verify_out_cache_loc_buf is not None assert self._draft_block_end_buf is not None assert self._draft_seq_lens_cpu_buf is not None block_ids = self._draft_block_ids_buf[:bs] prefix_lens = batch.seq_lens positions_2d = self._draft_block_positions_buf[:bs] verify_out_cache_loc_2d = self._draft_verify_out_cache_loc_buf[:bs] if self._use_triton_prepare_block: try: _prepare_dflash_draft_block_unchecked( bonus_tokens=draft_input.bonus_tokens.view(-1), prefix_lens=prefix_lens.view(-1), req_pool_indices=batch.req_pool_indices.view(-1), req_to_token=self.model_runner.req_to_token_pool.req_to_token, block_ids_out=block_ids, positions_out=positions_2d, cache_loc_out=verify_out_cache_loc_2d, mask_token_id=int(self._mask_token_id), ) except Exception as e: self._use_triton_prepare_block = False logger.warning( "DFLASH Triton prepare_block failed; falling back to eager path: %s", e, ) block_ids.fill_(int(self._mask_token_id)) block_ids[:, 0].copy_(draft_input.bonus_tokens) torch.add( prefix_lens.unsqueeze(1), self._block_pos_offsets, out=positions_2d, ) end_offset = prefix_lens + block_size verify_out_cache_loc = assign_extend_cache_locs_func( req_pool_indices=batch.req_pool_indices, req_to_token=self.model_runner.req_to_token_pool.req_to_token, start_offset=prefix_lens, end_offset=end_offset, batch_size=bs, draft_token_num=block_size, device=device, ) verify_out_cache_loc_2d.copy_(verify_out_cache_loc.view(bs, block_size)) else: block_ids.fill_(int(self._mask_token_id)) block_ids[:, 0].copy_(draft_input.bonus_tokens) torch.add( prefix_lens.unsqueeze(1), self._block_pos_offsets, out=positions_2d, ) end_offset = prefix_lens + block_size verify_out_cache_loc = assign_extend_cache_locs_func( req_pool_indices=batch.req_pool_indices, req_to_token=self.model_runner.req_to_token_pool.req_to_token, start_offset=prefix_lens, end_offset=end_offset, batch_size=bs, draft_token_num=block_size, device=device, ) verify_out_cache_loc_2d.copy_(verify_out_cache_loc.view(bs, block_size)) noise_embedding = embed_module(block_ids) input_embeds = noise_embedding.view(-1, noise_embedding.shape[-1]) positions = positions_2d.reshape(-1) verify_out_cache_loc = verify_out_cache_loc_2d.reshape(-1) seq_lens_cpu = self._draft_seq_lens_cpu_buf[:bs] if self.use_compact_draft_cache: # Rebuild the draft-local sliding-window view from committed target state. draft_prefix_lens = self._compute_compact_draft_seq_lens(prefix_lens) seq_lens_cpu.copy_(draft_prefix_lens.to(device="cpu", dtype=torch.int32)) suffix_start = prefix_lens.to(torch.int64) - draft_prefix_lens.to( torch.int64 ) suffix_cache_loc = self._gather_req_to_token_segments( req_to_token=self.model_runner.req_to_token_pool.req_to_token, req_pool_indices=batch.req_pool_indices, start=suffix_start, lengths=draft_prefix_lens, ) assign_req_to_token_pool_func( batch.req_pool_indices, self.draft_model_runner.req_to_token_pool.req_to_token, torch.zeros_like(draft_prefix_lens), draft_prefix_lens, suffix_cache_loc, bs, ) block_end = self._draft_block_end_buf[:bs] torch.add(draft_prefix_lens, block_size, out=block_end) assign_req_to_token_pool_func( batch.req_pool_indices, self.draft_model_runner.req_to_token_pool.req_to_token, draft_prefix_lens, block_end, verify_out_cache_loc, bs, ) draft_seq_lens = draft_prefix_lens draft_seq_lens_sum = int(seq_lens_cpu.sum().item()) else: # Non-windowed path uses the shared overallocated mapping directly. # Backend planning only needs a safe upper bound for the committed # prefix lengths, not the full allocator reservation length. draft_seq_lens = prefix_lens if batch.seq_lens_cpu is not None: # Host bound = committed prefix + one verify block. seq_lens_cpu.copy_(batch.seq_lens_cpu) seq_lens_cpu.add_(block_size) draft_seq_lens_sum = int(seq_lens_cpu.sum()) elif draft_input.reserved_seq_lens_cpu is not None: # GPU-only backend: reserved is a safe over-estimate. seq_lens_cpu.copy_(draft_input.reserved_seq_lens_cpu) draft_seq_lens_sum = int(draft_input.reserved_seq_lens_sum) else: seq_lens_cpu.copy_(prefix_lens.to("cpu", dtype=torch.int32)) draft_seq_lens_sum = int(prefix_lens.sum().item()) forward_batch = ForwardBatch( forward_mode=ForwardMode.TARGET_VERIFY, batch_size=bs, input_ids=block_ids.flatten(), req_pool_indices=batch.req_pool_indices, seq_lens=draft_seq_lens, out_cache_loc=verify_out_cache_loc, seq_lens_sum=draft_seq_lens_sum, seq_lens_cpu=seq_lens_cpu, positions=positions, input_embeds=input_embeds, spec_algorithm=SpeculativeAlgorithm.DFLASH, spec_info=self._draft_block_spec_info, capture_hidden_mode=CaptureHiddenMode.NULL, ) with torch.inference_mode(): draft_out = self.draft_model_runner.forward(forward_batch) draft_logits_output = draft_out.logits_output if self._draft_sampler is not None and draft_out.can_run_graph: draft_next = self._draft_sampler.out[ : bs * (int(self.block_size) - 1) ].view(bs, int(self.block_size) - 1) else: draft_hidden = draft_logits_output.hidden_states if draft_hidden is None: raise RuntimeError("DFLASH draft model returned no hidden states.") draft_hidden = draft_hidden.view(bs, int(self.block_size), -1) draft_next = self._greedy_sample_from_vocab_parallel_head( hidden_states=draft_hidden[:, 1:, :].reshape( -1, draft_hidden.shape[-1] ), lm_head=lm_head, ).view(bs, int(self.block_size) - 1) draft_tokens = self._draft_block_tokens_buf[:bs] draft_tokens[:, 0].copy_(block_ids[:, 0]) draft_tokens[:, 1:].copy_(draft_next) # --- 2) Target verify. # TARGET_VERIFY uses standard causal masking; custom masks are unnecessary here. custom_mask = None verify_input_ids = draft_tokens.reshape(-1) verify_input = DFlashVerifyInput( draft_token=verify_input_ids, positions=positions, draft_token_num=int(self.block_size), custom_mask=custom_mask, capture_hidden_mode=CaptureHiddenMode.FULL, ) batch.out_cache_loc = verify_out_cache_loc sampling_info = batch.sampling_info seq_lens_pre_verify = ( batch.seq_lens.clone() if self._need_mamba_verify_commit else None ) seq_lens_cpu_backup = batch.seq_lens_cpu seq_lens_sum_backup = batch.seq_lens_sum if seq_lens_cpu_backup is not None: # Verify host bound = committed prefix + one verify block (matches draft). verify_host_seq_lens = seq_lens_cpu_backup + block_size batch.seq_lens_cpu = verify_host_seq_lens batch.seq_lens_sum = int(verify_host_seq_lens.sum()) elif draft_input.reserved_seq_lens_cpu is not None: batch.seq_lens_cpu = draft_input.reserved_seq_lens_cpu batch.seq_lens_sum = int(draft_input.reserved_seq_lens_sum) verify_forward_batch, _ = verify_input.prepare_for_verify( batch, self.target_worker ) batch.seq_lens_cpu = seq_lens_cpu_backup batch.seq_lens_sum = seq_lens_sum_backup target_out = self.target_worker.forward_batch_generation( batch=None, forward_batch=verify_forward_batch, is_verify=True, skip_attn_backend_init=True, ) logits_output = target_out.logits_output can_run_cuda_graph = target_out.can_run_cuda_graph if sampling_info is not None: apply_dflash_verify_logits_adjustments( next_token_logits=logits_output.next_token_logits, sampling_info=sampling_info, draft_token_num=int(self.block_size), ) candidates = draft_tokens new_seq_lens = None if ( sampling_info is not None and not sampling_info.is_all_greedy and is_dflash_sampling_verify_available() ): accept_len, bonus = compute_dflash_sampling_correct_drafts_and_bonus( candidates=candidates, next_token_logits=logits_output.next_token_logits, sampling_info=sampling_info, max_top_k=draft_input.max_top_k, uniform_top_k_value=draft_input.uniform_top_k_value, ) commit_lens = accept_len.to(torch.int32) + 1 # [bs] out_tokens = torch.empty( (bs, int(self.block_size)), dtype=torch.int64, device=device ) if int(self.block_size) > 1: out_tokens[:, : int(self.block_size) - 1].copy_(candidates[:, 1:]) out_tokens[:, int(self.block_size) - 1].fill_(0) out_tokens.scatter_(1, accept_len.to(torch.int64)[:, None], bonus[:, None]) else: target_predict = torch.argmax(logits_output.next_token_logits, dim=-1).view( bs, int(self.block_size) ) if self._use_triton_accept_bonus: try: ( accept_len, commit_lens, bonus, out_tokens, new_seq_lens, ) = self._next_accept_bonus_buffers(bs) _compute_dflash_accept_bonus_triton_unchecked( candidates=candidates, target_top1=target_predict, accept_lens_out=accept_len, commit_lens_out=commit_lens, bonus_ids_out=bonus, out_tokens_out=out_tokens, prefix_lens=prefix_lens, new_seq_lens_out=new_seq_lens, ) except Exception as e: self._use_triton_accept_bonus = False logger.warning( "DFLASH Triton accept/bonus failed; falling back to eager path: %s", e, ) accept_len, bonus = compute_dflash_correct_drafts_and_bonus( candidates=candidates, target_predict=target_predict, ) commit_lens = accept_len.to(torch.int32) + 1 # [bs] out_tokens = torch.empty( (bs, int(self.block_size)), dtype=torch.int64, device=device, ) if int(self.block_size) > 1: out_tokens[:, : int(self.block_size) - 1].copy_( candidates[:, 1:] ) out_tokens[:, int(self.block_size) - 1].fill_(0) out_tokens.scatter_( 1, accept_len.to(torch.int64)[:, None], bonus[:, None] ) else: accept_len, bonus = compute_dflash_correct_drafts_and_bonus( candidates=candidates, target_predict=target_predict, ) commit_lens = accept_len.to(torch.int32) + 1 # [bs] out_tokens = torch.empty( (bs, int(self.block_size)), dtype=torch.int64, device=device ) if int(self.block_size) > 1: out_tokens[:, : int(self.block_size) - 1].copy_(candidates[:, 1:]) out_tokens[:, int(self.block_size) - 1].fill_(0) out_tokens.scatter_( 1, accept_len.to(torch.int64)[:, None], bonus[:, None] ) if self._need_mamba_verify_commit: assert seq_lens_pre_verify is not None self._update_target_mamba_state_after_verify( batch=batch, seq_lens_pre_verify=seq_lens_pre_verify, commit_lens=commit_lens, ) if new_seq_lens is None: new_seq_lens = prefix_lens + commit_lens.to(prefix_lens.dtype) if on_publish is not None: on_publish(new_seq_lens) # --- 3) Materialize committed verify-input tokens into draft KV cache. hidden = logits_output.hidden_states if hidden is None: raise RuntimeError( "DFLASH verify requires target hidden states, but got None." ) hidden = hidden.view(bs, int(self.block_size), -1) self._append_target_hidden_to_draft_kv_by_loc( target_hidden=hidden.reshape(-1, hidden.shape[-1]), cache_loc=verify_out_cache_loc, cache_loc_2d=verify_out_cache_loc_2d, positions=positions, commit_lens=commit_lens, ) # Avoid copying large hidden-state buffers to CPU in overlap scheduling. logits_output.hidden_states = None next_draft_input = self._make_next_draft_input_decode( bonus_tokens=bonus, new_seq_lens=new_seq_lens, ) return GenerationBatchResult( logits_output=logits_output, next_token_ids=out_tokens.reshape(-1), accept_lens=commit_lens, can_run_cuda_graph=can_run_cuda_graph, next_draft_input=next_draft_input, speculative_num_draft_tokens=int(self.block_size), # The non-overlap (sync) scheduler path advances batch.seq_lens # from the result; overlap carries it via next_draft_input instead. new_seq_lens=new_seq_lens, )