From bcbd2aabce537560e64e027537b218b922097fa6 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 22 Dec 2025 10:55:48 -0800 Subject: [PATCH 01/38] I seem that I can call my cuda wrapper from python --- csrc/flashinfer_mamba_binding.cu | 35 ++ csrc/selective_state_update.cu | 69 ++++ flashinfer/__init__.py | 1 + flashinfer/jit/mamba/__init__.py | 19 ++ flashinfer/mamba/__init__.py | 19 ++ flashinfer/mamba/selective_state_update.py | 134 ++++++++ .../mamba/selective_state_update.cuh | 53 +++ tests/mamba/test_selective_state_update.py | 304 ++++++++++++++++++ 8 files changed, 634 insertions(+) create mode 100644 csrc/flashinfer_mamba_binding.cu create mode 100644 csrc/selective_state_update.cu create mode 100644 flashinfer/jit/mamba/__init__.py create mode 100644 flashinfer/mamba/__init__.py create mode 100644 flashinfer/mamba/selective_state_update.py create mode 100644 include/flashinfer/mamba/selective_state_update.cuh create mode 100644 tests/mamba/test_selective_state_update.py diff --git a/csrc/flashinfer_mamba_binding.cu b/csrc/flashinfer_mamba_binding.cu new file mode 100644 index 0000000000..02198ce25c --- /dev/null +++ b/csrc/flashinfer_mamba_binding.cu @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "tvm_ffi_utils.h" + +// Declare the function(s) implemented in selective_state_update.cu +using tvm::ffi::Optional; + +namespace flashinfer::mamba { + +void selective_state_update(TensorView state, TensorView x, TensorView dt, TensorView output, + TensorView A, TensorView B, TensorView C, TensorView D, + Optional z, Optional dt_bias, + bool dt_softplus, Optional state_batch_indices, + int64_t pad_slot_id); + +} // namespace flashinfer::mamba + +// Export the function(s) via TVM-FFI +// This enables cross-language bindings (not just PyTorch) +TVM_FFI_DLL_EXPORT_TYPED_FUNC(selective_state_update, flashinfer::mamba::selective_state_update); + +// Add more mamba operations here as they are implemented \ No newline at end of file diff --git a/csrc/selective_state_update.cu b/csrc/selective_state_update.cu new file mode 100644 index 0000000000..d8b4958417 --- /dev/null +++ b/csrc/selective_state_update.cu @@ -0,0 +1,69 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include + +#include "tvm_ffi_utils.h" + +using namespace flashinfer; +using tvm::ffi::Optional; + +namespace flashinfer::mamba { + +// TODO: Implement the launcher function that: +// 1. Validates tensor dimensions and devices +// 2. Extracts raw pointers from TensorView +// 3. Calls the framework-agnostic kernel from include/flashinfer/mamba/selective_state_update.cuh +// +// Example signature (adjust based on your kernel requirements): +void selective_state_update(TensorView state, TensorView x, TensorView dt, TensorView output, + TensorView A, TensorView B, TensorView C, TensorView D, + Optional z, Optional dt_bias, + bool dt_softplus, Optional state_batch_indices, + int64_t pad_slot_id) { + throw std::runtime_error("selective_state_update is not implemented yet."); + + // TODO: Add input validation + // CHECK_LAST_DIM_CONTIGUOUS_INPUT(state); + // CHECK_LAST_DIM_CONTIGUOUS_INPUT(x); + // CHECK_DEVICE(state, x); + // TVM_FFI_ICHECK_EQ(state.ndim(), 3); // Example dimension check + + // TODO: Extract dimensions from tensors + // unsigned int batch_size = state.size(0); + // unsigned int dim_state = state.size(1); + // unsigned int dim_input = x.size(1); + + // TODO: Get CUDA device and stream + // ffi::CUDADeviceGuard device_guard(state.device().device_id); + // const cudaStream_t stream = get_stream(state.device()); + + // TODO: Dispatch based on dtype and call kernel + // DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(state.dtype(), c_type, [&] { + // cudaError_t status = mamba::SelectiveStateUpdate( + // static_cast(state.data_ptr()), + // static_cast(x.data_ptr()), + // static_cast(dt.data_ptr()), + // static_cast(A.data_ptr()), + // static_cast(B.data_ptr()), + // static_cast(C.data_ptr()), + // batch_size, dim_state, dim_input, stream); + // TVM_FFI_ICHECK(status == cudaSuccess) + // << "SelectiveStateUpdate failed with error code " << cudaGetErrorString(status); + // return true; + // }); +} + +} // namespace flashinfer::mamba diff --git a/flashinfer/__init__.py b/flashinfer/__init__.py index 07a913eb5f..55f7e48c8d 100644 --- a/flashinfer/__init__.py +++ b/flashinfer/__init__.py @@ -153,3 +153,4 @@ from .utils import next_positive_power_of_2 as next_positive_power_of_2 from .xqa import xqa as xqa from .xqa import xqa_mla as xqa_mla +from . import mamba as mamba diff --git a/flashinfer/jit/mamba/__init__.py b/flashinfer/jit/mamba/__init__.py new file mode 100644 index 0000000000..e71b2ea2f5 --- /dev/null +++ b/flashinfer/jit/mamba/__init__.py @@ -0,0 +1,19 @@ +""" +Copyright (c) 2025 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from .selective_state_update import gen_selective_state_update_module + +__all__ = ["gen_selective_state_update_module"] \ No newline at end of file diff --git a/flashinfer/mamba/__init__.py b/flashinfer/mamba/__init__.py new file mode 100644 index 0000000000..a201b0ff6d --- /dev/null +++ b/flashinfer/mamba/__init__.py @@ -0,0 +1,19 @@ +""" +Copyright (c) 2025 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from .selective_state_update import selective_state_update + +__all__ = ["selective_state_update"] \ No newline at end of file diff --git a/flashinfer/mamba/selective_state_update.py b/flashinfer/mamba/selective_state_update.py new file mode 100644 index 0000000000..4dbdd85b95 --- /dev/null +++ b/flashinfer/mamba/selective_state_update.py @@ -0,0 +1,134 @@ +""" +Copyright (c) 2025 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +import functools +from typing import Optional + +import torch + +from ..api_logging import flashinfer_api +from ..jit.mamba import gen_selective_state_update_module +from ..utils import register_custom_op, register_fake_op + + +@functools.cache +def get_selective_state_update_module(): + """Get cached JIT-compiled selective_state_update module. + + This uses @functools.cache for in-memory caching of loaded modules. + The JIT system also caches compiled .so files on disk in ~/.cache/flashinfer/ + """ + return gen_selective_state_update_module().build_and_load() + + +@flashinfer_api +def selective_state_update( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + z: Optional[torch.Tensor] = None, + dt_bias: Optional[torch.Tensor] = None, + dt_softplus: bool = False, + state_batch_indices: Optional[torch.Tensor] = None, + pad_slot_id: int = -1, +) -> torch.Tensor: + r"""Selective state update operation for Mamba layers. + + Parameters + ---------- + state : torch.Tensor + State tensor with shape (batch, dim, dstate) or (batch, nheads, dim, dstate) + x : torch.Tensor + Input tensor with shape (batch, dim) or (batch, nheads, dim) + dt : torch.Tensor + Delta time tensor with shape (batch, dim) or (batch, nheads, dim) + A : torch.Tensor + A matrix with shape (dim, dstate) or (nheads, dim, dstate) + B : torch.Tensor + B matrix with shape (batch, dstate) or (batch, ngroups, dstate) + C : torch.Tensor + C matrix with shape (batch, dstate) or (batch, ngroups, dstate) + D : torch.Tensor + D vector with shape (dim,) or (nheads, dim) + z : Optional[torch.Tensor] + Optional z tensor with shape (batch, dim) or (batch, nheads, dim) + dt_bias : Optional[torch.Tensor] + Optional dt bias with shape (dim,) or (nheads, dim) + dt_softplus : bool + Whether to apply softplus to dt + state_batch_indices : Optional[torch.Tensor] + Optional batch indices for cache processing + pad_slot_id : int + If state_batch_indices is passed, lets the kernel identify padded entries + that will not be processed. For example: state_batch_indices = [pad_slot_id, 1, 20, pad_slot_id] + in this case, the kernel will not process entries at indices 0 and 3 + + Returns + ------- + output : torch.Tensor + Output tensor with shape (batch, dim) or (batch, nheads, dim) + """ + output = torch.empty_like(x) + _selective_state_update( + state, x, dt, output, A, B, C, D, z, dt_bias, dt_softplus, state_batch_indices, pad_slot_id + ) + return output + + +@register_custom_op("flashinfer::selective_state_update", mutates_args=("state", "output")) +def _selective_state_update( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + output: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + z: Optional[torch.Tensor], + dt_bias: Optional[torch.Tensor], + dt_softplus: bool, + state_batch_indices: Optional[torch.Tensor], + pad_slot_id: int, +) -> None: + """Internal function registered with torch.library for torch.compile() support.""" + get_selective_state_update_module().selective_state_update( + state, x, dt, output, A, B, C, D, z, dt_bias, dt_softplus, state_batch_indices, pad_slot_id + ) + + +@register_fake_op("flashinfer::selective_state_update") +def _selective_state_update_fake( + state: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + output: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + z: Optional[torch.Tensor], + dt_bias: Optional[torch.Tensor], + dt_softplus: bool, + state_batch_indices: Optional[torch.Tensor], + pad_slot_id: int, +) -> None: + """Fake implementation for torch.compile() meta tensor propagation.""" + pass diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh new file mode 100644 index 0000000000..c1a272f6d1 --- /dev/null +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -0,0 +1,53 @@ +/* + * Copyright (c) 2025 by FlashInfer team. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#ifndef FLASHINFER_MAMBA_SELECTIVE_STATE_UPDATE_CUH_ +#define FLASHINFER_MAMBA_SELECTIVE_STATE_UPDATE_CUH_ + +#include "../utils.cuh" +#include "../vec_dtypes.cuh" + +namespace flashinfer { + +namespace mamba { + +// TODO: Add your kernel implementation here +// This should be framework-agnostic and use raw pointers only (no PyTorch/TensorFlow types) +// +// Example: +// template +// __global__ void SelectiveStateUpdateKernel( +// T* state, const T* x, const T* dt, const T* A, const T* B, const T* C, +// int batch_size, int dim_state, int dim_input) { +// // Your kernel implementation +// } + +// TODO: Add host function that launches the kernel +// Example: +// template +// cudaError_t SelectiveStateUpdate( +// T* state, const T* x, const T* dt, const T* A, const T* B, const T* C, +// int batch_size, int dim_state, int dim_input, +// cudaStream_t stream = 0) { +// // Compute grid/block dimensions +// // Launch kernel +// // Return cudaGetLastError() +// } + +} // namespace mamba + +} // namespace flashinfer + +#endif // FLASHINFER_MAMBA_SELECTIVE_STATE_UPDATE_CUH_ \ No newline at end of file diff --git a/tests/mamba/test_selective_state_update.py b/tests/mamba/test_selective_state_update.py new file mode 100644 index 0000000000..465d824215 --- /dev/null +++ b/tests/mamba/test_selective_state_update.py @@ -0,0 +1,304 @@ +import numpy as np +import pytest +import torch +import torch.nn.functional as F +from einops import rearrange, repeat + +import flashinfer + + +def selective_state_update_ref( + state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False +): + """ + Argument: + state: (batch, dstate, dim) or (batch, nheads, dstate, dim) + x: (batch, dim) or (batch, nheads, dim) + dt: (batch, dim) or (batch, nheads, dim) + A: (dstate, dim) or (nheads, dstate, dim) + B: (batch, dstate) or (batch, ngroups, dstate) + C: (batch, dstate) or (batch, ngroups, dstate) + D: (dim,) or (nheads, dim) + z: (batch, dim) or (batch, nheads, dim) + dt_bias: (dim,) or (nheads, dim) + Return: + out: (batch, dim) or (batch, nheads, dim) + """ + 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) + batch, nheads, dstate, dim = state.shape + + assert x.shape == (batch, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dstate, dim) + ngroups = B.shape[1] + assert nheads % ngroups == 0, "nheads must be divisible by ngroups" + assert B.shape == (batch, ngroups, dstate) + assert C.shape == B.shape + + if D is not None: + assert D.shape == (nheads, dim) + if z is not None: + assert z.shape == x.shape + if dt_bias is not None: + assert dt_bias.shape == (nheads, dim) + dt = dt + dt_bias + dt = F.softplus(dt) if dt_softplus else dt + dA = torch.exp( + rearrange(dt, "b h d -> b h 1 d") * A + ) # (batch, nheads, dstate, dim) + B = repeat(B, "b g n -> b (g h) n", h=nheads // ngroups) # (batch, nheads, dstate) + C = repeat(C, "b g n -> b (g h) n", h=nheads // ngroups) # (batch, nheads, dstate) + dB = rearrange(dt, "b h d -> b h 1 d") * rearrange( + B.float(), "b h n -> b h n 1" + ) # (batch, nheads, dstate, dim) + state_new = state.float() * dA + dB * rearrange( + x.float(), "b h d -> b h 1 d" + ) # (batch, nheads, dstate, dim) + state.copy_(state_new.to(state.dtype)) + out = torch.einsum("bhnd,bhn->bhd", state_new, C.float()) + if D is not None: + out += x.float() * D + out = (out if z is None else out * F.silu(z.float())).to(x.dtype) + if not has_heads: + out = out.squeeze(1) + return out + + +def create_test_inputs( + batch_size, nheads, dim, dstate, ngroups, input_dtype, weight_dtype, matrixA_dtype +): + # Set seed for reproducibility + torch.manual_seed(0) + + device = torch.device("cuda") + + # if we use the cache, then the state indices are taken from a specific slot + # so the state in the kernel will have batch as the first dimension, but it will + # only come from a particular slot; the full tensor first dim is larger + ssm_state_cache_size = max(384, int(2 * batch_size)) + + state_cache = torch.randn( + ssm_state_cache_size, nheads, dim, dstate, dtype=input_dtype, device=device + ) + + x = torch.randn(batch_size, nheads, dim, dtype=input_dtype, device=device) + + dt = torch.randn(batch_size, nheads, dtype=weight_dtype, device=device).as_strided( + (batch_size, nheads, dim), (nheads, 1, 0) + ) + + A = -torch.rand(nheads, dtype=torch.float32, device=device).as_strided( + (nheads, dim, dstate), (1, 0, 0) + ) + + # B and C - (batch_size, ngroups, dstate) + B = torch.randn(batch_size, ngroups, dstate, dtype=input_dtype, device=device) + C = torch.randn(batch_size, ngroups, dstate, dtype=input_dtype, device=device) + + # D - (nheads, dim) with strides (1, 0) - one value per head + D = torch.randn(nheads, dtype=weight_dtype, device=device).as_strided( + (nheads, dim), (1, 0) + ) + + dt_bias = torch.randn(nheads, dtype=weight_dtype, device=device).as_strided( + (nheads, dim), (1, 0) + ) + + # Slot indices for state batching - (batch_size,) + slot_idx = torch.randperm(ssm_state_cache_size, dtype=torch.int32, device=device)[ + :batch_size + ] + + return { + "state_cache": state_cache, + "x": x, + "dt": dt, + "A": A, + "B": B, + "C": C, + "D": D, + "dt_bias": dt_bias, + "slot_idx": slot_idx, + } + +# @pytest.mark.parametrize("batch", [1, 32, 64]) +# @pytest.mark.parametrize("nheads", [8, 64]) +# @pytest.mark.parametrize("dim", [64]) +# @pytest.mark.parametrize("dstate", [128]) +# @pytest.mark.parametrize("ngroups", [8]) +# @pytest.mark.parametrize("delta_softplus", [True]) +# @pytest.mark.parametrize("input_dtype", [torch.bfloat16]) +# @pytest.mark.parametrize("weight_dtype", [torch.bfloat16]) +# @pytest.mark.parametrize("matrixA_dtype", [torch.float32]) +@pytest.mark.parametrize("batch", [1]) +@pytest.mark.parametrize("nheads", [64]) +@pytest.mark.parametrize("dim", [64]) +@pytest.mark.parametrize("dstate", [128]) +@pytest.mark.parametrize("ngroups", [8]) +@pytest.mark.parametrize("delta_softplus", [True]) +@pytest.mark.parametrize("input_dtype", [torch.bfloat16]) +@pytest.mark.parametrize("weight_dtype", [torch.bfloat16]) +@pytest.mark.parametrize("matrixA_dtype", [torch.float32]) +def test_selective_state_update( + batch, + nheads, + dim, + dstate, + ngroups, + delta_softplus, + input_dtype, + weight_dtype, + matrixA_dtype, +): + """Test selective_state_update correctness against reference implementation.""" + # TODO: Create input tensors based on your kernel requirements + # Example: + inputs = create_test_inputs( + batch, + nheads, + dim, + dstate, + ngroups, + input_dtype, + weight_dtype, + matrixA_dtype, + ) + + state = inputs["state_cache"][inputs["slot_idx"]] + state_ref = rearrange(state, "... p n -> ... n p").detach().clone() + A = inputs["A"] + A_ref = (rearrange(A, "... p n -> ... n p") if A.ndim == 3 else A).detach().clone() + y_ref = selective_state_update_ref( + state_ref, + inputs["x"], + inputs["dt"], + A_ref, + inputs["B"], + inputs["C"], + D=inputs["D"], + z=None, + dt_bias=inputs["dt_bias"], + dt_softplus=delta_softplus, + ) + + y_test = flashinfer.mamba.selective_state_update( + state, + inputs["x"], + inputs["dt"], + A, + inputs["B"], + inputs["C"], + D=inputs["D"], + z=None, + dt_bias=inputs["dt_bias"], + dt_softplus=delta_softplus, + state_batch_indices=inputs["slot_idx"], + pad_slot_id=-1, + ) + + atol=1e-3 + rtol=1e-2 + outputs_match = torch.allclose(y_ref, y_test, atol=atol, rtol=rtol) + + if outputs_match: + print(f"✓ Outputs match within tolerance (atol={atol}, rtol={rtol})") + else: + print(f"✗ Outputs do NOT match within tolerance (atol={atol}, rtol={rtol})") + all_passed = False + + # Detailed comparison using numpy testing + y_ref_np = y_ref.detach().cpu().float().numpy() + y_test_np = y_test.detach().cpu().float().numpy() + print(f"dtypes: ref {y_ref_np.dtype}, test {y_test_np.dtype}") + + print("\nDetailed mismatch analysis:") + mismatch_mask = ~np.isclose(y_ref_np, y_test_np, atol=atol, rtol=rtol) + num_mismatches = np.sum(mismatch_mask) + total_elements = y_ref_np.size + + print( + f"Number of mismatched elements: {num_mismatches} / {total_elements} ({100 * num_mismatches / total_elements:.2f}%)" + ) + + mismatch_indices = np.argwhere(mismatch_mask) + print(f"First few mismatch locations (up to 10):") + for i, idx in enumerate(mismatch_indices[:10]): + idx_tuple = tuple(idx) + ref_val = y_ref_np[idx_tuple] + test_val = y_test_np[idx_tuple] + diff = abs(ref_val - test_val) + rel_diff = diff / (abs(ref_val) + 1e-8) + print( + f" Index {idx_tuple}: ref={ref_val:.6f}, test={test_val:.6f}, diff={diff:.6e}, rel_diff={rel_diff:.6e}" + ) + + assert(outputs_match) + + # Compare states (updated in-place) + print("\nComparing states with reference...") + + + state_test = state[inputs["slot_idx"]] + state_diff = torch.abs(state_ref - state_test) + max_state_diff = state_diff.max().item() + mean_state_diff = state_diff.mean().item() + + print(f"Max absolute state difference: {max_state_diff:.6e}") + print(f"Mean absolute state difference: {mean_state_diff:.6e}") + + # Check if states match within tolerance + states_match = torch.allclose( state_ref, state_test, atol=atol, rtol=rtol ) + + if states_match: + print(f"✓ States match within tolerance (atol={atol}, rtol={rtol})") + return True + else: + print(f"✗ States do NOT match within tolerance (atol={atol}, rtol={rtol})") + + # Detailed comparison using numpy testing + state_ref_np = state_ref.detach().cpu().float().numpy() + state_test_np = state_test.detach().cpu().float().numpy() + + print("\nDetailed state mismatch analysis:") + state_mismatch_mask = ~np.isclose( + state_ref_np, state_test_np, atol=atol, rtol=rtol + ) + num_state_mismatches = np.sum(state_mismatch_mask) + total_state_elements = state_ref_np.size + + print( + f"Number of mismatched state elements: {num_state_mismatches} / {total_state_elements} ({100 * num_state_mismatches / total_state_elements:.2f}%)" + ) + + state_mismatch_indices = np.argwhere(state_mismatch_mask) + print(f"First few state mismatch locations (up to 10):") + for i, idx in enumerate(state_mismatch_indices[:10]): + idx_tuple = tuple(idx) + ref_val = state_ref_np[idx_tuple] + test_val = state_test_np[idx_tuple] + diff = abs(ref_val - test_val) + rel_diff = diff / (abs(ref_val) + 1e-8) + print( + f" Index {idx_tuple}: ref={ref_val:.6f}, test={test_val:.6f}, diff={diff:.6e}, rel_diff={rel_diff:.6e}" + ) + return False + + assert(states_match) From 44d2b41e38a3e1f0d72e2c0cfbb39c30cee5b434 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 22 Dec 2025 13:44:15 -0800 Subject: [PATCH 02/38] check all the kernel inputs. ready to develop dispatch code. --- csrc/selective_state_update.cu | 208 +++++++++++++++--- .../mamba/selective_state_update.cuh | 67 +++--- tests/mamba/test_selective_state_update.py | 13 +- 3 files changed, 217 insertions(+), 71 deletions(-) diff --git a/csrc/selective_state_update.cu b/csrc/selective_state_update.cu index d8b4958417..31e5859e84 100644 --- a/csrc/selective_state_update.cu +++ b/csrc/selective_state_update.cu @@ -27,43 +27,181 @@ namespace flashinfer::mamba { // 2. Extracts raw pointers from TensorView // 3. Calls the framework-agnostic kernel from include/flashinfer/mamba/selective_state_update.cuh // -// Example signature (adjust based on your kernel requirements): void selective_state_update(TensorView state, TensorView x, TensorView dt, TensorView output, TensorView A, TensorView B, TensorView C, TensorView D, - Optional z, Optional dt_bias, - bool dt_softplus, Optional state_batch_indices, - int64_t pad_slot_id) { - throw std::runtime_error("selective_state_update is not implemented yet."); - - // TODO: Add input validation - // CHECK_LAST_DIM_CONTIGUOUS_INPUT(state); - // CHECK_LAST_DIM_CONTIGUOUS_INPUT(x); - // CHECK_DEVICE(state, x); - // TVM_FFI_ICHECK_EQ(state.ndim(), 3); // Example dimension check - - // TODO: Extract dimensions from tensors - // unsigned int batch_size = state.size(0); - // unsigned int dim_state = state.size(1); - // unsigned int dim_input = x.size(1); - - // TODO: Get CUDA device and stream - // ffi::CUDADeviceGuard device_guard(state.device().device_id); - // const cudaStream_t stream = get_stream(state.device()); - - // TODO: Dispatch based on dtype and call kernel - // DISPATCH_DLPACK_DTYPE_TO_CTYPE_FP16(state.dtype(), c_type, [&] { - // cudaError_t status = mamba::SelectiveStateUpdate( - // static_cast(state.data_ptr()), - // static_cast(x.data_ptr()), - // static_cast(dt.data_ptr()), - // static_cast(A.data_ptr()), - // static_cast(B.data_ptr()), - // static_cast(C.data_ptr()), - // batch_size, dim_state, dim_input, stream); - // TVM_FFI_ICHECK(status == cudaSuccess) - // << "SelectiveStateUpdate failed with error code " << cudaGetErrorString(status); - // return true; - // }); + Optional z, Optional dt_bias, bool dt_softplus, + Optional state_batch_indices, int64_t pad_slot_id) { + auto const batch = x.size(0); + auto const state_cache_size = state.size(0); + auto const nheads = state.size(1); + auto const dim = state.size(2); + auto const dstate = state.size(3); + auto const ngroups = B.size(1); + + FLASHINFER_CHECK(nheads % ngroups == 0, "nheads must be divisible by ngroups"); + + // Check x shape and strides + CHECK_DIM(3, x); + FLASHINFER_CHECK(x.size(1) == nheads, "x.size(1) must equal nheads"); + FLASHINFER_CHECK(x.size(2) == dim, "x.size(2) must equal dim"); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(x); + FLASHINFER_CHECK(x.stride(1) == dim, "x.stride(1) must equal dim, got ", x.stride(1), + " expected ", x.size(2)); + + // Check output shape and strides + CHECK_DIM(3, output); + CHECK_LAST_DIM_CONTIGUOUS(output); + FLASHINFER_CHECK(output.size(1) == nheads, "output.size(1) must equal nheads"); + FLASHINFER_CHECK(output.size(2) == dim, "output.size(2) must equal dim"); + FLASHINFER_CHECK(output.stride(1) == dim, "output.stride(1) must equal dim"); + + // Check dt shape and strides + CHECK_CUDA(dt); + CHECK_DIM(3, dt); // dt: {batch, nheads, dim} + FLASHINFER_CHECK(dt.size(1) == nheads, "dt.size(1) must equal nheads"); + FLASHINFER_CHECK(dt.size(2) == dim, "dt.size(2) must equal dim"); + FLASHINFER_CHECK(dt.stride(1) == 1, "dt.stride(1) must be 1, got ", dt.stride(1)); + FLASHINFER_CHECK(dt.stride(2) == 0, "dt.stride(2) must be 0 (broadcasted), got ", dt.stride(2)); + + // Check state - fully contiguous + CHECK_INPUT(state); // CUDA + fully contiguous (uses TVM FFI) + CHECK_DIM(4, state); // state: {state_cache_size, nheads, dim, dstate} + + // Check B shape and strides + CHECK_CUDA(B); + CHECK_DIM(3, B); // B: {batch, B.size(1), dstate} + FLASHINFER_CHECK(B.size(0) == batch, "B.size(0) must equal batch"); + FLASHINFER_CHECK(B.size(1) == ngroups, "B.size(1) must equal ngroups"); + FLASHINFER_CHECK(B.size(2) == dstate, "B.size(2) must equal dstate"); + CHECK_LAST_DIM_CONTIGUOUS(B); // stride(2) == 1 + FLASHINFER_CHECK(B.stride(1) == B.size(2), "B.stride(1) must equal dstate, got ", B.stride(1), + " expected ", B.size(2)); + + // Check C shape and strides + CHECK_CUDA(C); + CHECK_LAST_DIM_CONTIGUOUS(C); // stride(2) == 1 + CHECK_DIM(3, C); // C: {batch, C.size(1), dstate} + FLASHINFER_CHECK(C.stride(1) == C.size(2), "C.stride(1) must equal dstate, got ", C.stride(1), + " expected ", C.size(2)); + FLASHINFER_CHECK(C.size(0) == batch, "C.size(0) must equal batch"); + FLASHINFER_CHECK(C.size(1) == ngroups, "C.size(1) must equal ngroups"); + FLASHINFER_CHECK(C.size(2) == dstate, "C.size(2) must equal dstate"); + + // Check D - specific stride patterns indicating broadcasting + CHECK_CUDA(D); + CHECK_DIM(2, D); // D: {nheads, dim} + FLASHINFER_CHECK(D.size(0) == nheads, "D.size(0) must equal nheads"); + FLASHINFER_CHECK(D.size(1) == dim, "D.size(1) must equal dim"); + FLASHINFER_CHECK(D.stride(0) == 1, "D.stride(0) must be 1, got ", D.stride(0)); + FLASHINFER_CHECK(D.stride(1) == 0, "D.stride(1) must be 0 (broadcasted), got ", D.stride(1)); + + // Check A - specific stride patterns indicating broadcasting + CHECK_CUDA(A); + CHECK_DIM(3, A); // A: {nheads, dim, dstate} + FLASHINFER_CHECK(A.size(0) == nheads, "A.size(0) must equal nheads"); + FLASHINFER_CHECK(A.size(1) == dim, "A.size(1) must equal dim"); + FLASHINFER_CHECK(A.size(2) == dstate, "A.size(2) must equal dstate"); + FLASHINFER_CHECK(A.stride(1) == 0, "A.stride(1) must be 0 (broadcasted), got ", A.stride(1)); + FLASHINFER_CHECK(A.stride(2) == 0, "A.stride(2) must be 0 (broadcasted), got ", A.stride(2)); + + // Optional dt_bias check + if (dt_bias.has_value()) { + auto& bias = dt_bias.value(); + CHECK_CUDA(bias); + CHECK_DIM(2, bias); // dt_bias: {nheads, dim} + FLASHINFER_CHECK(bias.size(0) == nheads, "dt_bias.size(0) must equal nheads"); + FLASHINFER_CHECK(bias.size(1) == dim, "dt_bias.size(1) must equal dim"); + FLASHINFER_CHECK(bias.stride(0) == 1, "dt_bias.stride(0) must be 1, got ", bias.stride(0)); + FLASHINFER_CHECK(bias.stride(1) == 0, "dt_bias.stride(1) must be 0 (broadcasted), got ", + bias.stride(1)); + } + + // if(z.has_value()) + FLASHINFER_CHECK(!z.has_value() && "z is not supported yet"); + + if (state_batch_indices) { + CHECK_DIM(1, (*state_batch_indices)); + FLASHINFER_CHECK(state_batch_indices.value().size(0) == batch, + "state_batch_indices.shape must be (", batch, ")"); + } + + SelectiveStateUpdateParams p; + + // copy dimensions + p.batch = batch; + p.nheads = nheads; + p.dim = dim; + p.dstate = dstate; + p.ngroups = ngroups; + p.state_cache_size = state_cache_size; + p.dt_softplus = dt_softplus; + p.pad_slot_id = pad_slot_id; + + // Copy strides + p.x_stride_batch = x.stride(0); + p.dt_stride_batch = dt.stride(0); + p.B_stride_batch = B.stride(0); + p.C_stride_batch = C.stride(0); + p.out_stride_batch = output.stride(0); + if (state_batch_indices) p.state_batch_indices = state_batch_indices.value().data_ptr(); + + // Copy pointers + p.state = state.data_ptr(); + p.x = x.data_ptr(); + p.dt = dt.data_ptr(); + p.output = output.data_ptr(); + if (dt_bias) { + p.dt_bias = dt_bias.value().data_ptr(); + } + if (z) { + p.z = z.value().data_ptr(); + } + p.A = A.data_ptr(); + p.B = B.data_ptr(); + p.C = C.data_ptr(); + p.D = D.data_ptr(); + + // Set device and get stream + ffi::CUDADeviceGuard device_guard(state.device().device_id); + const cudaStream_t stream = get_stream(state.device()); + + // Dispatch based on dtype combination + DLDataType state_dtype = state.dtype(); + DLDataType input_dtype = x.dtype(); + DLDataType weight_dtype = dt.dtype(); + DLDataType matrixA_dtype = A.dtype(); + + int64_t state_dtype_code = encode_dlpack_dtype(state_dtype); + int64_t input_dtype_code = encode_dlpack_dtype(input_dtype); + int64_t weight_dtype_code = encode_dlpack_dtype(weight_dtype); + int64_t matrixA_dtype_code = encode_dlpack_dtype(matrixA_dtype); + + // Pack all dtype codes into a single value for switching + auto dtype_key = + std::make_tuple(state_dtype_code, input_dtype_code, weight_dtype_code, matrixA_dtype_code); + + // Currently only support: input_t = weight_t = state_t = bfloat16, matrixA_t = float + if (dtype_key == std::make_tuple(bfloat16_code, bfloat16_code, bfloat16_code, float32_code)) { + using state_t = nv_bfloat16; + using input_t = nv_bfloat16; + using weight_t = nv_bfloat16; + using matrixA_t = float; + + invokeSelectiveStateUpdate(p, stream); + } else { + // Default case: unsupported dtype combination + TVM_FFI_ICHECK(false) << "Unsupported dtype combination for selective_state_update: " + << "state_dtype=" << state_dtype.code << ":" << state_dtype.bits << ", " + << "input_dtype=" << input_dtype.code << ":" << input_dtype.bits << ", " + << "weight_dtype=" << weight_dtype.code << ":" << weight_dtype.bits + << ", " + << "matrixA_dtype=" << matrixA_dtype.code << ":" << matrixA_dtype.bits + << ". Currently only support: " + << "state=bfloat16, input=bfloat16, weight=bfloat16, matrixA=float32"; + + throw std::runtime_error("selective_state_update is not implemented yet."); + } + } } // namespace flashinfer::mamba diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh index c1a272f6d1..d3ca3ff640 100644 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -19,35 +19,38 @@ #include "../utils.cuh" #include "../vec_dtypes.cuh" -namespace flashinfer { - -namespace mamba { - -// TODO: Add your kernel implementation here -// This should be framework-agnostic and use raw pointers only (no PyTorch/TensorFlow types) -// -// Example: -// template -// __global__ void SelectiveStateUpdateKernel( -// T* state, const T* x, const T* dt, const T* A, const T* B, const T* C, -// int batch_size, int dim_state, int dim_input) { -// // Your kernel implementation -// } - -// TODO: Add host function that launches the kernel -// Example: -// template -// cudaError_t SelectiveStateUpdate( -// T* state, const T* x, const T* dt, const T* A, const T* B, const T* C, -// int batch_size, int dim_state, int dim_input, -// cudaStream_t stream = 0) { -// // Compute grid/block dimensions -// // Launch kernel -// // Return cudaGetLastError() -// } - -} // namespace mamba - -} // namespace flashinfer - -#endif // FLASHINFER_MAMBA_SELECTIVE_STATE_UPDATE_CUH_ \ No newline at end of file +namespace flashinfer::mamba { + +struct SelectiveStateUpdateParams { + uint32_t batch{}, nheads{}, dim{}, dstate{}, ngroups{}, state_cache_size{}; + int32_t pad_slot_id{-1}; + bool dt_softplus{false}; + + int64_t x_stride_batch{}, dt_stride_batch{}, B_stride_batch{}, C_stride_batch{}, + out_stride_batch{}; + + void* __restrict__ state{nullptr}; // state_t: (state_cache_size, nheads, dim, dstate) + void* __restrict__ x{nullptr}; // input_t: (batch, nheads, dim) + void* __restrict__ dt{nullptr}; // weight_t: (batch, nheads) but pretends to be (batch, nheads, dim) + void* __restrict__ dt_bias{nullptr}; // weight_t (nheads) but pretends to be (nheads, dim) + void* __restrict__ A{nullptr}; // matrixA_t: (nheads) but pretends to be (nheads, dim, dstate) + void* __restrict__ B{nullptr}; // input_t: (batch, ngroups, dstate) + void* __restrict__ C{nullptr}; // input_t: (batch, ngroups, dstate) + void* __restrict__ D{nullptr}; // weight_t: (nheads) but pretends to be (nheads, dim) + void* __restrict__ z{nullptr}; // input_t: (batch, nheads, dim) + void* __restrict__ output{nullptr}; // input_t: (batch, nheads, dim) + void* __restrict__ state_batch_indices{nullptr}; // state_batch_indices: (batch,) +}; + +template +void invokeSelectiveStateUpdate( SelectiveStateUpdateParams& params, cudaStream_t stream) +{ + // This function is implemented in selective_state_update_kernel.cu + throw std::runtime_error( + "invokeSelectiveStateUpdate is not implemented for the given data types."); +} + + +} // namespace flashinfer::mamba + +#endif // FLASHINFER_MAMBA_SELECTIVE_STATE_UPDATE_CUH_ diff --git a/tests/mamba/test_selective_state_update.py b/tests/mamba/test_selective_state_update.py index 465d824215..fbe0443565 100644 --- a/tests/mamba/test_selective_state_update.py +++ b/tests/mamba/test_selective_state_update.py @@ -105,9 +105,14 @@ def create_test_inputs( (batch_size, nheads, dim), (nheads, 1, 0) ) - A = -torch.rand(nheads, dtype=torch.float32, device=device).as_strided( - (nheads, dim, dstate), (1, 0, 0) - ) + A_base = -torch.rand(nheads, dtype=torch.float32, device=device) + A = A_base.as_strided((nheads, dim, dstate), (1, 0, 0)) + + # A = -torch.rand(nheads, dtype=torch.float32, device=device).as_strided( + # (nheads, dim, dstate), (1, 0, 0) + # ) + # print(f"A dtype: {A.dtype}, shape: {A.shape}, stride {A.stride()}") + assert(A.stride() == (1, 0, 0)) # B and C - (batch_size, ngroups, dstate) B = torch.randn(batch_size, ngroups, dstate, dtype=input_dtype, device=device) @@ -203,7 +208,7 @@ def test_selective_state_update( state, inputs["x"], inputs["dt"], - A, + inputs["A"], inputs["B"], inputs["C"], D=inputs["D"], From 3d00fdf30a8874cb4d312b44e6327b3075bfbd7d Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 5 Jan 2026 11:31:44 -0800 Subject: [PATCH 03/38] simple implementation of selective_state_update is working --- .../jit/mamba/selective_state_update.py | 66 ++++ include/flashinfer/mamba/conversion.cuh | 31 ++ .../mamba/selective_state_update.cuh | 216 ++++++++++- tests/mamba/__init__.py | 15 + tests/mamba/selective_state_update_triton.py | 345 ++++++++++++++++++ tests/mamba/test_selective_state_update.py | 50 ++- 6 files changed, 701 insertions(+), 22 deletions(-) create mode 100644 flashinfer/jit/mamba/selective_state_update.py create mode 100644 include/flashinfer/mamba/conversion.cuh create mode 100644 tests/mamba/__init__.py create mode 100644 tests/mamba/selective_state_update_triton.py diff --git a/flashinfer/jit/mamba/selective_state_update.py b/flashinfer/jit/mamba/selective_state_update.py new file mode 100644 index 0000000000..59b26e1ef5 --- /dev/null +++ b/flashinfer/jit/mamba/selective_state_update.py @@ -0,0 +1,66 @@ +""" +Copyright (c) 2025 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" + +from .. import env as jit_env +from ..core import JitSpec, gen_jit_spec + + +def gen_selective_state_update_module() -> JitSpec: + """Generate JIT module for selective_state_update operation. + + This function: + 1. Computes unique identifier from parameters (if needed) + 2. Copies source files to generation directory + 3. Returns JitSpec with compilation metadata + + Note: This is a simple example without type specialization. + If you need different code for different dtypes/parameters, + you can add Jinja template rendering here (see flashinfer/jit/sampling.py). + """ + + # TODO: If you need type specialization, add parameters like: + # def gen_selective_state_update_module(dtype_in, dtype_out, ...) -> JitSpec: + # uri = get_selective_state_update_uri(dtype_in, dtype_out, ...) + # gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri + # + # # Render Jinja template for type-specialized config + # with open(jit_env.FLASHINFER_CSRC_DIR / "selective_state_update_customize_config.jinja") as f: + # template = jinja2.Template(f.read()) + # config_content = template.render(dtype_in=dtype_map[dtype_in], ...) + # write_if_different(gen_directory / "selective_state_update_config.inc", config_content) + # + # # Copy source files + # sources = [] + # for fname in ["selective_state_update.cu", "flashinfer_mamba_binding.cu"]: + # shutil.copy(jit_env.FLASHINFER_CSRC_DIR / fname, gen_directory / fname) + # sources.append(gen_directory / fname) + # + # return gen_jit_spec(uri, sources, extra_cuda_cflags=[...]) + + # Simple version without type specialization (like norm.py): + nvcc_flags = [ + "-DENABLE_BF16", + "-DENABLE_FP8", + ] + + return gen_jit_spec( + "mamba_selective_state_update", + [ + jit_env.FLASHINFER_CSRC_DIR / "selective_state_update.cu", + jit_env.FLASHINFER_CSRC_DIR / "flashinfer_mamba_binding.cu", + ], + extra_cuda_cflags=nvcc_flags, + ) \ No newline at end of file diff --git a/include/flashinfer/mamba/conversion.cuh b/include/flashinfer/mamba/conversion.cuh new file mode 100644 index 0000000000..7e4fba558e --- /dev/null +++ b/include/flashinfer/mamba/conversion.cuh @@ -0,0 +1,31 @@ +#pragma once +#include +#include +#include + +#ifdef ENABLE_FP8 +#include +#endif + +namespace flashinfer::mamba::conversion { + +inline __device__ float toFloat(float f) { return f; } + +inline __device__ float toFloat(__half h) { return __half2float(h); } +#ifdef ENABLE_BF16 +inline __device__ float toFloat(__nv_bfloat16 val) { return __bfloat162float(val); } +#endif + +inline __device__ void convertAndStore(float* output, float input) { *output = input; } + +inline __device__ void convertAndStore(__half* output, float input) { + *output = __float2half(input); +} + +#ifdef ENABLE_BF16 +inline __device__ void convertAndStore(__nv_bfloat16* output, float input) { + *output = __float2bfloat16(input); +} +#endif + +} // namespace flashinfer::mamba::conversion diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh index d3ca3ff640..3434868ba1 100644 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -18,9 +18,12 @@ #include "../utils.cuh" #include "../vec_dtypes.cuh" +#include "conversion.cuh" namespace flashinfer::mamba { +using namespace conversion; + struct SelectiveStateUpdateParams { uint32_t batch{}, nheads{}, dim{}, dstate{}, ngroups{}, state_cache_size{}; int32_t pad_slot_id{-1}; @@ -31,7 +34,8 @@ struct SelectiveStateUpdateParams { void* __restrict__ state{nullptr}; // state_t: (state_cache_size, nheads, dim, dstate) void* __restrict__ x{nullptr}; // input_t: (batch, nheads, dim) - void* __restrict__ dt{nullptr}; // weight_t: (batch, nheads) but pretends to be (batch, nheads, dim) + void* __restrict__ dt{ + nullptr}; // weight_t: (batch, nheads) but pretends to be (batch, nheads, dim) void* __restrict__ dt_bias{nullptr}; // weight_t (nheads) but pretends to be (nheads, dim) void* __restrict__ A{nullptr}; // matrixA_t: (nheads) but pretends to be (nheads, dim, dstate) void* __restrict__ B{nullptr}; // input_t: (batch, ngroups, dstate) @@ -42,14 +46,214 @@ struct SelectiveStateUpdateParams { void* __restrict__ state_batch_indices{nullptr}; // state_batch_indices: (batch,) }; -template -void invokeSelectiveStateUpdate( SelectiveStateUpdateParams& params, cudaStream_t stream) +__forceinline__ __device__ float softplus(float x) +{ + return __logf(1.f + __expf(x)); + // return log1pf(exp(x)); +} + +__device__ __forceinline__ float fast_exp(float x) +{ + constexpr float log2_E = 1.4426950408889634f; + float y; + asm("ex2.approx.f32 %0,%1;" : "=f"(y) : "f"(x * log2_E)); + return y; +} + +__device__ __forceinline__ float thresholded_softplus(float dt_value) +{ + constexpr float threshold = 20.f; + return (dt_value <= threshold) ? softplus(dt_value) : dt_value; +} + +template +__device__ inline auto make_zero() -> T; + +template <> +__device__ inline auto make_zero() -> float2 { - // This function is implemented in selective_state_update_kernel.cu - throw std::runtime_error( - "invokeSelectiveStateUpdate is not implemented for the given data types."); + return make_float2(0.f, 0.f); } +template +__device__ inline auto make_zeros() -> load_t +{ + load_t rValue; +#pragma unroll + for (int i = 0; i < sizeof(load_t) / sizeof(compute_t); i++) + { + auto* dst = reinterpret_cast(&rValue) + i; + convertAndStore(dst, 0.f); + } + return rValue; +} + +__device__ __forceinline__ float warpReduceSum(float val) +{ + constexpr auto warpSize = 32; + for (int s = warpSize / 2; s > 0; s /= 2) + { + val += __shfl_down_sync(UINT32_MAX, val, s); + } + return val; +} + +template +struct VectorizedLoadTraits{}; + +template <> +struct VectorizedLoadTraits<__nv_bfloat16, __nv_bfloat16, __nv_bfloat16> { + using input = float2; + using weight = float2; + using state = float2; + static constexpr auto chunk_size = sizeof(input) / sizeof(__nv_bfloat16); +}; + +template +__global__ void selective_state_update_kernel_simple(SelectiveStateUpdateParams params) { + auto* __restrict__ output = + reinterpret_cast(params.output); // output: (batch, nheads, dim) + auto* __restrict__ state = + reinterpret_cast(params.state); // state: (batch, nheads, dim, dstate) + + auto const* __restrict__ x = + reinterpret_cast(params.x); // x: (batch, nheads, dim) + auto const* __restrict__ dt = + reinterpret_cast(params.dt); // dt: (batch, nheads) + auto const* __restrict__ A = reinterpret_cast(params.A); // A: (nheads) + auto const* __restrict__ B = + reinterpret_cast(params.B); // B: (batch, ngroups, dstate) + auto const* __restrict__ C = + reinterpret_cast(params.C); // C: (batch, ngroups, dstate) + auto const* __restrict__ D = reinterpret_cast(params.D); // D: (nheads, dim) + auto const* __restrict__ dt_bias = reinterpret_cast(params.dt_bias); // (nheads) + auto const* __restrict__ z = reinterpret_cast(params.z); + auto const* __restrict__ state_batch_indices = + reinterpret_cast(params.state_batch_indices); + bool const dt_softplus = params.dt_softplus; + + int const nheads = params.nheads; + int const ngroups = params.ngroups; + int const dim = params.dim; + + constexpr auto warpSize = 32; + auto const dim_offset = blockIdx.x * warpSize * numWarps; + auto const batch = blockIdx.y; + auto const head = blockIdx.z; + auto const group = head / (nheads / ngroups); + auto lane = threadIdx.x % warpSize; + auto warp = threadIdx.y; + + auto const state_batch = (state_batch_indices) ? state_batch_indices[batch] : batch; + state += state_batch * nheads * dim * DSTATE + head * dim * DSTATE; + + __shared__ input_t sx[numWarps * warpSize]; + __shared__ float sdt[numWarps * warpSize]; + __shared__ weight_t sz[numWarps * warpSize]; + + auto const A_value = toFloat(A[head]); + + auto dt_value = toFloat(dt[batch * params.dt_stride_batch + head]); + if (dt_bias) dt_value += toFloat(dt_bias[head]); + if (dt_softplus) { + dt_value = thresholded_softplus(dt_value); + } + + // auto const dA = fast_exp(A_value * dt_value); + auto const dA = __expf(A_value * dt_value); + // auto const dA = exp(A_value * dt_value); + + auto d_value = D ? toFloat(D[head]) : 0.f; + + auto _d = warp * warpSize + lane; + auto d = dim_offset + _d; + if (d < dim) { + sx[_d] = x[batch * params.x_stride_batch + head * dim + d]; + if (z) { + sz[_d] = z[batch * nheads * dim + head * dim + d]; + } else { + convertAndStore(&sz[_d], 0.f); + } + } else { + convertAndStore(&sx[_d], 0.f); + convertAndStore(&sz[_d], 0.f); + } + + using Load = VectorizedLoadTraits; + + for (auto _d = warp * warpSize; _d < (warp + 1) * warpSize; _d++) { + auto d = dim_offset + _d; + if (d >= dim) break; + + float x_value = toFloat(sx[_d]); + float out_value = d_value * x_value * int(lane == 0); // first lane has the value + + for (int i = threadIdx.x * Load::chunk_size; i < DSTATE; i += warpSize * Load::chunk_size) { + auto rState = make_zeros(); + if (state_batch != params.pad_slot_id) + rState = *reinterpret_cast(&state[d * DSTATE + i]); + auto rB = *reinterpret_cast( + &B[batch * params.B_stride_batch + group * DSTATE + i]); + auto rC = *reinterpret_cast( + &C[batch * params.C_stride_batch + group * DSTATE + i]); + + auto* state_vals = reinterpret_cast(&rState); + auto const* B_vals = reinterpret_cast(&rB); + auto const* C_vals = reinterpret_cast(&rC); + + for (int ii = 0; ii < Load::chunk_size; ii++) { + auto state_value = toFloat(state_vals[ii]); + auto B_value = toFloat(B_vals[ii]); + auto C_value = toFloat(C_vals[ii]); + + auto const dB = B_value * dt_value; + auto const new_state = state_value * dA + dB * x_value; + convertAndStore(&state_vals[ii], new_state); + + out_value += new_state * C_value; + } + if (state_batch != params.pad_slot_id) + *reinterpret_cast(&state[d * DSTATE + i]) = rState; + } + + // warpReduce the out_value + out_value = warpReduceSum(out_value); + if (lane == 0) { + sdt[_d] = out_value; + } + } + + if (d < dim) { + auto out_value = sdt[_d]; + if (z) { + float z_value = toFloat(sz[_d]); + float sig_z = __fdividef(1.f, (1.f + __expf(0.f - z_value))); + float silu_z = z_value * sig_z; + out_value *= silu_z; + } + convertAndStore(&output[batch * params.out_stride_batch + head * dim + d], out_value); + } +} + +template +void invokeSelectiveStateUpdate(SelectiveStateUpdateParams& params, cudaStream_t stream) { + + constexpr int DSTATE = 128; + FLASHINFER_CHECK(params.dstate == DSTATE); + + // #if (__CUDA_ARCH__ < 900) // pre-Hopper + { + constexpr int numWarps = 2; + int const blocks_per_dim = (params.dim + 32 * numWarps - 1) / (32 * numWarps); + dim3 block(32, numWarps); + dim3 grid(blocks_per_dim, params.batch, params.nheads); + selective_state_update_kernel_simple + <<>>(params); + } + // #else + // #endif +} } // namespace flashinfer::mamba diff --git a/tests/mamba/__init__.py b/tests/mamba/__init__.py new file mode 100644 index 0000000000..0a2157a5eb --- /dev/null +++ b/tests/mamba/__init__.py @@ -0,0 +1,15 @@ +""" +Copyright (c) 2025 by FlashInfer team. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +""" \ No newline at end of file diff --git a/tests/mamba/selective_state_update_triton.py b/tests/mamba/selective_state_update_triton.py new file mode 100644 index 0000000000..232c748edd --- /dev/null +++ b/tests/mamba/selective_state_update_triton.py @@ -0,0 +1,345 @@ +# Adapted from https://github.com/state-spaces/mamba/blob/v2.2.4/mamba_ssm/ops/triton/selective_state_update.py +# Copyright (c) 2024, Tri Dao, Albert Gu. +# +# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch +import triton +import triton.language as tl +from packaging import version + +TRITON3 = version.parse(triton.__version__) >= version.parse("3.0.0") + +if TRITON3: + + @triton.jit + def softplus(dt): + return tl.math.log(tl.math.exp(dt) + 1) + +else: + + @triton.jit + def softplus(dt): + return tl.math.log1p(tl.exp(dt)) + +PAD_SLOT_ID = -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_scan_update_kernel( + # Pointers to matrices + 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, + pad_slot_id, + # Matrix dimensions + batch, + nheads, + dim, + dstate, + nheads_ngroups_ratio, + # Strides + stride_state_batch, + stride_state_head, + stride_state_dim, + stride_state_dstate, + stride_x_batch, + stride_x_head, + stride_x_dim, + stride_dt_batch, + stride_dt_head, + stride_dt_dim, + stride_dt_bias_head, + stride_dt_bias_dim, + stride_A_head, + stride_A_dim, + stride_A_dstate, + stride_B_batch, + stride_B_group, + stride_B_dstate, + stride_C_batch, + stride_C_group, + stride_C_dstate, + stride_D_head, + stride_D_dim, + stride_z_batch, + stride_z_head, + stride_z_dim, + stride_out_batch, + stride_out_head, + stride_out_dim, + # Meta-parameters + 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) + + # If HAS_STATE_BATCH_INDICES is true, then the ssm state's batch coordinate + # is taken from the state_batch_indices_ptr Otherwise, the state coordinate + # is the same as the batch id. + if HAS_STATE_BATCH_INDICES: + state_batch_indices_ptr += pid_b + state_batch_idx = tl.load(state_batch_indices_ptr) + state_ptr += state_batch_idx * stride_state_batch + pid_h * stride_state_head + else: + state_ptr += pid_b * stride_state_batch + pid_h * stride_state_head + + x_ptr += pid_b * stride_x_batch + pid_h * stride_x_head + dt_ptr += pid_b * stride_dt_batch + pid_h * stride_dt_head + if HAS_DT_BIAS: + dt_bias_ptr += pid_h * stride_dt_bias_head + A_ptr += pid_h * stride_A_head + B_ptr += pid_b * stride_B_batch + (pid_h // + nheads_ngroups_ratio) * stride_B_group + C_ptr += pid_b * stride_C_batch + (pid_h // + nheads_ngroups_ratio) * stride_C_group + if HAS_Z: + z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head + out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head + + offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) + state_ptrs = state_ptr + (offs_m[:, None] * stride_state_dim + + offs_n[None, :] * stride_state_dstate) + x_ptrs = x_ptr + offs_m * stride_x_dim + dt_ptrs = dt_ptr + offs_m * stride_dt_dim + if HAS_DT_BIAS: + dt_bias_ptrs = dt_bias_ptr + offs_m * stride_dt_bias_dim + if HAS_D: + D_ptr += pid_h * stride_D_head + A_ptrs = A_ptr + (offs_m[:, None] * stride_A_dim + + offs_n[None, :] * stride_A_dstate) + B_ptrs = B_ptr + offs_n * stride_B_dstate + C_ptrs = C_ptr + offs_n * stride_C_dstate + if HAS_D: + D_ptrs = D_ptr + offs_m * stride_D_dim + if HAS_Z: + z_ptrs = z_ptr + offs_m * stride_z_dim + out_ptrs = out_ptr + offs_m * stride_out_dim + mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate) + if HAS_STATE_BATCH_INDICES: + mask &= (state_batch_idx != pad_slot_id) + state = tl.load(state_ptrs, mask=mask, 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) # scalar, not a matrix + + 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 # vector of size (dstate,) + state = state * dA + dB * x[:, None] + + mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate) + if HAS_STATE_BATCH_INDICES: + mask &= (state_batch_idx != pad_slot_id) + tl.store(state_ptrs, state, mask=mask) + 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) + + +def selective_state_update_triton(state, + x, + dt, + A, + B, + C, + D=None, + z=None, + dt_bias=None, + dt_softplus=False, + state_batch_indices=None, + pad_slot_id=PAD_SLOT_ID): + """ + Argument: + state: (batch, dim, dstate) or (batch, nheads, dim, dstate) + x: (batch, dim) or (batch, nheads, dim) + dt: (batch, dim) or (batch, nheads, dim) + A: (dim, dstate) or (nheads, dim, dstate) + B: (batch, dstate) or (batch, ngroups, dstate) + C: (batch, dstate) or (batch, ngroups, dstate) + D: (dim,) or (nheads, dim) + z: (batch, dim) or (batch, nheads, dim) + dt_bias: (dim,) or (nheads, dim) + pad_slot_id: int + if cache_indices is passed, lets the kernel identify padded + entries that will not be processed, + for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id] + in this case, the kernel will not process entries at + indices 0 and 3 + Return: + out: (batch, dim) or (batch, nheads, dim) + """ + 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) + + _, nheads, dim, dstate = state.shape + batch = x.shape[0] + + assert x.shape == (batch, nheads, dim) + assert dt.shape == x.shape + assert A.shape == (nheads, dim, dstate) + ngroups = B.shape[1] + assert nheads % ngroups == 0, "nheads must be divisible by ngroups" + assert B.shape == (batch, ngroups, dstate) + assert C.shape == B.shape + if D is not None: + assert D.shape == (nheads, dim) + if z is not None: + assert z.shape == x.shape + if dt_bias is not None: + assert dt_bias.shape == (nheads, dim) + if state_batch_indices is not None: + assert state_batch_indices.shape == (batch, ) + out = torch.empty_like(x) + grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) + z_strides = (z.stride(0), z.stride(1), + z.stride(2)) if z is not None else (0, 0, 0) + # We don't want autotune since it will overwrite the state + # We instead tune by hand. + BLOCK_SIZE_M, num_warps = ((32, 4) if dstate <= 16 else + ((16, 4) if dstate <= 32 else + ((8, 4) if dstate <= 64 else + ((4, 4) if dstate <= 128 else ((4, 8)))))) + tie_hdim = (A.stride(-1) == 0 and A.stride(-2) == 0 and dt.stride(-1) == 0 + and dt_bias.stride(-1) == 0) + with torch.cuda.device(x.device.index): + _selective_scan_update_kernel[grid]( + state, + x, + dt, + dt_bias, + A, + B, + C, + D, + z, + out, + state_batch_indices, + pad_slot_id, + batch, + nheads, + dim, + dstate, + nheads // ngroups, + 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), + 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), 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_SIZE_M, + num_warps=num_warps, + ) + if not has_heads: + out = out.squeeze(1) + return out diff --git a/tests/mamba/test_selective_state_update.py b/tests/mamba/test_selective_state_update.py index fbe0443565..e675b99726 100644 --- a/tests/mamba/test_selective_state_update.py +++ b/tests/mamba/test_selective_state_update.py @@ -3,6 +3,7 @@ import torch import torch.nn.functional as F from einops import rearrange, repeat +from .selective_state_update_triton import selective_state_update_triton import flashinfer @@ -187,22 +188,39 @@ def test_selective_state_update( matrixA_dtype, ) - state = inputs["state_cache"][inputs["slot_idx"]] - state_ref = rearrange(state, "... p n -> ... n p").detach().clone() - A = inputs["A"] - A_ref = (rearrange(A, "... p n -> ... n p") if A.ndim == 3 else A).detach().clone() - y_ref = selective_state_update_ref( - state_ref, - inputs["x"], - inputs["dt"], - A_ref, - inputs["B"], - inputs["C"], - D=inputs["D"], - z=None, - dt_bias=inputs["dt_bias"], - dt_softplus=delta_softplus, - ) + # state = inputs["state_cache"][inputs["slot_idx"]] + # state_ref = rearrange(state, "... p n -> ... n p").detach().clone() + state = inputs["state_cache"] + state_ref = state.clone() + # A = inputs["A"] + + # A_ref = (rearrange(A, "... p n -> ... n p") if A.ndim == 3 else A).detach().clone() + # y_ref = selective_state_update_ref( + # state_ref, + # inputs["x"], + # inputs["dt"], + # A_ref, + # inputs["B"], + # inputs["C"], + # D=inputs["D"], + # z=None, + # dt_bias=inputs["dt_bias"], + # dt_softplus=delta_softplus, + # ) + y_ref = selective_state_update_triton( + state_ref, + inputs["x"], + inputs["dt"], + inputs["A"], + inputs["B"], + inputs["C"], + D=inputs["D"], + z=None, + dt_bias=inputs["dt_bias"], + dt_softplus=delta_softplus, + state_batch_indices=inputs["slot_idx"], + pad_slot_id=-1, + ) y_test = flashinfer.mamba.selective_state_update( state, From ebe907027204eb2aeda068201b1c2620b7ac1571 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 5 Jan 2026 14:19:32 -0800 Subject: [PATCH 04/38] ported the hopper version + runtime dispatch check --- .../flashinfer/mamba/create_tensor_map.cuh | 86 +++++ .../mamba/selective_state_update.cuh | 321 +++++++++++++++--- 2 files changed, 368 insertions(+), 39 deletions(-) create mode 100644 include/flashinfer/mamba/create_tensor_map.cuh diff --git a/include/flashinfer/mamba/create_tensor_map.cuh b/include/flashinfer/mamba/create_tensor_map.cuh new file mode 100644 index 0000000000..aa2e9fef46 --- /dev/null +++ b/include/flashinfer/mamba/create_tensor_map.cuh @@ -0,0 +1,86 @@ +#pragma once +#include +#include // PFN_cuTensorMapEncodeTiled, CUtensorMap +#include + +#include +#include +#include +#include + +#ifndef gpuErrchk +#define gpuErrchk(ans) \ + { gpuAssert((ans), __FILE__, __LINE__); } +#endif + +static inline void gpuAssert(cudaError_t code, const char* file, int line, bool abort = true) { + if (code != cudaSuccess) { + fprintf(stderr, "GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line); + std::cout << "GPU assert failed" << std::endl; + if (abort) exit(code); + } +} +namespace flashinfer::mamba::tma { + +// namespace cde = cuda::device::experimental; + +static inline PFN_cuTensorMapEncodeTiled_v12000 get_cuTensorMapEncodeTiled() { + // Get pointer to cuTensorMapEncodeTiled + cudaDriverEntryPointQueryResult driver_status; + void* cuTensorMapEncodeTiled_ptr = nullptr; + gpuErrchk(cudaGetDriverEntryPointByVersion("cuTensorMapEncodeTiled", &cuTensorMapEncodeTiled_ptr, + 12000, cudaEnableDefault, &driver_status)); + + if (driver_status != cudaDriverEntryPointSuccess) { + std::cerr << "Could not get cuTensorMapEncodeTiled driver entry point" << std::endl; + abort(); + } + + return reinterpret_cast(cuTensorMapEncodeTiled_ptr); +} + +template +inline CUtensorMap createTensorMap(void* matrix_ptr, uint32_t matrix_height, uint32_t matrix_width, + uint32_t tile_height, uint32_t tile_width) { + CUtensorMap tensor_map{}; + constexpr uint32_t rank = 2; + + std::array matrix_dim = {matrix_width, matrix_height}; + std::array stride = {matrix_width * sizeof(Dtype)}; + std::array box_size = {tile_width, tile_height}; + std::array elem_stride = {1, 1}; + + // CUtensorMapDataType dtype_format = CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_FLOAT16; + CUtensorMapDataType dtype_format; + if constexpr (std::is_same_v) { + dtype_format = CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_FLOAT16; + } else if constexpr (std::is_same_v) { + dtype_format = CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_FLOAT32; + } else if constexpr (std::is_same_v) { + dtype_format = CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_BFLOAT16; + } else { + static_assert([]() { return false; }(), "Unsupported data type for TMA tensor map"); + return tensor_map; // shut the compiler up + } + + auto cuTensorMapEncodeTiled = get_cuTensorMapEncodeTiled(); + CUresult res = cuTensorMapEncodeTiled( + &tensor_map, dtype_format, rank, matrix_ptr, matrix_dim.data(), stride.data(), + box_size.data(), elem_stride.data(), CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE, + CU_TENSOR_MAP_SWIZZLE_NONE, CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE, + CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE); + + if (res != CUDA_SUCCESS) { + const char* err_name = nullptr; + const char* err_str = nullptr; + cuGetErrorName(res, &err_name); + cuGetErrorString(res, &err_str); + std::cerr << "Could not create a tensor map" << std::endl; + std::cerr << "Error is: " << err_name << ": " << err_str << std::endl; + abort(); + } + + return tensor_map; +} + +} // namespace flashinfer::mamba::tma diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh index 3434868ba1..c200aa8d88 100644 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -16,9 +16,18 @@ #ifndef FLASHINFER_MAMBA_SELECTIVE_STATE_UPDATE_CUH_ #define FLASHINFER_MAMBA_SELECTIVE_STATE_UPDATE_CUH_ +#include +#include +#include + +#include +#include +#include + #include "../utils.cuh" #include "../vec_dtypes.cuh" #include "conversion.cuh" +#include "create_tensor_map.cuh" namespace flashinfer::mamba { @@ -46,60 +55,52 @@ struct SelectiveStateUpdateParams { void* __restrict__ state_batch_indices{nullptr}; // state_batch_indices: (batch,) }; -__forceinline__ __device__ float softplus(float x) -{ - return __logf(1.f + __expf(x)); - // return log1pf(exp(x)); +__forceinline__ __device__ float softplus(float x) { + return __logf(1.f + __expf(x)); + // return log1pf(exp(x)); } -__device__ __forceinline__ float fast_exp(float x) -{ - constexpr float log2_E = 1.4426950408889634f; - float y; - asm("ex2.approx.f32 %0,%1;" : "=f"(y) : "f"(x * log2_E)); - return y; +__device__ __forceinline__ float fast_exp(float x) { + constexpr float log2_E = 1.4426950408889634f; + float y; + asm("ex2.approx.f32 %0,%1;" : "=f"(y) : "f"(x * log2_E)); + return y; } -__device__ __forceinline__ float thresholded_softplus(float dt_value) -{ - constexpr float threshold = 20.f; - return (dt_value <= threshold) ? softplus(dt_value) : dt_value; +__device__ __forceinline__ float thresholded_softplus(float dt_value) { + constexpr float threshold = 20.f; + return (dt_value <= threshold) ? softplus(dt_value) : dt_value; } template __device__ inline auto make_zero() -> T; template <> -__device__ inline auto make_zero() -> float2 -{ - return make_float2(0.f, 0.f); +__device__ inline auto make_zero() -> float2 { + return make_float2(0.f, 0.f); } template -__device__ inline auto make_zeros() -> load_t -{ - load_t rValue; +__device__ inline auto make_zeros() -> load_t { + load_t rValue; #pragma unroll - for (int i = 0; i < sizeof(load_t) / sizeof(compute_t); i++) - { - auto* dst = reinterpret_cast(&rValue) + i; - convertAndStore(dst, 0.f); - } - return rValue; + for (int i = 0; i < sizeof(load_t) / sizeof(compute_t); i++) { + auto* dst = reinterpret_cast(&rValue) + i; + convertAndStore(dst, 0.f); + } + return rValue; } -__device__ __forceinline__ float warpReduceSum(float val) -{ - constexpr auto warpSize = 32; - for (int s = warpSize / 2; s > 0; s /= 2) - { - val += __shfl_down_sync(UINT32_MAX, val, s); - } - return val; +__device__ __forceinline__ float warpReduceSum(float val) { + constexpr auto warpSize = 32; + for (int s = warpSize / 2; s > 0; s /= 2) { + val += __shfl_down_sync(UINT32_MAX, val, s); + } + return val; } template -struct VectorizedLoadTraits{}; +struct VectorizedLoadTraits {}; template <> struct VectorizedLoadTraits<__nv_bfloat16, __nv_bfloat16, __nv_bfloat16> { @@ -236,13 +237,228 @@ __global__ void selective_state_update_kernel_simple(SelectiveStateUpdateParams } } +template +struct SharedStorage { + alignas(128) state_t state[numStages][rowsPerStage * dstate]; + input_t x[dim]; + float out[dim]; // dt is special cause we're gonna store input in there as well + input_t z[dim]; + input_t B[dstate]; + input_t C[dstate]; + + using barrier_t = cuda::barrier; + barrier_t bar_empty[numStages]; + barrier_t bar_full[numStages]; + barrier_t bar_consumers; +}; + +template +__global__ void selective_state_update_kernel_producer_consumer_vertical( + SelectiveStateUpdateParams params, __grid_constant__ CUtensorMap const tensorState) { + auto* __restrict__ output = + reinterpret_cast(params.output); // output: (batch, nheads, dim) + + auto const* __restrict__ x = + reinterpret_cast(params.x); // x: (batch, nheads, dim) + auto const* __restrict__ dt = + reinterpret_cast(params.dt); // dt: (batch, nheads, dim) + auto const* __restrict__ A = reinterpret_cast(params.A); // A: (nheads) + auto const* __restrict__ B = + reinterpret_cast(params.B); // B: (batch, ngroups, dstate) + auto const* __restrict__ C = + reinterpret_cast(params.C); // C: (batch, ngroups, dstate) + auto const* __restrict__ D = reinterpret_cast(params.D); // D: (nheads, dim) + auto const* __restrict__ dt_bias = reinterpret_cast(params.dt_bias); + auto const* __restrict__ z = reinterpret_cast(params.z); + auto const* __restrict__ state_batch_indices = + reinterpret_cast(params.state_batch_indices); + + int const nheads = params.nheads; + int const ngroups = params.ngroups; + int const dim = params.dim; + + constexpr auto warpSize = 32; + constexpr auto numWarps = 1 + consumerWarps; + + auto const batch = blockIdx.x; + auto const head = blockIdx.y; + auto const group = head / (nheads / ngroups); + auto lane = threadIdx.x % warpSize; + auto warp = threadIdx.y; + + auto const state_batch = (state_batch_indices) ? state_batch_indices[batch] : batch; + + using sram_t = + SharedStorage; +#pragma nv_diag_suppress 20054 + __shared__ sram_t sram; +#pragma nv_diag_default 20054 + + namespace cde = cuda::device::experimental; + namespace cg = cooperative_groups; + + for (int stage = warp; stage < numStages; stage += numWarps) { + if (lane > 0) continue; + constexpr auto num_arrivals = 1 + consumerWarps * warpSize; + init(&sram.bar_empty[stage], num_arrivals); + init(&sram.bar_full[stage], num_arrivals); + // signal to async proxy that barriers are initilized + cde::fence_proxy_async_shared_cta(); + } + if (lane == 0 && warp == 0) { + init(&sram.bar_consumers, warpSize * consumerWarps); + } + __syncthreads(); + + if (warp == consumerWarps) { + auto const state_offset = (state_batch * nheads + head) * dim; + + for (int d = 0, stage = 0; d < dim + rowsPerStage * numStages; + d += rowsPerStage, stage = (stage + 1) % numStages) { + if (lane == 0) { + cg::invoke_one(cg::coalesced_threads(), [&]() { + sram.bar_empty[stage].wait(sram.bar_empty[stage].arrive()); + + if (state_batch != params.pad_slot_id) { + // Writeback + if (d >= rowsPerStage * numStages) { + cde::cp_async_bulk_tensor_2d_shared_to_global( + &tensorState, + /*x*/ 0, + /*y*/ state_offset + d - rowsPerStage * numStages, &sram.state[stage][0]); + cde::cp_async_bulk_commit_group(); + cde::cp_async_bulk_wait_group_read<0>(); + } + + if (d < dim) { + cde::cp_async_bulk_tensor_2d_global_to_shared(&sram.state[stage][0], &tensorState, + /*x*/ 0, /*y*/ state_offset + d, + sram.bar_full[stage]); + + // Unblock the consumers + auto constexpr bytesState = rowsPerStage * DSTATE * sizeof(state_t); + auto constexpr bytesToArrive = bytesState; + auto const _ = + cuda::device::barrier_arrive_tx(sram.bar_full[stage], 1, bytesToArrive); + } + } else { + auto const _ = sram.bar_full[stage].arrive(); + } + }); + } + } + } else { // consumers + + using load_t = float2; + static constexpr auto vectorizedLoadSize = sizeof(load_t) / sizeof(weight_t); + +#pragma unroll + // Unblock the producer + for (uint8_t stage = 0; stage < numStages; ++stage) { + auto const _ = sram.bar_empty[stage].arrive(); + } + + // Load A + auto const A_value = toFloat(A[head]); + + // Load D + auto const d_value = D ? toFloat(D[head]) : 0.f; + + // load dt_value + auto dt_value = toFloat(dt[batch * params.dt_stride_batch + head]); + if (dt_bias) dt_value += toFloat(dt_bias[head]); + if (params.dt_softplus) { + dt_value = thresholded_softplus(dt_value); + } + + if (warp == 0) { // Load x, B + for (auto d = lane * vectorizedLoadSize; d < dim; d += warpSize * vectorizedLoadSize) { + auto* dst = reinterpret_cast(&sram.x[d]); + *dst = *reinterpret_cast(&x[batch * params.x_stride_batch + head * dim + d]); + } + for (auto i = lane * vectorizedLoadSize; i < DSTATE; i += warpSize * vectorizedLoadSize) { + auto* dst = reinterpret_cast(&sram.B[i]); + *dst = *reinterpret_cast( + &B[batch * params.B_stride_batch + group * DSTATE + i]); + } + } else if (warp == 1) { // Load z, C + for (auto d = lane * vectorizedLoadSize; d < dim; d += warpSize * vectorizedLoadSize) { + auto* dst = reinterpret_cast(&sram.z[d]); + *dst = z ? *reinterpret_cast(&z[batch * nheads * dim + head * dim + d]) + : make_zero(); + } + for (auto i = lane * vectorizedLoadSize; i < DSTATE; i += warpSize * vectorizedLoadSize) { + auto* dst = reinterpret_cast(&sram.C[i]); + *dst = *reinterpret_cast( + &C[batch * params.C_stride_batch + group * DSTATE + i]); + } + } + + sram.bar_consumers.wait(sram.bar_consumers.arrive()); + + for (auto dBegin = 0, stage = 0; dBegin < dim; + dBegin += rowsPerStage, stage = (stage + 1) % numStages) { + // wait for the producer + sram.bar_full[stage].wait(sram.bar_full[stage].arrive()); + +#pragma unroll + for (auto dd = warp; dd < rowsPerStage; dd += consumerWarps) { + auto d = dBegin + dd; + float const x_value = toFloat(sram.x[d]); + float out_value = toFloat(d_value) * x_value * int(lane == 0); // first lane has the value + + for (int i = lane; i < DSTATE; i += warpSize) { + auto const state_value = (state_batch != params.pad_slot_id) + ? toFloat(sram.state[stage][dd * DSTATE + i]) + : 0.f; + auto const B_value = toFloat(sram.B[i]); + auto const C_value = toFloat(sram.C[i]); + + auto const dA = fast_exp(A_value * dt_value); + auto const dB = B_value * dt_value; + auto const new_state = state_value * dA + dB * x_value; + + convertAndStore(&sram.state[stage][dd * DSTATE + i], new_state); + out_value += new_state * C_value; + } + + out_value = warpReduceSum(out_value); + if (lane == 0) { + sram.out[d] = out_value; + } + } + + // Unblock producer + cde::fence_proxy_async_shared_cta(); + auto _ = sram.bar_empty[stage].arrive(); + } + + // Write output + sram.bar_consumers.wait(sram.bar_consumers.arrive()); + auto d = warp * warpSize + lane; + if (d < dim) { + auto out_value = sram.out[d]; + if (z) { + float z_value = toFloat(sram.z[d]); + float sig_z = __fdividef(1.f, (1.f + __expf(0.f - z_value))); + float silu_z = z_value * sig_z; + out_value *= silu_z; + } + convertAndStore(&output[batch * params.out_stride_batch + head * dim + d], out_value); + } + } +} + template void invokeSelectiveStateUpdate(SelectiveStateUpdateParams& params, cudaStream_t stream) { - constexpr int DSTATE = 128; FLASHINFER_CHECK(params.dstate == DSTATE); - // #if (__CUDA_ARCH__ < 900) // pre-Hopper + auto [sm_major, sm_minor] = GetCudaComputeCapability(); + + if (sm_major < 9) // pre-Hopper { constexpr int numWarps = 2; int const blocks_per_dim = (params.dim + 32 * numWarps - 1) / (32 * numWarps); @@ -251,8 +467,35 @@ void invokeSelectiveStateUpdate(SelectiveStateUpdateParams& params, cudaStream_t selective_state_update_kernel_simple <<>>(params); } - // #else - // #endif + else + { + constexpr auto numConsumers = 4; + constexpr auto numWarps = 1 + numConsumers; + constexpr auto numStages = 3; + constexpr auto rowsPerStage = 4 * numConsumers; + constexpr auto DIM = 64; + FLASHINFER_CHECK(params.dim == DIM); + auto scan_func = selective_state_update_kernel_producer_consumer_vertical< + input_t, weight_t, matrixA_t, state_t, DIM, DSTATE, numConsumers, rowsPerStage, numStages>; + + dim3 block(32, numWarps); + dim3 grid(params.batch, params.nheads); + + auto nh = params.nheads; + auto dim = params.dim; + auto B = params.state_cache_size; + + FLASHINFER_CHECK(reinterpret_cast(params.state) % 128 == 0); // TMA requires 128B aligned + auto tensorState = tma::createTensorMap(params.state, B * nh * dim, DSTATE, rowsPerStage, DSTATE); + FLASHINFER_CHECK(params.dim % rowsPerStage == 0); + + scan_func<<>>(params, tensorState); + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) + { + printf("Kernel launch failed: %s\n", cudaGetErrorString(err)); + } + } } } // namespace flashinfer::mamba From 8e80802f5f66a5f32f83bdca6b072b150dd51080 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Tue, 6 Jan 2026 22:49:20 +0000 Subject: [PATCH 05/38] Passed pre-commit checks --- csrc/flashinfer_mamba_binding.cu | 7 +- csrc/selective_state_update.cu | 1 - flashinfer/jit/mamba/__init__.py | 2 +- .../jit/mamba/selective_state_update.py | 36 +--- flashinfer/mamba/__init__.py | 2 +- flashinfer/mamba/selective_state_update.py | 32 ++- .../flashinfer/mamba/create_tensor_map.cuh | 6 +- .../mamba/selective_state_update.cuh | 15 +- tests/mamba/__init__.py | 2 +- tests/mamba/selective_state_update_triton.py | 101 +++++---- tests/mamba/test_selective_state_update.py | 194 ++++-------------- 11 files changed, 142 insertions(+), 256 deletions(-) diff --git a/csrc/flashinfer_mamba_binding.cu b/csrc/flashinfer_mamba_binding.cu index 02198ce25c..347fd0fb4f 100644 --- a/csrc/flashinfer_mamba_binding.cu +++ b/csrc/flashinfer_mamba_binding.cu @@ -22,9 +22,8 @@ namespace flashinfer::mamba { void selective_state_update(TensorView state, TensorView x, TensorView dt, TensorView output, TensorView A, TensorView B, TensorView C, TensorView D, - Optional z, Optional dt_bias, - bool dt_softplus, Optional state_batch_indices, - int64_t pad_slot_id); + Optional z, Optional dt_bias, bool dt_softplus, + Optional state_batch_indices, int64_t pad_slot_id); } // namespace flashinfer::mamba @@ -32,4 +31,4 @@ void selective_state_update(TensorView state, TensorView x, TensorView dt, Tenso // This enables cross-language bindings (not just PyTorch) TVM_FFI_DLL_EXPORT_TYPED_FUNC(selective_state_update, flashinfer::mamba::selective_state_update); -// Add more mamba operations here as they are implemented \ No newline at end of file +// Add more mamba operations here as they are implemented diff --git a/csrc/selective_state_update.cu b/csrc/selective_state_update.cu index 31e5859e84..4928cfbae9 100644 --- a/csrc/selective_state_update.cu +++ b/csrc/selective_state_update.cu @@ -201,7 +201,6 @@ void selective_state_update(TensorView state, TensorView x, TensorView dt, Tenso throw std::runtime_error("selective_state_update is not implemented yet."); } - } } // namespace flashinfer::mamba diff --git a/flashinfer/jit/mamba/__init__.py b/flashinfer/jit/mamba/__init__.py index e71b2ea2f5..a17f20af35 100644 --- a/flashinfer/jit/mamba/__init__.py +++ b/flashinfer/jit/mamba/__init__.py @@ -16,4 +16,4 @@ from .selective_state_update import gen_selective_state_update_module -__all__ = ["gen_selective_state_update_module"] \ No newline at end of file +__all__ = ["gen_selective_state_update_module"] diff --git a/flashinfer/jit/mamba/selective_state_update.py b/flashinfer/jit/mamba/selective_state_update.py index 59b26e1ef5..266b50edb8 100644 --- a/flashinfer/jit/mamba/selective_state_update.py +++ b/flashinfer/jit/mamba/selective_state_update.py @@ -19,43 +19,11 @@ def gen_selective_state_update_module() -> JitSpec: - """Generate JIT module for selective_state_update operation. - - This function: - 1. Computes unique identifier from parameters (if needed) - 2. Copies source files to generation directory - 3. Returns JitSpec with compilation metadata - - Note: This is a simple example without type specialization. - If you need different code for different dtypes/parameters, - you can add Jinja template rendering here (see flashinfer/jit/sampling.py). - """ - - # TODO: If you need type specialization, add parameters like: - # def gen_selective_state_update_module(dtype_in, dtype_out, ...) -> JitSpec: - # uri = get_selective_state_update_uri(dtype_in, dtype_out, ...) - # gen_directory = jit_env.FLASHINFER_GEN_SRC_DIR / uri - # - # # Render Jinja template for type-specialized config - # with open(jit_env.FLASHINFER_CSRC_DIR / "selective_state_update_customize_config.jinja") as f: - # template = jinja2.Template(f.read()) - # config_content = template.render(dtype_in=dtype_map[dtype_in], ...) - # write_if_different(gen_directory / "selective_state_update_config.inc", config_content) - # - # # Copy source files - # sources = [] - # for fname in ["selective_state_update.cu", "flashinfer_mamba_binding.cu"]: - # shutil.copy(jit_env.FLASHINFER_CSRC_DIR / fname, gen_directory / fname) - # sources.append(gen_directory / fname) - # - # return gen_jit_spec(uri, sources, extra_cuda_cflags=[...]) - - # Simple version without type specialization (like norm.py): nvcc_flags = [ "-DENABLE_BF16", "-DENABLE_FP8", ] - + return gen_jit_spec( "mamba_selective_state_update", [ @@ -63,4 +31,4 @@ def gen_selective_state_update_module() -> JitSpec: jit_env.FLASHINFER_CSRC_DIR / "flashinfer_mamba_binding.cu", ], extra_cuda_cflags=nvcc_flags, - ) \ No newline at end of file + ) diff --git a/flashinfer/mamba/__init__.py b/flashinfer/mamba/__init__.py index a201b0ff6d..4ecfe7e667 100644 --- a/flashinfer/mamba/__init__.py +++ b/flashinfer/mamba/__init__.py @@ -16,4 +16,4 @@ from .selective_state_update import selective_state_update -__all__ = ["selective_state_update"] \ No newline at end of file +__all__ = ["selective_state_update"] diff --git a/flashinfer/mamba/selective_state_update.py b/flashinfer/mamba/selective_state_update.py index 4dbdd85b95..99cce6dc36 100644 --- a/flashinfer/mamba/selective_state_update.py +++ b/flashinfer/mamba/selective_state_update.py @@ -87,12 +87,26 @@ def selective_state_update( """ output = torch.empty_like(x) _selective_state_update( - state, x, dt, output, A, B, C, D, z, dt_bias, dt_softplus, state_batch_indices, pad_slot_id + state, + x, + dt, + output, + A, + B, + C, + D, + z, + dt_bias, + dt_softplus, + state_batch_indices, + pad_slot_id, ) return output -@register_custom_op("flashinfer::selective_state_update", mutates_args=("state", "output")) +@register_custom_op( + "flashinfer::selective_state_update", mutates_args=("state", "output") +) def _selective_state_update( state: torch.Tensor, x: torch.Tensor, @@ -110,7 +124,19 @@ def _selective_state_update( ) -> None: """Internal function registered with torch.library for torch.compile() support.""" get_selective_state_update_module().selective_state_update( - state, x, dt, output, A, B, C, D, z, dt_bias, dt_softplus, state_batch_indices, pad_slot_id + state, + x, + dt, + output, + A, + B, + C, + D, + z, + dt_bias, + dt_softplus, + state_batch_indices, + pad_slot_id, ) diff --git a/include/flashinfer/mamba/create_tensor_map.cuh b/include/flashinfer/mamba/create_tensor_map.cuh index aa2e9fef46..4208622dfa 100644 --- a/include/flashinfer/mamba/create_tensor_map.cuh +++ b/include/flashinfer/mamba/create_tensor_map.cuh @@ -9,8 +9,10 @@ #include #ifndef gpuErrchk -#define gpuErrchk(ans) \ - { gpuAssert((ans), __FILE__, __LINE__); } +#define gpuErrchk(ans) \ + { \ + gpuAssert((ans), __FILE__, __LINE__); \ + } #endif static inline void gpuAssert(cudaError_t code, const char* file, int line, bool abort = true) { diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh index c200aa8d88..f5320a72f8 100644 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -466,9 +466,7 @@ void invokeSelectiveStateUpdate(SelectiveStateUpdateParams& params, cudaStream_t dim3 grid(blocks_per_dim, params.batch, params.nheads); selective_state_update_kernel_simple <<>>(params); - } - else - { + } else { constexpr auto numConsumers = 4; constexpr auto numWarps = 1 + numConsumers; constexpr auto numStages = 3; @@ -485,15 +483,16 @@ void invokeSelectiveStateUpdate(SelectiveStateUpdateParams& params, cudaStream_t auto dim = params.dim; auto B = params.state_cache_size; - FLASHINFER_CHECK(reinterpret_cast(params.state) % 128 == 0); // TMA requires 128B aligned - auto tensorState = tma::createTensorMap(params.state, B * nh * dim, DSTATE, rowsPerStage, DSTATE); + FLASHINFER_CHECK(reinterpret_cast(params.state) % 128 == + 0); // TMA requires 128B aligned + auto tensorState = + tma::createTensorMap(params.state, B * nh * dim, DSTATE, rowsPerStage, DSTATE); FLASHINFER_CHECK(params.dim % rowsPerStage == 0); scan_func<<>>(params, tensorState); cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) - { - printf("Kernel launch failed: %s\n", cudaGetErrorString(err)); + if (err != cudaSuccess) { + printf("Kernel launch failed: %s\n", cudaGetErrorString(err)); } } } diff --git a/tests/mamba/__init__.py b/tests/mamba/__init__.py index 0a2157a5eb..2132b72828 100644 --- a/tests/mamba/__init__.py +++ b/tests/mamba/__init__.py @@ -12,4 +12,4 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the specific language governing permissions and limitations under the License. -""" \ No newline at end of file +""" diff --git a/tests/mamba/selective_state_update_triton.py b/tests/mamba/selective_state_update_triton.py index 232c748edd..cee4c0f75d 100644 --- a/tests/mamba/selective_state_update_triton.py +++ b/tests/mamba/selective_state_update_triton.py @@ -35,19 +35,22 @@ def softplus(dt): def softplus(dt): return tl.math.log1p(tl.exp(dt)) + PAD_SLOT_ID = -1 -@triton.heuristics( - {"HAS_DT_BIAS": lambda args: args["dt_bias_ptr"] is not None}) +@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"])}) + { + "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_scan_update_kernel( # Pointers to matrices @@ -128,26 +131,26 @@ def _selective_scan_update_kernel( if HAS_DT_BIAS: dt_bias_ptr += pid_h * stride_dt_bias_head A_ptr += pid_h * stride_A_head - B_ptr += pid_b * stride_B_batch + (pid_h // - nheads_ngroups_ratio) * stride_B_group - C_ptr += pid_b * stride_C_batch + (pid_h // - nheads_ngroups_ratio) * stride_C_group + B_ptr += pid_b * stride_B_batch + (pid_h // nheads_ngroups_ratio) * stride_B_group + C_ptr += pid_b * stride_C_batch + (pid_h // nheads_ngroups_ratio) * stride_C_group if HAS_Z: z_ptr += pid_b * stride_z_batch + pid_h * stride_z_head out_ptr += pid_b * stride_out_batch + pid_h * stride_out_head offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) offs_n = tl.arange(0, BLOCK_SIZE_DSTATE) - state_ptrs = state_ptr + (offs_m[:, None] * stride_state_dim + - offs_n[None, :] * stride_state_dstate) + state_ptrs = state_ptr + ( + offs_m[:, None] * stride_state_dim + offs_n[None, :] * stride_state_dstate + ) x_ptrs = x_ptr + offs_m * stride_x_dim dt_ptrs = dt_ptr + offs_m * stride_dt_dim if HAS_DT_BIAS: dt_bias_ptrs = dt_bias_ptr + offs_m * stride_dt_bias_dim if HAS_D: D_ptr += pid_h * stride_D_head - A_ptrs = A_ptr + (offs_m[:, None] * stride_A_dim + - offs_n[None, :] * stride_A_dstate) + A_ptrs = A_ptr + ( + offs_m[:, None] * stride_A_dim + offs_n[None, :] * stride_A_dstate + ) B_ptrs = B_ptr + offs_n * stride_B_dstate C_ptrs = C_ptr + offs_n * stride_C_dstate if HAS_D: @@ -157,20 +160,19 @@ def _selective_scan_update_kernel( out_ptrs = out_ptr + offs_m * stride_out_dim mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate) if HAS_STATE_BATCH_INDICES: - mask &= (state_batch_idx != pad_slot_id) + mask &= state_batch_idx != pad_slot_id state = tl.load(state_ptrs, mask=mask, 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) + 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) + 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) @@ -196,7 +198,7 @@ def _selective_scan_update_kernel( mask = (offs_m[:, None] < dim) & (offs_n[None, :] < dstate) if HAS_STATE_BATCH_INDICES: - mask &= (state_batch_idx != pad_slot_id) + mask &= state_batch_idx != pad_slot_id tl.store(state_ptrs, state, mask=mask) out = tl.sum(state * C[None, :], axis=1) if HAS_D: @@ -206,18 +208,20 @@ def _selective_scan_update_kernel( tl.store(out_ptrs, out, mask=offs_m < dim) -def selective_state_update_triton(state, - x, - dt, - A, - B, - C, - D=None, - z=None, - dt_bias=None, - dt_softplus=False, - state_batch_indices=None, - pad_slot_id=PAD_SLOT_ID): +def selective_state_update_triton( + state, + x, + dt, + A, + B, + C, + D=None, + z=None, + dt_bias=None, + dt_softplus=False, + state_batch_indices=None, + pad_slot_id=PAD_SLOT_ID, +): """ Argument: state: (batch, dim, dstate) or (batch, nheads, dim, dstate) @@ -275,19 +279,27 @@ def selective_state_update_triton(state, if dt_bias is not None: assert dt_bias.shape == (nheads, dim) if state_batch_indices is not None: - assert state_batch_indices.shape == (batch, ) + assert state_batch_indices.shape == (batch,) out = torch.empty_like(x) grid = lambda META: (triton.cdiv(dim, META["BLOCK_SIZE_M"]), batch, nheads) - z_strides = (z.stride(0), z.stride(1), - z.stride(2)) if z is not None else (0, 0, 0) + z_strides = (z.stride(0), z.stride(1), z.stride(2)) if z is not None else (0, 0, 0) # We don't want autotune since it will overwrite the state # We instead tune by hand. - BLOCK_SIZE_M, num_warps = ((32, 4) if dstate <= 16 else - ((16, 4) if dstate <= 32 else - ((8, 4) if dstate <= 64 else - ((4, 4) if dstate <= 128 else ((4, 8)))))) - tie_hdim = (A.stride(-1) == 0 and A.stride(-2) == 0 and dt.stride(-1) == 0 - and dt_bias.stride(-1) == 0) + BLOCK_SIZE_M, num_warps = ( + (32, 4) + if dstate <= 16 + else ( + (16, 4) + if dstate <= 32 + else ((8, 4) if dstate <= 64 else ((4, 4) if dstate <= 128 else ((4, 8)))) + ) + ) + tie_hdim = ( + A.stride(-1) == 0 + and A.stride(-2) == 0 + and dt.stride(-1) == 0 + and dt_bias.stride(-1) == 0 + ) with torch.cuda.device(x.device.index): _selective_scan_update_kernel[grid]( state, @@ -317,8 +329,7 @@ def selective_state_update_triton(state, dt.stride(0), dt.stride(1), dt.stride(2), - *(dt_bias.stride(0), - dt_bias.stride(1)) if dt_bias is not None else 0, + *(dt_bias.stride(0), dt_bias.stride(1)) if dt_bias is not None else 0, A.stride(0), A.stride(1), A.stride(2), diff --git a/tests/mamba/test_selective_state_update.py b/tests/mamba/test_selective_state_update.py index e675b99726..0ec381fa10 100644 --- a/tests/mamba/test_selective_state_update.py +++ b/tests/mamba/test_selective_state_update.py @@ -1,86 +1,10 @@ import numpy as np import pytest import torch -import torch.nn.functional as F -from einops import rearrange, repeat -from .selective_state_update_triton import selective_state_update_triton import flashinfer - -def selective_state_update_ref( - state, x, dt, A, B, C, D=None, z=None, dt_bias=None, dt_softplus=False -): - """ - Argument: - state: (batch, dstate, dim) or (batch, nheads, dstate, dim) - x: (batch, dim) or (batch, nheads, dim) - dt: (batch, dim) or (batch, nheads, dim) - A: (dstate, dim) or (nheads, dstate, dim) - B: (batch, dstate) or (batch, ngroups, dstate) - C: (batch, dstate) or (batch, ngroups, dstate) - D: (dim,) or (nheads, dim) - z: (batch, dim) or (batch, nheads, dim) - dt_bias: (dim,) or (nheads, dim) - Return: - out: (batch, dim) or (batch, nheads, dim) - """ - 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) - batch, nheads, dstate, dim = state.shape - - assert x.shape == (batch, nheads, dim) - assert dt.shape == x.shape - assert A.shape == (nheads, dstate, dim) - ngroups = B.shape[1] - assert nheads % ngroups == 0, "nheads must be divisible by ngroups" - assert B.shape == (batch, ngroups, dstate) - assert C.shape == B.shape - - if D is not None: - assert D.shape == (nheads, dim) - if z is not None: - assert z.shape == x.shape - if dt_bias is not None: - assert dt_bias.shape == (nheads, dim) - dt = dt + dt_bias - dt = F.softplus(dt) if dt_softplus else dt - dA = torch.exp( - rearrange(dt, "b h d -> b h 1 d") * A - ) # (batch, nheads, dstate, dim) - B = repeat(B, "b g n -> b (g h) n", h=nheads // ngroups) # (batch, nheads, dstate) - C = repeat(C, "b g n -> b (g h) n", h=nheads // ngroups) # (batch, nheads, dstate) - dB = rearrange(dt, "b h d -> b h 1 d") * rearrange( - B.float(), "b h n -> b h n 1" - ) # (batch, nheads, dstate, dim) - state_new = state.float() * dA + dB * rearrange( - x.float(), "b h d -> b h 1 d" - ) # (batch, nheads, dstate, dim) - state.copy_(state_new.to(state.dtype)) - out = torch.einsum("bhnd,bhn->bhd", state_new, C.float()) - if D is not None: - out += x.float() * D - out = (out if z is None else out * F.silu(z.float())).to(x.dtype) - if not has_heads: - out = out.squeeze(1) - return out +from .selective_state_update_triton import selective_state_update_triton def create_test_inputs( @@ -108,12 +32,7 @@ def create_test_inputs( A_base = -torch.rand(nheads, dtype=torch.float32, device=device) A = A_base.as_strided((nheads, dim, dstate), (1, 0, 0)) - - # A = -torch.rand(nheads, dtype=torch.float32, device=device).as_strided( - # (nheads, dim, dstate), (1, 0, 0) - # ) - # print(f"A dtype: {A.dtype}, shape: {A.shape}, stride {A.stride()}") - assert(A.stride() == (1, 0, 0)) + assert A.stride() == (1, 0, 0) # B and C - (batch_size, ngroups, dstate) B = torch.randn(batch_size, ngroups, dstate, dtype=input_dtype, device=device) @@ -145,17 +64,9 @@ def create_test_inputs( "slot_idx": slot_idx, } -# @pytest.mark.parametrize("batch", [1, 32, 64]) -# @pytest.mark.parametrize("nheads", [8, 64]) -# @pytest.mark.parametrize("dim", [64]) -# @pytest.mark.parametrize("dstate", [128]) -# @pytest.mark.parametrize("ngroups", [8]) -# @pytest.mark.parametrize("delta_softplus", [True]) -# @pytest.mark.parametrize("input_dtype", [torch.bfloat16]) -# @pytest.mark.parametrize("weight_dtype", [torch.bfloat16]) -# @pytest.mark.parametrize("matrixA_dtype", [torch.float32]) -@pytest.mark.parametrize("batch", [1]) -@pytest.mark.parametrize("nheads", [64]) + +@pytest.mark.parametrize("batch", [1, 32, 64]) +@pytest.mark.parametrize("nheads", [8, 64]) @pytest.mark.parametrize("dim", [64]) @pytest.mark.parametrize("dstate", [128]) @pytest.mark.parametrize("ngroups", [8]) @@ -175,8 +86,6 @@ def test_selective_state_update( matrixA_dtype, ): """Test selective_state_update correctness against reference implementation.""" - # TODO: Create input tensors based on your kernel requirements - # Example: inputs = create_test_inputs( batch, nheads, @@ -188,41 +97,24 @@ def test_selective_state_update( matrixA_dtype, ) - # state = inputs["state_cache"][inputs["slot_idx"]] - # state_ref = rearrange(state, "... p n -> ... n p").detach().clone() state = inputs["state_cache"] state_ref = state.clone() - # A = inputs["A"] - - # A_ref = (rearrange(A, "... p n -> ... n p") if A.ndim == 3 else A).detach().clone() - # y_ref = selective_state_update_ref( - # state_ref, - # inputs["x"], - # inputs["dt"], - # A_ref, - # inputs["B"], - # inputs["C"], - # D=inputs["D"], - # z=None, - # dt_bias=inputs["dt_bias"], - # dt_softplus=delta_softplus, - # ) y_ref = selective_state_update_triton( - state_ref, - inputs["x"], - inputs["dt"], - inputs["A"], - inputs["B"], - inputs["C"], - D=inputs["D"], - z=None, - dt_bias=inputs["dt_bias"], - dt_softplus=delta_softplus, - state_batch_indices=inputs["slot_idx"], - pad_slot_id=-1, - ) + state_ref, + inputs["x"], + inputs["dt"], + inputs["A"], + inputs["B"], + inputs["C"], + D=inputs["D"], + z=None, + dt_bias=inputs["dt_bias"], + dt_softplus=delta_softplus, + state_batch_indices=inputs["slot_idx"], + pad_slot_id=-1, + ) - y_test = flashinfer.mamba.selective_state_update( + y_test = flashinfer.mamba.selective_state_update( state, inputs["x"], inputs["dt"], @@ -237,15 +129,14 @@ def test_selective_state_update( pad_slot_id=-1, ) - atol=1e-3 - rtol=1e-2 + atol = 1e-3 + rtol = 1e-2 outputs_match = torch.allclose(y_ref, y_test, atol=atol, rtol=rtol) if outputs_match: print(f"✓ Outputs match within tolerance (atol={atol}, rtol={rtol})") else: print(f"✗ Outputs do NOT match within tolerance (atol={atol}, rtol={rtol})") - all_passed = False # Detailed comparison using numpy testing y_ref_np = y_ref.detach().cpu().float().numpy() @@ -262,9 +153,9 @@ def test_selective_state_update( ) mismatch_indices = np.argwhere(mismatch_mask) - print(f"First few mismatch locations (up to 10):") - for i, idx in enumerate(mismatch_indices[:10]): - idx_tuple = tuple(idx) + print("First few mismatch locations (up to 10):") + for idx in mismatch_indices[:10]: + idx_tuple = tuple(int(i) for i in idx) ref_val = y_ref_np[idx_tuple] test_val = y_test_np[idx_tuple] diff = abs(ref_val - test_val) @@ -273,32 +164,24 @@ def test_selective_state_update( f" Index {idx_tuple}: ref={ref_val:.6f}, test={test_val:.6f}, diff={diff:.6e}, rel_diff={rel_diff:.6e}" ) - assert(outputs_match) - - # Compare states (updated in-place) - print("\nComparing states with reference...") - - - state_test = state[inputs["slot_idx"]] - state_diff = torch.abs(state_ref - state_test) - max_state_diff = state_diff.max().item() - mean_state_diff = state_diff.mean().item() - - print(f"Max absolute state difference: {max_state_diff:.6e}") - print(f"Mean absolute state difference: {mean_state_diff:.6e}") + assert outputs_match # Check if states match within tolerance - states_match = torch.allclose( state_ref, state_test, atol=atol, rtol=rtol ) + states_match = torch.allclose( + state_ref[inputs["slot_idx"]], + state[inputs["slot_idx"]], + atol=atol, + rtol=rtol, + ) if states_match: print(f"✓ States match within tolerance (atol={atol}, rtol={rtol})") - return True else: print(f"✗ States do NOT match within tolerance (atol={atol}, rtol={rtol})") # Detailed comparison using numpy testing - state_ref_np = state_ref.detach().cpu().float().numpy() - state_test_np = state_test.detach().cpu().float().numpy() + state_ref_np = state_ref[inputs["slot_idx"]].detach().cpu().float().numpy() + state_test_np = state[inputs["slot_idx"]].detach().cpu().float().numpy() print("\nDetailed state mismatch analysis:") state_mismatch_mask = ~np.isclose( @@ -312,16 +195,15 @@ def test_selective_state_update( ) state_mismatch_indices = np.argwhere(state_mismatch_mask) - print(f"First few state mismatch locations (up to 10):") - for i, idx in enumerate(state_mismatch_indices[:10]): - idx_tuple = tuple(idx) + print("First few state mismatch locations (up to 10):") + for idx in state_mismatch_indices[:10]: + idx_tuple = tuple(int(i) for i in idx) ref_val = state_ref_np[idx_tuple] test_val = state_test_np[idx_tuple] diff = abs(ref_val - test_val) rel_diff = diff / (abs(ref_val) + 1e-8) print( - f" Index {idx_tuple}: ref={ref_val:.6f}, test={test_val:.6f}, diff={diff:.6e}, rel_diff={rel_diff:.6e}" + f" Index{idx_tuple}: ref={ref_val:.6f}, test={test_val:.6f}, diff={diff:.6e}, rel_diff={rel_diff:.6e}" ) - return False - assert(states_match) + assert states_match From 8568d67321a226642a131d3e4aee39bc59dfd188 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 09:47:53 -0800 Subject: [PATCH 06/38] Update flashinfer/mamba/selective_state_update.py make input tensors always have batch and nheads as dimensions. Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- flashinfer/mamba/selective_state_update.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/flashinfer/mamba/selective_state_update.py b/flashinfer/mamba/selective_state_update.py index 99cce6dc36..ef7d0ad653 100644 --- a/flashinfer/mamba/selective_state_update.py +++ b/flashinfer/mamba/selective_state_update.py @@ -85,6 +85,24 @@ def selective_state_update( output : torch.Tensor Output tensor with shape (batch, dim) or (batch, nheads, dim) """ + 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) output = torch.empty_like(x) _selective_state_update( state, From 7bb48ccaa1aeff09a8f5e15e04254608c1457e3f Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 10:43:29 -0800 Subject: [PATCH 07/38] remove unreachable code --- csrc/selective_state_update.cu | 7 ------- 1 file changed, 7 deletions(-) diff --git a/csrc/selective_state_update.cu b/csrc/selective_state_update.cu index 4928cfbae9..6032449625 100644 --- a/csrc/selective_state_update.cu +++ b/csrc/selective_state_update.cu @@ -22,11 +22,6 @@ using tvm::ffi::Optional; namespace flashinfer::mamba { -// TODO: Implement the launcher function that: -// 1. Validates tensor dimensions and devices -// 2. Extracts raw pointers from TensorView -// 3. Calls the framework-agnostic kernel from include/flashinfer/mamba/selective_state_update.cuh -// void selective_state_update(TensorView state, TensorView x, TensorView dt, TensorView output, TensorView A, TensorView B, TensorView C, TensorView D, Optional z, Optional dt_bias, bool dt_softplus, @@ -198,8 +193,6 @@ void selective_state_update(TensorView state, TensorView x, TensorView dt, Tenso << "matrixA_dtype=" << matrixA_dtype.code << ":" << matrixA_dtype.bits << ". Currently only support: " << "state=bfloat16, input=bfloat16, weight=bfloat16, matrixA=float32"; - - throw std::runtime_error("selective_state_update is not implemented yet."); } } From 5812938ed6407dd79495efd4d75ebe7eb432aaf3 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 10:49:25 -0800 Subject: [PATCH 08/38] Improve docstring for input state shape Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> --- flashinfer/mamba/selective_state_update.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flashinfer/mamba/selective_state_update.py b/flashinfer/mamba/selective_state_update.py index ef7d0ad653..541eaff596 100644 --- a/flashinfer/mamba/selective_state_update.py +++ b/flashinfer/mamba/selective_state_update.py @@ -54,7 +54,7 @@ def selective_state_update( Parameters ---------- state : torch.Tensor - State tensor with shape (batch, dim, dstate) or (batch, nheads, dim, dstate) + State tensor with shape (state_cache_size, dim, dstate) or (state_cache_size, nheads, dim, dstate) x : torch.Tensor Input tensor with shape (batch, dim) or (batch, nheads, dim) dt : torch.Tensor From 9883eab7e27cb61568c8f1e62b1fd519e717e4e7 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 11:06:07 -0800 Subject: [PATCH 09/38] Simple kernel also uses fast_exp --- include/flashinfer/mamba/selective_state_update.cuh | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) mode change 100644 => 100755 include/flashinfer/mamba/selective_state_update.cuh diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh old mode 100644 new mode 100755 index f5320a72f8..84ea429615 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -161,9 +161,7 @@ __global__ void selective_state_update_kernel_simple(SelectiveStateUpdateParams dt_value = thresholded_softplus(dt_value); } - // auto const dA = fast_exp(A_value * dt_value); - auto const dA = __expf(A_value * dt_value); - // auto const dA = exp(A_value * dt_value); + auto const dA = fast_exp(A_value * dt_value); auto d_value = D ? toFloat(D[head]) : 0.f; From 8aed46837ff955bf519f88440e02f682d5e89e5c Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 11:09:55 -0800 Subject: [PATCH 10/38] Remove error check of the kernel launch (it was a debugging leftover). --- include/flashinfer/mamba/selective_state_update.cuh | 4 ---- 1 file changed, 4 deletions(-) diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh index 84ea429615..b8c6f1479b 100755 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -488,10 +488,6 @@ void invokeSelectiveStateUpdate(SelectiveStateUpdateParams& params, cudaStream_t FLASHINFER_CHECK(params.dim % rowsPerStage == 0); scan_func<<>>(params, tensorState); - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - printf("Kernel launch failed: %s\n", cudaGetErrorString(err)); - } } } From 5c65e532a91b4ad8464501affccbf172f8b44a99 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 13:32:03 -0800 Subject: [PATCH 11/38] no need for the fast_exp function --- include/flashinfer/mamba/selective_state_update.cuh | 12 ++---------- 1 file changed, 2 insertions(+), 10 deletions(-) diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh index b8c6f1479b..0b68a572ff 100755 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -57,14 +57,6 @@ struct SelectiveStateUpdateParams { __forceinline__ __device__ float softplus(float x) { return __logf(1.f + __expf(x)); - // return log1pf(exp(x)); -} - -__device__ __forceinline__ float fast_exp(float x) { - constexpr float log2_E = 1.4426950408889634f; - float y; - asm("ex2.approx.f32 %0,%1;" : "=f"(y) : "f"(x * log2_E)); - return y; } __device__ __forceinline__ float thresholded_softplus(float dt_value) { @@ -161,7 +153,7 @@ __global__ void selective_state_update_kernel_simple(SelectiveStateUpdateParams dt_value = thresholded_softplus(dt_value); } - auto const dA = fast_exp(A_value * dt_value); + auto const dA = __expf(A_value * dt_value); auto d_value = D ? toFloat(D[head]) : 0.f; @@ -414,7 +406,7 @@ __global__ void selective_state_update_kernel_producer_consumer_vertical( auto const B_value = toFloat(sram.B[i]); auto const C_value = toFloat(sram.C[i]); - auto const dA = fast_exp(A_value * dt_value); + auto const dA = __expf(A_value * dt_value); auto const dB = B_value * dt_value; auto const new_state = state_value * dA + dB * x_value; From 681a5bac4963295b95a3e8bbb8fcd69ddaa01975 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 13:40:36 -0800 Subject: [PATCH 12/38] Remove unnecessary None check for D before unsqueeze --- flashinfer/mamba/selective_state_update.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/flashinfer/mamba/selective_state_update.py b/flashinfer/mamba/selective_state_update.py index 541eaff596..d8853061a5 100644 --- a/flashinfer/mamba/selective_state_update.py +++ b/flashinfer/mamba/selective_state_update.py @@ -97,7 +97,7 @@ def selective_state_update( B = B.unsqueeze(1) if C.dim() == 2: C = C.unsqueeze(1) - if D is not None and D.dim() == 1: + if D.dim() == 1: D = D.unsqueeze(0) if z is not None and z.dim() == 2: z = z.unsqueeze(1) From 5e8536e2a41b395e0ab9b01ed6aebd755de34877 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 13:46:26 -0800 Subject: [PATCH 13/38] Hoist dA computation outside the innermost loop. --- include/flashinfer/mamba/selective_state_update.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh index 0b68a572ff..2675a159f8 100755 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -362,6 +362,7 @@ __global__ void selective_state_update_kernel_producer_consumer_vertical( if (params.dt_softplus) { dt_value = thresholded_softplus(dt_value); } + auto const dA = __expf(A_value * dt_value); if (warp == 0) { // Load x, B for (auto d = lane * vectorizedLoadSize; d < dim; d += warpSize * vectorizedLoadSize) { @@ -406,7 +407,6 @@ __global__ void selective_state_update_kernel_producer_consumer_vertical( auto const B_value = toFloat(sram.B[i]); auto const C_value = toFloat(sram.C[i]); - auto const dA = __expf(A_value * dt_value); auto const dB = B_value * dt_value; auto const new_state = state_value * dA + dB * x_value; From acab6bd5f13fd2994903545dd45933f0f7f03358 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 14:09:26 -0800 Subject: [PATCH 14/38] support for non-none z --- csrc/selective_state_update.cu | 14 ++- .../mamba/selective_state_update.cuh | 6 +- tests/mamba/test_selective_state_update.py | 93 ++++++++++++++++++- 3 files changed, 105 insertions(+), 8 deletions(-) diff --git a/csrc/selective_state_update.cu b/csrc/selective_state_update.cu index 6032449625..d422064162 100644 --- a/csrc/selective_state_update.cu +++ b/csrc/selective_state_update.cu @@ -111,8 +111,17 @@ void selective_state_update(TensorView state, TensorView x, TensorView dt, Tenso bias.stride(1)); } - // if(z.has_value()) - FLASHINFER_CHECK(!z.has_value() && "z is not supported yet"); + if (z.has_value()) { + auto& z_tensor = z.value(); + CHECK_CUDA(z_tensor); + CHECK_DIM(3, z_tensor); // z: {batch, nheads, dim} + FLASHINFER_CHECK(z_tensor.size(0) == batch, "z.size(0) must equal batch"); + FLASHINFER_CHECK(z_tensor.size(1) == nheads, "z.size(1) must equal nheads"); + FLASHINFER_CHECK(z_tensor.size(2) == dim, "z.size(2) must equal dim"); + CHECK_LAST_DIM_CONTIGUOUS_INPUT(z_tensor); + FLASHINFER_CHECK(z_tensor.stride(1) == dim, "z.stride(1) must equal dim, got ", + z_tensor.stride(1), " expected ", z_tensor.size(2)); + } if (state_batch_indices) { CHECK_DIM(1, (*state_batch_indices)); @@ -150,6 +159,7 @@ void selective_state_update(TensorView state, TensorView x, TensorView dt, Tenso } if (z) { p.z = z.value().data_ptr(); + p.z_stride_batch = z.value().stride(0); } p.A = A.data_ptr(); p.B = B.data_ptr(); diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh index 2675a159f8..793078d438 100755 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -39,7 +39,7 @@ struct SelectiveStateUpdateParams { bool dt_softplus{false}; int64_t x_stride_batch{}, dt_stride_batch{}, B_stride_batch{}, C_stride_batch{}, - out_stride_batch{}; + out_stride_batch{}, z_stride_batch{}; void* __restrict__ state{nullptr}; // state_t: (state_cache_size, nheads, dim, dstate) void* __restrict__ x{nullptr}; // input_t: (batch, nheads, dim) @@ -162,7 +162,7 @@ __global__ void selective_state_update_kernel_simple(SelectiveStateUpdateParams if (d < dim) { sx[_d] = x[batch * params.x_stride_batch + head * dim + d]; if (z) { - sz[_d] = z[batch * nheads * dim + head * dim + d]; + sz[_d] = z[batch * params.z_stride_batch + head * dim + d]; } else { convertAndStore(&sz[_d], 0.f); } @@ -377,7 +377,7 @@ __global__ void selective_state_update_kernel_producer_consumer_vertical( } else if (warp == 1) { // Load z, C for (auto d = lane * vectorizedLoadSize; d < dim; d += warpSize * vectorizedLoadSize) { auto* dst = reinterpret_cast(&sram.z[d]); - *dst = z ? *reinterpret_cast(&z[batch * nheads * dim + head * dim + d]) + *dst = z ? *reinterpret_cast(&z[batch * params.z_stride_batch + head * dim + d]) : make_zero(); } for (auto i = lane * vectorizedLoadSize; i < DSTATE; i += warpSize * vectorizedLoadSize) { diff --git a/tests/mamba/test_selective_state_update.py b/tests/mamba/test_selective_state_update.py index 0ec381fa10..ff1ff4b6df 100644 --- a/tests/mamba/test_selective_state_update.py +++ b/tests/mamba/test_selective_state_update.py @@ -8,7 +8,7 @@ def create_test_inputs( - batch_size, nheads, dim, dstate, ngroups, input_dtype, weight_dtype, matrixA_dtype + batch_size, nheads, dim, dstate, ngroups, input_dtype, weight_dtype, matrixA_dtype, z_none=True ): # Set seed for reproducibility torch.manual_seed(0) @@ -52,6 +52,9 @@ def create_test_inputs( :batch_size ] + # Create z tensor if z_none is False + z = None if z_none else torch.randn(batch_size, nheads, dim, dtype=input_dtype, device=device) + return { "state_cache": state_cache, "x": x, @@ -95,6 +98,7 @@ def test_selective_state_update( input_dtype, weight_dtype, matrixA_dtype, + z_none=True, ) state = inputs["state_cache"] @@ -107,7 +111,7 @@ def test_selective_state_update( inputs["B"], inputs["C"], D=inputs["D"], - z=None, + z=inputs["z"], dt_bias=inputs["dt_bias"], dt_softplus=delta_softplus, state_batch_indices=inputs["slot_idx"], @@ -122,7 +126,7 @@ def test_selective_state_update( inputs["B"], inputs["C"], D=inputs["D"], - z=None, + z=inputs["z"], dt_bias=inputs["dt_bias"], dt_softplus=delta_softplus, state_batch_indices=inputs["slot_idx"], @@ -207,3 +211,86 @@ def test_selective_state_update( ) assert states_match + + +def test_selective_state_update_with_z(): + """Test selective_state_update with z tensor (not None).""" + batch = 1 + nheads = 8 + dim = 64 + dstate = 128 + ngroups = 8 + delta_softplus = True + input_dtype = torch.bfloat16 + weight_dtype = torch.bfloat16 + matrixA_dtype = torch.float32 + + inputs = create_test_inputs( + batch, + nheads, + dim, + dstate, + ngroups, + input_dtype, + weight_dtype, + matrixA_dtype, + z_none=False, + ) + + state = inputs["state_cache"] + state_ref = state.clone() + y_ref = selective_state_update_triton( + state_ref, + inputs["x"], + inputs["dt"], + inputs["A"], + inputs["B"], + inputs["C"], + D=inputs["D"], + z=inputs["z"], + dt_bias=inputs["dt_bias"], + dt_softplus=delta_softplus, + state_batch_indices=inputs["slot_idx"], + pad_slot_id=-1, + ) + + y_test = flashinfer.mamba.selective_state_update( + state, + inputs["x"], + inputs["dt"], + inputs["A"], + inputs["B"], + inputs["C"], + D=inputs["D"], + z=inputs["z"], + dt_bias=inputs["dt_bias"], + dt_softplus=delta_softplus, + state_batch_indices=inputs["slot_idx"], + pad_slot_id=-1, + ) + + atol = 1e-3 + rtol = 1e-2 + outputs_match = torch.allclose(y_ref, y_test, atol=atol, rtol=rtol) + + if outputs_match: + print(f"✓ Outputs match within tolerance (atol={atol}, rtol={rtol})") + else: + print(f"✗ Outputs do NOT match within tolerance (atol={atol}, rtol={rtol})") + + assert outputs_match + + # Check if states match within tolerance + states_match = torch.allclose( + state_ref[inputs["slot_idx"]], + state[inputs["slot_idx"]], + atol=atol, + rtol=rtol, + ) + + if states_match: + print(f"✓ States match within tolerance (atol={atol}, rtol={rtol})") + else: + print(f"✗ States do NOT match within tolerance (atol={atol}, rtol={rtol})") + + assert states_match From 84bdc07a05ed2ca883c3e75d6d1abc29c8c63876 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 14:15:21 -0800 Subject: [PATCH 15/38] stage forgotten z handling in the test --- tests/mamba/test_selective_state_update.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/mamba/test_selective_state_update.py b/tests/mamba/test_selective_state_update.py index ff1ff4b6df..47ff214169 100644 --- a/tests/mamba/test_selective_state_update.py +++ b/tests/mamba/test_selective_state_update.py @@ -224,7 +224,7 @@ def test_selective_state_update_with_z(): input_dtype = torch.bfloat16 weight_dtype = torch.bfloat16 matrixA_dtype = torch.float32 - + inputs = create_test_inputs( batch, nheads, From edb52692d0ea41f6d00a5a1c0f36c5065dba2978 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 14:17:00 -0800 Subject: [PATCH 16/38] test: handle z --- tests/mamba/test_selective_state_update.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/mamba/test_selective_state_update.py b/tests/mamba/test_selective_state_update.py index 47ff214169..a040b945dc 100644 --- a/tests/mamba/test_selective_state_update.py +++ b/tests/mamba/test_selective_state_update.py @@ -63,6 +63,7 @@ def create_test_inputs( "B": B, "C": C, "D": D, + "z": z, "dt_bias": dt_bias, "slot_idx": slot_idx, } From 3e9735474ad7eacfe8f57ff2a8cdfcd07d2dc71d Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 14:34:04 -0800 Subject: [PATCH 17/38] make sure that batch size does not exceed the state cache size --- csrc/selective_state_update.cu | 2 ++ 1 file changed, 2 insertions(+) diff --git a/csrc/selective_state_update.cu b/csrc/selective_state_update.cu index d422064162..a5e58df924 100644 --- a/csrc/selective_state_update.cu +++ b/csrc/selective_state_update.cu @@ -33,6 +33,8 @@ void selective_state_update(TensorView state, TensorView x, TensorView dt, Tenso auto const dstate = state.size(3); auto const ngroups = B.size(1); + FLASHINFER_CHECK(state_cache_size >= batch, "state.size(0) must be >= x.size(0)"); + FLASHINFER_CHECK(nheads % ngroups == 0, "nheads must be divisible by ngroups"); // Check x shape and strides From 7ba913dc7894b7cb2b4d568da474d378136d4764 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 14:38:18 -0800 Subject: [PATCH 18/38] do use matrixA_dtype in the test --- tests/mamba/test_selective_state_update.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/mamba/test_selective_state_update.py b/tests/mamba/test_selective_state_update.py index a040b945dc..deb76d2a39 100644 --- a/tests/mamba/test_selective_state_update.py +++ b/tests/mamba/test_selective_state_update.py @@ -30,7 +30,7 @@ def create_test_inputs( (batch_size, nheads, dim), (nheads, 1, 0) ) - A_base = -torch.rand(nheads, dtype=torch.float32, device=device) + A_base = -torch.rand(nheads, dtype=matrixA_dtype, device=device) A = A_base.as_strided((nheads, dim, dstate), (1, 0, 0)) assert A.stride() == (1, 0, 0) From 83c5ecb2dca7df214e9e179a99001eef72d08325 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 14:49:27 -0800 Subject: [PATCH 19/38] Add selective state update module to JIT specs --- flashinfer/aot.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/flashinfer/aot.py b/flashinfer/aot.py index 34096af940..d7aebdf117 100644 --- a/flashinfer/aot.py +++ b/flashinfer/aot.py @@ -62,6 +62,7 @@ ) from .jit.spdlog import gen_spdlog_module from .jit.mla import gen_mla_module +from .jit.mamba import gen_selective_state_update_module from .jit.norm import gen_norm_module from .jit.page import gen_page_module from .jit.quantization import gen_quantization_module @@ -533,6 +534,7 @@ def gen_all_modules( gen_sampling_module(), gen_topk_module(), ] + jit_specs.append(gen_selective_state_update_module()) if has_sm90: jit_specs.append(gen_trtllm_utils_module()) From c0f96aae79434130fa9136b894b0f65147f95673 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 15:01:27 -0800 Subject: [PATCH 20/38] Fix dt_bias stride check to handle None value --- tests/mamba/selective_state_update_triton.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/mamba/selective_state_update_triton.py b/tests/mamba/selective_state_update_triton.py index cee4c0f75d..f27886212b 100644 --- a/tests/mamba/selective_state_update_triton.py +++ b/tests/mamba/selective_state_update_triton.py @@ -298,7 +298,7 @@ def selective_state_update_triton( A.stride(-1) == 0 and A.stride(-2) == 0 and dt.stride(-1) == 0 - and dt_bias.stride(-1) == 0 + and (dt_bias is not None and dt_bias.stride(-1) == 0) ) with torch.cuda.device(x.device.index): _selective_scan_update_kernel[grid]( From 03a8570e6ab318ee914300a0c69a18cf606d7b4d Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 15:21:36 -0800 Subject: [PATCH 21/38] handle a few different DIM and DSTATE --- .../mamba/selective_state_update.cuh | 135 +++++++++++------- 1 file changed, 85 insertions(+), 50 deletions(-) diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh index 793078d438..2b72115512 100755 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -55,9 +55,7 @@ struct SelectiveStateUpdateParams { void* __restrict__ state_batch_indices{nullptr}; // state_batch_indices: (batch,) }; -__forceinline__ __device__ float softplus(float x) { - return __logf(1.f + __expf(x)); -} +__forceinline__ __device__ float softplus(float x) { return __logf(1.f + __expf(x)); } __device__ __forceinline__ float thresholded_softplus(float dt_value) { constexpr float threshold = 20.f; @@ -247,19 +245,14 @@ template __global__ void selective_state_update_kernel_producer_consumer_vertical( SelectiveStateUpdateParams params, __grid_constant__ CUtensorMap const tensorState) { - auto* __restrict__ output = - reinterpret_cast(params.output); // output: (batch, nheads, dim) - - auto const* __restrict__ x = - reinterpret_cast(params.x); // x: (batch, nheads, dim) - auto const* __restrict__ dt = - reinterpret_cast(params.dt); // dt: (batch, nheads, dim) - auto const* __restrict__ A = reinterpret_cast(params.A); // A: (nheads) - auto const* __restrict__ B = - reinterpret_cast(params.B); // B: (batch, ngroups, dstate) - auto const* __restrict__ C = - reinterpret_cast(params.C); // C: (batch, ngroups, dstate) - auto const* __restrict__ D = reinterpret_cast(params.D); // D: (nheads, dim) + auto* __restrict__ output = reinterpret_cast(params.output); + + auto const* __restrict__ x = reinterpret_cast(params.x); + auto const* __restrict__ dt = reinterpret_cast(params.dt); + auto const* __restrict__ A = reinterpret_cast(params.A); + auto const* __restrict__ B = reinterpret_cast(params.B); + auto const* __restrict__ C = reinterpret_cast(params.C); + auto const* __restrict__ D = reinterpret_cast(params.D); auto const* __restrict__ dt_bias = reinterpret_cast(params.dt_bias); auto const* __restrict__ z = reinterpret_cast(params.z); auto const* __restrict__ state_batch_indices = @@ -377,8 +370,9 @@ __global__ void selective_state_update_kernel_producer_consumer_vertical( } else if (warp == 1) { // Load z, C for (auto d = lane * vectorizedLoadSize; d < dim; d += warpSize * vectorizedLoadSize) { auto* dst = reinterpret_cast(&sram.z[d]); - *dst = z ? *reinterpret_cast(&z[batch * params.z_stride_batch + head * dim + d]) - : make_zero(); + *dst = + z ? *reinterpret_cast(&z[batch * params.z_stride_batch + head * dim + d]) + : make_zero(); } for (auto i = lane * vectorizedLoadSize; i < DSTATE; i += warpSize * vectorizedLoadSize) { auto* dst = reinterpret_cast(&sram.C[i]); @@ -443,43 +437,84 @@ __global__ void selective_state_update_kernel_producer_consumer_vertical( template void invokeSelectiveStateUpdate(SelectiveStateUpdateParams& params, cudaStream_t stream) { - constexpr int DSTATE = 128; - FLASHINFER_CHECK(params.dstate == DSTATE); - auto [sm_major, sm_minor] = GetCudaComputeCapability(); if (sm_major < 9) // pre-Hopper { - constexpr int numWarps = 2; - int const blocks_per_dim = (params.dim + 32 * numWarps - 1) / (32 * numWarps); - dim3 block(32, numWarps); - dim3 grid(blocks_per_dim, params.batch, params.nheads); - selective_state_update_kernel_simple - <<>>(params); + auto dispatch_dstate = [&]() { + constexpr int numWarps = 2; + int const blocks_per_dim = (params.dim + 32 * numWarps - 1) / (32 * numWarps); + dim3 block(32, numWarps); + dim3 grid(blocks_per_dim, params.batch, params.nheads); + selective_state_update_kernel_simple + <<>>(params); + }; + + switch (params.dstate) { + case 64: + dispatch_dstate.template operator()<64>(); + break; + case 128: + dispatch_dstate.template operator()<128>(); + break; + case 256: + dispatch_dstate.template operator()<256>(); + break; + default: + FLASHINFER_CHECK(false, "Unsupported dstate value. Supported values are: 64, 128, 256"); + } } else { - constexpr auto numConsumers = 4; - constexpr auto numWarps = 1 + numConsumers; - constexpr auto numStages = 3; - constexpr auto rowsPerStage = 4 * numConsumers; - constexpr auto DIM = 64; - FLASHINFER_CHECK(params.dim == DIM); - auto scan_func = selective_state_update_kernel_producer_consumer_vertical< - input_t, weight_t, matrixA_t, state_t, DIM, DSTATE, numConsumers, rowsPerStage, numStages>; - - dim3 block(32, numWarps); - dim3 grid(params.batch, params.nheads); - - auto nh = params.nheads; - auto dim = params.dim; - auto B = params.state_cache_size; - - FLASHINFER_CHECK(reinterpret_cast(params.state) % 128 == - 0); // TMA requires 128B aligned - auto tensorState = - tma::createTensorMap(params.state, B * nh * dim, DSTATE, rowsPerStage, DSTATE); - FLASHINFER_CHECK(params.dim % rowsPerStage == 0); - - scan_func<<>>(params, tensorState); + auto dispatch_dim_dstate = [&]() { + constexpr auto numConsumers = 4; + constexpr auto numWarps = 1 + numConsumers; + constexpr auto numStages = 3; + constexpr auto rowsPerStage = 4 * numConsumers; + FLASHINFER_CHECK(params.dim % rowsPerStage == 0); + auto scan_func = selective_state_update_kernel_producer_consumer_vertical< + input_t, weight_t, matrixA_t, state_t, DIM, DSTATE, numConsumers, rowsPerStage, + numStages>; + + dim3 block(32, numWarps); + dim3 grid(params.batch, params.nheads); + + auto nh = params.nheads; + auto dim = params.dim; + auto B = params.state_cache_size; + + FLASHINFER_CHECK(reinterpret_cast(params.state) % 128 == + 0); // TMA requires 128B aligned + auto tensorState = + tma::createTensorMap(params.state, B * nh * dim, DSTATE, rowsPerStage, DSTATE); + + scan_func<<>>(params, tensorState); + }; + + auto dispatch_dstate = [&]() { + switch (params.dstate) { + case 64: + dispatch_dim_dstate.template operator()(); + break; + case 128: + dispatch_dim_dstate.template operator()(); + break; + case 256: + dispatch_dim_dstate.template operator()(); + break; + default: + FLASHINFER_CHECK(false, "Unsupported dstate value. Supported values are: 64, 128, 256"); + } + }; + + switch (params.dim) { + case 64: + dispatch_dstate.template operator()<64>(); + break; + case 128: + dispatch_dstate.template operator()<128>(); + break; + default: + FLASHINFER_CHECK(false, "Unsupported dim value. Supported values are: 64, 128"); + } } } From ce2e8682a04d68ed35cc0658ddd4af40099d7543 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 15:22:07 -0800 Subject: [PATCH 22/38] selective state: test various dims and state sizes --- tests/mamba/test_selective_state_update.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/mamba/test_selective_state_update.py b/tests/mamba/test_selective_state_update.py index deb76d2a39..a5a1f80bd8 100644 --- a/tests/mamba/test_selective_state_update.py +++ b/tests/mamba/test_selective_state_update.py @@ -69,10 +69,10 @@ def create_test_inputs( } -@pytest.mark.parametrize("batch", [1, 32, 64]) +@pytest.mark.parametrize("batch", [1, 64]) @pytest.mark.parametrize("nheads", [8, 64]) -@pytest.mark.parametrize("dim", [64]) -@pytest.mark.parametrize("dstate", [128]) +@pytest.mark.parametrize("dim", [64, 128]) +@pytest.mark.parametrize("dstate", [64, 128, 256]) @pytest.mark.parametrize("ngroups", [8]) @pytest.mark.parametrize("delta_softplus", [True]) @pytest.mark.parametrize("input_dtype", [torch.bfloat16]) From d0cd9da609f65f62e72945425a8a8d61c410245c Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Wed, 7 Jan 2026 15:41:16 -0800 Subject: [PATCH 23/38] Fix https://github.com/flashinfer-ai/flashinfer/pull/2301#pullrequestreview-3637176283 --- flashinfer/mamba/selective_state_update.py | 2 +- tests/mamba/selective_state_update_triton.py | 4 ++-- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/flashinfer/mamba/selective_state_update.py b/flashinfer/mamba/selective_state_update.py index d8853061a5..0a045197a5 100644 --- a/flashinfer/mamba/selective_state_update.py +++ b/flashinfer/mamba/selective_state_update.py @@ -49,7 +49,7 @@ def selective_state_update( state_batch_indices: Optional[torch.Tensor] = None, pad_slot_id: int = -1, ) -> torch.Tensor: - r"""Selective state update operation for Mamba layers. + r"""Selective state update operation for Mamba layers (the generation phase). Parameters ---------- diff --git a/tests/mamba/selective_state_update_triton.py b/tests/mamba/selective_state_update_triton.py index f27886212b..2d46ffe0c9 100644 --- a/tests/mamba/selective_state_update_triton.py +++ b/tests/mamba/selective_state_update_triton.py @@ -329,7 +329,7 @@ def selective_state_update_triton( dt.stride(0), dt.stride(1), dt.stride(2), - *(dt_bias.stride(0), dt_bias.stride(1)) if dt_bias is not None else 0, + *(dt_bias.stride(0), dt_bias.stride(1)) if dt_bias is not None else (0, 0), A.stride(0), A.stride(1), A.stride(2), @@ -339,7 +339,7 @@ def selective_state_update_triton( C.stride(0), C.stride(1), C.stride(2), - *(D.stride(0), D.stride(1)) if D is not None else 0, + *(D.stride(0), D.stride(1)) if D is not None else (0, 0), z_strides[0], z_strides[1], z_strides[2], From 0f5f4b8c8d12c0cf68035e8c0a8d215475f8967d Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 8 Jan 2026 16:52:31 +0000 Subject: [PATCH 24/38] formatting --- tests/mamba/test_selective_state_update.py | 16 ++++++++++++++-- 1 file changed, 14 insertions(+), 2 deletions(-) diff --git a/tests/mamba/test_selective_state_update.py b/tests/mamba/test_selective_state_update.py index a5a1f80bd8..65849c436f 100644 --- a/tests/mamba/test_selective_state_update.py +++ b/tests/mamba/test_selective_state_update.py @@ -8,7 +8,15 @@ def create_test_inputs( - batch_size, nheads, dim, dstate, ngroups, input_dtype, weight_dtype, matrixA_dtype, z_none=True + batch_size, + nheads, + dim, + dstate, + ngroups, + input_dtype, + weight_dtype, + matrixA_dtype, + z_none=True, ): # Set seed for reproducibility torch.manual_seed(0) @@ -53,7 +61,11 @@ def create_test_inputs( ] # Create z tensor if z_none is False - z = None if z_none else torch.randn(batch_size, nheads, dim, dtype=input_dtype, device=device) + z = ( + None + if z_none + else torch.randn(batch_size, nheads, dim, dtype=input_dtype, device=device) + ) return { "state_cache": state_cache, From a35c4081d9e64534e8e961a1c6e86853c1da66e9 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 8 Jan 2026 22:29:03 +0000 Subject: [PATCH 25/38] Add SM90 Mamba selective state update JIT module --- flashinfer/aot.py | 8 ++++++-- .../jit/mamba/selective_state_update.py | 19 +++++++++++++++++-- 2 files changed, 23 insertions(+), 4 deletions(-) diff --git a/flashinfer/aot.py b/flashinfer/aot.py index a7566d8fc3..242fadbe81 100644 --- a/flashinfer/aot.py +++ b/flashinfer/aot.py @@ -63,7 +63,10 @@ ) from .jit.spdlog import gen_spdlog_module from .jit.mla import gen_mla_module -from .jit.mamba import gen_selective_state_update_module +from .jit.mamba import ( + gen_selective_state_update_module, + gen_selective_state_update_sm90_module, +) from .jit.norm import gen_norm_module from .jit.page import gen_page_module from .jit.quantization import gen_quantization_module @@ -534,9 +537,10 @@ def gen_all_modules( gen_rope_module(), gen_sampling_module(), gen_topk_module(), + gen_selective_state_update_module(), ] - jit_specs.append(gen_selective_state_update_module()) if has_sm90: + jit_specs.append(gen_selective_state_update_sm90_module()) jit_specs.append(gen_trtllm_utils_module()) jit_specs.append(gen_gdn_prefill_sm90_module()) diff --git a/flashinfer/jit/mamba/selective_state_update.py b/flashinfer/jit/mamba/selective_state_update.py index 266b50edb8..0431c51587 100644 --- a/flashinfer/jit/mamba/selective_state_update.py +++ b/flashinfer/jit/mamba/selective_state_update.py @@ -15,13 +15,12 @@ """ from .. import env as jit_env -from ..core import JitSpec, gen_jit_spec +from ..core import JitSpec, gen_jit_spec, sm90a_nvcc_flags def gen_selective_state_update_module() -> JitSpec: nvcc_flags = [ "-DENABLE_BF16", - "-DENABLE_FP8", ] return gen_jit_spec( @@ -32,3 +31,19 @@ def gen_selective_state_update_module() -> JitSpec: ], extra_cuda_cflags=nvcc_flags, ) + + +def gen_selective_state_update_sm90_module() -> JitSpec: + nvcc_flags = sm90a_nvcc_flags + [ + "-DENABLE_BF16", + "-DFLASHINFER_MAMBA_ENABLE_SM90", + ] + + return gen_jit_spec( + "mamba_selective_state_update_sm90", + [ + jit_env.FLASHINFER_CSRC_DIR / "selective_state_update.cu", + jit_env.FLASHINFER_CSRC_DIR / "flashinfer_mamba_binding.cu", + ], + extra_cuda_cflags=nvcc_flags, + ) From 81163a2268baa4e1a4f4fb54cd36f9b1237e2ea9 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 8 Jan 2026 22:46:12 +0000 Subject: [PATCH 26/38] ifdef guards for the hopper+ implementation for aot --- include/flashinfer/mamba/selective_state_update.cuh | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh index 2b72115512..3701e13fb0 100755 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -245,6 +245,7 @@ template __global__ void selective_state_update_kernel_producer_consumer_vertical( SelectiveStateUpdateParams params, __grid_constant__ CUtensorMap const tensorState) { +#ifdef FLASHINFER_MAMBA_ENABLE_SM90 auto* __restrict__ output = reinterpret_cast(params.output); auto const* __restrict__ x = reinterpret_cast(params.x); @@ -433,13 +434,16 @@ __global__ void selective_state_update_kernel_producer_consumer_vertical( convertAndStore(&output[batch * params.out_stride_batch + head * dim + d], out_value); } } +#endif } template void invokeSelectiveStateUpdate(SelectiveStateUpdateParams& params, cudaStream_t stream) { auto [sm_major, sm_minor] = GetCudaComputeCapability(); +#ifdef FLASHINFER_MAMBA_ENABLE_SM90 if (sm_major < 9) // pre-Hopper +#endif { auto dispatch_dstate = [&]() { constexpr int numWarps = 2; @@ -463,7 +467,10 @@ void invokeSelectiveStateUpdate(SelectiveStateUpdateParams& params, cudaStream_t default: FLASHINFER_CHECK(false, "Unsupported dstate value. Supported values are: 64, 128, 256"); } - } else { + } +#ifdef FLASHINFER_MAMBA_ENABLE_SM90 + else { + auto dispatch_dim_dstate = [&]() { constexpr auto numConsumers = 4; constexpr auto numWarps = 1 + numConsumers; @@ -516,6 +523,7 @@ void invokeSelectiveStateUpdate(SelectiveStateUpdateParams& params, cudaStream_t FLASHINFER_CHECK(false, "Unsupported dim value. Supported values are: 64, 128"); } } +#endif } } // namespace flashinfer::mamba From d2a147aec424a8cccb58d1db608aa8b3020fc27c Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 8 Jan 2026 23:20:37 +0000 Subject: [PATCH 27/38] export both selective_state_update and selective_state_update_sm90 modules --- flashinfer/jit/mamba/__init__.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/flashinfer/jit/mamba/__init__.py b/flashinfer/jit/mamba/__init__.py index a17f20af35..f6a2628b43 100644 --- a/flashinfer/jit/mamba/__init__.py +++ b/flashinfer/jit/mamba/__init__.py @@ -14,6 +14,12 @@ limitations under the License. """ -from .selective_state_update import gen_selective_state_update_module +from .selective_state_update import ( + gen_selective_state_update_module, + gen_selective_state_update_sm90_module, +) -__all__ = ["gen_selective_state_update_module"] +__all__ = [ + "gen_selective_state_update_module", + "gen_selective_state_update_sm90_module", +] From 040734ea0e243e1da538e7cc2e30900b61c8d7c2 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 8 Jan 2026 22:34:35 -0800 Subject: [PATCH 28/38] Add alignment checks for Mamba state update This ensures pointers and strides meet vectorization requirements on pre-Hopper architectures (SM < 90) to prevent illegal memory accesses. Fixes: https://github.com/flashinfer-ai/flashinfer/pull/2301#discussion_r2674245145 --- .../mamba/selective_state_update.cuh | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh index 3701e13fb0..f128adbcb9 100755 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -445,6 +445,24 @@ void invokeSelectiveStateUpdate(SelectiveStateUpdateParams& params, cudaStream_t if (sm_major < 9) // pre-Hopper #endif { + using Load = VectorizedLoadTraits; + constexpr size_t vec_size = sizeof(typename Load::input); + constexpr size_t state_vec_size = sizeof(typename Load::state); + if (Load::chunk_size > 1) { + FLASHINFER_CHECK(reinterpret_cast(params.state) % state_vec_size == 0, + "state pointer must be aligned to ", state_vec_size, " bytes"); + FLASHINFER_CHECK(reinterpret_cast(params.B) % vec_size == 0, + "B pointer must be aligned to ", vec_size, " bytes"); + FLASHINFER_CHECK(reinterpret_cast(params.C) % vec_size == 0, + "C pointer must be aligned to ", vec_size, " bytes"); + FLASHINFER_CHECK((params.B_stride_batch * sizeof(input_t)) % vec_size == 0, + "B batch stride must be aligned to ", vec_size, " bytes"); + FLASHINFER_CHECK((params.C_stride_batch * sizeof(input_t)) % vec_size == 0, + "C batch stride must be aligned to ", vec_size, " bytes"); + FLASHINFER_CHECK((params.dim * params.dstate * sizeof(state_t)) % state_vec_size == 0, + "state head stride must be aligned to ", state_vec_size, " bytes"); + } + auto dispatch_dstate = [&]() { constexpr int numWarps = 2; int const blocks_per_dim = (params.dim + 32 * numWarps - 1) / (32 * numWarps); From 6d9aa1b8cbf90e2c4eee08f6074bd71e7fc3daea Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Thu, 8 Jan 2026 22:39:50 -0800 Subject: [PATCH 29/38] Fix `toFloat` usage in Mamba selective state update kernel Removes an unnecessary `toFloat` conversion on `d_value` in the vertical producer-consumer kernel, as `d_value` is already a float type. Fixes: https://github.com/flashinfer-ai/flashinfer/pull/2301#discussion_r2674245149 --- include/flashinfer/mamba/selective_state_update.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh index f128adbcb9..5234046c75 100755 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -393,7 +393,7 @@ __global__ void selective_state_update_kernel_producer_consumer_vertical( for (auto dd = warp; dd < rowsPerStage; dd += consumerWarps) { auto d = dBegin + dd; float const x_value = toFloat(sram.x[d]); - float out_value = toFloat(d_value) * x_value * int(lane == 0); // first lane has the value + float out_value = d_value * x_value * int(lane == 0); // first lane has the value for (int i = lane; i < DSTATE; i += warpSize) { auto const state_value = (state_batch != params.pad_slot_id) From d22eec49f7264f51d38bf8ee96478c9f1306ee8c Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Fri, 9 Jan 2026 09:10:41 -0800 Subject: [PATCH 30/38] avoid an ambiguous variable name --- include/flashinfer/mamba/selective_state_update.cuh | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) mode change 100755 => 100644 include/flashinfer/mamba/selective_state_update.cuh diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh old mode 100755 new mode 100644 index 5234046c75..e3bff89033 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -504,12 +504,11 @@ void invokeSelectiveStateUpdate(SelectiveStateUpdateParams& params, cudaStream_t auto nh = params.nheads; auto dim = params.dim; - auto B = params.state_cache_size; FLASHINFER_CHECK(reinterpret_cast(params.state) % 128 == 0); // TMA requires 128B aligned - auto tensorState = - tma::createTensorMap(params.state, B * nh * dim, DSTATE, rowsPerStage, DSTATE); + auto tensorState = tma::createTensorMap( + params.state, params.state_cache_size * nh * dim, DSTATE, rowsPerStage, DSTATE); scan_func<<>>(params, tensorState); }; From 59c9f85bb8ae40d1f43a4536b925690ba9276547 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Fri, 9 Jan 2026 09:32:47 -0800 Subject: [PATCH 31/38] Support SM90+ for Mamba selective state update Updates JIT flags to use generic compute_90 PTX for forward compatibility with SM100. Adds runtime dispatch logic to select the optimized SM90 kernel on devices with compute capability >= 9.0. Addresses https://github.com/flashinfer-ai/flashinfer/pull/2301#discussion_r2674245153 --- .../jit/mamba/selective_state_update.py | 9 ++++-- flashinfer/mamba/selective_state_update.py | 30 +++++++++++++------ 2 files changed, 27 insertions(+), 12 deletions(-) diff --git a/flashinfer/jit/mamba/selective_state_update.py b/flashinfer/jit/mamba/selective_state_update.py index 0431c51587..5384bd5c5b 100644 --- a/flashinfer/jit/mamba/selective_state_update.py +++ b/flashinfer/jit/mamba/selective_state_update.py @@ -15,7 +15,7 @@ """ from .. import env as jit_env -from ..core import JitSpec, gen_jit_spec, sm90a_nvcc_flags +from ..core import JitSpec, common_nvcc_flags, gen_jit_spec def gen_selective_state_update_module() -> JitSpec: @@ -34,10 +34,13 @@ def gen_selective_state_update_module() -> JitSpec: def gen_selective_state_update_sm90_module() -> JitSpec: - nvcc_flags = sm90a_nvcc_flags + [ + # Use generic SM90 flags to support SM90, SM100 and future architectures + # code=compute_90 embeds PTX for forward compatibility + nvcc_flags = [ + "-gencode=arch=compute_90,code=[sm_90,compute_90]", "-DENABLE_BF16", "-DFLASHINFER_MAMBA_ENABLE_SM90", - ] + ] + common_nvcc_flags return gen_jit_spec( "mamba_selective_state_update_sm90", diff --git a/flashinfer/mamba/selective_state_update.py b/flashinfer/mamba/selective_state_update.py index 0a045197a5..e435423570 100644 --- a/flashinfer/mamba/selective_state_update.py +++ b/flashinfer/mamba/selective_state_update.py @@ -20,20 +20,32 @@ import torch from ..api_logging import flashinfer_api -from ..jit.mamba import gen_selective_state_update_module -from ..utils import register_custom_op, register_fake_op +from ..jit.mamba import ( + gen_selective_state_update_module, + gen_selective_state_update_sm90_module, +) +from ..utils import get_compute_capability, register_custom_op, register_fake_op @functools.cache -def get_selective_state_update_module(): - """Get cached JIT-compiled selective_state_update module. - - This uses @functools.cache for in-memory caching of loaded modules. - The JIT system also caches compiled .so files on disk in ~/.cache/flashinfer/ - """ +def get_selective_state_update_module_base(): + """Get cached JIT-compiled selective_state_update module (base version).""" return gen_selective_state_update_module().build_and_load() +@functools.cache +def get_selective_state_update_module_sm90(): + """Get cached JIT-compiled selective_state_update module (SM90+ version).""" + return gen_selective_state_update_sm90_module().build_and_load() + + +def get_selective_state_update_module(device: torch.device): + if get_compute_capability(device)[0] >= 9: + return get_selective_state_update_module_sm90() + else: + return get_selective_state_update_module_base() + + @flashinfer_api def selective_state_update( state: torch.Tensor, @@ -141,7 +153,7 @@ def _selective_state_update( pad_slot_id: int, ) -> None: """Internal function registered with torch.library for torch.compile() support.""" - get_selective_state_update_module().selective_state_update( + get_selective_state_update_module(state.device).selective_state_update( state, x, dt, From c443c54566c38ff19634e27e22c0478c447e5c3c Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Fri, 9 Jan 2026 11:23:54 -0800 Subject: [PATCH 32/38] Improve FLASHINFER_CHECK error message in selective_state_update Include the required divisor and kernel variant in the error message for easier debugging when dim is not properly aligned. Address: https://github.com/flashinfer-ai/flashinfer/pull/2301#pullrequestreview-3645264047 --- include/flashinfer/mamba/selective_state_update.cuh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/include/flashinfer/mamba/selective_state_update.cuh b/include/flashinfer/mamba/selective_state_update.cuh index e3bff89033..1ea1dd3182 100644 --- a/include/flashinfer/mamba/selective_state_update.cuh +++ b/include/flashinfer/mamba/selective_state_update.cuh @@ -494,7 +494,8 @@ void invokeSelectiveStateUpdate(SelectiveStateUpdateParams& params, cudaStream_t constexpr auto numWarps = 1 + numConsumers; constexpr auto numStages = 3; constexpr auto rowsPerStage = 4 * numConsumers; - FLASHINFER_CHECK(params.dim % rowsPerStage == 0); + FLASHINFER_CHECK(params.dim % rowsPerStage == 0, "dim must be divisible by ", rowsPerStage, + " for SM90+ kernel"); auto scan_func = selective_state_update_kernel_producer_consumer_vertical< input_t, weight_t, matrixA_t, state_t, DIM, DSTATE, numConsumers, rowsPerStage, numStages>; From 95fa14dd84091dfdee67d23420ed719bed133039 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Fri, 9 Jan 2026 12:05:20 -0800 Subject: [PATCH 33/38] Refactor SM90 module to use CompilationContext for nvcc flags This enables architecture filtering via supported_major_versions and removes reliance on hardcoded nvcc flags. Future architectures are now supported automatically. address https://github.com/flashinfer-ai/flashinfer/pull/2301#discussion_r2677343561 --- flashinfer/jit/mamba/selective_state_update.py | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) diff --git a/flashinfer/jit/mamba/selective_state_update.py b/flashinfer/jit/mamba/selective_state_update.py index 5384bd5c5b..fd121d8e0a 100644 --- a/flashinfer/jit/mamba/selective_state_update.py +++ b/flashinfer/jit/mamba/selective_state_update.py @@ -15,7 +15,8 @@ """ from .. import env as jit_env -from ..core import JitSpec, common_nvcc_flags, gen_jit_spec +from ...compilation_context import CompilationContext +from ..core import JitSpec, gen_jit_spec def gen_selective_state_update_module() -> JitSpec: @@ -34,13 +35,16 @@ def gen_selective_state_update_module() -> JitSpec: def gen_selective_state_update_sm90_module() -> JitSpec: - # Use generic SM90 flags to support SM90, SM100 and future architectures - # code=compute_90 embeds PTX for forward compatibility - nvcc_flags = [ - "-gencode=arch=compute_90,code=[sm_90,compute_90]", + # Use CompilationContext to get nvcc flags filtered by supported major versions + # This supports SM90 (Hopper) and future architectures + compilation_context = CompilationContext() + nvcc_flags = compilation_context.get_nvcc_flags_list( + supported_major_versions=[9, 10, 11, 12] + ) + nvcc_flags += [ "-DENABLE_BF16", "-DFLASHINFER_MAMBA_ENABLE_SM90", - ] + common_nvcc_flags + ] return gen_jit_spec( "mamba_selective_state_update_sm90", From dfc0a38c538a85e89d96674990014e2272f5a774 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 12 Jan 2026 09:29:06 -0800 Subject: [PATCH 34/38] Exclude an unnecessary compiler flag as it is part of the default flags. addresses https://github.com/flashinfer-ai/flashinfer/pull/2301#discussion_r2680026505 --- flashinfer/jit/mamba/selective_state_update.py | 8 +------- 1 file changed, 1 insertion(+), 7 deletions(-) diff --git a/flashinfer/jit/mamba/selective_state_update.py b/flashinfer/jit/mamba/selective_state_update.py index fd121d8e0a..03d2791b5c 100644 --- a/flashinfer/jit/mamba/selective_state_update.py +++ b/flashinfer/jit/mamba/selective_state_update.py @@ -14,23 +14,18 @@ limitations under the License. """ -from .. import env as jit_env from ...compilation_context import CompilationContext +from .. import env as jit_env from ..core import JitSpec, gen_jit_spec def gen_selective_state_update_module() -> JitSpec: - nvcc_flags = [ - "-DENABLE_BF16", - ] - return gen_jit_spec( "mamba_selective_state_update", [ jit_env.FLASHINFER_CSRC_DIR / "selective_state_update.cu", jit_env.FLASHINFER_CSRC_DIR / "flashinfer_mamba_binding.cu", ], - extra_cuda_cflags=nvcc_flags, ) @@ -42,7 +37,6 @@ def gen_selective_state_update_sm90_module() -> JitSpec: supported_major_versions=[9, 10, 11, 12] ) nvcc_flags += [ - "-DENABLE_BF16", "-DFLASHINFER_MAMBA_ENABLE_SM90", ] From baa55a1ff391f7afdf9caba3e8eabed78cb7c2cb Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 12 Jan 2026 09:33:06 -0800 Subject: [PATCH 35/38] comment about the use of TMA to justify the neccesity of sm90 module --- flashinfer/jit/mamba/selective_state_update.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/flashinfer/jit/mamba/selective_state_update.py b/flashinfer/jit/mamba/selective_state_update.py index 03d2791b5c..116ecf96c8 100644 --- a/flashinfer/jit/mamba/selective_state_update.py +++ b/flashinfer/jit/mamba/selective_state_update.py @@ -30,7 +30,8 @@ def gen_selective_state_update_module() -> JitSpec: def gen_selective_state_update_sm90_module() -> JitSpec: - # Use CompilationContext to get nvcc flags filtered by supported major versions + # We use a specialized model for Hopper+ GPUs due to the explicit use + # of TMA device functions. # This supports SM90 (Hopper) and future architectures compilation_context = CompilationContext() nvcc_flags = compilation_context.get_nvcc_flags_list( From 5afc7f6527916c2d0ec83c61558a877011ef1ac4 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 12 Jan 2026 09:38:16 -0800 Subject: [PATCH 36/38] a comment about the choice of A tensor in the unit test --- tests/mamba/test_selective_state_update.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/tests/mamba/test_selective_state_update.py b/tests/mamba/test_selective_state_update.py index 65849c436f..8d47523753 100644 --- a/tests/mamba/test_selective_state_update.py +++ b/tests/mamba/test_selective_state_update.py @@ -37,8 +37,9 @@ def create_test_inputs( dt = torch.randn(batch_size, nheads, dtype=weight_dtype, device=device).as_strided( (batch_size, nheads, dim), (nheads, 1, 0) ) - - A_base = -torch.rand(nheads, dtype=matrixA_dtype, device=device) + # The dtype of A is separate in nemotron nano v3. + # A only has one value per head (as discussed in mamba 2 block) hence the strides. + A_base = torch.rand(nheads, dtype=matrixA_dtype, device=device) A = A_base.as_strided((nheads, dim, dstate), (1, 0, 0)) assert A.stride() == (1, 0, 0) From 0d9c71b9f54479a6dc4a717a83eb166c83020f69 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 12 Jan 2026 09:43:39 -0800 Subject: [PATCH 37/38] Use torch.testing.assert_allclose instead of torch.allclose as the latter can display the number of mismatch elements and the location of mismatch elements. --- tests/mamba/test_selective_state_update.py | 20 ++------------------ 1 file changed, 2 insertions(+), 18 deletions(-) diff --git a/tests/mamba/test_selective_state_update.py b/tests/mamba/test_selective_state_update.py index 8d47523753..e50994646a 100644 --- a/tests/mamba/test_selective_state_update.py +++ b/tests/mamba/test_selective_state_update.py @@ -285,26 +285,10 @@ def test_selective_state_update_with_z(): atol = 1e-3 rtol = 1e-2 - outputs_match = torch.allclose(y_ref, y_test, atol=atol, rtol=rtol) - - if outputs_match: - print(f"✓ Outputs match within tolerance (atol={atol}, rtol={rtol})") - else: - print(f"✗ Outputs do NOT match within tolerance (atol={atol}, rtol={rtol})") - - assert outputs_match - - # Check if states match within tolerance - states_match = torch.allclose( + torch.testing.assert_allclose(y_ref, y_test, atol=atol, rtol=rtol) + torch.testing.assert_allclose( state_ref[inputs["slot_idx"]], state[inputs["slot_idx"]], atol=atol, rtol=rtol, ) - - if states_match: - print(f"✓ States match within tolerance (atol={atol}, rtol={rtol})") - else: - print(f"✗ States do NOT match within tolerance (atol={atol}, rtol={rtol})") - - assert states_match From d99923c231cb95f04843b12acfde8267523216c8 Mon Sep 17 00:00:00 2001 From: Igor Shovkun Date: Mon, 12 Jan 2026 15:24:53 -0800 Subject: [PATCH 38/38] fix obsoleted bf16 ifdefs --- include/flashinfer/mamba/conversion.cuh | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/include/flashinfer/mamba/conversion.cuh b/include/flashinfer/mamba/conversion.cuh index 7e4fba558e..7e8a84bf8a 100644 --- a/include/flashinfer/mamba/conversion.cuh +++ b/include/flashinfer/mamba/conversion.cuh @@ -1,8 +1,9 @@ #pragma once #include -#include #include - +#ifdef FLASHINFER_ENABLE_BF16 +#include +#endif #ifdef ENABLE_FP8 #include #endif @@ -12,7 +13,8 @@ namespace flashinfer::mamba::conversion { inline __device__ float toFloat(float f) { return f; } inline __device__ float toFloat(__half h) { return __half2float(h); } -#ifdef ENABLE_BF16 + +#ifdef FLASHINFER_ENABLE_BF16 inline __device__ float toFloat(__nv_bfloat16 val) { return __bfloat162float(val); } #endif @@ -22,7 +24,7 @@ inline __device__ void convertAndStore(__half* output, float input) { *output = __float2half(input); } -#ifdef ENABLE_BF16 +#ifdef FLASHINFER_ENABLE_BF16 inline __device__ void convertAndStore(__nv_bfloat16* output, float input) { *output = __float2bfloat16(input); }