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 13, 04:30 UTC
[Hybrid] [Kernel] Fix chunk scan kernel when BLOCK_SIZE_DSTATE > 128 (#28295)
Signed-off-by: Thomas Parnell <tpa@zurich.ibm.com>
Thomas Parnell committed
Nov 14, 2025 at 23:55 UTC
e0c910bb89e45f4a2a976dc3c76248bbdea854e0
1 file changed
+1
-1
vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py
+1
-1
@@ -245,7 +245,7 @@ def _chunk_scan_fwd_kernel(
245
)
246
if not HAS_INITSTATES and (seq_idx != seq_idx_prev):
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
)
250
else:
251
prev_states = tl.load(