iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf] Batch Mamba2 prefill SSM state saves, removing GPU<->CPU syncs (#49371)

Browse Source 
Signed-off-by: samuelkim7 <samuelmwkim@gmail.com>
Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Samuel Kim
2026-09-25 01:11:19 +00:00
committed by GitHub
co-authored by Claude
parent a811738a60
commit 34ee85d5ca
4 changed files with 113 additions and 62 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_mamba_update_block_table.py
+37
View File
@@ -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}"
)
vllm/model_executor/layers/mamba/mamba_mixer2.py
+42 -56
View File
@@ -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
vllm/v1/attention/backends/mamba2_attn.py
+4 -1
View File
@@ -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())
vllm/v1/attention/backends/mamba_attn.py
+30 -5
View File
@@ -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