singularity-llm / singularity_llm /splitbit /splitbit_tokens.py
hermescures1's picture
Upload folder using huggingface_hub
32112fa verified
Raw
History Blame Contribute Delete
6.61 kB
"""Singularity Token System — sub-byte token encoding for Singularity LLM.
Replaces standard 32-bit token IDs with compressed Singularity 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 SingularityTokenConfig:
"""Configuration for Singularity token encoding."""
format: str = "q4_k_m"
block_size: int = 64
cache_enabled: bool = True
cache_max_size: int = 10000
class SingularityTokenizer:
"""Encodes/decodes token IDs using Singularity 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: SingularityTokenConfig | None = None) -> None:
self.config = config or SingularityTokenConfig()
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),
}