| """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__) |
|
|
| |
| TOKEN_BPW: dict[str, float] = { |
| "ternary": math.log2(3), |
| "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 |
|
|
| 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 = 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 = (4 - len(flat) % 4) % 4 |
| if pad > 0: |
| flat = np.pad(flat, (0, pad)) |
| |
| 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() |
| |
| 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), |
| } |
|
|
