// Go语言实现：简单的对称量化函数
package quantization

import (
	"fmt"
	"math"
)

// QuantizeWeights performs symmetric INT8 quantization on weight matrix
// Input: weights in float32, scale factor calculated from max absolute value
// Output: quantized weights in int8, and the scale factor
func QuantizeWeights(weights []float32) ([]int8, float32) {
	// Step 1: Find the maximum absolute value in weights
	maxAbs := float32(0.0)
	for _, w := range weights {
		absW := float32(math.Abs(float64(w)))
		if absW > maxAbs {
			maxAbs = absW
		}
	}

	// Step 2: Calculate scale factor (127 / maxAbs)
	// This ensures the quantized range [-127, 127] covers the full weight range
	scale := 127.0 / maxAbs

	// Step 3: Quantize each weight
	quantized := make([]int8, len(weights))
	for i, w := range weights {
		// Clamp to [-127, 127] to avoid overflow
		qVal := int(math.Round(float64(w * scale)))
		if qVal > 127 {
			qVal = 127
		} else if qVal < -127 {
			qVal = -127
		}
		quantized[i] = int8(qVal)
	}

	return quantized, scale
}

// DequantizeWeights converts INT8 weights back to float32 for inference
func DequantizeWeights(quantized []int8, scale float32) []float32 {
	dequantized := make([]float32, len(quantized))
	for i, q := range quantized {
		dequantized[i] = float32(q) / scale
	}
	return dequantized
}

// Example usage
func main() {
	// Simulate original weights (FP32)
	originalWeights := []float32{0.5, -1.2, 3.4, -0.8, 2.1, -0.3}
	
	// Quantize
	qWeights, scale := QuantizeWeights(originalWeights)
	fmt.Printf("Original weights: %v\n", originalWeights)
	fmt.Printf("Quantized weights (INT8): %v\n", qWeights)
	fmt.Printf("Scale factor: %f\n", scale)
	
	// Dequantize for comparison
	dqWeights := DequantizeWeights(qWeights, scale)
	fmt.Printf("Dequantized weights: %v\n", dqWeights)
	
	// Calculate quantization error
	var totalError float32
	for i := range originalWeights {
		err := originalWeights[i] - dqWeights[i]
		totalError += err * err
	}
	fmt.Printf("Mean squared error: %f\n", totalError/float32(len(originalWeights)))
}

// Go语言实现：基于幅度的权重剪枝
package pruning

import (
	"math"
	"sort"
)

// WeightMagnitudePruner implements magnitude-based weight pruning
type WeightMagnitudePruner struct {
	Sparsity float64 // Target sparsity (0.0 to 1.0)
}

// NewWeightMagnitudePruner creates a new pruner with target sparsity
func NewWeightMagnitudePruner(sparsity float64) *WeightMagnitudePruner {
	return &WeightMagnitudePruner{Sparsity: sparsity}
}

// Prune performs weight pruning based on magnitude
// Returns pruned weights and a mask indicating which weights are kept
func (p *WeightMagnitudePruner) Prune(weights []float32) ([]float32, []bool) {
	n := len(weights)
	if n == 0 {
		return nil, nil
	}

	// Step 1: Calculate absolute values
	absWeights := make([]float64, n)
	for i, w := range weights {
		absWeights[i] = math.Abs(float64(w))
	}

	// Step 2: Sort absolute values to find threshold
	sortedAbs := make([]float64, n)
	copy(sortedAbs, absWeights)
	sort.Float64s(sortedAbs)

	// Step 3: Find threshold based on target sparsity
	thresholdIndex := int(float64(n) * p.Sparsity)
	if thresholdIndex >= n {
		thresholdIndex = n - 1
	}
	threshold := sortedAbs[thresholdIndex]

	// Step 4: Create mask and prune
	prunedWeights := make([]float32, n)
	mask := make([]bool, n)
	for i, w := range weights {
		if absWeights[i] >= threshold {
			prunedWeights[i] = w
			mask[i] = true
		} else {
			prunedWeights[i] = 0 // Zero out pruned weights
			mask[i] = false
		}
	}

	return prunedWeights, mask
}

