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 5, 18:30 UTC
[Bugfix] clamp dA_cumsum differences to prevent Inf in Mamba2 SSD kernels (#37501)
Signed-off-by: Jingu Kang <jg.k@navercorp.com>
Jingu Kang committed
Apr 1, 2026 at 00:35 UTC
f1ff50c86cfac67b68ddef67336e96a1b6e424b6
2 files changed
+2
-2
vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py
+1
-1
@@ -356,7 +356,7 @@ def _chunk_scan_fwd_kernel(
356
)
357
# If there's seq_idx, we already set cb[i, j] = 0 for seq_idx[i] != seq_idx[j].
358
# So we don't need masking wrt seq_idx here.
359
- cb *= fast_exp(dA_cs_m[:, None] - dA_cs_k[None, :])
359
+ cb *= fast_exp(tl.minimum(dA_cs_m[:, None] - dA_cs_k[None, :], 0.0))
360
dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size - k, other=0.0).to(tl.float32)
361
cb *= dt_k
362
if IS_CAUSAL:
vllm/model_executor/layers/mamba/ops/ssd_chunk_state.py
+1
-1
@@ -280,7 +280,7 @@ def _chunk_state_fwd_kernel(
280
dt_k = tl.load(dt_ptrs, mask=offs_k < chunk_size_limit - k, other=0.0).to(
281
tl.float32
282
)
283
- scale = fast_exp(dA_cs_last - dA_cs_k) * dt_k
283
+ scale = fast_exp(tl.minimum(dA_cs_last - dA_cs_k, 0.0)) * dt_k
284
b *= scale[:, None]
285
b = b.to(x_ptr.dtype.element_ty)
286
acc += tl.dot(x, b)