Most advice on scaling LLMs comes down to buying more VRAM. That works until the last step of training, the cross-entropy loss, where the memory cost stops scaling with your model and starts scaling with your vocabulary. It is where the out-of-memory errors show up, usually after everything else has already run fine. The cause is the logit tensor and the fix is fused kernels.
The logit bottleneck
Llama 3 has a vocabulary of 128,256 tokens. Predicting the next token means projecting the hidden state into that entire space, and even at a modest batch size the intermediate logit tensor can pass 80GB of VRAM. The waste is worse than the number suggests: you write billions of logits out to VRAM and then read them straight back to compute the loss.
I have spent years optimizing WooCommerce checkouts and custom APIs, and the rule there is the same one: cut the round trips. On a GPU you do that by fusing operations into a single kernel. I covered a related idea in my post on mastering FP8 performance.
Why fused kernels help
A fused kernel puts the linear projection and the cross-entropy calculation in one GPU program, so the full logit matrix never lands in global memory. Tiling does the arithmetic in blocks small enough to sit in registers, which takes the O(V) term out of the memory cost. Triton’s fused softmax tutorial is the clearest place to see the pattern.
The naive approach
# This is what PyTorch does by default
logits = input @ weight.T # Huge memory allocation here
loss = cross_entropy(logits, targets)The fused Triton version
The custom kernel computes the target logit and the log-sum-exp (LSE) as it goes, walking the hidden and vocabulary dimensions in tiles. Online softmax keeps the numbers stable without ever holding the full tensor, and that is the same trick behind the memory savings in libraries like Unsloth.
@triton.jit
def bbioon_fused_linear_ce_kernel(X_ptr, W_ptr, Y_ptr, LSE_ptr, ...):
# Program ID handles one row (token) at a time
row_idx = tl.program_id(0)
# Tiled Dot Product for Target Logit
# We only load small blocks into registers
for v_idx in range(0, V, V_BLOCK):
w_tile = tl.load(W_ptr + col_offsets)
logits = tl.dot(x_tile, w_tile)
# Online Max and Sum updates happen here
# No massive O(V) matrix in VRAMA war story: the atomic traffic jam
My first version leaned on tl.atomic_add for the backward pass gradients. On an A100 that turned into a traffic jam: thousands of threads hammering the same weight gradient addresses at once. The hardware serializes those updates, so my “optimized” kernel ended up slower than plain PyTorch. Fusing the forward pass is the easy half. The backward pass usually needs a dedicated kernel for the weight gradients and a grid layout that keeps threads off each other’s addresses.
If kernel work like this is eating your dev hours, I can take it off your plate. I have been building WordPress and AI integrations for years, and I know where these bottlenecks tend to hide.
What this buys you
On consumer hardware, custom Triton kernels are closer to a requirement than an optimization. Cutting peak memory by 84% lets you raise the batch size and get through training faster, and doing the backward pass math yourself keeps the gradients precise. The Liger-Kernel implementation is worth reading as a production reference. Profile before you rewrite anything: the bottleneck is rarely where you assume it is.
