Add mamba kernel (#10234)

2025-09-10 03:58:43 +08:00
parent 8471e5e616
commit 8cbe1538ef
8 changed files with 1418 additions and 0 deletions
--- a/sgl-kernel/python/sgl_kernel/init.py
+++ b/sgl-kernel/python/sgl_kernel/init.py
@@ -34,6 +34,7 @@ from sgl_kernel.elementwise import (
    rmsnorm,
    silu_and_mul,
 )
+from sgl_kernel.mamba import causal_conv1d_fwd, causal_conv1d_update

 if torch.version.hip is not None:
    from sgl_kernel.elementwise import gelu_quick