chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,467 @@
|
||||
# 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()
|
||||
Reference in New Issue
Block a user