diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py index 3ab354cd19..7632b26ea4 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py @@ -31,11 +31,13 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import ( from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import ( OffloadingConnectorScheduler, RequestOffloadState, + get_sliding_window_size_in_chunks, is_store_reachable_swa_chunk, ) from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, FullAttentionSpec, KVCacheGroupSpec, SlidingWindowSpec, @@ -3311,3 +3313,19 @@ def test_request_finished_mixed_full_attn_and_sliding_window( # Verify fence is empty after full lifecycle (cleanup happened). assert runner.connector_scheduler._block_id_to_pending_jobs == {} assert len(runner.connector_scheduler._jobs) == 0 + + +def test_chunked_local_attention_reports_its_chunk_window(): + """Llama 4 uses chunked local attention, which used to trip the + FullAttentionSpec assert and kill the engine at startup.""" + spec = ChunkedLocalAttentionSpec( + block_size=16, + num_kv_heads=2, + head_size=64, + dtype=torch.bfloat16, + attention_chunk_size=8192, + ) + + assert get_sliding_window_size_in_chunks(spec, tokens_per_chunk=1024) == 8 + # Partial chunks round up, so the reachable tail is never understated. + assert get_sliding_window_size_in_chunks(spec, tokens_per_chunk=3000) == 3 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py index 8bc4451967..1644d9c72e 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py @@ -31,6 +31,7 @@ from vllm.utils.math_utils import cdiv, round_down from vllm.v1.core.kv_cache_manager import KVCacheBlocks from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.kv_cache_interface import ( + ChunkedLocalAttentionSpec, FullAttentionSpec, KVCacheConfig, KVCacheSpec, @@ -112,6 +113,11 @@ def get_sliding_window_size_in_chunks( assert kv_cache_spec.sliding_window > 0 return cdiv(kv_cache_spec.sliding_window, tokens_per_chunk) + if isinstance(kv_cache_spec, ChunkedLocalAttentionSpec): + # Attention never reaches back past one chunk + assert kv_cache_spec.attention_chunk_size > 0 + return cdiv(kv_cache_spec.attention_chunk_size, tokens_per_chunk) + if isinstance(kv_cache_spec, MambaSpec): # Mamba depends on a single state return 1