[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>
Almog Tavor committed
Aug 9, 2026 at 10:25 UTC
eb24bc38cfc6818d6dc6d8440faf6e7062396f17
2 files changed
+24
tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py
+18
@@ -31,11 +31,13 @@ from vllm.distributed.kv_transfer.kv_connector.v1.offloading.metrics import (
31
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
32
OffloadingConnectorScheduler,
33
RequestOffloadState,
34
+ get_sliding_window_size_in_chunks,
35
is_store_reachable_swa_chunk,
36
)
37
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
38
from vllm.v1.core.kv_cache_utils import BlockHash, KVCacheBlock
39
from vllm.v1.kv_cache_interface import (
40
+ ChunkedLocalAttentionSpec,
41
FullAttentionSpec,
42
KVCacheGroupSpec,
43
SlidingWindowSpec,
@@ -3311,3 +3313,19 @@ def test_request_finished_mixed_full_attn_and_sliding_window(
3313
# Verify fence is empty after full lifecycle (cleanup happened).
3314
assert runner.connector_scheduler._block_id_to_pending_jobs == {}
3315
assert len(runner.connector_scheduler._jobs) == 0
3316
+
3317
+
3318
+def test_chunked_local_attention_reports_its_chunk_window():
3319
+ """Llama 4 uses chunked local attention, which used to trip the
3320
+ FullAttentionSpec assert and kill the engine at startup."""
3321
+ spec = ChunkedLocalAttentionSpec(
3322
+ block_size=16,
3323
+ num_kv_heads=2,
3324
+ head_size=64,
3325
+ dtype=torch.bfloat16,
3326
+ attention_chunk_size=8192,
3327
+ )
3328
+
3329
+ assert get_sliding_window_size_in_chunks(spec, tokens_per_chunk=1024) == 8
3330
+ # Partial chunks round up, so the reachable tail is never understated.
3331
+ assert get_sliding_window_size_in_chunks(spec, tokens_per_chunk=3000) == 3
vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py
+6
@@ -31,6 +31,7 @@ from vllm.utils.math_utils import cdiv, round_down
31
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
32
from vllm.v1.core.sched.output import SchedulerOutput
33
from vllm.v1.kv_cache_interface import (
34
+ ChunkedLocalAttentionSpec,
35
FullAttentionSpec,
36
KVCacheConfig,
37
KVCacheSpec,
@@ -112,6 +113,11 @@ def get_sliding_window_size_in_chunks(
113
assert kv_cache_spec.sliding_window > 0
114
return cdiv(kv_cache_spec.sliding_window, tokens_per_chunk)
115
116
+ if isinstance(kv_cache_spec, ChunkedLocalAttentionSpec):
117
+ # Attention never reaches back past one chunk
118
+ assert kv_cache_spec.attention_chunk_size > 0
119
+ return cdiv(kv_cache_spec.attention_chunk_size, tokens_per_chunk)
120
+
121
if isinstance(kv_cache_spec, MambaSpec):
122
# Mamba depends on a single state
123
return 1