diff --git a/cmake/external_projects/flashkda.cmake b/cmake/external_projects/flashkda.cmake index 988a27d7c7..b70042f2f4 100644 --- a/cmake/external_projects/flashkda.cmake +++ b/cmake/external_projects/flashkda.cmake @@ -13,7 +13,7 @@ else() FetchContent_Declare( flashkda GIT_REPOSITORY https://github.com/vllm-project/FlashKDA.git - GIT_TAG 053de1b716ef3255873e02d2d28f4adf09951978 + GIT_TAG ee0be888cd0e972f9409bf53756f8c38c6652173 GIT_PROGRESS TRUE GIT_SUBMODULES cutlass ) diff --git a/csrc/flashkda_registration.cpp b/csrc/flashkda_registration.cpp index f3da713a09..6a9f56b8d2 100644 --- a/csrc/flashkda_registration.cpp +++ b/csrc/flashkda_registration.cpp @@ -10,7 +10,8 @@ STABLE_TORCH_LIBRARY(_flashkda_C, m) { "Tensor(a!) out, Tensor(c!) workspace, Tensor A_log, Tensor dt_bias, " "float lower_bound, " "Tensor? initial_state=None, Tensor(b!)? final_state=None, " - "Tensor? cu_seqlens=None) -> ()"); + "Tensor? cu_seqlens=None, Tensor(d!)? checkpoint_state=None, " + "Tensor? checkpoint_offsets=None) -> ()"); } STABLE_TORCH_LIBRARY_IMPL(_flashkda_C, CompositeExplicitAutograd, m) { diff --git a/tests/models/kimi_k3/test_kda.py b/tests/models/kimi_k3/test_kda.py index 7416cca02a..c8454e9e6f 100644 --- a/tests/models/kimi_k3/test_kda.py +++ b/tests/models/kimi_k3/test_kda.py @@ -21,6 +21,8 @@ from vllm.models.kimi_k3.amd.ops.third_party.kda import ( fused_recurrent_kda_packed_decode as fused_recurrent_kda_packed_decode_amd, ) from vllm.models.kimi_k3.nvidia.kda import ( + _flashkda_prefill, + _store_cache_checkpoints_kernel, is_flashkda_supported, is_fused_kda_decode_supported, ) @@ -40,6 +42,7 @@ from vllm.models.kimi_k3.nvidia.ops.third_party.kda import ( ) from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.l2norm import l2norm_fwd +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID DEVICE = current_platform.device_type @@ -1088,6 +1091,17 @@ def test_flashkda_correctness(): expected_states.append(final_state) expected_out = torch.cat(expected_outputs, dim=1) expected_state = torch.cat(expected_states).transpose(-1, -2).contiguous() + _, expected_checkpoint = naive_recurrent_kda( + q_norm[:, :16], + k_norm[:, :16], + v[:, :16], + gate[:, :16], + beta[:, :16], + initial_state=initial_state[0:1].transpose(-1, -2), + output_final_state=True, + ) + assert expected_checkpoint is not None + expected_checkpoint = expected_checkpoint.transpose(-1, -2).contiguous() actual_out = torch.empty_like(v) actual_state = torch.empty_like(initial_state) @@ -1115,3 +1129,70 @@ def test_flashkda_correctness(): assert_close("o", expected_out, actual_out, 0.01) assert_close("ht", expected_state, actual_state, 0.01) + + checkpoint_out = torch.empty_like(v) + checkpoint_final_state = torch.empty_like(initial_state) + checkpoint_state = torch.empty_like(initial_state) + checkpoint_offsets = torch.tensor([16, 31], dtype=torch.int32, device=DEVICE) + _flashkda_prefill( + q=q, + k=k, + v=v, + g=raw_g, + beta=beta_logits, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + initial_state=initial_state, + cu_seqlens=cu_seqlens, + out=checkpoint_out, + final_state=checkpoint_final_state, + workspace=workspace, + checkpoint_state=checkpoint_state, + checkpoint_offsets=checkpoint_offsets, + ) + + assert_close("checkpoint_o", expected_out, checkpoint_out, 0.01) + assert_close("checkpoint_ht", expected_state, checkpoint_final_state, 0.01) + assert_close("checkpoint", expected_checkpoint, checkpoint_state[:1], 0.01) + + conv_state = torch.zeros(2, H * D, 3, dtype=q.dtype, device=DEVICE) + recurrent_storage = torch.zeros(2, H * D * D + 8, device=DEVICE) + recurrent_state = recurrent_storage[:, : H * D * D].view(2, H, D, D) + conv_input = q[0].flatten(1) + checkpoint_state_indices = torch.tensor( + [1, NULL_BLOCK_ID], dtype=torch.int32, device=DEVICE + ) + state_len = conv_state.shape[-1] + width = H * D + recurrent_row_size = checkpoint_state[0].numel() + block_size = 256 + _store_cache_checkpoints_kernel[ + ( + checkpoint_state_indices.numel(), + (max(width * state_len, recurrent_row_size) + block_size - 1) // block_size, + ) + ]( + conv_input, + conv_state, + checkpoint_state, + recurrent_state, + cu_seqlens, + checkpoint_offsets, + checkpoint_state_indices, + conv_input.stride(0), + conv_input.stride(1), + conv_state.stride(0), + conv_state.stride(1), + conv_state.stride(2), + checkpoint_state.stride(0), + recurrent_state.stride(0), + checkpoint_offsets.stride(0), + state_len, + width, + recurrent_row_size, + NULL_BLOCK_ID, + block_size, + ) + torch.testing.assert_close(conv_state[1], q[0, 13:16].flatten(1).transpose(0, 1)) + torch.testing.assert_close(recurrent_state[1], checkpoint_state[0]) diff --git a/tests/models/kimi_k3/test_kda_metadata.py b/tests/models/kimi_k3/test_kda_metadata.py index 069c9a4d7a..ec00467746 100644 --- a/tests/models/kimi_k3/test_kda_metadata.py +++ b/tests/models/kimi_k3/test_kda_metadata.py @@ -87,6 +87,7 @@ def _make_builder( device: torch.device = DEVICE, mamba_cache_mode: str = "none", use_recoverssm: bool = False, + num_prefill_checkpoint_blocks: int = 0, ) -> AttentionMetadataBuilder: vllm_config = create_vllm_config( model_name="Qwen/Qwen3.5-0.8B", @@ -110,6 +111,7 @@ def _make_builder( dtypes=(torch.float16,), mamba_cache_mode=mamba_cache_mode, num_speculative_blocks=(0 if use_recoverssm else num_speculative_tokens), + num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks, ), layer_names=["layer.0"], vllm_config=vllm_config, @@ -121,6 +123,54 @@ def _make_builder( return builder +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_internal_checkpoint_metadata_targets_last_aligned_boundary(): + device = torch.device("cuda") + batch = BatchSpec(seq_lens=[50, 32], query_lens=[50, 16]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, device, arange_block_indices=True + ).replace(is_prefilling=torch.tensor([True, True])) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=0, + full_cuda_graph=False, + mamba_cache_mode="align", + num_prefill_checkpoint_blocks=1, + device=device, + ).build(0, common_attn_metadata) + + assert actual.checkpoint is not None + torch.testing.assert_close( + actual.checkpoint.state_indices, + torch.tensor([2, NULL_BLOCK_ID], dtype=torch.int32, device=device), + ) + torch.testing.assert_close( + actual.checkpoint.checkpoint_offsets, + torch.tensor([48, 0], dtype=torch.int32, device=device), + ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") +def test_internal_checkpoint_metadata_skips_unaligned_offset(): + device = torch.device("cuda") + # The checkpoint block boundary is 48, but this query starts at token 1, + # making the real checkpoint offset 47, which is not FlashKDA-aligned. + batch = BatchSpec(seq_lens=[50], query_lens=[49]) + common_attn_metadata = create_common_attn_metadata( + batch, BLOCK_SIZE, device, arange_block_indices=True + ).replace(is_prefilling=torch.tensor([True])) + actual = _make_builder( + KimiK3KDAMetadataBuilder, + num_speculative_tokens=0, + full_cuda_graph=False, + mamba_cache_mode="align", + num_prefill_checkpoint_blocks=1, + device=device, + ).build(0, common_attn_metadata) + + assert actual.checkpoint is None + + @pytest.mark.parametrize( ( "batch", diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index da5bd5b5a1..eff889b987 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -77,6 +77,7 @@ def make_full_mamba_manager( mamba_block_size: int = 4, num_blocks: int = 32, use_eagle: bool = False, + num_prefill_checkpoint_blocks: int = 0, ): kv_cache_config = KVCacheConfig( num_blocks=num_blocks, @@ -98,6 +99,7 @@ def make_full_mamba_manager( shapes=(1, 1), dtypes=(torch.float32,), mamba_cache_mode="align", + num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks, ), ), ], @@ -135,6 +137,7 @@ def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int): dcp_world_size=dcp_world_size, scheduler_block_size=scheduler_block_size, mamba_partial_cache_hit=True, + mamba_has_prefill_checkpoint_blocks=False, ) split = Scheduler._mamba_block_aligned_split @@ -180,6 +183,7 @@ def test_mamba_align_split_when_block_exceeds_scheduling_budget(): use_eagle=False, hash_block_size=32, mamba_partial_cache_hit=False, + mamba_has_prefill_checkpoint_blocks=False, ) req = make_request("0", [0] * prompt_length, 32, sha256) split = Scheduler._mamba_block_aligned_split @@ -218,6 +222,7 @@ def test_mamba_align_split_when_block_exceeds_long_prefill_threshold(): use_eagle=False, hash_block_size=32, mamba_partial_cache_hit=False, + mamba_has_prefill_checkpoint_blocks=False, ) req = make_request("0", [0] * prompt_length, 32, sha256) split = Scheduler._mamba_block_aligned_split @@ -462,6 +467,85 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue(): assert moved[0].block_hash_num_tokens == 6 +def test_partial_hit_then_internal_checkpoint_uses_distinct_mamba_blocks(): + hash_block_size = 2 + mamba_block_size = 4 + manager = make_full_mamba_manager( + dcp_world_size=1, + hash_block_size=hash_block_size, + full_block_size=hash_block_size, + mamba_block_size=mamba_block_size, + num_prefill_checkpoint_blocks=1, + ) + + owner = make_request("owner", [0, 0, 1, 1, 2, 2], hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(owner) + assert manager.allocate_slots(owner, 6, num_computed, computed_blocks) is not None + manager.free(owner) + manager.new_step_starts() + + partial_hash = owner.block_hashes[2] + partial_block = manager.block_pool.get_cached_block(partial_hash, [1]) + assert partial_block is not None + partial_block_id = partial_block[0].block_id + + replay = make_request( + "replay", + [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6], + hash_block_size, + sha256, + ) + computed_blocks, num_computed, _ = manager.get_computed_blocks(replay) + assert num_computed == 6 + + realigned_blocks = manager.allocate_slots(replay, 2, num_computed, computed_blocks) + assert realigned_blocks is not None + mamba_cow_block_id = realigned_blocks.get_block_ids()[1][0] + copies, retained = manager.take_kv_cache_block_copies() + assert KVCacheBlockCopy(partial_block_id, mamba_cow_block_id) in copies + manager.block_pool.free_blocks(retained) + + replay.num_computed_tokens = 8 + manager.new_step_starts() + final_blocks = manager.allocate_slots(replay, 6) + assert final_blocks is not None + + checkpoint_block_id, running_block_id = final_blocks.get_block_ids()[1] + assert len({mamba_cow_block_id, checkpoint_block_id, running_block_id}) == 3 + mamba_blocks = manager.get_blocks(replay.request_id).blocks[1] + assert mamba_blocks[2].block_id == checkpoint_block_id + assert mamba_blocks[3].block_id == running_block_id + + +def test_internal_checkpoint_requires_block_aligned_start(): + hash_block_size = 2 + mamba_block_size = 16 + manager = make_full_mamba_manager( + dcp_world_size=1, + hash_block_size=hash_block_size, + full_block_size=hash_block_size, + mamba_block_size=mamba_block_size, + num_prefill_checkpoint_blocks=1, + ) + request = make_request("producer", list(range(50)), hash_block_size, sha256) + + # Compute token 0 first, so the next query starts at token 1, which is not + # aligned to the Mamba block size. + assert manager.allocate_slots(request, 1) is not None + request.num_computed_tokens = 1 + manager.new_step_starts() + + new_blocks = manager.allocate_slots(request, 49) + + assert new_blocks is not None + mamba_blocks = manager.get_blocks(request.request_id).blocks[1] + checkpoint_block_idx = 48 // mamba_block_size - 1 + assert mamba_blocks[checkpoint_block_idx].is_null + assert not mamba_blocks[-1].is_null + checkpoint_hash = request.block_hashes[48 // hash_block_size - 1] + assert manager.block_pool.get_cached_block(checkpoint_hash, [1]) is None + + def test_external_mamba_hit_same_block_uses_running_cow_on_continue(): """An external mid-block hit must become a running request even when its first continuation does not need another Mamba block.""" @@ -560,6 +644,7 @@ def test_take_partial_tail_offloads_returns_cow_target(): shapes=(1, 1), dtypes=(torch.float32,), mamba_cache_mode="align", + num_prefill_checkpoint_blocks=1, ), ), ], diff --git a/tests/v1/core/test_mamba_align_chunk_split.py b/tests/v1/core/test_mamba_align_chunk_split.py index 31671caf76..bc809fc3b5 100644 --- a/tests/v1/core/test_mamba_align_chunk_split.py +++ b/tests/v1/core/test_mamba_align_chunk_split.py @@ -37,7 +37,9 @@ PROMPT_LEN = 2002 MAMBA_GROUP_ID = 1 -def _make_hybrid_kv_cache_manager() -> KVCacheManager: +def _make_hybrid_kv_cache_manager( + num_prefill_checkpoint_blocks: int = 0, +) -> KVCacheManager: config = KVCacheConfig( num_blocks=10000, kv_cache_tensors=[], @@ -59,6 +61,7 @@ def _make_hybrid_kv_cache_manager() -> KVCacheManager: dtypes=(torch.float32,), mamba_cache_mode="align", num_speculative_blocks=NUM_SPEC, + num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks, ), ), ], @@ -78,6 +81,7 @@ def _split( num_new_tokens: int, use_eagle: bool = True, partial_hit: bool = False, + num_prefill_checkpoint_blocks: int = 0, ) -> int: """Call the real `Scheduler._mamba_block_aligned_split` on a stub self.""" stub = SimpleNamespace( @@ -88,10 +92,42 @@ def _split( # `prefix_match_unit` finer than the block size (#46384). mamba_partial_cache_hit=partial_hit, hash_block_size=ATTN_BLOCK_SIZE, + mamba_has_prefill_checkpoint_blocks=( + num_prefill_checkpoint_blocks > 0 and not use_eagle + ), ) return Scheduler._mamba_block_aligned_split(stub, request, num_new_tokens) +@pytest.mark.parametrize( + ("prompt_len", "num_new_tokens", "use_eagle", "expected"), + [ + (2002, 2002, False, 2002), + (3602, 2000, False, 2000), + (3602, 3602, True, MAMBA_BLOCK_SIZE), + ], +) +def test_internal_checkpoint_split( + prompt_len: int, num_new_tokens: int, use_eagle: bool, expected: int +) -> None: + (request,) = create_requests(1, num_tokens=prompt_len, block_size=ATTN_BLOCK_SIZE) + assert ( + _split( + request, + num_new_tokens, + use_eagle=use_eagle, + num_prefill_checkpoint_blocks=1, + ) + == expected + ) + if not use_eagle: + manager = _make_hybrid_kv_cache_manager(num_prefill_checkpoint_blocks=1) + assert manager.allocate_slots(request, expected) is not None + mamba_manager = manager.coordinator.single_type_managers[MAMBA_GROUP_ID] + blocks = mamba_manager.req_to_blocks[request.request_id] + assert all(not block.is_null for block in blocks) # checkpoint + running state + + def _run_chunked_prefill( manager: KVCacheManager, request: Request, budgets: list[int] ) -> dict[int, int]: @@ -211,8 +247,15 @@ def test_poisoning_is_block_size_independent( @pytest.mark.parametrize("partial_hit", [False, True]) @pytest.mark.parametrize("resume_at", [331, 1599, 1601, 2531, 3011]) +@pytest.mark.parametrize( + ("num_prefill_checkpoint_blocks", "use_eagle"), + [(0, True), (1, False)], +) def test_unaligned_resume_never_runs_past_its_block( - partial_hit: bool, resume_at: int + partial_hit: bool, + resume_at: int, + num_prefill_checkpoint_blocks: int, + use_eagle: bool, ) -> None: """A prefill resuming mid-block must re-align before crossing a boundary. @@ -227,7 +270,13 @@ def test_unaligned_resume_never_runs_past_its_block( pos, ends = resume_at, [] while pos < prompt_len: request.num_computed_tokens = pos - num_new = _split(request, prompt_len - pos, partial_hit=partial_hit) + num_new = _split( + request, + prompt_len - pos, + use_eagle=use_eagle, + partial_hit=partial_hit, + num_prefill_checkpoint_blocks=num_prefill_checkpoint_blocks, + ) assert num_new > 0, f"no progress at {pos}" if pos % MAMBA_BLOCK_SIZE != 0: block_end = (pos // MAMBA_BLOCK_SIZE + 1) * MAMBA_BLOCK_SIZE diff --git a/tests/v1/core/test_prefix_caching.py b/tests/v1/core/test_prefix_caching.py index 038d45be30..6fc2131b6e 100644 --- a/tests/v1/core/test_prefix_caching.py +++ b/tests/v1/core/test_prefix_caching.py @@ -1089,6 +1089,7 @@ def test_hybrid_cache_mamba_align_shared_prefix_detection(): use_eagle=False, hash_block_size=block_size, mamba_partial_cache_hit=False, + mamba_has_prefill_checkpoint_blocks=False, ) req_2.shared_prefix_boundary = shared_prefix_boundary num_new_tokens_adjusted = Scheduler._mamba_block_aligned_split( diff --git a/vllm/models/kimi_k3/nvidia/kda.py b/vllm/models/kimi_k3/nvidia/kda.py index ec6b514f88..d2ca6f42f5 100644 --- a/vllm/models/kimi_k3/nvidia/kda.py +++ b/vllm/models/kimi_k3/nvidia/kda.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable +from dataclasses import replace import torch from einops import rearrange @@ -45,7 +46,10 @@ from vllm.models.kimi_k3.nvidia.kda_metadata import ( from vllm.platforms import current_platform from vllm.third_party.flash_linear_attention.ops.kda import FusedRMSNormGated from vllm.transformers_utils.configs.kimi_linear import KimiLinearConfig +from vllm.triton_utils import tl, triton from vllm.v1.attention.backend import AttentionBackend +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID +from vllm.v1.kv_cache_interface import MambaSpec from vllm.v1.worker.workspace import current_workspace_manager logger = init_logger(__name__) @@ -199,6 +203,8 @@ def _flashkda_prefill( out: torch.Tensor, final_state: torch.Tensor, workspace: torch.Tensor, + checkpoint_state: torch.Tensor | None = None, + checkpoint_offsets: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: import vllm._flashkda_C # noqa: F401 @@ -221,10 +227,75 @@ def _flashkda_prefill( initial_state.contiguous(), final_state, cu_seqlens.contiguous(), + checkpoint_state, + checkpoint_offsets.contiguous() if checkpoint_offsets is not None else None, ) return out, final_state +@triton.jit +def _store_cache_checkpoints_kernel( + x_ptr, + conv_state_ptr, + recurrent_checkpoint_ptr, + recurrent_state_ptr, + query_start_loc_ptr, + checkpoint_offsets_ptr, + checkpoint_state_indices_ptr, + x_stride_0: tl.constexpr, + x_stride_1: tl.constexpr, + state_stride_0: tl.constexpr, + state_stride_1: tl.constexpr, + state_stride_2: tl.constexpr, + checkpoint_stride_0: tl.constexpr, + recurrent_state_stride_0: tl.constexpr, + checkpoint_offset_stride: tl.constexpr, + STATE_LEN: tl.constexpr, + WIDTH: tl.constexpr, + RECURRENT_ROW_SIZE: tl.constexpr, + NULL_STATE_IDX: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + # store checkpoints to cache + seq_idx = tl.program_id(0) + cols = tl.program_id(1) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + state_idx = tl.load(checkpoint_state_indices_ptr + seq_idx) + checkpoint_offset = tl.load( + checkpoint_offsets_ptr + seq_idx * checkpoint_offset_stride + ) + valid_checkpoint = (state_idx != NULL_STATE_IDX) & (checkpoint_offset > 0) + valid_conv = ( + (cols < WIDTH * STATE_LEN) & valid_checkpoint & (checkpoint_offset >= STATE_LEN) + ) + width_idx = cols // STATE_LEN + history_idx = cols % STATE_LEN + checkpoint_end = tl.load(query_start_loc_ptr + seq_idx) + checkpoint_offset + token_idx = checkpoint_end - STATE_LEN + history_idx + values = tl.load( + x_ptr + token_idx * x_stride_0 + width_idx * x_stride_1, + mask=valid_conv, + ) + tl.store( + conv_state_ptr + + state_idx * state_stride_0 + + width_idx * state_stride_1 + + history_idx * state_stride_2, + values, + mask=valid_conv, + ) + + valid_recurrent = (cols < RECURRENT_ROW_SIZE) & valid_checkpoint + recurrent = tl.load( + recurrent_checkpoint_ptr + seq_idx * checkpoint_stride_0 + cols, + mask=valid_recurrent, + ) + tl.store( + recurrent_state_ptr + state_idx * recurrent_state_stride_0 + cols, + recurrent, + mask=valid_recurrent, + ) + + def resolve_kda_prefill_backend( backend: str, head_dim: int, @@ -471,6 +542,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): self._flashkda_buffer_specs = ( ((1, T, H, D), self.model_config.dtype), ((N, H, D, D), self.get_state_dtype()[1]), + ((N, H, D, D), self.get_state_dtype()[1]), ((workspace_size,), torch.uint8), ) @@ -518,6 +590,14 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self + def get_kv_cache_spec(self, vllm_config: VllmConfig) -> MambaSpec: + spec = super().get_kv_cache_spec(vllm_config) + assert isinstance(spec, MambaSpec) + return replace( + spec, + num_prefill_checkpoint_blocks=int(self.kda_prefill_backend == "flashkda"), + ) + def forward( self, hidden_states: torch.Tensor, @@ -590,6 +670,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): spec_query_start_loc = m.spec_query_start_loc num_accepted_tokens = m.num_accepted_tokens num_actual_tokens = m.num_actual_tokens + checkpoint = m.checkpoint has_spec_decode = m.num_spec_decodes > 0 mixed_qkv = mixed_qkv[:num_actual_tokens] g1 = g1[:, :num_actual_tokens] @@ -780,7 +861,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): if self.kda_prefill_backend == "flashkda": assert self.gate_lower_bound is not None assert self._flashkda_buffer_specs is not None - workspace_out, final_state, workspace = ( + workspace_out, final_state, checkpoint_state, workspace = ( current_workspace_manager().get_simultaneous( *self._flashkda_buffer_specs ) @@ -788,24 +869,85 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): flashkda_out = ( workspace_out if has_spec_decode else core_attn_out )[:, : q_ns.shape[1]] - ( - core_attn_out_non_spec, - last_recurrent_state, - ) = _flashkda_prefill( - q=q_ns, - k=k_ns, - v=v_ns, - g=g1_ns, - beta=beta_ns, - A_log=self.A_log, - dt_bias=self.dt_bias, - lower_bound=self.gate_lower_bound, - initial_state=initial_state, - cu_seqlens=non_spec_query_start_loc, - out=flashkda_out, - final_state=final_state[: initial_state.shape[0]], - workspace=workspace, - ) + if checkpoint is not None: + assert non_spec_query_start_loc is not None + num_sequences = initial_state.shape[0] + assert checkpoint.checkpoint_offsets.shape == (num_sequences,) + final_state = final_state[:num_sequences] + checkpoint_state = checkpoint_state[:num_sequences] + checkpoint_offsets = checkpoint.checkpoint_offsets + _flashkda_prefill( + q=q_ns, + k=k_ns, + v=v_ns, + g=g1_ns, + beta=beta_ns, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=initial_state, + cu_seqlens=non_spec_query_start_loc, + out=flashkda_out, + final_state=final_state, + workspace=workspace, + checkpoint_state=checkpoint_state, + checkpoint_offsets=checkpoint_offsets, + ) + core_attn_out_non_spec = flashkda_out + last_recurrent_state = final_state + state_len = conv_state.shape[-1] + width = mixed_qkv_ns.shape[-1] + recurrent_row_size = checkpoint_state[0].numel() + block_size = 256 + _store_cache_checkpoints_kernel[ + ( + checkpoint_offsets.numel(), + triton.cdiv( + max(width * state_len, recurrent_row_size), + block_size, + ), + ) + ]( + mixed_qkv_ns, + conv_state, + checkpoint_state, + recurrent_state, + non_spec_query_start_loc, + checkpoint_offsets, + checkpoint.state_indices, + mixed_qkv_ns.stride(0), + mixed_qkv_ns.stride(1), + conv_state.stride(0), + conv_state.stride(1), + conv_state.stride(2), + checkpoint_state.stride(0), + recurrent_state.stride(0), + checkpoint_offsets.stride(0), + state_len, + width, + recurrent_row_size, + NULL_BLOCK_ID, + block_size, + ) + else: + ( + core_attn_out_non_spec, + last_recurrent_state, + ) = _flashkda_prefill( + q=q_ns, + k=k_ns, + v=v_ns, + g=g1_ns, + beta=beta_ns, + A_log=self.A_log, + dt_bias=self.dt_bias, + lower_bound=self.gate_lower_bound, + initial_state=initial_state, + cu_seqlens=non_spec_query_start_loc, + out=flashkda_out, + final_state=final_state[: initial_state.shape[0]], + workspace=workspace, + ) else: ( core_attn_out_non_spec, diff --git a/vllm/models/kimi_k3/nvidia/kda_metadata.py b/vllm/models/kimi_k3/nvidia/kda_metadata.py index 2154eaa3c5..b903beac24 100644 --- a/vllm/models/kimi_k3/nvidia/kda_metadata.py +++ b/vllm/models/kimi_k3/nvidia/kda_metadata.py @@ -44,6 +44,9 @@ if TYPE_CHECKING: ) +FLASHKDA_CHUNK_SIZE = 16 + + @cache def _metadata_launch_pdl() -> bool: return current_platform.is_arch_support_pdl() @@ -257,12 +260,19 @@ class KDARecoverSSMCommitMetadata: align: KDARecoverSSMAlignMetadata | None +@dataclass +class KDACheckpointMetadata: + checkpoint_offsets: torch.Tensor + state_indices: torch.Tensor + + @dataclass class KimiK3KDAMetadata(GDNAttentionMetadata, RecoverSSMMetadata): recoverssm_commit: KDARecoverSSMCommitMetadata | None = None recoverssm_context: "KDARecoverSSMCommitContext | None" = field( default=None, repr=False, compare=False ) + checkpoint: KDACheckpointMetadata | None = None def commit_recoverssm_state( self, num_accepted_tokens: torch.Tensor @@ -557,6 +567,62 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): else: has_initial_state = None + checkpoint = None + if ( + num_prefills > 0 + and self.kv_cache_spec.num_prefill_checkpoint_blocks > 0 + and self.vllm_config.cache_config.mamba_cache_mode == "align" + ): + # prepare checkpoint metadata + assert m.seq_lens_cpu_upper_bound is not None + request_rows = list(range(m.num_reqs)) + if spec_sequence_masks_cpu is not None: + # get non spec request rows + request_rows = active_non_spec_mask_cpu.nonzero().flatten().tolist() + all_query_lens = query_start_loc_cpu.diff().tolist() + query_lens = [all_query_lens[row] for row in request_rows] + seq_lens = m.seq_lens_cpu_upper_bound.tolist() + block_size = self.kv_cache_spec.block_size + checkpoint_splits = [] + checkpoint_cols = [] + for row, query_len in zip(request_rows, query_lens): + seq_len = seq_lens[row] + offset = seq_len // block_size * block_size - (seq_len - query_len) + # offset should be less than query_len + valid = ( + seq_len % block_size != 0 + and 0 < offset < query_len + and offset % FLASHKDA_CHUNK_SIZE == 0 + ) + offset = offset if valid else 0 + first_len = offset or query_len + checkpoint_splits.append((first_len, query_len - first_len)) + checkpoint_cols.append(seq_len // block_size - 1 if valid else -1) + if any(tail for _, tail in checkpoint_splits): + checkpoint_offsets_tensor = async_tensor_h2d( + [first if tail else 0 for first, tail in checkpoint_splits], + dtype=torch.int32, + device=query_start_loc.device, + ) + request_rows_tensor = async_tensor_h2d( + request_rows, dtype=torch.int64, device=query_start_loc.device + ) + checkpoint_cols_tensor = async_tensor_h2d( + checkpoint_cols, dtype=torch.int64, device=query_start_loc.device + ) + checkpoint_state_indices = m.block_table_tensor[ + request_rows_tensor, checkpoint_cols_tensor + ] + checkpoint_state_indices = torch.where( + checkpoint_cols_tensor >= 0, + checkpoint_state_indices, + NULL_BLOCK_ID, + ) + checkpoint = KDACheckpointMetadata( + checkpoint_offsets_tensor, + checkpoint_state_indices, + ) + # Prepare per-request tensors for cudagraph replay. num_actual_tokens # may be token-padded, while state/query/acceptance metadata is indexed # by request. @@ -647,6 +713,7 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder): nums_dict=nums_dict, batch_ptr=batch_ptr, token_chunk_offset_ptr=token_chunk_offset_ptr, + checkpoint=checkpoint, ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index b41727f273..c06760b806 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -53,7 +53,7 @@ from vllm.v1.core.sched.request_queue import ( ) from vllm.v1.core.sched.utils import check_stop, remove_all from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs -from vllm.v1.kv_cache_interface import KVCacheConfig +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec from vllm.v1.metrics.perf import ModelMetrics, PerfStats from vllm.v1.metrics.stats import ( PrefixCacheStats, @@ -324,6 +324,16 @@ class Scheduler(SchedulerInterface): self.need_mamba_block_aligned_split = ( self.has_mamba_layers and self.cache_config.mamba_cache_mode == "align" ) + self.mamba_has_prefill_checkpoint_blocks = ( + self.has_mamba_layers + # TODO: support spec decoding + and not self.use_eagle + and all( + not isinstance(group.kv_cache_spec, MambaSpec) + or group.kv_cache_spec.num_prefill_checkpoint_blocks > 0 + for group in kv_cache_config.kv_cache_groups + ) + ) # A finer prefix_match_unit is configured: a mamba partial tail entry # can only be registered by a step ending exactly at the prompt's last # hash boundary, so the split adds that stop. @@ -406,12 +416,17 @@ class Scheduler(SchedulerInterface): last_cache_position = max(last_cache_position - block_size, 0) end = start + num_new_tokens + use_internal_checkpoint = ( + self.mamba_has_prefill_checkpoint_blocks and start % block_size == 0 + ) + if use_internal_checkpoint: + last_cache_position = 0 # Invariant: slot p holds the state after exactly (p + 1) * block_size # tokens. State is written at chunk ends, so chunk ends must be block # aligned. Exempt: the prompt's last chunk, whose slot decode advances # to the boundary. A block too wide for one chunk advances sub-block # and re-aligns at the next boundary. - if end < prefill_end: + if end < prefill_end and not use_internal_checkpoint: max_prefill_tokens = self.max_num_scheduled_tokens long_prefill_threshold = self.scheduler_config.long_prefill_token_threshold if long_prefill_threshold > 0: diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index 75acb7f497..ceb56a5987 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1285,6 +1285,9 @@ class MambaManager(SingleTypeKVCacheManager): self.last_state_block_idx: dict[str, int] = {} # The set of the requests that have been allocated blocks self._allocated_block_reqs: set[str] = set() + # Number of internal checkpoint blocks required by each request's + # current allocation. + self._num_checkpoint_blocks: dict[str, int] = {} # Requests that registered their own last-prompt-boundary partial # tail (producers). On the next step's CoW the boundary state moves # into a private cow_block; we record that block for connector @@ -1464,6 +1467,23 @@ class MambaManager(SingleTypeKVCacheManager): """ return 0 + def _needs_internal_checkpoint( + self, + request_id: str, + num_tokens: int, + num_computed_tokens: int, + ) -> bool: + assert isinstance(self.kv_cache_spec, MambaSpec) + checkpoint_idx = cdiv(num_tokens, self.block_size) - 2 + blocks = self.req_to_blocks[request_id] + return ( + self.kv_cache_spec.num_prefill_checkpoint_blocks > 0 + and num_tokens % self.block_size != 0 + and num_computed_tokens % self.block_size == 0 + and checkpoint_idx >= 0 + and (checkpoint_idx >= len(blocks) or blocks[checkpoint_idx].is_null) + ) + def get_num_blocks_to_allocate( self, request_id: str, @@ -1526,18 +1546,17 @@ class MambaManager(SingleTypeKVCacheManager): ) if has_partial_hit: num_new_blocks = max(num_new_blocks, 0) + 1 + checkpoint_block = int( + self._needs_internal_checkpoint( + request_id, num_tokens, total_computed_tokens + ) + ) + if not apply_admission_cap: + self._num_checkpoint_blocks[request_id] = checkpoint_block if num_new_blocks > 0: - if request_id in self._allocated_block_reqs: - # Old request. Needs at most 1 more blocks as we can reuse the - # speculative blocks in previous step. - num_new_blocks = 1 + int(has_partial_hit) - else: - # First prefill. Allocate 1 block for running state, the - # speculative blocks, and one extra block if a partial cache - # hit must be copy-on-written before the new tokens run. - num_new_blocks = ( - 1 + self.num_speculative_blocks + int(has_partial_hit) - ) + num_new_blocks = 1 + int(has_partial_hit) + checkpoint_block + if request_id not in self._allocated_block_reqs: + num_new_blocks += self.num_speculative_blocks num_evictable_computed_blocks = self._get_num_evictable_blocks( new_computed_blocks @@ -1569,6 +1588,7 @@ class MambaManager(SingleTypeKVCacheManager): num_required_blocks = ( cdiv(num_tokens, self.block_size) + self.num_speculative_blocks ) + checkpoint_block = self._num_checkpoint_blocks.get(request_id, 0) partial_hit = self._partial_hit_reqs.get(request_id) has_partial_hit = partial_hit is not None # `num_required_blocks` might be less than `len(req_blocks)` if blocks are @@ -1596,11 +1616,11 @@ class MambaManager(SingleTypeKVCacheManager): ) # null blocks if prev_block_len < num_skipped_blocks: + # minus the internal checkpoint block + # so we don't set null for that block + null_end = num_skipped_blocks - checkpoint_block req_blocks.extend( - [ - self._null_block - for _ in range(prev_block_len, num_skipped_blocks) - ] + [self._null_block for _ in range(prev_block_len, null_end)] ) if blocks_allocated: @@ -1616,12 +1636,10 @@ class MambaManager(SingleTypeKVCacheManager): num_new_blocks = num_required_blocks - len(req_blocks) if has_partial_hit: num_new_blocks = max(num_new_blocks, 0) + 1 - if blocks_allocated: - assert num_new_blocks <= 1 + int(has_partial_hit) - else: - assert num_new_blocks <= self.num_speculative_blocks + 1 + int( - has_partial_hit - ) + max_new_blocks = 1 + int(has_partial_hit) + checkpoint_block + if not blocks_allocated: + max_new_blocks += self.num_speculative_blocks + assert num_new_blocks <= max_new_blocks new_blocks = self.block_pool.get_new_blocks(num_new_blocks) returned_blocks = req_blocks[prev_block_len:] if partial_hit is not None: @@ -1670,6 +1688,7 @@ class MambaManager(SingleTypeKVCacheManager): if self.mamba_cache_mode == "align": self._allocated_block_reqs.discard(request_id) self.last_state_block_idx.pop(request_id, None) + self._num_checkpoint_blocks.pop(request_id, None) self._producer_partial_tail_reqs.pop(request_id, None) # A hand-off whose request died in this same scheduling pass must # not reach the connector: its unpin hook (free) has already run. diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index b660f49317..55c9e806a5 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -815,6 +815,7 @@ class MambaSpec(KVCacheSpec): mamba_type: MambaAttentionBackendEnum = MambaAttentionBackendEnum.MAMBA2 mamba_cache_mode: str = "none" num_speculative_blocks: int = 0 + num_prefill_checkpoint_blocks: int = 0 num_heads: int = 1 tokens_per_state: int = -1 @@ -843,7 +844,9 @@ class MambaSpec(KVCacheSpec): cdiv(max_model_len, self.block_size) + self.num_speculative_blocks ) * self.page_size_bytes elif vllm_config.cache_config.mamba_cache_mode == "align": - return self.page_size_bytes * (2 + self.num_speculative_blocks) + return self.page_size_bytes * ( + 2 + self.num_speculative_blocks + self.num_prefill_checkpoint_blocks + ) else: return self.page_size_bytes * (1 + self.num_speculative_blocks) @@ -865,6 +868,7 @@ class MambaSpec(KVCacheSpec): return all( isinstance(spec, MambaSpec) and spec.num_speculative_blocks == self.num_speculative_blocks + and spec.num_prefill_checkpoint_blocks == self.num_prefill_checkpoint_blocks for spec in kv_cache_specs.values() )