diff --git a/README.md b/README.md index ab01e5f..cc923a6 100644 --- a/README.md +++ b/README.md @@ -114,6 +114,16 @@ The Anthropic-compatible API is available at `http://127.0.0.1:11500` and send the token from `rec auth token` as a bearer token or as `x-api-key`. The model is `flint`. +By default the server binds to loopback only. To expose it on your network for +other machines to reach, start it on all interfaces: + +```bash +rec serve --host 0.0.0.0 +``` + +The endpoint then listens on every interface; connect through the machine's IP +address, for example `http://192.168.1.42:11500/v1`. Keep the API key required +and restrict the interface with a firewall or VPN when binding to `0.0.0.0`. ## Commands diff --git a/cli/cli.go b/cli/cli.go index 49e42a6..7729222 100644 --- a/cli/cli.go +++ b/cli/cli.go @@ -91,11 +91,25 @@ func run(args []string) error { flags := flag.NewFlagSet(command, flag.ContinueOnError) manifest := flags.String("manifest", "", "") runtimeDir := flags.String("runtime-dir", "", "") + quant := flags.String("quant", QuantInt4, "") var contextTokens int + var tempVal float64 + var presenceVal float64 + var repeatVal float64 + var tempPtr *float64 + var presencePtr *float64 + var repeatPtr *float64 port := 11500 + host := "127.0.0.1" if command == "serve" { flags.IntVar(&contextTokens, "context", 0, "") flags.IntVar(&port, "port", port, "") + flags.StringVar(&host, "host", host, "") + flags.Float64Var(&tempVal, "temperature", 1.0, "") + flags.Float64Var(&presenceVal, "presence-penalty", 0.0, "") + flags.Float64Var(&repeatVal, "repeat-penalty", 1.0, "") + flags.Float64Var(&repeatVal, "repeat_penalty", 1.0, "") + flags.Float64Var(&repeatVal, "repetition-penalty", 1.0, "") } flags.SetOutput(io.Discard) if e := flags.Parse(rest); e != nil { @@ -107,14 +121,36 @@ func run(args []string) error { if port < 1 || port > 65535 { return fmt.Errorf("port must be between 1 and 65535") } - m, e := loadManifest(*manifest) + normalizedQuant, quantErr := validateQuant(*quant) + if quantErr != nil { + return quantErr + } + m, e := loadManifest(*manifest, normalizedQuant) if e == nil && command == "serve" { flags.Visit(func(f *flag.Flag) { if f.Name == "context" { m, e = withContext(m, contextTokens) } + if f.Name == "temperature" { + tempPtr = &tempVal + } + if f.Name == "presence-penalty" { + presencePtr = &presenceVal + } + if f.Name == "repeat-penalty" || f.Name == "repeat_penalty" || f.Name == "repetition-penalty" { + repeatPtr = &repeatVal + } }) } + if tempPtr != nil && (*tempPtr < 0 || *tempPtr > 2) { + return fmt.Errorf("temperature must be between 0 and 2") + } + if presencePtr != nil && (*presencePtr < -2 || *presencePtr > 2) { + return fmt.Errorf("presence-penalty must be between -2 and 2") + } + if repeatPtr != nil && (*repeatPtr <= 0 || *repeatPtr > 5) { + return fmt.Errorf("repeat-penalty must be greater than 0 and at most 5") + } if command == "doctor" { return doctor(p, m, e) } @@ -139,7 +175,7 @@ func run(args []string) error { } return e } - return serve(ctx, p, m, *runtimeDir, port, os.Stderr, background) + return serve(ctx, p, m, *runtimeDir, host, port, tempPtr, presencePtr, repeatPtr, os.Stderr, background) } func doctor(p Paths, m Manifest, manifestErr error) error { diff --git a/cli/context_test.go b/cli/context_test.go index 1d7a4b8..fe77c6c 100644 --- a/cli/context_test.go +++ b/cli/context_test.go @@ -53,32 +53,32 @@ func TestExistingInstanceUsesEffectiveProfile(t *testing.T) { m := Manifest{Model: "flint", Revision: "test", RuntimeSHA: "runtime", Context: 4096} s := Status{Ready: true, Model: m.Model, Revision: m.Revision, Runtime: m.RuntimeSHA, Context: m.Context, Instance: instanceKey(m), Endpoint: "http://127.0.0.1:11500/v1"} - if !sameInstance(s, m, 11500) { + if !sameInstance(s, m, "127.0.0.1", 11500) { t.Fatal("identical instance rejected") } explicit, _ := withContext(m, 4096) - if !sameInstance(s, explicit, 11500) { + if !sameInstance(s, explicit, "127.0.0.1", 11500) { t.Fatal("explicit default rejected") } changed := m changed.Context = 16384 - if sameInstance(s, changed, 11500) { + if sameInstance(s, changed, "127.0.0.1", 11500) { t.Fatal("reused a smaller context") } changed = m changed.MinRAM++ - if sameInstance(s, changed, 11500) { + if sameInstance(s, changed, "127.0.0.1", 11500) { t.Fatal("reused different serving settings") } s.Instance = "" - if sameInstance(s, m, 11500) { + if sameInstance(s, m, "127.0.0.1", 11500) { t.Fatal("legacy server without profile must restart") } } func TestContextReachesEngine(t *testing.T) { m := releaseManifest(t) - _, args, _ := launchCommand(Paths{}, m, "/runtime", "/model/flint-int4.json", 11500) + _, args, _ := launchCommand(Paths{}, m, "/runtime", "/model/flint-int4.json", "127.0.0.1", 11500, nil, nil, nil) found := false for i, arg := range args { if arg == "--context-limit" && i+1 < len(args) { diff --git a/cli/engine.go b/cli/engine.go index f0b7938..d0f2900 100644 --- a/cli/engine.go +++ b/cli/engine.go @@ -6,12 +6,23 @@ import ( "strconv" ) -func launchCommand(p Paths, m Manifest, runtimeDir, model string, port int) (string, []string, []string) { - return filepath.Join(runtimeDir, runtimeServer), []string{ +func launchCommand(p Paths, m Manifest, runtimeDir, model string, host string, port int, temp *float64, presence *float64, repeat *float64) (string, []string, []string) { + args := []string{ "serve", m.Model, "--model-dir", filepath.Dir(model), "--manifest", model, "--port", strconv.Itoa(port), "--context-limit", strconv.Itoa(m.Context), "--api-key-file", filepath.Join(p.Config, "api-key"), - }, []string{"HF_HUB_OFFLINE=1", "TOKENIZERS_PARALLELISM=false"} + "--host", host, + } + if temp != nil { + args = append(args, "--temperature", fmt.Sprintf("%g", *temp)) + } + if presence != nil { + args = append(args, "--presence-penalty", fmt.Sprintf("%g", *presence)) + } + if repeat != nil { + args = append(args, "--repeat-penalty", fmt.Sprintf("%g", *repeat)) + } + return filepath.Join(runtimeDir, runtimeServer), args, []string{"HF_HUB_OFFLINE=1", "TOKENIZERS_PARALLELISM=false"} } func checkAppleChip(name string) error { diff --git a/cli/help.go b/cli/help.go index 2bc4474..40a3642 100644 --- a/cli/help.go +++ b/cli/help.go @@ -14,6 +14,11 @@ type commandHelp struct { options [][2]string } +var quantOption = [2]string{"--quant FORMAT", "Model quantization: int4 (default, 48 GiB+ RAM) or int6 (64 GiB+ RAM)"} +var temperatureOption = [2]string{"--temperature NUM", "Sampling temperature between 0 and 2 (default: 1.0)"} +var presenceOption = [2]string{"--presence-penalty NUM", "Presence penalty between -2 and 2 (default: 0.0)"} +var repeatOption = [2]string{"--repeat-penalty NUM", "Repetition penalty greater than 0 (default: 1.0)"} + var developmentOptions = [][2]string{ {"--manifest FILE", "Use this manifest instead of the bundled one"}, {"--runtime-dir DIR", "Use this runtime instead of the bundled one; needs --manifest"}, @@ -22,17 +27,26 @@ var developmentOptions = [][2]string{ var commands = []commandHelp{ {"launch", "pi|opencode", "Start the server if needed and open a harness", [][2]string{ {"--args", "Print the harness command instead of running it; needs a running server"}, + temperatureOption, + presenceOption, + repeatOption, + quantOption, }}, {"exec", "pi|opencode ARGS", "Run a harness command, such as installing a plugin", nil}, {"serve", "[MODEL]", "Start the model server in this terminal", append([][2]string{ {"--context TOKENS", "Total input and output tokens (default: set by the release)"}, {"--port PORT", "Local port (default: 11500)"}, + {"--host HOST", "Network interface to bind the server to (default: 127.0.0.1)"}, + temperatureOption, + presenceOption, + repeatOption, + quantOption, }, developmentOptions...)}, {"status", "", "Show whether the server is running", nil}, {"stop", "", "Stop the server", nil}, {"auth token", "", "Print the local API key", nil}, - {"pull", "[MODEL]", "Download the model without starting the server", developmentOptions}, - {"doctor", "", "Check your hardware and installation", developmentOptions}, + {"pull", "[MODEL]", "Download the model without starting the server", append([][2]string{quantOption}, developmentOptions...)}, + {"doctor", "", "Check your hardware and installation", append([][2]string{quantOption}, developmentOptions...)}, {"update", "[-y]", "Update rec to the latest release", [][2]string{{"-y", "Do not ask for confirmation"}}}, {"uninstall", "[-y]", "Remove rec, its models and its settings", [][2]string{{"-y", "Do not ask for confirmation"}}}, {"--version", "", "Print the version", nil}, diff --git a/cli/launch.go b/cli/launch.go index 57b6e59..f35b05b 100644 --- a/cli/launch.go +++ b/cli/launch.go @@ -15,16 +15,132 @@ import ( "path/filepath" "regexp" "slices" + "strconv" "strings" "syscall" "time" ) -func launch(p Paths, args []string) error { - printOnly := len(args) > 0 && args[0] == "--args" - if printOnly { - args = args[1:] +func parseLaunchArgs(args []string) (printOnly bool, quant string, manifest string, temp *float64, presence *float64, repeat *float64, harness string, rest []string, err error) { + quant = QuantInt4 + beforeDash := args + afterDash := []string{} + for i, a := range args { + if a == "--" { + beforeDash = args[:i] + afterDash = args[i+1:] + break + } + } + var positional []string + for i := 0; i < len(beforeDash); i++ { + arg := beforeDash[i] + switch { + case arg == "--args": + printOnly = true + case arg == "--quant": + if i+1 >= len(beforeDash) { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("flag needs an argument: --quant") + } + i++ + quant = beforeDash[i] + case strings.HasPrefix(arg, "--quant="): + quant = strings.TrimPrefix(arg, "--quant=") + case arg == "--manifest": + if i+1 >= len(beforeDash) { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("flag needs an argument: --manifest") + } + i++ + manifest = beforeDash[i] + case strings.HasPrefix(arg, "--manifest="): + manifest = strings.TrimPrefix(arg, "--manifest=") + case arg == "--temperature": + if i+1 >= len(beforeDash) { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("flag needs an argument: --temperature") + } + i++ + val, err := strconv.ParseFloat(beforeDash[i], 64) + if err != nil || val < 0 || val > 2 { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("temperature must be between 0 and 2") + } + temp = &val + case strings.HasPrefix(arg, "--temperature="): + valStr := strings.TrimPrefix(arg, "--temperature=") + val, err := strconv.ParseFloat(valStr, 64) + if err != nil || val < 0 || val > 2 { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("temperature must be between 0 and 2") + } + temp = &val + case arg == "--presence-penalty": + if i+1 >= len(beforeDash) { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("flag needs an argument: --presence-penalty") + } + i++ + val, err := strconv.ParseFloat(beforeDash[i], 64) + if err != nil || val < -2 || val > 2 { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("presence-penalty must be between -2 and 2") + } + presence = &val + case strings.HasPrefix(arg, "--presence-penalty="): + valStr := strings.TrimPrefix(arg, "--presence-penalty=") + val, err := strconv.ParseFloat(valStr, 64) + if err != nil || val < -2 || val > 2 { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("presence-penalty must be between -2 and 2") + } + presence = &val + case arg == "--repeat-penalty" || arg == "--repeat_penalty" || arg == "--repetition-penalty": + if i+1 >= len(beforeDash) { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("flag needs an argument: %s", arg) + } + i++ + val, err := strconv.ParseFloat(beforeDash[i], 64) + if err != nil || val <= 0 || val > 5 { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("repeat-penalty must be greater than 0 and at most 5") + } + repeat = &val + case strings.HasPrefix(arg, "--repeat-penalty="): + valStr := strings.TrimPrefix(arg, "--repeat-penalty=") + val, err := strconv.ParseFloat(valStr, 64) + if err != nil || val <= 0 || val > 5 { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("repeat-penalty must be greater than 0 and at most 5") + } + repeat = &val + case strings.HasPrefix(arg, "--repeat_penalty="): + valStr := strings.TrimPrefix(arg, "--repeat_penalty=") + val, err := strconv.ParseFloat(valStr, 64) + if err != nil || val <= 0 || val > 5 { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("repeat-penalty must be greater than 0 and at most 5") + } + repeat = &val + case strings.HasPrefix(arg, "--repetition-penalty="): + valStr := strings.TrimPrefix(arg, "--repetition-penalty=") + val, err := strconv.ParseFloat(valStr, 64) + if err != nil || val <= 0 || val > 5 { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("repeat-penalty must be greater than 0 and at most 5") + } + repeat = &val + case strings.HasPrefix(arg, "-"): + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("flag provided but not defined: %s", arg) + default: + positional = append(positional, arg) + } + } + if len(positional) == 0 { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("missing harness; use rec launch pi or rec launch opencode") + } + harness = positional[0] + rest = append(positional[1:], afterDash...) + quant, err = validateQuant(quant) + if err != nil { + return false, "", "", nil, nil, nil, "", nil, err } + if harness != "opencode" && harness != "pi" { + return false, "", "", nil, nil, nil, "", nil, fmt.Errorf("unsupported harness %q; use rec launch pi or rec launch opencode", harness) + } + return printOnly, quant, manifest, temp, presence, repeat, harness, rest, nil +} + +func launch(p Paths, args []string) error { if len(args) == 0 || args[0] == "--help" || args[0] == "-h" || args[0] == "help" { printUsage(os.Stdout, "launch") fmt.Println("\nArguments after -- go to the harness. Reuses a running server, which keeps") @@ -32,21 +148,22 @@ func launch(p Paths, args []string) error { fmt.Println("and npm. The first launch installs the pinned harness.") return nil } - if args[0] != "opencode" && args[0] != "pi" { - return fmt.Errorf("unsupported harness %q; use rec launch pi or rec launch opencode", args[0]) - } - rest := args[1:] - if len(rest) > 0 && rest[0] == "--" { - rest = rest[1:] + printOnly, quant, manifest, temp, presence, repeat, targetHarness, rest, err := parseLaunchArgs(args) + if err != nil { + return err } if printOnly { - return printHarness(p, args[0], rest) + return printHarness(p, targetHarness, rest) } - if _, err := installClient(p, args[0]); err != nil { + if _, err := installClient(p, targetHarness); err != nil { return err } ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM, syscall.SIGHUP) - s, err := launchReady(ctx, p, os.Stderr) + folder := "" + if m, err := loadManifest(manifest, quant); err == nil && len(m.Files) > 0 { + folder = modelDigest(m) + } + s, err := launchReady(ctx, p, os.Stdout, quant, manifest, folder, temp, presence, repeat) cancel() if err != nil { return err @@ -55,11 +172,11 @@ func launch(p Paths, args []string) error { if err != nil { return err } - h, err := harnessCommand(p, args[0], s, os.Environ(), key, rest) + h, err := harnessCommand(p, targetHarness, s, os.Environ(), key, rest) if err != nil { return err } - openingClient(map[string]string{"pi": "Pi", "opencode": "OpenCode"}[args[0]], s) + openingClient(map[string]string{"pi": "Pi", "opencode": "OpenCode"}[targetHarness], s, quant, folder) return syscall.Exec(h.argv[0], h.argv, h.environ(os.Environ())) } @@ -163,8 +280,12 @@ func shellQuote(value string) string { return "'" + strings.ReplaceAll(value, "'", `'\''`) + "'" } -func openingClient(name string, s Status) { - newScreen(os.Stderr).step(accentStyle.Render("→"), "Opening "+name, fmt.Sprintf("Flint at %s, %s tokens", s.Endpoint, formatTokens(s.Context))) +func openingClient(name string, s Status, quant, folder string) { + detail := fmt.Sprintf("Flint (%s) at %s, %s tokens", quant, s.Endpoint, formatTokens(s.Context)) + if folder != "" { + detail = fmt.Sprintf("Flint (%s) (%s) at %s, %s tokens", quant, folder, s.Endpoint, formatTokens(s.Context)) + } + newScreen(os.Stdout).step(accentStyle.Render("→"), "Opening "+name, detail) } func openCodeArgs(args []string) []string { @@ -177,7 +298,7 @@ func openCodeArgs(args []string) []string { return append([]string{"--standalone"}, args...) } -func launchReady(ctx context.Context, p Paths, out io.Writer) (Status, error) { +func launchReady(ctx context.Context, p Paths, out io.Writer, quant, manifest, folder string, temp *float64, presence *float64, repeat *float64) (Status, error) { l, err := lock(filepath.Join(p.State, "launch.lock")) if err != nil { return Status{}, fmt.Errorf("another launcher is starting Flint; retry when it is ready") @@ -187,7 +308,25 @@ func launchReady(ctx context.Context, p Paths, out io.Writer) (Status, error) { if err == nil && s.Model != "flint" { return s, fmt.Errorf("running model is %q, expected flint; use rec stop first", s.Model) } + expectedRevision := "flint-" + quant if err == nil && s.Ready { + if s.Revision != "" && s.Revision != expectedRevision { + return s, fmt.Errorf("running server uses %s, requested %s; use rec stop first", s.Revision, expectedRevision) + } + if temp != nil && s.Temperature != nil && *s.Temperature != *temp { + return s, fmt.Errorf("running server uses temperature %g, requested %g; use rec stop first", *s.Temperature, *temp) + } + if presence != nil && s.PresencePenalty != nil && *s.PresencePenalty != *presence { + return s, fmt.Errorf("running server uses presence penalty %g, requested %g; use rec stop first", *s.PresencePenalty, *presence) + } + if repeat != nil && s.RepeatPenalty != nil && *s.RepeatPenalty != *repeat { + return s, fmt.Errorf("running server uses repeat penalty %g, requested %g; use rec stop first", *s.RepeatPenalty, *repeat) + } + readyMsg := fmt.Sprintf("Flint (%s) is ready", quant) + if folder != "" { + readyMsg = fmt.Sprintf("Flint (%s) (%s) is ready", quant, folder) + } + newScreen(out).succeeded(readyMsg, "") return s, nil } var child *exec.Cmd @@ -213,7 +352,7 @@ func launchReady(ctx context.Context, p Paths, out io.Writer) (Status, error) { if err != nil { // Complete first-time downloads while we own the terminal, before // detaching the server. Serving still verifies every file. - m, e := loadManifest("") + m, e := loadManifest(manifest, quant) if e != nil { return s, e } @@ -237,7 +376,23 @@ func launchReady(ctx context.Context, p Paths, out io.Writer) (Status, error) { if e = log.Chmod(0600); e != nil { return s, e } - child = exec.Command(self, "_background") + cmdArgs := []string{"_background"} + if quant != QuantInt4 { + cmdArgs = append(cmdArgs, "--quant", quant) + } + if manifest != "" { + cmdArgs = append(cmdArgs, "--manifest", manifest) + } + if temp != nil { + cmdArgs = append(cmdArgs, "--temperature", fmt.Sprintf("%g", *temp)) + } + if presence != nil { + cmdArgs = append(cmdArgs, "--presence-penalty", fmt.Sprintf("%g", *presence)) + } + if repeat != nil { + cmdArgs = append(cmdArgs, "--repeat-penalty", fmt.Sprintf("%g", *repeat)) + } + child = exec.Command(self, cmdArgs...) child.Stdout, child.Stderr = log, log child.SysProcAttr = &syscall.SysProcAttr{Setsid: true} if e = child.Start(); e != nil { @@ -247,7 +402,13 @@ func launchReady(ctx context.Context, p Paths, out io.Writer) (Status, error) { done = make(chan error, 1) go func() { done <- child.Wait() }() } - starting := startActivity(out, "Starting Flint", "Flint is ready") + label := fmt.Sprintf("Starting Flint (%s)", quant) + doneLabel := fmt.Sprintf("Flint (%s) is ready", quant) + if folder != "" { + label = fmt.Sprintf("Starting Flint (%s) (%s)", quant, folder) + doneLabel = fmt.Sprintf("Flint (%s) (%s) is ready", quant, folder) + } + starting := startActivity(out, label, doneLabel) // ReadAt does not share the child's output offset. Show the server's latest // lifecycle line as the phase and keep the whole log for diagnosis. var server bytes.Buffer diff --git a/cli/manifest.go b/cli/manifest.go index c193378..3675672 100644 --- a/cli/manifest.go +++ b/cli/manifest.go @@ -19,6 +19,22 @@ import ( //go:embed assets/* var assets embed.FS +const ( + QuantInt4 = "int4" + QuantInt6 = "int6" +) + +func validateQuant(q string) (string, error) { + norm := strings.ToLower(strings.TrimSpace(q)) + if norm == "" || norm == QuantInt4 { + return QuantInt4, nil + } + if norm == QuantInt6 { + return QuantInt6, nil + } + return "", fmt.Errorf("invalid quantization %q; supported options are int4 (default, 48 GiB+ RAM) and int6 (64 GiB+ RAM)", q) +} + var Version = "dev" var identifier = regexp.MustCompile(`^[a-zA-Z0-9][a-zA-Z0-9._-]*$`) @@ -134,13 +150,29 @@ func decodeManifest(data []byte) (Manifest, error) { return m, nil } -func loadManifest(path string) (Manifest, error) { +func loadManifest(path string, quant ...string) (Manifest, error) { + q := QuantInt4 + if len(quant) > 0 && quant[0] != "" { + var err error + q, err = validateQuant(quant[0]) + if err != nil { + return Manifest{}, err + } + } var data []byte var err error if path != "" { data, err = os.ReadFile(path) + } else if q == QuantInt6 { + data, err = assets.ReadFile("assets/manifest-int6.json") + if os.IsNotExist(err) { + return Manifest{}, fmt.Errorf("this CLI has no INT6 release bundle; build with packaging/package.py or supply --manifest for development") + } } else { - data, err = assets.ReadFile("assets/manifest.json") + data, err = assets.ReadFile("assets/manifest-int4.json") + if os.IsNotExist(err) { + data, err = assets.ReadFile("assets/manifest.json") + } if os.IsNotExist(err) { return Manifest{}, fmt.Errorf("this CLI has no release bundle; build with packaging/package.py or supply --manifest and --runtime-dir for development") } diff --git a/cli/quant_test.go b/cli/quant_test.go new file mode 100644 index 0000000..08eefbc --- /dev/null +++ b/cli/quant_test.go @@ -0,0 +1,211 @@ +package main + +import ( + "encoding/json" + "os" + "strings" + "testing" +) + +func TestQuantValidation(t *testing.T) { + for _, tc := range []struct { + input string + expected string + ok bool + }{ + {"", "int4", true}, + {"int4", "int4", true}, + {"INT4", "int4", true}, + {" int4 ", "int4", true}, + {"int6", "int6", true}, + {"INT6", "int6", true}, + {" int6 ", "int6", true}, + {"int8", "", false}, + {"fp16", "", false}, + {"invalid", "", false}, + } { + got, err := validateQuant(tc.input) + if (err == nil) != tc.ok { + t.Errorf("validateQuant(%q): got error %v, expected ok=%v", tc.input, err, tc.ok) + } + if tc.ok && got != tc.expected { + t.Errorf("validateQuant(%q) = %q, expected %q", tc.input, got, tc.expected) + } + } +} + +func TestParseLaunchArgs(t *testing.T) { + printOnly, quant, manifest, temp, presence, repeat, harness, rest, err := parseLaunchArgs([]string{"pi"}) + if err != nil || printOnly || quant != "int4" || manifest != "" || temp != nil || presence != nil || repeat != nil || harness != "pi" || len(rest) != 0 { + t.Fatalf("unexpected defaults: printOnly=%v quant=%s temp=%v presence=%v repeat=%v harness=%s rest=%v err=%v", printOnly, quant, temp, presence, repeat, harness, rest, err) + } + + printOnly, quant, manifest, temp, presence, repeat, harness, rest, err = parseLaunchArgs([]string{"--args", "--quant", "int6", "--temperature", "0.6", "--presence-penalty", "0.2", "--repeat-penalty", "1.1", "opencode", "--", "arg1", "arg2"}) + if err != nil || !printOnly || quant != "int6" || manifest != "" || temp == nil || *temp != 0.6 || presence == nil || *presence != 0.2 || repeat == nil || *repeat != 1.1 || harness != "opencode" || len(rest) != 2 || rest[0] != "arg1" || rest[1] != "arg2" { + t.Fatalf("unexpected parsed args: printOnly=%v quant=%s temp=%v presence=%v repeat=%v harness=%s rest=%v err=%v", printOnly, quant, temp, presence, repeat, harness, rest, err) + } + + // Flag with equals and after harness name + _, quant, _, temp, presence, repeat, harness, _, err = parseLaunchArgs([]string{"pi", "--quant=int6", "--temperature=0.7", "--presence-penalty=0.3", "--repeat-penalty=1.2"}) + if err != nil || quant != "int6" || temp == nil || *temp != 0.7 || presence == nil || *presence != 0.3 || repeat == nil || *repeat != 1.2 || harness != "pi" { + t.Fatalf("unexpected trailing flag parse: quant=%s temp=%v presence=%v repeat=%v harness=%s err=%v", quant, temp, presence, repeat, harness, err) + } + + // Flag aliases: repeat_penalty and repetition-penalty + _, _, _, _, _, repeat, _, _, err = parseLaunchArgs([]string{"pi", "--repeat_penalty=1.3"}) + if err != nil || repeat == nil || *repeat != 1.3 { + t.Fatalf("unexpected repeat_penalty parse: repeat=%v err=%v", repeat, err) + } + _, _, _, _, _, repeat, _, _, err = parseLaunchArgs([]string{"pi", "--repetition-penalty", "1.4"}) + if err != nil || repeat == nil || *repeat != 1.4 { + t.Fatalf("unexpected repetition-penalty parse: repeat=%v err=%v", repeat, err) + } + + // Invalid temperature + _, _, _, _, _, _, _, _, err = parseLaunchArgs([]string{"--temperature", "2.5", "pi"}) + if err == nil || !strings.Contains(err.Error(), "temperature must be between 0 and 2") { + t.Fatalf("expected error for invalid temperature, got %v", err) + } + + // Invalid presence penalty + _, _, _, _, _, _, _, _, err = parseLaunchArgs([]string{"--presence-penalty", "3.0", "pi"}) + if err == nil || !strings.Contains(err.Error(), "presence-penalty must be between -2 and 2") { + t.Fatalf("expected error for invalid presence-penalty, got %v", err) + } + + // Invalid repeat penalty + _, _, _, _, _, _, _, _, err = parseLaunchArgs([]string{"--repeat-penalty", "0.0", "pi"}) + if err == nil || !strings.Contains(err.Error(), "repeat-penalty must be greater than 0 and at most 5") { + t.Fatalf("expected error for invalid repeat-penalty, got %v", err) + } + _, _, _, _, _, _, _, _, err = parseLaunchArgs([]string{"--repeat-penalty", "6.0", "pi"}) + if err == nil || !strings.Contains(err.Error(), "repeat-penalty must be greater than 0 and at most 5") { + t.Fatalf("expected error for invalid repeat-penalty, got %v", err) + } + + // Invalid quant + _, _, _, _, _, _, _, _, err = parseLaunchArgs([]string{"--quant", "int8", "pi"}) + if err == nil || !strings.Contains(err.Error(), "invalid quantization") { + t.Fatalf("expected error for invalid quant, got %v", err) + } + + // Missing harness + _, _, _, _, _, _, _, _, err = parseLaunchArgs([]string{"--quant", "int6"}) + if err == nil || !strings.Contains(err.Error(), "missing harness") { + t.Fatalf("expected missing harness error, got %v", err) + } +} + +func TestReleaseManifestInt6(t *testing.T) { + b, err := os.ReadFile("../packaging/flint-int6-mlx.json") + if err != nil { + t.Fatal(err) + } + var m Manifest + if err := json.Unmarshal(b, &m); err != nil { + t.Fatal(err) + } + m.RuntimeSHA = digest([]byte("runtime")) + m.RuntimeFiles = []File{{Name: runtimeServer, Size: 1, SHA256: digest([]byte("r")), Executable: true}} + decoded, err := decodeManifest(mustJSON(t, m)) + if err != nil { + t.Fatal(err) + } + if decoded.Revision != "flint-int6" { + t.Fatalf("expected revision flint-int6, got %s", decoded.Revision) + } + expectedMinRAM := int64(64 * 1024 * 1024 * 1024) + if decoded.MinRAM != expectedMinRAM { + t.Fatalf("expected 64 GiB min RAM (%d), got %d", expectedMinRAM, decoded.MinRAM) + } + if decoded.Context != 262144 || decoded.MaxContext != 262144 { + t.Fatalf("expected context 262144, got %d / %d", decoded.Context, decoded.MaxContext) + } +} + +func TestLoadManifestWithQuant(t *testing.T) { + // Custom manifest path should load regardless of quant + _, err := loadManifest("../packaging/flint-int6-mlx.json", "int6") + if err == nil { + // Needs runtime descriptor + t.Fatal("expected failure without runtime descriptor") + } + + // Invalid quant should fail early + _, err = loadManifest("", "int9") + if err == nil || !strings.Contains(err.Error(), "invalid quantization") { + t.Fatalf("expected invalid quant error, got %v", err) + } +} + +func TestLaunchQuantLabels(t *testing.T) { + for _, quant := range []string{"int4", "int6"} { + folder := "ccd8b361c654d3544e686c53a057f266135fa3851d18da32efaf5bddfb92b25a" + label := "Starting Flint (" + quant + ") (" + folder + ")" + doneLabel := "Flint (" + quant + ") (" + folder + ") is ready" + if !strings.Contains(label, quant) || !strings.Contains(label, folder) || !strings.Contains(doneLabel, quant) || !strings.Contains(doneLabel, folder) { + t.Fatalf("labels do not contain quant or folder: %s / %s", label, doneLabel) + } + } +} + +func TestLaunchCommandSamplingParameters(t *testing.T) { + temp := 0.6 + presence := 0.2 + repeat := 1.1 + m := Manifest{Model: "flint", Revision: "flint-int6", Context: 262144} + _, args, _ := launchCommand(Paths{}, m, "/runtime", "/model/flint-int6.json", "127.0.0.1", 11500, &temp, &presence, &repeat) + + var foundTemp, foundPresence, foundRepeat bool + for i, a := range args { + if a == "--temperature" && i+1 < len(args) && args[i+1] == "0.6" { + foundTemp = true + } + if a == "--presence-penalty" && i+1 < len(args) && args[i+1] == "0.2" { + foundPresence = true + } + if a == "--repeat-penalty" && i+1 < len(args) && args[i+1] == "1.1" { + foundRepeat = true + } + } + if !foundTemp { + t.Fatalf("expected --temperature 0.6 in launchCommand args, got %v", args) + } + if !foundPresence { + t.Fatalf("expected --presence-penalty 0.2 in launchCommand args, got %v", args) + } + if !foundRepeat { + t.Fatalf("expected --repeat-penalty 1.1 in launchCommand args, got %v", args) + } +} + +func TestSameInstanceSamplingParameters(t *testing.T) { + temp1 := 0.6 + temp2 := 0.8 + presence1 := 0.2 + repeat1 := 1.1 + repeat2 := 1.3 + m := Manifest{Model: "flint", Revision: "flint-int6", RuntimeSHA: "runtime", Context: 262144} + s := Status{ + Ready: true, + Model: m.Model, + Revision: m.Revision, + Runtime: m.RuntimeSHA, + Context: m.Context, + Instance: instanceKey(m), + Endpoint: "http://127.0.0.1:11500/v1", + Temperature: &temp1, + PresencePenalty: &presence1, + RepeatPenalty: &repeat1, + } + + if !sameInstance(s, m, "127.0.0.1", 11500, &temp1, &presence1, &repeat1) { + t.Fatal("expected matching instance to be accepted") + } + if sameInstance(s, m, "127.0.0.1", 11500, &temp2, &presence1, &repeat1) { + t.Fatal("expected different temperature to be rejected") + } + if sameInstance(s, m, "127.0.0.1", 11500, &temp1, &presence1, &repeat2) { + t.Fatal("expected different repeat penalty to be rejected") + } +} diff --git a/cli/serve.go b/cli/serve.go index c4def68..c86c8ba 100644 --- a/cli/serve.go +++ b/cli/serve.go @@ -21,16 +21,19 @@ import ( ) type Status struct { - Engine string `json:"engine"` - Ready bool `json:"ready"` - Version string `json:"version"` - Model string `json:"model"` - Revision string `json:"revision"` - Runtime string `json:"runtime_sha256"` - Endpoint string `json:"endpoint"` - Context int `json:"context"` - MaxContext int `json:"max_context"` - Instance string `json:"instance_sha256"` + Engine string `json:"engine"` + Ready bool `json:"ready"` + Version string `json:"version"` + Model string `json:"model"` + Revision string `json:"revision"` + Runtime string `json:"runtime_sha256"` + Endpoint string `json:"endpoint"` + Context int `json:"context"` + MaxContext int `json:"max_context"` + Instance string `json:"instance_sha256"` + Temperature *float64 `json:"temperature,omitempty"` + PresencePenalty *float64 `json:"presence_penalty,omitempty"` + RepeatPenalty *float64 `json:"repeat_penalty,omitempty"` } func instanceKey(m Manifest) string { @@ -39,10 +42,48 @@ func instanceKey(m Manifest) string { return digest(b) } -func sameInstance(s Status, m Manifest, port int) bool { - return s.Ready && s.Model == m.Model && s.Runtime == m.RuntimeSHA && +// endpointHost maps a bind address to the address clients can dial. The wildcard +// address reports loopback so local clients keep a stable endpoint. +func endpointHost(host string) string { + if host == "" || host == "0.0.0.0" { + return "127.0.0.1" + } + return host +} + +func instanceEndpoint(host string, port int) string { + return fmt.Sprintf("http://%s:%d/v1", endpointHost(host), port) +} + +func sameInstance(s Status, m Manifest, host string, port int, params ...*float64) bool { + matches := s.Ready && s.Model == m.Model && s.Runtime == m.RuntimeSHA && s.Revision == m.Revision && s.Context == m.Context && - s.Instance == instanceKey(m) && s.Endpoint == fmt.Sprintf("http://127.0.0.1:%d/v1", port) + s.Instance == instanceKey(m) && s.Endpoint == instanceEndpoint(host, port) + if !matches { + return false + } + var temp *float64 + var presence *float64 + var repeat *float64 + if len(params) > 0 { + temp = params[0] + } + if len(params) > 1 { + presence = params[1] + } + if len(params) > 2 { + repeat = params[2] + } + if temp != nil && (s.Temperature == nil || *s.Temperature != *temp) { + return false + } + if presence != nil && (s.PresencePenalty == nil || *s.PresencePenalty != *presence) { + return false + } + if repeat != nil && (s.RepeatPenalty == nil || *s.RepeatPenalty != *repeat) { + return false + } + return true } func control(p Paths, action string) (Status, error) { @@ -210,11 +251,11 @@ func probe(ctx context.Context, base, key, model string, phase ...func(string)) return nil } -func serve(parent context.Context, p Paths, m Manifest, override string, port int, out io.Writer, background ...bool) error { +func serve(parent context.Context, p Paths, m Manifest, override string, host string, port int, temp *float64, presence *float64, repeat *float64, out io.Writer, background ...bool) error { l, e := lock(filepath.Join(p.State, "instance.lock")) if e != nil { if s, err := control(p, "status"); err == nil { - if sameInstance(s, m, port) { + if sameInstance(s, m, host, port, temp, presence, repeat) { newScreen(out).succeeded(m.Model+" is already running", s.Endpoint) return nil } @@ -225,13 +266,26 @@ func serve(parent context.Context, p Paths, m Manifest, override string, port in defer unlock(l) ctx, cancel := context.WithCancel(parent) defer cancel() - address := fmt.Sprintf("127.0.0.1:%d", port) + address := fmt.Sprintf("%s:%d", host, port) check, e := net.Listen("tcp", address) if e != nil { return fmt.Errorf("port %d is unavailable; use --port with a free port", port) } check.Close() - status := Status{Engine: m.Engine, Version: Version, Model: m.Model, Revision: m.Revision, Runtime: m.RuntimeSHA, Endpoint: "http://" + address + "/v1", Context: m.Context, MaxContext: m.contextLimit(), Instance: instanceKey(m)} + status := Status{ + Engine: m.Engine, + Version: Version, + Model: m.Model, + Revision: m.Revision, + Runtime: m.RuntimeSHA, + Endpoint: instanceEndpoint(host, port), + Context: m.Context, + MaxContext: m.contextLimit(), + Instance: instanceKey(m), + Temperature: temp, + PresencePenalty: presence, + RepeatPenalty: repeat, + } var mu sync.Mutex socket := filepath.Join(p.State, "control.sock") os.Remove(socket) @@ -284,7 +338,7 @@ func serve(parent context.Context, p Paths, m Manifest, override string, port in if e != nil { return e } - executable, args, extra := launchCommand(p, m, runtimeDir, model, port) + executable, args, extra := launchCommand(p, m, runtimeDir, model, host, port, temp, presence, repeat) self, e := os.Executable() if e != nil { return e @@ -324,7 +378,9 @@ func serve(parent context.Context, p Paths, m Manifest, override string, port in probeCtx, stopProbe := context.WithTimeout(ctx, 10*time.Minute) defer stopProbe() ready := make(chan error, 1) - go func() { ready <- probe(probeCtx, "http://"+address, key, m.Model, loading.phase) }() + go func() { + ready <- probe(probeCtx, fmt.Sprintf("http://%s:%d", endpointHost(host), port), key, m.Model, loading.phase) + }() select { case e = <-done: loading.fail() @@ -350,7 +406,7 @@ func serve(parent context.Context, p Paths, m Manifest, override string, port in mu.Lock() status.Ready = true mu.Unlock() - printReady(ui, m, status.Endpoint, len(background) > 0 && background[0]) + printReady(ui, m, status, len(background) > 0 && background[0]) select { case <-ctx.Done(): return nil @@ -360,15 +416,25 @@ func serve(parent context.Context, p Paths, m Manifest, override string, port in } } -func printReady(ui *screen, m Manifest, endpoint string, background bool) { +func printReady(ui *screen, m Manifest, s Status, background bool) { var rows strings.Builder - for i, row := range [][2]string{ - {"Base URL", endpoint}, + rowItems := [][2]string{ + {"Base URL", s.Endpoint}, {"Model", m.Model}, {"Checkpoint", m.Revision}, {"Context", m.contextSummary()}, - {"API key", accentStyle.Render("rec auth token")}, - } { + } + if s.Temperature != nil { + rowItems = append(rowItems, [2]string{"Temperature", fmt.Sprintf("%g", *s.Temperature)}) + } + if s.PresencePenalty != nil { + rowItems = append(rowItems, [2]string{"Presence", fmt.Sprintf("%g", *s.PresencePenalty)}) + } + if s.RepeatPenalty != nil { + rowItems = append(rowItems, [2]string{"Repeat penalty", fmt.Sprintf("%g", *s.RepeatPenalty)}) + } + rowItems = append(rowItems, [2]string{"API key", accentStyle.Render("rec auth token")}) + for i, row := range rowItems { if i > 0 { rows.WriteString("\n") } diff --git a/cli/serve_test.go b/cli/serve_test.go index 050d204..4652c70 100644 --- a/cli/serve_test.go +++ b/cli/serve_test.go @@ -95,7 +95,7 @@ func TestServeLifecycle(t *testing.T) { var out lockedBuffer done := make(chan error, 1) - go func() { done <- serve(context.Background(), p, m, runtimeDir, port, &out) }() + go func() { done <- serve(context.Background(), p, m, runtimeDir, "127.0.0.1", port, nil, nil, nil, &out) }() deadline := time.Now().Add(2 * time.Minute) var s Status for !s.Ready { @@ -115,12 +115,12 @@ func TestServeLifecycle(t *testing.T) { } var again bytes.Buffer - if err := serve(context.Background(), p, m, runtimeDir, port, &again); err != nil || !strings.Contains(again.String(), "test is already running") { + if err := serve(context.Background(), p, m, runtimeDir, "127.0.0.1", port, nil, nil, nil, &again); err != nil || !strings.Contains(again.String(), "test is already running") { t.Fatalf("identical run did not reuse the server: %v %s", err, again.String()) } changed := m changed.Context, changed.MaxContext = 4096, 4096 - if err := serve(context.Background(), p, changed, runtimeDir, port, io.Discard); err == nil { + if err := serve(context.Background(), p, changed, runtimeDir, "127.0.0.1", port, nil, nil, nil, io.Discard); err == nil { t.Fatal("a different context reused the running server") } diff --git a/cli/testdata/fake-rec-mlx b/cli/testdata/fake-rec-mlx index 8a56828..0ebbef1 100755 --- a/cli/testdata/fake-rec-mlx +++ b/cli/testdata/fake-rec-mlx @@ -12,6 +12,10 @@ for flag in ("--model-dir", "--manifest", "--api-key-file"): parser.add_argument(flag, type=Path, required=True) parser.add_argument("--port", type=int, required=True) parser.add_argument("--context-limit", type=int, required=True) +parser.add_argument("--host", default="127.0.0.1") +parser.add_argument("--temperature", type=float) +parser.add_argument("--presence-penalty", type=float) +parser.add_argument("--repeat-penalty", type=float) args = parser.parse_args() assert (args.model_dir / args.manifest.name).is_file() key = args.api_key_file.read_text().strip() @@ -46,4 +50,4 @@ class Handler(http.server.BaseHTTPRequestHandler): self.reply({"choices": [{"index": 0, "message": message, "finish_reason": "stop"}]}) -http.server.ThreadingHTTPServer(("127.0.0.1", args.port), Handler).serve_forever() +http.server.ThreadingHTTPServer((args.host, args.port), Handler).serve_forever() diff --git a/engine/src/rec_mlx/chat.py b/engine/src/rec_mlx/chat.py index 868f0d4..c216907 100644 --- a/engine/src/rec_mlx/chat.py +++ b/engine/src/rec_mlx/chat.py @@ -108,6 +108,7 @@ def seeded_tool_parser(sampling, prompt_tokens): "top_p": float(sampling.top_p), "top_k": sampling.top_k, "presence_penalty": float(sampling.presence_penalty), + "repetition_penalty": float(getattr(sampling, "repetition_penalty", 1.0)), "max_tokens": sampling.max_tokens, }, separators=(",", ":"), diff --git a/engine/src/rec_mlx/checkpoint.py b/engine/src/rec_mlx/checkpoint.py index 9f341c9..f00ac5d 100644 --- a/engine/src/rec_mlx/checkpoint.py +++ b/engine/src/rec_mlx/checkpoint.py @@ -3,11 +3,40 @@ def validate(directory: Path, manifest_path: Path) -> dict: - manifest = json.loads(manifest_path.read_text()) - if manifest["weight_format"] != "mlx-affine-int4" or manifest["kv_format"] != "bf16": - raise ValueError("This engine requires affine INT4 weights and BF16 KV") + raw = json.loads(manifest_path.read_text()) + config = json.loads((directory / "config.json").read_text()) + if "weight_format" in raw: + manifest = raw + else: + # Standard Hugging Face MLX checkpoint where manifest is config.json + q = config.get("quantization", {}) + # Check if routed experts specify 6-bit or if base is 6-bit + bits = None + for k, v in q.items(): + if isinstance(v, dict) and "bits" in v and "switch_mlp" in k: + bits = v["bits"] + break + if bits is None: + bits = q.get("bits", 4) + shards = {p.name: "" for p in directory.glob("model-*.safetensors")} + manifest = { + "model_id": f"flint-int{bits}", + "weight_format": f"mlx-affine-int{bits}", + "kv_format": "bf16", + "max_context": config.get("max_position_embeddings", 262144), + "checkpoint_files": { + "config.json": "", + "tokenizer.json": "", + "tokenizer_config.json": "", + **shards, + }, + } + if manifest["weight_format"] not in ("mlx-affine-int4", "mlx-affine-int6") or manifest["kv_format"] != "bf16": + raise ValueError("This engine requires affine INT4 or INT6 weights and BF16 KV") files = manifest["checkpoint_files"] - required = {"config.json", "tokenizer.json", "tokenizer_config.json", "special_tokens_map.json"} + required = {"config.json", "tokenizer.json", "tokenizer_config.json"} + if "special_tokens_map.json" in files: + required.add("special_tokens_map.json") shards = {name for name in files if name.startswith("model-") and name.endswith(".safetensors")} if not required.issubset(files) or not shards: raise ValueError("Manifest must include configuration, tokenizer and weight shards") @@ -20,12 +49,29 @@ def validate(directory: Path, manifest_path: Path) -> dict: raise ValueError(f"Checkpoint filename must be a basename: {name}") if not (directory / name).is_file(): raise ValueError(f"Checkpoint is missing a manifest file: {name}") - config = json.loads((directory / "config.json").read_text()) if config["sliding_window"] != 513: raise ValueError("Flint requires sliding_window=513, including the current token") - export = config["mlx_export"] - if (export["bits"], export["routed_expert_bits"], export["embedding_bits"]) != (4, 4, 8): - raise ValueError("Expected INT4 projections/experts and INT8 embeddings") - if config["quantization"] != {"group_size": 64, "bits": 4}: - raise ValueError("Expected group-64 INT4 quantization") + expected_bits = 6 if manifest["weight_format"] == "mlx-affine-int6" else 4 + if expected_bits == 6: + # Check if routed experts are 6-bit or mlx_export specifies 6-bit routed experts + export = config.get("mlx_export") + if export is not None and export.get("routed_expert_bits") != 6: + raise ValueError("Expected INT6 routed experts") + q = config.get("quantization", {}) + has_6bit = any( + isinstance(v, dict) and v.get("bits") == 6 + for k, v in q.items() + ) or q.get("bits") == 6 + if not has_6bit: + raise ValueError("Expected INT6 quantization for routed experts or model") + else: + export = config.get("mlx_export", { + "bits": 4, + "routed_expert_bits": 4, + "embedding_bits": 8, + }) + if (export.get("bits"), export.get("routed_expert_bits"), export.get("embedding_bits")) != (4, 4, 8): + raise ValueError("Expected INT4 projections/experts and INT8 embeddings") + if config.get("quantization", {}).get("bits") != 4: + raise ValueError("Expected group-64 INT4 quantization") return manifest diff --git a/engine/src/rec_mlx/cli.py b/engine/src/rec_mlx/cli.py index 7b04258..6c169fb 100644 --- a/engine/src/rec_mlx/cli.py +++ b/engine/src/rec_mlx/cli.py @@ -33,7 +33,12 @@ def main(): help="total input and output tokens (default: %(default)s)", ) parser.add_argument( - "--port", type=int, default=8080, help="loopback port for serve (default: %(default)s)" + "--port", type=int, default=8080, help="port for serve (default: %(default)s)" + ) + parser.add_argument( + "--host", + default="127.0.0.1", + help="interface for the serve HTTP server to bind (default: %(default)s)", ) parser.add_argument( "--grammar", @@ -60,6 +65,21 @@ def main(): default=1.0, help="default sampling temperature (default: %(default)s)", ) + parser.add_argument( + "--presence-penalty", + type=float, + default=0.0, + help="default presence penalty (default: %(default)s)", + ) + parser.add_argument( + "--repeat-penalty", + "--repeat_penalty", + "--repetition-penalty", + type=float, + default=1.0, + dest="repetition_penalty", + help="default repetition penalty (default: %(default)s)", + ) parser.add_argument( "--top-p", type=float, @@ -88,6 +108,8 @@ def main(): temperature=args.temperature, top_p=args.top_p, top_k=args.top_k, + presence_penalty=args.presence_penalty, + repetition_penalty=args.repetition_penalty, max_tokens=args.max_tokens, ) except ValueError as error: @@ -123,13 +145,13 @@ def main(): server = make_server( engine, - ("127.0.0.1", args.port), + (args.host, args.port), args.metrics, defaults=settings, api_key=api_key, grammar_mode=args.grammar, ) - print(f"Listening on http://127.0.0.1:{server.server_port}/v1", flush=True) + print(f"Listening on http://{args.host}:{server.server_port}/v1", flush=True) def stop(signum, frame): server.generator.cancel() diff --git a/engine/src/rec_mlx/engine.py b/engine/src/rec_mlx/engine.py index 15ad765..d64276c 100644 --- a/engine/src/rec_mlx/engine.py +++ b/engine/src/rec_mlx/engine.py @@ -27,9 +27,10 @@ class Sampling: max_tokens: int = 1024 seed: int = 0 presence_penalty: float = 0.0 + repetition_penalty: float = 1.0 def __post_init__(self): - for name in ("temperature", "top_p", "presence_penalty"): + for name in ("temperature", "top_p", "presence_penalty", "repetition_penalty"): if type(getattr(self, name)) not in (int, float): raise ValueError(f"{name} must be a number") for name in ("top_k", "max_tokens", "seed"): @@ -39,6 +40,8 @@ def __post_init__(self): raise ValueError("temperature must be in [0, 2] and top_p in (0, 1]") if not -2 <= self.presence_penalty <= 2: raise ValueError("presence_penalty must be in [-2, 2]") + if not 0 < self.repetition_penalty <= 5: + raise ValueError("repetition_penalty must be in (0, 5]") if self.top_k < 0 or self.max_tokens < 1: raise ValueError("top_k must be nonnegative and max_tokens positive") if not 0 <= self.seed <= 2**64 - self.max_tokens: @@ -57,13 +60,14 @@ def __init__( if platform.system() != "Darwin" or platform.machine() != "arm64": raise RuntimeError("Flint requires an Apple Silicon Mac") memory = int(subprocess.check_output(["sysctl", "-n", "hw.memsize"], text=True)) - if memory < 48 * 1024**3: - raise RuntimeError("This build targets Macs with at least 48 GiB of memory") + self.identity = validate(directory, manifest_path) + min_memory = 64 * 1024**3 if self.identity.get("weight_format") == "mlx-affine-int6" else 48 * 1024**3 + if memory < min_memory: + raise RuntimeError(f"This model requires a Mac with at least {min_memory >> 30} GiB of memory; detected {memory >> 30} GiB") import mlx.core as mx device = mx.device_info() profile = profile_for_device(device) - self.identity = validate(directory, manifest_path) config = json.loads((directory / "config.json").read_text()) from .kernels import build_decoder from .quantized import QATModel @@ -87,7 +91,12 @@ def __init__( mx.set_cache_limit(0) mx.clear_cache() mx.reset_peak_memory() - self.runtime, self.resident_sampler = build_decoder(QATModel(directory), optimized=True) + if self.identity.get("weight_format") == "mlx-affine-int6": + from .native_engine import build_native_decoder + + self.runtime, self.resident_sampler = build_native_decoder(QATModel(directory)) + else: + self.runtime, self.resident_sampler = build_decoder(QATModel(directory), optimized=True) # Packing replaces native weights with lazy views. Materialize them on their creation # stream before HTTP threads use the native prefill model. mx.eval(self.runtime.model.native.parameters()) @@ -264,6 +273,8 @@ def _generate_tokens(self, ids, sampling, *, should_stop, grammar=None, return_l text_stream = StreamingText(self.tokenizer) if sampling.presence_penalty: session._state["presence_mask"] = mx.zeros(logits.size, mx.float32) + if sampling.repetition_penalty != 1.0: + session._state["recent_tokens"] = list(ids[-64:]) mx.eval(logits) prefill_end = time.perf_counter() prefill_peak = mx.get_peak_memory() @@ -287,6 +298,7 @@ def _generate_tokens(self, ids, sampling, *, should_stop, grammar=None, return_l top_p=sampling.top_p, top_k=sampling.top_k, presence_penalty=sampling.presence_penalty, + repetition_penalty=sampling.repetition_penalty, grammar=grammar, **options, ) @@ -296,6 +308,14 @@ def _generate_tokens(self, ids, sampling, *, should_stop, grammar=None, return_l if not bool(mx.all(mx.isfinite(logits)).item()): raise FloatingPointError("The inference engine produced nonfinite logits") flat = logits.reshape(-1).astype(mx.float32) + if sampling.repetition_penalty != 1.0 and session._state.get("recent_tokens"): + indices = mx.array(list(dict.fromkeys(session._state["recent_tokens"][-64:]))) + selected = flat[indices] + flat[indices] = mx.where( + selected < 0, + selected * sampling.repetition_penalty, + selected / sampling.repetition_penalty, + ) if grammar is not None: flat = grammar.first_logits(flat) if sampling.temperature == 0: @@ -315,6 +335,8 @@ def _generate_tokens(self, ids, sampling, *, should_stop, grammar=None, return_l computed_tokens += 1 if sampling.presence_penalty: session._state["presence_mask"][token] = 1.0 + if "recent_tokens" in session._state: + session._state["recent_tokens"].append(token) if first is None: first = time.perf_counter() generated.append(token) @@ -377,7 +399,9 @@ def _generate_tokens(self, ids, sampling, *, should_stop, grammar=None, return_l "peak_mlx_bytes": mx.get_peak_memory(), "process_peak_rss_bytes": resource.getrusage(resource.RUSAGE_SELF).ru_maxrss, "computed_tokens": computed_tokens, - "decode_strategy": "resident_sampling", + "decode_strategy": "native_sampling" + if self.identity.get("weight_format") == "mlx-affine-int6" + else "resident_sampling", **(grammar.metrics() if grammar is not None else {}), }, } diff --git a/engine/src/rec_mlx/kernels/sampling.py b/engine/src/rec_mlx/kernels/sampling.py index 6a83633..4d45adc 100644 --- a/engine/src/rec_mlx/kernels/sampling.py +++ b/engine/src/rec_mlx/kernels/sampling.py @@ -70,6 +70,7 @@ def block( top_p=1.0, top_k=0, presence_penalty=0.0, + repetition_penalty=1.0, grammar=None, logprobs=None, ): diff --git a/engine/src/rec_mlx/model.py b/engine/src/rec_mlx/model.py index 7e707f0..dc7ef6f 100644 --- a/engine/src/rec_mlx/model.py +++ b/engine/src/rec_mlx/model.py @@ -333,6 +333,10 @@ def layers(self): def sanitize(self, weights: dict[str, Any]) -> dict[str, Any]: out = {} for key, value in weights.items(): + if key.startswith("language_model."): + key = key[len("language_model.") :] + if key.endswith(".mlp.expert_bias"): + key = key.replace(".mlp.expert_bias", ".mlp.gate.expert_bias") if key.endswith(".self_attn.softmax_offset"): value = value.astype(mx.float32) elif key.endswith(".weight") and value.dtype == mx.uint8: diff --git a/engine/src/rec_mlx/native_engine.py b/engine/src/rec_mlx/native_engine.py new file mode 100644 index 0000000..6fcf110 --- /dev/null +++ b/engine/src/rec_mlx/native_engine.py @@ -0,0 +1,156 @@ +import time +from types import SimpleNamespace +import mlx.core as mx +from mlx_lm.sample_utils import make_sampler + +from .cache import KVStorage +from .kernels.device import profile_for_device +from .quantized import QATSession + + +class NativeEngine: + def __init__(self, model): + self.model = model + device = mx.device_info() + profile = profile_for_device(device) + self.device_profile = profile + self.prefill_chunk_size = profile.prefill_chunk_size + self.workers = profile.workers + + # Metadata marks global attention (is_global == 1) vs sliding attention (0) + self.metadata = [[0] * 64 for _ in range(50)] + for i, layer in enumerate(model.layers): + self.metadata[i][1] = int(layer.attention.is_global) + + def prefill_state(self, capacity, *, limit=None, storage=None): + if storage is None: + storage = KVStorage(self.metadata, capacity) + else: + storage.grow(capacity) + return { + "usable": True, + "kv": storage, + "metadata": storage.metadata, + "position": 0, + "capacity": storage.capacity, + "limit": capacity if limit is None else limit, + "growth_count": 0, + "growth_seconds": 0.0, + } + + def ensure_capacity(self, state, count): + if not state["usable"]: + raise RuntimeError("Session was invalidated by a GPU synchronization failure") + required = state["position"] + count + if required > state["limit"]: + raise ValueError("Decode context limit exhausted") + if required > state["capacity"]: + started = time.perf_counter() + state["kv"].grow(required) + state["capacity"] = state["kv"].capacity + state["growth_count"] += 1 + state["growth_seconds"] += time.perf_counter() - started + + def make_session(self, max_new_tokens): + return QATSession(self, max_new_tokens) + + def forward(self, state, token): + self.ensure_capacity(state, 1) + inputs = mx.array([[token]], dtype=mx.int32) + cache = state.get("cache") + hidden = self.model.native.model(inputs, cache=cache) + logits = self.model.native.lm_head(hidden[:, -1:]) + state["position"] += 1 + return logits, None + + +class NativeResidentSampler: + def __init__(self, engine): + self.engine = engine + + def block( + self, + state, + token, + count, + seed, + *, + temperature=1.0, + top_p=1.0, + top_k=0, + presence_penalty=0.0, + repetition_penalty=1.0, + grammar=None, + logprobs=None, + ): + packet = [] + unfiltered_t1 = (temperature, top_p, top_k) == (1.0, 1.0, 0) + sampler = ( + make_sampler(temp=temperature, top_p=top_p, top_k=top_k) + if not unfiltered_t1 and temperature != 0 + else None + ) + eos_tokens = set(self.engine.model.eos_tokens) + stops = eos_tokens if grammar is None else set(grammar.stops) + current = token + + for step in range(count): + logits = self.engine.forward(state, current)[0] + flat = logits.reshape(-1).astype(mx.float32) + + if repetition_penalty != 1.0 and "recent_tokens" in state and state["recent_tokens"]: + indices = mx.array(list(dict.fromkeys(state["recent_tokens"][-64:]))) + selected = flat[indices] + flat[indices] = mx.where( + selected < 0, + selected * repetition_penalty, + selected / repetition_penalty, + ) + + if presence_penalty and "presence_mask" in state: + flat = flat - presence_penalty * state["presence_mask"] + + if grammar is not None and "grammar_mask" in state: + gmask = state["grammar_mask"] + words = mx.arange(flat.size, dtype=mx.uint32) + allowed = ((gmask[words >> 5] >> (words & 31)) & 1) != 0 + flat = mx.where(allowed, flat, -float("inf")) + + if temperature == 0: + next_tok = int(mx.argmax(flat).item()) + elif unfiltered_t1: + next_tok = int( + mx.random.categorical(flat, key=mx.random.key(seed + step)).item() + ) + else: + next_tok = int(sampler(flat - mx.logsumexp(flat)).item()) + + if logprobs is not None: + report = flat / temperature if temperature else flat + logprob_val = float((report[next_tok] - mx.logsumexp(report)).item()) + logprobs.append(logprob_val) + + if presence_penalty and "presence_mask" in state: + state["presence_mask"][next_tok] = 1.0 + + if "recent_tokens" in state: + state["recent_tokens"].append(next_tok) + + if grammar is not None: + import numpy as np + + mask = grammar._advance(next_tok, True) + state["grammar_mask"] = mx.array(np.concatenate([mask, np.array([1], np.uint32)])) + + packet.append(next_tok) + current = next_tok + if next_tok in stops: + break + + return packet + + +def build_native_decoder(model): + engine = NativeEngine(model) + sampler = NativeResidentSampler(engine) + return engine, sampler diff --git a/engine/src/rec_mlx/quantized.py b/engine/src/rec_mlx/quantized.py index a204421..cf6da54 100644 --- a/engine/src/rec_mlx/quantized.py +++ b/engine/src/rec_mlx/quantized.py @@ -18,20 +18,45 @@ class QATModel: def __init__(self, directory): self.config = json.loads((directory / "config.json").read_text()) - export = self.config["mlx_export"] - assert export["bits"] == 4 and export["routed_expert_bits"] in (3, 4) + q_cfg = self.config.get("quantization", {}) + default_bits = q_cfg.get("bits", 4) + default_group_size = q_cfg.get("group_size", 64) + + routed_bits = default_bits + for k, v in q_cfg.items(): + if isinstance(v, dict) and "bits" in v and "switch_mlp" in k: + routed_bits = v["bits"] + break + + export = self.config.get("mlx_export", { + "bits": default_bits, + "routed_expert_bits": routed_bits, + "embedding_bits": 8, + }) + self.config["mlx_export"] = export + assert export["bits"] in (4, 6, 8) and export["routed_expert_bits"] in (3, 4, 6) assert export["embedding_bits"] == 8 native = Model(ModelArgs.from_dict(self.config)) def quantization(path, module): if not hasattr(module, "to_quantized"): return False - bits = 4 + lookup = path + if lookup not in q_cfg and f"language_model.{lookup}" in q_cfg: + lookup = f"language_model.{lookup}" + if lookup in q_cfg and isinstance(q_cfg[lookup], dict): + spec = q_cfg[lookup] + return { + "group_size": spec.get("group_size", default_group_size), + "bits": spec.get("bits", default_bits), + "mode": "affine", + } + bits = export["bits"] if path == "model.embed_tokens": bits = 8 elif ".switch_mlp." in path: bits = export["routed_expert_bits"] - return {"group_size": 64, "bits": bits, "mode": "affine"} + return {"group_size": default_group_size, "bits": bits, "mode": "affine"} nn.quantize(native, class_predicate=quantization) weights = {} @@ -144,6 +169,7 @@ def prefill( state["kv"].write(index, 0, keys, values) state["capacity"] = state["kv"].capacity state["position"] = state["initial"] = len(tokens) + state["cache"] = cache self._state = state return logits diff --git a/engine/src/rec_mlx/scheduler.py b/engine/src/rec_mlx/scheduler.py index 17e5a94..1213686 100644 --- a/engine/src/rec_mlx/scheduler.py +++ b/engine/src/rec_mlx/scheduler.py @@ -136,11 +136,15 @@ def _serve_single(self, queued_request): self.busy.clear() def _start_request(self, request, args): + rep = getattr(args.logits, "repetition_penalty", None) + if rep is None: + rep = self.defaults.repetition_penalty sampling = Sampling( temperature=args.sampling.temperature, top_p=args.sampling.top_p, top_k=args.sampling.top_k, presence_penalty=args.logits.presence_penalty, + repetition_penalty=rep, max_tokens=args.max_tokens, seed=self.defaults.seed if args.seed is None else args.seed, ) diff --git a/engine/src/rec_mlx/server.py b/engine/src/rec_mlx/server.py index ee4f3b8..d250380 100644 --- a/engine/src/rec_mlx/server.py +++ b/engine/src/rec_mlx/server.py @@ -198,6 +198,7 @@ def drop_unknown_fields(self): "stop", "frequency_penalty", "repetition_penalty", + "repeat_penalty", "presence_penalty", "logprobs", "min_p", @@ -208,6 +209,8 @@ def drop_unknown_fields(self): unknown = set(self.body) - allowed for field in unknown: del self.body[field] + if "repeat_penalty" in self.body and "repetition_penalty" not in self.body: + self.body["repetition_penalty"] = self.body.pop("repeat_penalty") with ignored_fields_lock: for field in sorted(unknown - ignored_fields): logging.warning("Ignoring unsupported request field %r", field) @@ -235,7 +238,8 @@ def validate_request(self): if self.body.get(key, 0) != 0: raise ValueError(f"Nonzero {key} is not supported by this decoder") if self.body.get("repetition_penalty", 1) != 1: - raise ValueError("Only neutral repetition_penalty=1 is supported") + if engine.identity.get("weight_format") != "mlx-affine-int6": + raise ValueError("repetition_penalty is only supported for INT6 models") if self.body.get("store", False) is not False: raise ValueError("Stored-completion retrieval is not supported; use store=false") stop = self.body.get("stop") @@ -302,11 +306,18 @@ def handle_completion(self, request, stop_words): self.max_tokens = ( remaining if explicit_maximum is None else min(self.max_tokens, remaining) ) + if "temperature" not in self.body and defaults.temperature is not None: + self.temperature = defaults.temperature + if "presence_penalty" not in self.body and defaults.presence_penalty is not None: + self.presence_penalty = defaults.presence_penalty + if "repetition_penalty" not in self.body and defaults.repetition_penalty is not None: + self.repetition_penalty = defaults.repetition_penalty Sampling( temperature=self.temperature, top_p=self.top_p, top_k=self.top_k, presence_penalty=self.presence_penalty, + repetition_penalty=getattr(self, "repetition_penalty", 1.0), max_tokens=self.max_tokens, seed=defaults.seed if self.seed is None else self.seed, ) diff --git a/engine/tests/test_chat.py b/engine/tests/test_chat.py index 8523123..113eb0e 100644 --- a/engine/tests/test_chat.py +++ b/engine/tests/test_chat.py @@ -185,6 +185,7 @@ def sample(**overrides): "top_p": 0.95, "top_k": 20, "presence_penalty": 0.0, + "repetition_penalty": 1.0, "max_tokens": 16, "seed": 42, } @@ -206,6 +207,7 @@ def ids(prompt, **overrides): ("top_p", 0.5), ("top_k", 4), ("presence_penalty", 1.5), + ("repetition_penalty", 1.2), ("max_tokens", 8), ): with self.subTest(name=name): diff --git a/engine/tests/test_checkpoint.py b/engine/tests/test_checkpoint.py new file mode 100644 index 0000000..b05fbfe --- /dev/null +++ b/engine/tests/test_checkpoint.py @@ -0,0 +1,96 @@ +import json +import tempfile +import unittest +from pathlib import Path + +from rec_mlx.checkpoint import validate + + +class CheckpointValidationTests(unittest.TestCase): + def setUp(self): + self.temp_dir = tempfile.TemporaryDirectory() + self.directory = Path(self.temp_dir.name) + + def tearDown(self): + self.temp_dir.cleanup() + + def _setup_files(self, weight_format, bits): + manifest_data = { + "model_id": f"flint-{weight_format.split('-')[-1]}", + "weight_format": weight_format, + "kv_format": "bf16", + "max_context": 262144, + "checkpoint_files": { + "config.json": "dummy", + "tokenizer.json": "dummy", + "tokenizer_config.json": "dummy", + "special_tokens_map.json": "dummy", + "model-00000.safetensors": "dummy", + }, + } + manifest_path = self.directory / "manifest.json" + manifest_path.write_text(json.dumps(manifest_data)) + + config_data = { + "sliding_window": 513, + "mlx_export": { + "bits": bits, + "routed_expert_bits": bits, + "embedding_bits": 8, + }, + "quantization": { + "group_size": 64, + "bits": bits, + }, + } + (self.directory / "config.json").write_text(json.dumps(config_data)) + (self.directory / "tokenizer.json").write_text("{}") + (self.directory / "tokenizer_config.json").write_text("{}") + (self.directory / "special_tokens_map.json").write_text("{}") + (self.directory / "model-00000.safetensors").write_text("") + return manifest_path + + def test_validate_int4(self): + manifest_path = self._setup_files("mlx-affine-int4", 4) + m = validate(self.directory, manifest_path) + self.assertEqual(m["weight_format"], "mlx-affine-int4") + + def test_validate_int6(self): + manifest_path = self._setup_files("mlx-affine-int6", 6) + m = validate(self.directory, manifest_path) + self.assertEqual(m["weight_format"], "mlx-affine-int6") + + def test_validate_unsupported_format(self): + manifest_path = self._setup_files("mlx-affine-int8", 8) + with self.assertRaises(ValueError) as ctx: + validate(self.directory, manifest_path) + self.assertIn("requires affine INT4 or INT6 weights", str(ctx.exception)) + + def test_validate_mismatched_bits(self): + manifest_path = self._setup_files("mlx-affine-int6", 4) + with self.assertRaises(ValueError) as ctx: + validate(self.directory, manifest_path) + self.assertIn("Expected INT6 routed experts", str(ctx.exception)) + + def test_validate_jang_int6(self): + config_data = { + "sliding_window": 513, + "quantization": { + "group_size": 64, + "bits": 8, + "model.embed_tokens": {"group_size": 64, "bits": 8, "mode": "affine"}, + "model.layers.2.mlp.switch_mlp.gate_proj": {"group_size": 64, "bits": 6, "mode": "affine"}, + }, + } + config_path = self.directory / "config.json" + config_path.write_text(json.dumps(config_data)) + (self.directory / "tokenizer.json").write_text("{}") + (self.directory / "tokenizer_config.json").write_text("{}") + (self.directory / "model-00000.safetensors").write_text("") + m = validate(self.directory, config_path) + self.assertEqual(m["weight_format"], "mlx-affine-int6") + self.assertEqual(m["model_id"], "flint-int6") + + +if __name__ == "__main__": + unittest.main() diff --git a/engine/tests/test_cli_defaults.py b/engine/tests/test_cli_defaults.py index 4fb343e..401fd62 100644 --- a/engine/tests/test_cli_defaults.py +++ b/engine/tests/test_cli_defaults.py @@ -51,8 +51,73 @@ def check_defaults(*args, _fast_sync=fast_sync, _engine=engine, **kwargs): main() self.assertEqual(make.call_args.kwargs["grammar_mode"], grammar) self.assertEqual(construct.call_args.kwargs["enforce_tool_schemas"], enforce) + self.assertEqual(make.call_args.args[1], ("127.0.0.1", 8080)) ready = json.loads(stdout.getvalue().splitlines()[0]) self.assertEqual(ready["grammar"], grammar) self.assertEqual(ready["fast_sync"], fast_sync) self.assertEqual(ready["enforce_tool_schemas"], enforce) self.assertEqual(ready["chat_template"], "flint") + + def test_serve_binds_configured_host(self): + engine = Mock(identity={"model_id": "flint"}, startup={"chat_template": "flint"}) + server = Mock(server_port=8080) + with ( + patch( + "sys.argv", + [ + "rec-mlx", + "serve", + "flint", + "--model-dir", + ".", + "--manifest", + "model.json", + "--host", + "0.0.0.0", + ], + ), + patch("rec_mlx.cli.Engine", return_value=engine), + patch("rec_mlx.server.make_server", return_value=server) as make, + patch("rec_mlx.cli.signal.signal"), + redirect_stdout(io.StringIO()) as stdout, + ): + main() + self.assertEqual(make.call_args.args[1], ("0.0.0.0", 8080)) + self.assertIn("Listening on http://0.0.0.0:8080/v1", stdout.getvalue()) + + def test_serve_sampling_parameters(self): + engine = Mock(identity={"model_id": "flint"}, startup={"chat_template": "flint"}) + server = Mock(server_port=8080) + with ( + patch( + "sys.argv", + [ + "rec-mlx", + "serve", + "flint", + "--model-dir", + ".", + "--manifest", + "model.json", + "--temperature", + "0.6", + "--presence-penalty", + "0.2", + "--repeat-penalty", + "1.1", + ], + ), + patch("rec_mlx.cli.Engine", return_value=engine), + patch("rec_mlx.server.make_server", return_value=server) as make, + patch("rec_mlx.cli.signal.signal"), + redirect_stdout(io.StringIO()) as stdout, + ): + main() + defaults = make.call_args.kwargs["defaults"] + self.assertEqual(defaults.temperature, 0.6) + self.assertEqual(defaults.presence_penalty, 0.2) + self.assertEqual(defaults.repetition_penalty, 1.1) + ready = json.loads(stdout.getvalue().splitlines()[0]) + self.assertEqual(ready["sampling"]["temperature"], 0.6) + self.assertEqual(ready["sampling"]["presence_penalty"], 0.2) + self.assertEqual(ready["sampling"]["repetition_penalty"], 1.1) diff --git a/engine/tests/test_serving.py b/engine/tests/test_serving.py index 71eab75..9655584 100644 --- a/engine/tests/test_serving.py +++ b/engine/tests/test_serving.py @@ -92,6 +92,19 @@ def generate_tokens(self, ids, sampling, *, should_stop, return_logprobs=False, class ServingTest(unittest.TestCase): + def test_server_binds_the_requested_address(self): + server = make_server(Engine(), ("0.0.0.0", 0), Path(self.temp.name) / "metrics-bind") + thread = threading.Thread(target=server.serve_forever) + thread.start() + try: + self.assertEqual(server.socket.getsockname()[0], "0.0.0.0") + finally: + server.shutdown() + server.generator.cancel() + server.generator.join() + server.server_close() + thread.join() + def test_idle_preparation_follows_response_and_yields_to_queued_request(self): preparing = threading.Event() proceed = threading.Event() diff --git a/packaging/flint-int6-mlx.json b/packaging/flint-int6-mlx.json new file mode 100644 index 0000000..16a702f --- /dev/null +++ b/packaging/flint-int6-mlx.json @@ -0,0 +1,94 @@ +{ + "schema": 1, + "model": "flint", + "revision": "flint-int6", + "development": false, + "engine": "mlx", + "os": "darwin", + "arch": "arm64", + "backend": "metal", + "context": 262144, + "max_context": 262144, + "min_ram_bytes": 68719476736, + "entry": "config.json", + "files": [ + { + "name": "chat_template.jinja", + "url": "https://huggingface.co/chimpanzeetaxidriver/flint-oQ6e/resolve/1504c5802a33e69337075a0f46895b9a9c9a32e8/chat_template.jinja", + "size": 2328, + "sha256": "7940f4ffa3275c36caa4a582f845aaad79da11018b3961e8fbcad6ae29733466" + }, + { + "name": "config.json", + "url": "https://huggingface.co/chimpanzeetaxidriver/flint-oQ6e/resolve/1504c5802a33e69337075a0f46895b9a9c9a32e8/config.json", + "size": 93877, + "sha256": "b6b94f77e664c8016e97308e488f37807b41f4a96779594f5e57e8aa25093ed7" + }, + { + "name": "generation_config.json", + "url": "https://huggingface.co/chimpanzeetaxidriver/flint-oQ6e/resolve/1504c5802a33e69337075a0f46895b9a9c9a32e8/generation_config.json", + "size": 202, + "sha256": "a499c71895e008abd572558a5ad5904feaed92cf2a3b6e2cb0e9cabf03b96ed3" + }, + { + "name": "model-00001-of-00007.safetensors", + "url": "https://huggingface.co/chimpanzeetaxidriver/flint-oQ6e/resolve/1504c5802a33e69337075a0f46895b9a9c9a32e8/model-00001-of-00007.safetensors", + "size": 5124583742, + "sha256": "3bf6e441093aa552fbb17c54446da33cc0fff22ec85d55ab696d7fdb8ac6e741" + }, + { + "name": "model-00002-of-00007.safetensors", + "url": "https://huggingface.co/chimpanzeetaxidriver/flint-oQ6e/resolve/1504c5802a33e69337075a0f46895b9a9c9a32e8/model-00002-of-00007.safetensors", + "size": 5016402320, + "sha256": "215166b13a86412d7f26e81ebf9c1c985b1a606d6dae0df218ca3f50d3baa506" + }, + { + "name": "model-00003-of-00007.safetensors", + "url": "https://huggingface.co/chimpanzeetaxidriver/flint-oQ6e/resolve/1504c5802a33e69337075a0f46895b9a9c9a32e8/model-00003-of-00007.safetensors", + "size": 5016401767, + "sha256": "46587824007ee48432ad17cdb930dc2fa1bbce801d1320d58ed4ed69d742b30a" + }, + { + "name": "model-00004-of-00007.safetensors", + "url": "https://huggingface.co/chimpanzeetaxidriver/flint-oQ6e/resolve/1504c5802a33e69337075a0f46895b9a9c9a32e8/model-00004-of-00007.safetensors", + "size": 5016402370, + "sha256": "8cd7221b042f9b17fa472992c6442eaa950f1420c462cb5c3c10d2c81537408a" + }, + { + "name": "model-00005-of-00007.safetensors", + "url": "https://huggingface.co/chimpanzeetaxidriver/flint-oQ6e/resolve/1504c5802a33e69337075a0f46895b9a9c9a32e8/model-00005-of-00007.safetensors", + "size": 5016402388, + "sha256": "ed0f1b4c95dfd7c6129b24a3d90e984bfadd0d9beeafaa667878ca62b65a9485" + }, + { + "name": "model-00006-of-00007.safetensors", + "url": "https://huggingface.co/chimpanzeetaxidriver/flint-oQ6e/resolve/1504c5802a33e69337075a0f46895b9a9c9a32e8/model-00006-of-00007.safetensors", + "size": 5016401751, + "sha256": "b5d4d759f943040b93464d62f2b4a52eaa0e43b9448792f718e7e129126569b2" + }, + { + "name": "model-00007-of-00007.safetensors", + "url": "https://huggingface.co/chimpanzeetaxidriver/flint-oQ6e/resolve/1504c5802a33e69337075a0f46895b9a9c9a32e8/model-00007-of-00007.safetensors", + "size": 2617253931, + "sha256": "47f8ba4eb02d499448508a67dd6168e71e3a582b20782294ba548b348175f99d" + }, + { + "name": "model.safetensors.index.json", + "url": "https://huggingface.co/chimpanzeetaxidriver/flint-oQ6e/resolve/1504c5802a33e69337075a0f46895b9a9c9a32e8/model.safetensors.index.json", + "size": 182839, + "sha256": "d5c00c73cc7f5d0207da0b7fd74cefdbd161567b207dcce67906a82679334c5a" + }, + { + "name": "tokenizer.json", + "url": "https://huggingface.co/chimpanzeetaxidriver/flint-oQ6e/resolve/1504c5802a33e69337075a0f46895b9a9c9a32e8/tokenizer.json", + "size": 10879128, + "sha256": "a263cce0b7e2b03f7ee07ce809783ef4b47f160c3f268a4ea1524f638da45b68" + }, + { + "name": "tokenizer_config.json", + "url": "https://huggingface.co/chimpanzeetaxidriver/flint-oQ6e/resolve/1504c5802a33e69337075a0f46895b9a9c9a32e8/tokenizer_config.json", + "size": 30439, + "sha256": "159ac88217e90ffd0a4a294726048bb302829e82e1958dc85bb203aa79138fbf" + } + ] +} diff --git a/packaging/package.py b/packaging/package.py index 78adce8..29021f0 100644 --- a/packaging/package.py +++ b/packaging/package.py @@ -65,6 +65,11 @@ def main(): required=True, help="model manifest, such as packaging/flint-int4-mlx.json", ) + p.add_argument( + "--manifest-int6", + type=Path, + help="optional INT6 model manifest, such as packaging/flint-int6-mlx.json", + ) p.add_argument( "--version", default="0.1.0-dev", @@ -90,6 +95,13 @@ def main(): shutil.rmtree(assets) assets.mkdir() (assets / "manifest.json").write_text(json.dumps(m, indent=2) + "\n") + (assets / "manifest-int4.json").write_text(json.dumps(m, indent=2) + "\n") + manifest_int6 = a.manifest_int6 or (root / "packaging/flint-int6-mlx.json") + if manifest_int6.exists(): + m_int6 = json.loads(manifest_int6.read_text()) + m_int6["runtime_files"] = m["runtime_files"] + m_int6["runtime_sha256"] = m["runtime_sha256"] + (assets / "manifest-int6.json").write_text(json.dumps(m_int6, indent=2) + "\n") (assets / "runtime.tar.gz").write_bytes(archive) subprocess.run( ["go", "test", "./cli", "-run", "TestEmbeddedManifest"], cwd=build, check=True diff --git a/tests/data/smoke-mlx-server.py b/tests/data/smoke-mlx-server.py index 830c10e..0f3ef94 100644 --- a/tests/data/smoke-mlx-server.py +++ b/tests/data/smoke-mlx-server.py @@ -21,6 +21,7 @@ parser.add_argument("--model-dir", type=Path, required=True) parser.add_argument("--manifest", type=Path, required=True) parser.add_argument("--port", type=int, required=True) +parser.add_argument("--host", default="127.0.0.1") parser.add_argument("--context-limit", type=int, required=True) parser.add_argument("--api-key-file", type=Path, required=True) args = parser.parse_args() @@ -126,5 +127,5 @@ def do_POST(self): ) -print(f"Listening on http://127.0.0.1:{args.port}/v1", flush=True) -http.server.HTTPServer(("127.0.0.1", args.port), Handler).serve_forever() +print(f"Listening on http://{args.host}:{args.port}/v1", flush=True) +http.server.HTTPServer((args.host, args.port), Handler).serve_forever()