diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index e473d3de352e..c935ce394890 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -1029,6 +1029,14 @@ 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 + and self._mapping.cp_config.get('cp_type') == CpType.HELIX): + # Promote like the encoder-decoder case so build_managers + # runs configure_kv_cache_capacity(), which sets the quota + # KVCacheManagerV2 requires at construction (V1 stays local). + # HELIX only: configure_kv_cache_capacity has no sizing path + # for other CP types and would hit its assertion. + self._skip_est = True model_config = self._model_engine.model.model_config if model_config.attn_backend == "VANILLA": estimating_kv_cache = False @@ -1071,6 +1079,41 @@ 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). + + Explicit quotas pass through; otherwise fraction sizing sets + ``max_gpu_total_bytes`` (a rank-local byte cap the manager consumes + as-is). Setting ``max_tokens`` here would overshoot the fraction: + the manager inflates that knob by 1 / max_util_for_resume. + """ + if (self._kv_cache_config.max_tokens is not None + and self._kv_cache_config.max_tokens <= 0): + raise ValueError( + "Helix CP: kv_cache_config.max_tokens must be positive when " + f"set, got {self._kv_cache_config.max_tokens}.") + if (self._kv_cache_config.max_gpu_total_bytes or 0) > 0 or \ + (self._kv_cache_config.max_tokens or 0) > 0: + 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() + budget_bytes = int(free_mem * fraction) + if budget_bytes <= 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_gpu_total_bytes={budget_bytes} (rank-local byte cap; the " + "manager min-syncs across ranks and converts to global tokens). " + "Set kv_cache_config.max_tokens or max_gpu_total_bytes to " + "override.") + self._kv_cache_config.max_gpu_total_bytes = budget_bytes + def configure_kv_cache_capacity(self, py_executor: PyExecutor = None) -> None: """Perform KV cache capacity estimation. @@ -1081,6 +1124,16 @@ 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: + if not self._is_kv_cache_manager_v2: + # The helix sizing below emits V2 ledger (global) quotas; + # V1 reads max_tokens as rank-local. Reject explicitly. + raise NotImplementedError( + "TRTLLM_SKIP_KV_CACHE_ESTIMATION with helix CP requires " + "the V2 KV cache manager " + "(kv_cache_config.use_kv_cache_manager_v2=True).") + 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 79731f60ce16..a6202947091f 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py @@ -842,6 +842,34 @@ def __init__( self.num_kv_heads = num_kv_heads self.head_dim = head_dim self.tokens_per_block = tokens_per_block + # Helix super-block ledger: bookkeeping runs in GLOBAL tokens, one + # ledger block = cp_size physical pages (one per rank), while + # self.tokens_per_block stays the PHYSICAL kernel page size. + self._has_cp_helix = mapping.has_cp_helix() + self._helix_cp_rank = mapping.cp_rank if self._has_cp_helix else 0 + self._helix_cp_size = mapping.cp_size if self._has_cp_helix else 1 + self._ledger_tokens_per_block = tokens_per_block * self._helix_cp_size + if self._has_cp_helix: + if kv_cache_config.enable_block_reuse: + raise ValueError( + "KVCacheManagerV2 does not support block reuse with " + "helix context parallelism: a ledger block's bytes span " + "all CP ranks, so publishing it requires a group-wide " + "protocol. Set kv_cache_config.enable_block_reuse=False." + ) + if is_draft: + raise ValueError( + "KVCacheManagerV2 does not support a draft cache " + "manager with helix context parallelism." + ) + if mapping.enable_attention_dp: + raise ValueError( + "KVCacheManagerV2 does not support attention-DP with " + "helix context parallelism: disagg transfer-completion " + "consensus is skipped under attention-DP, so the " + "scheduler's request view would not be rank-invariant " + "across the CP group." + ) self.max_seq_len = max_seq_len self.max_batch_size = max_batch_size self.max_num_tokens = max_num_tokens @@ -1179,7 +1207,9 @@ def append_to_kv_heads_per_layer( max_seq_capacity = ( self.max_seq_len + self.num_extra_kv_tokens + self._kv_reserve_draft_tokens + 1 ) - self.max_blocks_per_seq = (max_seq_capacity + tokens_per_block - 1) // tokens_per_block + self.max_blocks_per_seq = ( + max_seq_capacity + self._ledger_tokens_per_block - 1 + ) // self._ledger_tokens_per_block if self.max_blocks_per_seq % 4 != 0: self.max_blocks_per_seq = ((self.max_blocks_per_seq + 3) // 4) * 4 @@ -1453,6 +1483,17 @@ def _get_runtime_cache_size_layer_components(self) -> tuple[List[int], List[Opti return layer_sizes, attention_windows def _get_max_tokens_from_quota(self, quota: int) -> float: + """Rank-local byte quota -> token capacity (GLOBAL tokens under helix).""" + tokens = self._get_max_tokens_from_quota_impl(quota) + if self._has_cp_helix and not math.isinf(tokens): + # Floor to whole physical pages before scaling: a ledger block + # allocates one full page on every CP rank, so a partial + # trailing page in the rank-local budget is never usable. + tokens = int(tokens) // self.tokens_per_block * self.tokens_per_block + tokens *= self._helix_cp_size + return tokens + + def _get_max_tokens_from_quota_impl(self, quota: int) -> float: layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() full_attn_size_per_token = _estimate_full_attn_size_per_token( layer_sizes, attention_windows @@ -1486,6 +1527,16 @@ def _get_max_tokens_from_quota(self, quota: int) -> float: return self.max_num_tokens + (quota - context_limit_quota) / generation_size_per_token def _get_quota_from_max_tokens(self, max_tokens: int) -> int: + """Token capacity (GLOBAL tokens under helix) -> rank-local byte quota.""" + if self._has_cp_helix: + # Round up to whole ledger blocks first: allocation is page- + # granular on every rank, so a request of N global tokens costs + # ceil(N / ledger_tpb) full physical pages per rank. + blocks = -(-int(max_tokens) // self._ledger_tokens_per_block) + max_tokens = blocks * self.tokens_per_block + return self._get_quota_from_max_tokens_impl(max_tokens) + + def _get_quota_from_max_tokens_impl(self, max_tokens: int) -> int: layer_sizes, attention_windows = self._get_runtime_cache_size_layer_components() full_attn_size_per_token = _estimate_full_attn_size_per_token( layer_sizes, attention_windows @@ -1908,7 +1959,9 @@ def _build_base_config( ) return KVCacheManagerConfigPy( - tokens_per_block=tokens_per_block, + # Used by the backend only for token<->block arithmetic and + # radix hashing; BufferConfig.size above stays the physical page. + tokens_per_block=self._ledger_tokens_per_block, cache_tiers=cache_tiers, layers=layer_configs, typical_step=typical_step, @@ -2157,6 +2210,17 @@ def get_index_k_buffer( def get_num_available_tokens( self, *, token_num_upper_bound: int, batch_size: int = 1, max_num_draft_tokens: int = 0 ) -> int: + """Clamp ``token_num_upper_bound`` to the allocatable token capacity. + + Unit note: under helix the backend runs on the ledger + ``tokens_per_block``, so the returned capacity (like + ``token_num_upper_bound`` and ``max_seq_len``) is in GLOBAL ledger + tokens - the coordinate request lengths are expressed in. Callers + that additionally bound the result by per-forward budgets (e.g. + ``max_num_tokens``) stay consistent because a helix context forward + replicates all tokens on every rank, so both bounds constrain the + same request-length variable. + """ extra_tokens = self.num_extra_kv_tokens + max_num_draft_tokens # Token num upper bound is the maximum number of tokens that can be allocated in the kv cache manager. # We need to add extra tokens to the token num upper bound to account for the extra tokens. @@ -2231,6 +2295,33 @@ def _effective_draft_len(self, req: LlmRequest) -> int: draft_len = self.max_total_draft_tokens return draft_len + def _helix_local_len(self, global_len: int) -> int: + """Tokens of the first ``global_len`` owned by this CP rank + (continuation round-robin: page b lives on rank b %% cp).""" + phys = self.tokens_per_block + full, rem = divmod(global_len, self._ledger_tokens_per_block) + return full * phys + min(max(rem - self._helix_cp_rank * phys, 0), phys) + + def _set_helix_rank_fields(self, req: LlmRequest) -> None: + """Derive the per-rank helix fields from the global position. + + The decode-step index is manager-owned (committed in + ``try_allocate_generation`` on successful resize) rather than derived + from ``py_decoding_iter``: the sampler advances that counter after + scheduling under the overlap loop, so a schedule-time read is one + step behind and would repeat the first decode position, overwriting + the first generated token's KV. Assumes one new token per step + (draft-token modes are rejected under helix). + """ + step = req.py_helix_decode_group_index + 1 + pos = req.total_input_len_cp + step - 1 + owner = (pos // self.tokens_per_block) % self._helix_cp_size + active = owner == self._helix_cp_rank + req.py_helix_is_inactive_rank = not active + # Convention shared with model_engine: the active rank's seqlen + # includes the in-flight token (past_seen = seqlen - 1 there). + req.seqlen_this_rank_cp = self._helix_local_len(pos) + (1 if active else 0) + def _required_gen_capacity(self, req: LlmRequest, current_capacity: int) -> int: """Compute generation KV cache capacity for a request. @@ -2255,7 +2346,16 @@ def try_allocate_generation(self, req: LlmRequest) -> bool: 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)) + is_helix_req = self._has_cp_helix and not req.is_dummy_request + if is_helix_req: + self._set_helix_rank_fields(req) + if not kv_cache.resize(self._required_gen_capacity(req, kv_cache.capacity)): + return False + if is_helix_req: + # Commit only on success so a same-pass retry recomputes the + # same step instead of skipping one. + req.py_helix_decode_group_index += 1 + return True def revert_allocate_generation(self, req: LlmRequest) -> None: """Undo the capacity growth from try_allocate_generation. @@ -2273,6 +2373,9 @@ 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_decode_group_index > 0: + # The forward pass for this step is skipped; give the step back. + req.py_helix_decode_group_index -= 1 draft_len = self._allocated_draft_lens.pop( req.py_request_id, self._effective_draft_len(req) ) @@ -2374,7 +2477,10 @@ def _prepare_context_impl(self, req: LlmRequest) -> bool: tokens, cache_salt=req.cache_salt, is_dummy=req.is_dummy, - expected_prompt_length=req.prompt_len - 1, + expected_prompt_length=( + req.total_input_len_cp if self._has_cp_helix else req.prompt_len + ) + - 1, ) if kv_cache is None: return False @@ -2412,6 +2518,13 @@ 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: + 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." + ) kv_cache = self.kv_cache_map.get(req.py_request_id) if kv_cache is None: return False @@ -2445,11 +2558,14 @@ def prepare_disagg_gen_init(self, req: LlmRequest) -> bool: # prompt_len is the full incoming prompt length, robust to block # reuse (which may leave a non-zero context_current_position). - target = req.prompt_len + get_draft_token_length(req) + self.num_extra_kv_tokens + # Helix requests carry the rank-local strided slice in prompt_len; + # the global ledger sizes off the full prompt instead. + prompt_len = req.total_input_len_cp if self._has_cp_helix else req.prompt_len + target = prompt_len + get_draft_token_length(req) + self.num_extra_kv_tokens capacity = max(kv_cache.capacity, target) pre_cap = kv_cache.capacity - success = kv_cache.resize(capacity, req.prompt_len) + success = kv_cache.resize(capacity, prompt_len) if not success: if req.is_first_context_chunk: kv_cache.suspend() @@ -3166,6 +3282,10 @@ 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: + # Keep the frozen dummy fields self-consistent (the active + # rank's past_seen = seqlen - 1 must stay >= 1). + 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 and not materialize_history else None encoder_output_len = encoder_output_lens[i] if encoder_output_lens is not None else None @@ -3237,6 +3357,20 @@ 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 fields (V1 parity): dummies never pass the + # per-step derivation — 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) @@ -3671,7 +3805,11 @@ def update_resources( else kv_cache.capacity - rewind_len ) history_length = ( - None if self.kv_compression_manages_history else req.max_beam_num_tokens - 1 + None + # Reuse (history's consumer) is disabled under helix, and + # max_beam_num_tokens mixes rank-local and global counts. + if self.kv_compression_manages_history or self._has_cp_helix + else req.max_beam_num_tokens - 1 ) success = kv_cache.resize(new_capacity, history_length) if not success: diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py index 7b7d02951e8a..041c4933542a 100644 --- a/tensorrt_llm/_torch/pyexecutor/llm_request.py +++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py @@ -971,6 +971,9 @@ def __init__( self.py_end_id = self.end_id self.py_min_length = self.sampling_config.min_length self.py_helix_is_inactive_rank = False + # Manager-owned helix decode-step counter; see + # KVCacheManagerV2._set_helix_rank_fields. + self.py_helix_decode_group_index = 0 self.py_draft_logits = None self.py_target_probs = None self.py_per_pos_drafted = [0] * MAX_SPEC_DECODE_POSITIONS diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py index a4a298d16d94..cf12b4b91a7e 100644 --- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py +++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py @@ -193,6 +193,7 @@ def __init__( self.chunk_unit_size = 0 self.max_context_length = max_num_tokens self.tokens_per_block = kv_cache_manager.tokens_per_block + self.has_cp_helix = kv_cache_manager._has_cp_helix draft_mgr_name = ( type(draft_kv_cache_manager).__name__ if draft_kv_cache_manager is not None else "None" ) @@ -993,6 +994,18 @@ def _try_schedule_generation( success = self.kv_cache_manager.try_allocate_generation(req) if not success: + if self.has_cp_helix: + # No-evict stance: every validated helix run used + # GUARANTEED_NO_EVICT semantics; eviction stays off under + # helix until a KV-pressure e2e validates it. + raise RuntimeError( + f"[V2Scheduler] KV allocation failed for helix request " + f"{req.py_request_id}; eviction is disabled under helix " + f"CP pending end-to-end validation. Increase " + f"kv_cache_config.max_gpu_total_bytes (rank-local bytes) " + f"or max_tokens (GLOBAL tokens across CP ranks, not " + f"per-rank), or reduce concurrency." + ) req_it_end, success = self._try_evict_for_gen( req, requests_list, req_it, req_it_end, evicted, inflight_request_ids ) diff --git a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py index 673c7f666c22..059cb3b79738 100644 --- a/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py +++ b/tests/unittest/_torch/executor/test_dual_pool_kv_cache.py @@ -713,6 +713,7 @@ class TestKVCacheV2SchedulerCrossParam: def _make_mock_kv_mgr(self, tokens_per_block=64): mgr = Mock(spec=KVCacheManagerV2) mgr.tokens_per_block = tokens_per_block + mgr._has_cp_helix = False return mgr def test_default_cross_is_none(self): diff --git a/tests/unittest/_torch/executor/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/test_kv_cache_estimation.py index 03218c125e6d..9f0a4b9ac0df 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/test_kv_cache_estimation.py @@ -590,6 +590,7 @@ def get_cache_size_per_token(model_config, mapping, **kwargs): def test_v2_quota_from_max_tokens_models_context_swa_scratch(): manager = object.__new__(KVCacheManagerV2) + manager._has_cp_helix = False manager.num_local_layers = 3 manager.pp_layers = [0, 1, 2] manager.max_attention_window_vec = [128, 128, None] diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py index ddbcddfc2ee1..81f3c3c3a9d8 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py @@ -94,6 +94,9 @@ def _make_cache_config_for_test( cache_manager.max_num_tokens = max_num_tokens cache_manager.max_draft_len = max_draft_len cache_manager.get_layer_bytes_per_token = lambda **_: 128 + # Mirrors __init__: without helix the ledger block equals the physical + # page (the helper re-enacts construction for partial instances). + cache_manager._ledger_tokens_per_block = 128 return cache_manager._build_base_config( kv_cache_config, diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2_helix_superblock.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2_helix_superblock.py new file mode 100644 index 000000000000..e95682373fcd --- /dev/null +++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2_helix_superblock.py @@ -0,0 +1,366 @@ +# 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 the helix (decode-CP) super-block ledger in KVCacheManagerV2. + +Design under test: the request ledger runs in GLOBAL tokens with one ledger +block spanning ``cp_size`` physical pages (one per CP rank); the per-rank +view (owner of this step's token, tokens held by this rank) is a closed-form +function of the global position. The decode-step index feeding that position +is manager-owned (``py_helix_decode_group_index``), immune to when the +sampler advances ``py_decoding_iter``. +""" + +import math +from types import SimpleNamespace + +import pytest + +from tensorrt_llm._torch.pyexecutor.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState + + +def _mgr(cp_rank: int, cp_size: int, phys: int) -> SimpleNamespace: + m = SimpleNamespace( + tokens_per_block=phys, + _ledger_tokens_per_block=phys * cp_size, + _has_cp_helix=cp_size > 1, + _helix_cp_rank=cp_rank, + _helix_cp_size=cp_size, + ) + m._helix_local_len = lambda global_len: KVCacheManagerV2._helix_local_len(m, global_len) + return m + + +def _brute_local_len(global_len: int, cp_rank: int, cp_size: int, phys: int) -> int: + """Reference: token position p lives on rank (p // phys) % cp_size.""" + return sum(1 for p in range(global_len) if (p // phys) % cp_size == cp_rank) + + +def test_helix_local_len_matches_brute_force() -> None: + for cp_size in (1, 2, 4, 8): + for phys in (2, 4, 32): + for cp_rank in range(cp_size): + m = _mgr(cp_rank, cp_size, phys) + for global_len in range(0, 4 * cp_size * phys + 3): + assert KVCacheManagerV2._helix_local_len(m, global_len) == _brute_local_len( + global_len, cp_rank, cp_size, phys + ), (cp_size, phys, cp_rank, global_len) + + +def test_set_helix_rank_fields_cross_rank_consistency() -> None: + """For any (prompt_len, decode step): exactly one active rank, the + per-rank seqlens sum to the global in-flight length, and past_seen + (= seqlen - 0/1 per the model_engine convention) sums to the global + already-cached length.""" + cp_size, phys = 4, 4 + for total_input in (1, 5, 16, 17, 63): + for group_index in (0, 1, 2, 7, 40): # committed schedules so far + fields = [] + for r in range(cp_size): + req = SimpleNamespace( + total_input_len_cp=total_input, + py_helix_decode_group_index=group_index, + ) + KVCacheManagerV2._set_helix_rank_fields(_mgr(r, cp_size, phys), req) + fields.append(req) + pos = total_input + group_index + active = [r for r in range(cp_size) if not fields[r].py_helix_is_inactive_rank] + assert active == [(pos // phys) % cp_size] + assert sum(f.seqlen_this_rank_cp for f in fields) == pos + 1 + past_seen = [ + f.seqlen_this_rank_cp - (0 if f.py_helix_is_inactive_rank else 1) for f in fields + ] + assert sum(past_seen) == pos + # Prompt distribution matches the arrival striding convention. + for r in range(cp_size): + assert past_seen[r] >= 0 + + +def test_ledger_is_rank_invariant() -> None: + """The whole point of the super-block design: nothing the scheduler or + ledger consumes depends on cp_rank — only the derived per-rank fields + do. Verify the derivation never touches ledger quantities by checking + the same request object gives identical (owner-adjusted) views.""" + cp_size, phys = 8, 32 + req_proto = dict(total_input_len_cp=1000, py_helix_decode_group_index=17) + lens = [] + for r in range(cp_size): + req = SimpleNamespace(**req_proto) + KVCacheManagerV2._set_helix_rank_fields(_mgr(r, cp_size, phys), req) + lens.append(req.seqlen_this_rank_cp) + # Ledger-side numbers (global position, page count) are identical on + # every rank; per-rank lens differ by at most one physical page. + assert max(lens) - min(lens) <= phys + + +def test_ledger_position_immune_to_sampler_timing() -> None: + """Regression for the overlap-scheduler phase skew: the ledger position + must be L, L+1, L+2, ... no matter when the sampler advances + py_decoding_iter (the overlap loop updates it after scheduling; a stale + read would repeat the first position and overwrite that token's KV).""" + cp_size, phys, total_input = 4, 4, 17 + n_steps = 3 * cp_size * phys # cross several ownership rotations + + for sampler_timing in ("overlap", "non_overlap"): + req = SimpleNamespace( + total_input_len_cp=total_input, + py_decoding_iter=0, # disagg-gen seeding happens after scheduling + py_helix_decode_group_index=0, + ) + mgrs = [_mgr(r, cp_size, phys) for r in range(cp_size)] + positions = [] + for step in range(1, n_steps + 1): + if sampler_timing == "non_overlap": + req.py_decoding_iter = step # sampler already advanced + per_rank = [] + for r in range(cp_size): + view = SimpleNamespace(**vars(req)) + KVCacheManagerV2._set_helix_rank_fields(mgrs[r], view) + per_rank.append(view) + # Commit, mirroring try_allocate_generation's success path. + req.py_helix_decode_group_index += 1 + if sampler_timing == "overlap": + req.py_decoding_iter = step # sampler advances only now + + pos = total_input + step - 1 + active = [r for r in range(cp_size) if not per_rank[r].py_helix_is_inactive_rank] + assert active == [(pos // phys) % cp_size] + # In-flight convention: seqlens sum to pos + 1 (never repeats, + # so no write offset can collide with the previous step's). + assert sum(f.seqlen_this_rank_cp for f in per_rank) == pos + 1 + positions.append(pos) + assert positions == [total_input + n for n in range(n_steps)] + + +def test_decode_step_derivation_is_idempotent() -> None: + """A failed try_allocate (counter not committed) followed by a retry in + the same scheduling pass must derive the SAME fields; a revert (skipped + forward) must give the step back.""" + m = _mgr(cp_rank=1, cp_size=4, phys=4) + req = SimpleNamespace( + total_input_len_cp=10, + py_helix_decode_group_index=5, + ) + KVCacheManagerV2._set_helix_rank_fields(m, req) + first = (req.py_helix_is_inactive_rank, req.seqlen_this_rank_cp) + KVCacheManagerV2._set_helix_rank_fields(m, req) # retry, no commit + assert (req.py_helix_is_inactive_rank, req.seqlen_this_rank_cp) == first + # Commit then revert restores the same derivation. + req.py_helix_decode_group_index += 1 + req.py_helix_decode_group_index -= 1 + KVCacheManagerV2._set_helix_rank_fields(m, req) + assert (req.py_helix_is_inactive_rank, req.seqlen_this_rank_cp) == first + + +def test_quota_converters_scale_by_cp() -> None: + m = SimpleNamespace( + _has_cp_helix=True, + _helix_cp_size=4, + tokens_per_block=32, + _ledger_tokens_per_block=128, + _get_max_tokens_from_quota_impl=lambda quota: 100.0, + _get_quota_from_max_tokens_impl=lambda tokens: tokens * 7, + ) + # Rank-local byte quota buys 100 physical tokens, but only 96 (= 3 whole + # 32-token pages) are allocatable -> 384 global ledger tokens, not 400. + assert KVCacheManagerV2._get_max_tokens_from_quota(m, 12345) == 384.0 + # inf (all-SWA) passes through unscaled. + m_inf = SimpleNamespace( + _has_cp_helix=True, + _helix_cp_size=4, + tokens_per_block=32, + _ledger_tokens_per_block=128, + _get_max_tokens_from_quota_impl=lambda quota: float("inf"), + ) + assert math.isinf(KVCacheManagerV2._get_max_tokens_from_quota(m_inf, 1)) + # Global tokens -> whole ledger blocks (ceil) -> per-rank physical pages + # -> bytes: 401 global tokens need ceil(401/128) = 4 ledger blocks = + # 4 * 32 = 128 physical tokens per rank (not ceil(401/4) = 101). + assert KVCacheManagerV2._get_quota_from_max_tokens(m, 401) == 128 * 7 + # cp == 1 (non-helix) is the identity: no page rounding is applied. + m1 = SimpleNamespace( + _has_cp_helix=False, + _helix_cp_size=1, + tokens_per_block=32, + _ledger_tokens_per_block=32, + _get_max_tokens_from_quota_impl=lambda quota: 100.0, + _get_quota_from_max_tokens_impl=lambda tokens: tokens * 7, + ) + assert KVCacheManagerV2._get_max_tokens_from_quota(m1, 1) == 100.0 + assert KVCacheManagerV2._get_quota_from_max_tokens(m1, 400) == 2800 + + +def test_update_resources_leaves_history_untouched_under_helix() -> None: + resizes = [] + + kv = SimpleNamespace( + is_active=True, + capacity=100, + resize=lambda cap, hist: resizes.append((cap, hist)) or True, + ) + req = SimpleNamespace( + py_request_id=7, + py_rewind_len=0, + py_num_accepted_draft_tokens=0, + state=LlmRequestState.GENERATION_IN_PROGRESS, + max_beam_num_tokens=55, + ) + mgr = SimpleNamespace( + # is_draft=True skips the module-level draft-token relocation call; + # with zero reserve tokens the rewind math is unchanged. + is_draft=True, + _kv_reserve_draft_tokens=0, + kv_cache_map={7: kv}, + kv_compression_manages_history=False, + _has_cp_helix=True, + ) + batch = SimpleNamespace(generation_requests=[req]) + KVCacheManagerV2.update_resources(mgr, batch) + assert resizes == [(100, None)] + # Non-helix keeps the vanilla history commit. + mgr._has_cp_helix = False + resizes.clear() + KVCacheManagerV2.update_resources(mgr, batch) + assert resizes == [(100, 54)] + + +def test_helix_quota_fallback_sets_rank_local_bytes(monkeypatch: pytest.MonkeyPatch) -> None: + """Fraction sizing must set max_gpu_total_bytes (rank-local byte cap), + not max_tokens, which the manager inflates by 1/max_util_for_resume.""" + import torch + + from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator + + def creator(max_gpu_total_bytes, max_tokens): + return SimpleNamespace( + _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, + ), + ) + + 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 + # Explicit but non-positive max_tokens: rejected, not silently replaced + # by fraction sizing. + with pytest.raises(ValueError, match="must be positive"): + KvCacheCreator._configure_helix_kv_cache_capacity(creator(0, 0)) + # No quota: the rank-local byte budget (free * fraction) lands on + # max_gpu_total_bytes verbatim; max_tokens stays unset. + c = creator(0, None) + 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 + # 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_v1_helix_capacity_config_rejected() -> None: + """configure_kv_cache_capacity emits V2 super-block-ledger coordinates + under helix; V1 interprets max_tokens as rank-local. The V1 + helix + combination must be rejected loudly, not silently mis-sized.""" + from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator + from tensorrt_llm.mapping import CpType + + c = SimpleNamespace( + _mapping=SimpleNamespace(cp_config={"cp_type": CpType.HELIX}), + _is_kv_cache_manager_v2=False, + ) + with pytest.raises(NotImplementedError, match="V2 KV cache manager"): + KvCacheCreator.configure_kv_cache_capacity(c) + + +def test_estimation_prepare_promotes_skip_est_for_v2() -> None: + """Helix disables estimation; with a V2 manager it must also promote + _skip_est so build_managers() calls configure_kv_cache_capacity(). + Other CP types must NOT be promoted: configure_kv_cache_capacity has no + sizing path for them.""" + from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator + from tensorrt_llm.mapping import CpType + + def creator(is_v2, cp_type=CpType.HELIX): + return SimpleNamespace( + _skip_est=False, + _mapping=SimpleNamespace(cp_config={"cp_type": cp_type}), + _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 + c = creator(is_v2=True, cp_type=CpType.ULYSSES) + assert KvCacheCreator.try_prepare_estimation(c) is False + assert c._skip_est is False + + +def test_scheduler_allocation_failure_raises_under_helix() -> None: + """Precedent-consistent no-evict stance: allocation failure under helix + raises instead of entering the (unvalidated-under-helix) eviction path.""" + from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import ( + KVCacheV2Scheduler, + _RecomputePauseState, + ) + + sched = SimpleNamespace( + has_cp_helix=True, + kv_cache_manager=SimpleNamespace(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, + recompute_pause_state=_RecomputePauseState(1), + evicted=[], + recompute_paused=[], + inflight_request_ids=set(), + scheduled_beam_width=0, + ) + + +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 + the synthetic global length.""" + for cp_size in (2, 4, 8): + for token_num in (2, 5, 33): + seqlens = [token_num - 1 if r == cp_size - 1 else token_num for r in range(cp_size)] + total = token_num * cp_size - 1 + assert sum(seqlens) == total diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py index a67cafdc2b01..1d86b5d1129a 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_capacity_only.py @@ -24,6 +24,7 @@ def _manager( ) -> KVCacheManagerV2: manager = KVCacheManagerV2.__new__(KVCacheManagerV2) manager.is_draft = is_draft + manager._has_cp_helix = False manager.kv_compression_manages_history = kv_compression_manages_history manager._kv_reserve_draft_tokens = kv_reserve_draft_tokens manager.kv_cache_map = {} diff --git a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py index 934397f5ff67..e89bc6518405 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py +++ b/tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py @@ -169,6 +169,7 @@ def make_kv_cache_manager( mgr = Mock() mgr.tokens_per_block = tokens_per_block mgr.can_evict = can_evict + mgr._has_cp_helix = False mgr.kv_cache_map = _KVCacheMap() mgr.prepare_context.side_effect = prepare_context_fn or (lambda req: True) mgr.resize_context.side_effect = resize_context_fn or (lambda req, n: True) diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index 1d1feabd6114..4a986c22e5f4 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/tests/unittest/_torch/executor/test_mamba_cache_manager.py @@ -1759,6 +1759,7 @@ def test_v2_hybrid_warns_when_avg_seq_len_is_missing(monkeypatch): def test_v2_hybrid_rejects_quota_below_live_state_floor(): mgr = object.__new__(MambaHybridCacheManagerV2) + mgr._has_cp_helix = False mgr.max_batch_size = 2 mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) mgr.local_num_mamba_layers = 1 @@ -1788,6 +1789,7 @@ def test_v2_hybrid_rejects_quota_below_live_state_floor(): def test_v2_hybrid_pure_mamba_rank_does_not_reserve_attention_page(): mgr = object.__new__(MambaHybridCacheManagerV2) + mgr._has_cp_helix = False mgr.max_batch_size = 2 mgr.mapping = Mapping(world_size=1, rank=0, tp_size=1, pp_size=1) mgr.local_num_mamba_layers = 1 @@ -1986,6 +1988,7 @@ def test_expect_snapshot_points_binding_round_trip(): def test_v2_hybrid_pool_ratio_controls_allocated_memory(): def allocated_memory(pool_ratio): mgr = object.__new__(MambaHybridCacheManagerV2) + mgr._has_cp_helix = False mgr.kv_cache_type = CacheTypeCpp.SELF mgr.head_dim_per_layer = [64, 64] mgr.pp_layers = [0, 1]