# iq1_m_quantizer.py - IQ1_M 1bit混合精度量化器实现

import numpy as np
from typing import Tuple, List, Optional
import struct

class IQ1MQuantizer:
    """
    IQ1_M 1bit混合精度量化器
    核心思想：对权重矩阵按块分组，每组独立量化
    - 大部分权重采用1bit表示
    - 部分关键权重（如gate路由层）保留更高精度
    - 每块有独立的scale和offset
    """
    
    def __init__(self, 
                 block_size: int = 32,
                 outlier_threshold: float = 3.0,
                 outlier_bits: int = 8):
        """
        Args:
            block_size: 量化块大小（通常32）
            outlier_threshold: 异常值检测阈值（标准差倍数）
            outlier_bits: 异常值保留的精度
        """
        self.block_size = block_size
        self.outlier_threshold = outlier_threshold
        self.outlier_bits = outlier_bits
        
        # 量化统计
        self.stats = {
            "total_weights": 0,
            "quantized_1bit": 0,
            "quantized_outlier": 0,
            "compression_ratio": 0.0,
        }
    
    def quantize_block(self, weights: np.ndarray) -> Tuple[np.ndarray, float, float]:
        """
        量化一个权重块
        Args:
            weights: 权重块 [block_size]
        Returns:
            quantized: 量化后的1bit权重
            scale: 缩放因子
            offset: 偏移量
        """
        # 1. 计算统计量
        mean = np.mean(weights)
        std = np.std(weights)
        
        # 2. 检测异常值（保留高精度）
        outliers = np.abs(weights - mean) > self.outlier_threshold * std
        
        if np.any(outliers):
            # 对异常值保留8bit精度
            outlier_weights = weights[outliers]
            outlier_min = np.min(outlier_weights)
            outlier_max = np.max(outlier_weights)
            outlier_range = outlier_max - outlier_min
            
            if outlier_range > 0:
                outlier_quantized = np.round(
                    (outlier_weights - outlier_min) / outlier_range * 255
                ).astype(np.uint8)
            else:
                outlier_quantized = np.zeros_like(outlier_weights, dtype=np.uint8)
            
            # 存储异常值信息
            self.stats["quantized_outlier"] += np.sum(outliers)
        
        # 3. 对正常权重做1bit量化
        normal_weights = weights[~outliers]
        if len(normal_weights) > 0:
            scale = np.max(np.abs(normal_weights))
            if scale < 1e-10:
                scale = 1e-10
            
            # 1bit量化: 符号位
            quantized_1bit = (normal_weights > 0).astype(np.int8) * 2 - 1
            self.stats["quantized_1bit"] += len(normal_weights)
        else:
            scale = 1.0
            quantized_1bit = np.array([], dtype=np.int8)
        
        return quantized_1bit, scale, mean
    
    def dequantize_block(self, 
                         quantized: np.ndarray,
                         scale: float,
                         mean: float) -> np.ndarray:
        """
        反量化一个权重块
        """
        # 1bit反量化
        dequantized = quantized.astype(np.float32) * scale
        
        # 恢复均值
        dequantized += mean
        
        return dequantized
    
    def quantize_matrix(self, 
                        weight_matrix: np.ndarray,
                        layer_name: str = "") -> dict:
        """
        量化整个权重矩阵
        """
        rows, cols = weight_matrix.shape
        self.stats["total_weights"] = rows * cols
        
        compressed = {
            "shape": (rows, cols),
            "layer_name": layer_name,
            "blocks": [],
            "metadata": {}
        }
        
        # 按行分块量化
        for i in range(0, rows, self.block_size):
            block_end = min(i + self.block_size, rows)
            block = weight_matrix[i:block_end]
            
            # 按列处理
            block_quantized = []
            block_scales = []
            block_offsets = []
            
            for j in range(cols):
                col_block = block[:, j]
                q, s, o = self.quantize_block(col_block)
                block_quantized.append(q.tobytes())
                block_scales.append(s)
                block_offsets.append(o)
            
            compressed["blocks"].append({
                "row_start": i,
                "row_end": block_end,
                "quantized_data": block_quantized,
                "scales": block_scales,
                "offsets": block_offsets,
            })
        
        # 计算压缩率
        original_bytes = rows * cols * 2  # BF16 = 2 bytes
        compressed_bytes = self._calculate_compressed_size(compressed)
        self.stats["compression_ratio"] = original_bytes / compressed_bytes
        
        compressed["metadata"] = self.stats
        
        return compressed
    
    def _calculate_compressed_size(self, compressed: dict) -> int:
        """计算压缩后大小（字节）"""
        total = 0
        for block in compressed["blocks"]:
            # 1bit权重
            for q_data in block["quantized_data"]:
                total += len(q_data)
            # scales和offsets (FP32)
            total += len(block["scales"]) * 4
            total += len(block["offsets"]) * 4
        
        # 异常值（若有）
        total += self.stats["quantized_outlier"] * 1  # 8bit
        
        return total
    
    def print_stats(self):
        """打印量化统计"""
        total = self.stats["total_weights"]
        print(f"总权重数: {total:,}")
        print(f"1bit量化: {self.stats['quantized_1bit']:,} ({self.stats['quantized_1bit']/total*100:.1f}%)")
        print(f"异常值保留: {self.stats['quantized_outlier']:,} ({self.stats['quantized_outlier']/total*100:.1f}%)")
        print(f"压缩率: {self.stats['compression_ratio']:.1f}x")

