[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>

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