BitNet Model - Complete Inference Pipeline

Implements the complete BitNet transformer model for text generation.

Architecture

Input: token_ids [batch, seq_len]
  │
  ▼
[Embedding Layer]  ← TQ2_0 quantized
  │ [batch, seq_len, dim]
  │
  ├─► [Transformer Block 0]
  │       ├─► Attention + FFN
  │       └─► Residual connections
  │
  ├─► [Transformer Block 1]
  │       ...
  │
  ├─► [Transformer Block 31]
  │
  ▼ [batch, seq_len, dim]
[Final RMSNorm]
  │
  ▼
[LM Head] ← Project to vocabulary
  │ [batch, seq_len, vocab_size]
  │
  ▼
[Logits] → Sampling → Next token

Model Configurations

BitNet-3B

Vocabulary: 50,000 tokens
Embedding: 2560 dimensions
Layers: 32 transformer blocks
Heads: 32 attention heads
FFN hidden: 10,240 (4× expansion)
Context: 2048 tokens max

BitNet-1.3B (smaller variant)

Vocabulary: 50,000 tokens
Embedding: 2048 dimensions
Layers: 24 transformer blocks
Heads: 16 attention heads
FFN hidden: 8,192
Context: 2048 tokens max

Memory Footprint (BitNet-3B with TQ2_0)

Model weights:

Embedding:     50000 × 2560 × 0.25 = 32 MB
32 layers:     32 × 26 MB = 832 MB
Final norm:    2560 × 4 = 10 KB
LM head:       2560 × 50000 × 0.25 = 32 MB
─────────────────────────────────────────
Total:         ~896 MB

Compare to Float32: ~12 GB (13.4x savings!)

Activations (seq_len=2048):

Per layer peak: ~512 MB (attention scores)
With buffer reuse: ~512 MB total

Total inference memory: ~1.4 GB

Performance (BitNet-3B on A100)

Single token latency:

Embedding:          0.01 ms
32 transformer layers: 870 ms (27 ms/layer avg)
Final norm:         0.01 ms
LM head:            13 ms
Sampling:           0.1 ms
─────────────────────────────────────────
Total:              ~883 ms/token

Throughput:         1.13 tokens/sec (base)

With optimizations:

Flash Attention:    440 ms (2× speedup)
Tiled matmul:       290 ms (1.5× additional)
Kernel fusion:      223 ms (1.3× additional)
─────────────────────────────────────────
Target:             ~200-250 ms/token (4-5 tokens/sec)

Text Generation

Greedy Decoding

while len(tokens) < max_tokens:
  logits = model(tokens)
  next_token = argmax(logits[-1])
  tokens.append(next_token)

Top-k Sampling

logits = model(tokens)
top_k_logits, top_k_indices = topk(logits, k)
probs = softmax(top_k_logits / temperature)
next_token = sample(top_k_indices, probs)

Nucleus (Top-p) Sampling

logits = model(tokens)
sorted_logits, sorted_indices = sort(logits, descending=True)
cumulative_probs = cumsum(softmax(sorted_logits))
nucleus = sorted_indices[cumulative_probs <= p]
next_token = sample(nucleus)

References

• BitNet: https://arxiv.org/abs/2402.17764
• LLaMA: https://github.com/facebookresearch/llama
• llama.cpp: Main inference loop in llama.cpp

Configuration

structure Hesper.Models.BitNet.Config : Type

Model configuration

• vocabSize : Nat
• dim : Nat
• numLayers : Nat
• numHeads : Nat
• numKVHeads : Nat
• headDim : Nat
• ffnDim : Nat
• maxSeqLen : Nat
• ropeBase : Float
• temperature : Float
• topK : Nat
• topP : Float

instance Hesper.Models.BitNet.instReprConfig : Repr Config

Equations

• Hesper.Models.BitNet.instReprConfig = { reprPrec := Hesper.Models.BitNet.instReprConfig.repr }

def Hesper.Models.BitNet.instReprConfig.repr : Config → Nat → Std.Format

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Models.BitNet.Config.kvDim (cfg : Config) : Nat

KV dimension (numKVHeads * headDim)

Equations

• cfg.kvDim = cfg.numKVHeads * cfg.headDim

def Hesper.Models.BitNet.Config.bitnet2B : Config

Predefined configurations

Equations

• Hesper.Models.BitNet.Config.bitnet2B = { }

def Hesper.Models.BitNet.Config.bitnet3B : Config

Legacy alias

Equations

• Hesper.Models.BitNet.Config.bitnet3B = Hesper.Models.BitNet.Config.bitnet2B

def Hesper.Models.BitNet.Config.bitnet1_3B : Config

Equations

• Hesper.Models.BitNet.Config.bitnet1_3B = { vocabSize := 50000, dim := 2048, numLayers := 24, numHeads := 16, numKVHeads := 16, ffnDim := 8192 }

