Memory-Efficient LM-Head Inference for LLaDA

A few weeks ago, Udula and I came across the original paper on LLaDA. Starting with a sequence of masks, tokens are revealed over several denoising steps. Naturally, I wanted to look under the hood and see what LLaDA was actually doing during inference. I decided to implement the decoder, profile it, and work through the bottlenecks.

After reading the original paper, I decided to target the Language Model (LM) head. The LM head projects each final hidden state across LLaDA's 126,464-token vocabulary, producing a large tensor of logits. However, the decoder ultimately keeps only a predicted token, its confidence, and a few summary statistics, so almost all of that tensor is temporary.

That observation led to the work in this post: fusing the LM-head projection with the sampling reduction so the full logits tensor never needs to be materialized.

This post assumes basic familiarity with transformers.

How LLaDA Decoding Works

LLaDA is a masked diffusion language model: instead of generating one token at a time from left to right, as autoregressive models like GPT do, it fills in a masked sequence over several denoising steps. Because attention is bidirectional, each position can attend to tokens on both sides.

For conditional generation, LLaDA keeps the prompt fixed and initializes the generation region with mask tokens. It then repeatedly processes the full current sequence and predicts replacements for masked positions. In the zero-temperature configuration used here, decoding proceeds block by block. At each step, a scheduled number of masked positions in the active block are committed according to the softmax probability of their predicted token.

The LM-Head Bottleneck

The transformer produces a hidden vector for each position. LLaDA-8B's LM head then maps each 4,096-dimensional hidden vector across a vocabulary of 126,464 tokens:

where $X \in \mathbb{R}^{R \times 4096}$ contains the hidden states being scored, $W \in \mathbb{R}^{126464 \times 4096}$ is the output weight matrix, and $Z \in \mathbb{R}^{R \times 126464}$ is the logits tensor.

Here, $R$ is the number of positions sent through the LM head. Scoring a full batch of four 4,096-token sequences gives $R=16{,}384$.

A conventional implementation materializes all of $Z$ before launching separate operations for argmax, logsumexp, and confidence. The decoder ultimately keeps only a few values per row, but the intermediate tensor scales with $R \times V$.

At $R=16{,}384$, the fp16 logits alone occupy 3.86 GiB. PyTorch's logsumexp may create an additional 7.72 GiB fp32 working buffer, bringing these two temporary allocations to approximately 11.58 GiB.

After the reduction, however, the decoder retains only 16 bytes per row: one int64 token index and two fp32 values for logsumexp and confidence.

Scoring Only The Necessary Positions

Before applying the LM head, an optimized decoder should gather only the hidden states corresponding to unresolved positions in the active block. This reduces the row dimension $R$, often from thousands of sequence positions to only a few dozen.

Gathering and fusion solve different parts of the problem. Gathering reduces how many rows are projected over the vocabulary; the Fused LM Head avoids materializing and rereading the resulting $R \times V$ logits tensor. The kernels below therefore operate after masked-position gathering, optimizing the vocabulary projection that remains.

The Fused Sampler

When logits are materialized, the sampler must still compute argmax, logsumexp, and confidence. PyTorch typically does this with separate operations and may create an fp32 copy for logsumexp.

The Fused Sampler computes all three in one streaming pass, reaching 95.2% of the RTX 4070 Ti's memory-bandwidth roof. For each row, it tracks the maximum logit, its lowest vocabulary index, and the softmax denominator.

The larger memory saving comes from fusing this reduction directly into the LM-head matrix multiplication.

The Fused LM Head

The Fused LM Head combines the output projection with the sampler reduction. Each on-chip tile of vocabulary logits is immediately folded into the online-softmax state and discarded, so the full $R \times V$ tensor is never written to global memory.

For fp16 and bf16 inputs, Triton uses tensor-core matrix multiplication with fp32 accumulation. The kernel tracks a running maximum $m$ and denominator $d$. For a tile with local maximum $\tilde m$ and denominator $\tilde d$, the states merge as:

Both correction exponents are non-positive, which keeps the update numerically stable as additional vocabulary tiles are processed. Once the full vocabulary has been reduced,

and the maximum softmax probability is

Confidence therefore requires no additional vocabulary pass.

The core Triton reduction is short:

t_max = tl.max(acc, axis=1)
t_idx = tl.min(
    tl.where(acc == t_max[:, None], cols[None, :], V),
    axis=1,
)
t_d = tl.sum(tl.exp(acc - t_max[:, None]), axis=1)
 
new_m = tl.maximum(m, t_max)
d = d * tl.exp(m - new_m) + t_d * tl.exp(t_max - new_m)
idx = tl.where(t_max > m, t_idx, idx)  # Earlier tile keeps ties.
m = new_m

For the small number of positions typically scored during decoding, the vocabulary is divided across several parallel workers. Each processes part of the vocabulary, and their results are combined in a short final step.

The implementation consistently chooses the lowest-index token when scores are tied, matching torch.argmax. It was tested across different shapes, edge cases, and full decoding runs. In the real-model benchmark, the fused and standard implementations produced the same output tokens.

The underlying fusion strategy is established work. Cut Cross-Entropy (arXiv:2411.09009) and Liger (arXiv:2410.10989) fuse the LM head with the training loss. This library applies the same principle to diffusion-model inference, returning argmax, logsumexp, and confidence with an explicit tie-breaking contract.

Results