// StructuredChannelPruner implements structured pruning for convolutional layers
type StructuredChannelPruner struct {
	KeepRatio float64 // Ratio of channels to keep (0.0 to 1.0)
}

// PruneChannels prunes entire channels based on L2 norm
func (p *StructuredChannelPruner) PruneChannels(weightMatrix [][]float32) ([][]float32, []int) {
	if len(weightMatrix) == 0 {
		return nil, nil
	}

	numChannels := len(weightMatrix)
	// Step 1: Calculate L2 norm for each channel
	channelNorms := make([]float64, numChannels)
	for i, channel := range weightMatrix {
		var sumSquares float64
		for _, w := range channel {
			sumSquares += float64(w) * float64(w)
		}
		channelNorms[i] = math.Sqrt(sumSquares)
	}

	// Step 2: Sort channels by L2 norm (descending)
	type ChannelNorm struct {
		Index int
		Norm  float64
	}
	sortedNorms := make([]ChannelNorm, numChannels)
	for i, norm := range channelNorms {
		sortedNorms[i] = ChannelNorm{Index: i, Norm: norm}
	}
	sort.Slice(sortedNorms, func(i, j int) bool {
		return sortedNorms[i].Norm > sortedNorms[j].Norm
	})

	// Step 3: Keep top channels based on KeepRatio
	keepCount := int(float64(numChannels) * p.KeepRatio)
	if keepCount < 1 {
		keepCount = 1
	}

	// Step 4: Build pruned matrix and keep indices
	prunedMatrix := make([][]float32, keepCount)
	keptIndices := make([]int, keepCount)
	for i := 0; i < keepCount; i++ {
		origIdx := sortedNorms[i].Index
		prunedMatrix[i] = weightMatrix[origIdx]
		keptIndices[i] = origIdx
	}

	return prunedMatrix, keptIndices
}

// Go语言实现：知识蒸馏损失函数
package distillation

import (
	"math"
)

// SoftmaxWithTemperature applies temperature-scaled softmax
// Input: logits (raw scores), temperature (T)
// Output: probability distribution
func SoftmaxWithTemperature(logits []float32, temperature float32) []float32 {
	n := len(logits)
	if n == 0 {
		return nil
	}

	// Apply temperature scaling
	scaledLogits := make([]float64, n)
	var maxLogit float64 = -1e9
	for i, l := range logits {
		scaled := float64(l) / float64(temperature)
		scaledLogits[i] = scaled
		if scaled > maxLogit {
			maxLogit = scaled
		}
	}

	// Compute softmax
	var sumExp float64
	probs := make([]float32, n)
	for i, l := range scaledLogits {
		expVal := math.Exp(l - maxLogit) // Numerical stability
		probs[i] = float32(expVal)
		sumExp += expVal
	}

	// Normalize
	for i := range probs {
		probs[i] /= float32(sumExp)
	}

	return probs
}

// DistillationLoss computes the combined loss for knowledge distillation
// studentLogits: raw outputs from student model
// teacherLogits: raw outputs from teacher model
// hardLabels: ground truth labels (one-hot encoded)
// temperature: distillation temperature (T)
// alpha: weight for distillation loss (0.0 to 1.0)
func DistillationLoss(studentLogits, teacherLogits, hardLabels []float32, temperature float32, alpha float32) float32 {
	n := len(studentLogits)
	if n == 0 {
		return 0
	}

	// Step 1: Compute soft targets (teacher probabilities)
	teacherProbs := SoftmaxWithTemperature(teacherLogits, temperature)

	// Step 2: Compute student soft probabilities
	studentProbs := SoftmaxWithTemperature(studentLogits, temperature)

	// Step 3: Compute KL divergence between teacher and student distributions
	var klDivergence float64
	for i := 0; i < n; i++ {
		if teacherProbs[i] > 0 && studentProbs[i] > 0 {
			klDivergence += float64(teacherProbs[i]) * math.Log(float64(teacherProbs[i])/float64(studentProbs[i]))
		}
	}

	// Step 4: Compute cross-entropy loss with hard labels
	var crossEntropy float64
	for i := 0; i < n; i++ {
		if hardLabels[i] > 0 {
			crossEntropy -= float64(hardLabels[i]) * math.Log(float64(studentProbs[i]))
		}
	}

	// Step 5: Combine losses with alpha weighting
	distillationLoss := float32(klDivergence) * float32(temperature*temperature) // Scale by T^2
	hardLoss := float32(crossEntropy)
	totalLoss := alpha*distillationLoss + (1-alpha)*hardLoss

	return totalLoss
}

