[refactor] 调整命名规范、移除 patch_manager 等无用代码、优化代码可读性

docs/concepts.md

@@ -63,7 +63,7 @@
 **定义**：一组具有相同 provider 和相同 prefix_len 的复用关系形成的调试/统计视图。
 
 **说明**：
-- Prefix Group 不再是 Phase 1 的核心语义
+- Prefix Group 不再是 Phase 1 的核心语义，在prefix-sharing运行时中并没有实际用途
 - 同一个 provider 可出现在多个 Prefix Group 中，分别对应不同 `prefix_len`
 - 单个 `group_id` 不能完整表达 provider 的所有复用关系
 - 执行计划以 `PrefixReuseSpec`、`provider_index`、`prefix_lens` 为准

prefix-sharing/prefix_sharing/__init__.py

@@ -24,7 +24,6 @@
 from prefix_sharing.core.config import PrefixSharingConfig, PrefixSharingConfigError
 from prefix_sharing.core.prefix_detector import PrefixReuseSpec, TriePrefixDetector
 from prefix_sharing.core.planner import PrefixLastRestoreSpec, PrefixSharingPlan, PrefixSharingPlanner
-from prefix_sharing.integrations.verl_mcore import enable_prefix_sharing, prefix_sharing_enabled
 
 
 __all__ = [
@@ -35,8 +34,6 @@
     "PrefixReuseSpec",
     "PrefixSharingPlanner",
     "TriePrefixDetector",
-    "enable_prefix_sharing",
-    "prefix_sharing_enabled",
 ]
 
 # ── Monkey-patch auto-activation ──
@@ -45,10 +42,6 @@
 # This allows both ENABLE_PREFIX_SHARING env var and prefix_sharing_config yaml
 # key to independently control the feature — no separate "install gate" needed.
 #
-# For verl_v070 environments, the invasive import in megatron_actor.py
-# still works (enable_prefix_sharing / prefix_sharing_enabled), but the
-# setup module is not invoked — the old integration code handles everything.
-#
 # The setup.install() call is safe even when verl/Megatron are not present:
 # if the detected versions match no compat matrix entry, it raises
 # IncompatibleEnvironment which we catch and log as a warning (no patches

prefix-sharing/prefix_sharing/core/planner.py

@@ -206,42 +206,42 @@ def __post_init__(self) -> None:
 
     def plan(
         self,
-        input_ids: Sequence[Sequence[int]],
+        sequences: Sequence[Sequence[int]],
         *,
         forward_id: int | None = None,
         micro_batch_id: int | None = None,
     ) -> PrefixSharingPlan:
-        detection = self.detector.detect(input_ids)
+        detect_result = self.detector.detect(sequences)
         return self.plan_from_detection(
-            input_ids,
-            detection,
+            sequences,
+            detect_result,
             forward_id=forward_id,
             micro_batch_id=micro_batch_id,
         )
 
     def plan_from_detection(
         self,
-        input_ids: Sequence[Sequence[int]],
-        detection: PrefixDetectionResult,
+        sequences: Sequence[Sequence[int]],
+        detect_result: PrefixDetectionResult,
         *,
         forward_id: int | None = None,
         micro_batch_id: int | None = None,
     ) -> PrefixSharingPlan:
-        if len(input_ids) != detection.batch_size:
-            raise ValueError("input_ids batch size does not match detection result")
+        if len(sequences) != detect_result.batch_size:
+            raise ValueError("sequences batch size does not match detection result")
         if forward_id is None:
             forward_id = next(_forward_ids)
         if micro_batch_id is None:
             self._micro_batch_counter += 1
             micro_batch_id = self._micro_batch_counter
 
-        batch_size = len(input_ids)
-        original_lengths = [len(seq) for seq in input_ids]
-        group_ids = list(detection.group_ids)
-        is_provider = list(detection.is_provider)
-        provider_index = list(detection.provider_index)
-        prefix_lens = list(detection.prefix_lens)
-        reuse_specs = list(detection.reuse_specs)
+        batch_size = len(sequences)
+        original_lengths = [len(seq) for seq in sequences]
+        group_ids = list(detect_result.group_ids)
+        is_provider = list(detect_result.is_provider)
+        provider_index = list(detect_result.provider_index)
+        prefix_lens = list(detect_result.prefix_lens)
+        reuse_specs = list(detect_result.reuse_specs)
 
         suffix_lens: list[int] = []
         kept_lengths_q: list[int] = []
