[SM100] Fuse bias addition into fp8_blockwise_scaled_mm epilogue

[haowen-han]

Motivation

In FP8 blockwise linear layers the bias addition (output += bias) is currently a separate element-wise kernel launch after GEMM, which incurs extra global memory round-trip and kernel launch overhead on SM100 GPUs. Fusing bias into the CUTLASS epilogue eliminates this overhead.

Modifications

• sgl-kernel/csrc/gemm/fp8_blockwise_gemm_kernel.cu: Added launch_sm100_fp8_blockwise_scaled_mm_with_bias template function that uses ElementC = ElementD (non-void) in the CUTLASS CollectiveBuilder, passes bias pointer with broadcast stride (0, 1, 0) along M dimension, and sets alpha=1.0, beta=1.0 in epilogue arguments. Updated sm100_fp8_blockwise_dispatch_shape to dispatch to the bias-fused variant when bias is provided. Added explicit rejection of bias for SM90 and SM120+ code paths.
• sgl-kernel/csrc/common_extension.cc: Extended fp8_blockwise_scaled_mm op schema with optional Tensor? bias=None parameter.
• sgl-kernel/include/sgl_kernel_ops.h: Updated C++ declaration accordingly.
• sgl-kernel/python/sgl_kernel/gemm.py: Forwarded new bias kwarg through Python binding.
• python/sglang/srt/layers/quantization/fp8_utils.py: In cutlass_w8a8_block_fp8_linear_with_fallback, detect SM100 via device capability and pass bias into the fused path; non-SM100 GPUs fall back to the existing separate output += bias.
• sgl-kernel/tests/test_fp8_blockwise_gemm.py: Added _test_accuracy_once_with_bias and parametrized test_accuracy_with_bias covering M ∈ {1,3,5,127,128,512,1024,4096}, N ∈ {128..14080}, K ∈ {512..16384}, both bf16 and fp16, gated by SM100 skip condition.

Accuracy Tests

Added test_accuracy_with_bias test comparing fused-bias output against reference (baseline_scaled_mm with bias) at rtol=0.02, atol=1. Test is gated to run only on SM100 GPUs via @pytest.mark.skipif.

Speed Tests and Profiling

running the following benchmark code using B200 gpu:

"""Benchmark: SM100 FP8 blockwise GEMM with fused bias vs. separate bias-add.

Compares two paths:
  - Fused:   fp8_blockwise_scaled_mm(..., bias=bias)          [epilogue alpha=1,beta=1]
  - Separate: out = fp8_blockwise_scaled_mm(...); out += bias  [two kernel launches]

Usage:
    python bench_bias_fusion.py
"""

import torch
import triton

from sgl_kernel import fp8_blockwise_scaled_mm


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 cdiv(a, b):
    return -(a // -b)


def scale_shape(shape, group_shape):
    return tuple(cdiv(shape[i], group_shape[i]) for i in range(len(group_shape)))


@triton.testing.perf_report(
    triton.testing.Benchmark(
        x_names=["M"],
        x_vals=[256, 512, 1024, 2048, 4096, 8192, 16384],
        x_log=True,
        line_arg="mode",
        line_vals=["fused", "separate"],
        line_names=["Fused (bias in epilogue)", "Separate (gemm + add)"],
        styles=[("green", "-"), ("red", "--")],
        ylabel="ms",
        plot_name="fp8_blockwise_gemm_bias_fusion",
        args={},
    )
)
def benchmark(mode, M, N, K, out_dtype):
    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="cuda") - 0.5) * 2 * fp8_max
    a_fp8 = a_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)

    b_fp32 = (torch.rand(N, K, dtype=torch.float32, device="cuda") - 0.5) * 2 * fp8_max
    b_fp8 = b_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn).t()

    scale_a_group_shape = (1, 128)
    scale_b_group_shape = (128, 128)
    scale_a_shape = scale_shape(a_fp8.shape, scale_a_group_shape)
    scale_b_shape = scale_shape(b_fp8.shape, scale_b_group_shape)

    scale_a = torch.randn(scale_a_shape, device="cuda", dtype=torch.float32) * 0.001
    scale_b = torch.randn(scale_b_shape, device="cuda", dtype=torch.float32) * 0.001
    scale_a = scale_a.t().contiguous().t()
    scale_b = scale_b.t().contiguous().t()

    bias = torch.randn(N, device="cuda", dtype=out_dtype) * 0.01

    # Warmup
    _ = fp8_blockwise_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype)
    _ = fp8_blockwise_scaled_mm(a_fp8, b_fp8, scale_a, scale_b, out_dtype, bias=bias)
    torch.cuda.synchronize()

    quantiles = [0.5, 0.2, 0.8]

    def run_fused():
        return fp8_blockwise_scaled_mm(
            a_fp8, b_fp8, scale_a, scale_b, out_dtype, bias=bias
        )

    def run_separate():
        out = fp8_blockwise_scaled_mm(
            a_fp8, b_fp8, scale_a, scale_b, out_dtype
        )
        out += bias
        return out

    if mode == "fused":
        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
            run_fused, quantiles=quantiles
        )
    elif mode == "separate":
        ms, min_ms, max_ms = triton.testing.do_bench_cudagraph(
            run_separate, quantiles=quantiles
        )

    return ms * 1000, max_ms * 1000, min_ms * 1000


