refactor(models): share the short-conv decode helper and audit L=1 conv dispatch across SSM/hybrid decode paths

docs/benchmark_results/model_tests_gb10.md

@@ -178,7 +178,7 @@ Prefill/Decode are the measured-pass figures from `mlxcel-bench-decode`. Notes r
 | falcon-mamba-7b-4bit | ✅ | 92.54 | 22.06 | 2 tok |
 | jamba-v0.1-4bit | ✅ | 523.67 | 89.63 |  |
 | 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-130m | ✅ | 890.29 | 162.23 | decode conv audited (#752): runs on the fast `conv1d_c1_k1_nhwc` engine (576 launches = 24 layers x 24 tokens, one per op), not the `convolve_common_engine` per-channel path; the -10.4% sweep delta is environmental, not conv dispatch |
 | mamba2-1.3b-4bit | ✅ | 329.13 | 83.40 |  |
 | plamo-2-1b | ✅ | 199.12 | 44.54 |  |
 
@@ -374,7 +374,8 @@ Models that accept image input and generated tokens under the `"What is in this
 
 - **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, 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%).
+- **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%). These are NOT the #748 conv-dispatch pattern (correcting the PR #751 hypothesis on the record, per #752): neither model is a conv architecture (`gemma4` MoE and `glm4_moe_lite` respectively; neither file calls `conv1d`), so their decode conv cannot fall into cuDNN's grouped-conv engine. The drops need separate attribution and are out of scope for the conv audit.
+- **SSM / hybrid L=1 conv dispatch audited across the family (#752), no further regression found**: after #748/#751 fixed LFM2, every other decode-path depthwise-conv family was measured under `nsys` on a warm GB10 CUDA decode (`MLX_USE_CUDA_GRAPHS=0` for a complete kernel histogram). All run on the fast `conv1d_c1_k1_nhwc` cuDNN engine, not the per-channel `convolve_common_engine` that regressed LFM2: mamba2-130m (bf16), mamba2-1.3b-4bit, falcon-mamba-7b-4bit, falcon-h1-tiny-90m-4bit, granite-4.0-h-350m-4bit, jamba-v0.1-4bit, plamo-2-1b (f32 activations), qwen3.5-0.8b-4bit (covers the gated-delta conv), and nemotron-h-30b-4bit. The proof is the launch count: each shows conv instances equal to (conv layers x decode tokens), one launch per op, whereas the slow engine launches one kernel per channel (thousands per op). LFM2's slow-engine dispatch is specific to its `conv_L_cache = 3` / hidden-1024 shape and does not reproduce at the SSM `conv_kernel = 4` widths, so no family is adapted. All of this was measured on MLX pin `57c66cac` (0.32.1), the same pin whose CUDA conv dispatch heuristic produced the lfm2 slow path, so these REFUTED verdicts are point-in-time and should be re-checked against newer MLX pins. The #751 short-conv decode helper was still lifted into a shared `models::conv_decode` module so any future family that does regress can adopt it in one line.
 - **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.
 - **VLM (image-input) gains**: molmo-7b 23.81 → 35.89 (+51%), nemotron-omni-30b 32.86 → 70.09 (+113%), gemma3-4b family +13-18%.

src/models/conv_decode.rs

@@ -0,0 +1,76 @@
+// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Shared single-step (decode) depthwise short-conv fast path.
+//!
+//! 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`), which launches one kernel per channel
+//! and dominates decode. Issue #748 / PR #751 first hit this on LFM2's
+//! `ShortConv` and replaced the L=1 step with a broadcast multiply plus an axis
+//! sum against a decode weight precomputed at load. The identical rolling-window
+//! pad + depthwise `conv1d` call pattern executed every decode step exists in the
+//! SSM / hybrid families (mamba2, falcon-h1, granite-4.0-h, mamba, jamba, plamo2,
+//! nemotron-h, kimi linear, qwen3.5 linear-attention, …), so issue #752 lifts the
+//! two #751 helpers here for reuse.
+//!
+//! Both helpers are pure, depend only on `mlxcel-core`, and are checkpoint-free,
+//! so the parity tests in `conv_decode_tests.rs` (and the per-family test files)
+//! pin the elementwise step against `conv1d` without loading a model.
+
+use mlxcel_core::{MlxArray, UniquePtr};
+
+/// Materialize the time-major weight used by the decode fast path.
+///
+/// `conv_weight` is `[channels, kernel, 1]` (MLX depthwise layout); transposing
+/// to `[1, kernel, channels]` lets a single broadcast multiply-and-sum over the
+/// `kernel` 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> {
+    // [channels, kernel, 1] -> [1, kernel, channels].
+    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 / #752).
+///
+/// For `padded` of shape `[batch, kernel, channels]` and `decode_weight` of
+/// shape `[1, kernel, channels]` this returns `[batch, 1, channels]` 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 == channels` `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 and dominates decode.
+///
+/// `in_dtype` is the (possibly bf16/f16) activation dtype; the decode weight is
+/// cast to it only when it differs, so quantized checkpoints whose non-quantized
+/// conv weight is stored at a wider precision still multiply in the activation
+/// dtype (MLX widens the reduce accumulator for half dtypes).
+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, channels]
+}

