one of the most unintuitive problems to solve for in ml inference is the routing problem imo.
and i claim that this, and this alone, is responsible for the birth of every flavor of sparse attention (SWA, NSA, DSA, etc).
to set up why routing birthed sparse attention:
you have multiple nodes, and you have multiple requests.
1) can do pure load-balancing:
request 1 goes to node 1
request 2 to node 2 (node 1 is busy)
request 3 to node 3....
this would be perfect if the model did not depend on the kv cache. but it does, and so this is a terrible system, since you do prefill over the same tokens many times.
on the other extreme:
2) can do pure cache-aware routing.
request 1 goes to node 1 ( kv cache stored on node 1)
request 2 goes to node 1 (the kv cache is on node 1)
request 3 goes to ... node 1 (the kv cache ...)
this would be perfect if the system only depended on the kv cache. but it does not. decode time makes the requests queue, so this is also a terrible system.
and so it seems that the two pull in opposite directions (at least if we run both the cache aware and load based routing to their extreme case)
this is analogous to a somewhat dysfunctional engineering team. if ticket i+1 comes up, do you give it to the person who has context on ticket i, or do you route it to someone who has to onboard? the former is probably working on something else, but could finish it quickly. the latter can start on it immediately, but will take a while.
you have to do a balance of the 2.
and so the solution becomes a calculation that the router has to do:
what is the
a) estimated queue time
b) number of uncached tokens
c) speed of prefill (processing said uncached tokens)
and it assigns the total=a+b/c to each node, picking the node with the least total time at a per request level.
but, imo, this is a patch that exists because the kv cache stays on the node it was computed on, because we don't have the ability to keep one shared kv cache / one shared storage system for every node within a cluster...
but i'm obviously wrong. mooncake (and others) exists. maintaining a shared kv pool is a decade old idea...
but then:
1) you're still limited by the fact that you can't have a remote / distributed kv pool. you need every single block in the kv to be present on the node you're doing decode on (due to HBM speed and decode being mem bound), so you must pay the price of doing the transfer.
2) the global pool can only try to transfer some of the kv to a node, but cannot guarantee, because of evictions, and long transfers at scale.
even if transfers were guaranteed, even if evictions never happened, you still need (b)/(c) for the 'redo prefill' vs 'wait for kv transfer' arbitrage.
so then:
- what is the point of maintaining global kv pool, if only to track where distributed kv blocks are and do said arbitrage, if it cannot be used to directly draw kv blocks from?
- can this only be solved when network bandwidth equals that of HBM (ie unsolvable problem)?
and, on the premise that the answer to both points above is yes, the solution becomes to need the least amount of kv blocks per request, and this is why every open source lab is racing to show off a new and shiny sparse attention flavor.
もっと見る