| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Sparse attention indexer layers on aiter's paged MXFP4 kernels, for ROCm. |
| 4 | |
| 5 | With ``indexer_kv_dtype="mxfp4"`` on gfx950, the model writes the indexer K |
| 6 | cache in the order aiter's MQA-logits kernel reads (see |
| 7 | `rocm_paged_mxfp4_cache_layout`), and the kernel scores the cache in place. |
| 8 | Requires `RocmMxfp4IndexerMetadataBuilder` metadata. |
| 9 | """ |
| 10 | |
| 11 | import torch |
| 12 | from torch import nn |
| 13 | |
| 14 | from vllm.config import get_current_vllm_config |
| 15 | from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer |
| 16 | from vllm.v1.attention.ops.rocm_paged_mxfp4_indexer import ( |
| 17 | rocm_mxfp4_sparse_attn_indexer, |
| 18 | rocm_mxfp4_sparse_mqa_indexer, |
| 19 | ) |
| 20 | |
| 21 | |
| 22 | class RocmSparseAttnIndexer(SparseAttnIndexer): |
| 23 | """`SparseAttnIndexer` whose HIP path walks the paged MXFP4 cache with |
| 24 | aiter's kernel. A two-level indexer's source layer also publishes its |
| 25 | candidate pool here.""" |
| 26 | |
| 27 | def forward_hip( |
| 28 | self, |
| 29 | hidden_states: torch.Tensor, |
| 30 | q_quant: torch.Tensor | tuple[torch.Tensor, torch.Tensor], |
| 31 | k: torch.Tensor | None, |
| 32 | weights: torch.Tensor, |
| 33 | ): |
| 34 | assert isinstance(q_quant, tuple) and self.skip_k_cache_insert, ( |
| 35 | "the ROCm MXFP4 indexer takes (values, scales) Q and a K cache " |
| 36 | "the model already wrote" |
| 37 | ) |
| 38 | q_values, q_scale = q_quant |
| 39 | return rocm_mxfp4_sparse_attn_indexer( |
| 40 | hidden_states, |
| 41 | self.k_cache.prefix, |
| 42 | self.k_cache.kv_cache, |
| 43 | q_values, |
| 44 | q_scale, |
| 45 | weights, |
| 46 | self.topk_tokens, |
| 47 | self.head_dim, |
| 48 | self.max_model_len, |
| 49 | self.topk_indices_buffer, |
| 50 | compress_ratio=self.compress_ratio, |
| 51 | candidate_blocks=self.candidate_blocks, |
| 52 | candidate_block_size=self.candidate_block_size, |
| 53 | candidate_write=self.candidate_write, |
| 54 | ) |
| 55 | |
| 56 | |
| 57 | class RocmSparseMQAIndexer(nn.Module): |
| 58 | """DeepSeek-V4.1's candidate-consuming indexer on aiter's paged MXFP4 |
| 59 | MQA-logits kernel: the consumers score only the candidate pool the source |
| 60 | layer published, read straight from the paged cache. |
| 61 | |
| 62 | Built and called as `SparseMQAIndexer` is, so the model can take either. |
| 63 | """ |
| 64 | |
| 65 | weights_dtype = torch.float32 |
| 66 | """The scorer's weights and logits stay fp32: gfx950 has no packed bf16 |
| 67 | arithmetic, so its head reduce costs the same in either dtype.""" |
| 68 | |
| 69 | def __init__( |
| 70 | self, |
| 71 | k_cache, |
| 72 | topk_tokens: int, |
| 73 | head_dim: int, |
| 74 | max_total_seq_len: int, |
| 75 | topk_indices_buffer: torch.Tensor, |
| 76 | candidate_blocks: torch.Tensor, |
| 77 | candidate_block_size: int, |
| 78 | ): |
| 79 | super().__init__() |
| 80 | self.k_cache = k_cache |
| 81 | self.topk_tokens = topk_tokens |
| 82 | self.head_dim = head_dim |
| 83 | # aiter's kernel bounds a row by the model length in compressed |
| 84 | # positions, not by the prefill buffer SparseMQAIndexer sizes. |
| 85 | model_config = get_current_vllm_config().model_config |
| 86 | self.max_model_len = model_config.max_model_len // k_cache.compress_ratio |
| 87 | self.topk_indices_buffer = topk_indices_buffer |
| 88 | self.candidate_blocks = candidate_blocks |
| 89 | self.candidate_block_size = candidate_block_size |
| 90 | self.num_candidate_cols = candidate_blocks.shape[1] * candidate_block_size |
| 91 | |
| 92 | def forward( |
| 93 | self, |
| 94 | hidden_states: torch.Tensor, |
| 95 | q_quant: tuple[torch.Tensor, torch.Tensor], |
| 96 | k: torch.Tensor | None, |
| 97 | weights: torch.Tensor, |
| 98 | ) -> torch.Tensor: |
| 99 | assert k is None, "the model writes the indexer K cache" |
| 100 | q_values, q_scale = q_quant |
| 101 | return rocm_mxfp4_sparse_mqa_indexer( |
| 102 | hidden_states, |
| 103 | self.k_cache.prefix, |
| 104 | self.k_cache.kv_cache, |
| 105 | q_values, |
| 106 | q_scale, |
| 107 | weights, |
| 108 | self.topk_tokens, |
| 109 | self.head_dim, |
| 110 | self.max_model_len, |
| 111 | self.topk_indices_buffer, |
| 112 | self.k_cache.compress_ratio, |
| 113 | self.candidate_blocks, |
| 114 | self.candidate_block_size, |
| 115 | self.num_candidate_cols, |
| 116 | ) |