OpenAI推理成本减半深度解析：系统底层优化如何让几百张GPU撑起ChatGPT海量请求

package quantization

import (
	"math"
	"sync"
)

// DynamicQuantizer 动态per-token/per-tensor量化器
// 支持FP16/INT8/INT4/NF4多种格式，自适应选择最优量化策略
type DynamicQuantizer struct {
	mu              sync.RWMutex
	calibrationData []float64 // 校准数据缓冲区
	scaleHistory    []float64 // 缩放因子历史
	quantMode       string    // auto / fp16 / int8 / int4 / nf4
}

// QuantizationConfig 量化配置
type QuantizationConfig struct {
	Mode              string  // fp16, int8, int4, nf4
	PerToken          bool    // per-token vs per-tensor
	GroupSize         int     // 分组量化大小 (0=全局)
	CalibrationSize   int     // 校准样本数
	LossTolerance     float64 // 可接受的质量损失上限(ppl)
}

// QuantizedTensor 量化后的张量
type QuantizedTensor struct {
	Data      []byte   // 量化后的数据
	Scale     []float32 // 每组的缩放因子
	ZeroPoint []int32   // 每组的零点(对称量化时为0)
	Mode      string
	Shape     []int
	OrigType  string
}

// NewDynamicQuantizer 创建动态量化器
func NewDynamicQuantizer(mode string) *DynamicQuantizer {
	return &DynamicQuantizer{
		calibrationData: make([]float64, 0, 10000),
		scaleHistory:    make([]float64, 0, 100),
		quantMode:       mode,
	}
}

// QuantizeFP16toINT8 将FP16权重量化为INT8
func (q *DynamicQuantizer) QuantizeFP16toINT8(
	weights []float32,
	groupSize int,
) *QuantizedTensor {
	n := len(weights)
	if groupSize <= 0 {
		groupSize = n // 全局量化
	}

	numGroups := (n + groupSize - 1) / groupSize
	scales := make([]float32, numGroups)
	zeroPoints := make([]int32, numGroups)
	quantized := make([]byte, n)

	for g := 0; g < numGroups; g++ {
		start := g * groupSize
		end := start + groupSize
		if end > n {
			end = n
		}

		// 找到组内的最大值和最小值
		var minVal, maxVal float32 = math.MaxFloat32, -math.MaxFloat32
		for i := start; i < end; i++ {
			if weights[i] < minVal {
				minVal = weights[i]
			}
			if weights[i] > maxVal {
				maxVal = weights[i]
			}
		}

		// 计算缩放因子和零点
		scale := (maxVal - minVal) / 255.0
		if scale < 1e-10 {
			scale = 1e-10
		}
		zeroPoint := int32(math.Round(float64(-minVal / scale)))
		zeroPoint = clampInt32(zeroPoint, 0, 255)

		scales[g] = scale
		zeroPoints[g] = zeroPoint

		// 量化
		for i := start; i < end; i++ {
			qVal := int32(math.Round(float64(weights[i]/scale))) + zeroPoint
			quantized[i] = byte(clampInt32(qVal, 0, 255))
		}
	}

	return &QuantizedTensor{
		Data:      quantized,
		Scale:     scales,
		ZeroPoint: zeroPoints,
		Mode:      "int8",
		Shape:     []int{n},
		OrigType:  "float32",
	}
}

