internal/autogen/mtpdraft_test.go
Rama main · réplica de lectura del repositorio original.
package autogen
import (
"math"
"strings"
"testing"
)
// qwen38Meta is Qwen3.8-27B-UD-Q2_K_XL as its gguf reports it: a hybrid
// GatedDeltaNet with one full-attention layer in four, and ONE baked-in nextn
// (MTP) layer. It is the model that exposed the flat draft charge.
func qwen38Meta() Metadata {
return Metadata{
Architecture: "qwen35", BlockCount: 65, HeadCountKv: 4,
KeyLength: 256, ValueLength: 256, FullAttnInterval: 4,
FileSizeGB: 9.154, ContextLength: 262144,
IsMTP: true, NextnLayers: 1,
}
}
// The baked-in MTP drafter is a second context over the same model at the same
// n_ctx, so its KV is a per-token cost over its nextn layers, not a constant.
func TestMtpDraftSlopeGB(t *testing.T) {
meta := qwen38Meta()
// 1 nextn layer x 4 kv heads x (256 K + 256 V) x 2 bytes = 4096 B/token.
wantF16 := 4096.0 / gib
if got := mtpDraftSlopeGB(meta, "f16", "f16"); math.Abs(got-wantF16) > 1e-12 {
t.Fatalf("f16 slope=%g want %g", got, wantF16)
}
// Matching the draft cache to a q8_0 main cache is what -ctkd buys: at the
// 159744-token window this model sizes to, 0.61 GB becomes 0.32 GB.
q8 := mtpDraftSlopeGB(meta, "q8_0", "q8_0")
if ratio := q8 / wantF16; math.Abs(ratio-1.0625/2) > 1e-9 {
t.Fatalf("q8_0/f16 slope ratio=%g want %g", ratio, 1.0625/2)
}
if gb := wantF16 * 159744; math.Abs(gb-0.6094) > 1e-3 {
t.Fatalf("f16 draft KV at 159744 ctx = %.4f GB, want ~0.61", gb)
}
// No nextn layers (or no dims): nothing to charge.
plain := meta
plain.NextnLayers, plain.IsMTP = 0, false
if got := mtpDraftSlopeGB(plain, "f16", "f16"); got != 0 {
t.Fatalf("non-MTP model charged %g", got)
}
// The gate follows the emitter: a baked-in head scales, a separate draft
// gguf does not (it is charged its own weights by draftOverheadGB), and a
// spec chain without draft-mtp charges nothing at all.
if got := mtpDraftSlopeFor(meta, "draft-mtp+ngram-mod", "f16", "f16", 0); got != mtpDraftSlopeGB(meta, "f16", "f16") {
t.Fatalf("baked-in MTP not charged: %g", got)
}
if got := mtpDraftSlopeFor(meta, "draft-mtp+ngram-mod", "f16", "f16", 0.46); got != 0 {
t.Fatalf("sidecar draft charged the baked-in slope: %g", got)
}
if got := mtpDraftSlopeFor(meta, "ngram-mod", "f16", "f16", 0); got != 0 {
t.Fatalf("non-draft spec charged %g", got)
}
}
// The drafter's KV must scale with the window in the estimate, and must be
// reported under Draft rather than inflating the main KV segment.
func TestEstimatePlan_mtpDraftIsCtxScaled(t *testing.T) {
meta := qwen38Meta()
s := Settings{TargetVramGB: 40, VramOverheadGB: 1, ComputeBufFactor: 1}
s.applyDefaults()
slope := mtpDraftSlopeGB(meta, "q8_0", "q8_0") // -ctkd defaults to the main quant
small, err := EstimatePlan(s, meta, EstimateInput{Ctx: 32768, KvK: "q8_0", KvV: "q8_0"})
if err != nil {
t.Fatal(err)
}
big, err := EstimatePlan(s, meta, EstimateInput{Ctx: 131072, KvK: "q8_0", KvV: "q8_0"})
if err != nil {
t.Fatal(err)
}
if small.Ctx != 32768 || big.Ctx != 131072 {
t.Fatalf("ctx not honored: %d / %d", small.Ctx, big.Ctx)
}
wantDelta := slope * float64(big.Ctx-small.Ctx)
if got := big.DraftGB - small.DraftGB; math.Abs(got-wantDelta) > 1e-9 {
t.Fatalf("draft charge delta=%.4f want %.4f (flat charge would be 0)", got, wantDelta)
}
// Pad + KV, not one or the other.
wantSmall := mtpDraftPadGB + slope*float64(small.Ctx)
if math.Abs(small.DraftGB-wantSmall) > 1e-9 {
t.Fatalf("draft charge=%.4f want %.4f", small.DraftGB, wantSmall)
}
// The KV segment stays the MAIN cache: the drafter's share is split back out.
m := GetKvCostModel(meta, "q8_0", "q8_0")
if want := KvReserveGB(big.Ctx, m.SlopeGB, m.ConstGB); math.Abs(big.KvReserveGB-want) > 1e-9 {
t.Fatalf("kv reserve=%.4f want %.4f (drafter must not land in the KV segment)", big.KvReserveGB, want)
}
}
// -ctkd/-ctvd are emitted for ANY active draft spec, defaulted to the model's
// own -ctk/-ctv. A baked-in MTP head attaches no -md but still gets a draft
// context, and llama defaults that context's cache to f16 whatever -ctk says.
func TestBuildCmdLines_draftKvFollowsMain(t *testing.T) {
s := Settings{}
s.applyDefaults()
prof := profile{Name: "t", Ctx: 8192}
join := func(meta Metadata, ov *Override) string {
return strings.Join(buildCmdLines(s, meta, GgufRow{FullPath: "/m.gguf"}, prof, 8192, 99, 0, "q8_0", "q8_0", false, ov), " ")
}
got := join(qwen38Meta(), &Override{})
if !strings.Contains(got, "--spec-type draft-mtp") {
t.Fatalf("test assumption broken, no MTP spec: %s", got)
}
if !strings.Contains(got, "-ctkd q8_0 -ctvd q8_0") {
t.Fatalf("baked-in MTP draft KV not matched to the main quant: %s", got)
}
// An explicit override still wins on either side.
got = join(qwen38Meta(), &Override{KvKDraft: "f16", KvVDraft: "f16"})
if !strings.Contains(got, "-ctkd f16 -ctvd f16") {
t.Fatalf("draft KV override ignored: %s", got)
}
// No draft backend in the chain: no draft context, no flags.
if got := join(Metadata{Architecture: "qwen35"}, &Override{}); strings.Contains(got, "-ctkd") {
t.Fatalf("draft KV emitted without a draft spec: %s", got)
}
}
// qwen38MetaDims is the same model with the two dims a compute graph needs. The
// gguf carries them; qwen38Meta leaves them off so the KV tests above exercise
// the KV math alone.
func qwen38MetaDims() Metadata {
m := qwen38Meta()
m.EmbeddingLength, m.VocabSize = 5120, 248320
return m
}
// The baked-in drafter is a second llama_context over the same model, inheriting
// n_ctx and n_ubatch, so it allocates a second compute graph: same vocabulary,
// same physical batch, hence the same size. It is charged only when the drafter
// really is baked in - an external draft gguf is already priced by its weights.
func TestMtpDraftComputeGB(t *testing.T) {
meta := qwen38MetaDims()
graph, ok := computeGraphGB(meta, 512, 1.0)
if !ok || graph <= 0 {
t.Fatalf("test assumption broken: graph=%g ok=%v", graph, ok)
}
if got := mtpDraftComputeGB(meta, "draft-mtp", 0, 512, 1.0); got != graph {
t.Errorf("baked-in MTP: got %.4f, want the full graph %.4f", got, graph)
}
// A matched draft gguf: its context is a different, much smaller model.
if got := mtpDraftComputeGB(meta, "draft-mtp", 1.2, 512, 1.0); got != 0 {
t.Errorf("external draft model should not be charged the target graph, got %.4f", got)
}
// No draft backend at all.
if got := mtpDraftComputeGB(meta, "ngram", 0, 512, 1.0); got != 0 {
t.Errorf("no draft-mtp spec: got %.4f, want 0", got)
}
// Missing dims leave draftOverheadGB's flat pad as the whole charge.
if got := mtpDraftComputeGB(qwen38Meta(), "draft-mtp", 0, 512, 1.0); got != 0 {
t.Errorf("missing dims: got %.4f, want 0 (pad only)", got)
}
}
// The drafter's graph rides in DraftGB, not ComputeBufGB: the UI attributes it
// to the draft segment, and the process-wide runtime constant is charged once.
func TestEstimatePlan_mtpDraftChargesItsOwnGraph(t *testing.T) {
meta := qwen38MetaDims()
s := Settings{TargetVramGB: 40, VramOverheadGB: 1, ComputeBufFactor: 1}
s.applyDefaults()
got, err := EstimatePlan(s, meta, EstimateInput{Ctx: 32768, KvK: "q8_0", KvV: "q8_0"})
if err != nil {
t.Fatal(err)
}
// ComputeBufGB is runtime + ONE graph, so the graph the drafter copies is
// recoverable from it without having to re-derive the effective ubatch.
graph := got.ComputeBufGB - runtimeCtxGB()
if graph <= 0 {
t.Fatalf("compute buffer %.4f is below the runtime constant %.4f", got.ComputeBufGB, runtimeCtxGB())
}
slope := mtpDraftSlopeGB(meta, "q8_0", "q8_0")
want := mtpDraftPadGB + graph + slope*float64(got.Ctx)
if math.Abs(got.DraftGB-want) > 1e-9 {
t.Fatalf("draft charge=%.4f want %.4f (pad %.2f + graph %.4f + KV %.4f)",
got.DraftGB, want, mtpDraftPadGB, graph, slope*float64(got.Ctx))
}
}