From 216b750c8f5b80c3875b27b198b39a09bae871a0 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Sat, 18 Jul 2026 19:07:07 +0800 Subject: [PATCH] [Kernel] Sweep decoupled scattered kernels into sglang.kernels.ops (RFC #29630) (#31582) Co-authored-by: Claude Opus 4.8 (1M context) --- python/sglang/kernels/ops/attention/__init__.py | 3 +++ .../layers => kernels/ops}/attention/cute_utils/__init__.py | 0 .../layers => kernels/ops}/attention/cute_utils/_tcgen05.py | 0 .../{srt/layers => kernels/ops}/attention/cute_utils/cvt.py | 0 .../ops/attention/dsv4_attn_metadata_kernels.py} | 0 .../kernels/ops/attention/linear/gdn_blackwell/kernel_h.py | 2 +- .../ops/attention/linear/gdn_blackwell/kernel_kkt_inv_uw.py | 2 +- .../kernels/ops/attention/linear/gdn_blackwell/kernel_o.py | 2 +- .../kernels/ops/attention/linear/kda_blackwell/kernel_h.py | 2 +- .../ops/attention/linear/kda_blackwell/kernel_kkt_inv_uw.py | 2 +- .../kernels/ops/attention/linear/kda_blackwell/kernel_o.py | 2 +- python/sglang/kernels/ops/moe/__init__.py | 1 + .../ops/moe}/moe_fused_mul_sum.py | 0 python/sglang/kernels/ops/speculative/__init__.py | 3 +++ .../ops}/speculative/ragged_verify_kernels.py | 0 .../{srt => kernels/ops}/speculative/reject_sampling.py | 0 python/sglang/srt/layers/attention/deepseek_v4_backend.py | 6 +++--- python/sglang/srt/layers/moe/moe_runner/humming.py | 2 +- .../speculative/dspark_components/kernels/dspark_accept.py | 4 +++- python/sglang/srt/speculative/eagle_utils.py | 2 +- python/sglang/srt/speculative/ragged_verify.py | 4 ++-- test/registered/spec/dspark/test_dspark_kernel_parity.py | 6 ++++-- 22 files changed, 27 insertions(+), 16 deletions(-) rename python/sglang/{srt/layers => kernels/ops}/attention/cute_utils/__init__.py (100%) rename python/sglang/{srt/layers => kernels/ops}/attention/cute_utils/_tcgen05.py (100%) rename python/sglang/{srt/layers => kernels/ops}/attention/cute_utils/cvt.py (100%) rename python/sglang/{srt/layers/attention/dsv4/attn_metadata_kernels.py => kernels/ops/attention/dsv4_attn_metadata_kernels.py} (100%) rename python/sglang/{srt/layers/moe/fused_moe_triton => kernels/ops/moe}/moe_fused_mul_sum.py (100%) rename python/sglang/{srt => kernels/ops}/speculative/ragged_verify_kernels.py (100%) rename python/sglang/{srt => kernels/ops}/speculative/reject_sampling.py (100%) diff --git a/python/sglang/kernels/ops/attention/__init__.py b/python/sglang/kernels/ops/attention/__init__.py index f72763785..781522e36 100644 --- a/python/sglang/kernels/ops/attention/__init__.py +++ b/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( diff --git a/python/sglang/srt/layers/attention/cute_utils/__init__.py b/python/sglang/kernels/ops/attention/cute_utils/__init__.py similarity index 100% rename from python/sglang/srt/layers/attention/cute_utils/__init__.py rename to python/sglang/kernels/ops/attention/cute_utils/__init__.py diff --git a/python/sglang/srt/layers/attention/cute_utils/_tcgen05.py b/python/sglang/kernels/ops/attention/cute_utils/_tcgen05.py similarity index 100% rename from python/sglang/srt/layers/attention/cute_utils/_tcgen05.py rename to python/sglang/kernels/ops/attention/cute_utils/_tcgen05.py diff --git a/python/sglang/srt/layers/attention/cute_utils/cvt.py b/python/sglang/kernels/ops/attention/cute_utils/cvt.py similarity index 100% rename from python/sglang/srt/layers/attention/cute_utils/cvt.py rename to python/sglang/kernels/ops/attention/cute_utils/cvt.py diff --git a/python/sglang/srt/layers/attention/dsv4/attn_metadata_kernels.py b/python/sglang/kernels/ops/attention/dsv4_attn_metadata_kernels.py similarity index 100% rename from python/sglang/srt/layers/attention/dsv4/attn_metadata_kernels.py rename to python/sglang/kernels/ops/attention/dsv4_attn_metadata_kernels.py diff --git a/python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_h.py b/python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_h.py index 460381f45..9f1dc2c02 100644 --- a/python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_h.py +++ b/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, diff --git a/python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_kkt_inv_uw.py b/python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_kkt_inv_uw.py index 21fcd509a..c600018e4 100644 --- a/python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_kkt_inv_uw.py +++ b/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, diff --git a/python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_o.py b/python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_o.py index 3665cc19e..203e4e2fd 100644 --- a/python/sglang/kernels/ops/attention/linear/gdn_blackwell/kernel_o.py +++ b/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, diff --git a/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_h.py b/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_h.py index 035e851cd..078f48e7e 100644 --- a/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_h.py +++ b/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, diff --git a/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_kkt_inv_uw.py b/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_kkt_inv_uw.py index 7900edefb..68bb82062 100644 --- a/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_kkt_inv_uw.py +++ b/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, diff --git a/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_o.py b/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_o.py index 70c4634c3..f78c95690 100644 --- a/python/sglang/kernels/ops/attention/linear/kda_blackwell/kernel_o.py +++ b/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, diff --git a/python/sglang/kernels/ops/moe/__init__.py b/python/sglang/kernels/ops/moe/__init__.py index b4efb1219..447cb9615 100644 --- a/python/sglang/kernels/ops/moe/__init__.py +++ b/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( diff --git a/python/sglang/srt/layers/moe/fused_moe_triton/moe_fused_mul_sum.py b/python/sglang/kernels/ops/moe/moe_fused_mul_sum.py similarity index 100% rename from python/sglang/srt/layers/moe/fused_moe_triton/moe_fused_mul_sum.py rename to python/sglang/kernels/ops/moe/moe_fused_mul_sum.py diff --git a/python/sglang/kernels/ops/speculative/__init__.py b/python/sglang/kernels/ops/speculative/__init__.py index 827fac9e8..a67dab0cc 100644 --- a/python/sglang/kernels/ops/speculative/__init__.py +++ b/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( diff --git a/python/sglang/srt/speculative/ragged_verify_kernels.py b/python/sglang/kernels/ops/speculative/ragged_verify_kernels.py similarity index 100% rename from python/sglang/srt/speculative/ragged_verify_kernels.py rename to python/sglang/kernels/ops/speculative/ragged_verify_kernels.py diff --git a/python/sglang/srt/speculative/reject_sampling.py b/python/sglang/kernels/ops/speculative/reject_sampling.py similarity index 100% rename from python/sglang/srt/speculative/reject_sampling.py rename to python/sglang/kernels/ops/speculative/reject_sampling.py diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 306e60acc..cf5fa0183 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/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.srt.layers.attention.base_attn_backend import AttentionBackend -from sglang.srt.layers.attention.dsv4.attn_metadata_kernels import ( +from sglang.kernels.ops.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, diff --git a/python/sglang/srt/layers/moe/moe_runner/humming.py b/python/sglang/srt/layers/moe/moe_runner/humming.py index b72debc1b..845bf4e40 100644 --- a/python/sglang/srt/layers/moe/moe_runner/humming.py +++ b/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, diff --git a/python/sglang/srt/speculative/dspark_components/kernels/dspark_accept.py b/python/sglang/srt/speculative/dspark_components/kernels/dspark_accept.py index 560c17077..cbc5734f0 100644 --- a/python/sglang/srt/speculative/dspark_components/kernels/dspark_accept.py +++ b/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: diff --git a/python/sglang/srt/speculative/eagle_utils.py b/python/sglang/srt/speculative/eagle_utils.py index 6cd08b92b..56dedf7e8 100644 --- a/python/sglang/srt/speculative/eagle_utils.py +++ b/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, ) diff --git a/python/sglang/srt/speculative/ragged_verify.py b/python/sglang/srt/speculative/ragged_verify.py index 1dcb618ec..a2906681c 100644 --- a/python/sglang/srt/speculative/ragged_verify.py +++ b/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, ) diff --git a/test/registered/spec/dspark/test_dspark_kernel_parity.py b/test/registered/spec/dspark/test_dspark_kernel_parity.py index 9faa2f7a6..3554ab982 100644 --- a/test/registered/spec/dspark/test_dspark_kernel_parity.py +++ b/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.srt.speculative import ragged_verify_kernels +from sglang.kernels.ops.attention import ( + 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, )