[Kernel] Add B12X dense linear backends (#52016)
Signed-off-by: Luke Alonso <lalonso@gmail.com>
Luke Alonso committed
Aug 13, 2026 at 19:48 UTC
3c79b1a8bf7157d9c620c0c56779199c0721f196
17 files changed
+2753
-3
docs/features/quantization/b12x.md
new
+24
@@ -0,0 +1,24 @@
1
+# B12X Linear Backend
2
+
3
+[B12X](https://pypi.org/project/b12x/) provides optional CUDA kernels for
4
+NVIDIA SM120 and SM121 GPUs. Install the dependency with:
5
+
6
+```bash
7
+uv pip install "vllm[b12x]"
8
+```
9
+
10
+B12X participates in automatic kernel selection after established optimized
11
+backends and before emulation. Select it explicitly with:
12
+
13
+```bash
14
+vllm serve <model> --linear-backend b12x
15
+```
16
+
17
+## Supported Configurations
18
+
19
+| Backend | Supported configurations |
20
+| ------- | ------------------------ |
21
+| Linear | Per-tensor FP8, 128x128 block FP8, MXFP8, NVFP4, and MXFP4 |
22
+
23
+Dense W4A16 layers are not handled by B12X and continue to use another
24
+compatible backend such as Marlin.
setup.py
+1
@@ -1291,6 +1291,7 @@ setup(
1291
# only; also needs system GStreamer + libv4l (see docs).
1292
"deepstream": ["nvidia-deepstream-videodecode-cu13>=9.0.2"],
1293
"flashinfer": [], # Kept for backwards compatibility
1294
+ "b12x": ["b12x==1.2.4"],
1295
# Optional deps for Helion kernel development
1296
# NOTE: When updating helion version, also update CI files:
1297
# - .buildkite/test_areas/kernels.yaml
tests/kernels/quantization/test_block_fp8.py
+55
@@ -14,6 +14,10 @@ from tests.kernels.quant_utils import (
14
)
15
from tests.kernels.utils import fp8_ulp_distance
16
from vllm.config import VllmConfig
17
+from vllm.model_executor.kernels.linear.scaled_mm.b12x_block import (
18
+ B12xFp8BlockScaledMMKernel,
19
+ _run_b12x_fp8_block_scaled_mm,
20
+)
21
from vllm.model_executor.kernels.linear.scaled_mm.cutlass import cutlass_scaled_mm
22
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
23
per_token_group_quant_fp8,
@@ -353,3 +357,54 @@ def test_w8a8_block_fp8_flashinfer_matmul(M, N, K, block_size, out_dtype, seed):
357
torch.abs(out.to(torch.bfloat16) - ref_out.to(torch.bfloat16))
358
) / torch.mean(torch.abs(ref_out.to(torch.bfloat16)))
359
assert rel_diff < 0.001
360
+
361
+
362
+@pytest.mark.parametrize(
363
+ "M,N,K",
364
+ [(1, 128, 256), (8, 256, 512), (129, 256, 256), (2, 4096, 4096)],
365
+)
366
+@torch.inference_mode()
367
+def test_w8a8_block_fp8_b12x_matmul(M, N, K):
368
+ supported, reason = B12xFp8BlockScaledMMKernel.is_supported()
369
+ if not supported:
370
+ pytest.skip(reason)
371
+
372
+ torch.manual_seed(M)
373
+ fp8_max = torch.finfo(torch.float8_e4m3fn).max
374
+ A_bf16 = (torch.rand(M, K, dtype=torch.bfloat16) - 0.5) * 2 * fp8_max
375
+ B_bf16 = (torch.rand(N, K, dtype=torch.bfloat16) - 0.5) * 2 * fp8_max
376
+ A_fp8, As = per_token_group_quant_fp8(A_bf16, 128, use_ue8m0=False)
377
+ B_fp8, Bs = per_block_cast_to_fp8(
378
+ B_bf16,
379
+ block_size=[128, 128],
380
+ use_ue8m0=False,
381
+ )
382
+ As = As.float()
383
+ Bs = Bs.float()
384
+
385
+ ref_out = native_w8a8_block_matmul(
386
+ A_fp8,
387
+ B_fp8,
388
+ As,
389
+ Bs,
390
+ [128, 128],
391
+ torch.bfloat16,
392
+ )
393
+ out = _run_b12x_fp8_block_scaled_mm(
394
+ A_fp8,
395
+ B_fp8,
396
+ As,
397
+ Bs,
398
+ torch.bfloat16,
399
+ )
400
+
401
+ rel_diff = torch.mean(torch.abs(out.float() - ref_out.float())) / torch.mean(
402
+ torch.abs(ref_out.float())
403
+ )
404
+ cosine = torch.nn.functional.cosine_similarity(
405
+ out.float().flatten(),
406
+ ref_out.float().flatten(),
407
+ dim=0,
408
+ )
409
+ assert rel_diff < 0.002
410
+ assert cosine >= 0.9999
tests/model_executor/kernels/test_b12x_mxfp4_linear.py
new
+199
@@ -0,0 +1,199 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from __future__ import annotations
5
+
6
+import types
7
+
8
+import torch
9
+
10
+from vllm.model_executor.kernels.linear import (
11
+ _LINEAR_BACKEND_KERNEL_MAP,
12
+ _POSSIBLE_MXFP4_KERNELS,
13
+ init_mxfp4_linear_kernel,
14
+)
15
+from vllm.model_executor.kernels.linear.mxfp4.b12x import (
16
+ B12xMxFp4LinearKernel,
17
+ warmup_b12x_mxfp4_linear,
18
+)
19
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
20
+ kMxfp4Dynamic,
21
+)
22
+from vllm.platforms import PlatformEnum
23
+
24
+
25
+def test_b12x_backend_maps_mxfp4_kernel() -> None:
26
+ assert B12xMxFp4LinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"]
27
+ assert B12xMxFp4LinearKernel in _POSSIBLE_MXFP4_KERNELS[PlatformEnum.CUDA]
28
+
29
+
30
+def test_b12x_mxfp4_fallback_priority() -> None:
31
+ kernels = _POSSIBLE_MXFP4_KERNELS[PlatformEnum.CUDA]
32
+ names = [kernel.__name__ for kernel in kernels]
33
+
34
+ assert (
35
+ names.index("HummingMxFp4LinearKernel")
36
+ < names.index("B12xMxFp4LinearKernel")
37
+ < names.index("EmulationMxfp4LinearKernel")
38
+ )
39
+
40
+
41
+def test_b12x_mxfp4_explicit_backend_selects_native_kernel(monkeypatch) -> None:
42
+ import vllm.model_executor.kernels.linear as linear_mod
43
+
44
+ monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
45
+ monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
46
+ monkeypatch.setattr(
47
+ B12xMxFp4LinearKernel,
48
+ "is_supported",
49
+ classmethod(lambda cls, compute_capability=None: (True, None)),
50
+ )
51
+ monkeypatch.setattr(
52
+ B12xMxFp4LinearKernel,
53
+ "can_implement",
54
+ classmethod(lambda cls, config: (True, None)),
55
+ )
56
+
57
+ kernel = init_mxfp4_linear_kernel(activation_quant_key=kMxfp4Dynamic)
58
+
59
+ assert isinstance(kernel, B12xMxFp4LinearKernel)
60
+
61
+
62
+def test_b12x_mxfp4_requires_dynamic_activations() -> None:
63
+ config = types.SimpleNamespace(activation_quant_key=kMxfp4Dynamic)
64
+ can_implement, reason = B12xMxFp4LinearKernel.can_implement(config)
65
+
66
+ assert can_implement
67
+ assert reason is None
68
+
69
+ config.activation_quant_key = None
70
+ can_implement, reason = B12xMxFp4LinearKernel.can_implement(config)
71
+
72
+ assert not can_implement
73
+ assert reason == "B12X MXFP4 GEMM requires dynamic MXFP4 activations"
74
+
75
+
76
+def test_b12x_mxfp4_processes_scale_and_preserves_loader(monkeypatch) -> None:
77
+ import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod
78
+
79
+ scale = torch.empty((48, 8), dtype=torch.uint8)
80
+ swizzled_scale = torch.empty((128, 8), dtype=torch.uint8)
81
+ monkeypatch.setattr(
82
+ b12x_mod,
83
+ "_import_b12x_intrinsics",
84
+ lambda: types.SimpleNamespace(swizzle_block_scale=lambda value: swizzled_scale),
85
+ )
86
+
87
+ layer = torch.nn.Module()
88
+ layer.prefix = "model.layers.0.mlp.shared_expert.down_proj"
89
+ layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False)
90
+ weight_loader = object()
91
+ layer.weight_scale.weight_loader = weight_loader
92
+ kernel = object.__new__(B12xMxFp4LinearKernel)
93
+
94
+ kernel.process_weights_after_loading(layer)
95
+
96
+ assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr()
97
+ assert layer.weight_scale.weight_loader is weight_loader
98
+ assert layer.b12x_mxfp4_linear
99
+
100
+
101
+def test_b12x_mxfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None:
102
+ import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod
103
+ import vllm.utils.flashinfer as flashinfer_utils
104
+
105
+ calls: list[tuple] = []
106
+ x_packed = torch.empty((6, 64), dtype=torch.uint8)
107
+ x_scale_storage = torch.empty((128, 4), dtype=torch.uint8)
108
+
109
+ def mm_mxfp4(*args, **kwargs):
110
+ calls.append((args, kwargs))
111
+ return torch.full((6, 48), 3.0, dtype=torch.bfloat16)
112
+
113
+ monkeypatch.setattr(
114
+ flashinfer_utils,
115
+ "flashinfer_mxfp4_quantize",
116
+ lambda *args, **kwargs: (x_packed, x_scale_storage),
117
+ )
118
+ monkeypatch.setattr(
119
+ b12x_mod,
120
+ "_import_b12x_blockscaled",
121
+ lambda: types.SimpleNamespace(mm_mxfp4=mm_mxfp4),
122
+ )
123
+
124
+ layer = torch.nn.Module()
125
+ layer.output_size_per_partition = 48
126
+ layer.weight = torch.empty((48, 64), dtype=torch.uint8)
127
+ layer.weight_scale = torch.empty((128, 4), dtype=torch.uint8)
128
+ x = torch.empty((2, 3, 128), dtype=torch.bfloat16)
129
+ bias = torch.ones(48, dtype=torch.bfloat16)
130
+ kernel = object.__new__(B12xMxFp4LinearKernel)
131
+
132
+ output = kernel.apply_weights(layer, x, bias)
133
+
134
+ assert output.shape == (2, 3, 48)
135
+ torch.testing.assert_close(output, torch.full_like(output, 4.0))
136
+ assert len(calls) == 1
137
+ args, kwargs = calls[0]
138
+ assert args == (
139
+ x_packed,
140
+ x_scale_storage,
141
+ layer.weight,
142
+ layer.weight_scale,
143
+ )
144
+ assert kwargs == {"out_dtype": torch.bfloat16}
145
+
146
+
147
+def test_warmup_b12x_mxfp4_dedupes_weight_signatures(monkeypatch) -> None:
148
+ import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod
149
+
150
+ calls = []
151
+
152
+ def apply(source, weight, weight_scale, bias):
153
+ calls.append((source.shape, weight, weight_scale, bias))
154
+ return source.new_empty((source.shape[0], weight.shape[0]))
155
+
156
+ platform = types.SimpleNamespace(
157
+ is_cuda=lambda: True,
158
+ is_device_capability_family=lambda family: family == 120,
159
+ )
160
+ monkeypatch.setattr(b12x_mod, "current_platform", platform)
161
+ monkeypatch.setattr(
162
+ b12x_mod,
163
+ "_import_b12x_blockscaled",
164
+ lambda: types.SimpleNamespace(),
165
+ )
166
+ monkeypatch.setattr(b12x_mod, "_apply_b12x_mxfp4_linear", apply)
167
+
168
+ def layer(n: int):
169
+ return types.SimpleNamespace(
170
+ b12x_mxfp4_linear=True,
171
+ weight=torch.empty((n, 64), dtype=torch.uint8),
172
+ weight_scale=torch.empty((128, 4), dtype=torch.uint8),
173
+ )
174
+
175
+ layer_a = layer(48)
176
+ layer_b = layer(48)
177
+ layer_c = layer(96)
178
+ model = types.SimpleNamespace(
179
+ modules=lambda: iter([layer_a, layer_b, layer_c, types.SimpleNamespace()])
180
+ )
181
+
182
+ warmed = warmup_b12x_mxfp4_linear(
183
+ model,
184
+ max_tokens=8,
185
+ cudagraph_capture_sizes=[1, 2],
186
+ )
187
+
188
+ assert warmed == 6
189
+ assert [call[0] for call in calls] == [
190
+ torch.Size([1, 128]),
191
+ torch.Size([2, 128]),
192
+ torch.Size([8, 128]),
193
+ torch.Size([1, 128]),
194
+ torch.Size([2, 128]),
195
+ torch.Size([8, 128]),
196
+ ]
197
+ assert calls[0][1] is layer_a.weight
198
+ assert calls[3][1] is layer_c.weight
199
+ assert all(call[3] is None for call in calls)
tests/model_executor/kernels/test_b12x_mxfp8_linear.py
new
+935
@@ -0,0 +1,935 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from __future__ import annotations
5
+
6
+import types
7
+from dataclasses import dataclass
8
+
9
+import pytest
10
+import torch
11
+
12
+from vllm.model_executor.kernels.linear import (
13
+ _LINEAR_BACKEND_KERNEL_MAP,
14
+ _POSSIBLE_FP8_BLOCK_KERNELS,
15
+ _POSSIBLE_FP8_KERNELS,
16
+ _POSSIBLE_MXFP8_KERNELS,
17
+ init_fp8_linear_kernel,
18
+ init_mxfp8_linear_kernel,
19
+)
20
+from vllm.model_executor.kernels.linear.mxfp8.b12x import (
21
+ B12xMxfp8LinearKernel,
22
+ _b12x_mxfp8_expected_m,
23
+ warmup_b12x_mxfp8_linear,
24
+)
25
+from vllm.model_executor.kernels.linear.mxfp8.Mxfp8LinearKernel import (
26
+ Mxfp8LinearLayerConfig,
27
+)
28
+from vllm.model_executor.kernels.linear.scaled_mm.b12x_block import (
29
+ B12xFp8BlockScaledMMKernel,
30
+ _run_b12x_fp8_block_scaled_mm,
31
+ warmup_b12x_block_fp8_linear,
32
+)
33
+from vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor import (
34
+ B12xTensorFP8ScaledMMLinearKernel,
35
+ warmup_b12x_tensor_fp8_linear,
36
+)
37
+from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
38
+ FP8ScaledMMLinearLayerConfig,
39
+)
40
+from vllm.model_executor.layers.linear import (
41
+ ReplicatedLinear,
42
+ UnquantizedLinearMethod,
43
+)
44
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
45
+ kFp8Dynamic128Sym,
46
+ kFp8Static128BlockSym,
47
+ kFp8StaticTensorSym,
48
+)
49
+from vllm.platforms import PlatformEnum
50
+from vllm.utils.b12x import b12x_warmup_token_counts
51
+
52
+
53
+def test_b12x_backend_maps_mxfp8_kernel() -> None:
54
+ assert B12xMxfp8LinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"]
55
+ assert B12xMxfp8LinearKernel in _POSSIBLE_MXFP8_KERNELS[PlatformEnum.CUDA]
56
+
57
+
58
+def test_b12x_backend_maps_tensor_fp8_kernel() -> None:
59
+ assert B12xTensorFP8ScaledMMLinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"]
60
+ assert B12xTensorFP8ScaledMMLinearKernel in _POSSIBLE_FP8_KERNELS[PlatformEnum.CUDA]
61
+
62
+
63
+@pytest.mark.parametrize(
64
+ ("kernels", "before", "b12x", "after"),
65
+ [
66
+ (
67
+ _POSSIBLE_FP8_KERNELS[PlatformEnum.CUDA],
68
+ "CutlassFP8ScaledMMLinearKernel",
69
+ "B12xTensorFP8ScaledMMLinearKernel",
70
+ "PerTensorTorchFP8ScaledMMLinearKernel",
71
+ ),
72
+ (
73
+ _POSSIBLE_FP8_BLOCK_KERNELS[PlatformEnum.CUDA],
74
+ "CutlassFp8BlockScaledMMKernel",
75
+ "B12xFp8BlockScaledMMKernel",
76
+ "MarlinFP8ScaledMMLinearKernel",
77
+ ),
78
+ (
79
+ _POSSIBLE_MXFP8_KERNELS[PlatformEnum.CUDA],
80
+ "MarlinMxfp8LinearKernel",
81
+ "B12xMxfp8LinearKernel",
82
+ "EmulationMxfp8LinearKernel",
83
+ ),
84
+ ],
85
+)
86
+def test_b12x_fp8_fallback_priority(
87
+ kernels: list[type],
88
+ before: str,
89
+ b12x: str,
90
+ after: str,
91
+) -> None:
92
+ names = [kernel.__name__ for kernel in kernels]
93
+
94
+ assert names.index(before) < names.index(b12x) < names.index(after)
95
+
96
+
97
+@torch.inference_mode()
98
+def test_b12x_backend_does_not_intercept_unquantized_bf16(
99
+ default_vllm_config,
100
+ dist_init,
101
+) -> None:
102
+ default_vllm_config.kernel_config.linear_backend = "b12x"
103
+ layer = ReplicatedLinear(
104
+ 128,
105
+ 64,
106
+ bias=False,
107
+ params_dtype=torch.bfloat16,
108
+ quant_config=None,
109
+ prefix="bf16_linear",
110
+ ).cuda()
111
+ layer.weight.data.normal_()
112
+ x = torch.randn(8, 128, device="cuda", dtype=torch.bfloat16)
113
+
114
+ output, output_bias = layer(x)
115
+ expected = torch.nn.functional.linear(x, layer.weight)
116
+
117
+ assert isinstance(layer.quant_method, UnquantizedLinearMethod)
118
+ assert output_bias is None
119
+ torch.testing.assert_close(output, expected)
120
+
121
+
122
+def test_b12x_explicit_backend_selects_per_tensor_fp8(
123
+ monkeypatch,
124
+ default_vllm_config,
125
+) -> None:
126
+ import vllm.model_executor.kernels.linear as linear_mod
127
+
128
+ monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
129
+ monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
130
+ monkeypatch.setattr(
131
+ B12xTensorFP8ScaledMMLinearKernel,
132
+ "is_supported",
133
+ classmethod(lambda cls, compute_capability=None: (True, None)),
134
+ )
135
+ monkeypatch.setattr(
136
+ B12xTensorFP8ScaledMMLinearKernel,
137
+ "can_implement",
138
+ classmethod(lambda cls, config: (True, None)),
139
+ )
140
+
141
+ kernel = init_fp8_linear_kernel(
142
+ activation_quant_key=kFp8StaticTensorSym,
143
+ weight_quant_key=kFp8StaticTensorSym,
144
+ input_dtype=torch.bfloat16,
145
+ out_dtype=torch.bfloat16,
146
+ weight_shape=(2048, 2048),
147
+ )
148
+
149
+ assert isinstance(kernel, B12xTensorFP8ScaledMMLinearKernel)
150
+
151
+
152
+def test_b12x_tensor_fp8_can_implement_supported_config() -> None:
153
+ config = FP8ScaledMMLinearLayerConfig(
154
+ activation_quant_key=kFp8StaticTensorSym,
155
+ weight_quant_key=kFp8StaticTensorSym,
156
+ weight_shape=(64, 128),
157
+ input_dtype=torch.bfloat16,
158
+ out_dtype=torch.bfloat16,
159
+ )
160
+
161
+ can_implement, reason = B12xTensorFP8ScaledMMLinearKernel.can_implement(config)
162
+
163
+ assert can_implement
164
+ assert reason is None
165
+
166
+
167
+def test_b12x_explicit_backend_selects_block_fp8(
168
+ monkeypatch,
169
+ default_vllm_config,
170
+) -> None:
171
+ import vllm.model_executor.kernels.linear as linear_mod
172
+
173
+ monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
174
+ monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
175
+ monkeypatch.setattr(
176
+ B12xFp8BlockScaledMMKernel,
177
+ "is_supported",
178
+ classmethod(lambda cls, compute_capability=None: (True, None)),
179
+ )
180
+
181
+ kernel = init_fp8_linear_kernel(
182
+ activation_quant_key=kFp8Dynamic128Sym,
183
+ weight_quant_key=kFp8Static128BlockSym,
184
+ input_dtype=torch.bfloat16,
185
+ out_dtype=torch.bfloat16,
186
+ weight_shape=(2048, 2048),
187
+ )
188
+
189
+ assert isinstance(kernel, B12xFp8BlockScaledMMKernel)
190
+
191
+
192
+def test_b12x_block_fp8_checks_runtime_support(monkeypatch) -> None:
193
+ import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod
194
+
195
+ platform = types.SimpleNamespace(
196
+ is_cuda=lambda: True,
197
+ is_device_capability_family=lambda family: family == 120,
198
+ )
199
+ monkeypatch.setattr(b12x_mod, "current_platform", platform)
200
+
201
+ monkeypatch.setattr(
202
+ b12x_mod,
203
+ "_import_b12x_blockscaled",
204
+ lambda: types.SimpleNamespace(is_supported=lambda: False),
205
+ )
206
+
207
+ supported, reason = B12xFp8BlockScaledMMKernel.is_supported()
208
+
209
+ assert not supported
210
+ assert reason == "B12X regular block-FP8 GEMM is not supported"
211
+
212
+
213
+def test_b12x_block_fp8_requires_matching_supported_dtypes() -> None:
214
+ def config(input_dtype: torch.dtype, out_dtype: torch.dtype):
215
+ return FP8ScaledMMLinearLayerConfig(
216
+ activation_quant_key=kFp8Dynamic128Sym,
217
+ weight_quant_key=kFp8Static128BlockSym,
218
+ weight_shape=(256, 128),
219
+ input_dtype=input_dtype,
220
+ out_dtype=out_dtype,
221
+ )
222
+
223
+ can_implement, reason = B12xFp8BlockScaledMMKernel.can_implement(
224
+ config(torch.float32, torch.float32)
225
+ )
226
+ assert not can_implement
227
+ assert reason == "Supports only bf16/fp16 input dtype"
228
+
229
+ can_implement, reason = B12xFp8BlockScaledMMKernel.can_implement(
230
+ config(torch.bfloat16, torch.float16)
231
+ )
232
+ assert not can_implement
233
+ assert reason == "Input and output dtype must match"
234
+
235
+ can_implement, reason = B12xFp8BlockScaledMMKernel.can_implement(
236
+ config(torch.float16, torch.float16)
237
+ )
238
+ assert can_implement
239
+ assert reason is None
240
+
241
+
242
+def test_b12x_block_fp8_requires_aligned_features() -> None:
243
+ def can_implement(weight_shape: tuple[int, int]):
244
+ config = FP8ScaledMMLinearLayerConfig(
245
+ activation_quant_key=kFp8Dynamic128Sym,
246
+ weight_quant_key=kFp8Static128BlockSym,
247
+ weight_shape=weight_shape,
248
+ input_dtype=torch.bfloat16,
249
+ out_dtype=torch.bfloat16,
250
+ )
251
+ return B12xFp8BlockScaledMMKernel.can_implement(config)
252
+
253
+ assert can_implement((256, 192)) == (
254
+ False,
255
+ "Input features must be a positive multiple of 128",
256
+ )
257
+ assert can_implement((192, 256)) == (
258
+ False,
259
+ "Output features must be a positive multiple of 128",
260
+ )
261
+
262
+
263
+def test_b12x_warmup_token_counts_cover_serving_regimes() -> None:
264
+ assert b12x_warmup_token_counts(
265
+ max_tokens=2048,
266
+ cudagraph_capture_sizes=[1, 2, 8, 128],
267
+ ) == (1, 2, 8, 128, 2048)
268
+
269
+
270
+def test_warmup_b12x_block_fp8_dedupes_weight_signatures(monkeypatch) -> None:
271
+ import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod
272
+
273
+ calls = []
274
+
275
+ def run(a, weight, a_scale, weight_scale, out_dtype):
276
+ calls.append((a.shape, weight, a_scale.shape, weight_scale, out_dtype))
277
+ return torch.empty((a.shape[0], weight.shape[0]), dtype=out_dtype)
278
+
279
+ platform = types.SimpleNamespace(
280
+ is_cuda=lambda: True,
281
+ is_device_capability_family=lambda family: family == 120,
282
+ )
283
+ monkeypatch.setattr(b12x_mod, "current_platform", platform)
284
+
285
+ monkeypatch.setattr(
286
+ b12x_mod,
287
+ "_import_b12x_blockscaled",
288
+ lambda: types.SimpleNamespace(),
289
+ )
290
+ monkeypatch.setattr(b12x_mod, "_run_b12x_fp8_block_scaled_mm", run)
291
+
292
+ def layer(in_features: int, out_features: int):
293
+ return types.SimpleNamespace(
294
+ b12x_block_fp8_linear=True,
295
+ weight=torch.empty((out_features, in_features), dtype=torch.float8_e4m3fn),
296
+ weight_scale_inv=torch.empty(
297
+ (out_features // 128, in_features // 128), dtype=torch.float32
298
+ ),
299
+ )
300
+
301
+ layer_a = layer(128, 256)
302
+ layer_b = layer(256, 128)
303
+ modules = [
304
+ layer_a,
305
+ layer_a,
306
+ layer_b,
307
+ types.SimpleNamespace(),
308
+ ]
309
+ model = types.SimpleNamespace(modules=lambda: iter(modules))
310
+
311
+ warmed = warmup_b12x_block_fp8_linear(
312
+ model,
313
+ max_tokens=32,
314
+ cudagraph_capture_sizes=[2, 8],
315
+ output_dtype=torch.bfloat16,
316
+ )
317
+
318
+ assert warmed == 8
319
+ assert [call[0][0] for call in calls] == [1, 2, 8, 32] * 2
320
+ assert [call[2][0] for call in calls] == [1, 2, 8, 32] * 2
321
+ assert calls[0][1] is layer_a.weight
322
+ assert calls[4][1] is layer_b.weight
323
+ assert calls[0][3] is layer_a.weight_scale_inv
324
+ assert calls[4][3] is layer_b.weight_scale_inv
325
+ assert all(call[4] == torch.bfloat16 for call in calls)
326
+
327
+
328
+def test_b12x_tensor_fp8_process_weights_packs_modelopt_layout(
329
+ monkeypatch,
330
+) -> None:
331
+ import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod
332
+
333
+ calls = []
334
+ packed = types.SimpleNamespace(out_features=64)
335
+
336
+ def pack(weight: torch.Tensor, output_scale: torch.Tensor):
337
+ calls.append((weight, output_scale))
338
+ return packed
339
+
340
+ monkeypatch.setattr(
341
+ b12x_mod,
342
+ "_import_b12x_tensor_fp8",
343
+ lambda: types.SimpleNamespace(pack_weight=pack),
344
+ )
345
+ layer = torch.nn.Module()
346
+ layer.prefix = "model.layers.0.self_attn.qkv_proj"
347
+ original_weight = (
348
+ torch.randn((128, 64), dtype=torch.float32).clamp(-4, 4).to(torch.float8_e4m3fn)
349
+ )
350
+ layer.weight = torch.nn.Parameter(original_weight, requires_grad=False)
351
+ layer.weight_scale = torch.nn.Parameter(torch.tensor(0.25), requires_grad=False)
352
+ layer.input_scale = torch.nn.Parameter(torch.tensor(0.5), requires_grad=False)
353
+ weight_loader = object()
354
+ scale_loader = object()
355
+ layer.weight.weight_loader = weight_loader
356
+ layer.weight_scale.weight_loader = scale_loader
357
+ kernel = object.__new__(B12xTensorFP8ScaledMMLinearKernel)
358
+ kernel.config = types.SimpleNamespace(weight_shape=(64, 128))
359
+ kernel.layer_param_names = (
360
+ "weight",
361
+ "weight_scale",
362
+ "input_scale",
363
+ "input_scale_ub",
364
+ )
365
+
366
+ kernel.process_weights_after_loading(layer)
367
+
368
+ assert layer.b12x_tensor_fp8_packed_weight is packed
369
+ assert len(calls) == 1
370
+ weight, output_scale = calls[0]
371
+ torch.testing.assert_close(weight, original_weight.T.contiguous())
372
+ torch.testing.assert_close(output_scale, torch.tensor([0.125]))
373
+ assert layer.weight.numel() == 0
374
+ assert layer.weight_scale.numel() == 0
375
+ assert layer.weight.weight_loader is weight_loader
376
+ assert layer.weight_scale.weight_loader is scale_loader
377
+ torch.testing.assert_close(layer.input_scale, torch.tensor(0.5))
378
+
379
+
380
+def test_warmup_b12x_tensor_fp8_dedupes_weight_signatures(monkeypatch) -> None:
381
+ import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod
382
+
383
+ calls = []
384
+
385
+ def prewarm(packed_weight, token_counts, *, out_dtype, stream):
386
+ del stream
387
+ calls.append((packed_weight, tuple(token_counts), out_dtype))
388
+ return len(tuple(token_counts))
389
+
390
+ platform = types.SimpleNamespace(
391
+ is_cuda=lambda: True,
392
+ is_device_capability_family=lambda family: family == 120,
393
+ )
394
+ monkeypatch.setattr(b12x_mod, "current_platform", platform)
395
+ monkeypatch.setattr(
396
+ b12x_mod,
397
+ "_import_b12x_tensor_fp8",
398
+ lambda: types.SimpleNamespace(prewarm=prewarm),
399
+ )
400
+ monkeypatch.setattr(
401
+ b12x_mod,
402
+ "current_stream",
403
+ lambda: types.SimpleNamespace(cuda_stream=object()),
404
+ )
405
+
406
+ def packed(in_features: int, padded_in_features: int, out_features: int):
407
+ return types.SimpleNamespace(
408
+ in_features=in_features,
409
+ padded_in_features=padded_in_features,
410
+ out_features=out_features,
411
+ values=torch.empty(1),
412
+ )
413
+
414
+ packed_a = packed(128, 128, 256)
415
+ packed_b = packed(160, 256, 512)
416
+ modules = [
417
+ types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_a),
418
+ types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_a),
419
+ types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_b),
420
+ types.SimpleNamespace(),
421
+ ]
422
+ model = types.SimpleNamespace(modules=lambda: iter(modules))
423
+
424
+ warmed = warmup_b12x_tensor_fp8_linear(
425
+ model,
426
+ max_tokens=2048,
427
+ cudagraph_capture_sizes=[1, 2],
428
+ )
429
+
430
+ assert warmed == 6
431
+ assert calls == [
432
+ (packed_a, (1, 2, 2048), torch.bfloat16),
433
+ (packed_b, (1, 2, 2048), torch.bfloat16),
434
+ ]
435
+
436
+
437
+def test_b12x_tensor_fp8_apply_quantizes_and_uses_packed_weight(
438
+ monkeypatch,
439
+) -> None:
440
+ import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod
441
+
442
+ calls = []
443
+
444
+ def mm(
445
+ source: torch.Tensor,
446
+ packed_weight,
447
+ *,
448
+ bias: torch.Tensor | None = None,
449
+ out_dtype: torch.dtype,
450
+ expected_m: int,
451
+ stream: object = None,
452
+ ) -> torch.Tensor:
453
+ del stream
454
+ calls.append((source, packed_weight, bias, out_dtype, expected_m))
455
+ return torch.full(
456
+ (source.shape[0], packed_weight.out_features),
457
+ 3.0,
458
+ dtype=out_dtype,
459
+ )
460
+
461
+ monkeypatch.setattr(
462
+ b12x_mod,
463
+ "_import_b12x_tensor_fp8",
464
+ lambda: types.SimpleNamespace(mm=mm),
465
+ )
466
+ monkeypatch.setattr(
467
+ b12x_mod,
468
+ "current_stream",
469
+ lambda: types.SimpleNamespace(cuda_stream=object()),
470
+ )
471
+ monkeypatch.setattr(torch.compiler, "is_compiling", lambda: False)
472
+
473
+ layer = torch.nn.Module()
474
+ packed = types.SimpleNamespace(out_features=48)
475
+ layer.b12x_tensor_fp8_packed_weight = packed
476
+ layer.weight = torch.nn.Parameter(
477
+ torch.empty((128, 48), dtype=torch.float8_e4m3fn),
478
+ requires_grad=False,
479
+ )
480
+ layer.weight_scale = torch.nn.Parameter(torch.tensor(0.25), requires_grad=False)
481
+ layer.input_scale = torch.nn.Parameter(torch.tensor(0.5), requires_grad=False)
482
+ x = torch.empty((2, 3, 128), dtype=torch.bfloat16)
483
+ x_q = torch.empty((6, 128), dtype=torch.float8_e4m3fn)
484
+ bias = torch.empty((48,), dtype=torch.bfloat16)
485
+ kernel = object.__new__(B12xTensorFP8ScaledMMLinearKernel)
486
+ kernel.config = types.SimpleNamespace(out_dtype=torch.bfloat16)
487
+ kernel.layer_param_names = (
488
+ "weight",
489
+ "weight_scale",
490
+ "input_scale",
491
+ "input_scale_ub",
492
+ )
493
+ kernel.quant_fp8 = lambda source, scale, scale_ub: (x_q, scale)
494
+
495
+ output = kernel.apply_weights(layer, x, bias)
496
+
497
+ assert output.shape == (2, 3, 48)
498
+ assert output.dtype == torch.bfloat16
499
+ assert len(calls) == 1
500
+ source, called_packed, called_bias, out_dtype, expected_m = calls[0]
501
+ assert source.data_ptr() == x_q.data_ptr()
502
+ assert called_packed is packed
503
+ assert called_bias is bias
504
+ assert out_dtype == torch.bfloat16
505
+ assert expected_m == 6
506
+
507
+
508
+def test_b12x_mxfp8_explicit_backend_selects_kernel(monkeypatch) -> None:
509
+ import vllm.model_executor.kernels.linear as linear_mod
510
+
511
+ monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
512
+ monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
513
+ monkeypatch.setattr(
514
+ B12xMxfp8LinearKernel,
515
+ "is_supported",
516
+ classmethod(lambda cls, compute_capability=None: (True, None)),
517
+ )
518
+ monkeypatch.setattr(
519
+ B12xMxfp8LinearKernel,
520
+ "can_implement",
521
+ classmethod(lambda cls, c: (True, None)),
522
+ )
523
+
524
+ kernel = init_mxfp8_linear_kernel()
525
+
526
+ assert isinstance(kernel, B12xMxfp8LinearKernel)
527
+
528
+
529
+def test_b12x_mxfp8_can_implement_supported_config() -> None:
530
+ can_implement, reason = B12xMxfp8LinearKernel.can_implement(
531
+ Mxfp8LinearLayerConfig()
532
+ )
533
+
534
+ assert can_implement
535
+ assert reason is None
536
+
537
+
538
+def test_b12x_mxfp8_expected_m_uses_live_m() -> None:
539
+ assert _b12x_mxfp8_expected_m(0) == 1
540
+ assert _b12x_mxfp8_expected_m(1) == 1
541
+ assert _b12x_mxfp8_expected_m(2) == 2
542
+ assert _b12x_mxfp8_expected_m(8) == 8
543
+ assert _b12x_mxfp8_expected_m(9) == 9
544
+ assert _b12x_mxfp8_expected_m(128) == 128
545
+ assert _b12x_mxfp8_expected_m(129) == 129
546
+ assert _b12x_mxfp8_expected_m(2048) == 2048
547
+
548
+
549
+def test_warmup_b12x_mxfp8_linear_dedupes_weight_signatures(
550
+ monkeypatch,
551
+) -> None:
552
+ import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod
553
+
554
+ calls = []
555
+
556
+ def mm(
557
+ source: torch.Tensor,
558
+ packed_weight,
559
+ *,
560
+ bias: torch.Tensor | None = None,
561
+ expected_m: int | None = None,
562
+ stream: object = None,
563
+ ) -> torch.Tensor:
564
+ del stream
565
+ calls.append((source.shape, packed_weight, bias, expected_m))
566
+ return source.new_empty((source.shape[0], packed_weight.out_features))
567
+
568
+ platform = types.SimpleNamespace(
569
+ is_cuda=lambda: True,
570
+ is_device_capability_family=lambda family: family == 120,
571
+ )
572
+ monkeypatch.setattr(b12x_mod, "current_platform", platform)
573
+ monkeypatch.setattr(
574
+ b12x_mod,
575
+ "_import_b12x_mxfp8",
576
+ lambda: types.SimpleNamespace(mm=mm),
577
+ )
578
+ monkeypatch.setattr(
579
+ b12x_mod,
580
+ "current_stream",
581
+ lambda: types.SimpleNamespace(cuda_stream=object()),
582
+ )
583
+
584
+ def packed(in_features: int, padded_in_features: int, out_features: int):
585
+ return types.SimpleNamespace(
586
+ in_features=in_features,
587
+ padded_in_features=padded_in_features,
588
+ out_features=out_features,
589
+ weight=types.SimpleNamespace(values=torch.empty(1)),
590
+ )
591
+
592
+ packed_a = packed(128, 128, 256)
593
+ packed_b = packed(128, 128, 512)
594
+ modules = [
595
+ types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_a),
596
+ types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_a),
597
+ types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_b),
598
+ types.SimpleNamespace(),
599
+ ]
600
+ model = types.SimpleNamespace(modules=lambda: iter(modules))
601
+
602
+ warmed = warmup_b12x_mxfp8_linear(
603
+ model,
604
+ max_tokens=2048,
605
+ cudagraph_capture_sizes=[1, 2],
606
+ )
607
+
608
+ assert warmed == 6
609
+ assert [call[0] for call in calls] == [
610
+ torch.Size([1, 128]),
611
+ torch.Size([2, 128]),
612
+ torch.Size([2048, 128]),
613
+ torch.Size([1, 128]),
614
+ torch.Size([2, 128]),
615
+ torch.Size([2048, 128]),
616
+ ]
617
+ assert [call[3] for call in calls] == [1, 2, 2048, 1, 2, 2048]
618
+ assert calls[0][1] is packed_a
619
+ assert calls[3][1] is packed_b
620
+
621
+
622
+def test_b12x_mxfp8_support_check_reports_missing_import(monkeypatch) -> None:
623
+ import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod
624
+
625
+ monkeypatch.setattr(b12x_mod.current_platform, "is_cuda", lambda: True)
626
+ monkeypatch.setattr(
627
+ b12x_mod.current_platform,
628
+ "is_device_capability_family",
629
+ lambda family: family == 120,
630
+ )
631
+ monkeypatch.setattr(b12x_mod, "_import_b12x_mxfp8", lambda: None)
632
+
633
+ is_supported, reason = B12xMxfp8LinearKernel.is_supported()
634
+
635
+ assert not is_supported
636
+ assert reason == "Install the B12X backend with `pip install vllm[b12x]`"
637
+
638
+
639
+def test_b12x_mxfp8_support_respects_runtime_probe(monkeypatch) -> None:
640
+ import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod
641
+
642
+ monkeypatch.setattr(b12x_mod.current_platform, "is_cuda", lambda: True)
643
+ monkeypatch.setattr(
644
+ b12x_mod.current_platform,
645
+ "is_device_capability_family",
646
+ lambda family: family == 120,
647
+ )
648
+ monkeypatch.setattr(
649
+ b12x_mod,
650
+ "_import_b12x_mxfp8",
651
+ lambda: types.SimpleNamespace(is_supported=lambda: False),
652
+ )
653
+
654
+ is_supported, reason = B12xMxfp8LinearKernel.is_supported()
655
+
656
+ assert not is_supported
657
+ assert reason == "b12x.gemm.mxfp8_linear is not supported"
658
+
659
+
660
+def test_b12x_mxfp8_process_weights_packs_modelopt_layout(monkeypatch) -> None:
661
+ import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod
662
+
663
+ calls = []
664
+ packed = types.SimpleNamespace(out_features=48)
665
+
666
+ def pack(weight: torch.Tensor, weight_scale: torch.Tensor):
667
+ calls.append((weight, weight_scale))
668
+ return packed
669
+
670
+ monkeypatch.setattr(
671
+ b12x_mod,
672
+ "_import_b12x_mxfp8",
673
+ lambda: types.SimpleNamespace(pack_weight=pack),
674
+ )
675
+
676
+ layer = torch.nn.Module()
677
+ layer.prefix = "model.layers.0.self_attn.qkv_proj"
678
+ layer.weight = torch.nn.Parameter(
679
+ torch.empty((48, 128), dtype=torch.float8_e4m3fn),
680
+ requires_grad=False,
681
+ )
682
+ layer.weight_scale = torch.nn.Parameter(
683
+ torch.empty((64, 8), dtype=torch.uint8),
684
+ requires_grad=False,
685
+ )
686
+ weight_loader = object()
687
+ scale_loader = object()
688
+ layer.weight.weight_loader = weight_loader
689
+ layer.weight_scale.weight_loader = scale_loader
690
+ kernel = object.__new__(B12xMxfp8LinearKernel)
691
+
692
+ kernel.process_weights_after_loading(layer)
693
+
694
+ assert layer.b12x_mxfp8_packed_weight is packed
695
+ assert len(calls) == 1
696
+ weight, weight_scale = calls[0]
697
+ assert weight.shape == (48, 128)
698
+ assert weight_scale.shape == (48, 4)
699
+ assert weight.dtype == torch.float8_e4m3fn
700
+ assert weight_scale.dtype == torch.uint8
701
+ assert layer.weight.numel() == 0
702
+ assert layer.weight_scale.numel() == 0
703
+ assert layer.weight.weight_loader is weight_loader
704
+ assert layer.weight_scale.weight_loader is scale_loader
705
+
706
+
707
+def test_b12x_mxfp8_reload_reuses_packed_tensor_addresses(monkeypatch) -> None:
708
+ import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod
709
+
710
+ @dataclass(frozen=True)
711
+ class PackedWeight:
712
+ values: torch.Tensor
713
+ scales: torch.Tensor
714
+ out_features: int
715
+
716
+ def pack(weight: torch.Tensor, weight_scale: torch.Tensor) -> PackedWeight:
717
+ return PackedWeight(
718
+ values=weight.clone(),
719
+ scales=weight_scale.clone(),
720
+ out_features=int(weight.shape[0]),
721
+ )
722
+
723
+ monkeypatch.setattr(
724
+ b12x_mod,
725
+ "_import_b12x_mxfp8",
726
+ lambda: types.SimpleNamespace(pack_weight=pack),
727
+ )
728
+ layer = torch.nn.Module()
729
+ layer.prefix = "model.layers.0.mlp.down_proj"
730
+ layer.weight = torch.nn.Parameter(
731
+ torch.zeros((48, 128), dtype=torch.float8_e4m3fn),
732
+ requires_grad=False,
733
+ )
734
+ layer.weight_scale = torch.nn.Parameter(
735
+ torch.zeros((48, 4), dtype=torch.uint8),
736
+ requires_grad=False,
737
+ )
738
+ kernel = object.__new__(B12xMxfp8LinearKernel)
739
+
740
+ kernel.process_weights_after_loading(layer)
741
+ packed = layer.b12x_mxfp8_packed_weight
742
+ values_ptr = packed.values.data_ptr()
743
+ scales_ptr = packed.scales.data_ptr()
744
+
745
+ layer.weight = torch.nn.Parameter(
746
+ torch.ones((48, 128), dtype=torch.float8_e4m3fn),
747
+ requires_grad=False,
748
+ )
749
+ layer.weight_scale = torch.nn.Parameter(
750
+ torch.full((48, 4), 3, dtype=torch.uint8),
751
+ requires_grad=False,
752
+ )
753
+ kernel.process_weights_after_loading(layer)
754
+
755
+ assert layer.b12x_mxfp8_packed_weight is packed
756
+ assert packed.values.data_ptr() == values_ptr
757
+ assert packed.scales.data_ptr() == scales_ptr
758
+ torch.testing.assert_close(
759
+ packed.values,
760
+ torch.ones((48, 128), dtype=torch.float8_e4m3fn),
761
+ )
762
+ torch.testing.assert_close(
763
+ packed.scales,
764
+ torch.full((48, 4), 3, dtype=torch.uint8),
765
+ )
766
+ assert layer.weight.numel() == 0
767
+ assert layer.weight_scale.numel() == 0
768
+
769
+
770
+def test_b12x_block_fp8_process_weights_keeps_native_block_layout() -> None:
771
+ layer = torch.nn.Module()
772
+ layer.weight = torch.nn.Parameter(
773
+ torch.empty((128, 128), dtype=torch.float8_e4m3fn),
774
+ requires_grad=False,
775
+ )
776
+ layer.weight_scale_inv = torch.nn.Parameter(
777
+ torch.empty((1, 1), dtype=torch.float32),
778
+ requires_grad=False,
779
+ )
780
+ layer.weight_block_size = [128, 128]
781
+ weight_loader = object()
782
+ scale_loader = object()
783
+ layer.weight.weight_loader = weight_loader
784
+ layer.weight_scale_inv.weight_loader = scale_loader
785
+ kernel = object.__new__(B12xFp8BlockScaledMMKernel)
786
+
787
+ kernel.process_weights_after_loading(layer)
788
+
789
+ assert layer.b12x_block_fp8_linear
790
+ assert layer.weight.shape == (128, 128)
791
+ assert layer.weight.dtype == torch.float8_e4m3fn
792
+ assert layer.weight_scale_inv.shape == (1, 1)
793
+ assert layer.weight_scale_inv.dtype == torch.float32
794
+ assert layer.weight.weight_loader is weight_loader
795
+ assert layer.weight_scale_inv.weight_loader is scale_loader
796
+
797
+
798
+@pytest.mark.parametrize("scale_dtype", [torch.float8_e8m0fnu, torch.uint8])
799
+def test_b12x_block_fp8_upcasts_e8m0_weight_scales(scale_dtype) -> None:
800
+ layer = torch.nn.Module()
801
+ layer.weight = torch.nn.Parameter(
802
+ torch.empty((128, 128), dtype=torch.float8_e4m3fn),
803
+ requires_grad=False,
804
+ )
805
+ scale_bytes = torch.tensor([[125]], dtype=torch.uint8)
806
+ layer.weight_scale_inv = torch.nn.Parameter(
807
+ scale_bytes.view(scale_dtype),
808
+ requires_grad=False,
809
+ )
810
+ layer.weight_block_size = [128, 128]
811
+ kernel = object.__new__(B12xFp8BlockScaledMMKernel)
812
+
813
+ kernel.process_weights_after_loading(layer)
814
+
815
+ assert layer.weight_scale_inv.dtype == torch.float32
816
+ torch.testing.assert_close(
817
+ layer.weight_scale_inv,
818
+ torch.tensor([[0.25]], dtype=torch.float32),
819
+ )
820
+
821
+
822
+def test_b12x_mxfp8_apply_uses_packed_weight(monkeypatch) -> None:
823
+ import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod
824
+
825
+ calls = []
826
+
827
+ def mxfp8_linear(
828
+ source: torch.Tensor,
829
+ packed_weight,
830
+ *,
831
+ bias: torch.Tensor | None = None,
832
+ expected_m: int | None = None,
833
+ stream: object = None,
834
+ ) -> torch.Tensor:
835
+ del stream
836
+ calls.append((source, packed_weight, bias, expected_m))
837
+ return source.new_full((source.shape[0], packed_weight.out_features), 3.0)
838
+
839
+ monkeypatch.setattr(
840
+ b12x_mod,
841
+ "_import_b12x_mxfp8",
842
+ lambda: types.SimpleNamespace(mm=mxfp8_linear),
843
+ )
844
+
845
+ layer = torch.nn.Module()
846
+ packed = types.SimpleNamespace(out_features=48)
847
+ layer.b12x_mxfp8_packed_weight = packed
848
+ x = torch.empty((2, 3, 128), dtype=torch.bfloat16)
849
+ bias = torch.empty((48,), dtype=torch.bfloat16)
850
+ kernel = object.__new__(B12xMxfp8LinearKernel)
851
+
852
+ output = kernel.apply_weights(layer, x, bias)
853
+
854
+ assert output.shape == (2, 3, 48)
855
+ assert output.dtype == x.dtype
856
+ assert len(calls) == 1
857
+ source, called_packed, called_bias, expected_m = calls[0]
858
+ assert source.shape == (6, 128)
859
+ assert called_packed is packed
860
+ assert called_bias is bias
861
+ assert expected_m == 6
862
+
863
+
864
+def test_b12x_block_fp8_apply_uses_b12x_recipe_api(monkeypatch) -> None:
865
+ import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod
866
+
867
+ calls = []
868
+
869
+ def mm_block_fp8(*args, **kwargs):
870
+ calls.append((args, kwargs))
871
+ return torch.full(
872
+ (args[0].shape[0], args[2].shape[0]),
873
+ 13.0,
874
+ dtype=kwargs["out_dtype"],
875
+ )
876
+
877
+ monkeypatch.setattr(
878
+ b12x_mod,
879
+ "_import_b12x_blockscaled",
880
+ lambda: types.SimpleNamespace(mm_block_fp8=mm_block_fp8),
881
+ )
882
+
883
+ a = torch.empty((6, 128), dtype=torch.float8_e4m3fn)
884
+ weight = torch.empty((256, 128), dtype=torch.float8_e4m3fn)
885
+ a_scale = torch.empty((6, 1), dtype=torch.float32)
886
+ weight_scale = torch.empty((2, 1), dtype=torch.float32)
887
+ kernel = object.__new__(B12xFp8BlockScaledMMKernel)
888
+ kernel.config = types.SimpleNamespace(out_dtype=torch.bfloat16)
889
+
890
+ output = kernel.apply_block_scaled_mm(a, weight, a_scale, weight_scale)
891
+
892
+ assert output.shape == (6, 256)
893
+ assert output.dtype == torch.bfloat16
894
+ assert len(calls) == 1
895
+ assert calls[0] == (
896
+ (a, a_scale, weight, weight_scale),
897
+ {"out_dtype": torch.bfloat16},
898
+ )
899
+ torch.testing.assert_close(output, torch.full_like(output, 13.0))
900
+
901
+
902
+def test_b12x_block_fp8_helper_uses_regular_compact_scale_api(monkeypatch) -> None:
903
+ import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod
904
+
905
+ calls = []
906
+
907
+ def mm_block_fp8(*args, **kwargs):
908
+ calls.append((args, kwargs))
909
+ return torch.full((6, 256), 17.0, dtype=torch.bfloat16)
910
+
911
+ monkeypatch.setattr(
912
+ b12x_mod,
913
+ "_import_b12x_blockscaled",
914
+ lambda: types.SimpleNamespace(mm_block_fp8=mm_block_fp8),
915
+ )
916
+
917
+ a = torch.empty((6, 128), dtype=torch.float8_e4m3fn)
918
+ weight = torch.empty((256, 128), dtype=torch.float8_e4m3fn)
919
+ a_scale = torch.empty((6, 1), dtype=torch.float32)
920
+ weight_scale = torch.empty((2, 1), dtype=torch.float32)
921
+ output = _run_b12x_fp8_block_scaled_mm(
922
+ a,
923
+ weight,
924
+ a_scale,
925
+ weight_scale,
926
+ torch.bfloat16,
927
+ )
928
+
929
+ assert output.shape == (6, 256)
930
+ args, kwargs = calls[0]
931
+ assert args == (a, a_scale, weight, weight_scale)
932
+ assert kwargs == {
933
+ "out_dtype": torch.bfloat16,
934
+ }
935
+ torch.testing.assert_close(output, torch.full_like(output, 17.0))
tests/model_executor/kernels/test_b12x_nvfp4_linear.py
new
+246
@@ -0,0 +1,246 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from __future__ import annotations
5
+
6
+import types
7
+
8
+import torch
9
+
10
+from vllm.model_executor.kernels.linear import (
11
+ _LINEAR_BACKEND_KERNEL_MAP,
12
+ _POSSIBLE_NVFP4_KERNELS,
13
+ _resolve_backend_kernels,
14
+ init_nvfp4_linear_kernel,
15
+)
16
+from vllm.model_executor.kernels.linear.nvfp4.b12x import (
17
+ B12xNvFp4LinearKernel,
18
+ warmup_b12x_nvfp4_linear,
19
+)
20
+from vllm.model_executor.kernels.linear.nvfp4.marlin import (
21
+ MarlinNvFp4LinearKernel,
22
+)
23
+from vllm.platforms import PlatformEnum
24
+
25
+
26
+def test_b12x_backend_maps_nvfp4_kernel() -> None:
27
+ assert B12xNvFp4LinearKernel in _LINEAR_BACKEND_KERNEL_MAP["b12x"]
28
+ assert B12xNvFp4LinearKernel in _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA]
29
+
30
+
31
+def test_b12x_nvfp4_fallback_priority() -> None:
32
+ kernels = _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA]
33
+ names = [kernel.__name__ for kernel in kernels]
34
+
35
+ assert (
36
+ names.index("FbgemmNvFp4LinearKernel")
37
+ < names.index("B12xNvFp4LinearKernel")
38
+ < names.index("EmulationNvFp4LinearKernel")
39
+ )
40
+
41
+
42
+def test_b12x_nvfp4_explicit_backend_selects_native_kernel(monkeypatch) -> None:
43
+ import vllm.model_executor.kernels.linear as linear_mod
44
+
45
+ monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
46
+ monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
47
+ monkeypatch.setattr(
48
+ B12xNvFp4LinearKernel,
49
+ "is_supported",
50
+ classmethod(lambda cls, compute_capability=None: (True, None)),
51
+ )
52
+ monkeypatch.setattr(
53
+ B12xNvFp4LinearKernel,
54
+ "can_implement",
55
+ classmethod(lambda cls, config: (True, None)),
56
+ )
57
+
58
+ kernel = init_nvfp4_linear_kernel()
59
+
60
+ assert isinstance(kernel, B12xNvFp4LinearKernel)
61
+
62
+
63
+def test_b12x_nvfp4_can_implement_supported_config() -> None:
64
+ can_implement, reason = B12xNvFp4LinearKernel.can_implement(None)
65
+
66
+ assert can_implement
67
+ assert reason is None
68
+
69
+
70
+def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None:
71
+ import vllm.model_executor.kernels.linear as linear_mod
72
+
73
+ monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
74
+ monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
75
+ monkeypatch.setattr(
76
+ MarlinNvFp4LinearKernel,
77
+ "is_supported",
78
+ classmethod(lambda cls, compute_capability=None: (True, None)),
79
+ )
80
+
81
+ kernel = init_nvfp4_linear_kernel(use_a16=True)
82
+
83
+ assert isinstance(kernel, MarlinNvFp4LinearKernel)
84
+
85
+
86
+def test_backend_without_w4a16_kernel_preserves_fallback(monkeypatch) -> None:
87
+ import vllm.model_executor.kernels.linear as linear_mod
88
+
89
+ kernels = [MarlinNvFp4LinearKernel]
90
+ monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
91
+ assert _resolve_backend_kernels(kernels, "NVFP4") == kernels
92
+
93
+
94
+def test_b12x_nvfp4_processes_scale_and_preserves_loader(monkeypatch) -> None:
95
+ import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod
96
+
97
+ scale = torch.empty((48, 8), dtype=torch.float8_e4m3fn)
98
+ swizzled_scale = torch.empty((128, 8), dtype=torch.float8_e4m3fn)
99
+ intrinsics = types.SimpleNamespace(swizzle_block_scale=lambda value: swizzled_scale)
100
+ monkeypatch.setattr(b12x_mod, "_import_b12x_intrinsics", lambda: intrinsics)
101
+
102
+ layer = torch.nn.Module()
103
+ layer.prefix = "model.layers.0.mlp.shared_expert.down_proj"
104
+ layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False)
105
+ weight_loader = object()
106
+ layer.weight_scale.weight_loader = weight_loader
107
+ kernel = object.__new__(B12xNvFp4LinearKernel)
108
+
109
+ kernel.process_weights_after_loading(layer)
110
+
111
+ assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr()
112
+ assert layer.weight_scale.weight_loader is weight_loader
113
+ assert layer.b12x_nvfp4_linear
114
+
115
+
116
+def test_b12x_nvfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None:
117
+ import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod
118
+
119
+ calls: list[tuple] = []
120
+ quant_calls: list[tuple] = []
121
+ x_packed = torch.empty((6, 64), dtype=torch.uint8)
122
+ x_scale_storage = torch.empty((128, 8), dtype=torch.float8_e4m3fn)
123
+
124
+ def quant(*args, **kwargs):
125
+ quant_calls.append((args, kwargs))
126
+ return x_packed, x_scale_storage
127
+
128
+ def mm_nvfp4(*args, **kwargs):
129
+ calls.append((args, kwargs))
130
+ return torch.full((6, 48), 3.0, dtype=torch.bfloat16)
131
+
132
+ monkeypatch.setattr(
133
+ b12x_mod,
134
+ "scaled_fp4_quant",
135
+ quant,
136
+ )
137
+ monkeypatch.setattr(
138
+ b12x_mod,
139
+ "_import_b12x_blockscaled",
140
+ lambda: types.SimpleNamespace(mm_nvfp4=mm_nvfp4),
141
+ )
142
+
143
+ layer = torch.nn.Module()
144
+ layer.output_size_per_partition = 48
145
+ layer.weight = torch.empty((48, 64), dtype=torch.uint8)
146
+ layer.weight_scale = torch.empty((128, 8), dtype=torch.float8_e4m3fn)
147
+ layer.input_global_scale_inv = torch.tensor(2.0)
148
+ layer.alpha = torch.tensor(0.25)
149
+ x = torch.empty((2, 3, 256), dtype=torch.bfloat16)[..., ::2]
150
+ bias = torch.ones(48, dtype=torch.bfloat16)
151
+ kernel = object.__new__(B12xNvFp4LinearKernel)
152
+
153
+ output = kernel.apply_weights(layer, x, bias)
154
+
155
+ assert output.shape == (2, 3, 48)
156
+ torch.testing.assert_close(output, torch.full_like(output, 4.0))
157
+ assert len(quant_calls) == 1
158
+ quant_args, quant_kwargs = quant_calls[0]
159
+ assert quant_args[0].shape == (6, 128)
160
+ assert quant_args[0].data_ptr() == x.data_ptr()
161
+ assert quant_args[1] is layer.input_global_scale_inv
162
+ assert not quant_args[0].is_contiguous()
163
+ assert quant_kwargs == {"is_sf_swizzled_layout": True}
164
+ assert len(calls) == 1
165
+ args, kwargs = calls[0]
166
+ assert args == (
167
+ x_packed,
168
+ x_scale_storage,
169
+ layer.weight,
170
+ layer.weight_scale,
171
+ layer.alpha,
172
+ )
173
+ assert kwargs == {"out_dtype": torch.bfloat16}
174
+
175
+
176
+def test_warmup_b12x_nvfp4_dedupes_weight_signatures(monkeypatch) -> None:
177
+ import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod
178
+
179
+ calls = []
180
+
181
+ def apply(
182
+ source,
183
+ weight,
184
+ weight_scale,
185
+ input_global_scale_inv,
186
+ alpha,
187
+ bias,
188
+ ):
189
+ calls.append(
190
+ (
191
+ source.shape,
192
+ weight,
193
+ weight_scale,
194
+ input_global_scale_inv,
195
+ alpha,
196
+ bias,
197
+ )
198
+ )
199
+ return source.new_empty((source.shape[0], weight.shape[0]))
200
+
201
+ platform = types.SimpleNamespace(
202
+ is_cuda=lambda: True,
203
+ is_device_capability_family=lambda family: family == 120,
204
+ )
205
+ monkeypatch.setattr(b12x_mod, "current_platform", platform)
206
+ monkeypatch.setattr(
207
+ b12x_mod,
208
+ "_import_b12x_blockscaled",
209
+ lambda: types.SimpleNamespace(),
210
+ )
211
+ monkeypatch.setattr(b12x_mod, "_apply_b12x_nvfp4_linear", apply)
212
+
213
+ def layer(n: int):
214
+ return types.SimpleNamespace(
215
+ b12x_nvfp4_linear=True,
216
+ weight=torch.empty((n, 64), dtype=torch.uint8),
217
+ weight_scale=torch.empty((128, 8), dtype=torch.float8_e4m3fn),
218
+ input_global_scale_inv=torch.tensor(2.0),
219
+ alpha=torch.tensor(0.25),
220
+ )
221
+
222
+ layer_a = layer(48)
223
+ layer_b = layer(48)
224
+ layer_c = layer(96)
225
+ model = types.SimpleNamespace(
226
+ modules=lambda: iter([layer_a, layer_b, layer_c, types.SimpleNamespace()])
227
+ )
228
+
229
+ warmed = warmup_b12x_nvfp4_linear(
230
+ model,
231
+ max_tokens=8,
232
+ cudagraph_capture_sizes=[1, 2],
233
+ )
234
+
235
+ assert warmed == 6
236
+ assert [call[0] for call in calls] == [
237
+ torch.Size([1, 128]),
238
+ torch.Size([2, 128]),
239
+ torch.Size([8, 128]),
240
+ torch.Size([1, 128]),
241
+ torch.Size([2, 128]),
242
+ torch.Size([8, 128]),
243
+ ]
244
+ assert calls[0][1] is layer_a.weight
245
+ assert calls[3][1] is layer_c.weight
246
+ assert all(call[5] is None for call in calls)
tests/model_executor/test_b12x_warmup.py
new
+100
@@ -0,0 +1,100 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import importlib
5
+from types import SimpleNamespace
6
+
7
+import pytest
8
+import torch
9
+
10
+from vllm.model_executor.warmup.b12x_warmup import b12x_warmup
11
+
12
+
13
+@pytest.mark.parametrize(
14
+ ("module_name", "warmup_name", "import_name"),
15
+ [
16
+ (
17
+ "vllm.model_executor.kernels.linear.scaled_mm.b12x_block",
18
+ "warmup_b12x_block_fp8_linear",
19
+ "_import_b12x_blockscaled",
20
+ ),
21
+ (
22
+ "vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor",
23
+ "warmup_b12x_tensor_fp8_linear",
24
+ "_import_b12x_tensor_fp8",
25
+ ),
26
+ (
27
+ "vllm.model_executor.kernels.linear.mxfp8.b12x",
28
+ "warmup_b12x_mxfp8_linear",
29
+ "_import_b12x_mxfp8",
30
+ ),
31
+ (
32
+ "vllm.model_executor.kernels.linear.mxfp4.b12x",
33
+ "warmup_b12x_mxfp4_linear",
34
+ "_import_b12x_blockscaled",
35
+ ),
36
+ (
37
+ "vllm.model_executor.kernels.linear.nvfp4.b12x",
38
+ "warmup_b12x_nvfp4_linear",
39
+ "_import_b12x_blockscaled",
40
+ ),
41
+ ],
42
+)
43
+def test_b12x_linear_warmup_skips_unused_provider(
44
+ monkeypatch,
45
+ module_name: str,
46
+ warmup_name: str,
47
+ import_name: str,
48
+) -> None:
49
+ module = importlib.import_module(module_name)
50
+ platform = SimpleNamespace(
51
+ is_cuda=lambda: True,
52
+ is_device_capability_family=lambda family: family == 120,
53
+ )
54
+ monkeypatch.setattr(module, "current_platform", platform)
55
+
56
+ def fail_import():
57
+ pytest.fail("unused B12X provider was imported")
58
+
59
+ monkeypatch.setattr(module, import_name, fail_import)
60
+
61
+ warmup = getattr(module, warmup_name)
62
+ assert warmup(torch.nn.Module(), max_tokens=128) == 0
63
+
64
+
65
+def test_b12x_warmup_covers_linear_serving_shapes(monkeypatch) -> None:
66
+ import vllm.model_executor.warmup.b12x_warmup as warmup_mod
67
+
68
+ model = torch.nn.Module()
69
+ linear_calls: list[tuple[torch.nn.Module, dict]] = []
70
+
71
+ def linear_warmup(model, **kwargs):
72
+ linear_calls.append((model, kwargs))
73
+ return 0
74
+
75
+ monkeypatch.setattr(warmup_mod, "warmup_b12x_block_fp8_linear", linear_warmup)
76
+ monkeypatch.setattr(warmup_mod, "warmup_b12x_mxfp4_linear", linear_warmup)
77
+ monkeypatch.setattr(warmup_mod, "warmup_b12x_mxfp8_linear", linear_warmup)
78
+ monkeypatch.setattr(warmup_mod, "warmup_b12x_nvfp4_linear", linear_warmup)
79
+ monkeypatch.setattr(warmup_mod, "warmup_b12x_tensor_fp8_linear", linear_warmup)
80
+
81
+ worker = SimpleNamespace(
82
+ get_model=lambda: model,
83
+ scheduler_config=SimpleNamespace(
84
+ max_num_batched_tokens=256,
85
+ max_num_scheduled_tokens=320,
86
+ ),
87
+ model_config=SimpleNamespace(dtype=torch.float16),
88
+ vllm_config=SimpleNamespace(
89
+ compilation_config=SimpleNamespace(compile_sizes=[32, "dynamic", 64])
90
+ ),
91
+ )
92
+
93
+ b12x_warmup(worker, [8, 16])
94
+
95
+ expected_linear_kwargs = {
96
+ "max_tokens": 256,
97
+ "cudagraph_capture_sizes": [8, 16],
98
+ "output_dtype": torch.float16,
99
+ }
100
+ assert linear_calls == [(model, expected_linear_kwargs)] * 5
vllm/config/kernel.py
+2
@@ -147,6 +147,7 @@ LinearBackend = Literal[
147
"flashinfer_trtllm",
148
"flashinfer_cudnn",
149
"flashinfer_b12x",
150
+ "b12x",
151
"marlin",
152
"humming",
153
"triton",
@@ -223,6 +224,7 @@ class KernelConfig:
224
- "flashinfer_trtllm": Use FlashInfer with TensorRT-LLM kernels
225
- "flashinfer_cudnn": Use FlashInfer with cuDNN kernels
226
- "flashinfer_b12x": Use FlashInfer b12x CuteDSL NVFP4 GEMM (SM120+)
227
+ - "b12x": Use native B12X FP8 and FP4 linear kernels on SM12x
228
- "marlin": Use Marlin kernels
229
- "triton": Use Triton-based kernels
230
- "deep_gemm": Use DeepGEMM kernels
vllm/model_executor/kernels/linear/__init__.py
+32
@@ -77,6 +77,9 @@ from vllm.model_executor.kernels.linear.mxfp4 import (
77
from vllm.model_executor.kernels.linear.mxfp4.aiter import (
78
AiterMxfp4LinearKernel,
79
)
80
+from vllm.model_executor.kernels.linear.mxfp4.b12x import (
81
+ B12xMxFp4LinearKernel,
82
+)
83
from vllm.model_executor.kernels.linear.mxfp4.emulation import (
84
EmulationMxfp4LinearKernel,
85
)
@@ -103,6 +106,9 @@ from vllm.model_executor.kernels.linear.mxfp8 import (
106
Mxfp8LinearKernel,
107
Mxfp8LinearLayerConfig,
108
)
109
+from vllm.model_executor.kernels.linear.mxfp8.b12x import (
110
+ B12xMxfp8LinearKernel,
111
+)
112
from vllm.model_executor.kernels.linear.mxfp8.emulation import (
113
EmulationMxfp8LinearKernel,
114
)
@@ -126,6 +132,9 @@ from vllm.model_executor.kernels.linear.nvfp4 import (
132
NvFp4LinearKernel,
133
NvFp4LinearLayerConfig,
134
)
135
+from vllm.model_executor.kernels.linear.nvfp4.b12x import (
136
+ B12xNvFp4LinearKernel,
137
+)
138
from vllm.model_executor.kernels.linear.nvfp4.cutlass import (
139
CutlassNvFp4LinearKernel,
140
)
@@ -163,6 +172,12 @@ from vllm.model_executor.kernels.linear.scaled_mm.aiter import (
172
AiterPerTokenFp8ScaledMMLinearKernel,
173
AiterPreshuffledPerTokenFp8ScaledMMLinearKernel,
174
)
175
+from vllm.model_executor.kernels.linear.scaled_mm.b12x_block import (
176
+ B12xFp8BlockScaledMMKernel,
177
+)
178
+from vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor import (
179
+ B12xTensorFP8ScaledMMLinearKernel,
180
+)
181
from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
182
CPUFp8BlockScaledMMKernel,
183
CPUInt8ScaledMMLinearKernel,
@@ -228,6 +243,13 @@ def _get_linear_backend() -> str:
243
# set are considered candidates. If none can implement the layer config,
244
# an error is raised to respect the user's explicit intent.
245
_LINEAR_BACKEND_KERNEL_MAP: dict[str, set[type]] = {
246
+ "b12x": {
247
+ B12xFp8BlockScaledMMKernel,
248
+ B12xMxFp4LinearKernel,
249
+ B12xMxfp8LinearKernel,
250
+ B12xNvFp4LinearKernel,
251
+ B12xTensorFP8ScaledMMLinearKernel,
252
+ },
253
"cutlass": {
254
CutlassInt8ScaledMMLinearKernel,
255
CutlassFP8ScaledMMLinearKernel,
@@ -377,6 +399,7 @@ _POSSIBLE_FP8_KERNELS: dict[PlatformEnum, list[type[FP8ScaledMMLinearKernel]]] =
399
MarlinFP8ScaledMMLinearKernel,
400
FlashInferFP8ScaledMMLinearKernel,
401
CutlassFP8ScaledMMLinearKernel,
402
+ B12xTensorFP8ScaledMMLinearKernel,
403
PerTensorTorchFP8ScaledMMLinearKernel,
404
ChannelWiseTorchFP8ScaledMMLinearKernel,
405
HummingFP8ScaledMMLinearKernel,
@@ -411,6 +434,7 @@ _POSSIBLE_FP8_BLOCK_KERNELS: dict[
434
FlashInferFp8DeepGEMMDynamicBlockScaledKernel,
435
DeepGemmFp8BlockScaledMMKernel,
436
CutlassFp8BlockScaledMMKernel,
437
+ B12xFp8BlockScaledMMKernel,
438
MarlinFP8ScaledMMLinearKernel,
439
TritonFp8BlockScaledMMKernel,
440
HummingFP8ScaledMMLinearKernel,
@@ -482,6 +506,7 @@ _POSSIBLE_MXFP8_KERNELS: dict[PlatformEnum, list[type[Mxfp8LinearKernel]]] = {
506
FlashInferCutedslMxfp8LinearKernel,
507
FlashInferCutlassMxfp8LinearKernel,
508
MarlinMxfp8LinearKernel,
509
+ B12xMxfp8LinearKernel,
510
EmulationMxfp8LinearKernel,
511
HummingMxfp8LinearKernel,
512
],
@@ -507,6 +532,7 @@ _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
532
FlashInferTrtllmNvFp4LinearKernel,
533
FlashInferCudnnNvFp4LinearKernel,
534
FbgemmNvFp4LinearKernel,
535
+ B12xNvFp4LinearKernel,
536
EmulationNvFp4LinearKernel,
537
HummingNvFp4LinearKernel,
538
],
@@ -529,6 +555,7 @@ _POSSIBLE_MXFP4_KERNELS: dict[PlatformEnum, list[type[MxFp4LinearKernel]]] = {
555
FlashInferMxFp4LinearKernel,
556
MarlinMxFp4LinearKernel,
557
HummingMxFp4LinearKernel,
558
+ B12xMxFp4LinearKernel,
559
EmulationMxfp4LinearKernel,
560
],
561
PlatformEnum.ROCM: [
@@ -1175,6 +1202,9 @@ __all__ = [
1202
"init_mxfp8_linear_kernel",
1203
"Mxfp8LinearKernel",
1204
"Mxfp8LinearLayerConfig",
1205
+ "B12xMxfp8LinearKernel",
1206
+ "B12xMxFp4LinearKernel",
1207
+ "B12xNvFp4LinearKernel",
1208
"init_mxfp4_linear_kernel",
1209
"MxFp4LinearKernel",
1210
"MxFp4LinearLayerConfig",
@@ -1203,4 +1233,6 @@ __all__ = [
1233
"_KernelT",
1234
"DeepGemmFp8BlockScaledMMKernel",
1235
"FlashInferFp8DeepGEMMDynamicBlockScaledKernel",
1236
+ "B12xFp8BlockScaledMMKernel",
1237
+ "B12xTensorFP8ScaledMMLinearKernel",
1238
]
vllm/model_executor/kernels/linear/mxfp4/b12x.py
new
+170
@@ -0,0 +1,170 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from __future__ import annotations
5
+
6
+from collections.abc import Iterable
7
+from typing import Any
8
+
9
+import torch
10
+
11
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
12
+ kMxfp4Dynamic,
13
+)
14
+from vllm.model_executor.utils import replace_parameter
15
+from vllm.platforms import current_platform
16
+from vllm.utils.b12x import (
17
+ b12x_warmup_token_counts,
18
+)
19
+from vllm.utils.b12x import (
20
+ get_b12x_blockscaled as _import_b12x_blockscaled,
21
+)
22
+from vllm.utils.b12x import get_b12x_intrinsics as _import_b12x_intrinsics
23
+
24
+from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
25
+
26
+
27
+def _apply_b12x_mxfp4_linear(
28
+ x: torch.Tensor,
29
+ weight: torch.Tensor,
30
+ weight_scale_storage: torch.Tensor,
31
+ bias: torch.Tensor | None,
32
+) -> torch.Tensor:
33
+ from vllm.utils.flashinfer import flashinfer_mxfp4_quantize
34
+
35
+ blockscaled = _import_b12x_blockscaled()
36
+ assert blockscaled is not None
37
+
38
+ output_size = int(weight.shape[0])
39
+ output_shape = [*x.shape[:-1], output_size]
40
+ x_2d = x.reshape(-1, x.shape[-1]).contiguous()
41
+ x_packed, x_scale_swizzled = flashinfer_mxfp4_quantize(x_2d, backend="cute-dsl")
42
+ output = blockscaled.mm_mxfp4(
43
+ x_packed,
44
+ x_scale_swizzled,
45
+ weight,
46
+ weight_scale_storage,
47
+ out_dtype=x.dtype,
48
+ )
49
+ if bias is not None:
50
+ output = output + bias
51
+ return output.view(*output_shape)
52
+
53
+
54
+def warmup_b12x_mxfp4_linear(
55
+ model: torch.nn.Module,
56
+ *,
57
+ max_tokens: int,
58
+ cudagraph_capture_sizes: Iterable[int] = (),
59
+ output_dtype: torch.dtype = torch.bfloat16,
60
+) -> int:
61
+ if not current_platform.is_cuda():
62
+ return 0
63
+ if not current_platform.is_device_capability_family(120):
64
+ return 0
65
+ if output_dtype not in (torch.bfloat16, torch.float16):
66
+ output_dtype = torch.bfloat16
67
+ layer_map: dict[tuple[Any, ...], torch.nn.Module] = {}
68
+ for layer in model.modules():
69
+ if not getattr(layer, "b12x_mxfp4_linear", False):
70
+ continue
71
+ weight = layer.weight
72
+ weight_scale = layer.weight_scale
73
+ n, packed_k = map(int, weight.shape)
74
+ signature = (
75
+ weight.device,
76
+ n,
77
+ packed_k * 2,
78
+ weight.dtype,
79
+ weight_scale.dtype,
80
+ output_dtype,
81
+ )
82
+ layer_map.setdefault(signature, layer)
83
+ if not layer_map:
84
+ return 0
85
+ if _import_b12x_blockscaled() is None:
86
+ return 0
87
+
88
+ token_counts = b12x_warmup_token_counts(
89
+ max_tokens=max_tokens,
90
+ cudagraph_capture_sizes=cudagraph_capture_sizes,
91
+ )
92
+ warmed = 0
93
+ last_device: torch.device | None = None
94
+
95
+ with torch.inference_mode():
96
+ for signature, layer in layer_map.items():
97
+ weight = layer.weight
98
+ weight_scale = layer.weight_scale
99
+ k = signature[2]
100
+ last_device = weight.device
101
+ for tokens in token_counts:
102
+ source = torch.zeros(
103
+ (tokens, k),
104
+ dtype=output_dtype,
105
+ device=weight.device,
106
+ )
107
+ _apply_b12x_mxfp4_linear(
108
+ source,
109
+ weight,
110
+ weight_scale,
111
+ None,
112
+ )
113
+ warmed += 1
114
+
115
+ if warmed > 0 and last_device is not None and last_device.type == "cuda":
116
+ torch.accelerator.synchronize(last_device)
117
+
118
+ return warmed
119
+
120
+
121
+class B12xMxFp4LinearKernel(MxFp4LinearKernel):
122
+ """MXFP4 linear through the native B12X SM120 dense GEMM."""
123
+
124
+ @classmethod
125
+ def is_supported(
126
+ cls, compute_capability: int | None = None
127
+ ) -> tuple[bool, str | None]:
128
+ del compute_capability
129
+ if not current_platform.is_cuda():
130
+ return False, "B12X MXFP4 kernels are only available on CUDA"
131
+ if not current_platform.is_device_capability_family(120):
132
+ return False, "B12X MXFP4 kernels require a Blackwell 12x device"
133
+ blockscaled = _import_b12x_blockscaled()
134
+ if blockscaled is None or _import_b12x_intrinsics() is None:
135
+ return False, "Install the B12X backend with `pip install vllm[b12x]`"
136
+ if not blockscaled.is_supported():
137
+ return False, "b12x native MXFP4 GEMM is not supported"
138
+ return True, None
139
+
140
+ @classmethod
141
+ def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
142
+ if config.activation_quant_key != kMxfp4Dynamic:
143
+ return False, "B12X MXFP4 GEMM requires dynamic MXFP4 activations"
144
+ return True, None
145
+
146
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
147
+ intrinsics = _import_b12x_intrinsics()
148
+ assert intrinsics is not None
149
+ replace_parameter(
150
+ layer,
151
+ "weight_scale",
152
+ intrinsics.swizzle_block_scale(layer.weight_scale.data),
153
+ )
154
+ layer.b12x_mxfp4_linear = True
155
+
156
+ def apply_weights(
157
+ self,
158
+ layer: torch.nn.Module,
159
+ x: torch.Tensor,
160
+ bias: torch.Tensor | None = None,
161
+ ) -> torch.Tensor:
162
+ return _apply_b12x_mxfp4_linear(
163
+ x,
164
+ layer.weight,
165
+ layer.weight_scale,
166
+ bias,
167
+ )
168
+
169
+
170
+__all__ = ["B12xMxFp4LinearKernel", "warmup_b12x_mxfp4_linear"]
vllm/model_executor/kernels/linear/mxfp8/b12x.py
new
+193
@@ -0,0 +1,193 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from __future__ import annotations
5
+
6
+from collections.abc import Iterable
7
+from typing import Any
8
+
9
+import torch
10
+
11
+from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
12
+ MXFP8_BLOCK_SIZE,
13
+ MXFP8_SCALE_DTYPE,
14
+ MXFP8_VALUE_DTYPE,
15
+)
16
+from vllm.model_executor.utils import replace_parameter
17
+from vllm.platforms import current_platform
18
+from vllm.utils.b12x import (
19
+ b12x_warmup_token_counts,
20
+ reuse_packed_weight_storage,
21
+)
22
+from vllm.utils.b12x import (
23
+ get_b12x_mxfp8_linear as _import_b12x_mxfp8,
24
+)
25
+from vllm.utils.torch_utils import current_stream
26
+
27
+from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
28
+
29
+
30
+def _b12x_mxfp8_expected_m(tokens: int) -> int:
31
+ return max(1, int(tokens))
32
+
33
+
34
+def _apply_b12x_mxfp8_packed_linear(
35
+ layer: torch.nn.Module,
36
+ x: torch.Tensor,
37
+ bias: torch.Tensor | None,
38
+) -> torch.Tensor:
39
+ packed_weight = getattr(layer, "b12x_mxfp8_packed_weight", None)
40
+ if packed_weight is None:
41
+ raise RuntimeError(
42
+ "b12x MXFP8 packed weights are missing; "
43
+ "process_weights_after_loading did not run for this layer"
44
+ )
45
+
46
+ input_2d = x.reshape(-1, x.shape[-1]).contiguous()
47
+ output_shape = [*x.shape[:-1], int(packed_weight.out_features)]
48
+
49
+ mxfp8 = _import_b12x_mxfp8()
50
+ assert mxfp8 is not None
51
+ output = mxfp8.mm(
52
+ input_2d,
53
+ packed_weight,
54
+ bias=bias,
55
+ expected_m=_b12x_mxfp8_expected_m(int(input_2d.shape[0])),
56
+ )
57
+ return output.view(*output_shape)
58
+
59
+
60
+def warmup_b12x_mxfp8_linear(
61
+ model: torch.nn.Module,
62
+ *,
63
+ max_tokens: int,
64
+ cudagraph_capture_sizes: Iterable[int] = (),
65
+ output_dtype: torch.dtype = torch.bfloat16,
66
+) -> int:
67
+ if not current_platform.is_cuda():
68
+ return 0
69
+ if not current_platform.is_device_capability_family(120):
70
+ return 0
71
+ if output_dtype not in (torch.bfloat16, torch.float16):
72
+ output_dtype = torch.bfloat16
73
+
74
+ layer_map: dict[tuple[Any, ...], Any] = {}
75
+ for layer in model.modules():
76
+ packed_weight = getattr(layer, "b12x_mxfp8_packed_weight", None)
77
+ if packed_weight is None:
78
+ continue
79
+ device = torch.device(packed_weight.weight.values.device)
80
+ signature = (
81
+ device,
82
+ int(packed_weight.in_features),
83
+ int(packed_weight.padded_in_features),
84
+ int(packed_weight.out_features),
85
+ output_dtype,
86
+ )
87
+ layer_map.setdefault(signature, packed_weight)
88
+ if not layer_map:
89
+ return 0
90
+
91
+ mxfp8 = _import_b12x_mxfp8()
92
+ if mxfp8 is None:
93
+ return 0
94
+
95
+ token_counts = b12x_warmup_token_counts(
96
+ max_tokens=max_tokens,
97
+ cudagraph_capture_sizes=cudagraph_capture_sizes,
98
+ )
99
+ warmed = 0
100
+ last_device: torch.device | None = None
101
+
102
+ with torch.inference_mode():
103
+ for signature, packed_weight in layer_map.items():
104
+ device = signature[0]
105
+ last_device = device
106
+ for tokens in token_counts:
107
+ source = torch.zeros(
108
+ (tokens, int(packed_weight.in_features)),
109
+ dtype=output_dtype,
110
+ device=device,
111
+ )
112
+ mxfp8.mm(
113
+ source,
114
+ packed_weight,
115
+ expected_m=_b12x_mxfp8_expected_m(tokens),
116
+ stream=current_stream().cuda_stream,
117
+ )
118
+ warmed += 1
119
+
120
+ if warmed > 0 and last_device is not None and last_device.type == "cuda":
121
+ torch.accelerator.synchronize(last_device)
122
+
123
+ return warmed
124
+
125
+
126
+class B12xMxfp8LinearKernel(Mxfp8LinearKernel):
127
+ """ModelOpt MXFP8 linear through the native b12x SM120 dense GEMM path."""
128
+
129
+ @classmethod
130
+ def is_supported(
131
+ cls,
132
+ compute_capability: int | None = None,
133
+ ) -> tuple[bool, str | None]:
134
+ del compute_capability
135
+ if not current_platform.is_cuda():
136
+ return False, "b12x MXFP8 kernels are only available on CUDA"
137
+ if not current_platform.is_device_capability_family(120):
138
+ return False, "b12x MXFP8 kernels require a Blackwell 12x device"
139
+ mxfp8 = _import_b12x_mxfp8()
140
+ if mxfp8 is None:
141
+ return False, "Install the B12X backend with `pip install vllm[b12x]`"
142
+ if not mxfp8.is_supported():
143
+ return False, "b12x.gemm.mxfp8_linear is not supported"
144
+ return True, None
145
+
146
+ @classmethod
147
+ def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]:
148
+ del c
149
+ return True, None
150
+
151
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
152
+ weight = layer.weight.data
153
+ assert weight.dtype == MXFP8_VALUE_DTYPE, (
154
+ f"b12x MXFP8 requires {MXFP8_VALUE_DTYPE}, got {weight.dtype}"
155
+ )
156
+ assert weight.ndim == 2, f"b12x MXFP8 weight must be 2D, got {weight.ndim}D"
157
+ assert hasattr(layer, "weight_scale"), "b12x MXFP8 linear requires weight_scale"
158
+
159
+ out_features, in_features = map(int, weight.shape)
160
+ assert in_features % MXFP8_BLOCK_SIZE == 0, (
161
+ "b12x MXFP8 requires input features divisible by "
162
+ f"{MXFP8_BLOCK_SIZE}, got {in_features}"
163
+ )
164
+ weight_scale = layer.weight_scale.data
165
+ assert weight_scale.dtype == MXFP8_SCALE_DTYPE, (
166
+ f"b12x MXFP8 requires {MXFP8_SCALE_DTYPE} weight_scale, "
167
+ f"got {weight_scale.dtype}"
168
+ )
169
+ assert weight_scale.ndim == 2, (
170
+ f"b12x MXFP8 weight_scale must be 2D, got {weight_scale.ndim}D"
171
+ )
172
+
173
+ mxfp8 = _import_b12x_mxfp8()
174
+ assert mxfp8 is not None
175
+ scale_k = in_features // MXFP8_BLOCK_SIZE
176
+ packed_weight = mxfp8.pack_weight(
177
+ weight[:out_features, :in_features].detach(),
178
+ weight_scale[:out_features, :scale_k].detach(),
179
+ )
180
+ layer.b12x_mxfp8_packed_weight = reuse_packed_weight_storage(
181
+ getattr(layer, "b12x_mxfp8_packed_weight", None),
182
+ packed_weight,
183
+ )
184
+ replace_parameter(layer, "weight", weight.new_empty((0,)))
185
+ replace_parameter(layer, "weight_scale", weight_scale.new_empty((0,)))
186
+
187
+ def apply_weights(
188
+ self,
189
+ layer: torch.nn.Module,
190
+ x: torch.Tensor,
191
+ bias: torch.Tensor | None = None,
192
+ ) -> torch.Tensor:
193
+ return _apply_b12x_mxfp8_packed_linear(layer, x, bias)
vllm/model_executor/kernels/linear/nvfp4/b12x.py
new
+176
@@ -0,0 +1,176 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from __future__ import annotations
5
+
6
+from collections.abc import Iterable
7
+from typing import Any
8
+
9
+import torch
10
+
11
+from vllm._custom_ops import scaled_fp4_quant
12
+from vllm.model_executor.utils import replace_parameter
13
+from vllm.platforms import current_platform
14
+from vllm.utils.b12x import (
15
+ b12x_warmup_token_counts,
16
+)
17
+from vllm.utils.b12x import (
18
+ get_b12x_blockscaled as _import_b12x_blockscaled,
19
+)
20
+from vllm.utils.b12x import get_b12x_intrinsics as _import_b12x_intrinsics
21
+
22
+from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
23
+
24
+
25
+def _apply_b12x_nvfp4_linear(
26
+ x: torch.Tensor,
27
+ weight: torch.Tensor,
28
+ weight_scale_storage: torch.Tensor,
29
+ input_global_scale_inv: torch.Tensor,
30
+ alpha: torch.Tensor,
31
+ bias: torch.Tensor | None,
32
+) -> torch.Tensor:
33
+ blockscaled = _import_b12x_blockscaled()
34
+ assert blockscaled is not None
35
+
36
+ output_size = int(weight.shape[0])
37
+ output_shape = [*x.shape[:-1], output_size]
38
+ x_2d = x.reshape(-1, x.shape[-1])
39
+ x_packed, x_scale_swizzled = scaled_fp4_quant(
40
+ x_2d,
41
+ input_global_scale_inv,
42
+ is_sf_swizzled_layout=True,
43
+ )
44
+ output = blockscaled.mm_nvfp4(
45
+ x_packed,
46
+ x_scale_swizzled,
47
+ weight,
48
+ weight_scale_storage,
49
+ alpha,
50
+ out_dtype=x.dtype,
51
+ )
52
+ if bias is not None:
53
+ output = output + bias
54
+ return output.view(*output_shape)
55
+
56
+
57
+def warmup_b12x_nvfp4_linear(
58
+ model: torch.nn.Module,
59
+ *,
60
+ max_tokens: int,
61
+ cudagraph_capture_sizes: Iterable[int] = (),
62
+ output_dtype: torch.dtype = torch.bfloat16,
63
+) -> int:
64
+ if not current_platform.is_cuda():
65
+ return 0
66
+ if not current_platform.is_device_capability_family(120):
67
+ return 0
68
+ if output_dtype not in (torch.bfloat16, torch.float16):
69
+ output_dtype = torch.bfloat16
70
+ layer_map: dict[tuple[Any, ...], torch.nn.Module] = {}
71
+ for layer in model.modules():
72
+ if not getattr(layer, "b12x_nvfp4_linear", False):
73
+ continue
74
+ weight = layer.weight
75
+ weight_scale = layer.weight_scale
76
+ n, packed_k = map(int, weight.shape)
77
+ signature = (
78
+ weight.device,
79
+ n,
80
+ packed_k * 2,
81
+ weight.dtype,
82
+ weight_scale.dtype,
83
+ output_dtype,
84
+ )
85
+ layer_map.setdefault(signature, layer)
86
+ if not layer_map:
87
+ return 0
88
+ if _import_b12x_blockscaled() is None:
89
+ return 0
90
+
91
+ token_counts = b12x_warmup_token_counts(
92
+ max_tokens=max_tokens,
93
+ cudagraph_capture_sizes=cudagraph_capture_sizes,
94
+ )
95
+ warmed = 0
96
+ last_device: torch.device | None = None
97
+
98
+ with torch.inference_mode():
99
+ for signature, layer in layer_map.items():
100
+ weight = layer.weight
101
+ weight_scale = layer.weight_scale
102
+ k = signature[2]
103
+ last_device = weight.device
104
+ for tokens in token_counts:
105
+ source = torch.zeros(
106
+ (tokens, k),
107
+ dtype=output_dtype,
108
+ device=weight.device,
109
+ )
110
+ _apply_b12x_nvfp4_linear(
111
+ source,
112
+ weight,
113
+ weight_scale,
114
+ layer.input_global_scale_inv,
115
+ layer.alpha,
116
+ None,
117
+ )
118
+ warmed += 1
119
+
120
+ if warmed > 0 and last_device is not None and last_device.type == "cuda":
121
+ torch.accelerator.synchronize(last_device)
122
+
123
+ return warmed
124
+
125
+
126
+class B12xNvFp4LinearKernel(NvFp4LinearKernel):
127
+ """ModelOpt NVFP4 linear through the native B12X SM120 dense GEMM."""
128
+
129
+ @classmethod
130
+ def is_supported(
131
+ cls, compute_capability: int | None = None
132
+ ) -> tuple[bool, str | None]:
133
+ del compute_capability
134
+ if not current_platform.is_cuda():
135
+ return False, "B12X NVFP4 kernels are only available on CUDA"
136
+ if not current_platform.is_device_capability_family(120):
137
+ return False, "B12X NVFP4 kernels require a Blackwell 12x device"
138
+ blockscaled = _import_b12x_blockscaled()
139
+ if blockscaled is None or _import_b12x_intrinsics() is None:
140
+ return False, "Install the B12X backend with `pip install vllm[b12x]`"
141
+ if not blockscaled.is_supported():
142
+ return False, "b12x native NVFP4 GEMM is not supported"
143
+ return True, None
144
+
145
+ @classmethod
146
+ def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
147
+ del config
148
+ return True, None
149
+
150
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
151
+ intrinsics = _import_b12x_intrinsics()
152
+ assert intrinsics is not None
153
+ replace_parameter(
154
+ layer,
155
+ "weight_scale",
156
+ intrinsics.swizzle_block_scale(layer.weight_scale.data),
157
+ )
158
+ layer.b12x_nvfp4_linear = True
159
+
160
+ def apply_weights(
161
+ self,
162
+ layer: torch.nn.Module,
163
+ x: torch.Tensor,
164
+ bias: torch.Tensor | None = None,
165
+ ) -> torch.Tensor:
166
+ return _apply_b12x_nvfp4_linear(
167
+ x,
168
+ layer.weight,
169
+ layer.weight_scale,
170
+ layer.input_global_scale_inv,
171
+ layer.alpha,
172
+ bias,
173
+ )
174
+
175
+
176
+__all__ = ["B12xNvFp4LinearKernel", "warmup_b12x_nvfp4_linear"]
vllm/model_executor/kernels/linear/scaled_mm/b12x_block.py
new
+214
@@ -0,0 +1,214 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from __future__ import annotations
5
+
6
+from collections.abc import Iterable
7
+from typing import Any
8
+
9
+import torch
10
+
11
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
12
+ _upcast_e8m0_to_fp32,
13
+)
14
+from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
15
+from vllm.model_executor.utils import replace_parameter
16
+from vllm.platforms import current_platform
17
+from vllm.utils.b12x import (
18
+ b12x_warmup_token_counts,
19
+)
20
+from vllm.utils.b12x import (
21
+ get_b12x_blockscaled as _import_b12x_blockscaled,
22
+)
23
+
24
+from .BlockScaledMMLinearKernel import (
25
+ Fp8BlockScaledMMLinearKernel,
26
+ FP8ScaledMMLinearLayerConfig,
27
+)
28
+
29
+
30
+def _run_b12x_fp8_block_scaled_mm(
31
+ a: torch.Tensor,
32
+ weight: torch.Tensor,
33
+ a_scale: torch.Tensor,
34
+ weight_scale: torch.Tensor,
35
+ out_dtype: torch.dtype,
36
+) -> torch.Tensor:
37
+ blockscaled = _import_b12x_blockscaled()
38
+ assert blockscaled is not None
39
+
40
+ return blockscaled.mm_block_fp8(
41
+ a,
42
+ a_scale,
43
+ weight,
44
+ weight_scale,
45
+ out_dtype=out_dtype,
46
+ )
47
+
48
+
49
+def warmup_b12x_block_fp8_linear(
50
+ model: torch.nn.Module,
51
+ *,
52
+ max_tokens: int,
53
+ cudagraph_capture_sizes: Iterable[int] = (),
54
+ output_dtype: torch.dtype = torch.bfloat16,
55
+) -> int:
56
+ if not current_platform.is_cuda():
57
+ return 0
58
+ if not current_platform.is_device_capability_family(120):
59
+ return 0
60
+ if output_dtype not in (torch.bfloat16, torch.float16):
61
+ output_dtype = torch.bfloat16
62
+
63
+ layer_map: dict[tuple[Any, ...], torch.nn.Module] = {}
64
+ for layer in model.modules():
65
+ if not getattr(layer, "b12x_block_fp8_linear", False):
66
+ continue
67
+ weight = layer.weight
68
+ weight_scale = getattr(layer, "weight_scale_inv", None)
69
+ if weight_scale is None:
70
+ weight_scale = layer.weight_scale
71
+ n, k = map(int, weight.shape)
72
+ signature = (
73
+ weight.device,
74
+ n,
75
+ k,
76
+ weight.dtype,
77
+ weight_scale.dtype,
78
+ output_dtype,
79
+ )
80
+ layer_map.setdefault(signature, layer)
81
+ if not layer_map:
82
+ return 0
83
+
84
+ blockscaled = _import_b12x_blockscaled()
85
+ if blockscaled is None:
86
+ return 0
87
+ token_counts = b12x_warmup_token_counts(
88
+ max_tokens=max_tokens,
89
+ cudagraph_capture_sizes=cudagraph_capture_sizes,
90
+ )
91
+ warmed = 0
92
+
93
+ with torch.inference_mode():
94
+ for signature, layer in layer_map.items():
95
+ weight = layer.weight
96
+ weight_scale = getattr(layer, "weight_scale_inv", None)
97
+ if weight_scale is None:
98
+ weight_scale = layer.weight_scale
99
+ k = signature[2]
100
+ for tokens in token_counts:
101
+ a = torch.empty(
102
+ (tokens, k),
103
+ dtype=weight.dtype,
104
+ device=weight.device,
105
+ )
106
+ a_scale = torch.empty(
107
+ (tokens, k // 128),
108
+ dtype=torch.float32,
109
+ device=weight.device,
110
+ )
111
+ _run_b12x_fp8_block_scaled_mm(
112
+ a,
113
+ weight,
114
+ a_scale,
115
+ weight_scale,
116
+ output_dtype,
117
+ )
118
+ warmed += 1
119
+ return warmed
120
+
121
+
122
+class B12xFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
123
+ """K128 block-FP8 linear through the native B12X SM120 dense GEMM."""
124
+
125
+ @classmethod
126
+ def is_supported(
127
+ cls,
128
+ compute_capability: int | None = None,
129
+ ) -> tuple[bool, str | None]:
130
+ del compute_capability
131
+ if not current_platform.is_cuda():
132
+ return False, "B12X FP8 kernels are only available on CUDA"
133
+ if not current_platform.is_device_capability_family(120):
134
+ return False, "B12X FP8 kernels require a Blackwell 12x device"
135
+ blockscaled = _import_b12x_blockscaled()
136
+ if blockscaled is None:
137
+ return False, "Install the B12X backend with `pip install vllm[b12x]`"
138
+ if not blockscaled.is_supported():
139
+ return False, "B12X regular block-FP8 GEMM is not supported"
140
+ return True, None
141
+
142
+ @classmethod
143
+ def can_implement(
144
+ cls,
145
+ config: FP8ScaledMMLinearLayerConfig,
146
+ ) -> tuple[bool, str | None]:
147
+ can_implement_base, reason = super().can_implement(config)
148
+ if not can_implement_base:
149
+ return can_implement_base, reason
150
+
151
+ if config.input_dtype not in (torch.bfloat16, torch.float16):
152
+ return False, "Supports only bf16/fp16 input dtype"
153
+ if config.input_dtype != config.out_dtype:
154
+ return False, "Input and output dtype must match"
155
+
156
+ act_group_shape = config.activation_quant_key.scale.group_shape
157
+ if act_group_shape != GroupShape(1, 128):
158
+ return (
159
+ False,
160
+ "Supports only dynamic per-token group activation quantization "
161
+ "with group_shape=(1,128)",
162
+ )
163
+ weight_group_shape = config.weight_quant_key.scale.group_shape
164
+ if weight_group_shape != GroupShape(128, 128):
165
+ return False, "Supports only 128x128 block-scaled FP8 weights"
166
+
167
+ out_features, in_features = config.weight_shape
168
+ if in_features <= 0 or in_features % 128 != 0:
169
+ return False, "Input features must be a positive multiple of 128"
170
+ if out_features <= 0 or out_features % 128 != 0:
171
+ return False, "Output features must be a positive multiple of 128"
172
+ return True, None
173
+
174
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
175
+ super().process_weights_after_loading(layer)
176
+ params = self._get_layer_params(layer)
177
+ if params.weight_scale_inv is not None:
178
+ weight_scale = params.weight_scale_inv
179
+ scale_attr = params.WEIGHT_SCALE_INV
180
+ else:
181
+ weight_scale = params.weight_scale
182
+ scale_attr = params.WEIGHT_SCALE
183
+ if weight_scale is not None and weight_scale.dtype in (
184
+ torch.float8_e8m0fnu,
185
+ torch.uint8,
186
+ ):
187
+ # TODO: Remove once B12X supports 128x128 UE8M0 block scales.
188
+ replace_parameter(
189
+ layer,
190
+ scale_attr,
191
+ _upcast_e8m0_to_fp32(weight_scale).contiguous(),
192
+ )
193
+ layer.b12x_block_fp8_linear = True
194
+
195
+ def apply_block_scaled_mm(
196
+ self,
197
+ A: torch.Tensor,
198
+ B: torch.Tensor,
199
+ As: torch.Tensor,
200
+ Bs: torch.Tensor,
201
+ ) -> torch.Tensor:
202
+ return _run_b12x_fp8_block_scaled_mm(
203
+ A,
204
+ B,
205
+ As,
206
+ Bs,
207
+ self.config.out_dtype,
208
+ )
209
+
210
+
211
+__all__ = [
212
+ "B12xFp8BlockScaledMMKernel",
213
+ "warmup_b12x_block_fp8_linear",
214
+]
vllm/model_executor/kernels/linear/scaled_mm/b12x_tensor.py
new
+238
@@ -0,0 +1,238 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from __future__ import annotations
5
+
6
+from collections.abc import Iterable
7
+from typing import Any
8
+
9
+import torch
10
+
11
+from vllm.model_executor.utils import replace_parameter
12
+from vllm.platforms import current_platform
13
+from vllm.utils.b12x import (
14
+ b12x_warmup_token_counts,
15
+ reuse_packed_weight_storage,
16
+)
17
+from vllm.utils.b12x import (
18
+ get_b12x_tensor_fp8_linear as _import_b12x_tensor_fp8,
19
+)
20
+from vllm.utils.torch_utils import current_stream
21
+
22
+from .ScaledMMLinearKernel import (
23
+ FP8ScaledMMLinearKernel,
24
+ FP8ScaledMMLinearLayerConfig,
25
+)
26
+
27
+
28
+def _apply_b12x_tensor_fp8_packed_linear(
29
+ layer: torch.nn.Module,
30
+ x_q: torch.Tensor,
31
+ bias: torch.Tensor | None,
32
+ out_dtype: torch.dtype,
33
+) -> torch.Tensor:
34
+ packed_weight = getattr(layer, "b12x_tensor_fp8_packed_weight", None)
35
+ if packed_weight is None:
36
+ raise RuntimeError(
37
+ "b12x tensor FP8 packed weights are missing; "
38
+ "process_weights_after_loading did not run for this layer"
39
+ )
40
+
41
+ tensor_fp8 = _import_b12x_tensor_fp8()
42
+ assert tensor_fp8 is not None
43
+
44
+ input_2d = x_q.reshape(-1, x_q.shape[-1]).contiguous()
45
+ output_shape = [*x_q.shape[:-1], int(packed_weight.out_features)]
46
+ output = tensor_fp8.mm(
47
+ input_2d,
48
+ packed_weight,
49
+ bias=bias,
50
+ out_dtype=out_dtype,
51
+ expected_m=max(1, int(input_2d.shape[0])),
52
+ )
53
+ return output.view(*output_shape)
54
+
55
+
56
+def warmup_b12x_tensor_fp8_linear(
57
+ model: torch.nn.Module,
58
+ *,
59
+ max_tokens: int,
60
+ cudagraph_capture_sizes: Iterable[int] = (),
61
+ output_dtype: torch.dtype = torch.bfloat16,
62
+) -> int:
63
+ if not current_platform.is_cuda():
64
+ return 0
65
+ if not current_platform.is_device_capability_family(120):
66
+ return 0
67
+
68
+ if output_dtype not in (torch.bfloat16, torch.float16):
69
+ output_dtype = torch.bfloat16
70
+
71
+ layer_map: dict[tuple[Any, ...], Any] = {}
72
+ for layer in model.modules():
73
+ packed_weight = getattr(
74
+ layer,
75
+ "b12x_tensor_fp8_packed_weight",
76
+ None,
77
+ )
78
+ if packed_weight is None:
79
+ continue
80
+ device = torch.device(packed_weight.values.device)
81
+ signature = (
82
+ device,
83
+ int(packed_weight.in_features),
84
+ int(packed_weight.padded_in_features),
85
+ int(packed_weight.out_features),
86
+ output_dtype,
87
+ )
88
+ layer_map.setdefault(signature, packed_weight)
89
+ if not layer_map:
90
+ return 0
91
+
92
+ tensor_fp8 = _import_b12x_tensor_fp8()
93
+ if tensor_fp8 is None:
94
+ return 0
95
+
96
+ token_counts = b12x_warmup_token_counts(
97
+ max_tokens=max_tokens,
98
+ cudagraph_capture_sizes=cudagraph_capture_sizes,
99
+ )
100
+ warmed = 0
101
+ last_device: torch.device | None = None
102
+
103
+ with torch.inference_mode():
104
+ for signature, packed_weight in layer_map.items():
105
+ last_device = signature[0]
106
+ warmed += int(
107
+ tensor_fp8.prewarm(
108
+ packed_weight,
109
+ token_counts,
110
+ out_dtype=output_dtype,
111
+ stream=current_stream().cuda_stream,
112
+ )
113
+ )
114
+
115
+ if warmed > 0 and last_device is not None and last_device.type == "cuda":
116
+ torch.accelerator.synchronize(last_device)
117
+
118
+ return warmed
119
+
120
+
121
+class B12xTensorFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
122
+ """Static per-tensor FP8 linear through the B12X SM12x dense GEMM."""
123
+
124
+ @classmethod
125
+ def is_supported(
126
+ cls,
127
+ compute_capability: int | None = None,
128
+ ) -> tuple[bool, str | None]:
129
+ del compute_capability
130
+ if not current_platform.is_cuda():
131
+ return False, "b12x tensor FP8 kernels are only available on CUDA"
132
+ if not current_platform.is_device_capability_family(120):
133
+ return False, "b12x tensor FP8 kernels require a Blackwell 12x device"
134
+ tensor_fp8 = _import_b12x_tensor_fp8()
135
+ if tensor_fp8 is None:
136
+ return False, "Install the B12X backend with `pip install vllm[b12x]`"
137
+ if not tensor_fp8.is_supported():
138
+ return False, "b12x.gemm.tensor_fp8_linear is not supported"
139
+ return True, None
140
+
141
+ @classmethod
142
+ def can_implement(
143
+ cls,
144
+ config: FP8ScaledMMLinearLayerConfig,
145
+ ) -> tuple[bool, str | None]:
146
+ activation_scale = config.activation_quant_key.scale
147
+ weight_scale = config.weight_quant_key.scale
148
+ if (
149
+ not activation_scale.static
150
+ or not activation_scale.group_shape.is_per_tensor()
151
+ ):
152
+ return False, "requires static per-tensor activation scales"
153
+ if not weight_scale.static or not weight_scale.group_shape.is_per_tensor():
154
+ return False, "requires static per-tensor weight scales"
155
+ if config.input_dtype not in (torch.bfloat16, torch.float16):
156
+ return False, "supports only bf16/fp16 input dtype"
157
+ if config.out_dtype not in (torch.bfloat16, torch.float16):
158
+ return False, "supports only bf16/fp16 output dtype"
159
+ out_features, in_features = config.weight_shape
160
+ if out_features <= 0 or in_features <= 0 or in_features % 32 != 0:
161
+ return False, "weight dimensions must be positive with K divisible by 32"
162
+ return True, None
163
+
164
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
165
+ weight, weight_scale, input_scale, _ = self._get_layer_params(layer)
166
+ if weight.dtype != torch.float8_e4m3fn:
167
+ raise ValueError(
168
+ f"b12x tensor FP8 requires float8_e4m3fn weight, got {weight.dtype}"
169
+ )
170
+ if weight_scale.numel() != 1 or input_scale is None or input_scale.numel() != 1:
171
+ raise ValueError(
172
+ "b12x tensor FP8 requires scalar weight and activation scales"
173
+ )
174
+
175
+ out_features, in_features = map(int, self.config.weight_shape)
176
+ if tuple(weight.shape) != (in_features, out_features):
177
+ raise ValueError(
178
+ "b12x tensor FP8 expects the processed weight in [K,N] layout, "
179
+ f"got {tuple(weight.shape)} for N={out_features}, K={in_features}"
180
+ )
181
+
182
+ tensor_fp8 = _import_b12x_tensor_fp8()
183
+ assert tensor_fp8 is not None
184
+ output_scale = (
185
+ input_scale.detach().to(torch.float32).reshape(1)
186
+ * weight_scale.detach().to(torch.float32).reshape(1)
187
+ ).contiguous()
188
+ packed_weight = tensor_fp8.pack_weight(
189
+ weight.detach().T.contiguous(),
190
+ output_scale,
191
+ )
192
+ layer.b12x_tensor_fp8_packed_weight = reuse_packed_weight_storage(
193
+ getattr(layer, "b12x_tensor_fp8_packed_weight", None),
194
+ packed_weight,
195
+ )
196
+ weight_name, weight_scale_name, _, _ = self.layer_param_names
197
+ replace_parameter(layer, weight_name, weight.new_empty((0,)))
198
+ replace_parameter(layer, weight_scale_name, weight_scale.new_empty((0,)))
199
+
200
+ def apply_weights(
201
+ self,
202
+ layer: torch.nn.Module,
203
+ x: torch.Tensor,
204
+ bias: torch.Tensor | None = None,
205
+ ) -> torch.Tensor:
206
+ if not isinstance(x, torch.Tensor):
207
+ raise TypeError("b12x tensor FP8 linear requires a Tensor input")
208
+ _, _, input_scale, input_scale_ub = self._get_layer_params(layer)
209
+ input_2d = x.reshape(-1, x.shape[-1])
210
+ x_q, _ = self.quant_fp8(input_2d, input_scale, input_scale_ub)
211
+ out_dtype = self.config.out_dtype
212
+ output = _apply_b12x_tensor_fp8_packed_linear(
213
+ layer,
214
+ x_q,
215
+ bias,
216
+ out_dtype,
217
+ )
218
+ return output.view(*x.shape[:-1], output.shape[-1])
219
+
220
+ def apply_scaled_mm(
221
+ self,
222
+ *,
223
+ A: torch.Tensor,
224
+ B: torch.Tensor,
225
+ out_dtype: torch.dtype,
226
+ As: torch.Tensor,
227
+ Bs: torch.Tensor,
228
+ bias: torch.Tensor | None,
229
+ output_shape: list,
230
+ ) -> torch.Tensor:
231
+ del A, B, out_dtype, As, Bs, bias, output_shape
232
+ raise NotImplementedError("b12x tensor FP8 linear overrides apply_weights")
233
+
234
+
235
+__all__ = [
236
+ "B12xTensorFP8ScaledMMLinearKernel",
237
+ "warmup_b12x_tensor_fp8_linear",
238
+]
vllm/model_executor/warmup/b12x_warmup.py
new
+60
@@ -0,0 +1,60 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Warm B12X JIT kernels used by a loaded model."""
4
+
5
+from typing import TYPE_CHECKING
6
+
7
+import torch
8
+
9
+from vllm.logger import init_logger
10
+from vllm.model_executor.kernels.linear.mxfp4.b12x import (
11
+ warmup_b12x_mxfp4_linear,
12
+)
13
+from vllm.model_executor.kernels.linear.mxfp8.b12x import (
14
+ warmup_b12x_mxfp8_linear,
15
+)
16
+from vllm.model_executor.kernels.linear.nvfp4.b12x import (
17
+ warmup_b12x_nvfp4_linear,
18
+)
19
+from vllm.model_executor.kernels.linear.scaled_mm.b12x_block import (
20
+ warmup_b12x_block_fp8_linear,
21
+)
22
+from vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor import (
23
+ warmup_b12x_tensor_fp8_linear,
24
+)
25
+
26
+if TYPE_CHECKING:
27
+ from vllm.v1.worker.gpu_worker import Worker
28
+
29
+logger = init_logger(__name__)
30
+
31
+
32
+def b12x_warmup(worker: "Worker", cudagraph_capture_sizes: list[int]) -> None:
33
+ model = worker.get_model()
34
+ max_tokens = worker.scheduler_config.max_num_batched_tokens
35
+ output_dtype = getattr(
36
+ getattr(worker, "model_config", None),
37
+ "dtype",
38
+ torch.bfloat16,
39
+ )
40
+
41
+ warmup_kwargs = {
42
+ "max_tokens": max_tokens,
43
+ "cudagraph_capture_sizes": cudagraph_capture_sizes,
44
+ "output_dtype": output_dtype,
45
+ }
46
+ providers = (
47
+ ("block-FP8", warmup_b12x_block_fp8_linear),
48
+ ("MXFP8", warmup_b12x_mxfp8_linear),
49
+ ("tensor FP8", warmup_b12x_tensor_fp8_linear),
50
+ ("MXFP4", warmup_b12x_mxfp4_linear),
51
+ ("NVFP4", warmup_b12x_nvfp4_linear),
52
+ )
53
+ for name, warmup in providers:
54
+ warmed = warmup(model, **warmup_kwargs)
55
+ if warmed:
56
+ logger.info_once(
57
+ "Warmed up %d B12X %s linear GEMM signatures.",
58
+ warmed,
59
+ name,
60
+ )
vllm/model_executor/warmup/kernel_warmup.py
+7
-3
@@ -12,6 +12,7 @@ import torch
12
13
import vllm.envs as envs
14
from vllm.logger import init_logger
15
+from vllm.model_executor.warmup.b12x_warmup import b12x_warmup
16
from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup
17
from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup
18
from vllm.model_executor.warmup.deepseek_v4_mhc_warmup import (
@@ -113,15 +114,16 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
114
115
qwen_triton_warmup(worker.model_runner, worker.vllm_config.model_config)
116
117
+ compilation_config = worker.vllm_config.compilation_config
118
+ cudagraph_capture_sizes = list(compilation_config.cudagraph_capture_sizes or [])
119
+
120
# DSv4 mHC TileLang kernels (hc_pre/hc_post/hc_head_op) run every decoder
121
# layer per token; warm them across token sizes first so the first real
122
# request doesn't pay JIT cost. No-op for non-DSv4 models (gated inside).
123
deepseek_v4_mhc_warmup(
124
worker.get_model(),
125
max_tokens=worker.scheduler_config.max_num_batched_tokens,
122
- cudagraph_capture_sizes=(
123
- worker.vllm_config.compilation_config.cudagraph_capture_sizes or []
124
- ),
126
+ cudagraph_capture_sizes=cudagraph_capture_sizes,
127
)
128
129
# Run next so input-prep kernels JIT against pristine runner state.
@@ -156,6 +158,8 @@ def kernel_warmup(worker: "Worker", *, process_local_only: bool = False):
158
max_tokens = worker.scheduler_config.max_num_batched_tokens
159
deep_gemm_warmup(model, max_tokens)
160
161
+ b12x_warmup(worker, cudagraph_capture_sizes)
162
+
163
minimax_m3_msa_warmup(worker)
164
165
enable_flashinfer_autotune = (
vllm/utils/b12x.py
new
+101
@@ -0,0 +1,101 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Lazy accessors for the optional ``b12x`` package."""
4
+
5
+import functools
6
+import importlib
7
+import importlib.util
8
+from collections.abc import Iterable
9
+from dataclasses import fields, is_dataclass
10
+from types import ModuleType
11
+from typing import Any
12
+
13
+import torch
14
+
15
+
16
+@functools.cache
17
+def has_b12x() -> bool:
18
+ """Return whether the B12X package is installed."""
19
+ return importlib.util.find_spec("b12x") is not None
20
+
21
+
22
+@functools.cache
23
+def _get_submodule(module_name: str) -> ModuleType | None:
24
+ if not has_b12x():
25
+ return None
26
+ try:
27
+ return importlib.import_module(module_name)
28
+ except (ImportError, ModuleNotFoundError):
29
+ return None
30
+
31
+
32
+def get_b12x_blockscaled() -> ModuleType | None:
33
+ return _get_submodule("b12x.gemm.blockscaled")
34
+
35
+
36
+def get_b12x_intrinsics() -> ModuleType | None:
37
+ return _get_submodule("b12x._lib.intrinsics")
38
+
39
+
40
+def get_b12x_mxfp8_linear() -> ModuleType | None:
41
+ return _get_submodule("b12x.gemm.mxfp8_linear")
42
+
43
+
44
+def get_b12x_tensor_fp8_linear() -> ModuleType | None:
45
+ return _get_submodule("b12x.gemm.tensor_fp8_linear")
46
+
47
+
48
+def b12x_warmup_token_counts(
49
+ *,
50
+ max_tokens: int,
51
+ cudagraph_capture_sizes: Iterable[int] = (),
52
+) -> tuple[int, ...]:
53
+ # B12X deduplicates shapes that select the same internal kernel policy.
54
+ # Keep the complete serving shape set here rather than duplicating its
55
+ # policy-selection heuristics in vLLM.
56
+ counts = {1}
57
+ counts.update(int(size) for size in cudagraph_capture_sizes if int(size) > 0)
58
+ if int(max_tokens) > 0:
59
+ counts.add(int(max_tokens))
60
+ return tuple(sorted(counts))
61
+
62
+
63
+def _same_packed_layout(current: Any, replacement: Any) -> bool:
64
+ if type(current) is not type(replacement):
65
+ return False
66
+ if isinstance(current, torch.Tensor):
67
+ return (
68
+ current.shape == replacement.shape
69
+ and current.stride() == replacement.stride()
70
+ and current.dtype == replacement.dtype
71
+ and current.device == replacement.device
72
+ )
73
+ if is_dataclass(current):
74
+ return all(
75
+ _same_packed_layout(
76
+ getattr(current, field.name),
77
+ getattr(replacement, field.name),
78
+ )
79
+ for field in fields(current)
80
+ )
81
+ return bool(current == replacement)
82
+
83
+
84
+def _copy_packed_tensors(current: Any, replacement: Any) -> None:
85
+ if isinstance(current, torch.Tensor):
86
+ current.copy_(replacement)
87
+ elif is_dataclass(current):
88
+ for field in fields(current):
89
+ _copy_packed_tensors(
90
+ getattr(current, field.name),
91
+ getattr(replacement, field.name),
92
+ )
93
+
94
+
95
+@torch.no_grad()
96
+def reuse_packed_weight_storage(current: Any, replacement: Any) -> Any:
97
+ """Reuse packed tensor addresses when a compatible weight is reloaded."""
98
+ if current is None or not _same_packed_layout(current, replacement):
99
+ return replacement
100
+ _copy_packed_tensors(current, replacement)
101
+ return current