Function quantize_to_fp4_tensor

tensorrt_llm/quantization/functional.py:1367–1393

Parameters: input : Tensor (On GPU) The input tensor. Its shape is [batch_size, seq_len, input_dim] or [num_tokens, input_dim] for remove_input_padding, should be fp16 sf_scale : Tensor (On GPU) The global per-tensor scaling factor. Its shape is [1,], sho

(input: Tensor, sf_scale: Tensor)

def quantize_to_fp4_tensor(input: Tensor, sf_scale: Tensor):
    '''
    Parameters:
        input : Tensor (On GPU)
            The input tensor. Its shape is [batch_size, seq_len, input_dim] or [num_tokens, input_dim] for remove_input_padding, should be fp16
        sf_scale : Tensor (On GPU)
            The global per-tensor scaling factor. Its shape is [1,], should be float32.
            used to scale SF from input range to fp8 range (448.f / (MaxVal of input / 6.f)).
        output : Tensor (On GPU)
            The output tensor. Its shape is [batch_size, seq_len, input_dim] or [num_tokens, input_dim] for remove_input_padding, should be FP4
        output_sf : Tensor (On GPU)
            The input scaling factor tensor. Its shape is [batch_size, seq_len, input_dim / scaling_vector_size] or [num_tokens, input_dim / scaling_vector_size] for remove_input_padding, should be FP8
    '''
    plg_creator = trt.get_plugin_registry().get_plugin_creator(
        'QuantizeToFP4', '1', TRT_LLM_PLUGIN_NAMESPACE)
    assert plg_creator is not None

    pfc = trt.PluginFieldCollection([])
    quantize_plug = plg_creator.create_plugin("quantize_to_fp4_plugin", pfc)

    plug_inputs = [input.trt_tensor, sf_scale.trt_tensor]
    layer = default_trtnet().add_plugin_v2(plug_inputs, quantize_plug)
    _add_plugin_info(layer, plg_creator, "quantize_to_fp4_plugin", pfc)

    quantized = _create_tensor(layer.get_output(0), layer)
    scales = _create_tensor(layer.get_output(1), layer)
    return quantized, scales


def dynamic_quantize(