[Feature] Support layerwise KV pool for Qwen3.5 hybrid GDN+full_attention models

[tyy0829]

Description

This PR adds layerwise KV pool support for Qwen3.5 (and Qwen3-Next) hybrid models that combine GDN (Gated Delta Net) linear attention layers with standard full attention layers.

Changes

1. GDN forward (ops/gdn.py): Add wait_for_kv_layer_from_connector(self.prefix) at the start of forward() and record_attention_compute_start() before the custom op call. GDN layers do not go through the @maybe_transfer_kv_layer decorator, so they must explicitly call these functions to participate in layerwise KV transfer.

2. Attention utils (attention/utils.py): Add connector.has_connector_metadata() guard to wait_for_kv_layer_from_connector and maybe_save_kv_layer_to_connector, matching the guard in the @maybe_transfer_kv_layer decorator. Prevents spurious current_layer counter increments during non-save steps (e.g. profile run).

3. Pool scheduler (pool_scheduler.py): Skip touch_sending_mamba_blocks when use_layerwise=True. The layerwise send thread (KVCacheStoreLayerSendingThread) does not use the completed_events mechanism, so touched mamba blocks would never be freed, leaking the block pool.

4. Postprocess kernel (ops/triton/mamba/postprocess.py): Update postprocess_mamba_fused_kernel signature to match vLLM v0.25.0, adding state_dim_row_count_ptr, state_dim_row_stride_ptr, idx_mapping_ptr, CONV_STATE_DIM_FIRST, HAS_IDX_MAPPING, and PRECOMPUTED_NEW_COMPUTED parameters. Preserves the triton-ascend pointer-type-cast optimization.

Testing

Verified on Atlas 800T A2 (Ascend 910B3, 8 NPU) with Qwen3.5-9B:

• Service starts successfully with use_layerwise=true and MemCache backend

• KV cache hit confirmed: 4 groups all hit 512 tokens, External prefix cache hit rate 22.9%

• 5 concurrent requests all succeed with cache hits

• MTP speculative decoding (num_speculative_tokens=3, method=qwen3_5_mtp) works correctly

• vLLM version: v0.25.1

• vLLM main: vllm-project/vllm@54503ec

[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 introduces support for layerwise KV pool functionality in Qwen3.5 hybrid models, which combine GDN and full attention layers. The changes ensure proper synchronization, prevent memory leaks in the block pool, and update the Mamba postprocess kernel to maintain compatibility with recent vLLM versions.

Highlights

• GDN Layer Integration: Enabled layerwise KV transfer for GDN layers by adding explicit connector wait and compute start signals.
• KV Pool Scheduler: Prevented memory leaks in the block pool by disabling touch_sending_mamba_blocks during layerwise operations.
• Triton Kernel Update: Updated the Mamba postprocess kernel to align with vLLM v0.25.0 and added support for advanced metadata handling.

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 [Feature] 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][Feature] Support layerwise KV transfer and DS conv layout in Mamba postprocessing

Suggested PR Summary:

### What this PR does / why we need it?
This PR introduces several enhancements and fixes to the Ascend backend:
1. Adds metadata checks in KV layer connector utilities to prevent redundant operations.
2. Integrates layerwise KV transfer wait and compute start recording into the GDN attention forward pass.
3. Updates the fused Mamba postprocessing Triton kernel to support index mapping, precomputed token counts, and a dimension-first layout for DS conv states, while hoisting pointer casts out of loops to avoid Triton compilation issues.
4. Adds an early return in `touch_sending_mamba_blocks` when layerwise transfer is enabled.

However, a critical issue was identified in `pool_scheduler.py`: returning early when `use_layerwise` is True prevents sending mamba blocks from being pinned. This can cause the block manager to free them immediately when the request finishes, leading to a race condition and data corruption if the background sending thread is still reading them.

### Does this PR introduce _any_ user-facing change?
No user-facing API changes are introduced. These are internal performance and compatibility improvements for the Ascend backend.

### How was this patch tested?
No specific tests were added in this PR. It is recommended to verify these changes with existing integration tests for the Ascend backend and add synchronization tests for the layerwise transfer.

[gemini-code-assist]

Returning early here when self.use_layerwise is True prevents the sending mamba blocks from being touched (i.e., pinned/referenced). Since request_finished and request_finished_all_groups also return False immediately when use_layerwise is True, these blocks will be freed immediately by the block manager when the request finishes. If the layerwise sending thread (KVCacheStoreLayerSendingThread) is still asynchronously reading from these blocks in the background, they can be reassigned to a new request and overwritten, leading to a critical race condition and data corruption. Consider implementing a proper synchronization or event-based completion mechanism for layerwise sending to safely defer freeing these blocks until the transfer is fully complete.