Initial kernel bundle upload: selective-update

.gitattributes

@@ -1,35 +1,3 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
+*.so filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,2 @@
+build
+*pyc

CARD.md

@@ -1,26 +1,61 @@
 ---
 library_name: kernels
-license: apache-2.0
----
+{% if license %}license: {{ license }}
+{% endif %}---
 
-This is the repository card of Ashiedu/selective-update that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated.
+This is the repository card of {{ repo_id }} that has been pushed on the Hub. It was built to be used with the [`kernels` library](https://github.com/huggingface/kernels). This card was automatically generated.
 
 ## How to use
+{% if functions %}
 
 ```python
 # make sure `kernels` is installed: `pip install -U kernels`
 from kernels import get_kernel
 
-kernel_module = get_kernel("Ashiedu/selective-update", version=1)
-selective_update = kernel_module.selective_update
+kernel_module = get_kernel("{{ repo_id }}", version={{ version }})
+{{ functions[0] }} = kernel_module.{{ functions[0] }}
 
-selective_update(...)
+{{ functions[0] }}(...)
 ```
+{% else %}
+
+Usage example not available.
+{% endif %}
 
 ## Available functions
-- `selective_update`
-- `selective_state_update`
+{% if functions %}
+{% for func in functions %}
+- `{{ func }}`
+{% endfor %}
+{% else %}
+
+Function list not available.
+{% endif %}
+{% if layers %}
+
+## Available layers
+{% for layer in layers %}
+- `{{ layer }}`
+{% endfor %}
+{% endif %}
 
 ## Benchmarks
+{% if has_benchmark %}
+
+Benchmarking script is available for this kernel. Run `kernels benchmark {{ repo_id }} --version {{ version }}`.
+{% else %}
+
+No benchmark available yet.
+{% endif %}
+{% if upstream %}
+
+## Upstream
+
+The original source code for this kernel comes from {{ upstream }}.
+{% endif %}
+{% if source %}
+
+## Source
 
-Benchmarking script is available for this kernel. Run `kernels benchmark Ashiedu/selective-update --version 1`.
+The kernel-builder formatted source for this kernel is available at {{ source }}.
+{% endif %}

build.toml

@@ -0,0 +1,13 @@
+[general]
+name = "selective-update"
+license = "Apache-2.0"
+backends = [
+  "rocm",
+]
+version = 1
+edition = 5
+
+[general.hub]
+repo-id = "Ashiedu/selective-update"
+
+[torch-noarch]

example.py

@@ -0,0 +1,42 @@
+# /// script
+# requires-python = ">=3.13"
+# dependencies = [
+#     "kernels",
+#     "numpy",
+#     "torch",
+# ]
+# ///
+
+import platform
+from pathlib import Path
+
+import kernels
+import torch
+
+# Load the locally built kernel
+kernel = kernels.get_local_kernel(Path("build"), "selective_update")
+
+# Select device
+if platform.system() == "Darwin":
+    device = torch.device("mps")
+elif hasattr(torch, "xpu") and torch.xpu.is_available():
+    device = torch.device("xpu")
+elif torch.version.cuda is not None and torch.cuda.is_available():
+    device = torch.device("cuda")
+else:
+    device = torch.device("cpu")
+
+print(f"Using device: {device}")
+
+# Create input tensor
+x = torch.tensor([1.0, 2.0, 3.0], device=device)
+print(f"Input:  {x}")
+
+# Run kernel (adds 1 to each element)
+result = kernel.selective_update(x)
+print(f"Output: {result}")
+
+# Verify result
+expected = x + 1.0
+assert torch.allclose(result, expected), "Kernel output doesn't match expected!"
+print("Success!")

flake.lock

@@ -0,0 +1,117 @@
+{
+  "nodes": {
+    "flake-compat": {
+      "locked": {
+        "lastModified": 1767039857,
+        "narHash": "sha256-vNpUSpF5Nuw8xvDLj2KCwwksIbjua2LZCqhV1LNRDns=",
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "rev": "5edf11c44bc78a0d334f6334cdaf7d60d732daab",
+        "type": "github"
+      },
+      "original": {
+        "owner": "edolstra",
+        "repo": "flake-compat",
+        "type": "github"
+      }
+    },
+    "flake-utils": {
+      "inputs": {
+        "systems": "systems"
+      },
+      "locked": {
+        "lastModified": 1731533236,
+        "narHash": "sha256-l0KFg5HjrsfsO/JpG+r7fRrqm12kzFHyUHqHCVpMMbI=",
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "rev": "11707dc2f618dd54ca8739b309ec4fc024de578b",
+        "type": "github"
+      },
+      "original": {
+        "owner": "numtide",
+        "repo": "flake-utils",
+        "type": "github"
+      }
+    },
+    "kernel-builder": {
+      "inputs": {
+        "flake-compat": "flake-compat",
+        "flake-utils": "flake-utils",
+        "nixpkgs": "nixpkgs",
+        "rust-overlay": "rust-overlay"
+      },
+      "locked": {
+        "lastModified": 1783738448,
+        "narHash": "sha256-vNb0C3LOAPlDPJH2OHbrVND5KrfvNXcl66nQaVp8rRU=",
+        "owner": "huggingface",
+        "repo": "kernels",
+        "rev": "570dcf4a810e6bd9d2433716c2bc6962c6ae46bb",
+        "type": "github"
+      },
+      "original": {
+        "owner": "huggingface",
+        "repo": "kernels",
+        "type": "github"
+      }
+    },
+    "nixpkgs": {
+      "locked": {
+        "lastModified": 1783284758,
+        "narHash": "sha256-tiQ8/qi8I45OOaBBYlVbXoAVkeQzvvTQOv5I45rMw5o=",
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d",
+        "type": "github"
+      },
+      "original": {
+        "owner": "NixOS",
+        "repo": "nixpkgs",
+        "rev": "ec1a11210589d294f0ac99d3290a27e6c73dfa1d",
+        "type": "github"
+      }
+    },
+    "root": {
+      "inputs": {
+        "kernel-builder": "kernel-builder"
+      }
+    },
+    "rust-overlay": {
+      "inputs": {
+        "nixpkgs": [
+          "kernel-builder",
+          "nixpkgs"
+        ]
+      },
+      "locked": {
+        "lastModified": 1783320166,
+        "narHash": "sha256-l7C/OsjcnWDOk2K3ssj+SBduwL67LashjBqis9+t468=",
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "rev": "20ee15370c9256669d66968b89ee20a4b0a4e673",
+        "type": "github"
+      },
+      "original": {
+        "owner": "oxalica",
+        "repo": "rust-overlay",
+        "type": "github"
+      }
+    },
+    "systems": {
+      "locked": {
+        "lastModified": 1681028828,
+        "narHash": "sha256-Vy1rq5AaRuLzOxct8nz4T6wlgyUR7zLU309k9mBC768=",
+        "owner": "nix-systems",
+        "repo": "default",
+        "rev": "da67096a3b9bf56a91d16901293e51ba5b49a27e",
+        "type": "github"
+      },
+      "original": {
+        "owner": "nix-systems",
+        "repo": "default",
+        "type": "github"
+      }
+    }
+  },
+  "root": "root",
+  "version": 7
+}

flake.nix

@@ -0,0 +1,11 @@
+{
+  inputs = {
+    kernel-builder.url = "github:huggingface/kernels";
+  };
+  outputs =
+    { self, kernel-builder, ... }:
+    kernel-builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

selective_update_cuda/selective_update.cu

@@ -0,0 +1,33 @@
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/all.h>
+
+__global__ void selective_update_kernel(float *__restrict__ out,
+                            float const *__restrict__ input, const int n) {
+  int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < n) {
+    out[idx] = input[idx] + 1.0f;
+  }
+}
+
+void selective_update(torch::Tensor &out, torch::Tensor const &input) {
+  TORCH_CHECK(input.device().is_cuda(), "input must be a CUDA tensor");
+  TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
+  TORCH_CHECK(input.scalar_type() == at::ScalarType::Float,
+              "selective_update only supports float32");
+  TORCH_CHECK(input.sizes() == out.sizes(),
+              "Tensors must have the same shape");
+  TORCH_CHECK(input.scalar_type() == out.scalar_type(),
+              "Tensors must have the same dtype");
+  TORCH_CHECK(input.device() == out.device(),
+              "Tensors must be on the same device");
+
+  int n = input.numel();
+  int threads = 256;
+  int blocks = (n + threads - 1) / threads;
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  selective_update_kernel<<<blocks, threads, 0, stream>>>(
+      out.data_ptr<float>(), input.data_ptr<float>(), n);
+}

tests/test_selective_update.py

@@ -0,0 +1,21 @@
+import platform
+
+import torch
+
+import selective_update
+
+
+def test_selective_update():
+    if platform.system() == "Darwin":
+        device = torch.device("mps")
+    elif hasattr(torch, "xpu") and torch.xpu.is_available():
+        device = torch.device("xpu")
+    elif torch.version.cuda is not None and torch.cuda.is_available():
+        device = torch.device("cuda")
+    else:
+        device = torch.device("cpu")
+
+    x = torch.randn(1024, 1024, dtype=torch.float32, device=device)
+    expected = x + 1.0
+    result = selective_update.selective_update(x)
+    torch.testing.assert_close(result, expected)

torch-ext/selective_update/__init__.py

@@ -0,0 +1,16 @@
+from typing import Optional
+
+import torch
+
+from .functional import selective_state_update
+
+
+def selective_update(x: torch.Tensor, out: Optional[torch.Tensor] = None) -> torch.Tensor:
+    """Legacy scaffold kernel — adds 1.0 to each element."""
+    if out is None:
+        out = torch.empty_like(x)
+    out.copy_(x + 1.0)
+    return out
+
+
+__all__ = ["selective_update", "selective_state_update"]

torch-ext/selective_update/_autotune_configs.py

@@ -0,0 +1,101 @@
+import os
+import warnings
+from packaging import version
+
+import torch
+
+try:
+    import triton
+    TRITON_VERSION = version.parse(triton.__version__)
+except ImportError:
+    TRITON_VERSION = version.parse("0.0.0")
+
+TRITON_HAS_CACHE_RESULTS = TRITON_VERSION >= version.parse("3.4.0")
+_autotune_warning_issued = False
+_deterministic_override = None
+
+
+def use_deterministic_mode():
+    if _deterministic_override is not None:
+        return _deterministic_override
+    env = os.environ.get("MAMBA_DETERMINISTIC")
+    if env:
+        return env[0] == "1"
+    return torch.are_deterministic_algorithms_enabled()
+
+
+def set_deterministic_mode(value):
+    global _deterministic_override
+    _deterministic_override = value
+
+
+def is_gfx1030():
+    """Identifies the HIP gfx1030 hardware."""
+    if not torch.cuda.is_available():
+        return False
+    try:
+        device_cap = torch.cuda.get_device_capability()
+        return device_cap == (10, 3)
+    except Exception:
+        return False
+
+
+def gfx1030_selective_update_configs(configs):
+    """Pin safe tile configs for selective-update on gfx1030."""
+    if not is_gfx1030():
+        return configs
+
+    expected = {
+        "BLOCK_SIZE_M": 32,
+        "BLOCK_SIZE_DSTATE": 64,
+    }
+    selected = [
+        config
+        for config in configs
+        if all(config.kwargs.get(key) == value for key, value in expected.items())
+    ]
+    if not selected:
+        return [configs[0]]
+    return selected[:1]
+
+
+def _estimate_config_cost(cfg):
+    block_product = 1
+    for key, val in cfg.kwargs.items():
+        if key.startswith("BLOCK_SIZE_"):
+            block_product *= val
+    return block_product * (getattr(cfg, "num_stages", 1) or 1)
+
+
+def _filter_configs_by_block_sizes(configs):
+    env_filters = {}
+    for suffix in ("M", "N", "K", "DSTATE"):
+        env_val = os.environ.get(f"TRITON_AUTOTUNE_BLOCK_SIZE_{suffix}")
+        if env_val is not None:
+            env_filters[f"BLOCK_SIZE_{suffix}"] = int(env_val)
+    if not env_filters:
+        return None
+    matching = configs
+    for key, target in env_filters.items():
+        matching = [c for c in matching if c.kwargs.get(key) == target]
+    return matching[:1] if matching else None
+
+
+def autotune_configs(configs):
+    if not configs or not use_deterministic_mode():
+        return configs
+    if TRITON_HAS_CACHE_RESULTS and os.environ.get("TRITON_CACHE_AUTOTUNING") == "1":
+        return configs
+    global _autotune_warning_issued
+    if not _autotune_warning_issued:
+        _autotune_warning_issued = True
+        msg = (
+            "Deterministic mode: set TRITON_CACHE_AUTOTUNING=1 for cached autotuning."
+            if TRITON_HAS_CACHE_RESULTS
+            else "Deterministic mode: upgrade to Triton >= 3.4.0 for cached autotuning."
+        )
+        warnings.warn(msg)
+    filtered = _filter_configs_by_block_sizes(configs)
+    if filtered:
+        return filtered
+    return [min(configs, key=_estimate_config_cost)]

torch-ext/selective_update/functional.py

@@ -0,0 +1,165 @@
+"""Python autograd wrapper for selective-update Triton kernel.
+
+Matches upstream ``mamba_ssm.ops.triton.selective_state_update`` ABI exactly:
+    state: (B, D, N) or (B, H, D, N)
+    x: (B, D) or (B, H, D)
+    dt: (B, D) or (B, H, D)
+    A: (D, N) or (H, D, N)
+    B: (B, N) or (B, G, N)
+    C: (B, N) or (B, G, N)
+    D: (D,) or (H, D) or None
+    z: (B, D) or (B, H, D) or None
+    dt_bias: (D,) or (H, D) or None
+    state_batch_indices: (B,) or None
+"""
+
+from __future__ import annotations
+
+import importlib.util
+from pathlib import Path
+
+import torch
+import triton
+
+_TRITON_KERNELS_PATH = (
+    Path(__file__).resolve().parent / "triton_kernels.py"
+)
+
+
+def _load_triton_kernel():
+    spec = importlib.util.spec_from_file_location(
+        "selective_update.triton_kernels",
+        _TRITON_KERNELS_PATH,
+    )
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.selective_state_update_kernel
+
+
+selective_state_update_kernel = None  # populated on first call
+
+
+def _get_kernel():
+    global selective_state_update_kernel
+    if selective_state_update_kernel is None:
+        selective_state_update_kernel = _load_triton_kernel()
+    return selective_state_update_kernel
+
+
+def _resolve_tie_hdim(A, dt, dt_bias):
+    if A.dim() == 2:
+        a0 = A[0]
+        return a0.stride(0) == 0 and a0.stride(1) == 0 and dt.stride(-1) == 0 and (dt_bias.stride(-1) == 0 if dt_bias is not None else True)
+    return A.stride(-1) == 0 and A.stride(-2) == 0 and dt.stride(-1) == 0 and (dt_bias.stride(-1) == 0 if dt_bias is not None else True)
+
+
+class SelectiveUpdateFunction(torch.autograd.Function):
+    """Differentiable single-token selective state update.
+
+    Mirrors upstream ``selective_state_update`` shape contract exactly.
+    Backward raises NotImplementedError — training uses the patched
+    selective_scan reverse path instead.
+    """
+
+    @staticmethod
+    def forward(ctx, state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False, state_batch_indices=None):
+        orig_state = state
+        has_heads = state.dim() > 3
+        if state.dim() == 3:
+            state = state.unsqueeze(1)
+        if x.dim() == 2:
+            x = x.unsqueeze(1)
+        if dt.dim() == 2:
+            dt = dt.unsqueeze(1)
+        if A.dim() == 2:
+            A = A.unsqueeze(0)
+        if B.dim() == 2:
+            B = B.unsqueeze(1)
+        if C.dim() == 2:
+            C = C.unsqueeze(1)
+        if D is not None and D.dim() == 1:
+            D = D.unsqueeze(0)
+        if z is not None and z.dim() == 2:
+            z = z.unsqueeze(1)
+        if dt_bias is not None and dt_bias.dim() == 1:
+            dt_bias = dt_bias.unsqueeze(0)
+
+        B_sz, H, D_sz, N_sz = state.shape
+        x_sz = x.shape[0]
+        if x.shape != (x_sz, H, D_sz):
+            raise ValueError(f"x shape {x.shape} does not match state batch/heads/dim")
+        if dt.shape != (x_sz, H, D_sz):
+            raise ValueError(f"dt shape {dt.shape} does not match x shape")
+        if A.shape != (H, D_sz, N_sz):
+            raise ValueError(f"A shape {A.shape} does not match (H, D, N)")
+        ngroups = B.shape[1]
+        if H % ngroups != 0:
+            raise ValueError(f"nheads {H} must be divisible by ngroups {ngroups}")
+        if B.shape[0] != x_sz or B.shape[2] != N_sz or C.shape[0] != x_sz or C.shape[2] != N_sz:
+            raise ValueError(f"B/C batch or dstate mismatch")
+        if D is not None and D.shape not in {(H, D_sz), (H,)}:
+            raise ValueError(f"D shape {D.shape} does not match (H, D) or (H,)")
+        if z is not None and z.shape != (x_sz, H, D_sz):
+            raise ValueError(f"z shape {z.shape} does not match x shape")
+        if dt_bias is not None and dt_bias.shape not in {(H, D_sz), (H,)}:
+            raise ValueError(f"dt_bias shape {dt_bias.shape} does not match (H, D) or (H,)")
+        if state_batch_indices is not None and state_batch_indices.shape != (x_sz,):
+            raise ValueError(f"state_batch_indices shape {state_batch_indices.shape} does not match (B,)")
+
+        tie_hdim = _resolve_tie_hdim(A, dt, dt_bias)
+        nheads_ratio = H // ngroups
+
+        out = torch.empty_like(x)
+        grid = lambda META: (triton.cdiv(D_sz, META["BLOCK_SIZE_M"]), x_sz, H)
+
+        BLOCK_M = 32 if N_sz <= 16 else (16 if N_sz <= 32 else (8 if N_sz <= 64 else (4 if N_sz <= 128 else 4)))
+        num_warps = 4 if N_sz <= 64 else 8
+
+        z_strides = ((z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0))
+
+        _get_kernel()[grid](
+            state, x, dt, dt_bias, A, B, C, D, z, out, state_batch_indices,
+            x_sz, H, D_sz, N_sz, nheads_ratio,
+            state.stride(0), state.stride(1), state.stride(2), state.stride(3),
+            x.stride(0), x.stride(1), x.stride(2),
+            dt.stride(0), dt.stride(1), dt.stride(2),
+            dt_bias.stride(0) if dt_bias is not None else 0, dt_bias.stride(1) if dt_bias is not None else 0,
+            A.stride(0), A.stride(1), A.stride(2),
+            B.stride(0), B.stride(1), B.stride(2),
+            C.stride(0), C.stride(1), C.stride(2),
+            D.stride(0) if D is not None else 0, D.stride(1) if D is not None else 0,
+            z_strides[0], z_strides[1], z_strides[2],
+            out.stride(0), out.stride(1), out.stride(2),
+            dt_softplus,
+            tie_hdim,
+            BLOCK_M,
+            num_warps=num_warps,
+            num_stages=2,
+        )
+
+        if not has_heads:
+            out = out.squeeze(1)
+            state = state.squeeze(1)
+        ctx.save_for_backward(orig_state, x, dt, A, B, C, D, z, dt_bias, state, state_batch_indices)
+        ctx.tie_hdim = tie_hdim
+        ctx.dt_softplus = dt_softplus
+        ctx.has_heads = has_heads
+        ctx.nheads_ratio = nheads_ratio
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_out):
+        raise NotImplementedError(
+            "SelectiveUpdateFunction backward is not implemented. "
+            "Use autograd through the patched selective_scan path for training."
+        )
+
+
+def selective_state_update(state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False, state_batch_indices=None, **kwargs):
+    """Dispatches single-token state updates to the custom Triton backend.
+
+    Shapes mirror upstream ``mamba_ssm.ops.triton.selective_state_update`` exactly.
+    """
+    return SelectiveUpdateFunction.apply(
+        state, x, dt, A, B, C, D, z, dt_bias, dt_softplus, state_batch_indices
+    )

