Introduction to NVIDIA latest Rubin GPU

⚠️ Disclaimer: statistics from this post all come from public contents.

Rubin is NVIDIA's latest Tesla-series GPU expected to be launched in 2026. In this post, I'm going to introduce some of the highlights in Rubin platform differing from Blackwell. It's more like a minor revision of Blackwell, not a game-changer, it's left in the next Feynman leap.

Hardware evolution overview

This table ^{1 2 3 4} highlights the hardware unit evolution compared to the previous two generations (the perf metrics are all Tensor Core dense based). We can see NVIDIA officially sacrifices FP64 perf numbers to yield more areas for AI favored low-precision computing (FP4 / FP8). Besides, there are also huge jumps in the HBM and scale-out NVLink(-C2C) bandwidth.

Software revision to match the hardware changes

Doubled K-dimension in tcgen05 MMA

Rubin doubles the MMA-K size of the tcgen05 instruction for supported low-precision dtypes. The motivation behind this is, even with tensor parallel, the K dimension in first GEMM of MLP or MoE layers is unchanged:

Image source: https://lightning.ai/docs/fabric/2.4.0/advanced/model_parallel/tp.html

So as the model size gets bigger, the large K dimension (i.e., hidden_dim) hinders the GEMM process in the compute-bound FFN / MoE layers.

From NVIDIA official blog ¹

PTX 9.4 docs mention the register descriptor tcgen05 MMA instruction now contains a bit indicating the K size in position 29.

For .kind::f16, .kind::tf32, .kind:f8f6f4:

• Bit 29 for K dimension is valid for architecture sm_107f or higher in the same family. ²

For supported dtypes, the number of MMA K-iterations can be reduced, thus Rubin is able to provide higher overall throughput in the same instruction cycles.

Increased SFU for Softmax

Exponential throughput rises from Blackwell to Rubin ¹:

The exponential computation in Softmax has become the bottleneck of attention, so NVIDIA spent more hardware units here. Another potential solution mentioned in FA4 ⁵ is harnessing both native EXP (SFU) and software-emulated EXP (CUDA core) via Taylor expansion interleaved ⁶ to compute Softmax.

Fine-grained block-level PDL

Since Hopper, CUDA offers PDL (Programmatic Dependent Launch) capability to issue the second dependent kernel in advance so that the prologue can overlap over the final exit of the first kernel. The first kernel informs its successor via cudaTriggerProgrammaticLaunchCompletion, and the second one firstly conducts some independent prologue then waits for the signal by cudaGridDependencySynchronize.

Block-level PDL.

Rubin extends this ability at a finer-grained level: the block level. Now each CTA in the second kernel can start immediately once the corresponding dependent blocks (CTAs) have released the flag.

This sounds more like the concept of MegaKernel, with this finer-grained PDL, it seems that we can achieve the goal of MegaKernel without messing the codes across the logic boundary.

In-flight TMA descriptor update

Currently (Hopper / Blackwell), to patch the TMA descriptor at the device side, we have to update a series of metadata in shared memory with additional memory barrier synchronization cost, like the following pseudo PTXs ²:

.global .align 128 .b8 gTensorMap[128];   // host-initialized descriptor
.shared .align 128 .b8 sTensorMap[128];

// 1. stage the 128 B descriptor into shared memory
cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes
    [sTensorMap], [gTensorMap], 128, [mbar];
// ... wait on mbar ...

// 2. patch the per-expert fields in place
tensormap.replace.tile.global_address.shared::cta.b1024.b64  [sTensorMap], newAddr;
tensormap.replace.tile.global_dim.shared::cta.b1024.b32      [sTensorMap], 0, newDim;
tensormap.replace.tile.global_stride.shared::cta.b1024.b64   [sTensorMap], 0, newStride;

// 3. copy back to global memory, releasing to the tensormap proxy
tensormap.cp_fenceproxy.global.shared::cta.tensormap::generic.release.gpu.sync.aligned
    [gTensorMap], [sTensorMap], 128;

// 4. acquire in the tensormap proxy before any TMA unit reads it
fence.proxy.tensormap::generic.acquire.gpu [gTensorMap], 128;

// 5. the descriptor is now consumable
cp.async.bulk.tensor.2d.shared::cluster.global.tile
    [sMem], [gTensorMap, {x, y}], [mbar];

Or there is a corresponding series of CuTe DSL tensormap manager APIs:

tensormap_manager = utils.TensorMapManager(
    utils.TensorMapUpdateMode.SMEM,      # or .GMEM
    bytes_per_tensormap,                 # 128
)
a_ptr = tensormap_manager.get_tensormap_ptr(tensormaps[(wsp_idx, 0, None)].iterator)

tensormap_manager.init_tensormap_from_atom(tma_atom_a, a_init_ptr, tma_warp_id)
tensormap_manager.fence_tensormap_initialization()

tensormap_manager.update_tensormap(
    (real_tensor_a, real_tensor_b),      # new gmem tensors for this group
    (tma_atom_a, tma_atom_b),
    (a_ptr, b_ptr),                      # gmem descriptor slots
    tma_warp_id,
    (a_smem_ptr, b_smem_ptr),
)
tensormap_manager.fence_tensormap_update(a_ptr)

Therefore, updating the TMA descriptor inside kernel hinders the end-to-end performance for latency-sensitive cases, e.g., a kernel in MoE layer processes multiple experts whose weights are located in different GMEM addresses.

Shared TMA descriptor via in-flight metadata update ¹.

Since Rubin, CUDA supports inline override for some fields of the tensormap struct while sharing other common fields. PTX ISA 9.4 adds the .override::global_address and .override_attribute qualifiers (requiring sm_107f) to:

• cp.async.bulk.tensor
• cp.reduce.async.bulk.tensor
• cp.async.bulk.prefetch.tensor

Note that all experts of the MoE layer basically share the same shape, stride and dtype, with only the base memory address varying. The overriding values are passed as extra operands right after tensorMap, so the descriptor itself stays untouched.

This could be combined with another feature introduced in Rubin: L2 cache control. Since Ampere, NVIDIA allows kernel authors to specify L2 cache line residual hints for normal LD instructions ⁷:

• .ca: cache at all levels
• .cg: cache at L2, not L1
• .cs: cache streaming, evict-first policy in L1 and L2
• .lu: last use, acting like .cs on global address space
• .cv: bypass cache

Rubin supports some of these cache residual hints by introducing qualifiers to cp.async.bulk.prefetch{.tensor} and applypriority.async.bulk{.tensor} instructions. PTX 9.4 exposes them through a .level::eviction_priority qualifier on the bulk prefetch instructions, plus a pair of new applypriority.async.bulk instructions to apply the policy on a certain memory range, all requiring sm_107f ²:

// prefetch into L2 and pin the lines
cp.async.bulk.prefetch.tensor.dim.L2.src{.load_mode}{.level::eviction_priority}
    [tensorMap, tensorCoords]{, im2colInfo};
cp.async.bulk.prefetch.L2.src{.level::eviction_priority}   [srcMem], size;

.level::eviction_priority = { .L2::evict_last }

// retag resident lines once they are no longer worth keeping
applypriority.async.bulk.tensor.dim{.src}.completion_mechanism{.load_mode}
    .level::eviction_priority   [tensorMap, tensorCoords]{, im2colInfo};
applypriority.async.bulk{.src}.completion_mechanism.level::eviction_priority   [a], size;

.level::eviction_priority = { .L2::evict_normal }

Take MoE layer as example again, if all tokens for an expert have finished execution, the kernel can mark it as "last used".

ACK-free GIN by counted writes

Hardware-supported counter-based peer write.

CUDA / NCCL now supports network operations from the device ⁸. To inform the peer side that the transmission has finished, the sender side has to launch another flag write to peer memory with release semantic.

Rubin now supports a hardware-based counter that records the number of bytes the device receives across the NVLink. Once the counted bytes have reached the total bytes, the receiver can answer the consumer polls without the atomic flag write from the sender.

FP8 GEMM B-LUT

From "Figure 280" in https://docs.nvidia.com/cuda/developer-preview/13.4/pdf/ptx_isa_9.4.pdf

Rubin also supports FP8 GEMM with the pre-compressed matrix B, where TMEM stores the 8-bit pre-compressed B values (acting like a LUT) with capacity equal to 8, indexed by 3 bits ($2^3 = 8$).

PTX 9.4 for Rubin adds a new .decompress::lut::b qualifier for the tcgen05.mma instruction: if specified, the tensor core will read 48 bytes from SMEM and leverage the first or second half of them ($3 \times 64 \div 8 = 24$ bytes) as the LUT for MMA.

Outro

There are some other new features like attention sparsity by dynamically reading the intermediate $\mathbf{Q}\mathbf{K}^\mathrm{T}$ result into the 2:4 sparsity format. But this affects the overall attention precision with the probability distribution. The 2:4 structure is too inflexible for runtime sparsity, so I think it's still hard to be utilized in real workloads.