| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Kimi-K3 specialization of GDN attention metadata. |
| 4 | |
| 5 | The request classification and cudagraph staging intentionally mirror |
| 6 | ``GDNAttentionMetadataBuilder``. Kimi-K3 builds the metadata required by its |
| 7 | prefill KDA kernel internally, so this builder omits the shared FLA chunk |
| 8 | metadata construction. |
| 9 | |
| 10 | For Kimi-K3 speculative decoding, ``--use-replayssm`` selects the simplified |
| 11 | RecoverSSM path implemented here instead of the Mamba2 ReplaySSM kernel. |
| 12 | """ |
| 13 | |
| 14 | from dataclasses import dataclass, field |
| 15 | from functools import cache |
| 16 | from typing import TYPE_CHECKING |
| 17 | |
| 18 | import torch |
| 19 | |
| 20 | from vllm.config import VllmConfig |
| 21 | from vllm.platforms import current_platform |
| 22 | from vllm.triton_utils import tl, triton |
| 23 | from vllm.utils.torch_utils import async_tensor_h2d |
| 24 | from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata |
| 25 | from vllm.v1.attention.backends.gdn_attn import ( |
| 26 | GDNAttentionBackend, |
| 27 | GDNAttentionMetadata, |
| 28 | GDNAttentionMetadataBuilder, |
| 29 | ) |
| 30 | from vllm.v1.attention.backends.recoverssm_metadata import ( |
| 31 | RecoverSSMMetadata, |
| 32 | RecoverSSMPostprocessMetadata, |
| 33 | ) |
| 34 | from vllm.v1.attention.backends.utils import ( |
| 35 | NULL_BLOCK_ID, |
| 36 | compute_causal_conv1d_metadata, |
| 37 | split_decodes_and_prefills, |
| 38 | ) |
| 39 | from vllm.v1.kv_cache_interface import MambaSpec |
| 40 | |
| 41 | if TYPE_CHECKING: |
| 42 | from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( |
| 43 | KDARecoverSSMCommitContext, |
| 44 | ) |
| 45 | |
| 46 | |
| 47 | @cache |
| 48 | def _metadata_launch_pdl() -> bool: |
| 49 | return current_platform.is_arch_support_pdl() |
| 50 | |
| 51 | |
| 52 | @triton.jit(do_not_specialize=["num_requests"]) |
| 53 | def _get_aligned_state_indices_kernel( |
| 54 | block_table_ptr, |
| 55 | seq_lens_ptr, |
| 56 | state_indices_ptr, |
| 57 | block_table_stride_0: tl.constexpr, |
| 58 | block_table_stride_1: tl.constexpr, |
| 59 | seq_lens_stride: tl.constexpr, |
| 60 | state_indices_stride_0: tl.constexpr, |
| 61 | state_indices_stride_1: tl.constexpr, |
| 62 | num_requests, |
| 63 | CACHE_BLOCK_SIZE: tl.constexpr, |
| 64 | NUM_STATE_SLOTS: tl.constexpr, |
| 65 | BLOCK_STATE_SLOTS: tl.constexpr, |
| 66 | BLOCK_ROWS: tl.constexpr, |
| 67 | launch_pdl: tl.constexpr, |
| 68 | ): |
| 69 | if launch_pdl: |
| 70 | tl.extra.cuda.gdc_wait() |
| 71 | tl.extra.cuda.gdc_launch_dependents() |
| 72 | |
| 73 | rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) |
| 74 | valid_row = rows < num_requests |
| 75 | seq_lens = tl.load( |
| 76 | seq_lens_ptr + rows * seq_lens_stride, |
| 77 | mask=valid_row, |
| 78 | other=1, |
| 79 | ) |
| 80 | # Triton truncates signed division toward zero, unlike PyTorch floor |
| 81 | # division. Clamping makes both semantics equivalent for seq_lens <= 0. |
| 82 | first_state_slot = tl.maximum((seq_lens - 1) // CACHE_BLOCK_SIZE, 0) |
| 83 | |
| 84 | state_slots = tl.arange(0, BLOCK_STATE_SLOTS) |
| 85 | valid_state_slot = state_slots < NUM_STATE_SLOTS |
| 86 | state_indices = tl.load( |
| 87 | block_table_ptr |
| 88 | + rows[:, None] * block_table_stride_0 |
| 89 | + (first_state_slot[:, None] + state_slots[None, :]) * block_table_stride_1, |
| 90 | mask=valid_row[:, None] & valid_state_slot[None, :], |
| 91 | ) |
| 92 | tl.store( |
| 93 | state_indices_ptr |
| 94 | + rows[:, None] * state_indices_stride_0 |
| 95 | + state_slots[None, :] * state_indices_stride_1, |
| 96 | state_indices, |
| 97 | mask=valid_row[:, None] & valid_state_slot[None, :], |
| 98 | ) |
| 99 | |
| 100 | |
| 101 | def _mamba_get_block_table_tensor( |
| 102 | block_table: torch.Tensor, |
| 103 | seq_lens: torch.Tensor, |
| 104 | kv_cache_spec: MambaSpec, |
| 105 | mamba_cache_mode: str, |
| 106 | ) -> torch.Tensor: |
| 107 | if mamba_cache_mode == "none": |
| 108 | return block_table |
| 109 | |
| 110 | assert block_table.is_cuda and seq_lens.is_cuda |
| 111 | num_requests = block_table.shape[0] |
| 112 | num_state_slots = 1 + kv_cache_spec.num_speculative_blocks |
| 113 | state_indices = torch.empty( |
| 114 | (num_requests, num_state_slots), |
| 115 | dtype=block_table.dtype, |
| 116 | device=block_table.device, |
| 117 | ) |
| 118 | BLOCK_ROWS = 32 |
| 119 | grid = (triton.cdiv(num_requests, BLOCK_ROWS),) |
| 120 | _get_aligned_state_indices_kernel[grid]( |
| 121 | block_table, |
| 122 | seq_lens, |
| 123 | state_indices, |
| 124 | block_table.stride(0), |
| 125 | block_table.stride(1), |
| 126 | seq_lens.stride(0), |
| 127 | state_indices.stride(0), |
| 128 | state_indices.stride(1), |
| 129 | num_requests, |
| 130 | CACHE_BLOCK_SIZE=kv_cache_spec.block_size, |
| 131 | NUM_STATE_SLOTS=num_state_slots, |
| 132 | BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots), |
| 133 | BLOCK_ROWS=BLOCK_ROWS, |
| 134 | num_warps=1, |
| 135 | launch_pdl=_metadata_launch_pdl(), |
| 136 | ) |
| 137 | return state_indices |
| 138 | |
| 139 | |
| 140 | @triton.jit(do_not_specialize=["num_spec_decodes", "batch_size"]) |
| 141 | def _stage_spec_decode_metadata_kernel( |
| 142 | state_indices_ptr, |
| 143 | query_start_loc_ptr, |
| 144 | num_accepted_tokens_ptr, |
| 145 | staged_state_indices_ptr, |
| 146 | staged_query_start_loc_ptr, |
| 147 | staged_num_accepted_tokens_ptr, |
| 148 | state_indices_stride_0: tl.constexpr, |
| 149 | state_indices_stride_1: tl.constexpr, |
| 150 | staged_state_indices_stride_0: tl.constexpr, |
| 151 | staged_state_indices_stride_1: tl.constexpr, |
| 152 | num_spec_decodes, |
| 153 | batch_size, |
| 154 | NUM_STATE_SLOTS: tl.constexpr, |
| 155 | BLOCK_STATE_SLOTS: tl.constexpr, |
| 156 | NULL_STATE_ID: tl.constexpr, |
| 157 | BLOCK_ROWS: tl.constexpr, |
| 158 | launch_pdl: tl.constexpr, |
| 159 | ): |
| 160 | if launch_pdl: |
| 161 | tl.extra.cuda.gdc_wait() |
| 162 | tl.extra.cuda.gdc_launch_dependents() |
| 163 | |
| 164 | rows = tl.program_id(0) * BLOCK_ROWS + tl.arange(0, BLOCK_ROWS) |
| 165 | real_request = rows < num_spec_decodes |
| 166 | |
| 167 | state_slots = tl.arange(0, BLOCK_STATE_SLOTS) |
| 168 | valid_state_slot = state_slots < NUM_STATE_SLOTS |
| 169 | state_indices = tl.load( |
| 170 | state_indices_ptr |
| 171 | + rows[:, None] * state_indices_stride_0 |
| 172 | + state_slots[None, :] * state_indices_stride_1, |
| 173 | mask=real_request[:, None] & valid_state_slot[None, :], |
| 174 | other=NULL_STATE_ID, |
| 175 | ) |
| 176 | tl.store( |
| 177 | staged_state_indices_ptr |
| 178 | + rows[:, None] * staged_state_indices_stride_0 |
| 179 | + state_slots[None, :] * staged_state_indices_stride_1, |
| 180 | state_indices, |
| 181 | mask=(rows < batch_size)[:, None] & valid_state_slot[None, :], |
| 182 | ) |
| 183 | |
| 184 | query_row = tl.minimum(rows, num_spec_decodes) |
| 185 | query_start_loc = tl.load( |
| 186 | query_start_loc_ptr + query_row, |
| 187 | mask=rows <= batch_size, |
| 188 | ) |
| 189 | tl.store( |
| 190 | staged_query_start_loc_ptr + rows, |
| 191 | query_start_loc, |
| 192 | mask=rows <= batch_size, |
| 193 | ) |
| 194 | |
| 195 | num_accepted_tokens = tl.load( |
| 196 | num_accepted_tokens_ptr + rows, |
| 197 | mask=real_request, |
| 198 | other=1, |
| 199 | ) |
| 200 | tl.store( |
| 201 | staged_num_accepted_tokens_ptr + rows, |
| 202 | num_accepted_tokens, |
| 203 | mask=rows < batch_size, |
| 204 | ) |
| 205 | |
| 206 | |
| 207 | def stage_spec_decode_metadata( |
| 208 | state_indices: torch.Tensor, |
| 209 | query_start_loc: torch.Tensor, |
| 210 | num_accepted_tokens: torch.Tensor, |
| 211 | staged_state_indices: torch.Tensor, |
| 212 | staged_query_start_loc: torch.Tensor, |
| 213 | staged_num_accepted_tokens: torch.Tensor, |
| 214 | *, |
| 215 | num_spec_decodes: int, |
| 216 | ) -> None: |
| 217 | """Stage speculative-decode metadata into CUDA-graph buffers.""" |
| 218 | assert state_indices.is_cuda |
| 219 | assert state_indices.ndim == 2 |
| 220 | batch_size, num_state_slots = staged_state_indices.shape |
| 221 | BLOCK_ROWS = 32 |
| 222 | grid = (triton.cdiv(batch_size + 1, BLOCK_ROWS),) |
| 223 | _stage_spec_decode_metadata_kernel[grid]( |
| 224 | state_indices, |
| 225 | query_start_loc, |
| 226 | num_accepted_tokens, |
| 227 | staged_state_indices, |
| 228 | staged_query_start_loc, |
| 229 | staged_num_accepted_tokens, |
| 230 | state_indices.stride(0), |
| 231 | state_indices.stride(1), |
| 232 | staged_state_indices.stride(0), |
| 233 | staged_state_indices.stride(1), |
| 234 | num_spec_decodes, |
| 235 | batch_size, |
| 236 | NUM_STATE_SLOTS=num_state_slots, |
| 237 | BLOCK_STATE_SLOTS=triton.next_power_of_2(num_state_slots), |
| 238 | NULL_STATE_ID=NULL_BLOCK_ID, |
| 239 | BLOCK_ROWS=BLOCK_ROWS, |
| 240 | num_warps=1, |
| 241 | launch_pdl=_metadata_launch_pdl(), |
| 242 | ) |
| 243 | |
| 244 | |
| 245 | @dataclass |
| 246 | class KDARecoverSSMAlignMetadata: |
| 247 | block_table: torch.Tensor |
| 248 | num_computed_tokens: torch.Tensor |
| 249 | block_size: int |
| 250 | |
| 251 | |
| 252 | @dataclass |
| 253 | class KDARecoverSSMCommitMetadata: |
| 254 | state_indices: torch.Tensor |
| 255 | query_start_loc: torch.Tensor |
| 256 | request_indices: torch.Tensor | None |
| 257 | align: KDARecoverSSMAlignMetadata | None |
| 258 | |
| 259 | |
| 260 | @dataclass |
| 261 | class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata): |
| 262 | spec_token_start: int | None = None |
| 263 | non_spec_token_start: int | None = None |
| 264 | flashinfer_prefill_query_start_loc: torch.Tensor | None = None |
| 265 | flashinfer_prefill_seq_order: torch.Tensor | None = None |
| 266 | recoverssm_commit: KDARecoverSSMCommitMetadata | None = None |
| 267 | recoverssm_context: "KDARecoverSSMCommitContext | None" = field( |
| 268 | default=None, repr=False, compare=False |
| 269 | ) |
| 270 | |
| 271 | def commit_recoverssm_state( |
| 272 | self, num_accepted_tokens: torch.Tensor |
| 273 | ) -> RecoverSSMPostprocessMetadata | None: |
| 274 | commit = self.recoverssm_commit |
| 275 | if commit is None: |
| 276 | return None |
| 277 | context = self.recoverssm_context |
| 278 | assert context is not None |
| 279 | align = commit.align |
| 280 | context.commit( |
| 281 | num_accepted_tokens, |
| 282 | commit.state_indices[: self.num_spec_decodes, 0], |
| 283 | commit.query_start_loc[: self.num_spec_decodes + 1], |
| 284 | request_indices=commit.request_indices, |
| 285 | block_table=align.block_table if align is not None else None, |
| 286 | num_computed_tokens=( |
| 287 | align.num_computed_tokens if align is not None else None |
| 288 | ), |
| 289 | mamba_block_size=align.block_size if align is not None else None, |
| 290 | ) |
| 291 | if align is None: |
| 292 | return None |
| 293 | return RecoverSSMPostprocessMetadata( |
| 294 | num_spec_decodes=self.num_spec_decodes, |
| 295 | request_indices=commit.request_indices, |
| 296 | block_table=align.block_table, |
| 297 | num_computed_tokens=align.num_computed_tokens, |
| 298 | block_size=align.block_size, |
| 299 | ) |
| 300 | |
| 301 | |
| 302 | class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): |
| 303 | # Overrides GDN's UNIFORM_BATCH: adaptive verification requires ALWAYS from every |
| 304 | # builder, and KDA reads per-request offsets off device within a fixed k+1 window, |
| 305 | # so one k+1 graph replays any 1..k+1 mix. |
| 306 | _cudagraph_support = AttentionCGSupport.ALWAYS |
| 307 | mamba_aligned_state_indices: torch.Tensor | None = None |
| 308 | |
| 309 | def __init__( |
| 310 | self, |
| 311 | kv_cache_spec: MambaSpec, |
| 312 | layer_names: list[str], |
| 313 | vllm_config: VllmConfig, |
| 314 | device: torch.device, |
| 315 | ) -> None: |
| 316 | super().__init__(kv_cache_spec, layer_names, vllm_config, device) |
| 317 | additional_config = vllm_config.additional_config |
| 318 | self.use_flashinfer_prefill = ( |
| 319 | isinstance(additional_config, dict) |
| 320 | and additional_config.get("kda_prefill_backend") == "flashinfer" |
| 321 | ) |
| 322 | self.use_recoverssm = vllm_config.cache_config.use_kda_recoverssm |
| 323 | self.spec_state_slots = 1 if self.use_recoverssm else self.num_spec + 1 |
| 324 | self.recoverssm_num_accepted_tokens: torch.Tensor | None = None |
| 325 | self.recoverssm_context: KDARecoverSSMCommitContext | None = None |
| 326 | if self.use_recoverssm: |
| 327 | max_num_reqs = vllm_config.scheduler_config.max_num_seqs |
| 328 | self.spec_state_indices_tensor = torch.empty( |
| 329 | (max_num_reqs, 1), |
| 330 | dtype=torch.int32, |
| 331 | device=device, |
| 332 | ) |
| 333 | self.recoverssm_num_accepted_tokens = torch.ones( |
| 334 | max_num_reqs, |
| 335 | dtype=torch.int32, |
| 336 | device=device, |
| 337 | ) |
| 338 | |
| 339 | def _get_recoverssm_context(self) -> "KDARecoverSSMCommitContext": |
| 340 | context = self.recoverssm_context |
| 341 | if context is not None: |
| 342 | return context |
| 343 | |
| 344 | from vllm.models.kimi_k3.nvidia.ops.recoverssm import ( |
| 345 | KDARecoverSSMCommitContext, |
| 346 | ) |
| 347 | |
| 348 | forward_context = self.vllm_config.compilation_config.static_forward_context |
| 349 | layers = [forward_context[layer_name] for layer_name in self.layer_names] |
| 350 | context = KDARecoverSSMCommitContext.create( |
| 351 | layers, |
| 352 | spec_query_len=1 + self.vllm_config.num_speculative_tokens, |
| 353 | max_num_reqs=self.vllm_config.scheduler_config.max_num_seqs, |
| 354 | ) |
| 355 | self.recoverssm_context = context |
| 356 | return context |
| 357 | |
| 358 | def build( # type: ignore[override] |
| 359 | self, |
| 360 | common_prefix_len: int, |
| 361 | common_attn_metadata: CommonAttentionMetadata, |
| 362 | num_accepted_tokens: torch.Tensor | None = None, |
| 363 | num_decode_draft_tokens_cpu: torch.Tensor | None = None, |
| 364 | fast_build: bool = False, |
| 365 | ) -> KimiK3KDAMetadata: |
| 366 | m = common_attn_metadata |
| 367 | query_start_loc = m.query_start_loc |
| 368 | query_start_loc_cpu = m.query_start_loc_cpu |
| 369 | assert isinstance(self.kv_cache_spec, MambaSpec) |
| 370 | # Equivalent PyTorch "align" path: |
| 371 | # start = ((seq_lens - 1) // block_size).clamp_(min=0) |
| 372 | # offsets = torch.arange(1 + num_speculative_blocks, dtype=torch.int32) |
| 373 | # indices = (start[:, None] + offsets).to(torch.int64) |
| 374 | # block_table_tensor = torch.gather(block_table, 1, indices) |
| 375 | if self.vllm_config.cache_config.mamba_cache_mode == "align": |
| 376 | if self.mamba_aligned_state_indices is not None: |
| 377 | block_table_tensor = self.mamba_aligned_state_indices[: m.num_reqs] |
| 378 | else: |
| 379 | assert not self.vllm_config.use_v2_model_runner, ( |
| 380 | "Aligned Mamba state indices must be precomputed" |
| 381 | ) |
| 382 | # TODO: remove this MRV1 fallback once MRV2 is the default runner. |
| 383 | block_table_tensor = _mamba_get_block_table_tensor( |
| 384 | m.block_table_tensor, |
| 385 | m.seq_lens, |
| 386 | self.kv_cache_spec, |
| 387 | self.vllm_config.cache_config.mamba_cache_mode, |
| 388 | ) |
| 389 | else: |
| 390 | block_table_tensor = m.block_table_tensor |
| 391 | |
| 392 | if not self.use_spec_decode or num_decode_draft_tokens_cpu is None: |
| 393 | spec_sequence_masks_cpu = None |
| 394 | num_spec_decodes = 0 |
| 395 | active_non_spec_mask_cpu = None |
| 396 | else: |
| 397 | spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0 |
| 398 | active_non_spec_mask_cpu = None |
| 399 | if self.use_recoverssm: |
| 400 | assert m.is_prefilling is not None |
| 401 | assert m.is_prefilling.device.type == "cpu" |
| 402 | active_decode_mask_cpu = (~m.is_prefilling) & ( |
| 403 | query_start_loc_cpu.diff() > 0 |
| 404 | ) |
| 405 | spec_sequence_masks_cpu |= active_decode_mask_cpu |
| 406 | # Adaptive's cost-table profiler probes eager dummy rows longer than |
| 407 | # RecoverSSM's num_spec+1 activation capacity; route them out instead |
| 408 | # of tripping the capacity check below. |
| 409 | spec_sequence_masks_cpu &= ( |
| 410 | query_start_loc_cpu.diff() <= self.num_spec + 1 |
| 411 | ) |
| 412 | # Native KDA can use its regular decode path when no draft token |
| 413 | # was scheduled. RecoverSSM must preserve its extended conv window. |
| 414 | if ( |
| 415 | not self.use_recoverssm |
| 416 | and num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() |
| 417 | == 0 |
| 418 | ): |
| 419 | spec_sequence_masks_cpu = None |
| 420 | num_spec_decodes = 0 |
| 421 | active_non_spec_mask_cpu = None |
| 422 | else: |
| 423 | num_spec_decodes = spec_sequence_masks_cpu.sum().item() |
| 424 | if num_spec_decodes == 0: |
| 425 | spec_sequence_masks_cpu = None |
| 426 | |
| 427 | spec_request_indices = None |
| 428 | spec_token_start = None |
| 429 | non_spec_token_start = None |
| 430 | if num_spec_decodes == 0: |
| 431 | # V2 already excludes prefills from full decode graphs via has_prefill. |
| 432 | # Classify first chunks as prefills to mask recycled state; |
| 433 | # resumed one-token chunks can still use the decode kernels. |
| 434 | assert m.seq_lens_cpu_upper_bound is not None |
| 435 | query_lens_cpu = query_start_loc_cpu.diff() |
| 436 | no_prior_state = (query_lens_cpu > 0) & ( |
| 437 | m.seq_lens_cpu_upper_bound <= query_lens_cpu |
| 438 | ) |
| 439 | # Capture batches also have seq_len == query_len, but are not prefills. |
| 440 | if m.is_prefilling is not None: |
| 441 | no_prior_state &= m.is_prefilling |
| 442 | else: |
| 443 | no_prior_state = torch.zeros_like(no_prior_state) |
| 444 | num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = ( |
| 445 | split_decodes_and_prefills( |
| 446 | m.replace(is_prefilling=no_prior_state), |
| 447 | decode_threshold=1, |
| 448 | treat_short_extends_as_decodes=False, |
| 449 | ) |
| 450 | ) |
| 451 | # Exclude trailing padding from both prefill counts. |
| 452 | if num_prefills: |
| 453 | num_prefills -= int((query_lens_cpu[num_decodes:] == 0).sum()) |
| 454 | num_prefill_tokens = ( |
| 455 | int(query_start_loc_cpu[num_decodes + num_prefills]) |
| 456 | - num_decode_tokens |
| 457 | ) |
| 458 | num_spec_decode_tokens = 0 |
| 459 | spec_token_indx = None |
| 460 | non_spec_token_indx = None |
| 461 | spec_state_indices_tensor = None |
| 462 | non_spec_state_indices_tensor = block_table_tensor[:, 0] |
| 463 | spec_query_start_loc = None |
| 464 | non_spec_query_start_loc = query_start_loc |
| 465 | non_spec_query_start_loc_cpu = query_start_loc_cpu |
| 466 | num_accepted_tokens = None |
| 467 | else: |
| 468 | assert spec_sequence_masks_cpu is not None |
| 469 | assert num_accepted_tokens is not None |
| 470 | query_lens_cpu = query_start_loc_cpu.diff() |
| 471 | if ( |
| 472 | self.use_recoverssm |
| 473 | and torch.any( |
| 474 | query_lens_cpu[spec_sequence_masks_cpu] > self.num_spec + 1 |
| 475 | ).item() |
| 476 | ): |
| 477 | raise ValueError( |
| 478 | "KDA RecoverSSM speculative decode query length exceeds " |
| 479 | f"its activation capacity ({self.num_spec + 1})" |
| 480 | ) |
| 481 | num_query_tokens = query_start_loc_cpu[-1].item() |
| 482 | |
| 483 | # Exclude zero-length cudagraph padding from request-indexed |
| 484 | # non-spec metadata. |
| 485 | active_non_spec_mask_cpu = (~spec_sequence_masks_cpu) & (query_lens_cpu > 0) |
| 486 | non_spec_query_lens_cpu = query_lens_cpu[active_non_spec_mask_cpu] |
| 487 | num_non_spec_requests = non_spec_query_lens_cpu.numel() |
| 488 | num_non_spec_tokens = non_spec_query_lens_cpu.sum().item() |
| 489 | |
| 490 | # Query length alone cannot distinguish a true decode from a |
| 491 | # one-token prefill chunk. Packed decode is safe only when every |
| 492 | # active non-spec request is a true, one-token decode. |
| 493 | assert m.is_prefilling is not None |
| 494 | assert m.is_prefilling.device.type == "cpu" |
| 495 | non_spec_is_prefilling = m.is_prefilling[active_non_spec_mask_cpu] |
| 496 | use_prefill = torch.any( |
| 497 | non_spec_is_prefilling | (non_spec_query_lens_cpu != 1) |
| 498 | ).item() |
| 499 | if use_prefill: |
| 500 | num_prefills = num_non_spec_requests |
| 501 | num_prefill_tokens = num_non_spec_tokens |
| 502 | num_decodes = 0 |
| 503 | num_decode_tokens = 0 |
| 504 | else: |
| 505 | num_prefills = 0 |
| 506 | num_prefill_tokens = 0 |
| 507 | num_decodes = num_non_spec_requests |
| 508 | num_decode_tokens = num_non_spec_tokens |
| 509 | |
| 510 | num_spec_decode_tokens = num_query_tokens - num_non_spec_tokens |
| 511 | |
| 512 | if num_prefills == 0 and num_decodes == 0: |
| 513 | spec_token_indx = None |
| 514 | non_spec_token_indx = None |
| 515 | # Real requests precede trailing cudagraph padding. |
| 516 | spec_state_indices_tensor = block_table_tensor[ |
| 517 | :num_spec_decodes, : self.spec_state_slots |
| 518 | ] |
| 519 | non_spec_state_indices_tensor = None |
| 520 | # Padding trails real requests, so this prefix already contains |
| 521 | # the correct cumulative token counts. |
| 522 | spec_query_start_loc = query_start_loc[: num_spec_decodes + 1] |
| 523 | non_spec_query_start_loc = None |
| 524 | non_spec_query_start_loc_cpu = None |
| 525 | num_accepted_tokens = num_accepted_tokens[:num_spec_decodes] |
| 526 | else: |
| 527 | query_lens = query_start_loc.diff() |
| 528 | spec_sequence_masks_gpu = async_tensor_h2d( |
| 529 | spec_sequence_masks_cpu, device=query_start_loc.device |
| 530 | ) |
| 531 | if self.use_recoverssm: |
| 532 | spec_request_indices = async_tensor_h2d( |
| 533 | spec_sequence_masks_cpu.nonzero(as_tuple=True)[0], |
| 534 | dtype=torch.int32, |
| 535 | device=query_start_loc.device, |
| 536 | ) |
| 537 | spec_token_masks = torch.repeat_interleave( |
| 538 | spec_sequence_masks_gpu, |
| 539 | query_lens, |
| 540 | output_size=num_query_tokens, |
| 541 | ) |
| 542 | # Stable partitioning preserves request-local token order in |
| 543 | # both subgroup tensors. |
| 544 | index = torch.argsort(spec_token_masks, stable=True) |
| 545 | num_non_spec_tokens = num_prefill_tokens + num_decode_tokens |
| 546 | non_spec_token_indx = index[:num_non_spec_tokens] |
| 547 | spec_token_indx = index[num_non_spec_tokens:] |
| 548 | |
| 549 | active_spec_mask = spec_sequence_masks_cpu[query_lens_cpu > 0] |
| 550 | # check if spec / non spec tokens are continuous |
| 551 | if (active_spec_mask[1:] != active_spec_mask[:-1]).sum().item() == 1: |
| 552 | spec_first = active_spec_mask[0].item() |
| 553 | spec_token_start = 0 if spec_first else num_non_spec_tokens |
| 554 | non_spec_token_start = num_spec_decode_tokens if spec_first else 0 |
| 555 | |
| 556 | # Native spec uses one state slot per step. RecoverSSM keeps |
| 557 | # only the current checkpoint slot. |
| 558 | spec_state_indices_tensor = block_table_tensor[ |
| 559 | spec_sequence_masks_cpu, : self.spec_state_slots |
| 560 | ] |
| 561 | non_spec_state_indices_tensor = block_table_tensor[ |
| 562 | active_non_spec_mask_cpu, 0 |
| 563 | ] |
| 564 | |
| 565 | spec_query_lens = query_lens[spec_sequence_masks_cpu] |
| 566 | spec_query_start_loc = torch.zeros( |
| 567 | num_spec_decodes + 1, |
| 568 | dtype=torch.int32, |
| 569 | device=query_start_loc.device, |
| 570 | ) |
| 571 | torch.cumsum( |
| 572 | spec_query_lens, |
| 573 | dim=0, |
| 574 | out=spec_query_start_loc[1:], |
| 575 | ) |
| 576 | if num_prefills > 0: |
| 577 | non_spec_query_lens = query_lens[active_non_spec_mask_cpu] |
| 578 | non_spec_query_start_loc = torch.zeros( |
| 579 | non_spec_query_lens.size(0) + 1, |
| 580 | dtype=torch.int32, |
| 581 | device=query_start_loc.device, |
| 582 | ) |
| 583 | torch.cumsum( |
| 584 | non_spec_query_lens, |
| 585 | dim=0, |
| 586 | out=non_spec_query_start_loc[1:], |
| 587 | ) |
| 588 | non_spec_query_start_loc_cpu = torch.zeros( |
| 589 | non_spec_query_lens_cpu.size(0) + 1, |
| 590 | dtype=torch.int32, |
| 591 | ) |
| 592 | torch.cumsum( |
| 593 | non_spec_query_lens_cpu, |
| 594 | dim=0, |
| 595 | out=non_spec_query_start_loc_cpu[1:], |
| 596 | ) |
| 597 | else: |
| 598 | # Packed decode consumes one row per request and does not |
| 599 | # use cumulative sequence lengths. |
| 600 | non_spec_query_start_loc = None |
| 601 | non_spec_query_start_loc_cpu = None |
| 602 | |
| 603 | num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu] |
| 604 | |
| 605 | if self.use_recoverssm: |
| 606 | assert self.recoverssm_num_accepted_tokens is not None |
| 607 | num_accepted_tokens = self.recoverssm_num_accepted_tokens[ |
| 608 | :num_spec_decodes |
| 609 | ] |
| 610 | |
| 611 | # Unlike the shared GDN layer, Kimi-K3's prefill KDA wrapper prepares |
| 612 | # its own chunk indices. Only causal-convolution metadata is needed here. |
| 613 | nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None |
| 614 | if num_prefills > 0: |
| 615 | has_initial_state = m.compute_num_computed_tokens() > 0 |
| 616 | if num_spec_decodes > 0: |
| 617 | has_initial_state = has_initial_state[active_non_spec_mask_cpu] |
| 618 | assert non_spec_query_start_loc_cpu is not None |
| 619 | nums_dict, batch_ptr, token_chunk_offset_ptr = ( |
| 620 | compute_causal_conv1d_metadata( |
| 621 | non_spec_query_start_loc_cpu, |
| 622 | device=query_start_loc.device, |
| 623 | ) |
| 624 | ) |
| 625 | else: |
| 626 | has_initial_state = None |
| 627 | |
| 628 | # Prepare per-request tensors for cudagraph replay. num_actual_tokens |
| 629 | # may be token-padded, while state/query/acceptance metadata is indexed |
| 630 | # by request. |
| 631 | batch_size = m.num_reqs |
| 632 | if ( |
| 633 | self.use_full_cuda_graph |
| 634 | and num_spec_decodes > 0 |
| 635 | and num_prefills == 0 |
| 636 | and num_decodes == 0 |
| 637 | and batch_size <= self.spec_state_indices_tensor.shape[0] |
| 638 | and num_spec_decode_tokens <= self.decode_cudagraph_max_bs |
| 639 | ): |
| 640 | # Equivalent PyTorch staging: |
| 641 | # state[:N].copy_(state_src); state[N:].fill_(NULL_BLOCK_ID) |
| 642 | # qsl[:N + 1].copy_(qsl_src); qsl[N + 1:].fill_(qsl_src[-1]) |
| 643 | # accepted[:N].copy_(accepted_src); accepted[N:].fill_(1) |
| 644 | stage_spec_decode_metadata( |
| 645 | state_indices=spec_state_indices_tensor, |
| 646 | query_start_loc=spec_query_start_loc, |
| 647 | num_accepted_tokens=num_accepted_tokens, |
| 648 | staged_state_indices=self.spec_state_indices_tensor[:batch_size], |
| 649 | staged_query_start_loc=self.spec_query_start_loc[: batch_size + 1], |
| 650 | staged_num_accepted_tokens=self.num_accepted_tokens[:batch_size], |
| 651 | num_spec_decodes=num_spec_decodes, |
| 652 | ) |
| 653 | |
| 654 | spec_state_indices_tensor = self.spec_state_indices_tensor[:batch_size] |
| 655 | spec_query_start_loc = self.spec_query_start_loc[: batch_size + 1] |
| 656 | num_accepted_tokens = self.num_accepted_tokens[:batch_size] |
| 657 | |
| 658 | if ( |
| 659 | self.use_full_cuda_graph |
| 660 | and num_prefills == 0 |
| 661 | and num_spec_decodes == 0 |
| 662 | and num_decodes <= self.decode_cudagraph_max_bs |
| 663 | ): |
| 664 | self.non_spec_state_indices_tensor[:num_decodes].copy_( |
| 665 | non_spec_state_indices_tensor, non_blocking=True |
| 666 | ) |
| 667 | self.non_spec_state_indices_tensor[num_decodes:batch_size].fill_( |
| 668 | NULL_BLOCK_ID |
| 669 | ) |
| 670 | non_spec_state_indices_tensor = self.non_spec_state_indices_tensor[ |
| 671 | :batch_size |
| 672 | ] |
| 673 | |
| 674 | recoverssm_commit = None |
| 675 | if self.use_recoverssm and num_spec_decodes > 0: |
| 676 | assert spec_state_indices_tensor is not None |
| 677 | assert spec_query_start_loc is not None |
| 678 | align = None |
| 679 | if self.kv_cache_spec.mamba_cache_mode == "align": |
| 680 | align = KDARecoverSSMAlignMetadata( |
| 681 | block_table=m.block_table_tensor, |
| 682 | num_computed_tokens=m.compute_num_computed_tokens(), |
| 683 | block_size=self.kv_cache_spec.block_size, |
| 684 | ) |
| 685 | recoverssm_commit = KDARecoverSSMCommitMetadata( |
| 686 | state_indices=spec_state_indices_tensor, |
| 687 | query_start_loc=spec_query_start_loc, |
| 688 | request_indices=spec_request_indices, |
| 689 | align=align, |
| 690 | ) |
| 691 | |
| 692 | flashinfer_prefill_query_start_loc = None |
| 693 | flashinfer_prefill_seq_order = None |
| 694 | if self.use_flashinfer_prefill and num_prefills > 0: |
| 695 | assert non_spec_query_start_loc is not None |
| 696 | flashinfer_prefill_query_start_loc = non_spec_query_start_loc.to( |
| 697 | torch.int64 |
| 698 | ) |
| 699 | num_non_spec_requests = non_spec_query_start_loc.shape[0] - 1 |
| 700 | num_non_spec_tokens = num_prefill_tokens + num_decode_tokens |
| 701 | if num_non_spec_tokens > num_non_spec_requests: |
| 702 | flashinfer_prefill_seq_order = torch.argsort( |
| 703 | flashinfer_prefill_query_start_loc.diff(), descending=True |
| 704 | ).to(torch.int32) |
| 705 | |
| 706 | checkpoint = None |
| 707 | if num_prefills > 0: |
| 708 | request_rows = list(range(m.num_reqs)) |
| 709 | if active_non_spec_mask_cpu is not None: |
| 710 | request_rows = active_non_spec_mask_cpu.nonzero().flatten().tolist() |
| 711 | checkpoint = self.checkpoint_builder.build(m, request_rows) |
| 712 | |
| 713 | return KimiK3KDAMetadata( |
| 714 | num_prefills=num_prefills, |
| 715 | num_prefill_tokens=num_prefill_tokens, |
| 716 | num_decodes=num_decodes, |
| 717 | num_decode_tokens=num_decode_tokens, |
| 718 | num_spec_decodes=num_spec_decodes, |
| 719 | num_spec_decode_tokens=num_spec_decode_tokens, |
| 720 | num_actual_tokens=m.num_actual_tokens, |
| 721 | has_initial_state=has_initial_state, |
| 722 | spec_query_start_loc=spec_query_start_loc, |
| 723 | non_spec_query_start_loc=non_spec_query_start_loc, |
| 724 | spec_state_indices_tensor=spec_state_indices_tensor, |
| 725 | non_spec_state_indices_tensor=non_spec_state_indices_tensor, |
| 726 | spec_sequence_masks=None, |
| 727 | spec_token_indx=spec_token_indx, |
| 728 | non_spec_token_indx=non_spec_token_indx, |
| 729 | num_accepted_tokens=num_accepted_tokens, |
| 730 | spec_token_start=spec_token_start, |
| 731 | non_spec_token_start=non_spec_token_start, |
| 732 | flashinfer_prefill_query_start_loc=flashinfer_prefill_query_start_loc, |
| 733 | flashinfer_prefill_seq_order=flashinfer_prefill_seq_order, |
| 734 | recoverssm_commit=recoverssm_commit, |
| 735 | recoverssm_context=( |
| 736 | self._get_recoverssm_context() |
| 737 | if recoverssm_commit is not None |
| 738 | else None |
| 739 | ), |
| 740 | nums_dict=nums_dict, |
| 741 | batch_ptr=batch_ptr, |
| 742 | token_chunk_offset_ptr=token_chunk_offset_ptr, |
| 743 | checkpoint=checkpoint, |
| 744 | ) |
| 745 | |
| 746 | |
| 747 | class KimiK3KDAAttentionBackend(GDNAttentionBackend): |
| 748 | @staticmethod |
| 749 | def get_name() -> str: |
| 750 | return "KIMI_K3_KDA" |
| 751 | |
| 752 | @staticmethod |
| 753 | def get_builder_cls() -> type[KimiK3KDAMetadataBuilder]: |
| 754 | return KimiK3KDAMetadataBuilder |
| 755 | |
| 756 | @classmethod |
| 757 | def supports_device_cpu_query_lens_mismatch(cls) -> bool: |
| 758 | # Pure spec-decode reads its plan from DEVICE offsets, so CPU/device |
| 759 | # mismatch is fine; the mixed branch still assumes CPU totals match. |
| 760 | return True |