• Dissecting NVIDIA Blackwell Tensor Cores

Blackwell is not just Hopper with bigger Tensor Cores. The software contract changed. On SM100, kernel performance now depends on explicitly managing Tensor Memory (TMEM), understanding CTA-scoped tensor instructions, and recognizing that the limiting resource is often no longer raw tensor math, but the bandwidth needed to feed it.

This also helps explain why several recent Blackwell-era kernels, including FlashAttention-4 and MXFP8 Training, had to substantially rethink their pipelines instead of simply porting Hopper code.

1. Blackwell changes who “owns” MMA execution

Three architectural shifts matter most.

First, Blackwell introduces Tensor Memory (TMEM), a new software-managed level in the memory hierarchy used to hold MMA accumulators. On Hopper, MMA results were tightly coupled to the issuing warpgroup and register file. On Blackwell, accumulators live in TMEM, which decouples result ownership from any particular thread. That sounds subtle, but it fundamentally changes kernel structure: epilogues, accumulator lifetimes, and overlap strategies all need to be redesigned around TMEM.

Second, tcgen05 instructions are CTA-scoped. A single thread issues the instruction on behalf of the entire CTA. This is a major departure from Hopper’s warpgroup-scoped wgmma model. The practical consequence is that threads are no longer symmetric participants in tensor-core issue. Some threads orchestrate, others move data, and the whole CTA acts as the execution unit.

Third, Blackwell adds TPC-scoped TMA and MMA instructions through cta_group::2 in PTX, or 2CTA in SASS. Two CTAs, spanning two SMs, can collaboratively execute the same tcgen05.mma. Combined with native support for sub-byte and microscaled datatypes, this gives Blackwell a much more flexible tensor-core pipeline, but only if the kernel is explicitly written to use it.

The broader pattern is clear: Blackwell rewards kernels that think in terms of asynchronous clusters and shared on-chip resources, not warp-synchronous loops.

2. Physical topology now leaks into software decisions

The execution hierarchy above the SM matters more than it used to.

CTAs grouped into clusters are guaranteed to co-schedule on the same Graphics Processing Cluster (GPC). That is essential for 2CTA execution and for efficient use of distributed shared memory. But there is a catch: if a persistent kernel launches one CTA per SM and the chosen cluster size does not evenly divide the SMs yielded by each GPC, the leftover CTAs can serialize. The result is that “one CTA per SM” is no longer automatically the right launch policy.

On B200, the topology is even more visible because the package spans two dies. Pointer-chasing measurements that intentionally fill L2 expose an average die-to-die latency penalty of roughly 300 cycles. For kernels relying on cluster-local reuse, this means that topology-blind scheduling can turn a theoretically local path into a very real latency tax.

In other words, Blackwell kernel tuning now includes placement as well as instruction selection.

3. LDGSTS and TMA serve different regimes

Blackwell offers two strong but very different ways to move data into shared memory.

LDGSTS: fast to ramp, fragile at scale

cp.async / LDGSTS writes directly into shared memory without staging through registers, which reduces register pressure and makes it attractive for irregular data movement. Its DRAM throughput saturates at roughly 6.6 TB/s with about 32 KiB in-flight per SM.

The problem is latency scaling. The baseline latency is about 600 ns, but once the in-flight footprint grows, the MIO subsystem becomes the bottleneck:

• At 8 KiB in-flight, latency rises to about 1229 cycles
• At 12 KiB in-flight, latency can spike to about 2177 cycles

So LDGSTS is excellent when the kernel needs responsiveness and flexibility, but it becomes increasingly fragile once too many warps are competing for the copy path.

TMA: higher ceiling, slower to fill

Tensor Memory Accelerator (TMA) is issued by a single thread, handles address generation in hardware, and can perform the swizzling needed by tensor-core layouts asynchronously. Its peak throughput is higher, approaching 7.2 TB/s, but it needs much more data in flight, typically more than 64 KiB, before it reaches that ceiling.

That makes TMA the better fit for large, regular, deeply pipelined tiles, while LDGSTS remains attractive for sparse or irregular patterns.

This tradeoff shows up in real kernels. A reasonable rule of thumb is:

• Use TMA for large, predictable tiles with enough buffering to hide setup cost
• Use LDGSTS for irregular or dynamic page fetches where responsiveness matters more than peak bandwidth

Even within LDGSTS-heavy kernels, adding more stages and more copy-participating threads continues to help until register allocation becomes the limiting factor.

4. Multicast and DSMEM are powerful, but not free

Blackwell’s cluster features are only as good as the access pattern driving them.

TMA multicast and the L2 Request Coalescer

With TMA multicast, a single load can populate shared memory on multiple SMs in a cluster. This is serviced through the L2 Request Coalescer (LRC). There is also an “implicit” form of multicast where multiple CTAs simply request the same data and rely on the hardware to merge requests.

Implicit multicast can reach roughly the same effective shared-memory fill throughput as explicit multicast, but the LRC stops saving much L2 traffic once the implicit requests exceed about 64 bytes in-flight. So if the objective is not just SMEM fill rate but also lower L2 pressure, explicit multicast remains the cleaner tool.

Remote shared memory is not local shared memory

The gap between local and remote shared-memory access is severe:

• Local ld.shared: about 128 B/clk
• Naive remote ld.shared::cluster: about 21 B/clk

The reason is painful but simple: the compiler often lowers remote loads to generic LD instructions rather than optimized LDS instructions. For high-throughput inter-CTA exchange, developers should rely on cp.async.bulk (UBLKCP in SASS), which pushes distributed shared-memory throughput up to about 32 B/clk.

The lesson is that Blackwell’s cluster features are not self-optimizing. The fast path usually has to be spelled out explicitly.

5. The real roofline is often shared-memory bandwidth

One of the most important Blackwell insights is that many MMA instructions are no longer math-bound.

For 1SM MMA, under-sized shapes are heavily penalized:

• M=64 uses only about half the datapath
• M=128 reaches near-full utilization

For 2SM MMA, M=256 is the sweet spot because it maps to 128 rows per SM, which keeps both SMs well utilized.

The deeper issue is operand movement. Blackwell supports:

• SS mode: both A and B come from shared memory
• TS mode: A comes from TMEM, B comes from shared memory

In SS mode, the instruction is entirely bound by shared-memory bandwidth for N < 128.

Consider an FP16 1SM MMA with shape M=128, N=64, K=16:

Assuming Blackwell shared memory sustains 128 B/clk, the shared-memory service time is:

If the effective Tensor Core throughput for this instruction regime is 8,192 FLOPs/clk, the math time is:

So shared memory still dominates:

That is the core result. For N=64, the instruction is physically SMEM-bound, not Tensor-Core-bound. Only when N=128 do the two sides align at roughly 64 cycles each, marking the transition into a math-limited regime.

This produces a distinctly sloped roofline at exactly 128 B/clk. On Blackwell, feeding the Tensor Cores is often harder than using them.

6. Why 2SM MMA can scale by more than 2x

This shared-memory bottleneck explains a counterintuitive result: in SS mode and for small shapes, 2SM MMA can achieve greater than 2x strong scaling over 1SM MMA.

That is not magic. It is bottleneck removal.

When the work is split across two SMs, each SM contributes its own shared-memory bandwidth. The kernel is no longer constrained by the single-SM SMEM ceiling that held back the 1SM path. In effect, the architecture doubles both the compute resources and the on-chip bandwidth feeding them, so the observed speedup can exceed the naive 2.0x expectation.

In TS mode, where operand A comes from TMEM rather than shared memory, scaling behaves much more cleanly and sits near the expected 2.0x.

7. Latency and data format still matter

Single-instruction latency reveals more of the underlying machine behavior.

Latency grows roughly linearly from N=64 to N=128, then spikes at N=256. Data format also changes the ordering:

The intuition is straightforward:

• S8 is fastest because integer tensor operations are power-efficient and simple
• Microscaled formats pay a small extra cost to derive and apply scale factors

Even if an instruction is well chosen, issue efficiency remains a separate problem. To truly approach speed-of-light throughput, a kernel likely needs on the order of 256 to 1024 in-flight MMA instructions so that issue overhead and commit waits are fully amortized.

Most real kernels are nowhere near that. They often carry only 1 to 4 in-flight MMAs, which artificially caps throughput around 78% to 80% of speed-of-light. That is why maximizing MMA instruction size per shared-memory tile is not optional on Blackwell; it is one of the few levers strong enough to move the roofline.

8. Practical rules for kernel writers

Blackwell tuning can be summarized in a few rules:

