| 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 | |
| 7 | import vllm.envs as envs |
| 8 | from vllm.compilation.breakable_cudagraph import eager_break_during_capture |
| 9 | from vllm.config import get_current_vllm_config_or_none |
| 10 | from vllm.forward_context import get_forward_context |
| 11 | from vllm.logger import init_logger |
| 12 | from vllm.model_executor.custom_op import CustomOp |
| 13 | from vllm.model_executor.layers.indexer_topk import get_indexer_topk |
| 14 | from vllm.models.glm5next.common.sparse_indexer import ( |
| 15 | RADIX_TOPK_WORKSPACE_SIZE, |
| 16 | _build_decode_scatter_indices, |
| 17 | _decode_topk_seq_lens, |
| 18 | _fill_causal_indices, |
| 19 | _fill_short_decode_causal_indices, |
| 20 | _gather_workspace_shapes, |
| 21 | _scatter_decode_tokens_by_request, |
| 22 | kv_cache_as_quant_view, |
| 23 | ) |
| 24 | from vllm.models.glm5next.nvidia.ops import kpool_compress as kpool_ops |
| 25 | from vllm.platforms import current_platform |
| 26 | from vllm.utils.deep_gemm import has_deep_gemm |
| 27 | from vllm.utils.torch_utils import ( |
| 28 | LayerNameType, |
| 29 | _resolve_layer_name, |
| 30 | ) |
| 31 | from vllm.v1.attention.backends.mla.indexer import ( |
| 32 | DeepseekV32IndexerMetadata, |
| 33 | ) |
| 34 | from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton |
| 35 | from vllm.v1.worker.workspace import current_workspace_manager |
| 36 | |
| 37 | if current_platform.is_cuda_alike(): |
| 38 | from vllm import _custom_ops as ops |
| 39 | |
| 40 | logger = init_logger(__name__) |
| 41 | |
| 42 | |
| 43 | # kpool write helper: form pools from the current token batch and compress them |
| 44 | # into the index K cache via the fused Triton kernel. |
| 45 | |
| 46 | |
| 47 | def _kpool_compress_insert( |
| 48 | k: torch.Tensor, |
| 49 | gate_score: torch.Tensor, |
| 50 | ape: torch.Tensor, |
| 51 | kv_cache: torch.Tensor, |
| 52 | slot_mapping: torch.Tensor, |
| 53 | kpool: int, |
| 54 | head_dim: int, |
| 55 | round_scale: bool, |
| 56 | ) -> None: |
| 57 | """Pool ``kpool`` consecutive tokens into one fp8 K and write at pool slots. |
| 58 | |
| 59 | ``slot_mapping`` is pool-granular (compress_ratio == kpool on the spec): |
| 60 | only the *last* token of each complete pool carries a valid (>=0) slot; |
| 61 | intra-pool tokens are -1. Every position is treated as a pool-completion |
| 62 | candidate and non-completions are masked off inside the kernel. Compacting |
| 63 | the valid rows first costs two device syncs on the eager prefill path and |
| 64 | buys nothing numerically. Assumes pool-aligned chunk starts. |
| 65 | """ |
| 66 | n = slot_mapping.shape[0] |
| 67 | # No pool can complete in a batch smaller than one pool; also keeps the |
| 68 | # clamped gather indices below in bounds. |
| 69 | if n < kpool: |
| 70 | return |
| 71 | pos = torch.arange(n, device=k.device) |
| 72 | valid = slot_mapping >= 0 |
| 73 | # Drop pools whose start falls before the batch (leading padding); their |
| 74 | # gate/k data is undefined anyway. |
| 75 | write_mask = valid & (pos >= kpool - 1) |
| 76 | offs = torch.arange(kpool, device=k.device) |
| 77 | idx = (pos - (kpool - 1)).clamp_min(0)[:, None] + offs[None, :] |
| 78 | kpool_ops.kpool_compress_and_write_cache( |
| 79 | kv_cache, |
| 80 | k[idx], # [n, kpool, head_dim] |
| 81 | gate_score[idx], |
| 82 | ape, |
| 83 | slot_mapping.to(torch.int64), |
| 84 | pool_size=kpool, |
| 85 | head_dim=head_dim, |
| 86 | write_mask=write_mask, |
| 87 | round_scale=round_scale, |
| 88 | write_cache=True, |
| 89 | return_compressed=False, |
| 90 | ) |
| 91 | |
| 92 | |
| 93 | @eager_break_during_capture |
| 94 | def sparse_attn_indexer_kpool( |
| 95 | hidden_states: torch.Tensor, |
| 96 | k_cache_prefix: LayerNameType, |
| 97 | kv_cache: torch.Tensor, |
| 98 | q_quant: torch.Tensor, |
| 99 | q_scale: torch.Tensor | None, |
| 100 | k: torch.Tensor, |
| 101 | weights: torch.Tensor, |
| 102 | quant_block_size: int, |
| 103 | scale_fmt: str | None, |
| 104 | topk_tokens: int, |
| 105 | head_dim: int, |
| 106 | max_pool_len: int, |
| 107 | total_seq_lens: int, |
| 108 | topk_indices_buffer: torch.Tensor, |
| 109 | skip_k_cache_insert: bool, |
| 110 | use_fp4_cache: bool = False, |
| 111 | # kpool params (Plan-A: gate is consumed at write time and read back at |
| 112 | # topk time to softmax-weight the pool). |
| 113 | gate_score: torch.Tensor | None = None, |
| 114 | compress_ape: torch.Tensor | None = None, |
| 115 | index_kpool: int = 1, |
| 116 | positions: torch.Tensor | None = None, |
| 117 | # Paged tail cache (in-progress pool's raw K + gate score), replacing the |
| 118 | # transient _DECODE_TAIL ring. tail_prefix resolves attn_metadata[tail_prefix] |
| 119 | # for the tail group's token-granular slot_mapping. None on the dummy/profiling |
| 120 | # path and when the tail cache is disabled. |
| 121 | tail_kv_cache: torch.Tensor | None = None, |
| 122 | tail_prefix: str | None = None, |
| 123 | topk_backend: str = "auto", |
| 124 | ) -> torch.Tensor: |
| 125 | # careful! this will be None in dummy run |
| 126 | attn_metadata = get_forward_context().attn_metadata |
| 127 | fp8_dtype = current_platform.fp8_dtype() |
| 128 | k_cache_prefix = _resolve_layer_name(k_cache_prefix) |
| 129 | |
| 130 | # assert isinstance(attn_metadata, dict) |
| 131 | if not isinstance(attn_metadata, dict): |
| 132 | # Reserve workspace for indexer during profiling run |
| 133 | values_spec, scales_spec = _gather_workspace_shapes( |
| 134 | total_seq_lens, head_dim, fp8_dtype, use_fp4_cache |
| 135 | ) |
| 136 | current_workspace_manager().get_simultaneous( |
| 137 | values_spec, |
| 138 | scales_spec, |
| 139 | ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8), |
| 140 | ) |
| 141 | |
| 142 | # Reserve profiler-visible memory for the worst-case decode logits, |
| 143 | # whose shape is [B * next_n, max_pool_len]. This profiling branch |
| 144 | # returns before invoking the logits kernel itself. |
| 145 | cfg = get_current_vllm_config_or_none() |
| 146 | worst_decode_tokens = 0 |
| 147 | if cfg is not None: |
| 148 | sched = cfg.scheduler_config |
| 149 | num_spec = ( |
| 150 | cfg.speculative_config.num_speculative_tokens |
| 151 | if cfg.speculative_config is not None |
| 152 | else 0 |
| 153 | ) |
| 154 | worst_decode_tokens = min( |
| 155 | sched.max_num_seqs * (num_spec + 1), |
| 156 | sched.max_num_batched_tokens, |
| 157 | ) |
| 158 | # float32 logits -> 4 bytes/element; uint8 sentinel so elems == bytes. |
| 159 | decode_logits_elems = worst_decode_tokens * max_pool_len * 4 |
| 160 | prefill_cap_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 |
| 161 | max_logits_elems = max(decode_logits_elems, prefill_cap_elems) |
| 162 | _ = torch.empty( |
| 163 | max_logits_elems, dtype=torch.uint8, device=hidden_states.device |
| 164 | ) |
| 165 | |
| 166 | return topk_indices_buffer |
| 167 | attn_metadata_narrowed = attn_metadata[k_cache_prefix] |
| 168 | assert isinstance(attn_metadata_narrowed, DeepseekV32IndexerMetadata) |
| 169 | slot_mapping = attn_metadata_narrowed.slot_mapping |
| 170 | has_decode = attn_metadata_narrowed.num_decodes > 0 |
| 171 | has_prefill = attn_metadata_narrowed.num_prefills > 0 |
| 172 | num_decode_tokens = attn_metadata_narrowed.num_decode_tokens |
| 173 | |
| 174 | # q_scale is required iff the FP4 cache path is enabled; the FP8 path |
| 175 | # folds the Q scale into `weights` inside fused_indexer_q_rope_quant. |
| 176 | if use_fp4_cache: |
| 177 | assert q_scale is not None, "use_fp4_cache=True requires q_scale" |
| 178 | else: |
| 179 | assert q_scale is None, "q_scale must be None when use_fp4_cache=False" |
| 180 | |
| 181 | # During speculative decoding, k may be padded to the CUDA graph batch |
| 182 | # size while slot_mapping only covers actual tokens. Truncate k to avoid |
| 183 | # out-of-bounds reads in the kernel. |
| 184 | num_tokens = slot_mapping.shape[0] |
| 185 | if k is not None: |
| 186 | k = k[:num_tokens] |
| 187 | |
| 188 | if not skip_k_cache_insert: |
| 189 | assert not use_fp4_cache, "Unfused FP4 Insert is not supported yet" |
| 190 | if index_kpool > 1 and gate_score is not None and compress_ape is not None: |
| 191 | # kpool prefill write: pool kpool consecutive prefill tokens via |
| 192 | # softmax(gate+ape)-weighted sum -> Hadamard -> fp8 -> pool slots. |
| 193 | # Decode tokens (the first num_decode_tokens in the batch) cannot be |
| 194 | # pooled here — their pool's earlier tokens are not in this batch — |
| 195 | # so they are deferred to the tail-buffer kernel in has_decode. |
| 196 | # compress_ratio == index_kpool makes slot_mapping pool-granular. |
| 197 | n_prefill = num_tokens - num_decode_tokens |
| 198 | if n_prefill > 0: |
| 199 | # decode tokens are batched first; prefill tokens follow. |
| 200 | prefill_slice = slice(num_decode_tokens, num_tokens) |
| 201 | _kpool_compress_insert( |
| 202 | k[prefill_slice], |
| 203 | gate_score[prefill_slice], |
| 204 | compress_ape, |
| 205 | kv_cache, |
| 206 | slot_mapping[prefill_slice], |
| 207 | index_kpool, |
| 208 | head_dim, |
| 209 | round_scale=(scale_fmt is not None), |
| 210 | ) |
| 211 | # Persist each request's incomplete prefill pool so decode can |
| 212 | # finish it, including after PD transfer. Tail slots use |
| 213 | # ``pos % kpool`` within the request's tail block. Processing |
| 214 | # only the batch's trailing tokens would miss all but the last |
| 215 | # request in a multi-request prefill. |
| 216 | if tail_kv_cache is not None and tail_prefix is not None: |
| 217 | tail_meta = attn_metadata.get(_resolve_layer_name(tail_prefix)) |
| 218 | if tail_meta is not None: |
| 219 | assert isinstance(tail_meta, DeepseekV32IndexerMetadata) |
| 220 | kpool_ops.kpool_seed_tail_cache( |
| 221 | tail_kv_cache, |
| 222 | k[prefill_slice], |
| 223 | gate_score[prefill_slice], |
| 224 | tail_meta.slot_mapping[prefill_slice], |
| 225 | index_kpool, |
| 226 | head_dim, |
| 227 | ) |
| 228 | else: |
| 229 | # standard: per-token fp8 quant + scatter (all tokens). |
| 230 | assert scale_fmt is not None |
| 231 | ops.indexer_k_quant_and_cache( |
| 232 | k, |
| 233 | kv_cache, |
| 234 | slot_mapping, |
| 235 | quant_block_size, |
| 236 | scale_fmt, |
| 237 | ) |
| 238 | |
| 239 | topk_indices_buffer[: hidden_states.shape[0]] = -1 |
| 240 | if has_prefill: |
| 241 | prefill_metadata = attn_metadata_narrowed.prefill |
| 242 | assert prefill_metadata is not None |
| 243 | |
| 244 | # Short sequences select every pool, so skip sparse scoring and fill |
| 245 | # the top-k buffer with all causal token indices. The index-K cache was |
| 246 | # already written above. |
| 247 | n_prefill_sf = num_tokens - num_decode_tokens |
| 248 | # Host-side short-prefill predicate: max_prefill_seq_len is computed |
| 249 | # in the metadata builder (exact for prefill rows) and equals |
| 250 | # positions[prefill_slice].max() + 1, so this replaces a |
| 251 | # positions.max().item() device sync per layer. -1 (unknown metadata) |
| 252 | # falls back to the device-side check. |
| 253 | if prefill_metadata.max_prefill_seq_len >= 0: |
| 254 | short_prefill = ( |
| 255 | n_prefill_sf > 0 |
| 256 | and positions is not None |
| 257 | and prefill_metadata.max_prefill_seq_len <= topk_tokens |
| 258 | ) |
| 259 | else: |
| 260 | short_prefill = ( |
| 261 | n_prefill_sf > 0 |
| 262 | and positions is not None |
| 263 | and int(positions[num_decode_tokens:num_tokens].max().item()) + 1 |
| 264 | <= topk_tokens |
| 265 | ) |
| 266 | if short_prefill: |
| 267 | # short_prefill is only True when positions is not None (above), |
| 268 | # but narrow explicitly for the indexer below. |
| 269 | assert positions is not None |
| 270 | _pos = positions[num_decode_tokens:num_tokens].to(torch.int32) |
| 271 | _buf = topk_indices_buffer[num_decode_tokens:num_tokens] |
| 272 | _fill_causal_indices(_buf, _pos) |
| 273 | |
| 274 | # Get the full shared workspace buffers once (will allocate on first use). |
| 275 | # Layout switches between FP8 (head_dim bytes + 4-byte fp32 scale) and |
| 276 | # MXFP4 (head_dim/2 bytes packed + head_dim/MXFP4_BLOCK_SIZE ue8m0 |
| 277 | # scales) based on use_fp4_cache. |
| 278 | workspace_manager = current_workspace_manager() |
| 279 | values_spec, scales_spec = _gather_workspace_shapes( |
| 280 | total_seq_lens, head_dim, fp8_dtype, use_fp4_cache |
| 281 | ) |
| 282 | k_quant_full, k_scale_full = workspace_manager.get_simultaneous( |
| 283 | values_spec, |
| 284 | scales_spec, |
| 285 | ) |
| 286 | for chunk in prefill_metadata.chunks if not short_prefill else (): |
| 287 | k_quant = k_quant_full[: chunk.total_seq_lens] |
| 288 | k_scale = k_scale_full[: chunk.total_seq_lens] |
| 289 | |
| 290 | if not chunk.skip_kv_gather: |
| 291 | ops.cp_gather_indexer_k_quant_cache( |
| 292 | kv_cache, |
| 293 | k_quant, |
| 294 | k_scale, |
| 295 | chunk.block_table, |
| 296 | chunk.cu_seq_lens, |
| 297 | ) |
| 298 | |
| 299 | q_slice = q_quant[chunk.token_start : chunk.token_end] |
| 300 | q_scale_slice = ( |
| 301 | q_scale[chunk.token_start : chunk.token_end] |
| 302 | if q_scale is not None |
| 303 | else None |
| 304 | ) |
| 305 | # DeepGEMM scalar-type tags (zero-copy): MXFP4 values → int8 |
| 306 | # (kPackedFP4), scales → int32 squeezed to 1-D kv_sf / 2-D q_sf. |
| 307 | if use_fp4_cache: |
| 308 | q_slice_cast = q_slice.view(torch.int8) |
| 309 | k_quant_cast = k_quant.view(torch.int8) |
| 310 | k_scale_cast = k_scale.view(torch.int32).squeeze(-1) |
| 311 | else: |
| 312 | q_slice_cast = q_slice |
| 313 | k_quant_cast = k_quant |
| 314 | k_scale_cast = k_scale.view(torch.float32).squeeze(-1) |
| 315 | from vllm.utils.deep_gemm import fp8_fp4_mqa_logits |
| 316 | |
| 317 | logits = fp8_fp4_mqa_logits( |
| 318 | (q_slice_cast, q_scale_slice), |
| 319 | (k_quant_cast, k_scale_cast), |
| 320 | weights[chunk.token_start : chunk.token_end], |
| 321 | chunk.cu_seqlen_ks, |
| 322 | chunk.cu_seqlen_ke, |
| 323 | clean_logits=False, |
| 324 | ) |
| 325 | num_rows = logits.shape[0] |
| 326 | |
| 327 | # kpool: logits are pool-granular (compress_ratio == index_kpool), |
| 328 | # so topk selects pools. We pick topk_tokens // kpool pools then |
| 329 | # expand each pool back to its kpool constituent tokens. |
| 330 | select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens |
| 331 | if index_kpool > 1: |
| 332 | pool_topk = torch.empty( |
| 333 | (num_rows, select_k), dtype=torch.int32, device=logits.device |
| 334 | ) |
| 335 | topk_dst = pool_topk |
| 336 | else: |
| 337 | topk_dst = topk_indices_buffer[ |
| 338 | chunk.token_start : chunk.token_end, :topk_tokens |
| 339 | ] |
| 340 | |
| 341 | torch.ops._C.top_k_per_row_prefill( |
| 342 | logits, |
| 343 | chunk.cu_seqlen_ks, |
| 344 | chunk.cu_seqlen_ke, |
| 345 | topk_dst, |
| 346 | num_rows, |
| 347 | logits.stride(0), |
| 348 | logits.stride(1), |
| 349 | select_k, |
| 350 | ) |
| 351 | |
| 352 | if index_kpool > 1: |
| 353 | pool_ids = pool_topk.to(torch.int64) |
| 354 | if positions is not None: |
| 355 | # Fused expand-pools + append-tail into one Triton kernel |
| 356 | # (replaces ~25 elementwise ops). seq_len is token-granular |
| 357 | # (pos+1); the kernel derives pool_len internally. |
| 358 | q_seq = ( |
| 359 | positions[chunk.token_start : chunk.token_end].to(torch.int32) |
| 360 | + 1 |
| 361 | ) |
| 362 | expanded = kpool_ops.expand_pools_and_append_tail( |
| 363 | pool_ids, q_seq, index_kpool |
| 364 | ) |
| 365 | else: |
| 366 | valid = pool_ids >= 0 |
| 367 | expanded = kpool_ops.expand_pools_to_tokens( |
| 368 | pool_ids, valid, topk_tokens, index_kpool |
| 369 | ) |
| 370 | topk_indices_buffer[ |
| 371 | chunk.token_start : chunk.token_end, : expanded.shape[-1] |
| 372 | ] = expanded |
| 373 | |
| 374 | if has_decode: |
| 375 | decode_metadata = attn_metadata_narrowed.decode |
| 376 | assert decode_metadata is not None |
| 377 | kv_cache_raw = kv_cache # raw [num_blocks, block_size, head_dim+4] for writes |
| 378 | kv_cache = kv_cache_as_quant_view(kv_cache, head_dim, use_fp4_cache) |
| 379 | |
| 380 | # Update the tail before reading logits; completed pools are compressed |
| 381 | # into the slot supplied by slot_mapping. |
| 382 | # Spec verification groups tokens by request and preserves position |
| 383 | # order so each token is stashed before the next completes its pool. |
| 384 | # Positions must remain token-granular because the kernel derives the |
| 385 | # pool phase and tail index from ``pos % kpool``. |
| 386 | if ( |
| 387 | index_kpool > 1 |
| 388 | and gate_score is not None |
| 389 | and compress_ape is not None |
| 390 | and positions is not None |
| 391 | and not skip_k_cache_insert |
| 392 | ): |
| 393 | num_requests = attn_metadata_narrowed.num_decodes |
| 394 | # Kpool writes must recover the original request grouping after the |
| 395 | # indexer's flattened decode path. Host metadata avoids a CUDA graph |
| 396 | # sync when choosing the uniform or padded layout. |
| 397 | per_req_lens = decode_metadata.per_req_decode_lens |
| 398 | if per_req_lens is not None: |
| 399 | use_uniform = ( |
| 400 | decode_metadata.decode_is_uniform |
| 401 | and num_decode_tokens |
| 402 | == num_requests * decode_metadata.write_max_decode_len |
| 403 | ) |
| 404 | group_lens = per_req_lens |
| 405 | lmax = decode_metadata.write_max_decode_len |
| 406 | else: |
| 407 | # Legacy metadata without per-request lens: fall back to the |
| 408 | # host-side requires_padding flag. Unreached now (per-request |
| 409 | # lens is always populated for decode), kept defensive. |
| 410 | use_uniform = not decode_metadata.requires_padding |
| 411 | group_lens = decode_metadata.decode_lens |
| 412 | lmax = int(decode_metadata.decode_lens.max().item()) |
| 413 | if not use_uniform: |
| 414 | # Non-uniform decode_lens (mixed plain-decode + spec-verify, or |
| 415 | # a variable MTP-verify batch): scatter actual tokens into a |
| 416 | # padded [B, lmax] layout. int32 tensors can't go through |
| 417 | # pack_seq_triton (float/uint8 only). The scatter indices are |
| 418 | # shared by all five scatters below (and the tail slot one). |
| 419 | scatter_idx = _build_decode_scatter_indices( |
| 420 | group_lens, num_requests, num_decode_tokens |
| 421 | ) |
| 422 | dec_k = _scatter_decode_tokens_by_request( |
| 423 | k[:num_decode_tokens], 0, num_requests, lmax, scatter_idx |
| 424 | ) |
| 425 | dec_gate = _scatter_decode_tokens_by_request( |
| 426 | gate_score[:num_decode_tokens], |
| 427 | 0, |
| 428 | num_requests, |
| 429 | lmax, |
| 430 | scatter_idx, |
| 431 | ) |
| 432 | dec_slot = _scatter_decode_tokens_by_request( |
| 433 | slot_mapping[:num_decode_tokens], |
| 434 | -1, |
| 435 | num_requests, |
| 436 | lmax, |
| 437 | scatter_idx, |
| 438 | ) |
| 439 | dec_pos = _scatter_decode_tokens_by_request( |
| 440 | positions[:num_decode_tokens].to(torch.int32), |
| 441 | -1, |
| 442 | num_requests, |
| 443 | lmax, |
| 444 | scatter_idx, |
| 445 | ) |
| 446 | else: |
| 447 | next_n = num_decode_tokens // num_requests |
| 448 | shape2 = (num_requests, next_n) |
| 449 | dec_k = k[:num_decode_tokens].view(*shape2, head_dim) |
| 450 | dec_gate = gate_score[:num_decode_tokens].view(*shape2, head_dim) |
| 451 | dec_slot = slot_mapping[:num_decode_tokens].view(shape2) |
| 452 | dec_pos = positions[:num_decode_tokens].to(torch.int32).view(shape2) |
| 453 | tail_meta = ( |
| 454 | attn_metadata.get(_resolve_layer_name(tail_prefix)) |
| 455 | if tail_prefix is not None |
| 456 | else None |
| 457 | ) |
| 458 | # Paged tail cache replaces the transient _DECODE_TAIL ring. Group |
| 459 | # the tail group's token-granular slot_mapping per-request, mirroring |
| 460 | # dec_slot / dec_pos, so the kernel gets each request's current-token |
| 461 | # tail slot (block * kpool + pos % kpool). |
| 462 | if tail_meta is not None: |
| 463 | assert isinstance(tail_meta, DeepseekV32IndexerMetadata) |
| 464 | if tail_meta is None or tail_kv_cache is None: |
| 465 | dec_tail_slot = None |
| 466 | elif not use_uniform: |
| 467 | dec_tail_slot = _scatter_decode_tokens_by_request( |
| 468 | tail_meta.slot_mapping[:num_decode_tokens], |
| 469 | -1, |
| 470 | num_requests, |
| 471 | lmax, |
| 472 | scatter_idx, |
| 473 | ) |
| 474 | else: |
| 475 | dec_tail_slot = tail_meta.slot_mapping[:num_decode_tokens].view(shape2) |
| 476 | # The compress kernel writes the raw fp8 cache (not the quant view); |
| 477 | # pass the underlying kv_cache, not kv_cache_quant_view. |
| 478 | if dec_tail_slot is not None: |
| 479 | # Single batched launch over [num_requests, next_n] replaces the |
| 480 | # per-token sequential loop. The kernel iterates each request's |
| 481 | # tokens in position order internally, preserving the |
| 482 | # pool-completion read-after-stash dependency that the loop |
| 483 | # provided. Inputs are already grouped per request (uniform: |
| 484 | # view; non-uniform: _scatter_decode_tokens_by_request padded to |
| 485 | # [B, lmax]) — no per-token .contiguous() copies needed. |
| 486 | kpool_ops.kpool_decode_update_and_maybe_write_cache_batched( |
| 487 | kv_cache_raw, |
| 488 | tail_kv_cache, |
| 489 | dec_tail_slot, |
| 490 | dec_k, |
| 491 | dec_gate, |
| 492 | compress_ape, |
| 493 | dec_slot, |
| 494 | dec_pos, |
| 495 | index_kpool, |
| 496 | head_dim, |
| 497 | round_scale=(scale_fmt is not None), |
| 498 | ) |
| 499 | if current_platform.is_cuda_alike() and _fill_short_decode_causal_indices( |
| 500 | topk_indices_buffer, |
| 501 | positions, |
| 502 | num_decode_tokens, |
| 503 | attn_metadata_narrowed.max_seq_len, |
| 504 | topk_tokens, |
| 505 | ): |
| 506 | return topk_indices_buffer |
| 507 | decode_lens = decode_metadata.decode_lens |
| 508 | if decode_metadata.requires_padding: |
| 509 | # Padding also covers short chunked prefills classified as decode. |
| 510 | # MXFP4 uses zero-byte padding so padded slots dequantize to zero. |
| 511 | if q_scale is not None: |
| 512 | padded_q_quant_decode_tokens = pack_seq_triton( |
| 513 | q_quant[:num_decode_tokens], decode_lens, pad_value=0 |
| 514 | ) |
| 515 | padded_q_scale = pack_seq_triton( |
| 516 | q_scale[:num_decode_tokens], decode_lens, pad_value=0 |
| 517 | ) |
| 518 | else: |
| 519 | padded_q_quant_decode_tokens = pack_seq_triton( |
| 520 | q_quant[:num_decode_tokens], decode_lens |
| 521 | ) |
| 522 | padded_q_scale = None |
| 523 | padded_weights = pack_seq_triton( |
| 524 | weights[:num_decode_tokens], decode_lens, pad_value=0 |
| 525 | ).reshape(-1, *weights.shape[1:]) |
| 526 | else: |
| 527 | padded_q_quant_decode_tokens = q_quant[:num_decode_tokens].reshape( |
| 528 | decode_lens.shape[0], -1, *q_quant.shape[1:] |
| 529 | ) |
| 530 | if q_scale is not None: |
| 531 | padded_q_scale = q_scale[:num_decode_tokens].reshape( |
| 532 | decode_lens.shape[0], -1, *q_scale.shape[1:] |
| 533 | ) |
| 534 | else: |
| 535 | padded_q_scale = None |
| 536 | padded_weights = weights[:num_decode_tokens] |
| 537 | # TODO: move and optimize below logic with triton kernels |
| 538 | batch_size = padded_q_quant_decode_tokens.shape[0] |
| 539 | next_n = padded_q_quant_decode_tokens.shape[1] |
| 540 | num_padded_tokens = batch_size * next_n |
| 541 | seq_lens = decode_metadata.seq_lens[:batch_size] |
| 542 | # seq_lens is always 2D: (B, next_n) for native spec decode, (B, 1) |
| 543 | # otherwise. deep_gemm fp8_fp4_paged_mqa_logits requires 2D context_lens; |
| 544 | # the downstream topk kernels accept both 1D and 2D. |
| 545 | padded_q_quant_cast = ( |
| 546 | padded_q_quant_decode_tokens.view(torch.int8) |
| 547 | if use_fp4_cache |
| 548 | else padded_q_quant_decode_tokens |
| 549 | ) |
| 550 | from vllm.utils.deep_gemm import fp8_fp4_paged_mqa_logits |
| 551 | |
| 552 | logits = fp8_fp4_paged_mqa_logits( |
| 553 | (padded_q_quant_cast, padded_q_scale), |
| 554 | kv_cache, |
| 555 | padded_weights[:num_padded_tokens], |
| 556 | seq_lens, |
| 557 | decode_metadata.block_table, |
| 558 | decode_metadata.schedule_metadata, |
| 559 | max_model_len=max_pool_len, |
| 560 | clean_logits=False, |
| 561 | indices=decode_metadata.indices, |
| 562 | ) |
| 563 | num_rows = logits.shape[0] |
| 564 | # kpool: logits are pool-granular -> select topk_tokens//kpool pools, |
| 565 | # then expand each pool back to its kpool tokens. |
| 566 | select_k = topk_tokens // index_kpool if index_kpool > 1 else topk_tokens |
| 567 | if index_kpool > 1: |
| 568 | pool_topk = torch.empty( |
| 569 | (num_rows, select_k), dtype=torch.int32, device=logits.device |
| 570 | ) |
| 571 | topk_dst = pool_topk |
| 572 | else: |
| 573 | topk_dst = topk_indices_buffer[:num_padded_tokens, :topk_tokens] |
| 574 | |
| 575 | # Shared dispatcher with the DSA sparse indexer, so the kpool select |
| 576 | # honors kernel_config.sparse_indexer_topk_backend and picks between |
| 577 | # cooperative/persistent by batch size instead of always taking the |
| 578 | # same kernel. |
| 579 | get_indexer_topk(topk_backend)( |
| 580 | logits, |
| 581 | seq_lens, |
| 582 | next_n, |
| 583 | topk_dst, |
| 584 | select_k, |
| 585 | attn_metadata_narrowed.max_seq_len, |
| 586 | ) |
| 587 | |
| 588 | # Resolve to token-level indices in the output buffer. |
| 589 | if index_kpool > 1: |
| 590 | pool_ids = pool_topk.to(torch.int64) |
| 591 | n = pool_topk.shape[0] |
| 592 | # Decode seq_lens are pool-granular; recover token lengths from |
| 593 | # positions using the padded [B, next_n] row layout when needed. |
| 594 | if positions is not None: |
| 595 | dec_seq = _decode_topk_seq_lens( |
| 596 | positions, |
| 597 | decode_lens, |
| 598 | num_decode_tokens, |
| 599 | batch_size, |
| 600 | next_n, |
| 601 | decode_metadata.requires_padding, |
| 602 | ) |
| 603 | else: |
| 604 | dec_seq = decode_metadata.seq_lens[:n] |
| 605 | if dec_seq.ndim == 2: |
| 606 | dec_seq = dec_seq[:, -1] |
| 607 | dec_seq = dec_seq.to(torch.int32) |
| 608 | out = kpool_ops.expand_pools_and_append_tail(pool_ids, dec_seq, index_kpool) |
| 609 | else: |
| 610 | out = topk_dst |
| 611 | |
| 612 | if decode_metadata.requires_padding: |
| 613 | # Drop padded query rows introduced by the next_n padding above. |
| 614 | out = unpack_seq_triton( |
| 615 | out.reshape(batch_size, -1, out.shape[-1]), decode_lens |
| 616 | ) |
| 617 | topk_indices_buffer[: out.shape[0], : out.shape[-1]] = out |
| 618 | |
| 619 | return topk_indices_buffer |
| 620 | |
| 621 | |
| 622 | @CustomOp.register("sparse_attn_indexer_kpool") |
| 623 | class SparseAttnIndexerKpool(CustomOp): |
| 624 | """Sparse Attention Indexer Custom Op Layer. This layer is extracted as a |
| 625 | separate custom op since it involves heavy custom kernels like `mqa_logits`, |
| 626 | `paged_mqa_logits` and `top_k_per_row`, etc. Those kernels maybe requires |
| 627 | specific memory layout or implementation for different hardware backends to |
| 628 | achieve optimal performance. |
| 629 | |
| 630 | For now, the default native path will use CUDA backend path. Other platform |
| 631 | may requires add the corresponding Custom Op name `sparse_attn_indexer` to |
| 632 | `custom_ops` in `CompilationConfig` to enable the platform specific path. |
| 633 | """ |
| 634 | |
| 635 | def __init__( |
| 636 | self, |
| 637 | k_cache, |
| 638 | quant_block_size: int, |
| 639 | scale_fmt: str, |
| 640 | topk_tokens: int, |
| 641 | head_dim: int, |
| 642 | max_pool_len: int, |
| 643 | max_total_seq_len: int, |
| 644 | topk_indices_buffer: torch.Tensor, |
| 645 | skip_k_cache_insert: bool = False, |
| 646 | use_fp4_cache: bool = False, |
| 647 | tail_cache=None, |
| 648 | ): |
| 649 | super().__init__() |
| 650 | self.k_cache = k_cache |
| 651 | self.tail_cache = tail_cache |
| 652 | self.quant_block_size = quant_block_size |
| 653 | self.scale_fmt = scale_fmt |
| 654 | self.topk_tokens = topk_tokens |
| 655 | self.head_dim = head_dim |
| 656 | self.max_pool_len = max_pool_len |
| 657 | self.max_total_seq_len = max_total_seq_len |
| 658 | self.topk_indices_buffer = topk_indices_buffer |
| 659 | self.skip_k_cache_insert = skip_k_cache_insert |
| 660 | self.use_fp4_cache = use_fp4_cache |
| 661 | if current_platform.is_cuda() and not has_deep_gemm(): |
| 662 | raise RuntimeError( |
| 663 | "Sparse Attention Indexer CUDA op requires DeepGEMM to be installed." |
| 664 | ) |
| 665 | _cfg = get_current_vllm_config_or_none() |
| 666 | self.topk_backend = ( |
| 667 | _cfg.kernel_config.sparse_indexer_topk_backend |
| 668 | if _cfg is not None |
| 669 | else "auto" |
| 670 | ) |
| 671 | _parallel = _cfg.parallel_config if _cfg is not None else None |
| 672 | if ( |
| 673 | _parallel is not None |
| 674 | and _parallel.prefill_context_parallel_size > 1 |
| 675 | and _parallel.decode_context_parallel_size > 1 |
| 676 | ): |
| 677 | raise NotImplementedError( |
| 678 | "SparseAttnIndexerKpool does not support PCP+DCP." |
| 679 | ) |
| 680 | |
| 681 | def forward_native( |
| 682 | self, |
| 683 | hidden_states: torch.Tensor, |
| 684 | q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 685 | k: torch.Tensor, |
| 686 | weights: torch.Tensor, |
| 687 | *, |
| 688 | gate_score: torch.Tensor | None = None, |
| 689 | compress_ape: torch.Tensor | None = None, |
| 690 | index_kpool: int = 1, |
| 691 | positions: torch.Tensor | None = None, |
| 692 | ): |
| 693 | return self.forward_cuda( |
| 694 | hidden_states, |
| 695 | q_quant, |
| 696 | k, |
| 697 | weights, |
| 698 | gate_score=gate_score, |
| 699 | compress_ape=compress_ape, |
| 700 | index_kpool=index_kpool, |
| 701 | positions=positions, |
| 702 | ) |
| 703 | |
| 704 | def forward_cuda( |
| 705 | self, |
| 706 | hidden_states: torch.Tensor, |
| 707 | q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 708 | k: torch.Tensor, |
| 709 | weights: torch.Tensor, |
| 710 | *, |
| 711 | gate_score: torch.Tensor | None = None, |
| 712 | compress_ape: torch.Tensor | None = None, |
| 713 | index_kpool: int = 1, |
| 714 | positions: torch.Tensor | None = None, |
| 715 | ): |
| 716 | # FP8 path: single tensor (per-token scale is folded into `weights`). |
| 717 | # FP4 path: (values, scales) tuple with scales required by the kernel. |
| 718 | if isinstance(q_quant, tuple): |
| 719 | q_values, q_scale = q_quant |
| 720 | else: |
| 721 | q_values, q_scale = q_quant, None |
| 722 | return sparse_attn_indexer_kpool( |
| 723 | hidden_states, |
| 724 | self.k_cache.prefix, |
| 725 | self.k_cache.kv_cache, |
| 726 | q_values, |
| 727 | q_scale, |
| 728 | k, |
| 729 | weights, |
| 730 | self.quant_block_size, |
| 731 | self.scale_fmt, |
| 732 | self.topk_tokens, |
| 733 | self.head_dim, |
| 734 | self.max_pool_len, |
| 735 | self.max_total_seq_len, |
| 736 | self.topk_indices_buffer, |
| 737 | self.skip_k_cache_insert, |
| 738 | self.use_fp4_cache, |
| 739 | gate_score, |
| 740 | compress_ape, |
| 741 | index_kpool, |
| 742 | positions, |
| 743 | self.tail_cache.kv_cache if self.tail_cache is not None else None, |
| 744 | self.tail_cache.prefix if self.tail_cache is not None else None, |
| 745 | self.topk_backend, |
| 746 | ) |