[Kernel] Sweep decoupled scattered kernels into sglang.kernels.ops (RFC #29630) (#31582)

Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

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

@@ -29,6 +29,9 @@ _TRITON_KERNELS = [
     ("pad", "unpad_draft_extend_output"),
     ("pad", "seqlens_expand_triton"),
     ("position", "compute_position_triton"),
+    ("dsv4_attn_metadata_kernels", "expand_prefill_causally"),
+    ("dsv4_attn_metadata_kernels", "build_page_table_positions"),
+    ("dsv4_attn_metadata_kernels", "build_causal_swa_page_indices"),
 ]
 for _mod, _fn in _TRITON_KERNELS:
     register_kernel(

python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_h.py

@@ -10,7 +10,7 @@ from cutlass import BFloat16, Float32, Int32, Int64, Uint32, cute
 from cutlass.cute.nvgpu import cpasync, warp
 from quack.compile_utils import make_fake_tensor
 
-from sglang.srt.layers.attention.cute_utils import (
+from sglang.kernels.ops.attention.cute_utils import (
     EVICT_FIRST,
     _tcgen05,
     cvt,

python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_kkt_inv_uw.py

@@ -10,7 +10,7 @@ from cutlass import BFloat16, Float32, Int32, Int64, Uint32, cute
 from cutlass.cute.nvgpu import cpasync, warp
 from quack.compile_utils import make_fake_tensor
 
-from sglang.srt.layers.attention.cute_utils import (
+from sglang.kernels.ops.attention.cute_utils import (
     EVICT_FIRST,
     _tcgen05,
     cvt,

python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_o.py

@@ -10,7 +10,7 @@ from cutlass import BFloat16, Float32, Int32, Int64, Uint32, cute
 from cutlass.cute.nvgpu import cpasync, warp
 from quack.compile_utils import make_fake_tensor
 
-from sglang.srt.layers.attention.cute_utils import (
+from sglang.kernels.ops.attention.cute_utils import (
     EVICT_FIRST,
     _tcgen05,
     cvt,

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

@@ -26,7 +26,7 @@ from cutlass import BFloat16, Float32, Int32, Int64, Uint32, cute
 from cutlass.cute.nvgpu import cpasync, warp
 from quack.compile_utils import make_fake_tensor
 
-from sglang.srt.layers.attention.cute_utils import (
+from sglang.kernels.ops.attention.cute_utils import (
     EVICT_FIRST,
     _tcgen05,
     cvt,

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

@@ -25,7 +25,7 @@ from cutlass import BFloat16, Float32, Int32, Int64, Uint32, cute
 from cutlass.cute.nvgpu import cpasync, warp
 from quack.compile_utils import make_fake_tensor
 
-from sglang.srt.layers.attention.cute_utils import (
+from sglang.kernels.ops.attention.cute_utils import (
     EVICT_FIRST,
     _tcgen05,
     cvt,

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

@@ -25,7 +25,7 @@ from cutlass import BFloat16, Int32, Int64, Uint32, cute
 from cutlass.cute.nvgpu import cpasync, warp
 from quack.compile_utils import make_fake_tensor
 
-from sglang.srt.layers.attention.cute_utils import (
+from sglang.kernels.ops.attention.cute_utils import (
     EVICT_FIRST,
     _tcgen05,
     cvt,

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

@@ -148,6 +148,7 @@ _PHASE25_TRITON_KERNELS = [
     ("router", "fused_moe_router_shim"),
     ("deepep_waterfill_kernels", "materialize_waterfill_dispatch_fused"),
     ("fill_padded_rows", "_fill_padded_rows"),
+    ("moe_fused_mul_sum", "moe_fused_mul_sum"),
 ]
 for _mod, _fn in _PHASE25_TRITON_KERNELS:
     register_kernel(

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

@@ -19,6 +19,9 @@ _TRITON_KERNELS = [
     ("spec_tree", "sgl_build_tree_kernel_efficient_triton"),
     ("spec_tree", "verify_tree_greedy_kernel_triton"),
     ("topk1", "draft_topk1_postprocess"),
+    ("ragged_verify_kernels", "pad_verify_lens_to_bucket"),
+    ("ragged_verify_kernels", "build_qo_indptr"),
+    ("reject_sampling", "chain_speculative_sampling_triton"),
 ]
 for _mod, _fn in _TRITON_KERNELS:
     register_kernel(

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

@@ -28,13 +28,13 @@ from sglang.kernels.ops.attention.dsv4.metadata_kernel import (
 from sglang.kernels.ops.attention.dsv4.quant_k_cache import (
     quant_to_nope_fp8_rope_bf16_pack_triton,
 )
-from sglang.srt.environ import envs
+from sglang.kernels.ops.attention.dsv4_attn_metadata_kernels import (
-from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
-from sglang.srt.layers.attention.dsv4.attn_metadata_kernels import (
     BuildCausalSwaPageIndices,
     BuildPageTablePositions,
     ExpandPrefillCausally,
 )
+from sglang.srt.environ import envs
+from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
 from sglang.srt.layers.attention.dsv4.compressor_v2 import (
     CompressorBackendMixin,
     FusedCompressMetadata,

python/sglang/srt/layers/moe/moe_runner/humming.py

@@ -10,8 +10,8 @@ from weakref import WeakValueDictionary
 import torch
 
 from sglang.kernels.ops.moe.ep_moe_kernels import moe_permute, moe_unpermute
+from sglang.kernels.ops.moe.moe_fused_mul_sum import moe_fused_mul_sum
 from sglang.srt.environ import envs
-from sglang.srt.layers.moe.fused_moe_triton.moe_fused_mul_sum import moe_fused_mul_sum
 from sglang.srt.layers.moe.moe_runner.base import (
     MoeQuantInfo,
     MoeRunnerConfig,

python/sglang/srt/speculative/dspark_components/kernels/dspark_accept.py

@@ -7,6 +7,9 @@ import torch
 import triton
 import triton.language as tl
 
+from sglang.kernels.ops.speculative.reject_sampling import (
+    chain_speculative_sampling_triton,
+)
 from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2
 from sglang.srt.speculative.dflash_utils import (
     _get_or_create_chain_verify_buffers,
@@ -14,7 +17,6 @@ from sglang.srt.speculative.dflash_utils import (
     compute_dflash_correct_drafts_and_bonus,
 )
 from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda
-from sglang.srt.speculative.reject_sampling import chain_speculative_sampling_triton
 
 
 class AcceptSampling:

python/sglang/srt/speculative/eagle_utils.py

@@ -668,7 +668,7 @@ def eagle_sample(
             tree_speculative_sampling_target_only,
         )
 
-        from sglang.srt.speculative.reject_sampling import (
+        from sglang.kernels.ops.speculative.reject_sampling import (
             chain_speculative_sampling_triton,
         )
 

python/sglang/srt/speculative/ragged_verify.py

@@ -90,7 +90,7 @@ class RaggedVerifyLayout(msgspec.Struct, frozen=True):
         verify_lens_cpu: Optional[list[int]] = None,
         total_verify_tokens: Optional[int] = None,
     ) -> RaggedVerifyLayout:
-        from sglang.srt.speculative.ragged_verify_kernels import (
+        from sglang.kernels.ops.speculative.ragged_verify_kernels import (
             BuildQoIndptr,
         )
 
@@ -155,7 +155,7 @@ class RaggedVerifyLayout(msgspec.Struct, frozen=True):
         )
 
     def padded_to_bucket(self, *, padded_bs: int) -> RaggedVerifyLayout:
-        from sglang.srt.speculative.ragged_verify_kernels import (
+        from sglang.kernels.ops.speculative.ragged_verify_kernels import (
             PaddedToBucket,
         )
 

test/registered/spec/dspark/test_dspark_kernel_parity.py

@@ -12,8 +12,10 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.attention.dsv4 import attn_metadata_kernels
+from sglang.kernels.ops.attention import (
-from sglang.srt.speculative import ragged_verify_kernels
+    dsv4_attn_metadata_kernels as attn_metadata_kernels,
+)
+from sglang.kernels.ops.speculative import ragged_verify_kernels
 from sglang.srt.speculative.dspark_components.dspark_planner import (
     DSparkScheduleConfig,
 )