From 247276f5444627e70d1e22ce3b075d5b6979b421 Mon Sep 17 00:00:00 2001 From: Liao Lanyu <108499334+lancelly@users.noreply.github.com> Date: Sun, 16 Aug 2026 20:45:37 -0700 Subject: [PATCH] [None][feat] KVCacheManagerV2: helix (decode-CP) support Add helix (decode context parallelism) round-robin bookkeeping to the V2 KV cache manager: a frontend gate with a rank-local contract, requiring zero changes to the cpp/Python V2 backend (each rank only declares its local tokens, so the backend sees an ordinary short sequence). - kv_cache_manager_v2.py: - __init__: helix fields (_has_cp_helix/_helix_cp_rank/_helix_cp_size) plus guards rejecting incompatible combinations (block reuse, draft manager). - try_allocate_generation: decode-block round-robin gate (max(1, py_decoding_iter) - 1) // tokens_per_block % cp == cp_rank. The max(1, .) clamp is essential: the V2 scheduler runs before the disagg transmission-complete handler seeds py_decoding_iter = 1, so a bare port would send the first step to rank cp_size - 1 via negative modulo. Inactive ranks return success with zero resource changes; active ranks do seqlen_this_rank_cp += 1 and resize(+1), rolled back on failure. - revert_allocate_generation: symmetric rollback on the active rank; a strict no-op on inactive ranks. - update_resources: capacity floors at the rank-local seqlen_this_rank_cp instead of the global sequence length, and history_length stays None under helix to preserve the backend's history monotonicity contract. - add_dummy_requests: token_num >= 2 and V1-parity helix fields frozen at creation (CUDA-graph padding dummies never pass through the scheduler, so the fields must remain permanently valid). - scheduler_v2.py: allocation failure raises under helix; eviction and preemption would need cross-CP-rank synchronization to keep round-robin ownership aligned, so they are disabled for now. - _util.py: helix skips capacity profiling (the estimator's dummy prefill is not CP-aware). An explicit quota (max_tokens / max_gpu_total_bytes) is honored as-is; otherwise fall back to V1-style fraction sizing of free memory, which KVCacheManagerV2 then min-syncs across ranks like V1's calculate_max_num_blocks. - mamba_cache_manager.py: hybrid state shards heads by tp_size * cp_size under helix (both manager variants), matching the repurposed CP-as-TP weight sharding of non-attention layers. - tests: add CPU-only unit tests for the round-robin gate bookkeeping. - try_prepare_estimation: under CP with a V2 manager, promote _skip_est (like the encoder-decoder case) so build_managers() actually calls configure_kv_cache_capacity(); the CP branch previously only cleared the local flag, leaving the helix quota logic unreachable and KVCacheManagerV2 asserting "Quota not set" at construction. V1 stays on the local flag (it sizes itself from the fraction internally). Locked by a new flow-contract unit test. Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com> --- tensorrt_llm/_torch/pyexecutor/_util.py | 39 +++ .../_torch/pyexecutor/kv_cache_manager_v2.py | 110 +++++++- .../_torch/pyexecutor/mamba_cache_manager.py | 10 + .../pyexecutor/scheduler/scheduler_v2.py | 13 + .../test_kv_cache_manager_v2_helix.py | 244 ++++++++++++++++++ 5 files changed, 407 insertions(+), 9 deletions(-) create mode 100644 tests/unittest/_torch/executor/test_kv_cache_manager_v2_helix.py 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, + )