Multi-Head Self-Attention

Implements the attention mechanism used in BitNet transformers.

Mathematical Definition

Self-attention allows each position to attend to all positions in the input:

Q = input @ W_q   # Query projection
K = input @ W_k   # Key projection
V = input @ W_v   # Value projection

scores = (Q @ K^T) / sqrt(d_k)
attn = softmax(scores)  # Attention weights
output = attn @ V
result = output @ W_o   # Output projection

Multi-Head Attention

Instead of single attention, split into H heads:

For each head h:
  Q_h = Q[:, h*d_h : (h+1)*d_h]
  K_h = K[:, h*d_h : (h+1)*d_h]
  V_h = V[:, h*d_h : (h+1)*d_h]

  attn_h = softmax(Q_h @ K_h^T / sqrt(d_h)) @ V_h

output = concat(attn_1, ..., attn_H) @ W_o

Why multiple heads?

• Each head learns different attention patterns
• Head 1: Local syntax (nearby words)
• Head 2: Long-range dependencies
• Head 3: Semantic relationships
• Etc.

BitNet Optimization

Standard Transformer:

• W_q, W_k, W_v, W_o: Float32 weights
• Memory: 4 × d² × 4 bytes per layer

BitNet:

• W_q, W_k, W_v, W_o: Ternary {-1, 0, 1} weights (TQ2_0)
• Memory: 4 × d² × 0.25 bytes per layer
• 16x memory savings
• Faster compute: additions instead of multiplications

Causal Masking

For autoregressive generation, prevent attending to future tokens:

mask[i,j] = 0   if j ≤ i  (can attend to past)
mask[i,j] = -∞  if j > i  (cannot attend to future)

attn = softmax(scores + mask)

After softmax, masked positions → 0 probability.

Performance

For BitNet-3B (hidden=2560, heads=32, seq=2048):

Per layer:
- Q,K,V projections: 3 × BitLinear(2560→2560) = 52.4 GFLOPS
- Attention scores: Q @ K^T = 107.4 GFLOPS (2048×80 @ 80×2048)
- Softmax: 0.1 GFLOPS (negligible)
- Attention output: attn @ V = 107.4 GFLOPS
- Output projection: BitLinear(2560→2560) = 17.5 GFLOPS

Total: ~285 GFLOPS per attention layer

References

• Attention is All You Need (Vaswani et al., 2017)
• LLaMA architecture: https://github.com/facebookresearch/llama
• llama.cpp: ggml-cuda/attention.cu
• Flash Attention: https://arxiv.org/abs/2205.14135

Configuration

structure Hesper.Layers.Attention.Config : Type

Multi-head attention configuration

• dim : Nat
• numHeads : Nat
• numKVHeads : Nat
• headDim : Nat
• maxSeqLen : Nat
• useCausalMask : Bool

instance Hesper.Layers.Attention.instReprConfig : Repr Config

Equations

• Hesper.Layers.Attention.instReprConfig = { reprPrec := Hesper.Layers.Attention.instReprConfig.repr }

def Hesper.Layers.Attention.instReprConfig.repr : Config → Nat → Std.Format

Equations

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

def Hesper.Layers.Attention.Config.effectiveKVHeads (c : Config) : Nat

Effective KV head count

Equations

• c.effectiveKVHeads = if (c.numKVHeads == 0) = true then c.numHeads else c.numKVHeads

def Hesper.Layers.Attention.Config.effectiveHeadDim (c : Config) : Nat

Effective head dimension

Equations

• c.effectiveHeadDim = if (c.headDim == 0) = true then c.dim / c.numHeads else c.headDim

def Hesper.Layers.Attention.Config.kvDim (c : Config) : Nat

KV dimension (numKVHeads * headDim)

Equations

• c.kvDim = c.effectiveKVHeads * c.effectiveHeadDim

def Hesper.Layers.Attention.Config.validate (config : Config) : Except String Unit

Validate configuration

Equations

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

Reshape Kernels for Multi-Head Attention

