Patrick Devine
15e6076d79
This change adds support for MTP (multi-token prediction) speculative decoding for the gemma4 model family. It includes: * support for importing safetensors based gemma4 draft models with `ollama create` * a new DRAFT command in the Modelfile for specifying draft models * a --quantize-draft flag for the ollama create command to quantize the draft model * cache support for speculation * changes to the rotating cache to be able to handle MTP correctly * sampling support for draft model token prediction --------- Co-authored-by: Daniel Hiltgen <daniel@ollama.com>
27 lines
560 B
Go
27 lines
560 B
Go
package mlx
|
|
|
|
// #include "generated.h"
|
|
import "C"
|
|
|
|
import "unsafe"
|
|
|
|
func (t *Array) Categorical(axis int) *Array {
|
|
key := New("")
|
|
out := New("")
|
|
C.mlx_random_categorical(&out.ctx, t.ctx, C.int(axis), key.ctx, DefaultStream().ctx)
|
|
return out
|
|
}
|
|
|
|
func Bernoulli(p *Array) *Array {
|
|
dims := p.Dims()
|
|
shape := make([]C.int, len(dims))
|
|
for i, d := range dims {
|
|
shape[i] = C.int(d)
|
|
}
|
|
|
|
key := New("")
|
|
out := New("BERNOULLI")
|
|
C.mlx_random_bernoulli(&out.ctx, p.ctx, unsafe.SliceData(shape), C.size_t(len(shape)), key.ctx, DefaultStream().ctx)
|
|
return out
|
|
}
|