推理优化技术

package quantization

import (
    "math"
)

// QuantizationType 量化类型
type QuantizationType int

const (
    QuantInt8 QuantizationType = iota
    QuantInt4
    QuantFP8
    QuantFP4
)

// QuantizationConfig 量化配置
type QuantizationConfig struct {
    Type          QuantizationType
    Symmetric     bool              // 对称量化
    PerChannel    bool              // 逐通道量化
    GroupSize     int               // 分组量化大小
    CalibrationMethod string        // 校准方法: minmax, percentile, entropy
}

// Quantizer 量化器
type Quantizer struct {
    config *QuantizationConfig
}

// QuantizedTensor 量化后的张量
type QuantizedTensor struct {
    Data      []int8           // 量化数据 (int8 示例)
    Scale     []float32        // 缩放因子
    ZeroPoint []int8           // 零点
    Shape     []int
    GroupSize int
}

// QuantizeSymmetric 对称量化
// x_q = round(x / scale)
// x = x_q * scale
func (q *Quantizer) QuantizeSymmetric(tensor []float32, bits int) *QuantizedTensor {
    maxVal := q.findAbsMax(tensor)
    qMax := float32(math.Pow(2, float64(bits-1)) - 1)

    scale := maxVal / qMax

    quantized := make([]int8, len(tensor))
    for i, v := range tensor {
        qVal := v / scale
        // 截断到范围 [-qMax, qMax]
        qVal = float32(math.Max(float64(-qMax), math.Min(float64(qMax), float64(qVal))))
        quantized[i] = int8(math.Round(float64(qVal)))
    }

    return &QuantizedTensor{
        Data:      quantized,
        Scale:     []float32{scale},
        ZeroPoint: []int8{0},
        Shape:     []int{len(tensor)},
    }
}

// QuantizeAsymmetric 非对称量化
// x_q = round(x / scale) + zero_point
// x = (x_q - zero_point) * scale
func (q *Quantizer) QuantizeAsymmetric(tensor []float32, bits int) *QuantizedTensor {
    minVal, maxVal := q.findMinMax(tensor)
    qMin := float32(0)
    qMax := float32(math.Pow(2, float64(bits)) - 1)

    scale := (maxVal - minVal) / (qMax - qMin)
    zeroPoint := int8(math.Round(float64(-minVal / scale)))

    quantized := make([]int8, len(tensor))
    for i, v := range tensor {
        qVal := v/scale + float32(zeroPoint)
        qVal = float32(math.Max(float64(qMin), math.Min(float64(qMax), float64(qVal))))
        quantized[i] = int8(math.Round(float64(qVal)))
    }

    return &QuantizedTensor{
        Data:      quantized,
        Scale:     []float32{scale},
        ZeroPoint: []int8{zeroPoint},
    }
}

// QuantizePerGroup 分组量化
// 每 group_size 个元素使用独立的 scale 和 zero_point
func (q *Quantizer) QuantizePerGroup(
    tensor []float32,
    groupSize int,
    bits int,
) *QuantizedTensor {

    numGroups := (len(tensor) + groupSize - 1) / groupSize
    scales := make([]float32, numGroups)
    zeroPoints := make([]int8, numGroups)
    quantized := make([]int8, len(tensor))

    qMax := float32(math.Pow(2, float64(bits-1)) - 1)

    for g := 0; g < numGroups; g++ {
        start := g * groupSize
        end := min(start+groupSize, len(tensor))
        group := tensor[start:end]

        // 计算该组的量化参数
        maxVal := q.findAbsMax(group)
        scale := maxVal / qMax
        scales[g] = scale

        // 量化该组
        for i, v := range group {
            qVal := v / scale
            qVal = float32(math.Max(float64(-qMax), math.Min(float64(qMax), float64(qVal))))
            quantized[start+i] = int8(math.Round(float64(qVal)))
        }
    }

    return &QuantizedTensor{
        Data:      quantized,
        Scale:     scales,
        ZeroPoint: zeroPoints,
        GroupSize: groupSize,
    }
}

// Dequantize 反量化
func (qt *QuantizedTensor) Dequantize() []float32 {
    result := make([]float32, len(qt.Data))

    if qt.GroupSize > 0 {
        // 分组量化反量化
        for i, qVal := range qt.Data {
            groupIdx := i / qt.GroupSize
            scale := qt.Scale[groupIdx]
            zp := qt.ZeroPoint[groupIdx]
            result[i] = float32(qVal-zp) * scale
        }
    } else {
        // 全局量化反量化
        scale := qt.Scale[0]
        zp := qt.ZeroPoint[0]
        for i, qVal := range qt.Data {
            result[i] = float32(qVal-zp) * scale
        }
    }

    return result
}

func (q *Quantizer) findAbsMax(tensor []float32) float32 {
    maxVal := float32(0)
    for _, v := range tensor {
        absV := float32(math.Abs(float64(v)))
        if absV > maxVal {
            maxVal = absV
        }
    }
    return maxVal
}

func (q *Quantizer) findMinMax(tensor []float32) (float32, float32) {
    minVal := tensor[0]
    maxVal := tensor[0]
    for _, v := range tensor {
        if v < minVal {
            minVal = v
        }
        if v > maxVal {
            maxVal = v
        }
    }
    return minVal, maxVal
}

package quantization

