[Perf] Batch Mamba2 prefill SSM state saves, removing GPU<->CPU syncs (#49371)
Signed-off-by: samuelkim7 <samuelmwkim@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
Samuel Kim committed
Sep 25, 2026 at 03:11 UTC
34ee85d5ca0728a36f687f5f0c7eb3531c2cdd12
4 files changed
+113
-62
tests/v1/attention/test_mamba_update_block_table.py
+37
@@ -427,3 +427,40 @@ def test_block_idx_prev_step_cudagraph_capture_uses_persistent_buffer():
427
428
# Tail values past num_decodes: zero-filled padding for cudagraph capture.
429
assert torch.all(out.block_idx_last_scheduled_token_prev_step[num_decodes:] == 0)
430
+
431
+
432
+def test_num_state_writes_p_matches_block_idx_tensors():
433
+ """The prefill state save sizes its batched copy from
434
+ num_state_writes_p instead of reading the block index tensors back from
435
+ the GPU, so the host count must equal what those tensors say."""
436
+ # (block_size, seq_lens, query_lens, is_prefilling)
437
+ cases = [
438
+ # block-aligned and unaligned computed tokens, and a chunked prefill
439
+ (16, [64, 33, 17], [48, 17, 1], [True, True, True]),
440
+ (32, [200, 100, 64], [200, 36, 64], [True, True, True]),
441
+ # decode rows ahead of the prefills
442
+ (16, [40, 41, 100, 200], [1, 1, 60, 200], [False, False, True, True]),
443
+ # prefills that stay inside their first block: nothing to write
444
+ (16, [65, 66, 67], [1, 2, 3], [True, True, True]),
445
+ ]
446
+
447
+ for block_size, seq_lens, query_lens, is_prefilling in cases:
448
+ metadata = MockMambaBuilder.build_mamba_metadata(
449
+ _make_vllm_config(4096, len(seq_lens), block_size=block_size),
450
+ seq_lens=seq_lens,
451
+ query_lens=query_lens,
452
+ is_prefilling=is_prefilling,
453
+ )
454
+ num_reqs, num_prefills = metadata.num_reqs, metadata.num_prefills
455
+ block_idx_last_p = metadata.block_idx_last_scheduled_token[
456
+ num_reqs - num_prefills : num_reqs
457
+ ]
458
+ expected = int(
459
+ (block_idx_last_p - metadata.block_idx_first_scheduled_token_p)
460
+ .clamp(min=0)
461
+ .sum()
462
+ )
463
+
464
+ assert metadata.num_state_writes_p == expected, (
465
+ f"seq_lens={seq_lens} query_lens={query_lens} block_size={block_size}"
466
+ )
vllm/model_executor/layers/mamba/mamba_mixer2.py
+42
-56
@@ -56,7 +56,6 @@ from vllm.model_executor.model_loader.weight_utils import (
56
from vllm.model_executor.parameter import BasevLLMParameter
57
from vllm.model_executor.utils import set_weight_attrs
58
from vllm.platforms import current_platform
59
-from vllm.utils.gpu_sync_debug import gpu_sync_allowed
59
from vllm.utils.torch_utils import (
60
LayerNameType,
61
_encode_layer_name,
@@ -927,68 +926,55 @@ class MambaMixer2(MambaBase, PluggableLayer):
926
# then chunk_stride = 2
927
chunk_stride = mamba_block_size // chunk_size
928
930
- # The per-sequence loop below uses these as Python scalars.
931
- # TODO avoid sync here?
932
- with gpu_sync_allowed():
933
- block_idx_first_cpu = block_idx_first_scheduled_token_p.tolist()
934
- block_idx_last_cpu = block_idx_last_scheduled_token_p.tolist()
935
- num_computed_tokens_p_cpu = num_computed_tokens_p.tolist()
936
- last_chunk_indices_p_cpu = last_chunk_indices_p.tolist()
937
-
938
- # Save state for sequences with more than just final state
939
- for seq_idx in range(num_prefills):
940
- # Block index for the first scheduled token
941
- block_idx_first_scheduled_token = block_idx_first_cpu[seq_idx]
942
-
943
- # Block index for the last scheduled token
944
- block_idx_last_scheduled_token = block_idx_last_cpu[seq_idx]
945
-
946
- # Number of blocks that need to be written
947
- n_blocks_to_fill = (
948
- block_idx_last_scheduled_token - block_idx_first_scheduled_token
929
+ # Save state for sequences with more than just final state,
930
+ # batched over sequences so that the number of kernel launches
931
+ # does not scale with the number of prefills.
932
+ n_blocks_to_fill = (
933
+ block_idx_last_scheduled_token_p - block_idx_first_scheduled_token_p
934
+ ).clamp(min=0)
935
+ total_fill = attn_metadata.num_state_writes_p
936
+ if total_fill > 0:
937
+ first_chunk = torch.zeros_like(last_chunk_indices_p)
938
+ first_chunk[1:] = last_chunk_indices_p[:-1] + 1
939
+
940
+ # The chunk that completes this sequence's first mamba block.
941
+ # A block spans chunk_stride chunks. When the sequence resumed
942
+ # mid-block, the chunks before this step are not in varlen_states,
943
+ # so subtract them: the boundary sits earlier by their count.
944
+ num_unaligned_computed_tokens = (
945
+ num_computed_tokens_p % mamba_block_size
946
+ )
947
+ unaligned_chunk_shift = num_unaligned_computed_tokens // chunk_size
948
+ first_aligned_chunk = (
949
+ first_chunk + (chunk_stride - 1) - unaligned_chunk_shift
950
)
951
951
- # Skip sequences that don't have any blocks to fill
952
- if n_blocks_to_fill == 0:
953
- continue
952
+ # Flatten all writes into one list: fill_counts=[2,1,2]
953
+ # -> seq_ids=[0,0,1,2,2], write_offsets=[0,1,0,0,1]
954
+ fill_counts = n_blocks_to_fill.long()
955
+ # output_size, or repeat_interleave reads the size off the device
956
+ seq_ids = torch.repeat_interleave(
957
+ torch.arange(num_prefills, device=ssm_state.device),
958
+ fill_counts,
959
+ output_size=total_fill,
960
+ )
961
+ write_starts = torch.cumsum(fill_counts, dim=0) - fill_counts
962
+ write_offsets = torch.arange(
963
+ total_fill, device=ssm_state.device
964
+ ) - torch.repeat_interleave(
965
+ write_starts, fill_counts, output_size=total_fill
966
+ )
967
955
- # Look up the state indices
968
cache_blocks_to_fill = state_indices_tensor_p[
957
- seq_idx,
958
- block_idx_first_scheduled_token:block_idx_last_scheduled_token,
969
+ seq_ids,
970
+ block_idx_first_scheduled_token_p[seq_ids] + write_offsets,
971
]
960
-
961
- # First chunk index for this sequence
962
- if seq_idx == 0:
963
- first_chunk = 0
964
- else:
965
- first_chunk = 1 + last_chunk_indices_p_cpu[seq_idx - 1]
966
-
967
- # First chunk that is aligned on the mamba block boundary
968
- first_aligned_chunk = first_chunk + chunk_stride - 1
969
-
970
- # Calculate the number of computed tokens that were not
971
- # already cached
972
- num_unaligned_computed_tokens = (
973
- num_computed_tokens_p_cpu[seq_idx] % mamba_block_size
974
- )
975
-
976
- if num_unaligned_computed_tokens > 0:
977
- # If the number of computed tokens is not block aligned,
978
- # then we need to shift the index accordingly
979
- first_aligned_chunk -= (
980
- num_unaligned_computed_tokens // chunk_size
981
- )
982
-
983
- # Get states to write
984
- from_where = varlen_states[
985
- first_aligned_chunk : first_aligned_chunk
986
- + n_blocks_to_fill * chunk_stride : chunk_stride
972
+ # Write the states. Each row goes to its own cache block,
973
+ # so the indexed copy has no duplicate destinations.
974
+ ssm_state[cache_blocks_to_fill] = varlen_states[
975
+ first_aligned_chunk[seq_ids] + write_offsets * chunk_stride
976
]
977
989
- # Write the states
990
- ssm_state[cache_blocks_to_fill] = from_where
991
-
978
# For all seqs, store the last state (note: might be partial):
979
assert state_indices_tensor_p is not None
980
ssm_state[
vllm/v1/attention/backends/mamba2_attn.py
+4
-1
@@ -156,7 +156,10 @@ class Mamba2AttentionMetadataBuilder(
156
# data so it needs no D2H. `seq_lens_cpu_upper_bound` is precise
157
# for prefill rows, which is all this slice covers.
158
num_computed_tokens_p_cpu, _ = self._prefill_cpu_metadata(
159
- common, common_attn_metadata
159
+ common_attn_metadata,
160
+ common.num_reqs,
161
+ common.num_prefills,
162
+ common.num_decode_tokens,
163
)
164
prep_initial_states = bool((num_computed_tokens_p_cpu > 0).any())
165
vllm/v1/attention/backends/mamba_attn.py
+30
-5
@@ -70,6 +70,9 @@ class BaseMambaAttentionMetadata:
70
# last_chunk_indices_p is a tensor of shape (batch,) that contains the
71
# index of the last chunk for every sequence in the (prefill) batch.
72
last_chunk_indices_p: torch.Tensor | None = None
73
+ # Number of block-aligned SSM states the prefill path writes, i.e.
74
+ # sum(block_idx_last_scheduled_token_p - block_idx_first_scheduled_token_p).
75
+ num_state_writes_p: int = 0
76
77
# The following attributes are for triton implementation of causal_conv1d
78
nums_dict: dict | None = None
@@ -348,8 +351,10 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
351
352
def _prefill_cpu_metadata(
353
self,
351
- common: M,
354
common_attn_metadata: CommonAttentionMetadata,
355
+ num_reqs: int,
356
+ num_prefills: int,
357
+ num_decode_tokens: int,
358
) -> tuple[torch.Tensor, torch.Tensor]:
359
"""Prefill context lengths and query offsets, from CPU data only.
360
@@ -361,11 +366,9 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
366
"""
367
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
368
assert seq_lens_cpu is not None
364
- num_reqs = common.num_reqs
365
- num_prefills = common.num_prefills
369
query_start_loc_p_cpu = (
370
common_attn_metadata.query_start_loc_cpu[-num_prefills - 1 :]
368
- - common.num_decode_tokens
371
+ - num_decode_tokens
372
)
373
prefill_query_lens_cpu = query_start_loc_p_cpu[1:] - query_start_loc_p_cpu[:-1]
374
num_computed_tokens_p_cpu = (
@@ -385,7 +388,10 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
388
num_prefills = common.num_prefills
389
390
num_computed_tokens_p_cpu, query_start_loc_p_cpu = self._prefill_cpu_metadata(
388
- common, common_attn_metadata
391
+ common_attn_metadata,
392
+ common.num_reqs,
393
+ num_prefills,
394
+ common.num_decode_tokens,
395
)
396
397
cu_chunk_seqlen, seq_idx, last_chunk_indices = self._compute_chunk_metadata(
@@ -533,6 +539,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
539
block_idx_last_computed_token = None
540
block_idx_last_scheduled_token = None
541
block_idx_last_scheduled_token_prev_step = None
542
+ num_state_writes_p = 0
543
544
# for causal_conv1d
545
nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None
@@ -625,6 +632,23 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
632
num_reqs - num_prefills : num_reqs
633
]
634
635
+ # How many block-aligned states prefill will write: the same
636
+ # (block_idx_last - block_idx_first) as above, from CPU data so
637
+ # the mixer needs no D2H read. The -1 in both indices cancels.
638
+ block_size = self.kv_cache_spec.block_size
639
+ seq_lens_p_cpu = seq_lens_cpu[num_reqs - num_prefills : num_reqs]
640
+ num_computed_tokens_p_cpu, _ = self._prefill_cpu_metadata(
641
+ common_attn_metadata, num_reqs, num_prefills, num_decode_tokens
642
+ )
643
+ num_state_writes_p = int(
644
+ (
645
+ (seq_lens_p_cpu + block_size - 1) // block_size
646
+ - (num_computed_tokens_p_cpu + block_size) // block_size
647
+ )
648
+ .clamp(min=0)
649
+ .sum()
650
+ )
651
+
652
if self.use_replayssm and not self.use_flashinfer_replayssm and num_decodes > 0:
653
decode_base_cpu = common_attn_metadata.replayssm_decode_base_cpu
654
seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound
@@ -728,6 +752,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
752
query_start_loc_d=query_start_loc_d,
753
block_idx_last_scheduled_token=block_idx_last_scheduled_token,
754
block_idx_first_scheduled_token_p=block_idx_first_scheduled_token_p,
755
+ num_state_writes_p=num_state_writes_p,
756
block_idx_last_computed_token=block_idx_last_computed_token,
757
block_idx_last_scheduled_token_prev_step=(
758
block_idx_last_scheduled_token_prev_step