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

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -96,6 +96,10 @@
 # accepts. Matches the KvCacheConfig.tokens_per_block default.
 _STANDALONE_DRAFT_TOKENS_PER_BLOCK = 32
 
+# Aggregated helix skips memory profiling, so its default KV fraction leaves
+# room for prefill activations. TODO(TRTLLM-17127): remove once helix profiles.
+_HELIX_AGG_FREE_GPU_MEMORY_FRACTION = 0.6
+
 
 def ceil_div(a: int, b: int) -> int:
     return (a + b - 1) // b
@@ -949,6 +953,24 @@ def __init__(
             model_engine)
         self._is_kv_cache_manager_v2 = issubclass(self._kv_cache_manager_cls,
                                                   KVCacheManagerV2)
+        if mapping.has_cp_helix() and not is_disagg:
+            if not self._is_kv_cache_manager_v2:
+                raise ValueError(
+                    "Aggregated helix context parallelism requires "
+                    f"KVCacheManagerV2, got {self._kv_cache_manager_cls.__name__}: "
+                    "only V2 sizes prefill KV in global ledger blocks shared "
+                    "by every CP rank. Set "
+                    "kv_cache_config.use_kv_cache_manager_v2=True and disable "
+                    "features that fall back to V1.")
+            # Read from the engine rather than llm_args: an SM-version fallback
+            # may have turned chunked prefill off after the args were parsed.
+            if not model_engine.attn_runtime_features.chunked_prefill:
+                raise ValueError(
+                    "Aggregated helix context parallelism requires chunked "
+                    "prefill: the context phase gathers cached KV across CP "
+                    "ranks only in the chunked-context path. Set "
+                    "enable_chunked_prefill=True on a GPU that supports it "
+                    "for MLA.")
         self._disable_overlap_scheduler = llm_args.disable_overlap_scheduler
         self._draft_config = draft_config
         self._skip_est = skip_est
@@ -1728,6 +1750,16 @@ def _configure_helix_kv_cache_capacity(self) -> None:
                         "the explicitly configured quota.")
             return
         fraction = self._kv_cache_config.free_gpu_memory_fraction
+        default_fraction = KvCacheConfig.model_fields[
+            "free_gpu_memory_fraction"].default
+        if fraction == default_fraction and not self._is_disagg:
+            fraction = _HELIX_AGG_FREE_GPU_MEMORY_FRACTION
+            logger.warning(
+                "Helix CP (aggregated): capacity profiling is unsupported, so "
+                f"the default free_gpu_memory_fraction {default_fraction} is "
+                f"lowered to {fraction} to leave room for prefill "
+                "activations. Set kv_cache_config.max_tokens or "
+                "max_gpu_total_bytes to override.")
         free_mem, _total = torch.cuda.mem_get_info()
         budget_bytes = int(free_mem * fraction)
         if budget_bytes <= 0:

tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py

@@ -3883,12 +3883,11 @@ def resize_context(self, req: LlmRequest, num_tokens: int) -> bool:
         assert not req.is_disagg_generation_init_state, (
             f"req {req.py_request_id}: use prepare_disagg_gen_init"
         )
-        if self._has_cp_helix and not req.is_dummy_request:
+        if self._has_cp_helix and self.is_disagg and not req.is_dummy_request:
             raise ValueError(
-                "resize_context is not helix-aware: its rank-local chunk "
-                "target would under-size the global ledger. Helix requests "
-                "are disagg-generation-only and must never take the context "
-                "path."
+                "resize_context is not available under disaggregated helix: "
+                "requests hold only this rank's slice of the prompt, which "
+                "would under-size the global ledger."
             )
         kv_cache = self.kv_cache_map.get(req.py_request_id)
         if kv_cache is None:

tensorrt_llm/_torch/pyexecutor/py_executor.py

@@ -6780,7 +6780,8 @@ def _fetch_new_requests(
                               cp_rank=self.dist.cp_rank,
                               cp_size=self.dist.cp_size,
                               exclude_last_generation_logits=self.
-                              _should_exclude_last_generation_logits())
+                              _should_exclude_last_generation_logits(),
+                              is_disagg=self.kv_cache_transceiver is not None)
 
     def _handle_special_queue_items(
             self,

tensorrt_llm/_torch/pyexecutor/py_executor_creator.py

@@ -675,6 +675,15 @@ def _create_py_executor(
     enable_sleep = sleep_config is not None
 
     cache_transceiver_config = llm_args.cache_transceiver_config
+    # TODO(TRTLLM-17127): remove once aggregated helix is validated end to end.
+    if (mapping.has_cp_helix()
+            and not is_disagg_enabled(cache_transceiver_config)
+            and os.environ.get("TRTLLM_HELIX_AGG_EXPERIMENTAL", "0") != "1"):
+        raise NotImplementedError(
+            "Helix context parallelism is supported only on disaggregated "
+            "generation servers; aggregated serving with Helix is not "
+            "supported yet (set TRTLLM_HELIX_AGG_EXPERIMENTAL=1 for "
+            "development only).")
 
     has_spec_drafter = False
     if spec_config is not None:

tensorrt_llm/_torch/pyexecutor/request_utils.py

@@ -334,6 +334,7 @@ def merge_requests(
     cp_rank: int,
     cp_size: int,
     exclude_last_generation_logits: bool,
+    is_disagg: bool,
 ) -> List[LlmRequest]:
     """Merge RequestQueueItems to LlmRequests based on CP configuration.
 
@@ -347,6 +348,9 @@ def merge_requests(
         cp_rank: Current CP rank.
         cp_size: Total number of CP ranks.
         exclude_last_generation_logits: Whether to exclude last generation logits.
+        is_disagg: Whether this executor serves disaggregated generation. Only
+            disaggregated Helix slices the prompt per rank; aggregated Helix
+            keeps the full prompt on every rank.
 
     Returns:
         List of LlmRequest objects.
@@ -357,13 +361,15 @@ def merge_requests(
     if "cp_type" in cp_config:
         cp_type = cp_config["cp_type"]
         if cp_type == CpType.HELIX:
-            return merge_helix_requests(
-                new_requests,
-                cp_rank=cp_rank,
-                cp_size=cp_size,
-                tokens_per_block=cp_config["tokens_per_block"],
-                exclude_last_generation_logits=exclude_last_generation_logits,
-            )
+            if is_disagg:
+                return merge_helix_requests(
+                    new_requests,
+                    cp_rank=cp_rank,
+                    cp_size=cp_size,
+                    tokens_per_block=cp_config["tokens_per_block"],
+                    exclude_last_generation_logits=exclude_last_generation_logits,
+                )
+            return merge_requests_to_llm_requests(new_requests, exclude_last_generation_logits)
         else:
             raise NotImplementedError(f"Unsupported cp type {cp_type.name}.")
 

tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py

@@ -252,13 +252,14 @@ def test_helix_quota_fallback_sets_rank_local_bytes(monkeypatch: pytest.MonkeyPa
 
     from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator
 
-    def creator(max_gpu_total_bytes, max_tokens):
+    def creator(max_gpu_total_bytes, max_tokens, *, is_disagg=False, fraction=0.5):
         return SimpleNamespace(
+            _is_disagg=is_disagg,
             _mapping=SimpleNamespace(cp_size=4),
             _kv_cache_config=SimpleNamespace(
                 max_gpu_total_bytes=max_gpu_total_bytes,
                 max_tokens=max_tokens,
-                free_gpu_memory_fraction=0.5,
+                free_gpu_memory_fraction=fraction,
             ),
         )
 
@@ -278,6 +279,14 @@ def creator(max_gpu_total_bytes, max_tokens):
     assert KvCacheCreator._configure_helix_kv_cache_capacity(c) is None
     assert c._kv_cache_config.max_gpu_total_bytes == 500_000
     assert c._kv_cache_config.max_tokens is None
+    # Default fraction: aggregated helix lowers it to leave room for prefill
+    # activations; disaggregated helix keeps it.
+    c = creator(0, None, fraction=0.9)
+    KvCacheCreator._configure_helix_kv_cache_capacity(c)
+    assert c._kv_cache_config.max_gpu_total_bytes == 600_000
+    c = creator(0, None, is_disagg=True, fraction=0.9)
+    KvCacheCreator._configure_helix_kv_cache_capacity(c)
+    assert c._kv_cache_config.max_gpu_total_bytes == 900_000
     # Degenerate free memory: actionable error instead of a deep assert.
     monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (0, 2_000_000))
     with pytest.raises(ValueError, match="free memory"):
@@ -369,6 +378,55 @@ def test_scheduler_allocation_failure_raises_under_helix() -> None:
         )
 
 
+def _context_resize_mgr(is_disagg: bool, resizes: list) -> SimpleNamespace:
+    kv = SimpleNamespace(capacity=0, resize=lambda cap: resizes.append(cap) or True)
+    return SimpleNamespace(
+        _has_cp_helix=True,
+        is_disagg=is_disagg,
+        kv_cache_map={7: kv},
+        num_extra_kv_tokens=0,
+        _connector_reservations_enabled=lambda: False,
+        _fill_fresh_kv_pages=lambda req_id: None,
+        _log_window_crossing=lambda *args: None,
+    )
+
+
+def _context_req(prompt_len: int, total_input_len_cp: int) -> SimpleNamespace:
+    return SimpleNamespace(
+        py_request_id=7,
+        is_disagg_generation_init_state=False,
+        is_dummy_request=False,
+        prompt_len=prompt_len,
+        total_input_len_cp=total_input_len_cp,
+        context_current_position=64,
+        is_first_context_chunk=False,
+    )
+
+
+def test_resize_context_allows_agg_helix_request() -> None:
+    """An agg helix request holds the full prompt on every rank, so its
+    context target is already in global ledger tokens."""
+    resizes = []
+    req = _context_req(prompt_len=200, total_input_len_cp=200)
+    assert KVCacheManagerV2.resize_context(
+        _context_resize_mgr(is_disagg=False, resizes=resizes), req, num_tokens=32
+    )
+    assert resizes == [64 + 32]
+
+
+@pytest.mark.parametrize("prompt_len", [50, 200])
+def test_resize_context_rejects_disagg_helix_request(prompt_len: int) -> None:
+    """A disagg helix request holds only this rank's slice of the prompt; a
+    rank-local target would under-size the global ledger. The guard keys on
+    disagg mode, so a slice that happens to equal the full prompt (a
+    single-block prompt on rank 0) is rejected too."""
+    req = _context_req(prompt_len=prompt_len, total_input_len_cp=200)
+    with pytest.raises(ValueError, match="disaggregated helix"):
+        KVCacheManagerV2.resize_context(
+            _context_resize_mgr(is_disagg=True, resizes=[]), req, num_tokens=32
+        )
+
+
 def test_dummy_frozen_fields_sum_invariant() -> None:
     """The frozen dummy fiction (last rank active, shared synthetic global
     length) keeps the same books as real requests: per-rank lengths sum to

tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager_v2_helix.py

@@ -0,0 +1,48 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Unit tests for MambaHybridCacheManagerV2 under helix context parallelism."""
+
+import pytest
+
+from tensorrt_llm._torch.pyexecutor.kv_cache.mamba_cache_manager import MambaHybridCacheManagerV2
+
+
+@pytest.mark.parametrize("cp_size", [2, 4, 8])
+def test_mamba_hybrid_quota_to_tokens_under_helix(cp_size: int) -> None:
+    """A rank-local byte quota converts to GLOBAL tokens in whole ledger
+    blocks; recurrent state is charged per request slot, not per token."""
+    phys, attn_bytes_per_token, state_bytes_per_slot = 32, 10, 1000
+    num_slots = 4 + 1  # resident sequences + reserved dummy slots
+    mgr = object.__new__(MambaHybridCacheManagerV2)
+    mgr._has_cp_helix = True
+    mgr._helix_cp_size = cp_size
+    mgr.tokens_per_block = phys
+    mgr._ledger_tokens_per_block = phys * cp_size
+    mgr.kv_cache_config = None
+    mgr._num_reserved_dummy_slots = 1
+    mgr._max_resident_sequences = lambda: 4
+    mgr._num_ssm_snapshots_for_capacity = lambda max_tokens, config: 0
+    mgr._mamba_state_bytes_per_slot = lambda: state_bytes_per_slot
+    mgr._attention_cache_bytes_per_token = lambda: attn_bytes_per_token
+    mgr._get_quota_from_max_tokens_impl = lambda tokens: tokens * attn_bytes_per_token
+
+    # State takes 5 slots; the rest buys 7 full physical pages plus a
+    # partial page that is never usable.
+    quota = num_slots * state_bytes_per_slot + 7 * phys * attn_bytes_per_token + 100
+    tokens = mgr._get_max_tokens_from_quota(quota)
+
+    assert tokens == 7 * phys * cp_size
+    assert mgr._get_quota_from_max_tokens(tokens) <= quota
+    assert mgr._get_quota_from_max_tokens(tokens + 1) > quota

tests/unittest/_torch/executor/test_py_executor.py

@@ -2739,6 +2739,7 @@ def _run_non_adp_fetch(active_requests, *, exclude_retiring_requests, max_num_ac
         enable_iter_perf_stats=False,
         num_fetch_requests=0,
         dist=Mock(cp_config={}, cp_rank=0, cp_size=1, rank=0),
+        kv_cache_transceiver=None,
         _should_exclude_last_generation_logits=lambda: False,
         _kv_load_tracker=KvLoadTracker(None, rank=0, enabled=False),
         _kv_cache_capacity={},
@@ -2824,6 +2825,7 @@ def test_non_adp_fetch_pops_a_full_batch_beside_retiring_requests(
         enable_iter_perf_stats=False,
         num_fetch_requests=0,
         dist=Mock(cp_config={}, cp_rank=0, cp_size=1, rank=0, tp_size=1),
+        kv_cache_transceiver=None,
         _should_exclude_last_generation_logits=lambda: False,
         _kv_load_tracker=KvLoadTracker(None, rank=0, enabled=False),
         _kv_cache_capacity={},

tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py

@@ -303,6 +303,7 @@ def _run_create_py_executor(
         tp_size=1,
         pp_size=1,
         has_pp=lambda: False,
+        has_cp_helix=lambda: False,
         enable_attention_dp=False,
         is_last_pp_rank=lambda: True,
     )

tests/unittest/_torch/executor/test_request_utils.py

@@ -445,7 +445,12 @@ def test_merge_requests_default(mock_convert):
 
     requests = [RequestQueueItem(1, Mock()), RequestQueueItem(2, Mock())]
     result = merge_requests(
-        requests, cp_config={}, cp_rank=0, cp_size=1, exclude_last_generation_logits=False
+        requests,
+        cp_config={},
+        cp_rank=0,
+        cp_size=1,
+        exclude_last_generation_logits=False,
+        is_disagg=False,
     )
 
     assert len(result) == 2
@@ -483,6 +488,7 @@ def test_merge_requests_with_helix_cp_config():
             cp_rank=rank,
             cp_size=4,
             exclude_last_generation_logits=False,
+            is_disagg=True,
         )
 
         # Verify the result.
@@ -507,6 +513,37 @@ def test_merge_requests_with_helix_cp_config():
             assert llm_request.get_tokens(0) == [7, 8]
 
 
+def test_merge_requests_agg_helix_keeps_full_prompt():
+    """Aggregated helix keeps the full prompt on every CP rank, so every rank
+    sees identical request lengths."""
+    input_tokens = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13]
+    executor_request = trtllm.Request(
+        input_token_ids=input_tokens,
+        max_tokens=5,
+        streaming=False,
+        sampling_config=trtllm.SamplingConfig(),
+        output_config=trtllm.OutputConfig(),
+    )
+    request_item = RequestQueueItem(id=1, request=executor_request)
+    cp_config = {"cp_type": CpType.HELIX, "tokens_per_block": 2}
+
+    for rank in range(4):
+        result = merge_requests(
+            [request_item],
+            cp_config=cp_config,
+            cp_rank=rank,
+            cp_size=4,
+            exclude_last_generation_logits=False,
+            is_disagg=False,
+        )
+
+        assert len(result) == 1
+        llm_request = result[0]
+        assert llm_request.get_tokens(0) == input_tokens
+        assert llm_request.prompt_len == len(input_tokens)
+        assert llm_request.total_input_len_cp == len(input_tokens)
+
+
 def test_get_from_waiting_queue():
     """Test getting items from waiting queue."""
     # Add items to waiting queue