fix(cuda): restore DeepSeek-V2 output by reverting the regressed RMSNorm kernel by inureyes · Pull Request #829 · lablup/mlxcel

inureyes · 2026-07-20T00:40:06Z

Summary

Bumps the vendored MLX pin to the latest upstream (b7c3dd6d) and adds a CUDA RMSNorm overlay that fixes the DeepSeek-V2-Lite generation regression tracked in #824. The MLX pin stays at the latest commit; only the one regressed kernel file is overlaid back to its last-known-good version.

Root cause (pinned by a deterministic bisect)

Upstream MLX #3792 ("Fix CUDA RMSNorm small-row dispatch", commit a5a684d, first shipped in the 0.32.1 pin) regressed the small-axis CUDA RMSNorm. DeepSeek-V2's MLA kv_a_layernorm normalizes over kv_lora_rank (512), which with N_READS == 8 for half-precision activations lands on the multi-warp block_dim == 64 dispatch config that #3792 changed (groups_per_block 2 → 1). The later #3850 rewrite ("RMSNorm forward speed up") kept the same axis broken. The result was a deterministically-wrong normalizer that degenerated DeepSeek-V2 generation into repeated punctuation. Models whose norm axis is larger (qwen3 hidden_size 1024+, etc.) take a different dispatch path and were unaffected, which is why only DeepSeek-V2 broke.

