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

benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm.py

@@ -12,7 +12,7 @@ from vllm.model_executor.layers.quantization.utils.fp8_utils import (
 )
 
 from sglang.benchmark.bench_utils import run_bench
-from sglang.srt.layers.quantization.fp8_kernel import (
+from sglang.kernels.ops.quantization.fp8_kernel import (
     w8a8_block_fp8_matmul_deepgemm as w8a8_block_fp8_matmul,
 )
 

benchmark/kernels/deepseek/benchmark_deepgemm_fp8_gemm_blackwell.py

@@ -7,7 +7,7 @@ from deep_gemm import ceil_div
 from flashinfer.gemm import gemm_fp8_nt_groupwise
 
 from sglang.benchmark.bench_utils import run_bench
-from sglang.srt.layers.quantization.fp8_kernel import (
+from sglang.kernels.ops.quantization.fp8_kernel import (
     sglang_per_token_group_quant_fp8,
     w8a8_block_fp8_matmul_deepgemm,
 )

benchmark/kernels/flashinfer_allreduce_fusion/benchmark_fused_collective.py

@@ -28,6 +28,8 @@ from typing import Optional
 import torch  # type: ignore
 import torch.distributed as dist  # type: ignore
 
+from sglang.kernels.ops.quantization.fp8_kernel import fp8_dtype as SGLANG_FP8_DTYPE
+from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
 from sglang.srt.distributed import get_tp_group, tensor_model_parallel_all_reduce
 from sglang.srt.distributed.parallel_state import (
     cleanup_dist_env_and_memory,
@@ -36,8 +38,6 @@ from sglang.srt.distributed.parallel_state import (
     initialize_model_parallel,
 )
 from sglang.srt.layers.layernorm import RMSNorm  # noqa
-from sglang.srt.layers.quantization.fp8_kernel import fp8_dtype as SGLANG_FP8_DTYPE
-from sglang.srt.layers.quantization.fp8_kernel import static_quant_fp8
 
 try:
     from sgl_kernel import fused_add_rmsnorm as SGL_FUSED_ADD_RMS_NORM

benchmark/kernels/quantization/bench_int8_quant.py

@@ -5,7 +5,7 @@ import triton
 from vllm._custom_ops import scaled_int8_quant as vllm_scaled_int8_quant
 
 from sglang.benchmark.bench_utils import run_bench
-from sglang.srt.layers.quantization.int8_kernel import per_token_quant_int8
+from sglang.kernels.ops.quantization.int8_kernel import per_token_quant_int8
 
 
 @torch.compile(backend="inductor")

benchmark/kernels/quantization/tuning_block_wise_kernel.py

@@ -27,11 +27,11 @@ from tqdm import tqdm
 
 mp.set_start_method("spawn", force=True)
 
-from sglang.srt.layers.quantization.fp8_kernel import (
+from sglang.kernels.ops.quantization.fp8_kernel import (
     _w8a8_block_fp8_matmul,
     _w8a8_block_fp8_matmul_unrolledx4,
 )
-from sglang.srt.layers.quantization.int8_kernel import _w8a8_block_int8_matmul
+from sglang.kernels.ops.quantization.int8_kernel import _w8a8_block_int8_matmul
 from sglang.srt.utils import (
     get_device,
     get_device_core_count,
@@ -521,7 +521,9 @@ if __name__ == "__main__":
     parser.add_argument("--block-k", type=int, default=128)
     parser.add_argument("--batch-sizes", nargs="+", type=int, required=False)
     parser.add_argument(
-        "--save-path", type=str, default="python/sglang/srt/layers/quantization/configs"
+        "--save-path",
+        type=str,
+        default="python/sglang/kernels/ops/quantization/configs",
     )
     args = parser.parse_args()
 

python/sglang/jit_kernel/tests/test_minimax_m3_mxfp8.py

@@ -27,7 +27,7 @@ if not is_hip():
 if not torch.cuda.is_available():
     pytest.skip("Requires a GPU.", allow_module_level=True)
 
-from sglang.srt.layers.quantization.mxfp8_amd_gfx95 import (  # noqa: E402
+from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import (  # noqa: E402
     _mxfp8_dot_scaled_linear,
     _mxfp8_e4m3_quantize_torch,
     _mxfp8_e4m3_quantize_triton,
@@ -90,7 +90,7 @@ def test_minimax_swiglu_mxfp8_quant_matches_unfused_fp32(m, inter):
         swiglu_oai_mxfp8_quant,
         swiglu_oai_split,
     )
-    from sglang.srt.layers.quantization.mxfp8_amd_gfx95 import mxfp8_e4m3_quantize
+    from sglang.kernels.ops.quantization.mxfp8_amd_gfx95 import mxfp8_e4m3_quantize
 
     torch.manual_seed(0)
     alpha, beta, limit = 1.702, 1.0, 7.0

python/sglang/jit_kernel/triton_store_cache.py

@@ -4,7 +4,7 @@ import torch
 import triton
 import triton.language as tl
 
-from sglang.srt.layers.quantization.fp8_kernel import is_fp8_fnuz
+from sglang.kernels.ops.quantization.fp8_kernel import is_fp8_fnuz
 
 _FP8_DTYPE = torch.float8_e4m3fnuz if is_fp8_fnuz() else torch.float8_e4m3fn
 _FP8_INFO = torch.finfo(_FP8_DTYPE)

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

@@ -117,3 +117,14 @@ for _mod, _fn in _TRITON_KERNELS:
         )
     )
 del _mod, _fn
+
+
+# Packed (topk_id << 16 | bf16-weight) kernel migrated from
+# srt/layers/quantization/mxfp4_flashinfer_trtllm_moe (RFC #29630, Phase 2.5).
+register_kernel(
+    KernelSpec(
+        op="moe.pack_topk_ids",
+        backend=KernelBackend.TRITON,
+        target="sglang.kernels.ops.moe.pack_topk_ids:PackTopkIds.triton",
+    )
+)

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

@@ -0,0 +1,86 @@
+"""Pack ``(topk_id, topk_weight)`` pairs into one int32 per entry.
+
+Migrated from ``sglang.srt.layers.quantization.mxfp4_flashinfer_trtllm_moe``
+(RFC #29630, Phase 2.5). Used by the FlashInfer TRT-LLM routed-MoE path, which
+consumes routing ids and bf16 weights packed as ``(id << 16) | weight_bits``.
+"""
+
+import torch
+import triton
+import triton.language as tl
+
+
+class PackTopkIds:
+
+    @classmethod
+    def execute(
+        cls, topk_ids: torch.Tensor, topk_weights: torch.Tensor
+    ) -> torch.Tensor:
+        return cls.triton(topk_ids, topk_weights)
+
+    @classmethod
+    def vanilla(
+        cls, topk_ids: torch.Tensor, topk_weights: torch.Tensor
+    ) -> torch.Tensor:
+        weight_bits = (
+            topk_weights.to(torch.bfloat16).view(torch.int16).to(torch.int32) & 0xFFFF
+        )
+        return (topk_ids.to(torch.int32) << 16) | weight_bits
+
+    @classmethod
+    def triton(cls, topk_ids: torch.Tensor, topk_weights: torch.Tensor) -> torch.Tensor:
+        assert (
+            topk_ids.shape == topk_weights.shape
+        ), f"shape mismatch: {topk_ids.shape=} vs {topk_weights.shape=}"
+        assert topk_ids.ndim >= 1, f"expected >=1D, got {topk_ids.shape=}"
+
+        assert (
+            topk_ids.dtype == torch.int32
+        ), f"topk_ids must be int32, got {topk_ids.dtype}"
+        assert (
+            topk_weights.dtype == torch.float32
+        ), f"topk_weights must be float32, got {topk_weights.dtype}"
+
+        assert topk_ids.is_contiguous(), "topk_ids must be contiguous"
+        assert topk_weights.is_contiguous(), "topk_weights must be contiguous"
+
+        out = torch.empty_like(topk_ids, dtype=torch.int32)
+        numel = out.numel()
+        if numel == 0:
+            return out
+
+        BLOCK_SIZE = 1024
+        grid = (triton.cdiv(numel, BLOCK_SIZE),)
+        _pack_topk_ids_triton_kernel[grid](
+            topk_ids,
+            topk_weights,
+            out,
+            numel,
+            BLOCK_SIZE=BLOCK_SIZE,
+        )
+        return out
+
+
+@triton.jit
+def _pack_topk_ids_triton_kernel(
+    topk_ids_ptr,
+    topk_weights_ptr,
+    out_ptr,
+    numel,
+    BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
+    mask = offsets < numel
+
+    ids = tl.load(topk_ids_ptr + offsets, mask=mask, other=0)
+    w = tl.load(topk_weights_ptr + offsets, mask=mask, other=0.0)
+
+    w_bf16 = w.to(tl.bfloat16)
+    w_i16 = w_bf16.to(tl.int16, bitcast=True)
+    w_i32 = w_i16.to(tl.int32) & 0xFFFF
+
+    ids_i32 = ids.to(tl.int32)
+    packed = (ids_i32 << 16) | w_i32
+
+    tl.store(out_ptr + offsets, packed, mask=mask)

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

@@ -120,3 +120,50 @@ __all__ = [
     "sgl_per_token_group_quant_fp8",
     "sgl_per_token_group_quant_int8",
 ]
+
+
+# Triton / CuTe DSL kernels migrated into this group from
+# srt/layers/quantization (RFC #29630, Phase 2.5); registered for inventory.
+# Import them from their modules.
+_TRITON_KERNELS = [
+    ("fp8_kernel", "per_token_group_quant_8bit"),
+    ("fp8_kernel", "sglang_per_token_group_quant_fp8"),
+    ("fp8_kernel", "sglang_per_token_group_quant_8bit"),
+    ("fp8_kernel", "sglang_per_token_quant_fp8"),
+    ("fp8_kernel", "static_quant_fp8"),
+    ("fp8_kernel", "w8a8_block_fp8_matmul"),
+    ("fp8_kernel", "mxfp8_block_scaled_matmul_triton"),
+    ("fp8_kernel", "per_tensor_quant_mla_fp8"),
+    ("fp8_kernel", "per_token_group_quant_mla_deep_gemm_masked_fp8"),
+    ("fp8_kernel", "per_token_group_quant_fp8_hopper_moe_mn_major"),
+    ("fp8_kernel", "per_group_transpose"),
+    ("fp8_kernel", "triton_scaled_mm"),
+    ("int8_kernel", "per_token_quant_int8"),
+    ("int8_kernel", "per_token_group_quant_int8"),
+    ("int8_kernel", "w8a8_block_int8_matmul"),
+    ("awq_triton", "awq_dequantize_triton"),
+    ("awq_triton", "awq_gemm_triton"),
+    ("mxfp8_amd_gfx95", "mxfp8_e4m3_quantize"),
+]
+for _mod, _fn in _TRITON_KERNELS:
+    register_kernel(
+        KernelSpec(
+            op=f"quantization.{_fn}",
+            backend=KernelBackend.TRITON,
+            target=f"sglang.kernels.ops.quantization.{_mod}:{_fn}",
+        )
+    )
+del _mod, _fn
+
+register_kernel(
+    KernelSpec(
+        op="quantization.nvfp4_gemm_swiglu_nvfp4_quant",
+        backend=KernelBackend.CUTE_DSL,
+        target=(
+            "sglang.kernels.ops.quantization.nvfp4_gemm_swiglu_nvfp4_quant"
+            ":nvfp4_gemm_swiglu_nvfp4_quant"
+        ),
+        capability=CapabilityRequirement(requires_cuda=True, min_cuda_arch=(10, 0)),
+        description="Fused NVFP4 GEMM + SwiGLU + NVFP4 quant (CuTe DSL, SM100).",
+    )
+)