perf(cuda/ssm): port the fused SSM decode kernel to CUDA by inureyes · Pull Request #727 · lablup/mlxcel

inureyes · 2026-07-10T07:31:40Z

Summary

Closes the hybrid-SSM decode gap on CUDA (0.29-0.36x of Metal M1 Ultra on GB10) by porting the fused single-token SSM update kernel from mx.fast.metal_kernel to mx.fast.cuda_kernel. granite-4.0-h-350m decode improves 62.8 to 281.7 tok/s (4.5x); the granite/falcon-h1 family now beats its Metal reference numbers, greedy parity is byte-identical, and the pure-mamba2 control is unchanged.

Closes #631

Root cause

The hybrid models (granitemoehybrid, falcon_h1, plamo2, nemotron_h) use a fused SSM decode kernel that replaces the ~55-op SSD scan graph with a single launch, but it existed only as a Metal kernel: ssm_kernel_available() hard-returned false off Apple, so on CUDA every decode step ran the ssm_step graph path per SSM layer. nsys on granite-4.0-h-350m (36 decode tokens): thousands of 1-2 us elementwise/copy launches per token (copy_v 9800 instances, binary_g_nd 4984, arange 4036), GPU busy only ~30% of decode wall. Pure mamba2 never used the fused kernel on either backend, which is why its CUDA/Metal ratio was ~1.0 while the hybrids collapsed: on Metal they ran the kernel, on CUDA they ran the graph.

Changes

src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp: added SSM_CUDA_SOURCE, a line-for-line mx.fast.cuda_kernel port of the Metal SSM kernel (one warp per (batch*head, head_dim row); each lane owns Ds/32 contiguous state columns; state = dA*state_in + x*dt*B, out = warp_sum(state*C) + x*D via __shfl_down_sync; d_idx guard for ceil-div grid padding, warp-uniform). Runtime selection by metal::is_available(), the same pattern as the BitLinear and fused-MoE CUDA ports. ssm_kernel_available() now returns cu::is_available() off Apple, with the MLXCEL_SSM_CUDA_KERNEL=0 kill switch forcing the graph path for A/B and rollback.
docs/environment-variables.md: documented the new env var.
docs/benchmark_results/hybrid-ssm-decode-cuda-kernel-gb10-2026-07-10.md: measurement note with nsys evidence, family table, parity results, and the two scope findings below.

No model-side changes: all four callers pick the kernel up through the existing seq_len == 1 && ssm_kernel_available() gate.

Results (GB10, 64 decode tokens, decode tok/s)

model	before (graph path)	after	speedup	Metal M1 Ultra ref	ratio
granite-4.0-h-350m-4bit	62.8	281.7	4.5x	219.5	1.28
granite-4.0-h-tiny-4bit	38.9	101.8	2.6x	96.3	1.06
falcon-h1-tiny-90m-4bit	112.0	319.7	2.9x	288.1	1.11
nemotron-h-30b-4bit	41.5	88.1	2.1x	91.5	0.96
plamo-2-1b	33.0	43.7	1.3x	107.1	0.41
mamba2-1.3b-4bit (control)	80.1	81.0	unchanged	79.2	1.02

Acceptance criteria: granite-4.0-h-350m >= 150 tok/s met at 281.7; mamba2 control flat; Metal untouched (runtime-gated); greedy 40-token parity byte-identical on granite-4.0-h-350m, falcon-h1-tiny, and plamo-2-1b (graph path vs fused kernel on the same binary).

Scope findings

