diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index ac2e76478..baac7d214 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1172,9 +1172,9 @@ class Scheduler( self.device_module = torch.get_device_module(self.device) # FutureMap is always-on: input_ids relay used in both modes. - # Workers not on BaseSpecWorker (e.g. NGRAM / DFLASH) lack the - # override; fall back to target-only so the helper still produces a - # safe decision (no accidental opt-out for unaudited shapes). + # Workers without the spec_v2_attn_backends override fall back to + # target-only so the helper still produces a safe decision (no + # accidental opt-out for unaudited shapes). if self.draft_worker is not None: attn_backends = getattr( self.draft_worker, diff --git a/python/sglang/srt/speculative/dflash_worker.py b/python/sglang/srt/speculative/dflash_worker.py deleted file mode 100644 index bc3d6cf0f..000000000 --- a/python/sglang/srt/speculative/dflash_worker.py +++ /dev/null @@ -1,1011 +0,0 @@ -import logging -import math -from copy import deepcopy -from typing import Optional - -import torch - -from sglang.srt.distributed import get_tp_group -from sglang.srt.managers.schedule_batch import ScheduleBatch -from sglang.srt.managers.tp_worker import TpModelWorker -from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode -from sglang.srt.server_args import ( - ServerArgs, - get_global_server_args, - set_global_server_args_for_scheduler, -) -from sglang.srt.speculative.dflash_info import DFlashVerifyInput -from sglang.srt.speculative.dflash_utils import ( - can_dflash_use_fused_qkv_proj, - parse_dflash_draft_config, -) -from sglang.srt.utils import is_cuda, 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.srt.speculative.triton_ops.fused_kv_materialize import ( - FusedKVMaterializeHelper, - ) - - _FusedKVMaterializeHelper = FusedKVMaterializeHelper - return _FusedKVMaterializeHelper - - -class DFlashWorker: - """Shared DFLASH infrastructure (draft model, draft KV materialization).""" - - 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.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 - - # Draft runner (separate KV cache + attention backend). - # 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. - target_req_to_token_pool, target_token_to_kv_pool_allocator = ( - target_worker.get_memory_pool() - ) - shared_req_to_token_pool = ( - None if self.use_compact_draft_cache else target_req_to_token_pool - ) - draft_server_args = deepcopy(server_args) - draft_server_args.skip_tokenizer_init = True - draft_backend = draft_server_args.speculative_draft_attention_backend - supported_draft_backends = ("flashinfer", "fa3", "fa4", "triton", "ascend") - if draft_backend is None: - draft_backend, _ = draft_server_args.get_attention_backends() - if draft_backend is None: - # Use triton on ROCm (no FlashInfer), flashinfer on CUDA - import torch as _torch - - draft_backend = "triton" if _torch.version.hip else "flashinfer" - elif draft_backend == "trtllm_mha": - import torch as _torch - - _fb = "triton" if _torch.version.hip else "flashinfer" - logger.warning( - "DFLASH draft worker does not support 'trtllm_mha' because the " - "draft path requires per-layer DFlash attention. Falling back to " - "'%s'.", - _fb, - ) - draft_backend = _fb - elif draft_backend not in supported_draft_backends: - import torch as _torch - - _fb = "triton" if _torch.version.hip else "flashinfer" - logger.warning( - "DFLASH draft worker only supports attention_backend in %s for now, " - "but got %r. Falling back to '%s'.", - supported_draft_backends, - draft_backend, - _fb, - ) - draft_backend = _fb - # Make the draft worker backend explicit and self-contained (no further overrides). - draft_server_args.speculative_draft_attention_backend = None - draft_server_args.prefill_attention_backend = None - draft_server_args.decode_attention_backend = None - draft_server_args.attention_backend = draft_backend - # Keep draft context length aligned with the target. - draft_server_args.context_length = ( - target_worker.model_runner.model_config.context_len - ) - saved_server_args = get_global_server_args() - self.draft_worker = TpModelWorker( - server_args=draft_server_args, - gpu_id=gpu_id, - tp_rank=tp_rank, - moe_ep_rank=moe_ep_rank, - pp_rank=0, - attn_cp_rank=attn_cp_rank, - moe_dp_rank=moe_dp_rank, - dp_rank=dp_rank, - nccl_port=nccl_port, - is_draft_worker=True, - req_to_token_pool=shared_req_to_token_pool, - token_to_kv_pool_allocator=target_token_to_kv_pool_allocator, - memory_pool_config=target_worker.model_runner.memory_pool_config, - ) - set_global_server_args_for_scheduler(saved_server_args) - self.draft_model_runner = self.draft_worker.model_runner - # Keep the same alias that other spec-v2 workers expose. - self.draft_worker.draft_runner = self.draft_model_runner - self.draft_model = self.draft_model_runner.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", - getattr(draft_server_args, "attention_backend", None), - 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 = torch.arange( - self.block_size, device=self.device, dtype=torch.int64 - ) - 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 = DFlashVerifyInput( - draft_token=torch.empty((0,), dtype=torch.long, device=self.device), - positions=torch.empty((0,), dtype=torch.int64, device=self.device), - draft_token_num=int(self.block_size), - custom_mask=None, - capture_hidden_mode=CaptureHiddenMode.NULL, - ) - 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() - self._fused_kv_helper: Optional[object] = None - if self._use_fused_kv_materialize: - self._init_fused_kv_helper() - - 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" - - 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. - return getattr(self.target_worker, name) - - def on_verify_complete_cpu( - self, num_correct_drafts_per_req: list[int], batch_size: int = 0 - ) -> None: - pass - - 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) - - tp_group = get_tp_group() - tp_size = int(tp_group.world_size) - - if not hasattr(lm_head, "weight") or not hasattr(lm_head, "shard_indices"): - raise RuntimeError( - "DFLASH greedy sampling requires a vocab-parallel head with `weight` and `shard_indices`." - ) - - shard = lm_head.shard_indices - weight = lm_head.weight # [local_vocab_padded, hidden] - weight_dtype = weight.dtype - - # 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) - - 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) - - 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 - k, v = attn.kv_proj_only(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 - if _is_npu: - _, k, v = attn.forward_prepare_npu(ctx_positions, ctx_hidden) - else: - k, v = attn.kv_proj_only(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. - """ - attn_backend = self.target_worker.model_runner.attn_backend - if not hasattr(attn_backend, "update_mamba_state_after_mtp_verify"): - return - - 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, - ) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index a44b5c637..d476a52fd 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -1,8 +1,11 @@ import logging +import math +from copy import deepcopy from typing import List, Optional import torch +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 @@ -12,16 +15,22 @@ from sglang.srt.model_executor.forward_batch_info import ( ForwardMode, compute_position, ) -from sglang.srt.server_args import ServerArgs +from sglang.srt.server_args import ( + ServerArgs, + get_global_server_args, + set_global_server_args_for_scheduler, +) +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.dflash_worker import DFlashWorker from sglang.srt.speculative.eagle_info_v2 import assign_extend_cache_locs_func from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.srt.speculative.spec_utils import assign_req_to_token_pool_func @@ -31,12 +40,28 @@ from sglang.srt.speculative.triton_ops.dflash_accept_bonus import ( from sglang.srt.speculative.triton_ops.dflash_prepare_block import ( _prepare_dflash_draft_block_unchecked, ) -from sglang.srt.utils import is_cuda, is_hip +from sglang.srt.utils import is_cuda, is_hip, is_npu + +_is_npu = is_npu() + logger = logging.getLogger(__name__) +_FusedKVMaterializeHelper = None -class DFlashWorkerV2(DFlashWorker): + +def _get_fused_kv_materialize_helper(): + global _FusedKVMaterializeHelper + if _FusedKVMaterializeHelper is None: + from sglang.srt.speculative.triton_ops.fused_kv_materialize import ( + FusedKVMaterializeHelper, + ) + + _FusedKVMaterializeHelper = FusedKVMaterializeHelper + return _FusedKVMaterializeHelper + + +class DFlashWorkerV2(BaseSpecWorker): """DFLASH speculative decoding worker (spec-v2). Drives both overlap and non-overlap scheduling, same as EAGLE: the @@ -55,17 +80,183 @@ class DFlashWorkerV2(DFlashWorker): nccl_port: int, target_worker: TpModelWorker, ): - super().__init__( - server_args=server_args, + 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.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 + + # Draft runner (separate KV cache + attention backend). + # 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. + target_req_to_token_pool, target_token_to_kv_pool_allocator = ( + target_worker.get_memory_pool() + ) + shared_req_to_token_pool = ( + None if self.use_compact_draft_cache else target_req_to_token_pool + ) + draft_server_args = deepcopy(server_args) + draft_server_args.skip_tokenizer_init = True + draft_backend = draft_server_args.speculative_draft_attention_backend + supported_draft_backends = ("flashinfer", "fa3", "fa4", "triton", "ascend") + if draft_backend is None: + draft_backend, _ = draft_server_args.get_attention_backends() + if draft_backend is None: + # Use triton on ROCm (no FlashInfer), flashinfer on CUDA + import torch as _torch + + draft_backend = "triton" if _torch.version.hip else "flashinfer" + elif draft_backend == "trtllm_mha": + import torch as _torch + + _fb = "triton" if _torch.version.hip else "flashinfer" + logger.warning( + "DFLASH draft worker does not support 'trtllm_mha' because the " + "draft path requires per-layer DFlash attention. Falling back to " + "'%s'.", + _fb, + ) + draft_backend = _fb + elif draft_backend not in supported_draft_backends: + import torch as _torch + + _fb = "triton" if _torch.version.hip else "flashinfer" + logger.warning( + "DFLASH draft worker only supports attention_backend in %s for now, " + "but got %r. Falling back to '%s'.", + supported_draft_backends, + draft_backend, + _fb, + ) + draft_backend = _fb + # Make the draft worker backend explicit and self-contained (no further overrides). + draft_server_args.speculative_draft_attention_backend = None + draft_server_args.prefill_attention_backend = None + draft_server_args.decode_attention_backend = None + draft_server_args.attention_backend = draft_backend + # Keep draft context length aligned with the target. + draft_server_args.context_length = ( + target_worker.model_runner.model_config.context_len + ) + saved_server_args = get_global_server_args() + self._draft_worker = TpModelWorker( + server_args=draft_server_args, gpu_id=gpu_id, tp_rank=tp_rank, - dp_rank=dp_rank, moe_ep_rank=moe_ep_rank, + pp_rank=0, attn_cp_rank=attn_cp_rank, moe_dp_rank=moe_dp_rank, + dp_rank=dp_rank, nccl_port=nccl_port, - target_worker=target_worker, + is_draft_worker=True, + req_to_token_pool=shared_req_to_token_pool, + token_to_kv_pool_allocator=target_token_to_kv_pool_allocator, + memory_pool_config=target_worker.model_runner.memory_pool_config, ) + set_global_server_args_for_scheduler(saved_server_args) + self.draft_model_runner = self._draft_worker.model_runner + # Keep the same alias that other spec-v2 workers expose. + self._draft_worker.draft_runner = self.draft_model_runner + self.draft_model = self.draft_model_runner.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", + getattr(draft_server_args, "attention_backend", None), + 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 = torch.arange( + self.block_size, device=self.device, dtype=torch.int64 + ) + 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 = DFlashVerifyInput( + draft_token=torch.empty((0,), dtype=torch.long, device=self.device), + positions=torch.empty((0,), dtype=torch.int64, device=self.device), + draft_token_num=int(self.block_size), + custom_mask=None, + capture_hidden_mode=CaptureHiddenMode.NULL, + ) + 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() + 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 @@ -77,6 +268,803 @@ class DFlashWorkerV2(DFlashWorker): 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 + # BaseDraftWorker 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 _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" + + 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) + + tp_group = get_tp_group() + tp_size = int(tp_group.world_size) + + if not hasattr(lm_head, "weight") or not hasattr(lm_head, "shard_indices"): + raise RuntimeError( + "DFLASH greedy sampling requires a vocab-parallel head with `weight` and `shard_indices`." + ) + + shard = lm_head.shard_indices + weight = lm_head.weight # [local_vocab_padded, hidden] + weight_dtype = weight.dtype + + # 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) + + 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) + + 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 + k, v = attn.kv_proj_only(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 + if _is_npu: + _, k, v = attn.forward_prepare_npu(ctx_positions, ctx_hidden) + else: + k, v = attn.kv_proj_only(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. + """ + attn_backend = self.target_worker.model_runner.attn_backend + if not hasattr(attn_backend, "update_mamba_state_after_mtp_verify"): + return + + 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