BitLinear Layer - Ternary Weight Matrix Multiplication (i2_s format) #
Implements BitNet's BitLinear layer with on-the-fly i2_s dequantization on GPU.
Key Innovation: Fused Kernel #
Instead of:
- Unpack i2_s → Float32 on CPU (slow, memory-intensive)
- Matrix multiply Float32 × Float32
We do:
- Upload raw packed i2_s bytes to GPU
- Read packed weights + compute matmul in same kernel
i2_s Packing Format #
Encoding table: | Ternary | 2-bit code | |---------|-----------| | -1 | 0b00 (0) | | 0 | 0b01 (1) | | +1 | 0b10 (2) |
Dequantization: float_value = (code - 1) * scale
Layout (groups of 128 elements per 32 bytes):
- Elements [0..31]: bytes[0..31] >> 6 & 3
- Elements [32..63]: bytes[0..31] >> 4 & 3
- Elements [64..95]: bytes[0..31] >> 2 & 3
- Elements [96..127]: bytes[0..31] >> 0 & 3
Scale: single F32 at the END of tensor data.
BitLinear Mathematics #
For ternary weights w in {-1, 0, 1}, the matrix-vector product:
y[i] = scale * sum_j( w[i,j] * x[j] )
= scale * (sum_{w=+1} x[j] - sum_{w=-1} x[j])
References #
- BitNet paper: https://arxiv.org/abs/2402.17764
- bitnet.cpp: i2_s format specification
Counters for PreparedDispatch fast-path vs slow-path
Runtime opt-in for the subgroup-matrix BitLinear kernel. Defaults
to on when the device also has ShaderF16 + SubgroupMatrix (see
dispatch logic in forward). Set to false to force the tiled
fallback for debugging.
Layer Configuration #
Equations
Equations
- One or more equations did not get rendered due to their size.
Instances For
Equations
Instances For
Fused BitLinear Kernel (i2_s format) #
Vectorized fused kernel: i2_s unpack + matrix-vector multiply
Algorithm (vectorized, 16 weights per u32 read):
for each output element y[out_idx]:
1. Load input[0..inDim] into shared memory (all threads cooperate)
2. acc = 0.0
3. for each u32 in packed weights (strided over threads):
Read 1 u32 = 4 bytes = 16 packed 2-bit weights
Extract all 16 weights via compile-time unrolled shifts
Accumulate 16 FMAs from shared memory input
4. Tree reduction of partial sums
5. y[out_idx] = scale * total_sum
Key optimizations over v1:
- 16x fewer weight buffer reads (1 u32 → 16 elements vs 1 u32 → 1 element)
- Input cached in shared memory (10KB for dim=2560, fast random access)
- Coalesced weight reads (consecutive threads read consecutive u32s)
- Compile-time unrolled inner loop (16 FMAs with no branch overhead)
- Tiled input loading for large dims (>3584 elements)
i2_s unpacking (vectorized per u32): All 4 bytes of a u32 are always within the same group-128 block because u32 boundaries (4 bytes) never span a 32-byte group boundary.
group128 = u32Idx / 8
baseGroupPos = (u32Idx % 8) * 4
For byte b in 0..4, shift s in [6,4,2,0]:
elemIdx = group128 * 128 + baseGroupPos + b + (3-s/2) * 32
code = (byte >> s) & 3
ternary = code - 1
@param config BitLinear layer configuration @param numRows Number of input rows (batch * seq_len) @param workgroupSize Threads per workgroup (default 256)
Equations
- One or more equations did not get rendered due to their size.
Instances For
Shared-Memory Tree Reduction Helper #
M=1 Warp-Cooperative BitLinear Kernel (Single-Token Inference) #
M=1 warp-cooperative kernel: one subgroup (32 threads) per output element
For single-token inference (M=1), uses subgroup-level cooperation:
- 32 threads in a subgroup cooperatively process one output row
- Weight reads are COALESCED: consecutive threads read consecutive u32s
- Input reads are L2-cached: all subgroups read the same 10-27KB input vector
- Reduction via hardware
subgroupAdd(no shared memory, no barriers)
vs tiled kernel (256 threads/workgroup):
- No shared memory for input (10KB × outDim redundant loads eliminated)
- No shared memory for reduction (subgroupAdd replaces tree reduction)
- No barriers (subgroup ops are implicit)
- 8x fewer threads (32 vs 256 per output)
Algorithm per subgroup (32 threads):
for u32Idx = tid; u32Idx < u32PerRow; u32Idx += 32:
// Coalesced: thread k reads u32 at rowBase + k, k+32, k+64...
packed = weights[outIdx * u32PerRow + u32Idx]
unpack 16 ternary weights
read 16 input values (L2-cached)
acc += 16 FMAs
total = subgroupAdd(acc)
if tid == 0: output[outIdx] = scale * total
@param config BitLinear layer configuration
Equations
- One or more equations did not get rendered due to their size.
Instances For
M=1 warp-cooperative kernel with fused residual add
Same as fusedBitLinearM1Kernel but outputs: output = residual + scale * dot_product
@param config BitLinear layer configuration
Equations
- One or more equations did not get rendered due to their size.
Instances For
Fused RMSNorm + BitLinear + Residual Add M=1 Kernel #
Fused RMSNorm + BitLinear + Residual M=1 kernel.
Combines three operations into one dispatch:
- RMSNorm: compute rmsInv = rsqrt(mean(input²) + eps)
- BitLinear dot product with inline normalization: dot = sum_j(ternary[outIdx,j] * input[j] * rmsInv * rmsScale[j])
- Residual add: output[outIdx] = residual[outIdx] + blScale * dot
Saves 1 dispatch per call (was: RMSNorm + BitLinear = 2 dispatches). Used for: attn_sub_norm + O projection, ffn_sub_norm + down projection.
Algorithm per subgroup (32 threads):
Phase 1: Compute RMS via subgroupAdd
partial_sq = 0
for elemIdx = tid; elemIdx < inDim; elemIdx += 32:
partial_sq += input[elemIdx]²
totalSq = subgroupAdd(partial_sq)
rmsInv = rsqrt(totalSq / dim + eps)
Phase 2: BitLinear dot product with inline normalization
for u32Idx = tid; u32Idx < u32PerRow; u32Idx += 32:
packed = weights[outIdx * u32PerRow + u32Idx]
for each of 16 elements:
normalized = input[elemIdx] * rmsInv * rmsScale[elemIdx]
acc += ternary * normalized
total = subgroupAdd(acc)
if tid == 0: output[outIdx] = residual[outIdx] + blScale * total
@param config BitLinear layer configuration @param eps RMSNorm epsilon (typically 1e-5)
Equations
- One or more equations did not get rendered due to their size.
Instances For
Fused Gate+Up+ReLU²×Mul M=1 Kernel #
Fused gate + up + ReLU²×mul M=1 kernel.
Combines three operations into one dispatch:
- gate_val = dot(gate_weights, input) * gate_scale
- up_val = dot(up_weights, input) * up_scale
- output = ReLU²(gate_val) * up_val = max(0, gate_val)² * up_val
Saves 2 dispatches per call (was: gate BitLinear + up BitLinear + ReluSqrMul = 3). Input is read once and used for both gate and up dot products.
Algorithm per subgroup (32 threads):
gate_acc = 0, up_acc = 0
for u32Idx = tid; u32Idx < u32PerRow; u32Idx += 32:
// Read input once, unpack both gate and up weights
gate_packed = gate_weights[outIdx * u32PerRow + u32Idx]
up_packed = up_weights[outIdx * u32PerRow + u32Idx]
for each of 16 elements:
input_val = input[elemIdx]
gate_acc += gate_ternary * input_val
up_acc += up_ternary * input_val
gate_total = subgroupAdd(gate_acc) * gate_scale
up_total = subgroupAdd(up_acc) * up_scale
if tid == 0:
relu_val = max(0, gate_total)
output[outIdx] = relu_val * relu_val * up_total
@param config BitLinear layer configuration (inDim=dim, outDim=ffnDim)
Equations
- One or more equations did not get rendered due to their size.
Instances For
Shared-Memory Fallback M=1 Kernels (No Subgroup Support) #
Fused BitLinear + Residual Add Kernel #
Fused kernel: i2_s unpack + matrix-vector multiply + residual add
Same as fusedBitLinearKernel but outputs: output[i] = residual[i] + scale * dot_product
Eliminates a separate elementwise add dispatch per layer. Used for attention O-projection and FFN down-projection residual connections.
@param config BitLinear layer configuration @param numRows Number of input rows @param workgroupSize Threads per workgroup (default 256)
Equations
- One or more equations did not get rendered due to their size.
Instances For
Subgroup-Matrix Multi-row BitLinear Kernel (cooperative matmul) #
Multi-row BitLinear kernel using cooperative matrix operations.
Targets the NVIDIA Ada / Ampere WMMA config (f16, f16) → f32 at 16×16×16.
One workgroup (= one subgroup of 32 threads) computes a single 16×16
output tile Y[rowBase .. rowBase+16, outBase .. outBase+16].
The math is Y = X @ W^T where X is the row-major [numRows, inDim]
input and W is the ternary weight matrix ([outDim, inDim]). So the
subgroup matrix layout is:
A (left, M×K = 16×16) = X[rowBase .. +16, kBase .. +16] B (right, K×N = 16×16) = W^T[kBase .. +16, outBase .. +16] = W[outBase .. +16, kBase .. +16]^T C (result, M×N = 16×16)
A is f16 (cast from f32 input), B is f16 (dequantized from i2_s
ternary), C accumulates as f32 for precision. Preconditions:
numRows % 16 == 0, outDim % 16 == 0, inDim % 16 == 0
(and inDim % 128 == 0 for i2_s). The caller (forward) must check
these and fall back to the existing tiled kernel otherwise.
Layout of shared memory: shared_A : array<f16, 256> row-major 16 × 16 (M × K) shared_B : array<f16, 256> row-major 16 × 16 (K × N)
The subgroup matrix load reads f16 tiles from these with stride=16,
does one 16×16×16 cooperative MAC per K-block, then at the end of
the K loop stores the f32 result into another shared buffer
shared_C of 256 f32 values, after which all 32 threads scale and
write 8 elements each to the output buffer.
Equations
- One or more equations did not get rendered due to their size.
Instances For
High-Level API #
Create BitLinear layer from i2_s packed data
@param device WebGPU device @param config Layer configuration @param packedWeights Raw i2_s packed byte data from GGUF @param scale Float32 scale factor for the ternary weights
Equations
- One or more equations did not get rendered due to their size.
Instances For
Create BitLinear layer from packed data + scale ByteArrays
Alternative constructor that takes pre-encoded scale bytes. Used when the caller has already extracted the raw data.
@param device WebGPU device @param config Layer configuration @param packedWeights Raw i2_s packed byte data @param scaleBytes 4-byte little-endian F32 scale
Equations
- One or more equations did not get rendered due to their size.
Instances For
Execute forward pass
@param device WebGPU device @param layer BitLinear layer @param inputBuf GPU buffer containing input (Float32) @param outputBuf GPU buffer for output (Float32)
Equations
- One or more equations did not get rendered due to their size.
Instances For
Execute forward pass with fused residual add: output = residual + scale * (weights @ input)
Saves one elementwise add dispatch per call (2 dispatches/layer × 30 layers = 60 saved).
@param device WebGPU device @param layer BitLinear layer @param inputBuf GPU buffer containing input (Float32) @param residualBuf GPU buffer containing residual to add (Float32) @param outputBuf GPU buffer for output (Float32) @param numRows Number of input rows
Equations
- One or more equations did not get rendered due to their size.
Instances For
Execute fused RMSNorm + BitLinear + Residual forward pass (M=1 only).
Combines: RMSNorm(input) → BitLinear dot product → residual add into a single GPU dispatch.
output = residual + blScale * (weights @ RMSNorm(input))
@param device WebGPU device @param layer BitLinear layer (weights + scale) @param rmsNorm RMSNorm layer (scale parameters) @param inputBuf Input buffer [inDim] @param residualBuf Residual buffer [outDim] @param outputBuf Output buffer [outDim] @param preparedRef Optional PreparedDispatch ref for fast-path replay
Equations
- One or more equations did not get rendered due to their size.
Instances For
Execute fused Gate+Up+ReLU²×Mul forward pass (M=1 only).
Combines: gate = BitLinear(input), up = BitLinear(input), output = ReLU²(gate) × up into a single GPU dispatch.
@param device WebGPU device @param gateLayer Gate BitLinear layer @param upLayer Up BitLinear layer @param inputBuf Input buffer [inDim] @param outputBuf Output buffer [outDim] @param preparedRef Optional PreparedDispatch ref for fast-path replay
Equations
- One or more equations did not get rendered due to their size.