| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inkling short-conv state managed as a sliding-window KV cache. |
| 4 | |
| 5 | Each decoder layer owns one ``InklingConvState`` (an ``AttentionLayerBase``) that |
| 6 | emits a single ``SlidingWindowSpec`` for the layer's 4 sconv streams (K, V, |
| 7 | attn-output, mlp-output), packed head-major into one block: |
| 8 | |
| 9 | H = num_kv_heads (per-rank), N = block_size = sconv_kernel_size, |
| 10 | D = head_dim(K) + head_dim(V) + hidden/H(attn) + hidden/H(mlp) |
| 11 | |
| 12 | ``D`` is TP-invariant; per rank we store ``H/TP`` heads of width ``D``. The conv |
| 13 | reads/writes this cache out-of-band via a custom backend; the (smaller) conv page |
| 14 | is padded up to the uniform attention page by ``unify_kv_cache_spec_page_size``. |
| 15 | """ |
| 16 | |
| 17 | from __future__ import annotations |
| 18 | |
| 19 | from dataclasses import dataclass |
| 20 | from typing import ClassVar |
| 21 | |
| 22 | import torch |
| 23 | from torch import nn |
| 24 | |
| 25 | from vllm.config import VllmConfig, get_current_vllm_config |
| 26 | from vllm.distributed import get_tensor_model_parallel_world_size |
| 27 | from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase |
| 28 | from vllm.v1.attention.backend import ( |
| 29 | AttentionBackend, |
| 30 | AttentionCGSupport, |
| 31 | AttentionMetadata, |
| 32 | AttentionMetadataBuilder, |
| 33 | CommonAttentionMetadata, |
| 34 | ) |
| 35 | from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec, SlidingWindowSpec |
| 36 | |
| 37 | from .ops.sconv import sconv_seq_metadata |
| 38 | |
| 39 | # Stream order within the per-head packed D (== contiguous sub-ranges). |
| 40 | _K, _V, _ATTN, _MLP = 0, 1, 2, 3 |
| 41 | |
| 42 | |
| 43 | @dataclass |
| 44 | class InklingSconvMetadata(AttentionMetadata): |
| 45 | block_table: torch.Tensor # [num_reqs, max_blocks] physical blocks per req |
| 46 | slot_mapping: torch.Tensor # [T] int64 flat slot of each token (-1 => skip) |
| 47 | seq_idx: torch.Tensor # [T] int32 token -> batch request |
| 48 | query_start: torch.Tensor # [T] int32 first x-row of each token's request |
| 49 | |
| 50 | |
| 51 | class InklingSconvMetadataBuilder(AttentionMetadataBuilder[InklingSconvMetadata]): |
| 52 | _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH |
| 53 | |
| 54 | def __init__( |
| 55 | self, |
| 56 | kv_cache_spec: AttentionSpec, |
| 57 | layer_names: list[str], |
| 58 | vllm_config: VllmConfig, |
| 59 | device: torch.device, |
| 60 | ) -> None: |
| 61 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 62 | assert isinstance(kv_cache_spec, SlidingWindowSpec) |
| 63 | # Persistent per-token buffers for CUDA graph capture. |
| 64 | max_num_tokens = vllm_config.scheduler_config.max_num_batched_tokens |
| 65 | self.seq_idx_buffer = torch.empty( |
| 66 | max_num_tokens, dtype=torch.int32, device=device |
| 67 | ) |
| 68 | self.query_start_buffer = torch.empty( |
| 69 | max_num_tokens, dtype=torch.int32, device=device |
| 70 | ) |
| 71 | |
| 72 | def build( |
| 73 | self, |
| 74 | common_prefix_len: int, |
| 75 | common_attn_metadata: CommonAttentionMetadata, |
| 76 | fast_build: bool = False, |
| 77 | ) -> InklingSconvMetadata: |
| 78 | num_reqs = common_attn_metadata.num_reqs |
| 79 | num_actual_tokens = int(common_attn_metadata.query_start_loc_cpu[-1]) |
| 80 | num_padded_tokens = common_attn_metadata.slot_mapping.shape[0] |
| 81 | assert num_padded_tokens >= num_actual_tokens |
| 82 | |
| 83 | # Per-token seq_idx (owning request) and query_start (first x-row of |
| 84 | # that request; the fused kernel uses it to tell same-forward taps, |
| 85 | # read from x, from pre-forward taps, read from cache) in one launch. |
| 86 | sconv_seq_metadata( |
| 87 | common_attn_metadata.query_start_loc, |
| 88 | num_reqs, |
| 89 | num_actual_tokens, |
| 90 | self.seq_idx_buffer, |
| 91 | self.query_start_buffer, |
| 92 | num_padded_tokens, |
| 93 | ) |
| 94 | |
| 95 | return InklingSconvMetadata( |
| 96 | block_table=common_attn_metadata.block_table_tensor, |
| 97 | slot_mapping=common_attn_metadata.slot_mapping[:num_padded_tokens], |
| 98 | seq_idx=self.seq_idx_buffer[:num_padded_tokens], |
| 99 | query_start=self.query_start_buffer[:num_padded_tokens], |
| 100 | ) |
| 101 | |
| 102 | |
| 103 | class InklingSconvBackend(AttentionBackend): |
| 104 | """Custom dummy backend for the sconv sliding-window cache management.""" |
| 105 | |
| 106 | @staticmethod |
| 107 | def get_name() -> str: |
| 108 | return "INKLING_SCONV_SWA" |
| 109 | |
| 110 | @staticmethod |
| 111 | def get_impl_cls(): |
| 112 | raise NotImplementedError( |
| 113 | "InklingSconvBackend has no attention impl; the conv runs out-of-band." |
| 114 | ) |
| 115 | |
| 116 | @staticmethod |
| 117 | def get_builder_cls() -> type[InklingSconvMetadataBuilder]: |
| 118 | return InklingSconvMetadataBuilder |
| 119 | |
| 120 | |
| 121 | class InklingConvState(nn.Module, AttentionLayerBase): |
| 122 | """Per-decoder-layer owner emitting one sliding-window conv-state spec.""" |
| 123 | |
| 124 | def __init__( |
| 125 | self, |
| 126 | *, |
| 127 | num_kv_heads: int, |
| 128 | head_dim: int, |
| 129 | hidden_size: int, |
| 130 | kernel_size: int, |
| 131 | prefix: str, |
| 132 | ) -> None: |
| 133 | super().__init__() |
| 134 | self.prefix = prefix |
| 135 | # Bound to the manager-allocated paged cache by bind_kv_cache; a |
| 136 | # placeholder until then. Read out-of-band by InklingShortConv. |
| 137 | self.kv_cache = torch.tensor([]) |
| 138 | tp_size = get_tensor_model_parallel_world_size() |
| 139 | # Guardrails for the conv-state layout below; only these are exercised. |
| 140 | # tp_size <= num_kv_heads keeps >=1 whole KV head per rank (no |
| 141 | # replication/clamping), so the per-head width stays TP-invariant. |
| 142 | assert tp_size <= num_kv_heads, ( |
| 143 | f"sconv SWA cache supports tp_size <= num_kv_heads ({num_kv_heads}), " |
| 144 | f"got {tp_size}" |
| 145 | ) |
| 146 | # Per-rank head count; D is TP-invariant (K/V heads and the hidden |
| 147 | # chunk both scale 1/TP together). The attn-/mlp-output sconv streams |
| 148 | # are hidden-sharded: each rank owns its H/tp chunk (the sublayer |
| 149 | # outputs are reduce-scattered / all-gathered around the conv). |
| 150 | self.num_kv_heads = num_kv_heads // tp_size |
| 151 | hidden_per_head = hidden_size // num_kv_heads |
| 152 | # Packed per-head width: K + V + attn-output chunk + mlp-output chunk, |
| 153 | # padded to a power of two so every layer's conv page is the same size |
| 154 | # and an exact multiple of the attention page (the page unifier then |
| 155 | # scales attention block sizes instead of padding). |
| 156 | raw_head_size = 2 * head_dim + 2 * hidden_per_head |
| 157 | self.head_size = 1 << (raw_head_size - 1).bit_length() |
| 158 | self.sliding_window = kernel_size |
| 159 | self.block_size = kernel_size |
| 160 | # Per-head D-sub-range (offset, width) for each stream. Streams share |
| 161 | # the cache; each writes/reads its own width across all H heads. |
| 162 | self.stream_ranges: tuple[tuple[int, int], ...] = ( |
| 163 | (0, head_dim), # _K |
| 164 | (head_dim, head_dim), # _V |
| 165 | (2 * head_dim, hidden_per_head), # _ATTN |
| 166 | (2 * head_dim + hidden_per_head, hidden_per_head), # _MLP |
| 167 | ) |
| 168 | vllm_config = get_current_vllm_config() |
| 169 | self._dtype = vllm_config.model_config.dtype |
| 170 | assert self._dtype == torch.bfloat16, ( |
| 171 | f"sconv SWA cache supports bfloat16 only, got {self._dtype}" |
| 172 | ) |
| 173 | # Register in the forward context so the runner enumerates this owner as |
| 174 | # an attention-like layer (get_kv_cache_spec / get_attn_backend). |
| 175 | compilation_config = vllm_config.compilation_config |
| 176 | if prefix in compilation_config.static_forward_context: |
| 177 | raise ValueError(f"Duplicate layer name: {prefix}") |
| 178 | compilation_config.static_forward_context[prefix] = self |
| 179 | |
| 180 | def forward(self): ... |
| 181 | |
| 182 | def get_attn_backend(self) -> type[AttentionBackend]: |
| 183 | return InklingSconvBackend |
| 184 | |
| 185 | def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec: |
| 186 | return SlidingWindowSpec( |
| 187 | block_size=self.block_size, |
| 188 | num_kv_heads=self.num_kv_heads, |
| 189 | head_size=self.head_size, |
| 190 | head_size_v=0, # all 4 streams packed into head_size |
| 191 | dtype=self._dtype, |
| 192 | sliding_window=self.sliding_window, |
| 193 | ) |