Selective State Update kernel (mamba)#2301

Merged

yzh119 merged 41 commits intoflashinfer-ai:mainfrom

Jan 13, 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