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 7, 12:45 UTC
[AMD][CI] Support Triton attention with ExampleConnector (#34931)
Signed-off-by: Ryan Rock <ryan.rock@amd.com>
Ryan Rock committed
Mar 1, 2026 at 01:58 UTC
87d319c52f22d3d08ef8ee49163aad9aad08f472
3 files changed
+28
-17
tests/v1/kv_connector/unit/test_example_connector.py
+11
-7
@@ -8,7 +8,7 @@ from PIL import Image
8
9
from vllm import LLM, EngineArgs, SamplingParams
10
from vllm.assets.image import ImageAsset
11
-from vllm.config import KVTransferConfig
11
+from vllm.config import AttentionConfig, KVTransferConfig
12
from vllm.multimodal.utils import encode_image_url
13
from vllm.platforms import current_platform
14
@@ -110,14 +110,17 @@ def process_prompt(processor, llm: LLM, question: str, image_urls: list[Image]):
110
print("-" * 50)
111
112
113
-@pytest.mark.skipif(
114
- current_platform.is_rocm(),
115
- reason=(
116
- "hipErrorLaunchFailure when running this test, see issue:"
117
- "https://github.com/ROCm/pytorch/issues/2822"
113
+@pytest.mark.parametrize(
114
+ "attn_backend",
115
+ (
116
+ ["FLASH_ATTN", "TRITON_ATTN"]
117
+ if current_platform.is_cuda()
118
+ else ["TRITON_ATTN"]
119
+ if current_platform.is_rocm()
120
+ else []
121
),
122
)
120
-def test_shared_storage_connector_hashes(tmp_path):
123
+def test_shared_storage_connector_hashes(tmp_path, attn_backend):
124
"""
125
Tests that ExampleConnector saves KV to the storage locations
126
with proper hashes; that are unique for inputs with identical text but
@@ -138,6 +141,7 @@ def test_shared_storage_connector_hashes(tmp_path):
141
max_model_len=8192,
142
max_num_seqs=1,
143
gpu_memory_utilization=0.4,
144
+ attention_config=AttentionConfig(backend=attn_backend),
145
enforce_eager=True,
146
kv_transfer_config=kv_transfer_config,
147
limit_mm_per_prompt={"image": 2},
tests/v1/kv_connector/unit/test_multi_connector.py
-8
@@ -20,7 +20,6 @@ from vllm.distributed.kv_transfer.kv_connector.v1.multi_connector import (
20
from vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector import (
21
NixlKVConnectorStats,
22
)
23
-from vllm.platforms import current_platform
23
24
MODEL_NAME = "meta-llama/Llama-3.2-1B-Instruct"
25
@@ -97,13 +96,6 @@ def _compare_directories(dir1: Path, dir2: Path) -> bool:
96
return True
97
98
100
-@pytest.mark.skipif(
101
- current_platform.is_rocm(),
102
- reason=(
103
- "hipErrorLaunchFailure when running this test, see issue:"
104
- "https://github.com/ROCm/pytorch/issues/2822"
105
- ),
106
-)
99
def test_multi_example_connector_consistency():
100
"""
101
Tests that MultiConnector with two ExampleConnectors saves
vllm/distributed/kv_transfer/kv_connector/v1/example_connector.py
+17
-2
@@ -17,6 +17,7 @@ from vllm.logger import init_logger
17
from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata
18
from vllm.utils.hashing import safe_hash
19
from vllm.v1.attention.backend import AttentionMetadata
20
+from vllm.v1.attention.backends.triton_attn import TritonAttentionMetadata
21
from vllm.v1.core.sched.output import SchedulerOutput
22
23
if TYPE_CHECKING:
@@ -118,12 +119,12 @@ class ExampleConnector(KVConnectorBase_V1):
119
The number of elements in kv_caches and layer_names should be
120
the same.
121
"""
121
- attn_metadata = forward_context.attn_metadata
122
123
def inject_kv_into_layer(
124
dst_kv_cache_layer: torch.Tensor,
125
src_kv_cache: torch.Tensor,
126
slot_mapping: torch.Tensor,
127
+ attn_metadata: AttentionMetadata,
128
) -> None:
129
"""Inject the KV cache into the layer.
130
@@ -145,6 +146,10 @@ class ExampleConnector(KVConnectorBase_V1):
146
num_pages * page_size, -1
147
)
148
dst_kv_cache_layer[slot_mapping, ...] = src_kv_cache
149
+ elif isinstance(attn_metadata, TritonAttentionMetadata):
150
+ block_idxs = slot_mapping // self._block_size
151
+ offsets = slot_mapping % self._block_size
152
+ dst_kv_cache_layer[block_idxs, :, offsets] = src_kv_cache
153
else:
154
num_pages = dst_kv_cache_layer_shape[1]
155
page_size = dst_kv_cache_layer_shape[2]
@@ -186,7 +191,13 @@ class ExampleConnector(KVConnectorBase_V1):
191
layer_name, request.token_ids, request.mm_hashes
192
)
193
kv_cache = safetensors.torch.load_file(filename)["kv_cache"].cuda()
189
- inject_kv_into_layer(kv_cache_layer, kv_cache, request.slot_mapping)
194
+ if isinstance(attn_metadata, dict):
195
+ inject_kv_into_layer(
196
+ kv_cache_layer,
197
+ kv_cache,
198
+ request.slot_mapping,
199
+ attn_metadata[layer_name],
200
+ )
201
202
def wait_for_layer_load(self, layer_name: str) -> None:
203
"""Blocking until the KV for a specific layer is loaded into vLLM's
@@ -229,6 +240,10 @@ class ExampleConnector(KVConnectorBase_V1):
240
if isinstance(attn_metadata, MLACommonMetadata):
241
num_pages, page_size = layer.shape[0], layer.shape[1]
242
return layer.reshape(num_pages * page_size, -1)[slot_mapping, ...]
243
+ elif isinstance(attn_metadata, TritonAttentionMetadata):
244
+ block_idxs = slot_mapping // self._block_size
245
+ offsets = slot_mapping % self._block_size
246
+ return layer[block_idxs, :, offsets]
247
num_pages, page_size = layer.shape[1], layer.shape[2]
248
return layer.reshape(2, num_pages * page_size, -1)[:, slot_mapping, ...]
249