[Hybrid] [Kernel] Fix chunk scan kernel when BLOCK_SIZE_DSTATE > 128 (#28295)

tdoublep · web-flow · commit e0c910bb89e4 · 2025-11-14T22:55:42.000Z
Signed-off-by: Thomas Parnell &lt;tpa@zurich.ibm.com&gt;
diff --git 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(

Original file line number	Diff line number	Diff line change
`@@ -245,7 +245,7 @@ def _chunk_scan_fwd_kernel(`
`245`	`245`	`)`
`246`	`246`	`if not HAS_INITSTATES and (seq_idx != seq_idx_prev):`
`247`	`247`	`prev_states = tl.zeros(`
`248`		`- (BLOCK_SIZE_DSTATE, BLOCK_SIZE_K), dtype=C_ptr.dtype.element_ty`
	`248`	`+ (BLOCK_SIZE_K, BLOCK_SIZE_N), dtype=C_ptr.dtype.element_ty`
`249`	`249`	`)`
`250`	`250`	`else:`
`251`	`251`	`prev_states = tl.load(`