Fix: Harden reductions, casts, and bindings against degenerate inputs

Several SIMD kernels and Python-binding paths mishandled degenerate shapes and
strides, ranging from a hang to out-of-bounds reads and writes. The reduce
kernels (skylake/icelake/haswell/sierra/alder/neonbfdot) hit an infinite loop,
SIGFPE, or stack overflow when handed stride_bytes == 0; their `!aligned`
serial fallbacks now also catch `stride_elements == 0`. The sub-byte cast
oracle in cast/serial.h hard-coded its pack/unpack loops to four bytes
regardless of count, reading and writing past the buffer on any odd i4/u4/e2m1
length -- every site now bounds the loop by nk_size_divide_round_up_(count, 2)
and guards the odd tail. Rank-0 and empty reductions no longer touch a
non-existent element: C++ moments()/minmax() indexed stride_bytes(SIZE_MAX) on
a 0-D view, the Python rank-0 path fabricated a zero stride, and empty minmax
primed its accumulators from element [0].

The Python bindings gained the missing input validation. Dense-metric `out=`
buffers are checked for rank and capacity before writing (an undersized out
overflowed the heap); parse_tensor now requires exact inner-axis contiguity,
rejecting the negative/zero strides (e.g. x[::-1]) that the old signed
`> itemsize` check let walk off the buffer; and DLPack import rejects negative
extents and a NULL data pointer. The Rust symmetric matrix verbs reject
non-contiguous-row (transposed) views, matching the packed and parallel paths.

The test suites were deduplicated and extended over the same edges. The Python
*_float/*_integer reduction and arithmetic pairs collapse into parametrized
tests (removing dead precise_*/baseline_sum helpers and a stale DLPack helper),
and new cases cover empty/0-D tensors, NaN/Inf preservation on dense casts, and
block-scaled round-trip / idempotence / byte-identity across all seven formats,
plus the degenerate inputs above. Net test lines shrink while coverage grows.

Comments were stripped of decorative banner separators and ephemeral numbered
"Phase"/"Step"/"Option" labels, whose ordering is implied by code order. Also
sets the rustfmt line width to 120 and bumps the Rust MSRV to 1.73 for
usize::div_ceil.

Author: ashvardanian

.cmake-format.py

@@ -1,6 +1,4 @@
-# -----------------------------
-# Options effecting formatting.
-# -----------------------------
+# Options affecting formatting.
 with section("format"):
     # How wide to allow formatted cmake files
     line_width = 120
@@ -29,9 +27,7 @@
     # one-per-line when word-wrapping exceeds 2 lines)
     max_lines_hwrap = 8
 
-# ----------------------------------
 # Options affecting comment handling.
-# ----------------------------------
 with section("markup"):
     # Do not reflow comment text
     enable_markup = False

Cargo.toml

@@ -18,7 +18,7 @@ links = "numkong"
 name = "numkong"
 readme = "rust/README.md"
 repository = "https://github.com/ashvardanian/NumKong" 
-rust-version = "1.64" # Introduced Core C FFI in stable Rust
+rust-version = "1.73" # `usize::div_ceil` (used for sub-byte / block-scaled size math)
 version = "7.7.0" 
 
 [lib]

include/numkong/attention/sapphireamx.h