// DistillationTrainer orchestrates the training process
type DistillationTrainer struct {
	Temperature float32
	Alpha       float32
}

// NewDistillationTrainer creates a new trainer with specified parameters
func NewDistillationTrainer(temp, alpha float32) *DistillationTrainer {
	return &DistillationTrainer{
		Temperature: temp,
		Alpha:       alpha,
	}
}

// TrainStep performs one training step for distillation
func (dt *DistillationTrainer) TrainStep(studentLogits, teacherLogits, hardLabels []float32) float32 {
	return DistillationLoss(studentLogits, teacherLogits, hardLabels, dt.Temperature, dt.Alpha)
}

// model_loader.go - CoreML模型加载与初始化
package engine

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

// QuantizedModel represents a quantized neural network model
type QuantizedModel struct {
	Weights      [][]int8  // Quantized weight matrices
	Biases       []float32 // Bias vectors (kept in FP32 for precision)
	ScaleFactors []float32 // Per-layer quantization scales
	ZeroPoints   []int8    // Per-layer zero points (for asymmetric quantization)
	LayerConfigs []LayerConfig
}

// LayerConfig describes a single layer's configuration
type LayerConfig struct {
	LayerType    string // "attention", "ffn", "layernorm", etc.
	InputDim     int
	OutputDim    int
	NumHeads     int    // For attention layers
	HeadDim      int    // For attention layers
	Activation   string // "relu", "gelu", "silu", etc.
	SparsityMask [][]bool // Optional pruning mask
}

// ModelLoader handles loading and initializing quantized models
type ModelLoader struct {
	modelPath string
	model     *QuantizedModel
	mu        sync.Mutex
}

// NewModelLoader creates a new model loader
func NewModelLoader(path string) *ModelLoader {
	return &ModelLoader{
		modelPath: path,
	}
}

// LoadModel loads a quantized model from binary file
// File format: [num_layers:4bytes] [layer_configs] [weights_data]
func (ml *ModelLoader) LoadModel() error {
	ml.mu.Lock()
	defer ml.mu.Unlock()

	file, err := os.Open(ml.modelPath)
	if err != nil {
		return fmt.Errorf("failed to open model file: %w", err)
	}
	defer file.Close()

	// Read number of layers
	var numLayers int32
	if err := binary.Read(file, binary.LittleEndian, &numLayers); err != nil {
		return fmt.Errorf("failed to read layer count: %w", err)
	}

	ml.model = &QuantizedModel{
		Weights:      make([][]int8, numLayers),
		Biases:       make([]float32, numLayers),
		ScaleFactors: make([]float32, numLayers),
		ZeroPoints:   make([]int8, numLayers),
		LayerConfigs: make([]LayerConfig, numLayers),
	}

	// Read layer configurations and weights
	for i := int32(0); i < numLayers; i++ {
		config, err := ml.readLayerConfig(file)
		if err != nil {
			return fmt.Errorf("failed to read layer %d config: %w", i, err)
		}
		ml.model.LayerConfigs[i] = config

		// Read weights
		weightSize := config.InputDim * config.OutputDim
		weights := make([]int8, weightSize)
		if err := binary.Read(file, binary.LittleEndian, &weights); err != nil {
			return fmt.Errorf("failed to read layer %d weights: %w", i, err)
		}
		ml.model.Weights[i] = weights

		// Read biases
		var bias float32
		if err := binary.Read(file, binary.LittleEndian, &bias); err != nil {
			return fmt.Errorf("failed to read layer %d bias: %w", i, err)
		}
		ml.model.Biases[i] = bias

		// Read quantization parameters
		var scale float32
		if err := binary.Read(file, binary.LittleEndian, &scale); err != nil {
			return fmt.Errorf("failed to read layer %d scale: %w", i, err)
		}
		ml.model.ScaleFactors[i] = scale

		// Read sparsity mask if present
		if config.SparsityMask != nil {
			maskSize := config.InputDim * config.OutputDim
			mask := make([]bool, maskSize)
			for j := 0; j < maskSize; j++ {
				var val byte
				if err := binary.Read(file, binary.LittleEndian, &val); err != nil {
					return fmt.Errorf("failed to read layer %d mask: %w", i, err)
				}
				mask[j] = val != 0
			}
			// Convert to 2D
			ml.model.LayerConfigs[i].SparsityMask = make([][]bool, config.OutputDim)
			for j := 0; j < config.OutputDim; j++ {
				ml.model.LayerConfigs[i].SparsityMask[j] = mask[j*config.InputDim : (j+1)*config.InputDim]
			}
		}
	}

	return nil
}

