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, 05:30 UTC
Avoid division by zero in cache DS MLA kernel (#26174)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Oct 3, 2025 at 13:35 UTC
300a59c4c313bcf2e5dd6254a647a10c89a37b1b
1 file changed
+2
-1
csrc/cache_kernels.cu
+2
-1
@@ -16,7 +16,7 @@
16
17
#include <algorithm>
18
#include <cassert>
19
-#include <cfloat> // FLT_MIN
19
+#include <cfloat>
20
21
#ifdef USE_ROCM
22
#include <hip/hip_bf16.h>
@@ -479,6 +479,7 @@ __global__ void concat_and_cache_ds_mla_kernel(
479
480
// Compute the scale for the tile
481
float tile_scale = max_abs / 448.f;
482
+ tile_scale = fmaxf(tile_scale, FLT_MIN);
483
484
// The first lane of each half-warp writes the scale to kv_cache
485
if ((lane_idx == 0) || (lane_idx == 16)) {