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 11, 22:30 UTC
[CI/Build][AMD] Use ROCM_ATTN instead of FLASH_ATTN test for test_register_kv_caches for ROCm and update test for TRITON_ATTN (#29985)
Signed-off-by: Randall Smith <ransmith@amd.com> Co-authored-by: Randall Smith <ransmith@amd.com> Co-authored-by: TJian <tunjian.tan@embeddedllm.com>
rasmith committed
Dec 5, 2025 at 22:57 UTC
b12f4a983077f0f085e3734d4d5b0c25f2576cec
1 file changed
+37
-11
tests/v1/kv_connector/unit/test_nixl_connector.py
+37
-11
@@ -41,6 +41,7 @@ from vllm.distributed.kv_transfer.kv_transfer_state import (
41
has_kv_transfer_group,
42
)
43
from vllm.forward_context import ForwardContext
44
+from vllm.platforms import current_platform
45
from vllm.platforms.interface import Platform
46
from vllm.sampling_params import SamplingParams
47
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
@@ -1111,7 +1112,26 @@ def _run_abort_timeout_test(llm: LLM, timeout: int):
1112
llm.llm_engine.engine_core.shutdown()
1113
1114
1114
-@pytest.mark.parametrize("attn_backend", ["FLASH_ATTN", "TRITON_ATTN"])
1115
+@pytest.mark.parametrize(
1116
+ "attn_backend",
1117
+ [
1118
+ pytest.param(
1119
+ "FLASH_ATTN",
1120
+ marks=pytest.mark.skipif(
1121
+ current_platform.is_rocm(),
1122
+ reason="Attention backend FLASH_ATTN is not supported on ROCm",
1123
+ ),
1124
+ ),
1125
+ pytest.param(
1126
+ "ROCM_ATTN",
1127
+ marks=pytest.mark.skipif(
1128
+ not current_platform.is_rocm(),
1129
+ reason="Attention backend ROCM_ATTN is only supported on ROCm",
1130
+ ),
1131
+ ),
1132
+ "TRITON_ATTN",
1133
+ ],
1134
+)
1135
def test_register_kv_caches(dist_init, attn_backend, monkeypatch):
1136
"""
1137
Test that register_kv_caches() properly calls nixl_wrapper methods with
@@ -1133,6 +1153,10 @@ def test_register_kv_caches(dist_init, attn_backend, monkeypatch):
1153
from vllm.v1.attention.backends.flash_attn import FlashAttentionBackend
1154
1155
backend_cls = FlashAttentionBackend
1156
+ elif attn_backend == "ROCM_ATTN":
1157
+ from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
1158
+
1159
+ backend_cls = RocmAttentionBackend
1160
else: # TRITON_ATTN
1161
from vllm.v1.attention.backends.triton_attn import TritonAttentionBackend
1162
@@ -1151,6 +1175,7 @@ def test_register_kv_caches(dist_init, attn_backend, monkeypatch):
1175
}
1176
1177
# Store tensor info for validation
1178
+
1179
test_shape = backend_cls.get_kv_cache_shape(
1180
num_blocks=1, block_size=16, num_kv_heads=1, head_size=1
1181
)
@@ -1175,17 +1200,18 @@ def test_register_kv_caches(dist_init, attn_backend, monkeypatch):
1200
]
1201
expected_num_entries = 4
1202
1203
+ nixl_module = "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector"
1204
with (
1179
- patch(
1180
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.NixlWrapper"
1181
- ) as mock_nixl_wrapper,
1182
- patch(
1183
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.threading.Event"
1184
- ),
1185
- patch(
1186
- "vllm.distributed.kv_transfer.kv_connector.v1.nixl_connector.threading.Thread"
1187
- ) as mock_thread,
1188
- ): # noqa: E501
1205
+ patch(f"{nixl_module}.NixlWrapper") as mock_nixl_wrapper,
1206
+ patch(f"{nixl_module}.threading.Event"),
1207
+ patch(f"{nixl_module}.threading.Thread") as mock_thread,
1208
+ patch(f"{nixl_module}.get_attn_backend") as mock_get_attn_backend,
1209
+ ):
1210
+ # Ensure get_attn_backend returns the correct value due to
1211
+ # _cached_get_attn_backend returning the backend from previous
1212
+ # test run if not mocking.
1213
+ mock_get_attn_backend.return_value = backend_cls
1214
+
1215
# Create connector
1216
connector = NixlConnector(vllm_config, KVConnectorRole.WORKER)
1217
connector.connector_worker = FakeNixlConnectorWorker(