From 58dd01f5db189a18d220b2286f15c9c3b74b7156 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Stanis=C5=82aw=20Szymczyk?= Date: Sat, 22 Aug 2026 13:23:58 +0200 Subject: [PATCH] webgpu : fix handling of infinity values during ARGSORT and TOP_K --- 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 46ed19fc7752..fa5d953572e9 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;