| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inkling short-convolution kernels backed by a paged sliding-window state cache. |
| 4 | |
| 5 | Each layer's 4 conv streams (K, V, attn-output, mlp-output) share one paged KV |
| 6 | cache ``[num_blocks, H, N, D]`` (head-major; see ``sconv_swa_attn.py``). A stream |
| 7 | occupies the contiguous D-sub-range ``[off_s, off_s + ws)`` across all ``H`` |
| 8 | heads, so its flat per-token width is ``H * ws`` and that is the conv channel |
| 9 | dim. The cache stores the conv *input* at every absolute position. |
| 10 | |
| 11 | ``fused_sconv`` is the single-launch path used by the model: per token it writes |
| 12 | the current input to its slot and convolves the ``W`` taps ending at its |
| 13 | absolute position. A tap landing inside the current forward is read from the |
| 14 | immutable input ``x`` (row ``src - pos + pid_t``); only pre-forward taps are read |
| 15 | from the paged cache (window position ``src`` -> physical block via |
| 16 | ``block_table[req, src // N]``). The just-written slot is never read back this |
| 17 | step, so there is no write/read hazard within or across programs -- which is |
| 18 | why this needs no decode-vs-prefill split and is valid for prefill / decode / |
| 19 | spec alike. |
| 20 | |
| 21 | All kernels address the cache purely by ``(slot, absolute_position)`` and |
| 22 | allocate nothing inside the captured region; their grids depend only on the |
| 23 | token count (``fused_sconv`` on a fixed token/channel tiling), so the same |
| 24 | path replays correctly under eager, breakable PIECEWISE, and FULL cudagraphs |
| 25 | without any data-dependent shape or branch. |
| 26 | """ |
| 27 | |
| 28 | from __future__ import annotations |
| 29 | |
| 30 | import torch |
| 31 | |
| 32 | from vllm.triton_utils import tl, triton |
| 33 | |
| 34 | |
| 35 | @triton.jit |
| 36 | def _fused_sconv_kernel( |
| 37 | x_ptr, # [T, H*WS] head-major current-token inputs (also residual) |
| 38 | cache_ptr, # [num_blocks, H, N, D] paged (page-strided view) |
| 39 | weight_ptr, # [H*WS, W] |
| 40 | out_ptr, # [T, H*WS] |
| 41 | pos_ptr, # [T] int64 absolute position per token |
| 42 | seq_idx_ptr, # [T] int32 token -> batch request |
| 43 | slot_ptr, # [T] int64 flat slot (block*N + blk_off); < 0 => PAD |
| 44 | block_table_ptr, # [num_reqs, max_blocks] int32 block_table |
| 45 | qstart_ptr, # [T] int32 first x-row of the token's request |
| 46 | T, # num tokens |
| 47 | stride_x_t, |
| 48 | stride_c_blk, |
| 49 | stride_c_h, |
| 50 | stride_c_n, |
| 51 | stride_c_d, |
| 52 | stride_w_d, |
| 53 | stride_w_w, |
| 54 | stride_bt_r, |
| 55 | MAX_BLOCKS, |
| 56 | N, # block_size |
| 57 | W: tl.constexpr, |
| 58 | USE_SILU: tl.constexpr, |
| 59 | USE_RESIDUAL: tl.constexpr, |
| 60 | OFF_S: tl.constexpr, |
| 61 | WS: tl.constexpr, |
| 62 | H: tl.constexpr, |
| 63 | BT: tl.constexpr, |
| 64 | BLOCK_C: tl.constexpr, |
| 65 | ): |
| 66 | # Each program owns a [BT tokens, BLOCK_C channels] tile. The flat channel |
| 67 | # index packs all H heads head-major (head = c // WS, in-stream offset = |
| 68 | # c % WS), so one program spans heads -- no per-head launch and no |
| 69 | # next_power_of_2(WS) lane waste. |
| 70 | pid_t = tl.program_id(0) |
| 71 | pid_c = tl.program_id(1) |
| 72 | toff = pid_t * BT + tl.arange(0, BT) # [BT] token rows |
| 73 | coff = pid_c * BLOCK_C + tl.arange(0, BLOCK_C) # [BLOCK_C] flat channels |
| 74 | C = H * WS |
| 75 | t_mask = toff < T |
| 76 | c_mask = coff < C |
| 77 | head = tl.minimum(coff // WS, H - 1) # clamp keeps masked lanes in-buffer |
| 78 | cd = OFF_S + coff % WS # cache D-index of the channel's stream slot |
| 79 | |
| 80 | slot = tl.load(slot_ptr + toff, mask=t_mask, other=-1) # [BT] |
| 81 | valid = slot >= 0 |
| 82 | pos = tl.load(pos_ptr + toff, mask=t_mask, other=0) |
| 83 | req = tl.load(seq_idx_ptr + toff, mask=t_mask, other=0) |
| 84 | qstart = tl.load(qstart_ptr + toff, mask=t_mask, other=0) |
| 85 | |
| 86 | tc_mask = t_mask[:, None] & c_mask[None, :] |
| 87 | |
| 88 | # 1) Insert each token's input into its paged slot (skip PAD rows). |
| 89 | xv = tl.load(x_ptr + toff[:, None] * stride_x_t + coff[None, :], mask=tc_mask) |
| 90 | safe_slot = tl.maximum(slot, 0) |
| 91 | dst = ( |
| 92 | cache_ptr |
| 93 | + (safe_slot // N)[:, None] * stride_c_blk |
| 94 | + head[None, :] * stride_c_h |
| 95 | + (safe_slot % N)[:, None] * stride_c_n |
| 96 | + cd[None, :] * stride_c_d |
| 97 | ) |
| 98 | tl.store(dst, xv, mask=tc_mask & valid[:, None]) |
| 99 | |
| 100 | # 2) Convolve the W taps ending at each token's `pos`. Each tap is read from |
| 101 | # `x` if it falls inside this forward (row >= the request's first row), else |
| 102 | # from the paged cache. Exactly one source is unmasked per tap, so we sum both. |
| 103 | acc = tl.zeros([BT, BLOCK_C], dtype=tl.float32) |
| 104 | for iw in tl.static_range(W): |
| 105 | src = pos - (W - 1) + iw # [BT] absolute window position |
| 106 | row = toff - (W - 1) + iw # [BT] x-row of `src` (== src - pos + token) |
| 107 | in_win = valid & (src >= 0) |
| 108 | intra = in_win & (row >= qstart) |
| 109 | cached = in_win & (row < qstart) |
| 110 | # intra-forward tap: read the immutable input x (never the slot we just |
| 111 | # wrote), so there is no write/read hazard. |
| 112 | safe_row = tl.maximum(row, 0) |
| 113 | xt = tl.load( |
| 114 | x_ptr + safe_row[:, None] * stride_x_t + coff[None, :], |
| 115 | mask=c_mask[None, :] & intra[:, None], |
| 116 | other=0.0, |
| 117 | ).to(tl.float32) |
| 118 | # pre-forward tap: read from the paged cache via the block table. Clamp |
| 119 | # addressing terms; the load is masked off when out of window. |
| 120 | safe_src = tl.maximum(src, 0) |
| 121 | safe_lblk = tl.minimum(safe_src // N, MAX_BLOCKS - 1) |
| 122 | blk = tl.load( |
| 123 | block_table_ptr + req * stride_bt_r + safe_lblk, mask=cached, other=0 |
| 124 | ).to(tl.int64) |
| 125 | cbase = ( |
| 126 | cache_ptr |
| 127 | + blk[:, None] * stride_c_blk |
| 128 | + head[None, :] * stride_c_h |
| 129 | + (safe_src % N)[:, None] * stride_c_n |
| 130 | + cd[None, :] * stride_c_d |
| 131 | ) |
| 132 | cv = tl.load(cbase, mask=c_mask[None, :] & cached[:, None], other=0.0).to( |
| 133 | tl.float32 |
| 134 | ) |
| 135 | wv = tl.load( |
| 136 | weight_ptr + coff * stride_w_d + iw * stride_w_w, mask=c_mask, other=0.0 |
| 137 | ).to(tl.float32) |
| 138 | acc += (xt + cv) * wv[None, :] |
| 139 | |
| 140 | if USE_SILU: |
| 141 | acc = acc * tl.sigmoid(acc) |
| 142 | if USE_RESIDUAL: |
| 143 | acc += xv.to(tl.float32) |
| 144 | |
| 145 | tl.store( |
| 146 | out_ptr + toff[:, None] * stride_x_t + coff[None, :], |
| 147 | acc.to(out_ptr.dtype.element_ty), |
| 148 | mask=tc_mask, |
| 149 | ) |
| 150 | |
| 151 | |
| 152 | def fused_sconv( |
| 153 | x: torch.Tensor, # [T, H*ws] head-major current-token inputs |
| 154 | weight: torch.Tensor, # [H*ws, W] |
| 155 | cache: torch.Tensor, # [num_blocks, H, N, D] paged |
| 156 | positions: torch.Tensor, # [T] int64 absolute position per token |
| 157 | block_table: torch.Tensor, # [num_reqs, max_blocks] int32 |
| 158 | seq_idx: torch.Tensor, # [T] int32 token -> batch request |
| 159 | slot_mapping: torch.Tensor, # [T] int64 flat slot (PAD = -1 => skip) |
| 160 | query_start: torch.Tensor, # [T] int32 first x-row of the token's request |
| 161 | off_s: int, |
| 162 | ws: int, |
| 163 | block_size: int, |
| 164 | activation: str | None = None, |
| 165 | use_residual: bool = True, |
| 166 | ) -> torch.Tensor: |
| 167 | """Single-launch insert + depthwise causal conv1d over the paged cache. |
| 168 | |
| 169 | Reads same-forward taps from ``x`` and pre-forward taps from the cache, so |
| 170 | it is race-free in one launch for prefill / decode / spec and cudagraph-safe |
| 171 | under eager / piecewise / full capture. |
| 172 | """ |
| 173 | T = x.shape[0] |
| 174 | out = torch.empty_like(x) |
| 175 | if T == 0: |
| 176 | return out |
| 177 | assert x.is_contiguous() |
| 178 | assert cache.stride(3) == 1, "cache D-dim must be contiguous" |
| 179 | H = cache.shape[1] |
| 180 | W = weight.shape[1] |
| 181 | C = H * ws # flat conv channel dim (all heads, head-major) |
| 182 | # Tile BT tokens x BLOCK_C channels per program: enough work per CTA to |
| 183 | # amortize the per-token addressing, while keeping the grid large for |
| 184 | # prefill. BLOCK_C spans heads so there is no per-head launch. A ~2K-element |
| 185 | # tile at 4 warps measured best on Blackwell; larger tiles spill registers. |
| 186 | BLOCK_C = min(triton.next_power_of_2(C), 256) |
| 187 | BT = 8 |
| 188 | grid = (triton.cdiv(T, BT), triton.cdiv(C, BLOCK_C)) |
| 189 | _fused_sconv_kernel[grid]( |
| 190 | x, |
| 191 | cache, |
| 192 | weight, |
| 193 | out, |
| 194 | positions, |
| 195 | seq_idx, |
| 196 | slot_mapping, |
| 197 | block_table, |
| 198 | query_start, |
| 199 | T, |
| 200 | x.stride(0), |
| 201 | cache.stride(0), |
| 202 | cache.stride(1), |
| 203 | cache.stride(2), |
| 204 | cache.stride(3), |
| 205 | weight.stride(0), |
| 206 | weight.stride(1), |
| 207 | block_table.stride(0), |
| 208 | block_table.shape[1], |
| 209 | block_size, |
| 210 | W=W, |
| 211 | USE_SILU=activation in ("silu", "swish"), |
| 212 | USE_RESIDUAL=use_residual, |
| 213 | OFF_S=off_s, |
| 214 | WS=ws, |
| 215 | H=H, |
| 216 | BT=BT, |
| 217 | BLOCK_C=BLOCK_C, |
| 218 | num_warps=4, |
| 219 | ) |
| 220 | return out |
| 221 | |
| 222 | |
| 223 | @triton.jit |
| 224 | def _seq_metadata_kernel( |
| 225 | qsl_ptr, # [num_reqs + 1] int32 cumulative query start rows |
| 226 | seq_idx_ptr, # [T] int32 out: token -> owning request |
| 227 | query_start_ptr, # [T] int32 out: first x-row of the token's request |
| 228 | num_reqs, |
| 229 | num_actual_tokens, |
| 230 | num_padded_tokens, |
| 231 | n_iters, # ceil(log2(num_reqs)): binary-search depth |
| 232 | BLOCK: tl.constexpr, |
| 233 | ): |
| 234 | offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) |
| 235 | tok = offs.to(tl.int32) |
| 236 | # Largest j in [0, num_reqs) with qsl[j] <= tok. |
| 237 | lo = tl.zeros([BLOCK], tl.int32) |
| 238 | hi = tl.full([BLOCK], num_reqs - 1, tl.int32) |
| 239 | for _ in range(n_iters): |
| 240 | mid = (lo + hi + 1) // 2 |
| 241 | below = tl.load(qsl_ptr + mid) <= tok |
| 242 | lo = tl.where(below, mid, lo) |
| 243 | hi = tl.where(below, hi, mid - 1) |
| 244 | actual = offs < num_actual_tokens |
| 245 | padded = offs < num_padded_tokens |
| 246 | query_start = tl.load(qsl_ptr + lo) |
| 247 | tl.store(seq_idx_ptr + offs, tl.where(actual, lo, 0), mask=padded) |
| 248 | tl.store( |
| 249 | query_start_ptr + offs, |
| 250 | tl.where(actual, query_start, 0), |
| 251 | mask=padded, |
| 252 | ) |
| 253 | |
| 254 | |
| 255 | def sconv_seq_metadata( |
| 256 | query_start_loc: torch.Tensor, |
| 257 | num_reqs: int, |
| 258 | num_actual_tokens: int, |
| 259 | seq_idx_out: torch.Tensor, |
| 260 | query_start_out: torch.Tensor, |
| 261 | num_padded_tokens: int | None = None, |
| 262 | ) -> None: |
| 263 | """Fill static per-token seq_idx / query_start buffers in one launch. |
| 264 | |
| 265 | Replaces the arange + searchsorted + clamp + gather + 2x copy chain of the |
| 266 | sconv metadata build with a single kernel writing both persistent buffers. |
| 267 | Padded rows are filled with zero and must have ``slot_mapping == -1``. |
| 268 | """ |
| 269 | if num_padded_tokens is None: |
| 270 | num_padded_tokens = num_actual_tokens |
| 271 | if num_padded_tokens < num_actual_tokens: |
| 272 | raise ValueError("num_padded_tokens must cover all actual tokens") |
| 273 | if num_padded_tokens > seq_idx_out.shape[0]: |
| 274 | raise ValueError("seq_idx_out is too small for the padded token count") |
| 275 | if num_padded_tokens > query_start_out.shape[0]: |
| 276 | raise ValueError("query_start_out is too small for the padded token count") |
| 277 | |
| 278 | BLOCK = 256 |
| 279 | n_iters = (num_reqs - 1).bit_length() |
| 280 | grid = (triton.cdiv(num_padded_tokens, BLOCK),) |
| 281 | _seq_metadata_kernel[grid]( |
| 282 | query_start_loc, |
| 283 | seq_idx_out, |
| 284 | query_start_out, |
| 285 | num_reqs, |
| 286 | num_actual_tokens, |
| 287 | num_padded_tokens, |
| 288 | n_iters, |
| 289 | BLOCK=BLOCK, |
| 290 | ) |