[Bugfix][MRV1] Fix KV cache buffer corruption for extracting hidden states

[PHOEBEMOON0802]

When a KV cache tensor's shared_by list contains both MambaSpec and HiddenStateCacheSpec layers, sharing one physical buffer causes float32 ssm_state writes to overwrite bfloat16 hidden-states data (same bytes, different interpretation). Allocate a separate tensor for the HiddenStateCacheSpec layers in that case.

Also:

• Route AscendExtractHiddenStatesProposer to cm_base when selecting spec_decode_common_attn_metadata so it uses the base KV cache group.
• Validate that extract_hidden_states drafter layers live in a single KV cache group during initialize_kv_cache.

What this PR does / why we need it?

Fix KV cache buffer corruption for hybrid model configurations that use both Mamba layers and extract-hidden-states layers in the same KV cache tensor group. Also wire AscendExtractHiddenStatesProposer into the spec-decode attention metadata plumbing.

• The added has_mamba / has_hidden scan runs only inside the existing hybrid/hidden-state branch, so the hot path for plain attention layers is untouched.
• Alignment handling (_align_memory with kv_transfer_config) is duplicated for the new tensor_hs; keeping the two allocations symmetric was intentional so KV-transfer setups behave identically for both buffers.
• No new public API. Behavior for non-hybrid configs is byte-for-byte identical.

Does this PR introduce any user-facing change?

No.

Signed-off-by: chenyue1122 oyoy7102@163.com

• vLLM version: v0.26.0
• vLLM main: vllm-project/vllm@0351e9a

[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 critical memory corruption issue occurring in hybrid model configurations where Mamba and hidden-state extraction layers share the same physical KV cache buffer. By isolating memory allocations for these specific layer types, the fix prevents incompatible data types from overwriting each other. Additionally, the PR improves the robustness of the speculative decoding pipeline by correctly wiring the Ascend-specific proposer and adding validation checks during KV cache initialization.

Highlights

• KV Cache Buffer Corruption Fix: Implemented logic to detect and prevent shared buffer usage between MambaSpec and HiddenStateCacheSpec layers, ensuring separate memory allocation to avoid data corruption.
• AscendExtractHiddenStatesProposer Integration: Updated the model runner to correctly route AscendExtractHiddenStatesProposer to the base KV cache group and added validation to ensure drafter layers reside in a single KV cache group.

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]

👋 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]

Code Review

Suggested PR Title:

[Ops][BugFix] Prevent memory corruption between Mamba and HiddenStateCache specs in KV cache allocation

Suggested PR Summary:

### What this PR does / why we need it?
This PR adds support for `AscendExtractHiddenStatesProposer` and prevents memory corruption when `shared_by` contains both `MambaSpec` and `HiddenStateCacheSpec` by allocating separate physical memory tensors. It also adds validation for speculative decoding configuration.

Feedback:
- Avoid code duplication in `_allocate_kv_cache_tensors` by using the existing helper method `_allocate_int8_cache_tensor`.
- Replace the runtime `assert` statement with an explicit type check and `TypeError` to ensure robustness in production environments.

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

### How was this patch tested?
CI passed with existing tests.

[gemini-code-assist]

There is an existing helper method _allocate_int8_cache_tensor defined in this class that handles the exact same allocation and alignment logic (including checking kv_transfer_config and performing _align_memory). Using this helper method here for both tensor and tensor_hs avoids code duplication, simplifies the logic, and makes the code much more maintainable.

Suggested change

	if self.vllm_config.kv_transfer_config is None:
	tensor = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=self.device)
	else:
	cache_size_aligned = kv_cache_tensor.size + alignment
	tensor = torch.zeros(cache_size_aligned, dtype=torch.int8, device=self.device)
	tensor = self._align_memory(tensor, alignment)[: kv_cache_tensor.size]
	for layer_name_inner in kv_cache_tensor.shared_by:
	# shared the kvcache for all shared layers
	kv_cache_raw_tensors[layer_name_inner] = tensor
	if has_mamba and has_hidden:
	# Allocate separate tensor for HiddenStateCacheSpec layers
	# so ssm_state writes don't corrupt hidden-states data
	if self.vllm_config.kv_transfer_config is None:
	tensor_hs = torch.zeros(kv_cache_tensor.size, dtype=torch.int8, device=self.device)
	else:
	cache_size_aligned = kv_cache_tensor.size + alignment
	tensor_hs = torch.zeros(cache_size_aligned, dtype=torch.int8, device=self.device)
	tensor_hs = self._align_memory(tensor_hs, alignment)[: kv_cache_tensor.size]
	tensor = self._allocate_int8_cache_tensor(kv_cache_tensor.size, alignment)
	if has_mamba and has_hidden:
	# Allocate separate tensor for HiddenStateCacheSpec layers
	# so ssm_state writes don't corrupt hidden-states data
	tensor_hs = self._allocate_int8_cache_tensor(kv_cache_tensor.size, alignment)

[gemini-code-assist]

Using assert statements for runtime type validation or configuration checks is discouraged in production code. Assertions can be globally disabled when Python is run with optimization flags (e.g., python -O), which would bypass this type check entirely and potentially lead to an unhandled AttributeError on the subsequent line. It is safer and more robust to explicitly check the type and raise a TypeError.

Suggested change

	assert isinstance(self.drafter, AscendExtractHiddenStatesProposer)
	if not isinstance(self.drafter, AscendExtractHiddenStatesProposer):
	raise TypeError(
	"Expected self.drafter to be an instance of "
	"AscendExtractHiddenStatesProposer when "
	"uses_extract_hidden_states is True."
	)

[realliujiaxu]

LGTM. Upstream vLLM doesn't have this problem because its KV cache config builder isolates incompatible spec types into separate KVCacheTensor entries before allocation ever runs, while the Ascend fork's rewritten _allocate_kv_cache_tensors lumps all non-attention layer types into a single shared buffer, breaking that isolation.