Liquid AI LFM2.5-2.6B——26亿参数端侧大模型越级碾压的架构革命与部署实践

"""
NAS搜索模拟：Liquid AI可能使用的架构搜索策略
使用进化算法在约束空间中找到最优层配置
"""
import numpy as np
from dataclasses import dataclass, field
from typing import List, Optional
import random
import math
import json

@dataclass
class NASConfig:
    """神经架构搜索配置空间"""
    total_layers: int = 30
    hidden_dim: int = 2048
    vocab_size: int = 128000
    
    # 搜索空间
    conv_kernel_sizes: List[int] = field(default_factory=lambda: [3, 5, 7])
    gqa_head_options: List[int] = field(default_factory=lambda: [8, 16, 32])
    ffn_scale_options: List[float] = field(default_factory=lambda: [2.5, 3.0, 3.5, 4.0])
    min_gqa_layers: int = 4
    max_gqa_layers: int = 12

@dataclass
class Architecture:
    """单个架构编码"""
    layer_types: List[str]  # 'conv' 或 'gqa'
    conv_kernel_size: int
    gqa_kv_heads: int
    ffn_scale: float
    
    def compute_kv_cache_mb(self, context_len: int = 131072) -> float:
        """估算KV缓存大小（MB）"""
        gqa_count = sum(1 for t in self.layer_types if t == 'gqa')
        if gqa_count == 0:
            return 0.0
        # 每个GQA层: 2 (K+V) * kv_heads * (hidden_dim//q_heads) * context_len * 2bytes
        head_dim = self.hidden_dim // 32  # 固定32个Q头
        bytes_per_layer = 2 * self.gqa_kv_heads * head_dim * context_len * 2
        return (bytes_per_layer * gqa_count) / (1024 * 1024)
    
    def estimate_compute_cost(self) -> float:
        """估算计算成本（相对值）"""
        conv_cost = sum(1 for t in self.layer_types if t == 'conv') * self.conv_kernel_size * 0.3
        gqa_cost = sum(1 for t in self.layer_types if t == 'gqa') * 1.0
        ffn_cost = self.total_layers * self.ffn_scale * 0.4
        return conv_cost + gqa_cost + ffn_cost

def random_architecture(config: NASConfig) -> Architecture:
    """随机生成一个架构"""
    num_gqa = random.randint(config.min_gqa_layers, config.max_gqa_layers)
    num_conv = config.total_layers - num_gqa
    
    # 生成层类型数组，GQA尽量均匀分布
    positions = sorted(random.sample(range(config.total_layers), num_gqa))
    layer_types = ['conv'] * config.total_layers
    for pos in positions:
        layer_types[pos] = 'gqa'
    
    arch = Architecture(
        layer_types=layer_types,
        conv_kernel_size=random.choice(config.conv_kernel_sizes),
        gqa_kv_heads=random.choice(config.gqa_head_options),
        ffn_scale=random.choice(config.ffn_scale_options),
        hidden_dim=config.hidden_dim,
        total_layers=config.total_layers
    )
    return arch

def mutate_architecture(arch: Architecture, config: NASConfig) -> Architecture:
    """变异架构"""
    new_types = arch.layer_types.copy()
    
    # 随机交换一个conv和一个gqa
    if random.random() < 0.3:
        conv_indices = [i for i, t in enumerate(new_types) if t == 'conv']
        gqa_indices = [i for i, t in enumerate(new_types) if t == 'gqa']
        if conv_indices and gqa_indices:
            ci = random.choice(conv_indices)
            gi = random.choice(gqa_indices)
            new_types[ci], new_types[gi] = new_types[gi], new_types[ci]
    
    # 随机修改超参数
    new_kernel = arch.conv_kernel_size
    if random.random() < 0.2:
        new_kernel = random.choice([k for k in config.conv_kernel_sizes if k != arch.conv_kernel_size] or config.conv_kernel_sizes)
    
    new_kv = arch.gqa_kv_heads
    if random.random() < 0.2:
        new_kv = random.choice([h for h in config.gqa_head_options if h != arch.gqa_kv_heads] or config.gqa_head_options)
    
    new_ffn = arch.ffn_scale
    if random.random() < 0.2:
        new_ffn = random.choice([s for s in config.ffn_scale_options if abs(s - arch.ffn_scale) > 0.1] or config.ffn_scale_options)
    
    return Architecture(
        layer_types=new_types,
        conv_kernel_size=new_kernel,
        gqa_kv_heads=new_kv,
        ffn_scale=new_ffn,
        hidden_dim=arch.hidden_dim,
        total_layers=arch.total_layers
    )

def crossover(a1: Architecture, a2: Architecture) -> Architecture:
    """交叉两个架构"""
    child_types = []
    for i in range(len(a1.layer_types)):
        child_types.append(random.choice([a1.layer_types[i], a2.layer_types[i]]))
    
    return Architecture(
        layer_types=child_types,
        conv_kernel_size=random.choice([a1.conv_kernel_size, a2.conv_kernel_size]),
        gqa_kv_heads=random.choice([a1.gqa_kv_heads, a2.gqa_kv_heads]),
        ffn_scale=random.choice([a1.ffn_scale, a2.ffn_scale]),
        hidden_dim=a1.hidden_dim,
        total_layers=a1.total_layers
    )

def fitness(arch: Architecture, target_kv_mb: float = 500.0) -> float:
    """适应度函数：平衡性能与资源约束"""
    kv_cache = arch.compute_kv_cache_mb()
    compute_cost = arch.estimate_compute_cost()
    
    # KV缓存不能超过目标
    if kv_cache > target_kv_mb * 1.5:
        return -float('inf')
    
    # GQA层数越多，长程能力越强（但成本越高）
    gqa_count = sum(1 for t in arch.layer_types if t == 'gqa')
    gqa_benefit = gqa_count * 1.5
    
    # 卷积层提供效率
    conv_count = sum(1 for t in arch.layer_types if t == 'gqa')
    conv_benefit = conv_count * 0.8
    
    # 总得分 = 能力 - 成本
    score = (gqa_benefit + conv_benefit) - compute_cost * 0.3
    
    # 偏好GQA均匀分布（避免所有注意力集中在开头或结尾）
    gqa_positions = [i for i, t in enumerate(arch.layer_types) if t == 'gqa']
    if gqa_positions:
        spread = max(gqa_positions) - min(gqa_positions)
        spread_score = spread / len(arch.layer_types) * 2.0
        score += spread_score
    
    return score

def evolutionary_search(config: NASConfig, 
                       population_size: int = 50,
                       generations: int = 100,
                       elite_ratio: float = 0.2) -> List[Architecture]:
    """
    进化算法搜索最优架构
    模拟Liquid AI的NAS过程
    """
    # 初始化种群
    population = [random_architecture(config) for _ in range(population_size)]
    
    best_archs = []
    
    for gen in range(generations):
        # 计算适应度
        scored = [(arch, fitness(arch)) for arch in population]
        scored.sort(key=lambda x: x[1], reverse=True)
        
        # 记录最优
        if scored[0][1] > -float('inf'):
            best_archs.append(scored[0][0])
        
        # 精英选择
        elite_count = int(population_size * elite_ratio)
        elites = [arch for arch, _ in scored[:elite_count]]
        
        # 填充下一代
        next_gen = elites.copy()
        while len(next_gen) < population_size:
            parent1 = random.choice(elites)
            if random.random() < 0.7:
                parent2 = random.choice(elites)
                child = crossover(parent1, parent2)
            else:
                child = parent1
            
            # 变异概率
            if random.random() < 0.4:
                child = mutate_architecture(child, config)
            
            next_gen.append(child)
        
        population = next_gen
        
        if (gen + 1) % 20 == 0:
            print(f"Generation {gen+1}: Best fitness = {scored[0][1]:.2f}, "
                  f"GQA layers = {sum(1 for t in scored[0][0].layer_types if t == 'gqa')}, "
                  f"KV cache = {scored[0][0].compute_kv_cache_mb():.1f}MB")
    
    return best_archs

