v8: GPUs, MXFP & NVFP4 for Rust, Python, and C++#350
Open
ashvardanian wants to merge 147 commits into
Open
Conversation
Adds a struct-of-arrays block-scaled tensor across C, C++, Rust, and Python: packed micro-float elements, a per-block scale tensor, and an optional per-tensor scale (NVFP4). The C kernel `nk_cast_block_scaled` encodes, decodes, and transcodes along the last axis across serial, Haswell, Skylake, Icelake, and NEON; the shared UE8M0 scale follows the OCP MX v1.0 floor formula so a block's own maximum is never clipped, with RNE element rounding and an E5M2 saturate clamp. Bit-to-byte sizes now use the round-up helpers and a sub-byte SIMD over-read was fixed. C++ adds `scaled_tensor`/`scaled_tensor_view`/`scaled_tensor_span` composing the existing tensor family; `cast()` encodes, decodes, and transcodes in place through a single block-scaled marshalling helper and is rank-general over the last axis. The cast backends were tidied too: the Haswell FP4 path now mirrors Icelake, and the SIMD loop counter and block-amax helper names were unified. Rust adds `ScaledTensor`/`View`/`Span` composing `Tensor`, with explicit per-format impls (no macros) and rank-general `try_cast_to_scaled` / `try_cast_dense` verbs funneled through one FFI helper. `TensorView` is now `Copy` (an immutable view is a shared borrow), removing the hand-rolled reborrow. Python adds a `ScaledTensor` produced and consumed via `astype`, with read-only `elements`/`block_scales`/`tensor_scale`/`block_size` attributes; the components are DLPack- and CUDA-array-interface-exportable for zero-copy GPU interop. Tests cover byte-identity against the C reference, rank-1/2/3 round-trips, slicing, and transcode in every binding.
The Python `Tensor(...)` constructor mishandled non-C-contiguous buffers — it
honored only the outermost stride and assumed packed inner axes, silently
corrupting Fortran-order, transposed, and strided NumPy inputs. It now routes
through `linearize_cast_into`, which walks every axis' stride.
`ScaledTensor.astype("<block-scaled>")` now transcodes between block-scaled
formats (e.g. NVFP4 -> MXFP8) instead of raising, matching the C++ and Rust
bindings. Adds coverage in both Python and Rust for non-contiguous / transposed
construction and for transcoding (verified equal to decode-then-encode).
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.
ashvardanian
force-pushed
the
main-v8-block-scaling
branch
from
June 12, 2026 11:57
a080325 to
aa7dddb
Compare
3 tasks
ashvardanian
changed the title
The block-scaling cast paths reference NULL (`from_tensor_scale != NULL`) in
cast/{serial,neon,icelake,skylake,haswell}.h, and capabilities.h passes NULL to
sysctlbyname. NumKong headers intentionally avoid <stddef.h>, so NULL resolves
only transitively through the C translation units' <stdlib.h>. The Swift module
and WASM/WASI clang builds compile the headers standalone and fail with
"use of undeclared identifier 'NULL'", which blocked the Swift, WASM, and wheel
CI jobs.
Add an NK_NULL macro to types.h (mirroring StringZilla's SZ_NULL: __null on
GCC/Clang, ((void *)0) otherwise) and use it at those call sites -- no header
dependency is introduced.
Also drop the stale float8_e8m0fnu case from test_ml_dtypes_incompatible_rejected:
e8m0 is the OCP UE8M0 scale format, now a first-class dtype via block scaling, so
it is accepted and round-trips exactly. A positive test covers it.
Add nk_size_mul_checked_ to types.h and route the shape-product and cdist size
computations in python/{each,distance}.c through it, so a buffer reporting an
overflowing shape raises OverflowError instead of wrapping into an undersized
allocation walked at the true extent.
The elementwise scalar-array paths returned NULL without setting an exception
when the buffer dtype is unsupported (a CPython protocol violation); they now
raise TypeError. The packed and symmetric matrix verbs leaked the owned output
tensor on an invalid row range; they now Py_DECREF it. The DLPack importer
propagates a PyObject_IsTrue error and nulls owner->managed if PyCapsule_SetName
fails, avoiding a double-free of the producer tensor.
cibuildwheel 4.0 removed the `cpython-freethreading` enable group (free-threaded CPython is built by default now), so the pinned `CIBW_ENABLE` aborts every wheel job with "Unknown enable group". Drop it from all seven wheel matrices; `cp31Xt-*` selectors still build the no-GIL wheels. Stop pinning the `Visual Studio 17 2022` generator in the Windows C test and release jobs; let CMake auto-detect whichever Visual Studio the runner ships (windows-2025 now carries VS18). `-A` still selects the target architecture. Intel's downloadmirror.intel.com now sits behind an AWS WAF challenge (HTTP 202, x-amzn-waf-action: challenge) that hands scripted clients an HTML page instead of the tarball, so `tar -xf` fails. Fetch SDE from the petarpetrovt/setup-sde GitHub-hosted mirror instead.
cibuildwheel 4.x removed support for Python 3.13 free-threading entirely (`cp313t` is no longer a build identifier and the tool errors out on the selector). Python 3.14 free-threading is the first supported tier and is built by default — no enable group needed. Drop the now-invalid `313t` matrix rows across every wheel platform; `314t` continues to ship.
`nk_size_mul_checked_` takes `nk_size_t *`, but the cdist and elementwise bindings passed `&size_t`. On x86_64 `nk_size_t` is `unsigned long long` while `size_t` is `unsigned long` (same width, distinct types), so GCC's `-Wincompatible-pointer-types` fires — and gcc-toolset-14 in the manylinux images treats it as a hard error, breaking every Linux wheel. clang and older gcc only warned, so it slipped through local builds. Do the count arithmetic in `nk_size_t` (the helper's and `nk_cast`'s domain) and convert to `size_t` only at the libc boundary, which also drops the now-redundant `(nk_size_t)total_elements` casts at the cast calls.
The ARM64 Windows wheels are cross-compiled on an x64 host. `/openmp:llvm` makes the extension import `libomp140.aarch64.dll`, which ships with Visual Studio's ARM64 cross-tools but is not on PATH, so delvewheel could not find it to bundle and the repair step failed. Add a step (ARM64 jobs only) that locates `libomp140.aarch64.dll` in the Visual Studio install and puts its directory on PATH. delvewheel searches PATH (matching the DLL by architecture), so it now vendors the runtime into the wheel. OpenMP stays enabled on every platform.
`tensor_slice_suffix_(all_t, rest...)` computed `inner.byte_data() - first_row.byte_data()` to locate the sliced sub-shape. When `rest...` over-slices the row to an empty tensor, `inner.byte_data()` is null, and a pointer subtraction with a null operand is undefined behavior. Detect the null inner and return an empty view first; a valid rank-0 scalar slice keeps a non-null pointer and is unaffected.
When the source point cloud has zero variance (all points identical), the similarity scale is mathematically undefined. The backends disagreed: the serial oracle and skylake's f32 path divided unguarded and yielded Inf/NaN, while Haswell, genoa, every NEON/RVV/v128 path, and skylake's f16/f64/bf16 paths guarded the division to return 0. Returning 0 silently produces a plausible-looking but wrong transform. Drop the guards so every backend propagates the Inf/NaN of the undefined division, matching the serial reference and making the degenerate case an explicit signal rather than a quiet 0. Non-degenerate inputs are unchanged (mesh suite passes at machine-epsilon error).
Construction sizes storage as `shape.product() / dimensions_per_value()`; for a rank-0 tensor the empty product is 1, so it allocates one element. `Drop` special-cased `ndim == 0` to a storage count of 0 and skipped the deallocation, leaking that element. Compute the count the same way construction does so the element is freed (rank-0 sub-byte stays at 0 via the same integer division, matching its dangling allocation).
The in-place ops (scale/add, scalar add/sub/mul, sin/cos/atan, and add/sub/mul-tensor) routed through `try_reborrow_tensor_inplace`, which fabricated a read-only TensorView aliasing the mutable TensorSpan over the same storage. The kernel walk then formed an overlapping `&[T]` + `&mut [T]` over those bytes — undefined behavior under Stacked/Tree Borrows (the C kernels tolerate the aliasing, but Rust's model does not). In-place mutation now lives on TensorSpan and operates on its own storage: each per-type `each_*_inplace` entry derives both raw pointers from the single `&mut` (source == dest, which the C kernels accept), and the binary ops keep `other` as a disjoint `&[T]`. No read-view of the target is ever constructed. `try_reborrow_tensor_inplace` is deleted; Tensor's in-place methods delegate to `self.span()`. Error handling is made consistent: unary in-place is infallible, binary validates shapes and returns Result. Adds in-place == out-of-place regression tests.
Run `cargo +nightly fmt` so the crate matches rustfmt.toml (fn_single_line, max_width=120, unstable_features). Pure formatting — no behavior change. Committing the reflow on its own keeps it out of logic diffs and makes future `fmt` runs no-ops.
The `all_t` slice overload guarded `inner.byte_data() - first_row.byte_data()` against a null inner, but the parallel `range` overload did the same pointer subtraction unguarded — UB when a nested slice over-slices to an empty tensor. Extract one `slice_inner_byte_offset_` helper (also rejecting a null first_row) and use it in both overloads.
The WASM-via-Wasmer job built `wasmer-cli` from `--git main` to get relaxed-SIMD support. Upstream's main now pulls a private SSH submodule (git@github.com:wasmerio/quickjs.git) that fails to authenticate in CI, so the install aborts. Relaxed-SIMD now ships in released Wasmer, so install the prebuilt latest release via get.wasmer.io — no source build, no submodule fetch.
Two RVV cast bugs, both verified against the serial oracle under qemu-riscv64 (CI only compiles RVV, so these were never caught): - f32→f16 (`nk_f32m2_to_f16m1_rvv_`) was a simplified conversion that clamped the exponent and rounded half-up, so overflow did not saturate to infinity (65520 → a finite value, 70000 → garbage), denormals and underflow were wrong, and rounding was not round-to-nearest-even. Rewritten to compute every IEEE bucket (zero, inf/nan, underflow, denormal, normal, overflow) with RNE and merge by exponent mask, matching nk_f32_to_f16_serial bit-for-bit (incl. 200 random values). - The packed i4/u4 ↔ i8/u8 casts sized their vector loop as `count / 2`, dropping the trailing element on odd counts (and skipping the whole cast at count == 1). Add a scalar tail for the odd high nibble, clamping packs to [-8,7]/[0,15] to match the vector helpers.
…erflow Two untrusted-producer ingestion paths in the Python binding computed byte extents from caller-supplied geometry without overflow or sign checks, so a hostile array could wrap a size into an undersized buffer walked out of bounds. `nk_get_buffer_via_array_interface` accepted negative `__array_interface__` shape entries (only a literal -1 was caught, as the PyLong error sentinel) and multiplied the dimensions into the C-contiguous strides and `buffer->len` with unchecked, signed-overflow-prone arithmetic. It now rejects negative dims with ValueError and routes both the stride accumulation and the length product through `nk_size_mul_checked_`, additionally rejecting a length that exceeds PY_SSIZE_T_MAX, raising OverflowError on wrap. `api_from_dlpack` formed each byte stride as `strides[i] * item_size` from a producer int64 with no guard. It now validates the stride magnitude (strides may be legitimately negative for reversed axes) against PY_SSIZE_T_MAX/item_size before any allocation, so the check fails cleanly with no owner/view to unwind. Verified: negative-dim and overflowing-shape arrays now raise; valid vectors and normal DLPack round-trips are unaffected; the dlpack + tensor Python suites pass (950 passed, 0 regressions).
Three long string literals had been chopped by clang-format into tiny adjacent fragments split mid-word (e.g. "IS_" "SUBBYTE_" "TYPE_" "PA" "DD" "ED"), because the source lines lacked the trailing `//` that the rest of the docstrings use to opt out of reflow. Reconstruct each into readable fragments and append the trailing `//` so clang-format leaves them intact (now a verified fixed point): - dlpack_interop.c: the FP4/FP6 max_version export error message. - tensor.c: the ScaledTensor `tp_doc`. - matrix.c: the euclideans_symmetric signature docstring tail. Rendered docstrings and error text are byte-identical; verified at runtime.
…esh kernels rmsd/kabsch/umeyama divided by the point count without guarding the empty case. kabsch and umeyama formed `inv_n = 1/n`, so an empty cloud produced Inf, then 0*Inf = NaN centroids, a NaN/garbage rotation out of the SVD, and a NaN metric; rmsd only accidentally returned 0 in serial (its `msd > 0 ? sqrt : 0` ternary) while every SIMD rmsd computed sqrt(0/0) = NaN. The serial oracle and the SIMD backends therefore disagreed on degenerate input. Every kernel now early-returns the neutral identity for n==0 — centroids 0, rotation = I, scale = 1, metric = 0 — matching the value rmsd already pre-fills. The guard is added to the serial macros, all SIMD backends (haswell, skylake, genoa, neon, neonbfdot, neonfhm, rvv, v128relaxed), and the C++ scalar fallbacks in mesh.hpp (the test oracle and the path used for non-native element types). The 6 rvv f16/bf16 kernels delegate to serial and inherit the guard unchanged. test_mesh.cpp exercises n==0 once per kernel against the serial reference. Runtime-verified: x86 (serial/haswell/skylake/genoa), ARM neon/neonbfdot/neonfhm under qemu-aarch64, and v128relaxed under wasmer all agree on the identity; rvv syntax-checked under clang.
`nk_angular_i4_icelake` reconstructs the signed dot product from the unsigned DPBUSD result with the correction `DPBUSD - 8*(∑ax + ∑bx) + 64*n`, computed in i32. Both `64 * (i32)n` and `8 * (i32)(∑ax + ∑bx)` overflow a 32-bit int around n ≈ 2^25 dimensions, corrupting the cosine for very high-dimensional i4 inputs. Promote the correction to i64 (the `∑` reductions are already i64); the final cast to f32 is unchanged. The i32 DPBUSD accumulators are left as-is per scope. Verified: angular_i4_icelake matches angular_i4_serial bit-for-tolerance on the default dimension sweep (native AVX-512 VNNI); no regression for normal n.
wasmtime / wasmtime-wasi 43 -> 44 closes GHSA-2r75-cxrj-cmph (WASI path_open(TRUNCATE) bypassing FilePerms::WRITE). mathjs 14 -> 15 closes GHSA-29qv-4j9f-fjw5 and GHSA-jvff-x2qm-6286 (prototype pollution). Both are dev-dependencies (WASM test harness / JS benchmark reference), so no library-facing change.
MaxSim scoring on Sapphire Rapids runs AMX tile instructions, which fault with SIGILL on Linux until the thread requests XTILEDATA state via `capabilities::configure_thread()` (arch_prctl ARCH_REQ_XCOMP_PERM). Every other family's tests already call it; the MaxSim scoring test did not, so it crashed the whole test binary. Call it in the test and show it in the module example.
The Emscripten branch set nk_shared's include directory to a raw source path, so CMake's generate step failed once a consumer (nk_shared_test, the sub-binding) pulled in its INTERFACE_INCLUDE_DIRECTORIES: a source-prefixed interface include must use a generator expression. Mirror the native branch's $<BUILD_INTERFACE:...>/$<INSTALL_INTERFACE:include> form. Fixes the WASM, WASM64, and Browser CI build jobs.
ashvardanian
force-pushed
the
main-v8-block-scaling
branch
from
July 20, 2026 16:39
6b4c047 to
a9ccd3f
Compare
Only the `first_task == 0` window writes the shared header and payload-offsets directory; every other window reads that directory to locate its plane. The parallel packer fanned all windows out with `for_n_dynamic`, so a worker could read the directory before task 0 populated it and scatter its plane to a garbage offset — a timing-dependent corruption that surfaced as an intermittent "parallel pack must match serial pack" failure on contended CI. Write task 0 serially first, then fan out the remaining windows over the now-populated directory. Also zero the blob before packing in both the serial and parallel paths so alignment padding is deterministic and packs are byte-reproducible.
The maxsim/attention naming pass renamed the attention_pack Python keyword head_dim to depth; the test call sites still passed head_dim=, which the extension now rejects. Update the four calls in test_attention.py.
Add validate_attention_views to validate the keys and values token views together and return their shared geometry once, replacing the asymmetric per-view checks duplicated in the serial and parallel pack paths, and flatten those paths with guard clauses. Extract the repeated packed-matrix/vector/tensor deallocation into tensor::dealloc_aligned, and invert the Drop and broadcast-window bodies to early returns. Rename the query-side locals to query_head_count to distinguish them from the stored KV head count.
The attention family (attention_pack, attention_packed, AttentionPackedMatrix) is implemented, registered, and tested, but was missing from the type stub and the README. Since the package ships py.typed, mypy/pyright flagged those symbols as unknown. Add the AttentionPackedMatrix stub (with the .heads/.depth/.segments/ .tokens/.shape getters), the two module functions, and a README example.
Add the kld/jsd f16+bf16 and jaccard u32 variants to the Emscripten EXPORTED_FUNCTIONS list, which had their siblings but not these. Drop the unused `architecture` matrix axis in test-python.yml (it was never consumed, doubling every Python job), remove the no-op `git submodule update` steps from the release workflows (there is no .gitmodules), and sync prerelease.yml's dry-run version list with release.yml's optionalDependencies lines.
ashvardanian
force-pushed
the
main-v8-block-scaling
branch
from
July 20, 2026 21:23
fe6ee9a to
e7fdddf
Compare
Each pack now zeros the metadata and alignment tails it owns at the source — the attention directory tail, the dots header `reserved[]` words across cross/AMX/SME/RVV, and the maxsim inter-region 64-byte slack — so packing the same input into a 0x00- or 0xFF-prefilled buffer yields byte-identical blobs. With the kernels hermetic, the Rust containers no longer pre-zero the whole buffer before packing, which spares the attention decode loop a full-blob memset on every KV-cache refill. New `pack_is_hermetic` tests across all three families guard the invariant.
`PackedMatrix.shape` (width, depth) and `MaxSimPackedMatrix.shape` (vectors, depth) were missing from the stub, though both are exposed at runtime — this completes the packed-container stubs to match.
Bump actions/checkout and actions/setup-python from v6 to v7 across every workflow; upload-artifact (v7) and download-artifact (v8) are already current. Add a concurrency group to the Pre-Release workflow — the one workflow that lacked one — so a new push to a branch cancels its superseded wheel-build runs, matching the test workflows.
ashvardanian
force-pushed
the
main-v8-block-scaling
branch
from
July 20, 2026 22:06
e7fdddf to
4d3b6ab
Compare
…ble-only Factor the byte blob every packed container owns — data, size, capacity, and allocator — into a single `PackedBuffer<Alloc>` in the tensor module, held as a private field the way `Vec` holds a `RawVec`. The dots, maxsim, and attention containers each embed one and delegate allocation, reuse, and teardown, replacing three copies of the grow/Drop/clone machinery. Drop every panicking convenience wrapper — `pack`, `pack_transposed`, `score`, `attention`, `attention_parallel` — leaving only the `try_*` forms. Packing is now always the typed constructor `XxxPackedMatrix:: try_pack`, which names the target dtype layout; the anomalous `MaxSimPackedOps` extension trait, which packed a bare tensor through `.try_pack()`, is removed so a tensor is never ambiguously "packed" for one family. Extension traits remain only for compute on an unpacked query. Attention gains parity with the other families: an instance `shape()`, a `from_packed_bytes_in` adopter, and the static header reader renamed `shape` to `peek_shape` so it no longer collides with the instance method.
Add `try_reserve` to the dots, maxsim, and attention containers, backed by a `PackedBuffer::try_reserve` that pre-grows the allocation while preserving the live bytes. A decode or serving loop can reserve for the largest expected geometry once, then repack every step with no further allocation and a stable pointer — the reuse the hermetic packs enable. Tests pin the contract: after reserving, capacity and the base pointer stay fixed across repacks of smaller inputs.
Adopt the bare `vectors` vocabulary for the maxsim packed matrix — the field, the `pack_size`/`try_reserve` parameters, the new `vectors()` getter, and the FFI/trait count parameters — replacing `vector_count`. The data-pointer parameter stays `data`, so `vectors` unambiguously names the count. Mechanical; no behavior change.
Parallelize dots packing by splitting the vector dimension into contiguous windows. Every nk_dots_pack_* kernel now takes a [columns_begin, columns_end) range and packs the tiles whose start column lands in its window, writing the shared header only from the window that begins at zero. This breaks the ABI of all dots pack kernels and adds DotsPackedMatrix::try_pack_parallel_into, which fans the work over a thread pool with for_slices. Attention's pack and compute windows move from first_task/task_count, with the task_count == 0 means all sentinel, to a half-open begin/end range, and the punned key_value_packed argument is spelled out in full. Maxsim packing stays serial and gains its own dispatch type, decoupled from the now-windowed dots pack type.
Add the archival DOI 10.5281/zenodo.21480520 to CITATION.cff and a Citation section to the README with a ready-to-use BibTeX entry, mirroring the StringZilla and ForkUnion layout. Align the CITATION title with the shared SDK description so the citation and package manifests read identically, and expand the abstract to cover the block-scaled MXFP & NVFP4 formats, scaled dot-product attention, and late-interaction MaxSim retrieval introduced for v8.
The fixes were already here as their originals, so the merge just adopts 7.7.1 across the manifests. Conflicts kept v8's side: rust-version 1.73 and the centered_norm_sq_a naming.
`nk_dot_u1` has shipped SIMD paths for NEON, Haswell, Icelake, RVV, Zvbb, WASM v128 and PowerVSX, and the C dispatcher has bound it to `nk_kernel_dot_k` for a while, yet the Rust surface never declared it. `u1x8` was the one packed type without `Dot`, so bit-packed embeddings could reach Hamming and Jaccard but not the raw intersection count that a ranking pass wants. Declare the kernel, implement `Dot` and `VDot` over it, and cover it with a test against hand-counted overlaps. `Dots` for `u1x8` was already implemented but never exercised — the batched checks enumerated every type down to `u4x2` and stopped — so extend those too. The transposed variant stays excluded because its `from_f32(2.0)` fill saturates to all-bits for a 1-bit type, the same reason `i4x2` and `u4x2` sit it out. Two feature gates disagreed with their users and warned in configurations CI does not usually build. `compute_thread_rows` was gated on `std` while all five call sites are `parallel`, and the `sets.rs` test module imported `align_depth` and `DIMS` unconditionally though only `parallel` code reads them. Every feature combination now builds warning-free.
`defer nk.ConfigureThread()()` needed a paragraph to explain that the first
`()` runs eagerly and only the second is deferred — and the production pool in
matrix.go never used it, preferring the explicit lock / defer-unlock form.
Advertise the same assign-and-defer idiom the code already follows:
unlock := nk.ConfigureThread()
defer unlock()
Also correct the pool comment: workers inline LockOSThread and the SIMD
configure call rather than calling the exported ConfigureThread.
`Tensor::<f32>::try_full(&[1 << 62, 1], 1.0)` segfaulted from entirely safe code. The element-count-to-byte-count multiply wrapped to zero in release, `Layout::from_size_align(0, 64)` succeeded, the allocator handed back a dangling pointer, and the constructor then wrote 2^62 elements through it. The shape product wrapped the same way, leaving the stored shape and the derived element count disagreeing so the stride arithmetic indexed outside the allocation. Route every owned allocation through `layout_for` / `layout_for_bytes` and every shape through `shape_product`, so the multiply is checked once and an unsatisfiable request reports `AllocationFailed` instead of wrapping into a small one. This also collapses the thirteen open-coded `from_size_align` sites, naming the `SIMD_ALIGNMENT` contract in a single place, and removes the `.unwrap()` that could panic while rebuilding a layout in `Vector::try_reserve`.
The in-house `unsafe trait Allocator` could not express an allocator passed by reference. A blanket bridge from `allocator_api2::alloc::Allocator` was the only way to reach the ecosystem, and being blanket over all `A` it foreclosed `impl Allocator for &A` forever — so an arena worked by value but not by `&arena`, which is precisely the escape a non-`Clone` arena needs against the twenty-one `Allocator + Clone` bounds. Adopt allocator-api2's trait as the bound everywhere and drop the in-house one. `Global` stays a NumKong type, now implementing that trait: taking allocator-api2's own would mean enabling its `alloc` feature, which adds a `Box`/`Vec`/`RawVec` reimplementation an order of magnitude larger than the 473-line trait and forces the `alloc` crate onto the heapless targets `--no-default-features` exists to serve. `allocate` now reports the block's actual size, so the slack a pool or arena already handed over is recorded as capacity instead of discarded. Freeing a tensor no longer calls `drop_in_place`, which was a guaranteed no-op under `StorageElement: Copy`, and the seven `extern crate alloc` declarations in files with no `alloc::` reference at all are gone. Two pre-commit gates were wrong and blocked this. The rustfmt gate ran stable `cargo fmt`, which silently drops the `unstable_features` and `fn_single_line` settings in `rustfmt.toml` and so reflows every file the repository formats with nightly. The sentence-spacing gate read the `!` of a `//!` doc comment as a sentence terminator, flagging every indented capitalised line in a doc example.
`sparse_dot` took four slices but passed only two lengths to the kernel, so `a_indices.len()` bounded the reads of `a_weights` as well. Nothing related the two, and the kernel reads `a_weights[i]` for every `i` below that length: eight indices against a one-element weight slice read seven `f32` past the end, from a plain safe call with no `unsafe` at the call site. A sparse vector carries one weight per index, so a shorter weights slice describes no vector at all and now yields the default instead.
Four inherent `Tensor::try_*_packed` methods pasted the body of `validate_packed_input` verbatim, while the very files holding them already imported that helper and called it from their trait-level counterparts. So each of these operations validated through two independent code paths, free to drift apart in which error variant they reported. Seventy-two lines go, and the two paths become one.
Each family repeated its whole operand-validation preamble once per scalar type — three identical copies for `ReduceRmsNorm`, and three for `EachSwiglu` differing only in the cast applied to the optional `up` pointer. Both now resolve their geometry through one function returning a plan, with the cast moved to the FFI boundary where a representation change belongs. The rmsnorm ordering is preserved deliberately rather than by luck: the empty-input return still precedes the gamma-length check, so a zero-row tensor never validates gamma.
Follows through on adopting the allocator-api2 trait, and closes the out-of-bounds accesses that reviewing that work turned up. `Global` gained `allocate_zeroed`, so a zero fill reaches `alloc_zeroed` and the lazily-zeroed pages the kernel already has rather than faulting the whole buffer in; a fill whose bytes all agree is now one `memset`. Reading the fill value's bytes rather than asking the type keeps the mini-floats honest, since `Ue8m0`'s all-zero encoding is 2^-127 and not zero. The three reserve paths that preserve their contents delegate to `Allocator::grow`, letting an arena extend a block in place; the packed buffer's grow keeps its explicit allocate-and-free, because packing overwrites and a copy there is waste. Attention's per-segment lengths were built in a `Vec` on the global heap, escaping the very allocator the container is parameterized over, in the pack path a decode loop runs every step. They now live on that allocator and are reused across packs, which also leaves `Global` as the crate's only remaining user of the `alloc` crate. The caller-supplied form of those lengths is gone: it was checked for count but never for value, while the kernel walks `lengths[i]` rows from `offsets[i]` knowing nothing of the tensor's height, so two tokens and a claimed hundred thousand read far past the allocation. They are derived from the offsets, which is all they ever were. Sub-byte scalars pack several logical dimensions into one storage value, and `try_flat`, `try_coords`, `row`, and slicing all bounds-checked against the logical count while offsetting by strides that span whole storage values. A 64-bit `u1x8` tensor occupies eight bytes, so index 63 passed the check and addressed byte 63; the `_mut` forms made that a write. `StorageElement` now locates a dimension in one place, matching how the C++ layer splits it, and `DimLocation` owns the read-modify-write that must preserve the neighbours sharing a storage value. Slicing keeps the representable case — taking the packed axis whole, which is what a leading-axis row slice does — and refuses to narrow it. Assembling a block-scaled tensor from parts validated nothing, though decoding sizes the scale reads from the elements shape alone; both parts must now agree. A rank past the fixed scratch buffers reports instead of panicking, a zero-width packed view is rejected where the geometry is established rather than dividing by zero later in the query path, and `sum` no longer reports a failed reduction as a zero tensor. The two attention pack and query preambles were identical, and the pack one mutated four geometry fields, so a fifth would have silently gone stale in the parallel path. Both are shared now. Tensor views and spans declare the `Send`/`Sync` their vector counterparts already did — a view is as shareable as the borrow it wraps — which the block-scaled views inherit. The raw-parts constructors are reconciled: one per container, and the allocator-taking ones carry the `_in` suffix every other such constructor does.
One vocabulary across the bindings, replacing names inherited from the BLAS `A · B = C` spelling and from single-letter FFI parameters. The result matrix was `c` in the dots, sets, and spatials families, `out` in the tensor module, and `output` in attention — one concept under three names, none of them discoverable. It is `output` everywhere now. The unpacked operand scored against a packed one becomes `queries`, matching the retrieval vocabulary maxsim and attention already use, and its companion stride follows. The matrix handed to a pack entry point is `matrix`, the packed right-hand operand `packed_right`, and `n_vectors` becomes `vector_count`, the shape the crate's other four count nouns already take. MaxSim's `q`/`d` become `queries`/`documents`, and the scalar dot products take `first`/`second`. Type parameters lose their initials: attention's `KIn`/`VIn`/`QIn` are `KeysTensor`/`ValuesTensor`/`QueriesTensor`, and the tuple-indexing arms `A0..A7` are `Axis0..Axis7`. Inside the slice machinery `s`, `st`, `nd`, and `off` become `sliced_shape`, `sliced_strides`, `sliced_ndim`, and `byte_offset`. Behavior is untouched: the tests pass unchanged apart from the identifiers themselves. Note that `cargo check` cannot police a rename at this boundary — `packed_b` is a substring of the `nk_dots_packed_bf16` symbol, and renaming it blindly produced a library that compiled and failed only at link time. The test link is what catches that class of mistake.
`Break: Window the dots packer` added `columns_begin`/`columns_end` to `nk_dots_pack_*` (and the punned typedef) but only updated the C core and the Rust caller, so the C++ bench and the Go, Swift, JavaScript, and Python bindings called the seven-parameter packer with five arguments and failed to compile — taking Swift, JS, Go, C, and every Python leg red on CI. Pass the full range `(0, width)` at each dots call site, which packs the whole matrix exactly as before the window was introduced. Document the two parameters on the `nk_dots_pack_bf16` doc block the others copy. maxsim does NOT share the dots packer: it has its own five-parameter `nk_maxsim_pack_*` family with no column window. Its Python and C++ callers were mistyped against the dots typedef, which compiled only while both were five parameters; the window split turned that into a real arity mismatch. Type them against `nk_maxsim_pack_punned_t` / `maxsim_pack_kernel_t` and keep the five-arg calls, rather than passing a window the packer ignores. Verified on Windows: the Python extension rebuilds and its packed, matmul, and maxsim tests pass; Go, Swift, and JavaScript compile; and the C++ bench translation units pass a C++20 syntax check.
The windowed attention pack and compute take a half-open `[begin, end)` task range over the segment × head grid. The Python bindings drove the parallel loop with `(task_index, 1)`, which is empty for every `task_index >= 1` — so only task 0 was ever packed and scored, and the remaining planes stayed at whatever the fresh allocation held. Attention over that uninitialized memory produced garbage (~1e34), failing all 24 bf16/e4m3 `test_attention_packed` cases at both thread counts. The C++ bench passed `(0, 0)` to both, an empty range that packed and measured nothing. Pass `(task_index, task_index + 1)` per iteration in Python and `(0, head_count)` for the single-shot bench, so each call covers exactly the tasks it should. Verified on Windows: all 28 attention tests pass, and the bench translation units pass a C++20 syntax check.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Uh oh!
There was an error while loading. Please reload this page.