mirror of
https://github.com/ollama/ollama.git
synced 2026-09-22 22:55:24 -04:00
Jesse Gross
2e036e7cdf
The MLX runner is the only Go inference runner left and is no longer experimental, so its packages leave x/. The bindings become a top-level mlx package beside the carried patches in mlx/compat, mirroring how llama/ holds the llama.cpp integration, and the runner becomes mlxrunner with the architectures nested under the package they implement. Subpackages move with their parent unless listed. x/mlxrunner/mlx mlx x/internal/mlxthread mlx/mlxthread x/internal/mlxthreadtest mlx/mlxthread/mlxthreadtest x/internal/mlxtest mlx/mlxtest x/quant mlx/quant mlx/compat/*.patch mlx/compat/mlx-c (MLX patches go in mlx/compat/mlx) x/mlxrunner mlxrunner x/models/nn mlxrunner/nn x/models/<arch> mlxrunner/model/<arch> x/mlxrunner/imports.go mlxrunner/model/architectures (new package) x/create create x/safetensors fs/safetensors x/tokenizer mlxrunner/tokenizer Every package keeps its name, so the Go changes are the import path rewrites the moves force, and the CMake, Dockerfile, CI cache keys, drift check and Darwin payload script follow the new paths. Four edits are not paths: the runner's blank architecture imports become the package mlxrunner/model/architectures, so the list to extend for a new model sits beside the architecture directories; a depguard rule keeps the two test harnesses out of non-test code, as the x/internal placement used to; the CI change filter's two entries for the long-deleted x/imagegen/mlx now name the bindings' CMake project and the carried patches, so a change to either builds the payload; and the tokenizer parity test reads its fixtures from its own testdata instead of walking out of x/. x/server and x/imagegen/manifest stay for the next two commits.
80 lines
2.5 KiB
Go
80 lines
2.5 KiB
Go
// Package quant holds the quantization format facts shared by the model
|
|
// importer, the runtime loader, and `ollama show`. It deliberately has no
|
|
// dependency on the MLX C library, so any package can use it without pulling
|
|
// in cgo — which is what keeps these facts from drifting between separate
|
|
// hand-maintained copies.
|
|
package quant
|
|
|
|
import "strings"
|
|
|
|
type params struct {
|
|
groupSize int
|
|
bits int
|
|
mode string
|
|
}
|
|
|
|
// byType maps each canonical quantization type to its parameters. Aliases are
|
|
// resolved to a canonical name by Canonical before lookup.
|
|
var byType = map[string]params{
|
|
"nvfp4": {groupSize: 16, bits: 4, mode: "nvfp4"},
|
|
"mxfp4": {groupSize: 32, bits: 4, mode: "mxfp4"},
|
|
"int4": {groupSize: 64, bits: 4, mode: "affine"},
|
|
"mxfp8": {groupSize: 32, bits: 8, mode: "mxfp8"},
|
|
"int8": {groupSize: 64, bits: 8, mode: "affine"},
|
|
}
|
|
|
|
// Canonical returns the canonical name for a quantization type, resolving
|
|
// aliases (for example "FP8" and "Q8" both map to "int8"). It returns "" for
|
|
// the empty string and for any type it does not recognize.
|
|
func Canonical(quantType string) string {
|
|
switch strings.ToUpper(strings.TrimSpace(quantType)) {
|
|
case "NVFP4":
|
|
return "nvfp4"
|
|
case "MXFP4":
|
|
return "mxfp4"
|
|
case "MXFP8":
|
|
return "mxfp8"
|
|
case "INT4", "FP4", "Q4":
|
|
return "int4"
|
|
case "INT8", "FP8", "Q8":
|
|
return "int8"
|
|
default:
|
|
return ""
|
|
}
|
|
}
|
|
|
|
// Params returns the default group size, bit width, and mode for a
|
|
// quantization type. The empty string returns zeros. An unrecognized
|
|
// non-empty type falls back to 8-bit affine, matching the runtime loader's
|
|
// historical leniency toward unexpected metadata.
|
|
func Params(quantType string) (groupSize, bits int, mode string) {
|
|
if strings.TrimSpace(quantType) == "" {
|
|
return 0, 0, ""
|
|
}
|
|
if p, ok := byType[Canonical(quantType)]; ok {
|
|
return p.groupSize, p.bits, p.mode
|
|
}
|
|
return 32, 8, "affine"
|
|
}
|
|
|
|
// Bits returns the bit width of a recognized quantization type, or 0 if the
|
|
// type is empty or unrecognized. Unlike Params it applies no fallback, so
|
|
// callers that size or display tensors never act on an unknown type.
|
|
func Bits(quantType string) int {
|
|
if p, ok := byType[Canonical(quantType)]; ok {
|
|
return p.bits
|
|
}
|
|
return 0
|
|
}
|
|
|
|
// PackFactor returns how many quantized values are packed into one 32-bit
|
|
// word, or 0 for an empty or unrecognized type. MLX stores quantized weights
|
|
// packed into U32 words, so a tensor's logical last dimension is its stored
|
|
// last dimension times this factor.
|
|
func PackFactor(quantType string) int {
|
|
if b := Bits(quantType); b > 0 {
|
|
return 32 / b
|
|
}
|
|
return 0
|
|
}
|