// readLayerConfig reads a single layer configuration from file
func (ml *ModelLoader) readLayerConfig(file *os.File) (LayerConfig, error) {
	var config LayerConfig

	// Read layer type length and string
	var typeLen int32
	if err := binary.Read(file, binary.LittleEndian, &typeLen); err != nil {
		return config, err
	}
	typeBytes := make([]byte, typeLen)
	if err := binary.Read(file, binary.LittleEndian, &typeBytes); err != nil {
		return config, err
	}
	config.LayerType = string(typeBytes)

	// Read dimensions
	if err := binary.Read(file, binary.LittleEndian, &config.InputDim); err != nil {
		return config, err
	}
	if err := binary.Read(file, binary.LittleEndian, &config.OutputDim); err != nil {
		return config, err
	}

	// Read attention-specific parameters
	if config.LayerType == "attention" {
		if err := binary.Read(file, binary.LittleEndian, &config.NumHeads); err != nil {
			return config, err
		}
		if err := binary.Read(file, binary.LittleEndian, &config.HeadDim); err != nil {
			return config, err
		}
	}

	// Read activation function
	var actLen int32
	if err := binary.Read(file, binary.LittleEndian, &actLen); err != nil {
		return config, err
	}
	actBytes := make([]byte, actLen)
	if err := binary.Read(file, binary.LittleEndian, &actBytes); err != nil {
		return config, err
	}
	config.Activation = string(actBytes)

	return config, nil
}

// GetModel returns the loaded model
func (ml *ModelLoader) GetModel() *QuantizedModel {
	return ml.model
}

// inference_engine.go - 端侧推理引擎核心
package engine

import (
	"math"
	"sync"
)

// InferenceEngine handles model inference with optimizations
type InferenceEngine struct {
	model      *QuantizedModel
	kvCache    *KVCache
	memoryPool *MemoryPool
	ops        *OperatorRegistry
}

// NewInferenceEngine creates a new inference engine
func NewInferenceEngine(model *QuantizedModel) *InferenceEngine {
	return &InferenceEngine{
		model:      model,
		kvCache:    NewKVCache(2048, 32), // Max sequence length, num layers
		memoryPool: NewMemoryPool(1024 * 1024 * 10), // 10MB pool
		ops:        NewOperatorRegistry(),
	}
}

// Forward performs a single forward pass through the model
// input: token embeddings [batch_size, hidden_dim]
// Returns: output logits [batch_size, vocab_size]
func (ie *InferenceEngine) Forward(input [][]float32) ([][]float32, error) {
	batchSize := len(input)
	if batchSize == 0 {
		return nil, nil
	}

	hiddenDim := len(input[0])
	current := ie.memoryPool.AllocateMatrix(batchSize, hiddenDim)
	copy(current, input)

	// Process each layer
	for layerIdx, config := range ie.model.LayerConfigs {
		switch config.LayerType {
		case "attention":
			// Multi-head attention with KV cache
			q, k, v := ie.computeQKV(current, layerIdx)
			ie.kvCache.Append(layerIdx, k, v)
			attnOutput := ie.scaledDotProductAttention(q, ie.kvCache.GetK(layerIdx), ie.kvCache.GetV(layerIdx), config)
			current = ie.addResidual(current, attnOutput)
			current = ie.layerNorm(current)

		case "ffn":
			// Feed-forward network with activation
			ffnOutput := ie.feedForward(current, layerIdx)
			current = ie.addResidual(current, ffnOutput)
			current = ie.layerNorm(current)

		case "embedding":
			// Token embedding lookup (handled by caller)
			continue

		case "output":
			// Final linear projection to vocabulary
			logits := ie.linearProjection(current, layerIdx)
			ie.memoryPool.FreeMatrix(current)
			return logits, nil
		}
	}

	return current, nil
}