def Hesper.Layers.Attention.reshapeToHeadsKernel (batchSize seqLen inputHeads outputHeads headDim : Nat) : WGSL.Monad.ShaderM Unit

Reshape [batch, seq, inputHeads, headDim] → [batch, outputHeads, seq, headDim]

For Q: inputHeads = outputHeads = numHeads (simple transpose) For K/V with GQA: inputHeads = numKVHeads, outputHeads = numHeads (repeat KV heads)

Each output element at [b, oh, s, d]: kv_head = oh / headsPerKVHead input_offset = b * seq * inputHeads * headDim + s * inputHeads * headDim + kv_head * headDim + d

Equations

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

def Hesper.Layers.Attention.reshapeFromHeadsKernel (batchSize seqLen numHeads headDim : Nat) : WGSL.Monad.ShaderM Unit

Reshape [batch, heads, seq, headDim] → [batch, seq, heads, headDim]

Reverse of reshapeToHeads. Used after attention output, before output projection.

Equations

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

def Hesper.Layers.Attention.executeReshapeToHeads {β : Type} [GPUBackend β] (ctx : β) (inBuf outBuf : GPUBackend.Buf β) (batchSize seqLen inputHeads outputHeads headDim : Nat) : IO Unit

Execute reshape to multi-head layout

Equations

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

def Hesper.Layers.Attention.executeReshapeFromHeads {β : Type} [GPUBackend β] (ctx : β) (inBuf outBuf : GPUBackend.Buf β) (batchSize seqLen numHeads headDim : Nat) : IO Unit

Execute reshape from multi-head layout

Equations

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

Pre-allocated Buffers

structure Hesper.Layers.Attention.AttentionBuffers (BufT : Type) : Type

Pre-allocated buffers for attention forward pass. Avoids ~12 GPU buffer allocations per layer per token.

• qBuf : BufT
• kBuf : BufT
• vBuf : BufT
• scoresBuf : BufT
• attnBuf : BufT
• qRotBuf : BufT
• kRotBuf : BufT
• qHeadBuf : BufT
• kHeadBuf : BufT
• vHeadBuf : BufT
• reshapedOutBuf : BufT
• subNormBuf : BufT
• rmsTempBuf : BufT

def Hesper.Layers.Attention.createAttentionBuffers {β : Type} [GPUBackend β] (ctx : β) (config : Config) (batchSize seqLen : Nat) : IO (AttentionBuffers (GPUBackend.Buf β))

Create pre-allocated attention buffers for given dimensions

Equations

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

Layer Structure

structure Hesper.Layers.Attention.Attention (BufT : Type) (CacheT KernelT : Type := Unit) : Type

Multi-head self-attention layer

• config : Config
• wQ : BitLinear.BitLinear BufT CacheT KernelT
• wK : BitLinear.BitLinear BufT CacheT KernelT
• wV : BitLinear.BitLinear BufT CacheT KernelT
• wO : BitLinear.BitLinear BufT CacheT KernelT
• rope : RoPE.RoPE
• softmax : Softmax.Softmax

Layer Creation

def Hesper.Layers.Attention.create {β : Type} [GPUBackend β] (ctx : β) (config : Config) (wqData wkData wvData woData : ByteArray) (qScale kScale vScale oScale : Float) : IO (Attention (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β))

Create attention layer from GGUF tensors

@param device WebGPU device @param config Attention configuration @param wqData Q projection weights (TQ2_0 packed) @param wkData K projection weights (TQ2_0 packed) @param wvData V projection weights (TQ2_0 packed) @param woData Output projection weights (TQ2_0 packed) @param qScales Q projection scales (FP16) @param kScales K projection scales (FP16) @param vScales V projection scales (FP16) @param oScales Output projection scales (FP16)

Equations

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

Forward Pass

