[Quantization][Decompression] Fix QDQ for dynamic quant; Update NVFP4 Compression Params

[dsikka]

• Add compression info for nvfp4 to support decompression with CompressedLinear
• Update to support batch_size > 1 when running QDQ for tensor_group / group dynamic activations - this is done by updatng the reshape command to be generic such that all dimensions are maintained apart from the last / group_dim which has to be reshaped for QDQ

[kylesayrs]

If this doesn't support 4d activations, you should add asserts making sure that the ndims matches expectations

[brian-dellabetta]

We should also get in

• #393

another dynamic quant fix

[rahul-tuli]

LGTM!

[kylesayrs]

I need to understand this first, sorry

[kylesayrs]

In the future it might be nice to take a step back and make a decision about when tensors need to be reshaped during the qdq process.

Maybe rather than reshaping all of in compute_dynamic_scales_and_zp, _process_quantization, and dequantize, we can have this logic just exist once in _process_quantization.

The implementation below is what it might look like for both activations and weights. Some of this is wrong (it's late for me) but this function ensures the last dim is the granularity you want to quantize by.

def reshape_for_groups(func):
   def wrapper(x, args, ...):
       assert x.ndim >= 2

       if args.strategy == "token":
           pass
       if args.strategy == "channel":
           x = x.unsqueeze(-1)
       if args.strategy in ("group", "tensor_group"):
           num_groups = x.size(-1) // args.group_size
           x = x.unflatten(-1, (num_groups, args.group_size))
       if args.strategy == "block":
           block_height, block1_width = args.block_structure
           x = x.unfold(-2, block_height, block_height)  # [num_horiz, x.dim[-1], block_height]
           x = x.unfold(-2, block1_width, block1_width)  # [num_horiz, num_vert, block_height, block_width]
           x = flatten(-4, -3)  # [num_blocks, block_height x block_width]

       x = func(x, args, ...)

       if args.strategy == "token":
           pass
       if args.strategy == "channel":
           x = x.squeeze(-1)
       if args.strategy in ("group", "tensor_group"):
           return x.flatten(-2, -1)
       if args.strategy == "block":
           x = torch.cat(x, dim=-2)
           x = torch.cat(x, dim=-2)

   return wrapper

@reshape_for_groups
def _process_quantization(x, args, ...):
   if do_quantize:
        ...
   if do_dequantize:
        ...

[kylesayrs]

Thank you!