// scaledDotProductAttention implements efficient attention with KV cache
func (ie *InferenceEngine) scaledDotProductAttention(q, k, v [][]float32, config LayerConfig) [][]float32 {
	batchSize := len(q)
	seqLen := len(k[0]) // KV cache length
	headDim := config.HeadDim
	numHeads := config.NumHeads

	// Allocate output
	output := ie.memoryPool.AllocateMatrix(batchSize, numHeads*headDim)

	// Quantized attention computation
	for b := 0; b < batchSize; b++ {
		// Compute attention scores: Q * K^T / sqrt(d_k)
		scores := make([][]float32, numHeads)
		for h := 0; h < numHeads; h++ {
			scores[h] = make([]float32, seqLen)
			headOffset := h * headDim
			for s := 0; s < seqLen; s++ {
				var score float32
				for d := 0; d < headDim; d++ {
					score += q[b][headOffset+d] * k[s][headOffset+d]
				}
				score /= float32(math.Sqrt(float64(headDim)))
				scores[h][s] = score
			}
		}

		// Softmax over sequence dimension
		for h := 0; h < numHeads; h++ {
			var maxScore float32 = -1e9
			for s := 0; s < seqLen; s++ {
				if scores[h][s] > maxScore {
					maxScore = scores[h][s]
				}
			}
			var sumExp float32
			for s := 0; s < seqLen; s++ {
				scores[h][s] = float32(math.Exp(float64(scores[h][s] - maxScore)))
				sumExp += scores[h][s]
			}
			for s := 0; s < seqLen; s++ {
				scores[h][s] /= sumExp
			}
		}

		// Weighted sum of values
		for h := 0; h < numHeads; h++ {
			headOffset := h * headDim
			for d := 0; d < headDim; d++ {
				var weightedSum float32
				for s := 0; s < seqLen; s++ {
					weightedSum += scores[h][s] * v[s][headOffset+d]
				}
				output[b][headOffset+d] = weightedSum
			}
		}
	}

	return output
}

// feedForward implements quantized feed-forward network
func (ie *InferenceEngine) feedForward(input [][]float32, layerIdx int) [][]float32 {
	config := ie.model.LayerConfigs[layerIdx]
	weights := ie.model.Weights[layerIdx]
	scale := ie.model.ScaleFactors[layerIdx]
	bias := ie.model.Biases[layerIdx]
	sparsityMask := config.SparsityMask

	batchSize := len(input)
	inputDim := config.InputDim
	outputDim := config.OutputDim

	output := ie.memoryPool.AllocateMatrix(batchSize, outputDim)

	for b := 0; b < batchSize; b++ {
		for o := 0; o < outputDim; o++ {
			var sum float32
			weightOffset := o * inputDim

			// Check sparsity mask for this output neuron
			if sparsityMask != nil && sparsityMask[o] != nil {
				for i := 0; i < inputDim; i++ {
					if sparsityMask[o][i] {
						// Dequantize weight and accumulate
						w := float32(weights[weightOffset+i]) / scale
						sum += w * input[b][i]
					}
				}
			} else {
				// Dense computation
				for i := 0; i < inputDim; i++ {
					w := float32(weights[weightOffset+i]) / scale
					sum += w * input[b][i]
				}
			}

			// Apply activation function
			switch config.Activation {
			case "gelu":
				output[b][o] = gelu(sum + bias)
			case "silu":
				output[b][o] = silu(sum + bias)
			case "relu":
				if sum+bias > 0 {
					output[b][o] = sum + bias
				} else {
					output[b][o] = 0
				}
			default:
				output[b][o] = sum + bias
			}
		}
	}

	return output
}

