Blackwell Brings NUMA Back, Inside One GPU

DeepSeek addressed the dual-die memory locality of Nvidia's Blackwell (SM100) B200 and B300 GPUs in a pull request to its DeepGEMM library (deepseek-ai/DeepGEMM#462), splitting MoE weights along the N dimension so each of the two dies joined by NV-HBI stores its own half in local HBM. Because CUDA exposes no API identifying which SM sits on which die, DeepGEMM measures it at startup with a pointer-chasing probe timed via clock64, accepting the result only if the far die is at least 1.25x slower and both SMs in every TPC agree, retrying up to 5 times before falling back to an even split. The code allocates die-local memory through MLOPart, an MPS mode, using cuMemAddressReserve and cuMemMap, pending direct per-die allocation that the post says should arrive in CUDA 13.4.

On Blackwell SM100 we are back to solving NUMA problems, inside a single GPU package. To software, a B200 or B300 looks like one big GPU. CUDA shows one device with a flat address space and a unified L2, and you launch kernels across all the SMs just like you did on Hopper. Physically, though, it is two separate silicon dies joined by NV-HBI. The link is fast, but HBM attached to your own die is still noticeably closer than HBM attached to the other one. If you ignore this, memory-heavy kernels like MoE layers pull half their weights across the link, which costs latency and power. DeepSeek recently handled this in DeepGEMM for their MegaMoE kernel deepseek-ai/DeepGEMM 462 https://github.com/deepseek-ai/DeepGEMM/pull/462 . All the code below comes from that implementation. Two locality domains two-locality-domains Each die is a locality domain with its own SMs and its own HBM stacks. CUDA has no API that tells you which SM sits on which die, so DeepGEMM measures it at startup. It puts one buffer on each die and has every SM chase a chain of pointers through both, timing the hops with clock64 : template <uint32 t kNumHops, uint32 t kNumChunksPerSM, uint32 t kNumChunkBytes global launch bounds 32, 1 void locality domain probe chase const uint32 t buf, uint16 t out { if threadIdx.x = 0 return; const auto sm idx = ptx::get sm idx ; for uint32 t chunk idx = 0; chunk idx < kNumChunksPerSM; ++ chunk idx { const auto chunk = buf + static cast<uint64 t sm idx kNumChunksPerSM + chunk idx kNumChunkBytes / 4 ; uint32 t idx = ptx::ld global cg chunk ; const auto start = clock64 ; pragma unroll for uint32 t i = 0; i < kNumHops; ++ i idx = ptx::ld global cg chunk + idx ; const auto cycles = clock64 - start; out sm idx kNumChunksPerSM + chunk idx = idx == 0xffffffffu ? 0u : static cast<uint16 t cycles / kNumHops ; } } The probe launches one 32-thread block per SM, and each block reserves the full capacity of shared memory so no SM can run two of them at once. Only thread 0 does the chase. Each 4 KB chunk is 32 cache lines, and each line points 7 lines ahead, so every load waits on the one before it. The loads use ld.global.cg to bypass L1, and L2 is flushed beforehand, so the cycles measure HBM round-trip latency. Each SM is assigned to whichever die it reads faster, taking the median over 8 chunks. The measurement is only accepted if the far die is at least 1.25x slower across the board and both SMs in every TPC agree. The two SMs in a TPC are physically paired, so disagreement means the measurement is noise. DeepGEMM retries up to 5 times and otherwise falls back to an even split. And because the scheduler expects an exact 50-50 balance, asymmetry from harvested SMs is fixed by reassigning a few TPCs to the smaller side. Slicing the weights across dies slicing-the-weights-across-dies Once you know where each SM lives, the next problem is memory placement. Instead of storing MoE weights as one big flat array, MegaMoE splits each tensor along N. So B, N, K becomes 2, B, N/2, K , with slice d stored in die d’s HBM: php def localize t: torch.Tensor, dim: int = -2 - torch.Tensor: num domains, dim = get num locality domains , dim % t.dim assert t.size dim % num domains == 0 localized = empty per domain t.shape :dim , t.size dim // num domains, t.shape dim + 1: , t.dtype localized.copy t.unflatten dim, num domains, -1 .movedim dim, 0 return localized Behind the scenes, empty per domain reserves one virtual address range and maps each half onto physical memory on its own die using the CUDA virtual memory APIs cuMemAddressReserve , cuMemMap . CUDA cannot allocate on a specific die directly yet—that should arrive in CUDA 13.4. Until then, DeepGEMM gets die-local memory through MLOPart, an MPS mode that exposes each die as a separate CUDA device. DeepGEMM spins up a private MPS server in a helper process, does the allocation there, and passes each chunk back to the main process as a shareable file descriptor. Affinity scheduling and work stealing affinity-scheduling-and-work-stealing Having localized data only helps if your thread blocks run on the right SMs. MegaMoE uses a persistent scheduler where each die gets its own task counter. A task pulled from die d’s counter maps to die d’s half of N, so by default an SM reads weights only from its own HBM. When its own queue runs dry, an SM can steal from the other die, but only if the layer has at least two waves of work. For smaller layers, stealing is skipped entirely: template <uint32 t kNumClusters static CUTLASS DEVICE bool get next task idx const uint32 t task count ptr, const uint32 t& locality domain idx, const uint32 t& num m blocks, uint32 t& m block idx, uint32 t& n cluster idx { constexpr uint32 t kNumLocalClusters = kNumClusters / kNumLocalityDomains; // One lane does an atomic add on this die's counter and broadcasts the result const uint32 t task idx = get next local task idx task count ptr, locality domain idx ; if task idx = num m blocks kNumLocalClusters return false; m block idx = task idx / kNumLocalClusters; // Clusters d kNumLocalClusters, d + 1 kNumLocalClusters cover die d's half of N n cluster idx = locality domain idx kNumLocalClusters + task idx % kNumLocalClusters; return true; } template <uint32 t kNumClusters CUTLASS DEVICE bool get next task idx const uint32 t task count ptr, const uint32 t& num m blocks, uint32 t& m block idx, uint32 t& n cluster idx const { // 1. This SM's own die first if get next task idx<kNumClusters task count ptr, sm locality domain idx, num m blocks, m block idx, n cluster idx return true; // 2. Steal from the other die only with 2+ waves of work one wave = kNumSMs / 2 two-SM clusters return kNumLocalityDomains 1 and num m blocks kNumClusters = kNumMinStealWaves kNumSMs / 2 and get next task idx<kNumClusters task count ptr, sm locality domain idx ^ 1, num m blocks, m block idx, n cluster idx ; } Stealing across dies means higher memory latency and a busier link. For a few leftover tiles it is faster to let the local die finish them; with two or more waves of work, the steal pays for its penalty. Hardware addressing with TMA hardware-addressing-with-tma The die split is also encoded in the Tensor Memory Accelerator TMA descriptor. The weight tensor map is 4D, and one dimension is the die index: K, N/2, die, expert . Instead of pointer math, each CTA picks the half with a coordinate: js // Each CTA of the 2-SM cluster loads its own BLOCK N slice const auto n block idx = task info.n cluster idx 2 + is leader cta ? 0u : 1u ; // half idx: which die's slice holds this N block const uint32 t half n = shape n / 2, n idx = n block idx BLOCK N, half idx = n idx = half n; const uint32 t half n idx = n idx - half idx half n, expert idx = task info.is shared ? 0 : task info.local expert idx; // TMA coordinates: k, n within the half, die, expert tma::copy nd<BLOCK K, LOAD BLOCK N, kSwizzleBMode, weight dtype t, 4 tensor map b ptr, &shared storage.full barriers stage idx , shared storage.smem b stage idx , 2, // 2-CTA load: completion is signaled to the leader CTA's barrier k idx, half n idx, half idx, expert idx ; The stride along the die dimension steps directly across the physical split in memory, so half idx decides which die serves the load. Normally that coordinate matches the SM’s physical die, so traffic stays local. Only stolen tasks, or SMs on reassigned TPCs, read across NV-HBI. The inner loop needs no branches or address recalculation. DeepSeek’s release notes mention this saves about 100 W of package power and gives 5-10% higher GPU clocks under heavy load. As chips scale by stitching dies together, the hardware gets less uniform, not more. Treating one GPU as two locality domains takes more engineering up front, and it pays for itself.