A DeepSeek-V3 transformer block
Before we can calculate any rooflines or look at profiles, we need a basic working understanding of how the DeepSeek-V3 architecture is put together. The classic path here would be to read the DeepSeek-V3 Technical Report and then look at some of the implementations of DeepSeek-V3: either the official inference reference or a trainable implementation in torchtitan or Megatron (config).
However, there are some aspects where I think the pre-existing literature could be clearer. For example, the official DeepSeek architecture diagram looks like this:
I personally find the MLA diagram quite difficult to parse! The difficulty is that the diagram focuses on the latent spaces you are projecting between (it doesn't even have any explicit indication where a matrix multiply is happening), whereas as an infra person I'm more interested in knowing what kernels I am going to be running and what the more expensive operations are (the ones we should be optimizing, c.f., Amdahl's law.) You could get this information by pointing a coding agent at a training codebase. But I think a schematic diagram is still quite helpful, especially for visualizing data dependencies in the transformer block. Fable helped me put together this diagram which highlights the most important operations in the DSv3 transformer block:
There are a few things I'm hoping to convey in this diagram:
- The bold boxes are entirely matrix multiplies (or attention). In roofline analysis, we often approximate FLOPs by only considering matmuls, as the quadratic cost of doing matmuls means they tend to dominate the overall FLOPs of a model. It's OK to squint your eyes and ignore the rest of the boxes!
- Unlike what you would see in a mathematical description of MLA, we have condensed the number of matmul blocks by applying some conventional horizontal fusions, e.g., fusing the K and V matmuls together. Any serious training implementation will have applied these fusions, as they are easy to implement and improve your matmul efficiency. It's good to know to expect five GEMM kernels in the MLA region!
- If you click the "sizes" button, you can toggle between seeing the internal structure versus the topline number that has it multiplied out. For example, the parenthetical numbers in grey all denote the parameter count associated with an operation. On the ffn gate/up, you can either see that it is 29M × 256 (the size of a single expert multiplied by the number of routed experts) or 7.5B (the total routed experts parameter count of a single layer). In the topline view you can see that FFN has an order of magnitude more parameters than any of the other operations in the transformer block.
- The intermediate activations are all named, and any that are saved for backwards are annotated with an arrow (e.g., ↓ or ↑) pointing at the operation that needs them saved for backwards. On occasion multiple operations save the same tensor for backwards (drawn as ⇅). Note that this diagram overstates how many activations you'd actually typically save for backwards--more on this in a future post.
That's it for this post. This diagram isn't earth shaking but we are going to return to it again and again as we start doing analyses. For example, we are going to come straight back to it when we do a roofline analysis of memory usage in the next post.