1. Design around TMEM explicitly. Accumulators no longer belong to registers or warps, so the pipeline has to be structured around TMEM residency and transfer boundaries.
2. Treat CTA clusters as first-class hardware. Launch geometry, cluster size, and GPC packing all affect whether the kernel actually runs in parallel.
3. Choose the copy path by access pattern, not ideology. TMA wins on large regular tiles; LDGSTS wins on responsiveness and irregularity.
4. Do not treat remote DSMEM like local SMEM. Use cp.async.bulk for real inter-CTA throughput.
5. Expect shared memory to be the bottleneck before tensor math. For SS-mode kernels, shape selection and operand staging dominate achievable performance.
6. Use larger MMA shapes and deeper in-flight pipelines whenever possible. Blackwell leaves a lot of performance stranded when kernels are too shallow.

The headline message is simple: Blackwell’s Tensor Cores got faster, but the software problem got harder. The best kernels are no longer the ones that merely maximize FLOPs. They are the ones that understand where the new bottlenecks moved, and then reorganize the entire pipeline around those new limits.

Gemma 4 represents a comprehensive effort in optimizing parameter efficiency, memory bandwidth, and long-context multimodal processing across edge and server deployments.

1. Model Lineup

Gemma 4 introduces four variants, categorized by structural paradigm:

2. Text Architecture

Interleaved Local and Global Attention

Instead of standard full attention, Gemma 4 interleaves local (sliding window) with global (full) attention:

• Sliding Window: 512 tokens for E2B/E4B, 1024 for 26B/31B.
• Interleaving Ratio: 5:1 (5 local layers per 1 global layer), except E2B which uses 4:1.
• Final Layer Constraint: The final layer is always global, ensuring full-sequence synthesis before output.

Shared KV Cache and K=V Trick

Two major optimizations for global attention memory overhead:

• GQA with K=V: Local layers use standard GQA (2 Query heads per 1 KV head). Global layers scale to 8 Query heads per KV head with doubled Key dimensions. For global attention layers, Keys are set equal to Values ($K=V$), collapsing the KV-cache into a K-cache only and halving the memory footprint for those states.
• Shared KV Cache: The last $N$ layers reuse $K$ and $V$ tensors from the previous non-shared layer of the same attention type (sliding or full), cutting both redundant compute and memory.

Proportional RoPE (p-RoPE)

Standard RoPE applies rotation across all embedding pairs. With 256K contexts, low-frequency pair rotations accumulate and add noise to semantic tracking.

Gemma 4 uses p-RoPE on global attention layers only: with $p = 0.25$, only the first 25% of coordinate pairs receive RoPE positional information, while the remaining 75% receive zero rotation. This isolates positional data to high-frequency dimensions, leaving low-frequency dimensions clean for semantic meaning.

3. Per-Layer Embeddings (PLE)

The “E” in E2B/E4B stands for “Effective” parameters. Instead of adding depth or width, these models use Per-Layer Embeddings.

• Architecture: Beyond the standard embedding lookup ($V \times d_{model}$), PLE adds a massive lookup table of dimensions $V \times d_{PLE} \times L$ (Vocabulary $\times$ PLE dim $\times$ Number of Layers). For E2B: $262{,}144 \times 256 \times 35$.
• Flow: At ingestion, the model fetches a per-layer $d_{PLE}$ embedding for every layer. During the forward pass at layer $L$, a gating function weights this layer-specific embedding, projects it back to $d_{model}$ (1,536 for E2B), and combines it with the residual stream via a lightweight block after attention and FFN.
• Hardware Insight: The PLE lookup table is queried only once per token at inference start, so it can reside in flash memory rather than VRAM. This allows a 5.1B parameter model to run at the speed and VRAM cost of a 2.3B model — hence the “effective” parameter metric.

4. Mixture of Experts (MoE)

The 26B A4B model uses sparse activation to run at 4B-model speed:

• Routing: FFN is split into 128 total experts; a router selects 8 experts per token. The token embedding is scaled by the router’s probability for each expert’s contribution.
• Shared Expert: 1 shared expert is always activated for every token, formulated at 3x the size of a standard expert to capture broad, general knowledge.

5. Multimodal Encoders

Vision Encoder

Based on a ViT (~150M params for E-models, ~550M for larger models):

• Variable Aspect Ratios: Adaptively resizes input while maintaining aspect ratio, applying padding only where the image doesn’t perfectly divide into 16x16 pixel patches (no warping into squares).
• 2D RoPE: Patch embeddings are split into two halves — one receives RoPE tracking width ($w$), the other tracks height ($h$), baking 2D coordinates into the transformer.
• Soft Token Budget & Spatial Pooling: Users define a “soft token budget” ($B$, range 70–1120). The image is divided into at most $B \times 9$ patches. Every $3 \times 3$ grid of neighboring patches is averaged (pooled) into a single embedding. A linear projection and RMSNorm then map vision space into text embedding space.

Audio Encoder (E2B and E4B only)

A ~300M parameter encoder processes raw audio (up to 30 seconds) into LLM-compatible tokens:

1. Extract Mel-spectrogram features (time vs. frequency).
2. Group features into chunks.
3. Overlap and downsample chunks via two 2D convolutional layers.
4. Process through a Conformer (Transformer Encoder with convolutional module).
5. Linear projection to align with the Gemma 4 embedding space.

6. Deployment Notes

• Sampling: DeepMind recommends temperature=1.0, top_p=0.95, top_k=64 across all models.
• Thinking Mode: Activated by placing a <|think|> token in the system prompt. The model outputs reasoning in <|channel>thought\n ... <channel|> tags. In multi-turn conversations, historical thoughts must be stripped — only final answers remain in context history.
• Modality Order: Image and audio soft tokens should always be placed before the text prompt. For PLE conditioning on multimodal inputs, audio/image positions use the pad token ID, passing neutral per-layer signals.

Core Motivation

The current paradigm of LLM scaling (Kaplan et al., 2020) relies heavily on expanding model size, data, and compute. However, the AI field is rapidly approaching the physical limits of available data and energy. While sparse modifications like Mixture of Experts (MoE) scale parameters without scaling per-example compute, the Residual Matrix Transformer (RMT) proposes an entirely new axis for scaling: the residual stream size.

In a standard transformer, the residual stream is a vector of dimension $D$ that acts as a memory bus across layers. Scaling $D$ linearly scales the size of all parameter matrices, inflating parameter and FLOP counts quadratically. The RMT solves this by replacing the residual stream vector with an outer product memory matrix, decoupling the residual stream’s bandwidth from the model’s compute and parameter footprint.

Mathematical Framework: Outer Product Memory

The architecture builds on outer product memory stores (Kohonen, 1972; Anderson, 1972).

Given a set of $N$ key vectors $q^{(p)} \in \mathbb{R}^{D_k}$ and data vectors $x^{(p)} \in \mathbb{R}^{D_v}$ for $p = 1, \ldots, N$, an outer product store $M \in \mathbb{R}^{D_k \times D_v}$ is constructed by summing their outer products:

where $u \otimes v = uv^T$, and Norm is LayerNorm.

To retrieve a specific data vector $x^{(r)}$ from $M$, we perform a tensor contraction over the first dimension using its associated key vector $q^{(r)}$:

RMT Architecture

In the RMT, the batched residual stream for $N$ tokens is represented as a tensor $X \in \mathbb{R}^{D_k \times D_v \times N}$, rather than the standard matrix $X \in \mathbb{R}^{D \times N}$.

Attention Layer

In a standard transformer, features are retrieved using linear transformations (e.g., $Q^{(h)} = W_Q^{(h)} X$). In the RMT, the expensive $W_Q, W_K, W_V \in \mathbb{R}^{D_h \times D}$ weight matrices are completely removed and replaced by learned key vectors $r_Q^{(h)}, r_K^{(h)}, r_V^{(h)} \in \mathbb{R}^{D_k}$.

The attention inputs are retrieved via tensor contraction:

These resulting matrices belong to $\mathbb{R}^{D_v \times N}$ (where $D_v$ acts as the attention head dimension $D_h$). Standard Head Attention (SHA) is applied normally, and the output is written back into the residual matrix using an output key vector $w_O^{(h)} \in \mathbb{R}^{D_k}$:

Feed-Forward Layer (FFN)

Unlike attention, the FFN retains its standard linear transformations $W_1$ and $W_2$, as evidence suggests these matrices store critical factual information rather than just routing it. The RMT uses key vector “adapters”: it retrieves $R$ data vectors from the matrix, concatenates them for the standard FFN operation, and then splits (un-vecs) the output to store it back into the residual matrix via outer products with $w_{FF}^{(h)}$.

Variance Propagation

For deep networks to initialize and train stably, the mean and variance of activations and gradients must propagate effectively (Glorot & Bengio, 2010). The paper provides a closed-form derivation proving that outer product storage and retrieval maintain healthy variance.

Let the forward storage operation for a single token be $X_{out} = \sum_{h=1}^R w^{(h)} \otimes x_{in}^{(h)}$, where weights are initialized independently with mean 0.

The variance propagates as:

Assuming independence and $\mu_w = 0$:

