[BugFix][Mamba] Overlap-safe state copies for align-mode block moves

[Liears]

Problem

The NPU path binds upstream's pointer-based collect_mamba_copy_meta and executes align-mode mamba state moves through the Ascend triton batch_memcpy port, which dropped upstream's is_left_overlap guard (vllm/v1/worker/mamba_utils.py:639-652). Same-pool moves can have overlapping src/dst ranges; without the guard the kernel is a program-order race whose corruption lands in the KV/mamba pools and is consumed by later kernels — the EZ9999 "MTE accesses an invalid GM address" crash family:

• faulting kernel drifts between runs (GatherPaKvCache / aiv_all_reduce) — the corruption is read by whoever runs next
• delayed manifestation (seconds after the corrupting write, often at the next host sync)
• block-reuse dependence (second round / next request after a long sequence completes)
• any host sync (per-step, prep-end, launch-blocking) masks it — classic async write-vs-read race

Fix

1. Collect tensor pairs alongside the pointer buffers (bind the torch collector on the NPU path too) and copy via dst.copy_(src.clone()) — memmove semantics identical to upstream's guarded kernel. GLM53_BATCH_MEMCPY=1 restores the raw pointer kernel for A/B.
2. Add upstream's is_left_overlap + tl.debug_barrier() guard to the fused postprocess copy loops (defensive; today's conv-state shift fits in a single load-then-store block iteration, so it is safe by construction).

Verification (GLM-5.3-Flash, hybrid KDA + spec=3 + align mode, TP8)

Case	Before	After
33.5k+15k concurrent, round-2 block reuse	crash (6/6 runs, EZ9999)	pass
3-way concurrent (33.5k+15k+30k)	crash	pass
2000-token decode at high context	crash	pass
prefix caching	had to be disabled	re-enabled

🤖 Generated with Claude Code

• vLLM main: vllm-project/vllm@ba07e4a

[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 critical data corruption issues occurring during Mamba state moves on the NPU path. By introducing overlap-safe copy mechanisms and adding necessary synchronization barriers in the Triton kernels, the changes resolve race conditions that previously led to intermittent kernel crashes and memory corruption. These fixes restore stability for complex concurrent workloads and prefix caching scenarios.

Highlights

• Overlap-Safe State Copies: Implemented a tensor-pair copy path using dst.copy_(src.clone()) to ensure safe memory movement when source and destination ranges overlap, effectively replicating upstream's guarded kernel behavior.
• Triton Kernel Guarding: Added is_left_overlap checks and tl.debug_barrier() to the fused postprocess Triton kernels to prevent race conditions during same-block state moves.
• Debugging Support: Introduced an environment variable GLM53_BATCH_MEMCPY=1 to allow toggling back to the raw pointer-based kernel for A/B testing and troubleshooting.

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

This pull request introduces overlap-safe state copies for Mamba block moves on Ascend NPUs by adding overlap guards and debug barriers in the Triton postprocess kernel, and implementing a fallback tensor-pair copy path using PyTorch's clone and copy operations to handle overlapping memory ranges safely. Feedback on the changes highlights a critical correctness bug in the CPU fallback path where using input_batch.num_accepted_tokens_cpu instead of num_accepted_tokens_cpu_tensor ignores actual GPU accepted counts and risks using stale data. Additionally, the reviewer noted a style guide violation regarding the PR title and summary format, providing the correct compliant templates.

[gemini-code-assist]

Critical Correctness Bug: Incorrect Accepted Tokens Used in CPU Fallback

In _postprocess_mamba_align_gpu_cpu_fallback, num_accepted_tokens has been changed to use input_batch.num_accepted_tokens_cpu instead of num_accepted_tokens_cpu_tensor.

This is a critical bug because:

1. num_accepted_tokens_cpu_tensor is the tensor that receives the actual accepted token counts copied from the GPU (num_accepted_tokens_gpu). By using input_batch.num_accepted_tokens_cpu, the actual accepted counts from the GPU are completely ignored, and stale/incorrect values are used instead.
2. As the deleted comment warned, input_batch rows may be condensed or reused by async scheduling before this fallback runs, making input_batch.num_accepted_tokens_cpu mutable and unsafe to use.

This will cause incorrect state copying and potential state corruption on 310P devices where this CPU fallback path is executed. Please revert this change to use num_accepted_tokens_cpu_tensor.

Suggested change

	num_accepted_tokens = input_batch.num_accepted_tokens_cpu
	num_accepted_tokens = num_accepted_tokens_cpu_tensor

[gemini-code-assist]

Style Guide Violation: PR Title and Summary Format

The current Pull Request title and summary do not adhere to the repository's style guide. Specifically:

1. The PR title format should be [Branch][Module][Action] Pull Request Title. The current title [BugFix][Mamba] ... places the Action before the Module, and uses an invalid Module [Mamba] instead of [Ops].
2. The PR summary should follow the specified markdown template.

Below are the suggested PR Title and PR Summary formatted according to the style guide:

Suggested PR Title:

[Ops][BugFix] Overlap-safe state copies for align-mode block moves

Suggested PR Summary:

### What this PR does / why we need it?

This PR fixes a potential data corruption and crash issue (EZ9999) in the Mamba state copy path on Ascend NPUs. The Triton-based `batch_memcpy` port dropped the upstream `is_left_overlap` guard, which is unsafe for same-pool moves with overlapping src/dst ranges. 

To fix this, the PR:
1. Binds the torch-based collector on the NPU path to collect tensor pairs and performs copies via `dst.copy_(src.clone())` to ensure safe memmove semantics.
2. Adds the `is_left_overlap` check and `tl.debug_barrier()` guard to the fused postprocess copy loops.

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

No.

### How was this patch tested?

Verified with GLM-5.3-Flash (hybrid KDA + spec=3 + align mode, TP8) under concurrent high-context decode and prefix caching scenarios.

References

1. The Pull Request Title must follow the format '[Branch][Module][Action] Pull Request Title'. The current title '[BugFix][Mamba] ...' violates this format by placing the Action before the Module, and using an invalid Module '[Mamba]' instead of '[Ops]'. ^(link)

[nu11ptr7]

This replaces the CPU snapshot (copied before the fallback runs) with a direct read of input_batch.num_accepted_tokens_cpu - the removed comment warned that input_batch rows can be condensed/reused by async scheduling before this fallback consumes them. Is the new field guaranteed to hold this step's values?

[Liears]

Good catch — no, it is not guaranteed. Async scheduling can condense/reuse input_batch rows before this fallback consumes them, so reading the live field was wrong.

Fixed in 87906e2: the fallback now reads the per-step snapshot staged into num_accepted_tokens_cpu_tensor (copied from the GPU accepted counts right above), which is the same source of truth the fused kernel consumes. The preprocess_mamba reads of input_batch.num_accepted_tokens_cpu are untouched — those run before the scheduler can mutate the rows, matching upstream semantics.