Benchmarks were run on an RTX 4070 Ti under WSL2.

Full methodology, raw results, and reproduction commands are available in benchmarks/results.md.

During normal decoding, only a small number of unresolved positions need to be scored. At $R=64$, the Fused LM Head runs in 2.287 ms, compared with 2.672 ms for torch.compile, a 1.17× speedup without allocating a temporary logits tensor.

On LLaDA-8B, this produces an end-to-end decoding improvement of about 1.1×.

The memory advantage becomes clearer at larger workloads:

Peak-memory measurements include the roughly 1 GB LM-head weight.

At $R=4{,}096$, the fused implementation matches torch.compile while using about half the memory. At $R=16{,}384$, it is 7% slower but uses 4.5× less memory. It is also substantially faster and more memory-efficient than processing the vocabulary in smaller PyTorch chunks.

When GPU memory is limited to 90% of the available 12 GB, the standard $R=16{,}384$ implementation runs out of memory. The fused version completes using approximately 1.1 GB.

The roofline plot compares measured performance with the GPU's theoretical bandwidth and compute limits. These percentages are estimates based on the amount of work and memory traffic required by the algorithm, rather than direct hardware-counter measurements.

Because PyTorch's memory statistics can include memory paged into system RAM, I also tested both implementations under a fixed GPU-memory limit. The materialized implementation ran out of memory, while the fused implementation completed using approximately 1.1 GB.

CUDA SIMT and WMMA Implementation

The repository includes three hand-written CUDA versions: naive SIMT, shared-memory tiled SIMT, and tensor-core WMMA.

SIMT kernels use groups of 32 threads, called warps, to process different values in parallel. WMMA instead lets a warp cooperatively multiply small matrix tiles on tensor cores. These kernels show the optimization steps behind Triton; at $R=4{,}096$, the WMMA version reaches 54.5% of Triton's performance.

PyTorch stores the LM-head weight as $W[V,H]$, while the multiplication requires $W^\top$. The WMMA kernel obtains this view without creating a transposed copy:

wm::fragment<wm::matrix_a, 16, 16, 16, __half, wm::row_major> fa;
wm::fragment<wm::matrix_b, 16, 16, 16, __half, wm::col_major> fb;
 
// Load the existing W[V,H] tile as column-major, giving WMMA the Wᵀ view.
wm::load_matrix_sync(fb, &ws[wc * 16][kk], W_LD);
 
#pragma unroll
for (int i = 0; i < 2; ++i) {
  // Load a hidden-state tile, multiply it by Wᵀ on tensor cores, and accumulate.
  wm::load_matrix_sync(fa, &hs[wr * 32 + i * 16][kk], W_LD);
  wm::mma_sync(fc[i], fa, fb, fc[i]);
}

Here, one warp loads a tile of hidden states and a tile of LM-head weights, multiplies them on tensor cores, and accumulates a block of logits. The loop handles 32 hidden-state rows in two 16-row WMMA operations.

Conclusion

This project targets the LM head, not the transformer trunk.

Bidirectional attention prevents the straightforward lossless KV-cache reuse used during autoregressive decoding, but it does not make caching impossible. Recent systems recover substantial trunk reuse through approximate caching, selective recomputation, and block-structured decoding.

Fast-dLLM introduces a block-wise approximate KV cache alongside confidence-aware parallel decoding. dLLM-Cache reuses intermediate computation through long-interval prompt caching and adaptive partial updates. EntropyCache uses the entropy of newly decoded token distributions as a signal for deciding when cached states should be refreshed.

These methods attack a different part of the inference budget. Caching reduces how often transformer states must be recomputed. The kernels in this post reduce the memory traffic and temporary storage required whenever positions are scored over the vocabulary. A practical inference stack can apply both.

Entropy-based cache policies also suggest a direct extension to the Fused LM Head. Entropy could be added to the streamed reduction by maintaining one additional rescaled statistic, such as $\sum_i e^{x_i - m} x_i$. Together with $m$ and $d$, this recovers

This kernel does not solve the full cost of diffusion-language-model inference. On large workloads, keeping only the required summary values and avoiding materialization of the logits tensor eliminates gigabytes of temporary memory while retaining comparable performance.

Using the Library

from dllm_kernels import fused_head
 
out = fused_head(hidden, weight)
# HeadOutput(argmax, lse, confidence)

The repository includes:

• the Fused Sampler;
• the Triton Fused LM Head;
• a LLaDA-compatible decode loop combining masked-position gathering with the fused head;
• hand-written CUDA versions of the fused head, including naive, shared-memory tiled, and tensor-core WMMA variants;
• benchmark suites.

See github.com/lsnchow/LLaDa-CUDA-fused-kernel.

Acknowledgements

This work was completed in collaboration with Udula Abeykoon, whose late-night discussions, advice, and support shaped the project throughout.

References

• LLaDA: Large Language Diffusion Models (arXiv:2502.09992)
• Online Normalizer Calculation for Softmax (arXiv:1805.02867)
• Liger Kernel: Efficient Triton Kernels for LLM Training (arXiv:2410.10989)
• Fast-dLLM (arXiv:2505.22618)
• dLLM-Cache (arXiv:2506.06295)
• EntropyCache (arXiv:2603.18489)

All results were measured on an RTX 4070 Ti (12 GB) under WSL2 with PyTorch 2.13 and Triton 3.7.