Fix slicing and get_plain() in GemLite

[mobicham]

Contributions

• GemLite 4.7.0 uses FMA mode by default to improve dequantization performance (Wq * s + z instead of (W_q - z ) * s), so we need to update get_plain() to make it compatible with both formats.
• Updated slicing to work directly on the packed data, since the older version that was using get_plain() was causing vLLM issues.

Notes

• gemlite.set_kernel_caching(True) gives wrong output with torchao but not when using gemlite as a module, not sure why, but that would impact perf for batch-size=1 by up to 10 tokens/sec.

Tests

End-2-End test

https://gist.github.com/mobicham/54fed6f18bee590f615f18391b45b71e

Slicing Test

import torch, gemlite
from torchao.quantization import GemliteUIntXWeightOnlyConfig, quantize_
device = 'cuda:0'
dtype = torch.float16
gemlite.set_autotune("default")

torch.manual_seed(0)
in_features, out_features, group_size = 256, 512, 64

orig_shape = [out_features, in_features]
layer = torch.nn.Linear(in_features, out_features, bias=False, dtype=dtype, device=device)
layer.weight.data /= 10.
weight = layer.weight.data.clone()

quantize_(layer, GemliteUIntXWeightOnlyConfig(bit_width=4, group_size=group_size))

meta_args =  layer.weight.tensor_impl.gemlite_kwargs['meta_args']
W_group_mode = meta_args[10]

#Test matmul
####################################################################################
torch.manual_seed(0)
x = torch.randn((1, layer.in_features), device=device, dtype=dtype) / 10.
y_ref = x @ weight.T
y_gem  = layer(x)
err = (y_ref - y_gem).abs().mean()
assert err < 5e-3, "Dot product mismatch. " + str(err)

#Test slicing 
####################################################################################
def dequant(input_layer, in_features, orig_shape):
    int_data = input_layer.tensor_impl.packed_weight
    scale = input_layer.tensor_impl.scale
    zero_point = input_layer.tensor_impl.zero_point

    W_q = (
        gemlite.bitpack.unpack_over_rows(
            int_data, W_nbits=4, num_output_rows=in_features, dtype=torch.uint8
        )
        .T.contiguous()
        .view([-1, group_size])
    )

    s = scale.t().contiguous().view(-1, 1)
    z = zero_point.t().contiguous().view(-1, 1)

    if W_group_mode == 4:  # FMA
        W_deq = (W_q * s + z).view(orig_shape)
    else:
        W_deq = ((W_q - z) * s).view(orig_shape)

    return W_deq

W_r = dequant(layer.weight, layer.in_features, orig_shape) #~weight


#Slicing in half
for slice_axis, start, end in [(0, 0, 256), (0, 256, 256), (1, 0, 128), (1, 128, 128)]:
    layer_sliced = layer.weight.narrow(slice_axis, start, end)

    if slice_axis == 0:
        num_rows, out_shape = layer.in_features, (orig_shape[0]//2, orig_shape[1]) 
    else:
        num_rows, out_shape = layer.in_features // 2, (orig_shape[0], orig_shape[1]//2)

    W_slice = dequant(layer_sliced, num_rows, out_shape)

    W_slice_ref = W_r[start:start+end, :] if slice_axis == 0 else W_r[:, start:start+end]
    assert (W_slice_ref - W_slice).abs().mean() == 0, f"Slicing {start}:{end} along axis={slice_axis} is incorrect"
    ```

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/2288

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 1 New Failure

As of commit b2892ce with merge base 35ffb26 ():

NEW FAILURE - The following job has failed:

• PR Label Check / Check PR Labels (gh)
  Process completed with exit code 1.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[jerryzh168]

could you incorporate test into

ao/test/dtypes/test_affine_quantized.py

Line 363 in 8366465

def test_slice_gemlite(self, device, dtype):

as well

[mobicham]

Updated the test and successfully tested on vLLM.

[mobicham]

In vLLM, I get _same_metadata() with models that have lm_head quantized, it seems to me that ao';s vLLM implementation doesn't support that?