From ef000003061d7ca0255c324ae30a6a87eaaaea4a Mon Sep 17 00:00:00 2001 From: Asfiya Baig Date: Mon, 19 Oct 2026 16:26:41 -0700 Subject: [PATCH 1/7] [TRTLLM-17124][feat] Allow resize_context for full-prompt Helix requests Narrow the Helix guard in KVCacheManagerV2.resize_context to requests sliced to this rank's tokens, so requests holding the full prompt on every rank can allocate KV during prefill. Reject Helix outside disaggregated generation servers at startup until aggregated Helix serving is complete. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Asfiya Baig --- .../kv_cache/kv_cache_manager_v2.py | 12 +++-- .../_torch/pyexecutor/py_executor_creator.py | 6 +++ ...st_kv_cache_manager_v2_helix_superblock.py | 47 +++++++++++++++++++ 3 files changed, 60 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 1e4ef77fdc2a..4070aab87a8b 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -3883,12 +3883,14 @@ 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 not req.is_dummy_request + and req.prompt_len != req.total_input_len_cp + ): 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 needs the full prompt under helix: a request " + "sliced to this rank's tokens would under-size the global ledger." ) kv_cache = self.kv_cache_map.get(req.py_request_id) if kv_cache is None: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 7a95ac8c2682..e7def00877c1 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -675,6 +675,12 @@ def _create_py_executor( enable_sleep = sleep_config is not None cache_transceiver_config = llm_args.cache_transceiver_config + if mapping.has_cp_helix( + ) and not is_disagg_enabled(cache_transceiver_config): + raise NotImplementedError( + "Helix context parallelism is supported only on disaggregated " + "generation servers; aggregated serving with Helix is not " + "supported yet.") has_spec_drafter = False if spec_config is not None: diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py index 9d3670966646..6225ed81abef 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py @@ -369,6 +369,53 @@ def test_scheduler_allocation_failure_raises_under_helix() -> None: ) +def _context_resize_mgr(cp_rank: int, cp_size: int, resizes: list) -> SimpleNamespace: + kv = SimpleNamespace(capacity=0, resize=lambda cap: resizes.append(cap) or True) + return SimpleNamespace( + _has_cp_helix=True, + _helix_cp_rank=cp_rank, + _helix_cp_size=cp_size, + 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_full_prompt_helix_request() -> None: + """An agg helix request holds the full prompt on every rank, so its + context target is already in global ledger tokens and is rank-invariant.""" + cp_size = 4 + resizes = [] + for r in range(cp_size): + req = _context_req(prompt_len=200, total_input_len_cp=200) + assert KVCacheManagerV2.resize_context( + _context_resize_mgr(r, cp_size, resizes), req, num_tokens=32 + ) + assert resizes == [64 + 32] * cp_size + + +def test_resize_context_rejects_sliced_helix_request() -> None: + """A disagg helix request holds only this rank's slice of the prompt; a + rank-local target would under-size the global ledger.""" + req = _context_req(prompt_len=50, total_input_len_cp=200) + with pytest.raises(ValueError, match="needs the full prompt"): + KVCacheManagerV2.resize_context(_context_resize_mgr(0, 4, []), 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 From ec3b8d9b535d43a61d049aec1439357e6f03ea5e Mon Sep 17 00:00:00 2001 From: Asfiya Baig Date: Tue, 20 Oct 2026 10:05:01 -0700 Subject: [PATCH 2/7] [TRTLLM-17124][feat] Allow aggregated Helix behind an experimental env var Keep rejecting Helix outside disaggregated generation servers at startup, but let developers bypass the check with TRTLLM_HELIX_AGG_EXPERIMENTAL=1 while aggregated Helix serving is being built. Signed-off-by: Asfiya Baig Co-Authored-By: Claude Opus 5.5 (1M context) --- tensorrt_llm/_torch/pyexecutor/py_executor_creator.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index e7def00877c1..735bb249dbf8 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -675,12 +675,14 @@ def _create_py_executor( enable_sleep = sleep_config is not None cache_transceiver_config = llm_args.cache_transceiver_config - if mapping.has_cp_helix( - ) and not is_disagg_enabled(cache_transceiver_config): + 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.") + "supported yet (set TRTLLM_HELIX_AGG_EXPERIMENTAL=1 for " + "development only).") has_spec_drafter = False if spec_config is not None: From 02004c4902c0ee509809a9251a28ad5ac67c60a5 Mon Sep 17 00:00:00 2001 From: Asfiya Baig Date: Tue, 20 Oct 2026 12:05:29 -0700 Subject: [PATCH 3/7] [TRTLLM-17126][feat] Helix KV cache sizing for aggregated serving Under aggregated Helix, lower the default free_gpu_memory_fraction to 0.6 so the unprofiled KV budget leaves room for prefill activations. Add MambaHybridCacheManagerV2 quota-to-token tests under Helix. Co-Authored-By: Claude Opus 5.5 (1M context) Signed-off-by: Asfiya Baig --- tensorrt_llm/_torch/pyexecutor/_util.py | 14 ++++++ .../_torch/pyexecutor/py_executor_creator.py | 1 + .../test_mamba_cache_manager_v2_helix.py | 48 +++++++++++++++++++ 3 files changed, 63 insertions(+) create mode 100644 tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager_v2_helix.py diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 9ad518e3db5f..30d3bcdfda6f 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/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 @@ -1728,6 +1732,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: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 735bb249dbf8..4c3daa85debf 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py @@ -675,6 +675,7 @@ 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"): diff --git a/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager_v2_helix.py b/tests/unittest/_torch/executor/kv_cache/test_mamba_cache_manager_v2_helix.py new file mode 100644 index 000000000000..6fcb160b6dc2 --- /dev/null +++ b/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 From cee194a735032d90ee704b85cebd88c0677cb9d9 Mon Sep 17 00:00:00 2001 From: Asfiya Baig Date: Wed, 21 Oct 2026 18:24:28 +0000 Subject: [PATCH 4/7] [TRTLLM-17124][feat] Keep the full prompt per rank for aggregated Helix Route aggregated Helix requests through merge_requests_to_llm_requests so every CP rank holds the full prompt; disaggregated Helix keeps slicing via merge_helix_requests. The resize_context guard now raises only under disaggregated Helix, so every rank reaches the same decision. Signed-off-by: Asfiya Baig --- .../kv_cache/kv_cache_manager_v2.py | 11 ++---- tensorrt_llm/_torch/pyexecutor/py_executor.py | 3 +- .../_torch/pyexecutor/request_utils.py | 20 ++++++---- ...st_kv_cache_manager_v2_helix_superblock.py | 36 +++++++++-------- .../_torch/executor/test_request_utils.py | 39 ++++++++++++++++++- 5 files changed, 76 insertions(+), 33 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py index 4070aab87a8b..80a9faf394b3 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache/kv_cache_manager_v2.py @@ -3883,14 +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 - and req.prompt_len != req.total_input_len_cp - ): + if self._has_cp_helix and self.is_disagg and not req.is_dummy_request: raise ValueError( - "resize_context needs the full prompt under helix: a request " - "sliced to this rank's tokens would under-size the global ledger." + "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: diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index a582900ec098..cbfaa0dad583 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/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, diff --git a/tensorrt_llm/_torch/pyexecutor/request_utils.py b/tensorrt_llm/_torch/pyexecutor/request_utils.py index 3193ea475c47..ee87534feb96 100644 --- a/tensorrt_llm/_torch/pyexecutor/request_utils.py +++ b/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}.") diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py index 6225ed81abef..7ea67a69526a 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py @@ -369,12 +369,11 @@ def test_scheduler_allocation_failure_raises_under_helix() -> None: ) -def _context_resize_mgr(cp_rank: int, cp_size: int, resizes: list) -> SimpleNamespace: +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, - _helix_cp_rank=cp_rank, - _helix_cp_size=cp_size, + is_disagg=is_disagg, kv_cache_map={7: kv}, num_extra_kv_tokens=0, _connector_reservations_enabled=lambda: False, @@ -395,25 +394,28 @@ def _context_req(prompt_len: int, total_input_len_cp: int) -> SimpleNamespace: ) -def test_resize_context_allows_full_prompt_helix_request() -> None: +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 and is rank-invariant.""" - cp_size = 4 + context target is already in global ledger tokens.""" resizes = [] - for r in range(cp_size): - req = _context_req(prompt_len=200, total_input_len_cp=200) - assert KVCacheManagerV2.resize_context( - _context_resize_mgr(r, cp_size, resizes), req, num_tokens=32 - ) - assert resizes == [64 + 32] * cp_size + 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] -def test_resize_context_rejects_sliced_helix_request() -> None: +@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.""" - req = _context_req(prompt_len=50, total_input_len_cp=200) - with pytest.raises(ValueError, match="needs the full prompt"): - KVCacheManagerV2.resize_context(_context_resize_mgr(0, 4, []), req, num_tokens=32) + 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: diff --git a/tests/unittest/_torch/executor/test_request_utils.py b/tests/unittest/_torch/executor/test_request_utils.py index e04064c1aa9b..3aaa454b4e16 100644 --- a/tests/unittest/_torch/executor/test_request_utils.py +++ b/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 From 72315c728922d403ece4a5883b8272b25d9c9b87 Mon Sep 17 00:00:00 2001 From: Asfiya Baig Date: Wed, 21 Oct 2026 18:36:37 +0000 Subject: [PATCH 5/7] [TRTLLM-17126][test] Cover the agg Helix KV fraction override Signed-off-by: Asfiya Baig --- .../test_kv_cache_manager_v2_helix_superblock.py | 13 +++++++++++-- 1 file changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py index 7ea67a69526a..f9bb9e6c20ce 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_manager_v2_helix_superblock.py +++ b/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"): From 9f0997091e5fa4f924991bd6eb1a7a7ed0352136 Mon Sep 17 00:00:00 2001 From: Asfiya Baig Date: Wed, 21 Oct 2026 23:32:27 +0000 Subject: [PATCH 6/7] [TRTLLM-17124][fix] Require KVCacheManagerV2 for aggregated Helix Aggregated Helix sizes prefill KV in global ledger blocks, which only KVCacheManagerV2 keeps. Reject any other resolved manager, including a V2 manager that fell back to V1, when the KvCacheCreator is built. Disaggregated Helix keeps V1. Add the Helix fields that the startup gate and the request routing now read to the test fakes that stand in for Mapping and PyExecutor. Signed-off-by: Asfiya Baig --- tensorrt_llm/_torch/pyexecutor/_util.py | 9 +++++++++ tests/unittest/_torch/executor/test_py_executor.py | 2 ++ .../test_py_executor_creator_mla_cache_reuse_sync.py | 1 + 3 files changed, 12 insertions(+) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 30d3bcdfda6f..aacf78479f45 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -953,6 +953,15 @@ 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 + and 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.") self._disable_overlap_scheduler = llm_args.disable_overlap_scheduler self._draft_config = draft_config self._skip_est = skip_est diff --git a/tests/unittest/_torch/executor/test_py_executor.py b/tests/unittest/_torch/executor/test_py_executor.py index 29df16273a38..2a19c5b6c532 100644 --- a/tests/unittest/_torch/executor/test_py_executor.py +++ b/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={}, diff --git a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py b/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py index a15caeec3595..38d60f2b7c7b 100644 --- a/tests/unittest/_torch/executor/test_py_executor_creator_mla_cache_reuse_sync.py +++ b/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, ) From 958d48363368076d322caa707fd1d388c13ec1bd Mon Sep 17 00:00:00 2001 From: Asfiya Baig Date: Thu, 22 Oct 2026 19:10:03 +0000 Subject: [PATCH 7/7] [TRTLLM-17124][fix] Require chunked prefill for aggregated Helix The aggregated Helix context phase gathers cached KV across CP ranks only in the chunked-context path, so reject aggregated Helix without chunked prefill when the KvCacheCreator is built. Read the setting from the model engine, which reflects the MLA SM-version fallback. Signed-off-by: Asfiya Baig --- tensorrt_llm/_torch/pyexecutor/_util.py | 27 ++++++++++++++++--------- 1 file changed, 18 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index aacf78479f45..d48581678400 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -953,15 +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 - and 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.") + 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