def Hesper.Layers.Attention.forward {β : Type} [GPUBackend β] (ctx : β) (layer : Attention (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β)) (inputBuf outputBuf : GPUBackend.Buf β) (batchSize seqLen : Nat) (subNorm : Option (RMSNorm.RMSNorm (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) := none) (preAllocBufs : Option (AttentionBuffers (GPUBackend.Buf β)) := none) (residualBuf : Option (GPUBackend.Buf β) := none) : IO Unit

Execute attention forward pass

Algorithm:

1. Project input to Q, K, V:
   Q = input @ W_q
   K = input @ W_k
   V = input @ W_v

2. Reshape to multi-head format:
   Q, K, V: [batch, seq, dim] → [batch, seq, heads, head_dim]
                              → [batch, heads, seq, head_dim]

3. Apply RoPE to Q and K:
   Q_rot = RoPE(Q)
   K_rot = RoPE(K)

4. Compute attention scores:
   scores = (Q_rot @ K_rot^T) / sqrt(head_dim)
   Shape: [batch, heads, seq, seq]

5. Apply softmax (with causal mask if enabled):
   attn_weights = softmax(scores)

6. Apply attention to values:
   attn_output = attn_weights @ V
   Shape: [batch, heads, seq, head_dim]

7. Reshape and project output:
   output = concat_heads(attn_output)
   result = output @ W_o

@param device WebGPU device @param layer Attention layer @param inputBuf Input tensor [batch, seq_len, dim] @param outputBuf Output tensor [batch, seq_len, dim] @param batchSize Batch size @param seqLen Sequence length

Equations

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

KV Cache Structures

structure Hesper.Layers.Attention.KVCache (BufT : Type) (CacheT : Type := Unit) : Type

KV cache for a single attention layer. Stores key and value vectors for all past positions.

• kBuf : BufT
• vBuf : BufT
• preparedCacheWriteK : IO.Ref (Option CacheT)
• preparedCacheWriteV : IO.Ref (Option CacheT)
• preparedCacheWriteKV : IO.Ref (Option CacheT)
• preparedScores : IO.Ref (Option CacheT)
• preparedApply : IO.Ref (Option CacheT)

def Hesper.Layers.Attention.createKVCache {β : Type} [GPUBackend β] (ctx : β) (config : Config) : IO (KVCache (GPUBackend.Buf β) (GPUBackend.CachedDispatch β))

Create KV cache for one attention layer

Equations

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

structure Hesper.Layers.Attention.CachedAttentionBuffers (BufT : Type) (CacheT : Type := Unit) : Type

Pre-allocated buffers for cached single-token attention

• qBuf : BufT
• kNewBuf : BufT
• vNewBuf : BufT
• qRotBuf : BufT
• kRotBuf : BufT
• scoresBuf : BufT
• attnBuf : BufT
• subNormBuf : BufT
• rmsTempBuf : BufT
• paramsBuf : BufT
• flashPartialBuf : BufT
• preparedSoftmax : IO.Ref (Option CacheT)
• preparedRopeQ : IO.Ref (Option CacheT)
• preparedRopeK : IO.Ref (Option CacheT)

def Hesper.Layers.Attention.createCachedAttentionBuffers {β : Type} [GPUBackend β] (ctx : β) (config : Config) : IO (CachedAttentionBuffers (GPUBackend.Buf β) (GPUBackend.CachedDispatch β))

Create pre-allocated buffers for cached attention

Equations

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

KV Cache Kernels

All cached attention kernels read pos and cacheLen from a params buffer instead of baking them as WGSL literals. This produces identical WGSL across tokens, enabling pipeline + bind group caching (97%+ hit rate).

Params buffer layout: [pos: u32, cacheLen: u32] (8 bytes)

def Hesper.Layers.Attention.cacheWriteKernel (numKVHeads maxSeqLen headDim kvDim : Nat) : WGSL.Monad.ShaderM Unit

Write new K or V data into cache at position read from params[0]. Input: [kvDim] = [numKVHeads * headDim] (flat) Cache: [numKVHeads, maxSeqLen, headDim] Params: [pos: u32, cacheLen: u32]

Equations

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

def Hesper.Layers.Attention.fusedRopeKAndCacheWriteKernel (numKVHeads maxSeqLen headDim kvDim : Nat) (ropeBase : Float) : WGSL.Monad.ShaderM Unit

Fused RoPE-on-K + KV cache write. Applies NeoX RoPE to the K tensor (using ropeWithFreqFactors's formula) and writes the rotated K plus the V tensor straight into the KV cache slot at position pos.

Dispatch: 1D over kvDim = numKVHeads * headDim threads. Each thread handles ONE (k_head, dim) element for V (plain copy) and a half of a rotation pair for K (the other half is read but only the local half written). The pair partner is at dim + halfDim if dim < halfDim.

Saves one dispatch per KV-bearing layer (RoPE-K + KV-write → 1 kernel). K is read from new_k (post-norm), V from new_v.

Equations

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

def Hesper.Layers.Attention.fusedRopeKAndCacheWriteKernelF16 (numKVHeads maxSeqLen headDim kvDim : Nat) (ropeBase : Float) : WGSL.Monad.ShaderM Unit

f16 variant of fusedRopeKAndCacheWriteKernel for the V11 attention path. Each thread processes ONE pair of dims (2 elements) and packs them into one u32 (pack2x16float) to write the f16 cache.

Grid: kvDim/2 threads. Each thread handles dim pair (2idx, 2idx+1) within its kvHead. The RoPE rotation needs the partner at dim ± halfDim (not dim ± 1), so each thread reads 4 source elements (xSelf for both dims of the pair, and xPair for both dims of the pair).

Equations

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

def Hesper.Layers.Attention.fusedRopeKAndCacheWriteBatchKernelF16 (numKVHeads maxSeqLen headDim seqLen : Nat) (ropeBase : Float) : WGSL.Monad.ShaderM Unit

Batched f16 RoPE-K + KV-write kernel for prefill.

Companion to fusedRopeKAndCacheWriteKernelF16 (single-position): processes seqLen query tokens at once and writes both the K cache (with NeoX RoPE rotation applied) and the V cache (plain copy) into the f16 packed half2 layout.

Inputs: new_k : f32[seqLen * kvDim] — K projections, column-major new_v : f32[seqLen * kvDim] — V projections, same layout params : u32[1] — params[0] = startPos freq_factors : f32[halfDim] Outputs: k_cache_f16 : u32[numKVHeads * maxSeqLen * (headDim/2)] v_cache_f16 : same shape

Grid: (seqLen * numKVHeads * (headDim/2)) threads, dispatch1D.

Cache slot for token col → pos = startPos + col.

Equations

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

def Hesper.Layers.Attention.kvWriteBatchKernelF16 (numKVHeads maxSeqLen headDim seqLen : Nat) : WGSL.Monad.ShaderM Unit

No-RoPE batched f16 KV-write. Used by SWA layers in forwardBlock / forwardPrefillBatch: K is already RoPE'd into new_k_roped upstream (because the SWA layers don't have freq_factors). Each thread handles one (col, kvHead, dPair) → reads 2 K dims + 2 V dims, packs each pair into a u32, writes to the f16 cache.

Equations

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

def Hesper.Layers.Attention.packVCacheF32ToF16Kernel (numKVHeads maxSeqLen headDim : Nat) : WGSL.Monad.ShaderM Unit

f32→f16 packed copy kernel for V cache. Each thread reads 2 f32 consecutive elements, packs them into one u32, writes to the f16 V cache. Used to mirror f32 V cache writes (which already happened via existing V-projection or fusedRopeKAndCacheWriteKernel's V branch) into the f16 cache for V11 reads.

Equations

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

def Hesper.Layers.Attention.fusedCacheWriteKVKernel (numKVHeads maxSeqLen headDim kvDim : Nat) : WGSL.Monad.ShaderM Unit

Equations

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

def Hesper.Layers.Attention.cachedScoresKernel (numHeads numKVHeads maxSeqLen headDim : Nat) (scale : Float) : WGSL.Monad.ShaderM Unit

Attention scores with dynamic cacheLen from params buffer. Q: [numHeads * headDim] K_cache: [numKVHeads, maxSeqLen, headDim] Scores: [numHeads * maxSeqLen] (only first cacheLen per head used) Params: [pos: u32, cacheLen: u32]

Equations

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

def Hesper.Layers.Attention.cachedSoftmaxKernel (numRows maxSeqLen : Nat) : WGSL.Monad.ShaderM Unit

Softmax with dynamic row size (cacheLen) from params buffer. Input/Output: [numRows * maxSeqLen] (only first cacheLen per row used) Params: [pos: u32, cacheLen: u32]

Equations

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

def Hesper.Layers.Attention.cachedApplyKernel (numHeads numKVHeads maxSeqLen headDim : Nat) : WGSL.Monad.ShaderM Unit

Apply attention weights to V_cache with dynamic cacheLen from params buffer. attn_weights: [numHeads * maxSeqLen] (only first cacheLen per head used) V_cache: [numKVHeads, maxSeqLen, headDim] Output: [numHeads * headDim] = [dim] (single token) Params: [pos: u32, cacheLen: u32]

Equations

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

Forward Pass with KV Cache

def Hesper.Layers.Attention.forwardWithCache {β : Type} [GPUBackend β] (ctx : β) (layer : Attention (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β)) (inputBuf outputBuf : GPUBackend.Buf β) (kvCache : KVCache (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) (pos : Nat) (subNorm : Option (RMSNorm.RMSNorm (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) := none) (preAllocBufs : Option (CachedAttentionBuffers (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) := none) (residualBuf : Option (GPUBackend.Buf β) := none) : IO Unit

Execute attention forward pass with KV cache (single-token inference).

For each new token:

1. Project Q, K_new, V_new (single row)
2. Apply RoPE at position pos
3. Append K_new, V_new to cache
4. Compute attention scores Q @ K_cache^T * scale (GQA)
5. Softmax
6. Apply attention to V_cache (GQA)
7. Sub-norm + O projection

@param pos Position of the new token (0-indexed)

Equations

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

def Hesper.Layers.Attention.forwardWithCacheLoRA (device : WebGPU.Device) (layer : Attention WebGPU.Buffer WGSL.Execute.PreparedDispatch WGSL.Execute.CompiledKernel) (inputBuf outputBuf : WebGPU.Buffer) (kvCache : KVCache WebGPU.Buffer WGSL.Execute.PreparedDispatch) (pos : Nat) (loraAdapter : LoRA.LayerAdapter) (loraScale : Float) (loraHBuf loraYBufQ loraYBufV : WebGPU.Buffer) (subNorm : Option (RMSNorm.RMSNorm WebGPU.Buffer WGSL.Execute.PreparedDispatch) := none) (preAllocBufs : Option (CachedAttentionBuffers WebGPU.Buffer WGSL.Execute.PreparedDispatch) := none) (residualBuf : Option WebGPU.Buffer := none) : IO Unit

Single-token cached attention forward WITH LoRA corrections on Q and V. Identical to forwardWithCache except it injects LoRA after Q/V BitLinear and before RoPE, so the LoRA contribution flows through the full attention.

Equations

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

Integration with GGUF

def Hesper.Layers.Attention.fromGGUF {β : Type} {α : Sort u_1} [GPUBackend β] (ctx : β) (gguf : α) (layerIdx : Nat) (config : Config) : IO (Attention (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β))

Create attention layer from GGUF file

Loads weight tensors for a specific transformer layer.

Example tensor names:

• blk.0.attn_q.weight - Q projection
• blk.0.attn_k.weight - K projection
• blk.0.attn_v.weight - V projection
• blk.0.attn_output.weight - Output projection

@param device WebGPU device @param gguf Loaded GGUF file @param layerIdx Layer index (0-31 for BitNet-3B) @param config Attention configuration

Equations

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