Breaking the {Hardware} Barrier: Software program FP8 for Older GPUs

{Hardware} Degree Assist

Identical to AVX-512 directions, that are supported solely on a restricted variety of {hardware} platforms, FP8 and FP16 directions and registers are additionally restricted by {hardware} and can be found solely on the current ones. If you’re on an RTX-30 or RTX-20 collection GPU from Nvidia, then you definately won’t be able to make the most of this decrease precision FP8 sort. That is precisely the issue that Feather makes an attempt to resolve.

Packing Technique

Utilizing bitwise operators, one can simply pack the FP16 sort right into a FP32. The algorithm is described beneath.

Packing FP16

• Solid the enter FP32 right into a FP16; this step could be carried out with ease utilizing numpy’s astype operate. 
• Solid them to U16 after which to U32; this units the higher 16 bits to 0s and decrease 16 bits to the precise FP16.
• Shift certainly one of them by 16 utilizing the bitwise LSHIFT operator, and mix each of them utilizing the bitwise OR operator. 

Unpacking FP16

• Extract the decrease 16 bits utilizing the bitwise AND operator and masks 0xFFFF.
• Extract the higher 16 bits utilizing the RSHIFT operation by 16 after which carry out a bitwise AND operation with the masks 0xFFFF. 
• Solid each U16 values again to FP16 and to FP32 if wanted.

Packing FP8 

FP8 has two broadly used codecs – E5M2 & E4M3. One can not use the identical algorithm used for packing two FP16 into FP32 as a result of the CPU doesn’t help FP8 varieties natively, however does for FP16 (half precision); that is the explanation that np.float8 doesn’t exist. 

Casting an FP16 to FP8-E5M2 is simple, as seen within the determine, as a result of each have the identical variety of exponent bits and differ solely of their fraction. 

FP8-E5M2 Packing

• Solid the enter FP32 right into a FP16; this step could be carried out with ease utilizing numpy’s astype operate, or get the enter itself as FP16.
• Solid to U16, LSHIFT by 8, then RSHIFT by 8 to isolate the higher 8 bits
• Do that for all 4 FP32s or FP16s.
• Now utilizing the LSHIFT operator, shift them by 0, 8, 16 and 24 items and mix them utilizing the bitwise OR operator.

As soon as once more, unpacking ought to be easy; it’s the actual reverse of packing.

Packing an FP8-E4M3 shouldn’t be as simple and simple as packing an FP16 or FP8-E5M2, as a result of exponent bits mismatch.

As a substitute of implementing it from scratch, the library makes use of the ml_dtypes library, which already does the casting math.

The ml_dtypes library supplies help for generally used FP8 requirements, corresponding to E5M2 and E4M3 casting, for NumPy arrays. Utilizing the identical astype operate, we will carry out casting simply as we did for FP16 varieties. The Algorithm is strictly an identical to how we pack FP16, so I’m skipping it right here. 

Triton GPU Kernels

After we pack, we’d like an algorithm (kernel) to utilise this packed datatype and carry out the computation. Passing the packed datatype to a kernel carried out for FP32 or FP64 will lead to undefined computation as a result of we have now already corrupted the FP32 or FP64 being handed. Writing a kernel that takes the packed datatype as enter in CUDA shouldn’t be an easy process and is error-prone. That is precisely the place Triton shines; it’s a Area-Particular Language library that leverages a customized intermediate illustration for GPU kernels. In layman’s phrases, it permits one to jot down GPU kernels in Python itself with out the necessity to write CUDA kernels in C. 

Triton kernels do precisely what was talked about beforehand; the algorithm is as follows:

• Load the packed array into reminiscence
• Unpack the reminiscence and upcast it to FP32 for accumulation duties
• Carry out the computation

The reader ought to observe that when performing the computation, upcasting is used to stop overflows. Subsequently, from a computational perspective, there isn’t any benefit. Nonetheless, from the angle of bandwidth, we’re loading reminiscence twice or 4 instances with out compromising the bandwidth. 

Triton Kernel Implementation (pseudocode)

@triton.jit
def gemv_fp8_kernel(packed_matrix_ptr, packed_vector_ptr, out_ptr): 
    # Get present row to course of
    row_id = get_program_id()
    
    # Initialize accumulator for dot product
    accumulator = 0
    
    # Iterate over row in blocks
    for every block in row:
        # Load packed FP32 values (every accommodates 4 FP8s)
        packed_matrix = load(packed_matrix_ptr)
        packed_vector = load(packed_vector_ptr)
        
        # Unpack the FP32 into 4 FP8 values
        m_a, m_b, m_c, m_d = unpack_fp8(packed_matrix)
        v_a, v_b, v_c, v_d = unpack_fp8(packed_vector)
        
        # Upcast to FP32 and compute partial dot merchandise
        accumulator += (m_a * v_a) + (m_b * v_b) + (m_c * v_c) + (m_d * v_d)
    
    # Retailer ultimate consequence
    retailer(out_ptr, accumulator)

Outcomes

{Hardware}: NVIDIA GeForce RTX 3050 6GB VRAM

CUDA Model: 13.0

Python Model: 3.13.9

GEMV Benchmark (M = 16384, N = 16384) (MxN matrix)

The theoretical efficiency increase that may be achieved is 4x; 3.3x is superb compared, with the remaining overhead primarily stemming from pack/unpack operations and kernel launch prices. 

E5M2 is quicker than E4M3 as a result of simpler unpacking, however E4M3 presents higher precision. Nonetheless, it’s considerably extra advanced to unpack (Feather makes use of a separate GPU kernel to unpack the E4M3 format).

Flash Consideration Benchmark (Sequence Size = 8192, Embedding Dimension = 512)

Accuracy & Precision

Testing with random matrices (integer distributions within the vary [-3, 3] and normal regular distributions) exhibits that each E4M3 and E5M2 keep numerical outcomes inside sensible tolerances for deep studying operations. The buildup errors stay manageable for typical workload sizes; nevertheless, customers requiring strict numerical precision ought to validate their particular use case.