kvcache-ai / kvcache-ai/Mooncake
[TE][UB]: Auto-chunk registrations larger than max_seg_size and split cross-segment transfers
- Dominant language
- C++
- Stars
- 6.6k
- Forks
- 1.2k
- Avg merge
- 3d 5h
- Merged PRs (30d)
- 312
Description
Parent: #3500
## Problem
UB is behind both EFA/CXI and RDMA.
1. **Registration still shrinks.** `UrmaContext::registerMemoryRegion` clamps `length` to `max_seg_size` instead of splitting into multiple segments. Metadata can advertise more bytes than were actually registered — the same class of bug #2644 fixed for RDMA.
2. **Submit path does not cap at `BufferDesc` seams.** `UbTransport::submitTransferTask` slices by `slice_size` only. The initial `selectDevice` uses the full `request.length`, so a request that spans two registered segments fails even if individual slices would fit.
Until both are fixed, Store must keep splitting UB segments at `max_seg_size`.
## Proposed work
1. Auto-chunk `registerLocalMemory` into `<= max_seg_size` segments and publish one `BufferDesc` per chunk (do not shrink).
2. Cap submit-time slices at source and target `BufferDesc` remainder, similar to RDMA `SliceLengthCalculator`.
## Test plan
Needs Kunpeng / UB hardware. Suggested coverage:
- Register a buffer larger than `max_seg_size` and confirm it is not truncated
- WRITE whose source or destination straddles a chunk seam
Contributor guide
Research direction
Start with UrmaContext::registerMemoryRegion and registerLocalMemory to trace how large registrations become BufferDesc entries, then compare the submit path in UbTransport::submitTransferTask with the RDMA SliceLengthCalculator. On Kunpeng or UB hardware, verify registrations are not truncated and that WRITE transfers crossing source or destination chunk seams complete successfully.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp
- Domain
- networking, performance
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Quiet
- Clarity
- Mostly clear
- Newbie friendliness
- 45/100