[HiCache] Add CP support for HiCache

[ShangmingCai]

Motivation

This PR is mostly for Qwen3 CP + Hicache. MLA models + Hicache will reuse cp 0's data for all ranks, so we don't need to distinguish the key.
CC: @whybeyoung

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review Process

1. Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
4. After green CI and required approvals, ask Merge Oncalls to merge.

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[vladnosiv]

Hi !
Don't you think that for full support, you also need to synchronize the cache state, as it is done here #20460 ?

[ShangmingCai]

Hi ! Don't you think that for full support, you also need to synchronize the cache state, as it is done here #20460 ?

Yeah, you are right, I thought it was only a tag issue, so I basically vibe this, but it turns out we need to handle the control plane coordinate as well.

[stmatengss]

Should we consider separating checking for enable_pp and enable_cp, or if only enable_pp is used, then attn_cp_rank should be the default value (_0)?

[ShangmingCai]

That depends on our design. If we want to support hetero setups, then maybe we should let the suffix always be f"{self.local_rank}_{self.pp_rank}_{self.attn_cp_rank}" ). They can all be 0 for most cases. But we might also need to put the pp size and cp size in the suffix.

[ShangmingCai]

/tag-and-rerun-ci

[ShangmingCai]

CI has passed. Let us make it more robust after hicache refactor.

[vladnosiv]

Hi !
I've been testing in the context of my PR (#20460), and I have questions about these changes.

Why is the information about CP entered in the kv cache key?
CP uses a strategy of splitting Q and replicating KV. That is, before each attention layer, it's guaranteed that the KV caches on the CP ranks are the same, and after the layer after all-gather, this guarantee is fulfilled again.

This means that each CP rank must receive the same cache when it is retrieved from the cache.
And when writing, it is enough for one rank to be written to the cache.

In other words, CP heterogeneity is automatically maintained and, in general, knowledge about CP in keys is not required.

I maybe missing something, but so far in my PR I have removed information about CP from the keys, review these changes, please

[ShangmingCai]

Hi ! I've been testing in the context of my PR (#20460), and I have questions about these changes.

Why is the information about CP entered in the kv cache key? CP uses a strategy of splitting Q and replicating KV. That is, before each attention layer, it's guaranteed that the KV caches on the CP ranks are the same, and after the layer after all-gather, this guarantee is fulfilled again.

This means that each CP rank must receive the same cache when it is retrieved from the cache. And when writing, it is enough for one rank to be written to the cache.

In other words, CP heterogeneity is automatically maintained and, in general, knowledge about CP in keys is not required.

I maybe missing something, but so far in my PR I have removed information about CP from the keys, review these changes, please

@vladnosiv Current CP impl makes each CP rank allocate the full KVCache length, and use allgather to fetch from peers. But we are thinking of supporting a ring-based solution without a full local copy, so I assume it for MHA/GQA now, since MHA/GQA are more memory-intensive than MLA. Anyway, I think it can be changed and optimized at that time, when we are supporting models like Qwen3.5 or others.