perf(cuda/ssm): close the hybrid-SSM decode gap on CUDA #631

Part of #623. Depends on #625: confirm persistence on the upgraded pin first.

Context

Hybrid SSM/attention models decode at 0.29-0.36x of Metal on GB10, far below the ~0.61 median CUDA/Metal ratio that pure bandwidth differences explain. Pure Mamba2 is fine, which localizes the problem to something the hybrids use and mamba2 does not:

(Source: benchmarks/cuda_gb10_2026-06-17.csv vs benchmarks/metal_m1ultra_2026-06-15.csv.)

These are small models, so the absolute cost is per-step overhead or a slow op, not bandwidth. Likely suspects: a per-step op in the hybrid recurrent block that has no fast CUDA path (conv1d state update, segmented scan variant, gating elementwise chain running as many small unfused kernels), a CPU fallback inside MLX for one op, or an eval/sync point in the hybrid cache update that the mamba2 path avoids.

Scope

Diagnose and fix decode throughput for the hybrid-SSM family on CUDA. Target: ratio vs Metal comparable to the control model (>= 0.8).

Implementation plan

1. Pick the smallest repro: granite-4.0-h-350m-4bit. Run with MLXCEL_PROFILE_PIPELINE_DETAIL=1 to split reshape/forward/sample/item-wait per token, then nsys profile to enumerate kernels per step and any cudaMemcpy/host sync inside the forward.
2. Locate the model implementations under src/models/ (granite hybrid, falcon-h1, plamo2, hunyuan) and the shared SSM building blocks in mlxcel-core (src/lib/mlxcel-core/src/ layers/cache: recurrent state caches, ModelStateSnapshot machinery in generate.rs:84-138). Diff the op composition of the hybrid recurrent block vs src/models/mamba2*.rs to find which op the fast control path avoids.
3. Compare against Python mlx-lm on the same checkpoint (scripts/bench_mlxlm.py) to decide kernel-level vs bridge-level.
4. Fix according to what profiling shows. Plausible shapes of the fix:
   • replace a hand-composed elementwise chain with an existing compiled window (add a compiled_* fused wrapper in src/lib/mlxcel-core/cpp/mlx_cxx_bridge.cpp, pattern: compiled_swiglu_activation);
   • port the hot op to fast::cuda_kernel following the fused-MoE port in src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp (use_cuda = !metal::is_available(), __shfl_down_sync reduction pattern);
   • remove an unnecessary per-step eval/host sync in the hybrid state cache update;
   • if it is an MLX CPU fallback, file/adopt upstream and overlay in src/lib/mlx-cpp/patches/mlx/backend/cuda/.
5. Re-run the whole hybrid family, plus mamba2 as a no-regression control.

Acceptance criteria

• granite-4.0-h-350m-4bit decode >= 150 tok/s on GB10 (from 64; ratio >= 0.68 as floor, target 0.8).
• Every model in the table improves; family table committed to a dated docs/benchmark_results/ note including the profiling evidence and the identified root cause.
• mamba2 control does not regress; Metal untouched (backend-conditional changes only).
• Greedy parity on 40-token continuations for granite-4.0-h and falcon-h1 vs pre-change output.
• cargo test --features cuda passes.

Validation

cargo build --release --features cuda
MLXCEL_PROFILE_PIPELINE_DETAIL=1 ./target/release/mlxcel-bench-decode --model ./models/granite-4.0-h-350m-4bit --max-tokens 64
nsys profile -o ssm_decode ./target/release/mlxcel-bench-decode --model ./models/granite-4.0-h-350m-4bit --max-tokens 16

References

• Snapshot/state machinery: src/lib/mlxcel-core/src/generate.rs:84-138,391-416; hybrid prompt-cache src/server/prompt_cache/hybrid_ssm.rs.
• cuda_kernel port pattern: src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp (fused MoE, perf(moe): implement the fused decode-MoE kernel on CUDA #319).
• Baselines: benchmarks/cuda_gb10_2026-06-17.csv, benchmarks/metal_m1ultra_2026-06-15.csv.