yangzhuxinyzx
234c56abb2
Assisted-by: OpenAI Codex Signed-off-by: yangzhuxinyzx <153831768+yangzhuxinyzx@users.noreply.github.com>
326 lines
11 KiB
Python
326 lines
11 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
from dataclasses import replace
|
|
from types import SimpleNamespace as NS
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
from vllm import envs
|
|
from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config
|
|
from vllm.model_executor.kernels import linear
|
|
from vllm.model_executor.kernels.linear import pre_ampere_qpn as qpn
|
|
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
|
kFp8Static128BlockSym,
|
|
kFp8StaticTensorSym,
|
|
)
|
|
from vllm.platforms import PlatformEnum
|
|
from vllm.platforms.interface import DeviceCapability
|
|
|
|
pytestmark = pytest.mark.cpu_test
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def engine(monkeypatch):
|
|
for name in (
|
|
"VLLM_SM70_QUANT_BACKEND",
|
|
"VLLM_LINEAR_BACKEND",
|
|
"VLLM_SM70_FP8_TURBOMIND",
|
|
"VLLM_SM70_NVFP4_TURBOMIND",
|
|
"VLLM_DISABLED_KERNELS",
|
|
"VLLM_NVFP4_GEMM_BACKEND",
|
|
):
|
|
monkeypatch.delenv(name, raising=False)
|
|
envs.disable_envs_cache()
|
|
platform = NS(
|
|
_enum=PlatformEnum.CUDA,
|
|
is_cuda=lambda: True,
|
|
get_device_capability=lambda *args: DeviceCapability(7, 5),
|
|
)
|
|
monkeypatch.setattr(linear, "current_platform", platform)
|
|
monkeypatch.setattr(qpn, "current_platform", platform)
|
|
monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 2)
|
|
monkeypatch.setattr(
|
|
torch.ops,
|
|
"_C",
|
|
NS(
|
|
nvfp4_qpn2_prepare_sm70=True,
|
|
nvfp4_qpn2_gemm_sm70_out=True,
|
|
fp8_qpn8_prepare_sm70=True,
|
|
fp8_qpn8_dispatch_sm70_out=True,
|
|
fp8_qpn8_prefill_sm70_out=True,
|
|
),
|
|
)
|
|
config = VllmConfig(device_config=DeviceConfig(device="cpu"))
|
|
with set_current_vllm_config(config):
|
|
yield config
|
|
envs.disable_envs_cache()
|
|
|
|
|
|
def nvfp4_config(n=40, k=384):
|
|
return qpn.TuringNvFp4LinearLayerConfig(
|
|
input_dtype=torch.float16,
|
|
weight_shape=(n, k // 2),
|
|
scale_shape=(n, k // 16),
|
|
weight_dtype=torch.uint8,
|
|
scale_dtype=torch.float8_e4m3fn,
|
|
)
|
|
|
|
|
|
def fp8_config(n=96, k=384):
|
|
return linear.FP8ScaledMMLinearLayerConfig(
|
|
weight_quant_key=kFp8StaticTensorSym,
|
|
activation_quant_key=kFp8StaticTensorSym,
|
|
weight_shape=(n, k),
|
|
input_dtype=torch.float16,
|
|
out_dtype=torch.float16,
|
|
)
|
|
|
|
|
|
def test_real_selectors_prefer_native_turing_kernels():
|
|
assert isinstance(
|
|
linear.init_nvfp4_linear_kernel(nvfp4_config()),
|
|
qpn.TuringQpn2NvFp4LinearKernel,
|
|
)
|
|
|
|
|
|
def test_modelopt_selects_with_metadata_for_explicit_provider(engine, monkeypatch):
|
|
from vllm.model_executor import parameter
|
|
from vllm.model_executor.layers.quantization import modelopt
|
|
|
|
engine.kernel_config.linear_backend = "turbomind"
|
|
engine.model_config = NS(dtype=torch.float16)
|
|
monkeypatch.setattr(modelopt.sm70_tm, "is_turing_cuda_platform", lambda: True)
|
|
monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0)
|
|
monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1)
|
|
method = modelopt.ModelOptNvFp4LinearMethod(
|
|
NS(is_checkpoint_nvfp4_serialized=True, group_size=16)
|
|
)
|
|
assert method.kernel is None
|
|
layer = torch.nn.Module()
|
|
method.create_weights(layer, 128, [40], 128, 40, torch.float16)
|
|
assert isinstance(method.kernel, qpn.TuringQpn2NvFp4LinearKernel)
|
|
|
|
|
|
def test_fp8_real_selector_prefers_native_turing_kernel():
|
|
assert (
|
|
linear.choose_scaled_mm_linear_kernel(
|
|
fp8_config(),
|
|
linear._POSSIBLE_FP8_KERNELS,
|
|
compute_capability=75,
|
|
)
|
|
is qpn.TuringQpn8Fp8LinearKernel
|
|
)
|
|
assert isinstance(
|
|
linear.init_fp8_linear_kernel(
|
|
activation_quant_key=kFp8StaticTensorSym,
|
|
weight_quant_key=kFp8StaticTensorSym,
|
|
weight_shape=(96, 384),
|
|
input_dtype=torch.float16,
|
|
out_dtype=torch.float16,
|
|
),
|
|
qpn.TuringQpn8Fp8LinearKernel,
|
|
)
|
|
|
|
|
|
def test_modelopt_fp8_preserves_checkpoint_orientation(engine, monkeypatch):
|
|
from vllm.model_executor import parameter
|
|
from vllm.model_executor.layers.quantization import modelopt
|
|
|
|
engine.model_config = NS(dtype=torch.float16)
|
|
monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0)
|
|
monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1)
|
|
monkeypatch.setattr(modelopt.sm70_tm, "is_exact_sm70_cuda_platform", lambda: False)
|
|
monkeypatch.setattr(torch, "get_default_dtype", lambda: torch.float16)
|
|
captured = []
|
|
monkeypatch.setattr(
|
|
qpn.tm,
|
|
"prepare_fp8_qpn8_dense_linear",
|
|
lambda layer, weight, scale: captured.append(weight.clone()),
|
|
)
|
|
method = modelopt.ModelOptFp8LinearMethod(NS(is_checkpoint_fp8_serialized=True))
|
|
layer = torch.nn.Module()
|
|
method.create_weights(layer, 128, [32], 128, 32, torch.float16)
|
|
checkpoint = (
|
|
torch.arange(32 * 128).reshape(32, 128).remainder(8).to(torch.float8_e4m3fn)
|
|
)
|
|
layer.weight.data.copy_(checkpoint)
|
|
layer.weight_scale.data.fill_(0.5)
|
|
layer.input_scale.data.fill_(1)
|
|
method.process_weights_after_loading(layer)
|
|
assert isinstance(method.fp8_linear, qpn.TuringQpn8Fp8LinearKernel)
|
|
assert len(captured) == 1
|
|
assert torch.equal(captured[0].float(), checkpoint.float())
|
|
assert layer.weight.numel() == 0
|
|
|
|
|
|
def test_modelopt_fp8_marlin_fallback_prepares_scales(engine, monkeypatch):
|
|
from vllm.model_executor import parameter
|
|
from vllm.model_executor.kernels.linear.scaled_mm.marlin import (
|
|
MarlinFP8ScaledMMLinearKernel,
|
|
)
|
|
from vllm.model_executor.layers.quantization import modelopt
|
|
from vllm.model_executor.layers.quantization.utils import marlin_utils_fp8
|
|
|
|
engine.model_config = NS(dtype=torch.float16)
|
|
monkeypatch.setattr(parameter, "get_tensor_model_parallel_rank", lambda: 0)
|
|
monkeypatch.setattr(parameter, "get_tensor_model_parallel_world_size", lambda: 1)
|
|
monkeypatch.setattr(modelopt.sm70_tm, "is_exact_sm70_cuda_platform", lambda: False)
|
|
monkeypatch.setattr(torch, "get_default_dtype", lambda: torch.float16)
|
|
kernel = object.__new__(MarlinFP8ScaledMMLinearKernel)
|
|
kernel.layer_param_names = [
|
|
"weight",
|
|
"weight_scale",
|
|
"input_scale",
|
|
"input_scale_ub",
|
|
]
|
|
kernel.block_quant = False
|
|
kernel.size_k_first = True
|
|
kernel.marlin_input_dtype = None
|
|
monkeypatch.setattr(modelopt, "init_fp8_linear_kernel", lambda **kwargs: kernel)
|
|
monkeypatch.setattr(
|
|
marlin_utils_fp8,
|
|
"marlin_make_workspace_new",
|
|
lambda device: torch.zeros(1, dtype=torch.int32, device=device),
|
|
)
|
|
monkeypatch.setattr(
|
|
marlin_utils_fp8.ops,
|
|
"gptq_marlin_repack",
|
|
lambda **kwargs: kwargs["b_q_weight"].clone(),
|
|
)
|
|
method = modelopt.ModelOptFp8LinearMethod(NS(is_checkpoint_fp8_serialized=True))
|
|
layer = torch.nn.Module()
|
|
method.create_weights(layer, 128, [64], 128, 64, torch.float16)
|
|
layer.weight.data.fill_(1)
|
|
layer.weight_scale.data.fill_(0.5)
|
|
layer.input_scale.data.fill_(1)
|
|
method.process_weights_after_loading(layer)
|
|
assert layer.weight_scale.dtype == torch.float16
|
|
assert layer.weight_scale.numel() == 64
|
|
|
|
|
|
@pytest.mark.parametrize("n,k", [(1, 128), (40, 384), (97, 768), (1024, 2048)])
|
|
def test_nvfp4_admits_padding_and_aligned_shapes(n, k):
|
|
assert qpn.TuringQpn2NvFp4LinearKernel.can_implement(nvfp4_config(n, k))[0]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"change,reason",
|
|
[
|
|
({"input_dtype": torch.bfloat16}, "FP16"),
|
|
({"weight_shape": (40, 24)}, "128"),
|
|
({"scale_shape": (40, 1)}, "16 weight"),
|
|
({"scale_dtype": torch.float32}, "E4M3"),
|
|
({"weight_shape": (40, 24, 1)}, "rank"),
|
|
],
|
|
)
|
|
def test_nvfp4_rejections_explain_fallback(change, reason):
|
|
accepted, why = qpn.TuringQpn2NvFp4LinearKernel.can_implement(
|
|
replace(nvfp4_config(), **change)
|
|
)
|
|
assert not accepted and reason in why
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"change,reason",
|
|
[
|
|
({"out_dtype": torch.bfloat16}, "FP16"),
|
|
({"weight_shape": (31, 128)}, "32"),
|
|
({"weight_shape": (32, 16)}, "128"),
|
|
({"weight_quant_key": kFp8Static128BlockSym}, "per-tensor"),
|
|
],
|
|
)
|
|
def test_fp8_rejections_explain_fallback(change, reason):
|
|
accepted, why = qpn.TuringQpn8Fp8LinearKernel.can_implement(
|
|
replace(fp8_config(), **change)
|
|
)
|
|
assert not accepted and reason in why
|
|
|
|
|
|
def test_kernelconfig_and_legacy_disable(engine, monkeypatch):
|
|
engine.kernel_config.sm70_nvfp4.dense_qpn2 = False
|
|
assert (
|
|
"KernelConfig"
|
|
in qpn.TuringQpn2NvFp4LinearKernel.can_implement(nvfp4_config())[1]
|
|
)
|
|
engine.kernel_config.sm70_fp8.enabled = False
|
|
assert not qpn.TuringQpn8Fp8LinearKernel.can_implement(fp8_config())[0]
|
|
engine.kernel_config.sm70_nvfp4.dense_qpn2 = True
|
|
monkeypatch.setenv("VLLM_SM70_NVFP4_TURBOMIND", "0")
|
|
assert qpn.TuringQpn2NvFp4LinearKernel.can_implement(nvfp4_config())[0]
|
|
engine.kernel_config.sm70_nvfp4.enabled = False
|
|
assert (
|
|
"override" in qpn.TuringQpn2NvFp4LinearKernel.can_implement(nvfp4_config())[1]
|
|
)
|
|
|
|
|
|
def test_native_missing_and_worker_device(monkeypatch):
|
|
asked = []
|
|
|
|
def worker_capability(device):
|
|
asked.append(device)
|
|
return DeviceCapability(7, 5)
|
|
|
|
monkeypatch.setattr(
|
|
qpn.current_platform,
|
|
"get_device_capability",
|
|
worker_capability,
|
|
)
|
|
|
|
assert qpn.TuringQpn2NvFp4LinearKernel.is_supported()[0]
|
|
assert asked == [2]
|
|
assert not qpn.TuringQpn2NvFp4LinearKernel.is_supported(70)[0]
|
|
assert not qpn.TuringQpn8Fp8LinearKernel.is_supported(72)[0]
|
|
monkeypatch.setattr(torch.ops, "_C", NS())
|
|
assert (
|
|
"missing native" in qpn.TuringQpn8Fp8LinearKernel.can_implement(fp8_config())[1]
|
|
)
|
|
|
|
|
|
def test_nvfp4_fallback_is_recorded_for_startup(engine, monkeypatch):
|
|
engine.kernel_config.sm70_nvfp4.dense_qpn2 = False
|
|
monkeypatch.setattr(
|
|
linear.MarlinNvFp4LinearKernel,
|
|
"is_supported",
|
|
classmethod(lambda cls, *args: (True, None)),
|
|
)
|
|
selected = linear.init_nvfp4_linear_kernel(nvfp4_config())
|
|
assert isinstance(selected, linear.MarlinNvFp4LinearKernel)
|
|
rows = engine.kernel_config.linear_kernel_selections.values()
|
|
assert any(
|
|
"KernelConfig" in row["paths"]["TuringQpn2NvFp4LinearKernel"]["reason"]
|
|
for row in rows
|
|
)
|
|
|
|
|
|
def test_fp8_scratch_is_owned_by_each_invocation(monkeypatch):
|
|
from vllm import _sm70_ops as ops
|
|
|
|
allocations = []
|
|
dispatched = []
|
|
original = torch.Tensor.new_empty
|
|
|
|
def allocate(tensor, shape, **kwargs):
|
|
result = original(tensor, shape, **kwargs)
|
|
allocations.append(result)
|
|
return result
|
|
|
|
monkeypatch.setattr(torch.Tensor, "new_empty", allocate)
|
|
monkeypatch.setattr(
|
|
ops,
|
|
"fp8_qpn8_dispatch_sm70_out",
|
|
lambda out, ptr, *args: dispatched.append(ptr),
|
|
)
|
|
x = torch.empty(9, 128, dtype=torch.float16)
|
|
for _ in range(2):
|
|
qpn._turing_fp8_qpn8_linear(x, torch.empty(0), torch.empty(0), 32, 8, 2, False)
|
|
workspaces = [tensor for tensor in allocations if tensor.shape == (128, 32)]
|
|
assert len(workspaces) == 2
|
|
assert dispatched == [tensor.data_ptr() for tensor in workspaces]
|
|
assert dispatched[0] != dispatched[1]
|
|
dispatched.clear()
|
|
assert qpn._turing_fp8_qpn8_linear(
|
|
x[:0], torch.empty(0), torch.empty(0), 32, 8, 2, False
|
|
).shape == (0, 32)
|
|
assert dispatched == []
|