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 15, 18:30 UTC
[Model] RE: Mamba2 Prefill Performance Tweaks: Fixing Flurry of Unnecessary Memory Copies (#14857)
Signed-off-by: Chih-Chieh-Yang <7364402+cyang49@users.noreply.github.com>
Chih-Chieh Yang committed
Mar 20, 2025 at 19:21 UTC
296f927f2493908984707354e3cc5d7b2e41650b
1 file changed
+7
-6
vllm/model_executor/layers/mamba/mamba_mixer2.py
+7
-6
@@ -470,10 +470,11 @@ class MambaMixer2(CustomOp):
470
if has_prefill:
471
472
initial_states = None
473
- if has_initial_states is not None and any(has_initial_states):
474
- for idx in mamba_cache_params.state_indices_tensor[
475
- ~has_initial_states]:
476
- mamba_cache_params.ssm_state[idx].zero_()
473
+ if has_initial_states is not None and torch.any(
474
+ has_initial_states):
475
+ zero_init_indices = mamba_cache_params.state_indices_tensor[
476
+ ~has_initial_states]
477
+ mamba_cache_params.ssm_state[zero_init_indices] = 0
478
initial_states = mamba_cache_params.ssm_state[
479
mamba_cache_params.state_indices_tensor]
480
@@ -499,8 +500,8 @@ class MambaMixer2(CustomOp):
500
501
# update ssm states
502
# - varlen state is a (batch, nheads, headdim, dstate) tensor
502
- for i, idx in enumerate(mamba_cache_params.state_indices_tensor):
503
- mamba_cache_params.ssm_state[idx].copy_(varlen_state[i])
503
+ mamba_cache_params.ssm_state[
504
+ mamba_cache_params.state_indices_tensor] = varlen_state
505
506
# - reshape
507
hidden_states = scan_output.view(seq_len, -1)