169 lines
7.9 KiB
Python
169 lines
7.9 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
|
"""Real A5 kernels: step selection, bias, shape, and reuse of prepared W4 weights."""
|
|
|
|
import pytest
|
|
import torch
|
|
|
|
pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion, pytest.mark.npu]
|
|
|
|
torch_npu = pytest.importorskip("torch_npu")
|
|
|
|
|
|
@torch.inference_mode()
|
|
def _exercise_mxfp4_fallback(with_bias, dtype, input_shape, mode, monkeypatch, *, policy, scale_alg=2):
|
|
from vllm_omni.diffusion.forward_context import (
|
|
ForwardContext,
|
|
override_forward_context,
|
|
set_forward_context_denoise_step_idx,
|
|
)
|
|
from vllm_omni.quantization.mxfp4_config import (
|
|
DiffusionMXFP4Config,
|
|
NPUMxfp4LinearMethod,
|
|
NPUMxfp4OnlineLinearMethod,
|
|
)
|
|
|
|
prefix = "blocks.10.attn1.to_qkv"
|
|
steps = [1] if policy in ("step", "layer_and_step") else []
|
|
layers = [prefix] if policy in ("layer", "layer_and_step") else []
|
|
fallback = {"w4a8_fallback_steps": steps, "w4a8_fallback_layers": layers, "mxfp4_scale_alg": scale_alg}
|
|
torch.manual_seed(42)
|
|
layer = torch.nn.Linear(1024, 128, bias=with_bias, device="npu", dtype=dtype)
|
|
if mode == "single":
|
|
method = NPUMxfp4OnlineLinearMethod(DiffusionMXFP4Config(**fallback), prefix=prefix)
|
|
elif mode == "single_offline":
|
|
method = NPUMxfp4LinearMethod(
|
|
DiffusionMXFP4Config(is_checkpoint_mxfp4_serialized=True, **fallback), prefix=prefix
|
|
)
|
|
layer.orig_dtype = dtype
|
|
values = torch.tensor([-4.0, -2.0, -1.0, -0.5, 0.0, 0.5, 1.0, 2.0])
|
|
layer.weight.copy_(values[torch.randint(8, (128, 1024))].to(layer.weight))
|
|
layer.register_parameter(
|
|
"weight_scale", torch.nn.Parameter(torch.full((128, 32), 124, dtype=torch.uint8, device="npu"), False)
|
|
)
|
|
layer.register_parameter("mul_scale", torch.nn.Parameter((torch.rand(1024) * 1.5 + 0.5).npu(), False))
|
|
else:
|
|
raise ValueError(f"Unsupported single-scale test mode: {mode}")
|
|
method.process_weights_after_loading(layer)
|
|
x = torch.randn(input_shape, device="npu", dtype=dtype)
|
|
pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
|
|
method.process_weights_after_loading(layer)
|
|
quantize = torch_npu.npu_dynamic_mx_quant
|
|
activation_types = []
|
|
|
|
def record_quantize(x, **kwargs):
|
|
dst = kwargs["dst_type"]
|
|
expected_alg = 0 if dst == torch_npu.float8_e4m3fn else scale_alg
|
|
assert kwargs["scale_alg"] == expected_alg
|
|
assert kwargs["axis"] == -1 and kwargs["block_size"] == 32 and kwargs["round_mode"] == "rint"
|
|
if expected_alg == 2:
|
|
assert kwargs["dst_type_max"] == 7.25
|
|
else:
|
|
assert "dst_type_max" not in kwargs
|
|
activation_types.append(dst)
|
|
return quantize(x, **kwargs)
|
|
|
|
monkeypatch.setattr(torch_npu, "npu_dynamic_mx_quant", record_quantize)
|
|
results = []
|
|
for _request in range(2):
|
|
with override_forward_context(ForwardContext()):
|
|
for step in range(3):
|
|
set_forward_context_denoise_step_idx(step)
|
|
output = method.apply(layer, x, layer.bias)
|
|
torch.npu.synchronize()
|
|
assert output.shape == (*input_shape[:-1], 128)
|
|
assert output.dtype == dtype
|
|
assert torch.isfinite(output).all()
|
|
results.append(output.cpu())
|
|
expected_types = [
|
|
torch_npu.float8_e4m3fn if layers or step in steps else torch_npu.float4_e2m1fn_x2 for step in range(3)
|
|
]
|
|
assert activation_types == expected_types * 2
|
|
assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == pointers
|
|
assert not hasattr(layer, "w4a8_weight")
|
|
assert not hasattr(layer, "w4a8_weight_scale")
|
|
for first, second in zip(results[:3], results[3:]):
|
|
torch.testing.assert_close(first, second, rtol=0, atol=0)
|
|
|
|
|
|
@pytest.mark.parametrize("with_bias", [False, True])
|
|
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
|
@pytest.mark.parametrize("input_shape", [(16, 1024), (2, 8, 1024)])
|
|
@pytest.mark.parametrize("mode", ["single", "single_offline"])
|
|
def test_mxfp4_step_fallback(with_bias, dtype, input_shape, mode, monkeypatch):
|
|
_exercise_mxfp4_fallback(with_bias, dtype, input_shape, mode, monkeypatch, policy="step")
|
|
|
|
|
|
@pytest.mark.cards_1
|
|
@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
|
|
@torch.inference_mode()
|
|
def test_single_scale_step_fallback_matches_independent_dense_reference(dtype):
|
|
if not torch.npu.is_available() or "Ascend950" not in torch.npu.get_device_name():
|
|
pytest.skip("MXFP4 kernels require an Ascend950 NPU")
|
|
|
|
from vllm_omni.diffusion.forward_context import (
|
|
ForwardContext,
|
|
override_forward_context,
|
|
set_forward_context_denoise_step_idx,
|
|
)
|
|
from vllm_omni.quantization.mxfp4_config import (
|
|
DiffusionMXFP4Config,
|
|
NPUMxfp4LinearMethod,
|
|
)
|
|
|
|
# Synthetic serialized single-scale values and E8M0 exponent bytes.
|
|
# Powers of two make both activation precisions exact: this independent
|
|
# CPU dense oracle needs no production packing, dequantization or apply.
|
|
# Real UOS-quantized checkpoint validation remains a separate checkpoint/E2E gate.
|
|
values = torch.tensor([1.0, -1.0, 2.0, -2.0]).repeat(32).unsqueeze(1).expand(128, 512)
|
|
weight = torch.cat((values, values), dim=1)
|
|
fine = torch.cat((torch.full((128, 16), 127), torch.full((128, 16), 128)), dim=1).to(torch.uint8)
|
|
smooth = torch.cat((torch.full((512,), 0.5), torch.full((512,), 2.0)))
|
|
bias = torch.tensor([0.25, -0.25]).repeat(64)
|
|
x = torch.tensor([1.0, -1.0, 0.5, -0.5]).repeat(4).unsqueeze(1).expand(16, 1024) / 64
|
|
expected = torch.nn.functional.linear(x * smooth, torch.cat((values, 2 * values), dim=1), bias)
|
|
|
|
layer = torch.nn.Linear(1024, 128, bias=True, device="npu", dtype=dtype)
|
|
layer.weight.copy_(weight)
|
|
layer.bias.copy_(bias)
|
|
layer.orig_dtype = dtype
|
|
for name, value in {"weight_scale": fine, "mul_scale": smooth}.items():
|
|
layer.register_parameter(name, torch.nn.Parameter(value.npu(), requires_grad=False))
|
|
method = NPUMxfp4LinearMethod(
|
|
DiffusionMXFP4Config(
|
|
is_checkpoint_mxfp4_serialized=True, require_smooth_scale=True, w4a8_fallback_steps=[1], mxfp4_scale_alg=2
|
|
)
|
|
)
|
|
method.process_weights_after_loading(layer)
|
|
prepared_pointers = (layer.weight.data_ptr(), layer.weight_scale.data_ptr())
|
|
assert set(layer.state_dict()) == {"weight", "bias", "weight_scale", "mul_scale"}
|
|
assert not hasattr(layer, "w4a8_weight")
|
|
assert not hasattr(layer, "w4a8_weight_scale")
|
|
x = x.to(device="npu", dtype=dtype)
|
|
|
|
with override_forward_context(ForwardContext()):
|
|
for step in range(3):
|
|
set_forward_context_denoise_step_idx(step)
|
|
actual = method.apply(layer, x, layer.bias).cpu()
|
|
# A8's documented GEMM output is BF16 before restoring the caller
|
|
# dtype; the reference does not reuse quantization or unpacking.
|
|
reference = expected.bfloat16().to(dtype) if step == 1 else expected.to(dtype)
|
|
torch.testing.assert_close(actual, reference, rtol=0, atol=0)
|
|
assert actual.dtype == dtype
|
|
assert (layer.weight.data_ptr(), layer.weight_scale.data_ptr()) == prepared_pointers
|
|
|
|
|
|
@pytest.mark.cards_1
|
|
@pytest.mark.parametrize("mode", ["single", "single_offline"])
|
|
@pytest.mark.parametrize("policy", ["layer", "layer_and_step"])
|
|
def test_mxfp4_layer_fallback(mode, policy, monkeypatch):
|
|
"""Reuse the kernel/shape/cache assertions without multiplying the full dtype/bias matrix."""
|
|
_exercise_mxfp4_fallback(False, torch.bfloat16, (16, 1024), mode, monkeypatch, policy=policy)
|
|
|
|
|
|
@pytest.mark.cards_1
|
|
@pytest.mark.parametrize("mode", ["single", "single_offline"])
|
|
def test_mxfp4_ocp_compatibility(mode, monkeypatch):
|
|
"""OCP compatibility: A4 -> A8 -> A4 with the same W4 across requests."""
|
|
_exercise_mxfp4_fallback(False, torch.bfloat16, (16, 1024), mode, monkeypatch, policy="step", scale_alg=0)
|