Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -48,8 +48,8 @@ if(NIXL_ROOT)
# whenever zmq is available, INDEPENDENT of ENABLE_UCX. cppzmq (header-only)
# is fetched at the top level when NIXL or UCX is on. Activates the
# TLLM_BOUNCE_V2 code paths; runtime use is still opt-in via
# CacheTransceiverConfig (agent_bounce_buffer_enable +
# kv_cache_bounce_size_mb, Python (v2) transceiver only).
# CacheTransceiverConfig (kv_cache_bounce_size_mb > 0, Python (v2) transceiver
# only).
# ------------------------------------------------------------------------
find_package(PkgConfig)
if(PKG_CONFIG_FOUND)
Expand Down Expand Up @@ -81,7 +81,7 @@ if(NIXL_ROOT)
else()
message(
STATUS
"NIXL bounce v2: disabled - libzmq/pkg-config not found; an explicit agent_bounce_buffer_enable opt-in fails at agent construction (on a bounce build, runtime setup failures degrade to standard NIXL with a WARNING)"
"NIXL bounce v2: disabled - libzmq/pkg-config not found; kv_cache_bounce_size_mb > 0 will fail at startup with this build"
)
endif()

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -131,20 +131,19 @@ namespace tensorrt_llm::executor::kv_cache::bounce
}

