DRAFT (DO NOT MERGE) [TRTLLM-12498][feat] Add support for beam search in disaggregated serving - #14470
Closed
athena-nv wants to merge 20 commits into
Closed
DRAFT (DO NOT MERGE) [TRTLLM-12498][feat] Add support for beam search in disaggregated serving#14470athena-nv wants to merge 20 commits into
athena-nv wants to merge 20 commits into
Conversation
…saggregated serving
athena-nv
requested review from
QiJune,
bo-nv,
kaiyux,
leslie-fang25 and
reasonsolo
May 22, 2026 19:42
Contributor
📝 WalkthroughWalkthroughThis PR implements beam search decoding support in disaggregated TensorRT-LLM serving. The changes remove KV cache formatting restrictions on beam width, propagate beam-structured block IDs through the cache-reuse and transceiver stack, and thread first-generation cumulative log probabilities from context to generation servers. ChangesBeam search in disaggregated serving
Design documentation and examples
🎯 4 (Complex) | ⏱️ ~60 minutes Suggested reviewers
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (1 warning, 1 inconclusive)
✅ Passed checks (3 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
⚔️ Resolve merge conflicts
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 14
Note
Due to the large number of review comments, Critical, Major severity comments were prioritized as inline comments.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/serve/openai_disagg_service.py (2)
343-376:⚠️ Potential issue | 🟠 Major | ⚡ Quick winMove these request-path dumps behind debug logging.
These
logger.infocalls emit per-request beam metadata, request IDs, andctx_info_endpointon every disaggregated request. That is a hot path, and it also exposes internal routing details in production logs. Please gate this behind debug/feature-flag logging or remove it before merge.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/serve/openai_disagg_service.py` around lines 343 - 376, The per-request logger.info calls in _get_gen_request that log disagg_request_id, ctx_info_endpoint and beam/request metadata should not run at INFO level on the hot path; change them to be guarded debug/feature-flag logs. Replace the two logger.info invocations with either logger.debug(...) or wrap them in an if logger.isEnabledFor(logging.DEBUG): block (or check a feature flag like self._verbose_logging) so that the messages from logger.info(...) in _get_gen_request (and any logging of request.disaggregated_params.ctx_info_endpoint, ctx_response.choices, CompletionRequest/ChatCompletionRequest prompt token fields, and DisaggregatedParams creation) are only emitted when debug/verbose logging is enabled.
532-559:⚠️ Potential issue | 🟠 Major | 🏗️ Heavy liftMulti-choice ctx responses are still reduced to
choices[0].Commenting out the
len(ctx_response.choices) != 1check changes the contract, but the rest of this flow still validates and forwards onlychoices[0](_verify_ctx_response,_need_gen,_get_gen_request). If beam search can produce multiple context choices here, the extra choices' disaggregation state and finish reasons are silently ignored.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/serve/openai_disagg_service.py` around lines 532 - 559, The code currently drops all but choices[0] even when ctx_response.choices has multiple entries; update the flow to handle multi-choice context responses instead of reducing to the first choice: in the validation routine (where you currently index into choice = ctx_response.choices[0]) iterate over ctx_response.choices and validate each choice.disaggregated_params, ctx_request_id and disagg_request_id (raising with contextual finish_reason and disagg/ctx ids for that specific choice if invalid), and then either (A) return the full ctx_response with all validated choices or (B) change the downstream callsites (_verify_ctx_response, _need_gen, _get_gen_request) to accept and process a list of choices (e.g., generate per-choice gen requests) so no disaggregation state or finish_reason is silently ignored. Ensure all places that previously assumed a single choice are updated to consume the list or handle branching per choice.
🟡 Minor comments (8)
examples/disaggregated/simpler_example/kill.sh-3-7 (1)
3-7:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winGuard
kill -9when no matching process exists.On Line 7,
kill -9runs even whenpgrepfinds nothing, which can fail the script path during normal “nothing to kill” runs.Suggested fix
- export CTX_PROCESS_ID=$(pgrep -f "trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8001 --config ./ctx_config.yaml") - export GEN_PROCESS_ID=$(pgrep -f "trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8002 --config ./gen_config.yaml") - export DISAGG_PROCESS_ID=$(pgrep -f "trtllm-serve disaggregated -c ./disagg_config.yaml") + CTX_PROCESS_ID="$(pgrep -f "trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8001 --config ./ctx_config.yaml" || true)" + GEN_PROCESS_ID="$(pgrep -f "trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8002 --config ./gen_config.yaml" || true)" + DISAGG_PROCESS_ID="$(pgrep -f "trtllm-serve disaggregated -c ./disagg_config.yaml" || true)" -kill -9 $CTX_PROCESS_ID $GEN_PROCESS_ID $DISAGG_PROCESS_ID +PIDS=() +[[ -n "${CTX_PROCESS_ID}" ]] && PIDS+=(${CTX_PROCESS_ID}) +[[ -n "${GEN_PROCESS_ID}" ]] && PIDS+=(${GEN_PROCESS_ID}) +[[ -n "${DISAGG_PROCESS_ID}" ]] && PIDS+=(${DISAGG_PROCESS_ID}) +(( ${`#PIDS`[@]} > 0 )) && kill -9 "${PIDS[@]}"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/disaggregated/simpler_example/kill.sh` around lines 3 - 7, The kill -9 invocation should be guarded so it only runs for PIDs that exist: check each captured variable (CTX_PROCESS_ID, GEN_PROCESS_ID, DISAGG_PROCESS_ID) for non-empty before calling kill, and invoke kill only with the subset of PIDs that are present (or skip if none). Update the script around the pgrep assignments and the kill line to build a list of non-empty PID variables (or test each with [ -n "$VAR" ]), then call kill on that list or do nothing if the list is empty; ensure this logic replaces the unconditional kill -9 $CTX_PROCESS_ID $GEN_PROCESS_ID $DISAGG_PROCESS_ID.examples/disaggregated/simpler_example/kill_agg.sh-3-4 (1)
3-4:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winHandle empty PID before force-kill.
On Line 4, this can fail when the process is already stopped. Make the script no-op in that case.
Suggested fix
-export AGG_PROCESS_ID=$(pgrep -f "trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8000 --config ./agg_config.yaml") -kill -9 $AGG_PROCESS_ID +AGG_PROCESS_ID="$(pgrep -f "trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8000 --config ./agg_config.yaml" || true)" +[[ -n "${AGG_PROCESS_ID}" ]] && kill -9 ${AGG_PROCESS_ID}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/disaggregated/simpler_example/kill_agg.sh` around lines 3 - 4, The script currently assigns AGG_PROCESS_ID with pgrep and always runs kill -9 $AGG_PROCESS_ID which will fail if no PID was found; update the script to check whether AGG_PROCESS_ID is non-empty (e.g., test -n "$AGG_PROCESS_ID" or [[ -n "$AGG_PROCESS_ID" ]]) before calling kill -9, and only invoke kill when AGG_PROCESS_ID contains a PID (ensure you quote the variable when used).examples/disaggregated/simpler_example/output_agg.json-1-4 (1)
1-4:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winKeep this artifact as valid JSON only.
Lines 1-3 make the file invalid JSON, which breaks tooling and parsers expecting a JSON document from
output_agg.json.Suggested fix
- % Total % Received % Xferd Average Speed Time Time Time Current - Dload Upload Total Spent Left Speed -...curl progress output... {"id":"cmpl-5cb435186c4542339c3168b176d51b08", ...}🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/disaggregated/simpler_example/output_agg.json` around lines 1 - 4, The file output_agg.json contains non-JSON progress lines before the JSON object which break parsers; remove the leading curl/progress text lines (the three lines starting with "% Total" and the two progress rows) so the file begins directly with the JSON object starting with {"id":...}; ensure no other stray characters precede or follow the JSON and that the remaining content is a single well-formed JSON document.examples/disaggregated/simpler_example/run.sh-6-10 (1)
6-10:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winGuard process termination when no matching server is running.
Lines 6-10 unconditionally call
kill -9with possibly empty PID vars. This is brittle and can fail depending on how the script is invoked.Suggested fix
-export CTX_PROCESS_ID=$(pgrep -f "trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8001 --config ./ctx_config.yaml") -export GEN_PROCESS_ID=$(pgrep -f "trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8002 --config ./gen_config.yaml") -export DISAGG_PROCESS_ID=$(pgrep -f "trtllm-serve disaggregated -c ./disagg_config.yaml") +CTX_PROCESS_ID=$(pgrep -f "trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8001 --config ./ctx_config.yaml" || true) +GEN_PROCESS_ID=$(pgrep -f "trtllm-serve TinyLlama/TinyLlama-1.1B-Chat-v1.0 --host localhost --port 8002 --config ./gen_config.yaml" || true) +DISAGG_PROCESS_ID=$(pgrep -f "trtllm-serve disaggregated -c ./disagg_config.yaml" || true) -kill -9 $CTX_PROCESS_ID $GEN_PROCESS_ID $DISAGG_PROCESS_ID +if [ -n "${CTX_PROCESS_ID}" ] || [ -n "${GEN_PROCESS_ID}" ] || [ -n "${DISAGG_PROCESS_ID}" ]; then + kill -9 ${CTX_PROCESS_ID} ${GEN_PROCESS_ID} ${DISAGG_PROCESS_ID} +fi🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@examples/disaggregated/simpler_example/run.sh` around lines 6 - 10, The script unconditionally runs kill -9 on variables CTX_PROCESS_ID, GEN_PROCESS_ID, DISAGG_PROCESS_ID which may be empty; update the shutdown logic to check each PID variable (CTX_PROCESS_ID, GEN_PROCESS_ID, DISAGG_PROCESS_ID) is non-empty (and optionally verify the PID is live) before calling kill, e.g. guard each kill with a conditional or loop over the three variables and only invoke kill for present PIDs, quoting variables to avoid word-splitting and handling multiple PIDs safely.implementation_plan.md-17-17 (1)
17-17:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winFix typo in a key implementation step (
demension→dimension).🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@implementation_plan.md` at line 17, Fix the typo "demension" → "dimension" in the implementation plan and update the code so beam_width is made accessible in _create_kv_slice and passed into get_block_ids and get_batch_cache_indices; remove any logic that flattens the beam dimension and delete the assertion that beam dimension == 1 in get_batch_cache_indices; update indexing in get_batch_cache_indices (and any related block index handling) to treat beam_width as an extra dimension in the block indices tensor so the functions correctly handle beam_width > 1.tensorrt_llm/_torch/pyexecutor/py_executor.py-3594-3601 (1)
3594-3601:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winFix the Flake8 indentation error in
_debug_request_tokens.Line 3601 currently trips
E126, so this helper will fail lint as written.As per coding guidelines `Python code should be indented with 4 spaces; do not use tabs`.Suggested fix
`@staticmethod` def _debug_request_tokens(req, beam_width, max_tokens=16): - return [{ - "beam": beam, - "num_tokens": len(tokens), - "tail": tokens[-max_tokens:], - } for beam in range(beam_width) - for tokens in [list(req.get_tokens(beam))]] + return [ + { + "beam": beam, + "num_tokens": len(tokens), + "tail": tokens[-max_tokens:], + } + for beam in range(beam_width) + for tokens in [list(req.get_tokens(beam))] + ]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 3594 - 3601, The list comprehension in _debug_request_tokens is mis-indented and triggers Flake8 E126; fix it by reformatting the comprehension so the continuation lines use 4-space indentation and align logically with the opening bracket — ensure the inner "for tokens in [list(req.get_tokens(beam))]" and the "for beam in range(beam_width)" parts are placed on their own properly indented lines under the opening "[" within the _debug_request_tokens staticmethod to satisfy PEP8/Flake8.tensorrt_llm/_torch/pyexecutor/resource_manager.py-1215-1274 (1)
1215-1274:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winMake both
zip()calls strict.These loops assume a 1:1 mapping between
request_ids,result, and the shaped result. Withoutstrict=True, a binding-side length mismatch truncates silently and the logged layout no longer matches the returned data.Suggested fix
- for req_id, req_blocks in zip(request_ids, result): + for req_id, req_blocks in zip(request_ids, result, strict=True): @@ - for req_id, shaped_blocks in zip(request_ids, result): + for req_id, shaped_blocks in zip(request_ids, result, strict=True):🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` around lines 1215 - 1274, The two loops using zip(request_ids, result) in get_batch_cache_indices should require equal lengths to avoid silent truncation; update both occurrences (the first loop starting with "for req_id, req_blocks in zip(request_ids, result):" and the later loop "for req_id, shaped_blocks in zip(request_ids, result):") to use zip(request_ids, result, strict=True) so a length mismatch raises an error and the logged token_layout/shaped_layout stay consistent with returned data.tensorrt_llm/_torch/disaggregation/native/transfer.py-728-748 (1)
728-748:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winReplace transfer-metadata
asserts with explicit exceptions in_build_kv_write_meta()
Theasserts guardingsrc_beam_block_ids.size/dst_beam_block_ids.sizeand the SWAtask._prompt_len is not Nonecheck protectsrc_start/dst_startcomputation and_align_kv_blocks()slicing before pointer extraction; replace them withValueError/RuntimeErrorso failures remain deterministic even if Python is run with assertions disabled.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/disaggregation/native/transfer.py` around lines 728 - 748, In _build_kv_write_meta(), replace the runtime assertions that guard src_beam_block_ids.size and dst_beam_block_ids.size and the SWA prompt-length check (currently using assert on task._prompt_len) with explicit exceptions: raise ValueError with clear messages when src_beam_block_ids.size or dst_beam_block_ids.size exceed total_blocks, and raise RuntimeError (or ValueError) when task._prompt_len is None; keep the same diagnostic text (including beam_idx, slice_end, tpb) and ensure these checks occur before computing src_start/dst_start and before calling _align_kv_blocks() so failures remain deterministic even when Python assertions are disabled.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/pyexecutor/resource_manager.py (1)
1196-1201: ⚡ Quick winFix the return type for the beam-aware shape.
When
beam_width > 1, this returns per-request per-beam block ID lists, notList[List[int]]. The current signature makes stale single-beam call sites look type-safe even when they are not. Please widen the return type or add overloads keyed onbeam_width. As per coding guidelines, "Use@overloadin Python when a return type depends on input type; alternatively useTypeVarif return type can be expressed using input type".🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` around lines 1196 - 1201, The return type of get_batch_cache_indices is incorrect for beam_width > 1 — it currently declares List[List[int]] but actually returns per-request per-beam lists (i.e., List[List[List[int]]]). Update the signature to reflect both shapes by adding `@overload` definitions: one overload for beam_width: int = 1 returning List[List[int]] and another for beam_width: int > 1 returning List[List[List[int]]>, or alternatively use a TypeVar/Union to express the conditional return type; change the concrete implementation signature to a widened return type (e.g., Union[List[List[int]], List[List[List[int]]]] or List[Any]) while keeping the overloads for callers and type checkers; adjust any callers if necessary to handle the beam-aware shape.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@examples/disaggregated/simpler_example/log_disagg`:
- Line 5: The committed runtime log file log_disagg contains sensitive local
environment details (local filesystem paths and internal endpoints/IPs); remove
this artifact from the commit or sanitize it by stripping or replacing
filesystem paths and any endpoint/IP addresses with generic placeholders (e.g.,
[REDACTED_PATH], [REDACTED_ENDPOINT]) and ensure no other runtime logs remain in
the file, then add log_disagg (or the logs directory) to .gitignore to prevent
future commits.
In `@examples/disaggregated/simpler_example/output.json`:
- Around line 1-4: The output.json file contains curl progress text mixed into
the JSON payload (curl progress lines injected into the JSON body), making it
invalid JSON; re-run the export/generation step using silent curl (use the
--silent/-s flag) or redirect curl's stderr to a separate file so only the raw
JSON response is written to output.json, then replace the file with the clean
JSON body and commit only that clean JSON; locate where the pipeline or script
writes to output.json (the command invoking curl) and change it to silent output
or separate stderr so no progress is embedded.
- Around line 3-4: The checked-in sample output exposes an internal endpoint
string in the JSON field disaggregated_params.ctx_info_endpoint; locate all
occurrences of "ctx_info_endpoint" in
examples/disaggregated/simpler_example/output.json and replace the actual
IP:port value (e.g. "tcp://10.176.16.44:40153") with a safe placeholder like
"tcp://<HOST>:<PORT>" or "REDACTED" so the artifact no longer leaks internal
topology details, ensuring every disaggregated_params.ctx_info_endpoint entry is
updated consistently.
In `@examples/disaggregated/simpler_example/parse_output.sh`:
- Around line 6-8: The Perl regex in the parse_output.sh one-liner is brittle
for JSON; replace the Perl extraction with a direct jq query that reads
output.json and extracts .choices[].text (e.g., use jq to pull .choices[].text,
then normalize newlines with gsub("[\r\n]+"; " ") and write to output.txt) so
the script uses jq end-to-end rather than the Perl regex.
In `@examples/disaggregated/simpler_example/run_agg.sh`:
- Around line 24-26: The health-check loop using curl against
"http://${HOST}:${PORT}/health" can hang forever; add a timeout by introducing a
deadline or max-wait variable (e.g., HEALTH_TIMEOUT or MAX_WAIT_SECS) and track
start time (or retry count) before entering the while loop that checks the curl
HTTP code, break out and exit non‑zero if the timeout is reached; update the
loop that references HOST and PORT so it fails fast with an error message and
non‑zero exit when the deadline expires.
In `@examples/disaggregated/simpler_example/run.sh`:
- Around line 33-35: The readiness loop in run.sh that polls
"http://localhost:8000/health" can hang indefinitely; modify the while loop that
currently checks curl's HTTP code to include a timeout mechanism (e.g., set a
MAX_WAIT_SECONDS variable, capture start time or decrement a counter each
iteration), break and exit non‑zero with a clear error message if the server
doesn't become healthy before timeout, and optionally print elapsed time or last
status before exiting so CI/automation can fail fast.
In `@implementation_plan.md`:
- Around line 39-40: The transfer pipeline is currently incorrect for beam_width
> 1 (Sender._build_kv_write_meta, _align_kv_blocks,
RecvReqInfo.to_bytes/from_bytes use block_ids_per_layer_groups as 1-D suffixes),
so either complete the sender/receiver metadata changes or add a hard runtime
guard; to fix quickly, add a clear runtime check at the request
construction/entry point (and mirror it in Sender._build_kv_write_meta and
RecvReqInfo construction) that rejects beam_width > 1 with a descriptive
error/exception (e.g., "beam_width > 1 not supported for KV transfer metadata
yet"), and ensure callers fail fast rather than proceeding into _align_kv_blocks
or RecvReqInfo.to_bytes/from_bytes; alternatively, if you implement full
support, update the src_start/dst_start math in Sender._build_kv_write_meta and
append per-beam tail block IDs on the destination side or move per-beam tails to
a dedicated channel and update RecvReqInfo.to_bytes/from_bytes to encode/decode
that format.
In `@tensorrt_llm/_torch/disaggregation/native/transfer.py`:
- Around line 64-69: The function _block_ids_debug_summary currently
materializes full block-id arrays via arr.tolist(), causing heavy allocations
and log spam; change the returned dict to include only metadata (e.g., "shape":
tuple(arr.shape) and "count": arr.size or len(arr)) and remove the "block_ids"
full dump, and if a full block-id list is still needed gate that materialization
behind an explicit debug-level check (e.g., only call arr.tolist() when
logger.isEnabledFor(logging.DEBUG) or when a debug flag is passed into
_block_ids_debug_summary). Apply the same change to the other similar debug
summary blocks mentioned (the sections around the second block, the large
transfer loop, and the later summary at 1421-1427), replacing full tolist()
dumps with counts/shapes and only producing full lists under the debug guard.
In `@tensorrt_llm/_torch/disaggregation/transceiver.py`:
- Line 160: The info-level logs in _create_kv_slice (e.g., the logger.info that
prints py_request_id and the places that call block_ids.tolist()) are
materializing potentially large block_ids arrays on the hot path; change these
to a lower level (debug/trace) and guard the tolist() calls behind a level check
so the array is only converted when that level is enabled (e.g., use
logger.isEnabledFor(logging.DEBUG) before calling block_ids.tolist()); apply the
same change for all similar logging sites in _create_kv_slice (the blocks around
the other info calls referenced in the comment, which reference block_ids and
req.py_request_id).
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 3571-3580: The info-level log in py_executor.py that prints
request IDs, token tails, draft tokens and cumulative log-probs (the logger.info
call that includes req.py_request_id, beam_width, first_gen_tokens, getattr(...
'first_gen_cum_log_probs'), req.py_draft_tokens and self._debug_request_tokens)
should be removed or demoted to debug/trace only; update that logger.info to
logger.debug (or remove the sensitive fields) and ensure similar prints at the
other transfer-complete sites (the nearby logger calls around the other blocks
you noted at the other two locations) are changed consistently so token/score
dumps are not emitted at info level. Ensure the change affects the
methods/blocks using self._debug_request_tokens and the getattr(...,
'first_gen_cum_log_probs') extraction.
In `@tensorrt_llm/_torch/pyexecutor/sampler.py`:
- Around line 2886-2896: The logger.info call in
TorchSampler._prepare_beam_search that dumps beam_search_store, seq_slots_long
and related buffers should be removed (or moved to after the buffers are fully
initialized) because it can leak uninitialized or stale data; locate the
logger.info block that references _debug_tensor_slice(seq_slots_long),
beam_search_store.cum_log_probs, beam_search_store.cache_indirection,
cache_indirection_buffer, and beam_search_store.predecessor_beams and delete
that pre-reset dump so no beam-store buffers are logged before they are
reset/initialized.
- Around line 119-197: The debug helpers (_debug_tensor_slice,
_debug_new_tokens_host, _debug_cache_indirection_prefix,
_debug_batch_next_tokens) eagerly detach/copy tensors and even perform
per-request tensor indexing inside loops; gate or rewrite them so they do not
cause D2H on the hot path. Fix by either short-circuiting when debug is disabled
(return a lightweight placeholder) or by converting whole batched tensors once
to host Python lists up-front (use tensor.detach().cpu().tolist() or
tensor.tolist() once outside any per-request loop), then iterate over those
lists (handle prefix_lens being int vs tensor by converting to a list similarly)
– remove per-iteration .detach()/cpu()/indexing and per-request
tensor.index_select usage in _debug_cache_indirection_prefix and
_debug_new_tokens_host, and ensure _debug_tensor_slice and
_debug_batch_next_tokens only materialize when explicitly enabled.
- Around line 3331-3343: The info-level logging in
TorchSampler._add_metadata_to_grouped_requests exposes sensitive request/token
state; change the logger.info call to debug-level (logger.debug) and avoid
printing raw token/cached state—either redact actual token/log-prob values
(replace content from _debug_tensor_slice and _debug_cache_indirection_prefix
with shapes/lengths or "<redacted>") or log only non-sensitive summaries (e.g.,
tensor shapes, seq lengths, and counts). Also wrap the debug emission with a
debug-enabled check (logger.isEnabledFor(logging.DEBUG)) if expensive to
compute. Apply the same change for the similar logging blocks referenced (the
logger calls around the other ranges that call
_debug_tensor_slice/_debug_cache_indirection_prefix and metadata fields such as
metadata.seq_slots, metadata.seq_lens, metadata.cum_log_probs,
metadata.cache_indirection, metadata.predecessor_beams,
metadata.finished_beams).
In `@tensorrt_llm/executor/result.py`:
- Around line 521-527: The code currently builds first_gen_cum_log_probs by
filtering out None values from self._outputs, which can compress and misalign
beam indices; instead, in the block that sets
self._disaggregated_params.first_gen_cum_log_probs (using self._outputs and
self._disaggregated_params), only assign a list of cumulative_logprob values
when every out in self._outputs has a non-None cumulative_logprob and the
resulting list length equals the number of beams (otherwise leave
first_gen_cum_log_probs unset/None); locate the creation in the method
referencing self._outputs and replace the filtered-comprehension logic with a
check for any None (or compare lengths) before mapping to a list of
out.cumulative_logprob and assigning it to
self._disaggregated_params.first_gen_cum_log_probs.
---
Outside diff comments:
In `@tensorrt_llm/serve/openai_disagg_service.py`:
- Around line 343-376: The per-request logger.info calls in _get_gen_request
that log disagg_request_id, ctx_info_endpoint and beam/request metadata should
not run at INFO level on the hot path; change them to be guarded
debug/feature-flag logs. Replace the two logger.info invocations with either
logger.debug(...) or wrap them in an if logger.isEnabledFor(logging.DEBUG):
block (or check a feature flag like self._verbose_logging) so that the messages
from logger.info(...) in _get_gen_request (and any logging of
request.disaggregated_params.ctx_info_endpoint, ctx_response.choices,
CompletionRequest/ChatCompletionRequest prompt token fields, and
DisaggregatedParams creation) are only emitted when debug/verbose logging is
enabled.
- Around line 532-559: The code currently drops all but choices[0] even when
ctx_response.choices has multiple entries; update the flow to handle
multi-choice context responses instead of reducing to the first choice: in the
validation routine (where you currently index into choice =
ctx_response.choices[0]) iterate over ctx_response.choices and validate each
choice.disaggregated_params, ctx_request_id and disagg_request_id (raising with
contextual finish_reason and disagg/ctx ids for that specific choice if
invalid), and then either (A) return the full ctx_response with all validated
choices or (B) change the downstream callsites (_verify_ctx_response, _need_gen,
_get_gen_request) to accept and process a list of choices (e.g., generate
per-choice gen requests) so no disaggregation state or finish_reason is silently
ignored. Ensure all places that previously assumed a single choice are updated
to consume the list or handle branching per choice.
---
Minor comments:
In `@examples/disaggregated/simpler_example/kill_agg.sh`:
- Around line 3-4: The script currently assigns AGG_PROCESS_ID with pgrep and
always runs kill -9 $AGG_PROCESS_ID which will fail if no PID was found; update
the script to check whether AGG_PROCESS_ID is non-empty (e.g., test -n
"$AGG_PROCESS_ID" or [[ -n "$AGG_PROCESS_ID" ]]) before calling kill -9, and
only invoke kill when AGG_PROCESS_ID contains a PID (ensure you quote the
variable when used).
In `@examples/disaggregated/simpler_example/kill.sh`:
- Around line 3-7: The kill -9 invocation should be guarded so it only runs for
PIDs that exist: check each captured variable (CTX_PROCESS_ID, GEN_PROCESS_ID,
DISAGG_PROCESS_ID) for non-empty before calling kill, and invoke kill only with
the subset of PIDs that are present (or skip if none). Update the script around
the pgrep assignments and the kill line to build a list of non-empty PID
variables (or test each with [ -n "$VAR" ]), then call kill on that list or do
nothing if the list is empty; ensure this logic replaces the unconditional kill
-9 $CTX_PROCESS_ID $GEN_PROCESS_ID $DISAGG_PROCESS_ID.
In `@examples/disaggregated/simpler_example/output_agg.json`:
- Around line 1-4: The file output_agg.json contains non-JSON progress lines
before the JSON object which break parsers; remove the leading curl/progress
text lines (the three lines starting with "% Total" and the two progress rows)
so the file begins directly with the JSON object starting with {"id":...};
ensure no other stray characters precede or follow the JSON and that the
remaining content is a single well-formed JSON document.
In `@examples/disaggregated/simpler_example/run.sh`:
- Around line 6-10: The script unconditionally runs kill -9 on variables
CTX_PROCESS_ID, GEN_PROCESS_ID, DISAGG_PROCESS_ID which may be empty; update the
shutdown logic to check each PID variable (CTX_PROCESS_ID, GEN_PROCESS_ID,
DISAGG_PROCESS_ID) is non-empty (and optionally verify the PID is live) before
calling kill, e.g. guard each kill with a conditional or loop over the three
variables and only invoke kill for present PIDs, quoting variables to avoid
word-splitting and handling multiple PIDs safely.
In `@implementation_plan.md`:
- Line 17: Fix the typo "demension" → "dimension" in the implementation plan and
update the code so beam_width is made accessible in _create_kv_slice and passed
into get_block_ids and get_batch_cache_indices; remove any logic that flattens
the beam dimension and delete the assertion that beam dimension == 1 in
get_batch_cache_indices; update indexing in get_batch_cache_indices (and any
related block index handling) to treat beam_width as an extra dimension in the
block indices tensor so the functions correctly handle beam_width > 1.
In `@tensorrt_llm/_torch/disaggregation/native/transfer.py`:
- Around line 728-748: In _build_kv_write_meta(), replace the runtime assertions
that guard src_beam_block_ids.size and dst_beam_block_ids.size and the SWA
prompt-length check (currently using assert on task._prompt_len) with explicit
exceptions: raise ValueError with clear messages when src_beam_block_ids.size or
dst_beam_block_ids.size exceed total_blocks, and raise RuntimeError (or
ValueError) when task._prompt_len is None; keep the same diagnostic text
(including beam_idx, slice_end, tpb) and ensure these checks occur before
computing src_start/dst_start and before calling _align_kv_blocks() so failures
remain deterministic even when Python assertions are disabled.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Around line 3594-3601: The list comprehension in _debug_request_tokens is
mis-indented and triggers Flake8 E126; fix it by reformatting the comprehension
so the continuation lines use 4-space indentation and align logically with the
opening bracket — ensure the inner "for tokens in [list(req.get_tokens(beam))]"
and the "for beam in range(beam_width)" parts are placed on their own properly
indented lines under the opening "[" within the _debug_request_tokens
staticmethod to satisfy PEP8/Flake8.
In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 1215-1274: The two loops using zip(request_ids, result) in
get_batch_cache_indices should require equal lengths to avoid silent truncation;
update both occurrences (the first loop starting with "for req_id, req_blocks in
zip(request_ids, result):" and the later loop "for req_id, shaped_blocks in
zip(request_ids, result):") to use zip(request_ids, result, strict=True) so a
length mismatch raises an error and the logged token_layout/shaped_layout stay
consistent with returned data.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 1196-1201: The return type of get_batch_cache_indices is incorrect
for beam_width > 1 — it currently declares List[List[int]] but actually returns
per-request per-beam lists (i.e., List[List[List[int]]]). Update the signature
to reflect both shapes by adding `@overload` definitions: one overload for
beam_width: int = 1 returning List[List[int]] and another for beam_width: int >
1 returning List[List[List[int]]>, or alternatively use a TypeVar/Union to
express the conditional return type; change the concrete implementation
signature to a widened return type (e.g., Union[List[List[int]],
List[List[List[int]]]] or List[Any]) while keeping the overloads for callers and
type checkers; adjust any callers if necessary to handle the beam-aware shape.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yaml
Review profile: CHILL
Plan: Enterprise
Run ID: b15de632-a9d3-42d9-97cc-409c44b0cbb2
📒 Files selected for processing (31)
cpp/tensorrt_llm/batch_manager/cacheFormatter.cppexamples/disaggregated/simpler_example/agg_config.yamlexamples/disaggregated/simpler_example/ctx_config.yamlexamples/disaggregated/simpler_example/disagg_config.yamlexamples/disaggregated/simpler_example/gen_config.yamlexamples/disaggregated/simpler_example/kill.shexamples/disaggregated/simpler_example/kill_agg.shexamples/disaggregated/simpler_example/log_aggexamples/disaggregated/simpler_example/log_ctx_0examples/disaggregated/simpler_example/log_disaggexamples/disaggregated/simpler_example/log_gen_0examples/disaggregated/simpler_example/output.jsonexamples/disaggregated/simpler_example/output.txtexamples/disaggregated/simpler_example/output_agg.jsonexamples/disaggregated/simpler_example/parse_output.shexamples/disaggregated/simpler_example/run.shexamples/disaggregated/simpler_example/run_agg.shimplementation_notes.mdimplementation_plan.mdtensorrt_llm/_torch/disaggregation/base/transfer.pytensorrt_llm/_torch/disaggregation/native/transfer.pytensorrt_llm/_torch/disaggregation/resource/cache_reuse.pytensorrt_llm/_torch/disaggregation/transceiver.pytensorrt_llm/_torch/pyexecutor/llm_request.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytensorrt_llm/_torch/pyexecutor/sampler.pytensorrt_llm/disaggregated_params.pytensorrt_llm/executor/result.pytensorrt_llm/serve/openai_disagg_service.pytensorrt_llm/serve/openai_protocol.py
| /usr/local/lib/python3.12/dist-packages/requests/__init__.py:113: RequestsDependencyWarning: urllib3 (2.6.3) or chardet (6.0.0.post1)/charset_normalizer (3.4.4) doesn't match a supported version! | ||
| warnings.warn( | ||
| Multiple distributions found for package modelopt. Picked distribution: nvidia-modelopt | ||
| /home/scratch.athenac_coreai/pyuser/lib/python3.12/site-packages/modelopt/torch/__init__.py:36: UserWarning: transformers version 5.3.0 is incompatible with nvidia-modelopt and may cause issues. Please install recommended version with `pip install nvidia-modelopt[hf]` if working with HF models. |
Contributor
There was a problem hiding this comment.
Do not commit raw runtime logs with environment details.
This file exposes local paths and internal endpoint/IP information (for example on Lines 5 and 40-42). Please remove this artifact from source control or sanitize it before committing.
Also applies to: 40-42
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/disaggregated/simpler_example/log_disagg` at line 5, The committed
runtime log file log_disagg contains sensitive local environment details (local
filesystem paths and internal endpoints/IPs); remove this artifact from the
commit or sanitize it by stripping or replacing filesystem paths and any
endpoint/IP addresses with generic placeholders (e.g., [REDACTED_PATH],
[REDACTED_ENDPOINT]) and ensure no other runtime logs remain in the file, then
add log_disagg (or the logs directory) to .gitignore to prevent future commits.
Comment on lines
+1
to
+4
| % Total % Received % Xferd Average Speed Time Time Time Current | ||
| Dload Upload Total Spent Left Speed | ||
| 0 0 0 0 0 0 0 0 --:--:-- --:--:-- --:--:-- 0100 219 0 0 100 219 0 182 0:00:01 0:00:01 --:--:-- 182100 219 0 0 100 219 0 99 0:00:02 0:00:02 --:--:-- 99100 219 0 0 100 219 0 68 0:00:03 0:00:03 --:--:-- 68{"id":"cmpl-103ca22caad247b789c73fccd5424a97","object":"text_completion","created":1779472663,"model":"TinyLlama/TinyLlama-1.1B-Chat-v1.0","choices":[{"index":0,"text":"their products are used in so many different industries. For example, NVIDIA's GPUs are used in the gaming industry, where they are used to create high-quality graphics for video games. NVIDIA's GPUs are also used in the automotive industry, where they are used to create high-performance engines. NVIDIA's GPUs are also used in the medical industry, where they are used to create high-quality medical imaging equipment. NVIDIA's GPUs are also used in the entertainment industry, where they are used to create high-quality visual effects for movies and TV shows. Overall, NVIDIA's products are used in a wide range of industries, and their products are used in many different applications.","token_ids":null,"logprobs":null,"context_logits":null,"finish_reason":"stop","stop_reason":null,"disaggregated_params":{"request_type":"generation_only","first_gen_tokens":[896,310,372,405],"first_gen_log_probs":null,"first_gen_cum_log_probs":[-0.8538165092468262,-1.5413165092468262,-1.6663165092468262,-3.791316509246826],"first_gen_logits":null,"ctx_request_id":15280142166110208,"encoded_opaque_state":null,"draft_tokens":null,"disagg_request_id":15280142166110208,"ctx_dp_rank":0,"ctx_info_endpoint":"tcp://10.176.16.44:40153","schedule_style":0,"conversation_id":null},"avg_decoded_tokens_per_iter":1.0},{"index":1,"text":"their products are used in so many different industries. For example, NVIDIA's GPUs are used in the gaming industry, where they are used to create high-quality graphics for video games. NVIDIA's GPUs are also used in the automotive industry, where they are used to create high-performance engines. NVIDIA's GPUs are also used in the medical industry, where they are used to create high-quality medical imaging equipment. NVIDIA's GPUs are also used in the entertainment industry, where they are used to create high-quality visual effects for movies and TV shows. Overall, NVIDIA's products are used in a wide variety of industries, and their products are used in many different applications.","token_ids":null,"logprobs":null,"context_logits":null,"finish_reason":"stop","stop_reason":null,"disaggregated_params":{"request_type":"generation_only","first_gen_tokens":[896,310,372,405],"first_gen_log_probs":null,"first_gen_cum_log_probs":[-0.8538165092468262,-1.5413165092468262,-1.6663165092468262,-3.791316509246826],"first_gen_logits":null,"ctx_request_id":15280142166110208,"encoded_opaque_state":null,"draft_tokens":null,"disagg_request_id":15280142166110208,"ctx_dp_rank":0,"ctx_info_endpoint":"tcp://10.176.16.44:40153","schedule_style":0,"conversation_id":null},"avg_decoded_tokens_per_iter":1.0},{"index":2,"text":"their products are used in so many different industries. For example, NVIDIA's GPUs are used in the gaming industry, where they are used to create high-quality graphics for video games. NVIDIA's GPUs are also used in the automotive industry, where they are used to create high-performance engines. NVIDIA's GPUs are also used in the medical industry, where they are used to create high-quality medical imaging equipment. NVIDIA's GPUs are also used in the entertainment industry, where they are used to create high-quality visual effects for movies and TV shows. Overall, NVIDIA's products are used in a wide range of industries, and their products are used in a variety of applications.","token_ids":null,"logprobs":null,"context_logits":null,"finish_reason":"stop","stop_reason":null,"disaggregated_params":{"request_type":"generation_only","first_gen_tokens":[896,310,372,405],"first_gen_log_probs":null,"first_gen_cum_log_probs":[-0.8538165092468262,-1.5413165092468262,-1.6663165092468262,-3.791316509246826],"first_gen_logits":null,"ctx_request_id":15280142166110208,"encoded_opaque_state":null,"draft_tokens":null,"disagg_request_id":15280142166110208,"ctx_dp_rank":0,"ctx_info_endpoint":"tcp://10.176.16.44:40153","schedule_style":0,"conversation_id":null},"avg_decoded_token100 5801 100 5582 100 219 1623 63 0:00:03 0:00:03 --:--:-- 1687 | ||
| s_per_iter":1.0},{"index":3,"text":"their products are used in so many different industries. For example, NVIDIA's GPUs are used in the gaming industry, where they are used to create high-quality graphics for video games. NVIDIA's GPUs are also used in the automotive industry, where they are used to create high-performance engines. NVIDIA's GPUs are also used in the medical industry, where they are used to create high-quality medical imaging equipment. NVIDIA's GPUs are also used in the entertainment industry, where they are used to create high-quality visual effects for movies and TV shows. Overall, NVIDIA's products are used in a wide range of industries, and their products are used in a variety of different applications.","token_ids":null,"logprobs":null,"context_logits":null,"finish_reason":"stop","stop_reason":null,"disaggregated_params":{"request_type":"generation_only","first_gen_tokens":[896,310,372,405],"first_gen_log_probs":null,"first_gen_cum_log_probs":[-0.8538165092468262,-1.5413165092468262,-1.6663165092468262,-3.791316509246826],"first_gen_logits":null,"ctx_request_id":15280142166110208,"encoded_opaque_state":null,"draft_tokens":null,"disagg_request_id":15280142166110208,"ctx_dp_rank":0,"ctx_info_endpoint":"tcp://10.176.16.44:40153","schedule_style":0,"conversation_id":null},"avg_decoded_tokens_per_iter":1.0}],"usage":{"prompt_tokens":9,"total_tokens":676,"completion_tokens":667,"prompt_tokens_details":{"cached_tokens":0}},"prompt_token_ids":null} |
Contributor
There was a problem hiding this comment.
output.json is malformed and not machine-readable JSON.
curl progress output is mixed into the payload (including an injected progress line inside a JSON field), so this file cannot be parsed as JSON. Please regenerate with silent curl output (or redirect stderr separately) and commit only the JSON body.
🧰 Tools
🪛 Biome (2.4.15)
[error] 1-1: unexpected character %
(parse)
[error] 1-1: String values must be double quoted.
(parse)
[error] 1-1: unexpected character %
(parse)
[error] 1-1: String values must be double quoted.
(parse)
[error] 1-1: unexpected character %
(parse)
[error] 1-1: String values must be double quoted.
(parse)
[error] 1-1: String values must be double quoted.
(parse)
[error] 1-1: String values must be double quoted.
(parse)
[error] 1-1: String values must be double quoted.
(parse)
[error] 1-1: String values must be double quoted.
(parse)
[error] 1-1: String values must be double quoted.
(parse)
[error] 1-1: String values must be double quoted.
(parse)
[error] 2-2: String values must be double quoted.
(parse)
[error] 2-2: String values must be double quoted.
(parse)
[error] 2-2: String values must be double quoted.
(parse)
[error] 2-2: String values must be double quoted.
(parse)
[error] 2-2: String values must be double quoted.
(parse)
[error] 2-2: String values must be double quoted.
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: End of file expected
(parse)
🪛 OpenGrep (1.21.0)
[ERROR] 3-3: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 3-3: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 3-3: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 4-4: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/disaggregated/simpler_example/output.json` around lines 1 - 4, The
output.json file contains curl progress text mixed into the JSON payload (curl
progress lines injected into the JSON body), making it invalid JSON; re-run the
export/generation step using silent curl (use the --silent/-s flag) or redirect
curl's stderr to a separate file so only the raw JSON response is written to
output.json, then replace the file with the clean JSON body and commit only that
clean JSON; locate where the pipeline or script writes to output.json (the
command invoking curl) and change it to silent output or separate stderr so no
progress is embedded.
Comment on lines
+3
to
+4
| 0 0 0 0 0 0 0 0 --:--:-- --:--:-- --:--:-- 0100 219 0 0 100 219 0 182 0:00:01 0:00:01 --:--:-- 182100 219 0 0 100 219 0 99 0:00:02 0:00:02 --:--:-- 99100 219 0 0 100 219 0 68 0:00:03 0:00:03 --:--:-- 68{"id":"cmpl-103ca22caad247b789c73fccd5424a97","object":"text_completion","created":1779472663,"model":"TinyLlama/TinyLlama-1.1B-Chat-v1.0","choices":[{"index":0,"text":"their products are used in so many different industries. For example, NVIDIA's GPUs are used in the gaming industry, where they are used to create high-quality graphics for video games. NVIDIA's GPUs are also used in the automotive industry, where they are used to create high-performance engines. NVIDIA's GPUs are also used in the medical industry, where they are used to create high-quality medical imaging equipment. NVIDIA's GPUs are also used in the entertainment industry, where they are used to create high-quality visual effects for movies and TV shows. Overall, NVIDIA's products are used in a wide range of industries, and their products are used in many different applications.","token_ids":null,"logprobs":null,"context_logits":null,"finish_reason":"stop","stop_reason":null,"disaggregated_params":{"request_type":"generation_only","first_gen_tokens":[896,310,372,405],"first_gen_log_probs":null,"first_gen_cum_log_probs":[-0.8538165092468262,-1.5413165092468262,-1.6663165092468262,-3.791316509246826],"first_gen_logits":null,"ctx_request_id":15280142166110208,"encoded_opaque_state":null,"draft_tokens":null,"disagg_request_id":15280142166110208,"ctx_dp_rank":0,"ctx_info_endpoint":"tcp://10.176.16.44:40153","schedule_style":0,"conversation_id":null},"avg_decoded_tokens_per_iter":1.0},{"index":1,"text":"their products are used in so many different industries. For example, NVIDIA's GPUs are used in the gaming industry, where they are used to create high-quality graphics for video games. NVIDIA's GPUs are also used in the automotive industry, where they are used to create high-performance engines. NVIDIA's GPUs are also used in the medical industry, where they are used to create high-quality medical imaging equipment. NVIDIA's GPUs are also used in the entertainment industry, where they are used to create high-quality visual effects for movies and TV shows. Overall, NVIDIA's products are used in a wide variety of industries, and their products are used in many different applications.","token_ids":null,"logprobs":null,"context_logits":null,"finish_reason":"stop","stop_reason":null,"disaggregated_params":{"request_type":"generation_only","first_gen_tokens":[896,310,372,405],"first_gen_log_probs":null,"first_gen_cum_log_probs":[-0.8538165092468262,-1.5413165092468262,-1.6663165092468262,-3.791316509246826],"first_gen_logits":null,"ctx_request_id":15280142166110208,"encoded_opaque_state":null,"draft_tokens":null,"disagg_request_id":15280142166110208,"ctx_dp_rank":0,"ctx_info_endpoint":"tcp://10.176.16.44:40153","schedule_style":0,"conversation_id":null},"avg_decoded_tokens_per_iter":1.0},{"index":2,"text":"their products are used in so many different industries. For example, NVIDIA's GPUs are used in the gaming industry, where they are used to create high-quality graphics for video games. NVIDIA's GPUs are also used in the automotive industry, where they are used to create high-performance engines. NVIDIA's GPUs are also used in the medical industry, where they are used to create high-quality medical imaging equipment. NVIDIA's GPUs are also used in the entertainment industry, where they are used to create high-quality visual effects for movies and TV shows. Overall, NVIDIA's products are used in a wide range of industries, and their products are used in a variety of applications.","token_ids":null,"logprobs":null,"context_logits":null,"finish_reason":"stop","stop_reason":null,"disaggregated_params":{"request_type":"generation_only","first_gen_tokens":[896,310,372,405],"first_gen_log_probs":null,"first_gen_cum_log_probs":[-0.8538165092468262,-1.5413165092468262,-1.6663165092468262,-3.791316509246826],"first_gen_logits":null,"ctx_request_id":15280142166110208,"encoded_opaque_state":null,"draft_tokens":null,"disagg_request_id":15280142166110208,"ctx_dp_rank":0,"ctx_info_endpoint":"tcp://10.176.16.44:40153","schedule_style":0,"conversation_id":null},"avg_decoded_token100 5801 100 5582 100 219 1623 63 0:00:03 0:00:03 --:--:-- 1687 | ||
| s_per_iter":1.0},{"index":3,"text":"their products are used in so many different industries. For example, NVIDIA's GPUs are used in the gaming industry, where they are used to create high-quality graphics for video games. NVIDIA's GPUs are also used in the automotive industry, where they are used to create high-performance engines. NVIDIA's GPUs are also used in the medical industry, where they are used to create high-quality medical imaging equipment. NVIDIA's GPUs are also used in the entertainment industry, where they are used to create high-quality visual effects for movies and TV shows. Overall, NVIDIA's products are used in a wide range of industries, and their products are used in a variety of different applications.","token_ids":null,"logprobs":null,"context_logits":null,"finish_reason":"stop","stop_reason":null,"disaggregated_params":{"request_type":"generation_only","first_gen_tokens":[896,310,372,405],"first_gen_log_probs":null,"first_gen_cum_log_probs":[-0.8538165092468262,-1.5413165092468262,-1.6663165092468262,-3.791316509246826],"first_gen_logits":null,"ctx_request_id":15280142166110208,"encoded_opaque_state":null,"draft_tokens":null,"disagg_request_id":15280142166110208,"ctx_dp_rank":0,"ctx_info_endpoint":"tcp://10.176.16.44:40153","schedule_style":0,"conversation_id":null},"avg_decoded_tokens_per_iter":1.0}],"usage":{"prompt_tokens":9,"total_tokens":676,"completion_tokens":667,"prompt_tokens_details":{"cached_tokens":0}},"prompt_token_ids":null} |
Contributor
There was a problem hiding this comment.
Redact internal endpoint details from committed sample output.
ctx_info_endpoint currently exposes an internal IP/port in a checked-in artifact. Please replace with a placeholder value in examples to avoid leaking internal topology details.
🧰 Tools
🪛 Biome (2.4.15)
[error] 4-4: End of file expected
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: End of file expected
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: Minus must be followed by a digit
(parse)
[error] 4-4: End of file expected
(parse)
🪛 OpenGrep (1.21.0)
[ERROR] 3-3: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 3-3: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 3-3: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
[ERROR] 4-4: Possible credit card number (PAN) detected in source code. Credit card numbers should never be hardcoded or stored in source files. Use a secrets manager or tokenization service instead.
(coderabbit.pii.credit-card-number)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/disaggregated/simpler_example/output.json` around lines 3 - 4, The
checked-in sample output exposes an internal endpoint string in the JSON field
disaggregated_params.ctx_info_endpoint; locate all occurrences of
"ctx_info_endpoint" in examples/disaggregated/simpler_example/output.json and
replace the actual IP:port value (e.g. "tcp://10.176.16.44:40153") with a safe
placeholder like "tcp://<HOST>:<PORT>" or "REDACTED" so the artifact no longer
leaks internal topology details, ensuring every
disaggregated_params.ctx_info_endpoint entry is updated consistently.
Comment on lines
+6
to
+8
| perl -0ne 'while (/"text"[[:space:]]*:[[:space:]]*"((?:\\.|[^"\\])*)"/gs) { print "\"$1\"\n"; }' \ | ||
| "${SCRIPT_DIR}/output.json" \ | ||
| | jq -r 'gsub("[\r\n]+"; " ")' > "${SCRIPT_DIR}/output.txt" No newline at end of file |
Contributor
There was a problem hiding this comment.
Avoid regex parsing for JSON extraction.
On Line 6, Perl regex parsing is brittle for escaped sequences and edge cases. Use jq to extract .choices[].text directly.
Suggested fix
-# extract text of each choice from output.json and output to output.txt
-perl -0ne 'while (/"text"[[:space:]]*:[[:space:]]*"((?:\\.|[^"\\])*)"/gs) { print "\"$1\"\n"; }' \
- "${SCRIPT_DIR}/output.json" \
- | jq -r 'gsub("[\r\n]+"; " ")' > "${SCRIPT_DIR}/output.txt"
+# extract text of each choice from output.json and output to output.txt
+jq -r '.choices[].text | gsub("[\r\n]+"; " ")' \
+ "${SCRIPT_DIR}/output.json" > "${SCRIPT_DIR}/output.txt"📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| perl -0ne 'while (/"text"[[:space:]]*:[[:space:]]*"((?:\\.|[^"\\])*)"/gs) { print "\"$1\"\n"; }' \ | |
| "${SCRIPT_DIR}/output.json" \ | |
| | jq -r 'gsub("[\r\n]+"; " ")' > "${SCRIPT_DIR}/output.txt" | |
| # extract text of each choice from output.json and output to output.txt | |
| jq -r '.choices[].text | gsub("[\r\n]+"; " ")' \ | |
| "${SCRIPT_DIR}/output.json" > "${SCRIPT_DIR}/output.txt" |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/disaggregated/simpler_example/parse_output.sh` around lines 6 - 8,
The Perl regex in the parse_output.sh one-liner is brittle for JSON; replace the
Perl extraction with a direct jq query that reads output.json and extracts
.choices[].text (e.g., use jq to pull .choices[].text, then normalize newlines
with gsub("[\r\n]+"; " ") and write to output.txt) so the script uses jq
end-to-end rather than the Perl regex.
Comment on lines
+24
to
+26
| while [ $(curl -s -o /dev/null -w "%{http_code}" "http://${HOST}:${PORT}/health") -ne 200 ]; do | ||
| sleep 1 | ||
| done |
Contributor
There was a problem hiding this comment.
Add a timeout to the health-check loop.
Line 24 can wait forever when startup fails, which makes this script hang indefinitely.
Suggested fix
-while [ $(curl -s -o /dev/null -w "%{http_code}" "http://${HOST}:${PORT}/health") -ne 200 ]; do
- sleep 1
-done
+MAX_RETRIES=180
+retries=0
+while [ "$(curl -s -o /dev/null -w "%{http_code}" "http://${HOST}:${PORT}/health")" -ne 200 ]; do
+ retries=$((retries + 1))
+ if [ "$retries" -ge "$MAX_RETRIES" ]; then
+ echo "Aggregate server did not become ready in time" >&2
+ exit 1
+ fi
+ sleep 1
+done📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
Suggested change
| while [ $(curl -s -o /dev/null -w "%{http_code}" "http://${HOST}:${PORT}/health") -ne 200 ]; do | |
| sleep 1 | |
| done | |
| MAX_RETRIES=180 | |
| retries=0 | |
| while [ "$(curl -s -o /dev/null -w "%{http_code}" "http://${HOST}:${PORT}/health")" -ne 200 ]; do | |
| retries=$((retries + 1)) | |
| if [ "$retries" -ge "$MAX_RETRIES" ]; then | |
| echo "Aggregate server did not become ready in time" >&2 | |
| exit 1 | |
| fi | |
| sleep 1 | |
| done |
🧰 Tools
🪛 Shellcheck (0.11.0)
[warning] 24-24: Quote this to prevent word splitting.
(SC2046)
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@examples/disaggregated/simpler_example/run_agg.sh` around lines 24 - 26, The
health-check loop using curl against "http://${HOST}:${PORT}/health" can hang
forever; add a timeout by introducing a deadline or max-wait variable (e.g.,
HEALTH_TIMEOUT or MAX_WAIT_SECS) and track start time (or retry count) before
entering the while loop that checks the curl HTTP code, break out and exit
non‑zero if the timeout is reached; update the loop that references HOST and
PORT so it fails fast with an error message and non‑zero exit when the deadline
expires.
Comment on lines
+3571
to
+3580
| logger.info( | ||
| "disagg generation transfer complete: " | ||
| f"request_id={req.py_request_id} beam_width={beam_width} " | ||
| f"first_gen_tokens={first_gen_tokens} " | ||
| f"first_gen_cum_log_probs=" | ||
| f"{getattr(getattr(req, 'py_disaggregated_params', None), 'first_gen_cum_log_probs', None)} " | ||
| f"draft_tokens={req.py_draft_tokens} " | ||
| f"ctx_request_id={req.context_phase_params.req_id} " | ||
| f"tokens_before=" | ||
| f"{self._debug_request_tokens(req, beam_width)}" |
Contributor
There was a problem hiding this comment.
Drop the info-level token and score dumps from this transfer-complete path.
This emits request IDs, token tails, draft tokens, and cumulative log probs for every completed disaggregated generation transfer. In a hot path that is both a throughput risk and a sensitive-data leak; please gate it behind a debug/trace-only path or remove it before merge.
Also applies to: 3585-3590, 3647-3650
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/pyexecutor/py_executor.py` around lines 3571 - 3580, The
info-level log in py_executor.py that prints request IDs, token tails, draft
tokens and cumulative log-probs (the logger.info call that includes
req.py_request_id, beam_width, first_gen_tokens, getattr(...
'first_gen_cum_log_probs'), req.py_draft_tokens and self._debug_request_tokens)
should be removed or demoted to debug/trace only; update that logger.info to
logger.debug (or remove the sensitive fields) and ensure similar prints at the
other transfer-complete sites (the nearby logger calls around the other blocks
you noted at the other two locations) are changed consistently so token/score
dumps are not emitted at info level. Ensure the change affects the
methods/blocks using self._debug_request_tokens and the getattr(...,
'first_gen_cum_log_probs') extraction.
Comment on lines
+119
to
+197
| def _debug_tensor_slice( | ||
| tensor: torch.Tensor, | ||
| indices: torch.Tensor | None = None, | ||
| dim: int = 0, | ||
| ): | ||
| try: | ||
| view = ( | ||
| tensor.index_select(dim, indices.to(tensor.device).long()) | ||
| if indices is not None | ||
| else tensor | ||
| ) | ||
| return view.detach().cpu().tolist() | ||
| except (RuntimeError, TypeError, ValueError) as err: | ||
| return f"<unavailable: {err}>" | ||
|
|
||
|
|
||
| def _debug_new_tokens_host( | ||
| new_tokens_host: torch.Tensor, | ||
| seq_slots: torch.Tensor, | ||
| req_num_generated_tokens: torch.Tensor, | ||
| ): | ||
| try: | ||
| new_tokens = new_tokens_host.tolist() | ||
| seq_slots_host = seq_slots.detach().cpu().tolist() | ||
| req_num_generated_tokens_host = req_num_generated_tokens.detach().cpu().tolist() | ||
| return [ | ||
| { | ||
| "seq_slot": seq_slot, | ||
| "tokens": [ | ||
| new_tokens[step][seq_slot] | ||
| for step in range(req_num_generated_tokens_host[req_idx]) | ||
| ], | ||
| } | ||
| for req_idx, seq_slot in enumerate(seq_slots_host) | ||
| ] | ||
| except (RuntimeError, TypeError, ValueError, IndexError) as err: | ||
| return f"<unavailable: {err}>" | ||
|
|
||
|
|
||
| def _debug_cache_indirection_prefix( | ||
| cache_indirection: torch.Tensor, | ||
| seq_slots: torch.Tensor, | ||
| prefix_lens: torch.Tensor | int, | ||
| max_positions: int = 16, | ||
| ): | ||
| try: | ||
| seq_slots_host = seq_slots.detach().cpu().tolist() | ||
| if isinstance(prefix_lens, int): | ||
| prefix_lens_host = [prefix_lens] * len(seq_slots_host) | ||
| else: | ||
| prefix_lens_host = prefix_lens.detach().cpu().tolist() | ||
|
|
||
| result = [] | ||
| for req_idx, seq_slot in enumerate(seq_slots_host): | ||
| prefix_len = int(prefix_lens_host[req_idx]) | ||
| start = max(0, prefix_len - max_positions) | ||
| values = cache_indirection[ | ||
| seq_slot, :, start:prefix_len | ||
| ].detach().cpu().tolist() | ||
| result.append({ | ||
| "seq_slot": seq_slot, | ||
| "prefix_len": prefix_len, | ||
| "shown_range": [start, prefix_len], | ||
| "values": values, | ||
| }) | ||
| return result | ||
| except (RuntimeError, TypeError, ValueError, IndexError) as err: | ||
| return f"<unavailable: {err}>" | ||
|
|
||
|
|
||
| def _debug_batch_next_tokens( | ||
| batch_next_tokens: torch.Tensor, | ||
| batch_req_indices: torch.Tensor, | ||
| ): | ||
| try: | ||
| valid_rows = batch_req_indices.numel() | ||
| return batch_next_tokens[:valid_rows].detach().cpu().tolist() | ||
| except (RuntimeError, TypeError, ValueError, IndexError) as err: | ||
| return f"<unavailable: {err}>" |
Contributor
There was a problem hiding this comment.
Gate or remove the debug tensor snapshot helpers.
These helpers do eager detach().cpu().tolist() materialization, and _debug_cache_indirection_prefix() additionally walks device-backed data inside a per-request loop. Since the call sites build f-strings unconditionally, beam-search decode now pays those D2H sync/copy costs on every iteration. That is too expensive for this hot path.
Based on learnings: In files under tensorrt_llm/_torch/pyexecutor, avoid accessing torch.Tensor objects inside for-loops when iterating over requests. Convert batched tensors to Python lists beforehand using tensor.tolist(), and then iterate over those lists. This improves performance by reducing tensor-bound operations inside hot loops.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/pyexecutor/sampler.py` around lines 119 - 197, The debug
helpers (_debug_tensor_slice, _debug_new_tokens_host,
_debug_cache_indirection_prefix, _debug_batch_next_tokens) eagerly detach/copy
tensors and even perform per-request tensor indexing inside loops; gate or
rewrite them so they do not cause D2H on the hot path. Fix by either
short-circuiting when debug is disabled (return a lightweight placeholder) or by
converting whole batched tensors once to host Python lists up-front (use
tensor.detach().cpu().tolist() or tensor.tolist() once outside any per-request
loop), then iterate over those lists (handle prefix_lens being int vs tensor by
converting to a list similarly) – remove per-iteration .detach()/cpu()/indexing
and per-request tensor.index_select usage in _debug_cache_indirection_prefix and
_debug_new_tokens_host, and ensure _debug_tensor_slice and
_debug_batch_next_tokens only materialize when explicitly enabled.
Comment on lines
+2886
to
+2896
| logger.info( | ||
| "TorchSampler._prepare_beam_search before reset: " | ||
| f"seq_slots={_debug_tensor_slice(seq_slots_long)} " | ||
| f"max_prompt_len={max_prompt_len} " | ||
| f"cum_log_probs={_debug_tensor_slice(beam_search_store.cum_log_probs, seq_slots_long)} " | ||
| f"cache_indirection_prefix=" | ||
| f"{_debug_cache_indirection_prefix(beam_search_store.cache_indirection, seq_slots_long, max_prompt_len)} " | ||
| f"cache_indirection_buffer_prefix=" | ||
| f"{_debug_cache_indirection_prefix(beam_search_store.cache_indirection_buffer, seq_slots_long, max_prompt_len)} " | ||
| f"predecessor_beams={_debug_tensor_slice(beam_search_store.predecessor_beams, seq_slots_long)}" | ||
| ) |
Contributor
There was a problem hiding this comment.
Remove the pre-reset beam-store dump.
This logs beam-search buffers before they are reset, but those buffers are backed by torch.empty(...) and reused seq slots. On a fresh allocation you'll emit uninitialized garbage; on slot reuse you'll emit stale state from a previous request. That's a real data-leak risk in a serving path.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/pyexecutor/sampler.py` around lines 2886 - 2896, The
logger.info call in TorchSampler._prepare_beam_search that dumps
beam_search_store, seq_slots_long and related buffers should be removed (or
moved to after the buffers are fully initialized) because it can leak
uninitialized or stale data; locate the logger.info block that references
_debug_tensor_slice(seq_slots_long), beam_search_store.cum_log_probs,
beam_search_store.cache_indirection, cache_indirection_buffer, and
beam_search_store.predecessor_beams and delete that pre-reset dump so no
beam-store buffers are logged before they are reset/initialized.
Comment on lines
+3331
to
+3343
| logger.info( | ||
| "TorchSampler._add_metadata_to_grouped_requests beam metadata: " | ||
| f"group_indices={_debug_tensor_slice(value.indices)} " | ||
| f"seq_slots={_debug_tensor_slice(metadata.seq_slots)} " | ||
| f"seq_lens={_debug_tensor_slice(metadata.seq_lens)} " | ||
| f"cum_log_probs={_debug_tensor_slice(metadata.cum_log_probs, metadata.seq_slots)} " | ||
| f"cache_indirection_prefix=" | ||
| f"{_debug_cache_indirection_prefix(metadata.cache_indirection, metadata.seq_slots, metadata.seq_lens)} " | ||
| f"cache_indirection_buffer_prefix=" | ||
| f"{_debug_cache_indirection_prefix(metadata.cache_indirection_buffer, metadata.seq_slots, metadata.seq_lens)} " | ||
| f"predecessor_beams={_debug_tensor_slice(metadata.predecessor_beams, metadata.seq_slots)} " | ||
| f"finished_beams={_debug_tensor_slice(metadata.finished_beams, metadata.seq_slots)}" | ||
| ) |
Contributor
There was a problem hiding this comment.
Don't log request/token state at info level in serving code.
These messages include request_id, full beam tokens before/after append, generated-token snapshots, cumulative log-probs, and cache-indirection state. Leaving that at normal info level will persist user data in routine service logs and create a compliance/support burden. Please keep this behind a short-lived debug-only trace with redaction, or remove it before merge.
Also applies to: 3444-3491, 3935-3950, 4087-4096
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/_torch/pyexecutor/sampler.py` around lines 3331 - 3343, The
info-level logging in TorchSampler._add_metadata_to_grouped_requests exposes
sensitive request/token state; change the logger.info call to debug-level
(logger.debug) and avoid printing raw token/cached state—either redact actual
token/log-prob values (replace content from _debug_tensor_slice and
_debug_cache_indirection_prefix with shapes/lengths or "<redacted>") or log only
non-sensitive summaries (e.g., tensor shapes, seq lengths, and counts). Also
wrap the debug emission with a debug-enabled check
(logger.isEnabledFor(logging.DEBUG)) if expensive to compute. Apply the same
change for the similar logging blocks referenced (the logger calls around the
other ranges that call _debug_tensor_slice/_debug_cache_indirection_prefix and
metadata fields such as metadata.seq_slots, metadata.seq_lens,
metadata.cum_log_probs, metadata.cache_indirection, metadata.predecessor_beams,
metadata.finished_beams).
Comment on lines
+521
to
+527
| first_gen_cum_log_probs = [ | ||
| out.cumulative_logprob for out in self._outputs | ||
| if out.cumulative_logprob is not None | ||
| ] | ||
| if first_gen_cum_log_probs: | ||
| self._disaggregated_params.first_gen_cum_log_probs = \ | ||
| list(first_gen_cum_log_probs) |
Contributor
There was a problem hiding this comment.
Fallback cumulatives can silently misalign beam indices.
The current fallback drops beams with None (if out.cumulative_logprob is not None), which can compress the list and shift beam-score alignment. Prefer assigning only when all beams have values (or keep it unset), rather than filtering by value.
Suggested fix
- if first_gen_cum_log_probs is None:
- first_gen_cum_log_probs = [
- out.cumulative_logprob for out in self._outputs
- if out.cumulative_logprob is not None
- ]
- if first_gen_cum_log_probs:
+ if first_gen_cum_log_probs is None:
+ fallback_cum = [out.cumulative_logprob for out in self._outputs]
+ if all(v is not None for v in fallback_cum):
+ first_gen_cum_log_probs = fallback_cum
+ if first_gen_cum_log_probs is not None:
self._disaggregated_params.first_gen_cum_log_probs = \
list(first_gen_cum_log_probs)🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@tensorrt_llm/executor/result.py` around lines 521 - 527, The code currently
builds first_gen_cum_log_probs by filtering out None values from self._outputs,
which can compress and misalign beam indices; instead, in the block that sets
self._disaggregated_params.first_gen_cum_log_probs (using self._outputs and
self._disaggregated_params), only assign a list of cumulative_logprob values
when every out in self._outputs has a non-None cumulative_logprob and the
resulting list length equals the number of beams (otherwise leave
first_gen_cum_log_probs unset/None); locate the creation in the method
referencing self._outputs and replace the filtered-comprehension logic with a
check for any None (or compare lengths) before mapping to a list of
out.cumulative_logprob and assigning it to
self._disaggregated_params.first_gen_cum_log_probs.
Tabrizian
reviewed
May 26, 2026
Tabrizian
left a comment
Tabrizian
left a comment
Member
There was a problem hiding this comment.
Please improve the CacheTransceiverTest to include beamWidth > 1: https://github.com/NVIDIA/TensorRT-LLM/blob/897c4bffd7bab93d8bd9252108af08b87e7b6478/cpp/tests/unit_tests/multi_gpu/cacheTransceiverTest.cpp
Member
There was a problem hiding this comment.
Please remove the examples/disaggregated/simpler_example folder
Tabrizian
reviewed
May 27, 2026
Tabrizian
left a comment
Tabrizian
left a comment
Member
There was a problem hiding this comment.
Please clean up the PR from the logging changes.
| mpi::MpiComm::world().getRank(), "Start sending KV cache for request ID: %ld.", llmRequest.mRequestId); | ||
|
|
||
| TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently, only beam width 1 is supported."); | ||
| // TLLM_CHECK_WITH_INFO(llmRequest.mSamplingConfig.beamWidth == 1, "Currently, only beam width 1 is supported."); |
Member
There was a problem hiding this comment.
This line should be reverted, since we are only adding support in Python Cache Transceiver.
jieli-matrix
approved these changes
Jun 1, 2026
leslie-fang25
removed their request for review
June 30, 2026 04:17
Collaborator
Author
|
Updated PR: #14876 |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary by CodeRabbit
Release Notes
New Features
Improvements
Description
Add beam search support to disaggregated serving for the Python Cache Transceiver and C++ KVCacheManager path.
TODO: add tests, remove logging
Test Coverage
This PR has not been tested yet.
See example scripts in
examples/disaggregated/simpler_example/PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.