By choosing standard initialization dimensions, the ratio $\frac{\sigma_{x_{out}}^2}{\sigma_{x_{in}}^2}$ can be kept close to 1, avoiding vanishing or exploding gradients.

Scaling Performance

The replacement of standard weight matrices with vectors transforms the economics of the network:

• Cost of Scaling: Increasing the residual stream size by 100% in a standard transformer yields ~94% more FLOPs and 100% more parameters. In the RMT, increasing residual matrix capacity ($D_k$) by 100% increases both parameters and FLOPs by < 1%.
• Efficiency: To reach the identical target loss, the RMT uses 58% fewer FLOPs, 25% fewer parameters, and 41% fewer training tokens compared to a Chinchilla-optimal baseline transformer.
• Zero-Shot Dominance: Evaluated on LAMBADA, PIQA, and ARC, an RMT trained on 28% fewer FLOPs outperformed a standard transformer that was 33% larger.
• “Free” Scaling Axis: When holding parameter count, dataset size, and compute budget constant, expanding the residual stream size $D_k$ monotonically decreased validation loss.

Caveats

• Memory: Roughly equivalent during training (larger residual activations for gradient checkpointing offset by fewer model parameters), but strictly more efficient during inference.
• Wall-Clock Time: In current PyTorch/JAX ecosystems, highly optimized GEMM kernels run significantly faster than unoptimized tensor contractions. Despite needing dramatically fewer FLOPs, the RMT takes ~43% longer per training step. A custom CUDA kernel for contracting over small key vectors could bridge this gap.

Summary

The RMT unlocks a new “free” scaling dimension by replacing the residual stream vector with an outer product memory matrix. Weight matrices in attention are replaced by learned key vectors, making residual stream scaling nearly cost-free in parameters and FLOPs. This promises significant reductions in the energy and data required for frontier LLM training, pending kernel-level optimization to close the wall-clock gap.

• Nvidia – The Inference Kingdom Expands (SemiAnalysis)

At GTC 2026, Nvidia aggressively expanded its hardware and inference ecosystem to address the emerging bottlenecks of memory-bound LLM decode phases, massive CPU demands in reinforcement learning, and KV cache storage limits. The company announced three entirely new system architectures: Groq LPX (integrating the LP30 chip), Vera ETL256, and STX, alongside major updates to the Kyber rack architecture (NVL144, NVL576, and NVL1152).

1. The Groq “Acquisition” and LP30 Architecture

Nvidia functionally acquired Groq via a $20B IP licensing and acqui-hire deal, bypassing drawn-out antitrust regulations. Groq’s hardware, specifically the LPU (Language Processing Unit), is utilized to create a disaggregated decode system.

LP30 (Groq 3 LPU) Hardware Details

• Silicon & Manufacturing: Designed on Samsung’s SF4X node, skipping the failed LPU 2 which suffered from 112G SerDes malfunctions. Using Samsung SF4X allows Nvidia to bypass TSMC N3 logic constraints and HBM allocation constraints, enabling incremental revenue scale-up.
• Memory Hierarchy: Features a single-level memory hierarchy with 500MB of on-chip SRAM (up from 230MB in Gen 1) providing an ultra-fast 150 TB/s memory bandwidth. It lacks HBM entirely.
• Compute: Dedicated to tensor-first compute with 1.2 PFLOPs of FP8, which is a fraction of standard GPU compute but highly optimized for deterministic execution. The chip pumps instructions vertically and streams data horizontally across functional slices (VXM, MEM, SXM, MXM).
• Form Factor: Deployed in the LPX Compute Tray, featuring a belly-to-belly PCB design (8 LPUs on top, 8 on bottom) to minimize X/Y trace distances, alongside 2 Altera “Fabric Expansion Logic” FPGAs, 1 Intel Granite Rapids CPU, and a BlueField-4 module.

2. Decoding Acceleration Techniques

The integration of the LPU is designed to accelerate the latency-sensitive, memory-bounded decode phase of LLM inference, leaving the compute-intensive prefill phase to GPUs.

Attention-FFN Disaggregation (AFD)

The Problem: During decode, GPU utilization for the Attention mechanism barely improves as batch sizes scale because it is bounded by loading KV cache. Conversely, Feed Forward Network (FFN) utilization scales effectively with larger batch sizes. In state-of-the-art sparse Mixture-of-Expert (MoE) models, utilization drops further as tokens route to a larger pool of experts.

The Solution: Attention operations are stateful (relying on dynamic KV cache) and are thus mapped to HBM-heavy Rubin GPUs. FFN operations are stateless (depending only on token inputs) and are mapped to the SRAM-heavy, deterministic LPUs.

Network Optimization: To hide the communication latency of dispatching tokens from GPU to LPU experts and combining them back, the system relies on ping-pong pipeline parallelism, allowing tokens to continuously bounce between GPUs and LPUs over Spectrum-X Ethernet.

The key insight here is a further decomposition beyond prefill-decode disaggregation. Within the decode phase itself, attention and FFN have fundamentally different computational profiles:

Speculative Decoding & Memory Management

• LPUs can host draft models or Multi-Token Prediction (MTP) layers to predict $k$ new tokens, which the main model verifies in a single “warm prefill” step.
• Unlike stateless FFNs, draft models/MTP layers require tens of gigabytes of dynamic KV cache. To support this, the Altera FPGAs provide up to 256GB of additional DDR5 memory per FPGA for the LPUs to access.

This is a natural evolution of the speculative decoding paradigm: instead of running the draft model on the same GPU as the verifier, offloading it to a dedicated LPU eliminates the resource contention entirely.

3. Networking Topologies and Bandwidth Math

Nvidia’s systems push the physical limits of copper to keep TCO down, orchestrating incredibly dense electrical networks before resorting to optical interconnects.

LPX Rack Network Math

• Intra-Tray: 16 LPUs connect via an all-to-all PCB mesh. Each LPU routes to 15 others via $4 \times 100G$ C2C links.
• Intra-Rack (Inter-Node): Each LPU routes $2 \times 100G$ to 15 other nodes. With FPGAs connecting at 25G/50G, a node features 1,020 differential pairs. Across 16 nodes, the copper backplane supports 8,160 differential pairs ($16 \times 1020 / 2$).
• Total intra-rack scale-up bandwidth:

Rubin Ultra NVL144 (Kyber Rack) Scale-up Math

The Kyber rack fits 144 Rubin Ultra GPUs and 72 NVLink 7 switches.

• GPU Bandwidth: Each GPU uses 72 Differential Pairs (DPs). At $200\text{ Gbit/s bi-di}$ per channel, each GPU achieves 14.4 Tbit/s uni-directional scale-up bandwidth.
• Switch Bandwidth: Each NVSwitch 7 uses 144 lanes of 200G, totaling 28.8 Tbit/s uni-directional bandwidth. Connecting these requires midplanes and copper flyover cables.
• NVL288 Constraints: Scaling to 288 GPUs across two racks via copper would require 20,736 additional DPs, acting as a massive upper bound on cable content, unless higher radix switches are introduced.

4. Co-Packaged Optics (CPO) vs. Copper Roadmap

A key architectural insight from GTC 2026: Nvidia uses copper where it can, and optics where it must.

• Rubin / Rubin Ultra: Scale-up within NVL72 and NVL144 Kyber racks remains strictly copper.
• NVL576 (Rubin Ultra): This 8-rack system will be the first introduction of CPO scale-up, utilizing a two-tier all-to-all network between racks, though intra-rack networking stays copper.
• Feynman NVL1152: Will fully adopt CPO for rack-to-rack scale-up, overcoming the physical reach/shoreline limits of bumping electrical SerDes from 224G to 448G.

The transition is driven by physics: as SerDes rates double, the reach of copper decreases and the shoreline (pin count per die edge) becomes the binding constraint. CPO sidesteps both by converting electrical signals to photons at the package boundary.

5. Ancillary Infrastructure: Vera ETL256 and STX

To prevent non-GPU components from bottlenecking system performance, Nvidia released auxiliary racks:

Vera ETL256

A standalone, liquid-cooled rack packing 256 Vera CPUs to handle the surging preprocessing and simulation demands of Reinforcement Learning workloads. It utilizes a single-tier Spectrum-6 multiplane topology.

This reflects a growing reality: RL training pipelines (reward model evaluation, environment simulation, data preprocessing) impose enormous CPU demands that steal GPU cycles if co-located. Dedicated CPU racks eliminate this contention.

CMX and STX (Context Memory Storage)

To combat the exponential growth of KV Cache, Nvidia introduced Tier G3.5 NVMe storage. The STX reference rack utilizes BlueField-4 DPUs (featuring a Vera CPU, 2x CX-9 NICs, and 2x SOCAMM modules) to offload “warm” KV cache from expensive GPU HBM and system DRAM, optimizing inference efficiency.

This creates a multi-tier memory hierarchy for KV cache:

