RFC: Add quantize_mx ATen Op for Microscaling Quantization

[avizon-aws]

Summary

Add a new ATen operator quantize_mx to support Microscaling (MX) format quantization in PyTorch. MX formats provide hardware-efficient quantization with shared exponents across blocks of values, enabling lower precision training and inference while maintaining model accuracy.

Motivation

Microscaling formats (MXFP8, MXFP4, MXINT8, NVFP4) are emerging quantization standards that:

• Provide better accuracy-performance tradeoffs than traditional INT8/FP8 quantization
• Use shared exponents per block (typically 32 elements)
• Are being adopted by hardware vendors (Nvidia, AWS TRN, AMD)

Currently, PyTorch lacks native support for quantizing into MX formats, implement a native aten would provide a cleaner option for users/libraries to enable quantization e.g. TorchAO to_mx functions

Proposed API

Core Operator

torch.quantize_mx(
    input: Tensor,
    dim: int, 
    block_size: int, 
    dtype: torch.dtype,  # "fp8_e5m2", "fp8_e4m3", "fp6_e3m2", "fp6_e2m3", "fp4", "int8"
    scale_calculation_mode: int  #{0: Round Down, 1: Round Up, 2: Round to Nearest ...}
) -> Tensor, Tensor

"""

    input: Tensor to be quantized
    dim: Dimension on perform quantization on i.e. reduction dimension e.g. 0,-1, has to be the innermost/outermost dimension
    dtype: data type of the quantized tensor
    scale_calculation_mode: the scale calculated can be rounded up/down/nearest etc.
    
"""

Storage Format

Tensor Layout for Input:
- input: [B, S, H], can be 2D/3D tensor

Tensor Layout for Output:
- quantized_input: [B, S, H], same as input
- scales: [B, S, H//block_size], depends on dimension of quantization

ATen Operator Signature

// aten/src/ATen/native/native_functions.yaml

- func: quantize_mx(Tensor self, int dim, int block_size, ScalarType dtype, int scale_calculation_mode) -> (Tensor output, Tensor scale)
  variants: function, method
  dispatch:
    CPU: quantize_mx_cpu
    Meta: quantize_mx_meta
  tags: [
    core,                 
    pt2_compliant_tag    
  ]

Implementation Plan (High level summary)

1. Will implement feature example in MXFP8 with CPU goldens and tests
2. Integration : Will include DTensor, torch.compile and eager mode.
3. Additional dtypes and backend/accelerators can be added by other contributors

[albanD]

cc @slayton58 @drisspg

[slayton58]

Managed to miss this, but very much relevant to my interests :)

I’ve actually proposed within meta a more general op, aping the API from torch.nn.functional.{scaled_mm,scaled_grouped_mm}, so provide an input tensor, and some enum (and other pertinent info), and get an appropriately scaled tensor (along with the corresponding scaling factors) back.

I can expound more on this when not on a phone, but I feel this idea both encapsulates your ask, along with being more general, covering nvfp4 and deepseek subtensor scalings, as well as more traditional tensor- and row-wise scalings.

I’d be interested in your thoughts :)

[avizon-aws]

Yes, this API is generic, currently I was of the opinion to keep this limited to MX formats, but if there are good enough reasons to encapsulate other scaling techniques (tensor-wise, row-wise etc.) then I am happy to consider that as well. My reasoning for keeping it limited to MX formats was that if we try to encapsulate more options, then the number of attributes to the API might grow significantly and since MX is quite different from the others (I think there already exist apis like quantize_per_tensor etc.), I thought to only consider using it for MX Formats, but happy to discuss further.

so provide an input tensor, and some enum (and other pertinent info), and get an appropriately scaled tensor (along with the corresponding scaling factors) back.

I think this is a good idea, but we will have to consider the scenarios where the user should be able to specify custom implementations for the quantize i.e. some hardware might have native support for quantization e.g. blackwell.

[slayton58]

good enough reasons to encapsulate other scaling techniques

I think we care at a minimum about MX and NV formats (and being block formats, the implementations are conceptually similar), and I do think there's worth in supporting all quantization schemes that we support for mm/grouped_mm.

My reasoning for keeping it limited to MX formats was that if we try to encapsulate more options, then the number of attributes to the API might grow significantly

I don't necessarily agree with this, consider something like:

torch.quantize_scaled_tensor(
    input: Tensor,
    dim: int,
    dtype: torch.dtype,  # "fp8_e5m2", "fp8_e4m3", "fp6_e3m2", "fp6_e2m3", "fp4", "int8"
    scale_type: List[ScalingType], # TENSORWISE, ROWWISE, BLOCK_1x32, BLOCK_1x16, etc.
    scale_tensors: List[Tensor] | None, # Additional metadata, i.e. statically calibrated per-tensor 2nd-level scale for NVFP4
) -> tuple[Tensor]

Implementation can be sparse (i.e. not the entire cross-product of options are actually useful, or would need implementing), but this can handle MX, NV, DeepSeek and less complex formats, is not (imo) materially more complex than your proposed API, and closely matches the API of torch.nn.functional.{scaled_mm,scaled_grouped_mm} methods that would be expected to consume the outputs.

consider the scenarios where the user should be able to specify custom implementations for the quantize i.e. some hardware might have native support for quantization e.g. blackwell.

I would anticipate this is handled the same way again as the scaled_mm operations - there's a top-level function that's specialized to each backend (i.e. CUDA, XPU, CPU, etc.) which implements its chosen subset of quantizations with custom implementations.

[avizon-aws]

Yeah, the API you presented above for generality is the one I initially thought of but then discarded, it does make sense to me, so I am aligned on it. Your reasoning for scaled_mm also makes sense and I think it would be covered in a generic way by the quantize API (will rename from quantize_mx -> quantize).

Let me know if there are more details needed for this, I will paste the final API (mostly same as above) for final review.

[HahTK]

If we are to fully generalize this I propose we consider additional parameters for scale_dtype and block_size. scale_type is doing too much and is vulnerable to combinatorial explosion if we see more innovation in quantization mechanisms.

Something along these lines but I'll leave @avizon-aws to make the final proposal.

torch.quantize_scaled_tensor(
    input: Tensor,
    dim: int,
    dtype: torch.dtype,  # "fp8_e5m2", "fp8_e4m3", "fp6_e3m2", "fp6_e2m3", "fp4", "uint8"
    block_size : List[int], # specify multi-dimensional quantization block with a list. "-1" for the entire dimension,
    scale_type: List[ScalingType], # TENSORWISE, ROWWISE, BLOCK, etc.
    scale_dtype: torch.dtype, # "fp8_e8m0", "fp8_e5m2", "fp8_e4m3", "fp6_e3m2", "fp6_e2m3", "fp4", "uint8"
    scale_tensors: List[Tensor] | None, # Additional metadata, i.e. statically calibrated per-tensor 2nd-level scale for NVFP4
) -> tuple[Tensor]

@avizon-aws : Extra notes for consideration in your final design.
IMHO : block_size has enough descriptive power to eliminate scale_type as defined previously, but we may still be need it to differentiate similar schemes with different underlying computation. Example, original (buggy) MX OCP MXFP8 vs NVFP8 both look exactly the same except scale factor rounding (torch.floor vs torch.ceil respectively). The right thing to do may be a redefinition

@slayton58 :
Any thoughts on the proposed modifications?
And why did you imaging that scale_type would need to be a list? Am I missing something here?

nits : uint8 instead of int8 and probably just fp8_e8m0 instead of uint8.

[avizon-aws]

I gave this some more thought, in order to be generic and also ensure that future quantization schemes are handled well, my proposal is as shown below:

torch.quantize(
    input: Tensor,
    dim: int,
    dtype: torch.dtype,  # "fp8_e5m2", "fp8_e4m3", "fp6_e3m2", "fp6_e2m3", "fp4", "uint8"
    block_size : List[int], # specify multi-dimensional quantization block with a list. "-1" for the entire dimension,
    scale_type: ScalingType, # TENSORWISE, ROWWISE, BLOCK, etc.
    scale_dtype: torch.dtype, # "fp8_e8m0", "fp8_e5m2", "fp8_e4m3", "fp6_e3m2", "fp6_e2m3", "fp4", "uint8", "fp8_e8m0",
    double_quantization: bool = False, # quantize the x and x.T
    extra_tensor_params: Optional[Dict[str, Tensor]], # Additional tensor parameters, can be used for enabling new schemes, add options e.g. outlier_threshold (clip outliers beyond a threshold), epsilon, rounding_mode (stochastic, round to nearest) etc.
) -> List[Tensor]

Changes to API from previous version:

1. scale_type: to support various quantization schemes
2. multidimensional block_size: scale_type is duplicate because block_size can convey all information but as mentioned previously by @HahTK , it might be used for cases like MX OCP MXFP8 vs NVFP8
3. double_quantization: this might be useful in most cases as the fwd and bwd pass require a quantization of the transposed tensor
4. extra_tensor_params: this is to ensure that attributes from future quantization schemes can be easily added without increasing the length of parameters in the API
5. variable length output: depending on the type of quantization, the number of outputs might differ e.g. we might get qdata, scales, zero points for INT quantization, qdata, fp8 scale, FP32 scale for NVFP4 etc.

[slayton58]

scale_type and block_size here are duplicating functionality - ScalingType encodes things like BLOCKWISE_128x128 where > 1 dimensions are being quantized.

double_quantization is, IMO, overloading a term too -- NVFP4 could be considered a "double quantized" type due to quant of both elements and per-block scaling factors. I'm also not necessarily convinced we want the additional complexity of quantizing both x and x.t in the same function - I understand the use-case, but it feels a little awkward here.

extra_tensor_params -- I'm not a fan of this approach, but I'm not 100% sure how to resolve this unease with something actually workable (as-in actually provide constructive feedback). In general, this API doesn't need to be all things to all people - it's ok to leave more esoteric functionality on the table in order to provide an API that does 99% of what's needed.

And why did you imaging that scale_type would need to be a list? Am I missing something here?
Because of multi-scale formats (like NVFP4) and future-proofing.

