wehub-resource-sync
78 lines
2.3 KiB
Python
78 lines
2.3 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
import torch
|
|
|
|
from vllm._custom_ops import (
|
|
cutlass_scaled_fp4_mm,
|
|
scaled_fp4_quant,
|
|
)
|
|
from vllm.model_executor.layers.quantization.utils.nvfp4_utils import (
|
|
cutlass_fp4_supported,
|
|
pad_nvfp4_weight_for_cutlass,
|
|
slice_nvfp4_output,
|
|
swizzle_blockscale,
|
|
)
|
|
|
|
from .base import NvFp4LinearKernel, NvFp4LinearLayerConfig
|
|
|
|
|
|
class CutlassNvFp4LinearKernel(NvFp4LinearKernel):
|
|
"""NVFP4 GEMM via the vLLM CUTLASS kernel."""
|
|
|
|
@classmethod
|
|
def is_supported(
|
|
cls, compute_capability: int | None = None
|
|
) -> tuple[bool, str | None]:
|
|
if cutlass_fp4_supported():
|
|
return True, None
|
|
return False, "CUTLASS FP4 kernels not available"
|
|
|
|
@classmethod
|
|
def can_implement(cls, config: NvFp4LinearLayerConfig) -> tuple[bool, str | None]:
|
|
return True, None
|
|
|
|
def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
|
|
layer.weight_scale = torch.nn.Parameter(
|
|
swizzle_blockscale(layer.weight_scale.data), requires_grad=False
|
|
)
|
|
padded_weight, weights_padding_cols = pad_nvfp4_weight_for_cutlass(
|
|
layer.weight.data
|
|
)
|
|
layer.weight = torch.nn.Parameter(padded_weight, requires_grad=False)
|
|
layer.weights_padding_cols = weights_padding_cols
|
|
|
|
def apply_weights(
|
|
self,
|
|
layer: torch.nn.Module,
|
|
x: torch.Tensor,
|
|
bias: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
output_size = layer.output_size_per_partition
|
|
output_dtype = x.dtype
|
|
output_shape = [*x.shape[:-1], output_size]
|
|
weights_padding_bytes = getattr(layer, "weights_padding_cols", 0)
|
|
|
|
x_fp4, x_blockscale = scaled_fp4_quant(
|
|
x,
|
|
layer.input_global_scale_inv,
|
|
is_sf_swizzled_layout=True,
|
|
backend="cutlass",
|
|
padded_n=x.shape[-1] + weights_padding_bytes * 2,
|
|
)
|
|
|
|
out = cutlass_scaled_fp4_mm(
|
|
x_fp4,
|
|
layer.weight,
|
|
x_blockscale,
|
|
layer.weight_scale,
|
|
layer.alpha,
|
|
output_dtype,
|
|
)
|
|
|
|
out = slice_nvfp4_output(out, output_size)
|
|
|
|
if bias is not None:
|
|
out = out + bias
|
|
return out.view(*output_shape)
|