wehub-resource-sync
864 lines
31 KiB
Python
864 lines
31 KiB
Python
# 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)
|