[asfiyab-nvidia]

Description: [TRTLLM-17086][feat] Agg Helix KV cache, part 1: prefill allocation and sizing

Epic: TRTLLM-17086 (part of TRTLLM-16624, Helix for aggregated serving), design doc
Tickets: TRTLLM-17124, TRTLLM-17126

Summary

First step toward running Helix context parallelism in aggregated serving, on the KV cache side. In aggregated serving, requests now keep the full prompt on every CP rank instead of being sliced per rank, so KVCacheManagerV2 can allocate prefill KV in global ledger blocks. Agg Helix stays off by default until it is validated end to end.

Changes

TRTLLM-17124: allow prefill KV allocation for agg Helix requests

• merge_requests: aggregated Helix requests are merged with merge_requests_to_llm_requests, so every CP rank holds the full prompt (prompt_len == total_input_len_cp). Disaggregated Helix still uses merge_helix_requests, which deals the prompt's blocks round-robin across ranks. PyExecutor passes is_disagg=self.kv_cache_transceiver is not None.
• KVCacheManagerV2.resize_context: the Helix guard now raises only under disaggregated serving (self.is_disagg), where requests hold just this rank's slice. Aggregated requests allocate as usual, sized in global ledger tokens. The check uses the serving mode rather than the two lengths, so every rank reaches the same decision.
• create_py_executor: Helix outside disaggregated serving fails at startup unless TRTLLM_HELIX_AGG_EXPERIMENTAL=1 (development only).

TRTLLM-17126: Helix KV cache sizing for aggregated serving

• _configure_helix_kv_cache_capacity: under agg Helix, a free_gpu_memory_fraction of 0.9 (the default) is lowered to 0.6, with a warning, to leave room for prefill activations, since Helix skips memory profiling. The check compares values because trtllm-serve always passes 0.9, so an explicit 0.9 is lowered too; any other fraction is used as given. An explicit max_tokens or max_gpu_total_bytes still overrides it.

Temporary pieces (removed in TRTLLM-17127)

• The startup guard and TRTLLM_HELIX_AGG_EXPERIMENTAL.
• The 0.6 fraction override (_HELIX_AGG_FREE_GPU_MEMORY_FRACTION), which profiling will replace.

Not affected

Disaggregated Helix: requests are still sliced by merge_helix_requests, generation servers never call resize_context, and the fraction override applies only when the server is not disaggregated.

Test Coverage

• test_request_utils.py: aggregated Helix keeps the full prompt on every CP rank (prompt_len == total_input_len_cp); the existing sliced case now runs as disaggregated.
• test_kv_cache_manager_v2_helix_superblock.py: aggregated Helix request accepted; disaggregated Helix request rejected, including a slice that equals the full prompt.
• test_mamba_cache_manager_v2_helix.py (new): MambaHybridCacheManagerV2 quota-to-token conversion under Helix for cp=2/4/8. It returns whole ledger blocks of global tokens and charges KDA state per request slot.

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.