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

[tlkahn]

Summary

Implements Stage D6 (#83): pure-dense MODEL_QWEN3_5 on the dense forward path.

• Full-attention deltas vs qwen3: attn_output_gate (doubled q_proj -> split queries/gate -> o_proj(attn * sigmoid(gate))) and partial_rotary_factor rope dims
• Support gate allows pure-dense qwen3_5 only; rejects MoE family, experts, linear heads, and hybrid layer_types
• weights_expected_names shares the qwen3 descriptor tables (gate is shape, not a new name)
• Tiny fixtures via gen_tiny_ckpt + GPU smoke (test_forward_qwen3_5_gpu)
• Parity table records mlx-community/Qwen3.5-0.8B-4bit; byte parity deferred to Stage E

Decision R1

1. Family discrimination stays MoE-field-based (num_experts > 0 => MODEL_QWEN3_5_MOE). Linear-bearing non-MoE configs are not reclassified as MoE.
2. Stage D gate rejects hybrid features on MODEL_QWEN3_5 (linear_num_*, hybrid layer_types, has_hybrid_layers).
3. D6 implements full-attn deltas only for pure-dense configs.
4. Real Qwen3.5 checkpoints are hybrid; load fails closed at the hybrid gate until Stage E. Byte-parity is an E-stage exit criterion.

Audit vs mlx-community/Qwen3.5-0.8B-4bit (config.json)

Field	qwen3 (typical)	qwen3_5 0.8B text_config	D6 handling
model_type	qwen3	root qwen3_5, text qwen3_5_text	already mapped (A1)
weight_prefix	model	language_model.model when nested	already set (A1)
head_dim	often hidden/heads	explicit 256	parsed; fixture sets explicitly
partial_rotary_factor	1.0 implicit	0.25 via rope_parameters	implemented
attn_output_gate	false	true	implemented
has_qk_norm	true	true	already on path
num_experts	0	0 (dense)	stays MODEL_QWEN3_5
linear_num_* / layer_types	absent	present (hybrid)	gate reject (Stage E)
rope_parameters.mrope_*	absent	present	text-only ignores; vision deferred

Deferred (Stage E / #56)

• GatedDeltaNet / linear-attention layers
• MoE router / experts
• Hybrid cache, vision / M-RoPE image path
• Token-level temp-0 parity on real Qwen3.5 checkpoints

Test plan

• make tests/test_emodel_gate && ./tests/test_emodel_gate
• make tests/test_model_config && ./tests/test_model_config
• make tests/test_weights && ./tests/test_weights
• make tests/test_forward_qwen3_5_gpu && ./tests/test_forward_qwen3_5_gpu
• llama / gemma4 GPU regression still green
• ASan/UBSan clean on qwen3_5 GPU smoke
• Parity wrapper skip-if-absent for qwen3_5
• make test / make test-gpu in CI (local: known pre-existing env flakes in registry discover / deepseek GGUF message / e2e-oracle-500)

Closes #83

[lgsunnyvale]

Code Review: PR #99 -- Stage D6 qwen3_5 dense full-attn deltas

I read the full diff (16 files, +728/-19) and examined the relevant source code in depth. Overall this is a well-structured, well-documented change that cleanly extends the existing qwen3 dense path with the qwen3_5 deltas. The gen_tiny_ckpt test infrastructure work is especially good.

Below are my findings, ordered by severity.

---

Issue 1 (Medium-Hard): Double mlx_array_eval + redundant mlx_array_new() in is_finite_f32 test helper

File: tests/test_forward_qwen3_5_gpu.c, lines 24-38

static int is_finite_f32(mlx_array a, mlx_stream s) {
    mlx_array f32 = mlx_array_new();          // (1) allocates empty array
    if (!MLXB_CHECK(mlx_astype(&f32, a, ...))) { ... } // (2) overwrites f32 ptr
    if (!MLXB_CHECK(mlx_array_eval(f32))) { ... }       // (3) eval
    size_t n = mlx_array_size(f32);
    const float *d = mlx_array_data_float32(f32);       // (4) second implicit eval?
    ...
}

Two issues:

1a. The mlx_array_new() at (1) is immediately overwritten by mlx_astype at (2). mlx_astype allocates internally and writes through the output pointer -- it does not free the previously-pointed-to (empty) array. This leaks one mlx_array handle per call.

Same pattern in fwd_attention: gate = mlx_array_new() at line 520 is followed by mlx_reshape(&gate, gate_4d, ...) on line 567. If the reshape succeeds, the original empty gate is overwritten without being freed. If it fails (goto cleanup), mlx_array_free(gate) at cleanup frees the original empty array, so the leak is only on the success path.

This pre-existing pattern exists throughout the codebase (e.g. mlx_reshape(&q_reshaped, q, ...) where q_reshaped was also mlx_array_new()'d), so it's not a new introduction. But since this PR adds a new instance via gate, it's worth noting. If the MLX-C API indeed frees the old array internally (I couldn't confirm from the headers), then disregard. If not, gate leaks on every attn_output_gate=true attention call.

1b. mlx_array_data_float32(f32) on line 33 may trigger its own lazy evaluation after mlx_array_eval on line 30. If mlx_array_eval is idempotent (as in most tensor frameworks), this is harmless but redundant.

Recommendation: Remove mlx_array_new() before mlx_astype calls (test helper and gate) to match the established convention of declaring and immediately assigning. Or confirm that mlx_astype/mlx_reshape frees the input array internally.

---

Issue 2 (Medium): rope_dims parity -- K/V cache re-use with partial_rotary_factor

File: src/engine/forward.c, lines 505-506 + 621-697

fwd_rope_dims_partial computes rotated dims from cfg->partial_rotary_factor and layer hd_l. This value is passed to fwd_rope_apply for both Q and K RoPE.

For KV-shared layers (line 621), only Q is rotated -- K/V come pre-rotated from the source layer's cache. If the source layer has a different partial_rotary_factor or head_dim than the consumer layer, the rope domains would not align.

This is not a bug now because:

1. qwen3_5 pure dense has no KV-shared layers (no num_kv_shared_layers set, no hybrid layers)
2. All layers are global/full-attention with identical head_dim

But it will become a correctness trap in Stage E when hybrid layers with different head_dim or PRF enter the picture. A KV-shared layer consuming K/V from a source layer with different rope_dims would silently produce wrong attention.

Recommendation: Add an assertion or early-return guard in Stage E's hybrid path that verifies rope_dims agreement between a KV-shared consumer and its source layer. Not actionable now, but worth documenting in the code.

---

Issue 3 (Medium-Hard): Missing forward test for attn_output_gate=false config

File: tests/test_forward_qwen3_5_gpu.c

The gate test (test_qwen3_5_pure_dense_passes) verifies that attn_output_gate=false with partial_rotary_factor=1.0 passes the emodel gate. But the forward test suite only exercises attn_output_gate=true fixtures (tiny_qwen3_5 and tiny_qwen3_5_tied). There is no GPU forward pass over a gate-disabled qwen3_5 config.

The gen_tiny_ckpt recipe also only defines a single QWEN3_5 recipe with attn_output_gate=true.

Recommendation: Either:

• Add a third fixture variant (gate-off dense) with a forward smoke test, or
• Document explicitly in the test file that attn_output_gate=false is covered by the existing qwen3 dense forward tests (which already test the same topology without the gate).

The second option is acceptable since qwen3_5 with gate-off is identical to qwen3 dense on the forward path.

---

Issue 4 (Low): partial_rotary_factor even-ness constraint for RoPE

File: src/engine/forward.c, lines 482-489

static int fwd_rope_dims_partial(const model_config_t *cfg, int head_dim) {
    float prf = cfg->partial_rotary_factor;
    ...
    int dims = (int)((float)head_dim * prf);
    if (dims <= 0 || dims > head_dim)
        return head_dim;
    return dims;
}

mlx_fast_rope operates on pairs of dimensions. If dims is odd (e.g., head_dim=17 * PRF=0.5 => 8.5 => 8 through truncation), the behavior depends on the underlying Metal kernel. For the real Qwen3.5-0.8B config (head_dim=256, PRF=0.25 => 64), this is always even. But there is no guard.

Recommendation: Add dims &= ~1; (floor to even) after the integer conversion to be safe against odd head_dims or odd PRF values. This is a zero-cost operation and prevents a silent kernel-level issue.

---

Issue 5 (Low): gen_tiny_ckpt emits flat rope_theta redundantly alongside rope_parameters.rope_theta

File: tools/gen_tiny_ckpt.c, lines 587 and 619-620

yyjson_mut_obj_add_real(doc, root, "rope_theta", r->rope_theta);   // line 587 (always)
...
yyjson_mut_obj_add_real(doc, rp, "rope_theta", r->rope_theta);     // line 620 (under rope_parameters)

The generated tiny_qwen3_5/config.json ends up with both a top-level rope_theta and rope_parameters.rope_theta. The parser reads the flat key first, then overwrites with the rope_parameters value. Since both carry the same value, it's harmless but redundant.

The hand-written model_config_qwen3_5_dense_full/config.json correctly omits the flat rope_theta key, which is the cleaner approach.

Recommendation: Guard the flat rope_theta emission behind a condition that skips it when rope_parameters will be emitted. Or simply accept the redundancy for test fixture generation.

---

Issue 6 (Low): Unused weight_prefix field in test_expected_names_qwen3_5

File: tests/test_weights.c, lines 632-633

cfg.family = MODEL_QWEN3_5;
cfg.weight_prefix = "model";

This duplicates the default. model_config_defaults() already sets weight_prefix = "model". The equivalent q3 struct does not set it explicitly either. Harmless.

---

Issue 7 (Nit): Error message drift in test_reject_all_other_families

File: tests/test_emodel_gate.c, lines 128-143

MODEL_QWEN3_5 was removed from the "other families" rejection list (which is correct: it's now supported). But this means the test no longer validates that other implementations (e.g. MODEL_DEEPSEEK_V3) are still properly rejected. This is very minor since each family generally has its own gate test.

---

Summary

#	Severity	Area	Description
1	Medium-Hard	forward.c + test	gate and is_finite_f32 leak empty mlx_array handle on mlx_reshape overwrite (pre-existing pattern, one new instance)
2	Medium	forward.c	KV-shared + partial_rotary_factor compatibility trap for Stage E (no bug now, document it)
3	Medium-Hard	Test coverage	No forward pass for attn_output_gate=false path (acceptable if documented as qwen3 coverage)
4	Low	forward.c	rope_dims should floor to even for safety
5	Low	gen_tiny_ckpt	Redundant flat rope_theta in generated configs
6	Low	test_weights.c	Redundant weight_prefix = "model"
7	Nit	test_emodel_gate.c	Test scope narrowing after family promotion

The PR is solid. Issues 1-2 are the ones I'd prioritize discussing. The structural decisions (no reject_dense_common, inlined gate, name-table reuse) are all correct and well-commented.

[lgsunnyvale]

Code review: PR #99 (Stage D6 - qwen3_5 dense full-attn deltas)

Thanks for the thorough PR description with the audit table and Decision R1 -- it made the review significantly easier to navigate. The implementation is clean and well-scoped to the pure-dense path. Below are findings organized by severity.

---

Blocking / needs verification

1. mlx_vector_array_get + mlx_vector_array_free lifetime (forward.c:561-571)

mlx_vector_array split_vec = mlx_vector_array_new();
// ... mlx_split fills split_vec with 2 new arrays ...
mlx_vector_array_get(&queries, split_vec, 0);
mlx_vector_array_get(&gate_4d, split_vec, 1);
mlx_vector_array_free(split_vec);   // <-- frees the container + its arrays

Then queries and gate_4d are used afterward (one is reshaped into gate for the downstream sigmoid multiply).

This pattern is new to the codebase -- there are zero prior callers of mlx_split or mlx_vector_array_get in src/. The safety depends on whether mlx_vector_array_free decrements refcounts of the contained arrays (safe, since get incremented refcounts) or eagerly destroys them (use-after-free).

Based on mlx's C++ semantics, arrays are shared_ptr-like (refcounted), and mlx_vector_array is a std::vector<array>. So free destroys the vector (decrementing each array's refcount), and get copies the handle (incrementing). This means the pattern is safe, but since this is the first use of these APIs in the forward path, I'd like explicit confirmation. If you've tested this under ASan and it's clean, that's sufficient -- please confirm.

2. rope_dims applies partial RoPE to all families, gated only by the emodel check (forward.c:482-490)

fwd_rope_dims_partial computes int(prf * head_dim) unconditionally. The emodel gate rejects partial_rotary_factor != 1.0f for qwen3/llama/gemma4, and the qwen3_5 gate allows (0, 1]. This means only qwen3_5 currently hits the non-1.0 path, but the coupling is silent. If a future family addition sets partial_rotary_factor and forgets the gate guard, partial RoPE silently activates. Consider adding an explicit family check here, or at minimum a comment on the gate side noting that fwd_rope_dims_partial relies on the gate to restrict which families reach it.

---

Medium

3. Redundant rope_theta emission in gen_tiny_ckpt.c (tools/gen_tiny_ckpt.c:612-626)

When partial_rotary_factor > 0, the tool emits rope_parameters as a nested object containing rope_theta. But the standard config-writing path always emits top-level rope_theta. This produces fixture configs with the value written twice:

"rope_theta": 10000000.0,
"rope_parameters": {
    "rope_theta": 10000000.0,
    "partial_rotary_factor": 0.25
}

Not a bug (the model.c parser reads top-level first, then rope_parameters overrides), but it's misleading for someone reading the fixture to understand the canonical config shape. Consider skipping the top-level rope_theta when rope_parameters is emitted.

4. No test-gpu Makefile target

The PR test plan lists "make test / make test-gpu in CI" as unchecked, noting pre-existing flakes. But there is no test-gpu target in the Makefile at all. The GPU tests (test_*_gpu.c) are excluded from make test (line 66) and must be built and run individually. This is pre-existing, not introduced by this PR, but the test plan's reference to make test-gpu implies a target that doesn't exist. Consider adding one (out of scope for this PR, but worth a follow-up issue).

5. is_finite_f32 is a copy of a pattern that exists in other GPU tests

test_forward_qwen3_5_gpu.c defines a static is_finite_f32 helper. test_forward_gemma3_gpu.c and others likely have similar or identical helpers. At some point this should be extracted to a shared test utility. Not blocking.

---

Minor

6. Gate redundancy: has_hybrid_layers vs has_explicit_layer_types loop (emodel.c:106-114)

The gate checks:

1. linear_num_key_heads > 0 || linear_num_value_heads > 0 (reject on field presence)
2. has_hybrid_layers (reject on parser-detected hybrid)
3. Loop over layer_is_global when has_explicit_layer_types (reject on any non-full-attention layer)

Check (3) is the most precise. Checks (1) and (2) are useful early-exits for malformed or unparseable configs, but their overlap with (3) for well-formed configs adds maintenance weight. Consider a single unified hybrid gate function when Stage E adds linear-attention support.

7. Error message drift in default case (emodel.c:137-138)

"unsupported model family (only qwen3/llama/gemma4/qwen3_5 dense supported)"

This string now requires updating every time a new family is added to the switch. Since it only fires for families that hit default, it's technically accurate, but the explicit list is misleading when families like deepseek_v3/bert/lfm2 have their own case branches. Consider a simpler message like "unsupported model family" and let the specific family be logged separately.

8. 5e-2f tolerance in incremental test (test_forward_qwen3_5_gpu.c:182)

The maxdiff tolerance for incremental_equals_full is 5e-2f. For a tiny bf16 model this may be genuinely needed due to numerical differences between prefill and decode attention paths. A brief comment explaining why the tolerance is large (e.g., "bf16 accumulation drift between prefill vs decode in tiny models") would help future readers.

9. Duplicated rope_theta in fixture configs (tests/fixtures/tiny_qwen3_5/config.json)

Same as item (3), but applied to the checked-in fixtures. Not a correctness issue but adds noise.

---

Nit

10. assert(0) in run_suite (test_forward_qwen3_5_gpu.c:198)

fprintf(stderr, "failed to load %s: %s\n", path, err);
assert(0);

The error message is lost in -DNDEBUG builds. The rest of the test suite uses bare assert() on expressions that self-document; mixing fprintf + assert(0) creates a false expectation that the message will always print. Consider using fprintf(stderr, ...) + abort() or just assert(engine_model_load(...) == 0).

---

Things I verified and found correct

• q_proj output shape doubling in gen_tiny_ckpt.c (q_out *= 2 when attn_output_gate) is consistent with the forward path split
• Weight descriptor reuse of qwen3's tables is correct -- the gate is a shape convention, not a separate tensor
• The model_config_load -> engine_model_load -> engine_model_check_supported pipeline correctly rejects hybrid dense at load time (not silently at forward time)
• qwen3_5 MoE family reclassification (num_experts > 0 -> MODEL_QWEN3_5_MOE) is preserved and tested
• KV-shared layer path reuses rope_dims correctly (Q roping uses the shared source's head_dim via hd_l = model_layer_head_dim(cfg, layer))
• Cleanup section frees gate (was mlx_array_new() at top, always safe even when not used)
• rope_dims is computed once from the layer's head_dim -- correct for homogeneous-head qwen3_5; would remain correct for future heterogeneous models since hd_l is per-layer
• Test fixtures cover untied/tied/quantized permutations, hybrid rejection, and boundary config shapes

---

Summary

The implementation is solid. The forward-path gate split logic is new API territory for the codebase and needs an explicit ASan confirmation on items (1) and (2). The remaining items are maintainability nits that don't block merge.

[tlkahn]

Reply to review (comment 5018437321)

Addressed findings 1-8 on issue-83 in ceeb350. Finding 9 was already self-retracted.

Per finding

1. High - parity wrapper never-fail-CI break: confirmed. scripts/parity_family.sh now hard-skips qwen3_5 when resolved via the canonical id / MLXD_PARITY_CKPT_ROOT path, even if the hybrid snapshot dir is present (skipped: qwen3_5 hybrid parity deferred to Stage E). MLXD_PARITY_CKPT_QWEN3_5 remains a deliberate escape hatch for a future dense checkpoint. Covered by wrapper-skip-present-qwen3_5-hybrid in tests/test_parity_script.sh. Stage E deletes the branch.

2. Medium - full_attention_interval unenforced: confirmed, with refined semantics. mlx-lm uses is_linear = (layer_idx+1) % interval != 0, so interval == 1 is pure dense (every layer full-attn). Gate now rejects > 1 and allows 0 (unset) and 1.

3. Medium - attn_output_gate default clobber: confirmed. Parse side now defaults true only when the key is absent (apply_family_defaults). Explicit false survives parse (new fixture model_config_qwen3_5_gate_off). Gate side requires gate-on: REJECT(!cfg->attn_output_gate, "qwen3_5 requires attn_output_gate"). Rationale: mlx-lm gates unconditionally; no oracle or real checkpoint for gate-off.

4. Medium - odd/zero partial rope dims: confirmed. After the prf-range check the gate computes dims = (int)(head_dim * prf) and rejects dims < 2 || (dims & 1). fwd_rope_dims_partial left mirroring mlx-lm; comment notes the gate invariant.

5. Medium - D6 deltas only self-consistency tested: confirmed. Added tests/test_qwen3_5_goldens_gpu.c loading both tiny fixtures via engine_model_load + model_forward on ids [1,2,3,4,5], comparing last-token first-8 logits + argmax against mlx-lm 0.31.3 (mlx_lm.models.qwen3_5) goldens. Dense bf16: exact match. Tied 4-bit: maxdiff ~2e-3 (atol 5e-2). Proved the oracle pins the deltas by temporarily skipping the sigmoid gate multiply (test failed) then restoring (test passed).

6. Low - qwen3_5 gate duplicates reject_dense_common: confirmed; refactored now. reject_dense_common(cfg, err, errlen, flags) takes REJECT_ALLOW_PARTIAL_ROPE and REJECT_ALLOW_ATTN_GATE. qwen3/llama pass 0; qwen3_5 passes both flags then keeps family-specific checks only.

7. Low - flat rope_theta beside rope_parameters: confirmed. Generator skips flat rope_theta when emitting rope_parameters, and emits full_attention_interval: 1 for the QWEN3_5 recipe. Regenerated fixtures: only qwen3_5 config.json files changed; safetensors hashes unchanged (seed 77).

8. Low - gate-off untested end-to-end: closed by Migrate chat module (jinja template rendering) #3. Explicit false is parsed then rejected as unsupported; no gate-off forward fixture.

9. Nit - mlx_array ownership: acknowledged as self-retracted; no code change.

Verification

• test_emodel_gate, test_model_config, test_forward_qwen3_5_gpu, test_qwen3_5_goldens_gpu: pass
• test-parity-script wrapper cases including the new hybrid present-skip: pass
• Unrelated pre-existing flakes left alone (test_registry_discover home HF hub, test_weights_gpu deepseek GGUF message, e2e-oracle-500)

[tlkahn]

Reply to second review (comment 5018410477)

Verified every finding against HEAD after ceeb350 (first-review fixes) and addressed the remaining valid ones in 2abb872.

Per finding

#	Verdict	Response
1	Confirmed safe	mlx-c mlx_vector_array_get copies the array handle (shared_ptr refcount +1); mlx_vector_array_free destroys only the container and drops its own refs. Empirical: full ASan+UBSan rebuild, test_forward_qwen3_5_gpu all 9 suites pass with zero sanitizer reports.
2	Accepted (narrowed to gemma4)	Real hole was gemma4 only: its gate case never checked plain partial_rotary_factor, while model.c parses a flat rope_parameters block for any family. Fixed gate-side: gemma4 now REJECTs partial_rotary_factor != 1.0f with message noting gemma4 uses rope_parameters.full_attention / partial_rotary_factor_global. New test test_gemma4_reject_partial_rotary_plain (TDD). Comment on fwd_rope_dims_partial updated to state every family's gate constrains the field so only qwen3_5 reaches the non-1.0 path.
3	Already fixed	ceeb350: gen_tiny_ckpt.c now skips top-level rope_theta when rope_parameters carries it. Crossed with this review in flight.
4	Refuted	test-gpu exists at Makefile:95 and predates this PR; it builds and runs all test_*_gpu binaries.
5	Fixed	Canonical is_finite_f32 extracted to tests/gpu_test_util.h; the 4 local copies in test_forward_{gemma3,gemma4,llama,qwen3_5}_gpu.c now include it.
6	Deferred	Agree with the Stage E suggestion. ceeb350 already moved has_hybrid_layers into shared reject_dense_common; full unification lands when Stage E adds linear attention.
7	Refuted	Premise is wrong: deepseek_v3/bert/lfm2 have no case branches in engine_model_check_supported (only qwen3/llama/gemma4/qwen3_5 + default). The message is accurate today and pinned by 3 assertions in test_emodel_gate.c (lines 34, 138, 218), so drift is test-caught.
8	Fixed	Added a gemma3-style comment above the maxdiff assert: argmax + max-abs-diff < 5e-2, bf16 prefill-vs-decode accumulation drift, same tolerance as sibling GPU tests. (5e-2f is the uniform tolerance across gemma3/gemma4/llama/qwen3_5, not qwen3_5-specific.)
9	Already fixed	ceeb350 removed flat rope_theta from both tiny_qwen3_5 fixtures (now 1 occurrence each).
10	Fixed	Replaced assert(0) with abort() after the fprintf at the 3 sites (test_forward_qwen3_5_gpu.c, test_forward_llama_gpu.c, test_qwen3_5_goldens_gpu.c) so the load-failure message survives NDEBUG.

Verification

• make test: CPU suite green on the changed surface (test_emodel_gate and friends). Two pre-existing unrelated aborts remain on this tree (test_registry_discover, test_weights_gpu) and reproduce on ceeb350 without these changes.
• make test-gpu: all suites touched by this PR pass (test_forward_{gemma3,gemma4,llama,qwen3_5}_gpu, test_qwen3_5_goldens_gpu).