ollama/create/draft.go

package create

import (
	"context"
	"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/"). store and fn must be non-nil.
func CreateDraftLayers(ctx context.Context, modelDir, tensorPrefix, configPrefix, quantize string, validation MLXValidationOptions, store BlobStore, fn func(status string)) ([]LayerInfo, error) {
	defer releaseMLXCache()
	return createDraftLayers(ctx, modelDir, tensorPrefix, configPrefix, quantize, validation, store, fn)
}

func createDraftLayers(ctx context.Context, modelDir, tensorPrefix, configPrefix, quantize string, validation MLXValidationOptions, 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")
	}
	if err := checkContext(ctx); err != nil {
		return nil, err
	}
	inv, err := ReadInventory(modelDir)
	if err != nil {
		return nil, fmt.Errorf("read draft model: %w", err)
	}
	if err := validateMLXSource(inv.Config, true, validation); err != nil {
		return nil, err
	}
	if err := checkContext(ctx); err != nil {
		return nil, 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(ctx, specs, modelDir, store)
	if err != nil {
		return nil, err
	}

	configLayers, _, err := importConfigBlobs(ctx, modelDir, configPrefix, nil, 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)
}