[v0.23.0][BugFix] Isolate layerwise GVA keys by parallel rank

[Pz1116]

What this PR does / why we need it?

Isolate AscendStore layerwise GVA entries by PCP, DCP, and PP rank, and require scheduler hit checks to find every parallel-rank and TP/head key.

With DSA CP enabled, each DCP rank owns a different local KV shard. The previous layerwise path omitted the DCP rank from its key and skipped non-leading ranks according to put_step, so DCP-local SFA caches could alias or never be saved. This change keeps the existing put_step deduplication when DCP is disabled, while allowing every DCP rank to allocate and save its local cache.

The PP rank is also part of the key because pipeline stages address dense local layer ranges from offset zero; sharing one GVA entry could otherwise make layers from different stages overwrite each other, including an MTP layer colliding with a target-model layer at the same local offset.

The change covers both single-group and multi-group KV cache layouts.

Does this PR introduce any user-facing change?

No. Existing layerwise GVA cache entries use the old key format and will be treated as cache misses.

How was this patch tested?

• AscendStore pool worker and scheduler unit tests: 137 passed.

• All changed-file pre-commit hooks passed.

• vLLM version: v0.23.0

• vLLM main: vllm-project/vllm@ee0da84

[github-actions]

👋 Hi! Thank you for contributing to the vLLM Ascend project. The following points will speed up your PR merge:‌‌

• A PR should do only one thing, smaller PRs enable faster reviews.
• Every PR should include unit tests and end-to-end tests ‌to ensure it works and is not broken by other future PRs.
• Write the commit message by fulfilling the PR description to help reviewer and future developers understand.

If CI fails, you can run linting and testing checks locally according Contributing and Testing.

---

Tip

💡 Consider Linking a Related Issue or RFC

Your PR title contains the [BugFix] tag, indicating a bug fix or new feature.

Linking a related issue or RFC in the PR description is strongly encouraged — it gives reviewers helpful context and speeds up the review. You can use any of these keywords:

• Fixes #<issue_number>
• Closes #<issue_number>
• Resolves #<issue_number>
• Refs #<rfc_or_issue_number> (for RFCs)

🙏 Thanks for helping us keep the project well-organized!

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request addresses a potential collision issue in AscendStore where different pipeline-parallel stages could overwrite each other's GVA cache entries due to shared local layer offsets. By incorporating the PP rank into the GVA key structure and updating the scheduler's hit-check logic to account for these unique keys, the system now correctly isolates cache entries across pipeline stages.

Highlights

• GVA Key Isolation: Updated AscendStore layerwise GVA key generation to include the pipeline-parallel (PP) rank, preventing collisions between different pipeline stages.
• Scheduler Hit Checks: Modified the scheduler to verify hits across all PP and TP/head ranks, ensuring comprehensive coverage for multi-stage pipeline models.
• Testing: Added a new unit test in test_pool_scheduler.py to verify that GVA hit checks correctly include all PP ranks.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

Suggested PR Title:

[Ops][Feature] Include PP rank in layerwise GVA keys for KV transfer

Suggested PR Summary:

### What this PR does / why we need it?
This PR updates the layerwise GVA key generation logic in both the pool scheduler and pool worker to include the pipeline parallel (PP) rank (`pp_rank`). This ensures that each PP worker, which stores a dense local layer range, has its own unique GVA allocation key.

Feedback from the review suggests:
- Reducing redundancy in `_make_layerwise_gva_keys_for_hit_check` and `_make_layerwise_gva_key` by dynamically constructing the group prefix.
- Refactoring the unit test to avoid using object attributes directly as loop variables, which is non-standard.

### Does this PR introduce _any_ user-facing change?
No.

### How was this patch tested?
A new unit test `test_gva_hit_check_includes_all_pp_ranks` was added to `test_pool_scheduler.py` to verify that the generated keys include all PP ranks.

[gemini-code-assist]

The conditional blocks for single-group and multi-group key generation are highly redundant. We can simplify this by dynamically constructing a group prefix, which reduces duplication and improves maintainability.

        group_prefix = f"@{group_id}" if len(self.kv_cache_group_ids) > 1 else ""
        return [
            f"{self.model_name}{group_prefix}@{block_hash_hex}@{h}@{p}"
            for h in range(head_or_tp_ranks)
            for p in range(self.pp_size)
        ]

[gemini-code-assist]

The conditional block here can be simplified by dynamically constructing a group prefix, reducing duplication and improving readability.

        group_prefix = f"@{group_id}" if self.num_kv_cache_groups > 1 else ""
        return f"{self.model_name}{group_prefix}@{block_hash_hex}@{self.head_or_tp_rank}@{self.pp_rank}"

[gemini-code-assist]

Using object attributes directly as loop variables (for worker.head_or_tp_rank in ...) is highly non-standard and can confuse static analysis tools, linters, and other developers. It is much cleaner and more idiomatic to use standard loop variables and assign them explicitly inside the loop body.

Suggested change

	for worker.head_or_tp_rank in range(scheduler.tp_size):
	for worker.pp_rank in range(scheduler.pp_size):
	self.assertIn(worker._make_layerwise_gva_key(group_id, "deadbeef"), keys)
	for h in range(scheduler.tp_size):
	for p in range(scheduler.pp_size):
	worker.head_or_tp_rank = h
	worker.pp_rank = p
	self.assertIn(worker._make_layerwise_gva_key(group_id, "deadbeef"), keys)