# 测试：量化一个模拟的权重矩阵
np.random.seed(42)
weight_matrix = np.random.randn(4096, 4096).astype(np.float32) * 0.1

# 添加一些异常值
weight_matrix[0, 0:100] = np.random.randn(100) * 2.0  # 大异常值

quantizer = IQ1MQuantizer(block_size=32, outlier_threshold=2.5)
compressed = quantizer.quantize_matrix(weight_matrix, "test_layer")
quantizer.print_stats()

# 验证量化质量
dequantized = np.zeros_like(weight_matrix)
for block in compressed["blocks"]:
    for j in range(weight_matrix.shape[1]):
        q_data = np.frombuffer(block["quantized_data"][j], dtype=np.int8)
        scale = block["scales"][j]
        offset = block["offsets"][j]
        dequantized_block = quantizer.dequantize_block(q_data, scale, offset)
        dequantized[block["row_start"]:block["row_end"], j] = dequantized_block

# 计算量化误差
mse = np.mean((weight_matrix - dequantized) ** 2)
snr = 10 * np.log10(np.var(weight_matrix) / mse)
print(f"量化MSE: {mse:.6f}")
print(f"信噪比SNR: {snr:.2f} dB")

// mixed_precision.go - Hy3层级混合精度量化策略

package main

import (
	"fmt"
	"math"
)

// LayerType 层类型
type LayerType int

const (
	AttentionQKV LayerType = iota
	AttentionOutput
	MLPGate
	MLPUp
	MLPDown
	MoERouter
	MoEExpertGate
	MoEExpertUp
	MoEExpertDown
	Embedding
	LayerNorm
)

// LayerQuantConfig 层量化配置
type LayerQuantConfig struct {
	LayerType   LayerType
	QuantBits   int     // 量化位数
	BlockSize   int     // 分块大小
	Importance  float64 // 重要性权重（0-1）
}

// Hy3QuantPlan Hy3量化计划
type Hy3QuantPlan struct {
	Configs []LayerQuantConfig
	Layers  map[string]LayerQuantConfig
}

// NewHy3QuantPlan 创建Hy3量化计划
func NewHy3QuantPlan() *Hy3QuantPlan {
	plan := &Hy3QuantPlan{
		Configs: []LayerQuantConfig{
			// MoE路由层 - 最高精度
			{LayerType: MoERouter, QuantBits: 8, BlockSize: 64, Importance: 1.0},
			// 注意力层 - 中等精度
			{LayerType: AttentionQKV, QuantBits: 4, BlockSize: 128, Importance: 0.8},
			{LayerType: AttentionOutput, QuantBits: 4, BlockSize: 128, Importance: 0.7},
			// MoE专家层 - 最低精度
			{LayerType: MoEExpertGate, QuantBits: 1, BlockSize: 32, Importance: 0.4},
			{LayerType: MoEExpertUp, QuantBits: 1, BlockSize: 32, Importance: 0.3},
			{LayerType: MoEExpertDown, QuantBits: 1, BlockSize: 32, Importance: 0.3},
			// MLP层 - 混合精度
			{LayerType: MLPGate, QuantBits: 4, BlockSize: 128, Importance: 0.6},
			{LayerType: MLPUp, QuantBits: 4, BlockSize: 128, Importance: 0.5},
			{LayerType: MLPDown, QuantBits: 4, BlockSize: 128, Importance: 0.5},
			// 嵌入层 - 高精度
			{LayerType: Embedding, QuantBits: 8, BlockSize: 64, Importance: 0.9},
			// LayerNorm - 不量化（FP32）
			{LayerType: LayerNorm, QuantBits: 32, BlockSize: 0, Importance: 1.0},
		},
		Layers: make(map[string]LayerQuantConfig),
	}
	return plan
}

// AssignLayerConfig 分配层的量化配置
func (p *Hy3QuantPlan) AssignLayerConfig(layerName string, layerType LayerType) LayerQuantConfig {
	// 查找匹配配置
	for _, cfg := range p.Configs {
		if cfg.LayerType == layerType {
			p.Layers[layerName] = cfg
			return cfg
		}
	}
	// 默认4bit
	return LayerQuantConfig{QuantBits: 4, BlockSize: 128, Importance: 0.5}
}

// CalculateCompressionRatio 计算压缩率
func (p *Hy3QuantPlan) CalculateCompressionRatio(
	layerParams map[string]int64,
) map[string]float64 {
	ratios := make(map[string]float64)
	
	for layerName, cfg := range p.Layers {
		params := layerParams[layerName]
		originalBytes := params * 2 // BF16
		quantizedBytes := int64(float64(params) * 2.0 / float64(cfg.QuantBits))
		ratios[layerName] = float64(originalBytes) / float64(quantizedBytes)
	}
	
	return ratios
}

// EstimateTotalSize 估算总大小
func (p *Hy3QuantPlan) EstimateTotalSize(
	layerParams map[string]int64,
) (float64, float64) {
	var totalOriginal, totalQuantized float64
	
	for layerName, cfg := range p.Layers {
		params := float64(layerParams[layerName])
		originalBytes := params * 2 // BF16
		
		var quantizedBytes float64
		if cfg.QuantBits == 32 {
			quantizedBytes = params * 4 // FP32
		} else {
			quantizedBytes = params * 2.0 / float64(cfg.QuantBits)
			// 加上scale/offset开销
			numBlocks := params / float64(cfg.BlockSize)
			if numBlocks < 1 {
				numBlocks = 1
			}
			quantizedBytes += numBlocks * 8 // 每个block 4byte scale + 4byte offset
		}
		
		totalOriginal += originalBytes
		totalQuantized += quantizedBytes
	}
	
	return totalOriginal, totalQuantized
}

func main() {
	plan := NewHy3QuantPlan()
	
	// 模拟Hy3各层参数分布
	layerParams := map[string]int64{
		"embedding":       152064 * 8192,     // 词汇表x隐藏维度
		"attention_qkv":   80 * 8192 * 8192 * 3, // 80层xQKV
		"attention_output": 80 * 8192 * 8192,
		"mlp_gate":        80 * 8192 * 32768,
		"mlp_up":          80 * 8192 * 32768,
		"mlp_down":        80 * 32768 * 8192,
		"moe_router":      80 * 8192 * 128,    // 128个专家路由
		"moe_expert_gate": 80 * 128 * 8192 * 16384,
		"moe_expert_up":   80 * 128 * 8192 * 16384,
		"moe_expert_down": 80 * 128 * 16384 * 8192,
	}
	
	// 分配配置
	for name := range layerParams {
		layerType := guessLayerType(name)
		plan.AssignLayerConfig(name, layerType)
	}
	
	orig, quant := plan.EstimateTotalSize(layerParams)
	
	fmt.Println("=== Hy3 295B 混合精度量化估算 ===")
	fmt.Printf("原始BF16大小: %.0f GB (%.0f MB)\n", orig/1e9, orig/1e6)
	fmt.Printf("量化后大小: %.0f GB (%.0f MB)\n", quant/1e9, quant/1e6)
	fmt.Printf("总压缩率: %.1fx\n", orig/quant)
	
	fmt.Println("\n各层压缩率:")
	for name, cfg := range plan.Layers {
		ratio := float64(layerParams[name]*2) / 
			(float64(layerParams[name]) * 2.0 / float64(cfg.QuantBits))
		fmt.Printf("  %-25s %dbits 压缩率 %.1fx\n", name, cfg.QuantBits, ratio)
	}
}

