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