[SM100] Fuse bias addition into fp8_blockwise_scaled_mm epilogue

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

@@ -59,6 +59,14 @@
 _use_aiter_gfx95 = _use_aiter and _is_gfx95_supported
 
 
+def _is_sm100_device(device: torch.device) -> bool:
+    if device.type != "cuda":
+        return False
+    device_id = torch.cuda.current_device() if device.index is None else device.index
+    major, _ = get_device_capability(device_id)
+    return major == 10
+
+
 def use_aiter_triton_gemm_w8a8_tuned_gfx950(n: int, k: int) -> bool:
     return (n, k) in [
         (1024, 8192),
@@ -104,7 +112,9 @@ def _fp8_scaled_mm_abstract(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=No
         return mat_a.new_empty((M, N), dtype=out_dtype)
 
     @register_fake_if_exists("sgl_kernel::fp8_blockwise_scaled_mm")
-    def _fp8_blockwise_scaled_mm_abstract(mat_a, mat_b, scales_a, scales_b, out_dtype):
+    def _fp8_blockwise_scaled_mm_abstract(
+        mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None
+    ):
         # mat_a: [M, K], mat_b: [K, N] or [N, K] depending on callsite layout; output is [M, N].
         M = mat_a.shape[-2]
         N = mat_b.shape[-1]
@@ -627,10 +637,16 @@ def cutlass_w8a8_block_fp8_linear_with_fallback(
     q_input, x_scale = per_token_group_quant_fp8(
         input_2d, block_size[1], column_major_scales=True
     )
+    fuse_bias = bias is not None and _is_sm100_device(input_2d.device)
     output = fp8_blockwise_scaled_mm(
-        q_input, weight.T, x_scale, weight_scale.T, out_dtype=input_2d.dtype
+        q_input,
+        weight.T,
+        x_scale,
+        weight_scale.T,
+        out_dtype=input_2d.dtype,
+        bias=bias if fuse_bias else None,
     )
-    if bias is not None:
+    if bias is not None and not fuse_bias:
         output += bias
     return output.to(dtype=input_2d.dtype).view(*output_shape)
 

sgl-kernel/csrc/common_extension.cc

@@ -124,7 +124,8 @@ TORCH_LIBRARY_FRAGMENT(sgl_kernel, m) {
   m.impl("fp8_scaled_mm", torch::kCUDA, &fp8_scaled_mm);
 
   m.def(
-      "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype) -> "
+      "fp8_blockwise_scaled_mm(Tensor mat_a, Tensor mat_b, Tensor scales_a, Tensor scales_b, ScalarType out_dtype, "
+      "Tensor? bias=None) -> "
       "Tensor");
   m.impl("fp8_blockwise_scaled_mm", torch::kCUDA, &fp8_blockwise_scaled_mm);
 

sgl-kernel/csrc/gemm/fp8_blockwise_gemm_kernel.cu

@@ -171,27 +171,187 @@ void launch_sm100_fp8_blockwise_scaled_mm(
   TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status))
 }
 
+template <
+    typename OutType,
+    typename MmaTileShape,
+    typename PerSmTileShape,
+    typename EpilogueTileShape,
+    typename ScalesPerTile,
+    int TileSizeM_ = 128,
+    class ClusterShape = Shape<_1, _1, _1>>
+void launch_sm100_fp8_blockwise_scaled_mm_with_bias(
+    torch::Tensor& out,
+    const torch::Tensor& a,
+    const torch::Tensor& b,
+    const torch::Tensor& scales_a,
+    const torch::Tensor& scales_b,
+    const torch::Tensor& bias) {
+  static constexpr int ScaleMsPerTile = size<0>(ScalesPerTile{});
+  static constexpr int ScaleGranularityM = size<0>(MmaTileShape{}) / ScaleMsPerTile;
+  static constexpr int ScaleGranularityN = size<1>(MmaTileShape{}) / size<1>(ScalesPerTile{});
+  static constexpr int ScaleGranularityK = size<2>(MmaTileShape{}) / size<2>(ScalesPerTile{});
+
+  using ElementAB = cutlass::float_e4m3_t;
+  using ElementA = ElementAB;
+  using ElementB = ElementAB;
+  using ElementD = OutType;
+  using ElementC = ElementD;
+  using LayoutA = cutlass::layout::RowMajor;
+  using LayoutB = cutlass::layout::ColumnMajor;
+  using LayoutD = cutlass::layout::RowMajor;
+  using LayoutC = LayoutD;
+  // This means both SFA and SFB are column-major.
+  using ScaleConfig = cutlass::detail::Sm100BlockwiseScaleConfig<
+      ScaleGranularityM,
+      ScaleGranularityN,
+      ScaleGranularityK,
+      cute::UMMA::Major::MN,
+      cute::UMMA::Major::K>;
+  using LayoutSFA = decltype(ScaleConfig::deduce_layoutSFA());
+  using LayoutSFB = decltype(ScaleConfig::deduce_layoutSFB());
+
+  static constexpr int AlignmentA = 128 / cutlass::sizeof_bits<ElementA>::value;
+  static constexpr int AlignmentB = 128 / cutlass::sizeof_bits<ElementB>::value;
+  static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
+  static constexpr int AlignmentC = AlignmentD;
+
+  using ElementAccumulator = float;
+  using ElementBlockScale = float;
+  using ElementCompute = float;
+  using ArchTag = cutlass::arch::Sm100;
+  using OperatorClass = cutlass::arch::OpClassTensorOp;
+
+  using CollectiveEpilogue = typename cutlass::epilogue::collective::CollectiveBuilder<
+      ArchTag,
+      cutlass::arch::OpClassTensorOp,
+      PerSmTileShape,
+      ClusterShape,
+      EpilogueTileShape,
+      ElementAccumulator,
+      ElementCompute,
+      ElementC,
+      LayoutC,
+      AlignmentC,
+      ElementD,
+      LayoutD,
+      AlignmentD,
+      cutlass::epilogue::TmaWarpSpecialized1Sm>::CollectiveOp;
+
+  using CollectiveMainloop = typename cutlass::gemm::collective::CollectiveBuilder<
+      ArchTag,
+      OperatorClass,
+      ElementA,
+      cute::tuple<LayoutA, LayoutSFA>,
+      AlignmentA,
+      ElementB,
+      cute::tuple<LayoutB, LayoutSFB>,
+      AlignmentB,
+      ElementAccumulator,
+      MmaTileShape,
+      ClusterShape,
+      cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
+          sizeof(typename CollectiveEpilogue::SharedStorage))>,
+      cutlass::gemm::KernelTmaWarpSpecializedBlockwise1SmSm100>::CollectiveOp;
+
+  using GemmKernel = cutlass::gemm::kernel::GemmUniversal<
+      Shape<int, int, int, int>,
+      CollectiveMainloop,
+      CollectiveEpilogue,
+      cutlass::gemm::PersistentScheduler>;
+  using Gemm = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
+
+  Gemm gemm_op;
+
+  int m = a.size(0);
+  int k = a.size(1);
+  int n = b.size(1);
+
+  auto a_ptr = static_cast<ElementAB*>(a.data_ptr());
+  auto b_ptr = static_cast<ElementAB*>(b.data_ptr());
+  auto scales_a_ptr = static_cast<float*>(scales_a.data_ptr());
+  auto scales_b_ptr = static_cast<float*>(scales_b.data_ptr());
+  auto c_ptr = static_cast<ElementD*>(out.data_ptr());
+  auto bias_ptr = static_cast<ElementD*>(bias.data_ptr());
+
+  using StrideA = typename GemmKernel::StrideA;
+  using StrideB = typename GemmKernel::StrideB;
+  using StrideD = typename GemmKernel::StrideD;
+  using StrideC = typename GemmKernel::StrideC;
+
+  StrideA a_stride = cutlass::make_cute_packed_stride(StrideA{}, cute::make_shape(m, k, 1));
+  StrideB b_stride = cutlass::make_cute_packed_stride(StrideB{}, cute::make_shape(n, k, 1));
+  StrideD d_stride = cutlass::make_cute_packed_stride(StrideD{}, cute::make_shape(m, n, 1));
+  // Broadcast bias along M dimension: stride(M)=0 means same bias for each row,
+  // stride(N)=1 reads consecutive elements, stride(L)=0 for batch
+  StrideC bias_stride = cute::make_stride(int64_t(0), cute::C<1>{}, int64_t(0));
+  LayoutSFA layout_SFA = ScaleConfig::tile_atom_to_shape_SFA(make_shape(m, n, k, 1));
+  LayoutSFB layout_SFB = ScaleConfig::tile_atom_to_shape_SFB(make_shape(m, n, k, 1));
+
+  typename GemmKernel::MainloopArguments mainloop_args{
+      a_ptr, a_stride, b_ptr, b_stride, scales_a_ptr, layout_SFA, scales_b_ptr, layout_SFB};
+
+  typename GemmKernel::EpilogueArguments epilogue_args{{}, bias_ptr, bias_stride, c_ptr, d_stride};
+  epilogue_args.thread.alpha = 1.0f;
+  epilogue_args.thread.beta = 1.0f;
+
+  typename GemmKernel::Arguments args = {
+      cutlass::gemm::GemmUniversalMode::kGemm, {m, n, k, 1}, mainloop_args, epilogue_args};
+
+  auto can_implement = gemm_op.can_implement(args);
+  TORCH_CHECK(can_implement == cutlass::Status::kSuccess, cutlassGetStatusString(can_implement))
+
+  size_t workspace_size = gemm_op.get_workspace_size(args);
+  cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
+
+  auto init_status = gemm_op.initialize(args, workspace.get());
+  TORCH_CHECK(init_status == cutlass::Status::kSuccess, cutlassGetStatusString(init_status));
+
+  auto stream = at::cuda::getCurrentCUDAStream(a.get_device());
+
+  auto status = gemm_op.run(stream);
+  TORCH_CHECK(status == cutlass::Status::kSuccess, cutlassGetStatusString(status))
+}
+
 template <typename OutType>
 void sm100_fp8_blockwise_dispatch_shape(
     torch::Tensor& out,
     const torch::Tensor& a,
     const torch::Tensor& b,
     const torch::Tensor& scales_a,
-    const torch::Tensor& scales_b) {
+    const torch::Tensor& scales_b,
+    const c10::optional<torch::Tensor>& bias) {
   if (a.size(0) <= 128) {
     using MmaTileShape = Shape<_64, _128, _128>;
     using PerSmTileShape = Shape<_64, _128, _128>;
     using EpilogueTileShape = Shape<_64, _64>;
     using ScalesPerTile = Shape<_64, _1, _1>;
-    launch_sm100_fp8_blockwise_scaled_mm<OutType, MmaTileShape, PerSmTileShape, EpilogueTileShape, ScalesPerTile>(
-        out, a, b, scales_a, scales_b);
+    if (bias.has_value()) {
+      launch_sm100_fp8_blockwise_scaled_mm_with_bias<
+          OutType,
+          MmaTileShape,
+          PerSmTileShape,
+          EpilogueTileShape,
+          ScalesPerTile>(out, a, b, scales_a, scales_b, *bias);
+    } else {
+      launch_sm100_fp8_blockwise_scaled_mm<OutType, MmaTileShape, PerSmTileShape, EpilogueTileShape, ScalesPerTile>(
+          out, a, b, scales_a, scales_b);
+    }
   } else {
     using MmaTileShape = Shape<_128, _128, _128>;
     using PerSmTileShape = Shape<_128, _128, _128>;
     using EpilogueTileShape = Shape<_128, _64>;
     using ScalesPerTile = Shape<_128, _1, _1>;
-    launch_sm100_fp8_blockwise_scaled_mm<OutType, MmaTileShape, PerSmTileShape, EpilogueTileShape, ScalesPerTile>(
-        out, a, b, scales_a, scales_b);
+    if (bias.has_value()) {
+      launch_sm100_fp8_blockwise_scaled_mm_with_bias<
+          OutType,
+          MmaTileShape,
+          PerSmTileShape,
+          EpilogueTileShape,
+          ScalesPerTile>(out, a, b, scales_a, scales_b, *bias);
+    } else {
+      launch_sm100_fp8_blockwise_scaled_mm<OutType, MmaTileShape, PerSmTileShape, EpilogueTileShape, ScalesPerTile>(
+          out, a, b, scales_a, scales_b);
+    }
   }
 }
 
@@ -427,7 +587,8 @@ torch::Tensor fp8_blockwise_scaled_mm(
     const torch::Tensor& mat_b,
     const torch::Tensor& scales_a,
     const torch::Tensor& scales_b,
-    const torch::Dtype& out_dtype) {
+    const torch::Dtype& out_dtype,
+    const c10::optional<torch::Tensor>& bias) {
   TORCH_CHECK(mat_a.is_cuda(), "mat_a must be a CUDA tensor");
   TORCH_CHECK(mat_b.is_cuda(), "mat_b must be a CUDA tensor");
   TORCH_CHECK(mat_a.dim() == 2, "mat_a must be a 2D tensor");
@@ -472,6 +633,7 @@ torch::Tensor fp8_blockwise_scaled_mm(
 #if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
 #if defined CUDA_VERSION && CUDA_VERSION >= 12000
   if (sm_version == 90) {
+    TORCH_CHECK(!bias.has_value(), "fp8_blockwise_scaled_mm with bias is only supported on SM100, got SM", sm_version);
     torch::Tensor scales_b_contiguous = scales_b.contiguous();
     if (out_dtype == torch::kBFloat16) {
       cutlass_gemm_blockwise_sm90_fp8_dispatch<cutlass::bfloat16_t>(
@@ -492,11 +654,19 @@ torch::Tensor fp8_blockwise_scaled_mm(
       || sm_version == 103
 #endif
   ) {
+    if (bias.has_value()) {
+      TORCH_CHECK(bias->is_cuda(), "bias must be a CUDA tensor");
+      TORCH_CHECK(bias->dim() == 1, "bias must be a 1D tensor");
+      TORCH_CHECK(bias->size(0) == mat_b.size(1), "bias size must match N dimension (mat_b columns)");
+      TORCH_CHECK(bias->scalar_type() == out_dtype, "bias dtype must match out_dtype");
+      TORCH_CHECK(bias->is_contiguous(), "bias must be contiguous");
+    }
     if (out_dtype == torch::kBFloat16) {
       sm100_fp8_blockwise_dispatch_shape<cutlass::bfloat16_t>(
-          out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b);
+          out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b, bias);
     } else {
-      sm100_fp8_blockwise_dispatch_shape<cutlass::half_t>(out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b);
+      sm100_fp8_blockwise_dispatch_shape<cutlass::half_t>(
+          out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b, bias);
     }
     return out_padded.slice(0, 0, original_rows);
   }
@@ -506,6 +676,7 @@ torch::Tensor fp8_blockwise_scaled_mm(
 #if defined(CUTLASS_ARCH_MMA_SM120A_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM120_SUPPORTED)
 #if defined(CUDA_VERSION) && CUDA_VERSION >= 12080
   if (sm_version >= 120) {
+    TORCH_CHECK(!bias.has_value(), "fp8_blockwise_scaled_mm with bias is only supported on SM100, got SM", sm_version);
     if (out_dtype == torch::kBFloat16) {
       sm120_fp8_blockwise_dispatch_shape<cutlass::bfloat16_t>(
           out_padded, mat_a_padded, mat_b, scales_a_padded, scales_b);

sgl-kernel/include/sgl_kernel_ops.h

@@ -191,7 +191,8 @@ torch::Tensor fp8_blockwise_scaled_mm(
     const torch::Tensor& mat_b,
     const torch::Tensor& scales_a,
     const torch::Tensor& scales_b,
-    const torch::Dtype& out_dtype);
+    const torch::Dtype& out_dtype,
+    const c10::optional<torch::Tensor>& bias);
 void sgl_per_token_group_quant_8bit(
     at::Tensor input,
     at::Tensor output_q,

sgl-kernel/python/sgl_kernel/gemm.py

@@ -21,13 +21,14 @@ def int8_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None):
     )
 
 
-def fp8_blockwise_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype):
+def fp8_blockwise_scaled_mm(mat_a, mat_b, scales_a, scales_b, out_dtype, bias=None):
     return torch.ops.sgl_kernel.fp8_blockwise_scaled_mm.default(
         mat_a,
         mat_b,
         scales_a,
         scales_b,
         out_dtype,
+        bias,
     )
 
 

sgl-kernel/tests/test_fp8_blockwise_gemm.py

@@ -60,7 +60,14 @@ def group_broadcast(t, shape):
     return output
 
 
-def _test_accuracy_once(M, N, K, out_dtype, device):
+def _is_sm100():
+    if not torch.cuda.is_available():
+        return False
+    major, minor = torch.cuda.get_device_capability()
+    return major == 10 and minor in (0, 3)
+
+
+def _test_accuracy_once(M, N, K, out_dtype, device, use_bias=False):
     fp8_info = torch.finfo(torch.float8_e4m3fn)
     fp8_max, fp8_min = fp8_info.max, fp8_info.min
     a_fp32 = (torch.rand(M, K, dtype=torch.float32, device=device) - 0.5) * 2 * fp8_max
@@ -75,19 +82,32 @@ def _test_accuracy_once(M, N, K, out_dtype, device):
     scale_b = torch.randn(scale_b_shape, device=device, dtype=torch.float32) * 0.001
     scale_a = scale_a.t().contiguous().t()
     scale_b = scale_b.t().contiguous().t()
-    o = baseline_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype)
-    o1 = fp8_blockwise_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype)
+    bias = torch.randn(N, device=device, dtype=out_dtype) * 0.01 if use_bias else None
+    o = baseline_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype, bias=bias)
+    o1 = fp8_blockwise_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype, bias=bias)
     rtol = 0.02
     atol = 1
     torch.testing.assert_close(o, o1, rtol=rtol, atol=atol)
 
 
+@pytest.mark.parametrize(
+    "use_bias",
+    [
+        False,
+        pytest.param(
+            True,
+            marks=pytest.mark.skipif(
+                not _is_sm100(), reason="bias only supported on sm100 GPUs"
+            ),
+        ),
+    ],
+)
 @pytest.mark.parametrize("M", [1, 3, 5, 127, 128, 512, 1024, 4096])
 @pytest.mark.parametrize("N", [128, 512, 1024, 4096, 8192, 14080])
 @pytest.mark.parametrize("K", [512, 1024, 4096, 8192, 14080, 16384])
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
-def test_accuracy(M, N, K, out_dtype):
-    _test_accuracy_once(M, N, K, out_dtype, "cuda")
+def test_accuracy(M, N, K, out_dtype, use_bias):
+    _test_accuracy_once(M, N, K, out_dtype, "cuda", use_bias=use_bias)
 
 
 if __name__ == "__main__":