Enable fp16/bf16/f32 support for selective_state_update (mamba)#2366

Merged

yzh119 merged 51 commits intoflashinfer-ai:mainfrom

Jan 19, 2026

Commits on Dec 22, 2025

I seem that I can call my cuda wrapper from python
Igor Shovkun
committed
check all the kernel inputs. ready to develop dispatch code.
Igor Shovkun
committed

Commits on Jan 5, 2026

simple implementation of selective_state_update is working
Igor Shovkun
committed
ported the hopper version + runtime dispatch check
Igor Shovkun
committed

Commits on Jan 6, 2026

Passed pre-commit checks
ishovkun
committed

Commits on Jan 7, 2026

Update flashinfer/mamba/selective_state_update.py

ishovkun
and
gemini-code-assist[bot]
authored
remove unreachable code
Igor Shovkun
committed
Improve docstring for input state shape

ishovkun
and
gemini-code-assist[bot]
authored
Simple kernel also uses fast_exp
Igor Shovkun
committed
Remove error check of the kernel launch (it was a debugging leftover).
Igor Shovkun
committed
no need for the fast_exp function
Igor Shovkun
committed
Remove unnecessary None check for D before unsqueeze
Igor Shovkun
committed
Hoist dA computation outside the innermost loop.
Igor Shovkun
committed
support for non-none z
Igor Shovkun
committed
stage forgotten z handling in the test
Igor Shovkun
committed
test: handle z
Igor Shovkun
committed
make sure that batch size does not exceed the state cache size
Igor Shovkun
committed
do use matrixA_dtype in the test
Igor Shovkun
committed
Add selective state update module to JIT specs
Igor Shovkun
committed
Fix dt_bias stride check to handle None value
Igor Shovkun
committed
handle a few different DIM and DSTATE
Igor Shovkun
committed
selective state: test various dims and state sizes
Igor Shovkun
committed
Fix
Igor Shovkun
committed

Commits on Jan 17, 2026

Commits on Jan 18, 2026