import (
    "math"
    "sync"
)

// GPTQQuantizer GPTQ 量化器
type GPTQQuantizer struct {
    bits       int
    groupSize  int
    symmetric  bool
    actOrder   bool     // 激活排序优化
    dampingPct float64  // 阻尼系数
}

func NewGPTQQuantizer(config *GPTQConfig) *GPTQQuantizer {
    return &GPTQQuantizer{
        bits:       config.Bits,
        groupSize:  config.GroupSize,
        symmetric:  config.Symmetric,
        actOrder:   config.ActOrder,
        dampingPct: config.DampingPct,
    }
}

// QuantizeLayer 量化单层权重
// 基于 OBS (Optimal Brain Surgeon) 框架
func (gq *GPTQQuantizer) QuantizeLayer(
    weight [][]float32,  // [out_features, in_features]
    hessian [][]float32, // Hessian 矩阵 [in_features, in_features]
) (*GPTQLayerResult, error) {

    outFeatures := len(weight)
    inFeatures := len(weight[0])

    // 复制权重（避免修改原始数据）
    W := make([][]float32, outFeatures)
    for i := range W {
        W[i] = make([]float32, inFeatures)
        copy(W[i], weight[i])
    }

    // 添加阻尼项到 Hessian 对角线
    H := gq.addDamping(hessian)

    // Cholesky 分解
    L, err := gq.choleskyDecomposition(H)
    if err != nil {
        return nil, err
    }

    // 计算 Hessian 逆
    Hinv := gq.invertLowerTriangular(L)

    // 量化参数存储
    numGroups := (inFeatures + gq.groupSize - 1) / gq.groupSize
    scales := make([][]float32, outFeatures)
    zeros := make([][]int32, outFeatures)
    quantizedWeight := make([][]int8, outFeatures)

    for i := range outFeatures {
        scales[i] = make([]float32, numGroups)
        zeros[i] = make([]int32, numGroups)
        quantizedWeight[i] = make([]int8, inFeatures)
    }

    // 获取量化顺序
    order := gq.getQuantOrder(H)

    // 逐列量化
    blockSize := 128 // 分块处理减少累积误差
    for blockStart := 0; blockStart < inFeatures; blockStart += blockSize {
        blockEnd := min(blockStart+blockSize, inFeatures)

        for j := blockStart; j < blockEnd; j++ {
            col := order[j]
            groupIdx := col / gq.groupSize

            for i := 0; i < outFeatures; i++ {
                // 量化当前元素
                w := W[i][col]

                // 计算该组的量化参数（如果是组的第一个元素）
                if col%gq.groupSize == 0 {
                    groupStart := col
                    groupEnd := min(col+gq.groupSize, inFeatures)
                    scale, zero := gq.computeGroupParams(W[i][groupStart:groupEnd])
                    scales[i][groupIdx] = scale
                    zeros[i][groupIdx] = zero
                }

                scale := scales[i][groupIdx]
                zero := zeros[i][groupIdx]

                // 量化
                qVal := gq.quantize(w, scale, zero)
                quantizedWeight[i][col] = qVal

                // 反量化
                wHat := gq.dequantize(qVal, scale, zero)

                // 计算量化误差
                err := w - wHat

                // 误差补偿（OBS 核心）
                // 将量化误差分配到后续未量化的列
                if j < blockEnd-1 {
                    errWeight := err / float32(Hinv[col][col])
                    for k := j + 1; k < blockEnd; k++ {
                        nextCol := order[k]
                        W[i][nextCol] -= errWeight * float32(Hinv[nextCol][col])
                    }
                }
            }
        }
    }

    return &GPTQLayerResult{
        QuantizedWeight: quantizedWeight,
        Scales:          scales,
        Zeros:           zeros,
        GroupSize:       gq.groupSize,
    }, nil
}

// ComputeHessian 计算 Hessian 矩阵
// H = X^T * X（对于线性层）
func (gq *GPTQQuantizer) ComputeHessian(
    calibrationData [][]float32, // 校准数据 [num_samples, in_features]
) [][]float32 {
    numSamples := len(calibrationData)
    inFeatures := len(calibrationData[0])

    H := make([][]float32, inFeatures)
    for i := range H {
        H[i] = make([]float32, inFeatures)
    }

    // 并行计算 Hessian
    var wg sync.WaitGroup
    numWorkers := 8
    rowsPerWorker := (inFeatures + numWorkers - 1) / numWorkers

    for w := 0; w < numWorkers; w++ {
        wg.Add(1)
        go func(workerID int) {
            defer wg.Done()
            startRow := workerID * rowsPerWorker
            endRow := min(startRow+rowsPerWorker, inFeatures)

            for i := startRow; i < endRow; i++ {
                for j := i; j < inFeatures; j++ {
                    var sum float32
                    for s := 0; s < numSamples; s++ {
                        sum += calibrationData[s][i] * calibrationData[s][j]
                    }
                    H[i][j] = sum / float32(numSamples)
                    if i != j {
                        H[j][i] = H[i][j] // 对称矩阵
                    }
                }
            }
        }(w)
    }
    wg.Wait()

    return H
}

