[Bugfix][KV Offload] Handle chunked local attention in offloading scheduler (#51161)

Signed-off-by: almogtavor <almogtavor@gmail.com>
Co-authored-by: Or Ozeri <oro@il.ibm.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

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

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