internal/autogen/estimate.go
Rama main · réplica de lectura del repositorio original.
package autogen
// estimate.go exposes the per-profile sizer as a one-shot preview so the web
// editor can show live VRAM/RAM use for a candidate tuning before it is saved.
// It mirrors the solo-profile path of emitModel without writing any config.
// EstimateInput is the subset of editor fields that affect placement/memory.
// Zero Ctx means "let the sizer pick". Zero TargetVramGB uses settings.
type EstimateInput struct {
Ctx int
// CtxExact marks Ctx as a custom launch pin (-c): the process runs that
// window whatever the sizer would have picked, so it is not rounded up to a
// 4096 multiple the way an auto pick is.
CtxExact bool
KvK string
KvV string
KvInRam bool
Spec string
// RopeScaling ("linear"/"yarn") lets Ctx exceed the model's trained length —
// the preview must know, or it silently sizes a clamped window and reports a
// KV reserve the real launch won't have.
RopeScaling string
TargetVramGB float64
CpuOffload int // >0 pins layers offloaded to CPU, overriding the sizer
CpuOffloadSet bool // CpuOffload == 0 is a pin too (every layer on GPU)
CtxCheckpoints *int // nil => llama default (32); 0 disables; reserves checkpoint VRAM
// CpuDevice previews a `device: cpu` pin (--device none): the same ngl-0
// footprint generate prices, all of it charged to RAM.
CpuDevice bool
// CheckpointMinStep pins -cms (checkpoint spacing in prompt tokens), which
// scales each snapshot's global-KV term. 0 => the arch default (wide on SWA).
CheckpointMinStep int
DraftGB float64 // separate MTP/DFlash draft gguf weights (GB); 0 => baked-in or none
MmprojGB float64 // "-vision" twin projector footprint (weights + CLIP reserve); 0 => none
// Ub pins -ub (physical batch). The compute buffer scales with it — ~0.5 GB
// between 512 and 1024 on a large-vocab model — so a preview that ignored a
// per-model Override.Ub charged a buffer the real launch never has. 0 => the
// same auto pick emit makes (effectiveUb).
Ub int
// KvKDraft/KvVDraft pin the draft context's cache type (-ctkd/-ctvd). Empty
// means "whatever the emitter would default to", i.e. the main K/V quant. A
// caller reconstructing a HAND-WRITTEN argv that carries -ctk but no -ctkd
// must pass "f16" explicitly, because that is what llama uses there.
KvKDraft string
KvVDraft string
// Parallel is llama-server's slot count (--parallel). Every slot gets its own
// window over one shared pool, so the KV cost the sizer solves against is
// Parallel x the per-token slope. 0/1 => a single slot.
Parallel int
// DraftKind is the paired draft sidecar's kind ("mtp"/"dflash", "" for none),
// as DraftSidecarForDir reports it. Only consulted when Spec is empty, to
// resolve the same auto spec the emitter picks — see EstimatePlan.
DraftKind string
}
// EstimateResult is the previewed load plan for a candidate tuning.
type EstimateResult struct {
Ctx int `json:"ctx"`
Ngl int `json:"ngl"`
NCpuMoe int `json:"nCpuMoe"`
EstVramGB float64 `json:"estVramGB"`
EstRamGB float64 `json:"estRamGB"`
TargetVramGB float64 `json:"targetVramGB"`
MaxRamGB float64 `json:"maxRamGB"`
KvReserveGB float64 `json:"kvReserveGB"`
// CheckpointGB is the VRAM reserved for context checkpoints (included in
// EstVramGB via overhead, broken out so the UI can attribute it separately
// from model weights).
CheckpointGB float64 `json:"checkpointGB"`
// DraftGB is the VRAM charged for the speculative draft / MTP nextn layer: a
// baked-in nextn head costs its compute pad, its own compute GRAPH and its
// OWN ctx-scaled KV cache (a second llama_context over the same model at the
// same window), a separate draft gguf its weights + pad. Folded into
// EstVramGB via overhead; broken out so the UI can attribute it separately
// from the main model weights. 0 when no draft-mtp spec is active.
DraftGB float64 `json:"draftGB"`
// ComputeBufGB is the GPU compute buffer: the main context's graph (logits +
// activations) plus the fixed per-process runtime constant for the backend in
// use (see runtimeCtxGB). A baked-in MTP drafter's own graph is NOT in here,
// it is charged to DraftGB. MmprojGB is a "-vision" twin's
// projector weights + CLIP reserve (0 for non-vision). OverheadGB is the global
// vramOverheadGB safety headroom. All three are folded into EstVramGB via
// overhead; broken out so the UI can label them separately from the model
// weights instead of lumping everything non-KV into "Weights".
ComputeBufGB float64 `json:"computeBufGB"`
MmprojGB float64 `json:"mmprojGB"`
OverheadGB float64 `json:"overheadGB"`
// FixedGB is the whole non-splittable share of EstVramGB: compute buffer,
// CUDA context, projector, draft, safety headroom. On a --tensor-split load
// these live on --main-gpu alone, so the spawn guard charges FixedGB to the
// main device before it re-derives the ratio.
FixedGB float64 `json:"fixedGB"`
RamExceeded bool `json:"ramExceeded"`
IsMoE bool `json:"isMoE"`
}
// EstimatePlan sizes one candidate tuning against the given settings + gguf
// metadata, returning the chosen ctx and the resulting VRAM/RAM estimate. It
// reuses sizeProfile/forceLowActiveMoE so the preview matches what a save would
// actually emit for the solo profile.
func EstimatePlan(s Settings, meta Metadata, in EstimateInput) (EstimateResult, error) {
// KV quant: the fleet default (f16, stepping down to q8_0 only under VRAM
// pressure) unless a valid matched override is given — mirrors emitModel, so
// the preview's KV reserve matches what a save would emit.
// An empty spec means AUTO, not "no drafter" — the emitter runs every model
// through effectiveSpec, which hands an MTP-capable gguf draft-mtp+ngram-mod
// and charges its drafter (the pad plus its ctx-scaled KV, or the sidecar's
// real weights). The preview used to read "" as no spec and charge 0, so a
// model left on auto spec previewed a whole drafter lighter than it bakes. On a tight budget that is the whole margin:
// Qwen3.8-27B at 106k previewed -ngl 99 / 0 GB RAM and emitted -ngl 64 / 0.32.
spec := in.Spec
if spec == "" {
spec = effectiveSpec(meta, nil, in.DraftKind)
}
estTarget := s.TargetVramGB
if in.TargetVramGB > 0 {
estTarget = in.TargetVramGB
}
// Only charge the sidecar when the active spec will actually attach it as -md
// (see matchedDraftSizeGB): a dir can pair a DFlash drafter to a model that
// runs on a baked-in MTP head, and that drafter never loads.
draftGB := matchedDraftSizeGB(spec, in.DraftKind, in.DraftGB)
kvDef := defaultKvQuant(s, meta, estTarget, s.VramOverheadGB+draftOverheadGB(spec, draftGB))
// The editor sends kvK/kvV independently, so a form that resolves only one
// side must not sink the preview back to the fleet default — see
// resolveKvPair.
kvK, kvV := resolveKvPair(in.KvK, in.KvV, kvDef, kvDef)
perTokGB, kvConstGB := 0.0, 0.0
if m := GetKvCostModel(meta, kvK, kvV); m.OK {
perTokGB, kvConstGB = m.SlopeGB, m.ConstGB
}
// N slots share one -c pool, each with its own KV, so charge N x. buildCmdLines
// multiplies the per-slot window back out for -c; charging one slot here
// under-sized the KV reserve and the admission figure by N.
slots := EffectiveParallel(&Override{Parallel: in.Parallel})
if slots > 1 {
perTokGB *= float64(slots)
kvConstGB *= float64(slots)
}
// Rope scaling lifts the trained-length ceiling; without it this is nativeCtx.
modelMax := ropeCeiling(meta, in.RopeScaling, in.Ctx)
target := s.TargetVramGB
if in.TargetVramGB > 0 {
target = in.TargetVramGB
}
// Draft overhead: a baked-in MTP nextn layer costs its compute pad here. A
// separate draft file (Gemma-4's MTP sidecar, or any DFlash drafter, always
// separate) instead charges its real on-disk weights + a small KV/compute pad,
// so big drafts scale up rather than under-counting.
specOh := draftOverheadGB(spec, draftGB)
// ...plus, for a baked-in nextn head, its own KV over the whole window. That
// part is a slope, so it has to reach the sizer as one: charging it as flat
// overhead would price it at a ctx the sizer had not picked yet.
dKvK, dKvV := draftKvPair(in.KvKDraft, in.KvVDraft, kvK, kvV)
draftSlopeGB := mtpDraftSlopeFor(meta, spec, dKvK, dKvV, draftGB)
if slots > 1 {
draftSlopeGB *= float64(slots)
}
prof := profile{
Name: "estimate",
Target: target,
Overhead: s.VramOverheadGB + specOh,
Ctx: in.Ctx,
CtxExact: in.CtxExact,
Spec: spec,
KvK: kvK,
KvV: kvV,
IsLong: in.Ctx >= longCtxThreshold,
CtxCheckpoints: in.CtxCheckpoints,
CheckpointMinStep: in.CheckpointMinStep,
DraftSlopeGB: draftSlopeGB,
}
// Charge the ub the launch will actually run with: effectiveUb reads the
// override's pinned value, and passing nil here made the preview size a
// different compute buffer than emit did for the same model.
computeUb := effectiveUb(meta, prof, &Override{Ub: in.Ub}, prof.Ctx, target)
computeBufGB := computeBufferGB(meta, computeUb, s.ComputeBufFactor)
// The baked-in drafter's second context allocates a graph of its own at the
// same ub. It is reported under Draft rather than folded into the compute
// buffer, so the UI keeps attributing it to the thing that caused it.
draftComputeGB := mtpDraftComputeGB(meta, spec, draftGB, computeUb, s.ComputeBufFactor)
prof.Overhead += computeBufGB + draftComputeGB
prof.Overhead += in.MmprojGB // "-vision" projector weights + CLIP compute reserve
// Everything charged so far lands on --main-gpu alone, so this is the figure
// the spawn guard re-derives the ratio from (FixedGB below). Every device
// past the main one pays its own runtime context on top; that part is not
// splittable either, but it belongs to the POOLED budget rather than to the
// main card, so it is added after FixedGB is taken. Generate does the same,
// and the two must agree or a spawn-time retune moves the split for no reason
// other than which code path computed it.
mainFixedGB := prof.Overhead
if gpus := s.GpuSetOrEmpty(); gpus.Multi() {
prof.Overhead += gpus.ExtraDeviceOverheadGB()
}
ctx, plan, kvReserve, planCkptGB, err := sizeProfile(meta, s, prof, perTokGB, kvConstGB, modelMax, in.KvInRam)
if err != nil {
return EstimateResult{}, err
}
// Checkpoints live wherever the KV cache does. When KV is in RAM they cost
// no VRAM, so report 0 to stay consistent with EstVramGB (which sizeProfile
// only inflates in the VRAM-KV branch).
checkpointGB := 0.0
if !in.KvInRam {
ckptCtxCeil := modelMax
if prof.Ctx != 0 {
ckptCtxCeil = min(ckptCtxCeil, prof.Ctx)
}
checkpointGB = checkpointReserveGB(prof, perTokGB, kvConstGB, ckptCtxCeil, meta.FullAttnInterval > 0)
}
ngl, ncpuMoe := forceLowActiveMoE(meta, plan, prof, kvReserve)
if in.CpuOffload > 0 || in.CpuOffloadSet {
ngl, ncpuMoe = applyForcedOffload(meta, in.CpuOffload)
}
if in.CpuDevice {
ngl, ncpuMoe = 0, 0
}
// Any placement that differs from the one GetLoadPlan priced has to be
// re-priced: forceLowActiveMoE rewrites -ngl/--n-cpu-moe in place and used to
// leave plan.EstVramGB describing the placement it just discarded.
if ngl != plan.Ngl || ncpuMoe != plan.NCpuMoe || in.CpuDevice {
plan.EstVramGB, plan.EstRamGB = estForOffload(meta, prof, kvReserve, planCkptGB, ngl, ncpuMoe)
if in.CpuDevice {
plan.EstVramGB, plan.EstRamGB = 0, plan.EstVramGB+plan.EstRamGB
}
plan.RamExceeded = s.MaxRamGB > 0 && plan.EstRamGB > s.MaxRamGB
}
// Dense checkpoints are split GPU/RAM by the layer placement (see sizeProfile),
// so the VRAM portion is only the GPU-resident fraction. Scale the reported
// figure to match EstVramGB. MoE keeps KV (and checkpoints) fully in VRAM.
if in.CpuDevice {
checkpointGB = 0
}
if checkpointGB > 0 && !meta.IsMoE && meta.BlockCount > 0 {
gpuFrac := float64(min(ngl, int(meta.BlockCount))) / float64(meta.BlockCount)
if gpuFrac < 0 {
gpuFrac = 0
}
checkpointGB *= gpuFrac
}
// sizeProfile folds the drafter's KV into kvReserve (it sizes against the
// combined slope). Split it back out so the UI's KV segment stays the main
// cache and the Draft segment shows what the drafter actually costs.
// KV-in-RAM is the exception: that branch sizes against the main slope alone
// (the reserve it reports is a flat RAM-side placeholder), so there is
// nothing folded in to take back out.
draftKvGB := 0.0
if !in.KvInRam {
draftKvGB = min(draftSlopeGB*float64(ctx), kvReserve)
kvReserve -= draftKvGB
}
res := EstimateResult{
Ctx: ctx,
Ngl: ngl,
NCpuMoe: ncpuMoe,
EstVramGB: plan.EstVramGB,
EstRamGB: plan.EstRamGB,
TargetVramGB: target,
MaxRamGB: s.MaxRamGB,
KvReserveGB: kvReserve,
CheckpointGB: checkpointGB,
DraftGB: specOh + draftComputeGB + draftKvGB,
ComputeBufGB: computeBufGB,
MmprojGB: in.MmprojGB,
OverheadGB: s.VramOverheadGB,
FixedGB: mainFixedGB,
RamExceeded: plan.RamExceeded,
IsMoE: meta.IsMoE,
}
// A CPU-device pin has no GPU, so every VRAM component is zero, not just the
// total: the UI draws each one as its own segment of the VRAM bar, and the
// spawn guard charges FixedGB to the main card.
if in.CpuDevice {
res.KvReserveGB, res.DraftGB, res.ComputeBufGB = 0, 0, 0
res.MmprojGB, res.OverheadGB, res.FixedGB = 0, 0, 0
}
return res, nil
}