feat(engine): stage D6 - qwen3_5 dense full-attn deltas

scripts/parity_family.sh

@@ -26,7 +26,11 @@ family_id() {
         llama)   echo "mlx-community/TinyLlama-1.1B-Chat-v1.0-4bit" ;;
         mistral) echo "" ;;
         gemma4)  echo "mlx-community/gemma-4-e4b-it-4bit" ;;
-        qwen3_5) echo "" ;;
+        # D6 records the smallest dense 4-bit Qwen3.5 id. Real checkpoints are
+        # hybrid (linear + full attention). mlxd Stage D rejects linear
+        # attention; byte parity is Stage E. skip-if-absent still applies when
+        # the snapshot is not on disk.
+        qwen3_5) echo "mlx-community/Qwen3.5-0.8B-4bit" ;;
         *)       return 1 ;;
     esac
 }
@@ -76,6 +80,17 @@ run_one() {
         return 0
     fi
 
+    # qwen3_5: real checkpoints are hybrid (linear + full attention). Stage D
+    # rejects linear attention at the support gate, so the recorded canonical
+    # id cannot run parity yet. Hard-skip when resolved via the family table /
+    # MLXD_PARITY_CKPT_ROOT path. MLXD_PARITY_CKPT_QWEN3_5 remains an escape
+    # hatch for a future dense checkpoint (checked above via _override).
+    # Stage E deletes this branch.
+    if [ "$_fam" = "qwen3_5" ] && [ -z "$_override" ]; then
+        printf 'skipped: qwen3_5 hybrid parity deferred to Stage E\n'
+        return 0
+    fi
+
     # gemma4-it: use chat-template parity. Raw completion on this instruct
     # checkpoint collapses into short token loops; mlx-serve's degenerate-loop
     # detector then stops early and creates length-only mismatches against

src/engine/emodel.c

@@ -13,8 +13,15 @@
     } \
 } while (0)
 