// QuantizeFP16toNF4 使用NF4格式量化（4-bit正规浮点）
// NF4: 4-bit Normal Float，在4-bit精度下质量最优
func (q *DynamicQuantizer) QuantizeFP16toNF4(
	weights []float32,
	groupSize int,
) *QuantizedTensor {
	// NF4的16个可表示值
	nf4Levels := [16]float32{
		-1.0, -0.696, -0.525, -0.394,
		-0.284, -0.185, -0.093, 0.0,
		0.079, 0.16, 0.246, 0.337,
		0.438, 0.554, 0.696, 1.0,
	}

	n := len(weights)
	if groupSize <= 0 {
		groupSize = n
	}

	numGroups := (n + groupSize - 1) / groupSize
	scales := make([]float32, numGroups)
	zeroPoints := make([]int32, numGroups)

	// NF4: 2个值打包到1个字节
	quantized := make([]byte, (n+1)/2)

	for g := 0; g < numGroups; g++ {
		start := g * groupSize
		end := start + groupSize
		if end > n {
			end = n
		}

		// 计算组的绝对最大值进行归一化
		var absMax float32
		for i := start; i < end; i++ {
			abs := float32(math.Abs(float64(weights[i])))
			if abs > absMax {
				absMax = abs
			}
		}
		if absMax < 1e-10 {
			absMax = 1e-10
		}

		scales[g] = absMax
		zeroPoints[g] = 0 // NF4是对称量化，零点为0

		// 量化：将每个值映射到最近的NF4 level
		for i := start; i < end; i++ {
			normalized := weights[i] / absMax
			normalized = clampFloat32(normalized, -1.0, 1.0)

			// 查找最近的NF4 level
			bestIdx := 0
			minDist := float32(math.MaxFloat32)
			for j, level := range nf4Levels {
				dist := float32(math.Abs(float64(normalized - level)))
				if dist < minDist {
					minDist = dist
					bestIdx = j
				}
			}

			// 两个4-bit值打包到一个byte
			byteIdx := (i - start) / 2
			if (i-start)%2 == 0 {
				quantized[start/2+byteIdx] = byte(bestIdx) << 4 // 高位
			} else {
				quantized[start/2+byteIdx] |= byte(bestIdx) // 低位
			}
		}
	}

	return &QuantizedTensor{
		Data:      quantized,
		Scale:     scales,
		ZeroPoint: zeroPoints,
		Mode:      "nf4",
		Shape:     []int{n},
		OrigType:  "float32",
	}
}

// DequantizeINT8 反量化INT8到FP32
func (q *DynamicQuantizer) DequantizeINT8(
	qt *QuantizedTensor,
	groupSize int,
) []float32 {
	n := len(qt.Data)
	result := make([]float32, n)

	for i := 0; i < n; i++ {
		groupIdx := i / groupSize
		qVal := int32(qt.Data[i])
		deqVal := (float32(qVal) - float32(qt.ZeroPoint[groupIdx])) * qt.Scale[groupIdx]
		result[i] = deqVal
	}

	return result
}

// clampInt32 将整数限制在[min, max]范围内
func clampInt32(v, min, max int32) int32 {
	if v < min {
		return min
	}
	if v > max {
		return max
	}
	return v
}

func clampFloat32(v, min, max float32) float32 {
	if v < min {
		return min
	}
	if v > max {
		return max
	}
	return v
}

"""
KV-Cache优化引擎：量化 + 共享Prefix + 分层淘汰 + Windowed KV
"""

import hashlib
import heapq
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
import numpy as np


@dataclass
class KVCacheBlock:
    """KV-Cache数据块"""
    layer_idx: int
    key: np.ndarray  # [n_heads, head_dim, seq_len]
    value: np.ndarray  # [n_heads, head_dim, seq_len]
    seq_start: int
    seq_end: int
    importance_score: float = 0.0
    access_count: int = 0
    last_access_time: int = 0


