Optimizing Large MoE Inference on NVIDIA Blackwell: NVFP4, ADP, and DualPipe Strat... Julien Demouth

This talk explores the optimization of Mixture of Experts (MOE) inference on NVIDIA's Blackwell architecture. The speaker discusses the importance of performance in MOE, specifically within transformer models, where data-dependent routing selects specific experts for token processing. Techniques such as reducing numerical precision with FP8 and introducing 4-bit data types are highlighted as methods to enhance speed and efficiency in computations. The presentation covers hardware/software co-design with a focus on balancing GPU utilization and optimizing communication for better expert handling. Additionally, the speaker discusses methods to implement attention data parallelism to minimize processing discrepancies across GPUs, crucial for the DeepSpeech model used in their benchmarks.

So, let me talk about about how we did optimize MOE inference on Blackwell. A couple of disclaimers. I was actually not supposed to be giving this talk. So, U Touch and Meg who were supposed to to be here to to talk about the work didn't get through the visa process. So, I'm sorry about that. So, I'll be doing the their talk with their slides on a computer

that is not my computer. >> [laughter] >> So, that sounds like a cool thing. So, forward. As you all know, the the very recent models, they all use MOE. And so, for us, optimizing MOE is very important. And we care a lot about the performance of MOE. And this is what this presentation is about. I don't know if you've slept for the past 2 years, let me

very quickly remind you of what MOE is. And especially MOE in the transformer architecture nowadays is this block that replaces the dense feed forward network. In an MOE, you typically different experts. Hence, the name mixture of experts. And given a token, you have an operation the router that will select the experts that will be best suited to deal with the specific token. So, you have this first

operation that is data dependent. For a given token, selects which experts to send the token to. And once the tokens are selected, this experts are selected, the token is sent to the experts. the token is computed by the different experts. And after the experts have done their job, their result is consolidated in a single set of values per token. So, that's the general idea MOE. And this

creates some challenges I want to talk about today. So, when we approach this problem, as Nvidia, we we do the usual stuff. It's not necessarily specific to MOEs. It's in general, what we do is that we have our hardware and we try to develop optimized kernels for this. We try to fuse as many things as possible. For the different kernels, we make sure we achieve the best

possible utilization of the hardware. And we work hard on this. But once I've said that, I want to give you more about the specific problem of MOE. So, our first idea play which is generic again, not necessarily apply only to MOE. It's a general technique we use which is to try to reduce the numerical precision as as much as possible. Because going to lower precision data types

allows us to speed things up. A smaller data type represents more compact data type and allows us to reduce the amount of data that we have to transfer over the memory system. So, we can better use the memory system or actually reduce the amount of data transfer that accelerates things. But also smaller data types are more amenable to hardware acceleration for computations which allows us to increase

the flops. One thing that I want to mention that is not necessarily super visible in the slide is that when we talk about FP8 in general, there are actually many different flavors of FP8 that we could be using for training but also for inference and especially inference in this case. When we talk about FP8, we talk about floating point 8-bit floating points. And we when we represent

8-bit floating points, there are different ways we can represent these floating points. And define scaling factors to quantize and dequantize these FP8 numbers. One way to quantize and dequantize FP8 numbers is for every tensor specify a single scaling factor that's going to be used to quantize the entire tensor and dequantize this same tensor. And this is the easiest way to optimize in hardware. When you want to

do inference, if you can maintain your accuracy with this format, this is the easiest way for us to optimize the code. Especially if the scaling factor is computed before inference and is static across the entire inference. Meaning that you have your model, you pre-calibrate your scaling factor before doing the inference. And then you reuse the scaling factor for the entire runtime. This is one one technique that

is the easiest to optimize. Then you can go and increase the granularity of the scaling factor. You can use different scaling factors per token. You can also imagine different strategies to compute the using dynamic calculation. So, every time you run inference, you're going to use compute new scaling factors based on the tokens you're seeing. And this might bring you more accuracy but it creates some problems for

