diff --git a/tests/kernels/mamba/test_causal_conv1d.py b/tests/kernels/mamba/test_causal_conv1d.py index ddd532ad13..06031842bd 100644 --- a/tests/kernels/mamba/test_causal_conv1d.py +++ b/tests/kernels/mamba/test_causal_conv1d.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterator + import pytest import torch import torch.nn.functional as F @@ -358,3 +360,113 @@ def test_causal_conv1d_varlen( ) unpadded_out = out[:, : out_ref_tensor.shape[-1]] assert torch.allclose(unpadded_out, out_ref_tensor, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("itype", [torch.bfloat16]) +@pytest.mark.parametrize("silu_activation", [True]) +@pytest.mark.parametrize("has_bias", [True]) +@pytest.mark.parametrize("width", [4]) +@pytest.mark.parametrize("dim", [64, 2048]) +@pytest.mark.parametrize("q_lens_cfg", [[3, 1, 0, 4, 2], [1, 5, 5, 1], [2, 0, 0, 3]]) +@pytest.mark.skipif( + current_platform.is_cpu(), + reason="CPU kernel for causal_conv1d_update ignores num_accepted_tokens.", +) +def test_causal_conv1d_update_ragged_varlen( + q_lens_cfg, dim, width, has_bias, silu_activation, itype +): + mk_inputs = lambda: _create_ragged_varlen_inputs( + q_lens_cfg, dim, width, has_bias, silu_activation, itype + ) + inputs, state = mk_inputs() + actual = causal_conv1d_update(conv_state=state, **inputs) + + ref_inputs, ref_state = mk_inputs() + ref_outs = _create_ragged_varlen_ref_outs(conv_state=ref_state, **ref_inputs) + + for i, ref_out in enumerate(ref_outs): + start = int(inputs["query_start_loc"][i]) + end = int(inputs["query_start_loc"][i + 1]) + entry = inputs["conv_state_indices"][i : i + 1] + assert torch.allclose(actual[start:end], ref_out, rtol=1e-2, atol=5e-2) + assert torch.equal(state[entry], ref_state[entry]) + + +def _create_ragged_varlen_inputs( + q_lens_cfg, dim, width, has_bias, silu_activation, itype +): + device = DEVICE + set_random_seed(0) + + batch_size = len(q_lens_cfg) + max_query_len = 5 + assert max(q_lens_cfg) <= max_query_len + state_len = width - 1 + max_query_len + total_entries = 10 * batch_size + + q_lens = torch.tensor(q_lens_cfg, dtype=torch.int32, device=device) + query_start_loc = torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=device), + torch.cumsum(q_lens, dim=0, dtype=torch.int32), + ] + ) + total_tokens = int(q_lens.sum()) + + x = torch.randn(total_tokens, dim, device=device, dtype=itype) + conv_state = torch.randn( + total_entries, state_len, dim, device=device, dtype=itype + ).transpose(1, 2) + weight = torch.randn(dim, width, device=device, dtype=itype) + bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None + activation = None if not silu_activation else "silu" + + conv_state_indices = ( + torch.randperm(total_entries - 1, device=device)[:batch_size] + 1 + ).to(torch.int32) + num_accepted_tokens = torch.randint( + 1, max_query_len + 1, (batch_size,), dtype=torch.int32, device=device + ) + inputs = dict( + x=x, + weight=weight, + bias=bias, + activation=activation, + conv_state_indices=conv_state_indices, + num_accepted_tokens=num_accepted_tokens, + query_start_loc=query_start_loc, + max_query_len=max_query_len, + ) + return inputs, conv_state + + +def _create_ragged_varlen_ref_outs( + x, + conv_state, + weight, + bias, + activation, + conv_state_indices, + num_accepted_tokens, + query_start_loc, + max_query_len, +) -> Iterator[torch.Tensor]: + for seq_idx in range(len(conv_state_indices)): + start = int(query_start_loc[seq_idx]) + end = int(query_start_loc[seq_idx + 1]) + entry = conv_state_indices[seq_idx : seq_idx + 1] + if start == end: + yield x[start:end] + continue + out_ref = causal_conv1d_update( + x[start:end], + conv_state, + weight, + bias, + activation=activation, + conv_state_indices=entry, + num_accepted_tokens=num_accepted_tokens[seq_idx : seq_idx + 1], + query_start_loc=query_start_loc.new_tensor([0, end - start]), + max_query_len=max_query_len, + ) + yield out_ref diff --git a/vllm/v1/attention/backends/mamba2_attn.py b/vllm/v1/attention/backends/mamba2_attn.py index fc81cf680e..1b9e573182 100644 --- a/vllm/v1/attention/backends/mamba2_attn.py +++ b/vllm/v1/attention/backends/mamba2_attn.py @@ -6,7 +6,7 @@ from typing import Any import torch -from vllm.config import VllmConfig +from vllm.config import VllmConfig, get_current_vllm_config_or_none from vllm.model_executor.layers.mamba.checkpoint import ( MambaPrefillCheckpointBuilder, MambaPrefillCheckpointMetadata, @@ -15,12 +15,13 @@ from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import ( AttentionBackend, CommonAttentionMetadata, + max_decode_query_len, ) from vllm.v1.attention.backends.mamba_attn import ( BaseMambaAttentionMetadata, BaseMambaAttentionMetadataBuilder, ) -from vllm.v1.kv_cache_interface import MambaSpec +from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec def compute_varlen_chunk_metadata( @@ -104,6 +105,14 @@ class Mamba2AttentionBackend(AttentionBackend): def is_ssm(cls) -> bool: return True + @classmethod + def supports_device_cpu_query_lens_mismatch(cls) -> bool: + # Decode kernels take device offsets; ReplaySSM plans off CPU counts. + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return True + return not vllm_config.cache_config.use_replayssm + @dataclass class Mamba2AttentionMetadata(BaseMambaAttentionMetadata): @@ -124,6 +133,14 @@ class Mamba2AttentionMetadataBuilder( ): metadata_cls = Mamba2AttentionMetadata + @classmethod + def get_varlen_cudagraph_max_query_len( + cls, + vllm_config: VllmConfig, + kv_cache_spec: KVCacheSpec, + ) -> int | None: + return max_decode_query_len(vllm_config) + def __init__( self, kv_cache_spec: MambaSpec,