[Kernel] Relocate vendored fla and mamba kernel trees to sglang.kernels (RFC #29630, Phase 2.5, 7/7) (#30795)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

benchmark/bench_linear_attention/bench_cutedsl_kda_decode.py

@@ -18,10 +18,10 @@ import torch
 import triton
 
 from sglang.jit_kernel.cutedsl_kda import cutedsl_fused_sigmoid_gating_kda_update
-from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
+from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
     fused_sigmoid_gating_delta_rule_update,
 )
-from sglang.srt.layers.attention.fla.kda import chunk_kda
+from sglang.kernels.ops.attention.fla.kda import chunk_kda
 
 
 def make_inputs(

benchmark/bench_linear_attention/bench_fused_gate_cumsum.py

@@ -22,9 +22,9 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python")
 import torch
 import triton
 
-from sglang.srt.layers.attention.fla.cumsum import chunk_local_cumsum
-from sglang.srt.layers.attention.fla.index import prepare_chunk_indices
-from sglang.srt.layers.attention.fla.kda import kda_gate_chunk_cumsum
+from sglang.kernels.ops.attention.fla.cumsum import chunk_local_cumsum
+from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices
+from sglang.kernels.ops.attention.fla.kda import kda_gate_chunk_cumsum
 
 CHUNK_SIZE = 64
 

benchmark/bench_linear_attention/bench_gdn_decode.py

@@ -29,10 +29,10 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "python"))
 import torch
 import triton
 
-from sglang.srt.layers.attention.fla.fused_recurrent import (
+from sglang.kernels.ops.attention.fla.fused_recurrent import (
     fused_recurrent_gated_delta_rule_packed_decode,
 )
-from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
+from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
     fused_sigmoid_gating_delta_rule_update,
 )
 

benchmark/bench_linear_attention/bench_gdn_prefill.py

@@ -29,10 +29,10 @@ from flashinfer.gdn_prefill import (
     chunk_gated_delta_rule as flashinfer_chunk_gated_delta_rule,
 )
 
-from sglang.srt.layers.attention.fla.chunk import (
+from sglang.kernels.ops.attention.fla.chunk import (
     chunk_gated_delta_rule as triton_chunk_gated_delta_rule,
 )
-from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
+from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd
 
 # ---------------------------------------------------------------------------
 # Helpers

benchmark/bench_linear_attention/bench_gdn_prefill_cutedsl.py

@@ -26,14 +26,14 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python")
 
 import torch
 
+from sglang.kernels.ops.attention.fla.chunk import (
+    chunk_gated_delta_rule as triton_chunk_gated_delta_rule,
+)
+from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd
 from sglang.kernels.ops.attention.linear.gdn_blackwell import (
     chunk_gated_delta_rule_cutedsl,
     prepare_metadata_cutedsl,
 )
-from sglang.srt.layers.attention.fla.chunk import (
-    chunk_gated_delta_rule as triton_chunk_gated_delta_rule,
-)
-from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
 
 # ---------------------------------------------------------------------------
 # Helpers (shared shape: pool layout [N, H, K, V] with K-last stride)

benchmark/bench_linear_attention/bench_kda_decode.py

@@ -22,10 +22,10 @@ import argparse
 
 import torch
 
-from sglang.srt.layers.attention.fla.fused_recurrent import (
+from sglang.kernels.ops.attention.fla.fused_recurrent import (
     fused_recurrent_kda_packed_decode,
 )
-from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
+from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
     fused_sigmoid_gating_delta_rule_update,
 )
 

benchmark/bench_linear_attention/bench_kda_prefill_cutedsl.py

@@ -29,6 +29,7 @@ sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "python")
 import torch
 import torch.nn.functional as F
 
+from sglang.kernels.ops.attention.fla.kda import chunk_kda, fused_recurrent_kda
 from sglang.kernels.ops.attention.linear.kda_blackwell import prepare_metadata
 from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_h import (
     kda_h_cutedsl,
@@ -42,7 +43,6 @@ from sglang.kernels.ops.attention.linear.kda_blackwell.kernel_o import (
 from sglang.kernels.ops.attention.linear.kda_blackwell.prologue import (
     kda_prologue,
 )
-from sglang.srt.layers.attention.fla.kda import chunk_kda, fused_recurrent_kda
 
 BT = 64  # chunk size
 

benchmark/fla/benchmark_layernorm_gated.py

@@ -4,10 +4,10 @@ import numpy as np
 import torch
 
 # Import the function to benchmark
-from sglang.srt.layers.attention.fla.layernorm_gated import (
+from sglang.kernels.ops.attention.fla.layernorm_gated import (
     _layer_norm_fwd as layer_norm_fwd,
 )
-from sglang.srt.layers.attention.fla.layernorm_gated import (
+from sglang.kernels.ops.attention.fla.layernorm_gated import (
     rms_norm_ref,
 )
 

benchmark/kernels/decoding_attention_triton/triton_flashinfer_cudnn.py

@@ -7,7 +7,7 @@ import torch.utils.benchmark as benchmark
 from flashinfer import BatchDecodeWithPagedKVCacheWrapper
 
 from sglang.kernels.ops.attention.decode_attention import decode_attention_fwd
-from sglang.srt.layers.attention.flashinfer_backend import should_use_tensor_core
+from sglang.kernels.ops.attention.flashinfer_backend import should_use_tensor_core
 
 
 def benchmark_forward(

python/sglang/kernels/ops/attention/__init__.py

@@ -43,6 +43,22 @@ del _mod, _fn
 __all__ = []
 
 
+# Vendored linear-attention (flash-linear-attention port) kernels relocated
+# in Phase 2.5 (RFC #29630); representative entry points for inventory.
+for _mod, _fn in [
+    ("fla.chunk", "chunk_gated_delta_rule"),
+    ("fla.fused_recurrent", "fused_recurrent_gated_delta_rule"),
+    ("fla.kda", "fused_recurrent_kda_fwd"),
+]:
+    register_kernel(
+        KernelSpec(
+            op=f"attention.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.attention.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn
+
 # Linear-attention / MiniMax-sparse / diffusion kernels migrated in Phase 2.5
 # (RFC #29630); registered for inventory.
 for _grp, _mod, _fn in [

python/sglang/kernels/ops/attention/fla/bench_gdn_replayssm_decode.py

@@ -15,7 +15,7 @@ traffic ratio is reported per L.
 
 Run::
 
-    python -m sglang.srt.layers.attention.fla.bench_gdn_replayssm_decode
+    python -m sglang.kernels.ops.attention.fla.bench_gdn_replayssm_decode
 
 Requires a GPU (Triton).
 """
@@ -27,10 +27,10 @@ import argparse
 import torch
 import triton
 
-from sglang.srt.layers.attention.fla.fused_recurrent import (
+from sglang.kernels.ops.attention.fla.fused_recurrent import (
     fused_recurrent_gated_delta_rule_packed_decode,
 )
-from sglang.srt.layers.attention.fla.fused_recurrent_linear_replayssm import (
+from sglang.kernels.ops.attention.fla.fused_recurrent_linear_replayssm import (
     fused_recurrent_gdn_replayssm_decode,
 )
 

python/sglang/kernels/ops/attention/fla/chunk.py

@@ -7,15 +7,15 @@ from typing import Optional
 import torch
 from einops import rearrange
 
-from sglang.srt.layers.attention.fla.chunk_delta_h import chunk_gated_delta_rule_fwd_h
-from sglang.srt.layers.attention.fla.chunk_fwd import chunk_gated_delta_rule_fwd_intra
-from sglang.srt.layers.attention.fla.chunk_o import chunk_fwd_o
-from sglang.srt.layers.attention.fla.cumsum import chunk_local_cumsum
-from sglang.srt.layers.attention.fla.index import (
+from sglang.kernels.ops.attention.fla.chunk_delta_h import chunk_gated_delta_rule_fwd_h
+from sglang.kernels.ops.attention.fla.chunk_fwd import chunk_gated_delta_rule_fwd_intra
+from sglang.kernels.ops.attention.fla.chunk_o import chunk_fwd_o
+from sglang.kernels.ops.attention.fla.cumsum import chunk_local_cumsum
+from sglang.kernels.ops.attention.fla.index import (
     prepare_chunk_indices,
 )
-from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
-from sglang.srt.layers.attention.fla.utils import (
+from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd
+from sglang.kernels.ops.attention.fla.utils import (
     SUPPRESS_LEVEL,
     autocast_custom_fwd,
     input_guard,

python/sglang/kernels/ops/attention/fla/chunk_delta_h.py

@@ -9,12 +9,12 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.fla.index import (
+from sglang.kernels.ops.attention.fla.index import (
     prepare_chunk_indices,
     prepare_chunk_offsets,
 )
-from sglang.srt.layers.attention.fla.op import exp, safe_exp
-from sglang.srt.layers.attention.fla.utils import (
+from sglang.kernels.ops.attention.fla.op import exp, safe_exp
+from sglang.kernels.ops.attention.fla.utils import (
     autotune_cache_kwargs,
     is_nvidia_hopper,
 )

python/sglang/kernels/ops/attention/fla/chunk_fwd.py

@@ -5,13 +5,13 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.fla.index import prepare_chunk_indices
-from sglang.srt.layers.attention.fla.op import safe_exp
-from sglang.srt.layers.attention.fla.utils import (
+from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices
+from sglang.kernels.ops.attention.fla.op import safe_exp
+from sglang.kernels.ops.attention.fla.utils import (
     autotune_cache_kwargs,
     is_tf32_supported,
 )
-from sglang.srt.layers.attention.fla.wy_fast import recompute_w_u_fwd
+from sglang.kernels.ops.attention.fla.wy_fast import recompute_w_u_fwd
 
 # TF32 for the block-merge dot products (16x16 matmuls) is safe and ~2x faster on SM90.
 # The numerically sensitive forward-substitution uses scalar ops, not tl.dot.

python/sglang/kernels/ops/attention/fla/chunk_intra.py

@@ -5,14 +5,14 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.fla.chunk_intra_token_parallel import (
+from sglang.kernels.ops.attention.fla.chunk_intra_token_parallel import (
     chunk_kda_fwd_intra_token_parallel,
 )
-from sglang.srt.layers.attention.fla.index import (
+from sglang.kernels.ops.attention.fla.index import (
     prepare_chunk_indices,
 )
-from sglang.srt.layers.attention.fla.op import exp, exp2, gather
-from sglang.srt.layers.attention.fla.utils import (
+from sglang.kernels.ops.attention.fla.op import exp, exp2, gather
+from sglang.kernels.ops.attention.fla.utils import (
     autotune_cache_kwargs,
     is_gather_supported,
     is_tf32_supported,
@@ -1036,7 +1036,7 @@ def chunk_kda_fwd_intra(
         FUSE_DIAGONAL=fuse_diagonal,
     )
 
-    from sglang.srt.layers.attention.fla.kda import (
+    from sglang.kernels.ops.attention.fla.kda import (
         recompute_w_u_fwd as kda_recompute_w_u_fwd,
     )
 

python/sglang/kernels/ops/attention/fla/chunk_intra_token_parallel.py

@@ -6,8 +6,8 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.fla.op import exp2
-from sglang.srt.layers.attention.fla.utils import autotune_cache_kwargs
+from sglang.kernels.ops.attention.fla.op import exp2
+from sglang.kernels.ops.attention.fla.utils import autotune_cache_kwargs
 
 
 @triton.heuristics(

python/sglang/kernels/ops/attention/fla/chunk_o.py

@@ -8,9 +8,9 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.fla.index import prepare_chunk_indices
-from sglang.srt.layers.attention.fla.op import exp, safe_exp
-from sglang.srt.layers.attention.fla.utils import check_shared_mem, is_nvidia_hopper
+from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices
+from sglang.kernels.ops.attention.fla.op import exp, safe_exp
+from sglang.kernels.ops.attention.fla.utils import check_shared_mem, is_nvidia_hopper
 
 BKV_LIST = [64, 128] if check_shared_mem() else [32, 64]
 NUM_WARPS = [2, 4] if is_nvidia_hopper else [2, 4, 8]

python/sglang/kernels/ops/attention/fla/chunk_scaled_dot_kkt.py

@@ -8,8 +8,8 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.fla.index import prepare_chunk_indices
-from sglang.srt.layers.attention.fla.op import safe_exp
+from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices
+from sglang.kernels.ops.attention.fla.op import safe_exp
 
 
 # @triton.autotune(

python/sglang/kernels/ops/attention/fla/cumsum.py

@@ -8,8 +8,8 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.fla.index import prepare_chunk_indices
-from sglang.srt.layers.attention.fla.utils import check_shared_mem, input_guard
+from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices
+from sglang.kernels.ops.attention.fla.utils import check_shared_mem, input_guard
 
 BS_LIST = [32, 64] if check_shared_mem() else [16, 32]
 

python/sglang/kernels/ops/attention/fla/fused_recurrent.py

@@ -8,8 +8,8 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.fla.op import exp
-from sglang.srt.layers.attention.fla.utils import input_guard
+from sglang.kernels.ops.attention.fla.op import exp
+from sglang.kernels.ops.attention.fla.utils import input_guard
 
 
 @triton.jit(do_not_specialize=["T"])

python/sglang/kernels/ops/attention/fla/index.py

@@ -5,7 +5,7 @@
 import torch
 import triton
 
-from sglang.srt.layers.attention.fla.utils import tensor_cache
+from sglang.kernels.ops.attention.fla.utils import tensor_cache
 
 
 @tensor_cache

python/sglang/kernels/ops/attention/fla/kda.py

@@ -12,19 +12,19 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.fla.chunk_delta_h import chunk_gated_delta_rule_fwd_h
-from sglang.srt.layers.attention.fla.chunk_intra import chunk_kda_fwd_intra
-from sglang.srt.layers.attention.fla.cumsum import chunk_local_cumsum
-from sglang.srt.layers.attention.fla.fused_norm_gate import layer_norm_gated_fwd
-from sglang.srt.layers.attention.fla.fused_recurrent import (
+from sglang.kernels.ops.attention.fla.chunk_delta_h import chunk_gated_delta_rule_fwd_h
+from sglang.kernels.ops.attention.fla.chunk_intra import chunk_kda_fwd_intra
+from sglang.kernels.ops.attention.fla.cumsum import chunk_local_cumsum
+from sglang.kernels.ops.attention.fla.fused_norm_gate import layer_norm_gated_fwd
+from sglang.kernels.ops.attention.fla.fused_recurrent import (
     fused_recurrent_gated_delta_rule_fwd_kernel,
 )
-from sglang.srt.layers.attention.fla.index import (
+from sglang.kernels.ops.attention.fla.index import (
     prepare_chunk_indices,
 )
-from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
-from sglang.srt.layers.attention.fla.op import exp, log
-from sglang.srt.layers.attention.fla.utils import (
+from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd
+from sglang.kernels.ops.attention.fla.op import exp, log
+from sglang.kernels.ops.attention.fla.utils import (
     check_shared_mem,
     is_intel,
 )

python/sglang/kernels/ops/attention/fla/l2norm.py

@@ -9,7 +9,7 @@ import torch.nn as nn
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.fla.utils import input_guard
+from sglang.kernels.ops.attention.fla.utils import input_guard
 
 BT_LIST = [8, 16, 32, 64, 128]
 

python/sglang/kernels/ops/attention/fla/op.py

@@ -8,7 +8,7 @@ import triton
 import triton.language as tl
 import triton.language.extra.libdevice as tldevice
 
-from sglang.srt.layers.attention.fla.utils import is_gather_supported
+from sglang.kernels.ops.attention.fla.utils import is_gather_supported
 
 if os.environ.get("FLA_USE_FAST_OPS", "0") == "1":
     exp = tldevice.fast_expf

python/sglang/kernels/ops/attention/fla/solve_tril.py

@@ -8,8 +8,8 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.fla.index import prepare_chunk_indices
-from sglang.srt.layers.attention.fla.utils import input_guard
+from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices
+from sglang.kernels.ops.attention.fla.utils import input_guard
 
 
 # @triton.autotune(

python/sglang/kernels/ops/attention/fla/wy_fast.py

@@ -8,7 +8,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.fla.index import prepare_chunk_indices
+from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices
 
 
 # @triton.autotune(

python/sglang/kernels/ops/attention/linear/kda_blackwell/__init__.py

@@ -163,7 +163,7 @@ def chunk_kda_cutedsl(
     # injected through the cutedsl KKT/Aqk MMAs as an identity-right-operand pass:
     # with kL'=M (M in the first 64 K-slots) and kR'=onehot(chunk-pos), the MMA
     # kL'@kR'.T == M, so kkt_inv_uw/kernel_o see the correct matrix without overflow.
-    from sglang.srt.layers.attention.fla.kda import chunk_kda_scaled_dot_kkt_fwd
+    from sglang.kernels.ops.attention.fla.kda import chunk_kda_scaled_dot_kkt_fwd
 
     ones_beta = q.new_ones(1, T, Hv, dtype=torch.float32)
     M_kk, M_qk = chunk_kda_scaled_dot_kkt_fwd(

python/sglang/kernels/ops/mamba/__init__.py

@@ -84,3 +84,20 @@ def causal_conv1d_update(
 
 
 __all__ = ["causal_conv1d_fwd", "causal_conv1d_update"]
+
+
+# Vendored mamba_ssm-derived kernels relocated in Phase 2.5 (RFC #29630).
+for _mod, _fn in [
+    ("triton_ops.ssd_combined", "mamba_chunk_scan_combined"),
+    ("triton_ops.mamba_ssm", "selective_state_update"),
+    ("causal_conv1d_triton", "causal_conv1d_fn"),
+    ("mamba_state_scatter_triton", "fused_mamba_state_scatter_with_mask"),
+]:
+    register_kernel(
+        KernelSpec(
+            op=f"mamba.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.mamba.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn

python/sglang/kernels/ops/mamba/triton_ops/ssu_dispatch.py

@@ -51,7 +51,7 @@ class TritonSSUBackend(MambaSSUBackend):
         enable_stochastic_rounding: bool = False,
         cache_philox_rounds: int = 0,
     ) -> None:
-        from sglang.srt.layers.attention.mamba.ops.mamba_ssm import (
+        from sglang.kernels.ops.mamba.triton_ops.mamba_ssm import (
             selective_state_update,
         )
 

python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_delta_h.py

@@ -4,12 +4,12 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.fla.index import (
+from sglang.kernels.ops.attention.fla.index import (
     prepare_chunk_indices,
     prepare_chunk_offsets,
 )
-from sglang.srt.layers.attention.fla.op import exp, make_tensor_descriptor, safe_exp
-from sglang.srt.layers.attention.fla.utils import (
+from sglang.kernels.ops.attention.fla.op import exp, make_tensor_descriptor, safe_exp
+from sglang.kernels.ops.attention.fla.utils import (
     autotune_cache_kwargs,
 )
 

python/sglang/srt/hardware_backend/xpu/kernels/fla/chunk_fwd.py

@@ -2,12 +2,12 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.attention.fla.index import prepare_chunk_indices
-from sglang.srt.layers.attention.fla.op import safe_exp
-from sglang.srt.layers.attention.fla.utils import (
+from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices
+from sglang.kernels.ops.attention.fla.op import safe_exp
+from sglang.kernels.ops.attention.fla.utils import (
     autotune_cache_kwargs,
 )
-from sglang.srt.layers.attention.fla.wy_fast import recompute_w_u_fwd
+from sglang.kernels.ops.attention.fla.wy_fast import recompute_w_u_fwd
 
 _MERGE_DOT_PRECISION = tl.constexpr("ieee")
 

python/sglang/srt/hardware_backend/xpu/kernels/fla/fused_sigmoid_gating_recurrent.py

@@ -3,7 +3,7 @@ from typing import Optional
 import torch
 import triton
 
-from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
+from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
     fused_sigmoid_gating_delta_rule_update_kernel,
 )
 

python/sglang/srt/layers/attention/attention_registry.py

@@ -43,7 +43,9 @@ def create_flashinfer_backend(runner):
     import torch
 
     if not runner.use_mla_backend:
-        from sglang.srt.layers.attention.flashinfer_backend import FlashInferAttnBackend
+        from sglang.srt.layers.attention.flashinfer_backend import (
+            FlashInferAttnBackend,
+        )
 
         # Init streams
         if runner.server_args.speculative_algorithm == "EAGLE":
@@ -219,7 +221,9 @@ def create_flashattention_v3_backend(runner):
 
 @register_attention_backend("fa4")
 def create_flashattention_v4_backend(runner):
-    from sglang.srt.layers.attention.flashattention_backend import FlashAttentionBackend
+    from sglang.srt.layers.attention.flashattention_backend import (
+        FlashAttentionBackend,
+    )
 
     return FlashAttentionBackend(runner, fa_impl_ver=4)
 
@@ -279,7 +283,7 @@ def attn_backend_wrapper(runner: "ModelRunner", full_attn_backend: "AttentionBac
         )
 
     if cfg := mambaish_config(runner.model_config):
-        from sglang.srt.layers.attention.fla.utils import check_environments
+        from sglang.kernels.ops.attention.fla.utils import check_environments
         from sglang.srt.layers.attention.linear.kda_backend import KDAAttnBackend
         from sglang.srt.layers.attention.linear.lightning_backend import (
             LightningAttentionBackend,

python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py

@@ -3,19 +3,19 @@ from typing import Optional, Union
 
 import torch
 
+from sglang.kernels.ops.mamba.causal_conv1d_triton import PAD_SLOT_ID
+from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
+    fused_conv_window_scatter_with_mask,
+    fused_mamba_state_scatter_with_mask,
+    track_mamba_states_if_needed,
+)
 from sglang.srt.configs.hybrid_arch import mamba2_config
 from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
-from sglang.srt.layers.attention.mamba.causal_conv1d_triton import PAD_SLOT_ID
 from sglang.srt.layers.attention.mamba.mamba import MambaMixer2
 from sglang.srt.layers.attention.mamba.mamba2_metadata import (
     ForwardMetadata,
     Mamba2Metadata,
 )
-from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import (
-    fused_conv_window_scatter_with_mask,
-    fused_mamba_state_scatter_with_mask,
-    track_mamba_states_if_needed,
-)
 from sglang.srt.layers.radix_attention import RadixAttention
 from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode

python/sglang/srt/layers/attention/linear/gdn_backend.py

@@ -2,8 +2,12 @@ from typing import Optional, Tuple, Union
 
 import torch
 
+from sglang.kernels.ops.attention.fla.fused_gdn_gating import fused_gdn_gating
+from sglang.kernels.ops.mamba.causal_conv1d_triton import (
+    causal_conv1d_fn,
+    causal_conv1d_update,
+)
 from sglang.srt.configs.hybrid_arch import hybrid_gdn_config
-from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating
 from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase
 from sglang.srt.layers.attention.linear.kernels.gdn_triton import TritonGDNKernel
 from sglang.srt.layers.attention.linear.utils import (
@@ -11,10 +15,6 @@ from sglang.srt.layers.attention.linear.utils import (
     get_linear_attn_decode_backend,
     get_linear_attn_prefill_backend,
 )
-from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
-    causal_conv1d_fn,
-    causal_conv1d_update,
-)
 from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
 from sglang.srt.mem_cache.memory_pool import MambaPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -23,7 +23,7 @@ from sglang.srt.utils import is_cpu, is_cuda, is_hip, is_npu
 from sglang.srt.utils.common import rank0_log
 
 if not is_cpu():
-    from sglang.srt.layers.attention.fla.chunk_delta_h import (
+    from sglang.kernels.ops.attention.fla.chunk_delta_h import (
         CHUNK_SIZE as FLA_CHUNK_SIZE,
     )
 

python/sglang/srt/layers/attention/linear/kda_backend.py

@@ -2,6 +2,10 @@ from typing import Optional, Tuple, Union
 
 import torch
 
+from sglang.kernels.ops.mamba.causal_conv1d_triton import (
+    causal_conv1d_fn,
+    causal_conv1d_update,
+)
 from sglang.srt.layers.attention.hybrid_linear_attn_backend import MambaAttnBackendBase
 from sglang.srt.layers.attention.linear.kernels.kda_triton import TritonKDAKernel
 from sglang.srt.layers.attention.linear.utils import (
@@ -9,10 +13,6 @@ from sglang.srt.layers.attention.linear.utils import (
     get_linear_attn_decode_backend,
     get_linear_attn_prefill_backend,
 )
-from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
-    causal_conv1d_fn,
-    causal_conv1d_update,
-)
 from sglang.srt.layers.radix_linear_attention import RadixLinearAttention
 from sglang.srt.utils import is_cpu, is_cuda, is_npu
 from sglang.srt.utils.common import rank0_log

python/sglang/srt/layers/attention/linear/kernels/gdn_cutedsl.py

@@ -65,11 +65,11 @@ class CuteDSLGDNKernel(LinearAttnKernelBase):
             raise RuntimeError(
                 f"CuTe DSL GDN prefill requires head_k_dim=128, got {head_k_dim}."
             )
+        from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd
         from sglang.kernels.ops.attention.linear.gdn_blackwell import (
             chunk_gated_delta_rule_cutedsl,
             prepare_metadata_cutedsl,
         )
-        from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
 
         self._extend_fn = chunk_gated_delta_rule_cutedsl
         self._prepare_meta_fn = prepare_metadata_cutedsl

python/sglang/srt/layers/attention/linear/kernels/gdn_flashinfer.py

@@ -233,7 +233,7 @@ class FlashInferGDNKernel(LinearAttnKernelBase):
         query_start_loc: torch.Tensor,
         **kwargs,
     ) -> tuple:
-        from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
+        from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd
 
         total_seq_len = q.shape[1]
         num_v_heads = v.shape[2]

python/sglang/srt/layers/attention/linear/kernels/gdn_triton.py

@@ -6,14 +6,14 @@ from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
 from sglang.srt.utils import is_cpu, is_npu, is_xpu
 
 if not is_cpu():
-    from sglang.srt.layers.attention.fla.chunk import chunk_gated_delta_rule
-    from sglang.srt.layers.attention.fla.fused_recurrent import (
+    from sglang.kernels.ops.attention.fla.chunk import chunk_gated_delta_rule
+    from sglang.kernels.ops.attention.fla.fused_recurrent import (
         fused_recurrent_gated_delta_rule_packed_decode,
     )
-    from sglang.srt.layers.attention.fla.fused_recurrent_linear_replayssm import (
+    from sglang.kernels.ops.attention.fla.fused_recurrent_linear_replayssm import (
         fused_recurrent_gdn_replayssm_decode,
     )
-    from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
+    from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
         fused_sigmoid_gating_delta_rule_update,
     )
 

python/sglang/srt/layers/attention/linear/kernels/kda_cutedsl.py

@@ -49,10 +49,10 @@ class CuteDSLKDAKernel(LinearAttnKernelBase):
             raise RuntimeError(
                 f"CuTe DSL KDA prefill requires head_k_dim=128, got {head_k_dim}."
             )
+        from sglang.kernels.ops.attention.fla.l2norm import l2norm_fwd
         from sglang.kernels.ops.attention.linear.kda_blackwell import (
             chunk_kda_cutedsl,
         )
-        from sglang.srt.layers.attention.fla.l2norm import l2norm_fwd
 
         self._extend_fn = chunk_kda_cutedsl
         self._l2norm_fn = l2norm_fwd

python/sglang/srt/layers/attention/linear/kernels/kda_flashkda.py

@@ -47,7 +47,7 @@ def _triton_fallback(
     -- otherwise the fallback silently skips activation. chunk_kda updates the
     ssm state in-place via cache_indices and returns only the output tensor.
     """
-    from sglang.srt.layers.attention.fla.kda import chunk_kda
+    from sglang.kernels.ops.attention.fla.kda import chunk_kda
 
     return chunk_kda(
         q=q,

python/sglang/srt/layers/attention/linear/kernels/kda_triton.py

@@ -8,16 +8,16 @@ from sglang.srt.layers.attention.linear.kernels.kernel_backend import (
 from sglang.srt.utils import is_cpu, is_npu
 
 if not is_cpu():
-    from sglang.srt.layers.attention.fla.fused_recurrent import (
+    from sglang.kernels.ops.attention.fla.fused_recurrent import (
         fused_recurrent_kda_packed_decode,
     )
-    from sglang.srt.layers.attention.fla.fused_recurrent_linear_replayssm import (
+    from sglang.kernels.ops.attention.fla.fused_recurrent_linear_replayssm import (
         fused_recurrent_linear_replayssm_decode,
     )
-    from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
+    from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
         fused_sigmoid_gating_delta_rule_update,
     )
-    from sglang.srt.layers.attention.fla.kda import chunk_kda
+    from sglang.kernels.ops.attention.fla.kda import chunk_kda
 
 
 class TritonKDAKernel(LinearAttnKernelBase):

python/sglang/srt/layers/attention/mamba/causal_conv1d.py

@@ -9,9 +9,15 @@ from typing import Optional
 
 import torch
 
-from .causal_conv1d_triton import PAD_SLOT_ID
-from .causal_conv1d_triton import causal_conv1d_fn as _causal_conv1d_fn_triton
-from .causal_conv1d_triton import causal_conv1d_update as _causal_conv1d_update_triton
+from sglang.kernels.ops.mamba.causal_conv1d_triton import (
+    PAD_SLOT_ID,
+)
+from sglang.kernels.ops.mamba.causal_conv1d_triton import (
+    causal_conv1d_fn as _causal_conv1d_fn_triton,
+)
+from sglang.kernels.ops.mamba.causal_conv1d_triton import (
+    causal_conv1d_update as _causal_conv1d_update_triton,
+)
 
 try:
     from sgl_kernel import causal_conv1d_fwd

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

@@ -4,6 +4,10 @@ from typing import Callable, List, Optional, Tuple
 import torch
 import torch.nn as nn
 
+from sglang.kernels.ops.mamba.triton_ops import (
+    mamba_chunk_scan_combined,
+    selective_state_update,
+)
 from sglang.srt.configs.mamba_utils import (
     Mamba2CacheParams,
     extra_groups_for_head_shards,
@@ -13,10 +17,6 @@ from sglang.srt.distributed import (
 )
 from sglang.srt.layers.attention.mamba.mamba2_metadata import Mamba2Metadata
 from sglang.srt.layers.attention.mamba.mixer2_rms_norm_gated import Mixer2RMSNormGated
-from sglang.srt.layers.attention.mamba.ops import (
-    mamba_chunk_scan_combined,
-    selective_state_update,
-)
 from sglang.srt.layers.dp_attention import (
     is_dp_attention_enabled,
 )
@@ -42,16 +42,16 @@ from sglang.srt.utils import (
 )
 
 if is_cuda():
+    from sglang.kernels.ops.mamba.causal_conv1d_triton import (
+        causal_conv1d_fn as causal_conv1d_fn_triton,
+    )
+    from sglang.kernels.ops.mamba.causal_conv1d_triton import (
+        causal_conv1d_update as causal_conv1d_update_triton,
+    )
     from sglang.srt.layers.attention.mamba.causal_conv1d import (
         causal_conv1d_fn,
         causal_conv1d_update,
     )
-    from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
-        causal_conv1d_fn as causal_conv1d_fn_triton,
-    )
-    from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
-        causal_conv1d_update as causal_conv1d_update_triton,
-    )
 elif is_npu():
     from sgl_kernel_npu.mamba.causal_conv1d import (
         causal_conv1d_fn_npu as causal_conv1d_fn,
@@ -63,16 +63,16 @@ elif is_xpu():
     # XPU has no native causal_conv1d kernel yet; use the portable Triton
     # implementation for both the "native" and the "_triton" entry points so
     # `causal_conv1d_fn` / `causal_conv1d_fn_triton` are always bound on XPU.
-    from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
+    from sglang.kernels.ops.mamba.causal_conv1d_triton import (
         causal_conv1d_fn as causal_conv1d_fn,
     )
-    from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
+    from sglang.kernels.ops.mamba.causal_conv1d_triton import (
         causal_conv1d_fn as causal_conv1d_fn_triton,
     )
-    from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
+    from sglang.kernels.ops.mamba.causal_conv1d_triton import (
         causal_conv1d_update as causal_conv1d_update,
     )
-    from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
+    from sglang.kernels.ops.mamba.causal_conv1d_triton import (
         causal_conv1d_update as causal_conv1d_update_triton,
     )
 

python/sglang/srt/layers/attention/mamba/mixer2_rms_norm_gated.py

@@ -2,11 +2,11 @@ from typing import Union
 
 import torch
 
+from sglang.kernels.ops.attention.fla.layernorm_gated import rms_norm_gated
 from sglang.srt.distributed.communication_op import (
     tensor_model_parallel_all_gather,
     tensor_model_parallel_all_reduce,
 )
-from sglang.srt.layers.attention.fla.layernorm_gated import rms_norm_gated
 from sglang.srt.layers.dp_attention import (
     attn_tp_all_reduce,
     is_dp_attention_enabled,

python/sglang/srt/managers/scheduler.py

@@ -38,6 +38,9 @@ import torch.distributed
 from torch.cuda import Stream as CudaStream
 from torch.distributed import barrier
 
+from sglang.kernels.ops.mamba.triton_ops import (
+    initialize_mamba_selective_state_update_backend,
+)
 from sglang.srt.configs.model_config import ModelConfig, ModelImpl, is_minimax_sparse
 from sglang.srt.constrained.grammar_manager import GrammarManager
 from sglang.srt.debug_utils.pr_fix_toggle import maybe_revert_pr_fix
@@ -69,9 +72,6 @@ from sglang.srt.distributed.parallel_state_wrapper import ParallelState
 from sglang.srt.dllm.mixin.scheduler import SchedulerDllmMixin
 from sglang.srt.environ import envs
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
-from sglang.srt.layers.attention.mamba.ops import (
-    initialize_mamba_selective_state_update_backend,
-)
 from sglang.srt.layers.dp_attention import compute_dp_attention_world_info
 from sglang.srt.layers.moe import initialize_moe_config
 from sglang.srt.layers.quantization.fp4_utils import initialize_fp4_gemm_config

python/sglang/srt/models/bailing_moe_linear.py

@@ -9,6 +9,8 @@ import torch.nn.functional as F
 from torch import nn
 from transformers import PretrainedConfig
 
+from sglang.kernels.ops.attention.fla.layernorm_gated import RMSNorm as RMSNormGated
+from sglang.kernels.ops.attention.fla.layernorm_gated import layernorm_fn
 from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
 from sglang.srt.distributed import (
     get_pp_group,
@@ -17,8 +19,6 @@ from sglang.srt.distributed import (
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.layers import deep_gemm_wrapper
 from sglang.srt.layers.activation import SiluAndMul
-from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGated
-from sglang.srt.layers.attention.fla.layernorm_gated import layernorm_fn
 from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes
 from sglang.srt.layers.dp_attention import (
     is_dp_attention_enabled,

python/sglang/srt/models/jet_nemotron.py

@@ -5,11 +5,11 @@ import einops
 import torch
 import torch.nn as nn
 
-from sglang.srt.configs.jet_nemotron import JetBlockConfig, JetNemotronConfig
-from sglang.srt.layers.attention.fla.fused_recurrent import (
+from sglang.kernels.ops.attention.fla.fused_recurrent import (
     fused_recurrent_gated_delta_rule_update,
 )
-from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGated
+from sglang.kernels.ops.attention.fla.layernorm_gated import RMSNorm as RMSNormGated
+from sglang.srt.configs.jet_nemotron import JetBlockConfig, JetNemotronConfig
 from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
     HybridLinearAttnBackend,
     MambaAttnBackendBase,

python/sglang/srt/models/kimi_linear.py

@@ -8,6 +8,7 @@ from typing import Optional
 import torch
 from torch import nn
 
+from sglang.kernels.ops.attention.fla.fused_norm_gate import FusedRMSNormGated
 from sglang.srt.configs.kimi_linear import KimiLinearConfig
 from sglang.srt.distributed import (
     divide,
@@ -15,7 +16,6 @@ from sglang.srt.distributed import (
     tensor_model_parallel_all_reduce,
 )
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
-from sglang.srt.layers.attention.fla.fused_norm_gate import FusedRMSNormGated
 from sglang.srt.layers.layernorm import RMSNorm
 from sglang.srt.layers.linear import (
     ColumnParallelBatchedLinear,

python/sglang/srt/models/qwen3_5.py

@@ -25,6 +25,9 @@ import triton
 from sglang.jit_kernel.triton.gdn_fused_proj import (
     fused_qkvzba_split_reshape_cat_contiguous,
 )
+
+# Layers - Attention
+from sglang.kernels.ops.attention.fla.layernorm_gated import RMSNorm as RMSNormGated
 from sglang.kernels.ops.layernorm.elementwise import fused_sigmoid_mul
 
 # Configs
@@ -38,9 +41,6 @@ from sglang.srt.configs.qwen3_5 import (
 from sglang.srt.distributed import get_pp_group
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
-
-# Layers - Attention
-from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGated
 from sglang.srt.layers.attention.mamba.mamba import mamba_v2_sharded_weight_loader
 from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes
 from sglang.srt.layers.dp_attention import (

python/sglang/srt/models/qwen3_next.py

@@ -7,12 +7,12 @@ import triton
 from torch import nn
 
 from sglang.jit_kernel.triton.gdn_fused_proj import fused_qkvzba_split_reshape_cat
+from sglang.kernels.ops.attention.fla.fused_norm_gate import FusedRMSNormGated
+from sglang.kernels.ops.attention.fla.layernorm_gated import RMSNorm as RMSNormGated
 from sglang.srt.configs.qwen3_next import Qwen3NextConfig
 from sglang.srt.distributed import get_pp_group
 from sglang.srt.eplb.expert_distribution import get_global_expert_distribution_recorder
 from sglang.srt.eplb.expert_location import ModelConfigForExpertLocation
-from sglang.srt.layers.attention.fla.fused_norm_gate import FusedRMSNormGated
-from sglang.srt.layers.attention.fla.layernorm_gated import RMSNorm as RMSNormGated
 from sglang.srt.layers.attention.mamba.mamba import mamba_v2_sharded_weight_loader
 from sglang.srt.layers.communicator import LayerCommunicator, LayerScatterModes
 from sglang.srt.layers.dp_attention import (

python/sglang/srt/server_args.py

@@ -32,6 +32,7 @@ from functools import cached_property
 from typing import Any, Callable, Dict, List, Literal, Optional, Union
 
 from sglang.jit_kernel.kv_canary.consts import RealKvHashMode
+from sglang.kernels.ops.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE
 from sglang.srt.arg_groups.arg_utils import A, Arg, add_cli_args_from_dataclass
 from sglang.srt.arg_groups.argparse_actions import (
     DeprecatedAction,
@@ -47,7 +48,6 @@ from sglang.srt.distributed.device_communicators.mooncake_transfer_engine import
 )
 from sglang.srt.environ import envs
 from sglang.srt.function_call.function_call_parser import FunctionCallParser
-from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE
 from sglang.srt.lora.lora_registry import LoRARef
 from sglang.srt.model_executor.cuda_graph_config import (
     ALLOWED_BACKENDS_PER_PHASE,

python/sglang/test/kits/attention_unittest/attention_methods/mamba2_attention.py

@@ -389,7 +389,7 @@ class MockMamba2ModelRunner(ModelRunner):
         # `MambaMixer2.forward_decode` requires. In production the
         # scheduler calls this during initialization; the fixture must
         # mirror that or DECODE crashes with a missing-backend error.
-        from sglang.srt.layers.attention.mamba.ops import (
+        from sglang.kernels.ops.mamba.triton_ops import (
             initialize_mamba_selective_state_update_backend,
         )
 

test/registered/attention/test_chunk_gated_delta_rule.py

@@ -2,8 +2,8 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.attention.fla.chunk import chunk_gated_delta_rule
-from sglang.srt.layers.attention.fla.fused_recurrent import (
+from sglang.kernels.ops.attention.fla.chunk import chunk_gated_delta_rule
+from sglang.kernels.ops.attention.fla.fused_recurrent import (
     fused_recurrent_gated_delta_rule,
 )
 from sglang.srt.utils import get_device

test/registered/attention/test_gdn_noncontiguous_stride.py

@@ -8,8 +8,8 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating
-from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
+from sglang.kernels.ops.attention.fla.fused_gdn_gating import fused_gdn_gating
+from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
     fused_sigmoid_gating_delta_rule_update,
 )
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci

test/registered/attention/test_gdn_prefill_cutedsl.py

@@ -23,17 +23,17 @@ if not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 1
         allow_module_level=True,
     )
 
+from sglang.kernels.ops.attention.fla.fused_recurrent import (  # noqa: E402
+    fused_recurrent_gated_delta_rule,
+)
+from sglang.kernels.ops.attention.fla.index import (  # noqa: E402
+    prepare_chunk_indices,
+    prepare_chunk_offsets,
+)
 from sglang.kernels.ops.attention.linear.gdn_blackwell import (  # noqa: E402
     chunk_gated_delta_rule_cutedsl,
     prepare_metadata_cutedsl,
 )
-from sglang.srt.layers.attention.fla.fused_recurrent import (  # noqa: E402
-    fused_recurrent_gated_delta_rule,
-)
-from sglang.srt.layers.attention.fla.index import (  # noqa: E402
-    prepare_chunk_indices,
-    prepare_chunk_offsets,
-)
 
 
 @pytest.mark.parametrize("num_seqs", [1, 5, 257])

test/registered/attention/test_kda_kernels.py

@@ -2,15 +2,15 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.attention.fla.cumsum import chunk_local_cumsum
-from sglang.srt.layers.attention.fla.fused_recurrent import (
+from sglang.kernels.ops.attention.fla.cumsum import chunk_local_cumsum
+from sglang.kernels.ops.attention.fla.fused_recurrent import (
     fused_recurrent_kda_packed_decode,
 )
-from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
+from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
     fused_sigmoid_gating_delta_rule_update,
 )
-from sglang.srt.layers.attention.fla.index import prepare_chunk_indices
-from sglang.srt.layers.attention.fla.kda import (
+from sglang.kernels.ops.attention.fla.index import prepare_chunk_indices
+from sglang.kernels.ops.attention.fla.kda import (
     fused_recurrent_kda,
     kda_gate_chunk_cumsum,
 )

test/registered/attention/test_kda_prefill_cutedsl.py

@@ -25,15 +25,15 @@ if not (torch.cuda.is_available() and torch.cuda.get_device_capability()[0] >= 1
         allow_module_level=True,
     )
 
+from sglang.kernels.ops.attention.fla.index import (  # noqa: E402
+    prepare_chunk_indices,
+    prepare_chunk_offsets,
+)
+from sglang.kernels.ops.attention.fla.kda import fused_recurrent_kda  # noqa: E402
 from sglang.kernels.ops.attention.linear.kda_blackwell import (  # noqa: E402
     chunk_kda_cutedsl,
     prepare_metadata,
 )
-from sglang.srt.layers.attention.fla.index import (  # noqa: E402
-    prepare_chunk_indices,
-    prepare_chunk_offsets,
-)
-from sglang.srt.layers.attention.fla.kda import fused_recurrent_kda  # noqa: E402
 
 
 def _l2norm(x: torch.Tensor) -> torch.Tensor:

test/registered/attention/test_kda_prefill_flashkda.py

@@ -39,7 +39,7 @@ except ImportError:
         allow_module_level=True,
     )
 
-from sglang.srt.layers.attention.fla.kda import chunk_kda  # noqa: E402
+from sglang.kernels.ops.attention.fla.kda import chunk_kda  # noqa: E402
 from sglang.srt.layers.attention.linear.kernels.kda_flashkda import (  # noqa: E402
     FlashKDAKernel,
 )

test/registered/attention/unittests/gdn/test_linear_replayssm_decode.py

@@ -98,11 +98,11 @@ class TestLinearReplaySSMDecode(CustomTestCase):
     L_SWEEP = (1, 4, 8, 16)
 
     def _run_one(self, cfg, L, dtype, force_flush_steps=(), is_kda=False):
-        from sglang.srt.layers.attention.fla.fused_recurrent import (
+        from sglang.kernels.ops.attention.fla.fused_recurrent import (
             fused_recurrent_gated_delta_rule_packed_decode,
             fused_recurrent_kda_packed_decode,
         )
-        from sglang.srt.layers.attention.fla.fused_recurrent_linear_replayssm import (
+        from sglang.kernels.ops.attention.fla.fused_recurrent_linear_replayssm import (
             fused_recurrent_linear_replayssm_decode,
         )
 

test/registered/jit/test_cutedsl_gdn.py

@@ -21,7 +21,7 @@ except ImportError:
     cutedsl_gdn = None
 
 try:
-    from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
+    from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
         fused_sigmoid_gating_delta_rule_update,
     )
 

test/registered/jit/test_fused_verify_triton_gdn.py

@@ -14,11 +14,11 @@ import torch
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 
 try:
-    from sglang.srt.layers.attention.fla.fused_gdn_gating import fused_gdn_gating
-    from sglang.srt.layers.attention.fla.fused_recurrent import (
+    from sglang.kernels.ops.attention.fla.fused_gdn_gating import fused_gdn_gating
+    from sglang.kernels.ops.attention.fla.fused_recurrent import (
         fused_recurrent_gated_delta_rule_update,
     )
-    from sglang.srt.layers.attention.fla.fused_sigmoid_gating_recurrent import (
+    from sglang.kernels.ops.attention.fla.fused_sigmoid_gating_recurrent import (
         fused_sigmoid_gating_delta_rule_update,
     )
 

test/registered/layers/mamba/conftest.py

@@ -1,7 +1,7 @@
 import pytest
 
-from sglang.srt.layers.attention.mamba.ops import ssu_dispatch
-from sglang.srt.layers.attention.mamba.ops.ssu_dispatch import (
+from sglang.kernels.ops.mamba.triton_ops import ssu_dispatch
+from sglang.kernels.ops.mamba.triton_ops.ssu_dispatch import (
     initialize_mamba_selective_state_update_backend,
 )
 from sglang.srt.server_args import ServerArgs

test/registered/layers/mamba/test_causal_conv1d.py

@@ -15,7 +15,7 @@ import torch
 import torch.nn.functional as F
 from einops import rearrange
 
-from sglang.srt.layers.attention.mamba.causal_conv1d_triton import (
+from sglang.kernels.ops.mamba.causal_conv1d_triton import (
     PAD_SLOT_ID,
     causal_conv1d_fn,
     causal_conv1d_update,

test/registered/layers/mamba/test_mamba_ssm.py

@@ -13,8 +13,8 @@ import torch
 import torch.nn.functional as F
 from einops import rearrange, repeat
 
-from sglang.srt.layers.attention.mamba.causal_conv1d_triton import PAD_SLOT_ID
-from sglang.srt.layers.attention.mamba.ops.mamba_ssm import selective_state_update
+from sglang.kernels.ops.mamba.causal_conv1d_triton import PAD_SLOT_ID
+from sglang.kernels.ops.mamba.triton_ops.mamba_ssm import selective_state_update
 from sglang.srt.utils import get_device, is_sm100_supported
 
 

test/registered/layers/mamba/test_mamba_ssm_ssd.py

@@ -14,8 +14,8 @@ import torch
 import torch.nn.functional as F
 from einops import rearrange, repeat
 
+from sglang.kernels.ops.mamba.triton_ops import mamba_chunk_scan_combined
 from sglang.srt.layers.attention.mamba.mamba2_metadata import Mamba2Metadata
-from sglang.srt.layers.attention.mamba.ops import mamba_chunk_scan_combined
 from sglang.srt.utils import get_device
 from sglang.srt.utils.common import is_hip
 from sglang.utils import is_in_ci

test/registered/layers/test_fla_layernorm_guard.py

@@ -8,10 +8,10 @@ import pytest
 import torch
 import torch.nn.functional as F
 
-from sglang.srt.layers.attention.fla.layernorm_gated import (
+from sglang.kernels.ops.attention.fla.layernorm_gated import (
     _layer_norm_fwd as layer_norm_fwd,
 )
-from sglang.srt.layers.attention.fla.layernorm_gated import (
+from sglang.kernels.ops.attention.fla.layernorm_gated import (
     layernorm_fn,
     rms_norm_ref,
 )

test/registered/mem_cache/test_int8_checkpoint_store.py

@@ -126,7 +126,7 @@ class TestInt8CheckpointCodec(unittest.TestCase):
 class TestInt8CheckpointDecodeError(unittest.TestCase):
     def test_decode_error_within_bound(self):
         try:
-            from sglang.srt.layers.attention.fla.kda import fused_recurrent_kda
+            from sglang.kernels.ops.attention.fla.kda import fused_recurrent_kda
         except (ImportError, ModuleNotFoundError) as e:
             self.skipTest(f"fla kernels unavailable: {e}")
 

test/registered/unit/layers/test_mamba_state_scatter_triton.py

@@ -8,7 +8,7 @@ import unittest
 import torch
 
 try:
-    from sglang.srt.layers.attention.mamba.mamba_state_scatter_triton import (
+    from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
         fused_mamba_state_scatter_with_mask,
     )
 

test/registered/unit/mem_cache/test_mamba_unittest.py

@@ -3,10 +3,10 @@ from array import array
 
 import torch
 
+from sglang.kernels.ops.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE
 from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
 from sglang.srt.disaggregation.kv_events import BlockRemoved, BlockStored
 from sglang.srt.environ import envs
-from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import (

test/registered/unit/mem_cache/test_unified_radix_cache_bench.py

@@ -19,9 +19,9 @@ from typing import Callable
 
 import torch
 
+from sglang.kernels.ops.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE
 from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
 from sglang.srt.environ import envs
-from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE
 from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import (
     DecLockRefParams,

test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py

@@ -12,6 +12,7 @@ from unittest import mock
 
 import torch
 
+from sglang.kernels.ops.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE
 from sglang.srt.configs.mamba_utils import Mamba2CacheParams, Mamba2StateShape
 from sglang.srt.disaggregation.kv_events import (
     BlockRemoved,
@@ -19,7 +20,6 @@ from sglang.srt.disaggregation.kv_events import (
     StorageMedium,
 )
 from sglang.srt.environ import envs
-from sglang.srt.layers.attention.fla.chunk_delta_h import CHUNK_SIZE as FLA_CHUNK_SIZE
 from sglang.srt.managers.schedule_batch import Req
 from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
 from sglang.srt.mem_cache.allocator.swa import SWATokenToKVPoolAllocator