Transformer Block

Implements a complete transformer layer combining attention and feed-forward network.

Architecture

Input
  │
  ├─────────────────────────┐
  │                         │
  ▼                         │
[RMSNorm]                   │  Pre-attention normalization
  │                         │
  ▼                         │
[Multi-Head Attention]      │  Self-attention mechanism
  │                         │
  ▼                         │
[Add] ←─────────────────────┘  Residual connection
  │
  ├─────────────────────────┐
  │                         │
  ▼                         │
[RMSNorm]                   │  Pre-FFN normalization
  │                         │
  ▼                         │
[FFN: Gate + Up/Down]       │  Feed-forward network
  │                         │
  ▼                         │
[Add] ←─────────────────────┘  Residual connection
  │
  ▼
Output

Feed-Forward Network (FFN)

BitNet uses a gated FFN similar to LLaMA:

gate = ReLU²(x @ W_gate)  # Gating signal (BitNet uses relu_sqr, NOT silu)
up = x @ W_up              # Value signal
hidden = gate × up         # Element-wise gating
output = hidden @ W_down   # Down-projection

Why gating?

• Non-linear mixing of features
• Selective information flow
• Better gradient properties than simple ReLU

Dimensions (BitNet-3B):

Input: [batch, seq, 2560]
Gate: [2560, 10240] (4× expansion)
Up: [2560, 10240]
Down: [10240, 2560] (back to original)

Memory Layout

Per-layer weights (TQ2_0 quantized):

Attention:
- W_q, W_k, W_v: 3 × 2560² × 0.25 bytes = 4.9 MB
- W_o: 2560² × 0.25 bytes = 1.6 MB

FFN:
- W_gate, W_up: 2 × (2560 × 10240) × 0.25 bytes = 13.1 MB
- W_down: (10240 × 2560) × 0.25 bytes = 6.5 MB

RMSNorm:
- attn_norm, ffn_norm: 2 × 2560 × 4 bytes = 20 KB

Total: ~26 MB per layer
32 layers: ~832 MB

Performance

Compute per token (BitNet-3B):

Attention: 285 GFLOPS
FFN:
- Gate projection: 52.4 GFLOPS
- Up projection: 52.4 GFLOPS
- SiLU + multiply: 0.05 GFLOPS
- Down projection: 104.8 GFLOPS
FFN Total: 209.7 GFLOPS

RMSNorm: 0.02 GFLOPS (negligible)

Total per layer: ~495 GFLOPS

References

• LLaMA: https://github.com/facebookresearch/llama
• BitNet: https://arxiv.org/abs/2402.17764
• llama.cpp: llama.cpp (llm_load_tensors, llm_build_graph)

Configuration

structure Hesper.Layers.TransformerBlock.Config : Type

Transformer block configuration

• layerIdx : Nat
• dim : Nat
• ffnDim : Nat
• numHeads : Nat
• maxSeqLen : Nat

instance Hesper.Layers.TransformerBlock.instReprConfig : Repr Config

Equations

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

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

Equations

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

Pre-allocated Buffers

structure Hesper.Layers.TransformerBlock.LayerBuffers (BufT : Type) : Type

Pre-allocated buffers for transformer block forward pass. Avoids ~21 GPU buffer allocations per layer per token (9 block + 12 attention).

• normedBuf : BufT
• attnOutBuf : BufT
• residual1Buf : BufT
• normed2Buf : BufT
• gateBuf : BufT
• upBuf : BufT
• hiddenBuf : BufT
• ffnOutBuf : BufT
• ffnNormedBuf : BufT
• rmsTempBuf : BufT
• attnBufs : Attention.AttentionBuffers BufT

def Hesper.Layers.TransformerBlock.createLayerBuffers {β : Type} [GPUBackend β] (ctx : β) (dim ffnDim : Nat) (attnConfig : Attention.Config) (batchSize seqLen : Nat) : IO (LayerBuffers (GPUBackend.Buf β))

Create pre-allocated buffers for one transformer block forward pass

Equations

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

Layer Structure

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

