| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import pytest |
| 5 | import torch |
| 6 | import torch.nn.functional as F |
| 7 | |
| 8 | from vllm.model_executor.kernels.linear import ( |
| 9 | FlashInferTrtllmMxfp8LinearKernel, |
| 10 | Mxfp8LinearLayerConfig, |
| 11 | ) |
| 12 | from vllm.platforms import current_platform |
| 13 | from vllm.utils import flashinfer as vllm_flashinfer |
| 14 | from vllm.utils.flashinfer import has_flashinfer |
| 15 | |
| 16 | if not ( |
| 17 | current_platform.is_cuda() |
| 18 | and current_platform.is_device_capability_family(100) |
| 19 | and has_flashinfer() |
| 20 | ): |
| 21 | pytest.skip( |
| 22 | reason="FlashInfer TRTLLM MXFP8 requires an SM100-family GPU", |
| 23 | allow_module_level=True, |
| 24 | ) |
| 25 | |
| 26 | |
| 27 | def _make_layer(weight: torch.Tensor) -> torch.nn.Module: |
| 28 | from flashinfer import SfLayout, mxfp8_quantize |
| 29 | |
| 30 | weight_mxfp8, weight_scale = mxfp8_quantize( |
| 31 | weight, |
| 32 | sf_swizzle_layout=SfLayout.layout_linear, |
| 33 | ) |
| 34 | layer = torch.nn.Module() |
| 35 | layer.weight = torch.nn.Parameter(weight_mxfp8, requires_grad=False) |
| 36 | layer.weight_scale = torch.nn.Parameter( |
| 37 | weight_scale.view(weight.shape[0], weight.shape[1] // 32), |
| 38 | requires_grad=False, |
| 39 | ) |
| 40 | return layer |
| 41 | |
| 42 | |
| 43 | @pytest.mark.parametrize("shape", [(1, 130, 256), (7, 256, 512), (128, 130, 768)]) |
| 44 | @torch.inference_mode() |
| 45 | def test_flashinfer_trtllm_mxfp8_linear_numerics( |
| 46 | shape: tuple[int, int, int], |
| 47 | ) -> None: |
| 48 | torch.manual_seed(0) |
| 49 | m, n, k = shape |
| 50 | x = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") |
| 51 | weight = torch.randn((n, k), dtype=torch.bfloat16, device="cuda") |
| 52 | layer = _make_layer(weight) |
| 53 | kernel = FlashInferTrtllmMxfp8LinearKernel( |
| 54 | Mxfp8LinearLayerConfig(weight_shape=(n, k)) |
| 55 | ) |
| 56 | kernel.process_weights_after_loading(layer) |
| 57 | |
| 58 | output = kernel.apply_weights(layer, x) |
| 59 | reference = torch.mm(x, weight.t()) |
| 60 | similarity = F.cosine_similarity( |
| 61 | output.float().flatten(), reference.float().flatten(), dim=0 |
| 62 | ) |
| 63 | |
| 64 | assert output.shape == (m, n) |
| 65 | assert output.is_contiguous() |
| 66 | assert similarity.item() > 0.98 |
| 67 | |
| 68 | |
| 69 | @torch.inference_mode() |
| 70 | def test_flashinfer_trtllm_mxfp8_custom_ops() -> None: |
| 71 | x = torch.randn((7, 512), dtype=torch.bfloat16, device="cuda") |
| 72 | weight = torch.randn((256, 512), dtype=torch.bfloat16, device="cuda") |
| 73 | layer = _make_layer(weight) |
| 74 | kernel = FlashInferTrtllmMxfp8LinearKernel( |
| 75 | Mxfp8LinearLayerConfig(weight_shape=(256, 512)) |
| 76 | ) |
| 77 | kernel.process_weights_after_loading(layer) |
| 78 | |
| 79 | torch.library.opcheck( |
| 80 | torch.ops.vllm.flashinfer_mxfp8_quantize_8x4.default, |
| 81 | (x,), |
| 82 | ) |
| 83 | x_mxfp8, x_scale = vllm_flashinfer.flashinfer_mxfp8_quantize_8x4(x) |
| 84 | # SchemaCheckMode compares inputs with allclose, which CUDA does not |
| 85 | # implement for float8. The numerical tests above guard input mutation. |
| 86 | torch.library.opcheck( |
| 87 | torch.ops.vllm.mm_mxfp8.default, |
| 88 | ( |
| 89 | x_mxfp8, |
| 90 | layer.weight.t(), |
| 91 | x_scale, |
| 92 | layer.weight_scale, |
| 93 | torch.bfloat16, |
| 94 | "trtllm", |
| 95 | True, |
| 96 | ), |
| 97 | test_utils=( |
| 98 | "test_autograd_registration", |
| 99 | "test_faketensor", |
| 100 | "test_aot_dispatch_dynamic", |
| 101 | ), |
| 102 | ) |
| 103 | |
| 104 | |
| 105 | @torch.inference_mode() |
| 106 | def test_flashinfer_trtllm_mxfp8_linear_cuda_graph() -> None: |
| 107 | torch.manual_seed(0) |
| 108 | m, n, k = 7, 130, 512 |
| 109 | weight = torch.randn((n, k), dtype=torch.bfloat16, device="cuda") |
| 110 | layer = _make_layer(weight) |
| 111 | kernel = FlashInferTrtllmMxfp8LinearKernel( |
| 112 | Mxfp8LinearLayerConfig(weight_shape=(n, k)) |
| 113 | ) |
| 114 | kernel.process_weights_after_loading(layer) |
| 115 | |
| 116 | static_x = torch.randn((m, k), dtype=torch.bfloat16, device="cuda") |
| 117 | kernel.apply_weights(layer, static_x) |
| 118 | torch.accelerator.synchronize() |
| 119 | |
| 120 | graph = torch.cuda.CUDAGraph() |
| 121 | with torch.cuda.graph(graph): |
| 122 | graph_output = kernel.apply_weights(layer, static_x) |
| 123 | |
| 124 | new_x = torch.randn_like(static_x) |
| 125 | static_x.copy_(new_x) |
| 126 | graph.replay() |
| 127 | eager_output = kernel.apply_weights(layer, new_x) |
| 128 | |
| 129 | torch.testing.assert_close(graph_output, eager_output, rtol=0, atol=0) |