[CRL] Add noncontiguous KV transfer benchmark mode and P2P in-flight flow control#501
Conversation
leoda1
requested review from
MC952-arch,
aoyulong and
mikethegoblin
as code owners
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cbad0299ae
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
Pull request overview
This PR extends the KV transfer microbenchmark to support a new noncontiguous (PD-step-like) mode that models fragmented KV-cache block pools and sweeps request length, while refactoring connector implementations (NIXL/Mooncake/FlagCX) to support one-time pool setup plus per-step address derivation.
Changes:
- Add
--mode {uniform,noncontig}and a noncontig driver that builds fragmented per-step WR patterns over request lengths. - Introduce
StepPattern+ block-pool fragmentation model and per-step verification for noncontig transfers. - Update NIXL/Mooncake/FlagCX benchmark backends to support
setup_pool()/prepare_step()and (for FlagCX) switch to a P2P engine + batched write API.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: dc405ce804
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| self._alloc_pool(pool_bytes) | ||
| self.sock = _init_zmq(self.remote_ip, self.zmq_port, self.role) | ||
| self.agent = self._agent_cls(self.role, self._cfg_cls(backends=[self.backend])) | ||
| self.reg = self.agent.register_memory(self.agent.get_reg_descs([self.buffer])) |
There was a problem hiding this comment.
For --connector=nixl this registers the entire pool as one tensor, while the new benchmark defaults to an 8 GiB pool and intentionally scatters offsets across that full range. The FlagCX NIXL backend in this repo forwards each NIXL memory descriptor to a single flagcxP2pEngineReg(...), so --nixl-backend=FLAGCX bypasses the _mr_chunks() workaround used by Mooncake/FlagCX and reintroduces the max-MR-size failure described in this file: offsets above the NIC-registered region can be advertised in later descriptors but not be reachable. Split the NIXL registration into chunk tensors/descriptors as well, or cap the NIXL pool to a safe MR size.
Useful? React with 👍 / 👎.
| import json, socket | ||
| self._alloc_pool(pool_bytes) | ||
| self.sock = _init_zmq(self.remote_ip, self.zmq_port, self.role) | ||
| self.engine = self.flagcx.flagcxP2pEngineCreate() |
There was a problem hiding this comment.
When running FlagCX with --device=gpu --local-gpu-idx set to a non-current CUDA device, this creates the P2P engine before switching the process current device. flagcxP2pEngineCreate() snapshots inferLocalGpuIdx() from the current device and later uses it for topology/device selection during registration, so a buffer allocated on cuda:1 can be registered/routed as GPU 0. Set torch.cuda.set_device(self.gpu_idx) before engine creation/registration in this path.
Useful? React with 👍 / 👎.
| # server (receiver) # client (sender / initiator) | ||
| python kv_transfer_noncontiguous.py --connector=mooncake --role=server \\ | ||
| --remote-ip=10.8.2.169 | ||
| python kv_transfer_noncontiguous.py --connector=mooncake --role=client \\ | ||
| --remote-ip=10.8.2.169 |
| pool_blocks = pool_bytes // block_bytes | ||
| n = max(1, total_bytes // block_bytes) | ||
| if n > pool_blocks: | ||
| raise ValueError(f"need {n} blocks but pool holds {pool_blocks}") |
| def verify(self) -> None: | ||
| if self.is_source: | ||
| return | ||
| pat = self._pat | ||
| for i, (off, tag) in enumerate(zip(pat.dst_offsets, pat.tags)): | ||
| region = self.buffer[off:off + pat.block_bytes] | ||
| if not torch.all(region == tag).item(): | ||
| bad = int((region != tag).sum().item()) | ||
| raise AssertionError( | ||
| f"[{self.role}] VERIFY FAIL: WR {i} @dst {off} " | ||
| f"(tag={tag}) has {bad}/{pat.block_bytes} bytes wrong.") | ||
|
|
|
@MC952-arch please take a look! |
Co-authored-by: Copilot Autofix powered by AI <175728472+Copilot@users.noreply.github.com>
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1ec969303e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| engines fragment differently); offsets are distinct (no overlap). | ||
| """ | ||
| pool_blocks = pool_bytes // block_bytes | ||
| n = max(1, total_bytes // block_bytes) |
There was a problem hiding this comment.
When --sizes is not an exact multiple of --block-bytes, this floors the requested byte count and then reports/benchmarks n * block_bytes instead of the user-requested size; for example --sizes=10K --block-bytes=8192 transfers only 8 KiB, while sizes below one block transfer a full block. Since the CLI describes --sizes as total transfer sizes, these runs can silently produce bandwidth and correctness results for a different payload; reject non-multiples or round/print them consistently.
Useful? React with 👍 / 👎.
| self._wr = ([base + o for o in pat.src_offsets], | ||
| [rb + o for o in pat.dst_offsets], | ||
| [bs] * pat.n) | ||
| self.sock.send(b"READY") # barrier: receiver zeroed its pool |
There was a problem hiding this comment.
This READY send is only one-way: the source can queue it and enter run_transfer() before the receiver has finished _refresh() zeroing the large pool. In Mooncake/FlagCX write mode, running with --warmup=0 can let the first measured RDMA write race with the receiver's zero fill, so the first timing is taken before the destination is actually ready and --iters=1 can be verified after the zero overwrote the write; have the receiver acknowledge after zeroing before the source starts. The same one-way pattern is repeated in FlagCXTransport.prepare_step.
Useful? React with 👍 / 👎.
PR Category
CRL (Communication Runtime Layer)
PR Types
New Features
PR Description
Add noncontiguous KV transfer benchmark mode