@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 70afdedc10

[K3] support recoverssm for K3 (#51855)

Signed-off-by: zjy0516 <riverclouds.zhu@qq.com> Signed-off-by: Benjamin Chislett <bchislett@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com> Co-authored-by: Benjamin Chislett <bchislett@nvidia.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Jiangyun Zhu committed Aug 17, 2026 at 19:34 UTC 70afdedc1081d28c3eaae53bece8292298484c86
19 files changed +2235 -75
tests/models/kimi_k3/test_kda.py
+342
@@ -6,6 +6,8 @@ Compares chunk_kda against a naive recurrent reference (float32).
6 Uses torch.rand for q/k/v to match FLA's test pattern.
7 """
8
9 +from types import SimpleNamespace
10 +
11 import pytest
12 import torch
13 import torch.nn.functional as F
@@ -22,6 +24,12 @@ from vllm.models.kimi_k3.nvidia.kda import (
24 is_flashkda_supported,
25 is_fused_kda_decode_supported,
26 )
27 +from vllm.models.kimi_k3.nvidia.model import KimiLinearForCausalLM
28 +from vllm.models.kimi_k3.nvidia.ops import recoverssm as recoverssm_ops
29 +from vllm.models.kimi_k3.nvidia.ops.recoverssm import (
30 + KDARecoverSSMCommitContext,
31 + kda_recoverssm_verify,
32 +)
33 from vllm.models.kimi_k3.nvidia.ops.third_party.kda import (
34 chunk_kda,
35 chunk_kda_with_fused_gate,
@@ -48,6 +56,38 @@ PACKED_DECODE_IMPLS = {
56 }
57
58
59 +def test_kda_recoverssm_config_state_layout():
60 + vllm_config = SimpleNamespace(
61 + model_config=SimpleNamespace(
62 + dtype=torch.bfloat16,
63 + hf_config=SimpleNamespace(
64 + linear_attn_config={
65 + "num_heads": 4,
66 + "head_dim": 32,
67 + "short_conv_kernel_size": 4,
68 + }
69 + ),
70 + ),
71 + cache_config=SimpleNamespace(
72 + mamba_cache_dtype="auto",
73 + use_kda_recoverssm=True,
74 + ),
75 + parallel_config=SimpleNamespace(tensor_parallel_size=1),
76 + speculative_config=SimpleNamespace(num_speculative_tokens=2),
77 + )
78 +
79 + assert KimiLinearForCausalLM.get_mamba_state_dtype_from_config(vllm_config) == (
80 + torch.bfloat16,
81 + torch.float32,
82 + torch.float32,
83 + torch.bfloat16,
84 + )
85 + assert KimiLinearForCausalLM.get_mamba_state_shape_from_config(vllm_config)[2:] == (
86 + (4, 3, 32),
87 + (4, 3, 64),
88 + )
89 +
90 +
91 @torch.inference_mode()
92 def test_gather_initial_states_correctness():
93 row_size = 8 * 128 * 128
@@ -535,6 +575,308 @@ def test_kda_spec_decode_correctness(
575 assert torch.isnan(output_storage[..., H * D :]).all()
576
577
578 +@pytest.mark.parametrize(
579 + (
580 + "conv_state_dim_first",
581 + "use_request_indices",
582 + "lower_bound",
583 + "align_mode",
584 + ),
585 + [
586 + pytest.param(False, False, None, False, id="baseline"),
587 + pytest.param(True, True, -5.0, True, id="all-features"),
588 + pytest.param(False, True, -5.0, False, id="request-indexed"),
589 + pytest.param(True, False, None, True, id="aligned"),
590 + ],
591 +)
592 +@torch.inference_mode()
593 +def test_kda_recoverssm_verify_and_group_commit(
594 + monkeypatch: pytest.MonkeyPatch,
595 + lower_bound: float | None,
596 + use_request_indices: bool,
597 + conv_state_dim_first: bool,
598 + align_mode: bool,
599 +):
600 + monkeypatch.setattr(
601 + recoverssm_ops,
602 + "is_conv_state_dim_first",
603 + lambda: conv_state_dim_first,
604 + )
605 + num_layers, num_seqs, query_len = 2, 2, 8
606 + num_blocks, num_heads, dim = (7 if align_mode else 3), 4, 128
607 + total_tokens = num_seqs * query_len
608 + torch.manual_seed(20260808)
609 +
610 + q, k, v, raw_g = [
611 + torch.randn(
612 + 1,
613 + total_tokens,
614 + num_heads,
615 + dim,
616 + dtype=torch.bfloat16,
617 + device=DEVICE,
618 + )
619 + for _ in range(4)
620 + ]
621 + raw_beta = torch.randn(
622 + 1,
623 + total_tokens,
624 + num_heads,
625 + dtype=torch.bfloat16,
626 + device=DEVICE,
627 + )
628 + query_start_loc = torch.arange(
629 + 0,
630 + total_tokens + 1,
631 + query_len,
632 + dtype=torch.int32,
633 + device=DEVICE,
634 + )
635 + state_indices = torch.tensor(
636 + [5, 6] if align_mode else [1, 2], dtype=torch.int32, device=DEVICE
637 + )
638 + accepted = [2, 8]
639 + if use_request_indices:
640 + global_num_accepted = torch.tensor(
641 + [0, accepted[0], 0, accepted[1]],
642 + dtype=torch.int32,
643 + device=DEVICE,
644 + )
645 + request_indices = torch.tensor([1, 3], dtype=torch.int32, device=DEVICE)
646 + else:
647 + global_num_accepted = torch.tensor(accepted, dtype=torch.int32, device=DEVICE)
648 + request_indices = None
649 +
650 + block_table = None
651 + num_computed_tokens = None
652 + mamba_block_size = None
653 + if align_mode:
654 + batch_size = 4 if use_request_indices else num_seqs
655 + block_table = torch.full((batch_size, 2), -1, dtype=torch.int32, device=DEVICE)
656 + rows = (
657 + request_indices
658 + if request_indices is not None
659 + else torch.arange(num_seqs, device=DEVICE)
660 + )
661 + block_table[rows] = torch.tensor(
662 + [[1, 5], [2, 6]],
663 + dtype=torch.int32,
664 + device=DEVICE,
665 + )
666 + num_computed_tokens = torch.zeros(batch_size, dtype=torch.int32, device=DEVICE)
667 + num_computed_tokens[rows] = 4
668 + mamba_block_size = 8
669 +
670 + layers = []
671 + expected_outputs = []
672 + expected_states = []
673 + initial_states = []
674 + initial_conv_states = []
675 + history_len, conv_dim = 3, 12
676 + for layer_idx in range(num_layers):
677 + A_log = (
678 + 0.2 * torch.randn(num_heads, dtype=torch.float32, device=DEVICE)
679 + + layer_idx * 0.03
680 + ).contiguous()
681 + dt_bias = (
682 + 0.1 * torch.randn(num_heads, dim, dtype=torch.float32, device=DEVICE)
683 + ).contiguous()
684 + checkpoint = 0.01 * torch.randn(
685 + num_blocks,
686 + num_heads,
687 + dim,
688 + dim,
689 + dtype=torch.float32,
690 + device=DEVICE,
691 + )
692 + conv_shape = (
693 + (num_blocks, conv_dim, history_len + query_len - 1)
694 + if conv_state_dim_first
695 + else (num_blocks, history_len + query_len - 1, conv_dim)
696 + )
697 + conv_state = torch.randn(conv_shape, dtype=torch.bfloat16, device=DEVICE)
698 + correction_cache = torch.empty(
699 + num_blocks,
700 + num_heads,
701 + query_len,
702 + dim,
703 + dtype=torch.float32,
704 + device=DEVICE,
705 + )
706 + kg_cache = torch.empty(
707 + num_blocks,
708 + num_heads,
709 + query_len,
710 + 2 * dim,
711 + dtype=torch.bfloat16,
712 + device=DEVICE,
713 + )
714 + layer = SimpleNamespace(
715 + kv_cache=(
716 + conv_state,
717 + checkpoint,
718 + correction_cache,
719 + kg_cache,
720 + ),
721 + A_log=A_log,
722 + dt_bias=dt_bias,
723 + local_num_heads=num_heads,
724 + head_dim=dim,
725 + gate_lower_bound=lower_bound,
726 + )
727 + layers.append(layer)
728 + initial_states.append(checkpoint.clone())
729 + initial_conv_states.append(conv_state.clone())
730 +
731 + actual_output = kda_recoverssm_verify(
732 + q=q,
733 + k=k,
734 + v=v,
735 + raw_g=raw_g,
736 + raw_beta=raw_beta,
737 + A_log=A_log,
738 + dt_bias=dt_bias,
739 + lower_bound=lower_bound,
740 + checkpoint_state=checkpoint,
741 + correction_cache=correction_cache,
742 + kg_cache=kg_cache,
743 + query_start_loc=query_start_loc,
744 + state_indices=state_indices,
745 + spec_query_len=query_len,
746 + )
747 +
748 + normalized_q = q.float() * torch.rsqrt(
749 + q.float().square().sum(dim=-1, keepdim=True) + 1e-6
750 + )
751 + normalized_k = k.float() * torch.rsqrt(
752 + k.float().square().sum(dim=-1, keepdim=True) + 1e-6
753 + )
754 + gate_input = raw_g.float() + dt_bias.view(1, 1, num_heads, dim)
755 + if lower_bound is None:
756 + gate = -A_log.exp().view(1, 1, num_heads, 1) * F.softplus(gate_input)
757 + else:
758 + gate = lower_bound * torch.sigmoid(
759 + A_log.exp().view(1, 1, num_heads, 1) * gate_input
760 + )
761 + beta = raw_beta.float().sigmoid()
762 +
763 + reference_output = []
764 + committed_states = checkpoint.clone()
765 + for seq_idx, commit_len in enumerate(accepted):
766 + start = seq_idx * query_len
767 + end = start + query_len
768 + output, _ = naive_recurrent_kda(
769 + normalized_q[:, start:end],
770 + normalized_k[:, start:end],
771 + v[:, start:end],
772 + gate[:, start:end],
773 + beta[:, start:end],
774 + initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2),
775 + )
776 + reference_output.append(output)
777 + _, committed_state = naive_recurrent_kda(
778 + normalized_q[:, start : start + commit_len],
779 + normalized_k[:, start : start + commit_len],
780 + v[:, start : start + commit_len],
781 + gate[:, start : start + commit_len],
782 + beta[:, start : start + commit_len],
783 + initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2),
784 + output_final_state=True,
785 + )
786 + assert committed_state is not None
787 + final_block = state_indices[seq_idx]
788 + if align_mode:
789 + assert block_table is not None
790 + row = request_indices[seq_idx] if use_request_indices else seq_idx
791 + final_block = block_table[row, (4 + commit_len) // 8]
792 + committed_states[final_block] = committed_state.transpose(-1, -2)
793 + if align_mode and 4 + commit_len >= 8:
794 + _, boundary_state = naive_recurrent_kda(
795 + normalized_q[:, start : start + 4],
796 + normalized_k[:, start : start + 4],
797 + v[:, start : start + 4],
798 + gate[:, start : start + 4],
799 + beta[:, start : start + 4],
800 + initial_state=checkpoint[state_indices[seq_idx]].transpose(-1, -2),
801 + output_final_state=True,
802 + )
803 + assert boundary_state is not None
804 + assert block_table is not None
805 + row = request_indices[seq_idx] if use_request_indices else seq_idx
806 + committed_states[block_table[row, 0]] = boundary_state.transpose(-1, -2)
807 + expected_outputs.append(torch.cat(reference_output, dim=1))
808 + expected_states.append(committed_states)
809 + torch.testing.assert_close(checkpoint, initial_states[-1])
810 + torch.testing.assert_close(
811 + actual_output,
812 + expected_outputs[-1],
813 + atol=3e-2,
814 + rtol=3e-2,
815 + )
816 +
817 + context = KDARecoverSSMCommitContext.create(
818 + layers,
819 + spec_query_len=query_len,
820 + max_num_reqs=global_num_accepted.shape[0],
821 + )
822 + context.commit(
823 + global_num_accepted,
824 + state_indices,
825 + query_start_loc,
826 + request_indices=request_indices,
827 + block_table=block_table,
828 + num_computed_tokens=num_computed_tokens,
829 + mamba_block_size=mamba_block_size,
830 + )
831 +
832 + for layer_idx, layer in enumerate(layers):
833 + torch.testing.assert_close(
834 + layer.kv_cache[1],
835 + expected_states[layer_idx],
836 + atol=3e-3,
837 + rtol=3e-3,
838 + )
839 + for seq_idx, commit_len in enumerate(accepted):
840 + block = state_indices[seq_idx]
841 + if align_mode:
842 + assert block_table is not None
843 + row = request_indices[seq_idx] if use_request_indices else seq_idx
844 + block = block_table[row, (4 + commit_len) // 8]
845 + source_block = state_indices[seq_idx] if align_mode else block
846 + if conv_state_dim_first:
847 + actual_conv = layer.kv_cache[0][block, :, :history_len]
848 + expected_conv = initial_conv_states[layer_idx][
849 + source_block,
850 + :,
851 + commit_len - 1 : commit_len - 1 + history_len,
852 + ]
853 + else:
854 + actual_conv = layer.kv_cache[0][block, :history_len]
855 + expected_conv = initial_conv_states[layer_idx][
856 + source_block,
857 + commit_len - 1 : commit_len - 1 + history_len,
858 + ]
859 + torch.testing.assert_close(actual_conv, expected_conv)
860 + if align_mode and 4 + commit_len >= 8:
861 + assert block_table is not None
862 + boundary_block = block_table[row, 0]
863 + if conv_state_dim_first:
864 + actual_boundary_conv = layer.kv_cache[0][
865 + boundary_block, :, :history_len
866 + ]
867 + expected_boundary_conv = initial_conv_states[layer_idx][
868 + state_indices[seq_idx], :, 3 : 3 + history_len
869 + ]
870 + else:
871 + actual_boundary_conv = layer.kv_cache[0][
872 + boundary_block, :history_len
873 + ]
874 + expected_boundary_conv = initial_conv_states[layer_idx][
875 + state_indices[seq_idx], 3 : 3 + history_len
876 + ]
877 + torch.testing.assert_close(actual_boundary_conv, expected_boundary_conv)
878 +
879 +
880 @pytest.mark.parametrize(
881 ("num_heads", "num_seqs", "lower_bound", "fuse_output_norm"),
882 [
tests/models/kimi_k3/test_kda_metadata.py
+145 -4
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from dataclasses import fields
5 +from unittest.mock import Mock
6
7 import pytest
8 import torch
@@ -26,6 +27,9 @@ from vllm.v1.attention.backends.gdn_attn import (
27 GDNAttentionMetadata,
28 GDNAttentionMetadataBuilder,
29 )
30 +from vllm.v1.attention.backends.recoverssm_metadata import (
31 + RecoverSSMPostprocessMetadata,
32 +)
33 from vllm.v1.attention.backends.utils import (
34 NULL_BLOCK_ID,
35 mamba_get_block_table_tensor,
@@ -44,8 +48,12 @@ PRUNED_METADATA_FIELDS = {
48 }
49
50
47 -def _assert_matches_shared_gdn(reference, actual: KimiK3KDAMetadata):
48 - for field in fields(KimiK3KDAMetadata):
51 +def _assert_matches_shared_gdn(
52 + reference: GDNAttentionMetadata, actual: KimiK3KDAMetadata
53 +):
54 + assert actual.recoverssm_commit is None
55 + assert actual.recoverssm_context is None
56 + for field in fields(GDNAttentionMetadata):
57 actual_value = getattr(actual, field.name)
58 expected_value = getattr(reference, field.name)
59 if field.name in PRUNED_METADATA_FIELDS:
@@ -78,6 +86,7 @@ def _make_builder(
86 full_cuda_graph: bool,
87 device: torch.device = DEVICE,
88 mamba_cache_mode: str = "none",
89 + use_recoverssm: bool = False,
90 ) -> AttentionMetadataBuilder:
91 vllm_config = create_vllm_config(
92 model_name="Qwen/Qwen3.5-0.8B",
@@ -92,17 +101,24 @@ def _make_builder(
101 CUDAGraphMode.FULL_AND_PIECEWISE if full_cuda_graph else CUDAGraphMode.NONE
102 )
103 vllm_config.cache_config.mamba_cache_mode = mamba_cache_mode
95 - return builder_cls(
104 + vllm_config.cache_config.use_replayssm = use_recoverssm
105 + vllm_config.cache_config.use_kda_recoverssm = use_recoverssm
106 + builder = builder_cls(
107 kv_cache_spec=MambaSpec(
108 block_size=BLOCK_SIZE,
109 shapes=((16, 64),),
110 dtypes=(torch.float16,),
100 - num_speculative_blocks=num_speculative_tokens,
111 + mamba_cache_mode=mamba_cache_mode,
112 + num_speculative_blocks=(0 if use_recoverssm else num_speculative_tokens),
113 ),
114 layer_names=["layer.0"],
115 vllm_config=vllm_config,
116 device=device,
117 )
118 + if use_recoverssm:
119 + assert isinstance(builder, KimiK3KDAMetadataBuilder)
120 + builder.recoverssm_context = Mock()
121 + return builder
122
123
124 @pytest.mark.parametrize(
@@ -244,6 +260,99 @@ def test_mixed_regular_and_spec_decode_excludes_request_padding():
260 torch.testing.assert_close(actual.spec_token_indx, torch.tensor([1, 2, 3]))
261
262
263 +@pytest.mark.parametrize("mamba_cache_mode", ["none", "align"])
264 +def test_recoverssm_spec_uses_one_state_slot_and_current_window(
265 + mamba_cache_mode: str,
266 +):
267 + if mamba_cache_mode == "align" and not torch.cuda.is_available():
268 + pytest.skip("align metadata construction requires CUDA")
269 + device = torch.device("cuda") if mamba_cache_mode == "align" else DEVICE
270 + batch = BatchSpec(seq_lens=[100, 65, 20], query_lens=[1, 1, 3])
271 + common_attn_metadata = create_common_attn_metadata(
272 + batch, BLOCK_SIZE, device
273 + ).replace(is_prefilling=torch.tensor([True, True, False]))
274 + builder = _make_builder(
275 + KimiK3KDAMetadataBuilder,
276 + num_speculative_tokens=2,
277 + full_cuda_graph=False,
278 + device=device,
279 + mamba_cache_mode=mamba_cache_mode,
280 + use_recoverssm=True,
281 + )
282 + assert isinstance(builder, KimiK3KDAMetadataBuilder)
283 + context = builder.recoverssm_context
284 + assert context is not None
285 + actual = builder.build(
286 + 0,
287 + common_attn_metadata,
288 + num_decode_draft_tokens_cpu=torch.tensor([-1, -1, 2], dtype=torch.int32),
289 + num_accepted_tokens=torch.tensor([3, 2, 2], dtype=torch.int32, device=device),
290 + )
291 +
292 + assert actual.spec_state_indices_tensor is not None
293 + assert actual.spec_state_indices_tensor.shape == (1, 1)
294 + torch.testing.assert_close(
295 + actual.num_accepted_tokens,
296 + torch.ones(1, dtype=torch.int32, device=device),
297 + )
298 + commit_metadata = actual.recoverssm_commit
299 + assert commit_metadata is not None
300 + torch.testing.assert_close(
301 + commit_metadata.request_indices,
302 + torch.tensor([2], dtype=torch.int32, device=device),
303 + )
304 + assert actual.recoverssm_context is context
305 + num_accepted_tokens = torch.tensor([3, 2, 1], dtype=torch.int32, device=device)
306 +
307 + postprocess = actual.commit_recoverssm_state(num_accepted_tokens)
308 +
309 + if mamba_cache_mode == "none":
310 + assert commit_metadata.align is None
311 + assert postprocess is None
312 + else:
313 + assert isinstance(postprocess, RecoverSSMPostprocessMetadata)
314 + assert postprocess.num_spec_decodes == 1
315 + assert postprocess.request_indices is commit_metadata.request_indices
316 + assert postprocess.block_table is common_attn_metadata.block_table_tensor
317 + assert (
318 + postprocess.num_computed_tokens
319 + is common_attn_metadata.compute_num_computed_tokens()
320 + )
321 + assert postprocess.block_size == BLOCK_SIZE
322 + args = context.commit.call_args.args
323 + assert args[0] is num_accepted_tokens
324 + torch.testing.assert_close(args[1], commit_metadata.state_indices[:, 0])
325 + torch.testing.assert_close(args[2], commit_metadata.query_start_loc)
326 +
327 +
328 +def test_recoverssm_distinguishes_draftless_decode_from_one_token_prefill():
329 + batch = BatchSpec(seq_lens=[40, 30], query_lens=[1, 1])
330 + common_attn_metadata = create_common_attn_metadata(
331 + batch, BLOCK_SIZE, DEVICE
332 + ).replace(is_prefilling=torch.tensor([False, True]))
333 + actual = _make_builder(
334 + KimiK3KDAMetadataBuilder,
335 + num_speculative_tokens=2,
336 + full_cuda_graph=False,
337 + use_recoverssm=True,
338 + ).build(
339 + 0,
340 + common_attn_metadata,
341 + num_decode_draft_tokens_cpu=torch.full((2,), -1, dtype=torch.int32),
342 + num_accepted_tokens=torch.ones(2, dtype=torch.int32),
343 + )
344 +
345 + assert actual.num_spec_decodes == 1
346 + assert actual.num_decodes == 0
347 + assert actual.num_prefills == 1
348 + assert actual.spec_state_indices_tensor is not None
349 + assert actual.spec_state_indices_tensor.shape == (1, 1)
350 + torch.testing.assert_close(
351 + actual.spec_query_start_loc,
352 + torch.tensor([0, 1], dtype=torch.int32),
353 + )
354 +
355 +
356 @pytest.mark.parametrize(
357 ("seq_len", "expected_has_initial_state"),
358 [
@@ -307,6 +416,38 @@ def test_kimi_k3_kda_cudagraph_capture_matches_shared_gdn():
416 _assert_matches_shared_gdn(reference, actual)
417
418
419 +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
420 +def test_recoverssm_spec_cudagraph_stages_one_checkpoint_per_request():
421 + device = torch.device("cuda")
422 + batch = BatchSpec(seq_lens=[50, 30], query_lens=[3, 3])
423 + common_attn_metadata = create_common_attn_metadata(
424 + batch, BLOCK_SIZE, device
425 + ).replace(is_prefilling=torch.tensor([False, False]))
426 + builder = _make_builder(
427 + KimiK3KDAMetadataBuilder,
428 + num_speculative_tokens=2,
429 + full_cuda_graph=True,
430 + device=device,
431 + use_recoverssm=True,
432 + )
433 + assert isinstance(builder, KimiK3KDAMetadataBuilder)
434 + assert builder.spec_state_indices_tensor.shape == (
435 + builder.vllm_config.scheduler_config.max_num_seqs,
436 + 1,
437 + )
438 + actual = builder.build_for_cudagraph_capture(common_attn_metadata)
439 +
440 + assert actual.spec_state_indices_tensor is not None
441 + assert actual.spec_state_indices_tensor.shape == (batch.batch_size, 1)
442 + assert actual.num_accepted_tokens is not None
443 + torch.testing.assert_close(
444 + actual.num_accepted_tokens,
445 + torch.ones(batch.batch_size, dtype=torch.int32, device=device),
446 + )
447 + assert actual.recoverssm_commit is not None
448 + assert actual.recoverssm_commit.request_indices is None
449 +
450 +
451 def test_kimi_k3_kda_backend_uses_private_metadata_builder():
452 assert KimiK3KDAAttentionBackend.get_builder_cls() is KimiK3KDAMetadataBuilder
453 assert KimiK3KDAAttentionBackend.is_ssm()
tests/models/test_registry.py
+3 -1
@@ -143,8 +143,10 @@ def test_registry_is_pp(model_arch, is_pp, init_cuda):
143 @pytest.mark.parametrize(
144 "model_arch,supported",
145 [
146 - # ReplaySSM is opt-in per model; only Nemotron-H sets the flag today.
146 + # ReplaySSM is opt-in per model.
147 ("NemotronHForCausalLM", True),
148 + ("KimiLinearForCausalLM", not current_platform.is_rocm()),
149 + ("KimiK3ForConditionalGeneration", not current_platform.is_rocm()),
150 ("Mamba2ForCausalLM", False),
151 ("Zamba2ForCausalLM", False),
152 ],
tests/test_config.py
+47
@@ -29,6 +29,7 @@ from vllm.config import (
29 from vllm.config.compilation import CompilationMode, CUDAGraphMode
30 from vllm.config.kernel import IrOpPriorityConfig
31 from vllm.config.load import LoadConfig
32 +from vllm.config.mamba import MambaBackendEnum
33 from vllm.config.utils import get_field
34 from vllm.config.vllm import OPTIMIZATION_LEVEL_TO_CONFIG, OptimizationLevel
35 from vllm.platforms import current_platform
@@ -37,6 +38,52 @@ from vllm.v1.attention.backend import AttentionCGSupport
38 DEVICE_TYPE = current_platform.device_type
39
40
41 +def test_kda_recoverssm_derivation_is_revalidated():
42 + config = SimpleNamespace(
43 + cache_config=SimpleNamespace(
44 + use_replayssm=True,
45 + use_kda_recoverssm=False,
46 + mamba_cache_mode="none",
47 + ),
48 + num_speculative_tokens=3,
49 + model_config=SimpleNamespace(
50 + supports_replayssm=True,
51 + architecture="KimiLinearForCausalLM",
52 + ),
53 + mamba_config=SimpleNamespace(
54 + backend=MambaBackendEnum.TRITON,
55 + enable_stochastic_rounding=False,
56 + ),
57 + parallel_config=SimpleNamespace(pipeline_parallel_size=1),
58 + kv_transfer_config=None,
59 + use_v2_model_runner=True,
60 + )
61 +
62 + VllmConfig.validate_mamba_cached_kernel(config)
63 + assert config.cache_config.use_replayssm
64 + assert config.cache_config.use_kda_recoverssm
65 +
66 + config.cache_config.mamba_cache_mode = "align"
67 + VllmConfig.validate_mamba_cached_kernel(config)
68 + config.use_v2_model_runner = False
69 + with pytest.raises(ValueError, match="VLLM_USE_V2_MODEL_RUNNER=1"):
70 + VllmConfig.validate_mamba_cached_kernel(config)
71 + config.use_v2_model_runner = True
72 + config.cache_config.mamba_cache_mode = "all"
73 + with pytest.raises(ValueError, match="only none and align"):
74 + VllmConfig.validate_mamba_cached_kernel(config)
75 + config.cache_config.mamba_cache_mode = "none"
76 +
77 + config.model_config.architecture = "NemotronHForCausalLM"
78 + with pytest.raises(ValueError, match="only supported for Kimi-K3 KDA"):
79 + VllmConfig.validate_mamba_cached_kernel(config)
80 +
81 + config.model_config.architecture = "KimiLinearForCausalLM"
82 + config.parallel_config.pipeline_parallel_size = 2
83 + with pytest.raises(ValueError, match="pipeline_parallel_size=1"):
84 + VllmConfig.validate_mamba_cached_kernel(config)
85 +
86 +
87 def test_compile_config_repr_succeeds():
88 # setup: VllmBackend mutates the config object
89 config = VllmConfig()
tests/v1/worker/test_mamba_hybrid_model_state.py
+67
@@ -1,11 +1,19 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +from types import SimpleNamespace
5 +from unittest.mock import Mock
6 +
7 import pytest
8 import torch
9
10 from vllm.platforms import current_platform
11 +from vllm.v1.attention.backends.recoverssm_metadata import (
12 + RecoverSSMMetadata,
13 + RecoverSSMPostprocessMetadata,
14 +)
15 from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState
16 +from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState
17
18
19 @pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
@@ -18,6 +26,7 @@ def test_postprocess_state_scalar_with_int32_mapping(
26 (4,), 9, dtype=torch.int32, device="cuda"
27 )
28 state._align_mode = False
29 + state.recoverssm = None
30 state._mamba_ctx = None
31 idx_mapping = torch.tensor([2, -1, 0], dtype=torch.int32, device="cuda")
32
@@ -27,3 +36,61 @@ def test_postprocess_state_scalar_with_int32_mapping(
36 [expected_value, 9, expected_value, 9], dtype=torch.int32, device="cuda"
37 )
38 torch.testing.assert_close(state.num_accepted_tokens_gpu, expected)
39 +
40 +
41 +def test_recoverssm_commits_accepted_window_after_v2_sampling() -> None:
42 + state = RecoverSSMState()
43 + metadata = Mock(spec=RecoverSSMMetadata)
44 + metadata.commit_recoverssm_state.return_value = None
45 + num_sampled = torch.tensor([3, 1], dtype=torch.int32)
46 + idx_mapping = torch.tensor([0, 1], dtype=torch.int32)
47 + num_accepted_tokens = torch.ones(2, dtype=torch.int32)
48 + group = SimpleNamespace(layer_names=["layer"])
49 +
50 + state.record_step({"layer": metadata}, [[group]], for_capture=False)
51 + state.commit_step(
52 + num_sampled,
53 + idx_mapping,
54 + state_indices=None,
55 + num_accepted_tokens=num_accepted_tokens,
56 + )
57 + state.commit_step(
58 + num_sampled,
59 + idx_mapping,
60 + state_indices=None,
61 + num_accepted_tokens=num_accepted_tokens,
62 + )
63 +
64 + metadata.commit_recoverssm_state.assert_called_once_with(num_sampled)
65 +
66 +
67 +@pytest.mark.skipif(not current_platform.is_cuda(), reason="Requires CUDA")
68 +def test_recoverssm_align_tracks_mixed_batch_state_and_neutralizes_copy_bias() -> None:
69 + state = object.__new__(MambaHybridModelState)
70 + state._align_mode = True
71 + state._mamba_ctx = None
72 + state._mamba_state_idx_gpu = torch.full((5,), -1, dtype=torch.int32, device="cuda")
73 + state.recoverssm = RecoverSSMState()
74 + state.num_accepted_tokens_gpu = torch.full(
75 + (5,), 9, dtype=torch.int32, device="cuda"
76 + )
77 + metadata = Mock(spec=RecoverSSMMetadata)
78 + metadata.commit_recoverssm_state.return_value = RecoverSSMPostprocessMetadata(
79 + num_spec_decodes=1,
80 + request_indices=torch.tensor([1], dtype=torch.int32, device="cuda"),
81 + num_computed_tokens=torch.tensor([6, 7], dtype=torch.int32, device="cuda"),
82 + block_size=8,
83 + block_table=torch.zeros((2, 4), dtype=torch.int32, device="cuda"),
84 + )
85 + num_sampled = torch.tensor([2, 3], dtype=torch.int32, device="cuda")
86 + idx_mapping = torch.tensor([3, 1], dtype=torch.int32, device="cuda")
87 + group = SimpleNamespace(layer_names=["layer"])
88 +
89 + state.recoverssm.record_step({"layer": metadata}, [[group]], for_capture=False)
90 +
91 + state.postprocess_state(idx_mapping, num_sampled)
92 +
93 + expected_state_indices = [-1, 1, -1, -1, -1]
94 + assert state._mamba_state_idx_gpu.tolist() == expected_state_indices
95 + expected_accepted = [9, 1, 9, 2, 9]
96 + assert state.num_accepted_tokens_gpu.tolist() == expected_accepted
tests/v1/worker/test_mamba_utils.py
+36
@@ -490,6 +490,42 @@ def test_stage_postprocess_inputs_to_gpu_asserts_on_missing_state_idx():
490 )
491
492
493 +def test_gpu_context_ignores_auxiliary_cache_tensors() -> None:
494 + device = torch.device("cpu")
495 + config = _TestConfig(num_layers=1)
496 + layer_names = ["layer_0"]
497 + kv_cache_config = _make_kv_cache_config(config, layer_names)
498 + conv_state = torch.empty(
499 + config.num_blocks,
500 + config.conv_width,
501 + config.conv_inner_dim,
502 + dtype=config.dtype,
503 + )
504 + temporal_state = torch.empty(
505 + config.num_blocks, config.temporal_state_dim, dtype=config.dtype
506 + )
507 + attention = MagicMock()
508 + attention.kv_cache = [
509 + conv_state,
510 + temporal_state,
511 + *(torch.empty(config.num_blocks, 1) for _ in range(4)),
512 + ]
513 + context = _make_gpu_ctx(config, kv_cache_config, device)
514 +
515 + context.initialize_from_forward_context(
516 + kv_cache_config,
517 + {"layer_0": attention},
518 + _COPY_FUNCS,
519 + [torch.zeros(1, 1, dtype=torch.int32)],
520 + )
521 +
522 + assert context.is_initialized
523 + assert context.state_base_addrs.tolist() == [
524 + conv_state.data_ptr(),
525 + temporal_state.data_ptr(),
526 + ]
527 +
528 +
529 def _run_gpu_postprocess(
530 gpu_ctx: MambaSpecDecodeGPUContext,
531 *,
vllm/config/cache.py
+4 -3
@@ -144,9 +144,8 @@ class CacheConfig:
144 caching is enabled.
145 """
146 replayssm_buffer_len: int = Field(default=16, gt=0)
147 - """ReplaySSM history buffer length B: with use_replayssm, standard decode
148 - caches recent SSM inputs in a size-B ring buffer and flushes the checkpoint
149 - state to HBM every B steps. Default 16."""
147 + """ReplaySSM history buffer length B for standard Mamba2 decode. Kimi-K3
148 + speculative decoding does not use B. Default 16."""
149 use_replayssm: bool = False
150 """Use the ReplaySSM Mamba2 decode kernel: cache recent SSM inputs and skip
151 the per-step full-state store, writing the checkpoint back only on flush.
@@ -154,6 +153,8 @@ class CacheConfig:
153 mamba backend; standard (non-speculative) decode only. In align mode flushes
154 are most efficient when mamba_block_size is a multiple of replayssm_buffer_len,
155 but this is not required."""
156 + use_kda_recoverssm: bool = field(default=False, init=False)
157 + """Whether Kimi-K3 KDA uses RecoverSSM speculative decode."""
158
159 # Will be set after profiling.
160 num_gpu_blocks: int | None = field(default=None, init=False)
vllm/config/vllm.py
+33 -8
@@ -2540,22 +2540,47 @@ class VllmConfig:
2540 @model_validator(mode="after")
2541 def validate_mamba_cached_kernel(self) -> "VllmConfig":
2542 if not self.cache_config.use_replayssm:
2543 + self.cache_config.use_kda_recoverssm = False
2544 return self
2544 - # ReplaySSM adds a 3-tensor ring to the mamba state; only models that
2545 - # opt in (supports_replayssm) build a consistent shape on both the layer
2546 - # and config paths. Reject others so the mamba page size cannot desync.
2545 + self.cache_config.use_kda_recoverssm = self.num_speculative_tokens > 0
2546 +
2547 if self.model_config is not None and not self.model_config.supports_replayssm:
2548 raise ValueError(
2549 - "--use-replayssm is only supported for Nemotron-H models "
2550 - f"(got architecture {self.model_config.architecture!r})"
2549 + "--use-replayssm is not supported for architecture "
2550 + f"{self.model_config.architecture!r}"
2551 )
2552 - if self.cache_config.mamba_cache_mode == "all":
2552 + if self.cache_config.use_kda_recoverssm:
2553 + if self.model_config is not None and self.model_config.architecture not in (
2554 + "KimiLinearForCausalLM",
2555 + "KimiK3ForConditionalGeneration",
2556 + ):
2557 + raise ValueError("RecoverSSM is only supported for Kimi-K3 KDA")
2558 + if self.mamba_config.enable_stochastic_rounding:
2559 + raise ValueError(
2560 + "RecoverSSM supports bfloat16/float32 "
2561 + "SSM state caches, not --enable-mamba-cache-stochastic-"
2562 + "rounding, which requires an explicit float16 cache"
2563 + )
2564 + if self.cache_config.mamba_cache_mode not in ("none", "align"):
2565 + raise ValueError(
2566 + "RecoverSSM supports only none and align Mamba cache modes"
2567 + )
2568 + if (
2569 + self.cache_config.mamba_cache_mode == "align"
2570 + and not self.use_v2_model_runner
2571 + ):
2572 + raise ValueError(
2573 + "RecoverSSM with align mode requires VLLM_USE_V2_MODEL_RUNNER=1"
2574 + )
2575 + if self.parallel_config.pipeline_parallel_size > 1:
2576 + raise ValueError(
2577 + "RecoverSSM currently requires pipeline_parallel_size=1"
2578 + )
2579 + elif self.cache_config.mamba_cache_mode == "all":
2580 raise ValueError(
2581 "--use-replayssm supports prefix caching only in align mode; "
2582 "pass --mamba-cache-mode align"
2583 )
2557 - if self.num_speculative_tokens > 0:
2558 - raise ValueError("--use-replayssm does not support speculative decoding")
2584 if self.mamba_config.backend != MambaBackendEnum.TRITON:
2585 raise ValueError("--use-replayssm requires --mamba-backend triton")
2586 if (
vllm/model_executor/layers/mamba/abstract.py
+5 -3
@@ -71,10 +71,12 @@ class MambaBase(AttentionLayerBase):
71 page_size_padded=page_size_padded,
72 mamba_type=self.mamba_type,
73 mamba_cache_mode=vllm_config.cache_config.mamba_cache_mode,
74 + # RecoverSSM verifies the whole window off one checkpoint, so it
75 + # never writes the baseline's per-draft-token state slots.
76 num_speculative_blocks=(
75 - vllm_config.speculative_config.num_speculative_tokens
76 - if vllm_config.speculative_config
77 - else 0
77 + 0
78 + if vllm_config.cache_config.use_kda_recoverssm
79 + else vllm_config.num_speculative_tokens
80 ),
81 )
82
vllm/model_executor/layers/mamba/mamba_utils.py
+30
@@ -136,6 +136,15 @@ class MambaStateDtypeCalculator:
136 state_dtype = get_kv_cache_torch_dtype(mamba_cache_dtype, model_dtype)
137 return (state_dtype, torch.float32)
138
139 + @classmethod
140 + def append_kda_recoverssm_record(
141 + cls,
142 + base_dtypes: tuple[torch.dtype, ...],
143 + model_dtype: ModelDType | torch.dtype,
144 + ) -> tuple[torch.dtype, ...]:
145 + activation_dtype = get_kv_cache_torch_dtype("auto", model_dtype)
146 + return (*base_dtypes, torch.float32, activation_dtype)
147 +
148
149 class MambaStateShapeCalculator:
150 @classmethod
@@ -293,6 +302,27 @@ class MambaStateShapeCalculator:
302 recurrent_state_shape = (divide(num_heads, tp_world_size), head_dim, head_dim)
303 return (conv_state_shape, recurrent_state_shape)
304
305 + @classmethod
306 + def append_kda_recoverssm_record(
307 + cls,
308 + base_shapes: tuple[tuple[int, int], tuple[int, int, int]],
309 + num_heads: int,
310 + head_dim: int,
311 + tp_world_size: int,
312 + spec_query_len: int,
313 + ) -> tuple[
314 + tuple[int, int],
315 + tuple[int, int, int],
316 + tuple[int, int, int],
317 + tuple[int, int, int],
318 + ]:
319 + local_num_heads = divide(num_heads, tp_world_size)
320 + return (
321 + *base_shapes,
322 + (local_num_heads, spec_query_len, head_dim),
323 + (local_num_heads, spec_query_len, 2 * head_dim),
324 + )
325 +
326
327 @dataclass
328 class MambaCopySpec:
vllm/model_executor/models/interfaces.py
+8 -2
@@ -75,6 +75,12 @@ MambaStateShapes: TypeAlias = (
75 | tuple[tuple[int, int, int]]
76 | tuple[tuple[int, int], tuple[int, int]]
77 | tuple[tuple[int, int], tuple[int, int, int]]
78 + | tuple[
79 + tuple[int, int],
80 + tuple[int, int, int],
81 + tuple[int, int, int],
82 + tuple[int, int, int],
83 + ]
84 )
85
86
@@ -1099,8 +1105,8 @@ def supports_mamba_prefix_caching(
1105
1106 @runtime_checkable
1107 class SupportsReplaySSM(Protocol):
1102 - """The interface for models whose Mamba2 layers support ReplaySSM cached
1103 - standard decode.
1108 + """The interface for models whose recurrent layers support ReplaySSM
1109 + cached decode.
1110
1111 This is currently experimental.
1112 """
vllm/models/kimi_k3/nvidia/kda.py
+73 -22
@@ -282,23 +282,37 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
282
283 def get_state_dtype(
284 self,
285 - ) -> tuple[torch.dtype, torch.dtype]:
285 + ) -> tuple[torch.dtype, ...]:
286 if self.model_config is None or self.cache_config is None:
287 raise ValueError("model_config and cache_config must be set")
288 - return MambaStateDtypeCalculator.kda_state_dtype(
288 + base_dtypes = MambaStateDtypeCalculator.kda_state_dtype(
289 self.model_config.dtype, self.cache_config.mamba_cache_dtype
290 )
291 + if self.cache_config.use_kda_recoverssm:
292 + return MambaStateDtypeCalculator.append_kda_recoverssm_record(
293 + base_dtypes, self.model_config.dtype
294 + )
295 + return base_dtypes
296
297 def get_state_shape(
298 self,
294 - ) -> tuple[tuple[int, ...], tuple[int, ...]]:
295 - return MambaStateShapeCalculator.kda_state_shape(
299 + ) -> tuple[tuple[int, ...], ...]:
300 + base_shapes = MambaStateShapeCalculator.kda_state_shape(
301 self.tp_size,
302 self.num_heads,
303 self.head_dim,
304 conv_kernel_size=self.conv_size,
305 num_spec=self.num_spec,
306 )
307 + if self.cache_config.use_kda_recoverssm:
308 + return MambaStateShapeCalculator.append_kda_recoverssm_record(
309 + base_shapes,
310 + self.num_heads,
311 + self.head_dim,
312 + tp_world_size=self.tp_size,
313 + spec_query_len=1 + self.num_spec,
314 + )
315 + return base_shapes
316
317 def __init__(
318 self,
@@ -308,6 +322,12 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
322 run_gemm_rs: bool = False,
323 ) -> None:
324 super().__init__(config, vllm_config, prefix)
325 + self.use_recoverssm = self.cache_config.use_kda_recoverssm
326 + if self.cache_config.use_replayssm and not self.use_recoverssm:
327 + raise ValueError(
328 + "Kimi-K3 supports --use-replayssm only with speculative decoding"
329 + )
330 + self.spec_query_len = 1 + self.num_spec
331
332 kda_config = config.linear_attn_config # type: ignore[attr-defined]
333 assert kda_config is not None, "linear_attn_config must be set"
@@ -371,7 +391,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
391 self.conv1d.weight.data = self.conv1d.weight.data.unsqueeze(1)
392 # Keep a width-major copy for fused decode without changing the layout
393 # consumed by the prefill and fallback decode kernels.
374 - conv_state_dtype, _ = self.get_state_dtype()
394 + conv_state_dtype = self.get_state_dtype()[0]
395 decode_conv1d_weight = None
396 if is_fused_kda_decode_supported(
397 self.local_num_heads,
@@ -567,7 +587,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
587 g1 = g1[:, :num_actual_tokens]
588 beta = beta[:, :num_actual_tokens]
589
570 - conv_state, recurrent_state = self.kv_cache
590 + conv_state, recurrent_state, *recoverssm_records = self.kv_cache
591 # The convolution kernels consume (..., dim, width - 1).
592 if not is_conv_state_dim_first():
593 conv_state = conv_state.transpose(-1, -2)
@@ -634,7 +654,11 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
654 assert spec_state_indices_tensor is not None
655 assert spec_query_start_loc is not None
656 spec_conv_indices = spec_state_indices_tensor[:, 0][: m.num_spec_decodes]
637 - spec_max_query_len = spec_state_indices_tensor.size(-1)
657 + spec_max_query_len = (
658 + self.spec_query_len
659 + if self.use_recoverssm
660 + else spec_state_indices_tensor.size(-1)
661 + )
662 spec_conv_out = torch.empty_like(mixed_qkv_spec)
663 mixed_qkv_spec = causal_conv1d_update(
664 mixed_qkv_spec,
@@ -659,21 +683,48 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention):
683 if m.num_prefills == 0 and m.num_decodes == 0
684 else None
685 )
662 - core_attn_out_spec, _ = fused_recurrent_kda(
663 - q=q_spec,
664 - k=k_spec,
665 - v=v_spec,
666 - raw_g=g1_spec,
667 - raw_beta=beta_spec,
668 - A_log=self.A_log,
669 - dt_bias=self.dt_bias,
670 - lower_bound=self.gate_lower_bound,
671 - initial_state=recurrent_state,
672 - cu_seqlens=spec_cu_seqlens,
673 - ssm_state_indices=spec_state_indices_tensor,
674 - num_accepted_tokens=num_accepted_tokens,
675 - out=spec_out,
676 - )
686 + if self.use_recoverssm:
687 + from vllm.models.kimi_k3.nvidia.ops.recoverssm import (
688 + kda_recoverssm_verify,
689 + )
690 +
691 + if len(recoverssm_records) != 2:
692 + raise ValueError(
693 + "KDA RecoverSSM requires correction and key/gate buffers"
694 + )
695 + core_attn_out_spec = kda_recoverssm_verify(
696 + q=q_spec,
697 + k=k_spec,
698 + v=v_spec,
699 + raw_g=g1_spec,
700 + raw_beta=beta_spec,
701 + A_log=self.A_log,
702 + dt_bias=self.dt_bias,
703 + lower_bound=self.gate_lower_bound,
704 + checkpoint_state=recurrent_state,
705 + correction_cache=recoverssm_records[0],
706 + kg_cache=recoverssm_records[1],
707 + query_start_loc=spec_cu_seqlens,
708 + state_indices=spec_state_indices_tensor[: m.num_spec_decodes, 0],
709 + spec_query_len=self.spec_query_len,
710 + out=spec_out,
711 + )
712 + else:
713 + core_attn_out_spec, _ = fused_recurrent_kda(
714 + q=q_spec,
715 + k=k_spec,
716 + v=v_spec,
717 + raw_g=g1_spec,
718 + raw_beta=beta_spec,
719 + A_log=self.A_log,
720 + dt_bias=self.dt_bias,
721 + lower_bound=self.gate_lower_bound,
722 + initial_state=recurrent_state,
723 + cu_seqlens=spec_cu_seqlens,
724 + ssm_state_indices=spec_state_indices_tensor,
725 + num_accepted_tokens=num_accepted_tokens,
726 + out=spec_out,
727 + )
728
729 # Prefill or plain-decode path.
730 core_attn_out_non_spec = None
vllm/models/kimi_k3/nvidia/kda_metadata.py
+176 -12
@@ -6,13 +6,18 @@ 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
11 -from dataclasses import dataclass
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
@@ -22,6 +27,10 @@ from vllm.v1.attention.backends.gdn_attn import (
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,
@@ -29,6 +38,11 @@ from vllm.v1.attention.backends.utils import (
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:
@@ -229,11 +243,103 @@ def stage_spec_decode_metadata(
243
244
245 @dataclass
232 -class KimiK3KDAMetadata(GDNAttentionMetadata):
233 - pass
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 + recoverssm_commit: KDARecoverSSMCommitMetadata | None = None
263 + recoverssm_context: "KDARecoverSSMCommitContext | None" = field(
264 + default=None, repr=False, compare=False
265 + )
266 +
267 + def commit_recoverssm_state(
268 + self, num_accepted_tokens: torch.Tensor
269 + ) -> RecoverSSMPostprocessMetadata | None:
270 + commit = self.recoverssm_commit
271 + if commit is None:
272 + return None
273 + context = self.recoverssm_context
274 + assert context is not None
275 + align = commit.align
276 + context.commit(
277 + num_accepted_tokens,
278 + commit.state_indices[: self.num_spec_decodes, 0],
279 + commit.query_start_loc[: self.num_spec_decodes + 1],
280 + request_indices=commit.request_indices,
281 + block_table=align.block_table if align is not None else None,
282 + num_computed_tokens=(
283 + align.num_computed_tokens if align is not None else None
284 + ),
285 + mamba_block_size=align.block_size if align is not None else None,
286 + )
287 + if align is None:
288 + return None
289 + return RecoverSSMPostprocessMetadata(
290 + num_spec_decodes=self.num_spec_decodes,
291 + request_indices=commit.request_indices,
292 + block_table=align.block_table,
293 + num_computed_tokens=align.num_computed_tokens,
294 + block_size=align.block_size,
295 + )
296
297
298 class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
299 + def __init__(
300 + self,
301 + kv_cache_spec: MambaSpec,
302 + layer_names: list[str],
303 + vllm_config: VllmConfig,
304 + device: torch.device,
305 + ) -> None:
306 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)
307 + self.use_recoverssm = vllm_config.cache_config.use_kda_recoverssm
308 + self.spec_state_slots = 1 if self.use_recoverssm else self.num_spec + 1
309 + self.recoverssm_num_accepted_tokens: torch.Tensor | None = None
310 + self.recoverssm_context: KDARecoverSSMCommitContext | None = None
311 + if self.use_recoverssm:
312 + max_num_reqs = vllm_config.scheduler_config.max_num_seqs
313 + self.spec_state_indices_tensor = torch.empty(
314 + (max_num_reqs, 1),
315 + dtype=torch.int32,
316 + device=device,
317 + )
318 + self.recoverssm_num_accepted_tokens = torch.ones(
319 + max_num_reqs,
320 + dtype=torch.int32,
321 + device=device,
322 + )
323 +
324 + def _get_recoverssm_context(self) -> "KDARecoverSSMCommitContext":
325 + context = self.recoverssm_context
326 + if context is not None:
327 + return context
328 +
329 + from vllm.models.kimi_k3.nvidia.ops.recoverssm import (
330 + KDARecoverSSMCommitContext,
331 + )
332 +
333 + forward_context = self.vllm_config.compilation_config.static_forward_context
334 + layers = [forward_context[layer_name] for layer_name in self.layer_names]
335 + context = KDARecoverSSMCommitContext.create(
336 + layers,
337 + spec_query_len=1 + self.vllm_config.num_speculative_tokens,
338 + max_num_reqs=self.vllm_config.scheduler_config.max_num_seqs,
339 + )
340 + self.recoverssm_context = context
341 + return context
342 +
343 def build( # type: ignore[override]
344 self,
345 common_prefix_len: int,
@@ -263,14 +369,26 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
369 num_spec_decodes = 0
370 else:
371 spec_sequence_masks_cpu = num_decode_draft_tokens_cpu >= 0
266 - # A nonnegative entry identifies a spec request. If no draft token
267 - # was scheduled, process the whole batch as non-spec instead.
268 - if num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item() == 0:
372 + if self.use_recoverssm:
373 + assert m.is_prefilling is not None
374 + assert m.is_prefilling.device.type == "cpu"
375 + active_decode_mask_cpu = (~m.is_prefilling) & (
376 + query_start_loc_cpu.diff() > 0
377 + )
378 + spec_sequence_masks_cpu |= active_decode_mask_cpu
379 + # Native KDA can use its regular decode path when no draft token
380 + # was scheduled. RecoverSSM must preserve its extended conv window.
381 + if (
382 + not self.use_recoverssm
383 + and num_decode_draft_tokens_cpu[spec_sequence_masks_cpu].sum().item()
384 + == 0
385 + ):
386 spec_sequence_masks_cpu = None
387 num_spec_decodes = 0
388 else:
389 num_spec_decodes = spec_sequence_masks_cpu.sum().item()
390
391 + spec_request_indices = None
392 if num_spec_decodes == 0:
393 # The runner orders ordinary decodes before prefills.
394 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
@@ -289,6 +407,16 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
407 assert spec_sequence_masks_cpu is not None
408 assert num_accepted_tokens is not None
409 query_lens_cpu = query_start_loc_cpu.diff()
410 + if (
411 + self.use_recoverssm
412 + and torch.any(
413 + query_lens_cpu[spec_sequence_masks_cpu] > self.num_spec + 1
414 + ).item()
415 + ):
416 + raise ValueError(
417 + "KDA RecoverSSM speculative decode query length exceeds "
418 + f"its activation capacity ({self.num_spec + 1})"
419 + )
420 num_query_tokens = query_start_loc_cpu[-1].item()
421
422 # Exclude zero-length cudagraph padding from request-indexed
@@ -325,7 +453,7 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
453 non_spec_token_indx = None
454 # Real requests precede trailing cudagraph padding.
455 spec_state_indices_tensor = block_table_tensor[
328 - :num_spec_decodes, : self.num_spec + 1
456 + :num_spec_decodes, : self.spec_state_slots
457 ]
458 non_spec_state_indices_tensor = None
459 # Padding trails real requests, so this prefix already contains
@@ -339,6 +467,12 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
467 spec_sequence_masks_gpu = async_tensor_h2d(
468 spec_sequence_masks_cpu, device=query_start_loc.device
469 )
470 + if self.use_recoverssm:
471 + spec_request_indices = async_tensor_h2d(
472 + spec_sequence_masks_cpu.nonzero(as_tuple=True)[0],
473 + dtype=torch.int32,
474 + device=query_start_loc.device,
475 + )
476 spec_token_masks = torch.repeat_interleave(
477 spec_sequence_masks_gpu,
478 query_lens,
@@ -351,10 +485,10 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
485 non_spec_token_indx = index[:num_non_spec_tokens]
486 spec_token_indx = index[num_non_spec_tokens:]
487
354 - # Spec requests carry one state slot per speculative step;
355 - # non-spec requests use only their current state slot.
488 + # Native spec uses one state slot per step. RecoverSSM keeps
489 + # only the current checkpoint slot.
490 spec_state_indices_tensor = block_table_tensor[
357 - spec_sequence_masks_cpu, : self.num_spec + 1
491 + spec_sequence_masks_cpu, : self.spec_state_slots
492 ]
493 non_spec_state_indices_tensor = block_table_tensor[
494 active_non_spec_mask_cpu, 0
@@ -400,12 +534,18 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
534
535 num_accepted_tokens = num_accepted_tokens[spec_sequence_masks_cpu]
536
537 + if self.use_recoverssm:
538 + assert self.recoverssm_num_accepted_tokens is not None
539 + num_accepted_tokens = self.recoverssm_num_accepted_tokens[
540 + :num_spec_decodes
541 + ]
542 +
543 # Unlike the shared GDN layer, Kimi-K3's prefill KDA wrapper prepares
544 # its own chunk indices. Only causal-convolution metadata is needed here.
545 nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
546 if num_prefills > 0:
547 has_initial_state = m.compute_num_computed_tokens() > 0
408 - if spec_sequence_masks_cpu is not None:
548 + if num_spec_decodes > 0:
549 has_initial_state = has_initial_state[active_non_spec_mask_cpu]
550 assert non_spec_query_start_loc_cpu is not None
551 nums_dict, batch_ptr, token_chunk_offset_ptr = (
@@ -426,7 +566,7 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
566 and num_spec_decodes > 0
567 and num_prefills == 0
568 and num_decodes == 0
429 - and num_spec_decodes <= self.decode_cudagraph_max_bs
569 + and batch_size <= self.spec_state_indices_tensor.shape[0]
570 and num_spec_decode_tokens <= self.decode_cudagraph_max_bs
571 ):
572 # Equivalent PyTorch staging:
@@ -463,6 +603,24 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
603 :batch_size
604 ]
605
606 + recoverssm_commit = None
607 + if self.use_recoverssm and num_spec_decodes > 0:
608 + assert spec_state_indices_tensor is not None
609 + assert spec_query_start_loc is not None
610 + align = None
611 + if self.kv_cache_spec.mamba_cache_mode == "align":
612 + align = KDARecoverSSMAlignMetadata(
613 + block_table=m.block_table_tensor,
614 + num_computed_tokens=m.compute_num_computed_tokens(),
615 + block_size=self.kv_cache_spec.block_size,
616 + )
617 + recoverssm_commit = KDARecoverSSMCommitMetadata(
618 + state_indices=spec_state_indices_tensor,
619 + query_start_loc=spec_query_start_loc,
620 + request_indices=spec_request_indices,
621 + align=align,
622 + )
623 +
624 return KimiK3KDAMetadata(
625 num_prefills=num_prefills,
626 num_prefill_tokens=num_prefill_tokens,
@@ -480,6 +638,12 @@ class KimiK3KDAMetadataBuilder(GDNAttentionMetadataBuilder):
638 spec_token_indx=spec_token_indx,
639 non_spec_token_indx=non_spec_token_indx,
640 num_accepted_tokens=num_accepted_tokens,
641 + recoverssm_commit=recoverssm_commit,
642 + recoverssm_context=(
643 + self._get_recoverssm_context()
644 + if recoverssm_commit is not None
645 + else None
646 + ),
647 nums_dict=nums_dict,
648 batch_ptr=batch_ptr,
649 token_chunk_offset_ptr=token_chunk_offset_ptr,
vllm/models/kimi_k3/nvidia/model.py
+28 -5
@@ -62,12 +62,14 @@ from vllm.model_executor.models.interfaces import (
62 EagleModelMixin,
63 HasInnerState,
64 IsHybrid,
65 + MambaStateShapes,
66 MixtureOfExperts,
67 SupportsEagle3,
68 SupportsEncoderCudaGraph,
69 SupportsMultiModal,
70 SupportsPP,
71 SupportsQuant,
72 + SupportsReplaySSM,
73 )
74 from vllm.model_executor.models.kimi_k25 import KimiK25MediaPixelInputs
75 from vllm.model_executor.models.kimi_k25_vit import (
@@ -1559,7 +1561,13 @@ class KimiLinearModel(nn.Module, EagleModelMixin, SupportsQuant):
1561
1562
1563 class KimiLinearForCausalLM(
1562 - nn.Module, HasInnerState, SupportsPP, MixtureOfExperts, IsHybrid, SupportsEagle3
1564 + nn.Module,
1565 + HasInnerState,
1566 + SupportsPP,
1567 + MixtureOfExperts,
1568 + IsHybrid,
1569 + SupportsEagle3,
1570 + SupportsReplaySSM,
1571 ):
1572 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
1573 super().__init__()
@@ -1613,15 +1621,20 @@ class KimiLinearForCausalLM(
1621 def get_mamba_state_dtype_from_config(
1622 cls,
1623 vllm_config: "VllmConfig",
1616 - ) -> tuple[torch.dtype, torch.dtype]:
1617 - return MambaStateDtypeCalculator.kda_state_dtype(
1624 + ) -> tuple[torch.dtype, ...]:
1625 + dtypes = MambaStateDtypeCalculator.kda_state_dtype(
1626 vllm_config.model_config.dtype, vllm_config.cache_config.mamba_cache_dtype
1627 )
1628 + if vllm_config.cache_config.use_kda_recoverssm:
1629 + dtypes = MambaStateDtypeCalculator.append_kda_recoverssm_record(
1630 + dtypes, vllm_config.model_config.dtype
1631 + )
1632 + return dtypes
1633
1634 @classmethod
1635 def get_mamba_state_shape_from_config(
1636 cls, vllm_config: "VllmConfig"
1624 - ) -> tuple[tuple[int, int], tuple[int, int, int]]:
1637 + ) -> MambaStateShapes:
1638 parallel_config = vllm_config.parallel_config
1639 hf_config = vllm_config.model_config.hf_config
1640 tp_size = parallel_config.tensor_parallel_size
@@ -1630,13 +1643,22 @@ class KimiLinearForCausalLM(
1643 if vllm_config.speculative_config
1644 else 0
1645 )
1633 - return MambaStateShapeCalculator.kda_state_shape(
1646 + shapes = MambaStateShapeCalculator.kda_state_shape(
1647 tp_size,
1648 hf_config.linear_attn_config["num_heads"],
1649 hf_config.linear_attn_config["head_dim"],
1650 conv_kernel_size=hf_config.linear_attn_config["short_conv_kernel_size"],
1651 num_spec=num_spec,
1652 )
1653 + if vllm_config.cache_config.use_kda_recoverssm:
1654 + return MambaStateShapeCalculator.append_kda_recoverssm_record(
1655 + shapes,
1656 + hf_config.linear_attn_config["num_heads"],
1657 + hf_config.linear_attn_config["head_dim"],
1658 + tp_world_size=tp_size,
1659 + spec_query_len=1 + num_spec,
1660 + )
1661 + return shapes
1662
1663 @classmethod
1664 def get_mamba_state_copy_func(
@@ -1697,6 +1719,7 @@ class KimiK3ForConditionalGeneration(
1719 SupportsEagle3,
1720 HasInnerState,
1721 IsHybrid,
1722 + SupportsReplaySSM,
1723 ):
1724 """Kimi-K3 model with Kimi-K2.5 vision and KimiLinear text."""
1725
vllm/models/kimi_k3/nvidia/ops/recoverssm.py new
+1067
@@ -0,0 +1,1067 @@
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"]
vllm/v1/attention/backends/recoverssm_metadata.py new
+26
@@ -0,0 +1,26 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import abc
5 +from dataclasses import dataclass
6 +
7 +import torch
8 +
9 +
10 +@dataclass(frozen=True)
11 +class RecoverSSMPostprocessMetadata:
12 + """Metadata used during postprocessing for align-mode prefix caching."""
13 +
14 + num_spec_decodes: int
15 + request_indices: torch.Tensor | None
16 + block_table: torch.Tensor
17 + num_computed_tokens: torch.Tensor
18 + block_size: int
19 +
20 +
21 +class RecoverSSMMetadata(abc.ABC):
22 + @abc.abstractmethod
23 + def commit_recoverssm_state(
24 + self, num_accepted_tokens: torch.Tensor
25 + ) -> RecoverSSMPostprocessMetadata | None:
26 + raise NotImplementedError
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+37 -13
@@ -20,6 +20,7 @@ from vllm.v1.worker.gpu.input_batch import InputBatch
20 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
21 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
22 from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
23 +from vllm.v1.worker.gpu.model_states.recoverssm import RecoverSSMState
24 from vllm.v1.worker.mamba_utils import (
25 MambaSpecDecodeGPUContext,
26 preprocess_mamba_align_fused_kernel,
@@ -80,6 +81,9 @@ class MambaHybridModelState(DefaultModelState):
81 # kernel reusing the postprocess copy machinery, so the per-step src
82 # columns and the running state_idx are kept GPU-resident.
83 self._align_mode = self.cache_config.mamba_cache_mode == "align"
84 + self.recoverssm = (
85 + RecoverSSMState() if self.cache_config.use_kda_recoverssm else None
86 + )
87 if self._align_mode:
88 self._mamba_state_idx_gpu = torch.zeros(
89 self.max_num_reqs, dtype=torch.int32, device=self.device
@@ -267,7 +271,7 @@ class MambaHybridModelState(DefaultModelState):
271 num_accepted_tokens=num_accepted_tokens,
272 num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
273 )
270 - return build_attn_metadata(
274 + attn_metadata = build_attn_metadata(
275 attn_groups=attn_groups,
276 num_reqs=num_reqs,
277 num_tokens=num_tokens,
@@ -285,6 +289,13 @@ class MambaHybridModelState(DefaultModelState):
289 for_cudagraph_capture=for_capture,
290 rswa_prefix_lens=input_batch.prompt_lens,
291 )
292 + if self.recoverssm is not None:
293 + self.recoverssm.record_step(
294 + attn_metadata,
295 + attn_groups,
296 + for_capture=for_capture,
297 + )
298 + return attn_metadata
299
300 def postprocess_state(
301 self,
@@ -295,21 +306,34 @@ class MambaHybridModelState(DefaultModelState):
306 # Chunked prefill does not sample a token, so num_sampled can be 0.
307 # Mamba treats num_accepted_tokens=1 as the neutral non-spec value.
308 num_reqs = idx_mapping.shape[0]
298 - if not num_reqs:
299 - return
309 + if num_reqs:
310 + if not isinstance(num_sampled, int):
311 + # idx_mapping may contain -1 sentinels (filtered rows) under PP; the
312 + # kernel skips them rather than scattering with a host-side gather.
313 + _scatter_num_accepted_kernel[(num_reqs,)](
314 + idx_mapping,
315 + num_sampled,
316 + self.num_accepted_tokens_gpu,
317 + )
318 + else:
319 + # Fill with single value.
320 + _fill_num_accepted_kernel[(num_reqs,)](
321 + idx_mapping,
322 + self.num_accepted_tokens_gpu,
323 + max(num_sampled, 1),
324 + )
325
301 - if not isinstance(num_sampled, int):
302 - # idx_mapping may contain -1 sentinels (filtered rows) under PP; the
303 - # kernel skips them rather than scattering with a host-side gather.
304 - _scatter_num_accepted_kernel[(num_reqs,)](
305 - idx_mapping, num_sampled, self.num_accepted_tokens_gpu
306 - )
307 - else:
308 - # Fill with single value.
309 - _fill_num_accepted_kernel[(num_reqs,)](
310 - idx_mapping, self.num_accepted_tokens_gpu, max(num_sampled, 1)
326 + if self.recoverssm is not None:
327 + self.recoverssm.commit_step(
328 + num_sampled,
329 + idx_mapping,
330 + state_indices=(self._mamba_state_idx_gpu if self._align_mode else None),
331 + num_accepted_tokens=self.num_accepted_tokens_gpu,
332 )
333
334 + if not num_reqs:
335 + return
336 +
337 # Align: save the running state to the block-aligned position when
338 # spec-decode acceptance leaves the sequence non-block-aligned (mirrors
339 # the V1 align postprocess). num_computed_tokens already holds the
vllm/v1/worker/gpu/model_states/recoverssm.py new
+101
@@ -0,0 +1,101 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Any
5 +
6 +import torch
7 +
8 +from vllm.triton_utils import tl, triton
9 +from vllm.v1.attention.backends.recoverssm_metadata import RecoverSSMMetadata
10 +from vllm.v1.worker.utils import AttentionGroup
11 +
12 +
13 +class RecoverSSMState:
14 + """Coordinates RecoverSSM metadata between attention and postprocessing."""
15 +
16 + def __init__(self) -> None:
17 + self._step: tuple[RecoverSSMMetadata, ...] | None = None
18 +
19 + def record_step(
20 + self,
21 + attn_metadata: dict[str, Any],
22 + attn_groups: list[list[AttentionGroup]],
23 + *,
24 + for_capture: bool,
25 + ) -> None:
26 + if for_capture:
27 + self._step = None
28 + return
29 +
30 + step: list[RecoverSSMMetadata] = []
31 + for group_list in attn_groups:
32 + for group in group_list:
33 + metadata = attn_metadata[group.layer_names[0]]
34 + if isinstance(metadata, RecoverSSMMetadata):
35 + step.append(metadata)
36 + self._step = tuple(step)
37 +
38 + def commit_step(
39 + self,
40 + num_sampled: torch.Tensor | int,
41 + idx_mapping: torch.Tensor,
42 + *,
43 + state_indices: torch.Tensor | None,
44 + num_accepted_tokens: torch.Tensor,
45 + ) -> None:
46 + step = self._step
47 + self._step = None
48 + if isinstance(num_sampled, int) or step is None:
49 + return
50 +
51 + for metadata in step:
52 + postprocess_meta = metadata.commit_recoverssm_state(num_sampled)
53 + if postprocess_meta is None:
54 + continue
55 + assert state_indices is not None
56 + # RecoverSSM already restored the accepted state. Update its running
57 + # column and reset the next-step copy bias to the neutral value.
58 + _postprocess_recoverssm_align_kernel[(postprocess_meta.num_spec_decodes,)](
59 + idx_mapping,
60 + num_sampled,
61 + postprocess_meta.request_indices,
62 + postprocess_meta.num_computed_tokens,
63 + state_indices,
64 + num_accepted_tokens,
65 + MAMBA_BLOCK_SIZE=postprocess_meta.block_size,
66 + BLOCK_TABLE_WIDTH=postprocess_meta.block_table.shape[1],
67 + )
68 +
69 +
70 +@triton.heuristics(
71 + {"HAS_REQUEST_INDICES": lambda args: args["request_indices_ptr"] is not None}
72 +)
73 +@triton.jit
74 +def _postprocess_recoverssm_align_kernel(
75 + idx_mapping_ptr,
76 + num_sampled_ptr,
77 + request_indices_ptr,
78 + num_computed_ptr,
79 + state_idx_ptr,
80 + num_accepted_ptr,
81 + HAS_REQUEST_INDICES: tl.constexpr,
82 + MAMBA_BLOCK_SIZE: tl.constexpr,
83 + BLOCK_TABLE_WIDTH: tl.constexpr,
84 +):
85 + spec_idx = tl.program_id(0)
86 + batch_idx = spec_idx
87 + if HAS_REQUEST_INDICES:
88 + batch_idx = tl.load(request_indices_ptr + spec_idx)
89 + req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
90 + if req_state_idx < 0:
91 + return
92 + num_sampled = tl.load(num_sampled_ptr + batch_idx)
93 + num_computed = tl.load(num_computed_ptr + batch_idx)
94 + tl.store(
95 + state_idx_ptr + req_state_idx,
96 + tl.minimum(
97 + (num_computed + num_sampled) // MAMBA_BLOCK_SIZE,
98 + BLOCK_TABLE_WIDTH - 1,
99 + ),
100 + )
101 + tl.store(num_accepted_ptr + req_state_idx, 1)
vllm/v1/worker/mamba_utils.py
+7 -2
@@ -828,7 +828,13 @@ class MambaSpecDecodeGPUContext:
828 attention = forward_context[layer_name]
829 kv_caches: list[torch.Tensor] = attention.kv_cache
830
831 - for state_type_idx, state in enumerate(kv_caches):
831 + if len(kv_caches) < self.num_state_types:
832 + raise ValueError(
833 + f"Expected at least {self.num_state_types} Mamba state "
834 + f"tensors, got {len(kv_caches)}"
835 + )
836 + for state_type_idx, copy_func in enumerate(mamba_state_copy_funcs):
837 + state = kv_caches[state_type_idx]
838 # Base address
839 self.state_base_addrs[idx] = state.data_ptr()
840
@@ -845,7 +851,6 @@ class MambaSpecDecodeGPUContext:
851 # Element size
852 self.state_elem_sizes[idx] = state.element_size()
853
848 - copy_func = mamba_state_copy_funcs[state_type_idx]
854 assert (
855 copy_func is get_conv_copy_spec
856 or copy_func is get_temporal_copy_spec