if __name__ == "__main__":
    config = NASConfig()
    print("开始NAS架构搜索模拟...")
    print(f"搜索空间: {config.total_layers}层, "
          f"卷积核={config.conv_kernel_sizes}, "
          f"GQA头数={config.gqa_head_options}")
    print()
    
    best_archs = evolutionary_search(config, population_size=60, generations=80)
    
    if best_archs:
        final = best_archs[-1]
        gqa_count = sum(1 for t in final.layer_types if t == 'gqa')
        conv_count = sum(1 for t in final.layer_types if t == 'conv')
        print(f"\n最佳架构:")
        print(f"  总层数: {final.total_layers}")
        print(f"  ConvBlock: {conv_count}层")
        print(f"  GQA: {gqa_count}层")
        print(f"  卷积核大小: {final.conv_kernel_size}")
        print(f"  KV头数: {final.gqa_kv_heads}")
        print(f"  FFN缩放: {final.ffn_scale}")
        print(f"  KV缓存: {final.compute_kv_cache_mb():.1f}MB")
        print(f"  层分布: {''.join('C' if t == 'conv' else 'A' for t in final.layer_types)}")

"""
MOPD (Multi-Domain On-Policy Distillation) 实现
多域同策略蒸馏的核心算法
"""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import Dict, List, Optional, Callable
from dataclasses import dataclass
import math

@dataclass
class MOPDConfig:
    """MOPD配置"""
    vocab_size: int = 128000
    hidden_dim: int = 2048
    num_layers: int = 30
    num_teachers: int = 4  # 数学、代码、工具、推理
    kl_weight: float = 0.5
    ce_weight: float = 1.0
    distill_temperature: float = 2.0
    domain_weights: List[float] = None

class MOPDDistiller:
    """
    多域同策略蒸馏器
    核心思想：在策略采样(on-policy)过程中，同时用教师和学生生成logits，
    然后通过KL散度+交叉熵的联合损失进行蒸馏
    """
    
    def __init__(self, config: MOPDConfig):
        self.config = config
        if config.domain_weights is None:
            self.config.domain_weights = [1.0, 1.0, 1.0, 1.0]
    
    def compute_distill_loss(
        self,
        student_logits: torch.Tensor,  # [batch, seq_len, vocab]
        teacher_logits_list: List[torch.Tensor],  # 4个教师的logits
        labels: torch.Tensor,  # [batch, seq_len]
        domain_ids: torch.Tensor,  # [batch], 每个样本所属领域
        attention_mask: Optional[torch.Tensor] = None
    ) -> Dict[str, torch.Tensor]:
        """
        计算蒸馏损失
        
        Args:
            student_logits: 学生模型输出logits
            teacher_logits_list: 四个教师模型的logits
            labels: 目标token ids
            domain_ids: 领域标签 (0=math, 1=code, 2=tool, 3=reasoning)
            attention_mask: 注意力掩码
        """
        batch_size, seq_len, vocab_size = student_logits.shape
        
        if attention_mask is None:
            attention_mask = torch.ones(batch_size, seq_len, dtype=torch.bool)
        
        # 1. 交叉熵损失（标准语言建模）
        ce_loss = F.cross_entropy(
            student_logits.view(-1, vocab_size),
            labels.view(-1),
            reduction='none'
        ).view(batch_size, seq_len)
        ce_loss = (ce_loss * attention_mask).sum() / attention_mask.sum()
        
        # 2. KL散度损失（蒸馏）
        # 对每个样本，只使用对应领域的教师
        kl_loss = 0.0
        student_log_probs = F.log_softmax(
            student_logits / self.config.distill_temperature, dim=-1
        )
        
        for domain_idx in range(self.config.num_teachers):
            domain_mask = (domain_ids == domain_idx)
            if domain_mask.sum() == 0:
                continue
            
            # 获取该领域教师logits
            teacher_logits = teacher_logits_list[domain_idx]
            
            # 教师概率分布
            teacher_probs = F.softmax(
                teacher_logits / self.config.distill_temperature, dim=-1
            )
            
            # KL(P_teacher || P_student)
            domain_kl = F.kl_div(
                student_log_probs[domain_mask],
                teacher_probs[domain_mask],
                reduction='sum',
                log_target=False
            )
            
            domain_weight = self.config.domain_weights[domain_idx]
            kl_loss += domain_weight * domain_kl
        
        kl_loss = kl_loss / attention_mask.sum()
        
        # 3. 联合损失
        total_loss = (self.config.ce_weight * ce_loss + 
                     self.config.kl_weight * kl_loss * 
                     (self.config.distill_temperature ** 2))
        
        return {
            'total_loss': total_loss,
            'ce_loss': ce_loss,
            'kl_loss': kl_loss,
        }

class GRPOTrainer:
    """
    GRPO (Group Relative Policy Optimization) Agent训练器
    用于Agentic RL阶段
    """
    
    def __init__(
        self,
        model: nn.Module,
        tokenizer: Callable,
        clip_epsilon: float = 0.2,
        kl_coeff: float = 0.01,
        group_size: int = 8
    ):
        self.model = model
        self.tokenizer = tokenizer
        self.clip_epsilon = clip_epsilon
        self.kl_coeff = kl_coeff
        self.group_size = group_size
    
    @dataclass
    class Trajectory:
        """单次Agent交互轨迹"""
        observations: List[str]
        actions: List[str]
        tool_calls: List[Dict]
        rewards: List[float]
        log_probs: List[float]
    
    def compute_grpo_loss(
        self,
        trajectories: List[Trajectory],
        old_log_probs: torch.Tensor,
        advantages: torch.Tensor
    ) -> torch.Tensor:
        """
        计算GRPO损失
        
        GRPO = -E[ min(r * A, clip(r, 1-ε, 1+ε) * A) ]
        其中 r = exp(log_prob_new - log_prob_old)
        """
        # 当前策略的log概率
        current_log_probs = self._compute_log_probs(trajectories)
        
        # 概率比
        ratios = torch.exp(current_log_probs - old_log_probs)
        
        # 裁剪后的替代目标
        surr1 = ratios * advantages
        surr2 = torch.clamp(ratios, 
                           1.0 - self.clip_epsilon, 
                           1.0 + self.clip_epsilon) * advantages
        
        policy_loss = -torch.min(surr1, surr2).mean()
        
        # KL惩罚（防止策略偏离太远）
        kl_div = (old_log_probs - current_log_probs).mean()
        
        return policy_loss + self.kl_coeff * kl_div
    
    def _compute_log_probs(self, trajectories: List[Trajectory]) -> torch.Tensor:
        """计算轨迹的log概率"""
        # 简化实现：实际中需要完整的模型前向传播
        log_probs = []
        for traj in trajectories:
            for log_prob in traj.log_probs:
                log_probs.append(log_prob)
        return torch.tensor(log_probs)


def run_mopd_pipeline():
    """
    演示完整的MOPD训练流程
    """
    config = MOPDConfig()
    distiller = MOPDDistiller(config)
    
    print("MOPD训练流程演示")
    print("=" * 60)
    print(f"词汇表大小: {config.vocab_size}")
    print(f"教师模型数量: {config.num_teachers}")
    print(f"蒸馏温度: {config.distill_temperature}")
    print(f"KL权重: {config.kl_weight}, CE权重: {config.ce_weight}")
    print()
    
    # 模拟训练数据
    batch_size = 4
    seq_len = 512
    dummy_student_logits = torch.randn(batch_size, seq_len, config.vocab_size)
    dummy_teacher_logits = [
        torch.randn(batch_size, seq_len, config.vocab_size) 
        for _ in range(config.num_teachers)
    ]
    dummy_labels = torch.randint(0, config.vocab_size, (batch_size, seq_len))
    dummy_domains = torch.randint(0, config.num_teachers, (batch_size,))
    
    # 计算损失
    losses = distiller.compute_distill_loss(
        dummy_student_logits,
        dummy_teacher_logits,
        dummy_labels,
        dummy_domains
    )
    
    print(f"总损失: {losses['total_loss']:.4f}")
    print(f"交叉熵损失: {losses['ce_loss']:.4f}")
    print(f"KL散度损失: {losses['kl_loss']:.4f}")
    print(f"蒸馏温度^2缩放: {config.distill_temperature ** 2:.2f}")
    
    # 演示GRPO
    grpo_trainer = GRPOTrainer(model=None, tokenizer=None)
    print(f"\nGRPO组大小: {grpo_trainer.group_size}")
    print(f"裁剪ε: {grpo_trainer.clip_epsilon}")
    print(f"KL系数: {grpo_trainer.kl_coeff}")

if __name__ == "__main__":
    run_mopd_pipeline()

// main.go - 端侧推理引擎核心实现
package main

import (
	"encoding/binary"
	"fmt"
	"math"
	"os"
	"sync"
	"time"
)

