From 4ab5e5012a27e2b751c679873f231bafa0f6b098 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Mon, 17 Aug 2026 12:45:45 -0400 Subject: [PATCH] [Refactor] Simplify B12X linear kernels and warmup (#52368) Signed-off-by: mgoin --- .buildkite/test_areas/kernels.yaml | 22 + tests/kernels/quantization/test_block_fp8.py | 2 +- ...2x_mxfp8_linear.py => test_b12x_linear.py} | 622 +++++++----------- .../kernels/test_b12x_mxfp4_linear.py | 199 ------ .../kernels/test_b12x_nvfp4_linear.py | 246 ------- tests/model_executor/test_b12x_warmup.py | 242 +++++-- .../model_executor/kernels/linear/__init__.py | 4 +- .../kernels/linear/mxfp4/b12x.py | 110 +--- .../kernels/linear/mxfp8/b12x.py | 126 ++-- .../kernels/linear/nvfp4/b12x.py | 119 ++-- .../kernels/linear/scaled_mm/b12x.py | 339 ++++++++++ .../kernels/linear/scaled_mm/b12x_block.py | 214 ------ .../kernels/linear/scaled_mm/b12x_tensor.py | 238 ------- vllm/model_executor/warmup/b12x_warmup.py | 92 +-- vllm/utils/b12x.py | 11 +- 15 files changed, 969 insertions(+), 1617 deletions(-) rename tests/model_executor/kernels/{test_b12x_mxfp8_linear.py => test_b12x_linear.py} (62%) delete mode 100644 tests/model_executor/kernels/test_b12x_mxfp4_linear.py delete mode 100644 tests/model_executor/kernels/test_b12x_nvfp4_linear.py create mode 100644 vllm/model_executor/kernels/linear/scaled_mm/b12x.py delete mode 100644 vllm/model_executor/kernels/linear/scaled_mm/b12x_block.py delete mode 100644 vllm/model_executor/kernels/linear/scaled_mm/b12x_tensor.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 4e73d7620e..887ed03e37 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -339,6 +339,28 @@ steps: # e2e - pytest -v -s tests/models/quantization/test_nvfp4.py +- label: B12X Linear Kernels (DGX Spark) Nightly + key: b12x-linear-kernels-dgx-spark-nightly + timeout_in_minutes: 30 + device: dgx-spark + optional: true + num_devices: 1 + depends_on: + - arm64-image-build + source_file_dependencies: + - setup.py + - vllm/model_executor/kernels/linear/ + - vllm/model_executor/warmup/b12x_warmup.py + - vllm/utils/b12x.py + - tests/model_executor/kernels/test_b12x_linear.py + - tests/model_executor/test_b12x_warmup.py + - tests/kernels/quantization/test_block_fp8.py + commands: + - uv pip install --system b12x==1.2.4 + - pytest -v -s model_executor/kernels/test_b12x_linear.py + model_executor/test_b12x_warmup.py + kernels/quantization/test_block_fp8.py -k b12x + - label: Kernels Helion Test key: kernels-helion-test timeout_in_minutes: 115 diff --git a/tests/kernels/quantization/test_block_fp8.py b/tests/kernels/quantization/test_block_fp8.py index 6f86aed59e..52f87c403e 100644 --- a/tests/kernels/quantization/test_block_fp8.py +++ b/tests/kernels/quantization/test_block_fp8.py @@ -14,7 +14,7 @@ from tests.kernels.quant_utils import ( ) from tests.kernels.utils import fp8_ulp_distance from vllm.config import VllmConfig -from vllm.model_executor.kernels.linear.scaled_mm.b12x_block import ( +from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( B12xFp8BlockScaledMMKernel, _run_b12x_fp8_block_scaled_mm, ) diff --git a/tests/model_executor/kernels/test_b12x_mxfp8_linear.py b/tests/model_executor/kernels/test_b12x_linear.py similarity index 62% rename from tests/model_executor/kernels/test_b12x_mxfp8_linear.py rename to tests/model_executor/kernels/test_b12x_linear.py index 74f74e2f4f..1efc521bd8 100644 --- a/tests/model_executor/kernels/test_b12x_mxfp8_linear.py +++ b/tests/model_executor/kernels/test_b12x_linear.py @@ -3,6 +3,7 @@ from __future__ import annotations +import importlib import types from dataclasses import dataclass @@ -13,141 +14,120 @@ from vllm.model_executor.kernels.linear import ( _LINEAR_BACKEND_KERNEL_MAP, _POSSIBLE_FP8_BLOCK_KERNELS, _POSSIBLE_FP8_KERNELS, + _POSSIBLE_MXFP4_KERNELS, _POSSIBLE_MXFP8_KERNELS, - init_fp8_linear_kernel, - init_mxfp8_linear_kernel, -) -from vllm.model_executor.kernels.linear.mxfp8.b12x import ( - B12xMxfp8LinearKernel, - _b12x_mxfp8_expected_m, - warmup_b12x_mxfp8_linear, -) -from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import ( - Mxfp8LinearLayerConfig, -) -from vllm.model_executor.kernels.linear.scaled_mm.b12x_block import ( + _POSSIBLE_NVFP4_KERNELS, B12xFp8BlockScaledMMKernel, - _run_b12x_fp8_block_scaled_mm, - warmup_b12x_block_fp8_linear, -) -from vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor import ( + B12xMxFp4LinearKernel, + B12xMxfp8LinearKernel, + B12xNvFp4LinearKernel, B12xTensorFP8ScaledMMLinearKernel, - warmup_b12x_tensor_fp8_linear, -) -from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import ( FP8ScaledMMLinearLayerConfig, + Mxfp8LinearLayerConfig, + init_fp8_linear_kernel, + init_mxfp4_linear_kernel, + init_mxfp8_linear_kernel, + init_nvfp4_linear_kernel, ) -from vllm.model_executor.layers.linear import ( - ReplicatedLinear, - UnquantizedLinearMethod, +from vllm.model_executor.kernels.linear.nvfp4.marlin import ( + MarlinNvFp4LinearKernel, ) from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8Dynamic128Sym, kFp8Static128BlockSym, kFp8StaticTensorSym, + kMxfp4Dynamic, ) -from vllm.platforms import PlatformEnum, current_platform -from vllm.utils.b12x import b12x_warmup_token_counts - - -def test_b12x_backend_maps_mxfp8_kernel() -> None: - assert B12xMxfp8LinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"] - assert B12xMxfp8LinearKernel in _POSSIBLE_MXFP8_KERNELS[PlatformEnum.CUDA] - - -def test_b12x_backend_maps_tensor_fp8_kernel() -> None: - assert B12xTensorFP8ScaledMMLinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"] - assert B12xTensorFP8ScaledMMLinearKernel in _POSSIBLE_FP8_KERNELS[PlatformEnum.CUDA] +from vllm.platforms import PlatformEnum @pytest.mark.parametrize( - ("kernels", "before", "b12x", "after"), + ("kernel_cls", "kernels", "before", "after", "initializer", "kwargs"), [ ( - _POSSIBLE_FP8_KERNELS[PlatformEnum.CUDA], - "CutlassFP8ScaledMMLinearKernel", - "B12xTensorFP8ScaledMMLinearKernel", - "PerTensorTorchFP8ScaledMMLinearKernel", + B12xMxFp4LinearKernel, + _POSSIBLE_MXFP4_KERNELS[PlatformEnum.CUDA], + "HummingMxFp4LinearKernel", + "EmulationMxfp4LinearKernel", + init_mxfp4_linear_kernel, + {"activation_quant_key": kMxfp4Dynamic}, ), ( - _POSSIBLE_FP8_BLOCK_KERNELS[PlatformEnum.CUDA], - "CutlassFp8BlockScaledMMKernel", - "B12xFp8BlockScaledMMKernel", - "MarlinFP8ScaledMMLinearKernel", + B12xNvFp4LinearKernel, + _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA], + "FbgemmNvFp4LinearKernel", + "EmulationNvFp4LinearKernel", + init_nvfp4_linear_kernel, + {}, ), ( + B12xMxfp8LinearKernel, _POSSIBLE_MXFP8_KERNELS[PlatformEnum.CUDA], "MarlinMxfp8LinearKernel", - "B12xMxfp8LinearKernel", "EmulationMxfp8LinearKernel", + init_mxfp8_linear_kernel, + {}, + ), + ( + B12xTensorFP8ScaledMMLinearKernel, + _POSSIBLE_FP8_KERNELS[PlatformEnum.CUDA], + "CutlassFP8ScaledMMLinearKernel", + "PerTensorTorchFP8ScaledMMLinearKernel", + init_fp8_linear_kernel, + { + "activation_quant_key": kFp8StaticTensorSym, + "weight_quant_key": kFp8StaticTensorSym, + "input_dtype": torch.bfloat16, + "out_dtype": torch.bfloat16, + "weight_shape": (2048, 2048), + }, + ), + ( + B12xFp8BlockScaledMMKernel, + _POSSIBLE_FP8_BLOCK_KERNELS[PlatformEnum.CUDA], + "CutlassFp8BlockScaledMMKernel", + "MarlinFP8ScaledMMLinearKernel", + init_fp8_linear_kernel, + { + "activation_quant_key": kFp8Dynamic128Sym, + "weight_quant_key": kFp8Static128BlockSym, + "input_dtype": torch.bfloat16, + "out_dtype": torch.bfloat16, + "weight_shape": (2048, 2048), + }, ), ], ) -def test_b12x_fp8_fallback_priority( - kernels: list[type], - before: str, - b12x: str, - after: str, -) -> None: - names = [kernel.__name__ for kernel in kernels] - - assert names.index(before) < names.index(b12x) < names.index(after) - - -@torch.inference_mode() -def test_b12x_backend_does_not_intercept_unquantized_bf16( - default_vllm_config, - dist_init, -) -> None: - default_vllm_config.kernel_config.linear_backend = "b12x" - device = current_platform.device_type - layer = ReplicatedLinear( - 128, - 64, - bias=False, - params_dtype=torch.bfloat16, - quant_config=None, - prefix="bf16_linear", - ).to(device) - layer.weight.data.normal_() - x = torch.randn(8, 128, device=device, dtype=torch.bfloat16) - - output, output_bias = layer(x) - expected = torch.nn.functional.linear(x, layer.weight) - - assert isinstance(layer.quant_method, UnquantizedLinearMethod) - assert output_bias is None - torch.testing.assert_close(output, expected) - - -def test_b12x_explicit_backend_selects_per_tensor_fp8( +def test_b12x_backend_registration_priority_and_selection( monkeypatch, default_vllm_config, + kernel_cls, + kernels, + before: str, + after: str, + initializer, + kwargs: dict, ) -> None: import vllm.model_executor.kernels.linear as linear_mod + assert kernel_cls in _LINEAR_BACKEND_KERNEL_MAP["b12x"] + names = [kernel.__name__ for kernel in kernels] + assert names.index(before) < names.index(kernel_cls.__name__) < names.index(after) + monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") monkeypatch.setattr( - B12xTensorFP8ScaledMMLinearKernel, + kernel_cls, "is_supported", classmethod(lambda cls, compute_capability=None: (True, None)), ) monkeypatch.setattr( - B12xTensorFP8ScaledMMLinearKernel, + kernel_cls, "can_implement", classmethod(lambda cls, config: (True, None)), ) - kernel = init_fp8_linear_kernel( - activation_quant_key=kFp8StaticTensorSym, - weight_quant_key=kFp8StaticTensorSym, - input_dtype=torch.bfloat16, - out_dtype=torch.bfloat16, - weight_shape=(2048, 2048), - ) - - assert isinstance(kernel, B12xTensorFP8ScaledMMLinearKernel) + assert isinstance(initializer(**kwargs), kernel_cls) def test_b12x_tensor_fp8_can_implement_supported_config() -> None: @@ -165,33 +145,8 @@ def test_b12x_tensor_fp8_can_implement_supported_config() -> None: assert reason is None -def test_b12x_explicit_backend_selects_block_fp8( - monkeypatch, - default_vllm_config, -) -> None: - import vllm.model_executor.kernels.linear as linear_mod - - monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") - monkeypatch.setattr( - B12xFp8BlockScaledMMKernel, - "is_supported", - classmethod(lambda cls, compute_capability=None: (True, None)), - ) - - kernel = init_fp8_linear_kernel( - activation_quant_key=kFp8Dynamic128Sym, - weight_quant_key=kFp8Static128BlockSym, - input_dtype=torch.bfloat16, - out_dtype=torch.bfloat16, - weight_shape=(2048, 2048), - ) - - assert isinstance(kernel, B12xFp8BlockScaledMMKernel) - - def test_b12x_block_fp8_checks_runtime_support(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod platform = types.SimpleNamespace( is_cuda=lambda: True, @@ -261,75 +216,10 @@ def test_b12x_block_fp8_requires_aligned_features() -> None: ) -def test_b12x_warmup_token_counts_cover_serving_regimes() -> None: - assert b12x_warmup_token_counts( - max_tokens=2048, - cudagraph_capture_sizes=[1, 2, 8, 128], - ) == (1, 2, 8, 128, 2048) - - -def test_warmup_b12x_block_fp8_dedupes_weight_signatures(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod - - calls = [] - - def run(a, weight, a_scale, weight_scale, out_dtype): - calls.append((a.shape, weight, a_scale.shape, weight_scale, out_dtype)) - return torch.empty((a.shape[0], weight.shape[0]), dtype=out_dtype) - - platform = types.SimpleNamespace( - is_cuda=lambda: True, - is_device_capability_family=lambda family: family == 120, - ) - monkeypatch.setattr(b12x_mod, "current_platform", platform) - - monkeypatch.setattr( - b12x_mod, - "_import_b12x_blockscaled", - lambda: types.SimpleNamespace(), - ) - monkeypatch.setattr(b12x_mod, "_run_b12x_fp8_block_scaled_mm", run) - - def layer(in_features: int, out_features: int): - return types.SimpleNamespace( - b12x_block_fp8_linear=True, - weight=torch.empty((out_features, in_features), dtype=torch.float8_e4m3fn), - weight_scale_inv=torch.empty( - (out_features // 128, in_features // 128), dtype=torch.float32 - ), - ) - - layer_a = layer(128, 256) - layer_b = layer(256, 128) - modules = [ - layer_a, - layer_a, - layer_b, - types.SimpleNamespace(), - ] - model = types.SimpleNamespace(modules=lambda: iter(modules)) - - warmed = warmup_b12x_block_fp8_linear( - model, - max_tokens=32, - cudagraph_capture_sizes=[2, 8], - output_dtype=torch.bfloat16, - ) - - assert warmed == 8 - assert [call[0][0] for call in calls] == [1, 2, 8, 32] * 2 - assert [call[2][0] for call in calls] == [1, 2, 8, 32] * 2 - assert calls[0][1] is layer_a.weight - assert calls[4][1] is layer_b.weight - assert calls[0][3] is layer_a.weight_scale_inv - assert calls[4][3] is layer_b.weight_scale_inv - assert all(call[4] == torch.bfloat16 for call in calls) - - def test_b12x_tensor_fp8_process_weights_packs_modelopt_layout( monkeypatch, ) -> None: - import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod calls = [] packed = types.SimpleNamespace(out_features=64) @@ -367,6 +257,7 @@ def test_b12x_tensor_fp8_process_weights_packs_modelopt_layout( kernel.process_weights_after_loading(layer) assert layer.b12x_tensor_fp8_packed_weight is packed + assert layer.b12x_warmup_provider is kernel assert len(calls) == 1 weight, output_scale = calls[0] torch.testing.assert_close(weight, original_weight.T.contiguous()) @@ -378,67 +269,10 @@ def test_b12x_tensor_fp8_process_weights_packs_modelopt_layout( torch.testing.assert_close(layer.input_scale, torch.tensor(0.5)) -def test_warmup_b12x_tensor_fp8_dedupes_weight_signatures(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod - - calls = [] - - def prewarm(packed_weight, token_counts, *, out_dtype, stream): - del stream - calls.append((packed_weight, tuple(token_counts), out_dtype)) - return len(tuple(token_counts)) - - platform = types.SimpleNamespace( - is_cuda=lambda: True, - is_device_capability_family=lambda family: family == 120, - ) - monkeypatch.setattr(b12x_mod, "current_platform", platform) - monkeypatch.setattr( - b12x_mod, - "_import_b12x_tensor_fp8", - lambda: types.SimpleNamespace(prewarm=prewarm), - ) - monkeypatch.setattr( - b12x_mod, - "current_stream", - lambda: types.SimpleNamespace(cuda_stream=object()), - ) - - def packed(in_features: int, padded_in_features: int, out_features: int): - return types.SimpleNamespace( - in_features=in_features, - padded_in_features=padded_in_features, - out_features=out_features, - values=torch.empty(1), - ) - - packed_a = packed(128, 128, 256) - packed_b = packed(160, 256, 512) - modules = [ - types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_a), - types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_a), - types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_b), - types.SimpleNamespace(), - ] - model = types.SimpleNamespace(modules=lambda: iter(modules)) - - warmed = warmup_b12x_tensor_fp8_linear( - model, - max_tokens=2048, - cudagraph_capture_sizes=[1, 2], - ) - - assert warmed == 6 - assert calls == [ - (packed_a, (1, 2, 2048), torch.bfloat16), - (packed_b, (1, 2, 2048), torch.bfloat16), - ] - - def test_b12x_tensor_fp8_apply_quantizes_and_uses_packed_weight( monkeypatch, ) -> None: - import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod calls = [] @@ -506,27 +340,6 @@ def test_b12x_tensor_fp8_apply_quantizes_and_uses_packed_weight( assert expected_m == 6 -def test_b12x_mxfp8_explicit_backend_selects_kernel(monkeypatch) -> None: - import vllm.model_executor.kernels.linear as linear_mod - - monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") - monkeypatch.setattr( - B12xMxfp8LinearKernel, - "is_supported", - classmethod(lambda cls, compute_capability=None: (True, None)), - ) - monkeypatch.setattr( - B12xMxfp8LinearKernel, - "can_implement", - classmethod(lambda cls, c: (True, None)), - ) - - kernel = init_mxfp8_linear_kernel() - - assert isinstance(kernel, B12xMxfp8LinearKernel) - - def test_b12x_mxfp8_can_implement_supported_config() -> None: can_implement, reason = B12xMxfp8LinearKernel.can_implement( Mxfp8LinearLayerConfig() @@ -536,90 +349,6 @@ def test_b12x_mxfp8_can_implement_supported_config() -> None: assert reason is None -def test_b12x_mxfp8_expected_m_uses_live_m() -> None: - assert _b12x_mxfp8_expected_m(0) == 1 - assert _b12x_mxfp8_expected_m(1) == 1 - assert _b12x_mxfp8_expected_m(2) == 2 - assert _b12x_mxfp8_expected_m(8) == 8 - assert _b12x_mxfp8_expected_m(9) == 9 - assert _b12x_mxfp8_expected_m(128) == 128 - assert _b12x_mxfp8_expected_m(129) == 129 - assert _b12x_mxfp8_expected_m(2048) == 2048 - - -def test_warmup_b12x_mxfp8_linear_dedupes_weight_signatures( - monkeypatch, -) -> None: - import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod - - calls = [] - - def mm( - source: torch.Tensor, - packed_weight, - *, - bias: torch.Tensor | None = None, - expected_m: int | None = None, - stream: object = None, - ) -> torch.Tensor: - del stream - calls.append((source.shape, packed_weight, bias, expected_m)) - return source.new_empty((source.shape[0], packed_weight.out_features)) - - platform = types.SimpleNamespace( - is_cuda=lambda: True, - is_device_capability_family=lambda family: family == 120, - ) - monkeypatch.setattr(b12x_mod, "current_platform", platform) - monkeypatch.setattr( - b12x_mod, - "_import_b12x_mxfp8", - lambda: types.SimpleNamespace(mm=mm), - ) - monkeypatch.setattr( - b12x_mod, - "current_stream", - lambda: types.SimpleNamespace(cuda_stream=object()), - ) - - def packed(in_features: int, padded_in_features: int, out_features: int): - return types.SimpleNamespace( - in_features=in_features, - padded_in_features=padded_in_features, - out_features=out_features, - weight=types.SimpleNamespace(values=torch.empty(1)), - ) - - packed_a = packed(128, 128, 256) - packed_b = packed(128, 128, 512) - modules = [ - types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_a), - types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_a), - types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_b), - types.SimpleNamespace(), - ] - model = types.SimpleNamespace(modules=lambda: iter(modules)) - - warmed = warmup_b12x_mxfp8_linear( - model, - max_tokens=2048, - cudagraph_capture_sizes=[1, 2], - ) - - assert warmed == 6 - assert [call[0] for call in calls] == [ - torch.Size([1, 128]), - torch.Size([2, 128]), - torch.Size([2048, 128]), - torch.Size([1, 128]), - torch.Size([2, 128]), - torch.Size([2048, 128]), - ] - assert [call[3] for call in calls] == [1, 2, 2048, 1, 2, 2048] - assert calls[0][1] is packed_a - assert calls[3][1] is packed_b - - def test_b12x_mxfp8_support_check_reports_missing_import(monkeypatch) -> None: import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod @@ -693,6 +422,7 @@ def test_b12x_mxfp8_process_weights_packs_modelopt_layout(monkeypatch) -> None: kernel.process_weights_after_loading(layer) assert layer.b12x_mxfp8_packed_weight is packed + assert layer.b12x_warmup_provider is kernel assert len(calls) == 1 weight, weight_scale = calls[0] assert weight.shape == (48, 128) @@ -805,7 +535,7 @@ def test_b12x_block_fp8_process_weights_keeps_native_block_layout() -> None: kernel.process_weights_after_loading(layer) - assert layer.b12x_block_fp8_linear + assert layer.b12x_warmup_provider is kernel assert layer.weight.shape == (128, 128) assert layer.weight.dtype == torch.float8_e4m3fn assert layer.weight_scale_inv.shape == (1, 1) @@ -882,7 +612,7 @@ def test_b12x_mxfp8_apply_uses_packed_weight(monkeypatch) -> None: def test_b12x_block_fp8_apply_uses_b12x_recipe_api(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod calls = [] @@ -919,37 +649,183 @@ def test_b12x_block_fp8_apply_uses_b12x_recipe_api(monkeypatch) -> None: torch.testing.assert_close(output, torch.full_like(output, 13.0)) -def test_b12x_block_fp8_helper_uses_regular_compact_scale_api(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod +def test_b12x_mxfp4_requires_dynamic_activations() -> None: + config = types.SimpleNamespace(activation_quant_key=kMxfp4Dynamic) + can_implement, reason = B12xMxFp4LinearKernel.can_implement(config) - calls = [] + assert can_implement + assert reason is None - def mm_block_fp8(*args, **kwargs): + config.activation_quant_key = None + can_implement, reason = B12xMxFp4LinearKernel.can_implement(config) + + assert not can_implement + assert reason == "B12X MXFP4 GEMM requires dynamic MXFP4 activations" + + +@pytest.mark.parametrize( + ("kernel_cls", "module_name", "scale_dtype"), + [ + ( + B12xMxFp4LinearKernel, + "vllm.model_executor.kernels.linear.mxfp4.b12x", + torch.uint8, + ), + ( + B12xNvFp4LinearKernel, + "vllm.model_executor.kernels.linear.nvfp4.b12x", + torch.float8_e4m3fn, + ), + ], +) +def test_b12x_fp4_processes_scale_and_preserves_loader( + monkeypatch, + kernel_cls, + module_name: str, + scale_dtype: torch.dtype, +) -> None: + scale = torch.empty((48, 8), dtype=scale_dtype) + swizzled_scale = torch.empty((128, 8), dtype=scale_dtype) + intrinsics = types.SimpleNamespace(swizzle_block_scale=lambda value: swizzled_scale) + monkeypatch.setattr( + importlib.import_module(module_name), + "_import_b12x_intrinsics", + lambda: intrinsics, + ) + layer = torch.nn.Module() + layer.prefix = "model.layers.0.mlp.shared_expert.down_proj" + layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False) + weight_loader = object() + layer.weight_scale.weight_loader = weight_loader + kernel = object.__new__(kernel_cls) + + kernel.process_weights_after_loading(layer) + + assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr() + assert layer.weight_scale.weight_loader is weight_loader + assert layer.b12x_warmup_provider is kernel + + +def test_b12x_mxfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod + import vllm.utils.flashinfer as flashinfer_utils + + calls: list[tuple] = [] + x_packed = torch.empty((6, 64), dtype=torch.uint8) + x_scale_storage = torch.empty((128, 4), dtype=torch.uint8) + + def mm_mxfp4(*args, **kwargs): calls.append((args, kwargs)) - return torch.full((6, 256), 17.0, dtype=torch.bfloat16) + return torch.full((6, 48), 3.0, dtype=torch.bfloat16) + monkeypatch.setattr( + flashinfer_utils, + "flashinfer_mxfp4_quantize", + lambda *args, **kwargs: (x_packed, x_scale_storage), + ) monkeypatch.setattr( b12x_mod, "_import_b12x_blockscaled", - lambda: types.SimpleNamespace(mm_block_fp8=mm_block_fp8), + lambda: types.SimpleNamespace(mm_mxfp4=mm_mxfp4), ) - a = torch.empty((6, 128), dtype=torch.float8_e4m3fn) - weight = torch.empty((256, 128), dtype=torch.float8_e4m3fn) - a_scale = torch.empty((6, 1), dtype=torch.float32) - weight_scale = torch.empty((2, 1), dtype=torch.float32) - output = _run_b12x_fp8_block_scaled_mm( - a, - weight, - a_scale, - weight_scale, - torch.bfloat16, - ) + layer = torch.nn.Module() + layer.output_size_per_partition = 48 + layer.weight = torch.empty((48, 64), dtype=torch.uint8) + layer.weight_scale = torch.empty((128, 4), dtype=torch.uint8) + x = torch.empty((2, 3, 128), dtype=torch.bfloat16) + bias = torch.ones(48, dtype=torch.bfloat16) + kernel = object.__new__(B12xMxFp4LinearKernel) - assert output.shape == (6, 256) + output = kernel.apply_weights(layer, x, bias) + + assert output.shape == (2, 3, 48) + torch.testing.assert_close(output, torch.full_like(output, 4.0)) + assert len(calls) == 1 args, kwargs = calls[0] - assert args == (a, a_scale, weight, weight_scale) - assert kwargs == { - "out_dtype": torch.bfloat16, - } - torch.testing.assert_close(output, torch.full_like(output, 17.0)) + assert args == ( + x_packed, + x_scale_storage, + layer.weight, + layer.weight_scale, + ) + assert kwargs == {"out_dtype": torch.bfloat16} + + +def test_b12x_nvfp4_can_implement_supported_config() -> None: + can_implement, reason = B12xNvFp4LinearKernel.can_implement(None) + + assert can_implement + assert reason is None + + +def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None: + import vllm.model_executor.kernels.linear as linear_mod + + monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) + monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") + monkeypatch.setattr( + MarlinNvFp4LinearKernel, + "is_supported", + classmethod(lambda cls, compute_capability=None: (True, None)), + ) + + kernel = init_nvfp4_linear_kernel(use_a16=True) + + assert isinstance(kernel, MarlinNvFp4LinearKernel) + + +def test_b12x_nvfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod + + calls: list[tuple] = [] + quant_calls: list[tuple] = [] + x_packed = torch.empty((6, 64), dtype=torch.uint8) + x_scale_storage = torch.empty((128, 8), dtype=torch.float8_e4m3fn) + + def quant(*args, **kwargs): + quant_calls.append((args, kwargs)) + return x_packed, x_scale_storage + + def mm_nvfp4(*args, **kwargs): + calls.append((args, kwargs)) + return torch.full((6, 48), 3.0, dtype=torch.bfloat16) + + monkeypatch.setattr(b12x_mod, "scaled_fp4_quant", quant) + monkeypatch.setattr( + b12x_mod, + "_import_b12x_blockscaled", + lambda: types.SimpleNamespace(mm_nvfp4=mm_nvfp4), + ) + + layer = torch.nn.Module() + layer.output_size_per_partition = 48 + layer.weight = torch.empty((48, 64), dtype=torch.uint8) + layer.weight_scale = torch.empty((128, 8), dtype=torch.float8_e4m3fn) + layer.input_global_scale_inv = torch.tensor(2.0) + layer.alpha = torch.tensor(0.25) + x = torch.empty((2, 3, 256), dtype=torch.bfloat16)[..., ::2] + bias = torch.ones(48, dtype=torch.bfloat16) + kernel = object.__new__(B12xNvFp4LinearKernel) + + output = kernel.apply_weights(layer, x, bias) + + assert output.shape == (2, 3, 48) + torch.testing.assert_close(output, torch.full_like(output, 4.0)) + assert len(quant_calls) == 1 + quant_args, quant_kwargs = quant_calls[0] + assert quant_args[0].shape == (6, 128) + assert quant_args[0].data_ptr() == x.data_ptr() + assert quant_args[1] is layer.input_global_scale_inv + assert not quant_args[0].is_contiguous() + assert quant_kwargs == {"is_sf_swizzled_layout": True} + assert len(calls) == 1 + args, kwargs = calls[0] + assert args == ( + x_packed, + x_scale_storage, + layer.weight, + layer.weight_scale, + layer.alpha, + ) + assert kwargs == {"out_dtype": torch.bfloat16} diff --git a/tests/model_executor/kernels/test_b12x_mxfp4_linear.py b/tests/model_executor/kernels/test_b12x_mxfp4_linear.py deleted file mode 100644 index 7ba1aad817..0000000000 --- a/tests/model_executor/kernels/test_b12x_mxfp4_linear.py +++ /dev/null @@ -1,199 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import types - -import torch - -from vllm.model_executor.kernels.linear import ( - _LINEAR_BACKEND_KERNEL_MAP, - _POSSIBLE_MXFP4_KERNELS, - init_mxfp4_linear_kernel, -) -from vllm.model_executor.kernels.linear.mxfp4.b12x import ( - B12xMxFp4LinearKernel, - warmup_b12x_mxfp4_linear, -) -from vllm.model_executor.layers.quantization.utils.quant_utils import ( - kMxfp4Dynamic, -) -from vllm.platforms import PlatformEnum - - -def test_b12x_backend_maps_mxfp4_kernel() -> None: - assert B12xMxFp4LinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"] - assert B12xMxFp4LinearKernel in _POSSIBLE_MXFP4_KERNELS[PlatformEnum.CUDA] - - -def test_b12x_mxfp4_fallback_priority() -> None: - kernels = _POSSIBLE_MXFP4_KERNELS[PlatformEnum.CUDA] - names = [kernel.__name__ for kernel in kernels] - - assert ( - names.index("HummingMxFp4LinearKernel") - < names.index("B12xMxFp4LinearKernel") - < names.index("EmulationMxfp4LinearKernel") - ) - - -def test_b12x_mxfp4_explicit_backend_selects_native_kernel(monkeypatch) -> None: - import vllm.model_executor.kernels.linear as linear_mod - - monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") - monkeypatch.setattr( - B12xMxFp4LinearKernel, - "is_supported", - classmethod(lambda cls, compute_capability=None: (True, None)), - ) - monkeypatch.setattr( - B12xMxFp4LinearKernel, - "can_implement", - classmethod(lambda cls, config: (True, None)), - ) - - kernel = init_mxfp4_linear_kernel(activation_quant_key=kMxfp4Dynamic) - - assert isinstance(kernel, B12xMxFp4LinearKernel) - - -def test_b12x_mxfp4_requires_dynamic_activations() -> None: - config = types.SimpleNamespace(activation_quant_key=kMxfp4Dynamic) - can_implement, reason = B12xMxFp4LinearKernel.can_implement(config) - - assert can_implement - assert reason is None - - config.activation_quant_key = None - can_implement, reason = B12xMxFp4LinearKernel.can_implement(config) - - assert not can_implement - assert reason == "B12X MXFP4 GEMM requires dynamic MXFP4 activations" - - -def test_b12x_mxfp4_processes_scale_and_preserves_loader(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod - - scale = torch.empty((48, 8), dtype=torch.uint8) - swizzled_scale = torch.empty((128, 8), dtype=torch.uint8) - monkeypatch.setattr( - b12x_mod, - "_import_b12x_intrinsics", - lambda: types.SimpleNamespace(swizzle_block_scale=lambda value: swizzled_scale), - ) - - layer = torch.nn.Module() - layer.prefix = "model.layers.0.mlp.shared_expert.down_proj" - layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False) - weight_loader = object() - layer.weight_scale.weight_loader = weight_loader - kernel = object.__new__(B12xMxFp4LinearKernel) - - kernel.process_weights_after_loading(layer) - - assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr() - assert layer.weight_scale.weight_loader is weight_loader - assert layer.b12x_mxfp4_linear - - -def test_b12x_mxfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod - import vllm.utils.flashinfer as flashinfer_utils - - calls: list[tuple] = [] - x_packed = torch.empty((6, 64), dtype=torch.uint8) - x_scale_storage = torch.empty((128, 4), dtype=torch.uint8) - - def mm_mxfp4(*args, **kwargs): - calls.append((args, kwargs)) - return torch.full((6, 48), 3.0, dtype=torch.bfloat16) - - monkeypatch.setattr( - flashinfer_utils, - "flashinfer_mxfp4_quantize", - lambda *args, **kwargs: (x_packed, x_scale_storage), - ) - monkeypatch.setattr( - b12x_mod, - "_import_b12x_blockscaled", - lambda: types.SimpleNamespace(mm_mxfp4=mm_mxfp4), - ) - - layer = torch.nn.Module() - layer.output_size_per_partition = 48 - layer.weight = torch.empty((48, 64), dtype=torch.uint8) - layer.weight_scale = torch.empty((128, 4), dtype=torch.uint8) - x = torch.empty((2, 3, 128), dtype=torch.bfloat16) - bias = torch.ones(48, dtype=torch.bfloat16) - kernel = object.__new__(B12xMxFp4LinearKernel) - - output = kernel.apply_weights(layer, x, bias) - - assert output.shape == (2, 3, 48) - torch.testing.assert_close(output, torch.full_like(output, 4.0)) - assert len(calls) == 1 - args, kwargs = calls[0] - assert args == ( - x_packed, - x_scale_storage, - layer.weight, - layer.weight_scale, - ) - assert kwargs == {"out_dtype": torch.bfloat16} - - -def test_warmup_b12x_mxfp4_dedupes_weight_signatures(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod - - calls = [] - - def apply(source, weight, weight_scale, bias): - calls.append((source.shape, weight, weight_scale, bias)) - return source.new_empty((source.shape[0], weight.shape[0])) - - platform = types.SimpleNamespace( - is_cuda=lambda: True, - is_device_capability_family=lambda family: family == 120, - ) - monkeypatch.setattr(b12x_mod, "current_platform", platform) - monkeypatch.setattr( - b12x_mod, - "_import_b12x_blockscaled", - lambda: types.SimpleNamespace(), - ) - monkeypatch.setattr(b12x_mod, "_apply_b12x_mxfp4_linear", apply) - - def layer(n: int): - return types.SimpleNamespace( - b12x_mxfp4_linear=True, - weight=torch.empty((n, 64), dtype=torch.uint8), - weight_scale=torch.empty((128, 4), dtype=torch.uint8), - ) - - layer_a = layer(48) - layer_b = layer(48) - layer_c = layer(96) - model = types.SimpleNamespace( - modules=lambda: iter([layer_a, layer_b, layer_c, types.SimpleNamespace()]) - ) - - warmed = warmup_b12x_mxfp4_linear( - model, - max_tokens=8, - cudagraph_capture_sizes=[1, 2], - ) - - assert warmed == 6 - assert [call[0] for call in calls] == [ - torch.Size([1, 128]), - torch.Size([2, 128]), - torch.Size([8, 128]), - torch.Size([1, 128]), - torch.Size([2, 128]), - torch.Size([8, 128]), - ] - assert calls[0][1] is layer_a.weight - assert calls[3][1] is layer_c.weight - assert all(call[3] is None for call in calls) diff --git a/tests/model_executor/kernels/test_b12x_nvfp4_linear.py b/tests/model_executor/kernels/test_b12x_nvfp4_linear.py deleted file mode 100644 index 1e2717c4e0..0000000000 --- a/tests/model_executor/kernels/test_b12x_nvfp4_linear.py +++ /dev/null @@ -1,246 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -import types - -import torch - -from vllm.model_executor.kernels.linear import ( - _LINEAR_BACKEND_KERNEL_MAP, - _POSSIBLE_NVFP4_KERNELS, - _resolve_backend_kernels, - init_nvfp4_linear_kernel, -) -from vllm.model_executor.kernels.linear.nvfp4.b12x import ( - B12xNvFp4LinearKernel, - warmup_b12x_nvfp4_linear, -) -from vllm.model_executor.kernels.linear.nvfp4.marlin import ( - MarlinNvFp4LinearKernel, -) -from vllm.platforms import PlatformEnum - - -def test_b12x_backend_maps_nvfp4_kernel() -> None: - assert B12xNvFp4LinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"] - assert B12xNvFp4LinearKernel in _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA] - - -def test_b12x_nvfp4_fallback_priority() -> None: - kernels = _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA] - names = [kernel.__name__ for kernel in kernels] - - assert ( - names.index("FbgemmNvFp4LinearKernel") - < names.index("B12xNvFp4LinearKernel") - < names.index("EmulationNvFp4LinearKernel") - ) - - -def test_b12x_nvfp4_explicit_backend_selects_native_kernel(monkeypatch) -> None: - import vllm.model_executor.kernels.linear as linear_mod - - monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") - monkeypatch.setattr( - B12xNvFp4LinearKernel, - "is_supported", - classmethod(lambda cls, compute_capability=None: (True, None)), - ) - monkeypatch.setattr( - B12xNvFp4LinearKernel, - "can_implement", - classmethod(lambda cls, config: (True, None)), - ) - - kernel = init_nvfp4_linear_kernel() - - assert isinstance(kernel, B12xNvFp4LinearKernel) - - -def test_b12x_nvfp4_can_implement_supported_config() -> None: - can_implement, reason = B12xNvFp4LinearKernel.can_implement(None) - - assert can_implement - assert reason is None - - -def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None: - import vllm.model_executor.kernels.linear as linear_mod - - monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA) - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") - monkeypatch.setattr( - MarlinNvFp4LinearKernel, - "is_supported", - classmethod(lambda cls, compute_capability=None: (True, None)), - ) - - kernel = init_nvfp4_linear_kernel(use_a16=True) - - assert isinstance(kernel, MarlinNvFp4LinearKernel) - - -def test_backend_without_w4a16_kernel_preserves_fallback(monkeypatch) -> None: - import vllm.model_executor.kernels.linear as linear_mod - - kernels = [MarlinNvFp4LinearKernel] - monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x") - assert _resolve_backend_kernels(kernels, "NVFP4") == kernels - - -def test_b12x_nvfp4_processes_scale_and_preserves_loader(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod - - scale = torch.empty((48, 8), dtype=torch.float8_e4m3fn) - swizzled_scale = torch.empty((128, 8), dtype=torch.float8_e4m3fn) - intrinsics = types.SimpleNamespace(swizzle_block_scale=lambda value: swizzled_scale) - monkeypatch.setattr(b12x_mod, "_import_b12x_intrinsics", lambda: intrinsics) - - layer = torch.nn.Module() - layer.prefix = "model.layers.0.mlp.shared_expert.down_proj" - layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False) - weight_loader = object() - layer.weight_scale.weight_loader = weight_loader - kernel = object.__new__(B12xNvFp4LinearKernel) - - kernel.process_weights_after_loading(layer) - - assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr() - assert layer.weight_scale.weight_loader is weight_loader - assert layer.b12x_nvfp4_linear - - -def test_b12x_nvfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod - - calls: list[tuple] = [] - quant_calls: list[tuple] = [] - x_packed = torch.empty((6, 64), dtype=torch.uint8) - x_scale_storage = torch.empty((128, 8), dtype=torch.float8_e4m3fn) - - def quant(*args, **kwargs): - quant_calls.append((args, kwargs)) - return x_packed, x_scale_storage - - def mm_nvfp4(*args, **kwargs): - calls.append((args, kwargs)) - return torch.full((6, 48), 3.0, dtype=torch.bfloat16) - - monkeypatch.setattr( - b12x_mod, - "scaled_fp4_quant", - quant, - ) - monkeypatch.setattr( - b12x_mod, - "_import_b12x_blockscaled", - lambda: types.SimpleNamespace(mm_nvfp4=mm_nvfp4), - ) - - layer = torch.nn.Module() - layer.output_size_per_partition = 48 - layer.weight = torch.empty((48, 64), dtype=torch.uint8) - layer.weight_scale = torch.empty((128, 8), dtype=torch.float8_e4m3fn) - layer.input_global_scale_inv = torch.tensor(2.0) - layer.alpha = torch.tensor(0.25) - x = torch.empty((2, 3, 256), dtype=torch.bfloat16)[..., ::2] - bias = torch.ones(48, dtype=torch.bfloat16) - kernel = object.__new__(B12xNvFp4LinearKernel) - - output = kernel.apply_weights(layer, x, bias) - - assert output.shape == (2, 3, 48) - torch.testing.assert_close(output, torch.full_like(output, 4.0)) - assert len(quant_calls) == 1 - quant_args, quant_kwargs = quant_calls[0] - assert quant_args[0].shape == (6, 128) - assert quant_args[0].data_ptr() == x.data_ptr() - assert quant_args[1] is layer.input_global_scale_inv - assert not quant_args[0].is_contiguous() - assert quant_kwargs == {"is_sf_swizzled_layout": True} - assert len(calls) == 1 - args, kwargs = calls[0] - assert args == ( - x_packed, - x_scale_storage, - layer.weight, - layer.weight_scale, - layer.alpha, - ) - assert kwargs == {"out_dtype": torch.bfloat16} - - -def test_warmup_b12x_nvfp4_dedupes_weight_signatures(monkeypatch) -> None: - import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod - - calls = [] - - def apply( - source, - weight, - weight_scale, - input_global_scale_inv, - alpha, - bias, - ): - calls.append( - ( - source.shape, - weight, - weight_scale, - input_global_scale_inv, - alpha, - bias, - ) - ) - return source.new_empty((source.shape[0], weight.shape[0])) - - platform = types.SimpleNamespace( - is_cuda=lambda: True, - is_device_capability_family=lambda family: family == 120, - ) - monkeypatch.setattr(b12x_mod, "current_platform", platform) - monkeypatch.setattr( - b12x_mod, - "_import_b12x_blockscaled", - lambda: types.SimpleNamespace(), - ) - monkeypatch.setattr(b12x_mod, "_apply_b12x_nvfp4_linear", apply) - - def layer(n: int): - return types.SimpleNamespace( - b12x_nvfp4_linear=True, - weight=torch.empty((n, 64), dtype=torch.uint8), - weight_scale=torch.empty((128, 8), dtype=torch.float8_e4m3fn), - input_global_scale_inv=torch.tensor(2.0), - alpha=torch.tensor(0.25), - ) - - layer_a = layer(48) - layer_b = layer(48) - layer_c = layer(96) - model = types.SimpleNamespace( - modules=lambda: iter([layer_a, layer_b, layer_c, types.SimpleNamespace()]) - ) - - warmed = warmup_b12x_nvfp4_linear( - model, - max_tokens=8, - cudagraph_capture_sizes=[1, 2], - ) - - assert warmed == 6 - assert [call[0] for call in calls] == [ - torch.Size([1, 128]), - torch.Size([2, 128]), - torch.Size([8, 128]), - torch.Size([1, 128]), - torch.Size([2, 128]), - torch.Size([8, 128]), - ] - assert calls[0][1] is layer_a.weight - assert calls[3][1] is layer_c.weight - assert all(call[5] is None for call in calls) diff --git a/tests/model_executor/test_b12x_warmup.py b/tests/model_executor/test_b12x_warmup.py index bd71f9bcd2..26069e0a9a 100644 --- a/tests/model_executor/test_b12x_warmup.py +++ b/tests/model_executor/test_b12x_warmup.py @@ -7,94 +7,200 @@ from types import SimpleNamespace import pytest import torch +from vllm.model_executor.kernels.linear import ( + B12xFp8BlockScaledMMKernel, + B12xMxFp4LinearKernel, + B12xMxfp8LinearKernel, + B12xNvFp4LinearKernel, + B12xTensorFP8ScaledMMLinearKernel, +) from vllm.model_executor.warmup.b12x_warmup import b12x_warmup +from vllm.utils.b12x import B12xWarmupUnit, b12x_warmup_token_counts + + +def test_b12x_warmup_token_counts_cover_serving_regimes() -> None: + assert b12x_warmup_token_counts( + max_tokens=2048, + cudagraph_capture_sizes=[1, 2, 8, 128], + ) == (1, 2, 8, 128, 2048) @pytest.mark.parametrize( - ("module_name", "warmup_name", "import_name"), + ("kernel_cls", "module_name", "call_name", "layer", "name"), [ ( - "vllm.model_executor.kernels.linear.scaled_mm.b12x_block", - "warmup_b12x_block_fp8_linear", - "_import_b12x_blockscaled", - ), - ( - "vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor", - "warmup_b12x_tensor_fp8_linear", - "_import_b12x_tensor_fp8", - ), - ( - "vllm.model_executor.kernels.linear.mxfp8.b12x", - "warmup_b12x_mxfp8_linear", - "_import_b12x_mxfp8", - ), - ( + B12xMxFp4LinearKernel, "vllm.model_executor.kernels.linear.mxfp4.b12x", - "warmup_b12x_mxfp4_linear", - "_import_b12x_blockscaled", + "_apply_b12x_mxfp4_linear", + SimpleNamespace( + weight=torch.empty((48, 64), dtype=torch.uint8), + weight_scale=torch.empty((128, 4), dtype=torch.uint8), + ), + "MXFP4", ), ( + B12xNvFp4LinearKernel, "vllm.model_executor.kernels.linear.nvfp4.b12x", - "warmup_b12x_nvfp4_linear", - "_import_b12x_blockscaled", + "_apply_b12x_nvfp4_linear", + SimpleNamespace( + weight=torch.empty((48, 64), dtype=torch.uint8), + weight_scale=torch.empty((128, 8), dtype=torch.float8_e4m3fn), + input_global_scale_inv=torch.tensor(2.0), + alpha=torch.tensor(0.25), + ), + "NVFP4", + ), + ( + B12xFp8BlockScaledMMKernel, + "vllm.model_executor.kernels.linear.scaled_mm.b12x", + "_run_b12x_fp8_block_scaled_mm", + SimpleNamespace( + weight=torch.empty((256, 128), dtype=torch.float8_e4m3fn), + weight_scale_inv=torch.empty((2, 1), dtype=torch.float32), + ), + "block-FP8", ), ], ) -def test_b12x_linear_warmup_skips_unused_provider( +def test_b12x_warmup_units_cover_token_counts( monkeypatch, + kernel_cls, module_name: str, - warmup_name: str, - import_name: str, + call_name: str, + layer, + name: str, ) -> None: - module = importlib.import_module(module_name) + calls = [] + monkeypatch.setattr( + importlib.import_module(module_name), + call_name, + lambda *args: calls.append(args), + ) + kernel = object.__new__(kernel_cls) + + unit = kernel.get_b12x_warmup_unit(layer, (1, 8), torch.bfloat16) + unit.compile() + + assert unit.name == name + assert [args[0].shape[0] for args in calls] == [1, 8] + assert unit.key[-1] == torch.bfloat16 + + +def test_b12x_mxfp8_warmup_unit(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod + + calls = [] + monkeypatch.setattr( + b12x_mod, + "_import_b12x_mxfp8", + lambda: SimpleNamespace( + mm=lambda *args, **kwargs: calls.append((args, kwargs)) + ), + ) + monkeypatch.setattr( + b12x_mod, + "current_stream", + lambda: SimpleNamespace(cuda_stream=object()), + ) + packed_weight = SimpleNamespace( + in_features=128, + padded_in_features=128, + out_features=256, + weight=SimpleNamespace(values=torch.empty(1)), + ) + layer = SimpleNamespace(b12x_mxfp8_packed_weight=packed_weight) + kernel = object.__new__(B12xMxfp8LinearKernel) + + unit = kernel.get_b12x_warmup_unit(layer, (1, 8), torch.float16) + unit.compile() + + assert [args[0].shape for args, _ in calls] == [(1, 128), (8, 128)] + assert [kwargs["expected_m"] for _, kwargs in calls] == [1, 8] + + +def test_b12x_tensor_fp8_warmup_unit(monkeypatch) -> None: + import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod + + calls = [] + monkeypatch.setattr( + b12x_mod, + "_import_b12x_tensor_fp8", + lambda: SimpleNamespace( + prewarm=lambda *args, **kwargs: calls.append((args, kwargs)) + ), + ) + monkeypatch.setattr( + b12x_mod, + "current_stream", + lambda: SimpleNamespace(cuda_stream=object()), + ) + packed_weight = SimpleNamespace( + in_features=128, + padded_in_features=128, + out_features=256, + values=torch.empty(1), + ) + layer = SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_weight) + kernel = object.__new__(B12xTensorFP8ScaledMMLinearKernel) + + unit = kernel.get_b12x_warmup_unit(layer, (1, 8), torch.bfloat16) + unit.compile() + + assert calls[0][0] == (packed_weight, (1, 8)) + assert calls[0][1]["out_dtype"] == torch.bfloat16 + + +def test_b12x_warmup_deduplicates_registered_signatures(monkeypatch) -> None: + import vllm.model_executor.warmup.b12x_warmup as warmup_mod + + calls: list[tuple[str, tuple[int, ...], torch.dtype]] = [] + + class Provider: + def get_b12x_warmup_unit(self, layer, token_counts, output_dtype): + return B12xWarmupUnit( + name="fake", + key=(type(self), layer.shape, output_dtype), + compile=lambda: calls.append((layer.name, token_counts, output_dtype)), + ) + + provider = Provider() + layers = [ + SimpleNamespace(name="first", shape=(128, 256), b12x_warmup_provider=provider), + SimpleNamespace( + name="duplicate", shape=(128, 256), b12x_warmup_provider=provider + ), + SimpleNamespace(name="second", shape=(256, 256), b12x_warmup_provider=provider), + SimpleNamespace(), + ] + scans = 0 + + def modules(): + nonlocal scans + scans += 1 + return iter(layers) + + worker = SimpleNamespace( + get_model=lambda: SimpleNamespace(modules=modules), + scheduler_config=SimpleNamespace(max_num_batched_tokens=8), + model_config=SimpleNamespace(dtype=torch.float32), + ) platform = SimpleNamespace( is_cuda=lambda: True, is_device_capability_family=lambda family: family == 120, ) - monkeypatch.setattr(module, "current_platform", platform) - - def fail_import(): - pytest.fail("unused B12X provider was imported") - - monkeypatch.setattr(module, import_name, fail_import) - - warmup = getattr(module, warmup_name) - assert warmup(torch.nn.Module(), max_tokens=128) == 0 - - -def test_b12x_warmup_covers_linear_serving_shapes(monkeypatch) -> None: - import vllm.model_executor.warmup.b12x_warmup as warmup_mod - - model = torch.nn.Module() - linear_calls: list[tuple[torch.nn.Module, dict]] = [] - - def linear_warmup(model, **kwargs): - linear_calls.append((model, kwargs)) - return 0 - - monkeypatch.setattr(warmup_mod, "warmup_b12x_block_fp8_linear", linear_warmup) - monkeypatch.setattr(warmup_mod, "warmup_b12x_mxfp4_linear", linear_warmup) - monkeypatch.setattr(warmup_mod, "warmup_b12x_mxfp8_linear", linear_warmup) - monkeypatch.setattr(warmup_mod, "warmup_b12x_nvfp4_linear", linear_warmup) - monkeypatch.setattr(warmup_mod, "warmup_b12x_tensor_fp8_linear", linear_warmup) - - worker = SimpleNamespace( - get_model=lambda: model, - scheduler_config=SimpleNamespace( - max_num_batched_tokens=256, - max_num_scheduled_tokens=320, - ), - model_config=SimpleNamespace(dtype=torch.float16), - vllm_config=SimpleNamespace( - compilation_config=SimpleNamespace(compile_sizes=[32, "dynamic", 64]) - ), + synchronized = [] + monkeypatch.setattr(warmup_mod, "current_platform", platform) + monkeypatch.setattr( + warmup_mod.torch.accelerator, + "synchronize", + lambda: synchronized.append(True), ) - b12x_warmup(worker, [8, 16]) + b12x_warmup(worker, [1, 2]) - expected_linear_kwargs = { - "max_tokens": 256, - "cudagraph_capture_sizes": [8, 16], - "output_dtype": torch.float16, - } - assert linear_calls == [(model, expected_linear_kwargs)] * 5 + assert scans == 1 + assert calls == [ + ("first", (1, 2, 8), torch.bfloat16), + ("second", (1, 2, 8), torch.bfloat16), + ] + assert synchronized == [True] diff --git a/vllm/model_executor/kernels/linear/__init__.py b/vllm/model_executor/kernels/linear/__init__.py index 3413d4a69f..3ba499155b 100644 --- a/vllm/model_executor/kernels/linear/__init__.py +++ b/vllm/model_executor/kernels/linear/__init__.py @@ -172,10 +172,8 @@ from vllm.model_executor.kernels.linear.scaled_mm.aiter import ( AiterPerTokenFp8ScaledMMLinearKernel, AiterPreshuffledPerTokenFp8ScaledMMLinearKernel, ) -from vllm.model_executor.kernels.linear.scaled_mm.b12x_block import ( +from vllm.model_executor.kernels.linear.scaled_mm.b12x import ( B12xFp8BlockScaledMMKernel, -) -from vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor import ( B12xTensorFP8ScaledMMLinearKernel, ) from vllm.model_executor.kernels.linear.scaled_mm.cpu import ( diff --git a/vllm/model_executor/kernels/linear/mxfp4/b12x.py b/vllm/model_executor/kernels/linear/mxfp4/b12x.py index 93e872c068..a4649d3c4f 100644 --- a/vllm/model_executor/kernels/linear/mxfp4/b12x.py +++ b/vllm/model_executor/kernels/linear/mxfp4/b12x.py @@ -3,9 +3,6 @@ from __future__ import annotations -from collections.abc import Iterable -from typing import Any - import torch from vllm.model_executor.layers.quantization.utils.quant_utils import ( @@ -13,9 +10,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( ) from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform -from vllm.utils.b12x import ( - b12x_warmup_token_counts, -) +from vllm.utils.b12x import B12xWarmupUnit from vllm.utils.b12x import ( get_b12x_blockscaled as _import_b12x_blockscaled, ) @@ -51,73 +46,6 @@ def _apply_b12x_mxfp4_linear( return output.view(*output_shape) -def warmup_b12x_mxfp4_linear( - model: torch.nn.Module, - *, - max_tokens: int, - cudagraph_capture_sizes: Iterable[int] = (), - output_dtype: torch.dtype = torch.bfloat16, -) -> int: - if not current_platform.is_cuda(): - return 0 - if not current_platform.is_device_capability_family(120): - return 0 - if output_dtype not in (torch.bfloat16, torch.float16): - output_dtype = torch.bfloat16 - layer_map: dict[tuple[Any, ...], torch.nn.Module] = {} - for layer in model.modules(): - if not getattr(layer, "b12x_mxfp4_linear", False): - continue - weight = layer.weight - weight_scale = layer.weight_scale - n, packed_k = map(int, weight.shape) - signature = ( - weight.device, - n, - packed_k * 2, - weight.dtype, - weight_scale.dtype, - output_dtype, - ) - layer_map.setdefault(signature, layer) - if not layer_map: - return 0 - if _import_b12x_blockscaled() is None: - return 0 - - token_counts = b12x_warmup_token_counts( - max_tokens=max_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes, - ) - warmed = 0 - last_device: torch.device | None = None - - with torch.inference_mode(): - for signature, layer in layer_map.items(): - weight = layer.weight - weight_scale = layer.weight_scale - k = signature[2] - last_device = weight.device - for tokens in token_counts: - source = torch.zeros( - (tokens, k), - dtype=output_dtype, - device=weight.device, - ) - _apply_b12x_mxfp4_linear( - source, - weight, - weight_scale, - None, - ) - warmed += 1 - - if warmed > 0 and last_device is not None and last_device.type == "cuda": - torch.accelerator.synchronize(last_device) - - return warmed - - class B12xMxFp4LinearKernel(MxFp4LinearKernel): """MXFP4 linear through the native B12X SM120 dense GEMM.""" @@ -151,7 +79,39 @@ class B12xMxFp4LinearKernel(MxFp4LinearKernel): "weight_scale", intrinsics.swizzle_block_scale(layer.weight_scale.data), ) - layer.b12x_mxfp4_linear = True + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + weight = layer.weight + weight_scale = layer.weight_scale + n, packed_k = map(int, weight.shape) + k = packed_k * 2 + + def compile() -> None: + for tokens in token_counts: + source = torch.zeros( + (tokens, k), dtype=output_dtype, device=weight.device + ) + _apply_b12x_mxfp4_linear(source, weight, weight_scale, None) + + return B12xWarmupUnit( + name="MXFP4", + key=( + type(self), + weight.device, + n, + k, + weight.dtype, + weight_scale.dtype, + output_dtype, + ), + compile=compile, + ) def apply_weights( self, @@ -167,4 +127,4 @@ class B12xMxFp4LinearKernel(MxFp4LinearKernel): ) -__all__ = ["B12xMxFp4LinearKernel", "warmup_b12x_mxfp4_linear"] +__all__ = ["B12xMxFp4LinearKernel"] diff --git a/vllm/model_executor/kernels/linear/mxfp8/b12x.py b/vllm/model_executor/kernels/linear/mxfp8/b12x.py index 4dd77dd6c0..48343f3d69 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/b12x.py +++ b/vllm/model_executor/kernels/linear/mxfp8/b12x.py @@ -3,9 +3,6 @@ from __future__ import annotations -from collections.abc import Iterable -from typing import Any - import torch from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( @@ -15,10 +12,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import ( ) from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform -from vllm.utils.b12x import ( - b12x_warmup_token_counts, - reuse_packed_weight_storage, -) +from vllm.utils.b12x import B12xWarmupUnit, reuse_packed_weight_storage from vllm.utils.b12x import ( get_b12x_mxfp8_linear as _import_b12x_mxfp8, ) @@ -27,21 +21,12 @@ from vllm.utils.torch_utils import current_stream from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig -def _b12x_mxfp8_expected_m(tokens: int) -> int: - return max(1, int(tokens)) - - def _apply_b12x_mxfp8_packed_linear( layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None, ) -> torch.Tensor: - packed_weight = getattr(layer, "b12x_mxfp8_packed_weight", None) - if packed_weight is None: - raise RuntimeError( - "b12x MXFP8 packed weights are missing; " - "process_weights_after_loading did not run for this layer" - ) + packed_weight = layer.b12x_mxfp8_packed_weight input_2d = x.reshape(-1, x.shape[-1]).contiguous() output_shape = [*x.shape[:-1], int(packed_weight.out_features)] @@ -52,77 +37,11 @@ def _apply_b12x_mxfp8_packed_linear( input_2d, packed_weight, bias=bias, - expected_m=_b12x_mxfp8_expected_m(int(input_2d.shape[0])), + expected_m=max(1, int(input_2d.shape[0])), ) return output.view(*output_shape) -def warmup_b12x_mxfp8_linear( - model: torch.nn.Module, - *, - max_tokens: int, - cudagraph_capture_sizes: Iterable[int] = (), - output_dtype: torch.dtype = torch.bfloat16, -) -> int: - if not current_platform.is_cuda(): - return 0 - if not current_platform.is_device_capability_family(120): - return 0 - if output_dtype not in (torch.bfloat16, torch.float16): - output_dtype = torch.bfloat16 - - layer_map: dict[tuple[Any, ...], Any] = {} - for layer in model.modules(): - packed_weight = getattr(layer, "b12x_mxfp8_packed_weight", None) - if packed_weight is None: - continue - device = torch.device(packed_weight.weight.values.device) - signature = ( - device, - int(packed_weight.in_features), - int(packed_weight.padded_in_features), - int(packed_weight.out_features), - output_dtype, - ) - layer_map.setdefault(signature, packed_weight) - if not layer_map: - return 0 - - mxfp8 = _import_b12x_mxfp8() - if mxfp8 is None: - return 0 - - token_counts = b12x_warmup_token_counts( - max_tokens=max_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes, - ) - warmed = 0 - last_device: torch.device | None = None - - with torch.inference_mode(): - for signature, packed_weight in layer_map.items(): - device = signature[0] - last_device = device - for tokens in token_counts: - source = torch.zeros( - (tokens, int(packed_weight.in_features)), - dtype=output_dtype, - device=device, - ) - mxfp8.mm( - source, - packed_weight, - expected_m=_b12x_mxfp8_expected_m(tokens), - stream=current_stream().cuda_stream, - ) - warmed += 1 - - if warmed > 0 and last_device is not None and last_device.type == "cuda": - torch.accelerator.synchronize(last_device) - - return warmed - - class B12xMxfp8LinearKernel(Mxfp8LinearKernel): """ModelOpt MXFP8 linear through the native b12x SM120 dense GEMM path.""" @@ -183,6 +102,45 @@ class B12xMxfp8LinearKernel(Mxfp8LinearKernel): ) replace_parameter(layer, "weight", weight.new_empty((0,))) replace_parameter(layer, "weight_scale", weight_scale.new_empty((0,))) + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + packed_weight = layer.b12x_mxfp8_packed_weight + device = torch.device(packed_weight.weight.values.device) + + def compile() -> None: + mxfp8 = _import_b12x_mxfp8() + assert mxfp8 is not None + for tokens in token_counts: + source = torch.zeros( + (tokens, int(packed_weight.in_features)), + dtype=output_dtype, + device=device, + ) + mxfp8.mm( + source, + packed_weight, + expected_m=max(1, int(tokens)), + stream=current_stream().cuda_stream, + ) + + return B12xWarmupUnit( + name="MXFP8", + key=( + type(self), + device, + int(packed_weight.in_features), + int(packed_weight.padded_in_features), + int(packed_weight.out_features), + output_dtype, + ), + compile=compile, + ) def apply_weights( self, diff --git a/vllm/model_executor/kernels/linear/nvfp4/b12x.py b/vllm/model_executor/kernels/linear/nvfp4/b12x.py index a9e5a69283..4ff531dd48 100644 --- a/vllm/model_executor/kernels/linear/nvfp4/b12x.py +++ b/vllm/model_executor/kernels/linear/nvfp4/b12x.py @@ -3,17 +3,12 @@ from __future__ import annotations -from collections.abc import Iterable -from typing import Any - import torch from vllm._custom_ops import scaled_fp4_quant from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform -from vllm.utils.b12x import ( - b12x_warmup_token_counts, -) +from vllm.utils.b12x import B12xWarmupUnit from vllm.utils.b12x import ( get_b12x_blockscaled as _import_b12x_blockscaled, ) @@ -54,75 +49,6 @@ def _apply_b12x_nvfp4_linear( return output.view(*output_shape) -def warmup_b12x_nvfp4_linear( - model: torch.nn.Module, - *, - max_tokens: int, - cudagraph_capture_sizes: Iterable[int] = (), - output_dtype: torch.dtype = torch.bfloat16, -) -> int: - if not current_platform.is_cuda(): - return 0 - if not current_platform.is_device_capability_family(120): - return 0 - if output_dtype not in (torch.bfloat16, torch.float16): - output_dtype = torch.bfloat16 - layer_map: dict[tuple[Any, ...], torch.nn.Module] = {} - for layer in model.modules(): - if not getattr(layer, "b12x_nvfp4_linear", False): - continue - weight = layer.weight - weight_scale = layer.weight_scale - n, packed_k = map(int, weight.shape) - signature = ( - weight.device, - n, - packed_k * 2, - weight.dtype, - weight_scale.dtype, - output_dtype, - ) - layer_map.setdefault(signature, layer) - if not layer_map: - return 0 - if _import_b12x_blockscaled() is None: - return 0 - - token_counts = b12x_warmup_token_counts( - max_tokens=max_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes, - ) - warmed = 0 - last_device: torch.device | None = None - - with torch.inference_mode(): - for signature, layer in layer_map.items(): - weight = layer.weight - weight_scale = layer.weight_scale - k = signature[2] - last_device = weight.device - for tokens in token_counts: - source = torch.zeros( - (tokens, k), - dtype=output_dtype, - device=weight.device, - ) - _apply_b12x_nvfp4_linear( - source, - weight, - weight_scale, - layer.input_global_scale_inv, - layer.alpha, - None, - ) - warmed += 1 - - if warmed > 0 and last_device is not None and last_device.type == "cuda": - torch.accelerator.synchronize(last_device) - - return warmed - - class B12xNvFp4LinearKernel(NvFp4LinearKernel): """ModelOpt NVFP4 linear through the native B12X SM120 dense GEMM.""" @@ -155,7 +81,46 @@ class B12xNvFp4LinearKernel(NvFp4LinearKernel): "weight_scale", intrinsics.swizzle_block_scale(layer.weight_scale.data), ) - layer.b12x_nvfp4_linear = True + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + weight = layer.weight + weight_scale = layer.weight_scale + n, packed_k = map(int, weight.shape) + k = packed_k * 2 + + def compile() -> None: + for tokens in token_counts: + source = torch.zeros( + (tokens, k), dtype=output_dtype, device=weight.device + ) + _apply_b12x_nvfp4_linear( + source, + weight, + weight_scale, + layer.input_global_scale_inv, + layer.alpha, + None, + ) + + return B12xWarmupUnit( + name="NVFP4", + key=( + type(self), + weight.device, + n, + k, + weight.dtype, + weight_scale.dtype, + output_dtype, + ), + compile=compile, + ) def apply_weights( self, @@ -173,4 +138,4 @@ class B12xNvFp4LinearKernel(NvFp4LinearKernel): ) -__all__ = ["B12xNvFp4LinearKernel", "warmup_b12x_nvfp4_linear"] +__all__ = ["B12xNvFp4LinearKernel"] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/b12x.py b/vllm/model_executor/kernels/linear/scaled_mm/b12x.py new file mode 100644 index 0000000000..62e839c729 --- /dev/null +++ b/vllm/model_executor/kernels/linear/scaled_mm/b12x.py @@ -0,0 +1,339 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from __future__ import annotations + +import torch + +from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + _upcast_e8m0_to_fp32, +) +from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape +from vllm.model_executor.utils import replace_parameter +from vllm.platforms import current_platform +from vllm.utils.b12x import B12xWarmupUnit, reuse_packed_weight_storage +from vllm.utils.b12x import ( + get_b12x_blockscaled as _import_b12x_blockscaled, +) +from vllm.utils.b12x import ( + get_b12x_tensor_fp8_linear as _import_b12x_tensor_fp8, +) +from vllm.utils.torch_utils import current_stream + +from .BlockScaledMMLinearKernel import ( + Fp8BlockScaledMMLinearKernel, + FP8ScaledMMLinearLayerConfig, +) +from .ScaledMMLinearKernel import FP8ScaledMMLinearKernel + + +def _run_b12x_fp8_block_scaled_mm( + a: torch.Tensor, + weight: torch.Tensor, + a_scale: torch.Tensor, + weight_scale: torch.Tensor, + out_dtype: torch.dtype, +) -> torch.Tensor: + blockscaled = _import_b12x_blockscaled() + assert blockscaled is not None + + return blockscaled.mm_block_fp8( + a, + a_scale, + weight, + weight_scale, + out_dtype=out_dtype, + ) + + +class B12xFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): + """K128 block-FP8 linear through the native B12X SM120 dense GEMM.""" + + @classmethod + def is_supported( + cls, + compute_capability: int | None = None, + ) -> tuple[bool, str | None]: + del compute_capability + if not current_platform.is_cuda(): + return False, "B12X FP8 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "B12X FP8 kernels require a Blackwell 12x device" + blockscaled = _import_b12x_blockscaled() + if blockscaled is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not blockscaled.is_supported(): + return False, "B12X regular block-FP8 GEMM is not supported" + return True, None + + @classmethod + def can_implement( + cls, + config: FP8ScaledMMLinearLayerConfig, + ) -> tuple[bool, str | None]: + can_implement_base, reason = super().can_implement(config) + if not can_implement_base: + return can_implement_base, reason + + if config.input_dtype not in (torch.bfloat16, torch.float16): + return False, "Supports only bf16/fp16 input dtype" + if config.input_dtype != config.out_dtype: + return False, "Input and output dtype must match" + + act_group_shape = config.activation_quant_key.scale.group_shape + if act_group_shape != GroupShape(1, 128): + return ( + False, + "Supports only dynamic per-token group activation quantization " + "with group_shape=(1,128)", + ) + weight_group_shape = config.weight_quant_key.scale.group_shape + if weight_group_shape != GroupShape(128, 128): + return False, "Supports only 128x128 block-scaled FP8 weights" + + out_features, in_features = config.weight_shape + if in_features <= 0 or in_features % 128 != 0: + return False, "Input features must be a positive multiple of 128" + if out_features <= 0 or out_features % 128 != 0: + return False, "Output features must be a positive multiple of 128" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + super().process_weights_after_loading(layer) + params = self._get_layer_params(layer) + if params.weight_scale_inv is not None: + weight_scale = params.weight_scale_inv + scale_attr = params.WEIGHT_SCALE_INV + else: + weight_scale = params.weight_scale + scale_attr = params.WEIGHT_SCALE + if weight_scale is not None and weight_scale.dtype in ( + torch.float8_e8m0fnu, + torch.uint8, + ): + # TODO: Remove once B12X supports 128x128 UE8M0 block scales. + replace_parameter( + layer, + scale_attr, + _upcast_e8m0_to_fp32(weight_scale).contiguous(), + ) + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + weight = layer.weight + weight_scale = getattr(layer, "weight_scale_inv", None) + if weight_scale is None: + weight_scale = layer.weight_scale + n, k = map(int, weight.shape) + + def compile() -> None: + for tokens in token_counts: + a = torch.empty((tokens, k), dtype=weight.dtype, device=weight.device) + a_scale = torch.empty( + (tokens, k // 128), + dtype=torch.float32, + device=weight.device, + ) + _run_b12x_fp8_block_scaled_mm( + a, weight, a_scale, weight_scale, output_dtype + ) + + return B12xWarmupUnit( + name="block-FP8", + key=( + type(self), + weight.device, + n, + k, + weight.dtype, + weight_scale.dtype, + output_dtype, + ), + compile=compile, + ) + + def apply_block_scaled_mm( + self, + A: torch.Tensor, + B: torch.Tensor, + As: torch.Tensor, + Bs: torch.Tensor, + ) -> torch.Tensor: + return _run_b12x_fp8_block_scaled_mm( + A, + B, + As, + Bs, + self.config.out_dtype, + ) + + +def _apply_b12x_tensor_fp8_packed_linear( + layer: torch.nn.Module, + x_q: torch.Tensor, + bias: torch.Tensor | None, + out_dtype: torch.dtype, +) -> torch.Tensor: + packed_weight = layer.b12x_tensor_fp8_packed_weight + + tensor_fp8 = _import_b12x_tensor_fp8() + assert tensor_fp8 is not None + + input_2d = x_q.reshape(-1, x_q.shape[-1]).contiguous() + output_shape = [*x_q.shape[:-1], int(packed_weight.out_features)] + output = tensor_fp8.mm( + input_2d, + packed_weight, + bias=bias, + out_dtype=out_dtype, + expected_m=max(1, int(input_2d.shape[0])), + ) + return output.view(*output_shape) + + +class B12xTensorFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): + """Static per-tensor FP8 linear through the B12X SM12x dense GEMM.""" + + @classmethod + def is_supported( + cls, + compute_capability: int | None = None, + ) -> tuple[bool, str | None]: + del compute_capability + if not current_platform.is_cuda(): + return False, "b12x tensor FP8 kernels are only available on CUDA" + if not current_platform.is_device_capability_family(120): + return False, "b12x tensor FP8 kernels require a Blackwell 12x device" + tensor_fp8 = _import_b12x_tensor_fp8() + if tensor_fp8 is None: + return False, "Install the B12X backend with `pip install vllm[b12x]`" + if not tensor_fp8.is_supported(): + return False, "b12x.gemm.tensor_fp8_linear is not supported" + return True, None + + @classmethod + def can_implement( + cls, + config: FP8ScaledMMLinearLayerConfig, + ) -> tuple[bool, str | None]: + activation_scale = config.activation_quant_key.scale + weight_scale = config.weight_quant_key.scale + if ( + not activation_scale.static + or not activation_scale.group_shape.is_per_tensor() + ): + return False, "requires static per-tensor activation scales" + if not weight_scale.static or not weight_scale.group_shape.is_per_tensor(): + return False, "requires static per-tensor weight scales" + if config.input_dtype not in (torch.bfloat16, torch.float16): + return False, "supports only bf16/fp16 input dtype" + if config.out_dtype not in (torch.bfloat16, torch.float16): + return False, "supports only bf16/fp16 output dtype" + out_features, in_features = config.weight_shape + if out_features <= 0 or in_features <= 0 or in_features % 32 != 0: + return False, "weight dimensions must be positive with K divisible by 32" + return True, None + + def process_weights_after_loading(self, layer: torch.nn.Module) -> None: + weight, weight_scale, input_scale, _ = self._get_layer_params(layer) + assert weight.dtype == torch.float8_e4m3fn + assert input_scale is not None + assert weight_scale.numel() == input_scale.numel() == 1 + + out_features, in_features = map(int, self.config.weight_shape) + assert tuple(weight.shape) == (in_features, out_features) + + tensor_fp8 = _import_b12x_tensor_fp8() + assert tensor_fp8 is not None + output_scale = ( + input_scale.detach().to(torch.float32).reshape(1) + * weight_scale.detach().to(torch.float32).reshape(1) + ).contiguous() + packed_weight = tensor_fp8.pack_weight( + weight.detach().T.contiguous(), + output_scale, + ) + layer.b12x_tensor_fp8_packed_weight = reuse_packed_weight_storage( + getattr(layer, "b12x_tensor_fp8_packed_weight", None), + packed_weight, + ) + weight_name, weight_scale_name, _, _ = self.layer_param_names + replace_parameter(layer, weight_name, weight.new_empty((0,))) + replace_parameter(layer, weight_scale_name, weight_scale.new_empty((0,))) + layer.b12x_warmup_provider = self + + def get_b12x_warmup_unit( + self, + layer: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, + ) -> B12xWarmupUnit: + packed_weight = layer.b12x_tensor_fp8_packed_weight + device = torch.device(packed_weight.values.device) + + def compile() -> None: + tensor_fp8 = _import_b12x_tensor_fp8() + assert tensor_fp8 is not None + tensor_fp8.prewarm( + packed_weight, + token_counts, + out_dtype=output_dtype, + stream=current_stream().cuda_stream, + ) + + return B12xWarmupUnit( + name="tensor FP8", + key=( + type(self), + device, + int(packed_weight.in_features), + int(packed_weight.padded_in_features), + int(packed_weight.out_features), + output_dtype, + ), + compile=compile, + ) + + def apply_weights( + self, + layer: torch.nn.Module, + x: torch.Tensor, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + assert isinstance(x, torch.Tensor) + _, _, input_scale, input_scale_ub = self._get_layer_params(layer) + input_2d = x.reshape(-1, x.shape[-1]) + x_q, _ = self.quant_fp8(input_2d, input_scale, input_scale_ub) + out_dtype = self.config.out_dtype + output = _apply_b12x_tensor_fp8_packed_linear( + layer, + x_q, + bias, + out_dtype, + ) + return output.view(*x.shape[:-1], output.shape[-1]) + + def apply_scaled_mm( + self, + *, + A: torch.Tensor, + B: torch.Tensor, + out_dtype: torch.dtype, + As: torch.Tensor, + Bs: torch.Tensor, + bias: torch.Tensor | None, + output_shape: list, + ) -> torch.Tensor: + del A, B, out_dtype, As, Bs, bias, output_shape + raise NotImplementedError("b12x tensor FP8 linear overrides apply_weights") + + +__all__ = [ + "B12xFp8BlockScaledMMKernel", + "B12xTensorFP8ScaledMMLinearKernel", +] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/b12x_block.py b/vllm/model_executor/kernels/linear/scaled_mm/b12x_block.py deleted file mode 100644 index c2d12549fb..0000000000 --- a/vllm/model_executor/kernels/linear/scaled_mm/b12x_block.py +++ /dev/null @@ -1,214 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -from collections.abc import Iterable -from typing import Any - -import torch - -from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - _upcast_e8m0_to_fp32, -) -from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape -from vllm.model_executor.utils import replace_parameter -from vllm.platforms import current_platform -from vllm.utils.b12x import ( - b12x_warmup_token_counts, -) -from vllm.utils.b12x import ( - get_b12x_blockscaled as _import_b12x_blockscaled, -) - -from .BlockScaledMMLinearKernel import ( - Fp8BlockScaledMMLinearKernel, - FP8ScaledMMLinearLayerConfig, -) - - -def _run_b12x_fp8_block_scaled_mm( - a: torch.Tensor, - weight: torch.Tensor, - a_scale: torch.Tensor, - weight_scale: torch.Tensor, - out_dtype: torch.dtype, -) -> torch.Tensor: - blockscaled = _import_b12x_blockscaled() - assert blockscaled is not None - - return blockscaled.mm_block_fp8( - a, - a_scale, - weight, - weight_scale, - out_dtype=out_dtype, - ) - - -def warmup_b12x_block_fp8_linear( - model: torch.nn.Module, - *, - max_tokens: int, - cudagraph_capture_sizes: Iterable[int] = (), - output_dtype: torch.dtype = torch.bfloat16, -) -> int: - if not current_platform.is_cuda(): - return 0 - if not current_platform.is_device_capability_family(120): - return 0 - if output_dtype not in (torch.bfloat16, torch.float16): - output_dtype = torch.bfloat16 - - layer_map: dict[tuple[Any, ...], torch.nn.Module] = {} - for layer in model.modules(): - if not getattr(layer, "b12x_block_fp8_linear", False): - continue - weight = layer.weight - weight_scale = getattr(layer, "weight_scale_inv", None) - if weight_scale is None: - weight_scale = layer.weight_scale - n, k = map(int, weight.shape) - signature = ( - weight.device, - n, - k, - weight.dtype, - weight_scale.dtype, - output_dtype, - ) - layer_map.setdefault(signature, layer) - if not layer_map: - return 0 - - blockscaled = _import_b12x_blockscaled() - if blockscaled is None: - return 0 - token_counts = b12x_warmup_token_counts( - max_tokens=max_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes, - ) - warmed = 0 - - with torch.inference_mode(): - for signature, layer in layer_map.items(): - weight = layer.weight - weight_scale = getattr(layer, "weight_scale_inv", None) - if weight_scale is None: - weight_scale = layer.weight_scale - k = signature[2] - for tokens in token_counts: - a = torch.empty( - (tokens, k), - dtype=weight.dtype, - device=weight.device, - ) - a_scale = torch.empty( - (tokens, k // 128), - dtype=torch.float32, - device=weight.device, - ) - _run_b12x_fp8_block_scaled_mm( - a, - weight, - a_scale, - weight_scale, - output_dtype, - ) - warmed += 1 - return warmed - - -class B12xFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel): - """K128 block-FP8 linear through the native B12X SM120 dense GEMM.""" - - @classmethod - def is_supported( - cls, - compute_capability: int | None = None, - ) -> tuple[bool, str | None]: - del compute_capability - if not current_platform.is_cuda(): - return False, "B12X FP8 kernels are only available on CUDA" - if not current_platform.is_device_capability_family(120): - return False, "B12X FP8 kernels require a Blackwell 12x device" - blockscaled = _import_b12x_blockscaled() - if blockscaled is None: - return False, "Install the B12X backend with `pip install vllm[b12x]`" - if not blockscaled.is_supported(): - return False, "B12X regular block-FP8 GEMM is not supported" - return True, None - - @classmethod - def can_implement( - cls, - config: FP8ScaledMMLinearLayerConfig, - ) -> tuple[bool, str | None]: - can_implement_base, reason = super().can_implement(config) - if not can_implement_base: - return can_implement_base, reason - - if config.input_dtype not in (torch.bfloat16, torch.float16): - return False, "Supports only bf16/fp16 input dtype" - if config.input_dtype != config.out_dtype: - return False, "Input and output dtype must match" - - act_group_shape = config.activation_quant_key.scale.group_shape - if act_group_shape != GroupShape(1, 128): - return ( - False, - "Supports only dynamic per-token group activation quantization " - "with group_shape=(1,128)", - ) - weight_group_shape = config.weight_quant_key.scale.group_shape - if weight_group_shape != GroupShape(128, 128): - return False, "Supports only 128x128 block-scaled FP8 weights" - - out_features, in_features = config.weight_shape - if in_features <= 0 or in_features % 128 != 0: - return False, "Input features must be a positive multiple of 128" - if out_features <= 0 or out_features % 128 != 0: - return False, "Output features must be a positive multiple of 128" - return True, None - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - super().process_weights_after_loading(layer) - params = self._get_layer_params(layer) - if params.weight_scale_inv is not None: - weight_scale = params.weight_scale_inv - scale_attr = params.WEIGHT_SCALE_INV - else: - weight_scale = params.weight_scale - scale_attr = params.WEIGHT_SCALE - if weight_scale is not None and weight_scale.dtype in ( - torch.float8_e8m0fnu, - torch.uint8, - ): - # TODO: Remove once B12X supports 128x128 UE8M0 block scales. - replace_parameter( - layer, - scale_attr, - _upcast_e8m0_to_fp32(weight_scale).contiguous(), - ) - layer.b12x_block_fp8_linear = True - - def apply_block_scaled_mm( - self, - A: torch.Tensor, - B: torch.Tensor, - As: torch.Tensor, - Bs: torch.Tensor, - ) -> torch.Tensor: - return _run_b12x_fp8_block_scaled_mm( - A, - B, - As, - Bs, - self.config.out_dtype, - ) - - -__all__ = [ - "B12xFp8BlockScaledMMKernel", - "warmup_b12x_block_fp8_linear", -] diff --git a/vllm/model_executor/kernels/linear/scaled_mm/b12x_tensor.py b/vllm/model_executor/kernels/linear/scaled_mm/b12x_tensor.py deleted file mode 100644 index 46b27a7d4b..0000000000 --- a/vllm/model_executor/kernels/linear/scaled_mm/b12x_tensor.py +++ /dev/null @@ -1,238 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -from __future__ import annotations - -from collections.abc import Iterable -from typing import Any - -import torch - -from vllm.model_executor.utils import replace_parameter -from vllm.platforms import current_platform -from vllm.utils.b12x import ( - b12x_warmup_token_counts, - reuse_packed_weight_storage, -) -from vllm.utils.b12x import ( - get_b12x_tensor_fp8_linear as _import_b12x_tensor_fp8, -) -from vllm.utils.torch_utils import current_stream - -from .ScaledMMLinearKernel import ( - FP8ScaledMMLinearKernel, - FP8ScaledMMLinearLayerConfig, -) - - -def _apply_b12x_tensor_fp8_packed_linear( - layer: torch.nn.Module, - x_q: torch.Tensor, - bias: torch.Tensor | None, - out_dtype: torch.dtype, -) -> torch.Tensor: - packed_weight = getattr(layer, "b12x_tensor_fp8_packed_weight", None) - if packed_weight is None: - raise RuntimeError( - "b12x tensor FP8 packed weights are missing; " - "process_weights_after_loading did not run for this layer" - ) - - tensor_fp8 = _import_b12x_tensor_fp8() - assert tensor_fp8 is not None - - input_2d = x_q.reshape(-1, x_q.shape[-1]).contiguous() - output_shape = [*x_q.shape[:-1], int(packed_weight.out_features)] - output = tensor_fp8.mm( - input_2d, - packed_weight, - bias=bias, - out_dtype=out_dtype, - expected_m=max(1, int(input_2d.shape[0])), - ) - return output.view(*output_shape) - - -def warmup_b12x_tensor_fp8_linear( - model: torch.nn.Module, - *, - max_tokens: int, - cudagraph_capture_sizes: Iterable[int] = (), - output_dtype: torch.dtype = torch.bfloat16, -) -> int: - if not current_platform.is_cuda(): - return 0 - if not current_platform.is_device_capability_family(120): - return 0 - - if output_dtype not in (torch.bfloat16, torch.float16): - output_dtype = torch.bfloat16 - - layer_map: dict[tuple[Any, ...], Any] = {} - for layer in model.modules(): - packed_weight = getattr( - layer, - "b12x_tensor_fp8_packed_weight", - None, - ) - if packed_weight is None: - continue - device = torch.device(packed_weight.values.device) - signature = ( - device, - int(packed_weight.in_features), - int(packed_weight.padded_in_features), - int(packed_weight.out_features), - output_dtype, - ) - layer_map.setdefault(signature, packed_weight) - if not layer_map: - return 0 - - tensor_fp8 = _import_b12x_tensor_fp8() - if tensor_fp8 is None: - return 0 - - token_counts = b12x_warmup_token_counts( - max_tokens=max_tokens, - cudagraph_capture_sizes=cudagraph_capture_sizes, - ) - warmed = 0 - last_device: torch.device | None = None - - with torch.inference_mode(): - for signature, packed_weight in layer_map.items(): - last_device = signature[0] - warmed += int( - tensor_fp8.prewarm( - packed_weight, - token_counts, - out_dtype=output_dtype, - stream=current_stream().cuda_stream, - ) - ) - - if warmed > 0 and last_device is not None and last_device.type == "cuda": - torch.accelerator.synchronize(last_device) - - return warmed - - -class B12xTensorFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel): - """Static per-tensor FP8 linear through the B12X SM12x dense GEMM.""" - - @classmethod - def is_supported( - cls, - compute_capability: int | None = None, - ) -> tuple[bool, str | None]: - del compute_capability - if not current_platform.is_cuda(): - return False, "b12x tensor FP8 kernels are only available on CUDA" - if not current_platform.is_device_capability_family(120): - return False, "b12x tensor FP8 kernels require a Blackwell 12x device" - tensor_fp8 = _import_b12x_tensor_fp8() - if tensor_fp8 is None: - return False, "Install the B12X backend with `pip install vllm[b12x]`" - if not tensor_fp8.is_supported(): - return False, "b12x.gemm.tensor_fp8_linear is not supported" - return True, None - - @classmethod - def can_implement( - cls, - config: FP8ScaledMMLinearLayerConfig, - ) -> tuple[bool, str | None]: - activation_scale = config.activation_quant_key.scale - weight_scale = config.weight_quant_key.scale - if ( - not activation_scale.static - or not activation_scale.group_shape.is_per_tensor() - ): - return False, "requires static per-tensor activation scales" - if not weight_scale.static or not weight_scale.group_shape.is_per_tensor(): - return False, "requires static per-tensor weight scales" - if config.input_dtype not in (torch.bfloat16, torch.float16): - return False, "supports only bf16/fp16 input dtype" - if config.out_dtype not in (torch.bfloat16, torch.float16): - return False, "supports only bf16/fp16 output dtype" - out_features, in_features = config.weight_shape - if out_features <= 0 or in_features <= 0 or in_features % 32 != 0: - return False, "weight dimensions must be positive with K divisible by 32" - return True, None - - def process_weights_after_loading(self, layer: torch.nn.Module) -> None: - weight, weight_scale, input_scale, _ = self._get_layer_params(layer) - if weight.dtype != torch.float8_e4m3fn: - raise ValueError( - f"b12x tensor FP8 requires float8_e4m3fn weight, got {weight.dtype}" - ) - if weight_scale.numel() != 1 or input_scale is None or input_scale.numel() != 1: - raise ValueError( - "b12x tensor FP8 requires scalar weight and activation scales" - ) - - out_features, in_features = map(int, self.config.weight_shape) - if tuple(weight.shape) != (in_features, out_features): - raise ValueError( - "b12x tensor FP8 expects the processed weight in [K,N] layout, " - f"got {tuple(weight.shape)} for N={out_features}, K={in_features}" - ) - - tensor_fp8 = _import_b12x_tensor_fp8() - assert tensor_fp8 is not None - output_scale = ( - input_scale.detach().to(torch.float32).reshape(1) - * weight_scale.detach().to(torch.float32).reshape(1) - ).contiguous() - packed_weight = tensor_fp8.pack_weight( - weight.detach().T.contiguous(), - output_scale, - ) - layer.b12x_tensor_fp8_packed_weight = reuse_packed_weight_storage( - getattr(layer, "b12x_tensor_fp8_packed_weight", None), - packed_weight, - ) - weight_name, weight_scale_name, _, _ = self.layer_param_names - replace_parameter(layer, weight_name, weight.new_empty((0,))) - replace_parameter(layer, weight_scale_name, weight_scale.new_empty((0,))) - - def apply_weights( - self, - layer: torch.nn.Module, - x: torch.Tensor, - bias: torch.Tensor | None = None, - ) -> torch.Tensor: - if not isinstance(x, torch.Tensor): - raise TypeError("b12x tensor FP8 linear requires a Tensor input") - _, _, input_scale, input_scale_ub = self._get_layer_params(layer) - input_2d = x.reshape(-1, x.shape[-1]) - x_q, _ = self.quant_fp8(input_2d, input_scale, input_scale_ub) - out_dtype = self.config.out_dtype - output = _apply_b12x_tensor_fp8_packed_linear( - layer, - x_q, - bias, - out_dtype, - ) - return output.view(*x.shape[:-1], output.shape[-1]) - - def apply_scaled_mm( - self, - *, - A: torch.Tensor, - B: torch.Tensor, - out_dtype: torch.dtype, - As: torch.Tensor, - Bs: torch.Tensor, - bias: torch.Tensor | None, - output_shape: list, - ) -> torch.Tensor: - del A, B, out_dtype, As, Bs, bias, output_shape - raise NotImplementedError("b12x tensor FP8 linear overrides apply_weights") - - -__all__ = [ - "B12xTensorFP8ScaledMMLinearKernel", - "warmup_b12x_tensor_fp8_linear", -] diff --git a/vllm/model_executor/warmup/b12x_warmup.py b/vllm/model_executor/warmup/b12x_warmup.py index daaa078628..5e0e4a0fa4 100644 --- a/vllm/model_executor/warmup/b12x_warmup.py +++ b/vllm/model_executor/warmup/b12x_warmup.py @@ -2,26 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Warm B12X JIT kernels used by a loaded model.""" +from collections import Counter +from collections.abc import Iterable from typing import TYPE_CHECKING import torch from vllm.logger import init_logger -from vllm.model_executor.kernels.linear.mxfp4.b12x import ( - warmup_b12x_mxfp4_linear, -) -from vllm.model_executor.kernels.linear.mxfp8.b12x import ( - warmup_b12x_mxfp8_linear, -) -from vllm.model_executor.kernels.linear.nvfp4.b12x import ( - warmup_b12x_nvfp4_linear, -) -from vllm.model_executor.kernels.linear.scaled_mm.b12x_block import ( - warmup_b12x_block_fp8_linear, -) -from vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor import ( - warmup_b12x_tensor_fp8_linear, -) +from vllm.platforms import current_platform +from vllm.utils.b12x import B12xWarmupUnit, b12x_warmup_token_counts if TYPE_CHECKING: from vllm.v1.worker.gpu_worker import Worker @@ -29,32 +18,61 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def _collect_warmup_units( + model: torch.nn.Module, + token_counts: tuple[int, ...], + output_dtype: torch.dtype, +) -> Iterable[B12xWarmupUnit]: + units: dict[object, B12xWarmupUnit] = {} + for layer in model.modules(): + provider = getattr(layer, "b12x_warmup_provider", None) + get_unit = getattr(provider, "get_b12x_warmup_unit", None) + if not callable(get_unit): + continue + unit = get_unit(layer, token_counts, output_dtype) + assert isinstance(unit, B12xWarmupUnit) + units.setdefault(unit.key, unit) + return units.values() + + +def _compile_warmup_units( + units: Iterable[B12xWarmupUnit], +) -> Counter[str]: + warmed: Counter[str] = Counter() + with torch.inference_mode(): + for unit in units: + unit.compile() + warmed[unit.name] += 1 + if warmed: + torch.accelerator.synchronize() + return warmed + + def b12x_warmup(worker: "Worker", cudagraph_capture_sizes: list[int]) -> None: - model = worker.get_model() - max_tokens = worker.scheduler_config.max_num_batched_tokens + if not current_platform.is_cuda(): + return + if not current_platform.is_device_capability_family(120): + return + output_dtype = getattr( getattr(worker, "model_config", None), "dtype", torch.bfloat16, ) - - warmup_kwargs = { - "max_tokens": max_tokens, - "cudagraph_capture_sizes": cudagraph_capture_sizes, - "output_dtype": output_dtype, - } - providers = ( - ("block-FP8", warmup_b12x_block_fp8_linear), - ("MXFP8", warmup_b12x_mxfp8_linear), - ("tensor FP8", warmup_b12x_tensor_fp8_linear), - ("MXFP4", warmup_b12x_mxfp4_linear), - ("NVFP4", warmup_b12x_nvfp4_linear), + if output_dtype not in (torch.bfloat16, torch.float16): + output_dtype = torch.bfloat16 + token_counts = b12x_warmup_token_counts( + max_tokens=worker.scheduler_config.max_num_batched_tokens, + cudagraph_capture_sizes=cudagraph_capture_sizes, ) - for name, warmup in providers: - warmed = warmup(model, **warmup_kwargs) - if warmed: - logger.info_once( - "Warmed up %d B12X %s linear GEMM signatures.", - warmed, - name, - ) + units = _collect_warmup_units( + worker.get_model(), + token_counts, + output_dtype, + ) + for name, count in _compile_warmup_units(units).items(): + logger.info_once( + "Warmed up %d B12X %s linear GEMM signatures.", + count, + name, + ) diff --git a/vllm/utils/b12x.py b/vllm/utils/b12x.py index 55afa35f09..8e2d7e938b 100644 --- a/vllm/utils/b12x.py +++ b/vllm/utils/b12x.py @@ -5,14 +5,21 @@ import functools import importlib import importlib.util -from collections.abc import Iterable -from dataclasses import fields, is_dataclass +from collections.abc import Callable, Hashable, Iterable +from dataclasses import dataclass, fields, is_dataclass from types import ModuleType from typing import Any import torch +@dataclass(frozen=True) +class B12xWarmupUnit: + name: str + key: Hashable + compile: Callable[[], None] + + @functools.cache def has_b12x() -> bool: """Return whether the B12X package is installed."""