[Bugfix][Hybrid] Fix cross-block race on num_accepted in MRv2 align prefix cache (#50432)
Signed-off-by: Francesco Fusco <ffu@zurich.ibm.com>
Francesco Fusco committed
Aug 4, 2026 at 00:54 UTC
c2881ce60302b5455867d2c29cdfae5fbeddecac
1 file changed
+15
-10
vllm/v1/worker/mamba_utils.py
+15
-10
@@ -247,12 +247,9 @@ def postprocess_mamba_fused_kernel(
247
dest_block_idx = aligned_new_computed // block_size - 1
248
249
# Update accepted-token count before early exits (per-request, so only
250
- # state_idx == 0 writes). V2 updates in place; V1 writes the _out buffer.
250
+ # state_idx == 0 writes).
251
if src_block_idx == dest_block_idx and state_idx == 0:
252
- if HAS_IDX_MAPPING:
253
- tl.store(num_accepted_tokens_ptr + req_idx, 1)
254
- else:
255
- tl.store(num_accepted_tokens_out_ptr + req_idx, 1)
252
+ tl.store(num_accepted_tokens_out_ptr + req_idx, 1)
253
254
# Skip no-op self-copy.
255
if src_block_idx == dest_block_idx and accept_token_bias == 0:
@@ -855,17 +852,25 @@ class MambaSpecDecodeGPUContext:
852
"""V2 align postprocess: save the running state to the block-aligned
853
position after spec-decode acceptance leaves the sequence non-aligned.
854
858
- ``num_accepted_tokens_gpu`` is updated in place (reset to 1 when the
859
- accepted position stays in the running block); ``new_num_computed_tokens``
860
- already holds the post-step computed count (PRECOMPUTED_NEW_COMPUTED).
855
+ ``num_accepted_tokens_gpu`` is updated in place while the kernel reads
856
+ from a snapshot to avoid cross-program races when the accepted position
857
+ stays in the running block and the count is reset to 1.
858
+ ``new_num_computed_tokens`` already holds the post-step computed count
859
+ (PRECOMPUTED_NEW_COMPUTED).
860
``idx_mapping`` maps batch row -> req-state slot (HAS_IDX_MAPPING).
861
"""
862
if num_reqs == 0 or not self.is_initialized:
863
return
864
+
865
+ # V2 reads non-contiguous idx_mapping positions, so snapshot the whole
866
+ # decision buffer rather than only [:num_reqs].
867
+ num_accepted_tokens_snapshot = self.num_accepted_tokens_out
868
+ num_accepted_tokens_snapshot.copy_(num_accepted_tokens_gpu)
869
+
870
total_states = self.num_layers * self.num_state_types
871
grid = (num_reqs, total_states)
872
postprocess_mamba_fused_kernel[grid](
868
- num_accepted_tokens_gpu,
873
+ num_accepted_tokens_snapshot,
874
state_idx_gpu,
875
None, # num_scheduled: unused under PRECOMPUTED_NEW_COMPUTED
876
new_num_computed_tokens_gpu,
@@ -880,7 +885,7 @@ class MambaSpecDecodeGPUContext:
885
self.state_group_indices,
886
self.state_dim_row_count,
887
self.state_dim_row_stride,
883
- None, # num_accepted_out: V2 updates num_accepted in place
888
+ num_accepted_tokens_gpu,
889
idx_mapping,
890
num_reqs,
891
block_size=self.block_size,