[None][feat] ADP conversation router: configurable least-queued placement for new conversations

[lancelly]

Summary by CodeRabbit

• New Features

  • Added configurable placement for new conversations in conversation-aware KV-cache routing.
  • Choose between round-robin assignment or least-queued placement.
  • Existing conversations remain pinned to their assigned rank.
  • Invalid placement values safely fall back to round-robin behavior.

• Tests

  • Added coverage for load balancing, batch leveling, sticky routing, deterministic results, and configuration fallback.

Description

kv_cache_routing_conversation_affinity pins each conversation to one attention-DP rank and round-robins first turns across ranks. Round-robin equalizes per-rank conversation counts, but per-conversation load (turn rate, agentic fan-out, prefill length) is far from uniform. On an agentic multi-turn workload (4×(TP8, attention-DP) prefill workers in a disaggregated deployment, concurrency 1000, ~95% KV block reuse) we measured:

• per-rank live conversation counts near-uniform: max/min ≤ 1.14×
• yet per-rank scheduler queue depth skewed 1.8–3.4× (hot rank mean 24 queued vs 8 on its neighbors)
• coldest ranks dummy-padded in up to 15.7% of iterations (lockstep ADP: an idle rank's iteration capacity is simply lost) while deep queues persisted elsewhere

This PR adds attention_dp_config.kv_cache_routing_new_conv_placement:

• "round_robin" (default) — existing behavior, unchanged.
• "least_queued" — place requests that have no pinned rank yet (first turn of a conversation, conversation-less requests, sticky overflow) on the rank with the fewest live requests, reusing the existing _least_loaded helper (previously only reachable as an RR fallback). Heavy pinned conversations keep their rank's count high, so new conversations steer away from hot ranks — emergent weight-awareness without predicting per-conversation weight. The shared in-batch count accumulator spreads bursts valley-first.

Sticky returns, soft/hard caps, and the padding invariant are untouched; routing remains a deterministic pure function of allgathered state (required by the replicated no-broadcast protocol — least-loaded is stateless, one less cross-iteration state than the RR cursor).

A/B on the workload above (only this knob changed):

metric	round_robin	least_queued
per-rank queue skew (max/min of rank means, worst instance)	3.4×	2.0×
dummy-padded iterations, worst rank	15.7%	9.9%
TTFT p90	35.5 s	30.0 s (−15.5%)
TTFT p99	105.6 s	87.5 s (−17.2%)
KV block reuse	95.0%	95.0%
request throughput	35.75 req/s	35.90 req/s

Remaining skew comes from already-pinned heavy conversations (not migrated, preserving prefix-cache residency by design).

[lancelly]

/bot run --disable-fail-fast

[coderabbitai]

Caution

Review failed

Failed to post review comments.

GitHub was unavailable or timed out while CodeRabbit was posting the review. Please request a new review later if the pull request still needs one. This happened while posting 1 inline comment. Use @coderabbitai full review to retry the review.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 10934ab4-dfe1-41e9-92a2-864c758d2f8e

📥 Commits

Reviewing files that changed from the base of the PR and between 6bd6bb8 and 01da19e.

📒 Files selected for processing (4)

• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py
• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/usage/llm_args_golden_manifest.json
• tests/unittest/_torch/executor/test_adp_router.py

🧰 Additional context used

📓 Path-based instructions (6)

**/*

📄 CodeRabbit inference engine (AGENTS.md)

Add an NVIDIA copyright header to all new files, and update the copyright year in modified files.

Files:

• tensorrt_llm/usage/llm_args_golden_manifest.json
• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py
• tests/unittest/_torch/executor/test_adp_router.py

**

⚙️ CodeRabbit configuration file

**:

TensorRT-LLM Coding Guidelines

C++ Coding Guidelines

The TensorRT-LLM C++ Coding Guidelines are mainly derived from the Google C++ Style Guide

General principle: Do not fight the tooling (clang-format, clang-tidy, pre-commit) without good reason. If the tool enforces a style, follow it.

Note: These guidelines have been inconsistently followed in the existing codebase. New code and modifications to existing code should adhere to these guidelines. Cleaning up existing code to match these guidelines as you touch it is encouraged, but bulk changes of surrounding code should be a separate PR.

---

Namespaces

1. Closing braces of namespaces should have a comment saying the namespace it closes:

namespace foo
{
...
} // namespace foo

2. Anonymous namespaces use the same convention:

namespace
{
...
} // namespace

Constants

1. Prefer const or constexpr variables over #defines whenever possible, as the latter are not visible to the compiler.
2. A variable that is not modified after its initialization should be declared as const. This applies to local and global variable declarations; it is not required for function parameters.
3. Use east-const style: place const to the right of the type it qualifies (e.g. int const x rather than const int x). This applies throughout, not just to constant declarations. This is enforced by clang-format (QualifierAlignment: Right). See Common Pitfalls for how this clarifies pointer-to-const vs const-pointer.
4. Non-POD constants (e.g. std::string, std::vector, std::unordered_map) must not be declared at file or namespace scope, as they are subject to the static initialization order fiasco. Wrap them in a function that returns a reference to a function-local static:

// Bad: non-POD at file scope — initialization order is undefined a...

Files:

• tensorrt_llm/usage/llm_args_golden_manifest.json
• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py
• tests/unittest/_torch/executor/test_adp_router.py

**/*.py

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

Python files are split into legacy and modern linting groups; when editing legacy files, keep changes compatible with the legacy toolchain and update the legacy baseline/configuration artifacts when required.

Files:

• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py
• tests/unittest/_torch/executor/test_adp_router.py

**/*.{cpp,h,cc,cuh,cu,hpp,py}

📄 CodeRabbit inference engine (CODING_GUIDELINES.md)

All TensorRT-LLM source files should contain the NVIDIA copyright header with the current year and Apache 2.0 license text.

Files:

• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py
• tests/unittest/_torch/executor/test_adp_router.py

**/*.{py,cpp,h,hpp,cc,cxx,cu,cuh}

📄 CodeRabbit inference engine (AGENTS.md)

Read and follow CODING_GUIDELINES.md for all C++ and Python code changes.

Files:

• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py
• tests/unittest/_torch/executor/test_adp_router.py

tests/**

⚙️ CodeRabbit configuration file

tests/**: Act as a QA engineer reviewing test changes and coverage for TensorRT-LLM.
Keep feedback actionable: suggest concrete list file names and whether
coverage is sufficient, insufficient, or needs follow-up outside the PR.

Files:

• tests/unittest/_torch/executor/test_adp_router.py

🧠 Learnings (14)

📚 Learning: 2026-02-13T10:15:37.120Z

Learnt from: ixlmar
Repo: NVIDIA/TensorRT-LLM PR: 11508
File: tests/unittest/_torch/sampler/test_beam_search_util.py:71-71
Timestamp: 2026-02-13T10:15:37.120Z
Learning: In TensorRT-LLM (Python requires >=3.10 and <4 as per setup.py), you can use Python 3.10+ features (e.g., PEP 585 generics like dict[str, int], list[str], etc.) throughout the codebase, and you do not need to add from __future__ import annotations. This applies to all Python files, including tests (e.g., tests/unittest/...); ensure tests and code consistently rely on Python 3.10+ features where applicable.

Applied to files:

• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py
• tests/unittest/_torch/executor/test_adp_router.py

📚 Learning: 2026-04-24T23:54:27.755Z

Learnt from: hnover-nv
Repo: NVIDIA/TensorRT-LLM PR: 13453
File: tensorrt_llm/_torch/modules/mamba/causal_conv1d_triton.py:1016-1018
Timestamp: 2026-04-24T23:54:27.755Z
Learning: In this TensorRT-LLM codebase, treat `get_sm_version()` as an intentional implementation detail: it always queries device 0 to determine the SM version. Since the project assumes homogeneous GPU deployments (all GPUs in a node share the same SM capability), code review should not flag `get_sm_version()` calls for not using the current device, a tensor’s `.device`, or an explicit device argument.

Applied to files:

• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py
• tests/unittest/_torch/executor/test_adp_router.py

📚 Learning: 2026-05-29T08:18:25.109Z

Learnt from: zhenhuaw-me
Repo: NVIDIA/TensorRT-LLM PR: 14733
File: tensorrt_llm/visual_gen/output.py:289-319
Timestamp: 2026-05-29T08:18:25.109Z
Learning: In this repository, Ruff rule A002 ("Function argument shadows a Python builtin") is not enabled in the configured Ruff rules. During code reviews, do not flag potential builtin-shadowing argument names (e.g., parameters like `format`) for A002, and do not recommend adding `# noqa: A002` suppressions—since the rule would not trigger here, such comments would be unnecessary noise, especially on public API surfaces.

Applied to files:

• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py
• tests/unittest/_torch/executor/test_adp_router.py

📚 Learning: 2026-05-29T08:18:27.934Z

Learnt from: zhenhuaw-me
Repo: NVIDIA/TensorRT-LLM PR: 14733
File: examples/visual_gen/serve/async_video_gen.py:36-37
Timestamp: 2026-05-29T08:18:27.934Z
Learning: In NVIDIA/TensorRT-LLM, Ruff’s `[tool.ruff.lint].select` enables only `D, E, F, I, PLE, W` and does NOT include `A` (flake8-builtins). Therefore, builtin-name shadowing issues like `A002` (e.g., using `format` as a function parameter) should not be treated as lint failures for this repo during code review.

Applied to files:

• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py
• tests/unittest/_torch/executor/test_adp_router.py

📚 Learning: 2026-06-25T13:04:55.284Z

Learnt from: taianz-nv
Repo: NVIDIA/TensorRT-LLM PR: 15555
File: tensorrt_llm/_torch/visual_gen/models/wan/wan_vae.py:23-24
Timestamp: 2026-06-25T13:04:55.284Z
Learning: For NVIDIA/TensorRT-LLM Python code, Ruff is configured in pyproject.toml to enable only {D, E, F, I, PLE, W}. In code review, do not request iterator/list-style cleanups tied to disabled Ruff rules B905, RUF005, or RUF007 (e.g., suggesting zip(..., strict=True), itertools.pairwise(), or rewriting `[x] + list_var` into unpacking) unless there is a separate, concrete correctness/performance/security reason. Prefer the repository’s existing plain forms for consistency.

Applied to files:

• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py
• tests/unittest/_torch/executor/test_adp_router.py

📚 Learning: 2026-03-06T11:45:02.068Z

Learnt from: ixlmar
Repo: NVIDIA/TensorRT-LLM PR: 11983
File: tensorrt_llm/_torch/pyexecutor/llm_request.py:237-240
Timestamp: 2026-03-06T11:45:02.068Z
Learning: In tensorrt_llm/_torch/pyexecutor/llm_request.py, LogProbStorage.append() relies on next(iter(prob.values())).logprob to accumulate cum_log_probs when cum_log_probs is None. This path assumes prob is a non-empty dict because TorchSampler places the sampled logprob first when num_logprobs >= 0. Therefore, no guard for empty dicts is needed here. If future changes may yield empty prob, consider adding a guard or a clearer invariant.

Applied to files:

• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py

📚 Learning: 2026-03-09T12:34:56.631Z

Learnt from: yibinl-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 12009
File: tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py:296-299
Timestamp: 2026-03-09T12:34:56.631Z
Learning: In tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py, BasicAVTransformerBlock should not be flagged for a config mismatch when config.parallel.dit_ulysses_size > 1. The function setup_sequence_parallelism() returns use_ulysses=True for dit_ulysses_size > 1, or raises a RuntimeError/ValueError/NotImplementedError; it never returns use_ulysses=False in that case. Treat this as intentional and correct; do not flag as a mismatch between raw config checks and setup_sequence_parallelism()'s result.

Applied to files:

• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py

📚 Learning: 2026-04-16T00:07:18.998Z

Learnt from: venkywonka
Repo: NVIDIA/TensorRT-LLM PR: 12944
File: tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py:654-679
Timestamp: 2026-04-16T00:07:18.998Z
Learning: When reviewing TensorRT-LLM multimodal handling, do not assume a bug/inconsistency just because `special_token_offsets` appears written in different locations for different backends. In the AutoDeploy path, `special_token_offsets` is written under `multimodal_data['layout_metadata']['special_token_offsets']` by the AD-specific input processor and is read by `ad_executor.py` via the nested `layout_metadata` lookup. In the PyTorch path, it’s written at the top level as `mm_data['special_token_offsets']` (e.g., via `tensorrt_llm/inputs/registry.py`) and is read by `model_engine.py` from the top level. These backends should not share requests, so the reviewer should verify reads/writes match the intended backend’s request structure rather than flaging cross-backend mismatches.

Applied to files:

• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py

📚 Learning: 2026-04-24T03:28:03.123Z

Learnt from: yihwang-nv
Repo: NVIDIA/TensorRT-LLM PR: 13373
File: tests/integration/defs/accuracy/test_llm_api_pytorch.py:51-61
Timestamp: 2026-04-24T03:28:03.123Z
Learning: When reviewing TensorRT-LLM attention backend code, treat `_TRTLLM_ENABLE_TRTLLM_GEN_ATTENTION` (in `tensorrt_llm._torch.attention_backend.trtllm`) as a *gate* that enables the `trtllm_gen` path, not as an identifier of the executed backend. The same `TrtllmAttention` class may handle both legacy and `trtllm_gen`, and the actual per-forward branch is also governed by checks like `trtllm_gen.is_supported(...)`. Do not expect public LLM/attention object attributes to reveal which path executed; if test assertions require confirming GEN usage, reviewers should require production-side instrumentation/counters (or other observable signals) rather than relying on existing public fields.

Applied to files:

• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py

📚 Learning: 2026-05-16T01:43:01.298Z

Learnt from: yibinl-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 14161
File: tensorrt_llm/serve/cluster_storage.py:5-5
Timestamp: 2026-05-16T01:43:01.298Z
Learning: In the NVIDIA/TensorRT-LLM codebase, do not raise a code review issue for missing NVIDIA Apache 2.0 license headers in Python files (they are intentionally not required/enforced in this repository).

Applied to files:

• tensorrt_llm/llmapi/llm_args.py
• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py

📚 Learning: 2025-12-12T03:27:08.565Z

Learnt from: tongyuantongyu
Repo: NVIDIA/TensorRT-LLM PR: 9655
File: tensorrt_llm/_torch/pyexecutor/sampler.py:3031-3031
Timestamp: 2025-12-12T03:27:08.565Z
Learning: 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. Apply this pattern to similar code paths that process batches to access simple Python data structures (lists) inside loops.

Applied to files:

• tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py

📚 Learning: 2026-02-25T01:48:05.078Z

Learnt from: yibinl-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 11683
File: tests/unittest/llmapi/_run_multi_mpi_comm_tasks.py:6-8
Timestamp: 2026-02-25T01:48:05.078Z
Learning: When reviewing cherry-pick PRs that backport commits (PRs that move fixes from one branch to another), do not suggest or require style/formatting changes in the backport. Cherry-picks should preserve the exact commit history and focus on functional correctness; style improvements are lower priority than preserving the integrity of the backported fixes. This guidance applies broadly to Python test files under tests/, ensuring reviews favor minimal, non-formatting changes that maintain correctness.

Applied to files:

• tests/unittest/_torch/executor/test_adp_router.py

📚 Learning: 2026-05-29T08:18:31.547Z

Learnt from: zhenhuaw-me
Repo: NVIDIA/TensorRT-LLM PR: 14733
File: tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py:1838-1840
Timestamp: 2026-05-29T08:18:31.547Z
Learning: In NVIDIA/TensorRT-LLM, Ruff’s `ruff select` does not enable the flake8-bandit (“S*”) rules, so Bandit checks like S108 are not enforced by the lint gate. During review of test Python code, do not flag `/tmp/...` string literals (or other situations) as Ruff S108 violations unless the repo’s Ruff configuration is updated to select flake8-bandit rules.

Applied to files:

• tests/unittest/_torch/executor/test_adp_router.py

📚 Learning: 2026-02-27T21:32:25.857Z

Learnt from: lucaslie
Repo: NVIDIA/TensorRT-LLM PR: 11796
File: tests/unittest/_torch/auto_deploy/unit/singlegpu/custom_ops/utils/test_block_table_ragged_conversion.py:713-747
Timestamp: 2026-02-27T21:32:25.857Z
Learning: When reviewing tests in this repository, lightweight benchmark tests (sub-second duration, ~100 iterations) that detect performance regressions should remain part of the standard test suite and should not be gated by environment variables or pytest markers. Only expensive/long-running benchmarks should be gated (e.g., via a pytest marker like pytest.mark.slow or a dedicated marker/flag) so they can be run selectively. This guidance applies to Python tests under tests/ and related subdirectories.

Applied to files:

• tests/unittest/_torch/executor/test_adp_router.py

📝 Walkthrough

Walkthrough

Adds a configurable policy for assigning new conversations to tensor-parallel ranks. The default remains round-robin; least_queued selects the least-loaded rank while preserving existing conversation-to-rank stickiness.

Changes

Conversation placement policy

Layer / File(s)	Summary
Placement configuration contract tensorrt_llm/llmapi/llm_args.py, tensorrt_llm/usage/llm_args_golden_manifest.json	Adds the typed round_robin/least_queued option and corresponding manifest schema entry.
Router policy wiring and selection tensorrt_llm/_torch/pyexecutor/scheduler/adp_router.py	Forwards and normalizes the policy, then selects least-loaded or round-robin ranks for new or unpinned conversations.
Placement behavior validation tests/unittest/_torch/executor/test_adp_router.py	Tests defaults, invalid-value fallback, least-loaded assignment, batch leveling, stickiness, factory forwarding, and deterministic routing.

Estimated code review effort: 3 (Moderate) | ~20 minutes

Sequence Diagram(s)

sequenceDiagram
  participant ADPRouter
  participant ConversationAwareADPRouter
  participant RankSelection
  ADPRouter->>ConversationAwareADPRouter: Pass new_conv_placement
  ConversationAwareADPRouter->>RankSelection: Select least-loaded rank when configured
  ConversationAwareADPRouter->>RankSelection: Select round-robin rank otherwise
  RankSelection-->>ConversationAwareADPRouter: Return target rank

Loading

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description explains the change, but it omits the required Test Coverage and PR Checklist sections.	Add a Test Coverage section with the relevant tests and complete the PR Checklist items required by the repository template.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the PR's main change and uses the required [None][feat] format.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

_{Comment @coderabbitai help to get the list of available commands.}

[tensorrt-cicd]

PR_Github #58882 [ run ] triggered by Bot. Commit: 01da19e Link to invocation

[tensorrt-cicd]

PR_Github #58882 [ run ] completed with state SUCCESS. Commit: 01da19e
/LLM/main/L0_MergeRequest_PR pipeline #47425 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

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

CI Agent Failure Analysis

Link to invocation

[lancelly]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58979 [ run ] triggered by Bot. Commit: 01da19e Link to invocation

[tensorrt-cicd]

PR_Github #58979 [ run ] completed with state SUCCESS. Commit: 01da19e
/LLM/main/L0_MergeRequest_PR pipeline #47510 completed with status: 'SUCCESS'

CI Report

Link to invocation

[YihuiLu512]

LGTM

[QiJune]

LGTM. The deterministic tie-break can create some cross-batch affinity skew for sparse or bursty traffic, as noted in the discussion, but this is non-blocking for the high-concurrency workload targeted by this change.