src/models/conv_decode_tests.rs

@@ -0,0 +1,205 @@
+// Copyright 2025-2026 Lablup Inc. and Jeongkyu Shin
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+//! Checkpoint-free parity tests for the shared single-step short-conv decode
+//! fast path (`build_conv_decode_weight` + `short_conv_decode_step`, issue
+//! #748 / #752).
+//!
+//! Each test builds a channel- and tap-asymmetric depthwise conv weight so a
+//! transpose or axis bug cannot pass by symmetry, then asserts the elementwise
+//! decode step is numerically identical to the stride-1 / no-pad / dilation-1 /
+//! `groups == channels` `conv1d` it replaces. Kernel sizes 3 (LFM2) and 4 (the
+//! mamba2 / falcon-h1 / granite-4.0-h / mamba / jamba / plamo2 / nemotron-h SSM
+//! conv width) are both covered, in f32 and bf16 (the dtype the fast path
+//! actually runs in on real checkpoints off Metal). The per-family test files
+//! add the family-specific post-conv shaping (bias, SiLU); the numeric core of
+//! the conv is pinned here.
+
+use super::conv_decode::{build_conv_decode_weight, short_conv_decode_step};
+use mlxcel_core::dtype;
+use mlxcel_core::utils::silu;
+
+/// Build an asymmetric `[channels, kernel, 1]` depthwise conv weight whose taps
+/// differ per channel and per position (no symmetry to hide a transpose bug).
+/// Values are kept in roughly `[-1, 1]` so that, like the LFM2 tests, bf16
+/// rounding stays small and the parity tolerances remain meaningful (large
+/// synthetic magnitudes would make bf16 error swamp any real logic bug).
+fn asymmetric_conv_weight(channels: usize, kernel: usize) -> Vec<f32> {
+    let mut data = Vec::with_capacity(channels * kernel);
+    for c in 0..channels {
+        for k in 0..kernel {
+            // Distinct, non-monotone value per (channel, tap), bounded in ~[-1, 1).
+            let v = (((c * 7 + k * 13 + 3) % 17) as f32) / 17.0 - 0.5;
+            data.push(if (c + k) % 2 == 0 { v } else { -v });
+        }
+    }
+    data
+}
+
+/// Deterministic asymmetric `[1, kernel, channels]` padded activation window,
+/// also bounded in ~`[-1, 1]`.
+fn asymmetric_padded(channels: usize, kernel: usize) -> Vec<f32> {
+    let mut data = Vec::with_capacity(channels * kernel);
+    for k in 0..kernel {
+        for c in 0..channels {
+            let v = (((k * 5 + c * 11 + 1) % 13) as f32) / 13.0 - 0.5;
+            data.push(if (k + c) % 3 == 0 { -v } else { v });
+        }
+    }
+    data
+}
+
+/// f32 parity: the elementwise decode step equals `conv1d` for a length-1
+/// output, for the given channel count and kernel width.
+fn assert_decode_matches_conv1d_f32(channels: usize, kernel: usize) {
+    let ch = channels as i32;
+    let k = kernel as i32;
+
+    let weight_data = asymmetric_conv_weight(channels, kernel);
+    let conv_weight = mlxcel_core::from_slice_f32(&weight_data, &[ch, k, 1]);
+
+    let padded_data = asymmetric_padded(channels, kernel);
+    let padded = mlxcel_core::from_slice_f32(&padded_data, &[1, k, ch]);
+
+    let reference = mlxcel_core::conv1d(&padded, &conv_weight, 1, 0, 1, ch);
+    assert_eq!(mlxcel_core::array_shape(&reference), vec![1, 1, ch]);
+
+    let decode_weight = build_conv_decode_weight(&conv_weight);
+    assert_eq!(mlxcel_core::array_shape(&decode_weight), vec![1, k, ch]);
+
+    let elementwise = short_conv_decode_step(&padded, &decode_weight, dtype::FLOAT32);
+    assert_eq!(mlxcel_core::array_shape(&elementwise), vec![1, 1, ch]);
+
+    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,
+        "f32 decode short-conv diverged from conv1d (channels={channels}, kernel={kernel}): max|diff| = {max_abs}"
+    );
+}
+
+/// bf16 parity: same as the f32 case but in the dtype the fast path runs in on
+/// real (bf16-activation) checkpoints. Values are built in f32 and cast to bf16
+/// for both the reference and the fast path, isolating dtype rounding from
+/// construction differences.
+fn assert_decode_matches_conv1d_bf16(channels: usize, kernel: usize) {
+    let ch = channels as i32;
+    let k = kernel as i32;
+
+    let weight_data = asymmetric_conv_weight(channels, kernel);
+    let conv_weight_f32 = mlxcel_core::from_slice_f32(&weight_data, &[ch, k, 1]);
+    let conv_weight = mlxcel_core::astype(&conv_weight_f32, dtype::BFLOAT16);
+
+    let padded_data = asymmetric_padded(channels, kernel);
+    let padded_f32 = mlxcel_core::from_slice_f32(&padded_data, &[1, k, ch]);
+    let padded = mlxcel_core::astype(&padded_f32, dtype::BFLOAT16);
+
+    let reference = mlxcel_core::conv1d(&padded, &conv_weight, 1, 0, 1, ch);
+    assert_eq!(mlxcel_core::array_dtype(&reference), dtype::BFLOAT16);
+
+    let decode_weight = build_conv_decode_weight(&conv_weight);
+    let elementwise = short_conv_decode_step(&padded, &decode_weight, dtype::BFLOAT16);
+    assert_eq!(mlxcel_core::array_shape(&elementwise), vec![1, 1, ch]);
+
+    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)));
+    // bf16 carries ~3 decimal digits; the assertion takes a max over all
+    // channels, and the expected magnitude of a max over more independent
+    // bf16 rounding errors grows roughly like sqrt(channels) (an
+    // order-statistic bound), so the tolerance scales with channel count too.
+    let tol = 3e-2 * (channels as f32).sqrt();
+    assert!(
+        max_abs < tol,
+        "bf16 decode short-conv diverged from bf16 conv1d (channels={channels}, kernel={kernel}): max|diff| = {max_abs} (tol {tol})"
+    );
+}
+
+#[test]
+fn conv_decode_matches_conv1d_kernel3_f32() {
+    // LFM2 `conv_L_cache = 3` width, many channels.
+    assert_decode_matches_conv1d_f32(48, 3);
+}
+
+#[test]
+fn conv_decode_matches_conv1d_kernel4_f32() {
+    // The mamba2 / falcon-h1 / granite-4.0-h / mamba / jamba / plamo2 /
+    // nemotron-h SSM conv width (`conv_kernel = 4`).
+    assert_decode_matches_conv1d_f32(64, 4);
+}
+
+#[test]
+fn conv_decode_matches_conv1d_kernel3_bf16() {
+    assert_decode_matches_conv1d_bf16(48, 3);
+}
+
+#[test]
+fn conv_decode_matches_conv1d_kernel4_bf16() {
+    assert_decode_matches_conv1d_bf16(64, 4);
+}
+
+#[test]
+fn conv_decode_with_bias_and_silu_matches_conv1d_bf16() {
+    // The SSM families post-process the conv as `silu(conv1d(x) + bias)`
+    // (mamba2 / falcon-h1 / granite-4.0-h / mamba / jamba / nemotron-h). Assert
+    // the whole shaped step is identical whether the conv came from the fast
+    // decode path or from `conv1d`, so the adaptation is safe end-to-end and not
+    // only at the bare-conv boundary.
+    let channels = 64usize;
+    let kernel = 4usize;
+    let ch = channels as i32;
+    let k = kernel as i32;
+
+    let weight_data = asymmetric_conv_weight(channels, kernel);
+    let conv_weight = mlxcel_core::astype(
+        &mlxcel_core::from_slice_f32(&weight_data, &[ch, k, 1]),
+        dtype::BFLOAT16,
+    );
+
+    // Asymmetric per-channel bias, bounded in ~[-0.5, 0.5].
+    let bias_data: Vec<f32> = (0..channels)
+        .map(|c| (((c * 3 + 2) % 7) as f32) / 7.0 - 0.5 + if c % 2 == 0 { 0.1 } else { -0.1 })
+        .collect();
+    let bias = mlxcel_core::astype(
+        &mlxcel_core::from_slice_f32(&bias_data, &[ch]),
+        dtype::BFLOAT16,
+    );
+    let bias_row = mlxcel_core::reshape(&bias, &[1, 1, -1]);
+
+    let padded = mlxcel_core::astype(
+        &mlxcel_core::from_slice_f32(&asymmetric_padded(channels, kernel), &[1, k, ch]),
+        dtype::BFLOAT16,
+    );
+
+    // Reference: conv1d -> + bias -> silu.
+    let ref_conv = mlxcel_core::conv1d(&padded, &conv_weight, 1, 0, 1, ch);
+    let ref_out = silu(&mlxcel_core::add(&ref_conv, &bias_row));
+
+    // Fast path: decode step -> + bias -> silu.
+    let decode_weight = build_conv_decode_weight(&conv_weight);
+    let fast_conv = short_conv_decode_step(&padded, &decode_weight, dtype::BFLOAT16);
+    let fast_out = silu(&mlxcel_core::add(&fast_conv, &bias_row));
+
+    let diff = mlxcel_core::subtract(
+        &mlxcel_core::astype(&ref_out, dtype::FLOAT32),
+        &mlxcel_core::astype(&fast_out, dtype::FLOAT32),
+    );
+    let max_abs = mlxcel_core::item_f32(&mlxcel_core::max_all(&mlxcel_core::abs(&diff)));
+    assert!(
+        max_abs < 3e-2,
+        "silu(bias + conv) diverged between fast decode and conv1d: max|diff| = {max_abs}"
+    );
+}