Are there any KV cache movement when using PP/TP?

[leizhenyuan]

I am not familiar with distributed. Below is my under standing:
For Pipeline parallel, we only need to transfer intermediate tensor between different node, so there are no KV cache communication.
For Tensor Parallel, different heads are allocated on different device, each head only need to calculate own output, there is also no KV cache communication.

But i also found need_recv_kvcache when worker is a driver worker, so i think my above understanding must be wrong, Can anyone answer this question? Thanks.

[comaniac]

Your understanding is correct. The API you found is used for prefill disaggregation. In this case we transfer the kv cache from a prefill worker to a decode worker.

[leizhenyuan]

Got it Thanks for you reply.