@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 30ae8248b6

[KVConnector][MoRIIO] Support K3 DSpark hybrid READ (#57700)

Signed-off-by: Yichao Zhu <Yichao.Zhu@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

YukioZzz committed Sep 30, 2026 at 03:14 UTC 30ae8248b663ea56398d44d30c0d8276a598c4a0

5 files changed +288 -45

tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py

+194 -6

@@ -21,6 +21,9 @@ import pytest

21 import torch

23 from vllm.platforms import current_platform

24 +from vllm.v1.core.block_pool import BlockPool

25 +from vllm.v1.core.single_type_kv_cache_manager import FullAttentionManager

26 +from vllm.v1.kv_cache_interface import FullAttentionSpec, UniformTypeKVCacheSpecs

28 mori_available = importlib.util.find_spec("mori") is not None

@@ -44,6 +47,16 @@ MoRIIOMode = moriio_connector.MoRIIOMode

47 moriio_common = importlib.import_module(

48 "vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common"

49 )

50 +get_role = moriio_common.get_role

51 +set_role = moriio_common.set_role

52 +

53 +

54 +@pytest.fixture

55 +def restore_moriio_role():

56 + """Undo the process-global role a full scheduler construction leaves behind."""

57 + role = get_role()

58 + yield

59 + set_role(role)

62 class _FakeScheduler(moriio_connector.MoRIIOConnectorScheduler): # type: ignore[name-defined]

@@ -53,6 +66,7 @@ class _FakeScheduler(moriio_connector.MoRIIOConnectorScheduler): # type: ignore

66 def __init__(self, **attrs):

67 self._mamba_group_ids: list[int] = []

68 self._attn_group_ids: list[int] = [0]

69 + self._num_ssm_scratch_blocks = 0

70 self._ssm_state_slots_are_positional = False

71 self._is_hma_required = False

72 self.kv_cache_config = SimpleNamespace(

@@ -110,22 +124,56 @@ def _gdn_split_info(conv_rows=3, key_dim=4, value_dim=8, dtype_size=2):

124 )

125

126

113 -def _mamba_spec(num_states: int = 2):

127 +def _mamba_spec(num_states: int = 2, *, num_speculative_blocks: int = 0):

128 shapes = tuple((1, 1) for _ in range(num_states))

129 return moriio_connector.MambaSpec(

130 block_size=16,

131 shapes=shapes,

132 dtypes=(torch.float32,) * num_states,

133 mamba_cache_mode="all",

134 + num_speculative_blocks=num_speculative_blocks,

135 )

136

137

123 -def _gate_vllm_config(*, read_mode: bool = True, speculative_config=None):

138 +def _spec_config(method: str):

139 + return SimpleNamespace(method=method, use_dspark=lambda: method == "dspark")

140 +

141 +

142 +def _gate_vllm_config(

143 + *,

144 + read_mode: bool = True,

145 + speculative_config=None,

146 + block_size: int = 16,

147 + num_lookahead_tokens: int = 0,

148 + num_speculative_tokens: int = 0,

149 + mamba_cache_mode: str = "all",

150 + dcp_size: int = 1,

151 +):

152 return SimpleNamespace(

153 kv_transfer_config=SimpleNamespace(

126 - kv_connector_extra_config={"read_mode": read_mode}

154 + kv_connector_extra_config={

155 + "read_mode": read_mode,

156 + "host_ip": "127.0.0.1",

157 + "handshake_port": 6001,

158 + "notify_port": 6002,

159 + },

160 + kv_role="kv_consumer",

161 ),

162 speculative_config=speculative_config,

163 + cache_config=SimpleNamespace(

164 + block_size=block_size,

165 + mamba_cache_mode=mamba_cache_mode,

166 + ),

167 + scheduler_config=SimpleNamespace(disable_hybrid_kv_cache_manager=False),

168 + num_lookahead_tokens=num_lookahead_tokens,

169 + num_speculative_tokens=num_speculative_tokens,

170 + parallel_config=SimpleNamespace(

171 + tensor_parallel_size=1,

172 + decode_context_parallel_size=dcp_size,

173 + data_parallel_rank=0,

174 + data_parallel_size_local=1,

175 + data_parallel_size=1,

176 + ),

177 )

178

179

@@ -196,6 +244,7 @@ def test_exchange_blocks_ignore_transfer_disabled_group():

244

245

246 def test_scheduler_rejects_multiple_attention_groups_with_mamba():

247 + """A drafter owning its own transfer group is refused, with a hint."""

248 config = SimpleNamespace(

249 kv_cache_groups=[

250 SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),

@@ -205,7 +254,7 @@ def test_scheduler_rejects_multiple_attention_groups_with_mamba():

254 )

255 config.transfer_groups = tuple(config.kv_cache_groups)

256

208 - with pytest.raises(moriio_common.MoRIIOError, match="exactly one transferable"):

257 + with pytest.raises(moriio_common.MoRIIOError, match="separate attention group"):

258 moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config)

