[JIT kernel] Apply jit per_tensor_quant_fp8 kernel (#15836)

python/sglang/jit_kernel/csrc/gemm/per_tensor_quant_fp8.cuh

@@ -104,31 +104,51 @@ template <bool kIsStatic>
 void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView output_q, tvm::ffi::TensorView output_s) {
   using namespace host;
 
-  SymbolicSize num_tokens = {"num_tokens"};
-  SymbolicSize hidden_dim = {"hidden_dim"};
-  SymbolicDevice device_;
-  SymbolicDType input_dtype;
+  const DLDevice device = input.device();
+  RuntimeCheck(device.device_type == kDLCUDA, "input must be on CUDA");
+  RuntimeCheck(input.is_contiguous(), "input must be contiguous");
 
-  TensorMatcher({num_tokens, hidden_dim})  //
-      .with_dtype<float, __half, __nv_bfloat16>(input_dtype)
-      .with_device<kDLCUDA>(device_)
-      .verify(input);
+  const int64_t ndim = input.dim();
+  RuntimeCheck(ndim >= 1, "input.ndim must be >= 1, but got ", ndim);
 
-  TensorMatcher({num_tokens, hidden_dim})  //
-      .with_dtype<__nv_fp8_e4m3>()
-      .with_device<kDLCUDA>(device_)
-      .verify(output_q);
+  RuntimeCheck(output_q.device() == device, "output_q must be on the same device as input");
+  RuntimeCheck(output_q.is_contiguous(), "output_q must be contiguous");
+  RuntimeCheck(output_q.dim() == ndim, "output_q.ndim must match input.ndim");
+  for (int64_t i = 0; i < ndim; ++i) {
+    RuntimeCheck(
+        output_q.size(i) == input.size(i),
+        "output_q.shape mismatch at dim ",
+        i,
+        ": expected ",
+        input.size(i),
+        " but got ",
+        output_q.size(i));
+  }
 
   TensorMatcher({1})  //
       .with_dtype<float>()
-      .with_device<kDLCUDA>(device_)
+      .with_device<kDLCUDA>()
       .verify(output_s);
+  RuntimeCheck(output_s.device() == device, "output_s must be on the same device as input");
 
-  const size_t total_elements = num_tokens.unwrap() * hidden_dim.unwrap();
+  const DLDataType in_dtype = input.dtype();
+  const bool in_ok = (in_dtype.code == kDLFloat && in_dtype.bits == 32) ||
+                     (in_dtype.code == kDLFloat && in_dtype.bits == 16) ||
+                     (in_dtype.code == kDLBfloat && in_dtype.bits == 16);
+  RuntimeCheck(in_ok, "input dtype must be fp32/fp16/bf16, but got ", in_dtype);
+
+  const DLDataType out_dtype = output_q.dtype();
+  RuntimeCheck(
+      out_dtype.code == kDLFloat8_e4m3fn && out_dtype.bits == 8,
+      "output_q dtype must be fp8_e4m3fn, but got ",
+      out_dtype);
+
+  size_t total_elements = 1;
+  for (const auto s : input.shape()) {
+    RuntimeCheck(s > 0, "Input tensor must be non-empty");
+    total_elements *= static_cast<size_t>(s);
+  }
   const size_t num_blocks = std::min((total_elements + kBlockSize - 1) / kBlockSize, size_t(1024));
-  const DLDevice device = device_.unwrap();
-
-  RuntimeCheck(total_elements > 0, "Input tensor must be non-empty");
 
   auto launch_kernels = [&]<typename T>() {
     if constexpr (!kIsStatic) {
@@ -147,12 +167,11 @@ void per_tensor_quant_fp8(tvm::ffi::TensorView input, tvm::ffi::TensorView outpu
         static_cast<int64_t>(total_elements));
   };
 
-  const DLDataType dtype = input_dtype.unwrap();
-  if (dtype.code == kDLFloat && dtype.bits == 32) {
+  if (in_dtype.code == kDLFloat && in_dtype.bits == 32) {
     launch_kernels.template operator()<float>();
-  } else if (dtype.code == kDLBfloat && dtype.bits == 16) {
+  } else if (in_dtype.code == kDLBfloat && in_dtype.bits == 16) {
     launch_kernels.template operator()<__nv_bfloat16>();
-  } else if (dtype.code == kDLFloat && dtype.bits == 16) {
+  } else if (in_dtype.code == kDLFloat && in_dtype.bits == 16) {
     launch_kernels.template operator()<__half>();
   }
 }

python/sglang/jit_kernel/per_tensor_quant_fp8.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-import functools
 import os
 from typing import TYPE_CHECKING
 
@@ -8,13 +7,14 @@ import flashinfer
 import torch
 from torch.utils.cpp_extension import CUDA_HOME
 
-from sglang.jit_kernel.utils import load_jit, make_cpp_args
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+from sglang.srt.utils.custom_op import register_custom_op
 
 if TYPE_CHECKING:
     from tvm_ffi.module import Module
 
 
-@functools.cache
+@cache_once
 def _jit_per_tensor_quant_fp8_module(is_static: bool) -> Module:
     args = make_cpp_args(is_static)
 
@@ -32,6 +32,10 @@ def _jit_per_tensor_quant_fp8_module(is_static: bool) -> Module:
     )
 
 
+@register_custom_op(
+    op_name="per_tensor_quant_fp8",
+    mutates_args=["output_q", "output_s"],
+)
 def per_tensor_quant_fp8(
     input: torch.Tensor,
     output_q: torch.Tensor,
@@ -44,8 +48,13 @@ def per_tensor_quant_fp8(
     Args:
         input: Input tensor to quantize (float, half, or bfloat16)
         output_q: Output quantized tensor (fp8_e4m3)
-        output_s: Output scale tensor (float scalar)
+        output_s: Output scale tensor (float scalar or 1D tensor with 1 element)
         is_static: If True, assumes scale is pre-computed and skips absmax computation
     """
+    # Ensure output_s has shape [1] instead of being a 0D scalar
+    # The JIT kernel expects a 1D tensor
+    if output_s.ndim == 0:
+        output_s = output_s.reshape(1)
+
     module = _jit_per_tensor_quant_fp8_module(is_static)
     module.per_tensor_quant_fp8(input, output_q, output_s)

python/sglang/jit_kernel/tests/test_per_tensor_quant_fp8.py

@@ -57,6 +57,22 @@ def test_jit_per_tensor_quant_compare_implementations(
         sglang_out.float(), torch_out.float(), rtol=1e-3, atol=1e-3
     )
 
+
+@pytest.mark.parametrize("shape", [(4, 8, 64), (2, 16, 128)])
+def test_jit_per_tensor_quant_supports_3d(shape):
+    device = torch.device("cuda")
+    x = torch.rand(shape, dtype=torch.bfloat16, device=device)
+    out = torch.empty_like(x, device=x.device, dtype=fp8_type_)
+    scale = torch.zeros(1, device=x.device, dtype=torch.float32)
+
+    per_tensor_quant_fp8(x, out, scale, is_static=False)
+
+    x_2d = x.flatten(0, -2)
+    out_ref_2d = torch_scaled_fp8_quant(x_2d, scale)
+    out_ref = out_ref_2d.reshape(shape)
+
+    torch.testing.assert_close(out.float(), out_ref.float(), rtol=1e-3, atol=1e-3)
+
     scale = torch.rand(1, dtype=torch.float32, device=device)
     sglang_out, sglang_scale = sglang_scaled_fp8_quant(x, scale)
     torch_out = torch_scaled_fp8_quant(x, scale)

python/sglang/srt/layers/moe/cutlass_w4a8_moe.py

@@ -3,13 +3,9 @@
 from typing import Optional
 
 import torch
-from sgl_kernel import (
-    cutlass_w4a8_moe_mm,
-    get_cutlass_w4a8_moe_mm_data,
-    sgl_per_tensor_quant_fp8,
-    silu_and_mul,
-)
+from sgl_kernel import cutlass_w4a8_moe_mm, get_cutlass_w4a8_moe_mm_data, silu_and_mul
 
+from sglang.jit_kernel.per_tensor_quant_fp8 import per_tensor_quant_fp8
 from sglang.srt.distributed import get_moe_expert_parallel_world_size
 from sglang.srt.layers.moe.ep_moe.kernels import (
     cutlass_w4_run_moe_ep_preproess,
@@ -321,9 +317,7 @@ def cutlass_w4a8_moe_deepep_normal(
     gateup_input = torch.empty(
         gateup_input_pre_reorder.shape, dtype=torch.float8_e4m3fn, device=device
     )
-    sgl_per_tensor_quant_fp8(
-        gateup_input_pre_reorder, gateup_input, a1_scale.float(), True
-    )
+    per_tensor_quant_fp8(gateup_input_pre_reorder, gateup_input, a1_scale.float(), True)
     del gateup_input_pre_reorder
     local_topk_ids = topk_ids_
     local_topk_ids = (
@@ -367,7 +361,7 @@ def cutlass_w4a8_moe_deepep_normal(
     intermediate_q = torch.empty(
         intermediate.shape, dtype=torch.float8_e4m3fn, device=device
     )
-    sgl_per_tensor_quant_fp8(intermediate, intermediate_q, a2_scale.float(), True)
+    per_tensor_quant_fp8(intermediate, intermediate_q, a2_scale.float(), True)
 
     cutlass_w4a8_moe_mm(
         c2,
@@ -495,7 +489,7 @@ def cutlass_w4a8_moe_deepep_ll(
     )
 
     gateup_input = torch.empty(a.shape, dtype=torch.float8_e4m3fn, device=device)
-    sgl_per_tensor_quant_fp8(a, gateup_input, a1_scale.float(), True)
+    per_tensor_quant_fp8(a, gateup_input, a1_scale.float(), True)
     c1 = torch.empty((num_experts, m, n * 2), device=device, dtype=torch.bfloat16)
     c2 = torch.empty((num_experts, m, k), device=device, dtype=torch.bfloat16)
 

python/sglang/srt/layers/quantization/fp8_kernel.py

@@ -42,7 +42,11 @@ _is_cpu = is_cpu()
 _use_aiter = get_bool_env_var("SGLANG_USE_AITER") and _is_hip
 
 if _is_cuda:
-    from sgl_kernel import sgl_per_tensor_quant_fp8, sgl_per_token_quant_fp8
+    from sgl_kernel import sgl_per_token_quant_fp8
+
+    from sglang.jit_kernel.per_tensor_quant_fp8 import (
+        per_tensor_quant_fp8 as sgl_per_tensor_quant_fp8,
+    )
 
     # Temporary
     try:
@@ -1864,7 +1868,3 @@ if _is_cuda:
         @torch.library.register_fake("sgl_kernel::sgl_per_token_quant_fp8")
         def _(input, output_q, output_s):
             return
-
-    @torch.library.register_fake("sgl_kernel::sgl_per_tensor_quant_fp8")
-    def _sgl_per_tensor_quant_fp8(input, output_q, output_s, is_static):
-        return