NVIDIA Blackwell UltraによるSoftmaxの効率化

LLM context lengths are exploding, and architectures are moving toward complex attention schemes like Multi-Head Latent Attention (MLA) and Grouped Query Attention (GQA). As a result, AI ”speed of thought” is increasingly governed not by the massive throughput of matrix multiplications, but by the transcendental math of the softmax function.

Transcendentals refer to functions that cannot be expressed as the root of a polynomial equation with rational coefficients. Subsequently, they “transcend” basic algebraic operations like addition and multiplication—the exact operations Tensor Cores excel at. In the specific context of softmax, the most computationally expensive of these transcendentals is the natural exponential function that is executed on Special Function Units (SFUs). In NVIDIA assembly instructions (SASS), this function is invoked via the MUFU.EX2 instruction. This architectural split creates a softmax bottleneck within the attention block, when powerful matrix engines are forced to idle while waiting for the SFU datapaths to normalize attention scores.

NVIDIA Blackwell Ultra alleviates this bottleneck by doubling SFU throughput over the standard NVIDIA Blackwell architecture.

This blog dives into the mechanics of softmax within the attention loop, explores how Blackwell Ultra’s hardware optimizations eliminate pipeline stalls, and provides a benchmark for you to measure the raw MUFU.EX2

How attention works

A foundational component of modern large language models is the attention mechanism, which allows a model to dynamically transform static token vectors into dynamic, context-aware representations. At its core, it is a process of re-weighting information by allowing tokens to adjust their importance to one another. To facilitate this interaction, every token in a sequence is projected into three functional roles:

Query: Represents what the current token is seeking to understand its own context.

Key: Represents a token’s profile that others use for matching. Tokens previous in the sequence have keys that signal their specific relevance to the query.

Value: This holds the actual informational content. Once a match is confirmed between a query and a key, the Value is the specific data that is transferred to the original token.

Figure 1 below shows attention in action. We have two sentences that utilize the word “dog” in two different definitions. Initially, we can see that the embeddings (the numerical vectors that capture meaning and nuance in a multidimensional space) of both “dog” mentions are identical.