quantization.py

"""
TurboLLM - Quantization Module
Python wrappers for quantization operations

Supports:
- 8-bit integer quantization (INT8)
- 4-bit integer quantization (INT4)
- 2-bit integer quantization (INT2)
- 1-bit binary quantization (BIT1)
- 1.58-bit ternary quantization (TERNARY)
"""

import numpy as np
from typing import Tuple, Dict, Optional
from dataclasses import dataclass
import struct
import warnings


@dataclass
class QuantData:
    """Container for quantized data"""
    qweight: np.ndarray
    scales: np.ndarray
    zeros: Optional[np.ndarray] = None
    rows: int = 0
    cols: int = 0
    quant_type: str = "int8"
    block_size: int = 32


class BaseQuantizer:
    """Base class for quantizers"""
    
    BLOCK_SIZE = 32
    
    @staticmethod
    def compute_block_stats(weights: np.ndarray, start: int, block_size: int) -> Tuple[float, float, float]:
        """Compute min, max, and scale for a block"""
        block_len = min(block_size, len(weights) - start)
        block = weights[start:start + block_len]
        
        min_val = float(np.min(block))
        max_val = float(np.max(block))
        scale = (max_val - min_val) / 255.0
        
        if scale == 0:
            scale = 1.0
        
        return min_val, max_val, scale


class Int8Quantizer(BaseQuantizer):
    """8-bit integer quantization"""
    
    @staticmethod
    def quantize(weights: np.ndarray, rows: int, cols: int) -> QuantData:
        """
        Quantize weights to 8-bit integers
        
        Args:
            weights: Flattened weight array
            rows: Number of rows
            cols: Number of columns
            
        Returns:
            QuantData with quantized weights and scales
        """
        if weights.dtype != np.float32:
            weights = weights.astype(np.float32)
        
        size = rows * cols
        num_blocks = (size + BaseQuantizer.BLOCK_SIZE - 1) // BaseQuantizer.BLOCK_SIZE
        
        qweight = np.zeros(size, dtype=np.int8)
        scales = np.zeros(num_blocks, dtype=np.float32)
        zeros = np.zeros(num_blocks, dtype=np.float32)
        
        block_idx = 0
        for b in range(0, size, BaseQuantizer.BLOCK_SIZE):
            block_len = min(BaseQuantizer.BLOCK_SIZE, size - b)
            block = weights[b:b + block_len]
            
            min_val = float(np.min(block))
            max_val = float(np.max(block))
            scale = (max_val - min_val) / 255.0
            
            if scale == 0:
                scale = 1.0
            
            scales[block_idx] = scale
            zeros[block_idx] = min_val
            
            # Quantize
            qvals = np.round((block - min_val) / scale).astype(np.int8)
            qweight[b:b + block_len] = np.clip(qvals, -128, 127)
            
            block_idx += 1
        
        print(f"  [INT8] Quantized {rows}x{cols} matrix ({num_blocks} blocks)")
        print(f"    ✓ Compression: {100.0 * 8 / 32:.1f}%")
        
        return QuantData(
            qweight=qweight,
            scales=scales,
            zeros=zeros,
            rows=rows,
            cols=cols,
            quant_type="int8"
        )
    
    @staticmethod
    def dequantize(qdata: QuantData) -> np.ndarray:
        """Dequantize INT8 weights back to float32"""
        size = qdata.rows * qdata.cols
        output = np.zeros(size, dtype=np.float32)
        
        block_idx = 0
        for i in range(0, size, BaseQuantizer.BLOCK_SIZE):
            block_len = min(BaseQuantizer.BLOCK_SIZE, size - i)
            scale = qdata.scales[block_idx]
            zero = qdata.zeros[block_idx]
            
            output[i:i + block_len] = qdata.qweight[i:i + block_len].astype(np.float32) * scale + zero
            
            block_idx += 1
        
        return output


