@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 203 lines 7.7 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import torch
5 from torch.nn import Module
6
7 from vllm._custom_ops import scaled_fp4_quant
8 from vllm.model_executor.layers.fused_moe import RoutedExperts
9 from vllm.model_executor.layers.fused_moe.config import (
10 FusedMoEConfig,
11 FusedMoEQuantConfig,
12 )
13 from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (
14 convert_to_nvfp4_moe_kernel_format,
15 make_nvfp4_moe_kernel,
16 make_nvfp4_moe_quant_config,
17 select_nvfp4_moe_backend,
18 )
19 from vllm.model_executor.layers.quantization.online.moe_base import (
20 OnlineMoEMethodBase,
21 )
22 from vllm.model_executor.layers.quantization.utils.nvfp4_emulation_utils import (
23 FLOAT4_E2M1_MAX,
24 )
25 from vllm.model_executor.layers.quantization.utils.quant_utils import (
26 amax_for_moe_weight_quant,
27 kNvfp4DynamicToken,
28 kNvfp4Static,
29 weight_amax,
30 )
31 from vllm.model_executor.utils import replace_parameter
32 from vllm.platforms import current_platform
33
34 FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
35
36
37 def _quantize_moe_weight_to_nvfp4(
38 weight: torch.Tensor,
39 moe_tp_size: int = 1,
40 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
41 """Quantize stacked MoE expert weights ``(E, N, K)`` to NVFP4.
42
43 One FP32 global scale per expert plus per-block (group-16) FP8 scales,
44 matching the ModelOpt NVFP4 checkpoint layout. Returns packed FP4 weights
45 ``(E, N, K // 2)``, block scales ``(E, N, K // 16)``, and the per-expert
46 global scale ``(E,)`` stored as ``amax / (fp4_max * fp8_max)``.
47 """
48 assert weight.dim() == 3, f"expected 3D expert weights, got {weight.shape}"
49 k = weight.shape[-1]
50 assert k % 16 == 0, f"last dim must be a multiple of 16, got {k}"
51
52 amax = weight_amax(weight.flatten(1), dim=-1).to(torch.float32)
53 amax = amax_for_moe_weight_quant(amax, moe_tp_size).clamp_min(1e-8)
54 global_scale = (FLOAT4_E2M1_MAX * FLOAT8_E4M3_MAX) / amax
55 weight_scale_2 = (1.0 / global_scale).to(torch.float32)
56
57 # Keep the original BF16/FP16 values as the quantizer input. Folding each
58 # expert's FP32 global scale into the weight would add a BF16/FP16 rounding
59 # before the group-16 scale and E2M1 values are selected.
60 weight = weight.contiguous()
61 quantized_experts = [
62 scaled_fp4_quant(
63 expert_weight,
64 expert_scale,
65 is_sf_swizzled_layout=False,
66 )
67 for expert_weight, expert_scale in zip(
68 weight,
69 global_scale,
70 strict=True,
71 )
72 ]
73 qweight = torch.stack([quantized for quantized, _ in quantized_experts])
74 block_scale = torch.stack([block_scale for _, block_scale in quantized_experts])
75 return (
76 qweight,
77 block_scale,
78 weight_scale_2,
79 )
80
81
82 class Nvfp4OnlineMoEMethod(OnlineMoEMethodBase):
83 """Online NVFP4 MoE quantization with per-token activation scales.
84
85 Quantizes fp16/bf16 expert weights to NVFP4 at load time; the FlashInfer
86 TRTLLM kernel computes per-token activation scales at runtime. Blackwell
87 (SM100) only.
88 """
89
90 def __init__(
91 self,
92 *,
93 moe: FusedMoEConfig,
94 ):
95 if not current_platform.is_device_capability_family(100):
96 raise ValueError(
97 "nvfp4_per_token online quantization requires a Blackwell (SM100) GPU."
98 )
99 super().__init__(moe)
100 self.nvfp4_backend, self.experts_cls = select_nvfp4_moe_backend(
101 config=self.moe,
102 weight_key=kNvfp4Static,
103 activation_key=kNvfp4DynamicToken,
104 )
105
106 def process_weights_after_loading(self, layer: Module) -> None:
107 if getattr(layer, "_already_called_process_weights_after_loading", False):
108 return
109
110 self._quantize_weights(layer)
111 self._setup_kernel(layer)
112
113 layer._already_called_process_weights_after_loading = True
114
115 def _quantize_weights(self, layer: Module) -> None:
116 moe_tp_size = self.moe.tp_size
117 w13, w13_scale, w13_scale_2 = _quantize_moe_weight_to_nvfp4(
118 layer.w13_weight, moe_tp_size
119 )
120 w2, w2_scale, w2_scale_2 = _quantize_moe_weight_to_nvfp4(
121 layer.w2_weight, moe_tp_size
122 )
123
124 replace_parameter(layer, "w13_weight", w13)
125 replace_parameter(layer, "w13_weight_scale", w13_scale)
126 replace_parameter(layer, "w13_weight_scale_2", w13_scale_2)
127 replace_parameter(layer, "w2_weight", w2)
128 replace_parameter(layer, "w2_weight_scale", w2_scale)
129 replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)
130
131 # Neutral (1.0) activation global scales: the kernel derives per-token
132 # scales at runtime, so the output scalars reduce to the weight scales.
133 ones = torch.ones(layer.num_experts, device=w13.device, dtype=torch.float32)
134 replace_parameter(layer, "w13_input_scale", ones)
135 replace_parameter(layer, "w2_input_scale", ones.clone())
136
137 def _setup_kernel(self, layer: RoutedExperts) -> None:
138 (
139 w13,
140 w13_scale,
141 w13_scale_2,
142 a13_scale,
143 w2,
144 w2_scale,
145 w2_scale_2,
146 a2_scale,
147 ) = convert_to_nvfp4_moe_kernel_format(
148 nvfp4_backend=self.nvfp4_backend,
149 layer=layer,
150 w13=layer.w13_weight,
151 w13_scale=layer.w13_weight_scale,
152 w13_scale_2=layer.w13_weight_scale_2,
153 a13_scale=layer.w13_input_scale,
154 w2=layer.w2_weight,
155 w2_scale=layer.w2_weight_scale,
156 w2_scale_2=layer.w2_weight_scale_2,
157 a2_scale=layer.w2_input_scale,
158 is_act_and_mul=self.moe.is_act_and_mul,
159 )
160
161 replace_parameter(layer, "w13_weight", w13)
162 replace_parameter(layer, "w13_weight_scale", w13_scale)
163 replace_parameter(layer, "w13_weight_scale_2", w13_scale_2)
164 replace_parameter(layer, "w13_input_scale", a13_scale)
165 replace_parameter(layer, "w2_weight", w2)
166 replace_parameter(layer, "w2_weight_scale", w2_scale)
167 replace_parameter(layer, "w2_weight_scale_2", w2_scale_2)
168 replace_parameter(layer, "w2_input_scale", a2_scale)
169
170 if self.moe_kernel is None:
171 self.moe_quant_config = self.get_fused_moe_quant_config(layer)
172 assert self.experts_cls is not None
173 self.moe_kernel = make_nvfp4_moe_kernel(
174 moe_quant_config=self.moe_quant_config,
175 moe_config=self.moe,
176 experts_cls=self.experts_cls,
177 backend=self.nvfp4_backend,
178 routing_tables=layer._expert_routing_tables(),
179 per_token_activation=True,
180 )
181 else:
182 # Reload creates new scale tensors; derived kernel scales must use
183 # their new values before layerwise reload restores captured storage.
184 assert self.moe_quant_config is not None
185 assert self.moe_quant_config.g1_alphas is not None
186 assert self.moe_quant_config.g2_alphas is not None
187 self.moe_quant_config.g1_alphas.copy_(w13_scale_2)
188 self.moe_quant_config.g2_alphas.copy_(w2_scale_2)
189
190 self.moe_kernel.fused_experts.process_weights_after_loading(layer)
191
192 def get_fused_moe_quant_config(self, layer: torch.nn.Module) -> FusedMoEQuantConfig:
193 return make_nvfp4_moe_quant_config(
194 backend=self.nvfp4_backend,
195 w13_scale=layer.w13_weight_scale,
196 w2_scale=layer.w2_weight_scale,
197 w13_scale_2=layer.w13_weight_scale_2,
198 w2_scale_2=layer.w2_weight_scale_2,
199 a13_scale=layer.w13_input_scale,
200 a2_scale=layer.w2_input_scale,
201 swiglu_limit=getattr(layer, "swiglu_limit", None),
202 layer=layer,
203 )