func guessLayerType(name string) LayerType {
	types := map[string]LayerType{
		"embedding":       Embedding,
		"attention_qkv":   AttentionQKV,
		"attention_output": AttentionOutput,
		"mlp_gate":        MLPGate,
		"mlp_up":          MLPUp,
		"mlp_down":        MLPDown,
		"moe_router":      MoERouter,
		"moe_expert_gate": MoEExpertGate,
		"moe_expert_up":   MoEExpertUp,
		"moe_expert_down": MoEExpertDown,
	}
	if t, ok := types[name]; ok {
		return t
	}
	return AttentionQKV
}

# mtp_speculative_decoder.py - Hy3 MTP投机解码实现

import numpy as np
from typing import List, Optional, Tuple
import time

class MTPDraftModel:
    """
    Hy3 MTP投机解码的draft模型
    轻量级Transformer，用于快速预测候选token
    """
    
    def __init__(self, 
                 vocab_size: int = 152064,
                 d_model: int = 2048,
                 num_layers: int = 4,
                 num_heads: int = 16):
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.num_layers = num_layers
        self.num_heads = num_heads
        
        # 模拟draft模型参数
        self.params_count = d_model * d_model * num_layers * 4
        print(f"Draft模型参数量: {self.params_count/1e6:.1f}M")
        print(f"约为Hy3主模型的: {self.params_count/295e9*100:.2f}%")
    
    def predict_tokens(self, 
                       input_ids: np.ndarray,
                       num_predict: int = 5) -> np.ndarray:
        """
        快速预测多个候选token
        Args:
            input_ids: 输入token序列 [seq_len]
            num_predict: 预测的token数量
        Returns:
            candidate_tokens: 候选token序列 [num_predict]
        """
        # 模拟draft模型快速解码
        # 在实际系统中，这里是一个4层小Transformer
        tokens = []
        current = input_ids[-1] if len(input_ids) > 0 else 0
        
        for _ in range(num_predict):
            # 模拟draft模型预测
            # 实际实现中会使用完整的Transformer前向
            next_token = (current * 31 + 7) % self.vocab_size
            tokens.append(next_token)
            current = next_token
        
        return np.array(tokens, dtype=np.int64)

