DeepSeek V4 正式版深度技术解析：MoE稀疏注意力 + DSpark推测解码 + 峰谷定价的技术经济学

package attention

import (
	"math"
	"sync"
)

// DenseSparseAttention DSA注意力机制
type DenseSparseAttention struct {
	HeadDim     int
	NumHeads    int
	CompressRatio float64 // token压缩比例
	SparseRatio float64   // 稀疏注意力比例
	mu          sync.Mutex
}

func NewDSA(headDim, numHeads int, compressRatio, sparseRatio float64) *DenseSparseAttention {
	return &DenseSparseAttention{
		HeadDim:       headDim,
		NumHeads:      numHeads,
		CompressRatio: compressRatio,
		SparseRatio:   sparseRatio,
	}
}

// TokenCompressionScore 计算每个token的重要性分数
func (dsa *DenseSparseAttention) TokenCompressionScore(hiddenStates [][]float32) []float64 {
	seqLen := len(hiddenStates)
	scores := make([]float64, seqLen)

	for i := 0; i < seqLen; i++ {
		var norm float64
		for _, v := range hiddenStates[i] {
			norm += float64(v) * float64(v)
		}
		scores[i] = math.Sqrt(norm / float64(len(hiddenStates[i])))
	}
	return scores
}

// CompressTokens 根据重要性分数压缩token序列
func (dsa *DenseSparseAttention) CompressTokens(
	hiddenStates [][]float32,
	scores []float64,
) ([][]float32, []int) {
	seqLen := len(hiddenStates)
	keepCount := int(float64(seqLen) * dsa.CompressRatio)

	// 创建索引并按分数排序
	type scoredIdx struct {
		idx   int
		score float64
	}
	pairs := make([]scoredIdx, seqLen)
	for i, s := range scores {
		pairs[i] = scoredIdx{idx: i, score: s}
	}

	// 快速选择：找到第keepCount大的分数阈值
	threshold := quickSelect(scores, keepCount)

	compressed := make([][]float32, 0, keepCount)
	indices := make([]int, 0, keepCount)
	for i, s := range scores {
		if s >= threshold && len(compressed) < keepCount {
			compressed = append(compressed, hiddenStates[i])
			indices = append(indices, i)
		}
	}
	return compressed, indices
}

// quickSelect 快速选择第k大的元素
func quickSelect(arr []float64, k int) float64 {
	if len(arr) <= k {
		return 0
	}
	// 简化实现：使用排序
	sorted := make([]float64, len(arr))
	copy(sorted, arr)
	// 降序排序
	for i := 0; i < len(sorted); i++ {
		for j := i + 1; j < len(sorted); j++ {
			if sorted[j] > sorted[i] {
				sorted[i], sorted[j] = sorted[j], sorted[i]
			}
		}
	}
	return sorted[k-1]
}

// DenseAttention 密集注意力：在压缩序列上计算全局注意力
func (dsa *DenseSparseAttention) DenseAttention(
	Q, K, V [][]float32,
) [][]float32 {
	seqLen := len(Q)
	output := make([][]float32, seqLen)
	for i := 0; i < seqLen; i++ {
		output[i] = make([]float32, dsa.HeadDim)
	}

	var wg sync.WaitGroup
	headSize := dsa.HeadDim

	for h := 0; h < dsa.NumHeads; h++ {
		wg.Add(1)
		go func(head int) {
			defer wg.Done()
			offset := head * headSize
			for i := 0; i < seqLen; i++ {
				// 计算注意力分数
				var maxScore float32 = -1e9
				scores := make([]float32, seqLen)
				for j := 0; j < seqLen; j++ {
					var dot float32
					for d := 0; d < headSize; d++ {
						dot += Q[i][offset+d] * K[j][offset+d]
					}
					scores[j] = dot / float32(math.Sqrt(float64(headSize)))
					if scores[j] > maxScore {
						maxScore = scores[j]
					}
				}
				// softmax
				var sum float32
				for j := 0; j < seqLen; j++ {
					scores[j] = float32(math.Exp(float64(scores[j] - maxScore)))
					sum += scores[j]
				}
				for j := 0; j < seqLen; j++ {
					scores[j] /= sum
				}
				// 加权求和
				for d := 0; d < headSize; d++ {
					var val float32
					for j := 0; j < seqLen; j++ {
						val += scores[j] * V[j][offset+d]
					}
					output[i][offset+d] = val
				}
			}
		}(h)
	}
	wg.Wait()
	return output
}

