腾讯混元Hy3正式版深度解析：295B MoE快慢思考融合架构，Apache 2.0开源，Agent任务解决率90%

import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List, Tuple, Optional

class SparseMoERouter(nn.Module):
    """
    Hy3 Sparse MoE Router with Top-2 routing and load balancing.
    
    实现思路：
    1. 对每个输入token，计算与所有N个专家的匹配分数
    2. 选择Top-2分数最高的专家进行激活
    3. 使用辅助损失确保专家负载均衡，避免"死专家"
    """
    def __init__(
        self,
        hidden_dim: int = 7168,        # Hy3 hidden dimension
        num_experts: int = 32,          # 总专家数
        top_k: int = 2,                 # 每次激活的专家数
        capacity_factor: float = 1.25,  # 容量因子，允许一定超额
    ):
        super().__init__()
        self.num_experts = num_experts
        self.top_k = top_k
        self.capacity_factor = capacity_factor
        
        # 路由网络：将hidden_dim映射到num_experts个分数
        self.gate = nn.Linear(hidden_dim, num_experts, bias=False)
        
        # 负载均衡的计数器（训练时维护）
        self.register_buffer("expert_counts", torch.zeros(num_experts))
        self.register_buffer("total_tokens", torch.tensor(0.0))
        
    def forward(
        self, 
        x: torch.Tensor,               # [batch, seq_len, hidden_dim]
        training: bool = True
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
        """
        Returns:
            - routing_weights: [batch, seq_len, top_k] 各专家的权重
            - expert_indices: [batch, seq_len, top_k] 选择的专家索引
            - aux_loss: 负载均衡辅助损失
        """
        batch, seq_len, _ = x.shape
        x_flat = x.view(-1, x.size(-1))  # [batch*seq_len, hidden_dim]
        n_tokens = x_flat.size(0)
        
        # 路由分数计算
        logits = self.gate(x_flat)  # [batch*seq_len, num_experts]
        
        # Softmax归一化
        routing_scores = F.softmax(logits, dim=-1)  # [batch*seq_len, num_experts]
        
        # Top-2选择
        top_k_scores, top_k_indices = torch.topk(
            routing_scores, k=self.top_k, dim=-1
        )  # 各 [batch*seq_len, top_k]
        
        # 权重归一化（确保sum=1）
        routing_weights = top_k_scores / (top_k_scores.sum(dim=-1, keepdim=True) + 1e-8)
        
        # 负载均衡辅助损失
        aux_loss = self._compute_load_balancing_loss(routing_scores, top_k_indices, training)
        
        # 更新统计（仅训练时）
        if training:
            with torch.no_grad():
                counts = torch.zeros(self.num_experts, device=x.device)
                for i in range(self.top_k):
                    indices = top_k_indices[:, i]
                    counts.scatter_add_(0, indices, torch.ones_like(indices, dtype=torch.float))
                self.expert_counts = self.expert_counts * 0.9 + counts * 0.1
                self.total_tokens = self.total_tokens * 0.9 + n_tokens * 0.1
        
        # 重塑输出
        routing_weights = routing_weights.view(batch, seq_len, self.top_k)
        expert_indices = top_k_indices.view(batch, seq_len, self.top_k)
        
        return routing_weights, expert_indices, aux_loss
    
    def _compute_load_balancing_loss(
        self,
        scores: torch.Tensor,       # [batch*seq_len, num_experts]
        indices: torch.Tensor,      # [batch*seq_len, top_k]
        training: bool
    ) -> torch.Tensor:
        """
        负载均衡辅助损失（Z-loss variant）
        
        核心思想：鼓励每个专家被选中的概率接近均匀分布
        loss = alpha * num_experts * sum(P_i * f_i)
        其中P_i是路由概率的平均，f_i是实际被选中的频率
        """
        if not training:
            return torch.tensor(0.0, device=scores.device)
        
        num_experts = self.num_experts
        n_tokens = scores.size(0)
        
        # P_i: 路由概率的平均
        P = scores.mean(dim=0)  # [num_experts]
        
        # f_i: 实际被选中频率
        f = torch.zeros(num_experts, device=scores.device)
        for i in range(self.top_k):
            f.scatter_add_(0, indices[:, i], torch.ones(n_tokens, device=scores.device))
        f = f / (n_tokens * self.top_k)
        
        # Z-loss: sum(P_i * f_i)
        aux_loss = torch.sum(P * f) * num_experts
        
        return aux_loss


class MoEFeedForward(nn.Module):
    """
    MoE专家前馈网络（每个专家独立）
    
    Hy3专家结构：SwiGLU激活的FFN
    """
    def __init__(
        self,
        hidden_dim: int = 7168,
        intermediate_dim: int = 20480,  # 约2.86x hidden_dim
    ):
        super().__init__()
        self.gate_proj = nn.Linear(hidden_dim, intermediate_dim, bias=False)
        self.up_proj = nn.Linear(hidden_dim, intermediate_dim, bias=False)
        self.down_proj = nn.Linear(intermediate_dim, hidden_dim, bias=False)
        
    def forward(self, x: torch.Tensor) -> torch.Tensor:
        # SwiGLU激活
        gate = F.silu(self.gate_proj(x))
        up = self.up_proj(x)
        return self.down_proj(gate * up)

package mtp

import (
	"math"
)

// MTPHead 多令牌预测头
// 在Hy3中，MTP层有3.8B参数，包含3个独立的预测头
// 每个头预测当前位置之后的第n个token
type MTPHead struct {
	// 3个预测头的线性变换
	Heads [3]*LinearTransform
	// 层归一化
	Norm *LayerNorm
	// 共享的embedding矩阵
	SharedEmbedding *EmbeddingMatrix
}

// LinearTransform 简化的线性变换
type LinearTransform struct {
	Weights [][]float32
	Bias    []float32
}

// LayerNorm 层归一化
type LayerNorm struct {
	Gamma []float32
	Beta  []float32
	Eps   float32
}

// EmbeddingMatrix 共享embedding矩阵
type EmbeddingMatrix struct {
	VocabSize int
	Dim       int
	Weights   [][]float32
}

// MTPOutput MTP多步预测输出
type MTPOutput struct {
	// 当前token的预测（主头）
	MainLogits []float32
	// 未来n步的预测
	FutureLogits [3][]float32
	// 置信度得分
	ConfidenceScores [3]float32
}

// Forward 执行MTP前向推理
func (m *MTPHead) Forward(hiddenState []float32) *MTPOutput {
	output := &MTPOutput{}

	// 1. 层归一化
	normalized := m.Norm.Forward(hiddenState)

	// 2. 主头预测（当前token）
	mainLogits := m.Heads[0].Forward(normalized)
	output.MainLogits = mainLogits

	// 3. 多步预测头
	for i := 0; i < 3; i++ {
		// 每个预测头有自己的参数
		futureLogits := m.Heads[i].Forward(normalized)
		output.FutureLogits[i] = futureLogits

		// 计算置信度（基于logits的熵）
		output.ConfidenceScores[i] = computeConfidence(futureLogits)
	}

	return output
}

// computeConfidence 基于logits计算预测置信度
// 使用softmax分布的负熵作为置信度指标
func computeConfidence(logits []float32) float32 {
	var maxLogit float32 = logits[0]
	for _, l := range logits[1:] {
		if l > maxLogit {
			maxLogit = l
		}
	}

	// Softmax数值稳定性处理
	var sum float32 = 0
	probs := make([]float32, len(logits))
	for i, l := range logits {
		probs[i] = float32(math.Exp(float64(l - maxLogit)))
		sum += probs[i]
	}

	// 归一化
	for i := range probs {
		probs[i] /= sum
	}

	// 计算熵的负值（越高越确定）
	var entropy float32 = 0
	for _, p := range probs {
		if p > 1e-10 {
			entropy -= p * float32(math.Log(float64(p)))
		}
	}

	// 归一化到[0,1]
	maxEntropy := float32(math.Log(float64(len(logits))))
	return 1.0 - entropy/maxEntropy
}

// BatchPredict 批量多步预测
// 同时处理批次中的多个序列
func (m *MTPHead) BatchPredict(hiddenStates [][]float32) []*MTPOutput {
	outputs := make([]*MTPOutput, len(hiddenStates))
	for i, hs := range hiddenStates {
		outputs[i] = m.Forward(hs)
	}
	return outputs
}

from enum import Enum
from dataclasses import dataclass
from typing import Optional, Callable
import numpy as np


class ThinkingMode(Enum):
    """推理模式枚举"""
    FAST = "fast"        # 系统1：快速直觉推理
    SLOW = "slow"        # 系统2：深度链式推理
    HYBRID = "hybrid"    # 混合模式：快慢结合


@dataclass
class ComplexityFeatures:
    """复杂度特征向量"""
    input_length: int                    # 输入长度
    num_entities: int                    # 实体数量
    has_math: bool                       # 是否包含数学
    has_code: bool                       # 是否包含代码
    has_multi_step: bool                 # 是否需要多步推理
    domain: str                          # 领域
    ambiguity_score: float               # 模糊度评分
    
    @property
    def complexity_score(self) -> float:
        """综合复杂度评分"""
        score = 0.0
        score += min(self.input_length / 2000, 1.0) * 0.2
        score += min(self.num_entities / 10, 1.0) * 0.15
        score += 0.2 if self.has_math else 0
        score += 0.15 if self.has_code else 0
        score += 0.2 if self.has_multi_step else 0
        score += self.ambiguity_score * 0.1
        return min(score, 1.0)


class ThinkingRouter:
    """
    快慢思考路由控制器
    
    根据输入复杂度自动选择推理模式：
    - 低复杂度（<0.3）：快速模式
    - 中复杂度（0.3-0.7）：混合模式
    - 高复杂度（>0.7）：慢思考模式
    """
    
    def __init__(self, fast_threshold: float = 0.3, slow_threshold: float = 0.7):
        self.fast_threshold = fast_threshold
        self.slow_threshold = slow_threshold
        self._history: list = []
        
    def select_mode(self, features: ComplexityFeatures) -> ThinkingMode:
        """
        根据复杂度特征选择推理模式
        """
        score = features.complexity_score
        
        # 记录历史
        self._history.append({
            "features": features,
            "score": score,
            "mode": None
        })
        
        if score < self.fast_threshold:
            mode = ThinkingMode.FAST
        elif score > self.slow_threshold:
            mode = ThinkingMode.SLOW
        else:
            mode = ThinkingMode.HYBRID
            
        self._history[-1]["mode"] = mode
        return mode
    
    def get_mode_stats(self) -> dict:
        """获取模式使用统计"""
        modes = [h["mode"] for h in self._history]
        return {
            "fast_ratio": modes.count(ThinkingMode.FAST) / len(modes) if modes else 0,
            "slow_ratio": modes.count(ThinkingMode.SLOW) / len(modes) if modes else 0,
            "hybrid_ratio": modes.count(ThinkingMode.HYBRID) / len(modes) if modes else 0,
        }


class FastReasoner:
    """
    快速推理引擎（系统1）
    
    执行1-2步直接推理，适合简单问答、常识判断、信息抽取等场景
    """
    
    def __init__(self, model_fn: Callable):
        self.model_fn = model_fn
        
    def reason(self, prompt: str, max_steps: int = 2) -> str:
        """
        快速推理：直接生成答案，不进行链式思考
        """
        # 使用特殊标记指示模型进行直接推理
        fast_prompt = f"[FAST_REASON]{prompt}\n直接回答："
        
        # 调用底层模型
        result = self.model_fn(fast_prompt, max_new_tokens=256)
        
        return result


class SlowReasoner:
    """
    深度推理引擎（系统2）
    
    执行多步Chain-of-Thought推理，适合数学、代码、Agent等复杂任务
    """
    
    def __init__(self, model_fn: Callable, max_chain_steps: int = 32):
        self.model_fn = model_fn
        self.max_chain_steps = max_chain_steps
        
    def reason(self, prompt: str, task_type: str = "general") -> str:
        """
        深度推理：多步链式思考
        """
        # 根据任务类型选择CoT模板
        templates = {
            "math": "[SLOW_REASON_MATH]{prompt}\n让我们一步一步思考：",
            "code": "[SLOW_REASON_CODE]{prompt}\n让我们分析需求并逐步实现：",
            "agent": "[SLOW_REASON_AGENT]{prompt}\n让我们逐步规划并执行：",
            "general": "[SLOW_REASON]{prompt}\n让我们仔细分析：",
        }
        
        template = templates.get(task_type, templates["general"])
        slow_prompt = template.format(prompt=prompt)
        
        # 多步推理，支持中间结果回溯
        result = self.model_fn(
            slow_prompt, 
            max_new_tokens=4096,
            temperature=0.7
        )
        
        return result

class HallucinationMitigationPipeline:
    """
    Hy3幻觉治理三阶段管线
    
    Phase 1: 数据清洗 - 移除训练数据中的事实错误
    Phase 2: 训练约束 - 引入幻觉感知损失
    Phase 3: 推理约束 - 事实一致性校验
    """
    
    def __init__(self):
        self.fact_checker = FactChecker()
        self.hallucination_detector = HallucinationDetector()
        
    def train_pipeline(self, dataset: list) -> list:
        """训练阶段数据清洗"""
        clean_data = []
        
        for sample in dataset:
            # 1. 事实性校验
            fact_score = self.fact_checker.verify(sample["text"])
            
            # 2. 幻觉检测
            hal_score = self.hallucination_detector.detect(sample["text"])
            
            # 3. 综合质量评分
            quality = (fact_score * 0.6 + (1 - hal_score) * 0.4)
            
            if quality > 0.7:
                sample["quality_weight"] = quality
                clean_data.append(sample)
                
        return clean_data


class FactChecker:
    """
    事实性校验器
    
    基于多源交叉验证，检测训练数据中的事实错误
    """
    
    def __init__(self):
        self.knowledge_base = {}
        self.entity_extractor = EntityExtractor()
        
    def verify(self, text: str) -> float:
        """返回事实性得分 [0, 1]"""
        entities = self.entity_extractor.extract(text)
        
        if not entities:
            return 1.0
        
        verified = 0
        total = 0
        
        for entity, claim in entities:
            total += 1
            if self._check_claim(entity, claim):
                verified += 1
                
        return verified / total if total > 0 else 1.0
    
    def _check_claim(self, entity: str, claim: str) -> bool:
        """检查单条声明的事实性"""
        # 基于知识库的交叉验证
        known_facts = self.knowledge_base.get(entity, [])
        for fact in known_facts:
            if self._semantic_match(claim, fact):
                return True
        return False
    
    def _semantic_match(self, claim: str, fact: str) -> bool:
        """语义级别的事实匹配"""
        # 简化实现：实际使用语义嵌入相似度
        return claim.lower() in fact.lower() or fact.lower() in claim.lower()

package agent

import (
	"context"
	"encoding/json"
	"fmt"
	"sync"
	"time"
)

// ToolCall Agent工具调用定义
type ToolCall struct {
	ID        string                 `json:"id"`
	Name      string                 `json:"name"`
	Arguments map[string]interface{} `json:"arguments"`
	Status    string                 `json:"status"` // pending, running, success, failed
	Result    string                 `json:"result,omitempty"`
	Error     string                 `json:"error,omitempty"`
	StartTime time.Time              `json:"start_time"`
	Duration  time.Duration          `json:"duration"`
}

// AgentContext Agent执行上下文
type AgentContext struct {
	SessionID    string
	Conversation []Message
	ToolRegistry map[string]Tool
	MaxSteps     int
	CurrentStep  int
	mu           sync.RWMutex
}

// Tool 工具接口
type Tool interface {
	Name() string
	Description() string
	Execute(ctx context.Context, args map[string]interface{}) (string, error)
}

// AgentExecutor Hy3 Agent执行器
type AgentExecutor struct {
	Model      *Hy3Model
	Context    *AgentContext
	MaxRetries int
}

// ExecuteAgent 执行Agent任务
func (ae *AgentExecutor) ExecuteAgent(ctx context.Context, task string) (*AgentResult, error) {
	startTime := time.Now()
	
	// 1. 任务分解
	steps, err := ae.decomposeTask(ctx, task)
	if err != nil {
		return nil, fmt.Errorf("task decomposition failed: %w", err)
	}
	
	result := &AgentResult{
		Task:       task,
		Steps:      steps,
		StartTime:  startTime,
		TotalSteps: len(steps),
	}
	
	// 2. 逐步执行
	for i, step := range steps {
		select {
		case <-ctx.Done():
			result.Status = "cancelled"
			result.Error = ctx.Err().Error()
			return result, nil
		default:
		}
		
		stepResult, err := ae.executeStep(ctx, &step, i)
		if err != nil {
			// 错误重试
			for retry := 0; retry < ae.MaxRetries; retry++ {
				stepResult, err = ae.executeStep(ctx, &step, i)
				if err == nil {
					break
				}
			}
		}
		
		step.Result = stepResult
		step.Status = "completed"
		if err != nil {
			step.Status = "failed"
			step.Error = err.Error()
		}
		steps[i] = step
		result.CompletedSteps++
	}
	
	// 3. 结果汇总
	result.EndTime = time.Now()
	result.Duration = result.EndTime.Sub(result.StartTime)
	
	if result.CompletedSteps == result.TotalSteps {
		result.Status = "success"
	} else {
		result.Status = "partial"
	}
	
	// 4. 结果验证
	result.Validation = ae.validateResult(ctx, task, result)
	
	return result, nil
}

// decomposeTask 将任务分解为子步骤
func (ae *AgentExecutor) decomposeTask(ctx context.Context, task string) ([]AgentStep, error) {
	prompt := fmt.Sprintf(`将以下任务分解为可执行的子步骤，每个步骤应该：
1. 明确描述需要做什么
2. 指定需要的工具
3. 定义成功标准

任务：%s

请以JSON数组格式输出每个步骤的详细描述。`, task)
	
	response, err := ae.Model.Generate(ctx, prompt)
	if err != nil {
		return nil, err
	}
	
	var steps []AgentStep
	if err := json.Unmarshal([]byte(response), &steps); err != nil {
		// 降级：简单分解
		steps = []AgentStep{
			{Description: task, Tools: []string{"search", "code"}},
		}
	}
	
	return steps, nil
}

// executeStep 执行单个步骤
func (ae *AgentExecutor) executeStep(ctx context.Context, step *AgentStep, stepIdx int) (string, error) {
	// 1. 构建步骤提示
	stepPrompt := fmt.Sprintf(
		"[AGENT_STEP %d/%d]\n任务描述：%s\n可用工具：%v\n历史上下文：%s\n请执行此步骤：",
		stepIdx+1, len(ae.Context.MaxSteps),
		step.Description,
		step.Tools,
		ae.Context.getRecentContext(),
	)
	
	// 2. 模型生成执行计划
	plan, err := ae.Model.Generate(ctx, stepPrompt)
	if err != nil {
		return "", fmt.Errorf("step planning failed: %w", err)
	}
	
	// 3. 解析并执行工具调用
	var toolCalls []ToolCall
	if err := json.Unmarshal([]byte(plan), &toolCalls); err == nil {
		for _, tc := range toolCalls {
			tool, ok := ae.Context.ToolRegistry[tc.Name]
			if !ok {
				continue
			}
			
			tc.Status = "running"
			tc.StartTime = time.Now()
			
			toolCtx, cancel := context.WithTimeout(ctx, 30*time.Second)
			result, err := tool.Execute(toolCtx, tc.Arguments)
			cancel()
			
			tc.Duration = time.Since(tc.StartTime)
			if err != nil {
				tc.Status = "failed"
				tc.Error = err.Error()
			} else {
				tc.Status = "success"
				tc.Result = result
			}
		}
	}
	
	// 4. 汇总结果
	return ae.summarizeStepResult(ctx, step, toolCalls)
}

// validateResult 验证Agent执行结果
func (ae *AgentExecutor) validateResult(ctx context.Context, task string, result *AgentResult) *ValidationReport {
	report := &ValidationReport{}
	
	// 检查是否所有步骤完成
	if result.CompletedSteps < result.TotalSteps {
		report.AddIssue("部分步骤未完成")
	}
	
	// 检查结果质量
	prompt := fmt.Sprintf(
		"验证以下Agent任务执行结果是否满足原始需求：\n原始任务：%s\n执行结果：%s\n质量评分（0-100）：",
		task, result.Summary(),
	)
	
	score, err := ae.Model.Generate(ctx, prompt)
	if err == nil {
		fmt.Sscanf(score, "%d", &report.QualityScore)
	}
	
	return report
}

package inference

import (
	"sync"
	"sync/atomic"
)

// Hy3InferenceEngine Hy3推理优化引擎
// 核心优化点：
// 1. KV Cache管理
// 2. PagedAttention
// 3. 连续批处理
// 4. 量化推理
type Hy3InferenceEngine struct {
	Model            *Hy3Model
	BatchSize        int
	MaxSeqLen        int
	KVcacheManager   *KVcacheManager
	BatchScheduler   *ContinuousBatchScheduler
	Quantization     *QuantizationConfig
	requestQueue     chan *InferenceRequest
	resultQueue      chan *InferenceResult
	activeRequests   sync.WaitGroup
	isRunning        atomic.Bool
}

// KVcacheManager KV缓存管理器
// 管理256K上下文的KV缓存，支持缓存复用和驱逐
type KVcacheManager struct {
	mu          sync.RWMutex
	cache       map[string]*KVcacheEntry
	maxEntries  int
	hitCount    atomic.Int64
	missCount   atomic.Int64
}

// KVcacheEntry KV缓存条目
type KVcacheEntry struct {
	SessionID    string
	PrefixHash   uint64
	KeyCache     [][][]float32 // [num_layers][num_heads, seq_len, head_dim]
	ValueCache   [][][]float32
	SeqLen       int
	LastAccess   int64
	RefCount     int32
}

// GetOrCreateCache 获取或创建KV缓存
func (m *KVcacheManager) GetOrCreateCache(sessionID string, prefixHash uint64) *KVcacheEntry {
	m.mu.RLock()
	entry, ok := m.cache[sessionID]
	m.mu.RUnlock()
	
	if ok {
		m.hitCount.Add(1)
		atomic.StoreInt64(&entry.LastAccess, time.Now().UnixNano())
		atomic.AddInt32(&entry.RefCount, 1)
		return entry
	}
	
	m.mu.Lock()
	defer m.mu.Unlock()
	
	// 二次检查
	if entry, ok := m.cache[sessionID]; ok {
		return entry
	}
	
	// 驱逐策略：LRU
	if len(m.cache) >= m.maxEntries {
		m.evictLRU()
	}
	
	entry = &KVcacheEntry{
		SessionID:  sessionID,
		PrefixHash: prefixHash,
		LastAccess: time.Now().UnixNano(),
		RefCount:   1,
	}
	
	m.cache[sessionID] = entry
	m.missCount.Add(1)
	return entry
}

// evictLRU LRU驱逐
func (m *KVcacheManager) evictLRU() {
	var oldestKey string
	var oldestTime int64 = math.MaxInt64
	
	for k, v := range m.cache {
		if atomic.LoadInt32(&v.RefCount) == 0 && v.LastAccess < oldestTime {
			oldestKey = k
			oldestTime = v.LastAccess
		}
	}
	
	if oldestKey != "" {
		delete(m.cache, oldestKey)
	}
}

// ContinuousBatchScheduler 连续批处理调度器
// 支持动态添加/移除请求，最大化GPU利用率
type ContinuousBatchScheduler struct {
	mu             sync.Mutex
	pendingBatch   []*InferenceRequest
	runningBatch   []*InferenceRequest
	maxBatchSize   int
	prefillLimit   int
	decodingLimit  int
}

// Schedule 调度批处理
func (s *ContinuousBatchScheduler) Schedule() []*InferenceRequest {
	s.mu.Lock()
	defer s.mu.Unlock()
	
	// 1. 完成当前批次中已结束的请求
	s.completeFinished()
	
	// 2. 从待处理队列中取出新请求
	slots := s.maxBatchSize - len(s.runningBatch)
	if slots > 0 && len(s.pendingBatch) > 0 {
		newCount := min(slots, len(s.pendingBatch))
		s.runningBatch = append(s.runningBatch, s.pendingBatch[:newCount]...)
		s.pendingBatch = s.pendingBatch[newCount:]
	}
	
	// 3. 返回当前批次
	batch := make([]*InferenceRequest, len(s.runningBatch))
	copy(batch, s.runningBatch)
	return batch
}

// completeFinished 移除已完成请求
func (s *ContinuousBatchScheduler) completeFinished() {
	active := make([]*InferenceRequest, 0, len(s.runningBatch))
	for _, req := range s.runningBatch {
		if !req.IsFinished() {
			active = append(active, req)
		}
	}
	s.runningBatch = active
}