relu^2 keeps a NaN: grouped experts, ck experts and Nemotron's pocket

CHANGELOG.md

@@ -4,6 +4,8 @@ Each release's page on GitHub has its notes and measurements. See [the 1.0.0 rel
 
 ## Unreleased
 
+- Nemotron on CUDA: squared ReLU in the grouped experts (MLX affine, GGUF and NVFP4), the NVFP4 `ck_experts` path and the pocket now keeps a NaN instead of turning it into 0, and an NVFP4 activation block holding a NaN gets a NaN scale instead of a finite one, so a NaN reaches the output rather than vanishing (#1628, step 1 of #1549). Finite values give the same bits as before: only the squared-ReLU functions' SASS changes, every Qwen and Gemma kernel compiles as before, and on a GB10 replies were token-identical on Nemotron 3.5 Lightning (NVFP4 and MLX 4-bit), Bolt and Super VL, greedy, sampled and with an image. Thanks to @BobClawblaw.
+
 - HIP: on gfx11 GPUs the prompt GEMM's matrix-core tile now runs its blocks in bands of four row blocks, as CUDA's tile order does, so neighbouring row blocks read each column block's weights while they are still in cache. Each tile computes the same sums in the same order, so replies are token-identical, and a new kernel test holds the tile's bytes to the dot2 tile's on the 27B's prompt shapes and the band edges. On a Ryzen AI Max+ 395 (gfx1151) with Qwen3.8-27B, a cold prompt takes 7.60 s instead of 8.64 s at 2k tokens (-12.0%), 31.4 s instead of 35.5 s at 8k, and 330.8 s instead of 365.7 s at 64k (-9.5%); decode is unchanged, and other GPUs never run the new order (#1648). Thanks to @AbrahamCete.
 
 - GLM on CUDA loads faster: 8 threads decode and pack each layer's routed experts instead of the loader's one (`TF_GLM_LOAD_PACK`, 1 to 16), the packed experts go up a piece at a time (at most 192 MiB) from pinned memory the GPU copies from, and 16 readers now read ahead of the load instead of 8. GLM-5.3-Flash XSFP4 loads in about 51 s instead of 155 s on two GB10s (TP=2) and 46 s instead of 124 s on three (TP=3), with the same replies and memory free at serving within about 1 GiB (#1636). Thanks to @Bizuayeu.

zig/kernels/cuda/ck_experts.cu

@@ -12,6 +12,12 @@ namespace tf_ck4 {
 
 __device__ __forceinline__ float bf(float v) { return __bfloat162float(__float2bfloat16_rn(v)); }
 
+// v < 0 ? 0 : v, not fmaxf, which would make a NaN 0.
+__device__ __forceinline__ float relu(float v) { return v < 0.0f ? 0.0f : v; }
+
+// The larger, or a NaN when either is (fmaxf drops it): a block holding a NaN gets a NaN scale.
+__device__ __forceinline__ float max_nan(float a, float b) { return (a > b || a != a) ? a : b; }
+
 __device__ __forceinline__ void unpack8(const uint4 u, float (&f)[8]) {
     const __nv_bfloat162* p = reinterpret_cast<const __nv_bfloat162*>(&u);
 #pragma unroll
@@ -37,7 +43,7 @@ extern "C" __global__ void __launch_bounds__(128)
     tf_ck4::unpack8(src[1], *reinterpret_cast<float(*)[8]>(f + 8));
     float amax = 0.0f;
 #pragma unroll
-    for (int i = 0; i < 16; ++i) amax = fmaxf(amax, fabsf(f[i]));
+    for (int i = 0; i < 16; ++i) amax = tf_ck4::max_nan(amax, fabsf(f[i]));
     const nvfp4q::Scale sc = nvfp4q::block_scale(amax, g);
     uint32_t w[2] = {0u, 0u};
 #pragma unroll
@@ -146,7 +152,7 @@ __device__ __forceinline__ void grouped(const uint8_t* __restrict__ codes, const
                 if constexpr (EPI == 0) {
                     *reinterpret_cast<float2*>(reinterpret_cast<float*>(out) + row + col) = make_float2(a0, a1);
                 } else {
-                    const float u0 = fmaxf(tf_ck4::bf(a0), 0.0f), u1 = fmaxf(tf_ck4::bf(a1), 0.0f);
+                    const float u0 = tf_ck4::relu(tf_ck4::bf(a0)), u1 = tf_ck4::relu(tf_ck4::bf(a1));
                     *reinterpret_cast<__nv_bfloat162*>(reinterpret_cast<__nv_bfloat16*>(out) + row + col) =
                         __floats2bfloat162_rn(u0 * u0, u1 * u1);
                 }
@@ -296,12 +302,12 @@ __device__ __forceinline__ void grouped_prompt(const uint8_t* __restrict__ codes
                             float f[4], amax = 0.0f;
 #pragma unroll
                             for (int c = 0; c < 4; ++c) {
-                                const float u = fmaxf(tf_ck4::bf(sum[i][2 * bk + c / 2][2 * h + c % 2]), 0.0f);
+                                const float u = tf_ck4::relu(tf_ck4::bf(sum[i][2 * bk + c / 2][2 * h + c % 2]));
                                 f[c] = tf_ck4::bf(u * u);
-                                amax = fmaxf(amax, fabsf(f[c]));
+                                amax = tf_ck4::max_nan(amax, fabsf(f[c]));
                             }
-                            amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 1));
-                            amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, 2));
+                            amax = tf_ck4::max_nan(amax, __shfl_xor_sync(0xffffffffu, amax, 1));
+                            amax = tf_ck4::max_nan(amax, __shfl_xor_sync(0xffffffffu, amax, 2));
                             const nvfp4q::Scale sc = nvfp4q::block_scale(amax, g_out);
                             // lane t holds inputs 2t, 2t + 1 of each 8: byte t of the block's two code words
                             uint32_t lo = (nvfp4q::e2m1(f[0] * sc.mul) | nvfp4q::e2m1(f[1] * sc.mul) << 4) << (8 * t);

zig/kernels/cuda/experts.cuh

@@ -61,8 +61,9 @@ __device__ __forceinline__ float swiglu(float g, float u, float limit) {
   return bf(gv / (1.f + expf(-gv))) * uv;
 }
 
+// a < 0 ? 0 : a, not fmaxf, which would make a NaN sum 0.
 __device__ __forceinline__ float relu2(float a) {
-  const float u = fmaxf(bf(a), 0.f);
+  const float b = bf(a), u = b < 0.f ? 0.f : b;
   return u * u;
 }
 

zig/kernels/cuda/nemotron_ops.cu

@@ -139,7 +139,7 @@ extern "C" __global__ void __launch_bounds__(256) tf_pocket_rows(const __nv_bflo
 #pragma unroll
         for (int o = 16; o > 0; o >>= 1) s += __shfl_xor_sync(0xffffffffu, s, o);
         if (lane == 0) {
-            const float z = fmaxf(s, 0.f);
+            const float z = s < 0.f ? 0.f : s;  // not fmaxf, which would make a NaN 0
             act[j] = z * z;
         }
     }
@@ -194,7 +194,7 @@ __device__ __forceinline__ void pocket_in(const __nv_bfloat16* __restrict__ y, i
 #pragma unroll
             for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o);
             if (lane == 0 && j0 + q < n && t < tile) {
-                const float z = fmaxf(v, 0.f);
+                const float z = v < 0.f ? 0.f : v;  // not fmaxf, which would make a NaN 0
                 act[static_cast<size_t>(r0 + t) * pocket_width + j0 + q] = z * z;
             }
         }