// Helper activation functions
func gelu(x float32) float32 {
	return 0.5 * x * (1 + float32(math.Tanh(float64(math.Sqrt(2/math.Pi)*(x+0.044715*x*x*x)))))
}

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

// Memory pool for efficient allocation
type MemoryPool struct {
	buffer []float32
	offset int
	mu     sync.Mutex
}

func NewMemoryPool(size int) *MemoryPool {
	return &MemoryPool{
		buffer: make([]float32, size),
	}
}

func (mp *MemoryPool) AllocateMatrix(rows, cols int) [][]float32 {
	mp.mu.Lock()
	defer mp.mu.Unlock()

	size := rows * cols
	if mp.offset+size > len(mp.buffer) {
		// Expand buffer
		newBuffer := make([]float32, len(mp.buffer)*2)
		copy(newBuffer, mp.buffer)
		mp.buffer = newBuffer
	}

	matrix := make([][]float32, rows)
	for i := 0; i < rows; i++ {
		matrix[i] = mp.buffer[mp.offset+i*cols : mp.offset+(i+1)*cols]
	}
	mp.offset += size
	return matrix
}

func (mp *MemoryPool) FreeMatrix(matrix [][]float32) {
	// In a real implementation, we'd track allocations and reuse
	// For simplicity, we just reset the offset periodically
}

// KV Cache for efficient autoregressive generation
type KVCache struct {
	kCache [][][]float32
	vCache [][][]float32
	maxLen int
	numLayers int
}

func NewKVCache(maxLen, numLayers int) *KVCache {
	return &KVCache{
		kCache: make([][][]float32, numLayers),
		vCache: make([][][]float32, numLayers),
		maxLen: maxLen,
		numLayers: numLayers,
	}
}

func (kc *KVCache) Append(layerIdx int, k, v [][]float32) {
	kc.kCache[layerIdx] = append(kc.kCache[layerIdx], k...)
	kc.vCache[layerIdx] = append(kc.vCache[layerIdx], v...)
	
	// Trim if exceeds max length
	if len(kc.kCache[layerIdx]) > kc.maxLen {
		kc.kCache[layerIdx] = kc.kCache[layerIdx][len(kc.kCache[layerIdx])-kc.maxLen:]
		kc.vCache[layerIdx] = kc.vCache[layerIdx][len(kc.vCache[layerIdx])-kc.maxLen:]
	}
}

func (kc *KVCache) GetK(layerIdx int) [][]float32 {
	return kc.kCache[layerIdx]
}

func (kc *KVCache) GetV(layerIdx int) [][]float32 {
	return kc.vCache[layerIdx]
}

// Go实现：LayerNorm + Attention融合
package fusion

import "math"

// FusedLayerNormAttention combines layer normalization and attention computation
// to reduce memory bandwidth usage
func FusedLayerNormAttention(input [][]float32, gamma, beta []float32, 
	numHeads, headDim int) [][]float32 {
	
	batchSize := len(input)
	hiddenDim := len(input[0])
	
	// Step 1: Compute layer norm statistics in-place
	// This avoids writing intermediate results to memory
	for b := 0; b < batchSize; b++ {
		var mean, variance float32
		
		// Compute mean
		for _, val := range input[b] {
			mean += val
		}
		mean /= float32(hiddenDim)
		
		// Compute variance
		for _, val := range input[b] {
			diff := val - mean
			variance += diff * diff
		}
		variance /= float32(hiddenDim)
		
		// Normalize and scale
		stdDev := float32(math.Sqrt(float64(variance + 1e-5)))
		for i := range input[b] {
			input[b][i] = gamma[i]*(input[b][i]-mean)/stdDev + beta[i]
		}
		
		// Step 2: Immediately use normalized values for attention QKV projection
		// This is where the fusion happens - we don't write to a separate buffer
		// but directly compute attention scores
		
		// Simplified attention computation
		for h := 0; h < numHeads; h++ {
			headStart

![](/images/blog/bc89dcf25fa2c67527546f481f11f6a6-202606101450.png)