func (gq *GPTQQuantizer) addDamping(H [][]float32) [][]float32 {
    n := len(H)
    result := make([][]float32, n)
    for i := range result {
        result[i] = make([]float32, n)
        copy(result[i], H[i])
    }

    // 计算对角线平均值
    var diagSum float32
    for i := 0; i < n; i++ {
        diagSum += result[i][i]
    }
    diagMean := diagSum / float32(n)

    // 添加阻尼
    damping := float32(gq.dampingPct) * diagMean
    for i := 0; i < n; i++ {
        result[i][i] += damping
    }

    return result
}

func (gq *GPTQQuantizer) getQuantOrder(H [][]float32) []int {
    n := len(H)
    order := make([]int, n)

    if gq.actOrder {
        // 按 Hessian 对角线元素排序（激活量大的先量化）
        diagValues := make([]float32, n)
        for i := 0; i < n; i++ {
            diagValues[i] = H[i][i]
        }

        // 获取排序索引
        indices := make([]int, n)
        for i := range indices {
            indices[i] = i
        }

        // 降序排序
        for i := 0; i < n-1; i++ {
            for j := i + 1; j < n; j++ {
                if diagValues[indices[i]] < diagValues[indices[j]] {
                    indices[i], indices[j] = indices[j], indices[i]
                }
            }
        }
        copy(order, indices)
    } else {
        // 默认顺序
        for i := range order {
            order[i] = i
        }
    }

    return order
}

func (gq *GPTQQuantizer) computeGroupParams(group []float32) (float32, int32) {
    maxVal := float32(0)
    for _, v := range group {
        absV := float32(math.Abs(float64(v)))
        if absV > maxVal {
            maxVal = absV
        }
    }

    qMax := float32(math.Pow(2, float64(gq.bits-1)) - 1)
    scale := maxVal / qMax
    if scale == 0 {
        scale = 1
    }

    return scale, 0
}

func (gq *GPTQQuantizer) quantize(val, scale float32, zero int32) int8 {
    qMax := float32(math.Pow(2, float64(gq.bits-1)) - 1)
    qVal := val/scale + float32(zero)
    qVal = float32(math.Max(float64(-qMax), math.Min(float64(qMax), float64(qVal))))
    return int8(math.Round(float64(qVal)))
}

func (gq *GPTQQuantizer) dequantize(qVal int8, scale float32, zero int32) float32 {
    return float32(int32(qVal)-zero) * scale
}

type GPTQConfig struct {
    Bits       int
    GroupSize  int
    Symmetric  bool
    ActOrder   bool
    DampingPct float64
}

type GPTQLayerResult struct {
    QuantizedWeight [][]int8
    Scales          [][]float32
    Zeros           [][]int32
    GroupSize       int
}

package quantization

import (
    "math"
)

// AWQQuantizer AWQ 量化器
type AWQQuantizer struct {
    bits       int
    groupSize  int
    searchScale bool
}

// AWQConfig AWQ 配置
type AWQConfig struct {
    Bits        int
    GroupSize   int
    SearchScale bool
    NumSamples  int
}

func NewAWQQuantizer(config *AWQConfig) *AWQQuantizer {
    return &AWQQuantizer{
        bits:        config.Bits,
        groupSize:   config.GroupSize,
        searchScale: config.SearchScale,
    }
}

// QuantizeLayer AWQ 量化层
func (aq *AWQQuantizer) QuantizeLayer(
    weight [][]float32,      // [out_features, in_features]
    activations [][]float32, // 校准激活 [num_samples, in_features]
) (*AWQLayerResult, error) {

    outFeatures := len(weight)
    inFeatures := len(weight[0])

    // 1. 计算激活通道重要性
    channelImportance := aq.computeChannelImportance(activations)

    // 2. 搜索最优缩放因子
    var scales []float32
    if aq.searchScale {
        scales = aq.searchOptimalScales(weight, channelImportance)
    } else {
        // 使用激活幅度作为缩放因子
        scales = channelImportance
    }

    // 3. 对权重应用缩放
    scaledWeight := aq.applyScales(weight, scales)

    // 4. 量化缩放后的权重
    numGroups := (inFeatures + aq.groupSize - 1) / aq.groupSize
    quantizedWeight := make([][]int8, outFeatures)
    qScales := make([][]float32, outFeatures)
    qZeros := make([][]int32, outFeatures)

    for i := 0; i < outFeatures; i++ {
        quantizedWeight[i] = make([]int8, inFeatures)
        qScales[i] = make([]float32, numGroups)
        qZeros[i] = make([]int32, numGroups)

        for g := 0; g < numGroups; g++ {
            start := g * aq.groupSize
            end := min(start+aq.groupSize, inFeatures)

            // 计算组内量化参数
            group := scaledWeight[i][start:end]
            scale, zero := aq.computeQuantParams(group)
            qScales[i][g] = scale
            qZeros[i][g] = zero

            // 量化
            for j := start; j < end; j++ {
                quantizedWeight[i][j] = aq.quantize(scaledWeight[i][j], scale, zero)
            }
        }
    }

    return &AWQLayerResult{
        QuantizedWeight:  quantizedWeight,
        Scales:           qScales,
        Zeros:            qZeros,
        ChannelScales:    scales,
        GroupSize:        aq.groupSize,
    }, nil
}