// ============================================================
// 核心数据类型
// ============================================================

// ModelConfig 模型配置
type ModelConfig struct {
	HiddenDim     int     // 2048
	NumLayers     int     // 30
	NumQHeads     int     // 32
	NumKVHeads    int     // 8
	ConvKernel    int     // 3
	VocabSize     int     // 128000
	MaxSeqLen     int     // 131072
	FFNScale      float64 // ~5.25 (10752/2048)
	RoPETheta     float64 // 10,000,000
	LayerTypes    []string // "conv" 或 "gqa"
}

// DefaultConfig LFM2.5-2.6B默认配置
func DefaultConfig() ModelConfig {
	// 30层配置：22个ConvBlock + 8个GQA
	layerTypes := make([]string, 30)
	for i := 0; i < 30; i++ {
		// 模拟LFM2.5-2.6B的层分布
		// 每2-3个ConvBlock插入1个GQA
		gqaPositions := map[int]bool{
			2: true, 5: true, 8: true, 11: true,
			14: true, 18: true, 22: true, 26: true,
		}
		if gqaPositions[i] {
			layerTypes[i] = "gqa"
		} else {
			layerTypes[i] = "conv"
		}
	}

	return ModelConfig{
		HiddenDim:  2048,
		NumLayers:  30,
		NumQHeads:  32,
		NumKVHeads: 8,
		ConvKernel: 3,
		VocabSize:  128000,
		MaxSeqLen:  131072,
		FFNScale:   5.25,
		RoPETheta:  10000000.0,
		LayerTypes: layerTypes,
	}
}

// KVCache KV缓存（GQA优化版）
type KVCache struct {
	keys   [][][]float32 // [layer][head][seq_len, head_dim]
	values [][][]float32
	mu     sync.RWMutex
}

// NewKVCache 创建KV缓存
func NewKVCache(config ModelConfig) *KVCache {
	numGQALayers := 0
	for _, t := range config.LayerTypes {
		if t == "gqa" {
			numGQALayers++
		}
	}

	headDim := config.HiddenDim / config.NumQHeads // 64

	kc := &KVCache{
		keys:   make([][][]float32, numGQALayers),
		values: make([][][]float32, numGQALayers),
	}

	for l := 0; l < numGQALayers; l++ {
		kc.keys[l] = make([][]float32, config.NumKVHeads)
		kc.values[l] = make([][]float32, config.NumKVHeads)
		for h := 0; h < config.NumKVHeads; h++ {
			// 初始容量，后续动态扩展
			kc.keys[l][h] = make([]float32, 0, config.MaxSeqLen*headDim)
			kc.values[l][h] = make([]float32, 0, config.MaxSeqLen*headDim)
		}
	}

	return kc
}

// Append 追加K/V到缓存
func (kc *KVCache) Append(layerIdx int, keys, values [][]float32) {
	kc.mu.Lock()
	defer kc.mu.Unlock()

	for h := 0; h < len(keys); h++ {
		kc.keys[layerIdx][h] = append(kc.keys[layerIdx][h], keys[h]...)
		kc.values[layerIdx][h] = append(kc.values[layerIdx][h], values[h]...)
	}
}

// Get 获取缓存中的K/V
func (kc *KVCache) Get(layerIdx int, numTokens int) ([][]float32, [][]float32) {
	kc.mu.RLock()
	defer kc.mu.RUnlock()

	keys := make([][]float32, len(kc.keys[layerIdx]))
	values := make([][]float32, len(kc.values[layerIdx]))

	for h := 0; h < len(kc.keys[layerIdx]); h++ {
		// 只返回最近的numTokens
		headDim := 64
		start := len(kc.keys[layerIdx][h]) - numTokens*headDim
		if start < 0 {
			start = 0
		}
		keys[h] = kc.keys[layerIdx][h][start:]
		values[h] = kc.values[layerIdx][h][start:]
	}

	return keys, values
}

// MemoryUsage 返回KV缓存内存使用量（MB）
func (kc *KVCache) MemoryUsage() float64 {
	kc.mu.RLock()
	defer kc.mu.RUnlock()

	var totalBytes int64
	for l := range kc.keys {
		for h := range kc.keys[l] {
			totalBytes += int64(len(kc.keys[l][h])) * 4 // float32
			totalBytes += int64(len(kc.values[l][h])) * 4
		}
	}
	return float64(totalBytes) / (1024 * 1024)
}

// ============================================================
// RoPE位置编码
// ============================================================

// RoPECache 旋转位置编码缓存
type RoPECache struct {
	sin [][]float32
	cos [][]float32
}

// NewRoPECache 预计算RoPE缓存
func NewRoPECache(config ModelConfig, maxSeqLen int) *RoPECache {
	headDim := config.HiddenDim / config.NumQHeads
	rope := &RoPECache{
		sin: make([][]float32, maxSeqLen),
		cos: make([][]float32, maxSeqLen),
	}

	for pos := 0; pos < maxSeqLen; pos++ {
		rope.sin[pos] = make([]float32, headDim)
		rope.cos[pos] = make([]float32, headDim)
		for d := 0; d < headDim; d += 2 {
			theta := math.Pow(config.RoPETheta, float64(-d)/float64(headDim))
			angle := float64(pos) * theta
			rope.sin[pos][d] = float32(math.Sin(angle))
			rope.cos[pos][d] = float32(math.Cos(angle))
			if d+1 < headDim {
				rope.sin[pos][d+1] = float32(math.Sin(angle))
				rope.cos[pos][d+1] = float32(math.Cos(angle))
			}
		}
	}

	return rope
}

// ApplyRoPE 对query/key应用旋转位置编码
func ApplyRoPE(x []float32, pos int, rope *RoPECache, headDim int) []float32 {
	result := make([]float32, len(x))
	copy(result, x)

	for h := 0; h < len(x)/headDim; h++ {
		offset := h * headDim
		for d := 0; d < headDim; d += 2 {
			i := offset + d
			j := offset + d + 1
			if j >= len(x) {
				break
			}
			sin := rope.sin[pos][d]
			cos := rope.cos[pos][d]
			result[i] = x[i]*cos - x[j]*sin
			result[j] = x[i]*sin + x[j]*cos
		}
	}

	return result
}

// ============================================================
// GQA注意力实现
// ============================================================

// GQAConfig GQA层配置
type GQAConfig struct {
	HiddenDim  int
	NumQHeads  int
	NumKVHeads int
	HeadDim    int
}

// GQALayer GQA注意力层
type GQALayer struct {
	config GQAConfig
	// 权重（简化表示）
	wQ, wK, wV, wO [][]float32
}

// NewGQALayer 创建GQA层
func NewGQALayer(cfg GQAConfig) *GQALayer {
	// 初始化权重（实际应该从模型加载）
	return &GQALayer{
		config: cfg,
	}
}

// Forward GQA前向传播
func (gqa *GQALayer) Forward(
	x []float32,
	pos int,
	kvCache *KVCache,
	layerIdx int,
	rope *RoPECache,
) []float32 {
	cfg := gqa.config
	headDim := cfg.HeadDim

	// 模拟QKV投影
	q := make([]float32, cfg.NumQHeads*headDim)
	k := make([]float32, cfg.NumKVHeads*headDim)
	v := make([]float32, cfg.NumKVHeads*headDim)
	for i := range q {
		q[i] = x[i%len(x)] * 0.1 // 简化的投影
	}
	for i := range k {
		k[i] = x[i%len(x)] * 0.1
	}
	for i := range v {
		v[i] = x[i%len(x)] * 0.1
	}

	// 应用RoPE
	q = ApplyRoPE(q, pos, rope, headDim)
	k = ApplyRoPE(k, pos, rope, headDim)

	// 更新KV缓存
	kHeads := make([][]float32, cfg.NumKVHeads)
	vHeads := make([][]float32, cfg.NumKVHeads)
	for h := 0; h < cfg.NumKVHeads; h++ {
		offset := h * headDim
		kHeads[h] = k[offset : offset+headDim]
		vHeads[h] = v[offset : offset+headDim]
	}
	kvCache.Append(layerIdx, kHeads, vHeads)

	// 获取完整KV缓存（用于注意力计算）
	allK, allV := kvCache.Get(layerIdx, pos+1)

	// GQA：每个Q头对应一个KV头
	groupSize := cfg.NumQHeads / cfg.NumKVHeads
	output := make([]float32, len(x))

	for qh := 0; qh < cfg.NumQHeads; qh++ {
		kvh := qh / groupSize
		qOffset := qh * headDim
		kHead := allK[kvh]
		vHead := allV[kvh]

		// 简化的注意力计算
		score := float32(0)
		for d := 0; d < headDim; d++ {
			score += q[qOffset+d] * kHead[(pos*headDim)+d]
		}
		score /= float32(math.Sqrt(float64(headDim)))

		attn := float32(math.Exp(float64(score)))
		for d := 0; d < headDim; d++ {
			output[qOffset+d] = attn * vHead[(pos*headDim)+d]
		}
	}

	return output
}

