diff --git a/scripts/parity_family.sh b/scripts/parity_family.sh index cffc195..16bdd82 100755 --- a/scripts/parity_family.sh +++ b/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 diff --git a/src/engine/emodel.c b/src/engine/emodel.c index 9783b4e..e7a1ca5 100644 --- a/src/engine/emodel.c +++ b/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)"); } } diff --git a/src/engine/forward.c b/src/engine/forward.c index 23ca3e4..3af30f0 100644 --- a/src/engine/forward.c +++ b/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,6 +533,7 @@ 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)); @@ -523,14 +541,52 @@ int fwd_attention(mlx_array *out, mlx_array x, int layer, 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); diff --git a/src/model/model.c b/src/model/model.c index 68f0cf6..46422d3 100644 --- a/src/model/model.c +++ b/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 && diff --git a/src/model/weights.c b/src/model/weights.c index e60c9d8..14ec550 100644 --- a/src/model/weights.c +++ b/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, diff --git a/tests/fixtures/model_config_qwen3_5_dense_full/config.json b/tests/fixtures/model_config_qwen3_5_dense_full/config.json new file mode 100644 index 0000000..d6ec395 --- /dev/null +++ b/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 + } +} diff --git a/tests/fixtures/model_config_qwen3_5_gate_off/config.json b/tests/fixtures/model_config_qwen3_5_gate_off/config.json new file mode 100644 index 0000000..102db91 --- /dev/null +++ b/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 + } +} diff --git a/tests/fixtures/model_config_qwen3_5_hybrid_dense/config.json b/tests/fixtures/model_config_qwen3_5_hybrid_dense/config.json new file mode 100644 index 0000000..8528dd8 --- /dev/null +++ b/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 + } +} diff --git a/tests/fixtures/tiny_qwen3_5/config.json b/tests/fixtures/tiny_qwen3_5/config.json new file mode 100644 index 0000000..92ba300 --- /dev/null +++ b/tests/fixtures/tiny_qwen3_5/config.json @@ -0,0 +1,19 @@ +{ + "model_type": "qwen3_5", + "vocab_size": 256, + "hidden_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 16, + "intermediate_size": 128, + "max_position_embeddings": 512, + "rms_norm_eps": 0.000001, + "tie_word_embeddings": false, + "attn_output_gate": true, + "full_attention_interval": 1, + "rope_parameters": { + "rope_theta": 10000000.0, + "partial_rotary_factor": 0.25 + } +} diff --git a/tests/fixtures/tiny_qwen3_5/model.safetensors b/tests/fixtures/tiny_qwen3_5/model.safetensors new file mode 100644 index 0000000..603396e Binary files /dev/null and b/tests/fixtures/tiny_qwen3_5/model.safetensors differ diff --git a/tests/fixtures/tiny_qwen3_5_tied/config.json b/tests/fixtures/tiny_qwen3_5_tied/config.json new file mode 100644 index 0000000..dd85e6c --- /dev/null +++ b/tests/fixtures/tiny_qwen3_5_tied/config.json @@ -0,0 +1,24 @@ +{ + "model_type": "qwen3_5", + "vocab_size": 256, + "hidden_size": 64, + "num_hidden_layers": 2, + "num_attention_heads": 4, + "num_key_value_heads": 2, + "head_dim": 16, + "intermediate_size": 128, + "max_position_embeddings": 512, + "rms_norm_eps": 0.000001, + "tie_word_embeddings": true, + "attn_output_gate": true, + "full_attention_interval": 1, + "rope_parameters": { + "rope_theta": 10000000.0, + "partial_rotary_factor": 0.25 + }, + "quantization": { + "bits": 4, + "group_size": 32, + "mode": "affine" + } +} diff --git a/tests/fixtures/tiny_qwen3_5_tied/model.safetensors b/tests/fixtures/tiny_qwen3_5_tied/model.safetensors new file mode 100644 index 0000000..e7e2990 Binary files /dev/null and b/tests/fixtures/tiny_qwen3_5_tied/model.safetensors differ diff --git a/tests/gpu_test_util.h b/tests/gpu_test_util.h new file mode 100644 index 0000000..b6a05a9 --- /dev/null +++ b/tests/gpu_test_util.h @@ -0,0 +1,29 @@ +#ifndef MLXD_GPU_TEST_UTIL_H +#define MLXD_GPU_TEST_UTIL_H + +#include "mlxbridge/mlxbridge.h" + +#include + +/* Shared finite-check helper for GPU forward tests. */ +static int is_finite_f32(mlx_array a, mlx_stream s) { + mlx_array f32 = mlx_array_new(); + if (!MLXB_CHECK(mlx_astype(&f32, a, MLX_FLOAT32, s))) { + mlx_array_free(f32); + return 0; + } + if (!MLXB_CHECK(mlx_array_eval(f32))) { + mlx_array_free(f32); + return 0; + } + size_t n = mlx_array_size(f32); + const float *d = mlx_array_data_float32(f32); + if (!d) { mlx_array_free(f32); return 0; } + for (size_t i = 0; i < n; i++) { + if (!isfinite(d[i])) { mlx_array_free(f32); return 0; } + } + mlx_array_free(f32); + return 1; +} + +#endif /* MLXD_GPU_TEST_UTIL_H */ diff --git a/tests/test_emodel_gate.c b/tests/test_emodel_gate.c index 781894f..567b767 100644 --- a/tests/test_emodel_gate.c +++ b/tests/test_emodel_gate.c @@ -30,7 +30,8 @@ static void test_reject_unsupported_family(void) { cfg.family = MODEL_GEMMA3; char err[256] = {0}; assert(engine_model_check_supported(&cfg, err, sizeof(err)) != 0); - assert(strcmp(err, "unsupported model family (only qwen3/llama/gemma4 dense supported)") == 0); + assert(strcmp(err, + "unsupported model family (only qwen3/llama/gemma4/qwen3_5 dense supported)") == 0); } static void test_reject_attention_bias(void) { @@ -124,7 +125,7 @@ static void test_reject_attn_output_gate(void) { static void test_reject_all_other_families(void) { static const model_family_t others[] = { MODEL_FAMILY_UNKNOWN, MODEL_GEMMA3, - MODEL_QWEN2, MODEL_QWEN3_5, MODEL_QWEN3_5_MOE, + MODEL_QWEN2, MODEL_QWEN3_5_MOE, MODEL_MISTRAL, MODEL_LFM2, MODEL_NEMOTRON_H, MODEL_DEEPSEEK_V4, MODEL_BERT, }; @@ -133,7 +134,158 @@ static void test_reject_all_other_families(void) { cfg.family = others[i]; char err[256] = {0}; assert(engine_model_check_supported(&cfg, err, sizeof(err)) != 0); - assert(strcmp(err, "unsupported model family (only qwen3/llama/gemma4 dense supported)") == 0); + assert(strcmp(err, + "unsupported model family (only qwen3/llama/gemma4/qwen3_5 dense supported)") == 0); + } +} + +/* --- qwen3_5 pure-dense gate tests (Stage D6 / R1) ----------------------- */ + +static model_config_t make_qwen3_5_supported(void) { + model_config_t cfg; + memset(&cfg, 0, sizeof(cfg)); + cfg.family = MODEL_QWEN3_5; + cfg.hidden_act = HIDDEN_ACT_SILU; + cfg.attn_output_gate = true; + cfg.partial_rotary_factor = 0.25f; + cfg.head_dim = 16; /* 16 * 0.25 = 4, even and >= 2 */ + cfg.has_qk_norm = true; + cfg.num_experts = 0; + cfg.linear_num_key_heads = 0; + cfg.linear_num_value_heads = 0; + cfg.full_attention_interval = 0; + return cfg; +} + +static void test_qwen3_5_pure_dense_passes(void) { + model_config_t cfg = make_qwen3_5_supported(); + char err[256] = {0}; + assert(cfg.attn_output_gate == true); + assert(engine_model_check_supported(&cfg, err, sizeof(err)) == 0); + assert(err[0] == '\0'); +} + +static void test_qwen3_5_rejects_gate_off(void) { + model_config_t cfg = make_qwen3_5_supported(); + cfg.attn_output_gate = false; + char err[256] = {0}; + assert(engine_model_check_supported(&cfg, err, sizeof(err)) != 0); + assert(strstr(err, "qwen3_5 requires attn_output_gate") != NULL); +} + +static void test_qwen3_5_full_attention_interval(void) { + model_config_t cfg = make_qwen3_5_supported(); + char err[256] = {0}; + + /* interval > 1 implies hybrid layers (is_linear when (idx+1)%interval != 0) */ + cfg.full_attention_interval = 4; + assert(engine_model_check_supported(&cfg, err, sizeof(err)) != 0); + assert(strstr(err, "full_attention_interval") != NULL); + + /* interval == 1 is pure dense (every layer full-attn) */ + cfg.full_attention_interval = 1; + memset(err, 0, sizeof(err)); + assert(engine_model_check_supported(&cfg, err, sizeof(err)) == 0); + + /* interval == 0 accepted (unset / no hybrid schedule) */ + cfg.full_attention_interval = 0; + memset(err, 0, sizeof(err)); + assert(engine_model_check_supported(&cfg, err, sizeof(err)) == 0); +} + +static void test_qwen3_5_rejects_odd_rope_dims(void) { + model_config_t cfg = make_qwen3_5_supported(); + /* head_dim=6, prf=0.5 -> dims=3 (odd); mx.fast.rope requires even >= 2 */ + cfg.head_dim = 6; + cfg.partial_rotary_factor = 0.5f; + char err[256] = {0}; + assert(engine_model_check_supported(&cfg, err, sizeof(err)) != 0); + assert(strstr(err, "partial_rotary_factor yields invalid rope dims") != NULL); + + /* even product stays accepted (16 * 0.25 = 4) */ + cfg.head_dim = 16; + cfg.partial_rotary_factor = 0.25f; + memset(err, 0, sizeof(err)); + assert(engine_model_check_supported(&cfg, err, sizeof(err)) == 0); +} + +static void test_qwen3_5_moe_family_still_rejected(void) { + model_config_t cfg = make_qwen3_5_supported(); + cfg.family = MODEL_QWEN3_5_MOE; + char err[256] = {0}; + assert(engine_model_check_supported(&cfg, err, sizeof(err)) != 0); + assert(strcmp(err, + "unsupported model family (only qwen3/llama/gemma4/qwen3_5 dense supported)") == 0); +} + +static void test_qwen3_5_rejects_linear_heads(void) { + model_config_t cfg = make_qwen3_5_supported(); + cfg.linear_num_key_heads = 16; + char err[256] = {0}; + assert(engine_model_check_supported(&cfg, err, sizeof(err)) != 0); + assert(strstr(err, "linear attention") != NULL); +} + +static void test_qwen3_5_rejects_experts(void) { + model_config_t cfg = make_qwen3_5_supported(); + cfg.num_experts = 8; + char err[256] = {0}; + assert(engine_model_check_supported(&cfg, err, sizeof(err)) != 0); + assert(strcmp(err, "MoE models not supported") == 0); +} + +static void test_qwen3_5_rejects_partial_rotary_out_of_range(void) { + model_config_t cfg = make_qwen3_5_supported(); + cfg.partial_rotary_factor = 0.0f; + char err[256] = {0}; + assert(engine_model_check_supported(&cfg, err, sizeof(err)) != 0); + assert(strstr(err, "partial_rotary_factor") != NULL); + + cfg.partial_rotary_factor = 1.5f; + memset(err, 0, sizeof(err)); + assert(engine_model_check_supported(&cfg, err, sizeof(err)) != 0); + assert(strstr(err, "partial_rotary_factor") != NULL); +} + +static void test_qwen3_5_loaded_boundary_fixtures(void) { + /* pure-dense full fixture -> accept */ + { + model_config_t cfg; + int rc = model_config_load( + &cfg, MLXD_FIXTURES_DIR "/model_config_qwen3_5_dense_full"); + assert(rc == 0); + assert(cfg.family == MODEL_QWEN3_5); + char err[256] = {0}; + assert(engine_model_check_supported(&cfg, err, sizeof(err)) == 0); + model_config_free(&cfg); + } + + /* moe fixture -> family MOE -> reject */ + { + model_config_t cfg; + int rc = model_config_load( + &cfg, MLXD_FIXTURES_DIR "/model_config_qwen3_5_moe"); + assert(rc == 0); + assert(cfg.family == MODEL_QWEN3_5_MOE); + char err[256] = {0}; + assert(engine_model_check_supported(&cfg, err, sizeof(err)) != 0); + model_config_free(&cfg); + } + + /* hybrid-dense (linear heads, num_experts==0) stays MODEL_QWEN3_5 under R1 + but gate rejects linear attention */ + { + model_config_t cfg; + int rc = model_config_load( + &cfg, MLXD_FIXTURES_DIR "/model_config_qwen3_5_hybrid_dense"); + assert(rc == 0); + assert(cfg.family == MODEL_QWEN3_5); + assert(cfg.num_experts == 0); + assert(cfg.linear_num_key_heads > 0); + char err[256] = {0}; + assert(engine_model_check_supported(&cfg, err, sizeof(err)) != 0); + assert(strstr(err, "linear attention") != NULL); + model_config_free(&cfg); } } @@ -285,6 +437,14 @@ static void test_gemma4_reject_proportional_factor_zero(void) { assert(strstr(err, "rope_proportional_factor") != NULL); } +static void test_gemma4_reject_partial_rotary_plain(void) { + model_config_t cfg = make_gemma4_supported(); + cfg.partial_rotary_factor = 0.5f; + char err[256] = {0}; + assert(engine_model_check_supported(&cfg, err, sizeof(err)) != 0); + assert(strstr(err, "partial rotary") != NULL); +} + static void test_gemma4_reject_partial_rotary_global_zero(void) { model_config_t cfg = make_gemma4_supported(); cfg.partial_rotary_factor_global = 0.0f; @@ -432,6 +592,9 @@ int main(void) { test_gemma4_reject_proportional_factor_zero(); printf(" test_gemma4_reject_proportional_factor_zero: passed\n"); + test_gemma4_reject_partial_rotary_plain(); + printf(" test_gemma4_reject_partial_rotary_plain: passed\n"); + test_gemma4_reject_partial_rotary_global_zero(); printf(" test_gemma4_reject_partial_rotary_global_zero: passed\n"); @@ -534,6 +697,33 @@ int main(void) { test_llama_llama3_valid_still_passes(); printf(" test_llama_llama3_valid_still_passes: passed\n"); + test_qwen3_5_pure_dense_passes(); + printf(" test_qwen3_5_pure_dense_passes: passed\n"); + + test_qwen3_5_rejects_gate_off(); + printf(" test_qwen3_5_rejects_gate_off: passed\n"); + + test_qwen3_5_full_attention_interval(); + printf(" test_qwen3_5_full_attention_interval: passed\n"); + + test_qwen3_5_rejects_odd_rope_dims(); + printf(" test_qwen3_5_rejects_odd_rope_dims: passed\n"); + + test_qwen3_5_moe_family_still_rejected(); + printf(" test_qwen3_5_moe_family_still_rejected: passed\n"); + + test_qwen3_5_rejects_linear_heads(); + printf(" test_qwen3_5_rejects_linear_heads: passed\n"); + + test_qwen3_5_rejects_experts(); + printf(" test_qwen3_5_rejects_experts: passed\n"); + + test_qwen3_5_rejects_partial_rotary_out_of_range(); + printf(" test_qwen3_5_rejects_partial_rotary_out_of_range: passed\n"); + + test_qwen3_5_loaded_boundary_fixtures(); + printf(" test_qwen3_5_loaded_boundary_fixtures: passed\n"); + printf("test_emodel_gate: all passed\n"); return 0; } diff --git a/tests/test_forward_gemma3_gpu.c b/tests/test_forward_gemma3_gpu.c index 65e057d..cc64691 100644 --- a/tests/test_forward_gemma3_gpu.c +++ b/tests/test_forward_gemma3_gpu.c @@ -4,6 +4,7 @@ #include "mlxbridge/mlxbridge.h" #include "model/model.h" #include "model/weights.h" +#include "gpu_test_util.h" #include #include @@ -14,26 +15,6 @@ static mlx_stream gpu; -static int is_finite_f32(mlx_array a, mlx_stream s) { - mlx_array f32 = mlx_array_new(); - if (!MLXB_CHECK(mlx_astype(&f32, a, MLX_FLOAT32, s))) { - mlx_array_free(f32); - return 0; - } - if (!MLXB_CHECK(mlx_array_eval(f32))) { - mlx_array_free(f32); - return 0; - } - size_t n = mlx_array_size(f32); - const float *d = mlx_array_data_float32(f32); - if (!d) { mlx_array_free(f32); return 0; } - for (size_t i = 0; i < n; i++) { - if (!isfinite(d[i])) { mlx_array_free(f32); return 0; } - } - mlx_array_free(f32); - return 1; -} - /* Load engine_model_t bypassing the support gate (gemma3 not yet whitelisted). TODO: replace with engine_model_load once gemma3 is in engine_model_check_supported (D1). */ static int load_gemma3(engine_model_t *em, const char *path) { diff --git a/tests/test_forward_gemma4_gpu.c b/tests/test_forward_gemma4_gpu.c index d8c1814..63ddfde 100644 --- a/tests/test_forward_gemma4_gpu.c +++ b/tests/test_forward_gemma4_gpu.c @@ -4,6 +4,7 @@ #include "mlxbridge/mlxbridge.h" #include "model/model.h" #include "model/weights.h" +#include "gpu_test_util.h" #include #include @@ -18,26 +19,6 @@ static mlx_stream gpu; -static int is_finite_f32(mlx_array a, mlx_stream s) { - mlx_array f32 = mlx_array_new(); - if (!MLXB_CHECK(mlx_astype(&f32, a, MLX_FLOAT32, s))) { - mlx_array_free(f32); - return 0; - } - if (!MLXB_CHECK(mlx_array_eval(f32))) { - mlx_array_free(f32); - return 0; - } - size_t n = mlx_array_size(f32); - const float *d = mlx_array_data_float32(f32); - if (!d) { mlx_array_free(f32); return 0; } - for (size_t i = 0; i < n; i++) { - if (!isfinite(d[i])) { mlx_array_free(f32); return 0; } - } - mlx_array_free(f32); - return 1; -} - static float max_abs_diff_arr(mlx_array a, mlx_array b, mlx_stream s) { mlx_array af = mlx_array_new(); mlx_array bf = mlx_array_new(); diff --git a/tests/test_forward_llama_gpu.c b/tests/test_forward_llama_gpu.c index aad7926..d581ce9 100644 --- a/tests/test_forward_llama_gpu.c +++ b/tests/test_forward_llama_gpu.c @@ -4,36 +4,18 @@ #include "mlxbridge/mlxbridge.h" #include "model/model.h" #include "model/weights.h" +#include "gpu_test_util.h" #include #include #include +#include #include #define FIXTURES MLXD_FIXTURES_DIR static mlx_stream gpu; -static int is_finite_f32(mlx_array a, mlx_stream s) { - mlx_array f32 = mlx_array_new(); - if (!MLXB_CHECK(mlx_astype(&f32, a, MLX_FLOAT32, s))) { - mlx_array_free(f32); - return 0; - } - if (!MLXB_CHECK(mlx_array_eval(f32))) { - mlx_array_free(f32); - return 0; - } - size_t n = mlx_array_size(f32); - const float *d = mlx_array_data_float32(f32); - if (!d) { mlx_array_free(f32); return 0; } - for (size_t i = 0; i < n; i++) { - if (!isfinite(d[i])) { mlx_array_free(f32); return 0; } - } - mlx_array_free(f32); - return 1; -} - /* ---- config flags ---- */ static void test_llama3_config_flags(engine_model_t *em) { @@ -253,7 +235,7 @@ static void run_suite(const char *label, const char *path, char err[256] = {0}; if (engine_model_load(&em, path, err, sizeof(err)) != 0) { fprintf(stderr, "failed to load %s: %s\n", path, err); - assert(0); + abort(); } config_test(&em); diff --git a/tests/test_forward_qwen3_5_gpu.c b/tests/test_forward_qwen3_5_gpu.c new file mode 100644 index 0000000..21a6d09 --- /dev/null +++ b/tests/test_forward_qwen3_5_gpu.c @@ -0,0 +1,214 @@ +#include "engine/emodel.h" +#include "engine/forward.h" +#include "engine/kvcache.h" +#include "mlxbridge/mlxbridge.h" +#include "model/model.h" +#include "model/weights.h" +#include "gpu_test_util.h" + +#include +#include +#include +#include +#include + +#define FIXTURES MLXD_FIXTURES_DIR + +static mlx_stream gpu; + +static void test_qwen3_5_config_flags(engine_model_t *em) { + assert(em->cfg.family == MODEL_QWEN3_5); + assert(em->cfg.attn_output_gate == true); + assert(em->cfg.partial_rotary_factor == 0.25f); + assert(em->cfg.has_qk_norm == true); + assert(em->cfg.hidden_act == HIDDEN_ACT_SILU); + assert(em->cfg.tie_word_embeddings == false); + assert(em->cfg.head_dim == 16); + assert(em->cfg.rope_theta == 10000000.0f); + assert(em->cfg.linear_num_key_heads == 0); + assert(em->cfg.num_experts == 0); + /* plain theta rope - no custom freqs array */ + assert(em->rope_freqs.ctx == NULL); +} + +static void test_qwen3_5_tied_config_flags(engine_model_t *em) { + assert(em->cfg.family == MODEL_QWEN3_5); + assert(em->cfg.tie_word_embeddings == true); + assert(em->cfg.attn_output_gate == true); + assert(em->cfg.partial_rotary_factor == 0.25f); + assert(em->cfg.has_qk_norm == true); + assert(em->cfg.quant_bits == 4); +} + +static void test_qwen3_5_weight_mapping(engine_model_t *em) { + /* Same tensor count as tiny_qwen3 dense (gate is shape, not extra name). */ + assert(em->w.count == 25); + + mlx_array probe = mlx_array_new(); + assert(weights_get(&probe, &em->w, + "model.layers.0.self_attn.q_proj.weight") == 0); + /* Doubled q_proj: [2*H*D, hidden] = [128, 64] */ + assert(mlx_array_ndim(probe) == 2); + assert(mlx_array_dim(probe, 0) == 2 * em->cfg.num_attention_heads * + em->cfg.head_dim); + assert(mlx_array_dim(probe, 1) == em->cfg.hidden_size); + mlx_array_free(probe); + + probe = mlx_array_new(); + assert(weights_get(&probe, &em->w, "lm_head.weight") == 0); + mlx_array_free(probe); + + probe = mlx_array_new(); + assert(weights_get(&probe, &em->w, + "model.layers.0.linear_attn.in_proj_qkv.weight") != 0); + mlx_array_free(probe); +} + +static void test_qwen3_5_tied_weight_mapping(engine_model_t *em) { + /* tied quantized with qk_norm: + embed(3) + norm(1) + 2*(7*3 matmul + 2 ln + 2 qk) = 4 + 2*25 = 54 */ + assert(em->w.count == 54); + + mlx_array probe = mlx_array_new(); + assert(weights_get(&probe, &em->w, "lm_head.weight") != 0); + mlx_array_free(probe); +} + +static void test_qwen3_5_forward_shape(engine_model_t *em) { + int32_t ids_data[] = {1, 2, 3, 4, 5, 6}; + int ids_shape[] = {1, 6}; + mlx_array ids = mlx_array_new_data(ids_data, ids_shape, 2, MLX_INT32); + + kvcache_t kv; + assert(kvcache_init(&kv, em->cfg.num_hidden_layers) == 0); + + mlx_array logits = mlx_array_new(); + assert(model_forward(em, ids, &kv, true, &logits) == 0); + assert(MLXB_CHECK(mlx_array_eval(logits))); + assert(mlx_array_ndim(logits) == 2); + assert(mlx_array_dim(logits, 0) == 1); + assert(mlx_array_dim(logits, 1) == 256); + assert(is_finite_f32(logits, gpu)); + + mlx_array_free(logits); + kvcache_free(&kv); + mlx_array_free(ids); +} + +static void test_qwen3_5_incremental_equals_full(engine_model_t *em) { + int32_t prompt[] = {10, 20, 30, 40, 50, 60}; + int P = 6; + + kvcache_t kv_a; + assert(kvcache_init(&kv_a, em->cfg.num_hidden_layers) == 0); + + int shape_all[] = {1, P}; + mlx_array ids_all = mlx_array_new_data(prompt, shape_all, 2, MLX_INT32); + mlx_array logits_a = mlx_array_new(); + assert(model_forward(em, ids_all, &kv_a, true, &logits_a) == 0); + assert(MLXB_CHECK(mlx_array_eval(logits_a))); + + kvcache_t kv_b; + assert(kvcache_init(&kv_b, em->cfg.num_hidden_layers) == 0); + + int shape_pre[] = {1, P - 1}; + mlx_array ids_pre = mlx_array_new_data(prompt, shape_pre, 2, MLX_INT32); + assert(model_forward(em, ids_pre, &kv_b, false, NULL) == 0); + + int shape_dec[] = {1, 1}; + mlx_array ids_dec = mlx_array_new_data(&prompt[P - 1], shape_dec, 2, MLX_INT32); + mlx_array logits_b = mlx_array_new(); + assert(model_forward(em, ids_dec, &kv_b, true, &logits_b) == 0); + assert(MLXB_CHECK(mlx_array_eval(logits_b))); + + mlx_array la_f32 = mlx_array_new(); + mlx_array lb_f32 = mlx_array_new(); + MLXB_CHECK(mlx_astype(&la_f32, logits_a, MLX_FLOAT32, gpu)); + MLXB_CHECK(mlx_astype(&lb_f32, logits_b, MLX_FLOAT32, gpu)); + assert(MLXB_CHECK(mlx_array_eval(la_f32))); + assert(MLXB_CHECK(mlx_array_eval(lb_f32))); + + const float *da = mlx_array_data_float32(la_f32); + const float *db = mlx_array_data_float32(lb_f32); + int vocab = em->cfg.vocab_size; + + int argmax_a = 0, argmax_b = 0; + for (int i = 1; i < vocab; i++) { + if (da[i] > da[argmax_a]) argmax_a = i; + if (db[i] > db[argmax_b]) argmax_b = i; + } + assert(argmax_a == argmax_b); + + /* Compare: argmax match + max-abs-diff < 5e-2. bf16 prefill-vs-decode + accumulation drift; same tolerance as sibling GPU tests. */ + float maxdiff = 0.0f; + for (int i = 0; i < vocab; i++) { + float d = fabsf(da[i] - db[i]); + if (d > maxdiff) maxdiff = d; + } + assert(maxdiff < 5e-2f); + + mlx_array_free(lb_f32); + mlx_array_free(la_f32); + mlx_array_free(logits_b); + mlx_array_free(ids_dec); + mlx_array_free(ids_pre); + kvcache_free(&kv_b); + mlx_array_free(logits_a); + mlx_array_free(ids_all); + kvcache_free(&kv_a); +} + +static void test_hybrid_config_load_rejected(void) { + engine_model_t em; + char err[256] = {0}; + int rc = engine_model_load( + &em, FIXTURES "/model_config_qwen3_5_hybrid_dense", err, sizeof(err)); + assert(rc != 0); + assert(strstr(err, "linear attention") != NULL); +} + +static void run_suite(const char *label, const char *path, + void (*config_test)(engine_model_t *), + void (*weight_test)(engine_model_t *)) { + engine_model_t em; + char err[256] = {0}; + if (engine_model_load(&em, path, err, sizeof(err)) != 0) { + fprintf(stderr, "failed to load %s: %s\n", path, err); + abort(); + } + + config_test(&em); + printf(" %s config_flags: passed\n", label); + + weight_test(&em); + printf(" %s weight_mapping: passed\n", label); + + test_qwen3_5_forward_shape(&em); + printf(" %s forward_shape: passed\n", label); + + test_qwen3_5_incremental_equals_full(&em); + printf(" %s incremental_equals_full: passed\n", label); + + engine_model_free(&em); +} + +int main(void) { + gpu = mlxbridge_gpu_stream(); + + test_hybrid_config_load_rejected(); + printf(" test_hybrid_config_load_rejected: passed\n"); + + char dense_path[512], tied_path[512]; + snprintf(dense_path, sizeof(dense_path), "%s/tiny_qwen3_5", FIXTURES); + snprintf(tied_path, sizeof(tied_path), "%s/tiny_qwen3_5_tied", FIXTURES); + + run_suite("qwen3_5_dense", dense_path, + test_qwen3_5_config_flags, test_qwen3_5_weight_mapping); + + run_suite("qwen3_5_tied", tied_path, + test_qwen3_5_tied_config_flags, test_qwen3_5_tied_weight_mapping); + + printf("test_forward_qwen3_5_gpu: all passed\n"); + return 0; +} diff --git a/tests/test_model_config.c b/tests/test_model_config.c index e74a4dd..4a6ffd1 100644 --- a/tests/test_model_config.c +++ b/tests/test_model_config.c @@ -435,6 +435,78 @@ static void test_qwen3_5_dense_vs_moe(void) { model_config_free(&cfg); } +/* Stage D6: pure-dense qwen3_5 shape + hybrid-dense R1 boundary */ +static void test_qwen3_5_dense_full_and_hybrid(void) { + model_config_t cfg; + int rc = model_config_load( + &cfg, MLXD_FIXTURES_DIR "/model_config_qwen3_5_dense_full"); + assert(rc == 0); + assert(cfg.family == MODEL_QWEN3_5); + assert(cfg.attn_output_gate == true); + assert(cfg.partial_rotary_factor == 0.25f); + assert(cfg.has_qk_norm == true); + assert(cfg.num_experts == 0); + assert(cfg.linear_num_key_heads == 0); + assert(cfg.linear_num_value_heads == 0); + assert(cfg.head_dim == 256); + assert(cfg.rope_theta == 10000000.0f); + assert(cfg.full_attention_interval == 1); + assert(cfg.has_explicit_layer_types == false); + model_config_free(&cfg); + + /* R1: hybrid linear fields do not reclassify as MOE */ + rc = model_config_load( + &cfg, MLXD_FIXTURES_DIR "/model_config_qwen3_5_hybrid_dense"); + assert(rc == 0); + assert(cfg.family == MODEL_QWEN3_5); + assert(cfg.num_experts == 0); + assert(cfg.linear_num_key_heads == 16); + assert(cfg.linear_num_value_heads == 16); + assert(cfg.attn_output_gate == true); + assert(cfg.partial_rotary_factor == 0.25f); + assert(cfg.has_explicit_layer_types == true); + assert(cfg.layer_is_global[0] == false); /* linear_attention */ + assert(cfg.layer_is_global[3] == true); /* full_attention */ + model_config_free(&cfg); +} + +/* Explicit attn_output_gate:false must survive parse (not clobbered by defaults). + Absent key still defaults to true (covered by model_config_qwen3_5). */ +/* tiny_qwen3_5 fixtures emit full_attention_interval:1 and rope_theta only + under rope_parameters (no flat duplicate). */ +static void test_tiny_qwen3_5_interval_and_rope(void) { + model_config_t cfg; + int rc = model_config_load(&cfg, MLXD_FIXTURES_DIR "/tiny_qwen3_5"); + assert(rc == 0); + assert(cfg.family == MODEL_QWEN3_5); + assert(cfg.full_attention_interval == 1); + assert(cfg.rope_theta == 10000000.0f); + assert(cfg.partial_rotary_factor == 0.25f); + model_config_free(&cfg); + + rc = model_config_load(&cfg, MLXD_FIXTURES_DIR "/tiny_qwen3_5_tied"); + assert(rc == 0); + assert(cfg.full_attention_interval == 1); + assert(cfg.rope_theta == 10000000.0f); + model_config_free(&cfg); +} + +static void test_qwen3_5_attn_output_gate_explicit_false(void) { + model_config_t cfg; + int rc = model_config_load( + &cfg, MLXD_FIXTURES_DIR "/model_config_qwen3_5_gate_off"); + assert(rc == 0); + assert(cfg.family == MODEL_QWEN3_5); + assert(cfg.attn_output_gate == false); + model_config_free(&cfg); + + /* Absent key: default stays true */ + rc = model_config_load(&cfg, MLXD_FIXTURES_DIR "/model_config_qwen3_5"); + assert(rc == 0); + assert(cfg.attn_output_gate == true); + model_config_free(&cfg); +} + /* --- A1 Cycle 7: nested text_config -------------------------------------- */ static void test_nested_text_config(void) { @@ -1099,6 +1171,9 @@ int main(void) { test_family_from_type(); test_family_defaults(); test_qwen3_5_dense_vs_moe(); + test_qwen3_5_dense_full_and_hybrid(); + test_qwen3_5_attn_output_gate_explicit_false(); + test_tiny_qwen3_5_interval_and_rope(); test_nested_text_config(); test_sliding_window_and_layers(); test_stage_e_dims(); diff --git a/tests/test_parity_script.sh b/tests/test_parity_script.sh index 366dccf..6503c75 100644 --- a/tests/test_parity_script.sh +++ b/tests/test_parity_script.sh @@ -625,13 +625,33 @@ else fail "wrapper-unknown-family-lfm2" "expected exit 2 + usage (rc=$rc)" fi -# --- wrapper: skip on TBD id (gemma4) --- +# --- wrapper: skip when canonical id present but dir absent (gemma4, qwen3_5) --- out=$(MLXD_PARITY_CKPT_ROOT=/tmp sh "$WRAPPER" gemma4 2>&1) && rc=0 || rc=$? -if [ "$rc" -eq 0 ] && printf '%s\n' "$out" | grep -q 'skipped.*no canonical checkpoint id'; then - pass "wrapper-skip-tbd-gemma4" +if [ "$rc" -eq 0 ] && printf '%s\n' "$out" | grep -q 'skipped.*checkpoint dir absent'; then + pass "wrapper-skip-absent-gemma4" +else + fail "wrapper-skip-absent-gemma4" "expected exit 0 + skipped:dir absent (rc=$rc)" +fi + +out=$(MLXD_PARITY_CKPT_ROOT=/tmp sh "$WRAPPER" qwen3_5 2>&1) && rc=0 || rc=$? +if [ "$rc" -eq 0 ] && printf '%s\n' "$out" | grep -q 'skipped.*checkpoint dir absent'; then + pass "wrapper-skip-absent-qwen3_5" +else + fail "wrapper-skip-absent-qwen3_5" "expected exit 0 + skipped:dir absent (rc=$rc)" +fi + +# --- wrapper: qwen3_5 hard-skips even when hybrid snapshot dir is present --- +# Real Qwen3.5 checkpoints are hybrid; Stage D rejects them. Keep the recorded +# canonical id but never invoke parity_temp0.sh via the root path. +W_Q35_DIR=$(mktemp -d) +mkdir -p "$W_Q35_DIR/root/mlx-community/Qwen3.5-0.8B-4bit" +out=$(MLXD_PARITY_CKPT_ROOT="$W_Q35_DIR/root" sh "$WRAPPER" qwen3_5 2>&1) && rc=0 || rc=$? +rm -rf "$W_Q35_DIR" +if [ "$rc" -eq 0 ] && printf '%s\n' "$out" | grep -q 'skipped: qwen3_5 hybrid parity deferred to Stage E'; then + pass "wrapper-skip-present-qwen3_5-hybrid" else - fail "wrapper-skip-tbd-gemma4" "expected exit 0 + skipped:no canonical (rc=$rc)" + fail "wrapper-skip-present-qwen3_5-hybrid" "expected exit 0 + hybrid deferred skip (rc=$rc out=$out)" fi # --- wrapper: per-family override beats MLXD_PARITY_CKPT_ROOT (N3) --- diff --git a/tests/test_qwen3_5_goldens_gpu.c b/tests/test_qwen3_5_goldens_gpu.c new file mode 100644 index 0000000..a31f63a --- /dev/null +++ b/tests/test_qwen3_5_goldens_gpu.c @@ -0,0 +1,133 @@ +#include "engine/emodel.h" +#include "engine/forward.h" +#include "engine/kvcache.h" +#include "mlxbridge/mlxbridge.h" + +#include +#include +#include +#include +#include + +#define FIXTURES MLXD_FIXTURES_DIR + +/* mlx-lm 0.31.3 (mlx_lm.models.qwen3_5) last-token logits on ids [1,2,3,4,5]. + Generated against the tiny fixtures after full_attention_interval:1 was + added. Pins the D6 deltas: sigmoid gate multiply + partial rope dims. */ + +/* tiny_qwen3_5 (bf16 dense, untied lm_head) */ +static const float dense_first8[] = { + -1.75781250e-01f, -1.33789062e-01f, 1.60156250e-01f, -5.55419922e-03f, + -2.51953125e-01f, -4.02343750e-01f, 1.23046875e-01f, 2.69775391e-02f +}; +static const int dense_argmax = 33; + +/* tiny_qwen3_5_tied (4-bit affine quant, tied embeddings) */ +static const float tied_first8[] = { + 7.41506517e-02f, -1.21642277e-01f, -1.48075938e-01f, 1.20444193e-01f, + -6.43087775e-02f, 1.13397074e+00f, -5.17259240e-02f, 1.73729137e-01f +}; +static const int tied_argmax = 5; + +static mlx_stream gpu; + +static void compare_last_logits(engine_model_t *em, + const float *expected_first8, + int expected_argmax, + float atol, + const char *label) { + int32_t ids_data[] = {1, 2, 3, 4, 5}; + int ids_shape[] = {1, 5}; + mlx_array ids = mlx_array_new_data(ids_data, ids_shape, 2, MLX_INT32); + + kvcache_t kv; + assert(kvcache_init(&kv, em->cfg.num_hidden_layers) == 0); + + mlx_array logits = mlx_array_new(); + assert(model_forward(em, ids, &kv, true, &logits) == 0); + assert(MLXB_CHECK(mlx_array_eval(logits))); + + /* model_forward returns last-token logits [1, vocab] */ + assert(mlx_array_ndim(logits) == 2); + assert(mlx_array_dim(logits, 0) == 1); + assert(mlx_array_dim(logits, 1) == em->cfg.vocab_size); + + mlx_array f32 = mlx_array_new(); + MLXB_CHECK(mlx_astype(&f32, logits, MLX_FLOAT32, gpu)); + assert(MLXB_CHECK(mlx_array_eval(f32))); + + const float *d = mlx_array_data_float32(f32); + assert(d != NULL); + + int vocab = em->cfg.vocab_size; + int argmax = 0; + for (int i = 1; i < vocab; i++) { + if (d[i] > d[argmax]) argmax = i; + } + if (argmax != expected_argmax) { + fprintf(stderr, "%s: argmax got %d expected %d (got_val=%f exp_val=%f)\n", + label, argmax, expected_argmax, + (double)d[argmax], (double)d[expected_argmax]); + } + assert(argmax == expected_argmax); + + float max_diff = 0.0f; + for (int i = 0; i < 8; i++) { + float diff = fabsf(d[i] - expected_first8[i]); + if (diff > max_diff) max_diff = diff; + if (diff > atol) { + fprintf(stderr, "%s first8[%d]: got %f expected %f (diff %f)\n", + label, i, (double)d[i], (double)expected_first8[i], + (double)diff); + } + assert(diff <= atol); + } + printf(" %s: argmax=%d max_first8_diff=%g (atol=%g)\n", + label, argmax, (double)max_diff, (double)atol); + + mlx_array_free(f32); + mlx_array_free(logits); + kvcache_free(&kv); + mlx_array_free(ids); +} + +static void test_fixture_golden(const char *relpath, + const float *expected_first8, + int expected_argmax, + float atol, + const char *label) { + char path[512]; + snprintf(path, sizeof(path), "%s/%s", FIXTURES, relpath); + + engine_model_t em; + char err[256] = {0}; + if (engine_model_load(&em, path, err, sizeof(err)) != 0) { + fprintf(stderr, "failed to load %s: %s\n", path, err); + abort(); + } + + assert(em.cfg.family == MODEL_QWEN3_5); + assert(em.cfg.attn_output_gate == true); + assert(em.cfg.partial_rotary_factor == 0.25f); + assert(em.cfg.full_attention_interval == 1); + + compare_last_logits(&em, expected_first8, expected_argmax, atol, label); + engine_model_free(&em); +} + +int main(void) { + gpu = mlxbridge_gpu_stream(); + + /* bf16 dense: tight tolerance */ + test_fixture_golden("tiny_qwen3_5", dense_first8, dense_argmax, + 1e-2f, "dense_bf16"); + printf(" test_qwen3_5_dense_golden: passed\n"); + + /* 4-bit tied: looser tolerance for quant noise */ + test_fixture_golden("tiny_qwen3_5_tied", tied_first8, tied_argmax, + 5e-2f, "tied_4bit"); + printf(" test_qwen3_5_tied_golden: passed\n"); + + printf("test_qwen3_5_goldens_gpu: all passed\n"); + return 0; +} diff --git a/tests/test_weights.c b/tests/test_weights.c index 17531b0..c2f8a33 100644 --- a/tests/test_weights.c +++ b/tests/test_weights.c @@ -627,10 +627,64 @@ static void test_expected_names_gemma4(void) { assert(weights_expected_names(&cfg, small, 2) == -1); } +static void test_expected_names_qwen3_5(void) { + model_config_t cfg = {0}; + cfg.family = MODEL_QWEN3_5; + cfg.weight_prefix = "model"; + cfg.num_hidden_layers = 2; + cfg.has_qk_norm = true; + cfg.tie_word_embeddings = false; + cfg.attn_output_gate = true; + + /* Same name set as qwen3; gate is a q_proj shape convention, not a name. */ + model_config_t q3 = {0}; + q3.family = MODEL_QWEN3; + q3.weight_prefix = "model"; + q3.num_hidden_layers = 2; + q3.has_qk_norm = true; + q3.tie_word_embeddings = false; + + int count_q35 = weights_expected_names(&cfg, NULL, 0); + int count_q3 = weights_expected_names(&q3, NULL, 0); + assert(count_q35 == count_q3); + assert(count_q35 == 25); + + weight_expected_t names[25]; + int rc = weights_expected_names(&cfg, names, 25); + assert(rc == 25); + + bool found_q0 = false, found_lm = false, found_lin = false; + for (int i = 0; i < rc; i++) { + if (strcmp(names[i].name, "model.layers.0.self_attn.q_proj") == 0) { + assert(names[i].kind == WEIGHT_KIND_MATMUL); + found_q0 = true; + } + if (strcmp(names[i].name, "lm_head") == 0) { + found_lm = true; + } + if (strstr(names[i].name, "linear_attn") != NULL) { + found_lin = true; + } + } + assert(found_q0); + assert(found_lm); + assert(!found_lin); + + /* Tied: no lm_head */ + cfg.tie_word_embeddings = true; + int tied = weights_expected_names(&cfg, NULL, 0); + assert(tied == 24); + weight_expected_t tied_names[24]; + rc = weights_expected_names(&cfg, tied_names, 24); + assert(rc == 24); + for (int i = 0; i < rc; i++) + assert(strcmp(tied_names[i].name, "lm_head") != 0); +} + static void test_expected_names_other_families_zero(void) { static const model_family_t others[] = { MODEL_FAMILY_UNKNOWN, MODEL_GEMMA3, - MODEL_QWEN2, MODEL_QWEN3_5, MODEL_QWEN3_5_MOE, + MODEL_QWEN2, MODEL_QWEN3_5_MOE, MODEL_MISTRAL, MODEL_LFM2, MODEL_NEMOTRON_H, MODEL_DEEPSEEK_V4, MODEL_BERT, }; @@ -847,6 +901,9 @@ int main(void) { test_expected_names_gemma4(); printf(" test_expected_names_gemma4: passed\n"); + test_expected_names_qwen3_5(); + printf(" test_expected_names_qwen3_5: passed\n"); + test_expected_names_other_families_zero(); printf(" test_expected_names_other_families_zero: passed\n"); diff --git a/tools/gen_tiny_ckpt.c b/tools/gen_tiny_ckpt.c index ceeaa09..1b9cd91 100644 --- a/tools/gen_tiny_ckpt.c +++ b/tools/gen_tiny_ckpt.c @@ -54,6 +54,11 @@ typedef struct { double rope_high_freq_factor; int rope_original_max_position_embeddings; + /* qwen3_5 full-attn deltas (0 = unset / default) */ + bool attn_output_gate; + double partial_rotary_factor; /* emit under rope_parameters when > 0 */ + int full_attention_interval; /* 0 = do not emit; 1 = pure dense */ + const tensor_spec_t *top; int n_top; const tensor_spec_t *layer; int n_layer; @@ -72,7 +77,10 @@ static void resolve_shape(const recipe_t *r, const char *suffix, } else if (strcmp(suffix, "lm_head.weight") == 0) { shape[0] = r->vocab; shape[1] = r->hidden; *ndim = 2; } else if (strcmp(suffix, "self_attn.q_proj.weight") == 0) { - shape[0] = r->heads * r->head_dim; shape[1] = r->hidden; *ndim = 2; + /* attn_output_gate doubles q_proj out-features: [2*H*D, hidden] */ + int q_out = r->heads * r->head_dim; + if (r->attn_output_gate) q_out *= 2; + shape[0] = q_out; shape[1] = r->hidden; *ndim = 2; } else if (strcmp(suffix, "self_attn.k_proj.weight") == 0) { shape[0] = r->kv_heads * r->head_dim; shape[1] = r->hidden; *ndim = 2; } else if (strcmp(suffix, "self_attn.v_proj.weight") == 0) { @@ -263,6 +271,27 @@ static const recipe_t LLAMA3 = { .fixture_base = "tiny_llama3", }; +/* Pure-dense qwen3_5: full-attn deltas only (gate + partial rope). + No linear_attn tensors / hybrid layer_types - those are Stage E. */ +static const recipe_t QWEN3_5 = { + .model_type = "qwen3_5", + .vocab = 256, .hidden = 64, .heads = 4, .kv_heads = 2, + .head_dim = 16, .inter = 128, .layers = 2, + .group_size = 32, .quant_bits = 4, + .max_position_embeddings = 512, + .rms_norm_eps = 1e-6, .rope_theta = 10000000.0, + .sliding_window = 0, .sliding_window_pattern = 0, + .query_pre_attn_scalar = 0, .rope_local_base_freq = 0, + .has_lm_head = true, .tie_word_embeddings = false, + .attn_output_gate = true, + .partial_rotary_factor = 0.25, + .full_attention_interval = 1, + .top = QWEN3_TOP, .n_top = LEN(QWEN3_TOP), + .layer = QWEN3_LAYER, .n_layer = LEN(QWEN3_LAYER), + .emit_dense = true, .emit_sharded = false, .emit_tied = true, + .fixture_base = "tiny_qwen3_5", +}; + /* --- gemma4 recipe (custom build - not table driven) --------------------- */ /* Plan shape: vocab=256, hidden=64, inter=128, layers=4, @@ -492,8 +521,8 @@ static void write_gemma4_config(const char *dir) { yyjson_mut_doc_free(doc); } -static const recipe_t *RECIPES[] = { &QWEN3, &GEMMA3, &LLAMA3 }; -static const int N_RECIPES = 3; +static const recipe_t *RECIPES[] = { &QWEN3, &GEMMA3, &LLAMA3, &QWEN3_5 }; +static const int N_RECIPES = 4; /* --- Build tensors (table-driven) ---------------------------------------- */ @@ -557,7 +586,10 @@ static void write_config_json(const recipe_t *r, const char *dir, yyjson_mut_obj_add_int(doc, root, "max_position_embeddings", r->max_position_embeddings); yyjson_mut_obj_add_real(doc, root, "rms_norm_eps", r->rms_norm_eps); - yyjson_mut_obj_add_real(doc, root, "rope_theta", r->rope_theta); + /* When rope_parameters carries rope_theta (partial_rotary path), skip the + flat top-level duplicate to match real qwen3_5 configs. */ + if (!(r->partial_rotary_factor > 0.0)) + yyjson_mut_obj_add_real(doc, root, "rope_theta", r->rope_theta); yyjson_mut_obj_add_bool(doc, root, "tie_word_embeddings", tied); if (r->sliding_window > 0) { @@ -585,6 +617,21 @@ static void write_config_json(const recipe_t *r, const char *dir, yyjson_mut_obj_add_val(doc, root, "rope_scaling", rs); } + if (r->attn_output_gate) + yyjson_mut_obj_add_bool(doc, root, "attn_output_gate", true); + + if (r->full_attention_interval > 0) + yyjson_mut_obj_add_int(doc, root, "full_attention_interval", + r->full_attention_interval); + + if (r->partial_rotary_factor > 0.0) { + yyjson_mut_val *rp = yyjson_mut_obj(doc); + yyjson_mut_obj_add_real(doc, rp, "rope_theta", r->rope_theta); + yyjson_mut_obj_add_real(doc, rp, "partial_rotary_factor", + r->partial_rotary_factor); + yyjson_mut_obj_add_val(doc, root, "rope_parameters", rp); + } + if (quantized) { yyjson_mut_val *qcfg = yyjson_mut_obj(doc); yyjson_mut_obj_add_int(doc, qcfg, "bits", r->quant_bits);