| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass |
| 4 | from typing import Any |
| 5 | |
| 6 | import numpy as np |
| 7 | import torch |
| 8 | |
| 9 | import vllm.envs as envs |
| 10 | from vllm.config import VllmConfig |
| 11 | from vllm.distributed import get_dcp_group, get_pcp_group |
| 12 | from vllm.logger import init_logger |
| 13 | from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs |
| 14 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 15 | LaunchSpec, |
| 16 | TritonPointerInputVariant, |
| 17 | TritonWarmupTensor, |
| 18 | VllmTritonJitKernel, |
| 19 | triton_scalar_specialization_rep, |
| 20 | ) |
| 21 | from vllm.platforms import current_platform |
| 22 | from vllm.triton_utils import tl, triton |
| 23 | from vllm.utils.deep_gemm import ( |
| 24 | get_paged_mqa_logits_metadata, |
| 25 | has_deep_gemm, |
| 26 | native_next_n_supported, |
| 27 | ) |
| 28 | from vllm.utils.math_utils import round_down |
| 29 | from vllm.utils.platform_utils import num_compute_units |
| 30 | from vllm.utils.torch_utils import PIN_MEMORY, async_tensor_h2d |
| 31 | from vllm.v1.attention.backend import ( |
| 32 | AttentionBackend, |
| 33 | AttentionCGSupport, |
| 34 | AttentionMetadataBuilder, |
| 35 | CommonAttentionMetadata, |
| 36 | MultipleOf, |
| 37 | ) |
| 38 | from vllm.v1.attention.backends.mla.compressor_utils import get_compressed_slot_mapping |
| 39 | from vllm.v1.attention.backends.mla.sparse_utils import request_row_bounds |
| 40 | from vllm.v1.attention.backends.utils import ( |
| 41 | get_dcp_local_seq_lens, |
| 42 | split_decodes_and_prefills, |
| 43 | ) |
| 44 | from vllm.v1.kv_cache_interface import ( |
| 45 | AttentionSpec, |
| 46 | KVCacheLayout, |
| 47 | KVCacheSpec, |
| 48 | MLAAttentionSpec, |
| 49 | ) |
| 50 | |
| 51 | logger = init_logger(__name__) |
| 52 | |
| 53 | # The DSA indexer K cache is always quantized; "auto" means fp8 (V3.2 layout) |
| 54 | # and mxfp4 is the opt-in Blackwell and gfx950 path. |
| 55 | DSA_INDEXER_KV_DTYPES = ("fp8", "mxfp4") |
| 56 | |
| 57 | |
| 58 | def dsa_indexer_uses_fp4(vllm_config: VllmConfig) -> bool: |
| 59 | """Whether the DeepSeek sparse indexer should use the MXFP4 K cache.""" |
| 60 | kv_dtype = vllm_config.attention_config.resolve_indexer_kv_dtype("fp8") |
| 61 | if kv_dtype not in DSA_INDEXER_KV_DTYPES: |
| 62 | raise ValueError( |
| 63 | f"indexer_kv_dtype={kv_dtype!r} is not supported by the DeepSeek " |
| 64 | f"sparse indexer (expected one of {DSA_INDEXER_KV_DTYPES})." |
| 65 | ) |
| 66 | use_fp4 = kv_dtype == "mxfp4" |
| 67 | if use_fp4 and current_platform.is_rocm(): |
| 68 | from vllm._aiter_ops import rocm_aiter_ops |
| 69 | from vllm.platforms.rocm import get_cdna_version |
| 70 | |
| 71 | # Only DeepSeek-V4.1 is wired to the ROCm MXFP4 cache; other DSA models |
| 72 | # would silently keep their FP8 one. |
| 73 | model_config = vllm_config.model_config |
| 74 | if model_config is None or model_config.hf_config.model_type != "deepseek_v41": |
| 75 | raise ValueError( |
| 76 | "indexer_kv_dtype='mxfp4' on ROCm is only supported for " |
| 77 | "DeepSeek-V4.1-Flash." |
| 78 | ) |
| 79 | if get_cdna_version() != 4: |
| 80 | raise ValueError( |
| 81 | "indexer_kv_dtype='mxfp4' on ROCm requires CDNA4 (MI350X/MI355X)." |
| 82 | ) |
| 83 | if not rocm_aiter_ops.is_enabled(): |
| 84 | raise ValueError( |
| 85 | "indexer_kv_dtype='mxfp4' on ROCm runs on aiter's kernels; enable " |
| 86 | "aiter with VLLM_ROCM_USE_AITER=1." |
| 87 | ) |
| 88 | return True |
| 89 | if use_fp4 and not current_platform.is_device_capability_family(100): |
| 90 | raise ValueError( |
| 91 | "indexer_kv_dtype='mxfp4' requires Blackwell datacenter GPUs " |
| 92 | "(sm_10x, e.g. B200/GB200); sm_120 (consumer Blackwell) and " |
| 93 | "earlier architectures are not supported." |
| 94 | ) |
| 95 | return use_fp4 |
| 96 | |
| 97 | |
| 98 | class PrepareUniformDecodeKernel( |
| 99 | VllmTritonJitKernel["PrepareUniformDecodeKernel.CompileKey"] |
| 100 | ): |
| 101 | BLOCK_SIZE = 1024 |
| 102 | |
| 103 | @dataclass(frozen=True) |
| 104 | class CompileKey: |
| 105 | block_size: int |
| 106 | block_table_stride: int |
| 107 | expanded_bt_stride: int |
| 108 | max_decode_len: int |
| 109 | |
| 110 | @staticmethod |
| 111 | @triton.jit |
| 112 | def kernel( |
| 113 | seq_lens_ptr, |
| 114 | decode_seq_lens_ptr, |
| 115 | block_table_ptr, |
| 116 | block_table_stride, |
| 117 | expanded_block_table_ptr, |
| 118 | expanded_bt_stride, |
| 119 | decode_lens_ptr, |
| 120 | max_decode_len, |
| 121 | BLOCK_SIZE: tl.constexpr, |
| 122 | ): |
| 123 | idx = tl.program_id(0) |
| 124 | req_id = idx // max_decode_len |
| 125 | local_idx = idx % max_decode_len |
| 126 | |
| 127 | # Compute number of KVs attended to by this token. Padding requests have |
| 128 | # seq_len == 0, which would otherwise make the first token of each padded |
| 129 | # request negative (e.g. next_n=2 gives 0-2+0+1 = -1). Downstream kernels |
| 130 | # read these as uint32, turning -1 into ~4e9. |
| 131 | seq_len = tl.load(seq_lens_ptr + req_id) |
| 132 | per_token_seq_len = tl.maximum(seq_len - max_decode_len + local_idx + 1, 0) |
| 133 | tl.store(decode_seq_lens_ptr + idx, per_token_seq_len) |
| 134 | |
| 135 | # Copy block table row. |
| 136 | src = block_table_ptr + req_id * block_table_stride |
| 137 | dst = expanded_block_table_ptr + idx * expanded_bt_stride |
| 138 | for i in tl.range(0, expanded_bt_stride, BLOCK_SIZE): |
| 139 | off = i + tl.arange(0, BLOCK_SIZE) |
| 140 | mask = off < expanded_bt_stride |
| 141 | src_block = tl.load(src + off, mask=mask) |
| 142 | tl.store(dst + off, src_block, mask=mask) |
| 143 | |
| 144 | # All reqs now have decode_len = 1. |
| 145 | tl.store(decode_lens_ptr + idx, 1) |
| 146 | |
| 147 | def dispatch( # type: ignore[override] |
| 148 | self, |
| 149 | *, |
| 150 | block_size: int, |
| 151 | block_table_stride: int, |
| 152 | expanded_bt_stride: int, |
| 153 | max_decode_len: int, |
| 154 | ) -> CompileKey: |
| 155 | return self.CompileKey( |
| 156 | block_size=block_size, |
| 157 | block_table_stride=triton_scalar_specialization_rep(block_table_stride), |
| 158 | expanded_bt_stride=triton_scalar_specialization_rep(expanded_bt_stride), |
| 159 | max_decode_len=triton_scalar_specialization_rep(max_decode_len), |
| 160 | ) |
| 161 | |
| 162 | def get_warmup_keys(self) -> list[CompileKey]: |
| 163 | return self._trace_dispatch(self.dispatch)( |
| 164 | block_size=self.BLOCK_SIZE, |
| 165 | block_table_stride=(1, 2, 16), |
| 166 | expanded_bt_stride=(1, 2, 16), |
| 167 | max_decode_len=(1, 2, 16), |
| 168 | ) |
| 169 | |
| 170 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 171 | int32_ptr = TritonWarmupTensor(torch.int32) |
| 172 | return dict( |
| 173 | seq_lens=int32_ptr, |
| 174 | decode_seq_lens=int32_ptr, |
| 175 | block_table=TritonWarmupTensor( |
| 176 | torch.int32, |
| 177 | shape=(1, compile_key.block_table_stride), |
| 178 | ), |
| 179 | expanded_block_table=TritonWarmupTensor( |
| 180 | torch.int32, |
| 181 | shape=(1, compile_key.expanded_bt_stride), |
| 182 | ), |
| 183 | decode_lens=int32_ptr, |
| 184 | num_decode_tokens=1, |
| 185 | max_decode_len=compile_key.max_decode_len, |
| 186 | ) |
| 187 | |
| 188 | @kernel_launcher |
| 189 | def __call__( |
| 190 | self, |
| 191 | seq_lens: torch.Tensor, |
| 192 | decode_seq_lens: torch.Tensor, |
| 193 | block_table: torch.Tensor, |
| 194 | expanded_block_table: torch.Tensor, |
| 195 | decode_lens: torch.Tensor, |
| 196 | num_decode_tokens: int, |
| 197 | max_decode_len: int, |
| 198 | ) -> LaunchSpec: |
| 199 | return (num_decode_tokens,), dict( |
| 200 | block_table_stride=block_table.stride(0), |
| 201 | expanded_bt_stride=expanded_block_table.stride(0), |
| 202 | BLOCK_SIZE=self.BLOCK_SIZE, |
| 203 | ) |
| 204 | |
| 205 | |
| 206 | class DeepseekV32IndexerBackend(AttentionBackend): |
| 207 | @classmethod |
| 208 | def supports_pcp(cls) -> bool: |
| 209 | return True |
| 210 | |
| 211 | @classmethod |
| 212 | def supports_device_cpu_query_lens_mismatch(cls) -> bool: |
| 213 | # Only the varlen paged MQA logits kernel takes per-request query |
| 214 | # lengths from device tensors natively. Hopper can instead flatten each |
| 215 | # query into a single-token row using device-built metadata. |
| 216 | return _supports_varlen_paged_mqa_logits() or ( |
| 217 | _supports_flattened_device_query_lens() |
| 218 | ) |
| 219 | |
| 220 | @staticmethod |
| 221 | def get_name() -> str: |
| 222 | return "DEEPSEEK_V32_INDEXER" |
| 223 | |
| 224 | @staticmethod |
| 225 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 226 | return [1, MultipleOf(16)] if current_platform.is_rocm() else [64] |
| 227 | |
| 228 | @classmethod |
| 229 | def get_supported_head_sizes(cls) -> list[int]: |
| 230 | return [32, 64, 128] |
| 231 | |
| 232 | @staticmethod |
| 233 | def get_builder_cls() -> type["DeepseekV32IndexerMetadataBuilder"]: |
| 234 | return DeepseekV32IndexerMetadataBuilder |
| 235 | |
| 236 | |
| 237 | class KpoolTailBackend(DeepseekV32IndexerBackend): |
| 238 | """Storage-only backend for the GLM-5.3-Flash kpool tail cache.""" |
| 239 | |
| 240 | @classmethod |
| 241 | def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: |
| 242 | return (KVCacheLayout.LBHNC,) |
| 243 | |
| 244 | @staticmethod |
| 245 | def get_name() -> str: |
| 246 | return "KPOOL_TAIL" |
| 247 | |
| 248 | @classmethod |
| 249 | def get_supported_head_sizes(cls) -> list[int]: |
| 250 | return [] |
| 251 | |
| 252 | @staticmethod |
| 253 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 254 | return [MultipleOf(1)] |
| 255 | |
| 256 | @staticmethod |
| 257 | def get_builder_cls() -> type["KpoolTailMetadataBuilder"]: # type: ignore[override] |
| 258 | return KpoolTailMetadataBuilder |
| 259 | |
| 260 | |
| 261 | class DeepseekV4IndexerBackend(DeepseekV32IndexerBackend): |
| 262 | @classmethod |
| 263 | def supports_device_cpu_query_lens_mismatch(cls) -> bool: |
| 264 | # ROCm runs adaptive verification through the per-token flattened |
| 265 | # indexer path, which derives row ownership from device decode lengths. |
| 266 | return ( |
| 267 | current_platform.is_rocm() |
| 268 | or super().supports_device_cpu_query_lens_mismatch() |
| 269 | ) |
| 270 | |
| 271 | @staticmethod |
| 272 | def get_name() -> str: |
| 273 | return "DEEPSEEK_V4_INDEXER" |
| 274 | |
| 275 | @classmethod |
| 276 | def supported_kv_cache_layouts(cls) -> tuple[KVCacheLayout, ...]: |
| 277 | # DeepSeek-V4 packs the indexer pages beside the MLA latent pages inside |
| 278 | # each block, so the layer dim must sit inside the block dim. |
| 279 | return (KVCacheLayout.BLHNC, KVCacheLayout.BLNHC) |
| 280 | |
| 281 | @staticmethod |
| 282 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 283 | # Block sizes count uncompressed tokens: C4 indexer pages hold 64 rows. |
| 284 | return [256] |
| 285 | |
| 286 | |
| 287 | class DeepseekV41IndexerBackend(DeepseekV4IndexerBackend): |
| 288 | @classmethod |
| 289 | def supports_device_cpu_query_lens_mismatch(cls) -> bool: |
| 290 | # The ROCm flattened-query support above is validated for the |
| 291 | # DeepSeek-V4 adaptive DSpark path only. |
| 292 | return DeepseekV32IndexerBackend.supports_device_cpu_query_lens_mismatch() |
| 293 | |
| 294 | @staticmethod |
| 295 | def get_name() -> str: |
| 296 | return "DEEPSEEK_V41_INDEXER" |
| 297 | |
| 298 | @staticmethod |
| 299 | def get_supported_kernel_block_sizes(kv_cache_spec=None) -> list[int | MultipleOf]: |
| 300 | return [64 if current_platform.is_device_capability_family(90) else 128] |
| 301 | |
| 302 | |
| 303 | @dataclass(frozen=True) |
| 304 | class PCPGlobalChunkPlan: |
| 305 | """PCP packing for one indexer prefill chunk under PCP + DCP.""" |
| 306 | |
| 307 | # [num_reqs+1] cumsum of the padded per-request context (shard rows x W): |
| 308 | # the global layout. |
| 309 | row_start_cu: torch.Tensor |
| 310 | # [num_reqs+1] like row_start_cu, but only a region's FIRST row carries its |
| 311 | # extent. |
| 312 | global_cu: torch.Tensor |
| 313 | # [num_reqs+1] cumsum of shard rows: this rank's padded layout. |
| 314 | padded_local_cu: torch.Tensor |
| 315 | padded_local_total: int |
| 316 | total: int |
| 317 | # [total] for each global row, its position in the rank-major gathered buffer. |
| 318 | deinterleave_idx: torch.Tensor |
| 319 | |
| 320 | |
| 321 | def build_pcp_global_chunk_plan( |
| 322 | row_req_idx: np.ndarray, |
| 323 | row_shard_rows: np.ndarray, |
| 324 | dcp_world_size: int, |
| 325 | device: torch.device, |
| 326 | interleave: int = 1, |
| 327 | ) -> PCPGlobalChunkPlan: |
| 328 | """Plan the PCP packing for one chunk from its requests' KV shard rows. |
| 329 | |
| 330 | The global layout is padded to ``shard rows x W`` per request. Positions |
| 331 | past a request's true extent lie beyond every token's causal bound, so |
| 332 | the kernel never reads them. |
| 333 | """ |
| 334 | shard_rows = np.ascontiguousarray(row_shard_rows, dtype=np.int64) |
| 335 | assert row_req_idx.shape == shard_rows.shape |
| 336 | num_rows = len(shard_rows) |
| 337 | |
| 338 | row_bounds = request_row_bounds(row_req_idx) |
| 339 | region_first_row = row_bounds[:-1] |
| 340 | region_of_row = np.repeat(np.arange(len(region_first_row)), np.diff(row_bounds)) |
| 341 | |
| 342 | region_padded = shard_rows[region_first_row] |
| 343 | assert np.all(region_padded > 0), ( |
| 344 | f"PCP+DCP prefill got an empty context: {region_padded.tolist()}" |
| 345 | ) |
| 346 | |
| 347 | # Pinned host buffers, filled in place so the copies below are async. |
| 348 | cu = torch.zeros((3, num_rows + 1), dtype=torch.int32, pin_memory=PIN_MEMORY) |
| 349 | row_start_cu, global_cu, padded_cu = cu |
| 350 | # Only a region's first row carries its extent. |
| 351 | first_row = torch.from_numpy(region_first_row) |
| 352 | row_padded = torch.zeros(num_rows, dtype=torch.int32) |
| 353 | row_padded[first_row] = torch.from_numpy(region_padded).int() |
| 354 | torch.cumsum(row_padded * dcp_world_size, 0, out=global_cu[1:]) |
| 355 | torch.cumsum(row_padded, 0, out=padded_cu[1:]) |
| 356 | row_start_cu[:num_rows] = global_cu[first_row][torch.from_numpy(region_of_row)] |
| 357 | row_start_cu[num_rows] = global_cu[num_rows] |
| 358 | total = int(global_cu[num_rows]) |
| 359 | padded_total = int(padded_cu[num_rows]) |
| 360 | |
| 361 | idx = torch.empty(total, dtype=torch.int64, pin_memory=PIN_MEMORY) |
| 362 | idx_np = idx.numpy() |
| 363 | region_start = global_cu[first_row].numpy() |
| 364 | region_padded_start = padded_cu[first_row].numpy() |
| 365 | for i in range(len(region_first_row)): |
| 366 | g = int(region_padded[i]) * dcp_world_size |
| 367 | t = np.arange(g, dtype=np.int64) |
| 368 | local = round_down(t // dcp_world_size, interleave) + t % interleave |
| 369 | # Padded positions are outside causal bounds but must remain in-bounds. |
| 370 | local = np.minimum(local, region_padded[i] - 1) |
| 371 | idx_np[region_start[i] : region_start[i] + g] = ( |
| 372 | ((t // interleave) % dcp_world_size) * padded_total |
| 373 | + region_padded_start[i] |
| 374 | + local |
| 375 | ) |
| 376 | |
| 377 | cu = async_tensor_h2d(cu, device=device) |
| 378 | return PCPGlobalChunkPlan( |
| 379 | row_start_cu=cu[0], |
| 380 | global_cu=cu[1], |
| 381 | padded_local_cu=cu[2], |
| 382 | padded_local_total=padded_total, |
| 383 | total=total, |
| 384 | deinterleave_idx=async_tensor_h2d(idx, device=device), |
| 385 | ) |
| 386 | |
| 387 | |
| 388 | @dataclass |
| 389 | class DeepseekV32IndexerPrefillChunkMetadata: |
| 390 | block_table: torch.Tensor |
| 391 | # Under DCP (dcp_world_size > 1) these hold this rank's local row bounds; |
| 392 | # otherwise they hold the global bounds. |
| 393 | cu_seqlen_ks: torch.Tensor |
| 394 | cu_seqlen_ke: torch.Tensor |
| 395 | cu_seq_lens: torch.Tensor |
| 396 | token_to_seq: torch.Tensor |
| 397 | total_seq_lens: int |
| 398 | token_start: int |
| 399 | token_end: int |
| 400 | num_reqs: int |
| 401 | skip_kv_gather: bool = False |
| 402 | local_cu_seq_lens: torch.Tensor | None = None |
| 403 | local_total_seq_lens: int = 0 |
| 404 | max_local_total_seq_lens: int = 0 |
| 405 | |
| 406 | pcp_deinterleave_idx: torch.Tensor | None = None |
| 407 | |
| 408 | |
| 409 | class BuildPrefillChunkMetadataKernel( |
| 410 | VllmTritonJitKernel["BuildPrefillChunkMetadataKernel.CompileKey"] |
| 411 | ): |
| 412 | BLOCK_SIZE = 1024 |
| 413 | |
| 414 | @dataclass(frozen=True) |
| 415 | class CompileKey: |
| 416 | query_slice_start: int |
| 417 | query_slice_stop: int |
| 418 | dcp_rank: int |
| 419 | dcp_world: int |
| 420 | dcp_interleave: int |
| 421 | block_size: int |
| 422 | compress_ratio: int |
| 423 | input_variant: TritonPointerInputVariant |
| 424 | |
| 425 | @staticmethod |
| 426 | @triton.jit |
| 427 | def kernel( |
| 428 | # Inputs |
| 429 | query_start_loc_ptr, |
| 430 | uncompressed_seq_lens_ptr, |
| 431 | cu_compressed_seq_lens_ptr, |
| 432 | # Row-start base for cu_seq_len_ks/ke: local cumulative lens under DCP, |
| 433 | # aliases cu_compressed_seq_lens_ptr otherwise. |
| 434 | row_start_cu_compressed_seq_lens_ptr, |
| 435 | # Outputs |
| 436 | token_to_seq_ptr, |
| 437 | cu_compressed_seq_len_ks_ptr, |
| 438 | cu_compressed_seq_len_ke_ptr, |
| 439 | query_slice_start, |
| 440 | query_slice_stop, |
| 441 | DCP_RANK, |
| 442 | DCP_WORLD, |
| 443 | DCP_INTERLEAVE, |
| 444 | BLOCK_SIZE: tl.constexpr, |
| 445 | COMPRESS_RATIO: tl.constexpr, |
| 446 | ): |
| 447 | batch_idx = tl.program_id(0) |
| 448 | |
| 449 | query_start = tl.load(query_start_loc_ptr + batch_idx) |
| 450 | query_end = tl.load(query_start_loc_ptr + batch_idx + 1) |
| 451 | query_len = query_end - query_start |
| 452 | |
| 453 | seq_start = tl.load(cu_compressed_seq_lens_ptr + batch_idx) |
| 454 | seq_end = tl.load(cu_compressed_seq_lens_ptr + batch_idx + 1) |
| 455 | compressed_seq_len = seq_end - seq_start |
| 456 | |
| 457 | # Row start for the (possibly localized) cu_seq_len_ks/ke. Equals seq_start |
| 458 | # when DCP is disabled (the pointer aliases cu_compressed_seq_lens_ptr). |
| 459 | row_start = tl.load(row_start_cu_compressed_seq_lens_ptr + batch_idx) |
| 460 | |
| 461 | uncompressed_seq_len = tl.load(uncompressed_seq_lens_ptr + batch_idx) |
| 462 | start_pos = uncompressed_seq_len - query_len |
| 463 | |
| 464 | for i in range(0, query_len, BLOCK_SIZE): |
| 465 | offset = i + tl.arange(0, BLOCK_SIZE) |
| 466 | abs_pos = query_start + offset |
| 467 | mask = ( |
| 468 | (offset < query_len) |
| 469 | & (abs_pos >= query_slice_start) |
| 470 | & (abs_pos < query_slice_stop) |
| 471 | ) |
| 472 | out_pos = abs_pos - query_slice_start |
| 473 | |
| 474 | # cu_seq_len_ks: row start in the gathered K buffer. |
| 475 | tl.store(cu_compressed_seq_len_ks_ptr + out_pos, row_start, mask=mask) |
| 476 | |
| 477 | # cu_seq_len_ke: row start + per-token context length. Under DCP the |
| 478 | # global per-token length is sharded across ranks. |
| 479 | global_ctx = start_pos + 1 + offset |
| 480 | len_per_token = global_ctx // COMPRESS_RATIO |
| 481 | if DCP_WORLD > 1: |
| 482 | # Per-rank local context length under interleave-aware DCP, matching |
| 483 | # get_dcp_local_seq_lens. K == 1 reduces to (len + world-1-rank)//world. |
| 484 | base = (len_per_token // DCP_INTERLEAVE // DCP_WORLD) * DCP_INTERLEAVE |
| 485 | remainder = len_per_token - base * DCP_WORLD |
| 486 | remainder = tl.minimum( |
| 487 | tl.maximum(remainder - DCP_RANK * DCP_INTERLEAVE, 0), DCP_INTERLEAVE |
| 488 | ) |
| 489 | len_per_token = base + remainder |
| 490 | tl.store( |
| 491 | cu_compressed_seq_len_ke_ptr + out_pos, |
| 492 | row_start + len_per_token, |
| 493 | mask=mask, |
| 494 | ) |
| 495 | |
| 496 | # Compute token_to_seq |
| 497 | for i in range(0, compressed_seq_len, BLOCK_SIZE): |
| 498 | offset = i + tl.arange(0, BLOCK_SIZE) |
| 499 | mask = offset < compressed_seq_len |
| 500 | tl.store(token_to_seq_ptr + seq_start + offset, batch_idx, mask=mask) |
| 501 | |
| 502 | def dispatch( # type: ignore[override] |
| 503 | self, |
| 504 | *, |
| 505 | query_slice_start: int, |
| 506 | query_slice_stop: int, |
| 507 | DCP_RANK: int, |
| 508 | DCP_WORLD: int, |
| 509 | DCP_INTERLEAVE: int, |
| 510 | BLOCK_SIZE: int, |
| 511 | COMPRESS_RATIO: int, |
| 512 | input_variant: TritonPointerInputVariant, |
| 513 | ) -> CompileKey: |
| 514 | return self.CompileKey( |
| 515 | query_slice_start=triton_scalar_specialization_rep(query_slice_start), |
| 516 | query_slice_stop=triton_scalar_specialization_rep(query_slice_stop), |
| 517 | dcp_rank=triton_scalar_specialization_rep(DCP_RANK), |
| 518 | dcp_world=triton_scalar_specialization_rep(DCP_WORLD), |
| 519 | dcp_interleave=triton_scalar_specialization_rep(DCP_INTERLEAVE), |
| 520 | block_size=BLOCK_SIZE, |
| 521 | compress_ratio=COMPRESS_RATIO, |
| 522 | input_variant=input_variant, |
| 523 | ) |
| 524 | |
| 525 | def get_warmup_keys(self, vllm_config: VllmConfig) -> list[CompileKey]: |
| 526 | max_tokens = max(1, min(vllm_config.scheduler_config.max_num_batched_tokens, 8)) |
| 527 | hf_text_config = vllm_config.model_config.hf_text_config |
| 528 | parallel_config = vllm_config.parallel_config |
| 529 | dcp_world = parallel_config.decode_context_parallel_size |
| 530 | dcp_interleave = parallel_config.cp_kv_cache_interleave_size |
| 531 | dcp_rank = get_dcp_group().rank_in_group if dcp_world > 1 else 0 |
| 532 | dcp_cases = [dict(DCP_RANK=dcp_rank, DCP_WORLD=dcp_world)] |
| 533 | if parallel_config.prefill_context_parallel_size > 1 and dcp_world > 1: |
| 534 | # PCP+DCP prefill constructs a global PCP chunk plan and therefore |
| 535 | # dispatches this kernel with already-localized row starts. The |
| 536 | # runtime key is normalized to rank 0/world 1 in that path; warm it |
| 537 | # explicitly so one PCP rank cannot enter a collective while |
| 538 | # another rank is still compiling the first real inference batch. |
| 539 | dcp_cases.append(dict(DCP_RANK=0, DCP_WORLD=1)) |
| 540 | compress_ratios = tuple( |
| 541 | dict.fromkeys( |
| 542 | max(1, int(ratio)) |
| 543 | for ratio in ( |
| 544 | *(getattr(hf_text_config, "compress_ratios", None) or (1,)), |
| 545 | getattr(hf_text_config, "index_kpool", 1) or 1, |
| 546 | ) |
| 547 | ) |
| 548 | ) |
| 549 | index_kpool = getattr(hf_text_config, "index_kpool", None) |
| 550 | if index_kpool and index_kpool > 1 and index_kpool not in compress_ratios: |
| 551 | compress_ratios = compress_ratios + (index_kpool,) |
| 552 | return self._trace_dispatch(self.dispatch)( |
| 553 | zip_inputs(*dcp_cases), |
| 554 | # Cover Triton's divisible, exact-one, and generic i32 classes. |
| 555 | query_slice_start=(0, 1, 2), |
| 556 | query_slice_stop=(1, 2 * max_tokens - 1, 2 * max_tokens), |
| 557 | DCP_INTERLEAVE=dcp_interleave, |
| 558 | BLOCK_SIZE=self.BLOCK_SIZE, |
| 559 | COMPRESS_RATIO=list(compress_ratios), |
| 560 | # PCP's global cumulative lengths are the second row of one packed |
| 561 | # allocation, so their pointer is not always 16-byte aligned. |
| 562 | # Prefill chunking can independently unalign uncompressed lengths. |
| 563 | input_variant=( |
| 564 | TritonPointerInputVariant.from_alignment( |
| 565 | uncompressed_seq_lens=True, |
| 566 | cu_compressed_seq_lens=True, |
| 567 | ), |
| 568 | TritonPointerInputVariant.from_alignment( |
| 569 | uncompressed_seq_lens=True, |
| 570 | cu_compressed_seq_lens=False, |
| 571 | ), |
| 572 | TritonPointerInputVariant.from_alignment( |
| 573 | uncompressed_seq_lens=False, |
| 574 | cu_compressed_seq_lens=True, |
| 575 | ), |
| 576 | TritonPointerInputVariant.from_alignment( |
| 577 | uncompressed_seq_lens=False, |
| 578 | cu_compressed_seq_lens=False, |
| 579 | ), |
| 580 | ), |
| 581 | ) |
| 582 | |
| 583 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 584 | int32_ptr = TritonWarmupTensor(torch.int32) |
| 585 | return dict( |
| 586 | query_start_loc=int32_ptr, |
| 587 | uncompressed_seq_lens=compile_key.input_variant.pointer( |
| 588 | "uncompressed_seq_lens", torch.int32 |
| 589 | ), |
| 590 | cu_compressed_seq_lens=compile_key.input_variant.pointer( |
| 591 | "cu_compressed_seq_lens", torch.int32 |
| 592 | ), |
| 593 | row_start_cu_compressed_seq_lens=int32_ptr, |
| 594 | token_to_seq=int32_ptr, |
| 595 | cu_compressed_seq_len_ks=int32_ptr, |
| 596 | cu_compressed_seq_len_ke=int32_ptr, |
| 597 | query_slice_start=compile_key.query_slice_start, |
| 598 | query_slice_stop=compile_key.query_slice_stop, |
| 599 | DCP_RANK=compile_key.dcp_rank, |
| 600 | DCP_WORLD=compile_key.dcp_world, |
| 601 | DCP_INTERLEAVE=compile_key.dcp_interleave, |
| 602 | num_reqs=1, |
| 603 | COMPRESS_RATIO=compile_key.compress_ratio, |
| 604 | ) |
| 605 | |
| 606 | @kernel_launcher |
| 607 | def __call__( |
| 608 | self, |
| 609 | query_start_loc: torch.Tensor, |
| 610 | uncompressed_seq_lens: torch.Tensor, |
| 611 | cu_compressed_seq_lens: torch.Tensor, |
| 612 | row_start_cu_compressed_seq_lens: torch.Tensor, |
| 613 | token_to_seq: torch.Tensor, |
| 614 | cu_compressed_seq_len_ks: torch.Tensor, |
| 615 | cu_compressed_seq_len_ke: torch.Tensor, |
| 616 | query_slice_start: int, |
| 617 | query_slice_stop: int, |
| 618 | DCP_RANK: int, |
| 619 | DCP_WORLD: int, |
| 620 | DCP_INTERLEAVE: int, |
| 621 | *, |
| 622 | num_reqs: int, |
| 623 | COMPRESS_RATIO: int, |
| 624 | ) -> LaunchSpec: |
| 625 | return (num_reqs,), dict( |
| 626 | BLOCK_SIZE=self.BLOCK_SIZE, |
| 627 | COMPRESS_RATIO=COMPRESS_RATIO, |
| 628 | ) |
| 629 | |
| 630 | |
| 631 | @dataclass |
| 632 | class DeepseekV32IndexerPrefillMetadata: |
| 633 | chunks: list[DeepseekV32IndexerPrefillChunkMetadata] |
| 634 | max_prefill_seq_len: int = -1 |
| 635 | |
| 636 | |
| 637 | @dataclass |
| 638 | class DeepSeekV32IndexerDecodeMetadata: |
| 639 | block_table: torch.Tensor |
| 640 | # seq_lens: per-token effective context lengths. |
| 641 | # - flatten path / plain decode: 1D (batch_size,) |
| 642 | # - native MTP path: 2D (B, next_n) where [b,j] = L_b - next_n + j + 1 |
| 643 | # Both fp8_fp4_paged_mqa_logits and the topk kernels accept both shapes. |
| 644 | seq_lens: torch.Tensor |
| 645 | decode_lens: torch.Tensor |
| 646 | requires_padding: bool |
| 647 | schedule_metadata: torch.Tensor |
| 648 | global_seq_lens: torch.Tensor | None = None |
| 649 | per_req_decode_lens: torch.Tensor | None = None |
| 650 | decode_is_uniform: bool = True |
| 651 | write_max_decode_len: int = 0 |
| 652 | indices: torch.Tensor | None = None |
| 653 | |
| 654 | |
| 655 | @dataclass |
| 656 | class DeepseekV32IndexerMetadata: |
| 657 | # FIXME (zyongye) |
| 658 | # hacky way to access the data now, need to be in chunked meta |
| 659 | seq_lens: torch.Tensor |
| 660 | max_seq_len: int |
| 661 | slot_mapping: torch.Tensor |
| 662 | |
| 663 | # New for MLA (compared to FlashAttention) |
| 664 | # For handling prefill decode split |
| 665 | num_decodes: int |
| 666 | num_decode_tokens: int |
| 667 | num_prefills: int |
| 668 | num_prefill_tokens: int |
| 669 | |
| 670 | decode: DeepSeekV32IndexerDecodeMetadata | None = None |
| 671 | prefill: DeepseekV32IndexerPrefillMetadata | None = None |
| 672 | |
| 673 | |
| 674 | @triton.jit(do_not_specialize=["num_reqs", "num_actual_tokens", "num_tokens"]) |
| 675 | def _kpool_tail_slot_mapping_kernel( |
| 676 | slot_mapping_ptr, |
| 677 | block_table_ptr, |
| 678 | block_table_stride, |
| 679 | query_start_loc_ptr, |
| 680 | positions_ptr, |
| 681 | out_ptr, |
| 682 | num_reqs, |
| 683 | num_actual_tokens, |
| 684 | num_tokens, |
| 685 | kpool, |
| 686 | BLOCK: tl.constexpr, |
| 687 | ): |
| 688 | pid = tl.program_id(0) |
| 689 | if pid < num_reqs: |
| 690 | start = tl.load(query_start_loc_ptr + pid) |
| 691 | # Tokens past the last request's boundary (if any) also map to it. |
| 692 | end = tl.load(query_start_loc_ptr + pid + 1) |
| 693 | end = tl.where(pid == num_reqs - 1, num_actual_tokens, end) |
| 694 | own_block = tl.load(block_table_ptr + pid * block_table_stride).to(tl.int64) |
| 695 | for i in range(start, end, BLOCK): |
| 696 | offs = i + tl.arange(0, BLOCK) |
| 697 | mask = offs < end |
| 698 | pos = tl.load(positions_ptr + offs, mask=mask, other=0).to(tl.int64) |
| 699 | tl.store(out_ptr + offs, own_block * kpool + pos % kpool, mask=mask) |
| 700 | else: |
| 701 | offs = num_actual_tokens + (pid - num_reqs) * BLOCK + tl.arange(0, BLOCK) |
| 702 | mask = offs < num_tokens |
| 703 | tl.store(out_ptr + offs, tl.load(slot_mapping_ptr + offs, mask=mask), mask=mask) |
| 704 | |
| 705 | |
| 706 | def compute_kpool_tail_slot_mapping( |
| 707 | slot_mapping: torch.Tensor, |
| 708 | block_table: torch.Tensor, |
| 709 | query_start_loc: torch.Tensor, |
| 710 | positions: torch.Tensor, |
| 711 | num_actual_tokens: int, |
| 712 | num_reqs: int, |
| 713 | kpool: int, |
| 714 | out: torch.Tensor | None = None, |
| 715 | ) -> torch.Tensor: |
| 716 | """Map every token to its request's one circular tail block.""" |
| 717 | if out is None: |
| 718 | out = torch.empty_like(slot_mapping) |
| 719 | else: |
| 720 | assert out.shape == slot_mapping.shape |
| 721 | if slot_mapping.is_cuda and slot_mapping.dim() == 1 and num_reqs > 0: |
| 722 | block = 256 |
| 723 | num_tokens = slot_mapping.shape[0] |
| 724 | num_actual_tokens = min(num_actual_tokens, num_tokens) |
| 725 | grid = (num_reqs + triton.cdiv(num_tokens - num_actual_tokens, block),) |
| 726 | _kpool_tail_slot_mapping_kernel[grid]( |
| 727 | slot_mapping, |
| 728 | block_table, |
| 729 | block_table.stride(0), |
| 730 | query_start_loc, |
| 731 | positions, |
| 732 | out, |
| 733 | num_reqs, |
| 734 | num_actual_tokens, |
| 735 | num_tokens, |
| 736 | kpool, |
| 737 | BLOCK=block, |
| 738 | num_warps=4, |
| 739 | ) |
| 740 | return out |
| 741 | # Torch fallback: CPU tensors (the CPU unit tests), non-1D or empty inputs. |
| 742 | # Production always takes the Triton path above — spec-decode tokens arrive |
| 743 | # flattened token-major, so slot_mapping is 1D there too. |
| 744 | out.copy_(slot_mapping) |
| 745 | if num_actual_tokens == 0: |
| 746 | return out |
| 747 | tokens = torch.arange(num_actual_tokens, device=slot_mapping.device) |
| 748 | req = torch.searchsorted(query_start_loc, tokens, right=True) - 1 |
| 749 | req = req.clamp_(min=0, max=num_reqs - 1) |
| 750 | own_block = block_table[:num_reqs, 0].index_select(0, req).to(torch.int64) |
| 751 | pos = positions[:num_actual_tokens].to(torch.int64) |
| 752 | out[:num_actual_tokens] = own_block * kpool + torch.remainder(pos, kpool) |
| 753 | return out |
| 754 | |
| 755 | |
| 756 | class KpoolTailMetadataBuilder(AttentionMetadataBuilder): |
| 757 | """Build only the circular slot mapping needed by the storage-only tail.""" |
| 758 | |
| 759 | _cudagraph_support = AttentionCGSupport.ALWAYS |
| 760 | supports_update_block_table = False |
| 761 | reorder_batch_threshold = None |
| 762 | |
| 763 | def __init__( |
| 764 | self, |
| 765 | kv_cache_spec: AttentionSpec, |
| 766 | layer_names: list[str], |
| 767 | vllm_config: VllmConfig, |
| 768 | device: torch.device, |
| 769 | ): |
| 770 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 771 | self.slot_mapping_buffer = torch.empty( |
| 772 | vllm_config.scheduler_config.max_num_batched_tokens, |
| 773 | dtype=torch.int64, |
| 774 | device=device, |
| 775 | ) |
| 776 | |
| 777 | def build( |
| 778 | self, |
| 779 | common_prefix_len: int, |
| 780 | common_attn_metadata: CommonAttentionMetadata, |
| 781 | fast_build: bool = False, |
| 782 | ) -> DeepseekV32IndexerMetadata: |
| 783 | num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( |
| 784 | split_decodes_and_prefills(common_attn_metadata) |
| 785 | ) |
| 786 | slot_mapping = common_attn_metadata.slot_mapping |
| 787 | positions = common_attn_metadata.positions |
| 788 | if positions is not None: |
| 789 | slot_mapping_buffer = self.slot_mapping_buffer[ |
| 790 | : slot_mapping.numel() |
| 791 | ].view_as(slot_mapping) |
| 792 | slot_mapping = compute_kpool_tail_slot_mapping( |
| 793 | slot_mapping, |
| 794 | common_attn_metadata.block_table_tensor, |
| 795 | common_attn_metadata.query_start_loc, |
| 796 | positions, |
| 797 | common_attn_metadata.num_actual_tokens, |
| 798 | common_attn_metadata.num_reqs, |
| 799 | self.kv_cache_spec.block_size, |
| 800 | out=slot_mapping_buffer, |
| 801 | ) |
| 802 | return DeepseekV32IndexerMetadata( |
| 803 | seq_lens=common_attn_metadata.seq_lens, |
| 804 | max_seq_len=common_attn_metadata.max_seq_len, |
| 805 | slot_mapping=slot_mapping, |
| 806 | num_decodes=num_decodes, |
| 807 | num_decode_tokens=num_decode_tokens, |
| 808 | num_prefills=num_prefills, |
| 809 | num_prefill_tokens=num_prefill_tokens, |
| 810 | ) |
| 811 | |
| 812 | |
| 813 | def get_max_prefill_buffer_size(vllm_config: VllmConfig): |
| 814 | max_model_len = vllm_config.model_config.max_model_len |
| 815 | # NOTE(Chen): 40 is a magic number for controlling the prefill buffer size. |
| 816 | # Each entry is 128 fp8 bytes and 4 scale bytes for a total of 132 bytes. |
| 817 | # The flashmla_sparse backend uses a workspace size of 5 * max_model_len. |
| 818 | # The memory usage of the workspace there is 576 * 2 bytes; so we size this as |
| 819 | # (576 * 2 // 132) * 5 = 40 to maximize this workspace size while still fitting |
| 820 | # within the flashmla_sparse workspace. |
| 821 | # For DeepSeek-V3.2, the max_model_len is 163840. |
| 822 | # 40 * 163840 * 132 = 865075200 bytes = 825 MB |
| 823 | return max_model_len * 40 |
| 824 | |
| 825 | |
| 826 | def _supports_varlen_paged_mqa_logits() -> bool: |
| 827 | return ( |
| 828 | current_platform.is_cuda() |
| 829 | and current_platform.is_device_capability_family(100) |
| 830 | and has_deep_gemm() |
| 831 | ) |
| 832 | |
| 833 | |
| 834 | def _supports_flattened_device_query_lens() -> bool: |
| 835 | return ( |
| 836 | current_platform.is_cuda() |
| 837 | and current_platform.is_device_capability_family(90) |
| 838 | and has_deep_gemm() |
| 839 | ) |
| 840 | |
| 841 | |
| 842 | def _rocm_supports_flattened_device_query_lens(vllm_config: VllmConfig) -> bool: |
| 843 | model_config = vllm_config.model_config |
| 844 | return ( |
| 845 | current_platform.is_rocm() |
| 846 | and model_config is not None |
| 847 | and "DeepseekV4ForCausalLM" in model_config.architectures |
| 848 | ) |
| 849 | |
| 850 | |
| 851 | def _supports_native_decode(next_n: int) -> bool: |
| 852 | """Whether decode can pass `next_n` Q rows per request to the kernel |
| 853 | instead of flattening to one single-token row per query, which re-reads |
| 854 | the KV tile once per row. |
| 855 | """ |
| 856 | if not (current_platform.is_cuda() and has_deep_gemm()): |
| 857 | return next_n in (1, 2) |
| 858 | if current_platform.is_device_capability_family(100): |
| 859 | return True |
| 860 | if current_platform.is_device_capability_family(90): |
| 861 | return native_next_n_supported(next_n) |
| 862 | return next_n in (1, 2) |
| 863 | |
| 864 | |
| 865 | def _use_flattening(vllm_config: VllmConfig) -> bool: |
| 866 | speculative_config = vllm_config.speculative_config |
| 867 | next_n = 1 + vllm_config.num_speculative_tokens |
| 868 | return not _supports_native_decode(next_n) or ( |
| 869 | speculative_config is not None |
| 870 | and speculative_config.enable_adaptive_verification |
| 871 | and ( |
| 872 | _supports_flattened_device_query_lens() |
| 873 | or _rocm_supports_flattened_device_query_lens(vllm_config) |
| 874 | ) |
| 875 | ) |
| 876 | |
| 877 | |
| 878 | class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder): |
| 879 | # The indexer opts out of the shared reorder-threshold vote (see __init__), |
| 880 | # so this is None; its own split uses self.decode_threshold. |
| 881 | reorder_batch_threshold: int | None = None |
| 882 | requires_block_table_width = True |
| 883 | |
| 884 | @classmethod |
| 885 | def get_cudagraph_support( |
| 886 | cls, |
| 887 | vllm_config: VllmConfig, |
| 888 | kv_cache_spec: KVCacheSpec, |
| 889 | ) -> AttentionCGSupport: |
| 890 | if _supports_varlen_paged_mqa_logits() or _use_flattening(vllm_config): |
| 891 | return AttentionCGSupport.ALWAYS |
| 892 | return AttentionCGSupport.UNIFORM_BATCH |
| 893 | |
| 894 | def __init__(self, *args, block_table_width: int, **kwargs) -> None: |
| 895 | super().__init__(*args, **kwargs) |
| 896 | scheduler_config = self.vllm_config.scheduler_config |
| 897 | parallel_config = self.vllm_config.parallel_config |
| 898 | self.dcp_world_size = parallel_config.decode_context_parallel_size |
| 899 | self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 |
| 900 | self.pcp_world_size = parallel_config.prefill_context_parallel_size |
| 901 | self.use_pcp = self.pcp_world_size > 1 |
| 902 | self.pcp_rank = get_pcp_group().rank_in_group if self.use_pcp else 0 |
| 903 | self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size |
| 904 | # KV compression (DeepseekV4). Default to 1 for no compression. |
| 905 | self.compress_ratio = 1 |
| 906 | if isinstance(self.kv_cache_spec, MLAAttentionSpec): |
| 907 | assert isinstance(self.kv_cache_spec.tokens_per_state, int) |
| 908 | self.compress_ratio = self.kv_cache_spec.tokens_per_state |
| 909 | # NOTE(Chen):an estimated max size of flattened_kv. Need to double check. |
| 910 | # Counted in compressed rows, like the chunker's seq_lens and the |
| 911 | # workspace. |
| 912 | self.max_prefill_buffer_size = ( |
| 913 | get_max_prefill_buffer_size(self.vllm_config) // self.compress_ratio |
| 914 | ) |
| 915 | self.num_speculative_tokens = ( |
| 916 | self.vllm_config.speculative_config.num_speculative_tokens |
| 917 | if self.vllm_config.speculative_config |
| 918 | else 0 |
| 919 | ) |
| 920 | self.indexer_uses_fp4 = dsa_indexer_uses_fp4(self.vllm_config) |
| 921 | |
| 922 | next_n = self.num_speculative_tokens + 1 |
| 923 | self.decode_threshold = next_n |
| 924 | self.reorder_batch_threshold = None |
| 925 | self.use_flattening = _use_flattening(self.vllm_config) |
| 926 | self.supports_varlen = _supports_varlen_paged_mqa_logits() |
| 927 | logger.info_once( |
| 928 | "DSA indexer decode path: use_flattening=%s supports_varlen=%s " |
| 929 | "(next_n=%d, use_fp4_cache=%s)", |
| 930 | self.use_flattening, |
| 931 | self.supports_varlen, |
| 932 | next_n, |
| 933 | self.indexer_uses_fp4, |
| 934 | ) |
| 935 | |
| 936 | sm_count = num_compute_units(self.device.index) |
| 937 | self.num_sms = sm_count |
| 938 | |
| 939 | self.offsets_buffer = torch.arange( |
| 940 | next_n, device=self.device, dtype=torch.int32 |
| 941 | ) |
| 942 | self.decode_lens_buffer = torch.zeros( |
| 943 | (scheduler_config.max_num_batched_tokens,), |
| 944 | dtype=torch.int32, |
| 945 | device=self.device, |
| 946 | ) |
| 947 | self.per_req_decode_lens_buffer = torch.zeros( |
| 948 | (scheduler_config.max_num_batched_tokens,), |
| 949 | dtype=torch.int32, |
| 950 | device=self.device, |
| 951 | ) |
| 952 | # Shared workspace for decode seq_lens. Native MTP views this as |
| 953 | # (B, max_decode_len) at runtime, keeping context_lens contiguous even |
| 954 | # when max_decode_len is smaller than next_n. |
| 955 | self.decode_seq_lens_buffer = torch.zeros( |
| 956 | (scheduler_config.max_num_batched_tokens,), |
| 957 | dtype=torch.int32, |
| 958 | device=self.device, |
| 959 | ) |
| 960 | self.global_decode_seq_lens_buffer = torch.zeros( |
| 961 | (scheduler_config.max_num_batched_tokens,), |
| 962 | dtype=torch.int32, |
| 963 | device=self.device, |
| 964 | ) |
| 965 | self.decode_indices_buffer = torch.zeros( |
| 966 | (scheduler_config.max_num_batched_tokens,), |
| 967 | dtype=torch.int32, |
| 968 | device=self.device, |
| 969 | ) |
| 970 | self.arange_buffer = torch.arange( |
| 971 | max( |
| 972 | scheduler_config.max_num_seqs * next_n, |
| 973 | scheduler_config.max_num_batched_tokens, |
| 974 | ), |
| 975 | dtype=torch.int32, |
| 976 | device=self.device, |
| 977 | ) |
| 978 | # Materialize the rank on device during builder initialization. Creating |
| 979 | # this scalar in build() would introduce a GPU<->CPU sync in the decode |
| 980 | # hot path. |
| 981 | self.dcp_rank_tensor = torch.tensor( |
| 982 | self.dcp_rank, dtype=torch.int32, device=self.device |
| 983 | ) |
| 984 | self.expanded_block_table_buffer = torch.zeros( |
| 985 | (scheduler_config.max_num_batched_tokens, block_table_width), |
| 986 | dtype=torch.int32, |
| 987 | device=self.device, |
| 988 | ) |
| 989 | |
| 990 | # See: DeepGMM/csrc/apis/attention.hpp. Sized for one slot per SM; |
| 991 | # build() narrows it to whatever the kernel actually schedules. |
| 992 | self.scheduler_metadata_buffer = torch.empty( |
| 993 | (self.num_sms + 1, 2), dtype=torch.int32, device=self.device |
| 994 | ) |
| 995 | |
| 996 | if self.dcp_world_size > 1 and self.compress_ratio > 1: |
| 997 | raise NotImplementedError( |
| 998 | "DCP is not supported with sparse indexer KV compression " |
| 999 | f"(compress_ratio={self.compress_ratio})." |
| 1000 | ) |
| 1001 | |
| 1002 | # Pre-allocate buffers for CUDA graph compatibility when |
| 1003 | if self.compress_ratio > 1: |
| 1004 | # compress_ratio > 1 (DeepseekV4) |
| 1005 | # Compressed slot mapping output buffer |
| 1006 | self.compressed_slot_mapping_buffer = torch.zeros( |
| 1007 | (scheduler_config.max_num_batched_tokens,), |
| 1008 | dtype=torch.int64, |
| 1009 | device=self.device, |
| 1010 | ) |
| 1011 | # Buffer for compressed seq_lens in decode path |
| 1012 | self.expanded_seq_lens_buffer = torch.zeros( |
| 1013 | (scheduler_config.max_num_batched_tokens,), |
| 1014 | dtype=torch.int32, |
| 1015 | device=self.device, |
| 1016 | ) |
| 1017 | self.indexer_decode_block_table_buffer: torch.Tensor | None = None |
| 1018 | self._max_num_batched_tokens = scheduler_config.max_num_batched_tokens |
| 1019 | |
| 1020 | def _dcp_localize_decode_seq_lens( |
| 1021 | self, |
| 1022 | seq_lens: torch.Tensor, |
| 1023 | num_decodes: int, |
| 1024 | seq_lens_is_buffer_view: bool, |
| 1025 | ) -> torch.Tensor: |
| 1026 | local_seq_lens = get_dcp_local_seq_lens( |
| 1027 | seq_lens, |
| 1028 | self.dcp_world_size, |
| 1029 | self.dcp_rank_tensor, |
| 1030 | self.cp_kv_cache_interleave_size, |
| 1031 | ) |
| 1032 | if seq_lens_is_buffer_view: |
| 1033 | seq_lens.copy_(local_seq_lens) |
| 1034 | return seq_lens |
| 1035 | |
| 1036 | out = self.decode_seq_lens_buffer[:num_decodes] |
| 1037 | out.copy_(local_seq_lens) |
| 1038 | return out |
| 1039 | |
| 1040 | def _prepare_decode_tensors( |
| 1041 | self, |
| 1042 | seq_lens: torch.Tensor, |
| 1043 | block_table: torch.Tensor, |
| 1044 | decode_lens: torch.Tensor, |
| 1045 | decode_lens_cpu: torch.Tensor, |
| 1046 | query_start_loc: torch.Tensor, |
| 1047 | num_decodes: int, |
| 1048 | num_decode_tokens: int, |
| 1049 | use_native: bool, |
| 1050 | next_n: int, |
| 1051 | max_decode_len: int, |
| 1052 | ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, int, bool]: |
| 1053 | """Prepare native or per-token flattened decode tensors.""" |
| 1054 | spec_config = self.vllm_config.speculative_config |
| 1055 | adaptive = bool(spec_config and spec_config.enable_adaptive_verification) |
| 1056 | min_decode_len = int(decode_lens_cpu.min().item()) |
| 1057 | if not use_native: |
| 1058 | assert self.decode_seq_lens_buffer.dim() == 1 |
| 1059 | if ( |
| 1060 | not self.supports_varlen |
| 1061 | and (num_decodes == 1 or not adaptive) |
| 1062 | and min_decode_len == max_decode_len |
| 1063 | and num_decodes * max_decode_len == num_decode_tokens |
| 1064 | ): |
| 1065 | # Uniform decode lengths with no cudagraph token padding. |
| 1066 | _PREPARE_UNIFORM_DECODE_KERNEL( |
| 1067 | seq_lens, |
| 1068 | self.decode_seq_lens_buffer, |
| 1069 | block_table, |
| 1070 | self.expanded_block_table_buffer, |
| 1071 | self.decode_lens_buffer, |
| 1072 | num_decode_tokens, |
| 1073 | max_decode_len, |
| 1074 | ) |
| 1075 | self.decode_seq_lens_buffer[num_decode_tokens:] = 0 |
| 1076 | seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens] |
| 1077 | block_table = self.expanded_block_table_buffer[:num_decode_tokens] |
| 1078 | decode_lens = self.decode_lens_buffer[:num_decode_tokens] |
| 1079 | return seq_lens, block_table, decode_lens, num_decode_tokens, False |
| 1080 | else: |
| 1081 | # Variable decode lengths. |
| 1082 | # Assume 4 requests with seq_lens [10, 7, 12, 0] (the final req is |
| 1083 | # padding) and decode_lens [3, 1, 4, 0] in the below example comments. |
| 1084 | # The context lengths are therefore |
| 1085 | # [10-3, 7-1, 12-4, 0-0] = [7, 6, 8, 0]. |
| 1086 | |
| 1087 | # 3 + 1 + 4 + 0 = 8 |
| 1088 | actual_expanded = int(decode_lens_cpu.sum().item()) |
| 1089 | |
| 1090 | # Fuse expanded_base and expanded_starts into a single |
| 1091 | # repeat_interleave: |
| 1092 | # seq_len_i = (context_start[b] - query_start_loc[b]) + arange[i] + 1 |
| 1093 | # where context_start[b] = seq_lens[b] - decode_lens[b]. |
| 1094 | # Example: offsets = [7-0, 6-3, 8-4, 0-8] = [7, 3, 4, -8] |
| 1095 | # expanded_offsets = [7, 7, 7, 3, 4, 4, 4, 4] |
| 1096 | # result = [8, 9, 10, 7, 9, 10, 11, 12] |
| 1097 | expanded_offsets = torch.repeat_interleave( |
| 1098 | seq_lens - decode_lens - query_start_loc, |
| 1099 | decode_lens, |
| 1100 | output_size=actual_expanded, |
| 1101 | ) |
| 1102 | |
| 1103 | # [8, 9, 10, 7, 9, 10, 11, 12, ...] where ... is unused buffer space |
| 1104 | self.decode_seq_lens_buffer[:actual_expanded] = ( |
| 1105 | expanded_offsets + self.arange_buffer[:actual_expanded] + 1 |
| 1106 | ) |
| 1107 | self.decode_seq_lens_buffer[actual_expanded:] = 0 |
| 1108 | seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens] |
| 1109 | |
| 1110 | # Give each of the flattened entries the same block table row as the |
| 1111 | # original request. |
| 1112 | self.expanded_block_table_buffer[:actual_expanded] = ( |
| 1113 | torch.repeat_interleave( |
| 1114 | block_table, decode_lens, dim=0, output_size=actual_expanded |
| 1115 | ) |
| 1116 | ) |
| 1117 | if actual_expanded < num_decode_tokens: |
| 1118 | self.expanded_block_table_buffer[ |
| 1119 | actual_expanded:num_decode_tokens, 0 |
| 1120 | ] = 0 |
| 1121 | block_table = self.expanded_block_table_buffer[:num_decode_tokens] |
| 1122 | |
| 1123 | # All reqs now have decode_len=1 |
| 1124 | self.decode_lens_buffer[:num_decode_tokens] = 1 |
| 1125 | decode_lens = self.decode_lens_buffer[:num_decode_tokens] |
| 1126 | return seq_lens, block_table, decode_lens, num_decode_tokens, False |
| 1127 | else: |
| 1128 | # Native path: plain decode (next_n==1) or spec decode |
| 1129 | # with 2D per-token context lengths (next_n > 1). |
| 1130 | # |
| 1131 | # When decode_lens are not truly uniform (e.g. some requests have |
| 1132 | # decode_len < next_n due to padding or short prefills), the simple |
| 1133 | # reshape in sparse_attn_indexer won't work. Use pack_seq_triton |
| 1134 | # (requires_padding) instead. |
| 1135 | requires_padding = min_decode_len != max_decode_len |
| 1136 | if use_native and next_n > 1: |
| 1137 | assert self.decode_seq_lens_buffer.dim() == 1 |
| 1138 | # (B, max_decode_len): token j attends to |
| 1139 | # L - max_decode_len + j + 1 KV tokens. |
| 1140 | seq_lens_buffer = self.decode_seq_lens_buffer[ |
| 1141 | : num_decodes * max_decode_len |
| 1142 | ].view(num_decodes, max_decode_len) |
| 1143 | # Clamp at 0: padding requests have seq_len == 0, which would |
| 1144 | # otherwise make token 0 negative (next_n=2 gives 0-2+1+0 = -1). |
| 1145 | # Downstream kernels read these as uint32, turning -1 into ~4e9. |
| 1146 | seq_lens_buffer[:] = ( |
| 1147 | seq_lens.unsqueeze(1) |
| 1148 | - max_decode_len |
| 1149 | + 1 |
| 1150 | + self.offsets_buffer[:max_decode_len] |
| 1151 | ).clamp_(min=0) |
| 1152 | seq_lens = seq_lens_buffer |
| 1153 | return seq_lens, block_table, decode_lens, num_decodes, requires_padding |
| 1154 | |
| 1155 | def _prepare_global_decode_seq_lens( |
| 1156 | self, |
| 1157 | global_seq_lens: torch.Tensor | None, |
| 1158 | decode_lens: torch.Tensor, |
| 1159 | decode_lens_cpu: torch.Tensor, |
| 1160 | query_start_loc: torch.Tensor, |
| 1161 | num_decode_tokens: int, |
| 1162 | use_native: bool, |
| 1163 | max_decode_len: int, |
| 1164 | ) -> torch.Tensor | None: |
| 1165 | if global_seq_lens is None: |
| 1166 | return None |
| 1167 | if use_native or max_decode_len <= 1: |
| 1168 | return global_seq_lens |
| 1169 | |
| 1170 | actual_expanded = int(decode_lens_cpu.sum().item()) |
| 1171 | if actual_expanded > 0: |
| 1172 | expanded_offsets = torch.repeat_interleave( |
| 1173 | global_seq_lens - decode_lens - query_start_loc, |
| 1174 | decode_lens, |
| 1175 | output_size=actual_expanded, |
| 1176 | ) |
| 1177 | self.global_decode_seq_lens_buffer[:actual_expanded] = ( |
| 1178 | expanded_offsets + self.arange_buffer[:actual_expanded] + 1 |
| 1179 | ) |
| 1180 | self.global_decode_seq_lens_buffer[actual_expanded:num_decode_tokens] = 0 |
| 1181 | return self.global_decode_seq_lens_buffer[:num_decode_tokens] |
| 1182 | |
| 1183 | def _split_pcp_dcp_prefill_chunks( |
| 1184 | self, |
| 1185 | row_req_idx: np.ndarray, |
| 1186 | row_shard_rows: np.ndarray, |
| 1187 | row_query_lens_cpu: torch.Tensor, |
| 1188 | max_logits_bytes: int, |
| 1189 | request_offset: int, |
| 1190 | ) -> list[tuple[slice, slice]]: |
| 1191 | """Chunk by request rather than by row, so a split prefill's two rows |
| 1192 | charge their shared context once, then widen each chunk back to rows. |
| 1193 | |
| 1194 | ``row_shard_rows`` holds the whole request's largest DCP shard on |
| 1195 | every row, so the plan is identical on every PCP rank. |
| 1196 | """ |
| 1197 | row_bounds = request_row_bounds(row_req_idx) |
| 1198 | first_rows = row_bounds[:-1] |
| 1199 | # Each request's context, padded to whole DCP shards. |
| 1200 | seq_lens = row_shard_rows[first_rows] * self.dcp_world_size |
| 1201 | # Every rank holds a full-size chunk of a split request (shorter ones |
| 1202 | # are replicated), so rows x longest row is the same on every rank, |
| 1203 | # whichever row holds the short tail. |
| 1204 | row_query_lens = row_query_lens_cpu.numpy() |
| 1205 | query_lens = np.diff(row_bounds) * np.maximum.reduceat( |
| 1206 | row_query_lens, first_rows |
| 1207 | ) |
| 1208 | chunk_specs = self._split_indexer_prefill_chunks( |
| 1209 | torch.from_numpy(seq_lens.astype(np.int32)), |
| 1210 | torch.from_numpy(query_lens.astype(np.int32)), |
| 1211 | self.max_prefill_buffer_size, |
| 1212 | max_logits_bytes, |
| 1213 | ) |
| 1214 | return [ |
| 1215 | ( |
| 1216 | slice( |
| 1217 | request_offset + int(row_bounds[request_slice.start]), |
| 1218 | request_offset + int(row_bounds[request_slice.stop]), |
| 1219 | ), |
| 1220 | query_slice, |
| 1221 | ) |
| 1222 | for request_slice, query_slice in chunk_specs |
| 1223 | ] |
| 1224 | |
| 1225 | def _prefill_split_seq_lens(self, seq_lens_cpu: torch.Tensor) -> torch.Tensor: |
| 1226 | """Per-request KV lengths the prefill chunker budgets logits with; |
| 1227 | subclasses whose logits rows are wider than the context override.""" |
| 1228 | return seq_lens_cpu |
| 1229 | |
| 1230 | @staticmethod |
| 1231 | def _split_indexer_prefill_chunks( |
| 1232 | compressed_seq_lens_cpu: torch.Tensor, |
| 1233 | prefill_query_lens_cpu: torch.Tensor, |
| 1234 | workspace_size: int, |
| 1235 | max_logits_bytes: int, |
| 1236 | request_offset: int = 0, |
| 1237 | ) -> list[tuple[slice, slice]]: |
| 1238 | """Split this step's prefill requests into chunks, respecting: |
| 1239 | - N constraint: total_seq_lens <= workspace_size (existing O(N) |
| 1240 | workspace) |
| 1241 | - Logits constraint: M * N * 4 <= max_logits_bytes |
| 1242 | |
| 1243 | When a single request-level chunk still exceeds the logits budget, |
| 1244 | sub-chunks on the query dimension (M) to bound peak memory. |
| 1245 | |
| 1246 | Returns list of (req_slice, query_slice) tuples. |
| 1247 | """ |
| 1248 | chunks: list[tuple[slice, slice]] = [] |
| 1249 | n = len(compressed_seq_lens_cpu) |
| 1250 | max_logits_elems = max_logits_bytes // 4 |
| 1251 | end = 0 |
| 1252 | |
| 1253 | while end < n: |
| 1254 | start, chunk_m, chunk_n = end, 0, 0 |
| 1255 | |
| 1256 | while end < n: |
| 1257 | q, s = ( |
| 1258 | prefill_query_lens_cpu[end].item(), |
| 1259 | compressed_seq_lens_cpu[end].item(), |
| 1260 | ) |
| 1261 | new_m, new_n = chunk_m + q, chunk_n + s |
| 1262 | if new_n <= workspace_size and new_m * new_n <= max_logits_elems: |
| 1263 | chunk_m, chunk_n = new_m, new_n |
| 1264 | end += 1 |
| 1265 | else: |
| 1266 | break |
| 1267 | |
| 1268 | # A single request can exceed the budget, requiring sub-chunking |
| 1269 | # on the query dimension. |
| 1270 | if end == start: |
| 1271 | chunk_m, chunk_n = ( |
| 1272 | prefill_query_lens_cpu[end].item(), |
| 1273 | compressed_seq_lens_cpu[end].item(), |
| 1274 | ) |
| 1275 | end += 1 |
| 1276 | |
| 1277 | req_slice = slice(start + request_offset, end + request_offset) |
| 1278 | max_q = ( |
| 1279 | max(1, max_logits_elems // chunk_n) if chunk_n > 0 else max(1, chunk_m) |
| 1280 | ) |
| 1281 | for q_off in range(0, chunk_m, max_q): |
| 1282 | sub_m = min(max_q, chunk_m - q_off) |
| 1283 | chunks.append((req_slice, slice(q_off, q_off + sub_m))) |
| 1284 | |
| 1285 | return chunks |
| 1286 | |
| 1287 | def build( |
| 1288 | self, |
| 1289 | common_prefix_len: int, |
| 1290 | common_attn_metadata: CommonAttentionMetadata, |
| 1291 | fast_build: bool = False, |
| 1292 | ) -> DeepseekV32IndexerMetadata: |
| 1293 | num_reqs = common_attn_metadata.num_reqs |
| 1294 | num_tokens = common_attn_metadata.num_actual_tokens |
| 1295 | query_start_loc = common_attn_metadata.query_start_loc |
| 1296 | query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu |
| 1297 | seq_lens = common_attn_metadata.seq_lens |
| 1298 | slot_mapping = common_attn_metadata.slot_mapping |
| 1299 | block_table = common_attn_metadata.block_table_tensor |
| 1300 | dcp_local_seq_lens = common_attn_metadata.dcp_local_seq_lens |
| 1301 | |
| 1302 | compressed_slot_mapping = slot_mapping |
| 1303 | indexer_block_table = block_table |
| 1304 | if self.compress_ratio > 1: |
| 1305 | kernel_block_size = self.kernel_block_size |
| 1306 | if ( |
| 1307 | kernel_block_size is not None |
| 1308 | and self.kv_cache_spec.block_size != kernel_block_size |
| 1309 | and self.kv_cache_spec.block_size % kernel_block_size == 0 |
| 1310 | ): |
| 1311 | factor = self.kv_cache_spec.block_size // kernel_block_size |
| 1312 | indexer_block_table = (block_table[:, ::factor] // factor).contiguous() |
| 1313 | padded_num_tokens = num_tokens |
| 1314 | local_slot_mapping = slot_mapping |
| 1315 | if self.use_pcp: |
| 1316 | # The gathered layout holds each rank's local tokens, padded, in |
| 1317 | # rank order, so this rank's segment lines up with query_start_loc. |
| 1318 | padded_num_tokens = slot_mapping.shape[0] // self.pcp_world_size |
| 1319 | local_slot_mapping = slot_mapping[ |
| 1320 | self.pcp_rank * padded_num_tokens : (self.pcp_rank + 1) |
| 1321 | * padded_num_tokens |
| 1322 | ] |
| 1323 | compressed_slot_mapping = get_compressed_slot_mapping( |
| 1324 | num_tokens, |
| 1325 | local_slot_mapping, |
| 1326 | query_start_loc, |
| 1327 | seq_lens, |
| 1328 | indexer_block_table, |
| 1329 | self.kv_cache_spec.num_states, |
| 1330 | self.compress_ratio, |
| 1331 | out=self.compressed_slot_mapping_buffer, |
| 1332 | ) |
| 1333 | if self.pcp_world_size > 1: |
| 1334 | compressed_slot_mapping = get_pcp_group().all_gather( |
| 1335 | self.compressed_slot_mapping_buffer[:padded_num_tokens], |
| 1336 | dim=0, |
| 1337 | ) |
| 1338 | |
| 1339 | # PCP decode sharding keeps a zero-token placeholder on ranks that own |
| 1340 | # no request in a step so collectives retain a uniform rank shape. Do |
| 1341 | # not turn that placeholder into a zero-length indexer decode request. |
| 1342 | if num_tokens == 0: |
| 1343 | return DeepseekV32IndexerMetadata( |
| 1344 | seq_lens=seq_lens, |
| 1345 | max_seq_len=common_attn_metadata.max_seq_len, |
| 1346 | slot_mapping=compressed_slot_mapping, |
| 1347 | num_decodes=0, |
| 1348 | num_decode_tokens=0, |
| 1349 | num_prefills=0, |
| 1350 | num_prefill_tokens=0, |
| 1351 | prefill=None, |
| 1352 | decode=None, |
| 1353 | ) |
| 1354 | |
| 1355 | num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( |
| 1356 | split_decodes_and_prefills( |
| 1357 | common_attn_metadata, |
| 1358 | decode_threshold=self.decode_threshold, |
| 1359 | require_uniform=not (self.use_flattening or self.supports_varlen), |
| 1360 | treat_short_extends_as_decodes=not self.use_pcp, |
| 1361 | ) |
| 1362 | ) |
| 1363 | |
| 1364 | assert num_decodes + num_prefills == num_reqs |
| 1365 | assert num_decode_tokens + num_prefill_tokens == num_tokens |
| 1366 | |
| 1367 | prefill_metadata = None |
| 1368 | if num_prefills > 0: |
| 1369 | compressed_seq_lens = ( |
| 1370 | seq_lens // self.compress_ratio if self.compress_ratio > 1 else seq_lens |
| 1371 | ) |
| 1372 | # This CPU value is an upper bound for async-spec extend rows. It |
| 1373 | # is safe for chunking/allocation because CUDA metadata below is |
| 1374 | # built from exact device seq_lens and gather ignores the tail. |
| 1375 | assert common_attn_metadata.seq_lens_cpu_upper_bound is not None |
| 1376 | seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound |
| 1377 | compressed_seq_lens_cpu = ( |
| 1378 | seq_lens_cpu // self.compress_ratio |
| 1379 | if self.compress_ratio > 1 |
| 1380 | else seq_lens_cpu |
| 1381 | ) |
| 1382 | prefill_query_lens_cpu = torch.diff( |
| 1383 | query_start_loc_cpu[num_decodes : num_decodes + num_prefills + 1] |
| 1384 | ) |
| 1385 | max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 |
| 1386 | # Upper bound is exact for prefill rows (the `[num_decodes:]` |
| 1387 | # slice below). |
| 1388 | assert common_attn_metadata.seq_lens_cpu_upper_bound is not None |
| 1389 | seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound |
| 1390 | req_idx = None |
| 1391 | shard_rows = None |
| 1392 | if self.use_pcp and self.dcp_world_size > 1: |
| 1393 | # The gathered KV must be packed identically on every PCP rank: |
| 1394 | # chunk by request from its DCP shard rows, which every rank |
| 1395 | # holds. A dummy batch bypasses the PCP manager and has one |
| 1396 | # row per request, so its own extent is the request's. |
| 1397 | req_idx = common_attn_metadata.req_idx |
| 1398 | if req_idx is None: |
| 1399 | req_idx = np.arange(num_reqs) |
| 1400 | shard_rows_cpu = common_attn_metadata.dcp_local_seq_lens_cpu_upper_bound |
| 1401 | if shard_rows_cpu is None: |
| 1402 | shard_rows_cpu = get_dcp_local_seq_lens( |
| 1403 | seq_lens_cpu, |
| 1404 | self.dcp_world_size, |
| 1405 | 0, |
| 1406 | self.cp_kv_cache_interleave_size, |
| 1407 | ) |
| 1408 | shard_rows = shard_rows_cpu.numpy() |
| 1409 | chunk_specs = self._split_pcp_dcp_prefill_chunks( |
| 1410 | req_idx[num_decodes:], |
| 1411 | shard_rows[num_decodes:], |
| 1412 | prefill_query_lens_cpu, |
| 1413 | max_logits_bytes, |
| 1414 | request_offset=num_decodes, |
| 1415 | ) |
| 1416 | else: |
| 1417 | chunk_specs = self._split_indexer_prefill_chunks( |
| 1418 | self._prefill_split_seq_lens(compressed_seq_lens_cpu[num_decodes:]), |
| 1419 | prefill_query_lens_cpu, |
| 1420 | self.max_prefill_buffer_size, |
| 1421 | max_logits_bytes, |
| 1422 | request_offset=num_decodes, |
| 1423 | ) |
| 1424 | |
| 1425 | chunks = [] |
| 1426 | for req_slice, query_slice in chunk_specs: |
| 1427 | pcp_plan = None |
| 1428 | if req_idx is not None: |
| 1429 | assert shard_rows is not None |
| 1430 | pcp_plan = build_pcp_global_chunk_plan( |
| 1431 | req_idx[req_slice], |
| 1432 | shard_rows[req_slice], |
| 1433 | self.dcp_world_size, |
| 1434 | self.device, |
| 1435 | self.cp_kv_cache_interleave_size, |
| 1436 | ) |
| 1437 | metadata = build_prefill_chunk_metadata( |
| 1438 | req_slice.start, |
| 1439 | req_slice.stop, |
| 1440 | query_start_loc, |
| 1441 | query_start_loc_cpu, |
| 1442 | seq_lens, |
| 1443 | compressed_seq_lens, |
| 1444 | compressed_seq_lens_cpu, |
| 1445 | indexer_block_table, |
| 1446 | self.compress_ratio, |
| 1447 | query_slice=query_slice, |
| 1448 | skip_kv_gather=query_slice.start > 0, |
| 1449 | dcp_rank=self.dcp_rank, |
| 1450 | dcp_world_size=self.dcp_world_size, |
| 1451 | cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size, |
| 1452 | pcp_plan=pcp_plan, |
| 1453 | ) |
| 1454 | # Skip when total_seq_lens is 0 (i.e., no compressed token). |
| 1455 | if metadata is not None: |
| 1456 | chunks.append(metadata) |
| 1457 | prefill_metadata = DeepseekV32IndexerPrefillMetadata( |
| 1458 | chunks, |
| 1459 | max_prefill_seq_len=( |
| 1460 | int(seq_lens_cpu[num_decodes:].max().item()) |
| 1461 | if num_prefills > 0 |
| 1462 | else 0 |
| 1463 | ), |
| 1464 | ) |
| 1465 | |
| 1466 | decode_metadata = None |
| 1467 | if num_decodes > 0: |
| 1468 | if not self.supports_varlen: |
| 1469 | torch.diff( |
| 1470 | common_attn_metadata.query_start_loc[: num_decodes + 1], |
| 1471 | out=self.decode_lens_buffer[:num_decodes], |
| 1472 | ) |
| 1473 | self.per_req_decode_lens_buffer[:num_decodes].copy_( |
| 1474 | self.decode_lens_buffer[:num_decodes] |
| 1475 | ) |
| 1476 | decode_lens = self.decode_lens_buffer[:num_decodes] |
| 1477 | decode_lens_cpu = torch.diff( |
| 1478 | common_attn_metadata.query_start_loc_cpu[: num_decodes + 1] |
| 1479 | ) |
| 1480 | |
| 1481 | # Under DCP the per-token decode bounds must be localized AFTER the |
| 1482 | # per-token expansion below, not before. Expanding from a |
| 1483 | # request-level localized length subtracts decode offsets in local |
| 1484 | # space and yields too-short bounds (e.g. world=2, rank=1, global |
| 1485 | # per-token bounds [8, 9, 10] -> [3, 4, 5] instead of [4, 4, 5]), so |
| 1486 | # the first decode token would run top-k against too short a local KV |
| 1487 | # range and miss valid tokens. Keep the global seq_lens here and |
| 1488 | # localize the expanded bounds further down. |
| 1489 | global_seq_lens_for_decode: torch.Tensor | None = None |
| 1490 | if dcp_local_seq_lens is not None: |
| 1491 | global_seq_lens_for_decode = common_attn_metadata.seq_lens[:num_decodes] |
| 1492 | seq_lens = common_attn_metadata.seq_lens[:num_decodes] |
| 1493 | block_table = common_attn_metadata.block_table_tensor[:num_decodes, ...] |
| 1494 | |
| 1495 | max_decode_len = int(decode_lens_cpu.max().item()) |
| 1496 | min_decode_len = int(decode_lens_cpu.min().item()) |
| 1497 | write_is_uniform = min_decode_len == max_decode_len |
| 1498 | next_n = 1 + self.num_speculative_tokens |
| 1499 | # The kernel sees max_decode_len Q rows, not the configured next_n, |
| 1500 | # so legality is per-step: on SM90 a uniformly 3-deep batch has no |
| 1501 | # native kernel. max_decode_len <= 1 always has one. |
| 1502 | step_next_n_ok = max_decode_len <= 1 or _supports_native_decode( |
| 1503 | max_decode_len |
| 1504 | ) |
| 1505 | use_native = ( |
| 1506 | not (self.use_flattening or self.supports_varlen) |
| 1507 | and max_decode_len <= next_n |
| 1508 | and step_next_n_ok |
| 1509 | ) |
| 1510 | |
| 1511 | if not self.supports_varlen: |
| 1512 | global_seq_lens_for_decode = self._prepare_global_decode_seq_lens( |
| 1513 | global_seq_lens=global_seq_lens_for_decode, |
| 1514 | decode_lens=decode_lens, |
| 1515 | decode_lens_cpu=decode_lens_cpu, |
| 1516 | query_start_loc=common_attn_metadata.query_start_loc[:num_decodes], |
| 1517 | num_decode_tokens=num_decode_tokens, |
| 1518 | use_native=use_native, |
| 1519 | max_decode_len=max_decode_len, |
| 1520 | ) |
| 1521 | |
| 1522 | decode_indices = None |
| 1523 | if self.supports_varlen: |
| 1524 | from vllm.v1.attention.ops.metadata import ( |
| 1525 | _indexer_decode_metadata_kernel, |
| 1526 | ) |
| 1527 | |
| 1528 | capacity = self.decode_seq_lens_buffer.numel() |
| 1529 | grid = max( |
| 1530 | num_decodes, |
| 1531 | num_decode_tokens + triton.cdiv(capacity - num_decode_tokens, 256), |
| 1532 | ) |
| 1533 | _indexer_decode_metadata_kernel[(grid,)]( |
| 1534 | query_start_loc, |
| 1535 | seq_lens, |
| 1536 | block_table, |
| 1537 | self.decode_seq_lens_buffer, |
| 1538 | self.expanded_block_table_buffer, |
| 1539 | self.decode_lens_buffer, |
| 1540 | self.decode_indices_buffer, |
| 1541 | self.per_req_decode_lens_buffer, |
| 1542 | num_decodes, |
| 1543 | num_decode_tokens, |
| 1544 | capacity, |
| 1545 | block_table.stride(0), |
| 1546 | self.expanded_block_table_buffer.stride(0), |
| 1547 | BLOCK_COLS=block_table.shape[1], |
| 1548 | num_warps=4, |
| 1549 | ) |
| 1550 | seq_lens = self.decode_seq_lens_buffer[:num_decode_tokens] |
| 1551 | block_table = self.expanded_block_table_buffer[:num_decode_tokens] |
| 1552 | decode_lens = self.decode_lens_buffer[:num_decode_tokens] |
| 1553 | decode_indices = self.decode_indices_buffer[:num_decode_tokens] |
| 1554 | requires_padding = False |
| 1555 | if global_seq_lens_for_decode is not None and max_decode_len > 1: |
| 1556 | self.global_decode_seq_lens_buffer[:num_decode_tokens].copy_( |
| 1557 | seq_lens |
| 1558 | ) |
| 1559 | global_seq_lens_for_decode = self.global_decode_seq_lens_buffer[ |
| 1560 | :num_decode_tokens |
| 1561 | ] |
| 1562 | else: |
| 1563 | seq_lens, block_table, decode_lens, batch_size, requires_padding = ( |
| 1564 | self._prepare_decode_tensors( |
| 1565 | seq_lens=seq_lens, |
| 1566 | block_table=block_table, |
| 1567 | decode_lens=decode_lens, |
| 1568 | decode_lens_cpu=decode_lens_cpu, |
| 1569 | query_start_loc=common_attn_metadata.query_start_loc[ |
| 1570 | :num_decodes |
| 1571 | ], |
| 1572 | num_decodes=num_decodes, |
| 1573 | num_decode_tokens=num_decode_tokens, |
| 1574 | use_native=use_native, |
| 1575 | next_n=next_n, |
| 1576 | max_decode_len=max_decode_len, |
| 1577 | ) |
| 1578 | ) |
| 1579 | |
| 1580 | if self.compress_ratio > 1: |
| 1581 | kernel_block_size = self.kernel_block_size |
| 1582 | if ( |
| 1583 | kernel_block_size is not None |
| 1584 | and self.kv_cache_spec.block_size != kernel_block_size |
| 1585 | and self.kv_cache_spec.block_size % kernel_block_size == 0 |
| 1586 | ): |
| 1587 | factor = self.kv_cache_spec.block_size // kernel_block_size |
| 1588 | compressed = block_table[:, ::factor] // factor |
| 1589 | rows, cols = compressed.shape |
| 1590 | if self.indexer_decode_block_table_buffer is None: |
| 1591 | self.indexer_decode_block_table_buffer = torch.zeros( |
| 1592 | (self._max_num_batched_tokens, cols), |
| 1593 | dtype=torch.int32, |
| 1594 | device=self.device, |
| 1595 | ) |
| 1596 | self.indexer_decode_block_table_buffer[:rows, :cols].copy_( |
| 1597 | compressed |
| 1598 | ) |
| 1599 | block_table = self.indexer_decode_block_table_buffer[:rows, :cols] |
| 1600 | |
| 1601 | # Flattening always returns a buffer view, including single-token |
| 1602 | # batches. Keep its address stable across varlen graph replays. |
| 1603 | seq_lens_is_buffer_view = not use_native or next_n > 1 |
| 1604 | |
| 1605 | # DCP: localize the now-expanded per-token global bounds to this |
| 1606 | # rank's owned KV. Done here (after expansion) so each token's global |
| 1607 | # causal length is localized individually; see the comment above. |
| 1608 | if dcp_local_seq_lens is not None: |
| 1609 | seq_lens = self._dcp_localize_decode_seq_lens( |
| 1610 | seq_lens, num_decodes, seq_lens_is_buffer_view |
| 1611 | ) |
| 1612 | |
| 1613 | # For DeepseekV4 (compress_ratio > 1), the indexer KV cache stores |
| 1614 | # compressed tokens. Convert uncompressed seq_lens to compressed. |
| 1615 | if self.compress_ratio > 1: |
| 1616 | if seq_lens_is_buffer_view: |
| 1617 | seq_lens //= self.compress_ratio |
| 1618 | else: |
| 1619 | # Copy to avoid mutating shared state; keeps CG address stable. |
| 1620 | self.expanded_seq_lens_buffer[:num_decodes] = ( |
| 1621 | seq_lens // self.compress_ratio |
| 1622 | ) |
| 1623 | self.expanded_seq_lens_buffer[num_decodes:num_decode_tokens] = 0 |
| 1624 | seq_lens = self.expanded_seq_lens_buffer[:num_decode_tokens] |
| 1625 | |
| 1626 | # Non-MTP: deep_gemm paged MQA logits requires 2D context_lens |
| 1627 | # (csrc/apis/attention.hpp). Unsqueeze to (B, 1) so downstream |
| 1628 | # kernels see the same (B, next_n) layout as the MTP path. |
| 1629 | if seq_lens.dim() == 1: |
| 1630 | seq_lens = seq_lens.unsqueeze(-1) |
| 1631 | |
| 1632 | # DeepGEMM is required for the paged MQA logits on CUDA devices |
| 1633 | schedule_metadata = self.scheduler_metadata_buffer |
| 1634 | if current_platform.is_cuda() and has_deep_gemm(): |
| 1635 | metadata = get_paged_mqa_logits_metadata( |
| 1636 | seq_lens, |
| 1637 | self.kv_cache_spec.num_states, |
| 1638 | self.num_sms, |
| 1639 | indices=decode_indices, |
| 1640 | ) |
| 1641 | schedule_metadata = self.scheduler_metadata_buffer[: metadata.shape[0]] |
| 1642 | schedule_metadata[:] = metadata |
| 1643 | |
| 1644 | decode_metadata = DeepSeekV32IndexerDecodeMetadata( |
| 1645 | block_table=block_table, |
| 1646 | seq_lens=seq_lens, |
| 1647 | decode_lens=decode_lens, |
| 1648 | requires_padding=requires_padding, |
| 1649 | schedule_metadata=schedule_metadata, |
| 1650 | indices=decode_indices, |
| 1651 | global_seq_lens=global_seq_lens_for_decode, |
| 1652 | per_req_decode_lens=self.per_req_decode_lens_buffer[:num_decodes], |
| 1653 | decode_is_uniform=write_is_uniform, |
| 1654 | write_max_decode_len=max_decode_len, |
| 1655 | ) |
| 1656 | |
| 1657 | attn_metadata = DeepseekV32IndexerMetadata( |
| 1658 | seq_lens=common_attn_metadata.seq_lens, |
| 1659 | max_seq_len=common_attn_metadata.max_seq_len, |
| 1660 | slot_mapping=compressed_slot_mapping, |
| 1661 | num_decodes=num_decodes, |
| 1662 | num_decode_tokens=num_decode_tokens, |
| 1663 | num_prefills=num_prefills, |
| 1664 | num_prefill_tokens=num_prefill_tokens, |
| 1665 | prefill=prefill_metadata, |
| 1666 | decode=decode_metadata, |
| 1667 | ) |
| 1668 | |
| 1669 | return attn_metadata |
| 1670 | |
| 1671 | |
| 1672 | def build_prefill_chunk_metadata( |
| 1673 | start_idx: int, |
| 1674 | end_idx: int, |
| 1675 | query_start_loc: torch.Tensor, |
| 1676 | query_start_loc_cpu: torch.Tensor, |
| 1677 | uncompressed_seq_lens: torch.Tensor, |
| 1678 | compressed_seq_lens: torch.Tensor, |
| 1679 | compressed_seq_lens_cpu: torch.Tensor, |
| 1680 | block_table: torch.Tensor, |
| 1681 | compress_ratio: int, |
| 1682 | query_slice: slice | None = None, |
| 1683 | skip_kv_gather: bool = False, |
| 1684 | dcp_rank: int = 0, |
| 1685 | dcp_world_size: int = 1, |
| 1686 | cp_kv_cache_interleave_size: int = 1, |
| 1687 | pcp_plan: PCPGlobalChunkPlan | None = None, |
| 1688 | ) -> DeepseekV32IndexerPrefillChunkMetadata | None: |
| 1689 | if pcp_plan is not None: |
| 1690 | total_seq_lens = pcp_plan.total |
| 1691 | else: |
| 1692 | total_seq_lens = compressed_seq_lens_cpu[start_idx:end_idx].sum().item() |
| 1693 | if total_seq_lens == 0: |
| 1694 | return None |
| 1695 | |
| 1696 | num_reqs = end_idx - start_idx |
| 1697 | device = block_table.device |
| 1698 | token_to_seq = torch.empty(total_seq_lens, dtype=torch.int32, device=device) |
| 1699 | |
| 1700 | if pcp_plan is not None: |
| 1701 | cu_seq_lens = pcp_plan.global_cu |
| 1702 | else: |
| 1703 | cu_seq_lens = torch.empty(num_reqs + 1, dtype=torch.int32, device=device) |
| 1704 | # Assigning to slice avoids cpu sync. |
| 1705 | cu_seq_lens[:1] = 0 |
| 1706 | torch.cumsum(compressed_seq_lens[start_idx:end_idx], dim=0, out=cu_seq_lens[1:]) |
| 1707 | |
| 1708 | local_cu_seq_lens = cu_seq_lens |
| 1709 | local_total_seq_lens = total_seq_lens |
| 1710 | max_local_total_seq_lens = total_seq_lens |
| 1711 | if pcp_plan is not None: |
| 1712 | local_cu_seq_lens = pcp_plan.padded_local_cu |
| 1713 | local_total_seq_lens = pcp_plan.padded_local_total |
| 1714 | max_local_total_seq_lens = pcp_plan.padded_local_total |
| 1715 | elif dcp_world_size > 1: |
| 1716 | # Per-rank local KV length under interleave-aware DCP sharding, shape |
| 1717 | # [num_reqs, dcp_world_size]. Reuse the canonical CP helper so the |
| 1718 | # sharding matches the rest of the DCP pipeline (decode/prefill). |
| 1719 | local_seq_lens = get_dcp_local_seq_lens( |
| 1720 | compressed_seq_lens[start_idx:end_idx], |
| 1721 | dcp_world_size, |
| 1722 | None, |
| 1723 | cp_kv_cache_interleave_size, |
| 1724 | ) |
| 1725 | this_rank_counts = local_seq_lens[:, dcp_rank].to(torch.int32) |
| 1726 | local_cu_seq_lens = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device) |
| 1727 | torch.cumsum(this_rank_counts, dim=0, out=local_cu_seq_lens[1:]) |
| 1728 | local_total_seq_lens = int(local_cu_seq_lens[-1].item()) |
| 1729 | max_local_total_seq_lens = int(local_seq_lens.sum(dim=0).max().item()) |
| 1730 | |
| 1731 | query_start_loc = ( |
| 1732 | query_start_loc[start_idx : end_idx + 1] - query_start_loc[start_idx] |
| 1733 | ) |
| 1734 | |
| 1735 | total_query_len = int( |
| 1736 | (query_start_loc_cpu[end_idx] - query_start_loc_cpu[start_idx]).item() |
| 1737 | ) |
| 1738 | if query_slice is not None: |
| 1739 | qs_start = min(query_slice.start, total_query_len) |
| 1740 | qs_stop = min(query_slice.stop, total_query_len) |
| 1741 | else: |
| 1742 | qs_start = 0 |
| 1743 | qs_stop = total_query_len |
| 1744 | output_query_len = qs_stop - qs_start |
| 1745 | assert 0 <= output_query_len <= total_query_len - qs_start, ( |
| 1746 | f"query slice [{qs_start}, {qs_stop}) does not fit this rank's " |
| 1747 | f"{total_query_len} tokens; the metadata tail would stay uninitialized" |
| 1748 | ) |
| 1749 | |
| 1750 | cu_seq_len_ks = torch.empty(output_query_len, dtype=torch.int32, device=device) |
| 1751 | cu_seq_len_ke = torch.empty(output_query_len, dtype=torch.int32, device=device) |
| 1752 | |
| 1753 | if pcp_plan is not None: |
| 1754 | row_start_cu = pcp_plan.row_start_cu |
| 1755 | kernel_dcp_rank, kernel_dcp_world = 0, 1 |
| 1756 | else: |
| 1757 | # Under DCP the kernel writes this rank's local row bounds into |
| 1758 | # cu_seq_len_ks/ke; otherwise local_cu_seq_lens aliases cu_seq_lens. |
| 1759 | row_start_cu = local_cu_seq_lens |
| 1760 | kernel_dcp_rank, kernel_dcp_world = dcp_rank, dcp_world_size |
| 1761 | _BUILD_PREFILL_CHUNK_METADATA_KERNEL( |
| 1762 | query_start_loc, |
| 1763 | uncompressed_seq_lens[start_idx:end_idx], |
| 1764 | cu_seq_lens, |
| 1765 | row_start_cu, |
| 1766 | token_to_seq, |
| 1767 | cu_seq_len_ks, |
| 1768 | cu_seq_len_ke, |
| 1769 | qs_start, |
| 1770 | qs_stop, |
| 1771 | kernel_dcp_rank, |
| 1772 | kernel_dcp_world, |
| 1773 | cp_kv_cache_interleave_size, |
| 1774 | num_reqs=num_reqs, |
| 1775 | COMPRESS_RATIO=compress_ratio, |
| 1776 | ) |
| 1777 | |
| 1778 | pcp_deinterleave_idx = pcp_plan.deinterleave_idx if pcp_plan is not None else None |
| 1779 | |
| 1780 | token_start = query_start_loc_cpu[start_idx].item() |
| 1781 | if query_slice is not None: |
| 1782 | token_end = token_start + qs_stop |
| 1783 | token_start = token_start + qs_start |
| 1784 | skip_kv_gather = skip_kv_gather or qs_start > 0 |
| 1785 | else: |
| 1786 | token_end = query_start_loc_cpu[end_idx].item() |
| 1787 | |
| 1788 | return DeepseekV32IndexerPrefillChunkMetadata( |
| 1789 | pcp_deinterleave_idx=pcp_deinterleave_idx, |
| 1790 | cu_seqlen_ks=cu_seq_len_ks, |
| 1791 | cu_seqlen_ke=cu_seq_len_ke, |
| 1792 | cu_seq_lens=cu_seq_lens, |
| 1793 | token_to_seq=token_to_seq, |
| 1794 | total_seq_lens=total_seq_lens, |
| 1795 | block_table=block_table[start_idx:end_idx], |
| 1796 | token_start=token_start, |
| 1797 | token_end=token_end, |
| 1798 | num_reqs=num_reqs, |
| 1799 | skip_kv_gather=skip_kv_gather, |
| 1800 | local_cu_seq_lens=local_cu_seq_lens, |
| 1801 | local_total_seq_lens=local_total_seq_lens, |
| 1802 | max_local_total_seq_lens=max_local_total_seq_lens, |
| 1803 | ) |
| 1804 | |
| 1805 | |
| 1806 | _PREPARE_UNIFORM_DECODE_KERNEL = PrepareUniformDecodeKernel() |
| 1807 | _BUILD_PREFILL_CHUNK_METADATA_KERNEL = BuildPrefillChunkMetadataKernel() |