class Hy3Model:
    """
    Hy3主模型（模拟）
    用于并行验证draft模型生成的候选token
    """
    
    def __init__(self, model_size_gb: float = 85.5):
        self.model_size_gb = model_size_gb
        self._simulated_latency = {
            "single_token_ms": 150.0,    # 单token解码延迟
            "verify_5_tokens_ms": 200.0,  # 并行验证5个token延迟
            "verify_10_tokens_ms": 280.0, # 并行验证10个token延迟
        }
    
    def verify_tokens(self, 
                      context: np.ndarray,
                      candidate_tokens: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """
        并行验证候选token
        Args:
            context: 上下文token [seq_len]
            candidate_tokens: 候选token [num_candidates]
        Returns:
            accepted: 被接受的token序列
            logits: 验证后的logits
        """
        # 模拟并行验证
        # 实际系统中，这里使用Hy3的主模型一次前向传播
        num_candidates = len(candidate_tokens)
        
        # 模拟接受率
        acceptance_rate = min(0.65, 0.4 + 0.05 * num_candidates)
        
        # 模拟验证结果
        accepted_mask = np.random.random(num_candidates) < acceptance_rate
        
        # 找到第一个被拒绝的位置
        reject_idx = len(accepted_mask)
        for i, accepted in enumerate(accepted_mask):
            if not accepted:
                reject_idx = i
                break
        
        accepted = candidate_tokens[:reject_idx]
        
        return accepted, np.random.randn(self.vocab_size)

class MTPDecoder:
    """
    MTP投机解码器
    协调draft模型和主模型，实现加速解码
    """
    
    def __init__(self, 
                 draft_model: MTPDraftModel,
                 main_model: Hy3Model,
                 num_speculate: int = 5,
                 max_seq_len: int = 4096):
        self.draft = draft_model
        self.main = main_model
        self.num_speculate = num_speculate
        self.max_seq_len = max_seq_len
        
        # 统计
        self.total_tokens = 0
        self.draft_tokens = 0
        self.accepted_tokens = 0
        self.rejected_tokens = 0
        self.total_time_ms = 0.0
    
    def generate(self, 
                 prompt: np.ndarray,
                 max_new_tokens: int = 256,
                 temperature: float = 0.7) -> List[int]:
        """
        使用MTP投机解码生成文本
        """
        generated = list(prompt)
        start_time = time.time()
        
        while len(generated) - len(prompt) < max_new_tokens:
            context = np.array(generated[-512:], dtype=np.int64)
            
            # 1. Draft模型快速预测
            candidates = self.draft.predict_tokens(
                context, self.num_speculate
            )
            self.draft_tokens += self.num_speculate
            
            # 2. 主模型并行验证
            accepted, _ = self.main.verify_tokens(context, candidates)
            
            self.accepted_tokens += len(accepted)
            self.rejected_tokens += self.num_speculate - len(accepted)
            
            # 3. 添加到生成序列
            for token in accepted:
                generated.append(int(token))
            
            self.total_tokens += len(accepted) + 1  # +1 for the verified token
            
            # 如果所有候选都被拒绝，至少接受一个
            if len(accepted) == 0:
                generated.append(int(candidates[0]))
                self.total_tokens += 1
        
        self.total_time_ms = (time.time() - start_time) * 1000
        
        return generated[len(prompt):]
    
    def get_stats(self) -> dict:
        """获取解码统计"""
        acceptance_rate = self.accepted_tokens / max(self.draft_tokens, 1)
        speedup = self.total_tokens / max(self.draft_tokens, 1) * self.num_speculate
        
        return {
            "total_tokens_generated": self.total_tokens,
            "draft_tokens_proposed": self.draft_tokens,
            "accepted_tokens": self.accepted_tokens,
            "rejected_tokens": self.rejected_tokens,
            "acceptance_rate": acceptance_rate,
            "theoretical_speedup": speedup,
            "total_time_ms": self.total_time_ms,
            "tokens_per_second": self.total_tokens / (self.total_time_ms / 1000),
        }

# 模拟MTP解码性能
draft = MTPDraftModel()
main = Hy3Model(model_size_gb=85.5)
decoder = MTPDecoder(draft, main, num_speculate=5)

prompt = np.array([1, 2, 3, 4, 5], dtype=np.int64)
output = decoder.generate(prompt, max_new_tokens=100)

stats = decoder.get_stats()
print(f"\n=== MTP投机解码性能 ===")
print(f"生成token数: {stats['total_tokens_generated']}")
print(f"Draft提出token数: {stats['draft_tokens_proposed']}")
print(f"接受率: {stats['acceptance_rate']:.1%}")
print(f"理论加速比: {stats['theoretical_speedup']:.1f}x")
print(f"生成速度: {stats['tokens_per_second']:.1f} tokens/s")

# hy3_mtp_patch.py - Hy3 MTP patch for llama.cpp

class Hy3MTPConfig:
    """Hy3 MTP配置"""
    
    def __init__(self):
        # MTP参数
        self.num_draft_tokens = 5       # 每次预测的token数
        self.draft_model_ratio = 0.02    # draft模型与主模型大小比
        self.draft_layers = 4            # draft模型层数
        self.draft_heads = 16            # draft模型注意力头数
        
        # 接受率相关
        self.target_acceptance = 0.60    # 目标接受率
        self.min_acceptance = 0.30       # 最低接受率
        self.dynamic_adjust = True       # 动态调整num_speculate
    
    def adjust_speculate_count(self, 
                               current_acceptance: float) -> int:
        """
        根据当前接受率动态调整预测数量
        """
        if not self.dynamic_adjust:
            return self.num_draft_tokens
        
        if current_acceptance < self.min_acceptance:
            return max(2, self.num_draft_tokens - 1)
        elif current_acceptance > self.target_acceptance:
            return min(10, self.num_draft_tokens + 1)
        else:
            return self.num_draft_tokens

# MTP解码整体流程
def mtp_decode_benchmark():
    """MTP解码性能基准测试"""
    
    configs = [
        ("无MTP (单token)", 0, 1.0, 150.0),
        ("MTP-3", 3, 0.55, 180.0), 
        ("MTP-5", 5, 0.60, 200.0),
        ("MTP-8", 8, 0.62, 260.0),
        ("MTP-10", 10, 0.63, 300.0),
    ]
    
    print(f"{'配置':<20s} {'接受率':<8s} {'验证延迟':<10s} {'总吞吐':<10s}")
    print("-" * 50)
    
    for name, n_spec, accept_rate, verify_latency in configs:
        if n_spec == 0:
            # 无MTP：每步150ms生成1个token
            throughput = 1000.0 / 150.0
        else:
            # 有MTP：每步平均生成 (1 + n_spec * accept_rate) 个token
            avg_tokens = 1 + n_spec * accept_rate
            # MTP延迟 = draft延迟 + 验证延迟
            draft_latency = 5 * n_spec  # 5ms per draft token
            total_latency = draft_latency + verify_latency
            throughput = avg_tokens / total_latency * 1000
        
        print(f"{name:<20s} {accept_rate:<8.1%} {verify_latency:<10.0f}ms {throughput:<10.1f} tok/s")

mtp_decode_benchmark()

#!/bin/bash
# setup_hy3_llama.sh - Hy3 llama.cpp 构建与部署脚本

set -e

# 配置
LLAMA_COMMIT="a1b2c3d4e5f6..."  # 经过验证的llama.cpp基线commit
HY3_MODEL_PATH="/models/Hy3-GGUF"
OUTPUT_DIR="./build"

# 颜色输出
GREEN='\033[0;32m'
BLUE='\033[0;34m'
NC='\033[0m'

echo -e "${BLUE}[Hy3 Deploy] 开始构建Hy3 llama.cpp环境${NC}"

# 1. 克隆llama.cpp
if [ ! -d "llama.cpp-hyv3" ]; then
    git clone https://github.com/ggerganov/llama.cpp.git llama.cpp-hyv3
    cd llama.cpp-hyv3
    git checkout ${LLAMA_COMMIT}
    cd ..
    echo -e "${GREEN}[OK] llama.cpp 已克隆${NC}"
fi

# 2. 应用Hy3架构patch
cd llama.cpp-hyv3
if [ ! -f "patches/01_hy3_arch.patch" ]; then
    mkdir -p patches
    
    # Patch 1: Hy3基础架构支持
    cat > patches/01_hy3_arch.patch << 'PATCH_EOF'
--- a/llama.cpp
+++ b/llama.cpp
@@ -1500,6 +1500,52 @@
     }
 };
 
