flashinfer-python 0.6.4__tar.gz → 0.6.5__tar.gz

{flashinfer_python-0.6.4 → flashinfer_python-0.6.5}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: flashinfer-python
-Version: 0.6.4
+Version: 0.6.5
 Summary: FlashInfer: Kernel Library for LLM Serving
 Author: FlashInfer team
 License-Expression: Apache-2.0
@@ -61,6 +61,7 @@ High-Performance GPU Kernels for Inference
 - **POD-Attention**: Fused prefill+decode for mixed batching
 
 ### GEMM & Linear Operations
+- **BF16 GEMM**: BF16 matrix multiplication for SM10.0+ GPUs.
 - **FP8 GEMM**: Per-tensor and groupwise scaling
 - **FP4 GEMM**: NVFP4 and MXFP4 matrix multiplication for Blackwell GPUs
 - **Grouped GEMM**: Efficient batched matrix operations for LoRA and multi-expert routing
@@ -165,6 +166,11 @@ python -m pip install -v .
 python -m pip install --no-build-isolation -e . -v
 ```
 
+> **Note:** When using `--no-build-isolation`, pip does not automatically install build dependencies. FlashInfer requires `setuptools>=77`. If you encounter an error like `AttributeError: module 'setuptools.build_meta' has no attribute 'prepare_metadata_for_build_editable'`, upgrade pip and setuptools first:
+> ```bash
+> python -m pip install --upgrade pip setuptools
+> ```
+
 Build optional packages:
 
 ```bash

{flashinfer_python-0.6.4 → flashinfer_python-0.6.5}/README.md

@@ -36,6 +36,7 @@ High-Performance GPU Kernels for Inference
 - **POD-Attention**: Fused prefill+decode for mixed batching
 
 ### GEMM & Linear Operations
+- **BF16 GEMM**: BF16 matrix multiplication for SM10.0+ GPUs.
 - **FP8 GEMM**: Per-tensor and groupwise scaling
 - **FP4 GEMM**: NVFP4 and MXFP4 matrix multiplication for Blackwell GPUs
 - **Grouped GEMM**: Efficient batched matrix operations for LoRA and multi-expert routing
@@ -140,6 +141,11 @@ python -m pip install -v .
 python -m pip install --no-build-isolation -e . -v
 ```
 
+> **Note:** When using `--no-build-isolation`, pip does not automatically install build dependencies. FlashInfer requires `setuptools>=77`. If you encounter an error like `AttributeError: module 'setuptools.build_meta' has no attribute 'prepare_metadata_for_build_editable'`, upgrade pip and setuptools first:
+> ```bash
+> python -m pip install --upgrade pip setuptools
+> ```
+
 Build optional packages:
 
 ```bash

{flashinfer_python-0.6.4 → flashinfer_python-0.6.5}/csrc/flashinfer_mamba_binding.cu

@@ -42,7 +42,8 @@ void selective_state_update(
     bool disable_state_update,
     Optional<TensorView> intermediate_states_buffer,  // (batch, cache_steps, nheads, dim, dstate)
     Optional<TensorView> intermediate_state_indices,  // (batch,)
-    int64_t cache_steps);
+    int64_t cache_steps,
+    int64_t algorithm);  // SSUAlgorithm: 0=auto, 1=simple, 2=vertical, 3=horizontal
 
 }  // namespace flashinfer::mamba
 

{flashinfer_python-0.6.4 → flashinfer_python-0.6.5}/csrc/fused_moe/cutlass_backend/flashinfer_cutlass_fused_moe_binding.cu

