A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 20:15 UTC
[MXFP4] Support for linear layers + compressed-tensors integration (#41664)
Dipika Sikka committed
May 12, 2026 at 07:49 UTC
a7b801e26d6b9d96bb49e939c0b6b3acf1d85796
11 files changed
+358
-26
tests/quantization/test_compressed_tensors.py
+29
@@ -24,6 +24,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.compressed_tenso
24
CompressedTensorsConfig,
25
CompressedTensorsLinearMethod,
26
CompressedTensorsW4A4Fp4,
27
+ CompressedTensorsW4A4Mxfp4,
28
CompressedTensorsW4A8Fp8,
29
CompressedTensorsW4A16Fp4,
30
CompressedTensorsW8A8Fp8,
@@ -689,3 +690,31 @@ def test_compressed_tensors_mxfp8_moe_setup(vllm_runner):
690
llm.apply_model(check_model)
691
output = llm.generate_greedy("Hello my name is", max_tokens=4)
692
assert output
693
+
694
+
695
+@pytest.mark.skipif(
696
+ not current_platform.is_cuda() or not current_platform.has_device_capability(80),
697
+ reason="MXFP4 requires ampere or newer",
698
+)
699
+def test_compressed_tensors_mxfp4(vllm_runner):
700
+ model_path = "nm-testing/TinyLlama-1.1B-Chat-v1.0-MXFP4"
701
+ with vllm_runner(model_path, enforce_eager=True) as llm:
702
+
703
+ def check_model(model):
704
+ layer = model.model.layers[0]
705
+
706
+ qkv_proj = layer.self_attn.qkv_proj
707
+ o_proj = layer.self_attn.o_proj
708
+ gate_up_proj = layer.mlp.gate_up_proj
709
+ down_proj = layer.mlp.down_proj
710
+
711
+ for proj in (qkv_proj, o_proj, gate_up_proj, down_proj):
712
+ assert isinstance(proj.quant_method, CompressedTensorsLinearMethod)
713
+ assert isinstance(proj.scheme, CompressedTensorsW4A4Mxfp4)
714
+
715
+ # Verify group size
716
+ assert proj.scheme.group_size == 32
717
+
718
+ llm.apply_model(check_model)
719
+ output = llm.generate_greedy("Hello my name is", max_tokens=4)
720
+ assert output
vllm/model_executor/kernels/linear/__init__.py
+68
@@ -58,6 +58,16 @@ from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
58
XPUW4A8IntLinearKernel,
59
XPUwNa16LinearKernel,
60
)
61
+from vllm.model_executor.kernels.linear.mxfp4 import (
62
+ MxFp4LinearKernel,
63
+ MxFp4LinearLayerConfig,
64
+)
65
+from vllm.model_executor.kernels.linear.mxfp4.flashinfer import (
66
+ FlashInferMxFp4LinearKernel,
67
+)
68
+from vllm.model_executor.kernels.linear.mxfp4.marlin import (
69
+ MarlinMxFp4LinearKernel,
70
+)
71
from vllm.model_executor.kernels.linear.mxfp8 import (
72
Mxfp8LinearKernel,
73
Mxfp8LinearLayerConfig,
@@ -276,6 +286,13 @@ _POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
286
],
287
}
288
289
+_POSSIBLE_MXFP4_KERNELS: dict[PlatformEnum, list[type[MxFp4LinearKernel]]] = {
290
+ PlatformEnum.CUDA: [
291
+ FlashInferMxFp4LinearKernel,
292
+ MarlinMxFp4LinearKernel,
293
+ ],
294
+}
295
+
296
# TODO make all kernels inherit from MMLinearKernel
297
# then bound _KernelT only to MMLinearKernel
298
_KernelT = TypeVar("_KernelT", bound=ScaledMMLinearKernel | MMLinearKernel)
@@ -570,6 +587,48 @@ def init_mxfp8_linear_kernel() -> Mxfp8LinearKernel:
587
)
588
589
590
+def init_mxfp4_linear_kernel() -> MxFp4LinearKernel:
591
+ """Select and instantiate the best MXFP4 linear kernel for the
592
+ current platform."""
593
+ force_kernel: type[MxFp4LinearKernel] | None = None
594
+ if envs.VLLM_MXFP4_USE_MARLIN:
595
+ force_kernel = MarlinMxFp4LinearKernel
596
+
597
+ if force_kernel is not None:
598
+ is_supported, reason = force_kernel.is_supported()
599
+ if not is_supported:
600
+ raise ValueError(
601
+ f"Forced MXFP4 kernel {force_kernel.__name__} is not "
602
+ f"supported: {reason}"
603
+ )
604
+ logger.info_once("Using %s for MXFP4 GEMM", force_kernel.__name__)
605
+ return force_kernel(MxFp4LinearLayerConfig())
606
+
607
+ platform = current_platform._enum
608
+ possible = _POSSIBLE_MXFP4_KERNELS.get(platform, [])
609
+
610
+ failure_reasons = []
611
+ for kernel_cls in possible:
612
+ if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS:
613
+ failure_reasons.append(
614
+ f" {kernel_cls.__name__} disabled by environment variable"
615
+ )
616
+ continue
617
+
618
+ is_supported, reason = kernel_cls.is_supported()
619
+ if not is_supported:
620
+ failure_reasons.append(f"{kernel_cls.__name__}: {reason}")
621
+ continue
622
+
623
+ logger.info_once("Using %s for MXFP4 GEMM", kernel_cls.__name__)
624
+ return kernel_cls(MxFp4LinearLayerConfig())
625
+
626
+ raise ValueError(
627
+ "Failed to find a kernel that can implement the "
628
+ "MXFP4 linear layer. Reasons: \n" + "\n".join(failure_reasons)
629
+ )
630
+
631
+
632
def init_wfp8_a16_linear_kernel(
633
weight_quant_key: QuantKey,
634
activation_quant_key: QuantKey,
@@ -730,6 +789,10 @@ def register_linear_kernel(
789
if platform not in _POSSIBLE_NVFP4_KERNELS:
790
_POSSIBLE_NVFP4_KERNELS[platform] = []
791
_POSSIBLE_NVFP4_KERNELS[platform].append(kernel_class)
792
+ elif kernel_type == "mxfp4":
793
+ if platform not in _POSSIBLE_MXFP4_KERNELS:
794
+ _POSSIBLE_MXFP4_KERNELS[platform] = []
795
+ _POSSIBLE_MXFP4_KERNELS[platform].append(kernel_class)
796
else:
797
raise ValueError(f"Unrecognized kernel type: {kernel_type}")
798
@@ -777,6 +840,11 @@ __all__ = [
840
"init_mxfp8_linear_kernel",
841
"Mxfp8LinearKernel",
842
"Mxfp8LinearLayerConfig",
843
+ "init_mxfp4_linear_kernel",
844
+ "MxFp4LinearKernel",
845
+ "MxFp4LinearLayerConfig",
846
+ "FlashInferMxFp4LinearKernel",
847
+ "MarlinMxFp4LinearKernel",
848
"FlashInferCutlassMxfp8LinearKernel",
849
"MarlinMxfp8LinearKernel",
850
"XPUMxFp8LinearKernel",
vllm/model_executor/kernels/linear/mxfp4/__init__.py
new
+12
@@ -0,0 +1,12 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from vllm.model_executor.kernels.linear.mxfp4.base import (
5
+ MxFp4LinearKernel,
6
+ MxFp4LinearLayerConfig,
7
+)
8
+
9
+__all__ = [
10
+ "MxFp4LinearKernel",
11
+ "MxFp4LinearLayerConfig",
12
+]
vllm/model_executor/kernels/linear/mxfp4/base.py
new
+67
@@ -0,0 +1,67 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from abc import ABC, abstractmethod
5
+from dataclasses import dataclass
6
+
7
+import torch
8
+
9
+
10
+@dataclass
11
+class MxFp4LinearLayerConfig:
12
+ """Configuration for an MXFP4 linear layer.
13
+
14
+ All MXFP4 layers share the same structure: packed uint8 weights (2 FP4 values per
15
+ byte) and per-block weight scales (group size 32).
16
+ """
17
+
18
+ pass
19
+
20
+
21
+class MxFp4LinearKernel(ABC):
22
+ """Base class for MXFP4 quantized linear kernels.
23
+
24
+ Each subclass implements a specific GEMM backend (CUTLASS, Marlin, etc).
25
+ The kernel selection mechanism iterates over registered subclasses in
26
+ priority order,calling ``is_supported`` and ``can_implement`` to find the best
27
+ match for the current hardware.
28
+ """
29
+
30
+ def __init__(self, config: MxFp4LinearLayerConfig) -> None:
31
+ assert self.can_implement(config)[0]
32
+ assert self.is_supported()[0]
33
+ self.config = config
34
+
35
+ @classmethod
36
+ @abstractmethod
37
+ def is_supported(
38
+ cls, compute_capability: int | None = None
39
+ ) -> tuple[bool, str | None]:
40
+ """Return whether this kernel can run on the current platform."""
41
+ raise NotImplementedError
42
+
43
+ @classmethod
44
+ @abstractmethod
45
+ def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
46
+ """Return whether this kernel can handle *config*."""
47
+ raise NotImplementedError
48
+
49
+ @abstractmethod
50
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
51
+ """Transform weights into the format required by this kernel.
52
+
53
+ Called once after checkpoint weights have been loaded onto the
54
+ device. Implementations should repack / swizzle / pad weights
55
+ and scales in-place on *layer*.
56
+ """
57
+ raise NotImplementedError
58
+
59
+ @abstractmethod
60
+ def apply_weights(
61
+ self,
62
+ layer: torch.nn.Module,
63
+ x: torch.Tensor,
64
+ bias: torch.Tensor | None = None,
65
+ ) -> torch.Tensor:
66
+ """Run the quantized GEMM."""
67
+ raise NotImplementedError
vllm/model_executor/kernels/linear/mxfp4/flashinfer.py
new
+74
@@ -0,0 +1,74 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+from torch.nn.parameter import Parameter
6
+
7
+from vllm.model_executor.layers.fused_moe.experts.cutlass_moe import (
8
+ swizzle_mxfp4_scales,
9
+)
10
+from vllm.platforms import current_platform
11
+from vllm.utils.flashinfer import has_flashinfer_cutedsl
12
+
13
+from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
14
+
15
+_MXFP4_GROUP_SIZE = 32
16
+
17
+
18
+class FlashInferMxFp4LinearKernel(MxFp4LinearKernel):
19
+ """MXFP4 W4A4 GEMM via FlashInfer CUTLASS (SM100+)."""
20
+
21
+ @classmethod
22
+ def is_supported(
23
+ cls, compute_capability: int | None = None
24
+ ) -> tuple[bool, str | None]:
25
+ if current_platform.has_device_capability(100) and has_flashinfer_cutedsl():
26
+ return True, None
27
+ return False, "FlashInfer + >=sm_100 (Blackwell) required"
28
+
29
+ @classmethod
30
+ def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
31
+ return True, None
32
+
33
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
34
+ N, scale_K = layer.weight_scale.shape
35
+ K = scale_K * _MXFP4_GROUP_SIZE
36
+
37
+ # swizzle pads N to the next multiple of 128 for CUTLASS tiling
38
+ padded_N = ((N + 127) // 128) * 128
39
+ layer.weight_scale = Parameter(
40
+ swizzle_mxfp4_scales(layer.weight_scale.data, N, K).reshape(padded_N, -1),
41
+ requires_grad=False,
42
+ )
43
+
44
+ def apply_weights(
45
+ self,
46
+ layer: torch.nn.Module,
47
+ x: torch.Tensor,
48
+ bias: torch.Tensor | None = None,
49
+ ) -> torch.Tensor:
50
+ from vllm.utils.flashinfer import (
51
+ flashinfer_mxfp4_quantize,
52
+ flashinfer_scaled_fp4_mm,
53
+ )
54
+
55
+ weight = layer.weight
56
+ out_shape = x.shape[:-1] + (layer.output_size_per_partition,)
57
+ x_2d = x.reshape(-1, x.shape[-1])
58
+
59
+ x_fp4, x_scale = flashinfer_mxfp4_quantize(x_2d)
60
+ out = flashinfer_scaled_fp4_mm(
61
+ x_fp4,
62
+ weight,
63
+ x_scale,
64
+ layer.weight_scale,
65
+ alpha=None,
66
+ out_dtype=x.dtype,
67
+ backend="cute-dsl",
68
+ block_size=_MXFP4_GROUP_SIZE,
69
+ use_nvfp4=False,
70
+ )
71
+
72
+ if bias is not None:
73
+ out = out + bias
74
+ return out.view(out_shape)
vllm/model_executor/kernels/linear/mxfp4/marlin.py
new
+52
@@ -0,0 +1,52 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import torch
5
+
6
+from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
7
+
8
+
9
+class MarlinMxFp4LinearKernel(MxFp4LinearKernel):
10
+ @classmethod
11
+ def is_supported(
12
+ cls, compute_capability: int | None = None
13
+ ) -> tuple[bool, str | None]:
14
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
15
+ is_fp4_marlin_supported,
16
+ )
17
+
18
+ if is_fp4_marlin_supported():
19
+ return True, None
20
+ return False, "Marlin FP4 not available"
21
+
22
+ @classmethod
23
+ def can_implement(cls, c: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
24
+ return True, None
25
+
26
+ def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
27
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
28
+ prepare_fp4_layer_for_marlin,
29
+ )
30
+
31
+ prepare_fp4_layer_for_marlin(layer)
32
+
33
+ def apply_weights(
34
+ self,
35
+ layer: torch.nn.Module,
36
+ x: torch.Tensor,
37
+ bias: torch.Tensor | None = None,
38
+ ) -> torch.Tensor:
39
+ from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
40
+ apply_fp4_marlin_linear,
41
+ )
42
+
43
+ return apply_fp4_marlin_linear(
44
+ input=x,
45
+ weight=layer.weight,
46
+ weight_scale=layer.weight_scale,
47
+ weight_global_scale=None,
48
+ workspace=layer.workspace,
49
+ size_n=layer.output_size_per_partition,
50
+ size_k=layer.input_size_per_partition,
51
+ bias=bias,
52
+ )
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+2
-2
@@ -42,10 +42,10 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
42
CompressedTensors24,
43
CompressedTensorsScheme,
44
CompressedTensorsW4A4Fp4,
45
+ CompressedTensorsW4A4Mxfp4,
46
CompressedTensorsW4A8Fp8,
47
CompressedTensorsW4A8Int,
48
CompressedTensorsW4A16Fp4,
48
- CompressedTensorsW4A16Mxfp4,
49
CompressedTensorsW8A8Fp8,
50
CompressedTensorsW8A8Int8,
51
CompressedTensorsW8A8Mxfp8,
@@ -625,7 +625,7 @@ class CompressedTensorsConfig(QuantizationConfig):
625
return CompressedTensorsW4A16Fp4()
626
627
if self._is_mxfp4(weight_quant):
628
- return CompressedTensorsW4A16Mxfp4()
628
+ return CompressedTensorsW4A4Mxfp4()
629
630
if self._is_mxfp8(weight_quant):
631
return CompressedTensorsW8A8Mxfp8()
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe/compressed_tensors_moe_w4a4_mxfp4.py
+1
@@ -41,6 +41,7 @@ class CompressedTensorsW4A4Mxfp4MoEMethod(CompressedTensorsMoEMethod):
41
super().__init__(moe)
42
self.group_size = 32
43
self.mxfp4_backend = Mxfp4MoeBackend.MARLIN
44
+ # use cutlass if supported, otherwise fallback to marlin for weight-only FP4
45
self.use_cutlass_mxfp4 = CutlassExpertsMxfp4._supports_current_device()
46
self.experts_cls: type[mk.FusedMoEExperts]
47
if self.use_cutlass_mxfp4:
vllm/model_executor/layers/quantization/compressed_tensors/schemes/__init__.py
+2
-2
@@ -2,10 +2,10 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
from .compressed_tensors_scheme import CompressedTensorsScheme
5
+from .compressed_tensors_w4a4_mxfp4 import CompressedTensorsW4A4Mxfp4
6
from .compressed_tensors_w4a4_nvfp4 import CompressedTensorsW4A4Fp4
7
from .compressed_tensors_w4a8_fp8 import CompressedTensorsW4A8Fp8
8
from .compressed_tensors_w4a8_int import CompressedTensorsW4A8Int
8
-from .compressed_tensors_w4a16_mxfp4 import CompressedTensorsW4A16Mxfp4
9
from .compressed_tensors_w4a16_nvfp4 import CompressedTensorsW4A16Fp4
10
from .compressed_tensors_w8a8_fp8 import CompressedTensorsW8A8Fp8
11
from .compressed_tensors_w8a8_int8 import CompressedTensorsW8A8Int8
@@ -25,7 +25,7 @@ __all__ = [
25
"WNA16_SUPPORTED_BITS",
26
"CompressedTensors24",
27
"CompressedTensorsW4A16Fp4",
28
- "CompressedTensorsW4A16Mxfp4",
28
+ "CompressedTensorsW4A4Mxfp4",
29
"CompressedTensorsW4A4Fp4",
30
"CompressedTensorsW4A8Int",
31
"CompressedTensorsW4A8Fp8",
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_mxfp4.py
renamed
+10
-20
@@ -5,24 +5,21 @@ from collections.abc import Callable
5
import torch
6
from torch.nn.parameter import Parameter
7
8
+from vllm.model_executor.kernels.linear import init_mxfp4_linear_kernel
9
from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
10
CompressedTensorsScheme,
11
)
11
-from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
12
- apply_fp4_marlin_linear,
13
- prepare_fp4_layer_for_marlin,
14
-)
12
from vllm.model_executor.parameter import (
13
GroupQuantScaleParameter,
14
ModelWeightParameter,
15
)
16
20
-__all__ = ["CompressedTensorsW4A16Mxfp4"]
17
+__all__ = ["CompressedTensorsW4A4Mxfp4"]
18
19
23
-class CompressedTensorsW4A16Mxfp4(CompressedTensorsScheme):
20
+class CompressedTensorsW4A4Mxfp4(CompressedTensorsScheme):
21
"""
25
- Compressed tensors scheme for MXFP4 weight-only quantization.
22
+ Compressed tensors scheme for MXFP4.
23
24
Supports models quantized with the compressed-tensors mxfp4-pack-quantized
25
format.
@@ -31,10 +28,14 @@ class CompressedTensorsW4A16Mxfp4(CompressedTensorsScheme):
28
- 4-bit float weights (E2M1) packed into uint8
29
- Per-group E8M0 scales with group_size=32
30
- No global scale (unlike NVFP4)
31
+
32
+ On SM100+ with FlashInfer: true W4A4 (activations dynamically quantized).
33
+ Otherwise: W4A16 weight-only via Marlin.
34
"""
35
36
def __init__(self):
37
self.group_size = 32
38
+ self.kernel = init_mxfp4_linear_kernel()
39
40
@classmethod
41
def get_min_capability(cls) -> int:
@@ -82,11 +83,9 @@ class CompressedTensorsW4A16Mxfp4(CompressedTensorsScheme):
83
layer.register_parameter("weight_scale", weight_scale)
84
85
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
85
- # Rename weight_packed to weight that marlin expects
86
layer.weight = Parameter(layer.weight_packed.data, requires_grad=False)
87
del layer.weight_packed
88
-
89
- prepare_fp4_layer_for_marlin(layer)
88
+ self.kernel.process_weights_after_loading(layer)
89
90
def apply_weights(
91
self,
@@ -94,13 +93,4 @@ class CompressedTensorsW4A16Mxfp4(CompressedTensorsScheme):
93
x: torch.Tensor,
94
bias: torch.Tensor | None = None,
95
) -> torch.Tensor:
97
- return apply_fp4_marlin_linear(
98
- input=x,
99
- weight=layer.weight,
100
- weight_scale=layer.weight_scale,
101
- weight_global_scale=None,
102
- workspace=layer.workspace,
103
- size_n=layer.output_size_per_partition,
104
- size_k=layer.input_size_per_partition,
105
- bias=bias,
106
- )
96
+ return self.kernel.apply_weights(layer, x, bias)
vllm/utils/flashinfer.py
+41
-2
@@ -20,6 +20,7 @@ import torch
20
import vllm.envs as envs
21
from vllm.logger import init_logger
22
from vllm.platforms import current_platform
23
+from vllm.utils.math_utils import cdiv
24
25
logger = init_logger(__name__)
26
@@ -480,6 +481,8 @@ if has_flashinfer():
481
dtype: torch.dtype,
482
use_8x4_sf_layout: bool,
483
backend: str,
484
+ block_size: int = 16,
485
+ use_nvfp4: bool = True,
486
) -> torch.Tensor:
487
from flashinfer import mm_fp4 as flashinfer_mm_fp4_
488
@@ -490,8 +493,9 @@ if has_flashinfer():
493
B_scale,
494
g_scale,
495
dtype,
493
- block_size=16,
496
+ block_size=block_size,
497
use_8x4_sf_layout=use_8x4_sf_layout,
498
+ use_nvfp4=use_nvfp4,
499
backend=backend,
500
)
501
@@ -507,9 +511,36 @@ if has_flashinfer():
511
dtype: torch.dtype,
512
use_8x4_sf_layout: bool,
513
backend: str,
514
+ block_size: int = 16,
515
+ use_nvfp4: bool = True,
516
) -> torch.Tensor:
517
return torch.empty(A.shape[0], B.shape[1], dtype=dtype, device=A.device)
518
519
+ @torch.library.custom_op(
520
+ "vllm::flashinfer_mxfp4_quantize",
521
+ mutates_args=[],
522
+ device_types="cuda",
523
+ )
524
+ def flashinfer_mxfp4_quantize(
525
+ a: torch.Tensor,
526
+ ) -> tuple[torch.Tensor, torch.Tensor]:
527
+ from flashinfer import mxfp4_quantize as _mxfp4_quantize
528
+
529
+ return _mxfp4_quantize(a)
530
+
531
+ @torch.library.register_fake("vllm::flashinfer_mxfp4_quantize")
532
+ def flashinfer_mxfp4_quantize_fake(
533
+ a: torch.Tensor,
534
+ ) -> tuple[torch.Tensor, torch.Tensor]:
535
+ m, k = a.shape
536
+ sf_vec_size = 32
537
+ padded_m = cdiv(m, 128) * 128
538
+ sf_cols = cdiv(k // sf_vec_size, 4) * 4
539
+ return (
540
+ torch.empty(m, k // 2, dtype=torch.uint8, device=a.device),
541
+ torch.empty(padded_m, sf_cols, dtype=torch.uint8, device=a.device),
542
+ )
543
+
544
@torch.library.custom_op(
545
"vllm::bmm_fp8",
546
mutates_args=[],
@@ -658,15 +689,20 @@ def flashinfer_scaled_fp4_mm(
689
b: torch.Tensor,
690
block_scale_a: torch.Tensor,
691
block_scale_b: torch.Tensor,
661
- alpha: torch.Tensor,
692
+ alpha: torch.Tensor | None,
693
out_dtype: torch.dtype,
694
backend: str,
695
+ block_size: int = 16,
696
+ use_nvfp4: bool = True,
697
) -> torch.Tensor:
698
assert a.ndim == 2 and b.ndim == 2
699
assert block_scale_a.ndim == 2 and block_scale_b.ndim == 2
700
assert a.stride(-1) == 1 and b.stride(-1) == 1
701
assert a.shape[1] == b.shape[1]
702
703
+ if alpha is None:
704
+ alpha = torch.ones(1, dtype=torch.float32, device=a.device)
705
+
706
if backend in ("cutlass", "cudnn"):
707
block_scale_a = block_scale_a.view(torch.uint8)
708
block_scale_b = block_scale_b.view(torch.uint8)
@@ -682,6 +718,8 @@ def flashinfer_scaled_fp4_mm(
718
out_dtype,
719
use_8x4_sf_layout=use_8x4_sf_layout,
720
backend=backend,
721
+ block_size=block_size,
722
+ use_nvfp4=use_nvfp4,
723
)
724
725
@@ -904,6 +942,7 @@ __all__ = [
942
"supports_trtllm_attention",
943
"can_use_trtllm_attention",
944
"use_trtllm_attention",
945
+ "flashinfer_mxfp4_quantize",
946
"flashinfer_scaled_fp4_mm",
947
"flashinfer_scaled_fp4_mm_out",
948
"flashinfer_scaled_fp8_mm",