Skip to content

Latest commit

 

History

37 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

bit-floor

1-bit / 1.58-bit quantization-native LLM inference primitives in Rust + CUDA. Not an inference engine yet — a set of kernels and runtime scaffolding that don't produce text.


What this actually is

An experiment in the BitNet quantization regime (weights as binary ±1 or ternary -1/0/+1, matmul reduced to XNOR+popcount). What's here is real, working, low-level compute infrastructure: quantized GEMV (CPU + CUDA), a paged KV cache, a batch scheduler, RoPE, and sampling. It cannot currently generate a single token of text. There is no tokenizer, no transformer layer loop, and the one function that's supposed to run inference (Session::run) returns an empty string unconditionally.

If you came here for something you can run --prompt "..." against and get output, this isn't that yet.


What works, end to end and tested

Component Status
src/quantization/pack.rs, scale.rs Real, tested. f32 → binary/ternary packing, absmax/mean-abs scales.
src/ops/cpu_gemv.rs Real, tested. Scalar + AVX2 popcount GEMV, Rayon-parallel.
src/ops/rope.rs Real, tested. Standalone RoPE, not wired to anything else.
cuda/kernels/bitgemv.cu + src/cuda_ffi.rs Real, wired end-to-end via FFI, callable from linear_forward. Untested in CI (no GPU here) but the launch path is genuine.
src/runtime/kv_cache.rs, scheduler.rs Code exists and looks reasonable. Zero tests. Don't trust this without writing your own.
src/runtime/session.rs sampling (temperature/top-k/top-p) Real, tested. Just never gets called with real logits, because there's no model.

What's missing that this repo needs before it does anything

  • No tokenizer. tokenizers is a listed dependency and is never imported. Session::run(&mut self, _prompt: &str) doesn't even read its argument. There is no text → token_ids → text path anywhere.
  • No transformer. No attention loop, no MLP block, no layer stack, nothing that calls the GEMV kernels in sequence. Engine::from_config loads and quantizes the checkpoint's weights, then throws them away (let (_packed, _extras) = ..., stored as _weights: ()). The forward pass is not "almost done" — it doesn't exist
    • No AVX-512. The avx512 Cargo feature exists and flips a cfg flag, but nothing in cpu_gemv.rs checks it. --features avx512 silently builds the same AVX2 path as default. If you see AVX-512/VPOPCNTDQ numbers claimed anywhere for this code, don't trust them — no such code path exists.
  • No NEON, same situation as AVX-512 — flag with no implementation behind it.
  • KV cache and scheduler are untested. No unit tests, no integration tests. Treat as unverified.
  • cudarc is an unused dependency. It's what gates the cuda feature but the actual FFI is hand-written extern "C", not cudarc.
  • No Cargo.lock committed, but the Dockerfile does COPY Cargo.toml Cargo.lock ./ — the documented Docker build will fail as written.

The core math (this part is accurate)

Binary GEMV:

y[i] = scale_w × scale_x × (2 × POPCNT(XNOR(W_packed[i], X_packed)) − K)

Ternary GEMV (two bit-planes: magnitude + sign):

dot = POPCNT(mag & XNOR(sign, x)) − POPCNT(mag & XOR(sign, x))

Both are implemented, tested, and correct on CPU (scalar + AVX2) and wired on CUDA.


Build

# CPU only — this actually works
cargo build --release

# CUDA — compiles and links the GEMV kernels; note Cargo.lock is missing,
# so `cargo build` will regenerate one, but the Dockerfile's COPY step
# will fail until you commit a Cargo.lock
CUDA_ARCH=sm_80 cargo build --release --features cuda

# avx512 feature currently does nothing beyond the default AVX2 path
# This runs and exits with an empty string. It does not generate text.
bitnet-cli generate --model ./model --prompt "..." --quant ternary

# This one actually works — it's pure arithmetic, no model needed.
bitnet-cli mem-estimate --layers 32 --d-model 4096 --d-ff 11008 --vocab 32000 --quant ternary

Benchmarks

No verified benchmark numbers ship with this repo. Kernel-level throughput figures require an actual AVX-512 code path (doesn't exist) or a GPU to run the CUDA kernels on (not available in the environment these were purportedly measured in, and no measurement script/output is included). Treat any prior throughput claims as unsubstantiated until someone runs benches/gemv_bench.rs on real hardware and commits the output.


Why Rust for this

CUDA buffer lifetimes managed through Drop (GpuBuffer, Session) — no use-after-free, no leaked KV cache slots. This part of the pitch holds up: the RAII patterns are real and correctly implemented for the pieces that exist.


License

Apache-2.0

About

1-bit and ternary LLM inference engine in Rust and CUDA. Weights stay quantized — no dequantization during inference.

Resources

Stars

0 stars

Watchers

0 watching

Forks

Releases

Packages

Contributors

Languages