if __name__ == "__main__":
    if not _is_sm100():
        print("SKIP: bias fusion is only supported on SM100/SM103 GPUs")
        exit(0)

    test_shapes = [
        # (N, K, out_dtype) — representative DeepSeek-V3 MoE / Qwen FP8 shapes
        (7168, 7168, torch.bfloat16),
        (7168, 16384, torch.bfloat16),
        (16384, 7168, torch.float16),
        (14080, 16384, torch.bfloat16),  # typical FFN dim
        (8192, 8192, torch.float16),
    ]

    for N, K, out_dtype in test_shapes:
        print(f"\n{'='*60}")
        print(f"N={N}, K={K}, dtype={'bf16' if out_dtype==torch.bfloat16 else 'fp16'}")
        print(f"{'='*60}")
        try:
            benchmark.run(print_data=True, N=N, K=K, out_dtype=out_dtype)
        except Exception as e:
            print(f"  ERROR: {e}")

    print("\nBenchmark finished!")

here is the result:

============================================================
N=7168, K=7168, dtype=bf16
============================================================
fp8_blockwise_gemm_bias_fusion:
         M  Fused (bias in epilogue)  Separate (gemm + add)
0    256.0                 19.376838              23.838453
1    512.0                 35.659245              42.251312
2   1024.0                 67.925394              78.150920
3   2048.0                131.872603             146.603412
4   4096.0                263.217060             286.696246
5   8192.0                468.621148             579.818849
6  16384.0                918.596848            1141.510010

============================================================
N=7168, K=16384, dtype=bf16
============================================================
fp8_blockwise_gemm_bias_fusion:
         M  Fused (bias in epilogue)  Separate (gemm + add)
0    256.0                 43.785433              47.810112
1    512.0                 81.634298              86.983676
2   1024.0                153.299782             168.448712
3   2048.0                285.402742             303.855991
4   4096.0                583.059466             608.137285
5   8192.0               1082.428631            1297.316297
6  16384.0               2441.153844            2524.587359

============================================================
N=16384, K=7168, dtype=fp16
============================================================
fp8_blockwise_gemm_bias_fusion:
         M  Fused (bias in epilogue)  Separate (gemm + add)
0    256.0                 42.899774              49.212395
1    512.0                 78.094634              90.390543
2   1024.0                146.660226             167.952741
3   2048.0                293.780782             326.080557
4   4096.0                599.295009             651.446079
5   8192.0               1280.669912            1328.287951
6  16384.0               2664.572001            3237.196732

============================================================
N=14080, K=16384, dtype=bf16
============================================================
fp8_blockwise_gemm_bias_fusion:
         M  Fused (bias in epilogue)  Separate (gemm + add)
0    256.0                 84.381665              91.406007
1    512.0                139.445754             150.312793
2   1024.0                262.667424             285.423491
3   2048.0                534.803549             565.675649
4   4096.0               1135.408000            1194.958985
5   8192.0               2607.092023            2585.369246
6  16384.0               5389.096022            5207.893213

============================================================
N=8192, K=8192, dtype=fp16
============================================================
fp8_blockwise_gemm_bias_fusion:
         M  Fused (bias in epilogue)  Separate (gemm + add)
0    256.0                 23.325837              26.356896
1    512.0                 42.567962              49.080000
2   1024.0                 87.857144              97.445689
3   2048.0                166.243331             181.424160
4   4096.0                333.237237             362.018457
5   8192.0                652.479536             810.797899
6  16384.0               1271.440955            1512.180010

Benchmark finished!

Checklist

• [ ✓] Format your code according to the Format code with pre-commit.
• [✓ ] Add unit tests according to the Run and add unit tests.
• [ ✓] Update documentation according to Write documentations.
• [ ✓] Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• [ ✓] Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

[gemini-code-assist]

Code Review

This pull request implements bias fusion for FP8 blockwise scaled matrix multiplication on SM100 devices. Key changes include the addition of a new CUDA kernel using CUTLASS to handle fused bias, updates to the C++ and Python interfaces to support an optional bias tensor, and logic in the quantization utilities to enable this fusion on compatible hardware. Review feedback highlights a performance bottleneck caused by synchronous memory allocation within the kernel launch path, which should be replaced with a pre-allocated workspace. Additionally, there are suggestions to remove unused template parameters and type aliases, and to refactor duplicated dispatch logic to improve maintainability.

[gemini-code-assist]

Using cutlass::device_memory::allocation inside the kernel launch path performs a synchronous cudaMalloc and cudaFree. This is a significant performance bottleneck in high-frequency inference scenarios as it triggers device synchronization. It is highly recommended to pass a workspace tensor from the Python layer (e.g., using a pre-allocated buffer pool) and use its data pointer here instead.

[gemini-code-assist]

The template parameter TileSizeM_ is defined but not used within the launch_sm100_fp8_blockwise_scaled_mm_with_bias function. It should be removed to clean up the template signature.

    class ClusterShape = Shape<_1, _1, _1>>

[gemini-code-assist]

The type alias ElementBlockScale is defined but never used in this function. It should be removed.

  using ElementAccumulator = float;
  using ElementCompute = float;

[gemini-code-assist]

The dispatch logic for the bias-fused variant is duplicated for both M <= 128 and M > 128 cases. This increases code verbosity and maintenance effort. Consider refactoring this to determine the template parameters first and then perform a single check for bias.has_value() to launch the appropriate kernel.

[hnyls2002]

Hi @haowen-han, the CUTLASS SM90/SM100 fp8_blockwise_gemm_kernel.cu this PR patches has been removed from the tree (see #30438, which moved the SM120 path to JIT and deleted the SM90/SM100 CUTLASS kernels). The bias fusion idea is still valid but this PR would need a full rewrite against the new JIT kernel layout (python/sglang/kernels/jit/csrc/gemm/fp8_blockwise/). Closing as obsolete - please reopen or resubmit against the new kernels.