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

{Hardware} Stage 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 latest ones. If you’re on an RTX-30 or RTX-20 sequence GPU from Nvidia, then you definitely will be unable to benefit from this decrease precision FP8 kind. That is precisely the issue that Feather makes an attempt to unravel.

Packing Technique

Utilizing bitwise operators, one can simply pack the FP16 kind 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 one in all 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 extensively used codecs – E5M2 & E4M3. One can’t 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 models and mix them utilizing the bitwise OR operator.

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

Packing an FP8-E4M3 will not be as simple and easy 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 offers help for generally used FP8 requirements, akin to E5M2 and E4M3 casting, for NumPy arrays. Utilizing the identical astype operate, we are able to carry out casting simply as we did for FP16 varieties. The Algorithm is precisely an identical to how we pack FP16, so I’m skipping it right here. 

Triton GPU Kernels

After we pack, we want an algorithm (kernel) to utilise this packed datatype and carry out the computation. Passing the packed datatype to a kernel applied for FP32 or FP64 will end in undefined computation as a result of we’ve already corrupted the FP32 or FP64 being handed. Writing a kernel that takes the packed datatype as enter in CUDA will not be a simple job 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 put in writing 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 notice that when performing the computation, upcasting is used to stop overflows. Subsequently, from a computational perspective, there is no such thing as a benefit. Nevertheless, from the attitude 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 incorporates 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 closing end result
    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 enhance that may be achieved is 4x; 3.3x is excellent as 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 affords higher precision. Nevertheless, it’s considerably extra complicated 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 customary 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; nonetheless, customers requiring strict numerical precision ought to validate their particular use case.