iHore/vLLM-4x-RTX-3090
8
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity
Files
main
vLLM-4x-RTX-3090/patches/0073-recoverssm-align-keep-a-boundary-ending-state-in-its.patch
T
MOVZX 777e910dd0 Initial commit 
Signed-off-by: MOVZX <movzx@yahoo.com>
2026-10-15 01:01:52 +07:00

245 lines
10 KiB
Diff
Raw Permalink Blame History

From 57a0d446925c54448a4bfbc7879d6556a77101cc Mon Sep 17 00:00:00 2001
From: halt95 <halt95@users.noreply.github.com>
Date: Sun, 27 Sep 2026 11:06:54 +0200
Subject: [PATCH 073/103] recoverssm: align: keep a boundary-ending state in
its block (port of upstream #58863, 5567cc1b25)
Port of the upstream vLLM change #58863 (Apache-2.0). A window whose accepted tokens end exactly on a block boundary now commits into
the block it ends in, not the next, possibly unallocated, column.
---
tests/kernels/mamba/test_recoverssm_gdn.py | 9 +-
tests/kernels/mamba/test_recoverssm_ple.py | 2 +-
.../kimi_k3/test_recoverssm_boundary.py | 129 ++++++++++++++++++
.../layers/mamba/ops/recoverssm.py | 6 +-
vllm/v1/worker/gpu/model_states/recoverssm.py | 4 +-
5 files changed, 143 insertions(+), 7 deletions(-)
create mode 100644 tests/models/kimi_k3/test_recoverssm_boundary.py
diff --git a/tests/kernels/mamba/test_recoverssm_gdn.py b/tests/kernels/mamba/test_recoverssm_gdn.py
index 1e0cbcf..ce16f79 100644
--- a/tests/kernels/mamba/test_recoverssm_gdn.py
+++ b/tests/kernels/mamba/test_recoverssm_gdn.py
@@ -299,14 +299,15 @@ def test_commit_maps_noncontiguous_requests_in_align_mode():
for i, r in enumerate(req_rows.tolist()):
n = min(int(acc[r]), qlens[i])
c0 = int(nc[r])
- final_blk = int(bt[r, (c0 + n) // bs])
+ final_blk = int(bt[r, max(c0 + n - 1, 0) // bs])
tok = int(inp["query_start_loc"][i]) + n - 1
torch.testing.assert_close(ckpt[final_blk], states[tok], rtol=1e-5, atol=1e-6)
# (num_computed, accepted, block size): boundary after 1 / after 2 / ending exactly on
# it / crossing and continuing / accepted = 1 without crossing / the full window
-# crossing after 3
+# crossing after 3. A window ending exactly on a boundary keeps its state in the block
+# holding its last token: the next block may still be unallocated.
ALIGN_CASES = [
(3, 2, 4),
(2, 3, 4),
@@ -338,7 +339,7 @@ def test_align_boundary_and_final_state_match_native(nc, accepted, bs):
mamba_block_size=bs,
)
next_boundary = (nc // bs + 1) * bs
- final_blk = int(bt[0, (nc + accepted) // bs])
+ final_blk = int(bt[0, max(nc + accepted - 1, 0) // bs])
if (
nc + accepted >= next_boundary
): # crosses: the boundary state goes into the source block
@@ -585,7 +586,7 @@ def test_align_boundary_bf16_state(nc, accepted, bs):
mamba_block_size=bs,
)
next_boundary = (nc // bs + 1) * bs
- final_blk = int(bt[0, (nc + accepted) // bs])
+ final_blk = int(bt[0, max(nc + accepted - 1, 0) // bs])
if nc + accepted >= next_boundary:
torch.testing.assert_close(
ckpt[int(src)].float(),
diff --git a/tests/kernels/mamba/test_recoverssm_ple.py b/tests/kernels/mamba/test_recoverssm_ple.py
index ec11ad3..6e2e1f0 100644
--- a/tests/kernels/mamba/test_recoverssm_ple.py
+++ b/tests/kernels/mamba/test_recoverssm_ple.py
@@ -124,7 +124,7 @@ def test_ple_commit_align_boundary(num_computed, accepted, crosses):
mamba_block_size=block_size,
)
- final_block = nxt if num_computed + accepted >= block_size else src
+ final_block = nxt if num_computed + accepted > block_size else src
assert torch.equal(view[final_block, :, :STATE_LEN], want_final)
if crosses:
assert torch.equal(view[src, :, :STATE_LEN], want_boundary)
diff --git a/tests/models/kimi_k3/test_recoverssm_boundary.py b/tests/models/kimi_k3/test_recoverssm_boundary.py
new file mode 100644
index 0000000..8023e44
--- /dev/null
+++ b/tests/models/kimi_k3/test_recoverssm_boundary.py
@@ -0,0 +1,129 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Align-mode RecoverSSM state columns at recurrent block boundaries.
+
+A speculative window whose accepted tokens end exactly on a recurrent block
+boundary must keep its state in the block holding the last committed token.
+The following block can still be unallocated, so selecting it would drop the
+commit while the bookkeeping recorded the state as advanced.
+"""
+
+import pytest
+import torch
+
+from vllm.models.kimi_k3.nvidia.ops.recoverssm import _prepare_commit_plan_kernel
+from vllm.platforms import current_platform
+from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
+from vllm.v1.worker.gpu.model_states.recoverssm import (
+ _postprocess_recoverssm_align_kernel,
+)
+
+pytestmark = pytest.mark.skipif(
+ not current_platform.is_cuda_alike(),
+ reason="The RecoverSSM kernels require a CUDA-alike device.",
+)
+
+BLOCK_SIZE = 2048
+SPEC_QUERY_LEN = 4
+BLOCKS = [100, 101, 102, 103, 104, 105]
+
+# (computed before the step, committed tokens, allocated block columns,
+# final column, boundary column or None)
+CASES = [
+ # 8,188 + 4 ends exactly on the 8,192 boundary; column 4 is unallocated.
+ (8188, 4, 4, 3, 3),
+ # The window crosses the boundary: state moves on, checkpoint stays behind.
+ (8190, 4, 5, 4, 3),
+ # The window stays inside column 3.
+ (8180, 4, 4, 3, None),
+ # The first token after a boundary belongs to the new column.
+ (8192, 1, 5, 4, None),
+ # Single-token commit ending on the first boundary.
+ (2047, 1, 1, 0, 0),
+]
+
+
+@pytest.mark.parametrize(
+ "num_computed,commit_len,allocated,final_col,boundary_col", CASES
+)
+def test_commit_plan_keeps_boundary_state_in_last_computed_block(
+ num_computed, commit_len, allocated, final_col, boundary_col
+):
+ device = current_platform.device_type
+ columns = [
+ block if col < allocated else NULL_BLOCK_ID for col, block in enumerate(BLOCKS)
+ ]
+ block_table = torch.tensor([columns], dtype=torch.int32, device=device)
+ num_accepted = torch.tensor([commit_len], dtype=torch.int32, device=device)
+ state_indices = torch.tensor(
+ [BLOCKS[(num_computed - 1) // BLOCK_SIZE]], dtype=torch.int32, device=device
+ )
+ query_start_loc = torch.tensor(
+ [0, SPEC_QUERY_LEN], dtype=torch.int32, device=device
+ )
+ computed = torch.tensor([num_computed], dtype=torch.int32, device=device)
+ commit_lens, final_indices, boundary_indices, recovery_lens = (
+ torch.full((1,), -7, dtype=torch.int32, device=device) for _ in range(4)
+ )
+
+ _prepare_commit_plan_kernel[(1,)](
+ num_accepted,
+ None,
+ state_indices,
+ query_start_loc,
+ block_table,
+ computed,
+ commit_lens,
+ final_indices,
+ boundary_indices,
+ recovery_lens,
+ NULL_BLOCK_ID,
+ BLOCK_SIZE,
+ block_table.shape[1],
+ num_accepted.stride(0),
+ 0,
+ state_indices.stride(0),
+ query_start_loc.stride(0),
+ block_table.stride(0),
+ block_table.stride(1),
+ computed.stride(0),
+ SPEC_QUERY_LEN=SPEC_QUERY_LEN,
+ num_warps=1,
+ )
+
+ assert commit_lens.item() == commit_len
+ assert final_indices.item() == BLOCKS[final_col]
+ expected_boundary = NULL_BLOCK_ID if boundary_col is None else BLOCKS[boundary_col]
+ assert boundary_indices.item() == expected_boundary
+ if boundary_col is not None:
+ next_boundary = (num_computed // BLOCK_SIZE + 1) * BLOCK_SIZE
+ assert recovery_lens.item() == next_boundary - num_computed
+
+
+@pytest.mark.parametrize(
+ "num_computed,commit_len,allocated,final_col,boundary_col", CASES
+)
+def test_postprocess_records_the_committed_state_column(
+ num_computed, commit_len, allocated, final_col, boundary_col
+):
+ device = current_platform.device_type
+ idx_mapping = torch.tensor([2], dtype=torch.int32, device=device)
+ num_sampled = torch.tensor([commit_len], dtype=torch.int32, device=device)
+ computed = torch.tensor([num_computed], dtype=torch.int32, device=device)
+ state_idx = torch.full((4,), -7, dtype=torch.int32, device=device)
+ num_accepted = torch.full((4,), -7, dtype=torch.int32, device=device)
+
+ _postprocess_recoverssm_align_kernel[(1,)](
+ idx_mapping,
+ num_sampled,
+ None,
+ computed,
+ state_idx,
+ num_accepted,
+ MAMBA_BLOCK_SIZE=BLOCK_SIZE,
+ BLOCK_TABLE_WIDTH=len(BLOCKS),
+ )
+
+ assert state_idx[2].item() == final_col
+ assert num_accepted[2].item() == 1
+ assert state_idx[[0, 1, 3]].tolist() == [-7, -7, -7]
diff --git a/vllm/model_executor/layers/mamba/ops/recoverssm.py b/vllm/model_executor/layers/mamba/ops/recoverssm.py
index 073bdd5..9379688 100644
--- a/vllm/model_executor/layers/mamba/ops/recoverssm.py
+++ b/vllm/model_executor/layers/mamba/ops/recoverssm.py
@@ -78,8 +78,12 @@ def prepare_commit_plan_kernel(
tl.int32
)
final_num_computed = num_computed + commit_len
+ # The committed state belongs to the block holding the last committed
+ # token. A window ending exactly on a boundary stays in that block; the
+ # next block may still be unallocated.
final_state_col = tl.minimum(
- final_num_computed // mamba_block_size, block_table_width - 1
+ tl.maximum(final_num_computed - 1, 0) // mamba_block_size,
+ block_table_width - 1,
)
final_state_idx = tl.load(
block_table_ptr
diff --git a/vllm/v1/worker/gpu/model_states/recoverssm.py b/vllm/v1/worker/gpu/model_states/recoverssm.py
index cfbe1a4..b02a6c1 100644
--- a/vllm/v1/worker/gpu/model_states/recoverssm.py
+++ b/vllm/v1/worker/gpu/model_states/recoverssm.py
@@ -91,10 +91,12 @@ def _postprocess_recoverssm_align_kernel(
return
num_sampled = tl.load(num_sampled_ptr + batch_idx)
num_computed = tl.load(num_computed_ptr + batch_idx)
+ # Match the commit plan: the running state lives in the block holding the
+ # last computed token, so the next step migrates it across a boundary.
tl.store(
state_idx_ptr + req_state_idx,
tl.minimum(
- (num_computed + num_sampled) // MAMBA_BLOCK_SIZE,
+ tl.maximum(num_computed + num_sampled - 1, 0) // MAMBA_BLOCK_SIZE,
BLOCK_TABLE_WIDTH - 1,
),
)
--
2.43.0
Reference in New Issue View Git Blame Copy Permalink