Delete CUTLASS FP8 blockwise for SM90 and SM100, move SM120 to JIT and add SwapAB (#30438)
Co-authored-by: Brayden Zhong <brayden.zhong@radixark.ai> Co-authored-by: Claude Sonnet 5 <noreply@anthropic.com> Co-authored-by: root <root@sgl-b300-inference.datacrunch.io> Co-authored-by: Brayden Zhong <brayden@radixark.ai>
This commit is contained in:
Brayden Zhong
co-authored by
Brayden Zhong
Claude Sonnet 5
root
Brayden Zhong
Brayden Zhong
Claude Sonnet 5
root
Brayden Zhong
parent
c124bec99d
commit
7431f35fd8
@@ -0,0 +1,103 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import sys
|
||||
|
||||
import torch
|
||||
import triton
|
||||
|
||||
from sglang.jit_kernel.benchmark.utils import get_benchmark_range, run_benchmark
|
||||
from sglang.jit_kernel.fp8_blockwise_gemm import fp8_blockwise_scaled_mm
|
||||
from sglang.srt.utils import is_sm120_supported
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=5,
|
||||
stage="base-b-kernel-benchmark",
|
||||
runner_config="1-gpu-large",
|
||||
)
|
||||
|
||||
|
||||
def _make_inputs(m: int, n: int, k: int, device: str = "cuda"):
|
||||
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
|
||||
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=device) - 0.5) * 2 * fp8_max
|
||||
b_fp8 = b_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn).t()
|
||||
|
||||
scale_a = torch.randn((m, k // 128), device=device, dtype=torch.float32) * 0.001
|
||||
scale_b = (
|
||||
torch.randn((k // 128, n // 128), device=device, dtype=torch.float32) * 0.001
|
||||
)
|
||||
scale_a = scale_a.t().contiguous().t()
|
||||
scale_b = scale_b.t().contiguous().t()
|
||||
return a_fp8, b_fp8, scale_a, scale_b
|
||||
|
||||
|
||||
def _torch_ref(a_fp8, b_fp8, scale_a, scale_b):
|
||||
def group_broadcast(t, shape):
|
||||
for i, s in enumerate(shape):
|
||||
if t.shape[i] != s and t.shape[i] != 1:
|
||||
assert s % t.shape[i] == 0
|
||||
t = (
|
||||
t.unsqueeze(i + 1)
|
||||
.expand(*t.shape[: i + 1], s // t.shape[i], *t.shape[i + 1 :])
|
||||
.flatten(i, i + 1)
|
||||
)
|
||||
return t
|
||||
|
||||
sa = group_broadcast(scale_a, a_fp8.shape)
|
||||
sb = group_broadcast(scale_b, b_fp8.shape)
|
||||
return torch.mm(sa * a_fp8.to(torch.float32), sb * b_fp8.to(torch.float32)).to(
|
||||
torch.bfloat16
|
||||
)
|
||||
|
||||
|
||||
shape_range = get_benchmark_range(
|
||||
full_range=[
|
||||
(16, 4096, 4096), # swapAB tile N=32
|
||||
(64, 4096, 4096), # swapAB tile N=64
|
||||
(128, 4096, 4096), # non-swap 128
|
||||
(512, 4096, 4096),
|
||||
(1024, 8192, 4096),
|
||||
],
|
||||
ci_range=[(16, 4096, 4096), (128, 4096, 4096)],
|
||||
)
|
||||
|
||||
|
||||
@triton.testing.perf_report(
|
||||
triton.testing.Benchmark(
|
||||
x_names=["m", "n", "k"],
|
||||
x_vals=shape_range,
|
||||
x_log=False,
|
||||
line_arg="provider",
|
||||
line_vals=["jit", "torch_ref"],
|
||||
line_names=["JIT FP8 Blockwise GEMM", "Torch Ref"],
|
||||
styles=[("green", "-"), ("blue", "-")],
|
||||
ylabel="us",
|
||||
plot_name="fp8-blockwise-scaled-mm-performance",
|
||||
args={},
|
||||
)
|
||||
)
|
||||
def benchmark(m, n, k, provider):
|
||||
a_fp8, b_fp8, scale_a, scale_b = _make_inputs(m, n, k)
|
||||
|
||||
if provider == "jit":
|
||||
fn = lambda: fp8_blockwise_scaled_mm(
|
||||
a_fp8, b_fp8, scale_a, scale_b, out_dtype=torch.bfloat16
|
||||
)
|
||||
elif provider == "torch_ref":
|
||||
fn = lambda: _torch_ref(a_fp8, b_fp8, scale_a, scale_b)
|
||||
else:
|
||||
raise ValueError(f"Unknown provider: {provider}")
|
||||
|
||||
return run_benchmark(fn)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
if not is_sm120_supported():
|
||||
print(
|
||||
"[skip] fp8_blockwise_scaled_mm benchmark requires SM120 with CUDA 12.8+."
|
||||
)
|
||||
sys.exit(0)
|
||||
benchmark.run(print_data=True)
|
||||
@@ -0,0 +1,90 @@
|
||||
import sys
|
||||
from typing import Optional, Type
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.jit_kernel.fp8_blockwise_gemm import fp8_blockwise_scaled_mm
|
||||
from sglang.srt.utils import is_sm120_supported
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=30,
|
||||
stage="base-b",
|
||||
runner_config="1-gpu-small",
|
||||
)
|
||||
|
||||
|
||||
def cdiv(a: int, b: int) -> int:
|
||||
return -(a // -b)
|
||||
|
||||
|
||||
def scale_shape(shape, group_shape):
|
||||
assert len(shape) == len(group_shape)
|
||||
return tuple(cdiv(shape[i], group_shape[i]) for i in range(len(group_shape)))
|
||||
|
||||
|
||||
def baseline_scaled_mm(
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
scale_a: torch.Tensor,
|
||||
scale_b: torch.Tensor,
|
||||
out_dtype: Type[torch.dtype],
|
||||
bias: Optional[torch.Tensor] = None,
|
||||
) -> torch.Tensor:
|
||||
def group_broadcast(t, shape):
|
||||
for i, s in enumerate(shape):
|
||||
if t.shape[i] != s and t.shape[i] != 1:
|
||||
assert s % t.shape[i] == 0
|
||||
t = (
|
||||
t.unsqueeze(i + 1)
|
||||
.expand(*t.shape[: i + 1], s // t.shape[i], *t.shape[i + 1 :])
|
||||
.flatten(i, i + 1)
|
||||
)
|
||||
return t
|
||||
|
||||
scale_a = group_broadcast(scale_a, a.shape)
|
||||
scale_b = group_broadcast(scale_b, b.shape)
|
||||
output = torch.mm(
|
||||
(scale_a * a.to(dtype=torch.float32)), (scale_b * b.to(dtype=torch.float32))
|
||||
).to(out_dtype)
|
||||
if bias is not None:
|
||||
output = output + bias
|
||||
return output
|
||||
|
||||
|
||||
def _test_accuracy_once(M, N, K, out_dtype, device):
|
||||
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
|
||||
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=device) - 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=device, dtype=torch.float32) * 0.001
|
||||
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)
|
||||
rtol = 0.02
|
||||
atol = 1
|
||||
torch.testing.assert_close(o, o1, rtol=rtol, atol=atol)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_sm120_supported(), reason="fp8_blockwise_scaled_mm requires SM120 (>= 12.0)"
|
||||
)
|
||||
@pytest.mark.parametrize("M", [1, 3, 5, 32, 48, 64, 127, 128, 512, 1024, 4096])
|
||||
@pytest.mark.parametrize("N", [128, 512, 1024, 4096, 8192])
|
||||
@pytest.mark.parametrize("K", [512, 1024, 4096, 8192])
|
||||
@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")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -1,136 +0,0 @@
|
||||
"""Unit tests for the row-padded quant path of the cutlass FP8 blockwise linear.
|
||||
|
||||
`cutlass_w8a8_block_fp8_linear_with_fallback` quantizes activations into
|
||||
row-aligned buffers (`sglang_per_token_group_quant_fp8_row_padded`) so the
|
||||
`fp8_blockwise_scaled_mm` wrapper's per-call mat_a/scales_a padding short-
|
||||
circuits. These tests pin the invariant that this is numerically identical to
|
||||
the legacy unpadded path, across both row-aligned and unaligned M.
|
||||
"""
|
||||
|
||||
import unittest
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.kernels.ops.quantization.fp8_kernel import (
|
||||
fp8_dtype,
|
||||
per_token_group_quant_fp8,
|
||||
sglang_per_token_group_quant_fp8_row_padded,
|
||||
)
|
||||
from sglang.srt.layers.quantization.fp8_utils import (
|
||||
_check_cutlass_block_fp8_hardware_support,
|
||||
cutlass_w8a8_block_fp8_linear_with_fallback,
|
||||
)
|
||||
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")
|
||||
|
||||
_FP8_MAX = torch.finfo(fp8_dtype).max
|
||||
_BLOCK = 128
|
||||
# Cover M == 1 (greedy decode), small unaligned M (speculative draft tokens),
|
||||
# the 4-row alignment boundary, and a large aligned batch.
|
||||
_M_VALUES = [1, 2, 3, 4, 5, 7, 13, 16, 31, 64, 256]
|
||||
|
||||
|
||||
def _quant_weight_blockwise(weight_bf16: torch.Tensor, block: int = _BLOCK):
|
||||
"""Block-quantize a (N, K) bf16 weight to fp8 with (N//block, K//block) fp32 scales."""
|
||||
n, k = weight_bf16.shape
|
||||
assert n % block == 0 and k % block == 0
|
||||
w = weight_bf16.float().reshape(n // block, block, k // block, block)
|
||||
amax = w.abs().amax(dim=(1, 3)).clamp(min=1e-12) # (N//block, K//block)
|
||||
scale = amax / _FP8_MAX
|
||||
wq = (w / scale[:, None, :, None]).clamp(-_FP8_MAX, _FP8_MAX).to(fp8_dtype)
|
||||
return wq.reshape(n, k), scale.to(torch.float32)
|
||||
|
||||
|
||||
def _legacy_cutlass_linear(x_2d, weight, weight_scale):
|
||||
"""The pre-optimization path: unpadded quant, wrapper pads internally."""
|
||||
from sgl_kernel import fp8_blockwise_scaled_mm
|
||||
|
||||
q_input, x_scale = per_token_group_quant_fp8(x_2d, _BLOCK, column_major_scales=True)
|
||||
return fp8_blockwise_scaled_mm(
|
||||
q_input, weight.T, x_scale, weight_scale.T, out_dtype=x_2d.dtype
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(
|
||||
_check_cutlass_block_fp8_hardware_support(),
|
||||
"cutlass block FP8 requires Hopper (SM90) or newer",
|
||||
)
|
||||
class TestFP8BlockwiseRowPadding(CustomTestCase):
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
cls.K = 512
|
||||
cls.N = 256
|
||||
torch.manual_seed(0)
|
||||
|
||||
def test_quant_buffers_row_aligned(self):
|
||||
"""Row-padded quant returns 4-aligned, M-major buffers whose live rows
|
||||
match the legacy column-major quant bit-for-bit."""
|
||||
for m in _M_VALUES:
|
||||
x = torch.randn(m, self.K, device="cuda", dtype=torch.bfloat16) * 0.1
|
||||
xq, xs = sglang_per_token_group_quant_fp8_row_padded(x, _BLOCK)
|
||||
m_pad = (m + 3) // 4 * 4
|
||||
|
||||
self.assertEqual(xq.shape, (m_pad, self.K), f"M={m}")
|
||||
self.assertEqual(xs.shape[0], m_pad, f"M={m}")
|
||||
# scales_a must stay M-major (stride(0) == 1) for the kernel contract.
|
||||
self.assertEqual(xs.stride(0), 1, f"M={m}")
|
||||
|
||||
xq_ref, xs_ref = per_token_group_quant_fp8(
|
||||
x, _BLOCK, column_major_scales=True
|
||||
)
|
||||
self.assertEqual(xq_ref.shape, (m, self.K), f"M={m}")
|
||||
# Live rows are produced by the same kernel, so they must be identical.
|
||||
self.assertTrue(
|
||||
torch.equal(xq[:m].view(torch.uint8), xq_ref.view(torch.uint8)),
|
||||
f"quantized activation mismatch at M={m}",
|
||||
)
|
||||
torch.testing.assert_close(xs[:m], xs_ref, atol=0.0, rtol=0.0)
|
||||
|
||||
def test_gemm_bit_exact_vs_legacy(self):
|
||||
"""The full linear (row-padded) is bit-identical to the legacy unpadded GEMM."""
|
||||
weight_bf16 = (
|
||||
torch.randn(self.N, self.K, device="cuda", dtype=torch.bfloat16) * 0.1
|
||||
)
|
||||
weight, weight_scale = _quant_weight_blockwise(weight_bf16)
|
||||
|
||||
for m in _M_VALUES:
|
||||
x = torch.randn(m, self.K, device="cuda", dtype=torch.bfloat16) * 0.1
|
||||
|
||||
out_ref = _legacy_cutlass_linear(x, weight, weight_scale)
|
||||
out_new = cutlass_w8a8_block_fp8_linear_with_fallback(
|
||||
input=x,
|
||||
weight=weight,
|
||||
block_size=[_BLOCK, _BLOCK],
|
||||
weight_scale=weight_scale,
|
||||
)
|
||||
|
||||
self.assertEqual(out_new.shape, (m, self.N), f"M={m}")
|
||||
self.assertTrue(
|
||||
torch.equal(out_ref, out_new),
|
||||
f"row-padded GEMM differs from legacy at M={m}: "
|
||||
f"max_abs_diff={(out_ref.float() - out_new.float()).abs().max().item()}",
|
||||
)
|
||||
|
||||
def test_linear_matches_bf16_reference(self):
|
||||
"""Sanity: the FP8 linear stays close to a bf16 reference matmul."""
|
||||
weight_bf16 = (
|
||||
torch.randn(self.N, self.K, device="cuda", dtype=torch.bfloat16) * 0.1
|
||||
)
|
||||
weight, weight_scale = _quant_weight_blockwise(weight_bf16)
|
||||
|
||||
for m in [1, 5, 64]:
|
||||
x = torch.randn(m, self.K, device="cuda", dtype=torch.bfloat16) * 0.1
|
||||
ref = (x.float() @ weight_bf16.float().T).to(torch.bfloat16)
|
||||
out = cutlass_w8a8_block_fp8_linear_with_fallback(
|
||||
input=x,
|
||||
weight=weight,
|
||||
block_size=[_BLOCK, _BLOCK],
|
||||
weight_scale=weight_scale,
|
||||
)
|
||||
torch.testing.assert_close(out, ref, atol=0.5, rtol=0.1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user