From 613fe76ef0d670a586ce4b79e8ada5227fc6ced6 Mon Sep 17 00:00:00 2001
From: Liao Lanyu <108499334+lancelly@users.noreply.github.com>
Date: Mon, 17 Aug 2026 21:32:32 -0700
Subject: [PATCH 1/5] [None][feat] KVCacheManagerV2: helix via a global
 super-block ledger

Alternative design to #17795 for helix (decode context parallelism)
support in KVCacheManagerV2: instead of running each rank's ledger in
rank-local tokens with a rotation gate, run the request LEDGER in GLOBAL
tokens with one ledger block spanning cp_size physical pages (one per CP
rank).

Physical layout is byte-identical to the existing helix block rotation
(page b lives on rank b % cp, every page a full tokens_per_block run),
so kernels, page tables, and the disagg transfer striding are untouched.
Only the accounting changes:

- kv_cache_manager_v2.py: _ledger_tokens_per_block = cp * physical tpb.
  The backend config gets the ledger constant (it only does token<->block
  arithmetic and radix hashing with it) while BufferConfig.size stays the
  physical page bytes, so the entire V2 backend (Python and C++) is
  unchanged: 1 ledger block == 1 local physical page on every rank.
- Every rank's ledger advances identically (+1 global token per step):
  no rotation gate, no inactive-rank allocation special case, no revert
  asymmetry, no rank-local capacity floor, and scheduler eviction stays
  enabled because all scheduling inputs are rank-invariant.
- Per-rank views (who owns this step's token, tokens held by this rank)
  become closed-form functions of the global position, derived into the
  request fields each step (_set_helix_rank_fields / _helix_local_len).
  Decode placement follows the global continuation round-robin, matching
  the [rank::cp] striding the context side ships.
- Quota converters interpret max_tokens as GLOBAL tokens (a rank-local
  byte budget buys cp times as many global tokens); max_blocks_per_seq
  counts ledger blocks (== per-rank physical pages).
- update_resources keeps vanilla capacity math (capacity is global);
  history stays untouched under helix (reuse is disabled and
  max_beam_num_tokens mixes the rank-local prompt slice with the global
  decode stream). Note: with the global ledger, radix hashing over
  ledger blocks becomes rank-consistent, so full-page reuse is a viable
  follow-up; it stays disabled here.
- Frozen dummy fields keep V1 parity (padding dummies never pass the
  per-step derivation).
- _util.py: promote _skip_est for helix+V2 so build_managers configures
  the quota, and add the fraction fallback emitting a global max_tokens.

Depends on #17811 (helix x overlap-scheduler position fix) for
overlap-on serving, like #17795.

8 unit tests: closed-form vs brute-force ownership, cross-rank
consistency (one active rank; seqlen/past_seen sum invariants), ledger
rank-invariance, quota scaling (incl. inf passthrough and cp==1
identity), history handling, quota fallback, estimation-skip promotion,
dummy sum invariant.

Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py       |  38 +++
 .../_torch/pyexecutor/kv_cache_manager_v2.py  | 111 +++++++-
 .../pyexecutor/scheduler/scheduler_v2.py      |  11 +
 .../executor/test_kv_cache_manager_v2.py      |   3 +
 ...st_kv_cache_manager_v2_helix_superblock.py | 267 ++++++++++++++++++
 5 files changed, 424 insertions(+), 6 deletions(-)
 create mode 100644 tests/unittest/_torch/executor/test_kv_cache_manager_v2_helix_superblock.py

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 6000aecbc1b3..84633f9f1d74 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1047,6 +1047,11 @@ 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:
+                # 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).
+                self._skip_est = True
         model_config = self._model_engine.model.model_config
         if model_config.attn_backend == "VANILLA":
             estimating_kv_cache = False
@@ -1089,6 +1094,36 @@ 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 V1-style fraction sizing.
+        Quotas are GLOBAL tokens (rank-local budget x cp_size); the manager
+        min-syncs 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()
+        local_tokens = cost.tokens_for_budget(int(free_mem * fraction))
+        max_tokens = int(local_tokens) * self._mapping.cp_size
+        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} global tokens (super-block ledger). "
+            "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.
@@ -1099,6 +1134,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 79731f60ce16..6bdb5b2d2858 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,13 @@ 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):
+            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 +1523,12 @@ 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:
+            max_tokens = -(-int(max_tokens) // self._helix_cp_size)
+        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 +1951,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,
@@ -2231,6 +2276,25 @@ 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 scheduler can run before py_decoding_iter is seeded; treating
+        # the first read as decode step 1 mirrors V1.
+        pos = req.total_input_len_cp + max(1, req.py_decoding_iter) - 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,6 +2319,8 @@ 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
+        if self._has_cp_helix and not req.is_dummy_request:
+            self._set_helix_rank_fields(req)
         return kv_cache.resize(self._required_gen_capacity(req, kv_cache.capacity))
 
     def revert_allocate_generation(self, req: LlmRequest) -> None:
@@ -2374,7 +2440,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 +2481,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"
         )
