diff --git a/server/sched.go b/server/sched.go index 980ad358..fcc3a61d 100644 --- a/server/sched.go +++ b/server/sched.go @@ -594,7 +594,7 @@ func (s *Scheduler) load(req *LlmRequest, systemInfo ml.SystemInfo, gpus []ml.De if slices.Contains(req.model.Config.Capabilities, "image") { llama, err = imagegen.NewServer(modelName) } else { - llama, err = mlxrunner.NewClient(modelName) + llama, err = mlxrunner.NewClient(modelName, req.opts.NumCtx) } } if err != nil { diff --git a/x/mlxrunner/client.go b/x/mlxrunner/client.go index 7633ec6d..f5f97df9 100644 --- a/x/mlxrunner/client.go +++ b/x/mlxrunner/client.go @@ -32,29 +32,31 @@ import ( // Client wraps an MLX runner subprocess to implement llm.LlamaServer for LLM models. type Client struct { - port int - modelName string - contextLength atomic.Int64 - memory atomic.Uint64 - done chan struct{} - doneErr error // valid after done is closed - client *http.Client - status *llm.StatusWriter - mu sync.Mutex - cmd *exec.Cmd + port int + modelName string + contextLength atomic.Int64 + softContextLength int // recommended limit to avoid poor performance + memory atomic.Uint64 + done chan struct{} + doneErr error // valid after done is closed + client *http.Client + status *llm.StatusWriter + mu sync.Mutex + cmd *exec.Cmd } // NewClient prepares a new MLX runner client for LLM models. // The subprocess is not started until Load() is called. -func NewClient(modelName string) (*Client, error) { +func NewClient(modelName string, softContextLength int) (*Client, error) { if err := imagegen.CheckPlatformSupport(); err != nil { return nil, err } c := &Client{ - modelName: modelName, - done: make(chan struct{}), - client: http.DefaultClient, + modelName: modelName, + softContextLength: softContextLength, + done: make(chan struct{}), + client: http.DefaultClient, } modelManifest, err := manifest.LoadManifest(modelName) @@ -223,6 +225,13 @@ func (c *Client) ContextLength() int { return int(c.contextLength.Load()) } +func (c *Client) reportedContextLength(modelContextLength int) int { + if c.softContextLength > 0 && (modelContextLength == 0 || c.softContextLength < modelContextLength) { + return c.softContextLength + } + return modelContextLength +} + // Detokenize implements llm.LlamaServer. func (c *Client) Detokenize(ctx context.Context, tokens []int) (string, error) { return "", errors.New("not supported") @@ -421,7 +430,7 @@ func (c *Client) Ping(ctx context.Context) error { return err } - c.contextLength.Store(int64(status.ContextLength)) + c.contextLength.Store(int64(c.reportedContextLength(status.ContextLength))) c.memory.Store(status.Memory) return nil