@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 357fddf614

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 2, 22:00 UTC

[kv_offload]: Add DSv4 support (#43142)

Signed-off-by: Or Ozeri <oro@il.ibm.com>

Or Ozeri committed May 24, 2026 at 11:10 UTC 357fddf6147780404cb07b5a7d58b8434c9e828d

4 files changed +214 -14

tests/v1/kv_connector/unit/offloading_connector/test_scheduler.py

+121

@@ -1041,3 +1041,124 @@ def test_reset_cache(request_runner, async_scheduling: bool):

1041 for req_status in runner.connector_scheduler._req_status.values():

1042 for group_state in req_status.group_states:

1043 assert group_state.next_stored_block_idx == 0

1044 +

1045 +

1046 +@pytest.mark.parametrize("async_scheduling", [True, False])

1047 +def test_swa_alignment_skip(request_runner, async_scheduling: bool):

1048 + """SWA blocks unreachable by the load path are skipped during store.

1049 +

1050 + Simulates a DeepSeek V4-like hybrid architecture where SWA groups have

1051 + much smaller block sizes than the full-attention (MLA) group, causing

1052 + most SWA blocks to be unreachable by the alignment-based load path.

1053 +

1054 + Setup:

1055 + - Group 0: full attention (MLA-like), block_size=16

1056 + - Group 1: SWA, block_size=4, sliding_window=8

1057 +

1058 + alignment_block_count = 16 / 4 = 4 SWA blocks per alignment segment.

1059 + sliding_window_size_in_blocks = ceil(8 / 4) = 2.

1060 + Within each segment of 4 SWA blocks, only the trailing 2 are stored.

1061 +

1062 + With 32 tokens (2 full-attn blocks, 8 SWA blocks):

1063 + - Group 0 stores: blocks 0, 1 (all full-attn blocks)

1064 + - Group 1 stores: blocks 2, 3, 6, 7 (skip 0,1,4,5)

1065 +

1066 + For real DeepSeek V4 (100K tokens), this reduces SWA stores by ~78%.

1067 + """

1068 + full_attn_block_size = 16

1069 + swa_block_size = 4

1070 + sliding_window = 8

1071 + num_gpu_blocks = 200

1072 +

1073 + kv_cache_groups = [

1074 + KVCacheGroupSpec(

1075 + ["layer0"],

1076 + FullAttentionSpec(

1077 + block_size=full_attn_block_size,

1078 + num_kv_heads=1,

1079 + head_size=1,

1080 + dtype=torch.float32,

1081 + ),

1082 + ),

1083 + KVCacheGroupSpec(

1084 + ["layer1"],

1085 + SlidingWindowSpec(

1086 + block_size=swa_block_size,

1087 + num_kv_heads=1,

1088 + head_size=1,

1089 + dtype=torch.float32,

1090 + sliding_window=sliding_window,

1091 + ),

1092 + ),

1093 + ]

1094 +

1095 + runner = request_runner(

1096 + block_size=swa_block_size,

1097 + num_gpu_blocks=num_gpu_blocks,

1098 + async_scheduling=async_scheduling,

1099 + kv_cache_groups=kv_cache_groups,

1100 + )

1101 +

1102 + # Verify config: alignment_block_count computed correctly

1103 + kv_group_configs = runner.connector_scheduler.config.kv_group_configs

1104 + assert len(kv_group_configs) == 2

1105 + # Group 0: full attention -> no alignment skip

1106 + assert kv_group_configs[0].alignment_block_count is None

1107 + assert kv_group_configs[0].sliding_window_size_in_blocks is None

1108 + assert kv_group_configs[0].offloaded_block_size == full_attn_block_size

1109 + # Group 1: SWA -> alignment_block_count = 16/4 = 4, tail = 2

1110 + assert kv_group_configs[1].alignment_block_count == 4

1111 + assert kv_group_configs[1].sliding_window_size_in_blocks == 2

1112 + assert kv_group_configs[1].offloaded_block_size == swa_block_size

1113 +

1114 + # Send 32 tokens = 2 full-attn blocks (block_size=16) = 8 SWA blocks

1115 + # (block_size=4). Decode 1 token to kick off processing (stores are

1116 + # deferred to next step).

1117 + num_tokens = 32

1118 + runner.new_request(token_ids=[0] * num_tokens)

1119 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

1120 + generate_store_output(keys)

1121 + )

1122 + runner.run(decoded_tokens=[0])

1123 +

1124 + # Decode 1 more token to complete the deferred stores from above.

1125 + runner.manager.prepare_store.side_effect = lambda keys, req_context: (

1126 + generate_store_output(keys)

1127 + )

1128 + runner.run(

1129 + decoded_tokens=[EOS_TOKEN_ID],

1130 + # Group 0 (full attn, block_size=16): 2 offloaded blocks

1131 + # -> GPU blocks (0, 0) and (0, 1)

1132 + # Group 1 (SWA, block_size=4): 8 offloaded blocks, skip first 2

1133 + # per segment of 4:

1134 + # Segment 0 (blocks 0-3): skip 0,1 -> store (1, 2), (1, 3)

1135 + # Segment 1 (blocks 4-7): skip 4,5 -> store (1, 6), (1, 7)

1136 + expected_stored=(

1137 + (0, 0),

1138 + (0, 1),

1139 + (1, 2),

1140 + (1, 3),

1141 + (1, 6),

1142 + (1, 7),

1143 + ),

1144 + )