us to optimize this. Typically, if you're doing dynamic scaling factor calculation for entire tensors, if you're distributing your calculation across multiple GPUs, then suddenly you have to synchronize these GPUs to compute the scaling factors. And that creates some potential challenges. That said, it's also possible to reduce the granularity of the scaling factor and use block formats. And one block format that you you're probably familiar with is

the Deep Seek FP8 block format where the computed scaling factor every 128 elements. Another block format is MXFP8 where you compute a scaling factor for every 32 The advantage of MXFP8 over the Deep Seek scaling format is that MXFP8 on Blackwell is hardware accelerated. So, that's one way you can accelerate the computation. And what do I mean by hardware accelerated? And what is the comparison with the

DSFP8? When you do matrix matrix multiplications, which is the key operation in MOE, if you have a a scaling factor every 32 element in the dimension that you're accumulating over. So, for a dot product, the accumulation dimension, your scaling fact your tensor core has to be aware of this operation or this scaling factor to operate efficiently. Especially on the Blackwell architecture where we have special memory to

store the accumulators. Because with a format like the Deep Seek which is not hardware accelerated, every time we have the scaling factor, we have to move the data between the special memory on the tensor core to register file and do the scaling and and send the data back to the the tensor memory which is potentially challenging and super optimal. I don't have a lot of time to

discuss about this now. But if you want to talk more about that, I'm happy to to discuss that after this presentation. going to FP8 is a way to accelerate. But obviously on Blackwell, we were were pushing this further. And we have 4-bit data types, floating point data types. Two 4-bit floating point data types, NVFP4 and MXFP4. These are very similar formats. You have scaling factor for every

block. With NVFP4, it's every 16 elements. For MXFP4, it's every 32 elements. And the scaling factors are slightly different. the the scaling factors are powers of twos, 8-bit powers of twos. NVFP4, the scaling factors are FP8 numbers. And we have developed a strategy to to do a two-stage quantization for NVFP4. I'm not really going to cover this in details right now. But again, let's talk more about

this if you're interested. using this quantization and the tools sorry, using this data type and the quantization tools that are available, you can quantize your And yeah, now the microphone is bad, right? Um Is there a way to switch the microphone Okay, I'm going to yell. Do you hear me? So, uh with this data type in mind, we have developed tools to quantize the model uh and

I invite you to check uh the model of tool from NVIDIA or uh we also uh there are also other solutions like LM LM quantization or compressor, LM compressor developed by uh our friends at Red Hat. Uh different solutions to to quantize the model. Uh the key thing is that this quantization for uh This example is for DeepSpeech R1. Uh for a bunch of uh key benchmarks,

we don't really see a degradation in terms of accuracy going from uh the reference FP8 implementation to FP4 and with this FP4 quantization we're capable of uh optimizing the performance. >> [cough and clears throat] >> Um that's about quantization, but optimizing uh model serving on on Blackwell is not just that. Uh when we created Blackwell, we did uh hardware software co-design and when we created the Blackwell

architecture, we were seeing MoEs as potential big candidates were being important in future models. So, we've created this uh architecture on the right that everyone knows as NVL 72, a combination of 72 GPUs uh tightly connected through an NVLink high bandwidth connection which delivers 1.8 TB per second between pairs of GPUs, bidirectional uh 1.8 uh TB. Uh what? 1.8 TB, sorry. Uh and obviously taking advantage of

this uh is also important when you optimize. Furthermore, the NVL 72 system is not GPUs with NVLink between the GPUs, but also two Grace CPUs and two Grace CPUs are Yeah, it's important tonight. Paris is playing, so I still need voice to >> Okay, uh so sorry about this one. Um so, we have high connection between the CPU and the GPU and we're taking advantage of this.

I'll going to cover this uh in the slide after this one, but uh to also optimize uh the the delivery of MoEs, uh we we we've done some optimization on the um implementations of communication primitives uh because as I said at the beginning, MoE you select the experts, you send the token to the different experts having an optimized uh implementation is also important to send the the

experts uh uh or the tokens as quickly as possible to the different For the interest of time, uh I'm I'm just going to keep it there, but let's discuss after the the meeting if you're interested. One thing that I want to talk about is thing about the high connection the high bandwidth connection between CPU and GPU uh because when you're doing MoE and you distribute your different

