bugfix: support disagg PD for MTP.

[RobbieLeung]

No description provided.

[zhang-minchao]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the KV cache transfer mechanism, seemingly to support disaggregated pipeline parallelism for MTP. The changes primarily involve removing the logic for transferring embeddings and modifying the KV cache allocation and deallocation to use direct aclrt calls instead of the llm_data_dist abstraction. The refactoring of push_kv_blocks to consolidate duplicated logic is a good improvement for maintainability.

My review includes suggestions to remove dead code, fix a potential type-safety issue with pointer casting, and improve code style and memory management. Overall, the changes are logical and well-structured.

[RobbieLeung]

/gemini review

[gemini-code-assist]

Code Review

This pull request refactors the KV cache transfer mechanism, particularly for speculative decoding on NPU, to support disaggregated pipeline parallelism. The key changes include:

• Modifying the KV cache allocation to use aclrtMalloc for direct device memory allocation followed by registration with llm_data_dist, instead of relying on llm_data_dist for allocation.
• Removing the logic for handling a separate embedding cache transfer, which simplifies the SpecKVCacheTransfer class by removing several functions and member variables.
• Refactoring push_kv_blocks and push_kv_blocks_spec to eliminate code duplication by introducing a common push_kv_blocks_internal function.

The changes are well-structured and improve code clarity and maintainability. I have a couple of suggestions to improve robustness and clarity: one regarding the use of a logging macro with a misleading error message, and another to prevent a potential null pointer dereference.