MLX_USE_CUDA_GRAPHS=0 exposes the fault as deterministic; compute-sanitizer initcheck and memcheck are both clean, so it is a numerical kernel error, not a memory-safety bug. Bisect boundary: 1700b39 (commit before #3792) is coherent, a5a684d (#3792) is garbage. This corrects #817's earlier "memory-safety nondeterminism, not bisectable" reading (the graphs-on layer masked the determinism).

Fix

Overlay only mlx/backend/cuda/rms_norm.cu with the last-known-good upstream version (1700b39a, immediately before #3792), via the existing patches/mlx/backend/cuda/ overlay mechanism. Its helper dependencies (load_vector, BlockBroadcastReduce, dispatch_float_types, contiguous_copy_gpu) are unchanged between that commit and the pin, so it compiles as-is against the latest MLX. This trades the #3850 RMSNorm speedup for correctness on the affected small-axis path. The pin is updated in all three synchronized locations (mlx-cpp/CMakeLists.txt GIT_TAG, mlxcel-core/build.rs MLX_EXPECTED_COMMIT, .github/workflows/release.yml).

Verification (GB10)

deepseek-v2-lite-4bit greedy, MLX_USE_CUDA_GRAPHS=0: 4/4 coherent (was 0/4 garbage before the fix).
No regression: qwen3-0.6b-4bit, qwen2.5-0.5b-4bit (answers "Paris"), llama-3.2-1b-4bit all coherent and unchanged.

Scope / follow-ups

The residual graphs-on nondeterminism for DeepSeek-V2 is the pre-existing CUDA-graph-capture layer documented in fix(models): dedicated root-cause debugging of the DeepSeek-V2-Lite greedy nondeterminism (follow-up to #817) #824 (present even before this regression, and at the healthy 2026-05-19 benchmark state); it is out of scope here.
Report the small-axis RMSNorm regression upstream to ml-explore/mlx and drop this overlay once fixed there.

Closes #824

Update: second commit adds the CUDA-graph-capture fix (required for the server)

The RMSNorm overlay above fixes the deterministic forward, but DeepSeek-V2 also has the same CUDA-graph decode hazard as Gemma 4 (#688): with graph capture on, incremental decode degenerates into repeated tokens even though the identical forward is coherent with capture off. This is pre-existing (present before the RMSNorm regression) and dominates the mlxcel-server path (graph capture on by default), so the kernel fix alone left the server returning garbage.

Second commit extends the established MLX_USE_CUDA_GRAPHS=0 startup lever to the DeepSeek-V2 family (maybe_disable_cuda_graphs_for_gemma4 generalized to maybe_disable_cuda_graphs_for_model, now covering Gemma 4 / DeepSeek-V2 / DeepSeek-VL2). An explicit operator MLX_USE_CUDA_GRAPHS still overrides.

End-to-end deployment verification (GB10): mlxcel-server -m deepseek-v2-lite-4bit at default settings now answers coherently 5/5 ("The capital of France is Paris."), where it returned The!!!! garbage before. Startup log: "DeepSeek-V2 detected: disabling CUDA graph capture ... issue #824".

Both layers of #824 are now resolved: the RMSNorm numerical regression (overlay) and the pre-existing graph-capture hazard (startup lever). The upstream small-axis RMSNorm bug report to ml-explore/mlx remains a follow-up.

…orm kernel Bump the vendored MLX pin to the latest upstream (b7c3dd6d) and add a CUDA RMSNorm overlay that fixes the DeepSeek-V2-Lite generation regression tracked in #824. Root cause, pinned by a deterministic bisect: upstream MLX #3792 ("Fix CUDA RMSNorm small-row dispatch", commit a5a684d, first shipped in the 0.32.1 pin) regressed the small-axis CUDA RMSNorm. DeepSeek-V2's MLA kv_a_layernorm normalizes over kv_lora_rank (512), which with N_READS == 8 for half-precision activations lands on the multi-warp block_dim == 64 dispatch config that #3792 changed (groups_per_block 2 -> 1); the later #3850 rewrite ("RMSNorm forward speed up") kept the same axis broken. The result was a deterministically-wrong normalizer that degenerated DeepSeek-V2 text generation into repeated punctuation. Models whose norm axis is larger (qwen3 hidden_size 1024+, etc.) take a different dispatch path and were unaffected, which is why only DeepSeek-V2 broke. The MLX_USE_CUDA_GRAPHS=0 run exposes it as deterministic (compute-sanitizer initcheck and memcheck are both clean, so it is a numerical kernel error, not a memory-safety bug). Bisect: 1700b39 (the commit before #3792) is coherent, a5a684d (#3792) is garbage. Fix: keep the MLX pin at the latest commit and overlay only mlx/backend/cuda/rms_norm.cu with the last-known-good upstream version (1700b39a, i.e. immediately before #3792). The overlay's helper dependencies (load_vector, BlockBroadcastReduce, dispatch_float_types, contiguous_copy_gpu) are unchanged between that commit and the pin, so it compiles as-is against the latest MLX. This trades the #3850 RMSNorm speedup for correctness on the affected small-axis path; re-sync or drop the overlay once the upstream small-axis RMSNorm bug is fixed (an upstream report is a follow-up). Verified on GB10: deepseek-v2-lite-4bit is deterministically coherent again with graphs off (4/4, was 0/4), and qwen3-0.6b / qwen2.5-0.5b / llama-3.2-1b are unchanged (no RMSNorm regression). The residual graphs-on nondeterminism for this model is the pre-existing CUDA-graph-capture layer noted in #824 (present even before this regression) and is out of scope here. The pin is updated in all three synchronized locations (mlx-cpp/CMakeLists.txt GIT_TAG, mlxcel-core/build.rs MLX_EXPECTED_COMMIT, .github/workflows/release.yml). Closes #824

…ecode (#824) The RMSNorm overlay in the previous commit fixes the deterministic forward, but DeepSeek-V2 also has the same CUDA-graph decode hazard that Gemma 4 has (issue #688): with graph capture on, incremental decode degenerates into repeated tokens even though the identical forward is deterministically coherent with capture off. This is pre-existing (present since before the RMSNorm regression) and dominates the server path, which runs graph capture on by default, so without this the model still returns garbage under `mlxcel-server` even with the kernel fixed. Extend the established `MLX_USE_CUDA_GRAPHS=0` startup lever to the DeepSeek-V2 family. Generalize `maybe_disable_cuda_graphs_for_gemma4` to `maybe_disable_cuda_graphs_for_model`, matching Gemma 4 (#688) and DeepSeek-V2 / DeepSeek-VL2 (#824) and logging the detected family and issue. An explicit operator `MLX_USE_CUDA_GRAPHS` still overrides. Unaffected families (Gemma 3, Qwen 3.5, etc.) keep graph capture on. Verified on GB10: `mlxcel-server -m deepseek-v2-lite-4bit` at its default settings now answers coherently (5/5 "The capital of France is Paris."), where it returned `The!!!!` garbage before. The startup log shows "DeepSeek-V2 detected: disabling CUDA graph capture ... issue #824". Refs #824

maybe_disable_cuda_graphs_for_model (src/loading/mod.rs) now covers both Gemma 4 (#688) and DeepSeek-V2 (#824), but its hot-swap Invariant paragraph and the two call-site comments in src/server/startup.rs and src/distributed/pipeline/stage_executor/mod.rs still read as Gemma-4-only, which both the implementation and security reviews flagged as stale. Reworded the Invariant paragraph and the two call-site comments to hazard-family / non-hazard-family wording and referenced issue #824 alongside #688. The detailed Gemma-4 mechanism prose in the loading/mod.rs doc comment is untouched since it remains accurate for that family; no behavior changed. Validation: - cargo fmt --all -- --check: clean - cargo clippy --lib --features cuda -- -D warnings: aborted per the watchdog constraint after it began a cold cutlass/MLX C++ rebuild under a fresh build-script fingerprint; killed the process tree before any object files landed and left the pre-existing green build untouched. Comment-only edits, so no warning surface changed. Refs #824

`expected_rows = (t * gh * gw) as i32` casts an already-i32 expression to i32, which trips `clippy::unnecessary_cast` under `-D warnings`. Pre-existing since the MiniMax-M3 processor landed; surfaced now because this branch touches `src/distributed/pipeline`, which triggers the pipeline-parallel-ci "Clippy on PP modules" step (the only CI job that lints this path). Drop the redundant cast.

…che (#995) ## Summary Adds `mlxcel_core::mla`, a shared matrix-absorbed MLA implementation over a compressed-latent KV cache, and wires `deepseek_v2` onto it behind `MLXCEL_MLA_ABSORBED` (default off). Stage 1 (cache format plus absorbed decode from composed MLX ops) and Stage 2 (split-KV reusing issue #898's merge kernel unchanged) both land here; Stage 2's partial producer is composed from MLX ops, not a fused kernel, and is not a speed path. ## A correction to the issue's premise, up front The issue says the MLA families cache the decompressed K/V. That is true for `deepseek_v2`, `minicpm3`, and `dots1`. It is **not** true for `deepseek_v3` and `deepseek_v32`: both have carried matrix absorption since before this work, decomposing `kv_b_proj` into `embed_q` / `unembed_out` in `sanitize_weights` and caching `(kv_latent, k_pe)` with the rope score passed to SDPA as an additive mask. So the actual gap was V2, MiniCPM3, and dots1, plus the fact that the transform existed twice with no shared home. This PR builds that shared home and moves V2 onto it. V3/V3.2 are untouched (converging them is a follow-up with real regression risk that deserves its own PR), and MiniCPM3/dots1 are untouched and still decompressed. ## What changed - `src/lib/mlxcel-core/src/mla/absorb.rs`: load-time fold of `kv_b_proj` into `W_UK [1, H, qk_nope, r]` for the query path and `W_UV^T [1, H, r, v_head]` for the output path. Dequantizes when quantized (including the NVFP4 `global_scale` sidecar, folded as a scalar since `(x @ W^T) * s2 == x @ (W * s2)^T`), refuses a biased layer because the identity has no term for a bias, and refuses a geometry that does not partition the matrix. Nothing on disk is touched. - `src/lib/mlxcel-core/src/mla/cache.rs`: `MlaLatentCache`, a checked view over `KVCache` that puts `ckv` in the key slot and `kpe` in the value slot. `supports()` is split out of `wrap()` so a caller can fall back (the borrow checker keeps the `&mut` alive through the `Err` arm of a `Result` carrying that borrow's lifetime, which would otherwise make the fallback unexpressible). Also the `latent_bytes_per_token` / `decompressed_bytes_per_token` accounting the KV-memory table is computed from. - `src/lib/mlxcel-core/src/mla/decode.rs`: Stage 1. `SDPA(q_absorbed, ckv, ckv, scale, mask = scale * q_pe @ kpe^T + causal)`. The rope half of the score is additive in exactly SDPA's mask position, so the identity is exact and MLX's fused attention kernel still does the softmax and the value accumulation. Same shape `deepseek_v3.rs` already uses. - `src/lib/mlxcel-core/src/mla/split_kv.rs`: Stage 2. `MlaSplitPlan` cuts the latent range; each chunk produces a normalized partial and its log2 LSE; `paged_attention_merge_states` merges them. - `src/lib/mlxcel-core/src/mla/stats.rs`: per-path call counters and a `path=count` summary. - `src/models/deepseek_v2.rs`: `MLAAttention` gains an optional fold and a `forward_absorbed` branch; `load_mla_attention` builds the fold only under the flag; `report_absorption` prints the load-time decision. - `examples/mla_absorbed_decode_bench.rs`: three-arm decode benchmark plus the KV-memory table. - `docs/mla-absorbed-decode.md`, indexed from `docs/README.md`, plus both variables in `docs/environment-variables.md`. ## Why the cache is a view over `KVCache` and not a new type `LanguageModel::make_caches` returns `Vec<KVCache>` and `forward` takes `&mut [KVCache]`, so a genuinely new cache type forces the family into the `ModelOwnedSequenceState` escape hatch the SSM hybrids use, which costs `supports_batching`, the paged decode backend, and the generic prompt-cache donation path. The FP16 update path already stores keys and values as independent `[B, H, L, D]` buffers reading each side's head dim from its own shape, so a 512-wide "key" beside a 64-wide "value" is already expressible. `offset`, `live_start`, `trim`, `trim_front`, `nbytes`, and `can_trim_prompt_cache` therefore keep their existing meanings on latent rows, and the slot assignment matches what `deepseek_v3.rs` has always done so both describe the same buffer. ## How #898's merge kernel is reused, and whether its contract held Reused verbatim. No change to `paged_attention_v2_merge.cpp`, to `paged_attention_v2.h`, to the `ffi::paged_attention_merge_states` signature, or to the C++ launcher. MLA fits the documented contract as `H = num_heads`, `D = kv_lora_rank`, `M = batch`, `N = batch * chunks`, and `o_indptr` request-major. The contract held unchanged; the only thing the new caller had to honour rather than discover is the **log2 LSE** clause, which is the one that fails silently: a natural-log LSE still merges and returns a plausible wrong weighted average. It is a named constant (`split_kv::LOG2_E`) with a comment, and `merge_rejects_natural_log_lse_units` is a negative control that asserts the natural-log variant does *not* match a host f64 reference while the log2 variant does. `split_kv_matches_the_decompressed_reference_at_every_chunk_length` restates, at the MLA caller, the regrouping property `paged_v2::launch_tests::merge_is_associative_across_regroupings` pins for the kernel. ## Verified vs unverified **Verified, on synthetic tensors with MLA geometry against a host-side f64 reference that performs the pre-absorption computation:** the absorption identity in f32 and f16; that the additive causal mask survives being merged into the rope score (with a positive control asserting the mask actually changes the answer, so the masked test cannot pass on a no-op); the up-projection direction of the fold; the `(ckv, kpe)` packing round-tripping through a real `KVCache` including the append boundary; that a latent cache stays trimmable so speculative decode's reject path removes one latent row per token; the split-and-merge decomposition at four chunk lengths and against the unsplit Stage 1 answer; the log2 LSE negative control; and, through `deepseek_v2`'s own `MLAAttention::forward`, that a 6-token prefill plus 4 decode steps produce the same block output (through `o_proj`) with and without the fold, step for step, and that the resulting cache is more than 3x smaller. **Not verified:** anything against a real MLA checkpoint. No DeepSeek-family checkpoint is available on this host (the locally available models are Llama-3.1-8B, Llama-3.2-1B, Qwen2.5-0.5B, Qwen3-0.6B, and gemma-3-1b, none of them MLA, and a DeepSeek-V3-class checkpoint does not fit in 128 GB). So the issue's token-parity gate on pinned DeepSeek oracles is **not** satisfied by this PR, and the flag stays default-off for exactly that reason. Nothing here should be read as real-model validation. **No benchmarks were run.** The harness is written and smoke-tested for correctness on a toy geometry only; all throughput numbers are for the reviewer to produce. ## CUDA No CUDA was written in this PR, so there is no new unvalidated CUDA to declare. Stage 2 reuses #898's merge kernel, which already ships a structurally parallel CUDA port that #898 declared unvalidated; that status is unchanged. The issue asked for a fused MLA partial kernel with CUDA first, and that is deliberately not in this PR: a fused partial kernel is the thing that would make Stage 2 faster, and it should be written after Stage 1's measurements justify it (as the issue itself scopes it) and against a host that can validate at least one of the two backends. The interactions the issue lists are unaffected: `kv_a_layernorm` is untouched, so the CUDA RMSNorm overlay (#829) sits where it sat; nothing here assumes CUDA graph replay. ## How to run the benchmark, and how to prove which path it took ``` export DEVELOPER_DIR=/Applications/Xcode-26.6.0.app/Contents/Developer cargo run --release --features metal,accelerate --example mla_absorbed_decode_bench -- \ --contexts 4096,16384,32768 --batches 1,4 --steps 64 --warmup 16 ``` Three arms per (context, batch): `decompressed`, `absorbed`, `split_kv`. Every row ends in a `paths=` field read from `mlxcel_core::mla::stats`, taken and reset around the measured region, and the harness prints a `WARNING` on any row whose counters disagree with its arm's label. A correct run looks like `decompressed=64 absorbed_prefill=0 absorbed_composed=0 absorbed_split_kv=0` on the first arm and `... absorbed_composed=64 ...` on the second; an "absorbed" row reading `absorbed_composed=0` is measuring the fallback and must not be reported. `split_kv` is expected to be slower than `absorbed` here, because its partials are composed from MLX ops rather than fused; that arm exists to show the decomposition running, not to be quoted. For a real checkpoint, `MLXCEL_MLA_ABSORBED=1 mlxcel generate -m models/<deepseek> ...` prints one stdout line at load stating how many layers folded and the bytes/token before and after. A run whose line reads `0/N layers` is running the fallback. Both signals go to stdout rather than through `tracing` because the `mlxcel` CLI installs no tracing subscriber, so `tracing::info!` on this path emits nothing at any `RUST_LOG`. ## Test plan - [x] `cargo test --release -p mlxcel-core --lib --features metal,accelerate mla::` — 33 passed, 0 failed - [x] `cargo test --release -p mlxcel-core --lib --features metal,accelerate paged_v2::` — 55 passed, 0 failed (the reused merge kernel and its regrouping property) - [x] `cargo test --release --lib --features metal,accelerate models::deepseek` — 32 passed, 0 failed (V2 wiring parity plus the untouched V3 / V3.2 suites) - [x] `cargo clippy -p mlxcel-core --lib --tests --features metal,accelerate -- -D warnings` — clean - [x] `cargo clippy --lib --tests --example mla_absorbed_decode_bench --features metal,accelerate -- -D warnings` — clean - [x] `cargo fmt --all -- --check` — clean - [x] `python3 scripts/ci/check_cross_repo_refs.py` — exit 0 - [x] Benchmark harness smoke-run on a toy geometry: all three arms execute and each reports its own path The full `cargo test -p mlxcel-core` sweep was not run: it exceeds this environment's watchdog. Its known pre-existing failure on this base commit is `fused_moe_parity_tests::fused_moe_geglu_kernel_matches_references_gemma4_shape` (tracked as #964), untouched here. ## Left out, deliberately Fused Stage 2 partial kernel (Metal and CUDA). MiniCPM3 and dots1 wiring. Converging `deepseek_v3` / `deepseek_v32` onto the shared module. Absorbed prefill (prefill up-projects through the still-loaded `kv_b_proj`, which keeps prefill arithmetic identical to the baseline but makes a chunked prefill re-up-project rows it already handled). Paged-backed and non-FP16 latent caches, both declined with a message. Any default-on change: the flag stays off until real-checkpoint token parity exists. Closes #907

## Summary Issue #830 asked for an upstream bug report against ml-explore/mlx: MLX #3792 was said to have regressed the small-axis CUDA RMSNorm, with #3850 keeping the same axis broken. Verifying that against real upstream source inverted it, so **no upstream report is filed and none should be**. The kernel at mlxcel's MLX pin is correct on the affected axis. The pre-#3792 kernel that mlxcel carries as `src/lib/mlx-cpp/patches/mlx/backend/cuda/rms_norm.cu` is the one that computes a wrong normalizer, and it reads out of bounds while doing it. ## What was verified, and against what The pin is not `0.32.1` as #830 states. `src/lib/mlx-cpp/CMakeLists.txt:95` pins `b7c3dd6d27f45b5365b08a840310187dc503f1db`, which is upstream PR #3850 itself. The on-disk MLX checkout under `target/` already has the overlay applied, so all three kernels were fetched from GitHub at their own commits instead: `1700b39a` (the direct parent of #3792, confirmed via the commit's `parents`), `a5a684d` (#3792), and the pin. Measured on GB10 (sm_121, CUDA 13) with the three kernel bodies copied verbatim, each launched with its own commit's dispatch configuration and diffed against a per-row float64 reference: | Kernel | Failing axis sizes, 8 to 8192, row counts 1 to 17 | |---|---| | `1700b39` (the overlay mlxcel ships) | all 34 swept sizes in `[257, 512]` | | `a5a684d` (#3792) | none | | `b7c3dd6d` (#3850, the pin) | none | `compute-sanitizer --tool memcheck` reports `Invalid __shared__ read of size 4 bytes` for the overlay's launch geometry only; the other two are clean. ## Root cause At `axis_size` in `(256, 512]` for 16-bit dtypes, the pre-#3792 dispatch pairs `block_dim == 64` with `groups_per_block == 2`. That launches a 128-thread block, but the reduction is `BlockBroadcastReduce<float, 64, 32>`, whose `TempStorage` is sized `BLOCK_DIM / WARP_SIZE == 2` while `cg::tiled_partition<32>` of the launched block reports `meta_group_size() == 4`. The second reduction stage indexes two floats past the end of the scratch array and folds in the partials belonging to the other row that shares the block. Two failure modes follow. Above one row the sum covers both rows, so the normalizer is deterministically wrong (0.55 to 0.94 relative error). At one row the `threadIdx.y == 1` threads take the row-bounds early return and never write their partials, so the normalizer picks up residual shared memory. Seeding that window with 0, 1, 100 and 10000 moves the applied normalizer to 99.99%, 92.79%, 24.19% and 2.49% of the float64 reference. #3792 changed that one value to 1 and added `static_assert(block_dim <= 32 || groups_per_block() == 1)` to pin the invariant. It is a fix, and its title says so. #3850 rewrote the forward path to one row per block, where block size and template parameter agree by construction. #830 got the band, the config, the changed value, and DeepSeek-V2's `kv_lora_rank == 512` landing on it all correct. Only the direction is backwards. ## Not reconciled #829 measured deepseek-v2-lite going 0/4 to 4/4 coherent with the overlay as the only change. Against these measurements the overlay should have made that axis worse. The document says so plainly rather than papering over it, and notes two unverified leads: #817's original "nondeterministic across fresh processes" symptom matches the uninitialized read exactly, and graph capture changes which kernels precede the norm, which is what determines the overlay's answer at one row. Settling it needs a build with the overlay removed, measured the way #829 measured. That is a separate change with real risk and is deliberately not attempted here, so **no runtime behavior changes in this PR**. ## What changed - `docs/upstream/mlx-cuda-rmsnorm-small-axis-regression.md` (new): the report, following the two existing `docs/upstream/` examples. Summary, evidence, root cause with quoted source at every cited line, minimal fix, repro, affected surface. The standalone CUDA repro was run; the `mlx.core` sketch and the model-level repro are labelled not run. - `src/lib/mlx-cpp/patches/mlx/backend/cuda/rms_norm.cu`: the tracking comment no longer dangles on "report filed as a follow-up". It points at the document, marks the superseded paragraph, and carries the corrected removal condition: nothing needs to land upstream first, so the overlay can be deleted once deepseek-v2-lite is re-verified on the upstream fast kernel, recovering the #3850 speedup. Comment lines above the `---- begin upstream ----` marker only; the vendored kernel body is untouched. The patch-file sanity check at `CMakeLists.txt:115-126` globs `patches/mlx/backend/metal/*.cpp` and does not match this file. ## Follow-up worth filing The overlay is wrong on an axis nine model files reach through `kv_a_layernorm`. Dropping it and re-verifying deepseek-v2-lite is the obvious next step, but it needs a real model run and is out of scope for a documentation change. ## Test plan - [x] `cargo fmt --all -- --check` - [x] `cargo clippy -p mlxcel-core -- -D warnings` (clean, 16.23s) - [x] `cargo check -p mlxcel-core --features cuda` (clean, 29.20s; re-applied the overlay and recompiled `rms_norm.cu.o` through nvcc, so the comment edit is confirmed to build) - [x] Overlay diff is comment-only: zero non-`//` lines changed - [x] No em dashes and no AI attribution in the diff or the commit message - [ ] `cargo clippy -p mlxcel-core --features cuda` **skipped**: it hashes to a fresh build directory and starts a cold full MLX CUDA rebuild. Aborted after two minutes and the partial directory removed. No Rust changed in this PR, so it carried no signal. Closes #830

#1037) ## Summary Root-causes the DeepSeek-V2 "CUDA-graph-captured incremental decode collapse" from #824/#829/#831: it was not a graph-capture hazard but mlxcel's own RMSNorm overlay (`patches/mlx/backend/cuda/rms_norm.cu`, upstream pre-#3792), whose two-stage reduction reads past its shared scratch on the exact `kv_a_layernorm` axis (`kv_lora_rank == 512` at bf16). Deletes the overlay, removes DeepSeek-V2 and DeepSeek-VL2 from `maybe_disable_cuda_graphs_for_model` so they run with CUDA graphs again, and pins the band with a numeric regression test. Gemma 4 stays on the lever: its norm axes never touch the broken dispatch geometry, so #688 is a separate, still-open bug, and #831's premise that the two families shared one hazard class is wrong. ## The decisive experiment 2x2 matrix on GB10 at `13741ae2`: deepseek-v2-lite-4bit, greedy `--temp 0`, 48 tokens, fresh process per rep, 6 reps per cell per prompt, two prompts (1-token "Hi", ~90-token summarization). "Identical" is by content hash of the generated text. | Cell | Kernel | `MLX_USE_CUDA_GRAPHS` | Short prompt | Long prompt | |---|---|---|---|---| | A | overlay `1700b39` | 0 | 6/6 coherent, identical | 6/6 coherent, identical | | B | overlay `1700b39` | 1 | 4/6 coherent, all four token-divergent; 2/6 collapse into repeated `!` | 0/6: every rep collapses after one generated token | | C | pinned `b7c3dd6d` | 0 | 6/6 coherent, identical | 6/6 coherent, identical | | D | pinned `b7c3dd6d` | 1 | 6/6 coherent, identical, byte-identical to C | 6/6 coherent, identical, byte-identical to C | Cells C and D being byte-identical means graph capture has no observable effect on this family once the kernel is correct. The overlay's single-row failure mode (normalizer depends on residual shared memory) explains every recorded symptom: greedy flakiness, coherent-prefix-then-collapse, prompt-length dependence, and `MLXCEL_FORCE_SYNC=1` not helping. The cell C/D binary was proven to contain the pinned kernel (the `configure_file(COPYONLY)` overlay mechanism makes a silent no-op rebuild possible): post-build source diff against `b7c3dd6d`, `rms_norm.cu.o` recompile mtime, and `cuobjdump`/`strings` mangling showing 128 `rms_norm_small` instantiations with #3850's `bool ALIGNED` parameter (`Lb0E`/`Lb1E`) and zero of the overlay's four-int form. ## What changed - `src/lib/mlx-cpp/patches/mlx/backend/cuda/rms_norm.cu`: deleted. The pinned kernel is correct on the affected axis and carries the #3850 forward speedup. - `src/loading/mod.rs`: `ModelType::DeepSeekV2 | ModelType::DeepSeekVL2` removed from the hazard match; function docs rewritten to record the corrected attribution. Gemma 4 arms unchanged. - `src/server/startup.rs`: main-thread hoist comment updated to the Gemma-4-only story. - `src/lib/mlxcel-core/src/rms_norm_small_axis_tests.rs` (new, registered in `lib.rs`): pins `fast::rms_norm` against a float64 host reference across the broken band, 16-bit `(256, 512]` and f32 `(128, 256]`, rows 8/3/2/1. Verified non-vacuous: with the overlay kernel linked back in, both tests fail at ~3e-1 worst-row relative error; on the pin they pass at rounding scale. - `docs/upstream/mlx-cuda-rmsnorm-small-axis-regression.md`: resolution section with the matrix, the binary proof, the non-reproduction of #829's 0/4 baseline, and the local checkpoints that sat in the band (five MLA models at `kv_lora_rank == 512` plus falcon-h1-tiny whose `hidden_size` itself is 512). ## Test plan - [x] `cargo fmt --all -- --check` clean - [x] `cargo clippy --release --features cuda -p mlxcel-core --lib --tests` and `-p mlxcel --bins`: no lints - [x] `cargo test --release --features cuda -p mlxcel-core --lib rms_norm_small_axis`: 2/2 (and 0/2 against the overlay kernel, restored afterwards) - [x] `cargo test --release --features cuda -p mlxcel-core --lib fused_norm_parity`: 8/8 - [x] `cargo test --release --features cuda --test deepseek_v2_greedy_parity`: 2/2, token-exact vs mlx-lm 0.31.3 with graphs now on by default - [x] Smoke: deepseek-v2-lite-4bit coherent with no env override (no disable log), deepseek-vl2-small-4bit coherent, gemma-4-e2b-it-4bit still logs the #688 graph disable - [x] 200-token decode: graphs-on 85.0-93.3 tok/s vs graphs-off 86.5-88.1 (n=3 each; GB10 single-run decode is bimodal, coherence is the gate) Closes #831

#1046) ## Summary Bumps the MLX pin from `b7c3dd6d` (2026-07-17) to `2c46b953` (2026-08-05 upstream main), 58 commits, and re-derives the three source overlays that upstream touched in that window. Closes #1042 ## Why Fixes in range that land on paths this repo ships: | Upstream | What | |---|---| | #3854 | incorrect nvfp4 `quantized_matmul` through the split-K path (NVFP4 direct transcode is mlxcel's default on all builds) | | #3893 | CUDA gemm conv unfold grid overflow past 65,535 output positions | | #3934 | Metal NVFP4 scale reduction per 16-lane group | | #3888 | Metal `gemv_wide` for few-row fp16/bf16 matmuls | | #3943 | head dimension 96 in Metal full attention | | #3898, #3813 | `prod` dtype promotion on a size-1 axis, shapeless matmul with dynamic batch dims | ## Overlay audit mlxcel applies whole-file overlays via `configure_file(... COPYONLY)`, so an overlay whose upstream counterpart moved silently reverts that upstream change. That is how a broken RMSNorm kernel shipped in #824 / #829 and took until #830 / #831 to undo. All 21 overlays were diffed against the upstream range; three conflicted. | Overlay | Upstream change | Resolution | |---|---|---| | `patches/mlx/backend/cuda/quantized/quantized.cpp` | #3757 inserted `const std::optional<array>& global_scale` into `qmm_naive` at position 5 | resolved by hand: mlxcel wraps both call sites in `run_row_chunked` / `run_batch_chunked`, which `git merge-file` cannot reconcile with an inline insertion. Both sites verified as 13 arguments in declaration order | | `patches/mlx/backend/metal/kernels/utils.h` | #3970 templated complex scalar lanes, #3963 explicit `thread` address space for metal 4.1 | clean three-way merge, no overlap with mlxcel's delta | | `patches-cuda/ops.cpp` | 189 lines | clean three-way merge, mlxcel's 57-line delta intact | `quantized/qmm/qmm.h` also changed but is not overlaid, so mlxcel takes the new declaration. Its full diff across the range is the single `global_scale` line, so `qmm_sm90`, `qmm_sm80`, `qmv`, `gather_qmv` and `fp_qmv` are unaffected. The `qmm_sm80.cuh` race-fix overlay from #910 is untouched upstream and still required. The stale `Synced to upstream e9463bb` header comment is updated to the real base, with an explicit instruction to re-verify these call sites on the next bump. ## The pin has two sources of truth `GIT_TAG` in `src/lib/mlx-cpp/CMakeLists.txt:95` and `MLX_EXPECTED_COMMIT` in `src/lib/mlxcel-core/build.rs:217`. The latter drives `purge_stale_mlx_cache`. Updating only the CMakeLists leaves the cache marker matching `MLX_EXPECTED_COMMIT`, so `_deps/` is judged valid, never purged, and the build silently compiles the OLD MLX while looking bumped. Both moved together, verified three ways with no trace of the old SHA: - binary embedded commit: `2c46b953…` - `_deps/.mlx-build-commit` marker: `2c46b953…` - fetched `_deps/mlx-src` HEAD: `2c46b953…` Worth a follow-up to collapse this to one source of truth or add a build-time assert. ## Validation (GB10, sm_121) - release CUDA build clean - `cargo fmt --all -- --check` clean - `fused_norm_parity` 8/8, `turbo_tests` 97/97 - greedy smoke on `llama-3.2-1b-4bit`, `qwen3-4b-4bit`, `deepseek-v2-lite-4bit`: coherent, and byte-identical run to run with a warm kernel cache. That determinism also confirms the `qmm_sm80` race-fix overlay still holds on the new base. ## Two pre-existing failures, separated by a control build A control build at the old pin (`_deps` re-fetched to `b7c3dd6d`, verified) reproduces both, so neither is caused by this bump. | Failure | New pin | Old pin (control) | |---|---|---| | `fused_rope_parity` 2 of 9 | offset 131069 rms 2.497e-3, offset 131071 rms 2.684e-3, tol 2.0e-3 | same digits, same 7 pass / 2 fail | | full lib suite abort under graph capture | `cuGraphAddKernelNode` invalid argument, then `cudaStreamEndCapture` previous error during capture (site and message vary per run) | aborts the same way | For the rope case the numbers match bit for bit, which also shows that path carries no exposure to this change. Upstream touched no rope file in the range. Both are left for separate issues rather than folded in here.

…orm kernel (lablup#829) ## Summary Bumps the vendored MLX pin to the latest upstream (`b7c3dd6d`) and adds a CUDA RMSNorm overlay that fixes the DeepSeek-V2-Lite generation regression tracked in lablup#824. The MLX pin stays at the latest commit; only the one regressed kernel file is overlaid back to its last-known-good version. ## Root cause (pinned by a deterministic bisect) Upstream MLX **#3792** ("Fix CUDA RMSNorm small-row dispatch", commit `a5a684d`, first shipped in the `0.32.1` pin) regressed the small-axis CUDA RMSNorm. DeepSeek-V2's MLA `kv_a_layernorm` normalizes over `kv_lora_rank` (512), which with `N_READS == 8` for half-precision activations lands on the multi-warp `block_dim == 64` dispatch config that #3792 changed (`groups_per_block` 2 → 1). The later **#3850** rewrite ("RMSNorm forward speed up") kept the same axis broken. The result was a deterministically-wrong normalizer that degenerated DeepSeek-V2 generation into repeated punctuation. Models whose norm axis is larger (qwen3 `hidden_size` 1024+, etc.) take a different dispatch path and were unaffected, which is why only DeepSeek-V2 broke. `MLX_USE_CUDA_GRAPHS=0` exposes the fault as deterministic; `compute-sanitizer` initcheck and memcheck are both clean, so it is a numerical kernel error, not a memory-safety bug. Bisect boundary: `1700b39` (commit before #3792) is coherent, `a5a684d` (#3792) is garbage. This corrects lablup#817's earlier "memory-safety nondeterminism, not bisectable" reading (the graphs-on layer masked the determinism). ## Fix Overlay only `mlx/backend/cuda/rms_norm.cu` with the last-known-good upstream version (`1700b39a`, immediately before #3792), via the existing `patches/mlx/backend/cuda/` overlay mechanism. Its helper dependencies (`load_vector`, `BlockBroadcastReduce`, `dispatch_float_types`, `contiguous_copy_gpu`) are unchanged between that commit and the pin, so it compiles as-is against the latest MLX. This trades the #3850 RMSNorm speedup for correctness on the affected small-axis path. The pin is updated in all three synchronized locations (`mlx-cpp/CMakeLists.txt` GIT_TAG, `mlxcel-core/build.rs` MLX_EXPECTED_COMMIT, `.github/workflows/release.yml`). ## Verification (GB10) - `deepseek-v2-lite-4bit` greedy, `MLX_USE_CUDA_GRAPHS=0`: **4/4 coherent** (was 0/4 garbage before the fix). - No regression: `qwen3-0.6b-4bit`, `qwen2.5-0.5b-4bit` (answers "Paris"), `llama-3.2-1b-4bit` all coherent and unchanged. ## Scope / follow-ups - The residual graphs-on nondeterminism for DeepSeek-V2 is the **pre-existing** CUDA-graph-capture layer documented in lablup#824 (present even before this regression, and at the healthy 2026-05-19 benchmark state); it is out of scope here. - Report the small-axis RMSNorm regression upstream to ml-explore/mlx and drop this overlay once fixed there. Closes lablup#824 --- ## Update: second commit adds the CUDA-graph-capture fix (required for the server) The RMSNorm overlay above fixes the deterministic forward, but DeepSeek-V2 also has the same CUDA-graph decode hazard as Gemma 4 (lablup#688): with graph capture on, incremental decode degenerates into repeated tokens even though the identical forward is coherent with capture off. This is pre-existing (present before the RMSNorm regression) and dominates the `mlxcel-server` path (graph capture on by default), so the kernel fix alone left the server returning garbage. Second commit extends the established `MLX_USE_CUDA_GRAPHS=0` startup lever to the DeepSeek-V2 family (`maybe_disable_cuda_graphs_for_gemma4` generalized to `maybe_disable_cuda_graphs_for_model`, now covering Gemma 4 / DeepSeek-V2 / DeepSeek-VL2). An explicit operator `MLX_USE_CUDA_GRAPHS` still overrides. **End-to-end deployment verification (GB10):** `mlxcel-server -m deepseek-v2-lite-4bit` at default settings now answers coherently **5/5** ("The capital of France is Paris."), where it returned `The!!!!` garbage before. Startup log: "DeepSeek-V2 detected: disabling CUDA graph capture ... issue lablup#824". Both layers of lablup#824 are now resolved: the RMSNorm numerical regression (overlay) and the pre-existing graph-capture hazard (startup lever). The upstream small-axis RMSNorm bug report to ml-explore/mlx remains a follow-up.

…che (lablup#995) ## Summary Adds `mlxcel_core::mla`, a shared matrix-absorbed MLA implementation over a compressed-latent KV cache, and wires `deepseek_v2` onto it behind `MLXCEL_MLA_ABSORBED` (default off). Stage 1 (cache format plus absorbed decode from composed MLX ops) and Stage 2 (split-KV reusing issue lablup#898's merge kernel unchanged) both land here; Stage 2's partial producer is composed from MLX ops, not a fused kernel, and is not a speed path. ## A correction to the issue's premise, up front The issue says the MLA families cache the decompressed K/V. That is true for `deepseek_v2`, `minicpm3`, and `dots1`. It is **not** true for `deepseek_v3` and `deepseek_v32`: both have carried matrix absorption since before this work, decomposing `kv_b_proj` into `embed_q` / `unembed_out` in `sanitize_weights` and caching `(kv_latent, k_pe)` with the rope score passed to SDPA as an additive mask. So the actual gap was V2, MiniCPM3, and dots1, plus the fact that the transform existed twice with no shared home. This PR builds that shared home and moves V2 onto it. V3/V3.2 are untouched (converging them is a follow-up with real regression risk that deserves its own PR), and MiniCPM3/dots1 are untouched and still decompressed. ## What changed - `src/lib/mlxcel-core/src/mla/absorb.rs`: load-time fold of `kv_b_proj` into `W_UK [1, H, qk_nope, r]` for the query path and `W_UV^T [1, H, r, v_head]` for the output path. Dequantizes when quantized (including the NVFP4 `global_scale` sidecar, folded as a scalar since `(x @ W^T) * s2 == x @ (W * s2)^T`), refuses a biased layer because the identity has no term for a bias, and refuses a geometry that does not partition the matrix. Nothing on disk is touched. - `src/lib/mlxcel-core/src/mla/cache.rs`: `MlaLatentCache`, a checked view over `KVCache` that puts `ckv` in the key slot and `kpe` in the value slot. `supports()` is split out of `wrap()` so a caller can fall back (the borrow checker keeps the `&mut` alive through the `Err` arm of a `Result` carrying that borrow's lifetime, which would otherwise make the fallback unexpressible). Also the `latent_bytes_per_token` / `decompressed_bytes_per_token` accounting the KV-memory table is computed from. - `src/lib/mlxcel-core/src/mla/decode.rs`: Stage 1. `SDPA(q_absorbed, ckv, ckv, scale, mask = scale * q_pe @ kpe^T + causal)`. The rope half of the score is additive in exactly SDPA's mask position, so the identity is exact and MLX's fused attention kernel still does the softmax and the value accumulation. Same shape `deepseek_v3.rs` already uses. - `src/lib/mlxcel-core/src/mla/split_kv.rs`: Stage 2. `MlaSplitPlan` cuts the latent range; each chunk produces a normalized partial and its log2 LSE; `paged_attention_merge_states` merges them. - `src/lib/mlxcel-core/src/mla/stats.rs`: per-path call counters and a `path=count` summary. - `src/models/deepseek_v2.rs`: `MLAAttention` gains an optional fold and a `forward_absorbed` branch; `load_mla_attention` builds the fold only under the flag; `report_absorption` prints the load-time decision. - `examples/mla_absorbed_decode_bench.rs`: three-arm decode benchmark plus the KV-memory table. - `docs/mla-absorbed-decode.md`, indexed from `docs/README.md`, plus both variables in `docs/environment-variables.md`. ## Why the cache is a view over `KVCache` and not a new type `LanguageModel::make_caches` returns `Vec<KVCache>` and `forward` takes `&mut [KVCache]`, so a genuinely new cache type forces the family into the `ModelOwnedSequenceState` escape hatch the SSM hybrids use, which costs `supports_batching`, the paged decode backend, and the generic prompt-cache donation path. The FP16 update path already stores keys and values as independent `[B, H, L, D]` buffers reading each side's head dim from its own shape, so a 512-wide "key" beside a 64-wide "value" is already expressible. `offset`, `live_start`, `trim`, `trim_front`, `nbytes`, and `can_trim_prompt_cache` therefore keep their existing meanings on latent rows, and the slot assignment matches what `deepseek_v3.rs` has always done so both describe the same buffer. ## How lablup#898's merge kernel is reused, and whether its contract held Reused verbatim. No change to `paged_attention_v2_merge.cpp`, to `paged_attention_v2.h`, to the `ffi::paged_attention_merge_states` signature, or to the C++ launcher. MLA fits the documented contract as `H = num_heads`, `D = kv_lora_rank`, `M = batch`, `N = batch * chunks`, and `o_indptr` request-major. The contract held unchanged; the only thing the new caller had to honour rather than discover is the **log2 LSE** clause, which is the one that fails silently: a natural-log LSE still merges and returns a plausible wrong weighted average. It is a named constant (`split_kv::LOG2_E`) with a comment, and `merge_rejects_natural_log_lse_units` is a negative control that asserts the natural-log variant does *not* match a host f64 reference while the log2 variant does. `split_kv_matches_the_decompressed_reference_at_every_chunk_length` restates, at the MLA caller, the regrouping property `paged_v2::launch_tests::merge_is_associative_across_regroupings` pins for the kernel. ## Verified vs unverified **Verified, on synthetic tensors with MLA geometry against a host-side f64 reference that performs the pre-absorption computation:** the absorption identity in f32 and f16; that the additive causal mask survives being merged into the rope score (with a positive control asserting the mask actually changes the answer, so the masked test cannot pass on a no-op); the up-projection direction of the fold; the `(ckv, kpe)` packing round-tripping through a real `KVCache` including the append boundary; that a latent cache stays trimmable so speculative decode's reject path removes one latent row per token; the split-and-merge decomposition at four chunk lengths and against the unsplit Stage 1 answer; the log2 LSE negative control; and, through `deepseek_v2`'s own `MLAAttention::forward`, that a 6-token prefill plus 4 decode steps produce the same block output (through `o_proj`) with and without the fold, step for step, and that the resulting cache is more than 3x smaller. **Not verified:** anything against a real MLA checkpoint. No DeepSeek-family checkpoint is available on this host (the locally available models are Llama-3.1-8B, Llama-3.2-1B, Qwen2.5-0.5B, Qwen3-0.6B, and gemma-3-1b, none of them MLA, and a DeepSeek-V3-class checkpoint does not fit in 128 GB). So the issue's token-parity gate on pinned DeepSeek oracles is **not** satisfied by this PR, and the flag stays default-off for exactly that reason. Nothing here should be read as real-model validation. **No benchmarks were run.** The harness is written and smoke-tested for correctness on a toy geometry only; all throughput numbers are for the reviewer to produce. ## CUDA No CUDA was written in this PR, so there is no new unvalidated CUDA to declare. Stage 2 reuses lablup#898's merge kernel, which already ships a structurally parallel CUDA port that lablup#898 declared unvalidated; that status is unchanged. The issue asked for a fused MLA partial kernel with CUDA first, and that is deliberately not in this PR: a fused partial kernel is the thing that would make Stage 2 faster, and it should be written after Stage 1's measurements justify it (as the issue itself scopes it) and against a host that can validate at least one of the two backends. The interactions the issue lists are unaffected: `kv_a_layernorm` is untouched, so the CUDA RMSNorm overlay (lablup#829) sits where it sat; nothing here assumes CUDA graph replay. ## How to run the benchmark, and how to prove which path it took ``` export DEVELOPER_DIR=/Applications/Xcode-26.6.0.app/Contents/Developer cargo run --release --features metal,accelerate --example mla_absorbed_decode_bench -- \ --contexts 4096,16384,32768 --batches 1,4 --steps 64 --warmup 16 ``` Three arms per (context, batch): `decompressed`, `absorbed`, `split_kv`. Every row ends in a `paths=` field read from `mlxcel_core::mla::stats`, taken and reset around the measured region, and the harness prints a `WARNING` on any row whose counters disagree with its arm's label. A correct run looks like `decompressed=64 absorbed_prefill=0 absorbed_composed=0 absorbed_split_kv=0` on the first arm and `... absorbed_composed=64 ...` on the second; an "absorbed" row reading `absorbed_composed=0` is measuring the fallback and must not be reported. `split_kv` is expected to be slower than `absorbed` here, because its partials are composed from MLX ops rather than fused; that arm exists to show the decomposition running, not to be quoted. For a real checkpoint, `MLXCEL_MLA_ABSORBED=1 mlxcel generate -m models/<deepseek> ...` prints one stdout line at load stating how many layers folded and the bytes/token before and after. A run whose line reads `0/N layers` is running the fallback. Both signals go to stdout rather than through `tracing` because the `mlxcel` CLI installs no tracing subscriber, so `tracing::info!` on this path emits nothing at any `RUST_LOG`. ## Test plan - [x] `cargo test --release -p mlxcel-core --lib --features metal,accelerate mla::` — 33 passed, 0 failed - [x] `cargo test --release -p mlxcel-core --lib --features metal,accelerate paged_v2::` — 55 passed, 0 failed (the reused merge kernel and its regrouping property) - [x] `cargo test --release --lib --features metal,accelerate models::deepseek` — 32 passed, 0 failed (V2 wiring parity plus the untouched V3 / V3.2 suites) - [x] `cargo clippy -p mlxcel-core --lib --tests --features metal,accelerate -- -D warnings` — clean - [x] `cargo clippy --lib --tests --example mla_absorbed_decode_bench --features metal,accelerate -- -D warnings` — clean - [x] `cargo fmt --all -- --check` — clean - [x] `python3 scripts/ci/check_cross_repo_refs.py` — exit 0 - [x] Benchmark harness smoke-run on a toy geometry: all three arms execute and each reports its own path The full `cargo test -p mlxcel-core` sweep was not run: it exceeds this environment's watchdog. Its known pre-existing failure on this base commit is `fused_moe_parity_tests::fused_moe_geglu_kernel_matches_references_gemma4_shape` (tracked as lablup#964), untouched here. ## Left out, deliberately Fused Stage 2 partial kernel (Metal and CUDA). MiniCPM3 and dots1 wiring. Converging `deepseek_v3` / `deepseek_v32` onto the shared module. Absorbed prefill (prefill up-projects through the still-loaded `kv_b_proj`, which keeps prefill arithmetic identical to the baseline but makes a chunked prefill re-up-project rows it already handled). Paged-backed and non-FP16 latent caches, both declined with a message. Any default-on change: the flag stays off until real-checkpoint token parity exists. Closes lablup#907

…up#1033) ## Summary Issue lablup#830 asked for an upstream bug report against ml-explore/mlx: MLX #3792 was said to have regressed the small-axis CUDA RMSNorm, with #3850 keeping the same axis broken. Verifying that against real upstream source inverted it, so **no upstream report is filed and none should be**. The kernel at mlxcel's MLX pin is correct on the affected axis. The pre-#3792 kernel that mlxcel carries as `src/lib/mlx-cpp/patches/mlx/backend/cuda/rms_norm.cu` is the one that computes a wrong normalizer, and it reads out of bounds while doing it. ## What was verified, and against what The pin is not `0.32.1` as lablup#830 states. `src/lib/mlx-cpp/CMakeLists.txt:95` pins `b7c3dd6d27f45b5365b08a840310187dc503f1db`, which is upstream PR #3850 itself. The on-disk MLX checkout under `target/` already has the overlay applied, so all three kernels were fetched from GitHub at their own commits instead: `1700b39a` (the direct parent of #3792, confirmed via the commit's `parents`), `a5a684d` (#3792), and the pin. Measured on GB10 (sm_121, CUDA 13) with the three kernel bodies copied verbatim, each launched with its own commit's dispatch configuration and diffed against a per-row float64 reference: | Kernel | Failing axis sizes, 8 to 8192, row counts 1 to 17 | |---|---| | `1700b39` (the overlay mlxcel ships) | all 34 swept sizes in `[257, 512]` | | `a5a684d` (#3792) | none | | `b7c3dd6d` (#3850, the pin) | none | `compute-sanitizer --tool memcheck` reports `Invalid __shared__ read of size 4 bytes` for the overlay's launch geometry only; the other two are clean. ## Root cause At `axis_size` in `(256, 512]` for 16-bit dtypes, the pre-#3792 dispatch pairs `block_dim == 64` with `groups_per_block == 2`. That launches a 128-thread block, but the reduction is `BlockBroadcastReduce<float, 64, 32>`, whose `TempStorage` is sized `BLOCK_DIM / WARP_SIZE == 2` while `cg::tiled_partition<32>` of the launched block reports `meta_group_size() == 4`. The second reduction stage indexes two floats past the end of the scratch array and folds in the partials belonging to the other row that shares the block. Two failure modes follow. Above one row the sum covers both rows, so the normalizer is deterministically wrong (0.55 to 0.94 relative error). At one row the `threadIdx.y == 1` threads take the row-bounds early return and never write their partials, so the normalizer picks up residual shared memory. Seeding that window with 0, 1, 100 and 10000 moves the applied normalizer to 99.99%, 92.79%, 24.19% and 2.49% of the float64 reference. #3792 changed that one value to 1 and added `static_assert(block_dim <= 32 || groups_per_block() == 1)` to pin the invariant. It is a fix, and its title says so. #3850 rewrote the forward path to one row per block, where block size and template parameter agree by construction. lablup#830 got the band, the config, the changed value, and DeepSeek-V2's `kv_lora_rank == 512` landing on it all correct. Only the direction is backwards. ## Not reconciled lablup#829 measured deepseek-v2-lite going 0/4 to 4/4 coherent with the overlay as the only change. Against these measurements the overlay should have made that axis worse. The document says so plainly rather than papering over it, and notes two unverified leads: lablup#817's original "nondeterministic across fresh processes" symptom matches the uninitialized read exactly, and graph capture changes which kernels precede the norm, which is what determines the overlay's answer at one row. Settling it needs a build with the overlay removed, measured the way lablup#829 measured. That is a separate change with real risk and is deliberately not attempted here, so **no runtime behavior changes in this PR**. ## What changed - `docs/upstream/mlx-cuda-rmsnorm-small-axis-regression.md` (new): the report, following the two existing `docs/upstream/` examples. Summary, evidence, root cause with quoted source at every cited line, minimal fix, repro, affected surface. The standalone CUDA repro was run; the `mlx.core` sketch and the model-level repro are labelled not run. - `src/lib/mlx-cpp/patches/mlx/backend/cuda/rms_norm.cu`: the tracking comment no longer dangles on "report filed as a follow-up". It points at the document, marks the superseded paragraph, and carries the corrected removal condition: nothing needs to land upstream first, so the overlay can be deleted once deepseek-v2-lite is re-verified on the upstream fast kernel, recovering the #3850 speedup. Comment lines above the `---- begin upstream ----` marker only; the vendored kernel body is untouched. The patch-file sanity check at `CMakeLists.txt:115-126` globs `patches/mlx/backend/metal/*.cpp` and does not match this file. ## Follow-up worth filing The overlay is wrong on an axis nine model files reach through `kv_a_layernorm`. Dropping it and re-verifying deepseek-v2-lite is the obvious next step, but it needs a real model run and is out of scope for a documentation change. ## Test plan - [x] `cargo fmt --all -- --check` - [x] `cargo clippy -p mlxcel-core -- -D warnings` (clean, 16.23s) - [x] `cargo check -p mlxcel-core --features cuda` (clean, 29.20s; re-applied the overlay and recompiled `rms_norm.cu.o` through nvcc, so the comment edit is confirmed to build) - [x] Overlay diff is comment-only: zero non-`//` lines changed - [x] No em dashes and no AI attribution in the diff or the commit message - [ ] `cargo clippy -p mlxcel-core --features cuda` **skipped**: it hashes to a fresh build directory and starts a cold full MLX CUDA rebuild. Aborted after two minutes and the partial directory removed. No Rust changed in this PR, so it carried no signal. Closes lablup#830

lablup#1037) ## Summary Root-causes the DeepSeek-V2 "CUDA-graph-captured incremental decode collapse" from lablup#824/lablup#829/lablup#831: it was not a graph-capture hazard but mlxcel's own RMSNorm overlay (`patches/mlx/backend/cuda/rms_norm.cu`, upstream pre-#3792), whose two-stage reduction reads past its shared scratch on the exact `kv_a_layernorm` axis (`kv_lora_rank == 512` at bf16). Deletes the overlay, removes DeepSeek-V2 and DeepSeek-VL2 from `maybe_disable_cuda_graphs_for_model` so they run with CUDA graphs again, and pins the band with a numeric regression test. Gemma 4 stays on the lever: its norm axes never touch the broken dispatch geometry, so lablup#688 is a separate, still-open bug, and lablup#831's premise that the two families shared one hazard class is wrong. ## The decisive experiment 2x2 matrix on GB10 at `13741ae2`: deepseek-v2-lite-4bit, greedy `--temp 0`, 48 tokens, fresh process per rep, 6 reps per cell per prompt, two prompts (1-token "Hi", ~90-token summarization). "Identical" is by content hash of the generated text. | Cell | Kernel | `MLX_USE_CUDA_GRAPHS` | Short prompt | Long prompt | |---|---|---|---|---| | A | overlay `1700b39` | 0 | 6/6 coherent, identical | 6/6 coherent, identical | | B | overlay `1700b39` | 1 | 4/6 coherent, all four token-divergent; 2/6 collapse into repeated `!` | 0/6: every rep collapses after one generated token | | C | pinned `b7c3dd6d` | 0 | 6/6 coherent, identical | 6/6 coherent, identical | | D | pinned `b7c3dd6d` | 1 | 6/6 coherent, identical, byte-identical to C | 6/6 coherent, identical, byte-identical to C | Cells C and D being byte-identical means graph capture has no observable effect on this family once the kernel is correct. The overlay's single-row failure mode (normalizer depends on residual shared memory) explains every recorded symptom: greedy flakiness, coherent-prefix-then-collapse, prompt-length dependence, and `MLXCEL_FORCE_SYNC=1` not helping. The cell C/D binary was proven to contain the pinned kernel (the `configure_file(COPYONLY)` overlay mechanism makes a silent no-op rebuild possible): post-build source diff against `b7c3dd6d`, `rms_norm.cu.o` recompile mtime, and `cuobjdump`/`strings` mangling showing 128 `rms_norm_small` instantiations with #3850's `bool ALIGNED` parameter (`Lb0E`/`Lb1E`) and zero of the overlay's four-int form. ## What changed - `src/lib/mlx-cpp/patches/mlx/backend/cuda/rms_norm.cu`: deleted. The pinned kernel is correct on the affected axis and carries the #3850 forward speedup. - `src/loading/mod.rs`: `ModelType::DeepSeekV2 | ModelType::DeepSeekVL2` removed from the hazard match; function docs rewritten to record the corrected attribution. Gemma 4 arms unchanged. - `src/server/startup.rs`: main-thread hoist comment updated to the Gemma-4-only story. - `src/lib/mlxcel-core/src/rms_norm_small_axis_tests.rs` (new, registered in `lib.rs`): pins `fast::rms_norm` against a float64 host reference across the broken band, 16-bit `(256, 512]` and f32 `(128, 256]`, rows 8/3/2/1. Verified non-vacuous: with the overlay kernel linked back in, both tests fail at ~3e-1 worst-row relative error; on the pin they pass at rounding scale. - `docs/upstream/mlx-cuda-rmsnorm-small-axis-regression.md`: resolution section with the matrix, the binary proof, the non-reproduction of lablup#829's 0/4 baseline, and the local checkpoints that sat in the band (five MLA models at `kv_lora_rank == 512` plus falcon-h1-tiny whose `hidden_size` itself is 512). ## Test plan - [x] `cargo fmt --all -- --check` clean - [x] `cargo clippy --release --features cuda -p mlxcel-core --lib --tests` and `-p mlxcel --bins`: no lints - [x] `cargo test --release --features cuda -p mlxcel-core --lib rms_norm_small_axis`: 2/2 (and 0/2 against the overlay kernel, restored afterwards) - [x] `cargo test --release --features cuda -p mlxcel-core --lib fused_norm_parity`: 8/8 - [x] `cargo test --release --features cuda --test deepseek_v2_greedy_parity`: 2/2, token-exact vs mlx-lm 0.31.3 with graphs now on by default - [x] Smoke: deepseek-v2-lite-4bit coherent with no env override (no disable log), deepseek-vl2-small-4bit coherent, gemma-4-e2b-it-4bit still logs the lablup#688 graph disable - [x] 200-token decode: graphs-on 85.0-93.3 tok/s vs graphs-off 86.5-88.1 (n=3 each; GB10 single-run decode is bimodal, coherence is the gate) Closes lablup#831

lablup#1046) ## Summary Bumps the MLX pin from `b7c3dd6d` (2026-07-17) to `2c46b953` (2026-08-05 upstream main), 58 commits, and re-derives the three source overlays that upstream touched in that window. Closes lablup#1042 ## Why Fixes in range that land on paths this repo ships: | Upstream | What | |---|---| | #3854 | incorrect nvfp4 `quantized_matmul` through the split-K path (NVFP4 direct transcode is mlxcel's default on all builds) | | #3893 | CUDA gemm conv unfold grid overflow past 65,535 output positions | | #3934 | Metal NVFP4 scale reduction per 16-lane group | | #3888 | Metal `gemv_wide` for few-row fp16/bf16 matmuls | | #3943 | head dimension 96 in Metal full attention | | #3898, #3813 | `prod` dtype promotion on a size-1 axis, shapeless matmul with dynamic batch dims | ## Overlay audit mlxcel applies whole-file overlays via `configure_file(... COPYONLY)`, so an overlay whose upstream counterpart moved silently reverts that upstream change. That is how a broken RMSNorm kernel shipped in lablup#824 / lablup#829 and took until lablup#830 / lablup#831 to undo. All 21 overlays were diffed against the upstream range; three conflicted. | Overlay | Upstream change | Resolution | |---|---|---| | `patches/mlx/backend/cuda/quantized/quantized.cpp` | #3757 inserted `const std::optional<array>& global_scale` into `qmm_naive` at position 5 | resolved by hand: mlxcel wraps both call sites in `run_row_chunked` / `run_batch_chunked`, which `git merge-file` cannot reconcile with an inline insertion. Both sites verified as 13 arguments in declaration order | | `patches/mlx/backend/metal/kernels/utils.h` | #3970 templated complex scalar lanes, #3963 explicit `thread` address space for metal 4.1 | clean three-way merge, no overlap with mlxcel's delta | | `patches-cuda/ops.cpp` | 189 lines | clean three-way merge, mlxcel's 57-line delta intact | `quantized/qmm/qmm.h` also changed but is not overlaid, so mlxcel takes the new declaration. Its full diff across the range is the single `global_scale` line, so `qmm_sm90`, `qmm_sm80`, `qmv`, `gather_qmv` and `fp_qmv` are unaffected. The `qmm_sm80.cuh` race-fix overlay from lablup#910 is untouched upstream and still required. The stale `Synced to upstream e9463bb` header comment is updated to the real base, with an explicit instruction to re-verify these call sites on the next bump. ## The pin has two sources of truth `GIT_TAG` in `src/lib/mlx-cpp/CMakeLists.txt:95` and `MLX_EXPECTED_COMMIT` in `src/lib/mlxcel-core/build.rs:217`. The latter drives `purge_stale_mlx_cache`. Updating only the CMakeLists leaves the cache marker matching `MLX_EXPECTED_COMMIT`, so `_deps/` is judged valid, never purged, and the build silently compiles the OLD MLX while looking bumped. Both moved together, verified three ways with no trace of the old SHA: - binary embedded commit: `2c46b953…` - `_deps/.mlx-build-commit` marker: `2c46b953…` - fetched `_deps/mlx-src` HEAD: `2c46b953…` Worth a follow-up to collapse this to one source of truth or add a build-time assert. ## Validation (GB10, sm_121) - release CUDA build clean - `cargo fmt --all -- --check` clean - `fused_norm_parity` 8/8, `turbo_tests` 97/97 - greedy smoke on `llama-3.2-1b-4bit`, `qwen3-4b-4bit`, `deepseek-v2-lite-4bit`: coherent, and byte-identical run to run with a warm kernel cache. That determinism also confirms the `qmm_sm80` race-fix overlay still holds on the new base. ## Two pre-existing failures, separated by a control build A control build at the old pin (`_deps` re-fetched to `b7c3dd6d`, verified) reproduces both, so neither is caused by this bump. | Failure | New pin | Old pin (control) | |---|---|---| | `fused_rope_parity` 2 of 9 | offset 131069 rms 2.497e-3, offset 131071 rms 2.684e-3, tol 2.0e-3 | same digits, same 7 pass / 2 fail | | full lib suite abort under graph capture | `cuGraphAddKernelNode` invalid argument, then `cudaStreamEndCapture` previous error during capture (site and message vary per run) | aborts the same way | For the rope case the numbers match bit for bit, which also shows that path carries no exposure to this change. Upstream touched no rope file in the range. Both are left for separate issues rather than folded in here.

inureyes added type:bug Bug fixes, error corrections, or issue resolutions priority:high High priority area:models Model architectures, weights, loading, metadata status:review Under review labels Jul 20, 2026

inureyes added 2 commits July 20, 2026 09:50

inureyes added status:done Completed and removed status:review Under review labels Jul 20, 2026

inureyes merged commit a9cc0c3 into main Jul 20, 2026
7 checks passed

inureyes deleted the fix/issue-824-mlx-bump-deepseek-v2 branch July 20, 2026 02:16

This was referenced Aug 2, 2026

fix(models): deepseek_v2 generates degenerate output on a real checkpoint while mlx-lm is correct on the same weights #991

Closed

feat(core): matrix-absorbed MLA decode over a compressed-latent KV cache #995

Merged

inureyes mentioned this pull request Aug 5, 2026

docs(upstream): correct the CUDA RMSNorm small-axis attribution #1033

Merged

5 of 6 tasks

inureyes mentioned this pull request Aug 5, 2026

fix(cuda): drop the broken RMSNorm overlay, restore DeepSeek-V2 graphs #1037

Merged

7 tasks done

inureyes self-assigned this Aug 31, 2026

inureyes mentioned this pull request Sep 19, 2026

chore(deps): bump the MLX pin and declare a base commit for each overlay #1932

Closed

9 tasks

inureyes mentioned this pull request Oct 6, 2026

fix(cuda): Qwen3 greedy decoding is not reproducible run to run on GB10 #2128

Closed

4 tasks

Navigation Menu

Repository navigation

fix(cuda): restore DeepSeek-V2 output by reverting the regressed RMSNorm kernel - #829

fix(cuda): restore DeepSeek-V2 output by reverting the regressed RMSNorm kernel#829
inureyes merged 4 commits into
mainfrom
fix/issue-824-mlx-bump-deepseek-v2

inureyes commented Jul 20, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Conversation

inureyes commented Jul 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Root cause (pinned by a deterministic bisect)

Fix

Verification (GB10)

Scope / follow-ups

Update: second commit adds the CUDA-graph-capture fix (required for the server)

Uh oh!

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

inureyes commented Jul 20, 2026 •

edited

Loading