From 62aa85d9aa2fabd339ee06077e09d3ce1ac4e608 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Thu, 23 Jul 2026 17:07:16 +0800 Subject: [PATCH] [Kernel] Sweep missed dedicated kernels into kernels.ops (moe/quant siblings + dspark) (RFC #29630) (#32160) Co-authored-by: Claude Opus 4.8 --- .../ops/moe}/gate_topk.py | 0 .../ops/moe}/inkling_moe.py | 0 .../ops/moe}/sigmoid_gate_topk_renorm.py | 10 +++--- .../ops}/quantization/mxfp8_interleave_sf.py | 0 .../ops}/quantization/mxfp8_quant.py | 0 .../ops/speculative/dspark}/__init__.py | 0 .../ops/speculative/dspark}/dispatch.py | 0 .../ops/speculative/dspark}/dspark_accept.py | 2 +- .../dspark}/dspark_attn_metadata.py | 2 +- .../speculative/dspark}/dspark_draft_model.py | 2 +- .../speculative/dspark}/dspark_schedule.py | 2 +- .../dspark}/dspark_verify_window.py | 2 +- .../layers/attention/deepseek_v4_backend.py | 10 +++--- python/sglang/srt/mem_cache/memory_pool.py | 6 ++-- .../sglang/srt/models/deepseek_v4_dspark.py | 8 ++--- .../sglang/srt/models/inkling_common/attn.py | 2 +- .../srt/models/inkling_common/dense_mlp.py | 2 +- .../sglang/srt/models/inkling_common/moe.py | 20 +++++------ .../dspark_components/dspark_draft.py | 6 ++-- .../dspark_components/dspark_kv_inject.py | 4 +-- .../dspark_components/dspark_planner.py | 8 ++--- .../dspark_components/dspark_verify.py | 34 +++++++++---------- .../kernels/ops/moe/test_moe_preprocess.py | 2 +- .../spec/dspark/test_dspark_kernel_parity.py | 8 ++--- .../spec/dspark/test_dspark_scheduler.py | 6 ++-- 25 files changed, 68 insertions(+), 68 deletions(-) rename python/sglang/{srt/layers/moe/moe_runner/triton_utils => kernels/ops/moe}/gate_topk.py (100%) rename python/sglang/{srt/layers/moe/moe_runner/triton_utils => kernels/ops/moe}/inkling_moe.py (100%) rename python/sglang/{srt/layers/moe/moe_runner/triton_utils => kernels/ops/moe}/sigmoid_gate_topk_renorm.py (99%) rename python/sglang/{srt/layers => kernels/ops}/quantization/mxfp8_interleave_sf.py (100%) rename python/sglang/{srt/layers => kernels/ops}/quantization/mxfp8_quant.py (100%) rename python/sglang/{srt/speculative/dspark_components/kernels => kernels/ops/speculative/dspark}/__init__.py (100%) rename python/sglang/{srt/speculative/dspark_components/kernels => kernels/ops/speculative/dspark}/dispatch.py (100%) rename python/sglang/{srt/speculative/dspark_components/kernels => kernels/ops/speculative/dspark}/dspark_accept.py (99%) rename python/sglang/{srt/speculative/dspark_components/kernels => kernels/ops/speculative/dspark}/dspark_attn_metadata.py (99%) rename python/sglang/{srt/speculative/dspark_components/kernels => kernels/ops/speculative/dspark}/dspark_draft_model.py (99%) rename python/sglang/{srt/speculative/dspark_components/kernels => kernels/ops/speculative/dspark}/dspark_schedule.py (99%) rename python/sglang/{srt/speculative/dspark_components/kernels => kernels/ops/speculative/dspark}/dspark_verify_window.py (99%) diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/gate_topk.py b/python/sglang/kernels/ops/moe/gate_topk.py similarity index 100% rename from python/sglang/srt/layers/moe/moe_runner/triton_utils/gate_topk.py rename to python/sglang/kernels/ops/moe/gate_topk.py diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/inkling_moe.py b/python/sglang/kernels/ops/moe/inkling_moe.py similarity index 100% rename from python/sglang/srt/layers/moe/moe_runner/triton_utils/inkling_moe.py rename to python/sglang/kernels/ops/moe/inkling_moe.py diff --git a/python/sglang/srt/layers/moe/moe_runner/triton_utils/sigmoid_gate_topk_renorm.py b/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py similarity index 99% rename from python/sglang/srt/layers/moe/moe_runner/triton_utils/sigmoid_gate_topk_renorm.py rename to python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py index e6a0a7cdf..9221086e9 100644 --- a/python/sglang/srt/layers/moe/moe_runner/triton_utils/sigmoid_gate_topk_renorm.py +++ b/python/sglang/kernels/ops/moe/sigmoid_gate_topk_renorm.py @@ -13,15 +13,15 @@ import triton import triton.language as tl from sglang.kernels.jit.utils import is_arch_support_pdl -from sglang.kernels.ops.moe.inkling_gate_topk_renorm import ( - inkling_gate_topk_renorm_v2, -) -from sglang.srt.environ import envs -from sglang.srt.layers.moe.moe_runner.triton_utils.gate_topk import ( +from sglang.kernels.ops.moe.gate_topk import ( fpval_to_key, indx_to_key, key_to_indx, ) +from sglang.kernels.ops.moe.inkling_gate_topk_renorm import ( + inkling_gate_topk_renorm_v2, +) +from sglang.srt.environ import envs @triton.jit diff --git a/python/sglang/srt/layers/quantization/mxfp8_interleave_sf.py b/python/sglang/kernels/ops/quantization/mxfp8_interleave_sf.py similarity index 100% rename from python/sglang/srt/layers/quantization/mxfp8_interleave_sf.py rename to python/sglang/kernels/ops/quantization/mxfp8_interleave_sf.py diff --git a/python/sglang/srt/layers/quantization/mxfp8_quant.py b/python/sglang/kernels/ops/quantization/mxfp8_quant.py similarity index 100% rename from python/sglang/srt/layers/quantization/mxfp8_quant.py rename to python/sglang/kernels/ops/quantization/mxfp8_quant.py diff --git a/python/sglang/srt/speculative/dspark_components/kernels/__init__.py b/python/sglang/kernels/ops/speculative/dspark/__init__.py similarity index 100% rename from python/sglang/srt/speculative/dspark_components/kernels/__init__.py rename to python/sglang/kernels/ops/speculative/dspark/__init__.py diff --git a/python/sglang/srt/speculative/dspark_components/kernels/dispatch.py b/python/sglang/kernels/ops/speculative/dspark/dispatch.py similarity index 100% rename from python/sglang/srt/speculative/dspark_components/kernels/dispatch.py rename to python/sglang/kernels/ops/speculative/dspark/dispatch.py diff --git a/python/sglang/srt/speculative/dspark_components/kernels/dspark_accept.py b/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py similarity index 99% rename from python/sglang/srt/speculative/dspark_components/kernels/dspark_accept.py rename to python/sglang/kernels/ops/speculative/dspark/dspark_accept.py index cbc5734f0..948a9d60a 100644 --- a/python/sglang/srt/speculative/dspark_components/kernels/dspark_accept.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_accept.py @@ -7,6 +7,7 @@ import torch import triton import triton.language as tl +from sglang.kernels.ops.speculative.dspark.dispatch import inputs_on_cuda from sglang.kernels.ops.speculative.reject_sampling import ( chain_speculative_sampling_triton, ) @@ -16,7 +17,6 @@ from sglang.srt.speculative.dflash_utils import ( build_dflash_verify_target_probs, compute_dflash_correct_drafts_and_bonus, ) -from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda class AcceptSampling: diff --git a/python/sglang/srt/speculative/dspark_components/kernels/dspark_attn_metadata.py b/python/sglang/kernels/ops/speculative/dspark/dspark_attn_metadata.py similarity index 99% rename from python/sglang/srt/speculative/dspark_components/kernels/dspark_attn_metadata.py rename to python/sglang/kernels/ops/speculative/dspark/dspark_attn_metadata.py index a18076a92..52177dce2 100644 --- a/python/sglang/srt/speculative/dspark_components/kernels/dspark_attn_metadata.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_attn_metadata.py @@ -7,7 +7,7 @@ import torch import triton import triton.language as tl -from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda +from sglang.kernels.ops.speculative.dspark.dispatch import inputs_on_cuda from sglang.srt.utils import ceil_align diff --git a/python/sglang/srt/speculative/dspark_components/kernels/dspark_draft_model.py b/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py similarity index 99% rename from python/sglang/srt/speculative/dspark_components/kernels/dspark_draft_model.py rename to python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py index d09c717e1..2ea446ce7 100644 --- a/python/sglang/srt/speculative/dspark_components/kernels/dspark_draft_model.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_draft_model.py @@ -8,7 +8,7 @@ import torch.nn.functional as F import triton import triton.language as tl -from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda +from sglang.kernels.ops.speculative.dspark.dispatch import inputs_on_cuda _BLOCK_V = 1024 _IDX_SENTINEL = tl.constexpr(2147483647) diff --git a/python/sglang/srt/speculative/dspark_components/kernels/dspark_schedule.py b/python/sglang/kernels/ops/speculative/dspark/dspark_schedule.py similarity index 99% rename from python/sglang/srt/speculative/dspark_components/kernels/dspark_schedule.py rename to python/sglang/kernels/ops/speculative/dspark/dspark_schedule.py index 52ffc74bb..52093ff79 100644 --- a/python/sglang/srt/speculative/dspark_components/kernels/dspark_schedule.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_schedule.py @@ -6,7 +6,7 @@ import torch import triton import triton.language as tl -from sglang.srt.speculative.dspark_components.kernels.dispatch import ( +from sglang.kernels.ops.speculative.dspark.dispatch import ( inputs_on_cuda, ) diff --git a/python/sglang/srt/speculative/dspark_components/kernels/dspark_verify_window.py b/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py similarity index 99% rename from python/sglang/srt/speculative/dspark_components/kernels/dspark_verify_window.py rename to python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py index 62f396d1a..c8070210e 100644 --- a/python/sglang/srt/speculative/dspark_components/kernels/dspark_verify_window.py +++ b/python/sglang/kernels/ops/speculative/dspark/dspark_verify_window.py @@ -6,8 +6,8 @@ import triton import triton.language as tl from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func +from sglang.kernels.ops.speculative.dspark.dispatch import inputs_on_cuda from sglang.srt.managers.schedule_batch import ScheduleBatch -from sglang.srt.speculative.dspark_components.kernels.dispatch import inputs_on_cuda from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 8820e95bf..659ae1438 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -33,6 +33,11 @@ from sglang.kernels.ops.attention.dsv4_attn_metadata_kernels import ( BuildPageTablePositions, ExpandPrefillCausally, ) +from sglang.kernels.ops.speculative.dspark.dspark_attn_metadata import ( + BuildBlockSeqLensCausal, + BuildDsparkSwaPageIndices, + ComputeDsparkWindowGather, +) 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 ( @@ -54,11 +59,6 @@ from sglang.srt.layers.attention.dsv4.sparse_prefill_utils import ( from sglang.srt.mem_cache.deepseek_v4_memory_pool import DeepSeekV4TokenToKVPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.runtime_context import get_parallel -from sglang.srt.speculative.dspark_components.kernels.dspark_attn_metadata import ( - BuildBlockSeqLensCausal, - BuildDsparkSwaPageIndices, - ComputeDsparkWindowGather, -) from sglang.srt.speculative.eagle_utils import per_step_draft_out_cache_loc from sglang.srt.speculative.ragged_verify import ( RaggedVerifyMode, diff --git a/python/sglang/srt/mem_cache/memory_pool.py b/python/sglang/srt/mem_cache/memory_pool.py index f0879929e..680d9e9d2 100644 --- a/python/sglang/srt/mem_cache/memory_pool.py +++ b/python/sglang/srt/mem_cache/memory_pool.py @@ -3361,7 +3361,7 @@ class MHATokenToKVPoolMXFP8(MHATokenToKVPool): # payload and the interleaved UE8M0 scales. if not self.mxfp8_sf_interleaved or cache_k.dtype == self.store_dtype: raise ValueError("MXFP8 KV cache requires K and V scale tensors.") - from sglang.srt.layers.quantization.mxfp8_quant import quant_store_kv_mxfp8 + from sglang.kernels.ops.quantization.mxfp8_quant import quant_store_kv_mxfp8 quant_store_kv_mxfp8( cache_k, @@ -3394,7 +3394,7 @@ class MHATokenToKVPoolMXFP8(MHATokenToKVPool): """Write per-token UE8M0 K/V scales — interleaved into the FA4 BlockScaledBasicChunk layout for page_size==128, flat otherwise.""" if self.mxfp8_sf_interleaved: - from sglang.srt.layers.quantization.mxfp8_interleave_sf import ( + from sglang.kernels.ops.quantization.mxfp8_interleave_sf import ( store_sf_interleaved, ) @@ -3436,7 +3436,7 @@ class MHATokenToKVPoolMXFP8(MHATokenToKVPool): # scale rows must travel with their fp8 payload or dequant reads # mismatched exponents. if self.mxfp8_sf_interleaved: - from sglang.srt.layers.quantization.mxfp8_interleave_sf import ( + from sglang.kernels.ops.quantization.mxfp8_interleave_sf import ( store_sf_interleaved, ) diff --git a/python/sglang/srt/models/deepseek_v4_dspark.py b/python/sglang/srt/models/deepseek_v4_dspark.py index ac7d60cf6..c82683bbe 100644 --- a/python/sglang/srt/models/deepseek_v4_dspark.py +++ b/python/sglang/srt/models/deepseek_v4_dspark.py @@ -9,6 +9,10 @@ import torch.nn.functional as F from torch import nn from sglang.kernels.ops.attention.dsv4 import fused_q_norm_rope, fused_rope_inplace +from sglang.kernels.ops.speculative.dspark.dspark_draft_model import ( + BuildStepLocal, + CommitKvProj, +) from sglang.srt.configs.deepseek_v4 import DeepSeekV4Config from sglang.srt.environ import envs from sglang.srt.layers.layernorm import RMSNorm @@ -40,10 +44,6 @@ from sglang.srt.runtime_context import get_parallel from sglang.srt.speculative.dspark_components.dspark_config import ( parse_dspark_draft_config, ) -from sglang.srt.speculative.dspark_components.kernels.dspark_draft_model import ( - BuildStepLocal, - CommitKvProj, -) from sglang.srt.speculative.ragged_verify import ( RaggedVerifyMode, read_ragged_verify_mode, diff --git a/python/sglang/srt/models/inkling_common/attn.py b/python/sglang/srt/models/inkling_common/attn.py index 6b6c7300b..8f4657e60 100644 --- a/python/sglang/srt/models/inkling_common/attn.py +++ b/python/sglang/srt/models/inkling_common/attn.py @@ -883,7 +883,7 @@ class InklingAttention(nn.Module): if prologue_q_descale is not None: extra_attn_kwargs["q_descale"] = prologue_q_descale else: - from sglang.srt.layers.quantization.mxfp8_quant import to_mxfp8 + from sglang.kernels.ops.quantization.mxfp8_quant import to_mxfp8 q_mxfp = to_mxfp8(q.view(num_tokens, self.num_tp_heads, self.head_dim)) q = q_mxfp.data.view(num_tokens, -1) diff --git a/python/sglang/srt/models/inkling_common/dense_mlp.py b/python/sglang/srt/models/inkling_common/dense_mlp.py index 0509cb29b..6d775a333 100644 --- a/python/sglang/srt/models/inkling_common/dense_mlp.py +++ b/python/sglang/srt/models/inkling_common/dense_mlp.py @@ -382,7 +382,7 @@ class InklingBatchDenseMLP(nn.Module, FusedMoELoadingMixin): def _swiglu(self, y_st2f: torch.Tensor, gammas_st: torch.Tensor) -> torch.Tensor: # Helion's kernel can produce NaNs for small shared-expert batches. - from sglang.srt.layers.moe.moe_runner.triton_utils.inkling_moe import ( + from sglang.kernels.ops.moe.inkling_moe import ( silu_and_mul_triton, ) diff --git a/python/sglang/srt/models/inkling_common/moe.py b/python/sglang/srt/models/inkling_common/moe.py index c4036069b..244476357 100644 --- a/python/sglang/srt/models/inkling_common/moe.py +++ b/python/sglang/srt/models/inkling_common/moe.py @@ -10,20 +10,13 @@ from torch import nn from triton.language.extra import libdevice from sglang.kernels.jit.utils import is_arch_support_pdl +from sglang.kernels.ops.moe.gate_topk import gate_topk from sglang.kernels.ops.moe.inkling_gate_topk_renorm import ( ensure_gate_gemv_fused_scratch, inkling_gate_gemv, inkling_gate_gemv_fused, ) -from sglang.srt.configs.inkling import InklingModelConfig -from sglang.srt.distributed import ( - get_tensor_model_parallel_group, -) -from sglang.srt.environ import GateGemvMode, envs -from sglang.srt.layers.moe import get_moe_runner_backend -from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE -from sglang.srt.layers.moe.moe_runner.triton_utils.gate_topk import gate_topk -from sglang.srt.layers.moe.moe_runner.triton_utils.inkling_moe import ( +from sglang.kernels.ops.moe.inkling_moe import ( FUSED_PREPROCESS_WIN_TOKENS, compute_grouped_gemm_metadata, fused_moe_preprocess, @@ -34,9 +27,16 @@ from sglang.srt.layers.moe.moe_runner.triton_utils.inkling_moe import ( select_grouped_gemm_block_m, silu_and_mul_helion, ) -from sglang.srt.layers.moe.moe_runner.triton_utils.sigmoid_gate_topk_renorm import ( +from sglang.kernels.ops.moe.sigmoid_gate_topk_renorm import ( sigmoid_gate_topk_renorm, ) +from sglang.srt.configs.inkling import InklingModelConfig +from sglang.srt.distributed import ( + get_tensor_model_parallel_group, +) +from sglang.srt.environ import GateGemvMode, envs +from sglang.srt.layers.moe import get_moe_runner_backend +from sglang.srt.layers.moe.fused_moe_triton.layer import FusedMoE from sglang.srt.layers.moe.topk import PackedTopKOutput, StandardTopKOutput from sglang.srt.layers.moe.utils import RoutingMethodType from sglang.srt.layers.quantization.base_config import QuantizationConfig diff --git a/python/sglang/srt/speculative/dspark_components/dspark_draft.py b/python/sglang/srt/speculative/dspark_components/dspark_draft.py index c4cb2e013..bc2e3cc56 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_draft.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_draft.py @@ -7,6 +7,9 @@ from typing import Optional import msgspec import torch +from sglang.kernels.ops.speculative.dspark.dspark_draft_model import ( + SampleStepTokens, +) from sglang.srt.environ import envs from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.model_executor.forward_batch_info import ( @@ -18,9 +21,6 @@ from sglang.srt.runtime_context import get_parallel from sglang.srt.speculative.dflash_info_v2 import DFlashDraftInputV2 from sglang.srt.speculative.draft_worker_common import make_draft_input_v2 from sglang.srt.speculative.dspark_components.dspark_planner import VerifyWindow -from sglang.srt.speculative.dspark_components.kernels.dspark_draft_model import ( - SampleStepTokens, -) from sglang.srt.speculative.spec_info import ( SpeculativeAlgorithm, spec_scale_global_num_tokens, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py index 329c3f455..d1e44145a 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_kv_inject.py @@ -3,10 +3,10 @@ from typing import Optional import torch from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func -from sglang.srt.managers.schedule_batch import ScheduleBatch -from sglang.srt.speculative.dspark_components.kernels.dspark_verify_window import ( +from sglang.kernels.ops.speculative.dspark.dspark_verify_window import ( BuildCommitInjectLayout, ) +from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout diff --git a/python/sglang/srt/speculative/dspark_components/dspark_planner.py b/python/sglang/srt/speculative/dspark_components/dspark_planner.py index e76206d5c..d2c257698 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_planner.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_planner.py @@ -7,6 +7,10 @@ import msgspec import torch from sglang.kernels.ops.speculative.cache_locs import assign_extend_cache_locs_func +from sglang.kernels.ops.speculative.dspark.dspark_schedule import ( + ScheduleVerifyLensTopk, + compute_sort_survival, +) from sglang.srt.distributed import get_tp_group from sglang.srt.environ import envs from sglang.srt.layers.dp_attention import is_dp_attention_enabled @@ -31,10 +35,6 @@ from sglang.srt.speculative.dspark_components.dspark_sps import ( from sglang.srt.speculative.dspark_components.dspark_sts import ( load_sts_calibration_from_path, ) -from sglang.srt.speculative.dspark_components.kernels.dspark_schedule import ( - ScheduleVerifyLensTopk, - compute_sort_survival, -) from sglang.srt.speculative.ragged_verify import ( RaggedVerifyLayout, RaggedVerifyMode, diff --git a/python/sglang/srt/speculative/dspark_components/dspark_verify.py b/python/sglang/srt/speculative/dspark_components/dspark_verify.py index 43099c9fe..0302d75b1 100644 --- a/python/sglang/srt/speculative/dspark_components/dspark_verify.py +++ b/python/sglang/srt/speculative/dspark_components/dspark_verify.py @@ -5,6 +5,23 @@ from typing import Optional import msgspec import torch +from sglang.kernels.ops.speculative.dspark.dspark_accept import ( + AcceptGreedy, + AcceptSampling, + FinalizeAcceptLens, + SelectMixedAccept, + SoftmaxTemp, + accept_greedy_triton, + finalize_accept_lens_triton, +) +from sglang.kernels.ops.speculative.dspark.dspark_verify_window import ( + BuildCommitInjectLayout, + BuildOutTokens, + BuildRaggedVerifyWindow, + RaggedVerifyWindow, + ScatterCompactToStrided, + scatter_compact_to_strided_into, +) from sglang.srt.layers.logits_processor import LogitsProcessorOutput from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.model_executor.forward_batch_info import CaptureHiddenMode, ForwardMode @@ -19,23 +36,6 @@ from sglang.srt.speculative.dspark_components.dspark_planner import ( VerifyWindow, apply_logits_adjustments_strided, ) -from sglang.srt.speculative.dspark_components.kernels.dspark_accept import ( - AcceptGreedy, - AcceptSampling, - FinalizeAcceptLens, - SelectMixedAccept, - SoftmaxTemp, - accept_greedy_triton, - finalize_accept_lens_triton, -) -from sglang.srt.speculative.dspark_components.kernels.dspark_verify_window import ( - BuildCommitInjectLayout, - BuildOutTokens, - BuildRaggedVerifyWindow, - RaggedVerifyWindow, - ScatterCompactToStrided, - scatter_compact_to_strided_into, -) from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout diff --git a/test/registered/kernels/ops/moe/test_moe_preprocess.py b/test/registered/kernels/ops/moe/test_moe_preprocess.py index 43e12c59e..6feb98441 100644 --- a/test/registered/kernels/ops/moe/test_moe_preprocess.py +++ b/test/registered/kernels/ops/moe/test_moe_preprocess.py @@ -6,7 +6,7 @@ configs (the block schedule and kernel config are chosen together). import pytest import torch -from sglang.srt.layers.moe.moe_runner.triton_utils.inkling_moe import ( +from sglang.kernels.ops.moe.inkling_moe import ( SMALL_M_BLOCK_SIZE_M, compute_grouped_gemm_metadata, fused_moe_preprocess, diff --git a/test/registered/spec/dspark/test_dspark_kernel_parity.py b/test/registered/spec/dspark/test_dspark_kernel_parity.py index 3554ab982..376ab2182 100644 --- a/test/registered/spec/dspark/test_dspark_kernel_parity.py +++ b/test/registered/spec/dspark/test_dspark_kernel_parity.py @@ -16,16 +16,16 @@ 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, -) -from sglang.srt.speculative.dspark_components.kernels import ( +from sglang.kernels.ops.speculative.dspark import ( dspark_accept, dspark_attn_metadata, dspark_draft_model, dspark_schedule, dspark_verify_window, ) +from sglang.srt.speculative.dspark_components.dspark_planner import ( + DSparkScheduleConfig, +) from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.test_utils import CustomTestCase diff --git a/test/registered/spec/dspark/test_dspark_scheduler.py b/test/registered/spec/dspark/test_dspark_scheduler.py index 1bb77b532..53715fa1d 100644 --- a/test/registered/spec/dspark/test_dspark_scheduler.py +++ b/test/registered/spec/dspark/test_dspark_scheduler.py @@ -4,6 +4,9 @@ import unittest import torch +from sglang.kernels.ops.speculative.dspark.dspark_schedule import ( + schedule_verify_lens_topk_from_survival, +) from sglang.srt.speculative.dspark_components.dspark_planner import ( DSparkScheduleConfig, HostConfidenceBudgetPlanner, @@ -15,9 +18,6 @@ from sglang.srt.speculative.dspark_components.dspark_sps import ( SpsAdditiveCostTable, SpsCostTable, ) -from sglang.srt.speculative.dspark_components.kernels.dspark_schedule import ( - schedule_verify_lens_topk_from_survival, -) from sglang.srt.speculative.ragged_verify import RaggedVerifyLayout from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase