[reasonsolo]

@coderabbitai summary

Description

Two opt-in context-scheduling policies for KVCacheV2Scheduler, aimed at attention-DP context servers with long, multi-turn prompts. Both are controlled by environment variables and are off by default. With the variables unset or 0, scheduling and routing are unchanged.

It also fixes a rank-divergent KV cache event flush in the overlap executor loop under attention DP (section 3), which these context servers hit with KV cache manager V2 and buffered KV events.

The GB300 Kimi K3 DWDP perf-sanity configs added in #20163 already set these variables (tests/scripts/perf-sanity/disaggregated/gb300_kimi-k3-fp4-dspark_agentx_*.yaml). Until now main had no reader for them. This PR keeps the variable names identical, so those configs take effect without changes.

Original design and implementation by @lancelly (Liao Lanyu); the KV event flush fix is originally by @liji-nv (Jin Li). Equal-cost chunking was first proposed in #18528. This PR ports it, the follow-ups and long-prefill demotion onto the current V2 scheduler.

1. Equal-cost context chunking by KV depth

Under attention DP every rank runs in lockstep, so an iteration lasts as long as its most expensive rank. A context chunk of n tokens at KV depth kv costs about (kv_cost_offset + kv) * n. A full-budget chunk of a deep request therefore takes much longer than the same number of tokens on a shallow rank, and the other ranks wait.

With equal-cost chunking, each rank caps its context chunks against a shared cost budget:

B = (kv_cost_offset + kv_depth_threshold) * max_num_tokens  [+ R * kv_depth_threshold]
chunk cost = (kv_cost_offset + kv_start) * n               [+ R * kv_start]

