[PyTorch][Common] Add MXFP8 and NVFP4 QDQ for mixed-precision GEMM

[negvet]

Description

Adds experimental MXFP8 and NVFP4 quantize–dequantize (QDQ) quantizers, similar to radixark/miles#2864.

• Fuses quantization and dequantization on SM100, avoiding intermediate quantized data and scale buffers. Native Q+DQ fallback supports additional NVFP4 modes and mixed input/output dtypes.
• Adds recipes for X[BF16] × W[NVFP4-QDQ] forward with BF16 backward, and X[MXFP8-QDQ] × W[NVFP4-QDQ] forward with MXFP8 backward.
• Introduces a shared QDQ quantizer base for custom formats. Includes numerical, distributed, and C++ tests and usage documentation.

def mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory(role):
    if role is None or role.module_type not in ("linear", "grouped_linear"):
        return IdentityQuantizer()

    if role.tensor_type == "input":
        return HybridQuantizer(
            rowwise_quantizer=MXFP8QDQQuantizer(),
            columnwise_quantizer=mxfp8_factory(role),
            columnwise_source="original",
        )

    if role.tensor_type == "weight":
        return HybridQuantizer(
            rowwise_quantizer=NVFP4QDQQuantizer(),
            columnwise_quantizer=mxfp8_factory(role),
            columnwise_source="rowwise_dequantized",
        )

    if role.tensor_type == "grad_output":
        return mxfp8_factory(role)

    return IdentityQuantizer()


recipe = CustomRecipe(
    qfactory=mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory,
)

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Changes

Please list the changes introduced in this PR:

• Change A
• Change B

Checklist:

• I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[rybakov]

can it support NVFP4 with per token second scale instead of global amax?

[negvet]

Added support

[rybakov]

X[MXFP8_QDQ] x W[NVFP4_QDQ] ?
the same comment is applied to all QDQ types used in the code and description

also in function name may be it is better to explicitly name it as
mxfp8_qdq_nvfp4_qdq_fwd_mxfp8_bwd_factory

I would suggest to add QDQ to all names of mxfp8 and nvfp4 when qdq is applied

[negvet]

Sure, fixed

[rybakov]

looks good, with comment about QDQ naming in the description and function names

[negvet]

/te-ci pytorch

[greptile-apps]

Confidence Score: 5/5

[High impact] The PR appears safe to merge; the older-CUDA test failure is fixed and no new blocking issue was found.

Summary

Adds experimental MXFP8 and NVFP4 QDQ quantizers that store decoded operands for high-precision GEMMs.

• MXFP8 and NVFP4 QDQ quantizers round values before high-precision computation.
• Linear layers can pair QDQ forward values with BF16 or MXFP8 backward formats.
• Custom quantizers can match stored values to a requested compute dtype.
• Custom recipes can enter autocast without an upfront FP8 hardware check.

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    Input[Original operand] --> Choice{QDQ backend}
    Choice -->|Supported| Fused[Fused quantize and decode]
    Choice -->|Reference| Native[Native quantize then decode]
    Fused --> HP[High-precision storage]
    Native --> HP
    HP --> Forward[High-precision forward GEMM]
    HP --> Backward[Optional backward quantizer]

Loading

_{Reviews (3) · Last reviewed commit: "Allow CustomRecipe without FP8 hardware ..." · Reviewed by Greptile}

[negvet]

/te-ci pytorch

[negvet]

/te-ci pytorch