Add an opt-in Helion kernel backend for the Mamba-2 SSD path

[shcho1118]

Summary

Adds a Helion implementation of the nine SSD kernels behind mamba_split_conv1d_scan_combined, plus the gated-LayerNorm backward, selected at runtime by MAMBA_USE_HELION=1. On a B200 it cuts the device time of those stages by 1.39x (fwd+bwd, seq 8192, bf16). Without the env var nothing changes: the Triton kernels stay the default and the numerical reference, and helion does not even have to be installed.

Design

The Triton path is untouched. The only changes to existing code are +38/-13 lines in two files, and they are pure indirection — the added lines are mostly a namespace listing the nine Triton kernels:

• ssd_combined.py — the nine Triton kernels are collected into _TRITON_SSD_KERNELS, and _ssd_kernel_impls() returns that namespace, or the Helion one under the opt-in. The fwd/bwd functions resolve it once and call through it.
• layernorm_gated.py — LayerNormFn.backward picks the Helion _layer_norm_bwd over the Triton one.

Dispatch is opt-in and lazy. mamba_ssm/ops/helion/dispatch.py is the only import of the new package from production code, and it imports the kernels — and therefore helion itself — only after MAMBA_USE_HELION=1 is confirmed. An install without helion imports and runs exactly as before.

Configs are files, never autotuned at runtime. Each kernel resolves its Helion config from MAMBA_HELION_CONFIG_DIR/<kernel>.json on first call; a missing config is a hard error unless MAMBA_HELION_AUTOTUNE is set explicitly. Autotuning must not happen under CUDA graph capture or in a distributed job, so it is an offline step run by a harness in benchmarks/helion/, which also verifies that every dispatched kernel left a config behind.

New code lives in mamba_ssm/ops/helion/ (runtime: dispatch, config loading, one file per kernel) and benchmarks/helion/ (offline: autotune harness, benchmark, example configs — excluded from the wheel, so nothing writes into site-packages).

Measurements

Device time per fwd+bwd iteration, MambaMixerMin (a Megatron-free replica of Megatron-LM's MambaMixer._ssm_training layout) at Nemotron3-Nano dimensions: d_model 2688, 64 heads x 64, d_state 128, 8 groups, chunk_size 128, seq_len 8192, batch 1, bf16, unpacked. NVIDIA B200, torch 2.11.0a0 (NGC 26.03, CUDA 13.2), triton 3.6.0, helion 1.4.0.

The left group is each dispatch stage's whole device time, including the aten helpers it launches; the right group is the generated Triton/Helion kernel alone.

Stage	Triton	Helion	Speedup	Triton (kernel)	Helion (kernel)	Speedup (kernel)
layer_norm_bwd	0.190	0.049	3.90x	0.182	0.035	5.24x
chunk_scan_fwd	0.174	0.138	1.26x	0.174	0.138	1.26x
chunk_scan_chunk_state_bwd_dx	0.169	0.146	1.16x	0.155	0.132	1.17x
chunk_scan_bwd_ddAcs_stable	0.142	0.112	1.27x	0.137	0.104	1.32x
chunk_state_bwd_db	0.122	0.099	1.22x	0.088	0.064	1.39x
chunk_state_fwd	0.121	0.115	1.05x	0.121	0.115	1.05x
chunk_scan_bwd_dC	0.114	0.083	1.38x	0.089	0.058	1.53x
state_passing_bwd	0.102	0.063	1.63x	0.099	0.060	1.65x
state_passing_fwd	0.101	0.078	1.30x	0.101	0.078	1.30x
chunk_scan_bwd_dstates	0.059	0.052	1.13x	0.059	0.052	1.13x
TOTAL (stages)	1.295	0.935	1.39x	1.206	0.836	1.44x

Attribution is per dispatch stage rather than per kernel name, because the two backends emit differently named kernels; the two column groups are split because they leave different amounts of work to aten (e.g. Triton reduces dD with a separate reduce_kernel where Helion uses in-kernel atomics).

Reproduce with python benchmarks/helion/bench_mamba_mixer.py (profiles both backends in separate subprocesses and prints the table).

Testing

pip install -e '.[helion,causal-conv1d]' --no-build-isolation
pytest tests/ops/helion/test_kernels.py

tests/ops/helion/test_kernels.py compares every kernel with its Triton reference inside a real pass: one MambaMixerMin fwd+bwd runs with the pipeline on Triton, and each dispatched kernel is wrapped so both implementations are invoked on the same production arguments and their outputs compared. The wrapper returns the Triton result, so every later stage still consumes exact production data and a mismatch is attributable to the kernel that produced it. Arguments are never rebuilt by hand — they are views into shared buffers (x/B/C/dx slice one packed xBC) and outputs of preceding stages, and Helion compiles per stride, so a contiguous stand-in would test a specialization that never runs.

Tolerances are one bf16 ulp for rtol plus an absolute floor of one ulp of the tensor's own max magnitude, since these outputs span five orders of magnitude in element size. With MAMBA_HELION_CONFIG_DIR unset the tests use the config set checked in under benchmarks/helion/configs/, so no environment setup is needed; if that set does not fit the GPU or toolchain, the failure carries the command to autotune a new one. The packed (THD) path is covered by MAMBA_HELION_TEST_PACKED=1 with its own config set.

Existing suites were run on this branch with MAMBA_USE_HELION unset, where the dispatch resolves to the same Triton functions as before: test_layernorm_gated.py 40 passed / 24 skipped, tests/test_determinism.py 35 passed, tests/ops/triton/test_ssd.py 23 passed with test_chunk_state_varlen[128-max-dtype2] failing. That one failure reproduces identically on the merge base without this branch, so it is pre-existing and unrelated.

Limitations, and questions for maintainers

• Determinism. The Helion backward kernels choose their reduction strategy at compile time and do not consult use_deterministic_mode(): chunk_scan_chunk_state_bwd_dx (ddt, 1-D dD) and chunk_scan_bwd_dC (ddA_cumsum_prev) always reduce with atomic adds. Enabling both warns once, and the Triton path remains the deterministic one. If you would rather the backend refuse to run in deterministic mode, that is a small change.
• The helion extra pins helion==1.4.0 rather than lower-bounding it, because a config is autotuned against one toolchain and Helion's config schema is still moving; bumping the pin means re-autotuning the checked-in set. Say the word if you would rather carry a range.
• Configs are environment-specific. A set is tied to (GPU, helion/triton version, dims, dtype, packing); the directory naming scheme encodes those axes. One example set is checked in (11 files, 48 KB) purely so the tests and benchmark are reproducible — happy to drop it and document generation instead if you would rather not carry generated data in-tree.
• CI. The tests need a CUDA device plus helion; they importorskip and skip cleanly otherwise, so they are inert in CPU CI. I can wire them into a workflow if there is a GPU runner.
• mamba_mixer_min.py sits in the package rather than in benchmarks/, because the harness, the benchmark and the tests all import it. Can move it wherever you prefer.

Commits

12 commits, each self-contained: the default path is Triton at every commit, and import mamba_ssm works without helion installed at every commit (verified by blocking the helion import and importing the package at each one). Kernels land one group at a time, each adding its entry to the dispatch namespace; the harness, tests, benchmark, example configs and docs follow.