Deep Dive - DeepGEMM

Four experts receive [64, 192, 32, 128] token rows. There are 416 useful rows, but rounding each group up to a multiple of 128 reserves 640 rows. Should we pay for that padding or launch four smaller matrix multiplications? Separate launches avoid this particular packing requirement, but still have edge tiles and launch overhead. Counting useful arithmetic alone won't tell us which wins.

This is the feed-forward counterpart to the previous chapter's irregular attention workload. A mixture-of-experts (MoE) layer in a large language model (LLM) routes tokens to selected expert networks. A rank, the process controlling one GPU in our example, must compute each expert's projection with however many rows arrive.

DeepGEMM is a CUDA library for exploring that trade-off. It isn't a model server: it doesn't admit requests, own a token budget, or expose HTTP. It supplies dense and grouped GEMM kernels, low-bit layouts, multi-query attention (MQA) scoring for DeepSeek's lightning indexer, and Mega MoE, which overlaps expert-parallel communication with expert compute.^{[1]Reference 1DeepGEMMhttps://github.com/deepseek-ai/DeepGEMM} Other serving systems call it. A kernel can win its microbenchmark and still leave end-to-end latency unchanged if the surrounding runtime spends the saved time elsewhere.

We'll keep the public Flash benchmark shape in view: 256 routed experts, top-k 6, hidden size 4096, and intermediate size 2048. Under eight-way expert parallelism (EP8), a balanced placement gives 32 experts per rank. Our four groups are a small subset used to understand packing, not the complete rank workload.^{[2]Reference 2Mega MoE Benchmark Resultshttps://github.com/deepseek-ai/DeepGEMM/pull/316}

