| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Kimi-K3 RecoverSSM speculative verify and accepted-state recovery.""" |
| 4 | |
| 5 | from collections.abc import Sequence |
| 6 | from dataclasses import dataclass |
| 7 | from typing import Any |
| 8 | |
| 9 | import torch |
| 10 | |
| 11 | from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first |
| 12 | from vllm.triton_utils import tl, triton |
| 13 | from vllm.v1.attention.backends.utils import NULL_BLOCK_ID |
| 14 | |
| 15 | |
| 16 | @triton.jit |
| 17 | def _kda_gate( |
| 18 | raw_g, |
| 19 | dt_bias, |
| 20 | A, |
| 21 | lower_bound, |
| 22 | USE_LOWER_BOUND: tl.constexpr, |
| 23 | ): |
| 24 | gate_input = raw_g + dt_bias |
| 25 | if USE_LOWER_BOUND: |
| 26 | return lower_bound * tl.sigmoid(A * gate_input) |
| 27 | softplus_gate = tl.where( |
| 28 | gate_input > 20.0, |
| 29 | gate_input, |
| 30 | tl.log(1.0 + tl.exp(gate_input)), |
| 31 | ) |
| 32 | return -A * softplus_gate |
| 33 | |
| 34 | |
| 35 | @triton.jit |
| 36 | def _kda_recurrent_step( |
| 37 | state, |
| 38 | k, |
| 39 | v, |
| 40 | raw_g, |
| 41 | raw_beta, |
| 42 | dt_bias, |
| 43 | A, |
| 44 | lower_bound, |
| 45 | USE_LOWER_BOUND: tl.constexpr, |
| 46 | ): |
| 47 | normalized_k = k * tl.rsqrt(tl.sum(k * k) + 1e-6) |
| 48 | gate = _kda_gate( |
| 49 | raw_g, |
| 50 | dt_bias, |
| 51 | A, |
| 52 | lower_bound, |
| 53 | USE_LOWER_BOUND, |
| 54 | ) |
| 55 | |
| 56 | state *= tl.exp(gate)[None, :] |
| 57 | correction = v - tl.sum(state * normalized_k[None, :], axis=1) |
| 58 | correction *= tl.sigmoid(raw_beta) |
| 59 | return state + correction[:, None] * normalized_k[None, :], correction |
| 60 | |
| 61 | |
| 62 | @triton.jit |
| 63 | def _kda_recoverssm_verify_kernel( |
| 64 | q_ptr, |
| 65 | k_ptr, |
| 66 | v_ptr, |
| 67 | raw_g_ptr, |
| 68 | raw_beta_ptr, |
| 69 | A_log_ptr, |
| 70 | dt_bias_ptr, |
| 71 | state_ptr, |
| 72 | correction_cache_ptr, |
| 73 | kg_cache_ptr, |
| 74 | out_ptr, |
| 75 | query_start_loc_ptr, |
| 76 | state_indices_ptr, |
| 77 | lower_bound, |
| 78 | null_block_id, |
| 79 | stride_q_token, |
| 80 | stride_k_token, |
| 81 | stride_v_token, |
| 82 | stride_g_token, |
| 83 | stride_beta_token, |
| 84 | stride_state_block, |
| 85 | stride_state_head, |
| 86 | stride_state_v, |
| 87 | stride_state_k, |
| 88 | stride_correction_block, |
| 89 | stride_correction_head, |
| 90 | stride_correction_pos, |
| 91 | stride_correction_dim, |
| 92 | stride_kg_block, |
| 93 | stride_kg_head, |
| 94 | stride_kg_pos, |
| 95 | stride_kg_dim, |
| 96 | stride_out_token, |
| 97 | stride_query_start_loc, |
| 98 | stride_state_indices, |
| 99 | K: tl.constexpr, |
| 100 | V: tl.constexpr, |
| 101 | BK: tl.constexpr, |
| 102 | BV: tl.constexpr, |
| 103 | SPEC_QUERY_LEN: tl.constexpr, |
| 104 | USE_LOWER_BOUND: tl.constexpr, |
| 105 | ): |
| 106 | pid_v = tl.program_id(0) |
| 107 | pid_b = tl.program_id(1) |
| 108 | pid_h = tl.program_id(2) |
| 109 | |
| 110 | bos = tl.load(query_start_loc_ptr + pid_b * stride_query_start_loc).to(tl.int64) |
| 111 | eos = tl.load(query_start_loc_ptr + (pid_b + 1) * stride_query_start_loc).to( |
| 112 | tl.int64 |
| 113 | ) |
| 114 | query_len = eos - bos |
| 115 | state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to(tl.int64) |
| 116 | |
| 117 | offs_k = tl.arange(0, BK) |
| 118 | offs_v = pid_v * BV + tl.arange(0, BV) |
| 119 | mask_k = offs_k < K |
| 120 | mask_v = offs_v < V |
| 121 | mask_state = mask_v[:, None] & mask_k[None, :] |
| 122 | |
| 123 | if state_idx <= null_block_id: |
| 124 | for token_offset in tl.static_range(SPEC_QUERY_LEN): |
| 125 | token_valid = token_offset < query_len |
| 126 | tl.store( |
| 127 | out_ptr + (bos + token_offset) * stride_out_token + pid_h * V + offs_v, |
| 128 | tl.zeros([BV], dtype=tl.float32), |
| 129 | mask=token_valid & mask_v, |
| 130 | ) |
| 131 | return |
| 132 | |
| 133 | state_ptrs = ( |
| 134 | state_ptr |
| 135 | + state_idx * stride_state_block |
| 136 | + pid_h * stride_state_head |
| 137 | + offs_v[:, None] * stride_state_v |
| 138 | + offs_k[None, :] * stride_state_k |
| 139 | ) |
| 140 | state = tl.load(state_ptrs, mask=mask_state, other=0.0).to(tl.float32) |
| 141 | A = tl.exp(tl.load(A_log_ptr + pid_h).to(tl.float32)) |
| 142 | |
| 143 | for token_offset in tl.static_range(SPEC_QUERY_LEN): |
| 144 | token_valid = token_offset < query_len |
| 145 | token = bos + token_offset |
| 146 | q = tl.load( |
| 147 | q_ptr + token * stride_q_token + pid_h * K + offs_k, |
| 148 | mask=token_valid & mask_k, |
| 149 | other=0.0, |
| 150 | ).to(tl.float32) |
| 151 | k = tl.load( |
| 152 | k_ptr + token * stride_k_token + pid_h * K + offs_k, |
| 153 | mask=token_valid & mask_k, |
| 154 | other=0.0, |
| 155 | ).to(tl.float32) |
| 156 | v = tl.load( |
| 157 | v_ptr + token * stride_v_token + pid_h * V + offs_v, |
| 158 | mask=token_valid & mask_v, |
| 159 | other=0.0, |
| 160 | ).to(tl.float32) |
| 161 | raw_g = tl.load( |
| 162 | raw_g_ptr + token * stride_g_token + pid_h * K + offs_k, |
| 163 | mask=token_valid & mask_k, |
| 164 | other=0.0, |
| 165 | ).to(tl.float32) |
| 166 | raw_beta = tl.load( |
| 167 | raw_beta_ptr + token * stride_beta_token + pid_h, |
| 168 | mask=token_valid, |
| 169 | other=0.0, |
| 170 | ).to(tl.float32) |
| 171 | |
| 172 | q *= tl.rsqrt(tl.sum(q * q) + 1e-6) * (K**-0.5) |
| 173 | dt_bias = tl.load(dt_bias_ptr + pid_h * K + offs_k, mask=mask_k, other=0.0).to( |
| 174 | tl.float32 |
| 175 | ) |
| 176 | updated_state, correction = _kda_recurrent_step( |
| 177 | state, |
| 178 | k, |
| 179 | v, |
| 180 | raw_g, |
| 181 | raw_beta, |
| 182 | dt_bias, |
| 183 | A, |
| 184 | lower_bound, |
| 185 | USE_LOWER_BOUND, |
| 186 | ) |
| 187 | state = tl.where(token_valid, updated_state, state) |
| 188 | |
| 189 | out = tl.sum(state * q[None, :], axis=1) |
| 190 | tl.store( |
| 191 | out_ptr + token * stride_out_token + pid_h * V + offs_v, |
| 192 | out, |
| 193 | mask=token_valid & mask_v, |
| 194 | ) |
| 195 | |
| 196 | correction_ptr = ( |
| 197 | correction_cache_ptr |
| 198 | + state_idx * stride_correction_block |
| 199 | + pid_h * stride_correction_head |
| 200 | + token_offset * stride_correction_pos |
| 201 | ) |
| 202 | tl.store( |
| 203 | correction_ptr + offs_v * stride_correction_dim, |
| 204 | correction, |
| 205 | mask=token_valid & mask_v, |
| 206 | ) |
| 207 | if pid_v == 0: |
| 208 | kg_ptr = ( |
| 209 | kg_cache_ptr |
| 210 | + state_idx * stride_kg_block |
| 211 | + pid_h * stride_kg_head |
| 212 | + token_offset * stride_kg_pos |
| 213 | ) |
| 214 | tl.store( |
| 215 | kg_ptr + offs_k * stride_kg_dim, |
| 216 | k, |
| 217 | mask=token_valid & mask_k, |
| 218 | ) |
| 219 | tl.store( |
| 220 | kg_ptr + (K + offs_k) * stride_kg_dim, |
| 221 | raw_g, |
| 222 | mask=token_valid & mask_k, |
| 223 | ) |
| 224 | |
| 225 | |
| 226 | @triton.heuristics( |
| 227 | { |
| 228 | "HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None, |
| 229 | "ALIGN_MODE": lambda args: args["block_table_ptr"] is not None, |
| 230 | } |
| 231 | ) |
| 232 | @triton.jit |
| 233 | def _prepare_commit_plan_kernel( |
| 234 | num_accepted_ptr, |
| 235 | request_indices_ptr, |
| 236 | state_indices_ptr, |
| 237 | query_start_loc_ptr, |
| 238 | block_table_ptr, |
| 239 | num_computed_ptr, |
| 240 | commit_lens_ptr, |
| 241 | final_state_indices_ptr, |
| 242 | boundary_state_indices_ptr, |
| 243 | boundary_recovery_lens_ptr, |
| 244 | null_block_id, |
| 245 | mamba_block_size, |
| 246 | block_table_width, |
| 247 | stride_num_accepted, |
| 248 | stride_request_indices, |
| 249 | stride_state_indices, |
| 250 | stride_query_start_loc, |
| 251 | stride_block_table_row, |
| 252 | stride_block_table_col, |
| 253 | stride_num_computed, |
| 254 | SPEC_QUERY_LEN: tl.constexpr, |
| 255 | HAS_REQUEST_INDICES: tl.constexpr, |
| 256 | ALIGN_MODE: tl.constexpr, |
| 257 | ): |
| 258 | spec_idx = tl.program_id(0) |
| 259 | source_state_idx = tl.load(state_indices_ptr + spec_idx * stride_state_indices).to( |
| 260 | tl.int64 |
| 261 | ) |
| 262 | request_idx = spec_idx |
| 263 | if HAS_REQUEST_INDICES: |
| 264 | request_idx = tl.load( |
| 265 | request_indices_ptr + spec_idx * stride_request_indices |
| 266 | ).to(tl.int64) |
| 267 | num_accepted = tl.load(num_accepted_ptr + request_idx * stride_num_accepted).to( |
| 268 | tl.int32 |
| 269 | ) |
| 270 | bos = tl.load(query_start_loc_ptr + spec_idx * stride_query_start_loc).to(tl.int64) |
| 271 | eos = tl.load(query_start_loc_ptr + (spec_idx + 1) * stride_query_start_loc).to( |
| 272 | tl.int64 |
| 273 | ) |
| 274 | query_len = (eos - bos).to(tl.int32) |
| 275 | commit_len = tl.minimum(tl.maximum(num_accepted, 0), query_len) |
| 276 | commit_len = tl.minimum(commit_len, SPEC_QUERY_LEN) |
| 277 | |
| 278 | final_state_idx = source_state_idx |
| 279 | boundary_state_idx = null_block_id |
| 280 | boundary_recovery_len = 0 |
| 281 | if ALIGN_MODE: |
| 282 | num_computed = tl.load(num_computed_ptr + request_idx * stride_num_computed).to( |
| 283 | tl.int32 |
| 284 | ) |
| 285 | final_num_computed = num_computed + commit_len |
| 286 | final_state_col = tl.minimum( |
| 287 | final_num_computed // mamba_block_size, block_table_width - 1 |
| 288 | ) |
| 289 | final_state_idx = tl.load( |
| 290 | block_table_ptr |
| 291 | + request_idx * stride_block_table_row |
| 292 | + final_state_col * stride_block_table_col |
| 293 | ).to(tl.int64) |
| 294 | next_boundary = (num_computed // mamba_block_size + 1) * mamba_block_size |
| 295 | crosses_boundary = final_num_computed >= next_boundary |
| 296 | boundary_recovery_len = next_boundary - num_computed |
| 297 | boundary_state_idx = tl.load( |
| 298 | block_table_ptr |
| 299 | + request_idx * stride_block_table_row |
| 300 | + (next_boundary // mamba_block_size - 1) * stride_block_table_col, |
| 301 | mask=crosses_boundary, |
| 302 | other=null_block_id, |
| 303 | ).to(tl.int64) |
| 304 | valid = (source_state_idx > null_block_id) & (commit_len > 0) |
| 305 | tl.store(commit_lens_ptr + spec_idx, tl.where(valid, commit_len, 0)) |
| 306 | tl.store( |
| 307 | final_state_indices_ptr + spec_idx, |
| 308 | tl.where(valid, final_state_idx, null_block_id), |
| 309 | ) |
| 310 | tl.store( |
| 311 | boundary_state_indices_ptr + spec_idx, |
| 312 | tl.where(valid, boundary_state_idx, null_block_id), |
| 313 | ) |
| 314 | tl.store( |
| 315 | boundary_recovery_lens_ptr + spec_idx, |
| 316 | tl.where(valid, boundary_recovery_len, 0), |
| 317 | ) |
| 318 | |
| 319 | |
| 320 | @triton.jit |
| 321 | def _compact_conv_state_kernel( |
| 322 | conv_state_ref_ptr, |
| 323 | conv_state_base_addrs_ptr, |
| 324 | conv_state_block_strides_ptr, |
| 325 | conv_state_dim_strides_ptr, |
| 326 | conv_state_token_strides_ptr, |
| 327 | state_indices_ptr, |
| 328 | commit_lens_ptr, |
| 329 | final_state_indices_ptr, |
| 330 | boundary_state_indices_ptr, |
| 331 | boundary_recovery_lens_ptr, |
| 332 | null_block_id, |
| 333 | conv_dim, |
| 334 | conv_history_len, |
| 335 | stride_state_indices, |
| 336 | BLOCK_D: tl.constexpr, |
| 337 | BLOCK_HISTORY: tl.constexpr, |
| 338 | ALIGN_MODE: tl.constexpr, |
| 339 | ): |
| 340 | pid_d = tl.program_id(0) |
| 341 | pid_b = tl.program_id(1) |
| 342 | pid_l = tl.program_id(2) |
| 343 | source_state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to( |
| 344 | tl.int64 |
| 345 | ) |
| 346 | if source_state_idx <= null_block_id: |
| 347 | return |
| 348 | |
| 349 | commit_len = tl.load(commit_lens_ptr + pid_b) |
| 350 | if commit_len == 0: |
| 351 | return |
| 352 | final_state_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64) |
| 353 | boundary_state_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) |
| 354 | boundary_recovery_len = tl.load(boundary_recovery_lens_ptr + pid_b) |
| 355 | |
| 356 | if final_state_idx <= null_block_id: |
| 357 | return |
| 358 | |
| 359 | base_addr = tl.load(conv_state_base_addrs_ptr + pid_l) |
| 360 | block_stride = tl.load(conv_state_block_strides_ptr + pid_l) |
| 361 | dim_stride = tl.load(conv_state_dim_strides_ptr + pid_l) |
| 362 | token_stride = tl.load(conv_state_token_strides_ptr + pid_l) |
| 363 | conv_state_ptr = base_addr.to(tl.pointer_type(conv_state_ref_ptr.dtype.element_ty)) |
| 364 | source_ptr = conv_state_ptr + source_state_idx * block_stride |
| 365 | final_ptr = conv_state_ptr + final_state_idx * block_stride |
| 366 | |
| 367 | offs_d = pid_d * BLOCK_D + tl.arange(0, BLOCK_D) |
| 368 | offs_h = tl.arange(0, BLOCK_HISTORY) |
| 369 | mask = (offs_d[:, None] < conv_dim) & (offs_h[None, :] < conv_history_len) |
| 370 | final_values = tl.load( |
| 371 | source_ptr |
| 372 | + offs_d[:, None] * dim_stride |
| 373 | + (commit_len - 1 + offs_h[None, :]) * token_stride, |
| 374 | mask=mask, |
| 375 | ) |
| 376 | if ALIGN_MODE: |
| 377 | boundary_values = tl.load( |
| 378 | source_ptr |
| 379 | + offs_d[:, None] * dim_stride |
| 380 | + (boundary_recovery_len - 1 + offs_h[None, :]) * token_stride, |
| 381 | mask=mask & (boundary_state_idx > null_block_id), |
| 382 | ) |
| 383 | boundary_ptr = conv_state_ptr + boundary_state_idx * block_stride |
| 384 | tl.store( |
| 385 | boundary_ptr |
| 386 | + offs_d[:, None] * dim_stride |
| 387 | + offs_h[None, :] * token_stride, |
| 388 | boundary_values, |
| 389 | mask=mask & (boundary_state_idx > null_block_id), |
| 390 | ) |
| 391 | tl.store( |
| 392 | final_ptr + offs_d[:, None] * dim_stride + offs_h[None, :] * token_stride, |
| 393 | final_values, |
| 394 | mask=mask, |
| 395 | ) |
| 396 | |
| 397 | |
| 398 | @triton.jit |
| 399 | def _commit_kda_state_kernel( |
| 400 | state_ref_ptr, |
| 401 | state_base_addrs_ptr, |
| 402 | state_block_strides_ptr, |
| 403 | correction_cache_ref_ptr, |
| 404 | correction_cache_base_addrs_ptr, |
| 405 | correction_cache_block_strides_ptr, |
| 406 | kg_cache_ref_ptr, |
| 407 | kg_cache_base_addrs_ptr, |
| 408 | kg_cache_block_strides_ptr, |
| 409 | A_log_ptr, |
| 410 | dt_bias_ptr, |
| 411 | state_indices_ptr, |
| 412 | commit_lens_ptr, |
| 413 | final_state_indices_ptr, |
| 414 | boundary_state_indices_ptr, |
| 415 | boundary_recovery_lens_ptr, |
| 416 | lower_bound, |
| 417 | null_block_id, |
| 418 | stride_state_head, |
| 419 | stride_state_v, |
| 420 | stride_state_k, |
| 421 | stride_correction_cache_head, |
| 422 | stride_correction_cache_pos, |
| 423 | stride_correction_cache_dim, |
| 424 | stride_kg_cache_head, |
| 425 | stride_kg_cache_pos, |
| 426 | stride_kg_cache_dim, |
| 427 | stride_A_layer, |
| 428 | stride_A_head, |
| 429 | stride_dt_bias_layer, |
| 430 | stride_dt_bias_head, |
| 431 | stride_dt_bias_dim, |
| 432 | stride_state_indices, |
| 433 | K: tl.constexpr, |
| 434 | V: tl.constexpr, |
| 435 | BK: tl.constexpr, |
| 436 | BV: tl.constexpr, |
| 437 | NUM_HEADS: tl.constexpr, |
| 438 | USE_LOWER_BOUND: tl.constexpr, |
| 439 | ALIGN_MODE: tl.constexpr, |
| 440 | ): |
| 441 | pid_v = tl.program_id(0) |
| 442 | pid_b = tl.program_id(1) |
| 443 | pid_lh = tl.program_id(2) |
| 444 | pid_l = pid_lh // NUM_HEADS |
| 445 | pid_h = pid_lh % NUM_HEADS |
| 446 | |
| 447 | source_state_idx = tl.load(state_indices_ptr + pid_b * stride_state_indices).to( |
| 448 | tl.int64 |
| 449 | ) |
| 450 | if source_state_idx <= null_block_id: |
| 451 | return |
| 452 | commit_len = tl.load(commit_lens_ptr + pid_b) |
| 453 | if commit_len == 0: |
| 454 | return |
| 455 | final_state_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64) |
| 456 | boundary_state_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) |
| 457 | boundary_recovery_len = tl.load(boundary_recovery_lens_ptr + pid_b) |
| 458 | |
| 459 | if final_state_idx <= null_block_id: |
| 460 | return |
| 461 | |
| 462 | state_base_addr = tl.load(state_base_addrs_ptr + pid_l) |
| 463 | state_block_stride = tl.load(state_block_strides_ptr + pid_l) |
| 464 | state_ptr = state_base_addr.to(tl.pointer_type(state_ref_ptr.dtype.element_ty)) |
| 465 | source_state_ptr = ( |
| 466 | state_ptr + source_state_idx * state_block_stride + pid_h * stride_state_head |
| 467 | ) |
| 468 | |
| 469 | correction_cache_base_addr = tl.load(correction_cache_base_addrs_ptr + pid_l) |
| 470 | correction_cache_block_stride = tl.load(correction_cache_block_strides_ptr + pid_l) |
| 471 | correction_cache_ptr = correction_cache_base_addr.to( |
| 472 | tl.pointer_type(correction_cache_ref_ptr.dtype.element_ty) |
| 473 | ) |
| 474 | correction_cache_ptr += ( |
| 475 | source_state_idx * correction_cache_block_stride |
| 476 | + pid_h * stride_correction_cache_head |
| 477 | ) |
| 478 | kg_cache_base_addr = tl.load(kg_cache_base_addrs_ptr + pid_l) |
| 479 | kg_cache_block_stride = tl.load(kg_cache_block_strides_ptr + pid_l) |
| 480 | kg_cache_ptr = kg_cache_base_addr.to( |
| 481 | tl.pointer_type(kg_cache_ref_ptr.dtype.element_ty) |
| 482 | ) |
| 483 | kg_cache_ptr += ( |
| 484 | source_state_idx * kg_cache_block_stride + pid_h * stride_kg_cache_head |
| 485 | ) |
| 486 | |
| 487 | offs_k = tl.arange(0, BK) |
| 488 | offs_v = pid_v * BV + tl.arange(0, BV) |
| 489 | mask_k = offs_k < K |
| 490 | mask_v = offs_v < V |
| 491 | mask_state = mask_v[:, None] & mask_k[None, :] |
| 492 | state_ptrs = ( |
| 493 | source_state_ptr |
| 494 | + offs_v[:, None] * stride_state_v |
| 495 | + offs_k[None, :] * stride_state_k |
| 496 | ) |
| 497 | initial_state = tl.load(state_ptrs, mask=mask_state, other=0.0).to(tl.float32) |
| 498 | A = tl.exp( |
| 499 | tl.load(A_log_ptr + pid_l * stride_A_layer + pid_h * stride_A_head).to( |
| 500 | tl.float32 |
| 501 | ) |
| 502 | ) |
| 503 | |
| 504 | dt_bias = tl.load( |
| 505 | dt_bias_ptr |
| 506 | + pid_l * stride_dt_bias_layer |
| 507 | + pid_h * stride_dt_bias_head |
| 508 | + offs_k * stride_dt_bias_dim, |
| 509 | mask=mask_k, |
| 510 | other=0.0, |
| 511 | ).to(tl.float32) |
| 512 | final_decay = tl.full([BK], 1.0, tl.float32) |
| 513 | final_correction = tl.zeros([BV, BK], tl.float32) |
| 514 | boundary_decay = tl.full([BK], 1.0, tl.float32) |
| 515 | boundary_correction = tl.zeros([BV, BK], tl.float32) |
| 516 | |
| 517 | for reverse_offset in range(commit_len): |
| 518 | token_offset = commit_len - reverse_offset - 1 |
| 519 | correction_ptr = ( |
| 520 | correction_cache_ptr + token_offset * stride_correction_cache_pos |
| 521 | ) |
| 522 | kg_ptr = kg_cache_ptr + token_offset * stride_kg_cache_pos |
| 523 | k = tl.load( |
| 524 | kg_ptr + offs_k * stride_kg_cache_dim, |
| 525 | mask=mask_k, |
| 526 | other=0.0, |
| 527 | ).to(tl.float32) |
| 528 | correction = tl.load( |
| 529 | correction_ptr + offs_v * stride_correction_cache_dim, |
| 530 | mask=mask_v, |
| 531 | other=0.0, |
| 532 | ).to(tl.float32) |
| 533 | raw_g = tl.load( |
| 534 | kg_ptr + (K + offs_k) * stride_kg_cache_dim, |
| 535 | mask=mask_k, |
| 536 | other=0.0, |
| 537 | ).to(tl.float32) |
| 538 | normalized_k = k * tl.rsqrt(tl.sum(k * k) + 1e-6) |
| 539 | gate = _kda_gate( |
| 540 | raw_g, |
| 541 | dt_bias, |
| 542 | A, |
| 543 | lower_bound, |
| 544 | USE_LOWER_BOUND, |
| 545 | ) |
| 546 | update = correction[:, None] * normalized_k[None, :] |
| 547 | decay = tl.exp(gate) |
| 548 | final_correction += update * final_decay[None, :] |
| 549 | final_decay *= decay |
| 550 | if ALIGN_MODE: |
| 551 | before_boundary = token_offset < boundary_recovery_len |
| 552 | boundary_correction += tl.where( |
| 553 | before_boundary, |
| 554 | update * boundary_decay[None, :], |
| 555 | 0.0, |
| 556 | ) |
| 557 | boundary_decay *= tl.where(before_boundary, decay, 1.0) |
| 558 | |
| 559 | state = initial_state * final_decay[None, :] + final_correction |
| 560 | if ALIGN_MODE: |
| 561 | boundary_ptrs = ( |
| 562 | state_ptr |
| 563 | + boundary_state_idx * state_block_stride |
| 564 | + pid_h * stride_state_head |
| 565 | + offs_v[:, None] * stride_state_v |
| 566 | + offs_k[None, :] * stride_state_k |
| 567 | ) |
| 568 | boundary_state = initial_state * boundary_decay[None, :] + boundary_correction |
| 569 | tl.store( |
| 570 | boundary_ptrs, |
| 571 | boundary_state, |
| 572 | mask=mask_state & (boundary_state_idx > null_block_id), |
| 573 | ) |
| 574 | |
| 575 | final_ptrs = ( |
| 576 | state_ptr |
| 577 | + final_state_idx * state_block_stride |
| 578 | + pid_h * stride_state_head |
| 579 | + offs_v[:, None] * stride_state_v |
| 580 | + offs_k[None, :] * stride_state_k |
| 581 | ) |
| 582 | tl.store(final_ptrs, state, mask=mask_state) |
| 583 | |
| 584 | |
| 585 | def kda_recoverssm_verify( |
| 586 | q: torch.Tensor, |
| 587 | k: torch.Tensor, |
| 588 | v: torch.Tensor, |
| 589 | raw_g: torch.Tensor, |
| 590 | raw_beta: torch.Tensor, |
| 591 | A_log: torch.Tensor, |
| 592 | dt_bias: torch.Tensor, |
| 593 | lower_bound: float | None, |
| 594 | checkpoint_state: torch.Tensor, |
| 595 | correction_cache: torch.Tensor, |
| 596 | kg_cache: torch.Tensor, |
| 597 | query_start_loc: torch.Tensor, |
| 598 | state_indices: torch.Tensor, |
| 599 | spec_query_len: int, |
| 600 | out: torch.Tensor | None = None, |
| 601 | ) -> torch.Tensor: |
| 602 | """Verify a KDA speculative window without modifying its checkpoint.""" |
| 603 | if q.ndim != 4 or q.shape[0] != 1: |
| 604 | raise ValueError("KDA RecoverSSM q must have shape [1, tokens, heads, dim]") |
| 605 | _, total_tokens, num_heads, key_dim = q.shape |
| 606 | value_dim = v.shape[-1] |
| 607 | if k.shape != q.shape or v.shape != (1, total_tokens, num_heads, value_dim): |
| 608 | raise ValueError("KDA RecoverSSM q, k, and v shapes are incompatible") |
| 609 | if raw_g.shape != q.shape or raw_beta.shape != (1, total_tokens, num_heads): |
| 610 | raise ValueError("KDA RecoverSSM gate or beta shape is incompatible") |
| 611 | if any(tensor.stride()[2:] != (key_dim, 1) for tensor in (q, k, raw_g)): |
| 612 | raise ValueError("KDA RecoverSSM q, k, and gate heads must be contiguous") |
| 613 | if v.stride()[2:] != (value_dim, 1) or raw_beta.stride(2) != 1: |
| 614 | raise ValueError("KDA RecoverSSM v and beta heads must be contiguous") |
| 615 | num_blocks = checkpoint_state.shape[0] |
| 616 | if checkpoint_state.shape[1:] != ( |
| 617 | num_heads, |
| 618 | value_dim, |
| 619 | key_dim, |
| 620 | ): |
| 621 | raise ValueError("KDA RecoverSSM checkpoint shape is incompatible") |
| 622 | expected_correction_shape = ( |
| 623 | num_blocks, |
| 624 | num_heads, |
| 625 | spec_query_len, |
| 626 | value_dim, |
| 627 | ) |
| 628 | if correction_cache.shape != expected_correction_shape: |
| 629 | raise ValueError( |
| 630 | f"KDA RecoverSSM correction buffer needs shape {expected_correction_shape}" |
| 631 | ) |
| 632 | expected_kg_shape = (num_blocks, num_heads, spec_query_len, 2 * key_dim) |
| 633 | if kg_cache.shape != expected_kg_shape: |
| 634 | raise ValueError( |
| 635 | f"KDA RecoverSSM key/gate buffer needs shape {expected_kg_shape}" |
| 636 | ) |
| 637 | if correction_cache.dtype != torch.float32: |
| 638 | raise ValueError("KDA RecoverSSM correction buffer must use float32") |
| 639 | if kg_cache.dtype != k.dtype: |
| 640 | raise ValueError("KDA RecoverSSM key/gate buffer must match activation dtype") |
| 641 | if A_log.shape != (num_heads,) or dt_bias.numel() != num_heads * key_dim: |
| 642 | raise ValueError("KDA RecoverSSM gate parameters are incompatible") |
| 643 | if not A_log.is_contiguous() or not dt_bias.is_contiguous(): |
| 644 | raise ValueError("KDA RecoverSSM gate parameters must be contiguous") |
| 645 | batch = state_indices.shape[0] |
| 646 | if query_start_loc.shape[0] != batch + 1: |
| 647 | raise ValueError("KDA RecoverSSM query metadata is incompatible") |
| 648 | if total_tokens > batch * spec_query_len: |
| 649 | raise ValueError( |
| 650 | "KDA RecoverSSM speculative decode input exceeds its activation capacity" |
| 651 | ) |
| 652 | if out is None: |
| 653 | out = torch.empty_like(v) |
| 654 | if out.shape != v.shape: |
| 655 | raise ValueError("KDA RecoverSSM output shape is incompatible") |
| 656 | if out.stride()[2:] != (value_dim, 1): |
| 657 | raise ValueError("KDA RecoverSSM output heads must be contiguous") |
| 658 | device = q.device |
| 659 | if any( |
| 660 | tensor.device != device |
| 661 | for tensor in ( |
| 662 | k, |
| 663 | v, |
| 664 | raw_g, |
| 665 | raw_beta, |
| 666 | A_log, |
| 667 | dt_bias, |
| 668 | checkpoint_state, |
| 669 | correction_cache, |
| 670 | kg_cache, |
| 671 | query_start_loc, |
| 672 | state_indices, |
| 673 | out, |
| 674 | ) |
| 675 | ): |
| 676 | raise ValueError("KDA RecoverSSM inputs must be on the same device") |
| 677 | if total_tokens == 0: |
| 678 | return out |
| 679 | |
| 680 | block_k = triton.next_power_of_2(key_dim) |
| 681 | block_v = min(triton.next_power_of_2(value_dim), 32) |
| 682 | grid = (triton.cdiv(value_dim, block_v), batch, num_heads) |
| 683 | _kda_recoverssm_verify_kernel[grid]( |
| 684 | q, |
| 685 | k, |
| 686 | v, |
| 687 | raw_g, |
| 688 | raw_beta, |
| 689 | A_log, |
| 690 | dt_bias, |
| 691 | checkpoint_state, |
| 692 | correction_cache, |
| 693 | kg_cache, |
| 694 | out, |
| 695 | query_start_loc, |
| 696 | state_indices, |
| 697 | lower_bound or 0.0, |
| 698 | NULL_BLOCK_ID, |
| 699 | q.stride(1), |
| 700 | k.stride(1), |
| 701 | v.stride(1), |
| 702 | raw_g.stride(1), |
| 703 | raw_beta.stride(1), |
| 704 | checkpoint_state.stride(0), |
| 705 | checkpoint_state.stride(1), |
| 706 | checkpoint_state.stride(2), |
| 707 | checkpoint_state.stride(3), |
| 708 | correction_cache.stride(0), |
| 709 | correction_cache.stride(1), |
| 710 | correction_cache.stride(2), |
| 711 | correction_cache.stride(3), |
| 712 | kg_cache.stride(0), |
| 713 | kg_cache.stride(1), |
| 714 | kg_cache.stride(2), |
| 715 | kg_cache.stride(3), |
| 716 | out.stride(1), |
| 717 | query_start_loc.stride(0), |
| 718 | state_indices.stride(0), |
| 719 | K=key_dim, |
| 720 | V=value_dim, |
| 721 | BK=block_k, |
| 722 | BV=block_v, |
| 723 | SPEC_QUERY_LEN=spec_query_len, |
| 724 | USE_LOWER_BOUND=lower_bound is not None, |
| 725 | num_warps=4, |
| 726 | num_stages=2, |
| 727 | ) |
| 728 | return out |
| 729 | |
| 730 | |
| 731 | @dataclass |
| 732 | class KDARecoverSSMCommitContext: |
| 733 | conv_states: tuple[torch.Tensor, ...] |
| 734 | conv_state_base_addrs: torch.Tensor |
| 735 | conv_state_block_strides: torch.Tensor |
| 736 | conv_state_dim_strides: torch.Tensor |
| 737 | conv_state_token_strides: torch.Tensor |
| 738 | conv_history_len: int |
| 739 | checkpoints: tuple[torch.Tensor, ...] |
| 740 | state_base_addrs: torch.Tensor |
| 741 | state_block_strides: torch.Tensor |
| 742 | correction_caches: tuple[torch.Tensor, ...] |
| 743 | correction_cache_base_addrs: torch.Tensor |
| 744 | correction_cache_block_strides: torch.Tensor |
| 745 | kg_caches: tuple[torch.Tensor, ...] |
| 746 | kg_cache_base_addrs: torch.Tensor |
| 747 | kg_cache_block_strides: torch.Tensor |
| 748 | commit_lens: torch.Tensor |
| 749 | final_state_indices: torch.Tensor |
| 750 | boundary_state_indices: torch.Tensor |
| 751 | boundary_recovery_lens: torch.Tensor |
| 752 | A_log: torch.Tensor |
| 753 | dt_bias: torch.Tensor |
| 754 | lower_bound: float | None |
| 755 | spec_query_len: int |
| 756 | |
| 757 | @classmethod |
| 758 | def create( |
| 759 | cls, |
| 760 | layers: Sequence[Any], |
| 761 | *, |
| 762 | spec_query_len: int, |
| 763 | max_num_reqs: int, |
| 764 | ) -> "KDARecoverSSMCommitContext": |
| 765 | if not layers: |
| 766 | raise ValueError("KDA RecoverSSM commit requires at least one layer") |
| 767 | if any(len(layer.kv_cache) != 4 for layer in layers): |
| 768 | raise ValueError( |
| 769 | "KDA RecoverSSM pages must contain conv, state, correction, " |
| 770 | "and key/gate" |
| 771 | ) |
| 772 | |
| 773 | conv_states = [layer.kv_cache[0] for layer in layers] |
| 774 | if not is_conv_state_dim_first(): |
| 775 | conv_states = [state.transpose(-1, -2) for state in conv_states] |
| 776 | checkpoints = [layer.kv_cache[1] for layer in layers] |
| 777 | correction_caches = [layer.kv_cache[2] for layer in layers] |
| 778 | kg_caches = [layer.kv_cache[3] for layer in layers] |
| 779 | A_log = [layer.A_log for layer in layers] |
| 780 | dt_bias = [ |
| 781 | layer.dt_bias.view(layer.local_num_heads, layer.head_dim) |
| 782 | for layer in layers |
| 783 | ] |
| 784 | lower_bounds = {layer.gate_lower_bound for layer in layers} |
| 785 | if len(lower_bounds) != 1: |
| 786 | raise ValueError("KDA RecoverSSM layers need matching gate bounds") |
| 787 | |
| 788 | state_ref = checkpoints[0] |
| 789 | if state_ref.ndim != 4: |
| 790 | raise ValueError("KDA RecoverSSM checkpoint must be four-dimensional") |
| 791 | num_blocks, num_heads, value_dim, key_dim = state_ref.shape |
| 792 | for state in checkpoints: |
| 793 | if ( |
| 794 | state.shape != state_ref.shape |
| 795 | or state.dtype != state_ref.dtype |
| 796 | or state.device != state_ref.device |
| 797 | or state.stride()[1:] != state_ref.stride()[1:] |
| 798 | ): |
| 799 | raise ValueError( |
| 800 | "KDA RecoverSSM layers need matching checkpoint layout" |
| 801 | ) |
| 802 | expected_correction_shape = ( |
| 803 | num_blocks, |
| 804 | num_heads, |
| 805 | spec_query_len, |
| 806 | value_dim, |
| 807 | ) |
| 808 | correction_ref = correction_caches[0] |
| 809 | for correction_cache in correction_caches: |
| 810 | if ( |
| 811 | correction_cache.shape != expected_correction_shape |
| 812 | or correction_cache.dtype != torch.float32 |
| 813 | or correction_cache.device != state_ref.device |
| 814 | or correction_cache.stride()[1:] != correction_ref.stride()[1:] |
| 815 | ): |
| 816 | raise ValueError( |
| 817 | "KDA RecoverSSM correction buffers need float32 shape " |
| 818 | f"{expected_correction_shape}" |
| 819 | ) |
| 820 | expected_kg_shape = (num_blocks, num_heads, spec_query_len, 2 * key_dim) |
| 821 | kg_ref = kg_caches[0] |
| 822 | for kg_cache in kg_caches: |
| 823 | if ( |
| 824 | kg_cache.shape != expected_kg_shape |
| 825 | or kg_cache.dtype != kg_ref.dtype |
| 826 | or kg_cache.device != state_ref.device |
| 827 | or kg_cache.stride()[1:] != kg_ref.stride()[1:] |
| 828 | ): |
| 829 | raise ValueError( |
| 830 | f"KDA RecoverSSM key/gate buffers need shape {expected_kg_shape}" |
| 831 | ) |
| 832 | if any(param.shape != (num_heads,) for param in A_log): |
| 833 | raise ValueError("KDA RecoverSSM A_log shape is incompatible") |
| 834 | if any(param.shape != (num_heads, key_dim) for param in dt_bias): |
| 835 | raise ValueError("KDA RecoverSSM dt_bias shape is incompatible") |
| 836 | |
| 837 | conv_ref = conv_states[0] |
| 838 | if conv_ref.ndim != 3: |
| 839 | raise ValueError("KDA RecoverSSM conv state must be three-dimensional") |
| 840 | conv_dim, conv_state_len = conv_ref.shape[1:] |
| 841 | conv_history_len = conv_state_len - spec_query_len + 1 |
| 842 | if conv_history_len <= 0: |
| 843 | raise ValueError("KDA RecoverSSM conv state is shorter than its window") |
| 844 | for conv_state in conv_states: |
| 845 | if ( |
| 846 | conv_state.shape != conv_ref.shape |
| 847 | or conv_state.dtype != conv_ref.dtype |
| 848 | or conv_state.device != state_ref.device |
| 849 | or conv_state.shape[0] != num_blocks |
| 850 | ): |
| 851 | raise ValueError("KDA RecoverSSM layers need matching conv state") |
| 852 | |
| 853 | device = state_ref.device |
| 854 | |
| 855 | def _base_addrs(tensors: Sequence[torch.Tensor]) -> torch.Tensor: |
| 856 | return torch.tensor( |
| 857 | [tensor.data_ptr() for tensor in tensors], |
| 858 | dtype=torch.int64, |
| 859 | device=device, |
| 860 | ) |
| 861 | |
| 862 | def _block_strides(tensors: Sequence[torch.Tensor]) -> torch.Tensor: |
| 863 | return torch.tensor( |
| 864 | [tensor.stride(0) for tensor in tensors], |
| 865 | dtype=torch.int64, |
| 866 | device=device, |
| 867 | ) |
| 868 | |
| 869 | return cls( |
| 870 | conv_states=tuple(conv_states), |
| 871 | conv_state_base_addrs=_base_addrs(conv_states), |
| 872 | conv_state_block_strides=_block_strides(conv_states), |
| 873 | conv_state_dim_strides=torch.tensor( |
| 874 | [state.stride(1) for state in conv_states], |
| 875 | dtype=torch.int64, |
| 876 | device=device, |
| 877 | ), |
| 878 | conv_state_token_strides=torch.tensor( |
| 879 | [state.stride(2) for state in conv_states], |
| 880 | dtype=torch.int64, |
| 881 | device=device, |
| 882 | ), |
| 883 | conv_history_len=conv_history_len, |
| 884 | checkpoints=tuple(checkpoints), |
| 885 | state_base_addrs=_base_addrs(checkpoints), |
| 886 | state_block_strides=_block_strides(checkpoints), |
| 887 | correction_caches=tuple(correction_caches), |
| 888 | correction_cache_base_addrs=_base_addrs(correction_caches), |
| 889 | correction_cache_block_strides=_block_strides(correction_caches), |
| 890 | kg_caches=tuple(kg_caches), |
| 891 | kg_cache_base_addrs=_base_addrs(kg_caches), |
| 892 | kg_cache_block_strides=_block_strides(kg_caches), |
| 893 | commit_lens=torch.empty(max_num_reqs, dtype=torch.int32, device=device), |
| 894 | final_state_indices=torch.empty( |
| 895 | max_num_reqs, dtype=torch.int32, device=device |
| 896 | ), |
| 897 | boundary_state_indices=torch.empty( |
| 898 | max_num_reqs, dtype=torch.int32, device=device |
| 899 | ), |
| 900 | boundary_recovery_lens=torch.empty( |
| 901 | max_num_reqs, dtype=torch.int32, device=device |
| 902 | ), |
| 903 | A_log=torch.stack(tuple(A_log)).contiguous(), |
| 904 | dt_bias=torch.stack(tuple(dt_bias)).contiguous(), |
| 905 | lower_bound=lower_bounds.pop(), |
| 906 | spec_query_len=spec_query_len, |
| 907 | ) |
| 908 | |
| 909 | def commit( |
| 910 | self, |
| 911 | num_accepted_tokens: torch.Tensor, |
| 912 | state_indices: torch.Tensor, |
| 913 | query_start_loc: torch.Tensor, |
| 914 | request_indices: torch.Tensor | None = None, |
| 915 | block_table: torch.Tensor | None = None, |
| 916 | num_computed_tokens: torch.Tensor | None = None, |
| 917 | mamba_block_size: int | None = None, |
| 918 | ) -> None: |
| 919 | """Fold accepted KDA and convolution inputs into every layer.""" |
| 920 | batch = state_indices.shape[0] |
| 921 | if batch == 0: |
| 922 | return |
| 923 | if batch > self.commit_lens.shape[0]: |
| 924 | raise ValueError("KDA RecoverSSM commit batch exceeds its plan capacity") |
| 925 | if query_start_loc.shape[0] != batch + 1: |
| 926 | raise ValueError("KDA RecoverSSM commit metadata is incompatible") |
| 927 | if request_indices is not None and request_indices.shape[0] < batch: |
| 928 | raise ValueError("KDA RecoverSSM request mapping is too short") |
| 929 | align_args = (block_table, num_computed_tokens, mamba_block_size) |
| 930 | if any(arg is not None for arg in align_args) and any( |
| 931 | arg is None for arg in align_args |
| 932 | ): |
| 933 | raise ValueError("KDA RecoverSSM align metadata is incomplete") |
| 934 | if mamba_block_size is not None and mamba_block_size < self.spec_query_len: |
| 935 | raise ValueError( |
| 936 | "KDA RecoverSSM align block size must cover one speculative window" |
| 937 | ) |
| 938 | if block_table is not None and block_table.ndim != 2: |
| 939 | raise ValueError("KDA RecoverSSM block table must be two-dimensional") |
| 940 | device = self.checkpoints[0].device |
| 941 | if ( |
| 942 | any( |
| 943 | tensor.device != device |
| 944 | for tensor in ( |
| 945 | num_accepted_tokens, |
| 946 | state_indices, |
| 947 | query_start_loc, |
| 948 | ) |
| 949 | ) |
| 950 | or (request_indices is not None and request_indices.device != device) |
| 951 | or (block_table is not None and block_table.device != device) |
| 952 | or ( |
| 953 | num_computed_tokens is not None and num_computed_tokens.device != device |
| 954 | ) |
| 955 | ): |
| 956 | raise ValueError("KDA RecoverSSM commit inputs must be on the same device") |
| 957 | |
| 958 | block_table_stride = (0, 0) if block_table is None else block_table.stride() |
| 959 | num_computed_stride = ( |
| 960 | 0 if num_computed_tokens is None else num_computed_tokens.stride(0) |
| 961 | ) |
| 962 | |
| 963 | num_layers = len(self.checkpoints) |
| 964 | conv_ref = self.conv_states[0] |
| 965 | conv_dim = conv_ref.shape[1] |
| 966 | block_history = triton.next_power_of_2(self.conv_history_len) |
| 967 | _prepare_commit_plan_kernel[(batch,)]( |
| 968 | num_accepted_tokens, |
| 969 | request_indices, |
| 970 | state_indices, |
| 971 | query_start_loc, |
| 972 | block_table, |
| 973 | num_computed_tokens, |
| 974 | self.commit_lens, |
| 975 | self.final_state_indices, |
| 976 | self.boundary_state_indices, |
| 977 | self.boundary_recovery_lens, |
| 978 | NULL_BLOCK_ID, |
| 979 | mamba_block_size or 1, |
| 980 | block_table.shape[1] if block_table is not None else 1, |
| 981 | num_accepted_tokens.stride(0), |
| 982 | request_indices.stride(0) if request_indices is not None else 0, |
| 983 | state_indices.stride(0), |
| 984 | query_start_loc.stride(0), |
| 985 | block_table_stride[0], |
| 986 | block_table_stride[1], |
| 987 | num_computed_stride, |
| 988 | SPEC_QUERY_LEN=self.spec_query_len, |
| 989 | num_warps=1, |
| 990 | ) |
| 991 | _compact_conv_state_kernel[(triton.cdiv(conv_dim, 256), batch, num_layers)]( |
| 992 | conv_ref, |
| 993 | self.conv_state_base_addrs, |
| 994 | self.conv_state_block_strides, |
| 995 | self.conv_state_dim_strides, |
| 996 | self.conv_state_token_strides, |
| 997 | state_indices, |
| 998 | self.commit_lens, |
| 999 | self.final_state_indices, |
| 1000 | self.boundary_state_indices, |
| 1001 | self.boundary_recovery_lens, |
| 1002 | NULL_BLOCK_ID, |
| 1003 | conv_dim, |
| 1004 | self.conv_history_len, |
| 1005 | state_indices.stride(0), |
| 1006 | BLOCK_D=256, |
| 1007 | BLOCK_HISTORY=block_history, |
| 1008 | ALIGN_MODE=block_table is not None, |
| 1009 | num_warps=4, |
| 1010 | ) |
| 1011 | |
| 1012 | state_ref = self.checkpoints[0] |
| 1013 | _, num_heads, value_dim, key_dim = state_ref.shape |
| 1014 | block_k = triton.next_power_of_2(key_dim) |
| 1015 | block_v = min(triton.next_power_of_2(value_dim), 32) |
| 1016 | grid = ( |
| 1017 | triton.cdiv(value_dim, block_v), |
| 1018 | batch, |
| 1019 | num_layers * num_heads, |
| 1020 | ) |
| 1021 | _commit_kda_state_kernel[grid]( |
| 1022 | state_ref, |
| 1023 | self.state_base_addrs, |
| 1024 | self.state_block_strides, |
| 1025 | self.correction_caches[0], |
| 1026 | self.correction_cache_base_addrs, |
| 1027 | self.correction_cache_block_strides, |
| 1028 | self.kg_caches[0], |
| 1029 | self.kg_cache_base_addrs, |
| 1030 | self.kg_cache_block_strides, |
| 1031 | self.A_log, |
| 1032 | self.dt_bias, |
| 1033 | state_indices, |
| 1034 | self.commit_lens, |
| 1035 | self.final_state_indices, |
| 1036 | self.boundary_state_indices, |
| 1037 | self.boundary_recovery_lens, |
| 1038 | self.lower_bound or 0.0, |
| 1039 | NULL_BLOCK_ID, |
| 1040 | state_ref.stride(1), |
| 1041 | state_ref.stride(2), |
| 1042 | state_ref.stride(3), |
| 1043 | self.correction_caches[0].stride(1), |
| 1044 | self.correction_caches[0].stride(2), |
| 1045 | self.correction_caches[0].stride(3), |
| 1046 | self.kg_caches[0].stride(1), |
| 1047 | self.kg_caches[0].stride(2), |
| 1048 | self.kg_caches[0].stride(3), |
| 1049 | self.A_log.stride(0), |
| 1050 | self.A_log.stride(1), |
| 1051 | self.dt_bias.stride(0), |
| 1052 | self.dt_bias.stride(1), |
| 1053 | self.dt_bias.stride(2), |
| 1054 | state_indices.stride(0), |
| 1055 | K=key_dim, |
| 1056 | V=value_dim, |
| 1057 | BK=block_k, |
| 1058 | BV=block_v, |
| 1059 | NUM_HEADS=num_heads, |
| 1060 | USE_LOWER_BOUND=self.lower_bound is not None, |
| 1061 | ALIGN_MODE=block_table is not None, |
| 1062 | num_warps=4, |
| 1063 | num_stages=2, |
| 1064 | ) |
| 1065 | |
| 1066 | |
| 1067 | __all__ = ["KDARecoverSSMCommitContext", "kda_recoverssm_verify"] |