Kimi linear 48b uses the same architecture as k3 and is public. We are using that to understand the trade offs of checkpointing recurrent state every N amount of tokens. Unlike caching KV, recurrent state has a forget gate so you need to continuously check point or the alternative is to redo prefill. No one is actually talking about any of this because they are retarded. You can’t run kimi k3 optimally if you don’t start somewhere else first. Since you lose information over time you need to continuously check point recurrent state and fetch it. The recurrent state will always be a fixed size regardless of context and is dense bf16, lower precision and sparsity won’t work. There is a trade off between how often you checkpoint versus redoing prefill. If you checkpoint all the time it actually would be larger than just doing pure GQA, we estimate the optimal amount is every 11k tokens but with flash being cheaper you can get away with 2k tokens. Additionally the recurrent state update at low batches (sub 8) gets bottlenecked by memory read, update, write back to HBM. Just thought I’d share as Nvidia doesn’t have some sort of scratchpad sram that makes this cost less. Hope that helps, we could share the implementation but honestly I don’t think you guys would get it. https://t.co/3Yhuc6vjtf
