From 7e9cd428a126b8be8d1d277986ffa6df3cb228e2 Mon Sep 17 00:00:00 2001 From: MOVZX Date: Wed, 30 Sep 2026 07:35:13 +0200 Subject: [PATCH] recoverssm: align commit: fold final states before boundary states (port of upstream #58863, 3388ba1a25) Port of the upstream vLLM change #58863 (Apache-2.0). The final-state folds run before the boundary-state folds, because in align mode a boundary state can be written into the slot a final fold still reads. Co-Authored-By: Hinata Signed-off-by: MOVZX --- .../layers/mamba/gdn/recoverssm_gdn.py | 51 ++++++++++--------- 1 file changed, 28 insertions(+), 23 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py index a95d579b35..1404edca08 100644 --- a/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py +++ b/vllm/model_executor/layers/mamba/gdn/recoverssm_gdn.py @@ -351,12 +351,12 @@ def _commit_gdn_state_kernel( BK: tl.constexpr, BV: tl.constexpr, NUM_HEADS: tl.constexpr, - ALIGN_MODE: tl.constexpr, + BOUNDARY: tl.constexpr, ): - # In align mode the second half of axis 0 folds the block-boundary state - # (the window's first boundary_len tokens) instead of the final state, so every - # program runs one fold and keeps one tile live. - NUM_V_TILES: tl.constexpr = (V + BV - 1) // BV + # Every program runs one fold and keeps one tile live. In align mode the + # block-boundary state (the window's first boundary_len tokens) is folded by a + # second launch (BOUNDARY): a crossing window's boundary state goes into its + # source block, which the first launch's final-state fold must still read. pid_v = tl.program_id(0) pid_b = tl.program_id(1) pid_lh = tl.program_id(2) @@ -372,14 +372,13 @@ def _commit_gdn_state_kernel( if final_idx <= null_block_id: return dst_idx = final_idx - # ALIGN_MODE is a constexpr: kept apart so non-align builds drop the branch. - if ALIGN_MODE: # noqa: SIM102 - if pid_v >= NUM_V_TILES: - pid_v -= NUM_V_TILES - dst_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) - if dst_idx <= null_block_id: - return - n = tl.load(boundary_recovery_lens_ptr + pid_b) + if BOUNDARY: + dst_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64) + # A window ending exactly on the boundary keeps its final state in that + # block: the first launch already wrote it (and may have overwritten src). + if dst_idx <= null_block_id or dst_idx == final_idx: + return + n = tl.load(boundary_recovery_lens_ptr + pid_b) state_ptr = tl.load(state_base_addrs_ptr + pid_l).to( tl.pointer_type(state_ref_ptr.dtype.element_ty) ) @@ -659,8 +658,7 @@ class GDNRecoverSSMCommitContext: BV = min(triton.next_power_of_2(V), 16) rr = self.replays[0] align = block_table is not None - num_v_tiles = triton.cdiv(V, BV) * (2 if align else 1) - _commit_gdn_state_kernel[(num_v_tiles, batch, num_layers * HV)]( + args = ( ref, self.state_base_addrs, self.state_block_strides, @@ -680,15 +678,22 @@ class GDNRecoverSSMCommitContext: rr.stride(2), rr.stride(3), state_indices.stride(0), - K=K, - V=V, - BK=BK, - BV=BV, - NUM_HEADS=HV, - ALIGN_MODE=block_table is not None, - num_warps=4, - num_stages=2, ) + grid = (triton.cdiv(V, BV), batch, num_layers * HV) + # Final states first: a boundary state can land in the window's source + # slot, which the final-state fold reads (stream order separates them). + for boundary in (False, True) if align else (False,): + _commit_gdn_state_kernel[grid]( + *args, + K=K, + V=V, + BK=BK, + BV=BV, + NUM_HEADS=HV, + BOUNDARY=boundary, + num_warps=4, + num_stages=2, + ) __all__ = ["GDNRecoverSSMCommitContext", "gdn_recoverssm_verify"]