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

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).",
+    )
+)