Clarification regarding the readme example

[sayakpaul]

from comfy_kitchen.tensor import QuantizedTensor, TensorCoreFP8Layout, TensorCoreNVFP4Layout

# Quantize a tensor
x = torch.randn(128, 256, device="cuda", dtype=torch.bfloat16)
qt = QuantizedTensor.from_float(x, TensorCoreFP8Layout)

# Operations dispatch to optimized kernels automatically
output = torch.nn.functional.linear(qt, weight_qt)

# Dequantize back to float
dq = qt.dequantize()

How is weight_qt computed in this case?

Also, is it recommended to quantize the input (which seems to be qt in this case) instead of the weights?

[contentis]

The from_float function calls into the quantize function from the provided LayoutCls. In this case it would call this

How is weight_qt computed in this case?

Here we assume that weight_qt is already an instance of QuantizedTensor. Usually we have a checkpoint that already has the weights quantized and load them as such directly by creating an instance. If you have a high-precision tensor/weight you could also leverage from_float to quantize it.

Also, is it recommended to quantize the input (which seems to be qt in this case) instead of the weights?

This entirely depends on your goal and dosen't have a short answer. The TLDR is that we can distinguish between Weight-Only (WoQ) and Weigth+Activation. When using WoQ methods we only quantize the weights and keep the input in high-precision float. When runnning an operation (e.g. MatMul) we upcast the weight into an intermediate buffer and perform the math in high-precision. out = x @ weight_qt.dequantize()
This saved VRAM but has no performance impact (usually a bit slower than high precision) but its easy to do.

For the other we can make use of dedicated Hardware like FP8 Tensor Cores. And this is where we see the performance gains, although often a bit harder to get good quality out of the box.

Here is some more in-depth explaination if you're interested: https://developer.nvidia.com/blog/model-quantization-concepts-methods-and-why-it-matters/

[sayakpaul]

Thanks so much!

For the other we can make use of dedicated Hardware like FP8 Tensor Cores. And this is where we see the performance gains, although often a bit harder to get good quality out of the box.

Oh? I was under the impression that even for weight-only FP8, there can still be kernels written to utilize the dedicated FP8 tensor cores. I guess you meant a workflow like so?

• Keep params in FP8, casts FP32/FP16/BF16 activations to FP8 on-the-fly
• Use FP8 tensor cores for the actual computation

# Pseudo-code for optimized weight-only FP8
out = fp8_matmul(
    x.cast_to_fp8(),           # Cast activation to FP8 on-the-fly
    weight_fp8,                 # Already stored in FP8
    output_dtype=fp16           # Accumulate in higher precision
)

Here we assume that weight_qt is already an instance of QuantizedTensor. Usually we have a checkpoint that already has the weights quantized and load them as such directly by creating an instance. If you have a high-precision tensor/weight you could also leverage from_float to quantize it.

That's what I had assumed but having a fully complete example (potentially with a checkpoint) would have been great as a starting point I believe.

[contentis]

Yeah, documentation and tooling could be improved and is something to be looked at. One last minor note, this:x.cast_to_fp8() is not a cast operation in the classcial sense but a scaled transformation where the tensor is first resaled to match the lower-precisions dynamic range and then catsted (round nearest).

[sayakpaul]

cast_to_fp8() is not a cast operation in the classcial sense but a scaled transformation where the tensor is first resaled to match the lower-precisions dynamic range and then catsted (round nearest).

Yes, of course. Lots of configurations there two in terms of how much granularity one wants (per-row vs. per-tensor, etc.).

Yeah, documentation and tooling could be improved and is something to be looked at.

Maybe we could keep this thread open till then?