KV cache refactor to decouple cache blocks and metadata about them #168
Conversation
| @@ -24,6 +24,7 @@ def get_default_mlc_serve_argparser(description="", allow_override=False): | |||
| parser.add_argument("--artifact-path", type=str, default="dist") | |||
| parser.add_argument("--use-sync-engine", action="store_true") | |||
| parser.add_argument("--max-num-batched-tokens", type=int, default=4096) | |||
| parser.add_argument("--num-sequences-to-sample", type=int, default=1) | |||
There was a problem hiding this comment.
This was accidentaly removed in #162 apparently.
|
I am wondering if it's still necessary to pass the |
|
I see what you mean. We could do another refactoring, where is changed to and all manipulations of However, going back to our initial design goals, we want to keep the model as pure as possible and put all states such as cache management into the engine side. The current interface is still good in that sense. Also I now realized that this PR goes against the model purity goal a bit, but I think it is reasonable since the cache blocks are updated in-place by the model anyway. |
Right now we are attaching both cache blocks and metadata about them such as
block_tableinto a single class which is passed back and forth between the engine and the model. While working on multi-gpu support for PT models, I learned that I need to use some RPC framework to manage multiple processes.The inputs from the engine need to cross the RPC boundary on each inference, and it is a bit awkward if the actual cache blocks also need to be passed from the engine always for no reason. For disco, only a handle to the cache blocks (
DRef) is communicated between the engine and the model, but even this is unnecessary.With this PR, the cache blocks are completely owned by the model and only
block_tablesetc need to be passed from the engine. I renamed the class toKVCacheInfoto reflect its new role. The interface inmodel_module.pydoesn't need to change since this change is impl details of the paged cache model.@yelite