// SparseAttention 稀疏注意力：在原始序列上按稀疏索引计算局部注意力
func (dsa *DenseSparseAttention) SparseAttention(
	Q, K, V [][]float32,
	sparseIndices []int,
	windowSize int,
) [][]float32 {
	seqLen := len(Q)
	output := make([][]float32, seqLen)
	for i := 0; i < seqLen; i++ {
		output[i] = make([]float32, dsa.HeadDim)
	}
	// 每个token只关注附近的稀疏索引
	sparseSet := make(map[int]bool)
	for _, idx := range sparseIndices {
		sparseSet[idx] = true
	}

	for i := 0; i < seqLen; i++ {
		start := i - windowSize
		if start < 0 {
			start = 0
		}
		end := i + windowSize
		if end > seqLen {
			end = seqLen
		}
		// 在窗口内找稀疏索引
		var sum float32
		scores := make(map[int]float32)
		for j := start; j < end; j++ {
			if sparseSet[j] {
				var dot float32
				for d := 0; d < dsa.HeadDim; d++ {
					dot += Q[i][d] * K[j][d]
				}
				scores[j] = dot / float32(math.Sqrt(float64(dsa.HeadDim)))
				if scores[j] > sum {
					sum = scores[j]
				}
			}
		}
		// softmax + 加权
		var total float32
		for j, s := range scores {
			scores[j] = float32(math.Exp(float64(s - sum)))
			total += scores[j]
		}
		for d := 0; d < dsa.HeadDim; d++ {
			var val float32
			for j, s := range scores {
				val += (s / total) * V[j][d]
			}
			output[i][d] = val
		}
	}
	return output
}

// Forward DSA前向传播
func (dsa *DenseSparseAttention) Forward(
	Q, K, V [][]float32,
) [][]float32 {
	// 1. Token压缩
	scores := dsa.TokenCompressionScore(Q)
	compressedK, indices := dsa.CompressTokens(K, scores)

	// 2. 密集注意力路径：压缩序列
	denseOut := dsa.DenseAttention(Q, compressedK, compressedK)

	// 3. 稀疏注意力路径：原始序列
	sparseOut := dsa.SparseAttention(Q, K, V, indices, 64)

	// 4. 融合
	alpha := float32(0.7) // 密集路径权重
	beta := float32(0.3)  // 稀疏路径权重
	seqLen := len(Q)
	fused := make([][]float32, seqLen)
	for i := 0; i < seqLen; i++ {
		fused[i] = make([]float32, dsa.HeadDim)
		for d := 0; d < dsa.HeadDim; d++ {
			fused[i][d] = alpha*denseOut[i][d] + beta*sparseOut[i][d]
		}
	}
	return fused
}

"""
DSpark: 推测性解码框架
DeepSeek × 北京大学联合实现
"""
import torch
import torch.nn.functional as F
from typing import List, Optional, Tuple
from dataclasses import dataclass


@dataclass
class SpeculativeResult:
    """推测性解码结果"""
    tokens: List[int]
    acceptance_rate: float
    speedup: float
    draft_steps: int
    target_steps: int
    rejected_positions: List[int]


