43. Fake-Quantized (INT8) Linear Layer

Problem

Cartesia serves real-time voice models where inference latency and memory bandwidth are first-order constraints, and one standard lever for cutting both is quantization: run the matmul-heavy layers at int8 instead of fp32/fp16. Before shipping true integer kernels, teams commonly prototype and train through a "fake quantization" version of the layer: round weights to int8 levels, then immediately dequantize back to float so the rest of the graph runs in float as normal. This lets you measure the accuracy impact of quantization, and — if used during training — lets gradients flow through the layer (quantization-aware training).

Implement a linear layer whose weight is fake-quantized to a symmetric num_bits integer grid before the matmul.

Source: src/43_quantization_aware_linear.py

def fake_quant_linear(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor | None,
    num_bits: int = 8,
) -> torch.Tensor:
    ...

Examples:
>>> x = torch.randn(2, 4)
>>> w = torch.randn(3, 4)
>>> out = fake_quant_linear(x, w, bias=None, num_bits=8)
>>> out.shape
torch.Size([2, 3])

Step-by-Step Approach

1. Pick a symmetric integer range for num_bits: qmin = -2^(num_bits-1), qmax = 2^(num_bits-1) - 1 (e.g. [-128, 127] for int8).
2. Compute a per-tensor scale from the weight's dynamic range: scale = weight.abs().max() / qmax, so the largest-magnitude weight maps to the edge of the integer grid.
3. Detach the scale from autograd — it's a statistic of the current weights, not a value you want gradients flowing into directly.
4. Quantize: divide by scale, round to the nearest integer, and clamp into [qmin, qmax] for safety against outliers.
5. Dequantize immediately by multiplying back by scale — the matmul itself still runs in float, but every weight value it sees has been snapped onto the int8 grid's precision.
6. Run a standard F.linear(x, w_dequant, bias) with the fake-quantized weight.

The key insight is that "fake" quantization never touches integer arithmetic at all — it only restricts float values to the precision a real int8 kernel would have, which is exactly what you want during training/calibration before switching to a genuine integer kernel at deploy time.

Reference solution

def fake_quant_linear(x, weight, bias, num_bits=8):
    # symmetric int8 grid: [-2^(b-1), 2^(b-1)-1], zero maps to zero exactly
    qmax = 2 ** (num_bits - 1) - 1
    qmin = -(2 ** (num_bits - 1))

    # per-tensor scale: largest magnitude weight maps to the top of the grid.
    # detach() so the scale itself isn't part of the autograd graph
    scale = weight.detach().abs().max().clamp(min=1e-8) / qmax

    # round-to-nearest onto the integer grid, then clamp for safety
    w_int = torch.clamp(torch.round(weight / scale), qmin, qmax)

    # dequantize immediately: the matmul still runs in float, but the
    # weight values are restricted to int8-grid precision
    w_dequant = w_int * scale

    return F.linear(x, w_dequant, bias)

Key Functions & Tricks

• torch.round — snaps a float onto the nearest integer of the quantization grid.
• torch.clamp — bounds both the quantized integers (grid range) and the scale (avoids divide-by-zero for all-zero weights).
• tensor.detach() — stops the scale computation from adding an extra path into the autograd graph.
• F.linear — the standard x @ w.T + b op, applied here to the dequantized weight.
• Symmetric quantization (zero-centered) — simpler than affine/asymmetric quantization since it needs no zero-point offset, at the cost of wasting one grid level when the tensor isn't zero-centered.

How to Recognize This Pattern