The implementation discussion is pinned to commit 559d79fb6994a58b8a15b4b93bf13ccc16edf247 (July 15, 2026), verified against the official repository on September 2, 2026.^{[3]Reference 3DeepGEMM Source Snapshot, July 15 2026https://github.com/deepseek-ai/DeepGEMM/tree/559d79fb6994a58b8a15b4b93bf13ccc16edf247} Runnable examples below use the CPU and standard library. They check arithmetic and dependencies, not CUDA execution, FP8 rounding, synchronization, or GPU speed.

What problem does DeepGEMM solve, and what problem does it leave to a serving runtime?

Answer

DeepGEMM solves selected GPU tensor operations and their data movement. A serving runtime still owns request scheduling, batching, admission, model loading, and end-to-end latency policy.

Project identity and system boundary

The repository is maintained in the DeepSeek AI open-source ecosystem. Its README describes a unified tensor-core kernel library and lists DeepSeek contributors. The Mega MoE release notes name a larger group because communication kernels, layouts, scheduling, and benchmarking cross subsystem boundaries.^{[1]Reference 1DeepGEMMhttps://github.com/deepseek-ai/DeepGEMM[4]Reference 4Mega MoE: Fusing Expert Parallel Communication with Computationhttps://github.com/deepseek-ai/DeepGEMM/pull/304}

That boundary also explains DeepGEMM's relationship to CUTLASS and CuTe. It borrows ideas from both while keeping a smaller set of public kernel functions. The narrow surface makes one generated path easier to inspect, but callers still have to prepare layouts and scales. No fallback can make a malformed tensor fast.

What callers actually use it for

Follow one routed token through the graph and four kernel families answer different bottlenecks:

1. Dense projection layers. Attention and feed-forward layers reduce to matrix products. DeepGEMM exposes BF16 and low-bit GEMMs for these hot loops.
2. Mixture-of-Experts. Grouped GEMMs process the variable number of tokens sent to each expert. Contiguous and masked layouts target prefill and decode differences.
3. DeepSeek indexer scoring. MQA scoring turns query vectors, compressed key/value vectors, and per-head weights into token-to-token logits. The kernel has non-paged and paged forms for prefill and decode, with an FP4 indexer path in later releases.^{[6]Reference 6DeepSeek-V3.2: Pushing the Frontier of Open Large Language Modelshttps://arxiv.org/abs/2512.02556[4]Reference 4Mega MoE: Fusing Expert Parallel Communication with Computationhttps://github.com/deepseek-ai/DeepGEMM/pull/304}
4. Expert-parallel inference. Mega MoE combines dispatch, two expert projections, SwiGLU, and combine so NVLink traffic can overlap tensor-core work.^{[4]Reference 4Mega MoE: Fusing Expert Parallel Communication with Computationhttps://github.com/deepseek-ai/DeepGEMM/pull/304}

The README also lists HyperConnection (HC) GEMMs. We won't trace every primitive: dense GEMM establishes the memory path, grouped GEMM adds uneven routing, and Mega MoE adds communication. The indexer is a useful contrast because its output isn't an expert activation or a full attention result.

Research foundations and source trail

No single paper explains every DeepGEMM kernel. Read the sources by the question they answer: DeepSeek-V3 motivates sparse expert routing, DeepSeek-V3.2 describes the lightning indexer, DeepSeek-V4 supplies the Flash and Pro shapes, and the repository plus release pull requests record the implementation.^{[7]Reference 7DeepSeek-V3 Technical Report.https://arxiv.org/abs/2412.19437[6]Reference 6DeepSeek-V3.2: Pushing the Frontier of Open Large Language Modelshttps://arxiv.org/abs/2512.02556[8]Reference 8DeepSeek-V4: Towards Highly Efficient Million-Token Context Intelligencehttps://arxiv.org/abs/2606.19348}

Use the sources in four layers:

A paper can motivate a workload without proving a kernel speedup. A benchmark can show a speedup without proving model quality. Keep those evidence types separate as we trace the code.

GEMM first: the contract behind the speed

General matrix-matrix multiplication, or GEMM, computes a matrix product and optionally adds an existing output. DeepGEMM's public naming convention is $D = C + A B$, and the common NT interface (non-transposed A, transposed B) implements $D = C + A B^{\top}$.^{[1]Reference 1DeepGEMMhttps://github.com/deepseek-ai/DeepGEMM}

For that NT layout, $A$ has shape $[M,K]$, $B$ has shape $[N,K]$, and $D$ has shape $[M,N]$. Each output element is a dot product of one row from A and one row from B. The operation count is about $2MNK$ floating-point operations because each multiply is paired with an add.

For one output, take A's row [1, 2] and B's row [3, 4]: their dot product is $1\times3+2\times4=11$. Pair A's next row [0, -1] with the same B row and the result is $-4$. A matrix multiplication repeats this operation over every row pair. B is stored as [N,K]; the transpose in the formula makes the inner K dimensions meet.

Now take a whole GEMM with $M=512$, $N=4096$, and $K=4096$. It carries $8{,}589{,}934{,}592$ multiply-accumulate pairs, or $17{,}179{,}869{,}184$ FLOPs when multiply and add each count as one. At a hypothetical 10 microseconds, that would be about 1,718 TFLOP/s. No timing was measured for this example. Also, these are whole-matrix dimensions, not one thread block's tile size.

The hardware sees more than arithmetic. It sees global-memory loads, shared-memory staging, scale-factor loads, register or tensor-memory pressure, barriers, and output stores. A kernel can have high arithmetic intensity and still lose if a layout forces extra transposes or if a small M dimension leaves tensor cores idle.

Why can two GEMMs with the same FLOP count have different latency?

Answer

Their tile shapes, memory layouts, scaling-factor loads, occupancy, and synchronization can differ. FLOP count describes arithmetic work, not data movement or parallelism.

Read the performance ladder before the kernel

Take the $512 \times 4096$ output from the previous example and implement it without a library. A naïve kernel gives one thread one output element. Each thread walks all $K$ values and repeatedly requests overlapping A and B data from global memory. The result can be correct, but redundant loads and launch overhead are now exposed.

A tiled kernel gives a thread block a rectangle of D. It loads matching A and B rectangles into shared memory, then reuses those values while the K loop advances. Threads keep one or more partial sums in registers, close to the arithmetic. Shared memory cuts repeated global loads; registers keep the accumulator nearby. Both are finite resources, so a larger tile or more accumulators can reduce the number of active blocks through shared-memory or register pressure. If the register budget is exceeded, values can spill to off-chip local memory and erase the reuse the tile was meant to create.^{[9]Reference 9CUDA C++ Best Practices Guidehttps://docs.nvidia.com/cuda/cuda-c-best-practices-guide/}

Tensor cores change the inner loop again. The kernel feeds matrix fragments to specialized matrix-multiply-accumulate units, carries partial sums across K tiles, and converts the accumulator when it stores D. Tile shape, operand layout, input precision, and accumulator type must match the architecture. DeepGEMM's TMA, TMEM, and tcgen05 path later in this chapter is this third stage made explicit for SM100, not a way around shape and layout checks.^{[10]Reference 10CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/[1]Reference 1DeepGEMMhttps://github.com/deepseek-ai/DeepGEMM}

Make one prediction before profiling. With $M=1$, much of a large output tile has no useful row work, so launch and staging overhead matter more. With $M=512$, a block can reuse more loaded data and keep more tensor-core work in flight. If a tile change raises register use, occupancy may fall without the kernel getting slower or faster in a predictable direction. Measure the change rather than treating occupancy as the target.^{[9]Reference 9CUDA C++ Best Practices Guidehttps://docs.nvidia.com/cuda/cuda-c-best-practices-guide/}

The performance receipt belongs beside the number. Record hardware and software, workload shape and token distribution, precision and algorithm, exact baseline, and the correctness method. Without those five labels, “faster” describes a run but not a result another engineer can reproduce.

Correctness comes before the receipt is useful. Compare the same inputs against a PyTorch or legacy reference with an explicit tolerance, and include one-token, ragged, edge-tile, and scale-boundary cases. Check output order and routed-token indices as well as the numerical error. A high TFLOPS value with a wrong expert row is a faster bug. NVIDIA's Compute Sanitizer can then check out-of-bounds or misaligned access, shared-memory hazards, uninitialized reads, and invalid synchronization while the numerical reference checks the answer.^{[11]Reference 11Compute Sanitizerhttps://docs.nvidia.com/compute-sanitizer/ComputeSanitizer/}

Precision and scaling

DeepGEMM's low-bit paths trade input precision for bandwidth and tensor-core throughput. FP8 here is E4M3: four exponent bits and three mantissa bits, commonly used for activations. FP4 here is E2M1 packed two values per byte, used for expert weights on the SM100 path. BF16 is a 16-bit operand and output format whose exponent range matches FP32. DeepGEMM's BF16 tensor-core paths keep matrix accumulators in FP32 before converting them to the configured BF16 or FP32 output.^{[1]Reference 1DeepGEMMhttps://github.com/deepseek-ai/DeepGEMM}

The kernel consumes scale factors in a specific layout. For the low-bit GEMM paths discussed here, SM90 uses FP32 scales and SM100's native representation packs four UE8M0 scales into one 32-bit torch.int. UE8M0 encodes unsigned powers of two with eight exponent bits and no mantissa. The API can prepare some scale layouts, so distinguish the native kernel representation from accepted wrapper inputs. Other operations, including the older FP8 indexer, have different scale contracts.^{[1]Reference 1DeepGEMMhttps://github.com/deepseek-ai/DeepGEMM}

Group size isn't universal either. SM90 FP8 paths commonly use K-group 128 with FP32 scales. The SM100 FP8×FP4 Mega MoE path uses K-group 32 with packed UE8M0. For a per-32 K-group, a row of K values splits into groups of 32 and each group gets one scale. A correct numerical tensor with the wrong stride can still fail an assertion or trigger a slow preparation path. The Tensor Memory Accelerator (TMA) stages compatible multidimensional tiles into shared memory, so its alignment requirements become part of the kernel contract too.

Treat casts and layout transforms as a pipeline stage before the kernel. DeepGEMM provides helpers such as transform_sf_into_required_layout, but the README warns that input transposition and FP8 casting may be slower utility operations. In a production path, fuse those operations into an earlier model kernel when profiling shows they dominate.

Why a scale belongs inside the K reduction

Suppose the stored rows are [1, 2, 3, 4] and [2, 1, -1, 2]. Use two K groups, each of width two, with A scales [0.5, 2] and B scales [2, 0.25]. The first group's raw dot is 4 and its scale product is 1. The second group's raw dot is 5 and its scale product is 0.5. Adding the scaled partial dots gives $4+2.5=6.5$, not the unscaled dot 9. A single scale applied after the whole reduction loses that distinction.

This CPU reference checks the grouping rule by comparing groupwise accumulation with explicit dequantization. Group width two is deliberately small for inspection; it isn't a supported DeepGEMM low-bit recipe.

Two architectures, two execution contracts

The repository supports NVIDIA SM90 (Hopper) and SM100 (Blackwell). Treat those as different execution contracts, not interchangeable names for a single GPU. A path selected for one family may use different scale types, layouts, instructions, and compiler flags on the other.^{[10]Reference 10CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/}

The dispatch layer checks the device architecture, input major layout, scale dtype, and output shape before selecting an implementation. On SM90, the current FP8 path uses FP32 scales and supports the NT layout for the primary GEMM interface. On SM100, the FP8 × FP4 path can select NT, TN, NN, or TT variants and uses packed UE8M0 scales.^{[1]Reference 1DeepGEMMhttps://github.com/deepseek-ai/DeepGEMM}

A single portable kernel could hide those differences behind extra branches or generalized iterators. DeepGEMM keeps architecture-specific implementations visible, then uses a thin API layer to check the contract and route to the matching implementation.

An architecture mismatch is a contract failure, not a tuning problem. If an SM100-only FP4 path is sent to SM90, no tile-size tweak can fix it. Check torch.cuda.get_device_capability(), CUDA version, scale dtype, and major layout before comparing benchmarks. At the pinned revision, the Mega MoE implementation dispatches to SM100 and its test contains a TODO about skipping SM90. That comment isn't an implemented automatic skip.

Inside an SM100 tile

The dispatch decision leaves one tile to follow. Global memory holds the source tensors, TMA stages a compatible rectangle, and the tensor-core path accumulates partial products before the epilogue stores D. Each handoff has a different lifetime and synchronization rule, so a correct matrix formula isn't enough.

TMA moves the right rectangle

TMA, the Tensor Memory Accelerator, copies multidimensional tensor regions between global memory and shared memory using a descriptor. A descriptor carries shape and stride information so the copy engine can move a tile without every thread computing addresses.

TMA reduces address instructions and stages a copy asynchronously while a previous tile is being consumed. That lets a kernel overlap global-memory movement with tensor-core math, provided barriers and shared-memory stages are sized correctly.^{[9]Reference 9CUDA C++ Best Practices Guidehttps://docs.nvidia.com/cuda/cuda-c-best-practices-guide/[10]Reference 10CUDA Programming Guide.https://docs.nvidia.com/cuda/cuda-programming-guide/}

The catch is alignment. DeepGEMM's API checks TMA-aligned strides for input and scale tensors. A caller who transposes a tensor with a view may produce mathematically correct values but a stride pattern that can't satisfy the descriptor. The safe path is to create the required major layout explicitly, then validate both shape and stride.

TMEM holds accumulators and block scales

TMEM, or tensor memory, is specialized SM100 storage used by the tcgen05 instruction family. This block-scaled implementation reserves TMEM columns for both FP32 accumulators and scale factors. After TMA stages scales in shared memory, a shared-to-TMEM copy supplies them to the matrix instruction. Treating TMEM as accumulator-only storage misses a real data dependency.^{[3]Reference 3DeepGEMM Source Snapshot, July 15 2026https://github.com/deepseek-ai/DeepGEMM/tree/559d79fb6994a58b8a15b4b93bf13ccc16edf247[12]Reference 12PTX ISA: Fifth-Generation TensorCore Instructionshttps://docs.nvidia.com/cuda/parallel-thread-execution/#tensorcore-5th-generation-instructions}

Keep those lifetimes separate when reading the code. Global memory holds model tensors. Shared memory stages TMA input. TMEM holds the matrix accumulator and the scales needed by this instruction. Each handoff has an ordering requirement.

tcgen05 issues MMA work

The SM100 implementation calls tcgen05 matrix multiply-accumulate instructions through CUTLASS and local PTX helpers. One instruction consumes a selected tile from shared memory and accumulates into TMEM. The epilogue then applies output conversion, activation-specific work, or a store.

The issuing CUDA warps follow instruction-specific participation and synchronization rules. Ordinary registers are thread-local, not a shared storage tier. DeepGEMM coordinates TMEM allocation, TMA completion, scale copies, asynchronous matrix completion, and epilogue reads. Some ordering is explicit and some is supplied by commit instructions, so adding or removing a fence by analogy to ordinary loads is unsafe.

Those five resources explain the JIT. Tile sizes, scale granularity, and architecture flags change generated code, so the library compiles a kernel for the signature it actually received instead of shipping every combination.

The runtime JIT and shape signature

DeepGEMM compiles kernels at runtime instead of requiring a complete CUDA build for every possible shape. The C++ JIT compiler creates source for a selected implementation, builds a compiled GPU binary (CUBIN), and loads it through a runtime handle. Installation can therefore skip a giant matrix of prebuilt kernels.^{[1]Reference 1DeepGEMMhttps://github.com/deepseek-ai/DeepGEMM}

The default compiler is NVCC. NVRTC is available behind DG_JIT_USE_NVRTC=1 and can compile faster, with a possible performance cost on some shapes. The cache key is a digest of kernel name, compiler signature, compiler flags, and generated source. A digest maps that signature to a directory under $HOME/.deep_gemm by default, or under DG_JIT_CACHE_DIR when configured. Compilation writes into a temporary directory, flushes files, and atomically renames the completed directory. That publish step matters when multiple distributed ranks try to compile the same shape on a shared filesystem. It prevents a reader from loading a half-built published directory; it doesn't prevent duplicate concurrent compilation or repair a filesystem that lacks the required rename semantics.^{[3]Reference 3DeepGEMM Source Snapshot, July 15 2026https://github.com/deepseek-ai/DeepGEMM/tree/559d79fb6994a58b8a15b4b93bf13ccc16edf247}

The JIT can specialize block sizes, scale granularity, architecture, and optional features. New generated signatures can therefore produce a compile tail. A new request shape doesn't necessarily mean a new binary: the library also exposes set_ignore_compile_dims to keep selected dimensions dynamic. Warm the signatures actually produced by your workload, not every imaginable (M,N,K) tuple.

What belongs in a signature

The signature should include every value that changes generated code or the memory contract: architecture, M/N/K tile choices, scale granularity, grouped versus dense mode, activation, and feature flags such as Programmatic Dependent Launch (PDL). If a value affects code but is missing from the key, a cached cubin can be reused for an incompatible layout.

DG_JIT_DEBUG, DG_JIT_PRINT_COMPILER_COMMAND, DG_JIT_DUMP_PTX, and DG_JIT_DUMP_SASS expose the compilation path. DG_JIT_PTXAS_CHECK=1 asks the build to reject local-memory usage. These switches are useful during kernel bring-up, but production launchers should control them deliberately because logs and compilation artifacts can be large.

Why is the JIT cache key more than a shape tuple like (M, N, K)?

Answer

Compiler version, flags, architecture, generated code options, scale format, activation, and layout can change the binary even when M, N, and K stay constant. The key must distinguish all code or ABI changes.

Dense GEMM still assumes one M for the whole launch. MoE breaks that assumption because each expert sees a different number of tokens. The grouped family has to carry those uneven M values without paying for one launch per expert.

Grouped GEMM: the MoE shape problem

An MoE layer routes each token to a small set of experts. Each expert sees a different number of tokens, so the straightforward choices are one GEMM per expert or padding every expert to the same M. Uneven routing makes each choice pay a different tax.

DeepGEMM's contiguous grouped API concatenates expert token segments into one M axis. N and K stay fixed because experts share weight dimensions. A compact layout tensor describes each segment, allowing one kernel to schedule multiple expert blocks while reusing one compiled implementation.

The segment boundaries still need alignment. The README requires each expert segment to meet the M block alignment returned by get_mk_alignment_for_contiguous_layout(). Padding belongs in the caller's token packing, and the resulting output must be unpacked to the original token order.

Masked grouped GEMM solves a different problem. During decode, the CPU may not know how many tokens each expert receives, especially when a CUDA graph fixes tensor shapes. A mask marks valid rows, so the GPU launches one shape while skipping invalid portions. That avoids a host synchronization, but it may execute less efficiently than a perfectly packed contiguous layout.

Worked routing example

Return to four of a Flash-shaped rank's 32 local experts. They receive [64, 192, 32, 128] routed rows. With M alignment 128, the contiguous pack reserves [128, 256, 128, 128] rows. Valid rows stay 416; padded capacity becomes 640. The first gated projection has $K=4096$, $N=2\times2048=4096$. The second has $K=2048$, $N=4096$.

Padding isn't free. It increases memory traffic and can lower arithmetic utilization. It can still beat four separate launches if one persistent kernel amortizes launch overhead and keeps tiles in flight. Measure total kernel time and useful-token throughput together. A low microsecond number with high padding can mislead.

Here M counts token-expert assignments, not unique source tokens. A token routed to six experts contributes six rows across the distributed expert workload. With hidden width H and intermediate width I, linear 1 costs $2M H(2I)$ FLOPs and linear 2 costs $2M I H$. Together that's $6MHI$. The next snippet uses that arithmetic to compare useful rows with a full-padded computation model. Actual kernel instruction counts require profiling; some paths skip invalid tiles.

Those 20.94 useful GFLOPs cover only the two expert GEMMs for 416 routed rows in our subset. They exclude dispatch, combine, SwiGLU, and the other local experts. Padding is 35% of reserved rows but 53.8% overhead relative to useful rows. These are different denominators, not conflicting results.

Grouped GEMM handles the expert projections, but it can't fuse the DeepSeek indexer's dequantization, ReLU, head weighting, and range mask. That different dataflow needs its own kernel.

MQA logits for the DeepSeek indexer

Multi-query attention uses many query heads with one shared key/value representation. DeepGEMM's indexer kernel computes a weighted ReLU score for each query token against a selected range of key/value tokens. For query i and key/value token j, the README gives this computation:

$s_{ij} = \sum_h \operatorname{ReLU}(q_{i,h} \cdot (k_j \cdot \operatorname{scale}_j)) \; w_{i,h}$

The shared key vector is dequantized with its scale, each head computes a dot product, ReLU clips negative head scores, and per-head weights combine them into one score. This is the README's FP8 path with no separate query-scale input, not a universal formula for every newer scaled-query variant. It matches the weighted-ReLU indexer in DeepSeek-V3.2.^{[6]Reference 6DeepSeek-V3.2: Pushing the Frontier of Open Large Language Modelshttps://arxiv.org/abs/2512.02556}

Walk one score by hand before looking at the paged path. Query heads are $(1, 0)$ and $(0, 1)$, weights are $0.4$ and $0.6$, the stored key is $(2, -1)$, and its scale is $0.5$. Dequantized $k$ is $(1, -0.5)$. The dots are $1.0$ and $-0.5$. ReLU keeps $1.0$ and zeros the negative head, so the logit is $0.4$.

The hand calculation explains what the specialized path adds. DeepGEMM offers non-paged scoring for prefill and paged scoring for decode. Paged mode reads key/value blocks through a block table, which matches a serving runtime's fragmented cache. The indexer can therefore sit beside, not inside, the main attention kernel: it selects useful context positions before the model performs a larger attention operation. FlashInfer already taught the paged-KV layout. This kernel consumes a similar table and writes a logit matrix instead of an attention output.

The bounds are part of correctness. A dense result can contain unwritten positions outside each query's valid key range. clean_logits=True fills those positions with negative infinity; without cleanup, the caller must exclude them explicitly. At this revision, the compressed non-paged mode requires clean_logits=False, so its bounds and output-index mapping remain the caller's responsibility. A zero-filled invalid position isn't safe either: head weights can make valid scores negative, and top-k could then select the invalid zero.^{[3]Reference 3DeepGEMM Source Snapshot, July 15 2026https://github.com/deepseek-ai/DeepGEMM/tree/559d79fb6994a58b8a15b4b93bf13ccc16edf247}

Check that failure without a GPU. Here only positions 1 and 2 are valid; position 0 must never win retrieval.

For an empty valid range, return no selected keys instead of applying top-k to a row of negative infinities. Neither this bounds check nor the earlier two-dimensional dot example tests an actual supported CUDA shape.

The MQA path makes the specialization boundary concrete. Generic GEMM can express a dot product, but it doesn't automatically fuse dequantization, ReLU, head weighting, range masking, and cache paging. The specialized path saves intermediate writes and keeps the score reduction close to the input tiles. The same question returns in Mega MoE: which data can move while the tensor cores are busy?

Mega MoE: overlap communication and compute

Mega MoE is DeepGEMM's most ambitious kernel. It fuses expert-parallel dispatch, linear 1, SwiGLU, linear 2, and expert-parallel combine into one SM100 kernel. It uses PyTorch symmetric memory so ranks can address each other's buffers and overlap NVLink traffic with tensor-core work.^{[4]Reference 4Mega MoE: Fusing Expert Parallel Communication with Computationhttps://github.com/deepseek-ai/DeepGEMM/pull/304}

The computation still follows ordinary MoE semantics:

1. Router selects top-k experts and weights for each token.
2. Expert parallelism dispatches token rows to owning ranks.
3. Linear 1 produces gate and up activations.
4. SwiGLU applies the gated activation.
5. Linear 2 projects back to the hidden size.
6. Combine returns weighted expert outputs to the source rank.

Schedule and storage provide the core mechanism. A symmetric ring buffer holds token data, scale data, routing indices, weights, and intermediate expert activations. Persistent schedulers claim blocks, move remote data into reusable slots, run L1 and L2 work, and release slots only after all consumers finish.

In this simplified schedule, each chunk needs one dispatch slot, two compute slots, and one combine slot. Two serial chunks take eight slots. With independent communication and compute resources, dispatch B can run during compute A, and combine A during compute B. Completion takes six slots. If both resources contend for the same bottleneck, or no second chunk is ready, that saving needn't occur.

This overlap is specialized expert-parallel traffic over NVLink and symmetric memory, not a general collective. The next chapter's NCCL library supplies the broader all-reduce, all-gather, and all-to-all contract that training and serving stacks compose on top of.

Persistent scheduling and ring reuse

Persistent scheduling means a resident set of cooperative thread arrays repeatedly claims work instead of launching a new grid for every expert block. Communication and compute phases can interleave, but more state lives in registers and shared memory. The source comments call out warmup waves, ring-block reuse, and ordering rules that prevent an L2 task from consuming an unfinished L1 block.

The ring buffer is bounded. A new dispatch can't overwrite a slot until the previous consumer signals that all N blocks are complete. That rule is a correctness invariant, not a performance hint. A missing release or early reuse can produce data races that look like occasional model-quality regressions rather than a clean CUDA error.

Symmetric memory contract

The Python API allocates a SymmBuffer for a distributed process group. For one rank it can use a regular CUDA allocation; for multiple ranks it uses PyTorch's symmetric-memory rendezvous. The buffer exposes typed views for input tokens, scale factors, top-k indices, top-k weights, and intermediate activations.

The documented low-bit Mega MoE path uses PyTorch 2.9 or newer, FP8 activations with FP4 weights, SwiGLU, SM100 hardware, and aligned dimensions.^{[4]Reference 4Mega MoE: Fusing Expert Parallel Communication with Computationhttps://github.com/deepseek-ai/DeepGEMM/pull/304} The pinned source also exposes a BF16 variant. Expert parallelism across ranks needs a multi-process launch; a single-process smoke test doesn't exercise remote buffer synchronization or NVLink overlap.

The pinned Python call accepts recipe=(1, 1, 32), and the C++ API asserts that exact recipe for the low-bit Mega MoE path. Callers pack activations and weights with K-group 32 UE8M0 scales, transform the weights, copy routing inputs into the symmetric buffer, then invoke fp8_fp4_mega_moe. Passing another group width isn't a tuning option.^{[3]Reference 3DeepGEMM Source Snapshot, July 15 2026https://github.com/deepseek-ai/DeepGEMM/tree/559d79fb6994a58b8a15b4b93bf13ccc16edf247}

This is an unexecuted API sketch, not a copy-runnable lab. The caller must initialize the process group, cast activations, pack scales, copy routing inputs into the symmetric buffer before each call, and size output tensors. The source's weight transform also interleaves gate/up rows. A buffer with the right byte count but the wrong row ordering isn't a valid input.

Read the Mega MoE benchmark as a bounded experiment

DeepGEMM's benchmark pull request answers a narrow question: how does Mega MoE compare with a legacy dispatch, GEMM, activation, and combine path under eight-way expert parallelism? Values are averaged across eight ranks and two named DeepSeek-V4 configurations. They are evidence for that harness, not a universal speedup.^{[2]Reference 2Mega MoE Benchmark Resultshttps://github.com/deepseek-ai/DeepGEMM/pull/316}

DeepSeek-V4-Flash shape

The reported Flash shape has 256 experts, top-k 6, hidden size 4096, and intermediate size 2048. Batch size means source tokens per rank, before each token's six expert assignments. These are historical results from PR #316, merged April 24, 2026, not measurements made for this lesson.^{[2]Reference 2Mega MoE Benchmark Resultshttps://github.com/deepseek-ai/DeepGEMM/pull/316}

DeepSeek-V4-Pro shape

The Pro benchmark shape uses 384 experts, top-k 6, hidden size 7168, and intermediate size 3072. Larger dimensions change both arithmetic work and communication volume.^{[2]Reference 2Mega MoE Benchmark Resultshttps://github.com/deepseek-ai/DeepGEMM/pull/316}

At 512 tokens per rank under EP8, the node holds $512\times8=4{,}096$ source tokens, or 24,576 token-expert assignments for top-k six. Balanced routing would average 3,072 received rows per rank; individual ranks can differ. The source-token denominator isn't the local expert-row denominator.^{[2]Reference 2Mega MoE Benchmark Resultshttps://github.com/deepseek-ai/DeepGEMM/pull/316}

The compute and bandwidth columns are work estimates divided by elapsed time, not hardware-counter measurements of traffic. For example, $6\times3072\times4096\times2048$ FLOPs divided by 146.5 microseconds gives about 1,055 TFLOP/s, close to the rounded 1,056 reported. The pinned harness estimates global-memory bytes and NVLink payload bytes; cache behavior, protocol traffic, and scale movement can make physical traffic differ.^{[3]Reference 3DeepGEMM Source Snapshot, July 15 2026https://github.com/deepseek-ai/DeepGEMM/tree/559d79fb6994a58b8a15b4b93bf13ccc16edf247}

The PR gives shapes and EP8 rank averaging but doesn't provide a complete GPU SKU, driver/toolkit, clock, topology, warmup, and correctness receipt for those historical rows. Don't fill that gap by guessing a Blackwell model. The pinned test's legacy comparison can also be skipped when baseline dependencies fail to import. A printed timing alone isn't evidence that correctness checks ran.

What does the 1.96x Mega MoE result prove, and what does it not prove?

Answer

For one Flash-shaped token per rank, the reported baseline-to-Mega time ratio was 1.96. Mega time was about 51% of baseline, a 49% reduction. It doesn't establish speedup on another topology or in a complete serving workload, and the historical environment receipt is incomplete.

Turn timing into a diagnosis

Start with Nsight Systems when the question is where time goes across the process. Its CUDA API and GPU workload traces show JIT compilation, host launches, memory transfers, kernels, and streams, so they can reveal preparation or serialization that a kernel stopwatch hides. Move to Nsight Compute for one selected kernel: Launch Statistics describes the grid, SpeedOfLight compares achieved compute and memory throughput with hardware ceilings, Memory Workload Analysis identifies the busy memory unit, and Occupancy or scheduler sections show whether resource pressure leaves warps waiting.^{[13]Reference 13Nsight Systems User Guidehttps://docs.nvidia.com/nsight-systems/UserGuide/[14]Reference 14Nsight Compute Profiling Guidehttps://docs.nvidia.com/nsight-compute/ProfilingGuide/}

Read those counters as clues, not verdicts. High occupancy doesn't guarantee high performance, and a low tensor-core percentage may be expected for a tiny M. A roofline point below the compute ceiling can mean low arithmetic intensity, poor reuse, or a launch that never feeds the units. Enable DG_JIT_WITH_LINEINFO=1 when source correlation matters, and use a narrow, repeatable profile range. Nsight Compute may replay a kernel to collect metrics, so keep its report for diagnosis and keep synchronized, unprofiled timings for the benchmark table.^{[9]Reference 9CUDA C++ Best Practices Guidehttps://docs.nvidia.com/cuda/cuda-c-best-practices-guide/[14]Reference 14Nsight Compute Profiling Guidehttps://docs.nvidia.com/nsight-compute/ProfilingGuide/}

For Mega MoE, the profile needs one more join: route histogram, padded rows, rank topology, dispatch and combine readiness, and complete-layer time. If compute looks busy while the layer gets slower, inspect whether communication actually overlaps or whether casts, packing, or a barrier sit on the critical path. That evidence tells you which contract to change next.

Strengths and weaknesses

The same choices that make DeepGEMM fast also define where it fits. Shape-specialized kernels reward a team that controls its GPU fleet and tensor layouts; they ask more from callers that need portability or arbitrary shapes.

Strengths

• Shape-aware performance. The JIT can specialize tile sizes, layouts, scale granularity, and architecture instead of forcing one generic iterator.
• Small public surface. A few kernel APIs make it practical to trace from Python call to C++ checks to generated CUDA.
• Model-specific fusion. MQA scoring and Mega MoE remove intermediate tensors that a generic graph would materialize.
• Distributed awareness. Symmetric memory, ring buffers, and persistent schedulers address expert-parallel traffic directly.
• Useful correctness references. Tests compare outputs against PyTorch or legacy paths across dense, grouped, paged, and fused operations.

Weaknesses

• Narrow hardware support. Current public paths target SM90 or SM100, and Mega MoE is SM100-first. A machine with another architecture needs a different kernel or fallback.
• Strict layout contracts. Scale dtype, stride, major order, and alignment are part of correctness. The caller must prepare them.
• Compilation latency. New signatures compile on first use. A wide shape distribution can create cold-start tail latency.
• Operational complexity. Mega MoE needs multi-process setup, symmetric memory, correct barriers, and a compatible PyTorch release.
• Limited portability. CUDA, CUTLASS, {fmt}, NVCC or NVRTC, and PyTorch versions all affect the supported path.
• Benchmark locality. PR tables compare one harness. They don't replace a workload-level test with real request lengths and rank placement.

DeepGEMM chooses explicit contracts and measured specialization over broad portability. That fits a model team that controls hardware, CUDA images, and tensor layouts, but fits less well as a drop-in library for arbitrary user tensors. Bring-up should make that boundary visible before timing begins.

Bring up the contract before timing

The README lists Python 3.8 or newer, C++20-capable compilers, PyTorch 2.1 or newer for the base library, CUTLASS 4.0 or newer, and CUDA 12.3 or newer for SM90. CUDA 12.9 or newer is recommended for best performance and required by current SM100 support.^{[1]Reference 1DeepGEMMhttps://github.com/deepseek-ai/DeepGEMM} Mega MoE adds a PyTorch 2.9 requirement. Record those versions with every benchmark receipt.

Before running a benchmark, record the environment and the shape it will exercise:

1. GPU model and compute capability.
2. CUDA toolkit and driver versions.
3. PyTorch and CUTLASS versions.
4. Kernel mode, scale format, and layout.
5. M, N, K or expert dimensions and token distribution.
6. Number of ranks, rank-to-GPU topology, and interconnect.
7. Warmup count, JIT cache state, synchronization method, and measurement window.
8. Correctness tolerance and reference implementation.

For a first local check, start with a dense BF16 or FP8 GEMM shape from the repository tests. Enable DG_JIT_DEBUG=1 only while tracing compilation. Use DG_JIT_CACHE_DIR on fast local storage, not a path with unreliable distributed-file semantics. For multi-rank Mega MoE, run the repository's distributed test harness and compare output to the non-fused baseline before reading timing output.

Shape constraints to expose

Make constraints part of the model runtime's validation layer. Check that hidden and intermediate dimensions meet the required multiples, expert count is divisible by rank count, token counts fit the symmetric buffer, and scale tensors have the required packed width. Reject a request with a clear error before it reaches a CUDA graph or a collective barrier.

Diagnose failures from their first symptom

Wrong output, no crash

Numerically wrong output with no crash usually points to scale packing, tensor major order, or output accumulation dtype. Compare a small shape against the PyTorch reference after dequantization. If the error appears only on SM100, inspect UE8M0 packing and TMA strides. If it appears only in grouped mode, print segment offsets and padded M counts.

First request is slow

If the first request is slow, check whether the JIT cache missed. Log kernel name, signature, compile duration, and cache directory. Warm expected signatures during startup or provide an explicit fallback. Report cold and warm latency separately, while keeping real cold requests in the product latency distribution.

Mega MoE hangs

If Mega MoE hangs, treat it as a synchronization or collective contract failure until proven otherwise. Confirm every rank entered the same process-group call, buffer sizes match, top-k indices are valid, and symmetric memory rendezvous completed. Enable communication-kernel debug only in a controlled reproduction because it zeros the buffer and requires callers to recopy inputs before each invocation.

Speedup disappears

If the speedup disappears, check whether casts, transposes, token packing, or network transport sit outside the timed kernel. Then compare useful-token throughput against padded-token throughput. A fused kernel can be faster while the full graph regresses if preparation dominates.

Thermal or topology drift

If results drift thermally or by topology, repeat after warmup and pin rank placement. A benchmark that runs eight ranks on a different NVLink topology from production can report a misleading interconnect rate. Record clocks and power policy when comparing machines.

A code-reading route through the repository

Use this order when studying the clone:

1. README.md defines public APIs, supported architectures, scale formats, and examples.
2. deep_gemm/__init__.py exposes Python functions and utility transforms.
3. csrc/apis/gemm.hpp checks shapes and dispatches dense and grouped GEMMs.
4. csrc/jit/compiler.hpp shows signatures, compiler flags, cache publication, and NVCC or NVRTC paths.
5. csrc/jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp is the host wrapper that JIT-compiles one SM100 tile path.
6. deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh traces TMA, TMEM, and tcgen05 stages for dense GEMM.
7. deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh explains expert task order, warmup, and ring capacity.
8. deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh connects remote loads, grouped GEMMs, barriers, and stores.
9. tests/test_attention.py and tests/test_mega_moe.py show reference calculations and shape sweeps.

Read one data structure at a time. First find the shape and stride checks. Then find where a scale tensor is transformed. Finally trace who writes each buffer and who releases its barrier. This route is more productive than starting with a 1,000-line CUDA kernel and guessing the invariants.

What strong answers show

• Derive NT output dimensions and count multiply-adds without confusing a whole GEMM with a block tile.
• Place each K-group's scale product inside the reduction and identify the required native scale layout.
• Trace A, B, scales, and accumulators through shared memory and TMEM, including completion dependencies.
• Distinguish source tokens, routed rows, padded capacity, and local versus node-wide denominators.
• Explain overlap across ready chunks while preserving each chunk's dispatch, compute, and combine order.
• Separate CPU reference agreement, GPU correctness, historical throughput estimates, and serving latency evidence.

Follow-up questions

Why might a faster grouped GEMM make the layer slower?

Packing, scale conversion, dispatch, and output restoration may cost more than the saved kernel time. Compare complete-layer latency on the same route histogram and include preparation inside the measured boundary. Useful-row throughput and padded capacity help explain the result; neither substitutes for elapsed time.

What if the benchmark prints a time but no correctness result?

Check whether the legacy dependencies imported and the correctness loop ran. At the pinned revision, the harness can continue when its baseline isn't available. Mark that run as timing-only, restore the reference path, and test routed row counts and numerical outputs before using the result to choose a production kernel.

Architectural summary

• DeepGEMM is a kernel library, not a serving runtime.
• Data layout, tensor-core instruction, schedule, and model shape are selected together. Change one without the others and the kernel can miss its path.
• SM90 and SM100 use different scale formats and implementation paths. Architecture checks are part of correctness.
• Runtime JIT makes shape specialization practical, but cold-start compilation and cache hygiene become operational concerns.
• Grouped GEMM handles uneven MoE token counts with contiguous or masked layouts, each with a different host-knowledge trade-off.
• MQA indexer scoring fuses dequantization, head weighting, ReLU, and range or page selection for DeepSeek-V3.2.
• Mega MoE fuses expert dispatch, FP8 × FP4 projections, SwiGLU, and combine, then overlaps NVLink traffic with tensor-core work on SM100.
• Benchmark numbers need shape, hardware, rank topology, warmup, and baseline details before they can guide a production choice.