@@ -788,7 +788,7 @@ NK_PUBLIC void nk_attention_bf16_sapphireamx(nk_bf16_t const *q, void const *kv_
                 // Extract K block: Kᵀ[head_dim, valid_kv] using bulk extraction
                 nk_attention_extract_k_block_(k_packed, k_block, kv_h, kvb, valid_kv, head_dim, kv_len);
 
-                // Phase 1: Compute S = Q × Kᵀ using AVX-512 FMA
+                // Compute S = Q × Kᵀ using AVX-512 FMA
                 for (nk_size_t qi = 0; qi < valid_q; qi++) {
                     for (nk_size_t ki = 0; ki < valid_kv; ki++) {
                         __m512 sum_v_f32x16 = _mm512_setzero_ps();
@@ -819,7 +819,7 @@ NK_PUBLIC void nk_attention_bf16_sapphireamx(nk_bf16_t const *q, void const *kv_
                     for (nk_size_t ki = 0; ki < 16; ki++) { scores[qi * 16 + ki] = NK_F32_MIN; }
                 }
 
-                // Phase 2: Online softmax update
+                // Online softmax update
                 __m512 old_max_f32x16 = softmax_state.row_max_f32x16;
                 nk_attention_softmax_update_(&softmax_state, scores, scale, weights);
 
@@ -829,7 +829,7 @@ NK_PUBLIC void nk_attention_bf16_sapphireamx(nk_bf16_t const *q, void const *kv_
                 // Extract V block: V[valid_kv, head_dim] using bulk extraction
                 nk_attention_extract_v_block_(v_packed, v_block, kv_h, kvb, valid_kv, head_dim, kv_len);
 
-                // Phase 3: Compute O += P × V using AVX-512 FMA
+                // Compute O += P × V using AVX-512 FMA
                 for (nk_size_t qi = 0; qi < valid_q; qi++) {
                     nk_size_t d = 0;
                     // Vectorized loop over head_dim
@@ -931,7 +931,7 @@ NK_PUBLIC void nk_attention_bf16_amx_bc32_sapphireamx(nk_bf16_t const *q, void c
             for (nk_size_t kvb = 0; kvb < kv_len; kvb += Bc) {
                 nk_size_t valid_kv = (kvb + Bc <= kv_len) ? Bc : (kv_len - kvb);
 
-                // Phase 1: S = Q × Kᵀ using AMX
+                // S = Q × Kᵀ using AMX
                 // Need 2 K tiles per block (each K tile has 16 columns)
                 nk_size_t k_tile_idx0 = kvb / 16;        // First K tile
                 nk_size_t k_tile_idx1 = (kvb + 16) / 16; // Second K tile
@@ -1038,12 +1038,12 @@ NK_PUBLIC void nk_attention_bf16_amx_bc32_sapphireamx(nk_bf16_t const *q, void c
                     for (nk_size_t qi = 0; qi < 16; qi++) { _mm512_store_ps(&scores[qi * 32 + 16], neg_inf_f32x16); }
                 }
 
-                // Phase 2: online softmax (fast degree-4 exp)
+                // online softmax (fast degree-4 exp)
                 __m512 old_max_f32x16 = softmax_state.row_max_f32x16;
                 nk_attention_softmax_update_bc32_fast_(&softmax_state, scores, scale, weights);
                 nk_attention_rescale_output_(o_acc, head_dim_padded, old_max_f32x16, softmax_state.row_max_f32x16);
 
-                // Phase 3: O += P × V using AMX
+                // O += P × V using AMX
                 // Convert P[16, 32] from F32 to BF16 and pack as A-tile
                 for (nk_size_t qi = 0; qi < 16; qi++) {
                     for (nk_size_t ki = 0; ki < 32; ki += 16) {
@@ -1213,7 +1213,7 @@ NK_PUBLIC void nk_attention_bf16_amx_optimized_sapphireamx(nk_bf16_t const *q, v
                 nk_size_t k_tile_idx0 = kvb / 16;
                 nk_size_t k_tile_idx1 = (kvb + 16) / 16;
 
-                // Phase 1: S = Q × Kᵀ using pre-packed Q tiles
+                // S = Q × Kᵀ using pre-packed Q tiles
                 _tile_zero(0); // Score cols 0:16
                 _tile_zero(3); // Score cols 16:32
 
@@ -1263,13 +1263,13 @@ NK_PUBLIC void nk_attention_bf16_amx_optimized_sapphireamx(nk_bf16_t const *q, v
                     }
                 }
 
-                // Phase 2: online softmax (fast degree-4 exp)
+                // online softmax (fast degree-4 exp)
                 __m512 old_max_f32x16 = softmax_state.row_max_f32x16;
                 nk_attention_softmax_update_bc32_fast_(&softmax_state, &scores[0][0], scale, &weights[0][0]);
                 nk_attention_rescale_output_(&o_acc[0][0], head_dim_padded, old_max_f32x16,
                                              softmax_state.row_max_f32x16);
 
-                // Phase 3: O += P × V with hoisted P tile load
+                // O += P × V with hoisted P tile load
                 // Convert F32 weights to BF16 P tile (once per KV block)
                 for (nk_size_t qi = 0; qi < 16; qi++) {
                     __m512 p0_f32x16 = _mm512_load_ps(&weights[qi][0]);

include/numkong/attention/sme.h

@@ -63,14 +63,14 @@
  *
  *  @section attention_sme_history Optimization History
  *
- *  Phase 1 (January 2026): Initial implementation using ZA staging transpose for Q×Kᵀ
+ *  January 2026: Initial implementation using ZA staging transpose for Q×Kᵀ
  *  and element-wise SVE for P×V. Q and K rows were loaded into ZA0/ZA1 horizontally,
  *  read back vertically to produce interleaved vectors for BFMOPA. The P×V phase used
  *  scalar `svmla_f32_x` loops over head_dim for each query-key pair. Softmax used
  *  degree-4 polynomial exp with per-row horizontal max/sum. Performance: ~25-50 GFLOP/s
  *  on Apple M4 (bf16, 8 heads, query_len=64, kv_len=4096, head_dim=128).
  *
- *  Phase 2 (February 2026): BFMOPA/FMOPA P×V with pre-packed V in interleaved format.
+ *  February 2026: BFMOPA/FMOPA P×V with pre-packed V in interleaved format.
  *  Key changes integrated:
  *  - Q pre-transposed once into a buffer, eliminating per-block ZA staging for Q
  *  - K pre-packed in interleaved format, enabling pure memory-to-BFMOPA Q×Kᵀ
@@ -1260,7 +1260,7 @@ __arm_new("za") static void nk_attention_f16_sme_streaming_( //
                       svread_ver_za32_f32_m(svdup_f32(0), predicate_all_b32x, 0, step));
     }
 
-    // === Bc=32 main loop (prefill only, skipped for decode) ===
+    // Bc=32 main loop, used for prefill only and skipped for decode.
     if (valid_query_count > 1) {
         for (; kv_start + 32 <= kv_len; kv_start += 32, kv_block_index += 2) {
             // Q×K^T: pure memory→FMOPA, no ZA staging for Q or K
@@ -1691,7 +1691,7 @@ __arm_new("za") static void nk_attention_f16_sme_streaming_( //
         }
     }
 
-    // === Bc=16 tail loop (handles remaining KV positions and decode path) ===
+    // Bc=16 tail loop handles the remaining KV positions and the decode path.
     for (; kv_start < kv_len; kv_start += 16, kv_block_index++) {
         nk_size_t const valid_kv = ((kv_start + 16) <= kv_len) ? 16 : (kv_len - kv_start);
 

include/numkong/cast/serial.h

@@ -2012,7 +2012,7 @@ NK_INTERNAL void nk_scalar_buffers_to_f64c_(                           //
     case nk_i4_k: {
         nk_i4x2_t const *pairs = (nk_i4x2_t const *)from_ptr;
         nk_i8_t unpacked[2];
-        for (i = 0; i < 4; ++i) {
+        for (i = 0; i < nk_size_divide_round_up_(from_count, 2); ++i) {
             nk_i4x2_to_i8x2_serial(&pairs[i], unpacked);
             to_buffers[i * 2].f64c.real = unpacked[0], to_buffers[i * 2].f64c.imag = 0;
             to_buffers[i * 2 + 1].f64c.real = unpacked[1], to_buffers[i * 2 + 1].f64c.imag = 0;
@@ -2022,7 +2022,7 @@ NK_INTERNAL void nk_scalar_buffers_to_f64c_(                           //
     case nk_u4_k: {
         nk_u4x2_t const *pairs = (nk_u4x2_t const *)from_ptr;
         nk_u8_t unpacked[2];
-        for (i = 0; i < 4; ++i) {
+        for (i = 0; i < nk_size_divide_round_up_(from_count, 2); ++i) {
             nk_u4x2_to_u8x2_serial(&pairs[i], unpacked);
             to_buffers[i * 2].f64c.real = unpacked[0], to_buffers[i * 2].f64c.imag = 0;
             to_buffers[i * 2 + 1].f64c.real = unpacked[1], to_buffers[i * 2 + 1].f64c.imag = 0;
@@ -2032,7 +2032,7 @@ NK_INTERNAL void nk_scalar_buffers_to_f64c_(                           //
     case nk_e2m1_k: {
         nk_e2m1x2_t const *pairs = (nk_e2m1x2_t const *)from_ptr;
         nk_f32_t unpacked[2];
-        for (i = 0; i < 4; ++i) {
+        for (i = 0; i < nk_size_divide_round_up_(from_count, 2); ++i) {
             nk_e2m1x2_to_f32x2_serial(&pairs[i], unpacked);
             to_buffers[i * 2].f64c.real = (nk_f64_t)unpacked[0], to_buffers[i * 2].f64c.imag = 0;
             to_buffers[i * 2 + 1].f64c.real = (nk_f64_t)unpacked[1], to_buffers[i * 2 + 1].f64c.imag = 0;
@@ -2149,8 +2149,9 @@ NK_INTERNAL void nk_scalar_buffers_from_f64c_(              //
     // Sub-byte: i4 - 8 nibbles to 4 bytes, high nibble = even index
     case nk_i4_k: {
         nk_u8_t *p = (nk_u8_t *)to_ptr;
-        for (i = 0; i < 4; ++i) {
-            nk_f64_t high = from_buffers[i * 2].f64c.real, low = from_buffers[i * 2 + 1].f64c.real;
+        for (i = 0; i < nk_size_divide_round_up_(to_count, 2); ++i) {
+            nk_f64_t high = from_buffers[i * 2].f64c.real;
+            nk_f64_t low = (i * 2 + 1 < to_count) ? from_buffers[i * 2 + 1].f64c.real : 0.0;
             high = high > 7 ? 7 : (high < -8 ? -8 : high);
             low = low > 7 ? 7 : (low < -8 ? -8 : low);
             p[i] = (nk_u8_t)((((nk_i8_t)high & 0x0F) << 4) | ((nk_i8_t)low & 0x0F));
@@ -2159,8 +2160,9 @@ NK_INTERNAL void nk_scalar_buffers_from_f64c_(              //
     // Sub-byte: u4 - 8 nibbles to 4 bytes, high nibble = even index
     case nk_u4_k: {
         nk_u8_t *p = (nk_u8_t *)to_ptr;
-        for (i = 0; i < 4; ++i) {
-            nk_f64_t high = from_buffers[i * 2].f64c.real, low = from_buffers[i * 2 + 1].f64c.real;
+        for (i = 0; i < nk_size_divide_round_up_(to_count, 2); ++i) {
+            nk_f64_t high = from_buffers[i * 2].f64c.real;
+            nk_f64_t low = (i * 2 + 1 < to_count) ? from_buffers[i * 2 + 1].f64c.real : 0.0;
             high = high > 15 ? 15 : (high < 0 ? 0 : high);
             low = low > 15 ? 15 : (low < 0 ? 0 : low);
             p[i] = (nk_u8_t)(((nk_u8_t)high << 4) | (nk_u8_t)low);
@@ -2169,10 +2171,10 @@ NK_INTERNAL void nk_scalar_buffers_from_f64c_(              //
     // Sub-byte: e2m1 - 8 nibbles to 4 bytes, high nibble = even index
     case nk_e2m1_k: {
         nk_e2m1x2_t *pairs = (nk_e2m1x2_t *)to_ptr;
-        for (i = 0; i < 4; ++i) {
+        for (i = 0; i < nk_size_divide_round_up_(to_count, 2); ++i) {
             nk_f32_t paired[2];
             paired[0] = (nk_f32_t)from_buffers[i * 2].f64c.real;
-            paired[1] = (nk_f32_t)from_buffers[i * 2 + 1].f64c.real;
+            paired[1] = (i * 2 + 1 < to_count) ? (nk_f32_t)from_buffers[i * 2 + 1].f64c.real : 0.0f;
             nk_f32x2_to_e2m1x2_serial(paired, &pairs[i]);
         }
     } break;
@@ -2215,7 +2217,7 @@ NK_INTERNAL void nk_scalar_buffers_to_i64_(                            //
     // Sub-byte: i4 - 4 bytes to 8 nibbles, sign-extend each nibble
     case nk_i4_k: {
         nk_i4x2_t const *pairs = (nk_i4x2_t const *)from_ptr;
-        for (i = 0; i < 4; ++i) {
+        for (i = 0; i < nk_size_divide_round_up_(from_count, 2); ++i) {
             nk_i8_t unpacked[2];
             nk_i4x2_to_i8x2_serial(&pairs[i], unpacked);
             to_buffers[i * 2].i64 = unpacked[0];
@@ -2240,7 +2242,7 @@ NK_INTERNAL void nk_scalar_buffers_to_i64_(                            //
     } break;
     case nk_u4_k: {
         nk_u8_t const *p = (nk_u8_t const *)from_ptr;
-        for (i = 0; i < 4; ++i) {
+        for (i = 0; i < nk_size_divide_round_up_(from_count, 2); ++i) {
             to_buffers[i * 2].i64 = (nk_i64_t)(p[i] >> 4);
             to_buffers[i * 2 + 1].i64 = (nk_i64_t)(p[i] & 0xF);
         }
@@ -2294,8 +2296,9 @@ NK_INTERNAL void nk_scalar_buffers_from_i64_(                //
     // Sub-byte: i4 - 8 nibbles to 4 bytes, clamp [-8,7]
     case nk_i4_k: {
         nk_i4x2_t *p = (nk_i4x2_t *)to_ptr;
-        for (i = 0; i < 4; ++i) {
-            nk_i64_t high = from_buffers[i * 2].i64, low = from_buffers[i * 2 + 1].i64;
+        for (i = 0; i < nk_size_divide_round_up_(to_count, 2); ++i) {
+            nk_i64_t high = from_buffers[i * 2].i64;
+            nk_i64_t low = (i * 2 + 1 < to_count) ? from_buffers[i * 2 + 1].i64 : 0;
             high = high > 7 ? 7 : (high < -8 ? -8 : high);
             low = low > 7 ? 7 : (low < -8 ? -8 : low);
             p[i] = (nk_u8_t)(((high & 0xF) << 4) | (low & 0xF));
@@ -2332,7 +2335,7 @@ NK_INTERNAL void nk_scalar_buffers_to_u64_(                            //
     // Sub-byte: u4 - 4 bytes to 8 nibbles, zero-extend
     case nk_u4_k: {
         nk_u4x2_t const *pairs = (nk_u4x2_t const *)from_ptr;
-        for (i = 0; i < 4; ++i) {
+        for (i = 0; i < nk_size_divide_round_up_(from_count, 2); ++i) {
             nk_u8_t unpacked[2];
             nk_u4x2_to_u8x2_serial(&pairs[i], unpacked);
             to_buffers[i * 2].u64 = unpacked[0];
@@ -2393,8 +2396,9 @@ NK_INTERNAL void nk_scalar_buffers_from_u64_(                //
     // Sub-byte: u4 - 8 nibbles to 4 bytes, clamp [0,15]
     case nk_u4_k: {
         nk_u4x2_t *p = (nk_u4x2_t *)to_ptr;
-        for (i = 0; i < 4; ++i) {
-            nk_u64_t high = from_buffers[i * 2].u64, low = from_buffers[i * 2 + 1].u64;
+        for (i = 0; i < nk_size_divide_round_up_(to_count, 2); ++i) {
+            nk_u64_t high = from_buffers[i * 2].u64;
+            nk_u64_t low = (i * 2 + 1 < to_count) ? from_buffers[i * 2 + 1].u64 : 0;
             high = high > 15 ? 15 : high;
             low = low > 15 ? 15 : low;
             p[i] = (nk_u8_t)((high << 4) | low);
@@ -2861,7 +2865,7 @@ NK_PUBLIC void nk_cast_block_scaled_serial(
     for (nk_size_t chunk_start = 0; chunk_start < count; chunk_start += chunk) {
         nk_size_t chunk_count = (chunk_start + chunk <= count) ? chunk : (count - chunk_start);
 
-        // --- Decode: source chunk → scratch[0..chunk_count) as f32 ---
+        // Decode the source chunk into scratch[0..chunk_count) as f32.
         if (from_plain) {
             void const *src = (nk_u8_t const *)from + (chunk_start * from_bits_per_element / NK_BITS_PER_BYTE);
             nk_cast_serial(src, from_format->element_dtype, chunk_count, scratch, nk_f32_k);
@@ -2880,7 +2884,7 @@ NK_PUBLIC void nk_cast_block_scaled_serial(
             }
         }
 
-        // --- Encode: scratch[0..chunk_count) as f32 → destination chunk ---
+        // Encode scratch[0..chunk_count) as f32 into the destination chunk.
         if (to_plain) {
             void *dst = (nk_u8_t *)to + (chunk_start * to_bits_per_element / NK_BITS_PER_BYTE);
             nk_cast_serial(scratch, nk_f32_k, chunk_count, dst, to_format->element_dtype);

include/numkong/dot/haswell.h

@@ -1505,7 +1505,7 @@ NK_INTERNAL void nk_dot_through_i32_finalize_haswell_(
     nk_size_t total_dimensions, nk_b128_vec_t *result) {
     nk_unused_(total_dimensions);
     // ILP-optimized 4-way horizontal reduction for i32 in AVX2
-    // Step 1: 8->4 for all 4 states
+    // 8->4 for all 4 states
     __m128i sum_a_i32x4 = _mm_add_epi32(_mm256_castsi256_si128(state_a->sum_i32x8),
                                         _mm256_extracti128_si256(state_a->sum_i32x8, 1));
     __m128i sum_b_i32x4 = _mm_add_epi32(_mm256_castsi256_si128(state_b->sum_i32x8),
@@ -1514,7 +1514,7 @@ NK_INTERNAL void nk_dot_through_i32_finalize_haswell_(
                                         _mm256_extracti128_si256(state_c->sum_i32x8, 1));
     __m128i sum_d_i32x4 = _mm_add_epi32(_mm256_castsi256_si128(state_d->sum_i32x8),
                                         _mm256_extracti128_si256(state_d->sum_i32x8, 1));
-    // Step 2: Transpose 4×4 matrix
+    // Transpose 4×4 matrix
     __m128i transpose_ab_low_i32x4 = _mm_unpacklo_epi32(sum_a_i32x4, sum_b_i32x4);
     __m128i transpose_cd_low_i32x4 = _mm_unpacklo_epi32(sum_c_i32x4, sum_d_i32x4);
     __m128i transpose_ab_high_i32x4 = _mm_unpackhi_epi32(sum_a_i32x4, sum_b_i32x4);
@@ -1523,7 +1523,7 @@ NK_INTERNAL void nk_dot_through_i32_finalize_haswell_(
     __m128i sum_lane1_i32x4 = _mm_unpackhi_epi64(transpose_ab_low_i32x4, transpose_cd_low_i32x4);
     __m128i sum_lane2_i32x4 = _mm_unpacklo_epi64(transpose_ab_high_i32x4, transpose_cd_high_i32x4);
     __m128i sum_lane3_i32x4 = _mm_unpackhi_epi64(transpose_ab_high_i32x4, transpose_cd_high_i32x4);
-    // Step 3: Vertical sum and store as i32
+    // Vertical sum and store as i32
     __m128i sum_i32x4 = _mm_add_epi32(_mm_add_epi32(sum_lane0_i32x4, sum_lane1_i32x4),
                                       _mm_add_epi32(sum_lane2_i32x4, sum_lane3_i32x4));
     result->xmm = sum_i32x4;