mirror of
https://github.com/ollama/ollama.git
synced 2026-09-05 02:38:47 +00:00
132e0ca25d
Draft token embeddings were kept at source precision. A draft that reuses its embedding as the output projection (the gemma4 assistant) then reads the whole 537MB bf16 tensor on every draft step — about half the step's cost. Draft quality only affects how many drafts are accepted, so the output head now takes the requested type instead of the 8-bit type that protects a target's output quality. gemma4:26b-mlx, M5 Max: MTP code decode 148 -> 157 tok/s (+26% -> +37% over plain); prose goes from roughly zero to +2-5%; acceptance unchanged.
89 lines
3.2 KiB
Go
89 lines
3.2 KiB
Go
package create
|
|
|
|
import (
|
|
"fmt"
|
|
"strings"
|
|
)
|
|
|
|
// CreateDraftLayers imports a draft (speculative-decoding / MTP assistant)
|
|
// safetensors model into prefixed tensor and config blobs and returns the
|
|
// layers WITHOUT writing a manifest — the caller folds them into the target
|
|
// model's manifest. A draft never stands alone; it always accompanies a target
|
|
// model named on the Modelfile's FROM line.
|
|
//
|
|
// It runs the same read → classify → plan → write pipeline as Create. Output
|
|
// tensor names keep their source form, namespaced by tensorPrefix (e.g.
|
|
// "draft.") so they cannot collide with the target's tensors; config blobs are
|
|
// named under configPrefix (e.g. "draft/").
|
|
func CreateDraftLayers(modelDir, tensorPrefix, configPrefix, quantize string, store BlobStore, fn func(status string)) ([]LayerInfo, error) {
|
|
if tensorPrefix == "" {
|
|
return nil, fmt.Errorf("draft tensor prefix must not be empty")
|
|
}
|
|
if configPrefix == "" {
|
|
return nil, fmt.Errorf("draft config prefix must not be empty")
|
|
}
|
|
defer sweepMLX()
|
|
|
|
inv, err := ReadInventory(modelDir)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("read draft model: %w", err)
|
|
}
|
|
class, err := Classify(inv, quantize)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
policy, err := newTensorImportTransform(inv)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("build draft quantization policy for %q: %w", inv.Config.Architecture(), err)
|
|
}
|
|
specs, err := Plan(inv, class, draftPolicy{policy})
|
|
if err != nil {
|
|
return nil, fmt.Errorf("plan draft model: %w", err)
|
|
}
|
|
specs = prefixSpecs(specs, tensorPrefix)
|
|
|
|
fn(fmt.Sprintf("importing draft (%d tensors%s)", len(inv.Tensors), quantizeStatus(class)))
|
|
layers, err := WriteBlobs(specs, modelDir, store)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
configLayers, _, err := importConfigBlobs(modelDir, configPrefix, store, fn)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return append(layers, configLayers...), nil
|
|
}
|
|
|
|
// prefixSpecs returns specs with prefix prepended to every output blob name and
|
|
// output tensor name, leaving the source references (which point at the source
|
|
// files) untouched. Scale/bias keys derive from the tensor name, so they inherit
|
|
// the prefix automatically.
|
|
func prefixSpecs(specs []BlobSpec, prefix string) []BlobSpec {
|
|
out := make([]BlobSpec, len(specs))
|
|
for i, spec := range specs {
|
|
tensors := make([]TensorSpec, len(spec.Tensors))
|
|
for j, ts := range spec.Tensors {
|
|
ts.Name = prefix + ts.Name
|
|
tensors[j] = ts
|
|
}
|
|
out[i] = BlobSpec{Name: prefix + spec.Name, Tensors: tensors, Metadata: spec.Metadata}
|
|
}
|
|
return out
|
|
}
|
|
|
|
// draftPolicy wraps an architecture policy to give a draft model's output head
|
|
// (tied token embedding or separate lm_head) the requested type directly: draft
|
|
// quality only affects acceptance, so the target head's 8-bit promotion buys
|
|
// nothing. It is given unprefixed source names; planning runs before prefixSpecs.
|
|
type draftPolicy struct{ inner quantizePolicy }
|
|
|
|
func (p draftPolicy) quantizationType(name string, shape []int32, requested string) string {
|
|
if isEmbedTokensWeight(name) || strings.HasSuffix(name, "lm_head.weight") {
|
|
if q := normalizeQuantType(requested); isAligned(shape, q) {
|
|
return q
|
|
}
|
|
return ""
|
|
}
|
|
return p.inner.quantizationType(name, shape, requested)
|
|
}
|