test_flashinfer_moe.py 5.31 KB
Newer Older
1
2
3
4
5
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import pytest
import torch

6
from tests.kernels.moe.utils import make_test_quant_config
7
8
9
10
11
from tests.kernels.quantization.nvfp4_utils import (
    FLOAT4_E2M1_MAX,
    FLOAT8_E4M3_MAX,
    dequantize_nvfp4_to_dtype,
)
12
13
14
from tests.kernels.utils import torch_moe
from vllm import _custom_ops as ops
from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
15
from vllm.model_executor.layers.fused_moe import fused_topk
16
17
18
19
20
from vllm.model_executor.layers.fused_moe.config import (
    FusedMoEConfig,
    FusedMoEParallelConfig,
    RoutingMethodType,
)
21
from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import (
22
23
24
25
    FlashInferExperts,
    is_valid_flashinfer_cutlass_fused_moe,
)
from vllm.model_executor.layers.fused_moe.modular_kernel import FusedMoEModularKernel
26
27
28
from vllm.model_executor.layers.fused_moe.prepare_finalize import (
    MoEPrepareAndFinalizeNoEP,
)
29
30
from vllm.platforms import current_platform
from vllm.utils.flashinfer import has_flashinfer_cutlass_fused_moe
31
from vllm.utils.torch_utils import set_random_seed
32

33
34
35
36
37
38
39
if not has_flashinfer_cutlass_fused_moe() or not current_platform.has_device_capability(
    100
):
    pytest.skip(
        "Requires flashinfer_cutlass_fused_moe and nvfp4 support",
        allow_module_level=True,
    )
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55

MNK_FACTORS = [
    (2, 1024, 1024),
    (2, 3072, 1024),
    (2, 3072, 1536),
    (64, 1024, 1536),
    (64, 3072, 1024),
    (64, 2048, 1536),
    (224, 1024, 1024),
    (224, 1024, 1536),
]


@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
@pytest.mark.parametrize("e", [40, 64, 256])
@pytest.mark.parametrize("topk", [1, 6, 8])
56
@pytest.mark.parametrize("dtype", [torch.bfloat16])
57
@pytest.mark.parametrize("activation", ["silu_and_mul", "relu2"])
58
@torch.inference_mode()
59
def test_flashinfer_fp4_moe_no_graph(
60
61
62
63
64
65
66
67
    m: int,
    n: int,
    k: int,
    e: int,
    topk: int,
    dtype: torch.dtype,
    activation: str,
    workspace_init,
68
):
69
    set_random_seed(7)
70
    with set_current_vllm_config(
71
72
        VllmConfig(parallel_config=ParallelConfig(pipeline_parallel_size=1))
    ):
73
74
75
        a = torch.randn((m, k), device="cuda", dtype=dtype) / 10

        quant_blocksize = 16
76
        is_gated_act = activation == "silu_and_mul"
77

78
79
80
81
82
83
84
85
        w1_q, w2_q, quant_config = make_test_quant_config(
            e,
            n,
            k,
            in_dtype=dtype,
            quant_dtype="nvfp4",
            block_shape=None,
            per_act_token_quant=False,
86
            make_gate=is_gated_act,
87
        )
88
89

        score = torch.randn((m, e), device="cuda", dtype=dtype)
90
        topk_weights, topk_ids, _ = fused_topk(a, score, topk, renormalize=False)
91
92
93

        assert is_valid_flashinfer_cutlass_fused_moe(a, w1_q, w2_q)

94
95
96
97
98
99
100
101
102
103
104
105
106
107
        moe_config = FusedMoEConfig(
            num_experts=e,
            experts_per_token=topk,
            hidden_dim=k,
            intermediate_size_per_partition=n,
            num_local_experts=e,
            activation=activation,
            device="cuda",
            moe_parallel_config=FusedMoEParallelConfig.make_no_parallel(),
            in_dtype=dtype,
            is_act_and_mul=is_gated_act,
            routing_method=RoutingMethodType.TopK,
        )

108
        flashinfer_experts = FusedMoEModularKernel(
109
            MoEPrepareAndFinalizeNoEP(),
110
            FlashInferExperts(moe_config=moe_config, quant_config=quant_config),
111
        )
112

113
114
        fi_activation = {"silu_and_mul": "silu", "relu2": "relu2_no_mul"}[activation]

115
116
117
118
119
120
        flashinfer_output = flashinfer_experts(
            hidden_states=a,
            w1=w1_q,
            w2=w2_q,
            topk_weights=topk_weights,
            topk_ids=topk_ids,
121
            activation=fi_activation,
122
123
124
        )

        # Reference check:
125
126
127
        a_global_scale = (
            (FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) / torch.amax(a.flatten(), dim=-1)
        ).to(torch.float32)
128
129
        a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale)
        _, m_k = a_fp4.shape
130
131
132
133
134
135
136
137
        a_in_dtype = dequantize_nvfp4_to_dtype(
            a_fp4,
            a_scale_interleaved,
            a_global_scale,
            dtype=a.dtype,
            device=a.device,
            block_size=quant_blocksize,
        )
138

139
140
141
        w1_d = torch.empty(
            (e, (2 if is_gated_act else 1) * n, k), device="cuda", dtype=dtype
        )
142
143
144
        w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)

        for idx in range(0, e):
145
146
            w1_d[idx] = dequantize_nvfp4_to_dtype(
                w1_q[idx],
147
148
                quant_config.w1_scale[idx],
                (1 / quant_config.g1_alphas[idx]),
149
150
                dtype=dtype,
                device=w1_q.device,
151
152
                block_size=quant_blocksize,
            )
153
154
            w2_d[idx] = dequantize_nvfp4_to_dtype(
                w2_q[idx],
155
156
                quant_config.w2_scale[idx],
                (1 / quant_config.g2_alphas[idx]),
157
158
                dtype=dtype,
                device=w2_q.device,
159
160
                block_size=quant_blocksize,
            )
161

162
163
164
        torch_output = torch_moe(
            a_in_dtype, w1_d, w2_d, score, topk, activation=activation
        )
165

166
167
168
        torch.testing.assert_close(
            torch_output, flashinfer_output, atol=1e-1, rtol=1e-1
        )
169
170
171
172


if __name__ == "__main__":
    test_flashinfer_fp4_moe_no_graph((2, 1024, 1024), 40, 1, torch.half)