# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from dataclasses import dataclass, field from typing import Any, ClassVar, cast import torch from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.warmup.jit_warmup import ( WarmupIntRange, zip_inputs, ) from vllm.model_executor.warmup.jit_warmup_triton_helper import ( LaunchSpec, TritonWarmupTensor, VllmTritonJitKernel, kernel_launcher, ) from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.math_utils import cdiv, next_power_of_2 from vllm.v1.attention.backend import ( AttentionBackend, AttentionCGSupport, AttentionMetadataBuilder, CommonAttentionMetadata, MultipleOf, ) from vllm.v1.attention.backends.mla.compressor_utils import ( get_dspark_swa_index_width, ) from vllm.v1.attention.backends.utils import split_decodes_and_prefills from vllm.v1.attention.ops.flashmla import FlashMLASchedMeta, get_mla_metadata from vllm.v1.kv_cache_interface import ( KVCacheSpec, MLAAttentionSpec, SlidingWindowMLASpec, get_kv_quant_mode, ) # DeepseekV4 decode layer types, keyed by compress_ratio. Each type has a distinct # (topk, extra_topk, extra_page_block_size) config, so they cannot share a # FlashMLA tile-scheduler plan. Within a type, all ~60 DeepseekV4 layers share one # plan per step because b / s_q / h_q / page_block_sizes / topks are identical. _LAYER_TYPE_SWAONLY = "swaonly" _LAYER_TYPE_C4A = "c4a" _LAYER_TYPE_C128A = "c128a" # v4.1 ratio-1 / ratio-2 indexer layers: same indexer-path shape as C4A, but # their compressed page block is block_size // ratio, so each gets its own plan. # v4.1 builders classify layer types themselves (ratio 0 is SWA-only there); see # deepseek_v4_1/sparse_mla.py. _LAYER_TYPE_C1A = "c1a" _LAYER_TYPE_C2A = "c2a" def _layer_type_for(compress_ratio: int) -> str: if compress_ratio <= 1: return _LAYER_TYPE_SWAONLY if compress_ratio == 2: return _LAYER_TYPE_C2A if compress_ratio == 4: return _LAYER_TYPE_C4A if compress_ratio == 128: return _LAYER_TYPE_C128A raise ValueError( f"Unsupported DeepseekV4 compress_ratio={compress_ratio}; " "expected 1, 2, 4, or 128." ) class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): def __init__( self, head_dim: int, window_size: int, dtype: torch.dtype, prefix: str, cache_config: CacheConfig, backend_cls: "type[AttentionBackend] | None" = None, block_size: int = 64, ): super().__init__() self.backend_cls = backend_cls or DeepseekSparseSWABackend self.kv_cache = torch.tensor([]) self.head_dim = head_dim self.window_size = window_size self.prefix = prefix self.cache_config = cache_config self.dtype = dtype compilation_config = get_current_vllm_config().compilation_config if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self # Any multiple of 32; the sparse decode kernels take the page size at # runtime. self.block_size = block_size # uint8: fp8_ds_mla UE8M0 paged layout. bfloat16 / float8_e4m3fn: # contiguous full-cache layout. assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn) def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: # [B, H=1, N, C] -> [B, N, C] self.kv_cache = kv_cache.squeeze(1) def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: # fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous # bf16/fp8 cache uses the natural element-size page. uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla" return SlidingWindowMLASpec( block_size=self.block_size, num_kv_heads=1, head_size=self.head_dim, dtype=self.dtype, sliding_window=self.window_size, cache_dtype_str=self.cache_config.cache_dtype, # DeepseekV4 fp8_ds_mla: 584B per token (448B NoPE + 128B RoPE + 8B scales) state_content_bytes=584 if uses_fp8_ds_mla_layout else None, # 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577). alignment=576 if uses_fp8_ds_mla_layout else 512, model_version="deepseek_v4", kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype), ) def forward(self): ... def get_attn_backend(self) -> type[AttentionBackend]: return self.backend_cls class DeepseekSparseSWABackend(AttentionBackend): @staticmethod def get_name() -> str: return "DEEPSEEK_SPARSE_SWA" @staticmethod def get_supported_kernel_block_sizes() -> list[int | MultipleOf]: return [MultipleOf(32)] @classmethod def get_swa_block_size(cls) -> int: """Tech2Wild/Kai 2026-09-10: tokens per SWA cache page handed to DeepseekV4SWACache. FlashMLA / TRTLLM-gen take the page size at runtime (32); kernels instantiated for one page size override this.""" return 32 @classmethod def get_preferred_block_size(cls, default_block_size: int) -> int: return 256 @classmethod def get_supported_head_sizes(cls) -> list[int]: return [512] @staticmethod def get_builder_cls() -> type["DeepseekSparseSWAMetadataBuilder"]: if current_platform.is_rocm(): from vllm.models.deepseek_v4.amd.rocm import ( DeepseekV4ROCMAiterSparseSWAMetadataBuilder, ) return DeepseekV4ROCMAiterSparseSWAMetadataBuilder return DeepseekSparseSWAMetadataBuilder @dataclass class DeepseekSparseSWAMetadata: block_table: torch.Tensor slot_mapping: torch.Tensor block_size: int seq_lens: torch.Tensor | None = None # [num_seqs] query_start_loc: torch.Tensor | None = None # [num_seqs + 1] query_start_loc_cpu: torch.Tensor | None = None # [num_seqs + 1] is_valid_token: torch.Tensor | None = None # [num_tokens] token_to_req_indices: torch.Tensor | None = None # [num_tokens] decode_swa_indices: torch.Tensor | None = None # [num_decode_tokens, width] decode_swa_lens: torch.Tensor | None = None # [num_decode_tokens] # window_size (causal) or noncausal_index_width (DSpark non-causal). decode_swa_width: int = 0 # Paged-coordinate prefill SWA indices/lens (FP8 paged-direct prefill). prefill_swa_indices: torch.Tensor | None = ( None # [num_prefill_tokens, 1, prefill_index_width] ) prefill_swa_lens: torch.Tensor | None = None # [num_prefill_tokens] # In-image bidirectional visibility (vision variant): per-token counts of # extra visible tokens to the left/right inside an image span. Indexed by # absolute (decode-first) token position; only prefill rows are written. # None when the model is text-only or the batch has no image spans. prefill_left_visible: torch.Tensor | None = None prefill_right_visible: torch.Tensor | None = None # Number of decode/prefill requests/tokens (batch is reordered: decodes first) num_decodes: int = 0 num_prefills: int = 0 num_decode_tokens: int = 0 num_prefill_tokens: int = 0 max_decode_query_len: int = 1 # Pre-computed prefill metadata shared across all DeepseekV4 attention layers. prefill_seq_lens: torch.Tensor | None = None prefill_seq_lens_cpu: torch.Tensor | None = None prefill_gather_lens: torch.Tensor | None = None prefill_query_lens_cpu: torch.Tensor | None = None prefill_window_size: int = 0 prefill_max_model_len: int = 0 prefill_max_num_batched_tokens: int = 0 # Per-layer-type FlashMLA tile-scheduler metadata. One FlashMLASchedMeta # per present DeepseekV4 layer type, shared across all ~60 layers of that type # within a decode step. The first forward call of a given type triggers # the in-kernel planner (which also allocates tile_scheduler_metadata and # num_splits via PyTorch's graph-aware allocator); subsequent same-type # calls skip planning and reuse the plan. Fresh instance per build(), so # have_initialized is always False at the start of a step and the plan # is re-derived from current seq_lens / topk_length on replay. # None for layer types the model does not use (or when num_decode_tokens # is zero). tile_sched_swaonly: "FlashMLASchedMeta | None" = None tile_sched_c4a: "FlashMLASchedMeta | None" = None tile_sched_c128a: "FlashMLASchedMeta | None" = None tile_sched_c1a: "FlashMLASchedMeta | None" = None tile_sched_c2a: "FlashMLASchedMeta | None" = None flashinfer_sparse_index_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = field( default_factory=dict ) def get_prefill_chunk_plan( self, compress_ratio: int, prefill_chunk_size: int, has_compressed: bool | None = None, ) -> list[tuple[int, int, int, int]]: if self.num_prefills == 0: return [] assert self.prefill_seq_lens_cpu is not None assert self.prefill_query_lens_cpu is not None # Whether the layer gathers a compressed-KV region into the prefill # workspace. v4.0 callers keep the legacy default (ratio <= 1 means # SWA-only); v4.1 passes has_compressed explicitly because its # compress_ratio==1 layers DO have a full-length compressed cache. if has_compressed is None: has_compressed = compress_ratio > 1 # query_len <= max_num_batched_tokens and # gather_len = query_len + min(prefix_len, window_size - 1), so the # worst-case gathered width is bounded by # max_num_batched_tokens + window_size - 1. The compressed prefix pool # is bounded by ceil(max_model_len / compress_ratio). max_workspace_area = prefill_chunk_size * ( (cdiv(self.prefill_max_model_len, compress_ratio) if has_compressed else 0) + self.prefill_window_size + self.prefill_max_num_batched_tokens ) prefix_lens_cpu = self.prefill_seq_lens_cpu - self.prefill_query_lens_cpu gather_lens_cpu = self.prefill_query_lens_cpu + torch.clamp( prefix_lens_cpu, min=0, max=self.prefill_window_size - 1 ) compressed_lens_cpu = ( torch.div( self.prefill_seq_lens_cpu, compress_ratio, rounding_mode="floor", ) if has_compressed else torch.zeros_like(self.prefill_seq_lens_cpu) ) chunk_plan: list[tuple[int, int, int, int]] = [] chunk_start = 0 while chunk_start < self.num_prefills: chunk_max_compressed = int(compressed_lens_cpu[chunk_start].item()) chunk_max_gather = int(gather_lens_cpu[chunk_start].item()) chunk_end = chunk_start + 1 while chunk_end < self.num_prefills: candidate_max_compressed = max( chunk_max_compressed, int(compressed_lens_cpu[chunk_end].item()), ) candidate_max_gather = max( chunk_max_gather, int(gather_lens_cpu[chunk_end].item()), ) candidate_width = candidate_max_compressed + candidate_max_gather candidate_area = (chunk_end - chunk_start + 1) * candidate_width if candidate_area > max_workspace_area: break chunk_max_compressed = candidate_max_compressed chunk_max_gather = candidate_max_gather chunk_end += 1 chunk_plan.append( ( chunk_start, chunk_end, chunk_max_compressed, chunk_max_compressed + chunk_max_gather, ) ) chunk_start = chunk_end return chunk_plan class ComputePrefillMetadataKernel( VllmTritonJitKernel["ComputePrefillMetadataKernel.CompileKey"] ): @dataclass(frozen=True) class CompileKey: block_size: int @staticmethod @triton.jit(do_not_specialize=["num_prefills", "num_decodes", "window_size"]) def kernel( # Outputs prefill_gather_lens_ptr, # Inputs seq_lens_ptr, query_start_loc_ptr, num_prefills, num_decodes, window_size, BLOCK_SIZE: tl.constexpr, ): """Compute prefill gather_lens in a single pass.""" offset = tl.arange(0, BLOCK_SIZE) mask = offset < num_prefills # SM12x + Triton 3.6 raises IMA on out-of-bounds address arithmetic for # masked-off lanes even though the load mask gates the actual read, so # clamp the offset. Caller guarantees num_prefills > 0. safe_offset = tl.minimum(offset, num_prefills - 1) seq_len = tl.load(seq_lens_ptr + num_decodes + safe_offset, mask=mask) qsl_start = tl.load(query_start_loc_ptr + num_decodes + safe_offset, mask=mask) qsl_end = tl.load( query_start_loc_ptr + num_decodes + safe_offset + 1, mask=mask ) query_len = qsl_end - qsl_start prefix_len = seq_len - query_len gather_len = query_len + tl.minimum(prefix_len, window_size - 1) tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask) def dispatch( # type: ignore[override] self, *, num_prefills: int, ) -> CompileKey: return self.CompileKey( block_size=next_power_of_2(num_prefills), ) def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: scheduler_config = vllm_config.scheduler_config max_prefills = max( 1, min( scheduler_config.max_num_seqs, scheduler_config.max_num_batched_tokens, ), ) return self._trace_dispatch(self.dispatch)( num_prefills=WarmupIntRange(1, max_prefills + 1), ) def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: int32_ptr = TritonWarmupTensor(torch.int32) return dict( prefill_gather_lens=int32_ptr, seq_lens=int32_ptr, query_start_loc=int32_ptr, num_prefills=compile_key.block_size, num_decodes=0, window_size=1, ) @kernel_launcher def __call__( self, prefill_gather_lens: torch.Tensor, seq_lens: torch.Tensor, query_start_loc: torch.Tensor, num_prefills: int, num_decodes: int, window_size: int, ) -> LaunchSpec: compile_key = self.dispatch(num_prefills=num_prefills) return (1,), dict( BLOCK_SIZE=compile_key.block_size, ) class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): """Builds metadata for DeepseekV4 SWA cache. Similar to the indexer, this handles mixed batches by: 1. Using split_decodes_and_prefills() to determine the boundary 2. Building separate metadata for decode and prefill portions Supports: - Mixed decode/prefill batches - MTP (Multi-Token Prediction) where decode has query_len > 1 - Chunked prefill (aligns with the indexer's chunking) """ reorder_batch_threshold: int | None = None _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH supports_draft_decode_metadata_update = True def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) assert isinstance(self.kv_cache_spec, SlidingWindowMLASpec | MLAAttentionSpec) mla_spec = cast(SlidingWindowMLASpec | MLAAttentionSpec, self.kv_cache_spec) self.head_size = mla_spec.head_size # Already considered quantization. assert isinstance(mla_spec.tokens_per_state, int) self.compress_ratio = mla_spec.tokens_per_state self.block_size = mla_spec.block_size self.max_model_len = self.vllm_config.model_config.max_model_len self.max_num_batched_tokens = ( self.vllm_config.scheduler_config.max_num_batched_tokens ) # Handle MTP: adjust decode_threshold like the indexer does spec_config = self.vllm_config.speculative_config self.num_speculative_tokens = ( spec_config.num_speculative_tokens if spec_config else 0 ) # Decode can have query_len up to # 1 + (2 if parallel drafting else 1) * num_speculative_tokens. # sparse_swa has no MQA-vs-dense-MHA routing, so multi-token queries take # the prefill path and the decode/prefill split stays at that width. spec_mult = ( 2 if (spec_config is not None and spec_config.parallel_drafting) else 1 ) self.decode_threshold = 1 + spec_mult * self.num_speculative_tokens self.reorder_batch_threshold = None hf_config = self.vllm_config.model_config.hf_config assert hasattr(hf_config, "sliding_window") self.window_size = hf_config.sliding_window # Vision variant: image spans (up to vision_max_n_token tokens) are # visible bidirectionally, so prefill index rows widen from # window_size to window_size + max_image_tokens. Text-only models keep # max_image_tokens == 0 and take the original code paths everywhere. self.max_image_tokens = ( getattr(hf_config, "vision_max_n_token", 0) if getattr(hf_config, "vision_n_layers", 0) > 0 else 0 ) self.prefill_index_width = self.window_size + self.max_image_tokens # Detect which DeepseekV4 layer types this model uses so we only build a # FlashMLA tile-scheduler plan for types that will actually be called. # Models without compress_ratios (pure SWA) fall back to swaonly. compress_ratios = getattr(hf_config, "compress_ratios", None) or [1] self._layer_types: set[str] = set() for ratio in compress_ratios: self._layer_types.add(_layer_type_for(int(ratio))) max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens self.token_to_req_indices = torch.zeros( max_tokens, dtype=torch.int32, device=self.device, ) self.decode_swa_indices = torch.zeros( max_tokens, 1, self.window_size, dtype=torch.int32, device=self.device, ) self.decode_swa_lens = torch.zeros( max_tokens, dtype=torch.int32, device=self.device, ) # Allocated unconditionally — consumer picks paged-direct vs dequant # at call time. self.prefill_swa_indices = torch.zeros( max_tokens, 1, self.prefill_index_width, dtype=torch.int32, device=self.device, ) self.prefill_swa_lens = torch.zeros( max_tokens, dtype=torch.int32, device=self.device, ) # In-image visibility side buffers (vision variant only). The span # CSR is uploaded per step from mm_req_doc_ranges; span capacity is # bounded by max_tokens since each span holds at least one token. if self.max_image_tokens > 0: self.left_visible = torch.zeros( max_tokens, dtype=torch.int32, device=self.device ) self.right_visible = torch.zeros( max_tokens, dtype=torch.int32, device=self.device ) self.span_indptr = torch.zeros( self.vllm_config.scheduler_config.max_num_seqs + 1, dtype=torch.int32, device=self.device, ) self.span_starts = torch.zeros( max_tokens, dtype=torch.int32, device=self.device ) self.span_ends = torch.zeros( max_tokens, dtype=torch.int32, device=self.device ) self.is_valid_token = torch.zeros( max_tokens, dtype=torch.bool, device=self.device, ) # DSpark draft: the block is non-causal (every query attends to the # trailing window of context PLUS all query tokens, including future ones), # so its per-token index list is wider than `window_size`. The kernel pads # the q-head count to B_TOPK. Pad to a kernel-supported width; the logical # SWA window remains unchanged when the padded matrix is built. self.is_dspark = spec_config is not None and spec_config.use_dspark() self.noncausal_index_width = ( get_dspark_swa_index_width( self.window_size, self.num_speculative_tokens, ) if self.is_dspark else 0 ) self.decode_swa_indices_noncausal: torch.Tensor | None = None self._max_tokens = max_tokens def build( self, common_prefix_len: int, common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> DeepseekSparseSWAMetadata: """Build SWA metadata for mixed decode/prefill batches. The batch is assumed to be reordered with decodes first (by vLLM scheduler). We use split_decodes_and_prefills() to find the boundary, then build separate window_topk_idxs for each portion. For prefill, we use chunked prefill to align with the indexer's chunking. """ seq_lens = common_attn_metadata.seq_lens seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound query_start_loc = common_attn_metadata.query_start_loc query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu block_table = common_attn_metadata.block_table_tensor slot_mapping = common_attn_metadata.slot_mapping # Split into decode and prefill portions using configurable threshold (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) = ( split_decodes_and_prefills( common_attn_metadata, decode_threshold=self.decode_threshold ) ) # NOTE: Ensure all metadata tensors maintain fixed memory addresses # for CUDA graph compatibility. token_to_req_indices = common_attn_metadata.token_to_req_indices( self.token_to_req_indices ) is_valid_token = self.is_valid_token[: slot_mapping.shape[0]] is_valid_token.copy_(slot_mapping >= 0) non_causal = not common_attn_metadata.causal decode_swa_width = ( self.noncausal_index_width if non_causal else self.window_size ) decode_swa_indices = self.decode_swa_indices if num_decode_tokens > 0: self.decode_swa_lens[num_decode_tokens:] = 0 if non_causal: assert self.is_dspark, ( "Non-causal DeepseekV4 SWA is only supported for the DSpark " "speculation mode, but causal=False was set without DSpark." ) if self.decode_swa_indices_noncausal is None: self.decode_swa_indices_noncausal = torch.zeros( self._max_tokens, 1, self.noncausal_index_width, dtype=torch.int32, device=self.device, ) decode_swa_indices = self.decode_swa_indices_noncausal _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL( decode_swa_indices, self.decode_swa_lens, self.window_size, self.noncausal_index_width, query_start_loc, seq_lens, token_to_req_indices, is_valid_token, block_table, self.block_size, num_tokens=num_decode_tokens, token_offset=0, ) else: _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( decode_swa_indices, self.decode_swa_lens, self.window_size, decode_swa_indices.shape[-1], self.decode_swa_lens, # unused (HAS_IMAGE=False) self.decode_swa_lens, # unused (HAS_IMAGE=False) query_start_loc, seq_lens, token_to_req_indices, is_valid_token, block_table, self.block_size, num_tokens=num_decode_tokens, token_offset=0, ) # Vision variant: per-token in-image visibility for prefill tokens. # Decode tokens are always past the image spans (spans are prefilled # atomically), so the decode path above never needs them. prefill_left_visible: torch.Tensor | None = None prefill_right_visible: torch.Tensor | None = None mm_ranges = common_attn_metadata.mm_req_doc_ranges if ( self.max_image_tokens > 0 and num_prefill_tokens > 0 and mm_ranges and any(mm_ranges.values()) ): prefill_left_visible, prefill_right_visible = self._build_image_visibility( common_attn_metadata.num_reqs, mm_ranges, num_decode_tokens, num_prefill_tokens, seq_lens, query_start_loc, token_to_req_indices, ) # Prefill SWA indices live in paged coordinates. `token_offset` lets # the kernel read is_valid_token / token_to_req_indices at absolute # prefill positions while writing output starting at index 0. if num_prefill_tokens > 0: has_image = prefill_left_visible is not None prefill_swa_indices = self.prefill_swa_indices[:num_prefill_tokens] prefill_swa_lens = self.prefill_swa_lens[:num_prefill_tokens] _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( prefill_swa_indices, prefill_swa_lens, self.window_size, self.prefill_index_width, prefill_left_visible if has_image else prefill_swa_lens, prefill_right_visible if has_image else prefill_swa_lens, query_start_loc, seq_lens, token_to_req_indices, is_valid_token, block_table, self.block_size, num_tokens=num_prefill_tokens, token_offset=num_decode_tokens, has_image=has_image, ) # Pre-compute DeepseekV4 prefill metadata shared across all attention layers. deepseek_v4_fields = self._build_deepseek_v4_metadata( num_decodes, num_prefills, seq_lens, seq_lens_cpu, query_start_loc, query_start_loc_cpu, ) # Per-layer-type tile-scheduler plan holders. Empty FlashMLASchedMeta # per present DeepseekV4 layer type; the first flash_mla_with_kvcache call of # each type triggers the planner and all same-type layers reuse the # resulting plan for the rest of the step. tile_sched = self.build_tile_scheduler(num_decode_tokens) return DeepseekSparseSWAMetadata( seq_lens=seq_lens, query_start_loc=query_start_loc, query_start_loc_cpu=query_start_loc_cpu, block_table=block_table, slot_mapping=slot_mapping, is_valid_token=is_valid_token, token_to_req_indices=token_to_req_indices, decode_swa_indices=decode_swa_indices[:num_decode_tokens], decode_swa_lens=self.decode_swa_lens[:num_decode_tokens], decode_swa_width=decode_swa_width, prefill_swa_indices=( self.prefill_swa_indices[:num_prefill_tokens] if num_prefill_tokens > 0 else None ), prefill_swa_lens=( self.prefill_swa_lens[:num_prefill_tokens] if num_prefill_tokens > 0 else None ), prefill_left_visible=prefill_left_visible, prefill_right_visible=prefill_right_visible, block_size=self.block_size, num_decodes=num_decodes, num_prefills=num_prefills, num_decode_tokens=num_decode_tokens, num_prefill_tokens=num_prefill_tokens, # Upper bound on decode-split rows for the kernel's max_q_len # hint. common max_query_len bounds every row (scheduled max under # adaptive verification), clamped to what the split can admit so a # mixed batch's prefill max does not inflate decode scheduling. max_decode_query_len=min( common_attn_metadata.max_query_len, self.decode_threshold ), tile_sched_swaonly=tile_sched[_LAYER_TYPE_SWAONLY], tile_sched_c4a=tile_sched[_LAYER_TYPE_C4A], tile_sched_c128a=tile_sched[_LAYER_TYPE_C128A], tile_sched_c1a=tile_sched[_LAYER_TYPE_C1A], tile_sched_c2a=tile_sched[_LAYER_TYPE_C2A], **deepseek_v4_fields, # type: ignore[arg-type] ) def _build_image_visibility( self, num_reqs: int, mm_ranges: dict[int, list[tuple[int, int]]], num_decode_tokens: int, num_prefill_tokens: int, seq_lens: torch.Tensor, query_start_loc: torch.Tensor, token_to_req_indices: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: """Upload image spans and compute per-token in-image visibility. Returns (left, right) int32 tensors indexed by absolute (decode-first) token position; only prefill rows are meaningful. """ indptr = [0] * (num_reqs + 1) starts: list[int] = [] ends: list[int] = [] for req_idx in range(num_reqs): for span_start, span_end in mm_ranges.get(req_idx, ()): starts.append(span_start) ends.append(span_end) indptr[req_idx + 1] = len(starts) num_spans = len(starts) assert num_spans <= self.span_starts.shape[0] self.span_indptr[: num_reqs + 1].copy_(torch.tensor(indptr, dtype=torch.int32)) if num_spans > 0: self.span_starts[:num_spans].copy_(torch.tensor(starts, dtype=torch.int32)) self.span_ends[:num_spans].copy_(torch.tensor(ends, dtype=torch.int32)) _compute_image_visibility_kernel[(num_prefill_tokens,)]( self.left_visible, self.right_visible, self.span_indptr, self.span_starts, self.span_ends, query_start_loc, seq_lens, token_to_req_indices, self.max_image_tokens, token_offset=num_decode_tokens, ) end = num_decode_tokens + num_prefill_tokens return self.left_visible[:end], self.right_visible[:end] def update_draft_decode_metadata( self, metadata: DeepseekSparseSWAMetadata, ) -> None: if metadata.num_decode_tokens == 0: return assert metadata.query_start_loc is not None assert metadata.seq_lens is not None assert metadata.token_to_req_indices is not None assert metadata.is_valid_token is not None assert metadata.decode_swa_indices is not None assert metadata.decode_swa_lens is not None _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( metadata.decode_swa_indices, metadata.decode_swa_lens, metadata.decode_swa_indices.shape[-1], metadata.decode_swa_indices.shape[-1], metadata.decode_swa_lens, # unused (HAS_IMAGE=False) metadata.decode_swa_lens, # unused (HAS_IMAGE=False) metadata.query_start_loc, metadata.seq_lens, metadata.token_to_req_indices, metadata.is_valid_token, metadata.block_table, self.block_size, num_tokens=metadata.num_decode_tokens, token_offset=0, ) tile_sched = self.build_tile_scheduler(metadata.num_decode_tokens) metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY] metadata.tile_sched_c4a = tile_sched[_LAYER_TYPE_C4A] metadata.tile_sched_c128a = tile_sched[_LAYER_TYPE_C128A] metadata.tile_sched_c1a = tile_sched[_LAYER_TYPE_C1A] metadata.tile_sched_c2a = tile_sched[_LAYER_TYPE_C2A] metadata.flashinfer_sparse_index_cache.clear() def build_tile_scheduler( self, num_decode_tokens: int ) -> dict[str, FlashMLASchedMeta | None]: """Allocate one empty ``FlashMLASchedMeta`` per present DeepseekV4 layer type. Returned instances have ``tile_scheduler_metadata`` / ``num_splits`` set to ``None``; the FlashMLA C++ decode path will allocate them and run the tile-scheduler planner on the first ``flash_mla_with_kvcache`` call of each type. Subsequent same-type calls reuse the plan because the tensors (and ``have_initialized``) are populated on the struct. Returns all-``None`` when there are no decode tokens this step, so ``_forward_decode`` sees a clean sentinel. """ out: dict[str, FlashMLASchedMeta | None] = { _LAYER_TYPE_SWAONLY: None, _LAYER_TYPE_C4A: None, _LAYER_TYPE_C128A: None, _LAYER_TYPE_C1A: None, _LAYER_TYPE_C2A: None, } if ( num_decode_tokens == 0 or current_platform.is_rocm() or current_platform.is_xpu() or current_platform.is_device_capability_family(120) ): return out for layer_type in self._layer_types: # get_mla_metadata() is the official FlashMLA entry point that # returns a fresh empty FlashMLASchedMeta; using it keeps this # call site aligned with the rest of the vLLM FlashMLA backends # that already go through the same stub. out[layer_type] = get_mla_metadata()[0] return out def _build_deepseek_v4_metadata( self, num_decodes: int, num_prefills: int, seq_lens: torch.Tensor, seq_lens_cpu: torch.Tensor | None, query_start_loc: torch.Tensor, query_start_loc_cpu: torch.Tensor, ) -> dict[str, torch.Tensor | int | None]: """Pre-compute DeepseekV4 prefill metadata during the metadata build phase. Returns a dict of keyword arguments to pass to the DeepseekSparseSWAMetadata constructor. Note: C128A sparse metadata is computed by the FlashMLASparse builder (which owns the C128A block_table), not here. """ result: dict[str, torch.Tensor | int | None] = {} # --- Prefill query metadata (single Triton kernel + CPU slicing) --- if num_prefills > 0: assert seq_lens_cpu is not None pfx_gather_lens = torch.empty( num_prefills, dtype=torch.int32, device=seq_lens.device ) _COMPUTE_PREFILL_METADATA_KERNEL( pfx_gather_lens, seq_lens, query_start_loc, num_prefills, num_decodes, self.window_size, ) result["prefill_seq_lens"] = seq_lens[num_decodes:] result["prefill_seq_lens_cpu"] = seq_lens_cpu[num_decodes:] result["prefill_gather_lens"] = pfx_gather_lens result["prefill_query_lens_cpu"] = ( query_start_loc_cpu[num_decodes + 1 : num_decodes + num_prefills + 1] - query_start_loc_cpu[num_decodes : num_decodes + num_prefills] ).to(dtype=torch.int32) result["prefill_window_size"] = self.window_size result["prefill_max_model_len"] = self.max_model_len result["prefill_max_num_batched_tokens"] = self.max_num_batched_tokens return result @triton.jit(do_not_specialize=["token_offset"]) def _compute_image_visibility_kernel( left_visible_ptr, right_visible_ptr, span_indptr_ptr, span_starts_ptr, span_ends_ptr, query_start_loc_ptr, seq_lens_ptr, token_to_req_indices_ptr, max_image_tokens, token_offset, ): """Per-token in-image visible counts (port of `get_image_visible`). One program per prefill token. A token at position pos inside a span [span_start, span_end] sees min(pos - span_start, max_image_tokens - 1) extra tokens to its left and min(span_end - pos, max_image_tokens) to its right; tokens outside every span get 0/0 (plain causal window). """ pid = tl.program_id(0) token_idx = pid + token_offset req_idx = tl.load(token_to_req_indices_ptr + token_idx) query_start = tl.load(query_start_loc_ptr + req_idx) query_end = tl.load(query_start_loc_ptr + req_idx + 1) seq_len = tl.load(seq_lens_ptr + req_idx) pos = seq_len - (query_end - query_start) + token_idx - query_start span_lo = tl.load(span_indptr_ptr + req_idx) span_hi = tl.load(span_indptr_ptr + req_idx + 1) left = tl.zeros((), dtype=tl.int32) right = tl.zeros((), dtype=tl.int32) for i in range(span_lo, span_hi): span_start = tl.load(span_starts_ptr + i) span_end = tl.load(span_ends_ptr + i) in_span = (pos >= span_start) & (pos <= span_end) left = tl.where( in_span, tl.minimum(pos - span_start, max_image_tokens - 1), left ) right = tl.where(in_span, tl.minimum(span_end - pos, max_image_tokens), right) tl.store(left_visible_ptr + token_idx, left) tl.store(right_visible_ptr + token_idx, right) # TODO(ben): unify this kernel to reduce duplication @triton.jit( do_not_specialize=[ "swa_indices_stride", "block_table_stride", "token_offset", ] ) def _compute_swa_indices_and_lens_kernel( swa_indices_ptr, swa_indices_stride, swa_lens_ptr, window_size, index_width, left_visible_ptr, right_visible_ptr, query_start_loc_ptr, seq_lens_ptr, token_to_req_indices_ptr, is_valid_token_ptr, block_table_ptr, block_table_stride, block_size, token_offset, HAS_IMAGE: tl.constexpr, TRITON_BLOCK_SIZE: tl.constexpr, ): pid = tl.program_id(0) token_idx = pid + token_offset is_valid = tl.load(is_valid_token_ptr + token_idx) if not is_valid: tl.store(swa_lens_ptr + pid, 0) # Clear the row so a padded token cannot gather through stale indices. for i in range(0, index_width, TRITON_BLOCK_SIZE): offset = i + tl.arange(0, TRITON_BLOCK_SIZE) tl.store( swa_indices_ptr + pid * swa_indices_stride + offset, -1, mask=offset < index_width, ) return req_idx = tl.load(token_to_req_indices_ptr + token_idx) query_start = tl.load(query_start_loc_ptr + req_idx) query_end = tl.load(query_start_loc_ptr + req_idx + 1) query_len = query_end - query_start seq_len = tl.load(seq_lens_ptr + req_idx) prefix_len = seq_len - query_len pos = prefix_len + token_idx - query_start if HAS_IMAGE: # In-image bidirectional visibility widens the window: the window # starts up to max(left - (window - 1), 0) positions earlier and # extends `right` positions past the query token. left = tl.load(left_visible_ptr + token_idx) right = tl.load(right_visible_ptr + token_idx) else: left = 0 right = 0 left_add = tl.maximum(left - (window_size - 1), 0) start_pos = tl.maximum(pos - (window_size - 1) - left_add, 0) end_pos = pos + right + 1 swa_len = end_pos - start_pos tl.store(swa_lens_ptr + pid, swa_len) for i in range(0, index_width, TRITON_BLOCK_SIZE): offset = i + tl.arange(0, TRITON_BLOCK_SIZE) pos_offset = start_pos + offset block_indices = pos_offset // block_size block_numbers = tl.load( block_table_ptr + req_idx * block_table_stride + block_indices, mask=pos_offset < end_pos, ) block_offsets = pos_offset % block_size slot_ids = block_numbers * block_size + block_offsets slot_ids = tl.where(offset < swa_len, slot_ids, -1) tl.store( swa_indices_ptr + pid * swa_indices_stride + offset, slot_ids, mask=offset < index_width, ) class ComputeSWAIndicesAndLensKernel( VllmTritonJitKernel["ComputeSWAIndicesAndLensKernel.CompileKey"] ): @dataclass(frozen=True) class CompileKey: window_size: int index_width: int has_image: int block_size: int triton_block_size: int kernel = staticmethod(_compute_swa_indices_and_lens_kernel) def dispatch( # type: ignore[override] self, **compile_key_fields: int, ) -> CompileKey: return self.CompileKey( **compile_key_fields, triton_block_size=1024, ) def get_warmup_keys( self, *, window_size: int, block_size: int, max_image_tokens: int = 0, ) -> list[CompileKey]: # Decode rows always use the plain causal window. Vision models widen # prefill rows by max_image_tokens and additionally need the in-image # bidirectional variant for batches that contain image spans. width_has_image = [(window_size, 0)] if max_image_tokens > 0: width_has_image += [ (window_size + max_image_tokens, 0), (window_size + max_image_tokens, 1), ] return self._trace_dispatch(self.dispatch)( zip_inputs( *( dict(index_width=index_width, has_image=has_image) for index_width, has_image in width_has_image ) ), window_size=window_size, block_size=block_size, ) def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: int32_ptr = TritonWarmupTensor(torch.int32) return dict( swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), swa_lens=int32_ptr, window_size=compile_key.window_size, index_width=compile_key.index_width, left_visible=int32_ptr, right_visible=int32_ptr, query_start_loc=int32_ptr, seq_lens=int32_ptr, token_to_req_indices=int32_ptr, is_valid_token=TritonWarmupTensor(torch.bool), block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), block_size=compile_key.block_size, num_tokens=1, token_offset=0, has_image=bool(compile_key.has_image), ) @kernel_launcher def __call__( self, swa_indices: torch.Tensor, swa_lens: torch.Tensor, window_size: int, index_width: int, left_visible: torch.Tensor, right_visible: torch.Tensor, query_start_loc: torch.Tensor, seq_lens: torch.Tensor, token_to_req_indices: torch.Tensor, is_valid_token: torch.Tensor, block_table: torch.Tensor, block_size: int, *, num_tokens: int, token_offset: int, has_image: bool = False, ) -> LaunchSpec: return (num_tokens,), dict( swa_indices_stride=swa_indices.stride(0), block_table_stride=block_table.stride(0), HAS_IMAGE=has_image, TRITON_BLOCK_SIZE=1024, ) class ComputeDSparkNoncausalSWAIndicesKernel( VllmTritonJitKernel["ComputeDSparkNoncausalSWAIndicesKernel.CompileKey"] ): @dataclass(frozen=True) class CompileKey: window_size: int index_width: int block_size: int triton_block_size: int @staticmethod @triton.jit( do_not_specialize=[ "swa_indices_stride", "block_table_stride", "token_offset", ] ) def kernel( swa_indices_ptr, swa_indices_stride, swa_lens_ptr, window_size, index_width, query_start_loc_ptr, seq_lens_ptr, token_to_req_indices_ptr, is_valid_token_ptr, block_table_ptr, block_table_stride, block_size, token_offset, TRITON_BLOCK_SIZE: tl.constexpr, ): """Non-causal per-token indices for the DSpark draft block. Here, we populate the topk indices with the trailing window of context tokens, plus all query tokens (including future ones). """ pid = tl.program_id(0) token_idx = pid + token_offset is_valid = tl.load(is_valid_token_ptr + token_idx) if not is_valid: tl.store(swa_lens_ptr + pid, 0) # Clear the row so a padded token cannot gather through stale indices. for i in range(0, index_width, TRITON_BLOCK_SIZE): offset = i + tl.arange(0, TRITON_BLOCK_SIZE) tl.store( swa_indices_ptr + pid * swa_indices_stride + offset, -1, mask=offset < index_width, ) return req_idx = tl.load(token_to_req_indices_ptr + token_idx) query_start = tl.load(query_start_loc_ptr + req_idx) query_end = tl.load(query_start_loc_ptr + req_idx + 1) query_len = query_end - query_start seq_len = tl.load(seq_lens_ptr + req_idx) prefix_len = seq_len - query_len # Block-anchored window (shared by every token in the block) + full block. start_pos = tl.maximum(prefix_len - window_size, 0) end_pos = seq_len swa_len = end_pos - start_pos tl.store(swa_lens_ptr + pid, swa_len) for i in range(0, index_width, TRITON_BLOCK_SIZE): offset = i + tl.arange(0, TRITON_BLOCK_SIZE) pos_offset = start_pos + offset block_indices = pos_offset // block_size block_numbers = tl.load( block_table_ptr + req_idx * block_table_stride + block_indices, mask=pos_offset < end_pos, ) block_offsets = pos_offset % block_size slot_ids = block_numbers * block_size + block_offsets slot_ids = tl.where(offset < swa_len, slot_ids, -1) tl.store( swa_indices_ptr + pid * swa_indices_stride + offset, slot_ids, mask=offset < index_width, ) def dispatch( # type: ignore[override] self, *, window_size: int, num_speculative_tokens: int, block_size: int, ) -> CompileKey: return self.CompileKey( window_size=window_size, index_width=get_dspark_swa_index_width(window_size, num_speculative_tokens), block_size=block_size, triton_block_size=1024, ) def get_warmup_keys( self, *, window_size: int, num_speculative_tokens: int, block_size: int, ) -> list[CompileKey]: return self._trace_dispatch(self.dispatch)( window_size=window_size, num_speculative_tokens=num_speculative_tokens, block_size=block_size, ) def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: int32_ptr = TritonWarmupTensor(torch.int32) return dict( swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), swa_lens=int32_ptr, window_size=compile_key.window_size, index_width=compile_key.index_width, query_start_loc=int32_ptr, seq_lens=int32_ptr, token_to_req_indices=int32_ptr, is_valid_token=TritonWarmupTensor(torch.bool), block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), block_size=compile_key.block_size, num_tokens=1, token_offset=0, ) @kernel_launcher def __call__( self, swa_indices: torch.Tensor, swa_lens: torch.Tensor, window_size: int, index_width: int, query_start_loc: torch.Tensor, seq_lens: torch.Tensor, token_to_req_indices: torch.Tensor, is_valid_token: torch.Tensor, block_table: torch.Tensor, block_size: int, *, num_tokens: int, token_offset: int, ) -> LaunchSpec: return (num_tokens,), dict( swa_indices_stride=swa_indices.stride(0), block_table_stride=block_table.stride(0), TRITON_BLOCK_SIZE=1024, ) _COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel() _COMPUTE_SWA_INDICES_AND_LENS_KERNEL = ComputeSWAIndicesAndLensKernel() _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL = ComputeDSparkNoncausalSWAIndicesKernel()