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

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

[P/D][NixlConnector] Enable FlashInfer backend (#19090)

Nicolò Lucchesi committed Jun 5, 2025 at 19:10 UTC 9ef9173cfa338ea7742f23ee9d0038e580081390

2 files changed +51 -15

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

+50 -15

@@ -15,6 +15,7 @@ import torch

15 import zmq

17 from vllm import envs

18 +from vllm.attention.selector import backend_name_to_enum, get_attn_backend

19 from vllm.config import VllmConfig

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

21 KVConnectorBase_V1, KVConnectorMetadata, KVConnectorRole)

@@ -22,6 +23,7 @@ from vllm.distributed.parallel_state import (

23 get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size,

24 get_tp_group)

25 from vllm.logger import init_logger

26 +from vllm.platforms import _Backend

27 from vllm.utils import make_zmq_path, make_zmq_socket, round_down

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

29 from vllm.v1.request import RequestStatus

@@ -57,6 +59,7 @@ class NixlAgentMetadata(

59 num_blocks: int

60 tp_size: int

61 block_len: int

62 + attn_backend_name: str

65 @dataclass

@@ -384,11 +387,25 @@ class NixlConnectorWorker:

387

388 self.vllm_config = vllm_config

389 self.block_size = vllm_config.cache_config.block_size

390 + self.model_config = vllm_config.model_config

391 + self.cache_config = vllm_config.cache_config

392

393 # TODO(mgoin): remove this once we have hybrid memory allocator

394 # Optimization for models with local attention (Llama 4)

395 # List of block window sizes for each layer for local attention

396 self.block_window_per_layer: list[Optional[int]] = []

397 + self.use_mla = self.model_config.use_mla

398 +

399 + backend = get_attn_backend(self.model_config.get_head_size(),

400 + self.model_config.dtype,

401 + self.cache_config.cache_dtype,

402 + self.block_size,

403 + self.model_config.is_attention_free,

404 + use_mla=self.use_mla)

405 + self.backend_name = backend.get_name()

406 + attn_backend = backend_name_to_enum(self.backend_name)

407 + self._use_flashinfer = attn_backend == _Backend.FLASHINFER_VLLM_V1

408 + logger.debug("Detected attention backend %s", self.backend_name)

409

410 self._tp_size: dict[str, int] = {self.engine_id: self.world_size}

411 # With heterogeneous TP, P must wait for all assigned D TP workers to

@@ -472,12 +489,16 @@ class NixlConnectorWorker:

489 kv_elem_size = first_kv_cache.element_size()

490

491 # TODO(tms): Find a more robust way to detect and handle MLA

475 - self.use_mla = len(first_kv_cache.shape) == 3

492 # NOTE (NickLucche) To move blocks efficiently with NIXL, the expected

493 # KV memory layout is HND, as opposed to the default NHD. Note that it

494 # will only affects the strides. For MLA instead, we make require no

495 # such thing and resort to the standard layout.

480 - if self.use_mla:

496 + use_mla = len(first_kv_cache.shape) == 3

497 + assert use_mla == self.use_mla

498 +

499 + # TODO (NickLucche) not compatible with hybrid allocator. Enforce check

500 + # once it goes live, as a single kv layout is expected for xfers.

501 + if use_mla:

502 # MLA case.

503 self.num_blocks = first_kv_cache.shape[0]

504 block_rank = 2 # [block_size, latent_dim]

@@ -485,11 +506,16 @@ class NixlConnectorWorker:

506 block_size, kv_latent_dim = block_shape

507 self.slot_size_bytes = kv_elem_size * kv_latent_dim

508 else:

488 - # [2 (k and v), num_blocks, block_size, kv_heads, head_dim]

489 - self.num_blocks = first_kv_cache.shape[1]

490 - block_rank = 3 # [block_size, kv_heads, head_dim]

509 + # [2 (k and v), num_blocks, ...]

510 + if self._use_flashinfer:

511 + # FlashInfer swaps 2<->num_blocks dimensions.

512 + self.num_blocks = first_kv_cache.shape[0]

513 + block_rank = 4 # [2, block_size, kv_heads, head_dim]

514 + else:

515 + self.num_blocks = first_kv_cache.shape[1]

516 + block_rank = 3 # [block_size, kv_heads, head_dim]

517 block_shape = first_kv_cache.shape[-block_rank:]

492 - block_size, n_kv_heads, head_dim = block_shape

518 + block_size, n_kv_heads, head_dim = block_shape[-3:]

519 # head size in bytes.

520 self.slot_size_bytes = kv_elem_size * n_kv_heads * head_dim

521 assert block_size == self.block_size

@@ -497,12 +523,10 @@ class NixlConnectorWorker:

523 # hybrid attn, etc

524 # block size in bytes

525 self.block_len = kv_elem_size * math.prod(block_shape)

500 -

501 - logger.debug("Registering KV_Caches. use_mla: %s, shape %s",

502 - self.use_mla, first_kv_cache.shape)

503 - logger.debug("num_blocks: %s, block_shape: %s", self.num_blocks,

504 - block_shape)

505 - logger.debug("Per layer kv cache size: %s", first_kv_cache.shape)

526 + logger.info(

527 + "Registering KV_Caches: use_mla: %s, num_blocks: %s, "

528 + "block_shape: %s, per_layer_kv_cache_shape: %s", use_mla,

529 + self.num_blocks, block_shape, first_kv_cache.shape)

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

531 self.kv_caches = kv_caches

532 kv_caches_base_addr = []

@@ -514,9 +538,12 @@ class NixlConnectorWorker:

538 # are non-contiguous (it's not locally guaranteed that they will be)

539 # Disadvantage is that the encoded NixlAgentMetadata is now larger

540 # (roughly 8KB vs 5KB).

541 + # Conversely for FlashInfer, K and V are transferred in the same tensor

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

543 for cache_or_caches in kv_caches.values():

544 # Normalize to always be a list of caches

519 - cache_list = [cache_or_caches] if self.use_mla else cache_or_caches

545 + cache_list = [cache_or_caches] if use_mla or self._use_flashinfer \

546 + else cache_or_caches

547 for cache in cache_list:

548 base_addr = cache.data_ptr()

549 region_len = self.num_blocks * self.block_len

@@ -581,7 +608,8 @@ class NixlConnectorWorker:

608 kv_caches_base_addr=self.kv_caches_base_addr[self.engine_id],

609 num_blocks=self.num_blocks,

610 tp_size=self.world_size,

584 - block_len=self.block_len)

