From ab0c89f5f34a813ef3da803d6cc5a5b64479bf46 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:08:19 +0800 Subject: [PATCH 01/11] [None][chore] Make kv_cache_bounce_size_mb alone select the C++ agent bounce The Python KV bounce buffer is no longer used now that the C++ transfer-agent bounce covers the E2E paths (#19588). Route kv_cache_bounce_size_mb straight to the C++ agent arena: - kv_cache_bounce_size_mb > 0 enables the C++ transfer-agent bounce buffer; 0 disables it. - agent_bounce_buffer_enable becomes a deprecated no-op that warns when set to True; True with size 0 no longer raises. - agent_bounce_params now requires kv_cache_bounce_size_mb > 0. - Lift the pipelined-transfer blocker that only applied to the Python bounce, and reword the C++ transceiver, pure-Python agent and FP4-MLA bridge messages for the new semantics. After this commit the Python bounce is unreachable; the next commit deletes it. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../disaggregation/kv_cache_transceiver.py | 8 +- .../_torch/disaggregation/native/transfer.py | 15 ++-- .../_torch/disaggregation/nixl/_agent_py.py | 11 ++- .../_torch/disaggregation/transceiver.py | 37 ++------ tensorrt_llm/llmapi/llm_args.py | 89 ++++++++++--------- .../schemas/trtllm-serve-config.schema.json | 9 +- .../trtllm-serve-disagg-config.schema.json | 9 +- .../test_cache_transceiver_single_process.py | 20 ++--- .../disaggregated/test_chunked_transfer.py | 29 +----- .../test_transceiver_bounded_polling.py | 37 ++------ tests/unittest/llmapi/test_llm_args.py | 62 +++++++++---- 11 files changed, 143 insertions(+), 183 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py b/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py index e93b8a819c59..978e2758efbb 100644 --- a/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/kv_cache_transceiver.py @@ -541,11 +541,11 @@ def __init__(self, f"RNN state transfer enabled: rnn_layer_num_per_pp={rnn_layer_num_per_pp_rank}" ) - if (cache_transceiver_config.kv_cache_bounce_size_mb > 0 - or cache_transceiver_config.agent_bounce_buffer_enable): + if cache_transceiver_config.kv_cache_bounce_size_mb > 0: logger.warning( - "bounce is only supported by the Python (v2) transceiver; " - "ignored on the C++ transceiver path") + "kv_cache_bounce_size_mb (the transfer-agent bounce buffer) is only supported by " + "the Python (v2) transceiver; ignored on the C++ transceiver path" + ) self.impl = CacheTransceiverCpp( kv_cache_manager.impl, total_num_kv_heads_per_layer, head_dim, diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 56c61cf1ab59..bc08639190fe 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -4504,10 +4504,8 @@ class TransferWorkerConfig: enforce_physical_ownership: bool = False # Internal activation gate: only the qualified containment policy supplies it. quiescence_fatal_callback: Optional[Callable[[QuiescenceFatalEvent], None]] = None - # Transfer-agent staging (bounce v2) buffer size in MiB; 0 disables the - # fast path. Derived upstream from CacheTransceiverConfig: it equals - # kv_cache_bounce_size_mb when agent_bounce_buffer_enable is set (in which - # case `bounce` above is off) and 0 otherwise. + # C++ transfer-agent bounce buffer size in MiB + # (CacheTransceiverConfig.kv_cache_bounce_size_mb); 0 disables it. agent_buffer_size_mb: int = 0 # Expert bounce knobs (TRTLLM_NIXL_BOUNCE_* names without the prefix and the # trailing _BYTES, lowercased); dict > env > default. Ignored when @@ -4621,11 +4619,10 @@ def _setup_transfer_engine(self): and not self._agent.bounce_enabled ): logger.warning( - f"TransferWorker: the C++ transfer-agent bounce was requested " - f"(agent_bounce_buffer_enable=True, kv_cache_bounce_size_mb=" - f"{self._config.agent_buffer_size_mb}) but is inactive on agent " - f"{self._agent.name} (init failed); standard per-descriptor NIXL " - "transfers are in use." + "TransferWorker: the transfer-agent bounce buffer was requested " + f"(kv_cache_bounce_size_mb={self._config.agent_buffer_size_mb}) but is " + f"inactive on agent {self._agent.name} (init failed); standard " + "per-descriptor NIXL transfers are in use." ) self._registered_mem: list = [] try: diff --git a/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py b/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py index b8c4f096b715..ac359aeacd58 100644 --- a/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py +++ b/tensorrt_llm/_torch/disaggregation/nixl/_agent_py.py @@ -102,17 +102,16 @@ def __init__( :param num_workers: Specify number of threads for the supported multi-threaded backends. :param rank: Process rank, used only to keep the shared agent interface consistent. :param world_size: Process count, used only to keep the shared agent interface consistent. - :param agent_buffer_size_mb: Accepted for interface consistency; the Python nixl-library - agent has no staging-buffer (bounce) support, so a positive size is ignored with a - warning. + :param agent_buffer_size_mb: Accepted for interface consistency; this agent has no + transfer-agent bounce buffer, so a positive size is ignored with a warning. :param agent_bounce_params: Accepted for interface consistency; ignored alongside agent_buffer_size_mb. """ if agent_buffer_size_mb > 0: logger.warning( - "The transfer-agent bounce buffer (agent_bounce_buffer_enable) is not " - "supported by the Python nixl-library agent; ignoring the buffer size " - "(and any agent_bounce_params)" + "kv_cache_bounce_size_mb > 0 requests the transfer-agent bounce buffer, which " + "the pure-Python NIXL agent does not implement; ignoring it and any " + "agent_bounce_params (standard NIXL transfers are used)" ) if (rank is None) != (world_size is None): raise ValueError("rank and world_size must be specified together") diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 3c9e4f245db7..81e59cb3744c 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -39,9 +39,6 @@ GenTransferStatus, KvCacheTransceiver, ) -from tensorrt_llm._torch.disaggregation.native.bounce import ( - config_from_size as bounce_config_from_size, -) from tensorrt_llm._torch.disaggregation.native.fetch import PeerFetch from tensorrt_llm._torch.disaggregation.native.perf_logger import perf_log_manager from tensorrt_llm._torch.disaggregation.native.publish import PeerPublish @@ -215,6 +212,9 @@ def _validate_fp4_mla_bridge_profile( and mapping.enable_attention_dp and mapping.pp_size == 1 and mapping.cp_size == 1 + # The C++ agent bounce scatters into the destination on the receiver agent's threads, + # which can still be writing after a failure is reported; the bridge cannot prove that + # those writes stopped. and cache_transceiver_config.kv_cache_bounce_size_mb == 0 and not cache_transceiver_config.enable_pipelined_transfer and os.getenv("TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP") != "1" @@ -225,7 +225,8 @@ def _validate_fp4_mla_bridge_profile( "FP4 MLA lifecycle bridge requires a disaggregated NVFP4 SELFKONLY " "KVCacheManagerV2, the C++ NIXL agent binding, a finite timeout, " "no-retry ADP, PP1/CP1, " - "async monolithic non-layerwise transfer, and bounce disabled" + "async monolithic non-layerwise transfer, and kv_cache_bounce_size_mb=0 " + "(transfer-agent bounce disabled)" ) return True @@ -300,22 +301,9 @@ def __init__( quiescence_fatal_callback=( self._fail_unproven_transfer if enforce_physical_ownership else None ), - # kv_cache_bounce_size_mb is the shared bounce capacity; agent_bounce_buffer_enable - # routes it to exactly one implementation: the Python bounce below (per-region, - # size 0 = off; the per-transfer size gates are internal, tuned via env: - # TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS for plain-KV payloads, - # TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES for recurrent-state payloads) or the C++ - # transfer-agent staging buffer (bounce v2, one buffer shared by send and recv). - bounce=bounce_config_from_size( - 0 - if cache_transceiver_config.agent_bounce_buffer_enable - else cache_transceiver_config.kv_cache_bounce_size_mb - ), - agent_buffer_size_mb=( - cache_transceiver_config.kv_cache_bounce_size_mb - if cache_transceiver_config.agent_bounce_buffer_enable - else 0 - ), + # kv_cache_bounce_size_mb sizes the C++ transfer-agent bounce buffer (one + # arena shared by send and receive); 0 disables it. + agent_buffer_size_mb=cache_transceiver_config.kv_cache_bounce_size_mb, agent_bounce_params=cache_transceiver_config.agent_bounce_params, ) ) @@ -1601,15 +1589,6 @@ def _resolve_pipelined_transfer(self, cfg: CacheTransceiverConfig) -> bool: if not cfg.enable_pipelined_transfer: return False blockers = [] - # The Python bounce reserves a receiver region for the whole request, not per chunk. - # The C++ transfer-agent bounce (agent_bounce_buffer_enable) stages each transfer request - # independently below the Python layer, so a pipelined chunk is just another request. - # TODO: This only sees the local config, so a chunking sender and a bouncing receiver pass. - if cfg.kv_cache_bounce_size_mb != 0 and not cfg.agent_bounce_buffer_enable: - blockers.append( - f"the Python bounce buffer (kv_cache_bounce_size_mb=" - f"{cfg.kv_cache_bounce_size_mb} without agent_bounce_buffer_enable)" - ) if isinstance(self._kv_cache_manager, (MambaHybridCacheManager, MambaHybridCacheManagerV2)): blockers.append("a Mamba/hybrid cache manager") # Refused for the configuration, not for the piece: the per-chunk guard downstream only sees diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 4d3b3b48de4d..4b03dec3b9a4 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -5369,15 +5369,25 @@ class CacheTransceiverConfig(StrictBaseModel, PybindMirror): default=0, ge=0, description= - "Capacity in MiB of the native-disagg KV-cache bounce buffer, which " - "coalesces a request's scattered per-block KV for a single multi-rail " - "NIXL write. The size doubles as the on/off switch: 0 (default) keeps " - "the per-block path, >0 enables bounce at that capacity. By default " - "two buffers of this size are allocated (one for send, one for recv); " - "with agent_bounce_buffer_enable a single shared buffer is allocated " - "instead and the size should be a power of two (256/512/1024). " - "Requires the Python (v2) transceiver (transceiver_runtime); the C++ " - "transceiver does not support bounce and ignores this field.") + "Size in MiB of the KV-cache bounce buffer in the C++ NIXL transfer " + "agent; 0 (default) turns it off and KV fragments are written " + "directly with standard NIXL. A positive value allocates one buffer " + "of this size per rank, shared by send and receive. Use a power of " + "two such as 256, 512 or 1024; for other values only the largest " + "power of two that fits is used. For an eligible write, the sender " + "gathers the fragments into the buffer and sends them in chunks, " + "and the receiver scatters them into its KV cache. By default a " + "write is eligible only with >= 1024 descriptors averaging <= 16 KiB " + "(typical when context and generation split attention heads " + "differently); agent_bounce_params changes that. Set it on both " + "context and generation servers; if either side has 0, the pair " + "uses standard NIXL. If either side's buffer is smaller than " + "max_chunk_size (32 MiB by default), both sides must use the same " + "size, or the pair uses standard NIXL. Requires the Python (v2) " + "transceiver with the C++ NIXL agent; on a build without bounce " + "support (libzmq) a positive value fails at startup. The C++ " + "transceiver and the pure-Python NIXL agent ignore it with a " + "warning.") enable_pipelined_transfer: bool = Field( default=False, @@ -5385,27 +5395,17 @@ class CacheTransceiverConfig(StrictBaseModel, PybindMirror): "Transfer each completed prefill chunk's KV cache while later chunks " "compute. Requires Python NIXL, generation-first scheduling, chunked " "prefill, pipeline_parallel_size=1, context_parallel_size=1 on both " - "peers, beam_width=1, no Python bounce buffer (the C++ transfer-agent " - "bounce selected by agent_bounce_buffer_enable is allowed) or " - "Mamba/hybrid cache, and block reuse disabled or set to all_reusable. " - "Invalid static settings fail at " - "startup; per-request constraints reject the request.") + "peers, beam_width=1, no Mamba/hybrid cache, and block reuse disabled " + "or set to all_reusable. Invalid static settings fail at startup; " + "per-request constraints reject the request.") agent_bounce_buffer_enable: bool = Field( default=False, + status="deprecated", description= - "Run the KV-cache bounce in the C++ transfer agent instead of the " - "Python transceiver, using one shared kv_cache_bounce_size_mb buffer " - "(use a power of two). Set this identically on context and " - "generation; kv_cache_bounce_size_mb matters only when its usable " - "capacity (rounded down to a power of two) is below max_chunk_size; " - "both sides must then clamp to the same chunk cap, otherwise the pair " - "falls back to standard NIXL. By default only writes with " - "many small descriptors (>= 1024, <= 16 KiB average, i.e. " - "head-mismatch layouts) take the path; head-matched layouts (MLA, " - "symmetric TP) stay on standard NIXL unless agent_bounce_params " - "relaxes the gate. Requires the Python (v2) transceiver; see the " - "disaggregated-serving docs.") + "DEPRECATED: no-op kept so existing configs still parse; it will be " + "removed in a future release. kv_cache_bounce_size_mb > 0 alone " + "enables the C++ transfer-agent bounce buffer, and 0 disables it.") agent_bounce_params: Optional[Dict[str, str]] = Field( default=None, @@ -5414,12 +5414,13 @@ class CacheTransceiverConfig(StrictBaseModel, PybindMirror): "tensorrt_llm/_torch/disaggregation/nixl/bounce_knobs.py for the valid " "keys. Byte-valued knobs accept a KB/MB/GB suffix (e.g. '32MB'). " "Precedence: this dict > environment variable > built-in default. " - "Requires agent_bounce_buffer_enable. request_timeout_ms and the " - "effective max_chunk_size (after the arena clamp) must match between " - "context and generation. " - "max_average_descriptor_size=0 stops outbound routing only (arena, " - "handshake and inbound scatter stay active); set " - "agent_bounce_buffer_enable=False to turn the feature off.") + "Requires kv_cache_bounce_size_mb > 0. request_timeout_ms and the " + "effective chunk size (max_chunk_size, capped by the buffer size) " + "must match between context and generation. " + "max_average_descriptor_size=0 stops this worker from sending " + "through the buffer, but it still allocates the buffer and receives " + "bounced writes; set kv_cache_bounce_size_mb=0 to turn the feature " + "off.") @field_validator('agent_bounce_params', mode='before') @classmethod @@ -5435,17 +5436,23 @@ def coerce_agent_bounce_params_to_str(cls, v): @model_validator(mode='after') def validate_bounce_config(self) -> 'CacheTransceiverConfig': - if self.agent_bounce_buffer_enable and self.kv_cache_bounce_size_mb == 0: - raise ValueError( - "agent_bounce_buffer_enable selects the C++ transfer-agent " - "bounce implementation, but kv_cache_bounce_size_mb is 0 " - "(bounce disabled); set a positive capacity.") + # Deprecated no-op: warn only on True. An explicit False looks like the + # default after apply_model_defaults_to_llm_args re-validates + # model_dump(), and warning_once keeps that re-validation from + # repeating the warning. + if self.agent_bounce_buffer_enable: + logger.warning_once( + "cache_transceiver_config.agent_bounce_buffer_enable is " + "deprecated and ignored: kv_cache_bounce_size_mb > 0 alone " + "enables the C++ transfer-agent bounce buffer. Remove the " + "field; it will be deleted in a future release.", + key="cache_transceiver_config.agent_bounce_buffer_enable") if self.agent_bounce_params: - if not self.agent_bounce_buffer_enable: + if self.kv_cache_bounce_size_mb == 0: raise ValueError( - "agent_bounce_params only applies to the C++ transfer-agent " - "bounce implementation; set agent_bounce_buffer_enable=True " - "or drop the params.") + "agent_bounce_params tunes the C++ transfer-agent bounce " + "buffer, but kv_cache_bounce_size_mb is 0 (bounce " + "disabled); set a positive capacity or drop the params.") # Lazy: importing tensorrt_llm._torch at module scope is circular. from tensorrt_llm._torch.disaggregation.nixl.bounce_knobs import \ AGENT_BOUNCE_PARAM_KEYS diff --git a/tensorrt_llm/schemas/trtllm-serve-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-config.schema.json index ce332dfae73a..b1271b6c540b 100644 --- a/tensorrt_llm/schemas/trtllm-serve-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-config.schema.json @@ -403,7 +403,8 @@ "properties": { "agent_bounce_buffer_enable": { "default": false, - "description": "Run the KV-cache bounce in the C++ transfer agent instead of the Python transceiver, using one shared kv_cache_bounce_size_mb buffer (use a power of two). Set this identically on context and generation; kv_cache_bounce_size_mb matters only when its usable capacity (rounded down to a power of two) is below max_chunk_size; both sides must then clamp to the same chunk cap, otherwise the pair falls back to standard NIXL. By default only writes with many small descriptors (>= 1024, <= 16 KiB average, i.e. head-mismatch layouts) take the path; head-matched layouts (MLA, symmetric TP) stay on standard NIXL unless agent_bounce_params relaxes the gate. Requires the Python (v2) transceiver; see the disaggregated-serving docs.", + "description": "DEPRECATED: no-op kept so existing configs still parse; it will be removed in a future release. kv_cache_bounce_size_mb > 0 alone enables the C++ transfer-agent bounce buffer, and 0 disables it.", + "status": "deprecated", "title": "Agent Bounce Buffer Enable", "type": "boolean" }, @@ -420,7 +421,7 @@ } ], "default": null, - "description": "Expert tuning knobs for the C++ transfer-agent bounce pipeline; see tensorrt_llm/_torch/disaggregation/nixl/bounce_knobs.py for the valid keys. Byte-valued knobs accept a KB/MB/GB suffix (e.g. '32MB'). Precedence: this dict > environment variable > built-in default. Requires agent_bounce_buffer_enable. request_timeout_ms and the effective max_chunk_size (after the arena clamp) must match between context and generation. max_average_descriptor_size=0 stops outbound routing only (arena, handshake and inbound scatter stay active); set agent_bounce_buffer_enable=False to turn the feature off.", + "description": "Expert tuning knobs for the C++ transfer-agent bounce pipeline; see tensorrt_llm/_torch/disaggregation/nixl/bounce_knobs.py for the valid keys. Byte-valued knobs accept a KB/MB/GB suffix (e.g. '32MB'). Precedence: this dict > environment variable > built-in default. Requires kv_cache_bounce_size_mb > 0. request_timeout_ms and the effective chunk size (max_chunk_size, capped by the buffer size) must match between context and generation. max_average_descriptor_size=0 stops this worker from sending through the buffer, but it still allocates the buffer and receives bounced writes; set kv_cache_bounce_size_mb=0 to turn the feature off.", "title": "Agent Bounce Params" }, "backend": { @@ -445,13 +446,13 @@ }, "enable_pipelined_transfer": { "default": false, - "description": "Transfer each completed prefill chunk's KV cache while later chunks compute. Requires Python NIXL, generation-first scheduling, chunked prefill, pipeline_parallel_size=1, context_parallel_size=1 on both peers, beam_width=1, no Python bounce buffer (the C++ transfer-agent bounce selected by agent_bounce_buffer_enable is allowed) or Mamba/hybrid cache, and block reuse disabled or set to all_reusable. Invalid static settings fail at startup; per-request constraints reject the request.", + "description": "Transfer each completed prefill chunk's KV cache while later chunks compute. Requires Python NIXL, generation-first scheduling, chunked prefill, pipeline_parallel_size=1, context_parallel_size=1 on both peers, beam_width=1, no Mamba/hybrid cache, and block reuse disabled or set to all_reusable. Invalid static settings fail at startup; per-request constraints reject the request.", "title": "Enable Pipelined Transfer", "type": "boolean" }, "kv_cache_bounce_size_mb": { "default": 0, - "description": "Capacity in MiB of the native-disagg KV-cache bounce buffer, which coalesces a request's scattered per-block KV for a single multi-rail NIXL write. The size doubles as the on/off switch: 0 (default) keeps the per-block path, >0 enables bounce at that capacity. By default two buffers of this size are allocated (one for send, one for recv); with agent_bounce_buffer_enable a single shared buffer is allocated instead and the size should be a power of two (256/512/1024). Requires the Python (v2) transceiver (transceiver_runtime); the C++ transceiver does not support bounce and ignores this field.", + "description": "Size in MiB of the KV-cache bounce buffer in the C++ NIXL transfer agent; 0 (default) turns it off and KV fragments are written directly with standard NIXL. A positive value allocates one buffer of this size per rank, shared by send and receive. Use a power of two such as 256, 512 or 1024; for other values only the largest power of two that fits is used. For an eligible write, the sender gathers the fragments into the buffer and sends them in chunks, and the receiver scatters them into its KV cache. By default a write is eligible only with >= 1024 descriptors averaging <= 16 KiB (typical when context and generation split attention heads differently); agent_bounce_params changes that. Set it on both context and generation servers; if either side has 0, the pair uses standard NIXL. If either side's buffer is smaller than max_chunk_size (32 MiB by default), both sides must use the same size, or the pair uses standard NIXL. Requires the Python (v2) transceiver with the C++ NIXL agent; on a build without bounce support (libzmq) a positive value fails at startup. The C++ transceiver and the pure-Python NIXL agent ignore it with a warning.", "minimum": 0, "title": "Kv Cache Bounce Size Mb", "type": "integer" diff --git a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json index e34b7f2adaa7..ad48aa452419 100644 --- a/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json +++ b/tensorrt_llm/schemas/trtllm-serve-disagg-config.schema.json @@ -403,7 +403,8 @@ "properties": { "agent_bounce_buffer_enable": { "default": false, - "description": "Run the KV-cache bounce in the C++ transfer agent instead of the Python transceiver, using one shared kv_cache_bounce_size_mb buffer (use a power of two). Set this identically on context and generation; kv_cache_bounce_size_mb matters only when its usable capacity (rounded down to a power of two) is below max_chunk_size; both sides must then clamp to the same chunk cap, otherwise the pair falls back to standard NIXL. By default only writes with many small descriptors (>= 1024, <= 16 KiB average, i.e. head-mismatch layouts) take the path; head-matched layouts (MLA, symmetric TP) stay on standard NIXL unless agent_bounce_params relaxes the gate. Requires the Python (v2) transceiver; see the disaggregated-serving docs.", + "description": "DEPRECATED: no-op kept so existing configs still parse; it will be removed in a future release. kv_cache_bounce_size_mb > 0 alone enables the C++ transfer-agent bounce buffer, and 0 disables it.", + "status": "deprecated", "title": "Agent Bounce Buffer Enable", "type": "boolean" }, @@ -420,7 +421,7 @@ } ], "default": null, - "description": "Expert tuning knobs for the C++ transfer-agent bounce pipeline; see tensorrt_llm/_torch/disaggregation/nixl/bounce_knobs.py for the valid keys. Byte-valued knobs accept a KB/MB/GB suffix (e.g. '32MB'). Precedence: this dict > environment variable > built-in default. Requires agent_bounce_buffer_enable. request_timeout_ms and the effective max_chunk_size (after the arena clamp) must match between context and generation. max_average_descriptor_size=0 stops outbound routing only (arena, handshake and inbound scatter stay active); set agent_bounce_buffer_enable=False to turn the feature off.", + "description": "Expert tuning knobs for the C++ transfer-agent bounce pipeline; see tensorrt_llm/_torch/disaggregation/nixl/bounce_knobs.py for the valid keys. Byte-valued knobs accept a KB/MB/GB suffix (e.g. '32MB'). Precedence: this dict > environment variable > built-in default. Requires kv_cache_bounce_size_mb > 0. request_timeout_ms and the effective chunk size (max_chunk_size, capped by the buffer size) must match between context and generation. max_average_descriptor_size=0 stops this worker from sending through the buffer, but it still allocates the buffer and receives bounced writes; set kv_cache_bounce_size_mb=0 to turn the feature off.", "title": "Agent Bounce Params" }, "backend": { @@ -445,13 +446,13 @@ }, "enable_pipelined_transfer": { "default": false, - "description": "Transfer each completed prefill chunk's KV cache while later chunks compute. Requires Python NIXL, generation-first scheduling, chunked prefill, pipeline_parallel_size=1, context_parallel_size=1 on both peers, beam_width=1, no Python bounce buffer (the C++ transfer-agent bounce selected by agent_bounce_buffer_enable is allowed) or Mamba/hybrid cache, and block reuse disabled or set to all_reusable. Invalid static settings fail at startup; per-request constraints reject the request.", + "description": "Transfer each completed prefill chunk's KV cache while later chunks compute. Requires Python NIXL, generation-first scheduling, chunked prefill, pipeline_parallel_size=1, context_parallel_size=1 on both peers, beam_width=1, no Mamba/hybrid cache, and block reuse disabled or set to all_reusable. Invalid static settings fail at startup; per-request constraints reject the request.", "title": "Enable Pipelined Transfer", "type": "boolean" }, "kv_cache_bounce_size_mb": { "default": 0, - "description": "Capacity in MiB of the native-disagg KV-cache bounce buffer, which coalesces a request's scattered per-block KV for a single multi-rail NIXL write. The size doubles as the on/off switch: 0 (default) keeps the per-block path, >0 enables bounce at that capacity. By default two buffers of this size are allocated (one for send, one for recv); with agent_bounce_buffer_enable a single shared buffer is allocated instead and the size should be a power of two (256/512/1024). Requires the Python (v2) transceiver (transceiver_runtime); the C++ transceiver does not support bounce and ignores this field.", + "description": "Size in MiB of the KV-cache bounce buffer in the C++ NIXL transfer agent; 0 (default) turns it off and KV fragments are written directly with standard NIXL. A positive value allocates one buffer of this size per rank, shared by send and receive. Use a power of two such as 256, 512 or 1024; for other values only the largest power of two that fits is used. For an eligible write, the sender gathers the fragments into the buffer and sends them in chunks, and the receiver scatters them into its KV cache. By default a write is eligible only with >= 1024 descriptors averaging <= 16 KiB (typical when context and generation split attention heads differently); agent_bounce_params changes that. Set it on both context and generation servers; if either side has 0, the pair uses standard NIXL. If either side's buffer is smaller than max_chunk_size (32 MiB by default), both sides must use the same size, or the pair uses standard NIXL. Requires the Python (v2) transceiver with the C++ NIXL agent; on a build without bounce support (libzmq) a positive value fails at startup. The C++ transceiver and the pure-Python NIXL agent ignore it with a warning.", "minimum": 0, "title": "Kv Cache Bounce Size Mb", "type": "integer" diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index 1aeb0bc55d25..e420884ce6d2 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py @@ -1295,12 +1295,10 @@ def run_transfer_test( backend="NIXL", transceiver_runtime="PYTHON", max_tokens_in_buffer=512, - # Keep the Python-native bounce layer disabled by default (size 0). - # Dedicated C++ bounce coverage enables the NIXL agent's bounce v2 via - # agent_bounce_buffer_enable + kv_cache_bounce_size_mb plus - # agent_bounce_params, which exercises the config plumbing end-to-end. + # Bounce is off unless expect_cpp_bounce. Then a positive kv_cache_bounce_size_mb + # enables the C++ agent bounce, and agent_bounce_params (a loose gate) exercises + # the config plumbing end to end. kv_cache_bounce_size_mb=64 if expect_cpp_bounce else 0, - agent_bounce_buffer_enable=expect_cpp_bounce, agent_bounce_params={ "min_descriptor_count": "1", "max_average_descriptor_size": "1MB", @@ -1496,11 +1494,6 @@ def run_transfer_test( # bounce v2 transport active, and at least one transfer must have been routed # through it (only KV senders submit WRITEs, so we assert on the total). if expect_cpp_bounce: - # The shared capacity must be routed to exactly one implementation: with - # agent_bounce_buffer_enable the Python-native bounce stays off. - assert all(tc._transfer_worker._config.bounce is None for tc in ctx_tcs + gen_tcs), ( - "Python-native bounce must stay disabled when the C++ agent bounce is selected" - ) agents = [tc._transfer_worker._agent for tc in ctx_tcs + gen_tcs] assert all(getattr(agent, "bounce_enabled", False) for agent in agents), ( "C++ bounce v2 transport is not active on every NIXL agent" @@ -2013,10 +2006,9 @@ def test_python_nixl_cache_transceiver_uses_cpp_bounce( ) -> None: """Exercise C++ bounce v2 through representative Python NIXL transceiver paths. - Bounce is enabled (agent_bounce_buffer_enable + kv_cache_bounce_size_mb) and tuned - (agent_bounce_params) via CacheTransceiverConfig inside run_transfer_test, covering - the config plumbing end-to-end; engagement is asserted programmatically there - (agent.bounce_enabled / bounce_submit_count). + Bounce is enabled by kv_cache_bounce_size_mb and tuned by agent_bounce_params via + CacheTransceiverConfig inside run_transfer_test, which covers the config plumbing end to + end and asserts engagement (agent.bounce_enabled / bounce_submit_count). """ # BindingsNixlTransferStatus.last_status_str() resolves this exact attribute on the C++ # status; if the binding name drifts, failure details silently degrade to "". diff --git a/tests/unittest/disaggregated/test_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index d75c51a6ca77..64315d468f70 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -495,31 +495,11 @@ def test_pipelined_transfer_allows_pipeline_parallelism_at_initialization(): assert KvCacheTransceiverV2._resolve_pipelined_transfer(transceiver, cache_transceiver_config) -def test_pipelined_transfer_rejects_bounce_buffer(): - """The Python bounce buffer stages whole requests rather than individual chunks.""" - from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 - - transceiver = object.__new__(KvCacheTransceiverV2) - transceiver._mapping = SimpleNamespace(pp_size=1, cp_size=1) - transceiver._kv_cache_manager = MagicMock() - cache_transceiver_config = CacheTransceiverConfig( - backend="NIXL", - enable_pipelined_transfer=True, - kv_cache_bounce_size_mb=1, - ) - - with pytest.raises( - ValueError, - match="not supported with the Python bounce buffer \\(kv_cache_bounce_size_mb=1", - ): - KvCacheTransceiverV2._resolve_pipelined_transfer(transceiver, cache_transceiver_config) - - -def test_pipelined_transfer_allows_agent_bounce_buffer(): - """The C++ transfer-agent bounce is compatible with per-chunk pipelined transfer. +def test_pipelined_transfer_allows_bounce_buffer(): + """The transfer-agent bounce buffer is compatible with per-chunk pipelined transfer. - It stages each transfer request on its own below the Python layer, so a pipelined - chunk is just another request. + The C++ agent stages each transfer request on its own below the Python layer, so a + pipelined chunk is just another request. """ from tensorrt_llm._torch.disaggregation.transceiver import KvCacheTransceiverV2 @@ -530,7 +510,6 @@ def test_pipelined_transfer_allows_agent_bounce_buffer(): backend="NIXL", enable_pipelined_transfer=True, kv_cache_bounce_size_mb=1, - agent_bounce_buffer_enable=True, ) assert KvCacheTransceiverV2._resolve_pipelined_transfer(transceiver, cache_transceiver_config) diff --git a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py index f9f86fed340b..56bb57220015 100644 --- a/tests/unittest/disaggregated/test_transceiver_bounded_polling.py +++ b/tests/unittest/disaggregated/test_transceiver_bounded_polling.py @@ -786,7 +786,6 @@ def _make_cache_config(**overrides) -> SimpleNamespace: kv_transfer_poll_interval_ms=5_000, kv_transfer_sender_future_timeout_ms=1_000, kv_cache_bounce_size_mb=0, - agent_bounce_buffer_enable=False, agent_bounce_params=None, enable_pipelined_transfer=False, ) @@ -809,12 +808,6 @@ def _construct_worker_config(monkeypatch, cache_config) -> TransferWorkerConfig: "tensorrt_llm._torch.disaggregation.transceiver.create_cache_reuse_adapter", Mock(return_value=Mock()), ) - # Echo the routed size: None for 0 (off, mirroring the real helper), a sentinel - # carrying the size otherwise, so tests can assert which implementation got it. - monkeypatch.setattr( - "tensorrt_llm._torch.disaggregation.transceiver.bounce_config_from_size", - Mock(side_effect=lambda size_mb: ("bounce", size_mb) if size_mb > 0 else None), - ) monkeypatch.setattr( "tensorrt_llm._torch.disaggregation.transceiver.torch.cuda.current_device", Mock(return_value=0), @@ -911,33 +904,17 @@ def test_transceiver_wires_separate_sender_slice_and_overall_timeout( assert worker_config.rx_timeout_s == expected_timeout_s -@pytest.mark.parametrize( - ("bounce_size_mb", "agent_enable", "expected_python_bounce", "expected_buffer_size_mb"), - [ - # Shared capacity, Python implementation (default): per-region bounce on, agent off. - (384, False, ("bounce", 384), 0), - # Shared capacity, C++ agent implementation: the agent bounce buffer gets the size, Python off. - (384, True, None, 384), - # Size 0 keeps both implementations off. - (0, False, None, 0), - ], -) -def test_transceiver_routes_bounce_capacity_to_one_implementation( - monkeypatch, - bounce_size_mb: int, - agent_enable: bool, - expected_python_bounce, - expected_buffer_size_mb: int, +@pytest.mark.parametrize("bounce_size_mb", [0, 384]) +def test_transceiver_routes_bounce_capacity_to_agent_buffer( + monkeypatch, bounce_size_mb: int ) -> None: + params = {"min_descriptor_count": "1"} if bounce_size_mb else None worker_config = _construct_worker_config( monkeypatch, - _make_cache_config( - kv_cache_bounce_size_mb=bounce_size_mb, - agent_bounce_buffer_enable=agent_enable, - ), + _make_cache_config(kv_cache_bounce_size_mb=bounce_size_mb, agent_bounce_params=params), ) - assert worker_config.bounce == expected_python_bounce - assert worker_config.agent_buffer_size_mb == expected_buffer_size_mb + assert worker_config.agent_buffer_size_mb == bounce_size_mb + assert worker_config.agent_bounce_params == params def test_transceiver_rejects_unset_transfer_timeout() -> None: diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 1fbc46e23b19..b44fda509b42 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -3251,24 +3251,18 @@ def test_cache_transceiver_config_arbitrary_args(self): with pytest.raises(pydantic_core._pydantic_core.ValidationError): CacheTransceiverConfig(kv_cache_bounce_size_mb=-1) - # agent_bounce_buffer_enable defaults to the Python implementation (False); enabling the - # C++ transfer-agent implementation with a zero capacity is a contradiction. + # The deprecated agent_bounce_buffer_enable defaults to False and is + # ignored: True with a zero capacity validates. assert config.agent_bounce_buffer_enable is False assert config.agent_bounce_params is None assert CacheTransceiverConfig( - kv_cache_bounce_size_mb=512, - agent_bounce_buffer_enable=True).agent_bounce_buffer_enable is True - with pytest.raises(pydantic_core._pydantic_core.ValidationError, - match="kv_cache_bounce_size_mb is 0"): - CacheTransceiverConfig(agent_bounce_buffer_enable=True) + agent_bounce_buffer_enable=True).kv_cache_bounce_size_mb == 0 # Bounce is Python-transceiver-only: the pybind (C++ transceiver) config # carries no bounce fields at all. backend must be set: from_string(None) # is a pre-existing _to_pybind limit. pybind_config = CacheTransceiverConfig( - backend="NIXL", - kv_cache_bounce_size_mb=384, - agent_bounce_buffer_enable=True)._to_pybind() + backend="NIXL", kv_cache_bounce_size_mb=384)._to_pybind() assert not any("bounce" in attr for attr in dir(pybind_config)) # Arbitrary arguments should be rejected @@ -3281,7 +3275,6 @@ def test_cache_transceiver_config_agent_bounce_params(self): """agent_bounce_params coercion, key/consistency validation, pybind passthrough.""" # Values are coerced to strings (YAML often yields ints/bools). config = CacheTransceiverConfig(kv_cache_bounce_size_mb=512, - agent_bounce_buffer_enable=True, agent_bounce_params={ "max_chunk_size": 4096, "enable_eager_gather": False @@ -3291,12 +3284,17 @@ def test_cache_transceiver_config_agent_bounce_params(self): "enable_eager_gather": "False" } - # Params without the C++ agent implementation are a contradiction: they - # would be silently ignored, so validation rejects them outright. + # Params with bounce disabled (capacity 0) are a contradiction: they + # would be silently ignored, so validation rejects them outright. The + # deprecated flag does not change that. with pytest.raises(pydantic_core._pydantic_core.ValidationError, - match="agent_bounce_buffer_enable"): + match="kv_cache_bounce_size_mb is 0"): CacheTransceiverConfig( - kv_cache_bounce_size_mb=512, + agent_bounce_params={"copy_stream_count": "2"}) + with pytest.raises(pydantic_core._pydantic_core.ValidationError, + match="kv_cache_bounce_size_mb is 0"): + CacheTransceiverConfig( + agent_bounce_buffer_enable=True, agent_bounce_params={"copy_stream_count": "2"}) # Unknown keys (here the env-var-style typo WITH the trailing _bytes) are @@ -3305,7 +3303,6 @@ def test_cache_transceiver_config_agent_bounce_params(self): match="max_chunk_size_bytes.*min_descriptor_count"): CacheTransceiverConfig( kv_cache_bounce_size_mb=512, - agent_bounce_buffer_enable=True, agent_bounce_params={"max_chunk_size_bytes": "4096"}) # Non-dict input fails cleanly in the coercion validator, not with a bare @@ -3313,9 +3310,40 @@ def test_cache_transceiver_config_agent_bounce_params(self): with pytest.raises(pydantic_core._pydantic_core.ValidationError, match="must be a dict"): CacheTransceiverConfig(kv_cache_bounce_size_mb=512, - agent_bounce_buffer_enable=True, agent_bounce_params="max_chunk_size=4096") + def test_cache_transceiver_config_agent_bounce_buffer_enable_is_deprecated_noop( + self, monkeypatch): + """agent_bounce_buffer_enable is a deprecated no-op that warns only when True.""" + warnings_seen = [] + monkeypatch.setattr( + llm_args_mod.logger, "warning_once", + lambda message, key: warnings_seen.append((message, key))) + + field = CacheTransceiverConfig.model_fields[ + "agent_bounce_buffer_enable"] + assert field.json_schema_extra["status"] == "deprecated" + + # True only warns: the config is otherwise identical to the flag-less one. + config = CacheTransceiverConfig(kv_cache_bounce_size_mb=512, + agent_bounce_buffer_enable=True) + assert len(warnings_seen) == 1 + message, key = warnings_seen[0] + assert "deprecated" in message + assert key == "cache_transceiver_config.agent_bounce_buffer_enable" + assert config.model_dump( + exclude={"agent_bounce_buffer_enable"}) == CacheTransceiverConfig( + kv_cache_bounce_size_mb=512).model_dump( + exclude={"agent_bounce_buffer_enable"}) + # Re-validating a full dump (as apply_model_defaults_to_llm_args does) + # and an explicit False stay silent. + warnings_seen.clear() + CacheTransceiverConfig(**CacheTransceiverConfig( + kv_cache_bounce_size_mb=512).model_dump()) + CacheTransceiverConfig(kv_cache_bounce_size_mb=512, + agent_bounce_buffer_enable=False) + assert warnings_seen == [] + def test_agent_bounce_param_keys_match_cpp_env_knobs(self): """AGENT_BOUNCE_PARAM_KEYS must mirror kEnvKnobs in BounceConfig.h (parsed here).""" from tensorrt_llm._torch.disaggregation.nixl.bounce_knobs import \ From 9161bc0aef1bd29df19aa22d2bc71d79eb6e9d28 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:17:56 +0800 Subject: [PATCH 02/11] [None][chore] Remove the Python KV bounce buffer from the Python transceiver Delete tensorrt_llm/_torch/disaggregation/native/bounce and every Python-transceiver adaptation that existed only to serve it: - bounce ctor args and TransferWorkerConfig.bounce; - the receiver reservation lifecycle (reserve, writer_base, idle-release, orphan, abort_publication, record_failure) and the fan-in gate; - the sender send-slot lifecycle and the KV_AGENT_RESULT bounce tail; - RecvReqInfo/WriteMeta.bounce_dst_base on the wire; - region sizing for recurrent state (mamba_receiver_payload_bytes); - the deferred scatter completion: the receiver now completes the task inline under RxSession.lock, as the non-bounced path already did, so wait_for_local_completion/finish_local_completion go away. RecvReqInfo no longer carries bounce_dst_base, so context and generation servers must be upgraded together. Tests that only covered the Python bounce are deleted; the KV result frame tests move to test_transfer_late_settlement.py and the helix MambaPolicy tests to test_rank_info.py. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../disaggregation/native/bounce/__init__.py | 55 - .../disaggregation/native/bounce/buffer.py | 192 --- .../disaggregation/native/bounce/config.py | 153 --- .../disaggregation/native/bounce/core.py | 275 ---- .../native/bounce/gather_scatter.py | 233 ---- .../disaggregation/native/bounce/impl.py | 649 ---------- .../disaggregation/native/mixers/ssm/peer.py | 43 +- .../_torch/disaggregation/native/transfer.py | 398 +----- .../_torch/disaggregation/nixl/_agent_cpp.py | 4 +- .../_torch/disaggregation/transceiver.py | 6 - .../connectors/mooncake_store/staging.py | 6 +- tests/integration/defs/.test_durations | 1 - .../integration/test_lists/test-db/l0_a10.yml | 1 - .../test_lists/test-db/l0_h100.yml | 1 - .../disaggregated/lifecycle_acceptance.md | 2 +- tests/unittest/disaggregated/test_bounce.py | 1121 ----------------- .../disaggregated/test_chunked_transfer.py | 1 - .../disaggregated/test_kda_mamba_transfer.py | 14 +- .../test_kv_transfer_fail_stop.py | 6 +- .../disaggregated/test_mamba_transfer.py | 88 -- .../unittest/disaggregated/test_rank_info.py | 54 + .../disaggregated/test_retirement_deadline.py | 19 +- .../disaggregated/test_task_handle.py | 66 +- .../test_transfer_late_settlement.py | 57 +- .../test_transfer_lifecycle_acceptance.py | 2 - .../test_transfer_ownership_regressions.py | 188 +-- .../disaggregated/test_unsubmitted_session.py | 14 +- .../test_verified_reuse_admission.py | 9 +- 28 files changed, 181 insertions(+), 3477 deletions(-) delete mode 100644 tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py delete mode 100644 tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py delete mode 100644 tensorrt_llm/_torch/disaggregation/native/bounce/config.py delete mode 100644 tensorrt_llm/_torch/disaggregation/native/bounce/core.py delete mode 100644 tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py delete mode 100644 tensorrt_llm/_torch/disaggregation/native/bounce/impl.py delete mode 100644 tests/unittest/disaggregated/test_bounce.py diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py b/tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py deleted file mode 100644 index f014842d5254..000000000000 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/__init__.py +++ /dev/null @@ -1,55 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Opt-in (TransferWorkerConfig.bounce) VRAM d2d KV bounce buffering: coalesce a -transfer's scattered per-block KV into ONE contiguous fabric-VMM WRITE (reliable -cuda_ipc/MNNVL) through the BounceTransport interface; default per-block path is -unchanged when no Config is given. config_from_size() is the on/off switch.""" - -from .buffer import Buffer, SlotAllocator -from .config import Config, FixedSizing, Sizing, SizingContext, config_from_size -from .core import BounceTransport, Disposition, ScatterState, TransferContext, TransferState -from .gather_scatter import Plan -from .impl import ( - NoBounceTransport, - VmmBounceTransport, - build_send_request, - create_bounce, - decode_result_tail, - encode_result_tail, - scatter_write_result, -) - -__all__ = [ - "BounceTransport", - "Buffer", - "Config", - "Disposition", - "FixedSizing", - "NoBounceTransport", - "Plan", - "ScatterState", - "Sizing", - "SizingContext", - "SlotAllocator", - "TransferContext", - "TransferState", - "VmmBounceTransport", - "build_send_request", - "config_from_size", - "create_bounce", - "decode_result_tail", - "encode_result_tail", - "scatter_write_result", -] diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py b/tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py deleted file mode 100644 index c45ff75297e8..000000000000 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/buffer.py +++ /dev/null @@ -1,192 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Fabric-VMM bounce buffers. A fabric region lets the write ride the fast intra-node fabric, which a -plain device allocation cannot; it is allocated once at setup and reused.""" - -import threading -import time -from typing import Dict, Optional, Tuple - -from tensorrt_llm import logger -from tensorrt_llm._torch.disaggregation.base.agent import RegMemoryDescs -from tensorrt_llm.runtime.kv_cache_manager_v2._introspection import PooledPhysMemAllocator, VirtMem - -_MIB = 1024 * 1024 - - -def _div_up(a: int, b: int) -> int: - return (a + b - 1) // b - - -class Buffer: - """One contiguous fabric region for coalescing cache data. The physical chunk size must match the - cache pool's chunk size, which the C++ splitter relies on.""" - - __slots__ = ("_device_id", "_name", "_size", "_vm") - - def __init__(self, capacity_bytes: int, phys_chunk_size: int, name: str = "kv_bounce"): - if capacity_bytes <= 0 or phys_chunk_size <= 0: - raise ValueError( - f"Buffer: capacity_bytes={capacity_bytes}, " - f"phys_chunk_size={phys_chunk_size} must both be > 0" - ) - vm_size = _div_up(capacity_bytes, phys_chunk_size) * phys_chunk_size - allocator = PooledPhysMemAllocator(phys_chunk_size) - # back the whole region up front so its address is writable and stable for life - self._vm = VirtMem(vm_size, allocator, init_num_phys_mem=vm_size // phys_chunk_size) - self._size = vm_size - self._device_id = allocator.device_id - self._name = name - logger.info( - f"[kv-bounce] allocated fabric bounce buffer '{name}': " - f"{vm_size / _MIB:.1f} MiB @ 0x{int(self._vm.address):x} " - f"(chunk={phys_chunk_size // _MIB}MiB, dev={self._device_id})" - ) - - @property - def base_ptr(self) -> int: - return int(self._vm.address) - - @property - def size(self) -> int: - return self._size - - @property - def device_id(self) -> int: - return self._device_id - - def reg_descs(self) -> "RegMemoryDescs": - # the type is the string "VRAM", not the enum, because the agent upper-cases it - return RegMemoryDescs("VRAM", [(self.base_ptr, self._size, self._device_id, self._name)]) - - def close(self) -> None: - vm = getattr(self, "_vm", None) - if vm is not None: - vm.destroy() - self._vm = None # type: ignore[assignment] - - def __del__(self): - # A destructor must never raise, but a leaked region should be visible, so log the failure. - try: - self.close() - except Exception as e: - logger.debug(f"[kv-bounce] buffer '{getattr(self, '_name', '?')}' cleanup failed: {e}") - - -# region starts are rounded to this for copy alignment (negligible waste) -_ALIGN = 512 - - -class SlotAllocator: - """First-fit allocator over one fabric buffer. Regions may be freed in any order, and first-fit - reuses a hole freed out of order rather than skipping it. Reserve is thread-safe and blocking. - The whole buffer is one registration, so a write can stripe across the network links.""" - - __slots__ = ("_buf", "_cap", "_cv", "_in_use", "_quarantine", "_next_slot_id") - - def __init__(self, capacity_bytes: int, phys_chunk_size: int, name: str = "kv_bounce"): - if capacity_bytes <= 0: - raise ValueError(f"SlotAllocator: capacity_bytes={capacity_bytes} must be > 0") - self._buf = Buffer(capacity_bytes, phys_chunk_size, name=name) - self._cap = self._buf.size # rounded up to a chunk multiple - self._in_use: Dict[int, Tuple[int, int]] = {} # each live slot maps to its start and size - # Quarantined slots not yet reusable: an orphaned writer's write may still be landing - # and cannot be aborted, so each is held out of the pool until its deadline passes. - self._quarantine: Dict[int, Tuple[int, int, float]] = {} - self._next_slot_id = 0 - self._cv = threading.Condition(threading.Lock()) - - @property - def capacity(self) -> int: - return self._cap - - def _occupied(self): - """Ranges that must not be handed out: live and quarantined, treated the same.""" - for s, n in self._in_use.values(): - yield s, n - for s, n, _dl in self._quarantine.values(): - yield s, n - - def _find_free_start(self, size: int) -> Optional[int]: - """Lowest free gap large enough, or None if none fits. Live and quarantined regions both - block reuse, so an out-of-order-freed hole is reused but a quarantined one is not.""" - cursor = 0 - for s, n in sorted(self._occupied()): - if s - cursor >= size: - return cursor - cursor = max(cursor, s + n) - return cursor if self._cap - cursor >= size else None - - def reserve(self, size: int, timeout: Optional[float] = None) -> Optional[Tuple[int, int]]: - """Reserve a contiguous region, or None if it can never fit or nothing frees within the - timeout.""" - size = _div_up(size, _ALIGN) * _ALIGN - if size <= 0 or size > self._cap: - return None - deadline = None if timeout is None else time.monotonic() + timeout - with self._cv: - while True: - start = self._find_free_start(size) - if start is not None: - slot_id = self._next_slot_id - self._next_slot_id += 1 - self._in_use[slot_id] = (start, size) - return slot_id, self._buf.base_ptr + start - if deadline is None: - self._cv.wait() - else: - remaining = deadline - time.monotonic() - if remaining <= 0 or not self._cv.wait(timeout=remaining): - return None - - def release(self, slot_id: int) -> None: - with self._cv: - self._in_use.pop(slot_id, None) - self._cv.notify_all() - - def quarantine(self, slot_id: int, grace_s: float) -> None: - """Hold a slot out of the free pool for the grace period instead of releasing it, because its - region may still be under an in-doubt write. An infinite grace holds it until close.""" - with self._cv: - entry = self._in_use.pop(slot_id, None) - if entry is not None: - start, size = entry - # a finite time plus infinity is infinity, so an infinite grace never expires - deadline = time.monotonic() + grace_s - self._quarantine[slot_id] = (start, size, deadline) - self._cv.notify_all() - - def reclaim_expired(self) -> int: - """Return quarantined slots past their deadline to the free pool and report how many. Runs - off a timer, not tied to reserve, so it makes progress even when the arena is full.""" - now = time.monotonic() - with self._cv: - expired = [sid for sid, (_s, _n, dl) in self._quarantine.items() if dl <= now] - for sid in expired: - del self._quarantine[sid] - if expired: - self._cv.notify_all() - return len(expired) - - @property - def quarantined_bytes(self) -> int: - """Bytes currently held in quarantine, for observability.""" - return sum(n for _s, n, _dl in self._quarantine.values()) - - def reg_descs(self) -> "RegMemoryDescs": - return self._buf.reg_descs() - - def close(self) -> None: - self._buf.close() diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/config.py b/tensorrt_llm/_torch/disaggregation/native/bounce/config.py deleted file mode 100644 index 6a6976ab6ae8..000000000000 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/config.py +++ /dev/null @@ -1,153 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Bounce configuration and pluggable sizing policy. A config enables bounce; leaving it unset keeps -the per-block path. The size knob doubles as the on and off switch.""" - -import os -from dataclasses import dataclass, field -from typing import Optional - -from tensorrt_llm import logger - -_MIB = 1024 * 1024 - -# Test/advanced overrides for the size gates below (users only tune the bounce size). Read on the -# generation side, so set them there; unset uses the defaults. -# - min_bytes gates payloads that carry recurrent (mamba/KDA) state: the fallback cost scales -# with bytes, not block count, so the gate is byte-denominated. -# - min_blocks is the legacy plain-KV gate, kept so existing bounce deployments see no behavior -# change. -# For Kimi K3 the byte gate never rejects: the fixed ~433 MiB per-request KDA payload always -# clears the 2 MiB default, so arena capacity plus reservation backpressure is the effective -# admission control. -_MIN_BYTES_ENV = "TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES" # byte gate for recurrent-state payloads -_MIN_BLOCKS_ENV = "TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS" # block-count gate for plain-KV payloads - - -def _env_int_gate(name: str, default: int) -> int: - """Read a gate from the env, defensively: unset or malformed falls back to the default (never - crashing), and the value is clamped to at least 1.""" - raw = os.environ.get(name) - if raw is None: - return default - try: - value = int(raw) - except ValueError: - logger.warning(f"{name}={raw!r} is not an integer; using default {default}") - return default - if value < 1: - logger.warning(f"{name}={value} < 1; clamping to 1 (bounce always clears the gate)") - return 1 - return value - - -def _round_up(a: int, b: int) -> int: - return (a + b - 1) // b * b - - -@dataclass(frozen=True) -class SizingContext: - free_bytes: int # free at setup, after the cache pool claimed its fraction - total_bytes: int - chunk_bytes: int - device_id: int - - -@dataclass(frozen=True) -class Sizing: - """Returns the byte size of one region; there are two, one for sending and one for receiving.""" - - def resolve(self, ctx: SizingContext) -> int: - raise NotImplementedError - - -# Default size in MiB per region. Raise it to bounce larger single transfers, lower it to save -# memory. It is clamped to the free-memory budget at setup. -DEFAULT_CAPACITY_MB = 384 - - -@dataclass(frozen=True) -class FixedSizing(Sizing): - """A fixed capacity per region, clamped to free memory at setup.""" - - capacity_mb: int = DEFAULT_CAPACITY_MB - - def resolve(self, ctx: SizingContext) -> int: - return max(_round_up(self.capacity_mb * _MIB, ctx.chunk_bytes), ctx.chunk_bytes) - - -# bounce takes at most this fraction of the free memory left after the cache pool -_HEADROOM_FRACTION = 0.5 - - -def fit_within_free( - capacity_bytes: int, - *, - free_bytes: int, - chunk_bytes: int, - max_free_fraction: float = _HEADROOM_FRACTION, -) -> Optional[int]: - """Clamp each region so the two together stay within the allowed fraction of free memory, rounded - to a chunk. Returns None if not even one chunk fits.""" - budget_per_dir = (int(free_bytes * max_free_fraction) // 2 // chunk_bytes) * chunk_bytes - if budget_per_dir < chunk_bytes: - return None - capacity_bytes = min(capacity_bytes, budget_per_dir) - capacity_bytes = max(capacity_bytes, chunk_bytes) - return capacity_bytes - - -# Byte gate for transfers that carry recurrent (mamba/KDA) state. The cost this gate guards scales -# with BYTES, not blocks: the block-count gate (96 blocks, calibrated for 128-token blocks) -# silently skipped a 433 MiB Kimi-K3 transfer (67 blocks of 32 tokens plus the non-paged KDA state) -# and dropped it onto the ~0.4 GB/s host-staged fallback, a ~1000x cliff. Break-even is small: -# bounce adds one gather plus one scatter copy (device-local, ~hundreds of GB/s) and ~0.1 ms of -# fixed launch/reservation overhead, while the in-place path can be as slow as ~0.4 GB/s -# inter-node — 2 MiB in-place at that rate is ~5 ms vs well under 1 ms bounced. Below 2 MiB the -# fixed overhead dominates and arena slots are better kept for large transfers. Heuristic, tunable -# via TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES. -DEFAULT_MIN_BYTES = 2 * _MIB - - -@dataclass -class Config: - sizing: Sizing = field(default_factory=FixedSizing) # how much memory to reserve (pluggable) - chunk_mb: int = 32 # physical chunk size; a large chunk keeps the write to a single descriptor - # Which gate applies depends on the payload (see VmmBounceTransport.reserve): transfers that - # carry recurrent state use min_bytes; plain-KV transfers keep the original min_blocks gate so - # pre-existing bounce deployments see no behavior change. - # byte gate for recurrent-state payloads (see DEFAULT_MIN_BYTES for the rationale) - min_bytes: int = DEFAULT_MIN_BYTES - # block-count gate for plain-KV payloads (roughly 12k tokens at 128 per block); heuristic, - # tunable via TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS - min_blocks: int = 96 - - -def config_from_size( - size_mb: int, min_blocks: Optional[int] = None, min_bytes: Optional[int] = None -) -> Optional[Config]: - """Build a bounce config from a per-region size in MiB, or None to leave bounce off (size <= 0). - Size is both the capacity and the on/off switch. min_bytes (recurrent-state payloads) and - min_blocks (plain-KV payloads) are the gates below which a transfer stays on the per-block - path; when unset they come from the env, else the defaults.""" - if size_mb is None or size_mb <= 0: - return None - if min_blocks is None: - min_blocks = _env_int_gate(_MIN_BLOCKS_ENV, Config.min_blocks) - if min_bytes is None: - min_bytes = _env_int_gate(_MIN_BYTES_ENV, Config.min_bytes) - return Config( - sizing=FixedSizing(capacity_mb=size_mb), min_bytes=min_bytes, min_blocks=min_blocks - ) diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/core.py b/tensorrt_llm/_torch/disaggregation/native/bounce/core.py deleted file mode 100644 index c193de0ce996..000000000000 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/core.py +++ /dev/null @@ -1,275 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Pure core of the KV bounce subsystem: the transport contract and the per-region lifetime state -machine. No CUDA or NIXL imports, so it is unit-testable on CPU.""" - -from abc import ABC, abstractmethod -from dataclasses import dataclass, field -from enum import Enum -from typing import Callable, Dict, List, Optional, Tuple - - -class TransferState(Enum): - """Lifecycle of one bounced receive region.""" - - INIT = "init" # leased, not yet advertised to any sender - ACTIVE = "active" # writers in flight - SCATTERING = "scattering" # all writers succeeded; scattering back into the cache - COMPLETED = "completed" # scattered cleanly, slot released - FAILED = "failed" # all writers done, at least one failed, slot released - QUARANTINED = "quarantined" # a writer is in doubt, slot held out of reuse - - -class ScatterState(Enum): - IDLE = "idle" - QUEUED = "queued" - DONE = "done" - FAILED = "failed" - - -class Disposition(Enum): - """What to do with the slot when the transfer finishes.""" - - RELEASE = "release" # drained: return it to the free pool - QUARANTINE = "quarantine" # in doubt: hold it out of the pool for a grace period - - -@dataclass -class Settlement: - """What the transport carries out on finish: the slot, whether to release or quarantine it, - whether the transfer succeeded, and the callback to fire.""" - - slot_id: int - disposition: Disposition - success: bool - on_done: Optional[Callable[[bool], None]] - - -@dataclass -class TransferContext: - """Lifetime state machine for one bounced receive region. The region is released only after every - writer reports (success or failure both mean its write has drained); a writer given up on (a - cancelled or timed-out transfer, whose one-sided write cannot be aborted) makes it quarantined - instead of reused.""" - - rid_slice: Tuple[int, int] - slot_id: int - base_addr: int - per_writer_bytes: int - num_writers: int - on_done: Optional[Callable[[bool], None]] = None - - # Whether each writer that reported succeeded, keyed by rank; presence means it reported. - _writer_ok: Dict[int, bool] = field(default_factory=dict) - # per successful writer: where it wrote, plus the fragments to scatter back - _scatter_descs: List[tuple] = field(default_factory=list) - _writer_cohort: Optional[frozenset[int]] = None - _publication_failed: bool = False - _orphaned: bool = False - scatter_state: ScatterState = ScatterState.IDLE - state: TransferState = TransferState.INIT - settled: bool = False - - def writer_base(self, writer_index: int) -> int: - """Where this fan-in writer writes in the region.""" - return self.base_addr + writer_index * self.per_writer_bytes - - def _writers_final(self) -> bool: - return self.settled or self.state in ( - TransferState.SCATTERING, - TransferState.COMPLETED, - TransferState.FAILED, - TransferState.QUARANTINED, - ) - - def _all_writers_reported(self) -> bool: - if self._writer_cohort is not None: - return self._writer_cohort.issubset(self._writer_ok) - return len(self._writer_ok) >= self.num_writers - - def _all_writers_succeeded(self) -> bool: - if not self._all_writers_reported(): - return False - if self._writer_cohort is not None: - return all(self._writer_ok[rank] for rank in self._writer_cohort) - return all(self._writer_ok.values()) - - # Mutations: call only while holding the transport's reservation lock. - def record_writer_result( - self, - peer_rank: int, - succeeded: bool, - *, - src_base: Optional[int] = None, - dst_ptrs=None, - sizes=None, - ) -> None: - """Record one writer's terminal report. Repeat or late reports are ignored, so duplicate or - out-of-order messages are harmless.""" - # A duplicate or late report can still arrive after the writer set is final, because - # scatter runs on another thread while the region stays live. For example a retransmitted - # notification, or a stray failure that would flip a good transfer to failed; drop it. - if self._writers_final() or peer_rank in self._writer_ok: - return - if self._writer_cohort is not None and peer_rank not in self._writer_cohort: - raise RuntimeError(f"writer {peer_rank} is outside the published bounce cohort") - self._writer_ok[peer_rank] = succeeded - if succeeded and dst_ptrs is not None and int(dst_ptrs.size) > 0: - self._scatter_descs.append( - (src_base if src_base is not None else self.base_addr, dst_ptrs, sizes) - ) - if self.state is TransferState.INIT: - self.state = TransferState.ACTIVE - - def mark_orphaned(self) -> None: - """Give up on writers that never reported (cancel, timeout, shutdown): a one-sided write - cannot be aborted, so the region is quarantined on settle. No-op once scattering or done.""" - if self._writers_final(): - return - self._orphaned = True - - def abort_publication(self, published_writers: set[int]) -> None: - """Close a failed fan-out around only the writers that were published. - - Published writers still have to report terminal evidence. Successful - partial data is intentionally not scattered because the request is - already incomplete. - """ - if self._writers_final(): - return - published = frozenset(published_writers) - if len(published) > self.num_writers or not self._writer_ok.keys() <= published: - raise RuntimeError( - "published bounce cohort is inconsistent with terminal evidence: " - f"reported={sorted(self._writer_ok)} published={sorted(published)} " - f"reserved_count={self.num_writers}" - ) - self._writer_cohort = published - self._publication_failed = True - - def begin_scatter(self) -> None: - self.state = TransferState.SCATTERING - self.scatter_state = ScatterState.QUEUED - - def sorted_scatter_descs(self) -> List[tuple]: - """Success fragments ordered by source, for a deterministic scatter.""" - return sorted(self._scatter_descs, key=lambda t: t[0]) - - def finish_scatter(self, ok: bool) -> None: - self.scatter_state = ScatterState.DONE if ok else ScatterState.FAILED - - def ready_to_scatter(self) -> bool: - """All writers succeeded, there is data to scatter, and scatter has not started.""" - return ( - self.state is TransferState.ACTIVE - and not self._orphaned - and not self._publication_failed - and self._all_writers_succeeded() - and bool(self._scatter_descs) - ) - - def ready_to_settle(self) -> bool: - if self.settled: - return False - if self._orphaned: - return True # in doubt: settle now and quarantine - if not self._all_writers_reported(): - return False # a writer has not reported yet - if self._publication_failed: - return True # every published writer drained; incomplete data is discarded - if self._all_writers_succeeded() and self._scatter_descs: - return self.scatter_state in (ScatterState.DONE, ScatterState.FAILED) - return True # nothing to scatter, or a failure among them: either way drained - - def settle(self) -> Optional[Settlement]: - """Finalize the transfer once and return the decision, or None if already done. Every - termination path calls this; only the first call has any effect.""" - if self.settled: - return None - self.settled = True - if self._orphaned: - self.state = TransferState.QUARANTINED - return Settlement(self.slot_id, Disposition.QUARANTINE, False, self.on_done) - success = ( - not self._publication_failed - and self._all_writers_succeeded() - and self.scatter_state is not ScatterState.FAILED - ) - self.state = TransferState.COMPLETED if success else TransferState.FAILED - return Settlement(self.slot_id, Disposition.RELEASE, success, self.on_done) - - -class BounceTransport(ABC): - """Contract both transports implement so they cannot drift; the factory returns one, and the rest - of the code depends only on this interface.""" - - #: True for a real transport, False for the disabled null object. - enabled: bool = False - - @abstractmethod - def build_request(self, write_meta): - """Gather the request's cache into a send region and build the coalesced write, or None to - fall back to the per-fragment path.""" - - @abstractmethod - def release_send(self, slot_id) -> None: - """Release a send region after its write completes.""" - - @abstractmethod - def reserve( - self, - recv_req, - num_writers: int = 1, - *, - timeout: Optional[float] = None, - extra_bytes: int = 0, - ) -> bool: - """Reserve a region and record its address for the senders. False falls back to per-fragment. - ``extra_bytes`` sizes the sender's non-paged payload (mamba/KDA recurrent state) that rides - the same coalesced write as the KV blocks.""" - - @abstractmethod - def writer_base(self, rid_slice, writer_index: int) -> Optional[int]: - """Where the given fan-in writer writes in the region.""" - - @abstractmethod - def is_bounced(self, rid_slice) -> bool: - """Whether this request and slice name a live bounced region.""" - - @abstractmethod - def release_idle_reservation(self, rid_slice) -> None: - """Immediately free a reservation cancelled before any address went out.""" - - @abstractmethod - def orphan_reservation(self, rid_slice) -> None: - """Give up on an in-flight reservation (cancel/timeout/lost result); quarantine, don't leak.""" - - @abstractmethod - def abort_publication(self, rid_slice, published_writers: set[int]) -> None: - """Limit a failed fan-out to writers whose REQUEST_DATA was queued.""" - - @abstractmethod - def record_result( - self, rid_slice, peer_rank, dst_ptrs=None, sizes=None, src_base=None, on_done=None - ) -> None: - """Handle a writer's success; scatter and finalize once all writers reported.""" - - @abstractmethod - def record_failure(self, rid_slice, peer_rank) -> None: - """Handle a writer's failure; free the region once all writers reported.""" - - @abstractmethod - def close(self) -> None: - """Stop the scatter worker, deregister memory, free the arenas.""" diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py b/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py deleted file mode 100644 index 21b81b250ff5..000000000000 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/gather_scatter.py +++ /dev/null @@ -1,233 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Device-to-device gather and scatter for KV bounce. It coalesces the scattered fragments into one -contiguous buffer, and the inverse, in a single batched kernel launch, and falls back to a -per-fragment copy loop when Triton or the GPU is unavailable.""" - -import threading -from dataclasses import dataclass - -import numpy as np - -try: - from cuda.bindings import runtime as cudart -except ImportError: - from cuda import cudart - -from tensorrt_llm._utils import CUASSERT, prefer_pinned - -try: - import torch - import triton - import triton.language as tl - - _HAVE_TRITON = True -except ImportError: # keep importable without GPU or Triton at import time - _HAVE_TRITON = False - -_D2D = cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice - -# copy in 128-bit elements, the widest per-thread transaction, as two 64-bit words -_ELEM_BYTES = 16 - -# elements per program; an architecture-agnostic default of about 64 KiB of copy -_BLOCK = 4096 -# fixed rather than autotuned, to avoid recompiling per shape on the hot path -_NUM_WARPS = 4 -_CHUNK_BYTES = _BLOCK * _ELEM_BYTES - - -@dataclass -class Plan: - """One transfer's coalescing plan: the source and destination fragment addresses, their sizes, - and the total. Both sides walk the fragments in the same order, and the offsets are the running - sum of the sizes.""" - - src_ptrs: np.ndarray # source fragment addresses - dst_ptrs: np.ndarray # destination fragment addresses - sizes: np.ndarray # bytes per fragment - total_size: int # total bytes; must not exceed the bounce buffer capacity - - @property - def offsets(self) -> np.ndarray: - if self.sizes.size == 0: - return np.zeros(0, dtype=np.int64) - off = np.empty(self.sizes.size, dtype=np.int64) - off[0] = 0 - np.cumsum(self.sizes[:-1], out=off[1:]) - return off - - -if _HAVE_TRITON: - - @triton.jit - def _batched_copy_kernel( - dst_ptrs_ptr, # destination address per fragment - src_ptrs_ptr, # source address per fragment - nvec_ptr, # 128-bit element count per fragment - n_frags, # number of fragments - BLOCK: tl.constexpr, # elements copied per program - ): - # each program copies one block-sized slice of one fragment - frag = tl.program_id(0) - chunk = tl.program_id(1) - if frag >= n_frags: - return - - dst_addr = tl.load(dst_ptrs_ptr + frag) - src_addr = tl.load(src_ptrs_ptr + frag) - nvec = tl.load(nvec_ptr + frag) - - vec = chunk * BLOCK + tl.arange(0, BLOCK) # element index within the fragment - mask = vec < nvec # masks off the fragment's tail - - # each row is two 64-bit words the compiler coalesces into one 128-bit access - i64 = vec[:, None] * 2 + tl.arange(0, 2)[None, :] - m2 = mask[:, None] - src_i64 = src_addr.to(tl.uint64).to(tl.pointer_type(tl.int64)) - dst_i64 = dst_addr.to(tl.uint64).to(tl.pointer_type(tl.int64)) - - vals = tl.load(src_i64 + i64, mask=m2) - tl.store(dst_i64 + i64, vals, mask=m2) - - -def _uniform_nelem(sizes: np.ndarray): - """Return the element counts, the maximum, and whether every fragment is 16-byte aligned, which - the 128-bit copy requires.""" - if sizes.size == 0: - return None, 0, False - if not np.all((sizes % _ELEM_BYTES) == 0): - return None, 0, False - nvec = (sizes // _ELEM_BYTES).astype(np.int64) - return nvec, int(nvec.max()), True - - -# reusable pinned staging for the metadata copy, one per stream, so the copy is a true async transfer -_meta_lock = threading.Lock() -_meta_buffers = {} # each stream maps to its pinned host buffer, device buffer, and capacity - - -def _get_meta_buffers(stream_handle: int, need: int, dev): - """Return the pinned host and device buffers for the stream, large enough for the request, - growing them under a lock.""" - with _meta_lock: - ent = _meta_buffers.get(stream_handle) - if ent is None or ent[2] < need: - new_cap = max(need, (ent[2] * 2 if ent else 0), 4096) - pinned = torch.empty(new_cap, dtype=torch.int64, pin_memory=prefer_pinned()) - devt = torch.empty(new_cap, dtype=torch.int64, device=dev) - ent = (pinned, devt, new_cap) - _meta_buffers[stream_handle] = ent - return ent[0], ent[1] - - -def _launch_batched_copy( - dst_addrs: np.ndarray, src_addrs: np.ndarray, sizes: np.ndarray, stream -) -> bool: - """Run the single batched copy on the stream. Returns False when the caller must use the loop - fallback.""" - if not _HAVE_TRITON or not torch.cuda.is_available(): - return False - - n = int(src_addrs.size) - if n == 0: - return True # nothing to copy, trivially done - - nvec, max_nvec, ok = _uniform_nelem(sizes) - if not ok or max_nvec == 0: - return False - - dev = torch.device("cuda", torch.cuda.current_device()) - # pack all three address arrays into one pinned buffer and one async copy, so the kernel is - # ordered after it - stream_handle = int(stream) - pinned, devt = _get_meta_buffers(stream_handle, 3 * n, dev) - host = pinned.numpy() - host[:n] = dst_addrs - host[n : 2 * n] = src_addrs - host[2 * n : 3 * n] = nvec - - n_chunks = triton.cdiv(max_nvec, _BLOCK) - grid = (n, n_chunks) - - # buffer reuse is safe: every caller syncs the stream between calls, so the previous copy and - # kernel finish before the refill - ext_stream = torch.cuda.ExternalStream(stream_handle) - with torch.cuda.stream(ext_stream): - devt[: 3 * n].copy_(pinned[: 3 * n], non_blocking=True) - _batched_copy_kernel[grid]( - devt[:n], - devt[n : 2 * n], - devt[2 * n : 3 * n], - n, - BLOCK=_BLOCK, - num_warps=_NUM_WARPS, - ) - return True - - -def _copy_frags(pairs, sizes: np.ndarray, stream) -> None: - """Fallback used when the batched copy is unavailable: one async copy per fragment.""" - # strict zip: a length mismatch must fail fast rather than silently drop part of the copy - for (dst, src), n in zip(pairs, sizes, strict=True): - CUASSERT(cudart.cudaMemcpyAsync(int(dst), int(src), int(n), _D2D, stream)) - - -def gather_contiguous( - dst_base: int, - src_ptrs: np.ndarray, - sizes: np.ndarray, - offsets: np.ndarray, - *, - stream, -) -> None: - """Gather each source fragment into its place in the contiguous buffer, asynchronously. The - caller syncs before issuing the write.""" - src_addrs = np.asarray(src_ptrs, dtype=np.int64) - dst_addrs = np.int64(dst_base) + np.asarray(offsets, dtype=np.int64) - sizes = np.asarray(sizes, dtype=np.int64) - - if _launch_batched_copy(dst_addrs, src_addrs, sizes, stream): - return - - _copy_frags( - ((int(dst_addrs[k]), int(src_addrs[k])) for k in range(src_addrs.size)), - sizes, - stream, - ) - - -def scatter_contiguous( - src_base: int, - dst_ptrs: np.ndarray, - sizes: np.ndarray, - offsets: np.ndarray, - *, - stream, -) -> None: - """The inverse of gather: scatter each piece of the contiguous buffer back to its destination - fragment, asynchronously. The caller syncs before signaling completion.""" - dst_addrs = np.asarray(dst_ptrs, dtype=np.int64) - src_addrs = np.int64(src_base) + np.asarray(offsets, dtype=np.int64) - sizes = np.asarray(sizes, dtype=np.int64) - - if _launch_batched_copy(dst_addrs, src_addrs, sizes, stream): - return - - _copy_frags( - ((int(dst_addrs[k]), int(src_addrs[k])) for k in range(dst_addrs.size)), - sizes, - stream, - ) diff --git a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py b/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py deleted file mode 100644 index b6bf6e81addc..000000000000 --- a/tensorrt_llm/_torch/disaggregation/native/bounce/impl.py +++ /dev/null @@ -1,649 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""The two KV bounce transports (the real fabric-VMM one and the disabled null object) implementing -the contract in core.py. Holds the buffers, the gather and scatter kernels, and the scatter worker, -and runs the side effects that drive each region's state machine. Never imports transfer.py.""" - -from __future__ import annotations - -import queue -import threading -from typing import TYPE_CHECKING, Callable, Dict, List, Optional - -import numpy as np - -try: - from cuda.bindings import runtime as cudart -except ImportError: - from cuda import cudart - -from tensorrt_llm import logger -from tensorrt_llm._torch.disaggregation.base.agent import ( - MemoryDescs, - MemoryType, - TransferOp, - TransferRequest, -) -from tensorrt_llm._utils import CUASSERT - -from .buffer import SlotAllocator -from .config import DEFAULT_MIN_BYTES, SizingContext, fit_within_free -from .core import BounceTransport, Disposition, Settlement, TransferContext -from .gather_scatter import Plan, gather_contiguous, scatter_contiguous - -if TYPE_CHECKING: - from tensorrt_llm._torch.disaggregation.resource.page import KVCachePageTable - -RidSlice = tuple # the request id and slice id a region serves -_MIB = 1024 * 1024 -_SCATTER_POLL_S = 0.5 # how often the scatter worker wakes to re-check the stop flag and reclaim -# Send side: a sender (KV transfer worker thread) may wait this long for bounce space before -# falling back to per-fragment writes. -_RESERVE_TIMEOUT_S = 0.2 -# Receive side: never wait. The receiver reserves on the executor thread while it admits -# generation requests, so any wait pauses the whole iteration (and, under attention DP, every rank -# of the group). If the recv region is full, use the per-fragment path right away; a timed-out -# wait ends in the same fallback, and the transfer stays asynchronous either way. -_RECV_RESERVE_TIMEOUT_S = 0.0 -_CLOSE_JOIN_S = 2.0 # max wait for the scatter thread to drain on close -_QUARANTINE_GRACE_S = 60.0 # how long an orphaned region is held out of reuse - - -class VmmBounceTransport(BounceTransport): - """The real transport: gather the request's cache into one fabric region, issue a single coalesced - multi-rail write, and scatter it back on the receiver.""" - - enabled = True - - @classmethod - def from_config( - cls, agent, cfg, *, device_id: int, block_bytes_per_group: list[int | None] - ) -> VmmBounceTransport | None: - """Build a transport sized from the config and clamped to free memory, or None if not even one - chunk fits.""" - chunk = cfg.chunk_mb * _MIB - free_b, total_b = CUASSERT(cudart.cudaMemGetInfo()) - want_capacity = cfg.sizing.resolve( - SizingContext( - free_bytes=free_b, total_bytes=total_b, chunk_bytes=chunk, device_id=device_id - ) - ) - capacity_bytes = fit_within_free(want_capacity, free_bytes=free_b, chunk_bytes=chunk) - if capacity_bytes is None: - logger.warning(f"[kv-bounce] disabled: only {free_b // _MIB}MiB free") - return None - if capacity_bytes != want_capacity: - logger.warning( - f"[kv-bounce] each region clamped to {capacity_bytes // _MIB}MiB " - f"(2x total) to fit {free_b // _MIB}MiB free" - ) - return cls( - agent, - device_id=device_id, - capacity_bytes=capacity_bytes, - phys_chunk_size=chunk, - block_bytes_per_group=block_bytes_per_group, - min_bytes=cfg.min_bytes, - min_blocks=cfg.min_blocks, - ) - - def __init__( - self, - agent, - *, - device_id: int, - capacity_bytes: int, - phys_chunk_size: int, - block_bytes_per_group: list[int | None], - min_bytes: int = DEFAULT_MIN_BYTES, - min_blocks: int = 96, - quarantine_grace_s: float = _QUARANTINE_GRACE_S, - name: str = "kv_bounce", - ): - self._agent = agent - self._device_id = device_id - # The byte size of one cache block, listed for each attention layer group. - self._block_bytes_per_group = list(block_bytes_per_group) - # Size gates below which bounce is skipped: coalescing only pays off once the transfer is - # large enough to beat the gather+scatter overhead (see config.DEFAULT_MIN_BYTES for the - # rationale). min_bytes applies to payloads carrying recurrent (mamba/KDA) state; - # min_blocks applies to plain-KV payloads (see the gate in reserve()). - self._min_bytes = min_bytes - self._min_blocks = min_blocks - # how long an orphaned region is held out of reuse; must outlast the worst in-flight write - self._quarantine_grace_s = quarantine_grace_s - - # one registered region each for sending and receiving - self._send_alloc = SlotAllocator(capacity_bytes, phys_chunk_size, name=f"{name}_send") - self._recv_alloc = SlotAllocator(capacity_bytes, phys_chunk_size, name=f"{name}_recv") - self._reg_descs = [self._send_alloc.reg_descs(), self._recv_alloc.reg_descs()] - for d in self._reg_descs: - self._agent.register_memory(d) - - self._send_stream = self._new_stream() - self._send_stream_lock = threading.Lock() - - self._init_recv_state() - self._start_scatter_worker(name) - - logger.info( - f"[kv-bounce] Transport: send+recv regions of " - f"{self._send_alloc.capacity / _MIB:.1f}MiB each" - ) - - def _init_recv_state(self) -> None: - # Live per-transfer state, guarded by a leaf lock: mutate and decide under it, then release - # it before any CUDA sync, allocator call, or callback. - self._reserved_map: Dict[RidSlice, TransferContext] = {} - self._reserved_map_lock = threading.Lock() - - def _start_scatter_worker(self, name: str) -> None: - # Scatter runs on its own thread: it ends in a blocking sync, so keeping it off the - # completion handler lets that handler keep draining other transfers. - self._scatter_q: "queue.Queue" = queue.Queue() - self._scatter_stream = self._new_stream() - self._stop = threading.Event() - self._scatter_thread = threading.Thread( - target=self._scatter_loop, name=f"{name}-scatter", daemon=True - ) - self._scatter_thread.start() - - def _new_stream(self): - return CUASSERT(cudart.cudaStreamCreate())[0] - - def _gather_blocking(self, src_addr: int, write_meta, total: int) -> None: - """Gather the scattered fragments into the send region and block until done. The whole gather - runs under the stream lock so a second sender thread can't overwrite the shared staging buffer - mid-copy and corrupt this region; only the fast gather serializes, the writes stay parallel.""" - plan = Plan(write_meta.src_ptrs, write_meta.dst_ptrs, write_meta.sizes, total) - with self._send_stream_lock: - gather_contiguous( - src_addr, plan.src_ptrs, plan.sizes, plan.offsets, stream=self._send_stream - ) - event = CUASSERT(cudart.cudaEventCreate())[0] - try: - CUASSERT(cudart.cudaEventRecord(event, self._send_stream)) - CUASSERT(cudart.cudaEventSynchronize(event)) - finally: - CUASSERT(cudart.cudaEventDestroy(event)) - - def _make_write(self, src_addr: int, write_meta, total: int): - # one coalesced descriptor spanning the whole region - sizes = np.array([total], dtype=np.int64) - src = MemoryDescs.from_arrays_uniform_device( - MemoryType.VRAM, np.array([src_addr], dtype=np.int64), sizes, self._device_id - ) - dst = MemoryDescs.from_arrays_uniform_device( - MemoryType.VRAM, - np.array([write_meta.bounce_dst_base], dtype=np.int64), - sizes, - write_meta.dst_device_id, - ) - return TransferRequest(TransferOp.WRITE, src, dst, write_meta.peer_name, None) - - def _reserve_and_gather(self, write_meta, *, timeout): - """Reserve a send slot and gather into it, or None on send-region backpressure. Eligibility - was already decided by the receiver, so the sender only falls back under backpressure.""" - total = int(write_meta.sizes.sum()) - res = self._send_alloc.reserve(total, timeout=timeout) - if res is None: - logger.warning_once( - f"[kv-bounce] in-place: no send region space for {total // _MIB}MiB within {timeout}s " - f"(sender backpressure); falling back", - key="kv-bounce-send-backpressure", - ) - return None - slot_id, src_addr = res - try: - self._gather_blocking(src_addr, write_meta, total) - except Exception: - self._send_alloc.release(slot_id) # free the slot if the gather raises - raise - return slot_id, src_addr, total - - def build_request(self, write_meta): - """Gather into a send slot and build the coalesced write, or None on backpressure. The gather - blocks (and frees the slot on failure) inside _reserve_and_gather.""" - gathered = self._reserve_and_gather(write_meta, timeout=_RESERVE_TIMEOUT_S) - if gathered is None: # backpressure: fall back - return None - slot_id, src_addr, total = gathered - return self._make_write(src_addr, write_meta, total), slot_id - - def release_send(self, slot_id) -> None: - """Release a send region after its write has completed.""" - self._send_alloc.release(slot_id) - - @staticmethod - def _skip_bounce(reason: str, *, warn_key: str) -> bool: - """Log why a transfer falls back to the per-fragment path and return False, so the guards - above stay one line each. Every reason logs at warning once per key: silently skipping - bounce can be a ~1000x bandwidth cliff (host-staged tcp vs cuda_ipc), so the first skip per - distinct reason must be visible at the default log level.""" - logger.warning_once(f"[kv-bounce] in-place: {reason}", key=warn_key) - return False - - def reserve( - self, - recv_req, - num_writers: int = 1, - *, - timeout: Optional[float] = _RECV_RESERVE_TIMEOUT_S, - extra_bytes: int = 0, - ) -> bool: - """Reserve a region and create its state, recording the address for the senders. Returns - False to fall back to the per-fragment path. A fan-in splits the region evenly, so the total - must divide across the writers. ``extra_bytes`` is the non-paged payload the sender appends - to the same coalesced write (mamba/KDA recurrent state, sized by the receiver via - ``mamba_receiver_payload_bytes``); the region must cover it or the write would overrun into the - neighboring slot.""" - total = 0 - has_state_group = False - # Block tables are positional with -1 holes (see Chunk); only the - # slots the receiver actually accepts are written into the region. - valid_counts = [ - int((block_ids >= 0).sum()) for block_ids in recv_req.block_ids_per_layer_groups - ] - for g, n_valid in enumerate(valid_counts): - if n_valid == 0: - continue - if g >= len(self._block_bytes_per_group): - return self._skip_bounce( - f"layer group {g} has blocks but no known slot size", - warn_key="kv-bounce-unknown-slot-size", - ) - if not self._block_bytes_per_group[g]: - # This group holds recurrent state (mamba/KDA), not paged KV blocks; - # its size is accounted for by extra_bytes below, not per-block math. - has_state_group = True - continue - total += n_valid * self._block_bytes_per_group[g] - if has_state_group and num_writers > 1: - # Fan-in with recurrent-state groups is unsafe: each writer may - # append its own state fragments (different PP stages hold different - # mamba layers), and the representative extra_bytes (computed from - # one sender's page table) may not reflect all participating senders. - # Even when extra_bytes == 0 (e.g. the representative rank has no - # mamba overlap), another PP sender may still contribute state bytes - # that were not reserved in the bounce region. - return self._skip_bounce( - f"fan-in across {num_writers} senders with STATE groups present; the " - f"equal split cannot account for per-writer state fragments", - warn_key="kv-bounce-mamba-fanin", - ) - total += int(extra_bytes) - if total <= 0: - return self._skip_bounce( - f"computed transfer size {total} <= 0", warn_key="kv-bounce-nonpositive-size" - ) - # Which size gate applies depends on the payload. Payloads carrying recurrent (mamba/KDA) - # state gate on BYTES: the cost the gate guards (falling back to the slow per-fragment - # path) scales with bytes, and a block count is meaningless for the non-paged state (the - # Kimi-K3 regression: a 433 MiB transfer of 67 small blocks plus KDA state failed a - # 96-block gate and fell onto the ~0.4 GB/s host-staged path). Plain-KV payloads keep the - # original block-count gate so pre-existing bounce deployments (opted in via - # kv_cache_bounce_size_mb) see no change in which transfers use the arena. - # TODO(TRTLLM-15194): investigate whether the byte-only gate is - # safe (or better) for plain-KV payloads too, so this special case can be removed and - # both payload kinds share one gate. - nblocks = sum(valid_counts) - if extra_bytes > 0: - if total < self._min_bytes: - return self._skip_bounce( - f"{total}B ({nblocks} blocks + recurrent state) < min {self._min_bytes}B " - f"(too small; tune TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES)", - warn_key="kv-bounce-below-min-bytes", - ) - elif nblocks < self._min_blocks: - return self._skip_bounce( - f"{nblocks} blocks < min {self._min_blocks} (too small; tune " - f"TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS)", - warn_key="kv-bounce-below-min-blocks", - ) - if num_writers > 1 and total % num_writers != 0: - return self._skip_bounce( - f"fan-in {total}B across {num_writers} senders is not an even split " - f"({total % num_writers}B remainder); head-mismatch explosion NOT mitigated", - warn_key="kv-bounce-uneven-fanin", - ) - if num_writers > 1: - # Fan-in gives each writer an equal share of the region, which only matches where it - # writes when all writers send the same bytes. Equal layer count guarantees that only - # when the per-block sizes match, so require that here, else fall back. - # Exclude None entries (STATE groups) — they are already guarded above. - present_slot_bytes = { - self._block_bytes_per_group[g] - for g, n_valid in enumerate(valid_counts) - if n_valid > 0 and self._block_bytes_per_group[g] - } - if len(present_slot_bytes) > 1: - return self._skip_bounce( - f"fan-in across {num_writers} senders with non-uniform layer-group slot bytes " - f"{sorted(present_slot_bytes)}; the equal split would overrun a sub-region", - warn_key="kv-bounce-heterogeneous-fanin", - ) - if total > self._recv_alloc.capacity: # too big to ever fit, unlike transient backpressure - return self._skip_bounce( - f"transfer {total // _MIB}MiB exceeds the {self._recv_alloc.capacity // _MIB}MiB bounce " - f"region; raise the bounce arena size to re-enable coalescing", - warn_key="kv-bounce-oversize", - ) - res = self._recv_alloc.reserve(total, timeout=timeout) - if res is None: - return self._skip_bounce( - f"no recv region space for {total // _MIB}MiB within {timeout}s (backpressure)", - warn_key="kv-bounce-recv-backpressure", - ) - slot_id, addr = res - recv_req.bounce_dst_base = addr - with self._reserved_map_lock: - ctx = TransferContext( - rid_slice=(recv_req.unique_rid, recv_req.slice_id), - slot_id=slot_id, - base_addr=addr, - per_writer_bytes=total // num_writers, - num_writers=num_writers, - ) - self._reserved_map[ctx.rid_slice] = ctx # inactive until the first writer reports - # Positive marker: all fall-back guards above passed, so this transfer provably takes the - # coalesced-bounce WRITE path. Logged once (per process) so an e2e test can assert that - # bounce actually engaged instead of silently falling back to the per-fragment path. - logger.info_once( - f"[kv-bounce] coalesced {nblocks} blocks / {total // _MIB}MiB into one region " - f"across {num_writers} writer(s)", - key="kv-bounce-coalesced", - ) - return True - - def writer_base(self, rid_slice: RidSlice, writer_index: int) -> Optional[int]: - """Where the given fan-in writer writes in the region.""" - with self._reserved_map_lock: - ctx = self._reserved_map.get(rid_slice) - return None if ctx is None else ctx.writer_base(writer_index) - - def is_bounced(self, rid_slice: RidSlice) -> bool: - with self._reserved_map_lock: - return rid_slice in self._reserved_map - - def release_idle_reservation(self, rid_slice: RidSlice) -> None: - """Immediately release a reservation cancelled before any address went out; no write can be - in flight. Idempotent. Drained transfers finalize through the result path instead.""" - with self._reserved_map_lock: - ctx = self._reserved_map.pop(rid_slice, None) - if ctx is not None: - self._recv_alloc.release(ctx.slot_id) - - def orphan_reservation(self, rid_slice: RidSlice) -> None: - """Give up on a reservation whose write may still be in flight (cancel/timeout/lost result). - The write can't be aborted, so quarantine the region (reclaimed later) rather than releasing - or leaking it. Idempotent; a no-op once the transfer has settled.""" - self._apply(rid_slice, lambda ctx: ctx.mark_orphaned()) - - def abort_publication(self, rid_slice: RidSlice, published_writers: set[int]) -> None: - """Retain a failed fan-out until every successfully published writer drains.""" - self._apply( - rid_slice, - lambda ctx: ctx.abort_publication(published_writers), - ) - - def _apply(self, rid_slice: RidSlice, mutate: Callable[[TransferContext], None]) -> None: - """Mutate the state under the lock, then do what it asks (scatter or settle) with the lock - released, never holding it across a CUDA sync, a queue put, or a callback. No-op if the - region is already gone.""" - scatter: Optional[tuple] = None - settlement: Optional[Settlement] = None - with self._reserved_map_lock: - ctx = self._reserved_map.get(rid_slice) - if ctx is None: - return - mutate(ctx) - if ctx.ready_to_scatter(): - ctx.begin_scatter() - scatter = (ctx, ctx.sorted_scatter_descs()) - elif ctx.ready_to_settle(): - settlement = ctx.settle() - if settlement is not None: - self._reserved_map.pop(rid_slice, None) - if scatter is not None: - self._enqueue_scatter(*scatter) - if settlement is not None: - self._commit(settlement) - - def _enqueue_scatter(self, ctx: TransferContext, descs: List[tuple]) -> None: - """Hand the per-writer fragments to the worker. Each is scattered from its own source, so a - writer that fell back to the in-place path cannot shift where the others are read from.""" - self._scatter_q.put((ctx, descs)) - - def _commit(self, settlement: Settlement) -> None: - """Carry out the decision: release or quarantine the slot, then fire the callback once. No - lock is held.""" - if settlement.disposition is Disposition.QUARANTINE: - self._recv_alloc.quarantine(settlement.slot_id, self._quarantine_grace_s) - else: - self._recv_alloc.release(settlement.slot_id) - if settlement.on_done is not None: - try: - settlement.on_done(settlement.success) - except Exception as e: # never let the callback strand the arena - logger.error( - f"[kv-bounce] completion callback failed (slot={settlement.slot_id}): {e}" - ) - - def record_result( - self, - rid_slice: RidSlice, - peer_rank: int, - dst_ptrs=None, - sizes=None, - src_base=None, - on_done: Optional[Callable[[bool], None]] = None, - ) -> None: - """A writer reported success. The completion callback fires only after the scatter lands, so - the reader never sees completion before the cache is in place.""" - - def mut(ctx: TransferContext) -> None: - if on_done is not None: - ctx.on_done = on_done - ctx.record_writer_result( - peer_rank, succeeded=True, src_base=src_base, dst_ptrs=dst_ptrs, sizes=sizes - ) - - self._apply(rid_slice, mut) - - def record_failure(self, rid_slice: RidSlice, peer_rank: int) -> None: - """A writer reported failure (it has drained). The region is freed only once every writer has - reported, not here.""" - self._apply(rid_slice, lambda ctx: ctx.record_writer_result(peer_rank, succeeded=False)) - - def _scatter_loop(self): - CUASSERT(cudart.cudaSetDevice(self._device_id)) - while not self._stop.is_set(): - try: - item = self._scatter_q.get(timeout=_SCATTER_POLL_S) - except queue.Empty: - # idle: reclaim quarantine past its grace period, independent of any reserve call - self._recv_alloc.reclaim_expired() - continue - if item is None: - break # poison pill from close: wake and exit - ctx, descs = item - ok = True - try: - # Scatter each writer's fragments from its own source, never one global offset, so a - # missing or fallback writer cannot shift where the others are read from. - for src_base, dst_ptrs, sizes in descs: - p = Plan(dst_ptrs, dst_ptrs, sizes, int(sizes.sum())) - scatter_contiguous( - src_base, p.dst_ptrs, p.sizes, p.offsets, stream=self._scatter_stream - ) - CUASSERT(cudart.cudaStreamSynchronize(self._scatter_stream)) - except Exception as e: - # a scatter failure must not kill the worker nor be reported as success - ok = False - logger.error(f"[kv-bounce] scatter failed (slot={ctx.slot_id}): {e}") - # record the outcome and settle; completion fires only after the sync above - self._apply(ctx.rid_slice, lambda c, ok=ok: c.finish_scatter(ok)) - - def close(self) -> None: - self._stop.set() - # poison pill: wake the worker now instead of waiting out its poll - self._scatter_q.put(None) - if self._scatter_thread.is_alive(): - self._scatter_thread.join(timeout=_CLOSE_JOIN_S) - for d in self._reg_descs: - try: - self._agent.deregister_memory(d) - except Exception: - pass - self._send_alloc.close() - self._recv_alloc.close() - - -class NoBounceTransport(BounceTransport): - """The disabled transport, used when bounce is off so callers need no None checks. Every method - is a no-op or a negative answer.""" - - enabled = False - _reg_descs = () - - def build_request(self, write_meta): - return None - - def release_send(self, slot_id) -> None: - pass - - def reserve( - self, - recv_req, - num_writers: int = 1, - *, - timeout: Optional[float] = _RECV_RESERVE_TIMEOUT_S, - extra_bytes: int = 0, - ) -> bool: - return False - - def writer_base(self, rid_slice, writer_index: int): - return None - - def is_bounced(self, rid_slice) -> bool: - return False - - def release_idle_reservation(self, rid_slice) -> None: - pass - - def orphan_reservation(self, rid_slice) -> None: - pass - - def abort_publication(self, rid_slice, published_writers: set[int]) -> None: - pass - - def record_result( - self, rid_slice, peer_rank, dst_ptrs=None, sizes=None, src_base=None, on_done=None - ): - pass - - def record_failure(self, rid_slice, peer_rank) -> None: - pass - - def close(self) -> None: - pass - - -def create_bounce(agent, cfg, *, device_id: int, page_table) -> BounceTransport: - """Build the real transport from the config, or the disabled one when bounce is off, it cannot - fit, or the fabric allocation races.""" - if cfg is None: - return NoBounceTransport() - try: - transport = VmmBounceTransport.from_config( - agent, cfg, device_id=device_id, block_bytes_per_group=block_bytes_per_group(page_table) - ) - return transport if transport is not None else NoBounceTransport() - except ( - Exception - ) as e: # rare race: memory taken between the free-memory query and the allocation - logger.warning(f"[kv-bounce] disabled (alloc failed: {e}); using in-place path") - return NoBounceTransport() - - -def build_send_request(bounce, write_meta, fallback): - """Build a coalesced bounce write when eligible (release the returned slot afterward), otherwise - fall back to the per-fragment request.""" - if write_meta.bounce_dst_base is not None: - built = bounce.build_request(write_meta) - if built is not None: - return built - return fallback(), None - - -def scatter_write_result( - bounce, rid_slice, peer_rank: int, dst_ptrs, sizes, src_base=None, on_done=None -) -> None: - """Handle a success result: a bounced transfer records the writer and scatters once all arrive; a - non-bounced transfer already landed in place, so fire the callback inline.""" - if bounce.is_bounced(rid_slice): - bounce.record_result(rid_slice, peer_rank, dst_ptrs, sizes, src_base, on_done) - elif on_done is not None: - on_done(True) - - -def encode_result_tail(write_meta) -> list: - """The binary tail appended to a bounced result: the destination fragment table and the source - this writer wrote to, so the receiver can scatter and order the fan-in writers.""" - sb = write_meta.bounce_dst_base if write_meta.bounce_dst_base is not None else 0 - return [ - write_meta.dst_ptrs.tobytes(), - write_meta.sizes.tobytes(), - np.array([sb], dtype=np.int64).tobytes(), - ] - - -def decode_result_tail(message): - """Recover the destination fragments, sizes, and source from the optional trailing frames, or - nothing if the tail is absent.""" - if len(message) >= 5: - return ( - np.frombuffer(message[2], dtype=np.int64), - np.frombuffer(message[3], dtype=np.int64), - int(np.frombuffer(message[4], dtype=np.int64)[0]), - ) - return None, None, None - - -def block_bytes_per_group(page_table: KVCachePageTable) -> list[int | None]: - """Return transferred bytes per cache block for each layer group. - - All distinct physical pools exposed by an attention group contribute to its - transfer size. Multiple logical views of the same physical pool contribute - only once. Non-attention groups retain a ``None`` placeholder so the result - remains aligned with receive-request layer-group indices. - """ - from tensorrt_llm._torch.disaggregation.resource.page import CacheKind - from tensorrt_llm._torch.disaggregation.resource.utils import get_physical_pool - - assert page_table is not None - out: list[int | None] = [] - for lg_idx, lg in enumerate(page_table.layer_groups): - if lg.kind != CacheKind.PAGED: - out.append(None) - continue - pool_indices = {pool_view.pool_idx for pool_view in lg.pool_views} - out.append( - sum( - int(get_physical_pool(page_table, lg_idx, pool_idx).slot_bytes) - for pool_idx in pool_indices - ) - ) - return out diff --git a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py index 5d7c993df49c..a5e4115280ea 100644 --- a/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py +++ b/tensorrt_llm/_torch/disaggregation/native/mixers/ssm/peer.py @@ -30,10 +30,7 @@ MapperKind, PoolView, ) -from tensorrt_llm._torch.disaggregation.resource.utils import ( - find_replicated_role_mismatch, - get_pool_view_global_layer_ids, -) +from tensorrt_llm._torch.disaggregation.resource.utils import find_replicated_role_mismatch from tensorrt_llm._utils import nvtx_range @@ -637,41 +634,3 @@ def is_paired(self_ri: RankInfo, peer_ri: RankInfo) -> bool: else: ratio = self_tp // peer_tp return (self_rank // ratio) == peer_rank - - -def mamba_receiver_payload_bytes( - sender_page_table: KVCachePageTable, - receiver_page_table: KVCachePageTable, - dst_slot: Optional[int], -) -> int: - """Recurrent-state bytes that will land in the receiver's slot. - - Receiver-local invariant: regardless of the sender-side shard pairing, - the slot receives exactly the receiver's own per-layer slot bytes over - the overlapping mamba layers. This avoids the RankInfoServer rank-0 - limitation that makes sender-side simulation return 0 for non-rank-0 - receivers. - """ - if dst_slot is None: - return 0 - sender_mlg = MambaPolicy._find_mamba_layer_group(sender_page_table) - receiver_mlg = MambaPolicy._find_mamba_layer_group(receiver_page_table) - if sender_mlg is None or receiver_mlg is None: - return 0 - - sender_views = { - (pool_view.pool_role, pool_view.mapper_kind): pool_view - for pool_view in sender_mlg.pool_views - } - total = 0 - for receiver_view in receiver_mlg.pool_views: - sender_view = sender_views.get((receiver_view.pool_role, receiver_view.mapper_kind)) - if sender_view is None: - continue - sender_layers = set(get_pool_view_global_layer_ids(sender_view, sender_mlg)) - receiver_layers = set(get_pool_view_global_layer_ids(receiver_view, receiver_mlg)) - overlap = sender_layers & receiver_layers - if not overlap: - continue - total += len(overlap) * receiver_view.bytes_per_layer - return total diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index bc08639190fe..4b2c4f431bfb 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -25,7 +25,7 @@ from contextlib import contextmanager, nullcontext from dataclasses import dataclass from enum import Enum -from typing import TYPE_CHECKING, Callable, Dict, Iterable, Iterator, List, Optional, Union +from typing import Callable, Dict, Iterable, Iterator, List, Optional, Union import msgpack import numpy as np @@ -63,11 +63,8 @@ ) from tensorrt_llm._torch.disaggregation.native.messenger import ZMQMessenger, decode_message from tensorrt_llm._torch.disaggregation.native.mixers.attention.peer import AttentionPolicy -from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( - MambaPolicy, - mamba_receiver_payload_bytes, -) -from tensorrt_llm._torch.disaggregation.native.peer import PeerOverlap, PeerRegistrar +from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy +from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.perf_logger import PerfTimer, perf_log_manager from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.native.retirement import ( @@ -78,16 +75,13 @@ from tensorrt_llm._torch.disaggregation.native.utils import get_local_ip from tensorrt_llm._torch.disaggregation.nixl.agent import NixlTransferAgent from tensorrt_llm._torch.disaggregation.resource.kv_extractor import KVRegionExtractorV1 -from tensorrt_llm._torch.disaggregation.resource.page import CacheKind, KVCachePageTable, MapperKind +from tensorrt_llm._torch.disaggregation.resource.page import CacheKind from tensorrt_llm._torch.disaggregation.resource.utils import get_unique_pool_memory_descs from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager from tensorrt_llm._utils import CUASSERT, nvtx_range from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle -if TYPE_CHECKING: - from .bounce import Config - AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType LlmRequestType = tensorrt_llm.bindings.internal.batch_manager.LlmRequestType @@ -121,7 +115,6 @@ class RecvReqInfo: dst_start_token: Optional[int] = None aux_slot: Optional[int] = None slice_id: Optional[int] = None - bounce_dst_base: Optional[int] = None def to_bytes(self) -> bytes: return msgpack.packb( @@ -136,7 +129,6 @@ def to_bytes(self) -> bytes: "dst_start_token": self.dst_start_token, "aux_slot": self.aux_slot, "slice_id": self.slice_id, - "bounce_dst_base": self.bounce_dst_base, } ) @@ -178,7 +170,6 @@ class WriteMeta: receiver_slice_id: int = 0 is_last_slice: bool = False meta_type: WriteMetaType = WriteMetaType.KV - bounce_dst_base: Optional[int] = None class MessageType: @@ -346,11 +337,10 @@ def __init__(self, retirement: Optional[RetirementDeadline] = None) -> None: self._in_doubt_writers: set[int] = set() self._settled_writers: set[int] = set() self._publication_failed = False - self._local_completion_pending = False self._invalid_evidence = False def _settle_if_drained_locked(self) -> None: - """Remove this claim only after the complete cohort and local work settle.""" + """Remove this claim only after the complete cohort and publication settle.""" if self._retirement is not None and self._resources_drained_locked(): self._retirement.settle(self) @@ -492,13 +482,7 @@ def record_writer_settlement(self, peer_rank: int) -> bool: self._settle_if_drained_locked() return True - def record_writer_result( - self, - peer_rank: int, - succeeded: bool, - *, - wait_for_local_completion: bool, - ) -> tuple[bool, bool]: + def record_writer_result(self, peer_rank: int, succeeded: bool) -> tuple[bool, bool]: """Record one writer and return ``(accepted, all_succeeded)``.""" with self._lock: if self._expected_writers is None: @@ -529,16 +513,9 @@ def record_writer_result( all_succeeded = ( all_reported and not self._publication_failed and all(self._writer_results.values()) ) - if all_succeeded and wait_for_local_completion: - self._local_completion_pending = True self._settle_if_drained_locked() return True, all_succeeded - def finish_local_completion(self) -> None: - with self._lock: - self._local_completion_pending = False - self._settle_if_drained_locked() - @property def all_writers_reported(self) -> bool: """Whether every writer of the sealed cohort has reported a terminal result.""" @@ -568,7 +545,6 @@ def _resources_drained_locked(self) -> bool: ) ) and not self._publication_pending - and not self._local_completion_pending and not self._in_doubt_writers and not self._invalid_evidence ) @@ -581,10 +557,9 @@ class AgentResult(Enum): FAILED_QUIESCED = "FAILED_QUIESCED" -# KV_AGENT_RESULT prefix in one struct frame (was ascii frames serialized/parsed under the -# GIL per slice per writer): instance_rank, unique_rid, slice_id, is_last, status, -# transfer_size. The optional bounce tail follows at message[2:]. -_KV_RESULT_PREFIX = struct.Struct(" list[bytes]: @@ -672,7 +644,6 @@ class _PhysicalOperation: class _PendingSettlement: write_meta: WriteMeta initial_report: Optional[list[bytes]] - send_slot_id: Optional[int] = None backend_done: bool = False report_error_logged: bool = False @@ -922,14 +893,12 @@ def __init__( self, peer_registrar: PeerRegistrar, agent: BaseTransferAgent, - bounce=None, enforce_physical_ownership: bool = False, retirement_watchdog: Optional[RetirementWatchdog] = None, ) -> None: self._registrar = peer_registrar self._device_id = peer_registrar.self_rank_info.device_id self._agent = agent - self._bounce = bounce self._enforce_physical_ownership = enforce_physical_ownership self._retirement_watchdog = retirement_watchdog self._peer_requests: dict = {} @@ -1272,17 +1241,12 @@ def _process_task_queue(self, thread_idx: int): ) dealers.clear() - def _retain_in_doubt_transfer( - self, - write_meta: WriteMeta, - initial_report: list[bytes], - send_slot_id: Optional[int] = None, - ) -> None: + def _retain_in_doubt_transfer(self, write_meta: WriteMeta, initial_report: list[bytes]) -> None: thread_idx = hash((write_meta.unique_rid, write_meta.peer_rank)) % self._num_threads key = (write_meta.task, write_meta.peer_rank) if key in self._pending_settlements[thread_idx]: return - pending = _PendingSettlement(write_meta, initial_report, send_slot_id) + pending = _PendingSettlement(write_meta, initial_report) self._pending_settlements[thread_idx][key] = pending try: self._get_result_dealer(write_meta.peer_endpoint).send(initial_report) @@ -1305,9 +1269,6 @@ def _poll_in_doubt_transfers(self, thread_idx: int) -> None: if not meta.task.poll_in_doubt_physical_operation(meta.peer_rank): continue pending.backend_done = True - if pending.send_slot_id is not None: - self._bounce.release_send(pending.send_slot_id) - pending.send_slot_id = None if meta.meta_type == WriteMetaType.AUX: message = _make_aux_result_msg( self._instance_rank, meta.unique_rid, AgentResult.FAILED_QUIESCED @@ -1432,28 +1393,21 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): ) return - from .bounce import build_send_request, encode_result_tail - agent_result = AgentResult.SUCCESS - send_slot_id = None submitted_and_completed = False if write_meta.src_ptrs.size > 0: try: - request, send_slot_id = build_send_request( - self._bounce, - write_meta, - lambda: Sender._make_agent_request(write_meta, device_id=self._device_id), - ) + request = Sender._make_agent_request(write_meta, device_id=self._device_id) except Exception as e: - # Don't let a gather fault escape: without a result the receiver would hang and its - # region leak. Tell the receiver it failed and fail the local task instead. + # Report the failure instead of letting it escape: without a result the receiver + # waits until its transfer timeout. Nothing was submitted yet. logger.error( - f"_deliver_kv_to_agent: failed to build the KV send request for " + f"_deliver_kv_to_agent: failed to build the KV transfer request for " f"{write_meta.unique_rid} slice={write_meta.slice_id}: {e}" ) if owned: task.retire_unsubmitted_physical_operation(write_meta.peer_rank) - task.fail(RuntimeError(f"build_send_request failed: {e}")) + task.fail(RuntimeError(f"failed to build the KV transfer request: {e}")) self._get_result_dealer(write_meta.peer_endpoint, legacy_shared=True).send( _make_kv_result_msg( self._instance_rank, @@ -1466,7 +1420,6 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): return if timer: timer.record_transfer_start(write_meta.peer_rank) - transfer_finished = False try: transfer_finished, last_status = self._submit_transfer( task, write_meta.peer_rank, request @@ -1474,7 +1427,7 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): if transfer_finished: submitted_and_completed = True del request - if not transfer_finished: + else: agent_result = AgentResult.IN_DOUBT if owned else AgentResult.FAILED agent_name = getattr(self._agent, "name", "") detail = ( @@ -1494,23 +1447,13 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): task.fail(error) except _TransferNotSubmittedError as error: # Backend admission failed, so local resources are safe to retire. - transfer_finished = True agent_result = AgentResult.FAILED task.fail(error) - finally: - if send_slot_id is not None and (not owned or transfer_finished): - self._bounce.release_send(send_slot_id) elif owned: task.retire_unsubmitted_physical_operation(write_meta.peer_rank) if timer: timer.record_transfer_end(write_meta.peer_rank) - # Report every chunk so failures reach the receiver immediately. - tail = ( - encode_result_tail(write_meta) - if send_slot_id is not None and agent_result == AgentResult.SUCCESS - else None - ) # Attested written bytes: only what was actually submitted to the # agent AND completed. The receiver checks the per-slice sum against # the byte total of its published destinations before admitting the @@ -1524,6 +1467,7 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): transfer_size = 0 if submitted_and_completed and agent_result == AgentResult.SUCCESS: transfer_size = int(write_meta.sizes.sum()) + # Report every chunk so failures reach the receiver immediately. result_msg = _make_kv_result_msg( self._instance_rank, write_meta.unique_rid, @@ -1531,10 +1475,9 @@ def _deliver_kv_to_agent(self, write_meta: WriteMeta): write_meta.is_last_slice, agent_result, transfer_size=transfer_size, - tail=tail, ) if agent_result == AgentResult.IN_DOUBT: - self._retain_in_doubt_transfer(write_meta, result_msg, send_slot_id) + self._retain_in_doubt_transfer(write_meta, result_msg) return self._get_result_dealer(write_meta.peer_endpoint).send(result_msg) @@ -1834,7 +1777,6 @@ def _build_kv_write_meta(self, task: KVSendTask, req_info: RecvReqInfo) -> Write slice_id=task.slice_id, receiver_slice_id=req_info.slice_id if req_info.slice_id is not None else 0, is_last_slice=task._chunk.is_last, - bounce_dst_base=req_info.bounce_dst_base, ) def _build_aux_write_meta(self, task: AuxSendTask, req_info: RecvReqInfo) -> WriteMeta: @@ -2937,21 +2879,8 @@ def record_writer_in_doubt(self, peer_rank: int) -> bool: def record_writer_settlement(self, peer_rank: int) -> bool: return self._get_physical_owner().record_writer_settlement(peer_rank) - def record_writer_result( - self, - peer_rank: int, - succeeded: bool, - *, - wait_for_local_completion: bool, - ) -> tuple[bool, bool]: - return self._get_physical_owner().record_writer_result( - peer_rank, - succeeded, - wait_for_local_completion=wait_for_local_completion, - ) - - def finish_local_completion(self) -> None: - self._get_physical_owner().finish_local_completion() + def record_writer_result(self, peer_rank: int, succeeded: bool) -> tuple[bool, bool]: + return self._get_physical_owner().record_writer_result(peer_rank, succeeded) @property def resources_drained(self) -> bool: @@ -2975,13 +2904,11 @@ def __init__( self, peer_registrar: PeerRegistrar, agent: BaseTransferAgent, - bounce=None, enforce_physical_ownership: bool = False, retirement_watchdog: Optional[RetirementWatchdog] = None, ): self._registrar = peer_registrar self._agent = agent - self._bounce = bounce # Internal component gate only. Production wiring is added after the # complete Phase 1 ownership path is available and qualified. self._enforce_physical_ownership = enforce_physical_ownership @@ -3132,56 +3059,6 @@ def _build_recv_req_info(self, task: KVRecvTask) -> RecvReqInfo: slice_id=task.slice_id, ) - @staticmethod - def _fanin_bounce_safe( - overlap: PeerOverlap, - peer_ri: RankInfo, - receiver_page_table: Optional[KVCachePageTable], - ) -> bool: - """Whether multi-writer bounce's equal total//num_writers split is valid for this overlap. - The split assumes every writer contributes the same size, which holds when: - * duplicate_head_factor == 1 -- else some ranks don't send KV yet still count in - expected_transfers, so the live writers overflow their slots; - * the PP layer split is even -- a single PP stage (overlap_pp_size <= 1) is trivially fine; - for PP fan-in, every overlapping stage must hold the same number of layers - (peer_ri.layer_num_per_pp all-equal) or per-writer sizes differ. If that full per-stage - list isn't available (shorter than the fan-in degree), be conservative and fall back. - Otherwise fall back to the per-fragment path (correct, just not coalesced). - Equal layer count means equal bytes only when the per-block sizes match; reserve() rejects - the mismatched case, so this only needs the count to split evenly.""" - if overlap.duplicate_head_factor != 1: - return False - if overlap.overlap_pp_size > 1: - lpp = getattr(peer_ri, "layer_num_per_pp", None) - if not lpp or len(lpp) < overlap.overlap_pp_size or len(set(lpp)) != 1: - return False - # Replicated pools (e.g. MiniMax M3 index-key) are sent by one elected - # fan-in owner only, so with multiple writers their contributions - # differ in size and the equal split is invalid. Inspect both endpoints: - # a masked PP stage may advertise no replicated view even though another - # stage owns one that is visible in the receiver's page table. - if len(overlap.ranks) > 1: - for page_table in (peer_ri.page_table, receiver_page_table): - if page_table is None: - continue - for layer_group in page_table.layer_groups: - for pool_view in getattr(layer_group, "pool_views", ()): - if pool_view.mapper_kind == MapperKind.REPLICATED: - return False - return True - - def _get_mamba_slot(self, req_info: RecvReqInfo) -> Optional[int]: - """Extract the mamba slot index from block_ids_per_layer_groups, or None.""" - pt = self._registrar.self_extractor.page_table - if pt is None: - return None - for lg_idx, lg in enumerate(pt.layer_groups): - if lg.kind == CacheKind.STATE: - block_ids = req_info.block_ids_per_layer_groups[lg_idx] - if block_ids.size > 0: - return int(block_ids[0]) - return None - def dispatch_task(self, task: KVRecvTask) -> None: params = task._params logger.debug( @@ -3200,9 +3077,9 @@ def dispatch_task(self, task: KVRecvTask) -> None: # globally once any rank fails here — ranks whose page table has # no recurrent layers (e.g. a hybrid-model PP stage) pass # validation and may still run a transfer that is then - # discarded. The task is already in the session's _kv_tasks, no - # bounce reservation exists yet, and session._sender_endpoints - # is still empty, so the unpublished owner can close locally. + # discarded. The task is already in the session's _kv_tasks and + # session._sender_endpoints is still empty, so the unpublished + # owner can close locally. logger.error( "dispatch_task: context peer incompatible, failing request " f"unique_rid={task._unique_rid}: {e}" @@ -3244,39 +3121,6 @@ def dispatch_task(self, task: KVRecvTask) -> None: task.expected_transfers = len(peer_overlap.ranks) else: task.expected_transfers = len(dp0_overlap.ranks) - # TP fan-in splits ONE region equally, so allow it only for a uniform writer set: - # _fanin_bounce_safe() (TP-by-head / even-PP), and never under ADP broadcast (sender_dp_rank - # None), where the real writer count exceeds expected_transfers and would overflow the slot. - topo_overlap = peer_overlap if sender_dp_rank is not None else dp0_overlap - # Helix writers contribute unequal shares (one KV owner, one state - # sender, the rest empty), so the fan-in equal split would overflow - # into neighboring reservations; use the per-fragment path. - cp_involved = self._registrar.self_rank_info.cp_size > 1 or peer_infos.cp_size > 1 - allow_bounce = task.expected_transfers == 1 or ( - sender_dp_rank is not None - and not cp_involved - and self._fanin_bounce_safe( - topo_overlap, - peer_infos, - self._registrar.self_extractor.page_table, - ) - ) - # Recurrent (mamba/KDA) state rides the SAME coalesced write as the KV blocks, - # so the bounce region must be sized for those bytes too. - extra_bytes = 0 - mamba_dst_slot = self._get_mamba_slot(receiver_req) - if mamba_dst_slot is not None: - if peer_infos.page_table is None: - allow_bounce = False # cannot size the recurrent-state payload - else: - extra_bytes = mamba_receiver_payload_bytes( - sender_page_table=peer_infos.page_table, - receiver_page_table=self._registrar.self_extractor.page_table, - dst_slot=mamba_dst_slot, - ) - bounced = allow_bounce and self._bounce.reserve( - receiver_req, task.expected_transfers, extra_bytes=extra_bytes - ) session = self._get_session(task._unique_rid) if session is None: raise RuntimeError( @@ -3284,43 +3128,27 @@ def dispatch_task(self, task: KVRecvTask) -> None: "session may have been closed before dispatch" ) sender_endpoints = {peer_infos.sender_endpoints[rank] for rank in peer_overlap.ranks} - # Fan-in: each sender gets its own sub-region base (writers must not overwrite); else serialize once. - fanin_bounce = bounced and task.expected_transfers > 1 - key = (receiver_req.unique_rid, receiver_req.slice_id) - receiver_req_bytes = None if fanin_bounce else receiver_req.to_bytes() + receiver_req_bytes = receiver_req.to_bytes() if not self._enforce_physical_ownership: session.mark_transferring(task.slice_id) # Cache sender endpoints so cancel() can send CANCEL_SESSION to them. session._sender_endpoints.update(sender_endpoints) - for i, rank in enumerate(peer_overlap.ranks): + for rank in peer_overlap.ranks: if task._perf_timer is not None: task._perf_timer.record_task_start(rank) - if fanin_bounce: - receiver_req.bounce_dst_base = self._bounce.writer_base(key, i) - receiver_req_bytes = receiver_req.to_bytes() self._request_sender_data(peer_infos.sender_endpoints[rank], receiver_req_bytes) return peer_ranks = list(peer_overlap.ranks) - serialized_requests: list[tuple[int, bytes]] = [] - for i, rank in enumerate(peer_ranks): - if fanin_bounce: - receiver_req.bounce_dst_base = self._bounce.writer_base(key, i) - payload = receiver_req.to_bytes() - else: - assert receiver_req_bytes is not None - payload = receiver_req_bytes - serialized_requests.append((rank, payload)) - published_writers: set[int] = set() def publish_requests() -> None: - for rank, payload in serialized_requests: + for rank in peer_ranks: if task._perf_timer is not None: task._perf_timer.record_task_start(rank) published_writers.add(rank) try: - self._request_sender_data(peer_infos.sender_endpoints[rank], payload) + self._request_sender_data(peer_infos.sender_endpoints[rank], receiver_req_bytes) except Exception: published_writers.discard(rank) raise @@ -3340,9 +3168,7 @@ def publish_requests() -> None: publish=publish_requests, published_writers=published_writers, ): - self._bounce.release_idle_reservation(key) task.cancel_unpublished() - return @staticmethod def _extract_info_endpoint(params: DisaggregatedParams) -> str: @@ -3504,11 +3330,8 @@ def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): ) return peer_rank, unique_rid, receiver_slice_id, is_last_slice, status_code, transfer_size = ( - _KV_RESULT_PREFIX.unpack(message[1]) + _KV_RESULT_FRAME.unpack(message[1]) ) - from .bounce import decode_result_tail - - dst_ptrs, sizes, src_base = decode_result_tail(message) session = self._get_session(unique_rid) if session is None: logger.warning( @@ -3520,9 +3343,6 @@ def _process_kv_agent_result(self, _send_id: bytes, message: list[bytes]): receiver_slice_id, is_last_slice, _AGENT_RESULT_BY_CODE[status_code], - dst_ptrs=dst_ptrs, - sizes=sizes, - src_base=src_base, transfer_size=transfer_size, ) @@ -3766,10 +3586,6 @@ def try_begin_transfer( # send transfers responsibility to that writer; a failed # send does not. Retain only the successfully queued prefix # and wait for each of those writers to settle. - self._receiver._bounce.abort_publication( - (self.disagg_request_id, task.slice_id), - published_writers, - ) task.abort_publication(published_writers) if self._aux_physical_owner is not None: self._aux_physical_owner.abort_publication(published_writers) @@ -3854,10 +3670,7 @@ def dispatch_prepared_receive(self, task: KVRecvTask) -> None: try: self._receiver.dispatch_task(task) except Exception as error: - if self.cancel_unpublished_task(task): - self._receiver._bounce.release_idle_reservation( - (self.disagg_request_id, task.slice_id) - ) + self.cancel_unpublished_task(task) self.fail_admission(error) raise @@ -3888,9 +3701,6 @@ def process_kv_agent_result( receiver_slice_id: int, is_last_slice: bool, status: AgentResult, - dst_ptrs=None, - sizes=None, - src_base=None, transfer_size: int = 0, ): with self._ownership_evidence_guard(), self.lock: @@ -3936,22 +3746,16 @@ def process_kv_agent_result( if not self._enforce_physical_ownership: raise RuntimeError("received physical settlement without ownership enabled") try: - accepted = task.record_writer_settlement(peer_rank) + task.record_writer_settlement(peer_rank) except Exception as error: self._record_ownership_evidence_error(error) raise - if accepted: - self._receiver._bounce.record_failure( - (self.disagg_request_id, task.slice_id), peer_rank - ) return if self._enforce_physical_ownership: if status == AgentResult.FAILED or is_last_slice: try: accepted, all_succeeded = task.record_writer_result( - peer_rank, - status == AgentResult.SUCCESS, - wait_for_local_completion=(status == AgentResult.SUCCESS), + peer_rank, status == AgentResult.SUCCESS ) except Exception as error: task.fail(error) @@ -3976,71 +3780,28 @@ def process_kv_agent_result( # attestation of destination bytes actually submitted and # completed for this chunk (0 unless the write completed). task.verified_write_bytes += transfer_size - - from .bounce import scatter_write_result - - on_done = None - if is_last_slice and all_succeeded: - # Completing message: defer task.complete()+perf until the scatter has actually - # landed. scatter_write_result fires this inline for the non-bounced path, or on - # the scatter worker (after cudaStreamSynchronize) for the bounced path, so the - # gen consumer never observes completion before the KV is scattered into place. - request_id = self.request_id + # The completing report: every expected writer succeeded. A concurrent failure that + # already failed the task stands. complete() sets status before _event, keeping + # wait_complete's status-first poll correct. + if is_last_slice and all_succeeded and task.status != TaskStatus.ERROR: ri = self._receiver._registrar.self_rank_info - instance_name, instance_rank = ri.instance_name, ri.instance_rank - - def on_done( - success, - task=task, - peer_rank=peer_rank, - receiver_slice_id=receiver_slice_id, - request_id=request_id, - instance_name=instance_name, - instance_rank=instance_rank, - ): - # Runs on the scatter worker thread for the bounced path. Do not acquire - # RxSession.lock here: the non-bounced path invokes this callback inline - # while already holding it. Task outcome/ownership locks are independent. - # complete() sets status before _event for wait_complete's status-first poll. - if self._enforce_physical_ownership: - task.finish_local_completion() - if not success: - task.fail( - RuntimeError( - f"KV bounce scatter failed for request {request_id} " - f"slice={receiver_slice_id}" - ) - ) - return - if task.status == TaskStatus.ERROR: - return # a concurrent FAILED writer already failed it; don't un-fail - task.complete() - # Record completion before best-effort diagnostics can block the worker. - if all(t.status == TaskStatus.TRANSFERRED for t in self._kv_tasks): - self.transfer_end_time = tensorrt_llm.bindings.global_steady_clock_now() - try: - if task._perf_timer is not None: - task._perf_timer.record_task_end(peer_rank) - task.print_perf_info(peer_rank, instance_name, instance_rank) - except Exception as e: # perf is best-effort; never block completion - logger.warning( - f"KV transfer perf logging failed for request {request_id} " - f"slice={receiver_slice_id}: {e}" - ) - logger.debug( - f"KV transfer complete for request {request_id} " - f"slice={receiver_slice_id}" + task.complete() + # Record completion before best-effort diagnostics. + if all(t.status == TaskStatus.TRANSFERRED for t in self._kv_tasks): + self.transfer_end_time = tensorrt_llm.bindings.global_steady_clock_now() + try: + if task._perf_timer is not None: + task._perf_timer.record_task_end(peer_rank) + task.print_perf_info(peer_rank, ri.instance_name, ri.instance_rank) + except Exception as e: # perf is best-effort; never block completion + logger.warning( + f"KV transfer perf logging failed for request {self.request_id} " + f"slice={receiver_slice_id}: {e}" ) - - scatter_write_result( - self._receiver._bounce, - (self.disagg_request_id, task.slice_id), - peer_rank, - dst_ptrs, - sizes, - src_base, - on_done, - ) + logger.debug( + f"KV transfer complete for request {self.request_id} " + f"slice={receiver_slice_id}" + ) elif status == AgentResult.FAILED: detail = ( f"KV transfer failed for request {self.request_id} " @@ -4049,11 +3810,6 @@ def on_done( f"(reported by remote agent; see sender-side log for nixl_status)" ) logger.error(detail) - # Drain-before-release: record this writer FAILED; the owner frees the shared region - # only once every fan-in writer is terminal (freeing now could race a sibling's RMA). - self._receiver._bounce.record_failure( - (self.disagg_request_id, task.slice_id), peer_rank - ) task.fail(RuntimeError(detail)) if self._terminal_status is None: # Don't overwrite CANCELLED with ERROR self._terminal_status = SessionStatus.ERROR @@ -4103,9 +3859,7 @@ def process_aux_agent_result(self, peer_rank: int, status: AgentResult): if self._aux_physical_owner is not None: try: accepted, all_succeeded = self._aux_physical_owner.record_writer_result( - peer_rank, - status == AgentResult.SUCCESS, - wait_for_local_completion=False, + peer_rank, status == AgentResult.SUCCESS ) except Exception as error: self._aux_status = TaskStatus.ERROR @@ -4208,20 +3962,10 @@ def cancel_local(self, by_peer: bool = False) -> bool: self.cancelled_by_peer = by_peer exc = RuntimeError(f"RxSession {self.disagg_request_id} cancelled") for task in self._kv_tasks: - rid_slice = (self.disagg_request_id, task.slice_id) if task.status == TaskStatus.INIT: if self._enforce_physical_ownership: task.cancel_unpublished() - self._receiver._bounce.release_idle_reservation(rid_slice) task.fail(exc) - else: - has_active_access = ( - not task.resources_drained - if self._enforce_physical_ownership - else task.status == TaskStatus.TRANSFERRING - ) - if has_active_access: - self._receiver._bounce.orphan_reservation(rid_slice) if self._aux_physical_owner is not None and not self._publication_may_have_escaped: self._aux_physical_owner.cancel_unpublished() return True @@ -4369,10 +4113,6 @@ def close(self) -> bool: self.aux_slot = None # Unregister from Receiver; keep fields alive for in-flight listener messages. if self._receiver is not None: - # Reclaim any bounce region still live at teardown (closed mid-transfer) so it isn't - # leaked; a no-op for finished or non-bounce transfers. - for task in self._kv_tasks: - self._receiver._bounce.orphan_reservation((self.disagg_request_id, task.slice_id)) self._receiver.clear_session(self.disagg_request_id) return True @@ -4499,7 +4239,6 @@ class TransferWorkerConfig: max_draft_len: Optional[int] = None tx_timeout_s: Optional[float] = None rx_timeout_s: Optional[float] = None - bounce: Optional["Config"] = None tx_overall_timeout_s: Optional[float] = None enforce_physical_ownership: bool = False # Internal activation gate: only the qualified containment policy supplies it. @@ -4629,28 +4368,15 @@ def _setup_transfer_engine(self): self._register_kv_cache() if self._aux_buffer is not None: self._register_aux_buffer() - from .bounce import create_bounce - - self._bounce = create_bounce( - self._agent, - self._config.bounce, - device_id=self._config.device_id, - page_table=self._rank_info.page_table, - ) - # The bounce object owns its own NIXL descriptors (deregistered by Transport.close in - # shutdown), so they are NOT added to _registered_mem — single-source to avoid a - # double deregister. self._sender = Sender( self._peer_registrar, self._agent, - bounce=self._bounce, enforce_physical_ownership=self._config.enforce_physical_ownership, retirement_watchdog=self._retirement_watchdog, ) self._receiver = Receiver( self._peer_registrar, self._agent, - bounce=self._bounce, enforce_physical_ownership=self._config.enforce_physical_ownership, retirement_watchdog=self._retirement_watchdog, ) @@ -4726,16 +4452,6 @@ def shutdown(self): if rank_info_server is not None: rank_info_server.shutdown() self._shutdown = True - # Close the bounce transport (stops its scatter thread, deregisters its own descriptors and - # frees the VMM buffers) once the receiver listener is stopped so no new scatter is enqueued. - # No-op for the non-bounced path (NoBounceTransport.close); without this the daemon thread + fabric - # buffers leak until process exit. - bounce = getattr(self, "_bounce", None) - if bounce is not None: - try: - bounce.close() - except Exception as e: - logger.warning(f"TransferWorker.shutdown: bounce close failed: {e}") # Deregister NIXL memory before agent.shutdown so pinned GPU memory is released # (e.g. when the KV cache manager is recreated after profiling). agent = getattr(self, "_agent", None) diff --git a/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py b/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py index 68a8dce2ce59..57668587a07d 100644 --- a/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py +++ b/tensorrt_llm/_torch/disaggregation/nixl/_agent_cpp.py @@ -92,8 +92,8 @@ def __init__( for key, value in backend_params.items(): backend_params[key] = str(value) backend_params["num_threads"] = str(num_threads) - # C++ reads "num_workers" (not "num_threads") for progress threads, default 1. - # Env-gated so the default is unchanged (more threads can hurt the coalesced bounce WRITE). + # C++ reads "num_workers" (not "num_threads") for progress threads. It is set only from + # TRTLLM_NIXL_NUM_WORKERS; otherwise the backend default (1) applies. # Validate at the boundary: only a positive int reaches the backend. nixl_num_workers = os.environ.get("TRTLLM_NIXL_NUM_WORKERS") if nixl_num_workers is not None: diff --git a/tensorrt_llm/_torch/disaggregation/transceiver.py b/tensorrt_llm/_torch/disaggregation/transceiver.py index 81e59cb3744c..586f0888b808 100644 --- a/tensorrt_llm/_torch/disaggregation/transceiver.py +++ b/tensorrt_llm/_torch/disaggregation/transceiver.py @@ -588,11 +588,6 @@ def _chunk_num_bytes(self, chunk: Chunk) -> int: sender bytes against, so it must count what the mappers move, never physical slot bytes: a coalesced pool carries one view per role class (summing slot_bytes per view would double count the slot) and ignored-role buffers occupy slot offsets no view transfers. - - Counterpart accounting: the bounce reserve sizing (bounce/impl.py block_bytes_per_group) - computes per-block bytes for the same layer groups but reads pool 0 only, while this sums - every pool view of a group. The pool-0-only sizing gap for multi-pool attention groups is - tracked under TRTLLM-15194; keep the two accountings in mind together when changing either. """ pt = self._page_table if pt is None: @@ -1210,7 +1205,6 @@ def request_and_receive_async(self, req: LlmRequest) -> None: raise finally: # The session exists even when publication failed, and the legacy sweep owns it. - # TODO: An idle bounce reservation is not handed back on the failure path. session = self._legacy_session(fetches) if session is not None: self._recv_sessions[rid] = session diff --git a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py index 513ed51b65c2..f06aaa164194 100644 --- a/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py +++ b/tensorrt_llm/_torch/pyexecutor/connectors/mooncake_store/staging.py @@ -24,9 +24,9 @@ payload a registered pool would produce, so a pool written by one path is readable by the other. -Copies go through `cudaMemcpyAsync` rather than the batched Triton kernel in -`disaggregation/native/bounce/gather_scatter.py`, since one side here is host -memory and the copy engines move that over the host link by DMA. +Copies go through `cudaMemcpyAsync` rather than a batched device gather kernel, +since one side here is host memory and the copy engines move that over the host +link by DMA. """ from typing import List, Optional, Sequence, Tuple diff --git a/tests/integration/defs/.test_durations b/tests/integration/defs/.test_durations index ce6f1bdc56a8..b1c5b7cb319e 100644 --- a/tests/integration/defs/.test_durations +++ b/tests/integration/defs/.test_durations @@ -1618,7 +1618,6 @@ "unittest/bindings": 25.701486743371685, "unittest/disaggregated": 181.68471248142646, "unittest/disaggregated/test_agent_multi_backends.py": 77.37030741410489, - "unittest/disaggregated/test_bounce.py": 23.578477232924694, "unittest/disaggregated/test_cache_transceiver_harness.py": 51.59372871287128, "unittest/disaggregated/test_cache_transceiver_harness_report.py": 20.445199616122842, "unittest/disaggregated/test_cache_transceiver_precheck_e2e.py": 108.125814453125, diff --git a/tests/integration/test_lists/test-db/l0_a10.yml b/tests/integration/test_lists/test-db/l0_a10.yml index 2b5b6f887860..cf16c18ac152 100644 --- a/tests/integration/test_lists/test-db/l0_a10.yml +++ b/tests/integration/test_lists/test-db/l0_a10.yml @@ -72,7 +72,6 @@ l0_a10: - unittest/disaggregated/test_coordinator_worker.py - unittest/disaggregated/test_agent_multi_backends.py - unittest/disaggregated/test_extractor.py - - unittest/disaggregated/test_bounce.py - unittest/disaggregated/test_mamba_transfer.py - unittest/disaggregated/test_msgspec_route.py - unittest/disaggregated/test_steady_clock_sync.py diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 3305a7c62845..921319c0934b 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -112,7 +112,6 @@ l0_h100: - unittest/_torch/modeling/test_modeling_qwen3_5_vl.py::test_qwen35_dense_vl_preserves_w4a16_nvfp4_behavior - unittest/disaggregated/test_agent_multi_backends.py - unittest/disaggregated/test_extractor.py - - unittest/disaggregated/test_bounce.py - unittest/disaggregated/test_sync_transfer_timing.py - unittest/disaggregated/test_kv_transfer.py - unittest/disaggregated/test_kv_transfer_mp.py diff --git a/tests/unittest/disaggregated/lifecycle_acceptance.md b/tests/unittest/disaggregated/lifecycle_acceptance.md index 949e831e9c3c..9cd59670e1db 100644 --- a/tests/unittest/disaggregated/lifecycle_acceptance.md +++ b/tests/unittest/disaggregated/lifecycle_acceptance.md @@ -116,7 +116,7 @@ has passed. Track those results separately for each exact revision/configuration | A-ADMISSION | Cancel before publication; retain partial publication | Ownership: `test_pre_cancelled_rx_session_never_publishes_destination`, `test_aborted_publication_cannot_complete_during_failure_unwind`, `test_fp4_mla_bridge_roots_send_and_receive_requests_before_admission` | Actual executor/allocator lifetime and publication-race qualification | | A-CANCEL-REUSE | Repeated cancellation/status polling cannot release active destinations | Acceptance: `test_cancelled_receive_retains_destinations_until_all_accessors_settle` (4 cases); default-off companion has 2 strict xfails | Real allocator reuse under delayed remote writes and actual executor cancellation/timeout cleanup | | A-SOURCE | Exact backend completion protects the source | Ownership: `test_sender_physical_operation_rejects_unproven_retirement`, `test_sender_retires_only_after_backend_done`, `test_ambiguous_sender_result_retains_source_and_reports_in_doubt` | Backend conformance and real source-page retention | -| A-COHORT | All writers and independent holds must settle | Ownership: `test_gen_first_no_retry_adp_count_seal_waits_for_one_writer_group`, `test_gen_first_aux_owner_uses_anonymous_no_retry_writer_count`, `test_failed_receive_consensus_waits_for_every_rank_to_drain` | Multi-rank cohort, local CUDA completion, and manager-hold composition | +| A-COHORT | All writers and independent holds must settle | Ownership: `test_gen_first_no_retry_adp_count_seal_waits_for_one_writer_group`, `test_gen_first_aux_owner_uses_anonymous_no_retry_writer_count`, `test_failed_receive_consensus_waits_for_every_rank_to_drain` | Multi-rank cohort and manager-hold composition | | B-OUTCOME | Stable event-time logical outcomes | Outcome: `test_receive_outcome_does_not_depend_on_polling_before_late_completion`, `test_receive_cancel_keeps_its_outcome_when_a_writer_later_fails`, `test_sender_cancel_survives_late_kv_and_aux_exceptions` | Same semantics through native executor/backend paths | | B-LATE-DONE | Exact retained-handle settlement, idempotence, no revival/reopening | Merged [PR #19378](https://github.com/NVIDIA/TensorRT-LLM/pull/19378); Late settlement: `test_late_settlement_requires_positive_retained_status`, `test_concurrent_late_done_retires_once_without_changing_failure`, `test_late_done_does_not_retire_active_sibling` | Qualify retained-handle completion and actual source/destination reuse with the real backend | | B-DEADLINE | Non-resettable quiescence clock, independent progress, sticky fatal expiry | Merged [PR #19673](https://github.com/NVIDIA/TensorRT-LLM/pull/19673); Deadline: `test_first_drain_trigger_cannot_be_extended`, `test_blocked_backend_cannot_delay_fatal_deadline`, `test_receiver_deadline_retains_kv_aux_and_registration` | Real outstanding GPU/NIXL access through deadline expiry; coordinator request timeout alone does not test the drain clock | diff --git a/tests/unittest/disaggregated/test_bounce.py b/tests/unittest/disaggregated/test_bounce.py deleted file mode 100644 index 5e7af40dd760..000000000000 --- a/tests/unittest/disaggregated/test_bounce.py +++ /dev/null @@ -1,1121 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""CPU-only unit tests for the native-disagg KV ``bounce`` module. - -Covers sizing/OOM-guard math, the KV_AGENT_RESULT wire format, the NoBounceTransport -no-op fallback, and the TP fan-in reserve/record/settle logic (GPU allocators/streams mocked). -""" - -import queue -from types import SimpleNamespace - -import numpy as np -import pytest - -from tensorrt_llm._torch.disaggregation.native.bounce import config as bcfg - -# core (the BounceTransport contract + the TransferContext state machine) is PURE (no CUDA / NIXL -# imports), so it is always importable on CPU. -from tensorrt_llm._torch.disaggregation.native.bounce import core as bcore - -# transport/transfer pull CUDA-binding + torch deps at import; skip gracefully when those are -# absent (CPU-only env). Catch only ImportError so a genuine bug in the module still fails CI -# instead of being silently turned into a skip. -try: - from tensorrt_llm._torch.disaggregation.native.bounce import buffer as bbuf - from tensorrt_llm._torch.disaggregation.native.bounce import impl as btr - from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( - MambaPolicy, - mamba_receiver_payload_bytes, - ) - - _HAVE_TRANSPORT = True -except ImportError: # pragma: no cover - CPU-only env without CUDA bindings - _HAVE_TRANSPORT = False - -# page.py is pure numpy dataclasses, always importable on CPU. -from tensorrt_llm._torch.disaggregation.resource.page import ( - BUFFER_ENTRY_DTYPE, - AttentionLayerGroup, - KVCachePageTable, - LocalLayer, - MambaLayerGroup, - PhysicalPool, - PhysicalPoolGroup, - PoolView, -) - -_MIB = 1024 * 1024 - - -# --------------------------------------------------------------------------- # -# config.py — sizing + OOM guard (pure math, always runnable) -# --------------------------------------------------------------------------- # -class TestSizing: - @pytest.mark.parametrize( - "a,b,want", [(0, 8, 0), (1, 8, 8), (8, 8, 8), (9, 8, 16), (33, 32, 64)] - ) - def test_round_up(self, a, b, want): - assert bcfg._round_up(a, b) == want - - def test_fixed_sizing_rounds_up_to_chunk(self): - chunk = 32 * _MIB - ctx = bcfg.SizingContext( - free_bytes=80 * _MIB, total_bytes=200 * _MIB, chunk_bytes=chunk, device_id=0 - ) - # 50 MiB requested -> rounded up to a 64 MiB (2-chunk) multiple. - assert bcfg.FixedSizing(capacity_mb=50).resolve(ctx) == 64 * _MIB - - def test_fixed_sizing_floor_is_one_chunk(self): - chunk = 32 * _MIB - ctx = bcfg.SizingContext( - free_bytes=80 * _MIB, total_bytes=200 * _MIB, chunk_bytes=chunk, device_id=0 - ) - # below one chunk still gets a whole chunk. - assert bcfg.FixedSizing(capacity_mb=1).resolve(ctx) == chunk - - def test_fit_within_free_clamps_to_half_split_two_regions(self): - chunk = 32 * _MIB - # max_free_fraction=0.5 of 1024 MiB = 512 MiB, /2 regions = 256 MiB. - got = bcfg.fit_within_free(want := 4096 * _MIB, free_bytes=1024 * _MIB, chunk_bytes=chunk) - assert got == 256 * _MIB - assert got < want - - def test_fit_within_free_keeps_request_when_it_fits(self): - chunk = 32 * _MIB - got = bcfg.fit_within_free(128 * _MIB, free_bytes=4096 * _MIB, chunk_bytes=chunk) - assert got == 128 * _MIB - - def test_fit_within_free_none_when_too_small(self): - chunk = 32 * _MIB - # only 32 MiB free -> half/2 = 8 MiB < one chunk -> cannot fit. - assert bcfg.fit_within_free(64 * _MIB, free_bytes=32 * _MIB, chunk_bytes=chunk) is None - - def test_config_defaults(self): - cfg = bcfg.Config() - assert isinstance(cfg.sizing, bcfg.FixedSizing) - assert cfg.chunk_mb == 32 - # Byte gate for recurrent-state payloads; block gate for plain-KV payloads. - assert cfg.min_bytes == bcfg.DEFAULT_MIN_BYTES == 2 * _MIB - assert cfg.min_blocks == 96 - - -# --------------------------------------------------------------------------- # -# config.py — config on/off switch (size knob doubles as the enable flag) -# --------------------------------------------------------------------------- # -class TestConfigFromSize: - @pytest.mark.parametrize("size_mb", [0, -1, None]) - def test_non_positive_is_off(self, size_mb): - assert bcfg.config_from_size(size_mb) is None - - def test_positive_enables_with_capacity(self): - cfg = bcfg.config_from_size(2048) - assert isinstance(cfg, bcfg.Config) - assert cfg.sizing.capacity_mb == 2048 - - def test_min_bytes_defaults_and_overrides(self, monkeypatch): - monkeypatch.delenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES", raising=False) - assert bcfg.config_from_size(2048).min_bytes == 2 * _MIB # keeps the Config default - assert bcfg.config_from_size(2048, min_bytes=1).min_bytes == 1 # explicit arg overrides - assert bcfg.config_from_size(2048, min_bytes=64 * _MIB).min_bytes == 64 * _MIB - # The gate is tuned for tests via the env override (no user-facing config field). - monkeypatch.setenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES", "1") - assert bcfg.config_from_size(2048).min_bytes == 1 # env override - assert bcfg.config_from_size(2048, min_bytes=250).min_bytes == 250 # arg beats env - - def test_min_blocks_defaults_and_overrides(self, monkeypatch): - # The block-count gate (plain-KV payloads) keeps its original default of 96 and honors - # the explicit arg and the env override. - monkeypatch.delenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", raising=False) - assert bcfg.config_from_size(2048).min_blocks == 96 # keeps the Config default - assert bcfg.config_from_size(2048, 250).min_blocks == 250 # explicit arg still overrides - monkeypatch.setenv("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", "48") - assert bcfg.config_from_size(2048).min_blocks == 48 # env override - assert bcfg.config_from_size(2048, 250).min_blocks == 250 # explicit arg beats env - - @pytest.mark.parametrize( - "env,attr,default", - [ - ("TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES", "min_bytes", 2 * _MIB), - ("TRTLLM_KV_CACHE_BOUNCE_MIN_BLOCKS", "min_blocks", 96), - ], - ) - def test_gate_env_is_parsed_defensively(self, monkeypatch, env, attr, default): - # A bad value must not crash setup (falls back to the default); a non-positive value clamps to 1. - for bad in ("", "auto", "1.5"): - monkeypatch.setenv(env, bad) - assert getattr(bcfg.config_from_size(2048), attr) == default - for nonpos in ("0", "-5"): - monkeypatch.setenv(env, nonpos) - assert getattr(bcfg.config_from_size(2048), attr) == 1 - - -# --------------------------------------------------------------------------- # -# wire format — KV_AGENT_RESULT struct prefix + bounce result tail -# --------------------------------------------------------------------------- # -@pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.transport import needs CUDA bindings") -class TestResultTail: - def _wm(self, dst, sizes, base): - return SimpleNamespace( - dst_ptrs=np.array(dst, dtype=np.int64), - sizes=np.array(sizes, dtype=np.int64), - bounce_dst_base=base, - ) - - def test_roundtrip_with_tail(self): - wm = self._wm([100, 200, 300], [16, 16, 16], 0xABCD) - # message = [msg_type, packed_prefix, *tail]; tail lives at [2:]. - msg = [b"KV_AGENT_RESULT", b"prefix"] + btr.encode_result_tail(wm) - assert len(msg) == 5 - dst, sizes, src_base = btr.decode_result_tail(msg) - assert np.array_equal(dst, wm.dst_ptrs) - assert np.array_equal(sizes, wm.sizes) - assert src_base == 0xABCD - - def test_no_tail_returns_none(self): - # non-bounced result: only [msg_type, prefix] -> no tail. - assert btr.decode_result_tail([b"KV_AGENT_RESULT", b"prefix"]) == (None, None, None) - - def test_encode_tail_handles_unset_base(self): - wm = self._wm([1, 2], [8, 8], None) - tail = btr.encode_result_tail(wm) - assert int(np.frombuffer(tail[2], dtype=np.int64)[0]) == 0 - - -def test_kv_result_prefix_roundtrip(): - """The KV_AGENT_RESULT binary prefix (transfer.py) must round-trip exactly.""" - tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") - for rank, rid, sl, last, status, size in [ - (7, 6925227277844486, 42, True, tfr.AgentResult.SUCCESS, 4096), - (0, 1, 0, False, tfr.AgentResult.FAILED, 0), - (31, 2**62, 9999, True, tfr.AgentResult.SUCCESS, 2**40), - ]: - packed = tfr._KV_RESULT_PREFIX.pack( - rank, rid, sl, last, tfr._AGENT_RESULT_CODE[status], size - ) - r, i, s, last_out, c, sz = tfr._KV_RESULT_PREFIX.unpack(packed) - assert (r, i, s, last_out, sz) == (rank, rid, sl, last, size) - assert tfr._AGENT_RESULT_BY_CODE[c] is status - - -@pytest.mark.parametrize("result_name", ["SUCCESS", "FAILED"]) -def test_make_kv_result_msg_uses_binary_frame(result_name): - """Every KV result (success and failure) uses the binary frame so the receiver can decode it.""" - tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") - result = getattr(tfr.AgentResult, result_name) - msg = tfr._make_kv_result_msg(3, 12345, 7, True, result, transfer_size=8192) - assert msg[0] == tfr.MessageType.KV_AGENT_RESULT - assert len(msg) == 2 # prefix only; no bounce tail when none is passed - r, rid, sl, last, code, size = tfr._KV_RESULT_PREFIX.unpack(msg[1]) - assert (r, rid, sl, last, size) == (3, 12345, 7, True, 8192) - assert tfr._AGENT_RESULT_BY_CODE[code] is result - - -# --------------------------------------------------------------------------- # -# fan-in safety gate — equal total//num_writers split only for uniform writers -# --------------------------------------------------------------------------- # -def test_fanin_bounce_safe_gate() -> None: - """Restrict multi-writer equal-split bounce to uniform TP-by-head. - - Uneven PP splits and duplicate_head_factor>1 (MLA / duplicate TP heads -> - some ranks don't send KV yet count in expected_transfers) must fall back to - the per-fragment path. - """ - tfr = pytest.importorskip("tensorrt_llm._torch.disaggregation.native.transfer") - from tensorrt_llm._torch.disaggregation.resource.page import MapperKind - - safe = tfr.Receiver._fanin_bounce_safe - - def ov(dup: int, pp: int, ranks: tuple[int, ...] = (0,)) -> SimpleNamespace: - return SimpleNamespace(duplicate_head_factor=dup, overlap_pp_size=pp, ranks=list(ranks)) - - def ri(lpp: list[int], page_table: SimpleNamespace | None = None) -> SimpleNamespace: - return SimpleNamespace(layer_num_per_pp=lpp, page_table=page_table) - - def pt(mapper_kind: MapperKind) -> SimpleNamespace: - view = SimpleNamespace(mapper_kind=mapper_kind) - return SimpleNamespace(layer_groups=[SimpleNamespace(pool_views=[view])]) - - # single PP stage (overlap_pp_size <= 1): only duplicate_head_factor matters - assert safe(ov(1, 1), ri([24]), None) is True - assert safe(ov(1, 0), ri([24]), None) is True - assert safe(ov(2, 1), ri([24]), None) is False # duplicate heads / MLA -> some don't send - # EVEN PP fan-in (equal layers per overlapping stage) -> allowed - assert safe(ov(1, 4), ri([20, 20, 20, 20]), None) is True - # UNEVEN PP fan-in -> per-writer sizes differ -> fall back - assert safe(ov(1, 4), ri([20, 20, 20, 19]), None) is False - # incomplete per-stage info (single element for a multi-stage fan-in) -> conservative fall back - assert safe(ov(1, 4), ri([20]), None) is False - # duplicate heads blocks even an otherwise-even PP split - assert safe(ov(2, 4), ri([20, 20, 20, 20]), None) is False - # replicated views (one elected sender per destination) make multi-writer - # contributions unequal -> fall back; single-writer overlap stays safe, - # and sharded-only view schemes are unaffected. Both page tables must be - # checked because the representative peer rank may be a fully masked PP - # stage while another sender stage owns the replicated rows. - assert safe(ov(1, 1, ranks=(0, 1)), ri([24], pt(MapperKind.REPLICATED)), None) is False - assert safe(ov(1, 1, ranks=(0,)), ri([24], pt(MapperKind.REPLICATED)), None) is True - assert ( - safe( - ov(1, 1, ranks=(0, 1)), - ri([24], pt(MapperKind.NHD)), - pt(MapperKind.REPLICATED), - ) - is False - ) - assert ( - safe( - ov(1, 1, ranks=(0, 1)), - ri([24], pt(MapperKind.NHD)), - pt(MapperKind.NHD), - ) - is True - ) - - -# --------------------------------------------------------------------------- # -# NoBounceTransport — disabled no-op fallback -# --------------------------------------------------------------------------- # -@pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.transport import needs CUDA bindings") -class TestNoBounce: - def test_noop_behaviour(self): - nb = btr.NoBounceTransport() - assert nb.enabled is False - assert nb.reserve(SimpleNamespace()) is False - assert nb.build_request(SimpleNamespace()) is None - assert nb.writer_base(("r", 0), 1) is None - assert nb.is_bounced(("r", 0)) is False - nb.record_result(("r", 0), 1) # no-op, must not raise - nb.record_failure(("r", 0), 1) # no-op, must not raise - nb.release_idle_reservation(("r", 0)) # no-op, must not raise - nb.release_send(0) - nb.close() - - def test_create_bounce_none_cfg(self): - assert isinstance( - btr.create_bounce(object(), None, device_id=0, page_table=None), btr.NoBounceTransport - ) - - def test_implements_bounce_transport_abc(self): - # Both implementations must satisfy the explicit contract; instantiating NoBounceTransport - # proves no abstractmethod is missing (guards the two impls against drifting apart). - assert issubclass(btr.NoBounceTransport, bcore.BounceTransport) - assert issubclass(btr.VmmBounceTransport, bcore.BounceTransport) - assert isinstance(btr.NoBounceTransport(), bcore.BounceTransport) - - -# --------------------------------------------------------------------------- # -# Transport TP fan-in logic — reserve gate / writer_base / accumulate ordering -# (GPU allocators, streams and the scatter worker are mocked out) -# --------------------------------------------------------------------------- # -class _FakeAlloc: - """Stand-in for SlotAllocator: hands out a fixed base, records reservations.""" - - def __init__(self, capacity_bytes, phys_chunk_size, name="kv_bounce"): - self._cap = capacity_bytes - self.base = 0x100000 - self.next_id = 0 - self.released = [] - self.quarantined = [] - self.reserved_sizes = [] - - @property - def capacity(self): - return self._cap - - def reserve(self, size, timeout=None): - if size > self._cap: - return None - sid = self.next_id - self.next_id += 1 - self.reserved_sizes.append(size) - return sid, self.base - - def release(self, slot_id): - self.released.append(slot_id) - - def quarantine(self, slot_id, grace_s): - self.quarantined.append(slot_id) - - def reclaim_expired(self): - return 0 - - def reg_descs(self): - return [] - - -def _make_transport( - monkeypatch, block_bytes_per_group, capacity=1 << 30, min_bytes=1, min_blocks=1 -): - monkeypatch.setattr(btr, "SlotAllocator", _FakeAlloc) - monkeypatch.setattr(btr.VmmBounceTransport, "_new_stream", lambda self: 0) - monkeypatch.setattr( - btr.VmmBounceTransport, - "_start_scatter_worker", - lambda self, name: setattr(self, "_scatter_q", queue.Queue()), - ) - agent = SimpleNamespace(register_memory=lambda d: None) - return btr.VmmBounceTransport( - agent, - device_id=0, - capacity_bytes=capacity, - phys_chunk_size=32 * _MIB, - block_bytes_per_group=block_bytes_per_group, - min_bytes=min_bytes, - min_blocks=min_blocks, - ) - - -def _recv_req(block_counts, rid=1, slice_id=0): - # Positional block tables: `n` held slots followed by one -1 hole, which - # the reserve sizing must not count. - return SimpleNamespace( - block_ids_per_layer_groups=[ - np.concatenate([np.arange(n, dtype=np.int64), np.array([-1], dtype=np.int64)]) - if n - else np.array([], dtype=np.int64) - for n in block_counts - ], - unique_rid=rid, - slice_id=slice_id, - bounce_dst_base=None, - ) - - -@pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.transport import needs CUDA bindings") -class TestFanInReserve: - def test_block_bytes_include_distinct_physical_pools_once(self) -> None: - entries = np.array([], dtype=BUFFER_ENTRY_DTYPE) - page_table = KVCachePageTable( - tokens_per_block=32, - layer_groups=[ - AttentionLayerGroup( - pool_group_idx=0, - local_layers=[LocalLayer(local_layer_id=0, global_layer_id=0)], - pool_views=[ - PoolView(pool_idx=0, buffer_entries=entries), - PoolView(pool_idx=1, buffer_entries=entries), - PoolView(pool_idx=1, buffer_entries=entries), - ], - ) - ], - pool_groups=[ - PhysicalPoolGroup( - pools=[ - PhysicalPool(base_address=0x200000, slot_bytes=100, num_slots=8), - PhysicalPool(base_address=0x300000, slot_bytes=25, num_slots=8), - ] - ) - ], - ) - - assert btr.block_bytes_per_group(page_table) == [125] - - def test_reserve_stamps_base_and_per_writer(self, monkeypatch): - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - req = _recv_req([2]) # total = 2 * 100 = 200 - assert t.reserve(req, num_writers=2) is True - assert req.bounce_dst_base == 0x100000 - # writer i lands at base + i * (total // num_writers); per_writer = 100. - assert t.writer_base((req.unique_rid, req.slice_id), 0) == 0x100000 - assert t.writer_base((req.unique_rid, req.slice_id), 1) == 0x100000 + 100 - - def test_reserve_uneven_fanin_falls_back(self, monkeypatch): - t = _make_transport(monkeypatch, block_bytes_per_group=[3]) - req = _recv_req([1]) # total = 3, not divisible by 2 - assert t.reserve(req, num_writers=2) is False - assert req.bounce_dst_base is None - - def test_reserve_heterogeneous_fanin_falls_back(self, monkeypatch): - # Two groups with different per-block sizes: the equal split would overrun a sub-region, so - # fall back (even though the total is divisible). - t = _make_transport(monkeypatch, block_bytes_per_group=[100, 200]) - req = _recv_req([2, 2]) # total = 2*100 + 2*200 = 600 - assert t.reserve(req, num_writers=2) is False - assert req.bounce_dst_base is None - - def test_reserve_uniform_multigroup_fanin_ok(self, monkeypatch): - # Uniform slot bytes across present groups -> even byte split -> bounce allowed. - t = _make_transport(monkeypatch, block_bytes_per_group=[100, 100]) - assert t.reserve(_recv_req([2, 2]), num_writers=2) is True - - def test_reserve_heterogeneous_single_writer_ok(self, monkeypatch): - # num_writers==1 has no split, so heterogeneous slot bytes are fine. - t = _make_transport(monkeypatch, block_bytes_per_group=[100, 200]) - assert t.reserve(_recv_req([2, 2]), num_writers=1) is True - - def test_reserve_single_writer_ok(self, monkeypatch): - t = _make_transport(monkeypatch, block_bytes_per_group=[3]) - req = _recv_req([1]) # total = 3, num_writers=1 -> no even-split requirement - assert t.reserve(req, num_writers=1) is True - - def test_reserve_below_min_bytes_falls_back(self, monkeypatch): - # Recurrent-state payloads gate on BYTES: 4 blocks x 100B + 100B state = 500B < 1000B - # falls back, regardless of how many blocks that is. - t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_bytes=1000) - req = _recv_req([4]) - assert t.reserve(req, num_writers=1, extra_bytes=100) is False - assert req.bounce_dst_base is None - - def test_reserve_at_min_bytes_bounces(self, monkeypatch): - # exactly at the threshold (9 x 100B + 100B state = 1000B >= 1000B) the transfer bounces - t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_bytes=1000) - req = _recv_req([9]) - assert t.reserve(req, num_writers=1, extra_bytes=100) is True - assert req.bounce_dst_base == 0x100000 - - def test_reserve_recurrent_payload_ignores_block_gate(self, monkeypatch): - # The K3 regression shape: FEW but LARGE blocks plus recurrent state must clear the byte - # gate even though a block-count gate of 96 fails (67 blocks x 6.5 MiB = 433 MiB). - t = _make_transport( - monkeypatch, - block_bytes_per_group=[int(6.5 * _MIB)], - min_bytes=2 * _MIB, - min_blocks=96, - ) - assert t.reserve(_recv_req([67]), num_writers=1, extra_bytes=1024) is True - - def test_reserve_plain_kv_uses_block_gate(self, monkeypatch): - # Plain-KV payloads (no recurrent state) keep the original block-count gate, and the byte - # gate does not apply to them (96 x 100B = 9600B passes despite min_bytes far above it). - t = _make_transport( - monkeypatch, block_bytes_per_group=[100], min_blocks=96, min_bytes=1 << 20 - ) - assert t.reserve(_recv_req([4]), num_writers=1) is False # 4 < 96 blocks - assert t.reserve(_recv_req([96]), num_writers=1) is True - - def test_reserve_unknown_slot_size_falls_back(self, monkeypatch): - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) # only 1 group known - assert t.reserve(_recv_req([2, 2]), num_writers=1) is False # 2nd group unknown - - def test_reserve_oversize_falls_back(self, monkeypatch): - t = _make_transport(monkeypatch, block_bytes_per_group=[1000], capacity=500) - assert t.reserve(_recv_req([2]), num_writers=1) is False # total 2000 > cap 500 - - def test_reserve_does_not_wait_for_recv_space(self, monkeypatch): - # The receiver reserves on the executor thread, so a full recv region must fall back to - # the per-fragment path at once instead of waiting for space to free up. - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - timeouts = [] - - def full_region(size, timeout=None): - timeouts.append(timeout) - return None - - monkeypatch.setattr(t._recv_alloc, "reserve", full_region) - req = _recv_req([2]) - assert t.reserve(req, num_writers=1) is False - assert timeouts == [0.0] - assert req.bounce_dst_base is None - - def test_reserve_state_group_fanin_falls_back(self, monkeypatch): - # STATE groups (None bytes) with num_writers > 1 must disable bounce even - # when extra_bytes == 0: the representative sender may have no mamba - # overlap, but another PP sender might still append state fragments that - # were not reserved in the bounce region. - t = _make_transport(monkeypatch, block_bytes_per_group=[100, None]) - # Two groups: PAGED (2 blocks x 100B) + STATE (1 slot ID, None bytes). - req = _recv_req([2, 1]) - assert t.reserve(req, num_writers=2, extra_bytes=0) is False - assert req.bounce_dst_base is None - - def test_reserve_state_group_single_writer_ok(self, monkeypatch): - # Single writer with STATE group is fine — no fan-in split concern. - t = _make_transport(monkeypatch, block_bytes_per_group=[100, None]) - req = _recv_req([2, 1]) - assert t.reserve(req, num_writers=1, extra_bytes=50) is True - assert req.bounce_dst_base is not None - - def test_fanin_scatters_ordered_by_src_base(self, monkeypatch): - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - req = _recv_req([2]) - assert t.reserve(req, num_writers=2) is True - rid_slice = (req.unique_rid, req.slice_id) - # writer for the HIGHER src_base reports first; scatter must reorder by src_base. - t.record_result( - rid_slice, - 7, - np.array([20], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=200, - ) - assert t._scatter_q.empty() # only 1 of 2 writers terminal -> no scatter - assert not t._recv_alloc.released # region NOT freed while a writer is still pending - t.record_result( - rid_slice, - 3, - np.array([10], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=100, - ) - ctx, descs = t._scatter_q.get_nowait() - # each tail carries its OWN src_base; sorted (100 before 200) so the scatter is deterministic. - assert [t[0] for t in descs] == [100, 200] # per-writer src_base preserved - assert [list(t[1]) for t in descs] == [[10], [20]] # dst_ptrs - assert [list(t[2]) for t in descs] == [[8], [8]] # sizes - - def test_fanin_fallback_writer_leaves_survivor_at_its_own_base(self, monkeypatch): - # Regression: if one fan-in writer falls back to in-place (SUCCESS, empty tail) while a - # sibling bounces, the survivor must be scattered from ITS OWN src_base, not packed to 0. - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - req = _recv_req([2]) - assert t.reserve(req, num_writers=2) is True - rid_slice = (req.unique_rid, req.slice_id) - t.record_result( - rid_slice, 7, None, None - ) # writer 0 fell back to in-place: SUCCESS, no tail - assert t._scatter_q.empty() # only 1 of 2 writers terminal - t.record_result( - rid_slice, - 3, - np.array([10], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=100, - ) # writer 1 bounced to base+100 - ctx, descs = t._scatter_q.get_nowait() - assert [t[0] for t in descs] == [100] # only the survivor, read from base+100 (NOT 0) - assert [list(t[1]) for t in descs] == [[10]] - - def test_fanin_failed_then_success_releases_only_after_both(self, monkeypatch): - # A FAILED writer must not free the shared region until every writer is terminal. - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - req = _recv_req([2]) - assert t.reserve(req, num_writers=2) is True - rid_slice = (req.unique_rid, req.slice_id) - t.record_failure(rid_slice, 7) # first writer fails - assert not t._recv_alloc.released # region held while a sibling may still be in flight - assert t.is_bounced(rid_slice) is True - t.record_result( - rid_slice, 3, np.array([10], dtype=np.int64), np.array([8], dtype=np.int64), src_base=0 - ) - # all terminal, >=1 FAILED -> no scatter, release (both drained), region freed. - assert t._scatter_q.empty() - assert t._recv_alloc.released - assert not t._recv_alloc.quarantined # FAILED has drained -> release, NOT quarantine - assert t.is_bounced(rid_slice) is False - - def test_on_done_fires_after_scatter_lands(self, monkeypatch): - # The completion cb rides the context and fires only when the worker records the scatter as - # done -> the gen never observes completion before the KV is scattered into place. - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - req = _recv_req([2]) - assert t.reserve(req, num_writers=1) is True - rid_slice = (req.unique_rid, req.slice_id) - calls = [] - t.record_result( - rid_slice, - 3, - np.array([10], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=0, - on_done=lambda ok: calls.append(ok), - ) - ctx, descs = t._scatter_q.get_nowait() - assert calls == [] # deferred, not fired inline - t._apply( - ctx.rid_slice, lambda c: c.finish_scatter(True) - ) # worker records scatter done -> settle - assert calls == [True] - assert t._recv_alloc.released - - def test_empty_acc_fires_on_done_inline_and_releases(self, monkeypatch): - # Bounced SUCCESS that carried no scatter tail: nothing to copy, but the task must still - # complete -> on_done(True) inline + slot released, nothing queued. - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - req = _recv_req([2]) - assert t.reserve(req, num_writers=1) is True - rid_slice = (req.unique_rid, req.slice_id) - calls = [] - t.record_result(rid_slice, 3, None, None, on_done=lambda ok: calls.append(ok)) - assert calls == [True] - assert t._scatter_q.empty() - assert t._recv_alloc.released # slot freed - - def test_missing_key_is_dropped(self, monkeypatch): - # A late/duplicate result for an already-settled (popped) rid_slice is dropped: the context's own - # settle already fired completion, so re-firing here would double-report. - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - calls = [] - t.record_result( - ("missing", 0), - 3, - np.array([10], dtype=np.int64), - np.array([8], dtype=np.int64), - src_base=0, - on_done=lambda ok: calls.append(ok), - ) - assert calls == [] # no-op, no callback - - def test_duplicate_writer_is_ignored(self, monkeypatch): - # A duplicate SUCCESS from the same peer_rank must not double-count toward all-terminal. - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - req = _recv_req([2]) - assert t.reserve(req, num_writers=2) is True - rid_slice = (req.unique_rid, req.slice_id) - arr = (np.array([10], dtype=np.int64), np.array([8], dtype=np.int64)) - t.record_result(rid_slice, 7, *arr, src_base=0) - t.record_result(rid_slice, 7, *arr, src_base=0) # duplicate of the SAME writer - assert t._scatter_q.empty() # still only 1 distinct writer -> not all terminal - assert not t._recv_alloc.released - - def test_scatter_write_result_non_bounce_fires_on_done(self): - # Non-bounced path completes inline (the in-place WRITE already landed the KV). - calls = [] - btr.scatter_write_result( - btr.NoBounceTransport(), ("r", 0), 1, None, None, on_done=lambda ok: calls.append(ok) - ) - assert calls == [True] - - def test_release_idle_reservation_frees_slot_and_is_idempotent(self, monkeypatch): - # INIT-cancel immediate release (nothing published) must free the recv slot; idempotent. - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - req = _recv_req([2]) - assert t.reserve(req, num_writers=1) is True - rid_slice = (req.unique_rid, req.slice_id) - assert t.is_bounced(rid_slice) is True - t.release_idle_reservation(rid_slice) - assert t.is_bounced(rid_slice) is False - assert t._recv_alloc.released # slot freed - t.release_idle_reservation(rid_slice) # already gone -> no-op, must not raise - - def test_orphan_reservation_quarantines_and_is_idempotent(self, monkeypatch): - # Giving up on an in-flight reservation must quarantine the region (a write may still land), - # not release or leak it; a second give-up is a no-op. - t = _make_transport(monkeypatch, block_bytes_per_group=[100]) - req = _recv_req([2]) - assert t.reserve(req, num_writers=1) is True - rid_slice = (req.unique_rid, req.slice_id) - t.orphan_reservation(rid_slice) - assert t._recv_alloc.quarantined == [0] # quarantined, not released - assert t._recv_alloc.released == [] - assert t.is_bounced(rid_slice) is False # settled and removed from the live map - t.orphan_reservation(rid_slice) # already gone -> no-op, must not raise - assert t._recv_alloc.quarantined == [0] - - -# --------------------------------------------------------------------------- # -# Hybrid layouts (Kimi K3: KDA mamba + MLA attention) — page-table gate, -# trailing-empty-group handling, and recurrent-state (extra_bytes) sizing -# --------------------------------------------------------------------------- # - -# Kimi K3 geometry: 24 MLA layers in one attention layer group; 69 KDA layers -# whose per-layer recurrent state is a bf16 short-conv slot + fp32 delta slot, -# replicated per rank. -_K3_MLA_BLOCK_BYTES = 32 * 576 * 2 * 24 # tpb x (kv_lora_rank+rope) x bf16 x 24 layers = 884,736 -_K3_KDA_LAYERS = 69 -_K3_CONV_SLOT_BYTES = 221_184 # [3*H*hd, W-1] bf16 per layer -_K3_SSM_SLOT_BYTES = 6_291_456 # [H, hd, hd] fp32 per layer (95.5% of the state) -_K3_KDA_PAYLOAD_BYTES = ( - 449_372_160 # 69 x (conv + delta) per request per rank, from the geometry above -) - - -def _k3_page_table() -> KVCachePageTable: - """A K3-shaped page table exactly as the builders produce it. - - Attention layer group(s) first and the mamba layer group appended LAST - (kv_extractor.py, both ``build_page_table`` and ``_build_page_table_v2``, - append it after every attention group). - """ - from tensorrt_llm._torch.disaggregation.resource.page import ( - MAMBA_CONV_ROLE, - MAMBA_SSM_ROLE, - MapperKind, - ) - - attn = AttentionLayerGroup( - pool_group_idx=0, - kv_head_num_per_rank=1, - sliding_window_size=None, - local_layers=[LocalLayer(i, i) for i in range(24)], - pool_views=[PoolView(pool_idx=0, buffer_entries=np.array([], dtype=BUFFER_ENTRY_DTYPE))], - ) - sorted_lids = list(range(_K3_KDA_LAYERS)) - local_layers = [ - LocalLayer(local_layer_id=i, global_layer_id=gl) - for i, gl in enumerate(range(24, 24 + _K3_KDA_LAYERS)) - ] - conv_pool = PhysicalPool(0x200000, _K3_CONV_SLOT_BYTES, 8) - ssm_pool = PhysicalPool(0x300000, _K3_SSM_SLOT_BYTES, 8) - pool_views = [ - PoolView( - pool_idx=0, - buffer_entries=np.array( - [(lid, 0, _K3_CONV_SLOT_BYTES) for lid in sorted_lids], dtype=BUFFER_ENTRY_DTYPE - ), - pool_role=MAMBA_CONV_ROLE, - mapper_kind=MapperKind.SECTIONED, - bytes_per_layer=_K3_CONV_SLOT_BYTES, - section_bytes=[_K3_CONV_SLOT_BYTES // 3] * 3, - ), - PoolView( - pool_idx=1, - buffer_entries=np.array( - [(lid, 0, _K3_SSM_SLOT_BYTES) for lid in sorted_lids], dtype=BUFFER_ENTRY_DTYPE - ), - pool_role=MAMBA_SSM_ROLE, - mapper_kind=MapperKind.INDEXED, - bytes_per_layer=_K3_SSM_SLOT_BYTES, - bytes_per_head=_K3_SSM_SLOT_BYTES // 4, - ), - ] - mamba = MambaLayerGroup( - pool_group_idx=1, - local_layers=local_layers, - pool_views=pool_views, - ) - return KVCachePageTable( - tokens_per_block=32, - layer_groups=[attn, mamba], - pool_groups=[ - PhysicalPoolGroup(pools=[PhysicalPool(0x100000, _K3_MLA_BLOCK_BYTES, 128)]), - PhysicalPoolGroup(pools=[conv_pool, ssm_pool]), - ], - ) - - -def _k3_rank_info(tp_size=1, tp_rank=0, cp_size=1, cp_rank=0): - # MambaPolicy reads tp_size / tp_rank / cp_size / cp_rank / - # attention.enable_attention_dp. - return SimpleNamespace( - tp_size=tp_size, tp_rank=tp_rank, cp_size=cp_size, cp_rank=cp_rank, attention=None - ) - - -@pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.transport import needs CUDA bindings") -class TestHybridK3Bounce: - def test_block_bytes_per_group_keeps_mamba_placeholder(self): - # The mamba group must stay in the list as a placeholder (None), aligned with the - # layer-group indices a recv request uses, instead of truncating the list. - assert btr.block_bytes_per_group(_k3_page_table()) == [_K3_MLA_BLOCK_BYTES, None] - - def test_reserve_engages_on_k3_mixed_layout(self, monkeypatch): - # Regression pin: a K3 recv request always carries a trailing EMPTY entry for - # the mamba layer group (transceiver._create_chunk), which used to trip the - # unknown-slot-size guard and silently push every K3 request onto the per-fragment - # (~0.4 GB/s host-staged) path. It must engage bounce, sized for MLA KV + KDA state. - t = _make_transport( - monkeypatch, - block_bytes_per_group=btr.block_bytes_per_group(_k3_page_table()), - min_bytes=2 * _MIB, - ) - req = _recv_req([67, 0]) # 67 MLA blocks (~2144 tokens) + the empty mamba entry - assert t.reserve(req, num_writers=1, extra_bytes=_K3_KDA_PAYLOAD_BYTES) is True - assert req.bounce_dst_base == 0x100000 - # the region covers the MLA KV blocks plus the appended KDA state - assert t._recv_alloc.reserved_sizes == [67 * _K3_MLA_BLOCK_BYTES + _K3_KDA_PAYLOAD_BYTES] - - def test_reserve_nonempty_slot_group_with_unknown_size_skipped(self, monkeypatch): - # Non-attention groups (STATE-based) carry a slot index but their bytes - # are in extra_bytes — bounce skips them rather than rejecting. - t = _make_transport(monkeypatch, block_bytes_per_group=[_K3_MLA_BLOCK_BYTES, None]) - assert t.reserve(_recv_req([2, 1]), num_writers=1) is True - - def test_reserve_extra_bytes_counts_toward_min_bytes(self, monkeypatch): - t = _make_transport(monkeypatch, block_bytes_per_group=[100], min_bytes=1000) - assert t.reserve(_recv_req([4]), num_writers=1, extra_bytes=599) is False # 999 < 1000 - assert t.reserve(_recv_req([4]), num_writers=1, extra_bytes=600) is True # exactly 1000 - - def test_reserve_extra_bytes_counts_toward_capacity(self, monkeypatch): - # extra_bytes must be part of the reservation, or the sender's coalesced write - # (KV + recurrent state) would overrun the region into the neighboring slot. - t = _make_transport(monkeypatch, block_bytes_per_group=[100], capacity=500) - assert t.reserve(_recv_req([2]), num_writers=1, extra_bytes=400) is False # 600 > 500 - - def test_reserve_extra_bytes_fanin_falls_back(self, monkeypatch): - # Fan-in with a STATE group present falls back: per-writer recurrent-state - # fragments can differ (PP stages hold different mamba layers). - t = _make_transport(monkeypatch, block_bytes_per_group=[100, None]) - assert t.reserve(_recv_req([2, 1]), num_writers=2, extra_bytes=64) is False - assert t.reserve(_recv_req([2, 0]), num_writers=2) is True # no state blocks -> fan-in fine - - def test_receiver_payload_bytes_matches_k3_geometry(self): - pt = _k3_page_table() - got = mamba_receiver_payload_bytes(pt, pt, dst_slot=3) - assert got == _K3_KDA_PAYLOAD_BYTES - - def test_receiver_payload_bytes_slot_independent(self): - pt = _k3_page_table() - got_3 = mamba_receiver_payload_bytes(pt, pt, dst_slot=3) - got_5 = mamba_receiver_payload_bytes(pt, pt, dst_slot=5) - assert got_3 == got_5 > 0 - - def test_receiver_payload_bytes_zero_without_slot_or_group(self): - pt = _k3_page_table() - assert mamba_receiver_payload_bytes(pt, pt, dst_slot=None) == 0 - pure_attn = KVCachePageTable( - tokens_per_block=32, - layer_groups=[pt.layer_groups[0]], - pool_groups=pt.pool_groups, - ) - assert mamba_receiver_payload_bytes(pure_attn, pure_attn, dst_slot=3) == 0 - - -# --------------------------------------------------------------------------- # -# MambaPolicy under decode-CP (helix): shard grid, pairing, receiver sizing -# --------------------------------------------------------------------------- # -class TestMambaHelixCP: - def test_mamba_tp_helix_grid_is_cp_minor(self): - # cp == 1 keeps the plain TP grid; cp > 1 flattens tp*cp CP-minor. - assert MambaPolicy._mamba_tp(_k3_rank_info()) == (1, 0) - assert MambaPolicy._mamba_tp(_k3_rank_info(cp_size=32, cp_rank=5)) == (32, 5) - assert MambaPolicy._mamba_tp(_k3_rank_info(tp_size=2, tp_rank=1, cp_size=4, cp_rank=3)) == ( - 8, - 7, - ) - - def test_equal_width_unpaired_sender_sends_nothing(self): - # Regression lock for the fan-in collapse: with equal shard widths the - # mapper is a whole-slot copy, so an unpaired sender must return zero - # bytes instead of overwriting the receiver slot with the wrong heads. - for sender_rank in range(2): - for gen_cp_rank in range(2): - sender_ri = _k3_rank_info(tp_size=2, tp_rank=sender_rank) - peer_ri = _k3_rank_info(cp_size=2, cp_rank=gen_cp_rank) - paired = MambaPolicy.is_paired(sender_ri, peer_ri) - if sender_rank == gen_cp_rank: - assert paired - else: - assert not paired - - def test_narrow_to_wide_pairing_follows_covering_tree(self): - # Sender grid 2, receiver grid 4: receiver g pairs with sender g // 2. - for sender_rank in range(2): - for gen_cp_rank in range(4): - sender_ri = _k3_rank_info(tp_size=2, tp_rank=sender_rank) - peer_ri = _k3_rank_info(cp_size=4, cp_rank=gen_cp_rank) - paired = MambaPolicy.is_paired(sender_ri, peer_ri) - if gen_cp_rank // 2 == sender_rank: - assert paired, (sender_rank, gen_cp_rank) - else: - assert not paired, (sender_rank, gen_cp_rank) - - def test_is_paired_unpaired_vs_paired(self): - # Unpaired: sender rank 0 is not paired with receiver cp_rank 1 - sender_ri = _k3_rank_info(tp_size=2, tp_rank=0) - peer_ri = _k3_rank_info(cp_size=2, cp_rank=1) - assert not MambaPolicy.is_paired(sender_ri, peer_ri) - # Paired sender - paired_ri = _k3_rank_info(tp_size=2, tp_rank=1) - assert MambaPolicy.is_paired(paired_ri, peer_ri) - - -# --------------------------------------------------------------------------- # -# SlotAllocator — first-fit reuses out-of-order freed holes (real Buffer mocked) -# --------------------------------------------------------------------------- # -@pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.transport import needs CUDA bindings") -class TestSlotAllocator: - def _alloc(self, monkeypatch, cap): - # SlotAllocator.__init__ resolves Buffer from the buffer module's globals, so patch it there. - monkeypatch.setattr( - bbuf, - "Buffer", - lambda *a, **k: SimpleNamespace( - size=cap, base_ptr=0x1000, reg_descs=lambda: [], close=lambda: None - ), - ) - return btr.SlotAllocator(cap, 512) - - def test_reuses_out_of_order_freed_hole(self, monkeypatch): - a = self._alloc(monkeypatch, cap=1024) # two 512-byte slots - _sa, _ = a.reserve(512) # [0, 512) stays live - sb, _ = a.reserve(512) # [512, 1024) - a.release(sb) # free the TAIL hole out of order - # A bump cursor sitting at 1024 would wrap to 0, hit the live [0,512) and block; first-fit - # must reuse [512, 1024) immediately. - res = a.reserve(512, timeout=0.5) - assert res is not None - _, addr = res - assert addr == 0x1000 + 512 - - def test_blocks_until_release_when_full(self, monkeypatch): - a = self._alloc(monkeypatch, cap=512) # single slot - a.reserve(512) - assert a.reserve(512, timeout=0.05) is None # full -> times out - - def test_zero_timeout_does_not_wait_when_full(self, monkeypatch): - a = self._alloc(monkeypatch, cap=512) # single slot - a.reserve(512) - monkeypatch.setattr(a._cv, "wait", lambda *args, **kwargs: pytest.fail("reserve waited")) - assert a.reserve(512, timeout=0.0) is None - - -# --------------------------------------------------------------------------- # -# core.TransferContext — the pure drain-before-release state machine. -# No CUDA / NIXL / allocator, so these run on any CPU (no skipif). -# --------------------------------------------------------------------------- # -class TestLifecycle: - def _ctx(self, num_writers, per_writer_bytes=100, base_addr=0x1000): - return bcore.TransferContext( - rid_slice=(1, 0), - slot_id=0, - base_addr=base_addr, - per_writer_bytes=per_writer_bytes, - num_writers=num_writers, - ) - - def _dst(self, v=10): - return dict(dst_ptrs=np.array([v], dtype=np.int64), sizes=np.array([8], dtype=np.int64)) - - def test_writer_base_layout(self): - c = self._ctx(3, per_writer_bytes=0x64, base_addr=0x1000) - assert [c.writer_base(i) for i in range(3)] == [0x1000, 0x1064, 0x10C8] - - def test_single_writer_success_scatters_then_releases(self): - c = self._ctx(1) - assert not c.ready_to_scatter() and not c.ready_to_settle() - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) - assert c.ready_to_scatter() - c.begin_scatter() - assert not c.ready_to_settle() # scatter not landed yet - c.finish_scatter(True) - assert c.ready_to_settle() - ret = c.settle() - assert ret.disposition is bcore.Disposition.RELEASE and ret.success is True - assert c.state is bcore.TransferState.COMPLETED - assert c.settle() is None # idempotent one-shot - - def test_fanin_holds_until_all_terminal(self): - c = self._ctx(2) - c.record_writer_result(7, succeeded=True, src_base=0, **self._dst()) - assert not c.ready_to_scatter() # 1/2 writers - assert not c.ready_to_settle() # drain-before-release - c.record_writer_result(3, succeeded=True, src_base=100, **self._dst()) - assert c.ready_to_scatter() # all success -> scatter - - def test_fanin_failed_then_success_releases(self): - c = self._ctx(2) - c.record_writer_result(7, succeeded=False) - assert not c.ready_to_settle() # a sibling is still pending -> hold - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) - assert not c.ready_to_scatter() # >=1 FAILED -> skip scatter - assert c.ready_to_settle() - ret = c.settle() - assert ret.disposition is bcore.Disposition.RELEASE # all drained -> free, not quarantine - assert ret.success is False - assert c.state is bcore.TransferState.FAILED - - def test_partial_publication_waits_for_published_writers_without_scatter(self): - c = self._ctx(2) - c.record_writer_result(7, succeeded=True, src_base=0, **self._dst()) - - c.abort_publication(published_writers={7}) - - assert not c.ready_to_scatter() - assert c.ready_to_settle() - ret = c.settle() - assert ret.disposition is bcore.Disposition.RELEASE - assert ret.success is False - assert c.state is bcore.TransferState.FAILED - - def test_failed_publication_with_no_published_writer_settles_immediately(self): - c = self._ctx(2) - - c.abort_publication(published_writers=set()) - - assert c.ready_to_settle() - ret = c.settle() - assert ret.disposition is bcore.Disposition.RELEASE - assert ret.success is False - - def test_orphan_quarantines(self): - c = self._ctx(2) - c.record_writer_result(7, succeeded=True, src_base=0, **self._dst()) - c.mark_orphaned() # the other writer is in-doubt - assert c.ready_to_settle() - ret = c.settle() - assert ret.disposition is bcore.Disposition.QUARANTINE and ret.success is False - assert c.state is bcore.TransferState.QUARANTINED - - def test_empty_tail_success_releases_without_scatter(self): - c = self._ctx(1) - c.record_writer_result(3, succeeded=True) # no dst tail - assert not c.ready_to_scatter() - assert c.ready_to_settle() - assert c.settle().disposition is bcore.Disposition.RELEASE - - def test_writers_locked_after_scatter_drops_late_writer(self): - c = self._ctx(1) - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) - c.begin_scatter() # SCATTERING -> frozen - c.record_writer_result(9, succeeded=False) # a late / reordered report - assert 9 not in c._writer_ok # dropped, cannot re-arm the state - - def test_duplicate_writer_dedup(self): - c = self._ctx(2) - c.record_writer_result(7, succeeded=True, src_base=0, **self._dst()) - c.record_writer_result(7, succeeded=False) # same rank again -> ignored - assert c._writer_ok[7] is True - assert not c.ready_to_settle() # still only 1 distinct writer of 2 - - def test_scatter_failure_releases_as_failed(self): - c = self._ctx(1) - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) - c.begin_scatter() - c.finish_scatter(False) # scatter kernel failed - ret = c.settle() - assert ret.disposition is bcore.Disposition.RELEASE and ret.success is False - - def test_orphan_after_scatter_is_ignored(self): - # once SCATTERING, all writers already reported SUCCESS -> nothing is in doubt, so a late - # orphan (e.g. a racing cancel) must NOT downgrade a clean transfer to quarantine. - c = self._ctx(1) - c.record_writer_result(3, succeeded=True, src_base=0, **self._dst()) - c.begin_scatter() - c.mark_orphaned() # no-op after SCATTERING - c.finish_scatter(True) - assert c.settle().disposition is bcore.Disposition.RELEASE - - -# --------------------------------------------------------------------------- # -# SlotAllocator quarantine — orphaned regions are held out of reuse for a grace -# period, then reclaimed off a timer (never handed out while a stray write could land) -# --------------------------------------------------------------------------- # -@pytest.mark.skipif(not _HAVE_TRANSPORT, reason="bounce.buffer import needs CUDA bindings") -class TestQuarantine: - def _alloc(self, monkeypatch, cap): - monkeypatch.setattr( - bbuf, - "Buffer", - lambda *a, **k: SimpleNamespace( - size=cap, base_ptr=0x1000, reg_descs=lambda: [], close=lambda: None - ), - ) - return btr.SlotAllocator(cap, 512) - - def test_quarantined_region_is_not_reused(self, monkeypatch): - a = self._alloc(monkeypatch, cap=1024) # two 512-byte slots - a.reserve(512) # [0, 512) stays live - sb, _ = a.reserve(512) # [512, 1024) - a.quarantine(sb, grace_s=float("inf")) # hold the tail out of reuse - assert a.quarantined_bytes == 512 - # live [0,512) + quarantined [512,1024) => arena full => next reserve can't fit. - assert a.reserve(512, timeout=0.05) is None - - def test_reclaim_returns_expired_to_free(self, monkeypatch): - a = self._alloc(monkeypatch, cap=512) # single slot - s, _ = a.reserve(512) - a.quarantine(s, grace_s=0.0) # deadline already in the past - assert a.reserve(512, timeout=0.05) is None # still held until reclaimed - assert a.reclaim_expired() == 1 - assert a.quarantined_bytes == 0 - assert a.reserve(512, timeout=0.5) is not None # now reusable - - def test_inf_grace_never_reclaims(self, monkeypatch): - a = self._alloc(monkeypatch, cap=512) - s, _ = a.reserve(512) - a.quarantine(s, grace_s=float("inf")) # close-only reclaim - assert a.reclaim_expired() == 0 - assert a.quarantined_bytes == 512 diff --git a/tests/unittest/disaggregated/test_chunked_transfer.py b/tests/unittest/disaggregated/test_chunked_transfer.py index 64315d468f70..39f832f6a520 100644 --- a/tests/unittest/disaggregated/test_chunked_transfer.py +++ b/tests/unittest/disaggregated/test_chunked_transfer.py @@ -325,7 +325,6 @@ def test_process_kv_agent_result_resolves_task_by_receiver_slice_id(): """The receiver resolves the task using its own slice id.""" session = _make_rx_session(1) session._kv_tasks[0].expected_transfers = 1 - session._receiver._bounce.is_bounced.return_value = False session.process_kv_agent_result( peer_rank=0, diff --git a/tests/unittest/disaggregated/test_kda_mamba_transfer.py b/tests/unittest/disaggregated/test_kda_mamba_transfer.py index 5606fb1a558b..484ba9d10c8e 100644 --- a/tests/unittest/disaggregated/test_kda_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_kda_mamba_transfer.py @@ -61,10 +61,7 @@ import tensorrt_llm.bindings import tensorrt_llm.tensorrt_llm_transfer_agent_binding # noqa: F401 from tensorrt_llm import DisaggregatedParams, Mapping, SamplingParams -from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( - MambaPolicy, - mamba_receiver_payload_bytes, -) +from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import MambaPolicy from tensorrt_llm._torch.disaggregation.native.peer import PeerRegistrar from tensorrt_llm._torch.disaggregation.native.rank_info import RankInfo from tensorrt_llm._torch.disaggregation.resource.kv_extractor import ( @@ -249,15 +246,6 @@ def test_kda_layer_group_descriptors(enable_attention_dp): assert sorted(ll.global_layer_id for ll in mlg.local_layers) == [ i for i, m in enumerate(_KDA_MASK) if m ] - - # Matched-parallelism payload: receiver_payload_bytes must equal - # NUM_KDA_LAYERS * (conv + ssm) for matched TP. - payload = mamba_receiver_payload_bytes( - sender_page_table=pt, - receiver_page_table=pt, - dst_slot=1, - ) - assert payload == NUM_KDA_LAYERS * (CONV_SLOT_BYTES + SSM_SLOT_BYTES) finally: mgr.shutdown() diff --git a/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py b/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py index 6ad0e67290ba..35a73b4cad1d 100644 --- a/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py +++ b/tests/unittest/disaggregated/test_kv_transfer_fail_stop.py @@ -520,7 +520,6 @@ def _shutdown_transceiver() -> tuple[KvCacheTransceiverV2, transfer_mod.RxSessio receiver._shutdown = False receiver._ownership_admission_lock = threading.Lock() receiver._ownership_poisoned = None - receiver._bounce = Mock() receiver._dealers = {} receiver._messenger = Mock() receiver.send_cancel_to_senders = Mock(return_value=None) @@ -539,7 +538,6 @@ def _shutdown_transceiver() -> tuple[KvCacheTransceiverV2, transfer_mod.RxSessio worker = transfer_mod.TransferWorker.__new__(transfer_mod.TransferWorker) worker._retirement_watchdog = watchdog worker._receiver = receiver - worker._bounce = receiver._bounce worker._registered_mem = [] instance = KvCacheTransceiverV2.__new__(KvCacheTransceiverV2) instance._transfer_worker = worker @@ -624,8 +622,8 @@ def dispatch() -> None: thread.join(timeout=1) assert not thread.is_alive() assert not errors - task.record_writer_result(7, True, wait_for_local_completion=False) - session._aux_physical_owner.record_writer_result(7, True, wait_for_local_completion=False) + task.record_writer_result(7, True) + session._aux_physical_owner.record_writer_result(7, True) instance.shutdown() assert session.status is SessionStatus.CANCELLED aux.free_slot.assert_called_once_with(3) diff --git a/tests/unittest/disaggregated/test_mamba_transfer.py b/tests/unittest/disaggregated/test_mamba_transfer.py index f31a4ca7a17c..dfc0e927c32c 100644 --- a/tests/unittest/disaggregated/test_mamba_transfer.py +++ b/tests/unittest/disaggregated/test_mamba_transfer.py @@ -299,94 +299,6 @@ def _zero_mamba_states(manager): state.zero_() -def test_mamba_receiver_payload_bytes_matched_tp(): - """mamba_receiver_payload_bytes returns correct size for matched-TP page tables.""" - from tensorrt_llm._torch.disaggregation.native.mixers.ssm.peer import ( - mamba_receiver_payload_bytes, - ) - from tensorrt_llm._torch.disaggregation.resource.page import ( - BUFFER_ENTRY_DTYPE, - MAMBA_CONV_ROLE, - MAMBA_SSM_ROLE, - KVCachePageTable, - LocalLayer, - MambaLayerGroup, - MapperKind, - PhysicalPool, - PhysicalPoolGroup, - PoolView, - ) - - conv_slot_bytes = 1024 - ssm_slot_bytes = 2048 - side_slot_bytes = 128 - sorted_lids = [0, 1] - - conv_pool = PhysicalPool(base_address=0xA000, slot_bytes=conv_slot_bytes, num_slots=8) - ssm_pool = PhysicalPool(base_address=0xB000, slot_bytes=ssm_slot_bytes, num_slots=8) - pool_views = [ - PoolView( - pool_idx=0, - buffer_entries=np.array( - [(lid, 0, conv_slot_bytes) for lid in sorted_lids], dtype=BUFFER_ENTRY_DTYPE - ), - pool_role=MAMBA_CONV_ROLE, - mapper_kind=MapperKind.SECTIONED, - bytes_per_layer=conv_slot_bytes, - section_bytes=[512, 256, 256], - ), - PoolView( - pool_idx=1, - buffer_entries=np.array( - [(lid, 0, ssm_slot_bytes) for lid in sorted_lids], dtype=BUFFER_ENTRY_DTYPE - ), - pool_role=MAMBA_SSM_ROLE, - mapper_kind=MapperKind.INDEXED, - bytes_per_layer=ssm_slot_bytes, - bytes_per_head=64, - ), - PoolView( - pool_idx=2, - buffer_entries=np.array([(1, 0, side_slot_bytes)], dtype=BUFFER_ENTRY_DTYPE), - pool_role=frozenset({"ple_ngram_context"}), - mapper_kind=MapperKind.REPLICATED, - bytes_per_layer=side_slot_bytes, - ), - ] - local_layers = [ - LocalLayer(local_layer_id=0, global_layer_id=1), - LocalLayer(local_layer_id=1, global_layer_id=2), - ] - - mlg = MambaLayerGroup( - pool_group_idx=1, - local_layers=local_layers, - pool_views=pool_views, - ) - pt = KVCachePageTable( - tokens_per_block=8, - layer_groups=[mlg], - pool_groups=[ - PhysicalPoolGroup(pools=[]), # idx 0 placeholder - PhysicalPoolGroup( - pools=[ - conv_pool, - ssm_pool, - PhysicalPool( - base_address=0xC000, - slot_bytes=side_slot_bytes, - num_slots=8, - ), - ] - ), - ], - ) - - got = mamba_receiver_payload_bytes(sender_page_table=pt, receiver_page_table=pt, dst_slot=3) - # Standard state covers both layers; the side role exists only on layer 2. - assert got == 2 * (conv_slot_bytes + ssm_slot_bytes) + side_slot_bytes - - def _make_rank_info(*, tp_size: int, tp_rank: int = 0) -> RankInfo: return RankInfo( instance_name="test", diff --git a/tests/unittest/disaggregated/test_rank_info.py b/tests/unittest/disaggregated/test_rank_info.py index a1d25e01c3f3..3a9652450918 100644 --- a/tests/unittest/disaggregated/test_rank_info.py +++ b/tests/unittest/disaggregated/test_rank_info.py @@ -203,3 +203,57 @@ def test_rank_info_represents_cache_element_bytes( restored = RankInfo.from_bytes(rank_info.to_bytes()) assert restored.attention.element_bytes == expected_element_bytes assert isinstance(restored.attention.element_bytes, expected_type) + + +def _mamba_rank_info(tp_size=1, tp_rank=0, cp_size=1, cp_rank=0): + # MambaPolicy reads tp_size / tp_rank / cp_size / cp_rank / + # attention.enable_attention_dp. + return SimpleNamespace( + tp_size=tp_size, tp_rank=tp_rank, cp_size=cp_size, cp_rank=cp_rank, attention=None + ) + + +class TestMambaHelixCP: + """MambaPolicy under decode-CP (helix): shard grid and pairing.""" + + def test_mamba_tp_helix_grid_is_cp_minor(self): + # cp == 1 keeps the plain TP grid; cp > 1 flattens tp*cp CP-minor. + assert MambaPolicy._mamba_tp(_mamba_rank_info()) == (1, 0) + assert MambaPolicy._mamba_tp(_mamba_rank_info(cp_size=32, cp_rank=5)) == (32, 5) + ri = _mamba_rank_info(tp_size=2, tp_rank=1, cp_size=4, cp_rank=3) + assert MambaPolicy._mamba_tp(ri) == (8, 7) + + def test_equal_width_pairs_only_matching_rank(self): + # Regression lock for the fan-in collapse: with equal shard widths the + # mapper is a whole-slot copy, so only the same-index sender may pair; + # any other sender would overwrite the receiver slot with the wrong heads. + for sender_rank in range(2): + for gen_cp_rank in range(2): + sender_ri = _mamba_rank_info(tp_size=2, tp_rank=sender_rank) + peer_ri = _mamba_rank_info(cp_size=2, cp_rank=gen_cp_rank) + paired = MambaPolicy.is_paired(sender_ri, peer_ri) + if sender_rank == gen_cp_rank: + assert paired + else: + assert not paired + + def test_narrow_to_wide_pairing_follows_covering_tree(self): + # Sender grid 2, receiver grid 4: receiver g pairs with sender g // 2. + for sender_rank in range(2): + for gen_cp_rank in range(4): + sender_ri = _mamba_rank_info(tp_size=2, tp_rank=sender_rank) + peer_ri = _mamba_rank_info(cp_size=4, cp_rank=gen_cp_rank) + paired = MambaPolicy.is_paired(sender_ri, peer_ri) + if gen_cp_rank // 2 == sender_rank: + assert paired, (sender_rank, gen_cp_rank) + else: + assert not paired, (sender_rank, gen_cp_rank) + + def test_is_paired_unpaired_vs_paired(self): + # Unpaired: sender rank 0 is not paired with receiver cp_rank 1 + sender_ri = _mamba_rank_info(tp_size=2, tp_rank=0) + peer_ri = _mamba_rank_info(cp_size=2, cp_rank=1) + assert not MambaPolicy.is_paired(sender_ri, peer_ri) + # Paired sender + paired_ri = _mamba_rank_info(tp_size=2, tp_rank=1) + assert MambaPolicy.is_paired(paired_ri, peer_ri) diff --git a/tests/unittest/disaggregated/test_retirement_deadline.py b/tests/unittest/disaggregated/test_retirement_deadline.py index ce530f83c5cc..ccd552b4f8e6 100644 --- a/tests/unittest/disaggregated/test_retirement_deadline.py +++ b/tests/unittest/disaggregated/test_retirement_deadline.py @@ -13,7 +13,6 @@ from tensorrt_llm._torch.disaggregation.base import Chunk, TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus from tensorrt_llm._torch.disaggregation.native import transfer as transfer_mod -from tensorrt_llm._torch.disaggregation.native.bounce import NoBounceTransport from tensorrt_llm._torch.disaggregation.native.retirement import RetirementWatchdog from tensorrt_llm.disaggregated_params import DisaggregatedParams, DisaggScheduleStyle @@ -350,7 +349,6 @@ def _receiver(watchdog: RetirementWatchdog) -> transfer_mod.Receiver: receiver._shutdown = False receiver._ownership_admission_lock = threading.Lock() receiver._ownership_poisoned = None - receiver._bounce = Mock() return receiver @@ -501,7 +499,6 @@ def test_final_completion_precedes_delayed_perf_logging(final_piece: str, monkey watchdog = RetirementWatchdog(Mock(), clock=clock) if final_piece == "receive_kv": receiver = _receiver(watchdog) - receiver._bounce = NoBounceTransport() receiver._registrar = SimpleNamespace( self_rank_info=SimpleNamespace(instance_name="receiver", instance_rank=0) ) @@ -526,7 +523,6 @@ def test_final_completion_precedes_delayed_perf_logging(final_piece: str, monkey task = session.send_aux() sender = session._sender sender._instance_rank = sender._device_id = 0 - sender._bounce = NoBounceTransport() sender._registrar = SimpleNamespace( self_rank_info=SimpleNamespace(instance_name="sender", instance_rank=0) ) @@ -612,7 +608,6 @@ def test_finished_nonfinal_piece_does_not_complete_session(direction: str, block _submit_piece(session, session.kv_tasks[0]) else: receiver = _receiver(watchdog) - receiver._bounce = NoBounceTransport() receiver._registrar = SimpleNamespace( self_rank_info=SimpleNamespace(instance_name="receiver", instance_rank=0) ) @@ -687,7 +682,6 @@ def test_expired_second_receive_preserves_existing_claims(need_aux: bool) -> Non clock = Mock(return_value=10.0) watchdog = RetirementWatchdog(Mock(), clock=clock) receiver = _receiver(watchdog) - receiver._bounce = NoBounceTransport() receiver._registrar = SimpleNamespace( self_rank_info=SimpleNamespace(instance_name="receiver", instance_rank=0) ) @@ -727,7 +721,6 @@ def _settled_rx_session( ) -> tuple[transfer_mod.RxSession, transfer_mod.KVRecvTask, Mock]: """Settle real KV and AUX result handlers before the executor closes the session.""" receiver = _receiver(watchdog) - receiver._bounce = NoBounceTransport() receiver._registrar = SimpleNamespace( self_rank_info=SimpleNamespace(instance_name="receiver", instance_rank=0) ) @@ -917,9 +910,9 @@ def seal(expected: int, cohort: set[int], *, published_writers: set[int] | None assert published == {7} assert session.status is SessionStatus.CANCELLED assert not session.resources_drained() - task.record_writer_result(7, False, wait_for_local_completion=False) + task.record_writer_result(7, False) assert not session.resources_drained() - session._aux_physical_owner.record_writer_result(7, False, wait_for_local_completion=False) + session._aux_physical_owner.record_writer_result(7, False) assert session.close() watchdog.stop() @@ -1013,7 +1006,6 @@ def test_receiver_deadline_retains_kv_aux_and_registration(trigger: str) -> None worker._receiver = receiver worker._agent = Mock() worker._registered_mem = [object()] - worker._bounce = Mock() if trigger == "cancel": session.cancel_local() elif trigger == "failure": @@ -1029,8 +1021,8 @@ def test_receiver_deadline_retains_kv_aux_and_registration(trigger: str) -> None clock.return_value = 20.0 if trigger == "timeout" else 15.0 watchdog.progress() assert watchdog.fatal is not None - task.record_writer_result(7, False, wait_for_local_completion=False) - session._aux_physical_owner.record_writer_result(7, False, wait_for_local_completion=False) + task.record_writer_result(7, False) + session._aux_physical_owner.record_writer_result(7, False) assert not session.resources_drained() assert not session.close() assert receiver._sessions[4] is session @@ -1038,7 +1030,6 @@ def test_receiver_deadline_retains_kv_aux_and_registration(trigger: str) -> None worker.shutdown() worker._agent.deregister_memory.assert_not_called() worker._agent.shutdown.assert_not_called() - worker._bounce.close.assert_not_called() assert len(worker._registered_mem) == 1 with pytest.raises(RuntimeError, match="admission is closed"): _rx_session(receiver) @@ -1169,7 +1160,7 @@ def complete_during_grace(*, timeout: float) -> bool: """Inject late backend evidence at the external wait boundary.""" assert timeout == 5.0 clock.return_value = 15.0 - task.record_writer_result(7, True, wait_for_local_completion=False) + task.record_writer_result(7, True) task.complete() return True diff --git a/tests/unittest/disaggregated/test_task_handle.py b/tests/unittest/disaggregated/test_task_handle.py index 43f79367cca1..24cb56ab86b1 100644 --- a/tests/unittest/disaggregated/test_task_handle.py +++ b/tests/unittest/disaggregated/test_task_handle.py @@ -82,7 +82,6 @@ def _stub_receiver(): receiver = MagicMock() receiver._enforce_physical_ownership = False receiver._ownership_poisoned = None - receiver._bounce.is_bounced.return_value = False return receiver @@ -338,11 +337,11 @@ def test_an_owned_task_answers_from_the_cohort_it_sealed(): ) handle = TaskHandle(session, task, TOKENS) - task.record_writer_result(0, False, wait_for_local_completion=False) + task.record_writer_result(0, False) task.fail(RuntimeError("first writer failed")) assert handle.poll().reports_pending is True - task.record_writer_result(1, True, wait_for_local_completion=False) + task.record_writer_result(1, True) assert _owes_a_report(handle.poll()) is False @@ -511,7 +510,6 @@ def _wired_receiver(owns_transfers: bool = False) -> Receiver: receiver._shutdown = False receiver._ownership_admission_lock = threading.Lock() receiver._ownership_poisoned = None - receiver._bounce = MagicMock() receiver._dealers = {} receiver._messenger = MagicMock() receiver.dispatch_task = MagicMock() @@ -880,63 +878,3 @@ def test_delivered_sender_piece_survives_a_direct_sibling_failure_without_pollin assert isinstance(TaskHandle(session, delivered, TOKENS).poll(), Delivered) assert isinstance(TaskHandle(session, failed, TOKENS).poll(), Failed) - - -@pytest.mark.parametrize("owns_transfers", [False, True]) -@pytest.mark.parametrize("scatter_succeeded", [False, True]) -@pytest.mark.parametrize("cancel_first", [False, True]) -def test_scatter_callback_and_cancel_commit_in_event_order( - monkeypatch: pytest.MonkeyPatch, - owns_transfers: bool, - scatter_succeeded: bool, - cancel_first: bool, -) -> None: - from tensorrt_llm._torch.disaggregation.native import bounce - - session, handle = _receiving_from(1, owns_transfers=owns_transfers) - deferred = [] - monkeypatch.setattr(bounce, "scatter_write_result", lambda *args: deferred.append(args[-1])) - _report(session, 0, AgentResult.SUCCESS) - assert len(deferred) == 1 - assert handle.poll() is None - if owns_transfers: - assert session.resources_drained() is False - - ready, resume, finished = threading.Event(), threading.Event(), threading.Event() - - def finish_scatter() -> None: - ready.set() - if resume.wait(timeout=5): - deferred[0](scatter_succeeded) - finished.set() - - worker = threading.Thread(target=finish_scatter) - worker.start() - try: - assert ready.wait(timeout=5) - if cancel_first: - session.cancel_local(by_peer=True) - assert isinstance(handle.poll(), Cancelled) - if owns_transfers: - assert session.resources_drained() is False - resume.set() - assert finished.wait(timeout=5) - if not cancel_first: - session.cancel_local(by_peer=True) - finally: - resume.set() - worker.join(timeout=5) - assert not worker.is_alive() - - for observer in (handle, TaskHandle(session, session._kv_tasks[0], TOKENS)): - outcome = observer.poll() - if cancel_first: - assert isinstance(outcome, Cancelled) - assert outcome.by_peer is True - elif scatter_succeeded: - assert isinstance(outcome, Delivered) - else: - assert isinstance(outcome, Failed) - assert "bounce scatter failed" in outcome.reason - if owns_transfers: - assert session.resources_drained() is True diff --git a/tests/unittest/disaggregated/test_transfer_late_settlement.py b/tests/unittest/disaggregated/test_transfer_late_settlement.py index 4dc9c93b1f40..61eb6d23130c 100644 --- a/tests/unittest/disaggregated/test_transfer_late_settlement.py +++ b/tests/unittest/disaggregated/test_transfer_late_settlement.py @@ -67,7 +67,6 @@ def _sender() -> transfer_mod.Sender: sender._device_id = 0 sender._num_threads = 1 sender._pending_settlements = [{}] - sender._bounce = Mock() return sender @@ -159,14 +158,14 @@ def test_only_explicit_settlement_closes_ambiguous_writer(count_only: bool) -> N owner.seal_writer_cohort(2, None if count_only else {7, 8}) owner.finish_publication() assert owner.record_writer_in_doubt(7) - assert owner.record_writer_result(7, False, wait_for_local_completion=False) == (False, False) + assert owner.record_writer_result(7, False) == (False, False) assert not owner.all_writers_reported assert not owner.resources_drained assert owner.record_writer_settlement(7) assert not owner.resources_drained assert not owner.record_writer_settlement(7) assert not owner.record_writer_in_doubt(7) - owner.record_writer_result(8, False, wait_for_local_completion=False) + owner.record_writer_result(8, False) assert owner.resources_drained @@ -186,7 +185,7 @@ def test_settlement_does_not_clear_contradictory_evidence() -> None: owner = _owner({7}) owner.record_writer_in_doubt(7) with pytest.raises(RuntimeError, match="success while in doubt"): - owner.record_writer_result(7, True, wait_for_local_completion=True) + owner.record_writer_result(7, True) assert owner.record_writer_settlement(7) assert not owner.resources_drained @@ -198,20 +197,17 @@ def test_abort_publication_cannot_exclude_ambiguous_writer() -> None: with pytest.raises(RuntimeError, match="unpublished writer"): owner.abort_publication({8}) assert owner.record_writer_settlement(7) - owner.record_writer_result(8, False, wait_for_local_completion=False) + owner.record_writer_result(8, False) assert not owner.resources_drained @pytest.mark.cpu_only -def test_settlement_does_not_clear_local_completion_or_publication() -> None: +def test_settlement_does_not_clear_publication() -> None: owner = _owner({7}) - owner._local_completion_pending = True owner._publication_pending = True owner.record_writer_in_doubt(7) owner.record_writer_settlement(7) assert not owner.resources_drained - owner.finish_local_completion() - assert not owner.resources_drained owner.finish_publication() assert owner.resources_drained @@ -292,7 +288,7 @@ def test_sender_late_done_reports_physical_failure_only( assert task._transfer_count == 1 session._claim_aux_terminal_result.assert_called_once_with(7) else: - decoded = [transfer_mod._KV_RESULT_PREFIX.unpack(message[1]) for message in messages] + decoded = [transfer_mod._KV_RESULT_FRAME.unpack(message[1]) for message in messages] assert [transfer_mod._AGENT_RESULT_BY_CODE[item[4]] for item in decoded] == [ transfer_mod.AgentResult.IN_DOUBT, transfer_mod.AgentResult.FAILED_QUIESCED, @@ -311,8 +307,6 @@ def test_receiver_late_settlement_retains_logical_failure_and_quarantine(auxilia receiver._shutdown = False receiver._ownership_admission_lock = threading.Lock() receiver._ownership_poisoned = None - receiver._bounce = Mock() - receiver._bounce.is_bounced.return_value = False session = transfer_mod.RxSession(request_id=401, params=_params(), receiver=receiver) task = session.prepare_receive(_chunk()) assert task is not None @@ -341,12 +335,11 @@ def report(status: transfer_mod.AgentResult) -> None: assert session.status is SessionStatus.ERROR assert session.exception is error assert receiver._ownership_poisoned is not None - receiver._bounce.record_failure.assert_called_once_with((401, 0), 7) assert session.close() @pytest.mark.cpu_only -def test_settlement_retries_in_order_without_releasing_twice() -> None: +def test_settlement_retries_in_order() -> None: sender = _sender() task = transfer_mod.KVSendTask(_chunk(), _params(), slice_id=0) status = Mock(is_completed=Mock(return_value=True)) @@ -371,21 +364,19 @@ def test_settlement_retries_in_order_without_releasing_twice() -> None: dealer = Mock() dealer.send.side_effect = [RuntimeError("initial send"), None, RuntimeError("late send"), None] sender._get_result_dealer = Mock(return_value=dealer) - sender._retain_in_doubt_transfer(meta, initial, send_slot_id=12) + sender._retain_in_doubt_transfer(meta, initial) status.is_completed.assert_not_called() sender._poll_in_doubt_transfers(0) assert task.resources_drained assert sender._pending_settlements[0] - sender._bounce.release_send.assert_called_once_with(12) with pytest.raises(RuntimeError, match="settlement reports are pending"): sender.shutdown() sender._poll_in_doubt_transfers(0) assert not sender._pending_settlements[0] - sender._bounce.release_send.assert_called_once_with(12) status.is_completed.assert_called_once() assert task.transferred_count == 1 codes = [ - transfer_mod._KV_RESULT_PREFIX.unpack(call.args[0][1])[4] + transfer_mod._KV_RESULT_FRAME.unpack(call.args[0][1])[4] for call in dealer.send.call_args_list ] assert codes == [2, 2, 3, 3] @@ -418,7 +409,6 @@ def test_committed_session_outcome_survives_late_settlement( receiver._enforce_physical_ownership = True receiver._ownership_poisoned = None receiver._get_ownership_admission_lock.return_value = threading.Lock() - receiver._bounce.is_bounced.return_value = False session = transfer_mod.RxSession(request_id=401, params=_params(), receiver=receiver) task = session.prepare_receive(_chunk()) assert task is not None @@ -522,3 +512,32 @@ def send(_message: list[bytes]) -> None: assert status.is_completed.call_count == 2 assert task.status is transfer_mod.TaskStatus.ERROR sender._shutdown = True + + +@pytest.mark.cpu_only +def test_kv_result_frame_roundtrip() -> None: + """The KV_AGENT_RESULT binary frame must round-trip exactly.""" + for rank, rid, sl, last, status, size in [ + (7, 6925227277844486, 42, True, transfer_mod.AgentResult.SUCCESS, 4096), + (0, 1, 0, False, transfer_mod.AgentResult.FAILED, 0), + (31, 2**62, 9999, True, transfer_mod.AgentResult.SUCCESS, 2**40), + ]: + packed = transfer_mod._KV_RESULT_FRAME.pack( + rank, rid, sl, last, transfer_mod._AGENT_RESULT_CODE[status], size + ) + r, i, s, last_out, c, sz = transfer_mod._KV_RESULT_FRAME.unpack(packed) + assert (r, i, s, last_out, sz) == (rank, rid, sl, last, size) + assert transfer_mod._AGENT_RESULT_BY_CODE[c] is status + + +@pytest.mark.cpu_only +@pytest.mark.parametrize("result_name", ["SUCCESS", "FAILED"]) +def test_make_kv_result_msg_uses_binary_frame(result_name: str) -> None: + """Every KV result (success and failure) uses the binary frame so the receiver can decode it.""" + result = getattr(transfer_mod.AgentResult, result_name) + msg = transfer_mod._make_kv_result_msg(3, 12345, 7, True, result, transfer_size=8192) + assert msg[0] == transfer_mod.MessageType.KV_AGENT_RESULT + assert len(msg) == 2 # type + frame + r, rid, sl, last, code, size = transfer_mod._KV_RESULT_FRAME.unpack(msg[1]) + assert (r, rid, sl, last, size) == (3, 12345, 7, True, 8192) + assert transfer_mod._AGENT_RESULT_BY_CODE[code] is result diff --git a/tests/unittest/disaggregated/test_transfer_lifecycle_acceptance.py b/tests/unittest/disaggregated/test_transfer_lifecycle_acceptance.py index a60d8f778a5a..8935ecca78f5 100644 --- a/tests/unittest/disaggregated/test_transfer_lifecycle_acceptance.py +++ b/tests/unittest/disaggregated/test_transfer_lifecycle_acceptance.py @@ -28,7 +28,6 @@ import tensorrt_llm._torch.disaggregation.transceiver as transceiver_mod from tensorrt_llm._torch.disaggregation.base import Cancelled, Chunk, TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus -from tensorrt_llm._torch.disaggregation.native.bounce.impl import NoBounceTransport from tensorrt_llm._torch.disaggregation.native.handle import TaskHandle from tensorrt_llm._torch.disaggregation.native.transfer import AgentResult, KVRecvTask, RxSession from tensorrt_llm._torch.disaggregation.orchestration import coordinator as coordinator_mod @@ -133,7 +132,6 @@ def __init__(self, monkeypatch: pytest.MonkeyPatch, *, bridge_enabled: bool) -> self.cleanup_count = 0 self.receiver = SimpleNamespace( _enforce_physical_ownership=ownership, - _bounce=NoBounceTransport(), _registrar=SimpleNamespace( self_rank_info=SimpleNamespace(instance_name="gen", instance_rank=0) ), diff --git a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py index cfa6213af139..81ebf9629174 100644 --- a/tests/unittest/disaggregated/test_transfer_ownership_regressions.py +++ b/tests/unittest/disaggregated/test_transfer_ownership_regressions.py @@ -30,7 +30,6 @@ from tensorrt_llm import DisaggregatedParams from tensorrt_llm._torch.disaggregation.base import CacheExtent, CacheKind, Chunk, TokenRange from tensorrt_llm._torch.disaggregation.base.transfer import SessionStatus, WaitResult -from tensorrt_llm._torch.disaggregation.native.bounce.core import TransferContext from tensorrt_llm._torch.disaggregation.native.transfer import ( AgentResult, KVRecvTask, @@ -56,123 +55,6 @@ def _sole_piece(block_ids_per_layer_groups=None, *, tokens=0, is_last=True) -> C ) -class _BounceProbe: - """No-bounce probe that records physical-owner cleanup decisions.""" - - def __init__(self) -> None: - self.failed_writers: list[tuple[tuple[int, int], int]] = [] - self.orphaned: list[tuple[int, int]] = [] - - def record_failure(self, rid_slice: tuple[int, int], peer_rank: int) -> None: - self.failed_writers.append((rid_slice, peer_rank)) - - def reserve(self, _receiver_req, _num_writers: int, *, extra_bytes: int = 0) -> bool: - del extra_bytes - return False - - def release_idle_reservation(self, _rid_slice: tuple[int, int]) -> None: - return - - def orphan_reservation(self, rid_slice: tuple[int, int]) -> None: - self.orphaned.append(rid_slice) - - def abort_publication( - self, - _rid_slice: tuple[int, int], - _published_writers: set[int], - ) -> None: - return - - def is_bounced(self, _rid_slice: tuple[int, int]) -> bool: - return False - - -class _LateReservationBounce(_BounceProbe): - """Track a reservation created after cancellation already ran cleanup.""" - - def __init__(self) -> None: - super().__init__() - self.active_reservations: set[tuple[int, int]] = set() - - def release_idle_reservation(self, rid_slice: tuple[int, int]) -> None: - self.active_reservations.discard(rid_slice) - - -class _PartialFanInBounce(_BounceProbe): - """CPU model of one bounced fan-in reservation.""" - - def __init__(self) -> None: - super().__init__() - self.context: TransferContext | None = None - self.release_count = 0 - self.scatter_count = 0 - - def reserve(self, receiver_req, num_writers: int, *, extra_bytes: int = 0) -> bool: - del extra_bytes - self.context = TransferContext( - rid_slice=(receiver_req.unique_rid, receiver_req.slice_id), - slot_id=0, - base_addr=0x1000, - per_writer_bytes=0x100, - num_writers=num_writers, - ) - return True - - def writer_base(self, rid_slice: tuple[int, int], writer_index: int) -> int | None: - if self.context is None or self.context.rid_slice != rid_slice: - return None - return self.context.writer_base(writer_index) - - def is_bounced(self, rid_slice: tuple[int, int]) -> bool: - return self.context is not None and self.context.rid_slice == rid_slice - - def _advance(self) -> None: - if self.context is None: - return - if self.context.ready_to_scatter(): - self.scatter_count += 1 - self.context.begin_scatter() - self.context.finish_scatter(True) - if not self.context.ready_to_settle(): - return - settlement = self.context.settle() - self.context = None - assert settlement is not None - self.release_count += 1 - if settlement.on_done is not None: - settlement.on_done(settlement.success) - - def abort_publication( - self, - rid_slice: tuple[int, int], - published_writers: set[int], - ) -> None: - assert self.context is not None and self.context.rid_slice == rid_slice - self.context.abort_publication(published_writers) - self._advance() - - def record_result( - self, - rid_slice: tuple[int, int], - peer_rank: int, - dst_ptrs=None, - sizes=None, - src_base=None, - on_done=None, - ) -> None: - assert self.context is not None and self.context.rid_slice == rid_slice - if on_done is not None: - self.context.on_done = on_done - self.context.record_writer_result( - peer_rank, - succeeded=True, - src_base=src_base, - dst_ptrs=dst_ptrs, - sizes=sizes, - ) - self._advance() - - def _start_checked_thread( target: Callable[[], None], results: queue.Queue[Exception | None], @@ -240,7 +122,6 @@ class _ReceiverProbe: """Minimal receiver that publishes one destination to two writers.""" def __init__(self) -> None: - self._bounce = _BounceProbe() self._enforce_physical_ownership = True self._session: RxSession | None = None self.clear_count = 0 @@ -295,7 +176,6 @@ def _make_owned_receiver() -> Receiver: receiver._shutdown = False receiver._ownership_admission_lock = threading.Lock() receiver._ownership_poisoned = None - receiver._bounce = _BounceProbe() return receiver @@ -384,7 +264,6 @@ def test_pre_cancelled_rx_session_never_publishes_destination( receiver._sessions_lock = threading.Lock() receiver._sessions = {} receiver._pre_cancelled_rids = {rid: True} - receiver._bounce = _BounceProbe() receiver._enforce_physical_ownership = True receiver._shutdown = False receiver.dispatch_task = Mock() @@ -413,7 +292,6 @@ def test_remote_cancel_resolves_strong_owned_session() -> None: receiver._sessions_lock = threading.Lock() receiver._sessions = {} receiver._pre_cancelled_rids = {} - receiver._bounce = _BounceProbe() receiver._enforce_physical_ownership = True receiver._shutdown = False receiver.send_cancel_to_senders = Mock(return_value=None) @@ -585,13 +463,11 @@ def test_gen_first_no_retry_adp_count_seal_waits_for_one_writer_group() -> None: receiver._sessions_lock = threading.Lock() receiver._sessions = {} receiver._pre_cancelled_rids = {} - receiver._bounce = _BounceProbe() receiver._enforce_physical_ownership = True receiver._shutdown = False receiver_req = SimpleNamespace( unique_rid=rid, slice_id=0, - bounce_dst_base=None, to_bytes=Mock(return_value=b"receiver-request"), ) receiver._build_recv_req_info = Mock(return_value=receiver_req) @@ -601,19 +477,12 @@ def test_gen_first_no_retry_adp_count_seal_waits_for_one_writer_group() -> None: } receiver._registrar = SimpleNamespace( get_peer_overlap=Mock(side_effect=lambda _peer, dp_rank: overlaps[dp_rank]), - self_extractor=SimpleNamespace(page_table=None), - self_rank_info=SimpleNamespace( - cp_size=1, - instance_name="gen", - instance_rank=0, - ), + self_rank_info=SimpleNamespace(instance_name="gen", instance_rank=0), ) receiver._get_sender_info = Mock( return_value=SimpleNamespace( sender_endpoints={rank: f"tcp://sender-{rank}" for rank in range(4)}, - page_table=None, dp_size=2, - cp_size=1, ) ) receiver._request_sender_data = Mock() @@ -663,18 +532,16 @@ def test_gen_first_no_retry_adp_count_seal_waits_for_one_writer_group() -> None: @pytest.mark.cpu_only @pytest.mark.parametrize("writer_settles_before_failure", [False, True]) -def test_partial_bounced_publication_waits_for_queued_writer_success( +def test_partial_publication_waits_for_queued_writer_success( monkeypatch: pytest.MonkeyPatch, writer_settles_before_failure: bool, ) -> None: rid = 86 queued_endpoints: list[str] = [] - bounce = _PartialFanInBounce() receiver = object.__new__(Receiver) receiver._sessions_lock = threading.Lock() receiver._sessions = {} receiver._pre_cancelled_rids = {} - receiver._bounce = bounce receiver._enforce_physical_ownership = True receiver._shutdown = False # This fixture delivers one result inline from the publication callback. @@ -686,25 +553,17 @@ def test_partial_bounced_publication_waits_for_queued_writer_success( unique_rid=rid, slice_id=0, mamba_state_index=None, - bounce_dst_base=None, - to_bytes=Mock(side_effect=[b"writer-0", b"writer-1"]), + to_bytes=Mock(return_value=b"receiver-request"), ) receiver._build_recv_req_info = Mock(return_value=receiver_req) - overlap = SimpleNamespace( - ranks=[0, 1], - duplicate_head_factor=1, - overlap_pp_size=1, - ) + overlap = SimpleNamespace(ranks=[0, 1]) receiver._registrar = SimpleNamespace( get_peer_overlap=Mock(return_value=overlap), - self_extractor=SimpleNamespace(page_table=None), - self_rank_info=SimpleNamespace(instance_name="gen", instance_rank=0, cp_size=1), + self_rank_info=SimpleNamespace(instance_name="gen", instance_rank=0), ) receiver._get_sender_info = Mock( return_value=SimpleNamespace( sender_endpoints={0: "tcp://sender-0", 1: "tcp://sender-1"}, - page_table=None, - cp_size=1, dp_size=1, ) ) @@ -719,9 +578,6 @@ def request_sender_data(endpoint: str, _payload: bytes) -> None: receiver_slice_id=0, is_last_slice=True, status=AgentResult.SUCCESS, - dst_ptrs=np.array([0x2000], dtype=np.int64), - sizes=np.array([0x100], dtype=np.int64), - src_base=0x1000, ) receiver._request_sender_data = request_sender_data @@ -745,8 +601,6 @@ def request_sender_data(endpoint: str, _payload: bytes) -> None: assert queued_endpoints == ["tcp://sender-0"] if not writer_settles_before_failure: - assert bounce.context is not None - assert bounce.release_count == 0 assert not session.resources_drained() # Only writer 0's REQUEST_DATA was successfully queued. The destination @@ -756,18 +610,12 @@ def request_sender_data(endpoint: str, _payload: bytes) -> None: receiver_slice_id=0, is_last_slice=True, status=AgentResult.SUCCESS, - dst_ptrs=np.array([0x2000], dtype=np.int64), - sizes=np.array([0x100], dtype=np.int64), - src_base=0x1000, ) assert session.status == SessionStatus.ERROR assert not session.resources_drained() session.process_aux_agent_result(0, AgentResult.FAILED) assert session.resources_drained() - assert bounce.context is None - assert bounce.scatter_count == 0 - assert bounce.release_count == 1 assert session.close() is True assert rid not in receiver._sessions @@ -854,7 +702,6 @@ def test_cancel_after_publication_cannot_overtake_request_data( receiver._sessions_lock = threading.Lock() receiver._sessions = {} receiver._pre_cancelled_rids = {} - receiver._bounce = _BounceProbe() receiver._enforce_physical_ownership = True receiver._shutdown = False receiver._dealers = {} @@ -863,23 +710,16 @@ def test_cancel_after_publication_cannot_overtake_request_data( unique_rid=rid, slice_id=0, mamba_state_index=None, - bounce_dst_base=None, to_bytes=Mock(return_value=b"receiver-request"), ) ) overlap = SimpleNamespace(ranks=[0]) - receiver._registrar = SimpleNamespace( - get_peer_overlap=Mock(return_value=overlap), - self_extractor=SimpleNamespace(page_table=None), - self_rank_info=SimpleNamespace(cp_size=1), - ) + receiver._registrar = SimpleNamespace(get_peer_overlap=Mock(return_value=overlap)) receiver._get_sender_info = Mock( return_value=SimpleNamespace( sender_endpoints={0: "tcp://sender-0"}, - page_table=None, tp_size=1, pp_size=1, - cp_size=1, dp_size=1, attention=None, ) @@ -963,24 +803,18 @@ def finish_writer() -> None: @pytest.mark.cpu_only -def test_cancel_before_dispatch_releases_late_idle_reservation() -> None: +def test_cancel_before_dispatch_closes_unpublished_task() -> None: rid = 81 receiver = _ReceiverProbe() - bounce = _LateReservationBounce() - receiver._bounce = bounce session = _make_rx_session(receiver, rid) task = session.prepare_receive(_sole_piece()) assert task is not None assert session.cancel_local() - bounce.active_reservations.add((rid, task.slice_id)) with pytest.raises(RuntimeError, match="became terminal before publication"): session.dispatch_prepared_receive(task) - assert bounce.active_reservations == set(), ( - "cancel-before-dispatch left a reservation allocated after cancellation cleanup" - ) assert session.status == SessionStatus.CANCELLED assert task.resources_drained @@ -1734,7 +1568,6 @@ def test_kv_build_failure_reports_safe_pre_submission_failure(monkeypatch) -> No sender = _make_owned_sender() sender._device_id = 0 sender._agent = Mock() - sender._bounce = Mock() dealer = Mock() sender._get_result_dealer = Mock(return_value=dealer) task = transfer_mod.KVSendTask( @@ -1778,7 +1611,7 @@ def test_kv_build_failure_reports_safe_pre_submission_failure(monkeypatch) -> No sender._get_result_dealer.assert_called_once_with("receiver", legacy_shared=True) dealer.send.assert_called_once() instance_rank, result_rid, result_slice, is_last, status_code, _ = ( - transfer_mod._KV_RESULT_PREFIX.unpack(dealer.send.call_args.args[0][1]) + transfer_mod._KV_RESULT_FRAME.unpack(dealer.send.call_args.args[0][1]) ) assert (instance_rank, result_rid, result_slice, is_last) == ( sender._instance_rank, @@ -1836,7 +1669,7 @@ def test_owned_kv_build_failure_settles_every_published_peer(monkeypatch) -> Non results = [sender._send_task_queues[0].get_nowait() for _ in peers] assert sender._send_task_queues[0].empty() assert {endpoint for endpoint, _ in results} == {"receiver"} - assert {transfer_mod._KV_RESULT_PREFIX.unpack(message[1])[2] for _, message in results} == { + assert {transfer_mod._KV_RESULT_FRAME.unpack(message[1])[2] for _, message in results} == { 1, 2, } @@ -1882,7 +1715,6 @@ def test_ambiguous_sender_result_retains_source_and_reports_in_doubt( peer_rank = 2 sender = _make_owned_sender() sender._device_id = 0 - sender._bounce = Mock() wait = Mock(return_value=False) detail = "ERROR" if failure_mode == "wait_exception": @@ -1947,7 +1779,7 @@ def submit(_request): if meta_type == transfer_mod.WriteMetaType.KV: sender._deliver_kv_to_agent(write_meta) - status_code = transfer_mod._KV_RESULT_PREFIX.unpack(dealer.send.call_args.args[0][1])[4] + status_code = transfer_mod._KV_RESULT_FRAME.unpack(dealer.send.call_args.args[0][1])[4] assert transfer_mod._AGENT_RESULT_BY_CODE[status_code] == AgentResult.IN_DOUBT else: sender._deliver_aux_to_agent(write_meta) diff --git a/tests/unittest/disaggregated/test_unsubmitted_session.py b/tests/unittest/disaggregated/test_unsubmitted_session.py index 09eb9bfa2296..fbda04972923 100644 --- a/tests/unittest/disaggregated/test_unsubmitted_session.py +++ b/tests/unittest/disaggregated/test_unsubmitted_session.py @@ -403,7 +403,6 @@ def _receive_session( receiver._shutdown = False receiver._ownership_admission_lock = threading.Lock() receiver._ownership_poisoned = None - receiver._bounce = Mock() receiver._dealers = {} receiver._messenger = Mock() receiver._registrar = SimpleNamespace( @@ -456,23 +455,16 @@ def test_adp_candidate_acks_do_not_replace_selected_writer_or_aux_evidence( assert session.close() -@pytest.mark.parametrize("unsettled", ["backend", "local_cuda"]) -def test_session_ack_cannot_bypass_existing_unsettled_access(unsettled: str) -> None: +def test_session_ack_cannot_bypass_existing_unsettled_access() -> None: owner = transfer_mod._ReceiveOperationOwner() owner.begin_publication() owner.seal_writer_cohort(1, published_writers={7, 8}) owner.finish_publication() - if unsettled == "backend": - owner.record_writer_in_doubt(7) - else: - owner.record_writer_result(7, True, wait_for_local_completion=True) + owner.record_writer_in_doubt(7) owner.record_session_quiesced(7) owner.record_session_quiesced(8) assert not owner.resources_drained - if unsettled == "backend": - owner.record_writer_settlement(7) - else: - owner.finish_local_completion() + owner.record_writer_settlement(7) assert owner.resources_drained diff --git a/tests/unittest/disaggregated/test_verified_reuse_admission.py b/tests/unittest/disaggregated/test_verified_reuse_admission.py index 1e2a130e0018..b45440231522 100644 --- a/tests/unittest/disaggregated/test_verified_reuse_admission.py +++ b/tests/unittest/disaggregated/test_verified_reuse_admission.py @@ -33,7 +33,6 @@ import numpy as np import pytest -import tensorrt_llm._torch.disaggregation.native.bounce as bounce_mod import tensorrt_llm._torch.disaggregation.native.transfer as transfer_mod from tensorrt_llm import DisaggregatedParams from tensorrt_llm._torch.disaggregation.base import CacheKind, Chunk, TokenRange @@ -103,7 +102,6 @@ def _make_rx_session( prompt_len=prompt_len, ) session.receive(_make_chunk(prompt_len), expected_write_bytes=expected_write_bytes) - session._receiver._bounce.is_bounced.return_value = False return session @@ -340,7 +338,6 @@ def _make_sender() -> transfer_mod.Sender: sender._instance_rank = 5 sender._device_id = 0 sender._agent = Mock() - sender._bounce = Mock() sender._registrar = SimpleNamespace( self_rank_info=SimpleNamespace(instance_name="ctx", instance_rank=5) ) @@ -379,13 +376,11 @@ def _deliver_kv(monkeypatch, submit_result): receiver_slice_id=0, is_last_slice=True, ) - monkeypatch.setattr(bounce_mod, "build_send_request", Mock(return_value=(Mock(), None))) + monkeypatch.setattr(transfer_mod.Sender, "_make_agent_request", Mock(return_value=Mock())) monkeypatch.setattr(transfer_mod.Sender, "_submit_transfer", Mock(return_value=submit_result)) sender._deliver_kv_to_agent(write_meta) dealer.send.assert_called_once() - return transfer_mod._KV_RESULT_PREFIX.unpack( - dealer.send.call_args.args[0][1][: transfer_mod._KV_RESULT_PREFIX.size] - ) + return transfer_mod._KV_RESULT_FRAME.unpack(dealer.send.call_args.args[0][1]) class TestSenderAttestation: From 82b2bb4c4c4801f20ea1ca889b0b9ad0d0929c5b Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:20:43 +0800 Subject: [PATCH 03/11] [None][chore] Drop agent_bounce_buffer_enable from configs and rewrite bounce docs kv_cache_bounce_size_mb alone now selects the C++ transfer-agent bounce, so remove agent_bounce_buffer_enable from the perf, perf-sanity and E2E YAMLs and from the MiniMax-M3 accuracy test, and drop the stale "Python bounce layer off" comment from the MiniMax-M3 perf YAMLs. Rewrite the docs and examples that described the Python bounce: - disagg-serving.md and kv-transfer.md describe the C++ transfer-agent bounce buffer. - The Kimi K3 examples are updated separately (see the K3 commit). - The cache transceiver test harness comment documents the gate and its TRTLLM_NIXL_BOUNCE_* env knobs. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- docs/source/developer-guide/kv-transfer.md | 2 +- docs/source/features/disagg-serving.md | 12 ++++++++++-- .../run_cache_transceiver_test.py | 16 ++++++++++------ .../defs/accuracy/test_disaggregated_serving.py | 1 - ...verlap_transceiver_runtime_python_bounce.yaml | 2 -- ...64_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ...24_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ...28_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ...n4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ...ctx1_dep4_gen1_dep32_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ..._ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ...ctx1_dep4_gen1_dep16_eplb0_mtp1_ccb-NIXL.yaml | 2 -- ..._ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ...ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ..._ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml | 2 -- ...24_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ...28_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ...n4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ...ctx1_dep4_gen1_dep32_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ..._ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ..._ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml | 2 -- ...ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ..._ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml | 2 -- ...tx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ..._ctx1_tp4_gen2_tp4_eplb0_eagle3_ccb-NIXL.yaml | 1 - ...tx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml | 2 -- ..._ctx1_tp4_gen2_tp4_eplb0_eagle3_ccb-NIXL.yaml | 1 - 27 files changed, 21 insertions(+), 54 deletions(-) diff --git a/docs/source/developer-guide/kv-transfer.md b/docs/source/developer-guide/kv-transfer.md index 603d976ffd59..290d018f2593 100644 --- a/docs/source/developer-guide/kv-transfer.md +++ b/docs/source/developer-guide/kv-transfer.md @@ -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. diff --git a/docs/source/features/disagg-serving.md b/docs/source/features/disagg-serving.md index cc583ac31d07..138385772b1d 100644 --- a/docs/source/features/disagg-serving.md +++ b/docs/source/features/disagg-serving.md @@ -157,9 +157,17 @@ cache_transceiver_config: `kv_transfer_timeout_ms` bounds how long a request may wait for its KV cache before it is cancelled and cleaned up. The default is `60000`. -`kv_cache_bounce_size_mb` is `0` by default, which sends each KV block separately. Setting it to a positive size enables a bounce buffer, which helps when a request's blocks are scattered. With the default Python implementation this is one contiguous buffer of that many MiB per direction and a single NIXL write per request; the C++ implementation (next paragraph) instead uses one buffer shared by both directions and chunked writes. The default Python implementation is designed for fabric (MNNVL) memory; the C++ implementation falls back to `cudaMalloc` where fabric memory is unavailable. +`kv_cache_bounce_size_mb` (default `0`) sets the size, in MiB, of the KV-cache bounce buffer; `0` turns it off. Without the buffer, each KV fragment is a separate descriptor in the NIXL write. With it, the sender gathers an eligible write's fragments into the buffer and sends them in chunks, and the receiver scatters them into its KV cache. This helps when a write consists of many small fragments, which is typical when context and generation split attention heads differently (for example, context TP != generation TP, or attention DP). -`agent_bounce_buffer_enable` (default `false`) switches to the C++ transfer-agent bounce implementation, which uses one buffer of `kv_cache_bounce_size_mb` MiB shared by send and receive (prefer a power of two; the usable capacity rounds down to one). By default it engages only for KV writes with many small descriptors (at least 1024, averaging at most 16 KiB), i.e. head-mismatch layouts such as context TP != generation TP or DP; head-matched layouts (MLA, symmetric TP) stay on the standard NIXL path unless `agent_bounce_params` relaxes the gate (lower `min_descriptor_count`, where `0` means no minimum; raise `max_average_descriptor_size`, e.g. `4MB`; `max_average_descriptor_size: 0` routes every outbound write to standard NIXL, acting as an outbound gate off switch — it disables outbound routing only: the arena is still allocated, the handshake still advertised and inbound transfers still served; set `agent_bounce_buffer_enable: false` to turn the feature off). Set it consistently on the context and generation workers; in particular `request_timeout_ms` and the effective chunk cap must match on both sides, otherwise the pair falls back to standard NIXL with a WARNING when the peer is loaded (transfers are then routed silently). The compared chunk cap is the effective one, min(`max_chunk_size`, the arena's usable capacity): `kv_cache_bounce_size_mb` matters only when its usable capacity (rounded down to a power of two) is below `max_chunk_size`; both sides must then clamp to the same chunk cap, otherwise the pair falls back to standard NIXL. +- The buffer lives in the C++ NIXL agent: one buffer per rank, shared by send and receive. It uses fabric (MNNVL) memory where available and `cudaMalloc` otherwise. Use a power of two such as `256`, `512` or `1024`; for other values only the largest power of two that fits is used. +- It requires the Python transceiver with the C++ NIXL agent. On a build without bounce support (libzmq), a positive value fails at worker startup. The C++ transceiver and the pure-Python NIXL agent ignore it with a warning. +- Set it on both the context and the generation workers. If either side has `0`, the pair uses standard NIXL. + +By default, only writes with at least 1024 descriptors averaging at most 16 KiB use the bounce buffer, so head-matched layouts (MLA, symmetric TP) stay on standard NIXL. To change this, set `agent_bounce_params`: lower `min_descriptor_count` (`0` means no minimum) or raise `max_average_descriptor_size` (for example `4MB`). Setting `max_average_descriptor_size: 0` stops this worker from sending through the bounce buffer; it still allocates the buffer and still receives bounced writes from peers. To turn the feature off entirely, set `kv_cache_bounce_size_mb: 0`. + +The two sides must also agree on `request_timeout_ms` (an `agent_bounce_params` key) and on the effective chunk size, `min(max_chunk_size, buffer size rounded down to a power of two)`. With the default `max_chunk_size` (32 MiB), the chunk sizes can only differ when a buffer is smaller than 32 MiB. If either value differs, the worker logs a WARNING when it first loads the peer's metadata, and then uses standard NIXL for that peer without further messages. + +`agent_bounce_buffer_enable` is deprecated and ignored. Remove it from existing configs. For example, you could launch two context servers and one generation server as follows: diff --git a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py index 397e4a8aa434..08dba396eebe 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py +++ b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py @@ -1033,12 +1033,16 @@ def arm_watchdog(combination_idx, reqlen_idx, what): backend=backend, transceiver_runtime=(None if runtime == "CPP" else "PYTHON"), max_tokens_in_buffer=cfg["kv_cache"]["max_tokens_in_buffer"], - # PYTHON (V2) only; 0 keeps bounce off. With bounce on, the KV data - # rides a fabric-VMM staging buffer (CU_MEM_HANDLE_TYPE_FABRIC), which - # is what lets UCX pick cuda_ipc across NVL72 nodes -- direct - # pool-to-pool transfers from non-fabric allocations fall back to - # much slower host-staged tcp, so enable bounce for cross-node - # transfers inside an NVLink domain. + # PYTHON (V2) only; 0 keeps bounce off. A positive size enables the C++ + # transfer-agent bounce buffer (a fabric-memory arena where available). Its + # default admission gate only admits writes with >= 1024 descriptors averaging + # <= 16 KiB (head-mismatch layouts). This harness is symmetric, so it sends one + # descriptor per block (far above 16 KiB) and at most a few hundred blocks per + # request: with the default gate every transfer is declined. To exercise the + # bounce, set TRTLLM_NIXL_BOUNCE_MIN_DESCRIPTOR_COUNT=1 and + # TRTLLM_NIXL_BOUNCE_MAX_AVERAGE_DESCRIPTOR_SIZE_BYTES=4MB on both sides (this + # harness has no agent_bounce_params passthrough); bounce_submit_count in the + # shutdown log shows engagement. kv_cache_bounce_size_mb=int(cfg["kv_cache"].get("bounce_size_mb", 0)), # For the Python sender, leave deterministic headroom for the # signal handler and final CTX/GEN ownership handshake after its diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 01cc40e03e4f..2e09040cedc4 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -2698,7 +2698,6 @@ def test_nvfp4_eagle3_disagg_head_mismatched_bounce(self) -> None: "cache_transceiver_config": { "backend": "NIXL", "transceiver_runtime": "PYTHON", - "agent_bounce_buffer_enable": True, "kv_cache_bounce_size_mb": 512, # Relax the admission gates for this short CI workload. # The production defaults remain 1024 descriptors / 16 KiB. diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml index 5453d5f88319..4decc5880f3a 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_overlap_transceiver_runtime_python_bounce.yaml @@ -25,7 +25,6 @@ context_servers: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 128 - agent_bounce_buffer_enable: true agent_bounce_params: min_descriptor_count: "1" max_average_descriptor_size: "4MB" @@ -46,7 +45,6 @@ generation_servers: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 128 - agent_bounce_buffer_enable: true agent_bounce_params: min_descriptor_count: "1" max_average_descriptor_size: "4MB" diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_1k1k_con64_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_1k1k_con64_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml index 9646158022c1..e11c4ab7ca1f 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_1k1k_con64_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_1k1k_con64_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml @@ -66,7 +66,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false num_postprocess_workers: 4 stream_interval: 20 @@ -92,5 +91,4 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml index 9430939984e8..50f6ecd73a03 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml @@ -78,7 +78,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false num_postprocess_workers: 4 stream_interval: 20 @@ -107,7 +106,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false num_postprocess_workers: 4 stream_interval: 20 diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml index c8d1841fd6df..0d9db251f398 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml @@ -66,7 +66,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false num_postprocess_workers: 4 stream_interval: 20 @@ -92,5 +91,4 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml index e2474d205b20..341574067d8f 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml @@ -66,7 +66,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false num_postprocess_workers: 4 stream_interval: 20 @@ -92,5 +91,4 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con2048_ctx1_dep4_gen1_dep32_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con2048_ctx1_dep4_gen1_dep32_eplb0_mtp0_ccb-NIXL.yaml index 319a6887e91f..61f41c1ba58f 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con2048_ctx1_dep4_gen1_dep32_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con2048_ctx1_dep4_gen1_dep32_eplb0_mtp0_ccb-NIXL.yaml @@ -68,7 +68,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false trust_remote_code: true num_postprocess_workers: 4 @@ -96,6 +95,5 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false trust_remote_code: true diff --git a/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con4_ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con4_ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml index 0534262f1cdf..97c9fd4cae86 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con4_ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con4_ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml @@ -66,7 +66,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false trust_remote_code: true num_postprocess_workers: 4 @@ -94,6 +93,5 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false trust_remote_code: true diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-r1-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp1_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-r1-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp1_ccb-NIXL.yaml index e149af64829a..acb5d42767dd 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-r1-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp1_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_deepseek-r1-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp1_ccb-NIXL.yaml @@ -68,7 +68,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 2048 - agent_bounce_buffer_enable: true kv_transfer_timeout_ms: 600000 disable_overlap_scheduler: false speculative_config: &id001 @@ -99,7 +98,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 2048 - agent_bounce_buffer_enable: true kv_transfer_timeout_ms: 600000 disable_overlap_scheduler: false speculative_config: *id001 diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_1k1k_con4_ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_1k1k_con4_ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml index b03eb065fee5..9b9f7288292d 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_1k1k_con4_ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_1k1k_con4_ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml @@ -66,7 +66,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false trust_remote_code: true num_postprocess_workers: 4 @@ -94,6 +93,5 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false trust_remote_code: true diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml index fddf7e101696..1cf5ef761a19 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml @@ -68,7 +68,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 2048 - agent_bounce_buffer_enable: true kv_transfer_timeout_ms: 600000 disable_overlap_scheduler: false trust_remote_code: true @@ -97,7 +96,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 2048 - agent_bounce_buffer_enable: true kv_transfer_timeout_ms: 600000 disable_overlap_scheduler: false trust_remote_code: true diff --git a/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml index 5c24767f2f39..55987597a023 100644 --- a/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf-sanity/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml @@ -66,7 +66,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false speculative_config: &id001 decoding_type: Eagle @@ -99,7 +98,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false speculative_config: *id001 trust_remote_code: true diff --git a/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml index 79785b570621..4b2081c527ad 100644 --- a/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con1024_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml @@ -78,7 +78,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false num_postprocess_workers: 4 stream_interval: 20 @@ -107,7 +106,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: true num_postprocess_workers: 4 stream_interval: 20 diff --git a/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml index 8cba8b225ce2..4e6d4854ddf2 100644 --- a/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con128_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml @@ -77,7 +77,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false num_postprocess_workers: 4 stream_interval: 20 @@ -103,5 +102,4 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false diff --git a/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml index 79a91064884d..4a73431aaaf6 100644 --- a/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb200_gpt-oss-120b-fp4_8k1k_con4_ctx1_tp1_gen1_tp4_eplb0_mtp0_ccb-NIXL.yaml @@ -77,7 +77,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false num_postprocess_workers: 4 stream_interval: 20 @@ -103,5 +102,4 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false diff --git a/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con2048_ctx1_dep4_gen1_dep32_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con2048_ctx1_dep4_gen1_dep32_eplb0_mtp0_ccb-NIXL.yaml index 319a6887e91f..61f41c1ba58f 100644 --- a/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con2048_ctx1_dep4_gen1_dep32_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con2048_ctx1_dep4_gen1_dep32_eplb0_mtp0_ccb-NIXL.yaml @@ -68,7 +68,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false trust_remote_code: true num_postprocess_workers: 4 @@ -96,6 +95,5 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false trust_remote_code: true diff --git a/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con4_ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con4_ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml index 0534262f1cdf..97c9fd4cae86 100644 --- a/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con4_ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb200_kimi-k25-thinking-fp4_1k1k_con4_ctx1_dep4_gen1_tep4_eplb0_mtp0_ccb-NIXL.yaml @@ -66,7 +66,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false trust_remote_code: true num_postprocess_workers: 4 @@ -94,6 +93,5 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false trust_remote_code: true diff --git a/tests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml index 54f432ed1432..f3d426ad03b3 100644 --- a/tests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb200_stress-gpt-oss-120b-fp4_8k1k_ctx1_tp1_gen1_tp4_eplb0_eagle3_ccb-NIXL.yaml @@ -83,7 +83,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false speculative_config: &id001 decoding_type: Eagle @@ -113,6 +112,5 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false speculative_config: *id001 diff --git a/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml index e913c63eb198..5a52e99526f2 100644 --- a/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4096_ctx1_dep4_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml @@ -79,7 +79,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 2048 - agent_bounce_buffer_enable: true kv_transfer_timeout_ms: 600000 disable_overlap_scheduler: false trust_remote_code: true @@ -108,7 +107,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 2048 - agent_bounce_buffer_enable: true kv_transfer_timeout_ms: 600000 disable_overlap_scheduler: true trust_remote_code: true diff --git a/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml index 88e1399ce3ce..f3ca6b459371 100644 --- a/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb300_kimi-k25-thinking-fp4_8k1k_con4_ctx1_dep4_gen1_tep8_eplb0_mtp3_ccb-NIXL.yaml @@ -77,7 +77,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false speculative_config: &id001 decoding_type: Eagle @@ -110,7 +109,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 512 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false speculative_config: *id001 trust_remote_code: true diff --git a/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml index 3bdbf964503a..045688195443 100644 --- a/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml @@ -99,7 +99,6 @@ worker_config: max_tokens_in_buffer: 8448 # Shared C++ agent arena for chunked gather/write/scatter transfers. kv_cache_bounce_size_mb: 1024 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false trust_remote_code: true ctx: @@ -126,6 +125,5 @@ worker_config: transceiver_runtime: PYTHON max_tokens_in_buffer: 8448 kv_cache_bounce_size_mb: 1024 - agent_bounce_buffer_enable: true disable_overlap_scheduler: true trust_remote_code: true diff --git a/tests/scripts/perf/disaggregated/gb300_minimax-m3-fp4_8k1k_con30_ctx1_tp4_gen2_tp4_eplb0_eagle3_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb300_minimax-m3-fp4_8k1k_con30_ctx1_tp4_gen2_tp4_eplb0_eagle3_ccb-NIXL.yaml index bfbb485aae3d..a9c4e8fbe675 100644 --- a/tests/scripts/perf/disaggregated/gb300_minimax-m3-fp4_8k1k_con30_ctx1_tp4_gen2_tp4_eplb0_eagle3_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb300_minimax-m3-fp4_8k1k_con30_ctx1_tp4_gen2_tp4_eplb0_eagle3_ccb-NIXL.yaml @@ -124,7 +124,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_transfer_timeout_ms: 600000 - # Keep the separate Python bounce layer off; the C++ NIXL agent owns it. kv_cache_bounce_size_mb: 0 disable_overlap_scheduler: false trust_remote_code: true diff --git a/tests/scripts/perf/disaggregated/vr200_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/vr200_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml index 208f4dc467fd..5ebf776f2ca1 100644 --- a/tests/scripts/perf/disaggregated/vr200_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/vr200_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml @@ -109,7 +109,6 @@ worker_config: max_tokens_in_buffer: 8448 # Shared C++ agent arena for chunked gather/write/scatter transfers. kv_cache_bounce_size_mb: 1024 - agent_bounce_buffer_enable: true disable_overlap_scheduler: false trust_remote_code: true ctx: @@ -136,6 +135,5 @@ worker_config: transceiver_runtime: PYTHON max_tokens_in_buffer: 8448 kv_cache_bounce_size_mb: 1024 - agent_bounce_buffer_enable: true disable_overlap_scheduler: true trust_remote_code: true diff --git a/tests/scripts/perf/disaggregated/vr200_minimax-m3-fp4_8k1k_con30_ctx1_tp4_gen2_tp4_eplb0_eagle3_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/vr200_minimax-m3-fp4_8k1k_con30_ctx1_tp4_gen2_tp4_eplb0_eagle3_ccb-NIXL.yaml index 59db7a3c3ed8..b408509bc6db 100644 --- a/tests/scripts/perf/disaggregated/vr200_minimax-m3-fp4_8k1k_con30_ctx1_tp4_gen2_tp4_eplb0_eagle3_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/vr200_minimax-m3-fp4_8k1k_con30_ctx1_tp4_gen2_tp4_eplb0_eagle3_ccb-NIXL.yaml @@ -132,7 +132,6 @@ worker_config: backend: NIXL transceiver_runtime: PYTHON kv_transfer_timeout_ms: 600000 - # Keep the separate Python bounce layer off; the C++ NIXL agent owns it. kv_cache_bounce_size_mb: 0 disable_overlap_scheduler: false trust_remote_code: true From 30c11bf683d4c29286bb99aff6531d8165338093 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:21:42 +0800 Subject: [PATCH 04/11] [None][chore] Reword C++ bounce messages for kv_cache_bounce_size_mb kv_cache_bounce_size_mb > 0 alone enables the C++ transfer-agent bounce, so the user-facing C++ strings and comments that named agent_bounce_buffer_enable now name kv_cache_bounce_size_mb: the no-libzmq THROW, the params-without-arena WARNING, the CMake status message and the config comments. The peer-handshake WARNING keeps the older-release cause (a peer that predates this change with agent_bounce_buffer_enable off). No behavior change. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../nixl_utils/CMakeLists.txt | 6 ++--- .../nixl_utils/bounce/BounceConfig.h | 27 +++++++++---------- .../nixl_utils/bounce/BounceTransport.cpp | 11 ++++---- .../nixl_utils/transferAgent.cpp | 19 +++++++------ .../nixl_utils/transferAgent.h | 7 +++-- .../executor/bounce/bounceAgentE2ETest.cpp | 5 ++-- 6 files changed, 36 insertions(+), 39 deletions(-) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt index d8dac95b8d99..12f5bf117d65 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/CMakeLists.txt @@ -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) @@ -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() diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h index f5acdefb0abe..75106754a93c 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceConfig.h @@ -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 @@ -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 @@ -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. diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp index 5657b249019a..ad5c7da39299 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/bounce/BounceTransport.cpp @@ -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(); @@ -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; } diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp index 61633d141244..811c31895ffa 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.cpp @@ -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 @@ -145,9 +145,8 @@ class BounceTransferStatus final : public TransferStatus void NixlTransferAgent::maybeInitBounce( std::size_t agentBufferSizeMb, std::unordered_map 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) { @@ -155,8 +154,7 @@ void NixlTransferAgent::maybeInitBounce( { 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; @@ -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()); } } diff --git a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h index 90794052d6f2..07a87acd862c 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/nixl_utils/transferAgent.h @@ -226,10 +226,9 @@ class NixlTransferAgent : public BaseTransferAgent /// Per-agent maps because different remote agents may have overlapping virtual addresses. std::unordered_map 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 mBounce; /// Lazily create the bounce transport (ctor, before any metadata exchange) when enabled. diff --git a/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp b/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp index 7fa215f170e6..45632b835dbc 100644 --- a/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp +++ b/cpp/tests/unit_tests/executor/bounce/bounceAgentE2ETest.cpp @@ -138,9 +138,8 @@ std::unique_ptr 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()) From 3a885e460b9bb435a61f45e5247981c351b77f9c Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:22:29 +0800 Subject: [PATCH 05/11] [None][chore] Remove the KVCacheManagerV2 VMM introspection hooks The PooledPhysMemAllocator/VirtMem nanobind classes and their Python _introspection forwarders existed only for the Python KV bounce buffer, which is gone. Remove them together with their lifetime test. cudaVirtMem.h still reaches its users through storage/core.h. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../batch_manager/kvCacheManagerV2.cpp | 15 -- .../kv_cache_manager_v2/_introspection.py | 18 +-- .../test_virt_mem_lifetime.py | 133 ------------------ 3 files changed, 1 insertion(+), 165 deletions(-) delete mode 100644 tests/unittest/kv_cache_manager_v2_tests/test_virt_mem_lifetime.py diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp index a601e1e100cd..823f00886e73 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManagerV2.cpp @@ -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" @@ -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()); - // 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_(mIntrospection, "PooledPhysMemAllocator") - .def(nb::init(), nb::arg("phys_mem_size")) - .def_prop_ro("device_id", &kv::PooledPhysMemAllocator::deviceId); - nb::class_(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(), 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_(m, "IKvCacheColdPageCodec"); m.def("create_default_kv_cache_cold_page_codec", &kv::createDefaultKvCacheColdPageCodec, diff --git a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py index da87028989ab..b1ea2b6a3c41 100644 --- a/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py +++ b/tensorrt_llm/runtime/kv_cache_manager_v2/_introspection.py @@ -18,9 +18,7 @@ Everything here forwards to the native ``_introspection`` submodule of the bindings. The indirection exists so callers import one Python path rather than reaching into ``tensorrt_llm.bindings`` directly, and so return values are normalised to plain Python -containers. Most hooks are white-box helpers for tests and accuracy harnesses, but a few -back production internals -- the stats report and the disaggregated bounce buffer -- so -treat this as internal API rather than test-only scaffolding. +containers. The hooks are white-box helpers for tests and accuracy harnesses. """ from __future__ import annotations @@ -46,20 +44,6 @@ def _cpp() -> Any: return cpp_introspection -#: CUDA virtual-memory primitives, forwarded rather than wrapped: they are classes the -#: caller constructs, so there is nothing to normalise. Used by the native disaggregated -#: bounce buffer to reserve one contiguous fabric region. -_FORWARDED_TYPES = ("PooledPhysMemAllocator", "VirtMem") - - -def __getattr__(name: str) -> Any: - # Resolved on first access, not at import, so importing this module never depends on - # the native submodule being present -- the same contract the hooks below follow. - if name in _FORWARDED_TYPES: - return getattr(_cpp(), name) - raise AttributeError(f"module {__name__!r} has no attribute {name!r}") - - def poison_for_testing(reason: str) -> None: """Set the poison latch directly, to exercise the refusal paths.""" _cpp().poison_for_testing(reason) diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_virt_mem_lifetime.py b/tests/unittest/kv_cache_manager_v2_tests/test_virt_mem_lifetime.py deleted file mode 100644 index 845ce2e7593e..000000000000 --- a/tests/unittest/kv_cache_manager_v2_tests/test_virt_mem_lifetime.py +++ /dev/null @@ -1,133 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""``VirtMem`` must outlive the Python handle to the allocator it borrows. - -``VirtMem`` holds ``PooledPhysMemAllocator&`` in C++ and dereferences it whenever it -unmaps physical memory. The binding's ``nb::keep_alive`` is what ties the allocator's -Python lifetime to the ``VirtMem``, and the native disaggregated bounce buffer depends -on it: ``bounce/buffer.py`` builds the allocator as a local and lets it go out of scope -while keeping only the ``VirtMem``. Without that tie the allocator is freed at the end -of that constructor and every later unmap reads freed memory. - -Only ``test_virt_mem_holds_a_reference_to_its_allocator`` guards the contract. Dropping -``keep_alive`` from the binding and rerunning this file leaves the end-to-end case below -passing, because the freed allocator reads back intact in a release build; the reference -count is what actually changes. Keep any new case here anchored on something observable -rather than on an unmap that happens to succeed. -""" - -import gc -import sys - -import pytest -import torch - -from tensorrt_llm.runtime.kv_cache_manager_v2._introspection import PooledPhysMemAllocator, VirtMem - -pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA") - -_CHUNK = 2 << 20 - - -@pytest.fixture(autouse=True) -def _cuda_context() -> None: - # PooledPhysMemAllocator drives the driver API directly, so it needs a current - # primary context; torch only creates one lazily on first use. - torch.cuda.init() - torch.zeros(1, device="cuda") - - -def _make_virt_mem_from_temporary_allocator() -> VirtMem: - """Build a VirtMem exactly as bounce/buffer.py does: allocator is a local.""" - allocator = PooledPhysMemAllocator(_CHUNK) - return VirtMem(_CHUNK, allocator, init_num_phys_mem=1) - - -def test_virt_mem_outlives_its_allocator_handle() -> None: - vm = _make_virt_mem_from_temporary_allocator() - # The only Python reference to the allocator went out of scope on return, so a - # missing keep_alive leaves the C++ reference dangling from here on. - gc.collect() - - assert vm.address != 0 - # destroy() unmaps the backing chunk, which is the path that dereferences the - # allocator. It must still be alive. - vm.destroy() - - -def test_virt_mem_holds_a_reference_to_its_allocator() -> None: - """The retention is asserted directly, not inferred from a surviving unmap. - - Dropping ``keep_alive`` leaves freed memory that a release build usually reads back - without faulting, so a test that only called ``destroy()`` could keep passing while - the contract is broken. The reference count is exact and fails the moment it goes. - """ - unheld = PooledPhysMemAllocator(_CHUNK) - baseline = sys.getrefcount(unheld) - - allocator = PooledPhysMemAllocator(_CHUNK) - vm = VirtMem(_CHUNK, allocator, init_num_phys_mem=1) - assert sys.getrefcount(allocator) == baseline + 1 - - del allocator - gc.collect() - assert vm.address != 0 - vm.destroy() - - -def _vm_size_bytes() -> int: - """Total virtual address space mapped by this process, from /proc/self/status.""" - with open("/proc/self/status", encoding="ascii") as status: - for line in status: - if line.startswith("VmSize:"): - return int(line.split()[1]) * 1024 - raise RuntimeError("VmSize missing from /proc/self/status") - - -@pytest.mark.skipif(not sys.platform.startswith("linux"), reason="reads /proc/self/status") -def test_failed_construction_releases_its_address_reservation() -> None: - """A VirtMem whose initial extend() throws must not strand the reservation. - - ``VirtMem`` reserves the address range, then maps the initial chunks. A constructor - that throws never runs ``~VirtMem``, so the release has to happen on that path - explicitly. ``create_bounce()`` treats an allocation failure as a soft fallback and - keeps the process running, so anything stranded here is held until exit. - - The reservation is charged to the process's virtual size, which makes the leak - directly measurable: without the rollback this grows by one chunk per attempt. - """ - # Sized off the device rather than a constant, so the first cuMemCreate cannot be - # satisfied on any GPU this runs on and the constructor throws right after the - # reservation succeeds. A fixed value would eventually be smaller than some card's - # memory, and the request would start succeeding instead. - granularity = 2 << 20 - total = torch.cuda.get_device_properties(0).total_memory - leak_chunk = -(-(2 * total) // granularity) * granularity - attempts = 8 - - allocator = PooledPhysMemAllocator(leak_chunk) - before = _vm_size_bytes() - for _ in range(attempts): - with pytest.raises(Exception): - VirtMem(leak_chunk, allocator, init_num_phys_mem=1) - growth = _vm_size_bytes() - before - - # One stranded reservation is already the bug, so allow only a fraction of one. - assert growth < leak_chunk // 2, ( - f"virtual size grew {growth / (1 << 30):.1f} GiB over {attempts} failed " - f"constructions; each strands {leak_chunk / (1 << 30):.0f} GiB when the " - f"constructor does not release its reservation" - ) From f2a4bd959f606e4da9fc6a94730add2922db4ae2 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Fri, 9 Oct 2026 17:22:59 +0800 Subject: [PATCH 06/11] [None][chore] Drop KVRecvTask's ownership state lock The lock serialized KVRecvTask.complete() against fail() because complete() could run on the Python bounce scatter worker without RxSession.lock. With the Python bounce gone, every complete() and every post-publication fail() runs under RxSession.lock, and the only lock-free fail() happens before publication. Keep the "a committed failure stands" check in complete() and drop the lock. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../_torch/disaggregation/native/transfer.py | 45 +++++++++---------- 1 file changed, 20 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index 4b2c4f431bfb..bcecbd2eaee7 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -2767,32 +2767,29 @@ def __init__( self._aux_slot = aux_slot self._perf_timer = PerfTimer() if perf_log_manager.enabled else None self._physical_owner: Optional[_ReceiveOperationOwner] = None - self._ownership_state_lock: Optional[threading.Lock] = None def fail(self, exc: Exception) -> None: + """Mark the task ERROR. + + Callers hold RxSession.lock, except dispatch_task's PeerIncompatibleError path, which + runs before publication. + """ self._logical_outcomes.fail(exc) - if self._ownership_state_lock is None: - self._exception = exc - self.status = TaskStatus.ERROR - self._event.set() - return - with self._ownership_state_lock: - self._exception = exc - self.status = TaskStatus.ERROR - self._event.set() + self._exception = exc + self.status = TaskStatus.ERROR + self._event.set() def complete(self) -> None: - if self._ownership_state_lock is None: - self._logical_outcomes.complete(self._logical_index) - self.status = TaskStatus.TRANSFERRED - self._event.set() + """Mark the task TRANSFERRED unless it already failed. + + Callers hold RxSession.lock, so this check-then-set is atomic against fail(). Status is + set before _event, so a thread woken by _event reads the status that complete() set. + """ + if self.status == TaskStatus.ERROR: return - with self._ownership_state_lock: - if self.status == TaskStatus.ERROR: - return - self._logical_outcomes.complete(self._logical_index) - self.status = TaskStatus.TRANSFERRED - self._event.set() + self._logical_outcomes.complete(self._logical_index) + self.status = TaskStatus.TRANSFERRED + self._event.set() def wait(self, timeout: Optional[float] = None) -> bool: """Block until terminal state. Returns True if done, False on timeout.""" @@ -2853,7 +2850,6 @@ def reports_outstanding(self) -> bool: def begin_publication(self) -> None: if self._physical_owner is None: self._physical_owner = _ReceiveOperationOwner(self._retirement) - self._ownership_state_lock = threading.Lock() self._physical_owner.begin_publication() def _get_physical_owner(self) -> _ReceiveOperationOwner: @@ -3780,15 +3776,14 @@ def process_kv_agent_result( # attestation of destination bytes actually submitted and # completed for this chunk (0 unless the write completed). task.verified_write_bytes += transfer_size - # The completing report: every expected writer succeeded. A concurrent failure that - # already failed the task stands. complete() sets status before _event, keeping - # wait_complete's status-first poll correct. + # Once every expected writer has reported success for the last slice, complete the + # task and log perf, unless the task has already failed. if is_last_slice and all_succeeded and task.status != TaskStatus.ERROR: - ri = self._receiver._registrar.self_rank_info task.complete() # Record completion before best-effort diagnostics. if all(t.status == TaskStatus.TRANSFERRED for t in self._kv_tasks): self.transfer_end_time = tensorrt_llm.bindings.global_steady_clock_now() + ri = self._receiver._registrar.self_rank_info try: if task._perf_timer is not None: task._perf_timer.record_task_end(peer_rank) From d00c14aa19e9a87a255eae97cefb202b0d0e97e0 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Fri, 9 Oct 2026 23:31:56 +0800 Subject: [PATCH 07/11] [None][chore] Drop the bounce buffer from the Kimi K3 configs and docs Kimi K3 checkpoints load KimiK3ForConditionalGeneration, which prefers KVCacheManagerV2. Its MLA KV cache and KDA state live in V2 pools, which are allocated as fabric memory when the platform supports it, so standard NIXL can write them pool to pool across nodes over MNNVL. The K3 examples were written for the Python bounce buffer, which this PR removes, and the C++ agent bounce's default gate declines K3's head-matched writes, so a 1 GiB per-GPU buffer would be unused. Drop kv_cache_bounce_size_mb from the K3 example and the two K3 perf configs, together with the Python-bounce sizing guidance and measurements. Drop max_tokens_in_buffer there (V2, PP1, asynchronous), except the benchmark config's generation key, which its gen_only mode still uses. Fix the stale claims that auto resolves K3 to the C++ transceiver and that the KV pool is non-fabric. The README now describes K3 on KVCacheManagerV2, calls the transceiver "Python transceiver", and notes that use_kv_cache_manager_v2: false (or "auto" with a text-only checkpoint) runs the V1 MixedMambaHybridCacheManager, whose KDA state cannot use cross-node cuda_ipc. The _util.py comment on the Mixed manager is scoped to that V1 path. On GB300 NVL72 the current configs match the previous transfer path in GSM8K and end-to-end throughput, with cross-node KV going pool to pool over MNNVL. The README and the K3 deployment guide report this, and the README gives an optional bounce setting that shortens the per-request KV transfer for 1 GiB per GPU. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../deployment-guide-for-kimi-k3-on-trtllm.md | 2 +- examples/kimi_k3/disagg/README.md | 102 +++++++++--------- .../disagg/benchmark_kimi_k3_dep16.yaml | 33 ++---- examples/kimi_k3/disagg/ctx_config.yaml | 21 +--- .../kimi_k3/disagg/disagg_proxy_config.yaml | 4 +- examples/kimi_k3/disagg/gen_config.yaml | 14 +-- examples/kimi_k3/disagg/gen_config_no_sa.yaml | 20 +--- ..._dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml | 11 +- ..._dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml | 11 +- 9 files changed, 77 insertions(+), 141 deletions(-) diff --git a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md index 5848892ffc9c..c8f1c04a37ef 100644 --- a/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-kimi-k3-on-trtllm.md @@ -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 diff --git a/examples/kimi_k3/disagg/README.md b/examples/kimi_k3/disagg/README.md index 9faa21eab53a..10daaaba994d 100644 --- a/examples/kimi_k3/disagg/README.md +++ b/examples/kimi_k3/disagg/README.md @@ -1,10 +1,10 @@ # Kimi K3 disaggregated serving (ctx/gen split) Configuration pair + deployment wiring for running Kimi K3 with separate -context (prefill) and generation (decode) servers. Status: **validated -end-to-end on hardware** (GB300 NVL72, 1 ctx + 1 gen, DEP16 both sides, -GSM8K accuracy parity with aggregated serving) — see the caveats section -for constraints. +context (prefill) and generation (decode) servers. Status: validated +end-to-end on GB300 NVL72 with the current configs (1 ctx + 1 gen, DEP16 +both sides, KVCacheManagerV2); GSM8K accuracy and throughput are on par +with the previous transfer path. See the caveats section for constraints. ## Files @@ -12,7 +12,7 @@ for constraints. |---|---| | `ctx_config.yaml` | Context-server extra LLM-API options (DEP16, overlap scheduler off, no spec decode) | | `gen_config.yaml` | Generation-server options WITH suffix-automaton (SA) speculative decoding (DEP16, eager) | -| `gen_config_no_sa.yaml` | Generation-server options WITHOUT spec decode — use this first (CUDA graphs ON by default: GSM8K 96.89, 765/2138 tok/s @c64/c256 vs aggregated 643/1972; null `cuda_graph_config` for token-parity debugging) | +| `gen_config_no_sa.yaml` | Generation-server options WITHOUT spec decode — use this first (CUDA graphs ON by default: GSM8K 96.5, ~950/2043 output tok/s @c64/c256 on 8k/1k random prompts; null `cuda_graph_config` for token-parity debugging) | | `disagg_proxy_config.yaml` | `trtllm-serve disaggregated` proxy config (1 ctx + 1 gen) | | `benchmark_kimi_k3_dep16.yaml` | Config for the SLURM benchmark harness (`examples/disaggregated/slurm/benchmark/submit.py`) | @@ -32,27 +32,29 @@ for constraints. replicated share ~113 GiB + experts/8), leaving no activation headroom on GB300 (288 GiB) and not fitting GB200 (186 GiB). Treat DEP8-ctx as ruled out on GB200 and an open (likely negative) question on GB300. -- **`transceiver_runtime: PYTHON` is mandatory**: `auto` resolves to the - C++ transceiver, which throws at construction for K3's - `MixedMambaHybridCacheManager`. +- **Python transceiver** (`backend: NIXL`, `transceiver_runtime: PYTHON`): + only it can move the KDA state. It is the default (and `auto` selects it + for K3 too); the configs set it explicitly. - `disable_overlap_scheduler: true` on the ctx server (disagg requirement) and on the gen server (SA runs eager; also keeps the SA-off smoke maximally comparable). - `enable_block_reuse: false`, no chunked prefill, beam width 1 (model requirements); `tokens_per_block: 64` is kept for historical reasons (older TRTLLM-Gen MLA kernels required it). -- `max_tokens_in_buffer: 8448` covers the target max ISL of 8192; raise - it together with `max_num_tokens`/`max_seq_len` for longer ISL. -- **`kv_cache_bounce_size_mb: 1024` on both sides**: the V2 transceiver's default - pool-to-pool path cannot use inter-node cuda_ipc on MNNVL (the KV pool - is a plain, non-fabric allocation) and falls back to ~0.4 GB/s - host-staged tcp; the fabric-VMM bounce buffer restores cuda_ipc/MNNVL - eligibility (measured ~455 GB/s/GPU). Bounce engages automatically for payloads above `TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES` (default 2 MiB) — always true for K3's ~433 MiB per-request state. - The region must fit ONE request's full payload: fixed 433 MiB KDA - state + ~27 KB/token MLA latent (649 MiB at 8k ISL). A - 512 MiB value makes every 8k transfer fall back to the per-fragment - tcp path (`[kv-bounce] in-place: transfer 649MiB exceeds the 512MiB - bounce region`). +- **No bounce buffer**: leave `kv_cache_bounce_size_mb` unset. K3 runs on + KVCacheManagerV2, which keeps the MLA KV cache and the KDA state in its + pools. On GB300 NVL72 these pools are fabric memory, and standard NIXL + writes them pool to pool across nodes over MNNVL (cuda_ipc). With its + default gate a bounce buffer would stay unused: the gate only accepts + writes between mismatched head layouts, and K3's matched DEP16 layout + produces none. +- **Optional, shorter KV transfer**: for latency-sensitive or short-ISL + serving, set `kv_cache_bounce_size_mb: 1024` and + `agent_bounce_params: {min_descriptor_count: "1", max_average_descriptor_size: "4MB"}` + on both servers to send K3's writes through the C++ transfer-agent + bounce buffer. On GB300 NVL72 this cut the per-request 8k KV transfer + about 4x (~40 ms to ~10 ms) for 1 GiB per GPU; end-to-end throughput + was unchanged. ## KDA state payload size @@ -71,7 +73,7 @@ prerequisites for every worker shell (see caveats below for why): ```bash export UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc # on clusters where verbs cannot # initialize; a container-default - # UCX_TLS=tcp breaks V2 NIXL + # UCX_TLS=tcp breaks NIXL setup ``` 1. Start the context server (16-rank MPI world across its 4 nodes): @@ -128,45 +130,46 @@ python3 examples/disaggregated/slurm/benchmark/submit.py \ ## Current caveats (read before running) -1. **V2 transceiver "MPI hang" — root-caused, environmental (RESOLVED +1. **Python-transceiver "MPI hang" — root-caused, environmental (RESOLVED with the env pins).** A reported multi-node hang in `KvCacheTransceiverV2._exchange_rank_info` → `mpi_allgather` was a downstream symptom of `UCX_TLS=all` on nodes where `ud_verbs` cannot initialize: the broken transport wedges native NIXL/UCX agent init asymmetrically per rank, and the healthy ranks park forever in the - setup MPI collectives. Not an MPI/pmix or V2 code bug; with - `UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc` V2 NIXL passes multi-node - with no code change. -2. **SA ships eager here.** SA speculative decoding in disagg is - validated for accuracy (GSM8K parity with aggregated serving) with - CUDA graphs disabled, as configured in `gen_config.yaml`. SA with - CUDA graphs is functional (the MLA latent-cache append under CUDA - graphs handles spec-dec verification), but the disagg SA + graphs - perf points have not been re-measured yet, so `gen_config.yaml` - keeps graphs off. Start with `gen_config_no_sa.yaml` for the first - bring-up on a new cluster, then switch to `gen_config.yaml`. + setup MPI collectives. Not an MPI/pmix or transceiver code bug; with + `UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc` the Python transceiver passes + multi-node with no code change. +2. **SA ships eager here.** SA speculative decoding in disagg was + validated for accuracy (GSM8K parity with aggregated serving) on the + earlier transfer path, with CUDA graphs disabled as configured in + `gen_config.yaml`. SA with CUDA graphs is functional (the MLA + latent-cache append under CUDA graphs handles spec-dec verification), + but the disagg SA + graphs perf points have not been re-measured yet, + so `gen_config.yaml` keeps graphs off. Start with + `gen_config_no_sa.yaml` for the first bring-up on a new cluster, then + switch to `gen_config.yaml`. 3. **Matched-DP only.** Keep ctx and gen at identical DEP16 with attention-DP on both sides; heterogeneous parallelism passes peer validation but is not validated end-to-end (see constraints above). 4. **Cluster environment** (NVL72 nodes): on clusters where verbs transports cannot initialize, pin `UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc` (`UCX_TLS=all` hangs setup, - see caveat 1) and never run V2 NIXL with a container-default - `UCX_TLS=tcp` (breaks V2 NIXL VRAM registration) — unset/override it. - No bounce env override is needed: the byte gate - (`TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES`, default 2 MiB) is always cleared - by K3 payloads (constraints section above). + see caveat 1) and never run the Python transceiver with a + container-default `UCX_TLS=tcp` (breaks NIXL VRAM registration) — + unset/override it. Keep `cuda_ipc` in any pin; UCX uses it for + transfers over MNNVL. 5. **Transfer payload**: each request moves a fixed 433.4 MiB (~454.5 MB) KDA state blob ctx → gen in addition to the MLA latent KV - (~27 KB/token). - Within an NVL72 domain this is ~0.9 ms/request (measured; not a - bottleneck), but off-fabric paths would pay 11–23 ms — keep ctx and - gen inside one NVL72 domain. -6. **Bounce-buffer sizing cliff (silent).** Size `kv_cache_bounce_size_mb` - to the largest single request's full KV payload (fixed KDA state plus - the per-token MLA latent; ≥1024 MB for 8k ISL). An undersized region - does not error — every transfer silently falls back to a much slower - host-staged TCP path. + (~27 KB/token). Keep ctx and gen inside one NVL72 domain so these + writes can go over MNNVL. +6. **V1 cache manager.** `use_kv_cache_manager_v2: false`, or `"auto"` with + a text-only K3 checkpoint (`KimiLinearForCausalLM`, which declares no + preference), runs the V1 `MixedMambaHybridCacheManager`. Its MLA KV + pool still uses fabric memory where supported, but its KDA state is + plain PyTorch memory, so cross-node KDA writes cannot use cuda_ipc and + fall back to a much slower transport (host-staged TCP under the + caveat-4 pin). Use the released checkpoint with the default KV cache + manager. 7. **SA caps gen-side batch size.** SA requires `max_batch_size` ≤ 8 on the generation server, which bounds per-instance concurrency at `8 × dp_size` (128 with DEP16). Plan instance counts accordingly. @@ -180,5 +183,6 @@ python3 examples/disaggregated/slurm/benchmark/submit.py \ job time limits accordingly. The disaggregated proxy does not serve `/v1/models` (404) — point readiness probes at a different endpoint. 10. **`max_num_tokens` coupling.** The generation side must cover - `max_batch_size × (1 + max_draft_len)`; the context side needs - `max_tokens_in_buffer` ≥ max ISL (see constraints above). + `max_batch_size × (1 + max_draft_len)`; with chunked prefill disabled, + the context side must fit the whole prompt in `max_num_tokens` and + `max_seq_len`. diff --git a/examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml b/examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml index 9fca02194a62..748caed3d372 100644 --- a/examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml +++ b/examples/kimi_k3/disagg/benchmark_kimi_k3_dep16.yaml @@ -2,7 +2,7 @@ # examples/disaggregated/slurm/benchmark/submit.py, modeled on # examples/wide_ep/slurm_scripts/kimi-k2-thinking.yaml with K3's # constraints baked in (EP-only DEP16 on both sides, no chunked prefill, -# no KV block reuse, NIXL + Python V2 transceiver) plus the historically +# no KV block reuse, NIXL + Python transceiver) plus the historically # validated tokens_per_block 64. # Fill in the <...> placeholders for your cluster before submitting. # @@ -59,20 +59,15 @@ environment: work_dir: "" # Cluster notes: # - UCX transports: the default is no pin. start_worker.sh clears any - # container-provided UCX_TLS (UCX_TLS=tcp breaks V2 NIXL VRAM + # container-provided UCX_TLS (UCX_TLS=tcp breaks NIXL VRAM # registration) and UCX then selects transports itself, RDMA/verbs # included. Only on clusters where verbs transports cannot # initialize on the compute nodes (symptom: UCX_TLS=all wedges - # native NIXL init asymmetrically, surviving ranks hang in the V2 - # setup MPI collectives), add + # native NIXL init asymmetrically, surviving ranks hang in the + # transceiver setup MPI collectives), add # TRTLLM_WORKER_UCX_TLS=tcp,self,sm,cuda_copy,cuda_ipc to # worker_env_var below; start_worker.sh re-exports it as UCX_TLS # after the clear. Note that list excludes RDMA/verbs. - # - pairs with the workers' - # kv_cache_bounce_size_mb (fabric-VMM bounce; measured - # ~455 GB/s/GPU). The default gate of 96 blocks would silently skip - # bounce for requests under 6144 tokens at tokens_per_block=64, - # falling back to ~0.4 GB/s host-staged tcp. # - pyxis/enroot reset image-defined PATH at container start, so the # in-place venv is injected via TRTLLM_PATH_PREPEND / # TRTLLM_PYTHONPATH_PREPEND, which the harness scripts apply inside @@ -102,19 +97,13 @@ worker_config: free_gpu_memory_fraction: 0.25 cache_transceiver_config: backend: NIXL - transceiver_runtime: PYTHON # mandatory: C++ path throws for K3 + # Only the Python transceiver can move the KDA state. + transceiver_runtime: PYTHON + # Caps prompt tokens received at once (gen_only mode). max_tokens_in_buffer: 8448 - # Fabric-VMM bounce: needed for cross-node V2 transfers on clusters - # where inter-node cuda_ipc requires fabric allocations (else - # ~0.4 GB/s host-staged tcp); measured ~455 GB/s/GPU with - # (worker_env_var above). - # 1024 fits one full request payload (433 MiB KDA state + ~27 KB/token - # MLA latent = 649 MiB at 8k ISL); at 512 every 8k transfer falls back - # to the per-fragment tcp path. - kv_cache_bounce_size_mb: 1024 trust_remote_code: true - # CTX: matched DEP16 (see ../disagg/ctx_config.yaml for the sizing - # rationale; hetero ctx/gen parallelism is not supported yet). + # CTX: matched DEP16 (see ../disagg/ctx_config.yaml; hetero ctx/gen + # parallelism is not supported yet). ctx: tensor_parallel_size: 16 moe_expert_parallel_size: 16 @@ -133,8 +122,6 @@ worker_config: free_gpu_memory_fraction: 0.25 cache_transceiver_config: backend: NIXL + # Only the Python transceiver can move the KDA state. transceiver_runtime: PYTHON - max_tokens_in_buffer: 8448 - # Must match the gen side; see the gen block's comment. - kv_cache_bounce_size_mb: 1024 trust_remote_code: true diff --git a/examples/kimi_k3/disagg/ctx_config.yaml b/examples/kimi_k3/disagg/ctx_config.yaml index 4b6cf9231d5f..ead31882646b 100644 --- a/examples/kimi_k3/disagg/ctx_config.yaml +++ b/examples/kimi_k3/disagg/ctx_config.yaml @@ -45,26 +45,7 @@ kv_cache_config: free_gpu_memory_fraction: 0.25 cache_transceiver_config: backend: NIXL - # K3 requires the Python V2 transceiver: `auto` resolves to the C++ - # runtime, which throws at construction for K3's - # MixedMambaHybridCacheManager (LinearAttentionMetadata not found). + # Only the Python transceiver can move the KDA state. transceiver_runtime: PYTHON - # Must be >= the target max ISL (8192 here). - max_tokens_in_buffer: 8448 - # Fabric-VMM bounce buffer, needed for cross-node transfers on NVLink - # domains: V2 sends pool-to-pool by default, but the KV pool is a - # plain (non-fabric) allocation, so inter-node cuda_ipc/MNNVL is - # impossible and UCX falls back to ~0.4 GB/s host-staged tcp. 512 MiB - # measured at ~455 GB/s/GPU. Engages automatically for payloads above - # TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES (default 2 MiB) — always true for - # K3's ~433 MiB per-request recurrent state; no env override needed. - # - # Sizing: the region must fit ONE request's full payload = fixed KDA - # state (433 MiB) + MLA latent for the ISL (~27 KB/token, ~216 MiB at - # 8k). At 512 every 8k request is 649 MiB and silently falls back to - # the per-fragment tcp path ("[kv-bounce] in-place: transfer 649MiB - # exceeds the 512MiB bounce region"). 1024 - # covers the 8448-token max ISL with headroom. - kv_cache_bounce_size_mb: 1024 # No speculative_config on the ctx server: SA drafting runs on the # generation side only (ctx produces the prompt KV + first token). diff --git a/examples/kimi_k3/disagg/disagg_proxy_config.yaml b/examples/kimi_k3/disagg/disagg_proxy_config.yaml index b8e83983e694..7485c816191c 100644 --- a/examples/kimi_k3/disagg/disagg_proxy_config.yaml +++ b/examples/kimi_k3/disagg/disagg_proxy_config.yaml @@ -4,8 +4,8 @@ # Modeled on examples/disaggregated/disagg_config.yaml. The proxy routes # each request to a context server for prefill, then to a generation # server for decode; the KV + KDA state moves ctx -> gen through the -# cache transceiver configured in the worker configs (NIXL, Python V2 -# runtime - see ctx_config.yaml / gen_config_no_sa.yaml). +# cache transceiver configured in the worker configs (NIXL, Python +# transceiver - see ctx_config.yaml / gen_config_no_sa.yaml). # # Replace hostnames/ports with the actual worker endpoints (each K3 # worker spans 16 GPUs = 4 NVL72 nodes; the URL is rank 0's node). diff --git a/examples/kimi_k3/disagg/gen_config.yaml b/examples/kimi_k3/disagg/gen_config.yaml index 4fb137b12e67..a42eef435867 100644 --- a/examples/kimi_k3/disagg/gen_config.yaml +++ b/examples/kimi_k3/disagg/gen_config.yaml @@ -28,20 +28,8 @@ kv_cache_config: free_gpu_memory_fraction: 0.25 cache_transceiver_config: backend: NIXL - # Mandatory for K3: `auto` resolves to the C++ transceiver, which - # throws for K3's MixedMambaHybridCacheManager. + # Only the Python transceiver can move the KDA state. transceiver_runtime: PYTHON - # Must be >= the target max ISL and match the ctx side. - max_tokens_in_buffer: 8448 - # Fabric-VMM bounce buffer (recv side); must match the ctx server. - # Without it, cross-node V2 pool-to-pool transfers fall back to - # ~0.4 GB/s host-staged tcp; 512 MiB measured at ~455 GB/s/GPU. - # Bounce engages automatically above TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES - # (default 2 MiB) — always true for K3 payloads. - # Sized to fit one full request payload (fixed 433 MiB KDA state + - # ~27 KB/token MLA latent; 649 MiB at 8k ISL) -- see ctx_config.yaml; - # at 512 every 8k transfer falls back to the per-fragment tcp path. - kv_cache_bounce_size_mb: 1024 speculative_config: decoding_type: SA max_draft_len: 2 diff --git a/examples/kimi_k3/disagg/gen_config_no_sa.yaml b/examples/kimi_k3/disagg/gen_config_no_sa.yaml index 66a94d344f82..7427d719bb23 100644 --- a/examples/kimi_k3/disagg/gen_config_no_sa.yaml +++ b/examples/kimi_k3/disagg/gen_config_no_sa.yaml @@ -2,10 +2,8 @@ # options WITHOUT speculative decoding # (`trtllm-serve --config gen_config_no_sa.yaml`). # -# Default: CUDA graphs ON. Validated end-to-end: GSM8K 96.89 (in the -# aggregated band) and 765/2138 out tok/s at c64/c256 — ~3x the eager -# TPOT and above aggregated serving throughput despite the overlap -# scheduler being off (a disagg requirement on the ctx side). +# Default: CUDA graphs ON. Validated end-to-end on GB300 NVL72: GSM8K +# 96.5 and ~950/2043 output tok/s at c64/c256 on 8k/1k random prompts. # For strict token-parity debugging (attribution-clean SA-on/SA-off or # disagg-vs-agg comparisons), null out the cuda graph config to match # the SA config's eager regime, and place ctx and gen in SEPARATE NVL72 @@ -34,17 +32,5 @@ kv_cache_config: free_gpu_memory_fraction: 0.25 cache_transceiver_config: backend: NIXL - # Mandatory for K3: `auto` resolves to the C++ transceiver, which - # throws for K3's MixedMambaHybridCacheManager. + # Only the Python transceiver can move the KDA state. transceiver_runtime: PYTHON - max_tokens_in_buffer: 8448 - # Fabric-VMM bounce buffer (recv side); must match the ctx server. - # Without it, cross-node V2 pool-to-pool transfers fall back to - # ~0.4 GB/s host-staged tcp; 512 MiB measured at ~455 GB/s/GPU. - # Bounce engages automatically above TRTLLM_KV_CACHE_BOUNCE_MIN_BYTES - # (default 2 MiB) — always true for K3 payloads. - # at tokens_per_block=64). See ctx_config.yaml. - # Sized to fit one full request payload (fixed 433 MiB KDA state + - # ~27 KB/token MLA latent; 649 MiB at 8k ISL) -- see ctx_config.yaml; - # at 512 every 8k transfer falls back to the per-fragment tcp path. - kv_cache_bounce_size_mb: 1024 diff --git a/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml index 045688195443..28012157886b 100644 --- a/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/gb300_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml @@ -63,7 +63,7 @@ accuracy: worker_config: # K3 constraints applied to both sides (see examples/kimi_k3/disagg/README.md): # EP-only DEP16 with matched ctx/gen parallelism, no chunked prefill, no KV - # block reuse, tokens_per_block 64, NIXL with the Python V2 transceiver. + # block reuse, tokens_per_block 64, NIXL with the Python transceiver. gen: print_iter_log: true max_batch_size: 32 @@ -93,12 +93,8 @@ worker_config: backend: MEGAMOE_DEEPGEMM cache_transceiver_config: backend: NIXL - # PYTHON is mandatory: the C++ transceiver throws at construction for - # K3's MixedMambaHybridCacheManager. + # Only the Python transceiver can move the KDA state. transceiver_runtime: PYTHON - max_tokens_in_buffer: 8448 - # Shared C++ agent arena for chunked gather/write/scatter transfers. - kv_cache_bounce_size_mb: 1024 disable_overlap_scheduler: false trust_remote_code: true ctx: @@ -122,8 +118,7 @@ worker_config: backend: MEGAMOE_DEEPGEMM cache_transceiver_config: backend: NIXL + # Only the Python transceiver can move the KDA state. transceiver_runtime: PYTHON - max_tokens_in_buffer: 8448 - kv_cache_bounce_size_mb: 1024 disable_overlap_scheduler: true trust_remote_code: true diff --git a/tests/scripts/perf/disaggregated/vr200_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml b/tests/scripts/perf/disaggregated/vr200_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml index 5ebf776f2ca1..9ebbc3070b29 100644 --- a/tests/scripts/perf/disaggregated/vr200_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml +++ b/tests/scripts/perf/disaggregated/vr200_kimi-k3-fp4_8k1k_con512_ctx1_dep16_gen1_dep16_eplb0_mtp0_ccb-NIXL.yaml @@ -73,7 +73,7 @@ accuracy: worker_config: # K3 constraints applied to both sides (see examples/kimi_k3/disagg/README.md): # EP-only DEP16 with matched ctx/gen parallelism, no chunked prefill, no KV - # block reuse, tokens_per_block 64, NIXL with the Python V2 transceiver. + # block reuse, tokens_per_block 64, NIXL with the Python transceiver. gen: print_iter_log: true max_batch_size: 32 @@ -103,12 +103,8 @@ worker_config: backend: MEGAMOE_DEEPGEMM cache_transceiver_config: backend: NIXL - # PYTHON is mandatory: the C++ transceiver throws at construction for - # K3's MixedMambaHybridCacheManager. + # Only the Python transceiver can move the KDA state. transceiver_runtime: PYTHON - max_tokens_in_buffer: 8448 - # Shared C++ agent arena for chunked gather/write/scatter transfers. - kv_cache_bounce_size_mb: 1024 disable_overlap_scheduler: false trust_remote_code: true ctx: @@ -132,8 +128,7 @@ worker_config: backend: MEGAMOE_DEEPGEMM cache_transceiver_config: backend: NIXL + # Only the Python transceiver can move the KDA state. transceiver_runtime: PYTHON - max_tokens_in_buffer: 8448 - kv_cache_bounce_size_mb: 1024 disable_overlap_scheduler: true trust_remote_code: true From d37362acd7d54e6baab200c3b0dd768655c6ebcd Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:27:29 +0800 Subject: [PATCH 08/11] [None][fix] Keep older generation servers working unless they use the Python bounce Generation servers from before the Python bounce buffer was removed put bounce_dst_base in every REQUEST_DATA. RecvReqInfo.from_bytes passed it to the dataclass, which raised TypeError; the context server only logged it, and the generation server waited out kv_transfer_timeout_ms. Drop the key when it is None: that server reserved no bounce region. When it is set, the server expects the KV in its Python bounce region, which this release no longer writes: raise LegacyBounceRequestError and answer the request with a FAILED KV result, so the generation server fails it at once. The context logs this once per generation server and each rejected request at DEBUG. Other unknown keys still fail decoding. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../_torch/disaggregation/native/transfer.py | 42 +++++++++++- .../disaggregated/test_unsubmitted_session.py | 66 ++++++++++++++++++- 2 files changed, 103 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index bcecbd2eaee7..ef88fde3bad7 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -101,6 +101,25 @@ class _TransferNotSubmittedError(RuntimeError): """The sender rejected an operation before the backend could access memory.""" +# Sent by generation servers from before the Python KV bounce buffer was removed: None, or the base +# of a bounce region the server reserved for this request and expects the KV in. +_LEGACY_BOUNCE_DST_BASE_KEY = "bounce_dst_base" + + +class LegacyBounceRequestError(ValueError): + """REQUEST_DATA from an older generation server that reserved a Python KV bounce region for the request.""" + + def __init__(self, req_info: RecvReqInfo): + super().__init__( + f"Failing request {req_info.unique_rid}: generation server {req_info.instance_name} " + f"rank {req_info.instance_rank} runs an older TensorRT-LLM release that expects the KV " + "cache in its Python bounce buffer, and this context server no longer writes there. " + "Run both servers on the same release, or set kv_cache_bounce_size_mb=0 on that " + "generation server." + ) + self.req_info = req_info + + @dataclass class RecvReqInfo: sender_req_id: int @@ -138,7 +157,11 @@ def from_bytes(cls, data: bytes) -> "RecvReqInfo": d["block_ids_per_layer_groups"] = [ np.frombuffer(b, dtype=np.int64).copy() for b in d["block_ids_per_layer_groups"] ] - return cls(**d) + legacy_bounce_dst_base = d.pop(_LEGACY_BOUNCE_DST_BASE_KEY, None) + info = cls(**d) + if legacy_bounce_dst_base is not None: + raise LegacyBounceRequestError(info) + return info @dataclass @@ -2003,7 +2026,22 @@ def _send_session_quiesced(self, marker: _SessionQuiescence) -> bool: def _respond_with_kv(self, _send_id: bytes, message: list[bytes]): # _sessions_lock prevents a race between session lookup and req_info save. # session.lock atomically saves peer info and snapshots tasks against send(). - info: RecvReqInfo = RecvReqInfo.from_bytes(message[1]) + try: + info: RecvReqInfo = RecvReqInfo.from_bytes(message[1]) + except LegacyBounceRequestError as error: + # Answer instead of dropping, or the generation server waits out kv_transfer_timeout_ms. + info = error.req_info + logger.error_once( + f"Failing every request from generation server {info.instance_name} rank " + f"{info.instance_rank}: it runs an older TensorRT-LLM release that expects the KV " + "cache in its Python bounce buffer, and this context server no longer writes there. " + "Run both servers on the same release, or set kv_cache_bounce_size_mb=0 on that " + "generation server. Each rejected request is logged at DEBUG.", + key=("legacy_bounce_request", info.instance_name, info.instance_rank), + ) + logger.debug(str(error)) + self._send_failed_result_to_receiver(info) + return with self._sessions_lock: session = self._get_session(info.unique_rid) fenced = ( diff --git a/tests/unittest/disaggregated/test_unsubmitted_session.py b/tests/unittest/disaggregated/test_unsubmitted_session.py index fbda04972923..52100efb2343 100644 --- a/tests/unittest/disaggregated/test_unsubmitted_session.py +++ b/tests/unittest/disaggregated/test_unsubmitted_session.py @@ -19,6 +19,7 @@ from types import SimpleNamespace from unittest.mock import Mock +import msgpack import numpy as np import pytest @@ -70,8 +71,8 @@ def _cancel(*, routing: bool = False) -> list[bytes]: return message + [b"gen", b"0"] if routing else message -def _request_data() -> list[bytes]: - info = transfer_mod.RecvReqInfo( +def _request_info() -> transfer_mod.RecvReqInfo: + return transfer_mod.RecvReqInfo( sender_req_id=401, instance_name="gen", instance_rank=0, @@ -80,7 +81,17 @@ def _request_data() -> list[bytes]: aux_slot=3, slice_id=0, ) - return [transfer_mod.MessageType.REQUEST_DATA, info.to_bytes()] + + +def _request_data() -> list[bytes]: + return [transfer_mod.MessageType.REQUEST_DATA, _request_info().to_bytes()] + + +def _legacy_request_info_bytes(bounce_dst_base) -> bytes: + """REQUEST_DATA payload of a generation server from before the Python bounce was removed.""" + fields = msgpack.unpackb(_request_info().to_bytes(), raw=False) + fields["bounce_dst_base"] = bounce_dst_base + return msgpack.packb(fields) def _assert_session_ack(sender: transfer_mod.Sender) -> None: @@ -556,3 +567,52 @@ def test_nonblocking_gen_progress_notifies_unsettled_senders_once(terminal: str) assert session.resources_drained() assert session.close() watchdog.stop() + + +def test_legacy_request_without_a_bounce_region_decodes() -> None: + info = transfer_mod.RecvReqInfo.from_bytes(_legacy_request_info_bytes(None)) + expected = _request_info() + assert (info.unique_rid, info.instance_name, info.instance_rank) == (401, "gen", 0) + assert (info.aux_slot, info.slice_id) == (expected.aux_slot, expected.slice_id) + + +def test_legacy_request_with_a_bounce_region_names_the_version_mismatch() -> None: + with pytest.raises( + transfer_mod.LegacyBounceRequestError, match="older TensorRT-LLM release" + ) as error: + transfer_mod.RecvReqInfo.from_bytes(_legacy_request_info_bytes(0x7F0000000000)) + assert error.value.req_info.unique_rid == 401 + + +@pytest.mark.parametrize("ownership", [False, True]) +def test_request_data_from_a_bouncing_legacy_generation_server_fails_it_at_once( + ownership: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + sender = _sender(ownership=ownership) + sender._get_or_connect_dealer = Mock() + # error_once deduplicates against the process-wide key set; start from an empty one. + monkeypatch.setattr(transfer_mod.logger, "_appeared_keys", set()) + log = Mock() + monkeypatch.setattr(transfer_mod.logger, "log", log) + legacy_request = [ + transfer_mod.MessageType.REQUEST_DATA, + _legacy_request_info_bytes(0x7F0000000000), + ] + + for _ in range(2): + sender._respond_with_kv(b"gen", legacy_request) + + failed = transfer_mod._make_kv_result_msg(7, 401, 0, True, transfer_mod.AgentResult.FAILED) + if ownership: + sent = [sender._send_task_queues[0].get_nowait() for _ in range(2)] + assert sent == [("tcp://gen:1234", failed)] * 2 + sender._get_or_connect_dealer.assert_not_called() + else: + dealer = sender._get_or_connect_dealer.return_value + assert [call.args for call in dealer.send.call_args_list] == [(failed,)] * 2 + assert sender._get_req_info(401) is None + errors = [ + call.args[1] for call in log.call_args_list if call.args[0] == transfer_mod.logger.ERROR + ] + assert len(errors) == 1 + assert errors[0].startswith("Failing every request from generation server gen rank 0") From 254ee3d42d99f5f4dca30d51eceea398fa2cad66 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:34:24 +0800 Subject: [PATCH 09/11] [None][chore] Fail an incompatible-peer receive task under its session lock KVRecvTask.complete() and fail() rely on the caller holding RxSession.lock, but dispatch_task's PeerIncompatibleError branch called fail() without it. No complete() could race there (REQUEST_DATA is not sent yet), so the rule was only documented. Take the lock there too when the session is still registered. If it is already closed, fail the task directly as before: no listener can reach a task of an unregistered session. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../_torch/disaggregation/native/transfer.py | 25 +++++---- .../disaggregated/test_task_handle.py | 53 +++++++++++++++++++ 2 files changed, 68 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/disaggregation/native/transfer.py b/tensorrt_llm/_torch/disaggregation/native/transfer.py index ef88fde3bad7..ac19741ed7df 100644 --- a/tensorrt_llm/_torch/disaggregation/native/transfer.py +++ b/tensorrt_llm/_torch/disaggregation/native/transfer.py @@ -2809,8 +2809,8 @@ def __init__( def fail(self, exc: Exception) -> None: """Mark the task ERROR. - Callers hold RxSession.lock, except dispatch_task's PeerIncompatibleError path, which - runs before publication. + Callers hold the owning RxSession.lock, unless that session is already closed and + unregistered (then no other thread can reach the task). """ self._logical_outcomes.fail(exc) self._exception = exc @@ -2820,8 +2820,9 @@ def fail(self, exc: Exception) -> None: def complete(self) -> None: """Mark the task TRANSFERRED unless it already failed. - Callers hold RxSession.lock, so this check-then-set is atomic against fail(). Status is - set before _event, so a thread woken by _event reads the status that complete() set. + Callers hold the owning RxSession.lock, which makes this check-then-set atomic against + fail(). Status is set before _event, so a thread woken by _event reads the status that + complete() set. """ if self.status == TaskStatus.ERROR: return @@ -3118,13 +3119,17 @@ def dispatch_task(self, task: KVRecvTask) -> None: "dispatch_task: context peer incompatible, failing request " f"unique_rid={task._unique_rid}: {e}" ) - if self._enforce_physical_ownership: - session = self._get_session(task._unique_rid) - if session is None: + session = self._get_session(task._unique_rid) + if session is None: + # Closed and unregistered: no listener or executor path can reach this task any more. + if self._enforce_physical_ownership: task.cancel_unpublished() - else: - session.cancel_unpublished_task(task) - task.fail(e) + task.fail(e) + return + if self._enforce_physical_ownership: + session.cancel_unpublished_task(task) + with session.lock: + task.fail(e) return if sender_dp_rank is not None: diff --git a/tests/unittest/disaggregated/test_task_handle.py b/tests/unittest/disaggregated/test_task_handle.py index 24cb56ab86b1..cf0cabdf2f8e 100644 --- a/tests/unittest/disaggregated/test_task_handle.py +++ b/tests/unittest/disaggregated/test_task_handle.py @@ -44,6 +44,7 @@ KVRecvTask, KVSendTask, MessageType, + PeerIncompatibleError, Receiver, RxSession, Sender, @@ -628,6 +629,58 @@ def test_a_fetch_the_cancel_beat_reports_the_peers_cancel_rather_than_a_failure( assert _owes_a_report(outcome) is False +def _receiver_with_an_incompatible_context_peer(owns_transfers: bool) -> Receiver: + """A receiver that dispatches for real, against a context peer whose layout it rejects.""" + receiver = _wired_receiver(owns_transfers) + del receiver.dispatch_task + receiver._registrar = SimpleNamespace( + self_rank_info=SimpleNamespace(instance_name="gen", instance_rank=0) + ) + receiver._get_sender_info = MagicMock( + side_effect=PeerIncompatibleError("synthetic layout mismatch") + ) + return receiver + + +@pytest.mark.parametrize("owns_transfers", [False, True]) +def test_an_incompatible_peer_fails_the_task_under_the_session_lock( + owns_transfers: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + """KVRecvTask.fail() relies on its session's lock, also on the path that never publishes.""" + receiver = _receiver_with_an_incompatible_context_peer(owns_transfers) + session = _receiving_over(receiver, 21) + lock_held_on_fail = [] + fail = KVRecvTask.fail + + def fail_recording_the_lock(task: KVRecvTask, exc: Exception) -> None: + lock_held_on_fail.append(session.lock.locked()) + fail(task, exc) + + monkeypatch.setattr(KVRecvTask, "fail", fail_recording_the_lock) + + session.receive(_sole_piece()) + + assert lock_held_on_fail == [True] + assert session.has_failed() + if owns_transfers: + # No writer was published, so the owner closed without waiting for one. + assert session._kv_tasks[0].resources_drained + + +@pytest.mark.parametrize("owns_transfers", [False, True]) +def test_an_incompatible_peer_fails_a_task_of_a_closed_session_without_raising( + owns_transfers: bool, +) -> None: + """Once the session is unregistered nothing else can reach the task, so no lock is needed.""" + receiver = _receiver_with_an_incompatible_context_peer(owns_transfers) + session = _receiving_over(receiver, 22) + receiver._sessions.pop(22) + + session.receive(_sole_piece()) + + assert session._kv_tasks[0].status is TaskStatus.ERROR + + def test_a_send_session_reads_a_parked_cancel_as_the_peers_too(): """The other direction parks cancels the same way, and loses the same thing.""" sender = _wired_sender() From d3312f255b71c97e632382e53c002fc0a1a62a70 Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 12 Oct 2026 15:36:37 +0800 Subject: [PATCH 10/11] [None][test] Cover pipelined KV transfer through the C++ agent bounce enable_pipelined_transfer is now allowed with kv_cache_bounce_size_mb > 0, but no GPU test sent a request in pieces. Add v2_mha_pipelined to test_python_nixl_cache_transceiver_uses_cpp_bounce: the context sends each prompt in 2-block pieces, the generation KV is checked byte for byte, and the context agent must route every piece through the bounce and decline none. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- .../test_cache_transceiver_single_process.py | 50 +++++++++++++++++-- 1 file changed, 45 insertions(+), 5 deletions(-) diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index e420884ce6d2..6509950fa499 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py @@ -1213,6 +1213,23 @@ def _verify_indexer_k_all_requests( # --------------------------------------------------------------------------- # Main test orchestrator # --------------------------------------------------------------------------- +def _send_prefill_in_pieces(transceiver, req: LlmRequest, piece_tokens: int) -> int: + """Send the prompt's KV in pieces of piece_tokens and return the number of pieces sent. + + One context chunk per piece, as the executor does with enable_pipelined_transfer. + """ + # A non-final piece without a complete block is not sent, which would break the count. + assert piece_tokens % TOKENS_PER_BLOCK == 0 + pieces = 0 + for start in range(0, req.prompt_len, piece_tokens): + end = min(start + piece_tokens, req.prompt_len) + req.py_last_context_chunk = (start, end) + req.context_current_position = end + transceiver.respond_and_send_async(req) + pieces += 1 + return pieces + + def run_transfer_test( ctx_tp: int, ctx_pp: int, @@ -1230,9 +1247,13 @@ def run_transfer_test( expect_cpp_bounce: bool = False, gen_scratch_tokens: int = 0, gen_cp: int = 1, + pipelined_piece_tokens: Optional[int] = None, ) -> None: """Run a full KV transfer test using KvCacheTransceiverV2. + pipelined_piece_tokens enables enable_pipelined_transfer: the context sends each prompt in + pieces of that many tokens instead of whole. + gen_scratch_tokens > 0 sizes the generation cache to prompt_len + gen_scratch_tokens (history still prompt_len), materialising an MTP speculative tail block past the prompt. Combined with a sliding window @@ -1305,6 +1326,7 @@ def run_transfer_test( } if expect_cpp_bounce else None, + enable_pipelined_transfer=pipelined_piece_tokens is not None, ) ctx_tcs = create_instance_transceivers( ctx_tp, ctx_pp, ctx_enable_dp, ctx_managers, config, is_mla @@ -1429,12 +1451,18 @@ def run_transfer_test( gen_tcs[rank].request_and_receive_async(req) # 6. ctx send after + pieces_sent = 0 for rank in range(ctx_world): for req_idx, req in ctx_handle_map[rank]: kv = _add_sequence(ctx_managers[rank], req.py_request_id, req.prompt_len, use_v2) if kv is not None: ctx_kv_caches[rank].append(kv) - ctx_tcs[rank].respond_and_send_async(req) + if pipelined_piece_tokens is None: + ctx_tcs[rank].respond_and_send_async(req) + else: + pieces_sent += _send_prefill_in_pieces( + ctx_tcs[rank], req, pipelined_piece_tokens + ) # 7. Wait for completion (threaded, dist calls inside) run_concurrent( @@ -1500,6 +1528,14 @@ def run_transfer_test( ) total_bounce_submits = sum(agent.bounce_submit_count for agent in agents) assert total_bounce_submits > 0, "no transfer was routed through the bounce fast path" + if pipelined_piece_tokens is not None: + # Exact for one context rank writing one generation rank: each piece is one + # transfer request, and every one of them must take the bounce. + ctx_agents = [tc._transfer_worker._agent for tc in ctx_tcs] + assert all(agent.bounce_reject_count == 0 for agent in ctx_agents), [ + agent.bounce_reject_counts for agent in ctx_agents + ] + assert sum(agent.bounce_submit_count for agent in ctx_agents) == pieces_sent finally: # Teardown must run on the failure path too, in this order: @@ -1987,12 +2023,14 @@ def test_cache_transceiver_v1_masked_dsa_indexer_across_asymmetric_pp() -> None: "is_mla", "use_v2", "enable_indexer_k_cache", + "pipelined_piece_tokens", ), [ - pytest.param(1, 1, 1, 1, False, True, False, id="v2_mha"), - pytest.param(2, 1, 2, 1, True, True, False, id="v2_mla_tp2"), - pytest.param(1, 1, 1, 1, True, False, True, id="v1_dsa_indexer"), - pytest.param(2, 1, 1, 2, False, True, False, id="v2_ctx_tp2_gen_pp2"), + pytest.param(1, 1, 1, 1, False, True, False, None, id="v2_mha"), + pytest.param(2, 1, 2, 1, True, True, False, None, id="v2_mla_tp2"), + pytest.param(1, 1, 1, 1, True, False, True, None, id="v1_dsa_indexer"), + pytest.param(2, 1, 1, 2, False, True, False, None, id="v2_ctx_tp2_gen_pp2"), + pytest.param(1, 1, 1, 1, False, True, False, 2 * TOKENS_PER_BLOCK, id="v2_mha_pipelined"), ], ) def test_python_nixl_cache_transceiver_uses_cpp_bounce( @@ -2003,6 +2041,7 @@ def test_python_nixl_cache_transceiver_uses_cpp_bounce( is_mla: bool, use_v2: bool, enable_indexer_k_cache: bool, + pipelined_piece_tokens: Optional[int], ) -> None: """Exercise C++ bounce v2 through representative Python NIXL transceiver paths. @@ -2028,6 +2067,7 @@ def test_python_nixl_cache_transceiver_uses_cpp_bounce( request_lengths=[30, 60], enable_indexer_k_cache=enable_indexer_k_cache, expect_cpp_bounce=True, + pipelined_piece_tokens=pipelined_piece_tokens, ) From 6f33685d262c672ebd93dcd79c6e13101a32b58c Mon Sep 17 00:00:00 2001 From: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> Date: Mon, 19 Oct 2026 04:21:28 +0000 Subject: [PATCH 11/11] [None][chore] Relax the C++ bounce gate for the Qwen3.5 GB300 agentic configs #20514's colocated 1P1D Qwen3.5 GB300 agentic configs set kv_cache_bounce_size_mb: 2048 without agent_bounce_buffer_enable, which selected the Python bounce buffer. With this PR the size selects the C++ transfer-agent bounce buffer instead. These profiles send KV from a TP1 context server to a TP2 generation server, so each attention KV descriptor holds one head of one block (16 KiB with fp8 KV, 64 tokens per block and head_dim 256): the many-small-descriptor case the bounce buffer is for. Each write also carries the GDN recurrent state, whose per-layer descriptors are much larger, so the write's average descriptor size exceeds the default 16 KiB limit and the default gate would send the whole write through standard NIXL. Short prompts also fall below the default minimum of 1024 descriptors. Relax the gate with agent_bounce_params (min_descriptor_count "1", max_average_descriptor_size "4MB") so these writes go through the bounce buffer, as they did with the Python bounce. The deployment guide names both settings. These profiles were not re-measured with the C++ bounce. Signed-off-by: Chuang Zhu <111838961+chuangz0@users.noreply.github.com> --- ...deployment-guide-for-qwen3.8-qwen3.5-on-trtllm.md | 12 ++++++++---- .../curated/qwen3.5-gb300-agentic-decode.yaml | 3 +++ .../curated/qwen3.5-gb300-agentic-prefill.yaml | 3 +++ 3 files changed, 14 insertions(+), 4 deletions(-) diff --git a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-qwen3.5-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-qwen3.5-on-trtllm.md index 98acba0ea8fe..495421861921 100644 --- a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-qwen3.5-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-qwen3.5-on-trtllm.md @@ -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 @@ -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 diff --git a/examples/configs/curated/qwen3.5-gb300-agentic-decode.yaml b/examples/configs/curated/qwen3.5-gb300-agentic-decode.yaml index 1e02ec5244b2..994ad10853fe 100644 --- a/examples/configs/curated/qwen3.5-gb300-agentic-decode.yaml +++ b/examples/configs/curated/qwen3.5-gb300-agentic-decode.yaml @@ -7,6 +7,9 @@ cache_transceiver_config: kv_transfer_timeout_ms: 600000 transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 2048 + agent_bounce_params: + min_descriptor_count: "1" + max_average_descriptor_size: "4MB" cuda_graph_config: enable_padding: true batch_sizes: diff --git a/examples/configs/curated/qwen3.5-gb300-agentic-prefill.yaml b/examples/configs/curated/qwen3.5-gb300-agentic-prefill.yaml index 928340c5a0fe..cf5b243791f7 100644 --- a/examples/configs/curated/qwen3.5-gb300-agentic-prefill.yaml +++ b/examples/configs/curated/qwen3.5-gb300-agentic-prefill.yaml @@ -9,6 +9,9 @@ cache_transceiver_config: kv_transfer_timeout_ms: 600000 transceiver_runtime: PYTHON kv_cache_bounce_size_mb: 2048 + agent_bounce_params: + min_descriptor_count: "1" + max_average_descriptor_size: "4MB" disable_overlap_scheduler: false enable_attention_dp: true enable_chunked_prefill: true