Revert "Add flashinfer rmsnorm + quant fusion support SM90, SM100, SM120" (#33455)

benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py

@@ -1,185 +0,0 @@
-"""Microbenchmark: fused RMSNorm + static per-tensor FP8 quant, comparing the
-flashinfer default kernels against the CuTe-DSL kernels and the unfused
-baseline (RMSNorm followed by a separate static FP8 quant).
-
-Providers:
-  unfused     RMSNorm.forward_cuda  +  static_quant_fp8
-  fused       flashinfer rmsnorm_quant / fused_add_rmsnorm_quant (default)
-  fused_cute  flashinfer rmsnorm_quant_cute / fused_add_rmsnorm_quant_cute
-
-All fused providers produce an ``(fp8, scale)`` activation (and updated residual
-when a residual is supplied), matching what a downstream FP8 static-per-tensor
-linear consumes. Covers the no-residual and residual (fused-add) cases across a
-few hidden sizes so you can pick the fastest kernel per shape.
-
-Run:
-    python benchmark/kernels/bench_fused_rmsnorm_fp8_quant.py
-"""
-
-import itertools
-
-import numpy as np
-import torch
-import triton
-from flashinfer.norm import fused_add_rmsnorm_quant, rmsnorm_quant
-from flashinfer.testing import bench_gpu_time
-
-from sglang.kernels.ops.quantization.fp8_kernel import static_quant_fp8
-from sglang.srt.layers.layernorm import RMSNorm, _flashinfer_rmsnorm_quant_available
-
-if not torch.cuda.is_available():
-    raise RuntimeError("CUDA is required for this benchmark")
-if not _flashinfer_rmsnorm_quant_available:
-    raise RuntimeError(
-        "flashinfer rmsnorm_quant / fused_add_rmsnorm_quant is not available; "
-        "install flashinfer to benchmark the fused path"
-    )
-
-try:
-    from flashinfer.norm import fused_add_rmsnorm_quant_cute, rmsnorm_quant_cute
-
-    _CUTE_AVAILABLE = True
-except ImportError:
-    _CUTE_AVAILABLE = False
-
-DEVICE = "cuda"
-DTYPE = torch.bfloat16
-FP8_DTYPE = torch.float8_e4m3fn
-HIDDEN_SIZES = [4096, 8192]
-# Per-tensor reciprocal scale (q = normed / scale); 0.05 keeps normed/scale well
-# within the e4m3 range for unit-scale activations.
-SCALE_VALUE = 0.05
-
-
-def make_layer(hidden_size):
-    layer = RMSNorm(hidden_size).to(device=DEVICE, dtype=DTYPE)
-    layer.weight.data.normal_(mean=1.0, std=0.1)
-    return layer
-
-
-def make_inputs(num_tokens, hidden_size, add_residual):
-    x = torch.randn(num_tokens, hidden_size, device=DEVICE, dtype=DTYPE)
-    residual = torch.randn_like(x) if add_residual else None
-    scale = torch.tensor([SCALE_VALUE], device=DEVICE, dtype=torch.float32)
-    return x, residual, scale
-
-
-def run_unfused(layer, x, residual, scale):
-    out = layer(x, residual)
-    if residual is not None:
-        normed, residual_out = out
-        q, q_scale = static_quant_fp8(normed, scale)
-        return (q, q_scale), residual_out
-    q, q_scale = static_quant_fp8(out, scale)
-    return q, q_scale
-
-
-def _run_fused(kernel, add_kernel, layer, x, residual, scale):
-    out = torch.empty_like(x, dtype=FP8_DTYPE)
-    if residual is not None:
-        # In-place: residual += x, then out = quant(rmsnorm(residual) * w).
-        add_kernel(out, x, residual, layer.weight.data, scale, layer.variance_epsilon)
-        return (out, scale), residual
-    kernel(out, x, layer.weight.data, scale, layer.variance_epsilon)
-    return out, scale
-
-
-def run_fused_default(layer, x, residual, scale):
-    return _run_fused(rmsnorm_quant, fused_add_rmsnorm_quant, layer, x, residual, scale)
-
-
-def run_fused_cute(layer, x, residual, scale):
-    return _run_fused(
-        rmsnorm_quant_cute, fused_add_rmsnorm_quant_cute, layer, x, residual, scale
-    )
-
-
-RUNNERS = {
-    "unfused": run_unfused,
-    "fused": run_fused_default,
-    "fused_cute": run_fused_cute,
-}
-
-# (provider key, plot label, style)
-_PROVIDERS = [
-    ("unfused", "rmsnorm + static_quant_fp8 (unfused)", ("blue", "-")),
-    ("fused", "rmsnorm_quant (fused, default)", ("green", "-")),
-]
-if _CUTE_AVAILABLE:
-    _PROVIDERS.append(
-        ("fused_cute", "rmsnorm_quant_cute (fused, cute-dsl)", ("red", "-"))
-    )
-
-
-def _bench_ms(fn, args, quantiles=(0.5, 0.2, 0.8)):
-    # Pass the GPU tensors as input_args so flashinfer's cold_l2_cache flush can
-    # find them; a zero-arg callable trips its "no GPU tensors found" warning and
-    # silently disables cold-L2 timing.
-    times = bench_gpu_time(
-        fn=fn,
-        input_args=args,
-        use_cuda_graph=True,
-        dry_run_time_ms=25,
-        repeat_time_ms=100,
-    )
-    return tuple(float(np.percentile(times, q * 100)) for q in quantiles)
-
-
-def _check_correctness():
-    """One-shot sanity check that every fused provider agrees with the unfused
-    baseline within FP8 precision."""
-    fused_providers = [p for p in RUNNERS if p != "unfused"]
-    for hidden_size, add_residual in itertools.product(HIDDEN_SIZES, [False, True]):
-        layer = make_layer(hidden_size)
-        x, residual, scale = make_inputs(64, hidden_size, add_residual)
-        with torch.inference_mode():
-            ref = run_unfused(
-                layer, x.clone(), residual.clone() if add_residual else None, scale
-            )
-        (uq, _), _ = ref if add_residual else (ref, None)
-        ref_deq = uq.float() * scale
-        for provider in fused_providers:
-            if provider == "fused_cute" and not _CUTE_AVAILABLE:
-                continue
-            with torch.inference_mode():
-                out = RUNNERS[provider](
-                    layer, x.clone(), residual.clone() if add_residual else None, scale
-                )
-            (q, _), _ = out if add_residual else (out, None)
-            cos = torch.nn.functional.cosine_similarity(
-                (q.float() * scale).flatten(), ref_deq.flatten(), dim=0
-            ).item()
-            assert (
-                cos > 0.99
-            ), f"{provider} h={hidden_size} residual={add_residual} cos={cos:.4f}"
-    print("correctness check passed (all fused providers vs unfused within FP8)")
-
-
-configs = [
-    triton.testing.Benchmark(
-        x_names=["num_tokens"],
-        x_vals=[512, 1024, 2048, 4096, 8192, 16384],
-        x_log=False,
-        line_arg="provider",
-        line_vals=[p[0] for p in _PROVIDERS],
-        line_names=[p[1] for p in _PROVIDERS],
-        styles=[p[2] for p in _PROVIDERS],
-        ylabel="latency (ms)",
-        plot_name=f"rmsnorm_fp8_quant_h{hidden_size}_residual{add_residual}",
-        args={"hidden_size": hidden_size, "add_residual": add_residual},
-    )
-    for hidden_size, add_residual in itertools.product(HIDDEN_SIZES, [False, True])
-]
-
-
-@triton.testing.perf_report(configs)
-def benchmark(num_tokens, hidden_size, add_residual, provider):
-    layer = make_layer(hidden_size)
-    x, residual, scale = make_inputs(num_tokens, hidden_size, add_residual)
-    return _bench_ms(RUNNERS[provider], (layer, x, residual, scale))
-
-
-if __name__ == "__main__":
-    torch.manual_seed(0)
-    _check_correctness()
-    benchmark.run(print_data=True, show_plots=False)

python/sglang/kernels/aot/benchmark/bench_fp8_gemm.py

@@ -8,7 +8,9 @@ import torch
 import triton
 from sgl_kernel import fp8_scaled_mm as sgl_scaled_mm
 
-from sglang.kernels.ops.quantization.per_tensor_quant_fp8 import per_tensor_quant_fp8
+from sglang.kernels.ops.quantization.per_tensor_quant_fp8 import (
+    per_tensor_quant_fp8,
+)
 from sglang.utils import is_in_ci
 
 # Optional vLLM import
@@ -104,7 +106,7 @@ def sglang_scaled_fp8_quant(
 if IS_CI:
     batch_sizes = [1]  # Single batch size for CI
 else:
-    batch_sizes = [1, 2, 8, 16, 64, 128, 256, 512, 1024, 2048]
+    batch_sizes = [1, 16, 64, 128, 256, 512, 1024, 2048]
 
 # Filter line_vals based on vLLM availability
 if VLLM_AVAILABLE:
@@ -113,39 +115,24 @@ if VLLM_AVAILABLE:
         "vllm-fp8-bf16",
         "sglang-fp8-fp16",
         "sglang-fp8-bf16",
-        "sglang-scalar-a-fp8-fp16",
-        "sglang-scalar-a-fp8-bf16",
     ]
     line_names = [
         "vllm-fp8-fp16",
         "vllm-fp8-bf16",
         "sglang-fp8-fp16",
         "sglang-fp8-bf16",
-        "sglang-scalar-a-fp8-fp16",
-        "sglang-scalar-a-fp8-bf16",
-    ]
-    styles = [
-        ("green", "-"),
-        ("green", "--"),
-        ("blue", "-"),
-        ("blue", "--"),
-        ("red", "-"),
-        ("red", "--"),
     ]
+    styles = [("green", "-"), ("green", "--"), ("blue", "-"), ("blue", "--")]
 else:
     line_vals = [
         "sglang-fp8-fp16",
         "sglang-fp8-bf16",
-        "sglang-scalar-a-fp8-fp16",
-        "sglang-scalar-a-fp8-bf16",
     ]
     line_names = [
         "sglang-fp8-fp16",
         "sglang-fp8-bf16",
-        "sglang-scalar-a-fp8-fp16",
-        "sglang-scalar-a-fp8-bf16",
     ]
-    styles = [("blue", "-"), ("blue", "--"), ("red", "-"), ("red", "--")]
+    styles = [("blue", "-"), ("blue", "--")]
 
 
 @triton.testing.perf_report(
@@ -187,9 +174,8 @@ def benchmark(batch_size, provider, N, K):
             lambda: vllm_scaled_mm(a_fp8, b_fp8, scale_a_fp8, scale_b_fp8, dtype),
             quantiles=quantiles,
         )
-    elif "sglang" in provider:
+    elif "sglang-fp8" in provider:
-        a_scale = scale_a_scalar if "scalar-a" in provider else scale_a
+        a_fp8, scale_a_fp8 = sglang_scaled_fp8_quant(a, scale_a)
-        a_fp8, scale_a_fp8 = sglang_scaled_fp8_quant(a, a_scale)
         b_fp8, scale_b_fp8 = sglang_scaled_fp8_quant(b, scale_b)
         b_fp8 = b_fp8.t()
         ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(

python/sglang/kernels/aot/csrc/gemm/fp8_gemm_kernel.cu

@@ -448,22 +448,19 @@ template <
     typename MainloopScheduleType,
     typename EpilogueScheduleType,
     typename TileSchedulerType = void,
-    bool WithBias = false,
+    bool WithBias = false>
-    bool ScalarA = false>
 struct DeviceGemmFp8RowwiseSm100 {
   static_assert(std::is_same_v<ElementType, cutlass::float_e4m3_t>, "ElementType must be FP8(e4m3)");
   using TileShape = CTAShape;
   using Accum = cutlass::epilogue::fusion::Sm90AccFetch;
 
   using ElementComputeEpilogue = float;
-  using VectorScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast<
+  using ScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast<
       0,
       TileShape,
       ElementComputeEpilogue,
       ElementComputeEpilogue,
       cute::Stride<cute::Int<1>, cute::Int<0>, cute::Int<0>>>;
-  using ScalarScaleA = cutlass::epilogue::fusion::Sm90ScalarBroadcast<float>;
-  using ScaleA = std::conditional_t<ScalarA, ScalarScaleA, VectorScaleA>;
 
   using ScaleB = cutlass::epilogue::fusion::Sm90RowBroadcast<
       0,
@@ -554,11 +551,7 @@ struct DeviceGemmFp8RowwiseSm100 {
     auto* data_ptr = static_cast<T*>(tensor.data_ptr());
     static_assert(
         std::is_same_v<Descriptor, ScaleA> || std::is_same_v<Descriptor, ScaleB> || std::is_same_v<Descriptor, Bias>);
-    if constexpr (std::is_same_v<Descriptor, ScalarScaleA>) {
+    return Arguments{data_ptr};
-      return Arguments{{}, {data_ptr}, {}};
-    } else {
-      return Arguments{data_ptr};
-    }
   }
 
  public:
@@ -664,7 +657,7 @@ void launch_sm100_fp8_scaled_mm(
   TORCH_CHECK(status == cutlass::Status::kSuccess)
 }
 
-template <typename OutType, bool ScalarA>
+template <typename OutType>
 void sm100_fp8_dispatch_bias(
     torch::Tensor& out,
     const torch::Tensor& a,
@@ -702,8 +695,7 @@ void sm100_fp8_dispatch_bias(
       MainloopScheduleType,
       EpilogueScheduleType,
       TileSchedulerType,
-      true,
+      true>;
-      ScalarA>;
   using BiasGemm256 = DeviceGemmFp8RowwiseSm100<
       ElementInput,
       ElementOutput,
@@ -713,8 +705,7 @@ void sm100_fp8_dispatch_bias(
       MainloopScheduleType,
       EpilogueScheduleType,
       TileSchedulerType,
-      true,
+      true>;
-      ScalarA>;
   using BiasGemm64 = DeviceGemmFp8RowwiseSm100<
       ElementInput,
       ElementOutput,
@@ -724,8 +715,7 @@ void sm100_fp8_dispatch_bias(
       MainloopScheduleType,
       EpilogueScheduleType,
       TileSchedulerType,
-      true,
+      true>;
-      ScalarA>;
   using BiasGemm16 = DeviceGemmFp8RowwiseSm100<
       ElementInput,
       ElementOutput,
@@ -735,8 +725,7 @@ void sm100_fp8_dispatch_bias(
       MainloopScheduleType,
       EpilogueScheduleType,
       TileSchedulerType,
-      true,
+      true>;
-      ScalarA>;
 
   // Gemm type without bias
   using GemmDefault = DeviceGemmFp8RowwiseSm100<
@@ -748,8 +737,7 @@ void sm100_fp8_dispatch_bias(
       MainloopScheduleType,
       EpilogueScheduleType,
       TileSchedulerType,
-      false,
+      false>;
-      ScalarA>;
   using Gemm256 = DeviceGemmFp8RowwiseSm100<
       ElementInput,
       ElementOutput,
@@ -759,8 +747,7 @@ void sm100_fp8_dispatch_bias(
       MainloopScheduleType,
       EpilogueScheduleType,
       TileSchedulerType,
-      false,
+      false>;
-      ScalarA>;
   using Gemm64 = DeviceGemmFp8RowwiseSm100<
       ElementInput,
       ElementOutput,
@@ -770,8 +757,7 @@ void sm100_fp8_dispatch_bias(
       MainloopScheduleType,
       EpilogueScheduleType,
       TileSchedulerType,
-      false,
+      false>;
-      ScalarA>;
   using Gemm16 = DeviceGemmFp8RowwiseSm100<
       ElementInput,
       ElementOutput,
@@ -781,8 +767,7 @@ void sm100_fp8_dispatch_bias(
       MainloopScheduleType,
       EpilogueScheduleType,
       TileSchedulerType,
-      false,
+      false>;
-      ScalarA>;
 
   // next power of 2 (minimum 16)
   uint32_t const m = a.size(0);
@@ -826,10 +811,7 @@ void sm100_fp8_dispatch_shape(
     const torch::Tensor& scales_a,
     const torch::Tensor& scales_b,
     const c10::optional<torch::Tensor>& bias) {
-  if (scales_a.numel() == 1) {
+  return sm100_fp8_dispatch_bias<OutType>(out, a, b, scales_a, scales_b, bias);
-    return sm100_fp8_dispatch_bias<OutType, true>(out, a, b, scales_a, scales_b, bias);
-  }
-  return sm100_fp8_dispatch_bias<OutType, false>(out, a, b, scales_a, scales_b, bias);
 }
 
 template <
@@ -841,22 +823,19 @@ template <
     typename MainloopScheduleType,
     typename EpilogueScheduleType,
     typename TileSchedulerType = void,
-    bool WithBias = false,
+    bool WithBias = false>
-    bool ScalarA = false>
 struct DeviceGemmFp8RowwiseSm120 {
   static_assert(std::is_same_v<ElementType, cutlass::float_e4m3_t>, "ElementType must be FP8(e4m3)");
   using TileShape = CTAShape;
   using Accum = cutlass::epilogue::fusion::Sm90AccFetch;
 
   using ElementComputeEpilogue = float;
-  using VectorScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast<
+  using ScaleA = cutlass::epilogue::fusion::Sm90ColBroadcast<
       0,
       TileShape,
       ElementComputeEpilogue,
       ElementComputeEpilogue,
       cute::Stride<cute::Int<1>, cute::Int<0>, cute::Int<0>>>;
-  using ScalarScaleA = cutlass::epilogue::fusion::Sm90ScalarBroadcast<float>;
-  using ScaleA = std::conditional_t<ScalarA, ScalarScaleA, VectorScaleA>;
 
   using ScaleB = cutlass::epilogue::fusion::Sm90RowBroadcast<
       0,
@@ -947,11 +926,7 @@ struct DeviceGemmFp8RowwiseSm120 {
     auto* data_ptr = static_cast<T*>(tensor.data_ptr());
     static_assert(
         std::is_same_v<Descriptor, ScaleA> || std::is_same_v<Descriptor, ScaleB> || std::is_same_v<Descriptor, Bias>);
-    if constexpr (std::is_same_v<Descriptor, ScalarScaleA>) {
+    return Arguments{data_ptr};
-      return Arguments{{}, {data_ptr}, {}};
-    } else {
-      return Arguments{data_ptr};
-    }
   }
 
  public:
@@ -1057,7 +1032,7 @@ void launch_sm120_fp8_scaled_mm(
   TORCH_CHECK(status == cutlass::Status::kSuccess)
 }
 
-template <typename OutType, bool ScalarA>
+template <typename OutType>
 void sm120_fp8_dispatch_bias(
     torch::Tensor& out,
     const torch::Tensor& a,
@@ -1085,8 +1060,7 @@ void sm120_fp8_dispatch_bias(
       MainloopScheduleType,
       EpilogueScheduleType,
       TileSchedulerType,
-      true,
+      true>;
-      ScalarA>;
 
   using GemmDefault = DeviceGemmFp8RowwiseSm120<
       ElementInput,
@@ -1097,8 +1071,7 @@ void sm120_fp8_dispatch_bias(
       MainloopScheduleType,
       EpilogueScheduleType,
       TileSchedulerType,
-      false,
+      false>;
-      ScalarA>;
 
   if (bias) {
     return launch_sm120_fp8_scaled_mm<BiasGemmDefault, true>(out, a, b, scales_a, scales_b, bias);
@@ -1115,10 +1088,7 @@ void sm120_fp8_dispatch_shape(
     const torch::Tensor& scales_a,
     const torch::Tensor& scales_b,
     const c10::optional<torch::Tensor>& bias) {
-  if (scales_a.numel() == 1) {
+  return sm120_fp8_dispatch_bias<OutType>(out, a, b, scales_a, scales_b, bias);
-    return sm120_fp8_dispatch_bias<OutType, true>(out, a, b, scales_a, scales_b, bias);
-  }
-  return sm120_fp8_dispatch_bias<OutType, false>(out, a, b, scales_a, scales_b, bias);
 }
 #endif
 
@@ -1145,26 +1115,7 @@ torch::Tensor fp8_scaled_mm(
   TORCH_CHECK(mat_b.scalar_type() == torch::kFloat8_e4m3fn, "mat_b must be Float8_e4m3fn");
   TORCH_CHECK(out_dtype == torch::kHalf || out_dtype == torch::kBFloat16, "out_dtype must be Half or BFloat16");
 
-  auto sm_version = getSMVersion();
+  TORCH_CHECK(scales_a.numel() == mat_a.size(0), "size of scales_a is not matched");
-  TORCH_CHECK(
-      scales_a.numel() == 1 || scales_a.numel() == mat_a.size(0),
-      "scales_a must contain either one scalar scale or one scale per row; got ",
-      scales_a.numel(),
-      " elements for M=",
-      mat_a.size(0));
-  bool scalar_a_scale_supported = false;
-#if defined CUDA_VERSION && CUDA_VERSION >= 12000
-  scalar_a_scale_supported = sm_version == 90;
-#endif
-#if defined CUDA_VERSION && CUDA_VERSION >= 12080
-  scalar_a_scale_supported = scalar_a_scale_supported || sm_version >= 100;
-#endif
-  TORCH_CHECK(
-      scales_a.numel() != 1 || mat_a.size(0) == 1 || scalar_a_scale_supported,
-      "scalar scales_a with M > 1 is unsupported on SM",
-      sm_version,
-      " for this build; got M=",
-      mat_a.size(0));
   TORCH_CHECK(scales_b.numel() == mat_b.size(1), "size of scales_b is not matched");
   TORCH_CHECK(scales_a.is_contiguous(), "scales_a must be contiguous");
   TORCH_CHECK(scales_b.is_contiguous(), "scales_b msut be contiguous");
@@ -1180,6 +1131,8 @@ torch::Tensor fp8_scaled_mm(
   torch::Tensor out = torch::empty({mat_a.size(0), mat_b.size(1)}, mat_a.options().dtype(out_dtype));
   TORCH_CHECK((out.size(1) * out.element_size()) % 16 == 0, "out must be multiple of 16 bytes for memory alignment");
 
+  auto sm_version = getSMVersion();
+
 #if defined CUDA_VERSION && CUDA_VERSION >= 12080
   if (sm_version >= 120) {
     if (out_dtype == torch::kBFloat16) {

python/sglang/kernels/aot/tests/test_fp8_gemm.py

@@ -5,22 +5,6 @@ import torch
 from sgl_kernel import fp8_scaled_mm
 
 
-def _cuda_version_at_least(major, minor):
-    if torch.version.cuda is None:
-        return False
-    version = tuple(int(component) for component in torch.version.cuda.split(".")[:2])
-    return version >= (major, minor)
-
-
-def _native_scalar_a_supported():
-    if not torch.cuda.is_available():
-        return False
-    capability = torch.cuda.get_device_capability()
-    if capability == (9, 0):
-        return _cuda_version_at_least(12, 0)
-    return capability[0] in (10, 12) and _cuda_version_at_least(12, 8)
-
-
 def torch_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias):
     o = torch.matmul(a.to(torch.float32), b.to(torch.float32))
     o = o.to(torch.float32)
@@ -54,40 +38,6 @@ def _test_accuracy_once(M, N, K, with_bias, out_dtype, device):
     print(f"M={M}, N={N}, K={K}, with_bias={with_bias}, out_dtype={out_dtype}: OK")
 
 
-def _test_scalar_a_accuracy_once(M, N, K, with_bias, out_dtype, device):
-    fp8_info = torch.finfo(torch.float8_e4m3fn)
-    a_fp8 = (
-        torch.randn(M, K, dtype=torch.float32, device=device)
-        .clamp(min=fp8_info.min, max=fp8_info.max)
-        .to(torch.float8_e4m3fn)
-    )
-    b_fp8 = (
-        torch.randn(N, K, dtype=torch.float32, device=device)
-        .clamp(min=fp8_info.min, max=fp8_info.max)
-        .to(torch.float8_e4m3fn)
-        .t()
-    )
-    scale_a = torch.tensor([0.03125], device=device, dtype=torch.float32)
-    scale_a_repeated = scale_a.repeat(M)
-
-    # Resemble merged projections whose component matrices were quantized with
-    # different tensorwise scales before concatenation.
-    scale_b = torch.empty(N, device=device, dtype=torch.float32)
-    first_boundary = N // 3
-    second_boundary = 2 * N // 3
-    scale_b[:first_boundary] = 0.015625
-    scale_b[first_boundary:second_boundary] = 0.03125
-    scale_b[second_boundary:] = 0.0625
-
-    bias = torch.randn(N, device=device, dtype=out_dtype) if with_bias else None
-    expected = torch_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype, bias)
-    actual = fp8_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype, bias)
-    repeated = fp8_scaled_mm(a_fp8, b_fp8, scale_a_repeated, scale_b, out_dtype, bias)
-
-    torch.testing.assert_close(expected, actual, rtol=0.02, atol=1)
-    torch.testing.assert_close(repeated, actual, rtol=0, atol=0)
-
-
 @pytest.mark.parametrize("M", [1, 128, 512, 1024, 4096])
 @pytest.mark.parametrize("N", [16, 128, 512, 1024, 4096])
 @pytest.mark.parametrize("K", [512, 1024, 4096, 8192, 16384])
@@ -142,43 +92,6 @@ def test_accuracy_sm90_swap_ab(shape_mn, K, with_bias, out_dtype):
     _test_accuracy_once(M, N, K, with_bias, out_dtype, "cuda")
 
 
-@pytest.mark.skipif(
-    not _native_scalar_a_supported(),
-    reason="native scalar A scales require a compatible SM90, SM100, or SM120 build",
-)
-@pytest.mark.parametrize("M", [1, 2, 8, 16, 64, 189])
-@pytest.mark.parametrize("with_bias", [True, False])
-@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
-def test_scalar_a_channelwise_b(M, with_bias, out_dtype):
-    _test_scalar_a_accuracy_once(M, 6144, 4096, with_bias, out_dtype, "cuda")
-
-
-def test_rejects_invalid_a_scale_count():
-    M, N, K = 8, 128, 512
-    a = torch.randn(M, K, device="cuda").to(torch.float8_e4m3fn)
-    b = torch.randn(N, K, device="cuda").to(torch.float8_e4m3fn).t()
-    scale_a = torch.ones(2, device="cuda", dtype=torch.float32)
-    scale_b = torch.ones(N, device="cuda", dtype=torch.float32)
-
-    with pytest.raises(RuntimeError, match="scales_a must contain either"):
-        fp8_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16, None)
-
-
-@pytest.mark.skipif(
-    not torch.cuda.is_available() or torch.cuda.get_device_capability() != (8, 9),
-    reason="SM89-specific scalar A validation",
-)
-def test_rejects_scalar_a_with_multiple_rows_on_sm89():
-    M, N, K = 8, 128, 512
-    a = torch.randn(M, K, device="cuda").to(torch.float8_e4m3fn)
-    b = torch.randn(N, K, device="cuda").to(torch.float8_e4m3fn).t()
-    scale_a = torch.ones(1, device="cuda", dtype=torch.float32)
-    scale_b = torch.ones(N, device="cuda", dtype=torch.float32)
-
-    with pytest.raises(RuntimeError, match="scalar scales_a with M > 1 is unsupported"):
-        fp8_scaled_mm(a, b, scale_a, scale_b, torch.bfloat16, None)
-
-
 PRODUCTION_LIKE_FP8_GEMM_CASES = [
     (189, 4608, 8192, False, torch.bfloat16),
     (3330, 256, 8192, False, torch.bfloat16),

python/sglang/srt/layers/layernorm.py

@@ -55,7 +55,6 @@ _is_cpu_amx_available = cpu_has_amx_support()
 _is_cpu = is_cpu()
 _is_xpu = is_xpu()
 _flashinfer_layernorm_available = False
-_flashinfer_rmsnorm_quant_available = False
 
 if _is_cuda or _is_xpu or _is_musa:
     if _is_flashinfer_available:
@@ -84,19 +83,8 @@ if _is_cuda or _is_xpu or _is_musa:
             _flashinfer_layernorm_available = True
         except (ImportError, AttributeError):
             _flashinfer_layernorm_available = False
-
-        try:
-            from flashinfer.norm import (
-                fused_add_rmsnorm_quant as _flashinfer_fused_add_rmsnorm_quant,
-            )
-            from flashinfer.norm import rmsnorm_quant as _flashinfer_rmsnorm_quant
-
-            _flashinfer_rmsnorm_quant_available = True
-        except (ImportError, AttributeError):
-            _flashinfer_rmsnorm_quant_available = False
     else:
         _flashinfer_layernorm_available = False
-        _flashinfer_rmsnorm_quant_available = False
 
     from sgl_kernel import (
         fused_add_rmsnorm,
@@ -169,7 +157,6 @@ if _is_cuda:
 
 logger = logging.getLogger(__name__)
 
-
 if _is_npu:
     import torch_npu
     from sgl_kernel_npu.norm.add_rmsnorm_bias import add_gemma_rms_norm
@@ -367,57 +354,6 @@ def _forward_with_allreduce_fusion_quant_per_group(
     return (bf16_out, fp8_out, scale_out), residual_out
 
 
-def _fp8_static_input_scale(linear) -> Optional[torch.Tensor]:
-    """Return the per-tensor static FP8 activation scale of ``linear`` if it is
-    an FP8 linear using static per-tensor activation scaling that can consume a
-    pre-quantized ``(fp8, scale)`` input; otherwise ``None``.
-
-    Recognizes both the native ``Fp8LinearMethod`` (non block/mxfp8/marlin) and
-    the compressed-tensors W8A8-FP8 scheme with a static per-tensor input scale
-    (e.g. RedHatAI ``*-FP8`` checkpoints). The flashinfer fused kernel only
-    supports per-tensor quant, hence the ``numel() == 1`` requirement.
-    """
-    if linear is None:
-        return None
-    quant_method = getattr(linear, "quant_method", None)
-    if quant_method is None:
-        return None
-    if not _is_static_per_tensor_fp8_linear(quant_method, linear):
-        return None
-    input_scale = getattr(linear, "input_scale", None)
-    if input_scale is None or input_scale.numel() != 1:
-        return None
-    return input_scale
-
-
-def _is_static_per_tensor_fp8_linear(quant_method, linear) -> bool:
-    try:
-        from sglang.srt.layers.quantization.fp8 import Fp8LinearMethod
-    except ImportError:
-        Fp8LinearMethod = ()
-    if isinstance(quant_method, Fp8LinearMethod):
-        return not (
-            getattr(quant_method, "block_quant", False)
-            or getattr(quant_method, "use_mxfp8", False)
-            or getattr(quant_method, "use_marlin", False)
-        )
-    try:
-        from sglang.srt.layers.quantization.compressed_tensors.compressed_tensors import (
-            CompressedTensorsLinearMethod,
-        )
-        from sglang.srt.layers.quantization.compressed_tensors.schemes import (
-            CompressedTensorsW8A8Fp8,
-        )
-    except ImportError:
-        return False
-    if isinstance(quant_method, CompressedTensorsLinearMethod):
-        scheme = getattr(linear, "scheme", None)
-        return isinstance(scheme, CompressedTensorsW8A8Fp8) and getattr(
-            scheme, "is_static_input_scheme", False
-        )
-    return False
-
-
 class RMSNorm(MultiPlatformOp):
     def __init__(
         self,
@@ -471,7 +407,6 @@ class RMSNorm(MultiPlatformOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
-        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         if x.numel() == 0:
             if residual is not None:
@@ -501,20 +436,6 @@ class RMSNorm(MultiPlatformOp):
             if needs_reshape:
                 out = out.reshape(original_shape)
             return out
-        # Fuse the downstream FP8 static per-tensor activation quant into the
-        # norm when supported. Placed after the empty / variance-override /
-        # batch-invariant guards above (all incompatible with the fused kernel)
-        # and gated on not-HF-cast, so it only runs on the standard RMSNorm path.
-        if (
-            quant_linear is not None
-            and not self.cast_x_before_out_mul
-            and _flashinfer_rmsnorm_quant_available
-        ):
-            scale = _fp8_static_input_scale(quant_linear)
-            if scale is not None:
-                return self.forward_with_per_tensor_quant_fusion(
-                    x, scale, residual, post_residual_addition
-                )
         if self.cast_x_before_out_mul and residual is None:
             # Use HF-semantics kernel (cast to dtype before weight multiply).
             if (
@@ -572,7 +493,6 @@ class RMSNorm(MultiPlatformOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
-        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         if residual is not None:
             if post_residual_addition is not None:
@@ -588,7 +508,6 @@ class RMSNorm(MultiPlatformOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
-        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         # Fix dsv4 dp attenton issue
         # the symptom is torch.AcceleratorError: HIP error: invalid configuration argument
@@ -665,7 +584,6 @@ class RMSNorm(MultiPlatformOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
-        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         # Fallback to native implementation if vllm is not available
         if not _has_vllm_rms_norm:
@@ -705,7 +623,6 @@ class RMSNorm(MultiPlatformOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
-        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         if check_cuda_graph_backend(Phase.PREFILL, Backend.TC_PIECEWISE):
             return self.forward_native(x, residual, post_residual_addition)
@@ -729,7 +646,6 @@ class RMSNorm(MultiPlatformOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
-        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         if not x.is_contiguous():
             x = x.contiguous()
@@ -780,7 +696,6 @@ class RMSNorm(MultiPlatformOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
-        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         if _is_cpu_amx_available:
             if residual is not None:
@@ -801,7 +716,6 @@ class RMSNorm(MultiPlatformOp):
         x: torch.Tensor,
         residual: Optional[torch.Tensor] = None,
         post_residual_addition: Optional[torch.Tensor] = None,
-        quant_linear: Optional[nn.Module] = None,
     ) -> Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]:
         if self.variance_size_override is not None:
             return self.forward_native(x, residual, post_residual_addition)
@@ -855,65 +769,6 @@ class RMSNorm(MultiPlatformOp):
             self, x, residual, self.weight, group_size, use_attn_tp_group, keep_bf16
         )
 
-    def forward_with_per_tensor_quant_fusion(
-        self,
-        x: torch.Tensor,
-        scale: torch.Tensor,
-        residual: Optional[torch.Tensor] = None,
-        post_residual_addition: Optional[torch.Tensor] = None,
-        fp8_dtype: torch.dtype = torch.float8_e4m3fn,
-    ) -> Union[
-        Tuple[torch.Tensor, torch.Tensor, torch.dtype],
-        Tuple[Tuple[torch.Tensor, torch.Tensor, torch.dtype], torch.Tensor],
-    ]:
-        """Fused RMSNorm + static per-tensor FP8 quantization.
-
-        The normed activation is quantized to ``fp8_dtype`` using the per-tensor
-        reciprocal ``scale`` (same convention as ``static_quant_fp8``:
-        ``q = normed / scale``), so a downstream FP8 linear carrying a matching
-        static ``input_scale`` can skip its own activation quant.
-
-        The quantized activation is emitted as a ``(fp8_out, scale, orig_dtype)``
-        tuple; ``orig_dtype`` (the un-quantized activation dtype) is carried so
-        the downstream FP8 GEMM produces its output in the model's dtype rather
-        than defaulting to bf16.
-
-        Return contract mirrors ``forward``:
-        * no residual -> ``(fp8_out, scale, orig_dtype)``
-        * w/ residual -> ``((fp8_out, scale, orig_dtype), residual_out)``
-        """
-        orig_dtype = x.dtype
-        needs_reshape = x.dim() != 2
-        if needs_reshape:
-            original_shape = x.shape
-            x = x.contiguous().reshape(-1, original_shape[-1])
-        elif not x.is_contiguous():
-            x = x.contiguous()
-
-        out = torch.empty_like(x, dtype=fp8_dtype)
-        if residual is not None:
-            if post_residual_addition is not None:
-                residual = residual + post_residual_addition
-            if residual.dim() != 2:
-                residual = residual.contiguous().reshape(-1, residual.shape[-1])
-            elif not residual.is_contiguous():
-                residual = residual.contiguous()
-            # In-place: residual += x, then out = quant(rmsnorm(residual) * w).
-            _flashinfer_fused_add_rmsnorm_quant(
-                out, x, residual, self.weight.data, scale, self.variance_epsilon
-            )
-            if needs_reshape:
-                out = out.reshape(original_shape)
-                residual = residual.reshape(original_shape)
-            return (out, scale, orig_dtype), residual
-
-        _flashinfer_rmsnorm_quant(
-            out, x, self.weight.data, scale, self.variance_epsilon
-        )
-        if needs_reshape:
-            out = out.reshape(original_shape)
-        return out, scale, orig_dtype
-
 
 class LayerNorm(MultiPlatformOp):
     def __init__(

python/sglang/srt/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py

@@ -231,23 +231,6 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsLinearScheme):
         x: torch.Tensor,
         bias: Optional[torch.Tensor] = None,
     ) -> torch.Tensor:
-        if isinstance(x, tuple):
-            # Pre-quantized activation from a fused RMSNorm+FP8 quant kernel:
-            # x = (fp8_input, per_tensor_input_scale[, orig_dtype]).
-            # apply_fp8_linear detects the fp8 dtype and skips re-quantizing;
-            # orig_dtype (when present) sets the GEMM output dtype.
-            qx, x_scale = x[0], x[1]
-            out_dtype = x[2] if len(x) > 2 else None
-            return apply_fp8_linear(
-                input=qx,
-                weight=layer.weight,
-                weight_scale=layer.weight_scale,
-                input_scale=x_scale,
-                bias=bias,
-                use_per_token_if_dynamic=True,
-                compressed_tensor_quant=True,
-                pre_quant_output_dtype=out_dtype,
-            )
         if self.weight_block_size is not None:
             return self.w8a8_block_fp8_linear(
                 input=x,

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

@@ -139,7 +139,10 @@ def _require_fp4_dtype():
 
 
 if _use_aiter or _use_hip_int4:
-    from aiter.ops.shuffle import shuffle_scale, shuffle_weight
+    from aiter.ops.shuffle import (
+        shuffle_scale,
+        shuffle_weight,
+    )
 
 if _use_aiter:
     from sglang.srt.layers.quantization.fp8_utils import (
@@ -1025,24 +1028,6 @@ class Fp8LinearMethod(LinearMethodBase):
                 bias=bias,
             )
 
-        if isinstance(x, tuple):
-            # Pre-quantized activation from a fused RMSNorm+FP8 quant kernel:
-            # x = (fp8_input, per_tensor_input_scale[, orig_dtype]).
-            # apply_fp8_linear detects the fp8 dtype and skips re-quantizing;
-            # orig_dtype (when present) sets the GEMM output dtype.
-            qx, x_scale = x[0], x[1]
-            out_dtype = x[2] if len(x) > 2 else None
-            return apply_fp8_linear(
-                input=qx,
-                weight=layer.weight,
-                weight_scale=layer.weight_scale,
-                input_scale=x_scale,
-                bias=bias,
-                cutlass_fp8_supported=self.cutlass_fp8_supported,
-                use_per_token_if_dynamic=self.use_per_token_if_dynamic,
-                pre_quant_output_dtype=out_dtype,
-            )
-
         return apply_fp8_linear(
             input=x,
             weight=layer.weight,
@@ -1839,7 +1824,9 @@ class Fp8MoEMethod(FusedMoEMethodBase):
             )
             return qweight.view_as(weight), scale_u8
 
-        from sglang.srt.layers.quantization.mxfp8_block_convert import _ue8m0_to_fp32
+        from sglang.srt.layers.quantization.mxfp8_block_convert import (
+            _ue8m0_to_fp32,
+        )
 
         def _quantize_for_deepgemm(weight: torch.Tensor):
             weight = weight.contiguous()

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

@@ -59,7 +59,6 @@ logger = logging.getLogger(__name__)
 _is_hip = is_hip()
 _is_cuda = is_cuda()
 _is_fp8_fnuz = is_fp8_fnuz()
-_is_sm90_supported = is_sm90_supported()
 _is_sm100_supported = is_sm100_supported()
 _is_sm120_supported = is_sm120_supported()
 _is_gfx95_supported = is_gfx95_supported()
@@ -313,9 +312,7 @@ FP8_GEMM_RUNNER_BACKEND: Fp8GemmRunnerBackend | None = None
 
 
 if is_blackwell_supported() and is_flashinfer_available():
-    from flashinfer import (
+    from flashinfer import SfLayout
-        SfLayout,
-    )
     from flashinfer import bmm_fp8 as _raw_flashinfer_bmm_fp8
     from flashinfer import mm_mxfp8 as _raw_flashinfer_mm_mxfp8
     from flashinfer import mxfp8_quantize as _raw_flashinfer_mxfp8_quantize
@@ -1553,7 +1550,9 @@ def requant_block_scale_ue8m0_for_deepgemm(
     scales are not already UE8M0, and DeepGEMM can run the layer (bf16 output,
     aligned shape). Returns True when it requantizes.
     """
-    from sglang.srt.model_loader.utils import should_deepgemm_weight_requant_ue8m0
+    from sglang.srt.model_loader.utils import (
+        should_deepgemm_weight_requant_ue8m0,
+    )
 
     if (
         not use_deepgemm_runner
@@ -1842,7 +1841,6 @@ def apply_fp8_linear(
     use_per_token_if_dynamic: bool = False,
     pad_output: Optional[bool] = None,
     compressed_tensor_quant: bool = False,
-    pre_quant_output_dtype: Optional[torch.dtype] = None,
 ) -> torch.Tensor:
     # Note: we pad the input because torch._scaled_mm is more performant
     # for matrices with batch dimension > 16.
@@ -1859,42 +1857,10 @@ def apply_fp8_linear(
     input_2d = input.view(-1, input.shape[-1])
     output_shape = [*input.shape[:-1], weight.shape[1]]
 
-    # A pre-quantized fp8 activation (e.g. from a fused RMSNorm+quant kernel)
+    if compressed_tensor_quant:
-    # carries no original dtype: skip re-quant, reuse the supplied per-tensor
-    # input_scale, and emit ``pre_quant_output_dtype`` (the model's activation
-    # dtype, propagated by the producer) or bf16 if it was not provided.
-    input_prequantized = input_2d.dtype in (
-        torch.float8_e4m3fn,
-        torch.float8_e4m3fnuz,
-    )
-    if input_prequantized:
-        output_dtype = pre_quant_output_dtype or torch.bfloat16
-    else:
-        output_dtype = input.dtype
-
-    channelwise_cutlass = (
-        cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]
-    )
-    cutlass_compatible_b = weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0
-    use_cutlass_channelwise_gemm = (
-        channelwise_cutlass and cutlass_compatible_b and not use_triton_w8a8_fp8_kernel
-    )
-    native_scalar_a_scale = use_cutlass_channelwise_gemm and (
-        _is_sm90_supported or _is_sm100_supported or _is_sm120_supported
-    )
-
-    if input_prequantized:
-        assert input_scale is not None and input_scale.numel() == 1
-        qinput = input_2d
-        if channelwise_cutlass and not native_scalar_a_scale:
-            # Unsupported CUTLASS epilogues require one A scale per row.
-            x_scale = input_scale.repeat(input_2d.shape[0]).view(-1, 1)
-        else:
-            x_scale = input_scale
-    elif compressed_tensor_quant:
         # Maybe apply padding to output, see comment in __init__
         num_token_padding = output_padding
-        if channelwise_cutlass:
+        if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]:
             num_token_padding = None
         # For static per-tensor activation scales when using inductor compiler,
         # use pure PyTorch ops instead of the opaque sgl_kernel quant kernel.
@@ -1923,19 +1889,13 @@ def apply_fp8_linear(
                 num_token_padding=num_token_padding,
                 use_per_token_if_dynamic=use_per_token_if_dynamic,
             )
-        if (
-            input_scale is not None
-            and channelwise_cutlass
-            and not native_scalar_a_scale
-        ):
-            x_scale = input_scale.repeat(input_2d.shape[0]).view(-1, 1)
     else:
+        # cutlass w8a8 fp8 sgl-kernel only supports per-token scale
         if input_scale is not None:
             assert input_scale.numel() == 1
+            # broadcast per-tensor scale to per-token scale when supporting cutlass
             qinput, x_scale = static_quant_fp8(
-                input_2d,
+                input_2d, input_scale, repeat_scale=cutlass_fp8_supported
-                input_scale,
-                repeat_scale=channelwise_cutlass and not native_scalar_a_scale,
             )
         else:
             # default use per-token quantization if dynamic
@@ -1956,12 +1916,13 @@ def apply_fp8_linear(
                         input_2d, group_size=input_2d.shape[1]
                     )
 
-    if channelwise_cutlass:
+    if cutlass_fp8_supported and weight_scale.numel() == weight.shape[1]:
-        if not use_cutlass_channelwise_gemm:
+        cutlass_compatible_b = weight.shape[0] % 16 == 0 and weight.shape[1] % 16 == 0
+        if not cutlass_compatible_b or use_triton_w8a8_fp8_kernel:
             # Massage the input to be 2D
             qinput = qinput.view(-1, qinput.shape[-1])
             output = triton_scaled_mm(
-                qinput, weight, x_scale, weight_scale, output_dtype, bias
+                qinput, weight, x_scale, weight_scale, input.dtype, bias
             )
         else:
             output = fp8_scaled_mm(
@@ -1969,7 +1930,7 @@ def apply_fp8_linear(
                 weight,
                 x_scale,
                 weight_scale,
-                out_dtype=output_dtype,
+                out_dtype=input.dtype,
                 bias=bias,
             )
         return output.view(*output_shape)
@@ -2002,7 +1963,7 @@ def apply_fp8_linear(
                 WQ=weight.T,
                 x_scale=x_scale,
                 w_scale=weight_scale,
-                dtype=output_dtype,
+                dtype=input.dtype,
             )
             if bias is not None:
                 output += bias
@@ -2018,7 +1979,7 @@ def apply_fp8_linear(
             output = torch._scaled_mm(
                 qinput,
                 weight,
-                out_dtype=output_dtype,
+                out_dtype=input.dtype,
                 scale_a=x_scale,
                 scale_b=weight_scale.t(),
                 bias=bias,
@@ -2032,7 +1993,7 @@ def apply_fp8_linear(
         output = torch._scaled_mm(
             qinput,
             weight,
-            out_dtype=output_dtype,
+            out_dtype=input.dtype,
             scale_a=x_scale,
             scale_b=weight_scale,
             bias=bias,
@@ -2061,7 +2022,7 @@ def apply_fp8_linear(
         input_2d.shape,
         output_shape,
         bias,
-        output_dtype,
+        input.dtype,
     )
 
 

python/sglang/srt/models/llama.py

@@ -348,13 +348,9 @@ class LlamaDecoderLayer(nn.Module):
         # Self Attention
         if residual is None:
             residual = hidden_states
-            hidden_states = self.input_layernorm(
+            hidden_states = self.input_layernorm(hidden_states)
-                hidden_states, quant_linear=self.self_attn.qkv_proj
-            )
         else:
-            hidden_states, residual = self.input_layernorm(
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
-                hidden_states, residual, quant_linear=self.self_attn.qkv_proj
-            )
         hidden_states = self.self_attn(
             positions=positions,
             hidden_states=hidden_states,
@@ -362,9 +358,7 @@ class LlamaDecoderLayer(nn.Module):
         )
 
         # Fully Connected
-        hidden_states, residual = self.post_attention_layernorm(
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
-            hidden_states, residual, quant_linear=self.mlp.gate_up_proj
-        )
         hidden_states = self.mlp(hidden_states)
         return hidden_states, residual
 

python/sglang/srt/models/llama_eagle.py

@@ -50,7 +50,7 @@ class LlamaDecoderLayer(LlamaDecoderLayer):
         # https://github.com/SafeAILab/EAGLE/blob/35c78f6cdc19a73e05cf5c330b4c358dad970c6a/eagle/model/cnets.py#L427
         if layer_id == 0:
             del self.input_layernorm
-            setattr(self, "input_layernorm", lambda x, quant_linear=None: x)
+            setattr(self, "input_layernorm", lambda x: x)
 
 
 class LlamaModel(nn.Module):

python/sglang/srt/models/qwen2.py

@@ -291,13 +291,9 @@ class Qwen2DecoderLayer(nn.Module):
         # Self Attention
         if residual is None:
             residual = hidden_states
-            hidden_states = self.input_layernorm(
+            hidden_states = self.input_layernorm(hidden_states)
-                hidden_states, quant_linear=self.self_attn.qkv_proj
-            )
         else:
-            hidden_states, residual = self.input_layernorm(
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
-                hidden_states, residual, quant_linear=self.self_attn.qkv_proj
-            )
         hidden_states = self.self_attn(
             positions=positions,
             hidden_states=hidden_states,
@@ -305,9 +301,7 @@ class Qwen2DecoderLayer(nn.Module):
         )
 
         # Fully Connected
-        hidden_states, residual = self.post_attention_layernorm(
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
-            hidden_states, residual, quant_linear=self.mlp.gate_up_proj
-        )
         hidden_states = self.mlp(hidden_states)
         return hidden_states, residual
 

python/sglang/srt/models/qwen2_eagle.py

@@ -51,7 +51,7 @@ class Qwen2DecoderLayer(Qwen2DecoderLayer):
         # https://github.com/SafeAILab/EAGLE/blob/35c78f6cdc19a73e05cf5c330b4c358dad970c6a/eagle/model/cnets.py#L427
         if layer_id == 0:
             del self.input_layernorm
-            setattr(self, "input_layernorm", lambda x, quant_linear=None: x)
+            setattr(self, "input_layernorm", lambda x: x)
 
 
 class Qwen2Model(nn.Module):

test/registered/layers/test_layernorm_fusion.py

@@ -1,142 +0,0 @@
-import itertools
-import unittest
-
-import torch
-
-from sglang.srt.layers.layernorm import RMSNorm
-from sglang.test.ci.ci_register import register_cuda_ci
-from sglang.test.test_utils import CustomTestCase
-
-register_cuda_ci(est_time=15, stage="base-b", runner_config="1-gpu-large")
-
-
-class TestRMSNormFp8QuantFusion(CustomTestCase):
-    DTYPES = [torch.bfloat16, torch.half]
-    NUM_TOKENS = [7, 83, 512]
-    HIDDEN_SIZES = [512, 4096]
-    ADD_RESIDUAL = [False, True]
-    SEED = 0
-    FP8_DTYPE = torch.float8_e4m3fn
-
-    @classmethod
-    def setUpClass(cls):
-        if not torch.cuda.is_available():
-            raise unittest.SkipTest("CUDA is not available")
-        from sglang.srt.layers.layernorm import _flashinfer_rmsnorm_quant_available
-
-        if not _flashinfer_rmsnorm_quant_available:
-            raise unittest.SkipTest("flashinfer rmsnorm_quant is not available")
-        torch.set_default_device("cuda")
-
-    def _run_fusion_test(self, num_tokens, hidden_size, add_residual, dtype):
-        torch.manual_seed(self.SEED)
-
-        layer = RMSNorm(hidden_size).to(dtype=dtype)
-        layer.weight.data.normal_(mean=1.0, std=0.1)
-        x = torch.randn(num_tokens, hidden_size, dtype=dtype)
-        residual = torch.randn_like(x) if add_residual else None
-        # Per-tensor reciprocal scale (as carried by a static FP8 linear).
-        scale = torch.tensor([0.05], dtype=torch.float32)
-
-        with torch.inference_mode():
-            ref = layer.forward_native(
-                x.clone(), residual.clone() if add_residual else None
-            )
-            normed_ref = ref[0] if add_residual else ref
-            residual_ref = ref[1] if add_residual else None
-
-            result = layer.forward_with_per_tensor_quant_fusion(
-                x.clone(), scale, residual.clone() if add_residual else None
-            )
-
-        if add_residual:
-            (q, s, out_dtype), r = result
-        else:
-            q, s, out_dtype = result
-            r = None
-
-        # Output contract.
-        self.assertEqual(q.dtype, self.FP8_DTYPE)
-        self.assertIs(s, scale)
-        self.assertEqual(out_dtype, dtype)
-        self.assertEqual(tuple(q.shape), (num_tokens, hidden_size))
-        if add_residual:
-            self.assertEqual(r.dtype, dtype)
-            self.assertTrue(
-                torch.allclose(r.float(), residual_ref.float(), atol=1e-2, rtol=1e-2)
-            )
-
-        # Numerical: dequantized (q * scale) matches the reference normed output
-        # within FP8 e4m3 precision.
-        deq = q.float() * scale
-        ref_flat = normed_ref.float().flatten()
-        cos = torch.nn.functional.cosine_similarity(deq.flatten(), ref_flat, dim=0)
-        self.assertGreater(cos.item(), 0.99)
-        rel_err = (
-            deq.flatten() - ref_flat
-        ).abs().mean() / ref_flat.abs().mean().clamp_min(1e-6)
-        self.assertLess(rel_err.item(), 0.1)
-
-    def test_rms_norm_fp8_quant_fusion(self):
-        for params in itertools.product(
-            self.NUM_TOKENS,
-            self.HIDDEN_SIZES,
-            self.ADD_RESIDUAL,
-            self.DTYPES,
-        ):
-            with self.subTest(
-                num_tokens=params[0],
-                hidden_size=params[1],
-                add_residual=params[2],
-                dtype=params[3],
-            ):
-                self._run_fusion_test(*params)
-
-    def test_forward_cuda_quant_linear_dispatch(self):
-        """forward_cuda routes to the fused path only when applicable."""
-        import sglang.srt.layers.layernorm as ln_mod
-
-        torch.manual_seed(self.SEED)
-        hidden_size, num_tokens = 512, 32
-        x = torch.randn(num_tokens, hidden_size, dtype=torch.bfloat16)
-        residual = torch.randn_like(x)
-        scale = torch.tensor([0.05], dtype=torch.float32)
-
-        orig_static_scale = ln_mod._fp8_static_input_scale
-        ln_mod._fp8_static_input_scale = lambda linear: scale
-        try:
-            plain = RMSNorm(hidden_size).to(dtype=torch.bfloat16)
-            plain.weight.data.normal_(mean=1.0, std=0.1)
-
-            with torch.inference_mode():
-                # Plain norm -> fused (fp8, scale, dtype) + bf16 residual.
-                (q, s, out_dtype), r = plain(
-                    x.clone(), residual.clone(), quant_linear=object()
-                )
-
-                # variance_size_override is incompatible -> must not fuse.
-                var_layer = RMSNorm(hidden_size, var_hidden_size=hidden_size // 2).to(
-                    dtype=torch.bfloat16
-                )
-                var_out = var_layer(x.clone(), residual.clone(), quant_linear=object())
-                # cast_x_before_out_mul (HF semantics) is incompatible -> must not fuse.
-                cast_layer = RMSNorm(hidden_size, cast_x_before_out_mul=True).to(
-                    dtype=torch.bfloat16
-                )
-                cast_out = cast_layer(
-                    x.clone(), residual.clone(), quant_linear=object()
-                )
-        finally:
-            ln_mod._fp8_static_input_scale = orig_static_scale
-
-        self.assertEqual(q.dtype, self.FP8_DTYPE)
-        self.assertIs(s, scale)
-        self.assertEqual(out_dtype, torch.bfloat16)
-        self.assertEqual(r.dtype, torch.bfloat16)
-
-        self.assertEqual(var_out[0].dtype, torch.bfloat16)
-        self.assertEqual(cast_out[0].dtype, torch.bfloat16)
-
-
-if __name__ == "__main__":
-    unittest.main()

test/registered/quant/test_fp8_utils.py

@@ -1,6 +1,4 @@
 import unittest
-from types import SimpleNamespace
-from unittest.mock import patch
 
 import torch
 
@@ -12,7 +10,7 @@ from sglang.srt.layers.quantization.fp8_utils import (
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
 
-register_cuda_ci(est_time=12, stage="base-b", runner_config="1-gpu-large")
+register_cuda_ci(est_time=9, stage="base-b", runner_config="1-gpu-large")
 
 
 class TestInverseTransformScaleUe8m0(CustomTestCase):
@@ -45,269 +43,5 @@ class TestInverseTransformScaleUe8m0(CustomTestCase):
             ), f"{sf_fp32_original=} {sf_fp32_recreated}"
 
 
-class TestApplyFp8LinearScaleDispatch(CustomTestCase):
-    @classmethod
-    def setUpClass(cls):
-        if not torch.cuda.is_available():
-            raise unittest.SkipTest("CUDA is not available")
-        torch.set_default_device("cuda")
-
-    @staticmethod
-    def _make_inputs(dtype=torch.bfloat16):
-        M, K, N = 8, 16, 32
-        input = torch.randn(M, K, dtype=dtype)
-        qinput = input.to(torch.float8_e4m3fn)
-        weight = torch.randn(N, K).to(torch.float8_e4m3fn).t()
-        input_scale = torch.tensor([0.05], dtype=torch.float32)
-        weight_scale = torch.linspace(0.01, 0.03, N, dtype=torch.float32)
-        return input, qinput, weight, input_scale, weight_scale
-
-    def test_native_scalar_a_static_prequant_and_dynamic_scale_shapes(self):
-        import sglang.srt.layers.quantization.fp8_utils as fp8_utils
-
-        exec_config = SimpleNamespace(
-            graph=SimpleNamespace(
-                cuda_graph_config=SimpleNamespace(
-                    prefill=SimpleNamespace(tc_compiler="none")
-                )
-            )
-        )
-        for capability in (
-            "_is_sm90_supported",
-            "_is_sm100_supported",
-            "_is_sm120_supported",
-        ):
-            with self.subTest(capability=capability):
-                input, qinput, weight, input_scale, weight_scale = self._make_inputs()
-                seen_scales = []
-
-                def fake_fp8_scaled_mm(
-                    mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None
-                ):
-                    seen_scales.append(scales_a)
-                    return torch.empty(
-                        (mat_a.shape[0], mat_b.shape[1]),
-                        dtype=out_dtype,
-                        device=mat_a.device,
-                    )
-
-                capabilities = {
-                    "_is_sm90_supported": False,
-                    "_is_sm100_supported": False,
-                    "_is_sm120_supported": False,
-                }
-                capabilities[capability] = True
-                with patch.multiple(fp8_utils, **capabilities), patch.object(
-                    fp8_utils, "fp8_scaled_mm", side_effect=fake_fp8_scaled_mm
-                ), patch.object(fp8_utils, "get_exec", return_value=exec_config):
-                    fp8_utils.apply_fp8_linear(
-                        input,
-                        weight,
-                        weight_scale,
-                        input_scale=input_scale,
-                        cutlass_fp8_supported=True,
-                    )
-                    fp8_utils.apply_fp8_linear(
-                        input,
-                        weight,
-                        weight_scale,
-                        input_scale=input_scale,
-                        cutlass_fp8_supported=True,
-                        use_per_token_if_dynamic=True,
-                        compressed_tensor_quant=True,
-                    )
-                    fp8_utils.apply_fp8_linear(
-                        qinput,
-                        weight,
-                        weight_scale,
-                        input_scale=input_scale,
-                        cutlass_fp8_supported=True,
-                        pre_quant_output_dtype=input.dtype,
-                    )
-                    fp8_utils.apply_fp8_linear(
-                        input,
-                        weight,
-                        weight_scale,
-                        input_scale=None,
-                        cutlass_fp8_supported=True,
-                        use_per_token_if_dynamic=True,
-                        compressed_tensor_quant=True,
-                    )
-
-                self.assertEqual(seen_scales[0].numel(), 1)
-                self.assertEqual(seen_scales[1].numel(), 1)
-                self.assertIs(seen_scales[2], input_scale)
-                self.assertEqual(tuple(seen_scales[3].shape), (input.shape[0], 1))
-
-    def test_without_native_scalar_a_static_scale_is_repeated(self):
-        import sglang.srt.layers.quantization.fp8_utils as fp8_utils
-
-        input, qinput, weight, input_scale, weight_scale = self._make_inputs()
-        seen_scales = []
-
-        def fake_fp8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None):
-            seen_scales.append(scales_a)
-            return torch.empty(
-                (mat_a.shape[0], mat_b.shape[1]), dtype=out_dtype, device=mat_a.device
-            )
-
-        with patch.multiple(
-            fp8_utils,
-            _is_sm90_supported=False,
-            _is_sm100_supported=False,
-            _is_sm120_supported=False,
-        ), patch.object(fp8_utils, "fp8_scaled_mm", side_effect=fake_fp8_scaled_mm):
-            fp8_utils.apply_fp8_linear(
-                input,
-                weight,
-                weight_scale,
-                input_scale=input_scale,
-                cutlass_fp8_supported=True,
-            )
-            fp8_utils.apply_fp8_linear(
-                qinput,
-                weight,
-                weight_scale,
-                input_scale=input_scale,
-                cutlass_fp8_supported=True,
-                pre_quant_output_dtype=input.dtype,
-            )
-
-        self.assertEqual(tuple(seen_scales[0].shape), (input.shape[0], 1))
-        self.assertEqual(tuple(seen_scales[1].shape), (input.shape[0], 1))
-
-    def test_linear_methods_forward_fused_scalar_tuple(self):
-        import sglang.srt.layers.quantization.compressed_tensors.schemes.compressed_tensors_w8a8_fp8 as compressed_fp8
-        import sglang.srt.layers.quantization.fp8 as native_fp8
-
-        input, qinput, weight, input_scale, weight_scale = self._make_inputs(
-            torch.float16
-        )
-
-        class Layer:
-            pass
-
-        layer = Layer()
-        layer.weight = weight
-        layer.weight_scale = weight_scale
-        layer.input_scale = input_scale
-
-        native_method = native_fp8.Fp8LinearMethod.__new__(native_fp8.Fp8LinearMethod)
-        native_method.use_marlin = False
-        native_method.use_mxfp8 = False
-        native_method.block_quant = False
-        native_method.cutlass_fp8_supported = True
-        native_method.use_per_token_if_dynamic = False
-
-        compressed_method = compressed_fp8.CompressedTensorsW8A8Fp8.__new__(
-            compressed_fp8.CompressedTensorsW8A8Fp8
-        )
-        compressed_method.weight_block_size = None
-
-        fused_input = (qinput, input_scale, input.dtype)
-        with patch.object(native_fp8, "apply_fp8_linear") as native_apply:
-            native_apply.return_value = torch.empty(
-                (qinput.shape[0], weight.shape[1]), dtype=input.dtype
-            )
-            native_method.apply(layer, fused_input)
-            self.assertIs(native_apply.call_args.kwargs["input_scale"], input_scale)
-            self.assertEqual(
-                native_apply.call_args.kwargs["pre_quant_output_dtype"], input.dtype
-            )
-
-        with patch.object(compressed_fp8, "apply_fp8_linear") as compressed_apply:
-            compressed_apply.return_value = torch.empty(
-                (qinput.shape[0], weight.shape[1]), dtype=input.dtype
-            )
-            compressed_method.apply_weights(layer, fused_input)
-            self.assertIs(compressed_apply.call_args.kwargs["input_scale"], input_scale)
-            self.assertEqual(
-                compressed_apply.call_args.kwargs["pre_quant_output_dtype"],
-                input.dtype,
-            )
-
-
-class TestApplyFp8LinearPrequantOutputDtype(CustomTestCase):
-    """apply_fp8_linear with a pre-quantized fp8 activation must emit the
-    caller-supplied ``pre_quant_output_dtype`` (the model's activation dtype),
-    not the fp8 input dtype. Regression test for FP16 models where hardcoding
-    bf16 caused a query/key dtype mismatch in attention."""
-
-    DTYPES = [torch.float16, torch.bfloat16]
-    FP8_DTYPE = torch.float8_e4m3fn
-
-    @classmethod
-    def setUpClass(cls):
-        if not torch.cuda.is_available():
-            raise unittest.SkipTest("CUDA is not available")
-        torch.set_default_device("cuda")
-
-    def _run(self, dtype):
-        from sglang.srt.layers.quantization.fp8_utils import (
-            apply_fp8_linear,
-            cutlass_fp8_supported,
-        )
-
-        torch.manual_seed(0)
-        M, K, N = 33, 512, 256
-        cf = cutlass_fp8_supported()
-        fp8_info = torch.finfo(self.FP8_DTYPE)
-
-        normed = torch.randn(M, K, dtype=dtype)
-        input_scale = torch.tensor([0.05], dtype=torch.float32)
-        # Per-channel fp8 weight in column-major (K, N) layout.
-        w = torch.randn(N, K, dtype=dtype) * 0.05
-        w_scale = (w.abs().amax(dim=1) / fp8_info.max).float()
-        weight = (
-            (w.float() / w_scale[:, None])
-            .clamp(fp8_info.min, fp8_info.max)
-            .to(self.FP8_DTYPE)
-            .t()
-        )
-
-        # Reference: non-pre-quantized input -> output dtype == input dtype.
-        ref = apply_fp8_linear(
-            input=normed,
-            weight=weight,
-            weight_scale=w_scale,
-            input_scale=input_scale,
-            cutlass_fp8_supported=cf,
-        )
-        self.assertEqual(ref.dtype, dtype)
-
-        qinput = (
-            (normed.float() * input_scale.reciprocal())
-            .clamp(fp8_info.min, fp8_info.max)
-            .to(self.FP8_DTYPE)
-        )
-
-        # Pre-quantized input with the dtype propagated -> output matches dtype.
-        out = apply_fp8_linear(
-            input=qinput,
-            weight=weight,
-            weight_scale=w_scale,
-            input_scale=input_scale,
-            cutlass_fp8_supported=cf,
-            pre_quant_output_dtype=dtype,
-        )
-        self.assertEqual(out.dtype, dtype)
-        self.assertTrue(torch.allclose(out.float(), ref.float(), atol=2e-2, rtol=2e-2))
-
-        # Without the dtype hint, the pre-quantized path falls back to bf16.
-        out_default = apply_fp8_linear(
-            input=qinput,
-            weight=weight,
-            weight_scale=w_scale,
-            input_scale=input_scale,
-            cutlass_fp8_supported=cf,
-        )
-        self.assertEqual(out_default.dtype, torch.bfloat16)
-
-    def test_prequant_output_dtype(self):
-        for dtype in self.DTYPES:
-            with self.subTest(dtype=dtype):
-                self._run(dtype)
-
-
 if __name__ == "__main__":
     unittest.main()