| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """ROCm entry point for the fused Kimi-K3 KDA chunk kernel. |
| 4 | |
| 5 | The kernel in ``csrc/libtorch_stable/kimi_k3/fused_kda_chunk_kernel_rocm.cu`` |
| 6 | replaces the chunk-state recurrence and the output GEMM of the Triton chunk |
| 7 | path with a single launch that keeps the per-chunk state in registers, so the |
| 8 | ``[chunks, H, V, K]`` state tensor and the recomputed values never reach HBM. |
| 9 | """ |
| 10 | |
| 11 | from functools import cache |
| 12 | |
| 13 | import torch |
| 14 | |
| 15 | from vllm.third_party.flash_linear_attention.ops.index import ( |
| 16 | prepare_chunk_indices, |
| 17 | prepare_chunk_offsets, |
| 18 | ) |
| 19 | from vllm.v1.attention.backends.utils import NULL_BLOCK_ID |
| 20 | |
| 21 | CHUNK_SIZE = 64 |
| 22 | HEAD_DIM = 128 |
| 23 | KDA_CHECKPOINT_ALIGNMENT = CHUNK_SIZE |
| 24 | # Chunk-group split policy constants; see _chunk_groups. |
| 25 | _BLOCKS_P1 = 2 # (kV + kK) / kBV |
| 26 | _BLOCKS_P2 = 1 # kV / kBV |
| 27 | _MIN_LEN = 4 |
| 28 | _DEEP_LEN = 12 |
| 29 | _MAX_GROUPS = 32 |
| 30 | |
| 31 | |
| 32 | def _hip_checkpoint_state_indices(state_indices: torch.Tensor) -> torch.Tensor: |
| 33 | """NULL_BLOCK_ID marks no checkpoint; the HIP kernel skips negative rows.""" |
| 34 | return ( |
| 35 | torch.where(state_indices != NULL_BLOCK_ID, state_indices, -1) |
| 36 | .to(torch.int32) |
| 37 | .contiguous() |
| 38 | ) |
| 39 | |
| 40 | |
| 41 | @cache |
| 42 | def is_fused_kda_chunk_supported() -> bool: |
| 43 | from vllm.platforms.rocm import on_gfx950 |
| 44 | |
| 45 | if not hasattr(torch.ops._C, "fused_kda_chunk"): |
| 46 | return False |
| 47 | # TODO: Verify on other archs, currently only validated on gfx950 |
| 48 | return on_gfx950() |
| 49 | |
| 50 | |
| 51 | def can_use_fused_kda_chunk( |
| 52 | head_k_dim: int, |
| 53 | head_v_dim: int, |
| 54 | dtype: torch.dtype, |
| 55 | chunk_size: int, |
| 56 | ) -> bool: |
| 57 | return ( |
| 58 | head_k_dim == HEAD_DIM |
| 59 | and head_v_dim == HEAD_DIM |
| 60 | and chunk_size == CHUNK_SIZE |
| 61 | and dtype == torch.bfloat16 |
| 62 | and is_fused_kda_chunk_supported() |
| 63 | ) |
| 64 | |
| 65 | |
| 66 | def fused_kda_prologue( |
| 67 | q: torch.Tensor, |
| 68 | k: torch.Tensor, |
| 69 | v: torch.Tensor, |
| 70 | raw_g: torch.Tensor, |
| 71 | raw_beta: torch.Tensor, |
| 72 | A_log: torch.Tensor, |
| 73 | dt_bias: torch.Tensor, |
| 74 | scale: float, |
| 75 | lower_bound: float, |
| 76 | cu_seqlens: torch.Tensor, |
| 77 | conv_weight: torch.Tensor | None = None, |
| 78 | conv_state: torch.Tensor | None = None, |
| 79 | conv_state_indices: torch.Tensor | None = None, |
| 80 | conv_has_initial_state: torch.Tensor | None = None, |
| 81 | chunk_indices: torch.Tensor | None = None, |
| 82 | ) -> dict[str, torch.Tensor]: |
| 83 | """Run the whole chunk-path prologue in one launch. |
| 84 | |
| 85 | Replaces the two L2 norms, the gate cumsum, both intra-chunk passes and the |
| 86 | w/u recompute. Returns the operands the fused chunk kernel consumes. |
| 87 | |
| 88 | The kernel can also apply the depthwise conv and its silu, in which case |
| 89 | ``q``/``k``/``v`` are the three raw bands of the QKV projection and |
| 90 | ``conv_state`` is updated in place. This path is currently not used. |
| 91 | """ |
| 92 | # `beta` and `g` reach the layer as last-dim slices of the fused QKVGFAB |
| 93 | # projection and carry its row stride. The kernel reads both with an |
| 94 | # explicit per-token stride, so neither is copied here. |
| 95 | _, t_total, num_heads, _ = q.shape |
| 96 | if chunk_indices is None: |
| 97 | chunk_indices = prepare_chunk_indices(cu_seqlens, CHUNK_SIZE) |
| 98 | chunk_indices = chunk_indices.to(torch.int32) |
| 99 | num_chunks = chunk_indices.shape[0] |
| 100 | dev = q.device |
| 101 | |
| 102 | # q may be a strided band view, so the workspaces are sized rather than |
| 103 | # cloned from it. |
| 104 | def _like(last: int) -> torch.Tensor: |
| 105 | return torch.empty(1, t_total, num_heads, last, dtype=q.dtype, device=q.device) |
| 106 | |
| 107 | ws = dict( |
| 108 | qg=_like(HEAD_DIM), |
| 109 | w=_like(HEAD_DIM), |
| 110 | u=_like(HEAD_DIM), |
| 111 | kg_t=torch.empty( |
| 112 | num_chunks, num_heads, HEAD_DIM, CHUNK_SIZE, dtype=q.dtype, device=dev |
| 113 | ), |
| 114 | aqk=torch.empty(1, t_total, num_heads, CHUNK_SIZE, dtype=q.dtype, device=dev), |
| 115 | decay=torch.empty( |
| 116 | num_chunks, num_heads, HEAD_DIM, dtype=torch.float32, device=dev |
| 117 | ), |
| 118 | ) |
| 119 | torch.ops._C.fused_kda_prologue( |
| 120 | q, |
| 121 | k, |
| 122 | v, |
| 123 | raw_g, |
| 124 | raw_beta, |
| 125 | A_log.reshape(-1), |
| 126 | dt_bias.reshape(-1), |
| 127 | ws["qg"], |
| 128 | ws["w"], |
| 129 | ws["u"], |
| 130 | ws["kg_t"], |
| 131 | ws["aqk"], |
| 132 | ws["decay"], |
| 133 | cu_seqlens.to(torch.int32), |
| 134 | chunk_indices, |
| 135 | conv_weight, |
| 136 | conv_state, |
| 137 | conv_state_indices, |
| 138 | conv_has_initial_state, |
| 139 | scale, |
| 140 | lower_bound, |
| 141 | ) |
| 142 | return ws |
| 143 | |
| 144 | |
| 145 | @cache |
| 146 | def _num_cus(device: int) -> int: |
| 147 | from vllm.platforms import current_platform |
| 148 | |
| 149 | return current_platform.get_device_properties(device).multi_processor_count |
| 150 | |
| 151 | |
| 152 | @cache |
| 153 | def _chunk_groups(chunks_per_seq: int, num_seqs: int, num_heads: int) -> int: |
| 154 | """How many parallel chunk groups to cut each sequence into. |
| 155 | |
| 156 | The scan composes the group operators with the exact transfer |
| 157 | ``M_g = prod_c (diag(d_c) - w_c^T kg_c)``, so the composition error does |
| 158 | not grow with group length and the choice is purely a machine fit: |
| 159 | |
| 160 | * ``fill1``/``fill2``: the largest G whose pass-one / pass-two grid still |
| 161 | fits in one scheduling batch. Rows are launched in blocks of at most |
| 162 | ``kBV = 128``, so pass one (``kV + kK`` rows) needs two blocks per |
| 163 | ``(n, h, g)`` and pass two one. Past a pass's fill point its block count |
| 164 | grows as fast as its depth shrinks and the rung stops paying. |
| 165 | * ``_DEEP_LEN``: unless the groups are still long, where pass two's halving |
| 166 | is worth letting pass one spill to a second batch. |
| 167 | * ``_MIN_LEN``: below this many chunks per group the per-group fixed cost |
| 168 | exceeds what the shorter walk saves. |
| 169 | |
| 170 | ``G == 2`` is depth-neutral -- two passes of ``nt/2`` -- and only doubles |
| 171 | staging, so the split is taken only from ``G >= 4``. |
| 172 | """ |
| 173 | cus = _num_cus(torch.accelerator.current_device_index()) |
| 174 | nh = max(num_seqs * num_heads, 1) |
| 175 | fill1 = cus // (_BLOCKS_P1 * nh) |
| 176 | fill2 = cus // (_BLOCKS_P2 * nh) |
| 177 | cand = fill1 |
| 178 | if fill2 > 0 and chunks_per_seq // fill2 >= _DEEP_LEN: |
| 179 | cand = fill2 |
| 180 | cand = min(cand, chunks_per_seq // _MIN_LEN, _MAX_GROUPS) |
| 181 | return cand if cand >= 4 else 1 |
| 182 | |
| 183 | |
| 184 | def _kda_group_workspace( |
| 185 | groups: int, nh: int, device: torch.device |
| 186 | ) -> torch.Tensor | None: |
| 187 | """One fp32 buffer the kernel carves into bg / sin_ / ag / mgT. |
| 188 | |
| 189 | Kept caller-owned so the memory stays inside the caching allocator and |
| 190 | vLLM's memory profiling accounts for it; the kernel's carve order must |
| 191 | match the layout here. |
| 192 | """ |
| 193 | if groups <= 1: |
| 194 | return None |
| 195 | planes = groups * nh |
| 196 | plane = HEAD_DIM * HEAD_DIM |
| 197 | floats = 2 * planes * plane + planes * HEAD_DIM # bg, sin_, ag |
| 198 | floats += (planes * plane + 1) // 2 # mgT, bf16 |
| 199 | return torch.empty(floats, dtype=torch.float32, device=device) |
| 200 | |
| 201 | |
| 202 | def fused_kda_chunk( |
| 203 | qg: torch.Tensor, |
| 204 | w: torch.Tensor, |
| 205 | u: torch.Tensor, |
| 206 | kg_t: torch.Tensor, |
| 207 | aqk: torch.Tensor, |
| 208 | decay: torch.Tensor, |
| 209 | out: torch.Tensor, |
| 210 | scale: float, |
| 211 | cu_seqlens: torch.Tensor, |
| 212 | initial_state: torch.Tensor | None, |
| 213 | output_final_state: bool, |
| 214 | chunk_offsets: torch.Tensor | None = None, |
| 215 | checkpoint_state: torch.Tensor | None = None, |
| 216 | checkpoint_offsets: torch.Tensor | None = None, |
| 217 | checkpoint_state_indices: torch.Tensor | None = None, |
| 218 | state_cache: torch.Tensor | None = None, |
| 219 | state_indices: torch.Tensor | None = None, |
| 220 | has_initial_state: torch.Tensor | None = None, |
| 221 | ) -> tuple[torch.Tensor, torch.Tensor | None]: |
| 222 | """Run the chunk recurrence and the output projection in one launch. |
| 223 | |
| 224 | Args: |
| 225 | qg: ``q * exp2(gk_cumsum)``, ``[1, T, H, 128]``. |
| 226 | w: chunk-local WY representation of the gated keys. |
| 227 | u: chunk-local WY representation of the values. |
| 228 | kg_t: chunk-major transposed gated keys, ``[chunks, H, 128, 64]``. |
| 229 | aqk: intra-chunk query-key attention, ``[chunks, H, 64, 64]``. |
| 230 | decay: ``exp2`` of each chunk's last gate row, ``[chunks, H, 128]``. |
| 231 | out: output buffer, ``[1, T, H, 128]``; may alias ``u``'s source. |
| 232 | scale: scale applied to the query-key products. |
| 233 | cu_seqlens: int32 cumulative sequence lengths. |
| 234 | initial_state: fp32 per-sequence initial recurrent state, or ``None``. |
| 235 | output_final_state: whether to return the final recurrent state. |
| 236 | chunk_offsets: int32 per-sequence first chunk index. |
| 237 | checkpoint_state: destination for the mid-prefill state snapshots, |
| 238 | fp32 ``[rows, H, 128, 128]``. Without |
| 239 | ``checkpoint_state_indices`` it is a staging buffer indexed by |
| 240 | sequence; with them it can be the paged state cache itself. |
| 241 | checkpoint_offsets: int32 ``[N]`` token offsets, each relative to its |
| 242 | sequence's first query token, at which to snapshot the recurrent |
| 243 | state. ``0`` disables the export for that sequence. A non-zero |
| 244 | offset must be a multiple of :data:`KDA_CHECKPOINT_ALIGNMENT`; |
| 245 | the kernel skips any that is not, so the caller has to filter |
| 246 | rather than rely on the store happening. |
| 247 | checkpoint_state_indices: optional int32 ``[N]`` destination row per |
| 248 | sequence. A negative row disables the export for that sequence. |
| 249 | state_cache: the paged recurrent state, fp32 |
| 250 | ``[slots, H, 128, 128]``. When given, the walk reads sequence |
| 251 | ``n``'s initial state from row ``state_indices[n]`` and writes its |
| 252 | final state back to the same row, which removes the gather and the |
| 253 | scatter that otherwise bracket this kernel. Replaces |
| 254 | ``initial_state`` / ``output_final_state``, and the returned final |
| 255 | state is then ``None`` because it is already in the cache. |
| 256 | state_indices: int32 ``[N]`` cache row per sequence. |
| 257 | has_initial_state: bool ``[N]``; a false entry starts that sequence |
| 258 | from a zero state and its cache row is read only, never before the |
| 259 | walk writes it. |
| 260 | |
| 261 | """ |
| 262 | if state_cache is not None: |
| 263 | if initial_state is not None or output_final_state: |
| 264 | raise ValueError("state_cache replaces initial_state/output_final_state") |
| 265 | if state_indices is None or has_initial_state is None: |
| 266 | raise ValueError("state_cache needs state_indices and has_initial_state") |
| 267 | num_seqs = cu_seqlens.numel() - 1 |
| 268 | final_state = None |
| 269 | if output_final_state: |
| 270 | final_state = torch.empty( |
| 271 | num_seqs, |
| 272 | u.shape[2], |
| 273 | u.shape[3], |
| 274 | qg.shape[3], |
| 275 | dtype=torch.float32, |
| 276 | device=u.device, |
| 277 | ) |
| 278 | cu_seqlens = cu_seqlens.to(torch.int32) |
| 279 | if chunk_offsets is None: |
| 280 | chunk_offsets = prepare_chunk_offsets(cu_seqlens, CHUNK_SIZE) |
| 281 | |
| 282 | # The chunk walk is serial in chunks. Splitting each sequence into `groups` |
| 283 | # stretches that run in parallel costs a second pass and a scan, and only |
| 284 | # pays once the chain is long enough to amortise both. |
| 285 | chunks_per_seq = kg_t.shape[0] // max(num_seqs, 1) |
| 286 | groups = _chunk_groups(chunks_per_seq, num_seqs, u.shape[2]) |
| 287 | group_state = _kda_group_workspace(groups, num_seqs * u.shape[2], u.device) |
| 288 | |
| 289 | torch.ops._C.fused_kda_chunk( |
| 290 | qg, |
| 291 | w, |
| 292 | u, |
| 293 | kg_t, |
| 294 | aqk, |
| 295 | decay, |
| 296 | initial_state, |
| 297 | final_state, |
| 298 | out, |
| 299 | cu_seqlens, |
| 300 | chunk_offsets, |
| 301 | scale, |
| 302 | group_state, |
| 303 | groups, |
| 304 | checkpoint_state, |
| 305 | None if checkpoint_offsets is None else checkpoint_offsets.to(torch.int32), |
| 306 | None |
| 307 | if checkpoint_state_indices is None |
| 308 | else _hip_checkpoint_state_indices(checkpoint_state_indices), |
| 309 | state_cache, |
| 310 | None if state_indices is None else state_indices.to(torch.int32).contiguous(), |
| 311 | has_initial_state, |
| 312 | ) |
| 313 | return out, final_state |