Correctness fixes from #149 (s_gemv barrier, bf16 NaN, QSA page mask, verify-window bounds, MSVC, test build)

CMakeLists.txt

@@ -494,11 +494,11 @@ if(STRATA_ENABLE_CUDA OR STRATA_ENABLE_HIP)
   add_test(NAME cvec_parity COMMAND cvec_parity)
 
   # Plan v0.3 P3: native MMVQ for 1..8 tokens (bitwise selftest + --bench timing); synthetic data, no model.
-  if(STRATA_BUILD_TESTS)
-  add_executable(native_mmvq_multi bench/micro/native_mmvq_multi.cpp)
-  target_include_directories(native_mmvq_multi PRIVATE ${_strata_gpu_include_directories})
-  target_link_libraries(native_mmvq_multi PRIVATE strata_kernels ${_strata_gpu_runtime_target})
-  add_test(NAME native_mmvq_multi COMMAND native_mmvq_multi --selftest)
+  if(STRATA_BUILD_TESTS AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/bench/micro/native_mmvq_multi.cpp")
+    add_executable(native_mmvq_multi bench/micro/native_mmvq_multi.cpp)
+    target_include_directories(native_mmvq_multi PRIVATE ${_strata_gpu_include_directories})
+    target_link_libraries(native_mmvq_multi PRIVATE strata_kernels ${_strata_gpu_runtime_target})
+    add_test(NAME native_mmvq_multi COMMAND native_mmvq_multi --selftest)
   endif()
 
   # Plan v0.3 P0.1: working-set lock of a host region (CPU only).
@@ -578,6 +578,10 @@ if(STRATA_BUILD_TESTS)
   add_executable(conv_cache_test src/program/conv_cache_test.cpp)
   target_include_directories(conv_cache_test PRIVATE include)
   add_test(NAME conv_cache_test COMMAND conv_cache_test)
+  # the f32 -> bf16 rounding against ggml's, NaN included, over every f32 bit pattern (header-only, CPU)
+  add_executable(bf16_bits_test src/kernels/bf16_bits_test.cpp)
+  target_include_directories(bf16_bits_test PRIVATE include)
+  add_test(NAME bf16_bits_test COMMAND bf16_bits_test)
 endif()
 
 add_library(strata_kernels_cpu STATIC src/kernels/cpu/expert.cpp src/kernels/cpu/pool.cpp
@@ -725,7 +729,7 @@ if(STRATA_NATIVE_EXPERTS)
   endif()
 endif()
 
-if(STRATA_ENABLE_CUDA AND STRATA_BUILD_TESTS)
+if(STRATA_ENABLE_CUDA AND STRATA_BUILD_TESTS AND EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/bench/micro/hit_cpu_order_parity.cu")
   add_executable(hit_cpu_order_parity bench/micro/hit_cpu_order_parity.cu)
   target_link_libraries(hit_cpu_order_parity PRIVATE strata_kernels strata_kernels_cpu ${_strata_gpu_runtime_target})
 endif()

include/strata/artifact/gguf_reader.hpp

@@ -32,6 +32,9 @@
 
 #ifdef _WIN32
 #define WIN32_LEAN_AND_MEAN
+#ifndef NOMINMAX
+#define NOMINMAX   // windows.h's min/max macros would break std::min/std::max in every file that includes this one
+#endif
 #include <windows.h>
 #else
 #include <fcntl.h>

include/strata/kernels/bf16_bits.hpp

@@ -28,13 +28,20 @@
 namespace strata::kernels {
 
 /// Round-to-nearest-even f32 -> bf16, returned as raw bits in the LOW half of the uint16.
+///
+/// **A NaN IS PASSED THROUGH, QUIETED, AS `ggml_compute_fp32_to_bf16` DOES.**  The rounding add alone
+/// carries out of the mantissa for a NaN: 0x7FFFFFFF became 0x8000 (-0) and 0x7F800001 became 0x7F80 (+inf), so
+/// a NaN upstream was masked instead of propagated.  The test is on the magnitude bits, the result keeps the
+/// sign and the top mantissa bits and sets the quiet bit (64 = bit 6 of the bf16).  Every non-NaN input,
+/// infinities included, rounds exactly as before.
 STRATA_BF16_HD inline uint16_t bf16_from_f32(float f) {
     uint32_t i;
 #if defined(__HIP_DEVICE_COMPILE__)
     __builtin_memcpy(&i, &f, 4);
 #else
     std::memcpy(&i, &f, 4);
 #endif
+    if ((i & 0x7FFFFFFFu) > 0x7F800000u) return (uint16_t) ((i >> 16) | 64u);
     i = (i + ((i >> 16) & 1u) + 0x7FFFu) & 0xFFFF0000u;
     return (uint16_t) (i >> 16);
 }

src/core/expert_source.cpp

@@ -772,6 +772,14 @@ void expert_pool_dispatch(void* user, const float* x_f, const int32_t* ids, cons
     d.experts += k;
 }
 
+namespace {
+// the verify window's per-entry tables in `expert_pool_dispatch_multi` (`kind`, `distinct`, `first_of`)
+// are fixed arrays of this many entries: MAXT tokens of the model's 10 routed experts must fit, and a larger k is
+// refused at run time rather than written past them.
+constexpr int64_t kMaxWindowEntries = 128;
+static_assert(strata::kernels::cpu::MAXT * 10 <= kMaxWindowEntries, "a verify window's entries overflow the tables");
+}  // namespace
+
 void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32_t* ids, int64_t n_tok, int64_t k,
                                 float* out) {
     using namespace strata::kernels::cpu;
@@ -782,6 +790,12 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32
         d.fail_layer = d.layers;
         return;
     }
+    if (k < 1 || n_tok * k > kMaxWindowEntries) {
+        d.failed = true;
+        d.fail = "a verify window routes more entries than the expert pool's window tables hold";
+        d.fail_layer = d.layers;
+        return;
+    }
     if ((int64_t) d.act_multi.size() < n_tok) d.act_multi.resize((size_t) MAXT);
     const ExpertLayout& lay = expert_layout();
     const bool native = lay.native;
@@ -802,9 +816,9 @@ void expert_pool_dispatch_multi(ExpertDispatch& d, const float* x_f, const int32
     // ---- plan v0.3 P6: the GPU's share, decided and published FIRST so the GPU starts while the CPU works.
     // Distinct experts in routing order; resident ones and the last pcie_num/256 of the missed ones go to the GPU.
     const int64_t n = n_tok * k;
-    int32_t kind[128];                     // per entry: -1 CPU, 0 VRAM, 1 PCIe
-    if (d.plan != nullptr && n <= 128 && n <= d.plan->cap) {
-        int64_t distinct[128], first_of[128];
+    int32_t kind[kMaxWindowEntries];       // per entry: -1 CPU, 0 VRAM, 1 PCIe
+    if (d.plan != nullptr && n <= kMaxWindowEntries && n <= d.plan->cap) {
+        int64_t distinct[kMaxWindowEntries], first_of[kMaxWindowEntries];
         int nd = 0, nmiss = 0;
         for (int64_t i = 0; i < n; ++i) {
             first_of[i] = i;

src/kernels/bf16_bits_test.cpp

@@ -0,0 +1,80 @@
+// src/kernels/bf16_bits_test.cpp - `bf16_bits.hpp` against ggml's rule, over EVERY f32 bit pattern (CPU only).
+//
+// the rounding add `i + 0x7FFF + bit16` carries out of the mantissa for a NaN, so 0x7FFFFFFF came back as
+// -0 and 0x7F800001 as +inf - an upstream NaN was masked instead of propagated.  `ggml_compute_fp32_to_bf16`
+// (ggml-impl.h at the pinned llama.cpp) tests the NaN first and returns `(i >> 16) | 64`.  Three checks, all
+// exhaustive (2^32 inputs, a few seconds):
+//
+//   1. the header equals ggml's function, transcribed below, on every input;
+//   2. on every NON-NaN input it equals the previous rounding, so the change cannot move a finite value or an
+//      infinity by a single bit;
+//   3. every NaN input gives a quiet NaN of the same sign.
+//
+// And the round trip over all 65,536 bf16 patterns: a non-NaN bf16 survives f32 and back unchanged.
+#include "strata/kernels/bf16_bits.hpp"
+
+#include <cstdint>
+#include <cstdio>
+#include <cstring>
+
+namespace {
+
+/// ggml_compute_fp32_to_bf16, ggml/src/ggml-impl.h (MIT, the ggml authors), on raw bits.
+uint16_t ggml_bf16(uint32_t u) {
+    if ((u & 0x7fffffffu) > 0x7f800000u) return (uint16_t) ((u >> 16) | 64);
+    return (uint16_t) ((u + (0x7fffu + ((u >> 16) & 1u))) >> 16);
+}
+
+/// The rule before this fix, kept only to prove the change is confined to NaN inputs.
+uint16_t previous_bf16(uint32_t u) {
+    u = (u + ((u >> 16) & 1u) + 0x7FFFu) & 0xFFFF0000u;
+    return (uint16_t) (u >> 16);
+}
+
+bool is_nan_bits(uint32_t u) { return (u & 0x7fffffffu) > 0x7f800000u; }
+
+}  // namespace
+
+int main() {
+    unsigned long long vs_ggml = 0, vs_previous = 0, nan_lost = 0, nan_inputs = 0, previous_masked = 0;
+    uint32_t first_bad = 0;
+    bool have_bad = false;
+    for (uint64_t k = 0; k <= 0xFFFFFFFFull; ++k) {
+        const uint32_t u = (uint32_t) k;
+        float f;
+        std::memcpy(&f, &u, 4);
+        const uint16_t got = strata::kernels::bf16_from_f32(f);
+        if (got != ggml_bf16(u)) {
+            ++vs_ggml;
+            if (!have_bad) { first_bad = u; have_bad = true; }
+        }
+        if (is_nan_bits(u)) {
+            ++nan_inputs;
+            const uint32_t back = (uint32_t) got << 16;
+            // a quiet NaN (bit 22 of the f32) with the input's sign
+            if (!is_nan_bits(back) || (back & 0x00400000u) == 0 || (back >> 31) != (u >> 31)) ++nan_lost;
+            if (!is_nan_bits((uint32_t) previous_bf16(u) << 16)) ++previous_masked;
+        } else if (got != previous_bf16(u)) {
+            ++vs_previous;
+            if (!have_bad) { first_bad = u; have_bad = true; }
+        }
+    }
+    std::printf("  f32 -> bf16 over 2^32 inputs: %llu differ from ggml, %llu non-NaN differ from the previous rule\n",
+                vs_ggml, vs_previous);
+    std::printf("  NaN inputs: %llu, not a quiet same-sign NaN now: %llu (the previous rule masked %llu of them)\n",
+                nan_inputs, nan_lost, previous_masked);
+    if (have_bad) std::printf("    first mismatch at input 0x%08x\n", first_bad);
+
+    unsigned long long trip_bad = 0;
+    for (uint32_t h = 0; h <= 0xFFFFu; ++h) {
+        const float f = strata::kernels::f32_from_bf16((uint16_t) h);
+        const uint16_t back = strata::kernels::bf16_from_f32(f);
+        const bool nan = is_nan_bits(h << 16);
+        if (back != (nan ? (uint16_t) (h | 64u) : (uint16_t) h)) ++trip_bad;
+    }
+    std::printf("  bf16 -> f32 -> bf16 over 65536 patterns: %llu wrong\n", trip_bad);
+
+    const bool ok = vs_ggml == 0 && vs_previous == 0 && nan_lost == 0 && trip_bad == 0 && previous_masked > 0;
+    std::printf("\nbf16_bits_test %s\n", ok ? "OK" : "FAILED");
+    return ok ? 0 : 1;
+}

src/kernels/cuda/dequant_bf16.cu

@@ -19,6 +19,9 @@ __device__ __forceinline__ float h2f(const uint8_t* p) {
 }
 __device__ __forceinline__ uint16_t f2bf(float f) {
     uint32_t u = __float_as_uint(f);
+    // a NaN (a NaN scale in the block) stays a quiet NaN, as in ggml_compute_fp32_to_bf16 and `bf16_from_f32`:
+    // the rounding add below would carry it into -0 or inf
+    if ((u & 0x7fffffffu) > 0x7f800000u) return (uint16_t) ((u >> 16) | 64u);
     u += 0x7fffu + ((u >> 16) & 1u);          // round to nearest even
     return (uint16_t) (u >> 16);
 }

src/kernels/cuda/kv_stream.cu

@@ -209,6 +209,14 @@ void kv_stream_resolve(const KvStreamMap& m, const QsaAttnPools& slots, const Kv
         std::fprintf(stderr, "kv_stream: a block's scale run must be a multiple of 16 bytes\n");
         std::exit(1);
     }
+    // one sweep step looks at RT consecutive slots `(hand + thread) % n_slots`; with fewer slots than RT two
+    // threads see the same slot and may both take it for two different misses.  The engine never streams with fewer
+    // than qsa_kv_resident_min() / page_size = 5,120 slots, so this is a guard, not a limit.
+    if (m.n_slots < RT) {
+        std::fprintf(stderr, "kv_stream: %lld slots is fewer than the resolve block (%d): the clock sweep would take a "
+                             "slot twice\n", (long long) m.n_slots, RT);
+        std::exit(1);
+    }
     resolve_kernel<<<1, RT, 0, (cudaStream_t) stream>>>(m, ids, steps, (int) n_q, (int) cap, (int) s.page_size);
     check("resolve");
     copy_kernel<<<96, 128, 0, (cudaStream_t) stream>>>(m, runs_of(slots, host, fmt, s));

src/kernels/cuda/qsa_decode_attn.cu

@@ -113,14 +113,16 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __rest
         if (t < n_here) {
             const int cell = ids[c0 + t];
             const long long page = (long long) p.page_table[cell / page_size];
-            r = (page * n_kv_heads + kvh) * page_size + (cell % page_size);
+            // a block the KV streaming could not make resident keeps page -1 (ctl[3]); its cells are masked
+            // (score -FLT_MAX, weight 0) instead of being read from before the pool.
+            if (page >= 0) r = (page * n_kv_heads + kvh) * page_size + (cell % page_size);
         }
         srow[t] = r;
     }
     __syncthreads();
     // scores: each warp takes cells warp, warp+8, ...; each lane holds 8 of the 256 dimensions.
     for (int c = warp; c < CHUNK; c += WARPS) {
-        if (c >= n_here) {
+        if (c >= n_here || srow[c] < 0) {
             if (lane < G) sp[lane][c] = -FLT_MAX;
             continue;
         }
@@ -141,8 +143,8 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __rest
     for (int h = warp; h < G; h += WARPS) {
         const float a = sp[h][lane], b = sp[h][lane + 32];
         const float m = warp_max(fmaxf(a, b));
-        const float ea = (lane < n_here) ? __expf(a - m) : 0.0f;
-        const float eb = (lane + 32 < n_here) ? __expf(b - m) : 0.0f;
+        const float ea = (lane < n_here && srow[lane] >= 0) ? __expf(a - m) : 0.0f;
+        const float eb = (lane + 32 < n_here && srow[lane + 32] >= 0) ? __expf(b - m) : 0.0f;
         sp[h][lane] = ea;
         sp[h][lane + 32] = eb;
         const float l = warp_sum(ea + eb);
@@ -154,6 +156,7 @@ __global__ void __launch_bounds__(THREADS) attn_chunk_kernel(const float* __rest
 #pragma unroll
     for (int h = 0; h < G; ++h) acc[h] = 0.0f;
     for (int c = 0; c < n_here; ++c) {
+        if (srow[c] < 0) continue;   // masked above, weight 0
         float v;
         if constexpr (KV_MODE == 0) {
             v = __half2float(__ushort_as_half(p.v_pool[srow[c] * HD + t]));

src/kernels/cuda/s2_expert_grouped.cu

@@ -16,6 +16,7 @@
 #include "strata/kernels/s2_expert_grouped.hpp"
 
 #include "strata/kernels/quantize_act.hpp"
+#include "strata/kernels/verify_kernels.hpp"
 
 #include <cuda_fp16.h>
 #include <cuda_runtime.h>
@@ -512,6 +513,9 @@ namespace {
 // hit kernel's.
 constexpr int GU_CHUNKS = (H / 32 + 31) / 32;   // 3: chunks of a gate/up row per lane (80 chunks / 32 lanes)
 constexpr int GMAX = 8;                          // entries per group (tokens routed to one expert in a window)
+// a group holds one entry per token of the window routed to its expert, and the kernels below keep at
+// most GMAX of them (`min(..., GMAX)`): a longer window would drop entries without a word.
+static_assert(GMAX >= kVerifyMaxT, "a verify window's group can exceed GMAX entries");
 constexpr int GU_ROWS = 32;                      // gate/up rows per block: 4 per warp
 constexpr int D_ROWS = 64;                       // down rows per block: 8 per warp
 

src/kernels/cuda/s_gemv.cu

@@ -147,14 +147,20 @@ __global__ void s_gemv_q8_split_kernel(const uint8_t* __restrict__ x, const uint
     constexpr int PER_BYTE = 8 / CODE_BITS;
     const int warps_per_block = (int) (blockDim.x >> 5);
     const long long o = (long long) blockIdx.x * warps_per_block + (threadIdx.x >> 5);
-    if (o >= n_out) return;
     const int lane = threadIdx.x & 31;
 
     // THE CODEBOOK, IN SHARED MEMORY.  16 bytes, loaded once per block by the first 16 threads, instead of a
     // divergent constant read per element - measured at 2.12x of this whole kernel (see the note on `kIq4Nl`).
+    //
+    // THE LOAD AND THE BARRIER COME BEFORE THE EARLY RETURN, as in `s_gemv_q8k_kernel`.  The return used to
+    // sit above them, so whenever `n_out % 8 != 0` the last block's surplus warps left before a barrier the
+    // others still waited at - undefined behaviour that every shape in the pack (a multiple of 8) happened to
+    // avoid.  `o` is per-WARP, so the return below keeps whole warps together and the shuffle reduction's
+    // full mask stays valid.
     __shared__ signed char s_iq4nl[16];
     if (threadIdx.x < 16) s_iq4nl[threadIdx.x] = kIq4Nl[threadIdx.x];
     __syncthreads();
+    if (o >= n_out) return;
 
     const long long n_groups = n_in >> group_shift;
     const long long codes_per_row = n_in / PER_BYTE;
@@ -655,8 +661,10 @@ void s_gemv_q8_0_split(const uint8_t* x_q8_0, const uint8_t* codes, const float*
                      form.group_elems);
         std::exit(1);
     }
-    if (form.group_elems % 4 != 0) {
-        std::fprintf(stderr, "s_gemv_q8_0_split: group_elems %d is not a multiple of 4\n", form.group_elems);
+    // SIXTEEN, NOT FOUR: a lane-iteration takes QE = 16 consecutive elements under ONE scale (see the kernel),
+    // so a group of 4 or 8 would read the wrong scale for most of them.  The Q8_K launcher already said 16.
+    if (form.group_elems % 16 != 0) {
+        std::fprintf(stderr, "s_gemv_q8_0_split: group_elems %d is not a multiple of 16\n", form.group_elems);
         std::exit(1);
     }
     int group_shift = 0;

src/kernels/elementwise_parity.cpp

@@ -12,6 +12,7 @@
 //      instead of decaying, so the sign is checked as a property and not assumed.
 //   3. `silu` IS COMPUTED IN DOUBLE then cast, because `ref/gdn.py`'s numpy does.  f32 `expf` differs in the
 //      last bits, and the test measures that rather than asserting it away.
+#include "strata/kernels/dequant_bf16.hpp"
 #include "strata/kernels/elementwise.hpp"
 #include "strata/kernels/f16_bits.hpp"
 
@@ -409,6 +410,64 @@ int main(int argc, char** argv) {
         if (mismatches || guards || fma_diff == 0) ++bad;
     }
 
+    // ---- THE f32 -> bf16 CONVERSIONS KEEP A NaN A NaN.  `f32_to_bf16_bulk` (the header's `bf16_from_f32`
+    // on the device) and the prompt path's dequantizer (its own `f2bf`) against ggml_compute_fp32_to_bf16's rule;
+    // before the fix 0x7FFFFFFF came back as -0 and 0x7F800001 as +inf.  `bf16_bits_test` covers every f32 on
+    // the host; this is the same header compiled by nvcc, plus the second copy in dequant_bf16.cu.
+    {
+        auto ggml_bf16 = [](uint32_t u) -> uint16_t {
+            if ((u & 0x7fffffffu) > 0x7f800000u) return (uint16_t) ((u >> 16) | 64);
+            return (uint16_t) ((u + (0x7fffu + ((u >> 16) & 1u))) >> 16);
+        };
+        std::vector<uint32_t> bits = {0x7FFFFFFFu, 0xFFFFFFFFu, 0x7F800001u, 0xFF800001u, 0x7FC00000u, 0x7FBFFFFFu,
+                                      0x7F800000u, 0xFF800000u, 0x7F7FFFFFu, 0x00000000u, 0x80000000u, 0x3F808000u,
+                                      0x3F818000u, 0x00008000u, 0x7F7F8000u};
+        std::mt19937 rng(13);
+        while (bits.size() < 4096) bits.push_back((uint32_t) rng());
+        const size_t n = bits.size();
+        float* dx = nullptr;
+        uint16_t* dy = nullptr;
+        check(cudaMalloc(&dx, n * 4), "bf16 x");
+        check(cudaMalloc(&dy, n * 2), "bf16 y");
+        check(cudaMemcpy(dx, bits.data(), n * 4, cudaMemcpyHostToDevice), "bf16 cx");
+        strata::kernels::f32_to_bf16_bulk(dx, dy, (int64_t) n, nullptr);
+        std::vector<uint16_t> got(n);
+        check(cudaMemcpy(got.data(), dy, n * 2, cudaMemcpyDeviceToHost), "bf16 cy");
+        int wrong = 0;
+        for (size_t i = 0; i < n; ++i) wrong += got[i] != ggml_bf16(bits[i]);
+
+        // A Q8_0 block (type 8) whose fp16 scale is a NaN: every dequantized value is NaN * q, a NaN, and it must
+        // still be one in BF16.  A second block with an ordinary scale checks the finite path did not move.
+        std::vector<uint8_t> blk(2 * 34, 0);
+        blk[0] = 0x00; blk[1] = 0x7E;                        // fp16 quiet NaN
+        blk[34] = 0x00; blk[35] = 0x3C;                      // fp16 1.0
+        for (int j = 0; j < 32; ++j) {
+            blk[(size_t) (2 + j)] = (uint8_t) (int8_t) (j - 16);
+            blk[(size_t) (36 + j)] = (uint8_t) (int8_t) (3 * j - 50);
+        }
+        uint8_t* db = nullptr;
+        uint16_t* dq = nullptr;
+        check(cudaMalloc(&db, blk.size()), "q8 blk");
+        check(cudaMalloc(&dq, 64 * 2), "q8 out");
+        check(cudaMemcpy(db, blk.data(), blk.size(), cudaMemcpyHostToDevice), "q8 cblk");
+        strata::kernels::dequant_bf16(8, db, 0, 2, 32, dq, nullptr);
+        check(cudaDeviceSynchronize(), "dequant_bf16");
+        std::vector<uint16_t> q(64);
+        check(cudaMemcpy(q.data(), dq, 64 * 2, cudaMemcpyDeviceToHost), "q8 cout");
+        int dq_wrong = 0;
+        for (int j = 0; j < 32; ++j) {
+            if (!((q[(size_t) j] & 0x7FFFu) > 0x7F80u)) ++dq_wrong;           // row 0: NaN scale -> NaN
+            const float v = (float) (3 * j - 50);                                // row 1: d = 1.0 -> the integer
+            uint32_t vb;
+            std::memcpy(&vb, &v, 4);
+            if (q[(size_t) (32 + j)] != ggml_bf16(vb)) ++dq_wrong;
+        }
+        std::printf("  f32 -> bf16 (NaN kept): bulk %d of %zu differ from ggml, dequant_bf16 %d of 64 wrong\n",
+                    wrong, n, dq_wrong);
+        if (wrong || dq_wrong) ++bad;
+        cudaFree(dx); cudaFree(dy); cudaFree(db); cudaFree(dq);
+    }
+
     std::printf("\nelementwise: %d failures\n", bad);
     if (bad) return 1;
     if (selftest) std::printf("elementwise_parity OK\n");