A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[V1] [Hybrid] Some additional clean-up in Mamba2 prefix caching (#26222)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Thomas Parnell committed
Oct 6, 2025 at 04:40 UTC
778f554157b7a486f62bf7ad869d9ace5876e197
4 files changed
+171
-136
vllm/model_executor/layers/mamba/mamba_mixer2.py
+96
-57
@@ -595,21 +595,32 @@ class MambaMixer2(MambaBase, CustomOp):
595
if prefix_caching_enabled:
596
# If prefix caching is enabled, retrieve the relevant variables
597
# for prefill and decode
598
- last_state_idx_d, last_state_idx_p = torch.split(
599
- attn_metadata.last_state_idx, [num_decodes, num_prefills], dim=0
598
+ block_idx_last_computed_token_d, block_idx_last_computed_token_p = (
599
+ torch.split(
600
+ attn_metadata.block_idx_last_computed_token,
601
+ [num_decodes, num_prefills],
602
+ dim=0,
603
+ )
604
)
601
- current_last_idx_d, current_last_idx_p = torch.split(
602
- attn_metadata.current_last_idx, [num_decodes, num_prefills], dim=0
605
+ block_idx_last_scheduled_token_d, block_idx_last_scheduled_token_p = (
606
+ torch.split(
607
+ attn_metadata.block_idx_last_scheduled_token,
608
+ [num_decodes, num_prefills],
609
+ dim=0,
610
+ )
611
)
612
# Prefill-only variables:
605
- current_first_idx_p = attn_metadata.current_first_idx_p
606
- context_lens_p = attn_metadata.context_lens_p
607
- last_computed_offset_p = attn_metadata.last_computed_offset_p
613
+ block_idx_first_scheduled_token_p = (
614
+ attn_metadata.block_idx_first_scheduled_token_p
615
+ )
616
+ num_computed_tokens_p = attn_metadata.num_computed_tokens_p
617
else:
609
- last_state_idx_d, last_state_idx_p = None, None
610
- current_last_idx_d, current_last_idx_p = None, None
611
- current_first_idx_p = None
612
- context_lens_p = None
618
+ block_idx_last_computed_token_d = None
619
+ block_idx_last_computed_token_p = None
620
+ block_idx_last_scheduled_token_d = None
621
+ block_idx_last_scheduled_token_p = None
622
+ block_idx_first_scheduled_token_p = None
623
+ num_computed_tokens_p = None
624
625
# Preallocate output tensor to avoid memcpy cost for merging prefill
626
# and decode outputs
@@ -637,7 +648,8 @@ class MambaMixer2(MambaBase, CustomOp):
648
# to by "state_indices_tensor_p".
649
# In particular, it will always write the state at the
650
# sequence end.
640
- # In addition, "current_first_idx_p" and "current_last_idx_p"
651
+ # In addition, "block_idx_first_scheduled_token_p" and
652
+ # "block_idx_last_scheduled_token_p"
653
# are provided (which are pointers into
654
# "state_indices_tensor_p"), it will write additional cache
655
# states aligned at "block_size_to_align".
@@ -652,10 +664,10 @@ class MambaMixer2(MambaBase, CustomOp):
664
conv_states=conv_state,
665
has_initial_state=has_initial_states_p,
666
cache_indices=state_indices_tensor_p,
655
- current_first_idx=current_first_idx_p,
656
- current_last_idx=current_last_idx_p,
657
- initial_state_idx=last_state_idx_p,
658
- context_lens=context_lens_p,
667
+ block_idx_first_scheduled_token=block_idx_first_scheduled_token_p,
668
+ block_idx_last_scheduled_token=block_idx_last_scheduled_token_p,
669
+ initial_state_idx=block_idx_last_computed_token_p,
670
+ num_computed_tokens=num_computed_tokens_p,
671
block_size_to_align=mamba_block_size,
672
metadata=attn_metadata,
673
query_start_loc=query_start_loc_p,
@@ -669,7 +681,7 @@ class MambaMixer2(MambaBase, CustomOp):
681
kernel_ssm_indices = state_indices_tensor_p
682
if prefix_caching_enabled:
683
kernel_ssm_indices = state_indices_tensor_p.gather(
672
- 1, last_state_idx_p.unsqueeze(1)
684
+ 1, block_idx_last_computed_token_p.unsqueeze(1)
685
).squeeze(1)
686
initial_states = torch.where(
687
has_initial_states_p[:, None, None, None],
@@ -703,52 +715,76 @@ class MambaMixer2(MambaBase, CustomOp):
715
)
716
717
if prefix_caching_enabled:
706
- # Save states for sequences with more than just the final state:
707
- n_blocks_to_fill = current_last_idx_p - current_first_idx_p
708
- for seq_idx in (n_blocks_to_fill > 0).nonzero().squeeze(1):
718
+ # The chunk_stride is the number of chunks per mamba block
719
+ # e.g., if mamba_block_size = 512 and chunk_size = 256,
720
+ # then chunk_stride = 2
721
+ chunk_stride = mamba_block_size // chunk_size
722
+
723
+ # Save state for sequences with more than just final state
724
+ for seq_idx in range(num_prefills):
725
+ # Block index for the first scheduled token
726
+ block_idx_first_scheduled_token = block_idx_first_scheduled_token_p[
727
+ seq_idx
728
+ ]
729
+
730
+ # Block index for the last scheduled token
731
+ block_idx_last_scheduled_token = block_idx_last_scheduled_token_p[
732
+ seq_idx
733
+ ]
734
+
735
+ # Number of blocks that need to be written
736
+ n_blocks_to_fill = (
737
+ block_idx_last_scheduled_token - block_idx_first_scheduled_token
738
+ )
739
+
740
+ # Skip sequences that don't have any blocks to fill
741
+ if n_blocks_to_fill == 0:
742
+ continue
743
+
744
+ # Look up the state indices
745
cache_blocks_to_fill = state_indices_tensor_p[
746
seq_idx,
711
- current_first_idx_p[seq_idx] : current_first_idx_p[seq_idx]
712
- + n_blocks_to_fill[seq_idx],
747
+ block_idx_first_scheduled_token:block_idx_last_scheduled_token,
748
]
714
- # chunks = [0 1 2 3 4 5 6 ...]
715
- # First aligned chunk would typically be:
716
- # mamba_block_size = 1024, chunk_size = 256
717
- # 1024 // 256 - 1 --> chunks[3]
718
- # But when last chunk wasn't block aligned:
719
- # - last_computed_offset_p[seq_idx] // chunk_size
720
- # e.g. 1000 // 256 -> 3 completed --> store chunk[0]
721
- # e.g. 513 // 256 -> 2 completed --> store chunk[1] (skip 1)
722
- # e.g. 256 // 256 -> 1 completed --> store chunk[2] (skip 2)
723
- # e.g. 10 // 256 -> 0 completed --> store chunk[3] (skip 3)
724
- chunk_stride = mamba_block_size // chunk_size
725
- first_aligned_chunk = (
726
- torch.concat(
727
- [
728
- torch.zeros(
729
- 1,
730
- dtype=last_chunk_indices_p.dtype,
731
- device=last_chunk_indices_p.device,
732
- ),
733
- last_chunk_indices_p + 1,
734
- ]
735
- )[seq_idx]
736
- + chunk_stride
737
- - 1
738
- - last_computed_offset_p[seq_idx] // chunk_size
749
+
750
+ # First chunk index for this sequence
751
+ if seq_idx == 0:
752
+ first_chunk = 0
753
+ else:
754
+ first_chunk = 1 + last_chunk_indices_p[seq_idx - 1]
755
+
756
+ # First chunk that is aligned on the mamba block boundary
757
+ first_aligned_chunk = first_chunk + chunk_stride - 1
758
+
759
+ # Calculate the number of computed tokens that were not
760
+ # already cached
761
+ num_unaligned_computed_tokens = (
762
+ num_computed_tokens_p[seq_idx] % mamba_block_size
763
)
764
+
765
+ if num_unaligned_computed_tokens > 0:
766
+ # If the number of computed tokens is not block aligned,
767
+ # then we need to shift the index accordingly
768
+ first_aligned_chunk -= (
769
+ num_unaligned_computed_tokens // chunk_size
770
+ )
771
+
772
+ # Get states to write
773
from_where = varlen_states[
774
first_aligned_chunk : first_aligned_chunk
742
- + n_blocks_to_fill[seq_idx] * chunk_stride : chunk_stride
775
+ + n_blocks_to_fill * chunk_stride : chunk_stride
776
]
777
+
778
+ # Write the states
779
ssm_state[cache_blocks_to_fill] = from_where
780
746
- # For all seqs, store the last state (Note: might be partial):
781
+ # For all seqs, store the last state (note: might be partial):
782
ssm_state[
783
state_indices_tensor_p.gather(
749
- 1, current_last_idx_p.unsqueeze(1)
784
+ 1, block_idx_last_scheduled_token_p.unsqueeze(1)
785
).squeeze(1)
786
] = varlen_states[last_chunk_indices_p]
787
+
788
else:
789
# update ssm states
790
# - varlen state is a (num_prefills, nheads, headdim, dstate)
@@ -759,14 +795,17 @@ class MambaMixer2(MambaBase, CustomOp):
795
if has_decode:
796
if prefix_caching_enabled:
797
state_indices_tensor_d_input = state_indices_tensor_d.gather(
762
- 1, last_state_idx_d.unsqueeze(1)
798
+ 1, block_idx_last_computed_token_d.unsqueeze(1)
799
).squeeze(1)
800
state_indices_tensor_d_output = state_indices_tensor_d.gather(
765
- 1, current_last_idx_d.unsqueeze(1)
801
+ 1, block_idx_last_scheduled_token_d.unsqueeze(1)
802
).squeeze(1)
767
- # Note:
768
- # for decode always: current_first_idx_d == current_last_idx_d
769
- # at block boundaries: current_first_idx_d > last_state_idx_d
803
+ # for decode:
804
+ # block_idx_first_scheduled_token_d ==
805
+ # block_idx_last_scheduled_token_d
806
+ # at block boundaries:
807
+ # block_idx_first_scheduled_token_d >
808
+ # block_idx_last_computed_token_d
809
else:
810
# Without caching, read and write in-place to the same blocks:
811
state_indices_tensor_d_input = state_indices_tensor_d
@@ -780,8 +819,8 @@ class MambaMixer2(MambaBase, CustomOp):
819
self.conv1d.bias,
820
self.activation,
821
conv_state_indices=state_indices_tensor_d,
783
- current_last_idx=current_last_idx_d,
784
- initial_state_idx=last_state_idx_d,
822
+ block_idx_last_scheduled_token=block_idx_last_scheduled_token_d,
823
+ initial_state_idx=block_idx_last_computed_token_d,
824
)
825
826
hidden_states_d, B_d, C_d = split_hidden_states_B_C_fn(hidden_states_B_C_d)
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
+22
-22
@@ -27,10 +27,10 @@ def _causal_conv1d_fwd_kernel( # continuous batching
27
query_start_loc_ptr,
28
batch_ptr,
29
token_chunk_offset_ptr,
30
- current_first_idx, # (batch,)
31
- current_last_idx, # (batch,)
30
+ block_idx_first_scheduled_token, # (batch,)
31
+ block_idx_last_scheduled_token, # (batch,)
32
initial_state_idx, # (batch,)
33
- context_lens, # (batch,)
33
+ num_computed_tokens, # (batch,)
34
o_ptr, # (dim, seqlen) - actually pointing to x_ptr
35
# Matrix dimensions
36
dim: tl.constexpr,
@@ -94,9 +94,9 @@ def _causal_conv1d_fwd_kernel( # continuous batching
94
# In particular, if prefix caching is enabled, the program write additional cache states to "cache_indices_ptr"
95
96
# Get the length of the completed sequence so far and compute the offset.
97
- current_first_index = tl.load(current_first_idx + idx_seq)
98
- current_last_index = tl.load(current_last_idx + idx_seq)
99
- sequence_completed_index = tl.load(context_lens + idx_seq)
97
+ current_first_index = tl.load(block_idx_first_scheduled_token + idx_seq)
98
+ current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq)
99
+ sequence_completed_index = tl.load(num_computed_tokens + idx_seq)
100
101
# Compute the offset where the first stride_block_m-aligned first full block is
102
# Value in "token-space"
@@ -476,10 +476,10 @@ def causal_conv1d_fn(
476
has_initial_state: Optional[torch.Tensor] = None,
477
activation: Optional[str] = "silu",
478
pad_slot_id: int = PAD_SLOT_ID,
479
- current_first_idx: Optional[torch.Tensor] = None,
480
- current_last_idx: Optional[torch.Tensor] = None,
479
+ block_idx_first_scheduled_token: Optional[torch.Tensor] = None,
480
+ block_idx_last_scheduled_token: Optional[torch.Tensor] = None,
481
initial_state_idx: Optional[torch.Tensor] = None,
482
- context_lens: Optional[torch.Tensor] = None,
482
+ num_computed_tokens: Optional[torch.Tensor] = None,
483
block_size_to_align=0,
484
metadata=None,
485
validate_data=False,
@@ -523,13 +523,13 @@ def causal_conv1d_fn(
523
for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id]
524
in this case, the kernel will not process entries at
525
indices 0 and 3
526
- current_first_idx: (batch,), dtype int32
526
+ block_idx_first_scheduled_token: (batch,), dtype int32
527
The pointer into cache_indices, where the first cache block to be filled is located.
528
- current_last_idx: (batch,), dtype int32
528
+ block_idx_last_scheduled_token: (batch,), dtype int32
529
The pointer into cache_indices, where the last cache block to be filled is located.
530
initial_state_idx: (batch,), dtype int32
531
The pointer into cache_indices, where the cache block containing the initial state is located.
532
- context_lens: (batch,), dtype int32
532
+ num_computed_tokens: (batch,), dtype int32
533
The number of tokens already completed for each sequence
534
block_size_to_align: int
535
The block size to align the cached states to
@@ -708,10 +708,10 @@ def causal_conv1d_fn(
708
query_start_loc,
709
batch_ptr,
710
token_chunk_offset_ptr,
711
- current_first_idx,
712
- current_last_idx,
711
+ block_idx_first_scheduled_token,
712
+ block_idx_last_scheduled_token,
713
initial_state_idx,
714
- context_lens,
714
+ num_computed_tokens,
715
out,
716
# Matrix dimensions
717
dim,
@@ -735,7 +735,7 @@ def causal_conv1d_fn(
735
HAS_BIAS=bias is not None,
736
KERNEL_WIDTH=width,
737
SILU_ACTIVATION=activation in ["silu", "swish"],
738
- IS_APC_ENABLED=current_last_idx is not None,
738
+ IS_APC_ENABLED=block_idx_last_scheduled_token is not None,
739
USE_PAD_SLOT=pad_slot_id is not None,
740
NP2_STATELEN=np2_statelen,
741
# launch_cooperative_grid=True
@@ -756,7 +756,7 @@ def _causal_conv1d_update_kernel(
756
conv_state_indices_ptr,
757
num_accepted_tokens_ptr,
758
query_start_loc_ptr, # (batch + 1)
759
- current_last_idx, # (batch,)
759
+ block_idx_last_scheduled_token, # (batch,)
760
initial_state_idx, # (batch,)
761
o_ptr, # (batch, dim, seqlen)
762
# Matrix dimensions
@@ -802,7 +802,7 @@ def _causal_conv1d_update_kernel(
802
if IS_APC_ENABLED:
803
# Get the state from the initial_state_idx
804
conv_state_init = tl.load(initial_state_idx + idx_seq)
805
- current_last_index = tl.load(current_last_idx + idx_seq)
805
+ current_last_index = tl.load(block_idx_last_scheduled_token + idx_seq)
806
else:
807
conv_state_init = 0
808
current_last_index = 0
@@ -1078,7 +1078,7 @@ def causal_conv1d_update(
1078
query_start_loc: Optional[torch.Tensor] = None,
1079
max_query_len: int = -1,
1080
pad_slot_id: int = PAD_SLOT_ID,
1081
- current_last_idx: Optional[torch.Tensor] = None,
1081
+ block_idx_last_scheduled_token: Optional[torch.Tensor] = None,
1082
initial_state_idx: Optional[torch.Tensor] = None,
1083
validate_data=False,
1084
):
@@ -1097,7 +1097,7 @@ def causal_conv1d_update(
1097
If not None, the conv_state is a larger tensor along the batch dim,
1098
and we are selecting the batch coords specified by conv_state_indices.
1099
Useful for a continuous batching scenario.
1100
- current_last_idx: (batch,), dtype int32
1100
+ block_idx_last_scheduled_token: (batch,), dtype int32
1101
The pointer into conv_state_indices, where the last cache block to be filled is located.
1102
initial_state_idx: (batch,), dtype int32
1103
The pointer into conv_state_indices, where the cache block containing the initial state is located.
@@ -1201,7 +1201,7 @@ def causal_conv1d_update(
1201
conv_state_indices,
1202
num_accepted_tokens,
1203
query_start_loc,
1204
- current_last_idx,
1204
+ block_idx_last_scheduled_token,
1205
initial_state_idx,
1206
out,
1207
# Matrix dimensions
@@ -1230,7 +1230,7 @@ def causal_conv1d_update(
1230
KERNEL_WIDTH=width,
1231
SILU_ACTIVATION=activation in ["silu", "swish"],
1232
IS_VARLEN=query_start_loc is not None,
1233
- IS_APC_ENABLED=current_last_idx is not None,
1233
+ IS_APC_ENABLED=block_idx_last_scheduled_token is not None,
1234
IS_SPEC_DECODING=num_accepted_tokens is not None,
1235
NP2_STATELEN=np2_statelen,
1236
USE_PAD_SLOT=pad_slot_id is not None,
vllm/v1/attention/backends/mamba2_attn.py
+52
-55
@@ -122,11 +122,10 @@ class Mamba2AttentionMetadata:
122
last_chunk_indices_p: Optional[torch.Tensor]
123
124
state_indices_tensor: torch.Tensor # shape: [batch,]
125
- current_last_idx: torch.Tensor
126
- current_first_idx_p: torch.Tensor
127
- last_state_idx: torch.Tensor
128
- context_lens_p: torch.Tensor
129
- last_computed_offset_p: torch.Tensor
125
+ block_idx_last_scheduled_token: torch.Tensor # shape: [batch,]
126
+ block_idx_first_scheduled_token_p: torch.Tensor # shape: [batch,]
127
+ block_idx_last_computed_token: torch.Tensor # shape: [batch,]
128
+ num_computed_tokens_p: torch.Tensor # shape: [batch,]
129
130
# The following attributes are for triton implementation of causal_conv1d
131
nums_dict: Optional[dict] = None
@@ -160,12 +159,12 @@ class Mamba2AttentionMetadataBuilder(
159
dtype=torch.int32,
160
device=device,
161
)
163
- self.current_last_idx = torch.empty(
162
+ self.block_idx_last_scheduled_token = torch.empty(
163
(self.decode_cudagraph_max_bs,),
164
dtype=torch.int32,
165
device=device,
166
)
168
- self.last_state_idx = torch.empty(
167
+ self.block_idx_last_computed_token = torch.empty(
168
(self.decode_cudagraph_max_bs,),
169
dtype=torch.int32,
170
device=device,
@@ -192,43 +191,38 @@ class Mamba2AttentionMetadataBuilder(
191
# for causal_conv1d
192
nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
193
195
- context_lens, context_lens_p = None, None
196
- current_first_idx, current_first_idx_p = None, None
197
- last_computed_offset, last_computed_offset_p = None, None
194
+ num_computed_tokens, num_computed_tokens_p = None, None
195
+ block_idx_first_scheduled_token = None
196
+ block_idx_first_scheduled_token_p = None
197
198
if self.vllm_config.cache_config.enable_prefix_caching:
199
# Return a tensor of shape (#requests, #max blocks)
200
state_indices_tensor = common_attn_metadata.block_table_tensor
202
-
201
# Additional cache-related varaiables:
202
mamba_block_size = self.kv_cache_spec.block_size
205
- seq_lens_pending = (
206
- torch.roll(common_attn_metadata.query_start_loc, -1, -1)
207
- - common_attn_metadata.query_start_loc
208
- )[:-1]
209
- context_lens = common_attn_metadata.seq_lens - seq_lens_pending
210
- last_computed_offset = context_lens % mamba_block_size
211
- # Indices: last_computed <= current_first <= current_last
212
- # Cases:
213
- # last_computed == current_first if last state was partially
214
- # computed and needs to be updated
215
- # current_first == current_last if no block crossing occurs, and
216
- # only one state will be stored
217
- # 0th based indexing leads to "-1" -> e.g. 16 computed -> state[15]:
218
- current_last_idx = (
219
- cdiv(context_lens + seq_lens_pending, mamba_block_size) - 1
203
+ num_computed_tokens = common_attn_metadata.num_computed_tokens_cpu.to(
204
+ self.device
205
+ )
206
+ # Block index of the last computed token
207
+ block_idx_last_computed_token = (
208
+ cdiv(num_computed_tokens, mamba_block_size) - 1
209
+ )
210
+ # which is <= block index for the first scheduled token
211
+ block_idx_first_scheduled_token = (
212
+ cdiv(num_computed_tokens + 1, mamba_block_size) - 1
213
+ )
214
+ # which is <= block index of the last scheduled token
215
+ block_idx_last_scheduled_token = (
216
+ cdiv(common_attn_metadata.seq_lens, mamba_block_size) - 1
217
)
221
- current_first_idx = cdiv(context_lens + 1, mamba_block_size) - 1
222
- last_state_idx = cdiv(context_lens, mamba_block_size) - 1
218
# -1 in case it's non-computed and causes later issues with indexing
224
- last_state_idx = last_state_idx.clamp(min=0)
225
-
219
+ block_idx_last_computed_token = block_idx_last_computed_token.clamp(min=0)
220
else:
221
# Always return just a single block per each request:
222
state_indices_tensor = common_attn_metadata.block_table_tensor[:, 0]
223
# Additional cache-related varaiables:
230
- current_last_idx = None
231
- last_state_idx = None
224
+ block_idx_last_scheduled_token = None
225
+ block_idx_last_computed_token = None
226
227
num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (
228
split_decodes_and_prefills(
@@ -256,18 +250,15 @@ class Mamba2AttentionMetadataBuilder(
250
)
251
252
if self.vllm_config.cache_config.enable_prefix_caching:
259
- assert context_lens is not None
260
- context_lens_p = context_lens[num_reqs - num_prefills : num_reqs]
261
- assert last_computed_offset is not None
262
- last_computed_offset_p = last_computed_offset[
253
+ assert num_computed_tokens is not None
254
+ num_computed_tokens_p = num_computed_tokens[
255
num_reqs - num_prefills : num_reqs
256
]
265
- assert current_first_idx is not None
266
- current_first_idx_p = current_first_idx[
257
+ assert block_idx_first_scheduled_token is not None
258
+ block_idx_first_scheduled_token_p = block_idx_first_scheduled_token[
259
num_reqs - num_prefills : num_reqs
260
]
269
-
270
- num_computed_tokens_p = common_attn_metadata.num_computed_tokens_cpu[
261
+ num_computed_tokens_p_cpu = common_attn_metadata.num_computed_tokens_cpu[
262
num_reqs - num_prefills : num_reqs
263
]
264
query_start_loc_p_cpu = (
@@ -290,7 +281,7 @@ class Mamba2AttentionMetadataBuilder(
281
last_chunk_indices = []
282
seqlen_pos = 0
283
for req_idx in range(num_prefills):
293
- this_num_computed = num_computed_tokens_p[req_idx].item()
284
+ this_num_computed = num_computed_tokens_p_cpu[req_idx].item()
285
this_new_tokens = (
286
query_start_loc_p_cpu[req_idx + 1].item()
287
- query_start_loc_p_cpu[req_idx].item()
@@ -338,7 +329,10 @@ class Mamba2AttentionMetadataBuilder(
329
compute_causal_conv1d_metadata(query_start_loc_p)
330
)
331
341
- elif num_decodes <= self.decode_cudagraph_max_bs:
332
+ elif (
333
+ num_decodes <= self.decode_cudagraph_max_bs
334
+ and self.compilation_config.full_cuda_graph
335
+ ):
336
# Pad state tensor for CUDA graph
337
num_input_tokens = self.vllm_config.pad_for_cudagraph(num_decodes)
338
self.state_indices_tensor[:num_decodes].copy_(
@@ -348,17 +342,21 @@ class Mamba2AttentionMetadataBuilder(
342
state_indices_tensor[num_decodes:] = PAD_SLOT_ID
343
344
if self.vllm_config.cache_config.enable_prefix_caching:
351
- self.current_last_idx[:num_decodes].copy_(
352
- current_last_idx, non_blocking=True
345
+ self.block_idx_last_scheduled_token[:num_decodes].copy_(
346
+ block_idx_last_scheduled_token, non_blocking=True
347
)
354
- current_last_idx = self.current_last_idx[:num_input_tokens]
355
- current_last_idx[num_decodes:] = 0
348
+ block_idx_last_scheduled_token = self.block_idx_last_scheduled_token[
349
+ :num_input_tokens
350
+ ]
351
+ block_idx_last_scheduled_token[num_decodes:] = 0
352
357
- self.last_state_idx[:num_decodes].copy_(
358
- last_state_idx, non_blocking=True
353
+ self.block_idx_last_computed_token[:num_decodes].copy_(
354
+ block_idx_last_computed_token, non_blocking=True
355
)
360
- last_state_idx = self.last_state_idx[:num_input_tokens]
361
- last_state_idx[num_decodes:] = 0
356
+ block_idx_last_computed_token = self.block_idx_last_computed_token[
357
+ :num_input_tokens
358
+ ]
359
+ block_idx_last_computed_token[num_decodes:] = 0
360
361
attn_metadata = Mamba2AttentionMetadata(
362
num_prefills=num_prefills,
@@ -377,10 +375,9 @@ class Mamba2AttentionMetadataBuilder(
375
nums_dict=nums_dict,
376
batch_ptr=batch_ptr,
377
token_chunk_offset_ptr=token_chunk_offset_ptr,
380
- current_last_idx=current_last_idx,
381
- current_first_idx_p=current_first_idx_p,
382
- last_state_idx=last_state_idx,
383
- context_lens_p=context_lens_p,
384
- last_computed_offset_p=last_computed_offset_p,
378
+ block_idx_last_scheduled_token=block_idx_last_scheduled_token,
379
+ block_idx_first_scheduled_token_p=block_idx_first_scheduled_token_p,
380
+ block_idx_last_computed_token=block_idx_last_computed_token,
381
+ num_computed_tokens_p=num_computed_tokens_p,
382
)
383
return attn_metadata
vllm/v1/core/single_type_kv_cache_manager.py
+1
-2
@@ -584,8 +584,7 @@ class MambaManager(SingleTypeKVCacheManager):
584
# hit_length = len(hit_blocks_other_attn[0])
585
# * self.other_block_size
586
# so we insert dummy blocks at the beginning:
587
- if i > 0:
588
- computed.extend([block_pool.null_block] * i)
587
+ computed.extend([block_pool.null_block] * i)
588
computed.append(cached)
589
break # we just need the last match - early stopping
590