// computeChannelImportance 计算通道重要性
func (aq *AWQQuantizer) computeChannelImportance(activations [][]float32) []float32 {
    numSamples := len(activations)
    inFeatures := len(activations[0])

    importance := make([]float32, inFeatures)

    for c := 0; c < inFeatures; c++ {
        var sum float32
        for s := 0; s < numSamples; s++ {
            sum += float32(math.Abs(float64(activations[s][c])))
        }
        importance[c] = sum / float32(numSamples)
    }

    // 归一化
    maxImp := importance[0]
    for _, v := range importance {
        if v > maxImp {
            maxImp = v
        }
    }
    if maxImp > 0 {
        for i := range importance {
            importance[i] /= maxImp
        }
    }

    return importance
}

// searchOptimalScales 搜索最优缩放因子
func (aq *AWQQuantizer) searchOptimalScales(
    weight [][]float32,
    channelImportance []float32,
) []float32 {

    inFeatures := len(weight[0])
    scales := make([]float32, inFeatures)

    // 网格搜索缩放因子
    searchRange := []float64{0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0}

    for c := 0; c < inFeatures; c++ {
        bestScale := float32(1.0)
        bestError := float32(math.MaxFloat32)

        baseScale := channelImportance[c]
        if baseScale < 0.01 {
            baseScale = 0.01
        }

        for _, ratio := range searchRange {
            scale := float32(ratio) * baseScale

            // 计算该缩放因子下的量化误差
            error := aq.computeQuantError(weight, c, scale)
            if error < bestError {
                bestError = error
                bestScale = scale
            }
        }

        scales[c] = bestScale
    }

    return scales
}

func (aq *AWQQuantizer) computeQuantError(weight [][]float32, channel int, scale float32) float32 {
    var totalError float32

    for i := range weight {
        w := weight[i][channel]
        scaledW := w * scale

        // 模拟量化
        qMax := float32(math.Pow(2, float64(aq.bits-1)) - 1)
        qScale := float32(math.Abs(float64(scaledW))) / qMax
        if qScale == 0 {
            qScale = 1
        }

        qVal := scaledW / qScale
        qVal = float32(math.Max(-float64(qMax), math.Min(float64(qMax), float64(qVal))))
        qVal = float32(math.Round(float64(qVal)))

        // 反量化
        deqVal := qVal * qScale / scale

        // 累计误差
        totalError += float32(math.Abs(float64(w - deqVal)))
    }

    return totalError
}

func (aq *AWQQuantizer) applyScales(weight [][]float32, scales []float32) [][]float32 {
    outFeatures := len(weight)
    inFeatures := len(weight[0])

    scaled := make([][]float32, outFeatures)
    for i := range scaled {
        scaled[i] = make([]float32, inFeatures)
        for j := 0; j < inFeatures; j++ {
            scaled[i][j] = weight[i][j] * scales[j]
        }
    }

    return scaled
}

func (aq *AWQQuantizer) computeQuantParams(group []float32) (float32, int32) {
    maxVal := float32(0)
    for _, v := range group {
        absV := float32(math.Abs(float64(v)))
        if absV > maxVal {
            maxVal = absV
        }
    }

    qMax := float32(math.Pow(2, float64(aq.bits-1)) - 1)
    scale := maxVal / qMax
    if scale == 0 {
        scale = 1
    }

    return scale, 0
}

func (aq *AWQQuantizer) quantize(val, scale float32, zero int32) int8 {
    qMax := float32(math.Pow(2, float64(aq.bits-1)) - 1)
    qVal := val/scale + float32(zero)
    qVal = float32(math.Max(float64(-qMax), math.Min(float64(qMax), float64(qVal))))
    return int8(math.Round(float64(qVal)))
}

type AWQLayerResult struct {
    QuantizedWeight [][]int8
    Scales          [][]float32
    Zeros           [][]int32
    ChannelScales   []float32
    GroupSize       int
}

package optimization

import (
    "fmt"
)

// GraphOptimizer 计算图优化器
type GraphOptimizer struct {
    graph       *ComputeGraph
    fusionRules []FusionRule
}

// ComputeGraph 计算图
type ComputeGraph struct {
    Nodes map[string]*GraphNode
    Edges []*GraphEdge
}

// GraphNode 图节点
type GraphNode struct {
    ID       string
    OpType   string
    Inputs   []string
    Outputs  []string
    Attrs    map[string]interface{}
}

// GraphEdge 图边
type GraphEdge struct {
    Source   string
    Target   string
    TensorID string
}

// FusionRule 融合规则
type FusionRule interface {
    Match(graph *ComputeGraph, nodeID string) []string
    Fuse(graph *ComputeGraph, nodeIDs []string) *GraphNode
    Name() string
}

// LinearGELUFusion Linear + GELU 融合
type LinearGELUFusion struct{}

func (f *LinearGELUFusion) Name() string {
    return "LinearGELUFusion"
}

func (f *LinearGELUFusion) Match(graph *ComputeGraph, nodeID string) []string {
    node := graph.Nodes[nodeID]
    if node.OpType != "Linear" {
        return nil
    }

    // 检查输出是否连接到 GELU
    for _, edge := range graph.Edges {
        if edge.Source == nodeID {
            targetNode := graph.Nodes[edge.Target]
            if targetNode.OpType == "GELU" {
                return []string{nodeID, edge.Target}
            }
        }
    }

    return nil
}

