internal/autogen/plan.go
Rama main · réplica de lectura del repositorio original.
package autogen
import (
"fmt"
"math"
"strings"
)
// moeExpertShare is the fraction of total weight in expert tensors, per arch.
// Approximate; mirrors the PowerShell planner table.
var moeExpertShare = map[string]float64{
"gemma4": 0.88,
"qwen3": 0.90,
"qwen3moe": 0.90,
"llama": 0.80,
"lfm2": 0.78,
"lfm2moe": 0.78,
}
const moeShareDefault = 0.85
// moeShareFor returns the expert-weight share for an architecture.
func moeShareFor(arch string) float64 {
if s, ok := moeExpertShare[strings.ToLower(arch)]; ok {
return s
}
return moeShareDefault
}
// effectiveShare prefers the exact expert-weight share derived from the gguf
// tensor section; it falls back to a per-arch heuristic only when the tensor
// section could not be sized (ExpertWeightShare == 0).
func effectiveShare(meta Metadata, fallback func(string) float64) float64 {
if meta.ExpertWeightShare > 0 {
return meta.ExpertWeightShare
}
return fallback(meta.Architecture)
}
// LoadPlan is the chosen llama-server placement plus diagnostics.
type LoadPlan struct {
Ngl int
NCpuMoe int
EstVramGB float64
EstRamGB float64
TargetVramGB float64
MaxRamGB float64
BlockCount int
Architecture string
IsMoE bool
FileSizeGB float64
RamExceeded bool
Reasons []string
}
// PlanOptions are the budgets fed to GetLoadPlan. Zero values fall back to the
// PowerShell defaults (KvCacheReserveGB=1.0, CudaOverheadGB=0.5).
type PlanOptions struct {
TargetVramGB float64
MaxRamGB float64 // 0 = no RAM check
KvCacheReserveGB float64
CudaOverheadGB float64
kvReserveSet bool // internal: distinguish an explicit 0 reserve
cudaSet bool
}
// densePlacement chooses -ngl for a model placed layer-by-layer. Each
// GPU-resident layer carries BOTH its weight slice and its KV-cache slice in
// VRAM; layers left on the CPU keep their KV in RAM. So the full-ctx KV reserve
// is spread across blocks (kvReserve/blocks per layer) rather than subtracted
// whole from the VRAM budget up front. The old "usableVram = target - kvReserve"
// form reserved every layer's KV in VRAM even for CPU-resident layers, which
// drove ngl toward 0 while the reserved VRAM sat unused (e.g. a 24B dense at 64k
// ctx pinned to -ngl 0 on a 7GB budget). Diverges from the PowerShell
// Get-LlamaLoadPlan, which shares the same flaw.
func densePlacement(size, kvReserve float64, blocks int, target, cudaOverhead float64) (ngl int, estVram, estRam float64) {
perLayer := (size + kvReserve) / float64(blocks)
budget := target - cudaOverhead
if perLayer > 0 && budget > 0 {
ngl = int(math.Floor(budget / perLayer))
}
if ngl > blocks+1 {
ngl = blocks + 1
}
if ngl < 0 {
ngl = 0
}
gpuFrac := float64(ngl) / float64(blocks)
if gpuFrac > 1 {
gpuFrac = 1
}
estVram = gpuFrac*(size+kvReserve) + cudaOverhead
estRam = (1 - gpuFrac) * (size + kvReserve)
return
}
// GetLoadPlan derives -ngl and --n-cpu-moe from a VRAM budget and gguf metadata.
// Port of Get-LlamaLoadPlan.ps1.
func GetLoadPlan(meta Metadata, opt PlanOptions) (LoadPlan, error) {
kvReserve := opt.KvCacheReserveGB
if !opt.kvReserveSet && kvReserve == 0 {
kvReserve = 1.0
}
cudaOverhead := opt.CudaOverheadGB
if !opt.cudaSet && cudaOverhead == 0 {
cudaOverhead = 0.5
}
arch := strings.ToLower(meta.Architecture)
size := meta.FileSizeGB
blocks := int(meta.BlockCount)
isMoE := meta.IsMoE
if blocks <= 0 {
return LoadPlan{}, fmt.Errorf("load plan: invalid block_count=%d in metadata", blocks)
}
if size <= 0 {
return LoadPlan{}, fmt.Errorf("load plan: invalid file size for %s", meta.Path)
}
usableVram := math.Max(0.0, opt.TargetVramGB-kvReserve-cudaOverhead)
ngl := 99
ncpuMoe := 0
var reasons []string
reasons = append(reasons, fmt.Sprintf("file=%.2fGB blocks=%d arch=%s moe=%v target_vram=%.1fGB usable_vram=%.1fGB",
size, blocks, arch, isMoE, opt.TargetVramGB, usableVram))
var estVram, estRam float64
switch {
case size <= usableVram:
reasons = append(reasons, "fits whole-model in VRAM budget -> -ngl 99")
estVram = size + kvReserve + cudaOverhead
estRam = 0.0
case isMoE:
share := effectiveShare(meta, moeShareFor)
nonExpertTotal := size * (1.0 - share)
usableForExperts := usableVram - nonExpertTotal
if usableForExperts <= 0 {
// Dense path can't even fit; fall back to -ngl reduction.
ngl, estVram, estRam = densePlacement(size, kvReserve, blocks, opt.TargetVramGB, cudaOverhead)
ncpuMoe = 0
reasons = append(reasons, fmt.Sprintf("MoE share=%.2f but non-expert weight %.2fGB exceeds usable VRAM -> dense fallback -ngl=%d", share, nonExpertTotal, ngl))
} else {
perMoeLayer := (size * share) / float64(blocks)
moeLayersOnGpu := math.Floor(usableForExperts / perMoeLayer)
if moeLayersOnGpu > float64(blocks) {
moeLayersOnGpu = float64(blocks)
}
ncpuMoeCandidate := int(math.Max(0, float64(blocks)-moeLayersOnGpu))
// Crossover: when more than ~half the blocks need CPU experts, the
// per-token PCIe round-trip cost exceeds the dense-on-GPU saving, so
// naive -ngl (one boundary transition) wins.
crossoverFrac := 0.5
if float64(ncpuMoeCandidate) > float64(blocks)*crossoverFrac {
ngl, estVram, estRam = densePlacement(size, kvReserve, blocks, opt.TargetVramGB, cudaOverhead)
ncpuMoe = 0
reasons = append(reasons, fmt.Sprintf("MoE share=%.2f but ncpumoe candidate %d/%d exceeds %.0f%% crossover -> naive -ngl=%d (avoid PCIe thrash)", share, ncpuMoeCandidate, blocks, crossoverFrac*100, ngl))
} else {
ncpuMoe = ncpuMoeCandidate
ngl = 99
expertGpuFrac := float64(blocks-ncpuMoe) / float64(blocks)
estVram = nonExpertTotal + (size * share * expertGpuFrac) + kvReserve + cudaOverhead
estRam = (size * share) * (float64(ncpuMoe) / float64(blocks))
reasons = append(reasons, fmt.Sprintf("MoE share=%.2f non_expert=%.2fGB per_moe_layer=%.3fGB -> --n-cpu-moe=%d", share, nonExpertTotal, perMoeLayer, ncpuMoe))
}
}
default:
ngl, estVram, estRam = densePlacement(size, kvReserve, blocks, opt.TargetVramGB, cudaOverhead)
ncpuMoe = 0
reasons = append(reasons, fmt.Sprintf("dense per_layer=%.3fGB(w)+%.3fGB(kv) -> -ngl=%d/%d", size/float64(blocks), kvReserve/float64(blocks), ngl, blocks))
}
ramExceeded := false
if opt.MaxRamGB > 0 && estRam > opt.MaxRamGB {
ramExceeded = true
reasons = append(reasons, fmt.Sprintf("WARN estimated RAM %.2fGB exceeds MaxRamGB %.1fGB", estRam, opt.MaxRamGB))
}
return LoadPlan{
Ngl: ngl,
NCpuMoe: ncpuMoe,
EstVramGB: round(estVram, 2),
EstRamGB: round(estRam, 2),
TargetVramGB: opt.TargetVramGB,
MaxRamGB: opt.MaxRamGB,
BlockCount: blocks,
Architecture: arch,
IsMoE: isMoE,
FileSizeGB: size,
RamExceeded: ramExceeded,
Reasons: reasons,
}, nil
}