Model Structure

structure Hesper.Models.BitNet.BitNetModel (BufT : Type) (CacheT KernelT : Type := Unit) : Type

Complete BitNet model

• config : Config
• embedding : Layers.Embedding.Embedding BufT
• layers : Array (Layers.TransformerBlock.TransformerBlock BufT CacheT KernelT)
• finalNorm : Layers.RMSNorm.RMSNorm BufT CacheT

def Hesper.Models.BitNet.resetPreparedDispatches {β : Type} [GPUBackend β] (model : BitNetModel (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β)) : IO Unit

Reset all PreparedDispatch caches in the model. Must be called when buffer bindings change (e.g., new generate call).

Equations

• One or more equations did not get rendered due to their size.

Model Creation

def Hesper.Models.BitNet.create {β : Type} [GPUBackend β] (_ctx : β) (config : Config) : IO (BitNetModel (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β))

Create BitNet model (placeholder - needs GGUF integration)

@param device WebGPU device @param config Model configuration

Equations

• One or more equations did not get rendered due to their size.

Forward Pass

def Hesper.Models.BitNet.forward {β : Type} [GPUBackend β] (ctx : β) (model : BitNetModel (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β)) (tokenIdsBuf outputBuf : GPUBackend.Buf β) (batchSize seqLen : Nat) : IO Unit

Execute forward pass through entire model

@param device WebGPU device @param model BitNet model @param tokenIdsBuf Input token IDs [batch, seq_len] @param outputBuf Output logits [batch, seq_len, vocab_size] @param batchSize Batch size @param seqLen Sequence length

Equations

• One or more equations did not get rendered due to their size.

KV Cache State

structure Hesper.Models.BitNet.KVCacheState (BufT : Type) (CacheT : Type := Unit) : Type

Full KV cache state for incremental inference

• kvCaches : Array (Layers.Attention.KVCache BufT CacheT)
• fusedRefs : Array (Layers.TransformerBlock.FusedLayerRefs CacheT)
• layerBufs : Layers.TransformerBlock.CachedLayerBuffers BufT CacheT
• buf1 : BufT
• buf2 : BufT
• logitsBuf : BufT
• argmaxBuf : BufT
• tokenBuf : BufT
• penaltyTokensBuf : BufT
• penaltyParamsBuf : BufT

def Hesper.Models.BitNet.createKVCacheState {β : Type} [GPUBackend β] (ctx : β) (model : BitNetModel (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β)) : IO (KVCacheState (GPUBackend.Buf β) (GPUBackend.CachedDispatch β))

