[ROCm][Performance] Coalesce contiguous SimpleCPUOffload DMA copies

[Fangzhou-Ai]

Purpose

SimpleCPUOffload submits one DMA descriptor per block per KV region even when consecutive mappings cover one contiguous span. A MiniMax-M3 TP4 rank with 118 KV regions and 256 blocks therefore submits 30,208 descriptors for about 0.99 GB of payload.

On ROCm, coalesce a run when source and destination block IDs both advance by one or both decrease by one. Split at gaps, repeated IDs, mismatched directions, and direction changes; normalize descending runs to their lower source and destination endpoints. Each merged range covers exactly the requested adjacent blocks. Payload bytes, destination mappings, region boundaries, streams, attributes, and the existing descriptor cap remain unchanged. CUDA retains its existing per-block descriptors.

Descending support matters for lazy offload: eviction order commonly produces paired descending physical block IDs, which the ascending-only implementation could not merge.

For every KV region, the merged run lengths partition the original block list, so sum(run_length * bytes_per_block) == block_count * bytes_per_block. A gap or turn starts a new run; the implementation never spans the gap. The byte-array tests also verify untouched destination sentinels. Coalescing therefore removes descriptors and does not add payload traffic.

Duplicate searches repeated on 2026-09-10 for SimpleCPUOffload coalescing, descending KV DMA, and hipMemcpyBatchAsync returned only this PR for this change. Related PRs operate at different layers: #54483 batches NIXL staging groups, #55841 changes block allocation preference, #53860 fixes memcpy attribute indices, #47324 fixes cross-stream synchronization, and #56170 changes a MiniMax-M3 scalar upload.

OpenAI Codex assisted with implementation, tests, benchmarking, analysis, and drafting. The submitting human will review every changed line and the validation evidence before merge.

Test Plan

.venv/bin/python -m pytest tests/v1/simple_kv_offload/test_hip_mem_ops.py -q
.venv/bin/python -m pre_commit run --files \
  vllm/v1/simple_kv_offload/cuda_mem_ops.py \
  tests/v1/simple_kv_offload/test_hip_mem_ops.py
.venv/bin/python -m pre_commit run mypy-3.12 --hook-stage manual --files \
  vllm/v1/simple_kv_offload/cuda_mem_ops.py \
  tests/v1/simple_kv_offload/test_hip_mem_ops.py

Reproducers, individual samples, byte checks, model outputs, and the sanitized C32 diagnostic are in the benchmark evidence.

Test Result

• All 49 copy-helper tests passed. They cover ROCm and unchanged CUDA behavior, paired ascending and descending runs, gaps, repeated IDs, direction changes, mismatched directions, heterogeneous region sizes, descriptor-cap chunking, and uint64 wrap boundaries. Byte-array tests verify every selected mapping and untouched destination sentinels.
• All applicable repository hooks passed, including Ruff, formatting, and the manual Python 3.12 mypy check.
• An additional 10,000 randomized mappings and an exhaustive emulation of 406,900 short mappings preserved the original mapping exactly.
• Real H2D/D2H byte and sentinel checks passed for every microbenchmark case and an additional heterogeneous-region case.
• OPT-125M on one MI355X with the current helper SHA produced 12/12 exact CPU-restore matches across three cache-reset rounds. Every round consumed 1,984 external KV tokens, with zero GPU prefix hits and zero preemptions. This is a serving correctness check, not a task-accuracy or throughput benchmark.
• CUDA hardware was not available in this AMD environment; the unchanged CUDA descriptor path is covered by the helper tests.

Copy microbenchmarks

One MI355X, ROCm 7.2.3 / HIP 7.2.53211, 118 regions, 256 blocks, 32 KiB per block per region, and 989,855,744 bytes per direction. Completion includes descriptor construction, submission, and stream synchronization; allocation, initialization, warmup, and validation are excluded.

Transfer pattern	Baseline completion	Coalesced completion	Speedup
H2D, ascending contiguous	409.41 ms	19.96 ms	20.51x
H2D, ascending runs of 8	402.52 ms	68.45 ms	5.88x
H2D, fragmented	414.31 ms	414.59 ms	1.00x
D2H, ascending contiguous	388.74 ms	20.96 ms	18.54x
D2H, ascending runs of 8	388.01 ms	68.46 ms	5.67x
D2H, fragmented	389.08 ms	389.19 ms	1.00x
H2D, descending contiguous	410.47 ms	19.81 ms	20.72x
D2H, descending contiguous	381.46 ms	20.60 ms	18.52x

The ascending rows compare the pre-PR helper with the intermediate ascending-only implementation. The descending rows compare that intermediate implementation with the current helper at commit ca560745bb / file SHA256 6c66dc25775e76aaa60a21e02151547c56a2a7380adcb6ebad41517cece6767e.

Bounded MiniMax-M3 C32 diagnostic

A separate TP4, 32-request natural-eviction diagnostic exercised the lazy descending restore path. Before and after transferred the same 16,128 blocks and 62,360,911,872 H2D bytes. Descending coalescing reduced submitted descriptors from 1,903,104 to 75,992, HIP calls from 252 to 32, and aggregate worker host-submission time from 23.367 s to 0.416 s. The restore phase fell from 8.141 s to 2.696 s and mean TTFT from 3.855 s to 0.853 s, with identical 516,096/521,128 CPU-cache recovery and zero preemptions. The fixed lazy restore duration was within 0.5% of eager (2.696 versus 2.684 s).

This was a bounded diagnostic run separately before and after on a broader experimental stack containing early-load/order changes. It is not a full InferenceX replay or model-quality evaluation; its answer-correctness diagnostic was false and was not used as an acceptance gate. The matched payload, restore counters, timing gates, and copy instrumentation passed.

[claude]

Claude Code Review

This pull request is from a fork — automated review is disabled. A repository maintainer can comment @claude review to run a one-time review.

[Fangzhou-Ai]

Hi @ivanium , please take a look at this PR and let's discuss how to effectively accelerate the ROCm MM3 cache offloading cc @ChuanLi1101