File size: 6,584 Bytes
0e3d4b8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 | """SplitBit Token System — sub-byte token encoding for SplitBit LLM.
Replaces standard 32-bit token IDs with compressed SplitBit encoding.
Reduces token memory by up to 20x, allowing larger context windows
on smaller hardware.
Encoding formats per hardware tier:
Tier Token Format Bits/Token Compression vs 32-bit
Mobile Ternary packed 1.58 20.25x
Minimal Q2_K packed 2.0 16.0x
Light Q3_K_S packed 3.0 10.67x
Standard Q4_K_M packed 4.0 8.0x
Full Q5_K_M packed 5.0 6.4x
Maximum Q8_0 packed 8.0 4.0x
Datacenter FP8 packed 8.0 4.0x
"""
from __future__ import annotations
import logging
import math
from dataclasses import dataclass, field
from typing import Any
import numpy as np
logger = logging.getLogger(__name__)
# Bits per token for each format
TOKEN_BPW: dict[str, float] = {
"ternary": math.log2(3), # 1.585
"q2_k": 2.0,
"q3_k_s": 3.0,
"q4_k_m": 4.0,
"q5_k_m": 5.0,
"q8_0": 8.0,
"fp16": 16.0,
"fp32": 32.0,
}
@dataclass
class SplitBitTokenConfig:
"""Configuration for SplitBit token encoding."""
format: str = "q4_k_m"
block_size: int = 64
cache_enabled: bool = True
cache_max_size: int = 10000
class SplitBitTokenizer:
"""Encodes/decodes token IDs using SplitBit sub-byte compression.
Token IDs (normally 32-bit integers) are packed into compressed
byte arrays using the configured quantization format. This reduces
memory usage for token sequences by 4-20x.
"""
def __init__(self, config: SplitBitTokenConfig | None = None) -> None:
self.config = config or SplitBitTokenConfig()
self.bpw = TOKEN_BPW.get(self.config.format, 4.0)
self._cache: dict[int, bytes] = {}
self._stats = {
"encoded": 0,
"decoded": 0,
"bytes_in": 0,
"bytes_out": 0,
"cache_hits": 0,
}
@property
def compression_ratio(self) -> float:
"""Compression ratio vs 32-bit token IDs."""
return 32.0 / self.bpw if self.bpw > 0 else 1.0
def encode(self, token_ids: list[int] | np.ndarray) -> bytes:
"""Encode token IDs to compressed bytes.
Uses block-based packing:
- Each block of `block_size` tokens is packed together
- Scale factor per block for quantized formats
"""
if isinstance(token_ids, np.ndarray):
token_ids = token_ids.tolist()
self._stats["encoded"] += 1
self._stats["bytes_in"] += len(token_ids) * 4 # 32-bit reference
arr = np.array(token_ids, dtype=np.int32)
block_size = self.config.block_size
n = len(arr)
n_blocks = (n + block_size - 1) // block_size
# Pad to block boundary
pad = n_blocks * block_size - n
if pad > 0:
arr = np.pad(arr, (0, pad))
blocks = arr.reshape(n_blocks, block_size)
if self.bpw <= 2.0:
packed = self._pack_2bit(blocks)
elif self.bpw <= 4.0:
packed = self._pack_4bit(blocks)
elif self.bpw <= 8.0:
packed = self._pack_8bit(blocks)
else:
packed = blocks.astype(np.int32).tobytes()
result = packed if isinstance(packed, bytes) else packed.tobytes()
self._stats["bytes_out"] += len(result)
return result
def decode(self, data: bytes, n_tokens: int) -> list[int]:
"""Decode compressed bytes back to token IDs."""
self._stats["decoded"] += 1
if self.bpw > 8.0:
arr = np.frombuffer(data, dtype=np.int32)
return arr[:n_tokens].tolist()
block_size = self.config.block_size
n_blocks = (n_tokens + block_size - 1) // block_size
if self.bpw <= 2.0:
blocks = self._unpack_2bit(data, n_blocks, block_size)
elif self.bpw <= 4.0:
blocks = self._unpack_4bit(data, n_blocks, block_size)
else:
blocks = self._unpack_8bit(data, n_blocks, block_size)
return blocks.flatten()[:n_tokens].tolist()
def _pack_2bit(self, blocks: np.ndarray) -> np.ndarray:
"""Pack 2-bit tokens: 4 tokens per byte."""
clamped = np.clip(blocks, 0, 3).astype(np.uint8)
flat = clamped.flatten()
# Pad to multiple of 4
pad = (4 - len(flat) % 4) % 4
if pad > 0:
flat = np.pad(flat, (0, pad))
# Pack 4 values per byte
n = len(flat) // 4
packed = np.zeros(n, dtype=np.uint8)
for i in range(4):
packed |= (flat[i::4][:n] & 0x03) << (i * 2)
return packed
def _unpack_2bit(self, data: bytes, n_blocks: int, block_size: int) -> np.ndarray:
"""Unpack 2-bit tokens."""
arr = np.frombuffer(data, dtype=np.uint8)
result = np.zeros(n_blocks * block_size, dtype=np.int32)
for i in range(4):
result[i::4] = (arr >> (i * 2)) & 0x03
return result.reshape(n_blocks, block_size)
def _pack_4bit(self, blocks: np.ndarray) -> np.ndarray:
"""Pack 4-bit tokens: 2 tokens per byte."""
clamped = np.clip(blocks, 0, 15).astype(np.uint8)
flat = clamped.flatten()
# Pad to even length
if len(flat) % 2 != 0:
flat = np.pad(flat, (0, 1))
packed = (flat[0::2] & 0x0F) | ((flat[1::2] & 0x0F) << 4)
return packed
def _unpack_4bit(self, data: bytes, n_blocks: int, block_size: int) -> np.ndarray:
"""Unpack 4-bit tokens."""
arr = np.frombuffer(data, dtype=np.uint8)
result = np.zeros(n_blocks * block_size, dtype=np.int32)
result[0::2] = arr & 0x0F
result[1::2] = (arr >> 4) & 0x0F
return result.reshape(n_blocks, block_size)
def _pack_8bit(self, blocks: np.ndarray) -> np.ndarray:
"""Pack 8-bit tokens: 1 token per byte."""
return np.clip(blocks, 0, 255).astype(np.uint8)
def _unpack_8bit(self, data: bytes, n_blocks: int, block_size: int) -> np.ndarray:
"""Unpack 8-bit tokens."""
arr = np.frombuffer(data, dtype=np.uint8)
return arr[:n_blocks * block_size].astype(np.int32).reshape(n_blocks, block_size)
def get_stats(self) -> dict[str, Any]:
return {
**self._stats,
"format": self.config.format,
"bpw": round(self.bpw, 4),
"compression_ratio": round(self.compression_ratio, 2),
"cache_size": len(self._cache),
}
|