[Bugfix][Hybrid] Fix cross-block race on num_accepted in MRv2 align prefix cache

[fuscof-ibm]

Purpose

PR #42406 introduced the support of aligned prefix cache in MRv2 by reusing the postprocess_mamba_fused_align kernels developed for MRv1. However there is a race condition that needs to be fixed.

run_fused_postprocess_align (the MRv2 align-mode postprocess path) passes num_accepted_tokens_gpu as both the postprocess_mamba_fused_align kernel's read source (num_accepted = tl.load(num_accepted_tokens_ptr + req_idx) at the top of every program) and the in-place write target (tl.store(num_accepted_tokens_ptr + req_idx, 1).

Programs are launched with grid (num_reqs, num_layers * num_state_types).
For a given req_idx, only state_idx == 0 performs the store, but every other state_idx for the same req_idx reads the same address.

Under PRECOMPUTED_NEW_COMPUTED:

num_tokens_running_state = new_num_computed - num_accepted + 1
accept_token_bias        = aligned_new_computed - num_tokens_running_state
= num_accepted - 1   (in the aligned case)

So different layers can compute different accept_token_bias values — some perform the state copy, some hit the accept_token_bias == 0 early return — leaving the mamba state inconsistent across layers.

The correctness of current code depends on GPU launch timing/ number of SMs/ number of hybrid layers rather than the memory model.

The race is visible when

Every program loads num_accepted in its first few instructions, while state_idx == 0's store happens only after several loads, arithmetic, and copy. When all programs for a given req_idx land in the same wave they all complete their load before the store can land, so every program sees the original value. If there are multiple waves, a later-wave can observe the (1) and compute accept_token_bias == 0, while an earlier-wave program already committed to a copy with the original num_accepted > 1. The two disagree, leaving the mamba state inconsistent across layers.

The fix: In this PR we apply the same pattern MRv1 already used: write to a distinct num_accepted_tokens_out buffer, then copy back into num_accepted_tokens_gpu after the kernel finishes (both operations are enqueued on the same stream, so the copy is strictly ordered after all programs complete). The kernel branch on HAS_IDX_MAPPING for the store is removed. It now unconditionally targets num_accepted_tokens_out_ptr. Read buffer and write buffer are now disjoint, so no cross-program race.

The difference with MRv1 is that we need to copy the entire tensor (of size max_num_reqs of int32) as the requests are not contiguous. However the copy size is negligible ( 512 bytes for the default 128 requests)

Claude discovered the race.

Test Plan

python3 -m pytest tests/v1/e2e/general/test_mamba_prefix_cache.py::test_mamba_prefix_cache_mrv2 -v -s

Test Result

The MRv2 tests are passing.

---

Essential Elements of an Effective PR Description Checklist

• [ x] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• [x ] The test plan, such as providing test command.
• [ x] The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.

[njhill]

cc @izhuhaoran

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[izhuhaoran]

@fuscof-ibm Thanks for identifying and fixing this subtle cross-program race. Separating the kernel’s read source from its write target is the correct approach. I left a suggested change to further simplify the implementation: reuse the existing output buffer as an immutable snapshot for the MRv2 path, while writing reset values directly to num_accepted_tokens_gpu. This keeps MRv1 unchanged, remains safe with non-contiguous idx_mapping, and eliminates the final full-buffer copy back. Also cc @njhill

[izhuhaoran]

Suggested change

	``num_accepted_tokens_gpu`` is updated with the aid of a scratch tensor
	(reset to 1 when the accepted position stays in the running block) to
	avoid races in the `postprocess_mamba_fused_kernel`.
	``new_num_computed_tokens`` already holds the post-step computed count
	(PRECOMPUTED_NEW_COMPUTED).
	``idx_mapping`` maps batch row -> req-state slot (HAS_IDX_MAPPING).
	"""
	if num_reqs == 0 or not self.is_initialized:
	return
	# Seed the whole scratch tensor: the kernel writes at idx_mapping
	# positions which are not contiguous in [0:num_reqs] under V2.
	self.num_accepted_tokens_out.copy_(num_accepted_tokens_gpu)
	total_states = self.num_layers * self.num_state_types
	grid = (num_reqs, total_states)
	postprocess_mamba_fused_kernel[grid](
	num_accepted_tokens_gpu,
	state_idx_gpu,
	``num_accepted_tokens_gpu`` is updated in place while the kernel reads
	from a snapshot to avoid cross-program races when the accepted position
	stays in the running block and the count is reset to 1.
	``new_num_computed_tokens`` already holds the post-step computed count
	(PRECOMPUTED_NEW_COMPUTED).
	``idx_mapping`` maps batch row -> req-state slot (HAS_IDX_MAPPING).
	"""
	if num_reqs == 0 or not self.is_initialized:
	return
	# V2 reads non-contiguous idx_mapping positions, so snapshot the whole
	# decision buffer rather than only [:num_reqs].
	num_accepted_tokens_snapshot = self.num_accepted_tokens_out
	num_accepted_tokens_snapshot.copy_(num_accepted_tokens_gpu)
	total_states = self.num_layers * self.num_state_types
	grid = (num_reqs, total_states)
	postprocess_mamba_fused_kernel[grid](
	num_accepted_tokens_snapshot,
	state_idx_gpu,

[fuscof-ibm]

Thanks @izhuhaoran . Indeed one can save a D2D copy. I pushed the update you suggested (and rerun the test).

CC @njhill

[izhuhaoran]

Suggested change

	num_accepted_tokens_gpu,
	idx_mapping,
	num_reqs,
	block_size=self.block_size,
	COPY_BLOCK_SIZE=1024,
	CONV_STATE_DIM_FIRST=is_conv_state_dim_first(),
	HAS_IDX_MAPPING=True,
	PRECOMPUTED_NEW_COMPUTED=True,
	)

[izhuhaoran]

Thanks for this fix. LGTM

[njhill]

Thanks @fuscof-ibm @izhuhaoran