// ============================================================
// ConvBlock实现
// ============================================================

// ConvBlock 短卷积块
type ConvBlock struct {
	hiddenDim int
	kernel    int
	weights   [][]float32 // 卷积权重
}

// NewConvBlock 创建卷积块
func NewConvBlock(hiddenDim, kernel int) *ConvBlock {
	return &ConvBlock{
		hiddenDim: hiddenDim,
		kernel:    kernel,
		weights:   make([][]float32, hiddenDim),
	}
}

// Forward 卷积前向传播
func (cb *ConvBlock) Forward(x []float32, cache [][]float32, step int) []float32 {
	output := make([]float32, len(x))
	copy(output, x)

	// 因果卷积：只使用当前位置及之前的位置
	for d := 0; d < cb.hiddenDim; d++ {
		var sum float32
		for k := 0; k < cb.kernel; k++ {
			pos := step - k
			if pos < 0 {
				continue
			}
			val := cache[pos][d]
			// 门控权重
			gate := float32(1.0 / float64(k+1))
			sum += val * gate
		}
		output[d] = sum * 0.1 // 简化缩放
	}

	return output
}

// ============================================================
// SwiGLU FFN实现
// ============================================================

// SwiGLUFFN SwiGLU激活的前馈网络
type SwiGLUFFN struct {
	hiddenDim    int
	intermediate int
}

// NewSwiGLUFFN 创建SwiGLU FFN
func NewSwiGLUFFN(hiddenDim, intermediate int) *SwiGLUFFN {
	return &SwiGLUFFN{
		hiddenDim:    hiddenDim,
		intermediate: intermediate,
	}
}

func swish(x float32) float32 {
	return x / (1 + float32(math.Exp(float64(-x))))
}

// Forward SwiGLU前向传播
func (ffn *SwiGLUFFN) Forward(x []float32) []float32 {
	// SwiGLU: output = (swish(xW1) * xW3) * W2
	// 这里简化为一个非线性变换
	output := make([]float32, len(x))
	for i, v := range x {
		gate := swish(v)
		output[i] = gate * v * 0.5
	}
	return output
}

// ============================================================
// 完整推理引擎
// ============================================================

// InferenceEngine 端侧推理引擎
type InferenceEngine struct {
	config  ModelConfig
	kvCache *KVCache
	rope    *RoPECache
	gqa     *GQALayer
	conv    *ConvBlock
	ffn     *SwiGLUFFN
}

// NewInferenceEngine 创建推理引擎
func NewInferenceEngine(config ModelConfig) *InferenceEngine {
	headDim := config.HiddenDim / config.NumQHeads
	intermediate := int(float64(config.HiddenDim) * config.FFNScale)

	return &InferenceEngine{
		config:  config,
		kvCache: NewKVCache(config),
		rope:    NewRoPECache(config, config.MaxSeqLen),
		gqa:     NewGQALayer(GQAConfig{
			HiddenDim:  config.HiddenDim,
			NumQHeads:  config.NumQHeads,
			NumKVHeads: config.NumKVHeads,
			HeadDim:    headDim,
		}),
		conv: NewConvBlock(config.HiddenDim, config.ConvKernel),
		ffn:  NewSwiGLUFFN(config.HiddenDim, intermediate),
	}
}

// Generate 生成文本
func (e *InferenceEngine) Generate(prompt []int, maxTokens int) []int {
	output := make([]int, 0, maxTokens)
	hidden := make([]float32, e.config.HiddenDim)

	// 缓存隐藏状态（模拟）
	convCache := make([][]float32, e.config.MaxSeqLen)
	for i := range convCache {
		convCache[i] = make([]float32, e.config.HiddenDim)
	}

	gqaLayerIdx := 0

	for step := 0; step < maxTokens; step++ {
		// 模拟embedding lookup
		for i := range hidden {
			hidden[i] = float32(step % 100) * 0.01
		}

		// 逐层推理
		for layer := 0; layer < e.config.NumLayers; layer++ {
			if e.config.LayerTypes[layer] == "conv" {
				hidden = e.conv.Forward(hidden, convCache, step)
				// 更新卷积缓存
				copy(convCache[step], hidden)
			} else {
				hidden = e.gqa.Forward(hidden, step, e.kvCache, gqaLayerIdx, e.rope)
				gqaLayerIdx++
			}

			// FFN
			hidden = e.ffn.Forward(hidden)
		}

		// 模拟logits采样
		nextToken := step % 1000
		output = append(output, nextToken)
	}

	return output
}

// GetMemoryStats 获取内存统计
func (e *InferenceEngine) GetMemoryStats() map[string]float64 {
	kvMem := e.kvCache.MemoryUsage()
	return map[string]float64{
		"kv_cache_mb": kvMem,
		"total_mb":    kvMem + 2500, // 模型权重约2.5GB
	}
}

// ============================================================
// 量化工具
// ============================================================

// Quantizer 量化器
type Quantizer struct {
	config ModelConfig
}

// Q4Weight Q4量化权重
type Q4Weight struct {
	scale   float32
	zero    float32
	indices []int8 // 4-bit存储，用int8低4位
}

// QuantizeWeights 量化权重到4-bit
func (q *Quantizer) QuantizeWeights(weights []float32, groupSize int) []Q4Weight {
	numGroups := (len(weights) + groupSize - 1) / groupSize
	result := make([]Q4Weight, numGroups)

	for g := 0; g < numGroups; g++ {
		start := g * groupSize
		end := start + groupSize
		if end > len(weights) {
			end = len(weights)
		}
		group := weights[start:end]

		// 计算min/max
		minVal := float32(math.Inf(1))
		maxVal := float32(math.Inf(-1))
		for _, w := range group {
			if w < minVal {
				minVal = w
			}
			if w > maxVal {
				maxVal = w
			}
		}

		scale := (maxVal - minVal) / 15.0
		zero := minVal

		indices := make([]int8, len(group))
		for i, w := range group {
			quantized := int8((w - zero) / scale)
			if quantized < 0 {
				quantized = 0
			}
			if quantized > 15 {
				quantized = 15
			}
			indices[i] = quantized
		}

		result[g] = Q4Weight{
			scale:   scale,
			zero:    zero,
			indices: indices,
		}
	}

	return result
}

// Dequantize 反量化
func (q *Quantizer) Dequantize(weights []Q4Weight) []float32 {
	totalLen := 0
	for _, w := range weights {
		totalLen += len(w.indices)
	}

	result := make([]float32, totalLen)
	idx := 0
	for _, w := range weights {
		for _, qi := range w.indices {
			result[idx] = w.zero + float32(qi)*w.scale
			idx++
		}
	}

	return result
}

// ============================================================
// 基准测试
// ============================================================