1145 +

1146 + # Verify that loads still work correctly for the stored SWA blocks.

1147 + runner.scheduler.reset_prefix_cache()

1148 + runner.new_request(token_ids=[0] * num_tokens + [1])

1149 + runner.manager.lookup.return_value = True

1150 + runner.connector_scheduler._maximal_prefix_lookup = lambda key, req_context: 2

1151 + runner.run(

1152 + decoded_tokens=[EOS_TOKEN_ID],

1153 + # Group 0: full prefix lookup hits 2 offloaded blocks

1154 + # -> loads GPU blocks (0, 0), (0, 1)

1155 + # Group 1: sliding window lookup finds trailing 2 from last segment

1156 + # (blocks 6, 7 which were stored)

1157 + # -> loads GPU blocks (1, 6), (1, 7)

1158 + expected_loaded=(

1159 + (0, 0),

1160 + (0, 1),

1161 + (1, 6),

1162 + (1, 7),

1163 + ),

1164 + )

vllm/distributed/kv_transfer/kv_connector/v1/offloading/scheduler.py

+70 -6

@@ -65,6 +65,12 @@ class GroupOffloadConfig(NamedTuple):

65 hash_block_size_factor: int

66 # None below means full attention

67 sliding_window_size_in_blocks: int | None

68 + # Number of this group's offloaded blocks per full-attention alignment

69 + # segment. Used to skip storing SWA blocks that can never serve a load

70 + # hit (e.g. DeepSeek V4 where SWA groups have much smaller block sizes

71 + # than the MLA full-attention group).

72 + # None for full-attention groups or when the optimization doesn't apply.

73 + alignment_block_count: int | None = None

