[Bugfix][V1][P/D]Fix the issue where repeated requests for the same input produce abnormal outputs for P2pNcclConnector

[Abatom]

Purpose

• A user reported — and I also observed — that after [V1][P/D]Enhance Performance and code readability for P2pNcclConnector #20906, when making repeated identical requests, the subsequent requests produce abnormal outputs, such as in issue [Bug]: [xPyD]Abnormal results when using v1 P2pNcclConnector as KV cache transport: repeated requests for the same input produce abnormal outputs #22965.
• Revert "The KVCache sender offloads the KVCache extraction and reshape operations to a dedicated sending thread" of [V1][P/D]Enhance Performance and code readability for P2pNcclConnector #20906,

Test Plan

Test Result

(Optional) Documentation Update

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[gemini-code-assist]

Code Review

This pull request addresses a bug where repeated identical requests could lead to abnormal outputs. The fix involves reverting a change that offloaded KVCache extraction to a background thread. By moving the extraction logic back into the main thread within P2pNcclConnector.save_kv_layer, it resolves a potential race condition where the KV cache could be modified before being sent, ensuring data integrity. The changes in P2pNcclEngine are consistent with this revert.

I've identified a potential critical issue where attn_metadata could be None, leading to a crash. Please see my detailed comment.

[gemini-code-assist]

There is a potential None value for attn_metadata that is not handled here, which could lead to a runtime crash.

The start_load_kv method in this class (line 122) checks if attn_metadata is None and returns early. However, save_kv_layer does not have this check. If attn_metadata is None, isinstance(attn_metadata, MLACommonMetadata) on line 257 will evaluate to False. For an MLA model, this will cause the code to incorrectly take the non-MLA path, resulting in an IndexError on line 261 because layer.shape will have 3 dimensions instead of the expected 4.

To prevent this crash, it's recommended to add a check for attn_metadata is None at the beginning of the save_kv_layer method, similar to the implementation in start_load_kv.

[Abatom]

@KuntaiDu PTAL!