/// POD config for the bounce v2 pipeline. There is no `enabled` field: the on/off switch is
/// CacheTransceiverConfig's agent_bounce_buffer_enable + kv_cache_bounce_size_mb (which the Python
/// frontend folds into the agent's arena size in MiB: 0 = off, >0 = arena size), and at runtime
/// "bounce on" simply means the owning agent built its bounce state (mBounce != nullptr).
/// `arenaSizeBytes` is NOT env-backed either — the owning agent sets it from the same knob. The
/// expert knobs resolve as: agent_bounce_params dict > TRTLLM_NIXL_BOUNCE_* env var > built-in
/// default (`fromEnv()` reads the env, `fromParams()` layers the dict on top). Byte-valued knobs
/// accept an optional case-insensitive binary suffix such as "256MB", "1gb", or "512KiB"
/// (K/M/G == KiB/MiB/GiB, powers of two).
/// CacheTransceiverConfig's kv_cache_bounce_size_mb (which the Python frontend passes as the
/// agent's arena size in MiB: 0 = off, >0 = arena size), and at runtime "bounce on" simply means
/// the owning agent built its bounce state (mBounce != nullptr). `arenaSizeBytes` is NOT env-backed
/// either — the owning agent sets it from the same knob. The expert knobs resolve as:
/// agent_bounce_params dict > TRTLLM_NIXL_BOUNCE_* env var > built-in default (`fromEnv()` reads
/// the env, `fromParams()` layers the dict on top). Byte-valued knobs accept an optional
/// case-insensitive binary suffix such as "256MB", "1gb", or "512KiB" (K/M/G == KiB/MiB/GiB, powers
/// of two).
struct BounceConfig
{
// Overwritten with the agent arena size (MiB << 20, derived from CacheTransceiverConfig's
// kv_cache_bounce_size_mb + agent_bounce_buffer_enable) on every production path
// (maybeInitBounce); the
// default only serves unit tests that construct a BounceConfig{} directly.
// Overwritten with the agent arena size (CacheTransceiverConfig's kv_cache_bounce_size_mb << 20)
// on every production path (maybeInitBounce); the default only serves unit tests that construct a
// BounceConfig{} directly.
std::size_t arenaSizeBytes{512ULL << 20};
std::size_t arenaAllocationGranularityBytes{1ULL << 20}; // arena_allocation_granularity
std::size_t maxChunkSizeBytes{32ULL << 20}; // max_chunk_size
Expand All @@ -159,7 +158,7 @@ struct BounceConfig
// min_descriptor_count lowered. max_average_descriptor_size = 0 routes EVERY write to standard
// NIXL (an outbound gate off switch, not "no limit") — it disables OUTBOUND routing only: the arena is still
// allocated and registered, the handshake still advertised, and inbound grants/scatter still
// served; to turn the feature off use agent_bounce_buffer_enable=false. min_descriptor_count = 0
// served; to turn the feature off set kv_cache_bounce_size_mb=0. min_descriptor_count = 0
// means no minimum.
std::size_t minDescriptorCount{1024}; // min_descriptor_count
std::size_t maxAverageDescriptorSizeBytes{16ULL << 10}; // max_average_descriptor_size
Expand Down Expand Up @@ -338,7 +337,7 @@ struct BounceConfig

/// Defaults overridden by any set TRTLLM_NIXL_BOUNCE_* env var. Each call reads the current
/// environment. The on/off switch and the arena size are NOT read here — they only come from
/// CacheTransceiverConfig (agent_bounce_buffer_enable + kv_cache_bounce_size_mb).
/// CacheTransceiverConfig (kv_cache_bounce_size_mb).
[[nodiscard]] static BounceConfig fromEnv()
{
// paramKey -> env var name. Byte-valued knobs keep the historical _BYTES env suffix.
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -674,8 +674,8 @@ void BounceReceiver::scatterWorkerLoop()
// by this flow, so bounding to it prevents any cross-flow read. dstAddr is NOT bounded here
// by decision: it is the receiver's own KV address round-tripped through the control plane;
// the standard NIXL path would reject an out-of-region address via covers() (createXferReq
// -> populate), so bounce drops that check (accepted cost/benefit; the Python bounce does
// not validate it either). Only the SOURCE side is bounded to prevent cross-flow reads.
// -> populate), so bounce drops that check (accepted cost/benefit). Only the SOURCE side is
// bounded to prevent cross-flow reads.
// Any bad entry -> skip launch, NACK, no ACK.
// regionBytes==0 means the region wasn't allocated (stale) -> reject the whole job.
std::uint64_t const arenaLo = mCtx.arena->baseAddr();
Expand Down Expand Up @@ -1940,9 +1940,10 @@ bool BounceTransport::registerPeerHandshake(std::string const& peer, std::string
if (blob.empty())
{
TLLM_LOG_WARNING(
"BounceTransport(%s): peer %s does not advertise bounce (agent_bounce_buffer_enable off, bounce not "
"built, or bounce init failed on that side) -> standard NIXL for this peer; this agent's bounce arena "
"stays idle for it",
"BounceTransport(%s): peer %s does not advertise bounce, so this agent uses standard NIXL for it and "
"its bounce arena stays idle for that peer. Possible causes on the peer: kv_cache_bounce_size_mb is 0 "
"(agent_bounce_buffer_enable off on older releases), the build has no bounce support (libzmq), or "
"bounce init failed",
mCtx.selfName.c_str(), peer.c_str());
return false;
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -57,8 +57,8 @@ namespace tensorrt_llm::executor::kv_cache
{

// ============================================================================
// Bounce v2 integration (opt-in via CacheTransceiverConfig: agent_bounce_buffer_enable +
// kv_cache_bounce_size_mb). When disabled, transfers remain on the standard NIXL path.
// Bounce v2 integration (opt-in via CacheTransceiverConfig.kv_cache_bounce_size_mb > 0). When disabled,
// transfers remain on the standard NIXL path.
// ============================================================================
#ifdef TLLM_BOUNCE_V2
namespace bounce
Expand Down Expand Up @@ -145,18 +145,16 @@ class BounceTransferStatus final : public TransferStatus
void NixlTransferAgent::maybeInitBounce(
std::size_t agentBufferSizeMb, std::unordered_map<std::string, std::string> const& bounceParams)
{
// The size doubles as the on/off switch (CacheTransceiverConfig.kv_cache_bounce_size_mb when
// agent_bounce_buffer_enable is set, 0 otherwise): 0 keeps
// bounce disabled, >0 enables it at that arena capacity. The expert knobs resolve as
// The size doubles as the on/off switch (CacheTransceiverConfig.kv_cache_bounce_size_mb): 0
// keeps bounce disabled, >0 enables it at that arena capacity. The expert knobs resolve as
// agent_bounce_params dict > TRTLLM_NIXL_BOUNCE_* env var > built-in default.
if (agentBufferSizeMb == 0)
{
if (!bounceParams.empty())
{
TLLM_LOG_WARNING(
"NixlTransferAgent(%s): agent_bounce_params set but the bounce arena size is 0 "
"(agent_bounce_buffer_enable off or kv_cache_bounce_size_mb 0) -> bounce stays "
"disabled and the params are ignored",
"(kv_cache_bounce_size_mb 0) -> bounce stays disabled and the params are ignored",
mName.c_str());
}
return;
Expand Down Expand Up @@ -411,15 +409,16 @@ void NixlTransferAgent::maybeInitBounce(
// with a WARNING.
if (agentBufferSizeMb > 0)
{
TLLM_THROW("agent_bounce_buffer_enable requires a build with libzmq; this build has no bounce support");
TLLM_THROW(
"kv_cache_bounce_size_mb > 0 requests the transfer-agent bounce buffer, but this build has no bounce "
"support (built without libzmq); set kv_cache_bounce_size_mb to 0 or use a build with libzmq");
}
// Same diagnostic as the bounce build for params without an arena.
if (!bounceParams.empty())
{
TLLM_LOG_WARNING(
"NixlTransferAgent(%s): agent_bounce_params set but the bounce arena size is 0 "
"(agent_bounce_buffer_enable off or kv_cache_bounce_size_mb 0) -> bounce stays "
"disabled and the params are ignored",
"(kv_cache_bounce_size_mb 0) -> bounce stays disabled and the params are ignored",
mName.c_str());
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -226,10 +226,9 @@ class NixlTransferAgent : public BaseTransferAgent
/// Per-agent maps because different remote agents may have overlapping virtual addresses.
std::unordered_map<std::string, VramRegionMap> mRemoteVramRegionInfo;

/// Bounce v2 transport (opt-in via CacheTransceiverConfig.agent_bounce_buffer_enable +
/// kv_cache_bounce_size_mb). Null unless
/// enabled & built; when null the agent behaves exactly as before. See the design overview at
/// the top of bounce/BounceTransport.h.
/// Bounce v2 transport (opt-in via CacheTransceiverConfig.kv_cache_bounce_size_mb > 0). Null
/// unless enabled & built; when null the agent behaves exactly as before. See the design overview
/// at the top of bounce/BounceTransport.h.
std::unique_ptr<bounce::NixlBounceState> mBounce;

/// Lazily create the bounce transport (ctor, before any metadata exchange) when enabled.
Expand Down
15 changes: 0 additions & 15 deletions cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,6 @@
#include "kv_cache_manager_v2/coldPageCodec.h"
#include "kv_cache_manager_v2/common.h"
#include "kv_cache_manager_v2/config.h"
#include "kv_cache_manager_v2/cudaVirtMem.h"
#include "kv_cache_manager_v2/eventManager.h"
#include "kv_cache_manager_v2/exceptions.h"
#include "kv_cache_manager_v2/introspection.h"
Expand Down Expand Up @@ -2544,20 +2543,6 @@ void KvCacheManagerV2Bindings::initBindings(nb::module_& m)
nb::arg("quota"), nb::arg("slot_size_lists"), nb::arg("ratio_list"), nb::arg("granularity"),
nb::arg("min_slots"), nb::call_guard<nb::gil_scoped_release>());

// CUDA virtual-memory primitives, reached through _introspection because they carry no
// stability promise: the native disaggregated bounce buffer reserves one contiguous fabric
// region with them and maps physical chunks into it up front.
nb::class_<kv::PooledPhysMemAllocator>(mIntrospection, "PooledPhysMemAllocator")
.def(nb::init<size_t>(), nb::arg("phys_mem_size"))
.def_prop_ro("device_id", &kv::PooledPhysMemAllocator::deviceId);
nb::class_<kv::VirtMem>(mIntrospection, "VirtMem")
// keep_alive<1, 3>: VirtMem holds PooledPhysMemAllocator by reference, so the allocator
// must outlive it. Argument 3 is the allocator (1 is self, 2 is vm_size).
.def(nb::init<size_t, kv::PooledPhysMemAllocator&, size_t>(), nb::arg("vm_size"), nb::arg("phys_mem_allocator"),
nb::arg("init_num_phys_mem") = 0, nb::keep_alive<1, 3>())
.def("destroy", &kv::VirtMem::destroy)
.def_prop_ro("address", &kv::VirtMem::address);

// ---- Cold-page codec --------------------------------------------------
nb::class_<kv::IKvCacheColdPageCodec>(m, "IKvCacheColdPageCodec");
m.def("create_default_kv_cache_cold_page_codec", &kv::createDefaultKvCacheColdPageCodec,
Expand Down
5 changes: 2 additions & 3 deletions cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -138,9 +138,8 @@ std::unique_ptr<kvc::NixlTransferAgent> tryMakeAgent(kvc::BaseAgentConfig cfg, s
}
} // namespace

// BaseAgentConfig::agentBufferSizeMb (derived from CacheTransceiverConfig's
// agent_bounce_buffer_enable + kv_cache_bounce_size_mb) is the
// ONLY on/off switch: 0 keeps bounce disabled, >0 enables it.
// BaseAgentConfig::agentBufferSizeMb (CacheTransceiverConfig.kv_cache_bounce_size_mb, passed
// through) is the ONLY on/off switch: 0 keeps bounce disabled, >0 enables it.
TEST(BounceAgentE2E, AgentBufferSizeControlsBounce)
{
if (!hasCuda())
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ The checkpoint and the configuration file must live on a shared filesystem visib
* **CUDA graphs and the overlap scheduler are enabled.** The performance-sweep recipes set `disable_overlap_scheduler: false` and enable CUDA graphs. DEP16 additionally sets `cuda_graph_config.enable_padding: true`.
* **Chunked prefill is supported and enabled** (`enable_chunked_prefill: true`), so prompts longer than `max_num_tokens` are scheduled across multiple steps.
* **`kv_cache_config.tokens_per_block` is `64`** — retained as the validated configuration. Earlier releases needed this page size to avoid a legacy MLA fallback that did not support Kimi K3's 96 query heads.
* **Disaggregated serving and suffix-automaton speculative decoding are available.** Disaggregated serving is validated end-to-end on GB300 with matched DEP16 context/generation servers; see `examples/kimi_k3/disagg/README.md`. Suffix-automaton speculation works for both aggregated and disaggregated serving. See the "Current limitations" section of `examples/kimi_k3/README.md` for constraints (e.g. only matched DEP16=DEP16 disagg geometry is validated at scale).
* **Disaggregated serving and suffix-automaton speculative decoding are available.** Disaggregated serving is validated end-to-end on GB300 NVL72 with matched DEP16 context/generation servers on KVCacheManagerV2; see `examples/kimi_k3/disagg/README.md`. Suffix-automaton speculation works for both aggregated and disaggregated serving. See the "Current limitations" section of `examples/kimi_k3/README.md` for constraints (e.g. only matched DEP16=DEP16 disagg geometry is validated at scale).

## Deployment Steps

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -306,7 +306,10 @@ GDN state snapshots, chunked prefill, full-graph compilation, and piecewise CUDA
graphs. The [decode YAML](../../../examples/configs/curated/qwen3.5-gb300-agentic-decode.yaml)
sets MTP6, padded CUDA graphs, a no-eviction scheduler, and `stream_interval: 20`.
Both use KV cache manager V2, native host offload, NIXL with the Python
transceiver, and a 2048 MiB bounce buffer for this colocated profile.
transceiver, and a 2048 MiB C++ transfer-agent bounce buffer for this
colocated profile. Its gate is relaxed (`agent_bounce_params`): each KV
write also carries the GDN state, whose large descriptors would otherwise
keep the write on standard NIXL.

Both workers use `--publish-metrics` with `enable_iter_perf_stats: true`,
`iter_perf_stats_interval: 40`, and
Expand Down Expand Up @@ -363,10 +366,11 @@ revisit it when your conversation-length distribution changes.
| 3P1D | 3 × P4-DEP | 1 × D16-B8 | 7 |
| 3P2D | 3 × P4-DEP | 2 × D4-B32 | 5 |

For the two colocated 1P1D profiles, retain
`cache_transceiver_config.kv_cache_bounce_size_mb: 2048` on both roles and
For the two colocated 1P1D profiles, retain the bounce settings
(`cache_transceiver_config.kv_cache_bounce_size_mb: 2048` and
`agent_bounce_params`) on both roles and
`TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL=1` on both workers. For the other
profiles, remove that bounce-size override and unset the weight-loading variable
profiles, remove those bounce settings and unset the weight-loading variable
to use the pinned runtime's defaults. Launch each TP8 worker across two
four-GPU nodes and each TP16 worker across four nodes, with one MPI rank per GPU.
Replicas share discovery endpoints, namespace, served model name, and the
Expand Down
2 changes: 1 addition & 1 deletion docs/source/developer-guide/kv-transfer.md
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,7 @@ The control plane — discovery, per-request metadata, and completion notices
1. Context phase completes computation, KV cache stays in device memory awaiting transmission. The context worker publishes only the endpoint where it can be reached, on `ContextPhaseParams.disagg_info_endpoint`.
2. On first contact with a context instance, the generation worker fetches that instance's `RankInfo` — topology, attention geometry, serialized page table, transfer agent descriptor — from the context rank 0, then registers its own `RankInfo` with every context rank. The result is cached per endpoint, so only the first request pays for the round trip.
3. Generation phase requests KV cache for specific tokens, sending its destination block IDs to each overlapping context rank.
4. Context computes both source and destination addresses locally and submits a one-sided NIXL transfer. On the default path that transfer carries one descriptor per cache fragment; only the optional bounce path coalesces a request into a single contiguous buffer and a single write. All address arithmetic happens on the context side; the generation worker's CPU is not involved.
4. Context computes both source and destination addresses locally and submits a one-sided NIXL write with one descriptor per cache fragment. On this default path all address arithmetic happens on the context side and the generation worker's CPU is not involved. With `kv_cache_bounce_size_mb > 0`, the C++ NIXL agents may route an eligible write through their bounce buffers instead: the context agent gathers the fragments and writes them in chunks into the generation agent's buffer, and the generation agent scatters them into its KV cache.
5. Context reports completion over the control plane. The receive session completes once every expected writer has reported.
6. Generation phase resumes computation, context releases KV cache. Both sides run an allgather first, so every rank agrees on the outcome before any state transition.

Expand Down
Loading
Loading