6. Strategic Synthesis

Nvidia is fundamentally transitioning from selling standalone AI accelerators to orchestrating entire datacenter macro-architectures. Several strategic threads emerge:

1. Supply chain arbitrage: By disaggregating decode tasks to SRAM-heavy LPUs manufactured on Samsung’s SF4X (not TSMC N3), Nvidia bypasses critical industry supply constraints (TSMC N3 capacity and HBM allocation) while preserving high-margin GPU allocations strictly for compute-heavy prefill.

2. Copper maximalism: The aggressive densification of copper flyover cables and midplanes in Kyber and LPX architectures proves that TCO optimization remains paramount. Nvidia delays the transition to expensive CPO interconnects until physical electrical bounds absolutely mandate it in the Feynman generation.

3. Full-stack lock-in: With dedicated GPU racks (Kyber), CPU racks (Vera ETL256), storage racks (STX), and decode accelerator racks (LPX), plus the networking fabric (Spectrum-X, NVLink 7) tying them together, Nvidia is selling complete datacenter blueprints rather than individual chips.

4. Inference economics: The LPU integration directly addresses the economic pain of decode. During decode, GPUs are massively underutilized on compute but bottlenecked on memory bandwidth. Offloading FFN to cheap, deterministic LPUs improves the cost-per-token by utilizing the right silicon for the right workload.

“Year after year, the flowers bloom the same; year after year, the people are not those of old.” (年年岁岁花相似，岁岁年年人不同)

My first encounter with these blossoms was in March 2020. My wife had just begun her doctoral studies, and the world was tilting into the surreal uncertainty of a global pandemic. I remember standing together under the pale canopy, our faces hidden behind masks, breathless at the sheer scale of the bloom. Even amidst the collective shadow of that year, the vitality of the trees felt like a promise—a glimmer of hope for a new chapter. Back then, I was still living a life in transit, commuting long distances for work. Those swirling petals felt like a fairy tale, a gentle welcome to the city we had decided to call home.

In the years that followed, the spring bloom became our annual pilgrimage. Each passing season brought new friends into our circle and saw our lives grow more settled and rooted. My phone turned into a digital archive of this evolution, filled with photos of my wife framed by those same pink clouds, year after year—until the masks were finally put away and the world moved on.

Then, life bloomed in a more literal sense. Two years ago, in the heart of spring, our son was born at a nearby hospital. When he was just two weeks old, the blossoms reached their peak. Though my wife was still navigating the fragile days of recovery, we walked together with my parents to capture a family portrait under the trees. Our son was a tiny, drowsy observer, taking in the bustling crowds with the wide, uncomprehending eyes of a newborn.

A year later, the scene shifted. He was a sturdy one-year-old in a blue sweater vest, sporting a look of adorable reluctance as we tried to capture the perfect photo. He had just learned to walk—steadying himself with our hands, stumbling, and picking himself back up. Though he couldn’t speak yet, he pointed at the falling petals with sheer excitement, narrating his joy in a language of babbles.

Now, at two years old, the trip to the grove has become a sacred family ritual—a living yardstick of our growth. This year, his excitement began the moment we left the house; he knew exactly where we were headed. Once we arrived, the roles reversed: instead of carrying him, we were the ones giving chase. He ran through the grass, his eyes darting from the crowds to the trees, then up to the drones and planes roaring in the distance. To him, the world is a gallery of wonders. I hope these moments settle into the foundation of his memory—a subconscious blueprint of beauty he can carry forever.

I do not know how much longer our path will keep us in this city. Life may eventually take us elsewhere, and a blooming season like this might one day become a rare luxury. But for now, I deeply cherish this seventh year.

A flowering tree can live for over a century. To the cherry blossoms, I am merely one of a million passing shadows in their long lives. But to me, they are a permanent part of my soul’s geography—a witness to the years we grew, the years we loved, and the beautiful, fleeting gift of being alive.

This paper asks a simple question: if attention replaced recurrence along the sequence dimension, why are we still using a fixed additive recurrence along the depth dimension?

Motivation: Residuals Behave Like Depth-Wise Recurrence

Standard residual connections can be written as:

Unrolling over depth gives:

That view makes the paper’s core observation clear: every layer sees a uniformly weighted sum of all earlier layer updates. The authors call this a time-depth duality. In an RNN, all past tokens are compressed into a single hidden state over time. In a Transformer with standard residuals, all past layer outputs are compressed into a single hidden state over depth.

The paper argues that this becomes especially problematic in PreNorm LLMs. Because the residual stream keeps accumulating unweighted updates, hidden-state magnitudes tend to grow with depth, leading to PreNorm dilution. Later layers then have to produce larger and larger updates just to maintain the same influence on the final representation.

Full Attention Residuals

The proposed fix is to replace fixed accumulation with learned softmax attention over previous layers:

Here, the values are the embedding and prior layer outputs, while the attention weights are:

Instead of deriving the query from the current hidden state, each layer uses a learned parameter vector:

This detail matters. The query is static, so depth-wise attention can still be computed in parallel across layers. The keys are RMS-normalized before scoring,

which prevents large-magnitude activations from dominating the attention distribution.

Conceptually, Full AttnRes turns the residual path into a learned depth mixer. A layer is no longer forced to treat all earlier layers equally; it can emphasize whichever layers are most useful.

Why the Full Version Does Not Scale Cleanly

From an arithmetic perspective, Full AttnRes is manageable. The paper quotes $O(L^2 d)$ work, which is acceptable for realistic layer counts. The real issue is systems cost: all previous layer outputs must remain available, so activation storage and cross-stage communication grow as $O(Ld)$. Under activation recomputation or pipeline parallelism, that becomes the bottleneck.

Block Attention Residuals

To make the idea practical, the paper introduces Block AttnRes. The $L$ layers are partitioned into $N$ blocks of size $S$, and each block is summarized by a simple sum:

A layer then attends to:

• completed block summaries $b_0, b_1, \dots, b_{n-1}$
• the running partial sum of its current block

This is the key compression step. Instead of exposing every earlier layer explicitly, the model exposes a smaller set of block summaries, which reduces memory and communication overhead from $O(Ld)$ to $O(Nd)$. The reported scaling results suggest that a small number of blocks, roughly $N \approx 8$, captures most of the gains of Full AttnRes.

Training and Inference Optimizations

The paper also does the systems work needed to make Block AttnRes usable at scale.

For training, the main trick is cross-stage caching under pipeline parallelism. Rather than repeatedly sending already-seen block summaries across virtual stages, each rank caches what it has received and only transmits the incremental blocks. That avoids redundant communication and reduces the peak communication cost enough to overlap it with computation.

For inference, the authors use a two-phase computation:

1. compute inter-block attention for all layers in a block together
2. compute intra-block attention sequentially and merge it with an online softmax

Because the layer queries are static parameters, this organization lowers memory traffic substantially. The paper reports total residual-path memory I/O of about $5.5d$ reads per layer, compared with $3d$ for standard residuals and much higher cost for multi-stream alternatives such as mHC.

Residual Connections as a Mixing Matrix

One of the nicest parts of the paper is the matrix view. If we define a depth mixing matrix $M \in \mathbb{R}^{L \times L}$, where $M_{i \to l}$ is the weight assigned by layer $l$ to layer $i$’s output, several architectures fall into the same template:

• Standard residuals: an all-ones lower-triangular matrix
• Highway networks: input-dependent scalar gates, but still low-complexity depth mixing
• Multi-stream methods such as mHC: depth-wise linear attention with a higher-rank structured state
• Attention residuals: full depth-wise softmax attention

Under this lens, Block AttnRes smoothly interpolates between standard residuals and Full AttnRes by changing how many block summaries are exposed.

Takeaways

The empirical story is straightforward:

• PreNorm dilution is reduced. Activation magnitudes stay bounded more cleanly, and gradients are distributed more evenly over depth.
• The model learns nontrivial skip patterns. Attention maps remain strongly local, but deeper layers sometimes jump back to very early layers or even the embedding.
• The preferred architecture shifts. In the paper’s iso-compute and iso-parameter sweeps, AttnRes favors deeper, narrower models than standard residual designs do.

That last point is especially interesting. If depth is no longer handicapped by uniform residual accumulation, then making a model deeper becomes a better tradeoff than it is in a baseline Transformer.

My main takeaway is that this paper treats residual connections as an architectural choice rather than a fixed law. Standard residuals hard-code a very specific depth mixing rule: every earlier layer contributes equally through simple addition. AttnRes replaces that rule with learned attention, then introduces a blockwise approximation that keeps the idea deployable at scale.

Whether this becomes a standard recipe will depend on implementation complexity and robustness across model families, but the framing is strong: residual paths are not just optimization scaffolding, they are a depth-wise information routing mechanism that can be redesigned.

