• Cody Yu's avatar
    [Kernel][FP8] Initial support with dynamic per-tensor scaling (#4118) · a22cdea3
    Cody Yu authored
    Provide an initial support to FP8 computation. This PR is inspired by HuggingFace TGI: huggingface/text-generation-inference#1726
    
    This feature can be enabled with --quantization fp8 or -q fp8 when launching an engine.
    
    Algorithm:
    We still load a model checkpoint in FP16/BF16. After the weights are loaded, Fp8LinearMethod calculates the per-tensor scaling factor of weights and quantizes the weights accordingly. The scaling factor will then be stored for future use. Meanwhile, the per-tensor scaling factor for activations is calculated in every forward pass.
    
    Initial Results:
    Currently tested Mistral-7B on 1xH100. With prompt length ~5 and decoding length 128:
    
    BF16: 1.47s
    FP8: 1.66s
    I'll try to use larger models and try to find more performance bottleneck. Meanwhile, you're welcome to try this code.
    a22cdea3
linear.py 24.8 KB