Feat (graph/gpxq): faster groupwise per-column quantization

[Giuseppe5]

Reason for this PR

When running algorithms in the GPxQ family, we need to quantize one column at the time.
For groupwise quantization, currently there is no way to quantize a single column without quantizing wastefully the entire matrix.

The main issue is related to the groupwise dependency on adjacent columns, and the shape manipulation we do when applying groupwise quantization.

This PR addresses that.
Some results for the following configuration:

End-to-end GPTQ

Groupwise INT8, eager

• Llama 8b

  • current: 30m
  • PR: 24m

• Qwen3 14B

  • current: 1h
  • PR: 45m

Groupwise quantization, layer shape [4096, 14436]

• MXINT8

  • torch.compile
    • current: 0.52 s
    • PR: 0.06 s
  • eager
    • current: 1.43 s
    • PR: 0.16 s

• MXFP8

  • torch.compile
    • current: 0.55 s
    • PR: 0.06 s
  • eager
    • current: 2.24 s
    • PR: 0.20 s

• Groupwise INT8 (with zero point)

  • torch.compile
    • current: 0.34 s
    • PR: 0.16 s
  • eager
    • current: 1.03 s
    • PR: 0.23 s

These were measured with the following utility

import argparse
import time

import torch

from brevitas.core.zero_point import ParameterFromStatsFromParameterZeroPoint
import brevitas.nn as qnn
from brevitas.nn.mixin import WeightRegion
from brevitas.quant.mx_quant_ocp import MXFloat8e4m3Weight
from brevitas.quant.mx_quant_ocp import MXInt8Weight
from brevitas.quant.shifted_scaled_int import ShiftedUint8WeightGroupQuantFloat


def synchronize(device):
    if device.type == 'cuda':
        torch.cuda.synchronize(device)


def measure(fn, repeats, device):
    samples = []
    for _ in range(repeats):
        synchronize(device)
        start = time.perf_counter()
        fn()
        synchronize(device)
        samples.append(time.perf_counter() - start)
    return torch.tensor(samples).median().item()


def main():
    parser = argparse.ArgumentParser(description="Benchmark GPxQ weight-region quantization")
    parser.add_argument('--out-features', type=int, default=4096)
    parser.add_argument('--in-features', type=int, default=4096)
    parser.add_argument('--group-size', type=int, default=32)
    parser.add_argument('--columns', type=int, default=128)
    parser.add_argument('--repeats', type=int, default=5)
    parser.add_argument('--format', choices=['int8', 'uint8', 'float8'], default='int8')
    parser.add_argument('--parameter-from-stats', action='store_true')
    parser.add_argument('--device', default='cuda')
    parser.add_argument('--compile', action='store_true')
    args = parser.parse_args()

    device = torch.device(args.device)
    quantizers = {
        'int8': MXInt8Weight,
        'uint8': ShiftedUint8WeightGroupQuantFloat,
        'float8': MXFloat8e4m3Weight}
    quantizer = quantizers[args.format]
    quantizer = quantizer.let(group_size=args.group_size)
    if args.parameter_from_stats:
        quantizer = quantizer.let(scaling_impl_type='parameter_from_stats')
        if args.format == 'uint8':
            quantizer = quantizer.let(
                zero_point_impl=ParameterFromStatsFromParameterZeroPoint)
    layer = qnn.QuantLinear(
        args.in_features,
        args.out_features,
        bias=False,
        weight_quant=quantizer,
        device=device)
    layer.eval()
    layer(torch.randn(1, args.in_features, device=device))
    if args.compile:
        layer.weight_quant.compile_quant()

    columns = torch.randperm(args.in_features)[:args.columns].tolist()

    def full_quantize_columns():
        for column in columns:
            layer.quant_weight().value[:, column:column + 1]

    def region_quantize_columns():
        for column in columns:
            layer.quant_weight_region(WeightRegion((None, (column, column + 1))))

    for column in columns[:min(4, len(columns))]:
        expected = layer.quant_weight().value[:, column:column + 1]
        actual = layer.quant_weight_region(WeightRegion((None, (column, column + 1))))
        torch.testing.assert_close(actual, expected)

    full_quantize_columns()
    region_quantize_columns()
    full_time = measure(full_quantize_columns, args.repeats, device)
    region_time = measure(region_quantize_columns, args.repeats, device)

    print(f"full median:   {full_time:.6f}s")
    print(f"region median: {region_time:.6f}s")
    print(f"speedup:       {full_time / region_time:.2f}x")


if __name__ == '__main__':
    main()