Skip to content

[BugFix][Attention] Fix C8 slot mappings under PCP and DCP - #16664

Open
recky-c wants to merge 2 commits into
vllm-project:mainfrom
recky-c:fix/indexer-c8-slot-dtype
Open

recky-c wants to merge 2 commits into
vllm-project:mainfrom
recky-c:fix/indexer-c8-slot-dtype

Conversation

@recky-c

@recky-c recky-c commented Sep 16, 2026 •

Copy link
Copy Markdown
Contributor

What this PR does / why we need it?

Fixes two C8 cache-write slot mapping issues exposed by PCP and PD disaggregation:

  1. PCP can supply an int64 indexer slot_mapping, while the native StoreKvBlockMetadata kernel reads it as int32_t*. Convert only the operator argument to int32, preserving the original metadata.
  2. In the combined PCP+DCP SFA path, PCP gathers a full token view before the deferred C8 cache write, but the inherited DCP helper truncated the slot mapping to the local token count. Override the helper in AscendSFAPCPDCPImpl so the gathered KV rows use the complete DCP slot mapping. The ordinary DCP path keeps its existing local slice.

The tests extend the existing indexer dtype coverage and verify that the combined PCP+DCP implementation keeps the complete SFA slot mapping.

Does this PR introduce any user-facing change?

Fixes incorrect indexer and SFA C8 cache writes that can produce incorrect completions when PCP, DCP, and PD disaggregation are combined. No configuration changes.

How was this patch tested?

  • Python compilation and git diff --check passed.
  • The indexer module tests passed 5/5 in the existing Ascend validation environment.
  • Hardware validation used main 28acc786 + [BugFix][Attention] Build SFA indexer DCP metadata independently #16325 (f56cb5b2) + [BugFix][Mooncake] Fix DCP mapping and support P/D-only DCP and PCP+DCP in PD disaggregation #16492 (037f5143), paired with vLLM 84030bbe, GLM-5.2 W4A8C8, MRV2, EP, async scheduling, FULL_DECODE_ONLY, MooncakeConnectorV1, and P TP8 x PCP2 -> D TP8.
  • A 48-sample probe confirmed the SFA failure: gathered KV had 4096 rows and the DCP context had 4096 slots, but only 2048 slots were passed to the C8 write. The typical long-request sample had 128 incorrect valid rows.
  • After the fixes, P DCP16 -> D DCP1 and P DCP16 -> D DCP8 each passed 28/28 requests with both SFA C8 and indexer C8 enabled. Each run transferred 87,767 external KV tokens. The isolated SFA C8 and combined-C8 boundary suites also passed 3/3.

The hardware results apply to the stated combined checkout and topology.

Signed-off-by: recky-c <ruiqicheng510@gmail.com>
@gemini-code-assist

Copy link
Copy Markdown
Contributor

Summary of Changes

Hello, I'm Gemini Code Assist1! 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 critical data type mismatch in the indexer C8 slot mapping logic when using PCP. By explicitly casting the slot_mapping to int32 before invoking the StoreKvBlockMetadata kernel, the fix ensures that write destinations are calculated correctly regardless of the input data type, preventing potential cache corruption.

Highlights

  • Data Type Mismatch Fix: Resolved an issue where slot_mapping was passed as int64 to a kernel expecting int32, which caused incorrect memory write destinations.
  • Regression Testing: Extended unit tests to verify correct handling of both int32 and int64 slot mapping inputs, ensuring robust behavior under PCP.
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. ↩

@github-actions

Copy link
Copy Markdown
Contributor

👋 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 gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request ensures that the slot mapping tensor is explicitly cast to int32 before being passed to the native metadata kernel in the Ascend attention backend, as the kernel expects int32 input. The change includes corresponding updates to the unit tests to verify this behavior across different input dtypes. As there were no review comments provided, I have no feedback to offer.

torch.ops._C_ascend.store_kv_block_metadata(
slot_mapping,
# The native metadata kernel reads slot indices as int32.
slot_mapping.to(torch.int32),

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not recommended to make changes in common code paths. Instead, add type conversion at the points where PCP introduces differences.

Signed-off-by: recky-c <ruiqicheng510@gmail.com>
@recky-c
recky-c force-pushed the fix/indexer-c8-slot-dtype branch from 840b73e to 0da6735 Compare September 16, 2026 07:49
@recky-c recky-c changed the title [BugFix][Attention] Fix indexer C8 slot mapping dtype under PCP [BugFix][Attention] Fix C8 slot mappings under PCP and DCP Sep 16, 2026
@github-actions

Copy link
Copy Markdown
Contributor

This pull request has conflicts, please resolve those before we can evaluate the pull request.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants