[None][feat] KVCacheManagerV2: helix (decode-CP) support

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -1027,6 +1027,13 @@ def try_prepare_estimation(self) -> bool:
             logger.info(
                 "KV cache size estimation is not supported for context parallelism, disable it."
             )
+            if self._is_kv_cache_manager_v2:
+                # Like the encoder-decoder case below: promote to _skip_est
+                # so build_managers runs configure_kv_cache_capacity(), which
+                # sets the explicit quota KVCacheManagerV2 requires at
+                # construction (V1 sizes itself from the fraction
+                # internally, so it stays on the local flag).
+                self._skip_est = True
         model_config = self._model_engine.model.model_config
         if model_config.attn_backend == "VANILLA":
             estimating_kv_cache = False
@@ -1069,6 +1076,35 @@ def try_prepare_estimation(self) -> bool:
                 self._kv_cache_config.max_tokens = max_tokens
         return estimating_kv_cache
 
+    def _configure_helix_kv_cache_capacity(self) -> None:
+        """Set the helix KV quota without profiling (not CP-aware).
+
+        An explicit quota is honored as-is; otherwise fall back to V1's
+        fraction sizing of free memory. KVCacheManagerV2 min-syncs the
+        derived max_tokens across ranks.
+        """
+        if (self._kv_cache_config.max_gpu_total_bytes or 0) > 0 or \
+                self._kv_cache_config.max_tokens:
+            logger.info("Helix CP: skipping KV cache capacity profiling; using "
+                        "the explicitly configured quota.")
+            return
+        fraction = self._kv_cache_config.free_gpu_memory_fraction
+        free_mem, _total = torch.cuda.mem_get_info()
+        cost = self._get_kv_size_per_token()
+        max_tokens = cost.tokens_for_budget(int(free_mem * fraction))
+        if max_tokens <= 0:
+            raise ValueError(
+                "Helix CP: fraction-based KV sizing found no usable free "
+                "memory; set kv_cache_config.max_tokens or "
+                "max_gpu_total_bytes.")
+        logger.warning(
+            "Helix CP: capacity profiling is unsupported; sizing the KV "
+            f"cache as fraction {fraction} of free memory -> "
+            f"max_tokens={max_tokens} (rank-local). Set "
+            "kv_cache_config.max_tokens or max_gpu_total_bytes to "
+            "override.")
+        self._kv_cache_config.max_tokens = max_tokens
+
     def configure_kv_cache_capacity(self,
                                     py_executor: PyExecutor = None) -> None:
         """Perform KV cache capacity estimation.
@@ -1079,6 +1115,9 @@ def configure_kv_cache_capacity(self,
         mapping = self._mapping
 
         # TODO: support CP by generating dummy requests for it.
+        if mapping.cp_config.get('cp_type') == CpType.HELIX:
+            self._configure_helix_kv_cache_capacity()
+            return
         assert 'cp_type' not in mapping.cp_config
 
         fraction = self._kv_cache_config.free_gpu_memory_fraction

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -783,6 +783,29 @@ def __init__(
             "Star attention is not supported for KVCacheManagerV2"
         )
 
+        # ---- Helix (decode context parallelism) bookkeeping --------------
+        # V1 rotates decode-KV ownership across CP ranks in
+        # prepare_resources; V2 grows KV at scheduling time, so the gate
+        # lives in try_allocate_generation / revert_allocate_generation. The
+        # backend stays helix-ignorant: it only sees rank-local token counts.
+        self._has_cp_helix = mapping.has_cp_helix()
+        if self._has_cp_helix:
+            self._helix_cp_rank = mapping.cp_rank
+            self._helix_cp_size = mapping.cp_size
+            if kv_cache_config.enable_block_reuse:
+                raise ValueError(
+                    "Helix CP requires enable_block_reuse=False: each rank "
+                    "holds an interleaved slice of the sequence, so radix "
+                    "prefix matching over the local token stream is "
+                    "meaningless."
+                )
+            if is_draft:
+                raise ValueError(
+                    "Helix CP does not support speculative decoding with "
+                    "KVCacheManagerV2 yet (draft-manager mirroring and "
+                    "rewind are not round-robin aware)."
+                )
+
         self.kv_cache_type = kv_cache_type
         self.pp_layers, self.num_layers = get_pp_layers(
             num_layers,
@@ -2227,9 +2250,32 @@ def try_allocate_generation(self, req: LlmRequest) -> bool:
                 return False
             self._restore_page_index_bufs(req.py_request_id, kv_cache)
 
+        if self._has_cp_helix and not req.is_dummy_request:
+            # Round-robin gate (V1 parity). The V2 scheduler runs before
+            # the disagg handler seeds py_decoding_iter = 1, so the first
+            # schedule reads 0 and the negative modulo would silently pick
+            # rank cp_size - 1; clamping to 1 reproduces V1's owner sequence.
+            decode_iter = max(1, req.py_decoding_iter)
+            decode_block_id = (decode_iter - 1) // self.tokens_per_block
+            if decode_block_id % self._helix_cp_size != self._helix_cp_rank:
+                # Inactive rank this step: materialize nothing but report
+                # success — False would trigger per-rank eviction hunting
+                # and desynchronize the CP group.
+                req.py_helix_is_inactive_rank = True
+                return True
+            req.py_helix_is_inactive_rank = False
+            req.seqlen_this_rank_cp += 1
+
         draft_len = self._effective_draft_len(req)
         self._allocated_draft_lens[req.py_request_id] = draft_len
-        return kv_cache.resize(self._required_gen_capacity(req, kv_cache.capacity))
+        if not kv_cache.resize(self._required_gen_capacity(req, kv_cache.capacity)):
+            if self._has_cp_helix and not req.is_dummy_request:
+                # Undo the bookkeeping above: the scheduler may retry this
+                # request in the same pass (after evicting a victim) and
+                # try_allocate_generation would increment again.
+                req.seqlen_this_rank_cp -= 1
+            return False
+        return True
 
     def revert_allocate_generation(self, req: LlmRequest) -> None:
         """Undo the capacity growth from try_allocate_generation.
