Move custom_ops under layers; move _custom_ops.py → custom_all_reduce_ops.py (#14326)

python/sglang/srt/layers/attention/mamba/ops/mamba_ssm.py

@@ -11,8 +11,6 @@ import triton
 import triton.language as tl
 from packaging import version
 
-from sglang.srt import _custom_ops as ops
-
 PAD_SLOT_ID = -1
 
 TRITON3 = version.parse(triton.__version__) >= version.parse("3.0.0")
@@ -344,99 +342,3 @@ def selective_state_update(
             BLOCK_SIZE_M,
             num_warps=num_warps,
         )
-
-
-def selective_scan_fn(
-    u,
-    ssm_states,
-    delta,
-    A,
-    B,
-    C,
-    D=None,
-    z=None,
-    delta_bias=None,
-    delta_softplus=False,
-    query_start_loc=None,
-    cache_indices=None,
-    has_initial_state=None,
-    pad_slot_id=PAD_SLOT_ID,
-) -> torch.Tensor:
-    """
-    u: (dim, total_length) for varlen or (batch, dim, seqlen)
-        applies changes in place.
-    ssm_states: (batch, dim, dstate) or (batch, nheads, dim, dstate)
-        applies changes in place.
-    delta: (dim, total_length) for varlen or (batch, dim, seqlen)
-    A: (dim, dstate)
-    B: (ngroups, dstate, total_length) for varlen or
-                                        (batch,ngroups,dstate,seqlen)
-    C: (ngroups, dstate, total_length) for varlen or
-                                        (batch,ngroups,dstate,seqlen)
-    D: (dim,)
-    z: (dim, total_length) for varlen or (batch, dim, seqlen)
-    dt_bias: (dim,) or (dim)
-    query_start_loc: (batch + 1) int32
-        The cumulative sequence lengths of the sequences in
-        the batch, used to index into sequence. prepended with 0.
-        for example: query_start_loc = torch.Tensor([0,10,16,17]),
-        x.shape=(dim,17)
-    cache_indices: (batch) int32
-        A tensor with each cell is a correspondent
-        input and output ssm_state index
-    has_initial_state: (batch) bool
-        A tensor populated with ones and zeros,
-        indicate if the ssm_state at the corresponding index should be
-        used as initial state. Not providing argument assumes
-        there's no initial state
-    pad_slot_id: int
-        if cache_indices is passed, lets the kernel identify padding entries
-        that will not be processed,
-        for example: cache_indices = [pad_slot_id, 1 ,20 ,pad_slot_id]
-        in this case, the kernel will not process entries at indices 0 and 3
-    returns
-        output: (dim, total_length) for varlen or (batch, dim, seqlen)
-                supports inplace replacement
-    """
-    if u.stride(-1) != 1:
-        u = u.contiguous()
-    if delta.stride(-1) != 1:
-        delta = delta.contiguous()
-    if D is not None:
-        D = D.contiguous()
-    if B.stride(-1) != 1:
-        B = B.contiguous()
-    if C.stride(-1) != 1:
-        C = C.contiguous()
-    if z is not None and z.stride(-1) != 1:
-        z = z.contiguous()
-    if B.dim() == 3 and query_start_loc is None:
-        B = B.unsqueeze(1)
-    if B.dim() == 2 and query_start_loc is not None:
-        B = B.unsqueeze(0)
-    if C.dim() == 3 and query_start_loc is None:
-        C = C.unsqueeze(1)
-    if C.dim() == 2 and query_start_loc is not None:
-        C = C.unsqueeze(0)
-
-    ops.selective_scan_fwd(
-        u,
-        delta,
-        A,
-        B,
-        C,
-        D,
-        z,
-        delta_bias,
-        delta_softplus,
-        query_start_loc,
-        cache_indices,
-        has_initial_state,
-        ssm_states,
-        pad_slot_id,
-    )
-
-    if z is None:
-        return delta  # output written inplace to delta
-    else:
-        return z  # output written inplace to z