feat: aarch64-unknown-linux-gnu (NVIDIA GB10 / DGX Spark) CUDA backend support - #71

…n3.6 classes) (#72)

## Summary

Collapses the **O(T) per-step gated-delta (GDN) recurrence** on the CUDA
(non-Metal) prefill path into **O(T/BT) chunk-serial batched matmuls**
(cuBLAS / tensor cores) — a device-agnostic, pure-`MxArray` port of the
in-tree Metal chunked kernel
(`crates/mlx-sys/src/metal/gated_delta_chunked.metal.inc`).

- **Zero build changes** (no nvcc/NVRTC) — matmuls route through cuBLAS.
- **Default-on** for the CUDA ops path; `MLX_GDN_KERNEL=perstep` reverts
for same-binary A/B.
- **No Metal impact** — `use_kernel=true` never reaches this path; the
Mac/Metal production path is byte-identical.

This attacks the *"GDN per-step recurrence is the prefill floor"*
bottleneck that [PR #71](#71
benchmark flagged as the **#1 dense lever**.

## Measured (GB10 / DGX Spark, Qwen3.6, warm, prefill TTFT vs per-step)

| model        | 1577-tok speedup | parity |
|--------------|------------------|--------|
| dense-Q4     | **1.62×**        | byte-identical / late-drift |
| dense-NVFP4  | 1.33×            | identical |
| MoE-Q4       | **1.75×**        | coherent |
| MoE-NVFP4    | 1.40×            | coherent |

Win **grows with prompt length** (dense-Q4: 1.06×@200 → ~1.58×@1577+;
the chunked inverse is fixed-cost while per-step is O(T)). chunked
prefill tok/s climbs to ~242 vs per-step's flat ~150. MoE wins more (GDN
is a larger prefill fraction); NVFP4 less (dequant dominates its
prefill).

## Numerical-stability fixes (each with a Mac regression test)

1. **Triangular inverse overflow.** `M = (I+A)⁻¹` by repeated squaring
overflows f32 at `BT=64` (`N³² ≈ 4e57` before nilpotency zeroes it at
`N⁶⁴`), producing garbage. Replaced with **row-iterative forward
substitution** (FLA / vLLM `solve_tril`): `M[i,:] = eᵢ − A[i,:]·M` — no
powers of A, stable for any `‖A‖`, serial depth independent of T. Test:
`chunked_ops_stable_with_correlated_unit_norm_keys`.
2. **Gate underflow (MoE-only garbage).** `g_log = g.log()` round-trips
through the exp-space gate; strong decay (which MoE has, dense doesn't)
underflows `g` to 0 → `log(0) = -inf` → chunked `gcum_i − gcum_j = inf −
inf = NaN`. Now compute `g_log = -exp(a_log)·softplus(a + dt_bias)`
**directly in log-space** (matches the native `g_log` the fused Metal
gating returns). Test: `compute_g_log_finite_under_strong_decay`.

## Validation

- 6/6 `gated_delta` Rust unit tests, `cargo clippy -p mlx-core
--all-targets`, `cargo fmt` — green.
- Correctness validated on the DGX across **all four Qwen3.6 classes**
(dense/MoE × Q4/NVFP4) — per-step vs chunked greedy A/B, coherent output
everywhere (garbage only before the two fixes above).
- Algorithm derivation in `docs/gdn-chunked-ops-spec.md`.

## Follow-ups (not blocking)

- FLA 16-block row-iterative + block-merge inverse to cut short-prompt
inverse depth 63→~16 (long-prompt asymptote is carry-bound, won't move).
- Runtime non-finite guard → fall back to per-step (overflow is
currently silent; the `Err` fallback doesn't catch `Inf`/`NaN`).
- Possibly raise `CHUNK_THRESHOLD`→256 (the 200-tok win is only 1.06×).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> **Medium Risk**
> Touches core Qwen3.5 inference recurrence and output numerics on CUDA,
though Metal is gated off, per-step fallback exists, and behavior is
covered by parity/stability tests.
> 
> **Overview**
> Adds a **default-on CUDA prefill fast path** for Qwen gated-delta
(GDN): long, unmasked sequences on the non-Metal ops branch now run
**`gated_delta_chunked_ops`**, a pure `MxArray` chunk-parallel port
(BT=64) that replaces the O(T) per-token loop with O(T/BT) chunk carries
and batched matmuls. Metal/`use_kernel=true` routing is unchanged;
decode and masked calls still use per-step ops.
**`MLX_GDN_KERNEL=perstep`** (and **`ForceChunkedOps`** / `chunked_ops`
aliases) support same-binary A/B.
> 
> Two **numerical fixes** ship with the chunked path:
**`compute_g_log`** computes the decay gate in log-space (avoids
`log(0)` → NaN on strong MoE decay), and
**`invert_i_plus_strict_lower`** builds `(I+A)⁻¹` via forward
substitution instead of f32 power squaring that overflows at BT=64.
Chunked ops errors fall back to per-step with a stderr warning.
> 
> Adds **`docs/gdn-chunked-ops-spec.md`** plus unit tests for env
parsing, chunked vs per-step parity across chunk boundaries,
correlated-key inverse stability, and strong-decay gating.
> 
> <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit
efe9c8f. Bugbot is set up for automated
code reviews on this repo. Configure
[here](https://www.cursor.com/dashboard/bugbot).</sup>
<!-- /CURSOR_SUMMARY -->

---------

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Sync the chat-engine refactor branch with origin/main's two CUDA commits.
All CUDA additions are Metal-gated, so the macOS path is unchanged.

Conflict resolution (refactor base kept; CUDA additions ported):
- The refactor relocated qwen3_5::persistence_common -> engine::persistence,
  so every CUDA-added compiled_forward_backend_available() guard (20 call
  sites across gemma4/lfm2/qwen3/qwen3_5/moe, auto-merged with the old path)
  is repointed to crate::engine::persistence.
- lfm2 keeps its compiled path: its Metal gate is folded onto the P5
  per-model-id deadlock-fix registration; persistence imports updated.
- qwen3_5 dense + qwen3_5_moe register_weights_with_cpp C++ compiled
  forwards were deleted by P6/P7; CUDA's edits to those blocks are dropped
  (0 callers on HEAD).
- gated_delta.rs: our GDN tape-replay test and #72's three chunked-ops
  tests coexist in one module (all 7 module tests green).

Gates (M5 Max): cargo build/clippy -D warnings/fmt --check clean;
mlx-core lib suite 1775 passed, 2 pre-existing f32-precision failures
(banded_attention / attention_vjp finite-diff) on code paths byte-identical
to origin/main — not introduced by this merge.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

…_config.json (#38) (#70)

## Chat-engine refactor — model-neutral engine → paging inversion →
compiled coexistence → pure-Rust eager/MTP → declarative convert

> **Scope.** This branch carries the **entire** chat-engine refactor
stacked on `main` — Phase 1 (engine extraction, S1–S12) plus follow-ups
**P1–P8**, the **MTP and vision genericizations**, and the `#38`
generation-config feature: **73 commits, 143 files, +39.5k / −54.0k (net
≈ −14.5k LOC)**. The phases were developed on one long-lived branch and
build on each other, so they cannot be split into separate PRs against
`main`. Opened for review.
>
> **Synced with `main`.** Merged `origin/main` (NVIDIA DGX/CUDA backend
[#71] + qwen3.5 GDN chunked-ops prefill [#72]) into the branch — **PR is
now conflict-free / mergeable**. The CUDA additions are Metal-gated, so
the macOS path is byte-for-byte unchanged; the only merge work was
repointing CUDA's `compiled_forward_backend_available()` Metal guards
from the pre-refactor `qwen3_5::persistence_common` to the relocated
`engine::persistence`, and keeping both our GDN tape-replay test and
#72's new chunked-ops tests.

### Why
Adding a new model family cost ~10–17k LOC because the production engine
(decode loop, session mgmt, paged-KV, compiled-graph path, MTP) was
built inside the first family (`qwen3_5/`) and copy-forked into each new
one. This refactor extracts a model-neutral engine and pushes every
production feature behind an engine-owned seam, so a new family becomes
a thin backend impl — hence the large net deletion.

### Phases on this branch

**Phase 1 — extract the model-neutral chat engine (S1–S12).** New
`crates/mlx-core/src/engine/`: `ChatConfig/ChatResult/ChatStreamChunk`
types, `ChatBackend`/`DecodeStep` traits, generic `run_decode_loop`,
generic session cores + `ChatCmd` + napi glue, compiled locks,
persistence helpers. All five families (qwen3, qwen3_5 dense,
qwen3_5_moe, gemma4, lfm2) migrated onto it; legacy per-family chat
machinery deleted (grep-gated).

**P1–P3 — engine seams.** `ThinkingPolicy` + ChatML delta
single-source-of-truth; paged-turn lifecycle dedup via a `PagedTurnPlan`
adapter.

**P4 — paging inversion (`engine::run_paged_turn<B: PagedBackend>`).**
The default paged auto-regressive whole turn now flows through one
generic engine path + `run_decode_loop` for all five families, behind a
paged-forced byte-equivalence gate.

**P5 — per-model-id compiled slot map.** Replaces the single
process-global `g_active_model_id` gate over the C++ compiled-decode
state with per-model-id slot maps, so **two compiled models coexist in
one process** instead of the 2nd load silently demoting the 1st to
eager. Includes the MTP draft/commit-graph residency owner-guard and a
deadlock fix on the per-`Inner` Drop.

**P6/P7 — delete the C++ compiled forwards, go pure-Rust.** Removed the
qwen3_5 dense + MoE C++ compiled forward paths (~9.4k LOC + 104 FFI
externs) in favor of a pure-Rust eager forward and a **pure-Rust eager
MTP** (drafter + batched verify + accept/rewind + a Rust GDN tape for
bit-exact acceptance). AR eager == former compiled (parity); eager MTP
keeps committed-history acceptance. lfm2/gemma4 compiled paths
preserved. A hard perf gate (eager-MTP must not regress vs compiled-MTP)
was enforced.

**P7 (collapse) — TS / NAPI / Cmd / training.** 6 streaming wrapper
classes → one `makeStreamingModel` factory + `MODEL_REGISTRY`; 35
hand-written `#[napi]` chat methods → one `chat_napi_surface!`
decl-macro; per-family training enums → `engine::cmd::TrainCmd` +
`handle_train_cmd<B>` (+ 12 fallible AdamW-apply `unwrap` → `?`).

**P8 — declarative convert (`ConversionRecipe` registry).** Collapsed
the 7 scattered per-family `convert.rs` dispatch sites into one
`recipe_for(model_type)` registry; **e2e `mlx convert` output is
byte-identical (sha256)** for qwen3_5 / lfm2 / privacy-filter.

**#38 — honor `generation_config.json` defaults engine-wide (vLLM
out-of-box).** When a request leaves a sampling field unspecified it
falls back to the checkpoint's `generation_config.json`, and that file's
`eos_token_id` list merges into the per-turn stop set. Precedence
`request > generation_config > sampler builtin`; eos is a **union**
(never an override). Default-ON. `index.d.cts` unchanged (defaults live
on each `Inner`).

**MTP genericization — `engine::run_mtp_turn<B: MtpBackend>` +
`run_mtp_cycle<S: MtpStepper>`.** Continuing P4's inversion for
speculative decoding: the MTP whole turn and its propose / verify /
accept-or-rewind cycle moved off the per-family `decode_loop_mtp!` macro
onto one engine-owned generic path. qwen3_5 dense (flat sync, flat
streaming, paged) and qwen3_5_moe (flat) all run it; the now-dead
`decode_loop_mtp!` macro + `MtpOps` closure-bundle were deleted.
Per-family T=0 byte-identity gates; the macro's unit tests were migrated
to the live engine path — Codex caught a real dense
`accept_with_residual` (T=1.0 dense-accept) coverage gap the migration
had dropped, fixed with exact-one-hot deterministic tests.

**Vision merge-contract genericization — `engine::vision::{VisionMerge,
run_vlm_prefill_layers}`.** The qwen3.5 dense + MoE image-prefill
already shared the vision encode + image-feature merge + M-RoPE step
(one free fn); its last forked piece — the prefill layer loop — is now
unified behind an engine-owned `VisionMerge` DTO + a `VlmPrefillLayer`
trait (associated cache type, so each family keeps its distinct layer
cache) + one generic `run_vlm_prefill_layers`. The sole divergence
(dense relies on the attention layer's internal causal SDPA; MoE passes
an explicit causal mask) is now a parameter, not a forked loop. gemma4
stays on its own core — it is structurally disjoint (masked_scatter + 4D
mask, no M-RoPE) and even rejects a text delta on an image session.
Env-gated T=0 image-chat goldens are byte-identical before / after for
both dense (Qwen3.5-4B) and MoE (Qwen3.6-35b-a3b).

### Docs & examples
- **Docs:** `ModelGenerationDefaults` rustdoc now documents the full
`generation_config.json` vs `config.json` override order (sampling
ladder + eos union + raw-`generate()`/Gemma4 divergences).
- **Examples cleanup:** removed ~30 internal
debug/probe/microbench/perf-A-B/smoke scripts (incl. the
`grpo/validation/**` real-API subtree); `examples/` now keeps only
user-facing examples — inference (`lm.ts`), VLM (`vlm-inference.ts`),
OCR (`paddle-ocr-pipeline.ts`), tool-use (`tool-use-example.ts`), and
the GRPO/SFT GitHub-tool training example. Reference-safe (`yarn
typecheck` + `lsp.test.ts` 75/75 green; dead `package.json` bench
scripts + doc references pruned).

### Verification
- `cargo test -p mlx-core --lib` green (only the 2–3 pre-existing
Metal-f32 failures, documented; on code paths byte-identical to
`origin/main`) · `cargo test -p mlx-paged-attn` green.
- Per-family real-checkpoint parity at each P4/P5 step (paged-vs-flat
byte-identical); P6/P7 AR eager==compiled + eager-MTP acceptance parity;
P8 byte-identical convert; #38 real-checkpoint smoke + 11 unit tests.
- `clippy --all -D warnings` + `cargo fmt --check` clean · `yarn
build:native` ok · `packages/core/index.d.cts` byte-stable · `yarn
typecheck` + TS lint clean.
- Post-sync re-gate on the merge commit: `build` + `clippy -D warnings`
+ `fmt --check` clean; `mlx-core` lib 1775 pass; the merged
`gated_delta` module runs both the GDN tape-replay test and #72's
chunked-ops tests (7/7).
- MTP + vision genericization: per-family T=0 byte-identity goldens
green before / after; `mlx-core` lib **1796 pass** / 3 pre-existing
Metal-f32; Codex adversarial-review **approve** (no material findings)
on both.
- Adversarial (Codex) review on each landed phase (all approve /
findings fixed).

### Post-review follow-ups & CI

- **Open review threads resolved (3/3), each verified against real
control flow first:** restored the native config constructor on the
streaming wrappers (`ConstructorParameters<C>` — `new
Gemma4Model(config)` / `new QianfanOCRModel(config)` type-check again);
reinserted the `.biases` sidecar on the vision-dequant fallback so a
failed dequant can't emit a corrupt affine quant group; the gemma4
paged-cap flag was confirmed **not** a regression (the clamp was
intentionally converged onto the shared engine contract — the KV pool is
sized to `max_position_embeddings` and `record_tokens` rolls back on
exhaustion, so over-budget turns abort cleanly rather than overrun).
- **Rust Lint fix:** moved the `VlmPrefillLayer` impl above the test
module (`clippy::items-after-test-module`, a `--all-targets`-only lint
that the local `--lib` check missed).
- **Removed T=0 parity scaffolding** (`scripts/t0-smoke.sh` + its
env-gated Rust capture harnesses, −1.6k LOC): dev-only, never ran in CI,
and the refactor it verified has landed. Kept the real CI coexistence
test.
- **Gated heavy model E2E, folded into `ci.yml` to reuse the single
Build:** real-checkpoint tests run as a label-gated `model-test` job
inside `.github/workflows/ci.yml` that **reuses the `build` job's
`mlx-node` + `mlx-metallib` artifacts — no second native build** (the
original standalone `model-e2e.yml` duplicated the entire build job
verbatim). A per-family matrix downloads the smallest checkpoint (qwen3
0.6B · qwen3.5 dense 0.8B · lfm2 1.2B · gemma4 E2B) and runs that
family's self-referential `--ignored` parity/session tests (paged≡flat,
reset/replay determinism, session warm-reuse, greedy MTP≡AR — no
host-specific golden). Runs only on `main` pushes, `workflow_dispatch`,
or a PR labeled `model-e2e` — keyed on the label that fired the event,
so adding an unrelated label to an already-labeled PR is a no-op. On a
`model-e2e` label event only `build` + `model-test` run (the regular
`test`/`lint` already ran on the PR's push). qwen3.5 MoE is excluded (no
checkpoint below 35B-A3B).

[#71]: #71
[#72]: #72

🤖 Generated with [Claude Code](https://claude.com/claude-code)

<!-- CURSOR_SUMMARY -->
---

> [!NOTE]
> **High Risk**
> The declaration file removes and reshapes large parts of the public
`@mlx-node/core` API (chat, Qwen3 paged paths, training hooks), and
publish now gates on optional E2E—both affect consumers and release
flow.
> 
> **Overview**
> Adds a **`model-test` CI job** inside `ci.yml` that reuses the build
artifacts, downloads small real checkpoints per family (qwen3, qwen3.5
dense, lfm2), and runs `#[ignore]` parity/session tests when pushing to
`main`, on `workflow_dispatch`, or when a PR carries the **`model-e2e`**
label—with guards so unrelated label events do not re-run the full
test/lint suite. **Publish** now depends on `model-test` succeeding (or
being skipped) so a failed E2E leg blocks release.
> 
> Regenerates **`packages/core/index.d.cts`** to match the refactored
native surface: unified **`ChatConfig` / `ChatResult` /
`ChatStreamChunk`**, session + streaming entry points on model classes,
**`ChatStreamHandle`**, block-paged cache helpers, expanded GRPO/SFT
engine factories, and new exports (Gemma4, LFM2, Harrier, PrivacyFilter,
QianfanOCR, ResponseStore, profiling/memory helpers). Several legacy
Qwen3 paged-attention and in-process `chat()`/`forward` APIs are
**removed** from the typings in favor of model-thread session APIs.
> 
> Small follow-ups: Gemma4 loader test asserts **transitive subclass**
via `makeStreamingModel`; README TUI example points at
`train-github-tool.ts`.
> 
> <sup>Reviewed by [Cursor Bugbot](https://cursor.com/bugbot) for commit
99cb197. Bugbot is set up for automated
code reviews on this repo. Configure
[here](https://www.cursor.com/dashboard/bugbot).</sup>
<!-- /CURSOR_SUMMARY -->

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Co-authored-by: p6 <p6@local>