Recently, I had the chance to catch up with an old high school classmate who had flown in from overseas for a conference. We both happened to be in the same city on a business trip. Looking back, our lives have always shared a curious, almost poetic symmetry: our mothers were high school classmates, and we spent our three years of high school together as friends and fellow students. In our small hometown school, we were both dedicated to our studies, often finding ourselves at the top of our class. They were undeniably brilliant, eventually heading off to a prestigious university. I, on the other hand, faced some setbacks during my entrance exams and followed a different path that, through a series of unexpected turns, led me to another university in the same city.

Our rare encounters during those university years now feel like warm memories from a past life—gathering near their campus to welcome a visiting friend, or taking a long walk through a historic park. After that, our lives became like two rays of light, diverging as we pursued our own careers, seemingly never to intersect again.

That changed a couple of years ago when they visited my city for work. We reunited over dinner for the first time in over a decade—a moment filled with excitement and the natural, quiet curiosity that comes with time. We had both started families by then. I remember us trying to explain our respective professional fields to each other, though our industries were so different it was a bit of a challenge to grasp the details. But little did we know, life was unfolding in parallel for us once again: shortly after that meeting, my first child was born, and just a few months later, they welcomed their first child into the world.

Our meeting this time around felt remarkably comfortable, filled with a quiet, steady ease. Sitting together over a warm meal, we realized it had been exactly twenty years since we first started high school together. Time truly flies. They have since pivoted into a new role in the tech space, and it was wonderful to see how sharp, curious, and passionate they remain about their work.

An old joke we’d shared online about introducing our kids someday took on a sense of genuine warmth over our meal. After dinner, we wandered into a toy store and bought two identical little plushies—sweet tokens of our friendship to bring back to our children. Our farewell was a classic modern mishap: their phone was nearly dead, and getting them back to their hotel became a bit of an adventure. I stayed with them until we got enough of a charge to book a ride, and I happily saw them off.

Watching their car disappear into the evening, I couldn’t help but smile. The spirited energy of our youth has finally been gently smoothed over by twenty years of living, learning, and growing. When old friends meet, we inevitably talk of the past, but it’s the present we truly cherish. I don’t know when our paths will cross next, but it reminded me of a simple truth: all we can do is treasure the wonderful people and the beautiful moments right in front of us.

• Scalable Training of Mixture-of-Experts Models with Megatron Core

This paper presents Megatron-Core MoE, the MoE training stack within NVIDIA’s Megatron-Core framework. It addresses the fundamental systems challenges of training trillion-parameter-class MoE models at high throughput on NVIDIA GPU clusters. The headline results: 1,233 TFLOPS/GPU on GB300 and 1,048 TFLOPS/GPU on GB200 for DeepSeek-V3-685B.

1. MoE Fundamentals

Architecture

Given an input token representation x, the router computes:

The MoE layer output is:

where $E_i$ is the $i$-th expert network. Three key advantages:

• Scalable capacity: model size grows independently of per-token compute
• Computational efficiency: only $K$ of $E$ experts activate per token
• Specialization: different experts learn different input patterns

The Parameter-Compute Mismatch

This is the central insight of the paper. In a dense transformer with $N_\text{total}$ parameters, FLOPs per token is approximately $6N_\text{total}$, so parameters and computation scale in lockstep.

For MoE, per-token computation is approximately $6N_\text{active}$ where $N_\text{active} \propto K$ while $N_\text{total} \propto E$, and $K \ll E$.

Concrete example: DeepSeek-V3 has 685B total parameters but only 37B active per token, an 18x gap.

This creates the Three Walls:

2. MoE Layer Architecture: Four-Stage Forward Pass

Stage 1: Route. A learned linear projection maps each token’s hidden state to $E$ logits: $\mathbf{l} = \mathbf{W}_r^\top \mathbf{x} \in \mathbb{R}^E$. A score function, either softmax or sigmoid as used in DeepSeek-V3, converts these into probabilities. Top-$k$ selection chooses the active experts.

Stage 2: Dispatch. Tokens are permuted so those destined for the same expert are contiguous. Three backends are described: AllGather, All-to-All, and Flex via DeepEP/HybridEP.

Stage 3: Expert Computation. All local experts run in a single Grouped GEMM call. Each expert is a two-layer MLP with optional gating: \(E_i(\mathbf{x}) = \mathbf{W}_2^{(i)} \phi(\mathbf{W}_1^{(i)} \mathbf{x})\)

Stage 4: Combine. Inverse communication returns tokens, unpermutation restores order, and shared expert output is added.

3. Parallel Folding and Multi-Dimensional Parallelism

The Dense-Sparse Mismatch

A single Transformer block contains two fundamentally different computation patterns:

Prior frameworks forced World Size = TP x CP x PP x DP, where EP <= DP. This creates three problems:

1. Multiplicative GPU requirements: EP=8 forces DP>=8, and with CP=8 the minimum becomes 64 GPUs.
2. Forced suboptimal parallelism: high TP fragments small experts, while low TP underparallelizes attention.
3. Cross-node communication: EP constrained within DP forces all-to-all across slow interconnects.

Parallel Folding Solution

Core idea: decouple attention and MoE parallelism mappings.

• Attention layers form groups over TP x CP x DP x PP
• MoE layers form groups over ETP x EP x EDP x PP
• Only constraint: PP must remain consistent

Key benefits:

1. Breaks EP <= DP: EP can fold across TP x CP groups. Example: attention TP=4, CP=2, DP=8, PP=4 on 256 GPUs. Traditionally EP<=8; with folding EP=64 becomes possible.
2. Reduces minimum GPUs: CP=8 and EP=8 traditionally requires 64 GPUs; with folding, only 8.
3. Independent optimization: attention uses high TP, while MoE uses ETP=1 for full expert width.
4. NVLink locality: both CP and EP all-to-all stay within the NVLink domain.

Gradient Handling

Expert gradients are scaled by edp_size / dp_size to account for the different effective batch sizes seen by experts versus dense layers.

4. Breaking the Memory Wall

Memory Anatomy (DeepSeek-V3, PP4 x VPP4 x EP64, 256 GPUs)

Key insight: activations dominate, exceeding weights and optimizer states combined.

Memory-Efficient Permutation (Zero Overhead)

Standard formulation applies routing weights after expert computation: \(y = \sum_{i \in \mathcal{T}(\mathbf{x})} p_i \cdot \mathbf{W}_2^{(i)} \phi(\mathbf{W}_1^{(i)} \mathbf{x})\)

Memory-efficient version absorbs $p_i$ before the second linear layer: \(y = \sum_{i \in \mathcal{T}(\mathbf{x})} \mathbf{W}_2^{(i)} \left(p_i \cdot \phi(\mathbf{W}_1^{(i)} \mathbf{x})\right)\)

Since $\mathbf{W}_2^{(i)}$ is a pure linear map with no bias, scalar multiplication commutes: \(p_i \cdot \mathbf{W}_2^{(i)} \mathbf{h} = \mathbf{W}_2^{(i)} (p_i \cdot \mathbf{h})\)

Why this saves memory: in the standard version, computing $\partial\mathcal{L}/\partial p_i$ requires retaining each expert output $E_i(\mathbf{x})$. In the efficient version, $p_i$ multiplies $\phi(\mathbf{z}_i)$ directly, so $\partial\mathcal{L}/\partial p_i$ only depends on $\phi(\mathbf{z}_i)$, which can be recomputed from $\mathbf{z}_i = \mathbf{W}_1^{(i)} \mathbf{x}$ already saved for SwiGLU backward. This saves roughly 26.3 GB per GPU for DeepSeek-V3 with essentially zero extra compute.

FP8/FP4 Activations

Linear layer inputs stored in FP8 instead of BF16 reduce memory by 50% per tensor. For DeepSeek-V3, that is roughly 16 GB saved. FP4 pushes this further to a 75% reduction.

Fine-Grained Recomputation

Two composable techniques:

1. Granular recomputation: selectively recompute only specific operations such as activation functions, LayerNorm, and MLA up-projection, typically with under 5% compute overhead.
2. Output-discarding recomputation: release checkpointed module outputs immediately after downstream consumption and restore them via recomputation during backward.

Critical insight: full-layer recomputation of MoE is especially expensive because it re-triggers EP all-to-all communication. Fine-grained recomputation avoids that penalty.

Fine-Grained Activation Offloading

Forward: input activations are offloaded to CPU through a dedicated D2H stream, overlapping with the next module’s computation.

Backward: Layer-Staggered Reload reloads the same module type from the next layer while gradients are computed for the current layer. Only one activation per module type resides on GPU at a time.

Peak memory advantage over full recomputation:

• Full recomputation: $L \times \text{layer_input} + 1 \times \text{layer_intermediate}$
• Offloading: $1 \times \text{layer_input} + 1 \times \text{layer_intermediate}$

Results: 10-18% memory reduction with only 1.6-2% throughput overhead. For Qwen3-235B, offloading enabled a lower TP degree and about 15% throughput improvement.

