diff --git a/docs/concepts.md b/docs/concepts.md index 89089836..111e84bc 100644 --- a/docs/concepts.md +++ b/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` 为准 diff --git a/prefix-sharing/prefix_sharing/__init__.py b/prefix-sharing/prefix_sharing/__init__.py index 3c90c41e..cac96883 100644 --- a/prefix-sharing/prefix_sharing/__init__.py +++ b/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 diff --git a/prefix-sharing/prefix_sharing/core/planner.py b/prefix-sharing/prefix_sharing/core/planner.py index 7a872a72..5676f1f4 100644 --- a/prefix-sharing/prefix_sharing/core/planner.py +++ b/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 diff --git a/prefix-sharing/prefix_sharing/core/prefix_detector.py b/prefix-sharing/prefix_sharing/core/prefix_detector.py index 3d40a37d..50d05f14 100644 --- a/prefix-sharing/prefix_sharing/core/prefix_detector.py +++ b/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), diff --git a/prefix-sharing/prefix_sharing/integrations/__init__.py b/prefix-sharing/prefix_sharing/integrations/__init__.py index ac73372c..70a2c753 100644 --- a/prefix-sharing/prefix_sharing/integrations/__init__.py +++ b/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", diff --git a/prefix-sharing/prefix_sharing/integrations/megatron_attention.py b/prefix-sharing/prefix_sharing/integrations/megatron_attention.py deleted file mode 100644 index f32f340d..00000000 --- a/prefix-sharing/prefix_sharing/integrations/megatron_attention.py +++ /dev/null @@ -1,40 +0,0 @@ -"""Megatron attention integration entrypoint for phase 1.""" - -from __future__ import annotations - -import importlib -from dataclasses import dataclass -from typing import Any - -from prefix_sharing.core.config import PrefixSharingConfig -from prefix_sharing.integrations.patch_manager import PatchHandle, PatchManager - - -class IntegrationUnavailable(RuntimeError): - pass - - -@dataclass -class MegatronAttentionIntegration: - config: PrefixSharingConfig - backend: Any - - def install(self, model_config: Any | None = None) -> PatchHandle: - self.config.validate(model_config=model_config, integrate_mode="verl_megatron_actor") - attention_mod = self._import_attention_module() - self_attention_cls = getattr(attention_mod, "SelfAttention", None) - if self_attention_cls is None: - raise IntegrationUnavailable("Megatron SelfAttention class was not found") - - original_forward = getattr(self_attention_cls, "forward", None) - if original_forward is None: - raise IntegrationUnavailable("Megatron SelfAttention.forward was not found") - - return PatchManager().handle() - - @staticmethod - def _import_attention_module() -> Any: - try: - return importlib.import_module("megatron.core.transformer.attention") - except ModuleNotFoundError as exc: - raise IntegrationUnavailable("Megatron is not importable in this environment") from exc diff --git a/prefix-sharing/prefix_sharing/integrations/patch_manager.py b/prefix-sharing/prefix_sharing/integrations/patch_manager.py deleted file mode 100644 index 92d3bc39..00000000 --- a/prefix-sharing/prefix_sharing/integrations/patch_manager.py +++ /dev/null @@ -1,84 +0,0 @@ -"""Small monkey-patch manager with rollback and idempotent handles.""" - -from __future__ import annotations - -from dataclasses import dataclass -from types import TracebackType -from typing import Any, Callable - - -@dataclass(frozen=True) -class _PatchRecord: - target: Any - attr_name: str - original: Any - replacement: Any - - -class PatchHandle: - def __init__(self, records: list[_PatchRecord]) -> None: - self._records = records - self._active = True - - @property - def active(self) -> bool: - return self._active - - def disable(self) -> None: - if not self._active: - return - for record in reversed(self._records): - setattr(record.target, record.attr_name, record.original) - self._active = False - - def __enter__(self) -> "PatchHandle": - return self - - def __exit__( - self, - exc_type: type[BaseException] | None, - exc: BaseException | None, - tb: TracebackType | None, - ) -> None: - self.disable() - - -class PatchManager: - """Install attribute patches and roll back if any step fails.""" - - def __init__(self) -> None: - self._records: list[_PatchRecord] = [] - - def patch_attr( - self, - target: Any, - attr_name: str, - replacement: Any, - *, - signature_check: Callable[[Any], None] | None = None, - ) -> None: - if not hasattr(target, attr_name): - raise AttributeError(f"{target!r} has no attribute {attr_name!r}") - original = getattr(target, attr_name) - if original is replacement: - return - if signature_check is not None: - signature_check(original) - setattr(target, attr_name, replacement) - self._records.append( - _PatchRecord( - target=target, - attr_name=attr_name, - original=original, - replacement=replacement, - ) - ) - - def handle(self) -> PatchHandle: - records = list(self._records) - self._records.clear() - return PatchHandle(records) - - def rollback(self) -> None: - handle = self.handle() - handle.disable() diff --git a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py index f968d29e..cdbedee4 100644 --- a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py +++ b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py @@ -8,21 +8,12 @@ handle the monkey-patch integration via ``setup/patches/``. Both paths share the same core logic (plan -> trim -> layout -> state). - -``VerlMCoreBatchAdapter`` is framework-light and testable locally. It turns a -verl-style micro-batch payload into prefix-sharing metadata plus trimmed -inputs/labels/masks, and it assembles restored logprobs after forward. -``VerlMCoreIntegration`` installs the Megatron attention patch. The real -Megatron QKV rewiring still requires the framework runtime and remains guarded -by optional integration tests. """ from __future__ import annotations -import importlib -from contextlib import contextmanager from dataclasses import dataclass -from typing import Any, Iterator, Mapping +from typing import Any, Mapping from prefix_sharing.backends.factory import get_backend_instance from prefix_sharing.backends.packed_layout import PackedBatchLayout @@ -30,10 +21,8 @@ from prefix_sharing.core.planner import PrefixSharingPlan from prefix_sharing.core.planner import PrefixSharingPlanner from prefix_sharing.integrations.context import current_prefix_sharing_context -from prefix_sharing.integrations.megatron_attention import IntegrationUnavailable, MegatronAttentionIntegration from prefix_sharing.integrations.parallel_info import MegatronParallelInfo from prefix_sharing.integrations.parallel_info import get_megatron_parallel_info -from prefix_sharing.integrations.patch_manager import PatchHandle @dataclass(frozen=True) @@ -45,54 +34,6 @@ class PrefixSharingRuntimeState: kept_position_ids: Any | None = None -@dataclass -class VerlMCoreIntegration: - config: PrefixSharingConfig - backend: Any | None = None - - def install(self, model_config: Any | None = None) -> PatchHandle: - self.config.validate(model_config=model_config, integrate_mode="verl_megatron_actor") - self._ensure_verl_importable() - backend = get_backend_instance(self.config, self.backend) - return MegatronAttentionIntegration(config=self.config, backend=backend).install( - model_config=model_config - ) - - @staticmethod - def _ensure_verl_importable() -> None: - try: - importlib.import_module("verl") - except ModuleNotFoundError as exc: - raise IntegrationUnavailable("verl is not importable in this environment") from exc - - -def enable_prefix_sharing( - config: PrefixSharingConfig, - *, - model_config: Any | None = None, - backend: Any | None = None, -) -> PatchHandle: - """Install Phase 1 prefix-sharing patches for the verl + Megatron path.""" - - return VerlMCoreIntegration(config=config, backend=backend).install(model_config=model_config) - - -@contextmanager -def prefix_sharing_enabled( - config: PrefixSharingConfig, - *, - model_config: Any | None = None, - backend: Any | None = None, -) -> Iterator[PatchHandle]: - """Context manager wrapper around :func:`enable_prefix_sharing`.""" - - handle = enable_prefix_sharing(config, model_config=model_config, backend=backend) - try: - yield handle - finally: - handle.disable() - - def build_prefix_sharing_micro_batch_verl070( batch: Any, actor_config: Any, diff --git a/prefix-sharing/prefix_sharing/setup/logged_patch.py b/prefix-sharing/prefix_sharing/setup/logged_patch.py index e9bc01d4..2d984427 100644 --- a/prefix-sharing/prefix_sharing/setup/logged_patch.py +++ b/prefix-sharing/prefix_sharing/setup/logged_patch.py @@ -1,12 +1,10 @@ """带日志的 monkey-patch manager — setup 专用。 -功能与 integrations/patch_manager.py 相同，但增加： +提供： - patch_attr 时 INFO 日志 - disable 时 INFO 日志（逐条打印恢复详情） - describe() 方法返回人类可读 patch 清单（含已应用和待挂起状态） - inspect_patch() 方法返回被替换函数的源码，供用户验证 - -此文件独立于 integrations/patch_manager.py，不影响原有代码。 """ from __future__ import annotations diff --git a/prefix-sharing/tests/integrated_test/test_patch_integrations.py b/prefix-sharing/tests/integrated_test/test_patch_integrations.py deleted file mode 100644 index b9b4ba62..00000000 --- a/prefix-sharing/tests/integrated_test/test_patch_integrations.py +++ /dev/null @@ -1,123 +0,0 @@ -import pytest - -from prefix_sharing.backends.torch_ref import TorchReferenceBackend -from prefix_sharing.core.config import PrefixSharingConfig -from prefix_sharing.integrations.megatron_attention import ( - IntegrationUnavailable, - MegatronAttentionIntegration, -) -from prefix_sharing.integrations.patch_manager import PatchManager -from prefix_sharing.integrations.verl_mcore import ( - VerlMCoreIntegration, - prefix_sharing_enabled, -) - - -class Target: - def method(self): - return "original" - - -def test_patch_manager_installs_and_disables_patch(): - target = Target() - manager = PatchManager() - - def replacement(instance): - return "patched" - - manager.patch_attr(Target, "method", replacement) - handle = manager.handle() - assert target.method() == "patched" - assert handle.active - - handle.disable() - assert target.method() == "original" - assert not handle.active - - -def test_patch_manager_context_manager_restores_original(): - target = Target() - manager = PatchManager() - manager.patch_attr(Target, "method", lambda instance: "patched") - - with manager.handle(): - assert target.method() == "patched" - assert target.method() == "original" - - -def test_megatron_integration_reports_missing_dependency_cleanly(monkeypatch): - import importlib - - _original_import = importlib.import_module - - def _mock_import(name, package=None): - if name == "megatron.core.transformer.attention": - raise ModuleNotFoundError("No module named 'megatron'") - return _original_import(name, package=package) - - monkeypatch.setattr(importlib, "import_module", _mock_import) - - config = PrefixSharingConfig(enable_prefix_sharing=True) - integration = MegatronAttentionIntegration(config=config, backend=TorchReferenceBackend()) - with pytest.raises(IntegrationUnavailable, match="Megatron"): - integration.install(model_config={}) - - -def test_verl_integration_reports_missing_dependency_cleanly(monkeypatch): - import importlib - - _original_import = importlib.import_module - - def _mock_import(name, package=None): - if name == "verl": - raise ModuleNotFoundError("No module named 'verl'") - return _original_import(name, package=package) - - monkeypatch.setattr(importlib, "import_module", _mock_import) - - config = PrefixSharingConfig(enable_prefix_sharing=True) - integration = VerlMCoreIntegration(config=config) - with pytest.raises(IntegrationUnavailable, match="verl"): - integration.install(model_config={}) - - -def test_prefix_sharing_config_from_raw_accepts_nested_config(): - config = PrefixSharingConfig.from_raw( - { - "enable_prefix_sharing": True, - "min_prefix_len": 4, - "min_group_size": 3, - "boundary_strategy": "prefix_last_restore", - } - ) - - assert config.enable_prefix_sharing is True - assert config.min_prefix_len == 4 - assert config.min_group_size == 3 - - -def test_prefix_sharing_config_from_raw_rejects_legacy_enabled_key(): - with pytest.raises(TypeError, match="enabled"): - PrefixSharingConfig.from_raw({"enabled": True}) - - -def test_prefix_sharing_config_from_raw_accepts_env_enable(monkeypatch): - monkeypatch.setenv("ENABLE_PREFIX_SHARING", "1") - - config = PrefixSharingConfig.from_raw(None) - - assert config.enable_prefix_sharing is True - - -def test_prefix_sharing_enabled_propagates_install_failure(monkeypatch): - class FakeIntegration: - def __init__(self, config, backend=None): - pass - - def install(self, model_config=None): - raise RuntimeError("install failed") - - monkeypatch.setattr("prefix_sharing.integrations.verl_mcore.VerlMCoreIntegration", FakeIntegration) - with pytest.raises(RuntimeError, match="install failed"): - with prefix_sharing_enabled(PrefixSharingConfig(enable_prefix_sharing=True)): - pass diff --git a/prefix-sharing/tests/unit_test/test_config.py b/prefix-sharing/tests/unit_test/test_config.py index 511fd211..bef4b283 100644 --- a/prefix-sharing/tests/unit_test/test_config.py +++ b/prefix-sharing/tests/unit_test/test_config.py @@ -106,3 +106,23 @@ def test_enabled_config_rejects_non_phase_one_modes(): PrefixSharingConfig(enable_prefix_sharing=True, boundary_strategy="restore_last_prefix_token").validate(ModelConfig()) with pytest.raises(PrefixSharingConfigError, match="integrate_mode"): PrefixSharingConfig(enable_prefix_sharing=True).validate(ModelConfig(), integrate_mode="verl_fsdp") + + +def test_prefix_sharing_config_from_raw_accepts_nested_config(): + config = PrefixSharingConfig.from_raw( + { + "enable_prefix_sharing": True, + "min_prefix_len": 4, + "min_group_size": 3, + "boundary_strategy": "prefix_last_restore", + } + ) + + assert config.enable_prefix_sharing is True + assert config.min_prefix_len == 4 + assert config.min_group_size == 3 + + +def test_prefix_sharing_config_from_raw_rejects_legacy_enabled_key(): + with pytest.raises(TypeError, match="enabled"): + PrefixSharingConfig.from_raw({"enabled": True}) diff --git a/prefix-sharing/tests/unit_test/test_detector.py b/prefix-sharing/tests/unit_test/test_detector.py index 4535ec02..b935d5b9 100644 --- a/prefix-sharing/tests/unit_test/test_detector.py +++ b/prefix-sharing/tests/unit_test/test_detector.py @@ -57,13 +57,13 @@ def test_trie_detector_builds_per_sample_reuse_relations(): assert result.provider_index == (0, 0, 0, 3, 3) assert result.prefix_lens == (0, 3, 5, 0, 2) assert result.is_provider == (True, False, False, True, False) - assert len(result.groups) == 3 - assert result.groups[0].member_indices == (0, 1) - assert result.groups[0].prefix_len == 3 - assert result.groups[1].member_indices == (0, 2) - assert result.groups[1].prefix_len == 5 - assert result.groups[2].member_indices == (3, 4) - assert result.groups[2].prefix_len == 2 + assert len(result.prefix_groups) == 3 + assert result.prefix_groups[0].member_indices == (0, 1) + assert result.prefix_groups[0].prefix_len == 3 + assert result.prefix_groups[1].member_indices == (0, 2) + assert result.prefix_groups[1].prefix_len == 5 + assert result.prefix_groups[2].member_indices == (3, 4) + assert result.prefix_groups[2].prefix_len == 2 def test_trie_detector_allows_reuser_to_provide_longer_prefix_later(): @@ -99,3 +99,31 @@ def test_trie_detector_respects_min_group_size_for_relation_threshold(): ] assert result.provider_index == (0, 1, 0) assert result.prefix_lens == (0, 0, 2) + + +def test_detect_handles_prefix_mismatch_then_inserts_suffix(): + """单次遍历实现必须确保：匹配阶段断裂后，后续 token 仍被完整插入 trie。 + + 场景： + - seq0: [1,2,3,4] (provider) + - seq1: [1,2,5,6] (匹配 [1,2] 后断裂，matched_prefix_len=2; [5,6] 仍需插入) + - seq2: [1,2,5,7] (应匹配到 seq1 的 [1,2,5]，prefix_len=3) + """ + # min_group_size=2 下，seq1 会 reuse seq0（prefix=2），seq2 会 reuse seq1（prefix=3） + # 关键验证点：seq2 能够匹配到 seq1 插入的 [1,2,5] 路径，证明"中途断裂后插入"成功 + detector = TriePrefixDetector(min_prefix_len=2, min_group_size=2) + result = detector.detect( + [ + [1, 2, 3, 4], # 0: provider + [1, 2, 5, 6], # 1: 匹配到 [1,2] 后断裂，matched_prefix_len=2 + [1, 2, 5, 7], # 2: 应 reuse seq1 的 [1,2,5]，prefix_len=3 + ] + ) + + # seq1 reuse seq0, seq2 reuse seq1 + assert [(s.reuse_idx_in_batch, s.provider_idx_in_batch, s.prefix_len) for s in result.reuse_specs] == [ + (1, 0, 2), + (2, 1, 3), + ] + assert result.provider_index == (0, 0, 1) + assert result.prefix_lens == (0, 2, 3)