[Kernel] Migrate scattered quantization kernels to sglang.kernels (RFC #29630, Phase 2.5, 1/7) (#30784)

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>

test/manual/quant/test_block_fp8.py

@@ -4,16 +4,16 @@ from functools import lru_cache
 
 import torch
 
-from sglang.srt.layers.activation import SiluAndMul
-from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_moe
-from sglang.srt.layers.moe.topk import TopKConfig, select_experts
-from sglang.srt.layers.quantization.fp8_kernel import (
+from sglang.kernels.ops.quantization.fp8_kernel import (
     per_tensor_quant_mla_fp8,
     per_token_group_quant_fp8,
     per_token_group_quant_mla_deep_gemm_masked_fp8,
     static_quant_fp8,
     w8a8_block_fp8_matmul,
 )
+from sglang.srt.layers.activation import SiluAndMul
+from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_moe
+from sglang.srt.layers.moe.topk import TopKConfig, select_experts
 from sglang.srt.layers.quantization.fp8_utils import (
     input_to_float8,
     mxfp8_group_quantize,

test/manual/quant/test_custom_ops.py

@@ -7,7 +7,7 @@ import sys
 import pytest
 import torch
 
-from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz, scaled_fp8_quant
+from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz, scaled_fp8_quant
 from sglang.srt.utils import is_cuda, is_hip
 
 _is_cuda = is_cuda()

test/registered/jit/benchmark/bench_per_token_group_quant_8bit.py

@@ -9,10 +9,10 @@ from sglang.jit_kernel.benchmark.utils import get_benchmark_range
 from sglang.jit_kernel.per_token_group_quant_8bit import (
     per_token_group_quant_8bit as sglang_per_token_group_quant_8bit,
 )
-from sglang.srt.layers.quantization.fp8_kernel import (
+from sglang.kernels.ops.quantization.fp8_kernel import (
     create_per_token_group_quant_fp8_output_scale,
 )
-from sglang.srt.layers.quantization.fp8_kernel import (
+from sglang.kernels.ops.quantization.fp8_kernel import (
     per_token_group_quant_8bit as triton_per_token_group_quant_8bit,
 )
 from sglang.srt.utils import is_hip

test/registered/jit/benchmark/bench_per_token_group_quant_8bit_v2.py

@@ -6,7 +6,7 @@ from sglang.jit_kernel.benchmark.utils import create_random
 from sglang.jit_kernel.per_token_group_quant_8bit_v2 import (
     per_token_group_quant_8bit_v2,
 )
-from sglang.srt.layers.quantization.fp8_kernel import (
+from sglang.kernels.ops.quantization.fp8_kernel import (
     create_per_token_group_quant_fp8_output_scale,
     fp8_dtype,
     fp8_max,

test/registered/jit/deepseek_v4/test_fp8_wo_a.py

@@ -11,7 +11,7 @@ import torch
 
 import sglang.jit_kernel.dsv4.fp8_wo_a as fp8_wo_a_module
 from sglang.jit_kernel.dsv4 import sglang_per_token_group_quant_fp8_dsv4_wo_a
-from sglang.srt.layers.quantization.fp8_kernel import (
+from sglang.kernels.ops.quantization.fp8_kernel import (
     fp8_dtype,
     sglang_per_token_group_quant_fp8,
 )

test/registered/jit/diffusion/test_diffusion_modelopt_fp8_scaled_mm.py

@@ -3,11 +3,11 @@ import sys
 import pytest
 import torch
 
+from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
 from sglang.multimodal_gen.runtime.layers.quantization.modelopt_quant import (
     ModelOptFp8Config,
     ModelOptFp8LinearMethod,
 )
-from sglang.srt.layers.quantization.fp8_kernel import static_quant_fp8
 from sglang.srt.layers.quantization.fp8_utils import (
     cutlass_fp8_supported,
     input_to_float8,

test/registered/jit/test_fused_store_index_cache.py

@@ -42,7 +42,7 @@ except ImportError:
     _is_hip = False
 
 try:
-    from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
+    from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
 
     _is_fp8_fnuz = is_fp8_fnuz()
 except ImportError:

test/registered/jit/test_per_token_group_quant_8bit.py

@@ -28,10 +28,10 @@ from sgl_kernel.test_utils import (  # noqa: E402
 from sglang.jit_kernel.per_token_group_quant_8bit import (  # noqa: E402
     per_token_group_quant_8bit as jit_per_token_group_quant_8bit,
 )
-from sglang.srt.layers.quantization.fp8_kernel import (  # noqa: E402
+from sglang.kernels.ops.quantization.fp8_kernel import (  # noqa: E402
     create_per_token_group_quant_fp8_output_scale,
 )
-from sglang.srt.layers.quantization.fp8_kernel import (  # noqa: E402
+from sglang.kernels.ops.quantization.fp8_kernel import (  # noqa: E402
     per_token_group_quant_8bit as triton_per_token_group_quant_8bit,
 )
 

test/registered/jit/test_per_token_group_quant_8bit_v2.py

@@ -21,7 +21,7 @@ if sgl_per_token_group_quant_8bit is None and not torch.cuda.is_available():
 if sgl_per_token_group_quant_8bit is None:
     raise ImportError("sgl_kernel AOT reference op is unavailable")
 
-from sglang.srt.layers.quantization.fp8_kernel import (  # noqa: E402
+from sglang.kernels.ops.quantization.fp8_kernel import (  # noqa: E402
     create_per_token_group_quant_fp8_output_scale,
     fp8_dtype,
     fp8_max,

test/registered/kernels/test_kernels_namespace.py

@@ -45,6 +45,12 @@ EXPECTED_OPS = {
     "moe.moe_align_block_size": {"cuda_aot", "cuda_jit"},
     "moe.topk_softmax": {"cuda_aot"},
     "quantization.sgl_per_token_quant_fp8": {"cuda_aot"},
+    # migrated from srt/layers/quantization (Phase 2.5)
+    "quantization.w8a8_block_fp8_matmul": {"triton"},
+    "quantization.per_token_quant_int8": {"triton"},
+    "quantization.awq_dequantize_triton": {"triton"},
+    "quantization.nvfp4_gemm_swiglu_nvfp4_quant": {"cute_dsl"},
+    "moe.pack_topk_ids": {"triton"},
     "quantization.sgl_per_token_group_quant_8bit": {"cuda_aot", "cuda_jit"},
     "quantization.sgl_per_token_group_quant_fp8": {"cuda_aot"},
     "quantization.sgl_per_token_group_quant_int8": {"cuda_aot"},

test/registered/moe/test_fused_moe.py

@@ -3,11 +3,11 @@ import unittest
 import torch
 from tqdm import tqdm
 
+from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.moe.moe_runner import MoeRunnerConfig
 from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_moe
 from sglang.srt.layers.moe.topk import TopKConfig, select_experts
-from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
 from sglang.srt.layers.quantization.fp8_utils import normalize_e4m3fn_to_e4m3fnuz
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.srt.utils import get_device, get_device_capability, is_hip

test/registered/moe/test_triton_moe_channel_fp8_kernel.py

@@ -3,10 +3,10 @@ import unittest
 
 import torch
 
+from sglang.kernels.ops.quantization.fp8_kernel import scaled_fp8_quant
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_moe
 from sglang.srt.layers.moe.topk import TopKConfig, select_experts
-from sglang.srt.layers.quantization.fp8_kernel import scaled_fp8_quant
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase

test/registered/quant/test_awq_dequant.py

@@ -12,7 +12,7 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.quantization.awq.awq_triton import (
+from sglang.kernels.ops.quantization.awq_triton import (
     AWQ_TRITON_SUPPORTED_GROUP_SIZES,
     awq_dequantize_triton,
     awq_gemm_triton,

test/registered/quant/test_fp8_blockwise_row_padding.py

@@ -11,7 +11,7 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.quantization.fp8_kernel import (
+from sglang.kernels.ops.quantization.fp8_kernel import (
     fp8_dtype,
     per_token_group_quant_fp8,
     sglang_per_token_group_quant_fp8_row_padded,

test/registered/quant/test_fp8_kernel.py

@@ -2,7 +2,7 @@ import unittest
 
 import torch
 
-from sglang.srt.layers.quantization.fp8_kernel import (
+from sglang.kernels.ops.quantization.fp8_kernel import (
     per_token_group_quant_fp8,
     w8a8_block_fp8_matmul,
 )

test/registered/quant/test_int8_kernel.py

@@ -3,10 +3,10 @@ import unittest
 
 import torch
 
+from sglang.kernels.ops.quantization.int8_kernel import per_token_quant_int8
 from sglang.srt.layers.activation import SiluAndMul
 from sglang.srt.layers.moe.moe_runner.triton_utils.fused_moe import fused_moe
 from sglang.srt.layers.moe.topk import TopKConfig, select_experts
-from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8
 from sglang.srt.server_args import ServerArgs, set_global_server_args_for_scheduler
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase

test/registered/quant/test_triton_scaled_mm.py

@@ -4,7 +4,7 @@ from typing import Optional
 import torch
 import torch.testing
 
-from sglang.srt.layers.quantization.fp8_kernel import triton_scaled_mm
+from sglang.kernels.ops.quantization.fp8_kernel import triton_scaled_mm
 from sglang.srt.utils.common import get_device
 from sglang.test.ci.ci_register import register_amd_ci, register_cuda_ci
 from sglang.test.test_utils import CustomTestCase