q8_1 activations: keep the block's fp16 scale and sum finite (#606)

Both q8_1 activation quantizers (iq_kernels.cu q8_1_store, native_mmvq.cu native_quantize_q8_1_kernel) stored the
sum of a block's 32 raw activations as fp16: past 65504 (one massive-activation dimension is enough) it rounded to
inf, and d = amax / 127 did past amax = 8.3M. The dot products then read inf * 0 = NaN, and a NaN in the residual
answers one token forever (ggml-org/llama.cpp#23606 is the same defect).

Now both are clamped to the largest finite half (and the values to the int8 range they then exceed). A block that
was finite before is stored bit for bit as before - values in (65504, 65520) round to 65504 anyway, and a finite
block's |x / d| is at most 127 - so the default output only changes where a block overflowed (and was NaN). A NaN
input stays NaN, so STRATA_DBG_NAN still finds it.

iq_multi_parity checks both quantizers against the 0.1.38 formula on a host transcription: every finite block as
before, and three overflowing blocks (sum 96,000, sum -69,984, amax 2e7) stored finite.

include/strata/kernels/q8_1_finite.hpp

@@ -0,0 +1,33 @@
+// include/strata/kernels/q8_1_finite.hpp - #606: the q8_1 activation block's scale and sum, kept finite.
+//
+// A q8_1 block stores its scale d = amax / 127 and the sum of its 32 raw activations as fp16. A sum past 65504 (one
+// massive-activation dimension among 32 values is enough) rounds to inf, and d does past amax = 8.3M; the dot
+// products then read inf * 0 = NaN, and a NaN in the residual ends as one token repeated (ggml-org/llama.cpp#23606
+// is the same defect). Clamped to the largest finite half, such a block gives an approximate product instead of NaN,
+// and the values are held to the int8 range they then exceed. Every block that was finite before is stored bit for
+// bit as before: values in (65504, 65520) round to 65504 anyway, and a finite block's |x / d| is at most 127. A NaN
+// stays NaN (the comparisons are false for it), so STRATA_DBG_NAN still sees one.
+#pragma once
+
+#include <cuda_fp16.h>
+
+namespace strata::kernels {
+
+constexpr float kHalfMax = 65504.0f;
+
+__device__ __forceinline__ float q8_1_finite(const float v) {
+    return fabsf(v) > kHalfMax ? copysignf(kHalfMax, v) : v;
+}
+
+/// The block's quantized value of `xi` for the (finite) scale `d`; `amax` = 0 is an all-zero block.
+__device__ __forceinline__ int8_t q8_1_quant(const float xi, const float d, const float amax) {
+    if (amax == 0.0f) return 0;
+    const float q = roundf(xi / d);
+    return (int8_t) (q > 127.0f ? 127.0f : q < -127.0f ? -127.0f : q);
+}
+
+__device__ __forceinline__ half2 q8_1_ds(const float d, const float sum) {
+    return make_half2(__float2half(q8_1_finite(d)), __float2half(q8_1_finite(sum)));
+}
+
+}  // namespace strata::kernels

src/kernels/cuda/iq_kernels.cu

@@ -6,6 +6,7 @@
 // included unchanged.
 #include "strata/kernels/iq_kernels.hpp"
 #include "strata/kernels/dp4a.hpp"
+#include "strata/kernels/q8_1_finite.hpp"
 
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
@@ -1026,11 +1027,11 @@ __device__ __forceinline__ void q8_1_store(const float xi, block_q8_1* __restric
         amax = fmaxf(amax, __shfl_xor_sync(0xffffffffu, amax, o));
         sum += __shfl_xor_sync(0xffffffffu, sum, o);
     }
-    const float d = amax / 127.0f;
-    const int8_t q = amax == 0.0f ? 0 : roundf(xi / d);
+    const float d = q8_1_finite(amax / 127.0f);   // #606: q8_1_finite.hpp - the same bits for every finite block
+    const int8_t q = q8_1_quant(xi, d, amax);
     const long long ib = i / 32, iqs = i % 32;
     y[ib].qs[iqs] = q;
-    if (iqs == 0) y[ib].ds = make_half2(d, sum);
+    if (iqs == 0) y[ib].ds = q8_1_ds(d, sum);
 }
 
 __global__ void quantize_q8_1_kernel(const float* __restrict__ x, block_q8_1* __restrict__ y, long long n) {

src/kernels/cuda/native_mmvq.cu

@@ -25,6 +25,7 @@
 
 #include "strata/kernels/native_mmvq.hpp"
 #include "strata/kernels/dp4a.hpp"
+#include "strata/kernels/q8_1_finite.hpp"
 #include "strata/kernels/iq_kernels.hpp"
 
 #include <cuda_fp16.h>
@@ -145,10 +146,10 @@ __global__ void native_quantize_q8_1_kernel(const float* __restrict__ x,
     const float xi = x[i];
     const float amax = warp_max(fabsf(xi));
     const float sum = warp_sum(xi);
-    const float d = amax / 127.0f;
-    const int8_t q = amax == 0.0f ? 0 : roundf(xi / d);
+    const float d = q8_1_finite(amax / 127.0f);   // #606: q8_1_finite.hpp - the same bits for every finite block
+    const int8_t q = q8_1_quant(xi, d, amax);
     y[i / Q8K].qs[i % Q8K] = q;
-    if (i % Q8K == 0) y[i / Q8K].ds = make_half2(d, sum);
+    if (i % Q8K == 0) y[i / Q8K].ds = q8_1_ds(d, sum);
 }
 
 // Exact pinned vec_dot_q5_K_q8_1_impl_vmmq expression and integer dot order.

src/kernels/iq_multi_parity.cpp

@@ -16,6 +16,7 @@
 // are set, small enough that the grouped path's SwiGLU output keeps a finite fp16 q8_1 scale.
 #include "strata/kernels/f16_bits.hpp"
 #include "strata/kernels/iq_kernels.hpp"
+#include "strata/kernels/native_mmvq.hpp"
 
 #include <cuda_runtime.h>
 
@@ -325,6 +326,71 @@ void bench(cudaStream_t s, std::mt19937& rng) {
     }
 }
 
+// #606: both q8_1 activation quantizers (quantize_q8_1_rows and native_quantize_q8_1) against a host transcription
+// of the 0.1.38 formula: every finite block bit for bit; a block whose sum (or scale) overflowed fp16 is now stored
+// finite (the largest half, its sign), with its values in the int8 range.
+void check_q8_1_finite(cudaStream_t s, std::mt19937& rng) {
+    const int n = 2560, rows = 3, nb = n / 32;
+    std::vector<float> x = random_x((size_t) rows * n, rng);
+    for (int i = 0; i < 32; ++i) x[(size_t) 5 * 32 + i] = 3000.0f;                  // sum 96,000: overflowed
+    for (int i = 0; i < 32; ++i) x[(size_t) 9 * 32 + i] = i == 3 ? -70000.0f : 0.5f;   // sum -69,984.5
+    for (int i = 0; i < 32; ++i) x[(size_t) 13 * 32 + i] = 2000.0f + i;            // sum 64,496: finite, kept
+    x[(size_t) n + 40] = 2.0e7f;                                                    // amax / 127 overflowed too
+    float* dx = dalloc<float>(x.size());
+    ck(cudaMemcpy(dx, x.data(), x.size() * 4, cudaMemcpyHostToDevice), "x");
+    uint8_t* dq = dalloc<uint8_t>((size_t) rows * nb * 36);
+    for (int path = 0; path < 2; ++path) {
+        if (path == 0) k::quantize_q8_1_rows(dx, rows, n, dq, s);
+        else k::native_quantize_q8_1(dx, dq, n, rows, s);
+        ck(cudaStreamSynchronize(s), "quantize");
+        std::vector<uint8_t> got((size_t) rows * nb * 36);
+        ck(cudaMemcpy(got.data(), dq, got.size(), cudaMemcpyDeviceToHost), "q");
+        int same = 0, clamped = 0, bad = 0;
+        for (int b = 0; b < rows * nb; ++b) {
+            const float* v = x.data() + (size_t) b * 32;
+            float lane[32], amax = 0.0f;
+            for (int i = 0; i < 32; ++i) { lane[i] = v[i]; amax = std::max(amax, std::fabs(v[i])); }
+            for (int o = 16; o > 0; o >>= 1) {   // the warp's butterfly: every lane ends with the same sum
+                float nx[32];
+                for (int i = 0; i < 32; ++i) nx[i] = lane[i] + lane[i ^ o];
+                std::memcpy(lane, nx, sizeof lane);
+            }
+            const float d = amax / 127.0f, sum = lane[0];
+            const uint8_t* g = got.data() + (size_t) b * 36;
+            uint16_t gd, gs;
+            std::memcpy(&gd, g, 2);
+            std::memcpy(&gs, g + 2, 2);
+            const uint16_t od = k::f16_from_f32(d), os = k::f16_from_f32(sum);
+            const bool finite_before = (od & 0x7c00) != 0x7c00 && (os & 0x7c00) != 0x7c00;
+            if (finite_before) {
+                uint8_t want[36];
+                std::memcpy(want, &od, 2);
+                std::memcpy(want + 2, &os, 2);
+                for (int i = 0; i < 32; ++i)   // CUDA's roundf: half away from zero, as std::round
+                    want[4 + i] = (uint8_t) (int8_t) (amax == 0.0f ? 0.0f : std::round(v[i] / d));
+                // native_mmvq.cu is built with --use_fast_math (its divisions are approximate): there the sum must be
+                // the same bits, the scale within one fp16 step and each value within one step
+                bool near = path == 1 && gs == os && std::abs((int) gd - (int) od) <= 1;
+                for (int i = 0; near && i < 32; ++i) near = std::abs((int) (int8_t) g[4 + i] - (int) (int8_t) want[4 + i]) <= 1;
+                if (std::memcmp(want, g, 36) == 0 || near) ++same;
+                else { ++bad; std::printf("  q8_1 path %d block %d: changed although it was finite\n", path, b); }
+            } else {
+                const float fd = k::f32_from_f16(gd), fs = k::f32_from_f16(gs);
+                bool ok = std::isfinite(fd) && std::isfinite(fs) && fd > 0.0f &&
+                          ((os & 0x7c00) != 0x7c00 ? gs == os : std::fabs(fs) == 65504.0f && (fs < 0) == (sum < 0));
+                for (int i = 0; i < 32; ++i) ok = ok && (int8_t) g[4 + i] >= -127 && (int8_t) g[4 + i] <= 127;
+                if (ok) ++clamped;
+                else { ++bad; std::printf("  q8_1 path %d block %d: d %g sum %g not finite and clamped\n", path, b, fd, fs); }
+            }
+        }
+        std::printf("q8_1 finite (%s): %d finite blocks as before, %d overflowed blocks clamped finite%s\n",
+                    path == 0 ? "quantize_q8_1_rows" : "native_quantize_q8_1", same, clamped, bad ? "  FAIL" : "");
+        g_fail += bad + (clamped != 3);
+    }
+    cudaFree(dx);
+    cudaFree(dq);
+}
+
 }  // namespace
 
 int main(int argc, char** argv) {
@@ -340,6 +406,7 @@ int main(int argc, char** argv) {
         for (int dt : {20, 42}) check_grouped(gu, dt, 2560, 640, s, rng);   // the model's shape
         check_grouped(gu, 23, 1024, 512, s, rng);                           // IQ4_XS down needs n_ff % 256 == 0
     }
+    check_q8_1_finite(s, rng);
     if (do_bench) bench(s, rng);
     std::printf("iq_multi_parity: %d failures\n", g_fail);
     cudaStreamDestroy(s);