+        assert not (self._has_cp_helix and not req.is_dummy_request), (
+            "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 +2519,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 +3243,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 +3318,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 +3766,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/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py
index a4a298d16d94..6abd501aef28 100644
--- a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py
@@ -993,6 +993,17 @@ def _try_schedule_generation(
         success = self.kv_cache_manager.try_allocate_generation(req)
 
         if not success:
+            if self.kv_cache_manager._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/max_tokens or "
+                    f"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_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..7c91960acd02
--- /dev/null
+++ b/tests/unittest/_torch/executor/test_kv_cache_manager_v2_helix_superblock.py
@@ -0,0 +1,267 @@
+# 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, so every rank's ledger advances identically
+and no rotation state exists.
+"""
+
+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, cp_size, phys):
+    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, cp_rank, cp_size, phys):
+    """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():
+    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():
+    """For any (prompt_len, decoding_iter): 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 decoding_iter in (0, 1, 2, 7, 40):  # 0 exercises the max(1,...) clamp
+            fields = []
+            for r in range(cp_size):
+                req = SimpleNamespace(
+                    total_input_len_cp=total_input,
+                    py_decoding_iter=decoding_iter,
+                )
+                KVCacheManagerV2._set_helix_rank_fields(_mgr(r, cp_size, phys), req)
+                fields.append(req)
+            pos = total_input + max(1, decoding_iter) - 1
+            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():
+    """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_decoding_iter=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_quota_converters_scale_by_cp():
+    m = SimpleNamespace(
+        _has_cp_helix=True,
+        _helix_cp_size=4,
+        _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 -> 400 global tokens.
+    assert KVCacheManagerV2._get_max_tokens_from_quota(m, 12345) == 400.0
+    # inf (all-SWA) passes through unscaled.
+    m_inf = SimpleNamespace(
+        _has_cp_helix=True,
+        _helix_cp_size=4,
+        _get_max_tokens_from_quota_impl=lambda quota: float("inf"),
+    )
+    assert math.isinf(KVCacheManagerV2._get_max_tokens_from_quota(m_inf, 1))
+    # Global tokens -> per-rank physical tokens (ceil) -> bytes.
+    assert KVCacheManagerV2._get_quota_from_max_tokens(m, 401) == 101 * 7
+    # cp == 1 (non-helix) is the identity.
+    m1 = SimpleNamespace(
+        _has_cp_helix=False,
+        _helix_cp_size=1,
+        _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():
+    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_emits_global_tokens(monkeypatch):
+    """Creator fallback: the rank-local byte budget buys N physical tokens,
+    i.e. N * cp_size global (super-block ledger) tokens."""
+    import torch
+
+    from tensorrt_llm._torch.pyexecutor._util import CacheCost, 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,
+            ),
+            _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: (1e6 * 0.5 - 8000) // 1000 = 492 physical -> 1968 global.
+    c = creator(0, None)
+    assert KvCacheCreator._configure_helix_kv_cache_capacity(c) is None
+    assert c._kv_cache_config.max_tokens == 492 * 4
+    # 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_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()."""
+    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():
+    """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
+
+    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,
+        )
+
+
+def test_dummy_frozen_fields_sum_invariant():
+    """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

From f084a358b814d275cb8e5c33510ed2d36bc94a10 Mon Sep 17 00:00:00 2001
From: Liao Lanyu <108499334+lancelly@users.noreply.github.com>
Date: Tue, 18 Aug 2026 19:53:40 -0700
Subject: [PATCH 2/5] [None][fix] Address review: page-granular helix quota
 rounding and guard polish

- Quota converters round at page/ledger-block granularity: rank-local
  capacity floors to whole physical pages before scaling to global
  tokens (a ledger block allocates one full page on every CP rank), and
  the reverse conversion rounds global tokens up to whole ledger blocks
  before deriving the rank-local byte quota.
- Fraction fallback floors to whole pages likewise, and an explicit
  non-positive kv_cache_config.max_tokens is rejected instead of
  silently falling through to fraction sizing.
- Restrict the _skip_est promotion to CpType.HELIX: other CP types have
  no sizing path in configure_kv_cache_capacity.
- resize_context helix guard raises ValueError instead of assert (per
  coding guidelines; asserts vanish under -O).
- get_num_available_tokens documents its unit (GLOBAL ledger tokens
  under helix) and why min() against per-forward budgets stays
  consistent (helix context forwards replicate all tokens per rank).
- V2Scheduler caches has_cp_helix at construction; the no-evict error
  message states that max_tokens is global across CP ranks.
- Tests updated for page-granular expectations; annotations added.

Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py       | 19 +++--
 .../_torch/pyexecutor/kv_cache_manager_v2.py  | 33 +++++++--
 .../pyexecutor/scheduler/scheduler_v2.py      |  8 ++-
 ...st_kv_cache_manager_v2_helix_superblock.py | 71 ++++++++++++-------
 4 files changed, 93 insertions(+), 38 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index 84633f9f1d74..8302a928fba6 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1047,10 +1047,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:
+            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":
@@ -1101,16 +1104,24 @@ def _configure_helix_kv_cache_capacity(self) -> None:
         Quotas are GLOBAL tokens (rank-local budget x cp_size); the manager
         min-syncs across ranks.
         """
+        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:
+                (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()
         cost = self._get_kv_size_per_token()
-        local_tokens = cost.tokens_for_budget(int(free_mem * fraction))
-        max_tokens = int(local_tokens) * self._mapping.cp_size
+        # Floor to whole physical pages: a ledger block allocates one full
+        # page on every CP rank, so a partial trailing page is never usable.
+        local_tokens = (int(cost.tokens_for_budget(int(free_mem * fraction))) //
+                        self._tokens_per_block * self._tokens_per_block)
+        max_tokens = local_tokens * self._mapping.cp_size
         if max_tokens <= 0:
             raise ValueError(
                 "Helix CP: fraction-based KV sizing found no usable free "
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
index 6bdb5b2d2858..3a6e27f72432 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
@@ -1486,6 +1486,10 @@ 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
 
@@ -1525,7 +1529,11 @@ def _get_max_tokens_from_quota_impl(self, quota: int) -> float:
     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:
-            max_tokens = -(-int(max_tokens) // self._helix_cp_size)
+            # 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:
@@ -2202,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.
@@ -2481,11 +2500,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"
         )
-        assert not (self._has_cp_helix and not req.is_dummy_request), (
-            "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."
-        )
+        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
diff --git a/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py b/tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py
index 6abd501aef28..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,7 +994,7 @@ def _try_schedule_generation(
         success = self.kv_cache_manager.try_allocate_generation(req)
 
         if not success:
-            if self.kv_cache_manager._has_cp_helix:
+            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.
@@ -1001,8 +1002,9 @@ def _try_schedule_generation(
                     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/max_tokens or "
-                    f"reduce concurrency."
+                    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_kv_cache_manager_v2_helix_superblock.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2_helix_superblock.py
index 7c91960acd02..da160022343c 100644
--- 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
@@ -30,7 +30,7 @@
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequestState
 
 
-def _mgr(cp_rank, cp_size, phys):
+def _mgr(cp_rank: int, cp_size: int, phys: int) -> SimpleNamespace:
     m = SimpleNamespace(
         tokens_per_block=phys,
         _ledger_tokens_per_block=phys * cp_size,
@@ -42,12 +42,12 @@ def _mgr(cp_rank, cp_size, phys):
     return m
 
 
-def _brute_local_len(global_len, cp_rank, cp_size, phys):
+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():
+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):
@@ -58,7 +58,7 @@ def test_helix_local_len_matches_brute_force():
                     ), (cp_size, phys, cp_rank, global_len)
 
 
-def test_set_helix_rank_fields_cross_rank_consistency():
+def test_set_helix_rank_fields_cross_rank_consistency() -> None:
     """For any (prompt_len, decoding_iter): 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
@@ -87,7 +87,7 @@ def test_set_helix_rank_fields_cross_rank_consistency():
                 assert past_seen[r] >= 0
 
 
-def test_ledger_is_rank_invariant():
+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
@@ -104,28 +104,37 @@ def test_ledger_is_rank_invariant():
     assert max(lens) - min(lens) <= phys
 
 
-def test_quota_converters_scale_by_cp():
+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 -> 400 global tokens.
-    assert KVCacheManagerV2._get_max_tokens_from_quota(m, 12345) == 400.0
+    # 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 -> per-rank physical tokens (ceil) -> bytes.
-    assert KVCacheManagerV2._get_quota_from_max_tokens(m, 401) == 101 * 7
-    # cp == 1 (non-helix) is the identity.
+    # 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,
     )
@@ -133,7 +142,7 @@ def test_quota_converters_scale_by_cp():
     assert KVCacheManagerV2._get_quota_from_max_tokens(m1, 400) == 2800
 
 
-def test_update_resources_leaves_history_untouched_under_helix():
+def test_update_resources_leaves_history_untouched_under_helix() -> None:
     resizes = []
 
     kv = SimpleNamespace(
@@ -167,9 +176,10 @@ def test_update_resources_leaves_history_untouched_under_helix():
     assert resizes == [(100, 54)]
 
 
-def test_helix_quota_fallback_emits_global_tokens(monkeypatch):
+def test_helix_quota_fallback_emits_global_tokens(monkeypatch: pytest.MonkeyPatch) -> None:
     """Creator fallback: the rank-local byte budget buys N physical tokens,
-    i.e. N * cp_size global (super-block ledger) tokens."""
+    floored to whole pages, i.e. (N // page) * page * cp_size global
+    (super-block ledger) tokens."""
     import torch
 
     from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator
@@ -177,6 +187,7 @@ def test_helix_quota_fallback_emits_global_tokens(monkeypatch):
     def creator(max_gpu_total_bytes, max_tokens):
         return SimpleNamespace(
             _mapping=SimpleNamespace(cp_size=4),
+            _tokens_per_block=32,
             _kv_cache_config=SimpleNamespace(
                 max_gpu_total_bytes=max_gpu_total_bytes,
                 max_tokens=max_tokens,
@@ -191,26 +202,34 @@ def creator(max_gpu_total_bytes, max_tokens):
     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: (1e6 * 0.5 - 8000) // 1000 = 492 physical -> 1968 global.
+    # 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: (1e6 * 0.5 - 8000) // 1000 = 492 physical tokens, floored to
+    # whole 32-token pages = 480 (a ledger block needs one full page on
+    # every rank) -> 1920 global.
     c = creator(0, None)
     assert KvCacheCreator._configure_helix_kv_cache_capacity(c) is None
-    assert c._kv_cache_config.max_tokens == 492 * 4
+    assert c._kv_cache_config.max_tokens == 480 * 4
     # 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_estimation_prepare_promotes_skip_est_for_v2():
+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()."""
+    _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):
+    def creator(is_v2, cp_type=CpType.HELIX):
         return SimpleNamespace(
             _skip_est=False,
-            _mapping=SimpleNamespace(cp_config={"cp_type": CpType.HELIX}),
+            _mapping=SimpleNamespace(cp_config={"cp_type": cp_type}),
             _is_kv_cache_manager_v2=is_v2,
             _model_engine=SimpleNamespace(
                 model=SimpleNamespace(
@@ -225,17 +244,19 @@ def creator(is_v2):
     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():
+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
 
     sched = SimpleNamespace(
-        kv_cache_manager=SimpleNamespace(
-            _has_cp_helix=True, try_allocate_generation=lambda req: False
-        ),
+        has_cp_helix=True,
+        kv_cache_manager=SimpleNamespace(try_allocate_generation=lambda req: False),
     )
     req = SimpleNamespace(
         py_request_id=7,
@@ -256,7 +277,7 @@ def test_scheduler_allocation_failure_raises_under_helix():
         )
 
 
-def test_dummy_frozen_fields_sum_invariant():
+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."""

From d619260aafbb3e0521e0173db34dc915c6c68910 Mon Sep 17 00:00:00 2001
From: Liao Lanyu <108499334+lancelly@users.noreply.github.com>
Date: Wed, 19 Aug 2026 20:08:17 -0700
Subject: [PATCH 3/5] [None][test] Set the helix flag in partial manager shells
 and mocks

The helix superblock ledger adds _has_cp_helix branches to shared V2
manager and scheduler paths. Existing executor tests hand-assemble
partial manager instances (object.__new__) and mock managers, so their
attribute lists must now include the flag: pin _has_cp_helix=False in
the scheduler and dual-pool mock factories and in the estimation,
capacity-only, and Mamba hybrid shells.

Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com>
---
 tests/unittest/_torch/executor/test_dual_pool_kv_cache.py      | 1 +
 tests/unittest/_torch/executor/test_kv_cache_estimation.py     | 1 +
 .../unittest/_torch/executor/test_kv_cache_v2_capacity_only.py | 1 +
 tests/unittest/_torch/executor/test_kv_cache_v2_scheduler.py   | 1 +
 tests/unittest/_torch/executor/test_mamba_cache_manager.py     | 3 +++
 5 files changed, 7 insertions(+)

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_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]

From de56cc90bad9243a736adc3d42357ce4753522ca Mon Sep 17 00:00:00 2001
From: Liao Lanyu <108499334+lancelly@users.noreply.github.com>
Date: Wed, 19 Aug 2026 23:29:55 -0700
Subject: [PATCH 4/5] [None][test] Match the new _try_schedule_generation
 signature

Main added recompute-pause scheduling params (recompute_pause_state,
recompute_paused, inflight_request_ids) to _try_schedule_generation;
the helix no-evict raise fires before any of them is used, so the
direct-call test passes fresh placeholders.

Signed-off-by: Liao Lanyu <108499334+lancelly@users.noreply.github.com>
---
 .../executor/test_kv_cache_manager_v2_helix_superblock.py | 8 +++++++-
 1 file changed, 7 insertions(+), 1 deletion(-)

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
index da160022343c..a5273bf6f452 100644
--- 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
@@ -252,7 +252,10 @@ def creator(is_v2, cp_type=CpType.HELIX):
 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
+    from tensorrt_llm._torch.pyexecutor.scheduler.scheduler_v2 import (
+        KVCacheV2Scheduler,
+        _RecomputePauseState,
+    )
 
     sched = SimpleNamespace(
         has_cp_helix=True,
@@ -272,7 +275,10 @@ def test_scheduler_allocation_failure_raises_under_helix() -> None:
             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,
         )
 

From 409560a96676cd710ce5864f9806ac37da079efc Mon Sep 17 00:00:00 2001
From: lancelly <108499334+lancelly@users.noreply.github.com>
Date: Fri, 21 Aug 2026 23:12:44 -0700
Subject: [PATCH 5/5] [None][fix] KVCM-V2 helix: address review round

- Derive the helix decode position from a manager-owned per-request step
  counter (committed on successful allocation, given back on revert)
  instead of py_decoding_iter, which the overlap loop advances only after
  scheduling; the stale read repeated the first decode position and
  overwrote the first generated token's KV.
- Reject V1 + helix + TRTLLM_SKIP_KV_CACHE_ESTIMATION explicitly instead
  of emitting V2-ledger (global) quotas that V1 reads as rank-local.
- Helix fraction fallback sets max_gpu_total_bytes (rank-local byte cap)
  instead of max_tokens, which gets inflated by 1/max_util_for_resume.
- Tests: ledger-position immunity to sampler timing, derivation
  idempotence on retry/revert, V1+helix rejection; fallback test updated
  to the bytes knob.

Signed-off-by: lancelly <108499334+lancelly@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/_util.py       |  28 +++--
 .../_torch/pyexecutor/kv_cache_manager_v2.py  |  30 ++++-
 tensorrt_llm/_torch/pyexecutor/llm_request.py |   3 +
 ...st_kv_cache_manager_v2_helix_superblock.py | 108 +++++++++++++++---
 4 files changed, 133 insertions(+), 36 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py
index ab4faf796e57..c935ce394890 100644
--- a/tensorrt_llm/_torch/pyexecutor/_util.py
+++ b/tensorrt_llm/_torch/pyexecutor/_util.py
@@ -1082,9 +1082,10 @@ def try_prepare_estimation(self) -> bool:
     def _configure_helix_kv_cache_capacity(self) -> None:
         """Set the helix KV quota without profiling (not CP-aware).
 
-        Explicit quotas pass through; otherwise V1-style fraction sizing.
-        Quotas are GLOBAL tokens (rank-local budget x cp_size); the manager
-        min-syncs across ranks.
+        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):
@@ -1098,13 +1099,8 @@ def _configure_helix_kv_cache_capacity(self) -> None:
             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()
-        # Floor to whole physical pages: a ledger block allocates one full
-        # page on every CP rank, so a partial trailing page is never usable.
-        local_tokens = (int(cost.tokens_for_budget(int(free_mem * fraction))) //
-                        self._tokens_per_block * self._tokens_per_block)
-        max_tokens = local_tokens * self._mapping.cp_size
-        if max_tokens <= 0:
+        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 "
@@ -1112,10 +1108,11 @@ def _configure_helix_kv_cache_capacity(self) -> None:
         logger.warning(
             "Helix CP: capacity profiling is unsupported; sizing the KV "
             f"cache as fraction {fraction} of free memory -> "
-            f"max_tokens={max_tokens} global tokens (super-block ledger). "
+            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_tokens = max_tokens
+        self._kv_cache_config.max_gpu_total_bytes = budget_bytes
 
     def configure_kv_cache_capacity(self,
                                     py_executor: PyExecutor = None) -> None:
@@ -1128,6 +1125,13 @@ def configure_kv_cache_capacity(self,
 
         # 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
diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
index 3a6e27f72432..a6202947091f 100644
--- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
+++ b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
@@ -2303,10 +2303,18 @@ def _helix_local_len(self, global_len: int) -> int:
         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 scheduler can run before py_decoding_iter is seeded; treating
-        # the first read as decode step 1 mirrors V1.
-        pos = req.total_input_len_cp + max(1, req.py_decoding_iter) - 1
+        """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
@@ -2338,9 +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
-        if self._has_cp_helix and not req.is_dummy_request:
+        is_helix_req = self._has_cp_helix and not req.is_dummy_request
+        if is_helix_req:
             self._set_helix_rank_fields(req)
-        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)):
+            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.
@@ -2358,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)
         )
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/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
index a5273bf6f452..e95682373fcd 100644
--- 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
@@ -17,8 +17,9 @@
 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, so every rank's ledger advances identically
-and no rotation state exists.
+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
@@ -59,22 +60,22 @@ def test_helix_local_len_matches_brute_force() -> None:
 
 
 def test_set_helix_rank_fields_cross_rank_consistency() -> None:
-    """For any (prompt_len, decoding_iter): exactly one active rank, the
+    """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 decoding_iter in (0, 1, 2, 7, 40):  # 0 exercises the max(1,...) clamp
+        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_decoding_iter=decoding_iter,
+                    py_helix_decode_group_index=group_index,
                 )
                 KVCacheManagerV2._set_helix_rank_fields(_mgr(r, cp_size, phys), req)
                 fields.append(req)
-            pos = total_input + max(1, decoding_iter) - 1
+            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
@@ -93,7 +94,7 @@ def test_ledger_is_rank_invariant() -> None:
     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_decoding_iter=17)
+    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)
@@ -104,6 +105,65 @@ def test_ledger_is_rank_invariant() -> None:
     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,
@@ -176,24 +236,21 @@ def test_update_resources_leaves_history_untouched_under_helix() -> None:
     assert resizes == [(100, 54)]
 
 
-def test_helix_quota_fallback_emits_global_tokens(monkeypatch: pytest.MonkeyPatch) -> None:
-    """Creator fallback: the rank-local byte budget buys N physical tokens,
-    floored to whole pages, i.e. (N // page) * page * cp_size global
-    (super-block ledger) tokens."""
+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 CacheCost, KvCacheCreator
+    from tensorrt_llm._torch.pyexecutor._util import KvCacheCreator
 
     def creator(max_gpu_total_bytes, max_tokens):
         return SimpleNamespace(
             _mapping=SimpleNamespace(cp_size=4),
-            _tokens_per_block=32,
             _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))
@@ -206,18 +263,33 @@ def creator(max_gpu_total_bytes, max_tokens):
     # by fraction sizing.
     with pytest.raises(ValueError, match="must be positive"):
         KvCacheCreator._configure_helix_kv_cache_capacity(creator(0, 0))
-    # No quota: (1e6 * 0.5 - 8000) // 1000 = 492 physical tokens, floored to
-    # whole 32-token pages = 480 (a ledger block needs one full page on
-    # every rank) -> 1920 global.
+    # 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_tokens == 480 * 4
+    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().