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 10, 06:30 UTC
[Misc][PD] Fix `get_attn_backend` usage in transfer connectors (#31988)
Signed-off-by: NickLucche <nlucches@redhat.com>
Nicolò Lucchesi committed
Jan 12, 2026 at 18:10 UTC
5b681074119b970c2f99f8baea43f856cafc0251
4 files changed
+39
-20
tests/v1/kv_connector/unit/test_nixl_connector.py
+1
-1
@@ -1439,7 +1439,7 @@ def test_register_kv_caches(default_vllm_config, dist_init, attn_backend):
1439
patch(f"{nixl_module}.NixlWrapper") as mock_nixl_wrapper,
1440
patch(f"{nixl_module}.threading.Event"),
1441
patch(f"{nixl_module}.threading.Thread") as mock_thread,
1442
- patch(f"{nixl_module}.get_attn_backend") as mock_get_attn_backend,
1442
+ patch(f"{nixl_module}.get_current_attn_backend") as mock_get_attn_backend,
1443
):
1444
# Ensure get_attn_backend returns the correct value due to
1445
# _cached_get_attn_backend returning the backend from previous
vllm/distributed/kv_transfer/kv_connector/utils.py
+26
-2
@@ -6,13 +6,14 @@ KV cache helper for store.
6
7
from collections.abc import Iterator
8
from dataclasses import dataclass
9
-from typing import TYPE_CHECKING, Literal
9
+from typing import TYPE_CHECKING, Any, Literal, cast
10
11
import torch
12
13
-from vllm.config import get_current_vllm_config
13
+from vllm.config import VllmConfig, get_current_vllm_config, get_layers_from_vllm_config
14
from vllm.distributed.kv_transfer.kv_connector.factory import KVConnectorFactory
15
from vllm.logger import init_logger
16
+from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase
17
from vllm.v1.attention.backend import AttentionBackend
18
from vllm.v1.outputs import KVConnectorOutput, ModelRunnerOutput
19
@@ -433,3 +434,26 @@ class TpKVTopology:
434
) -> list[int]:
435
remote_tp_size = self.remote_tp_size[remote_engine_id]
436
return self.get_target_remote_ranks(remote_tp_size)
437
+
438
+
439
+def get_current_attn_backend(vllm_config: VllmConfig):
440
+ layer_type = cast(type[Any], AttentionLayerBase)
441
+ layers = get_layers_from_vllm_config(vllm_config, layer_type, None)
442
+ if layers:
443
+ backend = next(iter(layers.values())).get_attn_backend()
444
+ else:
445
+ # Fallback for tests, when static_forward_context is empty.
446
+ logger.debug(
447
+ "No layers found in the vLLM config. "
448
+ "Falling back to default attention backend."
449
+ )
450
+ from vllm.v1.attention.selector import get_attn_backend
451
+
452
+ backend = get_attn_backend(
453
+ head_size=vllm_config.model_config.get_head_size(),
454
+ dtype=vllm_config.model_config.dtype,
455
+ kv_cache_dtype=vllm_config.cache_config.cache_dtype,
456
+ block_size=vllm_config.cache_config.block_size,
457
+ use_mla=vllm_config.model_config.use_mla,
458
+ )
459
+ return backend
vllm/distributed/kv_transfer/kv_connector/v1/mooncake_connector.py
+7
-9
@@ -16,7 +16,10 @@ import zmq.asyncio
16
17
from vllm import envs
18
from vllm.config import VllmConfig
19
-from vllm.distributed.kv_transfer.kv_connector.utils import TpKVTopology
19
+from vllm.distributed.kv_transfer.kv_connector.utils import (
20
+ TpKVTopology,
21
+ get_current_attn_backend,
22
+)
23
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
24
KVConnectorBase_V1,
25
KVConnectorMetadata,
@@ -32,7 +35,6 @@ from vllm.logger import init_logger
35
from vllm.utils.network_utils import get_ip, make_zmq_path, make_zmq_socket
36
from vllm.v1.attention.backend import AttentionMetadata
37
from vllm.v1.attention.backends.utils import get_kv_cache_layout
35
-from vllm.v1.attention.selector import get_attn_backend
38
from vllm.v1.core.sched.output import SchedulerOutput
39
from vllm.v1.request import RequestStatus
40
@@ -468,13 +470,9 @@ class MooncakeConnectorWorker:
470
self.cache_config = vllm_config.cache_config
471
self.use_mla = self.model_config.use_mla
472
471
- backend = get_attn_backend(
472
- self.model_config.get_head_size(),
473
- self.model_config.dtype,
474
- self.cache_config.cache_dtype,
475
- self.block_size,
476
- use_mla=self.use_mla,
477
- )
473
+ # Get the attention backend from the first layer
474
+ # NOTE (NickLucche) models with multiple backends are not supported yet
475
+ backend = get_current_attn_backend(vllm_config)
476
self.backend_name = backend.get_name()
477
self.kv_cache_layout = get_kv_cache_layout()
478
logger.debug("Detected attention backend %s", self.backend_name)
vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py
+5
-8
@@ -24,6 +24,7 @@ from vllm.config import VllmConfig
24
from vllm.distributed.kv_transfer.kv_connector.utils import (
25
EngineId,
26
TpKVTopology,
27
+ get_current_attn_backend,
28
kv_postprocess_blksize_and_layout_on_receive,
29
kv_postprocess_blksize_on_receive,
30
kv_postprocess_layout_on_receive,
@@ -53,7 +54,6 @@ from vllm.platforms import current_platform
54
from vllm.utils.network_utils import make_zmq_path, make_zmq_socket
55
from vllm.v1.attention.backend import AttentionMetadata
56
from vllm.v1.attention.backends.utils import get_kv_cache_layout
56
-from vllm.v1.attention.selector import get_attn_backend
57
from vllm.v1.core.sched.output import SchedulerOutput
58
from vllm.v1.worker.block_table import BlockTable
59
@@ -957,13 +957,10 @@ class NixlConnectorWorker:
957
self.block_window_per_layer: list[int | None] = []
958
self.use_mla = self.model_config.use_mla
959
960
- backend = get_attn_backend(
961
- self.model_config.get_head_size(),
962
- self.model_config.dtype,
963
- self.cache_config.cache_dtype,
964
- self.block_size,
965
- use_mla=self.use_mla,
966
- )
960
+ # Get the attention backend from the first layer
961
+ # NOTE (NickLucche) models with multiple backends are not supported yet
962
+ backend = get_current_attn_backend(vllm_config)
963
+
964
self.backend_name = backend.get_name()
965
self.kv_cache_layout = get_kv_cache_layout()
966
self.host_buffer_kv_cache_layout = self.kv_cache_layout