[Refactor] Simplify B12X linear kernels and warmup (#52368)
Signed-off-by: mgoin <mgoin64@gmail.com>
Misha Goin committed
Aug 17, 2026 at 12:45 UTC
4ab5e5012a27e2b751c679873f231bafa0f6b098
15 files changed
+959
-1607
.buildkite/test_areas/kernels.yaml
+22
@@ -339,6 +339,28 @@ steps:
339
# e2e
340
- pytest -v -s tests/models/quantization/test_nvfp4.py
341
342
+- label: B12X Linear Kernels (DGX Spark) Nightly
343
+ key: b12x-linear-kernels-dgx-spark-nightly
344
+ timeout_in_minutes: 30
345
+ device: dgx-spark
346
+ optional: true
347
+ num_devices: 1
348
+ depends_on:
349
+ - arm64-image-build
350
+ source_file_dependencies:
351
+ - setup.py
352
+ - vllm/model_executor/kernels/linear/
353
+ - vllm/model_executor/warmup/b12x_warmup.py
354
+ - vllm/utils/b12x.py
355
+ - tests/model_executor/kernels/test_b12x_linear.py
356
+ - tests/model_executor/test_b12x_warmup.py
357
+ - tests/kernels/quantization/test_block_fp8.py
358
+ commands:
359
+ - uv pip install --system b12x==1.2.4
360
+ - pytest -v -s model_executor/kernels/test_b12x_linear.py
361
+ model_executor/test_b12x_warmup.py
362
+ kernels/quantization/test_block_fp8.py -k b12x
363
+
364
- label: Kernels Helion Test
365
key: kernels-helion-test
366
timeout_in_minutes: 115
tests/kernels/quantization/test_block_fp8.py
+1
-1
@@ -14,7 +14,7 @@ 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 (
17
+from vllm.model_executor.kernels.linear.scaled_mm.b12x import (
18
B12xFp8BlockScaledMMKernel,
19
_run_b12x_fp8_block_scaled_mm,
20
)
tests/model_executor/kernels/test_b12x_linear.py
renamed
+247
-371
@@ -3,6 +3,7 @@
3
4
from __future__ import annotations
5
6
+import importlib
7
import types
8
from dataclasses import dataclass
9
@@ -13,141 +14,120 @@ from vllm.model_executor.kernels.linear import (
14
_LINEAR_BACKEND_KERNEL_MAP,
15
_POSSIBLE_FP8_BLOCK_KERNELS,
16
_POSSIBLE_FP8_KERNELS,
17
+ _POSSIBLE_MXFP4_KERNELS,
18
_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 (
19
+ _POSSIBLE_NVFP4_KERNELS,
20
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 (
21
+ B12xMxFp4LinearKernel,
22
+ B12xMxfp8LinearKernel,
23
+ B12xNvFp4LinearKernel,
24
B12xTensorFP8ScaledMMLinearKernel,
35
- warmup_b12x_tensor_fp8_linear,
36
-)
37
-from vllm.model_executor.kernels.linear.scaled_mm.ScaledMMLinearKernel import (
25
FP8ScaledMMLinearLayerConfig,
26
+ Mxfp8LinearLayerConfig,
27
+ init_fp8_linear_kernel,
28
+ init_mxfp4_linear_kernel,
29
+ init_mxfp8_linear_kernel,
30
+ init_nvfp4_linear_kernel,
31
)
40
-from vllm.model_executor.layers.linear import (
41
- ReplicatedLinear,
42
- UnquantizedLinearMethod,
32
+from vllm.model_executor.kernels.linear.nvfp4.marlin import (
33
+ MarlinNvFp4LinearKernel,
34
)
35
from vllm.model_executor.layers.quantization.utils.quant_utils import (
36
kFp8Dynamic128Sym,
37
kFp8Static128BlockSym,
38
kFp8StaticTensorSym,
39
+ kMxfp4Dynamic,
40
)
49
-from vllm.platforms import PlatformEnum, current_platform
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]
41
+from vllm.platforms import PlatformEnum
42
43
44
@pytest.mark.parametrize(
64
- ("kernels", "before", "b12x", "after"),
45
+ ("kernel_cls", "kernels", "before", "after", "initializer", "kwargs"),
46
[
47
(
48
+ B12xMxFp4LinearKernel,
49
+ _POSSIBLE_MXFP4_KERNELS[PlatformEnum.CUDA],
50
+ "HummingMxFp4LinearKernel",
51
+ "EmulationMxfp4LinearKernel",
52
+ init_mxfp4_linear_kernel,
53
+ {"activation_quant_key": kMxfp4Dynamic},
54
+ ),
55
+ (
56
+ B12xNvFp4LinearKernel,
57
+ _POSSIBLE_NVFP4_KERNELS[PlatformEnum.CUDA],
58
+ "FbgemmNvFp4LinearKernel",
59
+ "EmulationNvFp4LinearKernel",
60
+ init_nvfp4_linear_kernel,
61
+ {},
62
+ ),
63
+ (
64
+ B12xMxfp8LinearKernel,
65
+ _POSSIBLE_MXFP8_KERNELS[PlatformEnum.CUDA],
66
+ "MarlinMxfp8LinearKernel",
67
+ "EmulationMxfp8LinearKernel",
68
+ init_mxfp8_linear_kernel,
69
+ {},
70
+ ),
71
+ (
72
+ B12xTensorFP8ScaledMMLinearKernel,
73
_POSSIBLE_FP8_KERNELS[PlatformEnum.CUDA],
74
"CutlassFP8ScaledMMLinearKernel",
69
- "B12xTensorFP8ScaledMMLinearKernel",
75
"PerTensorTorchFP8ScaledMMLinearKernel",
76
+ init_fp8_linear_kernel,
77
+ {
78
+ "activation_quant_key": kFp8StaticTensorSym,
79
+ "weight_quant_key": kFp8StaticTensorSym,
80
+ "input_dtype": torch.bfloat16,
81
+ "out_dtype": torch.bfloat16,
82
+ "weight_shape": (2048, 2048),
83
+ },
84
),
85
(
86
+ B12xFp8BlockScaledMMKernel,
87
_POSSIBLE_FP8_BLOCK_KERNELS[PlatformEnum.CUDA],
88
"CutlassFp8BlockScaledMMKernel",
75
- "B12xFp8BlockScaledMMKernel",
89
"MarlinFP8ScaledMMLinearKernel",
77
- ),
78
- (
79
- _POSSIBLE_MXFP8_KERNELS[PlatformEnum.CUDA],
80
- "MarlinMxfp8LinearKernel",
81
- "B12xMxfp8LinearKernel",
82
- "EmulationMxfp8LinearKernel",
90
+ init_fp8_linear_kernel,
91
+ {
92
+ "activation_quant_key": kFp8Dynamic128Sym,
93
+ "weight_quant_key": kFp8Static128BlockSym,
94
+ "input_dtype": torch.bfloat16,
95
+ "out_dtype": torch.bfloat16,
96
+ "weight_shape": (2048, 2048),
97
+ },
98
),
99
],
100
)
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
- device = current_platform.device_type
104
- layer = ReplicatedLinear(
105
- 128,
106
- 64,
107
- bias=False,
108
- params_dtype=torch.bfloat16,
109
- quant_config=None,
110
- prefix="bf16_linear",
111
- ).to(device)
112
- layer.weight.data.normal_()
113
- x = torch.randn(8, 128, device=device, dtype=torch.bfloat16)
114
-
115
- output, output_bias = layer(x)
116
- expected = torch.nn.functional.linear(x, layer.weight)
117
-
118
- assert isinstance(layer.quant_method, UnquantizedLinearMethod)
119
- assert output_bias is None
120
- torch.testing.assert_close(output, expected)
121
-
122
-
123
-def test_b12x_explicit_backend_selects_per_tensor_fp8(
101
+def test_b12x_backend_registration_priority_and_selection(
102
monkeypatch,
103
default_vllm_config,
104
+ kernel_cls,
105
+ kernels,
106
+ before: str,
107
+ after: str,
108
+ initializer,
109
+ kwargs: dict,
110
) -> None:
111
import vllm.model_executor.kernels.linear as linear_mod
112
113
+ assert kernel_cls in _LINEAR_BACKEND_KERNEL_MAP["b12x"]
114
+ names = [kernel.__name__ for kernel in kernels]
115
+ assert names.index(before) < names.index(kernel_cls.__name__) < names.index(after)
116
+
117
monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
118
monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
119
monkeypatch.setattr(
132
- B12xTensorFP8ScaledMMLinearKernel,
120
+ kernel_cls,
121
"is_supported",
122
classmethod(lambda cls, compute_capability=None: (True, None)),
123
)
124
monkeypatch.setattr(
137
- B12xTensorFP8ScaledMMLinearKernel,
125
+ kernel_cls,
126
"can_implement",
127
classmethod(lambda cls, config: (True, None)),
128
)
129
142
- kernel = init_fp8_linear_kernel(
143
- activation_quant_key=kFp8StaticTensorSym,
144
- weight_quant_key=kFp8StaticTensorSym,
145
- input_dtype=torch.bfloat16,
146
- out_dtype=torch.bfloat16,
147
- weight_shape=(2048, 2048),
148
- )
149
-
150
- assert isinstance(kernel, B12xTensorFP8ScaledMMLinearKernel)
130
+ assert isinstance(initializer(**kwargs), kernel_cls)
131
132
133
def test_b12x_tensor_fp8_can_implement_supported_config() -> None:
@@ -165,33 +145,8 @@ def test_b12x_tensor_fp8_can_implement_supported_config() -> None:
145
assert reason is None
146
147
168
-def test_b12x_explicit_backend_selects_block_fp8(
169
- monkeypatch,
170
- default_vllm_config,
171
-) -> None:
172
- import vllm.model_executor.kernels.linear as linear_mod
173
-
174
- monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
175
- monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
176
- monkeypatch.setattr(
177
- B12xFp8BlockScaledMMKernel,
178
- "is_supported",
179
- classmethod(lambda cls, compute_capability=None: (True, None)),
180
- )
181
-
182
- kernel = init_fp8_linear_kernel(
183
- activation_quant_key=kFp8Dynamic128Sym,
184
- weight_quant_key=kFp8Static128BlockSym,
185
- input_dtype=torch.bfloat16,
186
- out_dtype=torch.bfloat16,
187
- weight_shape=(2048, 2048),
188
- )
189
-
190
- assert isinstance(kernel, B12xFp8BlockScaledMMKernel)
191
-
192
-
148
def test_b12x_block_fp8_checks_runtime_support(monkeypatch) -> None:
194
- import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod
149
+ import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod
150
151
platform = types.SimpleNamespace(
152
is_cuda=lambda: True,
@@ -261,75 +216,10 @@ def test_b12x_block_fp8_requires_aligned_features() -> None:
216
)
217
218
264
-def test_b12x_warmup_token_counts_cover_serving_regimes() -> None:
265
- assert b12x_warmup_token_counts(
266
- max_tokens=2048,
267
- cudagraph_capture_sizes=[1, 2, 8, 128],
268
- ) == (1, 2, 8, 128, 2048)
269
-
270
-
271
-def test_warmup_b12x_block_fp8_dedupes_weight_signatures(monkeypatch) -> None:
272
- import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod
273
-
274
- calls = []
275
-
276
- def run(a, weight, a_scale, weight_scale, out_dtype):
277
- calls.append((a.shape, weight, a_scale.shape, weight_scale, out_dtype))
278
- return torch.empty((a.shape[0], weight.shape[0]), dtype=out_dtype)
279
-
280
- platform = types.SimpleNamespace(
281
- is_cuda=lambda: True,
282
- is_device_capability_family=lambda family: family == 120,
283
- )
284
- monkeypatch.setattr(b12x_mod, "current_platform", platform)
285
-
286
- monkeypatch.setattr(
287
- b12x_mod,
288
- "_import_b12x_blockscaled",
289
- lambda: types.SimpleNamespace(),
290
- )
291
- monkeypatch.setattr(b12x_mod, "_run_b12x_fp8_block_scaled_mm", run)
292
-
293
- def layer(in_features: int, out_features: int):
294
- return types.SimpleNamespace(
295
- b12x_block_fp8_linear=True,
296
- weight=torch.empty((out_features, in_features), dtype=torch.float8_e4m3fn),
297
- weight_scale_inv=torch.empty(
298
- (out_features // 128, in_features // 128), dtype=torch.float32
299
- ),
300
- )
301
-
302
- layer_a = layer(128, 256)
303
- layer_b = layer(256, 128)
304
- modules = [
305
- layer_a,
306
- layer_a,
307
- layer_b,
308
- types.SimpleNamespace(),
309
- ]
310
- model = types.SimpleNamespace(modules=lambda: iter(modules))
311
-
312
- warmed = warmup_b12x_block_fp8_linear(
313
- model,
314
- max_tokens=32,
315
- cudagraph_capture_sizes=[2, 8],
316
- output_dtype=torch.bfloat16,
317
- )
318
-
319
- assert warmed == 8
320
- assert [call[0][0] for call in calls] == [1, 2, 8, 32] * 2
321
- assert [call[2][0] for call in calls] == [1, 2, 8, 32] * 2
322
- assert calls[0][1] is layer_a.weight
323
- assert calls[4][1] is layer_b.weight
324
- assert calls[0][3] is layer_a.weight_scale_inv
325
- assert calls[4][3] is layer_b.weight_scale_inv
326
- assert all(call[4] == torch.bfloat16 for call in calls)
327
-
328
-
219
def test_b12x_tensor_fp8_process_weights_packs_modelopt_layout(
220
monkeypatch,
221
) -> None:
332
- import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod
222
+ import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod
223
224
calls = []
225
packed = types.SimpleNamespace(out_features=64)
@@ -367,6 +257,7 @@ def test_b12x_tensor_fp8_process_weights_packs_modelopt_layout(
257
kernel.process_weights_after_loading(layer)
258
259
assert layer.b12x_tensor_fp8_packed_weight is packed
260
+ assert layer.b12x_warmup_provider is kernel
261
assert len(calls) == 1
262
weight, output_scale = calls[0]
263
torch.testing.assert_close(weight, original_weight.T.contiguous())
@@ -378,67 +269,10 @@ def test_b12x_tensor_fp8_process_weights_packs_modelopt_layout(
269
torch.testing.assert_close(layer.input_scale, torch.tensor(0.5))
270
271
381
-def test_warmup_b12x_tensor_fp8_dedupes_weight_signatures(monkeypatch) -> None:
382
- import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod
383
-
384
- calls = []
385
-
386
- def prewarm(packed_weight, token_counts, *, out_dtype, stream):
387
- del stream
388
- calls.append((packed_weight, tuple(token_counts), out_dtype))
389
- return len(tuple(token_counts))
390
-
391
- platform = types.SimpleNamespace(
392
- is_cuda=lambda: True,
393
- is_device_capability_family=lambda family: family == 120,
394
- )
395
- monkeypatch.setattr(b12x_mod, "current_platform", platform)
396
- monkeypatch.setattr(
397
- b12x_mod,
398
- "_import_b12x_tensor_fp8",
399
- lambda: types.SimpleNamespace(prewarm=prewarm),
400
- )
401
- monkeypatch.setattr(
402
- b12x_mod,
403
- "current_stream",
404
- lambda: types.SimpleNamespace(cuda_stream=object()),
405
- )
406
-
407
- def packed(in_features: int, padded_in_features: int, out_features: int):
408
- return types.SimpleNamespace(
409
- in_features=in_features,
410
- padded_in_features=padded_in_features,
411
- out_features=out_features,
412
- values=torch.empty(1),
413
- )
414
-
415
- packed_a = packed(128, 128, 256)
416
- packed_b = packed(160, 256, 512)
417
- modules = [
418
- types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_a),
419
- types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_a),
420
- types.SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_b),
421
- types.SimpleNamespace(),
422
- ]
423
- model = types.SimpleNamespace(modules=lambda: iter(modules))
424
-
425
- warmed = warmup_b12x_tensor_fp8_linear(
426
- model,
427
- max_tokens=2048,
428
- cudagraph_capture_sizes=[1, 2],
429
- )
430
-
431
- assert warmed == 6
432
- assert calls == [
433
- (packed_a, (1, 2, 2048), torch.bfloat16),
434
- (packed_b, (1, 2, 2048), torch.bfloat16),
435
- ]
436
-
437
-
272
def test_b12x_tensor_fp8_apply_quantizes_and_uses_packed_weight(
273
monkeypatch,
274
) -> None:
441
- import vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor as b12x_mod
275
+ import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod
276
277
calls = []
278
@@ -506,27 +340,6 @@ def test_b12x_tensor_fp8_apply_quantizes_and_uses_packed_weight(
340
assert expected_m == 6
341
342
509
-def test_b12x_mxfp8_explicit_backend_selects_kernel(monkeypatch) -> None:
510
- import vllm.model_executor.kernels.linear as linear_mod
511
-
512
- monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
513
- monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
514
- monkeypatch.setattr(
515
- B12xMxfp8LinearKernel,
516
- "is_supported",
517
- classmethod(lambda cls, compute_capability=None: (True, None)),
518
- )
519
- monkeypatch.setattr(
520
- B12xMxfp8LinearKernel,
521
- "can_implement",
522
- classmethod(lambda cls, c: (True, None)),
523
- )
524
-
525
- kernel = init_mxfp8_linear_kernel()
526
-
527
- assert isinstance(kernel, B12xMxfp8LinearKernel)
528
-
529
-
343
def test_b12x_mxfp8_can_implement_supported_config() -> None:
344
can_implement, reason = B12xMxfp8LinearKernel.can_implement(
345
Mxfp8LinearLayerConfig()
@@ -536,90 +349,6 @@ def test_b12x_mxfp8_can_implement_supported_config() -> None:
349
assert reason is None
350
351
539
-def test_b12x_mxfp8_expected_m_uses_live_m() -> None:
540
- assert _b12x_mxfp8_expected_m(0) == 1
541
- assert _b12x_mxfp8_expected_m(1) == 1
542
- assert _b12x_mxfp8_expected_m(2) == 2
543
- assert _b12x_mxfp8_expected_m(8) == 8
544
- assert _b12x_mxfp8_expected_m(9) == 9
545
- assert _b12x_mxfp8_expected_m(128) == 128
546
- assert _b12x_mxfp8_expected_m(129) == 129
547
- assert _b12x_mxfp8_expected_m(2048) == 2048
548
-
549
-
550
-def test_warmup_b12x_mxfp8_linear_dedupes_weight_signatures(
551
- monkeypatch,
552
-) -> None:
553
- import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod
554
-
555
- calls = []
556
-
557
- def mm(
558
- source: torch.Tensor,
559
- packed_weight,
560
- *,
561
- bias: torch.Tensor | None = None,
562
- expected_m: int | None = None,
563
- stream: object = None,
564
- ) -> torch.Tensor:
565
- del stream
566
- calls.append((source.shape, packed_weight, bias, expected_m))
567
- return source.new_empty((source.shape[0], packed_weight.out_features))
568
-
569
- platform = types.SimpleNamespace(
570
- is_cuda=lambda: True,
571
- is_device_capability_family=lambda family: family == 120,
572
- )
573
- monkeypatch.setattr(b12x_mod, "current_platform", platform)
574
- monkeypatch.setattr(
575
- b12x_mod,
576
- "_import_b12x_mxfp8",
577
- lambda: types.SimpleNamespace(mm=mm),
578
- )
579
- monkeypatch.setattr(
580
- b12x_mod,
581
- "current_stream",
582
- lambda: types.SimpleNamespace(cuda_stream=object()),
583
- )
584
-
585
- def packed(in_features: int, padded_in_features: int, out_features: int):
586
- return types.SimpleNamespace(
587
- in_features=in_features,
588
- padded_in_features=padded_in_features,
589
- out_features=out_features,
590
- weight=types.SimpleNamespace(values=torch.empty(1)),
591
- )
592
-
593
- packed_a = packed(128, 128, 256)
594
- packed_b = packed(128, 128, 512)
595
- modules = [
596
- types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_a),
597
- types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_a),
598
- types.SimpleNamespace(b12x_mxfp8_packed_weight=packed_b),
599
- types.SimpleNamespace(),
600
- ]
601
- model = types.SimpleNamespace(modules=lambda: iter(modules))
602
-
603
- warmed = warmup_b12x_mxfp8_linear(
604
- model,
605
- max_tokens=2048,
606
- cudagraph_capture_sizes=[1, 2],
607
- )
608
-
609
- assert warmed == 6
610
- assert [call[0] for call in calls] == [
611
- torch.Size([1, 128]),
612
- torch.Size([2, 128]),
613
- torch.Size([2048, 128]),
614
- torch.Size([1, 128]),
615
- torch.Size([2, 128]),
616
- torch.Size([2048, 128]),
617
- ]
618
- assert [call[3] for call in calls] == [1, 2, 2048, 1, 2, 2048]
619
- assert calls[0][1] is packed_a
620
- assert calls[3][1] is packed_b
621
-
622
-
352
def test_b12x_mxfp8_support_check_reports_missing_import(monkeypatch) -> None:
353
import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod
354
@@ -693,6 +422,7 @@ def test_b12x_mxfp8_process_weights_packs_modelopt_layout(monkeypatch) -> None:
422
kernel.process_weights_after_loading(layer)
423
424
assert layer.b12x_mxfp8_packed_weight is packed
425
+ assert layer.b12x_warmup_provider is kernel
426
assert len(calls) == 1
427
weight, weight_scale = calls[0]
428
assert weight.shape == (48, 128)
@@ -805,7 +535,7 @@ def test_b12x_block_fp8_process_weights_keeps_native_block_layout() -> None:
535
536
kernel.process_weights_after_loading(layer)
537
808
- assert layer.b12x_block_fp8_linear
538
+ assert layer.b12x_warmup_provider is kernel
539
assert layer.weight.shape == (128, 128)
540
assert layer.weight.dtype == torch.float8_e4m3fn
541
assert layer.weight_scale_inv.shape == (1, 1)
@@ -882,7 +612,7 @@ def test_b12x_mxfp8_apply_uses_packed_weight(monkeypatch) -> None:
612
613
614
def test_b12x_block_fp8_apply_uses_b12x_recipe_api(monkeypatch) -> None:
885
- import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod
615
+ import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod
616
617
calls = []
618
@@ -919,37 +649,183 @@ def test_b12x_block_fp8_apply_uses_b12x_recipe_api(monkeypatch) -> None:
649
torch.testing.assert_close(output, torch.full_like(output, 13.0))
650
651
922
-def test_b12x_block_fp8_helper_uses_regular_compact_scale_api(monkeypatch) -> None:
923
- import vllm.model_executor.kernels.linear.scaled_mm.b12x_block as b12x_mod
652
+def test_b12x_mxfp4_requires_dynamic_activations() -> None:
653
+ config = types.SimpleNamespace(activation_quant_key=kMxfp4Dynamic)
654
+ can_implement, reason = B12xMxFp4LinearKernel.can_implement(config)
655
925
- calls = []
656
+ assert can_implement
657
+ assert reason is None
658
927
- def mm_block_fp8(*args, **kwargs):
659
+ config.activation_quant_key = None
660
+ can_implement, reason = B12xMxFp4LinearKernel.can_implement(config)
661
+
662
+ assert not can_implement
663
+ assert reason == "B12X MXFP4 GEMM requires dynamic MXFP4 activations"
664
+
665
+
666
+@pytest.mark.parametrize(
667
+ ("kernel_cls", "module_name", "scale_dtype"),
668
+ [
669
+ (
670
+ B12xMxFp4LinearKernel,
671
+ "vllm.model_executor.kernels.linear.mxfp4.b12x",
672
+ torch.uint8,
673
+ ),
674
+ (
675
+ B12xNvFp4LinearKernel,
676
+ "vllm.model_executor.kernels.linear.nvfp4.b12x",
677
+ torch.float8_e4m3fn,
678
+ ),
679
+ ],
680
+)
681
+def test_b12x_fp4_processes_scale_and_preserves_loader(
682
+ monkeypatch,
683
+ kernel_cls,
684
+ module_name: str,
685
+ scale_dtype: torch.dtype,
686
+) -> None:
687
+ scale = torch.empty((48, 8), dtype=scale_dtype)
688
+ swizzled_scale = torch.empty((128, 8), dtype=scale_dtype)
689
+ intrinsics = types.SimpleNamespace(swizzle_block_scale=lambda value: swizzled_scale)
690
+ monkeypatch.setattr(
691
+ importlib.import_module(module_name),
692
+ "_import_b12x_intrinsics",
693
+ lambda: intrinsics,
694
+ )
695
+ layer = torch.nn.Module()
696
+ layer.prefix = "model.layers.0.mlp.shared_expert.down_proj"
697
+ layer.weight_scale = torch.nn.Parameter(scale, requires_grad=False)
698
+ weight_loader = object()
699
+ layer.weight_scale.weight_loader = weight_loader
700
+ kernel = object.__new__(kernel_cls)
701
+
702
+ kernel.process_weights_after_loading(layer)
703
+
704
+ assert layer.weight_scale.data_ptr() == swizzled_scale.data_ptr()
705
+ assert layer.weight_scale.weight_loader is weight_loader
706
+ assert layer.b12x_warmup_provider is kernel
707
+
708
+
709
+def test_b12x_mxfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None:
710
+ import vllm.model_executor.kernels.linear.mxfp4.b12x as b12x_mod
711
+ import vllm.utils.flashinfer as flashinfer_utils
712
+
713
+ calls: list[tuple] = []
714
+ x_packed = torch.empty((6, 64), dtype=torch.uint8)
715
+ x_scale_storage = torch.empty((128, 4), dtype=torch.uint8)
716
+
717
+ def mm_mxfp4(*args, **kwargs):
718
calls.append((args, kwargs))
929
- return torch.full((6, 256), 17.0, dtype=torch.bfloat16)
719
+ return torch.full((6, 48), 3.0, dtype=torch.bfloat16)
720
721
+ monkeypatch.setattr(
722
+ flashinfer_utils,
723
+ "flashinfer_mxfp4_quantize",
724
+ lambda *args, **kwargs: (x_packed, x_scale_storage),
725
+ )
726
monkeypatch.setattr(
727
b12x_mod,
728
"_import_b12x_blockscaled",
934
- lambda: types.SimpleNamespace(mm_block_fp8=mm_block_fp8),
729
+ lambda: types.SimpleNamespace(mm_mxfp4=mm_mxfp4),
730
)
731
937
- a = torch.empty((6, 128), dtype=torch.float8_e4m3fn)
938
- weight = torch.empty((256, 128), dtype=torch.float8_e4m3fn)
939
- a_scale = torch.empty((6, 1), dtype=torch.float32)
940
- weight_scale = torch.empty((2, 1), dtype=torch.float32)
941
- output = _run_b12x_fp8_block_scaled_mm(
942
- a,
943
- weight,
944
- a_scale,
945
- weight_scale,
946
- torch.bfloat16,
732
+ layer = torch.nn.Module()
733
+ layer.output_size_per_partition = 48
734
+ layer.weight = torch.empty((48, 64), dtype=torch.uint8)
735
+ layer.weight_scale = torch.empty((128, 4), dtype=torch.uint8)
736
+ x = torch.empty((2, 3, 128), dtype=torch.bfloat16)
737
+ bias = torch.ones(48, dtype=torch.bfloat16)
738
+ kernel = object.__new__(B12xMxFp4LinearKernel)
739
+
740
+ output = kernel.apply_weights(layer, x, bias)
741
+
742
+ assert output.shape == (2, 3, 48)
743
+ torch.testing.assert_close(output, torch.full_like(output, 4.0))
744
+ assert len(calls) == 1
745
+ args, kwargs = calls[0]
746
+ assert args == (
747
+ x_packed,
748
+ x_scale_storage,
749
+ layer.weight,
750
+ layer.weight_scale,
751
)
752
+ assert kwargs == {"out_dtype": torch.bfloat16}
753
949
- assert output.shape == (6, 256)
754
+
755
+def test_b12x_nvfp4_can_implement_supported_config() -> None:
756
+ can_implement, reason = B12xNvFp4LinearKernel.can_implement(None)
757
+
758
+ assert can_implement
759
+ assert reason is None
760
+
761
+
762
+def test_b12x_backend_preserves_w4a16_fallback(monkeypatch) -> None:
763
+ import vllm.model_executor.kernels.linear as linear_mod
764
+
765
+ monkeypatch.setattr(linear_mod.current_platform, "_enum", PlatformEnum.CUDA)
766
+ monkeypatch.setattr(linear_mod, "_get_linear_backend", lambda: "b12x")
767
+ monkeypatch.setattr(
768
+ MarlinNvFp4LinearKernel,
769
+ "is_supported",
770
+ classmethod(lambda cls, compute_capability=None: (True, None)),
771
+ )
772
+
773
+ kernel = init_nvfp4_linear_kernel(use_a16=True)
774
+
775
+ assert isinstance(kernel, MarlinNvFp4LinearKernel)
776
+
777
+
778
+def test_b12x_nvfp4_apply_calls_native_blockscaled_gemm(monkeypatch) -> None:
779
+ import vllm.model_executor.kernels.linear.nvfp4.b12x as b12x_mod
780
+
781
+ calls: list[tuple] = []
782
+ quant_calls: list[tuple] = []
783
+ x_packed = torch.empty((6, 64), dtype=torch.uint8)
784
+ x_scale_storage = torch.empty((128, 8), dtype=torch.float8_e4m3fn)
785
+
786
+ def quant(*args, **kwargs):
787
+ quant_calls.append((args, kwargs))
788
+ return x_packed, x_scale_storage
789
+
790
+ def mm_nvfp4(*args, **kwargs):
791
+ calls.append((args, kwargs))
792
+ return torch.full((6, 48), 3.0, dtype=torch.bfloat16)
793
+
794
+ monkeypatch.setattr(b12x_mod, "scaled_fp4_quant", quant)
795
+ monkeypatch.setattr(
796
+ b12x_mod,
797
+ "_import_b12x_blockscaled",
798
+ lambda: types.SimpleNamespace(mm_nvfp4=mm_nvfp4),
799
+ )
800
+
801
+ layer = torch.nn.Module()
802
+ layer.output_size_per_partition = 48
803
+ layer.weight = torch.empty((48, 64), dtype=torch.uint8)
804
+ layer.weight_scale = torch.empty((128, 8), dtype=torch.float8_e4m3fn)
805
+ layer.input_global_scale_inv = torch.tensor(2.0)
806
+ layer.alpha = torch.tensor(0.25)
807
+ x = torch.empty((2, 3, 256), dtype=torch.bfloat16)[..., ::2]
808
+ bias = torch.ones(48, dtype=torch.bfloat16)
809
+ kernel = object.__new__(B12xNvFp4LinearKernel)
810
+
811
+ output = kernel.apply_weights(layer, x, bias)
812
+
813
+ assert output.shape == (2, 3, 48)
814
+ torch.testing.assert_close(output, torch.full_like(output, 4.0))
815
+ assert len(quant_calls) == 1
816
+ quant_args, quant_kwargs = quant_calls[0]
817
+ assert quant_args[0].shape == (6, 128)
818
+ assert quant_args[0].data_ptr() == x.data_ptr()
819
+ assert quant_args[1] is layer.input_global_scale_inv
820
+ assert not quant_args[0].is_contiguous()
821
+ assert quant_kwargs == {"is_sf_swizzled_layout": True}
822
+ assert len(calls) == 1
823
args, kwargs = calls[0]
951
- assert args == (a, a_scale, weight, weight_scale)
952
- assert kwargs == {
953
- "out_dtype": torch.bfloat16,
954
- }
955
- torch.testing.assert_close(output, torch.full_like(output, 17.0))
824
+ assert args == (
825
+ x_packed,
826
+ x_scale_storage,
827
+ layer.weight,
828
+ layer.weight_scale,
829
+ layer.alpha,
830
+ )
831
+ assert kwargs == {"out_dtype": torch.bfloat16}
tests/model_executor/kernels/test_b12x_mxfp4_linear.py
deleted
-199
@@ -1,199 +0,0 @@
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_nvfp4_linear.py
deleted
-246
@@ -1,246 +0,0 @@
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
+166
-60
@@ -7,94 +7,200 @@ from types import SimpleNamespace
7
import pytest
8
import torch
9
10
+from vllm.model_executor.kernels.linear import (
11
+ B12xFp8BlockScaledMMKernel,
12
+ B12xMxFp4LinearKernel,
13
+ B12xMxfp8LinearKernel,
14
+ B12xNvFp4LinearKernel,
15
+ B12xTensorFP8ScaledMMLinearKernel,
16
+)
17
from vllm.model_executor.warmup.b12x_warmup import b12x_warmup
18
+from vllm.utils.b12x import B12xWarmupUnit, b12x_warmup_token_counts
19
+
20
+
21
+def test_b12x_warmup_token_counts_cover_serving_regimes() -> None:
22
+ assert b12x_warmup_token_counts(
23
+ max_tokens=2048,
24
+ cudagraph_capture_sizes=[1, 2, 8, 128],
25
+ ) == (1, 2, 8, 128, 2048)
26
27
28
@pytest.mark.parametrize(
14
- ("module_name", "warmup_name", "import_name"),
29
+ ("kernel_cls", "module_name", "call_name", "layer", "name"),
30
[
31
(
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
+ B12xMxFp4LinearKernel,
33
"vllm.model_executor.kernels.linear.mxfp4.b12x",
33
- "warmup_b12x_mxfp4_linear",
34
- "_import_b12x_blockscaled",
34
+ "_apply_b12x_mxfp4_linear",
35
+ SimpleNamespace(
36
+ weight=torch.empty((48, 64), dtype=torch.uint8),
37
+ weight_scale=torch.empty((128, 4), dtype=torch.uint8),
38
+ ),
39
+ "MXFP4",
40
),
41
(
42
+ B12xNvFp4LinearKernel,
43
"vllm.model_executor.kernels.linear.nvfp4.b12x",
38
- "warmup_b12x_nvfp4_linear",
39
- "_import_b12x_blockscaled",
44
+ "_apply_b12x_nvfp4_linear",
45
+ SimpleNamespace(
46
+ weight=torch.empty((48, 64), dtype=torch.uint8),
47
+ weight_scale=torch.empty((128, 8), dtype=torch.float8_e4m3fn),
48
+ input_global_scale_inv=torch.tensor(2.0),
49
+ alpha=torch.tensor(0.25),
50
+ ),
51
+ "NVFP4",
52
+ ),
53
+ (
54
+ B12xFp8BlockScaledMMKernel,
55
+ "vllm.model_executor.kernels.linear.scaled_mm.b12x",
56
+ "_run_b12x_fp8_block_scaled_mm",
57
+ SimpleNamespace(
58
+ weight=torch.empty((256, 128), dtype=torch.float8_e4m3fn),
59
+ weight_scale_inv=torch.empty((2, 1), dtype=torch.float32),
60
+ ),
61
+ "block-FP8",
62
),
63
],
64
)
43
-def test_b12x_linear_warmup_skips_unused_provider(
65
+def test_b12x_warmup_units_cover_token_counts(
66
monkeypatch,
67
+ kernel_cls,
68
module_name: str,
46
- warmup_name: str,
47
- import_name: str,
69
+ call_name: str,
70
+ layer,
71
+ name: str,
72
) -> 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,
73
+ calls = []
74
+ monkeypatch.setattr(
75
+ importlib.import_module(module_name),
76
+ call_name,
77
+ lambda *args: calls.append(args),
78
)
54
- monkeypatch.setattr(module, "current_platform", platform)
79
+ kernel = object.__new__(kernel_cls)
80
56
- def fail_import():
57
- pytest.fail("unused B12X provider was imported")
81
+ unit = kernel.get_b12x_warmup_unit(layer, (1, 8), torch.bfloat16)
82
+ unit.compile()
83
59
- monkeypatch.setattr(module, import_name, fail_import)
84
+ assert unit.name == name
85
+ assert [args[0].shape[0] for args in calls] == [1, 8]
86
+ assert unit.key[-1] == torch.bfloat16
87
61
- warmup = getattr(module, warmup_name)
62
- assert warmup(torch.nn.Module(), max_tokens=128) == 0
88
89
+def test_b12x_mxfp8_warmup_unit(monkeypatch) -> None:
90
+ import vllm.model_executor.kernels.linear.mxfp8.b12x as b12x_mod
91
65
-def test_b12x_warmup_covers_linear_serving_shapes(monkeypatch) -> None:
66
- import vllm.model_executor.warmup.b12x_warmup as warmup_mod
92
+ calls = []
93
+ monkeypatch.setattr(
94
+ b12x_mod,
95
+ "_import_b12x_mxfp8",
96
+ lambda: SimpleNamespace(
97
+ mm=lambda *args, **kwargs: calls.append((args, kwargs))
98
+ ),
99
+ )
100
+ monkeypatch.setattr(
101
+ b12x_mod,
102
+ "current_stream",
103
+ lambda: SimpleNamespace(cuda_stream=object()),
104
+ )
105
+ packed_weight = SimpleNamespace(
106
+ in_features=128,
107
+ padded_in_features=128,
108
+ out_features=256,
109
+ weight=SimpleNamespace(values=torch.empty(1)),
110
+ )
111
+ layer = SimpleNamespace(b12x_mxfp8_packed_weight=packed_weight)
112
+ kernel = object.__new__(B12xMxfp8LinearKernel)
113
68
- model = torch.nn.Module()
69
- linear_calls: list[tuple[torch.nn.Module, dict]] = []
114
+ unit = kernel.get_b12x_warmup_unit(layer, (1, 8), torch.float16)
115
+ unit.compile()
116
71
- def linear_warmup(model, **kwargs):
72
- linear_calls.append((model, kwargs))
73
- return 0
117
+ assert [args[0].shape for args, _ in calls] == [(1, 128), (8, 128)]
118
+ assert [kwargs["expected_m"] for _, kwargs in calls] == [1, 8]
119
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)
120
81
- worker = SimpleNamespace(
82
- get_model=lambda: model,
83
- scheduler_config=SimpleNamespace(
84
- max_num_batched_tokens=256,
85
- max_num_scheduled_tokens=320,
121
+def test_b12x_tensor_fp8_warmup_unit(monkeypatch) -> None:
122
+ import vllm.model_executor.kernels.linear.scaled_mm.b12x as b12x_mod
123
+
124
+ calls = []
125
+ monkeypatch.setattr(
126
+ b12x_mod,
127
+ "_import_b12x_tensor_fp8",
128
+ lambda: SimpleNamespace(
129
+ prewarm=lambda *args, **kwargs: calls.append((args, kwargs))
130
),
87
- model_config=SimpleNamespace(dtype=torch.float16),
88
- vllm_config=SimpleNamespace(
89
- compilation_config=SimpleNamespace(compile_sizes=[32, "dynamic", 64])
131
+ )
132
+ monkeypatch.setattr(
133
+ b12x_mod,
134
+ "current_stream",
135
+ lambda: SimpleNamespace(cuda_stream=object()),
136
+ )
137
+ packed_weight = SimpleNamespace(
138
+ in_features=128,
139
+ padded_in_features=128,
140
+ out_features=256,
141
+ values=torch.empty(1),
142
+ )
143
+ layer = SimpleNamespace(b12x_tensor_fp8_packed_weight=packed_weight)
144
+ kernel = object.__new__(B12xTensorFP8ScaledMMLinearKernel)
145
+
146
+ unit = kernel.get_b12x_warmup_unit(layer, (1, 8), torch.bfloat16)
147
+ unit.compile()
148
+
149
+ assert calls[0][0] == (packed_weight, (1, 8))
150
+ assert calls[0][1]["out_dtype"] == torch.bfloat16
151
+
152
+
153
+def test_b12x_warmup_deduplicates_registered_signatures(monkeypatch) -> None:
154
+ import vllm.model_executor.warmup.b12x_warmup as warmup_mod
155
+
156
+ calls: list[tuple[str, tuple[int, ...], torch.dtype]] = []
157
+
158
+ class Provider:
159
+ def get_b12x_warmup_unit(self, layer, token_counts, output_dtype):
160
+ return B12xWarmupUnit(
161
+ name="fake",
162
+ key=(type(self), layer.shape, output_dtype),
163
+ compile=lambda: calls.append((layer.name, token_counts, output_dtype)),
164
+ )
165
+
166
+ provider = Provider()
167
+ layers = [
168
+ SimpleNamespace(name="first", shape=(128, 256), b12x_warmup_provider=provider),
169
+ SimpleNamespace(
170
+ name="duplicate", shape=(128, 256), b12x_warmup_provider=provider
171
),
172
+ SimpleNamespace(name="second", shape=(256, 256), b12x_warmup_provider=provider),
173
+ SimpleNamespace(),
174
+ ]
175
+ scans = 0
176
+
177
+ def modules():
178
+ nonlocal scans
179
+ scans += 1
180
+ return iter(layers)
181
+
182
+ worker = SimpleNamespace(
183
+ get_model=lambda: SimpleNamespace(modules=modules),
184
+ scheduler_config=SimpleNamespace(max_num_batched_tokens=8),
185
+ model_config=SimpleNamespace(dtype=torch.float32),
186
+ )
187
+ platform = SimpleNamespace(
188
+ is_cuda=lambda: True,
189
+ is_device_capability_family=lambda family: family == 120,
190
+ )
191
+ synchronized = []
192
+ monkeypatch.setattr(warmup_mod, "current_platform", platform)
193
+ monkeypatch.setattr(
194
+ warmup_mod.torch.accelerator,
195
+ "synchronize",
196
+ lambda: synchronized.append(True),
197
)
198
93
- b12x_warmup(worker, [8, 16])
199
+ b12x_warmup(worker, [1, 2])
200
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
201
+ assert scans == 1
202
+ assert calls == [
203
+ ("first", (1, 2, 8), torch.bfloat16),
204
+ ("second", (1, 2, 8), torch.bfloat16),
205
+ ]
206
+ assert synchronized == [True]
vllm/model_executor/kernels/linear/__init__.py
+1
-3
@@ -172,10 +172,8 @@ 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 (
175
+from vllm.model_executor.kernels.linear.scaled_mm.b12x import (
176
B12xFp8BlockScaledMMKernel,
177
-)
178
-from vllm.model_executor.kernels.linear.scaled_mm.b12x_tensor import (
177
B12xTensorFP8ScaledMMLinearKernel,
178
)
179
from vllm.model_executor.kernels.linear.scaled_mm.cpu import (
vllm/model_executor/kernels/linear/mxfp4/b12x.py
+35
-75
@@ -3,9 +3,6 @@
3
4
from __future__ import annotations
5
6
-from collections.abc import Iterable
7
-from typing import Any
8
-
6
import torch
7
8
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -13,9 +10,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
10
)
11
from vllm.model_executor.utils import replace_parameter
12
from vllm.platforms import current_platform
16
-from vllm.utils.b12x import (
17
- b12x_warmup_token_counts,
18
-)
13
+from vllm.utils.b12x import B12xWarmupUnit
14
from vllm.utils.b12x import (
15
get_b12x_blockscaled as _import_b12x_blockscaled,
16
)
@@ -51,73 +46,6 @@ def _apply_b12x_mxfp4_linear(
46
return output.view(*output_shape)
47
48
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
-
49
class B12xMxFp4LinearKernel(MxFp4LinearKernel):
50
"""MXFP4 linear through the native B12X SM120 dense GEMM."""
51
@@ -151,7 +79,39 @@ class B12xMxFp4LinearKernel(MxFp4LinearKernel):
79
"weight_scale",
80
intrinsics.swizzle_block_scale(layer.weight_scale.data),
81
)
154
- layer.b12x_mxfp4_linear = True
82
+ layer.b12x_warmup_provider = self
83
+
84
+ def get_b12x_warmup_unit(
85
+ self,
86
+ layer: torch.nn.Module,
87
+ token_counts: tuple[int, ...],
88
+ output_dtype: torch.dtype,
89
+ ) -> B12xWarmupUnit:
90
+ weight = layer.weight
91
+ weight_scale = layer.weight_scale
92
+ n, packed_k = map(int, weight.shape)
93
+ k = packed_k * 2
94
+
95
+ def compile() -> None:
96
+ for tokens in token_counts:
97
+ source = torch.zeros(
98
+ (tokens, k), dtype=output_dtype, device=weight.device
99
+ )
100
+ _apply_b12x_mxfp4_linear(source, weight, weight_scale, None)
101
+
102
+ return B12xWarmupUnit(
103
+ name="MXFP4",
104
+ key=(
105
+ type(self),
106
+ weight.device,
107
+ n,
108
+ k,
109
+ weight.dtype,
110
+ weight_scale.dtype,
111
+ output_dtype,
112
+ ),
113
+ compile=compile,
114
+ )
115
116
def apply_weights(
117
self,
@@ -167,4 +127,4 @@ class B12xMxFp4LinearKernel(MxFp4LinearKernel):
127
)
128
129
170
-__all__ = ["B12xMxFp4LinearKernel", "warmup_b12x_mxfp4_linear"]
130
+__all__ = ["B12xMxFp4LinearKernel"]
vllm/model_executor/kernels/linear/mxfp8/b12x.py
+42
-84
@@ -3,9 +3,6 @@
3
4
from __future__ import annotations
5
6
-from collections.abc import Iterable
7
-from typing import Any
8
-
6
import torch
7
8
from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
@@ -15,10 +12,7 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
12
)
13
from vllm.model_executor.utils import replace_parameter
14
from vllm.platforms import current_platform
18
-from vllm.utils.b12x import (
19
- b12x_warmup_token_counts,
20
- reuse_packed_weight_storage,
21
-)
15
+from vllm.utils.b12x import B12xWarmupUnit, reuse_packed_weight_storage
16
from vllm.utils.b12x import (
17
get_b12x_mxfp8_linear as _import_b12x_mxfp8,
18
)
@@ -27,21 +21,12 @@ from vllm.utils.torch_utils import current_stream
21
from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig
22
23
30
-def _b12x_mxfp8_expected_m(tokens: int) -> int:
31
- return max(1, int(tokens))
32
-
33
-
24
def _apply_b12x_mxfp8_packed_linear(
25
layer: torch.nn.Module,
26
x: torch.Tensor,
27
bias: torch.Tensor | None,
28
) -> 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
- )
29
+ packed_weight = layer.b12x_mxfp8_packed_weight
30
31
input_2d = x.reshape(-1, x.shape[-1]).contiguous()
32
output_shape = [*x.shape[:-1], int(packed_weight.out_features)]
@@ -52,77 +37,11 @@ def _apply_b12x_mxfp8_packed_linear(
37
input_2d,
38
packed_weight,
39
bias=bias,
55
- expected_m=_b12x_mxfp8_expected_m(int(input_2d.shape[0])),
40
+ expected_m=max(1, int(input_2d.shape[0])),
41
)
42
return output.view(*output_shape)
43
44
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
-
45
class B12xMxfp8LinearKernel(Mxfp8LinearKernel):
46
"""ModelOpt MXFP8 linear through the native b12x SM120 dense GEMM path."""
47
@@ -183,6 +102,45 @@ class B12xMxfp8LinearKernel(Mxfp8LinearKernel):
102
)
103
replace_parameter(layer, "weight", weight.new_empty((0,)))
104
replace_parameter(layer, "weight_scale", weight_scale.new_empty((0,)))
105
+ layer.b12x_warmup_provider = self
106
+
107
+ def get_b12x_warmup_unit(
108
+ self,
109
+ layer: torch.nn.Module,
110
+ token_counts: tuple[int, ...],
111
+ output_dtype: torch.dtype,
112
+ ) -> B12xWarmupUnit:
113
+ packed_weight = layer.b12x_mxfp8_packed_weight
114
+ device = torch.device(packed_weight.weight.values.device)
115
+
116
+ def compile() -> None:
117
+ mxfp8 = _import_b12x_mxfp8()
118
+ assert mxfp8 is not None
119
+ for tokens in token_counts:
120
+ source = torch.zeros(
121
+ (tokens, int(packed_weight.in_features)),
122
+ dtype=output_dtype,
123
+ device=device,
124
+ )
125
+ mxfp8.mm(
126
+ source,
127
+ packed_weight,
128
+ expected_m=max(1, int(tokens)),
129
+ stream=current_stream().cuda_stream,
130
+ )
131
+
132
+ return B12xWarmupUnit(
133
+ name="MXFP8",
134
+ key=(
135
+ type(self),
136
+ device,
137
+ int(packed_weight.in_features),
138
+ int(packed_weight.padded_in_features),
139
+ int(packed_weight.out_features),
140
+ output_dtype,
141
+ ),
142
+ compile=compile,
143
+ )
144
145
def apply_weights(
146
self,
vllm/model_executor/kernels/linear/nvfp4/b12x.py
+42
-77
@@ -3,17 +3,12 @@
3
4
from __future__ import annotations
5
6
-from collections.abc import Iterable
7
-from typing import Any
8
-
6
import torch
7
8
from vllm._custom_ops import scaled_fp4_quant
9
from vllm.model_executor.utils import replace_parameter
10
from vllm.platforms import current_platform
14
-from vllm.utils.b12x import (
15
- b12x_warmup_token_counts,
16
-)
11
+from vllm.utils.b12x import B12xWarmupUnit
12
from vllm.utils.b12x import (
13
get_b12x_blockscaled as _import_b12x_blockscaled,
14
)
@@ -54,75 +49,6 @@ def _apply_b12x_nvfp4_linear(
49
return output.view(*output_shape)
50
51
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
-
52
class B12xNvFp4LinearKernel(NvFp4LinearKernel):
53
"""ModelOpt NVFP4 linear through the native B12X SM120 dense GEMM."""
54
@@ -155,7 +81,46 @@ class B12xNvFp4LinearKernel(NvFp4LinearKernel):
81
"weight_scale",
82
intrinsics.swizzle_block_scale(layer.weight_scale.data),
83
)
158
- layer.b12x_nvfp4_linear = True
84
+ layer.b12x_warmup_provider = self
85
+
86
+ def get_b12x_warmup_unit(
87
+ self,
88
+ layer: torch.nn.Module,
89
+ token_counts: tuple[int, ...],
90
+ output_dtype: torch.dtype,
91
+ ) -> B12xWarmupUnit:
92
+ weight = layer.weight
93
+ weight_scale = layer.weight_scale
94
+ n, packed_k = map(int, weight.shape)
95
+ k = packed_k * 2
96
+
97
+ def compile() -> None:
98
+ for tokens in token_counts:
99
+ source = torch.zeros(
100
+ (tokens, k), dtype=output_dtype, device=weight.device
101
+ )
102
+ _apply_b12x_nvfp4_linear(
103
+ source,
104
+ weight,
105
+ weight_scale,
106
+ layer.input_global_scale_inv,
107
+ layer.alpha,
108
+ None,
109
+ )
110
+
111
+ return B12xWarmupUnit(
112
+ name="NVFP4",
113
+ key=(
114
+ type(self),
115
+ weight.device,
116
+ n,
117
+ k,
118
+ weight.dtype,
119
+ weight_scale.dtype,
120
+ output_dtype,
121
+ ),
122
+ compile=compile,
123
+ )
124
125
def apply_weights(
126
self,
@@ -173,4 +138,4 @@ class B12xNvFp4LinearKernel(NvFp4LinearKernel):
138
)
139
140
176
-__all__ = ["B12xNvFp4LinearKernel", "warmup_b12x_nvfp4_linear"]
141
+__all__ = ["B12xNvFp4LinearKernel"]
vllm/model_executor/kernels/linear/scaled_mm/b12x.py
new
+339
@@ -0,0 +1,339 @@
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 torch
7
+
8
+from vllm.model_executor.layers.quantization.utils.fp8_utils import (
9
+ _upcast_e8m0_to_fp32,
10
+)
11
+from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape
12
+from vllm.model_executor.utils import replace_parameter
13
+from vllm.platforms import current_platform
14
+from vllm.utils.b12x import B12xWarmupUnit, reuse_packed_weight_storage
15
+from vllm.utils.b12x import (
16
+ get_b12x_blockscaled as _import_b12x_blockscaled,
17
+)
18
+from vllm.utils.b12x import (
19
+ get_b12x_tensor_fp8_linear as _import_b12x_tensor_fp8,
20
+)
21
+from vllm.utils.torch_utils import current_stream
22
+
23
+from .BlockScaledMMLinearKernel import (
24
+ Fp8BlockScaledMMLinearKernel,
25
+ FP8ScaledMMLinearLayerConfig,
26
+)
27
+from .ScaledMMLinearKernel import FP8ScaledMMLinearKernel
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
+class B12xFp8BlockScaledMMKernel(Fp8BlockScaledMMLinearKernel):
50
+ """K128 block-FP8 linear through the native B12X SM120 dense GEMM."""
51
+
52
+ @classmethod
53
+ def is_supported(
54
+ cls,
55
+ compute_capability: int | None = None,
56
+ ) -> tuple[bool, str | None]:
57
+ del compute_capability
58
+ if not current_platform.is_cuda():
59
+ return False, "B12X FP8 kernels are only available on CUDA"
60
+ if not current_platform.is_device_capability_family(120):
61
+ return False, "B12X FP8 kernels require a Blackwell 12x device"
62
+ blockscaled = _import_b12x_blockscaled()
63
+ if blockscaled is None:
64
+ return False, "Install the B12X backend with `pip install vllm[b12x]`"
65
+ if not blockscaled.is_supported():
66
+ return False, "B12X regular block-FP8 GEMM is not supported"
67
+ return True, None
68
+
69
+ @classmethod
70
+ def can_implement(
71
+ cls,
72
+ config: FP8ScaledMMLinearLayerConfig,
73
+ ) -> tuple[bool, str | None]:
74
+ can_implement_base, reason = super().can_implement(config)
75
+ if not can_implement_base:
76
+ return can_implement_base, reason
77
+
78
+ if config.input_dtype not in (torch.bfloat16, torch.float16):
79
+ return False, "Supports only bf16/fp16 input dtype"
80
+ if config.input_dtype != config.out_dtype:
81
+ return False, "Input and output dtype must match"
82
+
83
+ act_group_shape = config.activation_quant_key.scale.group_shape
84
+ if act_group_shape != GroupShape(1, 128):
85
+ return (
86
+ False,
87
+ "Supports only dynamic per-token group activation quantization "
88
+ "with group_shape=(1,128)",
89
+ )
90
+ weight_group_shape = config.weight_quant_key.scale.group_shape
91
+ if weight_group_shape != GroupShape(128, 128):
92
+ return False, "Supports only 128x128 block-scaled FP8 weights"
93
+
94
+ out_features, in_features = config.weight_shape
95
+ if in_features <= 0 or in_features % 128 != 0:
96
+ return False, "Input features must be a positive multiple of 128"
97
+ if out_features <= 0 or out_features % 128 != 0:
98
+ return False, "Output features must be a positive multiple of 128"
99
+ return True, None
100
+
101
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
102
+ super().process_weights_after_loading(layer)
103
+ params = self._get_layer_params(layer)
104
+ if params.weight_scale_inv is not None:
105
+ weight_scale = params.weight_scale_inv
106
+ scale_attr = params.WEIGHT_SCALE_INV
107
+ else:
108
+ weight_scale = params.weight_scale
109
+ scale_attr = params.WEIGHT_SCALE
110
+ if weight_scale is not None and weight_scale.dtype in (
111
+ torch.float8_e8m0fnu,
112
+ torch.uint8,
113
+ ):
114
+ # TODO: Remove once B12X supports 128x128 UE8M0 block scales.
115
+ replace_parameter(
116
+ layer,
117
+ scale_attr,
118
+ _upcast_e8m0_to_fp32(weight_scale).contiguous(),
119
+ )
120
+ layer.b12x_warmup_provider = self
121
+
122
+ def get_b12x_warmup_unit(
123
+ self,
124
+ layer: torch.nn.Module,
125
+ token_counts: tuple[int, ...],
126
+ output_dtype: torch.dtype,
127
+ ) -> B12xWarmupUnit:
128
+ weight = layer.weight
129
+ weight_scale = getattr(layer, "weight_scale_inv", None)
130
+ if weight_scale is None:
131
+ weight_scale = layer.weight_scale
132
+ n, k = map(int, weight.shape)
133
+
134
+ def compile() -> None:
135
+ for tokens in token_counts:
136
+ a = torch.empty((tokens, k), dtype=weight.dtype, device=weight.device)
137
+ a_scale = torch.empty(
138
+ (tokens, k // 128),
139
+ dtype=torch.float32,
140
+ device=weight.device,
141
+ )
142
+ _run_b12x_fp8_block_scaled_mm(
143
+ a, weight, a_scale, weight_scale, output_dtype
144
+ )
145
+
146
+ return B12xWarmupUnit(
147
+ name="block-FP8",
148
+ key=(
149
+ type(self),
150
+ weight.device,
151
+ n,
152
+ k,
153
+ weight.dtype,
154
+ weight_scale.dtype,
155
+ output_dtype,
156
+ ),
157
+ compile=compile,
158
+ )
159
+
160
+ def apply_block_scaled_mm(
161
+ self,
162
+ A: torch.Tensor,
163
+ B: torch.Tensor,
164
+ As: torch.Tensor,
165
+ Bs: torch.Tensor,
166
+ ) -> torch.Tensor:
167
+ return _run_b12x_fp8_block_scaled_mm(
168
+ A,
169
+ B,
170
+ As,
171
+ Bs,
172
+ self.config.out_dtype,
173
+ )
174
+
175
+
176
+def _apply_b12x_tensor_fp8_packed_linear(
177
+ layer: torch.nn.Module,
178
+ x_q: torch.Tensor,
179
+ bias: torch.Tensor | None,
180
+ out_dtype: torch.dtype,
181
+) -> torch.Tensor:
182
+ packed_weight = layer.b12x_tensor_fp8_packed_weight
183
+
184
+ tensor_fp8 = _import_b12x_tensor_fp8()
185
+ assert tensor_fp8 is not None
186
+
187
+ input_2d = x_q.reshape(-1, x_q.shape[-1]).contiguous()
188
+ output_shape = [*x_q.shape[:-1], int(packed_weight.out_features)]
189
+ output = tensor_fp8.mm(
190
+ input_2d,
191
+ packed_weight,
192
+ bias=bias,
193
+ out_dtype=out_dtype,
194
+ expected_m=max(1, int(input_2d.shape[0])),
195
+ )
196
+ return output.view(*output_shape)
197
+
198
+
199
+class B12xTensorFP8ScaledMMLinearKernel(FP8ScaledMMLinearKernel):
200
+ """Static per-tensor FP8 linear through the B12X SM12x dense GEMM."""
201
+
202
+ @classmethod
203
+ def is_supported(
204
+ cls,
205
+ compute_capability: int | None = None,
206
+ ) -> tuple[bool, str | None]:
207
+ del compute_capability
208
+ if not current_platform.is_cuda():
209
+ return False, "b12x tensor FP8 kernels are only available on CUDA"
210
+ if not current_platform.is_device_capability_family(120):
211
+ return False, "b12x tensor FP8 kernels require a Blackwell 12x device"
212
+ tensor_fp8 = _import_b12x_tensor_fp8()
213
+ if tensor_fp8 is None:
214
+ return False, "Install the B12X backend with `pip install vllm[b12x]`"
215
+ if not tensor_fp8.is_supported():
216
+ return False, "b12x.gemm.tensor_fp8_linear is not supported"
217
+ return True, None
218
+
219
+ @classmethod
220
+ def can_implement(
221
+ cls,
222
+ config: FP8ScaledMMLinearLayerConfig,
223
+ ) -> tuple[bool, str | None]:
224
+ activation_scale = config.activation_quant_key.scale
225
+ weight_scale = config.weight_quant_key.scale
226
+ if (
227
+ not activation_scale.static
228
+ or not activation_scale.group_shape.is_per_tensor()
229
+ ):
230
+ return False, "requires static per-tensor activation scales"
231
+ if not weight_scale.static or not weight_scale.group_shape.is_per_tensor():
232
+ return False, "requires static per-tensor weight scales"
233
+ if config.input_dtype not in (torch.bfloat16, torch.float16):
234
+ return False, "supports only bf16/fp16 input dtype"
235
+ if config.out_dtype not in (torch.bfloat16, torch.float16):
236
+ return False, "supports only bf16/fp16 output dtype"
237
+ out_features, in_features = config.weight_shape
238
+ if out_features <= 0 or in_features <= 0 or in_features % 32 != 0:
239
+ return False, "weight dimensions must be positive with K divisible by 32"
240
+ return True, None
241
+
242
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
243
+ weight, weight_scale, input_scale, _ = self._get_layer_params(layer)
244
+ assert weight.dtype == torch.float8_e4m3fn
245
+ assert input_scale is not None
246
+ assert weight_scale.numel() == input_scale.numel() == 1
247
+
248
+ out_features, in_features = map(int, self.config.weight_shape)
249
+ assert tuple(weight.shape) == (in_features, out_features)
250
+
251
+ tensor_fp8 = _import_b12x_tensor_fp8()
252
+ assert tensor_fp8 is not None
253
+ output_scale = (
254
+ input_scale.detach().to(torch.float32).reshape(1)
255
+ * weight_scale.detach().to(torch.float32).reshape(1)
256
+ ).contiguous()
257
+ packed_weight = tensor_fp8.pack_weight(
258
+ weight.detach().T.contiguous(),
259
+ output_scale,
260
+ )
261
+ layer.b12x_tensor_fp8_packed_weight = reuse_packed_weight_storage(
262
+ getattr(layer, "b12x_tensor_fp8_packed_weight", None),
263
+ packed_weight,
264
+ )
265
+ weight_name, weight_scale_name, _, _ = self.layer_param_names
266
+ replace_parameter(layer, weight_name, weight.new_empty((0,)))
267
+ replace_parameter(layer, weight_scale_name, weight_scale.new_empty((0,)))
268
+ layer.b12x_warmup_provider = self
269
+
270
+ def get_b12x_warmup_unit(
271
+ self,
272
+ layer: torch.nn.Module,
273
+ token_counts: tuple[int, ...],
274
+ output_dtype: torch.dtype,
275
+ ) -> B12xWarmupUnit:
276
+ packed_weight = layer.b12x_tensor_fp8_packed_weight
277
+ device = torch.device(packed_weight.values.device)
278
+
279
+ def compile() -> None:
280
+ tensor_fp8 = _import_b12x_tensor_fp8()
281
+ assert tensor_fp8 is not None
282
+ tensor_fp8.prewarm(
283
+ packed_weight,
284
+ token_counts,
285
+ out_dtype=output_dtype,
286
+ stream=current_stream().cuda_stream,
287
+ )
288
+
289
+ return B12xWarmupUnit(
290
+ name="tensor FP8",
291
+ key=(
292
+ type(self),
293
+ device,
294
+ int(packed_weight.in_features),
295
+ int(packed_weight.padded_in_features),
296
+ int(packed_weight.out_features),
297
+ output_dtype,
298
+ ),
299
+ compile=compile,
300
+ )
301
+
302
+ def apply_weights(
303
+ self,
304
+ layer: torch.nn.Module,
305
+ x: torch.Tensor,
306
+ bias: torch.Tensor | None = None,
307
+ ) -> torch.Tensor:
308
+ assert isinstance(x, torch.Tensor)
309
+ _, _, input_scale, input_scale_ub = self._get_layer_params(layer)
310
+ input_2d = x.reshape(-1, x.shape[-1])
311
+ x_q, _ = self.quant_fp8(input_2d, input_scale, input_scale_ub)
312
+ out_dtype = self.config.out_dtype
313
+ output = _apply_b12x_tensor_fp8_packed_linear(
314
+ layer,
315
+ x_q,
316
+ bias,
317
+ out_dtype,
318
+ )
319
+ return output.view(*x.shape[:-1], output.shape[-1])
320
+
321
+ def apply_scaled_mm(
322
+ self,
323
+ *,
324
+ A: torch.Tensor,
325
+ B: torch.Tensor,
326
+ out_dtype: torch.dtype,
327
+ As: torch.Tensor,
328
+ Bs: torch.Tensor,
329
+ bias: torch.Tensor | None,
330
+ output_shape: list,
331
+ ) -> torch.Tensor:
332
+ del A, B, out_dtype, As, Bs, bias, output_shape
333
+ raise NotImplementedError("b12x tensor FP8 linear overrides apply_weights")
334
+
335
+
336
+__all__ = [
337
+ "B12xFp8BlockScaledMMKernel",
338
+ "B12xTensorFP8ScaledMMLinearKernel",
339
+]
vllm/model_executor/kernels/linear/scaled_mm/b12x_block.py
deleted
-214
@@ -1,214 +0,0 @@
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
deleted
-238
@@ -1,238 +0,0 @@
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
+55
-37
@@ -2,26 +2,15 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""Warm B12X JIT kernels used by a loaded model."""
4
5
+from collections import Counter
6
+from collections.abc import Iterable
7
from typing import TYPE_CHECKING
8
9
import torch
10
11
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
-)
12
+from vllm.platforms import current_platform
13
+from vllm.utils.b12x import B12xWarmupUnit, b12x_warmup_token_counts
14
15
if TYPE_CHECKING:
16
from vllm.v1.worker.gpu_worker import Worker
@@ -29,32 +18,61 @@ if TYPE_CHECKING:
18
logger = init_logger(__name__)
19
20
21
+def _collect_warmup_units(
22
+ model: torch.nn.Module,
23
+ token_counts: tuple[int, ...],
24
+ output_dtype: torch.dtype,
25
+) -> Iterable[B12xWarmupUnit]:
26
+ units: dict[object, B12xWarmupUnit] = {}
27
+ for layer in model.modules():
28
+ provider = getattr(layer, "b12x_warmup_provider", None)
29
+ get_unit = getattr(provider, "get_b12x_warmup_unit", None)
30
+ if not callable(get_unit):
31
+ continue
32
+ unit = get_unit(layer, token_counts, output_dtype)
33
+ assert isinstance(unit, B12xWarmupUnit)
34
+ units.setdefault(unit.key, unit)
35
+ return units.values()
36
+
37
+
38
+def _compile_warmup_units(
39
+ units: Iterable[B12xWarmupUnit],
40
+) -> Counter[str]:
41
+ warmed: Counter[str] = Counter()
42
+ with torch.inference_mode():
43
+ for unit in units:
44
+ unit.compile()
45
+ warmed[unit.name] += 1
46
+ if warmed:
47
+ torch.accelerator.synchronize()
48
+ return warmed
49
+
50
+
51
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
52
+ if not current_platform.is_cuda():
53
+ return
54
+ if not current_platform.is_device_capability_family(120):
55
+ return
56
+
57
output_dtype = getattr(
58
getattr(worker, "model_config", None),
59
"dtype",
60
torch.bfloat16,
61
)
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),
62
+ if output_dtype not in (torch.bfloat16, torch.float16):
63
+ output_dtype = torch.bfloat16
64
+ token_counts = b12x_warmup_token_counts(
65
+ max_tokens=worker.scheduler_config.max_num_batched_tokens,
66
+ cudagraph_capture_sizes=cudagraph_capture_sizes,
67
)
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
- )
68
+ units = _collect_warmup_units(
69
+ worker.get_model(),
70
+ token_counts,
71
+ output_dtype,
72
+ )
73
+ for name, count in _compile_warmup_units(units).items():
74
+ logger.info_once(
75
+ "Warmed up %d B12X %s linear GEMM signatures.",
76
+ count,
77
+ name,
78
+ )
vllm/utils/b12x.py
+9
-2
@@ -5,14 +5,21 @@
5
import functools
6
import importlib
7
import importlib.util
8
-from collections.abc import Iterable
9
-from dataclasses import fields, is_dataclass
8
+from collections.abc import Callable, Hashable, Iterable
9
+from dataclasses import dataclass, fields, is_dataclass
10
from types import ModuleType
11
from typing import Any
12
13
import torch
14
15
16
+@dataclass(frozen=True)
17
+class B12xWarmupUnit:
18
+ name: str
19
+ key: Hashable
20
+ compile: Callable[[], None]
21
+
22
+
23
@functools.cache
24
def has_b12x() -> bool:
25
"""Return whether the B12X package is installed."""