From 34ee85d5ca0728a36f687f5f0c7eb3531c2cdd12 Mon Sep 17 00:00:00 2001 From: Samuel Kim <65876994+samuelkim7@users.noreply.github.com> Date: Fri, 25 Sep 2026 03:11:19 +0200 Subject: [PATCH] [Perf] Batch Mamba2 prefill SSM state saves, removing GPU<->CPU syncs (#49371) Signed-off-by: samuelkim7 Co-authored-by: Claude --- .../test_mamba_update_block_table.py | 37 +++++++ .../layers/mamba/mamba_mixer2.py | 98 ++++++++----------- vllm/v1/attention/backends/mamba2_attn.py | 5 +- vllm/v1/attention/backends/mamba_attn.py | 35 ++++++- 4 files changed, 113 insertions(+), 62 deletions(-) diff --git a/tests/v1/attention/test_mamba_update_block_table.py b/tests/v1/attention/test_mamba_update_block_table.py index f4451fe7af..49a298fe19 100644 --- a/tests/v1/attention/test_mamba_update_block_table.py +++ b/tests/v1/attention/test_mamba_update_block_table.py @@ -427,3 +427,40 @@ def test_block_idx_prev_step_cudagraph_capture_uses_persistent_buffer(): # Tail values past num_decodes: zero-filled padding for cudagraph capture. assert torch.all(out.block_idx_last_scheduled_token_prev_step[num_decodes:] == 0) + + +def test_num_state_writes_p_matches_block_idx_tensors(): + """The prefill state save sizes its batched copy from + num_state_writes_p instead of reading the block index tensors back from + the GPU, so the host count must equal what those tensors say.""" + # (block_size, seq_lens, query_lens, is_prefilling) + cases = [ + # block-aligned and unaligned computed tokens, and a chunked prefill + (16, [64, 33, 17], [48, 17, 1], [True, True, True]), + (32, [200, 100, 64], [200, 36, 64], [True, True, True]), + # decode rows ahead of the prefills + (16, [40, 41, 100, 200], [1, 1, 60, 200], [False, False, True, True]), + # prefills that stay inside their first block: nothing to write + (16, [65, 66, 67], [1, 2, 3], [True, True, True]), + ] + + for block_size, seq_lens, query_lens, is_prefilling in cases: + metadata = MockMambaBuilder.build_mamba_metadata( + _make_vllm_config(4096, len(seq_lens), block_size=block_size), + seq_lens=seq_lens, + query_lens=query_lens, + is_prefilling=is_prefilling, + ) + num_reqs, num_prefills = metadata.num_reqs, metadata.num_prefills + block_idx_last_p = metadata.block_idx_last_scheduled_token[ + num_reqs - num_prefills : num_reqs + ] + expected = int( + (block_idx_last_p - metadata.block_idx_first_scheduled_token_p) + .clamp(min=0) + .sum() + ) + + assert metadata.num_state_writes_p == expected, ( + f"seq_lens={seq_lens} query_lens={query_lens} block_size={block_size}" + ) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index db85ca0dee..b567d188eb 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -56,7 +56,6 @@ from vllm.model_executor.model_loader.weight_utils import ( from vllm.model_executor.parameter import BasevLLMParameter from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform -from vllm.utils.gpu_sync_debug import gpu_sync_allowed from vllm.utils.torch_utils import ( LayerNameType, _encode_layer_name, @@ -927,67 +926,54 @@ class MambaMixer2(MambaBase, PluggableLayer): # then chunk_stride = 2 chunk_stride = mamba_block_size // chunk_size - # The per-sequence loop below uses these as Python scalars. - # TODO avoid sync here? - with gpu_sync_allowed(): - block_idx_first_cpu = block_idx_first_scheduled_token_p.tolist() - block_idx_last_cpu = block_idx_last_scheduled_token_p.tolist() - num_computed_tokens_p_cpu = num_computed_tokens_p.tolist() - last_chunk_indices_p_cpu = last_chunk_indices_p.tolist() + # Save state for sequences with more than just final state, + # batched over sequences so that the number of kernel launches + # does not scale with the number of prefills. + n_blocks_to_fill = ( + block_idx_last_scheduled_token_p - block_idx_first_scheduled_token_p + ).clamp(min=0) + total_fill = attn_metadata.num_state_writes_p + if total_fill > 0: + first_chunk = torch.zeros_like(last_chunk_indices_p) + first_chunk[1:] = last_chunk_indices_p[:-1] + 1 - # Save state for sequences with more than just final state - for seq_idx in range(num_prefills): - # Block index for the first scheduled token - block_idx_first_scheduled_token = block_idx_first_cpu[seq_idx] - - # Block index for the last scheduled token - block_idx_last_scheduled_token = block_idx_last_cpu[seq_idx] - - # Number of blocks that need to be written - n_blocks_to_fill = ( - block_idx_last_scheduled_token - block_idx_first_scheduled_token - ) - - # Skip sequences that don't have any blocks to fill - if n_blocks_to_fill == 0: - continue - - # Look up the state indices - cache_blocks_to_fill = state_indices_tensor_p[ - seq_idx, - block_idx_first_scheduled_token:block_idx_last_scheduled_token, - ] - - # First chunk index for this sequence - if seq_idx == 0: - first_chunk = 0 - else: - first_chunk = 1 + last_chunk_indices_p_cpu[seq_idx - 1] - - # First chunk that is aligned on the mamba block boundary - first_aligned_chunk = first_chunk + chunk_stride - 1 - - # Calculate the number of computed tokens that were not - # already cached + # The chunk that completes this sequence's first mamba block. + # A block spans chunk_stride chunks. When the sequence resumed + # mid-block, the chunks before this step are not in varlen_states, + # so subtract them: the boundary sits earlier by their count. num_unaligned_computed_tokens = ( - num_computed_tokens_p_cpu[seq_idx] % mamba_block_size + num_computed_tokens_p % mamba_block_size + ) + unaligned_chunk_shift = num_unaligned_computed_tokens // chunk_size + first_aligned_chunk = ( + first_chunk + (chunk_stride - 1) - unaligned_chunk_shift ) - if num_unaligned_computed_tokens > 0: - # If the number of computed tokens is not block aligned, - # then we need to shift the index accordingly - first_aligned_chunk -= ( - num_unaligned_computed_tokens // chunk_size - ) + # Flatten all writes into one list: fill_counts=[2,1,2] + # -> seq_ids=[0,0,1,2,2], write_offsets=[0,1,0,0,1] + fill_counts = n_blocks_to_fill.long() + # output_size, or repeat_interleave reads the size off the device + seq_ids = torch.repeat_interleave( + torch.arange(num_prefills, device=ssm_state.device), + fill_counts, + output_size=total_fill, + ) + write_starts = torch.cumsum(fill_counts, dim=0) - fill_counts + write_offsets = torch.arange( + total_fill, device=ssm_state.device + ) - torch.repeat_interleave( + write_starts, fill_counts, output_size=total_fill + ) - # Get states to write - from_where = varlen_states[ - first_aligned_chunk : first_aligned_chunk - + n_blocks_to_fill * chunk_stride : chunk_stride + cache_blocks_to_fill = state_indices_tensor_p[ + seq_ids, + block_idx_first_scheduled_token_p[seq_ids] + write_offsets, + ] + # Write the states. Each row goes to its own cache block, + # so the indexed copy has no duplicate destinations. + ssm_state[cache_blocks_to_fill] = varlen_states[ + first_aligned_chunk[seq_ids] + write_offsets * chunk_stride ] - - # Write the states - ssm_state[cache_blocks_to_fill] = from_where # For all seqs, store the last state (note: might be partial): assert state_indices_tensor_p is not None diff --git a/vllm/v1/attention/backends/mamba2_attn.py b/vllm/v1/attention/backends/mamba2_attn.py index 0a4584632d..fc71e3485f 100644 --- a/vllm/v1/attention/backends/mamba2_attn.py +++ b/vllm/v1/attention/backends/mamba2_attn.py @@ -156,7 +156,10 @@ class Mamba2AttentionMetadataBuilder( # data so it needs no D2H. `seq_lens_cpu_upper_bound` is precise # for prefill rows, which is all this slice covers. num_computed_tokens_p_cpu, _ = self._prefill_cpu_metadata( - common, common_attn_metadata + common_attn_metadata, + common.num_reqs, + common.num_prefills, + common.num_decode_tokens, ) prep_initial_states = bool((num_computed_tokens_p_cpu > 0).any()) diff --git a/vllm/v1/attention/backends/mamba_attn.py b/vllm/v1/attention/backends/mamba_attn.py index b5f0b9b182..63f98f1786 100644 --- a/vllm/v1/attention/backends/mamba_attn.py +++ b/vllm/v1/attention/backends/mamba_attn.py @@ -70,6 +70,9 @@ class BaseMambaAttentionMetadata: # last_chunk_indices_p is a tensor of shape (batch,) that contains the # index of the last chunk for every sequence in the (prefill) batch. last_chunk_indices_p: torch.Tensor | None = None + # Number of block-aligned SSM states the prefill path writes, i.e. + # sum(block_idx_last_scheduled_token_p - block_idx_first_scheduled_token_p). + num_state_writes_p: int = 0 # The following attributes are for triton implementation of causal_conv1d nums_dict: dict | None = None @@ -348,8 +351,10 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): def _prefill_cpu_metadata( self, - common: M, common_attn_metadata: CommonAttentionMetadata, + num_reqs: int, + num_prefills: int, + num_decode_tokens: int, ) -> tuple[torch.Tensor, torch.Tensor]: """Prefill context lengths and query offsets, from CPU data only. @@ -361,11 +366,9 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): """ seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound assert seq_lens_cpu is not None - num_reqs = common.num_reqs - num_prefills = common.num_prefills query_start_loc_p_cpu = ( common_attn_metadata.query_start_loc_cpu[-num_prefills - 1 :] - - common.num_decode_tokens + - num_decode_tokens ) prefill_query_lens_cpu = query_start_loc_p_cpu[1:] - query_start_loc_p_cpu[:-1] num_computed_tokens_p_cpu = ( @@ -385,7 +388,10 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): num_prefills = common.num_prefills num_computed_tokens_p_cpu, query_start_loc_p_cpu = self._prefill_cpu_metadata( - common, common_attn_metadata + common_attn_metadata, + common.num_reqs, + num_prefills, + common.num_decode_tokens, ) cu_chunk_seqlen, seq_idx, last_chunk_indices = self._compute_chunk_metadata( @@ -533,6 +539,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): block_idx_last_computed_token = None block_idx_last_scheduled_token = None block_idx_last_scheduled_token_prev_step = None + num_state_writes_p = 0 # for causal_conv1d nums_dict, batch_ptr, token_chunk_offset_ptr = None, None, None @@ -625,6 +632,23 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): num_reqs - num_prefills : num_reqs ] + # How many block-aligned states prefill will write: the same + # (block_idx_last - block_idx_first) as above, from CPU data so + # the mixer needs no D2H read. The -1 in both indices cancels. + block_size = self.kv_cache_spec.block_size + seq_lens_p_cpu = seq_lens_cpu[num_reqs - num_prefills : num_reqs] + num_computed_tokens_p_cpu, _ = self._prefill_cpu_metadata( + common_attn_metadata, num_reqs, num_prefills, num_decode_tokens + ) + num_state_writes_p = int( + ( + (seq_lens_p_cpu + block_size - 1) // block_size + - (num_computed_tokens_p_cpu + block_size) // block_size + ) + .clamp(min=0) + .sum() + ) + if self.use_replayssm and not self.use_flashinfer_replayssm and num_decodes > 0: decode_base_cpu = common_attn_metadata.replayssm_decode_base_cpu seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound @@ -728,6 +752,7 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC): query_start_loc_d=query_start_loc_d, block_idx_last_scheduled_token=block_idx_last_scheduled_token, block_idx_first_scheduled_token_p=block_idx_first_scheduled_token_p, + num_state_writes_p=num_state_writes_p, block_idx_last_computed_token=block_idx_last_computed_token, block_idx_last_scheduled_token_prev_step=( block_idx_last_scheduled_token_prev_step