func (f *LinearGELUFusion) Fuse(graph *ComputeGraph, nodeIDs []string) *GraphNode {
    linearNode := graph.Nodes[nodeIDs[0]]
    geluNode := graph.Nodes[nodeIDs[1]]

    return &GraphNode{
        ID:      fmt.Sprintf("fused_%s_%s", linearNode.ID, geluNode.ID),
        OpType:  "LinearGELU",
        Inputs:  linearNode.Inputs,
        Outputs: geluNode.Outputs,
        Attrs: map[string]interface{}{
            "weight": linearNode.Attrs["weight"],
            "bias":   linearNode.Attrs["bias"],
        },
    }
}

// MultiHeadAttentionFusion 多头注意力融合
type MultiHeadAttentionFusion struct{}

func (f *MultiHeadAttentionFusion) Name() string {
    return "MultiHeadAttentionFusion"
}

func (f *MultiHeadAttentionFusion) Match(graph *ComputeGraph, nodeID string) []string {
    // 匹配 QKV projection -> Attention -> Output projection 模式
    node := graph.Nodes[nodeID]

    // 检查是否是 Q 投影
    if node.OpType != "Linear" || !isQProjection(node) {
        return nil
    }

    // 找到关联的 K, V 投影和后续操作
    kNode := findKProjection(graph, node)
    vNode := findVProjection(graph, node)

    if kNode == nil || vNode == nil {
        return nil
    }

    // 找到 attention 和 output 投影
    attnNode := findAttentionNode(graph, node, kNode, vNode)
    outNode := findOutputProjection(graph, attnNode)

    if attnNode == nil || outNode == nil {
        return nil
    }

    return []string{node.ID, kNode.ID, vNode.ID, attnNode.ID, outNode.ID}
}

func (f *MultiHeadAttentionFusion) Fuse(graph *ComputeGraph, nodeIDs []string) *GraphNode {
    qNode := graph.Nodes[nodeIDs[0]]

    return &GraphNode{
        ID:     fmt.Sprintf("fused_mha_%s", nodeIDs[0]),
        OpType: "FusedMultiHeadAttention",
        Inputs: qNode.Inputs,
        Outputs: []string{"mha_output"},
        Attrs: map[string]interface{}{
            "num_heads": extractNumHeads(graph.Nodes[nodeIDs[3]]),
            "qkv_fused": true,
        },
    }
}

// RMSNormFusion RMSNorm 融合
type RMSNormFusion struct{}

func (f *RMSNormFusion) Name() string {
    return "RMSNormFusion"
}

func (f *RMSNormFusion) Match(graph *ComputeGraph, nodeID string) []string {
    // 匹配 x -> pow(2) -> mean -> rsqrt -> mul -> mul(weight) 模式
    node := graph.Nodes[nodeID]

    if node.OpType != "Pow" {
        return nil
    }

    // 检查指数是否为 2
    if exp, ok := node.Attrs["exponent"].(float64); !ok || exp != 2 {
        return nil
    }

    // 查找后续节点链
    meanNode := findNextOp(graph, nodeID, "Mean")
    if meanNode == nil {
        return nil
    }

    rsqrtNode := findNextOp(graph, meanNode.ID, "Rsqrt")
    if rsqrtNode == nil {
        return nil
    }

    mul1Node := findNextOp(graph, rsqrtNode.ID, "Mul")
    if mul1Node == nil {
        return nil
    }

    mul2Node := findNextOp(graph, mul1Node.ID, "Mul")
    if mul2Node == nil {
        return nil
    }

    return []string{node.ID, meanNode.ID, rsqrtNode.ID, mul1Node.ID, mul2Node.ID}
}

func (f *RMSNormFusion) Fuse(graph *ComputeGraph, nodeIDs []string) *GraphNode {
    powNode := graph.Nodes[nodeIDs[0]]
    mul2Node := graph.Nodes[nodeIDs[4]]

    return &GraphNode{
        ID:      fmt.Sprintf("fused_rmsnorm_%s", nodeIDs[0]),
        OpType:  "FusedRMSNorm",
        Inputs:  powNode.Inputs,
        Outputs: mul2Node.Outputs,
        Attrs: map[string]interface{}{
            "eps":    1e-6,
            "weight": mul2Node.Attrs["other"],
        },
    }
}

// Optimize 执行图优化
func (opt *GraphOptimizer) Optimize() *ComputeGraph {
    optimizedGraph := opt.graph.Clone()

    changed := true
    for changed {
        changed = false
        for nodeID := range optimizedGraph.Nodes {
            for _, rule := range opt.fusionRules {
                matchedNodes := rule.Match(optimizedGraph, nodeID)
                if matchedNodes != nil {
                    fusedNode := rule.Fuse(optimizedGraph, matchedNodes)
                    opt.applyFusion(optimizedGraph, matchedNodes, fusedNode)
                    changed = true
                    break
                }
            }
            if changed {
                break
            }
        }
    }

    return optimizedGraph
}

