Skip to content

feat(engine): stage D2 - mistral sliding-window attention (#79)#103

Merged
lgsunnyvale merged 1 commit into
mainfrom
engine-inference-core-stage-d2-mistral-sliding-window-attention
Jul 20, 2026
Merged

feat(engine): stage D2 - mistral sliding-window attention (#79)#103
lgsunnyvale merged 1 commit into
mainfrom
engine-inference-core-stage-d2-mistral-sliding-window-attention

Conversation

@tlkahn

@tlkahn tlkahn commented Jul 20, 2026

Copy link
Copy Markdown
Owner

Summary

Stage D2 enablement for MODEL_MISTRAL (issue #79 / epic #55): llama-shaped dense path plus uniform sliding-window attention on every layer. No forward.c changes - D0 F7 mask/trim already covers the mistral config shape.

  • Support gate accepts windowed and null-window mistral; qwen3/llama still reject sliding window
  • weights_expected_names llama-shaped set (21 untied / 20 tied)
  • gen_tiny_ckpt mistral recipe (sliding_window=8, omit pattern when 0) + committed fixtures
  • GPU smoke + degenerate-window + window-boundary incremental==full
  • Parity table records v0.1-era checkpoint with real sliding_window=4096

Cycle 0 audit

Confirmed against HF mlx-community/Mistral-7B-Instruct-v0.1-4bit-mlx/config.json:

Field Value
model_type mistral
sliding_window 4096 (int; v0.2+ is null - do not use for window parity)
tie_word_embeddings false
rope_theta 10000.0
rope_scaling absent
weight names llama set (no qk-norm)

Config / forward delta vs llama

Field llama mistral v0.1 D2 impact
has_qk_norm false false same
attention_bias false false same
has_sliding_window false true when int allowed
sliding_window_pattern 0 absent (default 0) all layers local
rope_theta 500000 (llama3) 10000 plain rope

R1 decisions

  1. No forward.c changes unless RED proved a hole (none did).
  2. Gate allows sliding window only for mistral via REJECT_ALLOW_SLIDING_WINDOW.
  3. Both windowed and null-window mistral pass the gate.
  4. Generator omits sliding_window_pattern when 0 (get_dim_int rejects 0); gemma3 still emits pattern=2.
  5. Parity id = mlx-community/Mistral-7B-Instruct-v0.1-4bit-mlx. Short-context parity is family enablement; window-boundary correctness owned by tiny_mistral GPU tests (CI will not exceed 4096 tokens). Chat mode not forced preemptively (no dry-run evidence yet).
  6. Out of scope: engine: hoist prefill sliding-window mask build out of the per-layer attention call #92 prefill mask hoist, engine: KV cache backing store grows unbounded under sliding-window attention (ring-buffer follow-up) #93 ring-buffer KV, gemma3 interleaved pattern (D4).

Test plan

  • make test (43 passed)
  • ./tests/test_emodel_gate (windowed + null pass; llama/qwen3 SW rejects held)
  • ./tests/test_weights / ./tests/test_model_config (mistral expected names + tiny fixture)
  • ./tests/test_forward_mistral_gpu (load/map/shape, short incremental, degenerate window, boundary incremental==full)
  • make test-gpu (mistral + D0/kvcache/llama green; pre-existing test_weights_gpu deepseek GGUF assert unrelated)
  • make test-parity-script mistral absent-dir skip green (wrapper-skip-absent-mistral)
  • Manual: scripts/parity_family.sh mistral when checkpoint present

Closes #79

Enable MODEL_MISTRAL on the dense forward path: support gate allows
uniform sliding window (and null-window v0.2+), llama-shaped weights
descriptor, tiny_mistral fixture (W=8), GPU degenerate/boundary tests,
and parity id for Mistral-7B-Instruct-v0.1-4bit-mlx.
@tlkahn
tlkahn requested a review from lgsunnyvale as a code owner July 20, 2026 04:32
@tlkahn

tlkahn commented Jul 20, 2026

Copy link
Copy Markdown
Owner Author

TODO: mistral real-checkpoint parity on 32 GB+ Mac

Manual item from the test plan is still open:

  • scripts/parity_family.sh mistral with checkpoint present

Blocked on hardware, not on the D2 diff. The canonical id is mlx-community/Mistral-7B-Instruct-v0.1-4bit-mlx (~4.0 GB 4-bit). parity_temp0.sh cold-loads that checkpoint twice (oracle mlx_lm.server, then mlxd run). On the current 16 GB machine that double load is memory-bound and not a reliable gate.

Do this later on a 32 GB+ Apple Silicon Mac (local or cloud MacStadium / EC2 Mac / similar):

MLXD_MLX_SERVE_BIN=/path/to/scripts/mlx_lm_server_wrapper.sh \
MLXD_PARITY_CKPT_MISTRAL=$HOME/.cache/mlxd/models/mlx-community--Mistral-7B-Instruct-v0.1-4bit-mlx \
MLXD_PARITY_MAX_TOKENS=16 \
scripts/parity_family.sh mistral

Notes for whoever picks this up:

  • v0.1 with sliding_window=4096 is required for windowed family enablement. v0.2+ often has sliding_window: null and does not replace this check.
  • There is no smaller real HF mistral id that still exercises D2; window-boundary correctness stays with tiny_mistral / test_forward_mistral_gpu (already green).
  • Short-context text parity only (prompt + max_tokens << 4096). If raw instruct loops, retry with MLXD_PARITY_MODE=chat before blaming the forward path.
  • This TODO does not block merging D2 on the strength of the unit/GPU suite; it is residual real-quant confidence.

@lgsunnyvale
lgsunnyvale merged commit 01e7159 into main Jul 20, 2026
1 check passed
@lgsunnyvale
lgsunnyvale deleted the engine-inference-core-stage-d2-mistral-sliding-window-attention branch July 20, 2026 05:09
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Engine Inference Core Stage D2: mistral - sliding-window attention

2 participants