Function fp4_quantize

tensorrt_llm/quantization/quantize.py:539–552

(model, quant_config: QuantConfig)

def fp4_quantize(model, quant_config: QuantConfig):
    assert quant_config.quant_mode.has_nvfp4()
    quant_map = {
        ColumnLinear: FP4Linear,
        RowLinear: FP4RowLinear,
        MixtureOfExperts: MixtureOfExperts,
    }

    model = quantize_layers(
        model,
        quant_config,
        quant_map,
    )
    return model


# Now consider the kv cache is enabled for all layers