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

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 6, 10:45 UTC

[SSM/Mamba] Follow-up: N-1 prefill for P/D disaggregation (#37310)

zhanqiuhu committed Mar 19, 2026 at 03:22 UTC d49f27314454afa687fb0323cfce7819e123c1c9

5 files changed +263 -13

tests/v1/kv_connector/unit/test_nixl_connector.py

+1 -1

@@ -2007,7 +2007,7 @@ def test_transfer_failure_logging(

2007 connector = NixlConnector(

2008 vllm_config,

2009 KVConnectorRole.WORKER,

2010 - make_kv_cache_config(block_size=16, hma_enabled=enable_hma),

2010 + make_kv_cache_config(block_size=16, swa_enabled=enable_hma),

2011 )

2012 connector.connector_worker = FakeNixlConnectorWorker(

2013 vllm_config,

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

+113 -8

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 -"""Unit tests for NixlConnectorScheduler sw_sizes calculation with HMA."""

3 +"""Unit tests for NixlConnectorScheduler with HMA and Mamba N-1 prefill."""

5 from unittest.mock import patch

@@ -14,24 +14,26 @@ from vllm.v1.core.single_type_kv_cache_manager import (

14 )

16 from .utils import (

17 + create_request,

18 create_vllm_config,

19 make_kv_cache_config,

20 + make_nixl_scheduler,

21 )

24 @pytest.mark.cpu_test

25 @pytest.mark.parametrize(

24 - "hma_enabled,expected_sw_sizes",

26 + "swa_enabled,expected_sw_sizes",

27 [

26 - # HMA enabled: FullAttentionSpec (0) + SlidingWindowSpec (2048/16=128)

28 + # SWA enabled: FullAttentionSpec (0) + SlidingWindowSpec (2048/16=128)

29 (True, [0, 128 + 1]),

28 - # HMA disabled: only FullAttentionSpec (0)

30 + # SWA disabled: only FullAttentionSpec (0)

31 (False, [0]),

32 ],

33 )

34 @patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.current_platform")

33 -def test_sw_sizes(mock_platform, hma_enabled, expected_sw_sizes):

34 - """Test sw_sizes is correctly computed based on HMA enabled/disabled."""

35 +def test_sw_sizes(mock_platform, swa_enabled, expected_sw_sizes):

36 + """Test sw_sizes is correctly computed based on SWA enabled/disabled."""

37 from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (

38 NixlConnectorScheduler,

39 )

@@ -42,7 +44,7 @@ def test_sw_sizes(mock_platform, hma_enabled, expected_sw_sizes):

44 vllm_config = create_vllm_config(block_size=block_size)

45 # SW 2048 tokens=>128 blocks

46 kv_cache_config = make_kv_cache_config(

45 - block_size=block_size, hma_enabled=hma_enabled, sw_size=2048

47 + block_size=block_size, swa_enabled=swa_enabled, sw_size=2048

48 )

50 scheduler = NixlConnectorScheduler(

@@ -75,7 +77,7 @@ def test_logical_to_kernel_block_ids_with_hma():

77 # So each logical block maps to 2 kernel blocks eg [0]->[0,1]

78 worker._physical_blocks_per_logical_kv_block = 2

79 # FA + SW groups (neither is MambaSpec, so both get expanded)

78 - worker.kv_cache_config = make_kv_cache_config(block_size=16, hma_enabled=True)

80 + worker.kv_cache_config = make_kv_cache_config(block_size=16, swa_enabled=True)

82 # Test conversion: FA + SW group

83 logical_block_ids = [[0, 1, 2], [3, 4]]

@@ -313,3 +315,106 @@ def test_nixl_metadata_hybrid_ssm_block_ids():

315 assert list(req_meta.remote.block_ids[0]) == [10, 11, 12, 13, 14, 15, 16, 17]

316 assert list(req_meta.remote.block_ids[1]) == [20, 21]

317 assert len(req_meta.remote.block_ids[0]) != len(req_meta.remote.block_ids[1])

318 +

319 +

320 +# ── Mamba N-1 prefill tests ──────────────────────────────────────────────

321 +

322 +

323 +@pytest.mark.cpu_test

324 +@pytest.mark.parametrize(

325 + "has_mamba,is_hma_required,expected_count",

326 + [

327 + (True, True, 9),

328 + (False, False, 10),

329 + (False, True, 10),

330 + ],

331 + ids=["mamba", "fa_only", "swa_only"],

332 +)

333 +def test_mamba_n1_d_side(has_mamba, is_hma_required, expected_count):

334 + """D-side: Mamba gets N-1 matched tokens, non-Mamba gets N."""

335 + sched = make_nixl_scheduler(has_mamba=has_mamba, is_hma_required=is_hma_required)

336 + req = create_request(num_tokens=10, do_remote_prefill=True)

337 +

338 + count, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0)

339 + assert count == expected_count

340 + assert is_async is True

341 +

342 +

343 +@pytest.mark.cpu_test

344 +def test_mamba_n1_p_side_truncation():

345 + """P-side: Mamba truncates prompt to N-1, sets max_tokens=1.

346 +

347 + Also verifies idempotency (calling again is a no-op) which is

348 + needed for preemption safety via the _p_side_truncated guard,

349 + and that non-Mamba models skip truncation entirely.

350 + """

351 + sched = make_nixl_scheduler(has_mamba=True, is_hma_required=True)

352 + req = create_request(num_tokens=10, do_remote_decode=True)

353 + req.max_tokens = 128

354 + original_len = len(req.prompt_token_ids)

355 +

356 + count, is_async = sched.get_num_new_matched_tokens(req, num_computed_tokens=0)

357 +

358 + assert count == 0

359 + assert is_async is False

360 + assert len(req.prompt_token_ids) == original_len - 1

361 + assert req.num_prompt_tokens == original_len - 1

362 + assert req.max_tokens == 1

363 + assert req.kv_transfer_params["_p_side_truncated"] is True

364 +

365 + # Idempotency: second call must not truncate further

366 + sched.get_num_new_matched_tokens(req, num_computed_tokens=0)

367 + assert len(req.prompt_token_ids) == original_len - 1

368 +

369 + # Non-Mamba: truncation is skipped

370 + fa_sched = make_nixl_scheduler(has_mamba=False, is_hma_required=False)

371 + fa_req = create_request(num_tokens=10, do_remote_decode=True)

372 + fa_original = len(fa_req.prompt_token_ids)

373 +

374 + fa_sched.get_num_new_matched_tokens(fa_req, num_computed_tokens=0)

375 + assert len(fa_req.prompt_token_ids) == fa_original

376 +

377 +

378 +@pytest.mark.cpu_test

379 +@pytest.mark.parametrize(

380 + "swa_enabled,mamba_enabled,expected_has_mamba,expected_is_hma",

381 + [

382 + (True, True, True, True),

383 + (True, False, False, True),

384 + (False, False, False, False),

385 + ],

386 + ids=["fa_swa_mamba", "fa_swa_only", "fa_only"],

387 +)

388 +@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.current_platform")

389 +def test_has_mamba_init(

390 + mock_platform,

391 + swa_enabled,

392 + mamba_enabled,

393 + expected_has_mamba,

394 + expected_is_hma,

395 +):

396 + """Test _has_mamba / _is_hma_required derived from kv_cache_groups."""

397 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (

398 + NixlConnectorScheduler,

399 + )

400 +

401 + mock_platform.device_type = "cpu"

402 +

403 + block_size = 16

404 + vllm_config = create_vllm_config(block_size=block_size)

405 + # VllmConfig.__post_init__ auto-disables HMA when kv_transfer_config

406 + # is set; override so we can test the scheduler's own derivation.

407 + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False

408 + kv_cache_config = make_kv_cache_config(

409 + block_size=block_size,

410 + swa_enabled=swa_enabled,

411 + mamba_enabled=mamba_enabled,

412 + )

413 +

414 + scheduler = NixlConnectorScheduler(

415 + vllm_config=vllm_config,

416 + engine_id="test-engine",

417 + kv_cache_config=kv_cache_config,

418 + )

419 + assert scheduler._has_mamba is expected_has_mamba

420 + assert scheduler._is_hma_required is expected_is_hma

tests/v1/kv_connector/unit/test_remote_prefill_lifecycle.py

+77 -1

@@ -1,10 +1,15 @@

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import copy

4 +from unittest.mock import patch

6 import pytest

7 -from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT, KVConnectorOutput

8 +from vllm.v1.outputs import (

9 + EMPTY_MODEL_RUNNER_OUTPUT,

10 + KVConnectorOutput,

11 + ModelRunnerOutput,

12 +)

13 from vllm.v1.request import FinishReason, RequestStatus

15 from .utils import (

@@ -13,6 +18,7 @@ from .utils import (

18 create_request,

19 create_scheduler,

20 create_vllm_config,

21 + make_kv_cache_config,

22 )

24 pytestmark = pytest.mark.cpu_test

@@ -579,3 +585,73 @@ def test_cannot_recv():

585 scheduler.update_from_output(scheduler_output, model_runner_output)

586 _ = scheduler.schedule()

587 assert_scheduler_empty(scheduler)

588 +

589 +

590 +@patch("vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.current_platform")

591 +def test_p_side_chunked_prefill_mamba(mock_platform):

592 + """P-side integration: Mamba N-1 truncation + chunked prefill completes.

593 +

594 + A 64-token P-side request is truncated to 63 by the N-1 fix, then

595 + chunked into two prefill steps (32 + 31) and finishes with

596 + LENGTH_CAPPED because max_tokens is set to 1.

597 + """

598 + mock_platform.device_type = "cpu"

599 +

600 + BATCH_SIZE = 32

601 + NUM_TOKENS = 64

602 + BLOCK_SIZE = 16

603 +

604 + vllm_config = create_vllm_config(

605 + max_num_batched_tokens=BATCH_SIZE,

606 + block_size=BLOCK_SIZE,

607 + )

608 + vllm_config.scheduler_config.disable_hybrid_kv_cache_manager = False

609 +

610 + kv_cache_config = make_kv_cache_config(

611 + block_size=BLOCK_SIZE,

612 + mamba_enabled=True,

613 + num_blocks=10000,

614 + )

615 +

616 + scheduler = create_scheduler(vllm_config, kv_cache_config=kv_cache_config)

617 +

618 + request = create_request(

619 + num_tokens=NUM_TOKENS,

620 + do_remote_decode=True,

621 + block_size=BLOCK_SIZE,

622 + )

623 + request.max_tokens = 128

624 + scheduler.add_request(request)

625 + request_id = request.request_id

626 +

627 + # ── Step 1: first chunk ──

628 + scheduler_output = scheduler.schedule()

629 +

630 + assert len(request.prompt_token_ids) == NUM_TOKENS - 1

631 + assert request.max_tokens == 1

632 + assert scheduler_output.num_scheduled_tokens[request_id] == BATCH_SIZE

633 + assert request.num_computed_tokens == BATCH_SIZE

634 +

635 + # Model returns no tokens for intermediate prefill chunk

636 + intermediate_output = ModelRunnerOutput(

637 + req_ids=[request.request_id],

638 + req_id_to_index={request.request_id: 0},

639 + sampled_token_ids=[[]],

640 + )

641 + scheduler.update_from_output(scheduler_output, intermediate_output)

642 +

643 + # ── Step 2: remaining chunk ──

644 + scheduler_output = scheduler.schedule()

645 +

646 + remaining = NUM_TOKENS - 1 - BATCH_SIZE # 31

647 + assert scheduler_output.num_scheduled_tokens[request_id] == remaining

648 + assert request.num_computed_tokens == NUM_TOKENS - 1

649 +

650 + # Prefill complete: model generates 1 decode token

651 + final_output = create_model_runner_output([request])

652 + engine_core_outputs = scheduler.update_from_output(scheduler_output, final_output)

653 +

654 + # max_tokens=1 → request finishes with LENGTH

655 + outputs = engine_core_outputs[0].outputs

656 + assert len(outputs) == 1

657 + assert outputs[0].finish_reason == FinishReason.LENGTH

tests/v1/kv_connector/unit/utils.py

+30 -2

@@ -37,6 +37,7 @@ from vllm.v1.kv_cache_interface import (

37 FullAttentionSpec,

38 KVCacheConfig,

39 KVCacheGroupSpec,

40 + MambaSpec,

41 SlidingWindowSpec,

42 )

43 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput

@@ -423,7 +424,8 @@ KVConnectorFactory.register_connector(

424

425 def make_kv_cache_config(

426 block_size: int,

426 - hma_enabled: bool = False,

427 + swa_enabled: bool = False,

428 + mamba_enabled: bool = False,

429 sw_size: int = 128,

430 num_blocks: int = 100,

431 ) -> KVCacheConfig:

@@ -438,7 +440,7 @@ def make_kv_cache_config(

440 ),

441 )

442 ]

441 - if hma_enabled:

443 + if swa_enabled:

444 kv_cache_groups.append(

445 KVCacheGroupSpec(

446 ["layer1", "layer3"],

@@ -451,6 +453,32 @@ def make_kv_cache_config(

453 ),

454 )

455 )

456 + if mamba_enabled:

457 + kv_cache_groups.append(

458 + KVCacheGroupSpec(

459 + ["mamba0", "mamba1"],

460 + MambaSpec(

461 + block_size=block_size,

462 + shapes=((16,), (16,)),

463 + dtypes=(torch.float16,),

464 + ),

465 + )

466 + )

467 return KVCacheConfig(

468 num_blocks=num_blocks, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups

469 )

470 +

471 +

472 +def make_nixl_scheduler(has_mamba: bool = False, is_hma_required: bool = False):

473 + """Create a NixlConnectorScheduler via __new__ (skipping __init__).

474 +

475 + Only sets the two flags needed by the N-1 prefill logic.

476 + """

477 + from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (

478 + NixlConnectorScheduler,

479 + )

480 +

481 + sched = object.__new__(NixlConnectorScheduler)

482 + sched._has_mamba = has_mamba

483 + sched._is_hma_required = is_hma_required

484 + return sched

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+42 -1

@@ -572,6 +572,10 @@ class NixlConnectorScheduler:

572 for g in kv_cache_config.kv_cache_groups

573 )

574 )

575 + self._has_mamba = any(

576 + isinstance(g.kv_cache_spec, MambaSpec)

577 + for g in kv_cache_config.kv_cache_groups

578 + )

579

580 logger.info("Initializing NIXL Scheduler %s", engine_id)

581 if vllm_config.scheduler_config.disable_hybrid_kv_cache_manager:

@@ -717,6 +721,39 @@ class NixlConnectorScheduler:

721 logger.warning("Connection listener got unexpected message %s", msg)

722 sock.send_multipart((identity, b"", encoded_data[target_tp_rank]))

723

724 + def _mamba_prefill_token_count(self, num_prompt_tokens: int) -> int:

725 + """D-side only. Returns N-1 for Mamba models since the decoder

726 + always recomputes the last token and must start from h(N-1)."""

727 + if self._has_mamba and num_prompt_tokens > 1:

728 + return num_prompt_tokens - 1

729 + return num_prompt_tokens

730 +

731 + def _truncate_mamba_request_for_prefill(self, request: "Request") -> None:

732 + """P-side only: drop the last prompt token so the prefiller computes

733 + h(N-1) instead of h(N). The decoder recomputes the last token to

734 + derive h(N) correctly.

735 +

736 + Guarded by ``_p_side_truncated`` to avoid repeated truncation if the

737 + request is preempted and rescheduled."""

738 + params = request.kv_transfer_params

739 + if (

740 + params is not None

741 + # Guard against repeated truncation after preemption/reschedule.

742 + and not params.get("_p_side_truncated")

743 + and request.num_prompt_tokens > 1

744 + ):

745 + if request.prompt_token_ids is not None:

746 + request.prompt_token_ids.pop()

747 + elif request.prompt_embeds is not None:

748 + request.prompt_embeds = request.prompt_embeds[:-1]

749 + else:

750 + return

751 +

752 + request._all_token_ids.pop()

753 + request.num_prompt_tokens -= 1

754 + request.max_tokens = 1

755 + params["_p_side_truncated"] = True

756 +

757 def get_num_new_matched_tokens(

758 self, request: "Request", num_computed_tokens: int

759 ) -> tuple[int, bool]:

@@ -746,10 +783,14 @@ class NixlConnectorScheduler:

783 if params is not None and params.get("do_remote_prefill"):

784 # Remote prefill: get all prompt blocks from remote.

785 token_ids = request.prompt_token_ids or []

749 - count = len(token_ids) - num_computed_tokens

786 + actual = self._mamba_prefill_token_count(len(token_ids))

787 + count = actual - num_computed_tokens

788 if count > 0:

789 return count, True

790

791 + if params is not None and params.get("do_remote_decode") and self._has_mamba:

792 + self._truncate_mamba_request_for_prefill(request)

793 +

794 # No remote prefill for this request.

795 return 0, False

796