func runBenchmark() {
	fmt.Println("LFM2.5-2.6B 端侧推理引擎基准测试")
	fmt.Println("=" * 60)

	config := DefaultConfig()
	engine := NewInferenceEngine(config)

	// 统计架构信息
	gqaCount := 0
	convCount := 0
	for _, t := range config.LayerTypes {
		if t == "gqa" {
			gqaCount++
		} else {
			convCount++
		}
	}

	fmt.Printf("模型配置:\n")
	fmt.Printf("  总层数: %d\n", config.NumLayers)
	fmt.Printf("  ConvBlock: %d层\n", convCount)
	fmt.Printf("  GQA: %d层\n", gqaCount)
	fmt.Printf("  隐藏维度: %d\n", config.HiddenDim)
	fmt.Printf("  注意力头: %dQ / %dKV\n", config.NumQHeads, config.NumKVHeads)
	fmt.Printf("  词汇表: %d\n", config.VocabSize)
	fmt.Printf("  最大上下文: %d\n", config.MaxSeqLen)
	fmt.Println()

	// 推理速度测试
	fmt.Println("推理速度测试:")
	prompt := make([]int, 128) // 128 token prompt
	for i := range prompt {
		prompt[i] = i % 1000
	}

	numRuns := 5
	var totalTime float64
	var totalTokens int

	for run := 0; run < numRuns; run++ {
		start := time.Now()
		output := engine.Generate(prompt, 256)
		elapsed := time.Since(start).Seconds()
		totalTime += elapsed
		totalTokens += len(output)
	}

	avgTokensPerSec := float64(totalTokens) / totalTime
	fmt.Printf("  平均生成速度: %.1f tok/s\n", avgTokensPerSec)
	fmt.Println()

	// 内存分析
	fmt.Println("内存分析:")
	memStats := engine.GetMemoryStats()
	for k, v := range memStats {
		fmt.Printf("  %s: %.1f MB\n", k, v)
	}
	fmt.Println()

	// 量化测试
	fmt.Println("量化测试:")
	quantizer := &Quantizer{config: config}

	// 模拟权重
	numWeights := config.HiddenDim * config.HiddenDim * 4
	weights := make([]float32, numWeights)
	for i := range weights {
		weights[i] = float32(i) * 0.0001
	}

	origSize := len(weights) * 4 // float32

	groupSize := 32
	qweights := quantizer.QuantizeWeights(weights, groupSize)
	quantSize := len(qweights) * (4 + 4 + len(qweights[0].indices))

	fmt.Printf("  原始权重大小: %.2f MB\n", float64(origSize)/(1024*1024))
	fmt.Printf("  Q4量化后大小: %.2f MB\n", float64(quantSize)/(1024*1024))
	fmt.Printf("  压缩比: %.1f:1\n", float64(origSize)/float64(quantSize))
	fmt.Println()

	// 总内存估算
	modelWeightSize := 2.5 * 1024.0 // MB (BF16)
	quantModelSize := modelWeightSize / 4.0 // Q4
	fmt.Println("端侧部署内存估算:")
	fmt.Printf("  BF16模型权重: %.1f MB\n", modelWeightSize)
	fmt.Printf("  Q4量化模型: %.1f MB\n", quantModelSize)
	fmt.Printf("  KV缓存(128K上下文): ~%.0f MB\n", memStats["kv_cache_mb"])
	fmt.Printf("  总内存(Q4+KV): ~%.1f MB\n", quantModelSize+memStats["kv_cache_mb"])
}

func main() {
	runBenchmark()
}

"""
量化感知部署工具链
支持：动态量化、KV缓存优化、推理调度
"""
import numpy as np
from typing import List, Tuple, Optional, Dict
from dataclasses import dataclass, field
from enum import Enum
import struct
import math
import time

class QuantMethod(Enum):
    """量化方法"""
    Q4_0 = "q4_0"      # 4-bit, 对称
    Q4_K_M = "q4_k_m"  # 4-bit, 分组混合精度
    Q8_0 = "q8_0"      # 8-bit, 对称
    BF16 = "bf16"       # 无量化

@dataclass
class QuantConfig:
    """量化配置"""
    method: QuantMethod = QuantMethod.Q4_K_M
    group_size: int = 32
    use_sym: bool = False  # 是否对称量化
    
    # KV缓存量化
    kv_cache_quant: bool = True
    kv_cache_bits: int = 8
    
    # 内存预算
    memory_budget_mb: float = 2500.0

@dataclass
class LayerConfig:
    """单层配置"""
    layer_type: str  # 'conv' 或 'gqa'
    hidden_dim: int
    intermediate_dim: int
    num_heads: int = 0
    num_kv_heads: int = 0

