@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 163 lines 6.17 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """W4A16 GPTQ kernel for AMD RDNA3 (gfx1100) — fp16 + bf16.
4
5 Drop-in replacement for ExllamaLinearKernel on RDNA3 that adds native bf16
6 support. The HIP kernel lives in ``csrc/rocm/q_gemm_rdna3.cu``
7 and is exposed via ``torch.ops._rocm_C.gptq_gemm_rdna3``.
8
9 Registered ahead of TritonW4A16LinearKernel for the ROCm-RDNA3 path; falls
10 through to the Triton kernel on non-RDNA3 ROCm devices (e.g. CDNA/MI300).
11 """
12
13 import torch
14
15 from vllm import _custom_ops as ops
16 from vllm.model_executor.layers.quantization.utils.quant_utils import (
17 pack_quantized_values_into_int32,
18 )
19 from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_
20 from vllm.platforms import current_platform
21 from vllm.scalar_type import scalar_types
22
23 from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig
24
25
26 class RDNA3W4A16LinearKernel(MPLinearKernel):
27 SUPPORTED_QUANT_TYPES = [scalar_types.uint4b8]
28
29 @classmethod
30 def get_min_capability(cls) -> int:
31 # ROCm gates via on_gfx1100() in can_implement.
32 return 60
33
34 @classmethod
35 def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]:
36 if not current_platform.is_rocm():
37 return False, "RDNA3 W4A16 kernel is ROCm-only"
38
39 from vllm.platforms.rocm import on_gfx1100
40
41 if not on_gfx1100():
42 return False, "RDNA3 W4A16 kernel requires gfx1100"
43
44 # The HIP op is registered by the C++ extension; if a user is running
45 # against a vLLM build that doesn't include it (e.g. partial rebuild),
46 # fall through gracefully to the next kernel in the registry.
47 if not (
48 hasattr(torch.ops, "_rocm_C")
49 and hasattr(torch.ops._rocm_C, "gptq_gemm_rdna3")
50 ):
51 return (
52 False,
53 "torch.ops._rocm_C.gptq_gemm_rdna3 missing — rebuild C++ extension",
54 )
55
56 if c.act_type not in (torch.float16, torch.bfloat16):
57 return False, "RDNA3 W4A16 kernel only supports fp16 and bf16"
58
59 if c.weight_type not in cls.SUPPORTED_QUANT_TYPES:
60 return (
61 False,
62 f"Quant type ({c.weight_type}) not supported by "
63 f"RDNA3 W4A16 kernel; supported: {cls.SUPPORTED_QUANT_TYPES}",
64 )
65
66 if c.group_size <= 0:
67 return (
68 False,
69 "RDNA3 W4A16 kernel does not support channelwise quantization",
70 )
71
72 if c.full_weight_shape[0] % c.group_size != 0:
73 return (
74 False,
75 f"Group size ({c.group_size}) does not evenly divide K "
76 f"({c.full_weight_shape[0]})",
77 )
78
79 # Output features must be a multiple of the pack factor (8 nibbles per
80 # int32) and of 8 so that qzeros (packed 4-bit per col) align cleanly
81 # against the BLOCK_KN_SIZE*4 = 512 N-stride and per-thread 4 columns.
82 if c.partition_weight_shape[1] % 8 != 0:
83 return (
84 False,
85 "Output features must be a multiple of 8 for the RDNA3 "
86 "W4A16 kernel (qzeros packing)",
87 )
88
89 return True, None
90
91 # ----- Weight prep (identical layout/shuffle as ExllamaLinearKernel) -----
92
93 def process_weights_after_loading(self, layer: torch.nn.Module):
94 c = self.config
95 device = getattr(layer, self.w_q_name).device
96
97 # Synthesize zero points if the checkpoint doesn't carry them.
98 if not c.zero_points:
99 self.w_zp_name = "qzeros"
100 groups = c.partition_weight_shape[0] // c.group_size
101 out_features = c.partition_weight_shape[1]
102
103 if c.weight_type.has_bias():
104 # GPTQv1 quirk: the kernel adds 1 to the stored zero, so we
105 # encode (bias - 1) here. See exllama.py for the link to the
106 # documentation of this checkpoint-format wart.
107 zeros = torch.full(
108 (groups, out_features),
109 c.weight_type.bias - 1,
110 dtype=torch.int32,
111 device=device,
112 )
113 else:
114 raise NotImplementedError(
115 "RDNA3 W4A16 kernel: zero-bias 4-bit quant requires "
116 "explicit zero points (GPTQv1 +1 quirk)."
117 )
118 zeros = pack_quantized_values_into_int32(zeros, c.weight_type, packed_dim=1)
119 setattr(
120 layer, self.w_zp_name, torch.nn.Parameter(zeros, requires_grad=False)
121 )
122
123 def transform_w_q(x):
124 assert isinstance(x, BasevLLMParameter)
125 permute_param_layout_(x, input_dim=0, output_dim=1, packed_dim=0)
126 x_cont = x.data.contiguous()
127 # Same 4-bit shuffle as exllama. The RDNA3 kernel reads weights in
128 # the same shuffled int32 layout and uses the (qa & 0x000F000F)
129 # bit-trick on top.
130 ops.gptq_shuffle(x_cont, c.weight_type.size_bits)
131 return x_cont
132
133 def transform_w_s(x):
134 assert isinstance(x, BasevLLMParameter)
135 permute_param_layout_(x, input_dim=0, output_dim=1)
136 x.data = x.data.contiguous()
137 # Keep scales in the activation dtype (fp16 OR bf16) — the kernel
138 # branches on dtype internally.
139 return x.to(dtype=c.act_type)
140
141 self._transform_param(layer, self.w_q_name, transform_w_q)
142 self._transform_param(layer, self.w_s_name, transform_w_s)
143
144 # ----- Forward --------------------------------------------------------
145
146 def apply_weights(
147 self,
148 layer: torch.nn.Module,
149 x: torch.Tensor,
150 bias: torch.Tensor | None = None,
151 ) -> torch.Tensor:
152 c = self.config
153
154 x_2d = x.reshape(-1, x.shape[-1])
155 out_shape = x.shape[:-1] + (c.partition_weight_shape[1],)
156
157 w_q, w_s, w_zp = self._get_weight_params(layer)
158 assert w_zp is not None, "Zero points are required by RDNA3 W4A16"
159 output = ops.gptq_gemm_rdna3(x_2d, w_q, w_zp, w_s, False)
160
161 if bias is not None:
162 output.add_(bias)
163 return output.reshape(out_shape)