[ROCm][CI] Mirror the DSv4-Flash disaggregated DP EP group on MI355 (#58558)
Signed-off-by: Stefan Koncarevic <Stefan.Koncarevic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
stefankoncarevic committed
Sep 25, 2026 at 01:38 UTC
9ac3b7c69f8294ef0e80f03c2991b32ab8379f96
5 files changed
+164
-26
.buildkite/test-amd.yaml
+39
@@ -3814,6 +3814,45 @@ steps:
3814
not +quant_fp8 and not +rms_norm and DeepSeek-Coder and
3815
ROCM_AITER_MLA'
3816
3817
+- label: ":amd: (MI355) DSv4-Flash Disaggregated DP EP" # TBD
3818
+ timeout_in_minutes: 85
3819
+ mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355]
3820
+ dind: false
3821
+ agent_pool: mi355_8
3822
+ num_gpus: 8
3823
+ optional: true
3824
+ working_dir: "/vllm-workspace/tests"
3825
+ source_file_dependencies:
3826
+ - requirements/kv_connectors_rocm.txt
3827
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
3828
+ - vllm/model_executor/layers/fused_moe/
3829
+ - vllm/model_executor/layers/quantization/
3830
+ - vllm/models/deepseek_v4/
3831
+ - vllm/v1/attention/
3832
+ - tests/v1/kv_connector/nixl_integration/
3833
+ - vllm/_aiter_ops.py
3834
+ - vllm/platforms/rocm.py
3835
+ commands:
3836
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
3837
+ # A mirror does not inherit the parent env, so it is repeated here in
3838
+ # the parent's order. The layouts must stay empty: unset, the script
3839
+ # defaults to HND, which the V4 indexer rejects.
3840
+ - >-
3841
+ VLLM_ROCM_USE_AITER=1
3842
+ ENABLE_HMA_FLAG=1
3843
+ DP_EP=1
3844
+ GPU_MEMORY_UTILIZATION=0.85
3845
+ PREFILLER_TP_SIZE=4
3846
+ DECODER_TP_SIZE=4
3847
+ PREFILL_BLOCK_SIZE=256
3848
+ DECODE_BLOCK_SIZE=256
3849
+ MODEL_NAMES=deepseek-ai/DeepSeek-V4-Flash
3850
+ VLLM_ENGINE_READY_TIMEOUT_S=1800
3851
+ VLLM_SERVE_EXTRA_ARGS=--trust-remote-code,--kv-cache-dtype,fp8
3852
+ PREFILLER_KV_LAYOUT=
3853
+ DECODER_KV_LAYOUT=
3854
+ bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
3855
+
3856
#-------------------------------------------------------- mi355 · entrypoints --------------------------------------------------------#
3857
3858
- label: ":amd: (MI355) Entrypoints Integration (API Server) %N"
.buildkite/test_areas/disaggregated.yaml
+38
@@ -337,6 +337,44 @@ steps:
337
commands:
338
- bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh
339
- bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
340
+ mirror:
341
+ amd:
342
+ label: ":amd: (MI355) DSv4-Flash Disaggregated DP EP"
343
+ dind: false
344
+ device: mi355_8
345
+ timeout_in_minutes: 85
346
+ depends_on:
347
+ - image-build-amd
348
+ source_file_dependencies:
349
+ - requirements/kv_connectors_rocm.txt
350
+ - vllm/distributed/kv_transfer/kv_connector/v1/nixl/
351
+ - vllm/model_executor/layers/fused_moe/
352
+ - vllm/model_executor/layers/quantization/
353
+ - vllm/models/deepseek_v4/
354
+ - vllm/v1/attention/
355
+ - tests/v1/kv_connector/nixl_integration/
356
+ - vllm/_aiter_ops.py
357
+ - vllm/platforms/rocm.py
358
+ commands:
359
+ - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt
360
+ # A mirror does not inherit the parent env, so it is repeated here in
361
+ # the parent's order. The layouts must stay empty: unset, the script
362
+ # defaults to HND, which the V4 indexer rejects.
363
+ - >-
364
+ VLLM_ROCM_USE_AITER=1
365
+ ENABLE_HMA_FLAG=1
366
+ DP_EP=1
367
+ GPU_MEMORY_UTILIZATION=0.85
368
+ PREFILLER_TP_SIZE=4
369
+ DECODER_TP_SIZE=4
370
+ PREFILL_BLOCK_SIZE=256
371
+ DECODE_BLOCK_SIZE=256
372
+ MODEL_NAMES=deepseek-ai/DeepSeek-V4-Flash
373
+ VLLM_ENGINE_READY_TIMEOUT_S=1800
374
+ VLLM_SERVE_EXTRA_ARGS=--trust-remote-code,--kv-cache-dtype,fp8
375
+ PREFILLER_KV_LAYOUT=
376
+ DECODER_KV_LAYOUT=
377
+ bash v1/kv_connector/nixl_integration/run_accuracy_test.sh
378
379
- label: ":nvidia: (H200) Kimi-Linear-48B-A3B Disaggregated DP EP"
380
key: kimi-linear-disaggregated
tests/v1/kv_connector/unit/test_nixl_desc_geometry.py
+75
-22
@@ -13,6 +13,7 @@ mid-decode (silent corruption of an unrelated request).
13
14
from collections import defaultdict
15
from threading import Event, Lock
16
+from types import SimpleNamespace
17
from unittest.mock import MagicMock, patch
18
19
import numpy as np
@@ -120,16 +121,16 @@ def test_local_descriptors_follow_each_region_pool_capacity():
121
assert descriptors[:, 0].tolist() == [100, 116, 1000, 1016, 1032]
122
123
123
-@pytest.mark.cpu_test
124
-@pytest.mark.parametrize("push_pp", [False, True])
125
-def test_overlaid_transfer_groups_share_region_geometry(push_pp):
126
- """Groups overlaid on one allocation share its transfer region."""
127
- import msgspec
124
+def _register_overlaid_mla_worker(
125
+ *, push_pp: bool = False, tail_bytes: int = 0, page_covers_view: bool = True
126
+):
127
+ """Register two MLA layers overlaid on one allocation.
128
129
+ ``tail_bytes`` leaves spare bytes past the last block, the way an allocation
130
+ rounded up to a page boundary does. Clearing ``page_covers_view`` narrows
131
+ each layer's view below its page, leaving the block interior non-contiguous.
132
+ """
133
from vllm.distributed.kv_transfer.kv_connector.v1.nixl import base_worker as bw
130
- from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
131
- NixlAgentMetadata,
132
- )
134
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.push_worker import (
135
NixlPushConnectorWorker,
136
)
@@ -152,10 +153,12 @@ def test_overlaid_transfer_groups_share_region_geometry(push_pp):
153
)
154
page_size = spec.page_size_bytes
155
block_stride = 2 * page_size
155
- backing = torch.zeros(num_blocks, block_stride, dtype=torch.uint8)
156
+ view_size = page_size if page_covers_view else page_size - 8
157
+ allocation = torch.zeros(num_blocks * block_stride + tail_bytes, dtype=torch.uint8)
158
+ backing = allocation[: num_blocks * block_stride].view(num_blocks, block_stride)
159
caches = {
157
- "layer.0": backing[:, :page_size],
158
- "layer.1": backing[:, :page_size],
160
+ "layer.0": backing[:, :view_size],
161
+ "layer.1": backing[:, :view_size],
162
}
163
groups = [KVCacheGroupSpec([layer_name], spec) for layer_name in caches]
164
@@ -203,14 +206,9 @@ def test_overlaid_transfer_groups_share_region_geometry(push_pp):
206
worker._logical_num_blocks = num_blocks
207
worker.region_mem_types = []
208
worker.region_group_ids = []
206
- worker.region_mem_types = []
209
worker._mixed_mem_types = False
210
worker.region_names = []
211
worker.region_num_blocks = []
210
- worker._mixed_mem_types = False
211
- worker._desc_is_dram_by_block_size = {}
212
- worker._desc_pos_by_block_size = {}
213
- worker._dram_src_handles_by_block_size = {}
212
worker._transfer_layer_names = ()
213
worker._transfer_layer_region_indices = ()
214
worker._transfer_layer_group_ids = ()
@@ -231,7 +229,7 @@ def test_overlaid_transfer_groups_share_region_geometry(push_pp):
229
num_blocks=num_blocks,
230
kv_cache_tensors=[
231
KVCacheTensor(
234
- size=backing.nbytes,
232
+ size=allocation.nbytes,
233
layers=[name],
234
layer_stride=page_size,
235
block_stride=block_stride,
@@ -241,21 +239,43 @@ def test_overlaid_transfer_groups_share_region_geometry(push_pp):
239
kv_cache_groups=groups,
240
)
241
244
- transfer_topology = MagicMock()
245
-
242
with (
247
- patch.object(bw, "TransferTopology", return_value=transfer_topology),
243
+ patch.object(bw, "TransferTopology", return_value=MagicMock()),
244
patch.object(bw, "compute_nixl_compatibility_hash", return_value="hash"),
245
):
246
worker.register_kv_caches(caches)
247
248
+ return SimpleNamespace(
249
+ worker=worker,
250
+ allocation=allocation,
251
+ num_blocks=num_blocks,
252
+ block_stride=block_stride,
253
+ )
254
+
255
+
256
+@pytest.mark.cpu_test
257
+@pytest.mark.parametrize("push_pp", [False, True])
258
+def test_overlaid_transfer_groups_share_region_geometry(push_pp):
259
+ """Groups overlaid on one allocation share its transfer region."""
260
+ import msgspec
261
+
262
+ from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
263
+ NixlAgentMetadata,
264
+ )
265
+
266
+ registered = _register_overlaid_mla_worker(push_pp=push_pp)
267
+ worker = registered.worker
268
+ allocation = registered.allocation
269
+ num_blocks = registered.num_blocks
270
+ block_stride = registered.block_stride
271
+
272
assert worker.region_group_ids == [-1]
273
assert worker.block_stride_per_layer == [block_stride]
274
assert worker.nixl_wrapper.registered[0][0] == [
255
- (backing.data_ptr(), backing.nbytes, 0, "")
275
+ (allocation.data_ptr(), allocation.nbytes, 0, "")
276
]
277
expected_addrs = [
258
- backing.data_ptr() + block * block_stride for block in range(num_blocks)
278
+ allocation.data_ptr() + block * block_stride for block in range(num_blocks)
279
]
280
num_desc_regions = 2 if push_pp else 1
281
assert worker.src_blocks_data[:, 0].tolist() == expected_addrs * num_desc_regions
@@ -275,6 +295,39 @@ def test_overlaid_transfer_groups_share_region_geometry(push_pp):
295
assert worker._block_ids_by_region(([0], [2]), worker.region_group_ids) == [[0, 2]]
296
297
298
+def _descriptor_geometry(registered) -> dict:
299
+ """The geometry handed to NIXL, addressed relative to the allocation."""
300
+ worker = registered.worker
301
+ base = registered.allocation.data_ptr()
302
+ return {
303
+ "block_len_per_layer": list(worker.block_len_per_layer),
304
+ "block_stride_per_layer": list(worker.block_stride_per_layer),
305
+ "desc_offsets": [addr - base for addr in worker.src_blocks_data[:, 0].tolist()],
306
+ "desc_lens": worker.src_blocks_data[:, 1].tolist(),
307
+ }
308
+
309
+
310
+@pytest.mark.cpu_test
311
+@pytest.mark.parametrize(
312
+ "page_covers_view", [True, False], ids=["dense-page", "narrow-view"]
313
+)
314
+def test_registration_ignores_a_sub_block_padding_tail(page_covers_view):
315
+ """A tail shorter than one block must not move a single descriptor.
316
+
317
+ An allocation rounded up to a page boundary keeps a few spare bytes past the
318
+ last block, which cannot hold another block. The narrow view is the branch
319
+ DeepSeek-V4-Flash's sliding-window cache lands in.
320
+ """
321
+ unpadded = _register_overlaid_mla_worker(page_covers_view=page_covers_view)
322
+ # More spare bytes than blocks, so dividing the padded length by the block
323
+ # count overshoots the stride and would shift every descriptor.
324
+ padded = _register_overlaid_mla_worker(
325
+ tail_bytes=2 * unpadded.num_blocks, page_covers_view=page_covers_view
326
+ )
327
+
328
+ assert _descriptor_geometry(padded) == _descriptor_geometry(unpadded)
329
+
330
+
331
def _make_mla_hybrid_worker(local_block_size, kernel_block_size, num_logical_blocks):
332
"""Build a real pull worker with a hybrid MLA + 2xKDA HMA layout."""
333
from vllm.distributed.kv_transfer.kv_connector.v1.nixl import (
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py
+10
-2
@@ -1567,7 +1567,13 @@ class NixlBaseConnectorWorker:
1567
block_stride = physical_page_size
1568
else:
1569
block_stride = cache.stride(0) * cache.element_size()
1570
- storage_is_block_major = num_blocks * block_stride == registration_len
1570
+ # A connector that registers the backing as one memory region
1571
+ # may need its length rounded up to a page boundary. Such a tail
1572
+ # is shorter than one block, so it cannot make the storage
1573
+ # segmented.
1574
+ storage_is_block_major = (
1575
+ 0 <= registration_len - num_blocks * block_stride < block_stride
1576
+ )
1577
# A view with a contiguous page can own a strided region even
1578
# when sibling layers share the same backing allocation.
1579
page_contiguous = (
@@ -1589,7 +1595,9 @@ class NixlBaseConnectorWorker:
1595
):
1596
# TODO(Lucas): handle TP slicing for packed_storage; for now
1597
# restrict to MLA (DSv4) where kv is replicated.
1592
- storage_block_len = registration_len // num_blocks
1598
+ # The complete storage row is one block stride. Dividing the
1599
+ # registered length instead would fold the padding tail in.
1600
+ storage_block_len = block_stride
1601
region_specs = [
1602
(registration_base, storage_block_len, storage_block_len)
1603
]
vllm/v1/worker/utils.py
+2
-2
@@ -412,8 +412,8 @@ def allocate_kv_cache(
412
warmup_rocm_skinny_gemm_workspaces(device)
413
# Pad to the page granularity MoRIIO needs to register the shared
414
# backing as a single RDMA memory region. Other platforms keep the
415
- # exact-size allocation: NIXL and SimpleCPUOffload rely on
416
- # storage.nbytes() matching the logical KV size (see #53974).
415
+ # exact-size allocation (see #53974), so anything reading
416
+ # storage.nbytes() has to tolerate the tail on ROCm alone.
417
page_size = 4096
418
buf_size = ((raw_size + page_size - 1) // page_size) * page_size
419
else: