|
| 1 | +# ----------------------------------------------------------------------------- |
| 2 | +# |
| 3 | +# Copyright (c) Qualcomm Technologies, Inc. and/or its subsidiaries. |
| 4 | +# SPDX-License-Identifier: BSD-3-Clause |
| 5 | +# |
| 6 | +# ----------------------------------------------------------------------------- |
| 7 | + |
| 8 | +import onnxscript |
| 9 | +import torch |
| 10 | +from onnx import TensorProto |
| 11 | + |
| 12 | +from QEfficient.utils import constants |
| 13 | + |
| 14 | +ops = getattr(onnxscript, "opset" + str(constants.ONNX_EXPORT_OPSET)) |
| 15 | + |
| 16 | + |
| 17 | +@onnxscript.script(onnxscript.values.Opset("com.qti.aisw.onnx", 1)) |
| 18 | +def CastToUInt4(weight_packed: onnxscript.UINT8) -> onnxscript.UINT8: |
| 19 | + """ |
| 20 | + Unpack packed uint8 weights into uint4 values and cast output to UINT4. |
| 21 | + Supports N-D input: all leading dimensions are preserved; only the last |
| 22 | + dimension (in_features // 2) is doubled to (in_features). |
| 23 | +
|
| 24 | + Input: (..., in_features // 2) UINT8 |
| 25 | + Each byte holds two nibbles: byte = (w_y << 4) | (w_x & 0x0F) |
| 26 | + Output: (..., in_features) UINT4, values in [0, 15] |
| 27 | +
|
| 28 | + Operations: |
| 29 | + w_x = weight_packed % 16 (lower nibble) |
| 30 | + w_y = (weight_packed >> 4) % 16 (upper nibble) |
| 31 | + stacked = concat([w_x, w_y], axis=-1) after unsqueeze |
| 32 | + → (..., in//2, 2) |
| 33 | + leading_dims = shape[:-1] |
| 34 | + new_shape = [...leading_dims, last_dim * 2] |
| 35 | + reshaped = reshape(stacked, new_shape) |
| 36 | + output = Cast(reshaped, to=UINT4) |
| 37 | + """ |
| 38 | + sixteen = ops.CastLike(ops.Constant(value_ints=[16]), weight_packed) |
| 39 | + |
| 40 | + # Lower nibble: weight_packed & 0x0F = weight_packed % 16 |
| 41 | + w_x = ops.Mod(weight_packed, sixteen) |
| 42 | + |
| 43 | + # Upper nibble: (weight_packed >> 4) & 0x0F |
| 44 | + shift = ops.CastLike(ops.Constant(value_ints=[4]), weight_packed) |
| 45 | + w_shifted = ops.BitShift(weight_packed, shift, direction="RIGHT") |
| 46 | + w_y = ops.Mod(w_shifted, sixteen) |
| 47 | + |
| 48 | + # Stack along a new last dim → (..., in_features//2, 2) |
| 49 | + w_x_unsq = ops.Unsqueeze(w_x, [-1]) |
| 50 | + w_y_unsq = ops.Unsqueeze(w_y, [-1]) |
| 51 | + stacked = ops.Concat(w_x_unsq, w_y_unsq, axis=-1) |
| 52 | + |
| 53 | + # N-D aware reshape: preserve all leading dims, double the last dim. |
| 54 | + # packed_shape = [d0, d1, ..., last_dim] |
| 55 | + packed_shape = ops.Shape(weight_packed) |
| 56 | + # All dims except the last: [d0, d1, ...] |
| 57 | + leading_dims = ops.Slice(packed_shape, starts=[0], ends=[-1], axes=[0]) |
| 58 | + # Last dim only: [last_dim] |
| 59 | + last_dim = ops.Slice(packed_shape, starts=[-1], ends=[2147483647], axes=[0]) |
| 60 | + # Double the last dim: [last_dim * 2] |
| 61 | + last_dim_doubled = ops.Mul(last_dim, ops.Constant(value_ints=[2])) |
| 62 | + # New shape: [d0, d1, ..., last_dim * 2] |
| 63 | + new_shape = ops.Concat(leading_dims, last_dim_doubled, axis=0) |
| 64 | + reshaped = ops.Reshape(stacked, new_shape) |
| 65 | + |
| 66 | + # Cast to UINT4 — data_type value is version-dependent (21 in ONNX 1.18, 23 in newer) |
| 67 | + return ops.Cast(reshaped, to=int(TensorProto.UINT4)) |
| 68 | + |
| 69 | + |
| 70 | +class CastToUInt4Func(torch.autograd.Function): |
| 71 | + """ |
| 72 | + Custom op: unpacks packed uint8 → uint8 (values 0-15) in PyTorch. |
| 73 | + In ONNX the custom op subgraph includes a Cast → UINT4 as its last step. |
| 74 | + Supports N-D input: all leading dimensions are preserved. |
| 75 | +
|
| 76 | + PyTorch forward : packed uint8 (..., in//2) → uint8 (..., in), values [0, 15] |
| 77 | + ONNX symbolic : emits CastToUInt4 node (com.qti.aisw.onnx) |
| 78 | + The subgraph ends with Cast → UINT4. |
| 79 | + """ |
| 80 | + |
| 81 | + @staticmethod |
| 82 | + def forward(weight_packed: torch.Tensor) -> torch.Tensor: |
| 83 | + w_x = weight_packed & 0x0F # lower nibble, (..., in//2), range [0, 15] |
| 84 | + w_y = (weight_packed >> 4) & 0x0F # upper nibble, (..., in//2), range [0, 15] |
| 85 | + # New shape: all leading dims unchanged, last dim doubled |
| 86 | + new_shape = list(weight_packed.shape[:-1]) + [weight_packed.shape[-1] * 2] |
| 87 | + return torch.stack( |
| 88 | + [w_x, w_y], dim=-1 |
| 89 | + ).reshape( |
| 90 | + new_shape |
| 91 | + ) # Can't add a cast operation to uint4 here, as its not supported in pytorch; The ONNX export will handle the cast to IINT4 in the symbolic method. |
| 92 | + |
| 93 | + @staticmethod |
| 94 | + def setup_context(ctx, inputs, outputs): |
| 95 | + pass |
| 96 | + |
| 97 | + @staticmethod |
| 98 | + def symbolic(g: torch.Graph, weight_packed: torch.Value) -> torch.Value: |
| 99 | + output = g.onnxscript_op(CastToUInt4, weight_packed) |
| 100 | + return output |
| 101 | + |
| 102 | + |
| 103 | +class DequantizeLinearFunc(torch.autograd.Function): |
| 104 | + """ |
| 105 | + Emits a standard ONNX DequantizeLinear node (ai.onnx domain, not custom). |
| 106 | +
|
| 107 | + Symmetric blockwise quantization — no zero_point: |
| 108 | + output = x * scale (per block along the last axis) |
| 109 | +
|
| 110 | + Supports N-D input: |
| 111 | + weight_unpacked : (..., in_features) — quantized values |
| 112 | + scale : (..., num_blocks) — per-block scales |
| 113 | + block_size : int — elements per block |
| 114 | +
|
| 115 | + PyTorch forward : expand blockwise scale along last dim, multiply |
| 116 | + ONNX symbolic : DequantizeLinear(weight_unpacked, scale, |
| 117 | + axis=2, block_size=block_size) |
| 118 | + axis=2 for 3D input (2, out_features, in_features). |
| 119 | + No zero_point input (symmetric). |
| 120 | + """ |
| 121 | + |
| 122 | + @staticmethod |
| 123 | + def forward( |
| 124 | + weight_unpacked: torch.Tensor, scale: torch.Tensor, zeros: torch.Tensor, block_size: int |
| 125 | + ) -> torch.Tensor: |
| 126 | + # Expand per-block scale → per-element scale along last dim |
| 127 | + scale_expanded = scale.repeat_interleave(block_size, dim=-1) |
| 128 | + zeros_expanded = zeros.repeat_interleave(block_size, dim=-1) |
| 129 | + return (weight_unpacked.to(torch.int8) - zeros_expanded.to(torch.int8)) * scale_expanded |
| 130 | + |
| 131 | + @staticmethod |
| 132 | + def setup_context(ctx, inputs, outputs): |
| 133 | + pass |
| 134 | + |
| 135 | + @staticmethod |
| 136 | + def symbolic( |
| 137 | + g: torch.Graph, weight_unpacked: torch.Value, scale: torch.Value, zeros: torch.Value, block_size: int |
| 138 | + ) -> torch.Value: |
| 139 | + # Standard DequantizeLinear: symmetric (no zero_point), blockwise. |
| 140 | + # Input is 3D: (2, out_features, in_features) → axis=2 (last dim). |
| 141 | + # DequantizeLinear natively supports batch dimensions. |
| 142 | + return g.op( |
| 143 | + "DequantizeLinear", |
| 144 | + weight_unpacked, |
| 145 | + scale, |
| 146 | + zeros, |
| 147 | + axis_i=2, |
| 148 | + block_size_i=block_size, |
| 149 | + ) |
0 commit comments