karan

besides the regular serving side optimizations for llms, studying model architectures designs has been quite important as well to accelerating inference on gpus. i found this llm that recently released called Ling 3.0 flash by @AntLingAGI , and it seemed like an excellent case study to do since it has tons of the new architectural designs rather than the regular serving side optimization. ill do a 5 day case study on this model and towards the end youd have a more stronger idea to why this matters. breakdown 1 - ling 3.0 flash architecture deep dive this is a 124b moe model with a hybrid linear architecture. it has 512 routed experts, 8 active experts with 1 shared expert. sounds like a lot but lets unpack one at a time. whats an moe? moes are a type of sparse model where unlike the normal transformer block which has one FFN(feed forward layer), it replaces that FFN with multiple experts. these experts are basically FFN layers. so instead of every token going through one big FFN, the model has many FFNs sitting there and a router decides which ones should handle each token. the important word here is sparse. in a dense model, every token uses almost all the weights in every layer. in an moe model, the model can have a huge number of total parameters, but only a small subset of them are active for any given token. this is why you will see numbers like: 124b total parameters 5.5b active parameters at first this sounds weird. how can it be a 124b model but only use 5.5b? the reason is the model stores many experts, but it does not invoke all of them for every token. the core idea is: given a token, route it to a few experts, run only those experts, then combine their outputs. so what do we need to achieve this? - a router - the layer that looks at the token representation and decides which experts to use. - the experts themselves - these are the FFN blocks. - a combine step - after the selected experts run, their outputs are weighted and added back together. now training this is the tricky part. say we have 10 experts. early in training, the router does not really know which expert is useful for which token. if it keeps picking the same 1 or 2 experts, those experts get trained more, become better, and then get picked even more. some experts become overloaded, and some experts basically do nothing. that is bad because you paid for a large expert pool, but you are only using a tiny part of it. use a top-k routing method. instead of activating all experts, the router scores all experts and picks only the top-k. for Ling - 512 routed experts 8 active experts 1 shared expert so for each token, the model does not run all 512 routed experts. it chooses 8 of them. there is also 1 shared expert that is always available, which helps keep some common/general capacity around instead of forcing everything through routed experts. this is where the efficiency comes from. the model has the capacity of a large model because there are many experts. but the compute per token is closer to a much smaller model because only a few experts are active. so moe gives you this tradeoff - more total parameters for knowledge/capacity less active parameters for inference cost not free, obviously since you now have routing overhead, expert load balancing problems, and communication overhead during distributed inference. but if done well, moe is a very strong way to improve the capability per unit of active compute. cool, now lets focus on hybrid linear attention architecture. the model has 42 transformer layers - 35 KDA layers 7 gated MLA layers that is a 5:1 pattern. attention is another major inference bottleneck, especially for long context. at smaller contexts it usually memory bound but for the longer context it is compute bound. normal full attention needs to keep growing the KV cache as the context grows, but linear attention layers are quite different. for the linear layers, unlike full attention, the computed state in linear attention is superimposed on the previous historical intermediate matrix. this intermediate matrix is also called an SSM. Each new SSM can be directly added to all previous SSMs. the major difference is full attention scales with L(seq len), but in linear attention it stays fixed at d x d(hidden dim). do you see a problem here? because at each step, the attention just keeps getting added on top, you dont really know the state at each particular step is. no key points, no important information, no discard anything, everything is there just in a big bowl. so Ling mixes in MLA layers. the general idea for MLA is instead of sharing key and value heads like GQA, MLA compresses the key and value tensors into a lower-dimensional space before storing them in the KV cache. in the next posts, we will look at this with more detail and more practical numbers as well.