fix(cuda): restore lfm2-350m-8bit decode via elementwise short conv#751
Merged
Merged
Conversation
lfm2-350m-8bit decode on GB10 regressed ~10x between 0.3.1 and 0.4.0-rc.1 (409.01 -> ~40 tok/s) while prefill stayed healthy and even improved. The sibling lfm2-8b-a1b-4bit was unaffected. nsys localized the cause: on the single-step (L=1) decode path, MLX 0.32.1 dispatches the tiny bf16 depthwise short conv on CUDA to cuDNN's generic convolve_common_engine_float_NHWC, which launches one kernel per channel (~1024 for a 350M LFM2). That kernel accounted for 88.6% of decode GPU time (286,720 instances across the run, roughly 1024 per conv1d op). The 8b sibling runs f32 activations and its conv stays on the fast conv1d_c1_k1_nhwc kernel, so it never hit the slow path; the 350m runs bf16 and does. Prefill uses the same fast kernel because its sequence length is > 1, which is why only decode regressed. MLXCEL_QMV_MULTIROW=0 and the CUDA-graph toggle were both ruled out first (graphs-on is already fastest). The fix computes the single decode step of the depthwise causal short conv as a broadcast multiply-and-sum over the L_cache taps instead of calling conv1d: out[b,0,c] = sum_k padded[b,k,c] * weight[c,k,0], which is exactly what a stride-1/no-pad/dilation-1/groups==hidden conv1d produces for a length-1 output. The conv weight is transposed once at load to a time-major [1, L_cache, hidden] tensor so the decode step is two elementwise kernels (a broadcast multiply plus an axis-sum that accumulates in f32) rather than a grouped-conv dispatch. Prefill (out_len > 1) and Metal keep the proven conv1d path; the fast path is gated on !metal_is_available() and only fires for L=1. After the fix, nsys shows convolve_common_engine_float_NHWC gone entirely and decode correctly qmv-bound (75.3%). Decode returns to the 0.3.1 envelope: 393.84 tok/s via ./scripts/bench_decode.sh (evidence CSV added), 380-394 tok/s across repeats, with byte-identical greedy output (same 13-token early-EOS response). Controls unregressed: lfm2-8b-a1b-4bit 161.39, qwen2.5-7b-8bit 30.59 (issue baseline 28.96/30.18), gemma-4-e2b-it-8bit 58.49. A checkpoint-free unit test (lfm2_short_conv_decode_matches_conv1d) asserts the elementwise step equals conv1d. The gemma-4-26b/glm4-flash/mamma2-130m moderate drops noted in the sweep are a milder expression of the same conv dispatch on other short-conv/SSM families and are left for a separate follow-up. Closes #748
inureyes
added
type:bug
Reword the ShortConv fast-path comment to state precisely what gates it: any single-position output step (a true decode step, or a 1-token prefill) takes the elementwise path, while multi-position outputs (prefill, speculative verify) and Metal keep `conv1d`. The prior wording implied the split was strictly decode-vs-prefill, which is loose since a 1-token prefill also correctly takes the fast path. Add `lfm2_short_conv_decode_matches_conv1d_bf16`, a bf16 variant of the existing checkpoint-free equivalence test. It builds the same asymmetric kernel in f32, casts both the `conv1d` reference and the fast-path inputs to bf16 (matching how the decode weight is built from a bf16 checkpoint's conv weight), and compares in f32 with a loose tolerance. This is defense-in-depth against a future MLX change to how `sum_axis` accumulates for half dtypes, since the f32 test alone cannot see rounding differences that only appear in bf16 arithmetic. Validation: - cargo test --lib models::lfm2 --features cuda (11 passed, incl. new bf16 test) - cargo fmt on both touched files (no additional changes) - cargo clippy --lib --tests --features cuda -- -D warnings (clean) Refs #748
inureyes
added a commit
that referenced
this pull request
Jul 12, 2026
…nv dispatch across SSM/hybrid decode paths (#753) ## Summary Follow-up to #748 / PR #751 (`3fa2789`). That fix found MLX 0.32.1's CUDA backend dispatches a single-output-position (L=1) bf16 depthwise `conv1d` to cuDNN's generic per-channel engine (`convolve_common_engine_float_NHWC`, one launch per channel), which regressed LFM2 decode ~10x, and replaced the L=1 step with a broadcast multiply plus axis sum. This PR does two things: (1) lifts the #751 decode helpers into a shared module so any family can adopt them in one line, and (2) measures every other decode-path depthwise-conv family on GB10 to see whether the same slow dispatch recurs. It does not. Every SSM / hybrid family already runs its L=1 decode conv on the fast `conv1d_c1_k1_nhwc` cuDNN engine, so none is adapted. The result is a pure refactor plus a documented, measurement-driven negative result. ## What changed - `src/models/conv_decode.rs` (new): shared `build_conv_decode_weight` + `short_conv_decode_step`, lifted verbatim from `lfm2.rs`. `src/models/lfm2.rs` re-exports them at the original path (`pub(crate) use`), so `ShortConv` and the existing 11 lfm2 tests are unchanged. Pure move, behavior identical (lfm2-350m-8bit greedy output byte-identical: "Hello! I'm doing well, thanks. How about you?", 13 tokens, and decode still elementwise at 393-ish tok/s). - `src/models/conv_decode_tests.rs` (new): generalized channel- and tap-asymmetric parity tests (f32 + bf16, kernel widths 3 and 4, plus a `silu(bias + conv)` end-to-end variant) pinning the elementwise decode step against `conv1d` without a checkpoint. - `src/models/mod.rs`: wire `conv_decode` under the shared modules and `conv_decode_tests` under the test modules. - `docs/benchmark_results/model_tests_gb10.md`: record the audit outcome, annotate the mamba2-130m row, and correct the PR #751 hypothesis on the gemma-4 / glm4-flash drops. ## Method Each family with a local checkpoint was profiled under `nsys` on a warm GB10 CUDA decode (`mlxcel generate ... -n 24 --profile`, `MLX_USE_CUDA_GRAPHS=0` so the cuDNN engine choice is captured eagerly and the histogram is complete), then `nsys stats --report cuda_gpu_kern_sum` was grepped for conv kernels. The cuDNN engine choice is keyed on shape/dtype/alignment/groups and is independent of MLX graph capture, so graphs-off is representative of production. The discriminator is the launch count, not just the kernel name. The fast `conv1d_c1_k1_nhwc` engine launches one kernel per conv op; the slow `convolve_common_engine` launches one per channel (LFM2's 350M hit 286,720 launches, ~1024 per op). Every family measures at (conv layers x decode tokens) launches, i.e. one per op, which is direct positive proof of the fast engine rather than mere absence of the slow name. ## Per-family verdict No family was CONFIRMED, so there are no before/after tok/s (nothing was adapted). Kernel counts are from the warm 24-token decode. | Family | Checkpoint | Act. dtype | Conv engine | `conv1d_c1_k1_nhwc` launches | `convolve_common_engine` | conv % of decode | Verdict | |---|---|---|---|---:|---:|---:|---| | mamba2 | mamba2-130m | bf16 | fast | 576 (24 layers x 24 tok) | 0 | 1.6% | REFUTED | | mamba2 | mamba2-1.3b-4bit | fp16 | fast | 1152 | 0 | 1.7% | REFUTED | | mamba | falcon-mamba-7b-4bit | fp16 | fast | 256 (early EOS, 2 tok) | 0 | 0.9% | REFUTED | | falcon-h1 | falcon-h1-tiny-90m-instruct-4bit | bf16 | fast | 576 (24 x 24) | 0 | 3.0% | REFUTED | | granite-4.0-h | granite-4.0-h-350m-4bit | bf16 | fast | 672 | 0 | 2.1% | REFUTED | | jamba | jamba-v0.1-4bit | bf16 | fast | 624 | 0 | 1.0% | REFUTED | | plamo2 | plamo-2-1b | f32 | fast | 192 | 0 | 0.2% | REFUTED | | qwen3.5 / gated_delta | qwen3.5-0.8b-4bit | bf16 | fast | 432 | 0 | 1.9% | REFUTED | | nemotron-h | nemotron-h-30b-4bit | bf16 | fast | 552 (23 layers x 24 tok) | 0 | 0.7% | REFUTED | | kimi_linear (ShortConv1d) | kimi-vl-a3b-thinking-4bit | n/a | n/a | 0 (no conv runs) | 0 | n/a | NOT MEASURED | | recurrent_gemma | none on disk | n/a | n/a | n/a | n/a | n/a | NOT MEASURED | | qwen3-next | none on disk | n/a | n/a | n/a | n/a | n/a | NOT MEASURED | Notes on the NOT MEASURED rows: - kimi_linear: the listed `kimi-vl-a3b-thinking-4bit` is a `deepseek_v3` text backbone (MLA attention, no conv); it never exercises `kimi_linear.rs`'s `ShortConv1d`, so nsys shows no conv kernel at all. The actual Kimi-Linear architecture has no local checkpoint. Its `ShortConv1d` (conv1d -> silu, no bias) would fit the shared helper if a checkpoint appears. - recurrent_gemma: no local checkpoint, and its conv is structurally different (it transposes to `[B, C, L]`, convolves, transposes back), so it would need a small local adaptation rather than the shared helper. Documented, not adapted. - qwen3-next: no local checkpoint. It shares the gated-delta conv that qwen3.5-0.8b-4bit exercises and which is REFUTED above. ## Why LFM2 regressed but these do not LFM2's slow dispatch is specific to its `conv_L_cache = 3` (kernel width 3) at hidden 1024. cuDNN's engine heuristic selects the fast NHWC engine for the SSM families' `conv_kernel = 4` shapes and the slow generic per-channel engine only for LFM2's shape. The conv share of decode is 0.2-3.0% everywhere here, so even if a shape did tip over, the ceiling is small; adapting a refuted family would add the elementwise path for no benefit (and could be slower than the fast cuDNN kernel), so none is changed. ## PR #751 hypothesis correction (on the record) PR #751's body speculated that the gemma-4-26b-a4b-it(-qat) (-10 to -14%) and glm4-flash-4bit (-14%) sweep drops were the same conv-dispatch pattern. Code inspection refutes that: neither is a conv architecture (`gemma4` MoE and `glm4_moe_lite`; neither file calls `conv1d`), so their decode conv cannot fall into the grouped-conv engine. Those drops need separate attribution and are out of scope here. Similarly, mamba2-130m's -10.4% sweep delta is not conv dispatch (its conv is on the fast engine at 1.6% of decode); it is environmental sweep variance, as the issue anticipated. ## Test plan - [x] `cargo test --lib --features cuda models::conv_decode` (5 passed: f32/bf16 kernel-3 and kernel-4 parity + silu/bias variant) - [x] `cargo test --lib --features cuda models::lfm2` (11 passed, unmodified) - [x] `cargo clippy -p mlxcel --lib --tests --features cuda` (clean; only pre-existing mlxcel-core C++ build notes) - [x] `cargo fmt` (no changes) - [x] `cargo build --release --features cuda` then nsys audit of all 10 candidate checkpoints on GB10 - [x] lfm2-350m-8bit greedy output byte-identical after the refactor Closes #752
5 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
lfm2-350m-8bit decode on GB10 regressed ~10x between 0.3.1 and 0.4.0-rc.1 (409.01 -> ~40 tok/s). This computes the single-step depthwise short conv as a broadcast multiply-and-sum instead of a tiny bf16
conv1d, which MLX 0.32.1 routes to cuDNN's per-channel grouped-conv engine on CUDA. Decode returns to the 0.3.1 envelope (393.84 tok/s) with byte-identical greedy output.Root cause
Prefill was healthy and even improved, and the sibling
lfm2-8b-a1b-4bitwas never affected, so the regression is purely per-token decode latency on this one model.MLXCEL_QMV_MULTIROW=0(#740) changed nothing and the CUDA-graph toggle was not the lever (graphs-on is already fastest), so both were ruled out first.nsyson the warm decode localized it decisively. The single-step (L=1) bf16 depthwise short conv on CUDA dispatches to cuDNN's genericconvolve_common_engine_float_NHWC, which launches one kernel per channel (~1024 for a 350M LFM2). That kernel was 88.6% of decode GPU time: 286,720 instances across the run, roughly 1024 perconv1dop (10 conv layers x 28 tokens x 1024). The 8b sibling runs f32 activations and its conv stays on the fast customconv1d_c1_k1_nhwckernel (1 launch per op), which is why it was healthy; the 350m runs bf16 and falls into the slow engine. Prefill uses the same fast kernel because its sequence length is > 1, which is why only decode regressed. This is an MLX 0.32.1 CUDA conv-dispatch behavior, mitigated on the mlxcel side for the affected shape/dtype.What changed
src/models/lfm2.rs:ShortConvprecomputes a time-major[1, L_cache, hidden]decode weight at load (off Metal) and, for the L=1 decode step, computesout[b,0,c] = sum_k padded[b,k,c] * weight[c,k,0]as a broadcast multiply plus an axis sum (MLX widens the reduce accumulator for half dtypes), exactly what a stride-1/no-pad/dilation-1/groups==hiddenconv1dproduces for a length-1 output. Prefill (out_len > 1) and Metal keep the provenconv1dpath; the fast path is gated on!metal_is_available().src/models/lfm2_tests.rs:lfm2_short_conv_decode_matches_conv1d, a checkpoint-free unit test asserting the elementwise decode step is numerically identical toconv1d.docs/benchmark_results/model_tests_gb10.md: refresh thelfm2-350m-8bitrow (39.84 -> 393.84) and rewrite the regression note to reference the fix and root cause.benchmarks/cuda_gb10_2026-07-12_single_lfm2-350m-8bit_fix748.csv: post-fix evidence CSV from./scripts/bench_decode.sh(written to a distinct filename so the committed regression-evidence CSV is preserved).Measured (GB10, CUDA, release)
Greedy output is unchanged (same 13-token early-EOS response: "Hello! I'm doing well, thanks. How about you?"). Post-fix
nsysshows decode correctly qmv-bound (75.3%).The moderate sweep drops on gemma-4-26b-a4b-it(-qat), glm4-flash-4bit, and mamba2-130m are a milder expression of the same small-depthwise-conv dispatch on other short-conv/SSM families and are left for a separate follow-up.
Test plan
cargo test --lib --features cuda models::lfm2_tests(10 passed, incl.lfm2_short_conv_decode_matches_conv1d)cargo clippy --lib --tests --features cuda -- -D warnings(clean)./scripts/bench_decode.sh models/lfm2-350m-8bit --cooldown 0-> 393.84 tok/snsysbefore/after:convolve_common_engine_float_NHWCeliminatedCloses #748