From ca00b039aef885cf4f589d4dccd4028d1fe70196 Mon Sep 17 00:00:00 2001 From: Thomas Parnell Date: Fri, 7 Nov 2025 07:41:27 -0500 Subject: [PATCH] Fix chunk scan kernel when BLOCK_SIZE_DSTATE > 128 Signed-off-by: Thomas Parnell --- vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 e5a5c9dd6f71..661c884627b0 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py @@ -245,7 +245,7 @@ def _chunk_scan_fwd_kernel( ) if not HAS_INITSTATES and (seq_idx != seq_idx_prev): prev_states = tl.zeros( - (BLOCK_SIZE_DSTATE, BLOCK_SIZE_K), dtype=C_ptr.dtype.element_ty + (BLOCK_SIZE_K, BLOCK_SIZE_N), dtype=C_ptr.dtype.element_ty ) else: prev_states = tl.load(