@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 61ab70ec3b

[Model Runner V2] support mamba hybrid models align prefix cache (#42406)

Signed-off-by: zhuhaoran <zhuhaoran.zhr@alibaba-inc.com>

zhrrr committed Jun 30, 2026 at 05:09 UTC 61ab70ec3bd13dd422b86f3b80207d322994a5e7
9 files changed +1045 -132
tests/kernels/mamba/test_precopy_mamba_align.py new
+180
@@ -0,0 +1,180 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Equivalence test for ``precopy_mamba_align_fused_kernel``.
4 +
5 +The V2 "align" pre-copy must migrate mamba state across block boundaries with
6 +byte-identical semantics to the V1 copy specs (``get_conv_copy_spec`` /
7 +``get_temporal_copy_spec``):
8 +
9 +* conv state (SD layout, conv_width > 0): shift the sliding window by
10 + ``token_bias`` tokens -- ``state[bt[src_col], token_bias:]`` ->
11 + ``state[bt[dst_col], :conv_width - token_bias]``.
12 +* temporal state (conv_width == 0): ``token_bias`` selects the accepted
13 + speculative column -- ``state[bt[src_col + token_bias]]`` ->
14 + ``state[bt[dst_col]]``.
15 +
16 +The kernel must also no-op when ``src_col < 0`` (fresh request) or
17 +``src_col == dst_col`` (no boundary crossed).
18 +"""
19 +
20 +from __future__ import annotations
21 +
22 +import torch
23 +
24 +from vllm.platforms import current_platform
25 +from vllm.v1.worker.mamba_utils import precopy_mamba_align_fused_kernel
26 +
27 +try:
28 + import pytest
29 +
30 + pytestmark = pytest.mark.skipif(
31 + not current_platform.is_cuda(),
32 + reason="precopy_mamba_align_fused_kernel needs CUDA/Triton",
33 + )
34 + _parametrize = pytest.mark.parametrize
35 +except ModuleNotFoundError: # allow running directly as ``python <thisfile>``
36 + pytest = None
37 +
38 + def _parametrize(_name, _values):
39 + def _deco(fn):
40 + return fn
41 +
42 + return _deco
43 +
44 +
45 +NUM_LAYERS = 3
46 +CONV_WIDTH = 4 # conv_kernel - 1 + num_spec
47 +CONV_DIM = 96
48 +SSM_SHAPE = (4, 16, 16)
49 +MAX_COLS = 8
50 +
51 +
52 +def _build_state(num_blocks, device):
53 + """Per-layer (conv SD [nb, width, dim] bf16, ssm [nb, *shape] fp32) pools."""
54 + convs, ssms = [], []
55 + for _ in range(NUM_LAYERS):
56 + convs.append(
57 + torch.randn(
58 + num_blocks, CONV_WIDTH, CONV_DIM, dtype=torch.bfloat16, device=device
59 + )
60 + )
61 + ssms.append(
62 + torch.randn(num_blocks, *SSM_SHAPE, dtype=torch.float32, device=device)
63 + )
64 + return convs, ssms
65 +
66 +
67 +def _build_meta(convs, ssms, device):
68 + """Flattened per-(layer, state-type) metadata, ordered conv, ssm per layer."""
69 + n = NUM_LAYERS * 2
70 + base = torch.zeros(n, dtype=torch.int64, device=device)
71 + blk_stride = torch.zeros(n, dtype=torch.int64, device=device)
72 + elem = torch.zeros(n, dtype=torch.int32, device=device)
73 + inner = torch.zeros(n, dtype=torch.int64, device=device)
74 + width = torch.zeros(n, dtype=torch.int32, device=device)
75 + group = torch.zeros(n, dtype=torch.int32, device=device)
76 + drc = torch.zeros(n, dtype=torch.int32, device=device) # DS rows (unused, SD)
77 + drs = torch.zeros(n, dtype=torch.int64, device=device)
78 + i = 0
79 + for layer in range(NUM_LAYERS):
80 + conv, ssm = convs[layer], ssms[layer]
81 + # conv (SD): width = size(1), inner = stride(1)
82 + base[i] = conv.data_ptr()
83 + blk_stride[i] = conv.stride(0) * conv.element_size()
84 + elem[i] = conv.element_size()
85 + width[i] = conv.size(1)
86 + inner[i] = conv.stride(1)
87 + i += 1
88 + # ssm (temporal): width = 0, inner = elems per block
89 + base[i] = ssm.data_ptr()
90 + blk_stride[i] = ssm.stride(0) * ssm.element_size()
91 + elem[i] = ssm.element_size()
92 + width[i] = 0
93 + inner[i] = ssm[0].numel()
94 + i += 1
95 + return base, blk_stride, elem, inner, width, group, drc, drs
96 +
97 +
98 +def _reference(convs, ssms, bt, src_col, dst_col, bias, num_reqs):
99 + """Apply the V1 copy semantics on clones, reading from the pre-copy state."""
100 + conv_pre = [c.clone() for c in convs]
101 + ssm_pre = [s.clone() for s in ssms]
102 + conv_ref = [c.clone() for c in convs]
103 + ssm_ref = [s.clone() for s in ssms]
104 + for r in range(num_reqs):
105 + sc, dc, tb = int(src_col[r]), int(dst_col[r]), int(bias[r])
106 + if sc < 0 or sc == dc:
107 + continue
108 + sblk, dblk = int(bt[r, sc]), int(bt[r, dc])
109 + tblk = int(bt[r, sc + tb]) # temporal src column shifted by bias
110 + for layer in range(NUM_LAYERS):
111 + conv_ref[layer][dblk, : CONV_WIDTH - tb] = conv_pre[layer][sblk, tb:]
112 + ssm_ref[layer][dblk] = ssm_pre[layer][tblk]
113 + return conv_ref, ssm_ref
114 +
115 +
116 +@_parametrize("num_reqs", [1, 4, 16])
117 +@_parametrize("token_bias", [0, 1, 2])
118 +def test_precopy_matches_v1_copy_specs(num_reqs, token_bias):
119 + device = torch.device("cuda")
120 + torch.manual_seed(0)
121 + # Distinct physical block per (req, col) so copies never alias.
122 + num_blocks = num_reqs * MAX_COLS + 1
123 + bt = torch.empty(num_reqs, MAX_COLS, dtype=torch.int32, device=device)
124 + for r in range(num_reqs):
125 + bt[r] = torch.arange(
126 + 1 + r * MAX_COLS, 1 + (r + 1) * MAX_COLS, dtype=torch.int32, device=device
127 + )
128 +
129 + # Per-req columns: req 0 fresh (src=-1, skip), req 1 same block (skip),
130 + # the rest cross from col 1 -> col 0 with the given spec token bias.
131 + src_col = torch.full((num_reqs,), 1, dtype=torch.int32, device=device)
132 + dst_col = torch.zeros(num_reqs, dtype=torch.int32, device=device)
133 + bias = torch.full((num_reqs,), token_bias, dtype=torch.int32, device=device)
134 + if num_reqs >= 1:
135 + src_col[0] = -1 # fresh -> no copy
136 + if num_reqs >= 2:
137 + dst_col[1] = 1 # src_col == dst_col -> no copy
138 +
139 + convs, ssms = _build_state(num_blocks, device)
140 + conv_ref, ssm_ref = _reference(
141 + convs, ssms, bt.cpu(), src_col.cpu(), dst_col.cpu(), bias.cpu(), num_reqs
142 + )
143 +
144 + base, blk_stride, elem, inner, width, group, drc, drs = _build_meta(
145 + convs, ssms, device
146 + )
147 + bt_ptrs = torch.tensor([bt.data_ptr()], dtype=torch.int64, device=device)
148 + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device)
149 + grid = (num_reqs, NUM_LAYERS * 2)
150 + precopy_mamba_align_fused_kernel[grid](
151 + dst_col,
152 + src_col,
153 + bias,
154 + bt_ptrs,
155 + bt.stride(0),
156 + base,
157 + blk_stride,
158 + elem,
159 + inner,
160 + width,
161 + group,
162 + drc,
163 + drs,
164 + idx_mapping,
165 + num_reqs,
166 + COPY_BLOCK_SIZE=1024,
167 + CONV_STATE_DIM_FIRST=False,
168 + )
169 + torch.accelerator.synchronize()
170 +
171 + for layer in range(NUM_LAYERS):
172 + torch.testing.assert_close(convs[layer], conv_ref[layer], rtol=0, atol=0)
173 + torch.testing.assert_close(ssms[layer], ssm_ref[layer], rtol=0, atol=0)
174 +
175 +
176 +if __name__ == "__main__":
177 + for nr in (1, 4, 16):
178 + for tb in (0, 1, 2):
179 + test_precopy_matches_v1_copy_specs(nr, tb)
180 + print(f"OK num_reqs={nr} token_bias={tb}")
tests/v1/e2e/general/test_mamba_prefix_cache.py
+280 -16
@@ -11,6 +11,7 @@ import datasets
11 import pytest
12 import torch
13
14 +import vllm.envs as envs
15 from tests.utils import create_new_process_for_each_test
16 from vllm import LLM, SamplingParams, TokensPrompt
17 from vllm.config import CacheConfig
@@ -494,12 +495,7 @@ def apply_patch(monkeypatch: pytest.MonkeyPatch):
495 monkeypatch.setattr(mamba_utils, "do_mamba_copy_block", fake_copy_fn)
496
497
497 -@create_new_process_for_each_test()
498 -def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
499 - run_ref_mamba_state_in_subprocess()
500 - apply_patch(monkeypatch)
501 - prompt_dataset = datasets.load_dataset("heheda/a_long_article")
502 - full_prompt = prompt_dataset["train"][0]["text"]
498 +def get_mamba_prefix_cache_step_configs() -> dict[str, TestConfig]:
499 tests = {
500 "accept_1": TestConfig(
501 num_prompt_tokens=554,
@@ -731,6 +727,27 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
727 ),
728 }
729
730 + return tests
731 +
732 +
733 +def fill_following_kv_cache_block_ids(test_config: TestConfig) -> None:
734 + for step_action_prev, step_action_next in zip(
735 + test_config.step_actions[:-1], test_config.step_actions[1:]
736 + ):
737 + if len(step_action_next.kv_cache_block_ids) == 0:
738 + step_action_next.kv_cache_block_ids = (
739 + step_action_prev.kv_cache_block_ids.copy()
740 + )
741 +
742 +
743 +@create_new_process_for_each_test()
744 +def test_mamba_prefix_cache_mrv1(monkeypatch: pytest.MonkeyPatch):
745 + run_ref_mamba_state_in_subprocess()
746 + apply_patch(monkeypatch)
747 + prompt_dataset = datasets.load_dataset("heheda/a_long_article")
748 + full_prompt = prompt_dataset["train"][0]["text"]
749 + tests = get_mamba_prefix_cache_step_configs()
750 +
751 engine = LLM(
752 model=MODEL,
753 enable_prefix_caching=True,
@@ -758,16 +775,7 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
775 )
776 global cur_step_action_idx
777 cur_step_action_idx = 0
761 - for step_action_prev, step_action_next in zip(
762 - test_config.step_actions[:-1], test_config.step_actions[1:]
763 - ):
764 - if (
765 - step_action_next.kv_cache_block_ids is not None
766 - and len(step_action_next.kv_cache_block_ids) == 0
767 - ):
768 - prev_block_ids = step_action_prev.kv_cache_block_ids
769 - if prev_block_ids is not None:
770 - step_action_next.kv_cache_block_ids = prev_block_ids.copy()
778 + fill_following_kv_cache_block_ids(test_config)
779 global step_actions
780 step_actions = test_config.step_actions
781 _ = engine.generate(
@@ -787,3 +795,259 @@ def test_mamba_prefix_cache(monkeypatch: pytest.MonkeyPatch):
795 del engine
796 torch.accelerator.empty_cache()
797 cleanup_dist_env_and_memory()
798 +
799 +
800 +@create_new_process_for_each_test()
801 +def test_mamba_prefix_cache_mrv2(monkeypatch: pytest.MonkeyPatch):
802 + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
803 + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1")
804 + envs.disable_envs_cache()
805 +
806 + from vllm.v1.worker.gpu.model_runner import GPUModelRunner as MRV2GPUModelRunner
807 + from vllm.v1.worker.gpu.model_states.mamba_hybrid import (
808 + MambaHybridModelState,
809 + )
810 + from vllm.v1.worker.gpu.sample.output import SamplerOutput as MRV2SamplerOutput
811 +
812 + events: list[int] = []
813 + original_execute_model = MRV2GPUModelRunner.execute_model
814 + original_sample = MRV2GPUModelRunner.sample
815 + original_preprocess_state = MambaHybridModelState.preprocess_state
816 + original_postprocess_state = MambaHybridModelState.postprocess_state
817 + original_step_action_fn = InprocClient.get_output
818 + original_allocate_slots = KVCacheManager.allocate_slots
819 + captured: dict[str, Any] = {}
820 +
821 + def temporal_states(model_state, block_tables, kv_cache_config):
822 + # Qwen3-Next keeps the temporal (ssm) state as the last Mamba cache.
823 + forward_context = (
824 + model_state.vllm_config.compilation_config.static_forward_context
825 + )
826 + group_ids, _ = get_mamba_groups(kv_cache_config)
827 + for group_id in group_ids:
828 + block_table = block_tables[group_id]
829 + for layer_name in kv_cache_config.kv_cache_groups[group_id].layer_names:
830 + yield forward_context[layer_name].kv_cache[-1], block_table
831 +
832 + def temporal_block(temporal_state, block_table, col):
833 + return temporal_state[int(block_table[0, col].item())]
834 +
835 + def wrapped_preprocess_state(
836 + self: MambaHybridModelState,
837 + input_batch: Any,
838 + block_tables: tuple[torch.Tensor, ...],
839 + kv_cache_config: KVCacheConfig,
840 + num_computed_tokens: torch.Tensor,
841 + ) -> None:
842 + captured["block_tables"] = block_tables
843 + captured["kv_cache_config"] = kv_cache_config
844 + expected = (
845 + None if cur_step_action is None else cur_step_action.preprocess_copy_idx
846 + )
847 + snapshots = []
848 + if expected is not None and expected != (-1, -1):
849 + for temporal, bt in temporal_states(self, block_tables, kv_cache_config):
850 + snapshots.append(
851 + (temporal, bt, temporal_block(temporal, bt, expected[0]).clone())
852 + )
853 + ret = original_preprocess_state(
854 + self, input_batch, block_tables, kv_cache_config, num_computed_tokens
855 + )
856 + if cur_step_action is not None:
857 + req_idx = int(input_batch.idx_mapping[0].item())
858 + src_col = int(self._mamba_src_col_gpu[req_idx].item())
859 + off = int(self._mamba_src_off_gpu[req_idx].item())
860 + dst = int(self._mamba_state_idx_gpu[req_idx].item())
861 + actual = (-1, -1) if src_col < 0 or src_col == dst else (src_col + off, dst)
862 + assert actual == expected, (
863 + f"V2 align preprocess copy: expected={expected}, "
864 + f"actual={actual}, {cur_step_action=}"
865 + )
866 + for temporal, bt, src_state in snapshots:
867 + torch.testing.assert_close(
868 + temporal_block(temporal, bt, expected[1]), src_state
869 + )
870 + return ret
871 +
872 + def wrapped_postprocess_state(
873 + self: MambaHybridModelState,
874 + idx_mapping: torch.Tensor,
875 + num_sampled: torch.Tensor | int,
876 + num_computed_tokens: torch.Tensor | None = None,
877 + ) -> None:
878 + action = cur_step_action
879 + block_tables = captured.get("block_tables")
880 + kv_cache_config = captured.get("kv_cache_config")
881 + # The postprocess kernel does not expose its indices, so only the copy
882 + # case is checked, by effect: snapshot the src block, expect dst == src.
883 + if (
884 + action is None
885 + or num_computed_tokens is None
886 + or block_tables is None
887 + or action.postprocess_copy_idx == (-1, -1)
888 + ):
889 + return original_postprocess_state(
890 + self, idx_mapping, num_sampled, num_computed_tokens
891 + )
892 + expected = action.postprocess_copy_idx
893 + snapshots = [
894 + (temporal, bt, temporal_block(temporal, bt, expected[0]).clone())
895 + for temporal, bt in temporal_states(self, block_tables, kv_cache_config)
896 + ]
897 + ret = original_postprocess_state(
898 + self, idx_mapping, num_sampled, num_computed_tokens
899 + )
900 + for temporal, bt, src_state in snapshots:
901 + torch.testing.assert_close(
902 + temporal_block(temporal, bt, expected[1]), src_state
903 + )
904 + return ret
905 +
906 + def wrapped_execute_model(
907 + self: MRV2GPUModelRunner,
908 + scheduler_output: SchedulerOutput,
909 + *args: Any,
910 + **kwargs: Any,
911 + ):
912 + events.extend(
913 + req.num_computed_tokens for req in scheduler_output.scheduled_new_reqs
914 + )
915 + events.extend(scheduler_output.scheduled_cached_reqs.num_computed_tokens)
916 + if cur_step_action is not None:
917 + num_scheduled_tokens = next(
918 + iter(scheduler_output.num_scheduled_tokens.values())
919 + )
920 + assert num_scheduled_tokens == cur_step_action.num_scheduled_tokens
921 + ret = original_execute_model(self, scheduler_output, *args, **kwargs)
922 + if cur_step_action is not None and self.execute_model_state is not None:
923 + input_batch = self.execute_model_state.input_batch
924 + assert (
925 + cur_step_action.num_computed_tokens_start
926 + == input_batch.positions[input_batch.query_start_loc[0]].item()
927 + )
928 + return ret
929 +
930 + def fake_sample(
931 + self: MRV2GPUModelRunner,
932 + hidden_states: torch.Tensor,
933 + input_batch: Any,
934 + grammar_output: Any,
935 + ):
936 + if cur_step_action is None:
937 + return original_sample(self, hidden_states, input_batch, grammar_output)
938 +
939 + num_reqs = input_batch.num_reqs
940 + sampled_token_ids = torch.ones(
941 + (num_reqs, self.num_speculative_steps + 1),
942 + device=hidden_states.device,
943 + dtype=torch.int64,
944 + )
945 + num_logits = torch.tensor(
946 + input_batch.cu_num_logits_np[1 : num_reqs + 1]
947 + - input_batch.cu_num_logits_np[:num_reqs],
948 + device=hidden_states.device,
949 + dtype=torch.int32,
950 + )
951 + accepted = torch.full_like(num_logits, num_accepted_tokens)
952 + num_sampled = torch.minimum(accepted, num_logits)
953 + prefill_lens = self.req_states.prefill_len.gpu[input_batch.idx_mapping]
954 + is_chunked_prefill = input_batch.seq_lens[:num_reqs] < prefill_lens
955 + num_sampled = torch.where(is_chunked_prefill, 0, num_sampled)
956 + num_rejected = torch.where(is_chunked_prefill, 0, num_logits - num_sampled)
957 + sampler_output = MRV2SamplerOutput(
958 + sampled_token_ids=sampled_token_ids,
959 + logprobs_tensors=None,
960 + num_nans=None,
961 + num_sampled=num_sampled,
962 + )
963 + return sampler_output, num_sampled, num_rejected
964 +
965 + monkeypatch.setattr(
966 + InprocClient,
967 + "get_output",
968 + get_fake_step_action_fn(original_step_action_fn),
969 + )
970 + monkeypatch.setattr(
971 + KVCacheManager,
972 + "allocate_slots",
973 + get_fake_allocate_slots_fn(original_allocate_slots),
974 + )
975 + monkeypatch.setattr(MRV2GPUModelRunner, "execute_model", wrapped_execute_model)
976 + monkeypatch.setattr(MRV2GPUModelRunner, "sample", fake_sample)
977 + monkeypatch.setattr(
978 + MambaHybridModelState, "preprocess_state", wrapped_preprocess_state
979 + )
980 + monkeypatch.setattr(
981 + MambaHybridModelState, "postprocess_state", wrapped_postprocess_state
982 + )
983 +
984 + engine = LLM(
985 + model=MODEL,
986 + load_format="dummy",
987 + enforce_eager=True,
988 + skip_tokenizer_init=True,
989 + enable_prefix_caching=True,
990 + block_size=BLOCK_SIZE,
991 + mamba_cache_mode="align",
992 + speculative_config={
993 + "method": "qwen3_next_mtp",
994 + "num_speculative_tokens": num_speculative_tokens,
995 + },
996 + max_num_batched_tokens=3072,
997 + max_model_len=BLOCK_SIZE * 12,
998 + hf_overrides={"num_hidden_layers": NUM_HIDDEN_LAYERS},
999 + seed=42,
1000 + )
1001 +
1002 + try:
1003 + tests = get_mamba_prefix_cache_step_configs()
1004 +
1005 + global step_actions
1006 + global cur_step_action_idx
1007 + global num_accepted_tokens
1008 + for test_name, test_config in tests.items():
1009 + num_accepted_tokens = test_config.num_accepted_tokens
1010 + cur_step_action_idx = 0
1011 + fill_following_kv_cache_block_ids(test_config)
1012 + step_actions = test_config.step_actions
1013 + sampling_params = SamplingParams(
1014 + temperature=0.0,
1015 + max_tokens=test_config.num_generated_tokens,
1016 + ignore_eos=True,
1017 + )
1018 + _ = engine.generate(
1019 + [TokensPrompt(prompt_token_ids=[1] * test_config.num_prompt_tokens)],
1020 + sampling_params=sampling_params,
1021 + )
1022 + assert cur_step_action_idx == len(test_config.step_actions), test_name
1023 + assert (
1024 + engine.llm_engine.engine_core.engine_core.scheduler.reset_prefix_cache()
1025 + )
1026 +
1027 + step_actions = []
1028 + cur_step_action_idx = 0
1029 + num_accepted_tokens = 1
1030 + prompt = TokensPrompt(prompt_token_ids=[1] * (BLOCK_SIZE * 2))
1031 + sampling_params = SamplingParams(
1032 + temperature=0.0,
1033 + max_tokens=1,
1034 + ignore_eos=True,
1035 + )
1036 + _ = engine.generate([prompt], sampling_params=sampling_params)
1037 + first_event_count = len(events)
1038 + _ = engine.generate([prompt], sampling_params=sampling_params)
1039 + second_events = events[first_event_count:]
1040 + prefix_hits = [
1041 + num_computed_tokens
1042 + for num_computed_tokens in second_events
1043 + if num_computed_tokens >= BLOCK_SIZE
1044 + ]
1045 + assert prefix_hits, (
1046 + "Expected the second identical prompt to hit prefix cache, "
1047 + f"got events={second_events!r}"
1048 + )
1049 + assert engine.llm_engine.engine_core.engine_core.scheduler.reset_prefix_cache()
1050 + finally:
1051 + del engine
1052 + torch.accelerator.empty_cache()
1053 + cleanup_dist_env_and_memory()
vllm/config/vllm.py
-11
@@ -1997,13 +1997,6 @@ class VllmConfig:
1997 model_config = self.model_config
1998 speculative_config = self.speculative_config
1999
2000 - if (
2001 - model_config is not None
2002 - and model_config.has_inner_state
2003 - and self.cache_config.mamba_cache_mode == "align"
2004 - ):
2005 - unsupported.append("hybrid/mamba models with align cache mode")
2006 -
2000 if self.parallel_config.prefill_context_parallel_size > 1:
2001 unsupported.append("prefill context parallelism")
2002
@@ -2152,10 +2145,6 @@ class VllmConfig:
2145 "to schedule a multiple of block_size tokens even if they are "
2146 "in the middle of a mm input"
2147 )
2155 - # TODO: support align mamba cache mode for model runner v2
2156 - assert not envs.VLLM_USE_V2_MODEL_RUNNER, (
2157 - "Model Runner V2 has not yet supported mamba_cache_mode='align'. "
2158 - )
2148
2149 @model_validator(mode="after")
2150 def validate_nvfp4_kv_cache_with_mla(self) -> "VllmConfig":
vllm/model_executor/models/diffusion_gemma.py
+3 -1
@@ -973,7 +973,9 @@ class DiffusionGemmaModelState(ModelState):
973 # so the captured graph and runtime point to identical addresses.
974 return {"inputs_embeds": self._inputs_embeds_buf[:num_tokens]}
975
976 - def postprocess_state(self, idx_mapping, num_sampled) -> None:
976 + def postprocess_state(
977 + self, idx_mapping, num_sampled, num_computed_tokens=None
978 + ) -> None:
979 return None
980
981 def prepare_attn(
vllm/v1/worker/gpu/model_runner.py
+13 -1
@@ -1111,7 +1111,9 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1111 self.req_states.total_len.gpu,
1112 )
1113
1114 - self.model_state.postprocess_state(idx_mapping, num_sampled)
1114 + self.model_state.postprocess_state(
1115 + idx_mapping, num_sampled, self.req_states.num_computed_tokens.gpu
1116 + )
1117
1118 @torch.inference_mode()
1119 def execute_model(
@@ -1176,6 +1178,16 @@ class GPUModelRunner(LoRAModelRunnerMixin):
1178 # Prepare all the inputs and copy to the input buffers.
1179 input_batch = self.prepare_inputs(scheduler_output, batch_desc)
1180 block_tables, slot_mappings = self.prepare_attn(input_batch)
1181 + # Mamba "align" pre-copy: migrate recurrent state across block
1182 + # boundaries before the forward. Runs only on real batches, and
1183 + # before model_state.prepare_attn gathers num_accepted_tokens so the
1184 + # boundary reset is visible to the attention metadata.
1185 + self.model_state.preprocess_state(
1186 + input_batch,
1187 + block_tables,
1188 + self.kv_cache_config,
1189 + self.req_states.num_computed_tokens.gpu,
1190 + )
1191
1192 if self.lora_config:
1193 # Activate LoRA adapters.
vllm/v1/worker/gpu/model_states/interface.py
+16 -1
@@ -95,8 +95,23 @@ class ModelState(ABC):
95 def apply_staged_writes(self) -> None:
96 return None
97
98 + def preprocess_state(
99 + self,
100 + input_batch: InputBatch,
101 + block_tables: tuple[torch.Tensor, ...],
102 + kv_cache_config: KVCacheConfig,
103 + num_computed_tokens: torch.Tensor,
104 + ) -> None:
105 + """Hook run on real batches before the forward pass (after block tables
106 + are gathered). Used by mamba "align" prefix caching to pre-copy state
107 + across block boundaries. No-op by default."""
108 + return None
109 +
110 def postprocess_state(
99 - self, idx_mapping: torch.Tensor, num_sampled: torch.Tensor
111 + self,
112 + idx_mapping: torch.Tensor,
113 + num_sampled: torch.Tensor,
114 + num_computed_tokens: torch.Tensor | None = None,
115 ) -> None:
116 return None
117
vllm/v1/worker/gpu/model_states/mamba_hybrid.py
+174 -8
@@ -9,15 +9,25 @@ import torch.nn as nn
9
10 from vllm.config import VllmConfig
11 from vllm.config.compilation import CUDAGraphMode
12 +from vllm.model_executor.layers.mamba.mamba_utils import (
13 + get_conv_copy_spec,
14 + is_conv_state_dim_first,
15 +)
16 from vllm.triton_utils import tl, triton
17 from vllm.v1.attention.backends.gdn_attn import GDNAttentionMetadataBuilder
18 from vllm.v1.attention.backends.mamba2_attn import Mamba2AttentionMetadataBuilder
15 -from vllm.v1.kv_cache_interface import KVCacheConfig
19 +from vllm.v1.core.sched.output import NewRequestData
20 +from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec
21 +from vllm.v1.utils import CpuGpuBuffer
22 from vllm.v1.worker.gpu.attn_utils import build_attn_metadata
23 from vllm.v1.worker.gpu.input_batch import InputBatch
24 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
25 from vllm.v1.worker.gpu.model_states.default import DefaultModelState
26 from vllm.v1.worker.gpu.model_states.interface import ModelSpecificAttnMetadata
27 +from vllm.v1.worker.mamba_utils import (
28 + MambaSpecDecodeGPUContext,
29 + preprocess_mamba_align_fused_kernel,
30 +)
31 from vllm.v1.worker.utils import AttentionGroup
32
33
@@ -65,9 +75,142 @@ class MambaHybridModelState(DefaultModelState):
75 device: torch.device,
76 ) -> None:
77 super().__init__(vllm_config, model, encoder_cache, device)
78 + self.cache_config = vllm_config.cache_config
79 self.num_accepted_tokens_gpu = torch.ones(
80 self.max_num_reqs, dtype=torch.int32, device=self.device
81 )
82 + # Pre-copy "align" prefix-cache state (V2). The migration of each
83 + # request's mamba state across block boundaries runs as a fused GPU
84 + # kernel reusing the postprocess copy machinery, so the per-step src
85 + # columns and the running state_idx are kept GPU-resident.
86 + self._align_mode = self.cache_config.mamba_cache_mode == "align"
87 + if self._align_mode:
88 + self._mamba_state_idx_gpu = torch.zeros(
89 + self.max_num_reqs, dtype=torch.int32, device=self.device
90 + )
91 + self._mamba_src_col_gpu = torch.full(
92 + (self.max_num_reqs,), -1, dtype=torch.int32, device=self.device
93 + )
94 + self._mamba_src_off_gpu = torch.zeros(
95 + self.max_num_reqs, dtype=torch.int32, device=self.device
96 + )
97 + self._mamba_ctx: MambaSpecDecodeGPUContext | None = None
98 + self._mamba_group_ids: list[int] = []
99 + self._mamba_spec: MambaSpec | None = None
100 +
101 + def add_request(self, req_index: int, new_req_data: NewRequestData) -> None:
102 + super().add_request(req_index, new_req_data)
103 + if self._align_mode:
104 + # Seed the running state block from the resumed/prefilled position.
105 + self._mamba_state_idx_gpu[req_index] = (
106 + new_req_data.num_computed_tokens - 1
107 + ) // self.cache_config.block_size
108 + self.num_accepted_tokens_gpu[req_index] = 1
109 +
110 + def _get_mamba_group_info(
111 + self, kv_cache_config: KVCacheConfig
112 + ) -> tuple[list[int], MambaSpec]:
113 + if self._mamba_spec is None:
114 + group_ids: list[int] = []
115 + specs: list[MambaSpec] = []
116 + for i, group in enumerate(kv_cache_config.kv_cache_groups):
117 + spec = group.kv_cache_spec
118 + if isinstance(spec, MambaSpec):
119 + group_ids.append(i)
120 + specs.append(spec)
121 + assert specs, "no mamba layers in the model"
122 + assert all(specs[0] == s for s in specs)
123 + self._mamba_group_ids = group_ids
124 + self._mamba_spec = specs[0]
125 + return self._mamba_group_ids, self._mamba_spec
126 +
127 + def _ensure_align_ctx(
128 + self,
129 + kv_cache_config: KVCacheConfig,
130 + mamba_group_ids: list[int],
131 + block_tables: tuple[torch.Tensor, ...],
132 + ) -> MambaSpecDecodeGPUContext:
133 + if self._mamba_ctx is None:
134 + copy_funcs = self.model.get_mamba_state_copy_func()
135 + # The fused copy kernels shift conv windows assuming the SD layout;
136 + # the DS layout cannot express a >0 spec-decode shift as a single
137 + # contiguous copy (mirrors get_conv_copy_spec's NotImplementedError).
138 + if get_conv_copy_spec in copy_funcs and is_conv_state_dim_first():
139 + assert self.vllm_config.speculative_config is None, (
140 + "DS conv state layout does not support mamba align state "
141 + "copies with speculative decoding"
142 + )
143 + self._mamba_ctx = MambaSpecDecodeGPUContext.create(
144 + max_num_reqs=self.max_num_reqs,
145 + kv_cache_config=kv_cache_config,
146 + num_state_types=len(copy_funcs),
147 + device=self.device,
148 + make_buffer=lambda n, dtype: CpuGpuBuffer(
149 + n, dtype=dtype, device=self.device
150 + ),
151 + )
152 + ctx = self._mamba_ctx
153 + if not ctx.is_initialized:
154 + forward_context = self.vllm_config.compilation_config.static_forward_context
155 + # block_tables are batch-order slices of the persistent
156 + # input_block_tables (stable data_ptr), so the metadata is captured
157 + # once here and reused across steps.
158 + ctx.initialize_from_forward_context(
159 + kv_cache_config,
160 + forward_context,
161 + self.model.get_mamba_state_copy_func(),
162 + [block_tables[gid] for gid in mamba_group_ids],
163 + )
164 + return ctx
165 +
166 + def preprocess_state(
167 + self,
168 + input_batch: InputBatch,
169 + block_tables: tuple[torch.Tensor, ...],
170 + kv_cache_config: KVCacheConfig,
171 + num_computed_tokens: torch.Tensor,
172 + ) -> None:
173 + """Migrate each request's mamba state across block boundaries before the
174 + forward (V1 align semantics, done on GPU). Runs on real batches only
175 + (dummy DP/profiling runs skip preprocess_state), and before
176 + ``prepare_attn`` gathers ``num_accepted_tokens``, so the boundary reset
177 + is visible to the forward kernels.
178 + """
179 + if not self._align_mode:
180 + return
181 + num_reqs = input_batch.num_reqs
182 + if num_reqs == 0:
183 + return
184 + mamba_group_ids, mamba_spec = self._get_mamba_group_info(kv_cache_config)
185 + ctx = self._ensure_align_ctx(kv_cache_config, mamba_group_ids, block_tables)
186 +
187 + # The state-advance + pre-copy kernels run every step; they fast-exit per
188 + # request when src_col < 0 or src_col == dst_col, so no copy happens on
189 + # steps that don't cross a block boundary. (Skipping the launch entirely
190 + # would need a V1-style async-D2H of the actual num_computed, since
191 + # num_computed_tokens_np is an optimistic mirror under async scheduling;
192 + # the launch cost is ~0.3% of TPOT, so the GPU fast-exit suffices.)
193 + block = 256
194 + grid = (triton.cdiv(num_reqs, block),)
195 + preprocess_mamba_align_fused_kernel[grid](
196 + input_batch.idx_mapping,
197 + self._mamba_state_idx_gpu,
198 + num_computed_tokens,
199 + input_batch.query_start_loc,
200 + self.num_accepted_tokens_gpu,
201 + self._mamba_src_col_gpu,
202 + self._mamba_src_off_gpu,
203 + num_reqs,
204 + BLOCK_SIZE=block,
205 + MAMBA_BLOCK_SIZE=mamba_spec.block_size,
206 + )
207 + ctx.run_fused_precopy(
208 + num_reqs,
209 + self._mamba_state_idx_gpu,
210 + self._mamba_src_col_gpu,
211 + self._mamba_src_off_gpu,
212 + input_batch.idx_mapping,
213 + )
214
215 def prepare_attn(
216 self,
@@ -145,22 +288,45 @@ class MambaHybridModelState(DefaultModelState):
288 )
289
290 def postprocess_state(
148 - self, idx_mapping: torch.Tensor, num_sampled: torch.Tensor | int
291 + self,
292 + idx_mapping: torch.Tensor,
293 + num_sampled: torch.Tensor | int,
294 + num_computed_tokens: torch.Tensor | None = None,
295 ) -> None:
296 # Chunked prefill does not sample a token, so num_sampled can be 0.
297 # Mamba treats num_accepted_tokens=1 as the neutral non-spec value.
298 if not isinstance(num_sampled, int):
299 # idx_mapping may contain -1 sentinels (filtered rows) under PP; the
300 # kernel skips them rather than scattering with a host-side gather.
155 - num_reqs = idx_mapping.shape[0]
156 - if num_reqs:
157 - _scatter_num_accepted_kernel[(num_reqs,)](
301 + n = idx_mapping.shape[0]
302 + if n:
303 + _scatter_num_accepted_kernel[(n,)](
304 idx_mapping, num_sampled, self.num_accepted_tokens_gpu
305 )
160 - return
306 + else:
307 + # Fill with single value.
308 + self.num_accepted_tokens_gpu.index_fill_(
309 + 0, idx_mapping, max(num_sampled, 1)
310 + )
311
162 - # Fill with single value.
163 - self.num_accepted_tokens_gpu.index_fill_(0, idx_mapping, max(num_sampled, 1))
312 + # Align: save the running state to the block-aligned position when
313 + # spec-decode acceptance leaves the sequence non-block-aligned (mirrors
314 + # the V1 align postprocess). num_computed_tokens already holds the
315 + # post-step advanced count.
316 + if (
317 + self._align_mode
318 + and num_computed_tokens is not None
319 + and self._mamba_ctx is not None
320 + ):
321 + num_reqs = idx_mapping.shape[0]
322 + if num_reqs:
323 + self._mamba_ctx.run_fused_postprocess_align(
324 + num_reqs,
325 + self.num_accepted_tokens_gpu,
326 + self._mamba_state_idx_gpu,
327 + num_computed_tokens,
328 + idx_mapping,
329 + )
330
331
332 @triton.jit
vllm/v1/worker/gpu/warmup.py
+12 -3
@@ -17,6 +17,7 @@ from vllm.v1.core.sched.output import (
17 NewRequestData,
18 SchedulerOutput,
19 )
20 +from vllm.v1.kv_cache_interface import MambaSpec
21 from vllm.v1.request import Request
22 from vllm.v1.worker.gpu.model_runner import GPUModelRunner
23
@@ -177,9 +178,17 @@ def warmup_kernels(
178 num_kv_cache_groups = len(kv_cache_groups)
179
180 # Compute per-request block counts for each KV cache group.
180 - group_block_sizes = [g.kv_cache_spec.block_size for g in kv_cache_groups]
181 - prefill_block_counts = [cdiv(prompt_len, bs) for bs in group_block_sizes]
182 - decode_block_counts = [cdiv(decode_len, bs) for bs in group_block_sizes]
181 + def _warmup_block_count(num_tokens: int, spec: Any) -> int:
182 + num_blocks = cdiv(num_tokens, spec.block_size)
183 + if isinstance(spec, MambaSpec) and spec.mamba_cache_mode == "align":
184 + # Align mode reserves extra blocks beyond the token range for the
185 + # speculative-decode running-state snapshots.
186 + num_blocks += spec.num_speculative_blocks
187 + return num_blocks
188 +
189 + kv_cache_specs = [g.kv_cache_spec for g in kv_cache_groups]
190 + prefill_block_counts = [_warmup_block_count(prompt_len, s) for s in kv_cache_specs]
191 + decode_block_counts = [_warmup_block_count(decode_len, s) for s in kv_cache_specs]
192 decode_block_deltas = [
193 d - p for d, p in zip(decode_block_counts, prefill_block_counts)
194 ]
vllm/v1/worker/mamba_utils.py
+367 -91
@@ -23,6 +23,112 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState
23 from vllm.v1.worker.lora_model_runner_mixin import GPUInputBatch
24
25
26 +@triton.jit
27 +def _copy_mamba_state_block(
28 + state_idx,
29 + bt_row_idx,
30 + src_col,
31 + dst_col,
32 + token_bias,
33 + block_table_ptrs_ptr,
34 + block_table_stride_req,
35 + state_base_addrs_ptr,
36 + state_block_strides_ptr,
37 + state_elem_sizes_ptr,
38 + state_inner_sizes_ptr,
39 + state_conv_widths_ptr,
40 + state_group_indices_ptr,
41 + # DS conv row metadata. Zero keeps the single-region copy path.
42 + state_dim_row_count_ptr,
43 + state_dim_row_stride_ptr,
44 + COPY_BLOCK_SIZE: tl.constexpr,
45 + CONV_STATE_DIM_FIRST: tl.constexpr,
46 +):
47 + """Copy one (layer, state-type) mamba state block between block columns.
48 +
49 + Shared copy body of ``postprocess_mamba_fused_kernel`` and
50 + ``precopy_mamba_align_fused_kernel``, mirroring the V1 copy specs
51 + (``get_conv_copy_spec`` / ``get_temporal_copy_spec``):
52 + - conv state (conv_width > 0): shift the window by ``token_bias`` tokens,
53 + ``state[bt[src_col], token_bias:] ->
54 + state[bt[dst_col], :conv_width - token_bias]``
55 + - temporal state: ``token_bias`` selects the accepted speculative column,
56 + ``state[bt[src_col + token_bias]] -> state[bt[dst_col]]``
57 +
58 + The caller owns the decision logic (which columns, whether to copy); this
59 + device function only performs the byte copy for the given metadata slot.
60 + """
61 + state_base_addr = tl.load(state_base_addrs_ptr + state_idx)
62 + state_block_stride = tl.load(state_block_strides_ptr + state_idx)
63 + state_elem_size = tl.load(state_elem_sizes_ptr + state_idx)
64 + state_inner_size = tl.load(state_inner_sizes_ptr + state_idx)
65 + conv_width = tl.load(state_conv_widths_ptr + state_idx)
66 +
67 + # Load the group index for this state, then index into the correct
68 + # group's block table. Each mamba group has independently allocated
69 + # physical blocks. Reinterpret as int32* since block ids are int32.
70 + group_idx = tl.load(state_group_indices_ptr + state_idx).to(tl.int64)
71 + group_base_addr = tl.load(block_table_ptrs_ptr + group_idx)
72 + block_table_typed = group_base_addr.to(tl.pointer_type(tl.int32))
73 + block_table_base = block_table_typed + bt_row_idx * block_table_stride_req
74 +
75 + # Widen block ids to int64 before they reach `block_id * state_block_stride`
76 + # below: state_block_stride can exceed 2**31 bytes for large mamba caches,
77 + # and Triton would otherwise do the multiply in int32 and wrap.
78 + dest_block_id = tl.load(block_table_base + dst_col).to(tl.int64)
79 + dst_addr = state_base_addr + dest_block_id * state_block_stride
80 +
81 + is_conv_state = conv_width > 0
82 +
83 + if CONV_STATE_DIM_FIRST and is_conv_state:
84 + # DS conv layout: state_len is the slide axis; copy per dim row.
85 + src_block_id = tl.load(block_table_base + src_col).to(tl.int64)
86 + dim_rows = tl.load(state_dim_row_count_ptr + state_idx)
87 + row_stride = tl.load(state_dim_row_stride_ptr + state_idx)
88 + per_row_bytes = (conv_width - token_bias).to(tl.int64) * state_elem_size
89 + bias_bytes = token_bias.to(tl.int64) * state_elem_size
90 + src_block_addr = state_base_addr + src_block_id * state_block_stride
91 + offsets = tl.arange(0, COPY_BLOCK_SIZE)
92 + for d in range(0, dim_rows):
93 + row_src = src_block_addr + d * row_stride + bias_bytes
94 + row_dst = dst_addr + d * row_stride
95 + for i in range(0, per_row_bytes, COPY_BLOCK_SIZE):
96 + mask = (i + offsets) < per_row_bytes
97 + curr_src = (row_src + i + offsets).to(tl.pointer_type(tl.uint8))
98 + curr_dst = (row_dst + i + offsets).to(tl.pointer_type(tl.uint8))
99 + data = tl.load(curr_src, mask=mask)
100 + tl.store(curr_dst, data, mask=mask)
101 + return
102 +
103 + if is_conv_state:
104 + # SD conv: copy
105 + # state[bt[src_col], token_bias:] ->
106 + # state[bt[dst_col], :conv_width - token_bias]
107 + src_block_id = tl.load(block_table_base + src_col).to(tl.int64)
108 + src_offset = token_bias.to(tl.int64) * state_inner_size * state_elem_size
109 + src_addr = state_base_addr + src_block_id * state_block_stride + src_offset
110 + num_elems_to_copy = (conv_width - token_bias).to(tl.int64) * state_inner_size
111 + copy_size = num_elems_to_copy * state_elem_size
112 + else:
113 + # Temporal state: copy state[bt[src_col + token_bias]] -> state[bt[dst_col]]
114 + actual_src_block_id = tl.load(block_table_base + src_col + token_bias).to(
115 + tl.int64
116 + )
117 + src_addr = state_base_addr + actual_src_block_id * state_block_stride
118 + # Use natural block data size (inner_size * elem_size), NOT
119 + # state_block_stride which is the page stride and can exceed the
120 + # actual data when the state tensor uses as_strided page padding.
121 + copy_size = state_inner_size * state_elem_size
122 +
123 + offsets = tl.arange(0, COPY_BLOCK_SIZE)
124 + for i in range(0, copy_size, COPY_BLOCK_SIZE):
125 + mask = (i + offsets) < copy_size
126 + curr_src = (src_addr + i + offsets).to(tl.pointer_type(tl.uint8))
127 + curr_dst = (dst_addr + i + offsets).to(tl.pointer_type(tl.uint8))
128 + data = tl.load(curr_src, mask=mask)
129 + tl.store(curr_dst, data, mask=mask)
130 +
131 +
132 @triton.jit
133 def postprocess_mamba_fused_kernel(
134 # Decision inputs (per-request)
@@ -49,6 +155,10 @@ def postprocess_mamba_fused_kernel(
155 state_dim_row_stride_ptr, # int64: bytes between rows for DS conv
156 # Output: num_accepted_tokens update (for src==dst case)
157 num_accepted_tokens_out_ptr,
158 + # Optional: batch_idx -> req_idx mapping (V2 model runner / PP). The
159 + # per-request decision arrays are in req-state-slot order; the block table
160 + # is in batch order, so HAS_IDX_MAPPING splits the two indexings.
161 + idx_mapping_ptr,
162 # Runtime parameter (varies per batch - NOT constexpr to avoid recompilation)
163 num_reqs,
164 # Compile-time constants (fixed after model initialization)
@@ -57,35 +167,53 @@ def postprocess_mamba_fused_kernel(
167 # COPY_BLOCK_SIZE: fixed tuning parameter for memory copy loop
168 COPY_BLOCK_SIZE: tl.constexpr,
169 CONV_STATE_DIM_FIRST: tl.constexpr,
170 + # HAS_IDX_MAPPING: when True, program_id(0) is a batch index resolved to a
171 + # req-state slot via idx_mapping_ptr (V2). When False, it is the req index.
172 + HAS_IDX_MAPPING: tl.constexpr = False,
173 + # PRECOMPUTED_NEW_COMPUTED: when True, num_computed_tokens_ptr already holds
174 + # the post-step new_num_computed value (V2 supplies the advanced count).
175 + PRECOMPUTED_NEW_COMPUTED: tl.constexpr = False,
176 ):
177 """
178 Fused GPU kernel for postprocess_mamba that computes decisions AND performs
179 mamba state copies without any CPU-GPU synchronization.
180
181 Grid: (num_reqs, num_layers * num_state_types)
66 - - program_id(0) = request index
182 + - program_id(0) = request/batch index
183 - program_id(1) = state_idx (flattened index into layer/state_type metadata)
184
185 Note: num_layers and num_state_types are not passed as kernel parameters
186 because the kernel indexes directly into pre-flattened metadata arrays
187 using program_id(1). The grid dimensions encode the total state count.
188 """
73 - req_idx = tl.program_id(0)
189 + batch_idx = tl.program_id(0)
190 state_idx = tl.program_id(1)
191
192 # Bounds check
77 - if req_idx >= num_reqs:
193 + if batch_idx >= num_reqs:
194 return
195
196 + if HAS_IDX_MAPPING:
197 + req_idx = tl.load(idx_mapping_ptr + batch_idx)
198 + if req_idx < 0:
199 + return
200 + else:
201 + req_idx = batch_idx
202 +
203 # Compute decision logic (mirrors postprocess_mamba Python reference)
204 num_accepted = tl.load(num_accepted_tokens_ptr + req_idx)
205 src_block_idx = tl.load(mamba_state_idx_ptr + req_idx)
83 - num_scheduled = tl.load(num_scheduled_tokens_ptr + req_idx)
84 - num_computed = tl.load(num_computed_tokens_ptr + req_idx)
85 - num_draft = tl.load(num_draft_tokens_ptr + req_idx)
206
87 - num_tokens_running_state = num_computed + num_scheduled - num_draft
88 - new_num_computed = num_tokens_running_state + num_accepted - 1
207 + if PRECOMPUTED_NEW_COMPUTED:
208 + new_num_computed = tl.load(num_computed_tokens_ptr + req_idx)
209 + num_tokens_running_state = new_num_computed - num_accepted + 1
210 + else:
211 + num_scheduled = tl.load(num_scheduled_tokens_ptr + req_idx)
212 + num_computed = tl.load(num_computed_tokens_ptr + req_idx)
213 + num_draft = tl.load(num_draft_tokens_ptr + req_idx)
214 + num_tokens_running_state = num_computed + num_scheduled - num_draft
215 + new_num_computed = num_tokens_running_state + num_accepted - 1
216 +
217 aligned_new_computed = (new_num_computed // block_size) * block_size
218
219 needs_copy = aligned_new_computed >= num_tokens_running_state
@@ -97,99 +225,158 @@ def postprocess_mamba_fused_kernel(
225 accept_token_bias = aligned_new_computed - num_tokens_running_state
226 dest_block_idx = aligned_new_computed // block_size - 1
227
100 - # Load state metadata for this layer/state_type
101 - state_base_addr = tl.load(state_base_addrs_ptr + state_idx)
102 - state_block_stride = tl.load(state_block_strides_ptr + state_idx)
103 - state_elem_size = tl.load(state_elem_sizes_ptr + state_idx)
104 - state_inner_size = tl.load(state_inner_sizes_ptr + state_idx)
105 - conv_width = tl.load(state_conv_widths_ptr + state_idx)
228 + # Update accepted-token count before early exits (per-request, so only
229 + # state_idx == 0 writes). V2 updates in place; V1 writes the _out buffer.
230 + if src_block_idx == dest_block_idx and state_idx == 0:
231 + if HAS_IDX_MAPPING:
232 + tl.store(num_accepted_tokens_ptr + req_idx, 1)
233 + else:
234 + tl.store(num_accepted_tokens_out_ptr + req_idx, 1)
235
107 - # Load the group index for this state, then index into the correct
108 - # group's block table. Each mamba group has independently allocated
109 - # physical blocks.
110 - group_idx = tl.load(state_group_indices_ptr + state_idx).to(tl.int64)
236 + # Skip no-op self-copy.
237 + if src_block_idx == dest_block_idx and accept_token_bias == 0:
238 + return
239
112 - # block_table_ptrs_ptr holds one pointer per group (each group owns its own
113 - # block table). Reinterpret as int32* since block ids are int32.
114 - group_base_addr = tl.load(block_table_ptrs_ptr + group_idx)
115 - block_table_typed = group_base_addr.to(tl.pointer_type(tl.int32))
116 - block_table_base = block_table_typed + req_idx * block_table_stride_req
240 + bt_row_idx = batch_idx if HAS_IDX_MAPPING else req_idx
241 + _copy_mamba_state_block(
242 + state_idx,
243 + bt_row_idx,
244 + src_block_idx,
245 + dest_block_idx,
246 + accept_token_bias,
247 + block_table_ptrs_ptr,
248 + block_table_stride_req,
249 + state_base_addrs_ptr,
250 + state_block_strides_ptr,
251 + state_elem_sizes_ptr,
252 + state_inner_sizes_ptr,
253 + state_conv_widths_ptr,
254 + state_group_indices_ptr,
255 + state_dim_row_count_ptr,
256 + state_dim_row_stride_ptr,
257 + COPY_BLOCK_SIZE,
258 + CONV_STATE_DIM_FIRST,
259 + )
260
118 - # Widen block ids to int64 before they reach `block_id * state_block_stride`
119 - # below: state_block_stride can exceed 2**31 bytes for large mamba caches,
120 - # and Triton would otherwise do the multiply in int32 and wrap.
121 - src_block_id = tl.load(block_table_base + src_block_idx).to(tl.int64)
122 - dest_block_id = tl.load(block_table_base + dest_block_idx).to(tl.int64)
261
124 - # Compute source and destination addresses based on state type
125 - # conv_width > 0 means this is a conv state (get_conv_copy_spec logic)
126 - # conv_width == 0 means this is a temporal state (get_temporal_copy_spec logic)
127 - is_conv_state = conv_width > 0
262 +@triton.jit
263 +def preprocess_mamba_align_fused_kernel(
264 + idx_mapping_ptr,
265 + state_idx_ptr,
266 + num_computed_tokens_ptr,
267 + query_start_loc_ptr,
268 + num_accepted_tokens_ptr,
269 + src_col_ptr,
270 + src_off_ptr,
271 + num_reqs,
272 + BLOCK_SIZE: tl.constexpr,
273 + MAMBA_BLOCK_SIZE: tl.constexpr,
274 +):
275 + """Fused align preprocess: emit the pre-copy src column/offset AND advance
276 + state_idx (with accepted-token reset) in a single launch (V2 align).
277 +
278 + Per batch_idx (0..num_reqs-1), resolving req slot via idx_mapping:
279 + 1. Read pre-advance state_idx and num_accepted (last step's values).
280 + 2. Store the pre-copy src columns for ``precopy_mamba_align_fused_kernel``:
281 + - src_col = state_idx (the previous running block column)
282 + - src_off = max(num_accepted - 1, 0) (the accepted-token bias)
283 + 3. Advance state_idx to the new running block, and reset num_accepted to 1
284 + when a block boundary is crossed (so the migrated state, now at the
285 + start of the new block, is read with the neutral bias).
286 + """
287 + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
288 + mask = offsets < num_reqs
289 + req_indices = tl.load(idx_mapping_ptr + offsets, mask=mask, other=0)
290
129 - # Update accepted-token count before early exits.
130 - if src_block_idx == dest_block_idx and state_idx == 0:
131 - tl.store(num_accepted_tokens_out_ptr + req_idx, 1)
291 + state_idx = tl.load(state_idx_ptr + req_indices, mask=mask, other=-1)
292 + num_accepted = tl.load(num_accepted_tokens_ptr + req_indices, mask=mask, other=1)
293
133 - # Skip no-op self-copy.
134 - if src_block_idx == dest_block_idx and accept_token_bias == 0:
135 - return
294 + src_off = tl.maximum(num_accepted - 1, 0)
295 + tl.store(src_col_ptr + req_indices, state_idx, mask=mask)
296 + tl.store(src_off_ptr + req_indices, src_off, mask=mask)
297
137 - if CONV_STATE_DIM_FIRST and is_conv_state:
138 - dim_rows = tl.load(state_dim_row_count_ptr + state_idx)
139 - row_stride = tl.load(state_dim_row_stride_ptr + state_idx)
140 - per_row_bytes = (conv_width - accept_token_bias).to(tl.int64) * state_elem_size
141 - bias_bytes = accept_token_bias.to(tl.int64) * state_elem_size
142 - src_block_addr = state_base_addr + src_block_id * state_block_stride
143 - dst_block_addr = state_base_addr + dest_block_id * state_block_stride
144 - offsets = tl.arange(0, COPY_BLOCK_SIZE)
145 - for d in range(0, dim_rows):
146 - row_src = src_block_addr + d * row_stride + bias_bytes
147 - row_dst = dst_block_addr + d * row_stride
148 - for i in range(0, per_row_bytes, COPY_BLOCK_SIZE):
149 - mask = (i + offsets) < per_row_bytes
150 - curr_src = (row_src + i + offsets).to(tl.pointer_type(tl.uint8))
151 - curr_dst = (row_dst + i + offsets).to(tl.pointer_type(tl.uint8))
152 - data = tl.load(curr_src, mask=mask)
153 - tl.store(curr_dst, data, mask=mask)
298 + num_computed = tl.load(num_computed_tokens_ptr + req_indices, mask=mask, other=0)
299 + query_start = tl.load(query_start_loc_ptr + offsets, mask=mask, other=0)
300 + query_end = tl.load(query_start_loc_ptr + offsets + 1, mask=mask, other=0)
301 + computed_after = num_computed + query_end - query_start
302 + new_state_idx = (computed_after + MAMBA_BLOCK_SIZE - 1) // MAMBA_BLOCK_SIZE - 1
303 + tl.store(state_idx_ptr + req_indices, new_state_idx, mask=mask)
304 + should_reset = (state_idx >= 0) & (state_idx != new_state_idx)
305 + tl.store(num_accepted_tokens_ptr + req_indices, 1, mask=mask & should_reset)
306 +
307 +
308 +@triton.jit
309 +def precopy_mamba_align_fused_kernel(
310 + # Per-request-slot inputs (indexed by req_idx via idx_mapping), produced by
311 + # the V2 fused align preprocess kernel for the current step:
312 + mamba_state_idx_ptr, # post-advance dst block column
313 + src_col_ptr, # pre-advance src block column (-1 = fresh)
314 + token_bias_ptr, # accepted-token bias = num_accepted - 1 (pre-reset)
315 + # Same flattened state-layout metadata as postprocess_mamba_fused_kernel
316 + block_table_ptrs_ptr,
317 + block_table_stride_req: tl.int64,
318 + state_base_addrs_ptr,
319 + state_block_strides_ptr,
320 + state_elem_sizes_ptr,
321 + state_inner_sizes_ptr,
322 + state_conv_widths_ptr,
323 + state_group_indices_ptr,
324 + state_dim_row_count_ptr,
325 + state_dim_row_stride_ptr,
326 + idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx (-1 to skip)
327 + num_reqs,
328 + COPY_BLOCK_SIZE: tl.constexpr,
329 + CONV_STATE_DIM_FIRST: tl.constexpr,
330 +):
331 + """Pre-copy mamba "align" state across block boundaries on the V2 runner.
332 +
333 + Before the forward pass, copy each request's last SSM/conv state from its
334 + previous block column into the new window block column, so the kernels read
335 + the initial state from the write-side block as usual (V1 align semantics).
336 + Same per-(layer, state) copy semantics as ``postprocess_mamba_fused_kernel``
337 + (shared ``_copy_mamba_state_block`` body, i.e. the V1 ``preprocess_mamba``
338 + copy specs), but driven by the GPU-resident src columns so it needs no
339 + CPU-GPU sync (async-scheduling safe).
340 +
341 + Grid: (num_reqs, num_layers * num_state_types); block tables are indexed by
342 + batch row, per-request state by req_idx via idx_mapping (V2 layout).
343 + """
344 + batch_idx = tl.program_id(0)
345 + state_idx = tl.program_id(1)
346 + if batch_idx >= num_reqs:
347 + return
348 + req_idx = tl.load(idx_mapping_ptr + batch_idx)
349 + if req_idx < 0:
350 return
351
156 - if is_conv_state:
157 - # SD conv: copy
158 - # state[block_table[req_idx, src_block_idx], accept_token_bias:]
159 - # to
160 - # state[block_table[req_idx, dest_block_idx], :conv_width - accept_token_bias]
161 - src_offset = accept_token_bias.to(tl.int64) * state_inner_size * state_elem_size
162 - src_addr = state_base_addr + src_block_id * state_block_stride + src_offset
163 - dst_addr = state_base_addr + dest_block_id * state_block_stride
164 - # Number of elements to copy:
165 - # (conv_width - accept_token_bias) * inner_size
166 - num_elems_to_copy = (conv_width - accept_token_bias).to(
167 - tl.int64
168 - ) * state_inner_size
169 - copy_size = num_elems_to_copy * state_elem_size
170 - else:
171 - # Temporal state: copy
172 - # state[block_table[req_idx, src_block_idx + accept_token_bias]]
173 - # to
174 - # state[block_table[req_idx, dest_block_idx]]
175 - actual_src_block_idx = src_block_idx + accept_token_bias
176 - actual_src_block_id = tl.load(block_table_base + actual_src_block_idx).to(
177 - tl.int64
178 - )
179 - src_addr = state_base_addr + actual_src_block_id * state_block_stride
180 - dst_addr = state_base_addr + dest_block_id * state_block_stride
181 - # Use natural block data size (inner_size * elem_size), NOT
182 - # state_block_stride which is the page stride and can exceed the
183 - # actual data when the state tensor uses as_strided page padding.
184 - copy_size = state_inner_size * state_elem_size
352 + src_col = tl.load(src_col_ptr + req_idx)
353 + dst_col = tl.load(mamba_state_idx_ptr + req_idx)
354 + # Fresh state, or still writing the same block: kernels locate the initial
355 + # state in-block via num_accepted (preserved when no boundary is crossed),
356 + # so there is nothing to copy.
357 + if src_col < 0 or src_col == dst_col:
358 + return
359
186 - offsets = tl.arange(0, COPY_BLOCK_SIZE)
187 - for i in range(0, copy_size, COPY_BLOCK_SIZE):
188 - mask = (i + offsets) < copy_size
189 - curr_src = (src_addr + i + offsets).to(tl.pointer_type(tl.uint8))
190 - curr_dst = (dst_addr + i + offsets).to(tl.pointer_type(tl.uint8))
191 - data = tl.load(curr_src, mask=mask)
192 - tl.store(curr_dst, data, mask=mask)
360 + token_bias = tl.load(token_bias_ptr + req_idx)
361 + _copy_mamba_state_block(
362 + state_idx,
363 + batch_idx,
364 + src_col,
365 + dst_col,
366 + token_bias,
367 + block_table_ptrs_ptr,
368 + block_table_stride_req,
369 + state_base_addrs_ptr,
370 + state_block_strides_ptr,
371 + state_elem_sizes_ptr,
372 + state_inner_sizes_ptr,
373 + state_conv_widths_ptr,
374 + state_group_indices_ptr,
375 + state_dim_row_count_ptr,
376 + state_dim_row_stride_ptr,
377 + COPY_BLOCK_SIZE,
378 + CONV_STATE_DIM_FIRST,
379 + )
380
381
382 @triton.jit
@@ -559,10 +746,99 @@ class MambaSpecDecodeGPUContext:
746 self.state_dim_row_count,
747 self.state_dim_row_stride,
748 self.num_accepted_tokens_out,
749 + None, # idx_mapping: V1 decision arrays are already in req order
750 + num_reqs,
751 + block_size=self.block_size,
752 + COPY_BLOCK_SIZE=1024,
753 + CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
754 + )
755 +
756 + def run_fused_precopy(
757 + self,
758 + num_reqs: int,
759 + state_idx_gpu: torch.Tensor,
760 + src_col_gpu: torch.Tensor,
761 + token_bias_gpu: torch.Tensor,
762 + idx_mapping: torch.Tensor,
763 + ) -> None:
764 + """Pre-copy each request's previous running block into its new window
765 + block before the forward pass (V2 align boundary migration).
766 +
767 + Args:
768 + num_reqs: Number of active requests (batch order).
769 + state_idx_gpu: [max_reqs] post-advance dst block column per req slot.
770 + src_col_gpu: [max_reqs] pre-advance src block column (-1 = fresh).
771 + token_bias_gpu: [max_reqs] accepted-token bias (num_accepted - 1).
772 + idx_mapping: [num_reqs] batch_idx -> req_state_idx (-1 to skip).
773 + """
774 + if num_reqs == 0 or not self.is_initialized:
775 + return
776 + total_states = self.num_layers * self.num_state_types
777 + grid = (num_reqs, total_states)
778 + precopy_mamba_align_fused_kernel[grid](
779 + state_idx_gpu,
780 + src_col_gpu,
781 + token_bias_gpu,
782 + self.block_table_ptrs,
783 + self.block_table_stride_req,
784 + self.state_base_addrs,
785 + self.state_block_strides,
786 + self.state_elem_sizes,
787 + self.state_inner_sizes,
788 + self.state_conv_widths,
789 + self.state_group_indices,
790 + self.state_dim_row_count,
791 + self.state_dim_row_stride,
792 + idx_mapping,
793 + num_reqs,
794 + COPY_BLOCK_SIZE=1024,
795 + CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
796 + )
797 +
798 + def run_fused_postprocess_align(
799 + self,
800 + num_reqs: int,
801 + num_accepted_tokens_gpu: torch.Tensor,
802 + state_idx_gpu: torch.Tensor,
803 + new_num_computed_tokens_gpu: torch.Tensor,
804 + idx_mapping: torch.Tensor,
805 + ) -> None:
806 + """V2 align postprocess: save the running state to the block-aligned
807 + position after spec-decode acceptance leaves the sequence non-aligned.
808 +
809 + ``num_accepted_tokens_gpu`` is updated in place (reset to 1 when the
810 + accepted position stays in the running block); ``new_num_computed_tokens``
811 + already holds the post-step computed count (PRECOMPUTED_NEW_COMPUTED).
812 + ``idx_mapping`` maps batch row -> req-state slot (HAS_IDX_MAPPING).
813 + """
814 + if num_reqs == 0 or not self.is_initialized:
815 + return
816 + total_states = self.num_layers * self.num_state_types
817 + grid = (num_reqs, total_states)
818 + postprocess_mamba_fused_kernel[grid](
819 + num_accepted_tokens_gpu,
820 + state_idx_gpu,
821 + None, # num_scheduled: unused under PRECOMPUTED_NEW_COMPUTED
822 + new_num_computed_tokens_gpu,
823 + None, # num_draft: unused under PRECOMPUTED_NEW_COMPUTED
824 + self.block_table_ptrs,
825 + self.block_table_stride_req,
826 + self.state_base_addrs,
827 + self.state_block_strides,
828 + self.state_elem_sizes,
829 + self.state_inner_sizes,
830 + self.state_conv_widths,
831 + self.state_group_indices,
832 + self.state_dim_row_count,
833 + self.state_dim_row_stride,
834 + None, # num_accepted_out: V2 updates num_accepted in place
835 + idx_mapping,
836 num_reqs,
837 block_size=self.block_size,
838 COPY_BLOCK_SIZE=1024,
839 CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
840 + HAS_IDX_MAPPING=True,
841 + PRECOMPUTED_NEW_COMPUTED=True,
842 )
843
844