MLX: wire up scheduler selected context size for ps (#16918)

In the PS output, expose the scheduler selected size (clamped by model context size) instead of always reporting the model max context.  This will help provide a hint to clients to keep the context size below this value to avoid paging and poor performance on smaller VRAM systems.

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 {

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
+	port              int
-	modelName     string
+	modelName         string
-	contextLength atomic.Int64
+	contextLength     atomic.Int64
-	memory        atomic.Uint64
+	softContextLength int // recommended limit to avoid poor performance
-	done          chan struct{}
+	memory            atomic.Uint64
-	doneErr       error // valid after done is closed
+	done              chan struct{}
-	client        *http.Client
+	doneErr           error // valid after done is closed
-	status        *llm.StatusWriter
+	client            *http.Client
-	mu            sync.Mutex
+	status            *llm.StatusWriter
-	cmd           *exec.Cmd
+	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,
+		modelName:         modelName,
-		done:      make(chan struct{}),
+		softContextLength: softContextLength,
-		client:    http.DefaultClient,
+		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