From c64861d63c1a5362bfad443daf7a096f1bcfd1e4 Mon Sep 17 00:00:00 2001 From: Chih-Chieh Yang <7364402+cyang49@users.noreply.github.com> Date: Sat, 2 Aug 2025 02:55:57 -0400 Subject: [PATCH] [Bugfix] Mamba2 remove bugged initial state condition in chunk scan (#22034) Signed-off-by: Chih-Chieh-Yang <7364402+cyang49@users.noreply.github.com> --- .../model_executor/layers/mamba/ops/ssd_chunk_scan.py | 11 ++--------- 1 file changed, 2 insertions(+), 9 deletions(-) diff --git a/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py b/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py index 365e1c54b555a..61eff0c008f60 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py @@ -476,15 +476,8 @@ def _chunk_scan_fwd( # with initial states, we need to take care of how # seq_idx crosses the boundaries assert batch == 1, "chunk scan only supports initial states with batch 1" - - if initial_states.shape[0] == 1: - # no in this case no point to use initial states - initial_states = None - else: - assert chunk_indices is not None and chunk_offsets is not None, \ - ( - "chunk_indices and chunk_offsets should have been set" - ) + assert chunk_indices is not None and chunk_offsets is not None, \ + "chunk_indices and chunk_offsets should have been set" else: chunk_indices, chunk_offsets = None, None else: