TQ2_0 Quantization - GPU Unpacking Kernels

Implements on-the-fly dequantization of TQ2_0 (2-bit ternary) weights on GPU.

TQ2_0 Format

• Block size: 256 elements
• Packing: 4 ternary values per byte (2 bits each)
• Scale: 1 FP16 scale per block (stored as 2 bytes)
• Encoding: {-1, 0, 1} → {0b11, 0b00, 0b01}
  • 0b00 → 0
  • 0b01 → +1
  • 0b11 → -1 (two's complement)
• Layout: [packed_bytes (64)] [scale (FP16=2 bytes)] = 66 bytes per block

Memory Layout Example

Block 0: [byte0][byte1]...[byte63][scale_low][scale_high]
         └─4 values─┘                └────FP16────┘

Performance

• Bandwidth savings: 16x (2-bit vs 32-bit)
• PCIe transfer: ~23ms for 3B model (vs ~375ms for Float32)
• Compute: Fused unpack + compute in single kernel

References

• llama.cpp: ggml/src/ggml-quants.c (lines 2103-2271)
• GGUF spec: TQ2_0 = type ID 35

Constants

def Hesper.Quantization.TQ2_0.blockSize : Nat

TQ2_0 block size in elements

Equations

• Hesper.Quantization.TQ2_0.blockSize = 256

def Hesper.Quantization.TQ2_0.blockSizeBytes : Nat

TQ2_0 block size in bytes (64 packed + 2 scale)

Equations

• Hesper.Quantization.TQ2_0.blockSizeBytes = 66

def Hesper.Quantization.TQ2_0.valuesPerByte : Nat

Number of ternary values packed per byte

Equations

• Hesper.Quantization.TQ2_0.valuesPerByte = 4

DSL Helper Functions

def Hesper.Quantization.TQ2_0.unpackTernary2bit (packed idx : WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.u32)) : WGSL.Exp (WGSL.WGSLType.scalar WGSL.ScalarType.f32)

Unpack a single 2-bit ternary value from a packed byte (i2_s encoding) @param packed The packed u32 containing multiple bytes @param idx The index (0-3) of the value to extract within a byte @return Ternary value as f32: -1.0, 0.0, or 1.0

i2_s encoding:

• 0b00 (0) → -1
• 0b01 (1) → 0
• 0b10 (2) → +1 Formula: ternary = code - 1

Equations

• Hesper.Quantization.TQ2_0.unpackTernary2bit packed idx = ((packed.shiftRight (idx.mul (Hesper.WGSL.Exp.litU32 2))).bitAnd (Hesper.WGSL.Exp.litU32 3)).toF32.sub (Hesper.WGSL.Exp.litF32 1.0)

GPU Unpacking Kernel

def Hesper.Quantization.TQ2_0.unpackTQ2_0Kernel (numElements : Nat) : WGSL.Monad.ShaderM Unit

GPU kernel to unpack TQ2_0 quantized weights to Float32

This kernel reads packed ternary weights and FP16 scales from GPU buffers, unpacks them on-the-fly, and writes Float32 output.

Buffer Layout:

• Input buffer "packed": Array of u32 containing packed bytes (64 bytes per block)
• Input buffer "scales": Array of u32 containing FP16 scales (2 bytes per block, stored as u32)
• Output buffer "output": Array of f32 containing unpacked values

Workgroup size: 256 threads (one block per workgroup)

@param numElements Total number of elements to unpack

Equations

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

High-Level API

structure Hesper.Quantization.TQ2_0.UnpackConfig : Type

Configuration for TQ2_0 unpacking execution

• numElements : Nat
• workgroupSize : Nat

def Hesper.Quantization.TQ2_0.instReprUnpackConfig.repr : UnpackConfig → Nat → Std.Format

Equations

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

instance Hesper.Quantization.TQ2_0.instReprUnpackConfig : Repr UnpackConfig

Equations

• Hesper.Quantization.TQ2_0.instReprUnpackConfig = { reprPrec := Hesper.Quantization.TQ2_0.instReprUnpackConfig.repr }

def Hesper.Quantization.TQ2_0.executeUnpack (device : WebGPU.Device) (packedBuf scalesBuf outputBuf : WebGPU.Buffer) (config : UnpackConfig) : IO Unit

Execute TQ2_0 unpacking on GPU

@param device WebGPU device @param packedBuf GPU buffer containing packed TQ2_0 data @param scalesBuf GPU buffer containing FP16 scales (as u32) @param outputBuf GPU buffer for Float32 output @param config Unpacking configuration

Equations

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