Skip to content

feat(mlx): support Qwen3.8 Flash Next (qwen4_exp) - #2283

Open
tharalds wants to merge 1 commit into
exo-explore:mainfrom
tharalds:codex/qwen38-flash-next-exo
Open

feat(mlx): support Qwen3.8 Flash Next (qwen4_exp)#2283
tharalds wants to merge 1 commit into
exo-explore:mainfrom
tharalds:codex/qwen38-flash-next-exo

Conversation

@tharalds

Copy link
Copy Markdown

Motivation

Qwen3.8 Flash Next uses the qwen4_exp architecture, whose QSA/Gated
DeltaNet cache, hyper-connections, n-gram PLE storage, and weight sanitizer
currently live in MLX-VLM rather than MLX-LM. EXO could discover the model but
could not load or generate from a compatible checkpoint through its MLX engine.

Changes

  • Pin MLX-VLM to an immutable upstream revision containing Qwen4-Exp QSA/APC,
    continuous-batching, MTP, and external-PLE support.
  • Route qwen4_exp through MLX-VLM while preserving MLX-LM loading for every
    existing architecture.
  • Adapt MLX-VLM's structured language-model output to EXO's logits contract.
  • Add MLX-VLM cache types to EXO's cache model, including detached recurrent
    snapshots, QSA trimming, prefill rollback, and prefix-cache reuse.
  • Repair Qwen4-Exp hybrid layer/cache indices after pipeline slicing. Tensor
    parallel placement remains disabled until QSA/GDN/MoE/PLE sharding has a
    dedicated bit-exact implementation.
  • Normalize Qwen3.8 reasoning efforts to its supported low, medium, and
    xhigh template values and recognize its EOS tokens.
  • Add a conservative catalog card for the corrected
    sh0wie/Qwen3.8-Flash-Next-REAP-288-MLX-4bit checkpoint. Vision is not
    advertised because the checkpoint publisher has only evaluated text after
    pruning.
  • Add focused Metal-backed tests, model-card validation, a reusable full-weight
    smoke harness, and operator documentation.

The disaggregated-prefill wire path now rejects MLX-VLM caches explicitly
instead of attempting to serialize incomplete recurrent/QSA state.

Why It Works

MLX-VLM owns the current qwen4_exp implementation and sanitizes its official
conditional-generation weights. EXO retains only that model's language module,
wraps its structured output as logits, and continues to use EXO's existing chat,
generation, placement, and prefix-cache orchestration.

The recommended REAP-288 checkpoint corrects two source-conversion defects:
RMSNorm tensors use the zero-centered convention expected by MLX-VLM, and PLE
tensors plus per-tensor quantization overrides use the runtime's module names.
Its external PLE manifest also keeps active memory substantially below the
checkpoint's on-disk size.

Test Plan

Manual Testing

  • Hardware: Mac Studio, Apple M3 Ultra, 512 GB unified memory.
  • Downloaded all 131 indexed shards of the recommended checkpoint and verified
    that no indexed file was missing or empty.
  • Native MLX-VLM returned the exact requested coherent response, peaking at
    41.7 GB active memory.
  • EXO's loader, chat template, cache, and generator returned the exact requested
    coherent response, peaking at 44.5 GB active memory.
  • Thinking mode produced the correct answer with a properly delimited reasoning
    section.
  • A repeated full-weight EXO request reused 18 of 20 prompt tokens and dropped
    from 2.8 seconds to 0.7 seconds while returning identical correct output.
  • An older conversion was tested through both native MLX-VLM and EXO and
    produced the same incoherent output; it is deliberately not cataloged.

Automated Testing

  • 109 focused/regression tests passed (42 hardware/model-dependent tests
    deselected), covering Qwen4-Exp loading, real Metal prefill/decode, QSA and
    recurrent cache restore, pipeline handling, reasoning, API tool use, parsers,
    and disaggregated-prefill behavior.
  • ruff check .
  • ruff format --check .
  • basedpyright — 0 errors, 0 warnings.
  • uv lock --check

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.

1 participant