+/* Flags that relax reject_dense_common for families that need the feature. */
+enum {
+    REJECT_ALLOW_PARTIAL_ROPE = 1 << 0, /* skip partial_rotary_factor != 1 */
+    REJECT_ALLOW_ATTN_GATE    = 1 << 1, /* skip attn_output_gate reject */
+};
+
 static int reject_dense_common(const model_config_t *cfg,
-                                char *err, size_t errlen) {
+                                char *err, size_t errlen,
+                                unsigned flags) {
     REJECT(cfg->attention_bias,
            "attention_bias not supported");
     REJECT(cfg->has_sliding_window,
@@ -31,10 +38,12 @@ static int reject_dense_common(const model_config_t *cfg,
            "scale_embeddings not supported");
     REJECT(cfg->has_pre_ff_norm,
            "pre-feedforward norm not supported");
-    REJECT(cfg->partial_rotary_factor != 1.0f,
-           "partial rotary embedding not supported");
-    REJECT(cfg->attn_output_gate,
-           "attention output gate not supported");
+    if (!(flags & REJECT_ALLOW_PARTIAL_ROPE))
+        REJECT(cfg->partial_rotary_factor != 1.0f,
+               "partial rotary embedding not supported");
+    if (!(flags & REJECT_ALLOW_ATTN_GATE))
+        REJECT(cfg->attn_output_gate,
+               "attention output gate not supported");
     return 0;
 }
 
@@ -44,13 +53,13 @@ int engine_model_check_supported(const model_config_t *cfg,
 
     switch (cfg->family) {
     case MODEL_QWEN3:
-        if (reject_dense_common(cfg, err, errlen) != 0) return -1;
+        if (reject_dense_common(cfg, err, errlen, 0) != 0) return -1;
         REJECT(cfg->rope_scaling_type != NULL,
                "RoPE scaling not supported");
         return 0;
 
     case MODEL_LLAMA:
-        if (reject_dense_common(cfg, err, errlen) != 0) return -1;
+        if (reject_dense_common(cfg, err, errlen, 0) != 0) return -1;
         REJECT(cfg->rope_scaling_type != NULL &&
                strcmp(cfg->rope_scaling_type, "linear") != 0 &&
                strcmp(cfg->rope_scaling_type, "llama3") != 0 &&
@@ -86,6 +95,10 @@ int engine_model_check_supported(const model_config_t *cfg,
         REJECT(cfg->partial_rotary_factor_global <= 0.0f ||
                cfg->partial_rotary_factor_global > 1.0f,
                "partial_rotary_factor_global must be in (0, 1]");
+        /* Plain partial_rotary_factor is the full-attn dims path (qwen3_5).
+           gemma4 uses partial_rotary_factor_global via the freqs path instead. */
+        REJECT(cfg->partial_rotary_factor != 1.0f,
+               "partial rotary embedding not supported (gemma4 uses rope_parameters.full_attention)");
         /* Reject configs where a KV-shared layer has no same-type source */
         for (int i = 0; i < cfg->num_hidden_layers; i++) {
             if (model_layer_kv_shared(cfg, i) &&
@@ -96,9 +109,48 @@ int engine_model_check_supported(const model_config_t *cfg,
         }
         return 0;
 
+    case MODEL_QWEN3_5: {
+        /* Pure-dense qwen3_5 only. Common dense rejects apply, but partial rope
+           and attn_output_gate are required features here (not rejects).
+           Hybrid linear-attention / MoE stay Stage E. */
+        if (reject_dense_common(cfg, err, errlen,
+                                REJECT_ALLOW_PARTIAL_ROPE |
+                                REJECT_ALLOW_ATTN_GATE) != 0)
+            return -1;
+        REJECT(cfg->linear_num_key_heads > 0 || cfg->linear_num_value_heads > 0,
+               "linear attention / hybrid layers not supported");
+        if (cfg->has_explicit_layer_types) {
+            for (int i = 0; i < cfg->num_hidden_layers; i++) {
+                REJECT(!cfg->layer_is_global[i],
+                       "linear attention / hybrid layers not supported");
+            }
+        }
+        REJECT(cfg->partial_rotary_factor <= 0.0f ||
+               cfg->partial_rotary_factor > 1.0f,
+               "partial_rotary_factor must be in (0, 1]");
+        /* mlx.fast.rope requires dims even and positive. Mirror mlx-lm's
+           int(head_dim * prf) and fail fast on odd/degenerate products. */
+        {
+            int dims = (int)((float)cfg->head_dim * cfg->partial_rotary_factor);
+            REJECT(dims < 2 || (dims & 1),
+                   "partial_rotary_factor yields invalid rope dims (must be even, >= 2)");
+        }
+        /* interval==0 unset, interval==1 pure dense (every layer full-attn).
+           interval>1 implies hybrid: is_linear = (layer_idx+1) % interval != 0. */
+        REJECT(cfg->full_attention_interval > 1,
+               "full_attention_interval > 1 implies hybrid layers; not supported");
+        REJECT(cfg->rope_scaling_type != NULL,
+               "RoPE scaling not supported");
+        /* Gate is mandatory: mlx-lm gates unconditionally; no oracle or real
+           checkpoint for gate-off. Explicit false is parsed but unsupported. */
+        REJECT(!cfg->attn_output_gate,
+               "qwen3_5 requires attn_output_gate");
+        return 0;
+    }
+
     default:
         REJECT(true,
-               "unsupported model family (only qwen3/llama/gemma4 dense supported)");
+               "unsupported model family (only qwen3/llama/gemma4/qwen3_5 dense supported)");
     }
 }
 

src/engine/forward.c

@@ -476,6 +476,22 @@ static int add_proj_bias(mlx_array *io, const weights_t *w, const char *base,
     return rc;
 }
 
+/* RoPE dims for partial_rotary_factor on the full-attention path (qwen3_5).
+   Does not consult partial_rotary_factor_global (gemma4 freqs path).
+   Mirrors mlx-lm int(head_dim * prf). Every family's gate constrains
+   partial_rotary_factor (qwen3/llama via reject_dense_common, gemma4
+   explicitly), so only qwen3_5 reaches the non-1.0 path; its gate also
+   rejects odd/degenerate products so dims here are even and >= 2. */
+static int fwd_rope_dims_partial(const model_config_t *cfg, int head_dim) {
+    float prf = cfg->partial_rotary_factor;
+    if (prf <= 0.0f || prf > 1.0f)
+        return head_dim;
+    int dims = (int)((float)head_dim * prf);
+    if (dims <= 0 || dims > head_dim)
+        return head_dim;
+    return dims;
+}
+
 int fwd_attention(mlx_array *out, mlx_array x, int layer,
                   const weights_t *w, const model_config_t *cfg,
                   kvcache_t *kv, mlx_array rope_freqs, mlx_stream s) {
@@ -493,6 +509,7 @@ int fwd_attention(mlx_array *out, mlx_array x, int layer,
         cfg->attn_scale_one ? 1.0f : 1.0f / sqrtf((float)qpas);
     bool global = model_layer_is_global(cfg, layer);
     int kv_src = model_kv_source_layer(cfg, layer);
+    int rope_dims = fwd_rope_dims_partial(cfg, hd_l);
 
     mlx_array q = mlx_array_new();
     mlx_array k = mlx_array_new();
@@ -516,21 +533,60 @@ int fwd_attention(mlx_array *out, mlx_array x, int layer,
     mlx_array attn_reshaped = mlx_array_new();
     mlx_array result = mlx_array_new();
     mlx_array vnorm_w = mlx_array_new();
+    mlx_array gate = mlx_array_new();
 
     weight_triplet_t q_tri, k_tri, v_tri, o_tri;
     memset(&q_tri, 0, sizeof(q_tri));
     memset(&k_tri, 0, sizeof(k_tri));
     memset(&v_tri, 0, sizeof(v_tri));
     memset(&o_tri, 0, sizeof(o_tri));
 
-    /* Q projection (always present) */
+    /* Q projection (always present). With attn_output_gate, out features are
+       2*H*D and are split into queries + gate (mlx-lm Qwen3NextAttention). */
     weights_tensor_name(name, sizeof(name), cfg, layer, "self_attn.q_proj");
     if (weights_get_triplet(&q_tri, w, name) != 0) goto cleanup;
     if (fwd_linear(&q, x, &q_tri, cfg, s) != 0) goto cleanup;
     if (cfg->attention_bias && add_proj_bias(&q, w, name, s) != 0) goto cleanup;
 
-    int q_shape[] = {B, S, n_heads, hd_l};
-    if (!MLXB_CHECK(mlx_reshape(&q_reshaped, q, q_shape, 4, s))) goto cleanup;
+    if (cfg->attn_output_gate) {
+        int q_gate_shape[] = {B, S, n_heads, hd_l * 2};
+        if (!MLXB_CHECK(mlx_reshape(&q_reshaped, q, q_gate_shape, 4, s)))
+            goto cleanup;
+
+        mlx_vector_array split_vec = mlx_vector_array_new();
+        if (!MLXB_CHECK(mlx_split(&split_vec, q_reshaped, 2, -1, s))) {
+            mlx_vector_array_free(split_vec);
+            goto cleanup;
+        }
+        if (mlx_vector_array_size(split_vec) != 2) {
+            mlx_vector_array_free(split_vec);
+            goto cleanup;
+        }
+
+        mlx_array queries = mlx_array_new();
+        mlx_array gate_4d = mlx_array_new();
+        if (!MLXB_CHECK(mlx_vector_array_get(&queries, split_vec, 0)) ||
+            !MLXB_CHECK(mlx_vector_array_get(&gate_4d, split_vec, 1))) {
+            mlx_array_free(queries);
+            mlx_array_free(gate_4d);
+            mlx_vector_array_free(split_vec);
+            goto cleanup;
+        }
+        mlx_vector_array_free(split_vec);
+
+        mlx_array_free(q_reshaped);
+        q_reshaped = queries;
+
+        int gate_flat[] = {B, S, n_heads * hd_l};
+        if (!MLXB_CHECK(mlx_reshape(&gate, gate_4d, gate_flat, 3, s))) {
+            mlx_array_free(gate_4d);
+            goto cleanup;
+        }
+        mlx_array_free(gate_4d);
+    } else {
+        int q_shape[] = {B, S, n_heads, hd_l};
+        if (!MLXB_CHECK(mlx_reshape(&q_reshaped, q, q_shape, 4, s))) goto cleanup;
+    }
 
     /* Q norm */
     if (cfg->has_qk_norm) {
@@ -566,7 +622,7 @@ int fwd_attention(mlx_array *out, mlx_array x, int layer,
            Rope offset for Q = source's offset - S (source ran earlier this forward). */
         int src_offset = kvcache_layer_offset(kv, kv_src) - S;
         if (src_offset < 0) src_offset = 0;
-        if (fwd_rope_apply(&q_roped, q_transposed, hd_l, cfg, layer,
+        if (fwd_rope_apply(&q_roped, q_transposed, rope_dims, cfg, layer,
                            src_offset, rope_freqs, s) != 0)
             goto cleanup;
 
@@ -579,7 +635,7 @@ int fwd_attention(mlx_array *out, mlx_array x, int layer,
         /* K/V owning layer */
         int offset = kvcache_layer_offset(kv, layer);
         assert(offset >= 0);
-        if (fwd_rope_apply(&q_roped, q_transposed, hd_l, cfg, layer,
+        if (fwd_rope_apply(&q_roped, q_transposed, rope_dims, cfg, layer,
                            offset, rope_freqs, s) != 0)
             goto cleanup;
 
@@ -642,7 +698,7 @@ int fwd_attention(mlx_array *out, mlx_array x, int layer,
             goto cleanup;
 
         /* RoPE on K */
-        if (fwd_rope_apply(&k_roped, k_transposed, hd_l, cfg, layer,
+        if (fwd_rope_apply(&k_roped, k_transposed, rope_dims, cfg, layer,
                            offset, rope_freqs, s) != 0)
             goto cleanup;
 
@@ -682,6 +738,21 @@ int fwd_attention(mlx_array *out, mlx_array x, int layer,
     if (!MLXB_CHECK(mlx_reshape(&attn_reshaped, attn_back, out_shape, 3, s)))
         goto cleanup;
 
+    /* Optional output gate: o_proj(attn * sigmoid(gate)) */
+    if (cfg->attn_output_gate) {
+        mlx_array gate_sig = mlx_array_new();
+        mlx_array gated = mlx_array_new();
+        if (!MLXB_CHECK(mlx_sigmoid(&gate_sig, gate, s)) ||
+            !MLXB_CHECK(mlx_multiply(&gated, attn_reshaped, gate_sig, s))) {
+            mlx_array_free(gate_sig);
+            mlx_array_free(gated);
+            goto cleanup;
+        }
+        mlx_array_free(gate_sig);
+        mlx_array_free(attn_reshaped);
+        attn_reshaped = gated;
+    }
+
     /* O projection */
     weights_tensor_name(name, sizeof(name), cfg, layer, "self_attn.o_proj");
     if (weights_get_triplet(&o_tri, w, name) != 0) goto cleanup;
@@ -697,6 +768,7 @@ int fwd_attention(mlx_array *out, mlx_array x, int layer,
     weights_triplet_free(&v_tri);
     weights_triplet_free(&k_tri);
     weights_triplet_free(&q_tri);
+    mlx_array_free(gate);
     mlx_array_free(vnorm_w);
     mlx_array_free(result);
     mlx_array_free(attn_reshaped);

src/model/model.c

@@ -355,7 +355,10 @@ static int apply_family_defaults(model_config_t *cfg, yyjson_val *cfg_obj,
         cfg->has_qk_norm          = true;
         cfg->hidden_act           = HIDDEN_ACT_SILU;
         cfg->has_sliding_window   = false;
-        cfg->attn_output_gate     = true;
+        /* Default true only when key absent; honor explicit false. Gate rejects
+           false later (no oracle / real checkpoint for gate-off). */
+        if (yyjson_obj_get(cfg_obj, "attn_output_gate") == NULL)
+            cfg->attn_output_gate = true;
         cfg->rope_scaling_factor  = 1.0f;
         cfg->rope_local_base_freq = cfg->rope_theta;
         if (yyjson_obj_get(cfg_obj, "head_dim") == NULL &&

src/model/weights.c

@@ -262,6 +262,16 @@ static const weights_family_desc_t family_descs[] = {
         .layer_biases   = NULL,
         .extra_tensors  = NULL,
     },
+    {
+        /* qwen3_5 dense reuses the qwen3 name set; attn_output_gate doubles
+           q_proj out-features but does not add a separate tensor name. */
+        .family         = MODEL_QWEN3_5,
+        .layer_matmuls  = dense_std_layer_matmuls,
+        .layer_norms    = dense_std_layer_norms,
+        .layer_qk_norms = dense_std_layer_qk_norms,
+        .layer_biases   = NULL,
+        .extra_tensors  = NULL,
+    },
     {
         .family         = MODEL_LLAMA,
         .layer_matmuls  = dense_std_layer_matmuls,

tests/fixtures/model_config_qwen3_5_dense_full/config.json

@@ -0,0 +1,22 @@
+{
+  "model_type": "qwen3_5",
+  "vocab_size": 248320,
+  "hidden_size": 1024,
+  "num_hidden_layers": 4,
+  "num_attention_heads": 8,
+  "num_key_value_heads": 2,
+  "head_dim": 256,
+  "intermediate_size": 3584,
+  "max_position_embeddings": 262144,
+  "rms_norm_eps": 1e-06,
+  "hidden_act": "silu",
+  "attention_bias": false,
+  "attn_output_gate": true,
+  "full_attention_interval": 1,
+  "tie_word_embeddings": true,
+  "rope_parameters": {
+    "rope_type": "default",
+    "rope_theta": 10000000.0,
+    "partial_rotary_factor": 0.25
+  }
+}

tests/fixtures/model_config_qwen3_5_gate_off/config.json

@@ -0,0 +1,21 @@
+{
+  "model_type": "qwen3_5",
+  "vocab_size": 248320,
+  "hidden_size": 1024,
+  "num_hidden_layers": 4,
+  "num_attention_heads": 8,
+  "num_key_value_heads": 2,
+  "head_dim": 256,
+  "intermediate_size": 3584,
+  "max_position_embeddings": 262144,
+  "rms_norm_eps": 1e-06,
+  "hidden_act": "silu",
+  "attention_bias": false,
+  "attn_output_gate": false,
+  "tie_word_embeddings": true,
+  "rope_parameters": {
+    "rope_type": "default",
+    "rope_theta": 10000000.0,
+    "partial_rotary_factor": 0.25
+  }
+}

tests/fixtures/model_config_qwen3_5_hybrid_dense/config.json

@@ -0,0 +1,31 @@
+{
+  "model_type": "qwen3_5",
+  "vocab_size": 248320,
+  "hidden_size": 1024,
+  "num_hidden_layers": 4,
+  "num_attention_heads": 8,
+  "num_key_value_heads": 2,
+  "head_dim": 256,
+  "intermediate_size": 3584,
+  "max_position_embeddings": 262144,
+  "rms_norm_eps": 1e-06,
+  "hidden_act": "silu",
+  "attn_output_gate": true,
+  "full_attention_interval": 4,
+  "linear_conv_kernel_dim": 4,
+  "linear_key_head_dim": 128,
+  "linear_num_key_heads": 16,
+  "linear_num_value_heads": 16,
+  "linear_value_head_dim": 128,
+  "layer_types": [
+    "linear_attention",
+    "linear_attention",
+    "linear_attention",
+    "full_attention"
+  ],
+  "rope_parameters": {
+    "rope_type": "default",
+    "rope_theta": 10000000.0,
+    "partial_rotary_factor": 0.25
+  }
+}