[avizon-aws]

I see, extra_tensor_params is probably not the best design pattern for the API, but without this I think we should be able to satisfy the needs of most users without it. Regarding double_quantization , i can remove it as its not really a critical feature, users can just call quantize API twice.

So the final proposal becomes:

torch.quantize(
    input: Tensor,
    dim: int,
    dtype: torch.dtype,  # "fp8_e5m2", "fp8_e4m3", "fp6_e3m2", "fp6_e2m3", "fp4", "uint8"
    block_size : List[int], # specify multi-dimensional quantization block with a list. "-1" for the entire dimension,
    scale_type: List[ScalingType], # TENSORWISE, ROWWISE, BLOCK, etc.
    scale_dtype: torch.dtype, # "fp8_e8m0", "fp8_e5m2", "fp8_e4m3", "fp6_e3m2", "fp6_e2m3", "fp4", "uint8", "fp8_e8m0",
) -> List[Tensor]

@slayton58 could you clarify the use of :

scale_tensors: List[Tensor] | None, # Additional metadata, i.e. statically calibrated per-tensor 2nd-level scale for NVFP4

I have currently not added the above attribute, once i understand it, we can add it.

[slayton58]

For NVFP4 you generally statically calibrate on a per-tensor basis to provide a global amax to scale the e4m3 per-block scaling factors - this needs to be passed in. The list[Tensor] is to future-proof -- one can produce quantization schemes with arbitrary nested scalings (think per-1x32-block, per-128x128-block, per-tensor) as an example.

[avizon-aws]

Got it, but I am not sure why we need to pass the scales in to the API? I was planning that all the quantization logic lives inside this API, the user should not have to perform a partial scale calculation outside the API imo, so if the user specifies:

torch.quantize(
    input:
    dim: int,
    dtype: torch.fp4
    block_size =[16]
    scale_type: ["NVFP4_TENSOR_WISE, NVFP4_1X32, etc."] # TENSORWISE, ROW, COLUMN, MXFP, NVFP etc.
    scale_dtype: ["fp32", "e4m3"]
) -> List[Tensor]

All the per tensor(or block or whatever user wants) , followed by local block scale and quantization should be performed inside the API.

I think above gives the flexibility required for all cases.

scale_type and block_size here are duplicating functionality - ScalingType encodes things like BLOCKWISE_128x128 where > 1 dimensions are being quantized.

Yes, this is true to some extent, but I think having both is better for e.g. if some wants to use blockwise with 128x128, but someone else wants to use blockwise with 256x256 etc. there can be so many such combinations. Having the attributes block_size and scale_type make it simpler to handle within the API, we wont need to add multiple levels of if-else, we just need to add handling for "blockwise" and then block_size can directly be used as its an attribute.

[HahTK]

@avizon-aws :

1. Feedback on extra_tensor_params is aligned with internal review. so lgtm.
2. As currently designed, I concur with @slayton58 for scale_tensors: List[Tensor]. It is needed to pass in amax.
   • OTOH, I would prefer all quantization to be in torch.quantize but even in that case we need to pass in a dtype for the global scale (or nested scales).

@slayton58 : Thoughts on decoupling block_size and scale_type

1. Decoupling will help avoid combinatorial explosion while supporting a variety of quantization schemes. While there are indeed a set of more common schemes, we should not assume they are the only ones. ROW for example could technically be a variety of block sizes, and while NVFP4 is typically 16x1, they also published NVFP4 schemes with 16x16 block sizes (admittedly that trained rather poorly). I think we should have a flexible API that avoids potential combinatorial explosions.

2. Furthermore, one could write generic code for ROW and size but if you pass in ROW_128 and ROW_64, we'd have to map/extract size from ROW_128 and ROW_64 to some size=128 or 64 for a generic implementation.

3. Nested scaling actually support decoupling of block_size and scale_type. If you have 2 scales in scale_tensors where would we define block sizes for the each of the 2 scales?

[slayton58]

@HahTK Opinions on decoupling below.

1. This is not intended to be a library for truly generic quantization research - there exist codes for that already. The intention here is to support the subset of schemes that have existing traction / HW support - this to me is the subset we have support in torch.nn.functional.scaled_mm, and schemes that are easily implementable using that functionality (i.e. 16x16 NVFP4 and 32x32 MX formats). If another format becomes useful to implement it can be added as its own enum, but the number of combinations that are supported is deliberately small.
2. One could write generic code, and the 1x16 and 1x32 quant kernels will in fact be very similar, but that doesn't necessarily hold beyond those examples. For bs=64 or 128 for instance I might want a different mapping of resources -> problem.
3. scale_type : list[ScalingType] encodes the block-sizes.

[vkuzo]

torch.quantize(

I'd advocate for a significantly more specific name (for example, quantize_tensor_to_lowp), torch.quantize is very general since one can quantize many things (activations, weights, optimizer state, subgraphs, models, etc).

also, can we provide some examples of what expressing MXFP8 1d, NVFP4 1d, MXFP8 2d, etc is going to look like with the current WIP?