@@ -2247,6 +2293,10 @@ def revert_allocate_generation(self, req: LlmRequest) -> None:
         kv_cache = self.kv_cache_map.get(req.py_request_id)
         if kv_cache is None or not kv_cache.is_active:
             return
+        if self._has_cp_helix and not req.is_dummy_request and req.py_helix_is_inactive_rank:
+            # Inactive rank grew nothing this step; the default path would
+            # shrink capacity below the materialized history.
+            return
         draft_len = self._allocated_draft_lens.pop(
             req.py_request_id, self._effective_draft_len(req)
         )
@@ -2259,6 +2309,10 @@ def revert_allocate_generation(self, req: LlmRequest) -> None:
                 f"{req.py_request_id} from {kv_cache.capacity} to "
                 f"{reverted_cap}"
             )
+        if self._has_cp_helix and not req.is_dummy_request:
+            # Symmetric to the += 1 in try_allocate_generation (active rank
+            # only; inactive ranks returned above).
+            req.seqlen_this_rank_cp -= 1
 
     def revert_allocate_context(self, req: LlmRequest) -> None:
         """Undo the capacity growth from this iteration's context resize."""
@@ -3134,6 +3188,11 @@ def release_resources(
             # a non-zero number to skip illegal memory access issue in MLA kernel
             # during warmup.
             token_num = token_nums[i] if token_nums is not None else 1 + max_num_draft_tokens
+            if self._has_cp_helix:
+                # token_num >= 2 keeps the active rank's
+                # past_seen_token_num (= seqlen_this_rank_cp - 1)
+                # non-negative (V1 parity).
+                token_num = max(token_num, 2)
             # token_num - 1 is the past history length in generation.
             history_hint = max(0, token_num - 1) if is_gen else None
             encoder_output_len = encoder_output_lens[i] if encoder_output_lens is not None else None
@@ -3205,6 +3264,21 @@ def release_resources(
                 req.prompt_len = token_num - 1
                 req.py_prompt_len = req.prompt_len
                 req.py_draft_tokens = [1] * max_num_draft_tokens
+                if self._has_cp_helix:
+                    # Frozen helix fields (V1 parity): padding dummies are
+                    # appended inside forward and never pass the scheduling
+                    # gate, so these values must stay permanently valid —
+                    # last CP rank active, shared synthetic global length.
+                    if self._helix_cp_rank == self._helix_cp_size - 1:
+                        req.py_helix_is_inactive_rank = False
+                        req.prompt_len = token_num - 1
+                    else:
+                        req.py_helix_is_inactive_rank = True
+                        req.prompt_len = token_num
+                    req.py_prompt_len = req.prompt_len
+                    req.seqlen_this_rank_cp = req.prompt_len
+                    req.total_input_len_cp = token_num * self._helix_cp_size - 1
+                    req.py_decoding_iter = 1
                 if prepare_resource:
                     new_capacity = kv_cache.capacity + _kv_draft + 1
                     success = kv_cache.resize(new_capacity, history_length=history_hint)
@@ -3633,14 +3707,32 @@ def update_resources(
                 # it accumulates in the draft KV cache after every generation
                 # step. Target managers do not allocate this reserve slack.
                 rewind_len += max(self._kv_reserve_draft_tokens - runtime_draft_len, 0)
-            new_capacity = (
-                None
-                if req.state in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT)
-                else kv_cache.capacity - rewind_len
-            )
-            history_length = (
-                None if self.kv_compression_manages_history else req.max_beam_num_tokens - 1
-            )
+            if self._has_cp_helix and not req.is_dummy_request:
+                # max_beam_num_tokens is the GLOBAL length; on a rank
+                # holding ~1/cp of the tokens it would inflate capacity past
+                # the gate. Floor at the rank-local count; history stays
+                # untouched (its consumers are disabled under helix and it
+                # must never decrease per the backend resize contract).
+                new_capacity = (
+                    None
+                    if req.state
+                    in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT)
+                    else max(
+                        kv_cache.capacity - rewind_len,
+                        req.seqlen_this_rank_cp,
+                    )
+                )
+                history_length = None
+            else:
+                new_capacity = (
+                    None
+                    if req.state
+                    in (LlmRequestState.GENERATION_COMPLETE, LlmRequestState.CONTEXT_INIT)
+                    else kv_cache.capacity - rewind_len
+                )
+                history_length = (
+                    None if self.kv_compression_manages_history else req.max_beam_num_tokens - 1
+                )
             success = kv_cache.resize(new_capacity, history_length)
             if not success:
                 raise ValueError(

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -2201,6 +2201,11 @@ def __init__(
 
         # Derive ssm_state_shape and conv_state_shape from mamba params (same as MambaCacheManager)
         tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1
+        if mapping.has_cp_helix():
+            # Helix repurposes CP ranks as TP for non-attention layers:
+            # this rank's mamba state is a 1/(tp*cp) head slice; bare
+            # tp_size would allocate a full replica per CP rank.
+            tp_size = mapping.tp_size * mapping.cp_size
         d_inner = mamba_head_dim * mamba_num_heads
         conv_dim = d_inner + 2 * mamba_n_groups * mamba_d_state
         nheads = mamba_num_heads
@@ -2936,6 +2941,11 @@ def __init__(
 
         if self.local_num_mamba_layers > 0:
             tp_size = mapping.tp_size if not mapping.enable_attention_dp else 1
+            if mapping.has_cp_helix():
+                # Helix repurposes CP ranks as TP for non-attention layers:
+                # this rank's mamba state is a 1/(tp*cp) head slice; bare
+                # tp_size would allocate a full replica per CP rank.
+                tp_size = mapping.tp_size * mapping.cp_size
             d_inner = mamba_head_dim * mamba_num_heads
             grouped_state_dim = mamba_n_groups * mamba_d_state
             conv_dim = d_inner + 2 * grouped_state_dim

tensorrt_llm/_torch/pyexecutor/scheduler/scheduler_v2.py

@@ -961,6 +961,19 @@ def _try_schedule_generation(
         success = self.kv_cache_manager.try_allocate_generation(req)
 
         if not success:
+            if self.kv_cache_manager._has_cp_helix:
+                # Rank-local eviction decisions diverge across CP ranks
+                # (round-robin skew) and would desynchronize the group; fail
+                # loudly — capacity must be provisioned so scheduled
+                # requests always fit (GUARANTEED_NO_EVICT-like).
+                raise RuntimeError(
+                    f"[V2Scheduler] KV allocation failed for helix request "
+                    f"{req.py_request_id}; eviction is disabled under helix "
+                    f"CP because rank-local eviction decisions diverge "
+                    f"across the CP group. Increase "
+                    f"kv_cache_config.max_gpu_total_bytes/max_tokens or "
+                    f"reduce concurrency."
+                )
             req_it_end, success = self._try_evict_for_gen(
                 req, requests_list, req_it, req_it_end, evicted
             )