+// Hy3 MoE架构支持
+struct hy3_moe_layer {
+    struct ggml_tensor *gate_proj;
+    struct ggml_tensor *up_proj;
+    struct ggml_tensor *down_proj;
+    int num_experts;
+    int top_k;
+};
+
+static bool llm_load_tensors_hy3(
+    struct llama_model *model,
+    const struct llama_model_loader &ml) {
+    // Hy3专用张量加载逻辑
+    // 处理MoE专家的稀疏加载
+    fprintf(stderr, "hy3: loading MoE tensors with sparse strategy\n");
+    return true;
+}
PATCH_EOF
    echo -e "${GREEN}[OK] Patch 1: Hy3基础架构${NC}"
fi

# 3. 编译
mkdir -p ${OUTPUT_DIR}
cd ${OUTPUT_DIR}

cmake .. -DLLAMA_CUDA=ON \
         -DLLAMA_CUDA_F16=ON \
         -DLLAMA_K_QUANTS=ON \
         -DBUILD_SHARED_LIBS=OFF

make -j$(nproc) llama-cli llama-server

echo -e "${GREEN}[OK] 编译完成${NC}"

# 4. 验证部署
echo -e "${BLUE}[Hy3 Deploy] 验证量化模型${NC}"
./bin/llama-cli \
    -m ${HY3_MODEL_PATH}/Hy3-Q4_K_M-mtp.gguf \
    --prompt "Hello, Hy3!" \
    -n 32 \
    -t 8 \
    -c 4096 \
    --mtp 5 \
    --temp 0.7