@@ -253,18 +253,22 @@ def plan_from_detection(
         loss_mask_keep_ranges: list[tuple[int, int]] = []
         restore_specs: list[PrefixLastRestoreSpec] = []
 
-        for index, original_len in enumerate(original_lengths):
-            prefix_len = prefix_lens[index]
+        for seq_idx, original_len in enumerate(original_lengths):
+            prefix_len = prefix_lens[seq_idx]
             if prefix_len > original_len:
-                raise ValueError(f"prefix_len exceeds sequence length for batch index {index}")
-            is_reuser = provider_index[index] != index and prefix_len > 0
+                raise ValueError(f"prefix_len exceeds sequence length for batch index {seq_idx}")
+            is_reuser = provider_index[seq_idx] != seq_idx and prefix_len > 0
             suffix_len = original_len - prefix_len if is_reuser else original_len
             suffix_lens.append(suffix_len)
             expanded_lengths_kv.append(original_len)
 
             if is_reuser:
-                keep_start = prefix_len
-                keep_end = original_len
+                # 1 - 为输入数据的预处理做准备：输入阶段要对 reuser 进行序列裁剪
+                #   - 只需保留suffix部分
+                #   - prefix 部分
+                #       - KV 激活值：在 attention 时候读取缓存并进行拼接即可
+                #       - logp：可以在后处理阶段直接从provider复制得到
+                keep_start, keep_end = prefix_len, original_len
                 kept_len = suffix_len
                 q_offset = prefix_len
 
@@ -274,22 +278,21 @@ def plan_from_detection(
                 # provider's already-restored row (identical across the shared
                 # prefix), so the planner only emits the prefix-last token.
                 # The logits at position prefix_len-1 predict the first suffix
-                # token whose label is input_ids[prefix_len] (differs per
+                # token whose label is sequences[prefix_len] (differs per
                 # reuser); the restored logprob is written to 2D position
                 # prefix_len - 1 (the label slot for first-suffix prediction).
                 if prefix_len > 0 and suffix_len > 0:
                     restore_specs.append(
                         PrefixLastRestoreSpec(
-                            reuse_idx_in_batch=index,
-                            provider_idx_in_batch=provider_index[index],
-                            group_id=group_ids[index],
+                            reuse_idx_in_batch=seq_idx,
+                            provider_idx_in_batch=provider_index[seq_idx],
+                            group_id=group_ids[seq_idx],
                             target_2d_pos=prefix_len - 1,
-                            label_value=input_ids[index][prefix_len],
+                            label_value=sequences[seq_idx][prefix_len],
                         )
                     )
             else:
-                keep_start = 0
-                keep_end = original_len
+                keep_start, keep_end = 0, original_len
                 kept_len = original_len
                 q_offset = 0
 

prefix-sharing/prefix_sharing/core/prefix_detector.py

@@ -101,7 +101,7 @@ class PrefixReuseSpec:
 
 @dataclass(frozen=True)
 class PrefixDetectionResult:
-    """Per-batch detection output with per-sample reuse relations.
+    """Per-batch detection output with per-sample reuse relations (reuse_specs).
 
     ``reuse_specs`` is the semantic source of truth. The tuple fields
     ``group_ids``, ``provider_index``, ``prefix_lens``, and ``is_provider`` are
@@ -133,8 +133,8 @@ class PrefixDetectionResult:
     """
 
     batch_size: int
-    reuse_specs: tuple[PrefixReuseSpec, ...]
-    groups: tuple[PrefixGroup, ...]
+    reuse_specs: tuple[PrefixReuseSpec, ...] # 复用关系列表，每个元素是一对 reuser => provider 的复用关系
+    prefix_groups: tuple[PrefixGroup, ...]
     group_ids: tuple[int, ...]
     provider_index: tuple[int, ...]
     prefix_lens: tuple[int, ...]
@@ -169,66 +169,73 @@ def __init__(self, min_prefix_len: int = 1, min_group_size: int = 2) -> None:
         self.min_prefix_len = min_prefix_len
         self.min_group_size = min_group_size
 
-    def detect(self, input_ids: Sequence[TokenSequence]) -> PrefixDetectionResult:
-        batch_size = len(input_ids)
-        root = _TrieNode()
+    def detect(self, sequences: Sequence[TokenSequence]) -> PrefixDetectionResult:
+        batch_size = len(sequences)
         group_ids = [-1] * batch_size
-        provider_index = list(range(batch_size))
-        prefix_lens = [0] * batch_size
-        is_provider = [True] * batch_size
+        prefix_lens = [0] * batch_size # 没有找到可复用前缀前，所有序列的要复用的前缀长度都是 0
+        # 没有找到可复用前缀前，所有序列都是 provider，其 provider 索引是自己
+        is_provider, provider_index = [True] * batch_size, list(range(batch_size))
         reuse_specs: list[PrefixReuseSpec] = []
         group_key_to_id: dict[tuple[int, int], int] = {}
         group_members: dict[int, list[int]] = {}
 
-        for index, seq in enumerate(input_ids):
+        # 遍历所有序列，构建前缀树，一次Trie遍历同时完成：前缀匹配 + 新节点插入
+        root = _TrieNode()
+        for seq_idx, seq in enumerate(sequences):
             node = root
-            matched = 0
-            matched_provider = -1
+            matched_prefix_len = 0
+            matched_provider_idx = -1
             matched_group_size = 0
+            detect_status = "prefix_matching"  # 前缀匹配阶段
+
+            # 遍历序列中的每个 token，前缀树继续生长
             for token in seq:
                 child = node.children.get(int(token))
+
+                # 达到最长匹配
                 if child is None:
-                    break
+                    # 前缀匹配阶段结束，切换为 Trie 构建阶段
+                    detect_status = "trie_growing"
+                    child = _TrieNode(node.depth + 1)
+                    child.provider_index = seq_idx
+                    node.children[token] = child
+
+                # 还没达到最长匹配，继续匹配
                 node = child
-                matched += 1
-                if node.provider_index >= 0:
-                    matched_provider = node.provider_index
-                    matched_group_size = len(node.indices) + 1
+                if detect_status == "prefix_matching":
+                    matched_prefix_len += 1
+                    if node.provider_index >= 0:
+                        matched_provider_idx = node.provider_index # 将父亲的 provider 传递给儿子
+                        matched_group_size = len(node.indices) + 1 # 此时 node.indices 尚未包含当前 seq_idx
 
+                node.indices.append(seq_idx)
+
+            # 前缀检测已完成，记录复用关系
             if (
-                matched >= self.min_prefix_len
-                and matched_provider >= 0
+                matched_prefix_len >= self.min_prefix_len
                 and matched_group_size >= self.min_group_size
+                and matched_provider_idx >= 0
             ):
-                spec = PrefixReuseSpec(
-                    reuse_idx_in_batch=index,
-                    provider_idx_in_batch=matched_provider,
-                    prefix_len=matched,
+                # 为该序列记录复用关系
+                provider_index[seq_idx] = matched_provider_idx
+                prefix_lens[seq_idx] = matched_prefix_len
+                is_provider[seq_idx] = False
+                this_reuse_spec = PrefixReuseSpec(
+                    reuse_idx_in_batch=seq_idx,
+                    provider_idx_in_batch=matched_provider_idx,
+                    prefix_len=matched_prefix_len,
                 )
-                reuse_specs.append(spec)
-                provider_index[index] = matched_provider
-                prefix_lens[index] = matched
-                is_provider[index] = False
+                reuse_specs.append(this_reuse_spec)
 
-                group_key = (matched_provider, matched)
+                # 用于后续构造 PrefixGroup（在当前版本中暂时没有实际用途）
+                group_key = (matched_provider_idx, matched_prefix_len)
                 group_id = group_key_to_id.setdefault(group_key, len(group_key_to_id))
-                group_ids[index] = group_id
+                group_ids[seq_idx] = group_id
                 if group_id not in group_members:
-                    group_members[group_id] = [matched_provider]
-                group_members[group_id].append(index)
-
-            node = root
-            node.indices.append(index)
-            for token in seq:
-                token = int(token)
-                child = node.children.get(token)
-                if child is None:
-                    child = _TrieNode(node.depth + 1)
-                    child.provider_index = index
-                    node.children[token] = child
-                node = child
-                node.indices.append(index)
+                    group_members[group_id] = [matched_provider_idx]
+                group_members[group_id].append(seq_idx)
 
+        # 构造 PrefixGroup（在当前版本中暂时没有实际用途）
         groups = [
             PrefixGroup(
                 group_id=group_id,
@@ -242,10 +249,11 @@ def detect(self, input_ids: Sequence[TokenSequence]) -> PrefixDetectionResult:
             for members in (group_members[group_id],)
         ]
 
+        # 汇总整个 batch 的前缀复用关系
         return PrefixDetectionResult(
             batch_size=batch_size,
             reuse_specs=tuple(reuse_specs),
-            groups=tuple(groups),
+            prefix_groups=tuple(groups),
             group_ids=tuple(group_ids),
             provider_index=tuple(provider_index),
             prefix_lens=tuple(prefix_lens),

prefix-sharing/prefix_sharing/integrations/__init__.py

@@ -6,14 +6,10 @@
     current_prefix_sharing_context,
     prefix_sharing_runtime_context,
 )
-from prefix_sharing.integrations.patch_manager import PatchHandle, PatchManager
 from prefix_sharing.integrations.parallel_info import MegatronParallelInfo, get_megatron_parallel_info
 from prefix_sharing.backends.packed_layout import PackedBatchLayout
 from prefix_sharing.integrations.verl_mcore import (
     PrefixSharingRuntimeState,
-    VerlMCoreIntegration,
-    enable_prefix_sharing,
-    prefix_sharing_enabled,
     build_prefix_sharing_micro_batch_verl070,
     build_prefix_sharing_micro_batch_verl080,
     restore_reuser_prefix_columns_2d,
@@ -24,18 +20,13 @@
 )
 
 __all__ = [
-    "PatchHandle",
-    "PatchManager",
     "PackedBatchLayout",
     "PackedPrefixLastRestoreIndex",
     "MegatronParallelInfo",
     "PrefixSharingRuntimeContext",
     "PrefixSharingRuntimeState",
-    "VerlMCoreIntegration",
     "current_prefix_sharing_context",
-    "enable_prefix_sharing",
     "prefix_sharing_runtime_context",
-    "prefix_sharing_enabled",
     "build_prefix_sharing_micro_batch_verl070",
     "build_prefix_sharing_micro_batch_verl080",
     "restore_reuser_prefix_columns_2d",