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 14, 22:35 UTC
[Bugfix] Mamba2 remove bugged initial state condition in chunk scan (#22034)
Signed-off-by: Chih-Chieh-Yang <7364402+cyang49@users.noreply.github.com>
Chih-Chieh Yang committed
Aug 2, 2025 at 02:55 UTC
c64861d63c1a5362bfad443daf7a096f1bcfd1e4
1 file changed
+2
-9
vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py
+2
-9
@@ -476,15 +476,8 @@ def _chunk_scan_fwd(
476
# with initial states, we need to take care of how
477
# seq_idx crosses the boundaries
478
assert batch == 1, "chunk scan only supports initial states with batch 1"
479
-
480
- if initial_states.shape[0] == 1:
481
- # no in this case no point to use initial states
482
- initial_states = None
483
- else:
484
- assert chunk_indices is not None and chunk_offsets is not None, \
485
- (
486
- "chunk_indices and chunk_offsets should have been set"
487
- )
479
+ assert chunk_indices is not None and chunk_offsets is not None, \
480
+ "chunk_indices and chunk_offsets should have been set"
481
else:
482
chunk_indices, chunk_offsets = None, None
483
else: