Bug: Incorrect handling of odd sized dims in nvfp4 packing/unpacking

[fynnsu]

compressed-tensors/src/compressed_tensors/compressors/quantized_compressors/nvfp4_quantized.py

Lines 140 to 150 in 09b7ed4

	# Handle odd length by padding if necessary
	if indices.numel() % 2 != 0:
	indices = torch.cat([indices, torch.zeros(1, dtype=torch.long, device=device)])
	# Reshape to pair consecutive elements
	indices = indices.reshape(-1, 2)
	# Pack pairs of 4-bit values into 8-bit values
	packed = (indices[:, 0] | (indices[:, 1] << 4)).to(torch.uint8)
	return packed.reshape(m, n // 2)

Given a tensor of shape (m, n) this function packs pairs of values into uint8s. However, if there is an odd number of values the current logic handles this by concatenating an additional zero onto the values. This would work but the uint8 tensor is being reshaped to (m, n // 2) before returning, which fails when n is odd.

e.g.
m=5, n=3 
-> 15 values 
-> 15 + 1 = 16 values after padding
-> 8 uint8s after packing
-> reshape 8 values into (5, 1)

The options for fixes are:

1. Return the flattened uint8 tensor without reshaping
2. padding with m zeros instead of one and reshape to (m, (n+1)//2)
3. Explicitly remove support for this case

I'm happy to submit a patch fixing this.

[brian-dellabetta]

Discussed with @fynnsu that it is better to raise an error (3rd option) rather than having it error out with ValueError at line 150. Asked @fynnsu to provide an MRE on the issue and a unit test in the resolving PR to account for this. We probably have yet to hit this since all dimensions in the llm are usually a multiple of 2

[fynnsu]

Oh yes, sorry MRE is essentially the unit test I added in #402.

from compressed_tensors.compressors.quantized_compressors.nvfp4_quantized import (
    pack_fp4_to_uint8,
)
import torch

x = torch.Tensor(
    [
        [-0.5000, -6.0000, -0.5000, -1.5000, -1.0000, 6.0000, 0.0000],
        [-1.0000, -6.0000, -0.5000, -0.0000, 0.5000, 0.5000, -0.0000],
        [1.5000, 6.0000, -0.0000, -0.5000, 1.0000, 1.0000, -0.0000],
    ]
)

_packed = pack_fp4_to_uint8(x)