ggml-cuda: add mem check for fusion (#19916)

* ggml-cuda: add mem check for fusion

* Replace NaNs with -FLT_MAX

* fix typo

Co-authored-by: Johannes Gäßler <johannesg@5d6.de>

---------

Co-authored-by: Johannes Gäßler <johannesg@5d6.de>
This commit is contained in:
Aman Gupta
2026-03-07 00:05:43 +08:00
committed by GitHub
co-authored by Johannes Gäßler
parent ba2ff79e43
commit d48e876467
2 changed files with 79 additions and 2 deletions
+12
View File
@@ -119,6 +119,18 @@ __launch_bounds__(4 * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float *
}
}
// Sanitize NaN to -FLT_MAX so the iterative argmax produces unique expert IDs.
// NaN comparisons always return false, which would cause the same expert to be
// selected repeatedly. -FLT_MAX compares normally and is still excluded by the
// -INFINITY sentinel used after each selection round.
// More relevant for the cuBLAS path. See https://github.com/ggml-org/llama.cpp/issues/19659
#pragma unroll
for (int i = 0; i < experts_per_thread; i++) {
if (__isnanf(wt[i])) {
wt[i] = -FLT_MAX;
}
}
// selection_wt is only needed when bias is present (selection uses wt + bias)
// when no bias, we use wt directly for both selection and weight values
float selection_wt[has_bias ? experts_per_thread : 1];