echo -e "${GREEN}[OK] Hy3部署验证完成${NC}"

# hy3_server.py - Hy3量化模型推理服务

import subprocess
import json
import time
from typing import Optional, List, Dict, Any
from dataclasses import dataclass
from enum import Enum

class QuantVersion(Enum):
    IQ1_M = "iq1_m"      # 1bit, 85.5GB, 单卡
    Q4_K_M = "q4_k_m"    # 4bit, 169.9GB, 双卡
    GPTQ_INT4 = "gptq"   # GPTQ Int4, 适配vLLM

@dataclass
class Hy3Config:
    model_path: str
    quant_version: QuantVersion
    num_gpu_layers: int = 80
    context_length: int = 4096
    batch_size: int = 1
    num_threads: int = 8
    use_mtp: bool = True
    mtp_num: int = 5
    temperature: float = 0.7
    top_p: float = 0.9
    max_tokens: int = 2048

class Hy3InferenceServer:
    """Hy3推理服务"""
    
    def __init__(self, config: Hy3Config):
        self.config = config
        self.process: Optional[subprocess.Popen] = None
        self._start_time = 0
        self._request_count = 0
        self._total_latency = 0.0
    
    def start(self):
        """启动推理服务"""
        cmd = self._build_command()
        self.process = subprocess.Popen(
            cmd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            bufsize=1
        )
        self._start_time = time.time()
        print(f"[Hy3] 服务已启动 (量化: {self.config.quant_version.value})")
        print(f"[Hy3] 模型: {self.config.model_path}")
        print(f"[Hy3] MTP: {'开启' if self.config.use_mtp else '关闭'}")
    
    def _build_command(self) -> List[str]:
        """构建llama.cpp命令"""
        cmd = [
            "./llama.cpp-hyv3/build/bin/llama-cli",
            "-m", self.config.model_path,
            "--prompt", "",  # 通过stdin传入
            "-n", str(self.config.max_tokens),
            "-t", str(self.config.num_threads),
            "-c", str(self.config.context_length),
            "--temp", str(self.config.temperature),
            "--top-p", str(self.config.top_p),
            "--repeat-penalty", "1.1",
        ]
        
        if self.config.use_mtp:
            cmd.extend(["--mtp", str(self.config.mtp_num)])
        
        return cmd
    
    def generate(self, prompt: str) -> Dict[str, Any]:
        """
        生成文本
        Args:
            prompt: 输入提示词
        Returns:
            response: 包含生成结果和性能统计
        """
        if not self.process:
            raise RuntimeError("服务未启动")
        
        start = time.time()
        
        # 写入prompt
        self.process.stdin.write(prompt + "\n")
        self.process.stdin.flush()
        
        # 读取输出
        output = []
        while True:
            line = self.process.stdout.readline()
            if not line or line.strip() == "":
                break
            output.append(line.strip())
        
        elapsed = time.time() - start
        
        self._request_count += 1
        self._total_latency += elapsed
        
        return {
            "text": "\n".join(output),
            "latency_ms": elapsed * 1000,
            "tokens_per_second": len(output) / elapsed if elapsed > 0 else 0,
            "model": f"Hy3-{self.config.quant_version.value}",
            "mtp_enabled": self.config.use_mtp,
        }
    
    def get_stats(self) -> Dict[str, Any]:
        """获取服务统计"""
        uptime = time.time() - self._start_time
        avg_latency = self._total_latency / max(self._request_count, 1) * 1000
        
        return {
            "uptime_seconds": uptime,
            "total_requests": self._request_count,
            "avg_latency_ms": avg_latency,
            "quant_version": self.config.quant_version.value,
            "model_size_gb": 85.5 if self.config.quant_version == QuantVersion.IQ1_M else 169.9,
            "gpu_required": "1x96GB" if self.config.quant_version == QuantVersion.IQ1_M else "2x96GB",
        }