Create KV cache state for the model

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Models.BitNet.forwardSingleToken {β : Type} [GPUBackend β] (ctx : β) (model : BitNetModel (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β)) (tokenId pos : Nat) (cacheState : KVCacheState (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) : IO Unit

Run single-token forward pass with KV cache. Processes one token at position pos, using cached K/V from past tokens. Returns logits in cacheState.logitsBuf.

Equations

• One or more equations did not get rendered due to their size.

GPU Argmax

def Hesper.Models.BitNet.argmaxKernel (vocabSize : Nat) (workgroupSize : Nat := 256) : WGSL.Monad.ShaderM Unit

GPU argmax kernel: find index of maximum value using parallel reduction. Uses 1 workgroup of workgroupSize threads. Each thread scans a strided portion of the input, then shared memory reduction finds the global max. Output: single u32 (token index of max value).

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Models.BitNet.gpuArgmax {β : Type} [GPUBackend β] (ctx : β) (logitsBuf argmaxBuf : GPUBackend.Buf β) (vocabSize : Nat) : IO Nat

Execute GPU argmax on logits buffer. Returns the token index with maximum logit value. Downloads only 4 bytes instead of vocabSize × 4 bytes.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Models.BitNet.repetitionPenaltyKernel (maxTokens vocabSize : Nat) (penalty : Float) : WGSL.Monad.ShaderM Unit

GPU kernel: apply repetition penalty to logits in-place. One thread per previous token. Reads token ID, applies penalty to that logit. penalty is baked into the shader as a literal (changes rarely).

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Models.BitNet.appendPenaltyToken {β : Type} [GPUBackend β] (ctx : β) (cacheState : KVCacheState (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) (tokenId offset : Nat) : IO Unit

Append a single token ID to the penalty tokens buffer at the given offset.

Equations

• One or more equations did not get rendered due to their size.

def Hesper.Models.BitNet.gpuArgmaxWithPenalty {β : Type} [GPUBackend β] (ctx : β) (cacheState : KVCacheState (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) (vocabSize maxSeqLen numTokens : Nat) (penalty : Float) : IO Nat

Apply repetition penalty on GPU, then run GPU argmax. Assumes token IDs are already uploaded incrementally. Just updates numTokens param, runs penalty + argmax batched in a single GPU submit, downloads 4 bytes.

Equations

• One or more equations did not get rendered due to their size.

Text Generation

def Hesper.Models.BitNet.generate {β : Type} [GPUBackend β] (ctx : β) (model : BitNetModel (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β)) (promptTokens : Array Nat) (maxTokens : Nat) (strategy : Inference.Sampling.Strategy := Inference.Sampling.Strategy.Greedy) (eosToken : Option Nat := none) (showStats : Bool := false) (repetitionPenalty : Float := 1.1) : IO (Array Nat)

Generate text using greedy decoding

@param device WebGPU device @param model BitNet model @param promptTokens Initial prompt tokens @param maxTokens Maximum tokens to generate @return Generated token sequence

def Hesper.Models.BitNet.generateNaive (device : WebGPU.Device) (model : BitNetModel WebGPU.Buffer WGSL.Execute.PreparedDispatch WGSL.Execute.CompiledKernel) (promptTokens : Array Nat) (maxTokens : Nat) (strategy : Inference.Sampling.Strategy := Inference.Sampling.Strategy.Greedy) (eosToken : Option Nat := none) : IO (Array Nat)

Generate text WITHOUT KV cache (naive quadratic approach). Kept for validation/comparison with cached generate.

Equations

• One or more equations did not get rendered due to their size.

Sampling Strategies

def Hesper.Models.BitNet.sampleGreedy (logits : Array Float) : Nat

Greedy sampling: select token with highest probability

@param logits Logit scores [vocab_size] @return Selected token ID

Equations

• Hesper.Models.BitNet.sampleGreedy logits = Hesper.Inference.Sampling.sampleGreedy logits

def Hesper.Models.BitNet.sampleTopK (logits : Array Float) (k : Nat) (temperature : Float) (rng : Inference.Sampling.RNG) : Nat × Inference.Sampling.RNG

Top-k sampling: sample from k tokens with highest probability

@param logits Logit scores [vocab_size] @param k Number of top candidates @param temperature Sampling temperature (higher = more random) @param rng Random number generator @return (Selected token ID, new RNG)

Equations

• Hesper.Models.BitNet.sampleTopK logits k temperature rng = Hesper.Inference.Sampling.sampleWithRNG logits (Hesper.Inference.Sampling.Strategy.TopK k temperature) rng

def Hesper.Models.BitNet.sampleNucleus (logits : Array Float) (p temperature : Float) (rng : Inference.Sampling.RNG) : Nat × Inference.Sampling.RNG

Nucleus (top-p) sampling: sample from smallest set with cumulative probability >= p

@param logits Logit scores [vocab_size] @param p Cumulative probability threshold (typically 0.9) @param temperature Sampling temperature @param rng Random number generator @return (Selected token ID, new RNG)

Equations

• Hesper.Models.BitNet.sampleNucleus logits p temperature rng = Hesper.Inference.Sampling.sampleWithRNG logits (Hesper.Inference.Sampling.Strategy.Nucleus p temperature) rng

GGUF Integration

def Hesper.Models.BitNet.extractConfig (gguf : GGUF.GGUFFile) : IO Config

Extract model configuration from GGUF metadata

Reads metadata keys to determine model architecture.

@param gguf Parsed GGUF file @return Model configuration

Equations

• Hesper.Models.BitNet.extractConfig gguf = do IO.println "Using default BitNet-2B configuration" pure Hesper.Models.BitNet.Config.bitnet2B

def Hesper.Models.BitNet.fromGGUFObject {β : Type} [GPUBackend β] (ctx : β) (gguf : GGUF.GGUFFile) (config : Option Config := none) : IO (BitNetModel (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β))

Load BitNet model from pre-loaded GGUF object

@param device WebGPU device @param gguf Already parsed GGUF file object @param config Optional configuration (uses defaults if not provided) @return Loaded model

NOTE: This function prevents premature GC of the GGUF object during loading

Equations

• One or more equations did not get rendered due to their size.

Utilities

def Hesper.Models.BitNet.fromGGUF {β : Type} [GPUBackend β] (ctx : β) (ggufPath : String) (config : Option Config := none) : IO (BitNetModel (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β))

Load BitNet model from GGUF file path (convenience wrapper)

@param device WebGPU device @param ggufPath Path to GGUF file @param config Optional configuration (uses defaults if not provided) @return Loaded model

NOTE: This loads the GGUF file and calls fromGGUFObject

Equations

• Hesper.Models.BitNet.fromGGUF ctx ggufPath config = do let gguf ← Hesper.GGUF.loadGGUF ggufPath Hesper.Models.BitNet.fromGGUFObject ctx gguf config

def Hesper.Models.BitNet.printStats {BufT CacheT KernelT : Type} (model : BitNetModel BufT CacheT KernelT) : IO Unit

Print model statistics

@param model BitNet model

Equations

• One or more equations did not get rendered due to their size.