[Refactor] Simplify B12X linear kernels and warmup (#52368)
Signed-off-by: mgoin <mgoin64@gmail.com>
This commit is contained in:
Misha Goin
@@ -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
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
+249
-373
@@ -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}
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
@@ -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]
|
||||
|
||||
@@ -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 (
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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"]
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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",
|
||||
]
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
+9
-2
@@ -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."""
|
||||
|
||||
Reference in New Issue
Block a user