Upload 62 files

.gitattributes

@@ -33,3 +33,8 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+rscd/models/backbones/lib_mamba/kernels/selective_scan/build/lib.linux-x86_64-3.8/selective_scan_cuda_oflex.cpython-38-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/.ninja_deps filter=lfs diff=lfs merge=lfs -text
+rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/csrc/selective_scan/cusoflex/selective_scan_core_bwd.o filter=lfs diff=lfs merge=lfs -text
+rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/csrc/selective_scan/cusoflex/selective_scan_core_fwd.o filter=lfs diff=lfs merge=lfs -text
+rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/csrc/selective_scan/cusoflex/selective_scan_oflex.o filter=lfs diff=lfs merge=lfs -text

rscd/models/backbones/lib_mamba/__init__.py

@@ -0,0 +1,58 @@
+import os
+from functools import partial
+import torch
+
+from .vmamba import VSSM
+from .csms6s import flops_selective_scan_fn,flops_selective_scan_ref
+
+
+def build_vssm_model(config, **kwargs):
+    model_type = config.MODEL.TYPE
+    if model_type in ["vssm"]:
+        model = VSSM(
+            patch_size=config.MODEL.VSSM.PATCH_SIZE, 
+            in_chans=config.MODEL.VSSM.IN_CHANS, 
+            num_classes=config.MODEL.NUM_CLASSES, 
+            depths=config.MODEL.VSSM.DEPTHS, 
+            dims=config.MODEL.VSSM.EMBED_DIM, 
+            # ===================
+            ssm_d_state=config.MODEL.VSSM.SSM_D_STATE,
+            ssm_ratio=config.MODEL.VSSM.SSM_RATIO,
+            ssm_rank_ratio=config.MODEL.VSSM.SSM_RANK_RATIO,
+            ssm_dt_rank=("auto" if config.MODEL.VSSM.SSM_DT_RANK == "auto" else int(config.MODEL.VSSM.SSM_DT_RANK)),
+            ssm_act_layer=config.MODEL.VSSM.SSM_ACT_LAYER,
+            ssm_conv=config.MODEL.VSSM.SSM_CONV,
+            ssm_conv_bias=config.MODEL.VSSM.SSM_CONV_BIAS,
+            ssm_drop_rate=config.MODEL.VSSM.SSM_DROP_RATE,
+            ssm_init=config.MODEL.VSSM.SSM_INIT,
+            forward_type=config.MODEL.VSSM.SSM_FORWARDTYPE,
+            # ===================
+            mlp_ratio=config.MODEL.VSSM.MLP_RATIO,
+            mlp_act_layer=config.MODEL.VSSM.MLP_ACT_LAYER,
+            mlp_drop_rate=config.MODEL.VSSM.MLP_DROP_RATE,
+            # ===================
+            drop_path_rate=config.MODEL.DROP_PATH_RATE,
+            patch_norm=config.MODEL.VSSM.PATCH_NORM,
+            norm_layer=config.MODEL.VSSM.NORM_LAYER,
+            downsample_version=config.MODEL.VSSM.DOWNSAMPLE,
+            patchembed_version=config.MODEL.VSSM.PATCHEMBED,
+            gmlp=config.MODEL.VSSM.GMLP,
+            use_checkpoint=config.TRAIN.USE_CHECKPOINT,
+            # ===================
+            posembed=config.MODEL.VSSM.POSEMBED,
+            imgsize=config.DATA.IMG_SIZE,
+        )
+        return model
+
+    return None
+
+
+def build_model(config, is_pretrain=False):
+    model = None
+    if model is None:
+        model = build_vssm_model(config, is_pretrain)
+    return model
+
+
+
+

rscd/models/backbones/lib_mamba/csm_triton.py

@@ -0,0 +1,644 @@
+import torch
+import warnings
+
+WITH_TRITON = True
+# WITH_TRITON = False
+try:
+    import triton
+    import triton.language as tl
+except:
+    WITH_TRITON = False
+    warnings.warn("Triton not installed, fall back to pytorch implements.")
+
+# to make sure cached_property can be loaded for triton
+if WITH_TRITON:
+    try:
+        from functools import cached_property
+    except:
+        warnings.warn("if you are using py37, add this line to functools.py: "
+            "cached_property = lambda func: property(lru_cache()(func))")
+
+# torch implementation ========================================
+def cross_scan_fwd(x: torch.Tensor, in_channel_first=True, out_channel_first=True, scans=0):
+    if in_channel_first:
+        B, C, H, W = x.shape
+        if scans == 0:
+            y = x.new_empty((B, 4, C, H * W))
+            y[:, 0, :, :] = x.flatten(2, 3)
+            y[:, 1, :, :] = x.transpose(dim0=2, dim1=3).flatten(2, 3)
+            y[:, 2:4, :, :] = torch.flip(y[:, 0:2, :, :], dims=[-1])
+        elif scans == 1:
+            y = x.view(B, 1, C, H * W).repeat(1, 4, 1, 1)
+        elif scans == 2:
+            y = x.view(B, 1, C, H * W).repeat(1, 2, 1, 1)
+            y = torch.cat([y, y.flip(dims=[-1])], dim=1)
+    else:
+        B, H, W, C = x.shape
+        if scans == 0:
+            y = x.new_empty((B, H * W, 4, C))
+            y[:, :, 0, :] = x.flatten(1, 2)
+            y[:, :, 1, :] = x.transpose(dim0=1, dim1=2).flatten(1, 2)
+            y[:, :, 2:4, :] = torch.flip(y[:, :, 0:2, :], dims=[1])
+        elif scans == 1:
+            y = x.view(B, H * W, 1, C).repeat(1, 1, 4, 1)
+        elif scans == 2:
+            y = x.view(B, H * W, 1, C).repeat(1, 1, 2, 1)
+            y = torch.cat([y, y.flip(dims=[1])], dim=2)
+
+    if in_channel_first and (not out_channel_first):
+        y = y.permute(0, 3, 1, 2).contiguous()
+    elif (not in_channel_first) and out_channel_first:
+        y = y.permute(0, 2, 3, 1).contiguous()
+
+    return y
+
+
+def cross_merge_fwd(y: torch.Tensor, in_channel_first=True, out_channel_first=True, scans=0):
+    if out_channel_first:
+        B, K, D, H, W = y.shape
+        y = y.view(B, K, D, -1)
+        if scans == 0:
+            y = y[:, 0:2] + y[:, 2:4].flip(dims=[-1]).view(B, 2, D, -1)
+            y = y[:, 0] + y[:, 1].view(B, -1, W, H).transpose(dim0=2, dim1=3).contiguous().view(B, D, -1)
+        elif scans == 1:
+            y = y.sum(1)
+        elif scans == 2:
+            y = y[:, 0:2] + y[:, 2:4].flip(dims=[-1]).view(B, 2, D, -1)
+            y = y.sum(1)
+    else:
+        B, H, W, K, D = y.shape
+        y = y.view(B, -1, K, D)
+        if scans == 0:
+            y = y[:, :, 0:2] + y[:, :, 2:4].flip(dims=[1]).view(B, -1, 2, D)
+            y = y[:, :, 0] + y[:, :, 1].view(B, W, H, -1).transpose(dim0=1, dim1=2).contiguous().view(B, -1, D)        
+        elif scans == 1:
+            y = y.sum(2)
+        elif scans == 2:
+            y = y[:, :, 0:2] + y[:, :, 2:4].flip(dims=[1]).view(B, -1, 2, D)
+            y = y.sum(2)
+
+    if in_channel_first and (not out_channel_first):
+        y = y.permute(0, 2, 1).contiguous()
+    elif (not in_channel_first) and out_channel_first:
+        y = y.permute(0, 2, 1).contiguous()
+    
+    return y
+
+
+def cross_scan1b1_fwd(x: torch.Tensor, in_channel_first=True, out_channel_first=True, scans=0):
+    if in_channel_first:
+        B, _, C, H, W = x.shape
+        if scans == 0:
+            y = torch.stack([
+                x[:, 0].flatten(2, 3),
+                x[:, 1].transpose(dim0=2, dim1=3).flatten(2, 3),
+                torch.flip(x[:, 2].flatten(2, 3), dims=[-1]),
+                torch.flip(x[:, 3].transpose(dim0=2, dim1=3).flatten(2, 3), dims=[-1]),
+            ], dim=1)
+        elif scans == 1:
+            y = x.flatten(2, 3)
+        elif scans == 2:
+            y = torch.stack([
+                x[:, 0].flatten(2, 3),
+                x[:, 1].flatten(2, 3),
+                torch.flip(x[:, 2].flatten(2, 3), dims=[-1]),
+                torch.flip(x[:, 3].flatten(2, 3), dims=[-1]),
+            ], dim=1)
+    else:
+        B, H, W, _, C = x.shape
+        if scans == 0:
+            y = torch.stack([
+                x[:, :, :, 0].flatten(1, 2),
+                x[:, :, :, 1].transpose(dim0=1, dim1=2).flatten(1, 2),
+                torch.flip(x[:, :, :, 2].flatten(1, 2), dims=[1]),
+                torch.flip(x[:, :, :, 3].transpose(dim0=1, dim1=2).flatten(1, 2), dims=[1]),
+            ], dim=2)
+        elif scans == 1:
+            y = x.flatten(1, 2)
+        elif scans == 2:
+            y = torch.stack([
+                x[:, 0].flatten(1, 2),
+                x[:, 1].flatten(1, 2),
+                torch.flip(x[:, 2].flatten(1, 2), dims=[-1]),
+                torch.flip(x[:, 3].flatten(1, 2), dims=[-1]),
+            ], dim=2)
+
+    if in_channel_first and (not out_channel_first):
+        y = y.permute(0, 3, 1, 2).contiguous()
+    elif (not in_channel_first) and out_channel_first:
+        y = y.permute(0, 2, 3, 1).contiguous()
+
+    return y
+
+
+def cross_merge1b1_fwd(y: torch.Tensor, in_channel_first=True, out_channel_first=True, scans=0):
+    if out_channel_first:
+        B, K, D, H, W = y.shape
+        y = y.view(B, K, D, -1)
+        if scans == 0:
+            y = torch.stack([
+                y[:, 0],
+                y[:, 1].view(B, -1, W, H).transpose(dim0=2, dim1=3).flatten(2, 3),
+                torch.flip(y[:, 2], dims=[-1]),
+                torch.flip(y[:, 3].view(B, -1, W, H).transpose(dim0=2, dim1=3).flatten(2, 3), dims=[-1]),
+            ], dim=1)
+        elif scans == 1:
+            y = y
+        elif scans == 2:
+            y = torch.stack([
+                y[:, 0],
+                y[:, 1],
+                torch.flip(y[:, 2], dims=[-1]),
+                torch.flip(y[:, 3], dims=[-1]),
+            ], dim=1)
+    else:
+        B, H, W, _, D = y.shape
+        y = y.view(B, -1, K, D)
+        if scans == 0:
+            y = torch.stack([
+                y[:, :, 0],
+                y[:, :, 1].view(B, W, H, -1).transpose(dim0=1, dim1=2).flatten(1, 2),
+                torch.flip(y[:, :, 2], dims=[1]),
+                torch.flip(y[:, :, 3].view(B, W, H, -1).transpose(dim0=1, dim1=2).flatten(1, 2), dims=[1]),
+            ], dim=2)
+        elif scans == 1:
+            y = y
+        elif scans == 2:
+            y = torch.stack([
+                y[:, :, 0],
+                y[:, :, 1],
+                torch.flip(y[:, :, 2], dims=[1]),
+                torch.flip(y[:, :, 3], dims=[1]),
+            ], dim=2)
+
+    if out_channel_first and (not in_channel_first):
+        y = y.permute(0, 3, 1, 2).contiguous()
+    elif (not out_channel_first) and in_channel_first:
+        y = y.permute(0, 2, 3, 1).contiguous()
+
+    return y
+
+
+class CrossScanF(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=0):
+        # x: (B, C, H, W) | (B, H, W, C) | (B, 4, C, H, W) | (B, H, W, 4, C)
+        # y: (B, 4, C, H * W) | (B, H * W, 4, C)
+        ctx.in_channel_first = in_channel_first
+        ctx.out_channel_first = out_channel_first
+        ctx.one_by_one = one_by_one
+        ctx.scans = scans
+
+        if one_by_one:
+            B, K, C, H, W = x.shape
+            if not in_channel_first:
+                B, H, W, K, C = x.shape
+        else:
+            B, C, H, W = x.shape
+            if not in_channel_first:
+                B, H, W, C = x.shape
+        ctx.shape = (B, C, H, W)
+
+        _fn = cross_scan1b1_fwd if one_by_one else cross_scan_fwd
+        y = _fn(x, in_channel_first, out_channel_first, scans)
+
+        return y
+    
+    @staticmethod
+    def backward(ctx, ys: torch.Tensor):
+        # out: (b, k, d, l)
+        in_channel_first = ctx.in_channel_first
+        out_channel_first = ctx.out_channel_first
+        one_by_one = ctx.one_by_one
+        scans = ctx.scans
+        B, C, H, W = ctx.shape
+
+        ys = ys.view(B, -1, C, H, W) if out_channel_first else ys.view(B, H, W, -1, C)
+        _fn = cross_merge1b1_fwd if one_by_one else cross_merge_fwd
+        y = _fn(ys, in_channel_first, out_channel_first, scans)
+        
+        if one_by_one:
+            y = y.view(B, 4, -1, H, W) if in_channel_first else y.view(B, H, W, 4, -1)
+        else:
+            y = y.view(B, -1, H, W) if in_channel_first else y.view(B, H, W, -1)
+
+        return y, None, None, None, None
+
+
+class CrossMergeF(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, ys: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=0):
+        # x: (B, C, H, W) | (B, H, W, C) | (B, 4, C, H, W) | (B, H, W, 4, C)
+        # y: (B, 4, C, H * W) | (B, H * W, 4, C)
+        ctx.in_channel_first = in_channel_first
+        ctx.out_channel_first = out_channel_first
+        ctx.one_by_one = one_by_one
+        ctx.scans = scans
+
+        B, K, C, H, W = ys.shape
+        if not out_channel_first:
+            B, H, W, K, C = ys.shape
+        ctx.shape = (B, C, H, W)
+        
+        _fn = cross_merge1b1_fwd if one_by_one else cross_merge_fwd
+        y = _fn(ys, in_channel_first, out_channel_first, scans)
+
+        return y
+    
+    @staticmethod
+    def backward(ctx, x: torch.Tensor):
+        # B, D, L = x.shape
+        # out: (b, k, d, h, w)
+        in_channel_first = ctx.in_channel_first
+        out_channel_first = ctx.out_channel_first
+        one_by_one = ctx.one_by_one
+        scans = ctx.scans
+        B, C, H, W = ctx.shape
+    
+        if not one_by_one:
+            if in_channel_first:
+                x = x.view(B, C, H, W)
+            else:
+                x = x.view(B, H, W, C)
+        else:
+            if in_channel_first:
+                x = x.view(B, 4, C, H, W)
+            else:
+                x = x.view(B, H, W, 4, C)   
+                     
+        _fn = cross_scan1b1_fwd if one_by_one else cross_scan_fwd
+        x = _fn(x, in_channel_first, out_channel_first, scans)
+        x = x.view(B, 4, C, H, W) if out_channel_first else x.view(B, H, W, 4, C)
+    
+        return x, None, None, None, None
+
+
+# triton implements ========================================
+
+@triton.jit
+def triton_cross_scan_flex(
+    x, # (B, C, H, W) | (B, H, W, C) | (B, 4, C, H, W) | (B, H, W, 4, C)
+    y, # (B, 4, C, H, W) | (B, H, W, 4, C)
+    x_layout: tl.constexpr,
+    y_layout: tl.constexpr,
+    operation: tl.constexpr,
+    onebyone: tl.constexpr,
+    scans: tl.constexpr,
+    BC: tl.constexpr,
+    BH: tl.constexpr,
+    BW: tl.constexpr,
+    DC: tl.constexpr,
+    DH: tl.constexpr,
+    DW: tl.constexpr,
+    NH: tl.constexpr,
+    NW: tl.constexpr,
+):
+    # x_layout = 0
+    # y_layout = 1 # 0 BCHW, 1 BHWC
+    # operation = 0 # 0 scan, 1 merge
+    # onebyone = 0 # 0 false, 1 true
+    # scans = 0 # 0 cross scan, 1 unidirectional, 2 bidirectional
+
+    i_hw, i_c, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_h, i_w = (i_hw // NW), (i_hw % NW)
+    _mask_h = (i_h * BH + tl.arange(0, BH)) < DH
+    _mask_w = (i_w * BW + tl.arange(0, BW)) < DW
+    _mask_hw = _mask_h[:, None] & _mask_w[None, :]
+    _for_C = min(DC - i_c * BC, BC)
+
+    HWRoute0 = i_h * BH * DW  + tl.arange(0, BH)[:, None] * DW + i_w * BW + tl.arange(0, BW)[None, :]
+    HWRoute1 = i_w * BW * DH + tl.arange(0, BW)[None, :] * DH + i_h * BH + tl.arange(0, BH)[:, None]  # trans
+    HWRoute2 = (NH - i_h - 1) * BH * DW  + (BH - 1 - tl.arange(0, BH)[:, None]) * DW + (NW - i_w - 1) * BW + (BW - 1 - tl.arange(0, BW)[None, :]) + (DH - NH * BH) * DW + (DW - NW * BW) # flip
+    HWRoute3 = (NW - i_w - 1) * BW * DH  + (BW - 1 - tl.arange(0, BW)[None, :]) * DH + (NH - i_h - 1) * BH + (BH - 1 - tl.arange(0, BH)[:, None]) + (DH - NH * BH) + (DW - NW * BW) * DH  # trans + flip
+
+    if scans == 1:
+        HWRoute1 = HWRoute0
+        HWRoute2 = HWRoute0
+        HWRoute3 = HWRoute0
+    elif scans == 2:
+        HWRoute1 = HWRoute0
+        HWRoute3 = HWRoute2        
+
+    _tmp1 = DC * DH * DW
+
+    y_ptr_base = y + i_b * 4 * _tmp1 + (i_c * BC * DH * DW if y_layout == 0 else i_c * BC)
+    if y_layout == 0:
+        p_y1 = y_ptr_base + HWRoute0
+        p_y2 = y_ptr_base + _tmp1 + HWRoute1
+        p_y3 = y_ptr_base + 2 * _tmp1 + HWRoute2
+        p_y4 = y_ptr_base + 3 * _tmp1 + HWRoute3
+    else:
+        p_y1 = y_ptr_base + HWRoute0 * 4 * DC
+        p_y2 = y_ptr_base + DC + HWRoute1 * 4 * DC
+        p_y3 = y_ptr_base + 2 * DC + HWRoute2 * 4 * DC
+        p_y4 = y_ptr_base + 3 * DC + HWRoute3 * 4 * DC       
+    
+    if onebyone == 0:
+        x_ptr_base = x + i_b * _tmp1 + (i_c * BC * DH * DW if x_layout == 0 else i_c * BC)
+        if x_layout == 0:
+            p_x = x_ptr_base + HWRoute0
+        else:
+            p_x = x_ptr_base + HWRoute0 * DC
+
+        if operation == 0:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                _x = tl.load(p_x + _idx_x, mask=_mask_hw)
+                tl.store(p_y1 + _idx_y, _x, mask=_mask_hw)
+                tl.store(p_y2 + _idx_y, _x, mask=_mask_hw)
+                tl.store(p_y3 + _idx_y, _x, mask=_mask_hw)
+                tl.store(p_y4 + _idx_y, _x, mask=_mask_hw)
+        elif operation == 1:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                _y1 = tl.load(p_y1 + _idx_y, mask=_mask_hw)
+                _y2 = tl.load(p_y2 + _idx_y, mask=_mask_hw)
+                _y3 = tl.load(p_y3 + _idx_y, mask=_mask_hw)
+                _y4 = tl.load(p_y4 + _idx_y, mask=_mask_hw)
+                tl.store(p_x + _idx_x, _y1 + _y2 + _y3 + _y4, mask=_mask_hw)
+
+    else:
+        x_ptr_base = x + i_b * 4 * _tmp1 + (i_c * BC * DH * DW if x_layout == 0 else i_c * BC)
+        if x_layout == 0:
+            p_x1 = x_ptr_base + HWRoute0
+            p_x2 = p_x1 + _tmp1
+            p_x3 = p_x2 + _tmp1
+            p_x4 = p_x3 + _tmp1  
+        else:
+            p_x1 = x_ptr_base + HWRoute0 * 4 * DC
+            p_x2 = p_x1 + DC
+            p_x3 = p_x2 + DC
+            p_x4 = p_x3 + DC        
+    
+        if operation == 0:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                tl.store(p_y1 + _idx_y, tl.load(p_x1 + _idx_x, mask=_mask_hw), mask=_mask_hw)
+                tl.store(p_y2 + _idx_y, tl.load(p_x2 + _idx_x, mask=_mask_hw), mask=_mask_hw)
+                tl.store(p_y3 + _idx_y, tl.load(p_x3 + _idx_x, mask=_mask_hw), mask=_mask_hw)
+                tl.store(p_y4 + _idx_y, tl.load(p_x4 + _idx_x, mask=_mask_hw), mask=_mask_hw)
+        else:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                tl.store(p_x1 + _idx_x, tl.load(p_y1 + _idx_y), mask=_mask_hw)
+                tl.store(p_x2 + _idx_x, tl.load(p_y2 + _idx_y), mask=_mask_hw)
+                tl.store(p_x3 + _idx_x, tl.load(p_y3 + _idx_y), mask=_mask_hw)
+                tl.store(p_x4 + _idx_x, tl.load(p_y4 + _idx_y), mask=_mask_hw)
+
+
+class CrossScanTritonF(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=0):
+        if one_by_one:
+            if in_channel_first:
+                B, _, C, H, W = x.shape
+            else:
+                B, H, W, _, C = x.shape
+        else:
+            if in_channel_first:
+                B, C, H, W = x.shape
+            else:
+                B, H, W, C = x.shape
+        B, C, H, W = int(B), int(C), int(H), int(W)
+        BC, BH, BW = 1, 32, 32
+        NH, NW, NC = triton.cdiv(H, BH), triton.cdiv(W, BW), triton.cdiv(C, BC)
+        
+        ctx.in_channel_first = in_channel_first
+        ctx.out_channel_first = out_channel_first
+        ctx.one_by_one = one_by_one
+        ctx.scans = scans
+        ctx.shape = (B, C, H, W)
+        ctx.triton_shape = (BC, BH, BW, NC, NH, NW)
+
+        y = x.new_empty((B, 4, C, H * W)) if out_channel_first else x.new_empty((B, H * W, 4, C))
+        triton_cross_scan_flex[(NH * NW, NC, B)](
+            x.contiguous(), y, 
+            (0 if in_channel_first else 1), (0 if out_channel_first else 1), 0, (0 if not one_by_one else 1), scans, 
+            BC, BH, BW, C, H, W, NH, NW
+        )
+        return y
+        
+    @staticmethod
+    def backward(ctx, y: torch.Tensor):
+        in_channel_first = ctx.in_channel_first
+        out_channel_first = ctx.out_channel_first
+        one_by_one = ctx.one_by_one
+        scans = ctx.scans
+        B, C, H, W = ctx.shape
+        BC, BH, BW, NC, NH, NW = ctx.triton_shape
+        if one_by_one:
+            x = y.new_empty((B, 4, C, H, W)) if in_channel_first else y.new_empty((B, H, W, 4, C))
+        else:
+            x = y.new_empty((B, C, H, W)) if in_channel_first else y.new_empty((B, H, W, C))
+        
+        triton_cross_scan_flex[(NH * NW, NC, B)](
+            x, y.contiguous(), 
+            (0 if in_channel_first else 1), (0 if out_channel_first else 1), 1, (0 if not one_by_one else 1), scans,
+            BC, BH, BW, C, H, W, NH, NW
+        )
+        return x, None, None, None, None
+
+
+class CrossMergeTritonF(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, y: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=0):
+        if out_channel_first:
+            B, _, C, H, W = y.shape
+        else:
+            B, H, W, _, C = y.shape
+        B, C, H, W = int(B), int(C), int(H), int(W)
+        BC, BH, BW = 1, 32, 32
+        NH, NW, NC = triton.cdiv(H, BH), triton.cdiv(W, BW), triton.cdiv(C, BC)
+        ctx.in_channel_first = in_channel_first
+        ctx.out_channel_first = out_channel_first
+        ctx.one_by_one = one_by_one
+        ctx.scans = scans
+        ctx.shape = (B, C, H, W)
+        ctx.triton_shape = (BC, BH, BW, NC, NH, NW)
+        if one_by_one:
+            x = y.new_empty((B, 4, C, H * W)) if in_channel_first else y.new_empty((B, H * W, 4, C))
+        else:
+            x = y.new_empty((B, C, H * W)) if in_channel_first else y.new_empty((B, H * W, C))
+        triton_cross_scan_flex[(NH * NW, NC, B)](
+            x, y.contiguous(), 
+            (0 if in_channel_first else 1), (0 if out_channel_first else 1), 1, (0 if not one_by_one else 1), scans,
+            BC, BH, BW, C, H, W, NH, NW
+        )
+        return x
+        
+    @staticmethod
+    def backward(ctx, x: torch.Tensor):
+        in_channel_first = ctx.in_channel_first
+        out_channel_first = ctx.out_channel_first
+        one_by_one = ctx.one_by_one
+        scans = ctx.scans
+        B, C, H, W = ctx.shape
+        BC, BH, BW, NC, NH, NW = ctx.triton_shape
+        y = x.new_empty((B, 4, C, H, W)) if out_channel_first else x.new_empty((B, H, W, 4, C))
+        triton_cross_scan_flex[(NH * NW, NC, B)](
+            x.contiguous(), y, 
+            (0 if in_channel_first else 1), (0 if out_channel_first else 1), 0, (0 if not one_by_one else 1), scans,
+            BC, BH, BW, C, H, W, NH, NW
+        )
+        return y, None, None, None, None, None
+
+
+# @torch.compile(options={"triton.cudagraphs": True}, fullgraph=True)
+def cross_scan_fn(x: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=0, force_torch=False):
+    # x: (B, C, H, W) | (B, H, W, C) | (B, 4, C, H, W) | (B, H, W, 4, C)
+    # y: (B, 4, C, L) | (B, L, 4, C)
+    # scans: 0: cross scan; 1 unidirectional; 2: bidirectional;
+    CSF = CrossScanTritonF if WITH_TRITON and x.is_cuda and (not force_torch) else CrossScanF
+    return CSF.apply(x, in_channel_first, out_channel_first, one_by_one, scans)
+
+
+# @torch.compile(options={"triton.cudagraphs": True}, fullgraph=True)
+def cross_merge_fn(y: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=0, force_torch=False):
+    # y: (B, 4, C, L) | (B, L, 4, C)
+    # x: (B, C, H * W) | (B, H * W, C) | (B, 4, C, H * W) | (B, H * W, 4, C)
+    # scans: 0: cross scan; 1 unidirectional; 2: bidirectional;
+    CMF = CrossMergeTritonF if WITH_TRITON and y.is_cuda and (not force_torch) else CrossMergeF
+    return CMF.apply(y, in_channel_first, out_channel_first, one_by_one, scans)
+
+
+# checks =================================================================
+
+class CHECK:
+    def check_csm_triton():
+        B, C, H, W = 2, 192, 56, 57
+        dtype=torch.float16
+        dtype=torch.float32
+        x = torch.randn((B, C, H, W), dtype=dtype, device=torch.device("cuda")).requires_grad_(True)
+        y = torch.randn((B, 4, C, H, W), dtype=dtype, device=torch.device("cuda")).requires_grad_(True)
+        x1 = x.clone().detach().requires_grad_(True)
+        y1 = y.clone().detach().requires_grad_(True)
+
+        def cross_scan(x: torch.Tensor):
+            B, C, H, W = x.shape
+            L = H * W
+            xs = torch.stack([
+                x.view(B, C, L),
+                torch.transpose(x, dim0=2, dim1=3).contiguous().view(B, C, L),
+                torch.flip(x.contiguous().view(B, C, L), dims=[-1]),
+                torch.flip(torch.transpose(x, dim0=2, dim1=3).contiguous().view(B, C, L), dims=[-1]),
+            ], dim=1).view(B, 4, C, L)
+            return xs
+        
+        def cross_merge(out_y: torch.Tensor):
+            B, K, D, H, W = out_y.shape
+            L = H * W
+            out_y = out_y.view(B, K, D, L)
+            inv_y = torch.flip(out_y[:, 2:4], dims=[-1]).view(B, 2, -1, L)
+            wh_y = torch.transpose(out_y[:, 1].view(B, -1, W, H), dim0=2, dim1=3).contiguous().view(B, -1, L)
+            invwh_y = torch.transpose(inv_y[:, 1].view(B, -1, W, H), dim0=2, dim1=3).contiguous().view(B, -1, L)
+            y = out_y[:, 0] + inv_y[:, 0] + wh_y + invwh_y
+            return y
+
+        def cross_scan_1b1(x: torch.Tensor):
+            B, K, C, H, W = x.shape
+            L = H * W
+            xs = torch.stack([
+                x[:, 0].view(B, C, L),
+                torch.transpose(x[:, 1], dim0=2, dim1=3).contiguous().view(B, C, L),
+                torch.flip(x[:, 2].contiguous().view(B, C, L), dims=[-1]),
+                torch.flip(torch.transpose(x[:, 3], dim0=2, dim1=3).contiguous().view(B, C, L), dims=[-1]),
+            ], dim=1).view(B, 4, C, L)
+            return xs
+        
+        def unidi_scan(x):
+            B, C, H, W = x.shape
+            x = x.view(B, 1, C, H * W).repeat(1, 4, 1, 1)
+            return x
+        
+        def unidi_merge(ys):
+            B, K, C, H, W = ys.shape
+            return ys.view(B, 4, -1, H * W).sum(1)
+
+        def bidi_scan(x):
+            B, C, H, W = x.shape
+            x = x.view(B, 1, C, H * W).repeat(1, 2, 1, 1)
+            x = torch.cat([x, x.flip(dims=[-1])], dim=1)
+            return x
+        
+        def bidi_merge(ys):
+            B, K, D, H, W = ys.shape
+            ys = ys.view(B, K, D, -1)
+            ys = ys[:, 0:2] + ys[:, 2:4].flip(dims=[-1]).view(B, 2, D, -1)
+            return ys.contiguous().sum(1)
+
+        if True:
+            res0 = triton.testing.do_bench(lambda :cross_scan(x))
+            res1 = triton.testing.do_bench(lambda :cross_scan_fn(x, True, True, False))
+            # res2 = triton.testing.do_bench(lambda :CrossScanTriton.apply(x))
+            res3 = triton.testing.do_bench(lambda :cross_merge(y))
+            res4 = triton.testing.do_bench(lambda :cross_merge_fn(y, True, True, False))
+            # res5 = triton.testing.do_bench(lambda :CrossMergeTriton.apply(y))
+            # print(res0, res1, res2, res3, res4, res5)
+            print(res0, res1, res3, res4)
+            res0 = triton.testing.do_bench(lambda :cross_scan(x).sum().backward())
+            res1 = triton.testing.do_bench(lambda :cross_scan_fn(x, True, True, False).sum().backward())
+            # res2 = triton.testing.do_bench(lambda :CrossScanTriton.apply(x).sum().backward())
+            res3 = triton.testing.do_bench(lambda :cross_merge(y).sum().backward())
+            res4 = triton.testing.do_bench(lambda :cross_merge_fn(y, True, True, False).sum().backward())
+            # res5 = triton.testing.do_bench(lambda :CrossMergeTriton.apply(y).sum().backward())
+            # print(res0, res1, res2, res3, res4, res5)
+            print(res0, res1, res3, res4)
+
+        print("test cross scan")
+        for (cs0, cm0, cs1, cm1) in [
+            # channel_first -> channel_first
+            (cross_scan, cross_merge, cross_scan_fn, cross_merge_fn),
+            (unidi_scan, unidi_merge, lambda x: cross_scan_fn(x, scans=1), lambda x: cross_merge_fn(x, scans=1)),
+            (bidi_scan, bidi_merge, lambda x: cross_scan_fn(x, scans=2), lambda x: cross_merge_fn(x, scans=2)),
+            
+            # flex: BLC->BCL; BCL->BLC; BLC->BLC;
+            (cross_scan, cross_merge, lambda x: cross_scan_fn(x.permute(0, 2, 3, 1), in_channel_first=False), lambda x: cross_merge_fn(x, in_channel_first=False).permute(0, 2, 1)),
+            (cross_scan, cross_merge, lambda x: cross_scan_fn(x, out_channel_first=False).permute(0, 2, 3, 1), lambda x: cross_merge_fn(x.permute(0, 3, 4, 1, 2), out_channel_first=False)),
+            (cross_scan, cross_merge, lambda x: cross_scan_fn(x.permute(0, 2, 3, 1), in_channel_first=False, out_channel_first=False).permute(0, 2, 3, 1), lambda x: cross_merge_fn(x.permute(0, 3, 4, 1, 2), in_channel_first=False, out_channel_first=False).permute(0, 2, 1)),
+            
+            # previous
+            # (cross_scan, cross_merge, lambda x: CrossScanTriton.apply(x), lambda x: CrossMergeTriton.apply(x)),
+            # (unidi_scan, unidi_merge, lambda x: getCSM(1)[0].apply(x), lambda x: getCSM(1)[1].apply(x)),
+            # (bidi_scan, bidi_merge, lambda x: getCSM(2)[0].apply(x), lambda x: getCSM(2)[1].apply(x)),
+        ]:
+            x.grad, x1.grad, y.grad, y1.grad = None, None, None, None
+            o0 = cs0(x)
+            o1 = cs1(x1)
+            o0.backward(y.view(B, 4, C, H * W))
+            o1.backward(y.view(B, 4, C, H * W))
+            print((o0 - o1).abs().max())
+            print((x.grad - x1.grad).abs().max())
+            o0 = cm0(y)
+            o1 = cm1(y1)
+            o0.backward(x.view(B, C, H * W))
+            o1.backward(x.view(B, C, H * W))
+            print((o0 - o1).abs().max())
+            print((y.grad - y1.grad).abs().max())
+            x.grad, x1.grad, y.grad, y1.grad = None, None, None, None
+            print("===============", flush=True)
+
+        print("test cross scan one by one")
+        for (cs0, cs1) in [
+            (cross_scan_1b1, lambda x: cross_scan_fn(x, one_by_one=True)),
+            # (cross_scan_1b1, lambda x: CrossScanTriton1b1.apply(x)),
+        ]:
+            o0 = cs0(y)
+            o1 = cs1(y1)
+            o0.backward(y.view(B, 4, C, H * W))
+            o1.backward(y.view(B, 4, C, H * W))
+            print((o0 - o1).abs().max())
+            print((y.grad - y1.grad).abs().max())
+            x.grad, x1.grad, y.grad, y1.grad = None, None, None, None
+            print("===============", flush=True)
+
+
+if __name__ == "__main__":
+    CHECK.check_csm_triton()
+
+
+
+

rscd/models/backbones/lib_mamba/csm_tritonk2.py

@@ -0,0 +1,899 @@
+import torch
+import warnings
+import os
+os.environ["TRITON_INTERPRET"] = "1"
+
+WITH_TRITON = True
+# WITH_TRITON = False
+try:
+    import triton
+    import triton.language as tl
+except:
+    WITH_TRITON = False
+    warnings.warn("Triton not installed, fall back to pytorch implements.")
+
+# to make sure cached_property can be loaded for triton
+if WITH_TRITON:
+    try:
+        from functools import cached_property
+    except:
+        warnings.warn("if you are using py37, add this line to functools.py: "
+            "cached_property = lambda func: property(lru_cache()(func))")
+
+# torch implementation ========================================
+def cross_scan_fwd(x: torch.Tensor, in_channel_first=True, out_channel_first=True, scans=2):
+    if in_channel_first:
+        B, C, H, W = x.shape
+        if scans == 0:
+            y = x.new_empty((B, 4, C, H * W))
+            y[:, 0, :, :] = x.flatten(2, 3)
+            y[:, 1, :, :] = x.transpose(dim0=2, dim1=3).flatten(2, 3)
+            y[:, 2:4, :, :] = torch.flip(y[:, 0:2, :, :], dims=[-1])
+        elif scans == 1:
+            y = x.view(B, 1, C, H * W).repeat(1, 2, 1, 1)
+        elif scans == 2:
+            y = x.view(B, 1, C, H * W)
+            y = torch.cat([y, y.flip(dims=[-1])], dim=1)
+    else:
+        B, H, W, C = x.shape
+        if scans == 0:
+            y = x.new_empty((B, H * W, 4, C))
+            y[:, :, 0, :] = x.flatten(1, 2)
+            y[:, :, 1, :] = x.transpose(dim0=1, dim1=2).flatten(1, 2)
+            y[:, :, 2:4, :] = torch.flip(y[:, :, 0:2, :], dims=[1])
+        elif scans == 1:
+            y = x.view(B, H * W, 1, C).repeat(1, 1, 2, 1)
+        elif scans == 2:
+            y = x.view(B, H * W, 1, C)
+            y = torch.cat([y, y.flip(dims=[1])], dim=2)
+
+    if in_channel_first and (not out_channel_first):
+        y = y.permute(0, 3, 1, 2).contiguous()
+    elif (not in_channel_first) and out_channel_first:
+        y = y.permute(0, 2, 3, 1).contiguous()
+
+    return y
+
+
+def cross_merge_fwd(y: torch.Tensor, in_channel_first=True, out_channel_first=True, scans=2):
+    if out_channel_first:
+        B, K, D, H, W = y.shape
+        y = y.view(B, K, D, -1)
+        if scans == 0:
+            y = y[:, 0:2] + y[:, 2:4].flip(dims=[-1]).view(B, 2, D, -1)
+            y = y[:, 0] + y[:, 1].view(B, -1, W, H).transpose(dim0=2, dim1=3).contiguous().view(B, D, -1)
+        elif scans == 1:
+            y = y.sum(1)
+        elif scans == 2:
+            y = y[:, 0] + y[:, 1].flip(dims=[-1]).view(B, 1, D, -1)
+            y = y.sum(1)
+    else:
+        B, H, W, K, D = y.shape
+        y = y.view(B, -1, K, D)
+        if scans == 0:
+            y = y[:, :, 0:2] + y[:, :, 2:4].flip(dims=[1]).view(B, -1, 2, D)
+            y = y[:, :, 0] + y[:, :, 1].view(B, W, H, -1).transpose(dim0=1, dim1=2).contiguous().view(B, -1, D)
+        elif scans == 1:
+            y = y.sum(2)
+        elif scans == 2:
+            y = y[:, :, 0] + y[:, :, 1].flip(dims=[1]).view(B, -1, 1, D)
+            y = y.sum(2)
+
+    if in_channel_first and (not out_channel_first):
+        y = y.permute(0, 2, 1).contiguous()
+    elif (not in_channel_first) and out_channel_first:
+        y = y.permute(0, 2, 1).contiguous()
+
+    return y
+
+
+def cross_scan1b1_fwd(x: torch.Tensor, in_channel_first=True, out_channel_first=True, scans=2):
+    if in_channel_first:
+        B, _, C, H, W = x.shape
+        if scans == 0:
+            y = torch.stack([
+                x[:, 0].flatten(2, 3),
+                x[:, 1].transpose(dim0=2, dim1=3).flatten(2, 3),
+                torch.flip(x[:, 2].flatten(2, 3), dims=[-1]),
+                torch.flip(x[:, 3].transpose(dim0=2, dim1=3).flatten(2, 3), dims=[-1]),
+            ], dim=1)
+        elif scans == 1:
+            y = x.flatten(2, 3)
+        elif scans == 2:
+            y = torch.stack([
+                x[:, 0].flatten(2, 3),
+                x[:, 1].flatten(2, 3),
+                torch.flip(x[:, 2].flatten(2, 3), dims=[-1]),
+                torch.flip(x[:, 3].flatten(2, 3), dims=[-1]),
+            ], dim=1)
+    else:
+        B, H, W, _, C = x.shape
+        if scans == 0:
+            y = torch.stack([
+                x[:, :, :, 0].flatten(1, 2),
+                x[:, :, :, 1].transpose(dim0=1, dim1=2).flatten(1, 2),
+                torch.flip(x[:, :, :, 2].flatten(1, 2), dims=[1]),
+                torch.flip(x[:, :, :, 3].transpose(dim0=1, dim1=2).flatten(1, 2), dims=[1]),
+            ], dim=2)
+        elif scans == 1:
+            y = x.flatten(1, 2)
+        elif scans == 2:
+            y = torch.stack([
+                x[:, 0].flatten(1, 2),
+                x[:, 1].flatten(1, 2),
+                torch.flip(x[:, 2].flatten(1, 2), dims=[-1]),
+                torch.flip(x[:, 3].flatten(1, 2), dims=[-1]),
+            ], dim=2)
+
+    if in_channel_first and (not out_channel_first):
+        y = y.permute(0, 3, 1, 2).contiguous()
+    elif (not in_channel_first) and out_channel_first:
+        y = y.permute(0, 2, 3, 1).contiguous()
+
+    return y
+
+
+def cross_merge1b1_fwd(y: torch.Tensor, in_channel_first=True, out_channel_first=True, scans=2):
+    if out_channel_first:
+        B, K, D, H, W = y.shape
+        y = y.view(B, K, D, -1)
+        if scans == 0:
+            y = torch.stack([
+                y[:, 0],
+                y[:, 1].view(B, -1, W, H).transpose(dim0=2, dim1=3).flatten(2, 3),
+                torch.flip(y[:, 2], dims=[-1]),
+                torch.flip(y[:, 3].view(B, -1, W, H).transpose(dim0=2, dim1=3).flatten(2, 3), dims=[-1]),
+            ], dim=1)
+        elif scans == 1:
+            y = y
+        elif scans == 2:
+            y = torch.stack([
+                y[:, 0],
+                y[:, 1],
+                torch.flip(y[:, 2], dims=[-1]),
+                torch.flip(y[:, 3], dims=[-1]),
+            ], dim=1)
+    else:
+        B, H, W, _, D = y.shape
+        y = y.view(B, -1, 2, D)
+        if scans == 0:
+            y = torch.stack([
+                y[:, :, 0],
+                y[:, :, 1].view(B, W, H, -1).transpose(dim0=1, dim1=2).flatten(1, 2),
+                torch.flip(y[:, :, 2], dims=[1]),
+                torch.flip(y[:, :, 3].view(B, W, H, -1).transpose(dim0=1, dim1=2).flatten(1, 2), dims=[1]),
+            ], dim=2)
+        elif scans == 1:
+            y = y
+        elif scans == 2:
+            y = torch.stack([
+                y[:, :, 0],
+                y[:, :, 1],
+                torch.flip(y[:, :, 2], dims=[1]),
+                torch.flip(y[:, :, 3], dims=[1]),
+            ], dim=2)
+
+    if out_channel_first and (not in_channel_first):
+        y = y.permute(0, 3, 1, 2).contiguous()
+    elif (not out_channel_first) and in_channel_first:
+        y = y.permute(0, 2, 3, 1).contiguous()
+
+    return y
+
+class CrossScan(torch.nn.Module):
+    def __init__(self, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=2):
+        super(CrossScan, self).__init__()
+        self.in_channel_first = in_channel_first
+        self.out_channel_first = out_channel_first
+        self.one_by_one = one_by_one
+        self.scans = scans
+
+    def forward(self, x: torch.Tensor):
+        if self.one_by_one:
+            B, K, C, H, W = x.shape
+            if not self.in_channel_first:
+                B, H, W, K, C = x.shape
+        else:
+            B, C, H, W = x.shape
+            if not self.in_channel_first:
+                B, H, W, C = x.shape
+        self.shape = (B, C, H, W)
+
+        _fn = cross_scan1b1_fwd if self.one_by_one else cross_scan_fwd
+        y = _fn(x, self.in_channel_first, self.out_channel_first, self.scans)
+
+        return y
+
+    def backward(self, ys: torch.Tensor):
+        B, C, H, W = self.shape
+
+        ys = ys.view(B, -1, C, H, W) if self.out_channel_first else ys.view(B, H, W, -1, C)
+        _fn = cross_merge1b1_fwd if self.one_by_one else cross_merge_fwd
+        y = _fn(ys, self.in_channel_first, self.out_channel_first, self.scans)
+
+        if self.one_by_one:
+            y = y.view(B, 2, -1, H, W) if self.in_channel_first else y.view(B, H, W, 2, -1)
+        else:
+            y = y.view(B, -1, H, W) if self.in_channel_first else y.view(B, H, W, -1)
+
+        return y
+
+
+class CrossMerge(torch.nn.Module):
+    def __init__(self, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=2):
+        super(CrossMerge, self).__init__()
+        self.in_channel_first = in_channel_first
+        self.out_channel_first = out_channel_first
+        self.one_by_one = one_by_one
+        self.scans = scans
+
+    def forward(self, ys: torch.Tensor):
+        B, K, C, H, W = ys.shape
+        if not self.out_channel_first:
+            B, H, W, K, C = ys.shape
+        self.shape = (B, C, H, W)
+
+        _fn = cross_merge1b1_fwd if self.one_by_one else cross_merge_fwd
+        y = _fn(ys, self.in_channel_first, self.out_channel_first, self.scans)
+
+        return y
+
+    def backward(self, x: torch.Tensor):
+        B, C, H, W = self.shape
+
+        if not self.one_by_one:
+            if self.in_channel_first:
+                x = x.view(B, C, H, W)
+            else:
+                x = x.view(B, H, W, C)
+        else:
+            if self.in_channel_first:
+                x = x.view(B, 2, C, H, W)
+            else:
+                x = x.view(B, H, W, 2, C)
+
+        _fn = cross_scan1b1_fwd if self.one_by_one else cross_scan_fwd
+        x = _fn(x, self.in_channel_first, self.out_channel_first, self.scans)
+        x = x.view(B, 2, C, H, W) if self.out_channel_first else x.view(B, H, W, 2, C)
+
+        return x
+class CrossScanF(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=2):
+        # x: (B, C, H, W) | (B, H, W, C) | (B, 4, C, H, W) | (B, H, W, 2, C)
+        # y: (B, 2, C, H * W) | (B, H * W, 2, C)
+        ctx.in_channel_first = in_channel_first
+        ctx.out_channel_first = out_channel_first
+        ctx.one_by_one = one_by_one
+        ctx.scans = scans
+
+        if one_by_one:
+            B, K, C, H, W = x.shape
+            if not in_channel_first:
+                B, H, W, K, C = x.shape
+        else:
+            B, C, H, W = x.shape
+            if not in_channel_first:
+                B, H, W, C = x.shape
+        ctx.shape = (B, C, H, W)
+
+        _fn = cross_scan1b1_fwd if one_by_one else cross_scan_fwd
+        y = _fn(x, in_channel_first, out_channel_first, scans)
+
+        return y
+
+    @staticmethod
+    def backward(ctx, ys: torch.Tensor):
+        # out: (b, k, d, l)
+        in_channel_first = ctx.in_channel_first
+        out_channel_first = ctx.out_channel_first
+        one_by_one = ctx.one_by_one
+        scans = ctx.scans
+        B, C, H, W = ctx.shape
+
+        ys = ys.view(B, -1, C, H, W) if out_channel_first else ys.view(B, H, W, -1, C)
+        _fn = cross_merge1b1_fwd if one_by_one else cross_merge_fwd
+        y = _fn(ys, in_channel_first, out_channel_first, scans)
+
+        if one_by_one:
+            y = y.view(B, 2, -1, H, W) if in_channel_first else y.view(B, H, W, 2, -1)
+        else:
+            y = y.view(B, -1, H, W) if in_channel_first else y.view(B, H, W, -1)
+
+        return y, None, None, None, None
+
+
+class CrossMergeF(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, ys: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=2):
+        # x: (B, C, H, W) | (B, H, W, C) | (B, 2, C, H, W) | (B, H, W, 2, C)
+        # y: (B, 2, C, H * W) | (B, H * W, 4, C)
+        ctx.in_channel_first = in_channel_first
+        ctx.out_channel_first = out_channel_first
+        ctx.one_by_one = one_by_one
+        ctx.scans = scans
+
+        B, K, C, H, W = ys.shape
+        if not out_channel_first:
+            B, H, W, K, C = ys.shape
+        ctx.shape = (B, C, H, W)
+
+        _fn = cross_merge1b1_fwd if one_by_one else cross_merge_fwd
+        y = _fn(ys, in_channel_first, out_channel_first, scans)
+
+        return y
+
+    @staticmethod
+    def backward(ctx, x: torch.Tensor):
+        # B, D, L = x.shape
+        # out: (b, k, d, h, w)
+        in_channel_first = ctx.in_channel_first
+        out_channel_first = ctx.out_channel_first
+        one_by_one = ctx.one_by_one
+        scans = ctx.scans
+        B, C, H, W = ctx.shape
+
+        if not one_by_one:
+            if in_channel_first:
+                x = x.view(B, C, H, W)
+            else:
+                x = x.view(B, H, W, C)
+        else:
+            if in_channel_first:
+                x = x.view(B, 2, C, H, W)
+            else:
+                x = x.view(B, H, W, 2, C)
+
+        _fn = cross_scan1b1_fwd if one_by_one else cross_scan_fwd
+        x = _fn(x, in_channel_first, out_channel_first, scans)
+        x = x.view(B, 2, C, H, W) if out_channel_first else x.view(B, H, W, 2, C)
+
+        return x, None, None, None, None
+
+
+# triton implements ========================================
+
+@triton.jit
+def triton_cross_scan_flex_k2(
+    x, # (B, C, H, W) | (B, H, W, C) | (B, 4, C, H, W) | (B, H, W, 4, C)
+    y, # (B, 4, C, H, W) | (B, H, W, 4, C)
+    x_layout: tl.constexpr,
+    y_layout: tl.constexpr,
+    operation: tl.constexpr,
+    onebyone: tl.constexpr,
+    scans: tl.constexpr,
+    BC: tl.constexpr,
+    BH: tl.constexpr,
+    BW: tl.constexpr,
+    DC: tl.constexpr,
+    DH: tl.constexpr,
+    DW: tl.constexpr,
+    NH: tl.constexpr,
+    NW: tl.constexpr,
+):
+    # x_layout = 0
+    # y_layout = 1 # 0 BCHW, 1 BHWC
+    # operation = 0 # 0 scan, 1 merge
+    # onebyone = 0 # 0 false, 1 true
+    # scans = 0 # 0 cross scan, 1 unidirectional, 2 bidirectional
+
+    i_hw, i_c, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_h, i_w = (i_hw // NW), (i_hw % NW)
+    _mask_h = (i_h * BH + tl.arange(0, BH)) < DH
+    _mask_w = (i_w * BW + tl.arange(0, BW)) < DW
+    _mask_hw = _mask_h[:, None] & _mask_w[None, :]
+    _for_C = min(DC - i_c * BC, BC)
+
+    HWRoute0 = i_h * BH * DW  + tl.arange(0, BH)[:, None] * DW + i_w * BW + tl.arange(0, BW)[None, :]
+    # HWRoute1 = i_w * BW * DH + tl.arange(0, BW)[None, :] * DH + i_h * BH + tl.arange(0, BH)[:, None]  # trans
+    HWRoute2 = (NH - i_h - 1) * BH * DW  + (BH - 1 - tl.arange(0, BH)[:, None]) * DW + (NW - i_w - 1) * BW + (BW - 1 - tl.arange(0, BW)[None, :]) + (DH - NH * BH) * DW + (DW - NW * BW) # flip
+    # HWRoute3 = (NW - i_w - 1) * BW * DH  + (BW - 1 - tl.arange(0, BW)[None, :]) * DH + (NH - i_h - 1) * BH + (BH - 1 - tl.arange(0, BH)[:, None]) + (DH - NH * BH) + (DW - NW * BW) * DH  # trans + flip
+
+    if scans == 1:
+        HWRoute2 = HWRoute0
+
+
+    _tmp1 = DC * DH * DW
+
+    y_ptr_base = y + i_b * 4 * _tmp1 + (i_c * BC * DH * DW if y_layout == 0 else i_c * BC)
+    if y_layout == 0:
+        p_y1 = y_ptr_base + HWRoute0
+        # p_y2 = y_ptr_base + _tmp1 + HWRoute1
+        p_y3 = y_ptr_base + 2 * _tmp1 + HWRoute2
+        # p_y4 = y_ptr_base + 3 * _tmp1 + HWRoute3
+    else:
+        p_y1 = y_ptr_base + HWRoute0 * 4 * DC
+        # p_y2 = y_ptr_base + DC + HWRoute1 * 4 * DC
+        p_y3 = y_ptr_base + 2 * DC + HWRoute2 * 4 * DC
+        # p_y4 = y_ptr_base + 3 * DC + HWRoute3 * 4 * DC
+
+    if onebyone == 0:
+        x_ptr_base = x + i_b * _tmp1 + (i_c * BC * DH * DW if x_layout == 0 else i_c * BC)
+        if x_layout == 0:
+            p_x = x_ptr_base + HWRoute0
+        else:
+            p_x = x_ptr_base + HWRoute0 * DC
+
+        if operation == 0:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                _x = tl.load(p_x + _idx_x, mask=_mask_hw)
+                tl.store(p_y1 + _idx_y, _x, mask=_mask_hw)
+                # tl.store(p_y2 + _idx_y, _x, mask=_mask_hw)
+                tl.store(p_y3 + _idx_y, _x, mask=_mask_hw)
+                # tl.store(p_y4 + _idx_y, _x, mask=_mask_hw)
+        elif operation == 1:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                _y1 = tl.load(p_y1 + _idx_y, mask=_mask_hw)
+                # _y2 = tl.load(p_y2 + _idx_y, mask=_mask_hw)
+                _y3 = tl.load(p_y3 + _idx_y, mask=_mask_hw)
+                # _y4 = tl.load(p_y4 + _idx_y, mask=_mask_hw)
+                # tl.store(p_x + _idx_x, _y1 + _y2 + _y3 + _y4, mask=_mask_hw)
+                tl.store(p_x + _idx_x, _y1 + _y3, mask=_mask_hw)
+
+
+    else:
+        x_ptr_base = x + i_b * 4 * _tmp1 + (i_c * BC * DH * DW if x_layout == 0 else i_c * BC)
+        if x_layout == 0:
+            p_x1 = x_ptr_base + HWRoute0
+            p_x2 = p_x1 + _tmp1
+            p_x3 = p_x2 + _tmp1
+            p_x4 = p_x3 + _tmp1
+        else:
+            p_x1 = x_ptr_base + HWRoute0 * 4 * DC
+            p_x2 = p_x1 + DC
+            p_x3 = p_x2 + DC
+            p_x4 = p_x3 + DC
+
+        if operation == 0:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                tl.store(p_y1 + _idx_y, tl.load(p_x1 + _idx_x, mask=_mask_hw), mask=_mask_hw)
+                # tl.store(p_y2 + _idx_y, tl.load(p_x2 + _idx_x, mask=_mask_hw), mask=_mask_hw)
+                tl.store(p_y3 + _idx_y, tl.load(p_x3 + _idx_x, mask=_mask_hw), mask=_mask_hw)
+                # tl.store(p_y4 + _idx_y, tl.load(p_x4 + _idx_x, mask=_mask_hw), mask=_mask_hw)
+        else:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                tl.store(p_x1 + _idx_x, tl.load(p_y1 + _idx_y), mask=_mask_hw)
+                # tl.store(p_x2 + _idx_x, tl.load(p_y2 + _idx_y), mask=_mask_hw)
+                tl.store(p_x3 + _idx_x, tl.load(p_y3 + _idx_y), mask=_mask_hw)
+                # tl.store(p_x4 + _idx_x, tl.load(p_y4 + _idx_y), mask=_mask_hw)
+
+@triton.jit
+def triton_cross_scan_flex_k2(
+    x, # (B, C, H, W) | (B, H, W, C) | (B, 4, C, H, W) | (B, H, W, 4, C)
+    y, # (B, 4, C, H, W) | (B, H, W, 4, C)
+    x_layout: tl.constexpr,
+    y_layout: tl.constexpr,
+    operation: tl.constexpr,
+    onebyone: tl.constexpr,
+    scans: tl.constexpr,
+    BC: tl.constexpr,
+    BH: tl.constexpr,
+    BW: tl.constexpr,
+    DC: tl.constexpr,
+    DH: tl.constexpr,
+    DW: tl.constexpr,
+    NH: tl.constexpr,
+    NW: tl.constexpr,
+):
+    i_hw, i_c, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_h, i_w = (i_hw // NW), (i_hw % NW)
+    _mask_h = (i_h * BH + tl.arange(0, BH)) < DH
+    _mask_w = (i_w * BW + tl.arange(0, BW)) < DW
+    _mask_hw = _mask_h[:, None] & _mask_w[None, :]
+    _for_C = min(DC - i_c * BC, BC)
+
+    HWRoute0 = i_h * BH * DW  + tl.arange(0, BH)[:, None] * DW + i_w * BW + tl.arange(0, BW)[None, :]
+    HWRoute2 = (NH - i_h - 1) * BH * DW  + (BH - 1 - tl.arange(0, BH)[:, None]) * DW + (NW - i_w - 1) * BW + (BW - 1 - tl.arange(0, BW)[None, :]) + (DH - NH * BH) * DW + (DW - NW * BW) # flip
+
+    if scans == 1:
+        HWRoute2 = HWRoute0
+
+    _tmp1 = DC * DH * DW
+
+    y_ptr_base = y + i_b * 4 * _tmp1 + (i_c * BC * DH * DW if y_layout == 0 else i_c * BC)
+    if y_layout == 0:
+        p_y1 = y_ptr_base + HWRoute0
+        p_y2 = y_ptr_base + 2 * _tmp1 + HWRoute2
+    else:
+        p_y1 = y_ptr_base + HWRoute0 * 4 * DC
+        p_y2 = y_ptr_base + 2 * DC + HWRoute2 * 4 * DC
+
+    if onebyone == 0:
+        x_ptr_base = x + i_b * _tmp1 + (i_c * BC * DH * DW if x_layout == 0 else i_c * BC)
+        if x_layout == 0:
+            p_x = x_ptr_base + HWRoute0
+        else:
+            p_x = x_ptr_base + HWRoute0 * DC
+
+        if operation == 0:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                _x = tl.load(p_x + _idx_x, mask=_mask_hw)
+                tl.store(p_y1 + _idx_y, _x, mask=_mask_hw)
+                tl.store(p_y2 + _idx_y, _x, mask=_mask_hw)
+        elif operation == 1:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                _y1 = tl.load(p_y1 + _idx_y, mask=_mask_hw)
+                _y2 = tl.load(p_y2 + _idx_y, mask=_mask_hw)
+                tl.store(p_x + _idx_x, _y1 + _y2, mask=_mask_hw)
+
+    else:
+        x_ptr_base = x + i_b * 4 * _tmp1 + (i_c * BC * DH * DW if x_layout == 0 else i_c * BC)
+        if x_layout == 0:
+            p_x1 = x_ptr_base + HWRoute0
+            p_x2 = p_x1 + 2 * _tmp1
+        else:
+            p_x1 = x_ptr_base + HWRoute0 * 4 * DC
+            p_x2 = p_x1 + 2 * DC
+
+        if operation == 0:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                tl.store(p_y1 + _idx_y, tl.load(p_x1 + _idx_x, mask=_mask_hw), mask=_mask_hw)
+                tl.store(p_y2 + _idx_y, tl.load(p_x2 + _idx_x, mask=_mask_hw), mask=_mask_hw)
+        else:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                tl.store(p_x1 + _idx_x, tl.load(p_y1 + _idx_y), mask=_mask_hw)
+                tl.store(p_x2 + _idx_x, tl.load(p_y2 + _idx_y), mask=_mask_hw)
+
+@triton.jit
+def triton_cross_scan_flex_k2(
+    x, # (B, C, H, W) | (B, H, W, C) | (B, 4, C, H, W) | (B, H, W, 4, C)
+    y, # (B, 4, C, H, W) | (B, H, W, 4, C)
+    x_layout: tl.constexpr,
+    y_layout: tl.constexpr,
+    operation: tl.constexpr,
+    onebyone: tl.constexpr,
+    scans: tl.constexpr,
+    BC: tl.constexpr,
+    BH: tl.constexpr,
+    BW: tl.constexpr,
+    DC: tl.constexpr,
+    DH: tl.constexpr,
+    DW: tl.constexpr,
+    NH: tl.constexpr,
+    NW: tl.constexpr,
+):
+    i_hw, i_c, i_b = tl.program_id(0), tl.program_id(1), tl.program_id(2)
+    i_h, i_w = (i_hw // NW), (i_hw % NW)
+    _mask_h = (i_h * BH + tl.arange(0, BH)) < DH
+    _mask_w = (i_w * BW + tl.arange(0, BW)) < DW
+    _mask_hw = _mask_h[:, None] & _mask_w[None, :]
+    _for_C = min(DC - i_c * BC, BC)
+
+    HWRoute0 = i_h * BH * DW  + tl.arange(0, BH)[:, None] * DW + i_w * BW + tl.arange(0, BW)[None, :]
+    HWRoute2 = (NH - i_h - 1) * BH * DW  + (BH - 1 - tl.arange(0, BH)[:, None]) * DW + (NW - i_w - 1) * BW + (BW - 1 - tl.arange(0, BW)[None, :]) + (DH - NH * BH) * DW + (DW - NW * BW) # flip
+
+    if scans == 1:
+        HWRoute2 = HWRoute0
+
+    _tmp1 = DC * DH * DW
+
+    y_ptr_base = y + i_b * 2 * _tmp1 + (i_c * BC * DH * DW if y_layout == 0 else i_c * BC)
+    if y_layout == 0:
+        p_y1 = y_ptr_base + HWRoute0
+        p_y2 = y_ptr_base + 1 * _tmp1 + HWRoute2
+    else:
+        p_y1 = y_ptr_base + HWRoute0 * 4 * DC
+        p_y2 = y_ptr_base + 1 * DC + HWRoute2 * 4 * DC
+
+    if onebyone == 0:
+        x_ptr_base = x + i_b * _tmp1 + (i_c * BC * DH * DW if x_layout == 0 else i_c * BC)
+        if x_layout == 0:
+            p_x = x_ptr_base + HWRoute0
+        else:
+            p_x = x_ptr_base + HWRoute0 * DC
+
+        if operation == 0:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                _x = tl.load(p_x + _idx_x, mask=_mask_hw)
+                tl.store(p_y1 + _idx_y, _x, mask=_mask_hw)
+                tl.store(p_y2 + _idx_y, _x, mask=_mask_hw)
+        elif operation == 1:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                _y1 = tl.load(p_y1 + _idx_y, mask=_mask_hw)
+                _y2 = tl.load(p_y2 + _idx_y, mask=_mask_hw)
+                tl.store(p_x + _idx_x, _y1 + _y2, mask=_mask_hw)
+
+    else:
+        x_ptr_base = x + i_b * 4 * _tmp1 + (i_c * BC * DH * DW if x_layout == 0 else i_c * BC)
+        if x_layout == 0:
+            p_x1 = x_ptr_base + HWRoute0
+            p_x2 = p_x1 + 2 * _tmp1
+        else:
+            p_x1 = x_ptr_base + HWRoute0 * 4 * DC
+            p_x2 = p_x1 + 2 * DC
+
+        if operation == 0:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                _x1 = tl.load(p_x1 + _idx_x, mask=_mask_hw)
+                _x2 = tl.load(p_x2 + _idx_x, mask=_mask_hw)
+                tl.store(p_y1 + _idx_y, _x1, mask=_mask_hw)
+                tl.store(p_y2 + _idx_y, _x2, mask=_mask_hw)
+        else:
+            for idxc in range(_for_C):
+                _idx_x = idxc * DH * DW if x_layout == 0 else idxc
+                _idx_y = idxc * DH * DW if y_layout == 0 else idxc
+                _y1 = tl.load(p_y1 + _idx_y, mask=_mask_hw)
+                _y2 = tl.load(p_y2 + _idx_y, mask=_mask_hw)
+                tl.store(p_x1 + _idx_x, _y1, mask=_mask_hw)
+                tl.store(p_x2 + _idx_x, _y2, mask=_mask_hw)
+
+class CrossScanTritonFk2(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=2):
+        if one_by_one:
+            if in_channel_first:
+                B, _, C, H, W = x.shape
+            else:
+                B, H, W, _, C = x.shape
+        else:
+            if in_channel_first:
+                B, C, H, W = x.shape
+            else:
+                B, H, W, C = x.shape
+        B, C, H, W = int(B), int(C), int(H), int(W)
+        BC, BH, BW = 1, 32, 32
+        NH, NW, NC = triton.cdiv(H, BH), triton.cdiv(W, BW), triton.cdiv(C, BC)
+        
+        ctx.in_channel_first = in_channel_first
+        ctx.out_channel_first = out_channel_first
+        ctx.one_by_one = one_by_one
+        ctx.scans = scans
+        ctx.shape = (B, C, H, W)
+        ctx.triton_shape = (BC, BH, BW, NC, NH, NW)
+
+        y = x.new_empty((B, 2, C, H * W)) if out_channel_first else x.new_empty((B, H * W, 2, C))
+        triton_cross_scan_flex_k2[(NH * NW, NC, B)](
+            x.contiguous(), y, 
+            (0 if in_channel_first else 1), (0 if out_channel_first else 1), 0, (0 if not one_by_one else 1), scans, 
+            BC, BH, BW, C, H, W, NH, NW
+        )
+        return y
+        
+    @staticmethod
+    def backward(ctx, y: torch.Tensor):
+        in_channel_first = ctx.in_channel_first
+        out_channel_first = ctx.out_channel_first
+        one_by_one = ctx.one_by_one
+        scans = ctx.scans
+        B, C, H, W = ctx.shape
+        BC, BH, BW, NC, NH, NW = ctx.triton_shape
+        if one_by_one:
+            x = y.new_empty((B, 2, C, H, W)) if in_channel_first else y.new_empty((B, H, W, 2, C))
+        else:
+            x = y.new_empty((B, C, H, W)) if in_channel_first else y.new_empty((B, H, W, C))
+        
+        triton_cross_scan_flex_k2[(NH * NW, NC, B)](
+            x, y.contiguous(), 
+            (0 if in_channel_first else 1), (0 if out_channel_first else 1), 1, (0 if not one_by_one else 1), scans,
+            BC, BH, BW, C, H, W, NH, NW
+        )
+        return x, None, None, None, None
+
+
+class CrossMergeTritonFk2(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, y: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=2):
+        if out_channel_first:
+            B, _, C, H, W = y.shape
+        else:
+            B, H, W, _, C = y.shape
+        B, C, H, W = int(B), int(C), int(H), int(W)
+        BC, BH, BW = 1, 32, 32
+        NH, NW, NC = triton.cdiv(H, BH), triton.cdiv(W, BW), triton.cdiv(C, BC)
+        ctx.in_channel_first = in_channel_first
+        ctx.out_channel_first = out_channel_first
+        ctx.one_by_one = one_by_one
+        ctx.scans = scans
+        ctx.shape = (B, C, H, W)
+        ctx.triton_shape = (BC, BH, BW, NC, NH, NW)
+        if one_by_one:
+            x = y.new_empty((B, 2, C, H * W)) if in_channel_first else y.new_empty((B, H * W, 2, C))
+        else:
+            x = y.new_empty((B, C, H * W)) if in_channel_first else y.new_empty((B, H * W, C))
+        triton_cross_scan_flex_k2[(NH * NW, NC, B)](
+            x, y.contiguous(), 
+            (0 if in_channel_first else 1), (0 if out_channel_first else 1), 1, (0 if not one_by_one else 1), scans,
+            BC, BH, BW, C, H, W, NH, NW
+        )
+        return x
+        
+    @staticmethod
+    def backward(ctx, x: torch.Tensor):
+        in_channel_first = ctx.in_channel_first
+        out_channel_first = ctx.out_channel_first
+        one_by_one = ctx.one_by_one
+        scans = ctx.scans
+        B, C, H, W = ctx.shape
+        BC, BH, BW, NC, NH, NW = ctx.triton_shape
+        y = x.new_empty((B, 2, C, H, W)) if out_channel_first else x.new_empty((B, H, W, 2, C))
+        triton_cross_scan_flex_k2[(NH * NW, NC, B)](
+            x.contiguous(), y, 
+            (0 if in_channel_first else 1), (0 if out_channel_first else 1), 0, (0 if not one_by_one else 1), scans,
+            BC, BH, BW, C, H, W, NH, NW
+        )
+        return y, None, None, None, None, None
+
+
+# @torch.compile(options={"triton.cudagraphs": True}, fullgraph=True)
+def cross_scan_fn_k2(x: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=2, force_torch=False):
+    # x: (B, C, H, W) | (B, H, W, C) | (B, 2, C, H, W) | (B, H, W, 2, C)
+    # y: (B, 2, C, L) | (B, L, 2, C)
+    # scans: 0: cross scan; 1 unidirectional; 2: bidirectional;
+    CSF = CrossScanTritonFk2 if WITH_TRITON and x.is_cuda and (not force_torch) else CrossScanF
+    return CSF.apply(x, in_channel_first, out_channel_first, one_by_one, scans)
+
+# @torch.compile(options={"triton.cudagraphs": True}, fullgraph=True)
+def cross_merge_fn_k2(y: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=2, force_torch=False):
+    # y: (B, 2, C, L) | (B, L, 2, C)
+    # x: (B, C, H * W) | (B, H * W, C) | (B, 2, C, H * W) | (B, H * W, 2, C)
+    # scans: 0: cross scan; 1 unidirectional; 2: bidirectional;
+    CMF = CrossMergeTritonFk2 if WITH_TRITON and y.is_cuda and (not force_torch) else CrossMergeF
+    return CMF.apply(y, in_channel_first, out_channel_first, one_by_one, scans)
+
+def cross_scan_fn_k2_torch(x: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=2, force_torch=False):
+    cross_scan = CrossScan(in_channel_first, out_channel_first, one_by_one, scans)
+    return cross_scan(x)
+
+def cross_merge_fn_k2_torch(y: torch.Tensor, in_channel_first=True, out_channel_first=True, one_by_one=False, scans=2, force_torch=False):
+    cross_merge = CrossMerge(in_channel_first, out_channel_first, one_by_one, scans)
+    return cross_merge(y)
+
+# checks =================================================================
+
+class CHECK:
+    def check_csm_triton():
+        B, C, H, W = 2, 192, 56, 57
+        dtype=torch.float16
+        dtype=torch.float32
+        x = torch.randn((B, C, H, W), dtype=dtype, device=torch.device("cuda")).requires_grad_(True)
+        y = torch.randn((B, 2, C, H, W), dtype=dtype, device=torch.device("cuda")).requires_grad_(True)
+        x1 = x.clone().detach().requires_grad_(True)
+        y1 = y.clone().detach().requires_grad_(True)
+
+        def cross_scan(x: torch.Tensor):
+            B, C, H, W = x.shape
+            L = H * W
+            xs = torch.stack([
+                x.view(B, C, L),
+                torch.transpose(x, dim0=2, dim1=3).contiguous().view(B, C, L),
+                torch.flip(x.contiguous().view(B, C, L), dims=[-1]),
+                torch.flip(torch.transpose(x, dim0=2, dim1=3).contiguous().view(B, C, L), dims=[-1]),
+            ], dim=1).view(B, 4, C, L)
+            return xs
+        
+        def cross_merge(out_y: torch.Tensor):
+            B, K, D, H, W = out_y.shape
+            L = H * W
+            out_y = out_y.view(B, K, D, L)
+            inv_y = torch.flip(out_y[:, 2:4], dims=[-1]).view(B, 2, -1, L)
+            wh_y = torch.transpose(out_y[:, 1].view(B, -1, W, H), dim0=2, dim1=3).contiguous().view(B, -1, L)
+            invwh_y = torch.transpose(inv_y[:, 1].view(B, -1, W, H), dim0=2, dim1=3).contiguous().view(B, -1, L)
+            y = out_y[:, 0] + inv_y[:, 0] + wh_y + invwh_y
+            return y
+
+        def cross_scan_1b1(x: torch.Tensor):
+            B, K, C, H, W = x.shape
+            L = H * W
+            xs = torch.stack([
+                x[:, 0].view(B, C, L),
+                torch.transpose(x[:, 1], dim0=2, dim1=3).contiguous().view(B, C, L),
+                torch.flip(x[:, 2].contiguous().view(B, C, L), dims=[-1]),
+                torch.flip(torch.transpose(x[:, 3], dim0=2, dim1=3).contiguous().view(B, C, L), dims=[-1]),
+            ], dim=1).view(B, 2, C, L)
+            return xs
+        
+        def unidi_scan(x):
+            B, C, H, W = x.shape
+            x = x.view(B, 1, C, H * W).repeat(1, 4, 1, 1)
+            return x
+        
+        def unidi_merge(ys):
+            B, K, C, H, W = ys.shape
+            return ys.view(B, 4, -1, H * W).sum(1)
+
+        def bidi_scan(x):
+            B, C, H, W = x.shape
+            x = x.view(B, 1, C, H * W).repeat(1, 2, 1, 1)
+            x = torch.cat([x, x.flip(dims=[-1])], dim=1)
+            return x
+        
+        def bidi_merge(ys):
+            B, K, D, H, W = ys.shape
+            ys = ys.view(B, K, D, -1)
+            ys = ys[:, 0:2] + ys[:, 2:4].flip(dims=[-1]).view(B, 2, D, -1)
+            return ys.contiguous().sum(1)
+
+        if True:
+            # res0 = triton.testing.do_bench(lambda :cross_scan(x))
+            res1 = triton.testing.do_bench(lambda :cross_scan_fn_k2(x, True, True, False))
+            # res2 = triton.testing.do_bench(lambda :CrossScanTriton.apply(x))
+            # res3 = triton.testing.do_bench(lambda :cross_merge(y))
+            res4 = triton.testing.do_bench(lambda :cross_merge_fn_k2(y, True, True, False))
+            # res5 = triton.testing.do_bench(lambda :CrossMergeTriton.apply(y))
+            # print(res0, res1, res2, res3, res4, res5)
+            print(res0, res1, res3, res4)
+            res0 = triton.testing.do_bench(lambda :cross_scan(x).sum().backward())
+            res1 = triton.testing.do_bench(lambda :cross_scan_fn_k2(x, True, True, False).sum().backward())
+            # res2 = triton.testing.do_bench(lambda :CrossScanTriton.apply(x).sum().backward())
+            res3 = triton.testing.do_bench(lambda :cross_merge(y).sum().backward())
+            res4 = triton.testing.do_bench(lambda :cross_merge_fn_k2(y, True, True, False).sum().backward())
+            # res5 = triton.testing.do_bench(lambda :CrossMergeTriton.apply(y).sum().backward())
+            # print(res0, res1, res2, res3, res4, res5)
+            print(res0, res1, res3, res4)
+
+        print("test cross scan")
+        for (cs0, cm0, cs1, cm1) in [
+            # channel_first -> channel_first
+            (cross_scan, cross_merge, cross_scan_fn_k2, cross_merge_fn_k2),
+            (unidi_scan, unidi_merge, lambda x: cross_scan_fn_k2(x, scans=1), lambda x: cross_merge_fn_k2(x, scans=1)),
+            (bidi_scan, bidi_merge, lambda x: cross_scan_fn_k2(x, scans=2), lambda x: cross_merge_fn_k2(x, scans=2)),
+            
+            # flex: BLC->BCL; BCL->BLC; BLC->BLC;
+            (cross_scan, cross_merge, lambda x: cross_scan_fn_k2(x.permute(0, 2, 3, 1), in_channel_first=False), lambda x: cross_merge_fn_k2(x, in_channel_first=False).permute(0, 2, 1)),
+            (cross_scan, cross_merge, lambda x: cross_scan_fn_k2(x, out_channel_first=False).permute(0, 2, 3, 1), lambda x: cross_merge_fn_k2(x.permute(0, 3, 4, 1, 2), out_channel_first=False)),
+            (cross_scan, cross_merge, lambda x: cross_scan_fn_k2(x.permute(0, 2, 3, 1), in_channel_first=False, out_channel_first=False).permute(0, 2, 3, 1), lambda x: cross_merge_fn_k2(x.permute(0, 3, 4, 1, 2), in_channel_first=False, out_channel_first=False).permute(0, 2, 1)),
+            
+            # previous
+            # (cross_scan, cross_merge, lambda x: CrossScanTriton.apply(x), lambda x: CrossMergeTriton.apply(x)),
+            # (unidi_scan, unidi_merge, lambda x: getCSM(1)[0].apply(x), lambda x: getCSM(1)[1].apply(x)),
+            # (bidi_scan, bidi_merge, lambda x: getCSM(2)[0].apply(x), lambda x: getCSM(2)[1].apply(x)),
+        ]:
+            x.grad, x1.grad, y.grad, y1.grad = None, None, None, None
+            o0 = cs0(x)
+            o1 = cs1(x1)
+            o0.backward(y.view(B, 2, C, H * W))
+            o1.backward(y.view(B, 2, C, H * W))
+            print((o0 - o1).abs().max())
+            print((x.grad - x1.grad).abs().max())
+            o0 = cm0(y)
+            o1 = cm1(y1)
+            o0.backward(x.view(B, C, H * W))
+            o1.backward(x.view(B, C, H * W))
+            print((o0 - o1).abs().max())
+            print((y.grad - y1.grad).abs().max())
+            x.grad, x1.grad, y.grad, y1.grad = None, None, None, None
+            print("===============", flush=True)
+
+        print("test cross scan one by one")
+        for (cs0, cs1) in [
+            (cross_scan_1b1, lambda x: cross_scan_fn_k2(x, one_by_one=True)),
+            # (cross_scan_1b1, lambda x: CrossScanTriton1b1.apply(x)),
+        ]:
+            o0 = cs0(y)
+            o1 = cs1(y1)
+            o0.backward(y.view(B, 2, C, H * W))
+            o1.backward(y.view(B, 2, C, H * W))
+            print((o0 - o1).abs().max())
+            print((y.grad - y1.grad).abs().max())
+            x.grad, x1.grad, y.grad, y1.grad = None, None, None, None
+            print("===============", flush=True)
+
+
+if __name__ == "__main__":
+    CHECK.check_csm_triton()
+
+
+
+

rscd/models/backbones/lib_mamba/csms6s.py

@@ -0,0 +1,266 @@
+import time
+import torch
+import warnings
+
+
+WITH_SELECTIVESCAN_OFLEX = True
+WITH_SELECTIVESCAN_CORE = False
+WITH_SELECTIVESCAN_MAMBA = True
+try:
+    import selective_scan_cuda_oflex
+except ImportError:
+    WITH_SELECTIVESCAN_OFLEX = False
+    warnings.warn("Can not import selective_scan_cuda_oflex. This affects speed.")
+    print("Can not import selective_scan_cuda_oflex. This affects speed.", flush=True)
+try:
+    import selective_scan_cuda_core
+except ImportError:
+    WITH_SELECTIVESCAN_CORE = False
+try:
+    import selective_scan_cuda
+except ImportError:
+    WITH_SELECTIVESCAN_MAMBA = False
+
+
+def selective_scan_torch(
+    u: torch.Tensor, # (B, K * C, L)
+    delta: torch.Tensor, # (B, K * C, L)
+    A: torch.Tensor, # (K * C, N)
+    B: torch.Tensor, # (B, K, N, L)
+    C: torch.Tensor, # (B, K, N, L)
+    D: torch.Tensor = None, # (K * C)
+    delta_bias: torch.Tensor = None, # (K * C)
+    delta_softplus=True, 
+    oflex=True, 
+    *args,
+    **kwargs
+):
+    dtype_in = u.dtype
+    Batch, K, N, L = B.shape
+    KCdim = u.shape[1]
+    Cdim = int(KCdim / K)
+    assert u.shape == (Batch, KCdim, L)
+    assert delta.shape == (Batch, KCdim, L)
+    assert A.shape == (KCdim, N)
+    assert C.shape == B.shape
+
+    if delta_bias is not None:
+        delta = delta + delta_bias[..., None]
+    if delta_softplus:
+        delta = torch.nn.functional.softplus(delta)
+            
+    u, delta, A, B, C = u.float(), delta.float(), A.float(), B.float(), C.float()
+    B = B.view(Batch, K, 1, N, L).repeat(1, 1, Cdim, 1, 1).view(Batch, KCdim, N, L)
+    C = C.view(Batch, K, 1, N, L).repeat(1, 1, Cdim, 1, 1).view(Batch, KCdim, N, L)
+    deltaA = torch.exp(torch.einsum('bdl,dn->bdln', delta, A))
+    deltaB_u = torch.einsum('bdl,bdnl,bdl->bdln', delta, B, u)
+    
+    if True:
+        x = A.new_zeros((Batch, KCdim, N))
+        ys = []
+        for i in range(L):
+            x = deltaA[:, :, i, :] * x + deltaB_u[:, :, i, :]
+            y = torch.einsum('bdn,bdn->bd', x, C[:, :, :, i])
+            ys.append(y)
+        y = torch.stack(ys, dim=2) # (B, C, L)
+    
+    out = y if D is None else y + u * D.unsqueeze(-1)
+    return out if oflex else out.to(dtype=dtype_in)
+
+
+class SelectiveScanCuda(torch.autograd.Function):
+    @staticmethod
+    @torch.cuda.amp.custom_fwd
+    def forward(ctx, u, delta, A, B, C, D=None, delta_bias=None, delta_softplus=False, oflex=True, backend=None):
+        ctx.delta_softplus = delta_softplus
+        backend = "oflex" if WITH_SELECTIVESCAN_OFLEX and (backend is None) else backend
+        backend = "core" if WITH_SELECTIVESCAN_CORE and (backend is None) else backend
+        backend = "mamba" if WITH_SELECTIVESCAN_MAMBA and (backend is None) else backend
+        ctx.backend = backend
+        if backend == "oflex":
+            out, x, *rest = selective_scan_cuda_oflex.fwd(u, delta, A, B, C, D, delta_bias, delta_softplus, 1, oflex)
+        elif backend == "core":
+            out, x, *rest = selective_scan_cuda_core.fwd(u, delta, A, B, C, D, delta_bias, delta_softplus, 1)
+        elif backend == "mamba":
+            out, x, *rest = selective_scan_cuda.fwd(u, delta, A, B, C, D, None, delta_bias, delta_softplus)
+        ctx.save_for_backward(u, delta, A, B, C, D, delta_bias, x)
+        return out
+    
+    @staticmethod
+    @torch.cuda.amp.custom_bwd
+    def backward(ctx, dout, *args):
+        u, delta, A, B, C, D, delta_bias, x = ctx.saved_tensors
+        backend = ctx.backend
+        if dout.stride(-1) != 1:
+            dout = dout.contiguous()
+        if backend == "oflex":
+            du, ddelta, dA, dB, dC, dD, ddelta_bias, *rest = selective_scan_cuda_oflex.bwd(
+                u, delta, A, B, C, D, delta_bias, dout, x, ctx.delta_softplus, 1
+            )
+        elif backend == "core":
+            du, ddelta, dA, dB, dC, dD, ddelta_bias, *rest = selective_scan_cuda_core.bwd(
+                u, delta, A, B, C, D, delta_bias, dout, x, ctx.delta_softplus, 1
+            )
+        elif backend == "mamba":
+            du, ddelta, dA, dB, dC, dD, ddelta_bias, *rest = selective_scan_cuda.bwd(
+                u, delta, A, B, C, D, None, delta_bias, dout, x, None, None, ctx.delta_softplus,
+                False
+            )
+        return du, ddelta, dA, dB, dC, dD, ddelta_bias, None, None, None
+
+
+def selective_scan_fn(
+    u: torch.Tensor, # (B, K * C, L)
+    delta: torch.Tensor, # (B, K * C, L)
+    A: torch.Tensor, # (K * C, N)
+    B: torch.Tensor, # (B, K, N, L)
+    C: torch.Tensor, # (B, K, N, L)
+    D: torch.Tensor = None, # (K * C)
+    delta_bias: torch.Tensor = None, # (K * C)
+    delta_softplus=True, 
+    oflex=True,
+    backend=None,
+):
+    WITH_CUDA = (WITH_SELECTIVESCAN_OFLEX or WITH_SELECTIVESCAN_CORE or WITH_SELECTIVESCAN_MAMBA)
+    fn = selective_scan_torch if backend == "torch" or (not WITH_CUDA) else SelectiveScanCuda.apply
+    return fn(u, delta, A, B, C, D, delta_bias, delta_softplus, oflex, backend)
+
+
+# fvcore flops =======================================
+def print_jit_input_names(inputs):
+    print("input params: ", end=" ", flush=True)
+    try: 
+        for i in range(10):
+            print(inputs[i].debugName(), end=" ", flush=True)
+    except Exception as e:
+        pass
+    print("", flush=True)
+
+def flops_selective_scan_fn(B=1, L=256, D=768, N=16, with_D=True, with_Z=False, with_complex=False):
+    """
+    u: r(B D L)
+    delta: r(B D L)
+    A: r(D N)
+    B: r(B N L)
+    C: r(B N L)
+    D: r(D)
+    z: r(B D L)
+    delta_bias: r(D), fp32
+    
+    ignores:
+        [.float(), +, .softplus, .shape, new_zeros, repeat, stack, to(dtype), silu] 
+    """
+    assert not with_complex 
+    # https://github.com/state-spaces/mamba/issues/110
+    flops = 9 * B * L * D * N
+    if with_D:
+        flops += B * D * L
+    if with_Z:
+        flops += B * D * L    
+    return flops
+
+# this is only for selective_scan_ref...
+def flops_selective_scan_ref(B=1, L=256, D=768, N=16, with_D=True, with_Z=False, with_Group=True, with_complex=False):
+    """
+    u: r(B D L)
+    delta: r(B D L)
+    A: r(D N)
+    B: r(B N L)
+    C: r(B N L)
+    D: r(D)
+    z: r(B D L)
+    delta_bias: r(D), fp32
+    
+    ignores:
+        [.float(), +, .softplus, .shape, new_zeros, repeat, stack, to(dtype), silu] 
+    """
+    import numpy as np
+    
+    # fvcore.nn.jit_handles
+    def get_flops_einsum(input_shapes, equation):
+        np_arrs = [np.zeros(s) for s in input_shapes]
+        optim = np.einsum_path(equation, *np_arrs, optimize="optimal")[1]
+        for line in optim.split("\n"):
+            if "optimized flop" in line.lower():
+                # divided by 2 because we count MAC (multiply-add counted as one flop)
+                flop = float(np.floor(float(line.split(":")[-1]) / 2))
+                return flop
+    
+
+    assert not with_complex
+
+    flops = 0 # below code flops = 0
+
+    flops += get_flops_einsum([[B, D, L], [D, N]], "bdl,dn->bdln")
+    if with_Group:
+        flops += get_flops_einsum([[B, D, L], [B, N, L], [B, D, L]], "bdl,bnl,bdl->bdln")
+    else:
+        flops += get_flops_einsum([[B, D, L], [B, D, N, L], [B, D, L]], "bdl,bdnl,bdl->bdln")
+  
+    in_for_flops = B * D * N   
+    if with_Group:
+        in_for_flops += get_flops_einsum([[B, D, N], [B, D, N]], "bdn,bdn->bd")
+    else:
+        in_for_flops += get_flops_einsum([[B, D, N], [B, N]], "bdn,bn->bd")
+    flops += L * in_for_flops 
+    if with_D:
+        flops += B * D * L
+    if with_Z:
+        flops += B * D * L  
+    return flops
+
+def selective_scan_flop_jit(inputs, outputs, backend="prefixsum", verbose=True):
+    if verbose:
+        print_jit_input_names(inputs)
+    flops_fn = flops_selective_scan_ref if backend == "naive" else flops_selective_scan_fn
+    B, D, L = inputs[0].type().sizes()
+    N = inputs[2].type().sizes()[1]
+    flops = flops_fn(B=B, L=L, D=D, N=N, with_D=True, with_Z=False)
+    return flops
+
+
+if __name__ == "__main__":
+    def params(B, K, C, N, L, device = torch.device("cuda"), itype = torch.float):
+        As = (-0.5 * torch.rand(K * C, N, device=device, dtype=torch.float32)).requires_grad_()
+        Bs = torch.randn((B, K, N, L), device=device, dtype=itype).requires_grad_()
+        Cs = torch.randn((B, K, N, L), device=device, dtype=itype).requires_grad_()
+        Ds = torch.randn((K * C), device=device, dtype=torch.float32).requires_grad_()
+        u = torch.randn((B, K * C, L), device=device, dtype=itype).requires_grad_()
+        delta = (0.5 * torch.rand((B, K * C, L),  device=device, dtype=itype)).requires_grad_()
+        delta_bias = (0.5 * torch.rand((K * C), device=device, dtype=torch.float32)).requires_grad_()
+        return u, delta, As, Bs, Cs, Ds, delta_bias
+
+    def bench(func, xs, Warmup=30, NTimes=20):
+        import time
+        torch.cuda.synchronize()
+        for r in range(Warmup):
+            for x in xs:
+                func(x)
+        torch.cuda.synchronize()
+        tim0 = time.time()
+        for r in range(NTimes):
+            for x in xs:
+                func(x)
+        torch.cuda.synchronize()
+        return (time.time() - tim0) / NTimes
+
+    def check():
+        u, delta, As, Bs, Cs, Ds, delta_bias = params(1, 4, 16, 8, 512, itype=torch.float16)
+        u1, delta1, As1, Bs1, Cs1, Ds1, delta_bias1 = [x.clone().detach().requires_grad_() for x in [u, delta, As, Bs, Cs, Ds, delta_bias]]
+        
+        # out_ref = selective_scan_fn(u, delta, As, Bs, Cs, Ds, delta_bias, True, backend="torch")
+        out = selective_scan_fn(u1, delta1, As1, Bs1, Cs1, Ds1, delta_bias1, True, backend="oflex")
+        out_ref = selective_scan_fn(u, delta, As, Bs, Cs, Ds, delta_bias, True, backend="mamba")
+        print((out_ref - out).abs().max())
+        out.sum().backward()
+        out_ref.sum().backward()
+        for x, y in zip([u, As, Bs, Cs, Ds, delta, delta_bias], [u1, As1, Bs1, Cs1, Ds1, delta1, delta_bias1]):
+            print((x.grad - y.grad).abs().max())
+
+        u, delta, As, Bs, Cs, Ds, delta_bias = params(128, 4, 96, 8, 56 * 56)
+        print(bench(lambda x: selective_scan_fn(x[0], x[1], x[2], x[3], x[4], x[5], x[6], True, backend="oflex"), [(u, delta, As, Bs, Cs, Ds, delta_bias),]))
+        print(bench(lambda x: selective_scan_fn(x[0], x[1], x[2], x[3], x[4], x[5], x[6], True, backend="mamba"), [(u, delta, As, Bs, Cs, Ds, delta_bias),]))
+        print(bench(lambda x: selective_scan_fn(x[0], x[1], x[2], x[3], x[4], x[5], x[6], True, backend="torch"), [(u, delta, As, Bs, Cs, Ds, delta_bias),]))
+
+    check()
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/README.md

@@ -0,0 +1,97 @@
+# mamba-mini
+An efficient implementation of selective scan in one file, works with both cpu and gpu, with corresponding mathematical derivation. It is probably the code which is the most close to selective_scan_cuda in mamba.
+
+### mathematical derivation
+![image](../assets/derivation.png)
+
+### code
+```python
+import torch
+def selective_scan_easy(us, dts, As, Bs, Cs, Ds, delta_bias=None, delta_softplus=False, return_last_state=False, chunksize=64):
+    """
+    # B: batch_size, G: groups, D: dim, N: state dim, L: seqlen
+    us: B, G * D, L 
+    dts: B, G * D, L
+    As: G * D, N
+    Bs: B, G, N, L
+    Cs: B, G, N, L
+    Ds: G * D
+    delta_bias: G * D
+    # chunksize can be any as you like. But as the chunksize raises, hs may get None, as exp(sum(delta) A) is really small
+    """
+    def selective_scan_chunk(us, dts, As, Bs, Cs, hprefix):
+        """
+        partial(h) / partial(t) = Ah + Bu; y = Ch + Du;
+        => partial(h*exp(-At)) / partial(t) = Bu*exp(-At);
+        => h_t = h_0 + sum_{0}_{t}_{Bu*exp(A(t-v)) dv};
+        => h_b = exp(A(dt_a + ... + dt_{b-1})) * (h_a + sum_{a}_{b-1}_{Bu*exp(-A(dt_a + ... + dt_i)) dt_i});
+           y_i = C_i*h_i + D*u_i
+        """
+        """
+        us, dts: (L, B, G, D) # L is chunk_size
+        As: (G, D, N)
+        Bs, Cs: (L, B, G, N)
+        Ds: (G, D)
+        hprefix: (B, G, D, N)
+        """
+        ts = dts.cumsum(dim=0)
+        Ats = torch.einsum("gdn,lbgd->lbgdn", As, ts).exp()
+        scale = Ats[-1].detach()
+        rAts = Ats / scale
+        duts = dts * us
+        dtBus = torch.einsum("lbgd,lbgn->lbgdn", duts, Bs)
+        hs_tmp = rAts * (dtBus / rAts).cumsum(dim=0) 
+        hs = hs_tmp + Ats * hprefix.unsqueeze(0)
+        ys = torch.einsum("lbgn,lbgdn->lbgd", Cs, hs) 
+        return ys, hs
+    
+    inp_dtype = us.dtype
+    has_D = Ds is not None
+
+    dts = dts.float()
+    if delta_bias is not None:
+        dts = dts + delta_bias.view(1, -1, 1).float()
+    if delta_softplus:
+        dts = torch.nn.functional.softplus(dts)
+    
+    if len(Bs.shape) == 3:
+        Bs = Bs.unsqueeze(1)
+    if len(Cs.shape) == 3:
+        Cs = Cs.unsqueeze(1)
+    B, G, N, L = Bs.shape
+    us = us.view(B, G, -1, L).permute(3, 0, 1, 2).float()
+    dts = dts.view(B, G, -1, L).permute(3, 0, 1, 2).float()
+    As = As.view(G, -1, N).float()
+    Bs = Bs.permute(3, 0, 1, 2).float()
+    Cs = Cs.permute(3, 0, 1, 2).float()
+    Ds = Ds.view(G, -1).float() if has_D else None
+    D = As.shape[1]
+    
+    oys = []
+    # ohs = []
+    hprefix = us.new_zeros((B, G, D, N), dtype=torch.float)
+    for i in range(0, L - 1, chunksize):
+        ys, hs = selective_scan_chunk(
+            us[i:i + chunksize], dts[i:i + chunksize], 
+            As, Bs[i:i + chunksize], Cs[i:i + chunksize], hprefix, 
+        )
+        oys.append(ys)
+        # ohs.append(hs)
+        hprefix = hs[-1]
+
+    oys = torch.cat(oys, dim=0)
+    # ohs = torch.cat(ohs, dim=0)
+    if has_D:
+        oys = oys + Ds * us
+    oys = oys.permute(1, 2, 3, 0).view(B, -1, L)
+    oys = oys.to(inp_dtype)
+    # hprefix = hprefix.to(inp_dtype)
+
+    return oys if not return_last_state else (oys, hprefix.view(B, G * D, N))
+
+```
+
+### to test
+```bash
+pytest test_selective_scan.py
+```

rscd/models/backbones/lib_mamba/kernels/selective_scan/build/lib.linux-x86_64-3.8/selective_scan_cuda_oflex.cpython-38-x86_64-linux-gnu.so

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3fcf524eeaf71e641653c1aff1f8fac591a1e6916300d322830bd02476873ab1
+size 34969816

rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/.ninja_deps

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6e61c23bbef4f0b8f414187d9149e6e0818ce400c3351405d494b774b988bf6d
+size 501136

rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/.ninja_log

@@ -0,0 +1,4 @@
+# ninja log v5
+7	17272	1748140810026258300	/mnt/d/WORK/rschange-main/rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/csrc/selective_scan/cusoflex/selective_scan_core_fwd.o	ab3bac6bd7b8268f
+8	23832	1748140816810431800	/mnt/d/WORK/rschange-main/rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/csrc/selective_scan/cusoflex/selective_scan_oflex.o	7f9a77b388057fc6
+7	57431	1748140850419474900	/mnt/d/WORK/rschange-main/rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/csrc/selective_scan/cusoflex/selective_scan_core_bwd.o	3cffffbdd6b9fec1

rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/build.ninja

@@ -0,0 +1,35 @@
+ninja_required_version = 1.3
+cxx = c++
+nvcc = /usr/local/cuda-11.8/bin/nvcc
+
+cflags = -pthread -B /root/anaconda3/envs/rscd/compiler_compat -Wl,--sysroot=/ -Wsign-compare -DNDEBUG -g -fwrapv -O3 -Wall -Wstrict-prototypes -fPIC -I/mnt/d/WORK/rschange-main/rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan -I/root/anaconda3/envs/rscd/lib/python3.8/site-packages/torch/include -I/root/anaconda3/envs/rscd/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -I/root/anaconda3/envs/rscd/lib/python3.8/site-packages/torch/include/TH -I/root/anaconda3/envs/rscd/lib/python3.8/site-packages/torch/include/THC -I/usr/local/cuda-11.8/include -I/root/anaconda3/envs/rscd/include/python3.8 -c
+post_cflags = -O3 -std=c++17 -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=selective_scan_cuda_oflex -D_GLIBCXX_USE_CXX11_ABI=0
+cuda_cflags = -I/mnt/d/WORK/rschange-main/rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan -I/root/anaconda3/envs/rscd/lib/python3.8/site-packages/torch/include -I/root/anaconda3/envs/rscd/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -I/root/anaconda3/envs/rscd/lib/python3.8/site-packages/torch/include/TH -I/root/anaconda3/envs/rscd/lib/python3.8/site-packages/torch/include/THC -I/usr/local/cuda-11.8/include -I/root/anaconda3/envs/rscd/include/python3.8 -c
+cuda_post_cflags = -D__CUDA_NO_HALF_OPERATORS__ -D__CUDA_NO_HALF_CONVERSIONS__ -D__CUDA_NO_BFLOAT16_CONVERSIONS__ -D__CUDA_NO_HALF2_OPERATORS__ --expt-relaxed-constexpr --compiler-options ''"'"'-fPIC'"'"'' -O3 -std=c++17 -U__CUDA_NO_HALF_OPERATORS__ -U__CUDA_NO_HALF_CONVERSIONS__ -U__CUDA_NO_BFLOAT16_OPERATORS__ -U__CUDA_NO_BFLOAT16_CONVERSIONS__ -U__CUDA_NO_BFLOAT162_OPERATORS__ -U__CUDA_NO_BFLOAT162_CONVERSIONS__ --expt-relaxed-constexpr --expt-extended-lambda --use_fast_math --ptxas-options=-v -lineinfo -gencode arch=compute_70,code=sm_70 -gencode arch=compute_80,code=sm_80 -gencode arch=compute_90,code=sm_90 --threads 4 -DTORCH_API_INCLUDE_EXTENSION_H '-DPYBIND11_COMPILER_TYPE="_gcc"' '-DPYBIND11_STDLIB="_libstdcpp"' '-DPYBIND11_BUILD_ABI="_cxxabi1011"' -DTORCH_EXTENSION_NAME=selective_scan_cuda_oflex -D_GLIBCXX_USE_CXX11_ABI=0
+cuda_dlink_post_cflags = 
+ldflags = 
+
+rule compile
+  command = $cxx -MMD -MF $out.d $cflags -c $in -o $out $post_cflags
+  depfile = $out.d
+  deps = gcc
+
+rule cuda_compile
+  depfile = $out.d
+  deps = gcc
+  command = $nvcc  $cuda_cflags -c $in -o $out $cuda_post_cflags
+
+
+
+
+
+build /mnt/d/WORK/rschange-main/rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/csrc/selective_scan/cusoflex/selective_scan_core_bwd.o: cuda_compile /mnt/d/WORK/rschange-main/rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusoflex/selective_scan_core_bwd.cu
+build /mnt/d/WORK/rschange-main/rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/csrc/selective_scan/cusoflex/selective_scan_core_fwd.o: cuda_compile /mnt/d/WORK/rschange-main/rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusoflex/selective_scan_core_fwd.cu
+build /mnt/d/WORK/rschange-main/rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/csrc/selective_scan/cusoflex/selective_scan_oflex.o: compile /mnt/d/WORK/rschange-main/rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusoflex/selective_scan_oflex.cpp
+
+
+
+
+
+
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/csrc/selective_scan/cusoflex/selective_scan_core_bwd.o

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7a18749c332876fbeddf0356bbb0d4c979fffd89df7f4a27796e2dd39b523f2e
+size 12294744

rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/csrc/selective_scan/cusoflex/selective_scan_core_fwd.o

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:642208da032fd86584c184e13e9ed9d18a9c6e85d925770f7ce034e0d22774a9
+size 13211880

rscd/models/backbones/lib_mamba/kernels/selective_scan/build/temp.linux-x86_64-3.8/csrc/selective_scan/cusoflex/selective_scan_oflex.o

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0867e500c6352b2c0e1938ea0c8e6825aafbabc5699ec41d25a7793c56ed5d1e
+size 14839600

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cub_extra.cuh

@@ -0,0 +1,50 @@
+// WarpMask is copied from /usr/local/cuda-12.1/include/cub/util_ptx.cuh
+// PowerOfTwo is copied from /usr/local/cuda-12.1/include/cub/util_type.cuh
+
+#pragma once
+
+#include <cub/util_type.cuh>
+#include <cub/util_arch.cuh>
+#include <cub/util_namespace.cuh>
+#include <cub/util_debug.cuh>
+
+/**
+ * \brief Statically determine if N is a power-of-two
+ */
+ template <int N>
+ struct PowerOfTwo
+ {
+     enum { VALUE = ((N & (N - 1)) == 0) };
+ };
+ 
+
+/**
+ * @brief Returns the warp mask for a warp of @p LOGICAL_WARP_THREADS threads
+ *
+ * @par
+ * If the number of threads assigned to the virtual warp is not a power of two,
+ * it's assumed that only one virtual warp exists.
+ *
+ * @tparam LOGICAL_WARP_THREADS <b>[optional]</b> The number of threads per
+ *                              "logical" warp (may be less than the number of
+ *                              hardware warp threads).
+ * @param warp_id Id of virtual warp within architectural warp
+ */
+ template <int LOGICAL_WARP_THREADS, int LEGACY_PTX_ARCH = 0>
+ __host__ __device__ __forceinline__
+ unsigned int WarpMask(unsigned int warp_id)
+ {
+   constexpr bool is_pow_of_two = PowerOfTwo<LOGICAL_WARP_THREADS>::VALUE;
+   constexpr bool is_arch_warp  = LOGICAL_WARP_THREADS == CUB_WARP_THREADS(0);
+ 
+   unsigned int member_mask = 0xFFFFFFFFu >>
+                              (CUB_WARP_THREADS(0) - LOGICAL_WARP_THREADS);
+ 
+   if (is_pow_of_two && !is_arch_warp)
+   {
+     member_mask <<= warp_id * LOGICAL_WARP_THREADS;
+   }
+ 
+   return member_mask;
+ }
+ 

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cus/selective_scan.cpp

@@ -0,0 +1,354 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/extension.h>
+#include <vector>
+
+#include "selective_scan.h"
+#define MAX_DSTATE 256
+
+#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
+using weight_t = float;
+
+#define DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, NAME, ...)                    \
+    if (ITYPE == at::ScalarType::Half) {                                            \
+        using input_t = at::Half;                                                   \
+        __VA_ARGS__();                                                              \
+    } else if (ITYPE == at::ScalarType::BFloat16) {                                 \
+        using input_t = at::BFloat16;                                               \
+        __VA_ARGS__();                                                              \
+    } else if (ITYPE == at::ScalarType::Float)  {                                   \
+        using input_t = float;                                                      \
+        __VA_ARGS__();                                                              \
+    } else {                                                                        \
+        AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \
+    }
+
+template<int knrows, typename input_t, typename weight_t>
+void selective_scan_fwd_cuda(SSMParamsBase &params, cudaStream_t stream);
+
+template <int knrows, typename input_t, typename weight_t>
+void selective_scan_bwd_cuda(SSMParamsBwd &params, cudaStream_t stream);
+
+void set_ssm_params_fwd(SSMParamsBase &params,
+                        // sizes
+                        const size_t batch,
+                        const size_t dim,
+                        const size_t seqlen,
+                        const size_t dstate,
+                        const size_t n_groups,
+                        const size_t n_chunks,
+                        // device pointers
+                        const at::Tensor u,
+                        const at::Tensor delta,
+                        const at::Tensor A,
+                        const at::Tensor B,
+                        const at::Tensor C,
+                        const at::Tensor out,
+                        void* D_ptr,
+                        void* delta_bias_ptr,
+                        void* x_ptr,
+                        bool delta_softplus) {
+
+    // Reset the parameters
+    memset(&params, 0, sizeof(params));
+
+    params.batch = batch;
+    params.dim = dim;
+    params.seqlen = seqlen;
+    params.dstate = dstate;
+    params.n_groups = n_groups;
+    params.n_chunks = n_chunks;
+    params.dim_ngroups_ratio = dim / n_groups;
+
+    params.delta_softplus = delta_softplus;
+
+    // Set the pointers and strides.
+    params.u_ptr = u.data_ptr();
+    params.delta_ptr = delta.data_ptr();
+    params.A_ptr = A.data_ptr();
+    params.B_ptr = B.data_ptr();
+    params.C_ptr = C.data_ptr();
+    params.D_ptr = D_ptr;
+    params.delta_bias_ptr = delta_bias_ptr;
+    params.out_ptr = out.data_ptr();
+    params.x_ptr = x_ptr;
+
+    // All stride are in elements, not bytes.
+    params.A_d_stride = A.stride(0);
+    params.A_dstate_stride = A.stride(1);
+    params.B_batch_stride = B.stride(0);
+    params.B_group_stride = B.stride(1);
+    params.B_dstate_stride = B.stride(2);
+    params.C_batch_stride = C.stride(0);
+    params.C_group_stride = C.stride(1);
+    params.C_dstate_stride = C.stride(2);
+    params.u_batch_stride = u.stride(0);
+    params.u_d_stride = u.stride(1);
+    params.delta_batch_stride = delta.stride(0);
+    params.delta_d_stride = delta.stride(1);
+
+    params.out_batch_stride = out.stride(0);
+    params.out_d_stride = out.stride(1);
+}
+
+void set_ssm_params_bwd(SSMParamsBwd &params,
+                        // sizes
+                        const size_t batch,
+                        const size_t dim,
+                        const size_t seqlen,
+                        const size_t dstate,
+                        const size_t n_groups,
+                        const size_t n_chunks,
+                        // device pointers
+                        const at::Tensor u,
+                        const at::Tensor delta,
+                        const at::Tensor A,
+                        const at::Tensor B,
+                        const at::Tensor C,
+                        const at::Tensor out,
+                        void* D_ptr,
+                        void* delta_bias_ptr,
+                        void* x_ptr,
+                        const at::Tensor dout,
+                        const at::Tensor du,
+                        const at::Tensor ddelta,
+                        const at::Tensor dA,
+                        const at::Tensor dB,
+                        const at::Tensor dC,
+                        void* dD_ptr,
+                        void* ddelta_bias_ptr,
+                        bool delta_softplus) {
+    // Pass in "dout" instead of "out", we're not gonna use "out" unless we have z
+    set_ssm_params_fwd(params, batch, dim, seqlen, dstate, n_groups, n_chunks,
+                       u, delta, A, B, C, dout,
+                       D_ptr, delta_bias_ptr, x_ptr, delta_softplus);
+
+    // Set the pointers and strides.
+    params.dout_ptr = dout.data_ptr();
+    params.du_ptr = du.data_ptr();
+    params.dA_ptr = dA.data_ptr();
+    params.dB_ptr = dB.data_ptr();
+    params.dC_ptr = dC.data_ptr();
+    params.dD_ptr = dD_ptr;
+    params.ddelta_ptr = ddelta.data_ptr();
+    params.ddelta_bias_ptr = ddelta_bias_ptr;
+    // All stride are in elements, not bytes.
+    params.dout_batch_stride = dout.stride(0);
+    params.dout_d_stride = dout.stride(1);
+    params.dA_d_stride = dA.stride(0);
+    params.dA_dstate_stride = dA.stride(1);
+    params.dB_batch_stride = dB.stride(0);
+    params.dB_group_stride = dB.stride(1);
+    params.dB_dstate_stride = dB.stride(2);
+    params.dC_batch_stride = dC.stride(0);
+    params.dC_group_stride = dC.stride(1);
+    params.dC_dstate_stride = dC.stride(2);
+    params.du_batch_stride = du.stride(0);
+    params.du_d_stride = du.stride(1);
+    params.ddelta_batch_stride = ddelta.stride(0);
+    params.ddelta_d_stride = ddelta.stride(1);
+
+}
+
+std::vector<at::Tensor>
+selective_scan_fwd(const at::Tensor &u, const at::Tensor &delta,
+                  const at::Tensor &A, const at::Tensor &B, const at::Tensor &C,
+                  const c10::optional<at::Tensor> &D_,
+                  const c10::optional<at::Tensor> &delta_bias_,
+                  bool delta_softplus,
+                  int nrows
+                  ) {
+    auto input_type = u.scalar_type();
+    auto weight_type = A.scalar_type();
+    TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
+    TORCH_CHECK(weight_type == at::ScalarType::Float);
+
+    TORCH_CHECK(delta.scalar_type() == input_type);
+    TORCH_CHECK(B.scalar_type() == input_type);
+    TORCH_CHECK(C.scalar_type() == input_type);
+
+    TORCH_CHECK(u.is_cuda());
+    TORCH_CHECK(delta.is_cuda());
+    TORCH_CHECK(A.is_cuda());
+    TORCH_CHECK(B.is_cuda());
+    TORCH_CHECK(C.is_cuda());
+
+    TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
+    TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
+
+    const auto sizes = u.sizes();
+    const int batch_size = sizes[0];
+    const int dim = sizes[1];
+    const int seqlen = sizes[2];
+    const int dstate = A.size(1);
+    const int n_groups = B.size(1);
+
+    TORCH_CHECK(dim % n_groups == 0, "dims should be dividable by n_groups");
+    TORCH_CHECK(dstate <= MAX_DSTATE, "selective_scan only supports state dimension <= 256");
+
+    CHECK_SHAPE(u, batch_size, dim, seqlen);
+    CHECK_SHAPE(delta, batch_size, dim, seqlen);
+    CHECK_SHAPE(A, dim, dstate);
+    CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen);
+    TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
+    CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen);
+    TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
+
+    if (D_.has_value()) {
+        auto D = D_.value();
+        TORCH_CHECK(D.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(D.is_cuda());
+        TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
+        CHECK_SHAPE(D, dim);
+    }
+
+    if (delta_bias_.has_value()) {
+        auto delta_bias = delta_bias_.value();
+        TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(delta_bias.is_cuda());
+        TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
+        CHECK_SHAPE(delta_bias, dim);
+    }
+
+    const int n_chunks = (seqlen + 2048 - 1) / 2048; // max is 128 * 16 = 2048 in fwd_kernel
+    at::Tensor out = torch::empty_like(delta);
+    at::Tensor x;
+    x = torch::empty({batch_size, dim, n_chunks, dstate * 2}, u.options().dtype(weight_type));
+
+    SSMParamsBase params;
+    set_ssm_params_fwd(params, batch_size, dim, seqlen, dstate, n_groups, n_chunks,
+                       u, delta, A, B, C, out,
+                       D_.has_value() ? D_.value().data_ptr() : nullptr,
+                       delta_bias_.has_value() ? delta_bias_.value().data_ptr() : nullptr,
+                       x.data_ptr(),
+                       delta_softplus);
+
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)u.get_device()};
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), "selective_scan_fwd", [&] {
+        selective_scan_fwd_cuda<1, input_t, weight_t>(params, stream);
+    });
+    std::vector<at::Tensor> result = {out, x};
+    return result;
+}
+
+std::vector<at::Tensor>
+selective_scan_bwd(const at::Tensor &u, const at::Tensor &delta,
+                  const at::Tensor &A, const at::Tensor &B, const at::Tensor &C,
+                  const c10::optional<at::Tensor> &D_,
+                  const c10::optional<at::Tensor> &delta_bias_,
+                  const at::Tensor &dout,
+                  const c10::optional<at::Tensor> &x_,
+                  bool delta_softplus,
+                  int nrows
+                  ) {
+    auto input_type = u.scalar_type();
+    auto weight_type = A.scalar_type();
+    TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
+    TORCH_CHECK(weight_type == at::ScalarType::Float);
+
+    TORCH_CHECK(delta.scalar_type() == input_type);
+    TORCH_CHECK(B.scalar_type() == input_type);
+    TORCH_CHECK(C.scalar_type() == input_type);
+    TORCH_CHECK(dout.scalar_type() == input_type);
+
+    TORCH_CHECK(u.is_cuda());
+    TORCH_CHECK(delta.is_cuda());
+    TORCH_CHECK(A.is_cuda());
+    TORCH_CHECK(B.is_cuda());
+    TORCH_CHECK(C.is_cuda());
+    TORCH_CHECK(dout.is_cuda());
+
+    TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
+    TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
+    TORCH_CHECK(dout.stride(-1) == 1 || dout.size(-1) == 1);
+
+    const auto sizes = u.sizes();
+    const int batch_size = sizes[0];
+    const int dim = sizes[1];
+    const int seqlen = sizes[2];
+    const int dstate = A.size(1);
+    const int n_groups = B.size(1);
+
+    TORCH_CHECK(dim % n_groups == 0, "dims should be dividable by n_groups");
+    TORCH_CHECK(dstate <= MAX_DSTATE, "selective_scan only supports state dimension <= 256");
+    
+    CHECK_SHAPE(u, batch_size, dim, seqlen);
+    CHECK_SHAPE(delta, batch_size, dim, seqlen);
+    CHECK_SHAPE(A, dim, dstate);
+    CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen);
+    TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
+    CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen);
+    TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
+    CHECK_SHAPE(dout, batch_size, dim, seqlen);
+
+    if (D_.has_value()) {
+        auto D = D_.value();
+        TORCH_CHECK(D.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(D.is_cuda());
+        TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
+        CHECK_SHAPE(D, dim);
+    }
+
+    if (delta_bias_.has_value()) {
+        auto delta_bias = delta_bias_.value();
+        TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(delta_bias.is_cuda());
+        TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
+        CHECK_SHAPE(delta_bias, dim);
+    }
+
+    at::Tensor out;
+    const int n_chunks = (seqlen + 2048 - 1) / 2048;
+    // const int n_chunks = (seqlen + 1024 - 1) / 1024;
+    if (n_chunks > 1) { TORCH_CHECK(x_.has_value()); }
+    if (x_.has_value()) {
+        auto x = x_.value();
+        TORCH_CHECK(x.scalar_type() == weight_type);
+        TORCH_CHECK(x.is_cuda());
+        TORCH_CHECK(x.is_contiguous());
+        CHECK_SHAPE(x, batch_size, dim, n_chunks, 2 * dstate);
+    }
+
+    at::Tensor du = torch::empty_like(u);
+    at::Tensor ddelta = torch::empty_like(delta);
+    at::Tensor dA = torch::zeros_like(A);
+    at::Tensor dB = torch::zeros_like(B, B.options().dtype(torch::kFloat32));
+    at::Tensor dC = torch::zeros_like(C, C.options().dtype(torch::kFloat32));
+    at::Tensor dD;
+    if (D_.has_value()) { dD = torch::zeros_like(D_.value()); }
+    at::Tensor ddelta_bias;
+    if (delta_bias_.has_value()) { ddelta_bias = torch::zeros_like(delta_bias_.value()); }
+
+    SSMParamsBwd params;
+    set_ssm_params_bwd(params, batch_size, dim, seqlen, dstate, n_groups, n_chunks,
+                       u, delta, A, B, C, out,
+                       D_.has_value() ? D_.value().data_ptr() : nullptr,
+                       delta_bias_.has_value() ? delta_bias_.value().data_ptr() : nullptr,
+                       x_.has_value() ? x_.value().data_ptr() : nullptr,
+                       dout, du, ddelta, dA, dB, dC,
+                       D_.has_value() ? dD.data_ptr() : nullptr,
+                       delta_bias_.has_value() ? ddelta_bias.data_ptr() : nullptr,
+                       delta_softplus);
+
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)u.get_device()};
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), "selective_scan_bwd", [&] {
+        selective_scan_bwd_cuda<1, input_t, weight_t>(params, stream);
+    });
+    std::vector<at::Tensor> result = {du, ddelta, dA, dB.to(B.dtype()), dC.to(C.dtype()), dD, ddelta_bias};
+    return result;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("fwd", &selective_scan_fwd, "Selective scan forward");
+    m.def("bwd", &selective_scan_bwd, "Selective scan backward");
+}

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cus/selective_scan_bwd_kernel.cuh

@@ -0,0 +1,306 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <c10/util/BFloat16.h>
+#include <c10/util/Half.h>
+#include <c10/cuda/CUDAException.h>  // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
+#include <ATen/cuda/Atomic.cuh>  // For atomicAdd on complex
+
+#include <cub/block/block_load.cuh>
+#include <cub/block/block_store.cuh>
+#include <cub/block/block_scan.cuh>
+#include <cub/block/block_reduce.cuh>
+
+#include "selective_scan.h"
+#include "selective_scan_common.h"
+#include "reverse_scan.cuh"
+#include "static_switch.h"
+
+template<int kNThreads_, int kNItems_, bool kIsEvenLen_, bool kDeltaSoftplus_, typename input_t_, typename weight_t_>
+struct Selective_Scan_bwd_kernel_traits {
+    static_assert(kNItems_ % 4 == 0);
+    using input_t = input_t_;
+    using weight_t = weight_t_;
+    static constexpr int kNThreads = kNThreads_;
+    static constexpr int kNItems = kNItems_;
+    static constexpr int MaxDState = MAX_DSTATE;
+    static constexpr int kNBytes = sizeof(input_t);
+    static_assert(kNBytes == 2 || kNBytes == 4);
+    static constexpr int kNElts = kNBytes == 4 ? 4 : std::min(8, kNItems);
+    static_assert(kNItems % kNElts == 0);
+    static constexpr int kNLoads = kNItems / kNElts;
+    static constexpr bool kIsEvenLen = kIsEvenLen_;
+    static constexpr bool kDeltaSoftplus = kDeltaSoftplus_;
+    // Setting MinBlocksPerMP to be 3 (instead of 2) for 128 threads with float improves occupancy.
+    // For complex this would lead to massive register spilling, so we keep it at 2.
+    static constexpr int kMinBlocks = kNThreads == 128 && 3;
+    using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
+    using scan_t = float2;
+    using BlockLoadT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadWeightT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadWeightVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockStoreT = cub::BlockStore<input_t, kNThreads, kNItems, cub::BLOCK_STORE_WARP_TRANSPOSE>;
+    using BlockStoreVecT = cub::BlockStore<vec_t, kNThreads, kNLoads, cub::BLOCK_STORE_WARP_TRANSPOSE>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING_MEMOIZE>;
+    using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_WARP_SCANS>;
+    using BlockReverseScanT = BlockReverseScan<scan_t, kNThreads>;
+    using BlockReduceT = cub::BlockReduce<scan_t, kNThreads>;
+    using BlockReduceFloatT = cub::BlockReduce<float, kNThreads>;
+    using BlockExchangeT = cub::BlockExchange<float, kNThreads, kNItems>;
+    static constexpr int kSmemIOSize = std::max({sizeof(typename BlockLoadT::TempStorage),
+                                                 sizeof(typename BlockLoadVecT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightVecT::TempStorage),
+                                                 sizeof(typename BlockStoreT::TempStorage),
+                                                 sizeof(typename BlockStoreVecT::TempStorage)});
+    static constexpr int kSmemExchangeSize = 2 * sizeof(typename BlockExchangeT::TempStorage);
+    static constexpr int kSmemReduceSize = sizeof(typename BlockReduceT::TempStorage);
+    static constexpr int kSmemSize = kSmemIOSize + kSmemExchangeSize + kSmemReduceSize + sizeof(typename BlockScanT::TempStorage) + sizeof(typename BlockReverseScanT::TempStorage);
+};
+
+template<typename Ktraits>
+__global__ __launch_bounds__(Ktraits::kNThreads, Ktraits::kMinBlocks)
+void selective_scan_bwd_kernel(SSMParamsBwd params) {
+    constexpr bool kDeltaSoftplus = Ktraits::kDeltaSoftplus;
+    constexpr int kNThreads = Ktraits::kNThreads;
+    constexpr int kNItems = Ktraits::kNItems;
+    using input_t = typename Ktraits::input_t;
+    using weight_t = typename Ktraits::weight_t;
+    using scan_t = typename Ktraits::scan_t;
+
+    // Shared memory.
+    extern __shared__ char smem_[];
+    auto& smem_load = reinterpret_cast<typename Ktraits::BlockLoadT::TempStorage&>(smem_);
+    auto& smem_load_weight = reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage&>(smem_);
+    auto& smem_load_weight1 = *reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage*>(smem_ + sizeof(typename Ktraits::BlockLoadWeightT::TempStorage));
+    auto& smem_store = reinterpret_cast<typename Ktraits::BlockStoreT::TempStorage&>(smem_);
+    auto& smem_exchange = *reinterpret_cast<typename Ktraits::BlockExchangeT::TempStorage*>(smem_ + Ktraits::kSmemIOSize);
+    auto& smem_exchange1 = *reinterpret_cast<typename Ktraits::BlockExchangeT::TempStorage*>(smem_ + Ktraits::kSmemIOSize + sizeof(typename Ktraits::BlockExchangeT::TempStorage));
+    auto& smem_reduce = *reinterpret_cast<typename Ktraits::BlockReduceT::TempStorage*>(reinterpret_cast<char *>(&smem_exchange) + Ktraits::kSmemExchangeSize);
+    auto& smem_reduce_float = *reinterpret_cast<typename Ktraits::BlockReduceFloatT::TempStorage*>(&smem_reduce);
+    auto& smem_scan = *reinterpret_cast<typename Ktraits::BlockScanT::TempStorage*>(reinterpret_cast<char *>(&smem_reduce) + Ktraits::kSmemReduceSize);
+    auto& smem_reverse_scan = *reinterpret_cast<typename Ktraits::BlockReverseScanT::TempStorage*>(reinterpret_cast<char *>(&smem_scan) + sizeof(typename Ktraits::BlockScanT::TempStorage));
+    weight_t *smem_delta_a = reinterpret_cast<weight_t *>(smem_ + Ktraits::kSmemSize);
+    scan_t *smem_running_postfix = reinterpret_cast<scan_t *>(smem_delta_a + 2 * Ktraits::MaxDState + kNThreads);
+    weight_t *smem_da = reinterpret_cast<weight_t *>(smem_running_postfix + Ktraits::MaxDState);
+
+    const int batch_id = blockIdx.x;
+    const int dim_id = blockIdx.y;
+    const int group_id = dim_id / (params.dim_ngroups_ratio);
+    input_t *u = reinterpret_cast<input_t *>(params.u_ptr) + batch_id * params.u_batch_stride
+        + dim_id * params.u_d_stride;
+    input_t *delta = reinterpret_cast<input_t *>(params.delta_ptr) + batch_id * params.delta_batch_stride
+        + dim_id * params.delta_d_stride;
+    input_t *dout = reinterpret_cast<input_t *>(params.dout_ptr) + batch_id * params.dout_batch_stride
+        + dim_id * params.dout_d_stride;
+    weight_t *A = reinterpret_cast<weight_t *>(params.A_ptr) + dim_id * params.A_d_stride;
+    input_t *Bvar = reinterpret_cast<input_t *>(params.B_ptr) + batch_id * params.B_batch_stride + group_id * params.B_group_stride;
+    input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + batch_id * params.C_batch_stride + group_id * params.C_group_stride;
+    weight_t *dA = reinterpret_cast<weight_t *>(params.dA_ptr) + dim_id * params.dA_d_stride;
+    weight_t *dB = reinterpret_cast<weight_t *>(params.dB_ptr)
+        + (batch_id * params.dB_batch_stride + group_id * params.dB_group_stride);
+    weight_t *dC = reinterpret_cast<weight_t *>(params.dC_ptr)
+        + (batch_id * params.dC_batch_stride + group_id * params.dC_group_stride);
+    float *dD = params.dD_ptr == nullptr ? nullptr : reinterpret_cast<float *>(params.dD_ptr) + dim_id;
+    float D_val = params.D_ptr == nullptr ? 0 : reinterpret_cast<float *>(params.D_ptr)[dim_id];
+    float *ddelta_bias = params.ddelta_bias_ptr == nullptr ? nullptr : reinterpret_cast<float *>(params.ddelta_bias_ptr) + dim_id;
+    float delta_bias = params.delta_bias_ptr == nullptr ? 0 : reinterpret_cast<float *>(params.delta_bias_ptr)[dim_id];
+    scan_t *x = params.x_ptr == nullptr
+        ? nullptr
+        : reinterpret_cast<scan_t *>(params.x_ptr) + (batch_id * params.dim + dim_id) * (params.n_chunks) * params.dstate;
+    float dD_val = 0;
+    float ddelta_bias_val = 0;
+
+    constexpr int kChunkSize = kNThreads * kNItems;
+    u += (params.n_chunks - 1) * kChunkSize;
+    delta += (params.n_chunks - 1) * kChunkSize;
+    dout += (params.n_chunks - 1) * kChunkSize;
+    Bvar += (params.n_chunks - 1) * kChunkSize;
+    Cvar += (params.n_chunks - 1) * kChunkSize;
+    for (int chunk = params.n_chunks - 1; chunk >= 0; --chunk) {
+        input_t u_vals[kNItems];
+        input_t delta_vals_load[kNItems];
+        input_t dout_vals_load[kNItems];
+        __syncthreads();
+        load_input<Ktraits>(u, u_vals, smem_load, params.seqlen - chunk * kChunkSize);
+        __syncthreads();
+        load_input<Ktraits>(delta, delta_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
+        __syncthreads();
+        load_input<Ktraits>(dout, dout_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
+        u -= kChunkSize;
+        // Will reload delta at the same location if kDeltaSoftplus
+        if constexpr (!kDeltaSoftplus) { delta -= kChunkSize; }
+        dout -= kChunkSize;
+
+        float dout_vals[kNItems], delta_vals[kNItems];
+        float du_vals[kNItems];
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) {
+            dout_vals[i] = float(dout_vals_load[i]);
+            delta_vals[i] = float(delta_vals_load[i]) + delta_bias;
+            if constexpr (kDeltaSoftplus) {
+                delta_vals[i] = delta_vals[i] <= 20.f ? log1pf(expf(delta_vals[i])) : delta_vals[i];
+            }
+        }
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) { du_vals[i] = D_val * dout_vals[i]; }
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) { dD_val += dout_vals[i] * float(u_vals[i]); }
+
+        float ddelta_vals[kNItems] = {0};
+        __syncthreads();
+        for (int state_idx = 0; state_idx < params.dstate; ++state_idx) {
+            constexpr float kLog2e = M_LOG2E;
+            weight_t A_val = A[state_idx * params.A_dstate_stride];
+            weight_t A_scaled = A_val * kLog2e;
+            weight_t B_vals[kNItems], C_vals[kNItems];
+            load_weight<Ktraits>(Bvar + state_idx * params.B_dstate_stride, B_vals,
+                    smem_load_weight, (params.seqlen - chunk * kChunkSize));
+            auto &smem_load_weight_C = smem_load_weight1;
+            load_weight<Ktraits>(Cvar + state_idx * params.C_dstate_stride, C_vals,
+                    smem_load_weight_C, (params.seqlen - chunk * kChunkSize));
+            scan_t thread_data[kNItems], thread_reverse_data[kNItems];
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                const float delta_a_exp = exp2f(delta_vals[i] * A_scaled);
+                thread_data[i] = make_float2(delta_a_exp, delta_vals[i] * float(u_vals[i]) * B_vals[i]);
+                if (i == 0) {
+                    smem_delta_a[threadIdx.x == 0 ? state_idx + (chunk % 2) * Ktraits::MaxDState: threadIdx.x + 2 * Ktraits::MaxDState] = delta_a_exp;
+                } else {
+                    thread_reverse_data[i - 1].x = delta_a_exp;
+                }
+                thread_reverse_data[i].y = dout_vals[i] * C_vals[i];
+            }
+            __syncthreads();
+            thread_reverse_data[kNItems - 1].x = threadIdx.x == kNThreads - 1
+                ? (chunk == params.n_chunks - 1 ? 1.f : smem_delta_a[state_idx + ((chunk + 1) % 2) * Ktraits::MaxDState])
+                    : smem_delta_a[threadIdx.x + 1 + 2 * Ktraits::MaxDState];
+            // Initialize running total
+            scan_t running_prefix = chunk > 0 && threadIdx.x % 32 == 0 ? x[(chunk - 1) * params.dstate + state_idx] : make_float2(1.f, 0.f);
+            SSMScanPrefixCallbackOp<weight_t> prefix_op(running_prefix);
+            Ktraits::BlockScanT(smem_scan).InclusiveScan(
+                thread_data, thread_data, SSMScanOp<weight_t>(), prefix_op
+            );
+            scan_t running_postfix = chunk < params.n_chunks - 1 && threadIdx.x % 32 == 0 ? smem_running_postfix[state_idx] : make_float2(1.f, 0.f);
+            SSMScanPrefixCallbackOp<weight_t> postfix_op(running_postfix);
+            Ktraits::BlockReverseScanT(smem_reverse_scan).InclusiveReverseScan(
+                thread_reverse_data, thread_reverse_data, SSMScanOp<weight_t>(), postfix_op
+            );
+            if (threadIdx.x == 0) { smem_running_postfix[state_idx] = postfix_op.running_prefix; }
+            weight_t dA_val = 0;
+            weight_t dB_vals[kNItems], dC_vals[kNItems];
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                const float dx = thread_reverse_data[i].y;
+                const float ddelta_u = dx * B_vals[i];
+                du_vals[i] += ddelta_u * delta_vals[i];
+                const float a = thread_data[i].y - (delta_vals[i] * float(u_vals[i]) * B_vals[i]);
+                ddelta_vals[i] += ddelta_u * float(u_vals[i]) + dx * A_val * a;
+                dA_val += dx * delta_vals[i] * a;
+                dB_vals[i] = dx * delta_vals[i] * float(u_vals[i]);
+                dC_vals[i] = dout_vals[i] * thread_data[i].y;
+            }
+            // Block-exchange to make the atomicAdd's coalesced, otherwise they're much slower
+            Ktraits::BlockExchangeT(smem_exchange).BlockedToStriped(dB_vals, dB_vals);
+            auto &smem_exchange_C = smem_exchange1;
+            Ktraits::BlockExchangeT(smem_exchange_C).BlockedToStriped(dC_vals, dC_vals);
+            const int seqlen_remaining = params.seqlen - chunk * kChunkSize - threadIdx.x;
+            weight_t *dB_cur = dB + state_idx * params.dB_dstate_stride + chunk * kChunkSize + threadIdx.x;
+            weight_t *dC_cur = dC + state_idx * params.dC_dstate_stride + chunk * kChunkSize + threadIdx.x;
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                if (i * kNThreads < seqlen_remaining) {
+                    { gpuAtomicAdd(dB_cur + i * kNThreads, dB_vals[i]); }
+                    { gpuAtomicAdd(dC_cur + i * kNThreads, dC_vals[i]); }
+                }
+            }
+            dA_val = Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(dA_val);
+            if (threadIdx.x == 0) {
+                smem_da[state_idx] = chunk == params.n_chunks - 1 ? dA_val : dA_val + smem_da[state_idx];
+            }
+        }
+
+        if constexpr (kDeltaSoftplus) {
+            input_t delta_vals_load[kNItems];
+            __syncthreads();
+            load_input<Ktraits>(delta, delta_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
+            delta -= kChunkSize;
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                float delta_val = float(delta_vals_load[i]) + delta_bias;
+                float delta_val_neg_exp = expf(-delta_val);
+                ddelta_vals[i] = delta_val <= 20.f
+                    ? ddelta_vals[i] / (1.f + delta_val_neg_exp)
+                        : ddelta_vals[i];
+            }
+        }
+        
+        __syncthreads();
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) { ddelta_bias_val += ddelta_vals[i]; }
+
+        input_t *du = reinterpret_cast<input_t *>(params.du_ptr) + batch_id * params.du_batch_stride
+            + dim_id * params.du_d_stride + chunk * kChunkSize;
+        input_t *ddelta = reinterpret_cast<input_t *>(params.ddelta_ptr) + batch_id * params.ddelta_batch_stride
+            + dim_id * params.ddelta_d_stride + chunk * kChunkSize;
+        __syncthreads();
+        store_output<Ktraits>(du, du_vals, smem_store, params.seqlen - chunk * kChunkSize);
+        __syncthreads();
+        store_output<Ktraits>(ddelta, ddelta_vals, smem_store, params.seqlen - chunk * kChunkSize);
+        Bvar -= kChunkSize;
+        Cvar -= kChunkSize;
+    }
+
+    if (params.dD_ptr != nullptr) {
+        __syncthreads();
+        dD_val = Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(dD_val);
+        if (threadIdx.x == 0) { gpuAtomicAdd(dD, dD_val); }
+    }
+    if (params.ddelta_bias_ptr != nullptr) {
+        __syncthreads();
+        ddelta_bias_val = Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(ddelta_bias_val);
+        if (threadIdx.x == 0) { gpuAtomicAdd(ddelta_bias, ddelta_bias_val); }
+    }
+    __syncthreads();
+    for (int state_idx = threadIdx.x; state_idx < params.dstate; state_idx += blockDim.x) {
+        gpuAtomicAdd(&(dA[state_idx * params.dA_dstate_stride]), smem_da[state_idx]);
+    }
+}
+
+template<int kNThreads, int kNItems, typename input_t, typename weight_t>
+void selective_scan_bwd_launch(SSMParamsBwd &params, cudaStream_t stream) {
+    BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {
+        BOOL_SWITCH(params.delta_softplus, kDeltaSoftplus, [&] {
+            using Ktraits = Selective_Scan_bwd_kernel_traits<kNThreads, kNItems, kIsEvenLen, kDeltaSoftplus, input_t, weight_t>;
+            constexpr int kSmemSize = Ktraits::kSmemSize + Ktraits::MaxDState * sizeof(typename Ktraits::scan_t) + (kNThreads + 4 * Ktraits::MaxDState) * sizeof(typename Ktraits::weight_t);
+            // printf("smem_size = %d\n", kSmemSize);
+            dim3 grid(params.batch, params.dim);
+            auto kernel = &selective_scan_bwd_kernel<Ktraits>;
+            if (kSmemSize >= 48 * 1024) {
+                C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
+            }
+            kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
+            C10_CUDA_KERNEL_LAUNCH_CHECK();
+        });
+    });
+}
+
+template<int knrows, typename input_t, typename weight_t>
+void selective_scan_bwd_cuda(SSMParamsBwd &params, cudaStream_t stream) {
+    if (params.seqlen <= 128) {
+        selective_scan_bwd_launch<32, 4, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 256) {
+        selective_scan_bwd_launch<32, 8, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 512) {
+        selective_scan_bwd_launch<32, 16, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 1024) {
+        selective_scan_bwd_launch<64, 16, input_t, weight_t>(params, stream);
+    } else {
+        selective_scan_bwd_launch<128, 16, input_t, weight_t>(params, stream);
+    }
+}

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cus/selective_scan_core_bwd.cu

@@ -0,0 +1,9 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+#include "selective_scan_bwd_kernel.cuh"
+
+template void selective_scan_bwd_cuda<1, float, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<1, at::Half, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<1, at::BFloat16, float>(SSMParamsBwd &params, cudaStream_t stream);
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cus/selective_scan_core_fwd.cu

@@ -0,0 +1,9 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+#include "selective_scan_fwd_kernel.cuh"
+
+template void selective_scan_fwd_cuda<1, float, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<1, at::Half, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<1, at::BFloat16, float>(SSMParamsBase &params, cudaStream_t stream);
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cus/selective_scan_fwd_kernel.cuh

@@ -0,0 +1,203 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <c10/util/BFloat16.h>
+#include <c10/util/Half.h>
+#include <c10/cuda/CUDAException.h>  // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
+
+#include <cub/block/block_load.cuh>
+#include <cub/block/block_store.cuh>
+#include <cub/block/block_scan.cuh>
+
+#include "selective_scan.h"
+#include "selective_scan_common.h"
+#include "static_switch.h"
+
+template<int kNThreads_, int kNItems_, bool kIsEvenLen_, typename input_t_, typename weight_t_>
+struct Selective_Scan_fwd_kernel_traits {
+    static_assert(kNItems_ % 4 == 0);
+    using input_t = input_t_;
+    using weight_t = weight_t_;
+    static constexpr int kNThreads = kNThreads_;
+    // Setting MinBlocksPerMP to be 3 (instead of 2) for 128 threads improves occupancy.
+    static constexpr int kMinBlocks = kNThreads < 128 ? 5 : 3;
+    static constexpr int kNItems = kNItems_;
+    static constexpr int MaxDState = MAX_DSTATE;
+    static constexpr int kNBytes = sizeof(input_t);
+    static_assert(kNBytes == 2 || kNBytes == 4);
+    static constexpr int kNElts = kNBytes == 4 ? 4 : std::min(8, kNItems);
+    static_assert(kNItems % kNElts == 0);
+    static constexpr int kNLoads = kNItems / kNElts;
+    static constexpr bool kIsEvenLen = kIsEvenLen_;
+
+    static constexpr bool kDirectIO = kIsEvenLen && kNLoads == 1;
+
+    using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
+    using scan_t = float2;
+    using BlockLoadT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads,
+        !kDirectIO ? cub::BLOCK_LOAD_WARP_TRANSPOSE : cub::BLOCK_LOAD_DIRECT>;
+    using BlockLoadWeightT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadWeightVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads,
+        !kDirectIO ? cub::BLOCK_LOAD_WARP_TRANSPOSE  : cub::BLOCK_LOAD_DIRECT>;
+    using BlockStoreT = cub::BlockStore<input_t, kNThreads, kNItems, cub::BLOCK_STORE_WARP_TRANSPOSE>;
+    using BlockStoreVecT = cub::BlockStore<vec_t, kNThreads, kNLoads,
+        !kDirectIO ? cub::BLOCK_STORE_WARP_TRANSPOSE : cub::BLOCK_STORE_DIRECT>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING_MEMOIZE>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING>;
+    using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_WARP_SCANS>;
+    static constexpr int kSmemIOSize = std::max({sizeof(typename BlockLoadT::TempStorage),
+                                                 sizeof(typename BlockLoadVecT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightVecT::TempStorage),
+                                                 sizeof(typename BlockStoreT::TempStorage),
+                                                 sizeof(typename BlockStoreVecT::TempStorage)});
+    static constexpr int kSmemSize = kSmemIOSize + sizeof(typename BlockScanT::TempStorage);
+};
+
+template<typename Ktraits>
+__global__ __launch_bounds__(Ktraits::kNThreads, Ktraits::kMinBlocks)
+void selective_scan_fwd_kernel(SSMParamsBase params) {
+    constexpr int kNThreads = Ktraits::kNThreads;
+    constexpr int kNItems = Ktraits::kNItems;
+    constexpr bool kDirectIO = Ktraits::kDirectIO;
+    using input_t = typename Ktraits::input_t;
+    using weight_t = typename Ktraits::weight_t;
+    using scan_t = typename Ktraits::scan_t;
+
+    // Shared memory.
+    extern __shared__ char smem_[];
+    auto& smem_load = reinterpret_cast<typename Ktraits::BlockLoadT::TempStorage&>(smem_);
+    auto& smem_load_weight = reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage&>(smem_);
+    auto& smem_load_weight1 = *reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage*>(smem_ + sizeof(typename Ktraits::BlockLoadWeightT::TempStorage));
+    auto& smem_store = reinterpret_cast<typename Ktraits::BlockStoreT::TempStorage&>(smem_);
+    auto& smem_scan = *reinterpret_cast<typename Ktraits::BlockScanT::TempStorage*>(smem_ + Ktraits::kSmemIOSize);
+    scan_t *smem_running_prefix = reinterpret_cast<scan_t *>(smem_ + Ktraits::kSmemSize);
+
+    const int batch_id = blockIdx.x;
+    const int dim_id = blockIdx.y;
+    const int group_id = dim_id / (params.dim_ngroups_ratio);
+    input_t *u = reinterpret_cast<input_t *>(params.u_ptr) + batch_id * params.u_batch_stride
+        + dim_id * params.u_d_stride;
+    input_t *delta = reinterpret_cast<input_t *>(params.delta_ptr) + batch_id * params.delta_batch_stride
+        + dim_id * params.delta_d_stride;
+    weight_t *A = reinterpret_cast<weight_t *>(params.A_ptr) + dim_id * params.A_d_stride;
+    input_t *Bvar = reinterpret_cast<input_t *>(params.B_ptr) + batch_id * params.B_batch_stride + group_id * params.B_group_stride;
+    input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + batch_id * params.C_batch_stride + group_id * params.C_group_stride;
+    scan_t *x = reinterpret_cast<scan_t *>(params.x_ptr) + (batch_id * params.dim + dim_id) * params.n_chunks * params.dstate;
+
+    float D_val = 0; // attention!
+    if (params.D_ptr != nullptr) {
+        D_val = reinterpret_cast<float *>(params.D_ptr)[dim_id];
+    }
+    float delta_bias = 0;
+    if (params.delta_bias_ptr != nullptr) {
+        delta_bias = reinterpret_cast<float *>(params.delta_bias_ptr)[dim_id];
+    }
+
+    constexpr int kChunkSize = kNThreads * kNItems;
+    for (int chunk = 0; chunk < params.n_chunks; ++chunk) {
+        input_t u_vals[kNItems], delta_vals_load[kNItems];
+        __syncthreads();
+        load_input<Ktraits>(u, u_vals, smem_load, params.seqlen - chunk * kChunkSize);
+        if constexpr (!kDirectIO) { __syncthreads(); }
+        load_input<Ktraits>(delta, delta_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
+        u += kChunkSize;
+        delta += kChunkSize;
+
+        float delta_vals[kNItems], delta_u_vals[kNItems], out_vals[kNItems];
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) {
+            float u_val = float(u_vals[i]);
+            delta_vals[i] = float(delta_vals_load[i]) + delta_bias;
+            if (params.delta_softplus) {
+                delta_vals[i] = delta_vals[i] <= 20.f ? log1pf(expf(delta_vals[i])) : delta_vals[i];
+            }
+            delta_u_vals[i] = delta_vals[i] * u_val;
+            out_vals[i] = D_val * u_val;
+        }
+
+        __syncthreads();
+        for (int state_idx = 0; state_idx < params.dstate; ++state_idx) {
+            constexpr float kLog2e = M_LOG2E;
+            weight_t A_val = A[state_idx * params.A_dstate_stride];
+            A_val *= kLog2e;
+            weight_t B_vals[kNItems], C_vals[kNItems];
+            load_weight<Ktraits>(Bvar + state_idx * params.B_dstate_stride, B_vals,
+                    smem_load_weight, (params.seqlen - chunk * kChunkSize));
+            load_weight<Ktraits>(Cvar + state_idx * params.C_dstate_stride, C_vals,
+                    smem_load_weight1, (params.seqlen - chunk * kChunkSize));
+            __syncthreads();
+            scan_t thread_data[kNItems];
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                thread_data[i] = make_float2(exp2f(delta_vals[i] * A_val), B_vals[i] * delta_u_vals[i]);
+                if constexpr (!Ktraits::kIsEvenLen) {  // So that the last state is correct
+                    if (threadIdx.x * kNItems + i >= params.seqlen - chunk * kChunkSize) {
+                        thread_data[i] = make_float2(1.f, 0.f);
+                    }
+                }
+            }
+            // Initialize running total
+            scan_t running_prefix;
+            // If we use WARP_SCAN then all lane 0 of all warps (not just thread 0) needs to read
+            running_prefix = chunk > 0 && threadIdx.x % 32 == 0 ? smem_running_prefix[state_idx] : make_float2(1.f, 0.f);
+            // running_prefix = chunk > 0 && threadIdx.x == 0 ? smem_running_prefix[state_idx] : make_float2(1.f, 0.f);
+            SSMScanPrefixCallbackOp<weight_t> prefix_op(running_prefix);
+            Ktraits::BlockScanT(smem_scan).InclusiveScan(
+                thread_data, thread_data, SSMScanOp<weight_t>(), prefix_op
+            );
+            // There's a syncthreads in the scan op, so we don't need to sync here.
+            // Unless there's only 1 warp, but then it's the same thread (0) reading and writing.
+            if (threadIdx.x == 0) {
+                smem_running_prefix[state_idx] = prefix_op.running_prefix;
+                x[chunk * params.dstate + state_idx] = prefix_op.running_prefix;
+            }
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                out_vals[i] += thread_data[i].y * C_vals[i];
+            }
+        }
+
+        input_t *out = reinterpret_cast<input_t *>(params.out_ptr) + batch_id * params.out_batch_stride
+            + dim_id * params.out_d_stride + chunk * kChunkSize;
+        __syncthreads();
+        store_output<Ktraits>(out, out_vals, smem_store, params.seqlen - chunk * kChunkSize);
+        Bvar += kChunkSize;
+        Cvar += kChunkSize;
+    }
+}
+
+template<int kNThreads, int kNItems, typename input_t, typename weight_t>
+void selective_scan_fwd_launch(SSMParamsBase &params, cudaStream_t stream) {
+    BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {
+        using Ktraits = Selective_Scan_fwd_kernel_traits<kNThreads, kNItems, kIsEvenLen, input_t, weight_t>;
+        constexpr int kSmemSize = Ktraits::kSmemSize + Ktraits::MaxDState * sizeof(typename Ktraits::scan_t);
+        // printf("smem_size = %d\n", kSmemSize);
+        dim3 grid(params.batch, params.dim);
+        auto kernel = &selective_scan_fwd_kernel<Ktraits>;
+        if (kSmemSize >= 48 * 1024) {
+            C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
+        }
+        kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+    });
+}
+
+template<int knrows, typename input_t, typename weight_t>
+void selective_scan_fwd_cuda(SSMParamsBase &params, cudaStream_t stream) {
+    if (params.seqlen <= 128) {
+        selective_scan_fwd_launch<32, 4, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 256) {
+        selective_scan_fwd_launch<32, 8, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 512) {
+        selective_scan_fwd_launch<32, 16, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 1024) {
+        selective_scan_fwd_launch<64, 16, input_t, weight_t>(params, stream);
+    } else {
+        selective_scan_fwd_launch<128, 16, input_t, weight_t>(params, stream);
+    }
+}

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusndstate/selective_scan_bwd_kernel_ndstate.cuh

@@ -0,0 +1,302 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <c10/util/BFloat16.h>
+#include <c10/util/Half.h>
+#include <c10/cuda/CUDAException.h>  // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
+#include <ATen/cuda/Atomic.cuh>  // For atomicAdd on complex
+
+#include <cub/block/block_load.cuh>
+#include <cub/block/block_store.cuh>
+#include <cub/block/block_scan.cuh>
+#include <cub/block/block_reduce.cuh>
+
+#include "selective_scan_ndstate.h"
+#include "selective_scan_common.h"
+#include "reverse_scan.cuh"
+#include "static_switch.h"
+
+template<int kNThreads_, int kNItems_, bool kIsEvenLen_, bool kDeltaSoftplus_, typename input_t_, typename weight_t_>
+struct Selective_Scan_bwd_kernel_traits {
+    static_assert(kNItems_ % 4 == 0);
+    using input_t = input_t_;
+    using weight_t = weight_t_;
+    static constexpr int kNThreads = kNThreads_;
+    static constexpr int kNItems = kNItems_;
+    static constexpr int kNBytes = sizeof(input_t);
+    static_assert(kNBytes == 2 || kNBytes == 4);
+    static constexpr int kNElts = kNBytes == 4 ? 4 : std::min(8, kNItems);
+    static_assert(kNItems % kNElts == 0);
+    static constexpr int kNLoads = kNItems / kNElts;
+    static constexpr bool kIsEvenLen = kIsEvenLen_;
+    static constexpr bool kDeltaSoftplus = kDeltaSoftplus_;
+    // Setting MinBlocksPerMP to be 3 (instead of 2) for 128 threads with float improves occupancy.
+    // For complex this would lead to massive register spilling, so we keep it at 2.
+    static constexpr int kMinBlocks = kNThreads == 128 && 3;
+    using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
+    using scan_t = float2;
+    using BlockLoadT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadWeightT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadWeightVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockStoreT = cub::BlockStore<input_t, kNThreads, kNItems, cub::BLOCK_STORE_WARP_TRANSPOSE>;
+    using BlockStoreVecT = cub::BlockStore<vec_t, kNThreads, kNLoads, cub::BLOCK_STORE_WARP_TRANSPOSE>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING_MEMOIZE>;
+    using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_WARP_SCANS>;
+    using BlockReverseScanT = BlockReverseScan<scan_t, kNThreads>;
+    using BlockReduceT = cub::BlockReduce<scan_t, kNThreads>;
+    using BlockReduceFloatT = cub::BlockReduce<float, kNThreads>;
+    using BlockExchangeT = cub::BlockExchange<float, kNThreads, kNItems>;
+    static constexpr int kSmemIOSize = std::max({sizeof(typename BlockLoadT::TempStorage),
+                                                 sizeof(typename BlockLoadVecT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightVecT::TempStorage),
+                                                 sizeof(typename BlockStoreT::TempStorage),
+                                                 sizeof(typename BlockStoreVecT::TempStorage)});
+    static constexpr int kSmemExchangeSize = 2 * sizeof(typename BlockExchangeT::TempStorage);
+    static constexpr int kSmemReduceSize = sizeof(typename BlockReduceT::TempStorage);
+    static constexpr int kSmemSize = kSmemIOSize + kSmemExchangeSize + kSmemReduceSize + sizeof(typename BlockScanT::TempStorage) + sizeof(typename BlockReverseScanT::TempStorage);
+};
+
+template<typename Ktraits>
+__global__ __launch_bounds__(Ktraits::kNThreads, Ktraits::kMinBlocks)
+void selective_scan_bwd_kernel(SSMParamsBwd params) {
+    constexpr bool kDeltaSoftplus = Ktraits::kDeltaSoftplus;
+    constexpr int kNThreads = Ktraits::kNThreads;
+    constexpr int kNItems = Ktraits::kNItems;
+    using input_t = typename Ktraits::input_t;
+    using weight_t = typename Ktraits::weight_t;
+    using scan_t = typename Ktraits::scan_t;
+
+    // Shared memory.
+    extern __shared__ char smem_[];
+    auto& smem_load = reinterpret_cast<typename Ktraits::BlockLoadT::TempStorage&>(smem_);
+    auto& smem_load_weight = reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage&>(smem_);
+    auto& smem_load_weight1 = *reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage*>(smem_ + sizeof(typename Ktraits::BlockLoadWeightT::TempStorage));
+    auto& smem_store = reinterpret_cast<typename Ktraits::BlockStoreT::TempStorage&>(smem_);
+    auto& smem_exchange = *reinterpret_cast<typename Ktraits::BlockExchangeT::TempStorage*>(smem_ + Ktraits::kSmemIOSize);
+    auto& smem_exchange1 = *reinterpret_cast<typename Ktraits::BlockExchangeT::TempStorage*>(smem_ + Ktraits::kSmemIOSize + sizeof(typename Ktraits::BlockExchangeT::TempStorage));
+    auto& smem_reduce = *reinterpret_cast<typename Ktraits::BlockReduceT::TempStorage*>(reinterpret_cast<char *>(&smem_exchange) + Ktraits::kSmemExchangeSize);
+    auto& smem_reduce_float = *reinterpret_cast<typename Ktraits::BlockReduceFloatT::TempStorage*>(&smem_reduce);
+    auto& smem_scan = *reinterpret_cast<typename Ktraits::BlockScanT::TempStorage*>(reinterpret_cast<char *>(&smem_reduce) + Ktraits::kSmemReduceSize);
+    auto& smem_reverse_scan = *reinterpret_cast<typename Ktraits::BlockReverseScanT::TempStorage*>(reinterpret_cast<char *>(&smem_scan) + sizeof(typename Ktraits::BlockScanT::TempStorage));
+    weight_t *smem_delta_a = reinterpret_cast<weight_t *>(smem_ + Ktraits::kSmemSize);
+    scan_t *smem_running_postfix = reinterpret_cast<scan_t *>(smem_delta_a + 2 + kNThreads);
+    weight_t *smem_da = reinterpret_cast<weight_t *>(smem_running_postfix + 1);
+
+    const int batch_id = blockIdx.x;
+    const int dim_id = blockIdx.y;
+    const int group_id = dim_id / (params.dim_ngroups_ratio);
+    input_t *u = reinterpret_cast<input_t *>(params.u_ptr) + batch_id * params.u_batch_stride
+        + dim_id * params.u_d_stride;
+    input_t *delta = reinterpret_cast<input_t *>(params.delta_ptr) + batch_id * params.delta_batch_stride
+        + dim_id * params.delta_d_stride;
+    input_t *dout = reinterpret_cast<input_t *>(params.dout_ptr) + batch_id * params.dout_batch_stride
+        + dim_id * params.dout_d_stride;
+    weight_t A_val = reinterpret_cast<weight_t *>(params.A_ptr)[dim_id];
+    constexpr float kLog2e = M_LOG2E;
+    weight_t A_scaled = A_val * kLog2e;
+    input_t *Bvar = reinterpret_cast<input_t *>(params.B_ptr) + batch_id * params.B_batch_stride + group_id * params.B_group_stride;
+    input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + batch_id * params.C_batch_stride + group_id * params.C_group_stride;
+    weight_t *dA = reinterpret_cast<weight_t *>(params.dA_ptr) + dim_id;
+    weight_t *dB = reinterpret_cast<weight_t *>(params.dB_ptr)
+        + (batch_id * params.dB_batch_stride + group_id * params.dB_group_stride);
+    weight_t *dC = reinterpret_cast<weight_t *>(params.dC_ptr)
+        + (batch_id * params.dC_batch_stride + group_id * params.dC_group_stride);
+    float *dD = params.dD_ptr == nullptr ? nullptr : reinterpret_cast<float *>(params.dD_ptr) + dim_id;
+    float D_val = params.D_ptr == nullptr ? 0 : reinterpret_cast<float *>(params.D_ptr)[dim_id];
+    float *ddelta_bias = params.ddelta_bias_ptr == nullptr ? nullptr : reinterpret_cast<float *>(params.ddelta_bias_ptr) + dim_id;
+    float delta_bias = params.delta_bias_ptr == nullptr ? 0 : reinterpret_cast<float *>(params.delta_bias_ptr)[dim_id];
+    scan_t *x = params.x_ptr == nullptr
+        ? nullptr
+        : reinterpret_cast<scan_t *>(params.x_ptr) + (batch_id * params.dim + dim_id) * (params.n_chunks);
+    float dD_val = 0;
+    float ddelta_bias_val = 0;
+
+    constexpr int kChunkSize = kNThreads * kNItems;
+    u += (params.n_chunks - 1) * kChunkSize;
+    delta += (params.n_chunks - 1) * kChunkSize;
+    dout += (params.n_chunks - 1) * kChunkSize;
+    Bvar += (params.n_chunks - 1) * kChunkSize;
+    Cvar += (params.n_chunks - 1) * kChunkSize;
+    for (int chunk = params.n_chunks - 1; chunk >= 0; --chunk) {
+        input_t u_vals[kNItems];
+        input_t delta_vals_load[kNItems];
+        input_t dout_vals_load[kNItems];
+        __syncthreads();
+        load_input<Ktraits>(u, u_vals, smem_load, params.seqlen - chunk * kChunkSize);
+        __syncthreads();
+        load_input<Ktraits>(delta, delta_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
+        __syncthreads();
+        load_input<Ktraits>(dout, dout_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
+        u -= kChunkSize;
+        // Will reload delta at the same location if kDeltaSoftplus
+        if constexpr (!kDeltaSoftplus) { delta -= kChunkSize; }
+        dout -= kChunkSize;
+
+        float dout_vals[kNItems], delta_vals[kNItems];
+        float du_vals[kNItems];
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) {
+            dout_vals[i] = float(dout_vals_load[i]);
+            delta_vals[i] = float(delta_vals_load[i]) + delta_bias;
+            if constexpr (kDeltaSoftplus) {
+                delta_vals[i] = delta_vals[i] <= 20.f ? log1pf(expf(delta_vals[i])) : delta_vals[i];
+            }
+        }
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) { du_vals[i] = D_val * dout_vals[i]; }
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) { dD_val += dout_vals[i] * float(u_vals[i]); }
+
+        float ddelta_vals[kNItems] = {0};
+        __syncthreads();
+        {
+            weight_t B_vals[kNItems], C_vals[kNItems];
+            load_weight<Ktraits>(Bvar, B_vals,
+                    smem_load_weight, (params.seqlen - chunk * kChunkSize));
+            auto &smem_load_weight_C = smem_load_weight1;
+            load_weight<Ktraits>(Cvar, C_vals,
+                    smem_load_weight_C, (params.seqlen - chunk * kChunkSize));
+            scan_t thread_data[kNItems], thread_reverse_data[kNItems];
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                const float delta_a_exp = exp2f(delta_vals[i] * A_scaled);
+                thread_data[i] = make_float2(delta_a_exp, delta_vals[i] * float(u_vals[i]) * B_vals[i]);
+                if (i == 0) {
+                    smem_delta_a[threadIdx.x == 0 ? (chunk % 2): threadIdx.x + 2] = delta_a_exp;
+                } else {
+                    thread_reverse_data[i - 1].x = delta_a_exp;
+                }
+                thread_reverse_data[i].y = dout_vals[i] * C_vals[i];
+            }
+            __syncthreads();
+            thread_reverse_data[kNItems - 1].x = threadIdx.x == kNThreads - 1
+                ? (chunk == params.n_chunks - 1 ? 1.f : smem_delta_a[(chunk + 1) % 2])
+                    : smem_delta_a[threadIdx.x + 1 + 2];
+            // Initialize running total
+            scan_t running_prefix = chunk > 0 && threadIdx.x % 32 == 0 ? x[chunk - 1] : make_float2(1.f, 0.f);
+            SSMScanPrefixCallbackOp<weight_t> prefix_op(running_prefix);
+            Ktraits::BlockScanT(smem_scan).InclusiveScan(
+                thread_data, thread_data, SSMScanOp<weight_t>(), prefix_op
+            );
+            scan_t running_postfix = chunk < params.n_chunks - 1 && threadIdx.x % 32 == 0 ? smem_running_postfix[0] : make_float2(1.f, 0.f);
+            SSMScanPrefixCallbackOp<weight_t> postfix_op(running_postfix);
+            Ktraits::BlockReverseScanT(smem_reverse_scan).InclusiveReverseScan(
+                thread_reverse_data, thread_reverse_data, SSMScanOp<weight_t>(), postfix_op
+            );
+            if (threadIdx.x == 0) { smem_running_postfix[0] = postfix_op.running_prefix; }
+            weight_t dA_val = 0;
+            weight_t dB_vals[kNItems], dC_vals[kNItems];
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                const float dx = thread_reverse_data[i].y;
+                const float ddelta_u = dx * B_vals[i];
+                du_vals[i] += ddelta_u * delta_vals[i];
+                const float a = thread_data[i].y - (delta_vals[i] * float(u_vals[i]) * B_vals[i]);
+                ddelta_vals[i] += ddelta_u * float(u_vals[i]) + dx * A_val * a;
+                dA_val += dx * delta_vals[i] * a;
+                dB_vals[i] = dx * delta_vals[i] * float(u_vals[i]);
+                dC_vals[i] = dout_vals[i] * thread_data[i].y;
+            }
+            // Block-exchange to make the atomicAdd's coalesced, otherwise they're much slower
+            Ktraits::BlockExchangeT(smem_exchange).BlockedToStriped(dB_vals, dB_vals);
+            auto &smem_exchange_C = smem_exchange1;
+            Ktraits::BlockExchangeT(smem_exchange_C).BlockedToStriped(dC_vals, dC_vals);
+            const int seqlen_remaining = params.seqlen - chunk * kChunkSize - threadIdx.x;
+            weight_t *dB_cur = dB + chunk * kChunkSize + threadIdx.x;
+            weight_t *dC_cur = dC + chunk * kChunkSize + threadIdx.x;
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                if (i * kNThreads < seqlen_remaining) {
+                    { gpuAtomicAdd(dB_cur + i * kNThreads, dB_vals[i]); }
+                    { gpuAtomicAdd(dC_cur + i * kNThreads, dC_vals[i]); }
+                }
+            }
+            dA_val = Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(dA_val);
+            if (threadIdx.x == 0) {
+                smem_da[0] = chunk == params.n_chunks - 1 ? dA_val : dA_val + smem_da[0];
+            }
+        }
+
+        if constexpr (kDeltaSoftplus) {
+            input_t delta_vals_load[kNItems];
+            __syncthreads();
+            load_input<Ktraits>(delta, delta_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
+            delta -= kChunkSize;
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                float delta_val = float(delta_vals_load[i]) + delta_bias;
+                float delta_val_neg_exp = expf(-delta_val);
+                ddelta_vals[i] = delta_val <= 20.f
+                    ? ddelta_vals[i] / (1.f + delta_val_neg_exp)
+                        : ddelta_vals[i];
+            }
+        }
+        
+        __syncthreads();
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) { ddelta_bias_val += ddelta_vals[i]; }
+
+        input_t *du = reinterpret_cast<input_t *>(params.du_ptr) + batch_id * params.du_batch_stride
+            + dim_id * params.du_d_stride + chunk * kChunkSize;
+        input_t *ddelta = reinterpret_cast<input_t *>(params.ddelta_ptr) + batch_id * params.ddelta_batch_stride
+            + dim_id * params.ddelta_d_stride + chunk * kChunkSize;
+        __syncthreads();
+        store_output<Ktraits>(du, du_vals, smem_store, params.seqlen - chunk * kChunkSize);
+        __syncthreads();
+        store_output<Ktraits>(ddelta, ddelta_vals, smem_store, params.seqlen - chunk * kChunkSize);
+        Bvar -= kChunkSize;
+        Cvar -= kChunkSize;
+    }
+
+    if (params.dD_ptr != nullptr) {
+        __syncthreads();
+        dD_val = Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(dD_val);
+        if (threadIdx.x == 0) { gpuAtomicAdd(dD, dD_val); }
+    }
+    if (params.ddelta_bias_ptr != nullptr) {
+        __syncthreads();
+        ddelta_bias_val = Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(ddelta_bias_val);
+        if (threadIdx.x == 0) { gpuAtomicAdd(ddelta_bias, ddelta_bias_val); }
+    }
+    __syncthreads();
+    if (threadIdx.x == 0) { gpuAtomicAdd(dA, smem_da[0]); }
+}
+
+template<int kNThreads, int kNItems, typename input_t, typename weight_t>
+void selective_scan_bwd_launch(SSMParamsBwd &params, cudaStream_t stream) {
+    BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {
+        BOOL_SWITCH(params.delta_softplus, kDeltaSoftplus, [&] {
+            using Ktraits = Selective_Scan_bwd_kernel_traits<kNThreads, kNItems, kIsEvenLen, kDeltaSoftplus, input_t, weight_t>;
+            constexpr int kSmemSize = Ktraits::kSmemSize + sizeof(typename Ktraits::scan_t) + (kNThreads + 4) * sizeof(typename Ktraits::weight_t);
+            // printf("smem_size = %d\n", kSmemSize);
+            dim3 grid(params.batch, params.dim);
+            auto kernel = &selective_scan_bwd_kernel<Ktraits>;
+            if (kSmemSize >= 48 * 1024) {
+                C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
+            }
+            kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
+            C10_CUDA_KERNEL_LAUNCH_CHECK();
+        });
+    });
+}
+
+template<int knrows, typename input_t, typename weight_t>
+void selective_scan_bwd_cuda(SSMParamsBwd &params, cudaStream_t stream) {
+    if (params.seqlen <= 128) {
+        selective_scan_bwd_launch<32, 4, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 256) {
+        selective_scan_bwd_launch<32, 8, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 512) {
+        selective_scan_bwd_launch<32, 16, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 1024) {
+        selective_scan_bwd_launch<64, 16, input_t, weight_t>(params, stream);
+    } else {
+        selective_scan_bwd_launch<128, 16, input_t, weight_t>(params, stream);
+    }
+}

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusndstate/selective_scan_core_bwd.cu

@@ -0,0 +1,9 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+#include "selective_scan_bwd_kernel_ndstate.cuh"
+
+template void selective_scan_bwd_cuda<1, float, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<1, at::Half, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<1, at::BFloat16, float>(SSMParamsBwd &params, cudaStream_t stream);
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusndstate/selective_scan_core_fwd.cu

@@ -0,0 +1,9 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+#include "selective_scan_fwd_kernel_ndstate.cuh"
+
+template void selective_scan_fwd_cuda<1, float, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<1, at::Half, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<1, at::BFloat16, float>(SSMParamsBase &params, cudaStream_t stream);
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusndstate/selective_scan_fwd_kernel_ndstate.cuh

@@ -0,0 +1,200 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <c10/util/BFloat16.h>
+#include <c10/util/Half.h>
+#include <c10/cuda/CUDAException.h>  // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
+
+#include <cub/block/block_load.cuh>
+#include <cub/block/block_store.cuh>
+#include <cub/block/block_scan.cuh>
+
+#include "selective_scan_ndstate.h"
+#include "selective_scan_common.h"
+#include "static_switch.h"
+
+template<int kNThreads_, int kNItems_, bool kIsEvenLen_, typename input_t_, typename weight_t_>
+struct Selective_Scan_fwd_kernel_traits {
+    static_assert(kNItems_ % 4 == 0);
+    using input_t = input_t_;
+    using weight_t = weight_t_;
+    static constexpr int kNThreads = kNThreads_;
+    // Setting MinBlocksPerMP to be 3 (instead of 2) for 128 threads improves occupancy.
+    static constexpr int kMinBlocks = kNThreads < 128 ? 5 : 3;
+    static constexpr int kNItems = kNItems_;
+    static constexpr int kNBytes = sizeof(input_t);
+    static_assert(kNBytes == 2 || kNBytes == 4);
+    static constexpr int kNElts = kNBytes == 4 ? 4 : std::min(8, kNItems);
+    static_assert(kNItems % kNElts == 0);
+    static constexpr int kNLoads = kNItems / kNElts;
+    static constexpr bool kIsEvenLen = kIsEvenLen_;
+
+    static constexpr bool kDirectIO = kIsEvenLen && kNLoads == 1;
+
+    using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
+    using scan_t = float2;
+    using BlockLoadT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads,
+        !kDirectIO ? cub::BLOCK_LOAD_WARP_TRANSPOSE : cub::BLOCK_LOAD_DIRECT>;
+    using BlockLoadWeightT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadWeightVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads,
+        !kDirectIO ? cub::BLOCK_LOAD_WARP_TRANSPOSE  : cub::BLOCK_LOAD_DIRECT>;
+    using BlockStoreT = cub::BlockStore<input_t, kNThreads, kNItems, cub::BLOCK_STORE_WARP_TRANSPOSE>;
+    using BlockStoreVecT = cub::BlockStore<vec_t, kNThreads, kNLoads,
+        !kDirectIO ? cub::BLOCK_STORE_WARP_TRANSPOSE : cub::BLOCK_STORE_DIRECT>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING_MEMOIZE>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING>;
+    using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_WARP_SCANS>;
+    static constexpr int kSmemIOSize = std::max({sizeof(typename BlockLoadT::TempStorage),
+                                                 sizeof(typename BlockLoadVecT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightVecT::TempStorage),
+                                                 sizeof(typename BlockStoreT::TempStorage),
+                                                 sizeof(typename BlockStoreVecT::TempStorage)});
+    static constexpr int kSmemSize = kSmemIOSize + sizeof(typename BlockScanT::TempStorage);
+};
+
+template<typename Ktraits>
+__global__ __launch_bounds__(Ktraits::kNThreads, Ktraits::kMinBlocks)
+void selective_scan_fwd_kernel(SSMParamsBase params) {
+    constexpr int kNThreads = Ktraits::kNThreads;
+    constexpr int kNItems = Ktraits::kNItems;
+    constexpr bool kDirectIO = Ktraits::kDirectIO;
+    using input_t = typename Ktraits::input_t;
+    using weight_t = typename Ktraits::weight_t;
+    using scan_t = typename Ktraits::scan_t;
+
+    // Shared memory.
+    extern __shared__ char smem_[];
+    auto& smem_load = reinterpret_cast<typename Ktraits::BlockLoadT::TempStorage&>(smem_);
+    auto& smem_load_weight = reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage&>(smem_);
+    auto& smem_load_weight1 = *reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage*>(smem_ + sizeof(typename Ktraits::BlockLoadWeightT::TempStorage));
+    auto& smem_store = reinterpret_cast<typename Ktraits::BlockStoreT::TempStorage&>(smem_);
+    auto& smem_scan = *reinterpret_cast<typename Ktraits::BlockScanT::TempStorage*>(smem_ + Ktraits::kSmemIOSize);
+    scan_t *smem_running_prefix = reinterpret_cast<scan_t *>(smem_ + Ktraits::kSmemSize);
+
+    const int batch_id = blockIdx.x;
+    const int dim_id = blockIdx.y;
+    const int group_id = dim_id / (params.dim_ngroups_ratio);
+    input_t *u = reinterpret_cast<input_t *>(params.u_ptr) + batch_id * params.u_batch_stride
+        + dim_id * params.u_d_stride;
+    input_t *delta = reinterpret_cast<input_t *>(params.delta_ptr) + batch_id * params.delta_batch_stride
+        + dim_id * params.delta_d_stride;
+    constexpr float kLog2e = M_LOG2E;
+    weight_t A_val = reinterpret_cast<weight_t *>(params.A_ptr)[dim_id] * kLog2e;
+    input_t *Bvar = reinterpret_cast<input_t *>(params.B_ptr) + batch_id * params.B_batch_stride + group_id * params.B_group_stride;
+    input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + batch_id * params.C_batch_stride + group_id * params.C_group_stride;
+    scan_t *x = reinterpret_cast<scan_t *>(params.x_ptr) + (batch_id * params.dim + dim_id) * params.n_chunks;
+
+    float D_val = 0; // attention!
+    if (params.D_ptr != nullptr) {
+        D_val = reinterpret_cast<float *>(params.D_ptr)[dim_id];
+    }
+    float delta_bias = 0;
+    if (params.delta_bias_ptr != nullptr) {
+        delta_bias = reinterpret_cast<float *>(params.delta_bias_ptr)[dim_id];
+    }
+
+    constexpr int kChunkSize = kNThreads * kNItems;
+    for (int chunk = 0; chunk < params.n_chunks; ++chunk) {
+        input_t u_vals[kNItems], delta_vals_load[kNItems];
+        __syncthreads();
+        load_input<Ktraits>(u, u_vals, smem_load, params.seqlen - chunk * kChunkSize);
+        if constexpr (!kDirectIO) { __syncthreads(); }
+        load_input<Ktraits>(delta, delta_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
+        u += kChunkSize;
+        delta += kChunkSize;
+
+        float delta_vals[kNItems], delta_u_vals[kNItems], out_vals[kNItems];
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) {
+            float u_val = float(u_vals[i]);
+            delta_vals[i] = float(delta_vals_load[i]) + delta_bias;
+            if (params.delta_softplus) {
+                delta_vals[i] = delta_vals[i] <= 20.f ? log1pf(expf(delta_vals[i])) : delta_vals[i];
+            }
+            delta_u_vals[i] = delta_vals[i] * u_val;
+            out_vals[i] = D_val * u_val;
+        }
+
+        __syncthreads();
+        {
+            weight_t B_vals[kNItems], C_vals[kNItems];
+            load_weight<Ktraits>(Bvar, B_vals,
+                    smem_load_weight, (params.seqlen - chunk * kChunkSize));
+            auto &smem_load_weight_C = smem_load_weight1;
+            load_weight<Ktraits>(Cvar, C_vals,
+                    smem_load_weight_C, (params.seqlen - chunk * kChunkSize));
+            __syncthreads();
+            scan_t thread_data[kNItems];
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                thread_data[i] = make_float2(exp2f(delta_vals[i] * A_val), B_vals[i] * delta_u_vals[i]);
+                if constexpr (!Ktraits::kIsEvenLen) {  // So that the last state is correct
+                    if (threadIdx.x * kNItems + i >= params.seqlen - chunk * kChunkSize) {
+                        thread_data[i] = make_float2(1.f, 0.f);
+                    }
+                }
+            }
+            // Initialize running total
+            scan_t running_prefix;
+            // If we use WARP_SCAN then all lane 0 of all warps (not just thread 0) needs to read
+            running_prefix = chunk > 0 && threadIdx.x % 32 == 0 ? smem_running_prefix[0] : make_float2(1.f, 0.f);
+            SSMScanPrefixCallbackOp<weight_t> prefix_op(running_prefix);
+            Ktraits::BlockScanT(smem_scan).InclusiveScan(
+                thread_data, thread_data, SSMScanOp<weight_t>(), prefix_op
+            );
+            // There's a syncthreads in the scan op, so we don't need to sync here.
+            // Unless there's only 1 warp, but then it's the same thread (0) reading and writing.
+            if (threadIdx.x == 0) {
+                smem_running_prefix[0] = prefix_op.running_prefix;
+                x[chunk] = prefix_op.running_prefix;
+            }
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                out_vals[i] += thread_data[i].y * C_vals[i];
+            }
+        }
+
+        input_t *out = reinterpret_cast<input_t *>(params.out_ptr) + batch_id * params.out_batch_stride
+            + dim_id * params.out_d_stride + chunk * kChunkSize;
+        __syncthreads();
+        store_output<Ktraits>(out, out_vals, smem_store, params.seqlen - chunk * kChunkSize);
+        Bvar += kChunkSize;
+        Cvar += kChunkSize;
+    }
+}
+
+template<int kNThreads, int kNItems, typename input_t, typename weight_t>
+void selective_scan_fwd_launch(SSMParamsBase &params, cudaStream_t stream) {
+    BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {
+        using Ktraits = Selective_Scan_fwd_kernel_traits<kNThreads, kNItems, kIsEvenLen, input_t, weight_t>;
+        constexpr int kSmemSize = Ktraits::kSmemSize + sizeof(typename Ktraits::scan_t);
+        // printf("smem_size = %d\n", kSmemSize);
+        dim3 grid(params.batch, params.dim);
+        auto kernel = &selective_scan_fwd_kernel<Ktraits>;
+        if (kSmemSize >= 48 * 1024) {
+            C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
+        }
+        kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+    });
+}
+
+template<int knrows, typename input_t, typename weight_t>
+void selective_scan_fwd_cuda(SSMParamsBase &params, cudaStream_t stream) {
+    if (params.seqlen <= 128) {
+        selective_scan_fwd_launch<32, 4, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 256) {
+        selective_scan_fwd_launch<32, 8, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 512) {
+        selective_scan_fwd_launch<32, 16, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 1024) {
+        selective_scan_fwd_launch<64, 16, input_t, weight_t>(params, stream);
+    } else {
+        selective_scan_fwd_launch<128, 16, input_t, weight_t>(params, stream);
+    }
+}

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusndstate/selective_scan_ndstate.cpp

@@ -0,0 +1,341 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/extension.h>
+#include <vector>
+
+#include "selective_scan_ndstate.h"
+#define MAX_DSTATE 256
+
+#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
+using weight_t = float;
+
+#define DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, NAME, ...)                    \
+    if (ITYPE == at::ScalarType::Half) {                                            \
+        using input_t = at::Half;                                                   \
+        __VA_ARGS__();                                                              \
+    } else if (ITYPE == at::ScalarType::BFloat16) {                                 \
+        using input_t = at::BFloat16;                                               \
+        __VA_ARGS__();                                                              \
+    } else if (ITYPE == at::ScalarType::Float)  {                                   \
+        using input_t = float;                                                      \
+        __VA_ARGS__();                                                              \
+    } else {                                                                        \
+        AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \
+    }
+
+template<int knrows, typename input_t, typename weight_t>
+void selective_scan_fwd_cuda(SSMParamsBase &params, cudaStream_t stream);
+
+template <int knrows, typename input_t, typename weight_t>
+void selective_scan_bwd_cuda(SSMParamsBwd &params, cudaStream_t stream);
+
+void set_ssm_params_fwd(SSMParamsBase &params,
+                        // sizes
+                        const size_t batch,
+                        const size_t dim,
+                        const size_t seqlen,
+                        const size_t n_groups,
+                        const size_t n_chunks,
+                        // device pointers
+                        const at::Tensor u,
+                        const at::Tensor delta,
+                        const at::Tensor A,
+                        const at::Tensor B,
+                        const at::Tensor C,
+                        const at::Tensor out,
+                        void* D_ptr,
+                        void* delta_bias_ptr,
+                        void* x_ptr,
+                        bool delta_softplus) {
+
+    // Reset the parameters
+    memset(&params, 0, sizeof(params));
+
+    params.batch = batch;
+    params.dim = dim;
+    params.seqlen = seqlen;
+    params.n_groups = n_groups;
+    params.n_chunks = n_chunks;
+    params.dim_ngroups_ratio = dim / n_groups;
+
+    params.delta_softplus = delta_softplus;
+
+    // Set the pointers and strides.
+    params.u_ptr = u.data_ptr();
+    params.delta_ptr = delta.data_ptr();
+    params.A_ptr = A.data_ptr();
+    params.B_ptr = B.data_ptr();
+    params.C_ptr = C.data_ptr();
+    params.D_ptr = D_ptr;
+    params.delta_bias_ptr = delta_bias_ptr;
+    params.out_ptr = out.data_ptr();
+    params.x_ptr = x_ptr;
+
+    // All stride are in elements, not bytes.
+    params.A_d_stride = A.stride(0);
+    params.B_batch_stride = B.stride(0);
+    params.B_group_stride = B.stride(1);
+    params.C_batch_stride = C.stride(0);
+    params.C_group_stride = C.stride(1);
+    params.u_batch_stride = u.stride(0);
+    params.u_d_stride = u.stride(1);
+    params.delta_batch_stride = delta.stride(0);
+    params.delta_d_stride = delta.stride(1);
+
+    params.out_batch_stride = out.stride(0);
+    params.out_d_stride = out.stride(1);
+}
+
+void set_ssm_params_bwd(SSMParamsBwd &params,
+                        // sizes
+                        const size_t batch,
+                        const size_t dim,
+                        const size_t seqlen,
+                        const size_t n_groups,
+                        const size_t n_chunks,
+                        // device pointers
+                        const at::Tensor u,
+                        const at::Tensor delta,
+                        const at::Tensor A,
+                        const at::Tensor B,
+                        const at::Tensor C,
+                        const at::Tensor out,
+                        void* D_ptr,
+                        void* delta_bias_ptr,
+                        void* x_ptr,
+                        const at::Tensor dout,
+                        const at::Tensor du,
+                        const at::Tensor ddelta,
+                        const at::Tensor dA,
+                        const at::Tensor dB,
+                        const at::Tensor dC,
+                        void* dD_ptr,
+                        void* ddelta_bias_ptr,
+                        bool delta_softplus) {
+    // Pass in "dout" instead of "out", we're not gonna use "out" unless we have z
+    set_ssm_params_fwd(params, batch, dim, seqlen, n_groups, n_chunks,
+                       u, delta, A, B, C, dout,
+                       D_ptr, delta_bias_ptr, x_ptr, delta_softplus);
+
+    // Set the pointers and strides.
+    params.dout_ptr = dout.data_ptr();
+    params.du_ptr = du.data_ptr();
+    params.dA_ptr = dA.data_ptr();
+    params.dB_ptr = dB.data_ptr();
+    params.dC_ptr = dC.data_ptr();
+    params.dD_ptr = dD_ptr;
+    params.ddelta_ptr = ddelta.data_ptr();
+    params.ddelta_bias_ptr = ddelta_bias_ptr;
+    // All stride are in elements, not bytes.
+    params.dout_batch_stride = dout.stride(0);
+    params.dout_d_stride = dout.stride(1);
+    params.dA_d_stride = dA.stride(0);
+    params.dB_batch_stride = dB.stride(0);
+    params.dB_group_stride = dB.stride(1);
+    params.dC_batch_stride = dC.stride(0);
+    params.dC_group_stride = dC.stride(1);
+    params.du_batch_stride = du.stride(0);
+    params.du_d_stride = du.stride(1);
+    params.ddelta_batch_stride = ddelta.stride(0);
+    params.ddelta_d_stride = ddelta.stride(1);
+
+}
+
+std::vector<at::Tensor>
+selective_scan_fwd(const at::Tensor &u, const at::Tensor &delta,
+                  const at::Tensor &A, const at::Tensor &B, const at::Tensor &C,
+                  const c10::optional<at::Tensor> &D_,
+                  const c10::optional<at::Tensor> &delta_bias_,
+                  bool delta_softplus,
+                  int nrows
+                  ) {
+    auto input_type = u.scalar_type();
+    auto weight_type = A.scalar_type();
+    TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
+    TORCH_CHECK(weight_type == at::ScalarType::Float);
+
+    TORCH_CHECK(delta.scalar_type() == input_type);
+    TORCH_CHECK(B.scalar_type() == input_type);
+    TORCH_CHECK(C.scalar_type() == input_type);
+
+    TORCH_CHECK(u.is_cuda());
+    TORCH_CHECK(delta.is_cuda());
+    TORCH_CHECK(A.is_cuda());
+    TORCH_CHECK(B.is_cuda());
+    TORCH_CHECK(C.is_cuda());
+
+    TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
+    TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
+
+    const auto sizes = u.sizes();
+    const int batch_size = sizes[0];
+    const int dim = sizes[1];
+    const int seqlen = sizes[2];
+    const int n_groups = B.size(1);
+
+    TORCH_CHECK(dim % n_groups == 0, "dims should be dividable by n_groups");
+
+    CHECK_SHAPE(u, batch_size, dim, seqlen);
+    CHECK_SHAPE(delta, batch_size, dim, seqlen);
+    CHECK_SHAPE(A, dim);
+    CHECK_SHAPE(B, batch_size, n_groups, seqlen);
+    TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
+    CHECK_SHAPE(C, batch_size, n_groups, seqlen);
+    TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
+
+    if (D_.has_value()) {
+        auto D = D_.value();
+        TORCH_CHECK(D.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(D.is_cuda());
+        TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
+        CHECK_SHAPE(D, dim);
+    }
+
+    if (delta_bias_.has_value()) {
+        auto delta_bias = delta_bias_.value();
+        TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(delta_bias.is_cuda());
+        TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
+        CHECK_SHAPE(delta_bias, dim);
+    }
+
+    const int n_chunks = (seqlen + 2048 - 1) / 2048; // max is 128 * 16 = 2048 in fwd_kernel
+    at::Tensor out = torch::empty_like(delta);
+    at::Tensor x;
+    x = torch::empty({batch_size, dim, n_chunks, 1 * 2}, u.options().dtype(weight_type));
+
+    SSMParamsBase params;
+    set_ssm_params_fwd(params, batch_size, dim, seqlen, n_groups, n_chunks,
+                       u, delta, A, B, C, out,
+                       D_.has_value() ? D_.value().data_ptr() : nullptr,
+                       delta_bias_.has_value() ? delta_bias_.value().data_ptr() : nullptr,
+                       x.data_ptr(),
+                       delta_softplus);
+
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)u.get_device()};
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), "selective_scan_fwd", [&] {
+        selective_scan_fwd_cuda<1, input_t, weight_t>(params, stream);
+    });
+    std::vector<at::Tensor> result = {out, x};
+    return result;
+}
+
+std::vector<at::Tensor>
+selective_scan_bwd(const at::Tensor &u, const at::Tensor &delta,
+                  const at::Tensor &A, const at::Tensor &B, const at::Tensor &C,
+                  const c10::optional<at::Tensor> &D_,
+                  const c10::optional<at::Tensor> &delta_bias_,
+                  const at::Tensor &dout,
+                  const c10::optional<at::Tensor> &x_,
+                  bool delta_softplus,
+                  int nrows
+                  ) {
+    auto input_type = u.scalar_type();
+    auto weight_type = A.scalar_type();
+    TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
+    TORCH_CHECK(weight_type == at::ScalarType::Float);
+
+    TORCH_CHECK(delta.scalar_type() == input_type);
+    TORCH_CHECK(B.scalar_type() == input_type);
+    TORCH_CHECK(C.scalar_type() == input_type);
+    TORCH_CHECK(dout.scalar_type() == input_type);
+
+    TORCH_CHECK(u.is_cuda());
+    TORCH_CHECK(delta.is_cuda());
+    TORCH_CHECK(A.is_cuda());
+    TORCH_CHECK(B.is_cuda());
+    TORCH_CHECK(C.is_cuda());
+    TORCH_CHECK(dout.is_cuda());
+
+    TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
+    TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
+    TORCH_CHECK(dout.stride(-1) == 1 || dout.size(-1) == 1);
+
+    const auto sizes = u.sizes();
+    const int batch_size = sizes[0];
+    const int dim = sizes[1];
+    const int seqlen = sizes[2];
+    const int n_groups = B.size(1);
+
+    TORCH_CHECK(dim % n_groups == 0, "dims should be dividable by n_groups");
+    
+    CHECK_SHAPE(u, batch_size, dim, seqlen);
+    CHECK_SHAPE(delta, batch_size, dim, seqlen);
+    CHECK_SHAPE(A, dim);
+    CHECK_SHAPE(B, batch_size, n_groups, seqlen);
+    TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
+    CHECK_SHAPE(C, batch_size, n_groups, seqlen);
+    TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
+    CHECK_SHAPE(dout, batch_size, dim, seqlen);
+
+    if (D_.has_value()) {
+        auto D = D_.value();
+        TORCH_CHECK(D.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(D.is_cuda());
+        TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
+        CHECK_SHAPE(D, dim);
+    }
+
+    if (delta_bias_.has_value()) {
+        auto delta_bias = delta_bias_.value();
+        TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(delta_bias.is_cuda());
+        TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
+        CHECK_SHAPE(delta_bias, dim);
+    }
+
+    at::Tensor out;
+    const int n_chunks = (seqlen + 2048 - 1) / 2048;
+    // const int n_chunks = (seqlen + 1024 - 1) / 1024;
+    if (n_chunks > 1) { TORCH_CHECK(x_.has_value()); }
+    if (x_.has_value()) {
+        auto x = x_.value();
+        TORCH_CHECK(x.scalar_type() == weight_type);
+        TORCH_CHECK(x.is_cuda());
+        TORCH_CHECK(x.is_contiguous());
+        CHECK_SHAPE(x, batch_size, dim, n_chunks, 2 * 1);
+    }
+
+    at::Tensor du = torch::empty_like(u);
+    at::Tensor ddelta = torch::empty_like(delta);
+    at::Tensor dA = torch::zeros_like(A);
+    at::Tensor dB = torch::zeros_like(B, B.options().dtype(torch::kFloat32));
+    at::Tensor dC = torch::zeros_like(C, C.options().dtype(torch::kFloat32));
+    at::Tensor dD;
+    if (D_.has_value()) { dD = torch::zeros_like(D_.value()); }
+    at::Tensor ddelta_bias;
+    if (delta_bias_.has_value()) { ddelta_bias = torch::zeros_like(delta_bias_.value()); }
+
+    SSMParamsBwd params;
+    set_ssm_params_bwd(params, batch_size, dim, seqlen, n_groups, n_chunks,
+                       u, delta, A, B, C, out,
+                       D_.has_value() ? D_.value().data_ptr() : nullptr,
+                       delta_bias_.has_value() ? delta_bias_.value().data_ptr() : nullptr,
+                       x_.has_value() ? x_.value().data_ptr() : nullptr,
+                       dout, du, ddelta, dA, dB, dC,
+                       D_.has_value() ? dD.data_ptr() : nullptr,
+                       delta_bias_.has_value() ? ddelta_bias.data_ptr() : nullptr,
+                       delta_softplus);
+
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)u.get_device()};
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), "selective_scan_bwd", [&] {
+        selective_scan_bwd_cuda<1, input_t, weight_t>(params, stream);
+    });
+    std::vector<at::Tensor> result = {du, ddelta, dA, dB.to(B.dtype()), dC.to(C.dtype()), dD, ddelta_bias};
+    return result;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("fwd", &selective_scan_fwd, "Selective scan forward");
+    m.def("bwd", &selective_scan_bwd, "Selective scan backward");
+}

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusndstate/selective_scan_ndstate.h

@@ -0,0 +1,84 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct SSMScanParamsBase {
+    using index_t = uint32_t;
+
+    int batch, seqlen, n_chunks;
+    index_t a_batch_stride;
+    index_t b_batch_stride;
+    index_t out_batch_stride;
+
+    // Common data pointers.
+    void *__restrict__ a_ptr;
+    void *__restrict__ b_ptr;
+    void *__restrict__ out_ptr;
+    void *__restrict__ x_ptr;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct SSMParamsBase {
+    using index_t = uint32_t;
+
+    int batch, dim, seqlen, n_groups, n_chunks;
+    int dim_ngroups_ratio;
+
+    bool delta_softplus;
+
+    index_t A_d_stride;
+    index_t B_batch_stride;
+    index_t B_d_stride;
+    index_t B_group_stride;
+    index_t C_batch_stride;
+    index_t C_d_stride;
+    index_t C_group_stride;
+    index_t u_batch_stride;
+    index_t u_d_stride;
+    index_t delta_batch_stride;
+    index_t delta_d_stride;
+    index_t out_batch_stride;
+    index_t out_d_stride;
+
+    // Common data pointers.
+    void *__restrict__ A_ptr;
+    void *__restrict__ B_ptr;
+    void *__restrict__ C_ptr;
+    void *__restrict__ D_ptr;
+    void *__restrict__ u_ptr;
+    void *__restrict__ delta_ptr;
+    void *__restrict__ delta_bias_ptr;
+    void *__restrict__ out_ptr;
+    void *__restrict__ x_ptr;
+};
+
+struct SSMParamsBwd: public SSMParamsBase {
+    index_t dout_batch_stride;
+    index_t dout_d_stride;
+    index_t dA_d_stride;
+    index_t dB_batch_stride;
+    index_t dB_group_stride;
+    index_t dB_d_stride;
+    index_t dC_batch_stride;
+    index_t dC_group_stride;
+    index_t dC_d_stride;
+    index_t du_batch_stride;
+    index_t du_d_stride;
+    index_t ddelta_batch_stride;
+    index_t ddelta_d_stride;
+
+    // Common data pointers.
+    void *__restrict__ dout_ptr;
+    void *__restrict__ dA_ptr;
+    void *__restrict__ dB_ptr;
+    void *__restrict__ dC_ptr;
+    void *__restrict__ dD_ptr;
+    void *__restrict__ du_ptr;
+    void *__restrict__ ddelta_ptr;
+    void *__restrict__ ddelta_bias_ptr;
+};

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusnrow/selective_scan_bwd_kernel_nrow.cuh

@@ -0,0 +1,344 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <c10/util/BFloat16.h>
+#include <c10/util/Half.h>
+#include <c10/cuda/CUDAException.h>  // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
+#include <ATen/cuda/Atomic.cuh>  // For atomicAdd on complex
+
+#include <cub/block/block_load.cuh>
+#include <cub/block/block_store.cuh>
+#include <cub/block/block_scan.cuh>
+#include <cub/block/block_reduce.cuh>
+
+#include "selective_scan.h"
+#include "selective_scan_common.h"
+#include "reverse_scan.cuh"
+#include "static_switch.h"
+
+template<int kNThreads_, int kNItems_, int kNRows_, bool kIsEvenLen_, bool kDeltaSoftplus_, typename input_t_, typename weight_t_>
+struct Selective_Scan_bwd_kernel_traits {
+    static_assert(kNItems_ % 4 == 0);
+    using input_t = input_t_;
+    using weight_t = weight_t_;
+    static constexpr int kNThreads = kNThreads_;
+    static constexpr int kNItems = kNItems_;
+    static constexpr int kNRows = kNRows_;
+    static constexpr int MaxDState = MAX_DSTATE / kNRows_;
+    static constexpr int kNBytes = sizeof(input_t);
+    static_assert(kNBytes == 2 || kNBytes == 4);
+    static constexpr int kNElts = kNBytes == 4 ? 4 : std::min(8, kNItems);
+    static_assert(kNItems % kNElts == 0);
+    static constexpr int kNLoads = kNItems / kNElts;
+    static constexpr bool kIsEvenLen = kIsEvenLen_;
+    static constexpr bool kDeltaSoftplus = kDeltaSoftplus_;
+    // Setting MinBlocksPerMP to be 3 (instead of 2) for 128 threads with float improves occupancy.
+    // For complex this would lead to massive register spilling, so we keep it at 2.
+    static constexpr int kMinBlocks = kNThreads == 128 && 3;
+    using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
+    using scan_t = float2;
+    using BlockLoadT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadWeightT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadWeightVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockStoreT = cub::BlockStore<input_t, kNThreads, kNItems, cub::BLOCK_STORE_WARP_TRANSPOSE>;
+    using BlockStoreVecT = cub::BlockStore<vec_t, kNThreads, kNLoads, cub::BLOCK_STORE_WARP_TRANSPOSE>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING_MEMOIZE>;
+    using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_WARP_SCANS>;
+    using BlockReverseScanT = BlockReverseScan<scan_t, kNThreads>;
+    using BlockReduceT = cub::BlockReduce<scan_t, kNThreads>;
+    using BlockReduceFloatT = cub::BlockReduce<float, kNThreads>;
+    using BlockExchangeT = cub::BlockExchange<float, kNThreads, kNItems>;
+    static constexpr int kSmemIOSize = std::max({sizeof(typename BlockLoadT::TempStorage),
+                                                 sizeof(typename BlockLoadVecT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightVecT::TempStorage),
+                                                 sizeof(typename BlockStoreT::TempStorage),
+                                                 sizeof(typename BlockStoreVecT::TempStorage)});
+    static constexpr int kSmemExchangeSize = 2 * sizeof(typename BlockExchangeT::TempStorage);
+    static constexpr int kSmemReduceSize = sizeof(typename BlockReduceT::TempStorage);
+    static constexpr int kSmemSize = kSmemIOSize + kSmemExchangeSize + kSmemReduceSize + sizeof(typename BlockScanT::TempStorage) + sizeof(typename BlockReverseScanT::TempStorage);
+};
+
+template<typename Ktraits>
+__global__ __launch_bounds__(Ktraits::kNThreads, Ktraits::kMinBlocks)
+void selective_scan_bwd_kernel(SSMParamsBwd params) {
+    constexpr bool kDeltaSoftplus = Ktraits::kDeltaSoftplus;
+    constexpr int kNThreads = Ktraits::kNThreads;
+    constexpr int kNItems = Ktraits::kNItems;
+    constexpr int kNRows = Ktraits::kNRows;
+    using input_t = typename Ktraits::input_t;
+    using weight_t = typename Ktraits::weight_t;
+    using scan_t = typename Ktraits::scan_t;
+
+    // Shared memory.
+    extern __shared__ char smem_[];
+    auto& smem_load = reinterpret_cast<typename Ktraits::BlockLoadT::TempStorage&>(smem_);
+    auto& smem_load_weight = reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage&>(smem_);
+    auto& smem_load_weight1 = *reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage*>(smem_ + sizeof(typename Ktraits::BlockLoadWeightT::TempStorage));
+    auto& smem_store = reinterpret_cast<typename Ktraits::BlockStoreT::TempStorage&>(smem_);
+    auto& smem_exchange = *reinterpret_cast<typename Ktraits::BlockExchangeT::TempStorage*>(smem_ + Ktraits::kSmemIOSize);
+    auto& smem_exchange1 = *reinterpret_cast<typename Ktraits::BlockExchangeT::TempStorage*>(smem_ + Ktraits::kSmemIOSize + sizeof(typename Ktraits::BlockExchangeT::TempStorage));
+    auto& smem_reduce = *reinterpret_cast<typename Ktraits::BlockReduceT::TempStorage*>(reinterpret_cast<char *>(&smem_exchange) + Ktraits::kSmemExchangeSize);
+    auto& smem_reduce_float = *reinterpret_cast<typename Ktraits::BlockReduceFloatT::TempStorage*>(&smem_reduce);
+    auto& smem_scan = *reinterpret_cast<typename Ktraits::BlockScanT::TempStorage*>(reinterpret_cast<char *>(&smem_reduce) + Ktraits::kSmemReduceSize);
+    auto& smem_reverse_scan = *reinterpret_cast<typename Ktraits::BlockReverseScanT::TempStorage*>(reinterpret_cast<char *>(&smem_scan) + sizeof(typename Ktraits::BlockScanT::TempStorage));
+    weight_t *smem_delta_a = reinterpret_cast<weight_t *>(smem_ + Ktraits::kSmemSize);
+    // scan_t *smem_running_postfix = reinterpret_cast<scan_t *>(smem_delta_a + kNRows * (2 * Ktraits::MaxDState + kNThreads));
+    scan_t *smem_running_postfix = reinterpret_cast<scan_t *>(smem_delta_a + kNRows * 2 * Ktraits::MaxDState + kNThreads);
+    weight_t *smem_da = reinterpret_cast<weight_t *>(smem_running_postfix + kNRows * Ktraits::MaxDState);
+
+    const int batch_id = blockIdx.x;
+    const int dim_id = blockIdx.y;
+    const int dim_id_nrow = dim_id * kNRows;
+    const int group_id = dim_id_nrow / (params.dim_ngroups_ratio);
+    input_t *u = reinterpret_cast<input_t *>(params.u_ptr) + batch_id * params.u_batch_stride
+        + dim_id_nrow * params.u_d_stride;
+    input_t *delta = reinterpret_cast<input_t *>(params.delta_ptr) + batch_id * params.delta_batch_stride
+        + dim_id_nrow * params.delta_d_stride;
+    input_t *dout = reinterpret_cast<input_t *>(params.dout_ptr) + batch_id * params.dout_batch_stride
+        + dim_id_nrow * params.dout_d_stride;
+    weight_t *A = reinterpret_cast<weight_t *>(params.A_ptr) + dim_id_nrow * params.A_d_stride;
+    input_t *Bvar = reinterpret_cast<input_t *>(params.B_ptr) + batch_id * params.B_batch_stride + group_id * params.B_group_stride;
+    input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + batch_id * params.C_batch_stride + group_id * params.C_group_stride;
+    weight_t *dA = reinterpret_cast<weight_t *>(params.dA_ptr) + dim_id_nrow * params.dA_d_stride;
+    weight_t *dB = reinterpret_cast<weight_t *>(params.dB_ptr)
+        + (batch_id * params.dB_batch_stride + group_id * params.dB_group_stride);
+    weight_t *dC = reinterpret_cast<weight_t *>(params.dC_ptr)
+        + (batch_id * params.dC_batch_stride + group_id * params.dC_group_stride);
+    float *dD = params.dD_ptr == nullptr ? nullptr : reinterpret_cast<float *>(params.dD_ptr) + dim_id_nrow;
+    float *D_val = params.D_ptr == nullptr ? nullptr : reinterpret_cast<float *>(params.D_ptr) + dim_id_nrow;
+    float *ddelta_bias = params.ddelta_bias_ptr == nullptr ? nullptr : reinterpret_cast<float *>(params.ddelta_bias_ptr) + dim_id_nrow;
+    float *delta_bias = params.delta_bias_ptr == nullptr ? nullptr : reinterpret_cast<float *>(params.delta_bias_ptr) + dim_id_nrow;
+    scan_t *x = params.x_ptr == nullptr
+        ? nullptr
+        : reinterpret_cast<scan_t *>(params.x_ptr) + (batch_id * params.dim + dim_id_nrow) * (params.n_chunks) * params.dstate;
+    float dD_val[kNRows] = {0};
+    float ddelta_bias_val[kNRows] = {0};
+
+    constexpr int kChunkSize = kNThreads * kNItems;
+    u += (params.n_chunks - 1) * kChunkSize;
+    delta += (params.n_chunks - 1) * kChunkSize;
+    dout += (params.n_chunks - 1) * kChunkSize;
+    Bvar += (params.n_chunks - 1) * kChunkSize;
+    Cvar += (params.n_chunks - 1) * kChunkSize;
+    for (int chunk = params.n_chunks - 1; chunk >= 0; --chunk) {
+        input_t u_vals[kNRows][kNItems];
+        input_t delta_vals_load[kNRows][kNItems];
+        input_t dout_vals_load[kNRows][kNItems];
+        #pragma unroll
+        for (int r = 0; r < kNRows; ++r) {
+            __syncthreads();
+            load_input<Ktraits>(u + r * params.u_d_stride, u_vals[r], smem_load, params.seqlen - chunk * kChunkSize);
+            __syncthreads();
+            load_input<Ktraits>(delta + r * params.delta_d_stride, delta_vals_load[r], smem_load, params.seqlen - chunk * kChunkSize);
+            __syncthreads();
+            load_input<Ktraits>(dout + r * params.dout_d_stride, dout_vals_load[r], smem_load, params.seqlen - chunk * kChunkSize);
+        }
+        u -= kChunkSize;
+        // Will reload delta at the same location if kDeltaSoftplus
+        if constexpr (!kDeltaSoftplus) { delta -= kChunkSize; }
+        dout -= kChunkSize;
+
+        float dout_vals[kNRows][kNItems], delta_vals[kNRows][kNItems];
+        float du_vals[kNRows][kNItems];
+        #pragma unroll
+        for (int r = 0; r < kNRows; ++r) {
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                dout_vals[r][i] = float(dout_vals_load[r][i]);
+                delta_vals[r][i] = float(delta_vals_load[r][i]) + (delta_bias == nullptr? 0: delta_bias[r]);
+                if constexpr (kDeltaSoftplus) {
+                    delta_vals[r][i] = delta_vals[r][i] <= 20.f ? log1pf(expf(delta_vals[r][i])) : delta_vals[r][i];
+                }
+            }
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) { du_vals[r][i] = (D_val == nullptr? 0: D_val[r]) * dout_vals[r][i]; }
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) { dD_val[r] += dout_vals[r][i] * float(u_vals[r][i]); }
+        }
+
+        float ddelta_vals[kNRows][kNItems] = {0};
+        __syncthreads();
+        for (int state_idx = 0; state_idx < params.dstate; ++state_idx) {
+            weight_t A_val[kNRows];
+            weight_t A_scaled[kNRows];
+            #pragma unroll
+            for (int r = 0; r < kNRows; ++r) {
+                A_val[r] = A[state_idx * params.A_dstate_stride + r * params.A_d_stride];
+                constexpr float kLog2e = M_LOG2E;
+                A_scaled[r] = A_val[r] * kLog2e;
+            }
+            weight_t B_vals[kNItems], C_vals[kNItems];
+            load_weight<Ktraits>(Bvar + state_idx * params.B_dstate_stride, B_vals,
+                    smem_load_weight, (params.seqlen - chunk * kChunkSize));
+            auto &smem_load_weight_C = smem_load_weight1;
+            load_weight<Ktraits>(Cvar + state_idx * params.C_dstate_stride, C_vals,
+                    smem_load_weight_C, (params.seqlen - chunk * kChunkSize));
+            #pragma unroll
+            for (int r = 0; r < kNRows; ++r) {
+                scan_t thread_data[kNItems], thread_reverse_data[kNItems];
+                #pragma unroll
+                for (int i = 0; i < kNItems; ++i) {
+                    const float delta_a_exp = exp2f(delta_vals[r][i] * A_scaled[r]);
+                    thread_data[i] = make_float2(delta_a_exp, delta_vals[r][i] * float(u_vals[r][i]) * B_vals[i]);
+                    if (i == 0) {
+                        // smem_delta_a[threadIdx.x == 0 ? state_idx + (chunk % 2) * Ktraits::MaxDState + r * (2 * Ktraits::MaxDState + kNThreads) : threadIdx.x + 2 * Ktraits::MaxDState + r * (2 * Ktraits::MaxDState + kNThreads)] = delta_a_exp;
+                        smem_delta_a[threadIdx.x == 0 ? state_idx + (chunk % 2) * Ktraits::MaxDState + r * 2 * Ktraits::MaxDState : threadIdx.x + kNRows * 2 * Ktraits::MaxDState] = delta_a_exp;
+                        
+                    } else {
+                        thread_reverse_data[i - 1].x = delta_a_exp;
+                    }
+                    thread_reverse_data[i].y = dout_vals[r][i] * C_vals[i];
+                }
+                __syncthreads();
+                thread_reverse_data[kNItems - 1].x = threadIdx.x == kNThreads - 1
+                    // ? (chunk == params.n_chunks - 1 ? 1.f : smem_delta_a[state_idx + ((chunk + 1) % 2) * Ktraits::MaxDState + r * (2 * Ktraits::MaxDState + kNThreads)])
+                    // : smem_delta_a[threadIdx.x + 1 + 2 * Ktraits::MaxDState + r * (2 * Ktraits::MaxDState + kNThreads)];
+                    ? (chunk == params.n_chunks - 1 ? 1.f : smem_delta_a[state_idx + ((chunk + 1) % 2) * Ktraits::MaxDState + r * 2 * Ktraits::MaxDState])
+                    : smem_delta_a[threadIdx.x + 1 + kNRows * 2 * Ktraits::MaxDState];
+                // Initialize running total
+                scan_t running_prefix = chunk > 0 && threadIdx.x % 32 == 0 ? x[(r * params.n_chunks + chunk - 1) * params.dstate + state_idx] : make_float2(1.f, 0.f);
+                SSMScanPrefixCallbackOp<weight_t> prefix_op(running_prefix);
+                Ktraits::BlockScanT(smem_scan).InclusiveScan(
+                    thread_data, thread_data, SSMScanOp<weight_t>(), prefix_op
+                );
+                scan_t running_postfix = chunk < params.n_chunks - 1 && threadIdx.x % 32 == 0 ? smem_running_postfix[state_idx + r * Ktraits::MaxDState] : make_float2(1.f, 0.f);
+                SSMScanPrefixCallbackOp<weight_t> postfix_op(running_postfix);
+                Ktraits::BlockReverseScanT(smem_reverse_scan).InclusiveReverseScan(
+                    thread_reverse_data, thread_reverse_data, SSMScanOp<weight_t>(), postfix_op
+                );
+                if (threadIdx.x == 0) { smem_running_postfix[state_idx + r * Ktraits::MaxDState] = postfix_op.running_prefix; }
+                weight_t dA_val = 0;
+                weight_t dB_vals[kNItems], dC_vals[kNItems];
+                #pragma unroll
+                for (int i = 0; i < kNItems; ++i) {
+                    const float dx = thread_reverse_data[i].y;
+                    const float ddelta_u = dx * B_vals[i];
+                    du_vals[r][i] += ddelta_u * delta_vals[r][i];
+                    const float a = thread_data[i].y - (delta_vals[r][i] * float(u_vals[r][i]) * B_vals[i]);
+                    ddelta_vals[r][i] += ddelta_u * float(u_vals[r][i]) + dx * A_val[r] * a;
+                    dA_val += dx * delta_vals[r][i] * a;
+                    dB_vals[i] = dx * delta_vals[r][i] * float(u_vals[r][i]);
+                    dC_vals[i] = dout_vals[r][i] * thread_data[i].y;
+                }
+                // Block-exchange to make the atomicAdd's coalesced, otherwise they're much slower
+                    Ktraits::BlockExchangeT(smem_exchange).BlockedToStriped(dB_vals, dB_vals);
+                    auto &smem_exchange_C = smem_exchange1;
+                    Ktraits::BlockExchangeT(smem_exchange_C).BlockedToStriped(dC_vals, dC_vals);
+                    const int seqlen_remaining = params.seqlen - chunk * kChunkSize - threadIdx.x;
+                    weight_t *dB_cur = dB + state_idx * params.dB_dstate_stride + chunk * kChunkSize + threadIdx.x;
+                    weight_t *dC_cur = dC + state_idx * params.dC_dstate_stride + chunk * kChunkSize + threadIdx.x;
+                    #pragma unroll
+                    for (int i = 0; i < kNItems; ++i) {
+                        if (i * kNThreads < seqlen_remaining) {
+                            { gpuAtomicAdd(dB_cur + i * kNThreads, dB_vals[i]); }
+                            { gpuAtomicAdd(dC_cur + i * kNThreads, dC_vals[i]); }
+                        }
+                    }
+                dA_val = Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(dA_val);
+                if (threadIdx.x == 0) {
+                    smem_da[state_idx + r * Ktraits::MaxDState] = chunk == params.n_chunks - 1 ? dA_val : dA_val + smem_da[state_idx + r * Ktraits::MaxDState];
+                }
+            }
+        }
+
+        if constexpr (kDeltaSoftplus) {
+            input_t delta_vals_load[kNRows][kNItems];
+            #pragma unroll
+            for (int r = 0; r < kNRows; ++r) {
+                __syncthreads();
+                load_input<Ktraits>(delta + r * params.delta_d_stride, delta_vals_load[r], smem_load, params.seqlen - chunk * kChunkSize);
+            }
+            delta -= kChunkSize;
+            #pragma unroll
+            for (int r = 0; r < kNRows; ++r) {
+                #pragma unroll
+                for (int i = 0; i < kNItems; ++i) {
+                    float delta_val = float(delta_vals_load[r][i]) + (delta_bias == nullptr? 0: delta_bias[r]);
+                    float delta_val_neg_exp = expf(-delta_val);
+                    ddelta_vals[r][i] = delta_val <= 20.f
+                        ? ddelta_vals[r][i] / (1.f + delta_val_neg_exp)
+                        : ddelta_vals[r][i];
+                }
+            }
+        }
+        
+        __syncthreads();
+        #pragma unroll
+        for (int r = 0; r < kNRows; ++r) {
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) { ddelta_bias_val[r] += ddelta_vals[r][i]; }
+        }
+
+        input_t *du = reinterpret_cast<input_t *>(params.du_ptr) + batch_id * params.du_batch_stride
+            + dim_id_nrow * params.du_d_stride + chunk * kChunkSize;
+        input_t *ddelta = reinterpret_cast<input_t *>(params.ddelta_ptr) + batch_id * params.ddelta_batch_stride
+            + dim_id_nrow * params.ddelta_d_stride + chunk * kChunkSize;
+        #pragma unroll
+        for (int r = 0; r < kNRows; ++r) {
+            __syncthreads();
+            store_output<Ktraits>(du + r * params.du_d_stride, du_vals[r], smem_store, params.seqlen - chunk * kChunkSize);
+            __syncthreads();
+            store_output<Ktraits>(ddelta + r * params.ddelta_d_stride, ddelta_vals[r], smem_store, params.seqlen - chunk * kChunkSize);
+        }
+
+        Bvar -= kChunkSize;
+        Cvar -= kChunkSize;
+    }
+
+    #pragma unroll
+    for (int r = 0; r < kNRows; ++r) {
+        if (params.dD_ptr != nullptr) {
+            __syncthreads();
+            dD_val[r] = Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(dD_val[r]);
+            if (threadIdx.x == 0) { gpuAtomicAdd(&(dD[r]), dD_val[r]); }
+        }
+        if (params.ddelta_bias_ptr != nullptr) {
+            __syncthreads();
+            ddelta_bias_val[r] = Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(ddelta_bias_val[r]);
+            if (threadIdx.x == 0) { gpuAtomicAdd(&(ddelta_bias[r]), ddelta_bias_val[r]); }
+        }
+        __syncthreads();
+        for (int state_idx = threadIdx.x; state_idx < params.dstate; state_idx += blockDim.x) {
+            gpuAtomicAdd(&(dA[state_idx * params.dA_dstate_stride + r * params.dA_d_stride]), smem_da[state_idx + r * Ktraits::MaxDState]);
+        }
+    }
+}
+
+template<int kNThreads, int kNItems, int kNRows, typename input_t, typename weight_t>
+void selective_scan_bwd_launch(SSMParamsBwd &params, cudaStream_t stream) {
+    BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {
+        BOOL_SWITCH(params.delta_softplus, kDeltaSoftplus, [&] {
+            using Ktraits = Selective_Scan_bwd_kernel_traits<kNThreads, kNItems, kNRows, kIsEvenLen, kDeltaSoftplus, input_t, weight_t>;
+            constexpr int kSmemSize = Ktraits::kSmemSize + kNRows * Ktraits::MaxDState * sizeof(typename Ktraits::scan_t) + (kNThreads + 4 * kNRows * Ktraits::MaxDState) * sizeof(typename Ktraits::weight_t);
+            // printf("smem_size = %d\n", kSmemSize);
+            dim3 grid(params.batch, params.dim / kNRows);
+            auto kernel = &selective_scan_bwd_kernel<Ktraits>;
+            if (kSmemSize >= 48 * 1024) {
+                C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
+            }
+            kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
+            C10_CUDA_KERNEL_LAUNCH_CHECK();
+        });
+    });
+}
+
+template<int knrows, typename input_t, typename weight_t>
+void selective_scan_bwd_cuda(SSMParamsBwd &params, cudaStream_t stream) {
+    if (params.seqlen <= 128) {
+        selective_scan_bwd_launch<32, 4, knrows, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 256) {
+        selective_scan_bwd_launch<32, 8, knrows, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 512) {
+        selective_scan_bwd_launch<32, 16, knrows, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 1024) {
+        selective_scan_bwd_launch<64, 16, knrows, input_t, weight_t>(params, stream);
+    } else {
+        selective_scan_bwd_launch<128, 16, knrows, input_t, weight_t>(params, stream);
+    }
+}

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusnrow/selective_scan_core_bwd.cu

@@ -0,0 +1,9 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+#include "selective_scan_bwd_kernel_nrow.cuh"
+
+template void selective_scan_bwd_cuda<1, float, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<1, at::Half, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<1, at::BFloat16, float>(SSMParamsBwd &params, cudaStream_t stream);
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusnrow/selective_scan_core_bwd2.cu

@@ -0,0 +1,9 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+#include "selective_scan_bwd_kernel_nrow.cuh"
+
+template void selective_scan_bwd_cuda<2, float, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<2, at::Half, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<2, at::BFloat16, float>(SSMParamsBwd &params, cudaStream_t stream);
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusnrow/selective_scan_core_bwd3.cu

@@ -0,0 +1,8 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+#include "selective_scan_bwd_kernel_nrow.cuh"
+
+template void selective_scan_bwd_cuda<3, float, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<3, at::Half, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<3, at::BFloat16, float>(SSMParamsBwd &params, cudaStream_t stream);

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusnrow/selective_scan_core_bwd4.cu

@@ -0,0 +1,8 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+#include "selective_scan_bwd_kernel_nrow.cuh"
+
+template void selective_scan_bwd_cuda<4, float, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<4, at::Half, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<4, at::BFloat16, float>(SSMParamsBwd &params, cudaStream_t stream);

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusnrow/selective_scan_core_fwd.cu

@@ -0,0 +1,9 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+#include "selective_scan_fwd_kernel_nrow.cuh"
+
+template void selective_scan_fwd_cuda<1, float, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<1, at::Half, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<1, at::BFloat16, float>(SSMParamsBase &params, cudaStream_t stream);
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusnrow/selective_scan_core_fwd2.cu

@@ -0,0 +1,9 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+#include "selective_scan_fwd_kernel_nrow.cuh"
+
+template void selective_scan_fwd_cuda<2, float, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<2, at::Half, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<2, at::BFloat16, float>(SSMParamsBase &params, cudaStream_t stream);
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusnrow/selective_scan_core_fwd3.cu

@@ -0,0 +1,9 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+#include "selective_scan_fwd_kernel_nrow.cuh"
+
+template void selective_scan_fwd_cuda<3, float, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<3, at::Half, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<3, at::BFloat16, float>(SSMParamsBase &params, cudaStream_t stream);
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusnrow/selective_scan_core_fwd4.cu

@@ -0,0 +1,9 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+#include "selective_scan_fwd_kernel_nrow.cuh"
+
+template void selective_scan_fwd_cuda<4, float, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<4, at::Half, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<4, at::BFloat16, float>(SSMParamsBase &params, cudaStream_t stream);
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusnrow/selective_scan_fwd_kernel_nrow.cuh

@@ -0,0 +1,238 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <c10/util/BFloat16.h>
+#include <c10/util/Half.h>
+#include <c10/cuda/CUDAException.h>  // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
+
+#include <cub/block/block_load.cuh>
+#include <cub/block/block_store.cuh>
+#include <cub/block/block_scan.cuh>
+
+#include "selective_scan.h"
+#include "selective_scan_common.h"
+#include "static_switch.h"
+
+template<int kNThreads_, int kNItems_, int kNRows_, bool kIsEvenLen_, typename input_t_, typename weight_t_>
+struct Selective_Scan_fwd_kernel_traits {
+    static_assert(kNItems_ % 4 == 0);
+    using input_t = input_t_;
+    using weight_t = weight_t_;
+    static constexpr int kNThreads = kNThreads_;
+    // Setting MinBlocksPerMP to be 3 (instead of 2) for 128 threads improves occupancy.
+    static constexpr int kMinBlocks = kNThreads < 128 ? 5 : 3;
+    static constexpr int kNItems = kNItems_;
+    static constexpr int kNRows = kNRows_;
+    static constexpr int MaxDState = MAX_DSTATE / kNRows;
+    static constexpr int kNBytes = sizeof(input_t);
+    static_assert(kNBytes == 2 || kNBytes == 4);
+    static constexpr int kNElts = kNBytes == 4 ? 4 : std::min(8, kNItems);
+    static_assert(kNItems % kNElts == 0);
+    static constexpr int kNLoads = kNItems / kNElts;
+    static constexpr bool kIsEvenLen = kIsEvenLen_;
+
+    static constexpr bool kDirectIO = kIsEvenLen && kNLoads == 1;
+
+    using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
+    using scan_t = float2;
+    using BlockLoadT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads,
+        !kDirectIO ? cub::BLOCK_LOAD_WARP_TRANSPOSE : cub::BLOCK_LOAD_DIRECT>;
+    using BlockLoadWeightT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadWeightVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads,
+        !kDirectIO ? cub::BLOCK_LOAD_WARP_TRANSPOSE  : cub::BLOCK_LOAD_DIRECT>;
+    using BlockStoreT = cub::BlockStore<input_t, kNThreads, kNItems, cub::BLOCK_STORE_WARP_TRANSPOSE>;
+    using BlockStoreVecT = cub::BlockStore<vec_t, kNThreads, kNLoads,
+        !kDirectIO ? cub::BLOCK_STORE_WARP_TRANSPOSE : cub::BLOCK_STORE_DIRECT>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING_MEMOIZE>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING>;
+    using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_WARP_SCANS>;
+    static constexpr int kSmemIOSize = std::max({sizeof(typename BlockLoadT::TempStorage),
+                                                 sizeof(typename BlockLoadVecT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightVecT::TempStorage),
+                                                 sizeof(typename BlockStoreT::TempStorage),
+                                                 sizeof(typename BlockStoreVecT::TempStorage)});
+    static constexpr int kSmemSize = kSmemIOSize + sizeof(typename BlockScanT::TempStorage);
+};
+
+template<typename Ktraits>
+__global__ __launch_bounds__(Ktraits::kNThreads, Ktraits::kMinBlocks)
+void selective_scan_fwd_kernel(SSMParamsBase params) {
+    constexpr int kNThreads = Ktraits::kNThreads;
+    constexpr int kNItems = Ktraits::kNItems;
+    constexpr int kNRows = Ktraits::kNRows;
+    constexpr bool kDirectIO = Ktraits::kDirectIO;
+    using input_t = typename Ktraits::input_t;
+    using weight_t = typename Ktraits::weight_t;
+    using scan_t = typename Ktraits::scan_t;
+
+    // Shared memory.
+    extern __shared__ char smem_[];
+    auto& smem_load = reinterpret_cast<typename Ktraits::BlockLoadT::TempStorage&>(smem_);
+    auto& smem_load_weight = reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage&>(smem_);
+    auto& smem_load_weight1 = *reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage*>(smem_ + sizeof(typename Ktraits::BlockLoadWeightT::TempStorage));
+    auto& smem_store = reinterpret_cast<typename Ktraits::BlockStoreT::TempStorage&>(smem_);
+    auto& smem_scan = *reinterpret_cast<typename Ktraits::BlockScanT::TempStorage*>(smem_ + Ktraits::kSmemIOSize);
+    scan_t *smem_running_prefix = reinterpret_cast<scan_t *>(smem_ + Ktraits::kSmemSize);
+
+    const int batch_id = blockIdx.x;
+    const int dim_id = blockIdx.y;
+    const int dim_id_nrow = dim_id * kNRows;
+    const int group_id = dim_id_nrow / (params.dim_ngroups_ratio);
+    input_t *u = reinterpret_cast<input_t *>(params.u_ptr) + batch_id * params.u_batch_stride
+        + dim_id_nrow * params.u_d_stride;
+    input_t *delta = reinterpret_cast<input_t *>(params.delta_ptr) + batch_id * params.delta_batch_stride
+        + dim_id_nrow * params.delta_d_stride;
+    weight_t *A = reinterpret_cast<weight_t *>(params.A_ptr) + dim_id_nrow * params.A_d_stride;
+    input_t *Bvar = reinterpret_cast<input_t *>(params.B_ptr) + batch_id * params.B_batch_stride + group_id * params.B_group_stride;
+    input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + batch_id * params.C_batch_stride + group_id * params.C_group_stride;
+    scan_t *x = reinterpret_cast<scan_t *>(params.x_ptr) + (batch_id * params.dim + dim_id_nrow) * params.n_chunks * params.dstate;
+
+    float D_val[kNRows] = {0};
+    if (params.D_ptr != nullptr) {
+        #pragma unroll
+        for (int r = 0; r < kNRows; ++r) {
+            D_val[r] = reinterpret_cast<float *>(params.D_ptr)[dim_id_nrow + r];
+        }
+    }
+    float delta_bias[kNRows] = {0};
+    if (params.delta_bias_ptr != nullptr) {
+        #pragma unroll
+        for (int r = 0; r < kNRows; ++r) {
+            delta_bias[r] = reinterpret_cast<float *>(params.delta_bias_ptr)[dim_id_nrow + r];
+        }
+    }
+
+    constexpr int kChunkSize = kNThreads * kNItems;
+    for (int chunk = 0; chunk < params.n_chunks; ++chunk) {
+        input_t u_vals[kNRows][kNItems], delta_vals_load[kNRows][kNItems];
+        __syncthreads();
+        #pragma unroll
+        for (int r = 0; r < kNRows; ++r) {
+            if constexpr (!kDirectIO) {
+                if (r > 0) { __syncthreads(); }
+            }
+            load_input<Ktraits>(u + r * params.u_d_stride, u_vals[r], smem_load, params.seqlen - chunk * kChunkSize);
+            if constexpr (!kDirectIO) { __syncthreads(); }
+            load_input<Ktraits>(delta + r * params.delta_d_stride, delta_vals_load[r], smem_load, params.seqlen - chunk * kChunkSize);
+        }
+        u += kChunkSize;
+        delta += kChunkSize;
+
+        float delta_vals[kNRows][kNItems], delta_u_vals[kNRows][kNItems], out_vals[kNRows][kNItems];
+        #pragma unroll
+        for (int r = 0; r < kNRows; ++r) {
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                float u_val = float(u_vals[r][i]);
+                delta_vals[r][i] = float(delta_vals_load[r][i]) + delta_bias[r];
+                if (params.delta_softplus) {
+                    delta_vals[r][i] = delta_vals[r][i] <= 20.f ? log1pf(expf(delta_vals[r][i])) : delta_vals[r][i];
+                }
+                delta_u_vals[r][i] = delta_vals[r][i] * u_val;
+                out_vals[r][i] = D_val[r] * u_val;
+            }
+        }
+
+        __syncthreads();
+        for (int state_idx = 0; state_idx < params.dstate; ++state_idx) {
+            weight_t A_val[kNRows];
+            #pragma unroll
+            for (int r = 0; r < kNRows; ++r) {
+                A_val[r] = A[state_idx * params.A_dstate_stride + r * params.A_d_stride];
+                // Multiply the real part of A with LOG2E so we can use exp2f instead of expf.
+                constexpr float kLog2e = M_LOG2E;
+                A_val[r] *= kLog2e;
+            }
+            weight_t B_vals[kNItems], C_vals[kNItems];
+            load_weight<Ktraits>(Bvar + state_idx * params.B_dstate_stride, B_vals,
+                    smem_load_weight, (params.seqlen - chunk * kChunkSize));
+            auto &smem_load_weight_C = smem_load_weight1;
+            load_weight<Ktraits>(Cvar + state_idx * params.C_dstate_stride, C_vals,
+                    smem_load_weight_C, (params.seqlen - chunk * kChunkSize));
+            #pragma unroll
+            for (int r = 0; r < kNRows; ++r) {
+                if (r > 0) { __syncthreads(); }  // Scan could be using the same smem
+                scan_t thread_data[kNItems];
+                #pragma unroll
+                for (int i = 0; i < kNItems; ++i) {
+                    thread_data[i] = make_float2(exp2f(delta_vals[r][i] * A_val[r]),
+                                                 B_vals[i] * delta_u_vals[r][i]);
+                    if constexpr (!Ktraits::kIsEvenLen) {  // So that the last state is correct
+                        if (threadIdx.x * kNItems + i >= params.seqlen - chunk * kChunkSize) {
+                            thread_data[i] = make_float2(1.f, 0.f);
+                        }
+                    }
+                }
+                // Initialize running total
+                scan_t running_prefix;
+                // If we use WARP_SCAN then all lane 0 of all warps (not just thread 0) needs to read
+                running_prefix = chunk > 0 && threadIdx.x % 32 == 0 ? smem_running_prefix[state_idx + r * Ktraits::MaxDState] : make_float2(1.f, 0.f);
+                // running_prefix = chunk > 0 && threadIdx.x == 0 ? smem_running_prefix[state_idx] : make_float2(1.f, 0.f);
+                SSMScanPrefixCallbackOp<weight_t> prefix_op(running_prefix);
+                Ktraits::BlockScanT(smem_scan).InclusiveScan(
+                    thread_data, thread_data, SSMScanOp<weight_t>(), prefix_op
+                );
+                // There's a syncthreads in the scan op, so we don't need to sync here.
+                // Unless there's only 1 warp, but then it's the same thread (0) reading and writing.
+                if (threadIdx.x == 0) {
+                    smem_running_prefix[state_idx + r * Ktraits::MaxDState] = prefix_op.running_prefix;
+                    x[(r * params.n_chunks + chunk) * params.dstate + state_idx] = prefix_op.running_prefix;
+                }
+                #pragma unroll
+                for (int i = 0; i < kNItems; ++i) {
+                    out_vals[r][i] += thread_data[i].y * C_vals[i];
+                }
+            }
+        }
+
+        input_t *out = reinterpret_cast<input_t *>(params.out_ptr) + batch_id * params.out_batch_stride
+            + dim_id_nrow * params.out_d_stride + chunk * kChunkSize;
+        __syncthreads();
+        #pragma unroll
+        for (int r = 0; r < kNRows; ++r) {
+            if constexpr (!kDirectIO) {
+                if (r > 0) { __syncthreads(); }
+            }
+            store_output<Ktraits>(out + r * params.out_d_stride, out_vals[r], smem_store, params.seqlen - chunk * kChunkSize);
+        }
+
+        Bvar += kChunkSize;
+        Cvar += kChunkSize;
+    }
+}
+
+template<int kNThreads, int kNItems, int kNRows, typename input_t, typename weight_t>
+void selective_scan_fwd_launch(SSMParamsBase &params, cudaStream_t stream) {
+    BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {
+        using Ktraits = Selective_Scan_fwd_kernel_traits<kNThreads, kNItems, kNRows, kIsEvenLen, input_t, weight_t>;
+        constexpr int kSmemSize = Ktraits::kSmemSize + kNRows * Ktraits::MaxDState * sizeof(typename Ktraits::scan_t);
+        // printf("smem_size = %d\n", kSmemSize);
+        dim3 grid(params.batch, params.dim / kNRows);
+        auto kernel = &selective_scan_fwd_kernel<Ktraits>;
+        if (kSmemSize >= 48 * 1024) {
+            C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
+        }
+        kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+    });
+}
+
+template<int knrows, typename input_t, typename weight_t>
+void selective_scan_fwd_cuda(SSMParamsBase &params, cudaStream_t stream) {
+    if (params.seqlen <= 128) {
+        selective_scan_fwd_launch<32, 4, knrows, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 256) {
+        selective_scan_fwd_launch<32, 8, knrows, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 512) {
+        selective_scan_fwd_launch<32, 16, knrows, input_t, weight_t>(params, stream);
+    } else if (params.seqlen <= 1024) {
+        selective_scan_fwd_launch<64, 16, knrows, input_t, weight_t>(params, stream);
+    } else {
+        selective_scan_fwd_launch<128, 16, knrows, input_t, weight_t>(params, stream);
+    }
+}

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusnrow/selective_scan_nrow.cpp

@@ -0,0 +1,367 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/extension.h>
+#include <vector>
+
+#include "selective_scan.h"
+#define MAX_DSTATE 256
+
+#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
+using weight_t = float;
+
+#define DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, NAME, ...)                    \
+    if (ITYPE == at::ScalarType::Half) {                                            \
+        using input_t = at::Half;                                                   \
+        __VA_ARGS__();                                                              \
+    } else if (ITYPE == at::ScalarType::BFloat16) {                                 \
+        using input_t = at::BFloat16;                                               \
+        __VA_ARGS__();                                                              \
+    } else if (ITYPE == at::ScalarType::Float)  {                                   \
+        using input_t = float;                                                      \
+        __VA_ARGS__();                                                              \
+    } else {                                                                        \
+        AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \
+    }
+
+#define INT_SWITCH(INT, NAME, ...) [&] {                         \
+    if (INT == 2) {constexpr int NAME = 2; __VA_ARGS__(); }      \
+    else if (INT == 3) {constexpr int NAME = 3; __VA_ARGS__(); } \
+    else if (INT == 4) {constexpr int NAME = 4; __VA_ARGS__(); } \
+    else {constexpr int NAME = 1; __VA_ARGS__(); }               \
+}()                                                              \
+
+
+template<int knrows, typename input_t, typename weight_t>
+void selective_scan_fwd_cuda(SSMParamsBase &params, cudaStream_t stream);
+
+template <int knrows, typename input_t, typename weight_t>
+void selective_scan_bwd_cuda(SSMParamsBwd &params, cudaStream_t stream);
+
+void set_ssm_params_fwd(SSMParamsBase &params,
+                        // sizes
+                        const size_t batch,
+                        const size_t dim,
+                        const size_t seqlen,
+                        const size_t dstate,
+                        const size_t n_groups,
+                        const size_t n_chunks,
+                        // device pointers
+                        const at::Tensor u,
+                        const at::Tensor delta,
+                        const at::Tensor A,
+                        const at::Tensor B,
+                        const at::Tensor C,
+                        const at::Tensor out,
+                        void* D_ptr,
+                        void* delta_bias_ptr,
+                        void* x_ptr,
+                        bool delta_softplus) {
+
+    // Reset the parameters
+    memset(&params, 0, sizeof(params));
+
+    params.batch = batch;
+    params.dim = dim;
+    params.seqlen = seqlen;
+    params.dstate = dstate;
+    params.n_groups = n_groups;
+    params.n_chunks = n_chunks;
+    params.dim_ngroups_ratio = dim / n_groups;
+
+    params.delta_softplus = delta_softplus;
+
+    // Set the pointers and strides.
+    params.u_ptr = u.data_ptr();
+    params.delta_ptr = delta.data_ptr();
+    params.A_ptr = A.data_ptr();
+    params.B_ptr = B.data_ptr();
+    params.C_ptr = C.data_ptr();
+    params.D_ptr = D_ptr;
+    params.delta_bias_ptr = delta_bias_ptr;
+    params.out_ptr = out.data_ptr();
+    params.x_ptr = x_ptr;
+
+    // All stride are in elements, not bytes.
+    params.A_d_stride = A.stride(0);
+    params.A_dstate_stride = A.stride(1);
+    params.B_batch_stride = B.stride(0);
+    params.B_group_stride = B.stride(1);
+    params.B_dstate_stride = B.stride(2);
+    params.C_batch_stride = C.stride(0);
+    params.C_group_stride = C.stride(1);
+    params.C_dstate_stride = C.stride(2);
+    params.u_batch_stride = u.stride(0);
+    params.u_d_stride = u.stride(1);
+    params.delta_batch_stride = delta.stride(0);
+    params.delta_d_stride = delta.stride(1);
+
+    params.out_batch_stride = out.stride(0);
+    params.out_d_stride = out.stride(1);
+}
+
+void set_ssm_params_bwd(SSMParamsBwd &params,
+                        // sizes
+                        const size_t batch,
+                        const size_t dim,
+                        const size_t seqlen,
+                        const size_t dstate,
+                        const size_t n_groups,
+                        const size_t n_chunks,
+                        // device pointers
+                        const at::Tensor u,
+                        const at::Tensor delta,
+                        const at::Tensor A,
+                        const at::Tensor B,
+                        const at::Tensor C,
+                        const at::Tensor out,
+                        void* D_ptr,
+                        void* delta_bias_ptr,
+                        void* x_ptr,
+                        const at::Tensor dout,
+                        const at::Tensor du,
+                        const at::Tensor ddelta,
+                        const at::Tensor dA,
+                        const at::Tensor dB,
+                        const at::Tensor dC,
+                        void* dD_ptr,
+                        void* ddelta_bias_ptr,
+                        bool delta_softplus) {
+    // Pass in "dout" instead of "out", we're not gonna use "out" unless we have z
+    set_ssm_params_fwd(params, batch, dim, seqlen, dstate, n_groups, n_chunks,
+                       u, delta, A, B, C, dout,
+                       D_ptr, delta_bias_ptr, x_ptr, delta_softplus);
+
+    // Set the pointers and strides.
+    params.dout_ptr = dout.data_ptr();
+    params.du_ptr = du.data_ptr();
+    params.dA_ptr = dA.data_ptr();
+    params.dB_ptr = dB.data_ptr();
+    params.dC_ptr = dC.data_ptr();
+    params.dD_ptr = dD_ptr;
+    params.ddelta_ptr = ddelta.data_ptr();
+    params.ddelta_bias_ptr = ddelta_bias_ptr;
+    // All stride are in elements, not bytes.
+    params.dout_batch_stride = dout.stride(0);
+    params.dout_d_stride = dout.stride(1);
+    params.dA_d_stride = dA.stride(0);
+    params.dA_dstate_stride = dA.stride(1);
+    params.dB_batch_stride = dB.stride(0);
+    params.dB_group_stride = dB.stride(1);
+    params.dB_dstate_stride = dB.stride(2);
+    params.dC_batch_stride = dC.stride(0);
+    params.dC_group_stride = dC.stride(1);
+    params.dC_dstate_stride = dC.stride(2);
+    params.du_batch_stride = du.stride(0);
+    params.du_d_stride = du.stride(1);
+    params.ddelta_batch_stride = ddelta.stride(0);
+    params.ddelta_d_stride = ddelta.stride(1);
+
+}
+
+std::vector<at::Tensor>
+selective_scan_fwd(const at::Tensor &u, const at::Tensor &delta,
+                  const at::Tensor &A, const at::Tensor &B, const at::Tensor &C,
+                  const c10::optional<at::Tensor> &D_,
+                  const c10::optional<at::Tensor> &delta_bias_,
+                  bool delta_softplus,
+                  int nrows
+                  ) {
+    auto input_type = u.scalar_type();
+    auto weight_type = A.scalar_type();
+    TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
+    TORCH_CHECK(weight_type == at::ScalarType::Float);
+
+    TORCH_CHECK(delta.scalar_type() == input_type);
+    TORCH_CHECK(B.scalar_type() == input_type);
+    TORCH_CHECK(C.scalar_type() == input_type);
+
+    TORCH_CHECK(u.is_cuda());
+    TORCH_CHECK(delta.is_cuda());
+    TORCH_CHECK(A.is_cuda());
+    TORCH_CHECK(B.is_cuda());
+    TORCH_CHECK(C.is_cuda());
+
+    TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
+    TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
+
+    const auto sizes = u.sizes();
+    const int batch_size = sizes[0];
+    const int dim = sizes[1];
+    const int seqlen = sizes[2];
+    const int dstate = A.size(1);
+    const int n_groups = B.size(1);
+
+    TORCH_CHECK(dim % (n_groups * nrows) == 0, "dims should be dividable by n_groups * nrows");
+    TORCH_CHECK(dstate <= MAX_DSTATE / nrows, "selective_scan only supports state dimension <= 256 / nrows");
+
+    CHECK_SHAPE(u, batch_size, dim, seqlen);
+    CHECK_SHAPE(delta, batch_size, dim, seqlen);
+    CHECK_SHAPE(A, dim, dstate);
+    CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen);
+    TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
+    CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen);
+    TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
+
+    if (D_.has_value()) {
+        auto D = D_.value();
+        TORCH_CHECK(D.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(D.is_cuda());
+        TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
+        CHECK_SHAPE(D, dim);
+    }
+
+    if (delta_bias_.has_value()) {
+        auto delta_bias = delta_bias_.value();
+        TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(delta_bias.is_cuda());
+        TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
+        CHECK_SHAPE(delta_bias, dim);
+    }
+
+    const int n_chunks = (seqlen + 2048 - 1) / 2048; // max is 128 * 16 = 2048 in fwd_kernel
+    at::Tensor out = torch::empty_like(delta);
+    at::Tensor x;
+    x = torch::empty({batch_size, dim, n_chunks, dstate * 2}, u.options().dtype(weight_type));
+
+    SSMParamsBase params;
+    set_ssm_params_fwd(params, batch_size, dim, seqlen, dstate, n_groups, n_chunks,
+                       u, delta, A, B, C, out,
+                       D_.has_value() ? D_.value().data_ptr() : nullptr,
+                       delta_bias_.has_value() ? delta_bias_.value().data_ptr() : nullptr,
+                       x.data_ptr(),
+                       delta_softplus);
+
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)u.get_device()};
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), "selective_scan_fwd", [&] {
+        INT_SWITCH(nrows, kNRows, [&] {
+            selective_scan_fwd_cuda<kNRows, input_t, weight_t>(params, stream);
+        });
+    });
+    std::vector<at::Tensor> result = {out, x};
+    return result;
+}
+
+std::vector<at::Tensor>
+selective_scan_bwd(const at::Tensor &u, const at::Tensor &delta,
+                  const at::Tensor &A, const at::Tensor &B, const at::Tensor &C,
+                  const c10::optional<at::Tensor> &D_,
+                  const c10::optional<at::Tensor> &delta_bias_,
+                  const at::Tensor &dout,
+                  const c10::optional<at::Tensor> &x_,
+                  bool delta_softplus,
+                  int nrows
+                  ) {
+    auto input_type = u.scalar_type();
+    auto weight_type = A.scalar_type();
+    TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
+    TORCH_CHECK(weight_type == at::ScalarType::Float);
+
+    TORCH_CHECK(delta.scalar_type() == input_type);
+    TORCH_CHECK(B.scalar_type() == input_type);
+    TORCH_CHECK(C.scalar_type() == input_type);
+    TORCH_CHECK(dout.scalar_type() == input_type);
+
+    TORCH_CHECK(u.is_cuda());
+    TORCH_CHECK(delta.is_cuda());
+    TORCH_CHECK(A.is_cuda());
+    TORCH_CHECK(B.is_cuda());
+    TORCH_CHECK(C.is_cuda());
+    TORCH_CHECK(dout.is_cuda());
+
+    TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
+    TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
+    TORCH_CHECK(dout.stride(-1) == 1 || dout.size(-1) == 1);
+
+    const auto sizes = u.sizes();
+    const int batch_size = sizes[0];
+    const int dim = sizes[1];
+    const int seqlen = sizes[2];
+    const int dstate = A.size(1);
+    const int n_groups = B.size(1);
+
+    TORCH_CHECK(dim % (n_groups * nrows) == 0, "dims should be dividable by n_groups * nrows");
+    TORCH_CHECK(dstate <= MAX_DSTATE / nrows, "selective_scan only supports state dimension <= 256 / nrows");
+    
+    CHECK_SHAPE(u, batch_size, dim, seqlen);
+    CHECK_SHAPE(delta, batch_size, dim, seqlen);
+    CHECK_SHAPE(A, dim, dstate);
+    CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen);
+    TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
+    CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen);
+    TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
+    CHECK_SHAPE(dout, batch_size, dim, seqlen);
+
+    if (D_.has_value()) {
+        auto D = D_.value();
+        TORCH_CHECK(D.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(D.is_cuda());
+        TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
+        CHECK_SHAPE(D, dim);
+    }
+
+    if (delta_bias_.has_value()) {
+        auto delta_bias = delta_bias_.value();
+        TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(delta_bias.is_cuda());
+        TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
+        CHECK_SHAPE(delta_bias, dim);
+    }
+
+    at::Tensor out;
+    const int n_chunks = (seqlen + 2048 - 1) / 2048;
+    // const int n_chunks = (seqlen + 1024 - 1) / 1024;
+    if (n_chunks > 1) { TORCH_CHECK(x_.has_value()); }
+    if (x_.has_value()) {
+        auto x = x_.value();
+        TORCH_CHECK(x.scalar_type() == weight_type);
+        TORCH_CHECK(x.is_cuda());
+        TORCH_CHECK(x.is_contiguous());
+        CHECK_SHAPE(x, batch_size, dim, n_chunks, 2 * dstate);
+    }
+
+    at::Tensor du = torch::empty_like(u);
+    at::Tensor ddelta = torch::empty_like(delta);
+    at::Tensor dA = torch::zeros_like(A);
+    at::Tensor dB = torch::zeros_like(B, B.options().dtype(torch::kFloat32));
+    at::Tensor dC = torch::zeros_like(C, C.options().dtype(torch::kFloat32));
+    at::Tensor dD;
+    if (D_.has_value()) { dD = torch::zeros_like(D_.value()); }
+    at::Tensor ddelta_bias;
+    if (delta_bias_.has_value()) { ddelta_bias = torch::zeros_like(delta_bias_.value()); }
+
+    SSMParamsBwd params;
+    set_ssm_params_bwd(params, batch_size, dim, seqlen, dstate, n_groups, n_chunks,
+                       u, delta, A, B, C, out,
+                       D_.has_value() ? D_.value().data_ptr() : nullptr,
+                       delta_bias_.has_value() ? delta_bias_.value().data_ptr() : nullptr,
+                       x_.has_value() ? x_.value().data_ptr() : nullptr,
+                       dout, du, ddelta, dA, dB, dC,
+                       D_.has_value() ? dD.data_ptr() : nullptr,
+                       delta_bias_.has_value() ? ddelta_bias.data_ptr() : nullptr,
+                       delta_softplus);
+
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)u.get_device()};
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), "selective_scan_bwd", [&] {
+        // constexpr int kNRows = 1;
+        INT_SWITCH(nrows, kNRows, [&] {
+            selective_scan_bwd_cuda<kNRows, input_t, weight_t>(params, stream);
+        });
+    });
+    std::vector<at::Tensor> result = {du, ddelta, dA, dB.to(B.dtype()), dC.to(C.dtype()), dD, ddelta_bias};
+    return result;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("fwd", &selective_scan_fwd, "Selective scan forward");
+    m.def("bwd", &selective_scan_bwd, "Selective scan backward");
+}

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusoflex/selective_scan_bwd_kernel_oflex.cuh

@@ -0,0 +1,323 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <c10/util/BFloat16.h>
+#include <c10/util/Half.h>
+#include <c10/cuda/CUDAException.h>  // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
+#include <ATen/cuda/Atomic.cuh>  // For atomicAdd on complex
+
+#include <cub/block/block_load.cuh>
+#include <cub/block/block_store.cuh>
+#include <cub/block/block_scan.cuh>
+#include <cub/block/block_reduce.cuh>
+
+#include "selective_scan.h"
+#include "selective_scan_common.h"
+#include "reverse_scan.cuh"
+#include "static_switch.h"
+
+template<int kNThreads_, int kNItems_, bool kIsEvenLen_, bool kDeltaSoftplus_, typename input_t_, typename weight_t_, typename output_t_>
+struct Selective_Scan_bwd_kernel_traits {
+    static_assert(kNItems_ % 4 == 0);
+    using input_t = input_t_;
+    using weight_t = weight_t_;
+    using output_t = output_t_;
+
+    static constexpr int kNThreads = kNThreads_;
+    static constexpr int kNItems = kNItems_;
+    static constexpr int MaxDState = MAX_DSTATE;
+    static constexpr int kNBytes = sizeof(input_t);
+    static_assert(kNBytes == 2 || kNBytes == 4);
+    static constexpr int kNElts = kNBytes == 4 ? 4 : std::min(8, kNItems);
+    static_assert(kNItems % kNElts == 0);
+    static constexpr int kNLoads = kNItems / kNElts;
+    static constexpr bool kIsEvenLen = kIsEvenLen_;
+    static constexpr bool kDeltaSoftplus = kDeltaSoftplus_;
+    // Setting MinBlocksPerMP to be 3 (instead of 2) for 128 threads with float improves occupancy.
+    // For complex this would lead to massive register spilling, so we keep it at 2.
+    static constexpr int kMinBlocks = kNThreads == 128 && 3;
+    static constexpr int kNLoadsOutput = sizeof(output_t) * kNLoads / kNBytes;
+    using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
+    using scan_t = float2;
+    using BlockLoadT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadWeightT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadWeightVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadOutputT = cub::BlockLoad<output_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadOutputVecT = cub::BlockLoad<vec_t, kNThreads, kNLoadsOutput, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockStoreT = cub::BlockStore<input_t, kNThreads, kNItems, cub::BLOCK_STORE_WARP_TRANSPOSE>;
+    using BlockStoreVecT = cub::BlockStore<vec_t, kNThreads, kNLoads, cub::BLOCK_STORE_WARP_TRANSPOSE>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING_MEMOIZE>;
+    using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_WARP_SCANS>;
+    using BlockReverseScanT = BlockReverseScan<scan_t, kNThreads>;
+    using BlockReduceT = cub::BlockReduce<scan_t, kNThreads>;
+    using BlockReduceFloatT = cub::BlockReduce<float, kNThreads>;
+    using BlockExchangeT = cub::BlockExchange<float, kNThreads, kNItems>;
+    static constexpr int kSmemIOSize = std::max({sizeof(typename BlockLoadT::TempStorage),
+                                                 sizeof(typename BlockLoadVecT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightVecT::TempStorage),
+                                                 sizeof(typename BlockLoadOutputT::TempStorage),
+                                                 sizeof(typename BlockLoadOutputVecT::TempStorage),
+                                                 sizeof(typename BlockStoreT::TempStorage),
+                                                 sizeof(typename BlockStoreVecT::TempStorage)});
+    static constexpr int kSmemExchangeSize = 2 * sizeof(typename BlockExchangeT::TempStorage);
+    static constexpr int kSmemReduceSize = sizeof(typename BlockReduceT::TempStorage);
+    static constexpr int kSmemSize = kSmemIOSize + kSmemExchangeSize + kSmemReduceSize + sizeof(typename BlockScanT::TempStorage) + sizeof(typename BlockReverseScanT::TempStorage);
+};
+
+template<typename Ktraits>
+__global__ __launch_bounds__(Ktraits::kNThreads, Ktraits::kMinBlocks)
+void selective_scan_bwd_kernel(SSMParamsBwd params) {
+    constexpr bool kDeltaSoftplus = Ktraits::kDeltaSoftplus;
+    constexpr int kNThreads = Ktraits::kNThreads;
+    constexpr int kNItems = Ktraits::kNItems;
+    using input_t = typename Ktraits::input_t;
+    using weight_t = typename Ktraits::weight_t;
+    using output_t = typename Ktraits::output_t;
+    using scan_t = typename Ktraits::scan_t;
+
+    // Shared memory.
+    extern __shared__ char smem_[];
+    auto& smem_load = reinterpret_cast<typename Ktraits::BlockLoadT::TempStorage&>(smem_);
+    auto& smem_load1 = reinterpret_cast<typename Ktraits::BlockLoadOutputT::TempStorage&>(smem_);
+    auto& smem_load_weight = reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage&>(smem_);
+    auto& smem_load_weight1 = *reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage*>(smem_ + sizeof(typename Ktraits::BlockLoadWeightT::TempStorage));
+    auto& smem_store = reinterpret_cast<typename Ktraits::BlockStoreT::TempStorage&>(smem_);
+    auto& smem_exchange = *reinterpret_cast<typename Ktraits::BlockExchangeT::TempStorage*>(smem_ + Ktraits::kSmemIOSize);
+    auto& smem_exchange1 = *reinterpret_cast<typename Ktraits::BlockExchangeT::TempStorage*>(smem_ + Ktraits::kSmemIOSize + sizeof(typename Ktraits::BlockExchangeT::TempStorage));
+    auto& smem_reduce = *reinterpret_cast<typename Ktraits::BlockReduceT::TempStorage*>(reinterpret_cast<char *>(&smem_exchange) + Ktraits::kSmemExchangeSize);
+    auto& smem_reduce_float = *reinterpret_cast<typename Ktraits::BlockReduceFloatT::TempStorage*>(&smem_reduce);
+    auto& smem_scan = *reinterpret_cast<typename Ktraits::BlockScanT::TempStorage*>(reinterpret_cast<char *>(&smem_reduce) + Ktraits::kSmemReduceSize);
+    auto& smem_reverse_scan = *reinterpret_cast<typename Ktraits::BlockReverseScanT::TempStorage*>(reinterpret_cast<char *>(&smem_scan) + sizeof(typename Ktraits::BlockScanT::TempStorage));
+    weight_t *smem_delta_a = reinterpret_cast<weight_t *>(smem_ + Ktraits::kSmemSize);
+    scan_t *smem_running_postfix = reinterpret_cast<scan_t *>(smem_delta_a + 2 * Ktraits::MaxDState + kNThreads);
+    weight_t *smem_da = reinterpret_cast<weight_t *>(smem_running_postfix + Ktraits::MaxDState);
+
+    const int batch_id = blockIdx.x;
+    const int dim_id = blockIdx.y;
+    const int group_id = dim_id / (params.dim_ngroups_ratio);
+    input_t *u = reinterpret_cast<input_t *>(params.u_ptr) + batch_id * params.u_batch_stride
+        + dim_id * params.u_d_stride;
+    input_t *delta = reinterpret_cast<input_t *>(params.delta_ptr) + batch_id * params.delta_batch_stride
+        + dim_id * params.delta_d_stride;
+
+    weight_t *A = reinterpret_cast<weight_t *>(params.A_ptr) + dim_id * params.A_d_stride;
+    input_t *Bvar = reinterpret_cast<input_t *>(params.B_ptr) + batch_id * params.B_batch_stride + group_id * params.B_group_stride;
+    input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + batch_id * params.C_batch_stride + group_id * params.C_group_stride;
+    weight_t *dA = reinterpret_cast<weight_t *>(params.dA_ptr) + dim_id * params.dA_d_stride;
+    weight_t *dB = reinterpret_cast<weight_t *>(params.dB_ptr)
+        + (batch_id * params.dB_batch_stride + group_id * params.dB_group_stride);
+    weight_t *dC = reinterpret_cast<weight_t *>(params.dC_ptr)
+        + (batch_id * params.dC_batch_stride + group_id * params.dC_group_stride);
+    float *dD = params.dD_ptr == nullptr ? nullptr : reinterpret_cast<float *>(params.dD_ptr) + dim_id;
+    float D_val = params.D_ptr == nullptr ? 0 : reinterpret_cast<float *>(params.D_ptr)[dim_id];
+    float *ddelta_bias = params.ddelta_bias_ptr == nullptr ? nullptr : reinterpret_cast<float *>(params.ddelta_bias_ptr) + dim_id;
+    float delta_bias = params.delta_bias_ptr == nullptr ? 0 : reinterpret_cast<float *>(params.delta_bias_ptr)[dim_id];
+    scan_t *x = params.x_ptr == nullptr
+        ? nullptr
+        : reinterpret_cast<scan_t *>(params.x_ptr) + (batch_id * params.dim + dim_id) * (params.n_chunks) * params.dstate;
+    float dD_val = 0;
+    float ddelta_bias_val = 0;
+
+    output_t *dout = reinterpret_cast<output_t *>(params.dout_ptr) + batch_id * params.dout_batch_stride + dim_id * params.dout_d_stride;
+
+    constexpr int kChunkSize = kNThreads * kNItems;
+    u += (params.n_chunks - 1) * kChunkSize;
+    delta += (params.n_chunks - 1) * kChunkSize;
+    dout += (params.n_chunks - 1) * kChunkSize;
+    Bvar += (params.n_chunks - 1) * kChunkSize;
+    Cvar += (params.n_chunks - 1) * kChunkSize;
+    for (int chunk = params.n_chunks - 1; chunk >= 0; --chunk) {
+        input_t u_vals[kNItems];
+        input_t delta_vals_load[kNItems];
+        float dout_vals[kNItems];
+        __syncthreads();
+        load_input<Ktraits>(u, u_vals, smem_load, params.seqlen - chunk * kChunkSize);
+        __syncthreads();
+        load_input<Ktraits>(delta, delta_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
+        __syncthreads();
+        if constexpr (std::is_same_v<output_t, input_t>) {
+            input_t dout_vals_load[kNItems];
+            load_input<Ktraits>(reinterpret_cast<input_t *>(dout), dout_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
+            Converter<typename Ktraits::input_t, kNItems>::to_float(dout_vals_load, dout_vals);
+        } else {
+            static_assert(std::is_same_v<output_t, float>);
+            load_output<Ktraits>(dout, dout_vals, smem_load1, params.seqlen - chunk * kChunkSize);
+        }
+        u -= kChunkSize;
+        // Will reload delta at the same location if kDeltaSoftplus
+        if constexpr (!kDeltaSoftplus) { delta -= kChunkSize; }
+        dout -= kChunkSize;
+
+        float delta_vals[kNItems];
+        float du_vals[kNItems];
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) {
+            delta_vals[i] = float(delta_vals_load[i]) + delta_bias;
+            if constexpr (kDeltaSoftplus) {
+                delta_vals[i] = delta_vals[i] <= 20.f ? log1pf(expf(delta_vals[i])) : delta_vals[i];
+            }
+        }
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) { du_vals[i] = D_val * dout_vals[i]; }
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) { dD_val += dout_vals[i] * float(u_vals[i]); }
+
+        float ddelta_vals[kNItems] = {0};
+        __syncthreads();
+        for (int state_idx = 0; state_idx < params.dstate; ++state_idx) {
+            constexpr float kLog2e = M_LOG2E;
+            weight_t A_val = A[state_idx * params.A_dstate_stride];
+            weight_t A_scaled = A_val * kLog2e;
+            weight_t B_vals[kNItems], C_vals[kNItems];
+            load_weight<Ktraits>(Bvar + state_idx * params.B_dstate_stride, B_vals,
+                    smem_load_weight, (params.seqlen - chunk * kChunkSize));
+            auto &smem_load_weight_C = smem_load_weight1;
+            load_weight<Ktraits>(Cvar + state_idx * params.C_dstate_stride, C_vals,
+                    smem_load_weight_C, (params.seqlen - chunk * kChunkSize));
+            scan_t thread_data[kNItems], thread_reverse_data[kNItems];
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                const float delta_a_exp = exp2f(delta_vals[i] * A_scaled);
+                thread_data[i] = make_float2(delta_a_exp, delta_vals[i] * float(u_vals[i]) * B_vals[i]);
+                if (i == 0) {
+                    smem_delta_a[threadIdx.x == 0 ? state_idx + (chunk % 2) * Ktraits::MaxDState: threadIdx.x + 2 * Ktraits::MaxDState] = delta_a_exp;
+                } else {
+                    thread_reverse_data[i - 1].x = delta_a_exp;
+                }
+                thread_reverse_data[i].y = dout_vals[i] * C_vals[i];
+            }
+            __syncthreads();
+            thread_reverse_data[kNItems - 1].x = threadIdx.x == kNThreads - 1
+                ? (chunk == params.n_chunks - 1 ? 1.f : smem_delta_a[state_idx + ((chunk + 1) % 2) * Ktraits::MaxDState])
+                    : smem_delta_a[threadIdx.x + 1 + 2 * Ktraits::MaxDState];
+            // Initialize running total
+            scan_t running_prefix = chunk > 0 && threadIdx.x % 32 == 0 ? x[(chunk - 1) * params.dstate + state_idx] : make_float2(1.f, 0.f);
+            SSMScanPrefixCallbackOp<weight_t> prefix_op(running_prefix);
+            Ktraits::BlockScanT(smem_scan).InclusiveScan(
+                thread_data, thread_data, SSMScanOp<weight_t>(), prefix_op
+            );
+            scan_t running_postfix = chunk < params.n_chunks - 1 && threadIdx.x % 32 == 0 ? smem_running_postfix[state_idx] : make_float2(1.f, 0.f);
+            SSMScanPrefixCallbackOp<weight_t> postfix_op(running_postfix);
+            Ktraits::BlockReverseScanT(smem_reverse_scan).InclusiveReverseScan(
+                thread_reverse_data, thread_reverse_data, SSMScanOp<weight_t>(), postfix_op
+            );
+            if (threadIdx.x == 0) { smem_running_postfix[state_idx] = postfix_op.running_prefix; }
+            weight_t dA_val = 0;
+            weight_t dB_vals[kNItems], dC_vals[kNItems];
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                const float dx = thread_reverse_data[i].y;
+                const float ddelta_u = dx * B_vals[i];
+                du_vals[i] += ddelta_u * delta_vals[i];
+                const float a = thread_data[i].y - (delta_vals[i] * float(u_vals[i]) * B_vals[i]);
+                ddelta_vals[i] += ddelta_u * float(u_vals[i]) + dx * A_val * a;
+                dA_val += dx * delta_vals[i] * a;
+                dB_vals[i] = dx * delta_vals[i] * float(u_vals[i]);
+                dC_vals[i] = dout_vals[i] * thread_data[i].y;
+            }
+            // Block-exchange to make the atomicAdd's coalesced, otherwise they're much slower
+            Ktraits::BlockExchangeT(smem_exchange).BlockedToStriped(dB_vals, dB_vals);
+            auto &smem_exchange_C = smem_exchange1;
+            Ktraits::BlockExchangeT(smem_exchange_C).BlockedToStriped(dC_vals, dC_vals);
+            const int seqlen_remaining = params.seqlen - chunk * kChunkSize - threadIdx.x;
+            weight_t *dB_cur = dB + state_idx * params.dB_dstate_stride + chunk * kChunkSize + threadIdx.x;
+            weight_t *dC_cur = dC + state_idx * params.dC_dstate_stride + chunk * kChunkSize + threadIdx.x;
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                if (i * kNThreads < seqlen_remaining) {
+                    { gpuAtomicAdd(dB_cur + i * kNThreads, dB_vals[i]); }
+                    { gpuAtomicAdd(dC_cur + i * kNThreads, dC_vals[i]); }
+                }
+            }
+            dA_val = Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(dA_val);
+            if (threadIdx.x == 0) {
+                smem_da[state_idx] = chunk == params.n_chunks - 1 ? dA_val : dA_val + smem_da[state_idx];
+            }
+        }
+
+        if constexpr (kDeltaSoftplus) {
+            input_t delta_vals_load[kNItems];
+            __syncthreads();
+            load_input<Ktraits>(delta, delta_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
+            delta -= kChunkSize;
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                float delta_val = float(delta_vals_load[i]) + delta_bias;
+                float delta_val_neg_exp = expf(-delta_val);
+                ddelta_vals[i] = delta_val <= 20.f
+                    ? ddelta_vals[i] / (1.f + delta_val_neg_exp)
+                        : ddelta_vals[i];
+            }
+        }
+        
+        __syncthreads();
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) { ddelta_bias_val += ddelta_vals[i]; }
+
+        input_t *du = reinterpret_cast<input_t *>(params.du_ptr) + batch_id * params.du_batch_stride
+            + dim_id * params.du_d_stride + chunk * kChunkSize;
+        input_t *ddelta = reinterpret_cast<input_t *>(params.ddelta_ptr) + batch_id * params.ddelta_batch_stride
+            + dim_id * params.ddelta_d_stride + chunk * kChunkSize;
+        __syncthreads();
+        store_output<Ktraits>(du, du_vals, smem_store, params.seqlen - chunk * kChunkSize);
+        __syncthreads();
+        store_output<Ktraits>(ddelta, ddelta_vals, smem_store, params.seqlen - chunk * kChunkSize);
+        Bvar -= kChunkSize;
+        Cvar -= kChunkSize;
+    }
+
+    if (params.dD_ptr != nullptr) {
+        __syncthreads();
+        dD_val = Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(dD_val);
+        if (threadIdx.x == 0) { gpuAtomicAdd(dD, dD_val); }
+    }
+    if (params.ddelta_bias_ptr != nullptr) {
+        __syncthreads();
+        ddelta_bias_val = Ktraits::BlockReduceFloatT(smem_reduce_float).Sum(ddelta_bias_val);
+        if (threadIdx.x == 0) { gpuAtomicAdd(ddelta_bias, ddelta_bias_val); }
+    }
+    __syncthreads();
+    for (int state_idx = threadIdx.x; state_idx < params.dstate; state_idx += blockDim.x) {
+        gpuAtomicAdd(&(dA[state_idx * params.dA_dstate_stride]), smem_da[state_idx]);
+    }
+}
+
+template<int kNThreads, int kNItems, typename input_t, typename weight_t, typename output_t>
+void selective_scan_bwd_launch(SSMParamsBwd &params, cudaStream_t stream) {
+    BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {
+        BOOL_SWITCH(params.delta_softplus, kDeltaSoftplus, [&] {
+            using Ktraits = Selective_Scan_bwd_kernel_traits<kNThreads, kNItems, kIsEvenLen, kDeltaSoftplus, input_t, weight_t, output_t>;
+            constexpr int kSmemSize = Ktraits::kSmemSize + Ktraits::MaxDState * sizeof(typename Ktraits::scan_t) + (kNThreads + 4 * Ktraits::MaxDState) * sizeof(typename Ktraits::weight_t);
+            // printf("smem_size = %d\n", kSmemSize);
+            dim3 grid(params.batch, params.dim);
+            auto kernel = &selective_scan_bwd_kernel<Ktraits>;
+            if (kSmemSize >= 48 * 1024) {
+                C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
+            }
+            kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
+            C10_CUDA_KERNEL_LAUNCH_CHECK();
+        });
+    });
+}
+
+template<int knrows, typename input_t, typename weight_t, typename output_t>
+void selective_scan_bwd_cuda(SSMParamsBwd &params, cudaStream_t stream) {
+    if (params.seqlen <= 128) {
+        selective_scan_bwd_launch<32, 4, input_t, weight_t, output_t>(params, stream);
+    } else if (params.seqlen <= 256) {
+        selective_scan_bwd_launch<32, 8, input_t, weight_t, output_t>(params, stream);
+    } else if (params.seqlen <= 512) {
+        selective_scan_bwd_launch<32, 16, input_t, weight_t, output_t>(params, stream);
+    } else if (params.seqlen <= 1024) {
+        selective_scan_bwd_launch<64, 16, input_t, weight_t, output_t>(params, stream);
+    } else {
+        selective_scan_bwd_launch<128, 16, input_t, weight_t, output_t>(params, stream);
+    }
+}
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusoflex/selective_scan_core_bwd.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+#include "selective_scan_bwd_kernel_oflex.cuh"
+
+template void selective_scan_bwd_cuda<1, float, float, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<1, at::Half, float, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<1, at::BFloat16, float, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<1, at::Half, float, at::Half>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<1, at::BFloat16, float, at::BFloat16>(SSMParamsBwd &params, cudaStream_t stream);
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusoflex/selective_scan_core_fwd.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+#include "selective_scan_fwd_kernel_oflex.cuh"
+
+template void selective_scan_fwd_cuda<1, float, float, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<1, at::Half, float, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<1, at::BFloat16, float, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<1, at::Half, float, at::Half>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<1, at::BFloat16, float, at::BFloat16>(SSMParamsBase &params, cudaStream_t stream);
+

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusoflex/selective_scan_fwd_kernel_oflex.cuh

@@ -0,0 +1,211 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <c10/util/BFloat16.h>
+#include <c10/util/Half.h>
+#include <c10/cuda/CUDAException.h>  // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
+
+#include <cub/block/block_load.cuh>
+#include <cub/block/block_store.cuh>
+#include <cub/block/block_scan.cuh>
+
+#include "selective_scan.h"
+#include "selective_scan_common.h"
+#include "static_switch.h"
+
+template<int kNThreads_, int kNItems_, bool kIsEvenLen_, typename input_t_, typename weight_t_, typename output_t_>
+struct Selective_Scan_fwd_kernel_traits {
+    static_assert(kNItems_ % 4 == 0);
+    using input_t = input_t_;
+    using weight_t = weight_t_;
+    using output_t = output_t_;
+    static constexpr int kNThreads = kNThreads_;
+    // Setting MinBlocksPerMP to be 3 (instead of 2) for 128 threads improves occupancy.
+    static constexpr int kMinBlocks = kNThreads < 128 ? 5 : 3;
+    static constexpr int kNItems = kNItems_;
+    static constexpr int MaxDState = MAX_DSTATE;
+    static constexpr int kNBytes = sizeof(input_t);
+    static_assert(kNBytes == 2 || kNBytes == 4);
+    static constexpr int kNElts = kNBytes == 4 ? 4 : std::min(8, kNItems);
+    static_assert(kNItems % kNElts == 0);
+    static constexpr int kNLoads = kNItems / kNElts;
+    static constexpr bool kIsEvenLen = kIsEvenLen_;
+    static constexpr bool kDirectIO = kIsEvenLen && kNLoads == 1;
+    static constexpr int kNLoadsOutput = sizeof(output_t) * kNLoads / kNBytes;
+    static constexpr bool kDirectIOOutput = kDirectIO && (kNLoadsOutput == 1);
+    using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
+    using scan_t = float2;
+    using BlockLoadT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads,
+        !kDirectIO ? cub::BLOCK_LOAD_WARP_TRANSPOSE : cub::BLOCK_LOAD_DIRECT>;
+    using BlockLoadWeightT = cub::BlockLoad<input_t, kNThreads, kNItems, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
+    using BlockLoadWeightVecT = cub::BlockLoad<vec_t, kNThreads, kNLoads,
+        !kDirectIO ? cub::BLOCK_LOAD_WARP_TRANSPOSE  : cub::BLOCK_LOAD_DIRECT>;
+    using BlockStoreT = cub::BlockStore<input_t, kNThreads, kNItems, cub::BLOCK_STORE_WARP_TRANSPOSE>;
+    using BlockStoreVecT = cub::BlockStore<vec_t, kNThreads, kNLoads,
+        !kDirectIO ? cub::BLOCK_STORE_WARP_TRANSPOSE : cub::BLOCK_STORE_DIRECT>;
+    using BlockStoreOutputT = cub::BlockStore<output_t, kNThreads, kNItems, cub::BLOCK_STORE_WARP_TRANSPOSE>;
+    using BlockStoreOutputVecT = cub::BlockStore<vec_t, kNThreads, kNLoadsOutput,
+        !kDirectIOOutput ? cub::BLOCK_STORE_WARP_TRANSPOSE  : cub::BLOCK_STORE_DIRECT>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING_MEMOIZE>;
+    // using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_RAKING>;
+    using BlockScanT = cub::BlockScan<scan_t, kNThreads, cub::BLOCK_SCAN_WARP_SCANS>;
+    static constexpr int kSmemIOSize = std::max({sizeof(typename BlockLoadT::TempStorage),
+                                                 sizeof(typename BlockLoadVecT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightT::TempStorage),
+                                                 2 * sizeof(typename BlockLoadWeightVecT::TempStorage),
+                                                 sizeof(typename BlockStoreT::TempStorage),
+                                                 sizeof(typename BlockStoreVecT::TempStorage),
+                                                 sizeof(typename BlockStoreOutputT::TempStorage),
+                                                 sizeof(typename BlockStoreOutputVecT::TempStorage)});
+    static constexpr int kSmemSize = kSmemIOSize + sizeof(typename BlockScanT::TempStorage);
+};
+
+template<typename Ktraits>
+__global__ __launch_bounds__(Ktraits::kNThreads, Ktraits::kMinBlocks)
+void selective_scan_fwd_kernel(SSMParamsBase params) {
+    constexpr int kNThreads = Ktraits::kNThreads;
+    constexpr int kNItems = Ktraits::kNItems;
+    constexpr bool kDirectIO = Ktraits::kDirectIO;
+    using input_t = typename Ktraits::input_t;
+    using weight_t = typename Ktraits::weight_t;
+    using output_t = typename Ktraits::output_t;
+    using scan_t = typename Ktraits::scan_t;
+
+    // Shared memory.
+    extern __shared__ char smem_[];
+    auto& smem_load = reinterpret_cast<typename Ktraits::BlockLoadT::TempStorage&>(smem_);
+    auto& smem_load_weight = reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage&>(smem_);
+    auto& smem_load_weight1 = *reinterpret_cast<typename Ktraits::BlockLoadWeightT::TempStorage*>(smem_ + sizeof(typename Ktraits::BlockLoadWeightT::TempStorage));
+    auto& smem_store = reinterpret_cast<typename Ktraits::BlockStoreT::TempStorage&>(smem_);
+    auto& smem_store1 = reinterpret_cast<typename Ktraits::BlockStoreOutputT::TempStorage&>(smem_);
+    auto& smem_scan = *reinterpret_cast<typename Ktraits::BlockScanT::TempStorage*>(smem_ + Ktraits::kSmemIOSize);
+    scan_t *smem_running_prefix = reinterpret_cast<scan_t *>(smem_ + Ktraits::kSmemSize);
+
+    const int batch_id = blockIdx.x;
+    const int dim_id = blockIdx.y;
+    const int group_id = dim_id / (params.dim_ngroups_ratio);
+    input_t *u = reinterpret_cast<input_t *>(params.u_ptr) + batch_id * params.u_batch_stride
+        + dim_id * params.u_d_stride;
+    input_t *delta = reinterpret_cast<input_t *>(params.delta_ptr) + batch_id * params.delta_batch_stride
+        + dim_id * params.delta_d_stride;
+    weight_t *A = reinterpret_cast<weight_t *>(params.A_ptr) + dim_id * params.A_d_stride;
+    input_t *Bvar = reinterpret_cast<input_t *>(params.B_ptr) + batch_id * params.B_batch_stride + group_id * params.B_group_stride;
+    input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + batch_id * params.C_batch_stride + group_id * params.C_group_stride;
+    scan_t *x = reinterpret_cast<scan_t *>(params.x_ptr) + (batch_id * params.dim + dim_id) * params.n_chunks * params.dstate;
+
+    float D_val = 0; // attention!
+    if (params.D_ptr != nullptr) {
+        D_val = reinterpret_cast<float *>(params.D_ptr)[dim_id];
+    }
+    float delta_bias = 0;
+    if (params.delta_bias_ptr != nullptr) {
+        delta_bias = reinterpret_cast<float *>(params.delta_bias_ptr)[dim_id];
+    }
+
+    constexpr int kChunkSize = kNThreads * kNItems;
+    for (int chunk = 0; chunk < params.n_chunks; ++chunk) {
+        input_t u_vals[kNItems], delta_vals_load[kNItems];
+        __syncthreads();
+        load_input<Ktraits>(u, u_vals, smem_load, params.seqlen - chunk * kChunkSize);
+        if constexpr (!kDirectIO) { __syncthreads(); }
+        load_input<Ktraits>(delta, delta_vals_load, smem_load, params.seqlen - chunk * kChunkSize);
+        u += kChunkSize;
+        delta += kChunkSize;
+
+        float delta_vals[kNItems], delta_u_vals[kNItems], out_vals[kNItems];
+        #pragma unroll
+        for (int i = 0; i < kNItems; ++i) {
+            float u_val = float(u_vals[i]);
+            delta_vals[i] = float(delta_vals_load[i]) + delta_bias;
+            if (params.delta_softplus) {
+                delta_vals[i] = delta_vals[i] <= 20.f ? log1pf(expf(delta_vals[i])) : delta_vals[i];
+            }
+            delta_u_vals[i] = delta_vals[i] * u_val;
+            out_vals[i] = D_val * u_val;
+        }
+
+        __syncthreads();
+        for (int state_idx = 0; state_idx < params.dstate; ++state_idx) {
+            constexpr float kLog2e = M_LOG2E;
+            weight_t A_val = A[state_idx * params.A_dstate_stride];
+            A_val *= kLog2e;
+            weight_t B_vals[kNItems], C_vals[kNItems];
+            load_weight<Ktraits>(Bvar + state_idx * params.B_dstate_stride, B_vals,
+                    smem_load_weight, (params.seqlen - chunk * kChunkSize));
+            load_weight<Ktraits>(Cvar + state_idx * params.C_dstate_stride, C_vals,
+                    smem_load_weight1, (params.seqlen - chunk * kChunkSize));
+            __syncthreads();
+            scan_t thread_data[kNItems];
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                thread_data[i] = make_float2(exp2f(delta_vals[i] * A_val), B_vals[i] * delta_u_vals[i]);
+                if constexpr (!Ktraits::kIsEvenLen) {  // So that the last state is correct
+                    if (threadIdx.x * kNItems + i >= params.seqlen - chunk * kChunkSize) {
+                        thread_data[i] = make_float2(1.f, 0.f);
+                    }
+                }
+            }
+            // Initialize running total
+            scan_t running_prefix;
+            // If we use WARP_SCAN then all lane 0 of all warps (not just thread 0) needs to read
+            running_prefix = chunk > 0 && threadIdx.x % 32 == 0 ? smem_running_prefix[state_idx] : make_float2(1.f, 0.f);
+            // running_prefix = chunk > 0 && threadIdx.x == 0 ? smem_running_prefix[state_idx] : make_float2(1.f, 0.f);
+            SSMScanPrefixCallbackOp<weight_t> prefix_op(running_prefix);
+            Ktraits::BlockScanT(smem_scan).InclusiveScan(
+                thread_data, thread_data, SSMScanOp<weight_t>(), prefix_op
+            );
+            // There's a syncthreads in the scan op, so we don't need to sync here.
+            // Unless there's only 1 warp, but then it's the same thread (0) reading and writing.
+            if (threadIdx.x == 0) {
+                smem_running_prefix[state_idx] = prefix_op.running_prefix;
+                x[chunk * params.dstate + state_idx] = prefix_op.running_prefix;
+            }
+            #pragma unroll
+            for (int i = 0; i < kNItems; ++i) {
+                out_vals[i] += thread_data[i].y * C_vals[i];
+            }
+        }
+
+        output_t *out = reinterpret_cast<output_t *>(params.out_ptr) + batch_id * params.out_batch_stride
+            + dim_id * params.out_d_stride + chunk * kChunkSize;
+        __syncthreads();
+        store_output1<Ktraits>(out, out_vals, smem_store1, params.seqlen - chunk * kChunkSize);
+        Bvar += kChunkSize;
+        Cvar += kChunkSize;
+    }
+}
+
+template<int kNThreads, int kNItems, typename input_t, typename weight_t, typename output_t>
+void selective_scan_fwd_launch(SSMParamsBase &params, cudaStream_t stream) {
+    BOOL_SWITCH(params.seqlen % (kNThreads * kNItems) == 0, kIsEvenLen, [&] {
+        using Ktraits = Selective_Scan_fwd_kernel_traits<kNThreads, kNItems, kIsEvenLen, input_t, weight_t, output_t>;
+        constexpr int kSmemSize = Ktraits::kSmemSize + Ktraits::MaxDState * sizeof(typename Ktraits::scan_t);
+        // printf("smem_size = %d\n", kSmemSize);
+        dim3 grid(params.batch, params.dim);
+        auto kernel = &selective_scan_fwd_kernel<Ktraits>;
+        if (kSmemSize >= 48 * 1024) {
+            C10_CUDA_CHECK(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
+        }
+        kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
+        C10_CUDA_KERNEL_LAUNCH_CHECK();
+    });
+}
+
+template<int knrows, typename input_t, typename weight_t, typename output_t>
+void selective_scan_fwd_cuda(SSMParamsBase &params, cudaStream_t stream) {
+    if (params.seqlen <= 128) {
+        selective_scan_fwd_launch<32, 4, input_t, weight_t, output_t>(params, stream);
+    } else if (params.seqlen <= 256) {
+        selective_scan_fwd_launch<32, 8, input_t, weight_t, output_t>(params, stream);
+    } else if (params.seqlen <= 512) {
+        selective_scan_fwd_launch<32, 16, input_t, weight_t, output_t>(params, stream);
+    } else if (params.seqlen <= 1024) {
+        selective_scan_fwd_launch<64, 16, input_t, weight_t, output_t>(params, stream);
+    } else {
+        selective_scan_fwd_launch<128, 16, input_t, weight_t, output_t>(params, stream);
+    }
+}

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/cusoflex/selective_scan_oflex.cpp

@@ -0,0 +1,363 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <torch/extension.h>
+#include <vector>
+
+#include "selective_scan.h"
+#define MAX_DSTATE 256
+
+#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
+using weight_t = float;
+
+#define DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, NAME, ...)                    \
+    if (ITYPE == at::ScalarType::Half) {                                            \
+        using input_t = at::Half;                                                   \
+        __VA_ARGS__();                                                              \
+    } else if (ITYPE == at::ScalarType::BFloat16) {                                 \
+        using input_t = at::BFloat16;                                               \
+        __VA_ARGS__();                                                              \
+    } else if (ITYPE == at::ScalarType::Float)  {                                   \
+        using input_t = float;                                                      \
+        __VA_ARGS__();                                                              \
+    } else {                                                                        \
+        AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \
+    }
+
+template<int knrows, typename input_t, typename weight_t, typename output_t>
+void selective_scan_fwd_cuda(SSMParamsBase &params, cudaStream_t stream);
+
+template <int knrows, typename input_t, typename weight_t, typename output_t>
+void selective_scan_bwd_cuda(SSMParamsBwd &params, cudaStream_t stream);
+
+void set_ssm_params_fwd(SSMParamsBase &params,
+                        // sizes
+                        const size_t batch,
+                        const size_t dim,
+                        const size_t seqlen,
+                        const size_t dstate,
+                        const size_t n_groups,
+                        const size_t n_chunks,
+                        // device pointers
+                        const at::Tensor u,
+                        const at::Tensor delta,
+                        const at::Tensor A,
+                        const at::Tensor B,
+                        const at::Tensor C,
+                        const at::Tensor out,
+                        void* D_ptr,
+                        void* delta_bias_ptr,
+                        void* x_ptr,
+                        bool delta_softplus) {
+
+    // Reset the parameters
+    memset(&params, 0, sizeof(params));
+
+    params.batch = batch;
+    params.dim = dim;
+    params.seqlen = seqlen;
+    params.dstate = dstate;
+    params.n_groups = n_groups;
+    params.n_chunks = n_chunks;
+    params.dim_ngroups_ratio = dim / n_groups;
+
+    params.delta_softplus = delta_softplus;
+
+    // Set the pointers and strides.
+    params.u_ptr = u.data_ptr();
+    params.delta_ptr = delta.data_ptr();
+    params.A_ptr = A.data_ptr();
+    params.B_ptr = B.data_ptr();
+    params.C_ptr = C.data_ptr();
+    params.D_ptr = D_ptr;
+    params.delta_bias_ptr = delta_bias_ptr;
+    params.out_ptr = out.data_ptr();
+    params.x_ptr = x_ptr;
+
+    // All stride are in elements, not bytes.
+    params.A_d_stride = A.stride(0);
+    params.A_dstate_stride = A.stride(1);
+    params.B_batch_stride = B.stride(0);
+    params.B_group_stride = B.stride(1);
+    params.B_dstate_stride = B.stride(2);
+    params.C_batch_stride = C.stride(0);
+    params.C_group_stride = C.stride(1);
+    params.C_dstate_stride = C.stride(2);
+    params.u_batch_stride = u.stride(0);
+    params.u_d_stride = u.stride(1);
+    params.delta_batch_stride = delta.stride(0);
+    params.delta_d_stride = delta.stride(1);
+
+    params.out_batch_stride = out.stride(0);
+    params.out_d_stride = out.stride(1);
+}
+
+void set_ssm_params_bwd(SSMParamsBwd &params,
+                        // sizes
+                        const size_t batch,
+                        const size_t dim,
+                        const size_t seqlen,
+                        const size_t dstate,
+                        const size_t n_groups,
+                        const size_t n_chunks,
+                        // device pointers
+                        const at::Tensor u,
+                        const at::Tensor delta,
+                        const at::Tensor A,
+                        const at::Tensor B,
+                        const at::Tensor C,
+                        const at::Tensor out,
+                        void* D_ptr,
+                        void* delta_bias_ptr,
+                        void* x_ptr,
+                        const at::Tensor dout,
+                        const at::Tensor du,
+                        const at::Tensor ddelta,
+                        const at::Tensor dA,
+                        const at::Tensor dB,
+                        const at::Tensor dC,
+                        void* dD_ptr,
+                        void* ddelta_bias_ptr,
+                        bool delta_softplus) {
+    // Pass in "dout" instead of "out", we're not gonna use "out" unless we have z
+    set_ssm_params_fwd(params, batch, dim, seqlen, dstate, n_groups, n_chunks,
+                       u, delta, A, B, C, dout,
+                       D_ptr, delta_bias_ptr, x_ptr, delta_softplus);
+
+    // Set the pointers and strides.
+    params.dout_ptr = dout.data_ptr();
+    params.du_ptr = du.data_ptr();
+    params.dA_ptr = dA.data_ptr();
+    params.dB_ptr = dB.data_ptr();
+    params.dC_ptr = dC.data_ptr();
+    params.dD_ptr = dD_ptr;
+    params.ddelta_ptr = ddelta.data_ptr();
+    params.ddelta_bias_ptr = ddelta_bias_ptr;
+    // All stride are in elements, not bytes.
+    params.dout_batch_stride = dout.stride(0);
+    params.dout_d_stride = dout.stride(1);
+    params.dA_d_stride = dA.stride(0);
+    params.dA_dstate_stride = dA.stride(1);
+    params.dB_batch_stride = dB.stride(0);
+    params.dB_group_stride = dB.stride(1);
+    params.dB_dstate_stride = dB.stride(2);
+    params.dC_batch_stride = dC.stride(0);
+    params.dC_group_stride = dC.stride(1);
+    params.dC_dstate_stride = dC.stride(2);
+    params.du_batch_stride = du.stride(0);
+    params.du_d_stride = du.stride(1);
+    params.ddelta_batch_stride = ddelta.stride(0);
+    params.ddelta_d_stride = ddelta.stride(1);
+
+}
+
+std::vector<at::Tensor>
+selective_scan_fwd(const at::Tensor &u, const at::Tensor &delta,
+                  const at::Tensor &A, const at::Tensor &B, const at::Tensor &C,
+                  const c10::optional<at::Tensor> &D_,
+                  const c10::optional<at::Tensor> &delta_bias_,
+                  bool delta_softplus,
+                  int nrows,
+                  bool out_float
+                  ) {
+    auto input_type = u.scalar_type();
+    auto weight_type = A.scalar_type();
+    TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
+    TORCH_CHECK(weight_type == at::ScalarType::Float);
+
+    TORCH_CHECK(delta.scalar_type() == input_type);
+    TORCH_CHECK(B.scalar_type() == input_type);
+    TORCH_CHECK(C.scalar_type() == input_type);
+
+    TORCH_CHECK(u.is_cuda());
+    TORCH_CHECK(delta.is_cuda());
+    TORCH_CHECK(A.is_cuda());
+    TORCH_CHECK(B.is_cuda());
+    TORCH_CHECK(C.is_cuda());
+
+    TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
+    TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
+
+    const auto sizes = u.sizes();
+    const int batch_size = sizes[0];
+    const int dim = sizes[1];
+    const int seqlen = sizes[2];
+    const int dstate = A.size(1);
+    const int n_groups = B.size(1);
+
+    TORCH_CHECK(dim % n_groups == 0, "dims should be dividable by n_groups");
+    TORCH_CHECK(dstate <= MAX_DSTATE, "selective_scan only supports state dimension <= 256");
+
+    CHECK_SHAPE(u, batch_size, dim, seqlen);
+    CHECK_SHAPE(delta, batch_size, dim, seqlen);
+    CHECK_SHAPE(A, dim, dstate);
+    CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen);
+    TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
+    CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen);
+    TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
+
+    if (D_.has_value()) {
+        auto D = D_.value();
+        TORCH_CHECK(D.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(D.is_cuda());
+        TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
+        CHECK_SHAPE(D, dim);
+    }
+
+    if (delta_bias_.has_value()) {
+        auto delta_bias = delta_bias_.value();
+        TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(delta_bias.is_cuda());
+        TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
+        CHECK_SHAPE(delta_bias, dim);
+    }
+
+    const int n_chunks = (seqlen + 2048 - 1) / 2048; // max is 128 * 16 = 2048 in fwd_kernel
+    at::Tensor out = torch::empty({batch_size, dim, seqlen}, u.options().dtype(out_float? (at::ScalarType::Float): input_type));
+    at::Tensor x = torch::empty({batch_size, dim, n_chunks, dstate * 2}, u.options().dtype(weight_type));
+
+    SSMParamsBase params;
+    set_ssm_params_fwd(params, batch_size, dim, seqlen, dstate, n_groups, n_chunks,
+                       u, delta, A, B, C, out,
+                       D_.has_value() ? D_.value().data_ptr() : nullptr,
+                       delta_bias_.has_value() ? delta_bias_.value().data_ptr() : nullptr,
+                       x.data_ptr(),
+                       delta_softplus);
+
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)u.get_device()};
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), "selective_scan_fwd", [&] {
+        if (!out_float) {
+            selective_scan_fwd_cuda<1, input_t, weight_t, input_t>(params, stream);
+        } else {
+            selective_scan_fwd_cuda<1, input_t, weight_t, float>(params, stream);
+        }
+    });
+    std::vector<at::Tensor> result = {out, x};
+    return result;
+}
+
+std::vector<at::Tensor>
+selective_scan_bwd(const at::Tensor &u, const at::Tensor &delta,
+                  const at::Tensor &A, const at::Tensor &B, const at::Tensor &C,
+                  const c10::optional<at::Tensor> &D_,
+                  const c10::optional<at::Tensor> &delta_bias_,
+                  const at::Tensor &dout,
+                  const c10::optional<at::Tensor> &x_,
+                  bool delta_softplus,
+                  int nrows
+                  ) {
+    auto input_type = u.scalar_type();
+    auto weight_type = A.scalar_type();
+    auto output_type = dout.scalar_type();
+    TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
+    TORCH_CHECK(weight_type == at::ScalarType::Float);
+    TORCH_CHECK(output_type == input_type || output_type == at::ScalarType::Float);
+
+    TORCH_CHECK(delta.scalar_type() == input_type);
+    TORCH_CHECK(B.scalar_type() == input_type);
+    TORCH_CHECK(C.scalar_type() == input_type);
+
+    TORCH_CHECK(u.is_cuda());
+    TORCH_CHECK(delta.is_cuda());
+    TORCH_CHECK(A.is_cuda());
+    TORCH_CHECK(B.is_cuda());
+    TORCH_CHECK(C.is_cuda());
+    TORCH_CHECK(dout.is_cuda());
+
+    TORCH_CHECK(u.stride(-1) == 1 || u.size(-1) == 1);
+    TORCH_CHECK(delta.stride(-1) == 1 || delta.size(-1) == 1);
+    TORCH_CHECK(dout.stride(-1) == 1 || dout.size(-1) == 1);
+
+    const auto sizes = u.sizes();
+    const int batch_size = sizes[0];
+    const int dim = sizes[1];
+    const int seqlen = sizes[2];
+    const int dstate = A.size(1);
+    const int n_groups = B.size(1);
+
+    TORCH_CHECK(dim % n_groups == 0, "dims should be dividable by n_groups");
+    TORCH_CHECK(dstate <= MAX_DSTATE, "selective_scan only supports state dimension <= 256");
+    
+    CHECK_SHAPE(u, batch_size, dim, seqlen);
+    CHECK_SHAPE(delta, batch_size, dim, seqlen);
+    CHECK_SHAPE(A, dim, dstate);
+    CHECK_SHAPE(B, batch_size, n_groups, dstate, seqlen);
+    TORCH_CHECK(B.stride(-1) == 1 || B.size(-1) == 1);
+    CHECK_SHAPE(C, batch_size, n_groups, dstate, seqlen);
+    TORCH_CHECK(C.stride(-1) == 1 || C.size(-1) == 1);
+    CHECK_SHAPE(dout, batch_size, dim, seqlen);
+
+    if (D_.has_value()) {
+        auto D = D_.value();
+        TORCH_CHECK(D.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(D.is_cuda());
+        TORCH_CHECK(D.stride(-1) == 1 || D.size(-1) == 1);
+        CHECK_SHAPE(D, dim);
+    }
+
+    if (delta_bias_.has_value()) {
+        auto delta_bias = delta_bias_.value();
+        TORCH_CHECK(delta_bias.scalar_type() == at::ScalarType::Float);
+        TORCH_CHECK(delta_bias.is_cuda());
+        TORCH_CHECK(delta_bias.stride(-1) == 1 || delta_bias.size(-1) == 1);
+        CHECK_SHAPE(delta_bias, dim);
+    }
+
+    at::Tensor out;
+    const int n_chunks = (seqlen + 2048 - 1) / 2048;
+    // const int n_chunks = (seqlen + 1024 - 1) / 1024;
+    if (n_chunks > 1) { TORCH_CHECK(x_.has_value()); }
+    if (x_.has_value()) {
+        auto x = x_.value();
+        TORCH_CHECK(x.scalar_type() == weight_type);
+        TORCH_CHECK(x.is_cuda());
+        TORCH_CHECK(x.is_contiguous());
+        CHECK_SHAPE(x, batch_size, dim, n_chunks, 2 * dstate);
+    }
+
+    at::Tensor du = torch::empty_like(u);
+    at::Tensor ddelta = torch::empty_like(delta);
+    at::Tensor dA = torch::zeros_like(A);
+    at::Tensor dB = torch::zeros_like(B, B.options().dtype(torch::kFloat32));
+    at::Tensor dC = torch::zeros_like(C, C.options().dtype(torch::kFloat32));
+    at::Tensor dD;
+    if (D_.has_value()) { dD = torch::zeros_like(D_.value()); }
+    at::Tensor ddelta_bias;
+    if (delta_bias_.has_value()) { ddelta_bias = torch::zeros_like(delta_bias_.value()); }
+
+    SSMParamsBwd params;
+    set_ssm_params_bwd(params, batch_size, dim, seqlen, dstate, n_groups, n_chunks,
+                       u, delta, A, B, C, out,
+                       D_.has_value() ? D_.value().data_ptr() : nullptr,
+                       delta_bias_.has_value() ? delta_bias_.value().data_ptr() : nullptr,
+                       x_.has_value() ? x_.value().data_ptr() : nullptr,
+                       dout, du, ddelta, dA, dB, dC,
+                       D_.has_value() ? dD.data_ptr() : nullptr,
+                       delta_bias_.has_value() ? ddelta_bias.data_ptr() : nullptr,
+                       delta_softplus);
+
+    // Otherwise the kernel will be launched from cuda:0 device
+    // Cast to char to avoid compiler warning about narrowing
+    at::cuda::CUDAGuard device_guard{(char)u.get_device()};
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(u.scalar_type(), "selective_scan_bwd", [&] {
+        if (output_type == input_type) {
+            selective_scan_bwd_cuda<1, input_t, weight_t, input_t>(params, stream);
+        } else {
+            selective_scan_bwd_cuda<1, input_t, weight_t, float>(params, stream);
+        }
+    });
+    std::vector<at::Tensor> result = {du, ddelta, dA, dB.to(B.dtype()), dC.to(C.dtype()), dD, ddelta_bias};
+    return result;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("fwd", &selective_scan_fwd, "Selective scan forward");
+    m.def("bwd", &selective_scan_bwd, "Selective scan backward");
+}

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/reverse_scan.cuh

@@ -0,0 +1,403 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <cub/config.cuh>
+
+#include <cub/util_ptx.cuh>
+#include <cub/util_type.cuh>
+#include <cub/block/block_raking_layout.cuh>
+// #include <cub/detail/uninitialized_copy.cuh>
+#include "uninitialized_copy.cuh"
+#include "cub_extra.cuh"
+
+/**
+ * Perform a reverse sequential reduction over \p LENGTH elements of the \p input array.  The aggregate is returned.
+ */
+template <
+    int         LENGTH,
+    typename    T,
+    typename    ReductionOp>
+__device__ __forceinline__ T ThreadReverseReduce(const T (&input)[LENGTH], ReductionOp reduction_op) {
+    static_assert(LENGTH > 0);
+    T retval = input[LENGTH - 1];
+    #pragma unroll
+    for (int i = LENGTH - 2; i >= 0; --i) { retval = reduction_op(retval, input[i]); }
+    return retval;
+}
+
+/**
+ * Perform a sequential inclusive postfix reverse scan over the statically-sized \p input array, seeded with the specified \p postfix.  The aggregate is returned.
+ */
+template <
+    int         LENGTH,
+    typename    T,
+    typename    ScanOp>
+__device__ __forceinline__ T ThreadReverseScanInclusive(
+    const T (&input)[LENGTH],
+    T (&output)[LENGTH],
+    ScanOp scan_op,
+    const T postfix)
+{
+    T inclusive = postfix;
+    #pragma unroll
+    for (int i = LENGTH - 1; i >= 0; --i) {
+        inclusive = scan_op(inclusive, input[i]);
+        output[i] = inclusive;
+    }
+}
+
+/**
+ * Perform a sequential exclusive postfix reverse scan over the statically-sized \p input array, seeded with the specified \p postfix.  The aggregate is returned.
+ */
+template <
+    int         LENGTH,
+    typename    T,
+    typename    ScanOp>
+__device__ __forceinline__ T ThreadReverseScanExclusive(
+    const T (&input)[LENGTH],
+    T (&output)[LENGTH],
+    ScanOp scan_op,
+    const T postfix)
+{
+    // Careful, output maybe be aliased to input
+    T exclusive = postfix;
+    T inclusive;
+    #pragma unroll
+    for (int i = LENGTH - 1; i >= 0; --i) {
+        inclusive = scan_op(exclusive, input[i]);
+        output[i] = exclusive;
+        exclusive = inclusive;
+    }
+    return inclusive;
+}
+
+
+/**
+ * \brief WarpReverseScan provides SHFL-based variants of parallel postfix scan of items partitioned across a CUDA thread warp.
+ *
+ * LOGICAL_WARP_THREADS must be a power-of-two
+ */
+template <
+    typename    T,                      ///< Data type being scanned
+    int         LOGICAL_WARP_THREADS    ///< Number of threads per logical warp
+    >
+struct WarpReverseScan {
+    //---------------------------------------------------------------------
+    // Constants and type definitions
+    //---------------------------------------------------------------------
+
+    /// Whether the logical warp size and the PTX warp size coincide
+    static constexpr bool IS_ARCH_WARP = (LOGICAL_WARP_THREADS == CUB_WARP_THREADS(0));
+    /// The number of warp scan steps
+    static constexpr int STEPS = cub::Log2<LOGICAL_WARP_THREADS>::VALUE;
+    static_assert(LOGICAL_WARP_THREADS == 1 << STEPS);
+
+
+    //---------------------------------------------------------------------
+    // Thread fields
+    //---------------------------------------------------------------------
+
+    /// Lane index in logical warp
+    unsigned int lane_id;
+
+    /// Logical warp index in 32-thread physical warp
+    unsigned int warp_id;
+
+    /// 32-thread physical warp member mask of logical warp
+    unsigned int member_mask;
+
+    //---------------------------------------------------------------------
+    // Construction
+    //---------------------------------------------------------------------
+
+    /// Constructor
+    explicit __device__ __forceinline__
+    WarpReverseScan()
+        : lane_id(cub::LaneId())
+        , warp_id(IS_ARCH_WARP ? 0 : (lane_id / LOGICAL_WARP_THREADS))
+        // , member_mask(cub::WarpMask<LOGICAL_WARP_THREADS>(warp_id))
+        , member_mask(WarpMask<LOGICAL_WARP_THREADS>(warp_id))
+    {
+        if (!IS_ARCH_WARP) {
+            lane_id = lane_id % LOGICAL_WARP_THREADS;
+        }
+    }
+
+
+    /// Broadcast
+    __device__ __forceinline__ T Broadcast(
+        T               input,              ///< [in] The value to broadcast
+        int             src_lane)           ///< [in] Which warp lane is to do the broadcasting
+    {
+        return cub::ShuffleIndex<LOGICAL_WARP_THREADS>(input, src_lane, member_mask);
+    }
+
+
+    /// Inclusive scan
+    template <typename ScanOpT>
+    __device__ __forceinline__ void InclusiveReverseScan(
+        T               input,              ///< [in] Calling thread's input item.
+        T               &inclusive_output,  ///< [out] Calling thread's output item.  May be aliased with \p input.
+        ScanOpT         scan_op)            ///< [in] Binary scan operator
+    {
+        inclusive_output = input;
+        #pragma unroll
+        for (int STEP = 0; STEP < STEPS; STEP++) {
+            int offset = 1 << STEP;
+            T temp = cub::ShuffleDown<LOGICAL_WARP_THREADS>(
+                inclusive_output, offset, LOGICAL_WARP_THREADS - 1, member_mask
+            );
+            // Perform scan op if from a valid peer
+            inclusive_output = static_cast<int>(lane_id) >= LOGICAL_WARP_THREADS - offset
+                ? inclusive_output : scan_op(temp, inclusive_output);
+        }
+    }
+
+    /// Exclusive scan
+    // Get exclusive from inclusive
+    template <typename ScanOpT>
+    __device__ __forceinline__ void ExclusiveReverseScan(
+        T              input,              ///< [in] Calling thread's input item.
+        T              &exclusive_output,  ///< [out] Calling thread's output item.  May be aliased with \p input.
+        ScanOpT        scan_op,            ///< [in] Binary scan operator
+        T              &warp_aggregate)    ///< [out] Warp-wide aggregate reduction of input items.
+    {
+        T inclusive_output;
+        InclusiveReverseScan(input, inclusive_output, scan_op);
+        warp_aggregate = cub::ShuffleIndex<LOGICAL_WARP_THREADS>(inclusive_output, 0, member_mask);
+        // initial value unknown
+        exclusive_output = cub::ShuffleDown<LOGICAL_WARP_THREADS>(
+            inclusive_output, 1, LOGICAL_WARP_THREADS - 1, member_mask
+        );
+    }
+
+    /**
+     * \brief Computes both inclusive and exclusive reverse scans using the specified binary scan functor across the calling warp.  Because no initial value is supplied, the \p exclusive_output computed for the last <em>warp-lane</em> is undefined.
+     */
+    template <typename ScanOpT>
+    __device__ __forceinline__ void ReverseScan(
+        T               input,              ///< [in] Calling thread's input item.
+        T               &inclusive_output,  ///< [out] Calling thread's inclusive-scan output item.
+        T               &exclusive_output,  ///< [out] Calling thread's exclusive-scan output item.
+        ScanOpT         scan_op)            ///< [in] Binary scan operator
+    {
+        InclusiveReverseScan(input, inclusive_output, scan_op);
+        // initial value unknown
+        exclusive_output = cub::ShuffleDown<LOGICAL_WARP_THREADS>(
+            inclusive_output, 1, LOGICAL_WARP_THREADS - 1, member_mask
+        );
+    }
+
+};
+
+/**
+ * \brief BlockReverseScan provides variants of raking-based parallel postfix scan across a CUDA thread block.
+ */
+template <
+    typename    T,              ///< Data type being scanned
+    int         BLOCK_DIM_X,    ///< The thread block length in threads along the X dimension
+    bool        MEMOIZE=false   ///< Whether or not to buffer outer raking scan partials to incur fewer shared memory reads at the expense of higher register pressure
+    >
+struct BlockReverseScan {
+    //---------------------------------------------------------------------
+    // Types and constants
+    //---------------------------------------------------------------------
+
+    /// Constants
+    /// The thread block size in threads
+    static constexpr int BLOCK_THREADS = BLOCK_DIM_X;
+
+    /// Layout type for padded thread block raking grid
+    using BlockRakingLayout = cub::BlockRakingLayout<T, BLOCK_THREADS>;
+    // The number of reduction elements is not a multiple of the number of raking threads for now
+    static_assert(BlockRakingLayout::UNGUARDED);
+
+    /// Number of raking threads
+    static constexpr int RAKING_THREADS = BlockRakingLayout::RAKING_THREADS;
+    /// Number of raking elements per warp synchronous raking thread
+    static constexpr int SEGMENT_LENGTH = BlockRakingLayout::SEGMENT_LENGTH;
+    /// Cooperative work can be entirely warp synchronous
+    static constexpr bool WARP_SYNCHRONOUS = (int(BLOCK_THREADS) == int(RAKING_THREADS));
+
+    ///  WarpReverseScan utility type
+    using WarpReverseScan = WarpReverseScan<T, RAKING_THREADS>;
+
+    /// Shared memory storage layout type
+    struct _TempStorage {
+        typename BlockRakingLayout::TempStorage raking_grid;     ///< Padded thread block raking grid
+    };
+
+
+    /// Alias wrapper allowing storage to be unioned
+    struct TempStorage : cub::Uninitialized<_TempStorage> {};
+
+
+    //---------------------------------------------------------------------
+    // Per-thread fields
+    //---------------------------------------------------------------------
+
+    // Thread fields
+    _TempStorage    &temp_storage;
+    unsigned int    linear_tid;
+    T               cached_segment[SEGMENT_LENGTH];
+
+
+    //---------------------------------------------------------------------
+    // Utility methods
+    //---------------------------------------------------------------------
+
+    /// Performs upsweep raking reduction, returning the aggregate
+    template <typename ScanOp>
+    __device__ __forceinline__ T Upsweep(ScanOp scan_op) {
+        T *smem_raking_ptr = BlockRakingLayout::RakingPtr(temp_storage.raking_grid, linear_tid);
+        // Read data into registers
+        #pragma unroll
+        for (int i = 0; i < SEGMENT_LENGTH; ++i) { cached_segment[i] = smem_raking_ptr[i]; }
+        T raking_partial = cached_segment[SEGMENT_LENGTH - 1];
+        #pragma unroll
+        for (int i = SEGMENT_LENGTH - 2; i >= 0; --i) {
+            raking_partial = scan_op(raking_partial, cached_segment[i]);
+        }
+        return raking_partial;
+    }
+
+
+    /// Performs exclusive downsweep raking scan
+    template <typename ScanOp>
+    __device__ __forceinline__ void ExclusiveDownsweep(
+        ScanOp          scan_op,
+        T               raking_partial)
+    {
+        T *smem_raking_ptr = BlockRakingLayout::RakingPtr(temp_storage.raking_grid, linear_tid);
+        // Read data back into registers
+        if (!MEMOIZE) {
+            #pragma unroll
+            for (int i = 0; i < SEGMENT_LENGTH; ++i) { cached_segment[i] = smem_raking_ptr[i]; }
+        }
+        ThreadReverseScanExclusive(cached_segment, cached_segment, scan_op, raking_partial);
+        // Write data back to smem
+        #pragma unroll
+        for (int i = 0; i < SEGMENT_LENGTH; ++i) { smem_raking_ptr[i] = cached_segment[i]; }
+    }
+
+
+    //---------------------------------------------------------------------
+    // Constructors
+    //---------------------------------------------------------------------
+
+    /// Constructor
+    __device__ __forceinline__ BlockReverseScan(
+        TempStorage &temp_storage)
+    :
+        temp_storage(temp_storage.Alias()),
+        linear_tid(cub::RowMajorTid(BLOCK_DIM_X, 1, 1))
+    {}
+
+
+    /// Computes an exclusive thread block-wide postfix scan using the specified binary \p scan_op functor.  Each thread contributes one input element.  the call-back functor \p block_postfix_callback_op is invoked by the first warp in the block, and the value returned by <em>lane</em><sub>0</sub> in that warp is used as the "seed" value that logically postfixes the thread block's scan inputs.  Also provides every thread with the block-wide \p block_aggregate of all inputs.
+    template <
+        typename ScanOp,
+        typename BlockPostfixCallbackOp>
+    __device__ __forceinline__ void ExclusiveReverseScan(
+        T                       input,                          ///< [in] Calling thread's input item
+        T                       &exclusive_output,              ///< [out] Calling thread's output item (may be aliased to \p input)
+        ScanOp                  scan_op,                        ///< [in] Binary scan operator
+        BlockPostfixCallbackOp  &block_postfix_callback_op)     ///< [in-out] <b>[<em>warp</em><sub>0</sub> only]</b> Call-back functor for specifying a thread block-wide postfix to be applied to all inputs.
+    {
+        if (WARP_SYNCHRONOUS) {
+            // Short-circuit directly to warp-synchronous scan
+            T block_aggregate;
+            WarpReverseScan warp_scan;
+            warp_scan.ExclusiveReverseScan(input, exclusive_output, scan_op, block_aggregate);
+            // Obtain warp-wide postfix in lane0, then broadcast to other lanes
+            T block_postfix = block_postfix_callback_op(block_aggregate);
+            block_postfix = warp_scan.Broadcast(block_postfix, 0);
+            exclusive_output = linear_tid == BLOCK_THREADS - 1 ? block_postfix : scan_op(block_postfix, exclusive_output);
+        } else {
+            // Place thread partial into shared memory raking grid
+            T *placement_ptr = BlockRakingLayout::PlacementPtr(temp_storage.raking_grid, linear_tid);
+            detail::uninitialized_copy(placement_ptr, input);
+            cub::CTA_SYNC();
+            // Reduce parallelism down to just raking threads
+            if (linear_tid < RAKING_THREADS) {
+                WarpReverseScan warp_scan;
+                // Raking upsweep reduction across shared partials
+                T upsweep_partial = Upsweep(scan_op);
+                // Warp-synchronous scan
+                T exclusive_partial, block_aggregate;
+                warp_scan.ExclusiveReverseScan(upsweep_partial, exclusive_partial, scan_op, block_aggregate);
+                // Obtain block-wide postfix in lane0, then broadcast to other lanes
+                T block_postfix = block_postfix_callback_op(block_aggregate);
+                block_postfix = warp_scan.Broadcast(block_postfix, 0);
+                // Update postfix with warpscan exclusive partial
+                T downsweep_postfix = linear_tid == RAKING_THREADS - 1
+                    ? block_postfix : scan_op(block_postfix, exclusive_partial);
+                // Exclusive raking downsweep scan
+                ExclusiveDownsweep(scan_op, downsweep_postfix);
+            }
+            cub::CTA_SYNC();
+            // Grab thread postfix from shared memory
+            exclusive_output = *placement_ptr;
+
+            // // Compute warp scan in each warp.
+            // // The exclusive output from the last lane in each warp is invalid.
+            // T inclusive_output;
+            // WarpReverseScan warp_scan;
+            // warp_scan.ReverseScan(input, inclusive_output, exclusive_output, scan_op);
+
+            // // Compute the warp-wide postfix and block-wide aggregate for each warp.  Warp postfix for the last warp is invalid.
+            // T block_aggregate;
+            // T warp_postfix = ComputeWarpPostfix(scan_op, inclusive_output, block_aggregate);
+
+            // // Apply warp postfix to our lane's partial
+            // if (warp_id != 0) {
+            //     exclusive_output = scan_op(warp_postfix, exclusive_output);
+            //     if (lane_id == 0) { exclusive_output = warp_postfix; }
+            // }
+
+            // // Use the first warp to determine the thread block postfix, returning the result in lane0
+            // if (warp_id == 0) {
+            //     T block_postfix = block_postfix_callback_op(block_aggregate);
+            //     if (lane_id == 0) {
+            //         // Share the postfix with all threads
+            //         detail::uninitialized_copy(&temp_storage.block_postfix,
+            //                                   block_postfix);
+
+            //         exclusive_output = block_postfix; // The block postfix is the exclusive output for tid0
+            //     }
+            // }
+
+            // cub::CTA_SYNC();
+
+            // // Incorporate thread block postfix into outputs
+            // T block_postfix = temp_storage.block_postfix;
+            // if (linear_tid > 0) { exclusive_output = scan_op(block_postfix, exclusive_output); }
+        }
+    }
+
+
+    /**
+     * \brief Computes an inclusive block-wide postfix scan using the specified binary \p scan_op functor.  Each thread contributes an array of consecutive input elements.  the call-back functor \p block_postfix_callback_op is invoked by the first warp in the block, and the value returned by <em>lane</em><sub>0</sub> in that warp is used as the "seed" value that logically postfixes the thread block's scan inputs.  Also provides every thread with the block-wide \p block_aggregate of all inputs.
+     */
+    template <
+        int             ITEMS_PER_THREAD,
+        typename        ScanOp,
+        typename        BlockPostfixCallbackOp>
+    __device__ __forceinline__ void InclusiveReverseScan(
+        T                       (&input)[ITEMS_PER_THREAD],     ///< [in] Calling thread's input items
+        T                       (&output)[ITEMS_PER_THREAD],    ///< [out] Calling thread's output items (may be aliased to \p input)
+        ScanOp                  scan_op,                        ///< [in] Binary scan functor
+        BlockPostfixCallbackOp   &block_postfix_callback_op)    ///< [in-out] <b>[<em>warp</em><sub>0</sub> only]</b> Call-back functor for specifying a block-wide postfix to be applied to the logical input sequence.
+    {
+        // Reduce consecutive thread items in registers
+        T thread_postfix = ThreadReverseReduce(input, scan_op);
+        // Exclusive thread block-scan
+        ExclusiveReverseScan(thread_postfix, thread_postfix, scan_op, block_postfix_callback_op);
+        // Inclusive scan in registers with postfix as seed
+        ThreadReverseScanInclusive(input, output, scan_op, thread_postfix);
+    }
+
+};

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/selective_scan.h

@@ -0,0 +1,90 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct SSMScanParamsBase {
+    using index_t = uint32_t;
+
+    int batch, seqlen, n_chunks;
+    index_t a_batch_stride;
+    index_t b_batch_stride;
+    index_t out_batch_stride;
+
+    // Common data pointers.
+    void *__restrict__ a_ptr;
+    void *__restrict__ b_ptr;
+    void *__restrict__ out_ptr;
+    void *__restrict__ x_ptr;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct SSMParamsBase {
+    using index_t = uint32_t;
+
+    int batch, dim, seqlen, dstate, n_groups, n_chunks;
+    int dim_ngroups_ratio;
+
+    bool delta_softplus;
+
+    index_t A_d_stride;
+    index_t A_dstate_stride;
+    index_t B_batch_stride;
+    index_t B_d_stride;
+    index_t B_dstate_stride;
+    index_t B_group_stride;
+    index_t C_batch_stride;
+    index_t C_d_stride;
+    index_t C_dstate_stride;
+    index_t C_group_stride;
+    index_t u_batch_stride;
+    index_t u_d_stride;
+    index_t delta_batch_stride;
+    index_t delta_d_stride;
+    index_t out_batch_stride;
+    index_t out_d_stride;
+
+    // Common data pointers.
+    void *__restrict__ A_ptr;
+    void *__restrict__ B_ptr;
+    void *__restrict__ C_ptr;
+    void *__restrict__ D_ptr;
+    void *__restrict__ u_ptr;
+    void *__restrict__ delta_ptr;
+    void *__restrict__ delta_bias_ptr;
+    void *__restrict__ out_ptr;
+    void *__restrict__ x_ptr;
+};
+
+struct SSMParamsBwd: public SSMParamsBase {
+    index_t dout_batch_stride;
+    index_t dout_d_stride;
+    index_t dA_d_stride;
+    index_t dA_dstate_stride;
+    index_t dB_batch_stride;
+    index_t dB_group_stride;
+    index_t dB_d_stride;
+    index_t dB_dstate_stride;
+    index_t dC_batch_stride;
+    index_t dC_group_stride;
+    index_t dC_d_stride;
+    index_t dC_dstate_stride;
+    index_t du_batch_stride;
+    index_t du_d_stride;
+    index_t ddelta_batch_stride;
+    index_t ddelta_d_stride;
+
+    // Common data pointers.
+    void *__restrict__ dout_ptr;
+    void *__restrict__ dA_ptr;
+    void *__restrict__ dB_ptr;
+    void *__restrict__ dC_ptr;
+    void *__restrict__ dD_ptr;
+    void *__restrict__ du_ptr;
+    void *__restrict__ ddelta_ptr;
+    void *__restrict__ ddelta_bias_ptr;
+};

rscd/models/backbones/lib_mamba/kernels/selective_scan/csrc/selective_scan/selective_scan_common.h

@@ -0,0 +1,210 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+#pragma once
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+#include <c10/util/complex.h>  // For scalar_value_type
+
+#define MAX_DSTATE 256
+
+inline __device__ float2 operator+(const float2 & a, const float2 & b){
+    return {a.x + b.x, a.y + b.y};
+}
+
+inline __device__ float3 operator+(const float3 &a, const float3 &b) {
+  return {a.x + b.x, a.y + b.y, a.z + b.z};
+}
+
+inline __device__ float4 operator+(const float4 & a, const float4 & b){
+    return {a.x + b.x, a.y + b.y, a.z + b.z, a.w + b.w};
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<int BYTES> struct BytesToType {};
+
+template<> struct BytesToType<16> {
+    using Type = uint4;
+    static_assert(sizeof(Type) == 16);
+};
+
+template<> struct BytesToType<8> {
+    using Type = uint64_t;
+    static_assert(sizeof(Type) == 8);
+};
+
+template<> struct BytesToType<4> {
+    using Type = uint32_t;
+    static_assert(sizeof(Type) == 4);
+};
+
+template<> struct BytesToType<2> {
+    using Type = uint16_t;
+    static_assert(sizeof(Type) == 2);
+};
+
+template<> struct BytesToType<1> {
+    using Type = uint8_t;
+    static_assert(sizeof(Type) == 1);
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename scalar_t, int N>
+struct Converter{
+    static inline __device__ void to_float(const scalar_t (&src)[N], float (&dst)[N]) {
+        #pragma unroll
+        for (int i = 0; i < N; ++i) { dst[i] = src[i]; }
+    }
+};
+
+template<int N>
+struct Converter<at::Half, N>{
+    static inline __device__ void to_float(const at::Half (&src)[N], float (&dst)[N]) {
+        static_assert(N % 2 == 0);
+        auto &src2 = reinterpret_cast<const half2 (&)[N / 2]>(src);
+        auto &dst2 = reinterpret_cast<float2 (&)[N / 2]>(dst);
+        #pragma unroll
+        for (int i = 0; i < N / 2; ++i) { dst2[i] = __half22float2(src2[i]); }
+    }
+};
+
+#if __CUDA_ARCH__ >= 800
+template<int N>
+struct Converter<at::BFloat16, N>{
+    static inline __device__ void to_float(const at::BFloat16 (&src)[N], float (&dst)[N]) {
+        static_assert(N % 2 == 0);
+        auto &src2 = reinterpret_cast<const nv_bfloat162 (&)[N / 2]>(src);
+        auto &dst2 = reinterpret_cast<float2 (&)[N / 2]>(dst);
+        #pragma unroll
+        for (int i = 0; i < N / 2; ++i) { dst2[i] = __bfloat1622float2(src2[i]); }
+    }
+};
+#endif
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+template<typename scalar_t> struct SSMScanOp;
+
+template<>
+struct SSMScanOp<float> {
+    __device__ __forceinline__ float2 operator()(const float2 &ab0, const float2 &ab1) const {
+        return make_float2(ab1.x * ab0.x, ab1.x * ab0.y + ab1.y);
+    }
+};
+
+// A stateful callback functor that maintains a running prefix to be applied
+// during consecutive scan operations.
+template <typename scalar_t> struct SSMScanPrefixCallbackOp {
+    using scan_t = std::conditional_t<std::is_same_v<scalar_t, float>, float2, float4>;
+    scan_t running_prefix;
+    // Constructor
+    __device__ SSMScanPrefixCallbackOp(scan_t running_prefix_) : running_prefix(running_prefix_) {}
+    // Callback operator to be entered by the first warp of threads in the block.
+    // Thread-0 is responsible for returning a value for seeding the block-wide scan.
+    __device__ scan_t operator()(scan_t block_aggregate) {
+        scan_t old_prefix = running_prefix;
+        running_prefix = SSMScanOp<scalar_t>()(running_prefix, block_aggregate);
+        return old_prefix;
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Ktraits>
+inline __device__ void load_input(typename Ktraits::input_t *u,
+                                  typename Ktraits::input_t (&u_vals)[Ktraits::kNItems],
+                                  typename Ktraits::BlockLoadT::TempStorage &smem_load,
+                                  int seqlen) {
+    if constexpr (Ktraits::kIsEvenLen) {
+        auto& smem_load_vec = reinterpret_cast<typename Ktraits::BlockLoadVecT::TempStorage&>(smem_load);
+        using vec_t = typename Ktraits::vec_t;
+        Ktraits::BlockLoadVecT(smem_load_vec).Load(
+            reinterpret_cast<vec_t*>(u),
+            reinterpret_cast<vec_t(&)[Ktraits::kNLoads]>(u_vals)
+       );
+    } else {
+        Ktraits::BlockLoadT(smem_load).Load(u, u_vals, seqlen, 0.f);
+    }
+}
+
+template<typename Ktraits>
+inline __device__ void load_weight(typename Ktraits::input_t *Bvar,
+                                   typename Ktraits::weight_t (&B_vals)[Ktraits::kNItems],
+                                   typename Ktraits::BlockLoadWeightT::TempStorage &smem_load_weight,
+                                   int seqlen) {
+    constexpr int kNItems = Ktraits::kNItems;
+    typename Ktraits::input_t B_vals_load[kNItems];
+    if constexpr (Ktraits::kIsEvenLen) {
+        auto& smem_load_weight_vec = reinterpret_cast<typename Ktraits::BlockLoadWeightVecT::TempStorage&>(smem_load_weight);
+        using vec_t = typename Ktraits::vec_t;
+        Ktraits::BlockLoadWeightVecT(smem_load_weight_vec).Load(
+            reinterpret_cast<vec_t*>(Bvar),
+            reinterpret_cast<vec_t(&)[Ktraits::kNLoads]>(B_vals_load)
+      );
+    } else {
+        Ktraits::BlockLoadWeightT(smem_load_weight).Load(Bvar, B_vals_load, seqlen, 0.f);
+    }
+    // #pragma unroll
+    // for (int i = 0; i < kNItems; ++i) { B_vals[i] = B_vals_load[i]; }
+    Converter<typename Ktraits::input_t, kNItems>::to_float(B_vals_load, B_vals);
+}
+
+template<typename Ktraits>
+inline __device__ void store_output(typename Ktraits::input_t *out,
+                                    const float (&out_vals)[Ktraits::kNItems],
+                                    typename Ktraits::BlockStoreT::TempStorage &smem_store,
+                                    int seqlen) {
+    typename Ktraits::input_t write_vals[Ktraits::kNItems];
+    #pragma unroll
+    for (int i = 0; i < Ktraits::kNItems; ++i) { write_vals[i] = out_vals[i]; }
+    if constexpr (Ktraits::kIsEvenLen) {
+        auto& smem_store_vec = reinterpret_cast<typename Ktraits::BlockStoreVecT::TempStorage&>(smem_store);
+        using vec_t = typename Ktraits::vec_t;
+        Ktraits::BlockStoreVecT(smem_store_vec).Store(
+            reinterpret_cast<vec_t*>(out),
+            reinterpret_cast<vec_t(&)[Ktraits::kNLoads]>(write_vals)
+       );
+    } else {
+        Ktraits::BlockStoreT(smem_store).Store(out, write_vals, seqlen);
+    }
+}
+
+template<typename Ktraits>
+inline __device__ void store_output1(typename Ktraits::output_t *out,
+                                    const float (&out_vals)[Ktraits::kNItems],
+                                    typename Ktraits::BlockStoreOutputT::TempStorage &smem_store,
+                                    int seqlen) {
+    typename Ktraits::output_t write_vals[Ktraits::kNItems];
+    #pragma unroll
+    for (int i = 0; i < Ktraits::kNItems; ++i) { write_vals[i] = out_vals[i]; }
+    if constexpr (Ktraits::kIsEvenLen) {
+        auto& smem_store_vec = reinterpret_cast<typename Ktraits::BlockStoreOutputVecT::TempStorage&>(smem_store);
+        using vec_t = typename Ktraits::vec_t;
+        Ktraits::BlockStoreOutputVecT(smem_store_vec).Store(
+            reinterpret_cast<vec_t*>(out),
+            reinterpret_cast<vec_t(&)[Ktraits::kNLoadsOutput]>(write_vals)
+       );
+    } else {
+        Ktraits::BlockStoreOutputT(smem_store).Store(out, write_vals, seqlen);
+    }
+}
+
+template<typename Ktraits>
+inline __device__ void load_output(typename Ktraits::output_t *u,
+                                  typename Ktraits::output_t (&u_vals)[Ktraits::kNItems],
+                                  typename Ktraits::BlockLoadOutputT::TempStorage &smem_load,
+                                  int seqlen) {
+    if constexpr (Ktraits::kIsEvenLen) {
+        auto& smem_load_vec = reinterpret_cast<typename Ktraits::BlockLoadOutputVecT::TempStorage&>(smem_load);
+        using vec_t = typename Ktraits::vec_t;
+        Ktraits::BlockLoadOutputVecT(smem_load_vec).Load(
+            reinterpret_cast<vec_t*>(u),
+            reinterpret_cast<vec_t(&)[Ktraits::kNLoadsOutput]>(u_vals)
+       );
+    } else {
+        Ktraits::BlockLoadOutputT(smem_load).Load(u, u_vals, seqlen, 0.f);
+    }
+}