diff --git a/benchmarks/helion/autotune_mamba_mixer.py b/benchmarks/helion/autotune_mamba_mixer.py
new file mode 100644
index 000000000..f4c13f0c1
--- /dev/null
+++ b/benchmarks/helion/autotune_mamba_mixer.py
@@ -0,0 +1,300 @@
+"""Autotune the Helion kernels via a repo-internal MambaMixer.
+
+Drives :class:`MambaMixerMin` (a Megatron-free replica of MambaMixer's training
+path) through one forward+backward, hitting every Helion kernel in the production
+path on a single GPU with no Megatron-LM dependency. Dimensions default to the
+Nemotron3 Nano production config (seq_len 8k).
+
+``MAMBA_HELION_AUTOTUNE=1`` autotunes only kernels whose ``<kernel>.json`` is
+missing from ``MAMBA_HELION_CONFIG_DIR``; ``=2`` forces a re-autotune of all of
+them (e.g. for a new GPU / Helion version). Both default on: AUTOTUNE=1 and
+CONFIG_DIR=``<this dir>/configs``. The harness also enables
+``MAMBA_USE_HELION=1`` so production dispatch selects Helion.
+
+    python benchmarks/helion/autotune_mamba_mixer.py
+
+After the run the harness checks that every dispatched kernel (the nine SSD
+kernels plus the gated LayerNorm backward) left a ``<kernel>.json`` behind, and
+fails if one did not -- otherwise a kernel this configuration never reaches would
+only surface later as a runtime ``ensure_helion_config`` assertion.
+
+A config is specialized to the run's shapes/path (dims, dtype, ``--sequence-packing``,
+...), so autotune each distinct variant into its own ``MAMBA_HELION_CONFIG_DIR`` and
+point the runtime at the matching set.
+"""
+
+import argparse
+import os
+from pathlib import Path
+
+import torch
+
+# NOTE: mamba_ssm.ops.helion.mamba_mixer_min (and through it ssd_combined /
+# layernorm_gated) is imported inside main(), after configure_env() -- dispatch's
+# use_helion() is @cache'd, so any import-time read of MAMBA_USE_HELION would pin
+# the whole harness to Triton.
+
+_DTYPE_MAP = {"fp32": torch.float32, "bf16": torch.bfloat16, "fp16": torch.float16}
+
+
+def parse_args():
+    p = argparse.ArgumentParser(description="Autotune Helion kernels via MambaMixer")
+    # Dimensions — defaults match the Nemotron3 Nano production config (seq_len 8k).
+    p.add_argument("--hidden-size", type=int, default=2688)
+    p.add_argument("--mamba-num-heads", type=int, default=64)
+    p.add_argument("--mamba-head-dim", type=int, default=64)
+    p.add_argument("--mamba-state-dim", type=int, default=128)
+    p.add_argument("--mamba-num-groups", type=int, default=8)
+    p.add_argument("--chunk-size", type=int, default=128)
+    p.add_argument("--seq-len", type=int, default=8192)
+    p.add_argument("--batch-size", type=int, default=1)
+    p.add_argument("--dtype", choices=["fp32", "bf16", "fp16"], default="bf16")
+    p.add_argument("--seed", type=int, default=0)
+    # Sets HELION_AUTOTUNE_PRECOMPILE. "spawn" precompiles each candidate config in a
+    # subprocess, so a config that hard-crashes the CUDA context (e.g. Triton
+    # "misaligned address") is skipped instead of aborting the whole autotune run;
+    # it costs more memory and cold-start time. "fork" is Helion's default, "off"
+    # precompiles in-process.
+    p.add_argument(
+        "--precompile",
+        choices=["spawn", "fork", "off"],
+        default=None,
+        help="Helion autotune precompile mode. Explicit flag wins; otherwise an "
+        "exported HELION_AUTOTUNE_PRECOMPILE is kept, else 'fork'. Use 'spawn' to "
+        "isolate configs that crash the CUDA context.",
+    )
+    # Sequence packing (THD) — feeds a seq_idx tensor so the packed-sequence kernel
+    # path is autotuned. Requires --batch-size 1.
+    p.add_argument(
+        "--sequence-packing",
+        action="store_true",
+        help="Build a seq_idx and run the packed (THD) path.",
+    )
+    p.add_argument(
+        "--packed-num-seqs",
+        type=int,
+        default=4,
+        help="Number of sub-sequences to split --seq-len into (uniform).",
+    )
+    p.add_argument(
+        "--packed-seqlens",
+        type=str,
+        default=None,
+        help="Comma-separated sub-seq lengths summing to --seq-len. "
+        "Overrides --packed-num-seqs.",
+    )
+    return p.parse_args()
+
+
+def build_seq_idx(args):
+    """Return a (1, seq_len) int32 seq_idx mapping each token to its sub-seq index.
+
+    Mirrors Megatron MambaMixer._create_packed_seq_idx: for sub-seq lengths
+    [5, 2, 4] the result is [0,0,0,0,0, 1,1, 2,2,2,2]. Returns None when packing
+    is disabled.
+    """
+    if not args.sequence_packing:
+        return None
+    if args.batch_size != 1:
+        raise ValueError(
+            f"Sequence packing requires --batch-size 1 (got {args.batch_size})."
+        )
+
+    if args.packed_seqlens:
+        lengths = [int(s) for s in args.packed_seqlens.split(",") if s.strip()]
+        if sum(lengths) != args.seq_len:
+            raise ValueError(
+                f"--packed-seqlens sum ({sum(lengths)}) != --seq-len ({args.seq_len})."
+            )
+    else:
+        n = args.packed_num_seqs
+        if not 1 <= n <= args.seq_len:
+            raise ValueError(
+                f"--packed-num-seqs must be in [1, seq_len={args.seq_len}] (got {n})."
+            )
+        # Random sub-seq lengths (not a uniform split) so the packed path sees
+        # ragged boundaries like real data. Pick n-1 distinct interior cut points;
+        # the gaps give n lengths, each >= 1, summing to seq_len. Reproducible via
+        # --seed (main seeds torch).
+        cuts, _ = torch.sort(torch.randperm(args.seq_len - 1)[: n - 1] + 1)
+        bounds = torch.cat([cuts.new_zeros(1), cuts, cuts.new_full((1,), args.seq_len)])
+        lengths = torch.diff(bounds).tolist()
+
+    idx = torch.arange(len(lengths), device="cuda")
+    seq_idx = (
+        torch.repeat_interleave(idx, torch.tensor(lengths, device="cuda"))
+        .to(torch.int32)
+        .unsqueeze(0)
+    )  # (1, seq_len)
+    print(
+        f"sequence_packing : {len(lengths)} sub-seqs, lengths={lengths[:10]}"
+        f"{'...' if len(lengths) > 10 else ''}"
+    )
+    return seq_idx
+
+
+def expected_kernel_names():
+    """Return ``{dispatch wrapper name: [helion kernel name, ...]}``.
+
+    Walks the dispatch namespaces -- the nine SSD kernels plus the gated
+    LayerNorm backward -- and, for each wrapper, picks out the
+    ``json_cached_autotune``-wrapped kernels it references, so the list stays in
+    sync with dispatch.py instead of being restated here. The kernel names are
+    the ``<name>.json`` basenames ``ensure_helion_config`` looks for at runtime.
+    """
+    import helion
+
+    from mamba_ssm.ops.helion.dispatch import (
+        get_helion_layer_norm_bwd,
+        get_helion_ssd_kernels,
+    )
+
+    wrappers = list(vars(get_helion_ssd_kernels()).values())
+    wrappers.append(get_helion_layer_norm_bwd())
+
+    names = {}
+    for wrapper in wrappers:
+        referenced = [
+            obj.kernel.name
+            for obj in (wrapper.__globals__.get(n) for n in wrapper.__code__.co_names)
+            if isinstance(getattr(obj, "kernel", None), helion.Kernel)
+        ]
+        assert referenced, f"no Helion kernel found in {wrapper.__name__}"
+        names[wrapper.__name__] = referenced
+    return names
+
+
+def verify_configs(config_dir):
+    """Fail unless every dispatched kernel has a config file in ``config_dir``.
+
+    A kernel that the harness never reaches would otherwise leave the run
+    "successful" with its JSON missing, and only blow up later in
+    ``ensure_helion_config`` at runtime -- exactly the invariant this harness
+    exists to guarantee.
+    """
+    missing = []
+    for wrapper_name, kernel_names in expected_kernel_names().items():
+        for kernel_name in kernel_names:
+            found = (config_dir / f"{kernel_name}.json").is_file()
+            if not found:
+                missing.append(f"{kernel_name}.json (dispatched as {wrapper_name})")
+            print(f"  [{'ok' if found else 'MISSING'}] {kernel_name}.json")
+    assert not missing, (
+        f"{len(missing)} Helion config(s) missing from {config_dir} after the "
+        "autotune run -- the kernel(s) were not reached by this configuration:\n  "
+        + "\n  ".join(missing)
+    )
+
+
+def configure_env(args):
+    """Set the Helion env vars. Must run before anything imports the dispatch path.
+
+    ``dispatch.use_helion()`` is ``@cache``d, so the first read of
+    ``MAMBA_USE_HELION`` wins for the whole process: setting it after an import
+    that already asked would silently run this harness on Triton and autotune
+    nothing.
+    """
+    # Defaults so the script autotunes into the repo's configs dir out of the box;
+    # explicit env vars still win (setdefault). For the packed (THD) path you must set
+    # MAMBA_HELION_CONFIG_DIR yourself to a separate dir -- see the sequence-packing
+    # note above on why packed configs must stay separate from the non-packed ones.
+    os.environ.setdefault("MAMBA_USE_HELION", "1")
+    assert os.environ["MAMBA_USE_HELION"] == "1", (
+        "The Helion autotune harness requires MAMBA_USE_HELION=1"
+    )
+    os.environ.setdefault("MAMBA_HELION_AUTOTUNE", "1")
+    default_config_dir = Path(__file__).resolve().parent / "configs"
+    os.environ.setdefault("MAMBA_HELION_CONFIG_DIR", str(default_config_dir))
+    Path(os.environ["MAMBA_HELION_CONFIG_DIR"]).mkdir(parents=True, exist_ok=True)
+    # An explicit --precompile wins; otherwise keep an exported value, like the
+    # three above, and fall back to the harness default.
+    if args.precompile is not None:
+        os.environ["HELION_AUTOTUNE_PRECOMPILE"] = (
+            "0" if args.precompile == "off" else args.precompile
+        )
+    else:
+        os.environ.setdefault("HELION_AUTOTUNE_PRECOMPILE", "fork")
+
+    if os.environ.get("MAMBA_HELION_AUTOTUNE") not in ("1", "2"):
+        print(
+            "WARNING: MAMBA_HELION_AUTOTUNE not in {'1','2'}; existing configs are "
+            "loaded and nothing is autotuned."
+        )
+    print(
+        f"MAMBA_USE_HELION={os.environ['MAMBA_USE_HELION']} "
+        f"MAMBA_HELION_AUTOTUNE={os.environ['MAMBA_HELION_AUTOTUNE']} "
+        f"MAMBA_HELION_CONFIG_DIR={os.environ['MAMBA_HELION_CONFIG_DIR']} "
+        f"HELION_AUTOTUNE_PRECOMPILE={os.environ['HELION_AUTOTUNE_PRECOMPILE']}"
+    )
+
+    # Importing this module already pulled in the mamba_ssm package, so deferring
+    # the mixer import is not by itself proof that nothing read MAMBA_USE_HELION
+    # too early. Calling the @cache'd helper here either caches the right answer
+    # or reveals a stale cached one, instead of quietly running on Triton.
+    from mamba_ssm.ops.helion.dispatch import use_helion
+
+    assert use_helion(), (
+        "dispatch.use_helion() is already cached False -- something read "
+        "MAMBA_USE_HELION before configure_env(); the harness would run on Triton "
+        "and autotune nothing"
+    )
+
+
+def main():
+    assert torch.cuda.is_available(), "CUDA is required."
+    args = parse_args()
+    configure_env(args)
+
+    # Imported only now: this pulls in ssd_combined / layernorm_gated, which reach
+    # the @cache'd dispatch helpers.
+    from mamba_ssm.ops.helion.mamba_mixer_min import MambaMixerMin
+
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed_all(args.seed)
+    torch.cuda.set_device(0)
+
+    # Do not initialize torch.distributed here. These kernels do not declare
+    # distributed intent and normally use ordinary CUDA tensors, so Helion 1.4 treats
+    # them as ordinary kernels even in distributed jobs. The config helper enables
+    # Helion's timeout-protected subprocess benchmark path for autotuning.
+
+    dtype = _DTYPE_MAP[args.dtype]
+    mixer = MambaMixerMin(
+        d_model=args.hidden_size,
+        nheads=args.mamba_num_heads,
+        headdim=args.mamba_head_dim,
+        d_state=args.mamba_state_dim,
+        ngroups=args.mamba_num_groups,
+        chunk_size=args.chunk_size,
+        device="cuda",
+        dtype=dtype,
+    )
+    mixer.train()
+
+    print(
+        f"d_model={args.hidden_size} nheads={args.mamba_num_heads} headdim={args.mamba_head_dim} "
+        f"d_state={args.mamba_state_dim} ngroups={args.mamba_num_groups} "
+        f"chunk_size={args.chunk_size} seq_len={args.seq_len} batch={args.batch_size} dtype={args.dtype}"
+    )
+    seq_idx = build_seq_idx(args)
+    print("running one fwd+bwd to trigger Helion autotune for any missing configs...")
+
+    x = torch.randn(
+        args.seq_len,
+        args.batch_size,
+        args.hidden_size,
+        device="cuda",
+        dtype=dtype,
+        requires_grad=True,
+    )
+    out = mixer(x, seq_idx=seq_idx)
+    (out * torch.randn_like(out)).sum().backward()
+    torch.cuda.synchronize()
+
+    print("verifying every dispatched kernel has a config...")
+    verify_configs(Path(os.environ["MAMBA_HELION_CONFIG_DIR"]))
+    print("done.")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/helion/bench_mamba_mixer.py b/benchmarks/helion/bench_mamba_mixer.py
new file mode 100644
index 000000000..f4193dbb7
--- /dev/null
+++ b/benchmarks/helion/bench_mamba_mixer.py
@@ -0,0 +1,397 @@
+"""Compare Triton and Helion SSD kernel speed via torch profiler traces.
+
+Runs :class:`MambaMixerMin` (the same harness :mod:`autotune_mamba_mixer` uses,
+so the shapes and the kernel set match production) through fwd+bwd under
+``torch.profiler``, exports a chrome trace, and attributes every device kernel to
+the dispatch stage that launched it. Doing it per stage rather than per kernel
+name is the point: the two backends emit differently named kernels, so
+``chunk_scan_fwd`` on Triton and on Helion are only comparable through the
+dispatch wrapper they share.
+
+    python benchmarks/helion/bench_mamba_mixer.py                  # both backends, then compare
+    python benchmarks/helion/bench_mamba_mixer.py --backend helion # one trace only
+    python benchmarks/helion/bench_mamba_mixer.py --compare a.json b.json
+
+The default mode runs each backend in its own subprocess because
+``dispatch.use_helion()`` is ``@cache``d -- the first read of ``MAMBA_USE_HELION``
+pins the whole process to one backend. Autotuning is disabled (``MAMBA_HELION_AUTOTUNE``
+is removed from the environment), so a missing config in ``MAMBA_HELION_CONFIG_DIR``
+fails the run exactly as it would at runtime; generate configs with
+``autotune_mamba_mixer`` first.
+"""
+
+import argparse
+import bisect
+import json
+import os
+import subprocess
+import sys
+from collections import defaultdict
+from functools import wraps
+from pathlib import Path
+
+import torch
+
+# Sibling script, importable because it sits next to this one; both are run by path.
+from autotune_mamba_mixer import _DTYPE_MAP, build_seq_idx
+
+_REPO_ROOT = Path(__file__).resolve().parents[2]
+
+_STAGE_PREFIX = "ssd::"
+_ITER_MARK = "bench::iter"
+# Chrome-trace categories: device work, the launch APIs that correlate to it, and
+# record_function markers.
+_DEVICE_CATS = {"kernel", "gpu_memcpy", "gpu_memset"}
+_LAUNCH_CATS = {"cuda_runtime", "cuda_driver"}
+
+
+def parse_args():
+    p = argparse.ArgumentParser(description=__doc__.split("\n")[0])
+    # Dimensions -- keep in sync with autotune_mamba_mixer so a trace matches the
+    # configs that harness produced.
+    p.add_argument("--hidden-size", type=int, default=2688)
+    p.add_argument("--mamba-num-heads", type=int, default=64)
+    p.add_argument("--mamba-head-dim", type=int, default=64)
+    p.add_argument("--mamba-state-dim", type=int, default=128)
+    p.add_argument("--mamba-num-groups", type=int, default=8)
+    p.add_argument("--chunk-size", type=int, default=128)
+    p.add_argument("--seq-len", type=int, default=8192)
+    p.add_argument("--batch-size", type=int, default=1)
+    p.add_argument("--dtype", choices=["fp32", "bf16", "fp16"], default="bf16")
+    p.add_argument("--seed", type=int, default=0)
+    p.add_argument("--sequence-packing", action="store_true")
+    p.add_argument("--packed-num-seqs", type=int, default=4)
+    p.add_argument("--packed-seqlens", type=str, default=None)
+
+    p.add_argument("--warmup", type=int, default=3)
+    p.add_argument("--iters", type=int, default=10)
+    p.add_argument(
+        "--backend",
+        choices=["triton", "helion"],
+        default=None,
+        help="Profile this backend only. Omit to run both in subprocesses and compare.",
+    )
+    p.add_argument(
+        "--trace-dir",
+        type=Path,
+        default=_REPO_ROOT / "logs",
+        help="Where the chrome traces are written.",
+    )
+    p.add_argument(
+        "--compare",
+        type=Path,
+        nargs=2,
+        metavar=("TRITON_JSON", "HELION_JSON"),
+        default=None,
+        help="Skip profiling; report on two existing traces.",
+    )
+    return p.parse_args()
+
+
+def annotate_dispatch():
+    """Wrap the dispatch entry points in ``record_function`` markers.
+
+    Patches both the Triton namespace and the Helion getters in each consumer
+    module, so the same instrumentation works whichever backend
+    ``_ssd_kernel_impls()`` selects. Both call sites resolve these names at call
+    time (``ssd_combined._ssd_kernel_impls`` and the ternary in
+    ``LayerNormFn.backward``), which is what makes patching the module globals
+    enough.
+    """
+    from types import SimpleNamespace
+
+    from mamba_ssm.ops.triton import layernorm_gated, ssd_combined
+
+    def marked(name, fn):
+        @wraps(fn)
+        def wrapper(*args, **kwargs):
+            with torch.profiler.record_function(f"{_STAGE_PREFIX}{name}"):
+                return fn(*args, **kwargs)
+
+        return wrapper
+
+    kernels = ssd_combined._ssd_kernel_impls()
+    wrapped = SimpleNamespace(
+        **{name: marked(name, fn) for name, fn in vars(kernels).items()}
+    )
+    ssd_combined._TRITON_SSD_KERNELS = wrapped
+    ssd_combined.get_helion_ssd_kernels = lambda: wrapped
+
+    layer_norm_bwd = (
+        layernorm_gated.get_helion_layer_norm_bwd()
+        if layernorm_gated.use_helion()
+        else layernorm_gated._layer_norm_bwd
+    )
+    wrapped_ln = marked("layer_norm_bwd", layer_norm_bwd)
+    layernorm_gated._layer_norm_bwd = wrapped_ln
+    layernorm_gated.get_helion_layer_norm_bwd = lambda: wrapped_ln
+
+    return sorted(vars(wrapped)) + ["layer_norm_bwd"]
+
+
+def profile_backend(args, backend, trace_path):
+    """Profile ``args.iters`` fwd+bwd iterations and write a chrome trace."""
+    os.environ["MAMBA_USE_HELION"] = "1" if backend == "helion" else "0"
+    # Never autotune while benchmarking: it would both dwarf the measurement and
+    # install a config tuned for whatever shapes happen to be here.
+    os.environ.pop("MAMBA_HELION_AUTOTUNE", None)
+    os.environ.setdefault(
+        "MAMBA_HELION_CONFIG_DIR", str(Path(__file__).resolve().parent / "configs")
+    )
+
+    from mamba_ssm.ops.helion.dispatch import use_helion
+    from mamba_ssm.ops.helion.mamba_mixer_min import MambaMixerMin
+
+    assert use_helion() == (backend == "helion"), (
+        f"dispatch.use_helion()={use_helion()} does not match --backend {backend}; "
+        "something read MAMBA_USE_HELION before this ran"
+    )
+
+    torch.manual_seed(args.seed)
+    torch.cuda.manual_seed_all(args.seed)
+    torch.cuda.set_device(0)
+
+    dtype = _DTYPE_MAP[args.dtype]
+    mixer = MambaMixerMin(
+        d_model=args.hidden_size,
+        nheads=args.mamba_num_heads,
+        headdim=args.mamba_head_dim,
+        d_state=args.mamba_state_dim,
+        ngroups=args.mamba_num_groups,
+        chunk_size=args.chunk_size,
+        device="cuda",
+        dtype=dtype,
+    )
+    mixer.train()
+    seq_idx = build_seq_idx(args)
+    stages = annotate_dispatch()
+    print(f"backend={backend} instrumented stages: {', '.join(stages)}")
+
+    x = torch.randn(
+        args.seq_len,
+        args.batch_size,
+        args.hidden_size,
+        device="cuda",
+        dtype=dtype,
+        requires_grad=True,
+    )
+    # Fixed grad-output so no RNG kernels land inside the timed region.
+    dout = torch.randn_like(x)
+
+    def step():
+        mixer.zero_grad(set_to_none=True)
+        x.grad = None
+        mixer(x, seq_idx=seq_idx).backward(dout)
+
+    for _ in range(args.warmup):
+        step()
+    torch.cuda.synchronize()
+
+    with torch.profiler.profile(
+        activities=[
+            torch.profiler.ProfilerActivity.CPU,
+            torch.profiler.ProfilerActivity.CUDA,
+        ]
+    ) as prof:
+        for _ in range(args.iters):
+            with torch.profiler.record_function(_ITER_MARK):
+                step()
+        torch.cuda.synchronize()
+
+    trace_path.parent.mkdir(parents=True, exist_ok=True)
+    prof.export_chrome_trace(str(trace_path))
+    print(f"backend={backend} trace -> {trace_path}")
+
+
+def read_trace(path):
+    """Return ``(per-iteration stage totals in ms, diagnostics)`` for a trace.
+
+    Device work is tied to a stage through the launch API call that produced it:
+    kernels carry a ``correlation`` id shared with their ``cudaLaunchKernel`` /
+    ``cuLaunchKernel`` event, and that launch sits inside the ``record_function``
+    interval on the launching thread. Matching on the marker interval (rather
+    than on ``External id``, which points at the innermost aten op) keeps the
+    attribution correct for the backward pass too, where the launches happen on
+    the autograd thread.
+    """
+    with open(path) as f:
+        events = json.load(f)["traceEvents"]
+
+    markers = defaultdict(list)  # tid -> [(ts, end, name)]
+    iters = 0
+    launch_stage = {}  # correlation -> stage
+    device = []  # (correlation, dur, name)
+    for ev in events:
+        if ev.get("ph") != "X":
+            continue
+        cat, name = ev.get("cat"), ev.get("name", "")
+        if cat == "user_annotation":
+            if name == _ITER_MARK:
+                iters += 1
+            elif name.startswith(_STAGE_PREFIX):
+                markers[ev["tid"]].append(
+                    (ev["ts"], ev["ts"] + ev["dur"], name[len(_STAGE_PREFIX) :])
+                )
+        elif cat in _LAUNCH_CATS:
+            corr = ev.get("args", {}).get("correlation")
+            if corr is not None:
+                launch_stage[corr] = (ev["tid"], ev["ts"])
+        elif cat in _DEVICE_CATS:
+            device.append(
+                (ev.get("args", {}).get("correlation"), ev["dur"], name)
+            )
+
+    assert iters, f"no {_ITER_MARK} markers in {path} -- was it written by this script?"
+    for tid in markers:
+        markers[tid].sort()
+
+    def stage_of(corr):
+        launch = launch_stage.get(corr)
+        if launch is None:
+            return None
+        tid, ts = launch
+        intervals = markers.get(tid)
+        if not intervals:
+            return None
+        # Rightmost marker starting at or before the launch; markers of interest
+        # never nest, so containment in that one decides it.
+        i = bisect.bisect_right(intervals, (ts, float("inf"), "")) - 1
+        if i >= 0 and intervals[i][1] >= ts:
+            return intervals[i][2]
+        return None
+
+    stage_ms = defaultdict(float)
+    kernel_ms = defaultdict(lambda: defaultdict(float))
+    unattributed = 0.0
+    for corr, dur, name in device:
+        stage = stage_of(corr)
+        if stage is None:
+            unattributed += dur
+        else:
+            stage_ms[stage] += dur
+            kernel_ms[stage][name] += dur
+
+    scale = 1e-3 / iters  # us total -> ms per iteration
+    return (
+        {k: v * scale for k, v in stage_ms.items()},
+        {s: {n: v * scale for n, v in ks.items()} for s, ks in kernel_ms.items()},
+        {
+            "iters": iters,
+            "other_ms": unattributed * scale,
+            "total_ms": (sum(stage_ms.values()) + unattributed) * scale,
+        },
+    )
+
+
+def generated_only(kernels):
+    """Sum only the backend's own generated kernel(s) from a stage's kernels.
+
+    A Triton or Helion kernel is launched under its Python function name
+    (``_chunk_scan_fwd_kernel``, ``_helion__chunk_scan_fwd_kernel``), so its trace
+    name is a valid identifier; the helpers a stage also launches are not --
+    aten/cutlass kernels arrive as demangled C++ signatures (``void
+    at::native::...``) and memset/memcpy as ``Memset (Device)``. Dropping them
+    isolates the kernel comparison from the different amount of work the two
+    backends leave to aten (e.g. Triton reduces dD with a separate
+    ``reduce_kernel`` where Helion uses in-kernel atomics).
+    """
+    return sum(ms for name, ms in kernels.items() if name.isidentifier())
+
+
+def report(labelled_paths):
+    """Print the per-stage comparison table across the given traces."""
+    traces = [(label, *read_trace(path)) for label, path in labelled_paths]
+    for label, _, _, diag in traces:
+        print(
+            f"{label:>8}: {diag['iters']} iters, device time/iter "
+            f"{diag['total_ms']:.3f} ms (SSD+LN stages "
+            f"{diag['total_ms'] - diag['other_ms']:.3f} ms, "
+            f"other {diag['other_ms']:.3f} ms)"
+        )
+
+    stages = sorted(
+        {s for _, stage_ms, _, _ in traces for s in stage_ms},
+        key=lambda s: -max(stage_ms.get(s, 0.0) for _, stage_ms, _, _ in traces),
+    )
+    # In --compare mode the labels are file stems; keep the distinguishing tail so
+    # they still fit one numeric column.
+    labels = [label[-11:] for label, *_ in traces]
+    width = max(len(s) for s in stages) + 2
+    paired = len(traces) == 2
+    group = 12 * len(labels) + (9 if paired else 0)
+
+    def row(name, whole, kernel):
+        cells = "".join(f"{v:>12.3f}" for v in whole)
+        if paired:
+            cells += f"{(whole[0] / whole[1] if whole[1] else float('nan')):>8.2f}x"
+        cells += "   " + "".join(f"{v:>12.3f}" for v in kernel)
+        if paired:
+            cells += f"{(kernel[0] / kernel[1] if kernel[1] else float('nan')):>8.2f}x"
+        return f"{name:<{width}}{cells}"
+
+    groups = f"\n{'':<{width}}{'stage device time':^{group}}   {'generated kernel only':^{group}}"
+    header = f"{'stage':<{width}}" + 2 * (
+        "".join(f"{lb:>12}" for lb in labels) + (f"{'speedup':>9}" if paired else "") + "   "
+    )
+    print(groups)
+    print(header)
+    print("-" * len(header))
+    for stage in stages:
+        print(
+            row(
+                stage,
+                [stage_ms.get(stage, 0.0) for _, stage_ms, _, _ in traces],
+                [generated_only(km.get(stage, {})) for _, _, km, _ in traces],
+            )
+        )
+    print(
+        row(
+            "TOTAL (stages)",
+            [sum(stage_ms.values()) for _, stage_ms, _, _ in traces],
+            [sum(generated_only(ks) for ks in km.values()) for _, _, km, _ in traces],
+        )
+    )
+    print("\nall times are ms of device time per iteration; "
+          "speedup = first / second (>1 means the second is faster). "
+          "'stage device time' includes the aten helper kernels each stage launches; "
+          "'generated kernel only' is just the Triton/Helion kernel.")
+
+    print("\nper-kernel breakdown:")
+    for label, _, kernel_ms, _ in traces:
+        print(f"  [{label}]")
+        for stage in stages:
+            for name, ms in sorted(
+                kernel_ms.get(stage, {}).items(), key=lambda kv: -kv[1]
+            ):
+                print(f"    {stage:<32} {ms:8.3f}  {name}")
+
+
+def run_both(args):
+    """Profile each backend in its own subprocess, then compare the traces."""
+    paths = {}
+    for backend in ("triton", "helion"):
+        path = args.trace_dir / f"bench_mamba_mixer_{backend}.json"
+        cmd = [sys.executable, __file__, *sys.argv[1:], "--backend", backend]
+        print(f"$ {' '.join(cmd)}")
+        subprocess.run(cmd, check=True)
+        paths[backend] = path
+    report([("triton", paths["triton"]), ("helion", paths["helion"])])
+
+
+def main():
+    assert torch.cuda.is_available(), "CUDA is required."
+    args = parse_args()
+    if args.compare:
+        report([(p.stem, p) for p in args.compare])
+    elif args.backend:
+        profile_backend(
+            args,
+            args.backend,
+            args.trace_dir / f"bench_mamba_mixer_{args.backend}.json",
+        )
+    else:
+        run_both(args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/README.md b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/README.md
new file mode 100644
index 000000000..3e69e8385
--- /dev/null
+++ b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/README.md
@@ -0,0 +1,17 @@
+# b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k
+
+Example Helion configs, generated with:
+
+```sh
+MAMBA_HELION_CONFIG_DIR=$PWD/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k \
+    python benchmarks/helion/autotune_mamba_mixer.py
+```
+
+i.e. the harness defaults: Nemotron3 Nano dims (hidden 2688, 64 heads x 64 head
+dim, state 128, 8 groups, chunk 128), seq_len 8192, batch 1, bf16, unpacked path.
+
+Environment: NVIDIA B200, helion 1.4.0, triton 3.6.0,
+torch 2.11.0a0+a6c236b9fd.nv26.03.
+
+A config is specialized to all of the above, so reuse it only on a matching
+setup; otherwise re-autotune into a directory of your own.
diff --git a/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_bwd_dC_kernel.json b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_bwd_dC_kernel.json
new file mode 100644
index 000000000..7f3b590bd
--- /dev/null
+++ b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_bwd_dC_kernel.json
@@ -0,0 +1,56 @@
+{
+  "block_sizes": [
+    128,
+    128
+  ],
+  "loop_orders": [
+    [
+      4,
+      1,
+      5,
+      3,
+      0,
+      2
+    ]
+  ],
+  "l2_groupings": [
+    4
+  ],
+  "range_unroll_factors": [
+    0,
+    2
+  ],
+  "range_warp_specializes": [
+    true,
+    null
+  ],
+  "range_multi_buffers": [
+    false,
+    true
+  ],
+  "range_flattens": [
+    null,
+    null
+  ],
+  "load_eviction_policies": [
+    "first",
+    "last",
+    "last",
+    "last"
+  ],
+  "num_warps": 8,
+  "num_stages": 4,
+  "indexing": [
+    "pointer",
+    "tensor_descriptor",
+    "tensor_descriptor",
+    "pointer",
+    "tensor_descriptor"
+  ],
+  "atomic_indexing": [
+    "pointer"
+  ],
+  "pid_type": "persistent_interleaved",
+  "num_sm_multiplier": 1,
+  "nheads_per_program": 8
+}
\ No newline at end of file
diff --git a/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_bwd_ddAcs_stable_kernel.json b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_bwd_ddAcs_stable_kernel.json
new file mode 100644
index 000000000..33af4efb2
--- /dev/null
+++ b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_bwd_ddAcs_stable_kernel.json
@@ -0,0 +1,55 @@
+{
+  "block_sizes": [
+    32,
+    16
+  ],
+  "loop_orders": [
+    [
+      0,
+      1,
+      2,
+      3
+    ]
+  ],
+  "l2_groupings": [
+    1
+  ],
+  "range_unroll_factors": [
+    0,
+    0
+  ],
+  "range_warp_specializes": [
+    null,
+    null
+  ],
+  "range_num_stages": [],
+  "range_multi_buffers": [
+    null,
+    false
+  ],
+  "range_flattens": [
+    null,
+    null
+  ],
+  "load_eviction_policies": [
+    "",
+    "",
+    "",
+    "",
+    "",
+    ""
+  ],
+  "num_warps": 1,
+  "num_stages": 1,
+  "indexing": [
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer"
+  ],
+  "atomic_indexing": [],
+  "pid_type": "flat"
+}
\ No newline at end of file
diff --git a/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_bwd_dstates_kernel.json b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_bwd_dstates_kernel.json
new file mode 100644
index 000000000..c42beff6c
--- /dev/null
+++ b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_bwd_dstates_kernel.json
@@ -0,0 +1,51 @@
+{
+  "block_sizes": [
+    64,
+    64,
+    32
+  ],
+  "loop_orders": [
+    [
+      2,
+      1,
+      3,
+      4,
+      0
+    ]
+  ],
+  "l2_groupings": [
+    8
+  ],
+  "range_unroll_factors": [
+    0,
+    0
+  ],
+  "range_warp_specializes": [
+    null,
+    null
+  ],
+  "range_num_stages": [],
+  "range_multi_buffers": [
+    null,
+    null
+  ],
+  "range_flattens": [
+    null,
+    false
+  ],
+  "load_eviction_policies": [
+    "first",
+    "last",
+    ""
+  ],
+  "num_warps": 1,
+  "num_stages": 3,
+  "indexing": [
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer"
+  ],
+  "atomic_indexing": [],
+  "pid_type": "flat"
+}
\ No newline at end of file
diff --git a/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_chunk_state_bwd_dx_kernel.json b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_chunk_state_bwd_dx_kernel.json
new file mode 100644
index 000000000..12d877fa6
--- /dev/null
+++ b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_chunk_state_bwd_dx_kernel.json
@@ -0,0 +1,71 @@
+{
+  "block_sizes": [
+    64,
+    64,
+    32
+  ],
+  "loop_orders": [
+    [
+      1,
+      2,
+      4,
+      0,
+      3
+    ]
+  ],
+  "l2_groupings": [
+    4
+  ],
+  "range_unroll_factors": [
+    4,
+    1
+  ],
+  "range_warp_specializes": [
+    false,
+    false
+  ],
+  "range_multi_buffers": [
+    true,
+    false
+  ],
+  "range_flattens": [
+    false,
+    false
+  ],
+  "load_eviction_policies": [
+    "",
+    "",
+    "",
+    "last",
+    "first",
+    "last",
+    "",
+    "last",
+    "",
+    "last",
+    "first"
+  ],
+  "num_warps": 4,
+  "num_stages": 3,
+  "indexing": [
+    "pointer",
+    "tensor_descriptor",
+    "pointer",
+    "tensor_descriptor",
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer",
+    "tensor_descriptor",
+    "pointer",
+    "pointer"
+  ],
+  "atomic_indexing": [
+    "pointer",
+    "pointer"
+  ],
+  "pid_type": "persistent_interleaved",
+  "num_sm_multiplier": 8,
+  "maxnreg": 128
+}
\ No newline at end of file
diff --git a/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_fwd_kernel.json b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_fwd_kernel.json
new file mode 100644
index 000000000..ab9fc3d87
--- /dev/null
+++ b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_scan_fwd_kernel.json
@@ -0,0 +1,69 @@
+{
+  "block_sizes": [
+    64,
+    64,
+    16
+  ],
+  "loop_orders": [
+    [
+      1,
+      0,
+      2,
+      3,
+      4
+    ]
+  ],
+  "l2_groupings": [
+    4
+  ],
+  "range_unroll_factors": [
+    0,
+    3
+  ],
+  "range_warp_specializes": [
+    null,
+    null
+  ],
+  "range_num_stages": [
+    0,
+    4
+  ],
+  "range_multi_buffers": [
+    null,
+    false
+  ],
+  "range_flattens": [
+    null,
+    false
+  ],
+  "load_eviction_policies": [
+    "",
+    "first",
+    "last",
+    "last",
+    "",
+    "",
+    "last",
+    "",
+    "last",
+    "first"
+  ],
+  "num_warps": 2,
+  "num_stages": 3,
+  "indexing": [
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer",
+    "tensor_descriptor",
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer"
+  ],
+  "atomic_indexing": [],
+  "pid_type": "flat"
+}
\ No newline at end of file
diff --git a/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_state_bwd_db_kernel.json b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_state_bwd_db_kernel.json
new file mode 100644
index 000000000..a2fd13313
--- /dev/null
+++ b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_state_bwd_db_kernel.json
@@ -0,0 +1,59 @@
+{
+  "block_sizes": [
+    128,
+    64
+  ],
+  "loop_orders": [
+    [
+      4,
+      2,
+      5,
+      0,
+      3,
+      1
+    ]
+  ],
+  "l2_groupings": [
+    2
+  ],
+  "range_unroll_factors": [
+    0,
+    1
+  ],
+  "range_warp_specializes": [
+    null,
+    false
+  ],
+  "range_num_stages": [],
+  "range_multi_buffers": [
+    null,
+    false
+  ],
+  "range_flattens": [
+    null,
+    true
+  ],
+  "load_eviction_policies": [
+    "last",
+    "first",
+    "last",
+    "",
+    "first",
+    ""
+  ],
+  "num_warps": 4,
+  "num_stages": 3,
+  "indexing": [
+    "pointer",
+    "pointer",
+    "pointer",
+    "tensor_descriptor",
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer"
+  ],
+  "atomic_indexing": [],
+  "pid_type": "flat",
+  "nheads_per_program": 8
+}
\ No newline at end of file
diff --git a/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_state_fwd_kernel.json b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_state_fwd_kernel.json
new file mode 100644
index 000000000..7459355eb
--- /dev/null
+++ b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_chunk_state_fwd_kernel.json
@@ -0,0 +1,55 @@
+{
+  "block_sizes": [
+    64,
+    128,
+    64
+  ],
+  "loop_orders": [
+    [
+      2,
+      0,
+      3,
+      4,
+      1
+    ]
+  ],
+  "l2_groupings": [
+    1
+  ],
+  "range_unroll_factors": [
+    0,
+    4
+  ],
+  "range_warp_specializes": [
+    null,
+    null
+  ],
+  "range_num_stages": [],
+  "range_multi_buffers": [
+    null,
+    null
+  ],
+  "range_flattens": [
+    null,
+    true
+  ],
+  "load_eviction_policies": [
+    "first",
+    "first",
+    "",
+    "",
+    ""
+  ],
+  "num_warps": 4,
+  "num_stages": 2,
+  "indexing": [
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer"
+  ],
+  "atomic_indexing": [],
+  "pid_type": "flat"
+}
\ No newline at end of file
diff --git a/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_layer_norm_bwd_kernel.json b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_layer_norm_bwd_kernel.json
new file mode 100644
index 000000000..e786f0cdb
--- /dev/null
+++ b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_layer_norm_bwd_kernel.json
@@ -0,0 +1,51 @@
+{
+  "block_sizes": [],
+  "loop_orders": [
+    [
+      1,
+      0
+    ]
+  ],
+  "l2_groupings": [
+    8
+  ],
+  "range_unroll_factors": [
+    0,
+    4
+  ],
+  "range_warp_specializes": [
+    null,
+    false
+  ],
+  "range_num_stages": [
+    0,
+    4
+  ],
+  "range_multi_buffers": [
+    null,
+    true
+  ],
+  "range_flattens": [
+    null,
+    true
+  ],
+  "load_eviction_policies": [
+    "",
+    "first",
+    "first",
+    "last"
+  ],
+  "num_warps": 1,
+  "num_stages": 6,
+  "indexing": [
+    "pointer",
+    "pointer",
+    "tensor_descriptor",
+    "pointer",
+    "pointer",
+    "pointer"
+  ],
+  "atomic_indexing": [],
+  "pid_type": "flat",
+  "nrow_groups": 1113
+}
\ No newline at end of file
diff --git a/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_state_passing_bwd_kernel.json b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_state_passing_bwd_kernel.json
new file mode 100644
index 000000000..8350c9c80
--- /dev/null
+++ b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_state_passing_bwd_kernel.json
@@ -0,0 +1,60 @@
+{
+  "block_sizes": [
+    512,
+    4
+  ],
+  "loop_orders": [
+    [
+      1,
+      0,
+      2
+    ]
+  ],
+  "l2_groupings": [
+    2
+  ],
+  "range_unroll_factors": [
+    0,
+    3
+  ],
+  "range_warp_specializes": [
+    null,
+    false
+  ],
+  "range_num_stages": [],
+  "range_multi_buffers": [
+    null,
+    true
+  ],
+  "range_flattens": [
+    null,
+    true
+  ],
+  "load_eviction_policies": [
+    "last",
+    "first",
+    "first",
+    "first",
+    "",
+    "first"
+  ],
+  "num_warps": 1,
+  "num_stages": 8,
+  "indexing": [
+    "pointer",
+    "pointer",
+    "pointer",
+    "pointer",
+    "tensor_descriptor",
+    "pointer",
+    "pointer",
+    "pointer",
+    "tensor_descriptor",
+    "pointer",
+    "pointer",
+    "tensor_descriptor",
+    "pointer"
+  ],
+  "atomic_indexing": [],
+  "pid_type": "flat"
+}
\ No newline at end of file
diff --git a/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_state_passing_fwd_kernel_v2.json b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_state_passing_fwd_kernel_v2.json
new file mode 100644
index 000000000..fcc51bb54
--- /dev/null
+++ b/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/_state_passing_fwd_kernel_v2.json
@@ -0,0 +1,48 @@
+{
+  "block_sizes": [
+    256,
+    4
+  ],
+  "loop_orders": [
+    [
+      0,
+      2,
+      1
+    ]
+  ],
+  "l2_groupings": [
+    8
+  ],
+  "range_unroll_factors": [
+    0,
+    3
+  ],
+  "range_warp_specializes": [
+    null,
+    false
+  ],
+  "range_num_stages": [],
+  "range_multi_buffers": [
+    null,
+    true
+  ],
+  "range_flattens": [
+    null,
+    false
+  ],
+  "load_eviction_policies": [
+    "first",
+    "first"
+  ],
+  "num_warps": 1,
+  "num_stages": 4,
+  "indexing": [
+    "pointer",
+    "pointer",
+    "tensor_descriptor",
+    "tensor_descriptor",
+    "pointer"
+  ],
+  "atomic_indexing": [],
+  "pid_type": "flat"
+}
\ No newline at end of file
diff --git a/benchmarks/helion/run_autotune.sh b/benchmarks/helion/run_autotune.sh
new file mode 100755
index 000000000..5740c3115
--- /dev/null
+++ b/benchmarks/helion/run_autotune.sh
@@ -0,0 +1,22 @@
+#!/usr/bin/env bash
+# Launch the Helion autotune harness detached via nohup, logging to its own file.
+# Any arguments are forwarded to the harness, e.g.:
+#
+#   ./benchmarks/helion/run_autotune.sh --precompile fork
+#   LOG_DIR=/data/logs ./benchmarks/helion/run_autotune.sh --sequence-packing
+#
+# Env: PYTHON (interpreter, default python), LOG_DIR (default <repo>/logs).
+set -euo pipefail
+
+repo_root=$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)
+log_dir=${LOG_DIR:-"$repo_root/logs"}
+mkdir -p "$log_dir"
+log_file="$log_dir/autotune_mamba_mixer_$(date +%Y%m%d_%H%M%S).log"
+
+cd "$repo_root"
+# -u so the log is not block-buffered while the run is still in progress.
+nohup "${PYTHON:-python}" -u benchmarks/helion/autotune_mamba_mixer.py "$@" \
+    >"$log_file" 2>&1 &
+
+echo "pid $! logging to $log_file"
+echo "follow with: tail -f $log_file"
diff --git a/mamba_ssm/ops/helion/README.md b/mamba_ssm/ops/helion/README.md
new file mode 100644
index 000000000..5cadd2e87
--- /dev/null
+++ b/mamba_ssm/ops/helion/README.md
@@ -0,0 +1,250 @@
+# Helion kernel backend for the Mamba-2 SSD path
+
+An opt-in [Helion](https://github.com/pytorch/helion) implementation of the nine SSD
+kernels used by `mamba_split_conv1d_scan_combined` plus the gated-LayerNorm backward.
+The Triton kernels remain the default and the numerical reference; nothing changes
+unless `MAMBA_USE_HELION=1` is set.
+
+## What this branch adds
+
+Runtime code lives in this directory; the offline tooling and the example configs live
+outside the installed package, under `benchmarks/helion/`, so nothing autotunes or writes
+into site-packages.
+
+| File | Role |
+| --- | --- |
+| `dispatch.py` | The only import of this package from production code. Resolves the backend once per process and lazily imports the kernels **after** `MAMBA_USE_HELION=1` is confirmed, so an install without `helion` still works. |
+| `utils.py` | `json_cached_autotune` / `ensure_helion_config`: resolve each kernel's config from `MAMBA_HELION_CONFIG_DIR/<kernel.name>.json` on first call, autotune only when explicitly asked. |
+| `_chunk_state_fwd.py`, `_state_passing_fwd.py`, `_chunk_scan_fwd.py` | Forward stages. |
+| `_chunk_scan_bwd_dstates.py`, `_state_passing_bwd.py`, `_chunk_scan_chunk_state_bwd_dx.py`, `_chunk_state_bwd_db.py`, `_chunk_scan_bwd_dC.py`, `_chunk_scan_bwd_ddAcs_stable.py` | Backward stages. |
+| `_layer_norm_bwd.py` | Gated LayerNorm/RMSNorm backward (`layernorm_gated`). |
+| `mamba_mixer_min.py` | `MambaMixerMin`, a Megatron-free replica of Megatron-LM's `MambaMixer._ssm_training` layout. Imported by the harness, the benchmark and the tests, which is why it stays in the package. |
+
+Outside the package:
+
+| Path | Role |
+| --- | --- |
+| `benchmarks/helion/autotune_mamba_mixer.py` | Offline autotune harness — one fwd+bwd that reaches every kernel, then verifies each left a config behind. |
+| `benchmarks/helion/run_autotune.sh` | Runs the harness detached under `nohup`, logging to `logs/`. |
+| `benchmarks/helion/bench_mamba_mixer.py` | Triton-vs-Helion per-stage device-time comparison from `torch.profiler` traces. |
+| `benchmarks/helion/configs/` | Where the harness autotunes into, and where the checked-in example set lives (see below). Not shipped in the wheel. |
+| `tests/ops/helion/test_kernels.py` | Every kernel against its Triton reference, on production arguments. |
+
+Consumers in the Triton package:
+
+- [`ssd_combined.py`](../triton/ssd_combined.py) — `_ssd_kernel_impls()` returns either
+  `_TRITON_SSD_KERNELS` or `get_helion_ssd_kernels()`, both `SimpleNamespace`s with the
+  same nine attribute names; the fwd/bwd functions call through it.
+- [`layernorm_gated.py`](../triton/layernorm_gated.py) — `LayerNormFn.backward` picks
+  `get_helion_layer_norm_bwd()` over `_layer_norm_bwd`.
+
+Adding or renaming a kernel means keeping those two namespaces, `_SSD_KERNELS` in the
+test, and the per-kernel file in sync — `test_kernel_namespaces_match` guards exactly that.
+
+Requirements beyond the base install come from the `helion` extra (`helion==1.4.0` and
+`triton>=3.6`), plus `causal_conv1d` for the harness, benchmark and tests:
+
+```sh
+pip install -e '.[helion,causal-conv1d]' --no-build-isolation
+```
+
+`helion` is pinned exactly, not lower-bounded: a config is autotuned for one toolchain and
+Helion's config schema is still moving, so another version can need re-autotuning — or fail
+to load a config — even with the kernels unchanged. Raise the pin deliberately, together
+with a re-autotuned config set.
+
+## Validated environment
+
+Everything below — the checked-in config set, the test results and the benchmark numbers —
+was produced on a single node of:
+
+| | |
+| --- | --- |
+| GPU | NVIDIA B200 183GB (compute capability 10.0), driver 580.95.05; all runs single-GPU |
+| OS / Python | Linux 6.8.0 x86_64, glibc 2.39, CPython 3.12.3 |
+| torch | 2.11.0a0+a6c236b9fd.nv26.03 (NGC 26.03 container), CUDA 13.2 |
+| triton | 3.6.0 |
+| helion | 1.4.0 |
+| causal_conv1d | 1.6.2.post1 |
+| einops / pytest | 0.8.2 / 8.1.1 |
+
+Model shape for all of it, the autotune harness defaults (Nemotron3 Nano): `d_model` 2688,
+64 heads x 64 head dim, `d_state` 128, 8 groups, `chunk_size` 128, seq_len 8192, batch 1,
+bf16, `MambaMixerMin` un-fused layout.
+
+The kernels themselves are not Blackwell-specific, but a config set is (see
+[Autotuning configs](#autotuning-configs)): on any other GPU, toolchain version or model
+shape, re-autotune and re-run the tests before trusting the backend.
+
+## Enabling it
+
+```sh
+export MAMBA_USE_HELION=1
+export MAMBA_HELION_CONFIG_DIR=$PWD/benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k
+```
+
+| Env var | Meaning |
+| --- | --- |
+| `MAMBA_USE_HELION` | `1` selects the Helion kernels. Read once per process (`@cache`), so set it before importing anything that dispatches. |
+| `MAMBA_HELION_CONFIG_DIR` | Directory holding `<kernel>.json`. A missing config is an assertion failure at first call. |
+| `MAMBA_HELION_AUTOTUNE` | Offline only: `1` fills in missing configs (resumable), `2` forces a re-tune and overwrites. Unset at runtime. |
+
+Autotuning must never happen at runtime — not under CUDA graph capture, not in a
+distributed job — which is why a missing config fails loudly instead of tuning itself.
+The config loaded on first call becomes the kernel's single config, so every later
+specialization reuses it.
+
+Note: with `MAMBA_USE_HELION=1` and deterministic mode enabled, `dispatch.use_helion()`
+warns once — `chunk_scan_chunk_state_bwd_dx` (ddt, 1-D dD) and `chunk_scan_bwd_dC`
+(ddA_cumsum_prev) always reduce with atomic adds, so the backward is not bitwise
+reproducible. Use the Triton path when you need determinism.
+
+## Autotuning configs
+
+```sh
+python benchmarks/helion/autotune_mamba_mixer.py                    # unpacked, Nemotron3 Nano dims, seq 8k, bf16
+python benchmarks/helion/autotune_mamba_mixer.py --sequence-packing # packed (THD) path, needs --batch-size 1
+./benchmarks/helion/run_autotune.sh --precompile fork               # same, detached, logs to logs/
+```
+
+The harness sets `MAMBA_USE_HELION=1`, `MAMBA_HELION_AUTOTUNE=1` and
+`MAMBA_HELION_CONFIG_DIR=benchmarks/helion/configs` by default (explicit env vars win), runs one
+`MambaMixerMin` fwd+bwd, then asserts that every dispatched kernel left a
+`<kernel>.json` behind — a kernel this configuration never reached would otherwise
+surface only later, as a runtime `ensure_helion_config` assertion.
+
+Useful flags: the dimension flags (`--hidden-size`, `--mamba-num-heads`,
+`--mamba-head-dim`, `--mamba-state-dim`, `--mamba-num-groups`, `--chunk-size`,
+`--seq-len`, `--batch-size`, `--dtype`), the packing flags (`--sequence-packing`,
+`--packed-num-seqs`, `--packed-seqlens`), and `--precompile {spawn,fork,off}` — `spawn`
+precompiles each candidate config in a subprocess so a config that hard-crashes the CUDA
+context is skipped rather than aborting the whole run, at the cost of memory and
+cold-start time.
+
+**A config set is specialized to GPU, toolchain, dims, dtype and packing.** Give each
+variant its own directory; the naming scheme used here is
+`<gpu>-helion<minor>-triton<minor>-<model>-<dtype>-seq<len>[-packed]`, e.g. the
+checked-in [`benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/`](../../../benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k)
+(its `README.md` records the exact environment). Packed configs in particular *cannot*
+be reused: a config carries one `indexing` entry per load/store and the packed path adds
+the `seq_idx` loads, so an unpacked config fails to compile there.
+
+Configs are repo artifacts, deliberately not packaged: `benchmarks/` is excluded from the
+wheel, so a `pip install mamba-ssm` ships no config set and nothing can autotune into the
+installed package. Enabling the backend from an installed wheel therefore means bringing
+your own directory — autotuned from a source checkout, or copied from one — and pointing
+`MAMBA_HELION_CONFIG_DIR` at it.
+
+## Tests
+
+```sh
+pytest tests/ops/helion/test_kernels.py                       # uses the checked-in configs
+MAMBA_HELION_CONFIG_DIR=<unpacked configs> pytest tests/ops/helion/test_kernels.py
+MAMBA_HELION_TEST_PACKED=1 MAMBA_HELION_CONFIG_DIR=<packed configs> \
+    pytest tests/ops/helion/test_kernels.py
+pytest 'tests/ops/helion/test_kernels.py::test_ssd_kernel_against_triton[unpacked-chunk_scan_fwd]'
+```
+
+With `MAMBA_HELION_CONFIG_DIR` unset the tests fall back to the checked-in
+`benchmarks/helion/configs/b200-…-seq8k/` set, so a reviewer can run
+them with no setup. A config is a set of block sizes and indexing choices rather than
+anything GPU-specific, so that set usually compiles and passes on other hardware too — but
+if it does not, the failure carries the `autotune_mamba_mixer` command to generate your
+own. Autotuning never happens implicitly, in the tests as at runtime. The packed variant
+has no checked-in set and asks for one explicitly.
+
+One `MambaMixerMin` fwd+bwd runs with the pipeline on **Triton**, every dispatched kernel
+wrapped so both implementations see the same production arguments and their outputs are
+compared. The wrapper returns the Triton result, so every later stage still consumes
+exact production data and a difference is attributable to the kernel that produced it.
+Arguments are deliberately not rebuilt by hand — they are views into shared buffers
+(`x`/`B`/`C`/`dx` slice one packed `xBC`, `z` slices `zxbcdt`) and outputs of preceding
+stages; Helion compiles per stride, so a contiguous stand-in would test a specialization
+that never runs.
+
+Tolerances are one bf16 ulp for `rtol`, and an absolute floor of one ulp of the tensor's
+*own* max magnitude — these outputs span orders of magnitude, so any fixed atol is either
+impossibly tight for the large tensors or vacuous for the small ones.
+
+One process tests one variant (unpacked by default, packed under
+`MAMBA_HELION_TEST_PACKED=1`), because Helion pins a kernel's config on first call and the
+two variants need different config sets. `MAMBA_HELION_CONFIG_DIR` must already contain
+them: the fixture deletes `MAMBA_HELION_AUTOTUNE`, making a missing config a hard failure
+exactly as at runtime.
+
+On the environment above, with the checked-in config set, the unpacked variant is 11 passed
+(nine SSD kernels, the gated-norm backward, and the namespace guard) in ~34s — one module
+fixture runs the single fwd+bwd that all of them read.
+
+## Benchmark
+
+```sh
+python benchmarks/helion/bench_mamba_mixer.py                    # both backends in subprocesses, then compare
+python benchmarks/helion/bench_mamba_mixer.py --backend helion   # one trace only
+python benchmarks/helion/bench_mamba_mixer.py --compare logs/bench_mamba_mixer_triton.json logs/bench_mamba_mixer_helion.json
+```
+
+Runs `MambaMixerMin` fwd+bwd under `torch.profiler` (`--warmup`, `--iters`), writes a
+chrome trace per backend to `--trace-dir` (default `logs/`), and attributes every device
+kernel to the **dispatch stage** that launched it — per stage rather than per kernel name
+because the two backends emit differently named kernels, so `chunk_scan_fwd` on Triton and
+on Helion are only comparable through the dispatch wrapper they share. Attribution goes
+through each kernel's `correlation` id to its launch API call, and from there to the
+enclosing `record_function` interval, which keeps the backward pass (launched on the
+autograd thread) correct too.
+
+The report prints two columns per backend: *stage device time* (including the aten helper
+kernels a stage launches) and *generated kernel only* (just the Triton/Helion kernel) —
+the split matters because the backends leave different amounts of work to aten, e.g.
+Triton reduces dD with a separate `reduce_kernel` where Helion uses in-kernel atomics.
+
+Each backend runs in its own subprocess since `use_helion()` is cached per process, and
+autotuning is disabled while benchmarking (a missing config fails the run, as at runtime).
+Dimension flags mirror the autotune harness — keep them in sync with the config set you
+point at.
+
+### Measured example
+
+```sh
+CUDA_VISIBLE_DEVICES=1 python benchmarks/helion/bench_mamba_mixer.py
+```
+
+Harness defaults (warmup 3, 10 timed iterations) with the checked-in
+`benchmarks/helion/configs/b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k/` set, on
+the [environment above](#validated-environment).
+
+Device time per fwd+bwd iteration, in ms:
+
+| | Triton | Helion | Speedup |
+| --- | ---: | ---: | ---: |
+| whole iteration | 3.084 | 2.725 | 1.13x |
+| SSD + gated-norm stages | 1.295 | 0.935 | **1.39x** |
+| everything else (in/out projections, conv, aten) | 1.789 | 1.791 | 1.00x |
+
+Per stage — the left group is the stage's whole device time, the right group only the
+Triton/Helion kernel it generates, excluding the aten helpers the stage launches:
+
+| Stage | Triton | Helion | Speedup | Triton<br>(kernel) | Helion<br>(kernel) | Speedup<br>(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** |
+
+Speedup is Triton / Helion, so >1 means Helion is faster. Every stage is faster here, the
+gated LayerNorm backward by the largest margin; the untouched work around the stages is
+unchanged, as expected. The numbers are specific to this GPU, toolchain and config set —
+re-run the benchmark for any other combination rather than reading these as general
+figures. To re-print a table from traces already in `logs/`, without profiling again:
+
+```sh
+python benchmarks/helion/bench_mamba_mixer.py --compare \
+    logs/bench_mamba_mixer_triton.json logs/bench_mamba_mixer_helion.json
+```
diff --git a/mamba_ssm/ops/helion/__init__.py b/mamba_ssm/ops/helion/__init__.py
new file mode 100644
index 000000000..05fd2a179
--- /dev/null
+++ b/mamba_ssm/ops/helion/__init__.py
@@ -0,0 +1,8 @@
+"""Optional Helion kernels for Mamba SSM.
+
+The existing Triton kernels remain the default. Set ``MAMBA_USE_HELION=1`` and
+point ``MAMBA_HELION_CONFIG_DIR`` at matching pre-generated configs to opt in.
+Generate missing configs with::
+
+    python benchmarks/helion/autotune_mamba_mixer.py
+"""
diff --git a/mamba_ssm/ops/helion/_chunk_scan_bwd_dC.py b/mamba_ssm/ops/helion/_chunk_scan_bwd_dC.py
new file mode 100644
index 000000000..7510a2816
--- /dev/null
+++ b/mamba_ssm/ops/helion/_chunk_scan_bwd_dC.py
@@ -0,0 +1,181 @@
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+
+import helion
+import helion.language as hl
+from helion.autotuner import IntegerFragment
+
+from ..triton.ssd_chunk_scan import (
+    _chunk_scan_bwd_dC as _chunk_scan_bwd_dC_ref,
+)
+
+from .utils import json_cached_autotune
+
+
+@json_cached_autotune
+@helion.kernel(
+    autotune_accuracy_check=False,
+    autotune_baseline_fn=_chunk_scan_bwd_dC_ref,
+    ignore_warnings=[helion.exc.TensorOperationInWrapper],
+)
+def _chunk_scan_bwd_dC_kernel(
+    prev_states: torch.Tensor,
+    dA_cumsum: torch.Tensor,
+    dout: torch.Tensor,
+    seq_idx: Optional[torch.Tensor] = None,
+    C: Optional[torch.Tensor] = None,
+    ngroups: int = 1,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    batch, nchunks, nheads, headdim, dstate = prev_states.shape
+    _, seqlen, _, _ = dout.shape
+    _, _, _, chunk_size = dA_cumsum.shape
+    assert seqlen % chunk_size == 0
+    assert prev_states.shape == (batch, nchunks, nheads, headdim, dstate)
+    assert dA_cumsum.shape == (batch, nheads, nchunks, chunk_size)
+    assert dout.shape == (batch, seqlen, nheads, headdim)
+    if seq_idx is not None:
+        assert seq_idx.shape == (batch, seqlen)
+    if C is not None:
+        assert C.shape == (batch, seqlen, ngroups, dstate)
+        ddA_cumsum_prev = torch.zeros(
+            batch,
+            nheads,
+            nchunks,
+            chunk_size,
+            device=dout.device,
+            dtype=torch.float32,
+        )
+    else:
+        ddA_cumsum_prev = None
+
+    ngroups = hl.specialize(ngroups)
+    nheads_ngroups_ratio = nheads // ngroups
+    nheads_per_program = hl.register_tunable(
+        "nheads_per_program",
+        IntegerFragment(1, nheads_ngroups_ratio),
+    )
+    nsplits = helion.cdiv(nheads_ngroups_ratio, nheads_per_program)
+    dC = torch.empty(
+        batch,
+        seqlen,
+        nsplits,
+        ngroups,
+        dstate,
+        device=dout.device,
+        dtype=torch.float32,
+    )
+
+    HAS_DDA_CS = ddA_cumsum_prev is not None
+    HAS_SEQ_IDX = seq_idx is not None
+    block_m = hl.register_block_size(chunk_size)
+    block_n = hl.register_block_size(dstate)
+
+    for tile_c, tile_b, tile_s, tile_g, tile_m, tile_n in hl.tile(
+        [nchunks, batch, nsplits, ngroups, chunk_size, dstate],
+        block_size=[1, 1, 1, 1, block_m, block_n],
+    ):
+        acc = hl.zeros((block_m, block_n), dtype=torch.float32)
+        if HAS_DDA_CS:
+            c_local = C[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_m.index,
+                tile_g.begin,
+                tile_n,
+            ].to(torch.float32)
+        if HAS_SEQ_IDX:
+            if tile_c.id == 0:
+                seq_idx_prev_local = hl.zeros([], dtype=seq_idx.dtype)
+            else:
+                seq_idx_prev_local = seq_idx[
+                    tile_b.begin,
+                    tile_c.begin * chunk_size - 1,
+                ]
+            seq_idx_m_local = seq_idx[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_m.index,
+            ]
+
+        nheads_iter = min(
+            nheads_per_program,
+            nheads // ngroups - tile_s.id * nheads_per_program,
+        )
+        for tile_h in hl.tile(nheads_iter, block_size=1):
+            head_idx = (
+                tile_g.begin * (nheads // ngroups)
+                + tile_s.begin * nheads_per_program
+                + tile_h.begin
+            )
+            dout_local = dout[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_m.index,
+                head_idx,
+                :,
+            ]
+            prev_states_local = prev_states[
+                tile_b.begin,
+                tile_c.begin,
+                head_idx,
+                :,
+                tile_n,
+            ].to(dout.dtype)
+            dc = hl.dot(dout_local, prev_states_local)
+            dA_cs_m_local = dA_cumsum[
+                tile_b.begin,
+                head_idx,
+                tile_c.begin,
+                tile_m,
+            ].to(torch.float32)
+            scale = torch.exp(dA_cs_m_local)
+            if HAS_SEQ_IDX:
+                scale = torch.where(
+                    seq_idx_m_local == seq_idx_prev_local,
+                    scale,
+                    torch.zeros_like(scale),
+                )
+            dc *= scale[:, None]
+            if HAS_DDA_CS:
+                ddA_cs_local = torch.sum(dc * c_local, dim=1)
+                hl.atomic_add(
+                    ddA_cumsum_prev,
+                    [
+                        tile_b.begin,
+                        head_idx,
+                        tile_c.begin,
+                        tile_m,
+                    ],
+                    ddA_cs_local,
+                )
+            acc += dc
+
+        dC[
+            tile_b.begin,
+            tile_c.begin * chunk_size + tile_m.index,
+            tile_s.begin,
+            tile_g.begin,
+            tile_n,
+        ] = acc
+
+    dC = dC.sum(2)
+    return dC if C is None else (dC, ddA_cumsum_prev)
+
+
+def _chunk_scan_bwd_dC(
+    prev_states: torch.Tensor,
+    dA_cumsum: torch.Tensor,
+    dout: torch.Tensor,
+    seq_idx: Optional[torch.Tensor] = None,
+    C: Optional[torch.Tensor] = None,
+    ngroups: int = 1,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    return _chunk_scan_bwd_dC_kernel(
+        prev_states,
+        dA_cumsum,
+        dout,
+        seq_idx,
+        C,
+        ngroups,
+    )
diff --git a/mamba_ssm/ops/helion/_chunk_scan_bwd_ddAcs_stable.py b/mamba_ssm/ops/helion/_chunk_scan_bwd_ddAcs_stable.py
new file mode 100644
index 000000000..19acf65e5
--- /dev/null
+++ b/mamba_ssm/ops/helion/_chunk_scan_bwd_ddAcs_stable.py
@@ -0,0 +1,146 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+import helion
+import helion.language as hl
+
+from ..triton.ssd_chunk_scan import (
+    _chunk_scan_bwd_ddAcs_stable as _chunk_scan_bwd_ddAcs_stable_ref,
+)
+
+from .utils import json_cached_autotune
+
+
+@json_cached_autotune
+@helion.kernel(
+    autotune_accuracy_check=False,
+    autotune_baseline_fn=_chunk_scan_bwd_ddAcs_stable_ref,
+    ignore_warnings=[helion.exc.TensorOperationInWrapper],
+)
+def _chunk_scan_bwd_ddAcs_stable_kernel(
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    dA_cumsum: torch.Tensor,
+    dout: torch.Tensor,
+    cb: torch.Tensor,
+) -> torch.Tensor:
+    batch, seqlen, nheads, headdim = x.shape
+    _, _, nchunks, chunk_size = dt.shape
+    assert seqlen % chunk_size == 0
+    assert dt.shape == (batch, nheads, nchunks, chunk_size)
+    assert dout.shape == x.shape
+    assert dA_cumsum.shape == dt.shape
+    ngroups = cb.shape[2]
+    assert nheads % ngroups == 0
+    assert cb.shape == (batch, nchunks, ngroups, chunk_size, chunk_size)
+
+    BLOCK_SIZE_M_min = 32
+    ddA_cumsum = torch.zeros(
+        batch,
+        nheads,
+        nchunks,
+        helion.cdiv(chunk_size, BLOCK_SIZE_M_min),
+        chunk_size,
+        device=x.device,
+        dtype=torch.float32,
+    )
+
+    p = 1.44269504089  # 1 / ln(2)
+    block_m = hl.register_block_size(BLOCK_SIZE_M_min, chunk_size)
+    block_n = hl.register_block_size(chunk_size)
+
+    x = torch.transpose(x, 1, 3)
+
+    for tile_c, tile_b, tile_h, tile_m in hl.tile(
+        [nchunks, batch, nheads, chunk_size],
+        block_size=[1, 1, 1, block_m],
+    ):
+        rowsum_local = hl.zeros([block_m], dtype=torch.float32)
+        dout_local = dout[
+            tile_b.begin,
+            tile_c.begin * chunk_size + tile_m.index,
+            tile_h.begin,
+            :,
+        ]
+        dA_cs_m_local = (
+            dA_cumsum[
+                tile_b.begin,
+                tile_h.begin,
+                tile_c.begin,
+                tile_m,
+            ].to(torch.float32)
+            * p
+        )
+
+        lo, hi = 0, (tile_m.id + 1) * block_m
+        for tile_n in hl.tile(lo, hi, block_size=block_n):
+            x_local = x[
+                tile_b.begin,
+                :,
+                tile_h.begin,
+                tile_c.begin * chunk_size + tile_n.index,
+            ]
+            acc = hl.dot(dout_local, x_local)
+            dt_n_local = dt[
+                tile_b.begin,
+                tile_h.begin,
+                tile_c.begin,
+                tile_n,
+            ].to(torch.float32)
+            acc *= dt_n_local
+            cb_local = cb[
+                tile_b.begin,
+                tile_c.begin,
+                tile_h.begin // (nheads // ngroups),
+                tile_m,
+                tile_n,
+            ].to(torch.float32)
+            acc *= cb_local
+            dA_cs_n_local = (
+                dA_cumsum[
+                    tile_b.begin,
+                    tile_h.begin,
+                    tile_c.begin,
+                    tile_n,
+                ].to(torch.float32)
+                * p
+            )
+            acc *= torch.exp2(
+                (dA_cs_m_local[:, None] - dA_cs_n_local[None, :]).clamp(max=0.0)
+            )
+            mask = (tile_m.index + 0)[:, None] >= (tile_n.index + 0)[None, :] + 1
+            acc = torch.where(mask, acc, torch.zeros_like(acc))
+            rowsum_new_local = rowsum_local + torch.sum(acc, dim=1)
+            acc = rowsum_local[:, None] + torch.cumsum(acc, dim=1)
+            rowsum_local = rowsum_new_local
+            acc = torch.where(mask, acc, torch.zeros_like(acc))
+            ddA_cs_local = torch.sum(acc, dim=0)
+            hl.store(
+                ddA_cumsum,
+                (
+                    tile_b.begin,
+                    tile_h.begin,
+                    tile_c.begin,
+                    tile_m.id,
+                    tile_n.index + 1,
+                ),
+                ddA_cs_local,
+                extra_mask=tile_n.index + 1 < chunk_size,
+            )
+
+    n_valid_blocks = (chunk_size + block_m - 1) // block_m
+    ddA_cumsum = ddA_cumsum[:, :, :, :n_valid_blocks].sum(dim=3)
+    return ddA_cumsum
+
+
+def _chunk_scan_bwd_ddAcs_stable(
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    dA_cumsum: torch.Tensor,
+    dout: torch.Tensor,
+    cb: torch.Tensor,
+) -> torch.Tensor:
+    return _chunk_scan_bwd_ddAcs_stable_kernel(x, dt, dA_cumsum, dout, cb)
diff --git a/mamba_ssm/ops/helion/_chunk_scan_bwd_dstates.py b/mamba_ssm/ops/helion/_chunk_scan_bwd_dstates.py
new file mode 100644
index 000000000..9ea001b26
--- /dev/null
+++ b/mamba_ssm/ops/helion/_chunk_scan_bwd_dstates.py
@@ -0,0 +1,124 @@
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+from einops import rearrange
+import torch
+import triton
+
+import helion
+import helion.language as hl
+
+from ..triton.ssd_chunk_scan import (
+    _chunk_scan_bwd_dstates as _chunk_scan_bwd_dstates_ref,
+)
+
+from .utils import json_cached_autotune
+
+
+@json_cached_autotune
+@helion.kernel(
+    autotune_accuracy_check=False,
+    autotune_baseline_fn=_chunk_scan_bwd_dstates_ref,
+    ignore_warnings=[helion.exc.TensorOperationInWrapper],
+)
+def _chunk_scan_bwd_dstates_kernel(
+    C: torch.Tensor,
+    dA_cumsum: torch.Tensor,
+    dout: torch.Tensor,
+    seq_idx: Optional[torch.Tensor] = None,
+    dtype: Optional[torch.dtype] = None,
+) -> torch.Tensor:
+    batch, seqlen, nheads, headdim = dout.shape
+    _, _, nchunks, chunk_size = dA_cumsum.shape
+    _, _, ngroups, dstate = C.shape
+    assert seqlen % chunk_size == 0
+    assert nheads % ngroups == 0
+    assert C.shape == (batch, seqlen, ngroups, dstate)
+    assert dA_cumsum.shape == (batch, nheads, nchunks, chunk_size)
+    if seq_idx is not None:
+        assert seq_idx.shape == (batch, seqlen)
+    dtype = C.dtype if dtype is None else dtype
+    dprev_states = torch.empty(
+        batch,
+        nchunks,
+        nheads,
+        headdim,
+        dstate,
+        device=C.device,
+        dtype=dtype,
+    )
+
+    HAS_SEQ_IDX = seq_idx is not None
+
+    block_m = hl.register_block_size(headdim)
+    block_n = hl.register_block_size(dstate)
+    block_k = hl.register_block_size(chunk_size)
+
+    dout = torch.transpose(dout, 1, 3)
+
+    for tile_c, tile_b, tile_h, tile_m, tile_n in hl.tile(
+        [nchunks, batch, nheads, headdim, dstate],
+        block_size=[1, 1, 1, block_m, block_n],
+    ):
+        acc = hl.zeros([block_m, block_n], dtype=torch.float32)
+        if HAS_SEQ_IDX:
+            if tile_c.id == 0:
+                seq_idx_prev_local = hl.zeros([], dtype=seq_idx.dtype)
+            else:
+                seq_idx_prev_local = seq_idx[
+                    tile_b.begin,
+                    tile_c.begin * chunk_size - 1,
+                ]
+        for tile_k in hl.tile(chunk_size, block_size=block_k):
+            dout_local = dout[
+                tile_b.begin,
+                tile_m,
+                tile_h.begin,
+                tile_c.begin * chunk_size + tile_k.index,
+            ].to(torch.float32)
+            dA_cs_k_local = dA_cumsum[
+                tile_b.begin,
+                tile_h.begin,
+                tile_c.begin,
+                tile_k,
+            ].to(torch.float32)
+            scale_k = torch.exp(dA_cs_k_local)
+            if HAS_SEQ_IDX:
+                seq_idx_k_local = seq_idx[
+                    tile_b.begin,
+                    tile_c.begin * chunk_size + tile_k.index,
+                ]
+                scale_k = torch.where(
+                    seq_idx_k_local == seq_idx_prev_local,
+                    scale_k,
+                    torch.zeros_like(scale_k),
+                )
+            dout_local = (dout_local * scale_k).to(dout.dtype)
+            c_local = C[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_k.index,
+                tile_h.begin // (nheads // ngroups),
+                tile_n,
+            ].to(dout.dtype)
+            acc = hl.dot(dout_local, c_local, acc=acc)
+        dprev_states[
+            tile_b.begin,
+            tile_c.begin,
+            tile_h.begin,
+            tile_m,
+            tile_n,
+        ] = acc
+
+    return dprev_states
+
+
+def _chunk_scan_bwd_dstates(
+    C: torch.Tensor,
+    dA_cumsum: torch.Tensor,
+    dout: torch.Tensor,
+    seq_idx: Optional[torch.Tensor] = None,
+    dtype: Optional[torch.dtype] = None,
+) -> torch.Tensor:
+    return _chunk_scan_bwd_dstates_kernel(C, dA_cumsum, dout, seq_idx, dtype)
diff --git a/mamba_ssm/ops/helion/_chunk_scan_chunk_state_bwd_dx.py b/mamba_ssm/ops/helion/_chunk_scan_chunk_state_bwd_dx.py
new file mode 100644
index 000000000..e1fbe0057
--- /dev/null
+++ b/mamba_ssm/ops/helion/_chunk_scan_chunk_state_bwd_dx.py
@@ -0,0 +1,460 @@
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+import triton
+
+import helion
+import helion.language as hl
+
+from mamba_ssm.utils.determinism import (
+    alloc_tile_workspace,
+    finalize_tile_workspace,
+)
+
+from .utils import json_cached_autotune
+
+
+def _chunk_scan_chunk_state_bwd_dx_ref(
+    x, dt, dA_cumsum, B, CB, dout, dstates, D=None, seq_idx=None, dx=None
+):
+    # Autotune baseline for ``_chunk_scan_chunk_state_bwd_dx_kernel``. It drives the
+    # low-level Triton kernel directly (rather than ssd_combined's wrapper) so that
+    # its input/output exactly matches the Helion kernel: it returns ``dD`` at rank 2
+    # ``(nheads, dD_hdim)`` -- (nheads, headdim) for 2-D D and (nheads, 1) for 1-D D.
+    # Unlike ssd_combined's wrapper it does NOT squeeze the trailing dim for 1-D D;
+    # the Helion kernel keeps dD at rank 2 and defers that squeeze to its host
+    # wrapper, so the baseline must match that shape. Deterministic reduction is
+    # forced off so that 1-D D keeps dD_hdim == 1 and the reduction below yields
+    # (nheads, 1) directly, with no separate squeeze step.
+    #
+    # The ssd_combined import stays lazy to avoid a circular import: ssd_combined
+    # imports this module at top level, so importing it at module load time would
+    # fail with a partially-initialized-module error. It is only needed when
+    # autotuning runs, by which point ssd_combined is fully initialized.
+    from mamba_ssm.ops.triton.ssd_combined import (
+        _CHUNK_SCAN_CHUNK_STATE_BWD_DX_MIN_BLOCK_N,
+        TRITON_22,
+        _chunk_scan_chunk_state_bwd_dx_kernel,
+    )
+
+    batch, seqlen, nheads, headdim = x.shape
+    _, _, nchunks, chunk_size = dt.shape
+    _, _, ngroups, dstate = B.shape
+    assert nheads % ngroups == 0
+    assert B.shape == (batch, seqlen, ngroups, dstate)
+    assert CB.shape == (batch, nchunks, ngroups, chunk_size, chunk_size)
+    assert dt.shape == (batch, nheads, nchunks, chunk_size)
+    assert dA_cumsum.shape == dt.shape
+    assert dout.shape == x.shape
+    assert dstates.shape == (batch, nchunks, nheads, headdim, dstate)
+    if seq_idx is not None:
+        assert seq_idx.shape == (batch, seqlen)
+    deterministic = False
+    if D is not None:
+        assert D.shape == (nheads, headdim) or D.shape == (nheads,)
+        assert D.stride(-1) == 1
+        BLOCK_SIZE_min = 32
+        pid_m_tiles = triton.cdiv(chunk_size, BLOCK_SIZE_min)
+        pid_n_tiles = math.ceil(headdim / _CHUNK_SCAN_CHUNK_STATE_BWD_DX_MIN_BLOCK_N)
+        if D.dim() == 2:
+            dD_hdim = headdim
+        elif deterministic:
+            dD_hdim = pid_n_tiles
+        else:
+            dD_hdim = 1
+        dD = torch.zeros(
+            pid_m_tiles,
+            batch,
+            nchunks,
+            nheads,
+            dD_hdim,
+            device=D.device,
+            dtype=torch.float32,
+        )
+        dD_strides = (
+            dD.stride(0),
+            dD.stride(1),
+            dD.stride(2),
+            dD.stride(3),
+            dD.stride(4),
+        )
+    else:
+        dD = None
+        dD_strides = (0, 0, 0, 0, 0)
+    if dx is None:
+        dx = torch.empty_like(x)
+    else:
+        assert dx.shape == x.shape
+    tile_count = math.ceil(headdim / _CHUNK_SCAN_CHUNK_STATE_BWD_DX_MIN_BLOCK_N)
+    ddt, stride_ddt_tile = alloc_tile_workspace(
+        (batch, nheads, nchunks, chunk_size),
+        tile_count,
+        torch.float32,
+        dout.device,
+        deterministic,
+        zero_init=True,
+    )
+    grid_dx = lambda META: (
+        triton.cdiv(chunk_size, META["BLOCK_SIZE_M"])
+        * triton.cdiv(headdim, META["BLOCK_SIZE_N"]),
+        batch * nchunks,
+        nheads,
+    )
+    with torch.cuda.device(x.device.index):
+        _chunk_scan_chunk_state_bwd_dx_kernel[grid_dx](
+            x,
+            CB,
+            dout,
+            dt,
+            dA_cumsum,
+            seq_idx,
+            D,
+            B,
+            dstates,
+            dx,
+            ddt,
+            dD,
+            chunk_size,
+            headdim,
+            dstate,
+            batch,
+            seqlen,
+            nheads // ngroups,
+            x.stride(0),
+            x.stride(1),
+            x.stride(2),
+            x.stride(3),
+            CB.stride(0),
+            CB.stride(1),
+            CB.stride(2),
+            CB.stride(-1),
+            CB.stride(-2),
+            dout.stride(0),
+            dout.stride(1),
+            dout.stride(2),
+            dout.stride(3),
+            dt.stride(0),
+            dt.stride(2),
+            dt.stride(1),
+            dt.stride(3),
+            dA_cumsum.stride(0),
+            dA_cumsum.stride(2),
+            dA_cumsum.stride(1),
+            dA_cumsum.stride(3),
+            *(
+                (seq_idx.stride(0), seq_idx.stride(1))
+                if seq_idx is not None
+                else (0, 0)
+            ),
+            D.stride(0) if D is not None else 0,
+            B.stride(0),
+            B.stride(1),
+            B.stride(2),
+            B.stride(3),
+            dstates.stride(0),
+            dstates.stride(1),
+            dstates.stride(2),
+            dstates.stride(3),
+            dstates.stride(4),
+            dx.stride(0),
+            dx.stride(1),
+            dx.stride(2),
+            dx.stride(3),
+            ddt.stride(0),
+            ddt.stride(2),
+            ddt.stride(1),
+            ddt.stride(3),
+            stride_ddt_tile,
+            dD_strides[1],
+            dD_strides[2],
+            dD_strides[3],
+            dD_strides[0],
+            dD_strides[4],
+            D is not None,
+            D.dim() == 2 if D is not None else True,
+            HAS_SEQ_IDX=seq_idx is not None,
+            BLOCK_SIZE_DSTATE=max(triton.next_power_of_2(dstate), 16),
+            IS_TRITON_22=TRITON_22,
+            DETERMINISTIC_REDUCTION=deterministic,
+        )
+    if D is not None:
+        BLOCK_SIZE_actual = _chunk_scan_chunk_state_bwd_dx_kernel.best_config.kwargs[
+            "BLOCK_SIZE_M"
+        ]
+        n_valid_blocks = (chunk_size + BLOCK_SIZE_actual - 1) // BLOCK_SIZE_actual
+        dD = dD[:n_valid_blocks].sum(dim=(0, 1, 2))
+    ddt = finalize_tile_workspace(ddt, deterministic)
+    return dx, ddt, dD
+
+
+@json_cached_autotune
+@helion.kernel(
+    autotune_accuracy_check=False,
+    autotune_baseline_fn=_chunk_scan_chunk_state_bwd_dx_ref,
+    ignore_warnings=[helion.exc.TensorOperationInWrapper],
+)
+def _chunk_scan_chunk_state_bwd_dx_kernel(
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    dA_cumsum: torch.Tensor,
+    B: torch.Tensor,
+    CB: torch.Tensor,
+    dout: torch.Tensor,
+    dstates: torch.Tensor,
+    D: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+    dx: Optional[torch.Tensor] = None,  # required in practice; see host wrapper
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    batch, seqlen, nheads, headdim = x.shape
+    _, _, nchunks, chunk_size = dt.shape
+    _, _, ngroups, dstate = B.shape
+    assert seqlen % chunk_size == 0
+    assert dt.shape == (batch, nheads, nchunks, chunk_size)
+    assert dA_cumsum.shape == dt.shape
+    assert nheads % ngroups == 0
+    assert B.shape == (batch, seqlen, ngroups, dstate)
+    assert CB.shape == (batch, nchunks, ngroups, chunk_size, chunk_size)
+    assert dout.shape == x.shape
+    assert dstates.shape == (batch, nchunks, nheads, headdim, dstate)
+    if seq_idx is not None:
+        assert seq_idx.shape == (batch, seqlen)
+    BLOCK_M_MIN = 32
+    if D is not None:
+        assert D.shape == (nheads, headdim) or D.shape == (nheads,)
+        assert D.stride(-1) == 1
+        pid_m_tiles = helion.cdiv(chunk_size, BLOCK_M_MIN)
+        if D.dim() == 2:
+            dD_hdim = headdim
+        else:
+            dD_hdim = 1
+        dD = torch.zeros(
+            pid_m_tiles,
+            batch,
+            nchunks,
+            nheads,
+            dD_hdim,
+            device=D.device,
+            dtype=torch.float32,
+        )
+    else:
+        dD = None
+    # dx must already be a tensor: allocating it here would make Helion's type
+    # propagation merge LiteralType(None) with TensorType and fail at bind() time.
+    assert dx is not None, "dx must be preallocated by the caller"
+    assert dx.shape == x.shape
+    ddt = torch.zeros(
+        (batch, nheads, nchunks, chunk_size),
+        device=dt.device,
+        dtype=torch.float32,
+    )
+
+    accum_dtype = torch.float32
+    block_m = hl.register_block_size(BLOCK_M_MIN, chunk_size)
+    block_n = hl.register_block_size(headdim)
+    block_k = hl.register_block_size(chunk_size)
+
+    CB = torch.transpose(CB, 3, 4)
+    dstates = torch.transpose(dstates, 3, 4)
+
+    HAS_D = D is not None
+    if D is not None:
+        D_HAS_HDIM = D.dim() == 2
+    else:
+        D_HAS_HDIM = True
+    HAS_SEQ_IDX = seq_idx is not None
+
+    for tile_c, tile_b, tile_h, tile_m, tile_n in hl.tile(
+        [nchunks, batch, nheads, chunk_size, headdim],
+        block_size=[1, 1, 1, block_m, block_n],
+    ):
+        chunk_size_limit = min(chunk_size, seqlen - tile_c.id * chunk_size)
+
+        # [block_m, block_n]
+        acc = hl.zeros([block_m, block_n], dtype=accum_dtype)
+
+        # [block_m]
+        dA_cs_m_local = dA_cumsum[
+            tile_b.begin,
+            tile_h.begin,
+            tile_c.begin,
+            tile_m,
+        ].to(torch.float32)
+        # [1]
+        dA_cs_last_local = dA_cumsum[
+            tile_b.begin,
+            tile_h.begin,
+            tile_c.begin,
+            chunk_size - 1,
+        ].to(torch.float32)
+        scale = torch.exp((dA_cs_last_local - dA_cs_m_local).clamp(max=0.0))
+        if HAS_SEQ_IDX:
+            seq_idx_m_local = seq_idx[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_m.index,
+            ]
+            seq_idx_last_local = seq_idx[
+                tile_b.begin,
+                tile_c.begin * chunk_size + chunk_size - 1,
+            ]
+            scale_mask = seq_idx_m_local == seq_idx_last_local
+            scale = torch.where(scale_mask, scale, torch.zeros_like(scale))
+
+        # [block_m, dstate]
+        b_local = B[
+            tile_b.begin,
+            tile_c.begin * chunk_size + tile_m.index,
+            tile_h.begin // (nheads // ngroups),
+            :,
+        ]
+        # [dstate, block_n]
+        dstates_local = dstates[
+            tile_b.begin,
+            tile_c.begin,
+            tile_h.begin,
+            :,
+            tile_n,
+        ].to(B.dtype)
+        acc = hl.dot(b_local, dstates_local, acc=acc)
+        acc *= scale[:, None]
+
+        K_MIN = tile_m.id * block_m
+        K_MAX = chunk_size_limit
+        for tile_k in hl.tile(K_MIN, K_MAX, block_size=block_k):
+            # [block_m, block_k]
+            cb_local = CB[
+                tile_b.begin,
+                tile_c.begin,
+                tile_h.begin // (nheads // ngroups),
+                tile_m,
+                tile_k,
+            ]
+            # [block_k, block_n]
+            dout_local = dout[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_k.index,
+                tile_h.begin,
+                tile_n,
+            ]
+            # [block_k]
+            dA_cs_k_local = dA_cumsum[
+                tile_b.begin,
+                tile_h.begin,
+                tile_c.begin,
+                tile_k,
+            ].to(torch.float32)
+            cb_local *= torch.exp(
+                (dA_cs_k_local[None, :] - dA_cs_m_local[:, None]).clamp(max=0.0)
+            )
+            cb_local_mask = (tile_k.index + 0)[None, :] >= (tile_m.index + 0)[:, None]
+            cb_local = torch.where(cb_local_mask, cb_local, torch.zeros_like(cb_local))
+            cb_local = cb_local.to(dout.dtype)
+            acc = hl.dot(cb_local, dout_local, acc=acc)
+
+        # [block_m]
+        dt_m_local = dt[
+            tile_b.begin,
+            tile_h.begin,
+            tile_c.begin,
+            tile_m,
+        ].to(torch.float32)
+        dx_local = acc * dt_m_local[:, None]
+
+        if HAS_D:
+            # [block_m, block_n]
+            dout_res_local = dout[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_m.index,
+                tile_h.begin,
+                tile_n,
+            ].to(torch.float32)
+            if D_HAS_HDIM:
+                d_local = D[tile_h.begin, tile_n].to(torch.float32)
+            else:
+                d_local = D[tile_h.begin].to(torch.float32)
+            dx_local += dout_res_local * d_local
+        dx[
+            tile_b.begin,
+            tile_c.begin * chunk_size + tile_m.index,
+            tile_h.begin,
+            tile_n,
+        ] = dx_local
+
+        # [block_m, block_n]
+        x_local = x[
+            tile_b.begin,
+            tile_c.begin * chunk_size + tile_m.index,
+            tile_h.begin,
+            tile_n,
+        ].to(torch.float32)
+        if HAS_D:
+            if D_HAS_HDIM:
+                dd_local = torch.sum(dout_res_local * x_local, dim=0)
+                dD[
+                    tile_m.id,
+                    tile_b.begin,
+                    tile_c.begin,
+                    tile_h.begin,
+                    tile_n,
+                ] = dd_local
+            else:
+                dd_local = torch.sum(dout_res_local * x_local, dim=1)
+                dd_local = torch.sum(dd_local, dim=0)
+                hl.atomic_add(
+                    dD,
+                    [
+                        tile_m.id,
+                        tile_b.begin,
+                        tile_c.begin,
+                        tile_h.begin,
+                        0,
+                    ],
+                    dd_local,
+                )
+
+        # [block_m]
+        ddt_local = torch.sum(acc * x_local, dim=1)
+        hl.atomic_add(
+            ddt,
+            [
+                tile_b.begin,
+                tile_h.begin,
+                tile_c.begin,
+                tile_m,
+            ],
+            ddt_local,
+        )
+
+    if D is not None:
+        n_valid_blocks = (chunk_size + block_m - 1) // block_m
+        dD = dD[:n_valid_blocks].sum(dim=(0, 1, 2))
+
+    return dx, ddt, dD
+
+
+def _chunk_scan_chunk_state_bwd_dx(
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    dA_cumsum: torch.Tensor,
+    B: torch.Tensor,
+    CB: torch.Tensor,
+    dout: torch.Tensor,
+    dstates: torch.Tensor,
+    D: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+    dx: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    if dx is None:
+        dx = torch.empty_like(x)
+    else:
+        assert dx.shape == x.shape
+    dx, ddt, dD = _chunk_scan_chunk_state_bwd_dx_kernel(
+        x, dt, dA_cumsum, B, CB, dout, dstates, D, seq_idx, dx
+    )
+    if D is not None:
+        if D.dim() == 1:
+            dD = dD.sum(dim=-1)
+        dD = dD.to(dtype=D.dtype)
+    return dx, ddt, dD
diff --git a/mamba_ssm/ops/helion/_chunk_scan_fwd.py b/mamba_ssm/ops/helion/_chunk_scan_fwd.py
new file mode 100644
index 000000000..69fe48184
--- /dev/null
+++ b/mamba_ssm/ops/helion/_chunk_scan_fwd.py
@@ -0,0 +1,211 @@
+# https://github.com/pytorch/helion/blob/v1.0.0/examples/mamba2_chunk_scan.py
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+import helion
+import helion.language as hl
+
+from .utils import json_cached_autotune
+
+from ..triton.ssd_chunk_scan import _chunk_scan_fwd as _chunk_scan_fwd_ref
+
+
+@json_cached_autotune
+@helion.kernel(
+    autotune_accuracy_check=False,
+    autotune_baseline_fn=_chunk_scan_fwd_ref,
+    ignore_warnings=[helion.exc.TensorOperationInWrapper],
+)
+def _chunk_scan_fwd_kernel(
+    cb: torch.Tensor,
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    dA_cumsum: torch.Tensor,
+    C: torch.Tensor,
+    states: torch.Tensor,
+    D: Optional[torch.Tensor] = None,
+    z: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+    batch, seqlen, nheads, headdim = x.shape
+    _, _, nchunks, chunk_size = dt.shape
+    _, _, ngroups, dstate = C.shape
+    assert seqlen % chunk_size == 0
+    assert nheads % ngroups == 0
+    assert C.shape == (batch, seqlen, ngroups, dstate)
+    assert cb.shape == (batch, nchunks, ngroups, chunk_size, chunk_size)
+    if z is not None:
+        assert z.shape == x.shape
+    if D is not None:
+        assert D.shape == (nheads, headdim) or D.shape == (nheads,)
+    assert dt.shape == (batch, nheads, nchunks, chunk_size)
+    assert dA_cumsum.shape == (batch, nheads, nchunks, chunk_size)
+    assert states.shape == (batch, nchunks, nheads, headdim, dstate)
+    if seq_idx is not None:
+        assert seq_idx.shape == (batch, seqlen)
+
+    # Allocates output.
+    out = torch.empty(batch, seqlen, nheads, headdim, device=x.device, dtype=x.dtype)
+    if z is not None:
+        out_x = torch.empty(
+            batch, seqlen, nheads, headdim, device=x.device, dtype=x.dtype
+        )
+        assert out_x.stride() == out.stride()
+    else:
+        out_x = None
+
+    acc_dtype = torch.float32
+    # 1 / ln(2)
+    p = 1.44269504089
+    block_m = hl.register_block_size(chunk_size)
+    block_n = hl.register_block_size(headdim)
+    block_k = hl.register_block_size(chunk_size)
+
+    HAS_D = D is not None
+    D_HAS_HDIM = D.dim() == 2 if D is not None else False
+    HAS_Z = z is not None
+    HAS_SEQ_IDX = seq_idx is not None
+
+    for tile_h, tile_m, tile_n, tile_b, tile_c in hl.tile(
+        [nheads, chunk_size, headdim, batch, nchunks],
+        block_size=[1, block_m, block_n, 1, 1],
+    ):
+        acc = hl.zeros([tile_m, tile_n], dtype=acc_dtype)
+
+        dA_cs_m_local = (
+            dA_cumsum[
+                tile_b.begin,
+                tile_h.begin,
+                tile_c.begin,
+                tile_m,
+            ].to(torch.float32)
+            * p
+        )
+        scale_m_local = torch.exp2(dA_cs_m_local)
+        if HAS_SEQ_IDX:
+            if tile_c.id == 0:
+                seq_idx_prev_local = hl.zeros([], dtype=seq_idx.dtype)
+            else:
+                seq_idx_prev_local = seq_idx[
+                    tile_b.begin,
+                    tile_c.begin * chunk_size - 1,
+                ]
+            seq_idx_m_local = seq_idx[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_m.index,
+            ]
+            scale_m_local_mask = seq_idx_m_local == seq_idx_prev_local
+            scale_m_local = torch.where(
+                scale_m_local_mask, scale_m_local, torch.zeros_like(scale_m_local)
+            )
+
+        c_local = C[
+            tile_b.begin,
+            tile_m.index + tile_c.begin * chunk_size,
+            tile_h.begin // (nheads // ngroups),
+            :,
+        ]
+        prev_states_local = states[
+            tile_b.begin,
+            tile_c.begin,
+            tile_h.begin,
+            tile_n,
+            :,
+        ].to(C.dtype)
+        acc = hl.dot(c_local, prev_states_local.T, acc=acc)
+        acc *= scale_m_local[:, None]
+
+        for tile_k in hl.tile((tile_m.id + 1) * block_m, block_size=block_k):
+            cb_local = cb[
+                tile_b.begin,
+                tile_c.begin,
+                tile_h.begin // (nheads // ngroups),
+                tile_m,
+                tile_k,
+            ].to(torch.float32)
+            dA_cs_k_local = (
+                dA_cumsum[
+                    tile_b.begin,
+                    tile_h.begin,
+                    tile_c.begin,
+                    tile_k,
+                ].to(torch.float32)
+                * p
+            )
+            cb_local *= torch.exp2(
+                (dA_cs_m_local[:, None] - dA_cs_k_local[None, :]).clamp(max=0.0)
+            )
+            dt_local = dt[
+                tile_b.begin,
+                tile_h.begin,
+                tile_c.begin,
+                tile_k,
+            ].to(torch.float32)
+            cb_local = cb_local * dt_local[None, :]
+            cb_local_mask = (tile_m.index + 0)[:, None] >= (tile_k.index + 0)[None, :]
+            cb_local = torch.where(cb_local_mask, cb_local, torch.zeros_like(cb_local))
+            cb_local = cb_local.to(x.dtype)
+            x_local = x[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_k.index,
+                tile_h.begin,
+                tile_n,
+            ]
+            acc = hl.dot(cb_local, x_local, acc=acc)
+
+        if HAS_D:
+            if D_HAS_HDIM:
+                d_local = D[tile_h.begin, tile_n].to(torch.float32)
+            else:
+                d_local = D[tile_h.begin].to(torch.float32)
+            x_res_local = x[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_m.index,
+                tile_h.begin,
+                tile_n,
+            ].to(torch.float32)
+            acc += x_res_local * d_local
+
+        if HAS_Z:
+            out_x[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_m.index,
+                tile_h.begin,
+                tile_n,
+            ] = acc
+            z_local = z[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_m.index,
+                tile_h.begin,
+                tile_n,
+            ].to(torch.float32)
+            acc *= z_local * torch.sigmoid(z_local)
+
+        out[
+            tile_b.begin,
+            tile_c.begin * chunk_size + tile_m.index,
+            tile_h.begin,
+            tile_n,
+        ] = acc
+
+    return out, out_x
+
+
+def _chunk_scan_fwd(
+    cb: torch.Tensor,
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    dA_cumsum: torch.Tensor,
+    C: torch.Tensor,
+    states: torch.Tensor,
+    D: Optional[torch.Tensor] = None,
+    z: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
+    return _chunk_scan_fwd_kernel(
+        cb, x, dt, dA_cumsum, C, states, D=D, z=z, seq_idx=seq_idx
+    )
diff --git a/mamba_ssm/ops/helion/_chunk_state_bwd_db.py b/mamba_ssm/ops/helion/_chunk_state_bwd_db.py
new file mode 100644
index 000000000..5fcd5fcb9
--- /dev/null
+++ b/mamba_ssm/ops/helion/_chunk_state_bwd_db.py
@@ -0,0 +1,246 @@
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+
+import triton
+import triton.language as tl
+
+import helion
+import helion.language as hl
+from helion.autotuner import IntegerFragment
+
+from ..triton.ssd_chunk_state import (
+    _chunk_state_bwd_db as _chunk_state_bwd_db_ref,
+)
+
+from .utils import json_cached_autotune
+
+from mamba_ssm.utils.determinism import (
+    alloc_tile_workspace,
+    finalize_tile_workspace,
+)
+
+
+@triton.jit
+def _masked_atomic_add(out_ptr, offs, val, mask):
+    tl.atomic_add(out_ptr + offs, val, mask=mask)
+
+
+@json_cached_autotune
+@helion.kernel(
+    autotune_accuracy_check=False,
+    autotune_baseline_fn=_chunk_state_bwd_db_ref,
+    ignore_warnings=[helion.exc.TensorOperationInWrapper],
+)
+def _chunk_state_bwd_db_kernel(
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    dA_cumsum: torch.Tensor,
+    dstates: torch.Tensor,
+    seq_idx: Optional[torch.Tensor] = None,
+    B: Optional[torch.Tensor] = None,
+    ngroups: int = 1,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    batch, seqlen, nheads, headdim = x.shape
+    _, _, nchunks, chunk_size = dt.shape
+    dstate = dstates.shape[-1]
+    assert seqlen % chunk_size == 0
+    assert dt.shape == (batch, nheads, nchunks, chunk_size)
+    assert dA_cumsum.shape == dt.shape
+    assert dstates.shape == (batch, nchunks, nheads, headdim, dstate)
+    if seq_idx is not None:
+        assert seq_idx.shape == (batch, seqlen)
+
+    block_n = hl.register_block_size(dstate)
+    # Fixed, not autotuned: it selects between a per-tile ddA_cumsum workspace and the
+    # masked-atomic path below, so it also decides ddA_cumsum's shape.
+    # TODO: make it a tunable and let autotuning pick.
+    DETERMINISTIC_REDUCTION = True
+    DETERMINISTIC_REDUCTION = hl.specialize(DETERMINISTIC_REDUCTION)
+    if B is not None:
+        assert B.shape == (batch, seqlen, ngroups, dstate)
+        tile_count = math.ceil(dstate / block_n)
+        ddA_cumsum, _ = alloc_tile_workspace(
+            (batch, nheads, nchunks, chunk_size),
+            tile_count,
+            torch.float32,
+            x.device,
+            DETERMINISTIC_REDUCTION,
+            zero_init=True,
+        )
+    else:
+        ddA_cumsum = None
+
+    ngroups = hl.specialize(ngroups)
+    nheads_ngroups_ratio = nheads // ngroups
+    nheads_per_program = hl.register_tunable(
+        "nheads_per_program",
+        IntegerFragment(1, nheads_ngroups_ratio),
+    )
+    nsplits = helion.cdiv(nheads_ngroups_ratio, nheads_per_program)
+    dB = torch.empty(
+        batch,
+        seqlen,
+        nsplits,
+        ngroups,
+        dstate,
+        device=x.device,
+        dtype=torch.float32,
+    )
+
+    HAS_DDA_CS = ddA_cumsum is not None
+    HAS_SEQ_IDX = seq_idx is not None
+    p = 1.44269504089  # 1 / ln(2)
+    block_m = hl.register_block_size(chunk_size)
+
+    for tile_c, tile_b, tile_s, tile_g, tile_m, tile_n in hl.tile(
+        [nchunks, batch, nsplits, ngroups, chunk_size, dstate],
+        block_size=[1, 1, 1, 1, block_m, block_n],
+    ):
+        acc = hl.zeros((block_m, block_n), dtype=torch.float32)
+        if HAS_DDA_CS:
+            b_local = B[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_m.index,
+                tile_g.begin,
+                tile_n,
+            ].to(torch.float32)
+        if HAS_SEQ_IDX:
+            seq_idx_m_local = seq_idx[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_m.index,
+            ]
+            seq_idx_last_local = seq_idx[
+                tile_b.begin,
+                tile_c.begin * chunk_size + chunk_size - 1,
+            ]
+
+        nheads_iter = min(
+            nheads_per_program,
+            nheads // ngroups - tile_s.id * nheads_per_program,
+        )
+        for tile_h in hl.tile(nheads_iter, block_size=1):
+            head_idx = (
+                tile_g.begin * (nheads // ngroups)
+                + tile_s.begin * nheads_per_program
+                + tile_h.begin
+            )
+            x_local = x[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_m.index,
+                head_idx,
+                :,
+            ]
+            dstates_local = dstates[
+                tile_b.begin,
+                tile_c.begin,
+                head_idx,
+                :,
+                tile_n,
+            ].to(x.dtype)
+            db = hl.dot(x_local, dstates_local)
+            dA_cs_last_local = (
+                dA_cumsum[
+                    tile_b.begin,
+                    head_idx,
+                    tile_c.begin,
+                    chunk_size - 1,
+                ].to(torch.float32)
+                * p
+            )
+            dA_cs_m_local = (
+                dA_cumsum[
+                    tile_b.begin,
+                    head_idx,
+                    tile_c.begin,
+                    tile_m,
+                ].to(torch.float32)
+                * p
+            )
+            dt_m_local = dt[
+                tile_b.begin,
+                head_idx,
+                tile_c.begin,
+                tile_m,
+            ].to(torch.float32)
+            scale = torch.exp2((dA_cs_last_local - dA_cs_m_local).clamp(max=0.0))
+            if HAS_SEQ_IDX:
+                scale = torch.where(
+                    seq_idx_m_local == seq_idx_last_local,
+                    scale,
+                    torch.zeros_like(scale),
+                )
+            db *= (scale * dt_m_local)[:, None]
+            if HAS_DDA_CS:
+                ddA_cs_local = torch.sum(db * b_local, dim=1)
+                if DETERMINISTIC_REDUCTION:
+                    hl.store(
+                        ddA_cumsum,
+                        (
+                            tile_b.begin,
+                            head_idx,
+                            tile_c.begin,
+                            tile_m.index + 1,
+                            tile_n.id,
+                        ),
+                        ddA_cs_local,
+                        extra_mask=tile_m.index < (chunk_size - 1),
+                    )
+                else:
+                    # hl.atomic_add takes no extra_mask, so the masked atomic add goes
+                    # through a hand-written Triton kernel via hl.triton_kernel.
+                    ddA_off = (
+                        tile_b.begin * ddA_cumsum.stride(0)
+                        + head_idx * ddA_cumsum.stride(1)
+                        + tile_c.begin * ddA_cumsum.stride(2)
+                        + (tile_m.index + 1) * ddA_cumsum.stride(3)
+                    )
+                    mask = tile_m.index < (chunk_size - 1)
+                    hl.triton_kernel(
+                        _masked_atomic_add,
+                        args={
+                            "out_ptr": ddA_cumsum,
+                            "offs": ddA_off,
+                            "val": ddA_cs_local,
+                            "mask": mask,
+                        },
+                        output_like=None,
+                    )
+            acc += db
+
+        dB[
+            tile_b.begin,
+            tile_c.begin * chunk_size + tile_m.index,
+            tile_s.begin,
+            tile_g.begin,
+            tile_n,
+        ] = acc
+
+    dB = dB.sum(2)
+    if ddA_cumsum is not None:
+        ddA_cumsum = finalize_tile_workspace(ddA_cumsum, DETERMINISTIC_REDUCTION)
+        torch.cumsum(ddA_cumsum, dim=-1, out=ddA_cumsum)
+    return dB if B is None else (dB, ddA_cumsum)
+
+
+def _chunk_state_bwd_db(
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    dA_cumsum: torch.Tensor,
+    dstates: torch.Tensor,
+    seq_idx: Optional[torch.Tensor] = None,
+    B: Optional[torch.Tensor] = None,
+    ngroups: int = 1,
+) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+    return _chunk_state_bwd_db_kernel(
+        x,
+        dt,
+        dA_cumsum,
+        dstates,
+        seq_idx,
+        B,
+        ngroups,
+    )
diff --git a/mamba_ssm/ops/helion/_chunk_state_fwd.py b/mamba_ssm/ops/helion/_chunk_state_fwd.py
new file mode 100644
index 000000000..d53ec08bf
--- /dev/null
+++ b/mamba_ssm/ops/helion/_chunk_state_fwd.py
@@ -0,0 +1,147 @@
+# https://github.com/pytorch/helion/blob/v1.0.0/examples/mamba2_chunk_state.py
+
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+import helion
+import helion.language as hl
+
+from .utils import json_cached_autotune
+
+from ..triton.ssd_chunk_state import (
+    _chunk_state_fwd as _chunk_state_fwd_ref,
+)
+
+
+@json_cached_autotune
+@helion.kernel(
+    autotune_accuracy_check=False,
+    autotune_baseline_fn=_chunk_state_fwd_ref,
+    ignore_warnings=[helion.exc.TensorOperationInWrapper],
+)
+def _chunk_state_fwd_kernel(
+    B: torch.Tensor,
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    dA_cumsum: torch.Tensor,
+    seq_idx: Optional[torch.Tensor] = None,
+    states: Optional[torch.Tensor] = None,
+    states_in_fp32: hl.constexpr = True,
+) -> torch.Tensor:
+    batch, seqlen, nheads, headdim = x.shape
+    _, _, nchunks, chunk_size = dt.shape
+    _, _, ngroups, dstate = B.shape
+    assert seqlen % chunk_size == 0
+    assert nheads % ngroups == 0
+    assert B.shape == (batch, seqlen, ngroups, dstate)
+    assert dt.shape == (batch, nheads, nchunks, chunk_size)
+    assert dA_cumsum.shape == dt.shape
+    if seq_idx is not None:
+        assert seq_idx.shape == (batch, seqlen)
+    assert states is not None
+    assert states.shape == (batch, nchunks, nheads, headdim, dstate)
+
+    acc_dtype = torch.float32
+    # 1 / ln(2)
+    p = 1.44269504089
+    block_m = hl.register_block_size(headdim)
+    block_n = hl.register_block_size(dstate)
+    block_k = hl.register_block_size(chunk_size)
+
+    HAS_SEQ_IDX = seq_idx is not None
+
+    for tile_h, tile_m, tile_n, tile_b, tile_c in hl.tile(
+        [nheads, headdim, dstate, batch, nchunks],
+        block_size=[1, block_m, block_n, 1, 1],
+    ):
+        dA_cs_last_local = (
+            dA_cumsum[
+                tile_b.begin,
+                tile_h.begin,
+                tile_c.begin,
+                chunk_size - 1,
+            ].to(torch.float32)
+            * p
+        )
+        if HAS_SEQ_IDX:
+            seq_idx_last_local = seq_idx[
+                tile_b.begin,
+                tile_c.begin * chunk_size + chunk_size - 1,
+            ]
+
+        acc = hl.zeros([tile_m, tile_n], dtype=acc_dtype)
+        for tile_k in hl.tile(chunk_size, block_size=block_k):
+            x_local = x[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_k.index,
+                tile_h.begin,
+                tile_m,
+            ]
+            dA_cs_k_local = (
+                dA_cumsum[
+                    tile_b.begin,
+                    tile_h.begin,
+                    tile_c.begin,
+                    tile_k,
+                ].to(torch.float32)
+                * p
+            )
+            dt_k_local = dt[
+                tile_b.begin,
+                tile_h.begin,
+                tile_c.begin,
+                tile_k,
+            ].to(torch.float32)
+            scale = (
+                torch.exp2((dA_cs_last_local - dA_cs_k_local).clamp(max=0.0))
+                * dt_k_local
+            )
+            if HAS_SEQ_IDX:
+                seq_idx_k_local = seq_idx[
+                    tile_b.begin,
+                    tile_c.begin * chunk_size + tile_k.index,
+                ]
+                scale_mask = seq_idx_k_local == seq_idx_last_local
+                scale = torch.where(scale_mask, scale, torch.zeros_like(scale))
+            b_local = B[
+                tile_b.begin,
+                tile_c.begin * chunk_size + tile_k.index,
+                tile_h.begin // (nheads // ngroups),
+                tile_n,
+            ].to(torch.float32)
+            b_local *= scale[:, None]
+            b_local = b_local.to(x.dtype)
+            acc = hl.dot(x_local.T, b_local, acc=acc)
+
+        states[
+            tile_b.begin,
+            tile_c.begin,
+            tile_h.begin,
+            tile_m,
+            tile_n,
+        ] = acc
+
+    return states
+
+
+def _chunk_state_fwd(
+    B: torch.Tensor,
+    x: torch.Tensor,
+    dt: torch.Tensor,
+    dA_cumsum: torch.Tensor,
+    seq_idx: Optional[torch.Tensor] = None,
+    states: Optional[torch.Tensor] = None,
+    states_in_fp32: bool = True,
+) -> torch.Tensor:
+    batch, _, nheads, headdim = x.shape
+    _, _, nchunks, _ = dt.shape
+    _, _, _, dstate = B.shape
+    if states is not None:
+        assert states.shape == (batch, nchunks, nheads, headdim, dstate)
+    else:
+        states_dtype = torch.float32 if states_in_fp32 else B.dtype
+        states = torch.empty((batch, nchunks, nheads, headdim, dstate), device=x.device, dtype=states_dtype)
+    return _chunk_state_fwd_kernel(B, x, dt, dA_cumsum, seq_idx, states, states_in_fp32)
diff --git a/mamba_ssm/ops/helion/_layer_norm_bwd.py b/mamba_ssm/ops/helion/_layer_norm_bwd.py
new file mode 100644
index 000000000..0347a58e3
--- /dev/null
+++ b/mamba_ssm/ops/helion/_layer_norm_bwd.py
@@ -0,0 +1,301 @@
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import torch
+
+import helion
+import helion.language as hl
+from helion.autotuner import IntegerFragment
+
+from .utils import json_cached_autotune
+
+
+@json_cached_autotune
+@helion.kernel(
+    autotune_accuracy_check=False,
+    ignore_warnings=[helion.exc.TensorOperationInWrapper],
+)
+def _layer_norm_bwd_kernel(
+    dy: torch.Tensor,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    eps: float,
+    mean: torch.Tensor,
+    rstd: torch.Tensor,
+    BLOCK_N: hl.constexpr,
+    group_size: int,
+    z: Optional[torch.Tensor] = None,
+    norm_before_gate: hl.constexpr = True,
+    is_rms_norm: hl.constexpr = False,
+    recompute_output: hl.constexpr = False,
+    dz: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+) -> tuple[Optional[torch.Tensor], ...]:
+    M, N = x.shape
+    assert N % group_size == 0
+    ngroups = N // group_size
+    assert x.stride(-1) == 1
+    assert dy.stride(-1) == 1
+    assert dy.shape == (M, N)
+    if z is not None:
+        assert z.stride(-1) == 1
+        assert z.shape == (M, N)
+    assert weight.shape == (N,)
+    assert weight.stride(-1) == 1
+    if bias is not None:
+        assert bias.stride(-1) == 1
+        assert bias.shape == (N,)
+    # allocate output
+    dx = torch.empty_like(x)
+
+    nrow_groups = hl.register_tunable("nrow_groups", IntegerFragment(1, M))
+    _dw = torch.empty(
+        (nrow_groups, N),
+        dtype=torch.float32,
+        device=weight.device,
+    )
+    if bias is not None:
+        _db = torch.empty(
+            (nrow_groups, N),
+            dtype=torch.float32,
+            device=bias.device,
+        )
+    else:
+        _db = None
+    rows_per_program = math.ceil(M / nrow_groups)
+
+    NORM_BEFORE_GATE = norm_before_gate
+    IS_RMS_NORM = is_rms_norm
+    HAS_BIAS = bias is not None
+    HAS_Z = z is not None
+    RECOMPUTE_OUTPUT = out is not None
+
+    for tile_rb, tile_g in hl.tile(
+        [nrow_groups, ngroups],
+        block_size=[1, 1],
+    ):
+        cols = hl.arange(0, BLOCK_N)
+        mask = cols < group_size
+        w_local = hl.load(
+            weight,
+            [tile_g.begin * group_size + cols],
+            extra_mask=mask,
+        ).to(torch.float32)
+        if (RECOMPUTE_OUTPUT or HAS_Z) and HAS_BIAS:
+            b_local = hl.load(
+                bias,
+                [tile_g.begin * group_size + cols],
+                extra_mask=mask,
+            ).to(torch.float32)
+        dw_local = hl.zeros([BLOCK_N], dtype=torch.float32)
+        if HAS_BIAS:
+            db_local = hl.zeros([BLOCK_N], dtype=torch.float32)
+        row_start = tile_rb.begin * rows_per_program
+        row_end = min((tile_rb.begin + 1) * rows_per_program, M)
+        for tile_r in hl.tile(row_start, row_end, block_size=1):
+            x_local = hl.load(
+                x,
+                [
+                    tile_r.begin,
+                    tile_g.begin * group_size + cols,
+                ],
+                extra_mask=mask,
+            )
+            x_local = torch.where(mask, x_local, torch.zeros_like(x_local)).to(
+                torch.float32
+            )
+            dy_local = hl.load(
+                dy,
+                [
+                    tile_r.begin,
+                    tile_g.begin * group_size + cols,
+                ],
+                extra_mask=mask,
+            )
+            dy_local = torch.where(mask, dy_local, torch.zeros_like(dy_local)).to(
+                torch.float32
+            )
+            if not IS_RMS_NORM:
+                mean_local = mean[tile_g.begin * M + tile_r.begin]
+            if HAS_Z and not NORM_BEFORE_GATE:
+                z_local = hl.load(
+                    z,
+                    [
+                        tile_r.begin,
+                        tile_g.begin * group_size + cols,
+                    ],
+                    extra_mask=mask,
+                )
+                z_local = torch.where(mask, z_local, torch.zeros_like(z_local)).to(
+                    torch.float32
+                )
+                x_og = x_local
+                x_local = x_og * z_local * torch.sigmoid(z_local)
+            rstd_local = rstd[tile_g.begin * M + tile_r.begin]
+            if not IS_RMS_NORM:
+                xhat = (x_local - mean_local) * rstd_local
+            else:
+                xhat = x_local * rstd_local
+            xhat = torch.where(mask, xhat, torch.zeros_like(xhat))
+            if HAS_Z and NORM_BEFORE_GATE:
+                z_local = hl.load(
+                    z,
+                    [
+                        tile_r.begin,
+                        tile_g.begin * group_size + cols,
+                    ],
+                    extra_mask=mask,
+                )
+                z_local = torch.where(mask, z_local, torch.zeros_like(z_local)).to(
+                    torch.float32
+                )
+                z_sigmoid = torch.sigmoid(z_local)
+                if HAS_BIAS:
+                    y = xhat * w_local + b_local
+                else:
+                    y = xhat * w_local
+                if RECOMPUTE_OUTPUT:
+                    hl.store(
+                        out,
+                        [
+                            tile_r.begin,
+                            tile_g.begin * group_size + cols,
+                        ],
+                        y * z_local * z_sigmoid,
+                        extra_mask=mask,
+                    )
+                dz_local = dy_local * y * z_sigmoid * (1 + z_local * (1 - z_sigmoid))
+                hl.store(
+                    dz,
+                    [
+                        tile_r.begin,
+                        tile_g.begin * group_size + cols,
+                    ],
+                    dz_local,
+                    extra_mask=mask,
+                )
+                dy_local *= z_local * z_sigmoid
+            else:
+                if RECOMPUTE_OUTPUT:
+                    if HAS_BIAS:
+                        y = xhat * w_local + b_local
+                    else:
+                        y = xhat * w_local
+                    hl.store(
+                        out,
+                        [
+                            tile_r.begin,
+                            tile_g.begin * group_size + cols,
+                        ],
+                        y,
+                        extra_mask=mask,
+                    )
+            wdy = w_local * dy_local
+            c1 = torch.sum(xhat * wdy, dim=0) / group_size
+            if not IS_RMS_NORM:
+                c2 = torch.sum(wdy, dim=0) / group_size
+                dx_local = (wdy - (xhat * c1 + c2)) * rstd_local
+            else:
+                dx_local = (wdy - xhat * c1) * rstd_local
+            dw_local += dy_local * xhat
+            if HAS_BIAS:
+                db_local += dy_local
+            if HAS_Z and not NORM_BEFORE_GATE:
+                z_sigmoid = torch.sigmoid(z_local)
+                dz_local = dx_local * x_og * z_sigmoid * (1 + z_local * (1 - z_sigmoid))
+                hl.store(
+                    dz,
+                    [
+                        tile_r.begin,
+                        tile_g.begin * group_size + cols,
+                    ],
+                    dz_local,
+                    extra_mask=mask,
+                )
+                dx_local *= z_local * z_sigmoid
+            hl.store(
+                dx,
+                [
+                    tile_r.begin,
+                    tile_g.begin * group_size + cols,
+                ],
+                dx_local,
+                extra_mask=mask,
+            )
+        hl.store(
+            _dw,
+            [tile_rb.begin, tile_g.begin * group_size + cols],
+            dw_local,
+            extra_mask=mask,
+        )
+        if HAS_BIAS:
+            hl.store(
+                _db,
+                [tile_rb.begin, tile_g.begin * group_size + cols],
+                db_local,
+                extra_mask=mask,
+            )
+
+    dw = _dw.sum(0).to(weight.dtype)
+    if bias is not None:
+        db = _db.sum(0).to(bias.dtype)
+    else:
+        db = None
+    return (dx, dw, db, dz) if not recompute_output else (dx, dw, db, dz, out)
+
+
+def _layer_norm_bwd(
+    dy: torch.Tensor,
+    x: torch.Tensor,
+    weight: torch.Tensor,
+    bias: Optional[torch.Tensor],
+    eps: float,
+    mean: torch.Tensor,
+    rstd: torch.Tensor,
+    z: Optional[torch.Tensor] = None,
+    group_size: Optional[int] = None,
+    norm_before_gate: bool = True,
+    is_rms_norm: bool = False,
+    recompute_output: bool = False,
+    dz: Optional[torch.Tensor] = None,
+    out: Optional[torch.Tensor] = None,
+) -> tuple[Optional[torch.Tensor], ...]:
+    _, N = x.shape
+    if group_size is None:
+        group_size = N
+    if dz is not None:
+        assert z is not None
+        assert dz.shape == z.shape
+        assert dz.stride(-1) == 1
+    else:
+        dz = torch.empty_like(z) if z is not None else None
+    if recompute_output:
+        if out is None:
+            out = torch.empty_like(x)
+        assert out.shape == x.shape
+    # One group's features must fit in 64KB, as in the Triton kernel: BLOCK_N covers a
+    # whole group, so the kernel keeps a row of it live at a time.
+    MAX_FUSED_SIZE = 65536 // x.element_size()
+    BLOCK_N = min(MAX_FUSED_SIZE, helion.next_power_of_2(group_size))
+    if group_size > BLOCK_N:
+        raise RuntimeError("This layer norm doesn't support feature dim >= 64KB.")
+    return _layer_norm_bwd_kernel(
+        dy,
+        x,
+        weight,
+        bias,
+        eps,
+        mean,
+        rstd,
+        BLOCK_N,
+        group_size,
+        z,
+        norm_before_gate,
+        is_rms_norm,
+        recompute_output,
+        dz,
+        out,
+    )
diff --git a/mamba_ssm/ops/helion/_state_passing_bwd.py b/mamba_ssm/ops/helion/_state_passing_bwd.py
new file mode 100644
index 000000000..bb562060c
--- /dev/null
+++ b/mamba_ssm/ops/helion/_state_passing_bwd.py
@@ -0,0 +1,320 @@
+from __future__ import annotations
+
+import math
+from typing import Optional, Union
+
+import torch
+
+import helion
+import helion.language as hl
+
+from ..triton.ssd_state_passing import (
+    _state_passing_bwd as _state_passing_bwd_ref,
+)
+
+from .utils import json_cached_autotune
+
+
+def _affine_combine(left, right):
+    # The recurrence s_i = a_i * s_{i-1} + b_i as composition of the affine maps
+    # T_i(x) = a_i * x + b_i, which is associative and so scannable:
+    # T_right o T_left: (a_l, b_l), (a_r, b_r) -> (a_r*a_l, a_r*b_l + b_r)
+    a1, b1 = left
+    a2, b2 = right
+    return a1 * a2, a2 * b1 + b2
+
+
+@json_cached_autotune
+@helion.kernel(
+    autotune_accuracy_check=False,
+    autotune_baseline_fn=_state_passing_bwd_ref,
+    ignore_warnings=[helion.exc.TensorOperationInWrapper],
+)
+def _state_passing_bwd_kernel(
+    states: torch.Tensor,
+    dA_chunk_cumsum: torch.Tensor,
+    dout: torch.Tensor,
+    dfinal_states: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+    has_initial_states: hl.constexpr = False,
+    dstates_dtype: Optional[torch.dtype] = None,
+    states_dtype: Optional[torch.dtype] = None,
+    chunk_size: int = 0,
+) -> Union[
+    tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]],
+    tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor],
+]:
+    """
+    states contains the initial_states at index 0. The final states are not included in states.
+    """
+    batch, nchunks, nheads, dim = states.shape
+    assert dA_chunk_cumsum.shape == (batch, nheads, nchunks)
+    assert dout.shape == (batch, nchunks, nheads, dim)
+    if seq_idx is not None:
+        assert chunk_size > 0
+        chunk_size = hl.specialize(chunk_size)
+        seqlen = seq_idx.shape[-1]
+        assert seq_idx.shape == (batch, seqlen)
+
+    dstates = torch.empty_like(
+        dout,
+        dtype=dstates_dtype if dstates_dtype is not None else dout.dtype,
+    )
+    if states_dtype is not None and states_dtype != states.dtype:
+        states_converted = torch.empty_like(
+            states,
+            dtype=dstates_dtype if dstates_dtype is not None else dout.dtype,
+        )
+        assert states_converted.stride() == states.stride()
+    else:
+        states_converted = None
+
+    if dfinal_states is not None:
+        assert dfinal_states.shape == (batch, nheads, dim)
+
+    if has_initial_states:
+        dinitstates = torch.empty(
+            (batch, nheads, dim),
+            dtype=dstates_dtype if dstates_dtype is not None else dout.dtype,
+            device=dout.device,
+        )
+    else:
+        dinitstates = None
+
+    block_m = hl.register_block_size(dim)
+    n_blocks = (dim + block_m - 1) // block_m
+    ddA_chunk_cumsum = torch.empty(
+        batch,
+        nheads,
+        nchunks,
+        n_blocks,
+        dtype=torch.float32,
+        device=dA_chunk_cumsum.device,
+    )
+
+    CONVERT_STATES = states_converted is not None
+    HAS_DFINAL_STATES = dfinal_states is not None
+    HAS_DINITSTATES = has_initial_states
+    HAS_SEQ_IDX = seq_idx is not None
+
+    block_c = hl.register_block_size(nchunks)
+
+    for tile_b, tile_h, tile_m in hl.tile(
+        [batch, nheads, dim],
+        block_size=[1, 1, block_m],
+    ):
+        if HAS_DFINAL_STATES:
+            dstates_local = dfinal_states[tile_b.begin, tile_h.begin, tile_m].to(
+                torch.float32
+            )
+        else:
+            dstates_local = hl.zeros((block_m,), dtype=torch.float32)
+        # Chunk nchunks - 1 seeds the reverse scan: its dstates is dfinal_states (or
+        # zero), and the loop below walks the remaining chunks backwards.
+        dstates[
+            tile_b.begin,
+            nchunks - 1,
+            tile_h.begin,
+            tile_m,
+        ] = dstates_local
+        dA_cs_last_local = dA_chunk_cumsum[tile_b.begin, tile_h.begin, nchunks - 1].to(
+            torch.float32
+        )
+        scale_last = torch.exp(dA_cs_last_local)
+        if HAS_SEQ_IDX:
+            seq_idx_last_local = seq_idx[tile_b.begin, seqlen - 1]
+            seq_idx_new_last_local = seq_idx[tile_b.begin, seqlen - 1 - chunk_size]
+            scale_last = torch.where(
+                seq_idx_last_local == seq_idx_new_last_local,
+                scale_last,
+                torch.zeros_like(scale_last),
+            )
+        out_last_local = states[
+            tile_b.begin,
+            nchunks - 1,
+            tile_h.begin,
+            tile_m,
+        ].to(torch.float32)
+        if CONVERT_STATES:
+            states_converted[
+                tile_b.begin,
+                nchunks - 1,
+                tile_h.begin,
+                tile_m,
+            ] = out_last_local
+        ddA_last_local = torch.sum(out_last_local * dstates_local) * scale_last
+        ddA_chunk_cumsum[
+            tile_b.begin,
+            tile_h.begin,
+            nchunks - 1,
+            tile_m.id,
+        ] = ddA_last_local
+
+        for tile_c in hl.tile(1, nchunks, block_size=block_c):
+            # tile_c.index runs 1..nchunks-1, so this reads chunk
+            # (nchunks - 1) - (tile_c.index - 1): chunks nchunks - 1 down to 1.
+            dout_prev_local = dout[
+                tile_b.begin,
+                (nchunks - 1) - (tile_c.index - 1),
+                tile_h.begin,
+                tile_m,
+            ].to(torch.float32)
+            dA_cs_prev_local = dA_chunk_cumsum[
+                tile_b.begin,
+                tile_h.begin,
+                (nchunks - 1) - (tile_c.index - 1),
+            ].to(torch.float32)
+            scale_prev = torch.exp(dA_cs_prev_local)
+            if HAS_SEQ_IDX:
+                seq_idx_prev_local = seq_idx[
+                    tile_b.begin,
+                    (nchunks - (tile_c.index - 1)) * chunk_size - 1,
+                ]
+                seq_idx_new_prev_local = seq_idx[
+                    tile_b.begin,
+                    (nchunks - tile_c.index) * chunk_size - 1,
+                ]
+                scale_prev = torch.where(
+                    seq_idx_prev_local == seq_idx_new_prev_local,
+                    scale_prev,
+                    torch.zeros_like(scale_prev),
+                )
+            A_cum, B_cum = hl.associative_scan(
+                _affine_combine,
+                (scale_prev[:, None].expand(-1, block_m), dout_prev_local),
+                dim=0,
+            )
+            tmp_dstates_local = A_cum * dstates_local[None, :] + B_cum
+            # On a partial last tile (nchunks - 1 not a multiple of block_c) lane
+            # block_c - 1 is out of range and would carry garbage into the next tile;
+            # tile_c.end - 1 is this tile's last valid index.
+            last_mask = tile_c.index == (tile_c.end - 1)
+            dstates_local = torch.where(
+                last_mask[:, None],
+                tmp_dstates_local,
+                torch.zeros_like(tmp_dstates_local),
+            ).sum(dim=0)
+            # Chunk (nchunks - 1) - tile_c.index: chunks nchunks - 2 down to 0.
+            dstates[
+                tile_b.begin,
+                (nchunks - 1) - tile_c.index,
+                tile_h.begin,
+                tile_m,
+            ] = tmp_dstates_local
+            # Chunk (nchunks - 1) - tile_c.index: chunks nchunks - 2 down to 0.
+            dA_cs_local = dA_chunk_cumsum[
+                tile_b.begin,
+                tile_h.begin,
+                (nchunks - 1) - tile_c.index,
+            ].to(torch.float32)
+            scale = torch.exp(dA_cs_local)
+            if HAS_SEQ_IDX:
+                seq_idx_local = seq_idx_new_prev_local
+                seq_idx_new_local = hl.load(
+                    seq_idx,
+                    [
+                        tile_b.begin,
+                        (nchunks - (tile_c.index + 1)) * chunk_size - 1,
+                    ],
+                    extra_mask=tile_c.index + 1 < nchunks,
+                )
+                seq_idx_new_local = torch.where(
+                    tile_c.index + 1 < nchunks,
+                    seq_idx_new_local,
+                    torch.zeros_like(seq_idx_new_local),
+                )
+                scale = torch.where(
+                    seq_idx_new_local == seq_idx_local,
+                    scale,
+                    torch.zeros_like(scale),
+                )
+            out_local = states[
+                tile_b.begin,
+                (nchunks - 1) - tile_c.index,
+                tile_h.begin,
+                tile_m,
+            ].to(torch.float32)
+            if CONVERT_STATES:
+                states_converted[
+                    tile_b.begin,
+                    (nchunks - 1) - tile_c.index,
+                    tile_h.begin,
+                    tile_m,
+                ] = out_local
+            ddA_local = torch.sum(out_local * tmp_dstates_local, dim=1) * scale
+            ddA_chunk_cumsum[
+                tile_b.begin,
+                tile_h.begin,
+                (nchunks - 1) - tile_c.index,
+                tile_m.id,
+            ] = ddA_local
+
+        if not HAS_DINITSTATES:
+            ddA_chunk_cumsum[
+                tile_b.begin,
+                tile_h.begin,
+                0,
+                tile_m.id,
+            ] = hl.zeros((), dtype=torch.float32)
+        else:
+            # Triton applies the recurrence once more for chunk 0 after leaving the
+            # chunk loop (ssd_state_passing.py:182-193):
+            #     dinitial_states = scale_0 * dstates[0] + dout[0]
+            # Without this step dstates[0] would leave as initial_states' gradient.
+            # (chunk 0's ddA_chunk_cumsum is already produced by the loop's last
+            # iteration.)
+            dA_cs_first_local = dA_chunk_cumsum[tile_b.begin, tile_h.begin, 0].to(
+                torch.float32
+            )
+            scale_first = torch.exp(dA_cs_first_local)
+            if HAS_SEQ_IDX:
+                # The seq_idx Triton tracks is seq_idx[chunk_size - 1] at this point,
+                # compared against the value 0 it started the loop with.
+                seq_idx_first_local = seq_idx[tile_b.begin, chunk_size - 1]
+                scale_first = torch.where(
+                    seq_idx_first_local == 0,
+                    scale_first,
+                    torch.zeros_like(scale_first),
+                )
+            dout_first_local = dout[tile_b.begin, 0, tile_h.begin, tile_m].to(
+                torch.float32
+            )
+            dinitstates[tile_b.begin, tile_h.begin, tile_m] = (
+                scale_first * dstates_local + dout_first_local
+            )
+
+    ddA_chunk_cumsum = ddA_chunk_cumsum.sum(dim=-1).to(dtype=dA_chunk_cumsum.dtype)
+    if states_dtype is not None and states_dtype == states.dtype:
+        states_converted = states
+    return (
+        (dstates, ddA_chunk_cumsum, dinitstates)
+        if states_dtype is None
+        else (dstates, ddA_chunk_cumsum, dinitstates, states_converted)
+    )
+
+
+def _state_passing_bwd(
+    states: torch.Tensor,
+    dA_chunk_cumsum: torch.Tensor,
+    dout: torch.Tensor,
+    dfinal_states: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+    has_initial_states: bool = False,
+    dstates_dtype: Optional[torch.dtype] = None,
+    states_dtype: Optional[torch.dtype] = None,
+    chunk_size: int = 0,
+) -> Union[
+    tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]],
+    tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor], torch.Tensor],
+]:
+    return _state_passing_bwd_kernel(
+        states,
+        dA_chunk_cumsum,
+        dout,
+        dfinal_states=dfinal_states,
+        seq_idx=seq_idx,
+        has_initial_states=has_initial_states,
+        dstates_dtype=dstates_dtype,
+        states_dtype=states_dtype,
+        chunk_size=chunk_size,
+    )
diff --git a/mamba_ssm/ops/helion/_state_passing_fwd.py b/mamba_ssm/ops/helion/_state_passing_fwd.py
new file mode 100644
index 000000000..90d4641d1
--- /dev/null
+++ b/mamba_ssm/ops/helion/_state_passing_fwd.py
@@ -0,0 +1,261 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import torch
+
+import helion
+import helion.language as hl
+
+from ..triton.ssd_state_passing import (
+    _state_passing_fwd as _state_passing_fwd_ref,
+)
+
+from .utils import json_cached_autotune
+
+
+@json_cached_autotune
+@helion.kernel(
+    autotune_accuracy_check=False,
+    ignore_warnings=[helion.exc.TensorOperationInWrapper],
+)
+def _state_passing_fwd_kernel_v1(
+    states: torch.Tensor,
+    dA_chunk_cumsum: torch.Tensor,
+    initial_states: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+    chunk_size: Optional[int] = None,
+    out_dtype: Optional[torch.dtype] = None,
+):
+    batch, nchunks, nheads, dim = states.shape
+    assert dA_chunk_cumsum.shape == (batch, nheads, nchunks)
+
+    if initial_states is not None:
+        assert initial_states.shape == (batch, nheads, dim)
+
+    if seq_idx is not None:
+        assert chunk_size is not None
+        seqlen = seq_idx.shape[-1]
+        assert seq_idx.shape == (batch, seqlen)
+
+    out_dtype = states.dtype if out_dtype is None else out_dtype
+    out = torch.empty(
+        (batch, nchunks, nheads, dim),
+        device=states.device,
+        dtype=out_dtype,
+    )
+    final_states = torch.empty(
+        (batch, nheads, dim),
+        device=states.device,
+        dtype=torch.float32,
+    )
+
+    HAS_INITSTATES = initial_states is not None
+    HAS_SEQ_IDX = seq_idx is not None
+    if chunk_size is not None:
+        chunk_size = hl.specialize(chunk_size)
+
+    block = hl.register_block_size(dim)
+
+    for tile_b, tile_h, tile_m in hl.tile(
+        [batch, nheads, dim],
+        block_size=[1, 1, block],
+    ):
+        if not HAS_INITSTATES:
+            states_local = hl.zeros((block,), dtype=torch.float32)
+        else:
+            states_local = initial_states[tile_b.begin, tile_h.begin, tile_m].to(
+                torch.float32
+            )
+
+        out[tile_b.begin, 0, tile_h.begin, tile_m] = states_local
+
+        if HAS_SEQ_IDX:
+            seq_idx_local = hl.zeros([], dtype=seq_idx.dtype)
+
+        # One chunk per step: the recurrence is sequential in tile_c. The v2 kernel
+        # below does the same work with block_size > 1 via an associative scan.
+        for tile_c in hl.tile(nchunks, block_size=1):
+            new_states_local = states[
+                tile_b.begin, tile_c.begin, tile_h.begin, tile_m
+            ].to(torch.float32)
+            dA_cs_local = dA_chunk_cumsum[tile_b.begin, tile_h.begin, tile_c.begin].to(
+                torch.float32
+            )
+            scale = torch.exp(dA_cs_local)
+
+            if HAS_SEQ_IDX:
+                seq_idx_new_local = seq_idx[
+                    tile_b.begin, (tile_c.begin + 1) * chunk_size - 1
+                ]
+                scale_mask = seq_idx_new_local == seq_idx_local
+                scale = torch.where(scale_mask, scale, torch.zeros_like(scale))
+                seq_idx_local = seq_idx_new_local
+
+            states_local = scale * states_local + new_states_local
+
+            if tile_c.begin < nchunks - 1:
+                out[tile_b.begin, tile_c.begin + 1, tile_h.begin, tile_m] = states_local
+            else:
+                final_states[tile_b.begin, tile_h.begin, tile_m] = states_local
+
+    return out, final_states
+
+
+def _affine_combine(left, right):
+    # The recurrence s_i = a_i * s_{i-1} + b_i as composition of the affine maps
+    # T_i(x) = a_i * x + b_i, which is associative and so scannable:
+    # T_right o T_left: (a_l, b_l), (a_r, b_r) -> (a_r*a_l, a_r*b_l + b_r)
+    a1, b1 = left
+    a2, b2 = right
+    return a1 * a2, a2 * b1 + b2
+
+
+@json_cached_autotune
+@helion.kernel(
+    autotune_accuracy_check=False,
+    autotune_baseline_fn=_state_passing_fwd_ref,
+    ignore_warnings=[helion.exc.TensorOperationInWrapper],
+)
+def _state_passing_fwd_kernel_v2(
+    states: torch.Tensor,
+    dA_chunk_cumsum: torch.Tensor,
+    initial_states: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+    chunk_size: Optional[int] = None,
+    out_dtype: Optional[torch.dtype] = None,
+):
+    batch, nchunks, nheads, dim = states.shape
+    assert dA_chunk_cumsum.shape == (batch, nheads, nchunks)
+
+    if initial_states is not None:
+        assert initial_states.shape == (batch, nheads, dim)
+
+    if seq_idx is not None:
+        assert chunk_size is not None
+        seqlen = seq_idx.shape[-1]
+        assert seq_idx.shape == (batch, seqlen)
+
+    out_dtype = states.dtype if out_dtype is None else out_dtype
+    out = torch.empty(
+        (batch, nchunks, nheads, dim),
+        device=states.device,
+        dtype=out_dtype,
+    )
+    final_states = torch.empty(
+        (batch, nheads, dim),
+        device=states.device,
+        dtype=torch.float32,
+    )
+
+    HAS_INITSTATES = initial_states is not None
+    HAS_SEQ_IDX = seq_idx is not None
+    if chunk_size is not None:
+        chunk_size = hl.specialize(chunk_size)
+
+    block_m = hl.register_block_size(dim)
+    block_c = hl.register_block_size(nchunks)
+
+    for tile_b, tile_h, tile_m in hl.tile(
+        [batch, nheads, dim],
+        block_size=[1, 1, block_m],
+    ):
+        if not HAS_INITSTATES:
+            states_local = hl.zeros((block_m,), dtype=torch.float32)
+        else:
+            states_local = initial_states[tile_b.begin, tile_h.begin, tile_m].to(
+                torch.float32
+            )
+
+        out[tile_b.begin, 0, tile_h.begin, tile_m] = states_local
+
+        for tile_c in hl.tile(nchunks, block_size=block_c):
+            new_states_local = states[tile_b.begin, tile_c, tile_h.begin, tile_m].to(
+                torch.float32
+            )
+            dA_cs_local = dA_chunk_cumsum[tile_b.begin, tile_h.begin, tile_c].to(
+                torch.float32
+            )
+            scale = torch.exp(dA_cs_local)
+
+            if HAS_SEQ_IDX:
+                # On a partial last tile tile_c.index runs past nchunks, and with it
+                # the seq_idx index past seqlen, so the load has to be masked.
+                chunk_valid = tile_c.index < nchunks
+                seq_idx_new_local = hl.load(
+                    seq_idx,
+                    [tile_b.begin, (tile_c.index + 1) * chunk_size - 1],
+                    extra_mask=chunk_valid,
+                )
+                # Triton initializes seq_idx = 0 before the chunk loop
+                # (ssd_state_passing.py:71). Reading this slot as seq_idx[b, -1] at
+                # chunk 0 would go out of bounds (batch 0) or pick up the previous
+                # batch row's last value (batch > 0), so substitute the constant 0.
+                first_chunk = tile_c.index == 0
+                seq_idx_local = hl.load(
+                    seq_idx,
+                    [tile_b.begin, tile_c.index * chunk_size - 1],
+                    extra_mask=chunk_valid & (tile_c.index > 0),
+                )
+                seq_idx_local = torch.where(
+                    first_chunk, torch.zeros_like(seq_idx_local), seq_idx_local
+                )
+                scale_mask = seq_idx_new_local == seq_idx_local
+                scale = torch.where(scale_mask, scale, torch.zeros_like(scale))
+
+            # scale is broadcast to (block_c, block_m) explicitly: both elements of
+            # the scan tuple must have the same shape.
+            A_cum, B_cum = hl.associative_scan(
+                _affine_combine,
+                (scale[:, None].expand(-1, block_m), new_states_local),
+                dim=0,
+            )
+
+            tmp_states_local = A_cum * states_local[None, :] + B_cum
+            # Carry the scan's last row into the next tile. A kernel tensor cannot be
+            # indexed by an integer, so the row is extracted with a one-hot mask + sum.
+            # On a partial last tile lane block_c - 1 is out of range and would carry
+            # garbage, so key the mask on tile_c.end - 1, the last valid index.
+            last_mask = tile_c.index == (tile_c.end - 1)
+            states_local = torch.where(
+                last_mask[:, None], tmp_states_local, torch.zeros_like(tmp_states_local)
+            ).sum(dim=0)
+
+            # Shifted by one: out[c + 1] is the state entering chunk c + 1. The mask
+            # drops the last chunk's, which becomes final_states instead.
+            hl.store(
+                out,
+                (tile_b.begin, tile_c.index + 1, tile_h.begin, tile_m),
+                tmp_states_local,
+                extra_mask=(tile_c.index + 1 < nchunks)[:, None],
+            )
+
+        final_states[tile_b.begin, tile_h.begin, tile_m] = states_local
+
+    return out, final_states
+
+
+def _state_passing_fwd_v1(
+    states: torch.Tensor,
+    dA_chunk_cumsum: torch.Tensor,
+    initial_states: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+    chunk_size: Optional[int] = None,
+    out_dtype: Optional[torch.dtype] = None,
+):
+    return _state_passing_fwd_kernel_v1(
+        states, dA_chunk_cumsum, initial_states, seq_idx, chunk_size, out_dtype
+    )
+
+
+def _state_passing_fwd_v2(
+    states: torch.Tensor,
+    dA_chunk_cumsum: torch.Tensor,
+    initial_states: Optional[torch.Tensor] = None,
+    seq_idx: Optional[torch.Tensor] = None,
+    chunk_size: Optional[int] = None,
+    out_dtype: Optional[torch.dtype] = None,
+):
+    return _state_passing_fwd_kernel_v2(
+        states, dA_chunk_cumsum, initial_states, seq_idx, chunk_size, out_dtype
+    )
diff --git a/mamba_ssm/ops/helion/dispatch.py b/mamba_ssm/ops/helion/dispatch.py
new file mode 100644
index 000000000..25ca0a621
--- /dev/null
+++ b/mamba_ssm/ops/helion/dispatch.py
@@ -0,0 +1,86 @@
+import os
+import warnings
+from functools import cache
+from types import SimpleNamespace
+
+from ...utils.determinism import use_deterministic_mode
+
+
+MAMBA_USE_HELION = "MAMBA_USE_HELION"
+
+_deterministic_warning_issued = False
+
+
+@cache
+def _helion_opt_in() -> bool:
+    """Return the process-level Helion opt-in captured at first dispatch."""
+    return os.environ.get(MAMBA_USE_HELION, "0") == "1"
+
+
+def use_helion() -> bool:
+    """Return whether Helion kernels are selected, warning about determinism.
+
+    The env read stays cached, but the determinism check does not: deterministic
+    mode can be turned on after the first dispatch, and every call site asks on
+    every dispatch, so this catches that too. Warns at most once per process.
+    """
+    enabled = _helion_opt_in()
+    if enabled:
+        _warn_if_deterministic()
+    return enabled
+
+
+def _warn_if_deterministic() -> None:
+    """Warn that the Helion kernels ignore deterministic mode.
+
+    Unlike the Triton kernels, which route their reductions through
+    ``alloc_tile_workspace``/``finalize_tile_workspace`` when
+    ``use_deterministic_mode()`` is on, the Helion backward kernels decide at
+    compile time and never ask.
+    """
+    global _deterministic_warning_issued
+    if _deterministic_warning_issued or not use_deterministic_mode():
+        return
+    _deterministic_warning_issued = True
+    warnings.warn(
+        "MAMBA_USE_HELION=1 while deterministic mode is enabled, but the Helion "
+        "kernels do not honour it: chunk_scan_chunk_state_bwd_dx (ddt and 1-D dD) "
+        "and chunk_scan_bwd_dC (ddA_cumsum_prev) always reduce with atomic adds, "
+        "so the backward is not bitwise reproducible. Unset MAMBA_USE_HELION for "
+        "the deterministic Triton path.",
+        stacklevel=2,
+    )
+
+
+@cache
+def get_helion_ssd_kernels():
+    """Lazily import Helion SSD implementations only after explicit opt-in."""
+    from ._chunk_scan_bwd_dC import _chunk_scan_bwd_dC
+    from ._chunk_scan_bwd_ddAcs_stable import _chunk_scan_bwd_ddAcs_stable
+    from ._chunk_scan_bwd_dstates import _chunk_scan_bwd_dstates
+    from ._chunk_scan_chunk_state_bwd_dx import _chunk_scan_chunk_state_bwd_dx
+    from ._chunk_scan_fwd import _chunk_scan_fwd
+    from ._chunk_state_bwd_db import _chunk_state_bwd_db
+    from ._chunk_state_fwd import _chunk_state_fwd
+    from ._state_passing_bwd import _state_passing_bwd
+    from ._state_passing_fwd import _state_passing_fwd_v2
+
+    return SimpleNamespace(
+        chunk_scan_bwd_dC=_chunk_scan_bwd_dC,
+        chunk_scan_bwd_ddAcs_stable=_chunk_scan_bwd_ddAcs_stable,
+        chunk_scan_bwd_dstates=_chunk_scan_bwd_dstates,
+        chunk_scan_chunk_state_bwd_dx=_chunk_scan_chunk_state_bwd_dx,
+        chunk_scan_fwd=_chunk_scan_fwd,
+        chunk_state_bwd_db=_chunk_state_bwd_db,
+        chunk_state_fwd=_chunk_state_fwd,
+        state_passing_bwd=_state_passing_bwd,
+        state_passing_fwd=_state_passing_fwd_v2,
+    )
+
+
+@cache
+def get_helion_layer_norm_bwd():
+    """Lazily import the Helion gated LayerNorm backward implementation."""
+    from ._layer_norm_bwd import _layer_norm_bwd
+
+    return _layer_norm_bwd
diff --git a/mamba_ssm/ops/helion/mamba_mixer_min.py b/mamba_ssm/ops/helion/mamba_mixer_min.py
new file mode 100644
index 000000000..9b2d03ec5
--- /dev/null
+++ b/mamba_ssm/ops/helion/mamba_mixer_min.py
@@ -0,0 +1,197 @@
+"""Minimal, Megatron-free replica of Megatron-LM's ``MambaMixer`` training path,
+so the Helion kernels can be autotuned / profiled inside the mamba repo against
+the exact kernel set production uses.
+
+Unlike ``mamba_ssm.modules.mamba2.Mamba2`` -- which *fuses* the gated RMSNorm and
+out_proj into ``mamba_split_conv1d_scan_combined`` and so runs the plain Triton
+``_layer_norm_bwd`` -- this mirrors Megatron's ``MambaMixer._ssm_training``: the
+kernel is called WITHOUT ``rmsnorm_weight`` / ``outproj_weight`` and a separate
+``RMSNormGated`` runs afterwards, routing through Helion's ``_layer_norm_bwd``.
+The un-fused layout must be mirrored exactly to autotune the right kernels.
+
+Reference: ``MambaMixer._ssm_training`` in Megatron-LM at commit 571370c --
+https://github.com/NVIDIA/Megatron-LM/blob/571370c829ca768fe37244f4e2e7f28d8accc4ab/megatron/core/ssm/mamba_mixer.py#L683
+Verified bit-for-bit identical (fwd+bwd) to a ``Float16Module``-wrapped Megatron
+``MambaMixer`` at tp=cp=1, except for rounding-level in_proj/out_proj differences
+(``nn.Linear`` here vs Megatron's TE/fused bf16 GEMM) that don't affect the kernels.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+
+from einops import rearrange
+
+from mamba_ssm.ops.triton.layernorm_gated import RMSNorm as RMSNormGated
+from mamba_ssm.ops.triton.ssd_combined import mamba_split_conv1d_scan_combined
+
+
+class MambaMixerMin(nn.Module):
+    """Single-GPU MambaMixer mirroring Megatron's training-path layout.
+
+    Forward expects ``hidden_states`` of shape ``(seqlen, batch, d_model)`` (the
+    ``l b d`` layout Megatron uses) and returns the same shape.
+    """
+
+    def __init__(
+        self,
+        d_model: int,
+        nheads: int,
+        headdim: int = 64,
+        d_state: int = 128,
+        ngroups: int = 8,
+        d_conv: int = 4,
+        chunk_size: int = 128,
+        rmsnorm: bool = True,
+        norm_before_gate: bool = False,
+        D_has_hdim: bool = False,
+        bias: bool = False,
+        conv_bias: bool = True,
+        A_init_range: Tuple[float, float] = (1, 16),
+        dt_min: float = 0.001,
+        dt_max: float = 0.1,
+        dt_init_floor: float = 1e-4,
+        dt_limit: Tuple[float, float] = (0.0, float("inf")),
+        device=None,
+        dtype=None,
+    ):
+        factory_kwargs = {"device": device, "dtype": dtype}
+        super().__init__()
+
+        assert not bias
+        assert not norm_before_gate
+        assert nheads % ngroups == 0, "nheads must be evenly divisible by ngroups"
+
+        self.d_model = d_model
+        self.nheads = nheads
+        self.headdim = headdim
+        self.d_state = d_state
+        self.ngroups = ngroups
+        self.d_conv = d_conv
+        self.chunk_size = chunk_size
+        self.rmsnorm = rmsnorm
+        self.norm_before_gate = norm_before_gate
+        self.D_has_hdim = D_has_hdim
+        self.dt_limit = dt_limit
+        self.activation = "silu"
+        self.d_inner = nheads * headdim
+
+        # Order: [z, x, B, C, dt]
+        d_in_proj = 2 * self.d_inner + 2 * self.ngroups * self.d_state + self.nheads
+        self.in_proj = nn.Linear(self.d_model, d_in_proj, bias=bias, **factory_kwargs)
+
+        conv_dim = self.d_inner + 2 * self.ngroups * self.d_state  # x B C
+        self.conv1d = nn.Conv1d(
+            in_channels=conv_dim,
+            out_channels=conv_dim,
+            bias=conv_bias,
+            kernel_size=d_conv,
+            groups=conv_dim,
+            padding=d_conv - 1,
+            **factory_kwargs,
+        )
+
+        self.act = nn.SiLU()
+
+        # dt bias so that softplus(dt_bias) lands in (dt_min, dt_max).
+        dt = torch.exp(
+            torch.rand(self.nheads, **factory_kwargs)
+            * (math.log(dt_max) - math.log(dt_min))
+            + math.log(dt_min)
+        ).clamp(min=dt_init_floor)
+        inv_dt = dt + torch.log(-torch.expm1(-dt))  # inverse softplus
+        self.dt_bias = nn.Parameter(inv_dt)
+
+        assert A_init_range[0] > 0 and A_init_range[1] >= A_init_range[0]
+        A = torch.empty(self.nheads, dtype=torch.float32, device=device).uniform_(*A_init_range)
+        self.A_log = nn.Parameter(torch.log(A))  # keep A_log in fp32
+
+        self.D = nn.Parameter(
+            torch.ones(self.d_inner if self.D_has_hdim else self.nheads, device=device)
+        )  # keep in fp32
+
+        if self.rmsnorm:
+            self.norm = RMSNormGated(
+                self.d_inner,
+                eps=1e-5,
+                group_size=self.d_inner // self.ngroups,
+                norm_before_gate=self.norm_before_gate,
+                **factory_kwargs,
+            )
+
+        self.out_proj = nn.Linear(self.d_inner, self.d_model, bias=bias, **factory_kwargs)
+
+        # Megatron -- and so this constructor above -- initializes D and the gated-norm
+        # weight to all ones. Randomize them: with every element equal to 1, a kernel
+        # that indexes D by the wrong head, or normalizes with the wrong group, still
+        # produces the right answer, so those bugs would be invisible to the tests or to
+        # autotune_baseline_fn's accuracy check against Triton. A trained checkpoint has
+        # non-trivial values anyway; the norm weight is kept around the ones-init
+        # magnitude so the norm still behaves like a trained one.
+        with torch.no_grad():
+            self.D.normal_()
+            if self.rmsnorm:
+                self.norm.weight.uniform_(0.5, 1.5)
+
+        # Production wraps the whole model in Float16Module, whose .bfloat16()/.half()
+        # casts EVERY float param to params_dtype -- including A_log and D, which the
+        # constructor above deliberately created in fp32. The Megatron profiler
+        # replicates this with Float16Module(...).bfloat16(); without it, D reaches
+        # the Helion chunk-scan kernels as fp32 instead of bf16/fp16 (it is passed in
+        # without a .float() in the forward), diverging from real training. Mirror
+        # that whole-module cast here so the kernels see production dtypes.
+        if dtype in (torch.float16, torch.bfloat16):
+            self.to(dtype)
+
+    def forward(
+        self, hidden_states: torch.Tensor, seq_idx: Optional[torch.Tensor] = None
+    ) -> torch.Tensor:
+        """hidden_states: (seqlen, batch, d_model) -> (seqlen, batch, d_model)."""
+        _, _, dim = hidden_states.shape
+        assert dim == self.d_model
+
+        zxBCdt = self.in_proj(hidden_states)  # (l, b, d_in_proj)
+
+        # The Helion / Triton split-conv-scan kernel works in (b, l, d) layout.
+        zxBCdt = rearrange(zxBCdt, "l b d -> b l d").contiguous()
+
+        A = -torch.exp(self.A_log.float())  # (nheads,)
+
+        dt_limit_kwargs = (
+            {} if self.dt_limit == (0.0, float("inf")) else dict(dt_limit=self.dt_limit)
+        )
+
+        # NOTE: rmsnorm_weight / outproj_weight are intentionally NOT passed, so the
+        # gated RMSNorm and out_proj stay un-fused (the Megatron training layout).
+        y = mamba_split_conv1d_scan_combined(
+            zxBCdt,
+            rearrange(self.conv1d.weight, "d 1 w -> d w"),
+            self.conv1d.bias,
+            self.dt_bias.float(),
+            A,
+            D=(
+                rearrange(self.D.float(), "(h p) -> h p", p=self.headdim)
+                if self.D_has_hdim
+                else self.D
+            ),
+            chunk_size=self.chunk_size,
+            activation=self.activation,
+            headdim=None if self.D_has_hdim else self.headdim,
+            ngroups=self.ngroups,
+            norm_before_gate=self.norm_before_gate,
+            seq_idx=seq_idx,
+            **dt_limit_kwargs,
+        )
+
+        y = rearrange(y, "b l d -> l b d").contiguous()
+
+        if self.rmsnorm:
+            # Separate gated RMSNorm -> routes through Helion _layer_norm_bwd.
+            y = self.norm(y)
+
+        out = self.out_proj(y)
+        return out
diff --git a/mamba_ssm/ops/helion/utils.py b/mamba_ssm/ops/helion/utils.py
new file mode 100644
index 000000000..2d81155cf
--- /dev/null
+++ b/mamba_ssm/ops/helion/utils.py
@@ -0,0 +1,106 @@
+import os
+from functools import wraps
+from pathlib import Path
+from typing import Optional, Sequence
+
+import helion
+
+def save_helion_config(config: helion.Config, kernel_name: str) -> None:
+    """Save an autotuned Helion config as ``<kernel_name>.json``.
+
+    The output directory is read from the ``MAMBA_HELION_CONFIG_DIR``
+    environment variable, which must be set. Overwrites any existing file, which
+    is how a forced re-autotune (``MAMBA_HELION_AUTOTUNE=2``) replaces a config.
+    """
+    config_dir = os.environ.get("MAMBA_HELION_CONFIG_DIR")
+    assert config_dir, "MAMBA_HELION_CONFIG_DIR environment variable is not set"
+    path = Path(config_dir) / f"{kernel_name}.json"
+    config.save(path)
+    print(f"[helion] saved config for {kernel_name} to {path}")
+
+
+def load_helion_config(kernel_name: str) -> Optional[helion.Config]:
+    """Load a saved Helion config for ``<kernel_name>.json``.
+
+    The directory is read from the ``MAMBA_HELION_CONFIG_DIR`` environment
+    variable, which must be set. Returns ``None`` -- telling the caller to
+    autotune -- in two cases driven by ``MAMBA_HELION_AUTOTUNE``:
+
+    * ``=1`` (missing-only): no config file exists yet. An existing file is
+      loaded and reused, so a crashed multi-kernel run resumes where it left off.
+    * ``=2`` (force): always re-autotune and overwrite, even if a config file
+      already exists. Use this to re-tune for a new GPU / Helion version.
+
+    Otherwise a missing config is a hard error: autotuning is an offline,
+    single-process step and must never run at runtime (e.g. under CUDA graph
+    capture or distributed jobs).
+    """
+    config_dir = os.environ.get("MAMBA_HELION_CONFIG_DIR")
+    assert config_dir, "MAMBA_HELION_CONFIG_DIR environment variable is not set"
+    autotune = os.environ.get("MAMBA_HELION_AUTOTUNE")
+    path = Path(config_dir) / f"{kernel_name}.json"
+    if autotune == "2":
+        # Force: re-autotune and overwrite any existing config.
+        return None
+    if not path.is_file():
+        assert autotune == "1", (
+            f"[helion] config for {kernel_name} not found at {path}; "
+            f"run offline with MAMBA_HELION_AUTOTUNE=1 to generate it"
+        )
+        return None
+    config = helion.Config.load(path)
+    print(f"[helion] loaded config for {kernel_name} from {path}")
+    return config
+
+
+def ensure_helion_config(kernel: helion.Kernel, args: Sequence[object]) -> None:
+    """Make sure ``kernel`` has a config, loading from disk or autotuning.
+
+    Uses ``kernel.name`` as the config filename, so callers don't repeat the
+    kernel name as a string. On first call, load ``<kernel.name>.json`` and
+    install it as the kernel's single config so every specialization (e.g.
+    forward and backward calls with slightly different inputs) reuses it. A
+    missing config is a hard error unless ``MAMBA_HELION_AUTOTUNE`` is set, in
+    which case it autotunes with the example ``args`` and saves the result (an
+    offline, single-process step only): ``=1`` only fills in missing configs,
+    ``=2`` forces a re-autotune and overwrites any existing one. No-op once a
+    config is set.
+    """
+    if kernel.configs:
+        return
+    config = load_helion_config(kernel.name)
+    if config is None:
+        config = kernel.autotune(args)
+        save_helion_config(config, kernel.name)
+    kernel.configs = [config]
+
+
+def json_cached_autotune(kernel):
+    """Wrap a ``helion.Kernel`` to lazily resolve its config on first call: load
+    ``<kernel.name>.json`` from ``MAMBA_HELION_CONFIG_DIR``, or autotune and save
+    that JSON when none exists (``MAMBA_HELION_AUTOTUNE=1``) or when a re-autotune
+    is forced (``MAMBA_HELION_AUTOTUNE=2``, overwriting any existing config).
+
+    The autotune arguments come from ``kernel.normalize_args`` on the actual call, so
+    no call site has to restate them. The wrapper stays callable exactly like the
+    kernel and exposes it as ``.kernel``.
+    """
+
+    @wraps(kernel)
+    def wrapper(*args, **kwargs):
+        if not kernel.configs:
+            ensure_helion_config(kernel, kernel.normalize_args(*args, **kwargs))
+        # Triton's allocator is a ContextVar, so it does not cross thread boundaries:
+        # PyTorch autograd runs backward kernels on its own worker thread, which sees
+        # the default NullAllocator rather than whatever the forward thread installed.
+        # Helion only emits its own set_triton_allocator() at the top level of the
+        # generated module (so it runs once, on the compiling thread) and only for
+        # kernels taking tensor descriptors. Any kernel needing global scratch --
+        # TMA descriptors, or a persistent kernel's workspace -- would then fail at
+        # launch. Calling it here is idempotent: it no-ops unless the allocator is
+        # still NullAllocator.
+        helion.runtime.set_triton_allocator()
+        return kernel(*args, **kwargs)
+
+    wrapper.kernel = kernel
+    return wrapper
diff --git a/mamba_ssm/ops/triton/layernorm_gated.py b/mamba_ssm/ops/triton/layernorm_gated.py
index de4b2f481..952bf3414 100644
--- a/mamba_ssm/ops/triton/layernorm_gated.py
+++ b/mamba_ssm/ops/triton/layernorm_gated.py
@@ -14,6 +14,8 @@
 
 from einops import rearrange
 
+from ..helion.dispatch import get_helion_layer_norm_bwd, use_helion
+
 
 def rms_norm_ref(x, weight, bias, z=None, eps=1e-6, group_size=None, norm_before_gate=True, upcast=True):
     dtype = x.dtype
@@ -372,8 +374,9 @@ def backward(ctx, dy):
         if dy.stride(-1) != 1:
             dy = dy.contiguous()
         assert dy.shape == x.shape
-        dx, dw, db, dz = _layer_norm_bwd(dy, x, weight, bias, ctx.eps, mean, rstd, z, ctx.group_size,
-                                         ctx.norm_before_gate, ctx.is_rms_norm)
+        layer_norm_bwd = get_helion_layer_norm_bwd() if use_helion() else _layer_norm_bwd
+        dx, dw, db, dz = layer_norm_bwd(dy, x, weight, bias, ctx.eps, mean, rstd, z, ctx.group_size,
+                                        ctx.norm_before_gate, ctx.is_rms_norm)
         return dx.reshape(ctx.x_shape_og), dw, db, dz.reshape(ctx.x_shape_og) if dz is not None else None, None, None, None, None
 
 
diff --git a/mamba_ssm/ops/triton/ssd_combined.py b/mamba_ssm/ops/triton/ssd_combined.py
index af95537b9..834b7ae15 100644
--- a/mamba_ssm/ops/triton/ssd_combined.py
+++ b/mamba_ssm/ops/triton/ssd_combined.py
@@ -7,6 +7,7 @@
 
 import math
 from packaging import version
+from types import SimpleNamespace
 
 import torch
 import torch.nn.functional as F
@@ -49,6 +50,8 @@
     use_deterministic_mode,
 )
 
+from ..helion.dispatch import get_helion_ssd_kernels, use_helion
+
 TRITON_22 = version.parse(triton.__version__) >= version.parse('2.2.0')
 
 
@@ -340,7 +343,25 @@ def _chunk_scan_chunk_state_bwd_dx(x, dt, dA_cumsum, B, CB, dout, dstates, D=Non
     return dx, ddt, dD
 
 
+_TRITON_SSD_KERNELS = SimpleNamespace(
+    chunk_scan_bwd_dC=_chunk_scan_bwd_dC,
+    chunk_scan_bwd_ddAcs_stable=_chunk_scan_bwd_ddAcs_stable,
+    chunk_scan_bwd_dstates=_chunk_scan_bwd_dstates,
+    chunk_scan_chunk_state_bwd_dx=_chunk_scan_chunk_state_bwd_dx,
+    chunk_scan_fwd=_chunk_scan_fwd,
+    chunk_state_bwd_db=_chunk_state_bwd_db,
+    chunk_state_fwd=_chunk_state_fwd,
+    state_passing_bwd=_state_passing_bwd,
+    state_passing_fwd=_state_passing_fwd,
+)
+
+
+def _ssd_kernel_impls():
+    return get_helion_ssd_kernels() if use_helion() else _TRITON_SSD_KERNELS
+
+
 def _mamba_chunk_scan_combined_fwd(x, dt, A, B, C, chunk_size, D=None, z=None, dt_bias=None, initial_states=None, seq_idx=None, cu_seqlens=None, dt_softplus=False, dt_limit=(0.0, float("inf")), state_dtype=None):
+    kernels = _ssd_kernel_impls()
     batch, seqlen, nheads, headdim = x.shape
     _, _, ngroups, dstate = B.shape
     assert nheads % ngroups == 0
@@ -372,11 +393,11 @@ def _mamba_chunk_scan_combined_fwd(x, dt, A, B, C, chunk_size, D=None, z=None, d
     # dA_cumsum_tmp1, dt_tmp1 = _chunk_cumsum_fwd(dt[:, 147:], A, chunk_size, dt_bias=dt_bias, dt_softplus=dt_softplus)
     # dA_cumsum_tmp2, dt_tmp2 = _chunk_cumsum_fwd(dt[:, 147:256], A, chunk_size, dt_bias=dt_bias, dt_softplus=dt_softplus)
     dA_cumsum, dt = _chunk_cumsum_fwd(dt, A, chunk_size, dt_bias=dt_bias, dt_softplus=dt_softplus, dt_limit=dt_limit)
-    states = _chunk_state_fwd(B, x, dt, dA_cumsum, seq_idx=seq_idx, states_in_fp32=True)
+    states = kernels.chunk_state_fwd(B, x, dt, dA_cumsum, seq_idx=seq_idx, states_in_fp32=True)
     # states_tmp0 = _chunk_state_fwd(B[:, :147], x[:, :147], dt_tmp0, dA_cumsum_tmp0, states_in_fp32=True)
     # states_tmp1 = _chunk_state_fwd(B[:, 147:], x[:, 147:], dt_tmp1, dA_cumsum_tmp1, states_in_fp32=True)
     # states_tmp2 = _chunk_state_fwd(B[:, 147:256], x[:, 147:256], dt_tmp2, dA_cumsum_tmp2, states_in_fp32=True)
-    states, final_states = _state_passing_fwd(rearrange(states, "... p n -> ... (p n)"), dA_cumsum[:, :, :, -1],
+    states, final_states = kernels.state_passing_fwd(rearrange(states, "... p n -> ... (p n)"), dA_cumsum[:, :, :, -1],
                                               initial_states=rearrange(initial_states, "... p n -> ... (p n)") if initial_states is not None else None,
                                               seq_idx=seq_idx, chunk_size=chunk_size,
                                               out_dtype=state_dtype if state_dtype is not None else C.dtype)
@@ -384,7 +405,7 @@ def _mamba_chunk_scan_combined_fwd(x, dt, A, B, C, chunk_size, D=None, z=None, d
     # states_tmp0 = rearrange(_state_passing_fwd(rearrange(states_tmp0, "... p n -> ... (p n)"), dA_cumsum_tmp0[:, :, :, -1], chunk_size=chunk_size), "... (p n) -> ... p n", n=dstate)
     # states_tmp1 = rearrange(_state_passing_fwd(rearrange(states_tmp1, "... p n -> ... (p n)"), dA_cumsum_tmp1[:, :, :, -1], chunk_size=chunk_size), "... (p n) -> ... p n", n=dstate)
     CB = _bmm_chunk_fwd(C, B, chunk_size, seq_idx=seq_idx, output_dtype=torch.float32)
-    out, out_x = _chunk_scan_fwd(CB, x, dt, dA_cumsum, C, states, D=D, z=z, seq_idx=seq_idx)
+    out, out_x = kernels.chunk_scan_fwd(CB, x, dt, dA_cumsum, C, states, D=D, z=z, seq_idx=seq_idx)
     if cu_seqlens is None:
         return out, out_x, dt, dA_cumsum, states, final_states
     else:
@@ -399,6 +420,7 @@ def _mamba_chunk_scan_combined_bwd(dout, x, dt, A, B, C, out, chunk_size, D=None
                                    dt_limit=(0.0, float("inf")),
                                    dx=None, ddt=None, dB=None, dC=None, dz=None, recompute_output=False,
                                    state_dtype=None):
+    kernels = _ssd_kernel_impls()
     if dout.stride(-1) != 1:
         dout = dout.contiguous()
     batch, seqlen, nheads, headdim = x.shape
@@ -441,8 +463,8 @@ def _mamba_chunk_scan_combined_bwd(dout, x, dt, A, B, C, out, chunk_size, D=None
     dA_cumsum, dt = _chunk_cumsum_fwd(dt_in, A, chunk_size, dt_bias=dt_bias, dt_softplus=dt_softplus,
                                       dt_limit=dt_limit)
     CB = _bmm_chunk_fwd(C, B, chunk_size, seq_idx=seq_idx, output_dtype=torch.float32)
-    states = _chunk_state_fwd(B, x, dt, dA_cumsum, seq_idx=seq_idx, states_in_fp32=True)
-    states, _ = _state_passing_fwd(rearrange(states, "... p n -> ... (p n)"), dA_cumsum[:, :, :, -1],
+    states = kernels.chunk_state_fwd(B, x, dt, dA_cumsum, seq_idx=seq_idx, states_in_fp32=True)
+    states, _ = kernels.state_passing_fwd(rearrange(states, "... p n -> ... (p n)"), dA_cumsum[:, :, :, -1],
                                    initial_states=rearrange(initial_states, "... p n -> ... (p n)") if initial_states is not None else None,
                                    seq_idx=seq_idx, chunk_size=chunk_size)
     states = rearrange(states, "... (p n) -> ... p n", n=dstate)
@@ -452,13 +474,13 @@ def _mamba_chunk_scan_combined_bwd(dout, x, dt, A, B, C, out, chunk_size, D=None
     else:
         dz = None
         outz = out
-    dstates = _chunk_scan_bwd_dstates(C, dA_cumsum, dout, seq_idx=seq_idx, dtype=states.dtype)
+    dstates = kernels.chunk_scan_bwd_dstates(C, dA_cumsum, dout, seq_idx=seq_idx, dtype=states.dtype)
     # dstates has length nchunks, containing the gradient to initial states at index 0 and
     # gradient to the states of chunk (nchunks - 2) at index (nchunks - 1)
     # Do computation in fp32 but convert dstates and states to fp16/bf16 since dstates and states
     # will be used in matmul in the next kernels. When state_dtype is set, keep them in that
     # dtype instead so the backward consumes states at the same precision the forward used.
-    dstates, ddA_chunk_cumsum, dinitial_states, states = _state_passing_bwd(
+    dstates, ddA_chunk_cumsum, dinitial_states, states = kernels.state_passing_bwd(
         rearrange(states, "... p n -> ... (p n)"),
         dA_cumsum[:, :, :, -1],
         rearrange(dstates, "... p n -> ... (p n)"),
@@ -476,11 +498,11 @@ def _mamba_chunk_scan_combined_bwd(dout, x, dt, A, B, C, out, chunk_size, D=None
     states = rearrange(states, "... (p n) -> ... p n", n=dstate)
     dstates = rearrange(dstates, "... (p n) -> ... p n", n=dstate)
     dinitial_states = rearrange(dinitial_states, "... (p n) -> ... p n", n=dstate) if dinitial_states is not None else None
-    dx, ddt, dD_from_x = _chunk_scan_chunk_state_bwd_dx(x, dt, dA_cumsum, B, CB, dout, dstates, D=D, seq_idx=seq_idx, dx=dx)
+    dx, ddt, dD_from_x = kernels.chunk_scan_chunk_state_bwd_dx(x, dt, dA_cumsum, B, CB, dout, dstates, D=D, seq_idx=seq_idx, dx=dx)
     # dB = _chunk_state_bwd_db(x, dt, dA_cumsum, dstates, seq_idx=seq_idx, ngroups=ngroups)
-    dB, ddA_next = _chunk_state_bwd_db(x, dt, dA_cumsum, dstates, seq_idx=seq_idx, B=B, ngroups=ngroups)
+    dB, ddA_next = kernels.chunk_state_bwd_db(x, dt, dA_cumsum, dstates, seq_idx=seq_idx, B=B, ngroups=ngroups)
     # dC = _chunk_scan_bwd_dC(states[:, :-1].to(x.dtype), dA_cumsum, dout, seq_idx=seq_idx, ngroups=ngroups)
-    dC, ddA_cumsum_prev = _chunk_scan_bwd_dC(states.to(state_dtype if state_dtype is not None else x.dtype), dA_cumsum, dout, seq_idx=seq_idx, C=C, ngroups=ngroups)
+    dC, ddA_cumsum_prev = kernels.chunk_scan_bwd_dC(states.to(state_dtype if state_dtype is not None else x.dtype), dA_cumsum, dout, seq_idx=seq_idx, C=C, ngroups=ngroups)
     # Computing ddA with the dcb kernel is much slower, so we're not using it for now
     dCB = _chunk_scan_bwd_dcb(x, dt, dA_cumsum, dout, seq_idx=seq_idx, ngroups=ngroups)
     # dCB, ddA_tmp = _chunk_scan_bwd_dcb(x, dt, dA_cumsum, dout, seq_idx=seq_idx, CB=CB, ngroups=ngroups)
@@ -503,7 +525,7 @@ def _mamba_chunk_scan_combined_bwd(dout, x, dt, A, B, C, out, chunk_size, D=None
     # This is already done as part of bwd_dB kernel
     # ddA_next = _chunk_state_bwd_ddAcs_stable(B, x, dt, dA_cumsum, dstates, seq_idx=seq_idx)
     # We don't need to pass in seq_idx because CB also zeros out entries where seq_idx[i] != seq_idx[j]
-    ddA = _chunk_scan_bwd_ddAcs_stable(x, dt, dA_cumsum, dout, CB)
+    ddA = kernels.chunk_scan_bwd_ddAcs_stable(x, dt, dA_cumsum, dout, CB)
     ddA += ddA_next + ddA_prev
 
     ddt_given, dA, ddt_bias = _chunk_cumsum_bwd(ddA, ddt, dt_in, A, dt_bias=dt_bias, dt_softplus=dt_softplus, dt_limit=dt_limit, ddt=ddt_given)
diff --git a/pyproject.toml b/pyproject.toml
index c27c9d508..75cf48bd7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -34,6 +34,14 @@ Repository = "https://github.com/state-spaces/mamba"
 causal-conv1d = [
     "causal-conv1d>=1.2.0"
 ]
+# Opt-in Helion kernel backend (MAMBA_USE_HELION=1); see mamba_ssm/ops/helion/README.md.
+# helion is pinned rather than lower-bounded: a config is generated for one toolchain and
+# Helion's config schema is still moving, so a different version can need re-autotuning
+# (or fail to load a config) even when the kernels themselves are unchanged.
+helion = [
+    "helion==1.4.0",
+    "triton>=3.6.0",
+]
 dev = [
     "pytest"
 ]
diff --git a/tests/ops/helion/test_kernels.py b/tests/ops/helion/test_kernels.py
new file mode 100644
index 000000000..18c8da293
--- /dev/null
+++ b/tests/ops/helion/test_kernels.py
@@ -0,0 +1,377 @@
+"""Compare every Helion kernel with its Triton reference inside a real pass.
+
+One ``MambaMixerMin`` forward+backward runs with the pipeline on Triton, and each
+dispatched kernel is wrapped so that both implementations are invoked on the
+*production* arguments and their outputs compared. The wrapper returns the Triton
+result, so every later stage -- and therefore every other kernel's inputs -- stays
+exact production data and a difference is attributable to the kernel that
+produced it.
+
+The arguments are not rebuilt by hand on purpose. They are views into shared
+buffers (``x``/``B``/``C``/``dx`` slice one packed ``xBC``, ``z`` slices
+``zxbcdt``), several are outputs of preceding stages recomputed in a specific
+order (``state_passing_bwd`` consumes the backward's *recomputed* forward states,
+not ``chunk_state_fwd``'s output), and ``dout`` has already been through the
+z-gating backward. Hand-building all of that is how you end up testing a
+specialization, or a magnitude regime, that production never sees.
+
+``MAMBA_HELION_CONFIG_DIR`` selects the config set. With it unset the tests fall back
+to the checked-in example set so they run with no setup; autotuning stays disabled
+either way, so a config that is missing -- or that does not compile on this GPU --
+fails the run with a message pointing at ``autotune_mamba_mixer``.
+
+One process tests one variant -- unpacked by default, packed (a ragged
+``seq_idx``, exercising every kernel's ``HAS_SEQ_IDX`` branch) with
+``MAMBA_HELION_TEST_PACKED=1`` -- because each needs its own config set::
+
+    pytest tests/ops/helion/test_kernels.py                       # checked-in configs
+    MAMBA_HELION_CONFIG_DIR=<unpacked configs> pytest tests/ops/helion/test_kernels.py
+    MAMBA_HELION_TEST_PACKED=1 MAMBA_HELION_CONFIG_DIR=<packed configs> \
+        pytest tests/ops/helion/test_kernels.py
+
+A config carries one ``indexing`` entry per load/store, and the packed path adds
+the ``seq_idx`` loads, so an unpacked config does not merely run slower there --
+it fails to compile. Generate the packed set with
+``autotune_mamba_mixer --sequence-packing`` into a separate directory.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any
+
+import pytest
+import torch
+
+helion = pytest.importorskip("helion")
+pytest.importorskip("causal_conv1d")
+
+from mamba_ssm.ops.helion.dispatch import (
+    get_helion_layer_norm_bwd,
+    get_helion_ssd_kernels,
+)
+from mamba_ssm.ops.helion.mamba_mixer_min import MambaMixerMin
+from mamba_ssm.ops.triton import layernorm_gated, ssd_combined
+
+pytestmark = pytest.mark.skipif(
+    not torch.cuda.is_available(), reason="Helion kernel tests require CUDA"
+)
+
+
+# --- what this process tests -------------------------------------------------
+
+# The variant cannot be a parametrization: a config autotuned unpacked has one
+# `indexing` entry per load/store and the packed path adds the seq_idx loads, so it does
+# not even compile there -- and Helion pins a kernel's config on its first call, so one
+# process cannot use both config sets.
+_PACKED = os.environ.get("MAMBA_HELION_TEST_PACKED") == "1"
+
+_SSD_KERNELS = (
+    "chunk_state_fwd",
+    "state_passing_fwd",
+    "chunk_scan_fwd",
+    "chunk_scan_bwd_dstates",
+    "state_passing_bwd",
+    "chunk_scan_chunk_state_bwd_dx",
+    "chunk_state_bwd_db",
+    "chunk_scan_bwd_dC",
+    "chunk_scan_bwd_ddAcs_stable",
+)
+_LAYER_NORM_BWD = "layer_norm_bwd"
+
+
+# --- the model the kernels run inside ----------------------------------------
+
+# benchmarks/helion/autotune_mamba_mixer.py defaults (Nemotron3 Nano, seq 8k, bf16).
+_D_MODEL = 2688
+_BATCH = 1
+_SEQLEN = 8192
+_NHEADS = 64
+_HEADDIM = 64
+_DSTATE = 128
+_NGROUPS = 8
+_CHUNK = 128
+_DTYPE = torch.bfloat16
+
+# Ragged sub-sequence lengths summing to _SEQLEN, none of them a multiple of _CHUNK, so
+# every boundary lands mid-chunk -- that is what exercises the kernels' HAS_SEQ_IDX index
+# arithmetic rather than just switching it on.
+_PACKED_SEQLENS = (1531, 2048, 129, 4484)
+
+
+def _build_seq_idx() -> torch.Tensor:
+    """Return a (1, _SEQLEN) int32 seq_idx, as Megatron builds for the THD path."""
+    assert sum(_PACKED_SEQLENS) == _SEQLEN
+    idx = torch.arange(len(_PACKED_SEQLENS), device="cuda")
+    return (
+        torch.repeat_interleave(idx, torch.tensor(_PACKED_SEQLENS, device="cuda"))
+        .to(torch.int32)
+        .unsqueeze(0)
+    )
+
+
+# --- tolerances --------------------------------------------------------------
+
+# Both tolerances are one bf16 ulp. Every kernel here multiplies bf16 operands, so that is
+# the granularity at which two correct implementations may legitimately disagree: when
+# their fp32 accumulators straddle a rounding boundary, the bf16 results differ by exactly
+# 1 ulp with neither being wrong.
+_ULP = 2.0**-8
+_RTOL = _ULP
+
+
+def _atol_for(expected: torch.Tensor) -> float:
+    """One ulp of the tensor's *own* largest magnitude, rather than a constant.
+
+    These outputs span five orders of magnitude in typical element size -- 3e-3 for the
+    chunk states up to 4e2 for dD -- so any fixed atol is either impossibly tight for the
+    large tensors or vacuous for the small ones: at the atol=1e-1 this started with, a
+    deliberate 2% error injected into chunk_scan_fwd went undetected. Scaling by the max
+    rather than the mean is deliberate too, since the differences land on the large
+    elements. The measured worst case is 1.3e-3 of the max (layer_norm_bwd's dx), leaving
+    3x margin, and together the two tolerances detect a systematic error above
+    _RTOL + _ULP = 0.78% -- close to the floor of what a bf16-vs-bf16 comparison can do.
+    """
+    return _ULP * expected.abs().max().item()
+
+
+def _assert_close(name: str, actual: Any, expected: Any) -> None:
+    if actual is None or expected is None:
+        assert actual is expected, f"{name}: one implementation returned None"
+        return
+    if isinstance(actual, tuple) or isinstance(expected, tuple):
+        assert isinstance(actual, tuple) and isinstance(expected, tuple)
+        assert len(actual) == len(expected), f"{name}: output count differs"
+        for index, (actual_item, expected_item) in enumerate(zip(actual, expected)):
+            _assert_close(f"{name}[{index}]", actual_item, expected_item)
+        return
+    try:
+        torch.testing.assert_close(
+            actual, expected, rtol=_RTOL, atol=_atol_for(expected)
+        )
+    except AssertionError as error:
+        raise AssertionError(
+            f"{name}: Helion output differs from Triton\n{error}"
+        ) from error
+    else:
+        diff = (actual - expected).abs().float()
+        print(
+            f"{name}: diff abs mean={diff.mean().item():.3e} max={diff.max().item():.3e} | "
+            f"actual abs mean={actual.float().abs().mean().item():.3e} "
+            f"max={actual.float().abs().max().item():.3e} | "
+            f"expected abs mean={expected.float().abs().mean().item():.3e} "
+            f"max={expected.float().abs().max().item():.3e}"
+        )
+
+
+# --- shadowing the dispatched kernels ----------------------------------------
+
+# Arguments a kernel writes through in place, and the index of the corresponding return
+# value. Triton is given the caller's buffer, as production does, so later stages see the
+# reference result; Helion gets a private one, and the returned tensor must BE the buffer
+# it was handed -- dx is a view into the dxBC that causal_conv1d_bwd reads afterwards, so
+# writing through it is the contract, not returning the right values.
+_IN_PLACE_OUTPUTS = {
+    "chunk_scan_chunk_state_bwd_dx": {"dx": 0},
+}
+
+
+class _ShadowKernel:
+    """Run both implementations on one kernel's production arguments.
+
+    Inputs are passed through untouched rather than copied: they are views into
+    larger buffers and Helion compiles per stride, so a copy would test a
+    different specialization, and none of these kernels writes to an input except
+    through the declared in-place arguments. Returns the Triton result so the
+    surrounding pipeline stays on the reference path.
+
+    Comparison failures are recorded instead of raised -- every kernel is
+    shadowed in a single pass, and one kernel's mismatch must not stop the others
+    from being checked.
+    """
+
+    def __init__(self, name, triton_fn, helion_fn, in_place=None):
+        self.name = name
+        self.triton_fn = triton_fn
+        self.helion_fn = helion_fn
+        self.in_place = in_place or {}
+        self.calls = 0
+        self.errors: list[AssertionError] = []
+
+    def __call__(self, *args, **kwargs):
+        expected = self.triton_fn(*args, **kwargs)
+
+        # NaN-filled so a Helion kernel that leaves part of the buffer unwritten
+        # shows up as a mismatch rather than reading whatever was there.
+        scratch = {
+            param: torch.full_like(kwargs[param], float("nan"))
+            for param in self.in_place
+        }
+        actual = self.helion_fn(*args, **{**kwargs, **scratch})
+
+        label = f"{self.name}#{self.calls}"
+        try:
+            _assert_close(label, actual, expected)
+            for param, index in self.in_place.items():
+                assert expected[index] is kwargs[param], (
+                    f"{label}: Triton did not write the caller's {param}"
+                )
+                assert actual[index] is scratch[param], (
+                    f"{label}: Helion allocated its own {param} instead of "
+                    "writing the caller's buffer"
+                )
+        except AssertionError as error:
+            self.errors.append(error)
+
+        self.calls += 1
+        return expected
+
+
+def _run_shadowed(packed: bool, monkeypatch) -> dict[str, _ShadowKernel]:
+    """Drive one MambaMixerMin fwd+bwd with every dispatched kernel shadowed."""
+    triton_kernels = ssd_combined._TRITON_SSD_KERNELS
+    shadows = {
+        name: _ShadowKernel(
+            name,
+            getattr(triton_kernels, name),
+            getattr(get_helion_ssd_kernels(), name),
+            in_place=_IN_PLACE_OUTPUTS.get(name, {}),
+        )
+        for name in vars(triton_kernels)
+    }
+    monkeypatch.setattr(
+        ssd_combined, "_ssd_kernel_impls", lambda: SimpleNamespace(**shadows)
+    )
+
+    shadows[_LAYER_NORM_BWD] = _ShadowKernel(
+        _LAYER_NORM_BWD,
+        layernorm_gated._layer_norm_bwd,
+        get_helion_layer_norm_bwd(),
+    )
+    monkeypatch.setattr(layernorm_gated, "use_helion", lambda: True)
+    monkeypatch.setattr(
+        layernorm_gated, "get_helion_layer_norm_bwd", lambda: shadows[_LAYER_NORM_BWD]
+    )
+
+    torch.manual_seed(0)
+    mixer = MambaMixerMin(
+        d_model=_D_MODEL,
+        nheads=_NHEADS,
+        headdim=_HEADDIM,
+        d_state=_DSTATE,
+        ngroups=_NGROUPS,
+        chunk_size=_CHUNK,
+        device="cuda",
+        dtype=_DTYPE,
+    )
+    mixer.train()
+    hidden_states = torch.randn(
+        _SEQLEN, _BATCH, _D_MODEL, device="cuda", dtype=_DTYPE, requires_grad=True
+    )
+    out = mixer(hidden_states, seq_idx=_build_seq_idx() if packed else None)
+    (out * torch.randn_like(out)).sum().backward()
+    torch.cuda.synchronize()
+    return shadows
+
+
+# --- the config set the kernels compile with ---------------------------------
+
+# Used when MAMBA_HELION_CONFIG_DIR is unset, so the tests run with no setup. A config is
+# a set of block sizes and indexing choices rather than anything GPU-specific, so this set
+# usually compiles and passes off its own environment too. Update if it is renamed.
+_EXAMPLE_CONFIG_DIR = (
+    Path(__file__).resolve().parents[3]
+    / "benchmarks"
+    / "helion"
+    / "configs"
+    / "b200-helion1.4-triton3.6-nemotron3-nano-bf16-seq8k"
+)
+
+_AUTOTUNE_HINT = (
+    "Autotune a set of your own and point MAMBA_HELION_CONFIG_DIR at it:\n"
+    "    MAMBA_HELION_CONFIG_DIR=<dir> python "
+    "benchmarks/helion/autotune_mamba_mixer.py{packed}"
+)
+
+
+def _resolve_config_dir(monkeypatch):
+    """Point MAMBA_HELION_CONFIG_DIR at a config set, and forbid autotuning.
+
+    Must run before the kernels are first called, so it lives with the module-scoped
+    fixture that drives the pass rather than in an autouse one: pytest instantiates
+    higher-scoped fixtures first, so a function-scoped fixture would patch the
+    environment only after the configs had already been resolved.
+    """
+    if not os.environ.get("MAMBA_HELION_CONFIG_DIR"):
+        # The packed path needs its own set, and none is checked in.
+        assert not _PACKED, (
+            "MAMBA_HELION_TEST_PACKED=1 needs a packed config set of its own; there is "
+            "no checked-in example for it.\n"
+            + _AUTOTUNE_HINT.format(packed=" --sequence-packing --batch-size 1")
+        )
+        assert _EXAMPLE_CONFIG_DIR.is_dir(), (
+            f"no config dir given and the checked-in example set is missing at "
+            f"{_EXAMPLE_CONFIG_DIR} (it lives in the repo, not in the installed "
+            "package, so run from a source checkout).\n"
+            + _AUTOTUNE_HINT.format(packed="")
+        )
+        monkeypatch.setenv("MAMBA_HELION_CONFIG_DIR", str(_EXAMPLE_CONFIG_DIR))
+        print(f"\nMAMBA_HELION_CONFIG_DIR unset; using {_EXAMPLE_CONFIG_DIR}")
+    monkeypatch.delenv("MAMBA_HELION_AUTOTUNE", raising=False)
+
+
+# --- tests -------------------------------------------------------------------
+
+
+@pytest.fixture(
+    scope="module", params=[_PACKED], ids=["packed" if _PACKED else "unpacked"]
+)
+def shadows(request):
+    # monkeypatch is function-scoped, so undo the patches by hand here.
+    patched = pytest.MonkeyPatch()
+    try:
+        _resolve_config_dir(patched)
+        try:
+            yield _run_shadowed(request.param, patched)
+        except Exception as error:
+            # Anything raised out of the pass -- a missing config, or a Helion compile
+            # error from a config that does not fit -- means the config set is wrong for
+            # this environment, which is not obvious from the exception itself.
+            raise RuntimeError(
+                "the Helion pass failed with the config set in "
+                f"{os.environ['MAMBA_HELION_CONFIG_DIR']} -- it may not fit this GPU, "
+                "toolchain version or shape.\n"
+                + _AUTOTUNE_HINT.format(
+                    packed=" --sequence-packing --batch-size 1" if request.param else ""
+                )
+            ) from error
+    finally:
+        patched.undo()
+        torch.cuda.empty_cache()
+
+
+def _check(shadows, name):
+    shadow = shadows[name]
+    assert shadow.calls > 0, f"{name} was not reached by the pipeline"
+    if shadow.errors:
+        raise shadow.errors[0]
+
+
+@pytest.mark.parametrize("kernel_name", _SSD_KERNELS)
+def test_ssd_kernel_against_triton(shadows, kernel_name):
+    _check(shadows, kernel_name)
+
+
+def test_layer_norm_bwd_against_triton(shadows):
+    _check(shadows, _LAYER_NORM_BWD)
+
+
+def test_kernel_namespaces_match():
+    """The two dispatch namespaces and _SSD_KERNELS are maintained by hand in
+    three places; a kernel added to one and not the others would otherwise just
+    go untested, with nothing failing."""
+    triton_names = vars(ssd_combined._TRITON_SSD_KERNELS).keys()
+    assert vars(get_helion_ssd_kernels()).keys() == triton_names
+    assert set(_SSD_KERNELS) == triton_names