| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from dataclasses import dataclass |
| 4 | from typing import Any |
| 5 | |
| 6 | import torch |
| 7 | |
| 8 | from vllm.model_executor.warmup.jit_warmup import kernel_launcher, zip_inputs |
| 9 | from vllm.model_executor.warmup.jit_warmup_triton_helper import ( |
| 10 | LaunchSpec, |
| 11 | TritonWarmupTensor, |
| 12 | VllmTritonJitKernel, |
| 13 | triton_scalar_specialization_rep, |
| 14 | ) |
| 15 | from vllm.triton_utils import tl, triton |
| 16 | from vllm.utils.math_utils import cdiv |
| 17 | |
| 18 | _DSPARK_SWA_INDEX_ALIGNMENT = 64 |
| 19 | |
| 20 | |
| 21 | def get_dspark_swa_index_width( |
| 22 | window_size: int, |
| 23 | num_speculative_tokens: int, |
| 24 | ) -> int: |
| 25 | """Return the padded width of non-causal DSpark SWA indices.""" |
| 26 | width = max(int(window_size), 0) + max(int(num_speculative_tokens), 0) |
| 27 | return cdiv(width, _DSPARK_SWA_INDEX_ALIGNMENT) * _DSPARK_SWA_INDEX_ALIGNMENT |
| 28 | |
| 29 | |
| 30 | class CompressedSlotMappingKernel( |
| 31 | VllmTritonJitKernel["CompressedSlotMappingKernel.CompileKey"] |
| 32 | ): |
| 33 | TRITON_BLOCK_SIZE = 1024 |
| 34 | |
| 35 | @dataclass(frozen=True) |
| 36 | class CompileKey: |
| 37 | compress_ratio: int |
| 38 | triton_block_size: int |
| 39 | block_size: int |
| 40 | |
| 41 | @staticmethod |
| 42 | @triton.jit(do_not_specialize=["block_table_stride"]) |
| 43 | def kernel( |
| 44 | # [num_tokens] |
| 45 | compressed_slot_mapping_ptr, |
| 46 | # [num_tokens] the tokens' own slots |
| 47 | slot_mapping_ptr, |
| 48 | # [num_reqs + 1] |
| 49 | query_start_loc_ptr, |
| 50 | # [num_reqs] |
| 51 | seq_lens_ptr, |
| 52 | # [num_reqs, max_num_blocks] |
| 53 | block_table_ptr, |
| 54 | block_table_stride, |
| 55 | block_size, |
| 56 | COMPRESS_RATIO: tl.constexpr, |
| 57 | PAD_ID: tl.constexpr, |
| 58 | TRITON_BLOCK_SIZE: tl.constexpr, |
| 59 | ): |
| 60 | batch_idx = tl.program_id(0) |
| 61 | |
| 62 | query_start = tl.load(query_start_loc_ptr + batch_idx) |
| 63 | query_end = tl.load(query_start_loc_ptr + batch_idx + 1) |
| 64 | query_len = query_end - query_start |
| 65 | |
| 66 | seq_len = tl.load(seq_lens_ptr + batch_idx) |
| 67 | start_pos = seq_len - query_len |
| 68 | |
| 69 | for i in range(0, query_len, TRITON_BLOCK_SIZE): |
| 70 | offset = i + tl.arange(0, TRITON_BLOCK_SIZE) |
| 71 | mask = offset < query_len |
| 72 | |
| 73 | pos = start_pos + i + tl.arange(0, TRITON_BLOCK_SIZE) |
| 74 | # A replayed token has a PAD slot (SWA bounded replay: its KV is |
| 75 | # cached already); its compressed KV is cached too, so PAD here. |
| 76 | slot = tl.load( |
| 77 | slot_mapping_ptr + query_start + offset, mask=mask, other=PAD_ID |
| 78 | ) |
| 79 | is_valid = ((pos + 1) % COMPRESS_RATIO == 0) & (slot != PAD_ID) |
| 80 | pos_after_compress = pos // COMPRESS_RATIO |
| 81 | |
| 82 | block_ids = pos_after_compress // block_size |
| 83 | block_numbers = tl.load( |
| 84 | block_table_ptr + batch_idx * block_table_stride + block_ids, |
| 85 | mask=mask & is_valid, |
| 86 | ) |
| 87 | slot_ids = block_numbers * block_size + pos_after_compress % block_size |
| 88 | |
| 89 | # NOTE |
| 90 | slot_ids = tl.where(is_valid, slot_ids, PAD_ID) |
| 91 | tl.store( |
| 92 | compressed_slot_mapping_ptr + query_start + offset, slot_ids, mask=mask |
| 93 | ) |
| 94 | |
| 95 | def dispatch( # type: ignore[override] |
| 96 | self, |
| 97 | *, |
| 98 | compress_ratio: int, |
| 99 | block_size: int, |
| 100 | ) -> CompileKey: |
| 101 | return self.CompileKey( |
| 102 | compress_ratio=compress_ratio, |
| 103 | triton_block_size=self.TRITON_BLOCK_SIZE, |
| 104 | block_size=triton_scalar_specialization_rep(block_size), |
| 105 | ) |
| 106 | |
| 107 | def get_warmup_keys(self, vllm_config: Any) -> list[CompileKey]: |
| 108 | hf_text_config = vllm_config.model_config.hf_text_config |
| 109 | configured_ratios = ( |
| 110 | *(getattr(hf_text_config, "compress_ratios", None) or ()), |
| 111 | getattr(hf_text_config, "index_kpool", 1) or 1, |
| 112 | ) |
| 113 | compress_ratios = tuple( |
| 114 | dict.fromkeys(int(ratio) for ratio in configured_ratios if int(ratio) > 1) |
| 115 | ) |
| 116 | if not compress_ratios: |
| 117 | return [] |
| 118 | return self._trace_dispatch(self.dispatch)( |
| 119 | zip_inputs( |
| 120 | *( |
| 121 | dict( |
| 122 | compress_ratio=ratio, |
| 123 | block_size=vllm_config.cache_config.block_size // ratio, |
| 124 | ) |
| 125 | for ratio in compress_ratios |
| 126 | ) |
| 127 | ) |
| 128 | ) |
| 129 | |
| 130 | def warmup_inputs(self, compile_key: CompileKey) -> dict[str, Any]: |
| 131 | int32_ptr = TritonWarmupTensor(torch.int32) |
| 132 | return dict( |
| 133 | compressed_slot_mapping=TritonWarmupTensor(torch.int64), |
| 134 | slot_mapping=TritonWarmupTensor(torch.int64), |
| 135 | query_start_loc=int32_ptr, |
| 136 | seq_lens=int32_ptr, |
| 137 | block_table=int32_ptr, |
| 138 | block_size=compile_key.block_size, |
| 139 | compress_ratio=compile_key.compress_ratio, |
| 140 | ) |
| 141 | |
| 142 | @kernel_launcher |
| 143 | def __call__( |
| 144 | self, |
| 145 | compressed_slot_mapping: torch.Tensor, |
| 146 | slot_mapping: torch.Tensor, |
| 147 | query_start_loc: torch.Tensor, |
| 148 | seq_lens: torch.Tensor, |
| 149 | block_table: torch.Tensor, |
| 150 | block_size: int, |
| 151 | compress_ratio: int, |
| 152 | ) -> LaunchSpec: |
| 153 | return (block_table.shape[0],), dict( |
| 154 | block_table_stride=block_table.stride(0), |
| 155 | COMPRESS_RATIO=compress_ratio, |
| 156 | PAD_ID=-1, |
| 157 | TRITON_BLOCK_SIZE=self.TRITON_BLOCK_SIZE, |
| 158 | ) |
| 159 | |
| 160 | |
| 161 | def get_compressed_slot_mapping( |
| 162 | num_tokens: int, |
| 163 | slot_mapping: torch.Tensor, |
| 164 | query_start_loc: torch.Tensor, |
| 165 | seq_lens: torch.Tensor, |
| 166 | block_table: torch.Tensor, |
| 167 | block_size: int, |
| 168 | compress_ratio: int, |
| 169 | out: torch.Tensor | None = None, |
| 170 | ) -> torch.Tensor: |
| 171 | """Slot mapping for writing the compressed states of ``num_tokens`` tokens. |
| 172 | |
| 173 | Every ``compress_ratio`` tokens share one compressed state, written by the |
| 174 | last of them: that token maps to the state's slot, the others to PAD. A |
| 175 | token whose own ``slot_mapping`` entry is PAD maps to PAD too: SWA bounded |
| 176 | replay recomputes tokens whose KV is cached already, and their compressed |
| 177 | states must not be rewritten either. |
| 178 | """ |
| 179 | if out is not None: |
| 180 | # Guard: for padded / invalid sequences. |
| 181 | # Negative positions produce bogus block indices that lead to illegal memory |
| 182 | # accesses inside the block_table load. |
| 183 | # NOTE: Fill -1 to the whole tensor, not just the first `num_tokens`. |
| 184 | out.fill_(-1) |
| 185 | compressed_slot_mapping = out[:num_tokens] |
| 186 | else: |
| 187 | compressed_slot_mapping = torch.full( |
| 188 | (num_tokens,), -1, dtype=torch.int64, device=query_start_loc.device |
| 189 | ) |
| 190 | |
| 191 | _COMPRESSED_SLOT_MAPPING_KERNEL( |
| 192 | compressed_slot_mapping, |
| 193 | slot_mapping, |
| 194 | query_start_loc, |
| 195 | seq_lens, |
| 196 | block_table, |
| 197 | block_size, |
| 198 | compress_ratio, |
| 199 | ) |
| 200 | return compressed_slot_mapping |
| 201 | |
| 202 | |
| 203 | _COMPRESSED_SLOT_MAPPING_KERNEL = CompressedSlotMappingKernel() |