diff --git a/tests/kernels/mamba/test_checkpointing_ssu_stp.py b/tests/kernels/mamba/test_checkpointing_ssu_stp.py new file mode 100644 index 000000000000..7ffe595fa5c7 --- /dev/null +++ b/tests/kernels/mamba/test_checkpointing_ssu_stp.py @@ -0,0 +1,1340 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""STP parity tests for FlashInfer checkpointing SSU. + +These tests compare the new FlashInfer ``checkpointing_ssu`` path against the +old FlashInfer ``selective_state_update`` path in the single-token decode flow +that vLLM uses without speculative decoding. +""" + +from __future__ import annotations + +import pytest +import torch + +from vllm.config.mamba import MambaBackendEnum, MambaConfig +from vllm.model_executor.layers.mamba.ops.ssu_dispatch import FlashInferSSUBackend +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + +try: + import flashinfer.mamba # noqa: F401 + + HAS_FLASHINFER = True +except Exception: # pragma: no cover - environment-dependent + HAS_FLASHINFER = False + + +requires_flashinfer = pytest.mark.skipif( + not HAS_FLASHINFER or not torch.cuda.is_available(), + reason="FlashInfer Mamba kernels require CUDA and flashinfer", +) + + +def _make_backend( + *, + checkpoint_interval: int = 1, + enable_stochastic_rounding: bool = False, +) -> FlashInferSSUBackend: + return FlashInferSSUBackend( + MambaConfig( + backend=MambaBackendEnum.FLASHINFER, + checkpoint_interval=checkpoint_interval, + enable_stochastic_rounding=enable_stochastic_rounding, + stochastic_rounding_philox_rounds=10, + ) + ) + + +def _make_weights( + *, + nheads: int, + head_dim: int, + dstate: int, + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + # A is per-head and broadcast over the head_dim/dstate axes. + A = -torch.rand(nheads, device=device, dtype=torch.float32) + A = A[:, None, None].expand(nheads, head_dim, dstate) + D = torch.zeros(nheads, device=device, dtype=dtype)[:, None].expand( + nheads, head_dim + ) + dt_bias = torch.zeros(nheads, device=device, dtype=dtype)[:, None].expand( + nheads, head_dim + ) + return A, D, dt_bias + + +def _make_checkpointing_cache( + *, + cache_size: int, + nheads: int, + head_dim: int, + dstate: int, + ngroups: int, + max_window: int, + device: torch.device, + dtype: torch.dtype, +) -> dict[str, torch.Tensor]: + return { + "old_x": torch.zeros( + cache_size, + max_window, + nheads, + head_dim, + device=device, + dtype=dtype, + ), + "old_B": torch.zeros( + cache_size, + 2, + max_window, + ngroups, + dstate, + device=device, + dtype=dtype, + ), + "old_dt": torch.zeros( + cache_size, + 2, + nheads, + max_window, + device=device, + dtype=torch.float32, + ), + "old_cumAdt": torch.zeros( + cache_size, + 2, + nheads, + max_window, + device=device, + dtype=torch.float32, + ), + "cache_buf_idx": torch.zeros(cache_size, device=device, dtype=torch.int32), + "prev_num_accepted_tokens": torch.zeros( + cache_size, device=device, dtype=torch.int32 + ), + } + + +def _make_decode_inputs( + *, + batch_size: int, + nheads: int, + head_dim: int, + dstate: int, + ngroups: int, + device: torch.device, + dtype: torch.dtype, + seed: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + generator = torch.Generator(device=device).manual_seed(seed) + x = 0.1 * torch.randn( + batch_size, + nheads, + head_dim, + device=device, + dtype=dtype, + generator=generator, + ) + dt_base = 0.1 * torch.randn( + batch_size, + nheads, + device=device, + dtype=dtype, + generator=generator, + ) + dt = dt_base.unsqueeze(-1).expand(batch_size, nheads, head_dim) + B = 0.1 * torch.randn( + batch_size, + ngroups, + dstate, + device=device, + dtype=dtype, + generator=generator, + ) + C = 0.1 * torch.randn( + batch_size, + ngroups, + dstate, + device=device, + dtype=dtype, + generator=generator, + ) + return x, dt, B, C + + +def _call_backend( + backend: FlashInferSSUBackend, + *, + state: torch.Tensor, + cache: dict[str, torch.Tensor] | None, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + D: torch.Tensor, + dt_bias: torch.Tensor, + z: torch.Tensor | None = None, + state_batch_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + dst_state_batch_indices: torch.Tensor | None = None, + max_seqlen: int = 1, +) -> torch.Tensor: + out = torch.empty_like(x) + kwargs = {} + if cache is not None: + kwargs = { + "old_x": cache["old_x"], + "old_B": cache["old_B"], + "old_dt": cache["old_dt"], + "old_cumAdt": cache["old_cumAdt"], + "cache_buf_idx": cache["cache_buf_idx"], + "prev_num_accepted_tokens": cache["prev_num_accepted_tokens"], + } + backend( + state, + x, + dt, + A, + B, + C, + D, + dt_bias, + z=z, + dt_softplus=True, + state_batch_indices=state_batch_indices, + dst_state_batch_indices=dst_state_batch_indices, + null_block_id=NULL_BLOCK_ID, + out=out, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, + **kwargs, + ) + return out + + +@requires_flashinfer +@pytest.mark.parametrize("metadata_max_seqlen", [1, 8]) +@pytest.mark.parametrize("batch_size", [1, 4]) +@pytest.mark.parametrize("use_z", [False, True]) +@pytest.mark.parametrize("enable_stochastic_rounding", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_checkpointing_ssu_stp_outputs_match_old_flashinfer( + batch_size: int, + metadata_max_seqlen: int, + use_z: bool, + enable_stochastic_rounding: bool, + dtype: torch.dtype, +) -> None: + """The new STP path should produce the same token outputs as old FI SSU. + + With ``checkpoint_interval=1``, checkpointing SSU still keeps one replay + token between calls, so the HBM state is not expected to equal the old + materialized state at every step. The observable decode outputs should + match, and this test exercises several consecutive STP decode iterations + so the replay tracker is used. + """ + + torch.manual_seed(0) + device = torch.device("cuda") + cache_size = 8 + nheads = 4 + head_dim = 64 + dstate = 128 + ngroups = 1 + max_window = 1 + + old_backend = _make_backend( + enable_stochastic_rounding=enable_stochastic_rounding + ) + new_backend = _make_backend( + enable_stochastic_rounding=enable_stochastic_rounding + ) + + initial_state = 0.1 * torch.randn( + cache_size, + nheads, + head_dim, + dstate, + device=device, + dtype=torch.float16, + ) + old_state = initial_state.clone() + new_state = initial_state.clone() + cache = _make_checkpointing_cache( + cache_size=cache_size, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + max_window=max_window, + device=device, + dtype=dtype, + ) + A, D, dt_bias = _make_weights( + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + device=device, + dtype=dtype, + ) + + slots = torch.arange(batch_size, device=device, dtype=torch.int32) * 2 + 1 + state_index_table = torch.stack((slots - 1, slots), dim=1) + state_batch_indices = state_index_table[:, 1:2] + cu_seqlens = torch.arange(batch_size + 1, device=device, dtype=torch.int32) + + for step in range(5): + x, dt, B, C = _make_decode_inputs( + batch_size=batch_size, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + device=device, + dtype=dtype, + seed=100 + step, + ) + z = torch.tanh(2 * x) if use_z else None + torch.manual_seed(1000 + step) + old_out = _call_backend( + old_backend, + state=old_state, + cache=None, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + z=z, + state_batch_indices=state_batch_indices, + cu_seqlens=cu_seqlens, + max_seqlen=metadata_max_seqlen, + ) + torch.manual_seed(1000 + step) + new_out = _call_backend( + new_backend, + state=new_state, + cache=cache, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + z=z, + state_batch_indices=state_batch_indices, + cu_seqlens=cu_seqlens, + max_seqlen=metadata_max_seqlen, + ) + + torch.testing.assert_close( + new_out.float(), + old_out.float(), + atol=3e-2, + rtol=3e-2, + msg=f"STP output mismatch at decode step {step}", + ) + + +@requires_flashinfer +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_checkpointing_ssu_copies_replay_state_to_destination_slot( + dtype: torch.dtype, +) -> None: + """Slot moves must preserve checkpoint state and live replay buffers.""" + + torch.manual_seed(0) + device = torch.device("cuda") + cache_size = 5 + nheads = 4 + head_dim = 64 + dstate = 128 + ngroups = 1 + max_window = 1 + + old_backend = _make_backend() + new_backend = _make_backend() + initial_state = 0.1 * torch.randn( + cache_size, + nheads, + head_dim, + dstate, + device=device, + dtype=torch.float16, + ) + old_state = initial_state.clone() + new_state = initial_state.clone() + cache = _make_checkpointing_cache( + cache_size=cache_size, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + max_window=max_window, + device=device, + dtype=dtype, + ) + A, D, dt_bias = _make_weights( + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + device=device, + dtype=dtype, + ) + src = torch.tensor([1], device=device, dtype=torch.int32) + dst = torch.tensor([2], device=device, dtype=torch.int32) + cu_seqlens = torch.tensor([0, 1], device=device, dtype=torch.int32) + + x, dt, B, C = _make_decode_inputs( + batch_size=1, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + device=device, + dtype=dtype, + seed=99, + ) + _call_backend( + old_backend, + state=old_state, + cache=None, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=src, + cu_seqlens=cu_seqlens, + ) + _call_backend( + new_backend, + state=new_state, + cache=cache, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=src, + cu_seqlens=cu_seqlens, + ) + assert int(cache["prev_num_accepted_tokens"][1].item()) == 1 + src_slot = int(src.item()) + source_snapshot = ( + new_state[src_slot].clone(), + cache["old_x"][src_slot].clone(), + cache["old_B"][src_slot].clone(), + cache["old_dt"][src_slot].clone(), + cache["old_cumAdt"][src_slot].clone(), + cache["cache_buf_idx"][src_slot].clone(), + cache["prev_num_accepted_tokens"][src_slot].clone(), + ) + + x, dt, B, C = _make_decode_inputs( + batch_size=1, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + device=device, + dtype=dtype, + seed=100, + ) + old_out = _call_backend( + old_backend, + state=old_state, + cache=None, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=src, + dst_state_batch_indices=dst, + cu_seqlens=cu_seqlens, + ) + new_out = _call_backend( + new_backend, + state=new_state, + cache=cache, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=src, + dst_state_batch_indices=dst, + cu_seqlens=cu_seqlens, + ) + torch.testing.assert_close(new_out.float(), old_out.float(), atol=3e-2, rtol=3e-2) + for actual, expected in zip( + ( + new_state[src_slot], + cache["old_x"][src_slot], + cache["old_B"][src_slot], + cache["old_dt"][src_slot], + cache["old_cumAdt"][src_slot], + cache["cache_buf_idx"][src_slot], + cache["prev_num_accepted_tokens"][src_slot], + ), + source_snapshot, + ): + torch.testing.assert_close(actual, expected) + + +@requires_flashinfer +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_checkpointing_ssu_stp_large_batch_outputs_match_old_flashinfer( + dtype: torch.dtype, +) -> None: + """Large CUDA-graph decode batches must match the old FI STP path.""" + + device = torch.device("cuda") + batch_size = 24 + cache_size = batch_size * 2 + 8 + nheads = 128 + head_dim = 64 + dstate = 128 + ngroups = 8 + max_window = 6 + generator = torch.Generator(device=device).manual_seed(123) + + old_backend = _make_backend() + new_backend = _make_backend(checkpoint_interval=max_window) + initial_state = 0.01 * torch.randn( + cache_size, + nheads, + head_dim, + dstate, + device=device, + dtype=torch.float16, + generator=generator, + ) + old_state = initial_state.clone() + new_state = initial_state.clone() + cache = _make_checkpointing_cache( + cache_size=cache_size, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + max_window=max_window, + device=device, + dtype=dtype, + ) + A, D, dt_bias = _make_weights( + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + device=device, + dtype=dtype, + ) + slots = torch.arange(batch_size, device=device, dtype=torch.int32) * 2 + 1 + state_batch_indices = slots[:, None] + cu_seqlens = torch.arange(batch_size + 1, device=device, dtype=torch.int32) + + for step in range(12): + x, dt, B, C = _make_decode_inputs( + batch_size=batch_size, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + device=device, + dtype=dtype, + seed=200 + step, + ) + old_out = _call_backend( + old_backend, + state=old_state, + cache=None, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=state_batch_indices, + cu_seqlens=cu_seqlens, + max_seqlen=8, + ) + new_out = _call_backend( + new_backend, + state=new_state, + cache=cache, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=state_batch_indices, + cu_seqlens=cu_seqlens, + max_seqlen=8, + ) + max_abs_error = (new_out.float() - old_out.float()).abs().max() + assert max_abs_error <= 6e-2, ( + f"large-batch STP output mismatch at decode step {step}: " + f"max_abs_error={max_abs_error.item()}" + ) + tracked_tokens = cache["prev_num_accepted_tokens"][slots.to(torch.long)] + assert int(tracked_tokens.max().item()) == 1 + + +@requires_flashinfer +@pytest.mark.parametrize("enable_stochastic_rounding", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_checkpointing_ssu_stp_mixed_lifecycle_batch2_outputs_match_old_flashinfer( + enable_stochastic_rounding: bool, + dtype: torch.dtype, +) -> None: + """A newly joined request must batch with a replaying request correctly.""" + + device = torch.device("cuda") + cache_size = 16 + nheads = 128 + head_dim = 64 + dstate = 128 + ngroups = 8 + max_window = 1 + row_stride = 18560 + generator = torch.Generator(device=device).manual_seed(5678) + + old_backend = _make_backend( + enable_stochastic_rounding=enable_stochastic_rounding + ) + new_backend = _make_backend( + enable_stochastic_rounding=enable_stochastic_rounding + ) + initial_state = 0.01 * torch.randn( + cache_size, + nheads, + head_dim, + dstate, + device=device, + dtype=torch.float16, + generator=generator, + ) + old_state = initial_state.clone() + new_state = initial_state.clone() + cache = _make_checkpointing_cache( + cache_size=cache_size, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + max_window=max_window, + device=device, + dtype=dtype, + ) + A, D, dt_bias = _make_weights( + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + device=device, + dtype=dtype, + ) + slot_a = torch.tensor([1], device=device, dtype=torch.int32) + slots_ab = torch.tensor([1, 7], device=device, dtype=torch.int32) + cu_a = torch.tensor([0, 1], device=device, dtype=torch.int32) + cu_ab = torch.tensor([0, 1, 2], device=device, dtype=torch.int32) + + x_a, dt_a, B_a, C_a = _make_decode_inputs( + batch_size=1, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + device=device, + dtype=dtype, + seed=401, + ) + old_out_a = _call_backend( + old_backend, + state=old_state, + cache=None, + x=x_a, + dt=dt_a, + A=A, + B=B_a, + C=C_a, + D=D, + dt_bias=dt_bias, + state_batch_indices=slot_a, + cu_seqlens=cu_a, + ) + new_out_a = _call_backend( + new_backend, + state=new_state, + cache=cache, + x=x_a, + dt=dt_a, + A=A, + B=B_a, + C=C_a, + D=D, + dt_bias=dt_bias, + state_batch_indices=slot_a, + cu_seqlens=cu_a, + ) + torch.testing.assert_close( + new_out_a.float(), old_out_a.float(), atol=6e-2, rtol=6e-2 + ) + assert int(cache["prev_num_accepted_tokens"][1].item()) == 1 + assert int(cache["prev_num_accepted_tokens"][7].item()) == 0 + + x_dense, dt_dense, B_dense, C_dense = _make_decode_inputs( + batch_size=2, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + device=device, + dtype=dtype, + seed=402, + ) + x_base = torch.empty((2, row_stride), device=device, dtype=dtype) + dt_base = torch.empty((2, row_stride), device=device, dtype=dtype) + B_base = torch.empty((2, row_stride), device=device, dtype=dtype) + C_base = torch.empty((2, row_stride), device=device, dtype=dtype) + x = torch.as_strided(x_base, x_dense.shape, (row_stride, head_dim, 1)) + dt = torch.as_strided(dt_base, dt_dense.shape, (row_stride, 1, 0)) + B = torch.as_strided(B_base, B_dense.shape, (row_stride, dstate, 1)) + C = torch.as_strided(C_base, C_dense.shape, (row_stride, dstate, 1)) + x.copy_(x_dense) + dt_base[:, :nheads].copy_(dt_dense[..., 0]) + B.copy_(B_dense) + C.copy_(C_dense) + + old_out = _call_backend( + old_backend, + state=old_state, + cache=None, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=slots_ab, + cu_seqlens=cu_ab, + ) + new_out = _call_backend( + new_backend, + state=new_state, + cache=cache, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=slots_ab, + cu_seqlens=cu_ab, + ) + max_abs_error = (new_out.float() - old_out.float()).abs().max() + assert max_abs_error <= 6e-2, ( + "mixed-lifecycle batch-2 STP output mismatch: " + f"max_abs_error={max_abs_error.item()}" + ) + + +@requires_flashinfer +@pytest.mark.parametrize("enable_stochastic_rounding", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_checkpointing_ssu_stp_nonvarlen_mixed_lifecycle_batch2_matches_old_flashinfer( + enable_stochastic_rounding: bool, + dtype: torch.dtype, +) -> None: + """Non-varlen decode must handle a new request batched with replay.""" + + device = torch.device("cuda") + cache_size = 16 + nheads = 128 + head_dim = 64 + dstate = 128 + ngroups = 8 + max_window = 1 + row_stride = 18560 + generator = torch.Generator(device=device).manual_seed(8765) + + old_backend = _make_backend( + enable_stochastic_rounding=enable_stochastic_rounding + ) + new_backend = _make_backend( + enable_stochastic_rounding=enable_stochastic_rounding + ) + initial_state = 0.01 * torch.randn( + cache_size, + nheads, + head_dim, + dstate, + device=device, + dtype=torch.float16, + generator=generator, + ) + old_state = initial_state.clone() + new_state = initial_state.clone() + cache = _make_checkpointing_cache( + cache_size=cache_size, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + max_window=max_window, + device=device, + dtype=dtype, + ) + A, D, dt_bias = _make_weights( + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + device=device, + dtype=dtype, + ) + slot_a = torch.tensor([1], device=device, dtype=torch.int32) + slots_ab = torch.tensor([1, 7], device=device, dtype=torch.int32) + + x_a, dt_a, B_a, C_a = _make_decode_inputs( + batch_size=1, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + device=device, + dtype=dtype, + seed=801, + ) + old_out_a = _call_backend( + old_backend, + state=old_state, + cache=None, + x=x_a, + dt=dt_a, + A=A, + B=B_a, + C=C_a, + D=D, + dt_bias=dt_bias, + state_batch_indices=slot_a, + cu_seqlens=None, + ) + new_out_a = _call_backend( + new_backend, + state=new_state, + cache=cache, + x=x_a, + dt=dt_a, + A=A, + B=B_a, + C=C_a, + D=D, + dt_bias=dt_bias, + state_batch_indices=slot_a, + cu_seqlens=None, + ) + torch.testing.assert_close( + new_out_a.float(), old_out_a.float(), atol=6e-2, rtol=6e-2 + ) + assert int(cache["prev_num_accepted_tokens"][1].item()) == 1 + assert int(cache["prev_num_accepted_tokens"][7].item()) == 0 + + x_dense, dt_dense, B_dense, C_dense = _make_decode_inputs( + batch_size=2, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + device=device, + dtype=dtype, + seed=802, + ) + x_base = torch.empty((2, row_stride), device=device, dtype=dtype) + dt_base = torch.empty((2, row_stride), device=device, dtype=dtype) + B_base = torch.empty((2, row_stride), device=device, dtype=dtype) + C_base = torch.empty((2, row_stride), device=device, dtype=dtype) + x = torch.as_strided(x_base, x_dense.shape, (row_stride, head_dim, 1)) + dt = torch.as_strided(dt_base, dt_dense.shape, (row_stride, 1, 0)) + B = torch.as_strided(B_base, B_dense.shape, (row_stride, dstate, 1)) + C = torch.as_strided(C_base, C_dense.shape, (row_stride, dstate, 1)) + x.copy_(x_dense) + dt_base[:, :nheads].copy_(dt_dense[..., 0]) + B.copy_(B_dense) + C.copy_(C_dense) + + old_out = _call_backend( + old_backend, + state=old_state, + cache=None, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=slots_ab, + cu_seqlens=None, + ) + new_out = _call_backend( + new_backend, + state=new_state, + cache=cache, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=slots_ab, + cu_seqlens=None, + ) + max_abs_error = (new_out.float() - old_out.float()).abs().max() + assert max_abs_error <= 6e-2, ( + "non-varlen mixed-lifecycle batch-2 STP output mismatch: " + f"max_abs_error={max_abs_error.item()}" + ) + + +@requires_flashinfer +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_checkpointing_ssu_stp_vllm_padded_cache_layout_matches_old_flashinfer( + dtype: torch.dtype, +) -> None: + """Checkpointing cache tensors may be as_strided into padded vLLM pages.""" + + device = torch.device("cuda") + batch_size = 2 + cache_size = 16 + nheads = 128 + head_dim = 64 + dstate = 128 + ngroups = 8 + max_window = 1 + page_pad = 4096 + generator = torch.Generator(device=device).manual_seed(1357) + + old_backend = _make_backend() + new_backend = _make_backend() + initial_state = 0.01 * torch.randn( + cache_size, + nheads, + head_dim, + dstate, + device=device, + dtype=torch.float16, + generator=generator, + ) + old_state = initial_state.clone() + state_slot_el = nheads * head_dim * dstate + raw_state = torch.empty( + cache_size * (state_slot_el + page_pad), + device=device, + dtype=torch.float16, + ) + new_state = torch.as_strided( + raw_state, + (cache_size, nheads, head_dim, dstate), + (state_slot_el + page_pad, head_dim * dstate, dstate, 1), + ) + new_state.copy_(initial_state) + + old_x_slot_el = max_window * nheads * head_dim + old_B_slot_el = 2 * max_window * ngroups * dstate + old_dt_slot_el = 2 * nheads * max_window + raw_old_x = torch.zeros( + cache_size * (old_x_slot_el + page_pad), + device=device, + dtype=dtype, + ) + raw_old_B = torch.zeros( + cache_size * (old_B_slot_el + page_pad), + device=device, + dtype=dtype, + ) + raw_old_dt = torch.zeros( + cache_size * (old_dt_slot_el + page_pad), + device=device, + dtype=torch.float32, + ) + raw_old_cumAdt = torch.zeros_like(raw_old_dt) + cache = { + "old_x": torch.as_strided( + raw_old_x, + (cache_size, max_window, nheads, head_dim), + (old_x_slot_el + page_pad, nheads * head_dim, head_dim, 1), + ), + "old_B": torch.as_strided( + raw_old_B, + (cache_size, 2, max_window, ngroups, dstate), + (old_B_slot_el + page_pad, max_window * ngroups * dstate, + ngroups * dstate, dstate, 1), + ), + "old_dt": torch.as_strided( + raw_old_dt, + (cache_size, 2, nheads, max_window), + (old_dt_slot_el + page_pad, nheads * max_window, max_window, 1), + ), + "old_cumAdt": torch.as_strided( + raw_old_cumAdt, + (cache_size, 2, nheads, max_window), + (old_dt_slot_el + page_pad, nheads * max_window, max_window, 1), + ), + "cache_buf_idx": torch.zeros(cache_size, device=device, dtype=torch.int32), + "prev_num_accepted_tokens": torch.zeros( + cache_size, device=device, dtype=torch.int32 + ), + } + A, D, dt_bias = _make_weights( + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + device=device, + dtype=dtype, + ) + slots = torch.tensor([1, 7], device=device, dtype=torch.int32) + + for step in range(4): + dst_slots = slots + 1 + x, dt, B, C = _make_decode_inputs( + batch_size=batch_size, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + device=device, + dtype=dtype, + seed=901 + step, + ) + old_out = _call_backend( + old_backend, + state=old_state, + cache=None, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=slots, + dst_state_batch_indices=dst_slots, + cu_seqlens=None, + ) + new_out = _call_backend( + new_backend, + state=new_state, + cache=cache, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=slots, + dst_state_batch_indices=dst_slots, + cu_seqlens=None, + ) + max_abs_error = (new_out.float() - old_out.float()).abs().max() + assert max_abs_error <= 6e-2, ( + f"vLLM padded cache layout mismatch at decode step {step}: " + f"max_abs_error={max_abs_error.item()}" + ) + slots = dst_slots + + +@requires_flashinfer +@pytest.mark.parametrize("enable_stochastic_rounding", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_checkpointing_ssu_stp_fixed_padded_batch_outputs_match_old_flashinfer( + enable_stochastic_rounding: bool, + dtype: torch.dtype, +) -> None: + """Fixed CUDA-graph-shaped decode batches must ignore padded slots.""" + + device = torch.device("cuda") + active_batch_size = 2 + padded_batch_size = 8 + cache_size = 32 + nheads = 128 + head_dim = 64 + dstate = 128 + ngroups = 8 + max_window = 1 + generator = torch.Generator(device=device).manual_seed(1234) + + old_backend = _make_backend( + enable_stochastic_rounding=enable_stochastic_rounding + ) + new_backend = _make_backend( + enable_stochastic_rounding=enable_stochastic_rounding + ) + initial_state = 0.01 * torch.randn( + cache_size, + nheads, + head_dim, + dstate, + device=device, + dtype=torch.float16, + generator=generator, + ) + old_state = initial_state.clone() + new_state = initial_state.clone() + cache = _make_checkpointing_cache( + cache_size=cache_size, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + max_window=max_window, + device=device, + dtype=dtype, + ) + A, D, dt_bias = _make_weights( + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + device=device, + dtype=dtype, + ) + active_slots = torch.tensor([1, 5], device=device, dtype=torch.int32) + padded_slots = torch.full( + (padded_batch_size,), NULL_BLOCK_ID, device=device, dtype=torch.int32 + ) + padded_slots[:active_batch_size] = active_slots + + for step in range(12): + x, dt, B, C = _make_decode_inputs( + batch_size=padded_batch_size, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + device=device, + dtype=dtype, + seed=300 + step, + ) + torch.manual_seed(5000 + step) + old_out = _call_backend( + old_backend, + state=old_state, + cache=None, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=padded_slots, + cu_seqlens=None, + ) + torch.manual_seed(5000 + step) + new_out = _call_backend( + new_backend, + state=new_state, + cache=cache, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=padded_slots, + cu_seqlens=None, + ) + max_abs_error = ( + new_out[:active_batch_size].float() + - old_out[:active_batch_size].float() + ).abs().max() + assert max_abs_error <= 6e-2, ( + f"fixed padded-batch STP output mismatch at decode step {step}: " + f"max_abs_error={max_abs_error.item()}" + ) + + +@requires_flashinfer +@pytest.mark.parametrize("enable_stochastic_rounding", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16]) +def test_checkpointing_ssu_stp_overlapping_cache_slot_copy_matches_old_flashinfer( + enable_stochastic_rounding: bool, + dtype: torch.dtype, +) -> None: + """Batched cache-slot moves must be safe when source/destination overlap.""" + + device = torch.device("cuda") + cache_size = 16 + nheads = 128 + head_dim = 64 + dstate = 128 + ngroups = 8 + max_window = 1 + row_stride = 18560 + generator = torch.Generator(device=device).manual_seed(2468) + + old_backend = _make_backend( + enable_stochastic_rounding=enable_stochastic_rounding + ) + new_backend = _make_backend( + enable_stochastic_rounding=enable_stochastic_rounding + ) + initial_state = 0.01 * torch.randn( + cache_size, + nheads, + head_dim, + dstate, + device=device, + dtype=torch.float16, + generator=generator, + ) + old_state = initial_state.clone() + new_state = initial_state.clone() + cache = _make_checkpointing_cache( + cache_size=cache_size, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + max_window=max_window, + device=device, + dtype=dtype, + ) + A = -torch.rand(nheads, device=device, dtype=torch.float32, generator=generator) + A = A[:, None, None].expand(nheads, head_dim, dstate) + D = (0.1 * torch.randn( + nheads, device=device, dtype=dtype, generator=generator + ))[:, None].expand(nheads, head_dim) + dt_bias = (0.1 * torch.randn( + nheads, device=device, dtype=dtype, generator=generator + ))[:, None].expand(nheads, head_dim) + + def make_strided_inputs(seed: int) -> tuple[ + torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor + ]: + x_dense, dt_dense, B_dense, C_dense = _make_decode_inputs( + batch_size=2, + nheads=nheads, + head_dim=head_dim, + dstate=dstate, + ngroups=ngroups, + device=device, + dtype=dtype, + seed=seed, + ) + x_base = torch.empty((2, row_stride), device=device, dtype=dtype) + dt_base = torch.empty((2, row_stride), device=device, dtype=dtype) + B_base = torch.empty((2, row_stride), device=device, dtype=dtype) + C_base = torch.empty((2, row_stride), device=device, dtype=dtype) + x = torch.as_strided(x_base, x_dense.shape, (row_stride, head_dim, 1)) + dt = torch.as_strided(dt_base, dt_dense.shape, (row_stride, 1, 0)) + B = torch.as_strided(B_base, B_dense.shape, (row_stride, dstate, 1)) + C = torch.as_strided(C_base, C_dense.shape, (row_stride, dstate, 1)) + x.copy_(x_dense) + dt_base[:, :nheads].copy_(dt_dense[..., 0]) + B.copy_(B_dense) + C.copy_(C_dense) + return x, dt, B, C + + init_src = torch.tensor([0, 4], device=device, dtype=torch.int32) + init_dst = torch.tensor([1, 2], device=device, dtype=torch.int32) + x, dt, B, C = make_strided_inputs(701) + torch.manual_seed(9001) + old_init = _call_backend( + old_backend, + state=old_state, + cache=None, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=init_src, + dst_state_batch_indices=init_dst, + cu_seqlens=None, + ) + torch.manual_seed(9001) + new_init = _call_backend( + new_backend, + state=new_state, + cache=cache, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=init_src, + dst_state_batch_indices=init_dst, + cu_seqlens=None, + ) + torch.testing.assert_close( + new_init.float(), old_init.float(), atol=6e-2, rtol=6e-2 + ) + + # Slot 2 is both the destination of request 0 and the source of request 1. + # A batched copy must preserve slot 2's old contents until request 1 reads it. + overlap_src = torch.tensor([1, 2], device=device, dtype=torch.int32) + overlap_dst = torch.tensor([2, 3], device=device, dtype=torch.int32) + x, dt, B, C = make_strided_inputs(702) + torch.manual_seed(9002) + old_overlap = _call_backend( + old_backend, + state=old_state, + cache=None, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=overlap_src, + dst_state_batch_indices=overlap_dst, + cu_seqlens=None, + ) + torch.manual_seed(9002) + new_overlap = _call_backend( + new_backend, + state=new_state, + cache=cache, + x=x, + dt=dt, + A=A, + B=B, + C=C, + D=D, + dt_bias=dt_bias, + state_batch_indices=overlap_src, + dst_state_batch_indices=overlap_dst, + cu_seqlens=None, + ) + max_abs_error = (new_overlap.float() - old_overlap.float()).abs().max() + assert max_abs_error <= 6e-2, ( + "overlapping cache-slot STP output mismatch: " + f"max_abs_error={max_abs_error.item()}" + ) diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index c5d0661476e3..c740a8809a6c 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -2,6 +2,13 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from unittest.mock import MagicMock, patch +import torch + +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFuncCalculator, + get_stable_temporal_copy_spec, + get_temporal_copy_spec, +) from vllm.v1.core.sched.output import CachedRequestData, SchedulerOutput from vllm.v1.worker.mamba_utils import preprocess_mamba @@ -67,3 +74,29 @@ def test_resumed_req_ids_cleared_from_mamba_state_idx(): ) assert mamba_state_idx == {"keep": 99} + + +def test_checkpointing_temporal_copy_uses_stable_slot(): + state = torch.arange(10 * 2, dtype=torch.float32).view(10, 2) + block_ids = [3, 4, 5, 6] + + old_fi_spec = get_temporal_copy_spec( + state, block_ids, cur_block_idx=1, num_accepted_tokens=3 + ) + ckpt_spec = get_stable_temporal_copy_spec( + state, block_ids, cur_block_idx=1, num_accepted_tokens=3 + ) + + assert old_fi_spec.start_addr == state[block_ids[3]].data_ptr() + assert ckpt_spec.start_addr == state[block_ids[1]].data_ptr() + assert old_fi_spec.num_elements == ckpt_spec.num_elements == 2 + + +def test_mamba2_checkpointing_copy_funcs_keep_ssm_replay_state_stable(): + funcs = MambaStateCopyFuncCalculator.mamba2_checkpointing_state_copy_func() + + assert len(funcs) == 8 + assert funcs[1:] == (get_stable_temporal_copy_spec,) * 7 + assert MambaStateCopyFuncCalculator.mamba2_state_copy_func()[1] is ( + get_temporal_copy_spec + ) diff --git a/tests/v1/worker/test_utils.py b/tests/v1/worker/test_utils.py index 016b8aa5a635..07d32a042658 100644 --- a/tests/v1/worker/test_utils.py +++ b/tests/v1/worker/test_utils.py @@ -1,9 +1,15 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest import torch -from vllm.v1.worker.utils import bind_kv_cache +from vllm.model_executor.layers.mamba.mamba_utils import MambaStateShapeCalculator +from vllm.v1.kv_cache_interface import MambaSpec +from vllm.v1.worker.utils import AttentionGroup, KVBlockZeroer, bind_kv_cache def test_bind_kv_cache(default_vllm_config): @@ -90,3 +96,94 @@ def test_bind_kv_cache_draft_model(default_vllm_config): assert runner_kv_caches[1] is kv_cache["draft_model.layers.0.attn"] assert runner_kv_caches[2] is kv_cache["model.layers.1.attn"] assert runner_kv_caches[3] is kv_cache["draft_model.layers.1.attn"] + + +def test_kv_block_zeroer_zeros_mamba_cache_blocks(): + num_blocks = 4 + state_storage = torch.full((num_blocks * 16,), 1.0) + state = torch.as_strided( + state_storage, + size=(num_blocks, 2, 3), + stride=(16, 3, 1), + ) + tracker = torch.full((num_blocks,), 9, dtype=torch.int32) + spec = MambaSpec( + block_size=1, + shapes=((2, 3), ()), + dtypes=(torch.float32, torch.int32), + ) + group = AttentionGroup( + backend=Mock(), + layer_names=["layer.0"], + kv_cache_spec=spec, + kv_cache_group_id=0, + ) + zeroer = KVBlockZeroer(torch.device("cpu"), pin_memory=False) + zeroer.init_meta( + [group], + [1], + "auto", + set(), + {"layer.0": SimpleNamespace(kv_cache=(state, tracker))}, + ) + + zeroer.zero_block_ids([1, 3]) + + assert torch.all(state[1] == 0) + assert torch.all(state[3] == 0) + assert torch.all(tracker[torch.tensor([1, 3])] == 0) + assert torch.all(state[0] == 1) + assert torch.all(state[2] == 1) + assert torch.all(tracker[torch.tensor([0, 2])] == 9) + assert torch.all(state_storage[22:32] == 1) + + +@pytest.mark.parametrize("checkpoint_interval", [1, 6]) +def test_mamba2_state_shape_honors_checkpoint_interval(checkpoint_interval): + shapes = MambaStateShapeCalculator.mamba2_state_shape( + tp_world_size=1, + intermediate_size=8192, + n_groups=8, + num_heads=128, + head_dim=64, + state_size=128, + conv_kernel=4, + checkpoint_interval=checkpoint_interval, + ) + + old_x_shape = shapes[2] + old_B_shape = shapes[3] + old_dt_shape = shapes[4] + old_cumAdt_shape = shapes[5] + + assert old_x_shape[0] == checkpoint_interval + assert old_B_shape[1] == checkpoint_interval + assert old_dt_shape[-1] == checkpoint_interval + assert old_cumAdt_shape[-1] == checkpoint_interval + + +@pytest.mark.parametrize("checkpoint_interval", [1, 6]) +def test_nemotron_h_state_shape_honors_checkpoint_interval(checkpoint_interval): + from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM + + vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace(tensor_parallel_size=1), + model_config=SimpleNamespace( + hf_config=SimpleNamespace( + mamba_num_heads=128, + mamba_head_dim=64, + n_groups=8, + ssm_state_size=128, + conv_kernel=4, + ), + ), + num_speculative_tokens=0, + mamba_config=SimpleNamespace(checkpoint_interval=checkpoint_interval), + ) + + shapes = NemotronHForCausalLM.get_mamba_state_shape_from_config(vllm_config) + + assert shapes[2][0] == checkpoint_interval + assert shapes[3][1] == checkpoint_interval + assert shapes[4][-1] == checkpoint_interval + assert shapes[5][-1] == checkpoint_interval diff --git a/vllm/config/mamba.py b/vllm/config/mamba.py index 996478c36760..f9487f25d9bf 100644 --- a/vllm/config/mamba.py +++ b/vllm/config/mamba.py @@ -45,6 +45,10 @@ class MambaConfig: generation. 0 uses the Triton default. Higher values improve randomness quality at the cost of compute.""" + checkpoint_interval: int = 1 + """Number of decode tokens to keep in the FlashInfer checkpointing SSU + replay window before materializing a new SSM state checkpoint.""" + @field_validator("backend", mode="before") @classmethod def validate_backend_before(cls, value: Any) -> Any: @@ -53,6 +57,15 @@ def validate_backend_before(cls, value: Any) -> Any: return MambaBackendEnum[value.upper()] return value + @field_validator("checkpoint_interval") + @classmethod + def validate_checkpoint_interval(cls, value: int) -> int: + if value < 1: + raise ValueError("Mamba checkpoint interval must be >= 1") + if value > 16: + raise ValueError("FlashInfer checkpointing SSU supports interval <= 16") + return value + def __post_init__(self): if self.enable_stochastic_rounding: from vllm.platforms import current_platform diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 49189862eb7e..17c23cf88afe 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -676,6 +676,7 @@ class EngineArgs: MambaConfig.enable_stochastic_rounding ) mamba_cache_philox_rounds: int = MambaConfig.stochastic_rounding_philox_rounds + mamba_checkpoint_interval: int = MambaConfig.checkpoint_interval additional_config: dict[str, Any] = get_field(VllmConfig, "additional_config") @@ -907,6 +908,10 @@ def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser: "--mamba-cache-philox-rounds", **mamba_kwargs["stochastic_rounding_philox_rounds"], ) + mamba_group.add_argument( + "--mamba-checkpoint-interval", + **mamba_kwargs["checkpoint_interval"], + ) # Structured outputs arguments structured_outputs_kwargs = get_kwargs(StructuredOutputsConfig) @@ -2106,6 +2111,8 @@ def create_engine_config( mamba_config.stochastic_rounding_philox_rounds = ( self.mamba_cache_philox_rounds ) + if self.mamba_checkpoint_interval: + mamba_config.checkpoint_interval = self.mamba_checkpoint_interval # Kernel config overrides kernel_config = copy.deepcopy(self.kernel_config) diff --git a/vllm/model_executor/layers/mamba/mamba_mixer2.py b/vllm/model_executor/layers/mamba/mamba_mixer2.py index a6524961ea92..081dfc508d3d 100644 --- a/vllm/model_executor/layers/mamba/mamba_mixer2.py +++ b/vllm/model_executor/layers/mamba/mamba_mixer2.py @@ -5,7 +5,12 @@ import torch from torch import nn -from vllm.config import CacheConfig, ModelConfig, get_current_vllm_config +from vllm.config import ( + CUDAGraphMode, + CacheConfig, + ModelConfig, + get_current_vllm_config, +) from vllm.distributed import ( divide, get_tensor_model_parallel_rank, @@ -477,6 +482,7 @@ def __init__( ) self._ssd_kernels_warmed_up = False + self._warmup_ssd_kernels_in_profile = True # - get hidden_states, B and C after depthwise convolution. self.split_hidden_states_B_C_fn = lambda hidden_states_B_C: torch.split( @@ -491,17 +497,28 @@ def __init__( vllm_config = get_current_vllm_config() compilation_config = vllm_config.compilation_config + self._warmup_ssd_kernels_in_profile = ( + compilation_config.cudagraph_mode == CUDAGraphMode.NONE + ) if prefix in compilation_config.static_forward_context: raise ValueError(f"Duplicate layer name: {prefix}") compilation_config.static_forward_context[prefix] = self - # The tuple is (conv_state, ssm_state) - self.kv_cache = (torch.tensor([]), torch.tensor([])) + # The tuple is (conv_state, ssm_state, old_x, old_B, old_dt, + # old_cumAdt, cache_buf_idx, prev_num_accepted_tokens). + self.kv_cache = tuple(torch.tensor([]) for _ in range(8)) + self._checkpointing_old_x = torch.tensor([]) + self._checkpointing_old_B = torch.tensor([]) + self._checkpointing_old_dt = torch.tensor([]) + self._checkpointing_old_cumAdt = torch.tensor([]) + self._checkpointing_cache_buf_idx = torch.tensor([]) + self._checkpointing_prev_num_accepted_tokens = torch.tensor([]) self.model_config = model_config self.cache_config = cache_config self.prefix = prefix self.num_spec = vllm_config.num_speculative_tokens + self.mamba_checkpoint_interval = vllm_config.mamba_config.checkpoint_interval if self.num_spec > 0: self.register_buffer( "_decode_state_offsets", @@ -527,6 +544,25 @@ def __init__( # Check if running on Blackwell (SM100+) for kernel tuning self.is_blackwell = current_platform.is_device_capability_family(100) + def _get_contiguous_checkpointing_buffer( + self, + source: torch.Tensor, + attr_name: str, + *, + copy_source: bool = False, + ) -> torch.Tensor: + buffer = getattr(self, attr_name) + if ( + buffer.shape != source.shape + or buffer.device != source.device + or buffer.dtype != source.dtype + ): + buffer = torch.zeros_like(source, memory_format=torch.contiguous_format) + setattr(self, attr_name, buffer) + if copy_source: + buffer.copy_(source) + return buffer + def forward( self, hidden_states: torch.Tensor, @@ -591,7 +627,7 @@ def _warmup_ssd_kernels(self, projected_states: torch.Tensor) -> None: # Triton's autotuner includes tensor dtypes in its cache key, # so state_dtype must match what real inference uses. - _, ssm_state_dtype = self.get_state_dtype() + ssm_state_dtype = self.get_state_dtype()[1] # SSD kernel autotune keys depend on dtype and head dimensions, # not on sequence length or batch size, so a single shape suffices. @@ -702,6 +738,55 @@ def conv_ssm_forward( else self.kv_cache[0].transpose(-1, -2) ) ssm_state = self.kv_cache[1] + old_x = self.kv_cache[2] + old_B = self.kv_cache[3] + old_dt = self.kv_cache[4] + old_cumAdt = self.kv_cache[5] + cache_buf_idx = self.kv_cache[6] + prev_num_accepted_tokens = self.kv_cache[7] + old_x_src = old_x + old_B_src = old_B + old_dt_src = old_dt + old_cumAdt_src = old_cumAdt + cache_buf_idx_src = cache_buf_idx + prev_num_accepted_tokens_src = prev_num_accepted_tokens + if not old_x.is_contiguous(): + old_x = self._get_contiguous_checkpointing_buffer( + old_x, "_checkpointing_old_x", copy_source=True + ) + if not old_B.is_contiguous(): + old_B = self._get_contiguous_checkpointing_buffer( + old_B, "_checkpointing_old_B", copy_source=True + ) + if not old_dt.is_contiguous(): + old_dt = self._get_contiguous_checkpointing_buffer( + old_dt, "_checkpointing_old_dt", copy_source=True + ) + if not old_cumAdt.is_contiguous(): + old_cumAdt = self._get_contiguous_checkpointing_buffer( + old_cumAdt, "_checkpointing_old_cumAdt", copy_source=True + ) + if not cache_buf_idx.is_contiguous(): + cache_buf_idx = self._get_contiguous_checkpointing_buffer( + cache_buf_idx, + "_checkpointing_cache_buf_idx", + copy_source=True, + ) + if not prev_num_accepted_tokens.is_contiguous(): + prev_num_accepted_tokens = self._get_contiguous_checkpointing_buffer( + prev_num_accepted_tokens, + "_checkpointing_prev_num_accepted_tokens", + copy_source=True, + ) + + def reset_checkpointing_cache(block_indices: torch.Tensor) -> None: + old_x[block_indices] = 0 + old_B[block_indices] = 0 + old_dt[block_indices] = 0 + old_cumAdt[block_indices] = 0 + cache_buf_idx[block_indices] = 0 + prev_num_accepted_tokens[block_indices] = 0 + has_initial_states_p = attn_metadata.has_initial_states_p prep_initial_states = attn_metadata.prep_initial_states chunk_size = attn_metadata.chunk_size @@ -719,7 +804,8 @@ def conv_ssm_forward( if attn_metadata is None: # V1 profile run -- warm up SSD kernels so that autotuning # completes before SSM cache allocation. - self._warmup_ssd_kernels(projected_states) + if self._warmup_ssd_kernels_in_profile: + self._warmup_ssd_kernels(projected_states) hidden_states_B_C = ( hidden_states_B_C.transpose(0, 1).clone().transpose(0, 1) ).contiguous() @@ -939,14 +1025,15 @@ def conv_ssm_forward( # Write the states ssm_state[cache_blocks_to_fill] = from_where + reset_checkpointing_cache(cache_blocks_to_fill) # For all seqs, store the last state (note: might be partial): assert state_indices_tensor_p is not None - ssm_state[ - state_indices_tensor_p.gather( - 1, block_idx_last_scheduled_token_p.unsqueeze(1) - ).squeeze(1) - ] = varlen_states[last_chunk_indices_p] + last_state_indices = state_indices_tensor_p.gather( + 1, block_idx_last_scheduled_token_p.unsqueeze(1) + ).squeeze(1) + ssm_state[last_state_indices] = varlen_states[last_chunk_indices_p] + reset_checkpointing_cache(last_state_indices) else: # update ssm states @@ -954,6 +1041,7 @@ def conv_ssm_forward( # tensor assert state_indices_tensor_p is not None ssm_state[state_indices_tensor_p] = varlen_states + reset_checkpointing_cache(state_indices_tensor_p) # Process decode requests if has_decode: @@ -1042,10 +1130,29 @@ def conv_ssm_forward( out=preallocated_ssm_out_d.view(num_decode_tokens, -1, self.head_dim), num_accepted_tokens=num_accepted_tokens, cu_seqlens=query_start_loc_d, + max_seqlen=state_indices_tensor_d.size(-1), is_blackwell=self.is_blackwell, + old_x=old_x, + old_B=old_B, + old_dt=old_dt, + old_cumAdt=old_cumAdt, + cache_buf_idx=cache_buf_idx, + prev_num_accepted_tokens=prev_num_accepted_tokens, ) - - def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]: + if old_x is not old_x_src: + old_x_src.copy_(old_x) + if old_B is not old_B_src: + old_B_src.copy_(old_B) + if old_dt is not old_dt_src: + old_dt_src.copy_(old_dt) + if old_cumAdt is not old_cumAdt_src: + old_cumAdt_src.copy_(old_cumAdt) + if cache_buf_idx is not cache_buf_idx_src: + cache_buf_idx_src.copy_(cache_buf_idx) + if prev_num_accepted_tokens is not prev_num_accepted_tokens_src: + prev_num_accepted_tokens_src.copy_(prev_num_accepted_tokens) + + def get_state_dtype(self) -> tuple[torch.dtype, ...]: assert self.model_config is not None assert self.cache_config is not None return MambaStateDtypeCalculator.mamba2_state_dtype( @@ -1054,7 +1161,7 @@ def get_state_dtype(self) -> tuple[torch.dtype, torch.dtype]: self.cache_config.mamba_ssm_cache_dtype, ) - def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: + def get_state_shape(self) -> tuple[tuple[int, ...], ...]: return MambaStateShapeCalculator.mamba2_state_shape( intermediate_size=self.intermediate_size, tp_world_size=get_tensor_model_parallel_world_size(), @@ -1064,6 +1171,7 @@ def get_state_shape(self) -> tuple[tuple[int, ...], tuple[int, ...]]: state_size=self.ssm_state_size, conv_kernel=self.conv_kernel_size, num_spec=self.num_spec, + checkpoint_interval=self.mamba_checkpoint_interval, ) @property @@ -1079,7 +1187,11 @@ def mamba_mixer2( layer_name = _resolve_layer_name(layer_name) forward_context: ForwardContext = get_forward_context() self = forward_context.no_compile_layers[layer_name] - self.conv_ssm_forward(projected_states=projected_states, output=output) + ssm_output = self.conv_ssm_forward( + projected_states=projected_states, output=output + ) + if ssm_output is not None: + output.copy_(ssm_output) def mamba_mixer2_fake( diff --git a/vllm/model_executor/layers/mamba/mamba_utils.py b/vllm/model_executor/layers/mamba/mamba_utils.py index c1fd81e40e34..b997f970ec2c 100644 --- a/vllm/model_executor/layers/mamba/mamba_utils.py +++ b/vllm/model_executor/layers/mamba/mamba_utils.py @@ -76,9 +76,19 @@ def mamba2_state_dtype( mamba_cache_dtype: MambaDType, mamba_ssm_cache_dtype: MambaDType, ) -> tuple[torch.dtype, ...]: - return cls._mamba_state_dtype( + conv_state_dtype, temporal_state_dtype = cls._mamba_state_dtype( model_dtype, mamba_cache_dtype, mamba_ssm_cache_dtype ) + return ( + conv_state_dtype, + temporal_state_dtype, + torch.bfloat16, + torch.bfloat16, + torch.float32, + torch.float32, + torch.int32, + torch.int32, + ) @classmethod def _mamba_state_dtype( @@ -169,7 +179,8 @@ def mamba2_state_shape( state_size: int, conv_kernel: int, num_spec: int = 0, - ) -> tuple[tuple[int, int], tuple[int, int, int]]: + checkpoint_interval: int = 1, + ) -> tuple[tuple[int, ...], ...]: # if n_groups is not divisible by world_size, need to extend the shards # to ensure all groups needed by a head is sharded along with it n_groups = n_groups + cls.extra_groups_for_head_shards(n_groups, tp_world_size) @@ -184,7 +195,24 @@ def mamba2_state_shape( # - they are typically small # e.g., (h_heads, head_dim, state_size) = (128, 64, 128) temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, state_size) - return conv_state_shape, temporal_state_shape + nheads = divide(num_heads, tp_world_size) + ngroups = max(1, divide(n_groups, tp_world_size)) + old_x_shape = (checkpoint_interval, nheads, head_dim) + old_B_shape = (2, checkpoint_interval, ngroups, state_size) + old_dt_shape = (2, nheads, checkpoint_interval) + old_cumAdt_shape = (2, nheads, checkpoint_interval) + cache_buf_idx_shape = () + prev_num_accepted_tokens_shape = () + return ( + conv_state_shape, + temporal_state_shape, + old_x_shape, + old_B_shape, + old_dt_shape, + old_cumAdt_shape, + cache_buf_idx_shape, + prev_num_accepted_tokens_shape, + ) @classmethod def short_conv_state_shape( @@ -342,6 +370,21 @@ def get_temporal_copy_spec( ) +def get_stable_temporal_copy_spec( + state: torch.Tensor, + block_ids: list[int], + cur_block_idx: int, + num_accepted_tokens: int, +) -> MambaCopySpec: + """Return a copy spec for state kept on the stable per-request block.""" + del num_accepted_tokens + src_block_id = block_ids[cur_block_idx] + src_state = state[src_block_id] + return MambaCopySpec( + start_addr=src_state.data_ptr(), num_elements=src_state.numel() + ) + + class MambaStateCopyFuncCalculator: @classmethod def linear_attention_state_copy_func(cls): @@ -353,7 +396,31 @@ def mamba1_state_copy_func(cls): @classmethod def mamba2_state_copy_func(cls): - return get_conv_copy_spec, get_temporal_copy_spec + return ( + get_conv_copy_spec, + get_temporal_copy_spec, + get_stable_temporal_copy_spec, + get_stable_temporal_copy_spec, + get_stable_temporal_copy_spec, + get_stable_temporal_copy_spec, + get_stable_temporal_copy_spec, + get_stable_temporal_copy_spec, + ) + + @classmethod + def mamba2_checkpointing_state_copy_func(cls): + # Convolution state is still materialized by the conv update kernel. + # The SSM checkpoint state and replay tensors stay on the stable slot. + return ( + get_conv_copy_spec, + get_stable_temporal_copy_spec, + get_stable_temporal_copy_spec, + get_stable_temporal_copy_spec, + get_stable_temporal_copy_spec, + get_stable_temporal_copy_spec, + get_stable_temporal_copy_spec, + get_stable_temporal_copy_spec, + ) @classmethod def short_conv_state_copy_func(cls): diff --git a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py index 92258ef204bd..a729c33aa379 100644 --- a/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py +++ b/vllm/model_executor/layers/mamba/ops/ssu_dispatch.py @@ -14,6 +14,7 @@ from vllm.config.mamba import MambaBackendEnum, MambaConfig from vllm.logger import init_logger +from vllm.triton_utils import tl, triton from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum from vllm.v1.attention.backends.utils import NULL_BLOCK_ID from vllm.v1.kv_cache_interface import KVCacheConfig, MambaSpec @@ -21,6 +22,135 @@ logger = init_logger(__name__) +@triton.jit +def _update_checkpointing_trackers_kernel( + cache_buf_idx, + prev_num_accepted_tokens, + state_batch_indices, + cu_seqlens, + fixed_seq_len: tl.constexpr, + max_window: tl.constexpr, + pad_slot_id: tl.constexpr, + n_slots: tl.constexpr, + HAS_CU_SEQLENS: tl.constexpr, + BLOCK: tl.constexpr, +) -> None: + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < n_slots + slots = tl.load(state_batch_indices + offsets, mask=mask, other=pad_slot_id) + valid = mask & (slots != pad_slot_id) + if HAS_CU_SEQLENS: + seq_lens = tl.load(cu_seqlens + offsets + 1, mask=mask, other=0) - tl.load( + cu_seqlens + offsets, mask=mask, other=0 + ) + else: + seq_lens = tl.full((BLOCK,), fixed_seq_len, tl.int32) + prev = tl.load(prev_num_accepted_tokens + slots, mask=valid, other=0) + must_checkpoint = prev + seq_lens > max_window + old_buf = tl.load(cache_buf_idx + slots, mask=valid, other=0) + new_buf = tl.where(must_checkpoint, 1 - old_buf, old_buf) + new_prev = tl.minimum( + tl.where(must_checkpoint, seq_lens, prev + seq_lens), max_window + ) + tl.store(cache_buf_idx + slots, new_buf, mask=valid) + tl.store(prev_num_accepted_tokens + slots, new_prev, mask=valid) + + +@triton.jit +def _reset_checkpointing_trackers_kernel( + cache_buf_idx, + prev_num_accepted_tokens, + state_batch_indices, + pad_slot_id: tl.constexpr, + n_slots: tl.constexpr, + BLOCK: tl.constexpr, +) -> None: + offsets = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < n_slots + slots = tl.load(state_batch_indices + offsets, mask=mask, other=pad_slot_id) + valid = mask & (slots != pad_slot_id) + tl.store(cache_buf_idx + slots, 0, mask=valid) + tl.store(prev_num_accepted_tokens + slots, 0, mask=valid) + + +@triton.jit +def _copy_checkpointing_slots_kernel( + tensor, + src_indices, + dst_indices, + slot_size: tl.constexpr, + slot_stride: tl.constexpr, + pad_slot_id: tl.constexpr, + BLOCK: tl.constexpr, +) -> None: + slot = tl.program_id(0) + offsets = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < slot_size + src = tl.load(src_indices + slot) + dst = tl.load(dst_indices + slot) + valid = (src != pad_slot_id) & (dst != pad_slot_id) & (src != dst) + values = tl.load(tensor + src * slot_stride + offsets, mask=mask & valid) + tl.store(tensor + dst * slot_stride + offsets, values, mask=mask & valid) + + +@triton.jit +def _gather_checkpointing_slots_kernel( + tensor, + scratch, + src_indices, + dst_indices, + slot_size: tl.constexpr, + tensor_slot_stride: tl.constexpr, + scratch_slot_stride: tl.constexpr, + pad_slot_id: tl.constexpr, + BLOCK: tl.constexpr, +) -> None: + slot = tl.program_id(0) + offsets = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < slot_size + src = tl.load(src_indices + slot) + dst = tl.load(dst_indices + slot) + valid = (src != pad_slot_id) & (dst != pad_slot_id) & (src != dst) + values = tl.load( + tensor + src * tensor_slot_stride + offsets, + mask=mask & valid, + ) + tl.store( + scratch + slot * scratch_slot_stride + offsets, + values, + mask=mask & valid, + ) + + +@triton.jit +def _scatter_checkpointing_slots_kernel( + tensor, + scratch, + src_indices, + dst_indices, + slot_size: tl.constexpr, + tensor_slot_stride: tl.constexpr, + scratch_slot_stride: tl.constexpr, + pad_slot_id: tl.constexpr, + BLOCK: tl.constexpr, +) -> None: + slot = tl.program_id(0) + offsets = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK) + mask = offsets < slot_size + src = tl.load(src_indices + slot) + dst = tl.load(dst_indices + slot) + valid = (src != pad_slot_id) & (dst != pad_slot_id) & (src != dst) + values = tl.load( + scratch + slot * scratch_slot_stride + offsets, + mask=mask & valid, + ) + tl.store( + tensor + dst * tensor_slot_stride + offsets, + values, + mask=mask & valid, + ) + + class MambaSSUBackend(ABC): """Abstract base class for Mamba SSU backends.""" @@ -50,7 +180,14 @@ def __call__( out: torch.Tensor | None = None, num_accepted_tokens: torch.Tensor | None = None, cu_seqlens: torch.Tensor | None = None, + max_seqlen: int | None = None, is_blackwell: bool = False, + old_x: torch.Tensor | None = None, + old_B: torch.Tensor | None = None, + old_dt: torch.Tensor | None = None, + old_cumAdt: torch.Tensor | None = None, + cache_buf_idx: torch.Tensor | None = None, + prev_num_accepted_tokens: torch.Tensor | None = None, ) -> None: ... @@ -87,7 +224,14 @@ def __call__( out: torch.Tensor | None = None, num_accepted_tokens: torch.Tensor | None = None, cu_seqlens: torch.Tensor | None = None, + max_seqlen: int | None = None, is_blackwell: bool = False, + old_x: torch.Tensor | None = None, + old_B: torch.Tensor | None = None, + old_dt: torch.Tensor | None = None, + old_cumAdt: torch.Tensor | None = None, + cache_buf_idx: torch.Tensor | None = None, + prev_num_accepted_tokens: torch.Tensor | None = None, ) -> None: self._kernel( state, @@ -117,15 +261,17 @@ class FlashInferSSUBackend(MambaSSUBackend): def __init__(self, mamba_config: MambaConfig): super().__init__(mamba_config) + self._copy_scratch: dict[tuple[object, ...], torch.Tensor] = {} try: + from flashinfer.mamba import checkpointing_ssu as _fi_checkpointing_ssu from flashinfer.mamba import selective_state_update as _fi_ssu except ImportError as e: raise ImportError( "FlashInfer is required for the flashinfer Mamba SSU backend. " - "Please install flashinfer (>= 0.6.4): " - "pip install flashinfer-python" + "Please install a FlashInfer build with Mamba checkpointing SSU." ) from e self._kernel = _fi_ssu + self._checkpointing_kernel = _fi_checkpointing_ssu @property def name(self) -> str: @@ -149,14 +295,189 @@ def __call__( out: torch.Tensor | None = None, num_accepted_tokens: torch.Tensor | None = None, cu_seqlens: torch.Tensor | None = None, + max_seqlen: int | None = None, is_blackwell: bool = False, + old_x: torch.Tensor | None = None, + old_B: torch.Tensor | None = None, + old_dt: torch.Tensor | None = None, + old_cumAdt: torch.Tensor | None = None, + cache_buf_idx: torch.Tensor | None = None, + prev_num_accepted_tokens: torch.Tensor | None = None, ) -> None: rand_seed = ( - torch.randint(0, 2**32, (1,), device=state.device) + torch.randint(0, 2**32, (1,), dtype=torch.int64, device=state.device) if self._mamba_config.enable_stochastic_rounding else None ) + checkpointing_args = ( + old_x, + old_B, + old_dt, + old_cumAdt, + cache_buf_idx, + prev_num_accepted_tokens, + ) + state_indices = self._checkpointing_state_indices(state_batch_indices) + simple_decode = state_indices is not None and x.size(0) == state_indices.numel() + non_spec_varlen = state_indices is not None and cu_seqlens is not None + num_accepted_tokens_for_kernel = ( + None if simple_decode or non_spec_varlen else num_accepted_tokens + ) + can_checkpoint = ( + state_indices is not None + and (num_accepted_tokens is None or simple_decode or non_spec_varlen) + and all(arg is not None for arg in checkpointing_args) + and state.dtype in (torch.float16, torch.bfloat16, torch.float32) + ) + if can_checkpoint: + assert old_x is not None + assert old_B is not None + assert old_dt is not None + assert old_cumAdt is not None + assert cache_buf_idx is not None + assert prev_num_accepted_tokens is not None + assert state_indices is not None + kernel_state_indices = state_indices + self._ensure_copy_scratch( + ( + state, + old_x, + old_B, + old_dt, + old_cumAdt, + cache_buf_idx, + prev_num_accepted_tokens, + ), + kernel_state_indices.numel(), + ) + dst_indices = self._checkpointing_state_indices(dst_state_batch_indices) + if ( + dst_indices is not None + and dst_indices.numel() == state_indices.numel() + ): + copied_checkpointing_slots = True + self._copy_checkpointing_slots( + ( + state, + old_x, + old_B, + old_dt, + old_cumAdt, + cache_buf_idx, + prev_num_accepted_tokens, + ), + state_indices, + dst_indices, + null_block_id, + ) + kernel_state_indices = dst_indices + else: + copied_checkpointing_slots = False + ckpt_cu_seqlens = ( + None + if simple_decode + else self._checkpointing_cu_seqlens( + cu_seqlens, x, kernel_state_indices, max_seqlen + ) + ) + checkpoint_window = old_x.size(1) + kernel_old_x = old_x + kernel_old_B = old_B + kernel_old_dt = old_dt + kernel_old_cumAdt = old_cumAdt + if simple_decode and checkpoint_window > 1: + checkpoint_window = 1 + kernel_old_x = old_x[:, :1] + kernel_old_B = old_B[:, :, :1] + kernel_old_dt = old_dt[..., :1] + kernel_old_cumAdt = old_cumAdt[..., :1] + if ( + ckpt_cu_seqlens is not None + and max_seqlen is not None + and max_seqlen > checkpoint_window + and x.size(0) > kernel_state_indices.numel() + and not torch.cuda.is_current_stream_capturing() + ): + self._run_varlen_checkpointing_chunks( + state, + kernel_old_x, + kernel_old_B, + kernel_old_dt, + kernel_old_cumAdt, + cache_buf_idx, + prev_num_accepted_tokens, + x, + dt, + A, + B, + C, + out, + D, + z, + dt_bias, + dt_softplus, + kernel_state_indices, + ckpt_cu_seqlens, + max_seqlen, + checkpoint_window, + null_block_id, + rand_seed, + ) + return + x_ckpt, dt_ckpt, B_ckpt, C_ckpt, z_ckpt, out_ckpt, ckpt_max_seqlen = ( + self._reshape_checkpointing_inputs( + x, + dt, + B, + C, + z, + out, + kernel_state_indices, + ckpt_cu_seqlens, + max_seqlen, + checkpoint_window, + ) + ) + kernel_max_seqlen = ( + ckpt_max_seqlen if ckpt_cu_seqlens is not None else None + ) + self._run_checkpointing_kernel( + state, + kernel_old_x, + kernel_old_B, + kernel_old_dt, + kernel_old_cumAdt, + cache_buf_idx, + prev_num_accepted_tokens, + x_ckpt, + dt_ckpt, + A, + B_ckpt, + C_ckpt, + out_ckpt, + D, + z_ckpt, + dt_bias, + dt_softplus, + kernel_state_indices, + null_block_id, + rand_seed, + ckpt_cu_seqlens, + kernel_max_seqlen, + force_headwise_checkpointing=copied_checkpointing_slots, + ) + self._update_checkpointing_trackers( + cache_buf_idx, + prev_num_accepted_tokens, + kernel_state_indices, + ckpt_cu_seqlens, + ckpt_max_seqlen, + checkpoint_window, + null_block_id, + ) + return + self._kernel( state, x, @@ -171,7 +492,7 @@ def __call__( state_batch_indices=state_batch_indices, dst_state_batch_indices=dst_state_batch_indices, cu_seqlens=cu_seqlens, - num_accepted_tokens=num_accepted_tokens, + num_accepted_tokens=num_accepted_tokens_for_kernel, cache_steps=state_batch_indices.size(-1) if cu_seqlens is not None and state_batch_indices is not None else 0, @@ -179,7 +500,586 @@ def __call__( out=out, rand_seed=rand_seed, philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds or 10, + algorithm="simple", + ) + if ( + num_accepted_tokens is None + and cache_buf_idx is not None + and prev_num_accepted_tokens is not None + and dst_state_batch_indices is not None + and not self._same_state_indices( + state_batch_indices, dst_state_batch_indices + ) + ): + dst_indices = self._checkpointing_state_indices(dst_state_batch_indices) + if dst_indices is not None: + self._reset_checkpointing_trackers( + cache_buf_idx, + prev_num_accepted_tokens, + dst_indices, + null_block_id, + ) + + def _run_checkpointing_kernel( + self, + state: torch.Tensor, + old_x: torch.Tensor, + old_B: torch.Tensor, + old_dt: torch.Tensor, + old_cumAdt: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + x_ckpt: torch.Tensor, + dt_ckpt: torch.Tensor, + A: torch.Tensor, + B_ckpt: torch.Tensor, + C_ckpt: torch.Tensor, + out_ckpt: torch.Tensor, + D: torch.Tensor | None, + z_ckpt: torch.Tensor | None, + dt_bias: torch.Tensor | None, + dt_softplus: bool, + kernel_state_indices: torch.Tensor, + null_block_id: int, + rand_seed: torch.Tensor | None, + ckpt_cu_seqlens: torch.Tensor | None, + kernel_max_seqlen: int | None, + *, + force_headwise_checkpointing: bool = False, + ) -> None: + n_groups = old_B.size(-2) + n_heads = old_x.size(2) + if n_groups > 1 and n_heads % n_groups == 0: + heads_per_group = n_heads // n_groups + if force_headwise_checkpointing: + for group_idx in range(n_groups): + group_head_start = group_idx * heads_per_group + group_head_end = group_head_start + heads_per_group + group_slice = slice(group_idx, group_idx + 1) + old_B_snapshot = old_B[ + :, :, :, group_slice, : + ].contiguous() + old_B_group = torch.empty_like(old_B_snapshot) + for head_start in range(group_head_start, group_head_end): + head_end = head_start + 1 + head_slice = slice(head_start, head_end) + state_group = state[:, head_slice].contiguous() + old_x_group = old_x[:, :, head_slice].contiguous() + old_B_group.copy_(old_B_snapshot) + old_dt_group = old_dt[ + :, :, head_slice, : + ].contiguous() + old_cumAdt_group = old_cumAdt[ + :, :, head_slice, : + ].contiguous() + out_group = torch.empty_like( + x_ckpt[:, :, head_slice, :] + ) + self._checkpointing_kernel( + state_group, + old_x_group, + old_B_group, + old_dt_group, + old_cumAdt_group, + cache_buf_idx, + prev_num_accepted_tokens, + x_ckpt[:, :, head_slice, :].contiguous(), + dt_ckpt[:, :, head_slice, :], + A[head_slice], + B_ckpt[:, :, group_slice, :].contiguous(), + C_ckpt[:, :, group_slice, :].contiguous(), + out_group, + D=D[head_slice] if D is not None else None, + z=( + z_ckpt[:, :, head_slice, :].contiguous() + if z_ckpt is not None + else None + ), + dt_bias=( + dt_bias[head_slice] + if dt_bias is not None + else None + ), + dt_softplus=dt_softplus, + state_batch_indices=kernel_state_indices, + pad_slot_id=null_block_id, + rand_seed=rand_seed, + philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds + or 10, + cu_seqlens=ckpt_cu_seqlens, + max_seqlen=kernel_max_seqlen, + ) + state[:, head_slice].copy_(state_group) + old_x[:, :, head_slice].copy_(old_x_group) + old_dt[:, :, head_slice, :].copy_(old_dt_group) + old_cumAdt[:, :, head_slice, :].copy_(old_cumAdt_group) + out_ckpt[:, :, head_slice, :].copy_(out_group) + old_B[:, :, :, group_slice, :].copy_(old_B_group) + return + + for group_idx in range(n_groups): + group_head_start = group_idx * heads_per_group + group_head_end = group_head_start + heads_per_group + head_slice = slice(group_head_start, group_head_end) + group_slice = slice(group_idx, group_idx + 1) + state_group = state[:, head_slice].contiguous() + old_x_group = old_x[:, :, head_slice].contiguous() + old_B_group = old_B[:, :, :, group_slice, :].contiguous() + old_dt_group = old_dt[:, :, head_slice, :].contiguous() + old_cumAdt_group = old_cumAdt[ + :, :, head_slice, : + ].contiguous() + out_group = torch.empty_like(x_ckpt[:, :, head_slice, :]) + self._checkpointing_kernel( + state_group, + old_x_group, + old_B_group, + old_dt_group, + old_cumAdt_group, + cache_buf_idx, + prev_num_accepted_tokens, + x_ckpt[:, :, head_slice, :].contiguous(), + dt_ckpt[:, :, head_slice, :], + A[head_slice], + B_ckpt[:, :, group_slice, :].contiguous(), + C_ckpt[:, :, group_slice, :].contiguous(), + out_group, + D=D[head_slice] if D is not None else None, + z=( + z_ckpt[:, :, head_slice, :].contiguous() + if z_ckpt is not None + else None + ), + dt_bias=dt_bias[head_slice] if dt_bias is not None else None, + dt_softplus=dt_softplus, + state_batch_indices=kernel_state_indices, + pad_slot_id=null_block_id, + rand_seed=rand_seed, + philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds + or 10, + cu_seqlens=ckpt_cu_seqlens, + max_seqlen=kernel_max_seqlen, + ) + state[:, head_slice].copy_(state_group) + old_x[:, :, head_slice].copy_(old_x_group) + old_B[:, :, :, group_slice, :].copy_(old_B_group) + old_dt[:, :, head_slice, :].copy_(old_dt_group) + old_cumAdt[:, :, head_slice, :].copy_(old_cumAdt_group) + out_ckpt[:, :, head_slice, :].copy_(out_group) + return + + self._checkpointing_kernel( + state, + old_x, + old_B, + old_dt, + old_cumAdt, + cache_buf_idx, + prev_num_accepted_tokens, + x_ckpt, + dt_ckpt, + A, + B_ckpt, + C_ckpt, + out_ckpt, + D=D, + z=z_ckpt, + dt_bias=dt_bias, + dt_softplus=dt_softplus, + state_batch_indices=kernel_state_indices, + pad_slot_id=null_block_id, + rand_seed=rand_seed, + philox_rounds=self._mamba_config.stochastic_rounding_philox_rounds + or 10, + cu_seqlens=ckpt_cu_seqlens, + max_seqlen=kernel_max_seqlen, + ) + + @staticmethod + def _pack_varlen_chunk( + source: torch.Tensor, + chunk_ends: list[tuple[int, int]], + total_tokens: int, + ) -> torch.Tensor: + packed = torch.empty( + (total_tokens, *source.shape[1:]), + dtype=source.dtype, + device=source.device, + ) + offset = 0 + for start, end in chunk_ends: + length = end - start + packed[offset : offset + length].copy_(source[start:end]) + offset += length + return packed + + @staticmethod + def _pack_optional_varlen_chunk( + source: torch.Tensor | None, + chunk_ends: list[tuple[int, int]], + total_tokens: int, + ) -> torch.Tensor | None: + if source is None: + return None + return FlashInferSSUBackend._pack_varlen_chunk( + source, chunk_ends, total_tokens + ) + + def _run_varlen_checkpointing_chunks( + self, + state: torch.Tensor, + old_x: torch.Tensor, + old_B: torch.Tensor, + old_dt: torch.Tensor, + old_cumAdt: torch.Tensor, + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + x: torch.Tensor, + dt: torch.Tensor, + A: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + out: torch.Tensor | None, + D: torch.Tensor | None, + z: torch.Tensor | None, + dt_bias: torch.Tensor | None, + dt_softplus: bool, + kernel_state_indices: torch.Tensor, + cu_seqlens: torch.Tensor, + max_seqlen: int, + checkpoint_window: int, + null_block_id: int, + rand_seed: torch.Tensor | None, + ) -> None: + assert out is not None + boundaries = [int(v) for v in cu_seqlens.detach().cpu().tolist()] + batch = len(boundaries) - 1 + for chunk_start in range(0, max_seqlen, checkpoint_window): + chunk_ends: list[tuple[int, int]] = [] + chunk_lengths = [] + chunk_positions = [] + for seq_idx in range(batch): + seq_start = boundaries[seq_idx] + seq_end = boundaries[seq_idx + 1] + start = seq_start + chunk_start + if start >= seq_end: + continue + end = min(seq_end, start + checkpoint_window) + chunk_positions.append(seq_idx) + chunk_ends.append((start, end)) + chunk_lengths.append(end - start) + chunk_parts = len(chunk_ends) + if not chunk_parts: + continue + chunk_total = sum(chunk_lengths) + if chunk_parts == 1: + start, end = chunk_ends[0] + x_chunk = x[start:end] + dt_chunk = dt[start:end] + B_chunk = B[start:end] + C_chunk = C[start:end] + z_chunk = z[start:end] if z is not None else None + else: + x_chunk = self._pack_varlen_chunk(x, chunk_ends, chunk_total) + dt_chunk = self._pack_varlen_chunk(dt, chunk_ends, chunk_total) + B_chunk = self._pack_varlen_chunk(B, chunk_ends, chunk_total) + C_chunk = self._pack_varlen_chunk(C, chunk_ends, chunk_total) + z_chunk = self._pack_optional_varlen_chunk( + z, chunk_ends, chunk_total + ) + position_tensor = torch.tensor( + chunk_positions, + dtype=torch.long, + device=kernel_state_indices.device, + ) + chunk_state_indices = kernel_state_indices.index_select( + 0, position_tensor + ).to(torch.int32).contiguous() + chunk_cu = [0] + for length in chunk_lengths: + chunk_cu.append(chunk_cu[-1] + length) + chunk_cu_seqlens = torch.tensor( + chunk_cu, + dtype=torch.int32, + device=cu_seqlens.device, + ) + out_chunk = torch.empty_like(x_chunk) + chunk_max_seqlen = max(chunk_lengths) + ( + x_ckpt, + dt_ckpt, + B_ckpt, + C_ckpt, + z_ckpt, + out_ckpt, + ckpt_max_seqlen, + ) = self._reshape_checkpointing_inputs( + x_chunk, + dt_chunk, + B_chunk, + C_chunk, + z_chunk, + out_chunk, + chunk_state_indices, + chunk_cu_seqlens, + chunk_max_seqlen, + checkpoint_window, + ) + self._run_checkpointing_kernel( + state, + old_x, + old_B, + old_dt, + old_cumAdt, + cache_buf_idx, + prev_num_accepted_tokens, + x_ckpt, + dt_ckpt, + A, + B_ckpt, + C_ckpt, + out_ckpt, + D, + z_ckpt, + dt_bias, + dt_softplus, + chunk_state_indices, + null_block_id, + rand_seed, + chunk_cu_seqlens, + ckpt_max_seqlen, + ) + self._update_checkpointing_trackers( + cache_buf_idx, + prev_num_accepted_tokens, + chunk_state_indices, + chunk_cu_seqlens, + ckpt_max_seqlen, + checkpoint_window, + null_block_id, + ) + offset = 0 + for length, (start, end) in zip(chunk_lengths, chunk_ends): + out[start:end].copy_(out_chunk[offset : offset + length]) + offset += length + + @staticmethod + def _same_state_indices( + state_batch_indices: torch.Tensor | None, + dst_state_batch_indices: torch.Tensor | None, + ) -> bool: + if dst_state_batch_indices is None: + return True + if state_batch_indices is None: + return False + if state_batch_indices.shape != dst_state_batch_indices.shape: + return False + if ( + state_batch_indices.data_ptr() == dst_state_batch_indices.data_ptr() + and state_batch_indices.stride() == dst_state_batch_indices.stride() + ): + return True + return False + + @staticmethod + def _checkpointing_state_indices( + state_batch_indices: torch.Tensor | None, + ) -> torch.Tensor | None: + if state_batch_indices is None: + return None + if state_batch_indices.dim() == 1: + return state_batch_indices.to(torch.int32).contiguous() + if state_batch_indices.dim() == 2 and state_batch_indices.size(1) == 1: + return state_batch_indices[:, 0].to(torch.int32).contiguous() + return None + + @staticmethod + def _checkpointing_cu_seqlens( + cu_seqlens: torch.Tensor | None, + x: torch.Tensor, + state_batch_indices: torch.Tensor, + max_seqlen: int | None, + ) -> torch.Tensor | None: + del max_seqlen + return cu_seqlens + + @staticmethod + def _reshape_checkpointing_inputs( + x: torch.Tensor, + dt: torch.Tensor, + B: torch.Tensor, + C: torch.Tensor, + z: torch.Tensor | None, + out: torch.Tensor | None, + state_batch_indices: torch.Tensor, + cu_seqlens: torch.Tensor | None, + max_seqlen: int | None, + max_window: int, + ) -> tuple[ + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor, + torch.Tensor | None, + torch.Tensor, + int, + ]: + assert out is not None + if cu_seqlens is not None: + x = x.contiguous() + B = B.contiguous() + C = C.contiguous() + if z is not None: + z = z.contiguous() + return ( + x.unsqueeze(0), + dt.unsqueeze(0), + B.unsqueeze(0), + C.unsqueeze(0), + z.unsqueeze(0) if z is not None else None, + out.unsqueeze(0), + min(max_seqlen or max_window, max_window), + ) + batch = state_batch_indices.numel() + x = x.contiguous() + B = B.contiguous() + C = C.contiguous() + if z is not None: + z = z.contiguous() + tokens_per_batch = x.shape[0] // batch + z_ckpt = None + if z is not None: + z_ckpt = z.view(batch, tokens_per_batch, *z.shape[1:]) + return ( + x.view(batch, tokens_per_batch, *x.shape[1:]), + dt.view(batch, tokens_per_batch, *dt.shape[1:]), + B.view(batch, tokens_per_batch, *B.shape[1:]), + C.view(batch, tokens_per_batch, *C.shape[1:]), + z_ckpt, + out.view(batch, tokens_per_batch, *out.shape[1:]), + tokens_per_batch, + ) + + @staticmethod + def _update_checkpointing_trackers( + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + state_batch_indices: torch.Tensor, + cu_seqlens: torch.Tensor | None, + max_seqlen: int, + max_window: int, + pad_slot_id: int, + ) -> None: + block = 128 + n_slots = state_batch_indices.numel() + _update_checkpointing_trackers_kernel[(triton.cdiv(n_slots, block),)]( + cache_buf_idx, + prev_num_accepted_tokens, + state_batch_indices, + cu_seqlens, + max_seqlen, + max_window, + pad_slot_id, + n_slots, + cu_seqlens is not None, + BLOCK=block, + ) + + @staticmethod + def _reset_checkpointing_trackers( + cache_buf_idx: torch.Tensor, + prev_num_accepted_tokens: torch.Tensor, + state_batch_indices: torch.Tensor, + pad_slot_id: int, + ) -> None: + block = 128 + n_slots = state_batch_indices.numel() + _reset_checkpointing_trackers_kernel[(triton.cdiv(n_slots, block),)]( + cache_buf_idx, + prev_num_accepted_tokens, + state_batch_indices, + pad_slot_id, + n_slots, + BLOCK=block, + ) + + def _copy_checkpointing_slots( + self, + tensors: tuple[torch.Tensor, ...], + src_indices: torch.Tensor, + dst_indices: torch.Tensor, + pad_slot_id: int, + ) -> None: + block = 256 + n_slots = src_indices.numel() + for tensor in tensors: + slot_size = tensor[0].numel() + slot_stride = tensor.stride(0) + scratch = self._get_copy_scratch(tensor, n_slots) + scratch_stride = scratch.stride(0) + _gather_checkpointing_slots_kernel[ + (n_slots, triton.cdiv(slot_size, block)) + ]( + tensor, + scratch, + src_indices, + dst_indices, + slot_size, + slot_stride, + scratch_stride, + pad_slot_id, + BLOCK=block, + ) + _scatter_checkpointing_slots_kernel[ + (n_slots, triton.cdiv(slot_size, block)) + ]( + tensor, + scratch, + src_indices, + dst_indices, + slot_size, + slot_stride, + scratch_stride, + pad_slot_id, + BLOCK=block, + ) + + def _ensure_copy_scratch( + self, + tensors: tuple[torch.Tensor, ...], + n_slots: int, + ) -> None: + for tensor in tensors: + self._get_copy_scratch(tensor, n_slots) + + def _get_copy_scratch( + self, + tensor: torch.Tensor, + n_slots: int, + ) -> torch.Tensor: + key = ( + tensor.device.type, + tensor.device.index, + tensor.dtype, + tuple(tensor.shape[1:]), + n_slots, + ) + scratch = self._copy_scratch.get(key) + if scratch is not None: + return scratch + if tensor.is_cuda and torch.cuda.is_current_stream_capturing(): + raise RuntimeError( + "Checkpointing slot-copy scratch is unavailable during CUDA " + "graph capture. Warm up this tensor shape and n_slots before " + "capture." + ) + scratch = torch.empty( + (n_slots, *tensor.shape[1:]), + dtype=tensor.dtype, + device=tensor.device, ) + self._copy_scratch[key] = scratch + return scratch _BACKEND_REGISTRY: dict[MambaBackendEnum, type[MambaSSUBackend]] = { @@ -251,7 +1151,14 @@ def selective_state_update( out: torch.Tensor | None = None, num_accepted_tokens: torch.Tensor | None = None, cu_seqlens: torch.Tensor | None = None, + max_seqlen: int | None = None, is_blackwell: bool = False, + old_x: torch.Tensor | None = None, + old_B: torch.Tensor | None = None, + old_dt: torch.Tensor | None = None, + old_cumAdt: torch.Tensor | None = None, + cache_buf_idx: torch.Tensor | None = None, + prev_num_accepted_tokens: torch.Tensor | None = None, ) -> None: """Unified dispatch for Mamba selective state update. @@ -274,5 +1181,12 @@ def selective_state_update( out=out, num_accepted_tokens=num_accepted_tokens, cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, is_blackwell=is_blackwell, + old_x=old_x, + old_B=old_B, + old_dt=old_dt, + old_cumAdt=old_cumAdt, + cache_buf_idx=cache_buf_idx, + prev_num_accepted_tokens=prev_num_accepted_tokens, ) diff --git a/vllm/model_executor/models/bamba.py b/vllm/model_executor/models/bamba.py index d220b22ddae7..125f429159c1 100644 --- a/vllm/model_executor/models/bamba.py +++ b/vllm/model_executor/models/bamba.py @@ -455,6 +455,7 @@ def get_mamba_state_shape_from_config( head_dim=hf_config.mamba_d_head, state_size=hf_config.mamba_d_state, conv_kernel=hf_config.mamba_d_conv, + checkpoint_interval=vllm_config.mamba_config.checkpoint_interval, ) @classmethod diff --git a/vllm/model_executor/models/falcon_h1.py b/vllm/model_executor/models/falcon_h1.py index b837dc010dac..b4b8551e3fae 100644 --- a/vllm/model_executor/models/falcon_h1.py +++ b/vllm/model_executor/models/falcon_h1.py @@ -615,6 +615,7 @@ def get_mamba_state_shape_from_config( head_dim=hf_config.mamba_d_head, state_size=hf_config.mamba_d_state, conv_kernel=hf_config.mamba_d_conv, + checkpoint_interval=vllm_config.mamba_config.checkpoint_interval, ) @classmethod diff --git a/vllm/model_executor/models/granitemoehybrid.py b/vllm/model_executor/models/granitemoehybrid.py index 1ab069e3ba38..ca8b4708fd33 100644 --- a/vllm/model_executor/models/granitemoehybrid.py +++ b/vllm/model_executor/models/granitemoehybrid.py @@ -641,6 +641,7 @@ def get_mamba_state_shape_from_config( head_dim=hf_config.mamba_d_head, state_size=hf_config.mamba_d_state, conv_kernel=hf_config.mamba_d_conv, + checkpoint_interval=vllm_config.mamba_config.checkpoint_interval, ) @classmethod diff --git a/vllm/model_executor/models/mamba2.py b/vllm/model_executor/models/mamba2.py index deb20852a26a..909aab40a39d 100644 --- a/vllm/model_executor/models/mamba2.py +++ b/vllm/model_executor/models/mamba2.py @@ -229,6 +229,7 @@ def get_mamba_state_shape_from_config( state_size=hf_config.state_size, conv_kernel=hf_config.conv_kernel, num_spec=vllm_config.num_speculative_tokens, + checkpoint_interval=vllm_config.mamba_config.checkpoint_interval, ) @classmethod diff --git a/vllm/model_executor/models/nemotron_h.py b/vllm/model_executor/models/nemotron_h.py index 537e19afbcaf..a2960eeeb1b3 100644 --- a/vllm/model_executor/models/nemotron_h.py +++ b/vllm/model_executor/models/nemotron_h.py @@ -836,6 +836,7 @@ def get_mamba_state_shape_from_config( state_size=hf_config.ssm_state_size, conv_kernel=hf_config.conv_kernel, num_spec=vllm_config.num_speculative_tokens, + checkpoint_interval=vllm_config.mamba_config.checkpoint_interval, ) @classmethod diff --git a/vllm/model_executor/models/zamba2.py b/vllm/model_executor/models/zamba2.py index b4d844ba6d76..f3ca74cb8b8f 100644 --- a/vllm/model_executor/models/zamba2.py +++ b/vllm/model_executor/models/zamba2.py @@ -891,6 +891,7 @@ def get_mamba_state_shape_from_config( head_dim=hf_config.mamba_headdim, state_size=hf_config.mamba_d_state, conv_kernel=hf_config.mamba_d_conv, + checkpoint_interval=vllm_config.mamba_config.checkpoint_interval, ) @classmethod diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index 529bbc2d4a1e..936f4992dbae 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -37,6 +37,7 @@ update_config, ) from vllm.config.cache import CacheConfig +from vllm.config.mamba import MambaBackendEnum from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer from vllm.distributed.eplb.eplb_state import EplbState from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group @@ -60,6 +61,10 @@ from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( RoutedExpertsCapturer, ) +from vllm.model_executor.layers.mamba.mamba_utils import ( + MambaStateCopyFunc, + MambaStateCopyFuncCalculator, +) from vllm.model_executor.layers.mamba.ops.ssu_dispatch import ( initialize_mamba_ssu_backend, ) @@ -992,12 +997,23 @@ def _make_buffer( with_numpy=numpy, ) + def _get_mamba_state_copy_funcs(self) -> tuple[MambaStateCopyFunc, ...]: + copy_funcs = self.model.get_mamba_state_copy_func() + if ( + self.vllm_config.mamba_config.backend == MambaBackendEnum.FLASHINFER + and len(copy_funcs) == 8 + ): + return ( + MambaStateCopyFuncCalculator.mamba2_checkpointing_state_copy_func() + ) + return copy_funcs + def _get_mamba_copy_bufs(self) -> mamba_utils.MambaCopyBuffers: if self._mamba_copy_bufs is None: self._mamba_copy_bufs = mamba_utils.MambaCopyBuffers.create( self.max_num_reqs, self.kv_cache_config, - self.model.get_mamba_state_copy_func(), + self._get_mamba_state_copy_funcs(), self._make_buffer, ) return self._mamba_copy_bufs @@ -1524,7 +1540,7 @@ def _update_states_after_model_execute( self.compilation_config.static_forward_context if is_align else None ), mamba_state_copy_funcs=( - self.model.get_mamba_state_copy_func() if is_align else None + self._get_mamba_state_copy_funcs() if is_align else None ), copy_bufs=self._get_mamba_copy_bufs() if is_align else None, ) @@ -4109,7 +4125,7 @@ def execute_model( self.input_batch, self.requests, self.compilation_config.static_forward_context, - self.model.get_mamba_state_copy_func(), + self._get_mamba_state_copy_funcs(), self._get_mamba_copy_bufs(), ) # preprocess_mamba resets num_accepted_tokens_cpu to 1 diff --git a/vllm/v1/worker/utils.py b/vllm/v1/worker/utils.py index b693adbf2771..5c4ac9cd8326 100644 --- a/vllm/v1/worker/utils.py +++ b/vllm/v1/worker/utils.py @@ -89,6 +89,7 @@ def __init__(self, device: torch.device, pin_memory: bool): self.device = device self.pin_memory = pin_memory self._meta: tuple[torch.Tensor, int, int, int] | None = None + self._mamba_tensors: list[torch.Tensor] = [] self._id_cap: int = 0 self._ids_pinned: torch.Tensor | None = None self._ids_gpu: torch.Tensor | None = None @@ -112,18 +113,38 @@ def init_meta( PAGE_SIZE_EL accounts for this ratio so that ``block_id * PAGE_SIZE_EL`` lands at the correct offset. - Only AttentionSpec layers are processed; Mamba layers are skipped. + AttentionSpec layers are zeroed with a Triton kernel. MambaSpec + layers are tracked as state tensors and zeroed by block index. """ seen_ptrs: set[int] = set() seg_addrs: list[int] = [] page_size_el: int | None = None + self._mamba_tensors = [] for group in attn_groups_iter: spec = group.kv_cache_spec - if not isinstance(spec, FullAttentionSpec): - continue if group.kv_cache_group_id >= len(kernel_block_sizes): continue + + if isinstance(spec, MambaSpec): + for layer_name in group.layer_names: + if layer_name in runner_only_attn_layers: + continue + kv = static_forward_context[layer_name].kv_cache + if not isinstance(kv, (list, tuple)): + continue + for tensor in kv: + if not isinstance(tensor, torch.Tensor): + continue + dp = tensor.data_ptr() + if dp in seen_ptrs: + continue + seen_ptrs.add(dp) + self._mamba_tensors.append(tensor) + continue + + if not isinstance(spec, FullAttentionSpec): + continue kernel_bs = kernel_block_sizes[group.kv_cache_group_id] ratio = spec.block_size // kernel_bs block_dim = group.backend.get_kv_cache_block_dim( @@ -167,11 +188,20 @@ def init_meta( off_bytes = sum(i * s for i, s in zip(outer, outer_strides)) seg_addrs.append(dp + off_bytes) - if not seg_addrs or page_size_el is None: + if seg_addrs and page_size_el is not None: + blk_size = min(largest_power_of_2_divisor(page_size_el), 1024) + self._meta = ( + torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device), + page_size_el, + blk_size, + len(seg_addrs), + ) + else: self._meta = None + + if self._meta is None and not self._mamba_tensors: return - blk_size = min(largest_power_of_2_divisor(page_size_el), 1024) self._id_cap = 8192 self._ids_pinned = torch.empty( self._id_cap, @@ -179,18 +209,11 @@ def init_meta( pin_memory=self.pin_memory, ) self._ids_gpu = torch.empty(self._id_cap, dtype=torch.int64, device=self.device) - self._meta = ( - torch.tensor(seg_addrs, dtype=torch.uint64, device=self.device), - page_size_el, - blk_size, - len(seg_addrs), - ) def zero_block_ids(self, block_ids: list[int]) -> None: """Zero the KV cache memory for the given block IDs.""" - if not block_ids or self._meta is None: + if not block_ids or (self._meta is None and not self._mamba_tensors): return - seg_addrs, page_size_el, blk_size, n_segs = self._meta n_blocks = len(block_ids) if n_blocks > self._id_cap: self._id_cap = n_blocks * 2 @@ -206,15 +229,21 @@ def zero_block_ids(self, block_ids: list[int]) -> None: self._ids_pinned[:n_blocks].numpy()[:] = block_ids idx = self._ids_gpu[:n_blocks] idx.copy_(self._ids_pinned[:n_blocks], non_blocking=True) - grid = (n_blocks * n_segs * (page_size_el // blk_size),) - _zero_kv_blocks_kernel[grid]( - seg_addrs, - idx, - n_blocks, - N_SEGS=n_segs, - PAGE_SIZE_EL=page_size_el, - BLOCK_SIZE=blk_size, - ) + + if self._meta is not None: + seg_addrs, page_size_el, blk_size, n_segs = self._meta + grid = (n_blocks * n_segs * (page_size_el // blk_size),) + _zero_kv_blocks_kernel[grid]( + seg_addrs, + idx, + n_blocks, + N_SEGS=n_segs, + PAGE_SIZE_EL=page_size_el, + BLOCK_SIZE=blk_size, + ) + + for tensor in self._mamba_tensors: + tensor.index_fill_(0, idx, 0) @dataclass