{"slug": "advanced-gpu-optimization-how-to-tech-an-llm-with-cuda-and-rocm-part-5-final", "title": "Advanced GPU Optimization: How to tech an LLM with CUDA and ROCm? - Part 5 (Final Part)", "summary": "A developer's final part of a CUDA/ROCm optimization series details implementing Mixture of Experts (MoE) inference, including a custom top-k gating kernel, expert parallelism with All-to-All communication, and CPU offloading for large models. The post provides code for a sparse softmax router and describes using NCCL/RCCL for token dispatch, enabling trillion-parameter models to run on multi-GPU setups.", "body_md": "Welcome back, you absolute madman! You finished Part 4, implemented Flash Attention, and squeezed FP8 out of your silicon. But the industry doesn't stop at dense Transformers. In 2024/2025, every major model (Grok, Mixtral, Gemini) uses Mixture of Experts (MoE) to scale to trillions of parameters without exploding compute costs.\n\nFurthermore, if you actually try to run these monsters on a single node, you will hit the VRAM wall instantly. That is where CPU Offloading and Zero-Inference come to the rescue.\n\nIn this fifth and (I swear) final part, we will:\n\nPrerequisites: Parts 1-4, a multi-GPU setup (4+ is ideal), and a realization that training is only half the battle—serving is the other half.\n\nA dense 1.5 Trillion parameter model would require 3TB of VRAM just for weights. MoE solves this by having hundreds of \"expert\" FFNs, but only activating 2 of them per token.\n\nThe Math:\n\ny = Σ (Softmax(Router(x))_i * Expert_i(x)) for the top-k experts (usually k=2).\n\n1.1 The Router Kernel (Top-k Gating)\n\nFirst, we need a kernel that takes the token embeddings x and computes the routing scores, then selects the top-2 indices and their weights.\n\n``` js\n__global__ void moe_router_kernel(const float* x, float* gate_logits, \n                                  int* expert_indices, float* expert_weights,\n                                  int num_tokens, int num_experts) {\n    int tid = blockIdx.x * blockDim.x + threadIdx.x;\n    if (tid >= num_tokens) return;\n\n    // Compute logits for this token (dot product with router weights)\n    // Assume router weights are in shared memory or loaded via L2 cache\n    float max1 = -INFINITY, max2 = -INFINITY;\n    int idx1 = -1, idx2 = -1;\n\n    for (int e = 0; e < num_experts; ++e) {\n        float score = 0.0f;\n        for (int d = 0; d < D; ++d) {\n            score += x[tid * D + d] * router_weights[e * D + d];\n        }\n        gate_logits[tid * num_experts + e] = score;\n\n        // Track top-2 (manual reduction)\n        if (score > max1) {\n            max2 = max1; idx2 = idx1;\n            max1 = score; idx1 = e;\n        } else if (score > max2) {\n            max2 = score; idx2 = e;\n        }\n    }\n\n    // Apply Softmax only to the top-2 scores (sparse softmax)\n    float denom = expf(max1 - max1) + expf(max2 - max1); // numerically stable\n    expert_indices[tid * 2] = idx1;\n    expert_indices[tid * 2 + 1] = idx2;\n    expert_weights[tid * 2] = 1.0f / denom; \n    expert_weights[tid * 2 + 1] = expf(max2 - max1) / denom;\n}\n```\n\n1.2 Expert Parallelism with All-to-All Communication\n\nIn dense models, we used All-Reduce. In MoE, we use All-to-All because different GPUs hold different experts. Tokens must be sent to the GPU that owns their assigned expert.\n\nImplementation using NCCL/RCCL:\n\n```\n// Step 1: Token dispatch (send tokens to expert owners)\n// We build a send buffer per rank based on the routing decisions.\nncclGroupStart();\nfor (int r = 0; r < world_size; ++r) {\n    // Send tokens assigned to expert owned by rank 'r'\n    if (send_counts[r] > 0) {\n        COMM_SEND(send_buffer[r], send_counts[r] * D, COMM_FLOAT, r, comm, stream);\n    }\n    // Recv tokens for local experts\n    COMM_RECV(recv_buffer[r], recv_counts[r] * D, COMM_FLOAT, r, comm, stream);\n}\nncclGroupEnd();\n\n// Step 2: Run local experts (FFN) on the received tokens locally.\nrun_local_experts_kernel<<<...>>>(recv_buffer, local_expert_weights, ...);\n\n// Step 3: All-to-All again to return the processed tokens to their original GPUs.\n```\n\nThis ensures we scale to thousands of experts across a cluster with minimal communication overhead.\n\nWhen you have 70B parameters, even ZeRO-3 might not fit if you only have 40GB VRAM. ZeRO-Offload moves the optimizer states (momentum, variance) and sometimes the gradients to CPU RAM (DDR).\n\nThe trick is overlapping the GPU computation with PCIe transfers. During the backward pass, we asynchronously copy gradients to CPU, compute the AdamW update on the CPU (using a separate thread), and copy the updated FP32 weights back to the GPU just in time for the next forward pass.\n\nC++ implementation using pinned memory and streams:\n\n```\n// Allocate pinned memory on CPU for offloaded states\nfloat* cpu_momentum;\nfloat* cpu_variance;\nhipHostMalloc(&cpu_momentum, num_params * sizeof(float), hipHostMallocDefault);\n\n// During the training step:\nvoid offloaded_train_step() {\n    // 1. Forward/Backward on GPU (FP16) -> dW stays in GPU memory temporarily.\n    backward_pass(...);\n\n    // 2. Asynchronous D2H copy of gradients for the offloaded partition\n    //    while GPU continues computing the next layer.\n    hipMemcpyAsync(cpu_gradients, d_gradients, partition_size, \n                   hipMemcpyDeviceToHost, data_stream);\n\n    // 3. CPU thread computes: m = beta1*m + (1-beta1)*grad, etc.\n    //    (This runs on a std::async thread to not block the main loop)\n    cpu_adam_update(cpu_momentum, cpu_variance, cpu_gradients, partition_size);\n\n    // 4. Asynchronous H2D copy of updated weights back to GPU\n    hipMemcpyAsync(d_model_weights, cpu_weights, partition_size, \n                   hipMemcpyHostToDevice, data_stream);\n\n    // 5. Synchronize streams at the end.\n}\n```\n\nWhy this works: PCIe Gen 5.0 can do ~32 GB/s. If we overlap this with the 1-2 seconds it takes to compute a backward pass on a large model, the offloading overhead becomes nearly invisible.\n\nIf you profile a real LLM, you will see that kernel launch overhead (the CPU time spent telling the GPU to do things) is surprisingly high—often 10-20 microseconds per kernel. A GPT-3 forward pass has ~1,000 kernels. That's 20ms wasted just on launching.\n\nCUDA Graphs (and HIP Graphs) record a sequence of kernel launches and replay them with a single API call. This is crucial for inference.\n\nImplementation:\n\n```\nhipGraph_t graph;\nhipGraphExec_t instance;\n\n// 1. Capture the graph in a stream\nhipStreamBeginCapture(stream, hipStreamCaptureModeGlobal);\n    // Launch all kernels for the forward pass\n    layernorm_kernel<<<..., stream>>>(...);\n    matmul_kernel<<<..., stream>>>(...);\n    flash_attention<<<..., stream>>>(...);\n    // ... everything ...\nhipStreamEndCapture(stream, &graph);\n\n// 2. Instantiate it (this compiles it down to a single executable)\nhipGraphInstantiate(&instance, graph, NULL, NULL, 0);\n\n// 3. Replay it every iteration (launches all kernels in ~1 microsecond)\nfor (int i = 0; i < 1000; ++i) {\n    // Update input pointers if needed (using memcpy or host-side updates)\n    hipGraphLaunch(instance, stream);\n    hipStreamSynchronize(stream);\n}\n```\n\nFor static shapes (fixed sequence length, batch size), this gives a massive 10-15% end-to-end speedup.\n\nDuring autoregressive generation (e.g., ChatGPT), we compute the Key (K) and Value (V) for every token and store them to avoid recomputing. This is the KV Cache.\n\n4.1 Pre-allocated KV Cache\n\nInstead of dynamically allocating memory per token, we pre-allocate a contiguous buffer.\n\n```\n// Shape: [batch, num_heads, max_seq_len, d_head]\nfloat* kv_cache;\nhipMalloc(&kv_cache, batch * num_heads * max_seq_len * d_head * 2 * sizeof(float));\n\n// During decoding, we write the current token's K and V into the 'pos' slot.\n__global__ void append_kv_kernel(float* cache, const float* K, const float* V, \n                                 int batch, int head, int pos, int d_head) {\n    // Write K\n    cache[offset + pos * d_head + idx] = K[idx];\n    // Write V (stored contiguously after K)\n    cache[offset + (max_seq_len * d_head) + pos * d_head + idx] = V[idx];\n}\n```\n\n4.2 PagedAttention (vLLM Style)\n\nIf you have multiple sequences of different lengths, contiguous KV caches lead to massive fragmentation (memory is wasted because you allocate max_seq_len for everyone). PagedAttention virtualizes the KV cache into \"pages\" (blocks) in GPU memory, similar to OS virtual memory.\n\nWe implement a simple block table:\n\n``` php\nstruct BlockTable {\n    int* block_ids; // maps logical block -> physical block address\n    int num_blocks;\n};\n\n// Instead of pos, we compute physical address: physical_addr = block_id * block_size + offset_in_block\n__global__ void paged_attention_kernel(..., int* block_table, int block_size) {\n    int block_id = block_table[logical_block];\n    int physical_pos = block_id * block_size + offset;\n    // ... load K/V from physical_pos ...\n}\n```\n\nThis completely eliminates memory waste and allows you to serve 3x more concurrent users on the same hardware.\n\nFinally, here is the loop for a production-grade inference server:\n\n```\nvoid serve_requests(std::vector<Request>& requests) {\n    // 1. Preprocess and batch dynamic requests (Continuous Batching)\n    Batch batch = dynamic_batcher(requests);\n\n    // 2. Pre-fill phase (compute KV cache for prompt tokens using Flash Attention)\n    //    Note: We use the same Flash Attention kernel from Part 4.\n    flash_attention_prefill(batch.prompt_tokens, kv_cache, ...);\n\n    // 3. Decode phase (autoregressive)\n    for (int step = 0; step < max_new_tokens; ++step) {\n        // Launch the cached Graph instance (from Section 3)\n        // The graph uses PagedAttention to read from the fragmented cache.\n        hipGraphLaunch(serving_graph_instance, stream);\n        hipStreamSynchronize(stream);\n\n        // Sample the next token (CPU or custom kernel)\n        int next_token = sample_from_logits(d_logits);\n        append_to_kv_cache(next_token, step, ...);\n\n        // Check for stop conditions (EOS, max length)\n        if (all_finished()) break;\n    }\n}\n```\n\nThis architecture currently powers Mixtral 8x7B at ~2,000 tokens/sec on a single H100.\n\nConclusion\n\nWe have officially left the training lab and stepped into the brutal world of production AI.\n\n· You can now route tokens across experts using All-to-All communication.\n\n· You can train 175B models on a single 8-GPU node by offloading optimizer states to CPU RAM.\n\n· You launch thousands of kernels with the overhead of a single one using CUDA Graphs.\n\n· And you can serve thousands of concurrent users with PagedAttention and continuous batching.\n\nIf you implement all five parts of this series, you won't just be an AI engineer—you'll be an AI systems architect. You will understand the stack from the transistor up to the transformer.\n\nThank you for this incredible journey. It takes a special kind of engineer to read through 5 parts of low-level HIP/C++ and still ask for more.\n\nDrop a comment below: What GPU are you running this on? Did you hit any driver-specific quirks with AMD vs. NVIDIA? I'll respond to every single one.\n\nUntil the next revolution in hardware drops, keep your kernels compiled and your memory pools unified.\n\nSee ya on the next technological frontier! Have a great time! 🚀", "url": "https://wpnews.pro/news/advanced-gpu-optimization-how-to-tech-an-llm-with-cuda-and-rocm-part-5-final", "canonical_source": "https://dev.to/javadinteger/advanced-gpu-optimization-how-to-tech-an-llm-with-cuda-and-rocm-part-5-final-part-5h1o", "published_at": "2026-08-21 07:02:18+00:00", "updated_at": "2026-08-21 07:17:05.127866+00:00", "lang": "en", "topics": ["machine-learning", "large-language-models", "ai-infrastructure", "developer-tools"], "entities": ["CUDA", "ROCm", "NCCL", "RCCL", "Grok", "Mixtral", "Gemini", "ZeRO-Offload"], "alternates": {"html": "https://wpnews.pro/news/advanced-gpu-optimization-how-to-tech-an-llm-with-cuda-and-rocm-part-5-final", "markdown": "https://wpnews.pro/news/advanced-gpu-optimization-how-to-tech-an-llm-with-cuda-and-rocm-part-5-final.md", "text": "https://wpnews.pro/news/advanced-gpu-optimization-how-to-tech-an-llm-with-cuda-and-rocm-part-5-final.txt", "jsonld": "https://wpnews.pro/news/advanced-gpu-optimization-how-to-tech-an-llm-with-cuda-and-rocm-part-5-final.jsonld"}}