| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Custom Sparse Attention Indexer layers.""" |
| 4 | |
| 5 | import torch |
| 6 | import torch.distributed as dist |
| 7 | |
| 8 | import vllm.envs as envs |
| 9 | from vllm import _custom_ops as ops |
| 10 | from vllm._aiter_ops import rocm_aiter_ops |
| 11 | from vllm.compilation.breakable_cudagraph import eager_break_during_capture |
| 12 | from vllm.config import CUDAGraphMode, get_current_vllm_config |
| 13 | from vllm.distributed import get_dcp_group, get_pcp_group |
| 14 | from vllm.forward_context import get_forward_context |
| 15 | from vllm.logger import init_logger |
| 16 | from vllm.model_executor.custom_op import CustomOp |
| 17 | from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( |
| 18 | apply_candidate_mask as _apply_candidate_mask, |
| 19 | ) |
| 20 | from vllm.model_executor.kernels.attention.dsa.candidate_blocks import ( |
| 21 | select_candidate_blocks as _select_candidate_blocks, |
| 22 | ) |
| 23 | from vllm.model_executor.layers.indexer_topk import ( |
| 24 | RADIX_TOPK_WORKSPACE_SIZE, |
| 25 | get_indexer_topk, |
| 26 | ) |
| 27 | from vllm.model_executor.layers.quantization.utils.quant_utils import ( |
| 28 | get_fp8_min_max, |
| 29 | ) |
| 30 | from vllm.platforms import current_platform |
| 31 | from vllm.triton_utils import tl, triton |
| 32 | from vllm.utils.deep_gemm import ( |
| 33 | fp8_fp4_mqa_logits, |
| 34 | fp8_fp4_paged_mqa_logits, |
| 35 | has_deep_gemm, |
| 36 | ) |
| 37 | from vllm.utils.import_utils import has_cutedsl |
| 38 | from vllm.utils.torch_utils import ( |
| 39 | LayerNameType, |
| 40 | _encode_layer_name, |
| 41 | _resolve_layer_name, |
| 42 | direct_register_custom_op, |
| 43 | ) |
| 44 | from vllm.v1.attention.backends.mla.indexer import ( |
| 45 | DeepseekV32IndexerMetadata, |
| 46 | ) |
| 47 | from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton |
| 48 | from vllm.v1.attention.ops.pcp import maybe_gather_indexer_k |
| 49 | from vllm.v1.worker.workspace import current_workspace_manager |
| 50 | |
| 51 | logger = init_logger(__name__) |
| 52 | |
| 53 | # MXFP4 layout: 2 values packed per byte, ue8m0 (1-byte) scale per block of 32. |
| 54 | MXFP4_BLOCK_SIZE = 32 |
| 55 | |
| 56 | |
| 57 | def _assert_cutedsl_dcp_merge_supported( |
| 58 | logits: torch.Tensor, |
| 59 | topk_indices: torch.Tensor, |
| 60 | k: int, |
| 61 | ) -> None: |
| 62 | # The DCP merge only supports the CuteDSL path (Triton pack kernel + CuteDSL |
| 63 | # stable-topk selector); there is no PyTorch fallback. The first cut targets |
| 64 | # Blackwell/Hopper with index_topk in (512, 1024, 2048) (the selector's radix |
| 65 | # sizing); the Triton pack itself has no shape/topk constraints. |
| 66 | if not has_cutedsl(): |
| 67 | raise RuntimeError( |
| 68 | "DCP sparse-indexer merge requires CuteDSL; install it or disable DCP." |
| 69 | ) |
| 70 | if logits.device.type != "cuda": |
| 71 | raise RuntimeError("DCP sparse-indexer merge requires CUDA tensors.") |
| 72 | if logits.dtype != torch.float32 or topk_indices.dtype != torch.int32: |
| 73 | raise RuntimeError( |
| 74 | "DCP sparse-indexer merge requires fp32 logits and int32 indices." |
| 75 | ) |
| 76 | if k not in (512, 1024, 2048): |
| 77 | raise RuntimeError( |
| 78 | f"DCP sparse-indexer merge requires index_topk in (512, 1024, 2048); " |
| 79 | f"got {k}." |
| 80 | ) |
| 81 | |
| 82 | |
| 83 | def _merge_dcp_topk_global( |
| 84 | logits: torch.Tensor, |
| 85 | topk_indices: torch.Tensor, |
| 86 | topk_tokens: int, |
| 87 | dcp_rank: int, |
| 88 | dcp_world_size: int, |
| 89 | cp_interleave: int, |
| 90 | row_starts: torch.Tensor | None = None, |
| 91 | ) -> None: |
| 92 | """Merge each DCP rank's local top-K into the global top-K. |
| 93 | |
| 94 | ``topk_indices`` are this rank's local top-K positions into its 1/N KV |
| 95 | shard. A token in the global top-K must also be in its owning rank's local |
| 96 | top-K (at most ``topk_tokens - 1`` tokens rank globally above it, hence at |
| 97 | most that many on its own rank), so exchanging only the per-rank local |
| 98 | candidates is exact -- equivalent to all-gathering the full logit matrix, |
| 99 | but it ships ``dcp_world_size * topk_tokens`` candidates instead of the whole |
| 100 | score row. Overwrites ``topk_indices`` with global token ids (``-1`` for |
| 101 | padding); the attention backend localizes them back to physical slots per |
| 102 | rank. |
| 103 | """ |
| 104 | if dcp_world_size <= 1: |
| 105 | return |
| 106 | |
| 107 | # CuteDSL-only path (no PyTorch fallback): Triton-pack each rank's |
| 108 | # (score, global_id) candidates on-device, all-gather, then the CuteDSL |
| 109 | # stable-topk selector. |
| 110 | _assert_cutedsl_dcp_merge_supported(logits, topk_indices, topk_tokens) |
| 111 | from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import ( |
| 112 | pack_dcp_topk_candidates_cutedsl, |
| 113 | stable_topk_from_gathered_candidates_cutedsl, |
| 114 | ) |
| 115 | |
| 116 | packed = torch.empty( |
| 117 | (*topk_indices.shape, 2), |
| 118 | dtype=torch.float32, |
| 119 | device=topk_indices.device, |
| 120 | ) |
| 121 | pack_dcp_topk_candidates_cutedsl( |
| 122 | logits, |
| 123 | topk_indices, |
| 124 | packed, |
| 125 | dcp_rank, |
| 126 | dcp_world_size, |
| 127 | cp_interleave, |
| 128 | row_starts, |
| 129 | ) |
| 130 | gathered = get_dcp_group().all_gather(packed, dim=1) |
| 131 | stable_topk_from_gathered_candidates_cutedsl( |
| 132 | gathered, topk_tokens, out=topk_indices |
| 133 | ) |
| 134 | |
| 135 | |
| 136 | def dcp_gather_kv_rows( |
| 137 | local_padded: torch.Tensor, |
| 138 | deinterleave_idx: torch.Tensor, |
| 139 | gathered_buf: torch.Tensor, |
| 140 | out_buf: torch.Tensor, |
| 141 | ) -> torch.Tensor: |
| 142 | """All-gather this rank's KV shard and return it in global token order.""" |
| 143 | gathered = gathered_buf[: get_dcp_group().world_size * local_padded.shape[0]] |
| 144 | dist.all_gather_into_tensor( |
| 145 | gathered, local_padded.contiguous(), group=get_dcp_group().device_group |
| 146 | ) |
| 147 | return torch.index_select( |
| 148 | gathered, 0, deinterleave_idx, out=out_buf[: deinterleave_idx.shape[0]] |
| 149 | ) |
| 150 | |
| 151 | |
| 152 | @triton.jit |
| 153 | def _fused_indexer_q_rope_quant_kernel( |
| 154 | positions, |
| 155 | q, |
| 156 | q_s0, |
| 157 | q_s1, |
| 158 | cos_sin_cache, |
| 159 | cos_sin_s0, |
| 160 | q_fp8, |
| 161 | q_fp8_s0, |
| 162 | q_fp8_s1, |
| 163 | weights, |
| 164 | weights_s0, |
| 165 | weights_s1, |
| 166 | weights_out, |
| 167 | weights_out_s0, |
| 168 | weights_out_s1, |
| 169 | softmax_scale, |
| 170 | head_scale, |
| 171 | fp8_min: tl.constexpr, |
| 172 | fp8_max: tl.constexpr, |
| 173 | is_neox: tl.constexpr, |
| 174 | ): |
| 175 | token = tl.program_id(0) |
| 176 | head = tl.program_id(1) |
| 177 | offs32 = tl.arange(0, 32) |
| 178 | offs64 = tl.arange(0, 64) |
| 179 | |
| 180 | pos = tl.load(positions + token) |
| 181 | cos = tl.load(cos_sin_cache + pos * cos_sin_s0 + offs32).to(tl.float32) |
| 182 | sin = tl.load(cos_sin_cache + pos * cos_sin_s0 + 32 + offs32).to(tl.float32) |
| 183 | q_base = q + token * q_s0 + head * q_s1 |
| 184 | out_base = q_fp8 + token * q_fp8_s0 + head * q_fp8_s1 |
| 185 | |
| 186 | if is_neox: |
| 187 | # NeoX layout, x0 = q[0:32], x1 = q[32:64] |
| 188 | x0 = tl.load(q_base + offs32).to(tl.float32) |
| 189 | x1 = tl.load(q_base + 32 + offs32).to(tl.float32) |
| 190 | else: |
| 191 | # interleaved layout |
| 192 | # x0 = q[0, 2, 4, ...], x1 = q[1, 3, 5, ...] |
| 193 | x0 = tl.load(q_base + offs32 * 2).to(tl.float32) |
| 194 | x1 = tl.load(q_base + offs32 * 2 + 1).to(tl.float32) |
| 195 | r0 = (x0 * cos - x1 * sin).to(tl.bfloat16).to(tl.float32) |
| 196 | r1 = (x1 * cos + x0 * sin).to(tl.bfloat16).to(tl.float32) |
| 197 | amax = tl.maximum(tl.max(tl.abs(r0)), tl.max(tl.abs(r1))) |
| 198 | |
| 199 | q_nope = tl.load(q_base + 64 + offs64).to(tl.float32) |
| 200 | amax = tl.maximum(amax, tl.max(tl.abs(q_nope))) |
| 201 | scale_raw = tl.maximum(amax, 1e-10) * (1.0 / fp8_max) |
| 202 | # e8m0 format |
| 203 | q_scale = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) |
| 204 | |
| 205 | if is_neox: |
| 206 | tl.store( |
| 207 | out_base + offs32, |
| 208 | tl.clamp(r0 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), |
| 209 | ) |
| 210 | tl.store( |
| 211 | out_base + 32 + offs32, |
| 212 | tl.clamp(r1 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), |
| 213 | ) |
| 214 | else: |
| 215 | tl.store( |
| 216 | out_base + offs32 * 2, |
| 217 | tl.clamp(r0 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), |
| 218 | ) |
| 219 | tl.store( |
| 220 | out_base + offs32 * 2 + 1, |
| 221 | tl.clamp(r1 / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), |
| 222 | ) |
| 223 | tl.store( |
| 224 | out_base + 64 + offs64, |
| 225 | tl.clamp(q_nope / q_scale, fp8_min, fp8_max).to(q_fp8.dtype.element_ty), |
| 226 | ) |
| 227 | |
| 228 | weight = tl.load(weights + token * weights_s0 + head * weights_s1).to(tl.float32) |
| 229 | tl.store( |
| 230 | weights_out + token * weights_out_s0 + head * weights_out_s1, |
| 231 | weight * q_scale * softmax_scale * head_scale, |
| 232 | ) |
| 233 | |
| 234 | |
| 235 | def fused_indexer_q_rope_quant( |
| 236 | positions: torch.Tensor, |
| 237 | q: torch.Tensor, |
| 238 | cos_sin_cache: torch.Tensor, |
| 239 | weights: torch.Tensor, |
| 240 | softmax_scale: float, |
| 241 | head_scale: float, |
| 242 | is_neox: bool, |
| 243 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 244 | assert current_platform.is_cuda() |
| 245 | assert q.dtype == torch.bfloat16 |
| 246 | assert q.shape[-1] == 128 |
| 247 | assert cos_sin_cache.shape[-1] == 64 |
| 248 | assert weights.shape == q.shape[:2] |
| 249 | |
| 250 | q_fp8 = torch.empty_like(q, dtype=current_platform.fp8_dtype()) |
| 251 | weights_out = torch.empty_like(weights, dtype=torch.float32) |
| 252 | fp8_min, fp8_max = get_fp8_min_max() |
| 253 | _fused_indexer_q_rope_quant_kernel[(q.shape[0], q.shape[1])]( |
| 254 | positions, |
| 255 | q, |
| 256 | q.stride(0), |
| 257 | q.stride(1), |
| 258 | cos_sin_cache, |
| 259 | cos_sin_cache.stride(0), |
| 260 | q_fp8, |
| 261 | q_fp8.stride(0), |
| 262 | q_fp8.stride(1), |
| 263 | weights, |
| 264 | weights.stride(0), |
| 265 | weights.stride(1), |
| 266 | weights_out, |
| 267 | weights_out.stride(0), |
| 268 | weights_out.stride(1), |
| 269 | softmax_scale, |
| 270 | head_scale, |
| 271 | fp8_min=fp8_min, |
| 272 | fp8_max=fp8_max, |
| 273 | is_neox=is_neox, |
| 274 | num_warps=1, |
| 275 | ) |
| 276 | return q_fp8, weights_out |
| 277 | |
| 278 | |
| 279 | def _gather_workspace_shapes( |
| 280 | total_seq_lens: int, |
| 281 | head_dim: int, |
| 282 | fp8_dtype: torch.dtype, |
| 283 | use_fp4_cache: bool, |
| 284 | ) -> tuple[tuple[tuple[int, int], torch.dtype], tuple[tuple[int, int], torch.dtype]]: |
| 285 | """Return ((values_shape, values_dtype), (scales_shape, scales_dtype)) for |
| 286 | the K-gather workspace. FP8 path: (T, head_dim) fp8 + (T, 4) uint8 fp32 |
| 287 | scales. MXFP4 path: (T, head_dim // 2) uint8 packed mxfp4 + |
| 288 | (T, head_dim // MXFP4_BLOCK_SIZE) uint8 ue8m0 scales.""" |
| 289 | if use_fp4_cache: |
| 290 | return ( |
| 291 | ((total_seq_lens, head_dim // 2), torch.uint8), |
| 292 | ((total_seq_lens, head_dim // MXFP4_BLOCK_SIZE), torch.uint8), |
| 293 | ) |
| 294 | return ( |
| 295 | ((total_seq_lens, head_dim), fp8_dtype), |
| 296 | ((total_seq_lens, 4), torch.uint8), |
| 297 | ) |
| 298 | |
| 299 | |
| 300 | def kv_cache_as_quant_view( |
| 301 | kv_cache: torch.Tensor, |
| 302 | head_dim: int, |
| 303 | use_fp4_cache: bool, |
| 304 | ) -> torch.Tensor: |
| 305 | """4D ``[num_blocks, block_size, 1, head_width]`` view expected by |
| 306 | DeepGEMM, from the 3D indexer kv-cache allocation.""" |
| 307 | if use_fp4_cache: |
| 308 | assert kv_cache.ndim == 3 and kv_cache.dtype == torch.uint8 |
| 309 | num_blocks, block_size, _ = kv_cache.shape |
| 310 | page_bytes = int(kv_cache.stride(0)) |
| 311 | fp4_bytes = head_dim // 2 + head_dim // MXFP4_BLOCK_SIZE |
| 312 | return torch.as_strided( |
| 313 | kv_cache, |
| 314 | size=(num_blocks, block_size, 1, fp4_bytes), |
| 315 | stride=(page_bytes, fp4_bytes, fp4_bytes, 1), |
| 316 | ) |
| 317 | return kv_cache.unsqueeze(-2) |
| 318 | |
| 319 | |
| 320 | @eager_break_during_capture |
| 321 | def sparse_attn_indexer( |
| 322 | hidden_states: torch.Tensor, |
| 323 | k_cache_prefix: LayerNameType, |
| 324 | kv_cache: torch.Tensor, |
| 325 | q_quant: torch.Tensor, |
| 326 | q_scale: torch.Tensor | None, |
| 327 | k: torch.Tensor | None, |
| 328 | weights: torch.Tensor, |
| 329 | quant_block_size: int, |
| 330 | scale_fmt: str | None, |
| 331 | topk_tokens: int, |
| 332 | head_dim: int, |
| 333 | max_model_len: int, |
| 334 | total_seq_lens: int, |
| 335 | topk_indices_buffer: torch.Tensor, |
| 336 | skip_k_cache_insert: bool, |
| 337 | use_pcp: bool, |
| 338 | pcp_shard_decode_requests: bool, |
| 339 | dense_mha_metadata_layer_name: LayerNameType, |
| 340 | use_fp4_cache: bool = False, |
| 341 | dcp_rank: int = 0, |
| 342 | dcp_world_size: int = 1, |
| 343 | cp_kv_cache_interleave_size: int = 1, |
| 344 | skip_topk_buffer_clear: bool = False, |
| 345 | candidate_blocks: torch.Tensor | None = None, |
| 346 | candidate_block_size: int = 0, |
| 347 | candidate_write: bool = False, |
| 348 | topk_backend: str = "auto", |
| 349 | ) -> torch.Tensor: |
| 350 | # careful! this will be None in dummy run |
| 351 | forward_context = get_forward_context() |
| 352 | attn_metadata = forward_context.attn_metadata |
| 353 | fp8_dtype = current_platform.fp8_dtype() |
| 354 | k_cache_prefix = _resolve_layer_name(k_cache_prefix) |
| 355 | |
| 356 | if candidate_blocks is not None: |
| 357 | # Candidate blocks are request-local; the DCP-sharded logits layout |
| 358 | # would need per-rank translation that is not implemented. |
| 359 | assert dcp_world_size == 1, ( |
| 360 | "v4.1 two-level candidate filtering is not supported with DCP." |
| 361 | ) |
| 362 | assert candidate_block_size > 0 |
| 363 | |
| 364 | # assert isinstance(attn_metadata, dict) |
| 365 | if not isinstance(attn_metadata, dict): |
| 366 | # Reserve workspace for indexer during profiling run |
| 367 | values_spec, scales_spec = _gather_workspace_shapes( |
| 368 | total_seq_lens, head_dim, fp8_dtype, use_fp4_cache |
| 369 | ) |
| 370 | profile_specs: list[tuple[tuple[int, ...], torch.dtype]] = [ |
| 371 | values_spec, |
| 372 | scales_spec, |
| 373 | ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), |
| 374 | ] |
| 375 | if use_pcp and dcp_world_size > 1: |
| 376 | # The PCP+DCP path takes an all-gather destination and a |
| 377 | # de-interleaved result. |
| 378 | gather_spec = _gather_workspace_shapes( |
| 379 | total_seq_lens, head_dim, fp8_dtype, use_fp4_cache |
| 380 | ) |
| 381 | profile_specs.extend(gather_spec * 2) |
| 382 | current_workspace_manager().get_simultaneous(*profile_specs) |
| 383 | |
| 384 | # Dummy allocation to simulate for peak logits tensor memory during inference. |
| 385 | # FP8 elements so elements == bytes |
| 386 | max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 |
| 387 | _ = torch.empty( |
| 388 | max_logits_elems, dtype=torch.uint8, device=hidden_states.device |
| 389 | ) |
| 390 | |
| 391 | return sparse_attn_indexer_fake( |
| 392 | hidden_states, |
| 393 | k_cache_prefix, |
| 394 | kv_cache, |
| 395 | q_quant, |
| 396 | q_scale, |
| 397 | k, |
| 398 | weights, |
| 399 | quant_block_size, |
| 400 | scale_fmt, |
| 401 | topk_tokens, |
| 402 | head_dim, |
| 403 | max_model_len, |
| 404 | total_seq_lens, |
| 405 | topk_indices_buffer, |
| 406 | skip_k_cache_insert, |
| 407 | use_pcp, |
| 408 | pcp_shard_decode_requests, |
| 409 | dense_mha_metadata_layer_name, |
| 410 | use_fp4_cache, |
| 411 | candidate_blocks=candidate_blocks, |
| 412 | candidate_block_size=candidate_block_size, |
| 413 | candidate_write=candidate_write, |
| 414 | ) |
| 415 | attn_metadata_narrowed = attn_metadata[k_cache_prefix] |
| 416 | assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata) |
| 417 | slot_mapping = attn_metadata_narrowed.slot_mapping |
| 418 | has_decode = attn_metadata_narrowed.num_decodes > 0 |
| 419 | has_prefill = attn_metadata_narrowed.num_prefills > 0 |
| 420 | num_decode_tokens = attn_metadata_narrowed.num_decode_tokens |
| 421 | |
| 422 | # q_scale is required iff the FP4 cache path is enabled; the FP8 path |
| 423 | # folds the Q scale into `weights` inside fused_indexer_q_rope_quant. |
| 424 | if use_fp4_cache: |
| 425 | assert q_scale is not None, "use_fp4_cache=True requires q_scale" |
| 426 | else: |
| 427 | assert q_scale is None, "q_scale must be None when use_fp4_cache=False" |
| 428 | |
| 429 | # During speculative decoding, k may be padded to the CUDA graph batch |
| 430 | # size while slot_mapping only covers actual tokens. Truncate k to avoid |
| 431 | # out-of-bounds reads in the kernel. |
| 432 | # Keep PCP padding so every rank contributes the same all-gather shape. |
| 433 | num_tokens = slot_mapping.shape[0] |
| 434 | if k is not None: |
| 435 | # MTP draft decodes are replicated and use an unexpanded slot mapping. |
| 436 | if use_pcp and num_tokens > k.shape[0]: |
| 437 | num_tokens //= get_pcp_group().world_size |
| 438 | k = k[:num_tokens] |
| 439 | |
| 440 | if not skip_k_cache_insert: |
| 441 | assert k is not None |
| 442 | k, slot_mapping_for_cache = maybe_gather_indexer_k( |
| 443 | k, |
| 444 | slot_mapping, |
| 445 | num_decode_tokens, |
| 446 | use_pcp, |
| 447 | pcp_shard_decode_requests=pcp_shard_decode_requests, |
| 448 | ) |
| 449 | # scale_fmt can be None, but the function expects str |
| 450 | assert scale_fmt is not None |
| 451 | assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet" |
| 452 | ops.indexer_k_quant_and_cache( |
| 453 | k, |
| 454 | kv_cache, |
| 455 | slot_mapping_for_cache, |
| 456 | quant_block_size, |
| 457 | scale_fmt, |
| 458 | ) |
| 459 | # The indexer and main MLA may classify the same short extend differently |
| 460 | # because they use independent decode thresholds. Only the main MLA route |
| 461 | # can determine whether the top-k indices will be consumed. |
| 462 | if forward_context.cudagraph_runtime_mode != CUDAGraphMode.FULL: |
| 463 | dense_mha_layer = _resolve_layer_name(dense_mha_metadata_layer_name) |
| 464 | if dense_mha_layer: |
| 465 | mla_metadata = attn_metadata.get(dense_mha_layer) |
| 466 | prefill_metadata = getattr(mla_metadata, "prefill", None) |
| 467 | if ( |
| 468 | getattr(prefill_metadata, "use_dense_mha", False) |
| 469 | and getattr(mla_metadata, "num_decode_tokens", -1) == 0 |
| 470 | and not torch.cuda.is_current_stream_capturing() |
| 471 | ): |
| 472 | # Deliberately leave the buffer untouched. Dense MHA does not |
| 473 | # consume top-k indices for this batch; clearing it would be |
| 474 | # unnecessary work. |
| 475 | return topk_indices_buffer |
| 476 | |
| 477 | # The buffer must be pre-filled with -1 (the "no token" sentinel) before the |
| 478 | # top-k kernels scatter valid indices into it. On the fused deepseek_v32 |
| 479 | # nvidia path, _fused_norm_rope_kernel already cleared the same |
| 480 | # [:num_tokens, :topk] region earlier in this forward, so skip the redundant |
| 481 | # fill. |
| 482 | if not skip_topk_buffer_clear: |
| 483 | topk_indices_buffer[: hidden_states.shape[0]] = -1 |
| 484 | if has_prefill: |
| 485 | prefill_metadata = attn_metadata_narrowed.prefill |
| 486 | assert prefill_metadata is not None |
| 487 | |
| 488 | # Get the full shared workspace buffers once (will allocate on first use). |
| 489 | # Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and |
| 490 | # MXFP4 (head_dim/2 bytes packed + head_dim/MXFP4_BLOCK_SIZE ue8m0 |
| 491 | # scales) based on use_fp4_cache. |
| 492 | workspace_manager = current_workspace_manager() |
| 493 | values_spec, scales_spec = _gather_workspace_shapes( |
| 494 | total_seq_lens, head_dim, fp8_dtype, use_fp4_cache |
| 495 | ) |
| 496 | # PCP + DCP needs two more pairs: the rank-major all-gather destination |
| 497 | # and the de-interleaved result. |
| 498 | pcp_chunks = [ |
| 499 | (c, c.pcp_deinterleave_idx) |
| 500 | for c in prefill_metadata.chunks |
| 501 | if c.pcp_deinterleave_idx is not None |
| 502 | ] |
| 503 | gather_specs: list[tuple[tuple[int, int], torch.dtype]] = [] |
| 504 | if pcp_chunks: |
| 505 | gathered_rows = max( |
| 506 | dcp_world_size * c.local_total_seq_lens for c, _ in pcp_chunks |
| 507 | ) |
| 508 | deinterleaved_rows = max(idx.shape[0] for _, idx in pcp_chunks) |
| 509 | for rows in (gathered_rows, deinterleaved_rows): |
| 510 | gather_specs.extend( |
| 511 | _gather_workspace_shapes(rows, head_dim, fp8_dtype, use_fp4_cache) |
| 512 | ) |
| 513 | k_quant_full, k_scale_full, *gather_bufs = workspace_manager.get_simultaneous( |
| 514 | values_spec, |
| 515 | scales_spec, |
| 516 | *gather_specs, |
| 517 | ) |
| 518 | for chunk in prefill_metadata.chunks: |
| 519 | cu_seqlen_ks = chunk.cu_seqlen_ks |
| 520 | cu_seqlen_ke = chunk.cu_seqlen_ke |
| 521 | assert chunk.local_cu_seq_lens is not None |
| 522 | k_quant = k_quant_full[: chunk.max_local_total_seq_lens] |
| 523 | k_scale = k_scale_full[: chunk.max_local_total_seq_lens] |
| 524 | if not chunk.skip_kv_gather and chunk.local_total_seq_lens > 0: |
| 525 | ops.cp_gather_indexer_k_quant_cache( |
| 526 | kv_cache, |
| 527 | k_quant, |
| 528 | k_scale, |
| 529 | chunk.block_table, |
| 530 | chunk.local_cu_seq_lens, |
| 531 | ) |
| 532 | |
| 533 | # PCP + DCP KV all-gather. |
| 534 | deinterleave_idx = chunk.pcp_deinterleave_idx |
| 535 | if deinterleave_idx is not None: |
| 536 | # local_total_seq_lens is the PCP-padded extent here, so every |
| 537 | # rank contributes an identically shaped shard. |
| 538 | gathered_values, gathered_scales, out_values, out_scales = gather_bufs |
| 539 | shard = k_quant[: chunk.local_total_seq_lens] |
| 540 | k_quant = dcp_gather_kv_rows( |
| 541 | shard, deinterleave_idx, gathered_values, out_values |
| 542 | ) |
| 543 | k_scale = dcp_gather_kv_rows( |
| 544 | k_scale[: chunk.local_total_seq_lens], |
| 545 | deinterleave_idx, |
| 546 | gathered_scales, |
| 547 | out_scales, |
| 548 | ) |
| 549 | |
| 550 | q_slice = q_quant[chunk.token_start : chunk.token_end] |
| 551 | q_scale_slice = ( |
| 552 | q_scale[chunk.token_start : chunk.token_end] |
| 553 | if q_scale is not None |
| 554 | else None |
| 555 | ) |
| 556 | topk_indices = topk_indices_buffer[ |
| 557 | chunk.token_start : chunk.token_end, :topk_tokens |
| 558 | ] |
| 559 | |
| 560 | if chunk.local_total_seq_lens == 0 or q_slice.shape[0] == 0: |
| 561 | logits = q_slice.new_empty((q_slice.shape[0], 0), dtype=torch.float32) |
| 562 | topk_indices.fill_(-1) |
| 563 | else: |
| 564 | # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8 |
| 565 | # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf. |
| 566 | if use_fp4_cache: |
| 567 | q_slice_cast = q_slice.view(torch.int8) |
| 568 | k_quant_cast = k_quant.view(torch.int8) |
| 569 | k_scale_cast = k_scale.view(torch.int32).squeeze(-1) |
| 570 | else: |
| 571 | q_slice_cast = q_slice |
| 572 | k_quant_cast = k_quant |
| 573 | k_scale_cast = k_scale.view(torch.float32).squeeze(-1) |
| 574 | if current_platform.is_xpu(): |
| 575 | if q_scale_slice is not None: |
| 576 | raise RuntimeError("XPU fp8_mqa_logits does not support FP4 Q") |
| 577 | logits = torch.ops.vllm.xpu_fp8_mqa_logits( |
| 578 | q_slice_cast, |
| 579 | k_quant_cast, |
| 580 | k_scale_cast, |
| 581 | weights[chunk.token_start : chunk.token_end], |
| 582 | cu_seqlen_ks, |
| 583 | cu_seqlen_ke, |
| 584 | ) |
| 585 | else: |
| 586 | logits = fp8_fp4_mqa_logits( |
| 587 | (q_slice_cast, q_scale_slice), |
| 588 | (k_quant_cast, k_scale_cast), |
| 589 | weights[chunk.token_start : chunk.token_end], |
| 590 | cu_seqlen_ks, |
| 591 | cu_seqlen_ke, |
| 592 | clean_logits=False, |
| 593 | ) |
| 594 | num_rows = logits.shape[0] |
| 595 | if candidate_blocks is not None: |
| 596 | # Two-level selection (v4.1): the candidate source |
| 597 | # publishes its top blocks; later indexers mask their |
| 598 | # scores to them. Both before the row top-k. |
| 599 | chunk_candidates = candidate_blocks[ |
| 600 | chunk.token_start : chunk.token_end |
| 601 | ] |
| 602 | if candidate_write: |
| 603 | _select_candidate_blocks( |
| 604 | logits, |
| 605 | cu_seqlen_ks, |
| 606 | cu_seqlen_ke, |
| 607 | chunk_candidates.shape[1], |
| 608 | candidate_block_size, |
| 609 | chunk_candidates, |
| 610 | ) |
| 611 | else: |
| 612 | _apply_candidate_mask( |
| 613 | logits, |
| 614 | cu_seqlen_ks, |
| 615 | cu_seqlen_ke, |
| 616 | chunk_candidates, |
| 617 | candidate_block_size, |
| 618 | ) |
| 619 | ops.top_k_per_row_prefill( |
| 620 | logits, |
| 621 | cu_seqlen_ks, |
| 622 | cu_seqlen_ke, |
| 623 | topk_indices, |
| 624 | num_rows, |
| 625 | logits.stride(0), |
| 626 | logits.stride(1), |
| 627 | topk_tokens, |
| 628 | ) |
| 629 | |
| 630 | if deinterleave_idx is None: |
| 631 | # Under the PCP path the top-k already ran over the whole |
| 632 | # context. |
| 633 | _merge_dcp_topk_global( |
| 634 | logits, |
| 635 | topk_indices, |
| 636 | topk_tokens, |
| 637 | dcp_rank, |
| 638 | dcp_world_size, |
| 639 | cp_kv_cache_interleave_size, |
| 640 | row_starts=chunk.cu_seqlen_ks, |
| 641 | ) |
| 642 | |
| 643 | if has_decode: |
| 644 | decode_metadata = attn_metadata_narrowed.decode |
| 645 | assert decode_metadata is not None |
| 646 | kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache) |
| 647 | decode_lens = decode_metadata.decode_lens |
| 648 | # requires_padding can be computed False for ragged warmup/mixed |
| 649 | # batches (seen on SM120 TP=2: 8 tokens over 6 seqs -> uniform |
| 650 | # reshape below would crash). The token-count/batch-size divisibility |
| 651 | # check is a necessary-but-not-sufficient proxy for "uniform batch" |
| 652 | # (e.g. decode_lens=[1, 2, 3] divides evenly by 3 but isn't uniform); |
| 653 | # it catches the specific case seen in practice without requiring a |
| 654 | # device sync to inspect decode_lens directly, which would break |
| 655 | # CUDA-graph-capture safety on this path. Shared by both the pack |
| 656 | # branch below and the matching unpack branch further down — they |
| 657 | # must agree on which path was taken. |
| 658 | needs_padded_path = decode_metadata.requires_padding or ( |
| 659 | num_decode_tokens % decode_lens.shape[0] != 0 |
| 660 | ) |
| 661 | if num_decode_tokens == 0: |
| 662 | padded_q_quant_decode_tokens = q_quant[:1].reshape(1, 1, *q_quant.shape[1:]) |
| 663 | padded_q_scale = ( |
| 664 | q_scale[:1].reshape(1, 1, *q_scale.shape[1:]) |
| 665 | if q_scale is not None |
| 666 | else None |
| 667 | ) |
| 668 | elif needs_padded_path: |
| 669 | # pad in edge case where we have short chunked prefill length < |
| 670 | # decode_threshold since we unstrictly split |
| 671 | # prefill and decode by decode_threshold |
| 672 | # (currently set to 1 + speculative tokens). |
| 673 | # FP8 Q is float8_e4m3fn (pack_seq_triton's fp32 pad path is OK — |
| 674 | # downstream context_lens masks stale slots). MXFP4 Q is two |
| 675 | # uint8 tensors (values + ue8m0 scales) — use the dedicated uint8 |
| 676 | # packer with pad_byte=0 so padded slots dequantize to 0 and |
| 677 | # can't produce NaN/Inf in the logits kernel. |
| 678 | if q_scale is not None: |
| 679 | padded_q_quant_decode_tokens = pack_seq_triton( |
| 680 | q_quant[:num_decode_tokens], decode_lens, pad_value=0 |
| 681 | ) |
| 682 | padded_q_scale = pack_seq_triton( |
| 683 | q_scale[:num_decode_tokens], decode_lens, pad_value=0 |
| 684 | ) |
| 685 | else: |
| 686 | padded_q_quant_decode_tokens = pack_seq_triton( |
| 687 | q_quant[:num_decode_tokens], decode_lens |
| 688 | ) |
| 689 | padded_q_scale = None |
| 690 | else: |
| 691 | padded_q_quant_decode_tokens = q_quant[:num_decode_tokens].reshape( |
| 692 | decode_lens.shape[0], -1, *q_quant.shape[1:] |
| 693 | ) |
| 694 | if q_scale is not None: |
| 695 | padded_q_scale = q_scale[:num_decode_tokens].reshape( |
| 696 | decode_lens.shape[0], -1, *q_scale.shape[1:] |
| 697 | ) |
| 698 | else: |
| 699 | padded_q_scale = None |
| 700 | # TODO: move and optimize below logic with triton kernels |
| 701 | batch_size = padded_q_quant_decode_tokens.shape[0] |
| 702 | next_n = padded_q_quant_decode_tokens.shape[1] |
| 703 | num_padded_tokens = batch_size * next_n |
| 704 | seq_lens = decode_metadata.seq_lens[:batch_size] |
| 705 | # seq_lens is always 2D: (B, next_n) for native spec decode, (B, 1) |
| 706 | # otherwise. deep_gemm fp8_fp4_paged_mqa_logits requires 2D context_lens; |
| 707 | # the downstream topk kernels accept both 1D and 2D. |
| 708 | padded_q_quant_cast = ( |
| 709 | padded_q_quant_decode_tokens.view(torch.int8) |
| 710 | if use_fp4_cache |
| 711 | else padded_q_quant_decode_tokens |
| 712 | ) |
| 713 | if current_platform.is_xpu(): |
| 714 | if padded_q_scale is not None: |
| 715 | raise RuntimeError("XPU fp8_paged_mqa_logits does not support FP4 Q") |
| 716 | seq_lens_xpu = ( |
| 717 | seq_lens[:, -1].contiguous() if seq_lens.ndim == 2 else seq_lens |
| 718 | ) |
| 719 | logits = torch.ops.vllm.xpu_fp8_paged_mqa_logits( |
| 720 | padded_q_quant_cast, |
| 721 | kv_cache, |
| 722 | weights[:num_padded_tokens], |
| 723 | seq_lens_xpu, |
| 724 | decode_metadata.block_table, |
| 725 | decode_metadata.schedule_metadata, |
| 726 | max_model_len, |
| 727 | ) |
| 728 | else: |
| 729 | logits = fp8_fp4_paged_mqa_logits( |
| 730 | (padded_q_quant_cast, padded_q_scale), |
| 731 | kv_cache, |
| 732 | weights[:num_padded_tokens], |
| 733 | seq_lens, |
| 734 | decode_metadata.block_table, |
| 735 | decode_metadata.schedule_metadata, |
| 736 | max_model_len=max_model_len, |
| 737 | clean_logits=False, |
| 738 | indices=decode_metadata.indices, |
| 739 | ) |
| 740 | num_rows = logits.shape[0] |
| 741 | if candidate_blocks is not None: |
| 742 | # Two-level selection (v4.1) on the decode logits; columns are |
| 743 | # request-local compressed positions. seq_lens is (B, next_n) |
| 744 | # for native spec decode (per-row effective lens) and (B, 1) |
| 745 | # otherwise. |
| 746 | vis = seq_lens.reshape(-1) |
| 747 | row_repeat = next_n if vis.numel() != num_rows else 1 |
| 748 | vis = vis[:num_rows] |
| 749 | decode_candidates = candidate_blocks[:num_rows] |
| 750 | if candidate_write: |
| 751 | _select_candidate_blocks( |
| 752 | logits, |
| 753 | None, |
| 754 | vis, |
| 755 | decode_candidates.shape[1], |
| 756 | candidate_block_size, |
| 757 | decode_candidates, |
| 758 | row_repeat, |
| 759 | ) |
| 760 | else: |
| 761 | _apply_candidate_mask( |
| 762 | logits, |
| 763 | None, |
| 764 | vis, |
| 765 | decode_candidates, |
| 766 | candidate_block_size, |
| 767 | row_repeat, |
| 768 | ) |
| 769 | topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens] |
| 770 | |
| 771 | # The backend comes from the layer (config is only readable at model |
| 772 | # construction); dispatchers are cached per backend. |
| 773 | get_indexer_topk(topk_backend)( |
| 774 | logits, |
| 775 | seq_lens, |
| 776 | next_n, |
| 777 | topk_indices, |
| 778 | topk_tokens, |
| 779 | attn_metadata_narrowed.max_seq_len, |
| 780 | ) |
| 781 | |
| 782 | if decode_metadata.global_seq_lens is not None: |
| 783 | _merge_dcp_topk_global( |
| 784 | logits, |
| 785 | topk_indices, |
| 786 | topk_tokens, |
| 787 | dcp_rank, |
| 788 | dcp_world_size, |
| 789 | cp_kv_cache_interleave_size, |
| 790 | ) |
| 791 | |
| 792 | if needs_padded_path: |
| 793 | # if padded, we need to unpack |
| 794 | # the topk indices removing padded tokens |
| 795 | topk_indices = unpack_seq_triton( |
| 796 | topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]), |
| 797 | decode_lens, |
| 798 | ) |
| 799 | topk_indices_buffer[: topk_indices.shape[0], : topk_indices.shape[-1]] = ( |
| 800 | topk_indices |
| 801 | ) |
| 802 | |
| 803 | return topk_indices_buffer |
| 804 | |
| 805 | |
| 806 | def sparse_attn_indexer_fake( |
| 807 | hidden_states: torch.Tensor, |
| 808 | k_cache_prefix: LayerNameType, |
| 809 | kv_cache: torch.Tensor, |
| 810 | q_quant: torch.Tensor, |
| 811 | q_scale: torch.Tensor | None, |
| 812 | k: torch.Tensor | None, |
| 813 | weights: torch.Tensor, |
| 814 | quant_block_size: int, |
| 815 | scale_fmt: str | None, |
| 816 | topk_tokens: int, |
| 817 | head_dim: int, |
| 818 | max_model_len: int, |
| 819 | total_seq_lens: int, |
| 820 | topk_indices_buffer: torch.Tensor | None, |
| 821 | skip_k_cache_insert: bool, |
| 822 | use_pcp: bool, |
| 823 | pcp_shard_decode_requests: bool, |
| 824 | dense_mha_metadata_layer_name: LayerNameType, |
| 825 | use_fp4_cache: bool = False, |
| 826 | dcp_rank: int = 0, |
| 827 | dcp_world_size: int = 1, |
| 828 | cp_kv_cache_interleave_size: int = 1, |
| 829 | skip_topk_buffer_clear: bool = False, |
| 830 | candidate_blocks: torch.Tensor | None = None, |
| 831 | candidate_block_size: int = 0, |
| 832 | candidate_write: bool = False, |
| 833 | topk_backend: str = "auto", |
| 834 | ) -> torch.Tensor: |
| 835 | return topk_indices_buffer |
| 836 | |
| 837 | |
| 838 | direct_register_custom_op( |
| 839 | op_name="sparse_attn_indexer", |
| 840 | op_func=sparse_attn_indexer, |
| 841 | mutates_args=["topk_indices_buffer", "candidate_blocks"], |
| 842 | fake_impl=sparse_attn_indexer_fake, |
| 843 | dispatch_key=current_platform.dispatch_key, |
| 844 | ) |
| 845 | |
| 846 | |
| 847 | @CustomOp.register("sparse_attn_indexer") |
| 848 | class SparseAttnIndexer(CustomOp): |
| 849 | """Sparse Attention Indexer Custom Op Layer. This layer is extracted as a |
| 850 | separate custom op since it involves heavy custom kernels like `mqa_logits`, |
| 851 | `paged_mqa_logits` and `top_k_per_row`, etc. Those kernels maybe requires |
| 852 | specific memory layout or implementation for different hardware backends to |
| 853 | achieve optimal performance. |
| 854 | |
| 855 | For now, the default native path will use CUDA backend path. Other platform |
| 856 | may requires add the corresponding Custom Op name `sparse_attn_indexer` to |
| 857 | `custom_ops` in `CompilationConfig` to enable the platform specific path. |
| 858 | """ |
| 859 | |
| 860 | def __init__( |
| 861 | self, |
| 862 | k_cache, |
| 863 | quant_block_size: int, |
| 864 | scale_fmt: str, |
| 865 | topk_tokens: int, |
| 866 | head_dim: int, |
| 867 | max_model_len: int, |
| 868 | max_total_seq_len: int, |
| 869 | topk_indices_buffer: torch.Tensor, |
| 870 | skip_k_cache_insert: bool = False, |
| 871 | use_fp4_cache: bool = False, |
| 872 | compress_ratio: int = 1, |
| 873 | candidate_blocks: torch.Tensor | None = None, |
| 874 | candidate_block_size: int = 0, |
| 875 | candidate_write: bool = False, |
| 876 | ): |
| 877 | super().__init__() |
| 878 | self.k_cache = k_cache |
| 879 | self.quant_block_size = quant_block_size |
| 880 | self.scale_fmt = scale_fmt |
| 881 | self.topk_tokens = topk_tokens |
| 882 | self.head_dim = head_dim |
| 883 | self.max_model_len = max_model_len |
| 884 | self.max_total_seq_len = max_total_seq_len |
| 885 | self.topk_indices_buffer = topk_indices_buffer |
| 886 | self.skip_k_cache_insert = skip_k_cache_insert |
| 887 | self.use_fp4_cache = use_fp4_cache |
| 888 | self.compress_ratio = compress_ratio |
| 889 | # v4.1 two-level selection: the candidate source indexer writes the |
| 890 | # top candidate blocks here; later indexers mask their scores with it. |
| 891 | self.candidate_blocks = candidate_blocks |
| 892 | self.candidate_block_size = candidate_block_size |
| 893 | self.candidate_write = candidate_write |
| 894 | self.dense_mha_metadata_layer_name = "" |
| 895 | # DCP scalars are constant for the run; resolve them here (config is set |
| 896 | # during model construction) and pass them into the custom op, rather |
| 897 | # than threading them through per-step metadata. |
| 898 | vllm_config = get_current_vllm_config() |
| 899 | parallel_config = vllm_config.parallel_config |
| 900 | self.topk_backend = vllm_config.kernel_config.sparse_indexer_topk_backend |
| 901 | self._parallel_config = parallel_config |
| 902 | self.dcp_world_size = parallel_config.decode_context_parallel_size |
| 903 | self.dcp_rank = get_dcp_group().rank_in_group if self.dcp_world_size > 1 else 0 |
| 904 | self.use_pcp = parallel_config.prefill_context_parallel_size > 1 |
| 905 | self.pcp_shard_decode_requests = parallel_config.pcp_shard_decode_requests |
| 906 | self._cp_kv_cache_interleave_size: int | None = None |
| 907 | if current_platform.is_cuda() and not has_deep_gemm(): |
| 908 | raise RuntimeError( |
| 909 | "Sparse Attention Indexer CUDA op requires DeepGEMM support in " |
| 910 | "the current vLLM environment." |
| 911 | ) |
| 912 | |
| 913 | if vllm_config.kernel_config.enable_jit_warmup: |
| 914 | from vllm.v1.attention.ops.common import ( |
| 915 | _PACK_SEQ_TRITON_KERNEL, |
| 916 | _UNPACK_SEQ_TRITON_KERNEL, |
| 917 | ) |
| 918 | |
| 919 | pack_dtype = torch.uint8 if use_fp4_cache else current_platform.fp8_dtype() |
| 920 | _PACK_SEQ_TRITON_KERNEL.register_warmup( |
| 921 | dtype=pack_dtype, |
| 922 | pad_value=0 if use_fp4_cache else -float("inf"), |
| 923 | ) |
| 924 | _UNPACK_SEQ_TRITON_KERNEL.register_warmup() |
| 925 | |
| 926 | if self.dcp_world_size > 1 and current_platform.is_cuda() and has_cutedsl(): |
| 927 | from vllm.model_executor.kernels.attention.dsa.dcp_indexer_cutedsl import ( # noqa: E501 |
| 928 | _PACK_DCP_TOPK_CANDIDATES_KERNEL, |
| 929 | _STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL, |
| 930 | ) |
| 931 | |
| 932 | _PACK_DCP_TOPK_CANDIDATES_KERNEL.register_warmup() |
| 933 | _STABLE_TOPK_FROM_GATHERED_CANDIDATES_KERNEL.register_warmup() |
| 934 | |
| 935 | @property |
| 936 | def cp_kv_cache_interleave_size(self) -> int: |
| 937 | """With PD+DCP, the real value isn't known until block_size is finalized, |
| 938 | which happens after this layer is built. Safe to cache after the first access, |
| 939 | as long as the adjustment always runs before any forward pass |
| 940 | (it's set up in Worker.initialize_from_config, ahead of warmup/serving). |
| 941 | """ |
| 942 | if self._cp_kv_cache_interleave_size is None: |
| 943 | value = self._parallel_config.cp_kv_cache_interleave_size |
| 944 | if isinstance(get_forward_context().attn_metadata, dict): |
| 945 | self._cp_kv_cache_interleave_size = value |
| 946 | return value |
| 947 | return self._cp_kv_cache_interleave_size |
| 948 | |
| 949 | def forward_native( |
| 950 | self, |
| 951 | hidden_states: torch.Tensor, |
| 952 | q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 953 | k: torch.Tensor | None, |
| 954 | weights: torch.Tensor, |
| 955 | ): |
| 956 | if current_platform.is_cuda() or current_platform.is_xpu(): |
| 957 | return self.forward_cuda(hidden_states, q_quant, k, weights) |
| 958 | elif current_platform.is_rocm(): |
| 959 | if self.use_pcp and self.dcp_world_size > 1: |
| 960 | raise NotImplementedError( |
| 961 | "The ROCm sparse-indexer path does not support PCP+DCP." |
| 962 | ) |
| 963 | return self.forward_hip(hidden_states, q_quant, k, weights) |
| 964 | elif current_platform.is_cpu(): |
| 965 | return self.forward_cpu(hidden_states, q_quant, k, weights) |
| 966 | else: |
| 967 | raise NotImplementedError( |
| 968 | "SparseAttnIndexer native forward is only implemented for " |
| 969 | "CUDA, ROCm, XPU and CPU platforms." |
| 970 | ) |
| 971 | |
| 972 | def forward_cuda( |
| 973 | self, |
| 974 | hidden_states: torch.Tensor, |
| 975 | q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 976 | k: torch.Tensor | None, |
| 977 | weights: torch.Tensor, |
| 978 | ): |
| 979 | # FP8 path: single tensor (per-token scale is folded into `weights`). |
| 980 | # FP4 path: (values, scales) tuple with scales required by the kernel. |
| 981 | if isinstance(q_quant, tuple): |
| 982 | q_values, q_scale = q_quant |
| 983 | else: |
| 984 | q_values, q_scale = q_quant, None |
| 985 | return torch.ops.vllm.sparse_attn_indexer( |
| 986 | hidden_states, |
| 987 | _encode_layer_name(self.k_cache.prefix), |
| 988 | self.k_cache.kv_cache, |
| 989 | q_values, |
| 990 | q_scale, |
| 991 | k, |
| 992 | weights, |
| 993 | self.quant_block_size, |
| 994 | self.scale_fmt, |
| 995 | self.topk_tokens, |
| 996 | self.head_dim, |
| 997 | self.max_model_len, |
| 998 | self.max_total_seq_len, |
| 999 | self.topk_indices_buffer, |
| 1000 | self.skip_k_cache_insert, |
| 1001 | self.use_pcp, |
| 1002 | self.pcp_shard_decode_requests, |
| 1003 | _encode_layer_name(self.dense_mha_metadata_layer_name), |
| 1004 | self.use_fp4_cache, |
| 1005 | self.dcp_rank, |
| 1006 | self.dcp_world_size, |
| 1007 | self.cp_kv_cache_interleave_size, |
| 1008 | False, |
| 1009 | candidate_blocks=self.candidate_blocks, |
| 1010 | candidate_block_size=self.candidate_block_size, |
| 1011 | candidate_write=self.candidate_write, |
| 1012 | topk_backend=self.topk_backend, |
| 1013 | ) |
| 1014 | |
| 1015 | def forward_xpu( |
| 1016 | self, |
| 1017 | hidden_states: torch.Tensor, |
| 1018 | q_fp8: torch.Tensor, |
| 1019 | k: torch.Tensor | None, |
| 1020 | weights: torch.Tensor, |
| 1021 | ): |
| 1022 | return self.forward_cuda(hidden_states, q_fp8, k, weights) |
| 1023 | |
| 1024 | def forward_hip( |
| 1025 | self, |
| 1026 | hidden_states: torch.Tensor, |
| 1027 | q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 1028 | k: torch.Tensor | None, |
| 1029 | weights: torch.Tensor, |
| 1030 | ): |
| 1031 | assert not self.use_fp4_cache, "AMD platform doesn't support fp4 cache yet" |
| 1032 | assert isinstance(q_quant, torch.Tensor), ( |
| 1033 | "AMD sparse_attn_indexer expects a single FP8 q_quant tensor" |
| 1034 | ) |
| 1035 | from vllm.platforms.rocm import on_gfx11, on_gfx950 |
| 1036 | |
| 1037 | if ( |
| 1038 | rocm_aiter_ops.is_enabled() |
| 1039 | or rocm_aiter_ops.is_rdna_aiter_enabled() |
| 1040 | or on_gfx11() |
| 1041 | # The so-called AITER sparse indexer op has a native gfx950 path: |
| 1042 | # its cache insert, MQA logits, and top-k fallbacks are implemented |
| 1043 | # by local Triton/C++ kernels and do not require the aiter package. |
| 1044 | or on_gfx950() |
| 1045 | ): |
| 1046 | return torch.ops.vllm.rocm_aiter_sparse_attn_indexer( |
| 1047 | hidden_states, |
| 1048 | _encode_layer_name(self.k_cache.prefix), |
| 1049 | self.k_cache.kv_cache, |
| 1050 | q_quant, |
| 1051 | k, |
| 1052 | weights, |
| 1053 | self.quant_block_size, |
| 1054 | self.scale_fmt, |
| 1055 | self.topk_tokens, |
| 1056 | self.head_dim, |
| 1057 | self.max_model_len, |
| 1058 | self.max_total_seq_len, |
| 1059 | self.topk_indices_buffer, |
| 1060 | skip_k_cache_insert=self.skip_k_cache_insert, |
| 1061 | compress_ratio=self.compress_ratio, |
| 1062 | candidate_blocks=self.candidate_blocks, |
| 1063 | candidate_block_size=self.candidate_block_size, |
| 1064 | candidate_write=self.candidate_write, |
| 1065 | ) |
| 1066 | raise RuntimeError( |
| 1067 | "Sparse attention indexer ROCm path requires AITER or a supported " |
| 1068 | "native architecture (gfx950/gfx11)." |
| 1069 | ) |
| 1070 | |
| 1071 | def forward_cpu( |
| 1072 | self, |
| 1073 | hidden_states: torch.Tensor, |
| 1074 | q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 1075 | k: torch.Tensor | None, |
| 1076 | weights: torch.Tensor, |
| 1077 | ): |
| 1078 | """CPU sparse attention indexer: cache write stays eager Python glue |
| 1079 | (own K-cache layout, not shared with the main attention cache |
| 1080 | write). PREFILL and DECODE both call the ported |
| 1081 | ``fp8_paged_mqa_logits_cpu``/``topk_transform_512_cpu`` kernels, |
| 1082 | which read the paged K-cache directly via ``page_table`` -- no |
| 1083 | eager gather step, no per-request Python loop. |
| 1084 | |
| 1085 | ``prefill_metadata.chunks`` always has exactly one entry here: |
| 1086 | ``DeepseekV4CPUIndexerMetadataBuilder`` overrides the base chunk |
| 1087 | split to always return the whole step's prefill batch as one |
| 1088 | chunk, since the base chunking only bounds CUDA/XPU's dense M*N |
| 1089 | logits tensor and flat K-gather workspace, neither of which this |
| 1090 | paged kernel allocates. |
| 1091 | """ |
| 1092 | assert not self.use_fp4_cache, ( |
| 1093 | "CPU sparse indexer doesn't support fp4 cache yet" |
| 1094 | ) |
| 1095 | assert isinstance(q_quant, torch.Tensor), ( |
| 1096 | "CPU sparse_attn_indexer expects a single FP8 q_quant tensor" |
| 1097 | ) |
| 1098 | assert self.dcp_world_size <= 1 and not self.use_pcp, ( |
| 1099 | "CPU sparse indexer doesn't support decode/prefill context parallelism yet." |
| 1100 | ) |
| 1101 | |
| 1102 | forward_context = get_forward_context() |
| 1103 | attn_metadata = forward_context.attn_metadata |
| 1104 | attn_metadata_narrowed: DeepseekV32IndexerMetadata | None = None |
| 1105 | if isinstance(attn_metadata, dict): |
| 1106 | metadata = attn_metadata[self.k_cache.prefix] |
| 1107 | assert isinstance(metadata, DeepseekV32IndexerMetadata) |
| 1108 | attn_metadata_narrowed = metadata |
| 1109 | if attn_metadata_narrowed is None: |
| 1110 | # Profiling/dummy run: no real metadata to act on. |
| 1111 | return self.topk_indices_buffer |
| 1112 | |
| 1113 | kv_cache = self.k_cache.kv_cache |
| 1114 | topk_tokens = self.topk_tokens |
| 1115 | topk_indices_buffer = self.topk_indices_buffer |
| 1116 | slot_mapping = attn_metadata_narrowed.slot_mapping |
| 1117 | has_decode = attn_metadata_narrowed.num_decodes > 0 |
| 1118 | has_prefill = attn_metadata_narrowed.num_prefills > 0 |
| 1119 | num_decode_tokens = attn_metadata_narrowed.num_decode_tokens |
| 1120 | |
| 1121 | num_tokens = slot_mapping.shape[0] |
| 1122 | if k is not None: |
| 1123 | k = k[:num_tokens] |
| 1124 | |
| 1125 | if not self.skip_k_cache_insert: |
| 1126 | # Only reachable via DeepseekV32Attention with |
| 1127 | # prefill_context_parallel_size > 1 on CPU -- set_k_cpu/set_s_cpu |
| 1128 | # (the kernels this used to call) have been removed as unused/ |
| 1129 | # untested (csrc/cpu/sgl-kernels/store_cache.cpp). |
| 1130 | raise NotImplementedError( |
| 1131 | "SparseAttnIndexer.forward_cpu: skip_k_cache_insert=False " |
| 1132 | "(prefill context parallel on CPU) is not supported." |
| 1133 | ) |
| 1134 | |
| 1135 | topk_indices_buffer[: hidden_states.shape[0]] = -1 |
| 1136 | |
| 1137 | if has_prefill: |
| 1138 | assert topk_tokens == 512, ( |
| 1139 | "topk_transform_512_cpu only supports index_topk == 512." |
| 1140 | ) |
| 1141 | prefill_metadata = attn_metadata_narrowed.prefill |
| 1142 | assert prefill_metadata is not None |
| 1143 | assert len(prefill_metadata.chunks) == 1, ( |
| 1144 | "forward_cpu expects the prefill metadata builder to always " |
| 1145 | "produce a single chunk -- see " |
| 1146 | "DeepseekV4CPUIndexerMetadataBuilder._split_indexer_prefill_chunks." |
| 1147 | ) |
| 1148 | chunk = prefill_metadata.chunks[0] |
| 1149 | # kv_cache is a per-layer view into vLLM's shared multi-layer |
| 1150 | # cache allocation, so its block stride generally exceeds |
| 1151 | # block_size * page_width; fp8_paged_mqa_logits_cpu reads |
| 1152 | # kv_view.stride(0) explicitly, so no copy is needed here. |
| 1153 | kv_view = kv_cache.view(kv_cache.shape[0], -1) |
| 1154 | block_size = kv_cache.shape[1] |
| 1155 | q_slice = q_quant[chunk.token_start : chunk.token_end] |
| 1156 | topk_indices = topk_indices_buffer[ |
| 1157 | chunk.token_start : chunk.token_end, :topk_tokens |
| 1158 | ] |
| 1159 | if chunk.local_total_seq_lens == 0: |
| 1160 | topk_indices.fill_(-1) |
| 1161 | else: |
| 1162 | assert chunk.local_cu_seq_lens is not None |
| 1163 | # Each token's own (per-request, DCP-local) causal length. |
| 1164 | local_seq_lens = chunk.cu_seqlen_ke - chunk.cu_seqlen_ks |
| 1165 | # Recover each token's owning request from its row-start |
| 1166 | # tag. Ties (from zero-length requests) are harmless: those |
| 1167 | # tokens have local length 0 and never dereference |
| 1168 | # page_table. |
| 1169 | req_idx = ( |
| 1170 | torch.searchsorted( |
| 1171 | chunk.local_cu_seq_lens, chunk.cu_seqlen_ks, right=True |
| 1172 | ) |
| 1173 | - 1 |
| 1174 | ) |
| 1175 | page_table = chunk.block_table[req_idx] |
| 1176 | # The true max over this chunk's own rows, NOT |
| 1177 | # chunk.max_local_total_seq_lens (that field sums every |
| 1178 | # request's length in the chunk, bounding the old |
| 1179 | # flat-gather buffer this paged path no longer allocates). |
| 1180 | max_seq_len = int(local_seq_lens.max().item()) |
| 1181 | |
| 1182 | logits = ops.fp8_paged_mqa_logits_cpu( |
| 1183 | q_slice, |
| 1184 | kv_view, |
| 1185 | weights[chunk.token_start : chunk.token_end], |
| 1186 | local_seq_lens, |
| 1187 | page_table, |
| 1188 | block_size, |
| 1189 | max_seq_len, |
| 1190 | ) |
| 1191 | out_page_scratch = torch.empty( |
| 1192 | (q_slice.shape[0], topk_tokens), |
| 1193 | dtype=torch.int32, |
| 1194 | device=kv_cache.device, |
| 1195 | ) |
| 1196 | ops.topk_transform_512_cpu( |
| 1197 | logits, |
| 1198 | local_seq_lens, |
| 1199 | page_table, |
| 1200 | out_page_scratch, |
| 1201 | block_size, |
| 1202 | topk_indices, |
| 1203 | ) |
| 1204 | |
| 1205 | if has_decode: |
| 1206 | decode_metadata = attn_metadata_narrowed.decode |
| 1207 | assert decode_metadata is not None |
| 1208 | assert not decode_metadata.requires_padding, ( |
| 1209 | "CPU sparse indexer decode path does not support speculative " |
| 1210 | "decoding (native MTP) yet." |
| 1211 | ) |
| 1212 | batch_size = decode_metadata.decode_lens.shape[0] |
| 1213 | if batch_size > 0: |
| 1214 | # No native MTP on CPU (asserted above) => exactly one |
| 1215 | # query token per decode request, so the flat slice below |
| 1216 | # is already the batch-major layout |
| 1217 | # fp8_paged_mqa_logits_cpu wants. |
| 1218 | assert num_decode_tokens == batch_size, ( |
| 1219 | "CPU sparse indexer decode path expects exactly one query " |
| 1220 | "token per decode request." |
| 1221 | ) |
| 1222 | assert topk_tokens == 512, ( |
| 1223 | "topk_transform_512_cpu only supports index_topk == 512." |
| 1224 | ) |
| 1225 | seq_lens = decode_metadata.seq_lens |
| 1226 | seq_lens = ( |
| 1227 | seq_lens[:, -1].contiguous() if seq_lens.ndim == 2 else seq_lens |
| 1228 | ) |
| 1229 | block_table = decode_metadata.block_table[:batch_size] |
| 1230 | block_size = kv_cache.shape[1] |
| 1231 | kv_view = kv_cache.view(kv_cache.shape[0], -1) |
| 1232 | |
| 1233 | logits = ops.fp8_paged_mqa_logits_cpu( |
| 1234 | q_quant[:num_decode_tokens], |
| 1235 | kv_view, |
| 1236 | weights[:num_decode_tokens], |
| 1237 | seq_lens, |
| 1238 | block_table, |
| 1239 | block_size, |
| 1240 | attn_metadata_narrowed.max_seq_len, |
| 1241 | ) |
| 1242 | |
| 1243 | # out_page_indices is a required kernel output but unused: |
| 1244 | # the indexer's topk output must stay local/compressed- |
| 1245 | # context positions (resolved later by |
| 1246 | # DeepseekV4CPUAttention.forward_mqa via |
| 1247 | # map_local_to_global_slots_cpu). |
| 1248 | out_page_scratch = torch.empty( |
| 1249 | (batch_size, topk_tokens), |
| 1250 | dtype=torch.int32, |
| 1251 | device=kv_cache.device, |
| 1252 | ) |
| 1253 | ops.topk_transform_512_cpu( |
| 1254 | logits, |
| 1255 | seq_lens, |
| 1256 | block_table, |
| 1257 | out_page_scratch, |
| 1258 | block_size, |
| 1259 | topk_indices_buffer[:num_decode_tokens, :topk_tokens], |
| 1260 | ) |
| 1261 | |
| 1262 | return topk_indices_buffer |