@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 2800706f06
Fork 0

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 5, 09:45 UTC

[Refactor] Move NVFP4 GEMM management into NvFp4LinearKernel (#39129)

Signed-off-by: mgoin <mgoin64@gmail.com>

Michael Goin committed Apr 9, 2026 at 21:05 UTC 2800706f0649955a65334e8ccce35654cf988727
11 files changed +697 -377
vllm/model_executor/kernels/linear/__init__.py
+133
@@ -55,6 +55,27 @@ from vllm.model_executor.kernels.linear.mixed_precision.xpu import (
55 XPUW4A8IntLinearKernel,
56 XPUwNa16LinearKernel,
57 )
58 +from vllm.model_executor.kernels.linear.nvfp4 import (
59 + NvFp4LinearKernel,
60 + NvFp4LinearLayerConfig,
61 +)
62 +from vllm.model_executor.kernels.linear.nvfp4.cutlass import (
63 + CutlassNvFp4LinearKernel,
64 +)
65 +from vllm.model_executor.kernels.linear.nvfp4.emulation import (
66 + EmulationNvFp4LinearKernel,
67 +)
68 +from vllm.model_executor.kernels.linear.nvfp4.fbgemm import (
69 + FbgemmNvFp4LinearKernel,
70 +)
71 +from vllm.model_executor.kernels.linear.nvfp4.flashinfer import (
72 + FlashInferCudnnNvFp4LinearKernel,
73 + FlashInferCutlassNvFp4LinearKernel,
74 + FlashInferTrtllmNvFp4LinearKernel,
75 +)
76 +from vllm.model_executor.kernels.linear.nvfp4.marlin import (
77 + MarlinNvFp4LinearKernel,
78 +)
79 from vllm.model_executor.kernels.linear.scaled_mm import (
80 Fp8BlockScaledMMLinearKernel,
81 FP8ScaledMMLinearKernel,
@@ -180,6 +201,22 @@ _POSSIBLE_KERNELS: dict[PlatformEnum, list[type[MPLinearKernel]]] = {
201 ],
202 }
203
204 +# in priority/performance order (when available)
205 +_POSSIBLE_NVFP4_KERNELS: dict[PlatformEnum, list[type[NvFp4LinearKernel]]] = {
206 + PlatformEnum.CUDA: [
207 + FlashInferCutlassNvFp4LinearKernel,
208 + CutlassNvFp4LinearKernel,
209 + MarlinNvFp4LinearKernel,
210 + FlashInferTrtllmNvFp4LinearKernel,
211 + FlashInferCudnnNvFp4LinearKernel,
212 + FbgemmNvFp4LinearKernel,
213 + EmulationNvFp4LinearKernel,
214 + ],
215 + PlatformEnum.ROCM: [
216 + EmulationNvFp4LinearKernel,
217 + ],
218 +}
219 +
220 # TODO make all kernels inherit from MMLinearKernel
221 # then bound _KernelT only to MMLinearKernel
222 _KernelT = TypeVar("_KernelT", bound=ScaledMMLinearKernel | MMLinearKernel)
@@ -426,6 +463,88 @@ def choose_mp_linear_kernel(
463 )
464
465
466 +# Maps VLLM_NVFP4_GEMM_BACKEND env var values to kernel classes.
467 +_NVFP4_BACKEND_TO_KERNEL: dict[str, type[NvFp4LinearKernel]] = {
468 + "flashinfer-cutlass": FlashInferCutlassNvFp4LinearKernel,
469 + "cutlass": CutlassNvFp4LinearKernel,
470 + "marlin": MarlinNvFp4LinearKernel,
471 + "flashinfer-trtllm": FlashInferTrtllmNvFp4LinearKernel,
472 + "flashinfer-cudnn": FlashInferCudnnNvFp4LinearKernel,
473 + "emulation": EmulationNvFp4LinearKernel,
474 +}
475 +
476 +
477 +def init_nvfp4_linear_kernel() -> NvFp4LinearKernel:
478 + """Select and instantiate the best NVFP4 linear kernel for the
479 + current platform."""
480 + config = NvFp4LinearLayerConfig()
481 +
482 + # Env-var overrides.
483 + force_kernel: type[NvFp4LinearKernel] | None = None
484 + if envs.VLLM_USE_FBGEMM:
485 + force_kernel = FbgemmNvFp4LinearKernel
486 + elif envs.VLLM_USE_NVFP4_CT_EMULATIONS:
487 + force_kernel = EmulationNvFp4LinearKernel
488 + elif envs.VLLM_NVFP4_GEMM_BACKEND is not None:
489 + backend_name = envs.VLLM_NVFP4_GEMM_BACKEND
490 + force_kernel = _NVFP4_BACKEND_TO_KERNEL.get(backend_name)
491 + if force_kernel is None:
492 + raise ValueError(
493 + f"Unknown VLLM_NVFP4_GEMM_BACKEND={backend_name!r}. "
494 + f"Valid choices: {list(_NVFP4_BACKEND_TO_KERNEL.keys())}"
495 + )
496 +
497 + if force_kernel is not None:
498 + is_supported, reason = force_kernel.is_supported()
499 + if not is_supported:
500 + raise ValueError(
501 + f"Forced NVFP4 kernel {force_kernel.__name__} is not "
502 + f"supported: {reason}"
503 + )
504 + logger.info_once("Using %s for NVFP4 GEMM", force_kernel.__name__)
505 + return force_kernel(config)
506 +
507 + # Auto-select from registry.
508 + platform = current_platform._enum
509 + possible = _POSSIBLE_NVFP4_KERNELS.get(platform, [])
510 +
511 + failure_reasons = []
512 + for kernel_cls in possible:
513 + if kernel_cls.__name__ in envs.VLLM_DISABLED_KERNELS:
514 + failure_reasons.append(
515 + f" {kernel_cls.__name__} disabled by environment variable"
516 + )
517 + continue
518 +
519 + is_supported, reason = kernel_cls.is_supported()
520 + if not is_supported:
521 + failure_reasons.append(f"{kernel_cls.__name__}: {reason}")
522 + continue
523 +
524 + can_implement, reason = kernel_cls.can_implement(config)
525 + if not can_implement:
526 + failure_reasons.append(f"{kernel_cls.__name__}: {reason}")
527 + continue
528 +
529 + if kernel_cls is EmulationNvFp4LinearKernel and failure_reasons:
530 + logger.warning_once(
531 + "NVFP4 linear falling back to the slow and unoptimized "
532 + "emulation backend as no optimized backend is available "
533 + "(unavailable reasons:\n - %s\n). "
534 + "In case you expect one of these backends to be used, "
535 + "please verify your environment.",
536 + "\n - ".join(failure_reasons),
537 + )
538 +
539 + logger.info_once("Using %s for NVFP4 GEMM", kernel_cls.__name__)
540 + return kernel_cls(config)
541 +
542 + raise ValueError(
543 + "Failed to find a kernel that can implement the "
544 + "NVFP4 linear layer. Reasons: \n" + "\n".join(failure_reasons)
545 + )
546 +
547 +
548 def register_linear_kernel(
549 kernel_class: type,
550 platform: PlatformEnum,
@@ -455,6 +574,10 @@ def register_linear_kernel(
574 if platform not in _POSSIBLE_FP8_KERNELS:
575 _POSSIBLE_FP8_KERNELS[platform] = []
576 _POSSIBLE_FP8_KERNELS[platform].append(kernel_class)
577 + elif kernel_type == "nvfp4":
578 + if platform not in _POSSIBLE_NVFP4_KERNELS:
579 + _POSSIBLE_NVFP4_KERNELS[platform] = []
580 + _POSSIBLE_NVFP4_KERNELS[platform].append(kernel_class)
581 else:
582 raise ValueError(f"Unrecognized kernel type: {kernel_type}")
583
@@ -462,6 +585,7 @@ def register_linear_kernel(
585 __all__ = [
586 "init_fp8_linear_kernel",
587 "init_int8_linear_kernel",
588 + "init_nvfp4_linear_kernel",
589 "choose_mp_linear_kernel",
590 "register_linear_kernel",
591 "FP8ScaledMMLinearKernel",
@@ -470,6 +594,8 @@ __all__ = [
594 "FP8ScaledMMLinearLayerConfig",
595 "Int8ScaledMMLinearLayerConfig",
596 "ScaledMMLinearLayerConfig",
597 + "NvFp4LinearKernel",
598 + "NvFp4LinearLayerConfig",
599 "AiterInt8ScaledMMLinearKernel",
600 "CPUInt8ScaledMMLinearKernel",
601 "CutlassFP8ScaledMMLinearKernel",
@@ -492,6 +618,13 @@ __all__ = [
618 "MarlinLinearKernel",
619 "XPUW4A8IntLinearKernel",
620 "XPUwNa16LinearKernel",
621 + "CutlassNvFp4LinearKernel",
622 + "EmulationNvFp4LinearKernel",
623 + "FbgemmNvFp4LinearKernel",
624 + "FlashInferCutlassNvFp4LinearKernel",
625 + "FlashInferTrtllmNvFp4LinearKernel",
626 + "FlashInferCudnnNvFp4LinearKernel",
627 + "MarlinNvFp4LinearKernel",
628 "_KernelT",
629 "DeepGemmFp8BlockScaledMMKernel",
630 "FlashInferFp8DeepGEMMDynamicBlockScaledKernel",
vllm/model_executor/kernels/linear/nvfp4/__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.nvfp4.base import (
5 + NvFp4LinearKernel,
6 + NvFp4LinearLayerConfig,
7 +)
8 +
9 +__all__ = [
10 + "NvFp4LinearKernel",
11 + "NvFp4LinearLayerConfig",
12 +]
vllm/model_executor/kernels/linear/nvfp4/base.py new
+68
@@ -0,0 +1,68 @@
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 NvFp4LinearLayerConfig:
12 + """Configuration for an NVFP4 linear layer.
13 +
14 + All NVFP4 layers share the same structure: packed uint8 weights (2 FP4 values per
15 + byte), FP8-E4M3 per-block weight scales (group size 16), and scalar global
16 + scales for both weights and activations.
17 + """
18 +
19 + pass
20 +
21 +
22 +class NvFp4LinearKernel(ABC):
23 + """Base class for NVFP4 quantized linear kernels.
24 +
25 + Each subclass implements a specific GEMM backend (CUTLASS, Marlin, etc).
26 + The kernel selection mechanism iterates over registered subclasses in
27 + priority order,calling ``is_supported`` and ``can_implement`` to find the best
28 + match for the current hardware.
29 + """
30 +
31 + def __init__(self, config: NvFp4LinearLayerConfig) -> None:
32 + assert self.can_implement(config)[0]
33 + assert self.is_supported()[0]
34 + self.config = config
35 +
36 + @classmethod
37 + @abstractmethod
38 + def is_supported(
39 + cls, compute_capability: int | None = None
40 + ) -> tuple[bool, str | None]:
41 + """Return whether this kernel can run on the current platform."""
42 + raise NotImplementedError
43 +
44 + @classmethod
45 + @abstractmethod
46 + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
47 + """Return whether this kernel can handle *config*."""
48 + raise NotImplementedError
49 +
50 + @abstractmethod
51 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
52 + """Transform weights into the format required by this kernel.
53 +
54 + Called once after checkpoint weights have been loaded onto the
55 + device. Implementations should repack / swizzle / pad weights
56 + and scales in-place on *layer*.
57 + """
58 + raise NotImplementedError
59 +
60 + @abstractmethod
61 + def apply_weights(
62 + self,
63 + layer: torch.nn.Module,
64 + x: torch.Tensor,
65 + bias: torch.Tensor | None = None,
66 + ) -> torch.Tensor:
67 + """Run the quantized GEMM."""
68 + raise NotImplementedError
vllm/model_executor/kernels/linear/nvfp4/cutlass.py new
+80
@@ -0,0 +1,80 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm._custom_ops import (
7 + cutlass_scaled_fp4_mm,
8 + scaled_fp4_quant,
9 +)
10 +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
11 + cutlass_fp4_supported,
12 + pad_nvfp4_activation_for_cutlass,
13 + pad_nvfp4_weight_for_cutlass,
14 + slice_nvfp4_output,
15 + swizzle_blockscale,
16 +)
17 +
18 +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
19 +
20 +
21 +class CutlassNvFp4LinearKernel(NvFp4LinearKernel):
22 + """NVFP4 GEMM via the vLLM CUTLASS kernel."""
23 +
24 + @classmethod
25 + def is_supported(
26 + cls, compute_capability: int | None = None
27 + ) -> tuple[bool, str | None]:
28 + if cutlass_fp4_supported():
29 + return True, None
30 + return False, "CUTLASS FP4 kernels not available"
31 +
32 + @classmethod
33 + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
34 + return True, None
35 +
36 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
37 + layer.weight_scale = torch.nn.Parameter(
38 + swizzle_blockscale(layer.weight_scale.data), requires_grad=False
39 + )
40 + padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
41 + layer.weight.data
42 + )
43 + layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
44 + layer.weights_padding_cols = weights_padding_cols
45 +
46 + def apply_weights(
47 + self,
48 + layer: torch.nn.Module,
49 + x: torch.Tensor,
50 + bias: torch.Tensor | None = None,
51 + ) -> torch.Tensor:
52 + output_size = layer.output_size_per_partition
53 + output_dtype = x.dtype
54 + output_shape = [*x.shape[:-1], output_size]
55 +
56 + x_fp4, x_blockscale = scaled_fp4_quant(
57 + x,
58 + layer.input_global_scale_inv,
59 + is_sf_swizzled_layout=True,
60 + backend="cutlass",
61 + )
62 +
63 + x_fp4 = pad_nvfp4_activation_for_cutlass(
64 + x_fp4, getattr(layer, "weights_padding_cols", 0)
65 + )
66 +
67 + out = cutlass_scaled_fp4_mm(
68 + x_fp4,
69 + layer.weight,
70 + x_blockscale,
71 + layer.weight_scale,
72 + layer.alpha,
73 + output_dtype,
74 + )
75 +
76 + out = slice_nvfp4_output(out, output_size)
77 +
78 + if bias is not None:
79 + out = out + bias
80 + return out.view(*output_shape)
vllm/model_executor/kernels/linear/nvfp4/emulation.py new
+49
@@ -0,0 +1,49 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
7 + kE2M1ToFloat_handle,
8 + run_nvfp4_emulations,
9 +)
10 +
11 +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
12 +
13 +
14 +class EmulationNvFp4LinearKernel(NvFp4LinearKernel):
15 + """Software emulation fallback for NVFP4 (dequant → BF16 matmul)."""
16 +
17 + @classmethod
18 + def is_supported(
19 + cls, compute_capability: int | None = None
20 + ) -> tuple[bool, str | None]:
21 + # Always available as a last-resort fallback.
22 + return True, None
23 +
24 + @classmethod
25 + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
26 + return True, None
27 +
28 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
29 + # Move the E2M1 lookup table to the device now, because
30 + # `.to(device)` is not allowed during CUDA graph capture.
31 + kE2M1ToFloat_handle.val = kE2M1ToFloat_handle.val.to(layer.weight.device)
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 + out = run_nvfp4_emulations(
40 + x=x,
41 + input_global_scale=layer.input_global_scale_inv,
42 + weight=layer.weight,
43 + weight_scale_swizzled=layer.weight_scale,
44 + weight_global_scale=layer.weight_global_scale,
45 + swizzle=False,
46 + )
47 + if bias is not None:
48 + out = out + bias
49 + return out
vllm/model_executor/kernels/linear/nvfp4/fbgemm.py new
+69
@@ -0,0 +1,69 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm._custom_ops import scaled_fp4_quant
7 +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
8 + slice_nvfp4_output,
9 + swizzle_blockscale,
10 +)
11 +from vllm.utils.import_utils import has_fbgemm_gpu
12 +
13 +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
14 +
15 +
16 +class FbgemmNvFp4LinearKernel(NvFp4LinearKernel):
17 + """NVFP4 GEMM via FBGEMM."""
18 +
19 + @classmethod
20 + def is_supported(
21 + cls, compute_capability: int | None = None
22 + ) -> tuple[bool, str | None]:
23 + if has_fbgemm_gpu():
24 + return True, None
25 + return False, "fbgemm_gpu required"
26 +
27 + @classmethod
28 + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
29 + return True, None
30 +
31 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
32 + swizzled = swizzle_blockscale(layer.weight_scale.data)
33 + layer.weight_scale = torch.nn.Parameter(
34 + swizzled.view(-1).view(torch.uint8), requires_grad=False
35 + )
36 +
37 + def apply_weights(
38 + self,
39 + layer: torch.nn.Module,
40 + x: torch.Tensor,
41 + bias: torch.Tensor | None = None,
42 + ) -> torch.Tensor:
43 + import fbgemm_gpu # noqa: F401 - registers torch.ops.fbgemm.*
44 +
45 + output_size = layer.output_size_per_partition
46 + output_dtype = x.dtype
47 + output_shape = [*x.shape[:-1], output_size]
48 +
49 + x_fp4, x_blockscale = scaled_fp4_quant(
50 + x,
51 + layer.input_global_scale_inv,
52 + is_sf_swizzled_layout=True,
53 + backend="fbgemm",
54 + )
55 +
56 + out = torch.ops.fbgemm.f4f4bf16(
57 + x_fp4,
58 + layer.weight,
59 + x_blockscale.view(-1).view(torch.uint8),
60 + layer.weight_scale,
61 + layer.alpha,
62 + use_mx=False,
63 + ).to(output_dtype)
64 +
65 + out = slice_nvfp4_output(out, output_size)
66 +
67 + if bias is not None:
68 + out = out + bias
69 + return out.view(*output_shape)
vllm/model_executor/kernels/linear/nvfp4/flashinfer.py new
+218
@@ -0,0 +1,218 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm._custom_ops import scaled_fp4_quant
7 +from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
8 + pad_nvfp4_activation_for_cutlass,
9 + pad_nvfp4_weight_for_cutlass,
10 + slice_nvfp4_output,
11 + swizzle_blockscale,
12 +)
13 +from vllm.platforms import current_platform
14 +from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm, has_flashinfer
15 +
16 +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
17 +
18 +
19 +class FlashInferCutlassNvFp4LinearKernel(NvFp4LinearKernel):
20 + """NVFP4 GEMM via FlashInfer's CUTLASS wrapper."""
21 +
22 + @classmethod
23 + def is_supported(
24 + cls, compute_capability: int | None = None
25 + ) -> tuple[bool, str | None]:
26 + from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
27 + cutlass_fp4_supported,
28 + )
29 +
30 + if (
31 + cutlass_fp4_supported()
32 + and current_platform.has_device_capability(100)
33 + and has_flashinfer()
34 + ):
35 + return True, None
36 + return False, "FlashInfer + >=sm_100 required"
37 +
38 + @classmethod
39 + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
40 + return True, None
41 +
42 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
43 + layer.weight_scale = torch.nn.Parameter(
44 + swizzle_blockscale(layer.weight_scale.data), requires_grad=False
45 + )
46 + padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
47 + layer.weight.data
48 + )
49 + layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
50 + layer.weights_padding_cols = weights_padding_cols
51 +
52 + def apply_weights(
53 + self,
54 + layer: torch.nn.Module,
55 + x: torch.Tensor,
56 + bias: torch.Tensor | None = None,
57 + ) -> torch.Tensor:
58 + output_size = layer.output_size_per_partition
59 + output_dtype = x.dtype
60 + output_shape = [*x.shape[:-1], output_size]
61 +
62 + x_fp4, x_blockscale = scaled_fp4_quant(
63 + x,
64 + layer.input_global_scale_inv,
65 + is_sf_swizzled_layout=True,
66 + backend="flashinfer-cutlass",
67 + )
68 +
69 + x_fp4 = pad_nvfp4_activation_for_cutlass(
70 + x_fp4, getattr(layer, "weights_padding_cols", 0)
71 + )
72 +
73 + out = flashinfer_scaled_fp4_mm(
74 + x_fp4,
75 + layer.weight,
76 + x_blockscale,
77 + layer.weight_scale,
78 + layer.alpha,
79 + output_dtype,
80 + backend="cutlass",
81 + )
82 +
83 + out = slice_nvfp4_output(out, output_size)
84 +
85 + if bias is not None:
86 + out = out + bias
87 + return out.view(*output_shape)
88 +
89 +
90 +class FlashInferTrtllmNvFp4LinearKernel(NvFp4LinearKernel):
91 + """NVFP4 GEMM via FlashInfer's TensorRT-LLM wrapper."""
92 +
93 + @classmethod
94 + def is_supported(
95 + cls, compute_capability: int | None = None
96 + ) -> tuple[bool, str | None]:
97 + if has_flashinfer():
98 + return True, None
99 + return False, "FlashInfer required"
100 +
101 + @classmethod
102 + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
103 + return True, None
104 +
105 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
106 + from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a
107 +
108 + weight = layer.weight.data
109 + weight_scale = layer.weight_scale.data
110 + epilogue_tile_m = 128
111 +
112 + layer.weight = torch.nn.Parameter(
113 + shuffle_matrix_a(weight.view(torch.uint8), epilogue_tile_m),
114 + requires_grad=False,
115 + )
116 + layer.weight_scale = torch.nn.Parameter(
117 + shuffle_matrix_sf_a(weight_scale.view(torch.uint8), epilogue_tile_m)
118 + .reshape(weight_scale.shape)
119 + .view(torch.float8_e4m3fn),
120 + requires_grad=False,
121 + )
122 +
123 + def apply_weights(
124 + self,
125 + layer: torch.nn.Module,
126 + x: torch.Tensor,
127 + bias: torch.Tensor | None = None,
128 + ) -> torch.Tensor:
129 + output_size = layer.output_size_per_partition
130 + output_dtype = x.dtype
131 + output_shape = [*x.shape[:-1], output_size]
132 +
133 + x_fp4, x_blockscale = scaled_fp4_quant(
134 + x,
135 + layer.input_global_scale_inv,
136 + is_sf_swizzled_layout=True,
137 + backend="flashinfer-trtllm",
138 + )
139 +
140 + out = flashinfer_scaled_fp4_mm(
141 + x_fp4,
142 + layer.weight,
143 + x_blockscale,
144 + layer.weight_scale,
145 + layer.alpha,
146 + output_dtype,
147 + backend="trtllm",
148 + )
149 +
150 + out = slice_nvfp4_output(out, output_size)
151 +
152 + if bias is not None:
153 + out = out + bias
154 + return out.view(*output_shape)
155 +
156 +
157 +class FlashInferCudnnNvFp4LinearKernel(NvFp4LinearKernel):
158 + """NVFP4 GEMM via FlashInfer's cuDNN wrapper."""
159 +
160 + @classmethod
161 + def is_supported(
162 + cls, compute_capability: int | None = None
163 + ) -> tuple[bool, str | None]:
164 + if has_flashinfer():
165 + return True, None
166 + return False, "FlashInfer required"
167 +
168 + @classmethod
169 + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
170 + return True, None
171 +
172 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
173 + # cuDNN uses the same swizzled + padded layout as CUTLASS
174 + layer.weight_scale = torch.nn.Parameter(
175 + swizzle_blockscale(layer.weight_scale.data), requires_grad=False
176 + )
177 + padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
178 + layer.weight.data
179 + )
180 + layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
181 + layer.weights_padding_cols = weights_padding_cols
182 +
183 + def apply_weights(
184 + self,
185 + layer: torch.nn.Module,
186 + x: torch.Tensor,
187 + bias: torch.Tensor | None = None,
188 + ) -> torch.Tensor:
189 + output_size = layer.output_size_per_partition
190 + output_dtype = x.dtype
191 + output_shape = [*x.shape[:-1], output_size]
192 +
193 + x_fp4, x_blockscale = scaled_fp4_quant(
194 + x,
195 + layer.input_global_scale_inv,
196 + is_sf_swizzled_layout=True,
197 + backend="flashinfer-cudnn",
198 + )
199 +
200 + x_fp4 = pad_nvfp4_activation_for_cutlass(
201 + x_fp4, getattr(layer, "weights_padding_cols", 0)
202 + )
203 +
204 + out = flashinfer_scaled_fp4_mm(
205 + x_fp4,
206 + layer.weight,
207 + x_blockscale,
208 + layer.weight_scale,
209 + layer.alpha,
210 + output_dtype,
211 + backend="cudnn",
212 + )
213 +
214 + out = slice_nvfp4_output(out, output_size)
215 +
216 + if bias is not None:
217 + out = out + bias
218 + return out.view(*output_shape)
vllm/model_executor/kernels/linear/nvfp4/marlin.py new
+57
@@ -0,0 +1,57 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import torch
5 +
6 +from vllm.logger import init_logger
7 +from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
8 + apply_fp4_marlin_linear,
9 + is_fp4_marlin_supported,
10 + prepare_fp4_layer_for_marlin,
11 +)
12 +
13 +from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
14 +
15 +logger = init_logger(__name__)
16 +
17 +
18 +class MarlinNvFp4LinearKernel(NvFp4LinearKernel):
19 + """NVFP4 weight-only GEMM via Marlin (W4A16)."""
20 +
21 + @classmethod
22 + def is_supported(
23 + cls, compute_capability: int | None = None
24 + ) -> tuple[bool, str | None]:
25 + if is_fp4_marlin_supported():
26 + return True, None
27 + return False, "Marlin FP4 not available"
28 +
29 + @classmethod
30 + def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
31 + return True, None
32 +
33 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
34 + logger.warning_once(
35 + "Your GPU does not have native support for FP4 computation but "
36 + "FP4 quantization is being used. Weight-only FP4 compression "
37 + "will be used leveraging the Marlin kernel. This may degrade "
38 + "performance for compute-heavy workloads."
39 + )
40 + prepare_fp4_layer_for_marlin(layer)
41 +
42 + def apply_weights(
43 + self,
44 + layer: torch.nn.Module,
45 + x: torch.Tensor,
46 + bias: torch.Tensor | None = None,
47 + ) -> torch.Tensor:
48 + return apply_fp4_marlin_linear(
49 + input=x,
50 + weight=layer.weight,
51 + weight_scale=layer.weight_scale,
52 + weight_global_scale=layer.weight_global_scale,
53 + workspace=layer.workspace,
54 + size_n=layer.output_size_per_partition,
55 + size_k=layer.input_size_per_partition,
56 + bias=bias,
57 + )
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w4a4_nvfp4.py
+4 -19
@@ -6,15 +6,10 @@ import torch
6 from torch.nn.parameter import Parameter
7
8 from vllm.logger import init_logger
9 +from vllm.model_executor.kernels.linear import init_nvfp4_linear_kernel
10 from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
11 CompressedTensorsScheme,
12 )
12 -from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
13 - NvFp4LinearBackend,
14 - apply_nvfp4_linear,
15 - convert_to_nvfp4_linear_kernel_format,
16 - select_nvfp4_linear_backend,
17 -)
13 from vllm.model_executor.parameter import (
14 GroupQuantScaleParameter,
15 ModelWeightParameter,
@@ -29,13 +24,9 @@ __all__ = ["CompressedTensorsW4A4Fp4"]
24
25 class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
26 def __init__(self):
32 - self.backend = select_nvfp4_linear_backend()
27 + self.kernel = init_nvfp4_linear_kernel()
28 self.group_size = 16
29
35 - self.swizzle = None
36 - if self.backend == NvFp4LinearBackend.EMULATION:
37 - self.swizzle = False
38 -
30 @classmethod
31 def get_min_capability(cls) -> int:
32 return 75
@@ -130,7 +121,7 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
121 )
122
123 # Convert layer to NVFP4 linear kernel format
133 - convert_to_nvfp4_linear_kernel_format(self.backend, layer)
124 + self.kernel.process_weights_after_loading(layer)
125
126 def apply_weights(
127 self,
@@ -138,10 +129,4 @@ class CompressedTensorsW4A4Fp4(CompressedTensorsScheme):
129 x: torch.Tensor,
130 bias: torch.Tensor | None = None,
131 ) -> torch.Tensor:
141 - return apply_nvfp4_linear(
142 - backend=self.backend,
143 - layer=layer,
144 - x=x,
145 - bias=bias,
146 - swizzle=self.swizzle,
147 - )
132 + return self.kernel.apply_weights(layer=layer, x=x, bias=bias)
vllm/model_executor/layers/quantization/modelopt.py
+7 -20
@@ -10,7 +10,10 @@ from torch.nn.parameter import Parameter
10 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
11 from vllm.config import get_current_vllm_config
12 from vllm.logger import init_logger
13 -from vllm.model_executor.kernels.linear import init_fp8_linear_kernel
13 +from vllm.model_executor.kernels.linear import (
14 + init_fp8_linear_kernel,
15 + init_nvfp4_linear_kernel,
16 +)
17 from vllm.model_executor.layers.attention import Attention, MLAAttention
18 from vllm.model_executor.layers.fused_moe.activation import MoEActivation
19 from vllm.model_executor.layers.fused_moe.config import (
@@ -70,12 +73,6 @@ from vllm.model_executor.layers.quantization.utils.mxfp8_utils import (
73 Mxfp8LinearOp,
74 mxfp8_e4m3_quantize,
75 )
73 -from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
74 - NvFp4LinearBackend,
75 - apply_nvfp4_linear,
76 - convert_to_nvfp4_linear_kernel_format,
77 - select_nvfp4_linear_backend,
78 -)
76 from vllm.model_executor.layers.quantization.utils.quant_utils import (
77 GroupShape,
78 create_fp8_quant_key,
@@ -1090,11 +1087,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1087 def __init__(self, quant_config: ModelOptNvFp4Config) -> None:
1088 self.quant_config = quant_config
1089 self.marlin_input_dtype = None
1093 - self.backend = select_nvfp4_linear_backend()
1094 -
1095 - self.swizzle = None
1096 - if self.backend == NvFp4LinearBackend.EMULATION:
1097 - self.swizzle = False
1090 + self.kernel = init_nvfp4_linear_kernel()
1091
1092 def create_weights(
1093 self,
@@ -1201,7 +1194,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1194 )
1195
1196 # Convert layer to NVFP4 linear kernel format
1204 - convert_to_nvfp4_linear_kernel_format(self.backend, layer)
1197 + self.kernel.process_weights_after_loading(layer)
1198
1199 def apply(
1200 self,
@@ -1209,13 +1202,7 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
1202 x: torch.Tensor,
1203 bias: torch.Tensor | None = None,
1204 ) -> torch.Tensor:
1212 - return apply_nvfp4_linear(
1213 - backend=self.backend,
1214 - layer=layer,
1215 - x=x,
1216 - bias=bias,
1217 - swizzle=self.swizzle,
1218 - )
1205 + return self.kernel.apply_weights(layer=layer, x=x, bias=bias)
1206
1207
1208 class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
vllm/model_executor/layers/quantization/utils/nvfp4_utils.py
-338
@@ -1,352 +1,14 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -from enum import Enum
3
4 import torch
5
7 -import vllm.envs as envs
6 from vllm._custom_ops import (
9 - cutlass_scaled_fp4_mm,
7 cutlass_scaled_mm_supports_fp4,
11 - scaled_fp4_quant,
12 -)
13 -from vllm.logger import init_logger
14 -from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (
15 - apply_fp4_marlin_linear,
16 - is_fp4_marlin_supported,
17 - prepare_fp4_layer_for_marlin,
18 -)
19 -from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
20 - kE2M1ToFloat_handle,
21 - run_nvfp4_emulations,
8 )
9 from vllm.platforms import current_platform
24 -from vllm.utils.flashinfer import flashinfer_scaled_fp4_mm, has_flashinfer
25 -from vllm.utils.import_utils import has_fbgemm_gpu
10 from vllm.utils.math_utils import round_up
11
28 -logger = init_logger(__name__)
29 -
30 -
31 -# NOTE: This is ordered by preferred backend.
32 -# Example: if both are available, FLASHINFER_CUTLASS is preferred to VLLM_CUTLASS.
33 -class NvFp4LinearBackend(Enum):
34 - FLASHINFER_CUTLASS = "flashinfer-cutlass"
35 - VLLM_CUTLASS = "cutlass"
36 - MARLIN = "marlin"
37 - FLASHINFER_TRTLLM = "flashinfer-trtllm"
38 - FLASHINFER_CUDNN = "flashinfer-cudnn"
39 - FBGEMM = "fbgemm"
40 - EMULATION = "emulation"
41 -
42 -
43 -NVFP4_LINEAR_BACKENDS = list(NvFp4LinearBackend)
44 -
45 -
46 -def is_backend_supported(backend: NvFp4LinearBackend) -> tuple[bool, str | None]:
47 - reason = None
48 - supported = True
49 -
50 - if backend == NvFp4LinearBackend.FLASHINFER_CUTLASS:
51 - # cutlass_fp4_supported() checks that the vLLM NVFP4 kernels (both
52 - # quantization and GEMM) were compiled for the current SM version.
53 - # FlashInfer backends still rely on the vLLM quantization kernels,
54 - # so we gate them on the same check.
55 - supported = (
56 - cutlass_fp4_supported()
57 - and current_platform.has_device_capability(100)
58 - and has_flashinfer()
59 - )
60 -
61 - if not supported:
62 - reason = "FlashInfer is required, >=sm_100 is required"
63 - elif backend == NvFp4LinearBackend.VLLM_CUTLASS:
64 - supported = cutlass_fp4_supported()
65 - if not supported:
66 - reason = "Cutlass is required"
67 - elif backend == NvFp4LinearBackend.MARLIN:
68 - supported = is_fp4_marlin_supported()
69 - if not supported:
70 - reason = "Marlin is required"
71 - elif backend in [
72 - NvFp4LinearBackend.FLASHINFER_TRTLLM,
73 - NvFp4LinearBackend.FLASHINFER_CUDNN,
74 - ]:
75 - supported = has_flashinfer()
76 - if not supported:
77 - reason = "FlashInfer is required"
78 - elif backend == NvFp4LinearBackend.FBGEMM:
79 - supported = has_fbgemm_gpu()
80 - if not supported:
81 - reason = "fbgemm_gpu is required"
82 - elif backend == NvFp4LinearBackend.EMULATION:
83 - # e.g. AMD Instinct does not support native NVFP4.
84 - unsupported_reasons = {}
85 - for other_backend in NVFP4_LINEAR_BACKENDS:
86 - if other_backend == NvFp4LinearBackend.EMULATION:
87 - continue
88 - other_supported, other_reason = is_backend_supported(other_backend)
89 - if not other_supported:
90 - unsupported_reasons[other_backend] = other_reason
91 -
92 - if unsupported_reasons:
93 - unsupported_reasons_str = "\n - ".join(
94 - [f"{b.value}: {r}" for b, r in unsupported_reasons.items()]
95 - )
96 - logger.warning_once(
97 - f"NVFP4 linear falling back to the slow and unoptimized "
98 - f"backend=NvFp4LinearBackend.EMULATION as no optimized backend is "
99 - f"available (unavailable reasons:\n - {unsupported_reasons_str}\n). "
100 - "In case you expect one of these backend to be used, "
101 - "please verify your environment."
102 - )
103 -
104 - return supported, reason
105 -
106 -
107 -def select_nvfp4_linear_backend() -> NvFp4LinearBackend:
108 - """
109 - Select the best available NVFP4 GEMM backend based on environment
110 - configuration and platform capabilities.
111 - """
112 - if envs.VLLM_BATCH_INVARIANT:
113 - logger.info_once(
114 - "VLLM_BATCH_INVARIANT forces NVFP4 linear to use the emulation "
115 - "backend for deterministic execution."
116 - )
117 - return NvFp4LinearBackend.EMULATION
118 -
119 - selected_backend: NvFp4LinearBackend | None = None
120 -
121 - if envs.VLLM_USE_FBGEMM:
122 - try:
123 - import fbgemm_gpu # noqa: F401
124 - except ImportError as exc:
125 - raise ImportError(
126 - "Backend fbgemm requires fbgemm.f4f4bf16 operator, "
127 - "Please install with: pip install fbgemm-gpu-genai"
128 - ) from exc
129 - selected_backend = NvFp4LinearBackend.FBGEMM
130 - elif envs.VLLM_USE_NVFP4_CT_EMULATIONS:
131 - selected_backend = NvFp4LinearBackend.EMULATION
132 - elif envs.VLLM_NVFP4_GEMM_BACKEND is None:
133 - for backend in NVFP4_LINEAR_BACKENDS:
134 - supported, reason = is_backend_supported(backend)
135 - if supported:
136 - selected_backend = backend
137 - break
138 - else:
139 - selected_backend = NvFp4LinearBackend(envs.VLLM_NVFP4_GEMM_BACKEND)
140 -
141 - if selected_backend is None:
142 - raise ValueError(
143 - f"No NVFP4 GEMM backend selected, "
144 - f"available backends: {NVFP4_LINEAR_BACKENDS}"
145 - )
146 -
147 - supported, reason = is_backend_supported(selected_backend)
148 -
149 - if not supported:
150 - raise ValueError(
151 - f"The selected backend={selected_backend} is not supported in current "
152 - f"environment. Reason: {reason}. Current environment: "
153 - f"{envs.VLLM_USE_FBGEMM=}, {envs.VLLM_USE_NVFP4_CT_EMULATIONS=}, "
154 - f"{envs.VLLM_NVFP4_GEMM_BACKEND}."
155 - )
156 -
157 - logger.info_once(f"Using {selected_backend} for NVFP4 GEMM")
158 - return selected_backend
159 -
160 -
161 -def prepare_weights_for_nvfp4_flashinfer_trtllm(
162 - weight: torch.Tensor,
163 - weight_scale: torch.Tensor,
164 -) -> tuple[torch.Tensor, torch.Tensor]:
165 - """Prepare weights and scales for FlashInfer TRTLLM FP4 GEMM."""
166 - from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a
167 -
168 - epilogue_tile_m = 128
169 - shuffled_weight = shuffle_matrix_a(weight.view(torch.uint8), epilogue_tile_m)
170 - shuffled_weight_scale = (
171 - shuffle_matrix_sf_a(weight_scale.view(torch.uint8), epilogue_tile_m)
172 - .reshape(weight_scale.shape)
173 - .view(torch.float8_e4m3fn)
174 - )
175 -
176 - return shuffled_weight, shuffled_weight_scale
177 -
178 -
179 -def prepare_weights_for_nvfp4_cutlass(
180 - weight: torch.Tensor,
181 - weight_scale: torch.Tensor,
182 -) -> tuple[torch.Tensor, torch.Tensor, int]:
183 - """
184 - Prepare weights and scales for CUTLASS/FlashInfer-CUTLASS FP4 GEMM.
185 - This involves padding weights for alignment (K and N divisible by 32)
186 - """
187 - swizzled_weight_scale = swizzle_blockscale(weight_scale)
188 - padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(weight)
189 - return padded_weight, swizzled_weight_scale, weights_padding_cols
190 -
191 -
192 -def prepare_weights_for_nvfp4_fbgemm(
193 - weight: torch.Tensor,
194 - weight_scale: torch.Tensor,
195 -) -> tuple[torch.Tensor, torch.Tensor]:
196 - """Prepare weights and scales for FBGEMM FP4 GEMM."""
197 - swizzled_weight_scale = swizzle_blockscale(weight_scale)
198 - swizzled_weight_scale = swizzled_weight_scale.view(-1).view(torch.uint8)
199 - return weight, swizzled_weight_scale
200 -
201 -
202 -def convert_to_nvfp4_linear_kernel_format(
203 - backend: NvFp4LinearBackend,
204 - layer: torch.nn.Module,
205 -) -> None:
206 - """Convert layer to NVFP4 linear kernel format."""
207 -
208 - assert layer.weight_scale.dtype == torch.float8_e4m3fn, (
209 - "Weight Block scale must be represented as FP8-E4M3"
210 - )
211 -
212 - # Default to no padding
213 - layer.weights_padding_cols = 0
214 -
215 - if backend == NvFp4LinearBackend.MARLIN:
216 - logger.warning_once(
217 - "Your GPU does not have native support for FP4 computation but "
218 - "FP4 quantization is being used. Weight-only FP4 compression "
219 - "will be used leveraging the Marlin kernel. This may degrade "
220 - "performance for compute-heavy workloads."
221 - )
222 - prepare_fp4_layer_for_marlin(layer)
223 - elif backend == NvFp4LinearBackend.FLASHINFER_TRTLLM:
224 - weight, weight_scale = prepare_weights_for_nvfp4_flashinfer_trtllm(
225 - layer.weight.data, layer.weight_scale.data
226 - )
227 - layer.weight = torch.nn.Parameter(weight, requires_grad=False)
228 - layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
229 - elif backend == NvFp4LinearBackend.FBGEMM:
230 - weight, weight_scale = prepare_weights_for_nvfp4_fbgemm(
231 - layer.weight.data, layer.weight_scale.data
232 - )
233 - layer.weight = torch.nn.Parameter(weight, requires_grad=False)
234 - layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
235 - elif backend in (
236 - NvFp4LinearBackend.VLLM_CUTLASS,
237 - NvFp4LinearBackend.FLASHINFER_CUTLASS,
238 - NvFp4LinearBackend.FLASHINFER_CUDNN,
239 - ):
240 - weight, weight_scale, weights_padding_cols = prepare_weights_for_nvfp4_cutlass(
241 - layer.weight.data, layer.weight_scale.data
242 - )
243 - layer.weight = torch.nn.Parameter(weight, requires_grad=False)
244 - layer.weight_scale = torch.nn.Parameter(weight_scale, requires_grad=False)
245 - layer.weights_padding_cols = weights_padding_cols
246 - elif backend == NvFp4LinearBackend.EMULATION:
247 - # We can not call `.to(device)` during cuda graph capture - do it here instead.
248 - # (operation not permitted when stream is capturing)
249 - kE2M1ToFloat_handle.val = kE2M1ToFloat_handle.val.to(layer.weight.device)
250 -
251 -
252 -def apply_nvfp4_linear(
253 - backend: NvFp4LinearBackend,
254 - layer: torch.nn.Module,
255 - x: torch.Tensor,
256 - bias: torch.Tensor | None = None,
257 - swizzle: bool | None = None,
258 -) -> torch.Tensor:
259 - """
260 - Apply NVFP4 linear transformation using the specified backend.
261 - """
262 - weight = layer.weight
263 - weight_scale = layer.weight_scale
264 - weight_global_scale = layer.weight_global_scale
265 - input_global_scale_inv = layer.input_global_scale_inv
266 - alpha = layer.alpha
267 - output_size = layer.output_size_per_partition
268 - input_size = layer.input_size_per_partition
269 - output_dtype = x.dtype
270 - output_shape = [*x.shape[:-1], output_size]
271 -
272 - if backend == NvFp4LinearBackend.MARLIN:
273 - return apply_fp4_marlin_linear(
274 - input=x,
275 - weight=weight,
276 - weight_scale=weight_scale,
277 - weight_global_scale=weight_global_scale,
278 - workspace=layer.workspace,
279 - size_n=output_size,
280 - size_k=input_size,
281 - bias=bias,
282 - )
283 - elif backend == NvFp4LinearBackend.EMULATION:
284 - x_2d = x.reshape(-1, x.shape[-1])
285 - out = run_nvfp4_emulations(
286 - x=x_2d,
287 - input_global_scale=input_global_scale_inv,
288 - weight=weight,
289 - weight_scale_swizzled=weight_scale,
290 - weight_global_scale=weight_global_scale,
291 - swizzle=swizzle,
292 - )
293 - out = out[:, :output_size]
294 - if bias is not None:
295 - out = out + bias
296 - return out.view(*output_shape)
297 -
298 - # Quantize BF16 or FP16 to (FP4 and interleaved block scale)
299 - x_fp4, x_blockscale = scaled_fp4_quant(
300 - x, input_global_scale_inv, is_sf_swizzled_layout=True, backend=backend.value
301 - )
302 -
303 - # Validate dtypes
304 - assert x_fp4.dtype == torch.uint8
305 - assert weight.dtype == torch.uint8
306 - assert x_blockscale.dtype == torch.float8_e4m3fn
307 - # weight_scale is fp8 for most backends, but uint8 for fbgemm
308 - assert weight_scale.dtype in (torch.float8_e4m3fn, torch.uint8)
309 - assert alpha.dtype == torch.float32
310 -
311 - # Pad activations to match weight K-dimension padding
312 - weights_padding_cols = getattr(layer, "weights_padding_cols", 0)
313 - x_fp4 = pad_nvfp4_activation_for_cutlass(x_fp4, weights_padding_cols)
314 -
315 - # Prepare args for the matmul
316 - mm_args = (
317 - x_fp4,
318 - weight,
319 - x_blockscale,
320 - weight_scale,
321 - alpha,
322 - output_dtype,
323 - )
324 -
325 - # Call the appropriate backend
326 - if backend.value.startswith("flashinfer-"):
327 - backend_name = backend.value[len("flashinfer-") :]
328 - out = flashinfer_scaled_fp4_mm(*mm_args, backend=backend_name)
329 - elif backend == NvFp4LinearBackend.FBGEMM:
330 - out = torch.ops.fbgemm.f4f4bf16(
331 - x_fp4,
332 - weight,
333 - x_blockscale.view(-1).view(torch.uint8),
334 - weight_scale,
335 - alpha,
336 - use_mx=False,
337 - ).to(output_dtype)
338 - else:
339 - assert backend == NvFp4LinearBackend.VLLM_CUTLASS
340 - out = cutlass_scaled_fp4_mm(*mm_args)
341 -
342 - # Slice output to remove N-dimension padding
343 - out = slice_nvfp4_output(out, output_size)
344 -
345 - if bias is not None:
346 - out = out + bias
347 -
348 - return out.view(*output_shape)
349 -
12
13 def swizzle_blockscale(scale: torch.Tensor) -> torch.Tensor:
14 """