expert to different GPUs, which is called expert parallelism, then when you're sending tokens, it's not immediately obvious or it's pretty sure that the distribution will be unbalanced. You will have some experts that will more that will be way more activated than others. And so, if you have experts on some GPUs or some nodes that are more activated than others, you'll have this imbalance that may penalize your

performance. So, what you want is to have a as balanced distribution as possible. The distribution depends on the tokens and potentially on the training and and and so, it's not easy to to have this or I mean, it's not possible to balance this uh without changing your models. So, so what you do and what we did is to optimize the distribution of the experts. If you have,

for example, two experts on the same GPU that are highly utilized and on other GPUs experts that are less utilized, you'd better swap some of the experts, right? Rebalance because if you put one if you have two experts that are highly utilized, put each one of them on different GPUs and that's going to help you. Another thing is sometimes you have experts that are very, very much

utilized and even if you put them on their GPU, this GPU is going to be saturated and others may be underutilized. So, why not replicate this expert on multiple GPUs so that you you better balance your your traffic. And the connection I'm I want to make with what I was saying before is that we have high bandwidth connection between CPU and GPU on NVL 72. So, what

we do is that we keep the different experts in CPU memory and over time we analyze the traffic and the distribution of traffic to the different expert and send copies of those experts to the correct GPUs and rebalance the work. Stay with me. We keep the experts on the GPU, but having a copy of those experts on on the CPU and the high bandwidth connection allows us

to pretty frequently update uh the GPUs on uh the experts on the GPU and increase our utilization of GPUs. And this is the idea that uh is shown on this graph, which is pretty tough to read. Um So, uh the the the general idea is um the on this on this diagram there, uh you you have to compare um the two bars uh next to each other

where the the leftmost is the the the performance without um the balancing and the the the right one is with balancing. And lower is better. So, this shows with 16 experts parallelism of 16, this shows with expert parallelism of 32 and this shows with expert parallelism of uh 64. And as you see, the performance in increases because the time goes down. Um and with balancing, it's critical

to I mean, balancing plays a critical role in in helping performance uh for for for this uh these tokens. The the very last technique I want to talk about and and I'm going to be very brief because I'm pretty sure you're all looking forward to having lunch. Um is that we implemented uh attention data parallelism in in our implementation. when you want to increase the throughput, especially

with a model like DeepSpeech R1, um it doesn't make too much sense to to split the attention across multiple GPUs and do tensor parallelism across multiple GPUs. What makes more sense is actually to dispatch requests to the and have uh each GPU have its own KV cache representing the entire MLA attention. But the problem is that if you dispatch the different requests to the different ranks, you

may end [clears throat] up having some ranks doing prefill and some ranks doing decode. And decode obviously goes a lot faster. And so, uh you have some imbalance between some ranks doing the prefill step, which is costly in comparison with ranks doing the decode. And so, what we we've developed is a technique to uh try to group requests needing to do decode together uh so that when

we trigger a sort of decode phase, uh sorry, prefill, uh when we trigger a prefill phase, we trig- we trigger the prefill phase on all the ranks. So that we have no imbalancing uh between the ranks. All the ranks do the the prefill phase. This is equally costly on all the ranks. And so, you let's say uh 30 ranks doing decode and one rank doing prefill that

is slowing everyone down. And that works a lot better this way. This is in the context of aggregated serving, not disaggregated serving, but it helps a lot. this diagram shows that especially on the right, uh the or actually on those, this is the throughput per GPU and so higher is better. the red curve is when there's no So, the other curves that start doing balancing show that

the performance increases with with this balancing in part. And this uh this diagram there shows the time to first token. So, obviously balancing meaning that you you slow down or you wait for having enough prefill request to start doing prefill has an impact on the time to first token, but this diagram shows that it's rather limited. Um it's in the order of a microsecond or something like

that, which is not too too too bad in this context. I'm very sorry for uh for this rushed presentation, but if you have any question, I'm happy to answer the questions um right after this session and I'll be around, right? Thanks a lot. >> [applause]