| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Sparse Attention Indexer that scores only the candidate blocks. |
| 4 | |
| 5 | DeepSeek V4.1 two-level selection: the candidate-source indexer publishes |
| 6 | the top candidate blocks and later indexers pick their top-k inside them. |
| 7 | `SparseAttnIndexer` does that by computing dense logits over the whole |
| 8 | context and masking; this layer instead calls DeepGEMM's sparse MQA-logits |
| 9 | kernels on the candidate blocks only, so the work is O(candidate blocks) |
| 10 | instead of O(context). It requires the `DeepseekV41SparseIndexerBackend` |
| 11 | metadata (see ``AttentionConfig.indexer_sparse_logits``). |
| 12 | """ |
| 13 | |
| 14 | import torch |
| 15 | from torch import nn |
| 16 | |
| 17 | import vllm.envs as envs |
| 18 | from vllm import _custom_ops as ops |
| 19 | from vllm.forward_context import get_forward_context |
| 20 | from vllm.model_executor.kernels.attention.dsa.sparse_mqa_logits import ( |
| 21 | has_deep_select, |
| 22 | sparse_mqa_logits_paged_decode, |
| 23 | sparse_mqa_logits_prefill_chunk, |
| 24 | ) |
| 25 | from vllm.model_executor.layers.sparse_attn_indexer import ( |
| 26 | _gather_workspace_shapes, |
| 27 | kv_cache_as_quant_view, |
| 28 | ) |
| 29 | from vllm.platforms import current_platform |
| 30 | from vllm.utils.deep_gemm import has_deep_gemm_sparse_mqa |
| 31 | from vllm.v1.attention.backends.mla.indexer import ( |
| 32 | DeepseekV32IndexerPrefillChunkMetadata, |
| 33 | ) |
| 34 | from vllm.v1.attention.backends.mla.sparse_indexer import ( |
| 35 | DeepseekV41SparseIndexerMetadata, |
| 36 | ) |
| 37 | from vllm.v1.worker.workspace import current_workspace_manager |
| 38 | |
| 39 | |
| 40 | def _prefill_k_workspaces( |
| 41 | total_seq_lens: int, head_dim: int |
| 42 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 43 | """The packed MXFP4 K-gather workspace ``(values, scales)``, shared with |
| 44 | the dense indexer layers of the same model.""" |
| 45 | values_spec, scales_spec = _gather_workspace_shapes( |
| 46 | total_seq_lens, head_dim, current_platform.fp8_dtype(), use_fp4_cache=True |
| 47 | ) |
| 48 | k_quant, k_scale = current_workspace_manager().get_simultaneous( |
| 49 | values_spec, scales_spec |
| 50 | ) |
| 51 | return k_quant, k_scale |
| 52 | |
| 53 | |
| 54 | def _gather_prefill_chunk_k( |
| 55 | kv_cache: torch.Tensor, |
| 56 | k_quant_full: torch.Tensor, |
| 57 | k_scale_full: torch.Tensor, |
| 58 | chunk: DeepseekV32IndexerPrefillChunkMetadata, |
| 59 | ) -> tuple[torch.Tensor, torch.Tensor]: |
| 60 | """Gather one prefill chunk's paged K into the packed workspace.""" |
| 61 | assert chunk.local_cu_seq_lens is not None |
| 62 | k_quant = k_quant_full[: chunk.max_local_total_seq_lens] |
| 63 | k_scale = k_scale_full[: chunk.max_local_total_seq_lens] |
| 64 | if not chunk.skip_kv_gather and chunk.local_total_seq_lens > 0: |
| 65 | ops.cp_gather_indexer_k_quant_cache( |
| 66 | kv_cache, |
| 67 | k_quant, |
| 68 | k_scale, |
| 69 | chunk.block_table, |
| 70 | chunk.local_cu_seq_lens, |
| 71 | ) |
| 72 | return k_quant, k_scale |
| 73 | |
| 74 | |
| 75 | class SparseMQAIndexer(nn.Module): |
| 76 | """Candidate-consuming indexer on DeepGEMM's sparse MQA-logits kernels. |
| 77 | |
| 78 | Only valid for indexer layers that read candidate blocks with the MXFP4 |
| 79 | indexer cache on SM100. The K cache is written by the model before this |
| 80 | runs; ``forward`` takes the same arguments as `SparseAttnIndexer` so the |
| 81 | attention layer can call either. |
| 82 | """ |
| 83 | |
| 84 | weights_dtype = torch.bfloat16 |
| 85 | """Per-head weights dtype the sparse kernels take. The fused Q RoPE-quant |
| 86 | kernel writes it directly so no cast runs per step.""" |
| 87 | |
| 88 | def __init__( |
| 89 | self, |
| 90 | k_cache, |
| 91 | topk_tokens: int, |
| 92 | head_dim: int, |
| 93 | max_total_seq_len: int, |
| 94 | topk_indices_buffer: torch.Tensor, |
| 95 | candidate_blocks: torch.Tensor, |
| 96 | candidate_block_size: int, |
| 97 | ): |
| 98 | super().__init__() |
| 99 | if not ( |
| 100 | current_platform.is_cuda() |
| 101 | and current_platform.is_device_capability_family(100) |
| 102 | and has_deep_gemm_sparse_mqa() |
| 103 | ): |
| 104 | raise ValueError( |
| 105 | "SparseMQAIndexer requires an SM100-class GPU and DeepGEMM >= 2.8." |
| 106 | ) |
| 107 | if not has_deep_select(): |
| 108 | raise ValueError( |
| 109 | "SparseMQAIndexer requires the DeepSelect top-k extension " |
| 110 | "(vllm._deepselect_C)." |
| 111 | ) |
| 112 | self.k_cache = k_cache |
| 113 | self.topk_tokens = topk_tokens |
| 114 | self.head_dim = head_dim |
| 115 | self.max_total_seq_len = max_total_seq_len |
| 116 | self.topk_indices_buffer = topk_indices_buffer |
| 117 | self.candidate_blocks = candidate_blocks |
| 118 | self.candidate_block_size = candidate_block_size |
| 119 | |
| 120 | def _reserve_workspaces(self, device: torch.device) -> None: |
| 121 | """Profiling run: claim the K-gather workspace and the peak sparse |
| 122 | logits allocation so the memory estimate covers them.""" |
| 123 | _prefill_k_workspaces(self.max_total_seq_len, self.head_dim) |
| 124 | max_logits_bytes = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024 |
| 125 | torch.empty(max_logits_bytes, dtype=torch.uint8, device=device) |
| 126 | |
| 127 | def forward( |
| 128 | self, |
| 129 | hidden_states: torch.Tensor, |
| 130 | q_quant: tuple[torch.Tensor, torch.Tensor], |
| 131 | k: torch.Tensor | None, |
| 132 | weights: torch.Tensor, |
| 133 | ) -> torch.Tensor: |
| 134 | attn_metadata = get_forward_context().attn_metadata |
| 135 | if not isinstance(attn_metadata, dict): |
| 136 | self._reserve_workspaces(hidden_states.device) |
| 137 | return self.topk_indices_buffer |
| 138 | |
| 139 | metadata = attn_metadata[self.k_cache.prefix] |
| 140 | assert isinstance(metadata, DeepseekV41SparseIndexerMetadata), ( |
| 141 | "SparseMQAIndexer needs DeepseekV41SparseIndexerBackend metadata" |
| 142 | ) |
| 143 | assert k is None, "the model writes the indexer K cache" |
| 144 | q_values, q_scale = q_quant |
| 145 | kv_cache = self.k_cache.kv_cache |
| 146 | topk_tokens = self.topk_tokens |
| 147 | topk_indices_buffer = self.topk_indices_buffer |
| 148 | topk_indices_buffer[: hidden_states.shape[0]] = -1 |
| 149 | sparse_block_kv = metadata.sparse_block_kv |
| 150 | |
| 151 | if metadata.num_prefills > 0: |
| 152 | prefill = metadata.prefill |
| 153 | sparse_chunks = metadata.sparse_prefill |
| 154 | assert prefill is not None and sparse_chunks is not None |
| 155 | k_quant_full, k_scale_full = _prefill_k_workspaces( |
| 156 | self.max_total_seq_len, self.head_dim |
| 157 | ) |
| 158 | for chunk, rows in zip(prefill.chunks, sparse_chunks): |
| 159 | k_quant, k_scale = _gather_prefill_chunk_k( |
| 160 | kv_cache, k_quant_full, k_scale_full, chunk |
| 161 | ) |
| 162 | if chunk.local_total_seq_lens == 0: |
| 163 | continue # buffer already holds -1 |
| 164 | start, end = chunk.token_start, chunk.token_end |
| 165 | rows.kernel_metadata = sparse_mqa_logits_prefill_chunk( |
| 166 | q_values[start:end].view(torch.int8), |
| 167 | q_scale[start:end], |
| 168 | k_quant.view(torch.int8), |
| 169 | k_scale.view(torch.int32).squeeze(-1), |
| 170 | weights[start:end], |
| 171 | chunk.cu_seqlen_ks, |
| 172 | chunk.cu_seqlen_ke, |
| 173 | self.candidate_blocks[start:end], |
| 174 | self.candidate_block_size, |
| 175 | sparse_block_kv, |
| 176 | topk_tokens, |
| 177 | topk_indices_buffer[start:end, :topk_tokens], |
| 178 | sparse_indices=rows.sparse_indices, |
| 179 | end=rows.end, |
| 180 | col_indices=rows.col_indices, |
| 181 | kernel_metadata=rows.kernel_metadata, |
| 182 | ) |
| 183 | |
| 184 | if metadata.num_decodes > 0: |
| 185 | decode_rows = metadata.sparse_decode |
| 186 | assert decode_rows is not None and decode_rows.block_table is not None |
| 187 | assert decode_rows.row_indices is not None |
| 188 | num_rows = decode_rows.row_ke.shape[0] |
| 189 | if num_rows > 0: |
| 190 | decode_rows.kernel_metadata = sparse_mqa_logits_paged_decode( |
| 191 | q_values[:num_rows].view(torch.int8).unsqueeze(1), |
| 192 | q_scale[:num_rows].unsqueeze(1), |
| 193 | kv_cache_as_quant_view(kv_cache, self.head_dim, use_fp4_cache=True), |
| 194 | weights[:num_rows], |
| 195 | decode_rows.row_ke, |
| 196 | decode_rows.block_table, |
| 197 | decode_rows.row_indices, |
| 198 | self.candidate_blocks[:num_rows], |
| 199 | self.candidate_block_size, |
| 200 | sparse_block_kv, |
| 201 | topk_tokens, |
| 202 | topk_indices_buffer[:num_rows, :topk_tokens], |
| 203 | row_ks=decode_rows.row_ks, |
| 204 | sparse_indices=decode_rows.sparse_indices, |
| 205 | end=decode_rows.end, |
| 206 | col_indices=decode_rows.col_indices, |
| 207 | kernel_metadata=decode_rows.kernel_metadata, |
| 208 | ) |
| 209 | return topk_indices_buffer |