# 使用示例
config = Hy3Config(
    model_path="/models/Hy3-GGUF/Hy3-IQ1_M-mtp.gguf",
    quant_version=QuantVersion.IQ1_M,
    use_mtp=True,
    mtp_num=5,
    context_length=4096,
)
server = Hy3InferenceServer(config)

# 测试推理
test_prompts = [
    "用Python实现一个快速排序算法，并解释时间复杂度和空间复杂度。",
    "What is the difference between Mixture of Experts and Dense Transformer?",
    "写一个Go语言并发爬虫，从多个URL同时抓取数据。",
]

for prompt in test_prompts:
    result = server.generate(prompt)
    print(f"\nPrompt: {prompt[:50]}...")
    print(f"延迟: {result['latency_ms']:.0f}ms")
    print(f"速度: {result['tokens_per_second']:.1f} tok/s")
    print(f"输出预览: {result['text'][:100]}...")

stats = server.get_stats()
print(f"\n=== 服务统计 ===")
print(f"运行时间: {stats['uptime_seconds']:.0f}s")
print(f"请求数: {stats['total_requests']}")
print(f"平均延迟: {stats['avg_latency_ms']:.0f}ms")

# quality_eval.py - 量化质量评估

import numpy as np
from sklearn.metrics import mean_squared_error
from scipy.stats import pearsonr, spearmanr

class QuantizationQualityEval:
    """量化质量评估"""
    
    def __init__(self):
        self.metrics = {}
    
    def evaluate_layer(self, 
                       original: np.ndarray,
                       quantized: np.ndarray,
                       layer_name: str) -> dict:
        """评估单层量化质量"""
        
        # MSE
        mse = mean_squared_error(original.flatten(), quantized.flatten())
        
        # 信噪比
        signal_power = np.var(original.flatten())
        noise_power = mse
        snr = 10 * np.log10(signal_power / noise_power) if noise_power > 0 else float('inf')
        
        # 相关性
        orig_flat = original.flatten()
        quant_flat = quantized.flatten()
        
        # 采样（过大时）
        if len(orig_flat) > 100000:
            idx = np.random.choice(len(orig_flat), 100000, replace=False)
            orig_flat = orig_flat[idx]
            quant_flat = quant_flat[idx]
        
        pearson_r, _ = pearsonr(orig_flat, quant_flat)
        spearman_r, _ = spearmanr(orig_flat, quant_flat)
        
        # 最大相对误差
        max_rel_error = np.max(np.abs(orig_flat - quant_flat) / 
                              (np.abs(orig_flat) + 1e-10))
        
        result = {
            "layer": layer_name,
            "mse": float(mse),
            "snr_db": float(snr),
            "pearson_r": float(pearson_r),
            "spearman_r": float(spearman_r),
            "max_relative_error": float(max_rel_error),
        }
        
        self.metrics[layer_name] = result
        return result

# 模拟评估
eval = QuantizationQualityEval()

# 模拟各层量化前后的权重分布
layers = {
    "attention_qkv": (4096, 12288),
    "moe_router": (8192, 128),
    "moe_expert_gate_0": (8192, 16384),
    "moe_expert_down_0": (16384, 8192),
}

for name, shape in layers.items():
    original = np.random.randn(*shape).astype(np.float32) * 0.1
    # 模拟1bit量化后的效果
    quantized = np.sign(original) * np.max(np.abs(original)) * 0.01
    quantized += np.random.randn(*shape).astype(np.float32) * 0.005
    
    result = eval.evaluate_layer(original, quantized, name)
    print(f"{name:>25s} | MSE={result['mse']:.6f} | SNR={result['snr_db']:.1f}dB | "
          f"Pearson={result['pearson_r']:.4f}")