Study Plan

Int8 Matmul with Per-Channel Scales

JAX

Medium

Multiply a batch of activations by a weight matrix in int8 under jax.jit, with one symmetric scale for every output column of the weights and one for every row of the activations, and accumulate the integer products in int32.

The input x has shape (B, D_in), with one token per row, and the weight w has shape (D_in, D_out), with one output channel per column. The scale of column $j$ of w and the scale of row $b$ of x are

s^w_j = \frac{\max_i \lvert w_{ij} \rvert}{127}, \qquad s^x_b = \frac{\max_i \lvert x_{bi} \rvert}{127},

and a column or row that is entirely zero gets the scale 1 instead. The int8 codes are

q^w_{ij} = \operatorname{clip}\left( \operatorname{round}\left( \frac{w_{ij}}{s^w_j} \right), -127, 127 \right), \qquad q^x_{bi} = \operatorname{clip}\left( \operatorname{round}\left( \frac{x_{bi}}{s^x_b} \right), -127, 127 \right),

where round goes to the nearest integer. The output is

a_{bj} = \sum_{i} q^x_{bi} \, q^w_{ij}, \qquad y_{bj} = a_{bj} \, s^x_b \, s^w_j,

where $a_{bj}$ is the integer sum of code products, computed from the int8 codes with an int32 accumulator, and $y_{bj}$ is entry (b, j) of the float32 result, which approximates x times w.

Return a dict of JAX arrays with the keys w_q (int8, shape (D_in, D_out)), w_scales (float32, shape (D_out,)), x_q (int8, shape (B, D_in)), x_scales (float32, shape (B,)) and y (float32, shape (B, D_out)).

Examples

Input: x = [[1.2,2.0,-0.9]], w = [[0.5,-0.023],[-1.27,0.011],[0.3,0.041]]

Output: {"w_q":[[50,-71],[-127,34],[30,127]],"w_scales":[0.01,0.000323],"x_q":[[76,127,-57]],"x_scales":[0.015748],"y":[[-2.210866,-0.042284]]}

Explanation: The row of x has absmax 2.0, so its scale is 2 / 127, about 0.015748, and its codes are [76,127,-57]. Column 0 of w has absmax 1.27 and scale 0.01, while column 1 gets its own much smaller scale, about 0.000323, from its absmax of 0.041. The int32 sums of code products are -14039 and -8317, and multiplying each by the row scale and its column scale gives y, which is close to the float product [-2.21,-0.0425].

Input: x = [[0.0,0.0],[0.5,2.0]], w = [[0.8,0.0,-0.31],[-0.45,0.0,0.66]]

Output: {"w_q":[[127,0,-60],[-71,0,127]],"w_scales":[0.006299,1.0,0.005197],"x_q":[[0,0],[32,127]],"x_scales":[1.0,0.015748],"y":[[0.0,0.0,0.0],[-0.491339,0.0,1.162867]]}

Input: x = [[1.0,1.0,1.0,1.0],[0.7,-0.5,0.3,1.2]], w = [[0.5,0.25],[0.5,-0.25],[0.5,0.25],[0.5,0.1]]

Output: {"w_q":[[127,127],[127,-127],[127,127],[127,51]],"w_scales":[0.003937,0.001969],"x_q":[[127,127,127,127],[74,-53,32,127]],"x_scales":[0.007874,0.009449],"y":[[2.0,0.350394],[0.850394,0.496063]]}

Hint 1

jnp.max(jnp.abs(v), axis=k) reduces over axis k, so choose k to get one value per column or row.

Hint 2

jnp.round, jnp.clip(v, -127, 127) and .astype(jnp.int8) turn scaled values into int8 codes.

Hint 3

jnp.matmul(a, b, preferred_element_type=jnp.int32) multiplies two int8 arrays and keeps every sum in int32.

Requirements

Your function is called through jax.jit, so every step must use jax.numpy or jax.lax operations, and an all-zero row or column must be handled with jnp.where instead of a Python if.
Multiply x_q by w_q as int8 arrays with jax.lax.dot_general or jnp.matmul and preferred_element_type=jnp.int32, so the sums are accumulated in int32, and get y by converting those int32 sums to float32 and multiplying them by the row scale and the column scale.
y must be computed from that int8 matmul. The traced program may contain no other matmul, no convolution, no callback into Python such as jax.pure_callback, and no broadcast multiply and sum over all B × D_in × D_out products, so a matmul of dequantized float32 values is rejected even when its numbers come out close.
Form the outer product of the two scales by broadcasting, as in x_scales[:, None] * w_scales[None, :], because a matmul of the scales counts as a float matmul.
w_q and x_q must match exactly, and w_scales and x_scales must match to within a relative error of 1e-6.
Each entry of y is compared with $a_{bj} \, s^x_b \, s^w_j$ computed in float64 from the exact integer sum and the float32 scales, and it must be within a relative error of 1e-6 of that value, so an entry whose integer sum is 0 must be exactly 0.

Constraints

x is float32 with shape (B, D_in), and w is float32 with shape (D_in, D_out), where B and D_out are between 1 and 6 and D_in is between 1 and 64.
All values are finite, and a row of x or a column of w may be entirely zero.
No ratio of a value to its scale lies within 0.002 of a point halfway between two integers, so rounding is never ambiguous.
An integer sum $a_{bj}$ can reach 64 × 16,129 = 1,032,256 in magnitude, which fits in int32 but overflows int8 and int16.

Try Similar Problems

Mixed Precision with Loss ScalingMedium Rematerialized MLP BlockMedium Blockwise Attention with Online SoftmaxHard Pad and Mask to Bucket LengthMedium Hessian-Vector ProductMedium Per-Example Gradient ClippingMedium Straight-Through Estimator with Custom VJPMedium Traced Learning Rate ScheduleMedium Exponential Moving Average with ScanMedium Spectral Norm by Power IterationMedium Batched Attention with vmapMedium Pairwise Distances with Nested vmapMedium