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

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

[PD][Nixl] Add support for hybrid SSM-FA models (#36687)

Nicolò Lucchesi committed Mar 16, 2026 at 19:58 UTC f5c081d4325536975f79720125af48f200bcac75

7 files changed +587 -166

tests/v1/kv_connector/nixl_integration/config_sweep_accuracy_test.sh

@@ -18,11 +18,19 @@ dp_ep_configs=(

18 "DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=1 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP1, D-DPEP=2 (TP=1)

19 "DP_EP=1 GPU_MEMORY_UTILIZATION=0.8 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 MODEL_NAMES=deepseek-ai/deepseek-vl2-tiny" # MLA+P-TP2, D-DPEP=2 (TP=1)

20 )

21 +hybrid_ssm_configs=(

22 + "ENABLE_HMA_FLAG=1 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code"

23 + # TODO: (NickLucche) Address async scheduling issue with TP>1 separately as this may impact other models.

24 + "ENABLE_HMA_FLAG=1 PREFILLER_TP_SIZE=2 DECODER_TP_SIZE=2 GPU_MEMORY_UTILIZATION=0.8 MODEL_NAMES=nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8 VLLM_SERVE_EXTRA_ARGS=--max-model-len,8192,--trust-remote-code,--no-async-scheduling"

25 +)

27 # Select config array based on DP_EP env var

28 if [[ -n "${DP_EP:-}" ]]; then

29 configs=("${dp_ep_configs[@]}")

30 echo "DP_EP is set, using dp_ep_configs"

31 +elif [[ -n "${HYBRID_SSM:-}" ]]; then

32 + configs=("${hybrid_ssm_configs[@]}")

33 + echo "HYBRID_SSM is set, using hybrid_ssm_configs."

34 else

35 configs=("${tp_configs[@]}")

36 fi

tests/v1/kv_connector/nixl_integration/test_accuracy.py

@@ -18,6 +18,7 @@ EXPECTED_VALUES = {

18 "deepseek-ai/deepseek-vl2-tiny": 0.19,

19 "deepseek-ai/DeepSeek-V2-Lite-Chat": 0.65,

20 "google/gemma-3-4b-it": 0.74,

21 + "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8": 0.84,

22 }

24 SIMPLE_PROMPT = (

tests/v1/kv_connector/unit/test_nixl_connector.py

+76 -45

@@ -53,7 +53,13 @@ from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend

53 from vllm.v1.attention.backends.utils import set_kv_cache_layout

54 from vllm.v1.engine import EngineCoreRequest

55 from vllm.v1.engine.output_processor import OutputProcessor

56 -from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheConfig, KVCacheTensor

56 +from vllm.v1.kv_cache_interface import (

57 + AttentionSpec,

58 + FullAttentionSpec,

59 + KVCacheConfig,

60 + KVCacheGroupSpec,

61 + KVCacheTensor,

62 +)

63 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput

64 from vllm.v1.request import RequestStatus

65 from vllm.v1.worker.kv_connector_model_runner_mixin import KVConnectorModelRunnerMixin

@@ -332,8 +338,20 @@ def test_kv_transfer_handshake(dist_init):

338

339 # Prefill connector will register KV cache to populate proper handshake

340 # metadata.

335 - # TODO this must match with values used in kv cache config

336 - kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)

341 + kv_cache_groups = [

342 + KVCacheGroupSpec(

343 + ["layer0", "layer1", "layer2"],

344 + FullAttentionSpec(

345 + block_size=16,

346 + num_kv_heads=4,

347 + head_size=16,

348 + dtype=torch.float16,

349 + ),

350 + )

351 + ]

352 + kv_cache_config = KVCacheConfig(

353 + num_blocks=2, kv_cache_tensors=[], kv_cache_groups=kv_cache_groups

354 + )

355 prefill_connector = NixlConnector(

356 vllm_config, KVConnectorRole.WORKER, kv_cache_config

357 )

@@ -437,7 +455,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

455 self.kv_cache_layout = kv_cache_layout

456 # Mock register_kv_caches attribute needed for tests that do not call it.

457 self.src_xfer_handles_by_block_size = {self.block_size: 1}

440 - test_shape = self.attn_backend.get_kv_cache_shape(

458 + test_shape = self.attn_backends[0].get_kv_cache_shape(

459 num_blocks=1, block_size=16, num_kv_heads=1, head_size=1

460 )

461 self.kv_topo = TpKVTopology(

@@ -447,7 +465,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

465 remote_block_size=self._block_size, # shared state

466 is_mla=self.use_mla,

467 total_num_kv_heads=self.model_config.get_total_num_kv_heads(),

450 - attn_backend=self.attn_backend,

468 + attn_backends=self.attn_backends,

469 tensor_shape=test_shape,

470 )

471

@@ -501,6 +519,7 @@ class FakeNixlConnectorWorker(NixlConnectorWorker):

519 # is started. We mock HND here.

520 kv_cache_layout="HND",

521 block_size=self.block_size,

522 + ssm_sizes=(0, 0),

523 ),

524 remote_tp_rank=remote_tp_rank,

525 remote_tp_size=remote_tp_size,

@@ -951,6 +970,7 @@ class TestNixlHandshake:

970 block_lens=worker.block_len_per_layer,

971 kv_cache_layout=mismatched_layout,

972 block_size=worker.block_size,

973 + ssm_sizes=(0, 0),

974 )

975

976 with pytest.raises(RuntimeError):

@@ -1006,6 +1026,7 @@ class TestNixlHandshake:

1026 block_lens=[i * 2 for i in worker.block_len_per_layer],

1027 kv_cache_layout="HND",

1028 block_size=worker.block_size,

1029 + ssm_sizes=(0, 0),

1030 )

1031

1032 # We don't check layout for homogeneous TP and MLA for now, as the