class DSparkDecoder:
    """
    DSpark推测性解码器
    
    使用小型草稿模型生成候选序列，大型目标模型并行验证
    实现"一步生成多个token"的效果
    """
    
    def __init__(
        self,
        draft_model,      # 草稿模型（V4 Flash, 13B）
        target_model,     # 目标模型（V4 Pro, 49B）
        gamma: int = 5,   # 推测窗口大小
        temperature: float = 0.7,
        top_k: int = 40,
    ):
        self.draft_model = draft_model
        self.target_model = target_model
        self.gamma = gamma
        self.temperature = temperature
        self.top_k = top_k
    
    def generate(
        self,
        input_ids: torch.Tensor,
        max_new_tokens: int = 256,
    ) -> SpeculativeResult:
        """
        使用推测性解码生成文本
        
        Args:
            input_ids: 输入token序列
            max_new_tokens: 最大生成token数
            
        Returns:
            生成结果，包含token序列和性能指标
        """
        device = input_ids.device
        prefix = input_ids.clone()
        
        all_tokens = input_ids[0].tolist()
        total_draft_steps = 0
        total_target_steps = 0
        rejected_positions = []
        
        while len(all_tokens) < len(input_ids[0]) + max_new_tokens:
            # 1. 草稿阶段：用小模型快速生成gamma个候选token
            draft_tokens = self._draft_generate(prefix, self.gamma)
            total_draft_steps += 1
            
            # 2. 验证阶段：用大模型并行验证
            n_accepted, target_logits = self._verify(
                prefix, draft_tokens
            )
            total_target_steps += 1
            
            # 3. 接受/拒绝决策
            if n_accepted == 0:
                # 全部拒绝：从target采样一个token
                next_token = self._sample(target_logits[0])
                prefix = torch.cat([
                    prefix, 
                    torch.tensor([[next_token]], device=device)
                ], dim=1)
                all_tokens.append(next_token)
                rejected_positions.append(len(all_tokens) - 1)
            else:
                # 接受前n_accepted个token
                for i in range(n_accepted):
                    prefix = torch.cat([
                        prefix,
                        torch.tensor([[draft_tokens[i]]], device=device)
                    ], dim=1)
                    all_tokens.append(draft_tokens[i])
                
                # 如果需要拒绝后的补偿token
                if n_accepted < self.gamma:
                    bonus_token = self._sample(target_logits[n_accepted])
                    prefix = torch.cat([
                        prefix,
                        torch.tensor([[bonus_token]], device=device)
                    ], dim=1)
                    all_tokens.append(bonus_token)
                    rejected_positions.append(len(all_tokens) - 1)
        
        # 计算指标
        total_generated = len(all_tokens) - len(input_ids[0])
        acceptance_rate = 1 - len(rejected_positions) / max(
            1, total_generated
        )
        # 理论加速比
        speedup = self.gamma * total_target_steps / max(
            1, total_draft_steps + total_target_steps
        )
        
        return SpeculativeResult(
            tokens=all_tokens,
            acceptance_rate=acceptance_rate,
            speedup=speedup,
            draft_steps=total_draft_steps,
            target_steps=total_target_steps,
            rejected_positions=rejected_positions,
        )
    
    def _draft_generate(
        self, prefix: torch.Tensor, gamma: int
    ) -> List[int]:
        """草稿模型自回归生成候选序列"""
        draft_tokens = []
        current = prefix.clone()
        
        with torch.no_grad():
            for _ in range(gamma):
                logits = self.draft_model(current)
                next_logits = logits[0, -1, :]
                next_token = self._sample(next_logits)
                draft_tokens.append(next_token)
                current = torch.cat([
                    current,
                    torch.tensor([[next_token]], 
                                 device=current.device)
                ], dim=1)
        
        return draft_tokens
    
    def _verify(
        self, prefix: torch.Tensor, draft_tokens: List[int]
    ) -> Tuple[int, torch.Tensor]:
        """
        使用目标模型验证草稿token
        
        核心思想：目标模型并行计算每个位置的logits
        通过拒绝采样决定接受哪些草稿token
        """
        device = prefix.device
        gamma = len(draft_tokens)
        
        # 构建验证序列：prefix + draft_tokens
        draft_tensor = torch.tensor(
            [draft_tokens], device=device
        )
        verification_input = torch.cat(
            [prefix, draft_tensor], dim=1
        )
        
        # 目标模型前向传播（一次计算所有位置）
        with torch.no_grad():
            logits = self.target_model(verification_input)
        
        # 提取每个位置的目标分布
        target_probs = []
        prefix_len = prefix.shape[1]
        for i in range(gamma):
            pos_logits = logits[0, prefix_len + i, :]
            probs = F.softmax(pos_logits / self.temperature, dim=-1)
            target_probs.append(probs)
        
        # 拒绝采样决策
        n_accepted = 0
        for i in range(gamma):
            draft_token = draft_tokens[i]
            draft_prob = self._get_draft_prob(
                prefix if i == 0 else 
                torch.cat([prefix, draft_tensor[:, :i]], dim=1),
                draft_token
            )
            target_prob = target_probs[i][draft_token].item()
            
            if target_prob > draft_prob:
                # 接受概率 = target_prob / draft_prob
                acceptance_prob = target_prob / draft_prob
                if torch.rand(1).item() <= acceptance_prob:
                    n_accepted += 1
                else:
                    break
            else:
                n_accepted += 1
        
        # 返回接受数和补偿token的目标logits
        bonus_idx = min(n_accepted, gamma - 1)
        bonus_logits = logits[0, prefix_len + bonus_idx, :]
        
        return n_accepted, torch.stack(
            [logits[0, prefix_len + i, :] for i in range(gamma)]
        )
    
    def _get_draft_prob(
        self, prefix: torch.Tensor, token: int
    ) -> float:
        """获取草稿模型对指定token的概率"""
        with torch.no_grad():
            logits = self.draft_model(prefix)
            probs = F.softmax(
                logits[0, -1, :] / self.temperature, dim=-1
            )
        return probs[token].item()
    
    def _sample(self, logits: torch.Tensor) -> int:
        """从logits采样token"""
        # Top-k过滤
        top_k_logits, top_k_indices = torch.topk(
            logits, self.top_k, dim=-1
        )
        probs = F.softmax(
            top_k_logits / self.temperature, dim=-1
        )
        sampled_idx = torch.multinomial(probs, 1)
        return top_k_indices[sampled_idx].item()

package pricing

import "time"

// Tier 定价层级
type Tier int

const (
	OffPeak Tier = iota // 低谷
	Peak                // 高峰
)

// PriceConfig 定价配置
type PriceConfig struct {
	ProInputCacheHit  float64 // Pro缓存命中输入
	ProInputCacheMiss float64 // Pro缓存未命中输入
	ProOutput         float64 // Pro输出
	FlashInputCacheHit  float64
	FlashInputCacheMiss float64
	FlashOutput         float64
}

// PeakPriceManager 峰谷定价管理器
type PeakPriceManager struct {
	basePrices   PriceConfig
	peakPrices   PriceConfig
	peakStart1   int // 第一个高峰开始时间
	peakEnd1     int // 第一个高峰结束时间
	peakStart2   int // 第二个高峰开始时间
	peakEnd2     int // 第二个高峰结束时间
	peakFactor   float64 // 高峰倍数
}

func NewDefaultPeakPriceManager() *PeakPriceManager {
	return &PeakPriceManager{
		basePrices: PriceConfig{
			ProInputCacheHit:  0.025,
			ProInputCacheMiss: 3.0,
			ProOutput:         6.0,
			FlashInputCacheHit:  0.02,
			FlashInputCacheMiss: 1.0,
			FlashOutput:         2.0,
		},
		peakFactor: 2.0,
		peakStart1: 9,
		peakEnd1:   12,
		peakStart2: 14,
		peakEnd2:   18,
	}
}

// GetCurrentTier 判断当前时段
func (pm *PeakPriceManager) GetCurrentTier(t time.Time) Tier {
	h := t.Hour()
	if (h >= pm.peakStart1 && h < pm.peakEnd1) ||
		(h >= pm.peakStart2 && h < pm.peakEnd2) {
		return Peak
	}
	return OffPeak
}

// LoadAwarePricing 负载感知定价
type LoadAwarePricing struct {
	pm           *PeakPriceManager
	loadHistory  map[int]float64 // 小时→平均负载
	cacheHitRate float64         // 当前缓存命中率
}

func NewLoadAwarePricing(pm *PeakPriceManager) *LoadAwarePricing {
	return &LoadAwarePricing{
		pm:          pm,
		loadHistory: make(map[int]float64),
	}
}

// CalculateDynamicPrice 动态计算价格
func (lp *LoadAwarePricing) CalculateDynamicPrice(
	model string,     // "pro" 或 "flash"
	output bool,      // 是否输出
	cacheHit bool,    // 是否缓存命中
) float64 {
	tier := lp.pm.GetCurrentTier(time.Now())
	
	// 基础价格
	var basePrice float64
	switch {
	case model == "pro" && output:
		basePrice = lp.pm.basePrices.ProOutput
	case model == "pro" && cacheHit:
		basePrice = lp.pm.basePrices.ProInputCacheHit
	case model == "pro" && !cacheHit:
		basePrice = lp.pm.basePrices.ProInputCacheMiss
	case model == "flash" && output:
		basePrice = lp.pm.basePrices.FlashOutput
	case model == "flash" && cacheHit:
		basePrice = lp.pm.basePrices.FlashInputCacheHit
	default:
		basePrice = lp.pm.basePrices.FlashInputCacheMiss
	}

	if tier == Peak {
		basePrice *= lp.pm.peakFactor
	}

	// 缓存命中折扣
	if cacheHit {
		basePrice *= 0.008 // 缓存命中价格约为未命中的1/120
	}

	return basePrice
}

// OptimalSchedule 优化调度建议
type ScheduleAdvice struct {
	BestTime   time.Time
	CostSaving float64 // 相比立即执行的节省比例
	Reason     string
}

func (lp *LoadAwarePricing) AdviseSchedule(
	batchSize int,        // 批量任务大小
	maxDelay time.Duration, // 最大可接受延迟
) *ScheduleAdvice {
	now := time.Now()
	bestCost := 1e9
	var bestTime time.Time

	// 在maxDelay范围内扫描最优时间
	for t := now; t.Before(now.Add(maxDelay)); t = t.Add(time.Hour) {
		cost := lp.CalculateDynamicPrice("flash", true, false)
		if cost < bestCost {
			bestCost = cost
			bestTime = t
		}
	}

	saving := (lp.CalculateDynamicPrice("flash", true, false) - bestCost) /
		lp.CalculateDynamicPrice("flash", true, false) * 100

	return &ScheduleAdvice{
		BestTime:   bestTime,
		CostSaving: saving,
		Reason:     "低谷时段价格减半",
	}
}