torch-ext/selective_update/triton_kernels.py

@@ -0,0 +1,127 @@
+"""Triton kernels for Mamba-2 single-token selective state updates.
+
+Matches upstream ``mamba_ssm.ops.triton.selective_state_update`` ABI exactly:
+    state: (B, H, D, N)  x/dt: (B, H, D)  A: (H, D, N)  B/C: (B, G, N)
+"""
+
+import triton
+import triton.language as tl
+
+
+@triton.jit
+def _softplus(dt):
+    return tl.math.log(tl.math.exp(dt) + 1)
+
+
+@triton.heuristics({"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None})
+@triton.heuristics({"HAS_D": lambda args: args["D_ptr"] is not None})
+@triton.heuristics({"HAS_Z": lambda args: args["z_ptr"] is not None})
+@triton.heuristics({"HAS_STATE_BATCH_INDICES": lambda args: args["state_batch_indices_ptr"] is not None})
+@triton.heuristics({"BLOCK_SIZE_DSTATE": lambda args: triton.next_power_of_2(args["dstate"])})
+@triton.jit
+def selective_state_update_kernel(
+    state_ptr, x_ptr, dt_ptr, dt_bias_ptr, A_ptr, B_ptr, C_ptr, D_ptr, z_ptr, out_ptr, state_batch_indices_ptr,
+    batch, nheads, dim, dstate, nheads_ngroups_ratio,
+    s_state_batch, s_state_head, s_state_dim, s_state_dstate,
+    s_x_batch, s_x_head, s_x_dim,
+    s_dt_batch, s_dt_head, s_dt_dim,
+    s_dt_bias_head, s_dt_bias_dim,
+    s_A_head, s_A_dim, s_A_dstate,
+    s_B_batch, s_B_group, s_B_dstate,
+    s_C_batch, s_C_group, s_C_dstate,
+    s_D_head, s_D_dim,
+    s_z_batch, s_z_head, s_z_dim,
+    s_out_batch, s_out_head, s_out_dim,
+    DT_SOFTPLUS: tl.constexpr,
+    TIE_HDIM: tl.constexpr,
+    BLOCK_SIZE_M: tl.constexpr,
+    HAS_DT_BIAS: tl.constexpr,
+    HAS_D: tl.constexpr,
+    HAS_Z: tl.constexpr,
+    HAS_STATE_BATCH_INDICES: tl.constexpr,
+    BLOCK_SIZE_DSTATE: tl.constexpr,
+):
+    pid_m = tl.program_id(axis=0)
+    pid_b = tl.program_id(axis=1)
+    pid_h = tl.program_id(axis=2)
+
+    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)
+    out_ptr += pid_b * s_out_batch + pid_h * s_out_head
+    out_ptrs = out_ptr + offs_m * s_out_dim
+
+    if HAS_STATE_BATCH_INDICES:
+        state_batch_indices_ptr += pid_b
+        state_batch_idx = tl.load(state_batch_indices_ptr)
+        if state_batch_idx < 0:
+            tl.store(out_ptrs, 0.0, mask=offs_m < dim)
+            return
+        state_ptr += state_batch_idx * s_state_batch + pid_h * s_state_head
+    else:
+        state_ptr += pid_b * s_state_batch + pid_h * s_state_head
+
+    x_ptr += pid_b * s_x_batch + pid_h * s_x_head
+    dt_ptr += pid_b * s_dt_batch + pid_h * s_dt_head
+    if HAS_DT_BIAS:
+        dt_bias_ptr += pid_h * s_dt_bias_head
+    A_ptr += pid_h * s_A_head
+    B_ptr += pid_b * s_B_batch + (pid_h // nheads_ngroups_ratio) * s_B_group
+    C_ptr += pid_b * s_C_batch + (pid_h // nheads_ngroups_ratio) * s_C_group
+    if HAS_Z:
+        z_ptr += pid_b * s_z_batch + pid_h * s_z_head
+
+    offs_n = tl.arange(0, BLOCK_SIZE_DSTATE)
+    state_ptrs = state_ptr + (offs_m[:, None] * s_state_dim + offs_n[None, :] * s_state_dstate)
+    x_ptrs = x_ptr + offs_m * s_x_dim
+    dt_ptrs = dt_ptr + offs_m * s_dt_dim
+    if HAS_DT_BIAS:
+        dt_bias_ptrs = dt_bias_ptr + offs_m * s_dt_bias_dim
+    if HAS_D:
+        D_ptr += pid_h * s_D_head
+    A_ptrs = A_ptr + (offs_m[:, None] * s_A_dim + offs_n[None, :] * s_A_dstate)
+    B_ptrs = B_ptr + offs_n * s_B_dstate
+    C_ptrs = C_ptr + offs_n * s_C_dstate
+    if HAS_D:
+        D_ptrs = D_ptr + offs_m * s_D_dim
+    if HAS_Z:
+        z_ptrs = z_ptr + offs_m * s_z_dim
+
+    state = tl.load(state_ptrs, mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate), other=0.0)
+    x = tl.load(x_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
+    if not TIE_HDIM:
+        dt = tl.load(dt_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
+        if HAS_DT_BIAS:
+            dt += tl.load(dt_bias_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
+        if DT_SOFTPLUS:
+            dt = tl.where(dt <= 20.0, _softplus(dt), dt)
+        A = tl.load(A_ptrs, mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate), other=0.0).to(tl.float32)
+        dA = tl.exp(A * dt[:, None])
+    else:
+        dt = tl.load(dt_ptr).to(tl.float32)
+        if HAS_DT_BIAS:
+            dt += tl.load(dt_bias_ptr).to(tl.float32)
+        if DT_SOFTPLUS:
+            dt = tl.where(dt <= 20.0, _softplus(dt), dt)
+        A = tl.load(A_ptr).to(tl.float32)
+        dA = tl.exp(A * dt)
+
+    B = tl.load(B_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32)
+    C = tl.load(C_ptrs, mask=offs_n < dstate, other=0.0).to(tl.float32)
+    if HAS_D:
+        D = tl.load(D_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
+    if HAS_Z:
+        z = tl.load(z_ptrs, mask=offs_m < dim, other=0.0).to(tl.float32)
+
+    if not TIE_HDIM:
+        dB = B[None, :] * dt[:, None]
+    else:
+        dB = B * dt
+
+    state = state * dA + dB * x[:, None]
+    tl.store(state_ptrs, state, mask=(offs_m[:, None] < dim) & (offs_n[None, :] < dstate))
+    out = tl.sum(state * C[None, :], axis=1)
+    if HAS_D:
+        out += x * D
+    if HAS_Z:
+        out *= z * tl.sigmoid(z)
+
+    tl.store(out_ptrs, out, mask=offs_m < dim)

torch-ext/torch_binding.cpp

@@ -0,0 +1,19 @@
+#include <torch/library.h>
+
+#include "registration.h"
+#include "torch_binding.h"
+
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("selective_update(Tensor! out, Tensor input) -> ()");
+#if defined(CPU_KERNEL)
+  ops.impl("selective_update", torch::kCPU, &selective_update);
+#elif defined(CUDA_KERNEL) || defined(ROCM_KERNEL)
+  ops.impl("selective_update", torch::kCUDA, &selective_update);
+#elif defined(METAL_KERNEL)
+  ops.impl("selective_update", torch::kMPS, selective_update);
+#elif defined(XPU_KERNEL)
+  ops.impl("selective_update", torch::kXPU, &selective_update);
+#endif
+}
+
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h

@@ -0,0 +1,5 @@
+#pragma once
+
+#include <torch/torch.h>
+
+void selective_update(torch::Tensor &out, torch::Tensor const &input);