diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index 80419abbf404..283c1deabdcc 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1027,6 +1027,13 @@ def try_prepare_estimation(self) -> bool: logger.info( "KV cache size estimation is not supported for context parallelism, disable it." ) + if self._is_kv_cache_manager_v2: + # Like the encoder-decoder case below: promote to _skip_est + # so build_managers runs configure_kv_cache_capacity(), which + # sets the explicit quota KVCacheManagerV2 requires at + # construction (V1 sizes itself from the fraction + # internally, so it stays on the local flag). + self._skip_est = True model_config = self._model_engine.model.model_config if model_config.attn_backend == "VANILLA": estimating_kv_cache = False @@ -1069,6 +1076,35 @@ def try_prepare_estimation(self) -> bool: self._kv_cache_config.max_tokens = max_tokens return estimating_kv_cache + def _configure_helix_kv_cache_capacity(self) -> None: + """Set the helix KV quota without profiling (not CP-aware). + + An explicit quota is honored as-is; otherwise fall back to V1's + fraction sizing of free memory. KVCacheManagerV2 min-syncs the + derived max_tokens across ranks. + """ + if (self._kv_cache_config.max_gpu_total_bytes or 0) > 0 or \ + self._kv_cache_config.max_tokens: + logger.info("Helix CP: skipping KV cache capacity profiling; using " + "the explicitly configured quota.") + return + fraction = self._kv_cache_config.free_gpu_memory_fraction + free_mem, _total = torch.cuda.mem_get_info() + cost = self._get_kv_size_per_token() + max_tokens = cost.tokens_for_budget(int(free_mem * fraction)) + if max_tokens <= 0: + raise ValueError( + "Helix CP: fraction-based KV sizing found no usable free " + "memory; set kv_cache_config.max_tokens or " + "max_gpu_total_bytes.") + logger.warning( + "Helix CP: capacity profiling is unsupported; sizing the KV " + f"cache as fraction {fraction} of free memory -> " + f"max_tokens={max_tokens} (rank-local). Set " + "kv_cache_config.max_tokens or max_gpu_total_bytes to " + "override.") + self._kv_cache_config.max_tokens = max_tokens + def configure_kv_cache_capacity(self, py_executor: PyExecutor = None) -> None: """Perform KV cache capacity estimation. @@ -1079,6 +1115,9 @@ def configure_kv_cache_capacity(self, mapping = self._mapping # TODO: support CP by generating dummy requests for it. + if mapping.cp_config.get('cp_type') == CpType.HELIX: + self._configure_helix_kv_cache_capacity() + return assert 'cp_type' not in mapping.cp_config fraction = self._kv_cache_config.free_gpu_memory_fraction diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 02286c2707a1..2b80f203bd18 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -783,6 +783,29 @@ def __init__( "Star attention is not supported for KVCacheManagerV2" ) + # ---- Helix (decode context parallelism) bookkeeping -------------- + # V1 rotates decode-KV ownership across CP ranks in + # prepare_resources; V2 grows KV at scheduling time, so the gate + # lives in try_allocate_generation / revert_allocate_generation. The + # backend stays helix-ignorant: it only sees rank-local token counts. + self._has_cp_helix = mapping.has_cp_helix() + if self._has_cp_helix: + self._helix_cp_rank = mapping.cp_rank + self._helix_cp_size = mapping.cp_size + if kv_cache_config.enable_block_reuse: + raise ValueError( + "Helix CP requires enable_block_reuse=False: each rank " + "holds an interleaved slice of the sequence, so radix " + "prefix matching over the local token stream is " + "meaningless." + ) + if is_draft: + raise ValueError( + "Helix CP does not support speculative decoding with " + "KVCacheManagerV2 yet (draft-manager mirroring and " + "rewind are not round-robin aware)." + ) + self.kv_cache_type = kv_cache_type self.pp_layers, self.num_layers = get_pp_layers( num_layers, @@ -2227,9 +2250,32 @@ def try_allocate_generation(self, req: LlmRequest) -> bool: return False self._restore_page_index_bufs(req.py_request_id, kv_cache) + if self._has_cp_helix and not req.is_dummy_request: + # Round-robin gate (V1 parity). The V2 scheduler runs before + # the disagg handler seeds py_decoding_iter = 1, so the first + # schedule reads 0 and the negative modulo would silently pick + # rank cp_size - 1; clamping to 1 reproduces V1's owner sequence. + decode_iter = max(1, req.py_decoding_iter) + decode_block_id = (decode_iter - 1) // self.tokens_per_block + if decode_block_id % self._helix_cp_size != self._helix_cp_rank: + # Inactive rank this step: materialize nothing but report + # success — False would trigger per-rank eviction hunting + # and desynchronize the CP group. + req.py_helix_is_inactive_rank = True + return True + req.py_helix_is_inactive_rank = False + req.seqlen_this_rank_cp += 1 + draft_len = self._effective_draft_len(req) self._allocated_draft_lens[req.py_request_id] = draft_len - return kv_cache.resize(self._required_gen_capacity(req, kv_cache.capacity)) + if not kv_cache.resize(self._required_gen_capacity(req, kv_cache.capacity)): + if self._has_cp_helix and not req.is_dummy_request: + # Undo the bookkeeping above: the scheduler may retry this + # request in the same pass (after evicting a victim) and + # try_allocate_generation would increment again. + req.seqlen_this_rank_cp -= 1 + return False + return True def revert_allocate_generation(self, req: LlmRequest) -> None: """Undo the capacity growth from try_allocate_generation. @@ -2247,6 +2293,10 @@ def revert_allocate_generation(self, req: LlmRequest) -> None: kv_cache = self.kv_cache_map.get(req.py_request_id) if kv_cache is None or not kv_cache.is_active: return + if self._has_cp_helix and not req.is_dummy_request and req.py_helix_is_inactive_rank: + # Inactive rank grew nothing this step; the default path would + # shrink capacity below the materialized history. + return draft_len = self._allocated_draft_lens.pop( req.py_request_id, self._effective_draft_len(req) ) @@ -2259,6 +2309,10 @@ def revert_allocate_generation(self, req: LlmRequest) -> None: f"{req.py_request_id} from {kv_cache.capacity} to " f"{reverted_cap}" ) + if self._has_cp_helix and not req.is_dummy_request: + # Symmetric to the += 1 in try_allocate_generation (active rank + # only; inactive ranks returned above). + req.seqlen_this_rank_cp -= 1 def revert_allocate_context(self, req: LlmRequest) -> None: """Undo the capacity growth from this iteration's context resize.""" @@ -3134,6 +3188,11 @@ def release_resources( # a non-zero number to skip illegal memory access issue in MLA kernel # during warmup. token_num = token_nums[i] if token_nums is not None else 1 + max_num_draft_tokens + if self._has_cp_helix: + # token_num >= 2 keeps the active rank's + # past_seen_token_num (= seqlen_this_rank_cp - 1) + # non-negative (V1 parity). + token_num = max(token_num, 2) # token_num - 1 is the past history length in generation. history_hint = max(0, token_num - 1) if is_gen else None encoder_output_len = encoder_output_lens[i] if encoder_output_lens is not None else None @@ -3205,6 +3264,21 @@ def release_resources( req.prompt_len = token_num - 1 req.py_prompt_len = req.prompt_len req.py_draft_tokens = [1] * max_num_draft_tokens + if self._has_cp_helix: + # Frozen helix fields (V1 parity): padding dummies are + # appended inside forward and never pass the scheduling + # gate, so these values must stay permanently valid — + # last CP rank active, shared synthetic global length. + if self._helix_cp_rank == self._helix_cp_size - 1: + req.py_helix_is_inactive_rank = False + req.prompt_len = token_num - 1 + else: + req.py_helix_is_inactive_rank = True + req.prompt_len = token_num + req.py_prompt_len = req.prompt_len + req.seqlen_this_rank_cp = req.prompt_len + req.total_input_len_cp = token_num * self._helix_cp_size - 1 + req.py_decoding_iter = 1 if prepare_resource: new_capacity = kv_cache.capacity + _kv_draft + 1 success = kv_cache.resize(new_capacity, history_length=history_hint) @@ -3633,14 +3707,32 @@ def update_resources( # it accumulates in the draft KV cache after every generation # step. Target managers do not allocate this reserve slack. rewind_len += max(self._kv_reserve_draft_tokens - runtime_draft_len, 0) - new_capacity = ( - None - if req.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) - else kv_cache.capacity - rewind_len - ) - history_length = ( - None if self.kv_compression_manages_history else req.max_beam_num_tokens - 1 - ) + if self._has_cp_helix and not req.is_dummy_request: + # max_beam_num_tokens is the GLOBAL length; on a rank + # holding ~1/cp of the tokens it would inflate capacity past + # the gate. Floor at the rank-local count; history stays + # untouched (its consumers are disabled under helix and it + # must never decrease per the backend resize contract). + new_capacity = ( + None + if req.state + in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) + else max( + kv_cache.capacity - rewind_len, + req.seqlen_this_rank_cp, + ) + ) + history_length = None + else: + new_capacity = ( + None + if req.state + in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT) + else kv_cache.capacity - rewind_len + ) + history_length = ( + None if self.kv_compression_manages_history else req.max_beam_num_tokens - 1 + ) success = kv_cache.resize(new_capacity, history_length) if not success: raise ValueError( diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 10a7dba55540..1fb83eed8da9 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py @@ -2201,6 +2201,11 @@ def __init__( # Derive ssm_state_shape and conv_state_shape from mamba params (same as MambaCacheManager) tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1 + if mapping.has_cp_helix(): + # Helix repurposes CP ranks as TP for non-attention layers: + # this rank's mamba state is a 1/(tp*cp) head slice; bare + # tp_size would allocate a full replica per CP rank. + tp_size = mapping.tp_size * mapping.cp_size d_inner = mamba_head_dim * mamba_num_heads conv_dim = d_inner + 2 * mamba_n_groups * mamba_d_state nheads = mamba_num_heads @@ -2936,6 +2941,11 @@ def __init__( if self.local_num_mamba_layers > 0: tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1 + if mapping.has_cp_helix(): + # Helix repurposes CP ranks as TP for non-attention layers: + # this rank's mamba state is a 1/(tp*cp) head slice; bare + # tp_size would allocate a full replica per CP rank. + tp_size = mapping.tp_size * mapping.cp_size d_inner = mamba_head_dim * mamba_num_heads grouped_state_dim = mamba_n_groups * mamba_d_state conv_dim = d_inner + 2 * grouped_state_dim diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index 56f353f936c0..0bbd81dfb76d 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -961,6 +961,19 @@ def _try_schedule_generation( success = self.kv_cache_manager.try_allocate_generation(req) if not success: + if self.kv_cache_manager._has_cp_helix: + # Rank-local eviction decisions diverge across CP ranks + # (round-robin skew) and would desynchronize the group; fail + # loudly — capacity must be provisioned so scheduled + # requests always fit (GUARANTEED_NO_EVICT-like). + raise RuntimeError( + f"[V2Scheduler] KV allocation failed for helix request " + f"{req.py_request_id}; eviction is disabled under helix " + f"CP because rank-local eviction decisions diverge " + f"across the CP group. Increase " + f"kv_cache_config.max_gpu_total_bytes/max_tokens or " + f"reduce concurrency." + ) req_it_end, success = self._try_evict_for_gen( req, requests_list, req_it, req_it_end, evicted ) diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2_helix.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2_helix.py new file mode 100644 index 000000000000..65b597f291e4 --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2_helix.py @@ -0,0 +1,244 @@ +# 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. +"""Bookkeeping-only tests of the KVCacheManagerV2 helix round-robin gate. + +Binds the real KVCacheManagerV2 methods onto a stub object so the gate +arithmetic (try_allocate_generation / revert_allocate_generation) is +exercised without pools, GPUs or a model. +""" + +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 + +pytestmark = pytest.mark.cpu_only + +TPB = 4 # tokens_per_block: small so ownership rotates quickly +CP = 2 +INITIAL_CAPACITY = 10 + + +class _FakeKvCache: + def __init__(self, capacity): + self.capacity = capacity + self.is_active = True + self.history_length = 0 + self.resize_calls = [] + + def resume(self, _stream): + self.is_active = True + return True + + def resize(self, capacity, history_length=None): + self.resize_calls.append((capacity, history_length)) + if capacity is not None: + self.capacity = capacity + return True + + +def _make_mgr(cp_rank): + mgr = SimpleNamespace() + mgr._has_cp_helix = True + mgr._helix_cp_rank = cp_rank + mgr._helix_cp_size = CP + mgr.tokens_per_block = TPB + mgr.kv_cache_map = {} + mgr._allocated_draft_lens = {} + mgr._stream = SimpleNamespace(cuda_stream=None) + mgr._effective_draft_len = lambda req: 0 + mgr._required_gen_capacity = lambda req, cap: cap + 1 + mgr._restore_page_index_bufs = lambda rid, kv: None + return mgr + + +def _make_req(rid, iter0, seqlen): + return SimpleNamespace( + py_request_id=rid, + py_decoding_iter=iter0, + seqlen_this_rank_cp=seqlen, + py_helix_is_inactive_rank=False, + is_dummy_request=False, + ) + + +def _make_rank(cp_rank): + mgr = _make_mgr(cp_rank) + kv = _FakeKvCache(capacity=INITIAL_CAPACITY) + mgr.kv_cache_map[1] = kv + req = _make_req(1, 1, INITIAL_CAPACITY) + return mgr, kv, req + + +def _run_rank(cp_rank, steps): + """Simulate `steps` decode iterations on one rank; return per-step record.""" + mgr, kv, req = _make_rank(cp_rank) + trace = [] + for it in range(1, steps + 1): + req.py_decoding_iter = it + ok = KVCacheManagerV2.try_allocate_generation(mgr, req) + trace.append((it, ok, req.py_helix_is_inactive_rank, req.seqlen_this_rank_cp, kv.capacity)) + return trace, mgr, kv, req + + +@pytest.mark.parametrize("rank", range(CP)) +def test_ownership_rotation(rank): + """Decode block b is owned by rank b % CP; other ranks stay schedulable.""" + trace, _, kv, _ = _run_rank(rank, steps=TPB * 2) + for it, ok, inactive, _seq, _cap in trace: + owner = ((it - 1) // TPB) % CP + assert ok, f"iter {it} must stay schedulable on every rank" + assert inactive == (owner != rank), f"iter {it}: owner={owner} inactive={inactive}" + active_steps = sum(1 for _, _, inactive, _, _ in trace if not inactive) + assert kv.capacity == INITIAL_CAPACITY + active_steps + + +def test_seqlen_advances_only_on_active_steps(): + _, _, _, req = _run_rank(0, steps=TPB) + assert req.seqlen_this_rank_cp == INITIAL_CAPACITY + TPB + _, _, _, req = _run_rank(1, steps=TPB) + assert req.seqlen_this_rank_cp == INITIAL_CAPACITY + + +def test_first_schedule_before_decoding_iter_seeded(): + """First schedule reads py_decoding_iter == 0 (before the disagg + transmission-complete handler seeds it to 1) and must be treated as + decode step 1: owner = block 0 = rank 0, never rank cp_size - 1 via + Python's negative modulo.""" + mgr, kv, req = _make_rank(0) + req.py_decoding_iter = 0 + ok = KVCacheManagerV2.try_allocate_generation(mgr, req) + assert ok and not req.py_helix_is_inactive_rank + assert kv.capacity == INITIAL_CAPACITY + 1 + + mgr, kv, req = _make_rank(1) + req.py_decoding_iter = 0 + ok = KVCacheManagerV2.try_allocate_generation(mgr, req) + assert ok and req.py_helix_is_inactive_rank + assert kv.capacity == INITIAL_CAPACITY + + +def test_revert_symmetry_active_rank(): + mgr, kv, req = _make_rank(0) # rank 0 owns block 0 -> active + KVCacheManagerV2.try_allocate_generation(mgr, req) + KVCacheManagerV2.revert_allocate_generation(mgr, req) + assert kv.capacity == INITIAL_CAPACITY + assert req.seqlen_this_rank_cp == INITIAL_CAPACITY + + +def test_revert_is_noop_on_inactive_rank(): + mgr, kv, req = _make_rank(1) # rank 1 inactive during block 0 + KVCacheManagerV2.try_allocate_generation(mgr, req) + n_resizes = len(kv.resize_calls) + KVCacheManagerV2.revert_allocate_generation(mgr, req) + assert len(kv.resize_calls) == n_resizes + assert kv.capacity == INITIAL_CAPACITY + assert req.seqlen_this_rank_cp == INITIAL_CAPACITY + + +def test_helix_quota_fallback_uses_fraction_sizing(monkeypatch): + """Helix skips profiling (_configure_helix_kv_cache_capacity): an + explicit quota returns early untouched; without one it falls back to + V1-style fraction sizing of free memory (min-synced later by + KVCacheManagerV2).""" + import torch + + from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator + from tensorrt_llm.mapping import CpType + + def creator(max_gpu_total_bytes, max_tokens): + return SimpleNamespace( + _mapping=SimpleNamespace(cp_config={"cp_type": CpType.HELIX}), + _kv_cache_config=SimpleNamespace( + max_gpu_total_bytes=max_gpu_total_bytes, + max_tokens=max_tokens, + free_gpu_memory_fraction=0.5, + ), + _get_kv_size_per_token=lambda: CacheCost(slope=1000, intercept=8000), + ) + + monkeypatch.setattr(torch.cuda, "mem_get_info", lambda: (1_000_000, 2_000_000)) + + # Explicit quota: early return, config untouched. + c = creator(1 << 30, None) + assert KvCacheCreator._configure_helix_kv_cache_capacity(c) is None + assert c._kv_cache_config.max_tokens is None + # No quota: fraction fallback -> (1e6 * 0.5 - 8000) // 1000 = 492. + c = creator(0, None) + assert KvCacheCreator._configure_helix_kv_cache_capacity(c) is None + assert c._kv_cache_config.max_tokens == 492 + # 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"): + KvCacheCreator._configure_helix_kv_cache_capacity(creator(0, None)) + + +def test_helix_estimation_prepare_promotes_skip_est_for_v2(): + """Helix disables estimation; with a V2 manager it must also promote + _skip_est so build_managers() calls configure_kv_cache_capacity() — + otherwise KVCacheManagerV2 asserts "Quota not set" at construction. + V1 keeps sizing itself from the fraction, so it stays unpromoted.""" + from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator + from tensorrt_llm.mapping import CpType + + def creator(is_v2): + return SimpleNamespace( + _skip_est=False, + _mapping=SimpleNamespace(cp_config={"cp_type": CpType.HELIX}), + _is_kv_cache_manager_v2=is_v2, + _model_engine=SimpleNamespace( + model=SimpleNamespace( + model_config=SimpleNamespace(attn_backend="TRTLLM", is_encoder_decoder=False) + ) + ), + ) + + c = creator(is_v2=True) + assert KvCacheCreator.try_prepare_estimation(c) is False + assert c._skip_est is True + c = creator(is_v2=False) + assert KvCacheCreator.try_prepare_estimation(c) is False + assert c._skip_est is False + + +def test_scheduler_allocation_failure_raises_under_helix(): + """V2 scheduler: a failed generation allocation must raise under helix + instead of falling into rank-local eviction (which would desynchronize + the CP group).""" + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import KVCacheV2Scheduler + + sched = SimpleNamespace( + kv_cache_manager=SimpleNamespace( + _has_cp_helix=True, try_allocate_generation=lambda req: False + ), + ) + req = SimpleNamespace( + py_request_id=7, + get_beam_width_by_iter=lambda for_next_iteration: 1, + py_draft_tokens=None, + ) + budget = SimpleNamespace(can_fit_tokens=lambda n: True) + with pytest.raises(RuntimeError, match="eviction is disabled under helix"): + KVCacheV2Scheduler._try_schedule_generation( + sched, + req, + budget, + requests_list=[req], + req_it=0, + req_it_end=1, + evicted=[], + scheduled_beam_width=0, + )