lmstudio-ai / lmstudio-ai/mlx-engine

MLX auto-fit replaces the configured context length and disagrees with the prompt cache budget

Open
#366 1 comment 5 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.2k
Forks
133
Avg merge
21h 6m
Merged PRs (30d)
1

Description

## Summary

The MLX engine reports the auto-fit context length to the application and ignores the context length that the user configured. The prompt cache uses a different rule and keeps the larger of the two values. The two components then disagree about the context length of the same loaded model.

## Environment

- LM Studio 0.4.21+2
- Runtime `mlx-llm-mac-arm64-apple-metal-advsimd` 1.11.0
- Vendor package `app-mlx-generate-mac14-arm64@34`
- MacBook Pro M4, 24 GB unified memory, macOS (Darwin 25.5.0)
- Model `lmstudio-community/Qwen3.8-27B-MLX-4bit` (16.08 GB, arch `qwen3_5`)

## Steps to reproduce

1. Set the model loading guardrails to OFF.
2. Load the model with an explicit context length: `lms load qwen3.8-27b-mlx --context-length 24576`.
3. Run `lms ps`.
4. Send a prompt of about 7,000 tokens to `/v1/chat/completions`.

## Result

`lms ps` shows a context of 4096. The context field in the graphical interface also returns to 4096 after each load. The server rejects the 7,000-token prompt with this message: `number of tokens to keep from the initial prompt is greater than the context length`.

## Logs

```
[context_fit][WARNING]: Model context auto-fit calculated 0 tokens; using the 4,096 token minimum
[context_fit][INFO]: Model context auto-fit: family=qwen3_5 max=262,144 fitted=4,096
working_set=17.76GiB reserve=3.00GiB safe_ceiling=14.76GiB baseline=14.95GiB
full_kv=65536B/token prompt_inputs=10240B/token attention=98304B/token
rotating_peak=0.00GiB fixed_ssm=0.14GiB estimated_peak=15.76GiB
[cache_store][INFO]: VLM prompt cache context target: configured=24,576 fitted=4,096 effective=24,576
```

The last line shows the problem. The engine receives the configured value of 24,576. The prompt cache keeps 24,576. The engine still reports 4,096.

## Cause

Two code paths use the fitted value in different ways.

`BatchedModelKit._load_model()` stores only the fitted value:

```python
self._effective_context_length = fit_batched_vlm_context(
model=self.model,
prefill_step_size=self.prefill_step_size,
)
```

`get_runtime_load_info()` in `generate.py` reports that value to the application:

```python
context_length = getattr(model_kit, "effective_context_length", None)
```

`VlmPromptCacheStore.ensure_max_kv_size()` applies a different rule:

```python
self._max_kv_size = max(configured_max_kv_size or 0, max_kv_size)
```

`_load_model()` never reads `max_kv_size`. The constructor passes that value only to `VlmPromptCacheStore`. The reported context length and the budgeted context length differ by construction.

## The 4,096 floor hides a failed fit

On this hardware the fit computes zero tokens. The model baseline is 14.95 GiB. The safe ceiling is 14.76 GiB. `available_prompt_bytes` is zero, and `tokens_that_fit` is zero.

A result of zero means that the model does not fit at all. The floor at `MIN_FITTED_CONTEXT_TOKENS` turns this result into 4,096 tokens. The load then reports success. The user sees a model that loads but cannot accept a useful prompt.

## Suggested fix

Make the two paths agree. Two changes would fix this:

1. Fail the load when the fit computes zero tokens. Report a memory error instead of the 4,096 floor.
2. Report the explicit user setting, or reject it with an error. Do not replace it silently.

The second change needs care. On this hardware the memory arithmetic is correct. A 16 GB model on 24 GB of unified memory leaves about 2.8 GiB below the recommended working set. A larger context does not fit. A clear error helps the user more than a silent floor.

## Note on iogpu.wired_limit_mb

`calculate_context_fit()` reads `max_recommended_working_set_size` from Metal. This value does not follow `sysctl iogpu.wired_limit_mb`. A raised wired limit cannot change the fit.

## Related issues

These reports show the same disagreement between components:

- lmstudio-ai/lmstudio-bug-tracker#2212: the logs show `fitted=4,864` and `effective=6,144`
- lmstudio-ai/lmstudio-bug-tracker#2250: the fit overrides the configured value upward
- lmstudio-ai/lmstudio-bug-tracker#2191: the load reports success when memory is insufficient

Contributor guide

Open the contributing guide

Research direction

Start with BatchedModelKit._load_model(), get_runtime_load_info() in generate.py, and VlmPromptCacheStore.ensure_max_kv_size(); trace how configured, fitted, and maximum context values flow during loading. Reproduce the explicit-context case and verify that the application and prompt cache agree, while a zero-token fit is reported as a load failure rather than a successful 4,096-token load.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
backend, machine-learning
Issue type
Bug
Difficulty
4/5
Estimated time
3-5 days
Activity status
Quiet
Clarity
Mostly clear
Newbie friendliness
48/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.