chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,152 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import torch
|
||||
|
||||
from vllm._custom_ops import scaled_fp4_quant
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
|
||||
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
|
||||
|
||||
kE2M1ToFloat = torch.tensor(
|
||||
[0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0], dtype=torch.float32
|
||||
)
|
||||
|
||||
|
||||
def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size):
|
||||
m_tiles = (m + 128 - 1) // 128
|
||||
f = block_size * 4
|
||||
k_tiles = (k + f - 1) // f
|
||||
tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4))
|
||||
tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
|
||||
out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size)
|
||||
return out[0:m, 0:k]
|
||||
|
||||
|
||||
def convert_swizzled_8x4_layout_to_linear(
|
||||
a_sf_swizzled: torch.Tensor, m, k, block_size
|
||||
):
|
||||
m_tiles = (m + 8 - 1) // 8
|
||||
f = block_size * 4
|
||||
k_tiles = (k + f - 1) // f
|
||||
tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 8, 4))
|
||||
tmp = torch.permute(tmp, (0, 1, 3, 2, 4))
|
||||
out = tmp.reshape(m_tiles * 8, k_tiles * f // block_size)
|
||||
return out[0:m, 0:k]
|
||||
|
||||
|
||||
def dequantize_nvfp4_to_dtype(
|
||||
tensor_fp4,
|
||||
tensor_sf,
|
||||
global_scale,
|
||||
dtype,
|
||||
device,
|
||||
block_size=16,
|
||||
is_sf_128x4_layout=True,
|
||||
):
|
||||
"""Dequantize the fp4 tensor back to high precision."""
|
||||
# Two fp4 values are packed into one uint8.
|
||||
assert tensor_fp4.dtype == torch.uint8
|
||||
m, packed_k = tensor_fp4.shape
|
||||
k = packed_k * 2
|
||||
tensor_f32 = break_fp4_bytes(tensor_fp4, dtype)
|
||||
tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)
|
||||
tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
|
||||
if is_sf_128x4_layout:
|
||||
tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)
|
||||
else:
|
||||
tensor_sf = convert_swizzled_8x4_layout_to_linear(tensor_sf, m, k, block_size)
|
||||
|
||||
tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale
|
||||
|
||||
# scale the tensor
|
||||
out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)
|
||||
return out.to(dtype=dtype)
|
||||
|
||||
|
||||
def break_fp4_bytes(a, dtype):
|
||||
assert a.dtype == torch.uint8
|
||||
m, n = a.shape
|
||||
|
||||
# Vectorized nibble processing
|
||||
a_flat = a.flatten()
|
||||
high = (a_flat & 0xF0) >> 4 # Upper nibbles
|
||||
low = a_flat & 0x0F # Lower nibbles
|
||||
|
||||
# Combine nibbles for batch processing
|
||||
combined = torch.stack((low, high), dim=1).flatten()
|
||||
|
||||
# Vectorized sign and magnitude extraction
|
||||
signs = (combined & 0x08).to(torch.bool) # Sign bits
|
||||
abs_vals = (combined & 0x07).to(torch.long) # Magnitude indices
|
||||
|
||||
# Device-aware lookup and sign application
|
||||
kE2M1 = kE2M1ToFloat.to(device=a.device)
|
||||
values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0)
|
||||
|
||||
# Reshape to final form
|
||||
return values.reshape(m, n * 2).to(dtype=dtype)
|
||||
|
||||
|
||||
def dequant_nvfp4_kv_cache(
|
||||
fp4_data: torch.Tensor,
|
||||
block_scale: torch.Tensor,
|
||||
global_scale: float,
|
||||
head_size: int,
|
||||
block_size: int,
|
||||
) -> torch.Tensor:
|
||||
"""Dequantize an NVFP4 KV cache with 4x4-swizzled block scales.
|
||||
|
||||
The input must be in HND layout so that the last two dims are
|
||||
(block_size, last_dim). For NHD caches, permute to HND first.
|
||||
|
||||
Args:
|
||||
fp4_data: [..., num_heads, block_size, head_size//2] uint8 packed fp4.
|
||||
block_scale: [..., num_heads, block_size, head_size//16] fp8 block
|
||||
scales (as uint8 or float8_e4m3fn).
|
||||
global_scale: checkpoint dequant scale (k_scale or v_scale).
|
||||
head_size: head dimension.
|
||||
block_size: page size.
|
||||
|
||||
Returns:
|
||||
[..., num_heads, block_size, head_size] float32.
|
||||
"""
|
||||
data_dim = head_size // 2
|
||||
scale_dim = head_size // 16
|
||||
|
||||
fp4_packed = fp4_data
|
||||
sf_swizzled = block_scale.view(torch.uint8)
|
||||
|
||||
# Unswizzle 4x4 block scales on (block_size, scale_dim) plane.
|
||||
# [..., T, S] → [..., T//4, 4, sg, 4] → permute → [..., T, S]
|
||||
batch_shape = sf_swizzled.shape[:-2]
|
||||
T, S = block_size, scale_dim
|
||||
sg = S // 4
|
||||
sf_reshape = sf_swizzled.reshape(*batch_shape, T // 4, 4, sg, 4)
|
||||
ndim = sf_reshape.ndim
|
||||
# Swap the last four dims: (..., T//4, 4, sg, 4) → (..., T//4, 4, 4, sg)
|
||||
perm = list(range(ndim - 4)) + [ndim - 4, ndim - 1, ndim - 3, ndim - 2]
|
||||
sf_linear = sf_reshape.permute(*perm).reshape(*batch_shape, T, S)
|
||||
sf_f32 = sf_linear.view(torch.float8_e4m3fn).to(torch.float32)
|
||||
|
||||
# Unpack fp4
|
||||
shape = fp4_packed.shape # [..., T, data_dim]
|
||||
fp4_flat = fp4_packed.reshape(-1, data_dim)
|
||||
fp4_vals = break_fp4_bytes(fp4_flat, torch.float32)
|
||||
fp4_vals = fp4_vals.reshape(*shape[:-1], head_size)
|
||||
|
||||
# Dequant: fp4_val * block_scale * global_scale per 16-element group
|
||||
return (
|
||||
fp4_vals.reshape(*shape[:-1], scale_dim, 16)
|
||||
* (sf_f32 * global_scale).unsqueeze(-1)
|
||||
).reshape(*shape[:-1], head_size)
|
||||
|
||||
|
||||
def get_nvfp4_global_scale(a: torch.Tensor):
|
||||
return (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.abs(a).max().to(torch.float32)
|
||||
|
||||
|
||||
def quant_nvfp4_tensor(a: torch.Tensor):
|
||||
a_global_scale = get_nvfp4_global_scale(a)
|
||||
a_quant, a_block_scale = scaled_fp4_quant(a, a_global_scale)
|
||||
return a_quant, a_block_scale, a_global_scale
|
||||
@@ -0,0 +1,128 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import DEFAULT_OPCHECK_TEST_UTILS, opcheck
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.model_executor.layers.quantization.utils.allspark_utils import (
|
||||
ALLSPARK_AMPERE_K_ALIGN,
|
||||
ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
|
||||
ALLSPARK_AMPERE_N_ALIGN,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import quantize_weights
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.scalar_type import scalar_types
|
||||
from vllm.utils.platform_utils import num_compute_units
|
||||
|
||||
|
||||
def is_gptq_allspark_supported(min_capability: int, max_capability: int) -> bool:
|
||||
if not current_platform.is_cuda():
|
||||
return False
|
||||
|
||||
capability = current_platform.get_device_capability()
|
||||
assert capability is not None
|
||||
|
||||
return (
|
||||
capability.to_int() >= min_capability and capability.to_int() <= max_capability
|
||||
)
|
||||
|
||||
|
||||
MNK_FACTORS = [
|
||||
(1, 4, 8),
|
||||
(13, 17, 67),
|
||||
(26, 37, 13),
|
||||
(48, 16, 24),
|
||||
(67, 13, 88),
|
||||
(257, 13, 11),
|
||||
(658, 13, 11),
|
||||
(1033, 9, 17),
|
||||
]
|
||||
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
HAS_ZP_OPTS = [False, True]
|
||||
|
||||
|
||||
def compute_max_diff(output, output_ref):
|
||||
return torch.mean(torch.abs(output - output_ref)) / torch.mean(
|
||||
torch.abs(output_ref)
|
||||
)
|
||||
|
||||
|
||||
def rand_data(shape, dtype=torch.float16):
|
||||
return torch.randn(shape, dtype=dtype, device="cuda")
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_gptq_allspark_supported(80, 89),
|
||||
reason="AllSpark Ampere kernel is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("mnk_factors", MNK_FACTORS)
|
||||
@pytest.mark.parametrize("group_size", [-1])
|
||||
@pytest.mark.parametrize("has_zp", HAS_ZP_OPTS)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
def test_gptq_allspark_gemm_ampere(mnk_factors, group_size, has_zp, dtype):
|
||||
m_factor, n_factor, k_factor = mnk_factors
|
||||
m = m_factor
|
||||
n = n_factor * ALLSPARK_AMPERE_N_ALIGN
|
||||
k = k_factor * ALLSPARK_AMPERE_K_ALIGN
|
||||
|
||||
input = rand_data((m, k), dtype=dtype)
|
||||
weight = rand_data((k, n), dtype=dtype)
|
||||
|
||||
# Quantize (and apply act_order if provided)
|
||||
w_ref, qw, s, zp = quantize_weights(
|
||||
weight, scalar_types.uint8b128, group_size, has_zp
|
||||
)
|
||||
|
||||
qw = qw.to(torch.uint8)
|
||||
if has_zp:
|
||||
zp = zp.to(dtype)
|
||||
properties = torch.cuda.get_device_properties(qw.device.index)
|
||||
sm_count = num_compute_units(qw.device.index)
|
||||
sm_version = properties.major * 10 + properties.minor
|
||||
|
||||
n_32align = (n + 32 - 1) // 32 * 32
|
||||
|
||||
qw_reorder, s_reorder, zp_reorder = ops.allspark_repack_weight(qw, s, zp, has_zp)
|
||||
opcheck(
|
||||
torch.ops._C.rearrange_kn_weight_as_n32k16_order,
|
||||
(qw, s, zp, has_zp, qw_reorder, s_reorder, zp_reorder, k, n, n_32align),
|
||||
)
|
||||
|
||||
opcheck(
|
||||
torch.ops._C.allspark_w8a16_gemm,
|
||||
(
|
||||
input,
|
||||
qw_reorder,
|
||||
s_reorder,
|
||||
zp_reorder,
|
||||
n,
|
||||
group_size,
|
||||
sm_count,
|
||||
sm_version,
|
||||
ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
|
||||
has_zp,
|
||||
True,
|
||||
),
|
||||
test_utils=DEFAULT_OPCHECK_TEST_UTILS,
|
||||
)
|
||||
output = ops.allspark_w8a16_gemm(
|
||||
input,
|
||||
qw_reorder,
|
||||
s_reorder,
|
||||
zp_reorder,
|
||||
n,
|
||||
group_size,
|
||||
sm_count,
|
||||
sm_version,
|
||||
ALLSPARK_AMPERE_M_CUBLAS_THRESHOLD,
|
||||
has_zp,
|
||||
True,
|
||||
)
|
||||
|
||||
output_ref = torch.matmul(input, w_ref)
|
||||
torch.accelerator.synchronize()
|
||||
max_diff = compute_max_diff(output, output_ref)
|
||||
|
||||
assert max_diff < 0.04
|
||||
@@ -0,0 +1,29 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm import _custom_ops as ops # noqa: F401
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not hasattr(torch.ops._C, "awq_dequantize"),
|
||||
reason="AWQ is not supported on this GPU type.",
|
||||
)
|
||||
def test_awq_dequantize_opcheck(monkeypatch: pytest.MonkeyPatch):
|
||||
with monkeypatch.context() as m:
|
||||
m.setenv("VLLM_USE_TRITON_AWQ", "0")
|
||||
qweight = torch.randint(
|
||||
-2000000000, 2000000000, (8192, 256), device="cuda", dtype=torch.int32
|
||||
)
|
||||
scales = torch.rand((64, 2048), device="cuda", dtype=torch.float16)
|
||||
zeros = torch.empty((64, 256), device="cuda", dtype=torch.int32)
|
||||
split_k_iters = 0
|
||||
thx = 0
|
||||
thy = 0
|
||||
opcheck(
|
||||
torch.ops._C.awq_dequantize,
|
||||
(qweight, scales, zeros, split_k_iters, thx, thy),
|
||||
)
|
||||
@@ -0,0 +1,177 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the AWQ Triton kernel.
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_awq_triton.py`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.quantization.awq_triton import (
|
||||
AWQ_TRITON_SUPPORTED_GROUP_SIZES,
|
||||
awq_dequantize_triton,
|
||||
awq_gemm_triton,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
pytestmark = pytest.mark.skipif(
|
||||
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
|
||||
reason="AWQ Triton kernels require CUDA/ROCm or XPU.",
|
||||
)
|
||||
|
||||
device = current_platform.device_type
|
||||
|
||||
|
||||
def reverse_awq_order(t: torch.Tensor):
|
||||
bits = 4
|
||||
AWQ_REVERSE_ORDER = [0, 4, 1, 5, 2, 6, 3, 7]
|
||||
reverse_order_tensor = torch.arange(
|
||||
t.shape[-1],
|
||||
dtype=torch.int32,
|
||||
device=t.device,
|
||||
)
|
||||
reverse_order_tensor = reverse_order_tensor.view(-1, 32 // bits)
|
||||
reverse_order_tensor = reverse_order_tensor[:, AWQ_REVERSE_ORDER]
|
||||
reverse_order_tensor = reverse_order_tensor.view(-1)
|
||||
|
||||
t = t[:, reverse_order_tensor] & 0xF
|
||||
return t
|
||||
|
||||
|
||||
# qweights - [R , C // 8], int32
|
||||
# scales - [R // G, C ], float16
|
||||
# zeros - [R // G, C // 8], int32
|
||||
def awq_dequantize_torch(
|
||||
qweight: torch.Tensor, scales: torch.Tensor, qzeros: torch.Tensor, group_size: int
|
||||
) -> torch.Tensor:
|
||||
if group_size == -1:
|
||||
group_size = qweight.shape[0]
|
||||
|
||||
bits = 4
|
||||
shifts = torch.arange(0, 32, bits, device=qzeros.device)
|
||||
|
||||
iweights = torch.bitwise_right_shift(qweight[:, :, None], shifts[None, None, :]).to(
|
||||
torch.int8
|
||||
)
|
||||
|
||||
iweights = iweights.view(iweights.shape[0], -1)
|
||||
|
||||
zeros = torch.bitwise_right_shift(qzeros[:, :, None], shifts[None, None, :]).to(
|
||||
torch.int8
|
||||
)
|
||||
zeros = zeros.view(qzeros.shape[0], -1)
|
||||
zeros = reverse_awq_order(zeros)
|
||||
|
||||
iweights = reverse_awq_order(iweights)
|
||||
|
||||
iweights = torch.bitwise_and(iweights, (2**bits) - 1)
|
||||
zeros = torch.bitwise_and(zeros, (2**bits) - 1)
|
||||
|
||||
scales = scales.repeat_interleave(group_size, dim=0)
|
||||
zeros = zeros.repeat_interleave(group_size, dim=0)
|
||||
return (iweights - zeros) * scales
|
||||
|
||||
|
||||
# qweights - [R , C // 8], int32
|
||||
# scales - [R // G, C ], float16
|
||||
# zeros - [R // G, C // 8], int32
|
||||
@pytest.mark.parametrize("qweight_rows", [3584, 18944, 128, 256, 512, 1024])
|
||||
@pytest.mark.parametrize("qweight_cols", [448, 576, 4736, 16, 32, 64, 128])
|
||||
@pytest.mark.parametrize("group_size", AWQ_TRITON_SUPPORTED_GROUP_SIZES)
|
||||
def test_dequantize(qweight_rows, qweight_cols, group_size):
|
||||
if group_size == -1:
|
||||
group_size = qweight_rows
|
||||
|
||||
qweight_dtype = torch.int32
|
||||
scales_rows = qweight_rows // group_size
|
||||
scales_cols = qweight_cols * 8
|
||||
scales_dtype = torch.float16
|
||||
zeros_rows = scales_rows
|
||||
zeros_cols = qweight_cols
|
||||
zeros_dtype = torch.int32
|
||||
|
||||
set_random_seed(0)
|
||||
|
||||
qweight = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(qweight_rows, qweight_cols),
|
||||
dtype=qweight_dtype,
|
||||
device=device,
|
||||
)
|
||||
scales = torch.rand(scales_rows, scales_cols, dtype=scales_dtype, device=device)
|
||||
zeros = torch.randint(
|
||||
0,
|
||||
torch.iinfo(torch.int32).max,
|
||||
(zeros_rows, zeros_cols),
|
||||
dtype=zeros_dtype,
|
||||
device=device,
|
||||
)
|
||||
|
||||
iweights_triton = awq_dequantize_triton(qweight, scales, zeros)
|
||||
|
||||
assert not torch.any(torch.isinf(iweights_triton)) and not torch.any(
|
||||
torch.isnan(iweights_triton)
|
||||
)
|
||||
|
||||
iweights_torch = awq_dequantize_torch(qweight, scales, zeros, group_size)
|
||||
|
||||
torch.testing.assert_close(iweights_triton, iweights_torch)
|
||||
|
||||
|
||||
# input - [N, K]
|
||||
# qweight - [K, M // 8]
|
||||
# qzeros - [K // G, M // 8]
|
||||
# scales - [K // G, M]
|
||||
@pytest.mark.parametrize("N", [1, 2, 4, 8, 14, 17, 23, 32])
|
||||
@pytest.mark.parametrize("K", [128])
|
||||
@pytest.mark.parametrize("M", [16, 24, 32])
|
||||
@pytest.mark.parametrize("group_size", AWQ_TRITON_SUPPORTED_GROUP_SIZES)
|
||||
@pytest.mark.parametrize("splitK", [1, 8])
|
||||
def test_gemm(N, K, M, splitK, group_size):
|
||||
if group_size == -1:
|
||||
group_size = K
|
||||
|
||||
split_k_iters = splitK
|
||||
|
||||
input_rows = N
|
||||
input_cols = K
|
||||
input_dtype = torch.float32
|
||||
qweight_rows = input_cols
|
||||
qweight_cols = M // 8
|
||||
scales_rows = qweight_rows // group_size
|
||||
scales_cols = M
|
||||
scales_dtype = torch.float32
|
||||
qzeros_rows = scales_rows
|
||||
qzeros_cols = qweight_cols
|
||||
|
||||
set_random_seed(0)
|
||||
|
||||
input = torch.rand((input_rows, input_cols), dtype=input_dtype, device=device)
|
||||
qweight = torch.randint(
|
||||
0, torch.iinfo(torch.int32).max, (qweight_rows, qweight_cols), device=device
|
||||
)
|
||||
qzeros = torch.randint(
|
||||
0, torch.iinfo(torch.int32).max, (qzeros_rows, qzeros_cols), device=device
|
||||
)
|
||||
scales = torch.rand((scales_rows, scales_cols), dtype=scales_dtype, device=device)
|
||||
|
||||
output_triton = awq_gemm_triton(input, qweight, scales, qzeros, split_k_iters)
|
||||
|
||||
assert not torch.any(torch.isinf(output_triton)) and not torch.any(
|
||||
torch.isnan(output_triton)
|
||||
)
|
||||
|
||||
dequantized_weights = awq_dequantize_triton(qweight, scales, qzeros)
|
||||
|
||||
output_torch = torch.matmul(input, dequantized_weights)
|
||||
|
||||
assert not torch.any(torch.isinf(output_torch)) and not torch.any(
|
||||
torch.isnan(output_torch)
|
||||
)
|
||||
|
||||
torch.testing.assert_close(
|
||||
output_triton.cpu(), output_torch.cpu(), atol=1e-1, rtol=1e-1
|
||||
)
|
||||
@@ -0,0 +1,295 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Adapted from https://github.com/sgl-project/sglang/pull/2575
|
||||
import itertools
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.quant_utils import (
|
||||
native_per_token_group_quant_fp8,
|
||||
native_w8a8_block_matmul,
|
||||
)
|
||||
from tests.kernels.utils import fp8_ulp_distance
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.model_executor.kernels.linear.scaled_mm.cutlass import cutlass_scaled_mm
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
per_token_group_quant_fp8,
|
||||
w8a8_triton_block_scaled_mm,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.deep_gemm import (
|
||||
fp8_gemm_nt,
|
||||
get_tma_aligned_size,
|
||||
per_block_cast_to_fp8,
|
||||
should_use_deepgemm_for_fp8_linear,
|
||||
)
|
||||
from vllm.utils.flashinfer import (
|
||||
flashinfer_fp8_blockscale_gemm,
|
||||
has_flashinfer_fp8_blockscale_gemm,
|
||||
)
|
||||
from vllm.utils.import_utils import has_deep_gemm
|
||||
|
||||
if current_platform.get_device_capability() < (9, 0):
|
||||
pytest.skip("FP8 Triton requires CUDA 9.0 or higher", allow_module_level=True)
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
|
||||
# Test configurations
|
||||
DTYPES = [torch.bfloat16] # [torch.half, torch.bfloat16, torch.float32]
|
||||
# Quantization test configs
|
||||
NUM_TOKENS = [7, 2050]
|
||||
D = [512, 4096, 5120, 13824]
|
||||
GROUP_SIZE = [64, 128, 512]
|
||||
COLUMN_MAJOR_SCALES = [True, False]
|
||||
TMA_ALIGNED_SCALES = [True, False]
|
||||
# Matmul test configs
|
||||
M = [1, 7, 8, 83, 4096]
|
||||
N = [128, 512, 576, 7168, 13824]
|
||||
K = [256, 3884, 4096, 13824, 16384]
|
||||
# Deepseek-V3's intermediate size 18432, so N is 18432*2/8=4608 at TP8
|
||||
# and its hidden size is 7168.
|
||||
BLOCK_SIZE = [[128, 128]]
|
||||
OUT_DTYPES = [torch.bfloat16] # [torch.float32, torch.half, torch.bfloat16]
|
||||
SEEDS = [0]
|
||||
|
||||
# Skip all tests if CUDA is not available
|
||||
pytest.importorskip("torch.cuda")
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def setup_cuda():
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_fp8_fnuz(),
|
||||
reason="This platform supports e4m3fnuz, not e4m3fn.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"num_tokens,d,dtype,group_size,column_major_scales,tma_aligned_scales,seed",
|
||||
itertools.product(
|
||||
NUM_TOKENS,
|
||||
D,
|
||||
DTYPES,
|
||||
GROUP_SIZE,
|
||||
COLUMN_MAJOR_SCALES,
|
||||
TMA_ALIGNED_SCALES,
|
||||
SEEDS,
|
||||
),
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_per_token_group_quant_fp8(
|
||||
num_tokens, d, dtype, group_size, column_major_scales, tma_aligned_scales, seed
|
||||
):
|
||||
torch.manual_seed(seed)
|
||||
x = torch.rand(num_tokens, d, dtype=dtype)
|
||||
|
||||
ref_out, ref_scale = native_per_token_group_quant_fp8(x, group_size)
|
||||
out, scale = per_token_group_quant_fp8(
|
||||
x,
|
||||
group_size,
|
||||
column_major_scales=column_major_scales,
|
||||
tma_aligned_scales=tma_aligned_scales,
|
||||
)
|
||||
|
||||
if current_platform.is_rocm():
|
||||
# On gfx950 the Triton and PyTorch FP8 kernels can round in opposite
|
||||
# directions when an element lands at the midpoint between two adjacent
|
||||
# e4m3fn values (1-ULP tie-breaking). Verify: (1) no element is more
|
||||
# than 1 FP8 ULP away, and (2) fewer than 0.05% of elements have any
|
||||
# mismatch. Observed worst case across all parameter combos: 0.049%,
|
||||
# max ULP = 1.
|
||||
ulp = fp8_ulp_distance(out, ref_out)
|
||||
assert (ulp <= 1).all(), (
|
||||
f"FP8 mismatch > 1 ULP: {int((ulp > 1).sum())} elements"
|
||||
)
|
||||
assert float((ulp > 0).float().mean()) < 5e-4, (
|
||||
f"Too many 1-ULP mismatches: {int((ulp > 0).sum())}/{ulp.numel()}"
|
||||
)
|
||||
else:
|
||||
assert torch.allclose(
|
||||
out.to(torch.float32), ref_out.to(torch.float32), rtol=0.15
|
||||
)
|
||||
assert torch.allclose(scale, ref_scale)
|
||||
|
||||
if column_major_scales:
|
||||
assert scale.stride()[-2] == 1
|
||||
if tma_aligned_scales:
|
||||
assert scale.stride()[-1] == get_tma_aligned_size(num_tokens, 4)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"M,N,K,block_size,out_dtype,seed",
|
||||
itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS),
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_w8a8_block_fp8_matmul(M, N, K, block_size, out_dtype, seed):
|
||||
torch.manual_seed(seed)
|
||||
factor_for_scale = 1e-2
|
||||
fp8_info = torch.finfo(current_platform.fp8_dtype())
|
||||
fp8_max, fp8_min = fp8_info.max, fp8_info.min
|
||||
|
||||
A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
|
||||
A_fp8 = A_fp32.clamp(min=fp8_min, max=fp8_max).to(current_platform.fp8_dtype())
|
||||
|
||||
B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
|
||||
B_fp8 = B_fp32.clamp(min=fp8_min, max=fp8_max).to(current_platform.fp8_dtype())
|
||||
|
||||
block_n, block_k = block_size[0], block_size[1]
|
||||
n_tiles = (N + block_n - 1) // block_n
|
||||
k_tiles = (K + block_k - 1) // block_k
|
||||
|
||||
As = torch.rand(M, k_tiles, dtype=torch.float32) * factor_for_scale
|
||||
Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale
|
||||
|
||||
ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
|
||||
out = w8a8_triton_block_scaled_mm(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
|
||||
|
||||
rel_diff = torch.mean(
|
||||
torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))
|
||||
) / torch.mean(torch.abs(ref_out.to(torch.float32)))
|
||||
assert rel_diff < 0.001
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(), reason="CUTLASS only supported on CUDA platform."
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_w8a8_block_fp8_cutlass_matmul():
|
||||
# Test simple case where weight.shape % 128 != 0,
|
||||
# like in DSV3 kv_a_proj_with_mqa
|
||||
M = 32
|
||||
N = 576
|
||||
K = 7168
|
||||
block_size = [128, 128]
|
||||
out_dtype = torch.bfloat16
|
||||
seed = 0
|
||||
|
||||
torch.manual_seed(seed)
|
||||
factor_for_scale = 1e-2
|
||||
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) - 0.5) * 2 * fp8_max
|
||||
|
||||
B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
|
||||
B_fp8 = B_fp32.clamp(min=fp8_min, max=fp8_max).to(torch.float8_e4m3fn)
|
||||
|
||||
block_n, block_k = block_size[0], block_size[1]
|
||||
n_tiles = (N + block_n - 1) // block_n
|
||||
k_tiles = (K + block_k - 1) // block_k
|
||||
|
||||
Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale
|
||||
|
||||
A_fp8, As = per_token_group_quant_fp8(
|
||||
A_fp32, block_size[1], column_major_scales=False
|
||||
)
|
||||
# CUTLASS uses column-major format for scales
|
||||
A_fp8_cutlass, As_cutlass = per_token_group_quant_fp8(
|
||||
A_fp32, block_size[1], column_major_scales=True
|
||||
)
|
||||
|
||||
ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
|
||||
out = cutlass_scaled_mm(A_fp8_cutlass, B_fp8, As_cutlass, Bs, block_size, out_dtype)
|
||||
|
||||
rel_diff = torch.mean(
|
||||
torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))
|
||||
) / torch.mean(torch.abs(ref_out.to(torch.float32)))
|
||||
assert rel_diff < 0.001
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_fp8_fnuz(),
|
||||
reason="This platform supports e4m3fnuz, not e4m3fn.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"M,N,K,block_size,out_dtype,seed",
|
||||
itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS),
|
||||
)
|
||||
@pytest.mark.skipif(not has_deep_gemm(), reason="DeepGemm kernels not available.")
|
||||
@torch.inference_mode()
|
||||
def test_w8a8_block_fp8_deep_gemm_matmul(M, N, K, block_size, out_dtype, seed):
|
||||
torch.manual_seed(seed)
|
||||
fp8_info = torch.finfo(torch.float8_e4m3fn)
|
||||
fp8_max = fp8_info.max
|
||||
|
||||
A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
|
||||
B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * fp8_max
|
||||
|
||||
# only aligned sizes are supported by deepgemm
|
||||
if not should_use_deepgemm_for_fp8_linear(
|
||||
output_dtype=out_dtype, weight_shape=B_fp32.shape, supports_deep_gemm=True
|
||||
):
|
||||
pytest.skip(f"Skipping test; invalid size {M}, {N}, {K}")
|
||||
|
||||
A_fp8, As_fp8 = per_token_group_quant_fp8(
|
||||
A_fp32, block_size[1], column_major_scales=True, tma_aligned_scales=True
|
||||
)
|
||||
B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_fp32, block_size=block_size)
|
||||
|
||||
As = As_fp8.to(torch.float32)
|
||||
Bs = Bs_fp8.to(torch.float32)
|
||||
|
||||
ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
|
||||
|
||||
out = torch.zeros((M, N), device="cuda", dtype=out_dtype)
|
||||
|
||||
assert As_fp8.shape == (M, (K + 127) // 128), (
|
||||
f"{As_fp8.shape} != {(M, (K + 127) // 128)}"
|
||||
)
|
||||
|
||||
fp8_gemm_nt((A_fp8, As_fp8), (B_fp8, Bs_fp8), out)
|
||||
|
||||
rel_diff = torch.mean(
|
||||
torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))
|
||||
) / torch.mean(torch.abs(ref_out.to(torch.float32)))
|
||||
assert rel_diff < 0.001
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_fp8_fnuz(),
|
||||
reason="This platform supports e4m3fnuz, not e4m3fn.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"M,N,K,block_size,out_dtype,seed",
|
||||
itertools.product(M, N, K, BLOCK_SIZE, OUT_DTYPES, SEEDS),
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_w8a8_block_fp8_flashinfer_matmul(M, N, K, block_size, out_dtype, seed):
|
||||
if not has_flashinfer_fp8_blockscale_gemm():
|
||||
pytest.skip(
|
||||
"FlashInfer block GEMM not available (requires SM90+ and FlashInfer)"
|
||||
)
|
||||
# only aligned sizes
|
||||
if K % 128 != 0 or N % 64 != 0:
|
||||
pytest.skip(f"Skipping test; invalid size {M}, {N}, {K}")
|
||||
|
||||
torch.manual_seed(seed)
|
||||
fp8_info = torch.finfo(torch.float8_e4m3fn)
|
||||
fp8_max = fp8_info.max
|
||||
|
||||
A_bf16 = (torch.rand(M, K, dtype=torch.bfloat16) - 0.5) * 2 * fp8_max
|
||||
B_bf16 = (torch.rand(N, K, dtype=torch.bfloat16) - 0.5) * 2 * fp8_max
|
||||
|
||||
A_fp8, As_fp8 = per_token_group_quant_fp8(A_bf16, block_size[1], use_ue8m0=False)
|
||||
B_fp8, Bs_fp8 = per_block_cast_to_fp8(B_bf16, block_size, use_ue8m0=False)
|
||||
|
||||
As = As_fp8.to(torch.float32)
|
||||
Bs = Bs_fp8.to(torch.float32)
|
||||
|
||||
ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
|
||||
|
||||
out = flashinfer_fp8_blockscale_gemm(
|
||||
input=A_bf16,
|
||||
weight=B_fp8,
|
||||
input_scale=None,
|
||||
weight_scale=Bs,
|
||||
out_dtype=out_dtype,
|
||||
)
|
||||
|
||||
rel_diff = torch.mean(
|
||||
torch.abs(out.to(torch.bfloat16) - ref_out.to(torch.bfloat16))
|
||||
) / torch.mean(torch.abs(ref_out.to(torch.bfloat16)))
|
||||
assert rel_diff < 0.001
|
||||
@@ -0,0 +1,67 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Adapted from https://github.com/sgl-project/sglang/blob/main/test/srt/test_block_int8.py
|
||||
import itertools
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.quant_utils import native_w8a8_block_matmul
|
||||
from vllm.config import VllmConfig
|
||||
from vllm.model_executor.layers.quantization.utils.int8_utils import (
|
||||
w8a8_block_int8_matmul,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if current_platform.get_device_capability() < (7, 0):
|
||||
pytest.skip("INT8 Triton requires CUDA 7.0 or higher", allow_module_level=True)
|
||||
|
||||
vllm_config = VllmConfig()
|
||||
|
||||
DTYPES = [torch.half, torch.bfloat16]
|
||||
M = [1, 33, 64, 222]
|
||||
N = [128, 1024]
|
||||
K = [256, 4096]
|
||||
# BLOCK_SIZE = [[64, 64], [64, 128], [128, 64], [128, 128]]
|
||||
BLOCK_SIZE = [[128, 128]]
|
||||
SEEDS = [0]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True, scope="module")
|
||||
def setup_cuda():
|
||||
"""Sets the default CUDA device for all tests in this module."""
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"M,N,K,block_size,out_dtype,seed",
|
||||
itertools.product(M, N, K, BLOCK_SIZE, DTYPES, SEEDS),
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_w8a8_block_int8_matmul(M, N, K, block_size, out_dtype, seed):
|
||||
torch.manual_seed(seed)
|
||||
factor_for_scale = 1e-2
|
||||
int8_info = torch.iinfo(torch.int8)
|
||||
int8_max, int8_min = int8_info.max, int8_info.min
|
||||
|
||||
A_fp32 = (torch.rand(M, K, dtype=torch.float32) - 0.5) * 2 * int8_max
|
||||
A_fp8 = A_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn)
|
||||
|
||||
B_fp32 = (torch.rand(N, K, dtype=torch.float32) - 0.5) * 2 * int8_max
|
||||
B_fp8 = B_fp32.clamp(min=int8_min, max=int8_max).to(torch.float8_e4m3fn)
|
||||
|
||||
block_n, block_k = block_size[0], block_size[1]
|
||||
n_tiles = (N + block_n - 1) // block_n
|
||||
k_tiles = (K + block_k - 1) // block_k
|
||||
|
||||
As = torch.rand(M, k_tiles, dtype=torch.float32) * factor_for_scale
|
||||
Bs = torch.rand(n_tiles, k_tiles, dtype=torch.float32) * factor_for_scale
|
||||
|
||||
ref_out = native_w8a8_block_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
|
||||
out = w8a8_block_int8_matmul(A_fp8, B_fp8, As, Bs, block_size, out_dtype)
|
||||
|
||||
rel_diff = torch.mean(
|
||||
torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))
|
||||
) / torch.mean(torch.abs(ref_out.to(torch.float32)))
|
||||
assert rel_diff < 0.001
|
||||
@@ -0,0 +1,162 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for CPU FP8 W8A16 block-scaled GEMM kernel (fp8_scaled_mm_cpu).
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_cpu_fp8_scaled_mm.py -v`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_cpu():
|
||||
pytest.skip("skipping CPU-only tests", allow_module_level=True)
|
||||
|
||||
if not ops._supports_cpu_fp8_w8a16:
|
||||
pytest.skip("fp8_scaled_mm_cpu op not available", allow_module_level=True)
|
||||
|
||||
BLOCK_SIZE = [128, 128]
|
||||
|
||||
|
||||
def cdiv(a: int, b: int) -> int:
|
||||
return -(a // -b)
|
||||
|
||||
|
||||
def quantize_weight_block_fp8(
|
||||
weight: torch.Tensor,
|
||||
block_size: list[int],
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
"""Quantize weight [N, K] to FP8 with block scales.
|
||||
|
||||
Returns:
|
||||
fp8_weight: [N, K] float8_e4m3fn
|
||||
scales: [n_tiles, k_tiles] float32
|
||||
"""
|
||||
N, K = weight.shape
|
||||
block_n, block_k = block_size
|
||||
fp8_max = torch.finfo(torch.float8_e4m3fn).max
|
||||
|
||||
n_tiles = cdiv(N, block_n)
|
||||
k_tiles = cdiv(K, block_k)
|
||||
|
||||
# Pad for even blocking
|
||||
pad_N = (block_n - (N % block_n)) % block_n
|
||||
pad_K = (block_k - (K % block_k)) % block_k
|
||||
if pad_N > 0 or pad_K > 0:
|
||||
weight = torch.nn.functional.pad(weight, (0, pad_K, 0, pad_N))
|
||||
|
||||
# Reshape into blocks
|
||||
w_blocks = weight.view(n_tiles, block_n, k_tiles, block_k)
|
||||
w_blocks = w_blocks.permute(0, 2, 1, 3).contiguous()
|
||||
|
||||
# Per-block scale
|
||||
abs_max = w_blocks.abs().amax(dim=(-2, -1), keepdim=True)
|
||||
scales = abs_max / fp8_max
|
||||
scales = torch.where(scales == 0, torch.ones_like(scales), scales)
|
||||
|
||||
# Quantize
|
||||
q_fp8 = (w_blocks / scales).clamp(-fp8_max, fp8_max).to(torch.float8_e4m3fn)
|
||||
|
||||
# Reshape back
|
||||
fp8_weight = (
|
||||
q_fp8.permute(0, 2, 1, 3)
|
||||
.contiguous()
|
||||
.view(N + pad_N, K + pad_K)[:N, :K]
|
||||
.contiguous()
|
||||
)
|
||||
|
||||
scales = scales.view(n_tiles, k_tiles)
|
||||
return fp8_weight, scales
|
||||
|
||||
|
||||
def dequant_weight_block_fp8(
|
||||
fp8_weight: torch.Tensor,
|
||||
scales: torch.Tensor,
|
||||
block_size: list[int],
|
||||
out_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Dequantize FP8 weight back to float for reference computation."""
|
||||
N, K = fp8_weight.shape
|
||||
block_n, block_k = block_size
|
||||
n_tiles, k_tiles = scales.shape
|
||||
|
||||
pad_N = (block_n - (N % block_n)) % block_n
|
||||
pad_K = (block_k - (K % block_k)) % block_k
|
||||
if pad_N > 0 or pad_K > 0:
|
||||
fp8_padded = torch.nn.functional.pad(fp8_weight.float(), (0, pad_K, 0, pad_N))
|
||||
else:
|
||||
fp8_padded = fp8_weight.float()
|
||||
|
||||
w_blocks = fp8_padded.view(n_tiles, block_n, k_tiles, block_k)
|
||||
w_blocks = w_blocks.permute(0, 2, 1, 3).contiguous()
|
||||
dq = w_blocks * scales.view(n_tiles, k_tiles, 1, 1)
|
||||
dq = dq.permute(0, 2, 1, 3).contiguous().view(N + pad_N, K + pad_K)
|
||||
return dq[:N, :K].to(out_dtype)
|
||||
|
||||
|
||||
def ref_fp8_block_scaled_mm(
|
||||
x: torch.Tensor,
|
||||
fp8_weight: torch.Tensor,
|
||||
scales: torch.Tensor,
|
||||
block_size: list[int],
|
||||
bias: torch.Tensor | None,
|
||||
out_dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
"""Reference: dequant FP8→float32, matmul in float32, cast to out_dtype."""
|
||||
w_dq = dequant_weight_block_fp8(fp8_weight, scales, block_size, torch.float32)
|
||||
out = torch.mm(x.float(), w_dq.t())
|
||||
if bias is not None:
|
||||
out = out + bias.float()
|
||||
return out.to(out_dtype)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test parameters
|
||||
# ---------------------------------------------------------------------------
|
||||
M_SIZES = [1, 4, 16, 64, 128]
|
||||
# (N, K) — weight shape is [N, K], output has N columns.
|
||||
NK_SIZES = [
|
||||
(128, 256),
|
||||
(256, 512),
|
||||
(512, 1024),
|
||||
(1024, 2048),
|
||||
(5120, 5120),
|
||||
(17408, 5120),
|
||||
(5120, 17408),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("M", M_SIZES)
|
||||
@pytest.mark.parametrize("N,K", NK_SIZES)
|
||||
@pytest.mark.parametrize("use_bias", [False, True])
|
||||
def test_cpu_fp8_scaled_mm(M: int, N: int, K: int, use_bias: bool):
|
||||
"""fp8_scaled_mm_cpu correctness against float reference."""
|
||||
torch.manual_seed(42)
|
||||
out_dtype = torch.bfloat16
|
||||
block_size = BLOCK_SIZE
|
||||
|
||||
x = torch.randn(M, K, dtype=out_dtype) / (K**0.5)
|
||||
w_f32 = torch.randn(N, K, dtype=torch.float32) / (K**0.5)
|
||||
fp8_weight, scales = quantize_weight_block_fp8(w_f32, block_size)
|
||||
|
||||
bias = torch.randn(N, dtype=torch.float32) * 0.1 if use_bias else None
|
||||
|
||||
ref_out = ref_fp8_block_scaled_mm(
|
||||
x, fp8_weight, scales, block_size, bias, out_dtype
|
||||
)
|
||||
|
||||
packed_weight = torch.ops._C.convert_weight_packed(fp8_weight)
|
||||
kernel_out = ops.fp8_scaled_mm_cpu(
|
||||
x,
|
||||
packed_weight,
|
||||
scales,
|
||||
block_size,
|
||||
bias,
|
||||
out_dtype,
|
||||
True,
|
||||
)
|
||||
|
||||
assert kernel_out.dtype == out_dtype
|
||||
torch.testing.assert_close(kernel_out, ref_out, rtol=0.02, atol=0.01)
|
||||
@@ -0,0 +1,759 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for cutlass kernels
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_cutlass_scaled_mm.py`.
|
||||
"""
|
||||
|
||||
import random
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import baseline_scaled_mm, opcheck, to_fp8, to_int8
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.math_utils import cdiv
|
||||
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip("These tests use CUTLASS which requires CUDA", allow_module_level=True)
|
||||
|
||||
MNK_FACTORS = [
|
||||
(1, 256, 128),
|
||||
(1, 16384, 1024),
|
||||
(1, 24576, 496),
|
||||
(16, 256, 496),
|
||||
(16, 16384, 128),
|
||||
(16, 24576, 4096),
|
||||
(32, 8192, 4096),
|
||||
(32, 16384, 4096),
|
||||
(33, 1024, 1024),
|
||||
(33, 8192, 128),
|
||||
(64, 2048, 496),
|
||||
(64, 16384, 1024),
|
||||
(100, 8192, 496),
|
||||
(128, 32768, 4096),
|
||||
(256, 4096, 4096),
|
||||
(512, 256, 1024),
|
||||
(512, 8192, 4096),
|
||||
(512, 16384, 128),
|
||||
(512, 24576, 128),
|
||||
]
|
||||
|
||||
# Shapes with N or K not divisible by 16. These exercise the padding path
|
||||
# inside CutlassFP8ScaledMMLinearKernel.apply_scaled_mm (e.g. Qwen2.5-VL
|
||||
# vision MLP dims).
|
||||
UNALIGNED_MNK_FACTORS = [
|
||||
(32, 3420, 1280),
|
||||
(32, 1280, 6840),
|
||||
(1, 3420, 1280),
|
||||
(64, 6840, 1280),
|
||||
(16, 100, 200),
|
||||
(33, 255, 513),
|
||||
]
|
||||
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
|
||||
# -1 means full extent in that dimension
|
||||
TENSORWISE_GROUP_SHAPE = (-1, -1)
|
||||
PER_TOKEN_GROUP_SHAPE = (1, -1)
|
||||
PER_OUT_CH_GROUP_SHAPE = (-1, 1)
|
||||
|
||||
capability = current_platform.get_device_capability()
|
||||
capability = capability[0] * 10 + capability[1]
|
||||
|
||||
|
||||
def rand_int8(shape: tuple, device: str = "cuda"):
|
||||
return to_int8(torch.rand(shape, device=device) * 255 - 128)
|
||||
|
||||
|
||||
def group_scale_helper(shape, group_shape):
|
||||
return [shape[i] if s < 0 else s for i, s in enumerate(group_shape)]
|
||||
|
||||
|
||||
def scale_shape(shape, group_shape):
|
||||
assert len(shape) == len(group_shape)
|
||||
group_shape = group_scale_helper(shape, group_shape)
|
||||
return tuple(cdiv(shape[i], group_shape[i]) for i in range(len(group_shape)))
|
||||
|
||||
|
||||
def cutlass_fp8_gemm_helper(
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
a_scale_group_shape: tuple,
|
||||
b_scale_group_shape: tuple,
|
||||
use_bias: bool,
|
||||
out_dtype: type[torch.dtype] = torch.bfloat16,
|
||||
device: str = "cuda",
|
||||
):
|
||||
# Test for a cutlass kernel with per-token activation quantization
|
||||
# and per-output channel weight quantization.
|
||||
a = to_fp8(torch.randn((m, k), device=device))
|
||||
b = to_fp8(torch.randn((n, k), device=device).t())
|
||||
|
||||
a_scales_shape = scale_shape(a.shape, a_scale_group_shape)
|
||||
b_scales_shape = scale_shape(b.shape, b_scale_group_shape)
|
||||
|
||||
scale_a = torch.randn(a_scales_shape, device=device, dtype=torch.float32)
|
||||
scale_b = torch.randn(b_scales_shape, device=device, dtype=torch.float32)
|
||||
|
||||
# make scales M-major for blockwise quant, doesn't affect 1D scales
|
||||
scale_a = scale_a.t().contiguous().t()
|
||||
# make scales K-major for blockwise quant, doesn't affect 1D scales
|
||||
scale_b = scale_b.t().contiguous().t()
|
||||
|
||||
bias = torch.rand((n,), device=device, dtype=out_dtype) * 10 if use_bias else None
|
||||
|
||||
out = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
|
||||
baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
|
||||
|
||||
torch.testing.assert_close(out, baseline, rtol=5e-1, atol=1.5e-1)
|
||||
|
||||
opcheck(torch.ops._C.cutlass_scaled_mm, (out, a, b, scale_a, scale_b, bias))
|
||||
|
||||
|
||||
def cutlass_int8_gemm_helper(
|
||||
m: int,
|
||||
n: int,
|
||||
k: int,
|
||||
a_scale_group_shape: tuple,
|
||||
b_scale_group_shape: tuple,
|
||||
use_bias: bool,
|
||||
out_dtype: type[torch.dtype] = torch.bfloat16,
|
||||
device: str = "cuda",
|
||||
):
|
||||
# Test for a cutlass kernel with per-token activation quantization
|
||||
# and per-output channel weight quantization.
|
||||
a = to_int8(torch.randn((m, k), device=device) * 5)
|
||||
b = to_int8(torch.randn((n, k), device=device).t() * 5)
|
||||
|
||||
a_scales_shape = scale_shape(a.shape, a_scale_group_shape)
|
||||
b_scales_shape = scale_shape(b.shape, b_scale_group_shape)
|
||||
|
||||
scale_a = torch.randn(a_scales_shape, device=device, dtype=torch.float32)
|
||||
scale_b = torch.randn(b_scales_shape, device=device, dtype=torch.float32)
|
||||
|
||||
bias = torch.rand((n,), device=device, dtype=out_dtype) * 10 if use_bias else None
|
||||
|
||||
out = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
|
||||
baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
|
||||
|
||||
torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
|
||||
|
||||
opcheck(torch.ops._C.cutlass_scaled_mm, (out, a, b, scale_a, scale_b, bias))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
|
||||
@pytest.mark.parametrize(
|
||||
"a_scale_group_shape", [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"b_scale_group_shape", [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(89),
|
||||
reason="FP8 is not supported on this GPU type.",
|
||||
)
|
||||
def test_cutlass_fp8_gemm(
|
||||
m: int, n: int, k: int, a_scale_group_shape, b_scale_group_shape, use_bias: bool
|
||||
):
|
||||
cutlass_fp8_gemm_helper(m, n, k, a_scale_group_shape, b_scale_group_shape, use_bias)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m,n,k", UNALIGNED_MNK_FACTORS)
|
||||
@pytest.mark.parametrize(
|
||||
"a_scale_group_shape", [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"b_scale_group_shape", [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(89),
|
||||
reason="FP8 is not supported on this GPU type.",
|
||||
)
|
||||
def test_cutlass_fp8_gemm_padded(
|
||||
m: int, n: int, k: int, a_scale_group_shape, b_scale_group_shape, use_bias: bool
|
||||
):
|
||||
"""Test CUTLASS FP8 GEMM with padding for non-16-aligned N/K dims.
|
||||
|
||||
Exercises CutlassFP8ScaledMMLinearKernel.apply_scaled_mm which pads
|
||||
inputs to satisfy CUTLASS alignment requirements — the path taken by
|
||||
models like Qwen2.5-VL whose vision MLP has non-16-aligned dims.
|
||||
"""
|
||||
from vllm.model_executor.kernels.linear.scaled_mm.cutlass import (
|
||||
CutlassFP8ScaledMMLinearKernel,
|
||||
)
|
||||
|
||||
a = to_fp8(torch.randn((m, k), device="cuda"))
|
||||
b = to_fp8(torch.randn((n, k), device="cuda").t())
|
||||
|
||||
a_scales_shape = scale_shape(a.shape, a_scale_group_shape)
|
||||
b_scales_shape = scale_shape(b.shape, b_scale_group_shape)
|
||||
|
||||
scale_a = torch.randn(a_scales_shape, device="cuda", dtype=torch.float32)
|
||||
scale_b = torch.randn(b_scales_shape, device="cuda", dtype=torch.float32)
|
||||
|
||||
scale_a = scale_a.t().contiguous().t()
|
||||
scale_b = scale_b.t().contiguous().t()
|
||||
|
||||
out_dtype = torch.bfloat16
|
||||
bias = torch.rand((n,), device="cuda", dtype=out_dtype) * 10 if use_bias else None
|
||||
|
||||
baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
|
||||
|
||||
# process_weights_after_loading pad b to 16
|
||||
pad_k = (16 - k % 16) % 16
|
||||
pad_n = (16 - n % 16) % 16
|
||||
if pad_k > 0 or pad_n > 0:
|
||||
b = torch.nn.functional.pad(b.t().contiguous(), (0, pad_k, 0, pad_n)).t()
|
||||
if pad_n > 0 and scale_b.numel() > 1:
|
||||
scale_b = torch.nn.functional.pad(scale_b, (0, pad_n), value=1.0)
|
||||
|
||||
kernel = object.__new__(CutlassFP8ScaledMMLinearKernel)
|
||||
kernel.logical_output_size = n
|
||||
out = kernel.apply_scaled_mm(
|
||||
A=a,
|
||||
B=b,
|
||||
out_dtype=out_dtype,
|
||||
As=scale_a,
|
||||
Bs=scale_b,
|
||||
bias=bias,
|
||||
output_shape=[m, n],
|
||||
)
|
||||
|
||||
torch.testing.assert_close(out, baseline, rtol=5e-1, atol=1.5e-1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
|
||||
@pytest.mark.parametrize(
|
||||
"a_scale_group_shape,b_scale_group_shape", [((1, 128), (128, 128))]
|
||||
)
|
||||
@pytest.mark.parametrize("use_bias", [False])
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(90),
|
||||
reason="FP8 blockwise is not supported on this GPU type.",
|
||||
)
|
||||
def test_cutlass_fp8_blockwise_scale_gemm(
|
||||
m: int, n: int, k: int, a_scale_group_shape, b_scale_group_shape, use_bias: bool
|
||||
):
|
||||
if k % b_scale_group_shape[0] != 0 or n % b_scale_group_shape[1] != 0:
|
||||
return
|
||||
if m % a_scale_group_shape[0] != 0 or k % a_scale_group_shape[1] != 0:
|
||||
return
|
||||
cutlass_fp8_gemm_helper(m, n, k, a_scale_group_shape, b_scale_group_shape, use_bias)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
|
||||
@pytest.mark.parametrize(
|
||||
"a_scale_group_shape", [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"b_scale_group_shape", [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
def test_cutlass_int8_gemm(
|
||||
m: int, n: int, k: int, a_scale_group_shape, b_scale_group_shape, use_bias: bool
|
||||
):
|
||||
cutlass_int8_gemm_helper(
|
||||
m, n, k, a_scale_group_shape, b_scale_group_shape, use_bias
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"a_scale_group_shape", [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"b_scale_group_shape", [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
def test_cutlass_int8_gemm_output_dtype(
|
||||
a_scale_group_shape,
|
||||
b_scale_group_shape,
|
||||
out_dtype: type[torch.dtype],
|
||||
use_bias: bool,
|
||||
):
|
||||
cutlass_int8_gemm_helper(
|
||||
512,
|
||||
512,
|
||||
512,
|
||||
a_scale_group_shape,
|
||||
b_scale_group_shape,
|
||||
use_bias,
|
||||
out_dtype=out_dtype,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"a_scale_group_shape", [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"b_scale_group_shape", [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(89),
|
||||
reason="FP8 is not supported on this GPU type.",
|
||||
)
|
||||
def test_cutlass_fp8_gemm_output_dtype(
|
||||
a_scale_group_shape,
|
||||
b_scale_group_shape,
|
||||
out_dtype: type[torch.dtype],
|
||||
use_bias: bool,
|
||||
):
|
||||
cutlass_fp8_gemm_helper(
|
||||
512,
|
||||
512,
|
||||
512,
|
||||
a_scale_group_shape,
|
||||
b_scale_group_shape,
|
||||
use_bias,
|
||||
out_dtype=out_dtype,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"a_scale_group_shape,b_scale_group_shape", [((1, 128), (128, 128))]
|
||||
)
|
||||
@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("use_bias", [False])
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(90),
|
||||
reason="FP8 blockwise is not supported on this GPU type.",
|
||||
)
|
||||
def test_cutlass_fp8_blockwise_scale_gemm_dtype(
|
||||
a_scale_group_shape,
|
||||
b_scale_group_shape,
|
||||
out_dtype: type[torch.dtype],
|
||||
use_bias: bool,
|
||||
):
|
||||
cutlass_fp8_gemm_helper(
|
||||
512,
|
||||
512,
|
||||
512,
|
||||
a_scale_group_shape,
|
||||
b_scale_group_shape,
|
||||
use_bias,
|
||||
out_dtype=out_dtype,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"a_scale_group_shape", [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"b_scale_group_shape", [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(89),
|
||||
reason="FP8 is not supported on this GPU type.",
|
||||
)
|
||||
def test_cutlass_fp8_gemm_devices(
|
||||
a_scale_group_shape, b_scale_group_shape, use_bias: bool, device: str
|
||||
):
|
||||
cutlass_fp8_gemm_helper(
|
||||
512,
|
||||
512,
|
||||
512,
|
||||
a_scale_group_shape,
|
||||
b_scale_group_shape,
|
||||
use_bias,
|
||||
torch.bfloat16,
|
||||
device,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"a_scale_group_shape", [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"b_scale_group_shape", [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
def test_cutlass_int8_gemm_devices(
|
||||
a_scale_group_shape, b_scale_group_shape, use_bias: bool, device: str
|
||||
):
|
||||
cutlass_int8_gemm_helper(
|
||||
512,
|
||||
512,
|
||||
512,
|
||||
a_scale_group_shape,
|
||||
b_scale_group_shape,
|
||||
use_bias,
|
||||
out_dtype=torch.bfloat16,
|
||||
device=device,
|
||||
)
|
||||
|
||||
|
||||
# For the following two tests:
|
||||
# N and K correspond to the size of the weight matrix and likely to be multiples
|
||||
# of a large power of two. In any case, the kernel will have a naive fallback
|
||||
# when N and K are not divisible by 16. But M is the number of tokens and the
|
||||
# kernel must handle any M thrown at it.
|
||||
@pytest.mark.parametrize(
|
||||
"a_scale_group_shape", [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"b_scale_group_shape", [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.has_device_capability(89),
|
||||
reason="FP8 is not supported on this GPU type.",
|
||||
)
|
||||
def test_cutlass_fp8_gemm_m_sweep(
|
||||
a_scale_group_shape, b_scale_group_shape, use_bias: bool
|
||||
):
|
||||
for nk in range(32, 128, 32):
|
||||
for m in range(1, 128):
|
||||
cutlass_fp8_gemm_helper(
|
||||
m, nk, nk, a_scale_group_shape, b_scale_group_shape, use_bias
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"a_scale_group_shape", [PER_TOKEN_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"b_scale_group_shape", [PER_OUT_CH_GROUP_SHAPE, TENSORWISE_GROUP_SHAPE]
|
||||
)
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
def test_cutlass_int8_gemm_m_sweep(
|
||||
a_scale_group_shape, b_scale_group_shape, use_bias: bool
|
||||
):
|
||||
for nk in range(32, 128, 32):
|
||||
for m in range(1, 128):
|
||||
cutlass_int8_gemm_helper(
|
||||
m, nk, nk, a_scale_group_shape, b_scale_group_shape, use_bias
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", [32, 64, 128])
|
||||
@pytest.mark.parametrize("n", [16, 32, 64])
|
||||
@pytest.mark.parametrize("k", [64, 128, 256])
|
||||
@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.skip
|
||||
def test_cutlass_int8_azp_bias_fold(m: int, n: int, k: int, out_dtype: torch.dtype):
|
||||
# Currently, the test is failing because folding azp into
|
||||
# 16-bit bias loses too much precision
|
||||
scale_a = torch.randn((1, 1), device="cuda", dtype=torch.float32) / 10
|
||||
scale_b = torch.randn((1, n), device="cuda", dtype=torch.float32) / 10
|
||||
|
||||
aq_i8 = rand_int8((m, k))
|
||||
bq_i8 = rand_int8((n, k)).t()
|
||||
|
||||
aq_i32 = aq_i8.to(dtype=torch.int32)
|
||||
bq_i32 = bq_i8.to(dtype=torch.int32)
|
||||
|
||||
aq_f32 = aq_i8.to(dtype=torch.float32)
|
||||
bq_f32 = bq_i8.to(dtype=torch.float32)
|
||||
|
||||
b_dq = scale_b * bq_f32
|
||||
|
||||
azp_a = torch.rand((1,), device="cuda", dtype=torch.float32) * 10 + 1.5
|
||||
azp_aq_i8 = (azp_a / scale_a).to(dtype=torch.int8)
|
||||
azp_a = azp_aq_i8.to(dtype=torch.float32) * scale_a # correct for rounding
|
||||
|
||||
a_dq = scale_a * (aq_i32 + azp_aq_i8).to(dtype=torch.float32)
|
||||
torch.testing.assert_close(a_dq, scale_a * aq_f32 + azp_a)
|
||||
|
||||
baseline_dq = torch.mm(a_dq, b_dq).to(out_dtype)
|
||||
|
||||
J = torch.ones((1, k), device="cuda", dtype=torch.float32)
|
||||
azp_bias = (azp_a * scale_b * (J @ bq_f32)).to(out_dtype)
|
||||
assert azp_bias.shape == (1, n)
|
||||
assert azp_bias[0, :].shape == (n,)
|
||||
|
||||
baseline_q = (
|
||||
scale_a.to(device="cpu")
|
||||
* scale_b.to(device="cpu")
|
||||
* ((aq_i32 + azp_aq_i8).to(device="cpu") @ bq_i32.to(device="cpu"))
|
||||
).to(dtype=out_dtype, device="cuda")
|
||||
|
||||
out = ops.cutlass_scaled_mm(
|
||||
aq_i8, bq_i8, scale_a, scale_b, out_dtype=out_dtype, bias=azp_bias[0, :]
|
||||
)
|
||||
torch.testing.assert_close(out, baseline_dq, rtol=1e-2, atol=1e0)
|
||||
torch.testing.assert_close(out, baseline_q, rtol=1e-2, atol=1e0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("m", [32, 64, 128])
|
||||
@pytest.mark.parametrize("n", [16, 32, 64])
|
||||
@pytest.mark.parametrize("k", [64, 128, 256])
|
||||
@pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float16])
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
@pytest.mark.parametrize("azp_per_token", [True, False])
|
||||
def test_cutlass_int8_azp(
|
||||
m: int, n: int, k: int, out_dtype: torch.dtype, use_bias: bool, azp_per_token: bool
|
||||
):
|
||||
m_azp = m if azp_per_token else 1
|
||||
scale_a = torch.randn((m_azp, 1), device="cuda", dtype=torch.float32) / 10
|
||||
scale_b = torch.randn((1, n), device="cuda", dtype=torch.float32) / 10
|
||||
|
||||
aq_i8 = rand_int8((m, k))
|
||||
aq_i32 = aq_i8.to(dtype=torch.int32)
|
||||
aq_f32 = aq_i8.to(dtype=torch.float32)
|
||||
|
||||
bq_i8 = rand_int8((n, k)).t()
|
||||
bq_i32 = bq_i8.to(dtype=torch.int32)
|
||||
bq_f32 = bq_i8.to(dtype=torch.float32)
|
||||
b_dq = scale_b * bq_f32
|
||||
|
||||
azp_a = torch.rand((m_azp, 1), device="cuda", dtype=torch.float32) * 10 + 1.5
|
||||
azp_aq_i8 = (azp_a / scale_a).to(dtype=torch.int8)
|
||||
azp_a = azp_aq_i8.to(dtype=torch.float32) * scale_a # correct for rounding
|
||||
|
||||
a_dq = scale_a * (aq_i32 - azp_aq_i8).to(dtype=torch.float32)
|
||||
torch.testing.assert_close(a_dq, scale_a * aq_f32 - azp_a, rtol=1e-4, atol=1e-3)
|
||||
|
||||
if use_bias:
|
||||
bias = torch.rand((1, n), device="cuda", dtype=out_dtype) * 10 + 2.5
|
||||
else:
|
||||
bias = torch.zeros((1, n), device="cuda", dtype=out_dtype)
|
||||
|
||||
baseline_dq = (torch.mm(a_dq, b_dq) + bias).to(out_dtype)
|
||||
|
||||
# int32 mm not supported on CUDA
|
||||
a_noazp_i32_cpu = (aq_i32 - azp_aq_i8).to(device="cpu")
|
||||
cq = (a_noazp_i32_cpu @ bq_i32.to(device="cpu")).to(device="cuda")
|
||||
baseline_q = (scale_a * scale_b * cq + bias).to(dtype=out_dtype)
|
||||
|
||||
# Hadamard is just the sum of the cols
|
||||
azp_adj_i32 = bq_i32.sum(dim=0, keepdim=True, dtype=torch.int32)
|
||||
azp_i32 = azp_aq_i8.to(dtype=torch.int32)
|
||||
func_bias = bias if use_bias else None
|
||||
|
||||
if azp_per_token:
|
||||
out = ops.cutlass_scaled_mm_azp(
|
||||
aq_i8, bq_i8, scale_a, scale_b, out_dtype, azp_adj_i32, azp_i32, func_bias
|
||||
)
|
||||
else:
|
||||
azp_with_adj_i32 = azp_i32 * azp_adj_i32
|
||||
out = ops.cutlass_scaled_mm_azp(
|
||||
aq_i8, bq_i8, scale_a, scale_b, out_dtype, azp_with_adj_i32, None, func_bias
|
||||
)
|
||||
|
||||
# bfloat16 precision is 7-bit mantissa -> 2^-8 ~ 0.4%
|
||||
# float16 precision is 10-bit mantissa -> 2^-11 ~ 0.05%
|
||||
rtol = 1e-2 if out_dtype == torch.bfloat16 else 1e-3
|
||||
atol = 1e-3
|
||||
torch.testing.assert_close(out, baseline_dq, rtol=rtol, atol=atol)
|
||||
torch.testing.assert_close(out, baseline_q, rtol=rtol, atol=atol)
|
||||
|
||||
if azp_per_token:
|
||||
opcheck(
|
||||
torch.ops._C.cutlass_scaled_mm_azp,
|
||||
(out, aq_i8, bq_i8, scale_a, scale_b, azp_adj_i32, azp_i32, func_bias),
|
||||
)
|
||||
else:
|
||||
opcheck(
|
||||
torch.ops._C.cutlass_scaled_mm_azp,
|
||||
(out, aq_i8, bq_i8, scale_a, scale_b, azp_with_adj_i32, None, func_bias),
|
||||
)
|
||||
|
||||
|
||||
# Test working with a subset of A and B
|
||||
def test_cutlass_subset():
|
||||
big_m, big_n, big_k = 1024, 1024, 1024
|
||||
m, n, k = 512, 512, 512
|
||||
|
||||
whole_a = to_int8(torch.randn((big_m, big_k), device="cuda") * 5)
|
||||
whole_b = to_int8(torch.randn((big_n, big_k), device="cuda").t() * 5)
|
||||
a = whole_a[0:m, 0:k]
|
||||
b = whole_b[0:k, 0:n]
|
||||
|
||||
scale_a = torch.randn((1, 1), device="cuda", dtype=torch.float32) / 10
|
||||
scale_b = torch.randn((1, 1), device="cuda", dtype=torch.float32) / 10
|
||||
|
||||
out = ops.cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype=torch.bfloat16)
|
||||
baseline = baseline_scaled_mm(a, b, scale_a, scale_b, out_dtype=torch.bfloat16)
|
||||
|
||||
torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
|
||||
|
||||
|
||||
# Test to make sure cuda graphs work
|
||||
class CutlassLayer(torch.nn.Module):
|
||||
def __init__(self, b, scale_a, scale_b, out_dtype):
|
||||
super().__init__()
|
||||
self.b = b
|
||||
self.scale_a = scale_a
|
||||
self.scale_b = scale_b
|
||||
self.out_dtype = out_dtype
|
||||
|
||||
def forward(self, a):
|
||||
return ops.cutlass_scaled_mm(
|
||||
a, self.b, self.scale_a, self.scale_b, self.out_dtype
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("per_act_token", [True, False])
|
||||
@pytest.mark.parametrize("per_out_ch", [True, False])
|
||||
def test_cutlass_cuda_graph(per_act_token: bool, per_out_ch: bool):
|
||||
m, n, k = 512, 512, 512
|
||||
|
||||
a = to_int8(torch.randn((m, k), device="cuda"))
|
||||
b = to_int8(torch.randn((n, k), device="cuda").t())
|
||||
|
||||
m_a_scales = m if per_act_token else 1
|
||||
n_b_scales = n if per_out_ch else 1
|
||||
|
||||
scale_a = torch.randn((m_a_scales, 1), device="cuda", dtype=torch.float32) / 10
|
||||
scale_b = torch.randn((1, n_b_scales), device="cuda", dtype=torch.float32) / 10
|
||||
|
||||
# Construct a trivial model with a single layer that calls a CUTLASS kernel
|
||||
model = CutlassLayer(b, scale_a, scale_b, torch.bfloat16)
|
||||
|
||||
# Run the model with a cuda graph
|
||||
stream = torch.cuda.Stream()
|
||||
with torch.cuda.stream(stream):
|
||||
g = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(g):
|
||||
out = model(a)
|
||||
out.zero_()
|
||||
g.replay()
|
||||
|
||||
baseline = torch.mm(
|
||||
scale_a * a.to(dtype=torch.float32), scale_b * b.to(dtype=torch.float32)
|
||||
).to(torch.bfloat16)
|
||||
torch.testing.assert_close(out, baseline, rtol=1e-1, atol=1e0)
|
||||
|
||||
|
||||
def test_cutlass_support_opcheck():
|
||||
opcheck(torch.ops._C.cutlass_scaled_mm_supports_fp8, (capability,))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_experts", [8, 64])
|
||||
@pytest.mark.parametrize("per_act_token", [True, False])
|
||||
@pytest.mark.parametrize("per_out_ch", [True, False])
|
||||
@pytest.mark.parametrize("use_bias", [False])
|
||||
@pytest.mark.skipif(
|
||||
(lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))(
|
||||
current_platform.get_device_capability()
|
||||
),
|
||||
reason="Grouped gemm is not supported on this GPU type.",
|
||||
)
|
||||
def test_cutlass_fp8_group_gemm(
|
||||
num_experts: int, per_act_token: bool, per_out_ch: bool, use_bias: bool
|
||||
):
|
||||
# Device and dtype setup
|
||||
device = "cuda"
|
||||
out_dtype = torch.half
|
||||
|
||||
# Create separate A, B, C tensors for each group
|
||||
a_tensors = []
|
||||
b_tensors = []
|
||||
a_scales_tensors = []
|
||||
b_scales_tensors = []
|
||||
baseline_tensors = []
|
||||
|
||||
expert_offsets = torch.zeros((num_experts + 1), device=device, dtype=torch.int64)
|
||||
|
||||
problem_sizes = torch.zeros((num_experts, 3), device=device, dtype=torch.int32)
|
||||
|
||||
if not per_act_token:
|
||||
one_scale_a = torch.randn((1, 1), device=device, dtype=torch.float32)
|
||||
|
||||
alignment = 16 # 128 // 8
|
||||
# For variation, each group has dimensions
|
||||
n_g = alignment * random.randint(1, 64)
|
||||
k_g = alignment * random.randint(1, 64)
|
||||
for g in range(num_experts):
|
||||
m_g = alignment * random.randint(1, 64)
|
||||
|
||||
expert_offsets[g + 1] = expert_offsets[g] + m_g
|
||||
problem_sizes[g][0] = m_g
|
||||
problem_sizes[g][1] = n_g
|
||||
problem_sizes[g][2] = k_g
|
||||
|
||||
m_a_scales = m_g if per_act_token else 1
|
||||
n_b_scales = n_g if per_out_ch else 1
|
||||
|
||||
# Create group-specific A and B (FP8) and output (FP16/FP32)
|
||||
a_g = to_fp8(torch.randn((m_g, k_g), device=device))
|
||||
b_g = to_fp8(torch.randn((n_g, k_g), device=device).t())
|
||||
a_tensors.append(a_g)
|
||||
b_tensors.append(b_g)
|
||||
|
||||
# Set up A/B scales
|
||||
scale_b = torch.randn((1, n_b_scales), device=device, dtype=torch.float32)
|
||||
b_scales_tensors.append(scale_b)
|
||||
|
||||
if per_act_token:
|
||||
scale_a = torch.randn((m_a_scales, 1), device=device, dtype=torch.float32)
|
||||
a_scales_tensors.append(scale_a)
|
||||
else:
|
||||
scale_a = one_scale_a
|
||||
|
||||
# Compute baseline result for this group
|
||||
baseline_g = baseline_scaled_mm(a_g, b_g, scale_a, scale_b, out_dtype, None)
|
||||
baseline_tensors.append(baseline_g)
|
||||
|
||||
a_tensors_stacked = torch.empty(
|
||||
(expert_offsets[num_experts], k_g), device=device, dtype=torch.float8_e4m3fn
|
||||
)
|
||||
b_tensors_stacked = torch.empty(
|
||||
(num_experts, n_g, k_g), device=device, dtype=torch.float8_e4m3fn
|
||||
)
|
||||
|
||||
for g in range(num_experts):
|
||||
a_tensors_stacked[expert_offsets[g] : expert_offsets[g + 1]] = a_tensors[g]
|
||||
b_tensors_stacked[g] = b_tensors[g].t()
|
||||
b_tensors_stacked = b_tensors_stacked.transpose(1, 2)
|
||||
|
||||
if per_act_token:
|
||||
a_scales_tensors_stacked = torch.empty(
|
||||
(expert_offsets[num_experts], 1), device=device, dtype=torch.float32
|
||||
)
|
||||
for g in range(num_experts):
|
||||
a_scales_tensors_stacked[expert_offsets[g] : expert_offsets[g + 1]] = (
|
||||
a_scales_tensors[g]
|
||||
)
|
||||
else:
|
||||
a_scales_tensors_stacked = one_scale_a
|
||||
|
||||
b_scales_tensors_stacked = torch.empty(
|
||||
(num_experts, n_b_scales), device=device, dtype=torch.float32
|
||||
)
|
||||
for g in range(num_experts):
|
||||
b_scales_tensors_stacked[g] = b_scales_tensors[g]
|
||||
|
||||
out_tensors_stacked = torch.zeros(
|
||||
(expert_offsets[num_experts], n_g), device=device, dtype=out_dtype
|
||||
)
|
||||
|
||||
ab_strides = torch.full(
|
||||
(num_experts,), a_tensors_stacked.stride(0), device="cuda", dtype=torch.int64
|
||||
)
|
||||
c_strides = torch.full(
|
||||
(num_experts,), out_tensors_stacked.stride(0), device="cuda", dtype=torch.int64
|
||||
)
|
||||
|
||||
ops.cutlass_moe_mm(
|
||||
out_tensors_stacked,
|
||||
a_tensors_stacked,
|
||||
b_tensors_stacked,
|
||||
a_scales_tensors_stacked,
|
||||
b_scales_tensors_stacked,
|
||||
expert_offsets[:-1],
|
||||
problem_sizes,
|
||||
ab_strides,
|
||||
ab_strides,
|
||||
c_strides,
|
||||
per_act_token,
|
||||
per_out_ch,
|
||||
)
|
||||
|
||||
# Validate each group's result against the baseline
|
||||
for g in range(num_experts):
|
||||
baseline = baseline_tensors[g]
|
||||
c = out_tensors_stacked[expert_offsets[g] : expert_offsets[g + 1]]
|
||||
torch.testing.assert_close(c, baseline, rtol=1e-2, atol=5e-4)
|
||||
@@ -0,0 +1,329 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the CUTLASS W4A8 kernel.
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_cutlass_w4a8.py`.
|
||||
"""
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
convert_packed_uint4b8_to_signed_int4_inplace,
|
||||
pack_cols,
|
||||
pack_rows,
|
||||
quantize_weights,
|
||||
unpack_quantized_values_into_int32,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.scalar_type import ScalarType, scalar_types
|
||||
|
||||
if not current_platform.is_cuda():
|
||||
pytest.skip("These tests use CUTLASS which requires CUDA", allow_module_level=True)
|
||||
|
||||
# TODO: in future PR refactor this and `is_quant_method_supported` in the kernel
|
||||
# unit tests to a common utility function. Currently the use of
|
||||
# `is_quant_method_supported` conflates kernels with quantization methods
|
||||
# an assumption which is breaking down as quantizations methods can have
|
||||
# have kernels and some kernels support multiple quantization methods.
|
||||
IS_SUPPORTED_BY_GPU = current_platform.get_device_capability()[0] >= 9
|
||||
|
||||
MNK_SHAPES = [
|
||||
(1, 128, 128),
|
||||
(1, 512, 1024),
|
||||
(1, 4096, 4096),
|
||||
(1, 8192, 28672),
|
||||
(13, 8192, 4096),
|
||||
(26, 4096, 8192),
|
||||
(64, 4096, 4096),
|
||||
(64, 8192, 28672),
|
||||
(257, 128, 4096),
|
||||
(257, 4096, 4096),
|
||||
(1024, 4096, 8192),
|
||||
(1024, 8192, 4096),
|
||||
]
|
||||
|
||||
# TODO(czhu): get supported schedules from fn
|
||||
SCHEDULES = [
|
||||
"128x16_1x1x1",
|
||||
"256x16_1x1x1",
|
||||
"128x32_1x1x1",
|
||||
"256x32_1x1x1",
|
||||
"128x64_1x1x1",
|
||||
"256x64_1x1x1",
|
||||
"128x128_1x1x1",
|
||||
"256x128_1x1x1",
|
||||
"128x256_1x1x1",
|
||||
"128x256_2x1x1",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class TypeConfig:
|
||||
act_type: torch.dtype
|
||||
weight_type: ScalarType
|
||||
output_type: torch.dtype | None
|
||||
group_scale_type: torch.dtype | None
|
||||
channel_scale_type: torch.dtype | None
|
||||
token_scale_type: torch.dtype | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Tensors:
|
||||
w_ref: torch.Tensor
|
||||
a_ref: torch.Tensor
|
||||
a: torch.Tensor
|
||||
w_q: torch.Tensor
|
||||
w_g_s: torch.Tensor
|
||||
w_ch_s: torch.Tensor
|
||||
w_tok_s: torch.Tensor
|
||||
|
||||
|
||||
# (Act Type, Weight Type, Output Type, Scale Type, ZeroPoints,
|
||||
# Ch Scales Type, Tok Scales Type)
|
||||
TestTypeTuple = tuple[
|
||||
list[torch.dtype], ScalarType, torch.dtype | None, torch.dtype | None, bool
|
||||
]
|
||||
TEST_TYPES = [
|
||||
*(
|
||||
TypeConfig(
|
||||
act_type=torch.float8_e4m3fn,
|
||||
weight_type=w_type,
|
||||
output_type=o_type,
|
||||
group_scale_type=torch.float8_e4m3fn,
|
||||
channel_scale_type=torch.float32,
|
||||
token_scale_type=torch.float32,
|
||||
)
|
||||
for w_type in [scalar_types.int4]
|
||||
# TODO(czhu): fp16 out type
|
||||
for o_type in [torch.bfloat16]
|
||||
),
|
||||
]
|
||||
|
||||
# TODO: in future PR refactor this and `is_quant_method_supported` in the kernel
|
||||
# unit tests to a common utility function. Currently the use of
|
||||
# `is_quant_method_supported` conflates kernels with quantization methods
|
||||
# an assumption which is breaking down as quantizations methods can have
|
||||
# have kernels and some kernels support multiple quantization methods.
|
||||
IS_SUPPORTED_BY_GPU = current_platform.has_device_capability(90)
|
||||
|
||||
|
||||
# For testing quantized linear kernels
|
||||
def to_fp8(tensor: torch.Tensor):
|
||||
finfo = torch.finfo(torch.float8_e4m3fn)
|
||||
return tensor.clamp(min=finfo.min, max=finfo.max).to(dtype=torch.float8_e4m3fn)
|
||||
|
||||
|
||||
def cutlass_quantize_and_pack(
|
||||
atype: torch.dtype,
|
||||
w: torch.Tensor,
|
||||
wtype: ScalarType,
|
||||
stype: torch.dtype | None,
|
||||
group_size: int | None,
|
||||
zero_points: bool = False,
|
||||
):
|
||||
assert wtype.is_integer(), "TODO: support floating point weights"
|
||||
|
||||
w_ref, w_q, w_s, w_zp = quantize_weights(
|
||||
w, wtype, group_size=group_size, zero_points=zero_points
|
||||
)
|
||||
|
||||
# since scales are cast to fp8, we need to compute w_ref this way
|
||||
w_ref = (
|
||||
(w_q).to(torch.float32)
|
||||
* w_s.to(atype).to(torch.float32).repeat_interleave(group_size, dim=0)
|
||||
).to(atype)
|
||||
|
||||
# bit mask prevents sign extending int4 when packing
|
||||
w_q = pack_rows(w_q & 0x0F, wtype.size_bits, *w_q.shape)
|
||||
w_q = w_q.t().contiguous().t() # convert to col major
|
||||
|
||||
w_q_packed = ops.cutlass_encode_and_reorder_int4b(w_q)
|
||||
w_s_packed = ops.cutlass_pack_scale_fp8(w_s.to(atype))
|
||||
|
||||
return w_ref, w_q_packed, w_s_packed, w_zp
|
||||
|
||||
|
||||
def create_test_tensors(
|
||||
shape: tuple[int, int, int], types: TypeConfig, group_size: int | None
|
||||
) -> Tensors:
|
||||
m, n, k = shape
|
||||
|
||||
print(
|
||||
"create_test_tensors, shape:", shape, "types:", types, "group_size:", group_size
|
||||
)
|
||||
|
||||
a = to_fp8(torch.randn((m, k), device="cuda"))
|
||||
w = to_fp8(torch.randn((k, n), device="cuda"))
|
||||
|
||||
if types.group_scale_type is not None:
|
||||
w = w.to(types.group_scale_type)
|
||||
if w.dtype.itemsize == 1:
|
||||
w = w.to(torch.float16)
|
||||
|
||||
w_ref, w_q_packed, w_s, _ = cutlass_quantize_and_pack(
|
||||
a.dtype, w, types.weight_type, types.group_scale_type, group_size, False
|
||||
)
|
||||
|
||||
a_ref = a.to(torch.float32)
|
||||
w_ref = w_ref.to(torch.float32)
|
||||
|
||||
# for the practical use case we need per-tok scales for fp8 activations
|
||||
w_tok_s = torch.randn((m,), device="cuda", dtype=types.token_scale_type)
|
||||
w_ch_s = torch.randn((n,), device="cuda", dtype=types.channel_scale_type)
|
||||
|
||||
return Tensors(
|
||||
w_ref=w_ref,
|
||||
a_ref=a_ref,
|
||||
a=a,
|
||||
w_q=w_q_packed,
|
||||
w_g_s=w_s,
|
||||
w_ch_s=w_ch_s,
|
||||
w_tok_s=w_tok_s,
|
||||
)
|
||||
|
||||
|
||||
def mm_test_helper(
|
||||
types: TypeConfig,
|
||||
tensors: Tensors,
|
||||
group_size: int | None = None,
|
||||
schedule: str | None = None,
|
||||
):
|
||||
# CUTLASS upstream uses fp8 with fastaccum as reference
|
||||
# https://github.com/NVIDIA/cutlass/blob/main/examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu#L406
|
||||
output_ref = torch._scaled_mm(
|
||||
tensors.a_ref.to(types.act_type),
|
||||
tensors.w_ref.to(types.act_type).t().contiguous().t(), # col major
|
||||
tensors.w_tok_s.unsqueeze(1),
|
||||
tensors.w_ch_s.unsqueeze(0),
|
||||
out_dtype=types.output_type,
|
||||
use_fast_accum=True,
|
||||
)
|
||||
|
||||
output = ops.cutlass_w4a8_mm(
|
||||
a=tensors.a,
|
||||
b_q=tensors.w_q,
|
||||
b_group_scales=tensors.w_g_s,
|
||||
b_group_size=group_size,
|
||||
b_channel_scales=tensors.w_ch_s,
|
||||
a_token_scales=tensors.w_tok_s,
|
||||
)
|
||||
|
||||
print(output)
|
||||
print(output_ref)
|
||||
|
||||
torch.testing.assert_close(
|
||||
output, output_ref.to(output.dtype), rtol=1e-2, atol=1e-2
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not IS_SUPPORTED_BY_GPU, reason="CUTLASS W4A8 is not supported on this GPU type."
|
||||
)
|
||||
@pytest.mark.parametrize("shape", MNK_SHAPES, ids=lambda x: "x".join(str(v) for v in x))
|
||||
@pytest.mark.parametrize("types", TEST_TYPES)
|
||||
@pytest.mark.parametrize("schedule", SCHEDULES)
|
||||
def test_cutlass_w4a8(shape, types: TypeConfig, schedule):
|
||||
group_sizes = [128]
|
||||
for group_size in group_sizes:
|
||||
tensors = create_test_tensors(shape, types, group_size)
|
||||
mm_test_helper(types, tensors, group_size, schedule)
|
||||
|
||||
|
||||
# Test to make sure cuda graphs work
|
||||
class W4A8Layer(torch.nn.Module):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__()
|
||||
self.kwargs = kwargs
|
||||
|
||||
def forward(self, a):
|
||||
return ops.cutlass_w4a8_mm(a=a, **self.kwargs)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not IS_SUPPORTED_BY_GPU, reason="CUTLASS W4A8 is not supported on this GPU type."
|
||||
)
|
||||
def test_w4a8_cuda_graph():
|
||||
m, n, k = 512, 4096, 4096
|
||||
|
||||
a = to_fp8(torch.randn((m, k), device="cuda"))
|
||||
b = to_fp8(torch.randn((k, n), device="cuda"))
|
||||
|
||||
wtype = scalar_types.int4
|
||||
stype = torch.float8_e4m3fn
|
||||
group_size = 128
|
||||
zero_points = False
|
||||
|
||||
w_ref, w_q_packed, w_s, _ = cutlass_quantize_and_pack(
|
||||
a.dtype, b.to(torch.float16), wtype, stype, group_size, zero_points
|
||||
)
|
||||
|
||||
w_tok_s = torch.randn((m,), device="cuda", dtype=torch.float32)
|
||||
w_ch_s = torch.randn((n,), device="cuda", dtype=torch.float32)
|
||||
|
||||
# Construct a trivial model with a single layer that calls the kernel
|
||||
model = W4A8Layer(
|
||||
b_q=w_q_packed,
|
||||
b_group_scales=w_s,
|
||||
b_group_size=group_size,
|
||||
b_channel_scales=w_ch_s,
|
||||
a_token_scales=w_tok_s,
|
||||
)
|
||||
|
||||
output_ref = torch._scaled_mm(
|
||||
a,
|
||||
w_ref.to(a.dtype).t().contiguous().t(), # col major
|
||||
w_tok_s.unsqueeze(1),
|
||||
w_ch_s.unsqueeze(0),
|
||||
out_dtype=torch.bfloat16,
|
||||
use_fast_accum=True,
|
||||
)
|
||||
|
||||
# Run the model with a cuda graph
|
||||
stream = torch.cuda.Stream()
|
||||
with torch.cuda.stream(stream):
|
||||
g = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(g):
|
||||
output = model(a)
|
||||
|
||||
output.zero_()
|
||||
g.replay()
|
||||
|
||||
torch.testing.assert_close(output, output_ref, rtol=1e-2, atol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not IS_SUPPORTED_BY_GPU, reason="CUTLASS W4A8 is not supported on this GPU type."
|
||||
)
|
||||
@pytest.mark.parametrize("shape", MNK_SHAPES)
|
||||
def test_convert_packed_uint4b8_to_signed_int4_inplace(shape):
|
||||
"""
|
||||
The W4A16 checkpoints encode the weights as int4b8 packed to int32.
|
||||
The CUTLASS kernels expect signed int4 packed to int32.
|
||||
This tests checks that the runtime int4b8 -> signed int4 conversion
|
||||
matches the offline conversion step exactly.
|
||||
"""
|
||||
_, N, K = shape
|
||||
# random weights packed to int32
|
||||
t = torch.randint(
|
||||
low=torch.iinfo(torch.int32).min,
|
||||
high=torch.iinfo(torch.int32).max + 1,
|
||||
size=(N, K // 8),
|
||||
dtype=torch.int32,
|
||||
device="cuda",
|
||||
)
|
||||
|
||||
# compute reference
|
||||
unpacked = unpack_quantized_values_into_int32(
|
||||
t.clone(), scalar_types.uint4b8, packed_dim=1
|
||||
)
|
||||
unpacked = unpacked - 8 # int4b8 -> signed int4
|
||||
ref = pack_cols(unpacked & 0x0F, 4, *unpacked.shape)
|
||||
|
||||
out = convert_packed_uint4b8_to_signed_int4_inplace(t.clone())
|
||||
|
||||
assert torch.equal(ref, out)
|
||||
assert not torch.equal(ref, t)
|
||||
@@ -0,0 +1,343 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests for the CUTLASS-based W4A8 grouped GEMM kernel and the full MoE layer.
|
||||
"""
|
||||
|
||||
import random
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
pack_rows,
|
||||
quantize_weights,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.scalar_type import ScalarType, scalar_types
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
IS_SUPPORTED_BY_GPU = (
|
||||
current_platform.is_cuda() and current_platform.get_device_capability()[0] >= 9
|
||||
)
|
||||
|
||||
|
||||
def to_fp8(tensor: torch.Tensor) -> torch.Tensor:
|
||||
finfo = torch.finfo(torch.float8_e4m3fn)
|
||||
return tensor.clamp(min=finfo.min, max=finfo.max).to(dtype=torch.float8_e4m3fn)
|
||||
|
||||
|
||||
def cutlass_quantize(
|
||||
atype: torch.dtype,
|
||||
w: torch.Tensor,
|
||||
wtype: ScalarType,
|
||||
stype: torch.dtype | None,
|
||||
group_size: int | None,
|
||||
zero_points: bool = False,
|
||||
):
|
||||
"""
|
||||
Quantize weights into W4 and compute reference dequantized weights.
|
||||
|
||||
Encoding/reordering of weights and packing of scales is deferred
|
||||
until after all experts are combined.
|
||||
"""
|
||||
assert wtype.is_integer(), "TODO: support floating point weights"
|
||||
|
||||
w_ref, w_q, w_s, w_zp = quantize_weights(
|
||||
w, wtype, group_size=group_size, zero_points=zero_points
|
||||
)
|
||||
|
||||
# Since scales are later cast to fp8, recompute w_ref in atype here.
|
||||
w_ref = (
|
||||
w_q.to(torch.float32)
|
||||
* w_s.to(atype).to(torch.float32).repeat_interleave(group_size, dim=0)
|
||||
).to(atype)
|
||||
|
||||
# Bit mask prevents sign extension of int4 when packing.
|
||||
w_q = pack_rows(w_q & 0x0F, wtype.size_bits, *w_q.shape)
|
||||
# Make weights row-major (N, K).
|
||||
w_q = w_q.t().contiguous()
|
||||
|
||||
return w_ref, w_q, w_s.to(atype), w_zp
|
||||
|
||||
|
||||
def cutlass_preprocess(
|
||||
w_q_experts: list[torch.Tensor], w_s_experts: list[torch.Tensor]
|
||||
):
|
||||
"""
|
||||
Reorder/encode expert weights and pack scales.
|
||||
|
||||
Returns:
|
||||
w_q_packed: Packed/encoded int4 weights for all experts.
|
||||
w_s_packed: Packed fp8 scales for all experts.
|
||||
packed_layout: Layout/stride metadata for grouped GEMM.
|
||||
"""
|
||||
w_s_packed = ops.cutlass_pack_scale_fp8(torch.stack(w_s_experts))
|
||||
w_q_packed, packed_layout = ops.cutlass_encode_and_reorder_int4b_grouped(
|
||||
torch.stack(w_q_experts)
|
||||
) # expects dim 3
|
||||
return w_q_packed, w_s_packed, packed_layout
|
||||
|
||||
|
||||
GROUP_SIZE = 128
|
||||
# (num_experts, N, K)
|
||||
TEST_SHAPES = [
|
||||
(8, 512, 2048),
|
||||
(8, 2048, 2048),
|
||||
(64, 512, 1024),
|
||||
(64, 2048, 2048),
|
||||
(4, 2048, 768),
|
||||
(8, 768, 2048),
|
||||
(64, 1536, 2048),
|
||||
(128, 8192, 4096), # test overflow int32
|
||||
]
|
||||
ALIGNMENT = 16 # torch._scaled_mm alignment for M, needed for reference check
|
||||
|
||||
|
||||
@dataclass
|
||||
class MoETestSetup:
|
||||
num_experts: int
|
||||
K: int
|
||||
N: int
|
||||
Ms: list[int]
|
||||
M_full: int
|
||||
a: torch.Tensor
|
||||
a_ref: torch.Tensor
|
||||
a_strides: torch.Tensor
|
||||
out: torch.Tensor
|
||||
c_strides: torch.Tensor
|
||||
per_tok_scales: torch.Tensor
|
||||
per_chan_scales: torch.Tensor
|
||||
w_refs: list[torch.Tensor]
|
||||
w_q_packed: torch.Tensor
|
||||
w_s_packed: torch.Tensor
|
||||
problem_sizes: torch.Tensor
|
||||
expert_offsets: torch.Tensor
|
||||
b_strides: torch.Tensor
|
||||
group_scale_strides: torch.Tensor
|
||||
|
||||
|
||||
def make_moe_test_setup(
|
||||
num_experts: int,
|
||||
K: int,
|
||||
N: int,
|
||||
*,
|
||||
alignment: int = ALIGNMENT,
|
||||
max_blocks: int = 64,
|
||||
device: str = "cuda",
|
||||
random_zero: bool = False,
|
||||
) -> MoETestSetup:
|
||||
"""Create a full set of tensors for testing cutlass_w4a8_moe_mm."""
|
||||
|
||||
assert K % GROUP_SIZE == 0
|
||||
# Token counts per expert (multiples of `alignment`).
|
||||
Ms = [alignment * random.randint(1, max_blocks) for _ in range(num_experts)]
|
||||
|
||||
# set random experts to 0 tokens
|
||||
if random_zero and num_experts > 1:
|
||||
num_zero = max(1, num_experts // 8)
|
||||
zero_indices = random.sample(range(num_experts), k=num_zero)
|
||||
for idx in zero_indices:
|
||||
Ms[idx] = 0
|
||||
|
||||
M_full = sum(Ms)
|
||||
assert M_full > 0
|
||||
|
||||
# Activations.
|
||||
a = to_fp8(torch.randn((M_full, K), device=device))
|
||||
a_ref = a.to(torch.float32)
|
||||
a_strides = torch.full((num_experts,), K, dtype=torch.int64, device=device)
|
||||
|
||||
# Output buffer.
|
||||
out = torch.empty((M_full, N), dtype=torch.bfloat16, device=device)
|
||||
c_strides = torch.full((num_experts,), N, dtype=torch.int64, device=device)
|
||||
|
||||
# Channel/token scales.
|
||||
per_tok_scales = torch.randn((M_full, 1), dtype=torch.float32, device=device)
|
||||
per_chan_scales = torch.randn(
|
||||
(num_experts, N, 1), dtype=torch.float32, device=device
|
||||
)
|
||||
|
||||
# Expert weights and scales.
|
||||
wtype = scalar_types.int4
|
||||
atype = stype = torch.float8_e4m3fn
|
||||
w_refs, w_qs, w_ss = [], [], []
|
||||
for _ in range(num_experts):
|
||||
b = to_fp8(torch.randn((K, N), device=device))
|
||||
w_ref, w_q, w_s, _ = cutlass_quantize(
|
||||
atype, b.to(torch.float16), wtype, stype, GROUP_SIZE, zero_points=False
|
||||
)
|
||||
w_refs.append(w_ref)
|
||||
w_qs.append(w_q)
|
||||
w_ss.append(w_s)
|
||||
|
||||
w_q_packed, w_s_packed, packed_layout = cutlass_preprocess(w_qs, w_ss)
|
||||
|
||||
problem_sizes = torch.tensor(
|
||||
[[N, M, K] for M in Ms], dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
expert_offsets = torch.cat(
|
||||
[
|
||||
torch.tensor([0], dtype=torch.int64),
|
||||
torch.cumsum(torch.tensor(Ms, dtype=torch.int64), dim=0)[:-1],
|
||||
]
|
||||
).to(device=device)
|
||||
|
||||
# B strides and group scale strides.
|
||||
b_strides = packed_layout
|
||||
group_scale_strides = torch.zeros(
|
||||
(num_experts, 2), dtype=torch.int64, device=device
|
||||
)
|
||||
group_scale_strides[:, 0] = N
|
||||
|
||||
return MoETestSetup(
|
||||
num_experts=num_experts,
|
||||
K=K,
|
||||
N=N,
|
||||
Ms=Ms,
|
||||
M_full=M_full,
|
||||
a=a,
|
||||
a_ref=a_ref,
|
||||
a_strides=a_strides,
|
||||
out=out,
|
||||
c_strides=c_strides,
|
||||
per_tok_scales=per_tok_scales,
|
||||
per_chan_scales=per_chan_scales,
|
||||
w_refs=w_refs,
|
||||
w_q_packed=w_q_packed,
|
||||
w_s_packed=w_s_packed,
|
||||
problem_sizes=problem_sizes,
|
||||
expert_offsets=expert_offsets,
|
||||
b_strides=b_strides,
|
||||
group_scale_strides=group_scale_strides,
|
||||
)
|
||||
|
||||
|
||||
def compute_moe_reference_output(setup: MoETestSetup) -> torch.Tensor:
|
||||
"""Compute reference output using torch._scaled_mm per expert."""
|
||||
out_ref = torch.empty_like(setup.out)
|
||||
|
||||
ends = torch.cumsum(torch.tensor(setup.Ms), 0).tolist()
|
||||
starts = setup.expert_offsets.cpu().tolist()
|
||||
|
||||
for i in range(setup.num_experts):
|
||||
start, end = starts[i], ends[i]
|
||||
if start == end:
|
||||
continue
|
||||
|
||||
out_ref_i = torch._scaled_mm(
|
||||
setup.a_ref[start:end].to(torch.float8_e4m3fn),
|
||||
setup.w_refs[i].to(torch.float8_e4m3fn).t().contiguous().t(),
|
||||
setup.per_tok_scales[start:end], # (M, 1)
|
||||
setup.per_chan_scales[i].reshape(1, -1), # (1, N)
|
||||
out_dtype=torch.bfloat16,
|
||||
use_fast_accum=True,
|
||||
)
|
||||
out_ref[start:end] = out_ref_i
|
||||
|
||||
return out_ref
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not IS_SUPPORTED_BY_GPU,
|
||||
reason="W4A8 Grouped GEMM is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("shape", TEST_SHAPES)
|
||||
@pytest.mark.parametrize("random_zero", [True, False])
|
||||
def test_cutlass_w4a8_moe_mm_end_to_end(shape, random_zero):
|
||||
num_experts, N, K = shape
|
||||
set_random_seed(42)
|
||||
setup = make_moe_test_setup(
|
||||
num_experts=num_experts, K=K, N=N, max_blocks=64, random_zero=random_zero
|
||||
)
|
||||
|
||||
ops.cutlass_w4a8_moe_mm(
|
||||
setup.out,
|
||||
setup.a,
|
||||
setup.w_q_packed,
|
||||
setup.per_tok_scales,
|
||||
setup.per_chan_scales,
|
||||
setup.w_s_packed,
|
||||
GROUP_SIZE,
|
||||
setup.expert_offsets,
|
||||
setup.problem_sizes,
|
||||
setup.a_strides,
|
||||
setup.b_strides,
|
||||
setup.c_strides,
|
||||
setup.group_scale_strides,
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
out_ref = compute_moe_reference_output(setup)
|
||||
torch.testing.assert_close(setup.out, out_ref, rtol=1e-2, atol=1e-2)
|
||||
|
||||
|
||||
class W4A8MoELayer(torch.nn.Module):
|
||||
"""
|
||||
Minimal wrapper module to test cuda graphs
|
||||
"""
|
||||
|
||||
def __init__(self, setup: MoETestSetup):
|
||||
super().__init__()
|
||||
self.setup = setup
|
||||
|
||||
def forward(self, a: torch.Tensor) -> torch.Tensor:
|
||||
s = self.setup
|
||||
ops.cutlass_w4a8_moe_mm(
|
||||
s.out,
|
||||
a,
|
||||
s.w_q_packed,
|
||||
s.per_tok_scales,
|
||||
s.per_chan_scales,
|
||||
s.w_s_packed,
|
||||
GROUP_SIZE,
|
||||
s.expert_offsets,
|
||||
s.problem_sizes,
|
||||
s.a_strides,
|
||||
s.b_strides,
|
||||
s.c_strides,
|
||||
s.group_scale_strides,
|
||||
)
|
||||
return s.out
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not IS_SUPPORTED_BY_GPU,
|
||||
reason="W4A8 Grouped GEMM is not supported on this GPU type.",
|
||||
)
|
||||
def test_cutlass_w4a8_moe_mm_cuda_graph():
|
||||
set_random_seed(42)
|
||||
# Fixed config for CUDA graph test (single parameter point).
|
||||
num_experts = 8
|
||||
K = 512
|
||||
N = 2048
|
||||
|
||||
setup = make_moe_test_setup(
|
||||
num_experts=num_experts,
|
||||
K=K,
|
||||
N=N,
|
||||
max_blocks=32,
|
||||
)
|
||||
|
||||
# Construct model that calls the grouped GEMM kernel.
|
||||
model = W4A8MoELayer(setup)
|
||||
|
||||
# Build reference output once.
|
||||
out_ref = compute_moe_reference_output(setup)
|
||||
|
||||
# Capture and run the model in a CUDA graph.
|
||||
a_static = setup.a.clone() # static input tensor for graph replay
|
||||
|
||||
stream = torch.cuda.Stream()
|
||||
with torch.cuda.stream(stream):
|
||||
g = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(g):
|
||||
out_static = model(a_static)
|
||||
|
||||
out_static.zero_()
|
||||
g.replay()
|
||||
|
||||
torch.testing.assert_close(out_static, out_ref, rtol=1e-2, atol=1e-2)
|
||||
@@ -0,0 +1,168 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
from nvfp4_utils import (
|
||||
FLOAT4_E2M1_MAX,
|
||||
FLOAT8_E4M3_MAX,
|
||||
convert_swizzled_to_linear,
|
||||
dequantize_nvfp4_to_dtype,
|
||||
)
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.flashinfer import (
|
||||
flashinfer_scaled_fp4_mm,
|
||||
has_flashinfer_b12x_gemm,
|
||||
)
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not current_platform.has_device_capability(100):
|
||||
pytest.skip(
|
||||
reason="Nvfp4 Requires compute capability of 10 or above.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
# m, n, k
|
||||
SHAPES = [
|
||||
(128, 128, 64),
|
||||
(128, 128, 128),
|
||||
(256, 128, 64),
|
||||
(128, 256, 128),
|
||||
(1, 128, 128),
|
||||
]
|
||||
PAD_SHAPES = [(150, 128, 64), (128, 128, 96), (2, 128, 64), (3, 128, 96)]
|
||||
SHAPES.extend(PAD_SHAPES)
|
||||
|
||||
SEEDS = [42]
|
||||
CUDA_DEVICES = ["cuda:0"]
|
||||
|
||||
|
||||
def get_ref_results(
|
||||
a_fp4,
|
||||
b_fp4,
|
||||
a_sf,
|
||||
b_sf,
|
||||
a_global_scale,
|
||||
b_global_scale,
|
||||
m,
|
||||
n,
|
||||
dtype,
|
||||
block_size,
|
||||
device,
|
||||
is_sf_128x4_layout,
|
||||
):
|
||||
_, m_k = a_fp4.shape
|
||||
_, n_k = b_fp4.shape
|
||||
assert m_k == n_k
|
||||
a_in_dtype = dequantize_nvfp4_to_dtype(
|
||||
a_fp4,
|
||||
a_sf,
|
||||
a_global_scale,
|
||||
dtype=dtype,
|
||||
device=device,
|
||||
block_size=block_size,
|
||||
is_sf_128x4_layout=is_sf_128x4_layout,
|
||||
)
|
||||
b_in_dtype = dequantize_nvfp4_to_dtype(
|
||||
b_fp4, b_sf, b_global_scale, dtype=dtype, device=device, block_size=block_size
|
||||
)
|
||||
return torch.matmul(a_in_dtype, b_in_dtype.t())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("shape", SHAPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("backend", ["cute-dsl", "cutlass", "cudnn", "trtllm", "b12x"])
|
||||
@pytest.mark.parametrize("autotune", [False, True])
|
||||
@torch.inference_mode()
|
||||
def test_flashinfer_nvfp4_gemm(
|
||||
dtype: torch.dtype,
|
||||
shape: tuple[int, int, int],
|
||||
seed: int,
|
||||
device: str,
|
||||
backend: str,
|
||||
autotune: bool,
|
||||
) -> None:
|
||||
if "trtllm" in backend and dtype == torch.float16:
|
||||
pytest.skip("Only torch.bfloat16 is supported for TRTLLM FP4 GEMM operations")
|
||||
if backend == "cute-dsl" and not current_platform.is_device_capability_family(100):
|
||||
pytest.skip("FlashInfer cutedsl backend is only supported on SM10x")
|
||||
if backend == "b12x" and not current_platform.has_device_capability(120):
|
||||
pytest.skip("b12x FP4 GEMM requires SM120+ (CC 12.0+)")
|
||||
if backend == "b12x" and not has_flashinfer_b12x_gemm():
|
||||
pytest.skip("b12x FP4 GEMM backend not available in installed FlashInfer")
|
||||
|
||||
set_random_seed(seed)
|
||||
m, n, packed_k = shape
|
||||
k = packed_k * 2
|
||||
block_size = 16
|
||||
a_dtype = torch.randn((m, k), dtype=dtype, device=device)
|
||||
b_dtype = torch.randn((n, k), dtype=dtype, device=device)
|
||||
|
||||
a_global_scale = (
|
||||
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1)
|
||||
).to(torch.float32)
|
||||
b_global_scale = (
|
||||
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1)
|
||||
).to(torch.float32)
|
||||
alpha = 1.0 / (a_global_scale * b_global_scale)
|
||||
|
||||
# ops.scaled_fp4_quant returns swizzled scales, while weights
|
||||
# from checkpoints are in linear scales.
|
||||
# cutlass and b12x use swizzled scales directly; trtllm needs them unswizzled.
|
||||
a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(
|
||||
a_dtype, a_global_scale, is_sf_swizzled_layout=True, backend=backend
|
||||
)
|
||||
is_sf_128x4_layout = not (backend == "trtllm" and m <= 32)
|
||||
|
||||
b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(
|
||||
b_dtype, b_global_scale, is_sf_swizzled_layout=True
|
||||
)
|
||||
|
||||
# get_ref_results unswizzles the scales internally.
|
||||
expected_out = get_ref_results(
|
||||
a_fp4,
|
||||
b_fp4,
|
||||
a_scale_interleaved,
|
||||
b_scale_interleaved,
|
||||
a_global_scale,
|
||||
b_global_scale,
|
||||
m,
|
||||
n,
|
||||
dtype,
|
||||
block_size,
|
||||
device,
|
||||
is_sf_128x4_layout,
|
||||
)
|
||||
|
||||
import flashinfer
|
||||
|
||||
if "trtllm" in backend:
|
||||
epilogue_tile_m = 128
|
||||
b_fp4 = flashinfer.shuffle_matrix_a(b_fp4.view(torch.uint8), epilogue_tile_m)
|
||||
b_scale_interleaved = convert_swizzled_to_linear(
|
||||
b_scale_interleaved, n, k, block_size
|
||||
)
|
||||
b_scale_interleaved = (
|
||||
flashinfer.shuffle_matrix_sf_a(
|
||||
b_scale_interleaved.view(torch.uint8), epilogue_tile_m
|
||||
)
|
||||
.reshape(b_scale_interleaved.shape)
|
||||
.view(torch.float8_e4m3fn)
|
||||
)
|
||||
|
||||
with flashinfer.autotune(autotune):
|
||||
out = flashinfer_scaled_fp4_mm(
|
||||
a_fp4,
|
||||
b_fp4,
|
||||
a_scale_interleaved,
|
||||
b_scale_interleaved,
|
||||
alpha,
|
||||
dtype,
|
||||
backend=backend,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(out, expected_out.to(dtype=dtype), atol=1e-1, rtol=1e-1)
|
||||
@@ -0,0 +1,73 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.flashinfer import flashinfer_scaled_fp8_mm
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not current_platform.has_device_capability(100):
|
||||
pytest.skip(
|
||||
reason="Flashinfer FP8 gemms requires compute capability of 10.0 or above.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
# m, n, k
|
||||
SHAPES = [(128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128)]
|
||||
PAD_SHAPES = [(150, 128, 64), (128, 128, 96)]
|
||||
SHAPES.extend(PAD_SHAPES)
|
||||
|
||||
SEEDS = [42]
|
||||
CUDA_DEVICES = ["cuda:0"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("shape", SHAPES)
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@pytest.mark.parametrize("autotune", [False, True])
|
||||
@torch.inference_mode()
|
||||
def test_flashinfer_fp8_gemm(
|
||||
dtype: torch.dtype,
|
||||
shape: tuple[int, int, int],
|
||||
use_bias: bool,
|
||||
seed: int,
|
||||
device: str,
|
||||
autotune: bool,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
m, n, k = shape
|
||||
a = torch.randn((m, k), dtype=dtype, device=device)
|
||||
b = torch.randn((n, k), dtype=dtype, device=device) / k
|
||||
|
||||
a_fp8, a_scale = ops.scaled_fp8_quant(a)
|
||||
b_fp8, b_scale = ops.scaled_fp8_quant(b)
|
||||
|
||||
expected_out = torch.mm(
|
||||
a_scale * a_fp8.to(dtype=torch.float32),
|
||||
b_scale * b_fp8.to(dtype=torch.float32).t(),
|
||||
).to(dtype=dtype)
|
||||
|
||||
if use_bias:
|
||||
bias = torch.randn((n,), dtype=dtype, device=device)
|
||||
expected_out = expected_out + bias
|
||||
else:
|
||||
bias = None
|
||||
|
||||
import flashinfer
|
||||
|
||||
with flashinfer.autotune(autotune):
|
||||
out = flashinfer_scaled_fp8_mm(
|
||||
a_fp8,
|
||||
b_fp8.t(),
|
||||
a_scale,
|
||||
b_scale,
|
||||
dtype,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(out, expected_out, atol=1e-2, rtol=1e-2)
|
||||
@@ -0,0 +1,65 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Unit tests for the get_fp8_min_max() helper function.
|
||||
|
||||
These tests verify the FP8 min/max value logic for both standard
|
||||
and fnuz (ROCm MI300) dtype handling.
|
||||
"""
|
||||
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
get_fp8_min_max,
|
||||
)
|
||||
|
||||
|
||||
class TestGetFp8MinMax:
|
||||
"""Test cases for get_fp8_min_max() function."""
|
||||
|
||||
@patch("vllm.model_executor.layers.quantization.utils.quant_utils.current_platform")
|
||||
def test_standard_fp8_platform(self, mock_platform):
|
||||
"""Test that standard FP8 platform uses PyTorch's finfo values."""
|
||||
mock_platform.is_fp8_fnuz.return_value = False
|
||||
mock_platform.fp8_dtype.return_value = torch.float8_e4m3fn
|
||||
|
||||
fp8_min, fp8_max = get_fp8_min_max()
|
||||
finfo = torch.finfo(torch.float8_e4m3fn)
|
||||
|
||||
# Standard FP8 max is 448.0 for e4m3fn
|
||||
assert fp8_max == finfo.max, f"Expected finfo.max={finfo.max}, got {fp8_max}"
|
||||
assert fp8_min == finfo.min, f"Expected finfo.min={finfo.min}, got {fp8_min}"
|
||||
|
||||
@patch("vllm.model_executor.layers.quantization.utils.quant_utils.current_platform")
|
||||
def test_fnuz_platform_returns_224(self, mock_platform):
|
||||
"""Test that fnuz platform returns 224.0."""
|
||||
mock_platform.is_fp8_fnuz.return_value = True
|
||||
|
||||
fp8_min, fp8_max = get_fp8_min_max()
|
||||
|
||||
# fnuz on ROCm MI300 should return 224.0, not 240.0
|
||||
assert fp8_max == 224.0, f"Expected 224.0 for fnuz platform, got {fp8_max}"
|
||||
assert fp8_min == -224.0, f"Expected -224.0 for fnuz platform, got {fp8_min}"
|
||||
|
||||
@patch("vllm.model_executor.layers.quantization.utils.quant_utils.current_platform")
|
||||
def test_non_fnuz_platform_uses_finfo(self, mock_platform):
|
||||
"""Test that non-fnuz platform uses finfo values."""
|
||||
mock_platform.is_fp8_fnuz.return_value = False
|
||||
mock_platform.fp8_dtype.return_value = torch.float8_e4m3fn
|
||||
|
||||
fp8_min, fp8_max = get_fp8_min_max()
|
||||
finfo = torch.finfo(torch.float8_e4m3fn)
|
||||
|
||||
assert fp8_max == finfo.max, (
|
||||
f"Non-fnuz platform should use finfo.max={finfo.max}, got {fp8_max}"
|
||||
)
|
||||
assert fp8_min == finfo.min, (
|
||||
f"Non-fnuz platform should use finfo.min={finfo.min}, got {fp8_min}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v"])
|
||||
@@ -0,0 +1,221 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
from tests.kernels.quant_utils import (
|
||||
FP8_DTYPE,
|
||||
ref_dynamic_per_tensor_fp8_quant,
|
||||
ref_dynamic_per_token_quant,
|
||||
)
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
scaled_quantize,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
DTYPES = [torch.bfloat16, torch.float]
|
||||
HIDDEN_SIZES = [17, 1024, 1025, 1026, 5137, 8193]
|
||||
NUM_TOKENS = [1, 7, 4096]
|
||||
SCALE_UBS = [True, False]
|
||||
SEEDS = [0]
|
||||
|
||||
|
||||
def opcheck_fp8_quant(
|
||||
output,
|
||||
input,
|
||||
scale=None,
|
||||
scale_ub=None,
|
||||
use_per_token_if_dynamic=False,
|
||||
group_shape=None,
|
||||
):
|
||||
if scale is not None:
|
||||
opcheck(
|
||||
torch.ops._C.static_scaled_fp8_quant,
|
||||
(output, input, scale, group_shape),
|
||||
)
|
||||
elif use_per_token_if_dynamic:
|
||||
scale = torch.empty(
|
||||
(input.shape[0], 1), device=input.device, dtype=torch.float32
|
||||
)
|
||||
opcheck(
|
||||
torch.ops._C.dynamic_per_token_scaled_fp8_quant,
|
||||
(output, input, scale, scale_ub),
|
||||
)
|
||||
else:
|
||||
scale = torch.empty(
|
||||
(input.numel() // input.shape[-1], 1),
|
||||
device=input.device,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
opcheck(torch.ops._C.dynamic_scaled_fp8_quant, (output, input, scale))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("do_scale_ub", SCALE_UBS)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_dynamic_per_token_fp8_quant(
|
||||
num_tokens: int, hidden_size: int, dtype: torch.dtype, do_scale_ub: bool, seed: int
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
|
||||
x = (
|
||||
torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") + 1e-6
|
||||
) # avoid nans
|
||||
|
||||
scale_ub = (
|
||||
torch.mean(x).to(dtype=torch.float32, device="cuda") if do_scale_ub else None
|
||||
)
|
||||
ref_out, ref_scales = ref_dynamic_per_token_quant(x, FP8_DTYPE, scale_ub)
|
||||
ops_out, ops_scales = ops.scaled_fp8_quant(
|
||||
x, scale_ub=scale_ub, use_per_token_if_dynamic=True
|
||||
)
|
||||
|
||||
torch.testing.assert_close(ref_scales, ops_scales)
|
||||
torch.testing.assert_close(
|
||||
ref_out.to(dtype=torch.float32), ops_out.to(dtype=torch.float32)
|
||||
)
|
||||
|
||||
opcheck_fp8_quant(ops_out, x, None, scale_ub, use_per_token_if_dynamic=True)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_dynamic_per_tensor_fp8_quant(
|
||||
num_tokens: int, hidden_size: int, dtype: torch.dtype, seed: int
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
|
||||
x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda")
|
||||
|
||||
ref_out, ref_scale = ref_dynamic_per_tensor_fp8_quant(x)
|
||||
ops_out, ops_scale = ops.scaled_fp8_quant(x)
|
||||
|
||||
torch.testing.assert_close(ref_scale, ops_scale)
|
||||
torch.testing.assert_close(
|
||||
ref_out.to(dtype=torch.float32), ops_out.to(dtype=torch.float32)
|
||||
)
|
||||
|
||||
opcheck_fp8_quant(ops_out, x)
|
||||
|
||||
|
||||
# Regression test for a case with large activations where an int32 index cannot
|
||||
# represent the number of elements.
|
||||
@torch.inference_mode()
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
def test_fp8_quant_large(seed: int) -> None:
|
||||
set_random_seed(seed)
|
||||
|
||||
num_tokens = 1024000 # Mistral-Nemo's max_position_embeddings
|
||||
hidden_size = 1152 # Smallest hidden_size to reproduce the error
|
||||
dtype = torch.bfloat16
|
||||
|
||||
x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda")
|
||||
ref_out, scale = ref_dynamic_per_tensor_fp8_quant(x)
|
||||
ops_out, _ = ops.scaled_fp8_quant(x, scale)
|
||||
|
||||
# Minimize memory footprint in this test by freeing x and upconverting
|
||||
# the outputs in place. (torch.allclose does not support fp8)
|
||||
del x
|
||||
ref_out = ref_out.to(dtype=dtype)
|
||||
ops_out = ops_out.to(dtype=dtype)
|
||||
|
||||
torch.testing.assert_close(ref_out, ops_out)
|
||||
|
||||
|
||||
# Test static FP8 quantization with 2D group scales
|
||||
GROUP_SHAPES_2D = [
|
||||
(-1, -1), # Per-tensor
|
||||
(-1, 1), # Per-channel
|
||||
(1, -1), # Per-token
|
||||
(-1, 128), # Per-head quantization
|
||||
(1, 128), # DeepSeek-style per-token-per-group (group_m=1, group_n=128)
|
||||
(128, 128), # DeepSeek-style block quantization
|
||||
(1, 64), # Smaller group size
|
||||
(1, 16), # Small group (scalar path in kernel)
|
||||
(4, 256), # Non-trivial both dimensions
|
||||
]
|
||||
# Use sizes divisible by all group shapes
|
||||
NUM_TOKENS_GROUP = [128, 512]
|
||||
HIDDEN_SIZES_GROUP = [256, 1024, 2048]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS_GROUP)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES_GROUP)
|
||||
@pytest.mark.parametrize("group_shape", GROUP_SHAPES_2D)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_static_fp8_quant_group_2d(
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
group_shape: tuple[int, int],
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
) -> None:
|
||||
"""Test static FP8 quantization with 2D group scales using scaled_quantize."""
|
||||
# Normalize group_shape (-1 means full extent)
|
||||
norm_group_m = num_tokens if group_shape[0] == -1 else group_shape[0]
|
||||
norm_group_n = hidden_size if group_shape[1] == -1 else group_shape[1]
|
||||
|
||||
# Skip if sizes are not divisible by group shape
|
||||
if num_tokens % norm_group_m != 0 or hidden_size % norm_group_n != 0:
|
||||
pytest.skip(
|
||||
f"Skipping: ({num_tokens}, {hidden_size}) not divisible by "
|
||||
f"group_shape ({group_shape[0]}, {group_shape[1]})"
|
||||
)
|
||||
|
||||
set_random_seed(seed)
|
||||
|
||||
x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda")
|
||||
ref_out, scale = scaled_quantize(
|
||||
x, group_shape, current_platform.fp8_dtype(), compute_dtype=torch.float32
|
||||
)
|
||||
ops_out, ops_scale = ops.scaled_fp8_quant(x, scale=scale, group_shape=group_shape)
|
||||
|
||||
torch.testing.assert_close(scale, ops_scale)
|
||||
torch.testing.assert_close(ref_out.float(), ops_out.float(), rtol=1.2e-1, atol=1e-5)
|
||||
|
||||
opcheck_fp8_quant(ops_out, x, scale=scale)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS_GROUP)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES_GROUP)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("group_shape", [(1, -1), (-1, 1)]) # per-token, per-channel
|
||||
@torch.inference_mode()
|
||||
def test_static_fp8_quant_1d_scale(
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
group_shape: tuple[int, int],
|
||||
) -> None:
|
||||
"""Test static FP8 quantization with 1D scale (per-token or per-channel)."""
|
||||
set_random_seed(seed)
|
||||
|
||||
x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda")
|
||||
ref_out, scale_2d = scaled_quantize(
|
||||
x, group_shape, FP8_DTYPE, compute_dtype=torch.float32
|
||||
)
|
||||
|
||||
# Flatten scale to 1D for testing 1D scale path
|
||||
scale_1d = scale_2d.flatten()
|
||||
ops_out, ops_scale = ops.scaled_fp8_quant(
|
||||
x, scale=scale_1d, group_shape=group_shape
|
||||
)
|
||||
|
||||
torch.testing.assert_close(scale_1d, ops_scale)
|
||||
torch.testing.assert_close(ref_out.float(), ops_out.float(), rtol=0.12, atol=0.0)
|
||||
|
||||
opcheck_fp8_quant(ops_out, x, scale=scale_1d, group_shape=group_shape)
|
||||
@@ -0,0 +1,173 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for QuantFP8 Group Quantization implementation."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"batch_size,hidden_dim,group_size",
|
||||
[
|
||||
(16, 256, 32), # Small
|
||||
(64, 1024, 64), # Medium
|
||||
(128, 2048, 128), # Large
|
||||
(8, 513, 64), # Non-divisible (native only)
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("seed", [42])
|
||||
@pytest.mark.parametrize("use_ue8m0", [True, False])
|
||||
@torch.inference_mode()
|
||||
def test_quantfp8_group_functionality(
|
||||
default_vllm_config,
|
||||
batch_size: int,
|
||||
hidden_dim: int,
|
||||
group_size: int,
|
||||
seed: int,
|
||||
use_ue8m0: bool,
|
||||
) -> None:
|
||||
"""Test QuantFP8 group quantization with various configurations.
|
||||
|
||||
Tests both CUDA and native implementations, column-major scales,
|
||||
and verifies consistency between implementations.
|
||||
"""
|
||||
set_random_seed(seed)
|
||||
|
||||
x = torch.randn((batch_size, hidden_dim), dtype=torch.bfloat16, device="cuda") * 8
|
||||
expected_num_groups = (hidden_dim + group_size - 1) // group_size
|
||||
is_divisible = hidden_dim % group_size == 0
|
||||
|
||||
group_shape = GroupShape(1, group_size)
|
||||
quant_op = QuantFP8(
|
||||
static=False,
|
||||
group_shape=group_shape,
|
||||
column_major_scales=False,
|
||||
use_ue8m0=use_ue8m0,
|
||||
)
|
||||
|
||||
# 1. Test native implementation (always available)
|
||||
x_quant_native, scales_native = quant_op.forward_native(x.clone())
|
||||
assert x_quant_native.shape == x.shape
|
||||
assert scales_native.shape == (batch_size, expected_num_groups)
|
||||
|
||||
# 2. Test column-major scales configuration
|
||||
quant_op_col = QuantFP8(
|
||||
static=False,
|
||||
group_shape=group_shape,
|
||||
column_major_scales=True,
|
||||
use_ue8m0=use_ue8m0,
|
||||
)
|
||||
_, scales_col = quant_op_col.forward_native(x.clone())
|
||||
assert scales_col.shape == (batch_size, expected_num_groups)
|
||||
assert scales_col.stride(0) == 1
|
||||
assert scales_col.stride(1) == batch_size
|
||||
|
||||
# Test column-major scales consistency
|
||||
torch.testing.assert_close(scales_col, scales_native, rtol=1e-9, atol=1e-8)
|
||||
|
||||
# 3. Test CUDA implementation (only for divisible dimensions)
|
||||
if is_divisible:
|
||||
x_quant_cuda, scales_cuda = quant_op.forward_cuda(x.clone())
|
||||
assert x_quant_cuda.shape == x.shape
|
||||
assert scales_cuda.shape == (batch_size, expected_num_groups)
|
||||
|
||||
# Verify CUDA/native consistency
|
||||
torch.testing.assert_close(scales_cuda, scales_native, rtol=2e-7, atol=2e-8)
|
||||
|
||||
# Quantized values should mostly match
|
||||
diff_count = (x_quant_cuda != x_quant_native).sum().item()
|
||||
diff_ratio = diff_count / x_quant_cuda.numel()
|
||||
assert diff_ratio < 0.002, f"Too many differences: {diff_ratio:.4%}"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seed", [42])
|
||||
@pytest.mark.parametrize("use_ue8m0", [True, False])
|
||||
@torch.inference_mode()
|
||||
def test_quantfp8_group_multidimensional(
|
||||
default_vllm_config, seed: int, use_ue8m0: bool
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
|
||||
group_size = 64
|
||||
|
||||
# Test with 3D input
|
||||
batch1, batch2, hidden_dim = 4, 8, 1024
|
||||
x_3d = (
|
||||
torch.randn((batch1, batch2, hidden_dim), dtype=torch.bfloat16, device="cuda")
|
||||
* 8
|
||||
)
|
||||
|
||||
group_shape = GroupShape(1, group_size)
|
||||
quant_op = QuantFP8(
|
||||
static=False,
|
||||
group_shape=group_shape,
|
||||
column_major_scales=False,
|
||||
use_ue8m0=use_ue8m0,
|
||||
)
|
||||
|
||||
x_quant, scales = quant_op.forward_native(x_3d.clone())
|
||||
assert x_quant.shape == x_3d.shape
|
||||
assert scales.shape == (batch1, batch2, hidden_dim // group_size)
|
||||
|
||||
# Test column_major_scales with multi-dim
|
||||
quant_op_col = QuantFP8(
|
||||
static=False,
|
||||
group_shape=group_shape,
|
||||
column_major_scales=True,
|
||||
use_ue8m0=use_ue8m0,
|
||||
)
|
||||
_, scales_col = quant_op_col.forward_native(x_3d.clone())
|
||||
assert scales_col.shape == (batch1, batch2, hidden_dim // group_size)
|
||||
|
||||
# Test with 4D input
|
||||
batch1, batch2, batch3, hidden_dim = 2, 3, 4, 256
|
||||
x_4d = (
|
||||
torch.randn(
|
||||
(batch1, batch2, batch3, hidden_dim), dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
* 8
|
||||
)
|
||||
|
||||
x_quant_4d, scales_4d = quant_op.forward_native(x_4d.clone())
|
||||
assert x_quant_4d.shape == x_4d.shape
|
||||
assert scales_4d.shape == (batch1, batch2, batch3, hidden_dim // group_size)
|
||||
|
||||
_, scales_4d_col = quant_op_col.forward_native(x_4d.clone())
|
||||
assert scales_4d_col.shape == (batch1, batch2, hidden_dim // group_size, batch3)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("seed", [42])
|
||||
@torch.inference_mode()
|
||||
def test_quantfp8_group_edge_cases(default_vllm_config, seed: int) -> None:
|
||||
set_random_seed(seed)
|
||||
|
||||
batch_size = 16
|
||||
group_size = 64
|
||||
|
||||
# Test with single group (group_size >= hidden_dim)
|
||||
x_small = torch.randn((batch_size, 32), dtype=torch.bfloat16, device="cuda") * 8
|
||||
group_shape = GroupShape(1, group_size)
|
||||
quant_op = QuantFP8(
|
||||
static=False, group_shape=group_shape, column_major_scales=False
|
||||
)
|
||||
|
||||
x_quant_small, scales_small = quant_op.forward_native(x_small.clone())
|
||||
assert x_quant_small.shape == x_small.shape
|
||||
assert scales_small.shape == (batch_size, 1)
|
||||
|
||||
# Test with zero inputs
|
||||
x_zero = torch.zeros((batch_size, 256), dtype=torch.bfloat16, device="cuda")
|
||||
x_quant_zero, scales_zero = quant_op.forward_native(x_zero.clone())
|
||||
assert x_quant_zero.shape == x_zero.shape
|
||||
assert (scales_zero > 0).all(), "Scales should be clamped to minimum"
|
||||
|
||||
# Test very large values
|
||||
x_large = torch.full((batch_size, 256), 1000.0, dtype=torch.bfloat16, device="cuda")
|
||||
x_quant_large, scales_large = quant_op.forward_native(x_large.clone())
|
||||
assert x_quant_large.shape == x_large.shape
|
||||
# FP8 max is typically 448 or 224, so scales should be > 1
|
||||
assert (scales_large > 1.0).all(), "Large values should have scales > 1"
|
||||
@@ -0,0 +1,35 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm import _custom_ops as ops # noqa: F401
|
||||
|
||||
|
||||
def test_gptq_shuffle_opcheck():
|
||||
weight = torch.randint(
|
||||
-2000000, 2000000, (1792, 4096), device="cuda", dtype=torch.int32
|
||||
)
|
||||
perm = torch.empty((0,), device="cuda", dtype=torch.int32)
|
||||
bit = 4
|
||||
opcheck(torch.ops._C.gptq_shuffle, (weight, perm, bit))
|
||||
|
||||
|
||||
def test_gptq_gemm_opcheck():
|
||||
a = torch.rand((240, 4096), device="cuda", dtype=torch.float16)
|
||||
weight = torch.randint(
|
||||
-2000000, 2000000, (512, 6144), device="cuda", dtype=torch.int32
|
||||
)
|
||||
zeros = torch.zeros((32, 768), device="cuda", dtype=torch.int32)
|
||||
scales = torch.rand((32, 6144), device="cuda", dtype=torch.float16)
|
||||
idx = torch.empty((0,), device="cuda", dtype=torch.int32)
|
||||
use_exllama = True
|
||||
bit = 4
|
||||
# Test both GPTQv1 and GPTQv2 format
|
||||
opcheck(
|
||||
torch.ops._C.gptq_gemm, (a, weight, zeros, scales, idx, use_exllama, True, bit)
|
||||
)
|
||||
opcheck(
|
||||
torch.ops._C.gptq_gemm, (a, weight, zeros, scales, idx, use_exllama, False, bit)
|
||||
)
|
||||
@@ -0,0 +1,33 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import math
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from compressed_tensors.transform import deterministic_hadamard_matrix
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if current_platform.is_rocm():
|
||||
pytest.skip(
|
||||
"These tests require hadacore_transform, not supported on ROCm.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch_size", [1, 32])
|
||||
@pytest.mark.parametrize("hidden_dim", [2**n for n in range(10)])
|
||||
def test_hadacore(batch_size, hidden_dim, dtype=torch.bfloat16, device="cuda"):
|
||||
x = torch.eye(hidden_dim, dtype=dtype, device=device)
|
||||
hadamard = deterministic_hadamard_matrix(
|
||||
hidden_dim, dtype=torch.float64, device="cuda"
|
||||
) / math.sqrt(hidden_dim)
|
||||
|
||||
y = ops.hadacore_transform(x.clone())
|
||||
y_true = (x.to(hadamard.dtype) @ hadamard.T).to(y.dtype)
|
||||
assert torch.allclose(y, y_true)
|
||||
|
||||
y = ops.hadacore_transform(y)
|
||||
assert torch.allclose(y, x)
|
||||
@@ -0,0 +1,155 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
# Adapted from https://github.com/sgl-project/sglang/blob/main/test/srt/test_int8_kernel.py
|
||||
import itertools
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.activation import SiluAndMul
|
||||
from vllm.model_executor.layers.fused_moe import fused_experts
|
||||
from vllm.model_executor.layers.fused_moe.config import FusedMoEQuantConfig
|
||||
from vllm.model_executor.layers.quantization.utils.int8_utils import (
|
||||
per_token_quant_int8,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if current_platform.get_device_capability() < (7, 0):
|
||||
pytest.skip("INT8 Triton requires CUDA 7.0 or higher", allow_module_level=True)
|
||||
|
||||
|
||||
def native_w8a8_per_token_matmul(A, B, As, Bs, output_dtype=torch.float16):
|
||||
"""Matrix multiplication function that supports per-token input
|
||||
quantization and per-column weight quantization"""
|
||||
A = A.to(torch.float32)
|
||||
B = B.to(torch.float32)
|
||||
|
||||
assert A.shape[-1] == B.shape[-1], "Dimension mismatch"
|
||||
assert B.ndim == 2 and B.is_contiguous(), "B must be a 2D contiguous tensor"
|
||||
|
||||
# Reshape input
|
||||
M = A.numel() // A.shape[-1]
|
||||
B = B.t() # Transpose weight matrix
|
||||
N, K = B.shape
|
||||
origin_C_shape = A.shape[:-1] + (K,)
|
||||
A = A.reshape(M, N)
|
||||
|
||||
# As is per-token [M, 1], Bs is per-column [1, K]
|
||||
C = torch.matmul(A, B) # [M, K]
|
||||
C = As * C * Bs.view(1, -1) # Broadcast per-column scale
|
||||
|
||||
return C.reshape(origin_C_shape).to(output_dtype)
|
||||
|
||||
|
||||
def torch_w8a8_per_column_moe(a, w1, w2, w1_s, w2_s, topk, topk_weight, topk_ids):
|
||||
"""This function performs fused moe with per-column int8 quantization
|
||||
using native torch."""
|
||||
|
||||
B, D = a.shape
|
||||
# Perform per-token quantization
|
||||
a_q, a_s = per_token_quant_int8(a)
|
||||
# Repeat tokens to match topk
|
||||
a_q = a_q.view(B, -1, D).repeat(1, topk, 1).reshape(-1, D)
|
||||
# Also repeat the scale
|
||||
a_s = a_s.view(B, -1, 1).repeat(1, topk, 1).reshape(-1, 1) # [B*topk, 1]
|
||||
|
||||
out = torch.zeros(B * topk, w2.shape[1], dtype=a.dtype, device=a.device)
|
||||
|
||||
# Calculate routing
|
||||
topk_weight = topk_weight.view(-1)
|
||||
topk_ids = topk_ids.view(-1)
|
||||
# Process each expert
|
||||
for i in range(w1.shape[0]):
|
||||
mask = topk_ids == i
|
||||
if mask.sum():
|
||||
# First MLP layer: note that a_s is now per-token
|
||||
inter_out = native_w8a8_per_token_matmul(
|
||||
a_q[mask], w1[i], a_s[mask], w1_s[i], output_dtype=a.dtype
|
||||
)
|
||||
# Activation function
|
||||
act_out = SiluAndMul().forward_native(inter_out)
|
||||
# Quantize activation output with per-token
|
||||
act_out_q, act_out_s = per_token_quant_int8(act_out)
|
||||
|
||||
# Second MLP layer
|
||||
out[mask] = native_w8a8_per_token_matmul(
|
||||
act_out_q, w2[i], act_out_s, w2_s[i], output_dtype=a.dtype
|
||||
)
|
||||
# Apply routing weights and sum
|
||||
return (
|
||||
out.view(B, -1, w2.shape[1]) * topk_weight.view(B, -1, 1).to(out.dtype)
|
||||
).sum(dim=1)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True, scope="module")
|
||||
def setup_cuda():
|
||||
"""Sets the default CUDA device for all tests in this module."""
|
||||
torch.set_default_device("cuda")
|
||||
|
||||
|
||||
DTYPES = [torch.half, torch.bfloat16]
|
||||
M = [1, 33]
|
||||
N = [128, 1024]
|
||||
K = [256, 4096]
|
||||
E = [8]
|
||||
TOP_KS = [2, 6]
|
||||
SEEDS = [0]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"M, N, K, E, topk, dtype, seed",
|
||||
itertools.product(M, N, K, E, TOP_KS, DTYPES, SEEDS),
|
||||
)
|
||||
@torch.inference_mode()
|
||||
def test_w8a8_fp8_fused_moe(default_vllm_config, M, N, K, E, topk, dtype, seed):
|
||||
torch.manual_seed(seed)
|
||||
# Initialize int8 quantization parameters
|
||||
factor_for_scale = 1e-2
|
||||
int8_max = 127
|
||||
int8_min = -128
|
||||
|
||||
# Input tensor
|
||||
# M * K
|
||||
a = torch.randn((M, K), dtype=dtype) / 10
|
||||
|
||||
# Generate int8 weights
|
||||
w1_fp32 = (torch.rand((E, 2 * N, K), dtype=torch.float32) - 0.5) * 2
|
||||
w1 = (w1_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
|
||||
|
||||
w2_fp32 = (torch.rand((E, K, N), dtype=torch.float32) - 0.5) * 2
|
||||
w2 = (w2_fp32 * int8_max).clamp(min=int8_min, max=int8_max).to(torch.int8)
|
||||
|
||||
# Generate scale for each column (per-column quantization)
|
||||
w1_s = torch.rand(E, 2 * N, device=w1_fp32.device) * factor_for_scale
|
||||
w2_s = torch.rand(E, K, device=w2_fp32.device) * factor_for_scale
|
||||
score = torch.randn((M, E), dtype=dtype)
|
||||
score = torch.softmax(score, dim=-1, dtype=torch.float32)
|
||||
topk_weights, topk_ids = torch.topk(score, topk)
|
||||
|
||||
ref_out = torch_w8a8_per_column_moe(
|
||||
a, w1, w2, w1_s, w2_s, topk, topk_weights, topk_ids
|
||||
)
|
||||
|
||||
quant_config = FusedMoEQuantConfig.make(
|
||||
torch.int8,
|
||||
per_act_token_quant=True,
|
||||
block_shape=None,
|
||||
w1_scale=w1_s,
|
||||
w2_scale=w2_s,
|
||||
)
|
||||
|
||||
out = fused_experts(
|
||||
a,
|
||||
w1,
|
||||
w2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
quant_config=quant_config,
|
||||
)
|
||||
|
||||
# Check results
|
||||
rel_diff = torch.mean(
|
||||
torch.abs(out.to(torch.float32) - ref_out.to(torch.float32))
|
||||
) / torch.mean(torch.abs(ref_out.to(torch.float32)))
|
||||
assert rel_diff < 0.05
|
||||
@@ -0,0 +1,195 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.quant_utils import ref_dynamic_per_token_quant
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm._custom_ops import scaled_int8_quant
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
DTYPES = [torch.bfloat16, torch.float]
|
||||
HIDDEN_SIZES = [17, 1024, 1025, 1026, 5137, 8193]
|
||||
NUM_TOKENS = [1, 7, 4096]
|
||||
SEEDS = [0]
|
||||
SCALE = [0.1, 2.1]
|
||||
|
||||
|
||||
def opcheck_int8_quant_static(output, input, scale, azp=None):
|
||||
if azp is None:
|
||||
opcheck(torch.ops._C.static_scaled_int8_quant, (output, input, scale, None))
|
||||
else:
|
||||
opcheck(torch.ops._C.static_scaled_int8_quant, (output, input, scale, azp))
|
||||
|
||||
|
||||
def opcheck_int8_quant_dynamic(output, input, symmetric=True):
|
||||
scale = torch.empty(
|
||||
(input.numel() // input.shape[-1], 1), device=input.device, dtype=torch.float32
|
||||
)
|
||||
if symmetric:
|
||||
opcheck(torch.ops._C.dynamic_scaled_int8_quant, (output, input, scale, None))
|
||||
else:
|
||||
azp = torch.empty(
|
||||
(input.numel() // input.shape[-1], 1),
|
||||
device=input.device,
|
||||
dtype=torch.int32,
|
||||
)
|
||||
opcheck(torch.ops._C.dynamic_scaled_int8_quant, (output, input, scale, azp))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_dynamic_scaled_int8_quant(
|
||||
num_tokens: int, hidden_size: int, dtype: torch.dtype, seed: int
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
|
||||
x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") * 1000
|
||||
|
||||
# reference
|
||||
ref_out, ref_scales = ref_dynamic_per_token_quant(x, torch.int8)
|
||||
# kernel
|
||||
ops_out, ops_scales, _ = scaled_int8_quant(x)
|
||||
|
||||
torch.testing.assert_close(ops_scales, ref_scales)
|
||||
# big atol to account for rounding errors
|
||||
torch.testing.assert_close(ops_out, ref_out, atol=1, rtol=0.0)
|
||||
|
||||
opcheck_int8_quant_dynamic(ops_out, x)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@torch.inference_mode()
|
||||
def test_dynamic_scaled_int8_azp_quant(
|
||||
num_tokens: int, hidden_size: int, dtype: torch.dtype, seed: int
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
int8_traits = torch.iinfo(torch.int8)
|
||||
|
||||
x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") * 1000 - 300
|
||||
|
||||
x_token_max, _ = x.to(dtype=torch.float32).max(dim=1, keepdim=True)
|
||||
x_token_min, _ = x.to(dtype=torch.float32).min(dim=1, keepdim=True)
|
||||
|
||||
# calculate scale and azp, and adjust the range
|
||||
scales = (x_token_max - x_token_min) / torch.tensor(255.0)
|
||||
azps = torch.round(torch.tensor(-128.0) - x_token_min / scales).to(torch.int32)
|
||||
|
||||
torch_out = (
|
||||
((x / scales).round() + azps)
|
||||
.clamp(int8_traits.min, int8_traits.max)
|
||||
.to(torch.int8)
|
||||
)
|
||||
assert torch_out.min() >= int8_traits.min and torch_out.max() <= int8_traits.max
|
||||
|
||||
ops_out, scales_out, azp_out = scaled_int8_quant(x, symmetric=False)
|
||||
|
||||
if not torch.allclose(scales_out, scales):
|
||||
print(torch.argmax(torch.abs(scales_out - scales)))
|
||||
torch.testing.assert_close(scales_out, scales)
|
||||
# big atol to account for rounding errors
|
||||
torch.testing.assert_close(azp_out, azps, atol=1, rtol=0.0)
|
||||
# if AZP is off by 1, after rounding-to-even, the output may be off by 2
|
||||
torch.testing.assert_close(ops_out, torch_out, atol=2, rtol=0.0)
|
||||
|
||||
opcheck_int8_quant_dynamic(ops_out, x, False)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("scale", SCALE)
|
||||
@torch.inference_mode()
|
||||
def test_static_scaled_int8_quant(
|
||||
num_tokens: int, hidden_size: int, dtype: torch.dtype, seed: int, scale: float
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
int8_traits = torch.iinfo(torch.int8)
|
||||
|
||||
x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") * 1000
|
||||
scale_arg = torch.tensor([scale], dtype=torch.float32, device="cuda")
|
||||
|
||||
out1 = (
|
||||
(x / scale_arg).round().clamp(int8_traits.min, int8_traits.max).to(torch.int8)
|
||||
)
|
||||
out2, scale2, _ = scaled_int8_quant(x, scale_arg)
|
||||
assert scale2 is scale_arg
|
||||
|
||||
# big atol to account for rounding errors
|
||||
torch.testing.assert_close(out1, out2, atol=1, rtol=0.0)
|
||||
|
||||
opcheck_int8_quant_static(out2, x, scale_arg)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("scale", SCALE)
|
||||
@pytest.mark.parametrize("azp", [-255, 54])
|
||||
@torch.inference_mode()
|
||||
def test_static_scaled_int8_azp_quant(
|
||||
num_tokens: int,
|
||||
hidden_size: int,
|
||||
dtype: torch.dtype,
|
||||
seed: int,
|
||||
scale: float,
|
||||
azp: int,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
int8_traits = torch.iinfo(torch.int8)
|
||||
|
||||
x = torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") * 1000 - 300
|
||||
|
||||
out1 = (
|
||||
((x / scale).round() + azp)
|
||||
.clamp(int8_traits.min, int8_traits.max)
|
||||
.to(torch.int8)
|
||||
)
|
||||
scale_arg = torch.tensor([scale], dtype=torch.float32, device="cuda")
|
||||
azp_arg = torch.tensor([azp], dtype=torch.int32, device="cuda")
|
||||
|
||||
out2, scale2, azp2 = scaled_int8_quant(x, scale_arg, azp_arg, symmetric=False)
|
||||
assert scale2 is scale_arg
|
||||
assert azp2 is azp_arg
|
||||
|
||||
# big atol to account for rounding errors
|
||||
torch.testing.assert_close(out1, out2, atol=1, rtol=0.0)
|
||||
|
||||
opcheck_int8_quant_static(out2, x, scale_arg, azp_arg)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("is_max", [True, False])
|
||||
@torch.inference_mode()
|
||||
def test_static_scaled_int8_azp_quant_saturating_cast(is_max: bool) -> None:
|
||||
# Test that the saturating cast works correctly for values near i32 max/min
|
||||
|
||||
from numpy import inf, nextafter
|
||||
|
||||
int32_traits = torch.iinfo(torch.int32)
|
||||
val = float(int32_traits.max if is_max else int32_traits.min)
|
||||
|
||||
x_vals = [[nextafter(val, inf), val + 1, val, val - 1, nextafter(val, -inf)]]
|
||||
x = torch.tensor(x_vals, dtype=torch.float32, device="cuda")
|
||||
|
||||
# The calculation in the kernel is: cast<int8>(cast<int32>(x / scale) + azp)
|
||||
# where cast<T> is a saturating cast to type T.
|
||||
# Scale is set to 1.0 so that the input values are the ones that are cast.
|
||||
# AZP is set to 0 to make sure the int8 saturating cast is tested as well.
|
||||
scale = torch.scalar_tensor(1.0, dtype=torch.float32, device="cuda")
|
||||
azp = torch.scalar_tensor(0, dtype=torch.int32, device="cuda")
|
||||
|
||||
int8_traits = torch.iinfo(torch.int8)
|
||||
val_i8 = int8_traits.max if is_max else int8_traits.min
|
||||
expected = torch.full((1, 5), val_i8, dtype=torch.int8, device="cuda")
|
||||
|
||||
out, _, _ = scaled_int8_quant(x, scale, azp, symmetric=False)
|
||||
torch.testing.assert_close(expected, out, atol=0, rtol=0)
|
||||
@@ -0,0 +1,449 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the machete kernel.
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_machete_mm.py`.
|
||||
"""
|
||||
|
||||
import math
|
||||
from dataclasses import dataclass, fields
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.model_executor.layers.quantization.utils.machete_utils import (
|
||||
query_machete_supported_group_sizes,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
pack_rows,
|
||||
quantize_weights,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.scalar_type import ScalarType, scalar_types
|
||||
|
||||
if current_platform.is_rocm():
|
||||
pytest.skip(
|
||||
"These tests require machete_prepack_B, not supported on ROCm.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
CUDA_DEVICES = [
|
||||
f"cuda:{i}" for i in range(1 if torch.accelerator.device_count() == 1 else 2)
|
||||
]
|
||||
|
||||
# TODO: in future PR refactor this and `is_quant_method_supported` in the kernel
|
||||
# unit tests to a common utility function. Currently the use of
|
||||
# `is_quant_method_supported` conflates kernels with quantization methods
|
||||
# an assumption which is breaking down as quantizations methods can have
|
||||
# have kernels and some kernels support multiple quantization methods.
|
||||
IS_SUPPORTED_BY_GPU = current_platform.get_device_capability()[0] >= 9
|
||||
|
||||
MNK_SHAPES = [
|
||||
(1, 128, 128),
|
||||
(1, 8192, 28672),
|
||||
(13, 8192, 4096),
|
||||
(26, 4096, 8192),
|
||||
(64, 4096, 4096),
|
||||
(64, 8192, 28672),
|
||||
(257, 128, 4096),
|
||||
(257, 4224, 4160),
|
||||
(1024, 8192, 4096),
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class TypeConfig:
|
||||
act_type: torch.dtype
|
||||
weight_type: ScalarType
|
||||
output_type: torch.dtype | None
|
||||
group_scale_type: torch.dtype | None
|
||||
group_zero_type: torch.dtype | None
|
||||
channel_scale_type: torch.dtype | None
|
||||
token_scale_type: torch.dtype | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class Tensors:
|
||||
w_ref: torch.Tensor
|
||||
a_ref: torch.Tensor
|
||||
a: torch.Tensor
|
||||
w_q: torch.Tensor
|
||||
w_g_s: torch.Tensor | None
|
||||
w_g_zp: torch.Tensor | None
|
||||
w_ch_s: torch.Tensor | None
|
||||
w_tok_s: torch.Tensor | None
|
||||
|
||||
|
||||
# (Act Type, Weight Type, Output Type, Scale Type, ZeroPoints,
|
||||
# Ch Scales Type, Tok Scales Type)
|
||||
# NOTE: None "Scale Type" means the act type is floating point
|
||||
# None "Output Type" means the output type is the same as the act type
|
||||
TestTypeTuple = tuple[
|
||||
list[torch.dtype], ScalarType, torch.dtype | None, torch.dtype | None, bool
|
||||
]
|
||||
TEST_TYPES = [
|
||||
# GPTQ style
|
||||
*(
|
||||
TypeConfig(
|
||||
act_type=a_type,
|
||||
weight_type=w_type,
|
||||
output_type=None,
|
||||
group_scale_type=a_type,
|
||||
group_zero_type=None,
|
||||
channel_scale_type=None,
|
||||
token_scale_type=None,
|
||||
)
|
||||
for w_type in [scalar_types.uint4b8, scalar_types.uint8b128]
|
||||
for a_type in [torch.float16, torch.bfloat16]
|
||||
),
|
||||
# AWQ style
|
||||
*(
|
||||
TypeConfig(
|
||||
act_type=a_type,
|
||||
weight_type=w_type,
|
||||
output_type=None,
|
||||
group_scale_type=a_type,
|
||||
group_zero_type=a_type,
|
||||
channel_scale_type=None,
|
||||
token_scale_type=None,
|
||||
)
|
||||
for w_type in [scalar_types.uint4, scalar_types.uint8]
|
||||
for a_type in [torch.float16, torch.bfloat16]
|
||||
),
|
||||
# # QQQ style
|
||||
# *(TypeConfig(act_type=torch.int8,
|
||||
# weight_type=scalar_types.uint4b8,
|
||||
# output_type=torch.float16,
|
||||
# group_scale_type=group_scale_type,
|
||||
# group_zero_type=None,
|
||||
# channel_scale_type=torch.float,
|
||||
# token_scale_type=torch.float)
|
||||
# for group_scale_type in [None, torch.float16]),
|
||||
# *(TypeConfig(act_type=torch.float8_e4m3fn,
|
||||
# weight_type=scalar_types.uint4b8,
|
||||
# output_type=torch.float16,
|
||||
# group_scale_type=group_scale_type,
|
||||
# group_zero_type=None,
|
||||
# channel_scale_type=torch.float,
|
||||
# token_scale_type=torch.float)
|
||||
# for group_scale_type in [None, torch.float16]),
|
||||
]
|
||||
|
||||
# TODO: in future PR refactor this and `is_quant_method_supported` in the kernel
|
||||
# unit tests to a common utility function. Currently the use of
|
||||
# `is_quant_method_supported` conflates kernels with quantization methods
|
||||
# an assumption which is breaking down as quantizations methods can have
|
||||
# have kernels and some kernels support multiple quantization methods.
|
||||
IS_SUPPORTED_BY_GPU = current_platform.has_device_capability(90)
|
||||
|
||||
|
||||
def rand_data(shape, dtype=torch.float16, scale=1, offset=0):
|
||||
if dtype.is_floating_point:
|
||||
return (scale * torch.rand(shape, device="cuda") - offset).to(dtype)
|
||||
else:
|
||||
return torch.randint(-8, 7, shape, dtype=dtype, device="cuda")
|
||||
|
||||
|
||||
def maybe_convert_zeropoints(zps: torch.Tensor | None, s: torch.Tensor):
|
||||
return zps if zps is None else -1 * s * (zps.to(s.dtype))
|
||||
|
||||
|
||||
def group_size_valid(shape: tuple[int, int, int], group_size: int | None) -> bool:
|
||||
return group_size is None or group_size == -1 or shape[2] % group_size == 0
|
||||
|
||||
|
||||
def machete_quantize_and_pack(
|
||||
atype: torch.dtype,
|
||||
w: torch.Tensor,
|
||||
wtype: ScalarType,
|
||||
stype: torch.dtype | None,
|
||||
group_size: int | None,
|
||||
zero_points: bool = False,
|
||||
):
|
||||
assert wtype.is_integer(), "TODO: support floating point weights"
|
||||
|
||||
w_ref, w_q, w_s, w_zp = quantize_weights(
|
||||
w,
|
||||
wtype,
|
||||
group_size=group_size,
|
||||
zero_points=zero_points,
|
||||
# to match how the kernel applies zps
|
||||
ref_zero_points_after_scales=True,
|
||||
)
|
||||
|
||||
w_q = pack_rows(w_q, wtype.size_bits, *w_q.shape)
|
||||
w_q = w_q.t().contiguous().t() # convert to col major
|
||||
|
||||
w_q_machete = ops.machete_prepack_B(w_q, atype, wtype, stype)
|
||||
opcheck(torch.ops._C.machete_prepack_B, (w_q, atype, wtype.id, stype))
|
||||
|
||||
return w_ref, w_q_machete, w_s, w_zp
|
||||
|
||||
|
||||
def create_test_tensors(
|
||||
shape: tuple[int, int, int],
|
||||
types: TypeConfig,
|
||||
group_size: int | None,
|
||||
subset_stride_factor: int | None = None,
|
||||
) -> Tensors:
|
||||
m, n, k = shape
|
||||
factor = subset_stride_factor or 1
|
||||
|
||||
print(
|
||||
"create_test_tensors, shape:", shape, "types:", types, "group_size:", group_size
|
||||
)
|
||||
|
||||
a = rand_data((m * factor, k * factor), types.act_type, scale=3, offset=2)
|
||||
w = rand_data((k * factor, n * factor), types.act_type, scale=3, offset=1)
|
||||
|
||||
if factor > 1:
|
||||
a = a[0:m, 0:k]
|
||||
w = w[0:k, 0:n]
|
||||
|
||||
if types.group_scale_type is not None:
|
||||
w = w.to(types.group_scale_type)
|
||||
if w.dtype.itemsize == 1:
|
||||
w = w.to(torch.float16)
|
||||
|
||||
w_ref, w_q_packed, w_s, w_zp = machete_quantize_and_pack(
|
||||
a.dtype,
|
||||
w,
|
||||
types.weight_type,
|
||||
types.group_scale_type,
|
||||
group_size,
|
||||
types.group_zero_type is not None,
|
||||
)
|
||||
|
||||
if not a.dtype.is_floating_point:
|
||||
aiinfo = torch.iinfo(a.dtype)
|
||||
w_ref = w_ref.round().clamp(aiinfo.min, aiinfo.max)
|
||||
|
||||
a_ref = a.to(torch.float32)
|
||||
w_ref = w_ref.to(torch.float32)
|
||||
|
||||
w_ch_s = (
|
||||
None
|
||||
if types.channel_scale_type is None
|
||||
else rand_data((n,), types.channel_scale_type)
|
||||
)
|
||||
w_tok_s = (
|
||||
None
|
||||
if types.token_scale_type is None
|
||||
else rand_data((m,), types.token_scale_type)
|
||||
)
|
||||
|
||||
return Tensors(
|
||||
w_ref=w_ref,
|
||||
a_ref=a_ref,
|
||||
a=a,
|
||||
w_q=w_q_packed,
|
||||
w_g_s=w_s,
|
||||
w_g_zp=maybe_convert_zeropoints(w_zp, w_s),
|
||||
w_ch_s=w_ch_s,
|
||||
w_tok_s=w_tok_s,
|
||||
)
|
||||
|
||||
|
||||
# None stype means scales use the same dtype as a
|
||||
def machete_mm_test_helper(
|
||||
types: TypeConfig,
|
||||
tensors: Tensors,
|
||||
group_size: int | None = None,
|
||||
schedule: str | None = None,
|
||||
):
|
||||
output_ref = torch.matmul(tensors.a_ref, tensors.w_ref)
|
||||
output_ref_type = output_ref.dtype
|
||||
|
||||
if tensors.w_ch_s is not None:
|
||||
output_ref = (
|
||||
output_ref.to(tensors.w_ch_s.dtype) * tensors.w_ch_s.unsqueeze(0)
|
||||
).to(output_ref_type)
|
||||
if tensors.w_tok_s is not None:
|
||||
output_ref = (
|
||||
output_ref.to(tensors.w_tok_s.dtype) * tensors.w_tok_s.unsqueeze(1)
|
||||
).to(output_ref_type)
|
||||
|
||||
output = ops.machete_mm(
|
||||
a=tensors.a,
|
||||
b_q=tensors.w_q,
|
||||
b_type=types.weight_type,
|
||||
b_group_scales=tensors.w_g_s,
|
||||
b_group_zeros=tensors.w_g_zp,
|
||||
b_group_size=group_size,
|
||||
b_channel_scales=tensors.w_ch_s,
|
||||
a_token_scales=tensors.w_tok_s,
|
||||
out_type=types.output_type,
|
||||
schedule=schedule,
|
||||
)
|
||||
|
||||
print(output)
|
||||
print(output_ref)
|
||||
|
||||
# Relax atol as our reduction dim becomes larger (more rounding error)
|
||||
# Relax atol when we have zeropoints since the way machete applies
|
||||
# zeropoints (after scales) causes noise around 0
|
||||
atol = (
|
||||
1
|
||||
if tensors.w_g_zp is not None
|
||||
else min(5e-2 * math.sqrt(tensors.a.shape[1]), 1)
|
||||
)
|
||||
rtol = 1e-1 if tensors.a.element_size() >= 2 else 2e-1
|
||||
torch.testing.assert_close(
|
||||
output, output_ref.to(output.dtype), rtol=rtol, atol=atol
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not IS_SUPPORTED_BY_GPU, reason="Machete is not supported on this GPU type."
|
||||
)
|
||||
@pytest.mark.parametrize("shape", MNK_SHAPES, ids=lambda x: "x".join(str(v) for v in x))
|
||||
@pytest.mark.parametrize("types", TEST_TYPES)
|
||||
def test_machete_all_schedules(shape, types: TypeConfig):
|
||||
group_sizes: list[int | None] = []
|
||||
if types.group_scale_type is None:
|
||||
group_sizes = [None]
|
||||
else:
|
||||
group_sizes = query_machete_supported_group_sizes(types.act_type)
|
||||
|
||||
for group_size in group_sizes:
|
||||
if not group_size_valid(shape, group_size):
|
||||
continue
|
||||
|
||||
tensors = create_test_tensors(shape, types, group_size)
|
||||
print(f"MNK = {shape}")
|
||||
for schedule in ops.machete_supported_schedules(
|
||||
types.act_type,
|
||||
types.weight_type,
|
||||
group_scales_type=types.group_scale_type,
|
||||
group_zeros_type=types.group_scale_type,
|
||||
out_type=types.output_type,
|
||||
):
|
||||
print(f"Testing schedule {schedule}")
|
||||
machete_mm_test_helper(types, tensors, group_size, schedule)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not IS_SUPPORTED_BY_GPU, reason="Machete is not supported on this GPU type."
|
||||
)
|
||||
@pytest.mark.parametrize("shape", MNK_SHAPES, ids=lambda x: "x".join(str(v) for v in x))
|
||||
@pytest.mark.parametrize("types", TEST_TYPES)
|
||||
def test_machete_heuristic(shape, types: TypeConfig):
|
||||
group_sizes: list[int | None] = []
|
||||
if types.group_scale_type is None:
|
||||
group_sizes = [None]
|
||||
else:
|
||||
group_sizes = query_machete_supported_group_sizes(types.act_type)
|
||||
|
||||
for group_size in group_sizes:
|
||||
if not group_size_valid(shape, group_size):
|
||||
continue
|
||||
|
||||
tensors = create_test_tensors(shape, types, group_size)
|
||||
machete_mm_test_helper(types, tensors, group_size)
|
||||
|
||||
|
||||
# Test working on other devices
|
||||
@pytest.mark.skipif(
|
||||
not IS_SUPPORTED_BY_GPU, reason="Machete is not supported on this GPU type."
|
||||
)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
def test_machete_devices(device: str):
|
||||
group_size = 128
|
||||
|
||||
type_config = TypeConfig(
|
||||
act_type=torch.float16,
|
||||
weight_type=scalar_types.uint4b8,
|
||||
output_type=None,
|
||||
group_scale_type=torch.float16,
|
||||
group_zero_type=None,
|
||||
channel_scale_type=None,
|
||||
token_scale_type=None,
|
||||
)
|
||||
|
||||
tensors = create_test_tensors((512, 4096, 4096), type_config, group_size)
|
||||
|
||||
for field in fields(Tensors):
|
||||
tensor = getattr(tensors, field.name)
|
||||
if isinstance(tensor, torch.Tensor):
|
||||
setattr(tensors, field.name, tensor.to(device))
|
||||
|
||||
machete_mm_test_helper(type_config, tensors, group_size)
|
||||
|
||||
|
||||
# Test working with a subset of A and B
|
||||
@pytest.mark.skipif(
|
||||
not IS_SUPPORTED_BY_GPU, reason="Machete is not supported on this GPU type."
|
||||
)
|
||||
def test_machete_subset():
|
||||
group_size = 128
|
||||
|
||||
type_config = TypeConfig(
|
||||
act_type=torch.float16,
|
||||
weight_type=scalar_types.uint4b8,
|
||||
output_type=None,
|
||||
group_scale_type=torch.float16,
|
||||
group_zero_type=None,
|
||||
channel_scale_type=None,
|
||||
token_scale_type=None,
|
||||
)
|
||||
|
||||
tensors = create_test_tensors(
|
||||
(512, 4096, 4096), type_config, group_size, subset_stride_factor=2
|
||||
)
|
||||
machete_mm_test_helper(type_config, tensors, group_size)
|
||||
|
||||
|
||||
# Test to make sure cuda graphs work
|
||||
class MacheteLayer(torch.nn.Module):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__()
|
||||
self.kwargs = kwargs
|
||||
|
||||
def forward(self, a):
|
||||
return ops.machete_mm(a=a, **self.kwargs)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not IS_SUPPORTED_BY_GPU, reason="Machete is not supported on this GPU type."
|
||||
)
|
||||
def test_machete_cuda_graph():
|
||||
m, n, k = 512, 4096, 4096
|
||||
|
||||
a = rand_data((m, k), torch.float16)
|
||||
b = rand_data((k, n), torch.float16)
|
||||
wtype = scalar_types.uint4b8
|
||||
stype = torch.float16
|
||||
group_size = 128
|
||||
zero_points = False
|
||||
|
||||
w_ref, w_q_packed, w_s, w_zp = machete_quantize_and_pack(
|
||||
a.dtype, b, wtype, stype, group_size, zero_points
|
||||
)
|
||||
|
||||
# Construct a trivial model with a single layer that calls a machete kernel
|
||||
model = MacheteLayer(
|
||||
b_q=w_q_packed,
|
||||
b_type=wtype,
|
||||
b_group_scales=w_s,
|
||||
b_group_zeros=maybe_convert_zeropoints(w_zp, w_s),
|
||||
b_group_size=group_size,
|
||||
)
|
||||
|
||||
output_ref = torch.matmul(a, w_ref)
|
||||
|
||||
# Run the model with a cuda graph
|
||||
stream = torch.cuda.Stream()
|
||||
with torch.cuda.stream(stream):
|
||||
g = torch.cuda.CUDAGraph()
|
||||
with torch.cuda.graph(g):
|
||||
output = model(a)
|
||||
output.zero_()
|
||||
g.replay()
|
||||
|
||||
# Relax atol as our reduction dim becomes larger (more rounding error)
|
||||
# Relax atol when we have zeropoints since the way machete applies
|
||||
# zeropoints (after scales) causes noise around 0
|
||||
atol = 1 if zero_points else min(5e-2 * math.sqrt(k), 1)
|
||||
torch.testing.assert_close(output, output_ref, rtol=1e-1, atol=atol)
|
||||
@@ -0,0 +1,621 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the marlin kernel.
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_marlin_gemm.py`.
|
||||
"""
|
||||
|
||||
import itertools
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.utils import opcheck
|
||||
from tests.quantization.utils import is_quant_method_supported
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.model_executor.layers.quantization.utils.int8_utils import (
|
||||
per_token_quant_int8,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
marlin_make_empty_g_idx,
|
||||
marlin_make_workspace_new,
|
||||
marlin_permute_bias,
|
||||
query_marlin_supported_quant_types,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
|
||||
rand_marlin_weight_mxfp4_like,
|
||||
rand_marlin_weight_nvfp4_like,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
|
||||
marlin_quant_fp8_torch,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils_test import (
|
||||
awq_marlin_quantize,
|
||||
get_weight_perm,
|
||||
marlin_quantize,
|
||||
marlin_weights,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
awq_pack,
|
||||
gptq_pack,
|
||||
gptq_quantize_weights,
|
||||
quantize_weights,
|
||||
sort_weights,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
if current_platform.is_rocm():
|
||||
pytest.skip(
|
||||
"These tests require marlin, which is not supported on ROCm.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
ACT_ORDER_OPTS = [False, True]
|
||||
K_FULL_OPTS = [False, True]
|
||||
USE_ATOMIC_ADD_OPTS = [False, True]
|
||||
USE_FP32_REDUCE_OPTS = [True]
|
||||
|
||||
MARLIN_K_CHUNKS = [128]
|
||||
MARLIN_N_CHUNKS = [64, 256]
|
||||
|
||||
MARLIN_REPACK_NK_FACTORS = [
|
||||
(4, 8),
|
||||
(7, 5),
|
||||
(13, 11),
|
||||
]
|
||||
|
||||
MNK_FACTORS = [
|
||||
(1, 1, 1),
|
||||
(1, 4, 8),
|
||||
(26, 37, 13),
|
||||
(257, 13, 11),
|
||||
]
|
||||
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
|
||||
DENSE_MARLIN_QUANT_TEST_CONFIGS = [
|
||||
# AWQ-INT4
|
||||
{"b_type": scalar_types.uint4, "group_blocks": [-1, 2, 4, 8]},
|
||||
# GPTQ-INT4
|
||||
{
|
||||
"b_type": scalar_types.uint4b8,
|
||||
"support_act_order": True,
|
||||
"group_blocks": [-1, 2, 4, 8],
|
||||
},
|
||||
# GPTQ-INT8
|
||||
{
|
||||
"b_type": scalar_types.uint8b128,
|
||||
"support_act_order": True,
|
||||
"group_blocks": [-1, 2, 4, 8],
|
||||
},
|
||||
# FP8
|
||||
{"b_type": scalar_types.float8_e4m3fn, "group_blocks": [-1, 8]},
|
||||
# NVFP4
|
||||
{"b_type": scalar_types.float4_e2m1f, "group_blocks": [1]},
|
||||
# MXFP4
|
||||
{
|
||||
"a_type": [scalar_types.bfloat16],
|
||||
"b_type": scalar_types.float4_e2m1f,
|
||||
"group_blocks": [2],
|
||||
},
|
||||
# AWQ-INT4 with INT8 activation
|
||||
{
|
||||
"a_type": [scalar_types.int8],
|
||||
"b_type": scalar_types.uint4,
|
||||
"group_blocks": [-1, 2, 4, 8],
|
||||
},
|
||||
# GPTQ-INT4 with INT8 activation
|
||||
{
|
||||
"a_type": [scalar_types.int8],
|
||||
"b_type": scalar_types.uint4b8,
|
||||
"group_blocks": [-1, 2, 4, 8],
|
||||
},
|
||||
# GPTQ-INT4 with FP8 activation
|
||||
{
|
||||
"a_type": [scalar_types.float8_e4m3fn],
|
||||
"b_type": scalar_types.uint4b8,
|
||||
"group_blocks": [-1, 2, 4, 8],
|
||||
},
|
||||
# AWQ-INT4 with FP8 activation
|
||||
{
|
||||
"a_type": [scalar_types.float8_e4m3fn],
|
||||
"b_type": scalar_types.uint4,
|
||||
"group_blocks": [-1, 2, 4, 8],
|
||||
},
|
||||
# MXFP4 with FP8 activation
|
||||
{
|
||||
"a_type": [scalar_types.float8_e4m3fn],
|
||||
"b_type": scalar_types.float4_e2m1f,
|
||||
"c_type": [scalar_types.bfloat16],
|
||||
"group_blocks": [2],
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def compute_max_diff(output, output_ref):
|
||||
return torch.mean(torch.abs(output - output_ref)) / torch.mean(
|
||||
torch.abs(output_ref)
|
||||
)
|
||||
|
||||
|
||||
def rand_data(shape, dtype=torch.float16):
|
||||
return torch.randn(shape, dtype=dtype, device="cuda")
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_quant_method_supported("gptq_marlin"),
|
||||
reason="Marlin is not supported on this GPU type.",
|
||||
)
|
||||
def test_marlin_int4_fp8_preprocess_without_zp():
|
||||
qweight_unpacked = torch.randint(
|
||||
0, 16, size=(2048, 2048), dtype=torch.int32, device="cuda"
|
||||
)
|
||||
qweight_packed = qweight_unpacked[:, ::2] * 16 + qweight_unpacked[:, 1::2]
|
||||
qweight_packed = qweight_packed.to(torch.int8).view(torch.int32)
|
||||
|
||||
cuda_res = ops.marlin_int4_fp8_preprocess(qweight_packed)
|
||||
|
||||
torch_res = torch.where(
|
||||
qweight_unpacked >= 8, qweight_unpacked - 8, 15 - qweight_unpacked
|
||||
)
|
||||
torch_res = torch_res[:, ::2] * 16 + torch_res[:, 1::2]
|
||||
torch_res = torch_res.to(torch.int8).view(torch.int32)
|
||||
|
||||
assert (cuda_res == torch_res).all()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_quant_method_supported("gptq_marlin"),
|
||||
reason="Marlin is not supported on this GPU type.",
|
||||
)
|
||||
def test_marlin_int4_fp8_preprocess_awq():
|
||||
group_size = 128
|
||||
|
||||
qweight_unpacked = torch.randint(
|
||||
0, 16, size=(2048, 2048), dtype=torch.int32, device="cuda"
|
||||
)
|
||||
qzeros_unpacked = torch.randint(
|
||||
0, 16, size=(2048 // group_size, 2048), dtype=torch.int32, device="cuda"
|
||||
)
|
||||
|
||||
qweight_packed = qweight_unpacked[:, ::2] * 16 + qweight_unpacked[:, 1::2]
|
||||
qweight_packed = qweight_packed.to(torch.int8).view(torch.int32)
|
||||
qzeros_packed = qzeros_unpacked[:, ::2] * 16 + qzeros_unpacked[:, 1::2]
|
||||
qzeros_packed = qzeros_packed.to(torch.int8).view(torch.int32)
|
||||
|
||||
cuda_res = ops.marlin_int4_fp8_preprocess(qweight_packed, qzeros_packed)
|
||||
|
||||
repeated_zp = qzeros_unpacked.repeat_interleave(group_size, 0)
|
||||
torch_res = qweight_unpacked - repeated_zp
|
||||
torch_res[torch_res < 0] = 15 - qweight_unpacked[torch_res < 0]
|
||||
torch_res = torch_res[:, ::2] * 16 + torch_res[:, 1::2]
|
||||
torch_res = torch_res.to(torch.int8).view(torch.int32)
|
||||
|
||||
assert (cuda_res == torch_res).all()
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_quant_method_supported("gptq_marlin"),
|
||||
reason="Marlin is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
|
||||
@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
|
||||
@pytest.mark.parametrize("quant_type", query_marlin_supported_quant_types(False, False))
|
||||
@pytest.mark.parametrize("act_order", ACT_ORDER_OPTS)
|
||||
@pytest.mark.parametrize("is_a_8bit", [True, False])
|
||||
@pytest.mark.parametrize("nk_factors", MARLIN_REPACK_NK_FACTORS)
|
||||
def test_gptq_marlin_repack(
|
||||
k_chunk, n_chunk, quant_type, act_order, is_a_8bit, nk_factors
|
||||
):
|
||||
n_factor, k_factor = nk_factors
|
||||
|
||||
size_k = k_chunk * k_factor
|
||||
size_n = n_chunk * n_factor
|
||||
group_size = 128
|
||||
|
||||
# Filter act_order
|
||||
if act_order:
|
||||
if group_size == -1:
|
||||
return
|
||||
if group_size == size_k:
|
||||
return
|
||||
if is_a_8bit:
|
||||
return
|
||||
|
||||
# Normalize group_size
|
||||
if group_size == -1:
|
||||
group_size = size_k
|
||||
assert group_size <= size_k
|
||||
|
||||
# Create input
|
||||
b_weight = rand_data((size_k, size_n))
|
||||
|
||||
# Quantize (and apply act_order if provided)
|
||||
w_ref, q_w, s, g_idx, rand_perm = gptq_quantize_weights(
|
||||
b_weight, quant_type, group_size, act_order
|
||||
)
|
||||
|
||||
# Pack to GPTQ format
|
||||
q_w_gptq = gptq_pack(q_w, quant_type.size_bits, size_k, size_n)
|
||||
|
||||
# For act_order, sort the "weights" and "g_idx" so that group ids are
|
||||
# increasing
|
||||
sort_indices = torch.empty(0, dtype=torch.int, device=b_weight.device)
|
||||
if act_order:
|
||||
q_w, g_idx, sort_indices = sort_weights(q_w, g_idx)
|
||||
|
||||
# Pack to Marlin format
|
||||
weight_perm = get_weight_perm(quant_type.size_bits, is_a_8bit)
|
||||
marlin_q_w_1 = marlin_weights(
|
||||
q_w, size_k, size_n, quant_type.size_bits, weight_perm, is_a_8bit
|
||||
)
|
||||
|
||||
opcheck(
|
||||
torch.ops._C.gptq_marlin_repack,
|
||||
(q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits, is_a_8bit),
|
||||
)
|
||||
|
||||
# Run Marlin repack GPU kernel
|
||||
marlin_q_w_2 = ops.gptq_marlin_repack(
|
||||
q_w_gptq, sort_indices, size_k, size_n, quant_type.size_bits, is_a_8bit
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
torch.testing.assert_close(marlin_q_w_1, marlin_q_w_2)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_quant_method_supported("gptq_marlin"),
|
||||
reason="Marlin is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("k_chunk", MARLIN_K_CHUNKS)
|
||||
@pytest.mark.parametrize("n_chunk", MARLIN_N_CHUNKS)
|
||||
@pytest.mark.parametrize("quant_type", query_marlin_supported_quant_types(True))
|
||||
@pytest.mark.parametrize("is_a_8bit", [True, False])
|
||||
@pytest.mark.parametrize("nk_factors", MARLIN_REPACK_NK_FACTORS)
|
||||
def test_awq_marlin_repack(k_chunk, n_chunk, quant_type, is_a_8bit, nk_factors):
|
||||
n_factor, k_factor = nk_factors
|
||||
|
||||
size_k = k_chunk * k_factor
|
||||
size_n = n_chunk * n_factor
|
||||
|
||||
group_size = 128
|
||||
|
||||
# Create input
|
||||
b_weight = rand_data((size_k, size_n))
|
||||
|
||||
# Quantize
|
||||
w_ref, q_w, s, zp = quantize_weights(
|
||||
b_weight, quant_type, group_size, zero_points=True
|
||||
)
|
||||
|
||||
# Pack to AWQ format
|
||||
q_w_awq = awq_pack(q_w, quant_type.size_bits, size_k, size_n)
|
||||
|
||||
# Pack to Marlin format
|
||||
weight_perm = get_weight_perm(quant_type.size_bits, is_a_8bit)
|
||||
marlin_q_w_1 = marlin_weights(
|
||||
q_w, size_k, size_n, quant_type.size_bits, weight_perm, is_a_8bit
|
||||
)
|
||||
|
||||
opcheck(
|
||||
torch.ops._C.awq_marlin_repack,
|
||||
(q_w_awq, size_k, size_n, quant_type.size_bits, is_a_8bit),
|
||||
)
|
||||
|
||||
# Run Marlin repack GPU kernel
|
||||
marlin_q_w_2 = ops.awq_marlin_repack(
|
||||
q_w_awq, size_k, size_n, quant_type.size_bits, is_a_8bit
|
||||
)
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
torch.testing.assert_close(marlin_q_w_1, marlin_q_w_2)
|
||||
|
||||
|
||||
def marlin_generate_valid_test_cases():
|
||||
all_combinations = itertools.product(
|
||||
DENSE_MARLIN_QUANT_TEST_CONFIGS,
|
||||
MNK_FACTORS,
|
||||
MARLIN_N_CHUNKS,
|
||||
MARLIN_K_CHUNKS,
|
||||
ACT_ORDER_OPTS,
|
||||
K_FULL_OPTS,
|
||||
USE_ATOMIC_ADD_OPTS,
|
||||
USE_FP32_REDUCE_OPTS,
|
||||
)
|
||||
|
||||
def is_invalid(
|
||||
a_type,
|
||||
b_type,
|
||||
c_type,
|
||||
group_blocks,
|
||||
size_m,
|
||||
size_n,
|
||||
size_k,
|
||||
act_order,
|
||||
is_k_full,
|
||||
use_atomic_add,
|
||||
use_fp32_reduce,
|
||||
):
|
||||
if use_atomic_add:
|
||||
if use_fp32_reduce:
|
||||
return False
|
||||
if (
|
||||
c_type == scalar_types.bfloat16
|
||||
and torch.cuda.get_device_capability()[0] < 9
|
||||
):
|
||||
return False
|
||||
|
||||
group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
|
||||
if group_size > 0 and size_k % group_size != 0:
|
||||
return False
|
||||
|
||||
if act_order and group_size in [-1, size_k]:
|
||||
return False
|
||||
if group_size == size_k:
|
||||
return False
|
||||
if not act_order and is_k_full:
|
||||
return False
|
||||
|
||||
return a_type.size_bits < 16 or a_type is c_type
|
||||
|
||||
cases = []
|
||||
for case in all_combinations:
|
||||
quant_test_config, mnk_factors, n_chunk, k_chunk, act_order, *_ = case
|
||||
size_m = mnk_factors[0]
|
||||
size_n = mnk_factors[1] * n_chunk
|
||||
size_k = mnk_factors[2] * k_chunk
|
||||
|
||||
if act_order and not quant_test_config.get("support_act_order", False):
|
||||
continue
|
||||
|
||||
f16_types = [scalar_types.float16, scalar_types.bfloat16]
|
||||
inner_combinations = itertools.product(
|
||||
quant_test_config.get("a_type", f16_types),
|
||||
[quant_test_config["b_type"]],
|
||||
quant_test_config.get("c_type", f16_types),
|
||||
quant_test_config["group_blocks"],
|
||||
)
|
||||
|
||||
for sub_case in inner_combinations:
|
||||
if (
|
||||
sub_case[0] == scalar_types.float8_e4m3fn
|
||||
and not current_platform.is_device_capability(89)
|
||||
and not current_platform.is_device_capability_family(120)
|
||||
):
|
||||
continue
|
||||
args = sub_case + (size_m, size_n, size_k) + case[4:]
|
||||
if is_invalid(*args):
|
||||
cases.append(args)
|
||||
return cases
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not is_quant_method_supported("gptq_marlin"),
|
||||
reason="Marlin is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
(
|
||||
"a_type, b_type, c_type, group_blocks,"
|
||||
"size_m, size_n, size_k, act_order, is_k_full,"
|
||||
"use_atomic_add, use_fp32_reduce"
|
||||
),
|
||||
marlin_generate_valid_test_cases(),
|
||||
)
|
||||
def test_marlin_gemm(
|
||||
a_type,
|
||||
b_type,
|
||||
c_type,
|
||||
group_blocks,
|
||||
size_m,
|
||||
size_n,
|
||||
size_k,
|
||||
act_order,
|
||||
is_k_full,
|
||||
use_atomic_add,
|
||||
use_fp32_reduce,
|
||||
):
|
||||
has_zp = b_type in [scalar_types.uint4, scalar_types.uint8]
|
||||
|
||||
group_size = group_blocks if group_blocks <= 0 else group_blocks * 16
|
||||
|
||||
if c_type == scalar_types.float16:
|
||||
dtype = torch.float16
|
||||
elif c_type == scalar_types.bfloat16:
|
||||
dtype = torch.bfloat16
|
||||
else:
|
||||
raise RuntimeError("unsupported c_type")
|
||||
|
||||
if a_type == scalar_types.int8:
|
||||
a_dtype = torch.int8
|
||||
elif a_type == scalar_types.float8_e4m3fn:
|
||||
a_dtype = torch.float8_e4m3fn
|
||||
else:
|
||||
a_dtype = dtype
|
||||
|
||||
a_input = rand_data((size_m, size_k), dtype=dtype)
|
||||
b_weight = rand_data((size_k, size_n), dtype=dtype)
|
||||
|
||||
if b_type == scalar_types.float4_e2m1f:
|
||||
if group_size == 16:
|
||||
w_ref, marlin_q_w, marlin_s, marlin_s2 = rand_marlin_weight_nvfp4_like(
|
||||
b_weight.T, group_size, input_dtype=a_dtype
|
||||
)
|
||||
else:
|
||||
w_ref, marlin_q_w, marlin_s = rand_marlin_weight_mxfp4_like(
|
||||
b_weight.T, group_size, input_dtype=a_dtype
|
||||
)
|
||||
marlin_s2 = None
|
||||
|
||||
g_idx = None
|
||||
sort_indices = None
|
||||
marlin_zp = None
|
||||
elif b_type == scalar_types.float8_e4m3fn:
|
||||
w_ref, marlin_q_w, marlin_s = marlin_quant_fp8_torch(
|
||||
b_weight.T, group_size, input_dtype=a_dtype
|
||||
)
|
||||
g_idx = None
|
||||
sort_indices = None
|
||||
marlin_zp = None
|
||||
marlin_s2 = None
|
||||
elif has_zp:
|
||||
w_ref, marlin_q_w, marlin_s, marlin_zp = awq_marlin_quantize(
|
||||
b_weight, b_type, group_size, input_dtype=a_dtype
|
||||
)
|
||||
g_idx = None
|
||||
sort_indices = None
|
||||
marlin_s2 = None
|
||||
else:
|
||||
w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
|
||||
b_weight, b_type, group_size, act_order, input_dtype=a_dtype
|
||||
)
|
||||
|
||||
marlin_zp = None
|
||||
marlin_s2 = None
|
||||
|
||||
workspace = marlin_make_workspace_new(w_ref.device)
|
||||
|
||||
if a_type == scalar_types.int8:
|
||||
a_input, a_scales = per_token_quant_int8(a_input)
|
||||
a_input_ref = a_input.to(a_scales.dtype) * a_scales.view(-1, 1)
|
||||
a_input_ref = a_input_ref.to(dtype)
|
||||
|
||||
if group_size != -1:
|
||||
a_scales = a_scales / 4096 * marlin_s.max()
|
||||
a_scales = a_scales.float()
|
||||
marlin_s = marlin_s / marlin_s.max() * 4096
|
||||
marlin_s = marlin_s.round().to(torch.int16).view(dtype)
|
||||
elif a_type == scalar_types.float8_e4m3fn:
|
||||
a_input, a_scales = ops.scaled_fp8_quant(a_input, use_per_token_if_dynamic=True)
|
||||
a_input_ref = a_input.to(a_scales.dtype) * a_scales.view(-1, 1)
|
||||
a_input_ref = a_input_ref.to(dtype)
|
||||
else:
|
||||
assert a_type.size_bits == 16
|
||||
a_input_ref = a_input
|
||||
a_scales = None
|
||||
|
||||
output = torch.empty((size_m, size_n), dtype=dtype, device=a_input.device)
|
||||
|
||||
output = ops.marlin_gemm(
|
||||
a_input,
|
||||
output,
|
||||
marlin_q_w,
|
||||
None,
|
||||
marlin_s,
|
||||
a_scales,
|
||||
marlin_s2,
|
||||
marlin_zp,
|
||||
g_idx,
|
||||
sort_indices,
|
||||
workspace,
|
||||
b_type,
|
||||
a_input.shape[0],
|
||||
b_weight.shape[1],
|
||||
a_input.shape[1],
|
||||
is_k_full=is_k_full,
|
||||
use_atomic_add=use_atomic_add,
|
||||
use_fp32_reduce=use_fp32_reduce,
|
||||
is_zp_float=False,
|
||||
)
|
||||
output_ref = torch.matmul(a_input_ref, w_ref)
|
||||
|
||||
max_diff = compute_max_diff(output, output_ref)
|
||||
assert max_diff < 0.04
|
||||
|
||||
|
||||
def test_marlin_gemm_subset_input():
|
||||
quant_type = scalar_types.uint4b8
|
||||
group_size = 128
|
||||
|
||||
size_m, size_k, size_n = 32, 1024, 2048
|
||||
big_m = size_m * 2
|
||||
big_k = size_k * 2
|
||||
|
||||
a_input = rand_data((big_m, big_k))[8 : size_m + 8, 8 : size_k + 8]
|
||||
b_weight = rand_data((size_k, size_n))
|
||||
|
||||
w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
|
||||
b_weight, quant_type, group_size, False
|
||||
)
|
||||
|
||||
marlin_zp = marlin_make_empty_g_idx(marlin_s.device)
|
||||
workspace = marlin_make_workspace_new(a_input.device)
|
||||
|
||||
output = ops.marlin_gemm(
|
||||
a_input,
|
||||
None,
|
||||
marlin_q_w,
|
||||
None,
|
||||
marlin_s,
|
||||
None,
|
||||
None,
|
||||
marlin_zp,
|
||||
g_idx,
|
||||
sort_indices,
|
||||
workspace,
|
||||
quant_type,
|
||||
a_input.shape[0],
|
||||
b_weight.shape[1],
|
||||
a_input.shape[1],
|
||||
is_k_full=True,
|
||||
use_atomic_add=False,
|
||||
use_fp32_reduce=True,
|
||||
is_zp_float=False,
|
||||
)
|
||||
output_ref = torch.matmul(a_input, w_ref)
|
||||
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
max_diff = compute_max_diff(output, output_ref)
|
||||
|
||||
assert max_diff < 0.04
|
||||
|
||||
|
||||
@pytest.mark.parametrize("size_m", [1, 256])
|
||||
def test_marlin_gemm_with_bias(size_m):
|
||||
quant_type = scalar_types.uint4b8
|
||||
group_size = 128
|
||||
|
||||
size_k, size_n = 1024, 2048
|
||||
a_input = rand_data((size_m, size_k))
|
||||
b_weight = rand_data((size_k, size_n))
|
||||
b_bias = rand_data((size_n,)) * 10
|
||||
|
||||
marlin_bias = marlin_permute_bias(b_bias)
|
||||
|
||||
w_ref, marlin_q_w, marlin_s, g_idx, sort_indices, _ = marlin_quantize(
|
||||
b_weight, quant_type, group_size, False
|
||||
)
|
||||
|
||||
marlin_zp = marlin_make_empty_g_idx(marlin_s.device)
|
||||
workspace = marlin_make_workspace_new(a_input.device)
|
||||
|
||||
output = ops.marlin_gemm(
|
||||
a_input,
|
||||
None,
|
||||
marlin_q_w,
|
||||
marlin_bias,
|
||||
marlin_s,
|
||||
None,
|
||||
None,
|
||||
marlin_zp,
|
||||
g_idx,
|
||||
sort_indices,
|
||||
workspace,
|
||||
quant_type,
|
||||
a_input.shape[0],
|
||||
b_weight.shape[1],
|
||||
a_input.shape[1],
|
||||
is_k_full=True,
|
||||
use_atomic_add=False,
|
||||
use_fp32_reduce=True,
|
||||
is_zp_float=False,
|
||||
)
|
||||
output_ref = torch.matmul(a_input, w_ref) + b_bias.view(1, -1)
|
||||
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
max_diff = compute_max_diff(output, output_ref)
|
||||
|
||||
assert max_diff < 0.04
|
||||
@@ -0,0 +1,863 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for Marlin thread-tile padding of TP-sharded weight shapes.
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_marlin_tile_padding.py`.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
GPTQ_MARLIN_TILE,
|
||||
apply_gptq_marlin_linear,
|
||||
marlin_make_empty_g_idx,
|
||||
marlin_make_workspace_new,
|
||||
marlin_moe_padded_intermediate,
|
||||
marlin_pad_qweight,
|
||||
marlin_pad_scales,
|
||||
marlin_padded_nk,
|
||||
marlin_permute_scales,
|
||||
marlin_repacked_nk,
|
||||
marlin_zero_points,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
|
||||
apply_fp4_marlin_linear,
|
||||
is_fp4_marlin_supported,
|
||||
prepare_fp4_layer_for_marlin,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
|
||||
apply_fp8_marlin_linear,
|
||||
apply_mxfp8_marlin_linear,
|
||||
is_fp8_marlin_supported,
|
||||
prepare_fp8_layer_for_marlin,
|
||||
prepare_mxfp8_layer_for_marlin,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
gptq_pack,
|
||||
gptq_quantize_weights,
|
||||
quantize_weights,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
# (size_n, size_k) rank-local shapes that violate Marlin tile alignment,
|
||||
# e.g. produced by TP-sharding dims that are valid at TP=1.
|
||||
ODD_SHAPES = [
|
||||
(200, 288), # N padded
|
||||
(256, 208), # K padded
|
||||
(200, 208), # both padded
|
||||
(4640, 512), # Nemotron-Super-120B q_proj shard at TP=4
|
||||
]
|
||||
ALIGNED_SHAPES = [(64, 128), (128, 64), (256, 256), (4608, 4096)]
|
||||
|
||||
|
||||
def _is_tile_aligned(size_n: int, size_k: int) -> bool:
|
||||
return (size_n % 64 == 0 and size_k % 128 == 0) or (
|
||||
size_n % 128 == 0 and size_k % 64 == 0
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", ODD_SHAPES + ALIGNED_SHAPES)
|
||||
@pytest.mark.parametrize("group_size", [-1, 16, 32, 64, 128])
|
||||
def test_marlin_padded_nk(shape, group_size):
|
||||
size_n, size_k = shape
|
||||
padded_n, padded_k = marlin_padded_nk(size_n, size_k, group_size)
|
||||
|
||||
assert padded_n >= size_n and padded_k >= size_k
|
||||
assert _is_tile_aligned(padded_n, padded_k)
|
||||
if group_size > 0:
|
||||
assert padded_k % group_size == 0
|
||||
|
||||
# Aligned shapes must pass through unchanged (zero hot-path cost).
|
||||
if _is_tile_aligned(size_n, size_k) and (
|
||||
group_size <= 0 or size_k % group_size == 0
|
||||
):
|
||||
assert (padded_n, padded_k) == (size_n, size_k)
|
||||
|
||||
# Minimal: no valid shape with a smaller padded area exists.
|
||||
area = padded_n * padded_k
|
||||
for cand_n in range(size_n, padded_n + 1):
|
||||
for cand_k in range(size_k, padded_k + 1):
|
||||
if (
|
||||
_is_tile_aligned(cand_n, cand_k)
|
||||
and (group_size <= 0 or cand_k % group_size == 0)
|
||||
and cand_n * cand_k < area
|
||||
):
|
||||
pytest.fail(f"({cand_n}, {cand_k}) beats ({padded_n}, {padded_k})")
|
||||
|
||||
# Apply-time derivation from the repacked-tensor shape must round-trip.
|
||||
for num_bits in (4, 8):
|
||||
pack_factor = 32 // num_bits
|
||||
repacked_shape = (
|
||||
padded_k // GPTQ_MARLIN_TILE,
|
||||
padded_n * GPTQ_MARLIN_TILE // pack_factor,
|
||||
)
|
||||
repacked = torch.empty(repacked_shape, device="meta")
|
||||
assert marlin_repacked_nk(repacked, num_bits) == (padded_n, padded_k)
|
||||
|
||||
|
||||
def test_marlin_pad_helpers_shapes():
|
||||
size_n, size_k, group_size = 200, 208, 16
|
||||
padded_n, padded_k = marlin_padded_nk(size_n, size_k, group_size)
|
||||
|
||||
qweight = torch.zeros(size_k // 8, size_n, dtype=torch.int32)
|
||||
padded = marlin_pad_qweight(qweight, size_n, size_k, padded_n, padded_k)
|
||||
assert padded.shape == (padded_k // 8, padded_n)
|
||||
|
||||
scales = torch.ones(size_k // group_size, size_n)
|
||||
padded = marlin_pad_scales(scales, size_n, size_k, padded_n, padded_k, group_size)
|
||||
assert padded.shape == (padded_k // group_size, padded_n)
|
||||
assert padded[:, size_n:].abs().sum() == 0
|
||||
|
||||
channelwise = torch.ones(1, size_n)
|
||||
padded = marlin_pad_scales(channelwise, size_n, size_k, padded_n, padded_k, -1)
|
||||
assert padded.shape == (1, padded_n)
|
||||
|
||||
|
||||
# Rank-local MoE intermediate sizes. group<=0 / 32 with a non-multiple-of-64
|
||||
# size is where tile padding triggers; 64/128 are already tile-aligned.
|
||||
MOE_INTERMEDIATE_SIZES = [64, 96, 100, 176, 192, 256, 2816]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("intermediate", MOE_INTERMEDIATE_SIZES)
|
||||
@pytest.mark.parametrize("group_size", [-1, 32, 64, 128])
|
||||
def test_marlin_moe_padded_intermediate(intermediate, group_size):
|
||||
# The MoE gate only admits shapes where the group does not straddle the
|
||||
# boundary, i.e. group divides the intermediate size.
|
||||
if group_size > 0 and intermediate % group_size != 0:
|
||||
pytest.skip("group straddles the boundary; rejected by the MoE gate")
|
||||
|
||||
padded = marlin_moe_padded_intermediate(intermediate, group_size)
|
||||
assert padded >= intermediate
|
||||
# Valid MoE thread tile: gate-up n = 2*intermediate % 128, down k % 64.
|
||||
assert (2 * padded) % 128 == 0
|
||||
assert padded % 64 == 0
|
||||
if group_size > 0:
|
||||
assert padded % group_size == 0
|
||||
|
||||
# Minimal: no smaller valid intermediate exists.
|
||||
for cand in range(intermediate, padded):
|
||||
if (
|
||||
(2 * cand) % 128 == 0
|
||||
and cand % 64 == 0
|
||||
and (group_size <= 0 or cand % group_size == 0)
|
||||
):
|
||||
pytest.fail(f"{cand} beats {padded}")
|
||||
|
||||
# Already-tile-aligned sizes pass through unchanged (zero hot-path cost).
|
||||
if intermediate % 64 == 0:
|
||||
assert padded == intermediate
|
||||
|
||||
|
||||
def test_marlin_moe_pad_helpers_shapes():
|
||||
from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
|
||||
_pad_rows,
|
||||
_pad_w13_bias,
|
||||
_pad_w13_shard_cols,
|
||||
)
|
||||
|
||||
E, rows, N, padded_N = 2, 8, 96, 128
|
||||
|
||||
# w13 stores the two gate/up shards along the last dim; padding each shard
|
||||
# must preserve the loaded values and zero the padded columns.
|
||||
w13 = torch.arange(E * rows * 2 * N).reshape(E, rows, 2 * N).float()
|
||||
padded = _pad_w13_shard_cols(w13, N, padded_N)
|
||||
assert padded.shape == (E, rows, 2 * padded_N)
|
||||
shards = padded.view(E, rows, 2, padded_N)
|
||||
orig = w13.view(E, rows, 2, N)
|
||||
assert torch.equal(shards[..., :N], orig)
|
||||
assert shards[..., N:].abs().sum() == 0
|
||||
|
||||
# w2 stores the intermediate dim in the rows.
|
||||
w2 = torch.ones(E, N // 32, 16)
|
||||
padded = _pad_rows(w2, padded_N // 32)
|
||||
assert padded.shape == (E, padded_N // 32, 16)
|
||||
assert padded[:, N // 32 :, :].abs().sum() == 0
|
||||
|
||||
bias = torch.arange(E * 2 * N).reshape(E, 2 * N).float()
|
||||
padded = _pad_w13_bias(bias, N, padded_N)
|
||||
assert padded.shape == (E, 2 * padded_N)
|
||||
bias_shards = padded.view(E, 2, padded_N)
|
||||
assert torch.equal(bias_shards[..., :N], bias.view(E, 2, N))
|
||||
assert bias_shards[..., N:].abs().sum() == 0
|
||||
|
||||
|
||||
def _gpu_marlin_unsupported() -> bool:
|
||||
return not (
|
||||
current_platform.is_cuda() and current_platform.has_device_capability(80)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
_gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
|
||||
reason="FP8 Marlin is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("shape", ODD_SHAPES)
|
||||
@pytest.mark.parametrize("use_bias", [False, True])
|
||||
def test_fp8_marlin_padded_round_trip(shape, use_bias):
|
||||
size_n, size_k = shape
|
||||
dtype = torch.float16
|
||||
layer = torch.nn.Module()
|
||||
layer.output_size_per_partition = size_n
|
||||
layer.input_size_per_partition = size_k
|
||||
layer.orig_dtype = dtype
|
||||
|
||||
weight = torch.randn(size_k, size_n, dtype=dtype, device="cuda") / size_k**0.5
|
||||
scale = weight.abs().max() / 448
|
||||
weight_fp8 = (weight / scale).to(torch.float8_e4m3fn)
|
||||
layer.weight = torch.nn.Parameter(weight_fp8, requires_grad=False)
|
||||
layer.weight_scale = torch.nn.Parameter(
|
||||
scale.to(torch.float32), requires_grad=False
|
||||
)
|
||||
bias = None
|
||||
if use_bias:
|
||||
bias = torch.randn(size_n, dtype=dtype, device="cuda")
|
||||
layer.bias = torch.nn.Parameter(bias.clone(), requires_grad=False)
|
||||
|
||||
prepare_fp8_layer_for_marlin(layer, size_k_first=True)
|
||||
|
||||
x = torch.randn(8, size_k, dtype=dtype, device="cuda")
|
||||
output = apply_fp8_marlin_linear(
|
||||
input=x,
|
||||
weight=layer.weight,
|
||||
weight_scale=layer.weight_scale,
|
||||
workspace=layer.workspace,
|
||||
size_n=size_n,
|
||||
size_k=size_k,
|
||||
bias=layer.bias if use_bias else None,
|
||||
)
|
||||
ref = x @ (weight_fp8.to(dtype) * scale.to(dtype))
|
||||
if use_bias:
|
||||
ref = ref + bias
|
||||
|
||||
assert output.shape == (8, size_n)
|
||||
torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)
|
||||
|
||||
|
||||
def _dequant_fp4(packed: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
|
||||
"""Dequantize packed e2m1 nibbles (N, K // 2) -> (N, K) in dtype."""
|
||||
lo = (packed & 0b10000000) | ((packed & 0b01110000) >> 2)
|
||||
lo = lo.view(torch.float8_e4m3fn).to(dtype) * (2**6)
|
||||
hi_bits = packed << 4
|
||||
hi = (hi_bits & 0b10000000) | ((hi_bits & 0b01110000) >> 2)
|
||||
hi = hi.view(torch.float8_e4m3fn).to(dtype) * (2**6)
|
||||
return torch.cat([hi.unsqueeze(2), lo.unsqueeze(2)], 2).view(packed.size(0), -1)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
_gpu_marlin_unsupported() or not is_fp4_marlin_supported(),
|
||||
reason="FP4 Marlin is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("shape", ODD_SHAPES)
|
||||
def test_nvfp4_marlin_padded_round_trip(shape):
|
||||
size_n, size_k = shape
|
||||
group_size = 16
|
||||
dtype = torch.float16
|
||||
layer = torch.nn.Module()
|
||||
layer.output_size_per_partition = size_n
|
||||
layer.input_size_per_partition = size_k
|
||||
layer.params_dtype = dtype
|
||||
|
||||
packed = torch.randint(
|
||||
0, 256, (size_n, size_k // 2), dtype=torch.uint8, device="cuda"
|
||||
)
|
||||
scales = (torch.rand(size_n, size_k // group_size, device="cuda") + 0.25).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
global_scale = torch.tensor([0.002], dtype=torch.float32, device="cuda")
|
||||
|
||||
ref_weight = (
|
||||
_dequant_fp4(packed, dtype)
|
||||
* scales.to(dtype).repeat_interleave(group_size, 1)
|
||||
* global_scale.to(dtype)
|
||||
)
|
||||
|
||||
layer.weight = torch.nn.Parameter(packed, requires_grad=False)
|
||||
layer.weight_scale = torch.nn.Parameter(scales, requires_grad=False)
|
||||
layer.weight_global_scale = torch.nn.Parameter(global_scale, requires_grad=False)
|
||||
|
||||
prepare_fp4_layer_for_marlin(layer)
|
||||
|
||||
x = torch.randn(8, size_k, dtype=dtype, device="cuda") / size_k**0.5
|
||||
output = apply_fp4_marlin_linear(
|
||||
input=x,
|
||||
weight=layer.weight,
|
||||
weight_scale=layer.weight_scale,
|
||||
weight_global_scale=layer.weight_global_scale,
|
||||
workspace=layer.workspace,
|
||||
size_n=size_n,
|
||||
size_k=size_k,
|
||||
)
|
||||
ref = x @ ref_weight.T
|
||||
|
||||
assert output.shape == (8, size_n)
|
||||
torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
_gpu_marlin_unsupported(),
|
||||
reason="Marlin is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("shape", ODD_SHAPES)
|
||||
@pytest.mark.parametrize("group_size", [-1, 128])
|
||||
def test_gptq_marlin_padded_round_trip(shape, group_size):
|
||||
"""Pad-then-repack a GPTQ int4 weight the way MarlinLinearKernel does and
|
||||
check the GEMM against the dequantized reference.
|
||||
|
||||
Symmetric int4's quantized zero decodes to -8, so this exercises the
|
||||
zero-padded-scales cancellation, not just zero weights.
|
||||
"""
|
||||
size_n, size_k = shape
|
||||
if group_size > 0 and size_k % group_size != 0:
|
||||
pytest.skip("group must divide the rank-local K (not fixable by padding)")
|
||||
dtype = torch.float16
|
||||
quant_type = scalar_types.uint4b8
|
||||
device = torch.device("cuda")
|
||||
|
||||
weight = torch.randn(size_k, size_n, dtype=dtype, device=device) / size_k**0.5
|
||||
w_ref, q_w, s, _, _ = gptq_quantize_weights(
|
||||
weight, quant_type, group_size, act_order=False
|
||||
)
|
||||
qweight = gptq_pack(q_w, quant_type.size_bits, size_k, size_n)
|
||||
|
||||
padded_n, padded_k = marlin_padded_nk(size_n, size_k, group_size)
|
||||
qweight = marlin_pad_qweight(qweight, size_n, size_k, padded_n, padded_k)
|
||||
marlin_qweight = ops.gptq_marlin_repack(
|
||||
b_q_weight=qweight,
|
||||
perm=torch.empty(0, dtype=torch.int, device=device),
|
||||
size_k=padded_k,
|
||||
size_n=padded_n,
|
||||
num_bits=quant_type.size_bits,
|
||||
)
|
||||
s = marlin_pad_scales(s, size_n, size_k, padded_n, padded_k, group_size)
|
||||
marlin_s = marlin_permute_scales(
|
||||
s, size_k=padded_k, size_n=padded_n, group_size=group_size
|
||||
)
|
||||
|
||||
x = torch.randn(8, size_k, dtype=dtype, device=device)
|
||||
output = apply_gptq_marlin_linear(
|
||||
input=x,
|
||||
weight=marlin_qweight,
|
||||
weight_scale=marlin_s,
|
||||
weight_zp=marlin_make_empty_g_idx(device),
|
||||
g_idx=marlin_make_empty_g_idx(device),
|
||||
g_idx_sort_indices=marlin_make_empty_g_idx(device),
|
||||
workspace=marlin_make_workspace_new(device),
|
||||
wtype=quant_type,
|
||||
output_size_per_partition=size_n,
|
||||
input_size_per_partition=size_k,
|
||||
is_k_full=True,
|
||||
)
|
||||
ref = x @ w_ref
|
||||
|
||||
assert output.shape == (8, size_n)
|
||||
torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
_gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
|
||||
reason="FP8 Marlin is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("shape", [(200, 512), (4640, 512)])
|
||||
def test_fp8_block_marlin_padded_round_trip(shape):
|
||||
"""Block-quantized FP8 (e.g. Nemotron NVFP4 checkpoints' FP8 layers):
|
||||
group_size=128 exercises the lcm K-alignment in marlin_padded_nk and the
|
||||
weight_scale_inv group-wise scale padding."""
|
||||
size_n, size_k = shape
|
||||
block = 128
|
||||
dtype = torch.float16
|
||||
layer = torch.nn.Module()
|
||||
layer.output_size_per_partition = size_n
|
||||
layer.input_size_per_partition = size_k
|
||||
layer.orig_dtype = dtype
|
||||
layer.weight_block_size = [block, block]
|
||||
|
||||
weight = torch.randn(size_n, size_k, dtype=dtype, device="cuda") / size_k**0.5
|
||||
n_blocks, k_blocks = (size_n + block - 1) // block, size_k // block
|
||||
padded = torch.zeros(n_blocks * block, size_k, dtype=dtype, device="cuda")
|
||||
padded[:size_n] = weight
|
||||
scales = padded.view(n_blocks, block, k_blocks, block).abs().amax(dim=(1, 3)) / 448
|
||||
scales_expanded = scales.repeat_interleave(block, 0)[:size_n].repeat_interleave(
|
||||
block, 1
|
||||
)
|
||||
weight_fp8 = (weight / scales_expanded).to(torch.float8_e4m3fn)
|
||||
|
||||
layer.weight = torch.nn.Parameter(weight_fp8, requires_grad=False)
|
||||
layer.weight_scale_inv = torch.nn.Parameter(
|
||||
scales.to(torch.float32), requires_grad=False
|
||||
)
|
||||
|
||||
prepare_fp8_layer_for_marlin(layer, size_k_first=False)
|
||||
|
||||
x = torch.randn(8, size_k, dtype=dtype, device="cuda")
|
||||
output = apply_fp8_marlin_linear(
|
||||
input=x,
|
||||
weight=layer.weight,
|
||||
weight_scale=layer.weight_scale_inv,
|
||||
workspace=layer.workspace,
|
||||
size_n=size_n,
|
||||
size_k=size_k,
|
||||
bias=None,
|
||||
)
|
||||
ref = x @ (weight_fp8.to(dtype) * scales_expanded.to(dtype)).T
|
||||
|
||||
assert output.shape == (8, size_n)
|
||||
torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
_gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
|
||||
reason="FP8 Marlin is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("shape", [(200, 288), (4640, 512)])
|
||||
def test_mxfp8_marlin_padded_round_trip(shape):
|
||||
"""MXFP8 exercises the e8m0 scale path, where padded 0.0 scales clamp to
|
||||
2^-127 instead of zero and must still contribute nothing."""
|
||||
size_n, size_k = shape
|
||||
group_size = 32
|
||||
# The e8m0-scale Marlin kernels are only instantiated for bf16 activations.
|
||||
dtype = torch.bfloat16
|
||||
layer = torch.nn.Module()
|
||||
layer.output_size_per_partition = size_n
|
||||
layer.input_size_per_partition = size_k
|
||||
|
||||
weight_fp8 = (torch.randn(size_n, size_k, dtype=dtype, device="cuda") / 4).to(
|
||||
torch.float8_e4m3fn
|
||||
)
|
||||
# e8m0 exponents around 1.0 (127): scales in [2^-6, 2^0]
|
||||
scales = torch.randint(
|
||||
121, 128, (size_n, size_k // group_size), dtype=torch.uint8, device="cuda"
|
||||
)
|
||||
ref_weight = weight_fp8.to(dtype) * (
|
||||
2.0 ** (scales.to(dtype) - 127)
|
||||
).repeat_interleave(group_size, 1)
|
||||
|
||||
layer.weight = torch.nn.Parameter(weight_fp8, requires_grad=False)
|
||||
layer.weight_scale = torch.nn.Parameter(scales, requires_grad=False)
|
||||
|
||||
prepare_mxfp8_layer_for_marlin(layer)
|
||||
|
||||
x = torch.randn(8, size_k, dtype=dtype, device="cuda") / size_k**0.5
|
||||
output = apply_mxfp8_marlin_linear(
|
||||
input=x,
|
||||
weight=layer.weight,
|
||||
weight_scale=layer.weight_scale,
|
||||
workspace=layer.workspace,
|
||||
size_n=size_n,
|
||||
size_k=size_k,
|
||||
)
|
||||
ref = x @ ref_weight.T
|
||||
|
||||
assert output.shape == (8, size_n)
|
||||
torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
_gpu_marlin_unsupported(),
|
||||
reason="Marlin is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("shape", [(200, 512), (4640, 512)])
|
||||
def test_awq_zp_marlin_padded_round_trip(shape):
|
||||
"""AWQ-style uint4 with runtime zero-points, padded the way
|
||||
MarlinLinearKernel does: padded columns rely on (q=0 - zp=0) * scale=0."""
|
||||
size_n, size_k = shape
|
||||
group_size = 128
|
||||
dtype = torch.float16
|
||||
quant_type = scalar_types.uint4
|
||||
device = torch.device("cuda")
|
||||
|
||||
weight = torch.randn(size_k, size_n, dtype=dtype, device=device) / size_k**0.5
|
||||
w_ref, q_w, s, zp = quantize_weights(
|
||||
weight, quant_type, group_size, zero_points=True
|
||||
)
|
||||
qweight = gptq_pack(q_w, quant_type.size_bits, size_k, size_n)
|
||||
|
||||
padded_n, padded_k = marlin_padded_nk(size_n, size_k, group_size)
|
||||
qweight = marlin_pad_qweight(qweight, size_n, size_k, padded_n, padded_k)
|
||||
marlin_qweight = ops.gptq_marlin_repack(
|
||||
b_q_weight=qweight,
|
||||
perm=torch.empty(0, dtype=torch.int, device=device),
|
||||
size_k=padded_k,
|
||||
size_n=padded_n,
|
||||
num_bits=quant_type.size_bits,
|
||||
)
|
||||
s = marlin_pad_scales(s, size_n, size_k, padded_n, padded_k, group_size)
|
||||
marlin_s = marlin_permute_scales(
|
||||
s, size_k=padded_k, size_n=padded_n, group_size=group_size
|
||||
)
|
||||
zp = marlin_pad_scales(zp, size_n, size_k, padded_n, padded_k, group_size)
|
||||
marlin_zp = marlin_zero_points(
|
||||
zp,
|
||||
size_k=padded_k // group_size,
|
||||
size_n=padded_n,
|
||||
num_bits=quant_type.size_bits,
|
||||
)
|
||||
|
||||
x = torch.randn(8, size_k, dtype=dtype, device=device)
|
||||
output = apply_gptq_marlin_linear(
|
||||
input=x,
|
||||
weight=marlin_qweight,
|
||||
weight_scale=marlin_s,
|
||||
weight_zp=marlin_zp,
|
||||
g_idx=marlin_make_empty_g_idx(device),
|
||||
g_idx_sort_indices=marlin_make_empty_g_idx(device),
|
||||
workspace=marlin_make_workspace_new(device),
|
||||
wtype=quant_type,
|
||||
output_size_per_partition=size_n,
|
||||
input_size_per_partition=size_k,
|
||||
is_k_full=True,
|
||||
)
|
||||
ref = x @ w_ref
|
||||
|
||||
assert output.shape == (8, size_n)
|
||||
torch.testing.assert_close(output, ref, rtol=2e-2, atol=2e-2)
|
||||
|
||||
|
||||
class _FakeLinear:
|
||||
def __init__(self, size_n, size_k, input_size=None):
|
||||
self.output_size_per_partition = size_n
|
||||
self.input_size_per_partition = size_k
|
||||
self.output_size = size_n
|
||||
self.input_size = input_size if input_size is not None else size_k
|
||||
|
||||
|
||||
def test_check_marlin_supports_layer_allow_tile_padding():
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
check_marlin_supports_layer,
|
||||
)
|
||||
|
||||
# Tile-misaligned but group-aligned: rejected strictly, allowed w/ padding
|
||||
layer = _FakeLinear(4640, 512, input_size=2048)
|
||||
assert not check_marlin_supports_layer(layer, 128)
|
||||
assert check_marlin_supports_layer(layer, 128, allow_tile_padding=True)
|
||||
assert check_marlin_supports_layer(layer, -1, allow_tile_padding=True)
|
||||
|
||||
# A group straddling the TP shard cannot be fixed by padding
|
||||
layer = _FakeLinear(4608, 4672, input_size=18688)
|
||||
assert not check_marlin_supports_layer(layer, 128, allow_tile_padding=True)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
_gpu_marlin_unsupported(),
|
||||
reason="Marlin is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("group_size", [-1, 32])
|
||||
@pytest.mark.parametrize("shape", [(96, 256, 8), (160, 512, 4)])
|
||||
def test_gptq_marlin_moe_padded_round_trip(shape, group_size):
|
||||
"""Pad a tile-misaligned MoE intermediate the way the WNA16 Marlin MoE prep
|
||||
does, run the real repack + fused_marlin_moe, and check against the
|
||||
dequantized reference. Symmetric int4's quantized zero decodes to -8, so the
|
||||
padded region only stays out of the output via the zero-padded scales.
|
||||
"""
|
||||
from tests.kernels.utils import torch_experts
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.fused_moe import fused_topk
|
||||
from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
|
||||
fused_marlin_moe,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.oracle.int_wna16 import (
|
||||
_pad_rows,
|
||||
_pad_w13_shard_cols,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
marlin_moe_padded_intermediate,
|
||||
marlin_moe_permute_scales,
|
||||
)
|
||||
|
||||
n, k, e = shape
|
||||
topk, m = 2, 33
|
||||
padded_n = marlin_moe_padded_intermediate(n, group_size)
|
||||
assert padded_n != n, "test should exercise padding"
|
||||
|
||||
dtype = torch.float16
|
||||
device = torch.device("cuda")
|
||||
quant_type = scalar_types.uint4b8
|
||||
bits = quant_type.size_bits
|
||||
pack = 32 // bits
|
||||
|
||||
a = torch.randn((m, k), device=device, dtype=dtype) / 10
|
||||
w1 = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / k**0.5
|
||||
w2 = torch.randn((e, k, n), device=device, dtype=dtype) / n**0.5
|
||||
|
||||
def quant(w, size_k, size_n):
|
||||
# w is (size_n, size_k); gptq expects (size_k, size_n).
|
||||
ref, q_w, s, _, _ = gptq_quantize_weights(
|
||||
w.T, quant_type, group_size, act_order=False
|
||||
)
|
||||
return ref, gptq_pack(q_w, bits, size_k, size_n), s
|
||||
|
||||
w13_qw, w13_s, w13_ref = [], [], []
|
||||
w2_qw, w2_s, w2_ref = [], [], []
|
||||
for i in range(e):
|
||||
ref, qw, s = quant(w1[i], k, 2 * n)
|
||||
w13_ref.append(ref.T) # (2n, k)
|
||||
w13_qw.append(qw)
|
||||
w13_s.append(s)
|
||||
ref, qw, s = quant(w2[i], n, k)
|
||||
w2_ref.append(ref.T) # (k, n)
|
||||
w2_qw.append(qw)
|
||||
w2_s.append(s)
|
||||
|
||||
w13_qweight = torch.stack(w13_qw)
|
||||
w2_qweight = torch.stack(w2_qw)
|
||||
w13_scales = torch.stack(w13_s)
|
||||
w2_scales = torch.stack(w2_s)
|
||||
w1_ref = torch.stack(w13_ref) # (e, 2n, k)
|
||||
w2_ref = torch.stack(w2_ref) # (e, k, n)
|
||||
|
||||
# Pad the intermediate via the production helpers.
|
||||
w13_qweight = _pad_w13_shard_cols(w13_qweight, n, padded_n)
|
||||
w2_qweight = _pad_rows(w2_qweight, padded_n // pack)
|
||||
w13_scales = _pad_w13_shard_cols(w13_scales, n, padded_n)
|
||||
if group_size > 0:
|
||||
w2_scales = _pad_rows(w2_scales, padded_n // group_size)
|
||||
|
||||
sort_idx = torch.empty((e, 0), dtype=torch.int32, device=device)
|
||||
marlin_w13 = ops.gptq_marlin_moe_repack(
|
||||
w13_qweight, sort_idx, w13_qweight.shape[1] * pack, w13_qweight.shape[2], bits
|
||||
)
|
||||
marlin_w2 = ops.gptq_marlin_moe_repack(
|
||||
w2_qweight, sort_idx, w2_qweight.shape[1] * pack, w2_qweight.shape[2], bits
|
||||
)
|
||||
group_or_pack = group_size if group_size != -1 else pack
|
||||
marlin_w13_s = marlin_moe_permute_scales(
|
||||
s=w13_scales, size_k=n, size_n=w13_scales.shape[2], group_size=group_size
|
||||
)
|
||||
marlin_w2_s = marlin_moe_permute_scales(
|
||||
s=w2_scales,
|
||||
size_k=w2_scales.shape[1] * group_or_pack,
|
||||
size_n=w2_scales.shape[2],
|
||||
group_size=group_size,
|
||||
)
|
||||
|
||||
score = torch.randn((m, e), device=device, dtype=dtype)
|
||||
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
|
||||
|
||||
marlin_out = fused_marlin_moe(
|
||||
a,
|
||||
marlin_w13,
|
||||
marlin_w2,
|
||||
None,
|
||||
None,
|
||||
marlin_w13_s,
|
||||
marlin_w2_s,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
quant_type_id=quant_type.id,
|
||||
global_num_experts=e,
|
||||
is_k_full=True,
|
||||
)
|
||||
with set_current_vllm_config(VllmConfig()):
|
||||
ref = torch_experts(
|
||||
a,
|
||||
w1_ref,
|
||||
w2_ref,
|
||||
topk_weight=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
global_num_experts=e,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(marlin_out, ref, atol=5e-2, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
current_platform.is_rocm(),
|
||||
reason="MoE Marlin is not selected on ROCm.",
|
||||
)
|
||||
def test_check_moe_marlin_supports_layer_padding():
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
check_moe_marlin_supports_layer,
|
||||
)
|
||||
|
||||
def make_layer(hidden, intermediate):
|
||||
layer = SimpleNamespace()
|
||||
layer.hidden_size = hidden
|
||||
layer.apply_router_weight_on_input = False
|
||||
layer.moe_config = SimpleNamespace(
|
||||
intermediate_size_per_partition_unpadded=intermediate
|
||||
)
|
||||
return layer
|
||||
|
||||
# group=32 with intermediate % 64 != 0: rejected strictly, accepted w/ padding
|
||||
layer = make_layer(4096, 96)
|
||||
assert not check_moe_marlin_supports_layer(layer, 32)
|
||||
assert check_moe_marlin_supports_layer(layer, 32, allow_tile_padding=True)
|
||||
# channelwise misaligned intermediate is paddable
|
||||
assert check_moe_marlin_supports_layer(layer, -1, allow_tile_padding=True)
|
||||
|
||||
# A group straddling the boundary cannot be fixed by padding
|
||||
layer = make_layer(4096, 176)
|
||||
assert not check_moe_marlin_supports_layer(layer, 128, allow_tile_padding=True)
|
||||
|
||||
# hidden_size is the MoE I/O extent and is never padded
|
||||
layer = make_layer(4090, 128)
|
||||
assert not check_moe_marlin_supports_layer(layer, 64, allow_tile_padding=True)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
_gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
|
||||
reason="FP8 Marlin is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("quant", ["channel", "tensor"])
|
||||
@pytest.mark.parametrize("shape", [(96, 256, 8), (160, 512, 4)])
|
||||
def test_fp8_marlin_moe_padded_round_trip(shape, quant):
|
||||
"""FP8 weight-only MoE: pad a tile-misaligned intermediate and check the
|
||||
real prepare + fused_marlin_moe against the dequantized reference."""
|
||||
from tests.kernels.utils import torch_experts
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.fused_moe import fused_topk
|
||||
from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
|
||||
fused_marlin_moe,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
marlin_moe_intermediate_size,
|
||||
marlin_moe_padded_intermediate,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
|
||||
prepare_fp8_moe_layer_for_marlin,
|
||||
)
|
||||
|
||||
n, k, e = shape
|
||||
topk, m = 2, 33
|
||||
fp8 = torch.float8_e4m3fn
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device("cuda")
|
||||
padded_n = marlin_moe_padded_intermediate(n, -1)
|
||||
assert padded_n != n
|
||||
|
||||
def q(w): # (out, in) -> fp8 weight, scale, dequant reference
|
||||
dim = None if quant == "tensor" else 1
|
||||
s = (w.abs().amax(dim, keepdim=dim is not None) / 448.0).clamp(min=1e-8)
|
||||
wq = (w / s).clamp(-448, 448).to(fp8)
|
||||
ref = wq.to(dtype) * s.to(dtype)
|
||||
s = s.reshape(1) if quant == "tensor" else s.squeeze(1)
|
||||
return wq, s, ref
|
||||
|
||||
a = torch.randn((m, k), device=device, dtype=dtype) / 10
|
||||
w1 = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / k**0.5
|
||||
w2 = torch.randn((e, k, n), device=device, dtype=dtype) / n**0.5
|
||||
w13_q, w13_s, w1_ref = zip(*(q(w1[i]) for i in range(e)))
|
||||
w2_q, w2_s, w2_ref = zip(*(q(w2[i]) for i in range(e)))
|
||||
|
||||
w13_weight, w2_weight = torch.stack(w13_q), torch.stack(w2_q)
|
||||
layer = SimpleNamespace(
|
||||
num_experts=e,
|
||||
hidden_size=k,
|
||||
intermediate_size_per_partition=n,
|
||||
orig_dtype=dtype,
|
||||
w13_weight=w13_weight,
|
||||
)
|
||||
pw13, pw2, ps13, ps2 = prepare_fp8_moe_layer_for_marlin(
|
||||
layer, w13_weight, w2_weight, torch.stack(w13_s), torch.stack(w2_s)
|
||||
)
|
||||
assert marlin_moe_intermediate_size(pw13, pw2) == padded_n
|
||||
|
||||
score = torch.randn((m, e), device=device, dtype=dtype)
|
||||
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
|
||||
out = fused_marlin_moe(
|
||||
a,
|
||||
pw13,
|
||||
pw2,
|
||||
None,
|
||||
None,
|
||||
ps13,
|
||||
ps2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
quant_type_id=scalar_types.float8_e4m3fn.id,
|
||||
global_num_experts=e,
|
||||
is_k_full=True,
|
||||
workspace=layer.workspace,
|
||||
)
|
||||
with set_current_vllm_config(VllmConfig()):
|
||||
ref = torch_experts(
|
||||
a,
|
||||
torch.stack(w1_ref),
|
||||
torch.stack(w2_ref),
|
||||
topk_weight=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
global_num_experts=e,
|
||||
)
|
||||
torch.testing.assert_close(out, ref, atol=8e-2, rtol=0)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
_gpu_marlin_unsupported() or not is_fp8_marlin_supported(),
|
||||
reason="FP8 Marlin is not supported on this GPU type.",
|
||||
)
|
||||
@pytest.mark.parametrize("shape", [(96, 256, 8), (160, 512, 4)])
|
||||
def test_mxfp8_marlin_moe_padded_round_trip(shape):
|
||||
"""MXFP8 weight-only MoE round-trip at a tile-misaligned intermediate, with
|
||||
unit e8m0 scales so the reference is the exact fp8 dequant."""
|
||||
from tests.kernels.utils import torch_experts
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.model_executor.layers.fused_moe import fused_topk
|
||||
from vllm.model_executor.layers.fused_moe.experts.marlin_moe import (
|
||||
fused_marlin_moe,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils import (
|
||||
marlin_moe_intermediate_size,
|
||||
marlin_moe_padded_intermediate,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.marlin_utils_fp8 import (
|
||||
prepare_mxfp8_moe_layer_for_marlin,
|
||||
)
|
||||
|
||||
n, k, e = shape
|
||||
topk, m, gs, e8m0_one = 2, 33, 32, 127
|
||||
fp8 = torch.float8_e4m3fn
|
||||
dtype = torch.bfloat16
|
||||
device = torch.device("cuda")
|
||||
padded_n = marlin_moe_padded_intermediate(n, gs)
|
||||
assert padded_n != n
|
||||
|
||||
a = torch.randn((m, k), device=device, dtype=dtype) / 10
|
||||
w13_weight = torch.randn((e, 2 * n, k), device=device, dtype=dtype) / k**0.5
|
||||
w2_weight = torch.randn((e, k, n), device=device, dtype=dtype) / n**0.5
|
||||
w13_weight = w13_weight.clamp(-448, 448).to(fp8)
|
||||
w2_weight = w2_weight.clamp(-448, 448).to(fp8)
|
||||
w13_scale = torch.full(
|
||||
(e, 2 * n, k // gs), e8m0_one, dtype=torch.uint8, device=device
|
||||
)
|
||||
w2_scale = torch.full((e, k, n // gs), e8m0_one, dtype=torch.uint8, device=device)
|
||||
|
||||
layer = SimpleNamespace(
|
||||
num_experts=e, hidden_size=k, intermediate_size_per_partition=n
|
||||
)
|
||||
with set_current_vllm_config(VllmConfig()):
|
||||
pw13, pw2, ps13, ps2 = prepare_mxfp8_moe_layer_for_marlin(
|
||||
layer, w13_weight, w2_weight, w13_scale, w2_scale
|
||||
)
|
||||
assert marlin_moe_intermediate_size(pw13, pw2) == padded_n
|
||||
|
||||
score = torch.randn((m, e), device=device, dtype=dtype)
|
||||
topk_weights, topk_ids, _ = fused_topk(a, score, topk, False)
|
||||
out = fused_marlin_moe(
|
||||
a,
|
||||
pw13,
|
||||
pw2,
|
||||
None,
|
||||
None,
|
||||
ps13,
|
||||
ps2,
|
||||
topk_weights,
|
||||
topk_ids,
|
||||
quant_type_id=scalar_types.float8_e4m3fn.id,
|
||||
global_num_experts=e,
|
||||
is_k_full=True,
|
||||
workspace=layer.workspace,
|
||||
)
|
||||
with set_current_vllm_config(VllmConfig()):
|
||||
ref = torch_experts(
|
||||
a,
|
||||
w13_weight.to(dtype),
|
||||
w2_weight.to(dtype),
|
||||
topk_weight=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
global_num_experts=e,
|
||||
)
|
||||
torch.testing.assert_close(out, ref, atol=8e-2, rtol=0)
|
||||
@@ -0,0 +1,304 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Copyright (C) 2025 Roberto L. Castro (Roberto.LopezCastro@ist.ac.at).
|
||||
# All Rights Reserved.
|
||||
#
|
||||
# Licensed 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.
|
||||
#
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from compressed_tensors.transform.utils.hadamard import deterministic_hadamard_matrix
|
||||
|
||||
from vllm._custom_ops import fusedQuantizeMx, matmul_mxf4_bf16_tn
|
||||
from vllm.model_executor.layers.quantization.qutlass_utils import to_blocked
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA required for these tests.", allow_module_level=True)
|
||||
|
||||
if not (
|
||||
current_platform.has_device_capability(100)
|
||||
or current_platform.has_device_capability(120)
|
||||
):
|
||||
pytest.skip(
|
||||
reason="Tests require compute capability 10.0 (100) or 12.0 (120).",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
# ----- Helpers -----
|
||||
def get_hadamard_matrix(group_size: int, dtype: torch.dtype, device: torch.device):
|
||||
return (
|
||||
deterministic_hadamard_matrix(group_size, dtype=dtype, device=device)
|
||||
* group_size**-0.5
|
||||
)
|
||||
|
||||
|
||||
def _rtne_fp4(x: torch.Tensor):
|
||||
device = x.device
|
||||
grid = torch.tensor(
|
||||
[
|
||||
-6.0,
|
||||
-4.0,
|
||||
-3.0,
|
||||
-2.0,
|
||||
-1.5,
|
||||
-1.0,
|
||||
-0.5,
|
||||
-0.0,
|
||||
0.0,
|
||||
0.5,
|
||||
1.0,
|
||||
1.5,
|
||||
2.0,
|
||||
3.0,
|
||||
4.0,
|
||||
6.0,
|
||||
],
|
||||
dtype=x.dtype,
|
||||
device=x.device,
|
||||
)
|
||||
grid_int = torch.tensor(
|
||||
[-1, -2, -3, -4, -5, -6, -7, -8, 0, 1, 2, 3, 4, 5, 6, 7],
|
||||
dtype=torch.uint8,
|
||||
device=device,
|
||||
)
|
||||
inds = torch.bucketize(x, grid)
|
||||
lo, hi = (inds - 1).clamp(min=0, max=15), inds.clamp(min=0, max=15)
|
||||
g_lo, g_hi = grid[lo], grid[hi]
|
||||
pick_hi = (g_hi - x < x - g_lo) | (g_hi - x == x - g_lo) & (grid_int[hi] % 2 == 0)
|
||||
y = torch.where(pick_hi, g_hi, g_lo)
|
||||
y_int = torch.where(pick_hi, grid_int[hi], grid_int[lo])
|
||||
y_int_packed = (y_int[..., 1::2] & 0xF) << 4 | y_int[..., ::2] & 0xF
|
||||
return y, y_int_packed
|
||||
|
||||
|
||||
def _dq_fp4(x_e2m1: torch.Tensor, x_e8m0: torch.Tensor, alpha: float):
|
||||
device = x_e2m1.device
|
||||
|
||||
x_e2m1_i32 = x_e2m1.view(dtype=torch.uint8).to(dtype=torch.int32)
|
||||
x_e2m1_unpacked = torch.stack(
|
||||
[x_e2m1_i32 & 0xF, (x_e2m1_i32 >> 4) & 0xF], dim=-1
|
||||
).flatten(start_dim=-2)
|
||||
|
||||
grid_dq = torch.tensor(
|
||||
[
|
||||
0.0,
|
||||
0.5,
|
||||
1.0,
|
||||
1.5,
|
||||
2.0,
|
||||
3.0,
|
||||
4.0,
|
||||
6.0,
|
||||
-0.0,
|
||||
-0.5,
|
||||
-1.0,
|
||||
-1.5,
|
||||
-2.0,
|
||||
-3.0,
|
||||
-4.0,
|
||||
-6.0,
|
||||
],
|
||||
dtype=torch.float64,
|
||||
device=device,
|
||||
)
|
||||
x_fp4_dq = grid_dq[x_e2m1_unpacked]
|
||||
scales_dq = x_e8m0.to(torch.float64)
|
||||
|
||||
x_dq = (x_fp4_dq.unflatten(dim=-1, sizes=(-1, 32)) * scales_dq[..., None]).flatten(
|
||||
start_dim=-2
|
||||
) / alpha
|
||||
return x_dq, x_fp4_dq, scales_dq
|
||||
|
||||
|
||||
def _unpack_mask(clip_mask: torch.Tensor) -> torch.Tensor:
|
||||
clip_mask_unpacked_dq = torch.zeros(
|
||||
*clip_mask.shape[:-1],
|
||||
clip_mask.size(-1) * 8,
|
||||
dtype=torch.bool,
|
||||
device=clip_mask.device,
|
||||
)
|
||||
for i in range(8):
|
||||
clip_mask_unpacked_dq[..., i::8] = (clip_mask >> i) & 1
|
||||
return clip_mask_unpacked_dq
|
||||
|
||||
|
||||
def _forward_quantize_ref(
|
||||
x: torch.Tensor, h: torch.Tensor, rot_size: int, quest: bool = True
|
||||
):
|
||||
device = x.device
|
||||
xh_ref64 = (
|
||||
x.unflatten(dim=-1, sizes=(-1, rot_size)).to(dtype=torch.float64)
|
||||
@ h.reshape(rot_size, rot_size).to(dtype=torch.float64)
|
||||
).flatten(start_dim=-2)
|
||||
|
||||
if quest:
|
||||
scales_ref64_ = (
|
||||
xh_ref64.unflatten(dim=-1, sizes=(-1, 32)).std(dim=-1, correction=0)
|
||||
* (2.92247856 / 6.0)
|
||||
+ 1e-8
|
||||
)
|
||||
else:
|
||||
abs_max = xh_ref64.unflatten(dim=-1, sizes=(-1, 32)).abs().amax(dim=-1)
|
||||
scales_ref64_ = abs_max + 1e-8
|
||||
|
||||
xh_e8m0_ref = scales_ref64_.log2().floor().exp2().to(dtype=torch.float8_e8m0fnu)
|
||||
scales_ref64 = xh_e8m0_ref.to(dtype=torch.float64)
|
||||
|
||||
xh_scaled_ref64 = (
|
||||
xh_ref64.unflatten(dim=-1, sizes=(-1, 32)) / scales_ref64[..., None]
|
||||
).flatten(start_dim=-2)
|
||||
if not quest:
|
||||
xh_scaled_ref64 *= 3
|
||||
|
||||
clip_mask_unpacked_ref = xh_scaled_ref64.abs() < 6.0
|
||||
clip_mask_ref = torch.zeros(
|
||||
*x.shape[:-1], x.size(-1) // 8, dtype=torch.uint8, device=device
|
||||
)
|
||||
for i in range(8):
|
||||
clip_mask_ref |= clip_mask_unpacked_ref[..., i::8].to(dtype=torch.uint8) << i
|
||||
|
||||
xh_fp4_ref, xh_e2m1_ref = _rtne_fp4(xh_scaled_ref64)
|
||||
xh_dq, xh_fp4_dq, scales_dq = _dq_fp4(
|
||||
xh_e2m1_ref, xh_e8m0_ref, alpha=1.0 if quest else 3.0
|
||||
)
|
||||
clip_mask_unpacked_dq = _unpack_mask(clip_mask_ref)
|
||||
|
||||
assert xh_fp4_dq.equal(xh_fp4_ref)
|
||||
assert scales_dq.equal(scales_ref64)
|
||||
assert clip_mask_unpacked_dq.equal(clip_mask_unpacked_ref)
|
||||
|
||||
return (
|
||||
xh_dq,
|
||||
clip_mask_unpacked_ref,
|
||||
(xh_e2m1_ref, xh_e8m0_ref, clip_mask_ref),
|
||||
)
|
||||
|
||||
|
||||
DTYPE = torch.bfloat16
|
||||
DEVICE = torch.device("cuda:0")
|
||||
|
||||
ROT_SIZES = [32, 64, 128]
|
||||
SEEDS = [0]
|
||||
BATCHES = [1, 16]
|
||||
|
||||
LLAMA_MODELS = {
|
||||
"7B": [(4096, 3 * 4096), (4096, 4096), (4096, 2 * 10752), (10752, 4096)],
|
||||
"13B": [(5120, 3 * 5120), (5120, 5120), (5120, 2 * 13568), (13568, 5120)],
|
||||
"33B": [(6656, 3 * 6656), (6656, 6656), (6656, 2 * 17664), (17664, 6656)],
|
||||
"70B": [(8192, 3 * 8192), (8192, 8192), (8192, 2 * 21760), (21760, 8192)],
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _seed_each_test():
|
||||
set_random_seed(0)
|
||||
np.random.seed(0)
|
||||
torch.random.manual_seed(0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rot_size", ROT_SIZES)
|
||||
@torch.inference_mode()
|
||||
def test_fused_quantization_absmax(rot_size: int):
|
||||
dtype, device = DTYPE, DEVICE
|
||||
h = get_hadamard_matrix(rot_size, dtype, device)
|
||||
x = torch.randn(2, 4096, 4096, dtype=dtype, device=device) * 25.0
|
||||
|
||||
xh_dq_ref, _, _ = _forward_quantize_ref(x, h, rot_size, quest=False)
|
||||
xh_e2m1, xh_e8m0 = fusedQuantizeMx(x, h, method="abs_max")
|
||||
xh_e8m0 = xh_e8m0.reshape(2, 4096, 4096 // 32)
|
||||
xh_dq, *_ = _dq_fp4(xh_e2m1, xh_e8m0, alpha=3.0)
|
||||
|
||||
torch.testing.assert_close(xh_dq, xh_dq_ref, rtol=0.34, atol=100)
|
||||
assert (xh_dq != xh_dq_ref).float().mean() <= 1e-4
|
||||
|
||||
m, n, k = 1, 504, 4096
|
||||
a = torch.randn(m, k, dtype=dtype, device=device) * 25.0
|
||||
b = torch.randn(n, k, dtype=dtype, device=device) * 25.0
|
||||
|
||||
a_e2m1, a_e8m0 = fusedQuantizeMx(a, h, method="abs_max")
|
||||
b_e2m1, b_e8m0 = fusedQuantizeMx(b, h, method="abs_max")
|
||||
a_dq, *_ = _dq_fp4(a_e2m1, a_e8m0[:m, :k], alpha=1.0)
|
||||
b_dq, *_ = _dq_fp4(b_e2m1, b_e8m0[:n, :k], alpha=1.0)
|
||||
out_ref = a_dq @ b_dq.transpose(-2, -1)
|
||||
|
||||
a_scale_block = to_blocked(a_e8m0, backend="triton")
|
||||
b_scale_block = to_blocked(b_e8m0, backend="triton")
|
||||
alpha = torch.tensor([1.0], device=device)
|
||||
out = matmul_mxf4_bf16_tn(a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha)
|
||||
assert out.equal(out_ref.to(dtype=out.dtype))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rot_size", ROT_SIZES)
|
||||
@torch.inference_mode()
|
||||
def test_fused_quantization_quest(rot_size: int):
|
||||
dtype, device = DTYPE, DEVICE
|
||||
h = get_hadamard_matrix(rot_size, dtype, device)
|
||||
x = torch.randn(2, 4096, 4096, dtype=dtype, device=device) * 25.0
|
||||
|
||||
xh_dq_ref, _, _ = _forward_quantize_ref(x, h, rot_size, quest=True)
|
||||
xh_e2m1, xh_e8m0 = fusedQuantizeMx(x, h, method="quest")
|
||||
xh_e8m0 = xh_e8m0.reshape(2, 4096, 4096 // 32)
|
||||
xh_dq, *_ = _dq_fp4(xh_e2m1, xh_e8m0, alpha=1.0)
|
||||
|
||||
torch.testing.assert_close(xh_dq, xh_dq_ref, rtol=0.34, atol=100)
|
||||
assert (xh_dq != xh_dq_ref).float().mean() <= 1e-4
|
||||
|
||||
m, n, k = 504, 504, 2048
|
||||
a = torch.randn(m, k, dtype=dtype, device=device) * 25.0
|
||||
b = torch.randn(n, k, dtype=dtype, device=device) * 25.0
|
||||
|
||||
a_e2m1, a_e8m0 = fusedQuantizeMx(a, h, method="quest")
|
||||
b_e2m1, b_e8m0 = fusedQuantizeMx(b, h, method="quest")
|
||||
a_dq, *_ = _dq_fp4(a_e2m1, a_e8m0[:m, :k], alpha=1.0)
|
||||
b_dq, *_ = _dq_fp4(b_e2m1, b_e8m0[:n, :k], alpha=1.0)
|
||||
out_ref = a_dq @ b_dq.transpose(-2, -1)
|
||||
|
||||
a_scale_block = to_blocked(a_e8m0, backend="triton")
|
||||
b_scale_block = to_blocked(b_e8m0, backend="triton")
|
||||
alpha = torch.tensor([1.0], device=device)
|
||||
out = matmul_mxf4_bf16_tn(a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha)
|
||||
assert out.equal(out_ref.to(dtype=out.dtype))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", list(LLAMA_MODELS.keys()))
|
||||
@pytest.mark.parametrize("layer_idx", [0, 1, 2, 3])
|
||||
@pytest.mark.parametrize("batch", [1, 16])
|
||||
@pytest.mark.parametrize("had_size", ROT_SIZES)
|
||||
@torch.inference_mode()
|
||||
def test_llama_shapes(model: str, layer_idx: int, batch: int, had_size: int):
|
||||
dtype, device = DTYPE, DEVICE
|
||||
m = batch
|
||||
k, n = LLAMA_MODELS[model][layer_idx]
|
||||
|
||||
h = get_hadamard_matrix(had_size, dtype, device)
|
||||
|
||||
a = torch.rand(m, k, dtype=dtype, device=device) * 25.0
|
||||
b = torch.rand(n, k, dtype=dtype, device=device) * 25.0
|
||||
|
||||
a_e2m1, a_e8m0 = fusedQuantizeMx(a, h, method="quest")
|
||||
b_e2m1, b_e8m0 = fusedQuantizeMx(b, h, method="quest")
|
||||
|
||||
a_dq, *_ = _dq_fp4(a_e2m1, a_e8m0[:m, :k], alpha=1.0)
|
||||
b_dq, *_ = _dq_fp4(b_e2m1, b_e8m0[:n, :k], alpha=1.0)
|
||||
out_ref = a_dq @ b_dq.transpose(-2, -1)
|
||||
|
||||
a_scale_block = to_blocked(a_e8m0, backend="triton")
|
||||
b_scale_block = to_blocked(b_e8m0, backend="triton")
|
||||
alpha = torch.tensor([1.0], device=device)
|
||||
out = matmul_mxf4_bf16_tn(a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha)
|
||||
assert out.equal(out_ref.to(dtype=out.dtype))
|
||||
@@ -0,0 +1,83 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
Tests that triton_kernel_moe_forward correctly applies expert_map
|
||||
remapping when expert parallelism (EP) is enabled.
|
||||
|
||||
When expert_map is provided, global expert IDs are remapped to local IDs
|
||||
via topk + expert_map remap + make_routing_data before building routing
|
||||
structures, and the expert_map passed downstream to triton_kernel_fused_experts
|
||||
is None (already applied).
|
||||
"""
|
||||
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
class TestTritonMoeForwardExpertMap:
|
||||
"""Test that triton_kernel_moe_forward applies expert_map remapping
|
||||
when expert_map is provided (EP active)."""
|
||||
|
||||
def test_expert_map_remap(self):
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
mock_expert_map = torch.tensor([0, -1, 1, -1], device=device)
|
||||
|
||||
from vllm.utils.import_utils import import_triton_kernels
|
||||
|
||||
import_triton_kernels()
|
||||
|
||||
mock_routing_data = MagicMock()
|
||||
mock_gather = MagicMock()
|
||||
mock_scatter = MagicMock()
|
||||
|
||||
with (
|
||||
patch("triton_kernels.topk.topk") as mock_topk,
|
||||
patch(
|
||||
"vllm.model_executor.layers.fused_moe.experts."
|
||||
"gpt_oss_triton_kernels_moe.make_routing_data"
|
||||
) as mock_make_routing,
|
||||
patch(
|
||||
"vllm.model_executor.layers.fused_moe.experts."
|
||||
"gpt_oss_triton_kernels_moe.triton_kernel_fused_experts"
|
||||
) as mock_fused_experts,
|
||||
):
|
||||
from vllm.model_executor.layers.fused_moe.experts.gpt_oss_triton_kernels_moe import ( # noqa: E501
|
||||
triton_kernel_moe_forward,
|
||||
)
|
||||
|
||||
sparse_result = MagicMock()
|
||||
sparse_result.indx = torch.tensor([[0, 2]], dtype=torch.int32)
|
||||
sparse_result.vals = torch.tensor([[0.6, 0.4]])
|
||||
mock_topk.return_value = sparse_result
|
||||
|
||||
mock_make_routing.return_value = (
|
||||
mock_routing_data,
|
||||
mock_gather,
|
||||
mock_scatter,
|
||||
)
|
||||
|
||||
mock_fused_experts.return_value = torch.zeros((1, 8), device=device)
|
||||
|
||||
hidden = torch.randn((1, 8), device=device)
|
||||
w1 = torch.randn((2, 8, 16), device=device)
|
||||
w2 = torch.randn((2, 8, 8), device=device)
|
||||
logits = torch.randn((1, 4), device=device)
|
||||
|
||||
triton_kernel_moe_forward(
|
||||
hidden_states=hidden,
|
||||
w1=w1,
|
||||
w2=w2,
|
||||
gating_output=logits,
|
||||
topk=2,
|
||||
renormalize=True,
|
||||
expert_map=mock_expert_map,
|
||||
)
|
||||
|
||||
mock_topk.assert_called_once()
|
||||
mock_make_routing.assert_called_once()
|
||||
|
||||
# expert_map should be None in the fused_experts call
|
||||
# (already applied).
|
||||
call_kwargs = mock_fused_experts.call_args
|
||||
assert call_kwargs[1].get("expert_map") is None or (len(call_kwargs[0]) > 0)
|
||||
@@ -0,0 +1,772 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from typing import cast
|
||||
|
||||
import huggingface_hub
|
||||
import pytest
|
||||
import torch
|
||||
from safetensors import safe_open
|
||||
|
||||
import vllm.model_executor.layers.fused_moe.modular_kernel as mk
|
||||
from vllm.model_executor.layers.fused_moe.activation import MoEActivation
|
||||
from vllm.model_executor.layers.fused_moe.config import (
|
||||
FusedMoEConfig,
|
||||
FusedMoEParallelConfig,
|
||||
FusedMoEQuantConfig,
|
||||
RoutingMethodType,
|
||||
nvfp4_moe_quant_config,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.nvfp4_emulation_moe import (
|
||||
Nvfp4QuantizationEmulationTritonExperts,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.experts.triton_moe import TritonExperts
|
||||
from vllm.model_executor.layers.quantization.utils import (
|
||||
nvfp4_emulation_utils,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
|
||||
dequantize_to_dtype,
|
||||
ref_nvfp4_quant_dequant,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
QuantKey,
|
||||
kNvfp4Dynamic,
|
||||
kNvfp4Static,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.triton_utils import triton
|
||||
|
||||
if current_platform.is_rocm():
|
||||
from vllm.platforms.rocm import on_gfx950
|
||||
else:
|
||||
|
||||
def on_gfx950() -> bool:
|
||||
return False
|
||||
|
||||
|
||||
MOE_MODEL_CONFIGS = {
|
||||
"nvidia/Qwen3-30B-A3B-NVFP4": {
|
||||
"shards": ["model-00001-of-00004.safetensors"],
|
||||
"expert_prefix": "model.layers.9.mlp.experts.",
|
||||
# Position of the expert index in the dot-split key.
|
||||
"expert_idx_pos": 5,
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def loaded_model_files():
|
||||
return {
|
||||
model_id: huggingface_hub.snapshot_download(
|
||||
repo_id=model_id, allow_patterns=config["shards"]
|
||||
)
|
||||
for model_id, config in MOE_MODEL_CONFIGS.items()
|
||||
}
|
||||
|
||||
|
||||
class Nvfp4QuantizationEmulationTritonExpertsReference(TritonExperts):
|
||||
"""
|
||||
Extension of TritonExperts to support emulated NVFP4 MoE experts.
|
||||
|
||||
It may be used for NVFP4 models when the device does not have
|
||||
native support for this dtype.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
moe_config: FusedMoEConfig,
|
||||
quant_config: FusedMoEQuantConfig,
|
||||
):
|
||||
super().__init__(moe_config, quant_config)
|
||||
|
||||
# `TritonExperts.apply` expects pre-dequantized weights,
|
||||
# which we handle in `apply` below.
|
||||
self.w1_scale_val = self.quant_config.w1_scale
|
||||
self.w2_scale_val = self.quant_config.w2_scale
|
||||
|
||||
self.quant_config._w1.scale = None
|
||||
self.quant_config._w2.scale = None
|
||||
|
||||
self.quantization_emulation = True
|
||||
|
||||
@property
|
||||
def quant_dtype(self) -> torch.dtype | str | None:
|
||||
return "nvfp4"
|
||||
|
||||
@property
|
||||
def a1_scale(self) -> torch.Tensor | None:
|
||||
return self.quant_config.a1_gscale
|
||||
|
||||
@property
|
||||
def expects_unquantized_inputs(self) -> bool:
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
def _supports_quant_scheme(
|
||||
weight_key: QuantKey | None,
|
||||
activation_key: QuantKey | None,
|
||||
) -> bool:
|
||||
return (weight_key, activation_key) == (kNvfp4Static, kNvfp4Dynamic)
|
||||
|
||||
def apply(
|
||||
self,
|
||||
output: torch.Tensor,
|
||||
hidden_states: torch.Tensor,
|
||||
w1: torch.Tensor,
|
||||
w2: torch.Tensor,
|
||||
topk_weights: torch.Tensor,
|
||||
topk_ids: torch.Tensor,
|
||||
activation: MoEActivation,
|
||||
global_num_experts: int,
|
||||
expert_map: torch.Tensor | None,
|
||||
a1q_scale: torch.Tensor | None,
|
||||
a2_scale: torch.Tensor | None,
|
||||
workspace13: torch.Tensor,
|
||||
workspace2: torch.Tensor,
|
||||
expert_tokens_meta: mk.ExpertTokensMetadata | None,
|
||||
apply_router_weight_on_input: bool,
|
||||
):
|
||||
assert w1.dtype == torch.uint8
|
||||
assert w2.dtype == torch.uint8
|
||||
|
||||
# Dequantize w1 from packed NVFP4 to fp16/bf16
|
||||
w13_global_scale = self.quant_config.g1_alphas
|
||||
|
||||
w1_dequant = dequantize_to_dtype(
|
||||
tensor_fp4=w1,
|
||||
tensor_sf=self.w1_scale_val,
|
||||
global_scale=w13_global_scale,
|
||||
dtype=hidden_states.dtype,
|
||||
block_size=16,
|
||||
swizzle=False,
|
||||
)
|
||||
|
||||
# Dequantize w2 from packed NVFP4 to fp16/bf16
|
||||
w2_global_scale = self.quant_config.g2_alphas
|
||||
|
||||
w2_dequant = dequantize_to_dtype(
|
||||
tensor_fp4=w2,
|
||||
tensor_sf=self.w2_scale_val,
|
||||
global_scale=w2_global_scale,
|
||||
dtype=hidden_states.dtype,
|
||||
block_size=16,
|
||||
swizzle=False,
|
||||
)
|
||||
|
||||
super().apply(
|
||||
output=output,
|
||||
hidden_states=hidden_states,
|
||||
w1=w1_dequant,
|
||||
w2=w2_dequant,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
activation=activation,
|
||||
global_num_experts=global_num_experts,
|
||||
expert_map=expert_map,
|
||||
a1q_scale=None,
|
||||
a2_scale=self.quant_config.a2_gscale,
|
||||
workspace13=workspace13,
|
||||
workspace2=workspace2,
|
||||
expert_tokens_meta=expert_tokens_meta,
|
||||
apply_router_weight_on_input=apply_router_weight_on_input,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("config_kwargs", "expected_reason"),
|
||||
[
|
||||
({"has_bias": True}, "kernel does not support bias"),
|
||||
({"is_lora_enabled": True}, "kernel does not support LoRA"),
|
||||
],
|
||||
)
|
||||
def test_nvfp4_emulation_support_check_rejects_bias_and_lora(
|
||||
config_kwargs: dict[str, bool],
|
||||
expected_reason: str,
|
||||
) -> None:
|
||||
moe_config = FusedMoEConfig(
|
||||
num_experts=2,
|
||||
experts_per_token=1,
|
||||
hidden_dim=16,
|
||||
intermediate_size=16,
|
||||
num_local_experts=2,
|
||||
num_logical_experts=2,
|
||||
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
|
||||
activation=MoEActivation.SILU,
|
||||
in_dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
routing_method=RoutingMethodType.TopK,
|
||||
**config_kwargs,
|
||||
)
|
||||
|
||||
supported, reason = Nvfp4QuantizationEmulationTritonExperts.is_supported_config(
|
||||
Nvfp4QuantizationEmulationTritonExperts,
|
||||
moe_config,
|
||||
kNvfp4Static,
|
||||
kNvfp4Dynamic,
|
||||
mk.FusedMoEActivationFormat.Standard,
|
||||
)
|
||||
|
||||
assert not supported
|
||||
assert reason == expected_reason
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="Triton NVFP4 kernel requires CUDA.",
|
||||
)
|
||||
def test_triton_dequantize_nvfp4(monkeypatch, loaded_model_files) -> None:
|
||||
"""Test the Triton dequantization kernel against the CPU reference
|
||||
using real NVFP4 weights from a checkpoint.
|
||||
|
||||
Tests both 2D (attention projection) and 3D (stacked MoE experts).
|
||||
"""
|
||||
checkpoint_path = loaded_model_files["nvidia/Qwen3-30B-A3B-NVFP4"]
|
||||
shards = cast(list[str], MOE_MODEL_CONFIGS["nvidia/Qwen3-30B-A3B-NVFP4"]["shards"])
|
||||
shard_path = f"{checkpoint_path}/{shards[0]}"
|
||||
block_size = 16
|
||||
|
||||
with safe_open(shard_path, framework="pt", device="cpu") as f:
|
||||
all_keys = list(f.keys())
|
||||
|
||||
# 2D case: attention projection
|
||||
tensor_fp4_2d = f.get_tensor("model.layers.9.self_attn.k_proj.weight")
|
||||
tensor_sf_2d = f.get_tensor("model.layers.9.self_attn.k_proj.weight_scale")
|
||||
global_scale_2d = f.get_tensor("model.layers.9.self_attn.k_proj.weight_scale_2")
|
||||
|
||||
# 3D case: stack ALL experts for layer 9 up_proj
|
||||
expert_prefix = "model.layers.9.mlp.experts."
|
||||
expert_indices = sorted(
|
||||
int(key.split(".")[5])
|
||||
for key in all_keys
|
||||
if key.startswith(expert_prefix) and key.endswith(".up_proj.weight")
|
||||
)
|
||||
assert len(expert_indices) > 0
|
||||
|
||||
all_fp4 = []
|
||||
all_sf = []
|
||||
all_global_scale = []
|
||||
for index in expert_indices:
|
||||
name = f"{expert_prefix}{index}.up_proj"
|
||||
all_fp4.append(f.get_tensor(f"{name}.weight"))
|
||||
all_sf.append(f.get_tensor(f"{name}.weight_scale"))
|
||||
all_global_scale.append(f.get_tensor(f"{name}.weight_scale_2"))
|
||||
|
||||
tensor_fp4_3d = torch.stack(all_fp4)
|
||||
tensor_sf_3d = torch.stack(all_sf)
|
||||
global_scale_3d = torch.stack(all_global_scale)
|
||||
|
||||
test_cases = [
|
||||
("2D base", tensor_fp4_2d, tensor_sf_2d, global_scale_2d),
|
||||
(
|
||||
"2D 2x rows",
|
||||
tensor_fp4_2d.repeat(2, 1),
|
||||
tensor_sf_2d.repeat(2, 1),
|
||||
global_scale_2d,
|
||||
),
|
||||
(
|
||||
"2D 4x rows",
|
||||
tensor_fp4_2d.repeat(4, 1),
|
||||
tensor_sf_2d.repeat(4, 1),
|
||||
global_scale_2d,
|
||||
),
|
||||
(
|
||||
"2D 2x cols",
|
||||
tensor_fp4_2d.repeat(1, 2),
|
||||
tensor_sf_2d.repeat(1, 2),
|
||||
global_scale_2d,
|
||||
),
|
||||
("3D base", tensor_fp4_3d, tensor_sf_3d, global_scale_3d),
|
||||
(
|
||||
"3D 2x experts",
|
||||
tensor_fp4_3d.repeat(2, 1, 1),
|
||||
tensor_sf_3d.repeat(2, 1, 1),
|
||||
global_scale_3d.repeat(2),
|
||||
),
|
||||
(
|
||||
"3D 2x rows",
|
||||
tensor_fp4_3d.repeat(1, 2, 1),
|
||||
tensor_sf_3d.repeat(1, 2, 1),
|
||||
global_scale_3d,
|
||||
),
|
||||
(
|
||||
"3D 2x cols",
|
||||
tensor_fp4_3d.repeat(1, 1, 2),
|
||||
tensor_sf_3d.repeat(1, 1, 2),
|
||||
global_scale_3d,
|
||||
),
|
||||
]
|
||||
|
||||
quantiles = [0.5, 0.001, 0.999]
|
||||
|
||||
# Move the E2M1 lookup table to CUDA ahead of time, as would normally
|
||||
# happen during model loading (process_weights_after_loading). Both the
|
||||
# Triton and PyTorch reference paths run on CUDA.
|
||||
nvfp4_emulation_utils.kE2M1ToFloat_handle.val = (
|
||||
nvfp4_emulation_utils.kE2M1ToFloat_handle.val.cuda()
|
||||
)
|
||||
|
||||
for label, tensor_fp4, tensor_sf, global_scale in test_cases:
|
||||
fp4_cuda = tensor_fp4.cuda()
|
||||
sf_cuda = tensor_sf.cuda()
|
||||
gs_cuda = global_scale.cuda()
|
||||
|
||||
# Triton path
|
||||
triton_result = dequantize_to_dtype(
|
||||
fp4_cuda,
|
||||
sf_cuda,
|
||||
gs_cuda,
|
||||
torch.bfloat16,
|
||||
block_size,
|
||||
swizzle=False,
|
||||
)
|
||||
|
||||
# Reference path (PyTorch ops on CUDA, Triton dispatch disabled)
|
||||
with monkeypatch.context() as m:
|
||||
m.setattr(
|
||||
nvfp4_emulation_utils.current_platform,
|
||||
"is_cuda_alike",
|
||||
lambda: False,
|
||||
)
|
||||
reference = dequantize_to_dtype(
|
||||
fp4_cuda,
|
||||
sf_cuda,
|
||||
gs_cuda,
|
||||
torch.bfloat16,
|
||||
block_size,
|
||||
swizzle=False,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(triton_result, reference, atol=0, rtol=0)
|
||||
|
||||
# Benchmark
|
||||
shape = list(tensor_fp4.shape)
|
||||
|
||||
def _triton_bench(
|
||||
fp4_cuda=fp4_cuda,
|
||||
scale_cuda=sf_cuda,
|
||||
global_scale_cuda=gs_cuda,
|
||||
block_size=block_size,
|
||||
):
|
||||
return dequantize_to_dtype(
|
||||
fp4_cuda,
|
||||
scale_cuda,
|
||||
global_scale_cuda,
|
||||
torch.bfloat16,
|
||||
block_size,
|
||||
swizzle=False,
|
||||
)
|
||||
|
||||
triton_ms, triton_min, triton_max = triton.testing.do_bench(
|
||||
_triton_bench, quantiles=quantiles
|
||||
)
|
||||
|
||||
def _reference_bench(
|
||||
fp4_cuda=fp4_cuda,
|
||||
scale_cuda=sf_cuda,
|
||||
global_scale_cuda=gs_cuda,
|
||||
block_size=block_size,
|
||||
):
|
||||
with monkeypatch.context() as m2:
|
||||
m2.setattr(
|
||||
nvfp4_emulation_utils.current_platform,
|
||||
"is_cuda_alike",
|
||||
lambda: False,
|
||||
)
|
||||
dequantize_to_dtype(
|
||||
fp4_cuda,
|
||||
scale_cuda,
|
||||
global_scale_cuda,
|
||||
torch.bfloat16,
|
||||
block_size,
|
||||
swizzle=False,
|
||||
)
|
||||
|
||||
ref_ms, ref_min, ref_max = triton.testing.do_bench(
|
||||
_reference_bench, quantiles=quantiles
|
||||
)
|
||||
|
||||
speedup = ref_ms / triton_ms if triton_ms > 0 else float("inf")
|
||||
print(f" dequantize {label} {shape}:")
|
||||
print(
|
||||
f" triton: median={triton_ms:.3f}ms, "
|
||||
f"min={triton_min:.3f}ms, max={triton_max:.3f}ms"
|
||||
)
|
||||
print(
|
||||
f" reference: median={ref_ms:.3f}ms, "
|
||||
f"min={ref_min:.3f}ms, max={ref_max:.3f}ms"
|
||||
)
|
||||
print(f" speedup: {speedup:.2f}x")
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="Triton NVFP4 kernel requires CUDA.",
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
"m, k",
|
||||
[
|
||||
(1, 16),
|
||||
(1, 4096),
|
||||
(2, 4096),
|
||||
(4, 4096),
|
||||
(8, 4096),
|
||||
(16, 4096),
|
||||
(24, 4096),
|
||||
(32, 4096),
|
||||
(1, 8192),
|
||||
(2, 8192),
|
||||
(4, 8192),
|
||||
(8, 8192),
|
||||
(16, 8192),
|
||||
(24, 8192),
|
||||
(32, 8192),
|
||||
(1, 32),
|
||||
(2, 48),
|
||||
(7, 64),
|
||||
(16, 128),
|
||||
(33, 160),
|
||||
(128, 256),
|
||||
(256, 512),
|
||||
(1024, 1024),
|
||||
(5120, 2048),
|
||||
(2048, 4096),
|
||||
(4096, 7168),
|
||||
(8192, 8192),
|
||||
(128, 16384),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("global_scale_value", [0.5, 1.0, 0.001])
|
||||
def test_triton_nvfp4_quant_dequant(
|
||||
monkeypatch, m: int, k: int, global_scale_value: float
|
||||
) -> None:
|
||||
"""Test the Triton quant-dequant kernel against the CPU reference."""
|
||||
block_size = 16
|
||||
x = torch.randn(m, k, dtype=torch.bfloat16, device="cuda")
|
||||
global_scale = torch.tensor(global_scale_value, dtype=torch.float32, device="cuda")
|
||||
|
||||
# Triton path
|
||||
triton_result = ref_nvfp4_quant_dequant(x, global_scale, block_size)
|
||||
|
||||
# CPU reference path
|
||||
with monkeypatch.context() as mp:
|
||||
mp.setattr(
|
||||
nvfp4_emulation_utils.current_platform,
|
||||
"is_cuda_alike",
|
||||
lambda: False,
|
||||
)
|
||||
reference = ref_nvfp4_quant_dequant(x.cpu(), global_scale.cpu(), block_size)
|
||||
|
||||
torch.testing.assert_close(triton_result.cpu(), reference, atol=0, rtol=0)
|
||||
|
||||
# Benchmark (both paths on CUDA tensors for fair comparison)
|
||||
quantiles = [0.5, 0.001, 0.999]
|
||||
|
||||
def _triton_bench(
|
||||
input_tensor=x, input_global_scale=global_scale, input_block_size=block_size
|
||||
):
|
||||
return ref_nvfp4_quant_dequant(
|
||||
input_tensor, input_global_scale, input_block_size
|
||||
)
|
||||
|
||||
triton_ms, triton_min, triton_max = triton.testing.do_bench(
|
||||
_triton_bench, quantiles=quantiles
|
||||
)
|
||||
|
||||
def _reference_bench(
|
||||
input_tensor=x, input_global_scale=global_scale, input_block_size=block_size
|
||||
):
|
||||
with monkeypatch.context() as mp2:
|
||||
mp2.setattr(
|
||||
nvfp4_emulation_utils.current_platform,
|
||||
"is_cuda_alike",
|
||||
lambda: False,
|
||||
)
|
||||
ref_nvfp4_quant_dequant(input_tensor, input_global_scale, input_block_size)
|
||||
|
||||
ref_ms, ref_min, ref_max = triton.testing.do_bench(
|
||||
_reference_bench, quantiles=quantiles
|
||||
)
|
||||
|
||||
speedup = ref_ms / triton_ms if triton_ms > 0 else float("inf")
|
||||
print(f" quant_dequant [{m}x{k}] gs={global_scale_value}:")
|
||||
print(
|
||||
f" triton: median={triton_ms:.3f}ms, "
|
||||
f"min={triton_min:.3f}ms, max={triton_max:.3f}ms"
|
||||
)
|
||||
print(
|
||||
f" reference: median={ref_ms:.3f}ms, "
|
||||
f"min={ref_min:.3f}ms, max={ref_max:.3f}ms"
|
||||
)
|
||||
print(f" speedup: {speedup:.2f}x")
|
||||
|
||||
|
||||
def _load_nvfp4_moe_weights(
|
||||
model_files: dict[str, str],
|
||||
model_id: str,
|
||||
tensor_parallel_size: int,
|
||||
max_experts: int | None = None,
|
||||
):
|
||||
"""Load and stack NVFP4 MoE weights from checkpoint shards.
|
||||
|
||||
Returns (w1, w1_scale, w1_gscale, w2, w2_scale, w2_gscale,
|
||||
a1_gscale, a2_gscale, num_experts, hidden_dim,
|
||||
intermediate_size).
|
||||
|
||||
When max_experts is set, only the first max_experts experts are loaded.
|
||||
|
||||
When tensor_parallel_size > 1, the N dimension of w1 and the K
|
||||
dimension of w2 are narrowed to the first TP shard (simulating
|
||||
column-parallel on w1 / row-parallel on w2).
|
||||
"""
|
||||
cfg = MOE_MODEL_CONFIGS[model_id]
|
||||
shards = cast(list[str], cfg["shards"])
|
||||
checkpoint_path = model_files[model_id]
|
||||
|
||||
expert_prefix = cfg["expert_prefix"]
|
||||
idx_pos = cast(int, cfg["expert_idx_pos"])
|
||||
|
||||
# Collect all tensors across shards into a flat dict — an expert's
|
||||
# tensors may be split across multiple shard files.
|
||||
all_tensors: dict[str, torch.Tensor] = {}
|
||||
for shard_name in shards:
|
||||
shard_path = f"{checkpoint_path}/{shard_name}"
|
||||
with safe_open(shard_path, framework="pt", device="cpu") as f:
|
||||
for key in f.keys(): # noqa: SIM118
|
||||
if key.startswith(expert_prefix):
|
||||
all_tensors[key] = f.get_tensor(key)
|
||||
|
||||
expert_indices = sorted(
|
||||
{
|
||||
int(key.split(".")[idx_pos])
|
||||
for key in all_tensors
|
||||
if key.endswith(".gate_proj.weight")
|
||||
}
|
||||
)
|
||||
if max_experts is not None:
|
||||
expert_indices = expert_indices[:max_experts]
|
||||
num_experts = len(expert_indices)
|
||||
|
||||
gate_weights, up_weights, down_weights = [], [], []
|
||||
gate_scales, up_scales, down_scales = [], [], []
|
||||
gate_gscales, up_gscales, down_gscales = [], [], []
|
||||
a1_scales, a2_scales = [], []
|
||||
|
||||
for idx in expert_indices:
|
||||
prefix = f"{expert_prefix}{idx}"
|
||||
gate_weights.append(all_tensors[f"{prefix}.gate_proj.weight"])
|
||||
gate_scales.append(all_tensors[f"{prefix}.gate_proj.weight_scale"])
|
||||
gate_gscales.append(all_tensors[f"{prefix}.gate_proj.weight_scale_2"])
|
||||
up_weights.append(all_tensors[f"{prefix}.up_proj.weight"])
|
||||
up_scales.append(all_tensors[f"{prefix}.up_proj.weight_scale"])
|
||||
up_gscales.append(all_tensors[f"{prefix}.up_proj.weight_scale_2"])
|
||||
down_weights.append(all_tensors[f"{prefix}.down_proj.weight"])
|
||||
down_scales.append(all_tensors[f"{prefix}.down_proj.weight_scale"])
|
||||
down_gscales.append(all_tensors[f"{prefix}.down_proj.weight_scale_2"])
|
||||
a1_scales.append(all_tensors[f"{prefix}.gate_proj.input_scale"])
|
||||
a2_scales.append(all_tensors[f"{prefix}.down_proj.input_scale"])
|
||||
|
||||
# Stack into MoE format.
|
||||
# w1 = [E, 2*intermediate, hidden//2] (gate + up concatenated)
|
||||
w1 = torch.stack(
|
||||
[torch.cat([g, u], dim=0) for g, u in zip(gate_weights, up_weights)]
|
||||
).cuda()
|
||||
w1_scale = torch.stack(
|
||||
[torch.cat([g, u], dim=0) for g, u in zip(gate_scales, up_scales)]
|
||||
).cuda()
|
||||
w1_gscale = torch.stack(gate_gscales).cuda()
|
||||
|
||||
# w2 = [E, hidden, intermediate//2]
|
||||
w2 = torch.stack(down_weights).cuda()
|
||||
w2_scale = torch.stack(down_scales).cuda()
|
||||
w2_gscale = torch.stack(down_gscales).cuda()
|
||||
|
||||
a13_scale_raw = torch.stack(a1_scales).cuda()
|
||||
a2_scale_raw = torch.stack(a2_scales).cuda()
|
||||
|
||||
# Apply EMULATION transforms (matches oracle/nvfp4.py).
|
||||
nvfp4_emulation_utils.kE2M1ToFloat_handle.val = (
|
||||
nvfp4_emulation_utils.kE2M1ToFloat_handle.val.cuda()
|
||||
)
|
||||
a1_gscale = 1.0 / a13_scale_raw.max().to(torch.float32)
|
||||
a2_gscale = 1.0 / a2_scale_raw.max().to(torch.float32)
|
||||
|
||||
# ── Simulate TP sharding ──
|
||||
# w1 (gate_up): column-parallel → shard the N dimension (dim 1).
|
||||
# w2 (down): row-parallel → shard the K dimension (dim 2,
|
||||
# which is the packed K//2 dim).
|
||||
# Scales follow the same sharding on the corresponding dimension.
|
||||
tp = tensor_parallel_size
|
||||
if tp > 1:
|
||||
n1 = w1.size(1) // tp
|
||||
w1 = w1[:, :n1, :].contiguous()
|
||||
w1_scale = w1_scale[:, :n1, :].contiguous()
|
||||
|
||||
k2_packed = w2.size(2) // tp
|
||||
k2_scale = w2_scale.size(2) // tp
|
||||
w2 = w2[:, :, :k2_packed].contiguous()
|
||||
w2_scale = w2_scale[:, :, :k2_scale].contiguous()
|
||||
|
||||
hidden_dim = w1.size(2) * 2
|
||||
intermediate_size = w1.size(1) // 2
|
||||
|
||||
return (
|
||||
w1,
|
||||
w1_scale,
|
||||
w1_gscale,
|
||||
w2,
|
||||
w2_scale,
|
||||
w2_gscale,
|
||||
a1_gscale,
|
||||
a2_gscale,
|
||||
num_experts,
|
||||
hidden_dim,
|
||||
intermediate_size,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="Triton NVFP4 kernel requires CUDA.",
|
||||
)
|
||||
@pytest.mark.parametrize("num_tokens", [1, 2, 4, 1024])
|
||||
@pytest.mark.parametrize("top_k", [4])
|
||||
@pytest.mark.parametrize("model_id", list(MOE_MODEL_CONFIGS.keys()))
|
||||
@pytest.mark.parametrize(
|
||||
"tensor_parallel_size",
|
||||
[pytest.param(val, id=f"tensor_parallel_size:{val}") for val in [1, 2, 4, 8]],
|
||||
)
|
||||
def test_nvfp4_moe_correctness(
|
||||
loaded_model_files,
|
||||
num_tokens: int,
|
||||
top_k: int,
|
||||
model_id: str,
|
||||
tensor_parallel_size: int,
|
||||
) -> None:
|
||||
"""Compare Nvfp4QuantizationEmulationTritonExperts (fused weight dequant + compute)
|
||||
against the unfused reference Nvfp4QuantizationEmulationTritonExpertsReference.
|
||||
|
||||
Both must produce bit-identical results.
|
||||
"""
|
||||
num_test_experts = max(8, top_k)
|
||||
(
|
||||
w1,
|
||||
w1_scale,
|
||||
w1_gscale,
|
||||
w2,
|
||||
w2_scale,
|
||||
w2_gscale,
|
||||
a1_gscale,
|
||||
a2_gscale,
|
||||
num_experts,
|
||||
hidden_dim,
|
||||
intermediate_size,
|
||||
) = _load_nvfp4_moe_weights(
|
||||
loaded_model_files,
|
||||
model_id,
|
||||
tensor_parallel_size,
|
||||
max_experts=num_test_experts,
|
||||
)
|
||||
|
||||
moe_config = FusedMoEConfig(
|
||||
num_experts=num_experts,
|
||||
experts_per_token=top_k,
|
||||
hidden_dim=hidden_dim,
|
||||
intermediate_size=intermediate_size,
|
||||
num_local_experts=num_experts,
|
||||
num_logical_experts=num_experts,
|
||||
moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
|
||||
activation=MoEActivation.SILU,
|
||||
in_dtype=torch.bfloat16,
|
||||
device="cuda",
|
||||
routing_method=RoutingMethodType.TopK,
|
||||
max_num_tokens=512,
|
||||
)
|
||||
|
||||
def _make_quant_config():
|
||||
return nvfp4_moe_quant_config(
|
||||
g1_alphas=w1_gscale.clone(),
|
||||
g2_alphas=w2_gscale.clone(),
|
||||
a1_gscale=a1_gscale.clone(),
|
||||
a2_gscale=a2_gscale.clone(),
|
||||
w1_scale=w1_scale.clone(),
|
||||
w2_scale=w2_scale.clone(),
|
||||
)
|
||||
|
||||
ref_experts = Nvfp4QuantizationEmulationTritonExpertsReference(
|
||||
moe_config=moe_config,
|
||||
quant_config=_make_quant_config(),
|
||||
)
|
||||
fused_experts = Nvfp4QuantizationEmulationTritonExperts(
|
||||
moe_config=moe_config,
|
||||
quant_config=_make_quant_config(),
|
||||
)
|
||||
|
||||
torch.manual_seed(42)
|
||||
hidden_states = torch.randn(
|
||||
num_tokens, hidden_dim, dtype=torch.bfloat16, device="cuda"
|
||||
)
|
||||
|
||||
topk_weights = torch.randn(
|
||||
num_tokens, top_k, dtype=torch.float32, device="cuda"
|
||||
).softmax(dim=-1)
|
||||
topk_ids = torch.stack(
|
||||
[torch.randperm(num_experts, device="cuda")[:top_k] for _ in range(num_tokens)]
|
||||
).to(torch.int32)
|
||||
|
||||
N = w1.size(1) # 2 * intermediate
|
||||
K = hidden_dim
|
||||
|
||||
ws13_size = num_tokens * top_k * max(intermediate_size, K)
|
||||
ws2_size = num_tokens * top_k * max(N, K)
|
||||
|
||||
workspace13_ref = torch.zeros(ws13_size, dtype=torch.bfloat16, device="cuda")
|
||||
workspace2_ref = torch.zeros(ws2_size, dtype=torch.bfloat16, device="cuda")
|
||||
output_ref = torch.zeros(num_tokens, K, dtype=torch.bfloat16, device="cuda")
|
||||
|
||||
workspace13_fused = torch.zeros_like(workspace13_ref)
|
||||
workspace2_fused = torch.zeros_like(workspace2_ref)
|
||||
output_fused = torch.zeros_like(output_ref)
|
||||
|
||||
apply_kwargs = dict(
|
||||
hidden_states=hidden_states,
|
||||
w1=w1,
|
||||
w2=w2,
|
||||
topk_weights=topk_weights,
|
||||
topk_ids=topk_ids,
|
||||
activation=MoEActivation.SILU,
|
||||
global_num_experts=num_experts,
|
||||
expert_map=None,
|
||||
a1q_scale=None,
|
||||
a2_scale=None,
|
||||
expert_tokens_meta=None,
|
||||
apply_router_weight_on_input=False,
|
||||
)
|
||||
|
||||
# Unfused reference.
|
||||
ref_experts.apply(
|
||||
output=output_ref,
|
||||
workspace13=workspace13_ref,
|
||||
workspace2=workspace2_ref,
|
||||
**apply_kwargs,
|
||||
)
|
||||
|
||||
# Fused implementation.
|
||||
fused_experts.apply(
|
||||
output=output_fused,
|
||||
workspace13=workspace13_fused,
|
||||
workspace2=workspace2_fused,
|
||||
**apply_kwargs,
|
||||
)
|
||||
|
||||
# Not strict equality on H100, MI325, MI300 (< 0.1% elements).
|
||||
# The fused on-the-fly dequant path can lower to a slightly
|
||||
# different Triton/MMA tiling than the pre-dequantized
|
||||
# reference; experiments with reference-like tiling/masking
|
||||
# reduced some diffs were not kept because they regress
|
||||
# the fused kernel speed.
|
||||
# Strict equality validated on MI355.
|
||||
torch.testing.assert_close(
|
||||
output_fused,
|
||||
output_ref,
|
||||
atol=0.0 if on_gfx950() else 0.02,
|
||||
rtol=0,
|
||||
)
|
||||
@@ -0,0 +1,308 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.scalar_type import scalar_types
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not current_platform.has_device_capability(100):
|
||||
pytest.skip(
|
||||
reason="Nvfp4 Requires compute capability of 10 or above.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
SHAPES = [(128, 64), (128, 128), (256, 64), (256, 128)]
|
||||
PAD_SHAPES = [
|
||||
(90, 64),
|
||||
(150, 64),
|
||||
(128, 48),
|
||||
(128, 80),
|
||||
(150, 80),
|
||||
(90, 48),
|
||||
(90, 128),
|
||||
(150, 128),
|
||||
(150, 48),
|
||||
(90, 80),
|
||||
(128, 512),
|
||||
(128, 1024),
|
||||
(128, 2048),
|
||||
(64, 7168),
|
||||
(64, 7152),
|
||||
(32, 14336),
|
||||
]
|
||||
PADDED_OUTPUT_SHAPES = [(128, 48), (128, 80), (150, 48), (150, 80), (64, 7152)]
|
||||
SEEDS = [42]
|
||||
CUDA_DEVICES = ["cuda:0"]
|
||||
|
||||
FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
|
||||
FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
|
||||
|
||||
# E2M1 to float
|
||||
# 0111 -> 6
|
||||
# 0110 -> 4
|
||||
# 0101 -> 3
|
||||
# 0100 -> 2
|
||||
# 0011 -> 1.5
|
||||
# 0010 -> 1
|
||||
# 0001 -> 0.5
|
||||
# 0000 -> 0
|
||||
E2M1_TO_FLOAT32 = [
|
||||
0.0,
|
||||
0.5,
|
||||
1.0,
|
||||
1.5,
|
||||
2.0,
|
||||
3.0,
|
||||
4.0,
|
||||
6.0,
|
||||
0.0,
|
||||
-0.5,
|
||||
-1.0,
|
||||
-1.5,
|
||||
-2.0,
|
||||
-3.0,
|
||||
-4.0,
|
||||
-6.0,
|
||||
]
|
||||
BLOCK_SIZE = 16
|
||||
|
||||
|
||||
def cast_from_fp4(x, m, n):
|
||||
# The fp4 values are packed in uint8 as [v_1st | v_2nd]
|
||||
v_2nd = x & 0xF
|
||||
v_1st = (x >> 4) & 0xF
|
||||
c = torch.stack((v_2nd, v_1st), dim=-1)
|
||||
out = torch.tensor([E2M1_TO_FLOAT32[x] for x in c.flatten()])
|
||||
out = out.reshape(m, n).to(torch.float32)
|
||||
return out
|
||||
|
||||
|
||||
def cast_to_fp4(x):
|
||||
sign = torch.sign(x)
|
||||
x = torch.abs(x)
|
||||
x[(x >= 0.0) & (x <= 0.25)] = 0.0
|
||||
x[(x > 0.25) & (x < 0.75)] = 0.5
|
||||
x[(x >= 0.75) & (x <= 1.25)] = 1.0
|
||||
x[(x > 1.25) & (x < 1.75)] = 1.5
|
||||
x[(x >= 1.75) & (x <= 2.5)] = 2.0
|
||||
x[(x > 2.5) & (x < 3.5)] = 3.0
|
||||
x[(x >= 3.5) & (x <= 5.0)] = 4.0
|
||||
x[x > 5.0] = 6.0
|
||||
return x * sign
|
||||
|
||||
|
||||
def get_reciprocal(x):
|
||||
if isinstance(x, torch.Tensor):
|
||||
return torch.where(x == 0, torch.tensor(0.0, dtype=x.dtype), 1.0 / x)
|
||||
elif isinstance(x, (float, int)):
|
||||
return 0.0 if x == 0 else 1.0 / x
|
||||
else:
|
||||
raise TypeError("Input must be a float, int, or a torch.Tensor.")
|
||||
|
||||
|
||||
def ref_nvfp4_quant(x, global_scale):
|
||||
assert global_scale.dtype == torch.float32
|
||||
assert x.ndim == 2
|
||||
m, n = x.shape
|
||||
x = torch.reshape(x, (m, n // BLOCK_SIZE, BLOCK_SIZE))
|
||||
vec_max = torch.max(torch.abs(x), dim=-1, keepdim=True)[0].to(torch.float32)
|
||||
scale = global_scale * (vec_max * get_reciprocal(FLOAT4_E2M1_MAX))
|
||||
scale = scale.to(torch.float8_e4m3fn).to(torch.float32)
|
||||
output_scale = get_reciprocal(scale * get_reciprocal(global_scale))
|
||||
|
||||
scaled_x = x.to(torch.float32) * output_scale
|
||||
clipped_x = torch.clamp(scaled_x, -6.0, 6.0).reshape(m, n)
|
||||
return cast_to_fp4(clipped_x), scale.squeeze(-1)
|
||||
|
||||
|
||||
def recover_swizzled_scales(scale, m, n):
|
||||
round_up = lambda x, y: (x + y - 1) // y * y
|
||||
rounded_m = round_up(m, 128)
|
||||
scale_n = n // BLOCK_SIZE
|
||||
rounded_n = round_up(scale_n, 4)
|
||||
# Recover the swizzled scaling factor to linear layout
|
||||
tmp = torch.reshape(scale, (1, rounded_m // 128, rounded_n // 4, 32, 4, 4))
|
||||
tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
|
||||
result = torch.reshape(tmp, (rounded_m, rounded_n)).to(torch.float32)
|
||||
return result[:m, :scale_n]
|
||||
|
||||
|
||||
def round_up(x: int, y: int) -> int:
|
||||
return (x + y - 1) // y * y
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("shape", SHAPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
def test_quantize_to_fp4(
|
||||
dtype: torch.dtype,
|
||||
shape: tuple[int, int],
|
||||
seed: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
torch.set_default_device(device)
|
||||
|
||||
m, n = shape
|
||||
|
||||
x = torch.randn((m, n), dtype=dtype)
|
||||
tensor_amax = torch.abs(x).max().to(torch.float32)
|
||||
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
|
||||
out_ref, scale_ref = ref_nvfp4_quant(x, global_scale)
|
||||
|
||||
out, out_scale = ops.scaled_fp4_quant(x, global_scale)
|
||||
scale_ans = recover_swizzled_scales(out_scale, m, n)
|
||||
out_ans = cast_from_fp4(out, m, n)
|
||||
|
||||
torch.testing.assert_close(out_ans, out_ref)
|
||||
torch.testing.assert_close(scale_ans, scale_ref)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shape",
|
||||
[(32, 4096), (128, 4096), (1, 64), (127, 1024), (256, 16384)],
|
||||
)
|
||||
@pytest.mark.parametrize("is_sf_swizzled_layout", [True, False])
|
||||
@torch.inference_mode()
|
||||
def test_python_util_matches_cpp_allocation(
|
||||
shape: tuple[int, int],
|
||||
is_sf_swizzled_layout: bool,
|
||||
) -> None:
|
||||
"""
|
||||
Verify that the Python utility (create_fp4_output_tensors) allocates
|
||||
tensors with the same shapes and dtypes as the C++ functional variant
|
||||
(scaled_fp4_quant_func).
|
||||
"""
|
||||
from vllm._custom_ops import create_fp4_output_tensors
|
||||
|
||||
torch.set_default_device("cuda:0")
|
||||
m, n = shape
|
||||
input_tensor = torch.randn((m, n), dtype=torch.bfloat16)
|
||||
input_scale = torch.tensor([1.0], dtype=torch.float32, device="cuda:0")
|
||||
|
||||
# C++ functional variant allocates internally
|
||||
cpp_out, cpp_scale = torch.ops._C.scaled_fp4_quant(
|
||||
input_tensor, input_scale, is_sf_swizzled_layout
|
||||
)
|
||||
|
||||
# Python utility
|
||||
py_out, py_scale = create_fp4_output_tensors(
|
||||
m, n, torch.device("cuda:0"), is_sf_swizzled_layout
|
||||
)
|
||||
|
||||
assert py_out.shape == cpp_out.shape, (
|
||||
f"Output shape mismatch: Python {py_out.shape} vs C++ {cpp_out.shape}"
|
||||
)
|
||||
assert py_out.dtype == cpp_out.dtype, (
|
||||
f"Output dtype mismatch: Python {py_out.dtype} vs C++ {cpp_out.dtype}"
|
||||
)
|
||||
assert py_scale.shape == cpp_scale.shape, (
|
||||
f"Scale shape mismatch: Python {py_scale.shape} vs C++ {cpp_scale.shape}"
|
||||
)
|
||||
assert py_scale.dtype == cpp_scale.dtype, (
|
||||
f"Scale dtype mismatch: Python {py_scale.dtype} vs C++ {cpp_scale.dtype}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", PADDED_OUTPUT_SHAPES)
|
||||
@pytest.mark.parametrize("is_sf_swizzled_layout", [True, False])
|
||||
@torch.inference_mode()
|
||||
def test_quantize_to_fp4_with_padded_output(
|
||||
shape: tuple[int, int],
|
||||
is_sf_swizzled_layout: bool,
|
||||
) -> None:
|
||||
from vllm._custom_ops import create_fp4_output_tensors
|
||||
|
||||
dtype = torch.float16
|
||||
set_random_seed(42)
|
||||
torch.set_default_device("cuda:0")
|
||||
|
||||
m, n = shape
|
||||
padded_n = round_up(n, 32)
|
||||
assert padded_n > n
|
||||
|
||||
x = torch.randn((m, n), dtype=dtype)
|
||||
tensor_amax = torch.abs(x).max().to(torch.float32)
|
||||
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
|
||||
out_ref, scale_ref = ref_nvfp4_quant(x, global_scale)
|
||||
|
||||
out, out_scale = ops.scaled_fp4_quant(
|
||||
x,
|
||||
global_scale,
|
||||
is_sf_swizzled_layout=is_sf_swizzled_layout,
|
||||
padded_n=padded_n,
|
||||
)
|
||||
py_out, py_scale = create_fp4_output_tensors(
|
||||
m,
|
||||
n,
|
||||
torch.device("cuda:0"),
|
||||
is_sf_swizzled_layout,
|
||||
padded_n=padded_n,
|
||||
)
|
||||
|
||||
assert out.shape == (m, padded_n // 2)
|
||||
assert out.shape == py_out.shape
|
||||
assert out_scale.shape == py_scale.view(torch.float8_e4m3fn).shape
|
||||
|
||||
out_ans = cast_from_fp4(out[:, : n // 2], m, n)
|
||||
torch.testing.assert_close(out_ans, out_ref)
|
||||
assert torch.count_nonzero(out[:, n // 2 :]) == 0
|
||||
|
||||
if is_sf_swizzled_layout:
|
||||
scale_ans = recover_swizzled_scales(out_scale, m, padded_n)
|
||||
torch.testing.assert_close(scale_ans[:, : n // BLOCK_SIZE], scale_ref)
|
||||
assert torch.count_nonzero(scale_ans[:, n // BLOCK_SIZE :]) == 0
|
||||
else:
|
||||
scale_ans = out_scale.to(torch.float32)
|
||||
torch.testing.assert_close(scale_ans[:, : n // BLOCK_SIZE], scale_ref)
|
||||
assert torch.count_nonzero(scale_ans[:, n // BLOCK_SIZE :]) == 0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pad_shape", PAD_SHAPES)
|
||||
@torch.inference_mode()
|
||||
def test_quantize_to_fp4_padded(pad_shape: tuple[int, int]) -> None:
|
||||
dtype = torch.float16
|
||||
set_random_seed(42)
|
||||
torch.set_default_device("cuda:0")
|
||||
|
||||
m, n = pad_shape
|
||||
|
||||
x = torch.randn((m, n), dtype=dtype)
|
||||
|
||||
tensor_amax = torch.abs(x).max().to(torch.float32)
|
||||
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
|
||||
out_ref, scale_ref = ref_nvfp4_quant(x, global_scale)
|
||||
|
||||
out, out_scale = ops.scaled_fp4_quant(x, global_scale)
|
||||
scale_ans = recover_swizzled_scales(out_scale, m, n)
|
||||
out_ans = cast_from_fp4(out, m, n)
|
||||
torch.testing.assert_close(out_ans, out_ref)
|
||||
torch.testing.assert_close(scale_ans, scale_ref)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pad_shape", PAD_SHAPES)
|
||||
@torch.inference_mode()
|
||||
def test_quantize_to_fp4_padded_no_sf_swizzled(pad_shape: tuple[int, int]) -> None:
|
||||
dtype = torch.float16
|
||||
set_random_seed(42)
|
||||
torch.set_default_device("cuda:0")
|
||||
|
||||
m, n = pad_shape
|
||||
|
||||
x = torch.randn((m, n), dtype=dtype)
|
||||
|
||||
tensor_amax = torch.abs(x).max().to(torch.float32)
|
||||
global_scale = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / tensor_amax
|
||||
out_ref, scale_ref = ref_nvfp4_quant(x, global_scale)
|
||||
|
||||
out, out_scale = ops.scaled_fp4_quant(x, global_scale, is_sf_swizzled_layout=False)
|
||||
scale_ans = out_scale.to(torch.float32)
|
||||
out_ans = cast_from_fp4(out, m, n)
|
||||
torch.testing.assert_close(out_ans, out_ref)
|
||||
torch.testing.assert_close(scale_ans, scale_ref)
|
||||
@@ -0,0 +1,269 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
#
|
||||
# Copyright (C) 2025 Roberto L. Castro (Roberto.LopezCastro@ist.ac.at).
|
||||
# All Rights Reserved.
|
||||
#
|
||||
# Licensed 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.
|
||||
#
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from compressed_tensors.transform.utils.hadamard import deterministic_hadamard_matrix
|
||||
|
||||
from vllm import _custom_ops as ops # use existing nvfp4 gemm in vllm
|
||||
from vllm._custom_ops import fusedQuantizeNv
|
||||
from vllm.model_executor.layers.quantization.qutlass_utils import to_blocked
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA required for these tests.", allow_module_level=True)
|
||||
|
||||
if not (
|
||||
current_platform.has_device_capability(100)
|
||||
or current_platform.has_device_capability(120)
|
||||
):
|
||||
pytest.skip(
|
||||
reason="Tests require compute capability 10.0 (100) or 12.0 (120).",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
|
||||
# ----- Helpers -----
|
||||
def get_hadamard_matrix(group_size: int, dtype: torch.dtype, device: torch.device):
|
||||
return (
|
||||
deterministic_hadamard_matrix(group_size, dtype=dtype, device=device)
|
||||
* group_size**-0.5
|
||||
)
|
||||
|
||||
|
||||
def _rtne_fp4(x: torch.Tensor):
|
||||
device = x.device
|
||||
grid = torch.tensor(
|
||||
[
|
||||
-6.0,
|
||||
-4.0,
|
||||
-3.0,
|
||||
-2.0,
|
||||
-1.5,
|
||||
-1.0,
|
||||
-0.5,
|
||||
-0.0,
|
||||
0.0,
|
||||
0.5,
|
||||
1.0,
|
||||
1.5,
|
||||
2.0,
|
||||
3.0,
|
||||
4.0,
|
||||
6.0,
|
||||
],
|
||||
dtype=x.dtype,
|
||||
device=x.device,
|
||||
)
|
||||
grid_int = torch.tensor(
|
||||
[-1, -2, -3, -4, -5, -6, -7, -8, 0, 1, 2, 3, 4, 5, 6, 7],
|
||||
dtype=torch.uint8,
|
||||
device=device,
|
||||
)
|
||||
inds = torch.bucketize(x, grid)
|
||||
lo, hi = (inds - 1).clamp(min=0, max=15), inds.clamp(min=0, max=15)
|
||||
g_lo, g_hi = grid[lo], grid[hi]
|
||||
pick_hi = (g_hi - x < x - g_lo) | (g_hi - x == x - g_lo) & (grid_int[hi] % 2 == 0)
|
||||
y = torch.where(pick_hi, g_hi, g_lo)
|
||||
y_int = torch.where(pick_hi, grid_int[hi], grid_int[lo])
|
||||
y_int_packed = (y_int[..., 1::2] & 0xF) << 4 | y_int[..., ::2] & 0xF
|
||||
return y, y_int_packed
|
||||
|
||||
|
||||
def _dq_fp4(x_e2m1: torch.Tensor, x_e4m3: torch.Tensor, alpha: float):
|
||||
device = x_e2m1.device
|
||||
|
||||
x_e2m1_i32 = x_e2m1.view(dtype=torch.uint8).to(dtype=torch.int32)
|
||||
x_e2m1_unpacked = torch.stack(
|
||||
[x_e2m1_i32 & 0xF, (x_e2m1_i32 >> 4) & 0xF], dim=-1
|
||||
).flatten(start_dim=-2)
|
||||
|
||||
grid_dq = torch.tensor(
|
||||
[
|
||||
0.0,
|
||||
0.5,
|
||||
1.0,
|
||||
1.5,
|
||||
2.0,
|
||||
3.0,
|
||||
4.0,
|
||||
6.0,
|
||||
-0.0,
|
||||
-0.5,
|
||||
-1.0,
|
||||
-1.5,
|
||||
-2.0,
|
||||
-3.0,
|
||||
-4.0,
|
||||
-6.0,
|
||||
],
|
||||
dtype=torch.float64,
|
||||
device=device,
|
||||
)
|
||||
x_fp4_dq = grid_dq[x_e2m1_unpacked]
|
||||
|
||||
scales_dq = x_e4m3.to(torch.float64)
|
||||
x_dq = (x_fp4_dq.unflatten(dim=-1, sizes=(-1, 16)) * scales_dq[..., None]).flatten(
|
||||
start_dim=-2
|
||||
) / alpha # * (4. / 3.)
|
||||
return x_dq, x_fp4_dq, scales_dq
|
||||
|
||||
|
||||
def _unpack_mask(clip_mask: torch.Tensor) -> torch.Tensor:
|
||||
clip_mask_unpacked_dq = torch.zeros(
|
||||
*clip_mask.shape[:-1],
|
||||
clip_mask.size(-1) * 8,
|
||||
dtype=torch.bool,
|
||||
device=clip_mask.device,
|
||||
)
|
||||
for i in range(8):
|
||||
clip_mask_unpacked_dq[..., i::8] = (clip_mask >> i) & 1
|
||||
return clip_mask_unpacked_dq
|
||||
|
||||
|
||||
def _forward_quantize_ref(x: torch.Tensor, h: torch.Tensor, rot_size: int):
|
||||
device = x.device
|
||||
|
||||
xh_ref64 = (
|
||||
x.unflatten(dim=-1, sizes=(-1, rot_size)).to(dtype=torch.float64)
|
||||
@ h.reshape(rot_size, rot_size).to(dtype=torch.float64)
|
||||
).flatten(start_dim=-2)
|
||||
|
||||
abs_max = xh_ref64.unflatten(dim=-1, sizes=(-1, 16)).abs().amax(dim=-1)
|
||||
scales_ref64_ = abs_max + 1e-8
|
||||
|
||||
xh_e4m3_ref = scales_ref64_.to(dtype=torch.float8_e4m3fn)
|
||||
scales_ref64 = xh_e4m3_ref.to(dtype=torch.float64)
|
||||
xh_scaled_ref64 = (
|
||||
xh_ref64.unflatten(dim=-1, sizes=(-1, 16)) / scales_ref64[..., None]
|
||||
).flatten(start_dim=-2)
|
||||
|
||||
xh_scaled_ref64 *= 6.0
|
||||
|
||||
clip_mask_unpacked_ref = xh_scaled_ref64.abs() < 6.0
|
||||
clip_mask_ref = torch.zeros(
|
||||
*x.shape[:-1], x.size(-1) // 8, dtype=torch.uint8, device=device
|
||||
)
|
||||
for i in range(8):
|
||||
clip_mask_ref |= clip_mask_unpacked_ref[..., i::8].to(dtype=torch.uint8) << i
|
||||
|
||||
xh_fp4_ref, xh_e2m1_ref = _rtne_fp4(xh_scaled_ref64)
|
||||
xh_dq, xh_fp4_dq, scales_dq = _dq_fp4(xh_e2m1_ref, xh_e4m3_ref, 6.0)
|
||||
clip_mask_unpacked_dq = _unpack_mask(clip_mask_ref)
|
||||
|
||||
assert xh_fp4_dq.equal(xh_fp4_ref)
|
||||
assert scales_dq.equal(scales_ref64)
|
||||
assert clip_mask_unpacked_dq.equal(clip_mask_unpacked_ref)
|
||||
|
||||
return (
|
||||
xh_dq,
|
||||
clip_mask_unpacked_ref,
|
||||
(xh_e2m1_ref, xh_e4m3_ref, clip_mask_ref),
|
||||
)
|
||||
|
||||
|
||||
DTYPE = torch.bfloat16
|
||||
DEVICE = torch.device("cuda:0")
|
||||
ROT_SIZES = [16, 32, 64, 128]
|
||||
GLOBAL_SCALES = [6.0]
|
||||
|
||||
LLAMA_MODELS = {
|
||||
"7B": [(4096, 3 * 4096), (4096, 4096), (4096, 2 * 10752), (10752, 4096)],
|
||||
"13B": [(5120, 3 * 5120), (5120, 5120), (5120, 2 * 13568), (13568, 5120)],
|
||||
"33B": [(6656, 3 * 6656), (6656, 6656), (6656, 2 * 17664), (17664, 6656)],
|
||||
"70B": [(8192, 3 * 8192), (8192, 8192), (8192, 2 * 21760), (21760, 8192)],
|
||||
}
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _seed_each_test():
|
||||
set_random_seed(0)
|
||||
np.random.seed(0)
|
||||
torch.random.manual_seed(0)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("rot_size", ROT_SIZES)
|
||||
@pytest.mark.parametrize("global_scale_value", GLOBAL_SCALES)
|
||||
@torch.inference_mode()
|
||||
def test_fused_quantization(rot_size: int, global_scale_value: float):
|
||||
dtype, device = DTYPE, DEVICE
|
||||
h = get_hadamard_matrix(rot_size, dtype, device)
|
||||
x = torch.randn(2, 4096, 4096, dtype=dtype, device=device) * 25.0
|
||||
global_scale = torch.tensor([global_scale_value], device=device)
|
||||
|
||||
xh_dq_ref, _, _ = _forward_quantize_ref(x, h, rot_size)
|
||||
xh_e2m1, xh_e4m3 = fusedQuantizeNv(x, h, global_scale)
|
||||
xh_e4m3 = xh_e4m3.reshape(2, 4096, 4096 // 16)
|
||||
xh_dq, *_ = _dq_fp4(xh_e2m1, xh_e4m3, alpha=global_scale_value)
|
||||
|
||||
torch.testing.assert_close(xh_dq, xh_dq_ref, rtol=0.34, atol=100)
|
||||
assert (xh_dq != xh_dq_ref).float().mean() <= 1e-1
|
||||
|
||||
m, n, k = 504, 4096 * 2, 4096
|
||||
a = torch.randn(m, k, dtype=dtype, device=device) * 25.0
|
||||
b = torch.randn(n, k, dtype=dtype, device=device) * 25.0
|
||||
|
||||
a_e2m1, a_e4m3 = fusedQuantizeNv(a, h, global_scale)
|
||||
b_e2m1, b_e4m3 = fusedQuantizeNv(b, h, global_scale)
|
||||
|
||||
a_dq, *_ = _dq_fp4(a_e2m1, a_e4m3[:m, :k], alpha=1.0)
|
||||
b_dq, *_ = _dq_fp4(b_e2m1, b_e4m3[:n, :k], alpha=1.0)
|
||||
out_ref = a_dq @ b_dq.transpose(-2, -1)
|
||||
|
||||
a_scale_block = to_blocked(a_e4m3, backend="triton").view(-1, k // 16)
|
||||
b_scale_block = to_blocked(b_e4m3, backend="triton").view(-1, k // 16)
|
||||
alpha = torch.tensor([1.0], device=device)
|
||||
out = ops.cutlass_scaled_fp4_mm(
|
||||
a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha, torch.bfloat16
|
||||
)
|
||||
assert out.equal(out_ref.to(dtype=out.dtype))
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model", list(LLAMA_MODELS.keys()))
|
||||
@pytest.mark.parametrize("layer_idx", [0, 1, 2, 3])
|
||||
@pytest.mark.parametrize("batch", [1, 16])
|
||||
@pytest.mark.parametrize("rot_size", ROT_SIZES)
|
||||
@torch.inference_mode()
|
||||
def test_llama_shapes(model: str, layer_idx: int, batch: int, rot_size: int):
|
||||
dtype, device = DTYPE, DEVICE
|
||||
m = batch
|
||||
k, n = LLAMA_MODELS[model][layer_idx]
|
||||
|
||||
h = get_hadamard_matrix(rot_size, dtype, device)
|
||||
|
||||
a = torch.randn(m, k, dtype=dtype, device=device) * 25.0
|
||||
b = torch.randn(n, k, dtype=dtype, device=device) * 25.0
|
||||
|
||||
global_scale = torch.tensor([1.0], device=device)
|
||||
|
||||
a_e2m1, a_e4m3 = fusedQuantizeNv(a, h, global_scale)
|
||||
b_e2m1, b_e4m3 = fusedQuantizeNv(b, h, global_scale)
|
||||
|
||||
a_dq, *_ = _dq_fp4(a_e2m1, a_e4m3[:m, :k], alpha=1.0)
|
||||
b_dq, *_ = _dq_fp4(b_e2m1, b_e4m3[:n, :k], alpha=1.0)
|
||||
out_ref = a_dq @ b_dq.transpose(-2, -1)
|
||||
|
||||
a_scale_block = to_blocked(a_e4m3, backend="triton").view(-1, k // 16)
|
||||
b_scale_block = to_blocked(b_e4m3, backend="triton").view(-1, k // 16)
|
||||
alpha = torch.tensor([1.0], device=device)
|
||||
out = ops.cutlass_scaled_fp4_mm(
|
||||
a_e2m1, b_e2m1, a_scale_block, b_scale_block, alpha, torch.bfloat16
|
||||
)
|
||||
assert out.equal(out_ref.to(dtype=out.dtype))
|
||||
@@ -0,0 +1,100 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
from nvfp4_utils import FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX, dequantize_nvfp4_to_dtype
|
||||
|
||||
from vllm import _custom_ops as ops
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not current_platform.has_device_capability(100):
|
||||
pytest.skip(
|
||||
reason="Nvfp4 Requires compute capability of 10 or above.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
# m, n, k
|
||||
SHAPES = [(128, 128, 64), (128, 128, 128), (256, 128, 64), (128, 256, 128)]
|
||||
PAD_SHAPES = [(150, 128, 64), (128, 128, 96)]
|
||||
SHAPES.extend(PAD_SHAPES)
|
||||
|
||||
SEEDS = [42]
|
||||
CUDA_DEVICES = ["cuda:0"]
|
||||
|
||||
|
||||
def get_ref_results(
|
||||
a_fp4,
|
||||
b_fp4,
|
||||
a_sf,
|
||||
b_sf,
|
||||
a_global_scale,
|
||||
b_global_scale,
|
||||
m,
|
||||
n,
|
||||
dtype,
|
||||
block_size,
|
||||
device,
|
||||
):
|
||||
_, m_k = a_fp4.shape
|
||||
_, n_k = b_fp4.shape
|
||||
assert m_k == n_k
|
||||
a_in_dtype = dequantize_nvfp4_to_dtype(
|
||||
a_fp4, a_sf, a_global_scale, dtype=dtype, device=device, block_size=block_size
|
||||
)
|
||||
b_in_dtype = dequantize_nvfp4_to_dtype(
|
||||
b_fp4, b_sf, b_global_scale, dtype=dtype, device=device, block_size=block_size
|
||||
)
|
||||
return torch.matmul(a_in_dtype, b_in_dtype.t())
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("shape", SHAPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("device", CUDA_DEVICES)
|
||||
@torch.inference_mode()
|
||||
def test_nvfp4_gemm(
|
||||
dtype: torch.dtype,
|
||||
shape: tuple[int, int, int],
|
||||
seed: int,
|
||||
device: str,
|
||||
) -> None:
|
||||
set_random_seed(seed)
|
||||
m, n, packed_k = shape
|
||||
k = packed_k * 2
|
||||
block_size = 16
|
||||
a_dtype = torch.randn((m, k), dtype=dtype, device=device)
|
||||
b_dtype = torch.randn((n, k), dtype=dtype, device=device)
|
||||
|
||||
a_global_scale = (
|
||||
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a_dtype.flatten(), dim=-1)
|
||||
).to(torch.float32)
|
||||
b_global_scale = (
|
||||
(FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(b_dtype.flatten(), dim=-1)
|
||||
).to(torch.float32)
|
||||
alpha = 1.0 / (a_global_scale * b_global_scale)
|
||||
# ops.scaled_fp4_quant returns swizzled scales, while weights
|
||||
# from checkpoints are in linear scales.
|
||||
a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a_dtype, a_global_scale)
|
||||
b_fp4, b_scale_interleaved = ops.scaled_fp4_quant(b_dtype, b_global_scale)
|
||||
|
||||
# get_ref_results unswizzles the scales internally.
|
||||
expected_out = get_ref_results(
|
||||
a_fp4,
|
||||
b_fp4,
|
||||
a_scale_interleaved,
|
||||
b_scale_interleaved,
|
||||
a_global_scale,
|
||||
b_global_scale,
|
||||
m,
|
||||
n,
|
||||
dtype,
|
||||
block_size,
|
||||
device,
|
||||
)
|
||||
out = ops.cutlass_scaled_fp4_mm(
|
||||
a_fp4, b_fp4, a_scale_interleaved, b_scale_interleaved, alpha, dtype
|
||||
)
|
||||
|
||||
torch.testing.assert_close(out, expected_out.to(dtype=dtype), atol=1e-1, rtol=1e-1)
|
||||
@@ -0,0 +1,430 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.quantization.utils import fp8_utils, int8_utils
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import get_fp8_min_max
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"shape", [(31, 128), (32, 128), (63, 256), (64, 256), (16, 512)]
|
||||
)
|
||||
@pytest.mark.parametrize("column_major", [False, True])
|
||||
@pytest.mark.parametrize("tma_aligned", [False, True])
|
||||
@pytest.mark.parametrize("scale_ue8m0", [False, True])
|
||||
@pytest.mark.parametrize("group_size", [64, 128])
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(), reason="Only test on CUDA/ROCm."
|
||||
)
|
||||
def test_per_token_group_quant_fp8(
|
||||
shape, column_major: bool, tma_aligned: bool, scale_ue8m0: bool, group_size: int
|
||||
):
|
||||
device = "cuda"
|
||||
|
||||
torch.manual_seed(42)
|
||||
num_tokens, hidden_dim = shape
|
||||
|
||||
x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
|
||||
|
||||
# cuda path
|
||||
out_q, scale = fp8_utils.per_token_group_quant_fp8(
|
||||
x,
|
||||
group_size,
|
||||
column_major_scales=column_major,
|
||||
tma_aligned_scales=tma_aligned,
|
||||
use_ue8m0=scale_ue8m0,
|
||||
)
|
||||
|
||||
# triton ref
|
||||
with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
|
||||
ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
|
||||
x,
|
||||
group_size,
|
||||
column_major_scales=column_major,
|
||||
use_ue8m0=scale_ue8m0,
|
||||
)
|
||||
|
||||
assert torch.allclose(out_q.float(), ref_q.float(), atol=0.15, rtol=0.15)
|
||||
assert torch.allclose(scale, ref_s, atol=0.01, rtol=0.01)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_tokens,hidden_dim,group_size",
|
||||
[
|
||||
# No padding: mn=4 (mult of 4), groups_per_row=56 (mult of 4)
|
||||
(4, 7168, 128),
|
||||
# MN padding only: mn=1, tma_aligned_mn=4
|
||||
(1, 7168, 128),
|
||||
# MN padding only: mn=3, tma_aligned_mn=4
|
||||
(3, 7168, 128),
|
||||
# K padding only: groups_per_row=5 (5%4=1)
|
||||
(4, 640, 128),
|
||||
# K padding only: groups_per_row=6 (6%4=2)
|
||||
(4, 768, 128),
|
||||
# Single packed column, no padding: k_num_packed=1, mn%4=0
|
||||
(4, 384, 128),
|
||||
# Both MN and K padding
|
||||
(1, 384, 128),
|
||||
(3, 640, 128),
|
||||
# Larger shapes with no padding
|
||||
(64, 7168, 128),
|
||||
(128, 14336, 128),
|
||||
# Larger shapes with padding
|
||||
(127, 7168, 128),
|
||||
(253, 640, 128),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("poisoned_scales", [False, True])
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="DeepGEMM not available on this platform",
|
||||
)
|
||||
def test_per_token_group_quant_fp8_packed(
|
||||
num_tokens, hidden_dim, group_size, poisoned_scales
|
||||
):
|
||||
"""Test the packed DeepGEMM quantization kernel against the Triton
|
||||
reference (row-major, UE8M0 scales)."""
|
||||
|
||||
device = "cuda"
|
||||
torch.manual_seed(42)
|
||||
|
||||
x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
|
||||
|
||||
mn = num_tokens
|
||||
groups_per_row = hidden_dim // group_size
|
||||
k_num_packed = (groups_per_row + 3) // 4
|
||||
tma_aligned_mn = ((mn + 3) // 4) * 4
|
||||
num_scale_elems = mn + (k_num_packed - 1) * tma_aligned_mn
|
||||
|
||||
if poisoned_scales:
|
||||
# Call the kernel with poisoned scale buffer to
|
||||
# ensure padded indices are correctly zeroed.
|
||||
fp8_dtype = current_platform.fp8_dtype()
|
||||
fp8_min, fp8_max = get_fp8_min_max()
|
||||
out_q = torch.empty_like(x, dtype=fp8_dtype)
|
||||
out_s_packed = torch.empty_strided(
|
||||
(mn, k_num_packed),
|
||||
(1, tma_aligned_mn),
|
||||
device=device,
|
||||
dtype=torch.int32,
|
||||
)
|
||||
torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).fill_(0x7F7F7F7F)
|
||||
torch.ops._C.per_token_group_fp8_quant_packed(
|
||||
x,
|
||||
out_q,
|
||||
out_s_packed,
|
||||
group_size,
|
||||
1e-10,
|
||||
fp8_min,
|
||||
fp8_max,
|
||||
)
|
||||
else:
|
||||
out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
|
||||
x,
|
||||
group_size=group_size,
|
||||
use_ue8m0=True,
|
||||
)
|
||||
|
||||
# Triton reference (row-major float32 scales, UE8M0)
|
||||
with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
|
||||
ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
|
||||
x,
|
||||
group_size,
|
||||
use_ue8m0=True,
|
||||
)
|
||||
|
||||
# Quantized values must match.
|
||||
assert torch.equal(out_q, ref_q), "Quantized output mismatch"
|
||||
|
||||
# Verify packed scales (valid exponents + padding zeros).
|
||||
ref_s_flat = ref_s.reshape(mn, groups_per_row)
|
||||
ref_exponents = (ref_s_flat.view(torch.int32) >> 23) & 0xFF
|
||||
|
||||
expected = torch.zeros(num_scale_elems, dtype=torch.int32, device="cpu")
|
||||
for row in range(mn):
|
||||
for g in range(groups_per_row):
|
||||
pack_col = g // 4
|
||||
pos = g % 4
|
||||
idx = pack_col * tma_aligned_mn + row
|
||||
expected[idx] |= int(ref_exponents[row, g].item()) << (pos * 8)
|
||||
|
||||
actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu()
|
||||
assert torch.equal(actual, expected), (
|
||||
f"Packed scale storage mismatch.\n"
|
||||
f"First diff at index "
|
||||
f"{(actual != expected).nonzero(as_tuple=True)[0][0].item()}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="DeepGEMM not available on this platform",
|
||||
)
|
||||
def test_per_token_group_quant_fp8_packed_all_zero():
|
||||
"""All-zero input must produce well-defined UE8M0 scale bytes via the eps
|
||||
floor in the kernel's UE8M0 path. Locks down the all-zero behavior before
|
||||
optimization.
|
||||
|
||||
The CUDA kernel computes:
|
||||
y_s = eps / fp8_max
|
||||
y_s = exp2(ceil(log2(fmax(y_s, 1e-10))))
|
||||
For all-zero input, eps/fp8_max < 1e-10, so the inner fmax clamps back to
|
||||
1e-10, giving exp2(ceil(log2(1e-10))) = exp2(-33) => UE8M0 byte 0x5E (94).
|
||||
"""
|
||||
|
||||
device = "cuda"
|
||||
num_tokens, hidden_dim, group_size = 4, 7168, 128
|
||||
x = torch.zeros((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16)
|
||||
|
||||
out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
|
||||
x,
|
||||
group_size=group_size,
|
||||
use_ue8m0=True,
|
||||
)
|
||||
|
||||
# Quantized values must be all zero.
|
||||
assert torch.equal(
|
||||
out_q.view(torch.uint8),
|
||||
torch.zeros_like(out_q, dtype=torch.uint8),
|
||||
), "All-zero input should produce all-zero FP8 output"
|
||||
|
||||
# UE8M0 byte produced by the kernel for all-zero input.
|
||||
# The kernel's inner fmax(y_s, 1e-10) clamps eps/fp8_max back to 1e-10.
|
||||
# 1e-10 as float32 has biased exponent 0x5D and a non-zero mantissa, so
|
||||
# the kernel's bit-twiddle (exp_bits + (mant_bits != 0)) rounds up to
|
||||
# 0x5E. This matches exp2(ceil(log2(1e-10))) = exp2(-33).
|
||||
expected_exp_byte = 0x5E
|
||||
|
||||
mn = num_tokens
|
||||
groups_per_row = hidden_dim // group_size
|
||||
k_num_packed = (groups_per_row + 3) // 4
|
||||
tma_aligned_mn = ((mn + 3) // 4) * 4
|
||||
num_scale_elems = mn + (k_num_packed - 1) * tma_aligned_mn
|
||||
|
||||
# All valid scale slots must contain the expected packed value.
|
||||
# Padding slots must be zero.
|
||||
actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu()
|
||||
|
||||
expected = torch.zeros(num_scale_elems, dtype=torch.int32, device="cpu")
|
||||
for row in range(mn):
|
||||
for g in range(groups_per_row):
|
||||
pack_col = g // 4
|
||||
pos = g % 4
|
||||
idx = pack_col * tma_aligned_mn + row
|
||||
expected[idx] |= expected_exp_byte << (pos * 8)
|
||||
|
||||
assert torch.equal(actual, expected), "All-zero scale bytes mismatch"
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="DeepGEMM not available on this platform",
|
||||
)
|
||||
def test_per_token_group_quant_fp8_packed_mantissa_rounds_up():
|
||||
"""Inputs whose absmax/max_8bit produces a non-power-of-2 force the
|
||||
mantissa-rounding-up branch (exp_byte += 1). Locks down this behavior
|
||||
before optimization."""
|
||||
|
||||
device = "cuda"
|
||||
num_tokens, hidden_dim, group_size = 4, 7168, 128
|
||||
|
||||
# Build a tensor whose per-group absmax = 1.5 * fp8_max * 2^k for various k.
|
||||
# fp8_max = torch.finfo(torch.float8_e4m3fn).max = 448.0.
|
||||
# Then absmax/fp8_max = 1.5 * 2^k -> non-zero mantissa, triggers ceil
|
||||
# rounding to 2^(k+1). Use k=0 for simplicity; the bf16 representation of
|
||||
# 1.5*448=672.0 is exact.
|
||||
x = torch.full(
|
||||
(num_tokens, hidden_dim),
|
||||
672.0,
|
||||
device=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
|
||||
x,
|
||||
group_size=group_size,
|
||||
use_ue8m0=True,
|
||||
)
|
||||
|
||||
with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
|
||||
ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
|
||||
x,
|
||||
group_size,
|
||||
use_ue8m0=True,
|
||||
)
|
||||
|
||||
assert torch.equal(out_q, ref_q), "Quantized output mismatch"
|
||||
|
||||
mn = num_tokens
|
||||
groups_per_row = hidden_dim // group_size
|
||||
k_num_packed = (groups_per_row + 3) // 4
|
||||
tma_aligned_mn = ((mn + 3) // 4) * 4
|
||||
num_scale_elems = mn + (k_num_packed - 1) * tma_aligned_mn
|
||||
|
||||
ref_s_flat = ref_s.reshape(mn, groups_per_row)
|
||||
ref_exponents = (ref_s_flat.view(torch.int32) >> 23) & 0xFF
|
||||
expected = torch.zeros(num_scale_elems, dtype=torch.int32, device="cpu")
|
||||
for row in range(mn):
|
||||
for g in range(groups_per_row):
|
||||
pack_col = g // 4
|
||||
pos = g % 4
|
||||
idx = pack_col * tma_aligned_mn + row
|
||||
expected[idx] |= int(ref_exponents[row, g].item()) << (pos * 8)
|
||||
|
||||
actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu()
|
||||
assert torch.equal(actual, expected), "Scale bytes mismatch"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"num_tokens,hidden_dim",
|
||||
[
|
||||
(1, 7168), # mn padded 1 -> 4
|
||||
(2, 7168), # mn padded 2 -> 4
|
||||
(3, 7168), # mn padded 3 -> 4
|
||||
(5, 7168), # mn padded 5 -> 8
|
||||
(127, 7168), # mn padded 127 -> 128
|
||||
(253, 640), # both mn and groups padded
|
||||
(1, 384), # extreme: 1 group, 1 mn row -> both axes padded
|
||||
],
|
||||
)
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="DeepGEMM not available on this platform",
|
||||
)
|
||||
def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q(
|
||||
num_tokens, hidden_dim
|
||||
):
|
||||
"""When output_q is allocated with shape (tma_aligned_mn, k) instead of
|
||||
(mn, k), the kernel must overwrite the padded mn rows with zeros so
|
||||
callers can use ``torch.empty`` instead of ``torch.zeros``."""
|
||||
|
||||
device = "cuda"
|
||||
group_size = 128
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
|
||||
|
||||
mn = num_tokens
|
||||
groups_per_row = hidden_dim // group_size
|
||||
k_num_packed = (groups_per_row + 3) // 4
|
||||
tma_aligned_mn = ((mn + 3) // 4) * 4
|
||||
|
||||
fp8_dtype = current_platform.fp8_dtype()
|
||||
fp8_min, fp8_max = get_fp8_min_max()
|
||||
# Allocate output_q with the padded mn extent and pre-fill with 0xFF
|
||||
# so the kernel cannot rely on a clean buffer.
|
||||
out_q = torch.empty((tma_aligned_mn, hidden_dim), device=device, dtype=fp8_dtype)
|
||||
out_q.view(torch.uint8).fill_(0xFF)
|
||||
|
||||
out_s_packed = torch.empty_strided(
|
||||
(mn, k_num_packed),
|
||||
(1, tma_aligned_mn),
|
||||
device=device,
|
||||
dtype=torch.int32,
|
||||
)
|
||||
|
||||
torch.ops._C.per_token_group_fp8_quant_packed(
|
||||
x, out_q, out_s_packed, group_size, 1e-10, fp8_min, fp8_max
|
||||
)
|
||||
|
||||
# Live rows must match the Triton reference.
|
||||
with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
|
||||
ref_q, _ = fp8_utils.per_token_group_quant_fp8(x, group_size, use_ue8m0=True)
|
||||
assert torch.equal(out_q[:mn], ref_q), "Live region mismatch"
|
||||
|
||||
# Padded rows must be all-zero; without this, downstream TMA loads would
|
||||
# see uninitialised data.
|
||||
if tma_aligned_mn > mn:
|
||||
padded_bytes = out_q[mn:tma_aligned_mn].view(torch.uint8)
|
||||
assert padded_bytes.eq(0).all(), (
|
||||
f"Padded rows [{mn}, {tma_aligned_mn}) not zeroed; "
|
||||
f"{padded_bytes.ne(0).sum().item()} non-zero bytes"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda_alike(),
|
||||
reason="packed FP8 per-token-group quant kernel requires a CUDA-alike GPU",
|
||||
)
|
||||
def test_per_token_group_quant_fp8_packed_large_mn():
|
||||
"""Regression test for https://github.com/vllm-project/vllm/issues/45099.
|
||||
|
||||
Some background: gridDim.x and gridDim.y have different limits of 2^31 - 1 and
|
||||
2^16 - 1, respectively.
|
||||
Prior code introduced a bug where it incorrectly assumed grid.x and y both have
|
||||
2^31 - 1 limits and mixed them up, which doesn't surface until the kernel is
|
||||
launched with a large mn that exceeds grid.y limit (2^16 - 1).
|
||||
|
||||
This issue doesn't surface often because each forward pass only processes a
|
||||
bounded token batch, not the full context.
|
||||
Quantizing tensors with more rows than that will fail at launch with
|
||||
"CUDA error: invalid argument".
|
||||
This is a differential test that compares fp8 output against Triton output
|
||||
reference when token size sits just above the gridDim.y 2^16 - 1 limit.
|
||||
"""
|
||||
|
||||
device = "cuda"
|
||||
group_size = 128
|
||||
# hidden 2048 -> 2048/128 = 16 groups per row -> kx=16, ry=1: one grid row per mn
|
||||
# row, so any mn > 65535 overflowed grid.y before the fix.
|
||||
num_tokens, hidden_dim = 65537, 2048
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
|
||||
|
||||
out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm(
|
||||
x,
|
||||
group_size=group_size,
|
||||
use_ue8m0=True,
|
||||
)
|
||||
|
||||
with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
|
||||
ref_q, ref_s = fp8_utils.per_token_group_quant_fp8(
|
||||
x, group_size, use_ue8m0=True
|
||||
)
|
||||
|
||||
assert torch.equal(out_q, ref_q), "Quantized output mismatch"
|
||||
|
||||
# Vectorized packed-scale check; the per-element loop used by the smaller
|
||||
# tests is too slow at this size. groups_per_row is a multiple of 4 here,
|
||||
# so there is no K padding and the packed view lines up.
|
||||
mn = num_tokens
|
||||
groups_per_row = hidden_dim // group_size
|
||||
k_num_packed = (groups_per_row + 3) // 4
|
||||
assert groups_per_row % 4 == 0
|
||||
ref_exponents = (ref_s.reshape(mn, groups_per_row).view(torch.int32) >> 23) & 0xFF
|
||||
exp = ref_exponents.view(mn, k_num_packed, 4)
|
||||
expected = (
|
||||
exp[..., 0] | (exp[..., 1] << 8) | (exp[..., 2] << 16) | (exp[..., 3] << 24)
|
||||
)
|
||||
assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
|
||||
@pytest.mark.parametrize("group_size", [64, 128])
|
||||
@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")
|
||||
def test_per_token_group_quant_int8(shape, group_size: int):
|
||||
device = "cuda"
|
||||
|
||||
torch.manual_seed(42)
|
||||
num_tokens, hidden_dim = shape
|
||||
|
||||
x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
|
||||
|
||||
# cuda path
|
||||
out_q, scale = int8_utils.per_token_group_quant_int8(
|
||||
x,
|
||||
group_size,
|
||||
)
|
||||
|
||||
# triton ref
|
||||
with patch("vllm.platforms.current_platform.is_cuda_alike", return_value=False):
|
||||
ref_q, ref_s = int8_utils.per_token_group_quant_int8(
|
||||
x,
|
||||
group_size,
|
||||
)
|
||||
|
||||
assert torch.allclose(out_q.float(), ref_q.float(), atol=0.15, rtol=0.15)
|
||||
assert torch.allclose(scale, ref_s, atol=0.01, rtol=0.01)
|
||||
@@ -0,0 +1,67 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the Triton dequant-gather kernel used by
|
||||
``CompressedTensorsEmbeddingWNA16Int`` (quantized embedding lookup)."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from compressed_tensors.compressors.pack_quantized.helpers import unpack_from_int32
|
||||
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_embedding import ( # noqa: E501
|
||||
_dequant_gather_triton,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
def _dequant_gather_torch(
|
||||
ids: torch.Tensor,
|
||||
weight_packed: torch.Tensor,
|
||||
weight_scale: torch.Tensor,
|
||||
hidden: int,
|
||||
num_bits: int,
|
||||
) -> torch.Tensor:
|
||||
"""Reference: gather packed rows by id, unpack int32-packed INT, dequant."""
|
||||
n = ids.shape[0]
|
||||
int8 = unpack_from_int32(weight_packed[ids], num_bits, torch.Size([n, hidden]))
|
||||
scale_rows = weight_scale[ids]
|
||||
w = int8.to(scale_rows.dtype)
|
||||
if scale_rows.shape[1] == 1:
|
||||
return w * scale_rows
|
||||
ng = scale_rows.shape[1]
|
||||
return (w.view(n, ng, hidden // ng) * scale_rows.unsqueeze(-1)).view(n, hidden)
|
||||
|
||||
|
||||
@pytest.mark.skipif(
|
||||
not current_platform.is_cuda(), reason="Triton dequant kernel requires CUDA"
|
||||
)
|
||||
@pytest.mark.parametrize("num_bits", [2, 4, 8])
|
||||
@pytest.mark.parametrize("group_size", [0, 256]) # 0 -> channel
|
||||
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
|
||||
@pytest.mark.parametrize("num_ids", [1, 17, 4096])
|
||||
def test_dequant_gather(num_bits, group_size, dtype, num_ids):
|
||||
torch.manual_seed(0)
|
||||
device = "cuda"
|
||||
vocab, hidden = 1000, 2048
|
||||
pack_factor = 32 // num_bits
|
||||
|
||||
# Random full-range int32 packed weights (covers the sign bit -> exercises the
|
||||
# arithmetic-shift + mask unpack path).
|
||||
weight_packed = torch.randint(
|
||||
-(2**31),
|
||||
2**31,
|
||||
(vocab, hidden // pack_factor),
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
|
||||
num_groups = 1 if group_size == 0 else hidden // group_size
|
||||
weight_scale = torch.rand(vocab, num_groups, dtype=dtype, device=device) + 0.01
|
||||
|
||||
ids = torch.randint(0, vocab, (num_ids,), dtype=torch.long, device=device)
|
||||
|
||||
out = _dequant_gather_triton(ids, weight_packed, weight_scale, hidden, num_bits)
|
||||
ref = _dequant_gather_torch(ids, weight_packed, weight_scale, hidden, num_bits)
|
||||
|
||||
assert out.shape == (num_ids, hidden)
|
||||
assert out.dtype == dtype
|
||||
torch.testing.assert_close(out, ref)
|
||||
@@ -0,0 +1,503 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Compile-guard tests for the ROCm RDNA3 W4A16 kernels (dense + MoE).
|
||||
|
||||
Verifies that the gfx1100 compilation and dispatch guards are hermetic:
|
||||
- On gfx1100: all ops exist, dispatch selects RDNA3 kernels.
|
||||
- On CDNA (gfx942/gfx950) or other non-gfx1100: ops must NOT exist,
|
||||
dispatch must fall through to Triton/Marlin, and no RDNA3 code
|
||||
path is reachable.
|
||||
|
||||
The negative (non-gfx1100) tests verify at three layers:
|
||||
1. Compile-level: on non-gfx1100 hardware, the RDNA3 ops are absent
|
||||
from the compiled _rocm_C extension — real binary verification.
|
||||
2. Static source analysis: parses CMakeLists.txt and torch_bindings.cpp
|
||||
to verify that all RDNA3 .cu files and op registrations are inside
|
||||
gfx1100-only guards.
|
||||
3. Runtime mock: patches on_gfx1100() to False and verifies that the
|
||||
Python dispatch chain rejects the RDNA3 path.
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_rdna3_compile_guards.py`.
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import regex as re
|
||||
import torch
|
||||
|
||||
import vllm
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_rocm():
|
||||
pytest.skip("RDNA3 compile-guard tests are ROCm-only", allow_module_level=True)
|
||||
|
||||
from vllm.platforms.rocm import on_gfx1100 # noqa: E402
|
||||
|
||||
gfx1100_only = pytest.mark.skipif(
|
||||
not on_gfx1100(),
|
||||
reason="Requires gfx1100 hardware",
|
||||
)
|
||||
|
||||
not_gfx1100 = pytest.mark.skipif(
|
||||
on_gfx1100(),
|
||||
reason="This test verifies non-gfx1100 builds — skip on gfx1100",
|
||||
)
|
||||
|
||||
RDNA3_OPS = ["gptq_gemm_rdna3", "gptq_gemm_rdna3_wmma", "moe_gptq_gemm_rdna3"]
|
||||
RDNA3_CU_FILES = [
|
||||
"q_gemm_rdna3.cu",
|
||||
"q_gemm_rdna3_wmma.cu",
|
||||
"moe_q_gemm_rdna3.cu",
|
||||
]
|
||||
|
||||
|
||||
def _find_repo_root() -> Path | None:
|
||||
"""Walk up from this file to find the repo root (has CMakeLists.txt)."""
|
||||
for parent in [Path(__file__).resolve(), *Path(__file__).resolve().parents]:
|
||||
if (parent / "CMakeLists.txt").exists() and (parent / "csrc").is_dir():
|
||||
return parent
|
||||
return None
|
||||
|
||||
|
||||
REPO_ROOT = _find_repo_root()
|
||||
|
||||
# Directory of the *installed* vllm python package. The .py guard checks read
|
||||
# from here so they verify the code that is actually imported at runtime — this
|
||||
# works even on CI images that ship the wheel instead of the python source tree
|
||||
# (where only csrc/ + CMakeLists.txt are checked out for building).
|
||||
VLLM_PKG_DIR: Path | None = (
|
||||
Path(vllm.__file__).parent if getattr(vllm, "__file__", None) else None
|
||||
)
|
||||
|
||||
needs_source = pytest.mark.skipif(
|
||||
REPO_ROOT is None,
|
||||
reason="C/CMake source tree not available (installed package only)",
|
||||
)
|
||||
|
||||
|
||||
def _read_source_or_skip(*relparts: str) -> str:
|
||||
"""Read a C/CMake source file from the repo tree, or skip if absent.
|
||||
|
||||
Used for csrc/ and CMakeLists.txt — these only exist in a source checkout,
|
||||
not in the installed wheel.
|
||||
"""
|
||||
assert REPO_ROOT is not None # callers are gated by @needs_source
|
||||
path = REPO_ROOT.joinpath(*relparts)
|
||||
if not path.exists():
|
||||
pytest.skip(f"{path} not present in this source tree")
|
||||
return path.read_text()
|
||||
|
||||
|
||||
def _read_pkg_source_or_skip(*relparts: str) -> str:
|
||||
"""Read a python source file from the installed vllm package.
|
||||
|
||||
Reflects the code actually loaded at runtime, so these guard checks run in
|
||||
CI against the wheel — no source checkout required. Only skips for an
|
||||
exotic install layout (namespace/zipimport) where __file__ is unavailable.
|
||||
"""
|
||||
if VLLM_PKG_DIR is None:
|
||||
pytest.skip("vllm package directory not resolvable (zip/namespace?)")
|
||||
assert VLLM_PKG_DIR is not None # narrow for mypy (skip above is NoReturn)
|
||||
path = VLLM_PKG_DIR.joinpath(*relparts)
|
||||
if not path.exists():
|
||||
pytest.skip(f"{path} not present in installed vllm package")
|
||||
return path.read_text()
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Part A: POSITIVE — on gfx1100, ops exist and dispatch works
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@gfx1100_only
|
||||
@pytest.mark.parametrize("op_name", RDNA3_OPS)
|
||||
def test_op_registered_on_gfx1100(op_name):
|
||||
"""On gfx1100, all RDNA3 ops must be registered in _rocm_C."""
|
||||
assert hasattr(torch.ops, "_rocm_C"), "_rocm_C module not loaded"
|
||||
assert hasattr(torch.ops._rocm_C, op_name), (
|
||||
f"_rocm_C.{op_name} not registered — "
|
||||
"check CMakeLists.txt VLLM_ROCM_HAS_GFX1100 "
|
||||
"and torch_bindings.cpp #ifdef VLLM_ROCM_GFX1100"
|
||||
)
|
||||
|
||||
|
||||
@gfx1100_only
|
||||
def test_all_ops_present_or_all_absent():
|
||||
"""The 3 RDNA3 ops are behind the same #ifdef — all present or all absent.
|
||||
|
||||
Catches someone accidentally moving an op outside the guard.
|
||||
"""
|
||||
has_rocm_c = hasattr(torch.ops, "_rocm_C")
|
||||
if not has_rocm_c:
|
||||
pytest.skip("_rocm_C not loaded")
|
||||
|
||||
present = {op: hasattr(torch.ops._rocm_C, op) for op in RDNA3_OPS}
|
||||
values = set(present.values())
|
||||
assert len(values) == 1, (
|
||||
f"Guard inconsistency — some RDNA3 ops registered, others not: "
|
||||
f"{present}. Check torch_bindings.cpp #ifdef VLLM_ROCM_GFX1100 block."
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Part B: NEGATIVE — compile-level verification on non-gfx1100
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@not_gfx1100
|
||||
@pytest.mark.parametrize("op_name", RDNA3_OPS)
|
||||
def test_op_absent_on_non_gfx1100(op_name):
|
||||
"""On non-gfx1100 (CDNA), RDNA3 ops must NOT exist in _rocm_C.
|
||||
|
||||
This is the real compile-level check: the binary was built without
|
||||
gfx1100 support, so the ops should not have been compiled or registered.
|
||||
"""
|
||||
if not hasattr(torch.ops, "_rocm_C"):
|
||||
return
|
||||
assert not hasattr(torch.ops._rocm_C, op_name), (
|
||||
f"_rocm_C.{op_name} is registered on non-gfx1100 hardware — "
|
||||
"compile guard is broken: check CMakeLists.txt "
|
||||
"VLLM_ROCM_HAS_GFX1100 and torch_bindings.cpp #ifdef"
|
||||
)
|
||||
|
||||
|
||||
@not_gfx1100
|
||||
def test_rocm_moe_not_supported_on_non_gfx1100():
|
||||
"""rocm_moe_rdna.is_supported() must return False on non-gfx1100 hardware."""
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
|
||||
rocm_moe_rdna,
|
||||
)
|
||||
|
||||
wq = type("WQ", (), {"num_bits": 4})()
|
||||
assert rocm_moe_rdna.is_supported(wq) is False, (
|
||||
"rocm_moe_rdna.is_supported() returned True on non-gfx1100 — "
|
||||
"dispatch guard is broken"
|
||||
)
|
||||
|
||||
|
||||
@not_gfx1100
|
||||
def test_dense_kernel_rejects_on_non_gfx1100():
|
||||
"""RDNA3W4A16LinearKernel.can_implement must reject on non-gfx1100."""
|
||||
from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E501
|
||||
MPLinearLayerConfig,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E501
|
||||
RDNA3W4A16LinearKernel,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=(1024, 256),
|
||||
partition_weight_shape=(1024, 256),
|
||||
weight_type=scalar_types.uint4b8,
|
||||
act_type=torch.float16,
|
||||
group_size=128,
|
||||
zero_points=False,
|
||||
has_g_idx=False,
|
||||
)
|
||||
ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
|
||||
assert ok is False, f"RDNA3 dense kernel accepted on non-gfx1100: {reason}"
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Part C: Static source analysis (build-level guards)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
@needs_source
|
||||
class TestCMakeGuards:
|
||||
"""Verify CMakeLists.txt only compiles RDNA3 .cu files for gfx1100."""
|
||||
|
||||
@staticmethod
|
||||
def _read_cmake():
|
||||
return _read_source_or_skip("CMakeLists.txt")
|
||||
|
||||
def test_rdna3_cu_files_inside_gfx1100_conditional(self):
|
||||
"""All RDNA3 .cu files must be listed inside the
|
||||
``if(VLLM_GPU_ARCHES MATCHES "gfx1100")`` block, not unconditionally.
|
||||
"""
|
||||
cmake = self._read_cmake()
|
||||
for cu_file in RDNA3_CU_FILES:
|
||||
assert cu_file in cmake, f"{cu_file} not found in CMakeLists.txt"
|
||||
|
||||
lines = cmake.splitlines()
|
||||
in_gfx1100_block = False
|
||||
for line in lines:
|
||||
if 'VLLM_GPU_ARCHES MATCHES "gfx1100"' in line:
|
||||
in_gfx1100_block = True
|
||||
if in_gfx1100_block and "endif()" in line:
|
||||
in_gfx1100_block = False
|
||||
if cu_file in line:
|
||||
assert in_gfx1100_block, (
|
||||
f"{cu_file} is listed OUTSIDE the gfx1100 "
|
||||
f"conditional in CMakeLists.txt — CDNA builds "
|
||||
f"would compile RDNA3 code. Line: {line.strip()}"
|
||||
)
|
||||
|
||||
def test_compile_definition_only_for_gfx1100(self):
|
||||
"""VLLM_ROCM_GFX1100 compile definition must be conditional."""
|
||||
cmake = self._read_cmake()
|
||||
lines = cmake.splitlines()
|
||||
in_gfx1100_block = False
|
||||
for line in lines:
|
||||
if "VLLM_ROCM_HAS_GFX1100)" in line:
|
||||
in_gfx1100_block = True
|
||||
if in_gfx1100_block and "endif()" in line:
|
||||
in_gfx1100_block = False
|
||||
if "VLLM_ROCM_GFX1100" in line and "target_compile_definitions" in line:
|
||||
assert in_gfx1100_block, (
|
||||
"VLLM_ROCM_GFX1100 compile definition is set outside "
|
||||
"the VLLM_ROCM_HAS_GFX1100 conditional — CDNA builds "
|
||||
f"would define it. Line: {line.strip()}"
|
||||
)
|
||||
|
||||
|
||||
@needs_source
|
||||
class TestTorchBindingsGuards:
|
||||
"""Verify torch_bindings.cpp gates all RDNA3 ops behind #ifdef."""
|
||||
|
||||
@staticmethod
|
||||
def _read_bindings():
|
||||
return _read_source_or_skip("csrc", "rocm", "torch_bindings.cpp")
|
||||
|
||||
def test_all_rdna3_ops_inside_ifdef(self):
|
||||
"""Every rdna3 op def/impl must be between #ifdef VLLM_ROCM_GFX1100
|
||||
and #endif. If any is outside, a CDNA build would try to register
|
||||
the op and link a symbol that doesn't exist.
|
||||
"""
|
||||
src = self._read_bindings()
|
||||
lines = src.splitlines()
|
||||
|
||||
inside_guard = False
|
||||
rdna3_lines_outside = []
|
||||
|
||||
for i, line in enumerate(lines, 1):
|
||||
if "#ifdef VLLM_ROCM_GFX1100" in line:
|
||||
inside_guard = True
|
||||
elif line.strip() == "#endif" and inside_guard:
|
||||
inside_guard = False
|
||||
|
||||
if (
|
||||
"rdna3" in line.lower()
|
||||
and not line.strip().startswith("//")
|
||||
and not inside_guard
|
||||
):
|
||||
rdna3_lines_outside.append((i, line.strip()))
|
||||
|
||||
assert not rdna3_lines_outside, (
|
||||
"RDNA3 op references found OUTSIDE #ifdef VLLM_ROCM_GFX1100 "
|
||||
"in torch_bindings.cpp — these would break CDNA builds:\n"
|
||||
+ "\n".join(f" L{n}: {s}" for n, s in rdna3_lines_outside)
|
||||
)
|
||||
|
||||
def test_no_unconditional_rdna3_includes(self):
|
||||
"""No #include of RDNA3-specific headers outside the guard."""
|
||||
src = self._read_bindings()
|
||||
lines = src.splitlines()
|
||||
|
||||
inside_guard = False
|
||||
for i, line in enumerate(lines, 1):
|
||||
if "#ifdef VLLM_ROCM_GFX1100" in line:
|
||||
inside_guard = True
|
||||
elif line.strip() == "#endif" and inside_guard:
|
||||
inside_guard = False
|
||||
|
||||
if "#include" in line and "rdna3" in line.lower():
|
||||
assert inside_guard, (
|
||||
f"L{i}: RDNA3 include outside gfx1100 guard: {line.strip()}"
|
||||
)
|
||||
|
||||
|
||||
class TestCustomOpsGuards:
|
||||
"""Verify _custom_ops.py gates register_fake behind hasattr checks."""
|
||||
|
||||
@staticmethod
|
||||
def _read_custom_ops():
|
||||
return _read_pkg_source_or_skip("_custom_ops.py")
|
||||
|
||||
def test_register_fake_guarded_by_hasattr(self):
|
||||
"""Every register_fake for an RDNA3 op must be preceded by a hasattr
|
||||
check — otherwise it would crash on import on CDNA where the ops
|
||||
don't exist.
|
||||
"""
|
||||
src = self._read_custom_ops()
|
||||
for op in RDNA3_OPS:
|
||||
pattern = rf'register_fake\(\s*"_rocm_C::{op}"\s*\)'
|
||||
match = re.search(pattern, src)
|
||||
if match is None:
|
||||
continue
|
||||
preceding = src[: match.start()]
|
||||
last_hasattr = preceding.rfind(f'hasattr(torch.ops._rocm_C, "{op}")')
|
||||
assert last_hasattr != -1, (
|
||||
f'register_fake("_rocm_C::{op}") is not preceded by a '
|
||||
f"hasattr check — would crash on CDNA import"
|
||||
)
|
||||
gap = preceding[last_hasattr:].count("\n")
|
||||
assert gap <= 5, (
|
||||
f"hasattr guard for {op} is {gap} lines before "
|
||||
f"register_fake — suspiciously far; verify it's the "
|
||||
f"actual guard and not a coincidence"
|
||||
)
|
||||
|
||||
def test_no_toplevel_rocm_c_import(self):
|
||||
"""No top-level ``from vllm._rocm_C import`` — would crash on CDNA."""
|
||||
src = self._read_custom_ops()
|
||||
for line in src.splitlines():
|
||||
stripped = line.strip()
|
||||
if stripped.startswith("#") or stripped.startswith("//"):
|
||||
continue
|
||||
assert "from vllm._rocm_C import" not in stripped, (
|
||||
f"Top-level import of _rocm_C in _custom_ops.py would "
|
||||
f"crash on CDNA: {stripped}"
|
||||
)
|
||||
|
||||
|
||||
# ============================================================================
|
||||
# Part D: Runtime mock (simulate CDNA on gfx1100 hardware)
|
||||
# ============================================================================
|
||||
|
||||
|
||||
class _FakeWeightQuant:
|
||||
"""Minimal stand-in for a weight quantization config."""
|
||||
|
||||
def __init__(self, num_bits):
|
||||
self.num_bits = num_bits
|
||||
|
||||
|
||||
class TestMoEDispatchMocked:
|
||||
"""Mock on_gfx1100() to False and verify RDNA3 MoE is unreachable."""
|
||||
|
||||
def test_is_supported_false_when_mocked_cdna(self):
|
||||
"""rocm_moe_rdna.is_supported() must return False when not on gfx1100."""
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
|
||||
rocm_moe_rdna,
|
||||
)
|
||||
|
||||
with patch("vllm.platforms.rocm.on_gfx1100", return_value=False):
|
||||
assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
|
||||
|
||||
@pytest.mark.parametrize("num_bits", [2, 3, 8, 16])
|
||||
def test_is_supported_rejects_non_w4(self, num_bits):
|
||||
"""is_supported() rejects non-4-bit even before checking arch."""
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
|
||||
rocm_moe_rdna,
|
||||
)
|
||||
|
||||
assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=num_bits)) is False
|
||||
|
||||
def test_is_supported_false_when_op_missing(self):
|
||||
"""is_supported() returns False when the C++ op doesn't exist."""
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
|
||||
rocm_moe_rdna,
|
||||
)
|
||||
|
||||
fake_rocm_c = type("FakeRocmC", (), {"gptq_gemm_rdna3": None})()
|
||||
with patch.object(torch, "ops", create=True) as mock_ops:
|
||||
mock_ops._rocm_C = fake_rocm_c
|
||||
assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
|
||||
|
||||
def test_is_supported_false_when_rocm_c_absent(self):
|
||||
"""is_supported() returns False when _rocm_C doesn't exist at all."""
|
||||
from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tensors_moe import ( # noqa: E501
|
||||
rocm_moe_rdna,
|
||||
)
|
||||
|
||||
fake_ops = type("FakeOps", (), {})()
|
||||
with patch.object(torch, "ops", fake_ops):
|
||||
assert rocm_moe_rdna.is_supported(_FakeWeightQuant(num_bits=4)) is False
|
||||
|
||||
|
||||
class TestDenseKernelSelectionMocked:
|
||||
"""Mock on_gfx1100() and verify dense RDNA3 kernel is not selected."""
|
||||
|
||||
@gfx1100_only
|
||||
def test_can_implement_rejects_when_mocked_cdna(self):
|
||||
"""RDNA3W4A16LinearKernel.can_implement must reject on mocked CDNA."""
|
||||
from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E501
|
||||
MPLinearLayerConfig,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E501
|
||||
RDNA3W4A16LinearKernel,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=(1024, 256),
|
||||
partition_weight_shape=(1024, 256),
|
||||
weight_type=scalar_types.uint4b8,
|
||||
act_type=torch.float16,
|
||||
group_size=128,
|
||||
zero_points=False,
|
||||
has_g_idx=False,
|
||||
)
|
||||
ok, _ = RDNA3W4A16LinearKernel.can_implement(config)
|
||||
assert ok is True
|
||||
|
||||
with (
|
||||
patch("vllm.platforms.rocm.on_gfx1100", return_value=False),
|
||||
patch("vllm.platforms.rocm._ON_GFX1100", False),
|
||||
):
|
||||
ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
|
||||
assert ok is False, f"RDNA3 kernel accepted on simulated CDNA: {reason}"
|
||||
|
||||
@gfx1100_only
|
||||
def test_chooser_skips_rdna3_when_mocked_cdna(self):
|
||||
"""choose_mp_linear_kernel must NOT return RDNA3 on mocked CDNA."""
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
choose_mp_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E501
|
||||
MPLinearLayerConfig,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=(1024, 256),
|
||||
partition_weight_shape=(1024, 256),
|
||||
weight_type=scalar_types.uint4b8,
|
||||
act_type=torch.float16,
|
||||
group_size=128,
|
||||
zero_points=False,
|
||||
has_g_idx=False,
|
||||
)
|
||||
with (
|
||||
patch("vllm.platforms.rocm.on_gfx1100", return_value=False),
|
||||
patch("vllm.platforms.rocm._ON_GFX1100", False),
|
||||
):
|
||||
chosen = choose_mp_linear_kernel(config)
|
||||
assert chosen.__name__ != "RDNA3W4A16LinearKernel", (
|
||||
"RDNA3 kernel was selected on simulated CDNA — "
|
||||
"choose_mp_linear_kernel guard is broken"
|
||||
)
|
||||
|
||||
|
||||
class TestCompressedTensorsMoEDispatchGuard:
|
||||
"""Verify compressed_tensors_moe.py only enters rocm_moe_rdna under is_rocm()."""
|
||||
|
||||
def test_rocm_guard_in_dispatch_source(self):
|
||||
"""The rocm_moe_rdna import and call must be inside an is_rocm() check."""
|
||||
src = _read_pkg_source_or_skip(
|
||||
"model_executor",
|
||||
"layers",
|
||||
"quantization",
|
||||
"compressed_tensors",
|
||||
"compressed_tensors_moe",
|
||||
"compressed_tensors_moe.py",
|
||||
)
|
||||
lines = src.splitlines()
|
||||
|
||||
for i, line in enumerate(lines, 1):
|
||||
stripped = line.strip()
|
||||
if "rocm_moe" in stripped and not stripped.startswith("#"):
|
||||
found_guard = False
|
||||
for j in range(i - 1, max(0, i - 15), -1):
|
||||
if "is_rocm()" in lines[j - 1]:
|
||||
found_guard = True
|
||||
break
|
||||
assert found_guard, (
|
||||
f"L{i}: rocm_moe_rdna reference not protected by "
|
||||
f"is_rocm() guard: {stripped}"
|
||||
)
|
||||
@@ -0,0 +1,367 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Correctness tests for the ROCm RDNA3 fused MoE W4A16 HIP kernel (gfx1100).
|
||||
|
||||
Tests ``moe_gptq_gemm_rdna3`` against the dense ``gptq_gemm_rdna3`` as
|
||||
reference: builds RDNA3-format weights (shuffled int32, synthesized qzeros),
|
||||
runs the fused MoE kernel, and compares per-expert results.
|
||||
|
||||
Model parameters taken from:
|
||||
- cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit
|
||||
(hidden=2048, inter=768, E=128, top_k=8, G=32)
|
||||
- Qwen3.6-35B-A3B-GPTQ-W4A16-G32
|
||||
(hidden=2048, inter=512, E=256, top_k=8, G=32)
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_rdna3_moe_w4a16.py`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_rocm():
|
||||
pytest.skip("RDNA3 MoE W4A16 kernel is ROCm-only", allow_module_level=True)
|
||||
|
||||
from vllm import _custom_ops as ops # noqa: E402
|
||||
from vllm.model_executor.layers.fused_moe.activation import ( # noqa: E402
|
||||
MoEActivation,
|
||||
apply_moe_activation,
|
||||
)
|
||||
from vllm.model_executor.layers.fused_moe.moe_align_block_size import ( # noqa: E402
|
||||
moe_align_block_size,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402
|
||||
pack_quantized_values_into_int32,
|
||||
)
|
||||
from vllm.platforms.rocm import on_gfx1100 # noqa: E402
|
||||
from vllm.scalar_type import scalar_types # noqa: E402
|
||||
|
||||
device = "cuda"
|
||||
|
||||
gfx1100_only = pytest.mark.skipif(
|
||||
not (
|
||||
on_gfx1100()
|
||||
and hasattr(torch.ops, "_rocm_C")
|
||||
and hasattr(torch.ops._rocm_C, "moe_gptq_gemm_rdna3")
|
||||
),
|
||||
reason="Requires gfx1100 with moe_gptq_gemm_rdna3 op",
|
||||
)
|
||||
|
||||
# Model configurations: real K/N/top_k/group_size dims, E capped at 16 to
|
||||
# fit in test GPU memory (full E=128/256 would need >20GB for weights alone).
|
||||
# Kernel behavior is E-independent (per-expert tiling), so E=16 is sufficient.
|
||||
MODEL_CONFIGS = [
|
||||
# cyankiwi/Qwen3-30B-A3B-Instruct-2507-AWQ-4bit dims (E capped)
|
||||
pytest.param(16, 2048, 768, 8, 32, id="Qwen3-30B-A3B"),
|
||||
# Qwen3.6-35B-A3B-GPTQ-W4A16-G32 dims (E capped)
|
||||
pytest.param(16, 2048, 512, 8, 32, id="Qwen3.6-35B-A3B"),
|
||||
]
|
||||
|
||||
# Token counts: decode (1), small batch (4), medium (16), prefill (64)
|
||||
NUM_TOKENS = [1, 4, 16, 64, 256, 512]
|
||||
|
||||
|
||||
def _make_packed_weights(E, K, N):
|
||||
"""Create random 4-bit packed weights [E, K/8, N] int32 + shuffle."""
|
||||
w = torch.randint(0, 16, (E, K, N), dtype=torch.int32, device=device)
|
||||
packed = torch.zeros(E, K // 8, N, dtype=torch.int32, device=device)
|
||||
for i in range(8):
|
||||
packed |= (w[:, i::8, :] & 0xF) << (i * 4)
|
||||
g_idx = torch.empty(0, dtype=torch.int32, device=device)
|
||||
for e in range(E):
|
||||
we = packed[e].contiguous()
|
||||
ops.gptq_shuffle(we, g_idx, 4)
|
||||
packed[e] = we
|
||||
return packed
|
||||
|
||||
|
||||
def _make_scales(E, groups, N, dtype):
|
||||
return torch.rand(E, groups, N, dtype=dtype, device=device) * 0.1
|
||||
|
||||
|
||||
def _make_qzeros(E, groups, N):
|
||||
zeros = torch.full(
|
||||
(groups, N),
|
||||
scalar_types.uint4b8.bias - 1,
|
||||
dtype=torch.int32,
|
||||
device=device,
|
||||
)
|
||||
qz = pack_quantized_values_into_int32(
|
||||
zeros,
|
||||
scalar_types.uint4b8,
|
||||
packed_dim=1,
|
||||
)
|
||||
return qz.unsqueeze(0).expand(E, -1, -1).contiguous()
|
||||
|
||||
|
||||
@gfx1100_only
|
||||
@pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS)
|
||||
@pytest.mark.parametrize("M", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("block_size_m", [1, 4])
|
||||
def test_fused_moe_w1_matches_dense(
|
||||
E, K, N_inter, top_k, group_size, M, dtype, block_size_m
|
||||
):
|
||||
"""w1 GEMM via fused kernel matches per-expert dense kernel."""
|
||||
N_gate_up = N_inter * 2
|
||||
groups = K // group_size
|
||||
|
||||
torch.manual_seed(42)
|
||||
x = torch.randn(M, K, dtype=dtype, device=device)
|
||||
w13 = _make_packed_weights(E, K, N_gate_up)
|
||||
w13_s = _make_scales(E, groups, N_gate_up, dtype)
|
||||
w13_z = _make_qzeros(E, groups, N_gate_up)
|
||||
g_idx = torch.empty(0, dtype=torch.int32, device=device)
|
||||
|
||||
topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32)
|
||||
si, ei, ntp = moe_align_block_size(topk_ids, block_size_m, E)
|
||||
|
||||
# Fused kernel
|
||||
fused_out = torch.zeros(M * top_k, N_gate_up, dtype=dtype, device=device)
|
||||
ops.moe_gptq_gemm_rdna3(
|
||||
x,
|
||||
fused_out,
|
||||
w13,
|
||||
w13_s,
|
||||
w13_z,
|
||||
torch.empty(0, device=device),
|
||||
si,
|
||||
ei,
|
||||
ntp,
|
||||
top_k,
|
||||
block_size_m,
|
||||
False,
|
||||
0,
|
||||
)
|
||||
|
||||
# Per-expert dense reference
|
||||
ref_out = torch.zeros(M * top_k, N_gate_up, dtype=dtype, device=device)
|
||||
for m in range(M):
|
||||
for k in range(top_k):
|
||||
e = topk_ids[m, k].item()
|
||||
flat = m * top_k + k
|
||||
ref = ops.gptq_gemm_rdna3(
|
||||
x[m : m + 1],
|
||||
w13[e],
|
||||
w13_z[e],
|
||||
w13_s[e],
|
||||
g_idx,
|
||||
False,
|
||||
)
|
||||
ref_out[flat] = ref.squeeze()
|
||||
|
||||
# Split-K atomics can cause minor fp16/bf16 rounding differences
|
||||
# at large K (e.g. K=2048 → 8 K-blocks). Use allclose, not equal.
|
||||
atol = 0.5 if dtype == torch.bfloat16 else 0.1
|
||||
assert torch.allclose(fused_out, ref_out, atol=atol, rtol=0.01), (
|
||||
f"max diff: {(fused_out - ref_out).abs().max().item()}"
|
||||
)
|
||||
|
||||
|
||||
@gfx1100_only
|
||||
@pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS)
|
||||
@pytest.mark.parametrize("M", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_fused_moe_output_topk_reduces(E, K, N_inter, top_k, group_size, M, dtype):
|
||||
"""output_topk fuses moe_sum: multiple experts write to same output row."""
|
||||
groups = K // group_size
|
||||
|
||||
torch.manual_seed(123)
|
||||
x = torch.randn(M * top_k, K, dtype=dtype, device=device)
|
||||
w = _make_packed_weights(E, K, N_inter)
|
||||
ws = _make_scales(E, groups, N_inter, dtype)
|
||||
wz = _make_qzeros(E, groups, N_inter)
|
||||
|
||||
topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32)
|
||||
topk_w = torch.softmax(
|
||||
torch.randn(M, top_k, device=device),
|
||||
dim=-1,
|
||||
).float()
|
||||
|
||||
si, ei, ntp = moe_align_block_size(topk_ids, 1, E)
|
||||
|
||||
# Without output_topk: write to [M*top_k, N] then moe_sum
|
||||
flat_out = torch.zeros(M * top_k, N_inter, dtype=dtype, device=device)
|
||||
ops.moe_gptq_gemm_rdna3(
|
||||
x,
|
||||
flat_out,
|
||||
w,
|
||||
ws,
|
||||
wz,
|
||||
topk_w.view(-1),
|
||||
si,
|
||||
ei,
|
||||
ntp,
|
||||
1,
|
||||
1,
|
||||
True,
|
||||
0,
|
||||
)
|
||||
ref = torch.zeros(M, N_inter, dtype=dtype, device=device)
|
||||
ops.moe_sum(flat_out.view(M, top_k, N_inter), ref)
|
||||
|
||||
# With output_topk: write directly to [M, N]
|
||||
fused = torch.zeros(M, N_inter, dtype=dtype, device=device)
|
||||
ops.moe_gptq_gemm_rdna3(
|
||||
x,
|
||||
fused,
|
||||
w,
|
||||
ws,
|
||||
wz,
|
||||
topk_w.view(-1),
|
||||
si,
|
||||
ei,
|
||||
ntp,
|
||||
1,
|
||||
1,
|
||||
True,
|
||||
top_k,
|
||||
)
|
||||
|
||||
atol = 1.0 if dtype == torch.bfloat16 else 0.1
|
||||
assert torch.allclose(fused, ref, atol=atol, rtol=0.01), (
|
||||
f"max diff: {(fused - ref).abs().max().item()}"
|
||||
)
|
||||
|
||||
|
||||
@gfx1100_only
|
||||
@pytest.mark.parametrize("E, K, N_inter, top_k, group_size", MODEL_CONFIGS)
|
||||
@pytest.mark.parametrize("M", NUM_TOKENS)
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_full_moe_e2e(E, K, N_inter, top_k, group_size, M, dtype):
|
||||
"""Full MoE forward: w1 + silu_and_mul + w2 with output_topk reduce."""
|
||||
N_gate_up = N_inter * 2
|
||||
hidden = K
|
||||
|
||||
torch.manual_seed(7)
|
||||
x = torch.randn(M, K, dtype=dtype, device=device)
|
||||
w13 = _make_packed_weights(E, K, N_gate_up)
|
||||
w13_s = _make_scales(E, K // group_size, N_gate_up, dtype)
|
||||
w13_z = _make_qzeros(E, K // group_size, N_gate_up)
|
||||
w2 = _make_packed_weights(E, N_inter, hidden)
|
||||
w2_s = _make_scales(E, N_inter // group_size, hidden, dtype)
|
||||
w2_z = _make_qzeros(E, N_inter // group_size, hidden)
|
||||
g_idx = torch.empty(0, dtype=torch.int32, device=device)
|
||||
|
||||
topk_ids = torch.randint(0, E, (M, top_k), device=device, dtype=torch.int32)
|
||||
topk_w = torch.softmax(
|
||||
torch.randn(M, top_k, device=device),
|
||||
dim=-1,
|
||||
).float()
|
||||
|
||||
si, ei, ntp = moe_align_block_size(topk_ids, 1, E)
|
||||
|
||||
# Fused path (what apply() does)
|
||||
w1_out = torch.zeros(M * top_k, N_gate_up, dtype=dtype, device=device)
|
||||
ops.moe_gptq_gemm_rdna3(
|
||||
x,
|
||||
w1_out,
|
||||
w13,
|
||||
w13_s,
|
||||
w13_z,
|
||||
torch.empty(0, device=device),
|
||||
si,
|
||||
ei,
|
||||
ntp,
|
||||
top_k,
|
||||
1,
|
||||
False,
|
||||
0,
|
||||
)
|
||||
act_out = torch.empty(M * top_k, N_inter, dtype=dtype, device=device)
|
||||
apply_moe_activation(MoEActivation.SILU, act_out, w1_out)
|
||||
fused = torch.zeros(M, hidden, dtype=dtype, device=device)
|
||||
ops.moe_gptq_gemm_rdna3(
|
||||
act_out,
|
||||
fused,
|
||||
w2,
|
||||
w2_s,
|
||||
w2_z,
|
||||
topk_w.view(-1),
|
||||
si,
|
||||
ei,
|
||||
ntp,
|
||||
1,
|
||||
1,
|
||||
True,
|
||||
top_k,
|
||||
)
|
||||
|
||||
# Per-expert reference
|
||||
ref = torch.zeros(M, hidden, dtype=dtype, device=device)
|
||||
for m_idx in range(M):
|
||||
for k_idx in range(top_k):
|
||||
e = topk_ids[m_idx, k_idx].item()
|
||||
w = topk_w[m_idx, k_idx].item()
|
||||
r1 = ops.gptq_gemm_rdna3(
|
||||
x[m_idx : m_idx + 1],
|
||||
w13[e],
|
||||
w13_z[e],
|
||||
w13_s[e],
|
||||
g_idx,
|
||||
False,
|
||||
)
|
||||
a = torch.empty(1, N_inter, dtype=dtype, device=device)
|
||||
apply_moe_activation(MoEActivation.SILU, a, r1)
|
||||
r2 = ops.gptq_gemm_rdna3(
|
||||
a,
|
||||
w2[e],
|
||||
w2_z[e],
|
||||
w2_s[e],
|
||||
g_idx,
|
||||
False,
|
||||
)
|
||||
ref[m_idx] += r2.squeeze() * w
|
||||
|
||||
# E2E chains w1 + activation + w2 + topk_w + output_topk reduce.
|
||||
# Each step accumulates rounding error (split-K atomics, topk_w
|
||||
# multiply order). Use relative L2 norm like the dense kernel test.
|
||||
diff_l2 = torch.norm(fused.float() - ref.float())
|
||||
ref_l2 = torch.norm(ref.float())
|
||||
rel_l2 = (diff_l2 / ref_l2).item() if ref_l2 > 0 else 0.0
|
||||
threshold = 0.05 if dtype == torch.float16 else 0.10
|
||||
assert rel_l2 < threshold, (
|
||||
f"rel L2 = {rel_l2:.4f} (threshold {threshold}), "
|
||||
f"max abs diff: {(fused - ref).abs().max().item()}"
|
||||
)
|
||||
|
||||
|
||||
@gfx1100_only
|
||||
def test_expert_id_minus_one():
|
||||
"""Kernel handles expert_id == -1 (expert parallelism) without crash."""
|
||||
# Qwen3-30B-A3B dims (E capped for memory)
|
||||
E, K, N = 16, 2048, 768
|
||||
groups = K // 32
|
||||
|
||||
w = _make_packed_weights(E, K, N)
|
||||
ws = _make_scales(E, groups, N, torch.bfloat16)
|
||||
wz = _make_qzeros(E, groups, N)
|
||||
x = torch.randn(1, K, dtype=torch.bfloat16, device=device)
|
||||
|
||||
# Manually create sorted_token_ids/expert_ids with -1
|
||||
sorted_ids = torch.tensor([0], dtype=torch.int32, device=device)
|
||||
expert_ids = torch.tensor([-1], dtype=torch.int32, device=device)
|
||||
ntp = torch.tensor([1], dtype=torch.int32, device=device)
|
||||
|
||||
out = torch.zeros(1, N, dtype=torch.bfloat16, device=device)
|
||||
ops.moe_gptq_gemm_rdna3(
|
||||
x,
|
||||
out,
|
||||
w,
|
||||
ws,
|
||||
wz,
|
||||
torch.empty(0, device=device),
|
||||
sorted_ids,
|
||||
expert_ids,
|
||||
ntp,
|
||||
1,
|
||||
1,
|
||||
False,
|
||||
0,
|
||||
)
|
||||
current_platform.synchronize()
|
||||
|
||||
# Output should remain zero (expert skipped)
|
||||
assert torch.equal(out, torch.zeros_like(out))
|
||||
@@ -0,0 +1,278 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Correctness tests for the ROCm RDNA3 W4A16 GPTQ kernel (gfx1100).
|
||||
|
||||
Exercises ``RDNA3W4A16LinearKernel`` end-to-end: it builds a layer with
|
||||
GPTQ-format checkpoint parameters, runs ``process_weights_after_loading``
|
||||
(weight shuffle + zero-point synthesis), then ``apply_weights``, and compares
|
||||
the result against an fp32 reference dequant-and-matmul.
|
||||
|
||||
The kernel is exposed via ``torch.ops._rocm_C.gptq_gemm_rdna3`` and is only
|
||||
built for gfx11; tests are skipped elsewhere.
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_rdna3_w4a16.py`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_rocm():
|
||||
pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True)
|
||||
|
||||
from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E402
|
||||
MPLinearLayerConfig,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E402
|
||||
RDNA3W4A16LinearKernel,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import ( # noqa: E402
|
||||
pack_quantized_values_into_int32,
|
||||
)
|
||||
from vllm.model_executor.parameter import ( # noqa: E402
|
||||
GroupQuantScaleParameter,
|
||||
PackedvLLMParameter,
|
||||
)
|
||||
from vllm.platforms.rocm import on_gfx1100 # noqa: E402
|
||||
from vllm.scalar_type import scalar_types # noqa: E402
|
||||
from vllm.utils.torch_utils import set_random_seed # noqa: E402
|
||||
|
||||
device = "cuda"
|
||||
|
||||
WEIGHT_TYPE = scalar_types.uint4b8 # symmetric int4, bias = 8
|
||||
PACK_FACTOR = 8 # 8 x 4-bit nibbles per int32
|
||||
|
||||
# Skip everything in this module unless we are on the only architecture the
|
||||
# kernel is built/registered for.
|
||||
gfx1100_only = pytest.mark.skipif(
|
||||
not (
|
||||
on_gfx1100()
|
||||
and hasattr(torch.ops, "_rocm_C")
|
||||
and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3")
|
||||
),
|
||||
reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in",
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Reference implementation
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _reference(
|
||||
x_mk: torch.Tensor,
|
||||
q_int4_kn: torch.Tensor,
|
||||
scales_gn: torch.Tensor,
|
||||
zeros_gn: torch.Tensor | None,
|
||||
group_size: int,
|
||||
bias: torch.Tensor | None,
|
||||
) -> torch.Tensor:
|
||||
"""fp32 reference for the RDNA3 W4A16 op.
|
||||
|
||||
x_mk: [M, K] fp16/bf16 activations.
|
||||
q_int4_kn: [K, N] int32 raw stored nibbles in [0, 15].
|
||||
scales_gn: [K//G, N] per-group scales (act dtype).
|
||||
zeros_gn: [K//G, N] int32 raw stored zero points in [0, 15], or None
|
||||
for the symmetric path (kernel synthesizes stored zero = 7).
|
||||
group_size: G.
|
||||
|
||||
The kernel applies the GPTQv1 "+1" zero-point quirk, so the effective
|
||||
zero is ``stored_zero + 1`` (symmetric path: 7 + 1 == bias == 8).
|
||||
"""
|
||||
K, N = q_int4_kn.shape
|
||||
s_full = scales_gn.repeat_interleave(group_size, dim=0).to(torch.float32) # [K,N]
|
||||
if zeros_gn is None:
|
||||
z_full = torch.full(
|
||||
(K, N), float(WEIGHT_TYPE.bias), device=x_mk.device, dtype=torch.float32
|
||||
)
|
||||
else:
|
||||
z_full = (zeros_gn + 1).repeat_interleave(group_size, dim=0).to(torch.float32)
|
||||
w_fp = (q_int4_kn.to(torch.float32) - z_full) * s_full # [K, N]
|
||||
out = x_mk.to(torch.float32) @ w_fp # [M, N]
|
||||
if bias is not None:
|
||||
out = out + bias.to(torch.float32)
|
||||
return out.to(x_mk.dtype)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Layer construction (GPTQ checkpoint format)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _build_layer(
|
||||
q_int4_kn: torch.Tensor,
|
||||
scales_gn: torch.Tensor,
|
||||
zeros_gn: torch.Tensor | None,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.nn.Module:
|
||||
"""Build a dummy layer carrying GPTQ-format params, as the loader would."""
|
||||
no_loader = lambda *args, **kwargs: None # noqa: E731
|
||||
|
||||
# qweight: int4 packed along K into int32 -> [K//8, N].
|
||||
qweight = pack_quantized_values_into_int32(q_int4_kn, WEIGHT_TYPE, packed_dim=0)
|
||||
|
||||
class DummyLayer(torch.nn.Module):
|
||||
pass
|
||||
|
||||
layer = DummyLayer()
|
||||
layer.register_parameter(
|
||||
"qweight",
|
||||
PackedvLLMParameter(
|
||||
data=qweight,
|
||||
weight_loader=no_loader,
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=0,
|
||||
packed_factor=PACK_FACTOR,
|
||||
),
|
||||
)
|
||||
layer.register_parameter(
|
||||
"scales",
|
||||
GroupQuantScaleParameter(
|
||||
data=scales_gn.to(dtype),
|
||||
weight_loader=no_loader,
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
),
|
||||
)
|
||||
if zeros_gn is not None:
|
||||
# qzeros: int4 packed along N into int32 -> [K//G, N//8].
|
||||
qzeros = pack_quantized_values_into_int32(zeros_gn, WEIGHT_TYPE, packed_dim=1)
|
||||
layer.register_parameter(
|
||||
"qzeros",
|
||||
PackedvLLMParameter(
|
||||
data=qzeros,
|
||||
weight_loader=no_loader,
|
||||
input_dim=0,
|
||||
output_dim=1,
|
||||
packed_dim=1,
|
||||
packed_factor=PACK_FACTOR,
|
||||
),
|
||||
)
|
||||
return layer
|
||||
|
||||
|
||||
def _run_kernel(
|
||||
x_mk: torch.Tensor,
|
||||
q_int4_kn: torch.Tensor,
|
||||
scales_gn: torch.Tensor,
|
||||
zeros_gn: torch.Tensor | None,
|
||||
group_size: int,
|
||||
bias: torch.Tensor | None,
|
||||
dtype: torch.dtype,
|
||||
) -> torch.Tensor:
|
||||
K, N = q_int4_kn.shape
|
||||
has_zp = zeros_gn is not None
|
||||
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=(K, N),
|
||||
partition_weight_shape=(K, N),
|
||||
weight_type=WEIGHT_TYPE,
|
||||
act_type=dtype,
|
||||
group_size=group_size,
|
||||
zero_points=has_zp,
|
||||
has_g_idx=False,
|
||||
)
|
||||
ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
|
||||
assert ok, f"can_implement rejected a supported config: {reason}"
|
||||
|
||||
layer = _build_layer(q_int4_kn, scales_gn, zeros_gn, dtype)
|
||||
kernel = RDNA3W4A16LinearKernel(
|
||||
config,
|
||||
w_q_param_name="qweight",
|
||||
w_s_param_name="scales",
|
||||
w_zp_param_name="qzeros" if has_zp else None,
|
||||
w_gidx_param_name=None,
|
||||
)
|
||||
kernel.process_weights_after_loading(layer)
|
||||
return kernel.apply_weights(layer, x_mk, bias=bias)
|
||||
|
||||
|
||||
# Relative-L2 tolerance per dtype. The bf16 path widens dequantized weights
|
||||
# to fp32 and accumulates in fp32, so it matches the reference almost exactly
|
||||
# (<0.4% incl. the WMMA prefill path). The fp16 path uses the exllamav2
|
||||
# "+1024" bit-trick (see qdq_4_rdna3.cuh): the dequantized weight is recovered
|
||||
# as the fp16 difference of two ~1024*scale magnitudes, which sheds low-order
|
||||
# mantissa bits and leaves ~2-3% relative noise that accumulates over K. We
|
||||
# compare on the relative Frobenius norm rather than elementwise, since the
|
||||
# bit-trick noise produces large *relative* errors on individual near-zero
|
||||
# outputs that carry negligible absolute weight.
|
||||
_REL_L2_TOL = {torch.float16: 5e-2, torch.bfloat16: 1e-2}
|
||||
|
||||
|
||||
def _assert_close(out: torch.Tensor, ref: torch.Tensor, dtype: torch.dtype):
|
||||
rel_l2 = (out.to(torch.float32) - ref.to(torch.float32)).norm() / ref.to(
|
||||
torch.float32
|
||||
).norm()
|
||||
tol = _REL_L2_TOL[dtype]
|
||||
assert rel_l2 < tol, f"relative L2 error {rel_l2:.4f} exceeds {tol} for {dtype}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Forward correctness
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# (M, K, N, group_size). M spans the scalar decode path (small M) and the
|
||||
# WMMA prefill path (M >= 16 on the bf16 dispatch). K/N satisfy the kernel's
|
||||
# divisibility constraints (K % G == 0, K % 8 == 0, N % 8 == 0).
|
||||
MKNG_SHAPES = [
|
||||
(1, 128, 128, 128), # single group, decode
|
||||
(2, 256, 256, 128), # two groups
|
||||
(8, 256, 512, 64), # M=8 scalar, smaller group
|
||||
(16, 512, 256, 128), # M=16 -> WMMA path for bf16
|
||||
(32, 512, 512, 64), # larger prefill
|
||||
]
|
||||
|
||||
|
||||
@gfx1100_only
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("has_zp", [False, True], ids=["no_zp", "with_zp"])
|
||||
@pytest.mark.parametrize(
|
||||
"M,K,N,G", MKNG_SHAPES, ids=[f"m{m}_k{k}_n{n}_g{g}" for m, k, n, g in MKNG_SHAPES]
|
||||
)
|
||||
def test_rdna3_w4a16_matches_reference(dtype, has_zp, M, K, N, G, dist_init):
|
||||
set_random_seed(0)
|
||||
assert K % G == 0 and K % PACK_FACTOR == 0 and N % PACK_FACTOR == 0
|
||||
|
||||
groups = K // G
|
||||
x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
|
||||
q_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
|
||||
scales_gn = (
|
||||
0.05 * torch.rand((groups, N), device=device, dtype=torch.float32) + 0.01
|
||||
).to(dtype)
|
||||
zeros_gn = (
|
||||
torch.randint(0, 16, (groups, N), device=device, dtype=torch.int32)
|
||||
if has_zp
|
||||
else None
|
||||
)
|
||||
|
||||
out = _run_kernel(x_mk, q_int4_kn, scales_gn, zeros_gn, G, None, dtype)
|
||||
ref = _reference(x_mk, q_int4_kn, scales_gn, zeros_gn, G, None)
|
||||
|
||||
assert out.shape == (M, N) and out.dtype == dtype
|
||||
_assert_close(out, ref, dtype)
|
||||
|
||||
|
||||
@gfx1100_only
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize("M", [1, 32], ids=["decode", "prefill"])
|
||||
def test_rdna3_w4a16_bias(dtype, M, dist_init):
|
||||
"""Bias is added on both the scalar (M=1) and WMMA (M=32) paths."""
|
||||
set_random_seed(0)
|
||||
K, N, G = 512, 256, 128
|
||||
groups = K // G
|
||||
|
||||
x_mk = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
|
||||
q_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
|
||||
scales_gn = (
|
||||
0.05 * torch.rand((groups, N), device=device, dtype=torch.float32) + 0.01
|
||||
).to(dtype)
|
||||
bias = (0.1 * torch.randn(N, device=device, dtype=torch.float32)).to(dtype)
|
||||
|
||||
out = _run_kernel(x_mk, q_int4_kn, scales_gn, None, G, bias, dtype)
|
||||
ref = _reference(x_mk, q_int4_kn, scales_gn, None, G, bias)
|
||||
|
||||
_assert_close(out, ref, dtype)
|
||||
@@ -0,0 +1,89 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Kernel-selection / gating tests for the ROCm RDNA3 W4A16 GPTQ kernel.
|
||||
|
||||
Verifies that ``choose_mp_linear_kernel`` resolves a supported W4A16 GPTQ
|
||||
config to ``RDNA3W4A16LinearKernel`` on gfx1100 (it is registered ahead of
|
||||
``TritonW4A16LinearKernel`` in the ROCm priority list), and that
|
||||
``RDNA3W4A16LinearKernel.can_implement`` rejects the configs it does not
|
||||
support so selection falls through to the next kernel.
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_rdna3_w4a16_selection.py`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
if not current_platform.is_rocm():
|
||||
pytest.skip("RDNA3 W4A16 kernel is ROCm-only", allow_module_level=True)
|
||||
|
||||
from vllm.model_executor.kernels.linear import ( # noqa: E402
|
||||
choose_mp_linear_kernel,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import ( # noqa: E402
|
||||
MPLinearLayerConfig,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear.mixed_precision.rdna3_w4a16 import ( # noqa: E402
|
||||
RDNA3W4A16LinearKernel,
|
||||
)
|
||||
from vllm.platforms.rocm import on_gfx1100 # noqa: E402
|
||||
from vllm.scalar_type import scalar_types # noqa: E402
|
||||
|
||||
WEIGHT_TYPE = scalar_types.uint4b8 # symmetric int4, bias = 8
|
||||
|
||||
# The kernel is only selectable when running on gfx1100 with the custom op
|
||||
# compiled in; otherwise can_implement rejects and selection falls through.
|
||||
gfx1100_only = pytest.mark.skipif(
|
||||
not (
|
||||
on_gfx1100()
|
||||
and hasattr(torch.ops, "_rocm_C")
|
||||
and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3")
|
||||
),
|
||||
reason="requires gfx1100 with the _rocm_C.gptq_gemm_rdna3 op built in",
|
||||
)
|
||||
|
||||
|
||||
@gfx1100_only
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_selection_prefers_rdna3(dtype):
|
||||
"""A supported W4A16 GPTQ config resolves to the RDNA3 kernel on gfx1100."""
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=(1024, 256),
|
||||
partition_weight_shape=(1024, 256),
|
||||
weight_type=WEIGHT_TYPE,
|
||||
act_type=dtype,
|
||||
group_size=128,
|
||||
zero_points=False,
|
||||
has_g_idx=False,
|
||||
)
|
||||
assert choose_mp_linear_kernel(config).__name__ == "RDNA3W4A16LinearKernel"
|
||||
|
||||
|
||||
@gfx1100_only
|
||||
@pytest.mark.parametrize(
|
||||
"weight_type,group_size,N,full_k,expected_ok",
|
||||
[
|
||||
(scalar_types.uint4b8, 128, 256, 1024, True), # nominal: supported
|
||||
(scalar_types.uint4b8, -1, 256, 1024, False), # channelwise unsupported
|
||||
(scalar_types.uint4b8, 128, 252, 1024, False), # N not a multiple of 8
|
||||
(scalar_types.uint4b8, 96, 256, 1024, False), # group does not divide K
|
||||
(scalar_types.uint8b128, 128, 256, 1024, False), # wrong quant type
|
||||
],
|
||||
ids=["ok", "channelwise", "bad_n", "group_ndiv_k", "wrong_qtype"],
|
||||
)
|
||||
def test_can_implement(weight_type, group_size, N, full_k, expected_ok):
|
||||
"""can_implement gates on quant type, group size, and N divisibility."""
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=(full_k, N),
|
||||
partition_weight_shape=(full_k, N),
|
||||
weight_type=weight_type,
|
||||
act_type=torch.float16,
|
||||
group_size=group_size,
|
||||
zero_points=False,
|
||||
has_g_idx=False,
|
||||
)
|
||||
ok, reason = RDNA3W4A16LinearKernel.can_implement(config)
|
||||
assert ok is expected_ok, reason
|
||||
@@ -0,0 +1,35 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""End-to-end smoke test for CT W4A16 models on ROCm.
|
||||
|
||||
This validates that a real compressed-tensors W4A16 model can run inference
|
||||
end-to-end (which will exercise the Triton W4A16 kernel when selected).
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_rocm_compressed_tensors_w4a16.py`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"model_path",
|
||||
[
|
||||
# Listed in tests/weight_loading/models.txt
|
||||
"nm-testing/tinyllama-oneshot-w4a16-group128-v2",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("max_tokens", [32])
|
||||
@pytest.mark.skipif(not current_platform.is_rocm(), reason="Should only run on ROCm")
|
||||
def test_rocm_compressed_tensors_w4a16_e2e(
|
||||
vllm_runner, example_prompts, model_path, max_tokens
|
||||
):
|
||||
# Use fp16 activations for maximum compatibility.
|
||||
# gpu_memory_utilization lowered to work on shared nodes.
|
||||
with vllm_runner(
|
||||
model_path, dtype="float16", gpu_memory_utilization=0.3
|
||||
) as vllm_model:
|
||||
# If the W4A16 kernel is broken, this will typically throw.
|
||||
vllm_model.generate_greedy(example_prompts, max_tokens=max_tokens)
|
||||
@@ -0,0 +1,274 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import math
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm._custom_ops as ops
|
||||
from tests.kernels.quant_utils import ref_dynamic_per_tensor_fp8_quant
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.platforms.rocm import on_gfx950
|
||||
from vllm.utils.platform_utils import num_compute_units
|
||||
|
||||
DTYPES = [torch.bfloat16, torch.float16]
|
||||
BIAS_MODES = [0, 1, 2]
|
||||
# Specific (N, K, M) combinations for targeted testing
|
||||
NKM_FACTORS_LLMM1 = [
|
||||
# Small, medium, large cases
|
||||
(1, 8, 16),
|
||||
(1, 32, 64),
|
||||
(1, 128, 256),
|
||||
(1, 512, 1024),
|
||||
(1, 2048, 4096),
|
||||
# Edge cases with specific K sizes
|
||||
(1, 6144, 1024),
|
||||
(1, 8192, 2048),
|
||||
# Very large case
|
||||
(1, 4096, 8192),
|
||||
]
|
||||
|
||||
NKM_FACTORS_WVSPLITK = [
|
||||
# Different batch sizes with key dimensions
|
||||
(1, 32, 16),
|
||||
(1, 64, 64),
|
||||
(2, 256, 256),
|
||||
(3, 1024, 1024),
|
||||
(4, 4096, 4096),
|
||||
(4, 4096, 4096 + 1),
|
||||
(4, 4096 + 16, 4096),
|
||||
(4, 4096 + 16, 4096 + 1),
|
||||
# Extended K values
|
||||
(1, 9216, 512),
|
||||
(2, 10240, 1024),
|
||||
(4, 16384, 8192),
|
||||
(4, 16384 * 2, 8192),
|
||||
(4, 16384 * 2, 8192 + 1),
|
||||
(4, 16384 * 2 + 16, 8192),
|
||||
(4, 16384 * 2 + 16, 8192 + 1),
|
||||
# Minimum M constraint validation (m >= 8)
|
||||
(1, 64, 8),
|
||||
(2, 128, 8),
|
||||
(4, 256, 8),
|
||||
]
|
||||
|
||||
N_FACTORS_WVSPLITKRC = [
|
||||
13,
|
||||
16,
|
||||
17,
|
||||
25,
|
||||
29,
|
||||
31,
|
||||
32,
|
||||
41,
|
||||
51,
|
||||
64,
|
||||
71,
|
||||
81,
|
||||
91,
|
||||
103,
|
||||
117,
|
||||
128,
|
||||
]
|
||||
K_FACTORS_WVSPLITKRC = [2880, 2880 + 8, 3072, 3072 + 8]
|
||||
M_FACTORS_WVSPLITKRC = [128, 128 + 16, 256, 256 + 16, 640, 640 + 16]
|
||||
|
||||
NKM_FACTORS_WVSPLITK_FP8 = [
|
||||
# FP8-specific cases with K % 16 == 0
|
||||
(1, 16, 16),
|
||||
(1, 32, 16 + 16),
|
||||
(1, 64, 64),
|
||||
(1, 64, 64 + 16),
|
||||
(1, 64 + 16, 64),
|
||||
(1, 64 + 16, 64 + 16),
|
||||
(4, 64, 64),
|
||||
(4, 64, 64 + 16),
|
||||
(4, 64 + 16, 64),
|
||||
(4, 64 + 16, 64 + 16),
|
||||
(2, 512, 512),
|
||||
(3, 512, 512),
|
||||
(3, 512, 512 + 16),
|
||||
(4, 512, 512),
|
||||
(3, 2048, 2048),
|
||||
(3, 2048, 2048 + 16),
|
||||
(4, 2048 + 16, 2048),
|
||||
(4, 2048 + 16, 2048 + 16),
|
||||
(4, 4096, 4096),
|
||||
(4, 16400, 2048),
|
||||
(4, 16400, 2048 + 16),
|
||||
# Extended FP8 dimensions not covered by WVSPLITK
|
||||
(1, 14336, 1024),
|
||||
(2, 24576, 2048),
|
||||
(4, 32768, 28672),
|
||||
(4, 32768 * 2, 28672),
|
||||
(4, 32768 * 2, 28672 + 16),
|
||||
(4, 32768 * 2 + 16, 28672),
|
||||
(4, 32768 * 2 + 16, 28672 + 16),
|
||||
]
|
||||
|
||||
SEEDS = [0]
|
||||
|
||||
|
||||
def pad_fp8(weight):
|
||||
num_pad = 256 // weight.element_size()
|
||||
import torch.nn.functional as F
|
||||
|
||||
return F.pad(weight, (0, num_pad), "constant", 0)[..., :-num_pad]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("xnorm", [False, True])
|
||||
@pytest.mark.parametrize("n", N_FACTORS_WVSPLITKRC)
|
||||
@pytest.mark.parametrize("k", K_FACTORS_WVSPLITKRC)
|
||||
@pytest.mark.parametrize("m", M_FACTORS_WVSPLITKRC)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("padded_a", [False, True])
|
||||
@pytest.mark.parametrize("bias_mode", BIAS_MODES)
|
||||
@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
|
||||
@pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950")
|
||||
def test_rocm_wvsplitkrc_kernel(xnorm, n, k, m, dtype, seed, padded_a, bias_mode):
|
||||
torch.manual_seed(seed)
|
||||
cu_count = num_compute_units()
|
||||
|
||||
# Next ^2 of n
|
||||
N_p2 = 1 << (n - 1).bit_length()
|
||||
# With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile),
|
||||
# and each working on a 512-shard of K, how many CUs would we need?
|
||||
rndup_cus = ((m + 64 - 1) // 64) * ((k + 512 - 1) // 512)
|
||||
# How many of 4 waves in a group can work on same 16 Ms at same time?
|
||||
# This reduces the Ms each group works on, i.e. increasing the number of CUs needed.
|
||||
GrpsShrB = min(N_p2 // 16, 4)
|
||||
# Given the above, how many CUs would we need?
|
||||
CuNeeded = rndup_cus * GrpsShrB
|
||||
# candidate for atomic reduce count splitk?
|
||||
fits_wvsplitkrc = (N_p2 * m * ((k + 512 - 1) // 512)) <= 128 * 1024 * 12
|
||||
fits_wvsplitkrc &= CuNeeded <= cu_count
|
||||
|
||||
if not fits_wvsplitkrc:
|
||||
pytest.skip("Too large for wvSplitKrc")
|
||||
|
||||
xavier = (
|
||||
math.sqrt(2 / k) if xnorm else 1
|
||||
) # normalize to avoid large output-bias deltas
|
||||
A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
|
||||
B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
|
||||
if padded_a:
|
||||
A = pad_fp8(A)
|
||||
|
||||
BIAS = None
|
||||
if bias_mode == 1:
|
||||
BIAS = torch.rand(m, dtype=dtype, device="cuda") * 2 - 1
|
||||
elif bias_mode == 2:
|
||||
BIAS = torch.rand(n, m, dtype=dtype, device="cuda") * 2 - 1
|
||||
elif bias_mode == 3:
|
||||
BIAS = torch.rand(1, m, dtype=dtype, device="cuda") * 2 - 1
|
||||
|
||||
ref_out = torch.nn.functional.linear(A, B, BIAS)
|
||||
out = ops.wvSplitKrc(A, B, cu_count, BIAS)
|
||||
|
||||
if xnorm:
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-8)
|
||||
else:
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_LLMM1)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("rows_per_block", [2, 4, 8, 16])
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
|
||||
@torch.inference_mode()
|
||||
def test_rocm_llmm1_kernel(n, k, m, dtype, rows_per_block, seed):
|
||||
torch.manual_seed(seed)
|
||||
# TODO: Zero-centering the inputs causes errors for LLMM1!
|
||||
# Without that the numbers quickly saturate, and may
|
||||
# be giving false matches.
|
||||
A = torch.rand(n, k, dtype=dtype, device="cuda")
|
||||
B = torch.rand(m, k, dtype=dtype, device="cuda")
|
||||
|
||||
ref_out = torch.matmul(A, B.t())
|
||||
out = ops.LLMM1(B, A, rows_per_block)
|
||||
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-8, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("xnorm", [False, True])
|
||||
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
|
||||
@pytest.mark.parametrize("bias_mode", BIAS_MODES)
|
||||
@pytest.mark.parametrize("padded_a", [False, True])
|
||||
@pytest.mark.parametrize("padded_b", [False, True])
|
||||
def test_rocm_wvsplitk_kernel(
|
||||
xnorm, n, k, m, dtype, seed, bias_mode, padded_a, padded_b
|
||||
):
|
||||
torch.manual_seed(seed)
|
||||
cu_count = num_compute_units()
|
||||
|
||||
xavier = (
|
||||
math.sqrt(2 / k) if xnorm else 1
|
||||
) # normalize to avoid large output-bias deltas
|
||||
A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
|
||||
B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
|
||||
|
||||
BIAS = None
|
||||
if bias_mode == 1:
|
||||
BIAS = torch.rand(m, dtype=dtype, device="cuda") * 2 - 1
|
||||
elif bias_mode == 2:
|
||||
BIAS = torch.rand(n, m, dtype=dtype, device="cuda") * 2 - 1
|
||||
|
||||
if padded_a:
|
||||
A = pad_fp8(A)
|
||||
if padded_b:
|
||||
B = pad_fp8(B)
|
||||
|
||||
ref_out = torch.nn.functional.linear(A, B, BIAS)
|
||||
out = ops.wvSplitK(B, A.view(-1, A.size(-1)), cu_count, BIAS)
|
||||
|
||||
# Accumulation error in fp16 GEMM scales with sqrt(K)
|
||||
atol = torch.finfo(dtype).eps * math.sqrt(k)
|
||||
torch.testing.assert_close(out, ref_out, atol=atol, rtol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("xnorm", [False, True])
|
||||
@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITK_FP8)
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("seed", SEEDS)
|
||||
@pytest.mark.parametrize("padded_a", [False, True])
|
||||
@pytest.mark.parametrize("padded_b", [False, True])
|
||||
@pytest.mark.parametrize("biased", [False, True])
|
||||
@pytest.mark.skipif(
|
||||
not (current_platform.is_rocm() and current_platform.supports_fp8()),
|
||||
reason="only test for rocm fp8",
|
||||
)
|
||||
def test_rocm_wvsplitk_fp8_kernel(
|
||||
xnorm, n, k, m, dtype, seed, padded_a, padded_b, biased
|
||||
):
|
||||
torch.manual_seed(seed)
|
||||
|
||||
xavier = math.sqrt(2 / k) if xnorm else 1 # normalize to avoid large deltas
|
||||
A = (torch.rand(n, k, device="cuda") * 2 - 1) * xavier
|
||||
B = (torch.rand(m, k, device="cuda") * 2 - 1) * xavier
|
||||
|
||||
A, scale_a = ref_dynamic_per_tensor_fp8_quant(A)
|
||||
B, scale_b = ref_dynamic_per_tensor_fp8_quant(B)
|
||||
if padded_b:
|
||||
B = pad_fp8(B)
|
||||
if padded_a:
|
||||
A = pad_fp8(A)
|
||||
|
||||
BIAS = None if (not biased) else (torch.rand(m, dtype=dtype, device="cuda") * 2 - 1)
|
||||
|
||||
ref_out = torch._scaled_mm(
|
||||
A, B.t(), out_dtype=dtype, scale_a=scale_a, scale_b=scale_b, bias=BIAS
|
||||
)
|
||||
out = ops.wvSplitKQ(B, A, dtype, scale_a, scale_b, num_compute_units(), BIAS)
|
||||
|
||||
if xnorm:
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-8)
|
||||
elif k >= 32 * 1024:
|
||||
# wider pytrch thresh for large-K & no xnorm
|
||||
torch.testing.assert_close(out, ref_out, atol=0.07, rtol=5e-2)
|
||||
else:
|
||||
torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
|
||||
@@ -0,0 +1,129 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for ScaledMM kernel selection logic (CPU-only)
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_scaled_mm_kernel_selection.py`.
|
||||
"""
|
||||
|
||||
import inspect
|
||||
from abc import ABC
|
||||
from unittest.mock import patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
AiterInt8ScaledMMLinearKernel,
|
||||
CPUInt8ScaledMMLinearKernel,
|
||||
Int8ScaledMMLinearKernel,
|
||||
Int8ScaledMMLinearLayerConfig,
|
||||
ScaledMMLinearKernel,
|
||||
init_int8_linear_kernel,
|
||||
register_linear_kernel,
|
||||
)
|
||||
from vllm.platforms import PlatformEnum
|
||||
|
||||
pytestmark = pytest.mark.cpu_test
|
||||
|
||||
|
||||
def test_is_supported_is_abstract():
|
||||
"""Test that is_supported() is properly defined as abstract."""
|
||||
assert issubclass(ScaledMMLinearKernel, ABC)
|
||||
assert hasattr(ScaledMMLinearKernel, "is_supported")
|
||||
|
||||
|
||||
def test_cpu_kernel_implements_is_supported():
|
||||
"""Test that CPUInt8ScaledMMLinearKernel implements is_supported() method."""
|
||||
assert hasattr(CPUInt8ScaledMMLinearKernel, "is_supported"), (
|
||||
"CPUInt8ScaledMMLinearKernel missing is_supported() method"
|
||||
)
|
||||
# Verify it's a classmethod by checking if it can be called with the class
|
||||
# and by checking the method type
|
||||
assert inspect.ismethod(
|
||||
CPUInt8ScaledMMLinearKernel.is_supported
|
||||
) or inspect.isfunction(CPUInt8ScaledMMLinearKernel.is_supported), (
|
||||
"CPUInt8ScaledMMLinearKernel.is_supported() should be a classmethod"
|
||||
)
|
||||
# Verify it can be called as a classmethod
|
||||
result, reason = CPUInt8ScaledMMLinearKernel.is_supported()
|
||||
assert isinstance(result, bool), "is_supported() should return a bool"
|
||||
assert reason is None or isinstance(reason, str), "reason should be str or None"
|
||||
|
||||
|
||||
def test_aiter_kernel_implements_is_supported():
|
||||
"""Test that AiterInt8ScaledMMLinearKernel implements is_supported() method."""
|
||||
assert hasattr(AiterInt8ScaledMMLinearKernel, "is_supported"), (
|
||||
"AiterInt8ScaledMMLinearKernel missing is_supported() method"
|
||||
)
|
||||
# Verify it's a classmethod by checking if it can be called with the class
|
||||
# and by checking the method type
|
||||
assert inspect.ismethod(
|
||||
AiterInt8ScaledMMLinearKernel.is_supported
|
||||
) or inspect.isfunction(AiterInt8ScaledMMLinearKernel.is_supported), (
|
||||
"AiterInt8ScaledMMLinearKernel.is_supported() should be a classmethod"
|
||||
)
|
||||
# Verify it can be called as a classmethod
|
||||
# (will return False on CPU, which is expected)
|
||||
result, reason = AiterInt8ScaledMMLinearKernel.is_supported()
|
||||
assert isinstance(result, bool), "is_supported() should return a bool"
|
||||
assert reason is None or isinstance(reason, str), "reason should be str or None"
|
||||
# On CPU, it should return False with a reason about requiring ROCm
|
||||
# This validates the method works correctly even on non-ROCm platforms
|
||||
|
||||
|
||||
def test_cpu_kernel_accepts_all_configs():
|
||||
"""Test that CPUInt8ScaledMMLinearKernel accepts all config combinations."""
|
||||
configs = [
|
||||
Int8ScaledMMLinearLayerConfig(
|
||||
is_channelwise=False,
|
||||
is_static_input_scheme=True,
|
||||
input_symmetric=True,
|
||||
),
|
||||
Int8ScaledMMLinearLayerConfig(
|
||||
is_channelwise=True,
|
||||
is_static_input_scheme=False,
|
||||
input_symmetric=False,
|
||||
),
|
||||
]
|
||||
|
||||
for config in configs:
|
||||
can_impl, reason = CPUInt8ScaledMMLinearKernel.can_implement(config)
|
||||
assert can_impl, (
|
||||
f"CPUInt8ScaledMMLinearKernel should accept config {config}: {reason}"
|
||||
)
|
||||
|
||||
|
||||
class OOTInt8ScaledMMLinearKernel(Int8ScaledMMLinearKernel):
|
||||
@classmethod
|
||||
def is_supported(
|
||||
cls, compute_capability: int | None = None
|
||||
) -> tuple[bool, str | None]:
|
||||
return True, None
|
||||
|
||||
@classmethod
|
||||
def can_implement(cls, c: Int8ScaledMMLinearLayerConfig) -> tuple[bool, str | None]:
|
||||
return True, None
|
||||
|
||||
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
||||
pass
|
||||
|
||||
def apply_weights(
|
||||
self,
|
||||
layer: torch.nn.Module,
|
||||
x: torch.Tensor,
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
pass
|
||||
|
||||
|
||||
@patch("vllm.model_executor.kernels.linear.current_platform")
|
||||
def test_register_oot_linear_kernel(platform_mock):
|
||||
"""Test that the linear kernel registration works correctly."""
|
||||
platform_mock._enum = PlatformEnum.OOT
|
||||
register_linear_kernel(OOTInt8ScaledMMLinearKernel, PlatformEnum.OOT, "int8")
|
||||
|
||||
kernel = init_int8_linear_kernel(True, True, True, "module")
|
||||
|
||||
assert isinstance(kernel, OOTInt8ScaledMMLinearKernel), (
|
||||
"init_int8_linear_kernel should return an instance of the registered kernel"
|
||||
)
|
||||
@@ -0,0 +1,78 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.kernels.quantization.nvfp4_utils import (
|
||||
FLOAT4_E2M1_MAX,
|
||||
FLOAT8_E4M3_MAX,
|
||||
dequantize_nvfp4_to_dtype,
|
||||
)
|
||||
from vllm._custom_ops import scaled_fp4_quant
|
||||
from vllm.model_executor.layers.activation import SiluAndMul
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
if not current_platform.has_device_capability(100):
|
||||
pytest.skip(
|
||||
reason="Nvfp4 Requires compute capability of 10 or above.",
|
||||
allow_module_level=True,
|
||||
)
|
||||
|
||||
FP4_DTYPE = torch.uint8
|
||||
FP8_DTYPE = current_platform.fp8_dtype()
|
||||
|
||||
DTYPES = [torch.float16, torch.bfloat16]
|
||||
SHAPES = [(128, 256), (128, 128), (256, 256), (256, 128)]
|
||||
BLOCK_SIZE = 16
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", DTYPES)
|
||||
@pytest.mark.parametrize("shape", SHAPES)
|
||||
@torch.inference_mode()
|
||||
def test_silu_mul_nvfp4_quant(
|
||||
default_vllm_config,
|
||||
dtype: torch.dtype,
|
||||
shape: tuple[int, int],
|
||||
) -> None:
|
||||
set_random_seed(42)
|
||||
device = "cuda:0"
|
||||
torch.set_default_device(device)
|
||||
|
||||
x = torch.randn(shape, dtype=dtype)
|
||||
|
||||
# ref op
|
||||
ref_output = SiluAndMul().forward_native(x)
|
||||
ref_global_scale = (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.abs(
|
||||
ref_output
|
||||
).max().to(torch.float32)
|
||||
ref_output_quant, ref_block_scale = scaled_fp4_quant(ref_output, ref_global_scale)
|
||||
|
||||
# fused op
|
||||
fused_output_quant = torch.empty_like(ref_output_quant)
|
||||
fused_block_scale = torch.empty_like(ref_block_scale)
|
||||
torch.ops._C.silu_and_mul_nvfp4_quant(
|
||||
fused_output_quant, fused_block_scale, x, ref_global_scale
|
||||
)
|
||||
|
||||
# check dtype
|
||||
assert ref_output_quant.dtype == FP4_DTYPE
|
||||
assert fused_output_quant.dtype == FP4_DTYPE
|
||||
assert ref_output_quant.shape == fused_output_quant.shape
|
||||
|
||||
assert ref_block_scale.dtype == FP8_DTYPE
|
||||
assert fused_block_scale.dtype == FP8_DTYPE
|
||||
assert ref_block_scale.shape == fused_block_scale.shape
|
||||
|
||||
# check dequantized output
|
||||
ref_output_dequant = dequantize_nvfp4_to_dtype(
|
||||
ref_output_quant, ref_block_scale, ref_global_scale, dtype, device
|
||||
)
|
||||
fused_output_dequant = dequantize_nvfp4_to_dtype(
|
||||
fused_output_quant, fused_block_scale, ref_global_scale, dtype, device
|
||||
)
|
||||
|
||||
atol, rtol = 3e-1, 3e-1
|
||||
torch.testing.assert_close(
|
||||
ref_output_dequant, fused_output_dequant, atol=atol, rtol=rtol
|
||||
)
|
||||
@@ -0,0 +1,125 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the triton_scaled_mm kernel
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_triton_scaled_mm.py`.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
device = "cuda"
|
||||
|
||||
triton_scaled_mm_module = importlib.import_module(
|
||||
"vllm.model_executor.layers.quantization.compressed_tensors.triton_scaled_mm"
|
||||
)
|
||||
triton_scaled_mm = triton_scaled_mm_module.triton_scaled_mm
|
||||
|
||||
|
||||
def torch_scaled_mm(
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
scale_a: torch.Tensor,
|
||||
scale_b: torch.Tensor,
|
||||
out_dtype: type[torch.dtype],
|
||||
bias: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
out = torch.mm(a.to(torch.float32), b.to(torch.float32))
|
||||
out = scale_a * out
|
||||
out = scale_b.T * out
|
||||
out = out.to(out_dtype)
|
||||
if bias is not None:
|
||||
out = out + bias
|
||||
|
||||
return out
|
||||
|
||||
|
||||
def get_8bit_types():
|
||||
types = [torch.int8]
|
||||
if current_platform.supports_fp8():
|
||||
types.append(current_platform.fp8_dtype())
|
||||
return types
|
||||
|
||||
|
||||
# This test is to check regressions for int8 support on ROCm.
|
||||
@pytest.mark.parametrize(
|
||||
"model_path",
|
||||
[
|
||||
"neuralmagic/Llama-3.2-1B-quantized.w8a8",
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("max_tokens", [32])
|
||||
@pytest.mark.parametrize("num_logprobs", [10])
|
||||
@pytest.mark.skipif(not current_platform.is_rocm(), reason="Should only run on ROCm")
|
||||
def test_rocm_compressed_tensors_w8a8(
|
||||
vllm_runner, example_prompts, model_path, max_tokens, num_logprobs
|
||||
):
|
||||
dtype = "bfloat16"
|
||||
|
||||
with vllm_runner(model_path, dtype=dtype) as vllm_model:
|
||||
vllm_model.generate_greedy_logprobs(example_prompts, max_tokens, num_logprobs)
|
||||
|
||||
|
||||
MNK_FACTORS = [
|
||||
(1, 256, 128),
|
||||
(33, 256, 496),
|
||||
(64, 971, 1024),
|
||||
(64, 20486, 128),
|
||||
(512, 256, 496),
|
||||
(512, 20486, 1024),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("M,N,K", MNK_FACTORS)
|
||||
@pytest.mark.parametrize("out_dtype", [torch.bfloat16])
|
||||
@pytest.mark.parametrize("in_dtype", get_8bit_types())
|
||||
@pytest.mark.parametrize("use_scalar_scale_a", [True, False])
|
||||
@pytest.mark.parametrize("use_scalar_scale_b", [True, False])
|
||||
@pytest.mark.parametrize("use_bias", [True, False])
|
||||
def test_scaled_mm(
|
||||
M, N, K, in_dtype, out_dtype, use_scalar_scale_a, use_scalar_scale_b, use_bias
|
||||
):
|
||||
is_floating_point_type = lambda t: torch.tensor([1, 1], dtype=t).is_floating_point()
|
||||
|
||||
set_random_seed(0)
|
||||
|
||||
# NOTE: There are cases, where if the matrix is large enough, an output
|
||||
# like 65504.4 can be produced, and can easily turn into inf when
|
||||
# multiplied when using float16/bfloat16. This means one function, e.g.,
|
||||
# testing function, and another function, e.g. golden function, can
|
||||
# produce a non-inf value while the other produces an inf value, and
|
||||
# will cause assert_close/allclose to fail, even though if overflow
|
||||
# wouldn't have occurred, the values would have been "close."
|
||||
#
|
||||
# So, the values here are kept small enough to avoid this situation.
|
||||
if is_floating_point_type(in_dtype):
|
||||
a = (0.25 * torch.rand((M, K), dtype=torch.float32, device=device)).to(in_dtype)
|
||||
b = (0.25 * torch.rand((K, N), dtype=torch.float32, device=device)).to(in_dtype)
|
||||
else:
|
||||
a = torch.randint(-32, 32, (M, K), dtype=in_dtype, device=device)
|
||||
b = torch.randint(-32, 32, (K, N), dtype=in_dtype, device=device)
|
||||
|
||||
if use_scalar_scale_a:
|
||||
scale_a = torch.rand((1, 1), device=device)
|
||||
else:
|
||||
scale_a = 0.25 * torch.rand((M, 1), device=device)
|
||||
|
||||
if use_scalar_scale_b:
|
||||
scale_b = torch.rand((1, 1), device=device)
|
||||
else:
|
||||
scale_b = 0.25 * torch.rand((N, 1), device=device)
|
||||
|
||||
bias = None
|
||||
if use_bias:
|
||||
bias = torch.rand((N,), device=device, dtype=out_dtype)
|
||||
|
||||
c_check = triton_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
|
||||
|
||||
c_actual = torch_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias)
|
||||
|
||||
torch.testing.assert_close(c_check, c_actual, rtol=1e-1, atol=1e-1)
|
||||
@@ -0,0 +1,304 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for the ROCm Triton W4A16 GEMM kernel.
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_triton_w4a16.py`.
|
||||
"""
|
||||
|
||||
import importlib
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.utils.torch_utils import set_random_seed
|
||||
|
||||
# This test module is ROCm/Triton specific. Avoid import-time failures on
|
||||
# non-ROCm or environments without Triton by skipping early.
|
||||
if not current_platform.is_rocm():
|
||||
pytest.skip("ROCm only", allow_module_level=True)
|
||||
|
||||
pytest.importorskip("triton")
|
||||
|
||||
device = "cuda"
|
||||
|
||||
triton_w4a16_module = importlib.import_module(
|
||||
"vllm.model_executor.kernels.linear.mixed_precision.triton_w4a16"
|
||||
)
|
||||
triton_w4a16_gemm = triton_w4a16_module.triton_w4a16_gemm
|
||||
TritonW4A16LinearKernel = triton_w4a16_module.TritonW4A16LinearKernel
|
||||
|
||||
|
||||
def _pack_int4_along_n(w_int4_kn: torch.Tensor) -> torch.Tensor:
|
||||
"""Pack int4 values along N: [K, N] -> [K, N//8] int32."""
|
||||
assert w_int4_kn.dtype == torch.int32
|
||||
K, N = w_int4_kn.shape
|
||||
assert N % 8 == 0
|
||||
shifts = torch.arange(8, device=w_int4_kn.device, dtype=torch.int32) * 4
|
||||
return torch.sum(
|
||||
(w_int4_kn.view(K, N // 8, 8) & 0xF) << shifts,
|
||||
dim=2,
|
||||
dtype=torch.int32,
|
||||
).contiguous()
|
||||
|
||||
|
||||
def _unpack_int4_along_n(w_packed_kn8: torch.Tensor) -> torch.Tensor:
|
||||
"""Unpack int4 values along N: [K, N//8] -> [K, N] int32."""
|
||||
assert w_packed_kn8.dtype == torch.int32
|
||||
K, N8 = w_packed_kn8.shape
|
||||
shifts = torch.arange(8, device=w_packed_kn8.device, dtype=torch.int32) * 4
|
||||
nibbles = (w_packed_kn8.unsqueeze(-1) >> shifts) & 0xF
|
||||
return nibbles.reshape(K, N8 * 8)
|
||||
|
||||
|
||||
def _pack_int4_along_k_to_ckpt(w_int4_kn: torch.Tensor) -> torch.Tensor:
|
||||
"""Pack int4 values along K into CT checkpoint layout: [K,N] -> [N, K//8]."""
|
||||
assert w_int4_kn.dtype == torch.int32
|
||||
K, N = w_int4_kn.shape
|
||||
assert K % 8 == 0
|
||||
out = torch.zeros((N, K // 8), dtype=torch.int32, device=w_int4_kn.device)
|
||||
for i in range(8):
|
||||
out |= (w_int4_kn[i::8, :].t() & 0xF) << (i * 4)
|
||||
return out.contiguous()
|
||||
|
||||
|
||||
def _w4a16_reference(
|
||||
a_mk: torch.Tensor,
|
||||
b_packed_kn8: torch.Tensor,
|
||||
scales_gn: torch.Tensor,
|
||||
*,
|
||||
group_size: int,
|
||||
qzeros_gn8: torch.Tensor | None,
|
||||
zp_bias: int,
|
||||
) -> torch.Tensor:
|
||||
"""Reference implementation for W4A16.
|
||||
|
||||
a_mk: [M,K] fp16/bf16
|
||||
b_packed_kn8: [K, N//8] int32, N-packed int4 weights
|
||||
scales_gn: [K//G, N] fp16/bf16
|
||||
qzeros_gn8: [K//G, N//8] int32, N-packed int4 zeros, or None
|
||||
"""
|
||||
assert a_mk.dtype in (torch.float16, torch.bfloat16)
|
||||
assert b_packed_kn8.dtype == torch.int32
|
||||
assert scales_gn.dtype == a_mk.dtype
|
||||
|
||||
M, K = a_mk.shape
|
||||
N = b_packed_kn8.shape[1] * 8
|
||||
assert b_packed_kn8.shape[0] == K
|
||||
|
||||
assert group_size > 0 and K % group_size == 0
|
||||
G = group_size
|
||||
num_groups = K // G
|
||||
assert scales_gn.shape == (num_groups, N)
|
||||
|
||||
w_int4 = _unpack_int4_along_n(b_packed_kn8) # [K,N]
|
||||
if qzeros_gn8 is None:
|
||||
z_full = torch.full((K, N), zp_bias, dtype=torch.int32, device=a_mk.device)
|
||||
else:
|
||||
assert qzeros_gn8.shape == (num_groups, N // 8)
|
||||
z_gn = _unpack_int4_along_n(qzeros_gn8) # [G,N] in groups
|
||||
z_full = z_gn.repeat_interleave(G, dim=0) # [K,N]
|
||||
|
||||
s_full = scales_gn.repeat_interleave(G, dim=0).to(torch.float32) # [K,N]
|
||||
w_fp = (w_int4 - z_full).to(torch.float32) * s_full # [K,N]
|
||||
|
||||
out = a_mk.to(torch.float32) @ w_fp # [M,N]
|
||||
return out.to(a_mk.dtype)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
@pytest.mark.parametrize(
|
||||
"M,K,N,G,has_zp",
|
||||
[
|
||||
(1, 256, 256, 32, False),
|
||||
(17, 256, 512, 32, False),
|
||||
(32, 512, 256, 64, False),
|
||||
(33, 512, 512, 128, False),
|
||||
(64, 1024, 256, 256, False),
|
||||
(128, 256, 1024, 32, True),
|
||||
(64, 512, 512, 64, True),
|
||||
],
|
||||
)
|
||||
def test_triton_w4a16_gemm_matches_reference(dtype, M, K, N, G, has_zp):
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA/HIP device not available")
|
||||
if N % 8 != 0 or K % G != 0:
|
||||
pytest.skip("Invalid test shape")
|
||||
|
||||
set_random_seed(0)
|
||||
|
||||
a = (0.25 * torch.randn((M, K), device=device, dtype=torch.float32)).to(dtype)
|
||||
w_int4 = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
|
||||
b_packed = _pack_int4_along_n(w_int4)
|
||||
|
||||
scales = (0.05 * torch.rand((K // G, N), device=device, dtype=torch.float32)).to(
|
||||
dtype
|
||||
)
|
||||
|
||||
qzeros = None
|
||||
if has_zp:
|
||||
zeros_int4 = torch.randint(0, 16, (K // G, N), device=device, dtype=torch.int32)
|
||||
qzeros = _pack_int4_along_n(zeros_int4)
|
||||
|
||||
out = triton_w4a16_gemm(
|
||||
a=a,
|
||||
b_q=b_packed,
|
||||
scales=scales,
|
||||
qzeros=qzeros,
|
||||
group_size=G,
|
||||
zp_bias=8,
|
||||
)
|
||||
ref = _w4a16_reference(
|
||||
a,
|
||||
b_packed,
|
||||
scales,
|
||||
group_size=G,
|
||||
qzeros_gn8=qzeros,
|
||||
zp_bias=8,
|
||||
)
|
||||
|
||||
torch.testing.assert_close(out, ref, rtol=1e-2, atol=1e-2)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
|
||||
def test_triton_w4a16_gemm_requires_contiguous_inputs():
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA/HIP device not available")
|
||||
|
||||
set_random_seed(0)
|
||||
M, K, N, G = 32, 256, 256, 32
|
||||
a = torch.randn((K, M), device=device, dtype=torch.float16).t() # non-contiguous
|
||||
w_int4 = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
|
||||
b_packed = _pack_int4_along_n(w_int4)
|
||||
scales = torch.rand((K // G, N), device=device, dtype=torch.float16)
|
||||
|
||||
with pytest.raises(AssertionError):
|
||||
triton_w4a16_gemm(
|
||||
a=a,
|
||||
b_q=b_packed,
|
||||
scales=scales,
|
||||
qzeros=None,
|
||||
group_size=G,
|
||||
zp_bias=8,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
|
||||
def test_triton_w4a16_process_weights_after_loading_repacks_layout():
|
||||
if not torch.cuda.is_available():
|
||||
pytest.skip("CUDA/HIP device not available")
|
||||
|
||||
from vllm.config import VllmConfig, set_current_vllm_config
|
||||
from vllm.distributed import (
|
||||
ensure_model_parallel_initialized,
|
||||
init_distributed_environment,
|
||||
)
|
||||
from vllm.model_executor.kernels.linear.mixed_precision.MPLinearKernel import (
|
||||
MPLinearLayerConfig,
|
||||
)
|
||||
from vllm.model_executor.parameter import (
|
||||
GroupQuantScaleParameter,
|
||||
PackedColumnParameter,
|
||||
PackedvLLMParameter,
|
||||
)
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
with set_current_vllm_config(VllmConfig()):
|
||||
init_distributed_environment(
|
||||
world_size=1,
|
||||
rank=0,
|
||||
distributed_init_method="tcp://127.0.0.1:0",
|
||||
local_rank=0,
|
||||
)
|
||||
ensure_model_parallel_initialized(1, 1)
|
||||
|
||||
set_random_seed(0)
|
||||
|
||||
# Small-but-nontrivial shapes.
|
||||
K, N = 256, 256
|
||||
G = 32
|
||||
assert K % 8 == 0 and N % 8 == 0 and K % G == 0
|
||||
|
||||
# Build a canonical int4 weight grid then pack into the CT checkpoint layout.
|
||||
w_int4_kn = torch.randint(0, 16, (K, N), device=device, dtype=torch.int32)
|
||||
w_ckpt_nk8 = _pack_int4_along_k_to_ckpt(w_int4_kn) # [N, K//8]
|
||||
|
||||
# Scales in CT checkpoint layout for WNA16: [N, K//G]
|
||||
scales_ckpt_nkg = 0.05 * torch.rand((N, K // G), device=device, dtype=torch.float16)
|
||||
|
||||
# Asymmetric case: zero points in CT checkpoint layout [N//8, K//G] (N-packed)
|
||||
zeros_int4_gn = torch.randint(0, 16, (K // G, N), device=device, dtype=torch.int32)
|
||||
zeros_packed_gn8 = _pack_int4_along_n(zeros_int4_gn) # [K//G, N//8]
|
||||
zeros_ckpt_n8kg = zeros_packed_gn8.t().contiguous() # [N//8, K//G]
|
||||
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=(K, N),
|
||||
partition_weight_shape=(K, N),
|
||||
weight_type=scalar_types.uint4, # asymmetric
|
||||
act_type=torch.float16,
|
||||
group_size=G,
|
||||
zero_points=True,
|
||||
has_g_idx=False,
|
||||
)
|
||||
kernel = TritonW4A16LinearKernel(
|
||||
config,
|
||||
w_q_param_name="weight_packed",
|
||||
w_s_param_name="weight_scale",
|
||||
w_zp_param_name="weight_zero_point",
|
||||
w_gidx_param_name=None,
|
||||
)
|
||||
|
||||
# Build dummy layer with vLLM parameter wrappers.
|
||||
weight_loader = lambda *args, **kwargs: None
|
||||
|
||||
class DummyLayer(torch.nn.Module):
|
||||
pass
|
||||
|
||||
layer = DummyLayer()
|
||||
layer.register_parameter(
|
||||
"weight_packed",
|
||||
PackedvLLMParameter(
|
||||
data=w_ckpt_nk8,
|
||||
weight_loader=weight_loader,
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
packed_factor=8,
|
||||
packed_dim=1,
|
||||
),
|
||||
)
|
||||
layer.register_parameter(
|
||||
"weight_scale",
|
||||
GroupQuantScaleParameter(
|
||||
data=scales_ckpt_nkg,
|
||||
weight_loader=weight_loader,
|
||||
input_dim=1,
|
||||
output_dim=0,
|
||||
),
|
||||
)
|
||||
layer.register_parameter(
|
||||
"weight_zero_point",
|
||||
PackedColumnParameter(
|
||||
data=zeros_ckpt_n8kg,
|
||||
weight_loader=weight_loader,
|
||||
output_dim=0,
|
||||
packed_factor=8,
|
||||
packed_dim=0,
|
||||
),
|
||||
)
|
||||
|
||||
kernel.process_weights_after_loading(layer)
|
||||
|
||||
# Expected transformed layouts.
|
||||
expected_w_kn8 = _pack_int4_along_n(w_int4_kn) # [K, N//8]
|
||||
expected_scales_gn = scales_ckpt_nkg.t().contiguous() # [K//G, N]
|
||||
expected_zeros_gn8 = zeros_ckpt_n8kg.t().contiguous() # [K//G, N//8]
|
||||
|
||||
assert tuple(layer.weight_packed.shape) == (K, N // 8)
|
||||
assert tuple(layer.weight_scale.shape) == (K // G, N)
|
||||
assert tuple(layer.weight_zero_point.shape) == (K // G, N // 8)
|
||||
|
||||
torch.testing.assert_close(layer.weight_packed, expected_w_kn8)
|
||||
torch.testing.assert_close(layer.weight_scale, expected_scales_gn)
|
||||
torch.testing.assert_close(layer.weight_zero_point, expected_zeros_gn8)
|
||||
@@ -0,0 +1,53 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for W4A16 kernel selection logic (ROCm).
|
||||
|
||||
Run `pytest tests/kernels/quantization/test_w4a16_kernel_selection.py`.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.kernels.linear import (
|
||||
MPLinearLayerConfig,
|
||||
choose_mp_linear_kernel,
|
||||
)
|
||||
from vllm.platforms import current_platform
|
||||
from vllm.scalar_type import scalar_types
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
|
||||
def test_choose_mp_linear_kernel_picks_triton_w4a16_for_uint4b8():
|
||||
# int4 weights, 16-bit activations (CT W4A16 typical config).
|
||||
K, N = 1024, 256
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=(K, N),
|
||||
partition_weight_shape=(K, N),
|
||||
weight_type=scalar_types.uint4b8, # symmetric int4 (bias=8)
|
||||
act_type=torch.float16,
|
||||
group_size=128,
|
||||
zero_points=False,
|
||||
has_g_idx=False,
|
||||
)
|
||||
|
||||
kernel_type = choose_mp_linear_kernel(config)
|
||||
assert kernel_type.__name__ == "TritonW4A16LinearKernel"
|
||||
|
||||
|
||||
@pytest.mark.skipif(not current_platform.is_rocm(), reason="ROCm only")
|
||||
def test_choose_mp_linear_kernel_picks_triton_w4a16_for_uint4_asymmetric():
|
||||
# Asymmetric int4 weights should also be supported (explicit zero points).
|
||||
K, N = 512, 512
|
||||
config = MPLinearLayerConfig(
|
||||
full_weight_shape=(K, N),
|
||||
partition_weight_shape=(K, N),
|
||||
weight_type=scalar_types.uint4, # asymmetric int4 (explicit zeros)
|
||||
act_type=torch.bfloat16,
|
||||
group_size=64,
|
||||
zero_points=True,
|
||||
has_g_idx=False,
|
||||
)
|
||||
|
||||
kernel_type = choose_mp_linear_kernel(config)
|
||||
assert kernel_type.__name__ == "TritonW4A16LinearKernel"
|
||||
Reference in New Issue
Block a user