[RFC][Store]Provide BatchPut() and BatchGet() interfaces for fine-grained KV Slices

[zhaoyongke]

Mooncake Store provides put()/get() interfaces with single key for now, it's ok for a couple of tokens' KV save/load.

While in vLLM and SGLang, an optimization is transfering KVs layer-by-layer for better compute-communication overlapping.

vllm offloading

sglang cache controller

However, when KVs split into layers, the value size will decreased sharply, and the number of requests increased accordingly.

To reduce the overhead of single key put()/get, we may provide batched python interfaces as:

def batch_put(key : list[str], value : list[byte]):
    ...
def batch_get(key : list[str]) -> list[byte]
    ...

To add batch_put(), batch_get() interfaces, some considerations are needed:

• Should be called asyncly
• Should provide detailed status for each key/value
• Should update metrics correctly
• (Optional)Better storage and network-bandwidth utilizing
• (Optional)Auto-batching, keep higher-level api unchanged

Welcome to codesign & review & contribue & PR!

[stmatengss]

Layer-wise KVCache has the potential to reduce TTFT as the transfer operations can be pipelined with prefill. However, careful consideration of batch size is required: processing all layers in a single batch would revert to a block-based put/get interface pattern. As an extension, is it feasible to implement cross-request layer batching (i.e., grouping layers from different requests into batches)?

[zhaoyongke]

Layer-wise KVCache has the potential to reduce TTFT as the transfer operations can be pipelined with prefill. However, careful consideration of batch size is required: processing all layers in a single batch would revert to a block-based put/get interface pattern. As an extension, is it feasible to implement cross-request layer batching (i.e., grouping layers from different requests into batches)?

At the first step, inference engines will choose a dynamic batch size as it needed, auto-batching in the mooncake store is preferred in the next step.

[zhaoyongke]

@xinranwang17

[xinranwang17]

I'd like to take this one.