func (opt *GraphOptimizer) applyFusion(
    graph *ComputeGraph,
    oldNodes []string,
    newNode *GraphNode,
) {
    // 添加新节点
    graph.Nodes[newNode.ID] = newNode

    // 更新边
    newEdges := make([]*GraphEdge, 0)
    for _, edge := range graph.Edges {
        // 更新指向旧节点的边
        sourceInOld := contains(oldNodes, edge.Source)
        targetInOld := contains(oldNodes, edge.Target)

        if sourceInOld && targetInOld {
            // 内部边，删除
            continue
        } else if sourceInOld {
            // 输出边，更新 source
            edge.Source = newNode.ID
        } else if targetInOld {
            // 输入边，更新 target
            edge.Target = newNode.ID
        }
        newEdges = append(newEdges, edge)
    }
    graph.Edges = newEdges

    // 删除旧节点
    for _, nodeID := range oldNodes {
        delete(graph.Nodes, nodeID)
    }
}

// 辅助函数
func contains(slice []string, item string) bool {
    for _, s := range slice {
        if s == item {
            return true
        }
    }
    return false
}

func findNextOp(graph *ComputeGraph, nodeID, opType string) *GraphNode {
    for _, edge := range graph.Edges {
        if edge.Source == nodeID {
            targetNode := graph.Nodes[edge.Target]
            if targetNode.OpType == opType {
                return targetNode
            }
        }
    }
    return nil
}

func isQProjection(node *GraphNode) bool {
    // 根据命名或属性判断
    return true // 简化实现
}

func findKProjection(graph *ComputeGraph, qNode *GraphNode) *GraphNode {
    return nil // 简化实现
}

func findVProjection(graph *ComputeGraph, qNode *GraphNode) *GraphNode {
    return nil // 简化实现
}

func findAttentionNode(graph *ComputeGraph, q, k, v *GraphNode) *GraphNode {
    return nil // 简化实现
}

func findOutputProjection(graph *ComputeGraph, attnNode *GraphNode) *GraphNode {
    return nil // 简化实现
}

func extractNumHeads(node *GraphNode) int {
    return 32 // 简化实现
}

func (g *ComputeGraph) Clone() *ComputeGraph {
    // 深拷贝
    return g // 简化实现
}

package optimization

import (
    "sort"
)

// MemoryPlanner 内存规划器
type MemoryPlanner struct {
    graph        *ComputeGraph
    tensorLifetimes map[string]*TensorLifetime
    memoryPlan   *MemoryPlan
}

// TensorLifetime 张量生命周期
type TensorLifetime struct {
    TensorID   string
    Size       int64
    FirstUse   int // 首次使用的操作索引
    LastUse    int // 最后使用的操作索引
    MemOffset  int64 // 分配的内存偏移
}

// MemoryPlan 内存计划
type MemoryPlan struct {
    TotalSize       int64
    Allocations     map[string]*Allocation
    ReusableBuffers []*BufferPool
}

// Allocation 内存分配
type Allocation struct {
    TensorID string
    Offset   int64
    Size     int64
    Reused   bool
}

// BufferPool 可重用缓冲池
type BufferPool struct {
    Size    int64
    Tensors []string
}

// Plan 规划内存
func (mp *MemoryPlanner) Plan(executionOrder []string) *MemoryPlan {
    // 1. 分析张量生命周期
    mp.analyzeTensorLifetimes(executionOrder)

    // 2. 按大小和生命周期排序
    tensors := mp.getSortedTensors()

    // 3. 使用最佳适配算法分配
    plan := &MemoryPlan{
        Allocations: make(map[string]*Allocation),
    }

    freeList := NewFreeList()

    for _, lt := range tensors {
        // 查找可重用的空间
        offset, reused := freeList.FindBestFit(lt.Size, lt.FirstUse)

        if !reused {
            // 分配新空间
            offset = plan.TotalSize
            plan.TotalSize += lt.Size
        }

        plan.Allocations[lt.TensorID] = &Allocation{
            TensorID: lt.TensorID,
            Offset:   offset,
            Size:     lt.Size,
            Reused:   reused,
        }

        lt.MemOffset = offset

        // 添加到空闲列表（在最后使用后释放）
        freeList.Add(&FreeBlock{
            Offset:    offset,
            Size:      lt.Size,
            FreeAfter: lt.LastUse,
        })
    }

    return plan
}

func (mp *MemoryPlanner) analyzeTensorLifetimes(executionOrder []string) {
    mp.tensorLifetimes = make(map[string]*TensorLifetime)

    for opIdx, nodeID := range executionOrder {
        node := mp.graph.Nodes[nodeID]

        // 处理输入张量
        for _, inputID := range node.Inputs {
            if lt, exists := mp.tensorLifetimes[inputID]; exists {
                lt.LastUse = opIdx
            }
        }

        // 处理输出张量
        for _, outputID := range node.Outputs {
            size := mp.estimateTensorSize(nodeID, outputID)
            mp.tensorLifetimes[outputID] = &TensorLifetime{
                TensorID: outputID,
                Size:     size,
                FirstUse: opIdx,
                LastUse:  opIdx,
            }
        }
    }
}

func (mp *MemoryPlanner) getSortedTensors() []*TensorLifetime {
    tensors := make([]*TensorLifetime, 0, len(mp.tensorLifetimes))
    for _, lt := range mp.tensorLifetimes {
        tensors = append(tensors, lt)
    }

    // 优先分配大张量和生命周期长的张量
    sort.Slice(tensors, func(i, j int) bool {
        // 首先按大小降序
        if tensors[i].Size != tensors[j].Size {
            return tensors[i].Size > tensors[j].Size
        }
        // 然后按生命周期降序
        lifeI := tensors[i].LastUse - tensors[i].FirstUse
        lifeJ := tensors[j].LastUse - tensors[j].FirstUse
        return lifeI > lifeJ
    })

    return tensors
}

