使用 JAX 和 XLA 优化推理工作负载的低延迟通信

// Fused One shot All Reduce + Root Mean Square Normalization kernel 

// peer_comm_buffer: 
//        -> thread_offsets
// 	|  [00 01 02 03,...]
// 	v  [10 11 12 13,....]
//   ranks
//
// Outputs : [00+10.., 01+11.., 02+12..., 03+13..]
__global__ void OneShotARNormKernel(std::vector<T*> peer_comm_buffer_ptrs, T* sum_vec, T* weight_buffer, float eps, int hidden_size)
{

for (int ii = 0; ii < NUM_RANKS; ++ii)
{
    // One-shot All Reduce sum
    sum_vec = add(sum_vec, peer_comm_buffer_ptrs[ii][thread_offset]);
}

......

// All Reduce and Norm fusion
// Compute x^2
squares = compute_square(sum_vec);
// Compute sum(x^2) across grid
summed_squares = block_reduce_sum(squares);
// Compute RMS denominator
float denom = __fsqrt_rn(__fdividef(summed_squares, hidden_size) + eps);
// Load per-element affine param if necessary
if learnable_affine
{
    weight_vec = weight_buffer + thread_offset;
}
// rms_norm = (sum_vec / (denom)) * weight_vec
sum_vec = rms_norm(denom, sum_vec, weight_vec);
}

// Define the C++ custom call to be invoked from JAX FFI

#include "xla/ffi/api/c_api.h"
#include "xla/ffi/api/ffi.h"

// Global ptr vector shared amongst all ranks or threads
std::vector<void*> peer_mem_ptr(MAX_RANKS);

void AllReduceDispatcher(...) {
 // Each thread/rank in the node populates its input buffer to global ptr map
 peer_mem_ptr[rank_id] = input;
 // Synchronizing all devices so the pointer vector is fully populated
 Barrier();
 
 // Launch kernel with peer_mem_ptr
 OneShotARNormKernel<<<...>>>(peer_mem_ptr, ...)
}

// Define fusion kernel launcher func
ffi::Error customAllReduce(cudaStream_t stream,
                       ffi::AnyBuffer input,
                       ffi::AnyBuffer weight_buffer,
                       ffi::Result<ffi::AnyBuffer> sum_vec,
                       int hidden_size,
                       float eps,
                       int rank_id)
{
    // Launch dispatcher...
    AllReduceDispatcher(...)
}


// Create symbol ArNorm with C linkage that can be loaded using Python ctypes
XLA_FFI_DEFINE_HANDLER_SYMBOL(
    ArNorm, customAllReduce,
    ffi::Ffi::Bind()
        .Ctx<ffi::PlatformStream<cudaStream_t>>()  // stream
        .Arg<ffi::AnyBuffer>()    // input
        .Arg<ffi::AnyBuffer>()    // weight_buffer
        .Ret<ffi::AnyBuffer>()    // sum_vec
        .Attr<int>("hidden_size")
        .Attr<float>("eps")
        .Attr<int>("rank_id")); 

# Invoke custom call in JAX application using JAX FFI

from jax.lib import xla_client
import ctypes
# Load the library built from the c++ funtion
SHARED_LIBRARY = os.path.join(os.path.dirname(__file__), "libcustom_ar_kernel.so")

library = ctypes.cdll.LoadLibrary(SHARED_LIBRARY)
# Register the FFI function in jax
XLA_CUSTOM_CALL_TARGET_AR_NORM = "ar-norm"
xla_client.register_custom_call_target(name=XLA_CUSTOM_CALL_TARGET_AR_NORM,
                                       fn=ffi.pycapsule(library.ArNorm),
                                       platform=XLA_PLATFORM,
                                       api_version=XLA_CUSTOM_CALL_API_VERSION)

# Invoke the custom call
output = ffi.ffi_call(
    XLA_CUSTOM_CALL_TARGET_AR_NORM,
    jax.ShapeDtypeStruct(input.shape, input.dtype),  # output type
    input,                                           # input buffer
    weight_buffer,                                   # weight buffer
    hidden_size=hidden_size,
    eps=eps,
    rank_id=rank_id)

// Creates symbol ArNorm with C linkage that can be loaded using Python ctypes
XLA_FFI_DEFINE_HANDLER_SYMBOL(
    ArNorm, customAllReduce,
    ffi::Ffi::Bind()
        .Ctx<ffi::PlatformStream<cudaStream_t>>()  // stream
        .Arg<ffi::AnyBuffer>()   			    // input
        .Arg<ffi::AnyBuffer>()                     // weight_buffer
        .Ret<ffi::AnyBuffer>()                     // sum_vec
        .Attr<int>("hidden_size")
        .Attr<float>("eps")
        .Attr<int>("rank_id"),
        {xla::ffi::Traits::kCmdBufferCompatible});