Precision-Aware Optimizer

Adam stores first and second moments. The optimization here is to store moments in BF16 or FP8, then cast to FP32 inside TransformerEngine’s FusedAdam kernel for the actual update.

Memory per parameter per DP rank decreases from $6 + 12/d$ bytes to $6 + 8/d$ bytes, saving on the order of 10-12 GB from the 32.1 GB optimizer-state budget.

FSDP for MoE

Dual DeviceMesh design: dense layers shard across the full DP group, while expert layers shard across the EDP group.

Two key optimizations:

1. Non-uniform sharding: flatten and concatenate module parameters, then shard non-uniformly so shard boundaries align with communication buffers for zero-copy collectives.
2. Persistent double buffers with NCCL User Buffer Registration: pre-allocate two persistent buffers and cycle between them. This reduces SM footprint from 8-32 SMs to 1-4 SMs.

5. Breaking the Communication Wall

Communication Anatomy

For DeepSeek-V3: 58 MoE layers times 2 operations per layer gives 116 dispatch/combine operations per forward pass. Backward doubles this. At 50 GB/s inter-node bandwidth, a single 200 MB dispatch already costs milliseconds, which compounds rapidly over an iteration.

HybridEP

Developed by NVIDIA for NVLink-rich topologies such as NVL72.

Dispatch: reads data from global memory into shared memory based on routing info, then writes to destinations via FIFO queues. For inter-node traffic, GPUs with the same local index across nodes exchange first, then forward within the node.

Combine: fuses reduction into the communication kernel itself. Cross-node data is reduced first, then the remaining intra-node reduction completes locally.

Performance on GB200 with EP64:

• HybridEP dispatch: 675 us vs all-to-all 930 us
• HybridEP combine: 744 us vs all-to-all 827 us

EP Communication Overlapping

1F1B forward-backward overlap merges the forward pass of one microbatch with the backward pass of another.

Two key optimizations:

1. Stream Separation: compute stream and communication stream run in parallel.
2. W/D Split: backward MLP is split into weight gradient (W/mlp) and data gradient (D/mlp). Only D/mlp depends on backward dispatch, which opens more room to hide communication.

Result: EP communication overhead drops from 30-40% to under 5% of iteration time for DeepSeek-V3 on H100.

6. Breaking the Compute Efficiency Wall

Grouped GEMM

DeepSeek-V3’s 256 small experts produce GEMMs with M dimensions around 128 tokens per expert, far below the regime needed for peak Tensor Core efficiency.

Four implementations are discussed:

1. Multi-stream cuBLASLt GEMMs: individual GEMMs launched into multiple CUDA streams
2. CUTLASS Grouped GEMM: a fused single-kernel path
3. cuBLASLt Grouped GEMM (device-initiated): reads shapes from device memory, making it CUDA-Graph-compatible
4. cuteDSL Grouped GEMM: fuses SwiGLU activation and FP8 quantization into the GEMM epilogue

Permutation Fusion

Three-stage pipeline:

• Preprocessing: generate the Row ID map once
• Permute: move tokens according to offset maps
• Unpermute: inverse permutation plus FP32 accumulation

Router and Aux-Loss Fusion

Three fused kernels are described:

• score computation with top-$k$ and softmax/sigmoid
• score computation for auxiliary loss
• auxiliary loss computation

CUDA Graphs

Full vs Partial: Full CUDA Graphs capture the entire forward-backward pass, but only for drop-and-pad MoE. Partial CUDA Graphs capture only the static components:

• Graphable: attention, router, EP preprocessing, shared experts, dense MLP
• Not graphable: token dispatch, expert GEMM with dynamic M dimension, token combine

Memory optimizations:

• Graph count reduction: without PP, microbatches share graphs (L x 2); with PP, each microbatch needs its own (L x M x 2). An is_first_microbatch GPU flag controls microbatch-specific behavior.
• Pool sharing: all graphs share one pool when captured in execution order.
• Buffer reuse: static I/O buffers are reused across graphs per PP execution order.

Result: about 10% end-to-end speedup with roughly 7 GB extra memory on DeepSeek-V3 GB200.

Full CUDA Graphs for Dropless MoE: Three Complementary Techniques

Challenge 1: kernel launch without knowing problem size. Solution: device-initiated Grouped GEMM, where cuBLASLt reads shapes directly from device memory. The cuteDSL path can also fuse SwiGLU and quantization into the epilogue.

Challenge 2: memory allocation without knowing actual size. Without mitigation, worst-case buffers waste $O(\text{EP_size})$ more memory.

ECHO (Elastic Cloning for Hot Experts)

Popular experts can receive far more tokens than others. ECHO dynamically clones hot experts onto spare slots on underutilized ranks.

Forward: the ECHO planner identifies hot experts, generates the hot expert map and updated routing map, copies weights to spare slots through HybridEP, and routes tokens to both the home and cloned experts.

Backward: expert-gradient dispatch collects gradients from cloned experts back to the home experts.

Paged Stashing

Decouples worst-case computation buffers from activation storage:

• Single tmp buffer sized for worst case and shared across all layers
• Paged stashing buffer stores only the actual tokens used per layer

This reduces memory from $O(\text{layers} \times \text{worst_case})$ to $O(\text{worst_case} + \text{actual_total})$.

Implementation detail: PagedStashBuffer uses 64 tokens per page with a circular-buffer free list. Stash and reload kernels are device-initiated and overlap with computation via dedicated Pack and Unpack CUDA streams.

7. Reduced-Precision Training (FP8/FP4)

Strategy: Selective Precision

Three principles:

1. Protect routing: keep the router in FP32
2. Preserve key components: embeddings, output layers, main gradients, master weights, and optimizer states stay high precision
3. Quantize bulk computation: expert GEMMs and activations go low precision

FP8 Recipes

Per-Tensor FP8 (Hopper and Blackwell): one scale per tensor. Two variants are discussed:

• Delayed scaling: uses historical amax, but is not the recommended path
• Current/live scaling: computes scale just in time

The hybrid format uses E4M3 for inputs and weights and E5M2 for gradients.

Blockwise FP8 (recommended on Hopper): uses E4M3 for all tensors, with activations and gradients quantized in 1 x 128 tiles and weights in 128 x 128 blocks. The paper notes this recipe has already been proven at scale on models such as DeepSeek-V3, Minimax-M2, and Ant Ling-2.0.

MXFP8 (recommended on Blackwell): uses 1 x 32 element granularity with native Blackwell Tensor Core support and hardware-accelerated scaling. The important caveat is that parameter AllGather still communicates in BF16 because MXFP8 uses different quantization directions for forward and backward.

NVFP4

Uses E2M1 with two-level microscaling:

• Per-tensor FP32 scale: remaps the overall distribution into a range compatible with block scaling
• Per-block E4M3 scale: blocks of 16 elements then map into FP4 range

Three critical algorithmic additions are required for stable training:

1. Random Hadamard Transforms (RHT): applied to weight-gradient computation to reduce outlier impact
2. 2D scaling: 16 x 16 weight-block scaling keeps forward and backward consistent
3. Stochastic rounding: applied on gradients to reduce rounding bias during FP4 conversion

FP8/FP4 Primary Weights

The framework eliminates a redundant BF16 copy by casting directly from FP32 to FP8 or FP4. The distributed-optimizer quantization flow is:

1. Get local abs-max from the master weights
2. AllReduce to get global abs-max
3. Use global abs-max plus master weights for the partial cast

For the blockwise recipe, specialized kernels aware of the 2D weight layout compute abs-max over 128 x 128 blocks.

MoE-Specific FP8/FP4 Challenges

Padding alignment: FP8 GEMMs require alignment to 16 for per-tensor and blockwise FP8, or 32 for MXFP8 and NVFP4. Because token dimension varies dynamically, the paper describes two solutions: routing-map padding and fusing padding directly into permutation.

Grouped quantization: multiple expert input tensors are quantized in a single fused kernel, and the implementation is CUDA-Graphable.

NVFP4 quantization fusion: RHT is fused with quantization to avoid extra BF16 traffic. Hadamard is still computed twice, once for amax and once inside the fused quantization path, but this remains faster than materializing a full BF16 buffer. Stochastic rounding uses cuRANDDx for in-kernel random-number generation.

8. Long-Context MoE Training

The Computational Shift

At 64K tokens, SDPA consumes 69% of FLOPs, versus roughly 10-15% at short sequence lengths. SDPA scales as $O(s^2)$ while MoE scales as $O(s)$, so attention becomes the dominant cost.

Key recommendation: do not recompute core attention at long sequence lengths. At 64K, SDPA recomputation adds about 18% compute overhead but saves only 9 GB of memory. Recomputing non-SDPA components instead saves 89.8 GB with lower performance impact.

CP vs TP Trade-offs