@@ -976,10 +976,10 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj {
       TVM_FFI_ICHECK_EQ(quant_scales.value().size(), 4)
       "Expecting 4 quant scales for W4A8_MXFP4_MXFP8 quantization";
 
-      TensorView fc1_weight_block = quant_scales.value()[0];
-      TensorView fc1_global = quant_scales.value()[1];
-      TensorView fc2_weight_block = quant_scales.value()[2];
-      TensorView fc2_global = quant_scales.value()[3];
+      auto const& fc1_weight_block = quant_scales.value()[0];
+      auto const& fc1_global = quant_scales.value()[1];
+      auto const& fc2_weight_block = quant_scales.value()[2];
+      auto const& fc2_global = quant_scales.value()[3];
 
       // The input for scale fc1_weight_block / fc2_weight_block is packed into INT32
       constexpr int FP8_PER_INT32 = 4;
@@ -1035,12 +1035,12 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj {
       TVM_FFI_ICHECK_EQ(quant_scales.value().size(), 6)
           << "Expecting 6 quant scales for nvfp4 quantization";
 
-      TensorView fc1_act_global = quant_scales.value()[0];
-      TensorView fc1_weight_block = quant_scales.value()[1];
-      TensorView fc1_global = quant_scales.value()[2];
-      TensorView fc2_act_global = quant_scales.value()[3];
-      TensorView fc2_weight_block = quant_scales.value()[4];
-      TensorView fc2_global = quant_scales.value()[5];
+      auto const& fc1_act_global = quant_scales.value()[0];
+      auto const& fc1_weight_block = quant_scales.value()[1];
+      auto const& fc1_global = quant_scales.value()[2];
+      auto const& fc2_act_global = quant_scales.value()[3];
+      auto const& fc2_weight_block = quant_scales.value()[4];
+      auto const& fc2_global = quant_scales.value()[5];
 
       // The input for scale fc1_weight_block / fc2_weight_block is packed into INT32
       constexpr int FP8_PER_INT32 = 4;
@@ -1118,8 +1118,8 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj {
           static_cast<float const*>(fc2_global.data_ptr()), fc1_act_global.ndim() == 1,
           fc2_act_global.ndim() == 1);
     } else if (mUseDeepSeekFP8BlockScaling) {
-      TensorView fc1_scales = quant_scales.value()[0];
-      TensorView fc2_scales = quant_scales.value()[1];
+      auto const& fc1_scales = quant_scales.value()[0];
+      auto const& fc2_scales = quant_scales.value()[1];
       return kernels::QuantParams::FP8BlockScaling(
           static_cast<float const*>(fc1_scales.data_ptr()),
           static_cast<float const*>(fc2_scales.data_ptr()));
@@ -1128,8 +1128,8 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj {
       TVM_FFI_ICHECK_EQ(quant_scales.value().size(), 2)
           << "Expecting 2 quant scales for W4A16 quantization";
 
-      TensorView fc1_weight_scales = quant_scales.value()[0];
-      TensorView fc2_weight_scales = quant_scales.value()[1];
+      auto const& fc1_weight_scales = quant_scales.value()[0];
+      auto const& fc2_weight_scales = quant_scales.value()[1];
       int group_size = TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::wfp4a16_group_size;
       return kernels::QuantParams::GroupWise(group_size,
                                              static_cast<void const*>(fc1_weight_scales.data_ptr()),
@@ -1139,14 +1139,14 @@ class FusedMoeRunner : public tvm::ffi::ModuleObj {
       TVM_FFI_ICHECK(quant_scales.has_value()) << "Expecting quant scales for INT4 quantization";
       TVM_FFI_ICHECK_EQ(quant_scales.value().size(), 8)
           << "Expecting 8 quant scales for INT4 quantization";
-      TensorView fc1_weight_scales = quant_scales.value()[0];
-      TensorView fc2_weight_scales = quant_scales.value()[1];
-      TensorView fc1_act_scales = quant_scales.value()[2];
-      TensorView fc2_act_scales = quant_scales.value()[3];
-      TensorView fc1_weight_zeros = quant_scales.value()[4];
-      TensorView fc2_weight_zeros = quant_scales.value()[5];
-      TensorView fc1_alpha = quant_scales.value()[6];
-      TensorView fc2_alpha = quant_scales.value()[7];
+      auto const& fc1_weight_scales = quant_scales.value()[0];
+      auto const& fc2_weight_scales = quant_scales.value()[1];
+      auto const& fc1_act_scales = quant_scales.value()[2];
+      auto const& fc2_act_scales = quant_scales.value()[3];
+      auto const& fc1_weight_zeros = quant_scales.value()[4];
+      auto const& fc2_weight_zeros = quant_scales.value()[5];
+      auto const& fc1_alpha = quant_scales.value()[6];
+      auto const& fc2_alpha = quant_scales.value()[7];
       int group_size = TmaWarpSpecializedGroupedGemmInput::INT4GroupwiseParams::int4_group_size;
       return kernels::QuantParams::GroupWise(
           group_size, static_cast<void const*>(fc1_weight_scales.data_ptr()),

{flashinfer_python-0.6.4 → flashinfer_python-0.6.5}/csrc/gemm_groupwise_sm120.cu

@@ -31,12 +31,13 @@ using namespace flashinfer;
                                    SCALE_GRANULARITY_M, SCALE_GRANULARITY_N, SCALE_GRANULARITY_K, \
                                    ...)                                                           \
   [&]() -> bool {                                                                                 \
-    /* SM120 Cooperative schedule uses 128x128x128 tile shape */                                  \
+    /* SM120/SM121 Cooperative schedule uses 128x128x128 tile shape */                            \
     /* TODO (yongwww): PingPong schedule (64x128x128) will need additional dispatch logic */      \
     constexpr int SCALE_GRANULARITY_K = 128;                                                      \
     if (scale_granularity_k != 128) {                                                             \
       TVM_FFI_ICHECK(false)                                                                       \
-          << "SM120 requires scale_granularity_k=128. CUTLASS enforces ScaleGranularityK must "   \
+          << "SM120/SM121 requires scale_granularity_k=128. CUTLASS enforces ScaleGranularityK "  \
+             "must "                                                                              \
              "equal tile shape K dimension (128 for both Cooperative and PingPong schedules).";   \
       return false;                                                                               \
     }                                                                                             \

{flashinfer_python-0.6.4 → flashinfer_python-0.6.5}/csrc/group_gemm_fp8_groupwise_sm120.cu

@@ -34,8 +34,9 @@ using namespace flashinfer;
     constexpr int SCALE_GRANULARITY_K = 128;                                                      \
     if (scale_granularity_k != 128) {                                                             \
       TVM_FFI_ICHECK(false)                                                                       \
-          << "SM120 requires scale_granularity_k=128. CUTLASS enforces ScaleGranularityK must "   \
-             "equal tile shape K dimension (128 for both Cooperative and PingPong schedules).";   \
+          << "SM120/SM121 requires scale_granularity_k=128. CUTLASS enforces ScaleGranularityK "  \
+             "must equal tile shape K dimension (128 for both Cooperative and PingPong "          \
+             "schedules).";                                                                       \
       return false;                                                                               \
     }                                                                                             \
     /* Match SM100's approach: support only (1,128,128) and (128,128,128) */                      \

{flashinfer_python-0.6.4 → flashinfer_python-0.6.5}/csrc/moe_utils_binding.cu

@@ -306,8 +306,9 @@ void moe_sort(
   routingData.mPtrCtaIdxXyToMnLimit = reinterpret_cast<int32_t*>(tile_idx_to_mn_limit_ptr);
   routingData.mPtrExpandedIdxToPermutedIdx =
       reinterpret_cast<int32_t*>(expanded_idx_to_permuted_idx_ptr);
-  routingData.mPtrPermutedIdxToTokenIdx =
+  routingData.mPtrPermutedIdxToExpandedIdx =
       reinterpret_cast<int32_t*>(permuted_idx_to_expanded_idx_ptr);
+  routingData.mPtrPermutedIdxToTokenIdx = nullptr;
   routingData.mPtrPermutedIdxSize = reinterpret_cast<int32_t*>(total_num_padded_tokens_ptr);
   routingData.mPtrNumNonExitingCtas = reinterpret_cast<int32_t*>(num_non_exiting_tiles_ptr);
 

{flashinfer_python-0.6.4 → flashinfer_python-0.6.5}/csrc/selective_state_update.cu

@@ -13,9 +13,13 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
+// clang-format off
+// config.inc MUST come before the header: it defines DIM, DSTATE, NTOKENS_MTP
+// constexprs that the header's function templates rely on. Reordering breaks compilation.
+// NOTE: the .inc file is generated from the jinja templates
+#include "selective_state_update_config.inc"
 #include <flashinfer/mamba/selective_state_update.cuh>
-#include <sstream>
-
+// clang-format on
 #include "tvm_ffi_utils.h"
 
 using namespace flashinfer;
@@ -124,87 +128,13 @@ inline void validate_dtype_consistency(
   }
 }
 
-// Helper to convert dtype code to string for error messages
-inline const char* dtype_code_to_string(int64_t code) {
-  if (code == bfloat16_code) return "bfloat16";
-  if (code == float16_code) return "float16";
-  if (code == float32_code) return "float32";
-  return "unknown";
-}
-
-// Type traits to map dtype codes to C++ types
-template <int64_t code>
-struct DTypeToType;
-
-template <>
-struct DTypeToType<bfloat16_code> {
-  using type = nv_bfloat16;
-};
-template <>
-struct DTypeToType<float16_code> {
-  using type = half;
-};
-template <>
-struct DTypeToType<float32_code> {
-  using type = float;
-};
-template <>
-struct DTypeToType<int32_code> {
-  using type = int32_t;
-};
-template <>
-struct DTypeToType<int64_code> {
-  using type = int64_t;
-};
-
-// Allowed dtype combinations: {state_code, input_code, weight_code, matrixA_code, stateIndex_code}
-constexpr std::tuple<int64_t, int64_t, int64_t, int64_t, int64_t> allowed_dtype_combos[] = {
-    {bfloat16_code, bfloat16_code, bfloat16_code, float32_code, int32_code},
-    {float16_code, bfloat16_code, bfloat16_code, float32_code, int32_code},
-    {float32_code, bfloat16_code, bfloat16_code, float32_code, int32_code},
-    {bfloat16_code, bfloat16_code, float32_code, float32_code, int32_code},
-    {float16_code, bfloat16_code, float32_code, float32_code, int32_code},
-    {float32_code, bfloat16_code, float32_code, float32_code, int32_code},
-    {bfloat16_code, bfloat16_code, bfloat16_code, float32_code, int64_code},
-    {float16_code, bfloat16_code, bfloat16_code, float32_code, int64_code},
-    {float32_code, bfloat16_code, bfloat16_code, float32_code, int64_code},
-    {bfloat16_code, bfloat16_code, float32_code, float32_code, int64_code},
-    {float16_code, bfloat16_code, float32_code, float32_code, int64_code},
-    {float32_code, bfloat16_code, float32_code, float32_code, int64_code},
-};
-
-// Helper to dispatch to the right template instantiation for STP
-template <int64_t state_code, int64_t input_code, int64_t weight_code, int64_t matrixA_code,
-          int64_t stateIndex_code>
-void dispatchCombo(SelectiveStateUpdateParams& p, cudaStream_t stream) {
-  using state_t = typename DTypeToType<state_code>::type;
-  using input_t = typename DTypeToType<input_code>::type;
-  using weight_t = typename DTypeToType<weight_code>::type;
-  using matrixA_t = typename DTypeToType<matrixA_code>::type;
-  using stateIndex_t = typename DTypeToType<stateIndex_code>::type;
-  invokeSelectiveStateUpdate<input_t, weight_t, matrixA_t, state_t, stateIndex_t>(p, stream);
-}
-
-// Helper to dispatch to the right template instantiation for MTP
-template <int64_t state_code, int64_t input_code, int64_t weight_code, int64_t matrixA_code,
-          int64_t stateIndex_code>
-void dispatchComboMTP(mtp::SelectiveStateMTPParams& p, cudaStream_t stream) {
-  using state_t = typename DTypeToType<state_code>::type;
-  using input_t = typename DTypeToType<input_code>::type;
-  using weight_t = typename DTypeToType<weight_code>::type;
-  using matrixA_t = typename DTypeToType<matrixA_code>::type;
-  using stateIndex_t = typename DTypeToType<stateIndex_code>::type;
-  mtp::invokeSelectiveStateUpdateMTP<input_t, weight_t, matrixA_t, state_t, stateIndex_t>(p,
-                                                                                          stream);
-}
-
 void run_selective_state_update_stp(TensorView const& state, TensorView const& x,
                                     TensorView const& dt, TensorView const& A, TensorView const& B,
                                     TensorView const& C, TensorView const& D,
                                     Optional<TensorView> z, Optional<TensorView> dt_bias,
                                     bool dt_softplus, Optional<TensorView> state_batch_indices,
                                     int64_t pad_slot_id, Optional<TensorView> out,
-                                    bool disable_state_update) {
+                                    bool disable_state_update, int64_t algorithm) {
   // Extract dimensions from input tensors
   auto const batch = x.size(0);
   auto const state_cache_size = state.size(0);
@@ -344,64 +274,8 @@ void run_selective_state_update_stp(TensorView const& state, TensorView const& x
   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);
-
-  // Get state_batch_indices dtype, default to int32 if not provided
-  int64_t stateIndex_dtype_code = int32_code;
-  if (state_batch_indices.has_value()) {
-    DLDataType stateIndex_dtype = state_batch_indices.value().dtype();
-    stateIndex_dtype_code = encode_dlpack_dtype(stateIndex_dtype);
-  }
-
-  // Dispatch kernel based on dtype combination
-  auto dtype_key = std::make_tuple(state_dtype_code, input_dtype_code, weight_dtype_code,
-                                   matrixA_dtype_code, stateIndex_dtype_code);
-
-  // Compile-time recursive dispatcher using Y-combinator pattern for lambda self-recursion
-  auto tryDispatch = [&](const auto& key, auto idx, auto& self) -> bool {
-    constexpr size_t I = decltype(idx)::value;
-    if constexpr (I < std::size(allowed_dtype_combos)) {
-      constexpr auto combo = allowed_dtype_combos[I];
-      if (key == combo) {
-        constexpr auto s = std::get<0>(combo);
-        constexpr auto i = std::get<1>(combo);
-        constexpr auto w = std::get<2>(combo);
-        constexpr auto m = std::get<3>(combo);
-        constexpr auto si = std::get<4>(combo);
-        dispatchCombo<s, i, w, m, si>(p, stream);
-        return true;
-      }
-      return self(key, std::integral_constant<size_t, I + 1>{}, self);
-    }
-    return false;
-  };
-
-  // Dispatch using compile-time type traits
-  if (!tryDispatch(dtype_key, std::integral_constant<size_t, 0>{}, tryDispatch)) {
-    // Unsupported dtype combination - build error message dynamically
-    std::ostringstream error_msg;
-    error_msg << "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
-              << ". Supported combos include:\n";
-    for (const auto& combo : allowed_dtype_combos) {
-      error_msg << "  (state=" << dtype_code_to_string(std::get<0>(combo))
-                << ", input=" << dtype_code_to_string(std::get<1>(combo))
-                << ", weight=" << dtype_code_to_string(std::get<2>(combo))
-                << ", matrixA=" << dtype_code_to_string(std::get<3>(combo)) << ")\n";
-    }
-    TVM_FFI_ICHECK(false) << error_msg.str();
-  }
+  auto algo = static_cast<SSUAlgorithm>(algorithm);
+  invokeSelectiveStateUpdate<input_t, weight_t, matrixA_t, state_t, stateIndex_t>(p, algo, stream);
 }
 
 void run_selective_state_update_mtp(
@@ -410,7 +284,7 @@ void run_selective_state_update_mtp(
     Optional<TensorView> dt_bias, bool dt_softplus, Optional<TensorView> state_batch_indices,
     int64_t pad_slot_id, Optional<TensorView> out, bool disable_state_update,
     Optional<TensorView> intermediate_states_buffer,
-    Optional<TensorView> intermediate_state_indices, int64_t cache_steps) {
+    Optional<TensorView> intermediate_state_indices, int64_t cache_steps, int64_t algorithm) {
   // Extract dimensions from input tensors
   auto const batch = x.size(0);
   auto const ntokens_mtp = x.size(1);
@@ -505,6 +379,15 @@ void run_selective_state_update_mtp(
   validate_intermediate_state_indices(intermediate_state_indices, batch);
   validate_intermediate_states_buffer(intermediate_states_buffer);
 
+  // Validate that state_batch_indices and intermediate_state_indices have the same dtype
+  if (state_batch_indices.has_value() && intermediate_state_indices.has_value()) {
+    DLDataType state_batch_idx_dtype = state_batch_indices.value().dtype();
+    DLDataType intermediate_idx_dtype = intermediate_state_indices.value().dtype();
+    FLASHINFER_CHECK(state_batch_idx_dtype.code == intermediate_idx_dtype.code &&
+                         state_batch_idx_dtype.bits == intermediate_idx_dtype.bits,
+                     "state_batch_indices and intermediate_state_indices must have the same dtype");
+  }
+
   // Validate cache_steps is non-negative
   FLASHINFER_CHECK(cache_steps >= 0, "cache_steps must be non-negative, got ", cache_steps);
 
@@ -588,75 +471,9 @@ void run_selective_state_update_mtp(
   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);
-
-  // Get stateIndex dtype from whichever index tensor is available
-  // If both are provided, they must have the same dtype
-  int64_t stateIndex_dtype_code = int32_code;  // default
-  if (state_batch_indices.has_value() && intermediate_state_indices.has_value()) {
-    DLDataType state_batch_idx_dtype = state_batch_indices.value().dtype();
-    DLDataType intermediate_idx_dtype = intermediate_state_indices.value().dtype();
-    FLASHINFER_CHECK(state_batch_idx_dtype.code == intermediate_idx_dtype.code &&
-                         state_batch_idx_dtype.bits == intermediate_idx_dtype.bits,
-                     "state_batch_indices and intermediate_state_indices must have the same dtype");
-    stateIndex_dtype_code = encode_dlpack_dtype(state_batch_idx_dtype);
-  } else if (state_batch_indices.has_value()) {
-    DLDataType state_batch_idx_dtype = state_batch_indices.value().dtype();
-    stateIndex_dtype_code = encode_dlpack_dtype(state_batch_idx_dtype);
-  } else if (intermediate_state_indices.has_value()) {
-    DLDataType intermediate_idx_dtype = intermediate_state_indices.value().dtype();
-    stateIndex_dtype_code = encode_dlpack_dtype(intermediate_idx_dtype);
-  }
-
-  // Dispatch kernel based on dtype combination
-  auto dtype_key = std::make_tuple(state_dtype_code, input_dtype_code, weight_dtype_code,
-                                   matrixA_dtype_code, stateIndex_dtype_code);
-
-  // Compile-time recursive dispatcher using Y-combinator pattern for lambda self-recursion
-  auto tryDispatch = [&](const auto& key, auto idx, auto& self) -> bool {
-    constexpr size_t I = decltype(idx)::value;
-    if constexpr (I < std::size(allowed_dtype_combos)) {
-      constexpr auto combo = allowed_dtype_combos[I];
-      if (key == combo) {
-        constexpr auto s = std::get<0>(combo);
-        constexpr auto i = std::get<1>(combo);
-        constexpr auto w = std::get<2>(combo);
-        constexpr auto m = std::get<3>(combo);
-        constexpr auto si = std::get<4>(combo);
-        dispatchComboMTP<s, i, w, m, si>(p, stream);
-        return true;
-      }
-      return self(key, std::integral_constant<size_t, I + 1>{}, self);
-    }
-    return false;
-  };
-
-  // Dispatch using compile-time type traits
-  if (!tryDispatch(dtype_key, std::integral_constant<size_t, 0>{}, tryDispatch)) {
-    // Unsupported dtype combination - build error message dynamically
-    std::ostringstream error_msg;
-    error_msg << "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
-              << ". Supported combos include:\n";
-    for (const auto& combo : allowed_dtype_combos) {
-      error_msg << "  (state=" << dtype_code_to_string(std::get<0>(combo))
-                << ", input=" << dtype_code_to_string(std::get<1>(combo))
-                << ", weight=" << dtype_code_to_string(std::get<2>(combo))
-                << ", matrixA=" << dtype_code_to_string(std::get<3>(combo)) << ")\n";
-    }
-    TVM_FFI_ICHECK(false) << error_msg.str();
-  }
+  auto algo = static_cast<SSUAlgorithm>(algorithm);
+  mtp::invokeSelectiveStateUpdateMTP<input_t, weight_t, matrixA_t, state_t, stateIndex_t>(p, algo,
+                                                                                          stream);
 }
 
 // =============================================================================
@@ -668,14 +485,17 @@ void selective_state_update(TensorView state, TensorView x, TensorView dt, Tenso
                             Optional<TensorView> state_batch_indices, int64_t pad_slot_id,
                             TensorView output, bool disable_state_update,
                             Optional<TensorView> intermediate_states_buffer,
-                            Optional<TensorView> intermediate_state_indices, int64_t cache_steps) {
+                            Optional<TensorView> intermediate_state_indices, int64_t cache_steps,
+                            int64_t algorithm) {
   if (x.dim() == 3) {
     run_selective_state_update_stp(state, x, dt, A, B, C, D, z, dt_bias, dt_softplus,
-                                   state_batch_indices, pad_slot_id, output, disable_state_update);
+                                   state_batch_indices, pad_slot_id, output, disable_state_update,
+                                   algorithm);
   } else if (x.dim() == 4) {
-    run_selective_state_update_mtp(
-        state, x, dt, A, B, C, D, z, dt_bias, dt_softplus, state_batch_indices, pad_slot_id, output,
-        disable_state_update, intermediate_states_buffer, intermediate_state_indices, cache_steps);
+    run_selective_state_update_mtp(state, x, dt, A, B, C, D, z, dt_bias, dt_softplus,
+                                   state_batch_indices, pad_slot_id, output, disable_state_update,
+                                   intermediate_states_buffer, intermediate_state_indices,
+                                   cache_steps, algorithm);
   } else {
     FLASHINFER_CHECK(false,
                      "x must have 3 dimensions (single-token) or 4 dimensions (multi-token), got ",

flashinfer_python-0.6.5/csrc/selective_state_update_customize_config.jinja

@@ -0,0 +1,14 @@
+#pragma once
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <cstdint>
+
+using state_t = {{ state_dtype }};
+using input_t = {{ input_dtype }};
+using weight_t = {{ weight_dtype }};
+using matrixA_t = {{ matrixA_dtype }};
+using stateIndex_t = {{ stateIndex_dtype }};
+
+constexpr int DIM = {{ dim }};
+constexpr int DSTATE = {{ dstate }};
+constexpr int NTOKENS_MTP = {{ ntokens_mtp }};

flashinfer_python-0.6.5/csrc/selective_state_update_dtype_inst.jinja

@@ -0,0 +1,32 @@
+/*
+ * 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.
+ */
+
+// Auto-generated file - do not edit directly.
+// Generated by flashinfer/jit/mamba/selective_state_update.py
+
+#include <flashinfer/mamba/selective_state_update.cuh>
+
+namespace flashinfer::mamba {
+
+template void invokeSelectiveStateUpdate<{{ input_dtype }}, {{ weight_dtype }}, {{ matrixA_dtype }}, {{ state_dtype }}, {{ stateIndex_dtype }}>(
+    SelectiveStateUpdateParams& params, SSUAlgorithm algo, cudaStream_t stream);
+
+namespace mtp {
+template void invokeSelectiveStateUpdateMTP<{{ input_dtype }}, {{ weight_dtype }}, {{ matrixA_dtype }}, {{ state_dtype }}, {{ stateIndex_dtype }}>(
+    SelectiveStateMTPParams& params, SSUAlgorithm algo, cudaStream_t stream);
+}  // namespace mtp
+
+}  // namespace flashinfer::mamba

flashinfer_python-0.6.5/csrc/selective_state_update_kernel_inst.cu

@@ -0,0 +1,18 @@
+// clang-format off
+// config.inc MUST come before the header: it defines DIM, DSTATE, NTOKENS_MTP
+// constexprs that the header's function templates rely on. Reordering breaks compilation.
+#include "selective_state_update_config.inc"
+#include <flashinfer/mamba/selective_state_update.cuh>
+// clang-format on
+
+namespace flashinfer::mamba {
+
+template void invokeSelectiveStateUpdate<input_t, weight_t, matrixA_t, state_t, stateIndex_t>(
+    SelectiveStateUpdateParams&, SSUAlgorithm, cudaStream_t);
+
+namespace mtp {
+template void invokeSelectiveStateUpdateMTP<input_t, weight_t, matrixA_t, state_t, stateIndex_t>(
+    SelectiveStateMTPParams&, SSUAlgorithm, cudaStream_t);
+}  // namespace mtp
+
+}  // namespace flashinfer::mamba