101 lines
4.2 KiB
Diff
101 lines
4.2 KiB
Diff
From c2a86c85f32040d36d0837657994478c731101a7 Mon Sep 17 00:00:00 2001
|
|
From: halt95 <halt95@users.noreply.github.com>
|
|
Date: Wed, 30 Sep 2026 07:35:13 +0200
|
|
Subject: [PATCH 078/103] 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.
|
|
---
|
|
.../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 a95d579..1404edc 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"]
|
|
--
|
|
2.43.0
|
|
|