@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 235 lines 8.41 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.parameter import Parameter
6
7 import vllm.envs as envs
8 from vllm._aiter_ops import is_aiter_found_and_supported, rocm_aiter_ops
9 from vllm.logger import init_logger
10 from vllm.model_executor.layers.quantization.utils.quant_utils import (
11 kMxfp4Dynamic,
12 )
13 from vllm.platforms import current_platform
14
15 from .base import MxFp4LinearKernel, MxFp4LinearLayerConfig
16
17 logger = init_logger(__name__)
18
19
20 _ASM_FP4_SCALE_ROW_MULTIPLE = 32
21 _ASM_FP4_SCALE_COL_MULTIPLE = 8
22
23
24 def _asm_fp4_scale_swizzle_supported(weight_scale: torch.Tensor) -> bool:
25 # The ASM swizzle reshapes weight_scale into a fixed tile layout, requiring
26 # scale rows divisible by 32 and scale columns divisible by 8
27 if weight_scale.ndim != 2:
28 return False
29 sm, sn = weight_scale.shape
30 return (
31 sm % _ASM_FP4_SCALE_ROW_MULTIPLE == 0 and sn % _ASM_FP4_SCALE_COL_MULTIPLE == 0
32 )
33
34
35 # NOTE: Do not import aiter at module scope. Importing aiter eagerly initializes HIP
36 # which can force the engine core to spawn instead of fork.
37 # is_aiter_found_and_supported() checks platform + arch + library availability via
38 # find_spec/amdsmi, so it stays HIP-free.
39 # Actual aiter imports are deferred to the functions/methods that need them,
40 # where HIP initialization is expected.
41 if is_aiter_found_and_supported():
42 from vllm.utils.torch_utils import direct_register_custom_op
43
44 def gemm_with_dynamic_quant(
45 x: torch.Tensor,
46 weight: torch.Tensor,
47 weight_scale: torch.Tensor,
48 rocm_use_aiter_fp4_asm_gemm: bool = False,
49 out_dtype: torch.dtype | None = torch.bfloat16,
50 x_scales: torch.Tensor | None = None,
51 ) -> torch.Tensor:
52 from aiter.ops.triton.gemm_afp4wfp4 import (
53 gemm_afp4wfp4,
54 gemm_afp4wfp4_preshuffle,
55 )
56 from aiter.ops.triton.quant import dynamic_mxfp4_quant
57
58 if rocm_use_aiter_fp4_asm_gemm:
59 from aiter import gemm_a4w4, per_1x32_f4_quant_hip
60
61 M = x.shape[0]
62 N = weight.shape[0]
63 K = weight.shape[1]
64 if rocm_use_aiter_fp4_asm_gemm:
65 if M <= 64 and rocm_aiter_ops.is_triton_gemm_afp4wfp4_presh_ws_tuned(N, K):
66 if x_scales is None:
67 # use hip quant kernel for performance
68 if M >= 32:
69 x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True)
70 else:
71 x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=False)
72 else:
73 x_q = x
74 x_s = x_scales
75
76 if M >= 32:
77 x_s = x_s.view(torch.uint8).view(x_s.shape[0] // 32, -1)
78 else:
79 x_s = x_s[:M, ...].view(torch.uint8)
80
81 y = torch.empty(M, N, device=x_q.device, dtype=out_dtype)
82 gemm_afp4wfp4_preshuffle(
83 x_q.view(torch.uint8),
84 weight.view(torch.uint8).view(weight.shape[0] // 16, -1),
85 x_s,
86 weight_scale.view(torch.uint8).view(
87 weight_scale.shape[0] // 32, -1
88 ),
89 out_dtype,
90 y,
91 )
92 else:
93 if x_scales is None:
94 # use hip quant kernel for performance
95 x_q, x_s = per_1x32_f4_quant_hip(x, shuffle=True)
96 else:
97 x_q = x
98 x_s = x_scales
99
100 y = gemm_a4w4(
101 x_q,
102 weight.view(x_q.dtype),
103 x_s,
104 weight_scale.view(x_s.dtype),
105 dtype=out_dtype,
106 bpreshuffle=True,
107 )
108 return y[:M]
109 else:
110 if x_scales is None:
111 x_q, x_s = dynamic_mxfp4_quant(x)
112 else:
113 x_q = x
114 x_s = x_scales
115 y = torch.empty(
116 x_q.shape[0], weight.shape[0], device=x_q.device, dtype=out_dtype
117 )
118
119 gemm_afp4wfp4(x_q, weight, x_s, weight_scale.T, out_dtype, y)
120 return y
121
122 def gemm_with_dynamic_quant_fake(
123 x: torch.Tensor,
124 weight: torch.Tensor,
125 weight_scale: torch.Tensor,
126 x_scales: torch.Tensor = None,
127 rocm_use_aiter_fp4_asm_gemm: bool = False,
128 out_dtype: torch.dtype | None = torch.bfloat16,
129 ) -> torch.Tensor:
130 return torch.empty(
131 (*x.shape[:-1], weight.shape[0]), dtype=out_dtype, device=x.device
132 )
133
134 direct_register_custom_op(
135 op_name="gemm_with_dynamic_quant",
136 op_func=gemm_with_dynamic_quant,
137 mutates_args=[],
138 fake_impl=gemm_with_dynamic_quant_fake,
139 dispatch_key=current_platform.dispatch_key,
140 )
141
142
143 class AiterMxfp4LinearKernel(MxFp4LinearKernel):
144 """AITER-based native MXFP4 GEMM kernel for ROCm."""
145
146 def __init__(self, config: MxFp4LinearLayerConfig) -> None:
147 super().__init__(config)
148 self.use_asm_gemm = rocm_aiter_ops.is_asm_fp4_gemm_dynamic_quant_enabled()
149 self.out_dtype = torch.get_default_dtype()
150
151 @classmethod
152 def is_supported(
153 cls, compute_capability: int | None = None
154 ) -> tuple[bool, str | None]:
155 if not current_platform.supports_mx():
156 return False, "current platform does not support native MXFP4 computation"
157
158 from vllm._aiter_ops import is_aiter_found_and_supported
159 from vllm.model_executor.kernels.linear import _get_linear_backend
160
161 linear_backend = _get_linear_backend(quantization="mxfp4")
162
163 if (
164 current_platform.is_rocm()
165 and current_platform.supports_mx()
166 and "AiterMxfp4LinearKernel" not in envs.VLLM_DISABLED_KERNELS
167 and linear_backend == "auto"
168 and not is_aiter_found_and_supported()
169 ):
170 logger.warning_once(
171 "This platform supports native MXFP4 W4A4 MOE "
172 "computation via AITER MOE backend, but AITER is not "
173 "found or not supported. Consider installing AITER: "
174 "https://github.com/ROCm/aiter."
175 )
176
177 if is_aiter_found_and_supported():
178 return True, None
179 return False, "AITER not found or not supported on the current platform"
180
181 @classmethod
182 def can_implement(cls, config: MxFp4LinearLayerConfig) -> tuple[bool, str | None]:
183 if config.activation_quant_key != kMxfp4Dynamic:
184 return False, "only supports MXFP4 dynamic activation"
185 return True, None
186
187 def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
188 if self.use_asm_gemm and not _asm_fp4_scale_swizzle_supported(
189 layer.weight_scale.data
190 ):
191 logger.warning_once(
192 "AITER ASM FP4 GEMM requires weight_scale dims divisible by "
193 "(%d, %d), but this layer has weight_scale shape %s. Falling "
194 "back to the AITER Triton FP4 GEMM for this layer.",
195 _ASM_FP4_SCALE_ROW_MULTIPLE,
196 _ASM_FP4_SCALE_COL_MULTIPLE,
197 tuple(layer.weight_scale.data.shape),
198 )
199 self.use_asm_gemm = False
200
201 if self.use_asm_gemm:
202 from aiter.ops.shuffle import shuffle_weight
203
204 weight_scale = layer.weight_scale.data
205 sm, sn = weight_scale.shape
206 weight_scale = weight_scale.view(sm // 32, 2, 16, sn // 8, 2, 4, 1)
207 weight_scale = weight_scale.permute(0, 3, 5, 2, 4, 1, 6).contiguous()
208 weight_scale = weight_scale.view(sm, sn)
209 layer.weight_scale = Parameter(weight_scale, requires_grad=False)
210
211 layer.weight = Parameter(
212 shuffle_weight(layer.weight.data, layout=(16, 16)),
213 requires_grad=False,
214 )
215 else:
216 layer.weight_scale = Parameter(
217 layer.weight_scale.data.T.contiguous(), requires_grad=False
218 )
219
220 def apply_weights(
221 self,
222 layer: torch.nn.Module,
223 x: torch.Tensor,
224 bias: torch.Tensor | None = None,
225 ) -> torch.Tensor:
226 y = torch.ops.vllm.gemm_with_dynamic_quant(
227 x,
228 layer.weight,
229 layer.weight_scale,
230 self.use_asm_gemm,
231 self.out_dtype,
232 )
233 if bias is not None:
234 y = y + bias
235 return y