SGLang: An Efficient Open-Source Framework for Large-Scale LLM Serving - Liangsheng Yin

Key Takeaways

Hi everyone. Uh today I will uh present a talk at an efficient open source framework for large scale serving and I'm leing I come from Shan University as a student and I also a member of the uh organization LMC's. Yeah. So uh yeah so uh let's introduce what is achil. I believe that most of you uh today may heard may hear uh heard of that our open source project. Yeah. uh we are fastest serving engine for and uh we are uh we are currently one of the state-of-the-art uh performance uh open source engine and we are the first open source implementation to nearly match the throughput reported in the official DeepS blog at large scale and meanwhile I want to stress that uh our elegant lightweight and customizable design has attracted a wide adoption uh from academics and big tech tech uh companies u such as XAI, Nvidia, AMD, based on Microsoft and linking and we also uh scale as a high performance solution uh in the RL side. Yeah. So today's talk I will today's talk I will part uh the outlines into four parts. The first part is u milestones and uh overall features overview and the second part is our efficient design of implementation of PDD aggregation. Uh the third part is large scale EP support and the deepseek blog reproduction and the last part is our ecosystem of the a strong community and the future development. Yeah. So uh let's uh let's take in the first part. Yeah. Uh this is our milestones and overall uh key features. uh uh on 2023 uh we released our initi in initial release uh with the structure LM programming and the prefix caching uh and also as well as constraint decoding. In this version we introduce uh the first open source and efficient prefix caching to reuse KV cache and to reduce the redundant commutation and u uh uh last year summer we introduced our 0.2 2 version uh we uh we we achieved the leading performance in among all the uh open source informance with deploy of the llama 3 and and uh with 0.3 release we uh we release our seven time faster deep MLA and uh faster with torch compile and uh uh we also support multi-image and and video uh uh video models and uh the the last year we released our 0.4 release. This is a this release contains zero overhead batch scheduler uh cache aware DP router and X grammar integration and we are um meanwhile we are the first to serve DeepSync W3 and uh just uh oh uh just uh last month uh we first open sourced implementation of large scale expert parism with pre and decode disagregation. uh we achieved like five time faster than Manila DP and we matched the performance uh uh mentioned in deepseek blog and we matched the API cost uh like 0.2 uh dollars uh 1 million token. Yeah. So uh the second part is uh uh is the efficient design and implementation of PD disagregation. Yeah. uh let's talk about the issue with nonPD disagregation non-disaggregation scheduling mode as as as we all know that uh uh without disagregation scheduling uh the decoding batches will always be preempted by uh prefilling uh this will introduce uh extra uh latency to the token generation and um the second problem is uh uh DP attention with non-disagation mode u it brings uh uh the computation and communication imbalance uh uh within the same DP DP attention group uh uh decoding and prefilling batches they may mix uh with each other and uh uh being executed simultaneously. This uh will increase the latency for decoded tokens and cause uh the load imbalance. Uh the third one is encounter with DP and this part I will also introduce later. uh prefill and decode uh they they they usually use a different dispatch mode uh without disagregation. DPB cannot support both in a same communication group. Uh so uh and we uh and we implement the PD disagregation and this is overview of our architecture. Uh uh for example the uh the purple uh ones is the preview instances or clusters and the orange ones is decoder instances. uh we use the load balancer as a request entry point and uh uh this this load balancer is uh for both preview and the decode path. So uh in in our design LB is decoupled from computational logic and only uh responsible for select PD pair and routes to this uh and return request is it's is it's is it's decoupled and uh we also support uh non-blocking and RDMA based KB transfer and uh uh also we we are we offers flexible API integration of the transfer engine uh for example like Nixo and Monake uh Just as uh uh Dr. John just uh mentioned in last talk. Yeah. Uh here is a demonstration of uh the PDD saggression timeline uh time time step. Yeah. Uh first we uh suppose that we have a preview instance and we have a decode instance. Uh the first thing is the load balancer they should to send a request. uh they they they should choose a PD pair and send the request uh to both prefer instance and decoder instance and uh there the first there is a handshake handshake process. The handshake exchange the metadata uh of like KV uh KV index and addresses and uh the sender the KV cache sender and the KV cache receivers they they are initialized during this phase. And the second thing is uh uh when you want to transfer the KV cache to the decoder engine, the decoder they must have enough memory to to to to to pull this KV cache. So the second uh the second event is the pre-allocation. uh uh after the decoder instance they have enough uh uh memory for uh the KV cache they they they pre-allocate the KV cache in the memory pool and then notify the prefill instance that yes uh we have enough memory so you can start the preing you you can start the prefilling phase and and send a KV cache to us and uh then the prefill they begin the prefill forward and uh after the prefilling uh they send a KV cache to the decoder instance Uh the last part is the decode instance to to to to do the decode. After the decode uh the result is returned to the load balancer and the load balancer return the result to to the user site and this is a PD disagregation time stamp is also a demonstration. Yeah. So uh the next part is uh the large scale expert par support and deepse block reproduction. Yeah. uh uh there even there is only like three layers of uh dense layers in deepse models uh we still need to figure out how to efficiently serve these layers with dense ffm uh we choose to use pure pure data parallel instead of tensor parallel or mixed uh of them uh this is for we we can avoid the TP fragmentation uh on large hidden sides uh we we don't need to cut the tensor uh into uh into different TP ranks and and uh the uh the the second advantage is uh we can uh optimize the memory efficiency. Uh for both preill and decode uh low TP degrees they can benefit from uh only DP uh equals one. Uh this can reduce per device memory. And uh the biggest and the third uh uh uh optimization is uh they can bring is to minimize the communication overhead uh by using purely uh uh purely data parallel we can avoid all the communication within the tensor paris group. Uh then for all layer for every layer is only one scatter and one or gather. uh the scatter is before the attention part and the all together is after the the uh the the attention part only twice communication uh and also for thee layers is sparse FFN uh we use the uh uh par strategy in deep block uh is it's a DP attention with sparse FFN with expert parism uh the reason is we can scale the model capacity different experts can be partitioned in different uh devices. So uh this can uh this can support more experts numbers and removing the memory bottleneck. And the the se uh the second uh uh advantage is uh we can uh uh no this is not advant this is the challenges uh brought by uh thee uh strategy uh by the expert par strategy uh uh when you deploy the uh uh layers with uh uh one DP attention uh dispatch uh EP sparse FFN and then combine they based pattern of uh dispatch expert combin uh combination. Uh this brings some uh communication overhead and underutilization of GPU resources. Uh luckily this can be optimized by DPP and the two batch overlap which I will also introduce later. And also this uh this also brings uh the load imbalance problem and this can be uh addressed by the EPLB. Yeah. So uh the first the the first problem is to uh solve the compat compatibility issue with DPP. Uh DPP they hold two dispatch mode for uh uh for for normal mode and the low latency mode. The normal mode uh they are pre uh they are profile friendly but uh they use uh symbolic tensor shape. Uh this this this doesn't support the CUDA graph and the loading mode they use static data uh tensor shape. is decode friendly uh and both of them they can support DP attention and there is uh auto mode to to automatically handle uh both input and output but uh as just just as I mentioned in uh in in in several slides before uh when you use uh preview and decode uh within the same uh communication group uh in DPP uh they are incompatible because they are uh you cannot put two dispatch mode in within the same communication group. So uh this is the reason uh this is solution how we solve the compatibility issue with DPB and this is also uh one of the reasons that we should use uh PD disagregation instead of unified scheduling. uh uh this part I will introduce the implementation details of uh two batch overlap uh for short tbo. Uh this is a example of uh uh improper launch order of TPO. uh you can notice that uh we want to implement TBO with communication and uh computation are execute uh at the same time uh simultaneously and if we dispatch the uh if we launch the dispatch kernel before the MLP kernel uh for example the uh the yellow ones dispatch kernel is launched before the the green batch uh MLP kernel. So uh the dispatch they will bring synchronization which means uh the CPU will be blocked until the GPU receive all the metadata from all other EP ranks then they can allocate the uh right shape tensored uh the the correctly sized tensors. So uh you cannot uh launch uh two kernel uh yeah so the dispatch kernel launch will block the MLP kernel launch which brings uh the waste of uh uh uh computation resources and so we need to solve this uh by by placing the computation kernel always before the communication kernel. So uh there is a slight difference is we put the green MLP before the yellow dispatch kernel. This is a launch order and for the execution side they are executed simultaneously. Yeah. So uh so in each round uh uh follow the pattern of computation and communication. Yeah. This is a launch uh order. then uh GPU can uh remain active uh during the communication. Uh so how to implement this uh in a clean manner? So uh we use the abstraction uh operation list and the y points. This uh enable cooperative scheduling and uh this uh eliminates code duplication and avoid um some possibly uh uh variable post fixes and this can um this can also manage the partial completion or between different uh attention layers uh decoder layers. Yeah. So uh this is our throughput performance and this this this benchmark uh result is also posted in our uh LMC's blog. You can just search at Google LMC.org and you can find the blog. Yeah. Uh this benchmark is uh uh evaluate on preview and decode uh uh independently assuming that unlimited resources for the other side. So uh and we can notice that uh for the preview mode we achieve uh three times uh faster throughput uh than vanilla tensorm of 16 uh TP uh ranks and we for the decode part we uh successfully enlarge the batch size uh from around uh maybe uh 15 to more than more than 200 uh 200. Yeah. And we also improved the throughput of decoder uh by 5.1 times and this can uh roughly match the uh blog uh posted by deepseek uh in this February. Uh also we we also implement expert parism uh load balancer. Uh in real world serving challenges uh uh the imbalance worsen at scale with uh some some node they will uh accept so many uh routed experts and uh some node they are they they are staying idle uh without uh routing expert to them. So uh two strategy to improve balance is uh we can improve the batch size uh and uh imp uh improve the uh clusters. Uh we can improve this while cluster scaling and uh uh speculative decoding like MTP. Uh with MTP we can forward like one or two three four tokens at the same decode time and then the batch size they can be multipled like three or four times. uh uh also SLAS implementation uh implement a uh a a re rebalance uh uh for for just the exchange expert ways with torch P2P uh operations and then this is ablation study of uh effects of uh scale and EPLB uh uh to balancedness. The balancedness is defined as the ratio between mean computation time as a maximum computation time fore layer among GPUs. Uh the balancedness uh decreases when the system scales uh with the number of nodes. Uh and uh we can uh we can find that uh enabling EPB significantly improves the balancedness. Yeah. And this is all part of uh our uh reproduction of deepseek blog. Yeah. So the last part is our ecosystem and astron community and future development. Uh also we we are the first we are the first one to open source uh the deepseek uh large scale EPU deploy and there are some future work we need to solve. uh is the first one is the latency optimization. uh uh we the the the first token latency and the inter token latency is is still too large and we need uh tuning for real-time use cases and the sequence length is constrained by uh uh by limited uh uh 96 GPU setup and MTP and the data parallel attention is not fully integrated uh together and the EPB is only uh evaluated in uh in simulated the workload and uh and uh the dense layers of FFN they can be benefit from small uh TPIS and we only support tensor parism full tensor parism and or full data parallel and uh we also plan to uh support black well architecture and uh uh if anyone here want to join the community to uh contribute to the blackwell support you can just uh join our slack Uh yeah so this is our team. Yeah we our team is incubated by LMS is uh is a corporation as non-benefit uh and the uh the there is a list of the uh major maintainers and we have uh contributors over 400. Yeah. Uh the last page is our community adoption. So we have uh yeah the our biggest uh client is XAI and Nvidia and AMD and also thanks to PyTorch to give me this chance to give this talk. Yeah, thank you. [Applause]