LLM inference, fused kernels and GPU networking

Modern large-scale AI inference combines orchestrated multi-GPU execution pipelines, SRAM-fused Triton/CUDA kernels, and ultra-low-latency fabric interconnects to generate tokens in milliseconds.

1. End-to-End Inference Walkthrough: From Prompt to Token

Suppose a user sends a 1,000-token prompt to a 70B-parameter LLM running across 4 GPUs using Tensor Parallelism inside an inference engine like vLLM.

• Step 1: Tokenization & Host-to-Device Transfer (CPU → GPU).

• The CPU tokenizer converts text into an array of 1,000 integer token IDs.
• The serving engine allocates non-contiguous physical memory blocks in GPU HBM via PagedAttention and copies the token IDs across the PCIe bus into GPU 0's VRAM.

• Step 2: Embedding Lookup & Broadcasting (NVLink).

• GPU 0 looks up the 1,000 embedding vectors (e.g., shape [1000, 8192]) and broadcasts the tensor across NVLink so all 4 GPUs hold identical starting activations.

• Step 3: Prefill Phase (Compute-Bound Tensor Core Execution).

• Inside each of the 80 Transformer layers, FlashAttention computes self-attention across all 1,000 prompt tokens simultaneously. The resulting Key (K) and Value (V) tensors are written into each GPU's PagedAttention KV cache blocks.
• Because the Feed-Forward Network (FFN) weights are split column-wise across the 4 GPUs (Tensor Parallelism), each GPU multiplies the [1000, 8192] activations against its local weight slice using FP8/FP4 Tensor Cores.
• At the end of the layer, the 4 GPUs perform an NVLink All-Reduce to sum their partial results, leaving every GPU with the full updated [1000, 8192] tensor before entering the next layer.

• Step 4: Sampling & First Token Output.

• After Layer 80, the final hidden state of token 1,000 is multiplied by the vocabulary matrix to produce logits over ~128,000 vocabulary words. A sampling kernel picks the first output token and sends it back to the CPU stream.

• Step 5: Autoregressive Decode Loop (CUDA Graphs + Memory-Bound Execution).

• To generate token 1,001, the CPU does not relaunch hundreds of individual kernels. Instead, it triggers a pre-recorded CUDA Graph in one API call.
• Now processing only a single [1, 8192] token vector, each GPU streams its local model weights and all 1,000 cached KV vectors from HBM into on-chip SRAM, computes the new attention output, executes the NVLink All-Reduce, and appends token 1,001's K and V vectors into the next free PagedAttention slot. This loop repeats until an End-of-Sequence (EOS) token is generated.

2. Building a Custom Fused Softmax Kernel in OpenAI Triton

A naive PyTorch Softmax (torch.exp(x) / torch.sum(torch.exp(x))) requires three separate HBM memory passes (one to find the row max for numerical stability, one to compute exponentials and sum them, and one to divide).

Below is a complete, production-style Fused Row-Wise Softmax kernel in OpenAI Triton. It loads an entire row into fast on-chip SRAM once, performs all three steps inside GPU registers, and writes the final output back to HBM in a single memory pass.

import torch
import triton
import triton.language as tl

@triton.jit
def fused_softmax_kernel(
    output_ptr, input_ptr,
    input_row_stride, output_row_stride,
    n_cols,
    BLOCK_SIZE: tl.constexpr
):
    # Each program instance processes one row of the matrix
    row_idx = tl.program_id(0)
    
    # Calculate memory pointers for the start of this specific row
    row_start_ptr = input_ptr + row_idx * input_row_stride
    col_offsets = tl.arange(0, BLOCK_SIZE)
    input_ptrs = row_start_ptr + col_offsets
    
    # Mask out-of-bounds threads if n_cols < BLOCK_SIZE (padding with -inf)
    mask = col_offsets < n_cols
    row = tl.load(input_ptrs, mask=mask, other=-float('inf'))
    
    # Step 1: Subtract row max inside fast SRAM/registers for numerical stability
    row_minus_max = row - tl.max(row, axis=0)
    
    # Step 2: Compute exponentials and sum in registers (zero HBM traffic)
    numerator = tl.exp(row_minus_max)
    denominator = tl.sum(numerator, axis=0)
    
    # Step 3: Normalize and write final probabilities back to HBM once
    softmax_output = numerator / denominator
    output_row_start_ptr = output_ptr + row_idx * output_row_stride
    tl.store(output_row_start_ptr + col_offsets, softmax_output, mask=mask)

def triton_softmax(x: torch.Tensor) -> torch.Tensor:
    n_rows, n_cols = x.shape
    # Round block size up to the next power of 2 to fit the entire row in SRAM
    BLOCK_SIZE = triton.next_power_of_2(n_cols)
    y = torch.empty_like(x)
    
    # Launch 1 Triton program block per matrix row
    fused_softmax_kernel[(n_rows,)](
        y, x,
        x.stride(0), y.stride(0),
        n_cols,
        BLOCK_SIZE=BLOCK_SIZE
    )
    return y

3. AI Cluster Networking: NVLink, InfiniBand, and RoCEv2

When scaling AI workloads across thousands of GPUs, network latency and congestion become just as critical as raw GPU FLOPS. Modern clusters use a two-tier scale-up and scale-out network hierarchy.

• How Collective Operations Use the Network.

• All-Reduce (Gradient & Tensor Sync): Every GPU starts with a partial tensor, sums all tensors across the cluster, and ends with the identical final sum. Using Ring or Tree algorithms—or In-Network Computing (SHARP) where switches add vectors directly inside switch silicon—bandwidth scales independently of cluster size.
• All-to-All (Mixture-of-Experts Routing): Unlike All-Reduce, where every GPU sends the same data to neighbors, All-to-All creates an all-to-all traffic burst where every GPU simultaneously sends unique token slices to every other GPU holding target experts. This causes severe incast congestion at receiver ports.

• Why AI Networks Require Lossless RDMA.

• Traditional TCP/IP networking relies on the CPU to process packets and drops packets when buffers fill up, triggering millisecond-scale retransmissions that stall an entire 10,000-GPU training cluster.
• Both InfiniBand and RoCEv2 use RDMA (Remote Direct Memory Access), allowing one GPU's Network Interface Card (NIC) to read/write directly into another GPU's HBM VRAM with zero CPU involvement. RoCEv2 enforces zero packet loss over standard Ethernet using PFC (Priority Flow Control) to pause upstream senders before switch buffers overflow and ECN (Explicit Congestion Notification) to throttle transmission rates smoothly.