[Bugfix][KV Offload] Keep Mamba block span unscaled under DCP (#49964)
Signed-off-by: Jonguk Cheong <jdal3031@snu.ac.kr> Co-authored-by: OpenAI Codex <codex@openai.com>
Jonguk Cheong committed
Jul 28, 2026 at 23:41 UTC
1e81853afc8701cb50b649e94a23c977da8c1ed0
2 files changed
+66
-2
tests/v1/kv_connector/unit/offloading_connector/test_config.py
+56
@@ -8,10 +8,15 @@ from unittest.mock import MagicMock, patch
8
import pytest
9
import torch
10
11
+from tests.v1.kv_connector.unit.offloading_connector.utils import MockOffloadingSpec
12
from vllm.config import KVTransferConfig, ParallelConfig, VllmConfig
13
from vllm.distributed.kv_transfer.kv_connector.v1.offloading.config import (
14
build_offloading_config,
15
)
16
+from vllm.distributed.kv_transfer.kv_connector.v1.offloading.scheduler import (
17
+ SchedulerOffloadConfig,
18
+ is_store_reachable_swa_chunk,
19
+)
20
from vllm.platforms import current_platform
21
from vllm.v1.kv_cache_interface import (
22
FullAttentionSpec,
@@ -179,6 +184,25 @@ def _make_hybrid_kv_cache_config() -> KVCacheConfig:
184
)
185
186
187
+def _make_mamba_hybrid_kv_cache_config() -> KVCacheConfig:
188
+ return KVCacheConfig(
189
+ num_blocks=4,
190
+ kv_cache_tensors=[],
191
+ kv_cache_groups=[
192
+ KVCacheGroupSpec(["full_layer"], _full_attention_spec()),
193
+ KVCacheGroupSpec(
194
+ ["mamba_layer"],
195
+ MambaSpec(
196
+ block_size=16,
197
+ shapes=((1, 1),),
198
+ dtypes=(torch.float32,),
199
+ mamba_cache_mode="align",
200
+ ),
201
+ ),
202
+ ],
203
+ )
204
+
205
+
206
def _parallelism_agnostic(kv_cache_groups: list[KVCacheGroupSpec]) -> bool:
207
config = _make_vllm_config()
208
kv_cache_config = KVCacheConfig(
@@ -267,6 +291,38 @@ def test_prefill_context_parallelism_does_not_scale_group_blocks():
291
assert offloading_config.cache.blocks_per_chunk == 4
292
293
294
+def test_dcp_scales_attention_but_not_mamba_group_blocks():
295
+ config = _make_vllm_config(tensor_parallel_size=2, decode_context_parallel_size=2)
296
+ config.speculative_config = None
297
+
298
+ offloading_config = build_offloading_config(
299
+ config, _make_mamba_hybrid_kv_cache_config()
300
+ )
301
+
302
+ assert tuple(group.tokens_per_block for group in offloading_config.groups) == (
303
+ 32,
304
+ 16,
305
+ )
306
+ scheduler_config = SchedulerOffloadConfig.from_spec(
307
+ MockOffloadingSpec(offloading_config),
308
+ config,
309
+ _make_mamba_hybrid_kv_cache_config(),
310
+ )
311
+ mamba_group = scheduler_config.kv_group_configs[1]
312
+ assert mamba_group.alignment_chunk_count == 2
313
+ assert [
314
+ chunk_idx
315
+ for chunk_idx in range(4)
316
+ if is_store_reachable_swa_chunk(
317
+ chunk_idx,
318
+ 4,
319
+ mamba_group.alignment_chunk_count,
320
+ mamba_group.sliding_window_size_in_chunks,
321
+ mamba_group.is_eagle_group,
322
+ )
323
+ ] == [1, 3]
324
+
325
+
326
def test_preserves_data_parallel_index():
327
config = _make_vllm_config()
328
config.parallel_config.data_parallel_index = 2
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+10
-2
@@ -5,7 +5,11 @@
5
from typing import TYPE_CHECKING
6
7
from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes
8
-from vllm.v1.kv_cache_interface import FullAttentionSpec, MLAAttentionSpec
8
+from vllm.v1.kv_cache_interface import (
9
+ AttentionSpec,
10
+ FullAttentionSpec,
11
+ MLAAttentionSpec,
12
+)
13
from vllm.v1.kv_offload.config import (
14
OffloadingCacheConfig,
15
OffloadingConfig,
@@ -40,7 +44,11 @@ def build_offloading_config(
44
OffloadingGroupConfig(
45
tokens_per_block=(
46
group.kv_cache_spec.block_size
43
- * parallel_config.decode_context_parallel_size
47
+ * (
48
+ parallel_config.decode_context_parallel_size
49
+ if isinstance(group.kv_cache_spec, AttentionSpec)
50
+ else 1
51
+ )
52
),
53
layer_names=tuple(group.layer_names),
54
)