Deep Dive - DeepGEMM

When a language model serves one token, it may execute thousands of matrix multiply instructions before a user sees the next word. A good kernel can save a few nanoseconds per tile, then repeat that saving across every layer, request, and GPU. DeepGEMM is a useful project to study because it exposes that work instead of hiding it behind a general-purpose framework.

The repository is a small CUDA library for tensor-core operations used by modern large language models. It includes dense and grouped GEMMs, low-bit formats, multi-query attention (MQA) scoring for the DeepSeek-V3.2 lightning indexer, and Mega MoE, which fuses expert-parallel communication with expert computation.^{[1]Reference 1DeepGEMMhttps://github.com/deepseek-ai/DeepGEMM} Its runtime just-in-time (JIT) compilation specializes kernels for a concrete shape. This lesson reads the code as an engineering argument: match data layout and schedule to a fixed hardware contract, then specialize only where measured shapes justify it.

The project isn't a model server. It doesn't own request admission, token budgets, or a public HTTP API. It supplies fast device kernels and a Python interface that other model and serving systems can call. That boundary is its first lesson: a kernel library wins at one layer, while the surrounding runtime decides whether that win improves end-to-end latency.

Project identity and system boundary

DeepGEMM 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 of contributors because communication kernels, layouts, scheduling, and benchmarking all cross subsystem boundaries.^{[1]Reference 1DeepGEMMhttps://github.com/deepseek-ai/DeepGEMM[2]Reference 2Mega MoE: Fusing Expert Parallel Communication with Computationhttps://github.com/deepseek-ai/DeepGEMM/pull/304}

The project sits beside, not above, libraries such as CUTLASS and CuTe. DeepGEMM borrows ideas from those projects while keeping a smaller set of public kernel functions. That smaller surface makes it easier to inspect the generated path for one shape, but it also means callers must prepare layouts and scaling factors correctly. There is no general-purpose fallback that can make every malformed tensor fast.

Applications

There are four main application paths:

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 (MoE). Grouped GEMMs process the variable number of tokens sent to each expert. The contiguous and masked layouts target prefill and decode differences.
3. DeepSeek-V3.2 indexing. 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.^{[4]Reference 4DeepSeek-V3.2: Pushing the Frontier of Open Large Language Modelshttps://arxiv.org/abs/2512.02556}
4. Expert-parallel inference. Mega MoE combines dispatch, two expert projections, SwiGLU (a gated feed-forward activation), and combine so communication can overlap tensor-core work.^{[2]Reference 2Mega MoE: Fusing Expert Parallel Communication with Computationhttps://github.com/deepseek-ai/DeepGEMM/pull/304}

DeepSeek-V3 is a clear motivation for the MoE paths. Its technical report describes 671B total parameters with roughly 37B active per token, 256 routed experts, and 8 routed experts selected per token.^{[5]Reference 5DeepSeek-V3 Technical Report.https://arxiv.org/abs/2412.19437} Active compute is smaller than total storage, but every expert still has to live somewhere and routed tokens still have to move between ranks. That gap is exactly where grouped kernels and Mega MoE spend their complexity budget.

Research foundations and source trail

DeepGEMM doesn't ship one standalone research paper that explains every kernel. Its source tree and release pull requests are the primary engineering record. The papers explain why the kernels matter: DeepSeek-V3 motivates sparse expert routing and expert-parallel placement, while DeepSeek-V3.2 describes the lightning indexer that needs weighted MQA logits.^{[5]Reference 5DeepSeek-V3 Technical Report.https://arxiv.org/abs/2412.19437[4]Reference 4DeepSeek-V3.2: Pushing the Frontier of Open Large Language Modelshttps://arxiv.org/abs/2512.02556}

Use the sources in four layers:

This distinction keeps the evidence honest. A paper can motivate a workload without proving a kernel speedup. A benchmark can show a speedup without proving model quality. Read each claim against the source that actually owns it.

GEMM first: the contract behind the speed

General Matrix-Matrix Multiplication, or GEMM, computes a matrix product and optionally adds an existing output:

$D = C + A B^T$

For the common non-transposed A, transposed B 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 approximately $2MNK$ floating-point operations because each multiply is paired with an add.

Consider a tile with $M=512$, $N=4096$, and $K=4096$. It carries $8{,}589{,}934{,}592$ multiply-accumulate pairs, or $17{,}179{,}869{,}184$ floating-point operations when multiply and add each count as one. If a kernel finishes in $10$ microseconds, its arithmetic rate is about $1{,}718$ teraFLOPS. That number is a derived rate for that shape, not a promise that every request will run at that rate.

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.

Precision and scaling

DeepGEMM's low-bit paths trade input precision for bandwidth and tensor-core throughput. FP8 is an 8-bit floating-point format used for activations or weights. FP4 stores two 4-bit values per byte and needs scale metadata to recover a useful numeric range. BF16 is a 16-bit format commonly used for accumulators and outputs because its exponent range is close to FP32.

The kernel doesn't cast a tensor and hope for the best. It consumes scale factors in a specific layout. For the current APIs, SM90 expects FP32 (32-bit floating point) scale factors, while SM100 expects packed UE8M0 scale values in an integer tensor. UE8M0 is a power-of-two style scale encoding with eight exponent bits and no mantissa bits. The format is compact and hardware-friendly, but it makes layout and packing part of the API contract.^{[1]Reference 1DeepGEMMhttps://github.com/deepseek-ai/DeepGEMM}

For a per-32 K-group scale, a row of K values is split into groups of 32. Each group gets one scale. The caller must align the scale tensor so the tensor-memory accelerator (TMA) can issue the expected multidimensional copy. A correct numerical tensor with the wrong stride can still fail an assertion or trigger a slow preparation path.

The practical rule is to 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.

Two architectures, two execution contracts

The repository supports NVIDIA SM90 and SM100 architectures. SM90 and SM100 are architecture families, not interchangeable names for a single GPU. A kernel selected for one family may use different scale types, layouts, instructions, and compiler flags on the other.

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 x FP4 path can select NT, TN, NN, or TT variants and uses packed UE8M0 scales.^{[1]Reference 1DeepGEMMhttps://github.com/deepseek-ai/DeepGEMM}

This split is deliberate. A single portable kernel would need to hide the differences behind extra branches or generalized iterators. DeepGEMM instead keeps architecture-specific implementations visible, then uses a thin API layer to check the contract and route to the right implementation.

The biggest mistake at this boundary is treating an architecture mismatch as a tuning issue. If a 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 the major layout before comparing benchmarks.

Inside an SM100 tile

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.

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 tensor-core accumulators

TMEM, or tensor memory, is the SM100 accumulator storage used by the tcgen05 instruction family. It differs from ordinary registers and shared memory in both access pattern and synchronization rules. The kernel allocates TMEM slices, issues matrix multiply-accumulate operations, then synchronizes before reading results back for an epilogue.

This split changes how a programmer thinks about a tile. Global memory holds model tensors. Shared memory stages TMA input. TMEM holds the matrix accumulator. The data path has three separate lifetimes and each has its own barrier or fence.

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 important invariant is ownership. Warp groups that issue tcgen05 work follow different access rules from ordinary CUDA threads reading and writing shared registers. DeepGEMM's implementation places thread synchronization around TMEM allocation, TMA completion, commit, and load operations. Removing one fence may produce a race that appears only under a particular shape or rank count.

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) with the NVIDIA CUDA Compiler (NVCC) or NVIDIA Runtime Compilation (NVRTC), 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 cache key includes the kernel name, compiler signature, compiler flags, and generated code. 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 atomic publish step matters when multiple distributed ranks try to compile the same shape on a shared filesystem.

The JIT can specialize block sizes, scale granularity, architecture, and optional features. Traffic with many new shapes exposes the corresponding weakness: a model serving system with unbounded sequence lengths can produce a long compile tail unless it warms likely signatures or constrains shapes.

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.

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 a naïve implementation launches one GEMM per expert or pads every expert to the same M. Both choices waste work when routing is uneven.

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

The segment boundaries need alignment. The README's API 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. This avoids a host synchronization, but it may execute less efficiently than a perfectly packed contiguous layout.

Worked routing example

Suppose four experts receive token counts [64, 192, 32, 128], and the kernel's M alignment is 128. The contiguous pack reserves aligned capacities such as [128, 256, 128, 128], while separate cumulative valid ends preserve each expert's true row boundary. Exact padding can be larger under the selected heuristic. The weights keep one shared N and K.

The padding isn't free. It increases memory traffic and can lower arithmetic utilization. Yet it can still win over four separate launches because one persistent kernel amortizes launch overhead and keeps tiles in flight. Measure both total kernel time and useful-token throughput; a low microsecond number with high padding can mislead.

MQA logits for the DeepSeek indexer

Multi-Query Attention (MQA) 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 key/value vector is dequantized with its scale, each head computes a dot product, ReLU clips negative scores, and the per-head weights sum the result into a token-to-token logit matrix. Callers can request a full matrix or a compressed range with per-query bounds.

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.^{[4]Reference 4DeepSeek-V3.2: Pushing the Frontier of Open Large Language Modelshttps://arxiv.org/abs/2512.02556}

The MQA path shows why a kernel library needs model-specific features. 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.

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 NVIDIA's NVLink GPU interconnect traffic with tensor-core work.^{[2]Reference 2Mega 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.

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. This makes communication and compute phases interleave, but it also places more state 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 an 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 current Mega MoE path requires PyTorch 2.9 or newer, FP8 x FP4 weights, SwiGLU activation, SM100 hardware, and dimensions that satisfy alignment checks.^{[2]Reference 2Mega MoE: Fusing Expert Parallel Communication with Computationhttps://github.com/deepseek-ai/DeepGEMM/pull/304} The operation also expects a multi-process launch when expert parallelism spans ranks. A single-process smoke test can validate API shape, but it doesn't prove remote buffer synchronization or NVLink overlap.

This snippet is intentionally a partial API sketch, not a copy-runnable lab. The caller must cast activations, pack scales, initialize the process group, copy routing inputs into the symmetric buffer, and size output tensors. Keeping those steps explicit prevents a common mistake: assuming the fused kernel will perform every model-graph conversion for you.

Benchmark claims without hype

DeepGEMM's benchmark pull request reports Mega MoE under eight-way expert parallelism, with values averaged across eight ranks. It compares a fused kernel with a legacy dispatch, GEMM, activation, and combine path. The results are shape-specific and use two named DeepSeek-V4 configurations, so they should be read as evidence for that harness rather than a universal speedup.^{[6]Reference 6Mega 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. In the pull request's labels, batch size means tokens per rank.

DeepSeek-V4-Pro shape

The Pro shape has 384 experts, top-k 6, hidden size 7168, and intermediate size 3072. Larger dimensions change both tensor-core occupancy and communication volume.

The table gives useful questions to ask. Why does one token show tiny interconnect throughput? Because elapsed time is short and the routed payload is small, not because the network is unimportant. Why does speedup vary with token count? Tile occupancy, buffer reuse, and the ratio of communication to compute change with M. A product benchmark should add model version, GPU type, CUDA version, rank topology, warmup count, correctness check, and p50/p99 latency before making a deployment decision.

Strengths and weaknesses

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. 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 choice 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.

Toolchain and bring-up checklist

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.

Before running a benchmark, record:

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 your 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.

Failure modes worth debugging

Wrong output, no crash

Start with scale packing, tensor major order, and 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

Check whether the JIT cache missed. Log kernel name, signature, compile duration, and cache directory. Warm the expected dimensions during process startup or route the first request to a fallback while compilation completes. Don't hide compile time inside p99 inference latency.

Mega MoE hangs

Treat a hang 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

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

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 to SM90 or SM100 implementations.
4. csrc/jit/compiler.hpp shows signatures, compiler flags, cache publication, and NVCC or NVRTC paths.
5. deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh traces TMA, TMEM, and tcgen05 stages for dense GEMM.
6. deep_gemm/include/deep_gemm/scheduler/mega_moe.cuh explains expert task order, warmup, and ring capacity.
7. deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh connects remote loads, grouped GEMMs, barriers, and stores.
8. 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.

Key takeaways

• DeepGEMM is a kernel library, not a serving runtime.
• Its main speed idea is co-design: data layout, tensor-core instruction, schedule, and model shape are selected together.
• 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 x FP4 projections, SwiGLU, and combine, then overlaps communication with tensor-core work.
• Benchmark numbers need shape, hardware, rank topology, warmup, and baseline details before they can guide a production choice.