Embedding Layer

Implements token embedding lookup with TQ2_0 quantization support.

Mathematical Definition

Embedding layer maps discrete token IDs to continuous vector representations:

Input: token_ids [batch, seq_len]  (integers: 0 to vocab_size-1)
Output: embeddings [batch, seq_len, dim]  (float vectors)

Operation:
  embedding[i,j] = embedding_table[token_ids[i,j]]

Quantized Embeddings

For memory efficiency, embeddings are stored in TQ2_0 format:

Standard: vocab_size × dim × 4 bytes (Float32)
TQ2_0:    vocab_size × dim × 0.25 bytes (2-bit ternary)
Savings:  16× reduction

BitNet-3B example:

Vocabulary: 50,000 tokens
Dimension: 2560
Standard: 50000 × 2560 × 4 = 512 MB
TQ2_0:    50000 × 2560 × 0.25 = 32 MB

GPU Implementation

Two approaches:

1. Dequantized Lookup (Simpler)

1. Precompute: unpack TQ2_0 → Float32 table (once at startup)
2. During inference: direct lookup from Float32 table
Pros: Fast lookup, simple implementation
Cons: Uses more memory (Float32 table)

2. On-the-Fly Lookup (Memory-efficient)

1. Store embeddings as TQ2_0 (packed)
2. During inference: lookup + unpack in single kernel
Pros: Minimal memory (16× savings)
Cons: Slightly more compute per lookup

We implement approach #1 for simplicity, but structure allows easy switch to #2.

Performance

Lookup operation:

Compute: ~0 FLOPs (just memory read)
Memory: batch × seq_len × dim × 4 bytes
Latency: ~1 memory transaction per token

For batch=1, seq=2048, dim=2560:
  Memory transfer: 2048 × 2560 × 4 = ~20 MB
  Time on A100 (2 TB/s): 20 MB / 2000 GB/s ≈ 0.01 ms
  (Negligible compared to attention/FFN)

References

• Word2Vec: "Efficient Estimation of Word Representations" (Mikolov et al., 2013)
• Transformer embeddings: "Attention is All You Need" (Vaswani et al., 2017)
• llama.cpp: llama.cpp (get_rows operation)

Configuration

structure Hesper.Layers.Embedding.Config : Type

Embedding layer configuration

• vocabSize : Nat
• dim : Nat

instance Hesper.Layers.Embedding.instReprConfig : Repr Config

Equations

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

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

Equations

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

GPU Kernel

def Hesper.Layers.Embedding.embeddingLookupKernel (config : Config) (batchSize seqLen : Nat) : WGSL.Monad.ShaderM Unit

Embedding lookup kernel

Maps token IDs to embedding vectors.

Input: token_ids [batch × seq_len] (UInt32) Output: embeddings [batch × seq_len × dim] (Float32)

Each thread processes one token, reading its embedding vector from the embedding table.

@param config Embedding configuration @param batchSize Batch size @param seqLen Sequence length

Equations

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

Layer Structure

structure Hesper.Layers.Embedding.Embedding (BufT : Type) : Type

Embedding layer

• config : Config
• embeddingTable : BufT
• f16Table : Option BufT

Layer Creation

def Hesper.Layers.Embedding.create {β : Type} [GPUBackend β] (ctx : β) (config : Config) (packedData scalesData : ByteArray) : IO (Embedding (GPUBackend.Buf β))

Create embedding layer from quantized data

Unpacks TQ2_0 embeddings to Float32 for fast lookup.

@param device WebGPU device @param config Embedding configuration @param packedData TQ2_0 packed embedding data @param scalesData FP16 scales for each block

Equations

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

def Hesper.Layers.Embedding.createFromFloat32 {β : Type} [GPUBackend β] (ctx : β) (config : Config) (float32Data : ByteArray) : IO (Embedding (GPUBackend.Buf β))

Create embedding layer from Float32 data (for testing)

@param device WebGPU device @param config Embedding configuration @param float32Data Raw Float32 embedding data

Equations

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

def Hesper.Layers.Embedding.unpackF16ToF32Kernel (numElements packedPerThread : Nat) : WGSL.Monad.ShaderM Unit

GPU kernel to unpack F16 → F32 using hardware instruction

Each thread processes packedPerThread packed U32 values (each containing 2 F16s), producing 2 * packedPerThread F32 outputs. This keeps workgroup count within the 65535 max per dimension limit.

Equations

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

def Hesper.Layers.Embedding.createFromF16 {β : Type} [GPUBackend β] (ctx : β) (config : Config) (f16Data : ByteArray) : IO (Embedding (GPUBackend.Buf β))

Create embedding layer from F16 data (GPU-optimized version)

Uploads raw F16 data (256 MB) and unpacks to F32 on GPU using hardware instruction. This is 6× faster than CPU conversion: 2× less bandwidth + GPU parallelism.

@param device WebGPU device @param config Embedding configuration @param f16Data Raw F16 (Float16) embedding data (packed as bytes)

Equations

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

Forward Pass

@[inline]

def Hesper.Layers.Embedding.forward {β : Type} [GPUBackend β] (ctx : β) (layer : Embedding (GPUBackend.Buf β)) (tokenIdsBuf outputBuf : GPUBackend.Buf β) (batchSize seqLen : Nat) : IO Unit

Execute embedding lookup

@param device WebGPU device @param layer Embedding layer @param tokenIdsBuf Input token IDs [batch, seq_len] (UInt32) @param outputBuf Output embeddings [batch, seq_len, dim] (Float32) @param batchSize Batch size @param seqLen Sequence length

Equations

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

Integration with GGUF

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

Create embedding layer from GGUF file

Typical tensor name: token_embd.weight

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

Equations

• Hesper.Layers.Embedding.fromGGUF ctx gguf config = do IO.println "[Embedding] Loading from GGUF" throw (IO.userError "fromGGUF not yet implemented - use create() or createFromFloat32()")