@@ -1496,9 +1517,47 @@ def test_register_kv_caches(

1517 # test run if not mocking.

1518 mock_get_attn_backend.return_value = backend_cls

1519 mock_get_attn_backends.return_value = [backend_cls]

1499 -

1520 + num_layers = 32

1521 + block_size = 16

1522 + num_blocks = 8

1523 + num_heads = 4

1524 + head_size = 16

1525 +

1526 + # TODO (NickLucche) the fact that connector depends on kv_cache_config for init

1527 + # but cross-layer preference cant be inferred prior to creating kv_cache_config

1528 + # is a bit awkward.

1529 + dummy_connector = NixlConnector(

1530 + vllm_config,

1531 + KVConnectorRole.WORKER,

1532 + make_kv_cache_config(block_size=block_size),

1533 + )

1534 + kv_cache_spec = FullAttentionSpec(

1535 + block_size=block_size,

1536 + num_kv_heads=num_heads,

1537 + head_size=head_size,

1538 + dtype=torch.float16,

1539 + )

1540 + if dummy_connector.prefer_cross_layer_blocks:

1541 + kv_cache_config = KVCacheConfig(

1542 + num_blocks=num_blocks,

1543 + kv_cache_tensors=[

1544 + KVCacheTensor(

1545 + size=kv_cache_spec.page_size_bytes * num_blocks,

1546 + shared_by=["all-layers"],

1547 + )

1548 + for _ in range(num_layers)

1549 + ],

1550 + kv_cache_groups=[KVCacheGroupSpec(["all-layers"], kv_cache_spec)],

1551 + )

1552 + else:

1553 + kv_cache_config = KVCacheConfig(

1554 + num_blocks=num_blocks,

1555 + kv_cache_tensors=[],

1556 + kv_cache_groups=[

1557 + KVCacheGroupSpec(["layer0", "layer1", "layer2"], kv_cache_spec)

1558 + ],

1559 + )

1560 # Create connector

1501 - kv_cache_config = make_kv_cache_config(block_size=16, num_blocks=2)

1561 connector = NixlConnector(vllm_config, KVConnectorRole.WORKER, kv_cache_config)

1562 connector.connector_worker = FakeNixlConnectorWorker(

1563 vllm_config,

@@ -1526,35 +1585,6 @@ def test_register_kv_caches(

1585 or connector.prefer_cross_layer_blocks

1586 )

1587 if connector.prefer_cross_layer_blocks:

1529 - num_layers = 32

1530 - block_size = 16

1531 - num_blocks = 8

1532 - # Keep the fake worker's expected num_blocks in sync with the

1533 - # cross-layer tensor we are about to register.

1534 - worker_kv_cache_config = make_kv_cache_config(

1535 - block_size=block_size, num_blocks=num_blocks

1536 - )

1537 - connector.connector_worker.kv_cache_config = worker_kv_cache_config

1538 - connector.connector_worker.num_blocks = worker_kv_cache_config.num_blocks

1539 - kv_cache_spec = AttentionSpec(

1540 - block_size=block_size,

1541 - num_kv_heads=4,

1542 - head_size=64,

1543 - dtype=torch.bfloat16,

1544 - )

1545 - kv_cache_config = KVCacheConfig(

1546 - num_blocks=num_blocks,

1547 - kv_cache_tensors=[

1548 - KVCacheTensor(

1549 - size=kv_cache_spec.page_size_bytes * num_blocks,

1550 - shared_by=["dummy-layer"],

1551 - )

1552 - for i in range(num_layers)

1553 - ],

1554 - # allocate_uniform_kv_caches does not use this

1555 - kv_cache_groups=[],

1556 - )

1557 -

1588 with set_current_vllm_config(vllm_config):

1589 _, cross_layers_kv_cache, _ = (

1590 KVConnectorModelRunnerMixin.allocate_uniform_kv_caches(

@@ -1586,12 +1616,8 @@ def test_register_kv_caches(

1616 expected_blocks_count = 8

1617

1618 kv_caches = {"all-layers": cross_layers_kv_cache}

1589 -

1619 else:

1620 # Create test kv cache tensors using proper backend shape

1592 - kv_cache_spec = cast(

1593 - AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec

1594 - )

1621 kv_cache_shape = backend_cls.get_kv_cache_shape(

1622 num_blocks=kv_cache_config.num_blocks,

1623 block_size=kv_cache_spec.block_size,

@@ -2261,7 +2287,7 @@ def test_compatibility_hash_validation(

2287 kv_cache_spec = cast(

2288 AttentionSpec, kv_cache_config.kv_cache_groups[0].kv_cache_spec

2289 )

2264 - kv_cache_shape = decode_worker.attn_backend.get_kv_cache_shape(

2290 + kv_cache_shape = decode_worker.attn_backends[0].get_kv_cache_shape(

2291 num_blocks=kv_cache_config.num_blocks,

2292 block_size=kv_cache_spec.block_size,

2293 num_kv_heads=kv_cache_spec.num_kv_heads,

@@ -2269,10 +2295,14 @@ def test_compatibility_hash_validation(

2295 )

2296 shared_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)

2297 unique_tensor = torch.zeros(*kv_cache_shape, dtype=kv_cache_spec.dtype)

2298 + # Build kv_caches from the actual layer names in kv_cache_config so that

2299 + # _layer_specs lookups in register_kv_caches always find a matching key.

2300 + layer_names = [

2301 + name for group in kv_cache_config.kv_cache_groups for name in group.layer_names

2302 + ]

2303 kv_caches = {

2273 - "layer0": shared_tensor,

2274 - "layer1": unique_tensor,

2275 - "layer2": shared_tensor,

2304 + name: shared_tensor if i % 2 == 0 else unique_tensor

2305 + for i, name in enumerate(layer_names)

2306 }

2307 decode_connector.register_kv_caches(kv_caches)

2308

@@ -2312,6 +2342,7 @@ def test_compatibility_hash_validation(

2342 block_lens=[4096 * prefill_block_size], # slot_size * block_size

2343 kv_cache_layout="HND",

2344 block_size=prefill_block_size,

2345 + ssm_sizes=(0, 0),

2346 )

2347 handshake_payload = NixlHandshakePayload(

2348 compatibility_hash=remote_hash,

@@ -2391,7 +2422,7 @@ def test_handshake_decode_errors(default_vllm_config, dist_init, error_scenario)

2422 remote_block_size=decode_worker._block_size, # shared state

2423 is_mla=decode_worker.use_mla,

2424 total_num_kv_heads=decode_worker.model_config.get_total_num_kv_heads(),

2394 - attn_backend=backend,

2425 + attn_backends=[backend],

2426 tensor_shape=test_shape,

2427 )

2428

tests/v1/kv_connector/unit/test_nixl_connector_hma.py

+112

@@ -74,6 +74,8 @@ def test_logical_to_kernel_block_ids_with_hma():

74 # Simulate HMA scenario: logical block size = 32, kernel block size = 16

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

76 worker._physical_blocks_per_logical_kv_block = 2

77 + # 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 # Test conversion: FA + SW group

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

@@ -201,3 +203,113 @@ def test_nixl_metadata_hma_block_ids_structure():

203 assert len(req_meta.remote.block_ids) == 2

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

205 assert list(req_meta.remote.block_ids[1]) == [18, 19, 20, 21]

206 +

207 +

208 +@pytest.mark.cpu_test

209 +def test_get_block_descs_ids_hybrid_ssm():

210 + """Test _get_block_descs_ids uses per-group strides for hybrid FA+SSM

211 + when ratio=1 (no kernel block size mismatch)."""

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

213 + NixlConnectorWorker,

214 + )

215 +

216 + worker = object.__new__(NixlConnectorWorker)

217 +

218 + num_blocks = 100

219 + engine_id = "test-engine"

220 + worker.num_regions = 2

221 + worker.dst_num_blocks = {engine_id: num_blocks}

222 + worker._has_mamba = True

223 + worker._is_mamba_group = [False, True]

224 + worker._physical_blocks_per_logical_kv_block = 1

225 + # num_descs = num_regions * num_blocks (no blocks_first doubling)

226 + worker.num_descs = 2 * num_blocks

227 +

228 + fa_blocks = [3, 5]

229 + ssm_blocks = [1, 2]

230 + result = worker._get_block_descs_ids(engine_id, (fa_blocks, ssm_blocks))

231 +

232 + # FA group: stride=num_blocks=100, offset=0

233 + # region0: [3, 5], region1: [103, 105]

234 + # SSM group: stride=logical_blocks=100 (=num_blocks/ratio=100/1),

235 + # offset=num_descs=200

236 + # region0: [201, 202], region1: [301, 302]

237 + expected = [3, 5, 103, 105, 201, 202, 301, 302]

238 + assert list(result) == expected, f"Expected {expected}, got {list(result)}"

239 +

240 +

241 +@pytest.mark.cpu_test

242 +def test_get_block_descs_ids_kernel_block_mismatch():

243 + """Test _get_block_descs_ids uses different strides for FA (kernel blocks)

244 + vs SSM (logical blocks) when ratio > 1."""

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

246 + NixlConnectorWorker,

247 + )

248 +

249 + worker = object.__new__(NixlConnectorWorker)

250 +

251 + ratio = 4

252 + logical_blocks = 100

253 + num_blocks = logical_blocks * ratio # 400 kernel blocks

254 + engine_id = "test-engine"

255 + worker.num_regions = 2

256 + worker.dst_num_blocks = {engine_id: num_blocks}

257 + worker._has_mamba = True

258 + worker._is_mamba_group = [False, True]

259 + worker._physical_blocks_per_logical_kv_block = ratio

260 + worker.num_descs = 2 * num_blocks # 800

261 +

262 + fa_blocks = [3, 7] # kernel-level block IDs

263 + ssm_blocks = [1, 2] # logical block IDs

264 + result = worker._get_block_descs_ids(engine_id, (fa_blocks, ssm_blocks))

265 +

266 + # FA group: stride=num_blocks=400, offset=0

267 + # region0: [3, 7], region1: [403, 407]

268 + # SSM group: stride=logical_blocks=400//4=100, offset=num_descs=800

269 + # region0: [801, 802], region1: [901, 902]

270 + expected = [3, 7, 403, 407, 801, 802, 901, 902]

271 + assert list(result) == expected, f"Expected {expected}, got {list(result)}"

272 +

273 +

274 +@pytest.mark.cpu_test

275 +def test_nixl_metadata_hybrid_ssm_block_ids():

276 + """Test NixlConnectorMetadata correctly stores block IDs for FA + SSM

277 + groups with different block counts (kernel mismatch active)."""

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

279 + NixlConnectorMetadata,

280 + )

281 +

282 + metadata = NixlConnectorMetadata()

283 +

284 + # FA: 8 kernel blocks (2 logical * ratio=4), SSM: 2 logical blocks

285 + fa_blocks = [0, 1, 2, 3, 4, 5, 6, 7]

286 + ssm_blocks = [0, 1]

287 +

288 + metadata.add_new_req_to_recv(

289 + request_id="test-req-hybrid",

290 + local_block_ids=(fa_blocks, ssm_blocks),

291 + kv_transfer_params={

292 + "remote_block_ids": ([10, 11, 12, 13, 14, 15, 16, 17], [20, 21]),

293 + "remote_engine_id": "remote-engine",

294 + "remote_request_id": "prefill-test-req-hybrid",

295 + "remote_host": "localhost",

296 + "remote_port": 1234,

297 + "tp_size": 1,

298 + },

299 + )

300 +

301 + assert "test-req-hybrid" in metadata.reqs_to_recv

302 + req_meta = metadata.reqs_to_recv["test-req-hybrid"]

303 +

304 + # Verify local block IDs: different lengths per group

305 + assert len(req_meta.local_block_ids) == 2

306 + assert list(req_meta.local_block_ids[0]) == fa_blocks

307 + assert list(req_meta.local_block_ids[1]) == ssm_blocks

308 + assert len(req_meta.local_block_ids[0]) != len(req_meta.local_block_ids[1])

309 +

310 + # Verify remote block IDs: same asymmetry preserved

311 + assert req_meta.remote is not None

312 + assert len(req_meta.remote.block_ids) == 2

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

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

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

vllm/distributed/kv_transfer/kv_connector/utils.py

+38 -8

@@ -16,10 +16,12 @@ from vllm.logger import init_logger

16 from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase

17 from vllm.platforms import current_platform

18 from vllm.v1.attention.backend import AttentionBackend

19 +from vllm.v1.kv_cache_interface import MambaSpec

20 from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput

22 if TYPE_CHECKING:

23 from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBase

24 + from vllm.v1.kv_cache_interface import KVCacheSpec

26 logger = init_logger(__name__)

@@ -328,22 +330,26 @@ class TpKVTopology:

330 remote_tp_size: dict[EngineId, int]

331 is_mla: bool

332 total_num_kv_heads: int

331 - attn_backend: type[AttentionBackend]

333 + attn_backends: list[type[AttentionBackend]]

334 engine_id: EngineId

335 remote_block_size: dict[EngineId, int]

336 tensor_shape: torch.Size | None = None

337 + is_mamba: bool = False

338

339 def __post_init__(self):

340 # Figure out whether the first dimension of the cache is K/V

341 # or num_blocks. This is used to register the memory regions correctly.

339 - _MOCK_BLOCK_SIZE = 16

340 - kv_cache_shape = self.attn_backend.get_kv_cache_shape(

341 - num_blocks=1, block_size=_MOCK_BLOCK_SIZE, num_kv_heads=1, head_size=1

342 - )

343 - logger.debug("Test kv_cache_shape: %s", kv_cache_shape)

342 + attn_backend = self.attn_backends[0]

343 + if not self.is_mamba:

344 + _MOCK_BLOCK_SIZE = 16

345 + kv_cache_shape: tuple[int, ...] = attn_backend.get_kv_cache_shape(

346 + num_blocks=1, block_size=_MOCK_BLOCK_SIZE, num_kv_heads=1, head_size=1

347 + )

348 + logger.debug("Test kv_cache_shape: %s", kv_cache_shape)

349 # Non-MLA backends caches have 5 dims [2, num_blocks, H,N,D],

350 # we just mock num_blocks to 1 for the dimension check below.

346 - self._is_kv_layout_blocks_first = (

351 + # Hybrid SSM models assume a single blocks_first layout

352 + self._is_kv_layout_blocks_first = self.is_mamba or (

353 len(kv_cache_shape) == 5 and kv_cache_shape[0] == 1

354 )

355

@@ -360,7 +366,7 @@ class TpKVTopology:

366 _MOCK_NUM_LAYERS = 80

367 kv_cache_shape = (_MOCK_NUM_LAYERS,) + kv_cache_shape

368 try:

363 - kv_cache_stride_order = self.attn_backend.get_kv_cache_stride_order(

369 + kv_cache_stride_order = attn_backend.get_kv_cache_stride_order(

370 include_num_layers_dimension=self._cross_layers_blocks

371 )

372 except (AttributeError, NotImplementedError):

@@ -483,6 +489,30 @@ class TpKVTopology:

489 remote_tp_size = self.remote_tp_size[remote_engine_id]

490 return self.get_target_remote_ranks(remote_tp_size)

491

492 + def get_transfer_cache_regions(

493 + self, cache: torch.Tensor, layer_spec: "KVCacheSpec"

494 + ) -> list[torch.Tensor] | torch.Tensor:

495 + """Return the cache tensor(s) to register as NIXL memory regions,

496 + also accounting for hybrid SSM models specificities.

497 + """

498 + if isinstance(layer_spec, MambaSpec):

499 + # Register the whole kv cache shared tensor, including SSM/Conv. This is

500 + # similar to FI with the difference that SSM/Conv have different sizes

501 + conv, ssm = cache

502 + return [conv]

503 +

504 + # Check may be hacky but it's matching `_update_hybrid_attention_mamba_layout`.

505 + if self.is_mamba and cache.shape[0] == 2:

506 + # When MAMBA is present, all backends are blocks first, so that blocks

507 + # can be shared between attention layers and mamba layers. Runner

508 + # `_update_hybrid_attention_mamba_layout` already adjusted strides

509 + # for FlashAttn-like backends so its num_blocks first.

510 + # Swap [2<>num_blocks] dims to get required layout for hybrid SSM.

511 + cache = cache.transpose(0, 1)

512 +

513 + # Regular case: backends like FA register K/V in separate regions

514 + return cache if self.split_k_and_v else [cache]

515 +

516

517 def get_current_attn_backends(

518 vllm_config: VllmConfig, layer_names: list[str] | None = None

vllm/distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py

+1 -1

@@ -564,7 +564,7 @@ class MooncakeConnectorWorker:

564 remote_block_size=self._block_size, # shared state

565 is_mla=self.use_mla,

566 total_num_kv_heads=self.model_config.get_total_num_kv_heads(),

567 - attn_backend=backend,

567 + attn_backends=[backend],

568 )

569

570 self.async_zmq_ctx = zmq.asyncio.Context()

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

+351 -112

@@ -59,7 +59,12 @@ from vllm.utils.network_utils import make_zmq_path, make_zmq_socket

59 from vllm.v1.attention.backend import AttentionBackend, AttentionMetadata

60 from vllm.v1.attention.backends.utils import get_kv_cache_layout

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

62 -from vllm.v1.kv_cache_interface import FullAttentionSpec, MambaSpec, SlidingWindowSpec

62 +from vllm.v1.kv_cache_interface import (

63 + FullAttentionSpec,

64 + MambaSpec,

65 + SlidingWindowSpec,

66 + UniformTypeKVCacheSpecs,

67 +)

68 from vllm.v1.worker.block_table import BlockTable

69 from vllm.v1.worker.utils import select_common_block_size

@@ -159,6 +164,7 @@ class NixlAgentMetadata:

164 block_lens: list[int]

165 kv_cache_layout: str

166 block_size: int

167 + ssm_sizes: tuple[int, int]

168

169

170 @dataclass

@@ -310,6 +316,15 @@ class NixlConnectorMetadata(KVConnectorMetadata):

316 class NixlConnector(KVConnectorBase_V1, SupportsHMA):

317 @property

318 def prefer_cross_layer_blocks(self) -> bool:

319 + if any(

320 + [

321 + isinstance(group.kv_cache_spec, MambaSpec)

322 + for group in self.kv_cache_config.kv_cache_groups

323 + ]

324 + ):

325 + # Hybrid SSM models do not yet support cross-layer layout

326 + return False

327 +

328 backend = get_current_attn_backend(self._vllm_config)

329 if backend.get_name() not in (

330 "FLASH_ATTN",

@@ -335,12 +350,9 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA):

350 kv_cache_config: "KVCacheConfig",

351 ):

352 super().__init__(vllm_config, role, kv_cache_config)

338 -

353 assert vllm_config.kv_transfer_config is not None

354 assert vllm_config.kv_transfer_config.engine_id is not None

341 - for group in kv_cache_config.kv_cache_groups:

342 - if isinstance(group.kv_cache_spec, MambaSpec):

343 - raise ValueError("NixlConnector does not support Mamba models.")

355 + self.kv_cache_config = kv_cache_config

356 self.engine_id: EngineId = vllm_config.kv_transfer_config.engine_id

357 self.kv_transfer_config = vllm_config.kv_transfer_config

358 if role == KVConnectorRole.SCHEDULER:

@@ -434,11 +446,7 @@ class NixlConnector(KVConnectorBase_V1, SupportsHMA):

446 self, kv_cache: torch.Tensor, attn_backend: type[AttentionBackend]

447 ):

448 assert self.connector_worker is not None

437 -

438 - cross_layer_name = "ALL_LAYERS"

439 - kv_caches = {cross_layer_name: kv_cache}

440 -

441 - self.connector_worker.register_kv_caches(kv_caches)

449 + self.connector_worker.register_cross_layers_kv_caches(kv_cache)

450

451 def set_host_xfer_buffer_ops(self, copy_operation: CopyBlocksOp):

452 assert self.connector_worker is not None

@@ -962,6 +970,40 @@ class NixlConnectorWorker:

970 )

971 )

972 self.kv_cache_config = kv_cache_config

973 + self._layer_specs = {

974 + layer: group.kv_cache_spec

975 + for group in kv_cache_config.kv_cache_groups

976 + for layer in group.layer_names

977 + }

978 + self.hma_group_size = len(kv_cache_config.kv_cache_tensors)

979 +

980 + # Mamba metadata

981 + self._is_mamba_group = [

982 + isinstance(group.kv_cache_spec, MambaSpec)

983 + for group in kv_cache_config.kv_cache_groups

984 + ]

985 + mamba_ssm_size = (0, 0)

986 + self._has_mamba = any(self._is_mamba_group)

987 + if self._has_mamba:

988 + assert self._is_hma_required

989 + mamba_spec = next(

990 + spec

991 + for spec in self._layer_specs.values()

992 + if isinstance(spec, MambaSpec)

993 + )

994 + conv_nbytes, ssm_nbytes = (

995 + torch.tensor([], dtype=mamba_spec.dtypes[0]).element_size(), # type: ignore[misc]

996 + torch.tensor([], dtype=mamba_spec.dtypes[1]).element_size(), # type: ignore[misc]

997 + )

998 + conv_shape, ssm_shape = (

999 + torch.Size(mamba_spec.shapes[0]),

1000 + torch.Size(mamba_spec.shapes[1]),

1001 + )

1002 + mamba_ssm_size = (

1003 + conv_shape.numel() * conv_nbytes,

1004 + ssm_shape.numel() * ssm_nbytes,

1005 + )

1006 + self._mamba_ssm_size = mamba_ssm_size

1007

1008 # Agent.

1009 non_ucx_backends = [b for b in self.nixl_backends if b != "UCX"]

@@ -1106,9 +1148,9 @@ class NixlConnectorWorker:

1148

1149 # Get the attention backend from the first layer

1150 # NOTE (NickLucche) models with multiple backends are not supported yet

1109 - self.attn_backend = get_current_attn_backend(vllm_config)

1151 + self.attn_backends = get_current_attn_backends(vllm_config)

1152 + self.backend_name = self.attn_backends[0].get_name()

1153

1111 - self.backend_name = self.attn_backend.get_name()

1154 self.kv_cache_layout = get_kv_cache_layout()

1155 self.host_buffer_kv_cache_layout = self.kv_cache_layout

1156 logger.info("Detected attention backend %s", self.backend_name)

@@ -1135,6 +1177,8 @@ class NixlConnectorWorker:

1177 def _sync_block_size_with_kernel(self) -> None:

1178 backends = get_current_attn_backends(self.vllm_config)

1179 kernel_block_size = select_common_block_size(self.block_size, backends)

1180 + # Number of blocks not accounting for kernel block mismatches

1181 + self._logical_num_blocks = self.num_blocks

1182 if self.block_size != kernel_block_size:

1183 logger.info_once(

1184 "User-specified logical block size (%s) does not match"

@@ -1428,9 +1472,19 @@ class NixlConnectorWorker:

1472

1473 fut.add_done_callback(request_ready)

1474

1475 + def register_cross_layers_kv_caches(self, kv_cache: torch.Tensor) -> None:

1476 + """Register a cross-layers KV cache tensor with NIXL.

1477 +

1478 + `use_uniform_kv_cache()` guarantees a single KV cache group whose

1479 + layers all share the same `AttentionSpec`, so any layer name from

1480 + `_layer_specs` yields the correct per-layer spec for `page_size_bytes`.

1481 + """

1482 + first_layer = next(iter(self._layer_specs))

1483 + # Forwarding a real layer name rather than a synthetic key

1484 + self.register_kv_caches({first_layer: kv_cache})

1485 +

1486 def register_kv_caches(self, kv_caches: dict[str, torch.Tensor]):

1487 """Register the KV Cache data in nixl."""

1433 -

1488 self.kv_topo = TpKVTopology(

1489 tp_rank=self.tp_rank,

1490 engine_id=self.engine_id,

@@ -1438,8 +1492,12 @@ class NixlConnectorWorker:

1492 remote_block_size=self._block_size, # shared state

1493 is_mla=self.use_mla,

1494 total_num_kv_heads=self.model_config.get_total_num_kv_heads(),

1441 - attn_backend=self.attn_backend,

1442 - tensor_shape=next(iter(kv_caches.values())).shape,

1495 + attn_backends=self.attn_backends,

1496 + # SSM States come in tuples (ssm, conv)

1497 + tensor_shape=next(iter(kv_caches.values())).shape

1498 + if not self._has_mamba

1499 + else None,

1500 + is_mamba=self._has_mamba,

1501 )

1502 self.compat_hash = compute_nixl_compatibility_hash(

1503 self.vllm_config, self.backend_name, self.kv_topo.cross_layers_blocks

@@ -1481,12 +1539,50 @@ class NixlConnectorWorker:

1539 # to better exploit the memory layout (ie num_blocks is the first dim).

1540 tensor_size_bytes = None

1541

1484 - # Enable different block lengths for different layers when MLA is used.

1542 + # Enable different block lengths for different layers *only* when MLA is used.

1543 + # This is not used for SSM layers, which use the counterpart `mamba_ssm_size`.

1544 self.block_len_per_layer = list[int]()

1545 for layer_name, cache_or_caches in xfer_buffers.items():

1487 - cache_list = (

1488 - cache_or_caches if self.kv_topo.split_k_and_v else [cache_or_caches]

1546 + # NOTE (NickLucche) Hybrid SSM models assume a layout that is similar to

1547 + # that of FI, with block laid out as in `get_backend_aware_kv_block_len`.

1548 + # However, physical page_size may differ when kernel requires a specific

1549 + # block size. This leads to SSM and FA layers having different num_blocks.

1550 + # `_physical_blocks_per_logical_kv_block` ratio is used to adjust for this.

1551 + layer_spec = self._layer_specs[layer_name]

1552 + if isinstance(layer_spec, UniformTypeKVCacheSpecs):

1553 + # MLA DSv32 Indexer case: UniformTypeKVCacheSpecs merges kv_cache_specs

1554 + layer_spec = layer_spec.kv_cache_specs[layer_name]

1555 + cache_list = self.kv_topo.get_transfer_cache_regions(

1556 + cache_or_caches, layer_spec

1557 + )

1558 + # `layer_spec.page_size_bytes` only accounts for logical page_size, that is

1559 + # the page_size assuming constant `self._logical_num_blocks`.

1560 + physical_page_size = (

1561 + layer_spec.page_size_bytes

1562 + if isinstance(layer_spec, MambaSpec)

1563 + else layer_spec.page_size_bytes

1564 + // self._physical_blocks_per_logical_kv_block

1565 )

1566 + # For when registering multiple tensors eg K/V in separate regions.

1567 + physical_page_size = physical_page_size // len(cache_list)

1568 + if self.kv_topo._cross_layers_blocks:

1569 + # When cross-layers blocks are used, multiply by number of layers

1570 + physical_page_size = physical_page_size * len(

1571 + self.kv_cache_config.kv_cache_tensors

1572 + )

1573 + num_blocks = (

1574 + self._logical_num_blocks

1575 + if isinstance(layer_spec, MambaSpec)

1576 + else self.num_blocks

1577 + )

1578 + # `page_size` accounts for physical blocks, st KVCache is always

1579 + # [`num_blocks` * `page_size`]

1580 + curr_tensor_size_bytes = num_blocks * physical_page_size

1581 + if tensor_size_bytes is None:

1582 + tensor_size_bytes = curr_tensor_size_bytes

1583 +

1584 + # TODO (NickLucche) we could eventually unify how we handle FA/FI regions,

1585 + # registering a single tensor for both K/V and splitting logically like FI.

1586 for cache in cache_list:

1587 base_addr = cache.data_ptr()

1588 if base_addr in seen_base_addresses:

@@ -1494,27 +1590,27 @@ class NixlConnectorWorker:

1590 # across groups. This results in skipping all tensors but the ones

1591 # pointed to by group0. Also, generally we will have more blocks

1592 # per tensor but fewer regions.

1593 + logger.debug("Skipping %s because it's already seen", layer_name)

1594 continue

1498 -

1595 logger.debug(

1596 "Registering layer %s with cache shape: %s", layer_name, cache.shape

1597 )

1598 seen_base_addresses.append(base_addr)

1503 - curr_tensor_size_bytes = cache.numel() * cache.element_size()

1504 -

1505 - if tensor_size_bytes is None:

1506 - tensor_size_bytes = curr_tensor_size_bytes

1599 + # Only record non-Mamba page sizes.

1600 + if isinstance(layer_spec, MambaSpec):

1601 + self.block_len_per_layer.append(

1602 + physical_page_size // self._physical_blocks_per_logical_kv_block

1603 + )

1604 + else:

1605 + self.block_len_per_layer.append(physical_page_size)

1606

1508 - assert cache.shape[0] == self.num_blocks, (

1607 + assert cache.shape[0] == num_blocks, (

1608 "All kv cache tensors must have the same number of blocks"

1609 )

1610

1512 - self.block_len_per_layer.append(

1513 - curr_tensor_size_bytes // self.num_blocks

1514 - )

1515 -

1611 if not self.use_mla:

1517 - # Different kv cache shape is not supported by HeteroTP

1612 + # Different kv cache shape is not supported by HeteroTP.

1613 + # This must also hold true for Mamba-like models.

1614 assert tensor_size_bytes == curr_tensor_size_bytes, (

1615 "All kv cache tensors must have the same size"

1616 )

@@ -1533,6 +1629,21 @@ class NixlConnectorWorker:

1629 self.kv_caches_base_addr[self.engine_id][self.tp_rank] = seen_base_addresses

1630 self.num_regions = len(caches_data)

1631

1632 + if self.kv_topo.is_kv_layout_blocks_first:

1633 + # NOTE (NickLucche) When FlashInfer is used, memory is registered

1634 + # with joint KV for each block. This minimizes the overhead in

1635 + # registerMem allowing faster descs queries. In order to be able to

1636 + # split on kv_heads dim as required by heterogeneous TP, one must

1637 + # be able to index K/V separately. Hence we double the number

1638 + # of 'virtual' regions here and halve `block_len` below.

1639 + # Similarly for Mamba layers, we register SSM+Conv as a single region and

1640 + # then duplicate it logically to be able to index SSM/Conv separately.

1641 + self.num_regions *= 2

1642 +

1643 + # TODO (NickLucche) Adapt to different descs views (engine_id->tp_rank) to

1644 + # support heterogeneous TP.

1645 + self.num_descs = self.num_regions * self.num_blocks

1646 +

1647 descs = self.nixl_wrapper.get_reg_descs(caches_data, self.nixl_memory_type)

1648 logger.debug("Registering descs: %s", caches_data)

1649 self.nixl_wrapper.register_memory(descs, backends=self.nixl_backends)

@@ -1542,17 +1653,21 @@ class NixlConnectorWorker:

1653 self.device_kv_caches = kv_caches

1654 self.dst_num_blocks[self.engine_id] = self.num_blocks

1655

1545 - if self.kv_topo.is_kv_layout_blocks_first:

1546 - # NOTE (NickLucche) When FlashInfer is used, memory is registered

1547 - # with joint KV for each block. This minimizes the overhead in

1548 - # registerMem allowing faster descs queries. In order to be able to

1549 - # split on kv_heads dim as required by heterogeneous TP, one must

1550 - # be able to index K/V separately. Hence we double the number

1551 - # of 'virtual' regions here and halve `block_len` below.

1552 - self.num_regions *= 2

1656 + if self._has_mamba:

1657 + logger.info(

1658 + "Hybrid SSM registration: num_blocks=%s, "

1659 + "logical_num_blocks=%s, ratio=%s, num_regions=%s, "

1660 + "num_descs=%s, mamba_ssm_size=%s, block_len_per_layer=%s",

1661 + self.num_blocks,

1662 + self._logical_num_blocks,

1663 + self._physical_blocks_per_logical_kv_block,

1664 + self.num_regions,

1665 + self.num_descs,

1666 + self._mamba_ssm_size,

1667 + set(self.block_len_per_layer),

1668 + )

1669

1670 # Register local/src descr for NIXL xfer.

1555 - self.seen_base_addresses = seen_base_addresses

1671 self.src_xfer_handles_by_block_size[self.block_size], self.src_blocks_data = (

1672 self.register_local_xfer_handler(self.block_size)

1673 )

@@ -1569,6 +1684,7 @@ class NixlConnectorWorker:

1684 if not self.use_host_buffer

1685 else self.host_buffer_kv_cache_layout,

1686 block_size=self.block_size,

1687 + ssm_sizes=self._mamba_ssm_size,

1688 )

1689 # Wrap metadata in payload with hash for defensive decoding

1690 assert self.compat_hash is not None

@@ -1594,40 +1710,65 @@ class NixlConnectorWorker:

1710 data copy correctness.

1711 """

1712 assert self.kv_topo is not None

1713 + kv_topo = self.kv_topo

1714

1715 block_size_ratio = self.block_size // block_size

1599 - blocks_data = []

1600 - for i, base_addr in enumerate(self.seen_base_addresses):

1601 - # The new block_len is using prefill block_len;

1602 - # and num_blocks is multiple with N

1603 - kv_block_len = (

1604 - self.get_backend_aware_kv_block_len(layer_idx=i) // block_size_ratio

1605 - )

1606 - block_len_per_layer = self.block_len_per_layer[i] // block_size_ratio

1607 - num_blocks = self.num_blocks * block_size_ratio

1608 - for block_id in range(num_blocks):

1609 - block_offset = block_id * block_len_per_layer

1610 - addr = base_addr + block_offset

1611 - # (addr, len, device id)

1612 - blocks_data.append((addr, kv_block_len, self.device_id))

1613 -

1614 - if self.kv_topo.is_kv_layout_blocks_first:

1615 - # Separate and interleave K/V regions to maintain the same

1616 - # descs ordering. This is needed for selecting contiguous heads

1617 - # when split across TP ranks.

1716 + blocks_data: list[tuple[int, int, int]] = []

1717 + local_base_addresses = self.kv_caches_base_addr[self.engine_id][self.tp_rank]

1718 +

1719 + def register_blocks(blocks_data: list[tuple[int, int, int]], mamba: bool):

1720 + for i, base_addr in enumerate(local_base_addresses):

1721 + # The new block_len is using prefill block_len;

1722 + # and num_blocks is multiple with N

1723 + kv_block_len = (

1724 + self.get_backend_aware_kv_block_len(

1725 + layer_idx=i, first_split=True, mamba_view=mamba

1726 + )

1727 + // block_size_ratio

1728 + )

1729 + # Jump one page_size, but ssm page_size may be bigger when kernel

1730 + # locks block size to a specific value.

1731 + block_len_per_layer = (

1732 + self.block_len_per_layer[i]

1733 + // block_size_ratio

1734 + * (1 if not mamba else self._physical_blocks_per_logical_kv_block)

1735 + )

1736 + num_blocks = self._logical_num_blocks if mamba else self.num_blocks

1737 + num_blocks = num_blocks * block_size_ratio

1738 for block_id in range(num_blocks):

1739 block_offset = block_id * block_len_per_layer

1740 addr = base_addr + block_offset

1621 - # Register addresses for V cache (K registered first).

1622 - v_addr = addr + kv_block_len

1623 - blocks_data.append((v_addr, kv_block_len, self.device_id))

1624 - logger.debug(

1625 - "Created %s blocks for src engine %s and rank %s on device id %s",

1626 - len(blocks_data),

1627 - self.engine_id,

1628 - self.tp_rank,

1629 - self.device_id,

1630 - )

1741 + # (addr, len, device id)

1742 + blocks_data.append((addr, kv_block_len, self.device_id))

1743 +

1744 + if kv_topo.is_kv_layout_blocks_first:

1745 + second_split = self.get_backend_aware_kv_block_len(

1746 + layer_idx=i, first_split=False, mamba_view=mamba

1747 + )

1748 + # Separate and interleave K/V regions to maintain the same

1749 + # descs ordering. This is needed for selecting contiguous heads

1750 + # when split across TP ranks.

1751 + for block_id in range(num_blocks):

1752 + block_offset = block_id * block_len_per_layer

1753 + addr = base_addr + block_offset

1754 + # Register addresses for V cache (K registered first).

1755 + v_addr = addr + kv_block_len

1756 + blocks_data.append((v_addr, second_split, self.device_id))

1757 + logger.debug(

1758 + "Created %s blocks for src engine %s and rank %s on device id %s",

1759 + len(blocks_data),

1760 + self.engine_id,

1761 + self.tp_rank,

1762 + self.device_id,

1763 + )

1764 +

1765 + register_blocks(blocks_data, mamba=False)

1766 + if self._has_mamba:

1767 + assert self.num_descs == len(blocks_data)

1768 + logger.debug(

1769 + "Registering additional %s local Mamba blocks", len(blocks_data)

1770 + )

1771 + register_blocks(blocks_data, mamba=True)

1772

1773 descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)

1774 # NIXL_INIT_AGENT to be used for preparations of local descs.

@@ -1708,7 +1849,8 @@ class NixlConnectorWorker:

1849 # local origin:| 0| 1| 8| 12|

1850 # local mapped:| 0| 1| 2| 3| 4| 5| 6| 7| 8| 9|10|11|12|13|14|15|

1851 assert self.kv_topo is not None

1711 - block_size_ratio = self.kv_topo.block_size_ratio_from_engine_id(engine_id)

1852 + kv_topo = self.kv_topo

1853 + block_size_ratio = kv_topo.block_size_ratio_from_engine_id(engine_id)

1854

1855 if engine_id not in self.dst_num_blocks:

1856 self.dst_num_blocks[engine_id] = nixl_agent_meta.num_blocks

@@ -1768,48 +1910,86 @@ class NixlConnectorWorker:

1910 # Eg. PTP1 DTP2 => P0 KV:[block0-KV_0 | block0-KV_1..].

1911

1912 # Register all remote blocks, but only the corresponding kv heads.

1771 - for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):

1772 - # Read our whole local region size from remote.

1773 - local_block_len = self.get_backend_aware_kv_block_len(layer_idx=i)

1774 - remote_kv_block_len = local_block_len // block_size_ratio

1775 - if block_size_ratio > 1:

1776 - # using remote kv_block_len as transfer unit

1777 - local_block_len = remote_kv_block_len

1913 + def register_remote_blocks(

1914 + blocks_data: list[tuple[int, int, int]], mamba: bool

1915 + ):

1916 + for i, base_addr in enumerate(nixl_agent_meta.kv_caches_base_addr):

1917 + # Read our whole local region size from remote.

1918 + local_block_len = self.get_backend_aware_kv_block_len(

1919 + layer_idx=i, first_split=True, mamba_view=mamba

1920 + )

1921 + remote_kv_block_len = local_block_len // block_size_ratio

1922 + if block_size_ratio > 1:

1923 + # using remote kv_block_len as transfer unit

1924 + local_block_len = remote_kv_block_len

1925 +

1926 + if tp_ratio < 0 and not self.use_mla:

1927 + # Remote tp is bigger: read a chunk of local region from remote

1928 + local_block_len = local_block_len // (-tp_ratio)

1929 + rank_offset = (

1930 + self.tp_rank % tp_ratio * remote_kv_block_len

1931 + if indexes_into_remote

1932 + else 0

1933 + )

1934

1779 - if tp_ratio < 0 and not self.use_mla:

1780 - # Remote tp is bigger: read a chunk of local region from remote

1781 - local_block_len = local_block_len // (-tp_ratio)

1782 - rank_offset = (

1783 - self.tp_rank % tp_ratio * remote_kv_block_len

1784 - if indexes_into_remote

1785 - else 0

1786 - )

1787 - for block_id in range(nixl_agent_meta.num_blocks):

1788 - block_offset = block_id * nixl_agent_meta.block_lens[i]

1789 - # For each block, grab the heads chunk belonging to rank_i

1790 - # of size remote_nheads // tp_ratio, which correspond to

1791 - # self.block_len == remote_block_len//tp_ratio bytes.

1792 - addr = base_addr + block_offset + rank_offset

1793 - # (addr, len, device id)

1794 - blocks_data.append((addr, local_block_len, nixl_agent_meta.device_id))

1795 -

1796 - if self.kv_topo.is_kv_layout_blocks_first:

1797 - # With FlashInfer index V separately to allow head splitting.

1798 - for block_id in range(nixl_agent_meta.num_blocks):

1799 - block_offset = block_id * nixl_agent_meta.block_lens[i]

1935 + # Assume same num_blocks for mamba and fa

1936 + num_blocks = (

1937 + nixl_agent_meta.num_blocks

1938 + if not mamba

1939 + else nixl_agent_meta.num_blocks

1940 + // self._physical_blocks_per_logical_kv_block

1941 + )

1942 + page_size = nixl_agent_meta.block_lens[i] * (

1943 + 1 if not mamba else self._physical_blocks_per_logical_kv_block

1944 + )

1945 + for block_id in range(num_blocks):

1946 + block_offset = block_id * page_size

1947 + # For each block, grab the heads chunk belonging to rank_i

1948 + # of size remote_nheads // tp_ratio, which correspond to

1949 + # self.block_len == remote_block_len//tp_ratio bytes.

1950 addr = base_addr + block_offset + rank_offset

1801 - v_addr = addr + nixl_agent_meta.block_lens[i] // 2

1951 + # (addr, len, device id)

1952 blocks_data.append(

1803 - (v_addr, local_block_len, nixl_agent_meta.device_id)

1953 + (addr, local_block_len, nixl_agent_meta.device_id)

1954 )

1955

1806 - logger.debug(

1807 - "Created %s blocks for dst engine %s with remote rank %s and local rank %s",

1808 - len(blocks_data),

1809 - engine_id,

1810 - remote_tp_rank,

1811 - self.tp_rank,

1812 - )

1956 + if kv_topo.is_kv_layout_blocks_first:

1957 + # With FlashInfer index V separately to allow head splitting.

1958 + second_split = self.get_backend_aware_kv_block_len(

1959 + layer_idx=i, first_split=False, mamba_view=mamba

1960 + )

1961 + # Apply the same scaling as local_block_len above for when we read

1962 + # a chunk of local V from `tp_ratio` separate remote workers.

1963 + if tp_ratio < 0 and not self.use_mla:

1964 + second_split = second_split // (-tp_ratio)

1965 + for block_id in range(num_blocks):

1966 + block_offset = block_id * page_size

1967 + addr = base_addr + block_offset + rank_offset

1968 + # Hop over the first split of remote page: either K or Conv.

1969 + if mamba:

1970 + v_addr = addr + nixl_agent_meta.ssm_sizes[0]

1971 + else:

1972 + v_addr = addr + nixl_agent_meta.block_lens[i] // 2

1973 + blocks_data.append(

1974 + (v_addr, second_split, nixl_agent_meta.device_id)

1975 + )

1976 +

1977 + logger.debug(

1978 + "Created %s blocks for dst engine %s"

1979 + " with remote rank %s and local rank %s",

1980 + len(blocks_data),

1981 + engine_id,

1982 + remote_tp_rank,

1983 + self.tp_rank,

1984 + )

1985 +

1986 + register_remote_blocks(blocks_data, mamba=False)

1987 + if self._has_mamba:

1988 + # Create extra descs for the Mamba "view" of the same KV cache tensors.

1989 + logger.debug(

1990 + "Registering additional %s remote Mamba blocks", len(blocks_data)

1991 + )

1992 + register_remote_blocks(blocks_data, mamba=True)

1993

1994 # Register with NIXL.

1995 descs = self.nixl_wrapper.get_xfer_descs(blocks_data, self.nixl_memory_type)

@@ -1849,6 +2029,9 @@ class NixlConnectorWorker:

2029 assert block_size_ratio == 1, (

2030 "HMA does not support different remote block size yet"

2031 )

2032 + # Mamba additional constraints

2033 + if self._has_mamba:

2034 + assert tp_ratio == 1, "Mamba does not support heterogeneous TP yet"

2035

2036 kv_cache_layout = (

2037 self.kv_cache_layout

@@ -2495,6 +2678,7 @@ class NixlConnectorWorker:

2678 A single flattened array is returned for all groups anyway.

2679 """

2680 region_ids = np.arange(self.num_regions)

2681 +

2682 # NOTE (NickLucche) With HMA, every kv group has the same number of layers and

2683 # layers from different groups share the same kv tensor.

2684 # eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be read across all regions,

@@ -2505,11 +2689,33 @@ class NixlConnectorWorker:

2689 if block_size_ratio is not None:

2690 num_blocks = int(num_blocks * block_size_ratio)

2691

2508 - # Compute the desc ids for each block.

2692 + # Compute desc ids per group using the right stride: FA descs have

2693 + # num_blocks entries per region (kernel granularity), SSM descs have

2694 + # logical_blocks entries per region (no kernel splitting).

2695 region_ids = region_ids[:, None]

2510 - block_ids = np.concatenate(block_ids)[None, :]

2511 - descs_ids = region_ids * num_blocks + block_ids

2512 - return descs_ids.flatten()

2696 + if not self._has_mamba:

2697 + block_ids = np.concatenate(block_ids)[None, :]

2698 + descs_ids = region_ids * num_blocks + block_ids

2699 + return descs_ids.flatten()

2700 + else:

2701 + # NOTE (NickLucche) SSM and Attention blocks regions can be exchanged

2702 + # arbitrarily by manager. Therefore, descs are duplicated for SSM and

2703 + # Attention like so:

2704 + # desc_handle->[descs_fa (all regions) | descs_ssm (all regions)].

2705 + # This is like having two "low-level views" of the same storage.

2706 + # `num_fa_descs` offset must be computed per-engine since P and D can

2707 + # have different num_blocks (and thus different FA descs counts).

2708 + ratio = self._physical_blocks_per_logical_kv_block

2709 + # SSM may register fewer num_blocks than FA

2710 + logical_blocks = num_blocks // ratio

2711 + num_fa_descs = self.num_regions * num_blocks

2712 + all_descs = []

2713 + for i, group in enumerate(block_ids):

2714 + stride = logical_blocks if self._is_mamba_group[i] else num_blocks

2715 + group_arr = np.asarray(group)[None, :]

2716 + offset = num_fa_descs if self._is_mamba_group[i] else 0

2717 + all_descs.append((region_ids * stride + group_arr + offset).flatten())

2718 + return np.concatenate(all_descs)

2719

2720 def _logical_to_kernel_block_ids(self, block_ids: BlockIds) -> BlockIds:

2721 """

@@ -2523,16 +2729,22 @@ class NixlConnectorWorker:

2729 block_arange = np.arange(0, self._physical_blocks_per_logical_kv_block).reshape(

2730 1, -1

2731 )

2732 + # Mamba blocks have no logical<>physical discrepancy

2733 + group_specs = self.kv_cache_config.kv_cache_groups

2734 return [

2735 BlockTable.map_to_kernel_blocks(

2736 np.array(group),

2737 self._physical_blocks_per_logical_kv_block,

2738 block_arange,

2739 ).tolist()

2532 - for group in block_ids

2740 + if not isinstance(group_specs[i].kv_cache_spec, MambaSpec)

2741 + else group

2742 + for i, group in enumerate(block_ids)

2743 ]

2744

2535 - def get_backend_aware_kv_block_len(self, layer_idx: int) -> int:

2745 + def get_backend_aware_kv_block_len(

2746 + self, layer_idx: int, first_split: bool = True, mamba_view: bool = False

2747 + ) -> int:

2748 """

2749 Get the block length for one K/V element (K and V have the same size).

2750

@@ -2540,11 +2752,38 @@ class NixlConnectorWorker:

2752 block, as K and V are in separate regions.

2753 For FlashInfer, this is half the length of the whole block, as K and V

2754 share the same region.

2755 + Similarly, for SSM-based models, state and conv are interleaved, but crucially

2756 + the their size differs.

2757 + Reference diagram:

2758 + KVCacheTensor (Shared)

2759 + / \

2760 + / \

2761 + / \

2762 + Attention (FlashInfer) View Mamba View

2763 + | |

2764 + | |

2765 + +-------------------+ +-------------------+

2766 + | KVCacheTensor | | KVCacheTensor |

2767 + | | | |

2768 + |<----- page ------>| |<----- page ------->|

2769 + | size | | size |

2770 + | Key 0 | Val 0 | |Conv 0 | SSM 0 |

2771 + | Key 1 | Val 1 | |Conv 1 | SSM 1 |

2772 + | ... | ... | | ... | ... |

2775 + +-------------------+ +--------------------+

2776 + |1st_split-2nd_split| |1st_split-2nd_split |

2777 """

2778 assert self.kv_topo is not None

2779 if self.kv_topo.is_kv_layout_blocks_first:

2780 # For indexing only half (either just the K or V part).

2547 - block_len = self.block_len_per_layer[layer_idx] // 2

2781 + if mamba_view:

2782 + # NOTE (NickLucche) Mamba Opt: this is already skipping the padding so

2783 + # we're only transferring the minimum required bytes.

2784 + block_len = self._mamba_ssm_size[not first_split]

2785 + else:

2786 + block_len = self.block_len_per_layer[layer_idx] // 2

2787 else:

2788 block_len = self.block_len_per_layer[layer_idx]

2789 return block_len