[TRTLLM-15344][feat] cache transceiver nixl bounce buffer

[chuangz0]

Dev Engineer Review

• Adds NIXL bounce-buffer v2 for disaggregated KV-cache transfers.
• Adds arena allocation, transfer planning, credit scheduling, CUDA gather/scatter, wire codecs, control channels, and NIXL transport integration.
• Adds capability handshakes, timeout handling, cancellation, peer recovery, and shutdown cleanup.
• Routes eligible transfers through Bounce v2 and falls back to standard NIXL when initialization fails.
• Extends AgentDesc serialization with bounce handshake metadata.
• Adds Python and C++ observability for bounce activation and submission counts.
• Review focus includes API consistency, CUDA resource ownership, thread safety, RDMA completion handling, shutdown ordering, fallback behavior, and configuration validation.
• CI failures require follow-up before merge.

QA Engineer Review

• Adds unit, CUDA, ZMQ, NIXL transport, RDMA, failure-recovery, production-agent, and Python integration tests.
• Covers arenas, configuration parsing, message codecs, transfer planning, allocation, scheduling, execution pools, gather/scatter, control channels, transport recovery, fairness, backpressure, handshakes, and data integrity.
• Adds production-path tests for one-way, bidirectional, concurrent, and multi-sender transfers.
• Adds test_python_nixl_cache_transceiver_uses_cpp_bounce.
• Registers the Python test in tests/integration/test_lists/test-db/l0_h100.yml.
• Coverage mapping for the C++ tests is not provided.
• Verdict: needs follow-up.

Description

NIXL Bounce-Buffer v2 — Design

Describes the current (as-built) implementation. Code lives in
cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/, integrated through
NixlTransferAgent (transferAgent.cpp) and gated at runtime by TRTLLM_NIXL_BOUNCE_ENABLE.

1. Motivation

A single disaggregated-KV submitTransferRequests often carries thousands to tens of thousands of
small (~4 KiB) scattered descriptors. Submitting them to NIXL/NIC one-by-one is dominated by
per-descriptor overhead, leaving the link far below line rate.

The data-path idea: the sender gathers the scattered small descriptors into one pre-registered
buffer → does a single RDMA write to a peer buffer → the receiver scatters back to the final
destinations:

src×N  ──gather──▶  region(sender)  ──RDMA write──▶  region(receiver)  ──scatter──▶  dst×N

The hard part is the control plane: how buffers are allocated across senders, flow-controlled,
recycled, and how errors are handled. Requirements:

• R1 Stream gracefully even when a single request exceeds the total buffer on both sides
  (bounded buffers carry an unbounded transfer).
• R2 Pipeline gather / RDMA write / scatter across GPU and NIC to saturate the NIC.
• R3 One receiver written concurrently by many senders: fair, starvation-free, no cross-talk.
• R4 Thread-safe under multi-threaded submit.
• R5 Any error / peer loss / teardown resolves wait() (SUCCESS/FAILURE) — never hangs.
• R6/R7 Modular, testable, readable; control and data planes are pluggable.
• R8 Pack many small concurrent requests (per-request bytes < buffer size) tightly into the shared
  buffer — no request reserves a whole buffer, which would waste memory and cap concurrency.

R1 and R8 pull in opposite directions — carry one over-sized request and pack many small ones into
the same fixed buffer — and together motivate the variable-size region arena (§2).

2. Overview

• One shared BounceArena (a device buffer registered once) is sliced into variable-size
  regions: each chunk gets a region of exactly its byte size. Many small requests pack tightly
  (high concurrency, no waste); a request larger than the arena streams chunk-by-chunk with recycling.
• Credit flow control: the receiver owns the buffer; a sender must obtain a credit
  (exclusive write permission to a region) before writing. Credits are granted incrementally per chunk
  (WANT/GRANT) and recycled on scatter completion (ACK) — bounded end to end.
• Eager gather (default on): submit() launches a chunk's gather before the GRANT arrives,
  overlapping the WANT→GRANT control round-trip with the gather kernel. Eager (credit-less) staging is
  capped at half the arena, so on a bidirectional deployment both sides can always still grant
  incoming regions (no mutual eager-starvation); credit-backed allocations are not capped.
• Pluggable control channel: ZMQ/TCP by default; optionally the control messages ride NIXL
  notifications (UCX active messages on the RDMA fabric, TRTLLM_NIXL_BOUNCE_USE_NIXL_NOTIFICATIONS),
  dropping a control hop from tens of microseconds to a few. Peers must use the same control kind —
  enforced by the capability handshake (§5).
• Single-IO-thread reactor: one IO thread per agent owns the credit / request state (nearly
  lock-free), plus M scatter workers.
• No notifMsg, no extra flush on the data plane: data-landed is decided purely by the sender
  polling poll==SUCCESS (the NIXL UCX backend already appends ucp_ep_flush_nbx per transfer, so
  SUCCESS ⇒ the data is visible at the remote target).
• Decoupled execution resources: the stream/event/scratch needed to run one gather/scatter kernel
  comes from a small ExecPool (copyStreamCount contexts), borrowed/returned per kernel — separate
  from a region's long lifetime (which lasts until ACK).

flowchart LR
    subgraph SND["Sender agent (IO thread)"]
      REQ["Request{ numChunks,<br/>nextPost, acked }"]
      OUT["shared BounceArena + ExecPool<br/>acquireLocal(bytes) → gather staging<br/>(eager: before GRANT, ≤ ½ arena)"]
    end
    subgraph RCV["Receiver agent (IO thread)"]
      SCH["CreditScheduler (single allocator)<br/>BuddyAllocator(arena)<br/>flows{ pending, held } + localHeld<br/>round-robin ring + drain mode"]
      IN["shared BounceArena<br/>variable-size regions (RDMA targets)"]
      SW["scatter workers ×M (borrow ExecPool)"]
    end
    REQ -- "① WANT(rid, chunk_bytes[])" --> SCH
    SCH -- "② GRANT(regionHandle, addr, len)" --> REQ
    REQ -- "③ gather→RDMA write→DATA(regionHandle, scatter runs)" --> IN
    IN --> SW
    SW -- "④ scatter_done" --> SCH
    SCH -- "⑤ ACK(regionHandle)" --> REQ
    SCH -. "owns/allocates arena regions" .-> IN
    REQ -. "③ occupy / ⑤ release" .-> OUT

Loading

One agent can be both sender and receiver, and both roles share the same arena (most disagg agents
only send or only receive, so one buffer halves memory; a dual-role agent still shares a single arena,
deadlock-free — the eager cap above is what makes the bidirectional case safe).

3. Terminology

Term	Meaning
region	A variable-size slice of BounceArena; its handle = byte offset within the arena (addr = baseAddr + offset).
chunk	A unit produced by BounceTransferPlan by bin-packing scattered (src,dst) descriptors; ≤ maxChunkSizeBytes, moved by one RDMA write.
credit	Exclusive write permission to a region, handed to exactly one sender as GRANT{regionHandle, addr, len}.
flow	One independent request stream, key = "peer\x1f rid" (concurrent requests from the same peer are distinct flows).
in-flight cap W	Per-flow cap on in-flight regions (= pipeline depth) — maxInflightChunksPerRequest.
arena	The single shared, registered-once device buffer (BounceArena).
ExecCtx	{stream, event, scratch, hostPinned} needed to run one gather/scatter, borrowed/returned from ExecPool.
gather / scatter	Batched D2D copy between many small descriptors and a contiguous region (GatherScatterKernel).
scatter run	A coalesced DATA plan entry (BounceScatterRun): adjacent descriptors merged when contiguous or uniformly strided.
handshake	Per-agent bounce capability blob in AgentDesc (§5); peers engage bounce only when compatible.

4. Modules

Pure logic (no GPU / threads / IO — unit-testable):

Module	Responsibility
BounceConfig	env → POD config snapshot (fromEnv); byte-valued vars accept binary K/M/G suffixes ("256MB", "1gb", "512KiB"); garbage values fall back to defaults instead of parsing to 0.
BounceTransferPlan	Bin-pack (src,dst) descriptors → chunks (each ≤ maxChunkSizeBytes, 32 B aligned, zero-length skipped); compute each descriptor's in-region offset and packedBytes; coalesce the scatter view into scatterRuns (contiguous or uniform-stride merge) so the DATA message shrinks from per-desc entries (hundreds of KB) to a handful of runs.
BounceMessage	Control-plane wire codec (fixed header + fixed-size entries); little-endian; includes the handshake codec (encodeHandshake/decodeHandshake) and the cancel codec (encodeCancel/isCancelWant).
BuddyAllocator	Pure power-of-two buddy allocator over byte offsets; alloc(bytes)/free(offset), coalesces buddies, no external fragmentation, internal ≤ 2×.
CreditScheduler	Receiver-side credit allocation + round-robin fairness + drain-mode anti-starvation (§5); embeds a BuddyAllocator over the arena; also serves local-sender acquireLocal(bytes) with the eager half-arena cap. Owned by the IO thread; the only cross-thread caller is acquireLocal() from submit() app threads (eager gather staging).

Device / IO / integration:

Module	Responsibility
BounceArena	An arenaSizeBytes device buffer (MNNVL via common::FabricMemory, else cudaMalloc), registered once; base()/baseAddr()/at(offset).
ExecPool	copyStreamCount ExecCtxs; tryAcquire() (non-blocking, nullptr when full) / release(), thread-safe.
GatherScatterKernel	Batched memcpy. Default: custom kernel (uint4-vectorized; byte path for misaligned) reading its `[srcs
TransferEngine (abstract)	The only data-plane ops: registerRegion / postWrite / poll / release. NixlTransferEngine (production, wraps one nixlAgent) is the sole implementation; transport tests run over real NIXL loopback.
ControlChannel (abstract)	addPeer (returns success) / removePeer / sendTo / recv / localEndpoint. Two implementations: ZmqControlChannel (default): a ROUTER for receive + one DEALER per peer for send; sendTo is non-blocking (drops on full kSendHwm, never blocks the IO thread). NixlNotifControlChannel: control messages as NIXL notifications on the RDMA fabric (no TCP sockets; "endpoint" is serialized NIXL metadata).
BounceTransport	Thin reactor: 1 IO thread + M scatter workers; holds BounceContext + BounceSender + BounceReceiver, routes control messages to the right role, drains both roles each tick; submit() / addPeer() / forgetPeer() / shutdown(); owns the capability handshake (localHandshakeBlob / registerPeerHandshake / hasPeerHandshake).
BounceContext	Dependencies shared by both roles and owned by the single IO thread: injected channel/engine/arena/exec, the single CreditScheduler (one arena serves both directions), sendGrants().
BounceSender	[S] role: submit→WANT (+ eager gather), GRANT→attach credits / gather+write, ACK→resolve; holds the request table + send-side deferred-cleanup state (mOrphanLocal / mPendingCancel).
BounceReceiver	[R] role: WANT→grant region, DATA→scatter, reply ACK; holds scatter workers + job/done queues + mScattering (orphaned flag for in-flight scatters).
BounceNvtx	NVTX ranges over the transfer pipeline (submit / gather / write / scatter / control hops) for nsys analysis.
NixlTransferAgent integration	maybeInitBounce (build arena+exec+transport, register arena; any init failure warns and falls back to the standard NIXL path — never fails agent construction), shouldUseBounce (routing decision, gated on the peer handshake), AgentDesc carries the local handshake blob, invalidateRemoteAgent→forgetPeer. Built only when NIXL + zmq are available (TLLM_BOUNCE_V2); decoupled from ENABLE_UCX.

5. Control Plane: Credit Flow Control + Fair Scheduling (R1/R3)

Messages (all over ControlChannel — zmq by default, NIXL notifications opt-in):

Message	Direction	Payload
WANT	sender → receiver	per-chunk byte sizes (empty = cancel) + the sender's own control endpoint.
GRANT	receiver → sender	credits {addr, len, devId(receiver), regionHandle}.
DATA	sender → receiver	regionHandle + coalesced scatter runs; sent only after poll==SUCCESS (the scatter trigger).
ACK	receiver → sender	regionHandle; scatter done, region recyclable.

Capability handshake (bootstrap + compatibility, key). Each agent's AgentDesc carries a bounce
handshake blob: {wireVersion, controlKind (ZMQ | NIXL_NOTIF), arenaUsableCapacityBytes, maxChunkSizeBytes, endpoint}. loadRemoteAgent → registerPeerHandshake validates it — version,
control kind, and maxChunkSizeBytes must match the local config, and the peer's endpoint must be
registrable — and only then marks the peer bounce-capable. shouldUseBounce requires
hasPeerHandshake(peer), so a peer with bounce disabled, a different control transport, or mismatched
chunking silently stays on the standard NIXL path (no WANT ever stalls to requestTimeoutMs).
An agent that cannot produce a usable local endpoint advertises no handshake (never breaks
metadata exchange).

Reverse-path bootstrap. Bounce needs a two-way control channel (sender sends WANT/DATA, receiver
replies GRANT/ACK), but the disagg metadata exchange is one-directional — the KV sender
loadRemoteAgents the receiver, the receiver never loads the sender. So WANT also carries the
sender's control endpoint, and the receiver addPeer(sender)s in onWant to bootstrap the
reverse path. A malformed endpoint in a WANT is rejected on the reactor thread (warn, no grant, no
exception escapes); a cancel is still honored even when endpoint registration fails, so it can reclaim
flow state left by an earlier valid WANT. Cancel/abort uses an empty WANT (still carrying the
endpoint) — no separate handshake / RETURN message.

Endpoints are routable IPs (multi-node, ZMQ). ZmqControlChannel must not bind 127.0.0.1
(unreachable cross-node). maybeInitBounce resolves the local routable IP via the shared
common::getLocalIp(getEnvNixlInterface(), rank) (TRTLLM_NIXL_INTERFACE picks the NIC, else
auto-detect by egress route / hostname — identical to UCX/NIXL addressing) and binds tcp://<ip>:*;
localEndpoint() reads the actual tcp://<ip>:<port> from zmq last_endpoint and advertises it via
the handshake / WANT. Unit tests that construct ZmqControlChannel directly keep the
tcp://127.0.0.1:* default. IPv6: zmq disables IPv6 by default, so an IPv6 bind address is
bracketed (tcp://[<ip>]:*) and the ROUTER sets ZMQ_IPV6; the DEALER (addPeer) sets ZMQ_IPV6
unconditionally (harmless for IPv4) — aligned with ucx_utils. The ROUTER additionally sets
ZMQ_ROUTER_HANDOVER so a peer that is forgotten (removePeer drops its DEALER) and later reconnects
with the same routing id is accepted, rather than having its messages silently dropped while the stale
connection is reaped. addPeer validates the endpoint before connecting and reports failure to the
caller. With NixlNotifControlChannel the "endpoint" is serialized NIXL metadata; no TCP is involved.

Receiver state (lives on the IO thread, lock-free): BuddyAllocator arena +
flows{ pending: per-chunk bytes, held: region offsets, blockedAtGrantSequence } + a round-robin
ring of active flow keys + a cursor + a grant sequence counter. Fixed per-flow cap
W = maxInflightChunksPerRequest.

schedule() — on-demand, round-robin fair, never poisons the queue, never deadlocks:

flowchart TD
    Start([event triggers schedule:<br/>onWant / onScatterDone / reclaimFlow / reclaimByPrefix / releaseLocal]) --> D0{drain mode active?<br/>a flow bypassed ≥2 full rounds}
    D0 -- yes --> DAlloc{"arena.alloc(drain flow's head)<br/>fits?"}
    DAlloc -- no --> Hold([no NEW remote grants until it fits<br/>existing regions keep freeing])
    DAlloc -- yes --> DGrant[grant it, exit drain mode] --> C1
    D0 -- no --> C1{ring non-empty?}
    C1 -- no --> Done([return accumulated GRANTs])
    C1 -- yes --> Sweep[one round-robin sweep from the cursor]
    Sweep --> Find{current flow:<br/>pending non-empty AND held.size &lt; W?}
    Find -- "none in the whole sweep" --> Done
    Find -- yes --> Alloc{"arena.alloc(pending.front())<br/>fits?"}
    Alloc -- no --> Mark[mark flow blocked at current grant sequence] --> Sweep
    Alloc -- yes --> Grant["off = the allocation<br/>pending.pop_front(), held += off<br/>accumulate GRANT{off, base+off, len}<br/>advance cursor past this flow"]
    Grant --> C1

Loading

Intuition. schedule() answers: many remote senders want to write into the same shared arena —
how to hand out arena space fairly and bounded. Three constraints:

1. Fair: flows take turns (round-robin); no early flow monopolizes.
2. Per-flow rate limit (W): a flow may hold at most W in-flight regions (held.size() < W);
   more must wait for an ACK to free one — this is the pipeline depth.
3. Aggregate rate limit (arena): all flows' in-flight regions share one arena; if the next chunk
   doesn't fit, skip it for now (backpressure, not an error).

Think of it as taking turns at a ticket counter: ring is the queue of flows, cursor is "who's
next". Each inner sweep hands out exactly one region — to the flow that is next, still wants more
(pending non-empty), is under its cap (held < W), and whose front chunk fits right now — then
advances cursor past it and starts a fresh sweep. So grants alternate across flows rather
than filling one flow's window first. The outer loop repeats until a whole sweep grants nothing
(all pending-empty / at-cap / can't-fit).

Example (W=2, arena currently fits 4 regions; flow A has 3 chunks c1/c2/c3, flow B has 2 chunks
d1/d2, cursor starts at A):

Step	cursor	Action	A.held	B.held
1	A	grant A/c1	1	0
2	B	grant B/d1	1	1
3	A	grant A/c2	2	1
4	B	grant B/d2	2	2
5	A	A at cap (2≥W), B at cap → no progress → stop	2	2

Grant order A,B,A,B (strict alternation); A's c3 stays in pending until one of A's regions is
ACKed (onScatterDone→schedule()) and the next schedule() grants it. Each schedule() returns
the new batch of GRANTs, which sendGrants splits by flow key and sends to the right peer.

Notes:

• Two-level rate limiting: per-flow W bounds a single flow's pipeline depth; arena capacity
  bounds aggregate concurrency (alloc failure ⇒ backpressure, not deadlock). There is no "divide cap
  by active flow count" logic.
• Large-chunk anti-starvation (drain mode): a flow whose head chunk fails to fit is stamped with
  the current grant sequence. When other grants have bypassed it for ≥2 full rounds
  (kBypassRounds), the receiver enters drain mode for the oldest such flow: no new remote grants
  are issued until that head chunk fits (existing regions keep progressing and freeing space).
  acquireLocal() is deliberately unaffected — this is a receiver-only admission barrier and cannot
  introduce a bidirectional circular wait. Config guarantees maxChunkSizeBytes ≤ arena usable capacity (clamped at init), so a drained arena can always fit any chunk.
• Flow lifecycle reclaim: the flow key contains a monotonic non-reused rid; when both pending
  and held are empty, eraseIfDone drops the flow immediately — otherwise flows/ring grow
  unbounded on a long-running server and schedule() degrades to O(historical requests).
• Local sender shares the arena: acquireLocal(bytes) takes a gather-staging region from the same
  BuddyAllocator; eager (credit-less) staging is capped at half the arena (§2); the conservation
  invariant holds across {free bytes, each flow's held, localHeld}.

6. Data Plane & Pipeline (R2)

• The receiver fills a flow's held up to W (when space allows, one GRANT message batches multiple
  credits) → the sender holds W credits at once → W chunks in flight; each ACK frees a region and
  the receiver refills to keep held = W.
• Eager gather hides the control RTT: with enableEagerGather (default), submit() immediately
  stages and launches gathers for the first chunks (up to the in-flight cap and the eager half-arena
  budget) before any GRANT arrives; attachCredits later binds arriving credits to
  already-gathered chunks in strict chunk order, promoting their staging regions out of the eager
  budget. The WANT→GRANT round-trip and the gather kernel overlap instead of serializing.
• W is sized by the round-trip: W ≥ ⌈(write + getXferStatus + DATA + scatter + ACK/GRANT return) / single-chunk write⌉ (bandwidth-delay product); gather/scatter (D2D ~TB/s) hide in the shadow of the
  previous chunk's RDMA (IB ~25 GB/s), so the NIC stays the bottleneck.
• Gather and scatter each borrow a separate ExecCtx stream, so different chunks' gather/write/scatter
  overlap in wall-clock.
• Small control messages: the DATA scatter plan is run-merged — adjacent descriptors whose
  (bounceOffset, dstPtr) advance contiguously or by a uniform stride collapse into one
  BounceScatterRun (the fully-dense case, e.g. ctx tp1 → gen tp4, collapses thousands of descs to a
  handful of runs). TRTLLM_NIXL_BOUNCE_DISABLE_SCATTER_RUN_MERGING restores per-desc entries for
  control-plane A/B debugging only.
• No notifMsg / no GPUDirect flush: getXferStatus==SUCCESS already includes NIXL's per-transfer
  ucp_ep_flush_nbx (complete at both origin and target). The sender sends DATA as soon as it polls
  SUCCESS; the receiver scatters on receipt.

Single-chunk timing across both planes (control = solid, data = dashed; eager gather runs ① before
②'s GRANT when enabled):

sequenceDiagram
    autonumber
    participant App as Sender app thread
    participant SIO as Sender IO thread
    participant NIC as Data plane RDMA/NIC
    participant RIO as Receiver IO thread
    participant SW as Receiver scatter worker
    App->>RIO: WANT(rid, chunk_bytes[])
    App-->>App: ① eager GATHER (D2D): N small src → arena.at(o) (launch + eventRecord, no sync)
    Note over RIO: schedule(): arena.alloc carves region s, held[flow]+=s
    RIO->>SIO: GRANT(rid, regionHandle=s, addr, len)
    Note over SIO: onGrant: attachCredits to the eager-gathered chunk (or gather now if not eager)
    Note over SIO: drainGatherReady: cudaEventQuery==success (gather done, no block)
    SIO-->>NIC: ② postWrite(arena.at(o) → addr) async RDMA, no notif, return ExecCtx on done
    NIC-->>RIO: data written into receiver region s
    Note over SIO: pollSenderHandles: poll==SUCCESS (incl. ucp_ep_flush_nbx ⇒ landed at remote target)
    SIO->>RIO: ③ DATA(rid, chunk, regionHandle=s, scatter runs)
    Note over RIO: onData: validate runs against the flow's region, enqueue ScatterJob(s, runs)
    RIO->>SW: ScatterJob(s)
    SW-->>SW: ④ SCATTER (D2D): arena.at(s) → final dst×N (borrow ExecCtx, launch + streamSync)
    SW->>RIO: scatter_done(s)
    Note over RIO: onScatterDone: arena.free(s), schedule() may GRANT the next chunk
    RIO->>SIO: ⑤ ACK(rid, chunk, regionHandle=s)
    Note over SIO: onAck: releaseLocal(region o), acked++, all ACKed ⇒ promise=SUCCESS
    SIO->>App: future ready (SUCCESS)

Loading

For a large request (K > in-flight cap / arena capacity), the loop above repeats with both sides'
memory staying at O(W) regions while the transfer size is unbounded (R1).

7. Threading (R4/R5)

• One IO thread per agent, the sole owner of the CreditScheduler + sender request table (the key
  to being lock-free; the one cross-thread entry is eager acquireLocal from submit()). Each tick:
  1. recv one control message and dispatch;
  2. drainGatherReady (gather event ready → postWrite + return ExecCtx);
  3. pollSenderHandles (poll==SUCCESS → send DATA);
  4. drainScatterDone (worker report → send ACK + free region + reschedule);
  5. drainForgets / drainPendingPosts (retry parked credits) / checkTimeouts.
• M scatter workers: take a ScatterJob → borrow ExecCtx → scatter kernel + sync → report to the IO
  thread.
• submit() never blocks: registers a Request + sends WANT (+ launches eager gathers,
  fire-and-forget), returns a shared_future; safe to call from multiple threads.
• Adaptive poll: when gather/scatter is in flight, recv uses a 0 ms timeout (low latency); fully
  idle uses 1 ms; after long 0 ms spinning it backs off ~50 µs so a long model-kernel gather delay
  can't busy-spin a core.

8. State Machines

Unified region lifecycle (one arena serves both roles; a region is held by exactly one kind of
owner at any time — the conservation invariant):

stateDiagram-v2
    [*] --> FREE: arena init
    FREE --> INCOMING_HELD: GRANT(region→remote flow) 【schedule()】
    FREE --> OUTGOING_HELD: acquireLocal(bytes) 【local gather staging — eager ≤ ½ arena】
    INCOMING_HELD --> QUEUED: DATA(regionHandle, runs) received
    QUEUED --> SCATTERING: scatter worker picks it up (borrow ExecPool ctx)
    SCATTERING --> FREE: scatter_done → onScatterDone (reply ACK + reschedule)
    INCOMING_HELD --> FREE: forgetPeer / reclaimByPrefix (no in-flight DATA)
    OUTGOING_HELD --> FREE: ACK / failure → releaseLocal (reschedule — freed bytes go to a waiting flow)
    FREE --> [*]: shutdown
    note right of OUTGOING_HELD
      Invariant: every region is in exactly one of
      { free arena bytes } ∪ { some remote flow's held } ∪ { localHeld }.
      INCOMING_HELD = remote RDMA-write target (QUEUED/SCATTERING are its sub-states).
      OUTGOING_HELD = local gather source. Never held by both.
    end note

Loading

Sender chunk/request lifecycle (never hangs). With eager gather a chunk may reach Gathered
(gather event signalled, ExecCtx returned) while its GRANT is still in flight; attachCredits then
promotes it straight to the write:

stateDiagram-v2
    [*] --> WANT_SENT: submit() registers Request + sends WANT(chunk_bytes[]) + eager gathers
    WANT_SENT --> GATHERING: eager acquireLocal + launch gather (no credit yet)
    GATHERING --> GATHERED: drainGatherReady sees cudaEventQuery==success (waiting for credit)
    WANT_SENT --> POSTING: GRANT received (onGrant, credits queued)
    GATHERED --> IN_FLIGHT: attachCredits → postWrite
    POSTING --> POSTING: arena/exec full → credit parked, retried by drainPendingPosts (no block)
    POSTING --> GATHERING: borrow ExecCtx + acquireLocal(bytes) → launch gather + eventRecord (no sync)
    GATHERING --> IN_FLIGHT: gather done + credit attached → postWrite + return ExecCtx
    IN_FLIGHT --> DATA_SENT: pollSenderHandles sees poll==SUCCESS → send DATA
    DATA_SENT --> POSTING: ACK received and chunks remain (releaseLocal that region)
    DATA_SENT --> SUCCESS: acked == numChunks
    WANT_SENT --> FAILURE: checkTimeouts, no progress beyond requestTimeoutMs
    POSTING --> FAILURE: forgetPeer / shutdown
    GATHERING --> FAILURE: gather launch/record/stream error / forgetPeer / shutdown
    IN_FLIGHT --> FAILURE: poll==kFailed / forgetPeer / shutdown
    SUCCESS --> [*]: promise=SUCCESS
    FAILURE --> [*]: promise=FAILURE (empty WANT retracts credits)
    note right of WANT_SENT
      Every terminal state resolves the promise →
      the caller's wait() always returns, never hangs (R5)
    end note

Loading

9. Error Handling & Lifecycle (R5)

Every request reaches a terminal state; wait() never hangs. Bounce failures degrade, never break:
maybeInitBounce catches any construction error (fabric alloc, zmq bind, NIXL registration) with a
warning and leaves the agent on the standard per-desc NIXL path.

Case	Trigger	Handling
Peer incompatible (bounce off, other control kind, chunk-size mismatch, bad handshake)	registerPeerHandshake at loadRemoteAgent	Peer not marked bounce-capable → shouldUseBounce routes to standard NIXL (no WANT is ever sent, no timeout burned).
Peer never GRANTs (unreachable/not ready)	checkTimeouts exceeds requestTimeoutMs	Request → FAILURE (empty WANT retracts credits).
Malformed sender endpoint in WANT	onWant addPeer fails/throws	WANT rejected with a warning on the reactor thread (no grant); a cancel is still honored to reclaim earlier flow state.
Transfer engine error	poll==kFailed	Request → FAILURE, release in-flight handle.
gather launch/record error	pumpRequest flags it → drainGatherReady	Request → FAILURE (an unrecorded event would be misread as complete, so fail explicitly).
scatter kernel/sync error	scatter worker	No ACK → sender times out; never falsely reports landed (which would silently corrupt KV). Region is still freed (no leak).
Peer loss	invalidateRemoteAgent → forgetPeer	forgetPeer drops the peer's DEALER + its handshake registration synchronously; the IO thread reclaimByPrefix("peer\x1f") reclaims all of the peer's flows + fails its in-flight requests. A fresh loadRemoteAgent must re-validate a new handshake.
Control message can't be sent (peer stalled, queue full kSendHwm)	non-blocking sendTo returns EAGAIN	Drop the message + WARNING; never block the IO thread. Affected request → FAILURE via requestTimeoutMs (no hang, no corruption).
shutdown	—	join threads → cudaDeviceSynchronize → fail all in-flight requests.

Concurrency-safety points:

• In-flight scatter vs. re-grant race: the receiver keeps a mScattering map (region offset → is
  orphaned) of all scattering regions. If forgetPeer reclaims a region still in mScattering (a
  worker is reading it), reclaimByPrefix defers the free (marks it orphaned) and only
  freeOrphanRegions it on scatter completion — preventing "new sender's RDMA write ⊥ worker's read".
• Gather on the failure path: before reclaiming a region still GATHERING, cudaStreamSynchronize
  its stream (so an abandoned gather can't write a re-granted region); a sync error WARNs and clears the
  sticky error.
• In-flight write on the failure path: a region in state Writing may still be read by the NIC as
  a source, so it cannot be freed immediately. It's recorded in mOrphanLocal;
  drainOrphanLocal() polls its xfer to a terminal state before release + releaseLocal (the
  send-side orphan mechanism, symmetric to the receive-side mScattering orphaned flag).
• Cancel (empty WANT): precise reclaim + late-DATA guard: on sender failure/abort it sends an empty
  WANT; the receiver reclaimFlow immediately frees the flow's granted-but-unwritten regions
  (otherwise they stay held until peer loss, leaking on a long-running receiver), deferring any
  scattering ones. Correspondingly onData validates with heldByFlow: if an empty WANT raced ahead of
  a DATA so the region was freed/re-granted, that late DATA is dropped (never scatter a region now
  owned by someone else).
• Scatter input validation (defense in depth): scatter runs come from the peer's DATA; before
  launch the receiver checks every run's source range lies within this flow's granted region, and
  that the expanded plan doesn't exceed local scratch capacity. Any out-of-bounds → no launch, no ACK.
• CUDA errors: gather/scatter/sync return codes are not swallowed; errors WARN (with
  cudaGetErrorString).

Entry routing (transparent to callers; disabling is byte-equivalent to the original NIXL path):

flowchart TD
    S[submitTransferRequests] --> E{shouldUseBounce?<br/>WRITE + both-VRAM + no syncMsg<br/>+ peer handshake OK + descCount/avg<br/>+ per-side uniform deviceId}
    E -- no --> N[standard NIXL path]
    E -- yes --> SUB[submit: register Request + send WANT + eager gathers + return future]
    SUB --> POST[IO thread: GRANT→attach/gather+postWrite→DATA, pipelined]
    POST --> ERR{poll==kFailed / scatter fail / peer gone / stalled beyond requestTimeoutMs?}
    ERR -- yes --> F[Request → FAILURE]
    ERR -- no --> OK{all chunks poll SUCCESS AND all ACKed?}
    OK -- yes --> SU[wait = SUCCESS]
    F --> W[wait = FAILURE → caller task.fail]

Loading

Bootstrap: the bounce handshake blob is serialized with the AgentDesc (getLocalAgentDesc /
loadRemoteAgent(AgentDesc) → registerPeerHandshake), i.e. the path production disagg already uses;
the first WANT starts directly with no separate handshake round-trip.

10. Configuration (env)

All prefixed TRTLLM_NIXL_BOUNCE_. Byte-valued variables accept case-insensitive binary suffixes
(K/KB/KiB, M/MB/MiB, G/GB/GiB; all powers of two), e.g. ARENA_SIZE_BYTES=512MB. Unparsable
values fall back to the default (never silently become 0).

Field	env (suffix after prefix)	Default	Meaning
enabled	ENABLE	off	Master switch.
arenaSizeBytes	ARENA_SIZE_BYTES	256 MiB	Shared region arena size.
arenaAllocationGranularityBytes	ARENA_ALLOCATION_GRANULARITY_BYTES	1 MiB	Buddy minimum block (allocation granularity).
maxChunkSizeBytes	MAX_CHUNK_SIZE_BYTES	32 MiB	Per-chunk byte cap (plan bin-pack cap); clamped to the arena's usable capacity if larger; must match the peer's (handshake-checked).
maxInflightChunksPerRequest	MAX_INFLIGHT_CHUNKS_PER_REQUEST	8	Per-flow in-flight cap W (pipeline depth).
copyStreamCount	COPY_STREAM_COUNT	8	ExecPool context count (GPU kernel concurrency cap).
scatterWorkerCount	SCATTER_WORKER_COUNT	4	Scatter worker thread count.
minDescriptorCount	MIN_DESCRIPTOR_COUNT	1024	Routing gate: minimum descriptor count.
maxAverageDescriptorSizeBytes	MAX_AVERAGE_DESCRIPTOR_SIZE_BYTES	16 KiB	Routing gate: maximum average descriptor bytes.
requestTimeoutMs	REQUEST_TIMEOUT_MS	30000	No-progress timeout.
disableFabricMemory	DISABLE_FABRIC_MEMORY	off	Use cudaMalloc instead of MNNVL fabric memory (CI/x86).
enableEagerGather	ENABLE_EAGER_GATHER	on	Launch gathers at submit() before GRANT (overlap control RTT); eager staging capped at ½ arena.
useNixlNotifications	USE_NIXL_NOTIFICATIONS	off	Control plane over NIXL notifications (UCX active messages) instead of ZMQ/TCP; must match the peer (handshake-checked).
useZeroCopyArguments	USE_ZERO_COPY_ARGUMENTS	on	Copy kernel reads its plan directly from pinned host memory (no H2D staging).
useCubCopy	USE_CUB_COPY	off	Use cub::DeviceMemcpy::Batched instead of the custom copy kernel (experimental).
disableScatterRunMerging	DISABLE_SCATTER_RUN_MERGING	off	DEBUG ONLY: per-desc DATA plan instead of coalesced runs.

shouldUseBounce fires when: op is WRITE, src/dst are both VRAM, no syncMessage, the peer passed
the capability handshake, descCount ≥ minDescriptorCount, all srcs are on this agent's device and
all dsts on one device, and average desc bytes ≤ maxAverageDescriptorSizeBytes; otherwise the
standard NIXL path is used.

11. Test Coverage

Tests live under cpp/tests/unit_tests/executor/bounce/.

• Pure logic (no GPU): buddyAllocatorTest (split/coalesce/fragmentation/boundaries/overflow),
  creditSchedulerTest (in-flight cap/fairness/drain-mode anti-starvation/reclaim/conservation/
  reclaim-defer/orphan/eager budget), bounceMessageCodecTest (round-trip/truncation/magic/
  cross-type-reject/large-count/handshake codec), bounceTransferPlanTest (bin-pack boundaries +
  scatter-run merging), bounceConfigTest (env parsing, byte suffixes, garbage fallback).
• GPU unit: bounceArenaTest, execPoolTest, gatherScatterKernelTest (custom kernel /
  zero-copy args / cub backend).
• Transport (real NIXL loopback + zmq): zmqControlChannelTest (incl. endpoint validation),
  bounceTransportTest (end-to-end, byte-exact; handshake compatibility; malformed-WANT rejection),
  bounceTransportFailureTest (no-GRANT timeout / engine failure / shutdown in-flight / forgetPeer
  in-flight / multi-peer shared-arena over-subscription no-deadlock / multi-threaded submit).
• Real NIXL RDMA e2e: nixlTransferEngineTest; bounceNixlE2ETest (RealRdmaLoopback single
  transfer / ConcurrentBidirectionalRealRdma 8-thread bidirectional /
  MultiAgentManySendersToOneReceiver / ForgetPeerInFlightRecovers); bounceAgentE2ETest
  (production submitTransferRequests path: single transfer + ConcurrentSubmitUsesBounce).
  All e2e tests verify byte-exactly (seed-distinct pattern per transfer, ruling out cross-talk).
• Python integration: test_cache_transceiver_single_process.py drives the NIXL bounce path
  through the Python cache transceiver (added to l0_h100.yml).

12. perf compare

https://docs.google.com/document/d/1J8ROqb1D-TQryIEyqLoYtP4Z3Hqk_7AreluczjdtH5w/edit?usp=sharing
gptoss gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL

con128

	gen_side p50/p90/p99/mean (ms)	max (ms)
py_bounce	10.1 / 108.6 / 116.7 / 25.1	2660
py_bounce_t1	9.8 / 113.6 / 119.3 / 40.5	2675
bounce_v2	4.2 / 5.0 / 6.0 / 4.3	201
bounce_v2_zmq	4.0 / 4.9 / 5.7 / 4.1	105

con1024

	gen_side p50/p90/p99/mean (ms)	max (ms)
py_bounce	10.2 / 113.6 / 119.1 / 38.0	2721
py_bounce_t1	10.1 / 116.8 / 120.3 / 45.0	2700
bounce_v2	4.5 / 5.5 / 6.3 / 4.6	509
bounce_v2_zmq	4.6 / 5.5 / 6.4 / 4.6	197

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

To see a list of available CI bot commands, please comment /bot help.

[chuangz0]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #56919 [ run ] triggered by Bot. Commit: 006072a Link to invocation

[tensorrt-cicd]

PR_Github #56919 [ run ] completed with state SUCCESS. Commit: 006072a
/LLM/main/L0_MergeRequest_PR pipeline #45725 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[chuangz0]

/bot run --disable-fail-fast

[chuangz0]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61555 [ run ] triggered by Bot. Commit: b7c648d Link to invocation

[tensorrt-cicd]

PR_Github #61555 [ run ] completed with state FAILURE. Commit: b7c648d
/LLM/main/L0_MergeRequest_PR pipeline #49766 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[chuangz0]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61779 [ run ] triggered by Bot. Commit: 2aef331 Link to invocation

[tensorrt-cicd]

PR_Github #61779 [ run ] completed with state FAILURE. Commit: 2aef331
/LLM/main/L0_MergeRequest_PR pipeline #49979 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[chuangz0]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #65860 [ run ] triggered by Bot. Commit: 4a4a256 Link to invocation

[tensorrt-cicd]

PR_Github #65860 [ run ] completed with state FAILURE. Commit: 4a4a256
/LLM/main/L0_MergeRequest_PR pipeline #53553 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[chuangz0]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #66193 [ run ] triggered by Bot. Commit: 2750bf4 Link to invocation

[tensorrt-cicd]

PR_Github #66193 [ run ] completed with state FAILURE. Commit: 2750bf4
/LLM/main/L0_MergeRequest_PR pipeline #53864 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation