[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