diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 9ad518e3db5f..d48581678400 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 @@ -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: 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..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,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: 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/py_executor_creator.py b/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py index 7a95ac8c2682..4c3daa85debf 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py +++ b/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: 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 9d3670966646..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"): @@ -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 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 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, ) 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