class Int4Quantizer(BaseQuantizer):
    """4-bit integer quantization (2 values per byte)"""
    
    @staticmethod
    def quantize(weights: np.ndarray, rows: int, cols: int) -> QuantData:
        """Quantize weights to 4-bit integers"""
        if weights.dtype != np.float32:
            weights = weights.astype(np.float32)
        
        size = rows * cols
        num_blocks = (size + BaseQuantizer.BLOCK_SIZE - 1) // BaseQuantizer.BLOCK_SIZE
        
        qweight = np.zeros((size + 1) // 2, dtype=np.uint8)
        scales = np.zeros(num_blocks, dtype=np.float32)
        zeros = np.zeros(num_blocks, dtype=np.float32)
        
        block_idx = 0
        for b in range(0, size, BaseQuantizer.BLOCK_SIZE):
            block_len = min(BaseQuantizer.BLOCK_SIZE, size - b)
            block = weights[b:b + block_len]
            
            min_val = float(np.min(block))
            max_val = float(np.max(block))
            scale = (max_val - min_val) / 15.0
            
            if scale == 0:
                scale = 1.0
            
            scales[block_idx] = scale
            zeros[block_idx] = min_val
            
            # Quantize to 4-bit
            qvals = np.round((block - min_val) / scale).astype(np.int8)
            qvals = np.clip(qvals, -8, 7)
            
            # Pack 2 values per byte
            for i in range(block_len):
                byte_idx = (b + i) // 2
                bit_pos = (b + i) % 2
                q = int(qvals[i]) + 8
                
                if bit_pos == 0:
                    qweight[byte_idx] = (q & 0x0F)
                else:
                    qweight[byte_idx] |= ((q & 0x0F) << 4)
            
            block_idx += 1
        
        print(f"  [INT4] Quantized {rows}x{cols} matrix ({num_blocks} blocks)")
        print(f"    ✓ Compression: {100.0 * 4 / 32:.1f}%")
        
        return QuantData(
            qweight=qweight,
            scales=scales,
            zeros=zeros,
            rows=rows,
            cols=cols,
            quant_type="int4"
        )
    
    @staticmethod
    def dequantize(qdata: QuantData) -> np.ndarray:
        """Dequantize INT4 weights"""
        size = qdata.rows * qdata.cols
        output = np.zeros(size, dtype=np.float32)
        
        block_idx = 0
        for i in range(0, size, BaseQuantizer.BLOCK_SIZE):
            block_len = min(BaseQuantizer.BLOCK_SIZE, size - i)
            scale = qdata.scales[block_idx]
            zero = qdata.zeros[block_idx]
            
            for j in range(block_len):
                byte_idx = (i + j) // 2
                bit_pos = (i + j) % 2
                
                if bit_pos == 0:
                    q = (qdata.qweight[byte_idx] & 0x0F) - 8
                else:
                    q = ((qdata.qweight[byte_idx] >> 4) & 0x0F) - 8
                
                output[i + j] = float(q) * scale + zero
            
            block_idx += 1
        
        return output


class Int2Quantizer(BaseQuantizer):
    """2-bit integer quantization (4 values per byte)"""
    
    @staticmethod
    def quantize(weights: np.ndarray, rows: int, cols: int) -> QuantData:
        """Quantize weights to 2-bit integers"""
        if weights.dtype != np.float32:
            weights = weights.astype(np.float32)
        
        size = rows * cols
        num_blocks = (size + BaseQuantizer.BLOCK_SIZE - 1) // BaseQuantizer.BLOCK_SIZE
        
        qweight = np.zeros((size + 3) // 4, dtype=np.uint8)
        scales = np.zeros(num_blocks, dtype=np.float32)
        zeros = np.zeros(num_blocks, dtype=np.float32)
        
        block_idx = 0
        for b in range(0, size, BaseQuantizer.BLOCK_SIZE):
            block_len = min(BaseQuantizer.BLOCK_SIZE, size - b)
            block = weights[b:b + block_len]
            
            min_val = float(np.min(block))
            max_val = float(np.max(block))
            scale = (max_val - min_val) / 3.0
            
            if scale == 0:
                scale = 1.0
            
            scales[block_idx] = scale
            zeros[block_idx] = min_val
            
            # Quantize to 2-bit
            qvals = np.round((block - min_val) / scale).astype(np.int8)
            qvals = np.clip(qvals, -2, 1)
            
            # Pack 4 values per byte
            for i in range(block_len):
                byte_idx = (b + i) // 4
                bit_pos = (b + i) % 4
                q = int(qvals[i]) + 2
                qweight[byte_idx] |= ((q & 0x03) << (bit_pos * 2))
            
            block_idx += 1
        
        print(f"  [INT2] Quantized {rows}x{cols} matrix ({num_blocks} blocks)")
        print(f"    ✓ Compression: {100.0 * 2 / 32:.1f}%")
        
        return QuantData(
            qweight=qweight,
            scales=scales,
            zeros=zeros,
            rows=rows,
            cols=cols,
            quant_type="int2"
        )
    
    @staticmethod
    def dequantize(qdata: QuantData) -> np.ndarray:
        """Dequantize INT2 weights"""
        size = qdata.rows * qdata.cols
        output = np.zeros(size, dtype=np.float32)
        
        block_idx = 0
        for i in range(0, size, BaseQuantizer.BLOCK_SIZE):
            block_len = min(BaseQuantizer.BLOCK_SIZE, size - i)
            scale = qdata.scales[block_idx]
            zero = qdata.zeros[block_idx]
            
            for j in range(block_len):
                byte_idx = (i + j) // 4
                bit_pos = (i + j) % 4
                
                q = ((qdata.qweight[byte_idx] >> (bit_pos * 2)) & 0x03) - 2
                output[i + j] = float(q) * scale + zero
            
            block_idx += 1
        
        return output


class Binary1bitQuantizer(BaseQuantizer):
    """1-bit binary quantization {-1, +1}"""
    
    @staticmethod
    def quantize(weights: np.ndarray, rows: int, cols: int) -> QuantData:
        """Quantize weights to 1-bit binary"""
        if weights.dtype != np.float32:
            weights = weights.astype(np.float32)
        
        size = rows * cols
        num_blocks = (size + BaseQuantizer.BLOCK_SIZE - 1) // BaseQuantizer.BLOCK_SIZE
        num_words = (size + 31) // 32
        
        qweight = np.zeros(num_words, dtype=np.uint32)
        scales = np.zeros(num_blocks, dtype=np.float32)
        
        block_idx = 0
        for b in range(0, size, BaseQuantizer.BLOCK_SIZE):
            block_len = min(BaseQuantizer.BLOCK_SIZE, size - b)
            block = weights[b:b + block_len]
            
            scale = float(np.mean(np.abs(block)))
            if scale == 0:
                scale = 1.0
            
            scales[block_idx] = scale
            
            # Quantize to {-1, +1}
            for i in range(block_len):
                word_idx = (b + i) // 32
                bit_idx = (b + i) % 32
                
                if block[i] >= 0:
                    qweight[word_idx] |= (1 << bit_idx)
            
            block_idx += 1
        
        print(f"  [1-BIT] Quantized {rows}x{cols} matrix ({num_blocks} blocks)")
        print(f"    ✓ Compression: {100.0 * 1 / 32:.1f}%")
        
        return QuantData(
            qweight=qweight,
            scales=scales,
            rows=rows,
            cols=cols,
            quant_type="bit1"
        )
    
    @staticmethod
    def dequantize(qdata: QuantData) -> np.ndarray:
        """Dequantize 1-bit weights"""
        size = qdata.rows * qdata.cols
        output = np.zeros(size, dtype=np.float32)
        
        block_idx = 0
        for i in range(0, size, BaseQuantizer.BLOCK_SIZE):
            block_len = min(BaseQuantizer.BLOCK_SIZE, size - i)
            scale = qdata.scales[block_idx]
            
            for j in range(block_len):
                word_idx = (i + j) // 32
                bit_idx = (i + j) % 32
                
                is_positive = bool((qdata.qweight[word_idx] >> bit_idx) & 1)
                output[i + j] = scale if is_positive else -scale
            
            block_idx += 1
        
        return output


class TernaryQuantizer(BaseQuantizer):
    """1.58-bit ternary quantization {-1, 0, +1}"""
    
    @staticmethod
    def quantize(weights: np.ndarray, rows: int, cols: int) -> QuantData:
        """Quantize weights to ternary {-1, 0, +1}"""
        if weights.dtype != np.float32:
            weights = weights.astype(np.float32)
        
        size = rows * cols
        num_blocks = (size + BaseQuantizer.BLOCK_SIZE - 1) // BaseQuantizer.BLOCK_SIZE
        
        qweight = np.zeros((size + 3) // 4, dtype=np.uint8)
        scales = np.zeros(num_blocks, dtype=np.float32)
        
        block_idx = 0
        for b in range(0, size, BaseQuantizer.BLOCK_SIZE):
            block_len = min(BaseQuantizer.BLOCK_SIZE, size - b)
            block = weights[b:b + block_len]
            
            scale = float(np.mean(np.abs(block)))
            if scale == 0:
                scale = 1.0
            
            threshold = 0.5 * scale
            scales[block_idx] = scale
            
            # Quantize to {-1, 0, +1}
            for i in range(block_len):
                byte_idx = (b + i) // 4
                bit_pos = (b + i) % 4
                
                if block[i] > threshold:
                    q = 1  # positive
                elif block[i] < -threshold:
                    q = 2  # negative
                else:
                    q = 0  # zero
                
                qweight[byte_idx] |= ((q & 0x03) << (bit_pos * 2))
            
            block_idx += 1
        
        print(f"  [TERNARY] Quantized {rows}x{cols} matrix ({num_blocks} blocks)")
        print(f"    ✓ Compression: {100.0 * 1.58 / 32:.1f}%")
        
        return QuantData(
            qweight=qweight,
            scales=scales,
            rows=rows,
            cols=cols,
            quant_type="ternary"
        )
    
    @staticmethod
    def dequantize(qdata: QuantData) -> np.ndarray:
        """Dequantize ternary weights"""
        size = qdata.rows * qdata.cols
        output = np.zeros(size, dtype=np.float32)
        
        block_idx = 0
        for i in range(0, size, BaseQuantizer.BLOCK_SIZE):
            block_len = min(BaseQuantizer.BLOCK_SIZE, size - i)
            scale = qdata.scales[block_idx]
            
            for j in range(block_len):
                byte_idx = (i + j) // 4
                bit_pos = (i + j) % 4
                
                q = (qdata.qweight[byte_idx] >> (bit_pos * 2)) & 0x03
                
                if q == 1:
                    output[i + j] = scale
                elif q == 2:
                    output[i + j] = -scale
                else:
                    output[i + j] = 0.0
            
            block_idx += 1
        
        return output


class QuantizationBenchmark:
    """Benchmark quantization methods"""
    
    @staticmethod
    def compute_metrics(original: np.ndarray, dequantized: np.ndarray) -> Dict[str, float]:
        """Compute quantization error metrics"""
        error = np.abs(original - dequantized)
        
        mse = float(np.mean(error ** 2))
        rmse = float(np.sqrt(mse))
        max_error = float(np.max(error))
        
        sum_orig = float(np.sum(np.abs(original)))
        snr = 20 * np.log10(sum_orig / (len(original) * rmse)) if rmse > 0 else 0
        
        return {
            "mse": mse,
            "rmse": rmse,
            "max_error": max_error,
            "snr_db": snr
        }
    
    @staticmethod
    def print_comparison(name: str, original: np.ndarray, dequantized: np.ndarray) -> None:
        """Print comparison metrics"""
        metrics = QuantizationBenchmark.compute_metrics(original, dequantized)
        
        print(f"    {name:12s} | RMSE: {metrics['rmse']:.2e} | "
              f"SNR: {metrics['snr_db']:6.1f} dB | "
              f"Max Error: {metrics['max_error']:.2e}")


# Compression ratios
COMPRESSION_RATIOS = {
    "fp32": 1.0,
    "fp16": 2.0,
    "int8": 4.0,
    "int4": 8.0,
    "int2": 16.0,
    "bit1": 32.0,
    "ternary": 20.3  # 32 / 1.58
}

QUANTIZERS = {
    "int8": Int8Quantizer,
    "int4": Int4Quantizer,
    "int2": Int2Quantizer,
    "bit1": Binary1bitQuantizer,
    "ternary": TernaryQuantizer,
}


def quantize(weights: np.ndarray, quant_type: str = "int4") -> QuantData:
    """
    Quantize weights using specified method
    
    Args:
        weights: Float32 weight array
        quant_type: Quantization type ("int8", "int4", "int2", "bit1", "ternary")
        
    Returns:
        QuantData with quantized weights
    """
    if quant_type not in QUANTIZERS:
        raise ValueError(f"Unknown quantization type: {quant_type}")
    
    if weights.ndim != 2:
        weights = weights.reshape(-1, 1)
    
    rows, cols = weights.shape
    return QUANTIZERS[quant_type].quantize(weights.flatten(), rows, cols)


def dequantize(qdata: QuantData) -> np.ndarray:
    """
    Dequantize quantized weights
    
    Args:
        qdata: QuantData object
        
    Returns:
        Float32 weight array
    """
    if qdata.quant_type not in QUANTIZERS:
        raise ValueError(f"Unknown quantization type: {qdata.quant_type}")
    
    output = QUANTIZERS[qdata.quant_type].dequantize(qdata)
    return output.reshape(qdata.rows, qdata.cols)


if __name__ == "__main__":
    # Test
    print("TurboLLM - Quantization Module Test")
    print("=" * 70)
    
    # Create random weights
    np.random.seed(42)
    weights = np.random.randn(128, 256).astype(np.float32) * 0.5
    
    # Test all quantization methods
    for quant_type in ["int8", "int4", "int2", "bit1", "ternary"]:
        print(f"\nTesting {quant_type.upper()}...")
        
        qdata = quantize(weights, quant_type)
        dequant = dequantize(qdata)
        
        QuantizationBenchmark.print_comparison(
            quant_type.upper(),
            weights.flatten(),
            dequant.flatten()
        )
    
    print("\n" + "=" * 70)
    print("✓ All quantization tests completed!")