259

260

@@ -251,7 +300,7 @@ def test_scheduler_rejects_mamba_group_without_two_states():

300 moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config)

301

302

254 -def test_scheduler_rejects_speculative_hybrid_read():

303 +def test_scheduler_rejects_non_dspark_speculative_hybrid_read():

304 config = SimpleNamespace(

305 kv_cache_groups=[

306 SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()),

@@ -262,7 +311,105 @@ def test_scheduler_rejects_speculative_hybrid_read():

311

312 with pytest.raises(moriio_common.MoRIIOError, match="speculative decoding"):

313 moriio_connector.MoRIIOConnectorScheduler(

265 - _gate_vllm_config(speculative_config=object()), "engine", config

314 + _gate_vllm_config(speculative_config=_spec_config("ngram")),

315 + "engine",

316 + config,

317 + )

318 +

319 +

320 +def _dspark_wrapped_attention_config(*, dcp_sharded: bool = True):

321 + """Mamba first, so a wrapped full-attention group at index 1 has to be found.

322 +

323 + A bare ``isinstance(FullAttentionSpec)`` misses the wrapper and leaves

324 + ``_full_attn_group_idx`` at its initial 0, which is the Mamba group.

325 + """

326 + full = FullAttentionSpec(

327 + block_size=4,

328 + num_kv_heads=1,

329 + head_size=1,

330 + dtype=torch.float32,

331 + dcp_sharded=dcp_sharded,

332 + )

333 + config = SimpleNamespace(

334 + kv_cache_groups=[

335 + SimpleNamespace(

336 + enable_kv_transfer=True,

337 + kv_cache_spec=_mamba_spec(num_speculative_blocks=2),

338 + ),

339 + SimpleNamespace(

340 + enable_kv_transfer=True,

341 + kv_cache_spec=UniformTypeKVCacheSpecs(

342 + block_size=4, kv_cache_specs={"target": full, "draft": full}

343 + ),

344 + ),

345 + ],

346 + )

347 + config.transfer_groups = tuple(config.kv_cache_groups)

348 + return config

349 +

350 +

351 +@pytest.mark.usefixtures("restore_moriio_role")

352 +@pytest.mark.parametrize(

353 + "dcp_size,dcp_sharded,expected_tail", [(1, True, 4), (4, True, 1), (4, False, 4)]

354 +)

355 +def test_scheduler_accepts_dspark_with_wrapped_full_attention(

356 + dcp_size, dcp_sharded, expected_tail

357 +):

358 + config = _dspark_wrapped_attention_config(dcp_sharded=dcp_sharded)

359 + vllm_config = _gate_vllm_config(

360 + speculative_config=_spec_config("dspark"),

361 + block_size=4,

362 + num_lookahead_tokens=7,

363 + num_speculative_tokens=7,

364 + mamba_cache_mode="align",

365 + dcp_size=dcp_size,

366 + )

367 + scheduler = moriio_connector.MoRIIOConnectorScheduler(vllm_config, "engine", config)

368 +

369 + assert scheduler._full_attn_group_idx == 1

370 + assert scheduler._num_ssm_scratch_blocks == 2

371 + assert scheduler._max_decode_tail_blocks == expected_tail

372 +

373 + # A first decode step reserves padded target rows AND drafter lookahead.

374 + spec = config.transfer_groups[1].kv_cache_spec

375 + manager = FullAttentionManager(

376 + spec,

377 + block_pool=BlockPool(32, enable_caching=False, hash_block_size=4),

378 + enable_caching=False,

379 + kv_cache_group_id=1,

380 + scheduler_block_size=16,

381 + dcp_world_size=dcp_size,

382 + )

383 + imported_tokens = manager.block_size

384 + target_tokens = 1 + vllm_config.num_speculative_tokens

385 + manager.allocate_new_blocks(

386 + "decode",

387 + num_tokens=imported_tokens + target_tokens + vllm_config.num_lookahead_tokens,

388 + num_tokens_main_model=imported_tokens + target_tokens,

389 + )

390 + local = [b.block_id for b in manager.req_to_blocks["decode"]]

391 + remote = [100]

392 + assert scheduler._align_read_blocks(

393 + local, remote, scheduler._max_decode_tail_blocks

394 + ) == (local[:1], remote)

395 + with pytest.raises(ValueError, match="tail"):

396 + scheduler._align_read_blocks(

397 + local + [99], remote, scheduler._max_decode_tail_blocks

398 + )

399 +

400 +

401 +def test_scheduler_rejects_dspark_with_positional_mamba_cache():

402 + """Positional mode counts lookahead in the Mamba allocation, so the

403 + one-block recurrent-state tail bound no longer holds."""

404 + with pytest.raises(moriio_common.MoRIIOError, match="mamba_cache_mode"):

405 + moriio_connector.MoRIIOConnectorScheduler(

406 + _gate_vllm_config(

407 + speculative_config=_spec_config("dspark"),

408 + block_size=4,

409 + mamba_cache_mode="all",

410 + ),

411 + "engine",

412 + _dspark_wrapped_attention_config(),

413 )

414

415

@@ -316,6 +463,30 @@ def test_split_block_groups_keeps_only_running_state_outside_all_mode():

463 assert sched.split_block_groups(([1], [40, 41])) == ([1], [[41]])

464

465

466 +@pytest.mark.parametrize(

467 + "positional,blocks,expected",

468 + [

469 + # Trailing scratch slots go in either cache mode.

470 + (True, [40, 41, 42], [40]),

471 + # Outside "all" the running state is picked only after they are gone.

472 + (False, [40, 41, 42, 43], [41]),

473 + # Never empty: a list shorter than the scratch count keeps one state.

474 + (True, [40], [40]),

475 + (False, [], []),

476 + ],

477 +)

478 +def test_split_block_groups_strips_dspark_scratch_slots(positional, blocks, expected):

479 + sched = _FakeScheduler(

480 + _has_mamba=True,

481 + _attn_group_ids=[0],

482 + _mamba_group_ids=[1],

483 + _num_ssm_scratch_blocks=2,

484 + _ssm_state_slots_are_positional=positional,

485 + )

486 +

487 + assert sched.split_block_groups(([1], blocks)) == ([1], [expected])

488 +

489 +

490 # --------------------------------------------------------------------------

491 # request_finished_all_groups

492 # --------------------------------------------------------------------------

@@ -434,6 +605,23 @@ def test_update_state_pairs_shorter_local_blocks_with_remote_suffix():

605 assert sched._req_kv_params["req"]["remote_block_ids"] == [[12], [92]]

606

607

608 +def test_update_state_pairs_trimmed_mamba_state_with_remote_state():

609 + """The DSpark shape end to end: scratch slots are stripped, the running

610 + state is then selected, and the attention list loses its lookahead tail."""

611 + sched = _make_read_scheduler()

612 + sched._ssm_state_slots_are_positional = False

613 + sched._num_ssm_scratch_blocks = 2

614 + sched._max_decode_tail_blocks = 2

615 + request = _make_read_request([[10, 11], [90]])

616 + # Mamba group: [null, running state, scratch, scratch].

617 + blocks = _FakeBlocks(all_groups=([100, 101, 102, 103], [200, 201, 202, 203]))

618 +

619 + sched.update_state_after_alloc(request, blocks, num_external_tokens=256)

620 +

621 + assert sched._reqs_need_recv["req"][1] == [[100, 101], [201]]

622 + assert sched._req_kv_params["req"]["remote_block_ids"] == [[10, 11], [90]]

623 +

624 +

625 def test_update_state_full_attention_hit_still_carries_mamba_state():

626 sched = _make_read_scheduler()

627 request = _make_read_request([[10], [90]])

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

@@ -2055,6 +2055,15 @@ def test_exchange_clipped_blocks_ssm_single_state():

2055 # Scratch slots not allocated: the state slot still survives.

2056 assert sched.get_exchange_clipped_blocks(([1], [5]))[1] == [5]

2057

2058 + assert sched.get_exchange_clipped_blocks(([1], [])) == ([1], [])

2059 + # Per-step partial lists must not be interpreted as state + scratch slots.

2060 + assert sched.get_exchange_clipped_blocks(([1], [7, 8]), clip_ssm=False) == (

2061 + [1],

2062 + [7, 8],

2063 + )

2064 + sched.blocks_per_sw = [2, 0]

2065 + assert sched.get_exchange_clipped_blocks(([1, 2, 3], [7, 8, 9])) == ([2, 3], [7])

2066 +

2067 # Non-mamba models pass through unchanged.

2068 fa_sched = make_nixl_scheduler(has_mamba=False)

2069 assert fa_sched.get_exchange_clipped_blocks(([1, 2],)) == ([1, 2],)

vllm/distributed/kv_transfer/kv_connector/utils.py

+12

@@ -33,6 +33,18 @@ EngineId = str

33 BlockIds = tuple[list[int], ...] | list[list[int]]

36 +def clip_ssm_state_blocks(

37 + blocks: list[int], num_spec_blocks: int, positional: bool

38 +) -> list[int]:

39 + """Drop speculative scratch slots, retaining position-indexed or running state."""

40 + if not blocks:

41 + return blocks

42 + # Keep at least the running-state slot, even before scratch allocation.

43 + if num_scratch := min(num_spec_blocks, len(blocks) - 1):

44 + blocks = blocks[:-num_scratch]

45 + return blocks if positional else blocks[-1:]

46 +

47 +

48 def get_kv_connector_cache_layout(vllm_config: VllmConfig | None = None):

49 # NOTE (NickLucche) When running disaggregated PD with NIXL, LBHNC layout is

50 # used for faster transfer.

vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py

+69 -34

@@ -17,7 +17,10 @@ import torch

17 import zmq

19 from vllm.config import CUDAGraphMode, VllmConfig

20 -from vllm.distributed.kv_transfer.kv_connector.utils import BlockIds

20 +from vllm.distributed.kv_transfer.kv_connector.utils import (

21 + BlockIds,

22 + clip_ssm_state_blocks,

23 +)

24 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

25 KVConnectorBase_V1,

26 KVConnectorMetadata,

@@ -88,10 +91,10 @@ from vllm.utils.network_utils import (

91 from vllm.v1.attention.selector import get_attn_backend

92 from vllm.v1.core.sched.output import SchedulerOutput

93 from vllm.v1.kv_cache_interface import (

91 - FullAttentionSpec,

94 KVCacheConfig,

95 MambaSpec,

96 SlidingWindowSpec,

97 + is_full_attention_spec,

98 )

99 from vllm.v1.outputs import KVConnectorOutput

100 from vllm.v1.request import RequestStatus

@@ -108,7 +111,7 @@ logger = init_logger(__name__)

111 # wait only needs to be long enough for it to make progress.

112 _SQ_FULL_BACKOFF_INITIAL_S = 0.001

113 _SQ_FULL_BACKOFF_MAX_S = 0.05

111 -_MAX_LOCAL_DECODE_TAIL_BLOCKS = 1

114 +_MAX_LOCAL_MAMBA_TAIL_BLOCKS = 1

115

116

117 try:

@@ -436,6 +439,26 @@ def _split_kv_cache_group_kinds(

439 return attn, mamba

440

441

442 +def _validate_hybrid_speculation(vllm_config: VllmConfig) -> None:

443 + speculative_config = vllm_config.speculative_config

444 + if speculative_config is None:

445 + return

446 + # DSpark is allowlisted because it is the only hybrid speculative method

447 + # validated end to end with MoRIIO READ.

448 + if not speculative_config.use_dspark():

449 + raise MoRIIOError(

450 + "MoRIIO hybrid READ supports DSpark speculative decoding only, got "

451 + f"method={speculative_config.method!r}"

452 + )

453 + if vllm_config.cache_config.mamba_cache_mode == "all":

454 + # Positional allocation counts lookahead, so its local tail can exceed

455 + # the single running-state block supported by hybrid READ.

456 + raise MoRIIOError(

457 + "MoRIIO hybrid READ does not support DSpark with "

458 + "mamba_cache_mode='all'; use 'align'"

459 + )

460 +

461 +

462 def _validate_mamba_specs(specs: Collection[MambaSpec]) -> MambaSpec | None:

463 """Validate the common packed-state contract and return its representative."""

464 specs = tuple(specs)

@@ -477,22 +500,22 @@ class MoRIIOConnectorScheduler:

500 kv_cache_config

501 )

502 self._has_mamba = bool(self._mamba_group_ids)

503 + mamba_spec = None

504 if self._has_mamba:

505 if len(self._attn_group_ids) != 1:

506 raise MoRIIOError(

507 "MoRIIO hybrid READ requires exactly one transferable "

508 "attention group, got "

485 - f"{len(self._attn_group_ids)} attention groups"

509 + f"{len(self._attn_group_ids)}; a drafter that owns a "

510 + "separate attention group is not supported, give it the "

511 + "target's group or disable KV transfer for it"

512 )

513 mamba_specs = [

514 cast(MambaSpec, kv_cache_config.transfer_groups[group_id].kv_cache_spec)

515 for group_id in self._mamba_group_ids

516 ]

491 - _validate_mamba_specs(mamba_specs)

492 - if vllm_config.speculative_config is not None:

493 - raise MoRIIOError(

494 - "MoRIIO hybrid READ does not support speculative decoding"

495 - )

517 + mamba_spec = _validate_mamba_specs(mamba_specs)

518 + _validate_hybrid_speculation(vllm_config)

519 if self.mode != MoRIIOMode.READ:

520 raise MoRIIOError(

521 "MoRIIO hybrid (mamba/KDA) transfer is implemented for READ "

@@ -503,18 +526,31 @@ class MoRIIOConnectorScheduler:

526 self._ssm_state_slots_are_positional = (

527 vllm_config.cache_config.mamba_cache_mode == "all"

528 )

506 - self.block_size = vllm_config.cache_config.block_size

507 - self._max_decode_tail_blocks = (

508 - _MAX_LOCAL_DECODE_TAIL_BLOCKS

509 - if self._has_mamba

510 - else math.ceil((vllm_config.num_lookahead_tokens + 1) / self.block_size)

529 + self._num_ssm_scratch_blocks = (

530 + mamba_spec.num_speculative_blocks if mamba_spec is not None else 0

531 )

532 + self.block_size = vllm_config.cache_config.block_size

533 + attn_block_size = self.block_size

534 + max_decode_tail_tokens = 1 + vllm_config.num_lookahead_tokens

535 + if self._has_mamba:

536 + attn_spec = kv_cache_config.transfer_groups[

537 + self._attn_group_ids[0]

538 + ].kv_cache_spec

539 + attn_block_size = attn_spec.block_size

540 + if attn_spec.dcp_sharded:

541 + attn_block_size *= (

542 + vllm_config.parallel_config.decode_context_parallel_size

543 + )

544 + # Admission can pad the final prompt token to 1 + num_spec rows;

545 + # allocate_slots reserves drafter lookahead in addition to those.

546 + max_decode_tail_tokens += vllm_config.num_speculative_tokens

547 + self._max_decode_tail_blocks = cdiv(max_decode_tail_tokens, attn_block_size)

548 self.engine_id: EngineId = engine_id

549

550 self._is_hma_required = (

551 not vllm_config.scheduler_config.disable_hybrid_kv_cache_manager

552 and any(

517 - not isinstance(g.kv_cache_spec, FullAttentionSpec)

553 + not is_full_attention_spec(g.kv_cache_spec)

554 for g in kv_cache_config.transfer_groups

555 )

556 )

@@ -524,9 +560,9 @@ class MoRIIOConnectorScheduler:

560 unsupported = [

561 type(g.kv_cache_spec).__name__

562 for g in kv_cache_config.transfer_groups

527 - if not isinstance(

528 - g.kv_cache_spec,

529 - (FullAttentionSpec, SlidingWindowSpec, MambaSpec),

563 + if not (

564 + is_full_attention_spec(g.kv_cache_spec)

565 + or isinstance(g.kv_cache_spec, (SlidingWindowSpec, MambaSpec))

566 )

567 ]

568 if unsupported:

@@ -565,7 +601,7 @@ class MoRIIOConnectorScheduler:

601 self._full_attn_group_idx = 0

602 self._full_attn_block_size = self.block_size

603 for gi, group in enumerate(kv_cache_config.kv_cache_groups):

568 - if isinstance(group.kv_cache_spec, FullAttentionSpec):

604 + if is_full_attention_spec(group.kv_cache_spec):

605 self._full_attn_group_idx = gi

606 self._full_attn_block_size = getattr(

607 group.kv_cache_spec, "block_size", self.block_size

@@ -953,6 +989,9 @@ class MoRIIOConnectorScheduler:

989 )

990 aligned_local_mamba: list[list[int]] = []

991 aligned_remote_mamba: list[list[int]] = []

992 + # Scratch slots were removed by

993 + # split_block_groups, so only one local running

994 + # state tail can remain.

995 for local_group, remote_group in zip(

996 mamba_block_groups,

997 remote_mamba_groups,

@@ -960,7 +999,9 @@ class MoRIIOConnectorScheduler:

999 ):

1000 aligned_local, aligned_remote = (

1001 self._align_read_blocks(

963 - local_group, remote_group, 1

1002 + local_group,

1003 + remote_group,

1004 + _MAX_LOCAL_MAMBA_TAIL_BLOCKS,

1005 )

1006 )

1007 aligned_local_mamba.append(aligned_local)

@@ -1257,21 +1298,15 @@ class MoRIIOConnectorScheduler:

1298 transfer_block_ids = self.kv_cache_config.select_transfer_block_ids(block_ids)

1299 attn = list(transfer_block_ids[self._attn_group_ids[0]])

1300 mamba_groups = [

1260 - self._clip_mamba_group(list(transfer_block_ids[group_id]))

1301 + clip_ssm_state_blocks(

1302 + list(transfer_block_ids[group_id]),

1303 + self._num_ssm_scratch_blocks,

1304 + self._ssm_state_slots_are_positional,

1305 + )

1306 for group_id in self._mamba_group_ids

1307 ]

1308 return attn, mamba_groups

1309

1265 - def _clip_mamba_group(self, blocks: list[int]) -> list[int]:

1266 - """Keep only the state-bearing slots of one mamba kv cache group."""

1267 - if not blocks:

1268 - return blocks

1269 - if not self._ssm_state_slots_are_positional:

1270 - # Single running state: everything before it is a null placeholder

1271 - # or the previous step's superseded state.

1272 - blocks = blocks[-1:]

1273 - return blocks

1274 -

1310 @staticmethod

1311 def _align_read_blocks(

1312 local_block_ids: list[int],

@@ -1282,9 +1317,9 @@ class MoRIIOConnectorScheduler:

1317

1318 Prefix hits can make the decode's unhashed suffix shorter than the

1319 prefill's full list, so read from the matching remote suffix. The

1285 - decoder can instead have one extra tail block for the locally

1286 - recomputed final token; that block has no remote source and is not

1287 - transferred.

1320 + decoder can instead have extra tail blocks for the locally recomputed

1321 + final token and speculative slots; these have no remote source and are

1322 + not transferred.

1323 """

1324 if len(local_block_ids) > len(remote_block_ids):

1325 extra_decode_blocks = len(local_block_ids) - len(remote_block_ids)

vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_scheduler.py

+4 -5

@@ -13,6 +13,7 @@ from vllm import envs

13 from vllm.distributed.kv_transfer.kv_connector.utils import (

14 BlockIds,

15 EngineId,

16 + clip_ssm_state_blocks,

17 yield_req_data,

18 )

19 from vllm.distributed.kv_transfer.kv_connector.v1.base import (

@@ -288,11 +289,9 @@ class NixlBaseConnectorScheduler:

289 and blocks

290 and (n_spec_blocks := self._ssm_spec_blocks[i]) is not None

291 ):

291 - if n_spec := min(n_spec_blocks, len(blocks) - 1):

292 - blocks = blocks[:-n_spec]

293 - if not self._ssm_state_slots_are_positional:

294 - # Never empty: downstream reads that as a full prefix hit.

295 - blocks = blocks[-1:]

292 + blocks = clip_ssm_state_blocks(

293 + blocks, n_spec_blocks, self._ssm_state_slots_are_positional

294 + )

295 clipped.append(blocks)

296 return tuple(clipped)

297