@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 163 lines 5.74 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Callable
5
6 import torch
7 from torch.nn.parameter import Parameter
8
9 from vllm.logger import init_logger
10 from vllm.model_executor.kernels.linear import init_nvfp4_linear_kernel
11 from vllm.model_executor.kernels.linear.nvfp4.emulation import (
12 EmulationNvFp4LinearKernel,
13 )
14 from vllm.model_executor.layers.quantization.quark.schemes.quark_scheme import (
15 QuarkScheme,
16 )
17 from vllm.model_executor.layers.quantization.utils.quant_utils import (
18 QuantKey,
19 kNvfp4Dynamic,
20 kNvfp4Static,
21 )
22 from vllm.model_executor.parameter import (
23 GroupQuantScaleParameter,
24 ModelWeightParameter,
25 PerTensorScaleParameter,
26 )
27
28 __all__ = ["QuarkNVFP4"]
29
30 logger = init_logger(__name__)
31
32
33 class QuarkNVFP4(QuarkScheme):
34 """Quark NVFP4 quantization scheme.
35
36 Supports loading NVFP4 checkpoints with the following structure:
37 - weight: uint8, shape [out_features, in_features // 2] (packed FP4)
38 - weight_scale: float8_e4m3fn, shape [out_features, in_features // group_size]
39 - weight_scale_2: bfloat16/float32, scalar (global weight scale)
40 - input_scale_2: bfloat16/float32, scalar (global input scale)
41 """
42
43 supported_activation_quant_keys = [kNvfp4Dynamic]
44 supported_weight_quant_keys = [kNvfp4Static]
45
46 def __init__(
47 self,
48 activation_quant_key: QuantKey | None,
49 ):
50 super().__init__(kNvfp4Static, activation_quant_key)
51 self.kernel = init_nvfp4_linear_kernel()
52 self.group_size = 16
53
54 if not isinstance(self.kernel, EmulationNvFp4LinearKernel):
55 logger.warning_once(
56 "Only EmulationNvFp4LinearKernel NVFP4 dense implementation is "
57 "tested with QuarkNVFP4, got kernel=%s. Correctness is not validated.",
58 type(self.kernel).__name__,
59 )
60
61 @classmethod
62 def get_min_capability(cls) -> int:
63 # FP4 requires Turing (75) or newer
64 return 75
65
66 def create_weights(
67 self,
68 layer: torch.nn.Module,
69 output_partition_sizes: list[int],
70 input_size_per_partition: int,
71 params_dtype: torch.dtype,
72 weight_loader: Callable,
73 **kwargs,
74 ):
75 output_size_per_partition = sum(output_partition_sizes)
76 layer.logical_widths = output_partition_sizes
77 layer.input_size_per_partition = input_size_per_partition
78 layer.output_size_per_partition = output_size_per_partition
79
80 if input_size_per_partition % self.group_size != 0:
81 raise ValueError(
82 f"Input size per partition ({input_size_per_partition}) must be "
83 f"divisible by group size ({self.group_size})"
84 )
85
86 # Weight: FP4 packed as uint8 (2 FP4 values per uint8)
87 weight = ModelWeightParameter(
88 data=torch.empty(
89 output_size_per_partition,
90 input_size_per_partition // 2,
91 dtype=torch.uint8,
92 ),
93 input_dim=1,
94 output_dim=0,
95 weight_loader=weight_loader,
96 )
97 layer.register_parameter("weight", weight)
98
99 # Per-group weight scale (FP8 E4M3)
100 weight_scale = GroupQuantScaleParameter(
101 data=torch.empty(
102 output_size_per_partition,
103 input_size_per_partition // self.group_size,
104 dtype=torch.float8_e4m3fn,
105 ),
106 input_dim=1,
107 output_dim=0,
108 weight_loader=weight_loader,
109 )
110 layer.register_parameter("weight_scale", weight_scale)
111
112 # Global weight scale (scalar, per partition)
113 weight_scale_2 = PerTensorScaleParameter(
114 data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
115 weight_loader=weight_loader,
116 )
117 layer.register_parameter("weight_scale_2", weight_scale_2)
118
119 # Global input scale (scalar, per partition)
120 input_scale_2 = PerTensorScaleParameter(
121 data=torch.empty(len(output_partition_sizes), dtype=torch.float32),
122 weight_loader=weight_loader,
123 )
124 layer.register_parameter("input_scale_2", input_scale_2)
125
126 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
127 input_global_scale = layer.input_scale_2.max().to(torch.float32)
128 layer.input_global_scale = Parameter(input_global_scale, requires_grad=False)
129 del layer.input_scale_2
130
131 weight_global_scale = layer.weight_scale_2.to(torch.float32)
132
133 if torch.unique(weight_global_scale).numel() != 1:
134 logger.warning_once(
135 "In NVFP4 linear, the global scale for weight are different"
136 " for parallel layers (e.g. q_proj, k_proj, v_proj). This"
137 " will likely result in reduced accuracy. Please verify the"
138 " model accuracy. Consider using a checkpoint with a shared"
139 " global NVFP4 scale for fused layers."
140 )
141
142 weight_global_scale = weight_global_scale.max()
143
144 layer.weight_global_scale = Parameter(weight_global_scale, requires_grad=False)
145 del layer.weight_scale_2
146
147 layer.alpha = Parameter(
148 layer.input_global_scale * layer.weight_global_scale, requires_grad=False
149 )
150 layer.input_global_scale_inv = Parameter(
151 (1.0 / layer.input_global_scale).to(torch.float32), requires_grad=False
152 )
153
154 # Convert layer to NVFP4 linear kernel format
155 self.kernel.process_weights_after_loading(layer)
156
157 def apply_weights(
158 self,
159 layer: torch.nn.Module,
160 x: torch.Tensor,
161 bias: torch.Tensor | None = None,
162 ) -> torch.Tensor:
163 return self.kernel.apply_weights(layer=layer, x=x, bias=bias)