wehub-resource-sync
468 lines
16 KiB
Python
468 lines
16 KiB
Python
# Licensed to the Apache Software Foundation (ASF) under one
|
|
# or more contributor license agreements. See the NOTICE file
|
|
# distributed with this work for additional information
|
|
# regarding copyright ownership. The ASF licenses this file
|
|
# to you under the Apache License, Version 2.0 (the
|
|
# "License"); you may not use this file except in compliance
|
|
# with the License. You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing,
|
|
# software distributed under the License is distributed on an
|
|
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
|
|
# KIND, either express or implied. See the License for the
|
|
# specific language governing permissions and limitations
|
|
# under the License.
|
|
# ruff: noqa: E501, F401, RUF005
|
|
|
|
"""Test for FlashInfer GroupedGemm TVM integration"""
|
|
|
|
import math
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
|
|
import tvm
|
|
import tvm.testing
|
|
from tvm import relax
|
|
|
|
DEFAULT_WORKSPACE_SIZE = 32 * 1024 * 1024
|
|
fp8_dtype = "float8_e4m3fn"
|
|
|
|
|
|
###########################################
|
|
################# Helpers #################
|
|
###########################################
|
|
def has_flashinfer():
|
|
"""Check if FlashInfer is available with the SM100 grouped-gemm symbol."""
|
|
try:
|
|
from flashinfer.gemm import ( # pylint: disable=import-outside-toplevel,unused-import
|
|
gen_gemm_sm100_module,
|
|
)
|
|
|
|
from tvm.relax.backend.cuda import ( # pylint: disable=import-outside-toplevel
|
|
flashinfer,
|
|
)
|
|
|
|
return True
|
|
except ImportError:
|
|
return False
|
|
|
|
|
|
def has_cutlass():
|
|
"""Check if CUTLASS is available for SM90+ operations"""
|
|
if not tvm.get_global_func("device_api.cuda", True):
|
|
return False
|
|
try:
|
|
import pynvml # pylint: disable=import-outside-toplevel
|
|
|
|
pynvml.nvmlInit()
|
|
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
|
|
major, minor = pynvml.nvmlDeviceGetCudaComputeCapability(handle)
|
|
return major >= 9 # SM90+
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def calc_diff(x: np.ndarray, y: np.ndarray):
|
|
denominator = (x * x + y * y).sum()
|
|
sim = 2 * (x * y).sum() / denominator
|
|
return 1 - sim
|
|
|
|
|
|
def quantize_fp8(x, scale_shape, tile_shape, scale_major_mode):
|
|
from einops import rearrange, reduce, repeat
|
|
|
|
"""
|
|
Quantizes a 2D or 3D tensor to FP8.
|
|
|
|
Args:
|
|
x (torch.Tensor): The 2D or 3D input tensor.
|
|
scale_shape (tuple): The shape of the scale tensor.
|
|
tile_shape (tuple): The shape of the tiles.
|
|
scale_major_mode (str): The tiling order, "K" for row-major like,
|
|
or another value for column-major like.
|
|
|
|
Returns:
|
|
tuple: A tuple containing the quantized FP8 tensor and the
|
|
calculated float32 scales.
|
|
"""
|
|
# 1. Assertions and Initial Setup
|
|
ndim = x.ndim
|
|
assert ndim == len(scale_shape) == len(tile_shape)
|
|
|
|
fp8_info = torch.finfo(torch.float8_e4m3fn)
|
|
fp8_amax = torch.tensor(fp8_info.max, device=x.device, dtype=torch.float32)
|
|
|
|
# 2. Tiling and Scale Calculation
|
|
if ndim == 2:
|
|
s0, s1 = scale_shape
|
|
t0, t1 = tile_shape
|
|
if scale_major_mode == "K":
|
|
# Tile x and find the max absolute value in each tile
|
|
x_tiled = rearrange(x, "(s0 t0) (s1 t1) -> s0 s1 t0 t1", s0=s0, s1=s1)
|
|
abs_max = reduce(x_tiled.abs(), "s0 s1 t0 t1 -> s0 s1", "max").clamp(1e-4)
|
|
x_scale = abs_max / fp8_amax
|
|
x_scale = torch.pow(2.0, torch.ceil(torch.log2(x_scale.abs())))
|
|
|
|
# Broadcast scales back to the original tensor shape
|
|
scales_repeated = repeat(x_scale, "s0 s1 -> (s0 t0) (s1 t1)", t0=t0, t1=t1)
|
|
else:
|
|
# Handle column-major tiling
|
|
x_tiled = rearrange(x, "(s1 t0) (s0 t1) -> s0 s1 t0 t1", s0=s0, s1=s1)
|
|
abs_max = reduce(x_tiled.abs(), "s0 s1 t0 t1 -> s0 s1", "max").clamp(1e-4)
|
|
x_scale = abs_max / fp8_amax
|
|
x_scale = torch.pow(2.0, torch.ceil(torch.log2(x_scale.abs())))
|
|
|
|
# Permute scale axes before repeating to match layout
|
|
scales_permuted = rearrange(x_scale, "s0 s1 -> s1 s0")
|
|
scales_repeated = repeat(scales_permuted, "s1 s0 -> (s1 t0) (s0 t1)", t0=t0, t1=t1)
|
|
|
|
elif ndim == 3:
|
|
s0, s1, s2 = scale_shape
|
|
t0, t1, t2 = tile_shape
|
|
if scale_major_mode == "K":
|
|
# Tile x and find the max absolute value in each tile
|
|
x_tiled = rearrange(
|
|
x, "(s0 t0) (s1 t1) (s2 t2) -> s0 s1 s2 t0 t1 t2", s0=s0, s1=s1, s2=s2
|
|
)
|
|
abs_max = reduce(x_tiled.abs(), "s0 s1 s2 t0 t1 t2 -> s0 s1 s2", "max").clamp(1e-4)
|
|
x_scale = abs_max / fp8_amax
|
|
x_scale = torch.pow(2.0, torch.ceil(torch.log2(x_scale.abs())))
|
|
|
|
# Broadcast scales back to the original tensor shape
|
|
scales_repeated = repeat(
|
|
x_scale, "s0 s1 s2 -> (s0 t0) (s1 t1) (s2 t2)", t0=t0, t1=t1, t2=t2
|
|
)
|
|
else:
|
|
# Handle layout where the last two axes are swapped
|
|
x_tiled = rearrange(
|
|
x, "(s0 t0) (s2 t1) (s1 t2) -> s0 s1 s2 t0 t1 t2", s0=s0, s1=s1, s2=s2
|
|
)
|
|
abs_max = reduce(x_tiled.abs(), "s0 s1 s2 t0 t1 t2 -> s0 s1 s2", "max").clamp(1e-4)
|
|
x_scale = abs_max / fp8_amax
|
|
x_scale = torch.pow(2.0, torch.ceil(torch.log2(x_scale.abs())))
|
|
# Permute scale axes before repeating to match layout
|
|
scales_permuted = rearrange(x_scale, "s0 s1 s2 -> s0 s2 s1")
|
|
scales_repeated = repeat(
|
|
scales_permuted,
|
|
"s0 s2 s1 -> (s0 t0) (s2 t1) (s1 t2)",
|
|
t0=t0,
|
|
t1=t1,
|
|
t2=t2,
|
|
)
|
|
# 3. Final Quantization
|
|
# Divide the original tensor by the broadcasted scales
|
|
x_fp32 = x / (scales_repeated + 1e-8)
|
|
|
|
# Convert the result to the target FP8 format
|
|
x_fp8 = x_fp32.to(torch.float8_e4m3fn)
|
|
|
|
return x_fp8, x_scale
|
|
|
|
|
|
def dequantize_fp8(x, x_scale, scale_major_mode):
|
|
from einops import rearrange
|
|
|
|
"""
|
|
Quantizes a 2D or 3D tensor to FP8.
|
|
|
|
Args:
|
|
x (torch.Tensor): The 2D or 3D input tensor.
|
|
scale_shape (tuple): The shape of the scale tensor.
|
|
tile_shape (tuple): The shape of the tiles.
|
|
scale_major_mode (str): The tiling order, "K" for row-major like,
|
|
or another value for column-major like.
|
|
|
|
Returns:
|
|
tuple: A tuple containing the quantized FP8 tensor and the
|
|
calculated float32 scales.
|
|
"""
|
|
# 1. Assertions and Initial Setup
|
|
ndim = x.ndim
|
|
assert ndim == len(x_scale.shape)
|
|
|
|
# 2. Tiling and Scale Calculation
|
|
if ndim == 2:
|
|
if scale_major_mode == "K":
|
|
s0, s1 = x_scale.shape
|
|
else:
|
|
s1, s0 = x_scale.shape
|
|
x = rearrange(x.to(torch.float32), "(s0 t0) (s1 t1) -> s0 s1 t0 t1", s0=s0, s1=s1)
|
|
if scale_major_mode == "K":
|
|
x_scale = rearrange(x_scale, "s0 s1 -> s0 s1 1 1")
|
|
else:
|
|
x_scale = rearrange(x_scale, "s0 s1 -> s1 s0 1 1")
|
|
out = rearrange(x * x_scale, "s0 s1 t0 t1 -> (s0 t0) (s1 t1)")
|
|
elif ndim == 3:
|
|
if scale_major_mode == "K":
|
|
s0, s1, s2 = x_scale.shape
|
|
else:
|
|
s0, s2, s1 = x_scale.shape
|
|
x = rearrange(
|
|
x.to(torch.float32),
|
|
"(s0 t0) (s1 t1) (s2 t2)-> s0 s1 s2 t0 t1 t2",
|
|
s0=s0,
|
|
s1=s1,
|
|
s2=s2,
|
|
)
|
|
if scale_major_mode == "K":
|
|
x_scale = rearrange(x_scale, "s0 s1 s2 -> s0 s1 s2 1 1 1")
|
|
else:
|
|
x_scale = rearrange(x_scale, "s0 s1 s2 -> s0 s2 s1 1 1 1")
|
|
out = rearrange(x * x_scale, "s0 s1 s2 t0 t1 t2 -> (s0 t0) (s1 t1) (s2 t2)")
|
|
|
|
return out
|
|
|
|
|
|
###########################################
|
|
########### Refernce generation ###########
|
|
###########################################
|
|
def compute_reference_grouped_gemm(
|
|
a_fp32: torch.Tensor, # (total_m, k)
|
|
b_fp32: torch.Tensor, # (batch_size, n, k)
|
|
m_indptr: torch.Tensor,
|
|
dtype_out: str, # (total_m, n)
|
|
):
|
|
"""Compute reference result using PyTorch operations"""
|
|
"""Compute reference result using original FP32 tensors"""
|
|
|
|
total_m, k = a_fp32.shape
|
|
batch_size, n, k2 = b_fp32.shape
|
|
assert k == k2
|
|
|
|
# Perform grouped GEMM computation directly on original FP32 data
|
|
results = []
|
|
|
|
for i in range(batch_size):
|
|
start_m = m_indptr[i].item()
|
|
end_m = m_indptr[i + 1].item()
|
|
|
|
# Extract group's portion of A
|
|
a_group = a_fp32[start_m:end_m, :] # [m_sizes[i], k]
|
|
b_group = b_fp32[i]
|
|
|
|
# Multiply with shared B matrix
|
|
result_group = torch.mm(a_group, b_group.T) # [m_sizes[i], n]
|
|
results.append(result_group)
|
|
|
|
result_fp32 = torch.cat(results, dim=0)
|
|
|
|
# Convert to output dtype
|
|
if dtype_out == "bfloat16":
|
|
result = result_fp32.to(torch.bfloat16)
|
|
elif dtype_out == "float16":
|
|
result = result_fp32.to(torch.float16)
|
|
else:
|
|
result = result_fp32
|
|
|
|
return result
|
|
|
|
|
|
###########################################
|
|
########### Test data generation ##########
|
|
###########################################
|
|
def generate_test_data(
|
|
m_sizes: list,
|
|
batch_size: int,
|
|
n: int,
|
|
k: int,
|
|
dtype_a: str,
|
|
dtype_b: str,
|
|
dtype_out: str,
|
|
scale_granularity_m: int,
|
|
scale_granularity_n: int,
|
|
scale_granularity_k: int,
|
|
scale_major_mode: str,
|
|
device: tvm.runtime.Device,
|
|
):
|
|
"""Generate test data for grouped GEMM operations"""
|
|
assert batch_size == len(m_sizes), (
|
|
f"batch_size ({batch_size}) must equal len(m_sizes) ({len(m_sizes)})"
|
|
)
|
|
|
|
# print(f"Device object: {device}")
|
|
torch_device = torch.device(f"cuda:{device.index}")
|
|
|
|
cum_m = [0] + list(np.cumsum(m_sizes))
|
|
total_m = cum_m[-1]
|
|
|
|
# Generate input matrices A and B (where we assert of form fp8) random data in fp32 first, then convert
|
|
assert dtype_a == "float8_e4m3fn"
|
|
a_fp32 = torch.randn(total_m, k, device=torch_device, dtype=torch.float32)
|
|
|
|
assert dtype_b == "float8_e4m3fn"
|
|
b_fp32 = torch.randn(batch_size, n, k, device=torch_device, dtype=torch.float32) / math.sqrt(k)
|
|
|
|
if scale_major_mode == "K": # K mode:
|
|
scale_a_shape = (total_m // scale_granularity_m, k // scale_granularity_k)
|
|
scale_b_shape = (batch_size, n // scale_granularity_n, k // scale_granularity_k)
|
|
|
|
else: # MN mode
|
|
scale_a_shape = (k // scale_granularity_k, total_m // scale_granularity_m)
|
|
scale_b_shape = (batch_size, k // scale_granularity_k, n // scale_granularity_n)
|
|
|
|
tile_a_shape = (scale_granularity_m, scale_granularity_k)
|
|
tile_b_shape = (1, scale_granularity_n, scale_granularity_k)
|
|
|
|
# quantize A, B
|
|
a_quantized, scale_a = quantize_fp8(a_fp32, scale_a_shape, tile_a_shape, scale_major_mode)
|
|
b_quantized, scale_b = quantize_fp8(b_fp32, scale_b_shape, tile_b_shape, scale_major_mode)
|
|
|
|
if dtype_a == "float8_e4m3fn":
|
|
a_tvm = tvm.runtime.tensor(
|
|
a_quantized.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device
|
|
)
|
|
else:
|
|
a_tvm = tvm.runtime.from_dlpack(a_quantized)
|
|
|
|
if dtype_b == "float8_e4m3fn":
|
|
b_tvm = tvm.runtime.tensor(
|
|
b_quantized.view(torch.uint8).cpu().numpy().view(fp8_dtype), device=device
|
|
)
|
|
else:
|
|
b_tvm = tvm.runtime.from_dlpack(b_quantized)
|
|
|
|
scale_a_tvm = tvm.runtime.from_dlpack(scale_a)
|
|
scale_b_tvm = tvm.runtime.from_dlpack(scale_b)
|
|
|
|
# Create m_indptr for grouped operation
|
|
m_indptr = torch.tensor(cum_m, device=torch_device, dtype=torch.int32)
|
|
m_indptr_tvm = tvm.runtime.tensor(m_indptr.cpu().numpy(), device)
|
|
|
|
return {
|
|
"a": a_tvm,
|
|
"b": b_tvm,
|
|
"torch_a": a_fp32,
|
|
"torch_b": b_fp32,
|
|
"scale_a": scale_a_tvm,
|
|
"scale_b": scale_b_tvm,
|
|
"m_indptr": m_indptr_tvm,
|
|
"m_sizes": m_sizes,
|
|
"n": n,
|
|
"k": k,
|
|
"total_m": total_m,
|
|
"torch_scale_a": scale_a,
|
|
"torch_scale_b": scale_b,
|
|
"torch_m_indptr": m_indptr,
|
|
}
|
|
|
|
|
|
###########################################
|
|
############### Test driver ###############
|
|
###########################################
|
|
@pytest.mark.skipif(not has_flashinfer(), reason="FlashInfer not available")
|
|
@pytest.mark.skipif(not has_cutlass(), reason="CUTLASS SM90+ not available")
|
|
@pytest.mark.parametrize(
|
|
"dtype_a,dtype_b,dtype_out",
|
|
[
|
|
("float8_e4m3fn", "float8_e4m3fn", "bfloat16"),
|
|
("float8_e4m3fn", "float8_e4m3fn", "float16"),
|
|
],
|
|
)
|
|
@pytest.mark.parametrize(
|
|
"scale_granularity_m,scale_granularity_n,scale_granularity_k",
|
|
[
|
|
(1, 128, 128), # Row-wise A, block-wise B
|
|
],
|
|
)
|
|
@pytest.mark.parametrize("scale_major_mode", ["K", "MN"])
|
|
@pytest.mark.parametrize("mma_sm", [1, 2])
|
|
@pytest.mark.parametrize(
|
|
"test_case",
|
|
[
|
|
{"batch_size": 4, "m_sizes": [128, 256, 192, 320], "n": 512, "k": 1024},
|
|
{"batch_size": 2, "m_sizes": [64, 128], "n": 256, "k": 512},
|
|
{"batch_size": 3, "m_sizes": [256, 256, 128], "n": 768, "k": 768},
|
|
{"batch_size": 2, "m_sizes": [20, 36], "n": 768, "k": 768},
|
|
],
|
|
)
|
|
def test_grouped_gemm_correctness(
|
|
dtype_a,
|
|
dtype_b,
|
|
dtype_out,
|
|
scale_granularity_m,
|
|
scale_granularity_n,
|
|
scale_granularity_k,
|
|
scale_major_mode,
|
|
mma_sm,
|
|
test_case,
|
|
):
|
|
"""Test correctness of GroupedGemm operations"""
|
|
target = tvm.target.Target.from_device(tvm.cuda(0))
|
|
|
|
# Generate the module
|
|
mod = relax.backend.cuda.flashinfer.gen_grouped_gemm_module(target=target)[0]
|
|
|
|
# Load the module
|
|
grouped_gemm_fn = mod["group_gemm_fp8_nt_groupwise"]
|
|
|
|
def run_and_check():
|
|
device = tvm.cuda(0)
|
|
test_data = generate_test_data(
|
|
batch_size=test_case["batch_size"],
|
|
m_sizes=test_case["m_sizes"],
|
|
n=test_case["n"],
|
|
k=test_case["k"],
|
|
dtype_a=dtype_a,
|
|
dtype_b=dtype_b,
|
|
dtype_out=dtype_out,
|
|
scale_granularity_m=scale_granularity_m,
|
|
scale_granularity_n=scale_granularity_n,
|
|
scale_granularity_k=scale_granularity_k,
|
|
scale_major_mode=scale_major_mode,
|
|
device=device,
|
|
)
|
|
|
|
output_shape = (test_data["total_m"], test_data["n"])
|
|
if dtype_out == "bfloat16":
|
|
output = tvm.runtime.empty(output_shape, dtype="bfloat16", device=device)
|
|
elif dtype_out == "float16":
|
|
output = tvm.runtime.empty(output_shape, dtype="float16", device=device)
|
|
else:
|
|
output = tvm.runtime.empty(output_shape, dtype="float32", device=device)
|
|
|
|
int_workspace = tvm.runtime.empty((DEFAULT_WORKSPACE_SIZE,), dtype="int32", device=device)
|
|
float_workspace = tvm.runtime.empty(
|
|
(DEFAULT_WORKSPACE_SIZE,), dtype="float32", device=device
|
|
)
|
|
|
|
grouped_gemm_fn(
|
|
int_workspace,
|
|
float_workspace,
|
|
test_data["a"],
|
|
test_data["b"],
|
|
test_data["scale_a"],
|
|
test_data["scale_b"],
|
|
output,
|
|
test_data["m_indptr"],
|
|
test_data["n"],
|
|
test_data["k"],
|
|
scale_granularity_m,
|
|
scale_granularity_n,
|
|
scale_granularity_k,
|
|
scale_major_mode,
|
|
mma_sm,
|
|
)
|
|
|
|
reference = compute_reference_grouped_gemm(
|
|
test_data["torch_a"],
|
|
test_data["torch_b"],
|
|
test_data["torch_m_indptr"],
|
|
dtype_out,
|
|
)
|
|
output_torch = torch.as_tensor(output, device=test_data["torch_a"].device)
|
|
assert output_torch.shape == reference.shape, (
|
|
f"Shape mismatch: got {output_torch.shape}, expected {reference.shape}"
|
|
)
|
|
diff = calc_diff(output_torch.cpu().double().numpy(), reference.cpu().double().numpy())
|
|
assert diff < 1e-3, f"diff too large {diff}"
|
|
|
|
tvm.testing.run_with_gpu_lock(run_and_check)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
tvm.testing.main()
|