| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass, field |
| 4 | from typing import Any, ClassVar, cast |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config |
| 9 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 10 | from vllm.model_executor.warmup.jit_warmup import ( |
| 11 | WarmupIntRange, |
| 12 | kernel_launcher, |
| 13 | zip_inputs, |
| 14 | ) |
| 15 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 16 | LaunchSpec, |
| 17 | TritonWarmupTensor, |
| 18 | VllmTritonJitKernel, |
| 19 | ) |
| 20 | from vllm.platforms import current_platform |
| 21 | from vllm.triton_utils import tl, triton |
| 22 | from vllm.utils.math_utils import cdiv, next_power_of_2 |
| 23 | from vllm.v1.attention.backend import ( |
| 24 | AttentionBackend, |
| 25 | AttentionCGSupport, |
| 26 | AttentionMetadataBuilder, |
| 27 | CommonAttentionMetadata, |
| 28 | MultipleOf, |
| 29 | max_decode_query_len, |
| 30 | ) |
| 31 | from vllm.v1.attention.backends.mla.compressor_utils import ( |
| 32 | get_dspark_swa_index_width, |
| 33 | ) |
| 34 | from vllm.v1.attention.backends.utils import split_decodes_and_prefills |
| 35 | from vllm.v1.attention.ops.flashmla import FlashMLASchedMeta, get_mla_metadata |
| 36 | from vllm.v1.kv_cache_interface import ( |
| 37 | KVCacheSpec, |
| 38 | MLAAttentionSpec, |
| 39 | SlidingWindowMLASpec, |
| 40 | get_kv_quant_mode, |
| 41 | ) |
| 42 | |
| 43 | # DeepseekV4 decode layer types, keyed by compress_ratio. Each type has a distinct |
| 44 | # (topk, extra_topk, extra_page_block_size) config, so they cannot share a |
| 45 | # FlashMLA tile-scheduler plan. Within a type, all ~60 DeepseekV4 layers share one |
| 46 | # plan per step because b / s_q / h_q / page_block_sizes / topks are identical. |
| 47 | _LAYER_TYPE_SWAONLY = "swaonly" |
| 48 | _LAYER_TYPE_C4A = "c4a" |
| 49 | _LAYER_TYPE_C128A = "c128a" |
| 50 | # v4.1 ratio-1 / ratio-2 indexer layers: same indexer-path shape as C4A, but |
| 51 | # their compressed page block is block_size // ratio, so each gets its own plan. |
| 52 | # v4.1 builders classify layer types themselves (ratio 0 is SWA-only there); see |
| 53 | # deepseek_v41/sparse_mla.py. |
| 54 | _LAYER_TYPE_C1A = "c1a" |
| 55 | _LAYER_TYPE_C2A = "c2a" |
| 56 | |
| 57 | |
| 58 | def _layer_type_for(compress_ratio: int) -> str: |
| 59 | if compress_ratio <= 1: |
| 60 | return _LAYER_TYPE_SWAONLY |
| 61 | if compress_ratio == 2: |
| 62 | return _LAYER_TYPE_C2A |
| 63 | if compress_ratio == 4: |
| 64 | return _LAYER_TYPE_C4A |
| 65 | if compress_ratio == 128: |
| 66 | return _LAYER_TYPE_C128A |
| 67 | raise ValueError( |
| 68 | f"Unsupported DeepseekV4 compress_ratio={compress_ratio}; " |
| 69 | "expected 1, 2, 4, or 128." |
| 70 | ) |
| 71 | |
| 72 | |
| 73 | class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase): |
| 74 | def __init__( |
| 75 | self, |
| 76 | head_dim: int, |
| 77 | window_size: int, |
| 78 | dtype: torch.dtype, |
| 79 | prefix: str, |
| 80 | cache_config: CacheConfig, |
| 81 | backend_cls: "type[AttentionBackend] | None" = None, |
| 82 | block_size: int = 64, |
| 83 | packed_bytes_per_token: int = 584, |
| 84 | packed_page_alignment: int = 576, |
| 85 | bounded_replay: bool = False, |
| 86 | ): |
| 87 | super().__init__() |
| 88 | self.backend_cls = backend_cls or DeepseekSparseSWABackend |
| 89 | # DeepseekV4.1 SWA bounded replay. |
| 90 | self.bounded_replay = bounded_replay |
| 91 | self.kv_cache = torch.tensor([]) |
| 92 | self.head_dim = head_dim |
| 93 | self.window_size = window_size |
| 94 | self.prefix = prefix |
| 95 | self.cache_config = cache_config |
| 96 | self.dtype = dtype |
| 97 | self.packed_bytes_per_token = packed_bytes_per_token |
| 98 | self.packed_page_alignment = packed_page_alignment |
| 99 | compilation_config = get_current_vllm_config().compilation_config |
| 100 | if prefix in compilation_config.static_forward_context: |
| 101 | raise ValueError(f"Duplicate layer name: {prefix}") |
| 102 | compilation_config.static_forward_context[prefix] = self |
| 103 | |
| 104 | # Any multiple of 32; the sparse decode kernels take the page size at |
| 105 | # runtime. |
| 106 | self.block_size = block_size |
| 107 | # uint8: fp8_ds_mla UE8M0 paged layout. bfloat16 / float8_e4m3fn: |
| 108 | # contiguous full-cache layout. |
| 109 | assert self.dtype in (torch.uint8, torch.bfloat16, torch.float8_e4m3fn) |
| 110 | |
| 111 | def bind_kv_cache(self, kv_cache: torch.Tensor) -> None: |
| 112 | # [B, H=1, N, C] -> [B, N, C] |
| 113 | self.kv_cache = kv_cache.squeeze(1) |
| 114 | |
| 115 | def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: |
| 116 | # fp8_ds_mla's UE8M0 paged layout rounds its page up to the decode |
| 117 | # kernel's TMA stride; contiguous bf16/fp8 cache uses the natural |
| 118 | # element-size page. |
| 119 | uses_fp8_ds_mla_layout = self.cache_config.cache_dtype in ( |
| 120 | "fp8_ds_mla", |
| 121 | "nvfp4_ds_mla", |
| 122 | ) |
| 123 | return SlidingWindowMLASpec( |
| 124 | bounded_replay=self.bounded_replay, |
| 125 | block_size=self.block_size, |
| 126 | num_kv_heads=1, |
| 127 | max_tp_shards=1, |
| 128 | head_size=self.head_dim, |
| 129 | dtype=self.dtype, |
| 130 | sliding_window=self.window_size, |
| 131 | cache_dtype_str=self.cache_config.cache_dtype, |
| 132 | state_content_bytes=( |
| 133 | self.packed_bytes_per_token if uses_fp8_ds_mla_layout else None |
| 134 | ), |
| 135 | # FlashMLA packing stride; 512B for FlashInfer sparse (#44577). |
| 136 | alignment=self.packed_page_alignment if uses_fp8_ds_mla_layout else 512, |
| 137 | model_version="deepseek_v4", |
| 138 | kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype), |
| 139 | ) |
| 140 | |
| 141 | def forward(self): ... |
| 142 | |
| 143 | def get_attn_backend(self) -> type[AttentionBackend]: |
| 144 | return self.backend_cls |
| 145 | |
| 146 | |
| 147 | class DeepseekSparseSWABackend(AttentionBackend): |
| 148 | @staticmethod |
| 149 | def get_name() -> str: |
| 150 | return "DEEPSEEK_SPARSE_SWA" |
| 151 | |
| 152 | @staticmethod |
| 153 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 154 | return [MultipleOf(32)] |
| 155 | |
| 156 | @classmethod |
| 157 | def get_preferred_block_size(cls, default_block_size: int) -> int: |
| 158 | return 256 |
| 159 | |
| 160 | @classmethod |
| 161 | def get_supported_head_sizes(cls) -> list[int]: |
| 162 | return [512] |
| 163 | |
| 164 | @staticmethod |
| 165 | def get_builder_cls() -> type["DeepseekSparseSWAMetadataBuilder"]: |
| 166 | if current_platform.is_rocm(): |
| 167 | from vllm.models.deepseek_v4.amd.rocm import ( |
| 168 | DeepseekV4ROCMAiterSparseSWAMetadataBuilder, |
| 169 | ) |
| 170 | |
| 171 | return DeepseekV4ROCMAiterSparseSWAMetadataBuilder |
| 172 | if current_platform.is_cpu(): |
| 173 | from vllm.models.deepseek_v4.cpu.cpu_mla import ( |
| 174 | DeepseekV4CPUSparseSWAMetadataBuilder, |
| 175 | ) |
| 176 | |
| 177 | return DeepseekV4CPUSparseSWAMetadataBuilder |
| 178 | return DeepseekSparseSWAMetadataBuilder |
| 179 | |
| 180 | |
| 181 | @dataclass |
| 182 | class DeepseekSparseSWAMetadata: |
| 183 | block_table: torch.Tensor |
| 184 | slot_mapping: torch.Tensor |
| 185 | block_size: int |
| 186 | seq_lens: torch.Tensor | None = None # [num_seqs] |
| 187 | query_start_loc: torch.Tensor | None = None # [num_seqs + 1] |
| 188 | query_start_loc_cpu: torch.Tensor | None = None # [num_seqs + 1] |
| 189 | |
| 190 | is_valid_token: torch.Tensor | None = None # [num_tokens] |
| 191 | token_to_req_indices: torch.Tensor | None = None # [num_tokens] |
| 192 | decode_swa_indices: torch.Tensor | None = None # [num_decode_tokens, width] |
| 193 | decode_swa_lens: torch.Tensor | None = None # [num_decode_tokens] |
| 194 | # window_size (causal) or noncausal_index_width (DSpark non-causal). |
| 195 | decode_swa_width: int = 0 |
| 196 | # Paged-coordinate prefill SWA indices/lens (FP8 paged-direct prefill). |
| 197 | prefill_swa_indices: torch.Tensor | None = ( |
| 198 | None # [num_prefill_tokens, 1, prefill_index_width] |
| 199 | ) |
| 200 | prefill_swa_lens: torch.Tensor | None = None # [num_prefill_tokens] |
| 201 | # In-image bidirectional visibility (vision variant): per-token counts of |
| 202 | # extra visible tokens to the left/right inside an image span. Indexed by |
| 203 | # absolute (decode-first) token position; only prefill rows are written. |
| 204 | # None when the model is text-only or the batch has no image spans. |
| 205 | prefill_left_visible: torch.Tensor | None = None |
| 206 | prefill_right_visible: torch.Tensor | None = None |
| 207 | # SWA bounded replay: [num_reqs] lower bound of the positions a request's |
| 208 | # window attention may read (SWA bounded replay); |
| 209 | # zeros when nothing replays. |
| 210 | replay_start: torch.Tensor | None = None |
| 211 | |
| 212 | # Number of decode/prefill requests/tokens (batch is reordered: decodes first) |
| 213 | num_decodes: int = 0 |
| 214 | num_prefills: int = 0 |
| 215 | num_decode_tokens: int = 0 |
| 216 | num_prefill_tokens: int = 0 |
| 217 | max_decode_query_len: int = 1 |
| 218 | flashinfer_decode_topk_lens: torch.Tensor | None = None |
| 219 | flashinfer_decode_seq_lens: torch.Tensor | None = None |
| 220 | |
| 221 | # Pre-computed prefill metadata shared across all DeepseekV4 attention layers. |
| 222 | prefill_seq_lens: torch.Tensor | None = None |
| 223 | prefill_seq_lens_cpu: torch.Tensor | None = None |
| 224 | prefill_gather_lens: torch.Tensor | None = None |
| 225 | prefill_query_lens_cpu: torch.Tensor | None = None |
| 226 | prefill_window_size: int = 0 |
| 227 | prefill_max_model_len: int = 0 |
| 228 | prefill_max_num_batched_tokens: int = 0 |
| 229 | |
| 230 | # Per-layer-type FlashMLA tile-scheduler metadata. One FlashMLASchedMeta |
| 231 | # per present DeepseekV4 layer type, shared across all ~60 layers of that type |
| 232 | # within a decode step. The first forward call of a given type triggers |
| 233 | # the in-kernel planner (which also allocates tile_scheduler_metadata and |
| 234 | # num_splits via PyTorch's graph-aware allocator); subsequent same-type |
| 235 | # calls skip planning and reuse the plan. Fresh instance per build(), so |
| 236 | # have_initialized is always False at the start of a step and the plan |
| 237 | # is re-derived from current seq_lens / topk_length on replay. |
| 238 | # None for layer types the model does not use (or when num_decode_tokens |
| 239 | # is zero). |
| 240 | tile_sched_swaonly: "FlashMLASchedMeta | None" = None |
| 241 | tile_sched_c4a: "FlashMLASchedMeta | None" = None |
| 242 | tile_sched_c128a: "FlashMLASchedMeta | None" = None |
| 243 | tile_sched_c1a: "FlashMLASchedMeta | None" = None |
| 244 | tile_sched_c2a: "FlashMLASchedMeta | None" = None |
| 245 | flashinfer_sparse_index_cache: dict[str, tuple[torch.Tensor, torch.Tensor]] = field( |
| 246 | default_factory=dict |
| 247 | ) |
| 248 | |
| 249 | def get_prefill_chunk_plan( |
| 250 | self, |
| 251 | compress_ratio: int, |
| 252 | prefill_chunk_size: int, |
| 253 | has_compressed: bool | None = None, |
| 254 | ) -> list[tuple[int, int, int, int]]: |
| 255 | if self.num_prefills == 0: |
| 256 | return [] |
| 257 | |
| 258 | assert self.prefill_seq_lens_cpu is not None |
| 259 | assert self.prefill_query_lens_cpu is not None |
| 260 | |
| 261 | # Whether the layer gathers a compressed-KV region into the prefill |
| 262 | # workspace. v4.0 callers keep the legacy default (ratio <= 1 means |
| 263 | # SWA-only); v4.1 passes has_compressed explicitly because its |
| 264 | # compress_ratio==1 layers DO have a full-length compressed cache. |
| 265 | if has_compressed is None: |
| 266 | has_compressed = compress_ratio > 1 |
| 267 | |
| 268 | # query_len <= max_num_batched_tokens and |
| 269 | # gather_len <= query_len + min(prefix_len, window_size - 1) (a replay |
| 270 | # boundary only shrinks it), so the worst-case gathered width is |
| 271 | # bounded by max_num_batched_tokens + window_size - 1. The compressed |
| 272 | # prefix pool is bounded by ceil(max_model_len / compress_ratio). The |
| 273 | # CPU gather lengths below are this upper bound, for chunk planning |
| 274 | # only; the kernels use the exact per-request values. |
| 275 | max_workspace_area = prefill_chunk_size * ( |
| 276 | (cdiv(self.prefill_max_model_len, compress_ratio) if has_compressed else 0) |
| 277 | + self.prefill_window_size |
| 278 | + self.prefill_max_num_batched_tokens |
| 279 | ) |
| 280 | prefix_lens_cpu = self.prefill_seq_lens_cpu - self.prefill_query_lens_cpu |
| 281 | gather_lens_cpu = self.prefill_query_lens_cpu + torch.clamp( |
| 282 | prefix_lens_cpu, min=0, max=self.prefill_window_size - 1 |
| 283 | ) |
| 284 | compressed_lens_cpu = ( |
| 285 | torch.div( |
| 286 | self.prefill_seq_lens_cpu, |
| 287 | compress_ratio, |
| 288 | rounding_mode="floor", |
| 289 | ) |
| 290 | if has_compressed |
| 291 | else torch.zeros_like(self.prefill_seq_lens_cpu) |
| 292 | ) |
| 293 | |
| 294 | chunk_plan: list[tuple[int, int, int, int]] = [] |
| 295 | chunk_start = 0 |
| 296 | while chunk_start < self.num_prefills: |
| 297 | chunk_max_compressed = int(compressed_lens_cpu[chunk_start].item()) |
| 298 | chunk_max_gather = int(gather_lens_cpu[chunk_start].item()) |
| 299 | chunk_end = chunk_start + 1 |
| 300 | |
| 301 | while chunk_end < self.num_prefills: |
| 302 | candidate_max_compressed = max( |
| 303 | chunk_max_compressed, |
| 304 | int(compressed_lens_cpu[chunk_end].item()), |
| 305 | ) |
| 306 | candidate_max_gather = max( |
| 307 | chunk_max_gather, |
| 308 | int(gather_lens_cpu[chunk_end].item()), |
| 309 | ) |
| 310 | candidate_width = candidate_max_compressed + candidate_max_gather |
| 311 | candidate_area = (chunk_end - chunk_start + 1) * candidate_width |
| 312 | if candidate_area > max_workspace_area: |
| 313 | break |
| 314 | chunk_max_compressed = candidate_max_compressed |
| 315 | chunk_max_gather = candidate_max_gather |
| 316 | chunk_end += 1 |
| 317 | |
| 318 | chunk_plan.append( |
| 319 | ( |
| 320 | chunk_start, |
| 321 | chunk_end, |
| 322 | chunk_max_compressed, |
| 323 | chunk_max_compressed + chunk_max_gather, |
| 324 | ) |
| 325 | ) |
| 326 | chunk_start = chunk_end |
| 327 | |
| 328 | return chunk_plan |
| 329 | |
| 330 | |
| 331 | class ComputePrefillMetadataKernel( |
| 332 | VllmTritonJitKernel["ComputePrefillMetadataKernel.CompileKey"] |
| 333 | ): |
| 334 | @dataclass(frozen=True) |
| 335 | class CompileKey: |
| 336 | block_size: int |
| 337 | |
| 338 | @staticmethod |
| 339 | @triton.jit(do_not_specialize=["num_prefills", "num_decodes", "window_size"]) |
| 340 | def kernel( |
| 341 | # Outputs |
| 342 | prefill_gather_lens_ptr, |
| 343 | # Inputs |
| 344 | seq_lens_ptr, |
| 345 | query_start_loc_ptr, |
| 346 | replay_start_ptr, |
| 347 | num_prefills, |
| 348 | num_decodes, |
| 349 | window_size, |
| 350 | BLOCK_SIZE: tl.constexpr, |
| 351 | ): |
| 352 | """Compute prefill gather_lens in a single pass.""" |
| 353 | offset = tl.arange(0, BLOCK_SIZE) |
| 354 | mask = offset < num_prefills |
| 355 | # SM12x + Triton 3.6 raises IMA on out-of-bounds address arithmetic for |
| 356 | # masked-off lanes even though the load mask gates the actual read, so |
| 357 | # clamp the offset. Caller guarantees num_prefills > 0. |
| 358 | safe_offset = tl.minimum(offset, num_prefills - 1) |
| 359 | |
| 360 | seq_len = tl.load(seq_lens_ptr + num_decodes + safe_offset, mask=mask) |
| 361 | qsl_start = tl.load(query_start_loc_ptr + num_decodes + safe_offset, mask=mask) |
| 362 | qsl_end = tl.load( |
| 363 | query_start_loc_ptr + num_decodes + safe_offset + 1, mask=mask |
| 364 | ) |
| 365 | |
| 366 | query_len = qsl_end - qsl_start |
| 367 | prefix_len = seq_len - query_len |
| 368 | # Context below replay_start has no window KV (SWA bounded replay). |
| 369 | replay_start = tl.load(replay_start_ptr + num_decodes + safe_offset, mask=mask) |
| 370 | visible_prefix = tl.maximum(prefix_len - replay_start, 0) |
| 371 | gather_len = query_len + tl.minimum(visible_prefix, window_size - 1) |
| 372 | |
| 373 | tl.store(prefill_gather_lens_ptr + offset, gather_len, mask=mask) |
| 374 | |
| 375 | def dispatch( # type: ignore[override] |
| 376 | self, |
| 377 | *, |
| 378 | num_prefills: int, |
| 379 | ) -> CompileKey: |
| 380 | return self.CompileKey( |
| 381 | block_size=next_power_of_2(num_prefills), |
| 382 | ) |
| 383 | |
| 384 | def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: |
| 385 | scheduler_config = vllm_config.scheduler_config |
| 386 | max_prefills = max( |
| 387 | 1, |
| 388 | min( |
| 389 | scheduler_config.max_num_seqs, |
| 390 | scheduler_config.max_num_batched_tokens, |
| 391 | ), |
| 392 | ) |
| 393 | return self._trace_dispatch(self.dispatch)( |
| 394 | num_prefills=WarmupIntRange(1, max_prefills + 1), |
| 395 | ) |
| 396 | |
| 397 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 398 | int32_ptr = TritonWarmupTensor(torch.int32) |
| 399 | return dict( |
| 400 | prefill_gather_lens=int32_ptr, |
| 401 | seq_lens=int32_ptr, |
| 402 | query_start_loc=int32_ptr, |
| 403 | replay_start=int32_ptr, |
| 404 | num_prefills=compile_key.block_size, |
| 405 | num_decodes=0, |
| 406 | window_size=1, |
| 407 | ) |
| 408 | |
| 409 | @kernel_launcher |
| 410 | def __call__( |
| 411 | self, |
| 412 | prefill_gather_lens: torch.Tensor, |
| 413 | seq_lens: torch.Tensor, |
| 414 | query_start_loc: torch.Tensor, |
| 415 | replay_start: torch.Tensor, |
| 416 | num_prefills: int, |
| 417 | num_decodes: int, |
| 418 | window_size: int, |
| 419 | ) -> LaunchSpec: |
| 420 | compile_key = self.dispatch(num_prefills=num_prefills) |
| 421 | return (1,), dict( |
| 422 | BLOCK_SIZE=compile_key.block_size, |
| 423 | ) |
| 424 | |
| 425 | |
| 426 | class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder): |
| 427 | """Builds metadata for DeepseekV4 SWA cache. |
| 428 | |
| 429 | Similar to the indexer, this handles mixed batches by: |
| 430 | 1. Using split_decodes_and_prefills() to determine the boundary |
| 431 | 2. Building separate metadata for decode and prefill portions |
| 432 | |
| 433 | Supports: |
| 434 | - Mixed decode/prefill batches |
| 435 | - MTP (Multi-Token Prediction) where decode has query_len > 1 |
| 436 | - Chunked prefill (aligns with the indexer's chunking) |
| 437 | """ |
| 438 | |
| 439 | reorder_batch_threshold: int | None = None |
| 440 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH |
| 441 | supports_draft_decode_metadata_update = True |
| 442 | |
| 443 | def __init__(self, *args, **kwargs): |
| 444 | super().__init__(*args, **kwargs) |
| 445 | assert isinstance(self.kv_cache_spec, SlidingWindowMLASpec | MLAAttentionSpec) |
| 446 | mla_spec = cast(SlidingWindowMLASpec | MLAAttentionSpec, self.kv_cache_spec) |
| 447 | self.head_size = mla_spec.head_size # Already considered quantization. |
| 448 | assert isinstance(mla_spec.tokens_per_state, int) |
| 449 | self.compress_ratio = mla_spec.tokens_per_state |
| 450 | self.block_size = mla_spec.block_size |
| 451 | self.max_model_len = self.vllm_config.model_config.max_model_len |
| 452 | self.max_num_batched_tokens = ( |
| 453 | self.vllm_config.scheduler_config.max_num_batched_tokens |
| 454 | ) |
| 455 | |
| 456 | # Handle MTP: adjust decode_threshold like the indexer does |
| 457 | spec_config = self.vllm_config.speculative_config |
| 458 | self.num_speculative_tokens = ( |
| 459 | spec_config.num_speculative_tokens if spec_config else 0 |
| 460 | ) |
| 461 | # sparse_swa has no MQA-vs-dense-MHA routing, so multi-token queries take |
| 462 | # the prefill path and the decode/prefill split stays at the widest |
| 463 | # decode. |
| 464 | self.decode_threshold = max_decode_query_len(self.vllm_config) |
| 465 | self.reorder_batch_threshold = None |
| 466 | |
| 467 | hf_config = self.vllm_config.model_config.hf_config |
| 468 | assert hasattr(hf_config, "sliding_window") |
| 469 | self.window_size = hf_config.sliding_window |
| 470 | |
| 471 | # V4 vision variant: image spans (up to vision_max_n_token tokens) are |
| 472 | # visible bidirectionally, so prefill index rows widen from |
| 473 | # window_size to window_size + max_image_tokens. The V4 config sets |
| 474 | # mm_prefix_clamp_sliding_window exactly for these in-kernel-widened |
| 475 | # ranges; V4.1 image tokens use the plain causal window, and text-only |
| 476 | # models keep max_image_tokens == 0 everywhere. |
| 477 | self.max_image_tokens = ( |
| 478 | getattr(hf_config, "vision_max_n_token", 0) |
| 479 | if getattr(hf_config, "mm_prefix_clamp_sliding_window", False) |
| 480 | else 0 |
| 481 | ) |
| 482 | self.prefill_index_width = self.window_size + self.max_image_tokens |
| 483 | |
| 484 | # Detect which DeepseekV4 layer types this model uses so we only build a |
| 485 | # FlashMLA tile-scheduler plan for types that will actually be called. |
| 486 | # Models without compress_ratios (pure SWA) fall back to swaonly. |
| 487 | compress_ratios = getattr(hf_config, "compress_ratios", None) or [1] |
| 488 | self._layer_types: set[str] = set() |
| 489 | for ratio in compress_ratios: |
| 490 | self._layer_types.add(_layer_type_for(int(ratio))) |
| 491 | |
| 492 | max_tokens = self.vllm_config.scheduler_config.max_num_batched_tokens |
| 493 | self.token_to_req_indices = torch.zeros( |
| 494 | max_tokens, |
| 495 | dtype=torch.int32, |
| 496 | device=self.device, |
| 497 | ) |
| 498 | self.decode_swa_indices = torch.zeros( |
| 499 | max_tokens, |
| 500 | 1, |
| 501 | self.window_size, |
| 502 | dtype=torch.int32, |
| 503 | device=self.device, |
| 504 | ) |
| 505 | self.decode_swa_lens = torch.zeros( |
| 506 | max_tokens, |
| 507 | dtype=torch.int32, |
| 508 | device=self.device, |
| 509 | ) |
| 510 | # Allocated unconditionally — consumer picks paged-direct vs dequant |
| 511 | # at call time. |
| 512 | self.prefill_swa_indices = torch.zeros( |
| 513 | max_tokens, |
| 514 | 1, |
| 515 | self.prefill_index_width, |
| 516 | dtype=torch.int32, |
| 517 | device=self.device, |
| 518 | ) |
| 519 | self.prefill_swa_lens = torch.zeros( |
| 520 | max_tokens, |
| 521 | dtype=torch.int32, |
| 522 | device=self.device, |
| 523 | ) |
| 524 | # In-image visibility side buffers (vision variant only). The span |
| 525 | # CSR is uploaded per step from mm_req_doc_ranges; span capacity is |
| 526 | # bounded by max_tokens since each span holds at least one token. |
| 527 | if self.max_image_tokens > 0: |
| 528 | self.left_visible = torch.zeros( |
| 529 | max_tokens, dtype=torch.int32, device=self.device |
| 530 | ) |
| 531 | self.right_visible = torch.zeros( |
| 532 | max_tokens, dtype=torch.int32, device=self.device |
| 533 | ) |
| 534 | self.span_indptr = torch.zeros( |
| 535 | self.vllm_config.scheduler_config.max_num_seqs + 1, |
| 536 | dtype=torch.int32, |
| 537 | device=self.device, |
| 538 | ) |
| 539 | self.span_starts = torch.zeros( |
| 540 | max_tokens, dtype=torch.int32, device=self.device |
| 541 | ) |
| 542 | self.span_ends = torch.zeros( |
| 543 | max_tokens, dtype=torch.int32, device=self.device |
| 544 | ) |
| 545 | self.is_valid_token = torch.zeros( |
| 546 | max_tokens, |
| 547 | dtype=torch.bool, |
| 548 | device=self.device, |
| 549 | ) |
| 550 | |
| 551 | # DSpark draft: the block is non-causal (every query attends to the |
| 552 | # trailing window of context PLUS all query tokens, including future ones), |
| 553 | # so its per-token index list is wider than `window_size`. The kernel pads |
| 554 | # the q-head count to B_TOPK. Pad to a kernel-supported width; the logical |
| 555 | # SWA window remains unchanged when the padded matrix is built. |
| 556 | self.is_dspark = spec_config is not None and spec_config.use_dspark() |
| 557 | self.noncausal_index_width = ( |
| 558 | get_dspark_swa_index_width( |
| 559 | self.window_size, |
| 560 | self.num_speculative_tokens, |
| 561 | ) |
| 562 | if self.is_dspark |
| 563 | else 0 |
| 564 | ) |
| 565 | self.decode_swa_indices_noncausal: torch.Tensor | None = None |
| 566 | self._max_tokens = max_tokens |
| 567 | # replay_start when no group replays: no lower bound. |
| 568 | self.no_replay_start = torch.zeros( |
| 569 | self.vllm_config.scheduler_config.max_num_seqs + 1, |
| 570 | dtype=torch.int32, |
| 571 | device=self.device, |
| 572 | ) |
| 573 | |
| 574 | def build( |
| 575 | self, |
| 576 | common_prefix_len: int, |
| 577 | common_attn_metadata: CommonAttentionMetadata, |
| 578 | fast_build: bool = False, |
| 579 | replay_start: torch.Tensor | None = None, |
| 580 | ) -> DeepseekSparseSWAMetadata: |
| 581 | """Build SWA metadata for mixed decode/prefill batches. |
| 582 | |
| 583 | The batch is assumed to be reordered with decodes first (by vLLM scheduler). |
| 584 | We use split_decodes_and_prefills() to find the boundary, then build |
| 585 | separate window_topk_idxs for each portion. |
| 586 | |
| 587 | For prefill, we use chunked prefill to align with the indexer's chunking. |
| 588 | |
| 589 | ``replay_start`` ([num_reqs], SWA bounded replay): the position from |
| 590 | which each request holds window KV; the model state passes it for every |
| 591 | batch of a replaying group, graph captures pass nothing. |
| 592 | """ |
| 593 | seq_lens = common_attn_metadata.seq_lens |
| 594 | seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound |
| 595 | query_start_loc = common_attn_metadata.query_start_loc |
| 596 | query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu |
| 597 | block_table = common_attn_metadata.block_table_tensor |
| 598 | slot_mapping = common_attn_metadata.slot_mapping |
| 599 | |
| 600 | # Split into decode and prefill portions using configurable threshold |
| 601 | (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens) = ( |
| 602 | split_decodes_and_prefills( |
| 603 | common_attn_metadata, decode_threshold=self.decode_threshold |
| 604 | ) |
| 605 | ) |
| 606 | |
| 607 | # NOTE: Ensure all metadata tensors maintain fixed memory addresses |
| 608 | # for CUDA graph compatibility. |
| 609 | token_to_req_indices = common_attn_metadata.token_to_req_indices( |
| 610 | self.token_to_req_indices |
| 611 | ) |
| 612 | |
| 613 | is_valid_token = self.is_valid_token[: slot_mapping.shape[0]] |
| 614 | torch.ge(slot_mapping, 0, out=is_valid_token) |
| 615 | if replay_start is None: |
| 616 | # Graph captures build without it; their dummy batches replay nothing. |
| 617 | replay_start = self.no_replay_start |
| 618 | |
| 619 | non_causal = not common_attn_metadata.causal |
| 620 | decode_swa_width = ( |
| 621 | self.noncausal_index_width if non_causal else self.window_size |
| 622 | ) |
| 623 | decode_swa_indices = self.decode_swa_indices |
| 624 | if num_decode_tokens > 0: |
| 625 | self.decode_swa_lens[num_decode_tokens:] = 0 |
| 626 | if non_causal: |
| 627 | assert self.is_dspark, ( |
| 628 | "Non-causal DeepseekV4 SWA is only supported for the DSpark " |
| 629 | "speculation mode, but causal=False was set without DSpark." |
| 630 | ) |
| 631 | if self.decode_swa_indices_noncausal is None: |
| 632 | self.decode_swa_indices_noncausal = torch.zeros( |
| 633 | self._max_tokens, |
| 634 | 1, |
| 635 | self.noncausal_index_width, |
| 636 | dtype=torch.int32, |
| 637 | device=self.device, |
| 638 | ) |
| 639 | decode_swa_indices = self.decode_swa_indices_noncausal |
| 640 | # No replay_start: decode rows sit above the hit, so a |
| 641 | # request's replayed window never reaches them. |
| 642 | _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL( |
| 643 | decode_swa_indices, |
| 644 | self.decode_swa_lens, |
| 645 | self.window_size, |
| 646 | self.noncausal_index_width, |
| 647 | query_start_loc, |
| 648 | seq_lens, |
| 649 | token_to_req_indices, |
| 650 | is_valid_token, |
| 651 | block_table, |
| 652 | self.block_size, |
| 653 | num_tokens=num_decode_tokens, |
| 654 | token_offset=0, |
| 655 | ) |
| 656 | else: |
| 657 | _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( |
| 658 | decode_swa_indices, |
| 659 | self.decode_swa_lens, |
| 660 | self.window_size, |
| 661 | decode_swa_indices.shape[-1], |
| 662 | self.decode_swa_lens, # unused (HAS_IMAGE=False) |
| 663 | self.decode_swa_lens, # unused (HAS_IMAGE=False) |
| 664 | query_start_loc, |
| 665 | seq_lens, |
| 666 | token_to_req_indices, |
| 667 | is_valid_token, |
| 668 | block_table, |
| 669 | self.block_size, |
| 670 | replay_start, |
| 671 | num_tokens=num_decode_tokens, |
| 672 | token_offset=0, |
| 673 | ) |
| 674 | |
| 675 | # Vision variant: per-token in-image visibility for prefill tokens. |
| 676 | # Decode tokens are always past the image spans (spans are prefilled |
| 677 | # atomically), so the decode path above never needs them. |
| 678 | prefill_left_visible: torch.Tensor | None = None |
| 679 | prefill_right_visible: torch.Tensor | None = None |
| 680 | mm_ranges = common_attn_metadata.mm_req_doc_ranges |
| 681 | if ( |
| 682 | self.max_image_tokens > 0 |
| 683 | and num_prefill_tokens > 0 |
| 684 | and mm_ranges |
| 685 | and any(mm_ranges.values()) |
| 686 | ): |
| 687 | prefill_left_visible, prefill_right_visible = self._build_image_visibility( |
| 688 | common_attn_metadata.num_reqs, |
| 689 | mm_ranges, |
| 690 | num_decode_tokens, |
| 691 | num_prefill_tokens, |
| 692 | seq_lens, |
| 693 | query_start_loc, |
| 694 | token_to_req_indices, |
| 695 | ) |
| 696 | |
| 697 | # Prefill SWA indices live in paged coordinates. `token_offset` lets |
| 698 | # the kernel read is_valid_token / token_to_req_indices at absolute |
| 699 | # prefill positions while writing output starting at index 0. |
| 700 | if num_prefill_tokens > 0: |
| 701 | has_image = prefill_left_visible is not None |
| 702 | prefill_swa_indices = self.prefill_swa_indices[:num_prefill_tokens] |
| 703 | prefill_swa_lens = self.prefill_swa_lens[:num_prefill_tokens] |
| 704 | _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( |
| 705 | prefill_swa_indices, |
| 706 | prefill_swa_lens, |
| 707 | self.window_size, |
| 708 | self.prefill_index_width, |
| 709 | prefill_left_visible if has_image else prefill_swa_lens, |
| 710 | prefill_right_visible if has_image else prefill_swa_lens, |
| 711 | query_start_loc, |
| 712 | seq_lens, |
| 713 | token_to_req_indices, |
| 714 | is_valid_token, |
| 715 | block_table, |
| 716 | self.block_size, |
| 717 | replay_start, |
| 718 | num_tokens=num_prefill_tokens, |
| 719 | token_offset=num_decode_tokens, |
| 720 | has_image=has_image, |
| 721 | ) |
| 722 | |
| 723 | # Pre-compute DeepseekV4 prefill metadata shared across all attention layers. |
| 724 | deepseek_v4_fields = self._build_deepseek_v4_metadata( |
| 725 | num_decodes, |
| 726 | num_prefills, |
| 727 | seq_lens, |
| 728 | seq_lens_cpu, |
| 729 | query_start_loc, |
| 730 | query_start_loc_cpu, |
| 731 | replay_start, |
| 732 | ) |
| 733 | |
| 734 | # Per-layer-type tile-scheduler plan holders. Empty FlashMLASchedMeta |
| 735 | # per present DeepseekV4 layer type; the first flash_mla_with_kvcache call of |
| 736 | # each type triggers the planner and all same-type layers reuse the |
| 737 | # resulting plan for the rest of the step. |
| 738 | tile_sched = self.build_tile_scheduler(num_decode_tokens) |
| 739 | |
| 740 | return DeepseekSparseSWAMetadata( |
| 741 | seq_lens=seq_lens, |
| 742 | query_start_loc=query_start_loc, |
| 743 | query_start_loc_cpu=query_start_loc_cpu, |
| 744 | block_table=block_table, |
| 745 | slot_mapping=slot_mapping, |
| 746 | is_valid_token=is_valid_token, |
| 747 | token_to_req_indices=token_to_req_indices, |
| 748 | decode_swa_indices=decode_swa_indices[:num_decode_tokens], |
| 749 | decode_swa_lens=self.decode_swa_lens[:num_decode_tokens], |
| 750 | decode_swa_width=decode_swa_width, |
| 751 | prefill_swa_indices=( |
| 752 | self.prefill_swa_indices[:num_prefill_tokens] |
| 753 | if num_prefill_tokens > 0 |
| 754 | else None |
| 755 | ), |
| 756 | prefill_swa_lens=( |
| 757 | self.prefill_swa_lens[:num_prefill_tokens] |
| 758 | if num_prefill_tokens > 0 |
| 759 | else None |
| 760 | ), |
| 761 | prefill_left_visible=prefill_left_visible, |
| 762 | prefill_right_visible=prefill_right_visible, |
| 763 | replay_start=replay_start, |
| 764 | block_size=self.block_size, |
| 765 | num_decodes=num_decodes, |
| 766 | num_prefills=num_prefills, |
| 767 | num_decode_tokens=num_decode_tokens, |
| 768 | num_prefill_tokens=num_prefill_tokens, |
| 769 | # Upper bound on decode-split rows for the kernel's max_q_len |
| 770 | # hint. common max_query_len bounds every row (scheduled max under |
| 771 | # adaptive verification), clamped to what the split can admit so a |
| 772 | # mixed batch's prefill max does not inflate decode scheduling. |
| 773 | max_decode_query_len=min( |
| 774 | common_attn_metadata.max_query_len, self.decode_threshold |
| 775 | ), |
| 776 | tile_sched_swaonly=tile_sched[_LAYER_TYPE_SWAONLY], |
| 777 | tile_sched_c4a=tile_sched[_LAYER_TYPE_C4A], |
| 778 | tile_sched_c128a=tile_sched[_LAYER_TYPE_C128A], |
| 779 | tile_sched_c1a=tile_sched[_LAYER_TYPE_C1A], |
| 780 | tile_sched_c2a=tile_sched[_LAYER_TYPE_C2A], |
| 781 | **deepseek_v4_fields, # type: ignore[arg-type] |
| 782 | ) |
| 783 | |
| 784 | def _build_image_visibility( |
| 785 | self, |
| 786 | num_reqs: int, |
| 787 | mm_ranges: dict[int, list[tuple[int, int]]], |
| 788 | num_decode_tokens: int, |
| 789 | num_prefill_tokens: int, |
| 790 | seq_lens: torch.Tensor, |
| 791 | query_start_loc: torch.Tensor, |
| 792 | token_to_req_indices: torch.Tensor, |
| 793 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 794 | """Upload image spans and compute per-token in-image visibility. |
| 795 | |
| 796 | Returns (left, right) int32 tensors indexed by absolute (decode-first) |
| 797 | token position; only prefill rows are meaningful. |
| 798 | """ |
| 799 | indptr = [0] * (num_reqs + 1) |
| 800 | starts: list[int] = [] |
| 801 | ends: list[int] = [] |
| 802 | for req_idx in range(num_reqs): |
| 803 | for span_start, span_end in mm_ranges.get(req_idx, ()): |
| 804 | starts.append(span_start) |
| 805 | ends.append(span_end) |
| 806 | indptr[req_idx + 1] = len(starts) |
| 807 | num_spans = len(starts) |
| 808 | assert num_spans <= self.span_starts.shape[0] |
| 809 | self.span_indptr[: num_reqs + 1].copy_(torch.tensor(indptr, dtype=torch.int32)) |
| 810 | if num_spans > 0: |
| 811 | self.span_starts[:num_spans].copy_(torch.tensor(starts, dtype=torch.int32)) |
| 812 | self.span_ends[:num_spans].copy_(torch.tensor(ends, dtype=torch.int32)) |
| 813 | _compute_image_visibility_kernel[(num_prefill_tokens,)]( |
| 814 | self.left_visible, |
| 815 | self.right_visible, |
| 816 | self.span_indptr, |
| 817 | self.span_starts, |
| 818 | self.span_ends, |
| 819 | query_start_loc, |
| 820 | seq_lens, |
| 821 | token_to_req_indices, |
| 822 | self.max_image_tokens, |
| 823 | token_offset=num_decode_tokens, |
| 824 | ) |
| 825 | end = num_decode_tokens + num_prefill_tokens |
| 826 | return self.left_visible[:end], self.right_visible[:end] |
| 827 | |
| 828 | def update_draft_decode_metadata( |
| 829 | self, |
| 830 | metadata: DeepseekSparseSWAMetadata, |
| 831 | ) -> None: |
| 832 | if metadata.num_decode_tokens == 0: |
| 833 | return |
| 834 | assert metadata.query_start_loc is not None |
| 835 | assert metadata.seq_lens is not None |
| 836 | assert metadata.token_to_req_indices is not None |
| 837 | assert metadata.is_valid_token is not None |
| 838 | assert metadata.decode_swa_indices is not None |
| 839 | assert metadata.decode_swa_lens is not None |
| 840 | assert metadata.replay_start is not None |
| 841 | |
| 842 | _COMPUTE_SWA_INDICES_AND_LENS_KERNEL( |
| 843 | metadata.decode_swa_indices, |
| 844 | metadata.decode_swa_lens, |
| 845 | metadata.decode_swa_indices.shape[-1], |
| 846 | metadata.decode_swa_indices.shape[-1], |
| 847 | metadata.decode_swa_lens, # unused (HAS_IMAGE=False) |
| 848 | metadata.decode_swa_lens, # unused (HAS_IMAGE=False) |
| 849 | metadata.query_start_loc, |
| 850 | metadata.seq_lens, |
| 851 | metadata.token_to_req_indices, |
| 852 | metadata.is_valid_token, |
| 853 | metadata.block_table, |
| 854 | self.block_size, |
| 855 | metadata.replay_start, |
| 856 | num_tokens=metadata.num_decode_tokens, |
| 857 | token_offset=0, |
| 858 | ) |
| 859 | tile_sched = self.build_tile_scheduler(metadata.num_decode_tokens) |
| 860 | metadata.tile_sched_swaonly = tile_sched[_LAYER_TYPE_SWAONLY] |
| 861 | metadata.tile_sched_c4a = tile_sched[_LAYER_TYPE_C4A] |
| 862 | metadata.tile_sched_c128a = tile_sched[_LAYER_TYPE_C128A] |
| 863 | metadata.tile_sched_c1a = tile_sched[_LAYER_TYPE_C1A] |
| 864 | metadata.tile_sched_c2a = tile_sched[_LAYER_TYPE_C2A] |
| 865 | metadata.flashinfer_sparse_index_cache.clear() |
| 866 | |
| 867 | def build_tile_scheduler( |
| 868 | self, num_decode_tokens: int |
| 869 | ) -> dict[str, FlashMLASchedMeta | None]: |
| 870 | """Allocate one empty ``FlashMLASchedMeta`` per present DeepseekV4 layer type. |
| 871 | |
| 872 | Returned instances have ``tile_scheduler_metadata`` / ``num_splits`` |
| 873 | set to ``None``; the FlashMLA C++ decode path will allocate them and |
| 874 | run the tile-scheduler planner on the first ``flash_mla_with_kvcache`` |
| 875 | call of each type. Subsequent same-type calls reuse the plan because |
| 876 | the tensors (and ``have_initialized``) are populated on the struct. |
| 877 | |
| 878 | Returns all-``None`` when there are no decode tokens this step, so |
| 879 | ``_forward_decode`` sees a clean sentinel. |
| 880 | """ |
| 881 | out: dict[str, FlashMLASchedMeta | None] = { |
| 882 | _LAYER_TYPE_SWAONLY: None, |
| 883 | _LAYER_TYPE_C4A: None, |
| 884 | _LAYER_TYPE_C128A: None, |
| 885 | _LAYER_TYPE_C1A: None, |
| 886 | _LAYER_TYPE_C2A: None, |
| 887 | } |
| 888 | if ( |
| 889 | num_decode_tokens == 0 |
| 890 | or current_platform.is_rocm() |
| 891 | or current_platform.is_xpu() |
| 892 | or current_platform.is_device_capability_family(120) |
| 893 | ): |
| 894 | return out |
| 895 | for layer_type in self._layer_types: |
| 896 | # get_mla_metadata() is the official FlashMLA entry point that |
| 897 | # returns a fresh empty FlashMLASchedMeta; using it keeps this |
| 898 | # call site aligned with the rest of the vLLM FlashMLA backends |
| 899 | # that already go through the same stub. |
| 900 | out[layer_type] = get_mla_metadata()[0] |
| 901 | return out |
| 902 | |
| 903 | def _build_deepseek_v4_metadata( |
| 904 | self, |
| 905 | num_decodes: int, |
| 906 | num_prefills: int, |
| 907 | seq_lens: torch.Tensor, |
| 908 | seq_lens_cpu: torch.Tensor | None, |
| 909 | query_start_loc: torch.Tensor, |
| 910 | query_start_loc_cpu: torch.Tensor, |
| 911 | replay_start: torch.Tensor, |
| 912 | ) -> dict[str, torch.Tensor | int | None]: |
| 913 | """Pre-compute DeepseekV4 prefill metadata during the metadata build phase. |
| 914 | |
| 915 | Returns a dict of keyword arguments to pass to the |
| 916 | DeepseekSparseSWAMetadata constructor. |
| 917 | |
| 918 | Note: C128A sparse metadata is computed by the FlashMLASparse builder |
| 919 | (which owns the C128A block_table), not here. |
| 920 | """ |
| 921 | result: dict[str, torch.Tensor | int | None] = {} |
| 922 | |
| 923 | # --- Prefill query metadata (single Triton kernel + CPU slicing) --- |
| 924 | if num_prefills > 0: |
| 925 | assert seq_lens_cpu is not None |
| 926 | pfx_gather_lens = torch.empty( |
| 927 | num_prefills, dtype=torch.int32, device=seq_lens.device |
| 928 | ) |
| 929 | _COMPUTE_PREFILL_METADATA_KERNEL( |
| 930 | pfx_gather_lens, |
| 931 | seq_lens, |
| 932 | query_start_loc, |
| 933 | replay_start, |
| 934 | num_prefills, |
| 935 | num_decodes, |
| 936 | self.window_size, |
| 937 | ) |
| 938 | |
| 939 | result["prefill_seq_lens"] = seq_lens[num_decodes:] |
| 940 | result["prefill_seq_lens_cpu"] = seq_lens_cpu[num_decodes:] |
| 941 | result["prefill_gather_lens"] = pfx_gather_lens |
| 942 | result["prefill_query_lens_cpu"] = ( |
| 943 | query_start_loc_cpu[num_decodes + 1 : num_decodes + num_prefills + 1] |
| 944 | - query_start_loc_cpu[num_decodes : num_decodes + num_prefills] |
| 945 | ).to(dtype=torch.int32) |
| 946 | result["prefill_window_size"] = self.window_size |
| 947 | result["prefill_max_model_len"] = self.max_model_len |
| 948 | result["prefill_max_num_batched_tokens"] = self.max_num_batched_tokens |
| 949 | |
| 950 | return result |
| 951 | |
| 952 | |
| 953 | @triton.jit(do_not_specialize=["token_offset"]) |
| 954 | def _compute_image_visibility_kernel( |
| 955 | left_visible_ptr, |
| 956 | right_visible_ptr, |
| 957 | span_indptr_ptr, |
| 958 | span_starts_ptr, |
| 959 | span_ends_ptr, |
| 960 | query_start_loc_ptr, |
| 961 | seq_lens_ptr, |
| 962 | token_to_req_indices_ptr, |
| 963 | max_image_tokens, |
| 964 | token_offset, |
| 965 | ): |
| 966 | """Per-token in-image visible counts (port of `get_image_visible`). |
| 967 | |
| 968 | One program per prefill token. A token at position pos inside a span |
| 969 | [span_start, span_end] sees min(pos - span_start, max_image_tokens - 1) |
| 970 | extra tokens to its left and min(span_end - pos, max_image_tokens) to its |
| 971 | right; tokens outside every span get 0/0 (plain causal window). |
| 972 | """ |
| 973 | pid = tl.program_id(0) |
| 974 | token_idx = pid + token_offset |
| 975 | req_idx = tl.load(token_to_req_indices_ptr + token_idx) |
| 976 | |
| 977 | query_start = tl.load(query_start_loc_ptr + req_idx) |
| 978 | query_end = tl.load(query_start_loc_ptr + req_idx + 1) |
| 979 | seq_len = tl.load(seq_lens_ptr + req_idx) |
| 980 | pos = seq_len - (query_end - query_start) + token_idx - query_start |
| 981 | |
| 982 | span_lo = tl.load(span_indptr_ptr + req_idx) |
| 983 | span_hi = tl.load(span_indptr_ptr + req_idx + 1) |
| 984 | left = tl.zeros((), dtype=tl.int32) |
| 985 | right = tl.zeros((), dtype=tl.int32) |
| 986 | for i in range(span_lo, span_hi): |
| 987 | span_start = tl.load(span_starts_ptr + i) |
| 988 | span_end = tl.load(span_ends_ptr + i) |
| 989 | in_span = (pos >= span_start) & (pos <= span_end) |
| 990 | left = tl.where( |
| 991 | in_span, tl.minimum(pos - span_start, max_image_tokens - 1), left |
| 992 | ) |
| 993 | right = tl.where(in_span, tl.minimum(span_end - pos, max_image_tokens), right) |
| 994 | tl.store(left_visible_ptr + token_idx, left) |
| 995 | tl.store(right_visible_ptr + token_idx, right) |
| 996 | |
| 997 | |
| 998 | # TODO(ben): unify this kernel to reduce duplication |
| 999 | |
| 1000 | |
| 1001 | @triton.jit( |
| 1002 | do_not_specialize=[ |
| 1003 | "swa_indices_stride", |
| 1004 | "block_table_stride", |
| 1005 | "token_offset", |
| 1006 | ] |
| 1007 | ) |
| 1008 | def _compute_swa_indices_and_lens_kernel( |
| 1009 | swa_indices_ptr, |
| 1010 | swa_indices_stride, |
| 1011 | swa_lens_ptr, |
| 1012 | window_size, |
| 1013 | index_width, |
| 1014 | left_visible_ptr, |
| 1015 | right_visible_ptr, |
| 1016 | query_start_loc_ptr, |
| 1017 | seq_lens_ptr, |
| 1018 | token_to_req_indices_ptr, |
| 1019 | is_valid_token_ptr, |
| 1020 | block_table_ptr, |
| 1021 | block_table_stride, |
| 1022 | block_size, |
| 1023 | replay_start_ptr, |
| 1024 | token_offset, |
| 1025 | HAS_IMAGE: tl.constexpr, |
| 1026 | TRITON_BLOCK_SIZE: tl.constexpr, |
| 1027 | ): |
| 1028 | pid = tl.program_id(0) |
| 1029 | token_idx = pid + token_offset |
| 1030 | is_valid = tl.load(is_valid_token_ptr + token_idx) |
| 1031 | if not is_valid: |
| 1032 | tl.store(swa_lens_ptr + pid, 0) |
| 1033 | # Clear the row so a padded token cannot gather through stale indices. |
| 1034 | for i in range(0, index_width, TRITON_BLOCK_SIZE): |
| 1035 | offset = i + tl.arange(0, TRITON_BLOCK_SIZE) |
| 1036 | tl.store( |
| 1037 | swa_indices_ptr + pid * swa_indices_stride + offset, |
| 1038 | -1, |
| 1039 | mask=offset < index_width, |
| 1040 | ) |
| 1041 | return |
| 1042 | |
| 1043 | req_idx = tl.load(token_to_req_indices_ptr + token_idx) |
| 1044 | |
| 1045 | query_start = tl.load(query_start_loc_ptr + req_idx) |
| 1046 | query_end = tl.load(query_start_loc_ptr + req_idx + 1) |
| 1047 | query_len = query_end - query_start |
| 1048 | |
| 1049 | seq_len = tl.load(seq_lens_ptr + req_idx) |
| 1050 | prefix_len = seq_len - query_len |
| 1051 | |
| 1052 | pos = prefix_len + token_idx - query_start |
| 1053 | if HAS_IMAGE: |
| 1054 | # In-image bidirectional visibility widens the window: the window |
| 1055 | # starts up to max(left - (window - 1), 0) positions earlier and |
| 1056 | # extends `right` positions past the query token. |
| 1057 | left = tl.load(left_visible_ptr + token_idx) |
| 1058 | right = tl.load(right_visible_ptr + token_idx) |
| 1059 | else: |
| 1060 | left = 0 |
| 1061 | right = 0 |
| 1062 | left_add = tl.maximum(left - (window_size - 1), 0) |
| 1063 | start_pos = tl.maximum(pos - (window_size - 1) - left_add, 0) |
| 1064 | # SWA bounded replay: no window KV exists below the request's replay start. |
| 1065 | start_pos = tl.maximum(start_pos, tl.load(replay_start_ptr + req_idx)) |
| 1066 | end_pos = pos + right + 1 |
| 1067 | |
| 1068 | swa_len = end_pos - start_pos |
| 1069 | tl.store(swa_lens_ptr + pid, swa_len) |
| 1070 | |
| 1071 | for i in range(0, index_width, TRITON_BLOCK_SIZE): |
| 1072 | offset = i + tl.arange(0, TRITON_BLOCK_SIZE) |
| 1073 | |
| 1074 | pos_offset = start_pos + offset |
| 1075 | block_indices = pos_offset // block_size |
| 1076 | block_numbers = tl.load( |
| 1077 | block_table_ptr + req_idx * block_table_stride + block_indices, |
| 1078 | mask=pos_offset < end_pos, |
| 1079 | ) |
| 1080 | block_offsets = pos_offset % block_size |
| 1081 | slot_ids = block_numbers * block_size + block_offsets |
| 1082 | |
| 1083 | slot_ids = tl.where(offset < swa_len, slot_ids, -1) |
| 1084 | tl.store( |
| 1085 | swa_indices_ptr + pid * swa_indices_stride + offset, |
| 1086 | slot_ids, |
| 1087 | mask=offset < index_width, |
| 1088 | ) |
| 1089 | |
| 1090 | |
| 1091 | class ComputeSWAIndicesAndLensKernel( |
| 1092 | VllmTritonJitKernel["ComputeSWAIndicesAndLensKernel.CompileKey"] |
| 1093 | ): |
| 1094 | @dataclass(frozen=True) |
| 1095 | class CompileKey: |
| 1096 | window_size: int |
| 1097 | index_width: int |
| 1098 | has_image: int |
| 1099 | block_size: int |
| 1100 | triton_block_size: int |
| 1101 | |
| 1102 | kernel = staticmethod(_compute_swa_indices_and_lens_kernel) |
| 1103 | |
| 1104 | def dispatch( # type: ignore[override] |
| 1105 | self, |
| 1106 | **compile_key_fields: int, |
| 1107 | ) -> CompileKey: |
| 1108 | return self.CompileKey( |
| 1109 | **compile_key_fields, |
| 1110 | triton_block_size=1024, |
| 1111 | ) |
| 1112 | |
| 1113 | def get_warmup_keys( |
| 1114 | self, |
| 1115 | *, |
| 1116 | window_size: int, |
| 1117 | block_size: int, |
| 1118 | max_image_tokens: int = 0, |
| 1119 | ) -> list[CompileKey]: |
| 1120 | # Decode rows always use the plain causal window. Vision models widen |
| 1121 | # prefill rows by max_image_tokens and additionally need the in-image |
| 1122 | # bidirectional variant for batches that contain image spans. |
| 1123 | width_has_image = [(window_size, 0)] |
| 1124 | if max_image_tokens > 0: |
| 1125 | width_has_image += [ |
| 1126 | (window_size + max_image_tokens, 0), |
| 1127 | (window_size + max_image_tokens, 1), |
| 1128 | ] |
| 1129 | return self._trace_dispatch(self.dispatch)( |
| 1130 | zip_inputs( |
| 1131 | *( |
| 1132 | dict(index_width=index_width, has_image=has_image) |
| 1133 | for index_width, has_image in width_has_image |
| 1134 | ) |
| 1135 | ), |
| 1136 | window_size=window_size, |
| 1137 | block_size=block_size, |
| 1138 | ) |
| 1139 | |
| 1140 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 1141 | int32_ptr = TritonWarmupTensor(torch.int32) |
| 1142 | return dict( |
| 1143 | swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), |
| 1144 | swa_lens=int32_ptr, |
| 1145 | window_size=compile_key.window_size, |
| 1146 | index_width=compile_key.index_width, |
| 1147 | left_visible=int32_ptr, |
| 1148 | right_visible=int32_ptr, |
| 1149 | query_start_loc=int32_ptr, |
| 1150 | seq_lens=int32_ptr, |
| 1151 | token_to_req_indices=int32_ptr, |
| 1152 | is_valid_token=TritonWarmupTensor(torch.bool), |
| 1153 | block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), |
| 1154 | block_size=compile_key.block_size, |
| 1155 | replay_start=int32_ptr, |
| 1156 | num_tokens=1, |
| 1157 | token_offset=0, |
| 1158 | has_image=bool(compile_key.has_image), |
| 1159 | ) |
| 1160 | |
| 1161 | @kernel_launcher |
| 1162 | def __call__( |
| 1163 | self, |
| 1164 | swa_indices: torch.Tensor, |
| 1165 | swa_lens: torch.Tensor, |
| 1166 | window_size: int, |
| 1167 | index_width: int, |
| 1168 | left_visible: torch.Tensor, |
| 1169 | right_visible: torch.Tensor, |
| 1170 | query_start_loc: torch.Tensor, |
| 1171 | seq_lens: torch.Tensor, |
| 1172 | token_to_req_indices: torch.Tensor, |
| 1173 | is_valid_token: torch.Tensor, |
| 1174 | block_table: torch.Tensor, |
| 1175 | block_size: int, |
| 1176 | replay_start: torch.Tensor, |
| 1177 | *, |
| 1178 | num_tokens: int, |
| 1179 | token_offset: int, |
| 1180 | has_image: bool = False, |
| 1181 | ) -> LaunchSpec: |
| 1182 | return (num_tokens,), dict( |
| 1183 | swa_indices_stride=swa_indices.stride(0), |
| 1184 | block_table_stride=block_table.stride(0), |
| 1185 | HAS_IMAGE=has_image, |
| 1186 | TRITON_BLOCK_SIZE=1024, |
| 1187 | ) |
| 1188 | |
| 1189 | |
| 1190 | class ComputeDSparkNoncausalSWAIndicesKernel( |
| 1191 | VllmTritonJitKernel["ComputeDSparkNoncausalSWAIndicesKernel.CompileKey"] |
| 1192 | ): |
| 1193 | @dataclass(frozen=True) |
| 1194 | class CompileKey: |
| 1195 | window_size: int |
| 1196 | index_width: int |
| 1197 | block_size: int |
| 1198 | triton_block_size: int |
| 1199 | |
| 1200 | @staticmethod |
| 1201 | @triton.jit( |
| 1202 | do_not_specialize=[ |
| 1203 | "swa_indices_stride", |
| 1204 | "block_table_stride", |
| 1205 | "token_offset", |
| 1206 | ] |
| 1207 | ) |
| 1208 | def kernel( |
| 1209 | swa_indices_ptr, |
| 1210 | swa_indices_stride, |
| 1211 | swa_lens_ptr, |
| 1212 | window_size, |
| 1213 | index_width, |
| 1214 | query_start_loc_ptr, |
| 1215 | seq_lens_ptr, |
| 1216 | token_to_req_indices_ptr, |
| 1217 | is_valid_token_ptr, |
| 1218 | block_table_ptr, |
| 1219 | block_table_stride, |
| 1220 | block_size, |
| 1221 | token_offset, |
| 1222 | TRITON_BLOCK_SIZE: tl.constexpr, |
| 1223 | ): |
| 1224 | """Non-causal per-token indices for the DSpark draft block. |
| 1225 | |
| 1226 | Here, we populate the topk indices with the trailing window of context tokens, |
| 1227 | plus all query tokens (including future ones). |
| 1228 | """ |
| 1229 | pid = tl.program_id(0) |
| 1230 | token_idx = pid + token_offset |
| 1231 | is_valid = tl.load(is_valid_token_ptr + token_idx) |
| 1232 | if not is_valid: |
| 1233 | tl.store(swa_lens_ptr + pid, 0) |
| 1234 | # Clear the row so a padded token cannot gather through stale indices. |
| 1235 | for i in range(0, index_width, TRITON_BLOCK_SIZE): |
| 1236 | offset = i + tl.arange(0, TRITON_BLOCK_SIZE) |
| 1237 | tl.store( |
| 1238 | swa_indices_ptr + pid * swa_indices_stride + offset, |
| 1239 | -1, |
| 1240 | mask=offset < index_width, |
| 1241 | ) |
| 1242 | return |
| 1243 | |
| 1244 | req_idx = tl.load(token_to_req_indices_ptr + token_idx) |
| 1245 | |
| 1246 | query_start = tl.load(query_start_loc_ptr + req_idx) |
| 1247 | query_end = tl.load(query_start_loc_ptr + req_idx + 1) |
| 1248 | query_len = query_end - query_start |
| 1249 | |
| 1250 | seq_len = tl.load(seq_lens_ptr + req_idx) |
| 1251 | prefix_len = seq_len - query_len |
| 1252 | |
| 1253 | # Block-anchored window (shared by every token in the block) + full block. |
| 1254 | start_pos = tl.maximum(prefix_len - window_size, 0) |
| 1255 | end_pos = seq_len |
| 1256 | |
| 1257 | swa_len = end_pos - start_pos |
| 1258 | tl.store(swa_lens_ptr + pid, swa_len) |
| 1259 | |
| 1260 | for i in range(0, index_width, TRITON_BLOCK_SIZE): |
| 1261 | offset = i + tl.arange(0, TRITON_BLOCK_SIZE) |
| 1262 | |
| 1263 | pos_offset = start_pos + offset |
| 1264 | block_indices = pos_offset // block_size |
| 1265 | block_numbers = tl.load( |
| 1266 | block_table_ptr + req_idx * block_table_stride + block_indices, |
| 1267 | mask=pos_offset < end_pos, |
| 1268 | ) |
| 1269 | block_offsets = pos_offset % block_size |
| 1270 | slot_ids = block_numbers * block_size + block_offsets |
| 1271 | |
| 1272 | slot_ids = tl.where(offset < swa_len, slot_ids, -1) |
| 1273 | tl.store( |
| 1274 | swa_indices_ptr + pid * swa_indices_stride + offset, |
| 1275 | slot_ids, |
| 1276 | mask=offset < index_width, |
| 1277 | ) |
| 1278 | |
| 1279 | def dispatch( # type: ignore[override] |
| 1280 | self, |
| 1281 | *, |
| 1282 | window_size: int, |
| 1283 | num_speculative_tokens: int, |
| 1284 | block_size: int, |
| 1285 | ) -> CompileKey: |
| 1286 | return self.CompileKey( |
| 1287 | window_size=window_size, |
| 1288 | index_width=get_dspark_swa_index_width(window_size, num_speculative_tokens), |
| 1289 | block_size=block_size, |
| 1290 | triton_block_size=1024, |
| 1291 | ) |
| 1292 | |
| 1293 | def get_warmup_keys( |
| 1294 | self, |
| 1295 | *, |
| 1296 | window_size: int, |
| 1297 | num_speculative_tokens: int, |
| 1298 | block_size: int, |
| 1299 | ) -> list[CompileKey]: |
| 1300 | return self._trace_dispatch(self.dispatch)( |
| 1301 | window_size=window_size, |
| 1302 | num_speculative_tokens=num_speculative_tokens, |
| 1303 | block_size=block_size, |
| 1304 | ) |
| 1305 | |
| 1306 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 1307 | int32_ptr = TritonWarmupTensor(torch.int32) |
| 1308 | return dict( |
| 1309 | swa_indices=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), |
| 1310 | swa_lens=int32_ptr, |
| 1311 | window_size=compile_key.window_size, |
| 1312 | index_width=compile_key.index_width, |
| 1313 | query_start_loc=int32_ptr, |
| 1314 | seq_lens=int32_ptr, |
| 1315 | token_to_req_indices=int32_ptr, |
| 1316 | is_valid_token=TritonWarmupTensor(torch.bool), |
| 1317 | block_table=TritonWarmupTensor(torch.int32, shape=(1, 1), strides=(1, 1)), |
| 1318 | block_size=compile_key.block_size, |
| 1319 | num_tokens=1, |
| 1320 | token_offset=0, |
| 1321 | ) |
| 1322 | |
| 1323 | @kernel_launcher |
| 1324 | def __call__( |
| 1325 | self, |
| 1326 | swa_indices: torch.Tensor, |
| 1327 | swa_lens: torch.Tensor, |
| 1328 | window_size: int, |
| 1329 | index_width: int, |
| 1330 | query_start_loc: torch.Tensor, |
| 1331 | seq_lens: torch.Tensor, |
| 1332 | token_to_req_indices: torch.Tensor, |
| 1333 | is_valid_token: torch.Tensor, |
| 1334 | block_table: torch.Tensor, |
| 1335 | block_size: int, |
| 1336 | *, |
| 1337 | num_tokens: int, |
| 1338 | token_offset: int, |
| 1339 | ) -> LaunchSpec: |
| 1340 | return (num_tokens,), dict( |
| 1341 | swa_indices_stride=swa_indices.stride(0), |
| 1342 | block_table_stride=block_table.stride(0), |
| 1343 | TRITON_BLOCK_SIZE=1024, |
| 1344 | ) |
| 1345 | |
| 1346 | |
| 1347 | _COMPUTE_PREFILL_METADATA_KERNEL = ComputePrefillMetadataKernel() |
| 1348 | _COMPUTE_SWA_INDICES_AND_LENS_KERNEL = ComputeSWAIndicesAndLensKernel() |
| 1349 | _COMPUTE_DSPARK_NONCAUSAL_SWA_INDICES_KERNEL = ComputeDSparkNoncausalSWAIndicesKernel() |