Complete transformer block

• config : Config
• attnNorm : RMSNorm.RMSNorm BufT CacheT
• attention : Attention.Attention BufT CacheT KernelT
• attnSubNorm : RMSNorm.RMSNorm BufT CacheT
• ffnNorm : RMSNorm.RMSNorm BufT CacheT
• ffnGate : BitLinear.BitLinear BufT CacheT KernelT
• ffnUp : BitLinear.BitLinear BufT CacheT KernelT
• ffnSubNorm : RMSNorm.RMSNorm BufT CacheT
• ffnDown : BitLinear.BitLinear BufT CacheT KernelT

Layer Creation

def Hesper.Layers.TransformerBlock.create {β : Type} [GPUBackend β] (ctx : β) (config : Config) (attnNormData attnSubNormData ffnNormData ffnSubNormData : ByteArray) (attnWeights : Attention.Attention (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β)) (ffnGateData ffnUpData ffnDownData : ByteArray × ByteArray) : IO (TransformerBlock (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β))

Create transformer block from GGUF tensors

@param device WebGPU device @param config Block configuration @param attnNormData Pre-attention RMSNorm scale parameters @param ffnNormData Pre-FFN RMSNorm scale parameters @param attnWeights Attention layer weight data (Q, K, V, O + scales) @param ffnGateData FFN gate projection weights + scales @param ffnUpData FFN up projection weights + scales @param ffnDownData FFN down projection weights + scales

Equations

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

def Hesper.Layers.TransformerBlock.createWithLayers {β : Type} [GPUBackend β] (ctx : β) (config : Config) (attnNormData attnSubNormData ffnNormData ffnSubNormData : ByteArray) (attnWeights : Attention.Attention (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β)) (ffnGateLayer ffnUpLayer ffnDownLayer : BitLinear.BitLinear (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β)) : IO (TransformerBlock (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β))

Create transformer block with pre-built BitLinear layers

@param device WebGPU device @param config Block configuration @param attnNormData Pre-attention RMSNorm scale parameters @param ffnNormData Pre-FFN RMSNorm scale parameters @param attnWeights Already-created Attention layer @param ffnGateLayer Already-created FFN gate BitLinear @param ffnUpLayer Already-created FFN up BitLinear @param ffnDownLayer Already-created FFN down BitLinear

Equations

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

Forward Pass

def Hesper.Layers.TransformerBlock.forward {β : Type} [GPUBackend β] (ctx : β) (block : TransformerBlock (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β)) (inputBuf outputBuf : GPUBackend.Buf β) (batchSize seqLen : Nat) (preAllocBufs : Option (LayerBuffers (GPUBackend.Buf β)) := none) : IO Unit

Execute transformer block forward pass

Algorithm:

# Attention sub-layer
1. normed = RMSNorm(input)
2. attn_out = MultiHeadAttention(normed)
3. residual1 = input + attn_out  # Residual connection

# FFN sub-layer
4. normed2 = RMSNorm(residual1)
5. gate = W_gate @ normed2
6. up = W_up @ normed2
7. hidden = ReLU²(gate) × up  # Gated activation (BitNet uses relu_sqr)
8. ffn_out = W_down @ hidden
9. residual2 = residual1 + ffn_out  # Residual connection

return residual2

@param device WebGPU device @param block Transformer block @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.

Pre-allocated Buffers for Cached Inference

structure Hesper.Layers.TransformerBlock.CachedLayerBuffers (BufT : Type) (CacheT : Type := Unit) : Type

Pre-allocated buffers for single-token forward pass with KV cache

• normedBuf : BufT
• attnOutBuf : BufT
• residual1Buf : BufT
• normed2Buf : BufT
• gateBuf : BufT
• upBuf : BufT
• hiddenBuf : BufT
• ffnOutBuf : BufT
• ffnNormedBuf : BufT
• rmsTempBuf : BufT
• attnBufs : Attention.CachedAttentionBuffers BufT CacheT
• preparedReluSqrMul : IO.Ref (Option CacheT)