plamo-2-1b is bandwidth-capped, not software-capped. It is the only f32 checkpoint in the family (torch_dtype: float32); post-change nsys shows 88.5% of decode in gemv_single<float>, i.e. pure f32 weight streaming. The GB10 ceiling for a ~1B f32 model is ~68 tok/s (273 GB/s), so the issue's 0.8x-of-Metal target (85.7 tok/s) exceeds the hardware: the M1 Ultra reference rides 800 GB/s. GB10 now runs at 64% of its own ceiling vs Metal's 53% of its ceiling.
hunyuan-13b was misclassified into the hybrid-SSM family. Its config has no SSM/mamba keys (HunYuanMoEV1ForCausalLM, 64-expert MoE, top-8 + 1 shared); it never touches the SSM path and is unchanged (15.4 tok/s). Its decode gap belongs to the MoE-decode work (perf(cuda): single-dtype decode graph: eliminate per-token AsType conversions #636).

Validation

cargo build --release --features cuda
./target/release/mlxcel-bench-decode --model ./models/granite-4.0-h-350m-4bit --prompt "..." --no-chat-template --max-tokens 64
MLXCEL_SSM_CUDA_KERNEL=0 ./target/release/mlxcel-bench-decode --model ./models/granite-4.0-h-350m-4bit --prompt "..." --no-chat-template --max-tokens 64   # graph path
cargo test --release --features cuda -- --test-threads=1

Hybrid SSM/attention models (granite-4.0-h, falcon-h1, plamo-2, nemotron-h) decoded at 0.29-0.36x of Metal M1 Ultra on GB10 because the fused single-token SSM update kernel was Metal-only: ssm_kernel_available() hard-returned false off Apple, so every CUDA decode step ran the ~55-op ssm_step graph path per SSM layer. nsys on granite-4.0-h-350m showed thousands of 1-2 us elementwise/copy kernel launches per token (copy_v 9800, binary_g_nd 4984, arange 4036 over 36 tokens) with the GPU busy only ~30% of decode wall time. Pure mamba2 was unaffected because it never used the fused kernel on either backend, which is what localized the gap to the hybrids. Port the Metal kernel line-for-line to mx.fast.cuda_kernel: one warp per (batch*head, head_dim row), each lane owns Ds/32 contiguous state columns, state = dA*state_in + x*dt*B, out = warp_sum(state*C) + x*D via __shfl_down_sync. Runtime-selected by metal::is_available(), the same pattern as the BitLinear and fused-MoE CUDA ports; all four model callers pick it up through the existing seq_len == 1 gate. Kill switch: MLXCEL_SSM_CUDA_KERNEL=0 forces the graph path. GB10 decode, 64 tokens: granite-4.0-h-350m 62.8 -> 281.7 tok/s (4.5x, 1.28x the Metal reference), granite-4.0-h-tiny 38.9 -> 101.8 (2.6x), falcon-h1-tiny 112.0 -> 319.7 (2.9x), plamo-2-1b 33.0 -> 43.7 (1.3x, bandwidth-capped: f32 checkpoint, ~68 tok/s GB10 ceiling), mamba2 control unchanged (80.1 -> 81.0). Greedy parity byte-identical on granite/falcon-h1/plamo-2 vs the graph path. hunyuan-13b turns out to carry no SSM layers at all (HunYuanMoEV1, 64-expert MoE) and was misclassified into this family; its gap is MoE-decode territory (#636). Full note in docs/benchmark_results/hybrid-ssm-decode-cuda-kernel-gb10-2026-07-10.md.

Review follow-up: the TL;DR claimed the nemotron-h class benefits without a measured row. A/B on GB10 (MLXCEL_SSM_CUDA_KERNEL=0 vs on, 64 decode tokens): 41.5 -> 88.1 tok/s (2.1x), 0.96x of the Metal M1 Ultra reference (91.5).

) ## Summary Port the fused paged-attention decode kernel from Metal-only (`mlx::core::fast::metal_kernel`) to `mx.fast.cuda_kernel` so CUDA no longer falls back to gather-then-SDPA for the pooled paged decode path, wire the runtime gate to allow native on CUDA behind the existing `MLXCEL_PAGED_ATTENTION_NATIVE` control, keep the gather fallback selectable, and add a native-vs-fallback numerical parity test across the issue's case matrix. ## What changed `src/lib/mlx-cpp/turbo/paged_attention.cpp`: add `PAGED_ATTENTION_DECODE_CUDA_SOURCE`, a `cuda_kernel` port of the split-K flash-decoding scheme. One CUDA block per (batch, query head), `(32, NumSplits)` threads where the 32 warp lanes partition the head dimension and NumSplits warps sweep strided token stripes; online-softmax accumulation in registers; a butterfly `__shfl_xor_sync` all-reduce replacing Metal `simd_sum` (every lane needs the score, so the reduction is an all-reduce, not a reduce-to-lane-0); `__syncthreads()` combine over `__shared__` scratch. Grid `(32, NumSplits, B*Hq)` over threadgroup `(32, NumSplits, 1)` ceil-divides to exact blocks, so no padded threads and no grid-padding guards. The launcher selects the CUDA body via `!metal::is_available()`, mirroring the #631/#727 SSM and MoE ports. The port matches the Metal kernel's contract exactly (f32 Q/out, f16 K/V pool via `const __half*`, `*_shape` metadata, template args, grid). `src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp`, `mlx_cxx_bridge.h`, `src/lib/mlxcel-core/src/lib.rs`: add a `cuda_is_available()` FFI probe over `mlx::core::cu::is_available()` (backend-agnostic, `no_cuda` stub off CUDA). `src/lib/mlxcel-core/src/layers.rs`: add `PagedDecodeBackend::Cuda`; `paged_decode_backend()` returns it when CUDA is available; `select_pooled_paged_dispatch` makes native the default on CUDA for any single-slab layer (the Metal-measured batch/context ceilings do not apply on CUDA); widen the dispatch cache's backend tag to two bits (decision bit moves to bit 50). The cached `MLXCEL_PAGED_ATTENTION_NATIVE` override is unchanged and still forces native (=1) or gather (=0) both ways; Metal dispatch is untouched. `src/lib/mlxcel-core/src/ffi_tests.rs`: add `test_fused_paged_decode_native_vs_fallback_matrix` (head dims 64/80/96/128, GQA 1:1/4:1/8:1, page-boundary cases, batch > 1); the older `test_fused_paged_decode_gqa_and_batched` no longer skips on non-Metal. `examples/paged_attention_kernel_bench.rs`: pick the backend at runtime so the `select=` column reflects CUDA; document the CUDA run command. ## Dtype and scope The Metal source is fp16-only (f16 K/V pool, f32 Q/out), so the port is fp16-only as well and says so. Quantized KV pages are out of scope here and interact with #635. ## Reachability (server vs library) The fused kernel is reached only via the library-only `paged_decode_attention_pooled` -> `PagedBlockPool::paged_decode_fused` path, exercised by `examples/paged_attention_kernel_bench.rs` and external mlxcel-core consumers. #720 retired the pooled decode from the server: `mlxcel-server --decode-storage paged` routes pool-backed layers through the per-sequence `update_and_fetch` gather+SDPA intercept (`src/models/llama3.rs` returns `None` for `is_paged_backed()` caches), and `mlxcel-bench-decode` uses dense single-stream caches. Neither reaches the fused kernel, so `MLXCEL_PAGED_ATTENTION_NATIVE` does not change their output. The single-slab constraint (#235) also caps the servable context at 1024 tokens; beyond that the kernel declines and falls back to gather, which bounds the achievable long-context win independently of this port. This PR delivers the CUDA kernel and gate; wiring the fused kernel onto the server decode path would reverse #720's deliberate library-only decision and is out of scope. ## Verification (GB10, CUDA) Numerical parity: `test_fused_paged_decode_native_vs_fallback_matrix` passes with worst RMS 5.45e-5, worst max-abs 2.35e-4 over 60 configs (well under the 5e-3 fp16 threshold). All 11 `ffi_tests` paged-decode tests and 53 `layers::tests` pass under `--features cuda`. Sanity A/B (`mlxcel-bench-decode`, llama-3.1-8b-4bit, 2048-token prompt, 32 decode tokens): `MLXCEL_PAGED_ATTENTION_NATIVE=0` decode 50.93 tok/s, `=1` decode 50.94 tok/s. Identical because bench_decode does not route through the fused kernel (dense single-stream caches), confirming no regression and no accidental path divergence. Kernel A/B (`examples/paged_attention_kernel_bench`, the path that does reach the fused kernel): `select=native` now fires on CUDA for single-slab shapes; fused vs gather speedup 1.31x (batch 1, ctx 512), 1.49x (batch 1, ctx 1024), and 2.55x / 6.20x / 1.27x on the single-slab batched island. Multi-slab shapes correctly decline and read gather. ## Test plan - [x] `cargo test --release --features cuda -p mlxcel-core --lib ffi_tests::test_fused_paged_decode_native_vs_fallback_matrix -- --test-threads=1` (worst RMS 5.45e-5) - [x] `cargo test --release --features cuda -p mlxcel-core --lib paged_decode -- --test-threads=1` (11 passed) - [x] `cargo test --release --features cuda -p mlxcel-core --lib layers::tests -- --test-threads=1` (53 passed) - [x] `cargo run --release --features cuda --example paged_attention_kernel_bench` (native fires on CUDA, 1.31x-6.20x) - [x] `MLXCEL_PAGED_ATTENTION_NATIVE=0 vs =1` bench_decode A/B (no crash, no regression) Closes #634

…measurement (#757) ## Summary Closes #755 with the post-reboot re-measurement it gated on, and a verdict: **none of the sweep drops is a code regression**. The gemma-4-26b pair was depressed by the stale ~5.5-day host/driver state and recovers decisively on a fresh boot; glm4-flash and mamba2-130m were never regressed, their sweep readings are single draws from a bimodal host-level throughput distribution whose fast mode reproduces the 0.3.1 records; and the same session's SSM-cluster re-verification shows the 1.4-3.4x "anomaly" is the intended #727 fused-SSM-kernel improvement, misdiagnosed as environmental in the sweep notes. Doc rows for all 13 re-measured models now carry the post-reboot values, and the 21 committed runs live in `benchmarks/cuda_gb10_2026-07-12_postreboot_single_*.csv`. Host: fresh boot 2026-07-12 19:25 (uptime 7 min at first check), mlxcel 0.4.0-rc.1 at `fcb4c20`, MLX pin `57c66cac` (0.32.1), incremental `cargo build --release --features cuda` rebuilt clean before measuring. All script runs used `--cooldown 30`. ## Measurements ### #755 subjects and controls (decode tok/s) | Model | 2026-06-17 (0.3.1) | 2026-07-12 sweep | Post-reboot | Verdict | |---|---|---|---|---| | gemma-4-26b-a4b-it-4bit | 58.59 | 50.19 | **59.88** (n=3: 59.57-60.07) | recovered, environmental | | gemma-4-26b-a4b-it-qat-4bit | 50.33 | 45.29 | **53.68** (n=3: 53.48-53.71) | recovered, environmental | | glm4-flash-4bit | 53.33 | 45.72 | **47.42** (n=3 templated: 46.42-50.07) | within noise floor (below) | | mamba2-130m | 181.05 | 162.23 | **202.98** (n=3: 147.98-204.77) | within noise floor | | qwen3-30b-a3b-4bit (control) | 90.70 | 92.41 | 94.39 | stable | | lfm2-8b-a1b-4bit (control) | 157.73 | 161.87 | 165.53 | stable | ### SSM-cluster re-verification (decode tok/s) | Model | 2026-06-17 | 2026-07-09 single | 2026-07-12 sweep | Post-reboot | |---|---|---|---|---| | granite-4.0-h-350m-4bit | 64.00 | 86.60 | 259.69 | 171.22 | | granite-4.0-h-tiny-4bit | 33.84 | - | 100.28 | 101.37 | | falcon-h1-tiny-90m-instruct-4bit | 102.99 | 110.42 | 413.00 | 354.43 | | nemotron-h-30b-4bit | 40.32 | - | 79.94 | 87.41 | | nemotron-nas-30b-4bit | 37.33 | - | 82.72 | 85.91 | | nemotron-3-nano-omni-30b (reasoning) | 38.45 | - | 80.86 | 82.88 | | plamo-2-1b | 34.36 | 35.14 | 44.54 | 46.84 | The cluster's gains persist on a fresh host, so they are release numbers, not an artifact. The attribution is a timeline fact the sweep notes missed: the low 2026-07-09 singles ran at 15:03, and #727 (fused single-token SSM decode kernel, CUDA port; its own PR measured granite-350m at 4.5x) merged the next day at 16:45 on 2026-07-10. The "no SSM-related code has landed since" premise behind the environmental explanation was simply false, and this PR corrects it on the record. ## Attribution evidence for the non-recovering pair glm4-flash decode over twelve **identical** 100-token greedy runs on the freshly booted host (same binary, same prompt, `--no-chat-template`, identical generated tokens): 39.23, 54.80, 42.33, 52.33, 51.48, 52.86, 40.51, 42.12, 41.28, 49.29, 52.24, 42.54. The distribution is bimodal (a ~41 mode and a ~52.5 mode), and the 0.3.1 record (53.33) sits inside the fast mode. mamba2-130m shows the same flapping (147.98, then 202.98 / 204.77 / 208.25 back-to-back, vs the 181.05 record). There is no stable deficit to attribute; a single sweep run draws one sample from this distribution. The code suspects named in the issue are each ruled out: - **#740 multirow qmv**: `MLXCEL_QMV_MULTIROW=0` reads 49.55, inside the baseline envelope; by design the path keeps `M*B == 1` classic decode on the stock kernel, so B=1 benches never exercise it. - **#727 fused SSM kernel**: `MLXCEL_SSM_CUDA_KERNEL=0` on mamba2-130m reads 208.25 (unchanged); pure mamba2 never uses the fused kernel, exactly as the #727 PR stated. - **#732 single-dtype decode graph**: shipped trace tooling (`MLXCEL_TRACE_ASTYPE`) plus a default-off opt-in normalization (`MLXCEL_CUDA_F16_NORMALIZE`); nothing changes by default. The slow mode is not SM clock capping: sampling `clocks.sm` at 2 Hz during runs caught slow runs (41.28, 42.54 tok/s) at a pinned 2411-2424 MHz, the same clocks as fast runs. The CPU governor is `performance` at max frequency. The host-level mechanism behind the two modes remains unidentified, but the fast mode reproducing the 0.3.1 number rules out a code-level regression, which is what #755 needed to establish. Two further observations recorded in the docs: - The 26B gemma MoE repeats within ±0.5% while the launch-bound small models flap by up to ±25%, so the noise floor is model-class-dependent: single-run deltas on small dense/SSM checkpoints and small MoEs should not be read as regressions without repeats. - glm4-flash's templated greedy output shortened from 100 tokens (0.3.1) to 18 (rc.1), a different greedy continuation rather than a failure, so its sweep-to-sweep decode averages additionally stopped being length-comparable. ## Changes - `benchmarks/cuda_gb10_2026-07-12_postreboot_single_*.csv` (13 files, 21 runs): the post-reboot singles, including the n=3 repeat sets for both gemma variants, glm4-flash, and mamba2-130m. Env-modified A/B runs are excluded from the CSVs and reported here instead. - `docs/benchmark_results/model_tests_gb10.md`: the 13 model rows now carry the post-reboot values (medians where n=3); the SSM-cluster notable-changes bullet is rewritten with the #727 attribution and the correction of the earlier environmental claim; the moderate-drops bullet is rewritten as resolved with the evidence above and the noise-floor guidance. - `docs/benchmark_results/model_tests.md`: Nemotron-H-30B cross-hardware row updated to 87.41 and its footnote now carries the #727/#755 attribution instead of the re-verify caveat. ## Acceptance criteria (from #755) - [x] Post-reboot single-model measurements recorded for the three subjects, mamba2-130m, and both controls (singles CSVs committed; numbers above) - [x] Verdict recorded: environmental for the gemma pair; no-regression (host run-to-run variance) for glm4-flash and mamba2-130m; the SSM-cluster side settled as a real #727 improvement - [x] Doc rows and the unattributed-drop notes updated to the post-reboot values, with this issue referenced - [x] Kill-switch A/B and clock-sampling evidence recorded for the persisting readings (no code fix needed; nothing to fix) - [x] The notes in `docs/benchmark_results/model_tests_gb10.md` reference #755 ## Test plan - [x] Incremental `cargo build --release --features cuda` clean before measuring (no source changes in this PR; docs and benchmark data only) - [x] 13 post-reboot singles + 8 repeat/variance runs + 3 kill-switch A/B runs + 6 clock-sampled runs on GB10 - [x] Doc tables cross-checked against the committed CSV rows

inureyes added type:performance Performance improvements priority:medium Medium priority area:core mlxcel-core: MLX FFI, primitives, KV cache, layers bug status:review Under review labels Jul 10, 2026

inureyes mentioned this pull request Jul 10, 2026

test(surgery): empty_surgery_pipeline_matches_baseline_output is flaky on CUDA (post-EOS tail nondeterminism) #728

Closed

3 tasks

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

inureyes merged commit 6ddcad7 into main Jul 10, 2026
5 checks passed

inureyes deleted the perf/631-cuda-ssm-decode-kernel branch July 10, 2026 07:45

inureyes mentioned this pull request Jul 10, 2026

perf(cuda/attn): port native paged-attention decode kernel to CUDA #731

Merged

5 tasks

inureyes mentioned this pull request Jul 12, 2026

docs(benchmarks): attribute the GB10 decode drops via post-reboot re-measurement #757

Merged

8 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

perf(cuda/ssm): port the fused SSM decode kernel to CUDA#727

perf(cuda/ssm): port the fused SSM decode kernel to CUDA#727
inureyes merged 2 commits into
mainfrom
perf/631-cuda-ssm-decode-kernel

inureyes commented Jul 10, 2026 •

edited

Loading

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

Uh oh!

Conversation

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

Uh oh!

Summary

Root cause

Changes

Results (GB10, 64 decode tokens, decode tok/s)

Scope findings

Validation

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant

inureyes commented Jul 10, 2026 •

edited

Loading