class BlockQuantizer:
    """分块量化器"""
    
    @staticmethod
    def quantize_q4_k_m(weights: np.ndarray, group_size: int = 32) -> Tuple[np.ndarray, np.ndarray]:
        """
        Q4_K_M量化：每个分组使用独立的scale和min
        
        Args:
            weights: 输入权重 [N,]
            group_size: 分组大小
            
        Returns:
            quantized: 量化后的4-bit数据
            metadata: (scales, mins) 每组scale和min
        """
        assert weights.ndim == 1
        n = len(weights)
        num_groups = (n + group_size - 1) // group_size
        
        quantized = np.zeros(num_groups * group_size // 2, dtype=np.uint8)
        scales = np.zeros(num_groups, dtype=np.float16)
        mins = np.zeros(num_groups, dtype=np.float16)
        
        for g in range(num_groups):
            start = g * group_size
            end = min(start + group_size, n)
            group = weights[start:end]
            
            # 找min/max
            w_min = group.min()
            w_max = group.max()
            scale = (w_max - w_min) / 15.0 if w_max != w_min else 1.0
            
            scales[g] = scale
            mins[g] = w_min
            
            # 量化到4-bit
            for i in range(0, len(group), 2):
                v0 = int((group[i] - w_min) / scale)
                v1 = int((group[i + 1] - w_min) / scale) if i + 1 < len(group) else 0
                v0 = max(0, min(15, v0))
                v1 = max(0, min(15, v1))
                quantized[g * group_size // 2 + i // 2] = (v0 & 0x0F) | ((v1 & 0x0F) << 4)
        
        return quantized, (scales, mins)
    
    @staticmethod
    def dequantize_q4_k_m(quantized: np.ndarray, metadata: Tuple[np.ndarray, np.ndarray], 
                          group_size: int = 32, total_len: int = None) -> np.ndarray:
        """反量化Q4_K_M"""
        scales, mins = metadata
        num_groups = len(scales)
        if total_len is None:
            total_len = num_groups * group_size
        
        result = np.zeros(total_len, dtype=np.float32)
        
        for g in range(num_groups):
            start = g * group_size
            end = min(start + group_size, total_len)
            scale = scales[g]
            w_min = mins[g]
            
            for i in range(start, end):
                byte_idx = g * group_size // 2 + i // 2
                if i % 2 == 0:
                    qi = quantized[byte_idx] & 0x0F
                else:
                    qi = (quantized[byte_idx] >> 4) & 0x0F
                result[i] = w_min + float(qi) * scale
        
        return result

class KVCacheQuantizer:
    """KV缓存量化器"""
    
    def __init__(self, bits: int = 8):
        self.bits = bits
        self.max_val = 2 ** (bits - 1) - 1
    
    def quantize(self, tensor: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
        """量化KV缓存张量"""
        abs_max = np.max(np.abs(tensor))
        if abs_max < 1e-10:
            abs_max = 1.0
        
        scale = self.max_val / abs_max
        quantized = np.clip(np.round(tensor * scale), -self.max_val - 1, self.max_val).astype(np.int8)
        
        return quantized, np.array([scale, abs_max])
    
    def dequantize(self, quantized: np.ndarray, metadata: np.ndarray) -> np.ndarray:
        """反量化KV缓存"""
        scale = metadata[0]
        return quantized.astype(np.float32) / scale


class MemoryOptimizer:
    """内存优化器"""
    
    def __init__(self, config: QuantConfig):
        self.config = config
        self.memory_tracker = {}
    
    def estimate_model_memory(self, num_params: int, bits: int) -> float:
        """估算模型权重内存"""
        bytes_per_param = bits / 8
        return num_params * bytes_per_param / (1024 ** 3)  # GB
    
    def estimate_kv_cache_memory(
        self,
        num_gqa_layers: int,
        num_kv_heads: int,
        head_dim: int,
        seq_len: int,
        batch_size: int,
        quant_bits: int = 8
    ) -> float:
        """
        估算KV缓存内存
        
        KV缓存 = 2 (K+V) * num_layers * num_kv_heads * head_dim * seq_len * batch_size * bytes_per_element
        """
        bytes_per_elem = quant_bits / 8
        total_bytes = (2 * num_gqa_layers * num_kv_heads * head_dim * 
                      seq_len * batch_size * bytes_per_elem)
        return total_bytes / (1024 ** 3)  # GB
    
    def optimize_sequence_length(
        self,
        available_memory_gb: float,
        model_memory_gb: float,
        num_gqa_layers: int,
        num_kv_heads: int,
        head_dim: int,
        batch_size: int = 1
    ) -> int:
        """根据可用内存计算最大序列长度"""
        kv_memory_budget = available_memory_gb - model_memory_gb
        if kv_memory_budget <= 0:
            return 0
        
        bytes_per_elem = self.config.kv_cache_bits / 8
        max_seq_len = int(
            kv_memory_budget * (1024 ** 3) / 
            (2 * num_gqa_layers * num_kv_heads * head_dim * batch_size * bytes_per_elem)
        )
        
        return max_seq_len
    
    def profile_memory(self, 
                       model_params: int,
                       quant_bits: int,
                       seq_len: int,
                       num_gqa: int,
                       num_kv_heads: int,
                       head_dim: int) -> Dict[str, float]:
        """完整内存分析"""
        model_mem = self.estimate_model_memory(model_params, quant_bits)
        kv_mem = self.estimate_kv_cache_memory(num_gqa, num_kv_heads, head_dim, seq_len, 1, 
                                               self.config.kv_cache_bits)
        
        total = model_mem + kv_mem
        
        self.memory_tracker = {
            "model_weights_gb": model_mem,
            "kv_cache_gb": kv_mem,
            "total_gb": total,
            "total_mb": total * 1024,
            "within_budget": total * 1024 <= self.config.memory_budget_mb
        }
        
        return self.memory_tracker


class LFM2_6B_DeploymentTool:
    """
    LFM2.5-2.6B部署工具链
    整合量化、内存优化、推理调度
    """
    
    def __init__(self, memory_budget_mb: float = 2500):
        self.config = QuantConfig(memory_budget_mb=memory_budget_mb)
        self.quantizer = BlockQuantizer()
        self.kv_quantizer = KVCacheQuantizer(bits=self.config.kv_cache_bits)
        self.mem_optimizer = MemoryOptimizer(self.config)
        
        # LFM2.5-2.6B架构参数
        self.arch = {
            "num_params": 2.69e9,
            "num_layers": 30,
            "num_gqa": 8,
            "num_conv": 22,
            "hidden_dim": 2048,
            "num_q_heads": 32,
            "num_kv_heads": 8,
            "head_dim": 64,
            "vocab_size": 128000,
            "max_seq_len": 131072,
        }
    
    def analyze_deployment(self):
        """分析部署可行性"""
        print("=" * 70)
        print("LFM2.5-2.6B 端侧部署分析")
        print("=" * 70)
        
        arch = self.arch
        
        print(f"\n📐 架构参数:")
        print(f"  总参数量: {arch['num_params']/1e9:.2f}B")
        print(f"  总层数: {arch['num_layers']}")
        print(f"    ├─ ConvBlock: {arch['num_conv']}层")
        print(f"    └─ GQA: {arch['num_gqa']}层")
        print(f"  隐藏维度: {arch['hidden_dim']}")
        print(f"  注意力头: {arch['num_q_heads']}Q / {arch['num_kv_heads']}KV")
        print(f"  词汇表: {arch['vocab_size']:,}")
        print(f"  最大上下文: {arch['max_seq_len']:,}")
        
        print(f"\n📊 量化方案对比:")
        quant_schemes = [
            ("BF16", 16),
            ("Q8_0", 8),
            ("Q4_K_M", 4),
            ("Q4_0", 4),
        ]
        
        memory_budget = self.config.memory_budget_mb
        
        for name, bits in quant_schemes:
            model_mem = self.mem_optimizer.estimate_model_memory(
                arch["num_params"], bits
            ) * 1024  # 转MB
            
            kv_mem = self.mem_optimizer.estimate_kv_cache_memory(
                arch["num_gqa"], arch["num_kv_heads"], arch["head_dim"],
                arch["max_seq_len"], 1, 8
            ) * 1024  # 转MB
            
            total = model_mem + kv_mem
            feasible = "✅" if total <= memory_budget else "❌"
            
            print(f"  {name:>8}: 模型={model_mem:>7.1f}MB, "
                  f"KV缓存={kv_mem:>7.1f}MB, "
                  f"总计={total:>7.1f}MB, "
                  f"预算={memory_budget:.0f}MB {feasible}")
        
        print(f"\n🔧 推荐配置: Q4_K_M + 8-bit KV缓存量化")
        print(f"  模型权重: {self.mem_optimizer.estimate_model_memory(arch['num_params'], 4) * 1024:.0f}MB")
        print(f"  KV缓存(128K): {self.mem_optimizer.estimate_kv_cache_memory(arch['num_gqa'], arch['num_kv_heads'], arch['head_dim'], arch['max_seq_len'], 1, 8) * 1024:.0f}MB")
        
        # 最大序列长度分析
        seq_lens = [4096, 8192, 16384, 32768, 65536, 131072]
        print(f"\n📏 不同上下文长度下的内存占用:")
        model_mem_q4 = self.mem_optimizer.estimate_model_memory(arch["num_params"], 4) * 1024
        
        for seq_len in seq_lens:
            kv_mem = self.mem_optimizer.estimate_kv_cache_memory(
                arch["num_gqa"], arch["num_kv_heads"], arch["head_dim"],
                seq_len, 1, 8
            ) * 1024
            total = model_mem_q4 + kv_mem
            flag = "✅" if total <= memory_budget else "❌"
            print(f"  {seq_len:>6,} tokens: 模型={model_mem_q4:>6.0f}MB + "
                  f"KV={kv_mem:>6.1f}MB = {total:>6.1f}MB {flag}")
    
    def simulate_quantization(self):
        """模拟量化过程"""
        print("\n" + "=" * 70)
        print("量化模拟")
        print("=" * 70)
        
        np.random.seed(42)
        sample_weights = np.random.randn(4096).astype(np.float32) * 0.1
        
        print(f"\n输入权重: 均值={sample_weights.mean():.6f}, "
              f"标准差={sample_weights.std():.6f}")
        
        # Q4_K_M量化
        start = time.time()
        q_data, meta = BlockQuantizer.quantize_q4_k_m(sample_weights, group_size=32)
        quant_time = time.time() - start
        
        start = time.time()
        deq = BlockQuantizer.dequantize_q4_k_m(q_data, meta, group_size=32, 
                                                total_len=len(sample_weights))
        dequant_time = time.time() - start
        
        # 计算误差
        mse = np.mean((sample_weights - deq) ** 2)
        max_err = np.max(np.abs(sample_weights - deq))
        
        # 压缩比
        orig_size = len(sample_weights) * 4  # float32
        meta_size = len(meta[0]) * 2 + len(meta[1]) * 2  # float16
        quant_size = len(q_data) + meta_size
        
        print(f"\nQ4_K_M 量化结果:")
        print(f"  分组大小: 32")
        print(f"  MSE: {mse:.8f}")
        print(f"  最大误差: {max_err:.6f}")
        print(f"  原始大小: {orig_size} bytes")
        print(f"  量化大小: {quant_size} bytes")
        print(f"  压缩比: {orig_size/quant_size:.1f}:1")
        print(f"  量化耗时: {quant_time*1000:.2f}ms")
        print(f"  反量化耗时: {dequant_time*1000:.2f}ms")
        
        return mse
    
    def kv_cache_optimization_demo(self):
        """KV缓存优化演示"""
        print("\n" + "=" * 70)
        print("KV缓存优化演示")
        print("=" * 70)
        
        arch = self.arch
        
        # 计算标准MHA vs GQA的KV缓存差异
        mha_kv_heads = arch["num_q_heads"]  # 32
        gqa_kv_heads = arch["num_kv_heads"]  # 8
        
        seq_len = 131072
        head_dim = arch["head_dim"]
        
        mha_kv_size = 2 * arch["num_gqa"] * mha_kv_heads * head_dim * seq_len * 2  # FP16
        gqa_kv_size = 2 * arch["num_gqa"] * gqa_kv_heads * head_dim * seq_len * 2  # FP16
        
        # 进一步量化到8-bit
        gqa_quant_size = gqa_kv_size / 2  # 8-bit
        
        print(f"\nKV缓存大小对比 (seq_len={seq_len:,}):")
        print(f"  MHA (32 KV heads): {mha_kv_size/(1024**3):.2f} GB")
        print(f"  GQA (8 KV heads):  {gqa_kv_size/(1024**3):.2f} GB")
        print(f"  GQA + 8-bit量化:  {gqa_quant_size/(1024**3):.2f} GB")
        print(f"  GQA优化缩减: {mha_kv_size/gqa_kv_size:.1f}x")
        print(f"  量化进一步缩减: {gqa_kv_size/gqa_quant_size:.1f}x")
        print(f"  总缩减: {mha_kv_size/gqa_quant_size:.1f}x")
        
        # 演示KV缓存量化
        print(f"\nKV缓存量化保存/加载演示:")
        sample_kv = np.random.randn(8, 64, 1024).astype(np.float32) * 0.5
        
        q_kv, meta = self.kv_quantizer.quantize(sample_kv)
        dq_kv = self.kv_quantizer.dequantize(q_kv, meta)
        
        kv_mse = np.mean((sample_kv - dq_kv) ** 2)
        kv_compress = sample_kv.nbytes / q_kv.nbytes
        
        print(f"  KV量化MSE: {kv_mse:.8f}")
        print(f"  KV压缩比: {kv_compress:.1f}:1 (FP32 → {self.config.kv_cache_bits}-bit)")


def main():
    """主函数：完整的部署分析"""
    tool = LFM2_6B_DeploymentTool(memory_budget_mb=2500)
    
    # 1. 部署可行性分析
    tool.analyze_deployment()
    
    # 2. 量化模拟
    mse = tool.simulate_quantization()
    
    # 3. KV缓存优化
    tool.kv_cache_optimization_demo()
    
    # 4. 总结
    print("\n" + "=" * 70)
    print("部署建议")
    print("=" * 70)
    print("""
LFM2.5-2.6B 端侧部署建议:

1. 智能手机端 (内存预算: 2.5GB):
   - 量化方案: Q4_K_M (4-bit)
   - KV缓存: 8-bit量化
   - 最大上下文: 32K-64K tokens
   - 预期速度: ~30 tok/s

2. 笔记本端 (M5 Max / Ryzen AI):
   - 量化方案: Q4_K_M 或 Q8_0
   - KV缓存: 8-bit或FP16
   - 最大上下文: 128K (完整支持)
   - 预期速度: 113-220 tok/s

3. 树莓派:
   - 量化方案: Q4_0 (最简量化)
   - KV缓存: 4-bit量化
   - 最大上下文: 8K-16K tokens
   - 需要额外的内存优化

4. 服务器端 (H100):
   - 量化方案: BF16或FP8
   - 最大吞吐: ~15,000 tok/s
   - 单卡日处理: ~1.3B tokens
    """)
    
    return tool

if __name__ == "__main__":
    tool = main()

// benchmark.go - 基准测试对比分析
package main

import (
	"fmt"
	"math"
	"sort"
)

// BenchmarkResult 基准测试结果
type BenchmarkResult struct {
	Name      string
	Category  string // "instruction", "tool", "agent", "stem", "coding"
	IsHigher  bool   // true=越高越好, false=越低越好
}

// ModelResult 模型在某个基准上的分数
type ModelResult struct {
	ModelName string
	Score     float64
}

// BenchmarkSuite 完整的基准测试套件
type BenchmarkSuite struct {
	Results map[string]map[string]float64 // benchmark -> model -> score
}

// LoadLFM25Benchmarks 加载LFM2.5-2.6B对比数据
func LoadLFM25Benchmarks() *BenchmarkSuite {
	bs := &BenchmarkSuite{
		Results: make(map[string]map[string]float64),
	}

	benchmarks := []struct {
		name     string
		category string
		isHigher bool
		data     map[string]float64
	}{
		{"AA Omniscience", "stem", false, map[string]float64{
			"LFM2.5-2.6B":   -29.50,
			"Gemma4-E2B":    -74.47,
			"Gemma4-E4B":    -49.03,
			"Qwen3.5-4B":    -54.30,
			"Qwen3.5-9B":    -50.43,
		}},
		{"AIME25", "stem", true, map[string]float64{
			"LFM2.5-2.6B": 51.87,
			"Gemma4-E2B":  26.33,
			"Gemma4-E4B":  34.27,
			"Qwen3.5-4B":  49.33,
			"Qwen3.5-9B":  56.07,
		}},
		{"LiveCodeBenchv6", "coding", true, map[string]float64{
			"LFM2.5-2.6B": 59.41,
			"Gemma4-E2B":  54.92,
			"Gemma4-E4B":  63.77,
			"Qwen3.5-4B":  60.85,
			"Qwen3.5-9B":  69.86,
		}},
		{"IFBench", "instruction", true, map[string]float64{
			"LFM2.5-2.6B": 59.17,
			"Gemma4-E2B":  34.08,
			"Gemma4-E4B":  39.24,
			"Qwen3.5-4B":  48.40,
			"Qwen3.5-9B":  56.47,
		}},
		{"Multi-IF", "instruction", true, map[string]float64{
			"LFM2.5-2.6B": 80.07,
			"Gemma4-E2B":  69.44,
			"Gemma4-E4B":  77.35,
			"Qwen3.5-4B":  55.67,
			"Qwen3.5-9B":  62.55,
		}},
		{"IFStruct", "instruction", true, map[string]float64{
			"LFM2.5-2.6B": 85.49,
			"Gemma4-E2B":  64.85,
			"Gemma4-E4B":  76.65,
			"Qwen3.5-4B":  36.25,
			"Qwen3.5-9B":  78.50,
		}},
		{"BFCLv4", "tool", true, map[string]float64{
			"LFM2.5-2.6B": 56.88,
			"Gemma4-E2B":  36.98,
			"Gemma4-E4B":  46.39,
			"Qwen3.5-4B":  50.56,
			"Qwen3.5-9B":  60.13,
		}},
		{"ToolSandbox", "tool", true, map[string]float64{
			"LFM2.5-2.6B": 77.83,
			"Gemma4-E2B":  52.40,
			"Gemma4-E4B":  65.00,
			"Qwen3.5-4B":  75.55,
			"Qwen3.5-9B":  76.44,
		}},
		{"τ³-Bench Banking", "agent", true, map[string]float64{
			"LFM2.5-2.6B": 5.67,
			"Gemma4-E2B":  3.35,
			"Gemma4-E4B":  4.12,
			"Qwen3.5-4B":  5.45,
			"Qwen3.5-9B":  5.15,
		}},
		{"Claw-Eval avg (EN)", "agent", true, map[string]float64{
			"LFM2.5-2.6B": 62.85,
			"Gemma4-E2B":  53.14,
			"Gemma4-E4B":  58.02,
			"Qwen3.5-4B":  62.28,
			"Qwen3.5-9B":  66.53,
		}},
		{"PinchBench", "agent", true, map[string]float64{
			"LFM2.5-2.6B": 68.22,
			"Gemma4-E2B":  44.24,
			"Gemma4-E4B":  55.09,
			"Qwen3.5-4B":  71.26,
			"Qwen3.5-9B":  71.45,
		}},
		{"BrowseComp+", "agent", true, map[string]float64{
			"LFM2.5-2.6B": 26.89,
			"Gemma4-E2B":  8.31,
			"Gemma4-E4B":  15.90,
			"Qwen3.5-4B":  24.46,
			"Qwen3.5-9B":  27.23,
		}},
	}

	for _, b := range benchmarks {
		bs.Results[b.name] = b.data
	}

	return bs
}

// ModelInfo 模型信息
type ModelInfo struct {
	Name       string
	Parameters float64 // 十亿
	Layers     int
	HiddenDim  int
	ContextLen int
}

// GetModelInfo 获取模型信息
func GetModelInfo() []ModelInfo {
	return []ModelInfo{
		{Name: "LFM2.5-2.6B", Parameters: 2.69, Layers: 30, HiddenDim: 2048, ContextLen: 131072},
		{Name: "Gemma4-E2B", Parameters: 5.1, Layers: 26, HiddenDim: 2048, ContextLen: 128000},
		{Name: "Gemma4-E4B", Parameters: 8.0, Layers: 32, HiddenDim: 2560, ContextLen: 128000},
		{Name: "Qwen3.5-4B", Parameters: 4.7, Layers: 32, HiddenDim: 2560, ContextLen: 262144},
		{Name: "Qwen3.5-9B", Parameters: 9.7, Layers: 36, HiddenDim: 3584, ContextLen: 262144},
	}
}

// ParameterEfficiencyScore 参数效率得分
type ParameterEfficiencyScore struct {
	ModelName      string
	ParametersB    float64
	BenchmarkScore float64
	ScorePerBParam float64
	Normalized     float64 // 相对于最佳模型的百分比
}

// ComputeParameterEfficiency 计算参数效率
func ComputeParameterEfficiency(bs *BenchmarkSuite, modelName string, 
                                 modelParams float64) []ParameterEfficiencyScore {
	var scores []ParameterEfficiencyScore

	for benchName, modelScores := range bs.Results {
		if score, ok := modelScores[modelName]; ok {
			scorePerB := score / modelParams
			scores = append(scores, ParameterEfficiencyScore{
				ModelName:      modelName,
				ParametersB:    modelParams,
				BenchmarkScore: score,
				ScorePerBParam: scorePerB,
			})
		}
		_ = benchName
	}

	return scores
}

// EfficiencyRanking 效率排名
type EfficiencyRanking struct {
	ModelName      string
	ParametersB    float64
	AvgEfficiency  float64 // 平均每B参数得分
	AvgNormalized  float64
	WinCount       int     // 在多少个基准上排名第一
}

// RankEfficiency 对所有模型进行效率排名
func RankEfficiency(bs *BenchmarkSuite) []EfficiencyRanking {
	models := []string{"LFM2.5-2.6B", "Gemma4-E2B", "Gemma4-E4B", "Qwen3.5-4B", "Qwen3.5-9B"}
	modelParams := map[string]float64{
		"LFM2.5-2.6B": 2.69,
		"Gemma4-E2B":  5.1,
		"Gemma4-E4B":  8.0,
		"Qwen3.5-4B":  4.7,
		"Qwen3.5-9B":  9.7,
	}

	rankings := make(map[string]*EfficiencyRanking)
	for _, m := range models {
		rankings[m] = &EfficiencyRanking{
			ModelName:   m,
			ParametersB: modelParams[m],
		}
	}

	// 对每个基准，计算每个模型的效率
	for benchName, modelScores := range bs.Results {
		// 计算基准中所有模型的效率
		type eff struct {
			model  string
			effVal float64
		}
		var efficiencies []eff

		for model, score := range modelScores {
			params := modelParams[model]
			e := eff{
				model:  model,
				effVal: score / params,
			}
			efficiencies = append(efficiencies, e)
		}

		// 找到最高效率
		sort.Slice(efficiencies, func(i, j int) bool {
			return efficiencies[i].effVal > efficiencies[j].effVal
		})

		bestEff := efficiencies[0].effVal

		for _, e := range efficiencies {
			r := rankings[e.model]
			r.AvgEfficiency += e.effVal
			r.AvgNormalized += (e.effVal / bestEff) * 100
			if e.effVal == bestEff {
				r.WinCount++
			}
		}
		_ = benchName
	}

	// 计算平均值
	numBenchmarks := len(bs.Results)
	var result []EfficiencyRanking
	for _, r := range rankings {
		r.AvgEfficiency /= float64(numBenchmarks)
		r.AvgNormalized /= float64(numBenchmarks)
		result = append(result, *r)
	}

	// 按平均标准化效率排序
	sort.Slice(result, func(i, j int) bool {
		return result[i].AvgNormalized > result[j].AvgNormalized
	})

	return result
}

// ComputeWinsAnalysis 胜场分析
type WinsAnalysis struct {
	Category string
	Model    string
	Wins     int
	Total    int
}

func ComputeCategoryWins(bs *BenchmarkSuite) []WinsAnalysis {
	benchmarkCategories := map[string]string{
		"AA Omniscience":  "STEM",
		"AIME25":         "STEM",
		"LiveCodeBenchv6": "Coding",
		"IFBench":        "Instruction",
		"Multi-IF":       "Instruction",
		"IFStruct":       "Instruction",
		"BFCLv4":         "Tool Use",
		"ToolSandbox":    "Tool Use",
		"τ³-Bench Banking": "Agent",
		"Claw-Eval avg (EN)": "Agent",
		"PinchBench":     "Agent",
		"BrowseComp+":    "Agent",
	}

	models := []string{"LFM2.5-2.6B", "Gemma4-E2B", "Gemma4-E4B", "Qwen3.5-4B", "Qwen3.5-9B"}
	categories := []string{"STEM", "Coding", "Instruction", "Tool Use", "Agent"}

	type catKey struct {
		category string
		model    string
	}
	wins := make(map[catKey]int)
	totals := make(map[string]int)

	for benchName, modelScores := range bs.Results {
		cat := benchmarkCategories[benchName]
		totals[cat]++

		// 找到最高分
		bestModel := ""
		bestScore := math.Inf(-1)
		for model, score := range modelScores {
			if score > bestScore {
				bestScore = score
				bestModel = model
			}
		}

		key := catKey{cat, bestModel}
		wins[key]++
	}

	var result []WinsAnalysis
	for _, cat := range categories {
		for _, model := range models {
			key := catKey{cat, model}
			w := wins[key]
			if w > 0 {
				result = append(result, WinsAnalysis{
					Category: cat,
					Model:    model,
					Wins:     w,
					Total:    totals[cat],
				})
			}
		}
	}

	return result
}

func main() {
	bs := LoadLFM25Benchmarks()
	models := GetModelInfo()

	fmt.Println("LFM2.5-2.6B 基准测试深度分析")
	fmt.Println("=" * 70)

	// 1. 模型信息总览
	fmt.Println("\n📋 模型参数对比:")
	fmt.Printf("%-20s %8s %8s %10s %10s\n", "Model", "Params(B)", "Layers", "Hidden", "Context")
	fmt.Println(string(repeatChar('-', 60)))
	for _, m := range models {
		fmt.Printf("%-20s %8.1f %8d %10d %10d\n",
			m.Name, m.Parameters, m.Layers, m.HiddenDim, m.ContextLen)
	}

	// 2. 参数效率排名
	fmt.Println("\n⚡ 参数效率排名 (每B参数得分):")
	rankings := RankEfficiency(bs)
	fmt.Printf("%-20s %10s %15s %15s %10s\n", "Model", "Params(B)", "Avg Efficiency", "Norm %", "Wins")
	fmt.Println(string(repeatChar('-', 75)))
	for _, r := range rankings {
		fmt.Printf("%-20s %10.1f %15.2f %14.1f%% %10d\n",
			r.ModelName, r.ParametersB, r.AvgEfficiency, r.AvgNormalized, r.WinCount)
	}

	// 3. 分类胜场分析
	fmt.Println("\n🏆 分类胜场分析:")
	wins := ComputeCategoryWins(bs)
	currentCat := ""
	for _, w := range wins {
		if w.Category != currentCat {
			fmt.Printf("\n  %s:\n", w.Category)
			currentCat = w.Category
		}
		fmt.Printf("    %-20s: %d/%d 胜场\n", w.Model, w.Wins, w.Total)
	}

	// 4. 效率比分析
	fmt.Println("\n📊 效率比 (LFM2.5-2.6B vs 对比模型):")
	fmt.Printf("%-20s %15s %15s\n", "Benchmark", "vs Gemma4-E4B", "vs Qwen3.5-9B")
	fmt.Println(string(repeatChar('-', 55)))

	lfmParams := 2.69
	gemma4Params := 8.0
	qwen9Params := 9.7

	benchmarks := []string{"IFBench", "Multi-IF", "IFStruct", "BFCLv4", "ToolSandbox", "AIME25"}
	for _, b := range benchmarks {
		lfmScore := bs.Results[b]["LFM2.5-2.6B"]
		gemmaScore := bs.Results[b]["Gemma4-E4B"]
		qwenScore := bs.Results[b]["Qwen3.5-9B"]

		lfmEff := lfmScore / lfmParams
		gemmaEff := gemmaScore / gemmaParams
		qwenEff := qwenScore / qwenParams

		vsGemma := lfmEff / gemmaEff
		vsQwen := lfmEff / qwenEff

		fmt.Printf("%-20s %14.1fx %14.1fx\n", b, vsGemma, vsQwen)
	}

	// 5. 总结
	fmt.Println("\n" + string(repeatChar('=', 70)))
	fmt.Println("核心结论")
	fmt.Println(string(repeatChar('=', 70)))
	fmt.Println(`
1. LFM2.5-2.6B在指令遵循基准上全面领先（IFBench 59.17, IFStruct 85.49）
2. 工具调用能力与Qwen3.5-9B持平，ToolSandbox甚至超越（77.83 vs 76.44）
3. Agent任务超越Gemma 4全系列，与Qwen3.5系列打平
4. 参数效率为Gemma 4-E4B的3-4倍，为Qwen3.5-9B的5-6倍
5. 代码生成是唯一短板，建议使用更大模型
6. 2.5GB内存占用+220 tok/s速度，端侧部署无竞品`)
}

func repeatChar(c byte, n int) string {
	b := make([]byte, n)
	for i := range b {
		b[i] = c
	}
	return string(b)
}