fix(cuda): restore lfm2-350m-8bit decode via elementwise short conv

benchmarks/cuda_gb10_2026-07-12_single_lfm2-350m-8bit_fix748.csv

@@ -0,0 +1,2 @@
+model,model_path,prompt_tokens,generated_tokens,prefill_ms,prefill_tok_s,decode_ms,decode_tok_s,date,hardware,mlx_version,build_type,max_tokens,prompt,prompt_target_len
+lfm2-350m-8bit,models/lfm2-350m-8bit,17,13,5.20,3270.66,33.01,393.84,2026-07-12,NVIDIA_GB10_122GB,0.4.0-rc.1,release,100,"Hello, how are you today?",

docs/benchmark_results/model_tests_gb10.md

@@ -177,7 +177,7 @@ Prefill/Decode are the measured-pass figures from `mlxcel-bench-decode`. Notes r
 | falcon-h1-tiny-90m-instruct-4bit | ✅ | 1419.44 | 413.00 | 30 tok |
 | falcon-mamba-7b-4bit | ✅ | 92.54 | 22.06 | 2 tok |
 | jamba-v0.1-4bit | ✅ | 523.67 | 89.63 |  |
-| lfm2-350m-8bit | ✅ | 3031.35 | 39.84 | 13 tok; decode regression vs 0.3.1 (409.01 -> ~40 tok/s, reproducible); tracked in #748 |
+| lfm2-350m-8bit | ✅ | 3270.66 | 393.84 | 13 tok; decode regression fixed (#748): was ~40 tok/s, restored to the 0.3.1 envelope by computing the single-step depthwise short conv as a broadcast multiply-sum instead of a tiny bf16 `conv1d` (MLX 0.32.1 routed that to cuDNN's per-channel grouped-conv engine on CUDA) |
 | mamba2-130m | ✅ | 890.29 | 162.23 |  |
 | mamba2-1.3b-4bit | ✅ | 329.13 | 83.40 |  |
 | plamo-2-1b | ✅ | 199.12 | 44.54 |  |
@@ -373,7 +373,7 @@ Models that accept image input and generated tokens under the `"What is in this
 ### Notable changes vs the 2026-06-17 (0.3.1) and 2026-07-09 (rc.1 subset) records
 
 - **SSM / hybrid / NAS decode reads 2-3x higher than every earlier record**: granite-4.0-h-350m 86.60 → 259.69, granite-4.0-h-tiny 33.84 → 100.28, falcon-h1-tiny 110.42 → 413.00, nemotron-h-30b 40.32 → 79.94, nemotron-nas-30b 37.33 → 82.72, nemotron-omni-30b 38.45 → 80.86, plamo-2-1b 35.14 → 44.54. Two of these (granite-350m, falcon-h1-tiny) were re-measured as recently as 2026-07-09 at the low values, and no SSM-related code has landed since, so the delta is environmental rather than a code change. The affected cluster is exactly the launch-latency-sensitive family. Re-verify after the planned pre-release reboot before treating these as the release numbers.
-- **lfm2-350m-8bit decode regressed ~10x**: 409.01 (0.3.1) → 39.84, and a standalone re-run reproduces it (42.47), so it is not sweep noise. The sibling `lfm2-8b-a1b-4bit` is unaffected (157.73 → 161.87). Tracked in #748.
+- **lfm2-350m-8bit decode regressed ~10x, now fixed (#748)**: 409.01 (0.3.1) → 39.84 at rc.1, restored to 393.84 by the fix. Root cause: MLX 0.32.1 dispatches the single-step (L=1) bf16 depthwise short conv on CUDA to cuDNN's generic `convolve_common_engine`, which launches one kernel per channel (~1024 for a 350M LFM2) and consumed 88.6% of decode GPU time; computing that decode step as a broadcast multiply-and-sum over the `L_cache` taps removes the grouped-conv dispatch. The sibling `lfm2-8b-a1b-4bit` was never affected (its conv runs on the fast `conv1d_c1_k1_nhwc` kernel; 157.73 → 161.87 → 161.39). Prefill still uses `conv1d` and was always healthy.
 - **Moderate decode drops worth watching**: gemma-4-26b-a4b-it-4bit 58.59 → 50.19 (-14%), gemma-4-26b-a4b-it-qat-4bit 50.33 → 45.29 (-10%), glm4-flash-4bit 53.33 → 45.72 (-14%).
 - **Text-only prefill for several VLM-capable models is an order of magnitude higher than 0.3.1** (aya-vision-8b 124.91 → 1354.53, pixtral-12b 35.97 → 120.39, youtu-vl 134.27 → 451.42, mistral-small-3.1-24b 63.68 → 891.89), consistent with the text-prompt path no longer paying vision-tower costs rather than with a kernel speedup.
 - **Decode improvements >10% on dense/MoE models**: apertus-8b +21%, gemma3-4b / gemma-3-4b-it +19%, gemma3-1b +15%, gemma3n-e2b +10%, qwen3-0.6b-4bit +13%, aya-expanse-8b +12%, plus the VLM-side gains below.

src/models/lfm2.rs

@@ -306,6 +306,50 @@ struct ShortConv {
     conv_bias: Option<UniquePtr<MlxArray>>,
     hidden_size: i32,
     l_cache: i32,
+    /// Time-major depthwise-conv weight for the single-step (decode) fast path
+    /// (issue #748), shape `[1, L_cache, hidden]` where
+    /// `decode_weight[0, k, c] == conv_weight[c, k, 0]`. `None` on Metal, where
+    /// MLX's `conv1d` already dispatches a fast small-conv kernel and the proven
+    /// path is kept as-is.
+    decode_weight: Option<UniquePtr<MlxArray>>,
+}
+
+/// Materialize the time-major weight used by the decode fast path.
+/// `conv_weight` is `[hidden, L_cache, 1]` (MLX depthwise layout); transposing
+/// to `[1, L_cache, hidden]` lets a single broadcast multiply-and-sum over the
+/// `L_cache` axis replace `conv1d`. Materialized once at load so decode never
+/// reshapes the weight.
+pub(crate) fn build_conv_decode_weight(conv_weight: &MlxArray) -> UniquePtr<MlxArray> {
+    // [hidden, L_cache, 1] -> [1, L_cache, hidden].
+    let w = mlxcel_core::transpose_axes(conv_weight, &[2, 1, 0]);
+    let w = mlxcel_core::contiguous(&w, false);
+    mlxcel_core::eval(&w);
+    w
+}
+
+/// Single decode step of the depthwise causal short conv, computed as a
+/// broadcast weighted sum instead of `conv1d` (issue #748).
+///
+/// For `padded` of shape `[batch, L_cache, hidden]` and `decode_weight` of
+/// shape `[1, L_cache, hidden]` this returns `[batch, 1, hidden]` where
+/// `out[b, 0, c] = sum_k padded[b, k, c] * decode_weight[0, k, c]`, which is
+/// exactly what a stride-1, no-pad, dilation-1, `groups == hidden` `conv1d`
+/// produces for a length-1 output. The two-kernel elementwise form avoids the
+/// `conv1d` CUDA dispatch (MLX 0.32.1) that sends this tiny bf16 depthwise conv
+/// to cuDNN's generic `convolve_common_engine`, which launches one kernel per
+/// channel (~1024 for a 350M LFM2) and dominates decode.
+pub(crate) fn short_conv_decode_step(
+    padded: &MlxArray,
+    decode_weight: &MlxArray,
+    in_dtype: i32,
+) -> UniquePtr<MlxArray> {
+    let prod = if mlxcel_core::array_dtype(decode_weight) == in_dtype {
+        mlxcel_core::multiply(padded, decode_weight)
+    } else {
+        let w = mlxcel_core::astype(decode_weight, in_dtype);
+        mlxcel_core::multiply(padded, &w)
+    };
+    mlxcel_core::sum_axis(&prod, 1, true) // [batch, 1, hidden]
 }
 
 impl ShortConv {
@@ -321,13 +365,23 @@ impl ShortConv {
             .get(&format!("{prefix}.conv.bias"))
             .map(|w| mlxcel_core::copy(w));
 
+        // Precompute the decode fast-path weight off Metal (issue #748). On
+        // Metal `conv1d` already dispatches a fast small-conv kernel, so keep
+        // its proven path and leave `decode_weight` unset there.
+        let decode_weight = if mlxcel_core::metal_is_available() {
+            None
+        } else {
+            Some(build_conv_decode_weight(&conv_weight))
+        };
+
         Ok(Self {
             in_proj,
             out_proj,
             conv_weight,
             conv_bias,
             hidden_size: args.hidden_size as i32,
             l_cache: args.conv_l_cache as i32,
+            decode_weight,
         })
     }
 
@@ -375,8 +429,18 @@ impl ShortConv {
         // Depthwise Conv1d (groups == hidden). Match the conv weight dtype to
         // the (possibly bf16) activations; quantized checkpoints leave the
         // non-quantized conv weight at its stored precision.
+        //
+        // Off Metal, any single-position output step (a decode step, or a
+        // 1-token prefill) computes the conv as an explicit weighted sum of
+        // the L_cache taps instead of calling `conv1d`: a tiny bf16 depthwise
+        // conv on CUDA (MLX 0.32.1) otherwise falls into cuDNN's generic
+        // `convolve_common_engine`, which launches one kernel per channel and
+        // regressed lfm2-350m-8bit decode ~10x (issue #748). Multi-position
+        // outputs (prefill, speculative verify) and Metal keep `conv1d`.
         let in_dtype = mlxcel_core::array_dtype(&padded);
-        let conv_out = if mlxcel_core::array_dtype(&self.conv_weight) == in_dtype {
+        let conv_out = if let (Some(decode_weight), 1) = (&self.decode_weight, bx_shape[1]) {
+            short_conv_decode_step(&padded, decode_weight, in_dtype)
+        } else if mlxcel_core::array_dtype(&self.conv_weight) == in_dtype {
             mlxcel_core::conv1d(&padded, &self.conv_weight, 1, 0, 1, h)
         } else {
             let w = mlxcel_core::astype(&self.conv_weight, in_dtype);

src/models/lfm2_tests.rs

@@ -197,6 +197,120 @@ fn lfm2_eos_token_id_handles_scalar_array_and_missing() {
     assert_eq!(missing.eos_token_ids(), vec![7]);
 }
 
+#[test]
+fn lfm2_short_conv_decode_matches_conv1d() {
+    // Guard the decode fast path (issue #748): the single-step depthwise short
+    // conv, computed as an explicit weighted sum of the L_cache taps, must be
+    // numerically identical to the stride-1/no-pad/depthwise `conv1d` it
+    // replaces. This is the checkpoint-free core of the regression fix; the
+    // CUDA kernel-dispatch win is measured end-to-end against real checkpoints.
+    use super::lfm2::{build_conv_decode_weight, short_conv_decode_step};
+
+    let hidden = 4;
+    let l_cache = 3;
+
+    // conv_weight is [hidden, L_cache, 1] (MLX depthwise layout): for each
+    // channel c, the three causal taps weight[c, 0..3, 0].
+    let weight_data: Vec<f32> = vec![
+        0.5, -0.25, 0.75, // channel 0
+        -1.0, 0.5, 0.25, // channel 1
+        0.1, 0.2, -0.3, // channel 2
+        2.0, -0.5, 1.5, // channel 3
+    ];
+    let conv_weight = mlxcel_core::from_slice_f32(&weight_data, &[hidden, l_cache, 1]);
+
+    // padded is [1, L_cache, hidden]: the cached conv-state tail prepended to
+    // the current Bx step, one row per time index.
+    let padded_data: Vec<f32> = vec![
+        1.0, 2.0, 3.0, 4.0, // t = 0
+        -1.0, 0.5, 2.0, -2.0, // t = 1
+        0.25, -0.75, 1.0, 0.5, // t = 2
+    ];
+    let padded = mlxcel_core::from_slice_f32(&padded_data, &[1, l_cache, hidden]);
+
+    // Ground truth: stride-1, no-pad, dilation-1, groups==hidden conv1d.
+    let reference = mlxcel_core::conv1d(&padded, &conv_weight, 1, 0, 1, hidden);
+    assert_eq!(mlxcel_core::array_shape(&reference), vec![1, 1, hidden]);
+
+    // Fast path: broadcast weighted sum over the precomputed time-major weight.
+    let decode_weight = build_conv_decode_weight(&conv_weight);
+    assert_eq!(
+        mlxcel_core::array_shape(&decode_weight),
+        vec![1, l_cache, hidden]
+    );
+    let elementwise = short_conv_decode_step(&padded, &decode_weight, mlxcel_core::dtype::FLOAT32);
+    assert_eq!(mlxcel_core::array_shape(&elementwise), vec![1, 1, hidden]);
+
+    let diff = mlxcel_core::subtract(&reference, &elementwise);
+    let max_abs = mlxcel_core::item_f32(&mlxcel_core::max_all(&mlxcel_core::abs(&diff)));
+    assert!(
+        max_abs < 1e-5,
+        "decode short-conv diverged from conv1d: max|diff| = {max_abs}"
+    );
+}
+
+#[test]
+fn lfm2_short_conv_decode_matches_conv1d_bf16() {
+    // bf16 variant of `lfm2_short_conv_decode_matches_conv1d`: this is the
+    // dtype the decode fast path actually runs in on real (bf16) LFM2
+    // checkpoints, and is defense-in-depth against a future MLX change to
+    // how `sum_axis` accumulates for half dtypes. Values are built in f32
+    // (same asymmetric kernel as the f32 test) and cast to bf16 for both
+    // the `conv1d` reference and the fast path, so the comparison isolates
+    // dtype-driven rounding rather than construction differences. bf16 has
+    // ~3 decimal digits of precision, so the tolerance is loose relative to
+    // the f32 test.
+    use super::lfm2::{build_conv_decode_weight, short_conv_decode_step};
+    use mlxcel_core::dtype;
+
+    let hidden = 4;
+    let l_cache = 3;
+
+    let weight_data: Vec<f32> = vec![
+        0.5, -0.25, 0.75, // channel 0
+        -1.0, 0.5, 0.25, // channel 1
+        0.1, 0.2, -0.3, // channel 2
+        2.0, -0.5, 1.5, // channel 3
+    ];
+    let conv_weight_f32 = mlxcel_core::from_slice_f32(&weight_data, &[hidden, l_cache, 1]);
+    let conv_weight = mlxcel_core::astype(&conv_weight_f32, dtype::BFLOAT16);
+
+    let padded_data: Vec<f32> = vec![
+        1.0, 2.0, 3.0, 4.0, // t = 0
+        -1.0, 0.5, 2.0, -2.0, // t = 1
+        0.25, -0.75, 1.0, 0.5, // t = 2
+    ];
+    let padded_f32 = mlxcel_core::from_slice_f32(&padded_data, &[1, l_cache, hidden]);
+    let padded = mlxcel_core::astype(&padded_f32, dtype::BFLOAT16);
+
+    // Ground truth: stride-1, no-pad, dilation-1, groups==hidden conv1d, run
+    // in bf16 (as it is on a bf16 checkpoint off Metal before this fix).
+    let reference = mlxcel_core::conv1d(&padded, &conv_weight, 1, 0, 1, hidden);
+    assert_eq!(mlxcel_core::array_shape(&reference), vec![1, 1, hidden]);
+    assert_eq!(mlxcel_core::array_dtype(&reference), dtype::BFLOAT16);
+
+    // Fast path, built from the bf16 conv weight exactly as `ShortConv`
+    // does for a bf16 checkpoint's non-quantized conv weight.
+    let decode_weight = build_conv_decode_weight(&conv_weight);
+    assert_eq!(
+        mlxcel_core::array_shape(&decode_weight),
+        vec![1, l_cache, hidden]
+    );
+    let elementwise = short_conv_decode_step(&padded, &decode_weight, dtype::BFLOAT16);
+    assert_eq!(mlxcel_core::array_shape(&elementwise), vec![1, 1, hidden]);
+
+    // Compare in f32 (bf16 subtraction/abs would itself be lossy).
+    let diff = mlxcel_core::subtract(
+        &mlxcel_core::astype(&reference, dtype::FLOAT32),
+        &mlxcel_core::astype(&elementwise, dtype::FLOAT32),
+    );
+    let max_abs = mlxcel_core::item_f32(&mlxcel_core::max_all(&mlxcel_core::abs(&diff)));
+    assert!(
+        max_abs < 2e-2,
+        "bf16 decode short-conv diverged from bf16 conv1d: max|diff| = {max_abs}"
+    );
+}
+
 #[test]
 fn lfm2_unquantized_config_uses_quantization_defaults() {
     // Drop the quantization block (a bf16 checkpoint) and confirm the defaults.