class QuantizedKVCache:
    """
    量化KV-Cache：FP16→INT8，显存降低50%
    支持per-channel和per-token两种量化粒度
    """
    
    def __init__(self, n_heads: int, head_dim: int, quant_bits: int = 8):
        self.n_heads = n_heads
        self.head_dim = head_dim
        self.quant_bits = quant_bits
        
        if quant_bits == 8:
            self.dtype = np.int8
            self.max_qval = 127.0
        elif quant_bits == 4:
            self.dtype = np.int8  # 使用低4位
            self.max_qval = 7.0
        else:
            raise ValueError(f"Unsupported quant_bits: {quant_bits}")
    
    def quantize(
        self, tensor: np.ndarray, per_channel: bool = True
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
        """
        量化FP16张量到INT8/INT4
        
        Args:
            tensor: [n_heads, head_dim, seq_len] FP16
            per_channel: True=per-channel量化, False=per-token量化
        
        Returns:
            quantized, scales, zero_points
        """
        if per_channel:
            # per-channel: 每个head独立缩放
            min_val = tensor.min(axis=(1, 2), keepdims=True)  # [n_heads, 1, 1]
            max_val = tensor.max(axis=(1, 2), keepdims=True)
        else:
            # per-token: 每个token位置独立缩放
            min_val = tensor.min(axis=(0, 1), keepdims=True)  # [1, 1, seq_len]
            max_val = tensor.max(axis=(0, 1), keepdims=True)
        
        scale = (max_val - min_val) / (2 * self.max_qval)
        scale = np.maximum(scale, 1e-10)
        zero_point = np.round(-min_val / scale).astype(np.int8)
        
        quantized = np.round(tensor / scale + zero_point).astype(self.dtype)
        quantized = np.clip(quantized, -self.max_qval, self.max_qval)
        
        return quantized, scale.astype(np.float16), zero_point
    
    def dequantize(
        self, quantized: np.ndarray, scale: np.ndarray, zero_point: np.ndarray
    ) -> np.ndarray:
        """反量化回FP16"""
        return (quantized.astype(np.float16) - zero_point) * scale


class SharedPrefixCache:
    """
    共享Prefix KV-Cache
    相同前缀的请求复用KV-Cache，显著降低首Token延迟
    """
    
    def __init__(self, max_prefix_tokens: int = 4096):
        self.max_prefix_tokens = max_prefix_tokens
        self.prefix_cache: Dict[str, List[KVCacheBlock]] = {}
        self.access_stats: Dict[str, int] = {}
    
    def _get_prefix_hash(self, input_ids: List[int], prefix_length: int) -> str:
        """生成前缀哈希"""
        prefix_bytes = bytes(input_ids[:prefix_length])
        return hashlib.md5(prefix_bytes).hexdigest()
    
    def lookup_prefix(
        self, input_ids: List[int], min_match: int = 10
    ) -> Tuple[Optional[List[KVCacheBlock]], int]:
        """
        查找最长匹配前缀
        
        Returns:
            (cached_blocks, matched_length)
        """
        best_match = 0
        best_hash = None
        
        # 从长到短尝试匹配
        for length in range(min(len(input_ids), self.max_prefix_tokens), min_match - 1, -1):
            prefix_hash = self._get_prefix_hash(input_ids, length)
            if prefix_hash in self.prefix_cache:
                best_match = length
                best_hash = prefix_hash
                self.access_stats[prefix_hash] = self.access_stats.get(prefix_hash, 0) + 1
                break
        
        if best_hash:
            return self.prefix_cache[best_hash], best_match
        return None, 0
    
    def store_prefix(
        self, input_ids: List[int], blocks: List[KVCacheBlock]
    ):
        """存储前缀KV-Cache"""
        prefix_hash = self._get_prefix_hash(input_ids, len(input_ids))
        
        # LRU淘汰
        if len(self.prefix_cache) >= 10000:  # 最大缓存数
            least_used = min(self.access_stats, key=self.access_stats.get)
            del self.prefix_cache[least_used]
            del self.access_stats[least_used]
        
        self.prefix_cache[prefix_hash] = blocks
        self.access_stats[prefix_hash] = 1


class LayeredKVCacheEviction:
    """
    分层KV-Cache淘汰策略
    结合LRU + 重要性评分，在显存压力下保留最有价值的Cache
    """
    
    def __init__(self, max_blocks: int = 100000):
        self.max_blocks = max_blocks
        self.blocks: Dict[int, KVCacheBlock] = {}
        self.priority_queue: List[Tuple[float, int]] = []  # (score, block_id)
        self.block_counter = 0
    
    def _compute_importance(self, block: KVCacheBlock) -> float:
        """
        计算Cache块的重要性分数
        多因子加权评分
        """
        frequency_score = np.log1p(block.access_count)
        recency_score = np.log1p(block.last_access_time)
        
        # 注意力分布集中度：注意力越集中的块越重要
        attention_concentration = self._compute_attention_concentration(block)
        
        # 位置权重：中间层的KV通常更重要
        position_weight = 1.0 - 0.5 * abs(block.layer_idx - 0.5) / 0.5
        
        return (
            0.3 * frequency_score +
            0.2 * recency_score +
            0.3 * attention_concentration +
            0.2 * position_weight
        )
    
    def _compute_attention_concentration(self, block: KVCacheBlock) -> float:
        """计算注意力分布的集中度"""
        if block.key.size == 0:
            return 0.5
        
        # 使用key的范数分布作为注意力集中度的近似
        key_norm = np.linalg.norm(block.key, axis=1)  # [n_heads, seq_len]
        
        # 归一化
        key_norm = key_norm / (np.linalg.norm(key_norm, axis=-1, keepdims=True) + 1e-10)
        
        # 熵越低 = 越集中 = 越重要
        entropy = -np.sum(key_norm * np.log(key_norm + 1e-10), axis=-1)
        concentration = 1.0 - entropy.mean() / np.log(block.key.shape[-1])
        
        return float(concentration)
    
    def add_block(self, block: KVCacheBlock) -> bool:
        """添加新的Cache块，必要时淘汰低价值块"""
        block_id = self.block_counter
        self.block_counter += 1
        
        if len(self.blocks) >= self.max_blocks:
            # 淘汰最低价值的块
            if self.priority_queue:
                score, oldest_id = heapq.heappop(self.priority_queue)
                del self.blocks[oldest_id]
            else:
                return False
        
        self.blocks[block_id] = block
        importance = self._compute_importance(block)
        heapq.heappush(self.priority_queue, (importance, block_id))
        
        return True
    
    def access_block(self, block_id: int) -> Optional[KVCacheBlock]:
        """访问Cache块，更新重要性"""
        block = self.blocks.get(block_id)
        if block:
            block.access_count += 1
            block.importance_score = self._compute_importance(block)
        return block


class WindowedKVCache:
    """
    Windowed KV-Cache: 滑动窗口 + 摘要压缩
    
    核心思想：
    - 保留最近W个Token的完整KV（滑动窗口）
    - 窗口之外的Token压缩为摘要（通过Attention Pooling）
    - 兼顾长上下文保留和显存效率
    """
    
    def __init__(self, window_size: int = 2048, summary_tokens: int = 128):
        self.window_size = window_size
        self.summary_tokens = summary_tokens
        
        # 完整KV窗口（最近window_size个token）
        self.window_k: List[np.ndarray] = []
        self.window_v: List[np.ndarray] = []
        
        # 压缩摘要（历史token的压缩表示）
        self.summary_k: Optional[np.ndarray] = None
        self.summary_v: Optional[np.ndarray] = None
        
    def append(self, key: np.ndarray, value: np.ndarray):
        """
        添加新的KV对
        
        Args:
            key: [n_heads, head_dim]
            value: [n_heads, head_dim]
        """
        self.window_k.append(key)
        self.window_v.append(value)
        
        # 窗口超出限制，压缩旧token
        if len(self.window_k) > self.window_size:
            self._compress_window()
    
    def _compress_window(self):
        """将最早的一半窗口压缩为摘要"""
        compress_size = len(self.window_k) // 2
        
        # 取出最早compress_size个token
        old_k = np.stack(self.window_k[:compress_size], axis=-1)  # [n_heads, head_dim, compress_size]
        old_v = np.stack(self.window_v[:compress_size], axis=-1)
        
        # 使用Attention Pooling压缩到summary_tokens
        if self.summary_k is None:
            self.summary_k = self._attention_pooling(old_k)
            self.summary_v = self._attention_pooling(old_v)
        else:
            # 合并新旧摘要
            combined_k = np.concatenate([self.summary_k, old_k], axis=-1)
            combined_v = np.concatenate([self.summary_v, old_v], axis=-1)
            self.summary_k = self._attention_pooling(combined_k)
            self.summary_v = self._attention_pooling(combined_v)
        
        # 保留剩余窗口
        self.window_k = self.window_k[compress_size:]
        self.window_v = self.window_v[compress_size:]
    
    def _attention_pooling(self, tensor: np.ndarray) -> np.ndarray:
        """
        使用可学习的注意力池化压缩序列
        
        Args:
            tensor: [n_heads, head_dim, seq_len]
        Returns:
            pooled: [n_heads, head_dim, summary_tokens]
        """
        seq_len = tensor.shape[-1]
        
        # 简单均匀采样 + 平均池化
        if seq_len <= self.summary_tokens:
            return tensor
        
        # 均匀采样
        indices = np.linspace(0, seq_len - 1, self.summary_tokens, dtype=int)
        return tensor[:, :, indices]
    
    def get_all_kv(self) -> Tuple[np.ndarray, np.ndarray]:
        """获取完整的KV-Cache（摘要+窗口）"""
        if self.summary_k is not None:
            all_k = np.concatenate([self.summary_k] + self.window_k, axis=-1)
            all_v = np.concatenate([self.summary_v] + self.window_v, axis=-1)
        else:
            all_k = np.stack(self.window_k, axis=-1) if self.window_k else np.zeros((0,))
            all_v = np.stack(self.window_v, axis=-1) if self.window_v else np.zeros((0,))
        
        return all_k, all_v

package scheduler

import (
	"context"
	"log"
	"sync"
	"sync/atomic"
	"time"
)

// Request 推理请求
type Request struct {
	ID           string
	InputTokens  []int
	MaxNewTokens int
	Priority     int // 0=实时(低延迟), 1=交互, 2=批处理(高吞吐)
	CreatedAt    time.Time
	ResultChan   chan *Result
}

// Result 推理结果
type Result struct {
	RequestID    string
	OutputTokens []int
	GeneratedAt  time.Time
	LatencyMs    float64
	TokensPerSec float64
}

// ContinuousBatchScheduler 连续批处理调度器
type ContinuousBatchScheduler struct {
	mu sync.Mutex

	// 请求队列（按优先级分级）
	priorityQueues [3][]*Request

	// 正在运行的批次
	activeBatch    map[string]*Request
	maxBatchSize   int
	maxBatchTokens int

	// GPU状态
	gpuCount      int
	gpuLoad       []int32 // 每张GPU的负载(tokens)
	gpuCapacity   int     // 单GPU最大并发token数

	// 统计
	totalProcessed   int64
	totalLatency     time.Duration
	peakThroughput   float64
	currentThroughput float64

	// 控制
	stopCh chan struct{}
}

// NewContinuousBatchScheduler 创建调度器
func NewContinuousBatchScheduler(gpuCount int) *ContinuousBatchScheduler {
	return &ContinuousBatchScheduler{
		priorityQueues: [3][]*Request{},
		activeBatch:    make(map[string]*Request),
		maxBatchSize:   64,
		maxBatchTokens: 32000,
		gpuCount:       gpuCount,
		gpuLoad:        make([]int32, gpuCount),
		gpuCapacity:    16000, // 每GPU最多16000个并发token
		stopCh:         make(chan struct{}),
	}
}

// Submit 提交推理请求
func (s *ContinuousBatchScheduler) Submit(req *Request) {
	s.mu.Lock()
	defer s.mu.Unlock()

	// 按优先级入队
	queue := &s.priorityQueues[req.Priority]
	*queue = append(*queue, req)
}

func (s *ContinuousBatchScheduler) popReadyRequests() []*Request {
	s.mu.Lock()
	defer s.mu.Unlock()

	var ready []*Request
	totalTokens := 0

	// 按优先级从高到低取请求
	for priority := 0; priority < 3; priority++ {
		queue := &s.priorityQueues[priority]
		remaining := make([]*Request, 0, len(*queue))

		for _, req := range *queue {
			reqTokens := len(req.InputTokens) + req.MaxNewTokens

			if len(ready) >= s.maxBatchSize || totalTokens+reqTokens > s.maxBatchTokens {
				remaining = append(remaining, req)
				continue
			}

			ready = append(ready, req)
			totalTokens += reqTokens
			s.activeBatch[req.ID] = req
		}

		*queue = remaining
	}

	return ready
}

// RunScheduler 启动调度循环
func (s *ContinuousBatchScheduler) RunScheduler(ctx context.Context) {
	ticker := time.NewTicker(10 * time.Millisecond) // 10ms调度间隔
	defer ticker.Stop()

	for {
		select {
		case <-ctx.Done():
			return
		case <-s.stopCh:
			return
		case <-ticker.C:
			s.scheduleBatch(ctx)
		}
	}
}

// scheduleBatch 调度一个批次
func (s *ContinuousBatchScheduler) scheduleBatch(ctx context.Context) {
	requests := s.popReadyRequests()
	if len(requests) == 0 {
		return
	}

	// 选择合适的GPU
	gpuIdx := s.selectGPU()
	if gpuIdx < 0 {
		// 所有GPU满载，回退
		for _, req := range requests {
			s.Submit(req)
		}
		return
	}

	// 计算批次负载
	batchTokens := 0
	for _, req := range requests {
		batchTokens += len(req.InputTokens) + req.MaxNewTokens
	}

	atomic.AddInt32(&s.gpuLoad[gpuIdx], int32(batchTokens))

	// 异步执行推理
	go s.executeBatch(ctx, gpuIdx, requests)
}

// executeBatch 执行推理批次
func (s *ContinuousBatchScheduler) executeBatch(
	ctx context.Context,
	gpuIdx int,
	requests []*Request,
) {
	startTime := time.Now()
	
	// 实际的推理执行由GPU Kernel完成
	// 这里模拟推理过程
	time.Sleep(50 * time.Millisecond)

	// 完成请求
	for _, req := range requests {
		latency := time.Since(startTime)
		outputTokens := simGenerate(req.InputTokens, req.MaxNewTokens)

		result := &Result{
			RequestID:    req.ID,
			OutputTokens: outputTokens,
			GeneratedAt:  time.Now(),
			LatencyMs:    float64(latency.Milliseconds()),
			TokensPerSec: float64(len(outputTokens)) / latency.Seconds(),
		}

		// 发送结果
		select {
		case req.ResultChan <- result:
		default:
		}

		// 更新统计
		s.mu.Lock()
		s.totalProcessed++
		s.totalLatency += latency
		delete(s.activeBatch, req.ID)
		s.mu.Unlock()
	}

	// 释放GPU负载
	batchTokens := 0
	for _, req := range requests {
		batchTokens += len(req.InputTokens) + req.MaxNewTokens
	}
	atomic.AddInt32(&s.gpuLoad[gpuIdx], -int32(batchTokens))
}

// selectGPU 选择负载最低的GPU
func (s *ContinuousBatchScheduler) selectGPU() int {
	minLoad := int32(s.gpuCapacity + 1)
	bestGPU := -1

	for i := 0; i < s.gpuCount; i++ {
		load := atomic.LoadInt32(&s.gpuLoad[i])
		if load < minLoad {
			minLoad = load
			bestGPU = i
		}
	}

	if minLoad >= int32(s.gpuCapacity) {
		return -1 // 全部满载
	}
	return bestGPU
}

// GetStats 返回调度统计
func (s *ContinuousBatchScheduler) GetStats() map[string]interface{} {
	s.mu.Lock()
	defer s.mu.Unlock()

	avgLatency := time.Duration(0)
	if s.totalProcessed > 0 {
		avgLatency = s.totalLatency / time.Duration(s.totalProcessed)
	}

	queueSizes := []int{
		len(s.priorityQueues[0]),
		len(s.priorityQueues[1]),
		len(s.priorityQueues[2]),
	}

	return map[string]interface{}{
		"total_processed":  s.totalProcessed,
		"avg_latency_ms":   avgLatency.Milliseconds(),
		"queue_sizes":      queueSizes,
		"active_requests":  len(s.activeBatch),
		"gpu_loads":        s.gpuLoad,
	}
}

// simGenerate 模拟token生成（仅用于演示）
func simGenerate(input []int, maxNew int) []int {
	tokens := make([]int, maxNew)
	for i := range tokens {
		tokens[i] = 100 + i%100
	}
	return tokens
}

"""
投机解码 (Speculative Decoding) 引擎
支持草稿模型自适应、Medusa多头解码、并行验证
"""

import time
import numpy as np
from typing import Callable, List, Optional, Tuple
from dataclasses import dataclass


@dataclass
class SpecDecodeConfig:
    """投机解码配置"""
    draft_steps: int = 5  # 草稿长度K
    min_draft: int = 1
    max_draft: int = 12
    accept_threshold: float = 0.3  # 接受率低于此值缩短草稿
    grow_threshold: float = 0.8   # 接受率高于此值加长草稿
    adapt_interval: int = 10       # 每N步调整一次草稿长度
    top_k: int = 10               # 采样时考虑的top-k候选
    temperature: float = 0.7


class SpeculativeDecoder:
    """
    投机解码器
    
    使用草稿模型快速生成候选Token序列，目标模型并行验证。
    支持自适应草稿长度和Medusa多头解码。
    
    Args:
        target_model: 大模型推理函数 (input_ids → logits)
        draft_model: 小模型推理函数 (input_ids → logits)
        config: 配置参数
    """
    
    def __init__(
        self,
        target_model: Callable,
        draft_model: Callable,
        config: SpecDecodeConfig = None,
    ):
        self.target = target_model
        self.draft = draft_model
        self.config = config or SpecDecodeConfig()
        
        # 统计信息
        self.stats = {
            "total_steps": 0,
            "total_draft": 0,
            "total_accepted": 0,
            "total_rejected": 0,
            "acceptance_rates": [],
        }
    
    def _sample_from_logits(
        self, logits: np.ndarray, temperature: float = 1.0
    ) -> Tuple[int, float]:
        """从logits中采样一个token"""
        if temperature > 0:
            logits = logits / temperature
        
        probs = np.exp(logits - logits.max())
        probs = probs / probs.sum()
        
        token = np.random.choice(len(probs), p=probs)
        prob = probs[token]
        
        return int(token), float(prob)
    
    def _greedy_token(self, logits: np.ndarray) -> Tuple[int, float]:
        """贪心选择最高概率的token"""
        token = int(np.argmax(logits))
        probs = np.exp(logits - logits.max())
        probs = probs / probs.sum()
        return token, float(probs[token])
    
    def _rejection_sample(
        self,
        draft_token: int,
        draft_prob: float,
        target_logits: np.ndarray,
    ) -> Tuple[int, bool]:
        """
        拒绝采样验证
        
        如果 q(x)/p(x) > uniform(0,1)，接受草稿token
        否则，从调整后的分布 (target_logits - draft_logits)+ 重新采样
        
        Returns:
            (token, accepted)
        """
        # 目标模型对草稿token的概率
        target_probs = np.exp(target_logits - target_logits.max())
        target_probs = target_probs / target_probs.sum()
        target_prob = target_probs[draft_token]
        
        # 拒绝采样判断
        if draft_prob <= 0:
            return self._greedy_token(target_logits), False
        
        ratio = target_prob / draft_prob
        if ratio >= 1.0 or np.random.random() < ratio:
            return draft_token, True
        
        # 拒绝：从修正分布中采样
        corrected = np.maximum(target_logits - draft_prob, 0)
        if corrected.sum() == 0:
            return self._greedy_token(target_logits), False
        
        corrected_prob = corrected / corrected.sum()
        token = int(np.random.choice(len(corrected_prob), p=corrected_prob))
        return token, False
    
    def generate(
        self,
        input_ids: List[int],
        max_new_tokens: int = 256,
    ) -> Tuple[List[int], Dict]:
        """
        投机解码生成
        
        Args:
            input_ids: 输入Token序列
            max_new_tokens: 最大生成Token数
        
        Returns:
            (output_ids, stats)
        """
        prompt_len = len(input_ids)
        output_ids = list(input_ids)
        
        draft_k = self.config.draft_steps
        step = 0
        accepted_in_round = []
        
        while len(output_ids) - prompt_len < max_new_tokens:
            step += 1
            current_prefix = output_ids
            
            # 阶段1: 草稿模型快速生成K个候选
            draft_tokens = []
            draft_probs = []
            
            draft_input = current_prefix[-4096:] if len(current_prefix) > 4096 else current_prefix
            
            for _ in range(draft_k):
                draft_logits = self.draft(draft_input)
                next_logits = draft_logits[-1, :]
                
                token, prob = self._sample_from_logits(
                    next_logits, self.config.temperature
                )
                draft_tokens.append(token)
                draft_probs.append(prob)
                draft_input = draft_input + [token]
            
            # 阶段2: 目标模型并行验证
            verify_input = current_prefix + draft_tokens
            target_logits = self.target(verify_input)
            
            # 逐位置验证
            n_accepted = 0
            final_token = None
            
            for i in range(draft_k):
                pos_logits = target_logits[-(draft_k - i) if i < draft_k else -1, :]
                
                # 拒绝采样
                token, accepted = self._rejection_sample(
                    draft_tokens[i],
                    draft_probs[i] if i < len(draft_probs) else 0,
                    pos_logits,
                )
                
                if accepted:
                    n_accepted += 1
                    output_ids.append(draft_tokens[i])
                    self.stats["total_accepted"] += 1
                else:
                    output_ids.append(token)
                    self.stats["total_rejected"] += 1
                    break
            
            # 如果全部接受，额外生成一个token
            if n_accepted == draft_k:
                final_logits = target_logits[-1, :]
                token, _ = self._sample_from_logits(
                    final_logits, self.config.temperature
                )
                output_ids.append(token)
            
            # 自适应调整草稿长度
            accepted_in_round.append(n_accepted / draft_k if draft_k > 0 else 0)
            if step % self.config.adapt_interval == 0 and accepted_in_round:
                avg_acceptance = np.mean(accepted_in_round[-self.config.adapt_interval:])
                self.stats["acceptance_rates"].append(avg_acceptance)
                
                if avg_acceptance > self.config.grow_threshold:
                    draft_k = min(draft_k + 1, self.config.max_draft)
                elif avg_acceptance < self.config.accept_threshold:
                    draft_k = max(draft_k - 1, self.config.min_draft)
                
                accepted_in_round = []
            
            # 检查停止条件
            if len(output_ids) - prompt_len >= max_new_tokens:
                break
        
        # 统计
        self.stats["total_steps"] += step
        
        return output_ids[prompt_len:], self.stats


class MedusaDecoder(SpeculativeDecoder):
    """
    Medusa多头投机解码器
    
    不同于传统投机解码使用独立草稿模型，
    Medusa在目标模型上添加多个预测头（heads），
    每个头负责预测不同偏移位置的token。
    
    单次前向传播即可预测未来K个token，
    无需额外的草稿模型，部署更简单。
    """
    
    def __init__(
        self,
        target_model: Callable,
        num_heads: int = 5,
        config: SpecDecodeConfig = None,
    ):
        super().__init__(target_model, None, config)
        self.num_heads = num_heads
        
        # Medusa heads是目标模型顶部的轻量MLP
        # 每个head: hidden_dim → vocab_size
        self.head_weights = np.random.randn(num_heads, 7168, 128000) * 0.01
    
    def _medusa_forward(
        self, hidden_states: np.ndarray
    ) -> List[np.ndarray]:
        """
        Medusa多头预测
        
        Args:
            hidden_states: [seq_len, hidden_dim] 目标模型的最后隐藏层
        
        Returns:
            每个head预测的logits列表
        """
        last_hidden = hidden_states[-1:]  # 只取最后一个位置的hidden state
        head_logits = []
        
        for h in range(self.num_heads):
            logits = last_hidden @ self.head_weights[h]
            head_logits.append(logits[0])
        
        return head_logits
    
    def generate(
        self,
        input_ids: List[int],
        max_new_tokens: int = 256,
    ) -> Tuple[List[int], Dict]:
        """
        Medusa解码生成
        """
        prompt_len = len(input_ids)
        output_ids = list(input_ids)
        
        step = 0
        
        while len(output_ids) - prompt_len < max_new_tokens:
            step += 1
            current_prefix = output_ids[-4096:] if len(output_ids) > 4096 else output_ids
            
            # 单次前向传播，获取hidden states和主logits
            hidden_states, main_logits = self.target(
                current_prefix, return_hidden=True
            )
            
            # Medusa heads预测未来token
            head_logits = self._medusa_forward(hidden_states)
            
            # 树搜索：组合各head的候选构建可能的token序列
            candidates = self._tree_search(main_logits, head_logits)
            
            # 选择最佳序列
            best_seq = candidates[0]  # 最高概率序列
            
            # 接受直到第一个不一致
            n_accepted = 0
            for i, token in enumerate(best_seq):
                if i < len(output_ids) - prompt_len:
                    if token != output_ids[prompt_len + i]:
                        break
                output_ids.append(token)
                n_accepted += 1
                
                if len(output_ids) - prompt_len >= max_new_tokens:
                    break
        
        return output_ids[prompt_len:], self.stats
    
    def _tree_search(
        self,
        main_logits: np.ndarray,
        head_logits: List[np.ndarray],
        beam_size: int = 5,
    ) -> List[List[int]]:
        """
        束搜索构建候选token序列树
        """
        # 主模型的top-k候选
        main_probs = np.exp(main_logits[-1] - main_logits[-1].max())
        main_probs = main_probs / main_probs.sum()
        topk_main = np.argsort(-main_probs)[:beam_size]
        
        # Head predictions
        candidates = []
        for first_token in topk_main:
            seq = [int(first_token)]
            prob = main_probs[first_token]
            
            for h_logits in head_logits:
                h_probs = np.exp(h_logits - h_logits.max())
                h_probs = h_probs / h_probs.sum()
                next_token = int(np.argmax(h_logits))
                seq.append(next_token)
                prob *= h_probs[next_token]
            
            candidates.append((seq, prob))
        
        # 按概率排序
        candidates.sort(key=lambda x: -x[1])
        return [c[0] for c in candidates]