feat(transfer)!: encode per-slot schema + NUMA source map in QueryBlocksForTransfer#354
feat(transfer)!: encode per-slot schema + NUMA source map in QueryBlocksForTransfer#354xiaguan wants to merge 8 commits into
Conversation
…sfer Introduce pegaflow-transfer-wire and transfer_plan module to run-encode block placements, materialize RDMA reads with coalescing, and shrink the query RPC metadata path. Co-authored-by: Cursor <cursoragent@cursor.com>
Coalesced transfers above ~1GiB hit LOC_LEN_ERR on mlx5; split each TransferDesc inline before building RdmaOps. Co-authored-by: Cursor <cursoragent@cursor.com>
Relax PCIe ordering for pinned host MRs on the rc_backend path. Co-authored-by: Cursor <cursoragent@cursor.com>
Thread holder-side per-slot NUMA from SealedBlock through the transfer plan (view/merge/schema/rebuild) so the requester rebuilds each block's slabs on the NUMA the holder used. Remote chunks carry numa_node and rebuild rejects UNKNOWN (crash early). Add p2p RDMA integration coverage for the slot->NUMA shapes the source map must reproduce: - normal MHA TP>1: cross-slot NUMA (each rank owns a slot column, saves with its own GPU's NUMA; a single block spans NUMA) - MLA replica (no DCP): cross-block NUMA via the server save-NUMA round-robin (one save RPC per block, rotating the candidate NUMA) Both self-adapt: on a single-NUMA box they still verify the multi-worker topology and the data roundtrip, logging that cross-NUMA was not exercised. Fixes uncovered while wiring the tests: - SealedBlock::from_slots builds with empty slot_numas again instead of routing through the strict-equal from_slots_with_numas (which panicked for every non-empty block). Matches the ssd_cache "empty OR full" contract; SSD-prefetch blocks keep their NUMA in SlotMeta. - transfer_plan round-trip fixture seeds a known NUMA so rebuild's UNKNOWN guard does not trip; deduped the two stride-block helpers. - kernel-backend integration test skips when the driver cannot load its PTX (older driver vs newer CUDA toolkit) instead of failing the suite. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The per-slot NUMA source map's crash-early guard (rebuild errors on an UNKNOWN remote-chunk NUMA) was only exercised incidentally before. Lock it in so it cannot silently regress into accepting UNKNOWN and materializing on an arbitrary node. Encode also warns (view.rs) and the production fetch boundary logs + meters the error (rdma_fetch), so the degradation is observable end to end, never silent. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The QueryBlocksForTransfer handler locked the found blocks inside the engine call, then encoded the TransferPlan. Encoding was fallible but the handler propagated its error with `?`, so on failure the requester never received the session id and could never release the lock — the blocks stayed pinned against LRU eviction until the 120s lock GC swept them. Release the lock explicitly when encoding fails. To make that release meaningful, the encode path is now genuinely panic-free: block_views_from_found and derive_slot_schemas return Result, and the reachable cross-block geometry mismatches (a namespace shared by incompatible KV layouts) become Err instead of asserts that would unwind the task and leak the lock anyway. This also fixes a latent out-of-bounds panic in derive_slot_schemas, which indexed every block by block 0's slot/segment shape before validating that the shapes matched. Remove dead code surfaced along the way: - SlotSchema.numa_node was write-only; NUMA lives on RemoteChunk. Drop the wire field and its sole writer. - assign_chunks ran a second coalesce pass whose `last.end == start` condition the first pass (`start <= last.end`) already subsumes. One pass coalesces. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
A coalesced QueryBlocksForTransfer plan hands the backend one big TransferDesc per NUMA node. batch_transfer_async mapped each desc to a single NIC (and, via the per-desc QP bucket, a single QP), so a multi-GB read saturated one NIC while the rest sat idle — large-segment p2p throughput collapsed to ~1/4 (one NUMA group's worth of NICs). Stripe each desc into NIC_STRIPE_BYTES (8 MiB) sub-ranges and round-robin them across its NUMA node's NICs, so one large read fans out over every NIC and QP in the group. The upper layer stays agnostic: it submits byte ranges; fanning them across the fabric is the transfer library's job. Single node, 8x400G NIC loopback, ~10 GiB/set (median GiB/s): segment master branch branch+fix 4 KiB 7.14 10.33 15.72 64 KiB 66.01 25.24 95.05 256 KiB 104.67 26.75 123.33 Adds a hardware-free unit test asserting a large desc touches every NIC in its group. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Perf: large-segment NIC-striping fix (cedfc25)While validating this branch's p2p RDMA fetch path I found (and fixed in Root cause. The coalesced Fix. Stripe each desc into 8 MiB sub-ranges, round-robin across its NUMA node's NICs (and their QPs), so one large coalesced read fans out over the whole group. The upper layer stays agnostic — it submits byte ranges; fanning them across the fabric is the transfer library's job. Measured — single node, 8×400G IB NICs, 2 NUMA sockets, NIC loopback, ~10 GiB/set, median of 3 measured sets (box under other load, so absolute numbers are loopback+contention-bound; read the relative A/B/C):
Pre-fix the branch fell to 0.26–0.38× of master at large segments; with the fix it's 1.2–2.2× master across the board — coalescing cuts WQE overhead at the IOPS end, striping restores full NIC parallelism at the bandwidth end. Caveat: single-node NIC loopback is not wire bandwidth; numbers are for relative comparison only. A unit test ( |
The QueryBlocksForTransfer encode path and the requester-side
rebuild_sealed_blocks both built nested per-cell Vec structures
(blocks * slots * 2 ~= 2.6M tiny allocations for a 4500-block batch),
then traversed them. Replace both with flat, column-major SoA:
- encode: BlockView{Vec<SlotView{Vec<SegmentView>}>} -> BlockMatrix with
contiguous [slot][segment][block] pointer/length buffers, so schema
derivation and chunk coalescing each scan a contiguous block-ordered
slice; one reused SegmentPoint scratch buffer across all columns.
- rebuild: drop the Vec<Vec<Vec<Option>>> matrix for two flat buffers
plus a per-chunk ChunkLoc that validates each placement run span once.
Plan-build CPU cost for 4500 blocks: encode 209ms -> 70ms (3.0x),
rebuild 383ms -> 101ms (3.7x). All validation/error paths preserved;
the segment-count consistency check moved earlier into the matrix builder.
Add a `bench` feature exposing transfer_plan::bench_support (RebuildFixture,
QueryFixture) built from the real encode -> materialize pipeline on a
vLLM-shaped fixture (slots = layers*tp, split K/V), gated so it never ships
in production. Include criterion benches (transfer_query, transfer_rebuild)
and encode_profile/rebuild_profile examples as isolated profiler drivers.
Also slim large_number_of_concurrent_sessions: it allocated one real
cudaHostAlloc mapped pool per dummy block (num_sessions * blocks_per_session
= 500 pools, each reserving ~100MB), OOM-killing the test box and blocking
the pre-commit hook. The lock manager only refcounts opaque Arc<SealedBlock>
handles, so all entries now share one block via Arc::clone.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Small-batch regression checkThe Performance numbers in the description are large-batch. This confirms tiny fetches (1-block warm hit, MLA single-block read) pay no fixed-overhead penalty from the TransferPlan encode / geometry / postcard work. Same 8-GPU, 2-NUMA, 8×400G node, single-host loopback, warm measure set (median of 4–5), byte-verified on all 8 ranks.
No regression — faster at every size. Even at 1 block the branch wins: master still emits per-block RDMA descriptors (1 block = 576 cells across 36 layers × TP8 × split-K/V) plus nested protobuf, whereas the TransferPlan coalesces to 2 descriptors, so Stage detail at the extremes (ms, measure set, connect amortized to 0):
The only line item where the branch is ever slightly higher is 🤖 Generated with Claude Code |
1 participant
What & why
QueryBlocksForTransferpreviously returned per-block / per-slot RDMA metadata as nested protobuf (TransferBlockInfo/TransferSlotInfo), with the wire size scaling asblocks × slots. This branch replaces that with a compact, run-encodedTransferPlan(newpegaflow-transfer-wirecrate, postcard-encoded into a singlebytes transfer_planfield): the response carries a per-slot segment schema plus a per-slot NUMA source map once, and the requester rebuilds each block's slabs on the NUMA the holder actually used.The "schema source map" point: the holder tells the requester not just where the bytes live, but which NUMA each slot is on, so the cross-node RDMA fetch materializes and rebuilds NUMA-correctly instead of landing everything on one node.
What's in it (5 commits)
pegaflow-transfer-wirecrate +transfer_plan/{encode,schema};QueryBlocksForTransferResponsecarriesbytes transfer_plan(postcard) instead of nested messages. Wire size scales with blocks, not blocks × slots.transfer_plan/{view,merge,schema,rebuild}+materialize; the requester allocates one slab per distinct chunk NUMA and rebuilds each slot on its source NUMA. Rebuild rejects UNKNOWN NUMA (crash-early; encode warns and the fetch boundary logs + meters the error, so the degradation is never silent). RestoresSealedBlock::from_slotsto build with emptyslot_numas(SSD-prefetch path; was regressed into a guaranteed panic). Adds the p2p RDMA integration coverage.Also: the kernel-backend integration test now self-skips when the driver can't JIT-load its PTX (older driver vs newer CUDA toolkit) instead of failing the suite — it still runs wherever the driver supports the kernel.
Breaking change
QueryBlocksForTransferResponsefields are renumbered andTransferBlockInfo/TransferSlotInfoare removed in favor ofbytes transfer_plan. Holder and requester must be upgraded together.Testing
cargo fmt/cargo clippy -D warnings/cargo test --releaseall pass.transfer_plan::tests: per-block NUMA preservation through rebuild, one slab per distinct chunk NUMA, loud UNKNOWN-NUMA rejection, encode→materialize→rebuild byte round-trip, wire-size scaling.distinct={0,1})distinct={0,1})Note
Related transfer work is in #350 (run-encode
QueryBlocksForTransferaddresses) on a separate branch — no shared commits, but both touch the transfer fetch path, so mind merge order.🤖 Generated with Claude Code