From 5ea87ddad22541a37053c7ba92b02ec1923617c6 Mon Sep 17 00:00:00 2001 From: fairydreaming <166155368+fairydreaming@users.noreply.github.com> Date: Tue, 25 Aug 2026 07:08:06 +0200 Subject: [PATCH] webgpu : fix handling of infinity values during ARGSORT and TOP_K (#27538) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Stanisław Szymczyk --- ggml/src/ggml-webgpu/wgsl-shaders/argsort.wgsl | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) diff --git a/ggml/src/ggml-webgpu/wgsl-shaders/argsort.wgsl b/ggml/src/ggml-webgpu/wgsl-shaders/argsort.wgsl index 46ed19fc77..fa5d953572 100644 --- a/ggml/src/ggml-webgpu/wgsl-shaders/argsort.wgsl +++ b/ggml/src/ggml-webgpu/wgsl-shaders/argsort.wgsl @@ -34,11 +34,9 @@ var params: Params; var shmem_idx: array; #if ORDER == 0 -#define EXTREME_VALUE 1e30 #define SWAP_COMPARE_UP > #define SWAP_COMPARE_DOWN < #else -#define EXTREME_VALUE -1e30 #define SWAP_COMPARE_UP < #define SWAP_COMPARE_DOWN > #endif @@ -78,11 +76,9 @@ fn main(@builtin(workgroup_id) wid: vec3, let dir_up = (lid.x & k) == 0; let a_idx = shmem_idx[lid.x]; let b_idx = shmem_idx[ixj]; - let a_val = select(EXTREME_VALUE, src[row_base + a_idx], a_idx < params.src_ne0); - let b_val = select(EXTREME_VALUE, src[row_base + b_idx], b_idx < params.src_ne0); let should_swap = select( - (a_val SWAP_COMPARE_DOWN b_val), - (a_val SWAP_COMPARE_UP b_val), + b_idx >= params.src_ne0 || (a_idx < params.src_ne0 && src[row_base + a_idx] SWAP_COMPARE_DOWN src[row_base + b_idx]), + a_idx >= params.src_ne0 || (b_idx < params.src_ne0 && src[row_base + a_idx] SWAP_COMPARE_UP src[row_base + b_idx]), dir_up); if (should_swap) { shmem_idx[lid.x] = b_idx;