| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from types import SimpleNamespace |
| 5 | |
| 6 | import pytest |
| 7 | import torch |
| 8 | |
| 9 | from vllm.model_executor.kernels.linear.scaled_mm.deep_gemm import ( |
| 10 | DeepGemmFp8BlockScaledMMKernel, |
| 11 | ) |
| 12 | from vllm.model_executor.warmup import deep_gemm_warmup |
| 13 | |
| 14 | |
| 15 | def _block_fp8_layer(n: int, k: int, scale_name: str = "weight_scale"): |
| 16 | layer = SimpleNamespace( |
| 17 | weight=torch.empty((n, k), dtype=torch.float8_e4m3fn), |
| 18 | weight_block_size=[128, 128], |
| 19 | deep_gemm_warmup_provider=object.__new__(DeepGemmFp8BlockScaledMMKernel), |
| 20 | ) |
| 21 | setattr(layer, scale_name, torch.empty((n // 128, k // 128), dtype=torch.float32)) |
| 22 | return layer |
| 23 | |
| 24 | |
| 25 | def _run_warmup(monkeypatch, layers) -> list[dict]: |
| 26 | calls: list[dict] = [] |
| 27 | monkeypatch.setattr( |
| 28 | deep_gemm_warmup, |
| 29 | "get_mk_alignment_for_contiguous_layout", |
| 30 | lambda: [128, 128], |
| 31 | ) |
| 32 | monkeypatch.setattr( |
| 33 | deep_gemm_warmup, |
| 34 | "_deepgemm_fp8_gemm_nt_warmup", |
| 35 | lambda **kwargs: calls.append(kwargs), |
| 36 | ) |
| 37 | model = SimpleNamespace(modules=lambda: iter(layers)) |
| 38 | deep_gemm_warmup.deepgemm_fp8_gemm_nt_warmup(model, max_tokens=16) |
| 39 | return calls |
| 40 | |
| 41 | |
| 42 | @pytest.mark.parametrize("scale_name", ["weight_scale", "weight_scale_inv"]) |
| 43 | def test_registered_deep_gemm_layer_is_warmed_up(monkeypatch, scale_name) -> None: |
| 44 | """Any layer stamped by the DeepGEMM kernel is warmed, regardless of the |
| 45 | quantization method that owns it or the name of its scale parameter.""" |
| 46 | layer = _block_fp8_layer(256, 128, scale_name) |
| 47 | |
| 48 | calls = _run_warmup(monkeypatch, [layer]) |
| 49 | |
| 50 | assert len(calls) == 1 |
| 51 | assert calls[0]["w"] is layer.weight |
| 52 | assert calls[0]["ws"] is getattr(layer, scale_name) |
| 53 | |
| 54 | |
| 55 | def test_n_multiple_of_64_matches_kernel_selection(monkeypatch) -> None: |
| 56 | """DeepGEMM accepts N % 64 == 0 (e.g. DeepSeek kv_a_proj_with_mqa, N=576), |
| 57 | so warmup must not require N % 128 == 0.""" |
| 58 | assert len(_run_warmup(monkeypatch, [_block_fp8_layer(576, 256)])) == 1 |
| 59 | |
| 60 | |
| 61 | def test_unstamped_and_mismatched_block_layers_are_skipped(monkeypatch) -> None: |
| 62 | mismatched = _block_fp8_layer(256, 128) |
| 63 | mismatched.weight_block_size = [1, 32] |
| 64 | layers = [SimpleNamespace(), mismatched, _block_fp8_layer(256, 128)] |
| 65 | |
| 66 | calls = _run_warmup(monkeypatch, layers) |
| 67 | |
| 68 | assert len(calls) == 1 |
| 69 | assert calls[0]["w"] is layers[-1].weight |
| 70 | |
| 71 | |
| 72 | @pytest.mark.parametrize("is_bmm", [False, True]) |
| 73 | def test_kernel_registers_itself_as_warmup_provider(is_bmm) -> None: |
| 74 | """Bmm layers bypass the kernel at runtime, so they must not be warmed.""" |
| 75 | kernel = object.__new__(DeepGemmFp8BlockScaledMMKernel) |
| 76 | kernel.is_deep_gemm_supported = False |
| 77 | layer = torch.nn.Module() |
| 78 | layer.weight = torch.nn.Parameter( |
| 79 | torch.empty((256, 128), dtype=torch.float8_e4m3fn), requires_grad=False |
| 80 | ) |
| 81 | layer.weight_scale = torch.nn.Parameter( |
| 82 | torch.empty((2, 1), dtype=torch.float32), requires_grad=False |
| 83 | ) |
| 84 | layer.weight_block_size = [128, 128] |
| 85 | layer.is_bmm = is_bmm |
| 86 | |
| 87 | kernel.process_weights_after_loading(layer) |
| 88 | |
| 89 | provider = getattr(layer, "deep_gemm_warmup_provider", None) |
| 90 | assert provider is (None if is_bmm else kernel) |