base

type DraftModel interface {
	// Draft fuses b.Hidden (the target hidden state) into its own forward and
	// returns the head's hidden plus the projected hidden seeding the next step.
	Draft(b *batch.Batch, caches []cache.Cache) (hidden, projected *mlx.Array)

	// Unembed projects a hidden state to vocabulary logits.
	Unembed(x *mlx.Array) *mlx.Array

	// DraftCaches selects the draft model's own KV caches from the full
	// per-request slice — any subset, or nil when the draft keeps no KV.
	DraftCaches(caches []cache.Cache) []cache.Cache

	// LoadWeights assigns manifest tensors to the draft head's fields.
	LoadWeights(tensors map[string]*mlx.Array) error
}

type Model interface {
	Forward(b *batch.Batch, cache []cache.Cache) *mlx.Array
	Unembed(x *mlx.Array) *mlx.Array
	NumLayers() int
	Tokenizer() *tokenizer.Tokenizer
	MaxContextLength() int

	// LoadWeights receives all tensors loaded from the manifest and assigns
	// them to model fields. Model-specific logic (MLA absorption, expert
	// stacking, quantized layer creation) happens here.
	LoadWeights(tensors map[string]*mlx.Array) error
}

type SelfDraft interface {
	SelfDraft() DraftModel
}