611 + block_len=self.block_len,

612 + attn_backend_name=self.backend_name)

613 ready_event = threading.Event()

614 self._nixl_handshake_listener_t = threading.Thread(

615 target=self._nixl_handshake_listener,

@@ -641,6 +669,10 @@ class NixlConnectorWorker:

669 assert self._tp_size[engine_id] == nixl_agent_meta.tp_size

670 else:

671 self._tp_size[engine_id] = nixl_agent_meta.tp_size

672 + # We may eventually enable this after asserting equality in cache

673 + # layout and close outputs.

674 + assert nixl_agent_meta.attn_backend_name == self.backend_name

675 +

676 self._remote_agents[engine_id][

677 remote_tp_rank] = self.nixl_wrapper.add_remote_agent(

678 nixl_agent_meta.agent_metadata)

@@ -659,13 +691,16 @@ class NixlConnectorWorker:

691 else:

692 remote_block_size = nixl_agent_meta.block_len // (

693 self.slot_size_bytes * tp_ratio)

694 + if self._use_flashinfer:

695 + # Account for joint KV in FlashInfer.

696 + remote_block_size //= 2

697

698 assert nixl_agent_meta.block_len == self.block_len * tp_ratio, (

699 "Remote P worker KV layer cache must be of shape [2, N, "

700 "local_kv_heads*tp_ratio, block_size, head_dim] and same dtype."

701 )

702

668 - assert self.block_size == remote_block_size, "Remote P worker with "

703 + assert self.block_size == remote_block_size, "Remote P worker with " \

704 "different block size is not supported"

705

706 assert self.num_blocks >= nixl_agent_meta.num_blocks

vllm/platforms/interface.py

@@ -47,6 +47,7 @@ class _Backend(enum.Enum):

47 ROCM_AITER_MLA_VLLM_V1 = enum.auto()

48 TORCH_SDPA = enum.auto()

49 FLASHINFER = enum.auto()

50 + FLASHINFER_VLLM_V1 = enum.auto()

51 TRITON_MLA = enum.auto() # Supported by V1

52 TRITON_MLA_VLLM_V1 = enum.auto()

53 FLASHMLA_VLLM_V1 = enum.auto()