| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Triton-based W4A16 GEMM kernel for ROCm MI300. |
| 4 | |
| 5 | Implements fused int4-weight dequantization + fp16 GEMM in a single kernel, |
| 6 | using GPTQ sequential packing (8 int4 values per int32, shifts [0,4,...,28]). |
| 7 | Plugs into the MPLinearKernel selection system and is preferred over |
| 8 | MarlinLinearKernel/ExllamaLinearKernel on ROCm. |
| 9 | |
| 10 | Weight layout expected by this kernel (post-process_weights_after_loading): |
| 11 | qweight: [K, N//8] int32 — rows=K (input), cols=N//8 (N is packed) |
| 12 | scales: [K//G, N] fp16/bf16 |
| 13 | qzeros: [K//G, N//8] int32 (optional; None for symmetric uint4b8) |
| 14 | |
| 15 | Checkpoint layout from compressed_tensors_wNa16 create_weights: |
| 16 | weight_packed: [N, K//8] int32 (output_dim=0, input_dim=1, packed_dim=1) |
| 17 | weight_scale: [N, K//G] fp16 (output_dim=0, input_dim=1) |
| 18 | weight_zero_point: [N//8, K//G] int32 (output_dim=0, packed_dim=0) |
| 19 | """ |
| 20 | |
| 21 | import torch |
| 22 | |
| 23 | from vllm.model_executor.layers.quantization.utils import replace_parameter |
| 24 | from vllm.model_executor.parameter import BasevLLMParameter, permute_param_layout_ |
| 25 | from vllm.platforms import current_platform |
| 26 | from vllm.scalar_type import scalar_types |
| 27 | from vllm.triton_utils import tl, triton |
| 28 | from vllm.utils.torch_utils import direct_register_custom_op |
| 29 | |
| 30 | from .MPLinearKernel import MPLinearKernel, MPLinearLayerConfig |
| 31 | |
| 32 | TRITON_W4A16_SUPPORTED_GROUP_SIZES = [-1, 32, 64, 128, 256] |
| 33 | TRITON_W4A16_SUPPORTED_QUANT_TYPES = [ |
| 34 | scalar_types.uint4b8, # symmetric GPTQ (bias=8) |
| 35 | scalar_types.uint4, # asymmetric with explicit zeros |
| 36 | ] |
| 37 | |
| 38 | |
| 39 | @triton.jit |
| 40 | def triton_w4a16_gemm_kernel( |
| 41 | # Pointers |
| 42 | a_ptr, # [M, K] fp16/bf16 activations |
| 43 | b_ptr, # [K, N//8] int32 packed 4-bit weights (N is the packed dim) |
| 44 | scales_ptr, # [K//G, N] fp16/bf16 scales |
| 45 | zeros_ptr, # [K//G, N//8] int32 packed zeros (unused when HAS_ZP=False) |
| 46 | c_ptr, # [M, N] fp16/bf16 output |
| 47 | # Dimensions |
| 48 | M, |
| 49 | N, |
| 50 | K, |
| 51 | # Strides |
| 52 | stride_am, |
| 53 | stride_ak, |
| 54 | stride_bk, |
| 55 | stride_bn, # stride in b along the packed N//8 dim |
| 56 | stride_cm, |
| 57 | stride_cn, |
| 58 | # Quantization parameters |
| 59 | group_size, |
| 60 | # Whether explicit zero points are provided |
| 61 | HAS_ZP: tl.constexpr, |
| 62 | # Zero bias used when HAS_ZP is False (e.g. 8 for uint4b8) |
| 63 | ZP_BIAS: tl.constexpr, |
| 64 | # Block sizes (tuned for MI300 wavefront=64) |
| 65 | BLOCK_M: tl.constexpr, |
| 66 | BLOCK_N: tl.constexpr, |
| 67 | BLOCK_K: tl.constexpr, |
| 68 | ): |
| 69 | """Fused W4A16 GEMM: C[M,N] = A[M,K] @ dequant(B)[K,N]. |
| 70 | |
| 71 | B is stored as [K, N//8] int32 using GPTQ sequential packing: |
| 72 | each int32 packs 8 consecutive N-values at bit offsets [0,4,8,12,16,20,24,28]. |
| 73 | |
| 74 | Dequant: w_fp = (w_int4 - zero) * scale |
| 75 | HAS_ZP=True: zero is loaded from zeros_ptr and unpacked |
| 76 | HAS_ZP=False: zero = ZP_BIAS constant (e.g. 8 for uint4b8 symmetric) |
| 77 | """ |
| 78 | pid_m = tl.program_id(0) |
| 79 | pid_n = tl.program_id(1) |
| 80 | |
| 81 | # Row/col offsets for this tile |
| 82 | offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) |
| 83 | offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) |
| 84 | |
| 85 | # b/zeros are stored with N packed: N//8 int32 columns per K row |
| 86 | offs_bn = pid_n * (BLOCK_N // 8) + tl.arange(0, BLOCK_N // 8) |
| 87 | |
| 88 | # GPTQ sequential shifts tiled across BLOCK_N: |
| 89 | # [0,4,8,...,28] repeating for every group of 8 N-values. |
| 90 | # Build 1D shifts_1d of length BLOCK_N: column j gets shift (j % 8) * 4. |
| 91 | shifts_row = tl.arange(0, 8) * 4 # [8] |
| 92 | shifts_1d_2d = tl.broadcast_to(shifts_row[None, :], (BLOCK_N // 8, 8)) |
| 93 | shifts_1d = tl.reshape(shifts_1d_2d, (BLOCK_N,)) # [BLOCK_N] |
| 94 | # Broadcast to [BLOCK_K, BLOCK_N] for weight unpacking |
| 95 | shifts = tl.broadcast_to(shifts_1d[None, :], (BLOCK_K, BLOCK_N)) |
| 96 | |
| 97 | # Scales column offsets: full N-width (one scale per output neuron) |
| 98 | offs_sn = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) |
| 99 | |
| 100 | accumulator = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) |
| 101 | |
| 102 | for k_start in range(0, tl.cdiv(K, BLOCK_K)): |
| 103 | offs_k = k_start * BLOCK_K + tl.arange(0, BLOCK_K) |
| 104 | mask_k = offs_k < K |
| 105 | |
| 106 | # ---- Load activations A: [BLOCK_M, BLOCK_K] ---- |
| 107 | a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak |
| 108 | mask_a = (offs_m[:, None] < M) & mask_k[None, :] |
| 109 | a = tl.load(a_ptrs, mask=mask_a, other=0.0) |
| 110 | |
| 111 | # ---- Load packed weights B: [BLOCK_K, BLOCK_N//8] int32 ---- |
| 112 | b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn |
| 113 | mask_b = mask_k[:, None] & (offs_bn[None, :] < N // 8) |
| 114 | b_packed = tl.load(b_ptrs, mask=mask_b, other=0) |
| 115 | |
| 116 | # ---- Unpack int4 weights → [BLOCK_K, BLOCK_N] ---- |
| 117 | # tl.interleave(x, x) doubles the last dim by interleaving. |
| 118 | # Starting from [BLOCK_K, BLOCK_N//8], three interleaves give |
| 119 | # [BLOCK_K, BLOCK_N], where each int32 is replicated 8 times. |
| 120 | b = tl.interleave(b_packed, b_packed) |
| 121 | b = tl.interleave(b, b) |
| 122 | b = tl.interleave(b, b) |
| 123 | # Extract the correct 4-bit nibble for each output column |
| 124 | b = (b >> shifts) & 0xF |
| 125 | |
| 126 | # ---- Compute scale/zero group row index ---- |
| 127 | group_idx = (k_start * BLOCK_K) // group_size |
| 128 | |
| 129 | # ---- Load scales: [BLOCK_N] → broadcast to [BLOCK_K, BLOCK_N] ---- |
| 130 | scale_offset = group_idx * N + offs_sn |
| 131 | scale_mask = offs_sn < N |
| 132 | scales = tl.load(scales_ptr + scale_offset, mask=scale_mask, other=1.0) |
| 133 | scales = tl.broadcast_to(scales[None, :], (BLOCK_K, BLOCK_N)) |
| 134 | |
| 135 | # ---- Load / compute zeros ---- |
| 136 | if HAS_ZP: |
| 137 | # Load packed zeros row: [BLOCK_N//8] int32 |
| 138 | zero_offset = group_idx * (N // 8) + offs_bn |
| 139 | zero_mask = offs_bn < N // 8 |
| 140 | z_packed = tl.load(zeros_ptr + zero_offset, mask=zero_mask, other=0) |
| 141 | # Unpack to [BLOCK_N] using same interleave+shift pattern |
| 142 | z = tl.interleave(z_packed, z_packed) |
| 143 | z = tl.interleave(z, z) |
| 144 | z = tl.interleave(z, z) |
| 145 | z = (z >> shifts_1d) & 0xF |
| 146 | z = tl.broadcast_to(z[None, :], (BLOCK_K, BLOCK_N)) |
| 147 | else: |
| 148 | z = tl.full((BLOCK_K, BLOCK_N), ZP_BIAS, dtype=tl.int32) |
| 149 | |
| 150 | # ---- Dequantize: (w - zero) * scale ---- |
| 151 | b_fp = (b - z).to(a.dtype) * scales |
| 152 | |
| 153 | # ---- Accumulate ---- |
| 154 | accumulator += tl.dot(a, b_fp, out_dtype=tl.float32) |
| 155 | |
| 156 | # ---- Store output C: [BLOCK_M, BLOCK_N] ---- |
| 157 | c = accumulator.to(c_ptr.type.element_ty) |
| 158 | c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn |
| 159 | mask_c = (offs_m[:, None] < M) & (offs_n[None, :] < N) |
| 160 | tl.store(c_ptrs, c, mask=mask_c) |
| 161 | |
| 162 | |
| 163 | def _triton_w4a16_gemm_impl( |
| 164 | a: torch.Tensor, # [M, K] fp16/bf16 |
| 165 | b_q: torch.Tensor, # [K, N//8] int32 |
| 166 | scales: torch.Tensor, # [K//G, N] fp16/bf16 |
| 167 | qzeros: torch.Tensor | None, # [K//G, N//8] int32, or None |
| 168 | group_size: int, |
| 169 | zp_bias: int = 8, # bias for uint4b8 when qzeros is None |
| 170 | ) -> torch.Tensor: |
| 171 | """Fused W4A16 GEMM using GPTQ-packed int4 weights. |
| 172 | |
| 173 | Args: |
| 174 | a: Activation matrix [M, K], float16 or bfloat16. |
| 175 | b_q: Packed weight matrix [K, N//8], int32 (GPTQ sequential). |
| 176 | scales: Per-group scales [K//G, N], same dtype as a. |
| 177 | qzeros: Per-group packed zero points [K//G, N//8] int32, or None |
| 178 | for symmetric quantization (uses zp_bias instead). |
| 179 | group_size: Quantization group size (resolved from -1 to K by caller). |
| 180 | zp_bias: Constant zero used when qzeros is None (default 8 for uint4b8). |
| 181 | |
| 182 | Returns: |
| 183 | Output matrix [M, N], same dtype as a. |
| 184 | |
| 185 | """ |
| 186 | assert a.is_contiguous(), "Activation matrix must be contiguous" |
| 187 | assert b_q.is_contiguous(), "Weight matrix must be contiguous" |
| 188 | assert scales.is_contiguous(), "Scales must be contiguous" |
| 189 | |
| 190 | M, K = a.shape |
| 191 | N = b_q.shape[1] * 8 |
| 192 | |
| 193 | assert b_q.shape == (K, N // 8), ( |
| 194 | f"b_q shape mismatch: {b_q.shape} vs ({K}, {N // 8})" |
| 195 | ) |
| 196 | assert scales.shape == (K // group_size, N), ( |
| 197 | f"scales shape mismatch: {scales.shape} vs ({K // group_size}, {N})" |
| 198 | ) |
| 199 | if qzeros is not None: |
| 200 | assert qzeros.shape == (K // group_size, N // 8), ( |
| 201 | f"qzeros shape mismatch: {qzeros.shape}" |
| 202 | ) |
| 203 | |
| 204 | c = torch.empty((M, N), dtype=a.dtype, device=a.device) |
| 205 | |
| 206 | has_zp = qzeros is not None |
| 207 | # Provide a dummy pointer when HAS_ZP=False (Triton requires a valid ptr) |
| 208 | zeros_ptr = qzeros if has_zp else b_q |
| 209 | |
| 210 | if current_platform.is_rocm(): |
| 211 | from vllm.platforms.rocm import on_gfx1x |
| 212 | |
| 213 | if on_gfx1x(): |
| 214 | # Tuned for RDNA 3.5 (gfx1151, 40 CUs, 32-wide wavefronts). |
| 215 | if M <= 32: |
| 216 | BLOCK_M, BLOCK_N, BLOCK_K = 32, 32, 64 |
| 217 | elif M <= 64: |
| 218 | BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32 |
| 219 | else: |
| 220 | BLOCK_M, BLOCK_N, BLOCK_K = 128, 32, 64 |
| 221 | else: |
| 222 | # Tuned for MI300 (gfx942, 304 CUs, 64-wide wavefronts). |
| 223 | if M <= 32: |
| 224 | BLOCK_M, BLOCK_N, BLOCK_K = 32, 64, 32 |
| 225 | elif M <= 64: |
| 226 | BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32 |
| 227 | else: |
| 228 | BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32 |
| 229 | else: |
| 230 | if M <= 32: |
| 231 | BLOCK_M, BLOCK_N, BLOCK_K = 32, 64, 32 |
| 232 | elif M <= 64: |
| 233 | BLOCK_M, BLOCK_N, BLOCK_K = 64, 64, 32 |
| 234 | else: |
| 235 | BLOCK_M, BLOCK_N, BLOCK_K = 128, 128, 32 |
| 236 | |
| 237 | # The kernel loads scales/zeros for a single group per BLOCK_K tile |
| 238 | # (one group index per iteration). If BLOCK_K > group_size, rows at the tail |
| 239 | # of the tile dequantize with the wrong group's scales, silently |
| 240 | # corrupting the output. Clamp BLOCK_K to group_size to keep one |
| 241 | # scale group per tile. |
| 242 | if group_size < BLOCK_K: |
| 243 | BLOCK_K = group_size |
| 244 | |
| 245 | grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N)) |
| 246 | |
| 247 | triton_w4a16_gemm_kernel[grid]( |
| 248 | a, |
| 249 | b_q, |
| 250 | scales, |
| 251 | zeros_ptr, |
| 252 | c, |
| 253 | M, |
| 254 | N, |
| 255 | K, |
| 256 | a.stride(0), |
| 257 | a.stride(1), |
| 258 | b_q.stride(0), |
| 259 | b_q.stride(1), |
| 260 | c.stride(0), |
| 261 | c.stride(1), |
| 262 | group_size=group_size, |
| 263 | HAS_ZP=has_zp, |
| 264 | ZP_BIAS=zp_bias, |
| 265 | BLOCK_M=BLOCK_M, |
| 266 | BLOCK_N=BLOCK_N, |
| 267 | BLOCK_K=BLOCK_K, |
| 268 | ) |
| 269 | return c |
| 270 | |
| 271 | |
| 272 | def _triton_w4a16_gemm_fake( |
| 273 | a: torch.Tensor, |
| 274 | b_q: torch.Tensor, |
| 275 | scales: torch.Tensor, |
| 276 | qzeros: torch.Tensor | None, |
| 277 | group_size: int, |
| 278 | zp_bias: int = 8, |
| 279 | ) -> torch.Tensor: |
| 280 | return torch.empty((a.size(0), b_q.size(1) * 8), dtype=a.dtype, device=a.device) |
| 281 | |
| 282 | |
| 283 | direct_register_custom_op( |
| 284 | op_name="triton_w4a16_gemm", |
| 285 | op_func=_triton_w4a16_gemm_impl, |
| 286 | mutates_args=[], |
| 287 | fake_impl=_triton_w4a16_gemm_fake, |
| 288 | ) |
| 289 | |
| 290 | |
| 291 | def triton_w4a16_gemm( |
| 292 | a: torch.Tensor, |
| 293 | b_q: torch.Tensor, |
| 294 | scales: torch.Tensor, |
| 295 | qzeros: torch.Tensor | None, |
| 296 | group_size: int, |
| 297 | zp_bias: int = 8, |
| 298 | ) -> torch.Tensor: |
| 299 | return torch.ops.vllm.triton_w4a16_gemm(a, b_q, scales, qzeros, group_size, zp_bias) |
| 300 | |
| 301 | |
| 302 | class TritonW4A16LinearKernel(MPLinearKernel): |
| 303 | """Triton-based W4A16 GEMM kernel for ROCm (MI300 and newer). |
| 304 | |
| 305 | Supports GPTQ-format int4 weights (uint4b8 symmetric, uint4 asymmetric) |
| 306 | with grouped quantization. Weight tensors are transposed from the |
| 307 | compressed-tensors checkpoint layout to the kernel's [K, N//8] layout. |
| 308 | """ |
| 309 | |
| 310 | SUPPORTED_QUANT_TYPES = TRITON_W4A16_SUPPORTED_QUANT_TYPES |
| 311 | |
| 312 | @classmethod |
| 313 | def get_min_capability(cls) -> int: |
| 314 | # Triton handles capability checks itself |
| 315 | return 0 |
| 316 | |
| 317 | @classmethod |
| 318 | def can_implement(cls, c: MPLinearLayerConfig) -> tuple[bool, str | None]: |
| 319 | if not (current_platform.is_rocm() or current_platform.is_cuda()): |
| 320 | return False, "TritonW4A16LinearKernel requires CUDA or ROCm" |
| 321 | |
| 322 | if c.weight_type not in cls.SUPPORTED_QUANT_TYPES: |
| 323 | return ( |
| 324 | False, |
| 325 | f"Quant type {c.weight_type} not supported; " |
| 326 | f"supported: {cls.SUPPORTED_QUANT_TYPES}", |
| 327 | ) |
| 328 | |
| 329 | if c.act_type not in (torch.float16, torch.bfloat16): |
| 330 | return False, "Only float16/bfloat16 activations are supported" |
| 331 | |
| 332 | N = c.partition_weight_shape[1] |
| 333 | if N % 8 != 0: |
| 334 | return ( |
| 335 | False, |
| 336 | f"Output features ({N}) must be divisible by 8 " |
| 337 | "(8 int4 values packed per int32)", |
| 338 | ) |
| 339 | |
| 340 | gs = c.group_size |
| 341 | if ( |
| 342 | gs not in TRITON_W4A16_SUPPORTED_GROUP_SIZES |
| 343 | and gs != c.full_weight_shape[0] |
| 344 | ): |
| 345 | return ( |
| 346 | False, |
| 347 | f"Group size {gs} not supported; " |
| 348 | f"supported: {TRITON_W4A16_SUPPORTED_GROUP_SIZES} " |
| 349 | f"or full K ({c.full_weight_shape[0]})", |
| 350 | ) |
| 351 | |
| 352 | K = c.partition_weight_shape[0] |
| 353 | eff_gs = gs if gs != -1 else K |
| 354 | if K % eff_gs != 0: |
| 355 | return (False, f"Input features {K} not divisible by group size {eff_gs}") |
| 356 | |
| 357 | return True, None |
| 358 | |
| 359 | def process_weights_after_loading(self, layer: torch.nn.Module) -> None: |
| 360 | """Convert compressed-tensors checkpoint layout to kernel layout. |
| 361 | |
| 362 | Checkpoint (from compressed_tensors_wNa16.create_weights): |
| 363 | weight_packed: [N, K//8] int32 input_dim=1, output_dim=0, packed_dim=1 |
| 364 | weight_scale: [N, K//G] fp16 input_dim=1, output_dim=0 |
| 365 | weight_zero_point: [N//8, K//G] int32 output_dim=0, packed_dim=0 |
| 366 | |
| 367 | Kernel needs: |
| 368 | qweight: [K, N//8] int32 (transpose weight_packed) |
| 369 | scales: [K//G, N] fp16 (transpose weight_scale) |
| 370 | qzeros: [K//G, N//8] int32 (transpose weight_zero_point) |
| 371 | """ |
| 372 | |
| 373 | # ---- Transform qweight: [N, K//8] → [K//8, N] → back to [K, N//8] ---- |
| 374 | # permute_param_layout_(x, input_dim=0, output_dim=1) rearranges so that |
| 375 | # the input(K) dimension is at physical dim 0 and output(N) at dim 1. |
| 376 | # Checkpoint has input_dim=1, output_dim=0, packed_dim=1 (K is packed). |
| 377 | # After permute we get [K//8, N] (K packed at dim 0, N at dim 1). |
| 378 | # The kernel wants [K, N//8] (K at dim 0, N packed at dim 1), so we |
| 379 | # then transpose: [K//8, N].T = [N, K//8] — that's not right. |
| 380 | # |
| 381 | # Actually we need to change WHAT is packed: |
| 382 | # Original packing: K packed into K//8 (8 K-values per int32) |
| 383 | # Kernel packing: N packed into N//8 (8 N-values per int32) |
| 384 | # These require a full repack, not just a transpose. |
| 385 | # |
| 386 | # Simple approach: unpack → transpose the full [N, K] → repack as [K, N//8]. |
| 387 | # This is done CPU-side at load time (one-time cost). |
| 388 | def repack_w_q(x: BasevLLMParameter) -> BasevLLMParameter: |
| 389 | # x.data is [N, K//8] int32, K packed (GPTQ checkpoint format) |
| 390 | # Step 1: bring to [N, K//8] with output(N) at dim 0 |
| 391 | permute_param_layout_(x, input_dim=1, output_dim=0, packed_dim=1) |
| 392 | w = x.data # [N, K//8] int32 |
| 393 | |
| 394 | N_dim, K8 = w.shape |
| 395 | K_dim = K8 * 8 |
| 396 | # Step 2: unpack to [N, K] int32 (vectorized) |
| 397 | shifts = torch.arange(8, device=w.device, dtype=torch.int32) * 4 |
| 398 | w_unpacked = ((w.unsqueeze(-1) >> shifts) & 0xF).reshape(N_dim, K_dim) |
| 399 | # Step 3: transpose to [K, N] int32 |
| 400 | w_KN = w_unpacked.t().contiguous() |
| 401 | # Step 4: repack N into N//8 int32 values → [K, N//8] (vectorized) |
| 402 | N8 = N_dim // 8 |
| 403 | w_repacked = torch.sum( |
| 404 | (w_KN.view(K_dim, N8, 8) & 0xF) << shifts, |
| 405 | dim=2, |
| 406 | dtype=torch.int32, |
| 407 | ) |
| 408 | x.data = w_repacked.contiguous() |
| 409 | return x |
| 410 | |
| 411 | def repack_w_s(x: BasevLLMParameter) -> BasevLLMParameter: |
| 412 | # x.data is [N, K//G] fp16, bring to [K//G, N] |
| 413 | permute_param_layout_(x, input_dim=1, output_dim=0) |
| 414 | x.data = x.data.t().contiguous() |
| 415 | return x |
| 416 | |
| 417 | self._transform_param(layer, self.w_q_name, repack_w_q) |
| 418 | self._transform_param(layer, self.w_s_name, repack_w_s) |
| 419 | |
| 420 | if self.w_zp_name is not None: |
| 421 | zp = getattr(layer, self.w_zp_name, None) |
| 422 | if zp is not None: |
| 423 | # Kernel needs [K//G, N//8]: |
| 424 | # input(K) at dim 0, output(N) packed at dim 1. |
| 425 | # AutoGPTQ: |
| 426 | # output_dim=1 -> already [K//G, N//8], no transpose. |
| 427 | # compressed-tensors: |
| 428 | # output_dim=0 -> [N//8, K//G], needs transpose. |
| 429 | # None (unknown): |
| 430 | # infer from shape; if square (ambiguous), default to transpose. |
| 431 | zp_output_dim = getattr(zp, "output_dim", None) |
| 432 | if zp_output_dim is not None: |
| 433 | needs_transpose = zp_output_dim != 1 |
| 434 | else: |
| 435 | # in case output_dim is None |
| 436 | c = self.config |
| 437 | K, N = c.partition_weight_shape |
| 438 | group_size = c.group_size if c.group_size != -1 else K |
| 439 | expected_shape = (K // group_size, N // 8) |
| 440 | transposed_shape = (N // 8, K // group_size) |
| 441 | if ( |
| 442 | tuple(zp.data.shape) == expected_shape |
| 443 | and expected_shape != transposed_shape |
| 444 | ): |
| 445 | needs_transpose = False |
| 446 | else: |
| 447 | needs_transpose = True |
| 448 | zp_data = ( |
| 449 | zp.data.t().contiguous() |
| 450 | if needs_transpose |
| 451 | else zp.data.contiguous() |
| 452 | ) |
| 453 | replace_parameter( |
| 454 | layer, |
| 455 | self.w_zp_name, |
| 456 | torch.nn.Parameter(zp_data, requires_grad=False), |
| 457 | ) |
| 458 | |
| 459 | def apply_weights( |
| 460 | self, layer: torch.nn.Module, x: torch.Tensor, bias: torch.Tensor | None = None |
| 461 | ) -> torch.Tensor: |
| 462 | c = self.config |
| 463 | w_q, w_s, w_zp = self._get_weight_params(layer) |
| 464 | |
| 465 | x_2d = x.reshape(-1, x.shape[-1]).contiguous() |
| 466 | out_shape = x.shape[:-1] + (c.partition_weight_shape[1],) |
| 467 | |
| 468 | K = c.partition_weight_shape[0] |
| 469 | group_size = c.group_size if c.group_size != -1 else K |
| 470 | |
| 471 | # For symmetric types (uint4b8), use the scalar bias; no zeros tensor |
| 472 | if c.weight_type.has_bias(): |
| 473 | zp_bias = c.weight_type.bias |
| 474 | w_zp = None # symmetric: ignore qzeros, use scalar bias instead |
| 475 | else: |
| 476 | zp_bias = 0 |
| 477 | |
| 478 | output = triton_w4a16_gemm( |
| 479 | a=x_2d, |
| 480 | b_q=w_q, |
| 481 | scales=w_s, |
| 482 | qzeros=w_zp, |
| 483 | group_size=group_size, |
| 484 | zp_bias=zp_bias, |
| 485 | ) |
| 486 | |
| 487 | if bias is not None: |
| 488 | output.add_(bias) |
| 489 | |
| 490 | return output.reshape(out_shape) |