• P2P CP: preferred across nodes because ring-style KV exchange overlaps naturally with SDPA
• All-to-all CP: converts sequence-sharded layouts to head-sharded layouts before SDPA
• TP: preferred within nodes for sharding linear weights

Practical guideline: all-to-all CP + TP inside nodes; P2P CP across nodes.

Dynamic Context Parallelism

For variable-length sequences, the system selects CP degree per microbatch based on actual sequence lengths. Multiple CP groups are pre-constructed per rank during initialization. The per-token loss is:

where $\mathcal{V}$ is the set of valid non-padding tokens.

9. Production Features

Load Balancing Strategies

Three approaches are discussed:

• Auxiliary loss: gradient-based, differentiable, soft balance
• Sinkhorn: assignment-based, non-differentiable, hard balance; iterates row and column normalization to convergence
• Aux-loss-free / Expert Bias: feedback-based, non-differentiable, adaptive; updates expert bias based on token-count feedback

Latent MoE

Latent MoE inserts a shared down-projection before expert dispatch and an up-projection after combine:

The compression ratio $\alpha = d / \ell$ reduces both all-to-all volume and per-expert weight size by a factor of $\alpha$.

Flexible Asymmetric VPP

Allows different numbers and types of layers per virtual pipeline stage. For DeepSeek-V3 with 61 decoder layers plus 1 MTP layer, PP=16, and VPP=2, the first stage holds the embedding plus 3 dense decoder layers to match the cost of 2 MoE layers, while the MTP layer sits in its own standalone stage and the loss is separated.

Upcycling

Converts a dense checkpoint into MoE via virtual-group initialization: shard MLP weights in the intermediate dimension (4h -> 2h), duplicate the shards, then initialize half the router weights and duplicate them. This guarantees Top-2 routing initially selects one expert from each shard pair, so the MoE output exactly matches the dense model at the start of training.

Multi-Token Prediction (MTP)

MTP optimizes the model to predict multiple consecutive future tokens at each position, densifying the supervision signal. Unlike parallel independent predictions, it preserves causal dependencies between predictions through hidden-state transitions, which improves convergence and generation quality. During inference, the model falls back to ordinary single-token prediction for compatibility.

Flexible pipeline parallelism allows MTP layers to be placed strategically within the VPP layout. In the DeepSeek-V3 example with PP=16 and VPP=2, the MTP layer is placed in a standalone pipeline stage on PP rank 14 to balance the workload.

Muon Optimizer

Unlike AdamW, which performs element-wise updates, Muon applies a matrix-aware optimization by orthogonalizing entire weight matrices. The production integration provides:

1. Split QKV support: efficient orthogonalization even when attention projection matrices are stored as separate Q, K, and V tensors
2. Distributed optimizer integration: optimizer states are sharded across data-parallel ranks while preserving correct orthogonalization semantics
3. CPU offloading: Muon’s orthogonalization buffers can be offloaded when GPU memory is tight

MuonClip addresses a separate stability problem in trillion-parameter training, where query-key dot products can grow without bound and trigger attention explosions. The paper notes hardware-accelerated implementations in cuDNN, cudnn-frontend, and Transformer Engine.

10. Performance Results

GB200 and GB300 deliver roughly 3x higher token throughput than H100. Long-context DeepSeek-V3 at 256K tokens still reaches 88% of short-context MFU.

11. Case Study: DeepSeek-V3 GB200 vs H100

Key insight: the same model requires fundamentally different strategies on different hardware. GB200’s 192 GB memory and NVL72 topology largely eliminate the communication wall, shifting the bottleneck toward CPU overhead and making CUDA Graphs essential. H100’s 80 GB memory and NVL8 topology instead make cross-node communication the main bottleneck, so EP overlap becomes essential, and FP8 is what frees enough memory to hold the extra overlap buffers.

12. Systematic Optimization Workflow

A three-phase workflow emerges from tuning Mixtral, DeepSeek-V3, and Qwen3 across GB200 and H100. The process is inherently iterative: solving one bottleneck often exposes the next.

Phase 1: Establish Memory-Feasible Parallelism

Memory feasibility is the first hard constraint. The paper explicitly frames the impact of each parallelism strategy on per-GPU memory:

* Requires distributed optimizer.

Practical tip: use --fake-init-process-group to emulate distributed training on a single GPU for rapid iteration on parallelism configurations without allocating a full cluster.

Phase 2: Select Optimal Parallelism Strategy

Five guidelines:

1. Minimize model parallelism, maximize data parallelism: keep TP, EP, PP, and CP as small as possible while still avoiding OOM. Use the distributed optimizer (--use-distributed-optimizer) to shard optimizer states across DP ranks.
2. Keep EP and TP within the NVLink domain: make sure EP x TP fits inside the local NVLink island, typically 8 GPUs per node or 72 GPUs for NVL72. When scaling beyond that, prefer PP over stretching TP or EP across nodes.
3. Use pipeline parallelism for multi-node scaling: PP distributes layers across nodes. Enable VPP to reduce pipeline bubbles when PP >= 2, and balance work across VPP ranks.
4. Prefer EP over TP for expert layers: EP yields larger local GEMMs, lower communication overhead, simpler computation graphs, and eliminates local token permutation when EP = num_experts. The paper gives a concrete example: Mixtral-8x7B with EP8 x TP1 outperforms EP4 x TP2.
5. Enable context parallelism for long sequences: use CP when sequence length is at least about 8K tokens. For sequences shorter than about 4K, the CP overhead can exceed the benefit.

Phase 3: Profile and Optimize Bottlenecks

Diagnose which wall dominates, then apply targeted fixes.

Memory bottleneck: symptom is forced full recomputation or overly aggressive parallelism just to avoid OOM.

Communication bottleneck: symptom is a profile dominated by collectives.

CPU overhead bottleneck: symptom is Nsight Systems showing gaps between GPU kernels.

Computation bottleneck: symptom is low SM utilization even after communication and CPU gaps are under control.

Key insight: the same model on different hardware needs different optimization priorities. On NVL8, where EP crosses nodes, the Communication Wall dominates and can consume 30-50% of step time in all-to-all. On NVL72, where EP stays inside NVLink, enabling FP8 often shifts the bottleneck to CPU overhead instead.

Iterative Nature

The ordering matters: memory in Phase 1, then parallelism in Phase 2, then profiling in Phase 3. But the process is cyclical. Memory optimizations may enable smaller parallelism degrees, which pushes you back to Phase 1. Phase 3 optimizations such as EP communication overlap and CUDA Graphs also consume memory, which can force you to revisit earlier choices.

13. RL Post-Training Insights

• Router Replay: logs expert assignments during inference and replays them during training for more stable optimization
• Packing-aware dynamic batch size: keeps total effective tokens per batch more consistent
• Attention cost metric: sorts microbatches by $\sum (\text{seq_len})^2$ in serpentine order to reduce synchronization bubbles
• Dynamic CP: selects CP degree per microbatch rather than provisioning for the worst-case sequence mix

14. Conclusion

MoE sparsity introduces two fundamental challenges:

1. Parameter-compute mismatch, which creates the Three Walls of memory, communication, and compute efficiency
2. Dense-sparse mismatch, which requires decoupled parallelism between attention and expert layers

Key contributions summarized:

Final performance: DeepSeek-V3 reaches 1,233 / 1,048 TFLOPS/GPU on GB300 and GB200 with 256 GPUs, versus 368 TFLOPS/GPU on H100 with 1,024 GPUs. GB200 and GB300 deliver about 3x higher token throughput than H100.

The broader takeaway is that Megatron-Core MoE is a full-stack systems response to these two mismatches. Memory savings from FP8, recomputation, and offloading are not isolated wins; they enable communication overlap, better parallelism choices, and eventually CUDA Graph coverage. Large-scale MoE training becomes feasible only when routing, parallelism, kernels, optimizer states, and long-context execution are all tuned together.

Time moves so quickly — before I could even begin to count the days, you are already two years old.

Looking back to the night before you were born, it feels like only yesterday. That morning, your mother felt unusually intense movement, and contractions began arriving in waves. I hurriedly packed our bags and drove her to the hospital on the west side. The night that followed was long and heart-wrenching: in the ward, your mother lay tethered to tubes and monitors, holding on with the help of anesthesia and an epidural.

It wasn’t until the early hours of the next morning that the doctors began the induction. Your mother summoned every ounce of her strength, and finally — you arrived. I cut the umbilical cord myself, witnessing the very first moment of your independent life. But before the joy could even settle, a complication arose, and your mother was rushed back to the operating room for an emergency procedure. Thankfully, she came through safely. Meanwhile, you announced your arrival with a thunderous wail, and we — guided by patient nurses — clumsily learned to feed you and change diapers every three hours. On the day we left the hospital, the moment you were buckled into your car seat, your constant crying miraculously stopped. It was your first truce with this world.

