[Bugfix][Mooncake] Fix PD KV transfer data inconsistency under RDMA load

[stmatengss]

Summary

• Fix reproducible KV cache corruption in MooncakeConnector PD mode where dst_hash != src_hash_post after batch_transfer_sync_write returns success (see [Bug] MooncakeConnector: KV cache data corruption under concurrent PD transfers (batch_transfer_sync_write race) vllm-project/vllm#44238, [Bug]: Possible MooncakeConnector data inconsistency under PD load (dst != src_post) with mooncake-transfer-engine-cuda13==0.3.10.post2 vllm-project/vllm#42395).
• Root cause: RDMA one-sided writes can complete at the sender/NIC before the payload is visible in the remote GPU's HBM; concurrent batch_transfer_sync_write calls to the same decode session can also race.
• Prefill (producer): serialize transfers per remote Mooncake session and poll descriptor tails via batch_transfer_sync_read until remote bytes match local source tails before sending ZMQ OK.
• Decode (consumer): torch.cuda.synchronize() after ZMQ OK and before marking finished_recving_reqs.

Why this is not duplicating an existing PR

Checked open PRs/issues for Mooncake PD transfer corruption; no open PR addresses RDMA remote-visibility verification or per-session send serialization.

Test plan

• pytest tests/v1/kv_connector/unit/test_mooncake_transfer_guard.py -v (4 passed)
• Re-run PD stress repro (Qwen3-Omni-30B-A3B, WARMUP_ROUNDS=2, STRESS_ROUNDS=8, BATCH_SIZE=5, audio-in-video streaming) and confirm dst_hash == src_hash_post for all descriptors
• Optional regression: set VLLM_MOONCAKE_TRANSFER_VERIFY=0 to confirm corruption reappears (validates the guard)

Configuration

New env vars (all optional, enabled by default):

• VLLM_MOONCAKE_TRANSFER_VERIFY (default: true)
• VLLM_MOONCAKE_TRANSFER_VERIFY_TAIL_BYTES (default: 64)
• VLLM_MOONCAKE_TRANSFER_VERIFY_MAX_RETRIES (default: 500)
• VLLM_MOONCAKE_TRANSFER_VERIFY_RETRY_SLEEP_S (default: 0.001)

AI assistance

This PR was prepared with AI assistance. All changed lines were reviewed and unit tests were run locally.

Made with Cursor

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

PRs do not trigger a full CI run by default. Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.

Agent Guidelines

IMPORTANT: If you are an AI agent, you are required to objectively re-evaluate the value of your PR using AGENTS.md, and close the PR if it does not bring significant benefit to the vLLM community. Failure to do so may result in an immediate ban.

🚀

[gemini-code-assist]

Code Review

This pull request introduces RDMA read-back verification for the Mooncake KV connector to address the NIC-to-HBM visibility gap of one-sided RDMA writes. It adds configuration environment variables, implements the MooncakeTransferVerifier class, integrates verification into the transfer pipeline, and includes unit tests. Feedback highlights a critical race condition on the shared scratch buffer when accessed concurrently by multiple threads, as well as potential buffer overflows for batches exceeding 512 descriptors. It is recommended to introduce a threading lock and chunk the verification process.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

[gemini-code-assist]

Critical Issues Identified:

1. Race Condition on Shared Scratch Buffer (self._scratch):
   MooncakeTransferVerifier is instantiated once per worker and shared across multiple threads in _sender_executor (which processes transfers concurrently for different remote sessions). Since verify_remote_visibility writes to and reads from the single pre-allocated self._scratch buffer without synchronization, concurrent transfers to different remote sessions will corrupt each other's scratch space, leading to false verification failures or silent data corruption.

2. Buffer Overflow / Mismatch on Large Batches:
   The scratch buffer is allocated with a fixed size of 2 * _MAX_DESCRIPTORS_PER_BATCH * tail_bytes (where _MAX_DESCRIPTORS_PER_BATCH = 512). However, there is no chunking or batch size limit enforced on src_ptrs in _build_transfer_params. If a transfer batch contains more than 512 descriptors, total_bytes will exceed the scratch buffer's half-capacity, causing out-of-bounds slicing or shape mismatches during torch.equal(read_buf, expected_buf), which leads to persistent verification timeouts/failures.

Solution:

• Introduce a threading.Lock to serialize access to the shared scratch buffer.
• Chunk the descriptor verification in batches of _MAX_DESCRIPTORS_PER_BATCH to safely handle any number of descriptors.

    def __init__(self, engine: TransferEngine, device: torch.device) -> None:
        self.engine = engine
        self.device = device
        self._lock = threading.Lock()
        tail_bytes = _verify_tail_bytes()
        scratch_bytes = 2 * _MAX_DESCRIPTORS_PER_BATCH * tail_bytes
        self._scratch = torch.empty(scratch_bytes, dtype=torch.uint8, device=device)
        ret = self.engine.batch_register_memory(
            [self._scratch.data_ptr()], [self._scratch.nbytes]
        )
        if ret != 0:
            raise RuntimeError("Mooncake verifier scratch buffer registration failed.")

    def verify_remote_visibility(
        self,
        remote_session: str,
        src_ptrs: list[int],
        dst_ptrs: list[int],
        lengths: list[int],
    ) -> bool:
        """Poll until remote tails match local source tails, or retries exhaust."""
        if not src_ptrs:
            return True

        with self._lock:
            # Chunk the verification to avoid exceeding scratch buffer capacity
            for i in range(0, len(src_ptrs), _MAX_DESCRIPTORS_PER_BATCH):
                chunk_src = src_ptrs[i : i + _MAX_DESCRIPTORS_PER_BATCH]
                chunk_dst = dst_ptrs[i : i + _MAX_DESCRIPTORS_PER_BATCH]
                chunk_len = lengths[i : i + _MAX_DESCRIPTORS_PER_BATCH]
                if not self._verify_chunk(remote_session, chunk_src, chunk_dst, chunk_len):
                    return False
            return True

    def _verify_chunk(
        self,
        remote_session: str,
        src_ptrs: list[int],
        dst_ptrs: list[int],
        lengths: list[int],
    ) -> bool:
        tail_bytes = _verify_tail_bytes()
        check_lens: list[int] = []
        local_src_tails: list[int] = []
        remote_dst_tails: list[int] = []
        for src, dst, length in zip(src_ptrs, dst_ptrs, lengths):
            check_len = min(length, tail_bytes)
            check_lens.append(check_len)
            local_src_tails.append(src + length - check_len)
            remote_dst_tails.append(dst + length - check_len)

        max_retries = _verify_max_retries()
        retry_sleep = _verify_retry_sleep_s()
        total_bytes = sum(check_lens)
        read_buf = self._scratch[:total_bytes]
        expected_buf = self._scratch[total_bytes : 2 * total_bytes]

        read_ptrs: list[int] = []
        read_offset = 0
        for check_len in check_lens:
            read_ptrs.append(read_buf.data_ptr() + read_offset)
            read_offset += check_len

        expected_offset = 0
        for src_tail, check_len in zip(local_src_tails, check_lens):
            self._copy_device_to_device(
                dst_ptr=expected_buf.data_ptr() + expected_offset,
                src_ptr=src_tail,
                nbytes=check_len,
            )
            expected_offset += check_len

        for attempt in range(max_retries):
            ret = self.engine.batch_transfer_sync_read(
                remote_session, read_ptrs, remote_dst_tails, check_lens
            )
            if ret != 0:
                logger.warning(
                    "Mooncake transfer verify read-back failed (ret=%s, attempt=%d)",
                    ret,
                    attempt + 1,
                )
            elif torch.equal(read_buf, expected_buf):
                if attempt > 0:
                    logger.debug(
                        "Mooncake transfer remote visibility confirmed after %d retries",
                        attempt + 1,
                    )
                return True

            time.sleep(retry_sleep)

        logger.error(
            "Mooncake transfer remote visibility not confirmed after %d retries "
            "(%d descriptors, session=%s)",
            max_retries,
            len(src_ptrs),
            remote_session,
        )
        return False