From 57a0d446925c54448a4bfbc7879d6556a77101cc Mon Sep 17 00:00:00 2001 From: halt95 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