func (mp *MemoryPlanner) estimateTensorSize(nodeID, tensorID string) int64 {
    // 根据节点类型和属性估算张量大小
    // 这里简化处理
    return 1024 * 1024 // 1MB
}

// FreeList 空闲列表
type FreeList struct {
    blocks []*FreeBlock
}

type FreeBlock struct {
    Offset    int64
    Size      int64
    FreeAfter int
}

func NewFreeList() *FreeList {
    return &FreeList{
        blocks: make([]*FreeBlock, 0),
    }
}

func (fl *FreeList) FindBestFit(size int64, currentOp int) (int64, bool) {
    var bestBlock *FreeBlock
    var bestIdx int = -1

    for i, block := range fl.blocks {
        // 检查是否已释放且大小足够
        if block.FreeAfter < currentOp && block.Size >= size {
            if bestBlock == nil || block.Size < bestBlock.Size {
                bestBlock = block
                bestIdx = i
            }
        }
    }

    if bestBlock != nil {
        // 找到合适的块
        offset := bestBlock.Offset

        // 如果块更大，拆分
        if bestBlock.Size > size {
            newBlock := &FreeBlock{
                Offset:    offset + size,
                Size:      bestBlock.Size - size,
                FreeAfter: bestBlock.FreeAfter,
            }
            fl.blocks[bestIdx] = newBlock
        } else {
            // 移除使用的块
            fl.blocks = append(fl.blocks[:bestIdx], fl.blocks[bestIdx+1:]...)
        }

        return offset, true
    }

    return 0, false
}

func (fl *FreeList) Add(block *FreeBlock) {
    fl.blocks = append(fl.blocks, block)
}

// InPlaceOptimizer 原地优化器
type InPlaceOptimizer struct {
    graph *ComputeGraph
}

// Optimize 执行原地优化
func (ipo *InPlaceOptimizer) Optimize() {
    for _, node := range ipo.graph.Nodes {
        if ipo.canBeInPlace(node) {
            node.Attrs["inplace"] = true
        }
    }
}

func (ipo *InPlaceOptimizer) canBeInPlace(node *GraphNode) bool {
    // 检查操作是否可以原地执行
    switch node.OpType {
    case "ReLU", "GELU", "Sigmoid", "Tanh":
        // 激活函数可以原地
        return true
    case "Add", "Mul":
        // 逐元素操作在形状相同时可以原地
        return ipo.hasMatchingShapes(node.Inputs)
    case "Dropout":
        // Dropout 可以原地
        return true
    default:
        return false
    }
}

func (ipo *InPlaceOptimizer) hasMatchingShapes(inputs []string) bool {
    // 检查输入形状是否匹配
    return true // 简化实现
}

package optimization

import (
    "errors"
    "sync"
)

// PagedKVCache 分页 KV Cache
type PagedKVCache struct {
    config      *PagedKVConfig
    blocks      []*KVBlock
    freeBlocks  []int
    blockTables map[int][]int // sequence_id -> block_indices
    mu          sync.RWMutex
}

type PagedKVConfig struct {
    NumLayers    int
    NumHeads     int
    HeadDim      int
    BlockSize    int   // 每块的 token 数
    NumGPUBlocks int
    NumCPUBlocks int
    DType        string // fp16, fp32
}

// KVBlock KV Cache 块
type KVBlock struct {
    KeyCache   [][]float16   // [num_tokens, num_heads * head_dim]
    ValueCache [][]float16   // [num_tokens, num_heads * head_dim]
    NumTokens  int
    RefCount   int
}

type float16 float32 // 简化表示

func NewPagedKVCache(config *PagedKVConfig) *PagedKVCache {
    cache := &PagedKVCache{
        config:      config,
        blocks:      make([]*KVBlock, config.NumGPUBlocks),
        freeBlocks:  make([]int, config.NumGPUBlocks),
        blockTables: make(map[int][]int),
    }

    // 初始化所有块
    for i := 0; i < config.NumGPUBlocks; i++ {
        cache.blocks[i] = &KVBlock{
            KeyCache:   make([][]float16, config.BlockSize),
            ValueCache: make([][]float16, config.BlockSize),
        }
        for j := 0; j < config.BlockSize; j++ {
            cache.blocks[i].KeyCache[j] = make([]float16, config.NumHeads*config.HeadDim)
            cache.blocks[i].ValueCache[j] = make([]float16, config.NumHeads*config.HeadDim)
        }
        cache.freeBlocks[i] = i
    }

    return cache
}

// AllocateBlocks 为序列分配块
func (c *PagedKVCache) AllocateBlocks(seqID, numBlocks int) ([]int, error) {
    c.mu.Lock()
    defer c.mu.Unlock()

    if len(c.freeBlocks) < numBlocks {
        return nil, errors.New("not enough free blocks")
    }

    // 分配块
    allocated := make([]int, numBlocks)
    copy(allocated, c.freeBlocks[:numBlocks])
    c.freeBlocks = c.freeBlocks[numBlocks:]

    // 增加引用计数
    for _, blockIdx := range allocated {
        c.blocks[blockIdx].RefCount = 1
    }

    // 更新块表
    c.blockTables[seqID] = allocated

    return allocated, nil
}

