diff --git a/tests/v1/attention/test_mla_noncausal.py b/tests/v1/attention/test_mla_noncausal.py index 040a2ad4be..c06ded5c13 100644 --- a/tests/v1/attention/test_mla_noncausal.py +++ b/tests/v1/attention/test_mla_noncausal.py @@ -20,7 +20,9 @@ class _NonCausalMLAMetadataBuilder(MLACommonMetadataBuilder[MLACommonMetadata]): def _metadata( - query_start_loc: list[int], num_tokens: int | None = None + query_start_loc: list[int], + num_tokens: int | None = None, + causal: bool = False, ) -> CommonAttentionMetadata: num_reqs = len(query_start_loc) - 1 num_tokens = query_start_loc[-1] if num_tokens is None else num_tokens @@ -38,7 +40,7 @@ def _metadata( num_reqs, 3 ), slot_mapping=torch.arange(num_tokens), - causal=False, + causal=causal, seq_lens_cpu_upper_bound=None, ) @@ -50,6 +52,7 @@ def _builder(marked: bool = True) -> _NonCausalMLAMetadataBuilder: builder.query_len_support = QueryLenSupport.SINGLE_ONLY builder.non_causal_multi_token_decode = marked builder.dcp_world_size = 1 + builder.use_pcp = False builder.metadata_cls = MLACommonMetadata builder.model_config = SimpleNamespace( dtype=torch.bfloat16, get_head_size=lambda: 576 @@ -57,9 +60,19 @@ def _builder(marked: bool = True) -> _NonCausalMLAMetadataBuilder: return builder -def test_noncausal_block_uses_decode_without_cpu_lengths(): +def test_group_capability_keeps_runtime_causality_per_step(): + builder = _builder() + + target_metadata = builder.build(0, _metadata([0, 1, 2], causal=True)) + assert ( + target_metadata.num_decodes, + target_metadata.num_decode_tokens, + target_metadata.num_prefills, + target_metadata.causal, + ) == (2, 2, 0, True) + common_metadata = _metadata([0, 8, 16]) - metadata = _builder().build(0, common_metadata) + metadata = builder.build(0, common_metadata) assert metadata.num_decodes == 2 assert metadata.num_decode_tokens == 16 @@ -109,7 +122,7 @@ def test_noncausal_decode_metadata_keeps_live_request_buffers(): ) -def test_mla_cache_marker_is_a_group_property(): +def test_mla_cache_marker_is_promoted_to_group_capability(): kwargs = { "block_size": 64, "num_kv_heads": 1, @@ -123,5 +136,4 @@ def test_mla_cache_marker_is_a_group_property(): assert not MLAAttentionSpec.merge( [unmarked, unmarked] ).non_causal_multi_token_decode - with pytest.raises(AssertionError, match="non_causal_multi_token_decode"): - MLAAttentionSpec.merge([marked, unmarked]) + assert MLAAttentionSpec.merge([unmarked, marked]).non_causal_multi_token_decode diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index 438b8753f8..f0bd296b71 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -3,8 +3,11 @@ import copy import hashlib import importlib +import subprocess +import sys from collections.abc import Callable from dataclasses import replace +from pathlib import Path from types import SimpleNamespace from typing import Any, cast @@ -1463,6 +1466,28 @@ def test_is_kv_cache_spec_uniform(): } assert not is_kv_cache_spec_uniform(kv_cache_spec) + script = """ +import sys + +from vllm.v1.core.kv_cache_utils import is_kv_cache_spec_uniform +from vllm.v1.kv_cache_interface import KVCacheSpec + +if sys.flags.optimize < 1: + raise RuntimeError("subprocess is not running with optimization enabled") +specs = { + "a": KVCacheSpec(block_size=1), + "b": KVCacheSpec(block_size=2), +} +if is_kv_cache_spec_uniform(specs): + raise RuntimeError("different specs were treated as uniform") +""" + subprocess.run( + [sys.executable, "-O", "-c", script], + cwd=Path(__file__).resolve().parents[3], + check=True, + timeout=60, + ) + @pytest.mark.parametrize( ("model_id", "max_model_len", "want_estimated_max_len"), diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index d322a98d1e..bb9494d15c 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -224,9 +224,10 @@ class KVCacheSpec: """ Merge a list of KVCacheSpec objects into a single KVCacheSpec object. """ - assert all(spec == specs[0] for spec in specs[1:]), ( - "All layers in the same KV cache group must be the same." - ) + if not all(spec == specs[0] for spec in specs[1:]): + raise AssertionError( + "All layers in the same KV cache group must be the same." + ) return copy.deepcopy(specs[0]) def is_uniform_with_collection( @@ -558,7 +559,8 @@ class MLAAttentionSpec(FullAttentionSpec): model_version: str | None = None storage_block_size: int | None = None """Token width used to view storage when it differs from the kernel block.""" - # Marks draft groups that flatten a non-causal query block into decode rows. + # Group capability enabled when any member flattens a non-causal query block + # into decode rows. Runtime metadata still selects causal vs. non-causal mode. non_causal_multi_token_decode: bool = False # MLA stores a single latent vector per state; there is no separate V. head_size_v: int = 0 @@ -586,11 +588,6 @@ class MLAAttentionSpec(FullAttentionSpec): "quantization method, tokens per state, model version, and storage " "block size." ) - non_causal_mtd_set = {spec.non_causal_multi_token_decode for spec in specs} - assert len(non_causal_mtd_set) == 1, ( - "All attention layers in the same KV cache group must agree on " - "non_causal_multi_token_decode." - ) merged_spec = cls( block_size=specs[0].block_size, num_kv_heads=specs[0].num_kv_heads, @@ -604,7 +601,9 @@ class MLAAttentionSpec(FullAttentionSpec): tokens_per_state=tokens_per_state_set.pop(), model_version=model_version_set.pop(), storage_block_size=storage_block_size_set.pop(), - non_causal_multi_token_decode=non_causal_mtd_set.pop(), + non_causal_multi_token_decode=any( + spec.non_causal_multi_token_decode for spec in specs + ), ) for spec in specs: for f in fields(AttentionSpec):