This evens out the per-rank iteration cost without any cross-rank coordination. Requests shallower than kv_depth_threshold keep the plain token cap because the token cap binds first. Leftover cost budget goes to the next context request (backfill). The first context request of an iteration always gets at least max(chunk_unit_size, max_num_tokens // 4) tokens, so a very deep request is never starved into unit-sized chunks that each re-read the whole prefix.

Variable	Default	Meaning
TLLM_V2_CTX_COST_KV_OFFSET	0 (off)	KV-token equivalents of the KV-independent per-token work (a model x kernel x precision x GPU property).
TLLM_V2_CTX_COST_KV_DEPTH_THRESHOLD	0 (off)	KV depth above which chunks start to shrink (a workload property).
TLLM_V2_CTX_COST_PER_CHUNK_DEPTH_TOKENS	0	Optional per-chunk prefix term R. Each chunk also costs R * kv_start, and the budget grows by R * kv_depth_threshold.

The feature is enabled only when both KV_OFFSET and KV_DEPTH_THRESHOLD are positive, max_num_tokens is set, and context chunking is enabled. Otherwise BudgetTracker gets no cost budget and the cost cap is skipped.

Calibration example. examples/scheduler_cost_calibration/fit_context_cost.py fits KV_OFFSET and KV_DEPTH_THRESHOLD from the per-rank iteration log of an ordinary token-budget run (TLLM_PROFILE_LOG_RANKS=all). It finds "monopoly spans", where one rank's single deep request determines the lockstep iteration time. It fits the time-vs-depth slope with Theil-Sen, which is robust to stall outliers, and derives kv_cost_offset from the intercept. It takes kv_depth_threshold as a percentile of the token-weighted depth distribution. It also reports a linearity check, the measured lockstep waste and the predicted gain, then prints ready-to-paste values. The README covers the workflow and the per-chunk term, which the script does not fit.

2. Long-prefill demotion

V2 context scheduling is FCFS. A prefill with hundreds of thousands of uncached tokens takes the whole token budget for every iteration it needs, and the short, mostly cached turns queued behind it wait. With demotion, a request whose estimated uncached prefill exceeds the threshold moves to the tail of the iteration's context order, keeping FCFS within each class. It takes whatever token budget the short requests leave, so it is reordered but never capped: alone in the queue, it still gets the full budget. Once it is mid-prefill, the exact remainder (context_remaining_length) releases it when no more than the threshold is left.

Variable	Default	Meaning
TLLM_V2_CTX_DEMOTE_UNCACHED_TOKENS	0 (off)	Demote a request whose estimated uncached prefill exceeds this many tokens.
TLLM_V2_CTX_DEMOTE_MAX_ITERS	0 (no bound)	After this many consecutive demoted iterations, a request is let through in arrival order until it is scheduled.

The estimate (LlmRequest.py_uncached_prefill_tokens) is published by ConversationAwareADPRouter while demotion is on:

• a conversation's first turn: the whole prompt;
• a later turn: the growth since the previous turn's prompt;
• a request without a conversation_id: the whole prompt.

The history is keyed by conversation_id, not by the routing key, so sub-agents that share a parent's affinity id keep their own history. It is LRU-bounded by max_sessions. A request without an estimate is never demoted, for example when another ADP router is in use or attention DP is off. The scheduler does not probe the radix tree, which would walk every prompt on every iteration. Demotion is also off when context chunking is disabled.

3. Fix: flush KV cache events outside the rank-divergent batch path under attention DP

Issue.

With attention DP and buffered KV cache events (event_buffer_max_size > 0, not streaming), KV cache manager V2 builds its KVCacheEventManager with attention_dp_gather=Distributed.get(mapping).allgather (kv_cache_manager_v2.py). EventManager::flushIterationEvents calls that gather synchronously (eventManager.cpp), so flush_iteration_events() is a collective across the attention-DP group.

In _executor_loop_overlap, the flush happened inside _process_previous_batch(). That function only runs under if self.previous_batch is not None and should_process_previous_batch:. Both conditions can differ between ranks:

• should_process_previous_batch = can_queue or not can_queue_this_rank. When some ranks scheduled work and others did not, can_queue is false on every rank. Ranks with local work then defer their previous batch, and idle ranks process theirs.
• An idle rank then clears previous_batch, while a busy rank keeps it. On the next iteration the ranks disagree on previous_batch is not None.

The loop already handles this divergence for its other collectives. It pairs _enqueue_responses([]) in the else branch, and calls _flush_pending_transfer_responses(), handle_timeouts_synced() and _flush_iter_stats_synced() outside the branch. The KV event flush was the remaining collective inside the branch. With attention DP, the overlap scheduler, KV cache manager V2 and buffered KV events, it could be issued on only some ranks. The result is a hang or an allgather that pairs with an unrelated collective.

The KV pool rebalance hook (_maybe_rebalance_kv_pools) has the same gap. It calls _process_previous_batch() through _consume_previous_batch_for_rebalance(), which returns early on ranks without a previous_batch.

Fix.

• _process_previous_batch() no longer flushes KV cache events.
• New _add_kv_cache_events_synced(processed_previous_batch) helper:
  • With attention DP, every rank flushes on every call. An empty flush only joins the gather and publishes nothing.
  • Without attention DP, it flushes only when this rank processed a previous batch, which is the same cadence as before.
  • With KV cache events disabled, it does nothing, as before.
• The overlap loop calls the helper right after the previous-batch if/else, before handle_timeouts_synced(). Every rank reaches that point in the same order.
• The rebalance hook calls it after _consume_previous_batch_for_rebalance(). Every rank reaches that point once _agreed_need_adjustment() agrees.
• The non-overlap _executor_loop already flushes under if can_queue, which is rank-agreed under attention DP. It is unchanged.

Behaviour without attention DP, or with KV cache events disabled, is unchanged. With attention DP, ranks also flush on iterations where they did not process a previous batch. Those flushes carry no events.

Interaction with existing V2 scheduler features

• Prefix-aware skip and encoder/cross requests. Demotion only reorders the phase-2 pending_ctx list. Contributed-block collection, the prefix-skip deferral, encoder scheduling (encoder requests are never demoted) and preempt_for_pages run unchanged on the new order.
• Remote-tail replay. The demotion counters are updated after the replay retry pass, so a retried iteration is counted once.
• Encoder replay and drafts. The cost ledger charges what budget.commit charges, including replay rows and last-chunk draft tokens. The cost cap is computed on the new-token chunk only.
• Multimodal block alignment. This runs after the cost cap. A snap-up for bidirectional MM attention may slightly exceed the cost budget, bounded by the token budget.
• Failed first-chunk admission. A first chunk that the cost budget cannot afford is skipped, and release_context_admission releases its prefix holds as for any other failed admission.

Test Coverage

• tests/unittest/_torch/executor/kv_cache/test_kv_cache_v2_scheduler.py
  • TestEqualCostChunking: off by default, off with one knob or without chunking, unlimited BudgetTracker without a cost budget, shallow requests unaffected, the deep-request formula, leftover backfill, the head-of-line floor, and the per-chunk depth term including a backfill that cannot pay its depth term.
  • TestLongContextDemotion: off by default and without chunking, estimated-long requests wait behind short ones from the first chunk, no estimate means no demotion, the exact remainder releases the tail, and the aging bound with count reset and no leaked entries.
• tests/unittest/_torch/executor/test_adp_router.py::TestConversationAwareADPRouter
  • The estimate is published only while demotion is on: opener, follow-up turn, and a request without a conversation id.
  • The history is per-conversation under sub-agent affinity.
  • The history is LRU-bounded.

The existing V2 scheduler and ADP router suites pass unchanged, which shows the defaults are a no-op.

• tests/unittest/_torch/executor/test_py_executor.py::TestKvCacheEventFlushRankSymmetry (CPU): drives the real _executor_loop_overlap / _process_previous_batch / _add_kv_cache_events (or _maybe_rebalance_kv_pools) with a mocked KV cache manager, and records the ordered trace of attention-DP collectives on each simulated rank.
  • test_adp_ranks_flush_kv_cache_events_in_lockstep: one rank defers its previous batch and one processes it, and their collective traces must match. Fails on main.
  • test_adp_rank_without_previous_batch_joins_kv_cache_event_gather: ranks disagree on previous_batch is not None. Fails on main.
  • test_rebalance_flushes_kv_cache_events_on_every_adp_rank: the rebalance hook flushes on every attention-DP rank, including one without a previous_batch (that case fails on main).
  • test_without_adp_flushes_once_per_processed_previous_batch: non-attention-DP cadence is unchanged.
  • test_kv_cache_events_disabled_never_flushes: no flush when KV cache events are off.
• tests/unittest/_torch/disaggregation/test_disagg_loop_transcript.py: the minimal executor stub now sets enable_kv_cache_events = False, which a real PyExecutor always sets.

PR Checklist

Please review the following before submitting your PR:

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

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

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

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

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

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

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

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

GitHub Bot Help

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