"tests/distributed/test_eplb_utils.py" did not exist on "6287e7fa208199da2dc2869646291a902c0d5caa"
test_fp8_quant.py 3.88 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
5
6
7
import pytest
import torch

import vllm._custom_ops as ops
8
9
10
11
12
from tests.kernels.quant_utils import (
    FP8_DTYPE,
    ref_dynamic_per_tensor_fp8_quant,
    ref_dynamic_per_token_quant,
)
13
from tests.kernels.utils import opcheck
14
from vllm.utils.torch_utils import set_random_seed
15

16
17
18
DTYPES = [torch.bfloat16, torch.float]
HIDDEN_SIZES = [17, 1024, 1025, 1026, 5137, 8193]
NUM_TOKENS = [1, 7, 4096]
19
SCALE_UBS = [True, False]
20
21
22
SEEDS = [0]


23
24
25
def opcheck_fp8_quant(
    output, input, scale=None, scale_ub=None, use_per_token_if_dynamic=False
):
26
27
28
    if scale is not None:
        opcheck(torch.ops._C.static_scaled_fp8_quant, (output, input, scale))
    elif use_per_token_if_dynamic:
29
30
31
32
33
34
35
        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),
        )
36
    else:
37
38
39
40
41
        scale = torch.empty(
            (input.numel() // input.shape[-1], 1),
            device=input.device,
            dtype=torch.float32,
        )
42
43
44
        opcheck(torch.ops._C.dynamic_scaled_fp8_quant, (output, input, scale))


45
46
47
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
@pytest.mark.parametrize("dtype", DTYPES)
48
@pytest.mark.parametrize("scale_ub", SCALE_UBS)
49
50
@pytest.mark.parametrize("seed", SEEDS)
@torch.inference_mode()
51
52
53
def test_dynamic_per_token_fp8_quant(
    num_tokens: int, hidden_size: int, dtype: torch.dtype, scale_ub: bool, seed: int
) -> None:
54
    set_random_seed(seed)
55

56
57
58
    x = (
        torch.rand(num_tokens, hidden_size, dtype=dtype, device="cuda") + 1e-6
    )  # avoid nans
59

60
61
62
    scale_ub = (
        torch.mean(x).to(dtype=torch.float32, device="cuda") if scale_ub else None
    )
63
    ref_out, ref_scales = ref_dynamic_per_token_quant(x, FP8_DTYPE, scale_ub)
64
65
66
    ops_out, ops_scales = ops.scaled_fp8_quant(
        x, scale_ub=scale_ub, use_per_token_if_dynamic=True
    )
67

68
    torch.testing.assert_close(ref_scales, ops_scales)
69
70
71
    torch.testing.assert_close(
        ref_out.to(dtype=torch.float32), ops_out.to(dtype=torch.float32)
    )
72

73
    opcheck_fp8_quant(ops_out, x, None, scale_ub, use_per_token_if_dynamic=True)
74

75
76
77
78
79
80

@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()
81
82
83
def test_dynamic_per_tensor_fp8_quant(
    num_tokens: int, hidden_size: int, dtype: torch.dtype, seed: int
) -> None:
84
    set_random_seed(seed)
85
86
87
88
89
90

    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)

91
    torch.testing.assert_close(ref_scale, ops_scale)
92
93
94
    torch.testing.assert_close(
        ref_out.to(dtype=torch.float32), ops_out.to(dtype=torch.float32)
    )
95

96
97
    opcheck_fp8_quant(ops_out, x)

98
99
100
101
102
103

# 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:
104
    set_random_seed(seed)
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119

    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)

120
    torch.testing.assert_close(ref_out, ops_out)