Chapter II: The Chaos of New Parenthood

That first night home was the start of a frantic scramble. You cried so hard it felt like our hearts were being torn apart, and as rookie parents, we didn’t even realize you were simply hungry. For the entire first month, you woke every three hours, and we bid farewell to unbroken sleep.

From there, life shifted rapidly through countless firsts: we took you to see the cherry blossoms in Seattle; we watched your first rollover, your first crawl, your first wobbly steps. We heard you call out “Mama” and “Papa” for the first time in your small, stumbling voice. That year, you accompanied your mother to her PhD graduation ceremony and gazed up at the Aurora Borealis with us. Grandparents took turns flying in from afar to help care for you. You went through weaning off the pacifier, doctor’s check-ups, and gleeful swinging at the amusement park. Life unfolded in small details, but it became vivid because of you.

What remains most unforgettable was the end of 2024. At only nine months old, you followed us on cross-continental flights between China and the U.S., spending a month in our home country. I also remember a deep winter night when the power went out. With no warm milk to drink, you simply sat there in the dark, laughing and content. In that moment, your effortless optimism dissolved all our exhaustion.

Chapter III: Waiting at the Window

On your first birthday, we invited friends and colleagues to celebrate. You sat in a little push-car, spinning happily around the living room.

After turning one, you started daycare. With it came the inevitable cycle of illnesses and trips to the urgent care, but it was also where you learned to socialize and play with children your age. Because of work, I began traveling frequently between California and Washington. Your mother tells me that midweek, you often press yourself against the window, watching and waiting for Papa to come home. Because of this, I hold our weekends sacred — even if it’s just a walk in the park or taking you to your little soccer class on Saturday mornings.

At home these days, you are a perpetual motion machine. Aside from when you’re sleeping, you almost never stop — running laps between the kitchen and the living room. You’ve grown especially attached to your mother lately, occasionally declaring in your baby voice, “Don’t want Papa, want Mama.” It stings a little, I won’t lie. But I’m deeply comforted to see you expressing your feelings so boldly.

Chapter IV: A Second Birthday and an Ordinary Day

Your second birthday was spent, once again, in a rush. I flew back from California just the day before and hurried home from work at 5 PM on your birthday.

Your mother had set up a backdrop with your favorite excavator balloons, and you dragged them around for ages. We gathered to sing “Happy Birthday,” and you asked with a perfectly serious face, “Where did the cake go?” When the cake finally appeared, you blew out the candles with delight alongside Papa, Mama, and Grandma. You ate until your face was covered in cream — a little frosted kitten.

The celebration was simple and grounded. After blowing out the candles, we went to Costco as usual to restock for the week. The day after your birthday, we took you to a nearby park and then out for hot pot. You didn’t wake from your afternoon nap until late. Grandma and Mama spent the afternoon playing with you, and after dinner, I took you for one last walk by the water. As the sun set and the sky darkened, watching you run and pause along the shore, I felt a tranquility that was almost unbearably precious.

Chapter V: The Cycles of Time

Time flies. I often think that in perhaps sixteen years, you will be as I once was — packing your bags and leaving this warm little nest to chase your own mountains and seas.

Today, I find myself having lived through another full cycle: seventeen years since I first left my hometown at eighteen. The world is changing fast, especially amid the surging wave of AI, where everything seems to accelerate. But I know that some things never change.

I will cherish this time together. I hope to walk beside you in health and happiness through all the years to come.

Not to squander our time; not to waste the beauty of these years.

• FlashAttention-4: Algorithm and Kernel Pipelining Co-Design for Asymmetric Hardware Scaling

Overview

FlashAttention-4 is the latest iteration of the wildly successful hardware-aware attention algorithm, designed specifically to tackle the architectural shifts in NVIDIA’s Blackwell (B200/GB200) GPUs. While FlashAttention-3 was heavily optimized for the Hopper architecture (H100), the transition to Blackwell introduces a phenomenon “asymmetric hardware scaling”. This paper provides a masterclass in algorithmic and kernel co-design, demonstrating how to shift compute paradigms when the hardware bottlenecks unexpectedly change.

The Core Problem: Asymmetric Hardware Scaling

The fundamental challenge addressed in FlashAttention-4 is that not all parts of the GPU got faster at the same rate. On the Blackwell B200, the Matrix Multiply-Accumulate (MMA) tensor core throughput doubled compared to Hopper, reaching a massive 8192 ops/clock/SM (2.25 PFLOPS for FP16/BF16). However, the shared memory (SMEM) bandwidth remained flat at 128 bytes/clock/SM, and the multi-function unit (MUFU), which handles exponential operations for softmax, remained at 16 ops/clock/SM.

Roofline analysis in the paper reveals a surprising reality for Blackwell: matrix multiplication is no longer the primary bottleneck for attention. Instead, shared memory traffic and exponential operations now dominate execution time, exceeding MMA compute by 25-60%.

Key Technical Innovations

1. Software-Emulated Exponential & Conditional Rescaling

Because the hardware MUFU cannot keep up with the doubled tensor core speed, the exponential calculation in the softmax becomes a major chokepoint.

• Polynomial Approximation: To bypass the MUFU bottleneck, it leveraged a software emulation of $2^x$ using floating-point FMA units via polynomial approximation. By distributing the exponential computations across both the MUFU and FMA units (using emulation for 10-25% of entries), they effectively increased the exponential throughput. The genius here is recognizing that while a degree-3 polynomial has a higher FP32-level error than the hardware MUFU, once the output is rounded to BF16 (the standard precision for attention), the quantization error dominates, making the software emulation virtually indistinguishable from hardware.

• Conditional Rescaling: FlashAttention relies on online softmax, which requires rescaling previous results when a new maximum value is encountered. FlashAttention-4 introduces a threshold-based skip mechanism: if the new maximum is not significantly larger than the old one, it skips the vector multiplication for rescaling and resolves the normalization at the very end.

2. Taming Shared Memory with 2-CTA MMA Mode

In the backward pass, SMEM bandwidth becomes an acute bottleneck because the algorithm requires five different MMA operations, forcing operands to be repeatedly read from shared memory.

To solve this, FlashAttention-4 leverages a new Blackwell feature: 2-CTA tensor core MMA mode. In this mode, two Cooperative Thread Arrays (CTAs) within the same cluster cooperatively execute a single MMA.

• Halving SMEM Traffic: Each CTA stages only half of operand B in its own shared memory, while the hardware consumes the combined B tile during the multiply.

• Halving Atomic Reductions: In the gradient accumulation step (dQ), the researchers use distributed shared memory (DSMEM) to exchange half of the gradient of the softmax (dS) between the two CTAs. This repacking allows each CTA to write only half of the dQ tile, cutting the number of expensive global atomic reductions in half.

3. Exploiting TMEM and Full Asynchrony

Blackwell introduces Tensor Memory (TMEM), a 256 KB on-chip memory per SM specifically for storing intermediate tensor core results. Unlike Hopper, where MMAs wrote to registers and caused massive register pressure, Blackwell’s MMAs write asynchronously directly to TMEM. FlashAttention-4 redesigns the software pipeline to aggressively overlap the fully asynchronous tensor core operations with softmax and memory operations, utilizing the larger 128x128 MMA tiles.

4. A Pythonic Shift: CuTe-DSL

From a developer ecosystem standpoint, one of the most exciting updates is the move away from notoriously slow C++ template metaprogramming. FlashAttention-4 is written entirely in CuTe-DSL embedded in Python. This framework lowers to PTX and compiles just-in-time (JIT), reducing compile times by 20-30x (down to ~1.4–2.5 seconds from 45–55 seconds) while retaining full low-level expressivity.

Performance Outcomes

By directly addressing the shifting hardware bottlenecks, FlashAttention-4 achieves impressive performance on B200 GPUs:

• Up to 1.3x speedup over cuDNN 9.13 and 2.7x over Triton for BF16.
• Achieves up to 1613 TFLOPs/s, utilizing 71% of the theoretical maximum compute.

Insights

1. The End of “Compute-Bound” Attention (For Now): We are entering an era where keeping the tensor cores fed is harder than the actual matrix multiplication. Kernel developers can no longer treat GPUs as uniform compute scaling machines; they must acutely monitor the ratio of MMA throughput to memory bandwidth and non-linear function units.

2. Software Emulation is Viable at Lower Precisions: The decision to emulate the exponential function using FMA units is a brilliant application of precision-aware optimization. It highlights a broader insight for AI system design: if your target data type (like BF16) has a high quantization error, you have “budget” to use cheaper, faster mathematical approximations without harming the end result.

3. Python for Low-Level GPU Kernels is Maturing: The use of CuTe-DSL proves that writing bare-metal, highly optimized GPU kernels no longer strictly requires the painful compilation cycles of complex C++ libraries like CUTLASS. This lowering of the barrier to entry will likely accelerate community experimentation with new attention variants.