// AppendToken 追加 token 的 KV Cache
func (c *PagedKVCache) AppendToken(
    seqID int,
    layer int,
    key, value []float16,
) error {
    c.mu.Lock()
    defer c.mu.Unlock()

    blockTable := c.blockTables[seqID]
    if blockTable == nil {
        return errors.New("sequence not found")
    }

    // 计算当前 token 位置
    totalTokens := c.getSequenceLength(seqID)
    blockIdx := totalTokens / c.config.BlockSize
    slotIdx := totalTokens % c.config.BlockSize

    // 检查是否需要新块
    if blockIdx >= len(blockTable) {
        // 需要分配新块
        if len(c.freeBlocks) == 0 {
            return errors.New("no free blocks")
        }
        newBlockIdx := c.freeBlocks[0]
        c.freeBlocks = c.freeBlocks[1:]
        c.blocks[newBlockIdx].RefCount = 1
        blockTable = append(blockTable, newBlockIdx)
        c.blockTables[seqID] = blockTable
    }

    // 写入 KV Cache
    physicalBlock := blockTable[blockIdx]
    copy(c.blocks[physicalBlock].KeyCache[slotIdx], key)
    copy(c.blocks[physicalBlock].ValueCache[slotIdx], value)
    c.blocks[physicalBlock].NumTokens++

    return nil
}

// ForkSequence 分叉序列（用于 beam search）
func (c *PagedKVCache) ForkSequence(parentSeqID, childSeqID int) error {
    c.mu.Lock()
    defer c.mu.Unlock()

    parentTable := c.blockTables[parentSeqID]
    if parentTable == nil {
        return errors.New("parent sequence not found")
    }

    // 创建子序列块表（共享块）
    childTable := make([]int, len(parentTable))
    copy(childTable, parentTable)

    // 增加引用计数（Copy-on-Write）
    for _, blockIdx := range childTable {
        c.blocks[blockIdx].RefCount++
    }

    c.blockTables[childSeqID] = childTable
    return nil
}

// CopyOnWrite 写时复制
func (c *PagedKVCache) CopyOnWrite(seqID, logicalBlockIdx int) (int, error) {
    c.mu.Lock()
    defer c.mu.Unlock()

    blockTable := c.blockTables[seqID]
    if blockTable == nil || logicalBlockIdx >= len(blockTable) {
        return -1, errors.New("invalid sequence or block index")
    }

    physicalBlockIdx := blockTable[logicalBlockIdx]

    // 检查是否需要复制
    if c.blocks[physicalBlockIdx].RefCount == 1 {
        return physicalBlockIdx, nil // 不需要复制
    }

    // 分配新块
    if len(c.freeBlocks) == 0 {
        return -1, errors.New("no free blocks for copy")
    }
    newBlockIdx := c.freeBlocks[0]
    c.freeBlocks = c.freeBlocks[1:]

    // 复制数据
    oldBlock := c.blocks[physicalBlockIdx]
    newBlock := c.blocks[newBlockIdx]

    for i := 0; i < c.config.BlockSize; i++ {
        copy(newBlock.KeyCache[i], oldBlock.KeyCache[i])
        copy(newBlock.ValueCache[i], oldBlock.ValueCache[i])
    }
    newBlock.NumTokens = oldBlock.NumTokens
    newBlock.RefCount = 1

    // 减少原块引用
    oldBlock.RefCount--

    // 更新块表
    blockTable[logicalBlockIdx] = newBlockIdx

    return newBlockIdx, nil
}

// FreeSequence 释放序列
func (c *PagedKVCache) FreeSequence(seqID int) {
    c.mu.Lock()
    defer c.mu.Unlock()

    blockTable := c.blockTables[seqID]
    if blockTable == nil {
        return
    }

    for _, blockIdx := range blockTable {
        c.blocks[blockIdx].RefCount--
        if c.blocks[blockIdx].RefCount == 0 {
            // 块空闲
            c.blocks[blockIdx].NumTokens = 0
            c.freeBlocks = append(c.freeBlocks, blockIdx)
        }
    }

    delete(c.blockTables, seqID)
}

// GetBlockTable 获取块表
func (c *PagedKVCache) GetBlockTable(seqID int) []int {
    c.mu.RLock()
    defer c.mu.RUnlock()

    table := c.blockTables[seqID]
    if table == nil {
        return nil
    }

    result := make([]int, len(table))
    copy(result, table)
    return result
}

func (c *PagedKVCache) getSequenceLength(seqID int) int {
    blockTable := c.blockTables[seqID]
    if len(blockTable) == 0 {
        return 0
    }

    totalTokens := 0
    for _, blockIdx := range blockTable {
        totalTokens += c.blocks[blockIdx].NumTokens
    }
    return totalTokens
}

// GetFreeBlockCount 获取空闲块数量
func (c *PagedKVCache) GetFreeBlockCount() int {
    c.mu.RLock()
    defer c.mu.RUnlock()
    return len(c.freeBlocks)
}

// GetUtilization 获取内存利用率
func (c *PagedKVCache) GetUtilization() float64 {
    c.mu.RLock()
    defer c.mu.RUnlock()

    usedBlocks := c.config.NumGPUBlocks - len(c.freeBlocks)
    return float64(usedBlocks) / float64(c.config.NumGPUBlocks)
}