Skip to content

feat(mlx-lm): add Qwen3.5 hybrid attention model support - #335

Open
Ogilthorp3 wants to merge 1 commit into
oxiglade:mainfrom
Ogilthorp3:feat/qwen3_5-model
Open

feat(mlx-lm): add Qwen3.5 hybrid attention model support#335
Ogilthorp3 wants to merge 1 commit into
oxiglade:mainfrom
Ogilthorp3:feat/qwen3_5-model

Conversation

@Ogilthorp3

Copy link
Copy Markdown

Summary

Adds Qwen3.5 (hybrid Mamba + Transformer) model support to mlx-lm.

Qwen3.5 is a novel architecture mixing linear attention (Mamba-style SSM) with standard full attention (GQA + RoPE) in a 48:16 ratio across 64 layers. This is the first Rust implementation.

Changes

  • **** (1,253 lines) — Full model implementation:
    • with nested and dispatch
    • : conv1d, SSM discretization (A_log, dt_bias), causal attention with B-scaled weighting, output gating
    • : GQA with partial RoPE (0.25 factor), q/k norms
    • enum dispatching between layer types
    • Quantized weight loading (4-bit affine, group_size=64)
    • Weight prefix stripping for VL-extracted checkpoints
  • **** — Added
  • **** — Added trait impl for qwen3_5

Tested with

  • Qwen3.5-27B-4bit-text (64 layers, 5120 hidden, 248K vocab)
  • Model loads in ~3s on M4 Pro (Mac Mini, 64 GB)
  • Compiles with zero warnings

Notes

  • SSM scan uses causal attention approximation — full exponential decay recurrence marked TODO
  • All linear layers use for quantization support
  • Follows existing qwen3.rs patterns exactly (derive macros, load functions, Module impl)

Built for Sanctum — a home AI platform running on Apple Silicon.

Implements the Qwen3.5 architecture — a hybrid model combining
Mamba-style linear attention (SSM) with standard full attention
(GQA + RoPE) in a 48:16 ratio across 64 layers.

Key components:
- ModelArgs with nested RopeParameters and layer_types dispatch
- LinearAttention: conv1d, SSM discretization (A_log, dt_bias),
  causal attention with B-scaled weighting, output gating
- FullAttention: GQA with partial RoPE (0.25 factor), q/k norms
- DecoderLayer enum dispatching between layer types
- Quantized weight loading (4-bit affine, group_size=64)
- Weight prefix stripping for VL-extracted checkpoints
- load_qwen3_5_model / load_qwen3_5_tokenizer functions

Tested with Qwen3.5-27B-4bit-text (64 layers, 5120 hidden,
248K vocab). Model loads in ~3s on M4 Pro.

SSM scan uses causal attention approximation — full exponential
decay recurrence marked TODO for production accuracy.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
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