76 def get_sliding_window_size_in_blocks(

@@ -89,6 +95,41 @@ class SchedulerOffloadConfig(NamedTuple):

96 @classmethod

97 def from_spec(cls, spec: OffloadingSpec) -> "SchedulerOffloadConfig":

98 + # Determine the alignment token count from the full-attention group(s).

99 + # This is the offloaded_block_size of the full-attention group; load

100 + # hits are always aligned to this boundary, so SWA blocks earlier in

101 + # each segment can never serve a load hit. Relevant for hybrid

102 + # architectures like DeepSeek V4 (MLA + SWA groups).

103 + full_attn_offloaded_block_sizes: set[int] = set()

104 + for idx, gpu_block_size in enumerate(spec.gpu_block_size):

105 + kv_spec = spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec

106 + sw = get_sliding_window_size_in_blocks(

107 + kv_spec, gpu_block_size * spec.block_size_factor

108 + )

109 + if sw is None:

110 + full_attn_offloaded_block_sizes.add(

111 + gpu_block_size * spec.block_size_factor

112 + )

113 +

114 + # Only apply the optimization if there's a single consistent

115 + # full-attention alignment size.

116 + alignment_tokens: int | None = None

117 + if len(full_attn_offloaded_block_sizes) == 1:

118 + alignment_tokens = full_attn_offloaded_block_sizes.pop()

119 +

120 + def _alignment_block_count(

121 + offloaded_block_size: int,

122 + sliding_window_size_in_blocks: int | None,

123 + ) -> int | None:

124 + if alignment_tokens is None or sliding_window_size_in_blocks is None:

125 + return None

126 + if alignment_tokens <= offloaded_block_size:

127 + return None

128 + per_segment = alignment_tokens // offloaded_block_size

129 + if sliding_window_size_in_blocks >= per_segment:

130 + return None

131 + return per_segment

132 +

133 return cls(

134 num_workers=spec.vllm_config.parallel_config.world_size,

135 kv_group_configs=tuple(

@@ -100,9 +141,14 @@ class SchedulerOffloadConfig(NamedTuple):

141 (gpu_block_size * spec.block_size_factor)

142 // spec.hash_block_size

143 ),

103 - sliding_window_size_in_blocks=get_sliding_window_size_in_blocks(

104 - spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec,

105 - gpu_block_size * spec.block_size_factor,

144 + sliding_window_size_in_blocks=(

145 + sw := get_sliding_window_size_in_blocks(

146 + spec.kv_cache_config.kv_cache_groups[idx].kv_cache_spec,

147 + gpu_block_size * spec.block_size_factor,

148 + )

149 + ),

150 + alignment_block_count=_alignment_block_count(

151 + gpu_block_size * spec.block_size_factor, sw

152 ),

153 )

154 for idx, gpu_block_size in enumerate(spec.gpu_block_size)

@@ -639,6 +685,7 @@ class OffloadingConnectorScheduler:

685 num_offloadable_tokens = min(num_tokens_after_batch, req.num_tokens)

686

687 # Filter out blocks skipped due to sliding window attention / SSM

688 + # or unreachable by the load path's alignment constraints.

689 new_offload_keys: list[OffloadKey] = []

690 for group_config, group_state in zip(

691 self.config.kv_group_configs, req_status.group_states

@@ -663,9 +710,26 @@ class OffloadingConnectorScheduler:

710 ]

711 assert len(offload_keys) == len(offload_block_ids)

712

666 - for offload_key, block_id in zip(offload_keys, offload_block_ids):

667 - if block_id != 0:

668 - new_offload_keys.append(offload_key)

713 + alignment_block_count = group_config.alignment_block_count

714 + tail = group_config.sliding_window_size_in_blocks

715 +

716 + for key_idx, (offload_key, block_id) in enumerate(

717 + zip(offload_keys, offload_block_ids)

718 + ):

719 + if block_id == 0:

720 + continue

721 + # Skip SWA blocks that can never serve a load hit:

722 + # within each full-attention alignment segment, only the

723 + # trailing `tail` blocks are reachable by

724 + # _sliding_window_lookup. For DeepSeek V4 with 100K

725 + # tokens this reduces SWA stores by ~78%.

726 + if alignment_block_count is not None:

727 + assert tail is not None

728 + abs_block_idx = start_block_idx + key_idx

729 + pos_in_segment = abs_block_idx % alignment_block_count

730 + if pos_in_segment < alignment_block_count - tail:

731 + continue

732 + new_offload_keys.append(offload_key)

733

734 if not new_offload_keys:

735 req_status.advance_stored_idx(num_offloadable_tokens)

vllm/distributed/kv_transfer/kv_connector/v1/offloading/worker.py

+16 -3

@@ -136,6 +136,9 @@ class OffloadingConnectorWorker:

136 page_size_bytes[layer_name] = (

137 layer_kv_cache_spec.page_size_bytes

138 )

139 + unpadded_page_size_bytes[layer_name] = (

140 + layer_kv_cache_spec.real_page_size_bytes

141 + )

142 else:

143 # Flash Attention case: (2, num_blocks, ...)

144 assert test_shape[0] == 2

@@ -157,8 +160,9 @@ class OffloadingConnectorWorker:

160 tensors_per_block[layer_name] = tuple(raw.unbind(0))

161

162 page_size_bytes[layer_name] = half_page_size

160 -

161 - unpadded_page_size_bytes[layer_name] = page_size_bytes[layer_name]

163 + unpadded_page_size_bytes[layer_name] = (

164 + layer_kv_cache_spec.real_page_size_bytes // 2

165 + )

166

167 elif isinstance(layer_kv_cache_spec, MambaSpec):

168 state_tensors = kv_caches[layer_name]

@@ -191,7 +195,16 @@ class OffloadingConnectorWorker:

195 block_tensors: list[CanonicalKVCacheTensor] = []

196 block_data_refs: dict[str, list[CanonicalKVCacheRef]] = defaultdict(list)

197 for kv_cache_tensor in self.spec.kv_cache_config.kv_cache_tensors:

194 - tensor_layer_names = kv_cache_tensor.shared_by

198 + # Filter to layers that were actually processed above.

199 + # _get_kv_cache_config_deepseek_v4 emits KVCacheTensor entries for

200 + # every (tuple_idx, page_size) slot; slots where no group has a

201 + # layer at that index produce an empty shared_by (reserved memory

202 + # with no corresponding model layer).

203 + tensor_layer_names = [

204 + n for n in kv_cache_tensor.shared_by if n in tensors_per_block

205 + ]

206 + if not tensor_layer_names:

207 + continue

208

209 # verify all layers in the group reference the exact same tensors

210 assert len({len(tensors_per_block[n]) for n in tensor_layer_names}) == 1

vllm/v1/kv_offload/base.py

+7 -5

@@ -13,6 +13,7 @@ import numpy as np

13 import torch

15 from vllm.logger import init_logger

16 +from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes

18 if TYPE_CHECKING:

19 from vllm.config import VllmConfig

@@ -348,17 +349,18 @@ class OffloadingSpec(ABC):

349 * parallel_config.prefill_context_parallel_size

350 )

351

351 - # block size used by vLLM for hashing request tokens for the sake

352 - # of enabling prefix caching

353 - self.hash_block_size = (

354 - vllm_config.cache_config.block_size * context_parallel_factor

355 - )

352 # gpu block size per group

353 self.gpu_block_size: tuple[int, ...] = tuple(

354 kv_cache_group.kv_cache_spec.block_size * context_parallel_factor

355 for kv_cache_group in kv_cache_config.kv_cache_groups

356 )

357

358 + # hash_block_size must match what the scheduler uses for

359 + # Request.block_hashes (resolved via resolve_kv_cache_block_sizes).

360 + _, self.hash_block_size = resolve_kv_cache_block_sizes(

361 + kv_cache_config, vllm_config

362 + )

363 +

364 for block_size in self.gpu_block_size:

365 assert block_size % self.hash_block_size == 0, (

366 f"gpu_block_size={block_size} not divisible by "