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 <hinata@idihore.id>
Signed-off-by: MOVZX <movzx@yahoo.com>

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"]