def Hesper.Layers.TransformerBlock.createCachedLayerBuffers {β : Type} [GPUBackend β] (ctx : β) (dim ffnDim : Nat) (attnConfig : Attention.Config) : IO (CachedLayerBuffers (GPUBackend.Buf β) (GPUBackend.CachedDispatch β))

Create pre-allocated buffers for cached single-token inference

Equations

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

Per-Layer Fused PreparedDispatch Refs

structure Hesper.Layers.TransformerBlock.FusedLayerRefs (CacheT : Type := Unit) : Type

Per-layer PreparedDispatch refs for fused kernels. Each layer needs its own refs because fused kernels bind per-layer weight buffers.

• fusedGateUpRelu : IO.Ref (Option CacheT)

def Hesper.Layers.TransformerBlock.createFusedLayerRefs {β : Type} [GPUBackend β] : IO (FusedLayerRefs (GPUBackend.CachedDispatch β))

Create per-layer fused PreparedDispatch refs

Equations

• Hesper.Layers.TransformerBlock.createFusedLayerRefs = do let __do_lift ← Hesper.GPUBackend.newCacheRef pure { fusedGateUpRelu := __do_lift }

Forward Pass with KV Cache

def Hesper.Layers.TransformerBlock.forwardWithCache {β : Type} [GPUBackend β] (ctx : β) (block : TransformerBlock (GPUBackend.Buf β) (GPUBackend.CachedDispatch β) (GPUBackend.CompiledKernel β)) (inputBuf outputBuf : GPUBackend.Buf β) (pos : Nat) (kvCache : Attention.KVCache (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) (preAllocBufs : Option (CachedLayerBuffers (GPUBackend.Buf β) (GPUBackend.CachedDispatch β)) := none) (fusedRefs : Option (FusedLayerRefs (GPUBackend.CachedDispatch β)) := none) : IO Unit

Execute single-token transformer block forward pass with KV cache.

Same algorithm as forward but uses cached attention for O(1) per token. All dimensions are for single token (batchSize=1, seqLen=1).

@param pos Current token position (0-indexed) @param kvCache KV cache for this layer's attention @param fusedRefs Per-layer fused PreparedDispatch refs (optional)

Equations

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

def Hesper.Layers.TransformerBlock.forwardWithCacheLoRA (device : WebGPU.Device) (block : TransformerBlock WebGPU.Buffer WGSL.Execute.PreparedDispatch WGSL.Execute.CompiledKernel) (inputBuf outputBuf : WebGPU.Buffer) (pos : Nat) (kvCache : Attention.KVCache WebGPU.Buffer WGSL.Execute.PreparedDispatch) (loraAdapter : LoRA.LayerAdapter) (loraScale : Float) (loraHBuf loraYBufQ loraYBufV : WebGPU.Buffer) (preAllocBufs : Option (CachedLayerBuffers WebGPU.Buffer WGSL.Execute.PreparedDispatch) := none) (fusedRefs : Option (FusedLayerRefs WGSL.Execute.PreparedDispatch) := none) : IO Unit

Cached forward pass WITH LoRA corrections on attention Q/V. Same as forwardWithCache but uses Attention.forwardWithCacheLoRA to inject LoRA before RoPE.

Equations

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

Integration with GGUF

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

Create transformer block from GGUF file

Loads all tensors for a specific transformer layer.

Example tensor names in GGUF:

• blk.{idx}.attn_norm.weight - Pre-attention RMSNorm
• blk.{idx}.attn_q.weight - Attention Q projection
• blk.{idx}.attn_k.weight - Attention K projection
• blk.{idx}.attn_v.weight - Attention V projection
• blk.{idx}.attn_output.weight - Attention output projection
• blk.{idx}.ffn_norm.weight - Pre-FFN RMSNorm
• blk.{idx}.ffn_gate.weight - FFN gate projection
• blk.{idx}.ffn_up.weight - FFN up projection
• blk.{idx}.ffn_down.weight - FFN down projection

@param device WebGPU device @param gguf Loaded GGUF file @param config Block configuration

Equations

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