From c2881ce60302b5455867d2c29cdfae5fbeddecac Mon Sep 17 00:00:00 2001 From: Francesco Fusco Date: Tue, 4 Aug 2026 00:54:24 +0200 Subject: [PATCH] [Bugfix][Hybrid] Fix cross-block race on num_accepted in MRv2 align prefix cache (#50432) Signed-off-by: Francesco Fusco --- vllm/v1/worker/mamba_utils.py | 25 +++++++++++++++---------- 1 file changed, 15 insertions(+), 10 deletions(-) diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index 131432fffb..ea0a9b8a1b 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -247,12 +247,9 @@ def postprocess_mamba_fused_kernel( dest_block_idx = aligned_new_computed // block_size - 1 # Update accepted-token count before early exits (per-request, so only - # state_idx == 0 writes). V2 updates in place; V1 writes the _out buffer. + # state_idx == 0 writes). if src_block_idx == dest_block_idx and state_idx == 0: - if HAS_IDX_MAPPING: - tl.store(num_accepted_tokens_ptr + req_idx, 1) - else: - tl.store(num_accepted_tokens_out_ptr + req_idx, 1) + tl.store(num_accepted_tokens_out_ptr + req_idx, 1) # Skip no-op self-copy. if src_block_idx == dest_block_idx and accept_token_bias == 0: @@ -855,17 +852,25 @@ class MambaSpecDecodeGPUContext: """V2 align postprocess: save the running state to the block-aligned position after spec-decode acceptance leaves the sequence non-aligned. - ``num_accepted_tokens_gpu`` is updated in place (reset to 1 when the - accepted position stays in the running block); ``new_num_computed_tokens`` - already holds the post-step computed count (PRECOMPUTED_NEW_COMPUTED). + ``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_gpu, + num_accepted_tokens_snapshot, state_idx_gpu, None, # num_scheduled: unused under PRECOMPUTED_NEW_COMPUTED new_num_computed_tokens_gpu, @@ -880,7 +885,7 @@ class MambaSpecDecodeGPUContext: self.state_group_indices, self.state_dim_row_count, self.state_dim_row_stride, - None, # num_accepted_out: V2 updates num_accepted in place + num_accepted_tokens_gpu, idx_mapping, num_reqs, block_size=self.block_size,