diff --git a/docs/developer-docs/impr-refactor.md b/docs/developer-docs/impr-refactor.md
index 619f3bd7..10d31146 100644
--- a/docs/developer-docs/impr-refactor.md
+++ b/docs/developer-docs/impr-refactor.md
@@ -270,11 +270,11 @@ integrations/verl_fsdp.py         # FSDP 专属流程
 - `build_prefix_sharing_micro_batch_fsdp()`
 - `PrefixSharingFSDPAttentionRuntime`
 - `restore_prefix_sharing_outputs_2d()`
-- `forward_prefix_sharing_fsdp_micro_batch()`
+- `forward_step_without_engine_prepare()`
 
 问题：
 
-- `forward_prefix_sharing_fsdp_micro_batch()` 更像测试/fake engine helper，不一定是合入 verl 的主路径，应避免让 reviewer 误以为这是生产接入方式。
+- `forward_step_without_engine_prepare()` 更像测试/fake engine helper，不一定是合入 verl 的主路径，应避免让 reviewer 误以为这是生产接入方式。
 - `PrefixSharingFSDPAttentionRuntime.forward()` 直接忽略 `attn_func/attention_mask/kwargs`，对 HF attention 接口兼容性说明不足。
 - dense `[B,L,H,D]` 与 packed `[1,T,H,D]` 两种路径混在同一个 runtime，缺少清晰的 input contract。
 - FSDP restore 的 logits/log_probs/entropy/attention_output copy 语义复杂，需要更强测试和更小函数。
@@ -2280,25 +2280,32 @@ PYTHONPATH=prefix-sharing pytest -q \
 - 需要固定数据和随机种子，形成可复现的精度对齐脚本。
 - optional GPU/NPU 环境测试不能作为每个本地 PR 的硬门槛，但 release 前必须复跑。
 
-### 6.4 Patch 与 import hook
+### 6.4 FSDP activation-checkpointing runtime context 生命周期
+
+- `origin/open-source` 已存在 FSDP PrefixSharing runtime context 生命周期缺口：`forward_only=True` 没有 backward 触发清理；普通训练在未开启诊断 dump 时也可能遗留 module 上的 prefix-sharing context。该问题起源于 activation-checkpointing 兼容改造，不是当前第三波重构引入。
+- 当前第三波重构以 `origin/open-source` 为行为基线闭环，不在本 PR 内修复此既有缺口；但所有新增或修改的测试不得掩盖它，也不得将其误归因为重构回归。
+- 后续应以独立修复 PR 处理：统一 ContextVar、store 与 attention module binding 的所有权和幂等 cleanup；覆盖普通训练、`forward_only`、activation-checkpointing recompute、异常路径及连续 micro-batch。
+- 在该独立修复完成前，FSDP 精度验证应固定 replay fixture，并明确记录所覆盖的执行模式；不得把单 micro-batch 的 ON/OFF 对齐外推为上述生命周期场景全部已验证。
+
+### 6.5 Patch 与 import hook
 
 - import hook 是否长期保留，需要等社区对 monkey patch 方式的反馈后再定。
 - 如果正式合入 verl，显式调用路径可能替代外部包 import auto patch。
 - 当前阶段必须保留 import auto patch，因为脚本化训练仍依赖 `VERL_USE_EXTERNAL_MODULES=prefix_sharing`。
 
-### 6.5 HybridAttention / Gated DeltaNet
+### 6.6 HybridAttention / Gated DeltaNet
 
 - 当前重构清理 Qwen3.5/GDN 专门化代码，不代表永久放弃 HybridAttention。
 - 后续需等待训练引擎侧真实接口稳定，再重新设计 activation/cache_param store。
 - 未来重新引入时，应以实际 mixer 类型命名，避免把 DeltaNet 泛化成所有 linear attention。
 
-### 6.6 Megatron / MCore / NPU
+### 6.7 Megatron / MCore / NPU
 
 - MCore path 保留 advanced/internal 定位。
 - NPU/MindSpeed/Megatron-Bridge 后续可继续支持，但不能阻塞 FSDP-first 开源主线。
 - 若后续重新提高 Megatron 优先级，需要单独补兼容矩阵、真实环境测试和文档。
 
-### 6.7 tools 与 diagnostics
+### 6.8 tools 与 diagnostics
 
 - tools 清理需要逐项判断，不能批量删除。
 - 保留工具必须补用途说明。
diff --git a/prefix-sharing/prefix_sharing/backends/packed_layout.py b/prefix-sharing/prefix_sharing/backends/packed_layout.py
index d30809cc..5eebf900 100644
--- a/prefix-sharing/prefix_sharing/backends/packed_layout.py
+++ b/prefix-sharing/prefix_sharing/backends/packed_layout.py
@@ -57,7 +57,7 @@ def from_kept_position_rows(
         cls,
         kept_position_rows: Sequence[Any],
         *,
-        align_size: int,
+        align_size: int = 1,
     ) -> "PackedBatchLayout":
         if align_size < 1:
             raise ValueError("align_size must be >= 1")
@@ -65,18 +65,19 @@ def from_kept_position_rows(
             return cls.from_valid_lengths([])
 
         first = kept_position_rows[0]
+        device = first.device
         valid_lengths = [int(row.shape[0]) for row in kept_position_rows]
         padded_lengths = [_pad_to_multiple(length, align_size) for length in valid_lengths]
         packed_position_rows = []
         valid_mask_rows = []
         for row, valid_length, padded_length in zip(kept_position_rows, valid_lengths, padded_lengths):
-            row = row.to(first.device)
+            row = row.to(device)
             pad_length = padded_length - valid_length
             valid_mask_rows.append(
                 torch.cat(
                     [
-                        torch.ones(valid_length, dtype=torch.bool, device=first.device),
-                        torch.zeros(pad_length, dtype=torch.bool, device=first.device),
+                        torch.ones(valid_length, dtype=torch.bool, device=device),
+                        torch.zeros(pad_length, dtype=torch.bool, device=device),
                     ],
                     dim=0,
                 )
@@ -84,7 +85,7 @@ def from_kept_position_rows(
             if pad_length == 0:
                 packed_position_rows.append(row)
                 continue
-            padding = torch.zeros(pad_length, dtype=row.dtype, device=first.device)
+            padding = torch.zeros(pad_length, dtype=row.dtype, device=device)
             packed_position_rows.append(torch.cat([row, padding], dim=0))
 
         return cls(
diff --git a/prefix-sharing/prefix_sharing/core/config.py b/prefix-sharing/prefix_sharing/core/config.py
index 35dc60a4..044bb25d 100644
--- a/prefix-sharing/prefix_sharing/core/config.py
+++ b/prefix-sharing/prefix_sharing/core/config.py
@@ -103,118 +103,117 @@ def from_raw(cls, raw: Any) -> "PrefixSharingConfig":
         return cls(**values)
 
     def validate(self, model_config: Any | None = None, integrate_mode: str | None = None) -> None:
-        """Validate phase-1 constraints against a model/config object.
+        """Validate constraints against model_config.
 
         Args:
-            model_config: Mapping or object with Megatron-like attributes.
+            model_config: Model attributes.
             integrate_mode: Optional integration mode name.
         """
 
         if not self.enable_prefix_sharing:
             return
-        if self.detector != "trie":
-            raise PrefixSharingConfigError("phase 1 supports only detector='trie'")
+
+        # detector guard
+        supported_detectors = {"trie"}
+        if self.detector not in supported_detectors:
+            raise PrefixSharingConfigError(
+                f"detector='{self.detector}' is not supported. Supported detectors: {supported_detectors}"
+            )
+
+        # backend guard
         supported_backends = {"torch_ref", "flash_atten_gpu", "flash_atten_npu"}
         if self.backend not in supported_backends:
             raise PrefixSharingConfigError(
-                f"backend='{self.backend}' is not supported. "
-                f"Supported backends: {supported_backends}"
+                f"backend='{self.backend}' is not supported. Supported backends: {supported_backends}"
             )
-        if self.boundary_strategy != "prefix_last_restore":
+
+        # boundary_strategy guard
+        supported_boundary_strategies = {"prefix_last_restore"}
+        if self.boundary_strategy not in supported_boundary_strategies:
             raise PrefixSharingConfigError(
-                "phase 1 currently implements only "
-                "boundary_strategy='prefix_last_restore'; future strategies may include "
-                "'boundary_token' and 'strict_suffix'"
+                f"boundary_strategy='{self.boundary_strategy}' is not supported. Supported strategies: {supported_boundary_strategies}"
             )
+
+        # min_prefix_len guard
         if self.min_prefix_len < 1:
-            raise PrefixSharingConfigError("min_prefix_len must be >= 1")
+            raise PrefixSharingConfigError(f"min_prefix_len='{self.min_prefix_len}' is not supported. Supported min_prefix_len: >= 1")
         if self.min_group_size < 2:
-            raise PrefixSharingConfigError("min_group_size must be >= 2")
+            raise PrefixSharingConfigError(f"min_group_size='{self.min_group_size}' is not supported. Supported min_group_size: >= 2")
 
-        active_mode = integrate_mode or self.integrate_mode
+        # integrate_mode guard
         supported_integrate_modes = {"verl_megatron_actor", "verl_fsdp"}
+        active_mode = integrate_mode or self.integrate_mode
         if active_mode not in supported_integrate_modes:
             raise PrefixSharingConfigError(
-                "phase 1 supports only integrate_mode in "
-                f"{sorted(supported_integrate_modes)}"
+                f"integrate_mode='{active_mode}' is not supported. Supported integrate_modes: {supported_integrate_modes}"
             )
 
+        # model_type guard
+        supported_model_types = {"text_only_causal_lm"}
         model_type = _read_config_value(model_config, "model_type", "text_only_causal_lm")
-        if self.model_type == "text_only_causal_lm" and model_type != "text_only_causal_lm":
+        if self.model_type in supported_model_types and model_type not in supported_model_types:
             raise PrefixSharingConfigError(
-                f"[Config Error] Current model_type '{model_type}' is not supported in this phase. "
-                f"Phase 1 only supports model_type='text_only_causal_lm' (text-only causal language model). "
-                f"Please use a supported model type or disable prefix sharing."
+                f"model_type='{model_type}' is not supported. Supported model_types: {supported_model_types}"
             )
+
+        # verl_fsdp guard
         if active_mode == "verl_fsdp":
             ulysses_sp_size = _read_config_value(model_config, "ulysses_sequence_parallel_size", 1)
             use_fused_kernels = _read_config_value(model_config, "use_fused_kernels", False)
             if int(ulysses_sp_size) != 1:
                 raise PrefixSharingConfigError(
-                    f"[Config Error] verl_fsdp does not currently support ulysses_sequence_parallel_size={ulysses_sp_size}. "
-                    "Please disable Ulysses SP or wait for a dedicated adaptation."
+                    f"ulysses_sequence_parallel_size={ulysses_sp_size} is not supported. "
+                    "Supported ulysses_sequence_parallel_size: 1"
                 )
             if use_fused_kernels:
                 raise PrefixSharingConfigError(
-                    "[Config Error] verl_fsdp does not currently support use_fused_kernels=True. "
-                    "Please disable fused kernels or wait for a dedicated adaptation."
+                    f"use_fused_kernels={use_fused_kernels} is not supported. "
+                    "Supported use_fused_kernels: False"
                 )
             return
 
+        # parallel strategy guards
         pp_size = _read_config_value(
             model_config,
             "pipeline_model_parallel_size",
             _read_config_value(model_config, "pipeline_parallel_size", 1),
         )
-        virtual_pp_size = _read_config_value(model_config, "virtual_pipeline_model_parallel_size", None)
-        num_layers_per_virtual_pipeline_stage = _read_config_value(
-            model_config,
-            "num_layers_per_virtual_pipeline_stage",
-            None,
-        )
-        cp_size = _read_config_value(
-            model_config,
-            "context_parallel_size",
-            self.supported_cp_size,
-        )
-        rope_fusion = _read_config_value(model_config, "apply_rope_fusion", False)
-        fused_qkv_rope = _read_config_value(model_config, "fused_single_qkv_rope", False)
-
         if int(pp_size) < 1:
             raise PrefixSharingConfigError(
-                f"[Config Error] pipeline_model_parallel_size={pp_size} is invalid. "
-                f"pipeline_model_parallel_size must be >= 1. "
-                f"Please set a valid physical PP size or disable prefix sharing."
+                f"pipeline_model_parallel_size={pp_size} is not supported. "
+                "Supported pipeline_model_parallel_size: >= 1"
             )
+        virtual_pp_size = _read_config_value(model_config, "virtual_pipeline_model_parallel_size", 1)
         if virtual_pp_size not in (None, 1):
             raise PrefixSharingConfigError(
-                f"[Config Error] virtual_pipeline_model_parallel_size={virtual_pp_size} is not supported in this phase. "
-                f"Only physical pipeline parallel is supported; virtual pipeline parallel is not. "
-                f"Please disable virtual PP or disable prefix sharing."
+                f"virtual_pipeline_model_parallel_size={virtual_pp_size} is not supported. "
+                "Supported virtual_pipeline_model_parallel_size: 1"
             )
+        num_layers_per_virtual_pipeline_stage = _read_config_value(
+            model_config, "num_layers_per_virtual_pipeline_stage", None
+        )
         if num_layers_per_virtual_pipeline_stage is not None:
             raise PrefixSharingConfigError(
-                "[Config Error] num_layers_per_virtual_pipeline_stage is not supported in this phase. "
-                "Only physical pipeline parallel is supported; virtual pipeline parallel is not. "
-                "Please disable virtual PP or disable prefix sharing."
+                f"num_layers_per_virtual_pipeline_stage={num_layers_per_virtual_pipeline_stage} "
+                "is not supported. Supported num_layers_per_virtual_pipeline_stage: None"
             )
+        cp_size = _read_config_value(model_config, "context_parallel_size", self.supported_cp_size)
         if cp_size != self.supported_cp_size:
             raise PrefixSharingConfigError(
-                f"[Config Error] context_parallel_size={cp_size} is not supported in this phase. "
-                f"Phase 1 only supports context_parallel_size=1 (no context parallelism). "
-                f"Please set CP size to 1 or disable prefix sharing."
+                f"context_parallel_size={cp_size} is not supported. "
+                f"Supported context_parallel_size: {self.supported_cp_size}"
             )
+        rope_fusion = _read_config_value(model_config, "apply_rope_fusion", False)
         if not self.supported_rope_fusion and rope_fusion:
             raise PrefixSharingConfigError(
-                "[Config Error] apply_rope_fusion=True is not supported in this phase. "
-                "Phase 1 requires rope fusion to be disabled (apply_rope_fusion=False). "
-                "Please update the configuration or disable prefix sharing."
+                f"apply_rope_fusion={rope_fusion} is not supported. "
+                "Supported apply_rope_fusion: False"
             )
+        fused_qkv_rope = _read_config_value(model_config, "fused_single_qkv_rope", False)
         if not self.supported_fused_qkv_rope and fused_qkv_rope:
             raise PrefixSharingConfigError(
-                "[Config Error] fused_single_qkv_rope=True is not supported in this phase. "
-                "Phase 1 requires fused QKV rope to be disabled (fused_single_qkv_rope=False). "
-                "Please update the configuration or disable prefix sharing."
+                f"fused_single_qkv_rope={fused_qkv_rope} is not supported. "
+                "Supported fused_single_qkv_rope: False"
             )
 
     def validate_for_engine(
@@ -222,17 +221,16 @@ def validate_for_engine(
         use_remove_padding: bool = True,
         integrate_mode: str = "verl_megatron_actor",
     ) -> None:
-        """Validate phase-1 constraints for verl engine architecture (verl 0.8.0+).
+        """Validate phase-1 constraints for the verl engine (verl 0.8.0+).
 
-        Unlike validate(), this method reads from engine_config rather than
-        model_config. Used in setup/patches for forward_step patching, where
-        only engine_config (self.engine_config) is available, not the Megatron
-        TransformerConfig.
+        Unlike ``validate()``, this reads constraints from engine-facing flags
+        rather than Megatron ``model_config``. Used by setup/patches
+        ``forward_step`` when only ``self.engine_config`` is available.
         """
         if not self.enable_prefix_sharing:
             return
 
-        # Basic validation
+        # Basic field checks.
         if self.detector != "trie":
             raise PrefixSharingConfigError("phase 1 supports only detector='trie'")
         if self.backend not in {"torch_ref", "flash_atten_gpu", "flash_atten_npu"}:
@@ -249,10 +247,10 @@ def validate_for_engine(
         if self.min_group_size < 2:
             raise PrefixSharingConfigError("min_group_size must be >= 2")
 
-        # THD packed layout requires use_remove_padding
+        # THD packed layout requires use_remove_padding.
         if not use_remove_padding:
             raise PrefixSharingConfigError(
                 "[Config Error] Phase 1 THD path requires use_remove_padding=True. "
-                "The BSHD path (use_remove_padding=False) is not yet supported in the current patch set. "
-                "Please enable use_remove_padding or use the BSHD-specific patch set."
+                "BSHD (use_remove_padding=False) is not supported by the current "
+                "patch set. Enable use_remove_padding or use a BSHD-specific patch set."
             )
diff --git a/prefix-sharing/prefix_sharing/integrations/__init__.py b/prefix-sharing/prefix_sharing/integrations/__init__.py
index 2b03ff05..4ab50dc1 100644
--- a/prefix-sharing/prefix_sharing/integrations/__init__.py
+++ b/prefix-sharing/prefix_sharing/integrations/__init__.py
@@ -9,15 +9,15 @@
 from prefix_sharing.integrations.parallel_info import MegatronParallelInfo, get_megatron_parallel_info
 from prefix_sharing.backends.packed_layout import PackedBatchLayout
 from prefix_sharing.integrations.runtime_state import PrefixSharingRuntimeState
-from prefix_sharing.integrations.verl_utils import read_ps_config_from_engine_config
-from prefix_sharing.integrations.verl_mcore import (
-    build_prefix_sharing_micro_batch_verl080,
+from prefix_sharing.integrations.verl_utils import (
+    read_ps_config_from_engine_config,
     restore_reuser_prefix_columns_2d,
 )
+from prefix_sharing.integrations.verl_mcore import build_prefix_sharing_micro_batch_verl080
 from prefix_sharing.integrations.verl_fsdp import (
     PrefixSharingFSDPAttentionRuntime,
-    build_prefix_sharing_micro_batch_fsdp,
-    forward_prefix_sharing_fsdp_micro_batch,
+    prepare_for_prefix_sharing_fsdp,
+    forward_step_without_engine_prepare,
     restore_prefix_sharing_outputs_2d,
 )
 from prefix_sharing.integrations.megatron_runtime import (
@@ -38,7 +38,7 @@
     "prefix_attention",
     "get_megatron_parallel_info",
     "PrefixSharingFSDPAttentionRuntime",
-    "build_prefix_sharing_micro_batch_fsdp",
-    "forward_prefix_sharing_fsdp_micro_batch",
+    "prepare_for_prefix_sharing_fsdp",
+    "forward_step_without_engine_prepare",
     "restore_prefix_sharing_outputs_2d",
 ]
diff --git a/prefix-sharing/prefix_sharing/integrations/context.py b/prefix-sharing/prefix_sharing/integrations/context.py
index b28498dc..b25f7937 100644
--- a/prefix-sharing/prefix_sharing/integrations/context.py
+++ b/prefix-sharing/prefix_sharing/integrations/context.py
@@ -151,7 +151,7 @@ def create_prefix_sharing_context(
     the ContextVar is left set until the caller invokes the returned cleanup
     function.  This is required for activation‑checkpointing compatibility: AC
     recompute runs inside ``backward()`` and reads the PS context from
-    ``module._ps_ctx`` (set independently by the caller), while the store must
+    ``module._prefix_sharing_context`` (set independently by the caller), while the store must
     still contain the per‑layer KV populated during the first forward.
 
     Returns:
@@ -164,12 +164,12 @@ def create_prefix_sharing_context(
     ctx = PrefixSharingRuntimeContext(prefix_sharing_runtime_state, store)
     ctxvar_token = _current_context.set(ctx)
 
-    def cleanup() -> None:
+    def cleanup_context() -> None:
         _current_context.reset(ctxvar_token)
         _log_prefix_sharing_audit(ctx)
         ctx.store.close()
 
-    return ctx, cleanup
+    return ctx, cleanup_context
 
 
 def _log_prefix_sharing_audit(ctx: PrefixSharingRuntimeContext) -> None:
diff --git a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py
index fb5b83aa..77fdc385 100644
--- a/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py
+++ b/prefix-sharing/prefix_sharing/integrations/verl_fsdp.py
@@ -1,7 +1,9 @@
-"""verl FSDP integration helpers for PrefixSharing.
+"""prefix_sharing.integrations.verl_fsdp
+
+verl FSDP integration helpers for PrefixSharing.
 
 The FSDP path follows the same public shape as the Megatron integration:
-``build_*`` returns ``(trimmed_micro_batch, PrefixSharingRuntimeState | None)``.
+``prepare_for_prefix_sharing_fsdp`` returns ``(trimmed_micro_batch, PrefixSharingRuntimeState | None)``.
 The helpers stay framework-light enough for CPU tests, while the explicit
 ``verl080_fsdp`` patch set wires them into ``FSDPEngineWithLMHead.forward_step``.
 """
@@ -20,12 +22,12 @@
 from prefix_sharing.integrations.context import prefix_sharing_runtime_context
 from prefix_sharing.integrations.parallel_info import MegatronParallelInfo
 from prefix_sharing.integrations.runtime_state import PrefixSharingRuntimeState
-from prefix_sharing.integrations.verl_utils import _clone_batch
-from prefix_sharing.integrations.verl_utils import _collect_kept_position_rows
-from prefix_sharing.integrations.verl_utils import _extract_seq_from_nested_tensor
-from prefix_sharing.integrations.verl_utils import _extract_sequences_async
-from prefix_sharing.integrations.verl_utils import _is_nested_tensor
-from prefix_sharing.integrations.verl_utils import _trim_nested_batch
+from prefix_sharing.integrations.verl_utils import clone_batch
+from prefix_sharing.integrations.verl_utils import extract_seq_from_dense_tensor
+from prefix_sharing.integrations.verl_utils import extract_seq_from_nested_tensor
+from prefix_sharing.integrations.verl_utils import is_nested_tensor
+from prefix_sharing.integrations.verl_utils import trim_redundant_prefix_in_nested_tensor
+from prefix_sharing.integrations.verl_utils import trim_redundant_prefix_in_dense_tensor
 
 class PrefixSharingFSDPAttentionRuntime:
     """Standalone FSDP attention runtime for PrefixSharing.
@@ -52,6 +54,11 @@ def forward(self, attn_func: Any, query: Any, key: Any, value: Any, *args: Any,
         from prefix_sharing.tools.perf_profiler import PerfProfiler
         profiler = PerfProfiler.current()
         # [PS-perf] end ——————————————————————————————————————
+
+        #########################################################
+        # handle THD format input
+        #########################################################
+
         if query.shape[0] == 1 and key.shape[0] == 1 and value.shape[0] == 1:
             packed_query = query.squeeze(0)
             packed_key = key.squeeze(0)
@@ -68,6 +75,10 @@ def forward(self, attn_func: Any, query: Any, key: Any, value: Any, *args: Any,
         if query.shape[:2] != key.shape[:2] or query.shape[:2] != value.shape[:2]:
             raise RuntimeError("query, key, and value must share dense batch/sequence dimensions")
 
+        #########################################################
+        # handle BSHD format input
+        #########################################################
+
         plan = ctx.prefix_sharing_plan
         # [PS-perf] start — attn.pack ———————————————————————
         _per_layer_ok = profiler is not None and getattr(profiler, "per_layer_enabled", False)
@@ -111,10 +122,10 @@ def forward(self, attn_func: Any, query: Any, key: Any, value: Any, *args: Any,
         return dense_output
 
 
-def forward_prefix_sharing_fsdp_micro_batch(
+def forward_step_without_engine_prepare(
     micro_batch: Any,
     model: Any,
-    config: PrefixSharingConfig,
+    ps_config: PrefixSharingConfig,
     *,
     model_config: Any | None = None,
     backend: Any | None = None,
@@ -124,39 +135,44 @@ def forward_prefix_sharing_fsdp_micro_batch(
     entropy_fn: Any | None = None,
     autocast_context: Any | None = None,
 ) -> dict[str, Any]:
-    """Run one dense verl/FSDP-style micro-batch with PrefixSharing.
-
-    This is the executable helper used by fake/local FSDP tests and by engines
-    that do not expose prepare hooks:
-    prepare the micro-batch, open the runtime context, run the model with a
-    PrefixSharing attention runtime, compute token-level outputs, then restore
-    reuser prefix columns. Real verl FSDP patching should prefer the engine's
-    own ``prepare_model_inputs`` / ``prepare_model_outputs`` path.
-    """
+    """Run one forward pass of a verl/FSDP-style micro-batch with PrefixSharing."""
 
-    trimmed_micro_batch, runtime_state = build_prefix_sharing_micro_batch_fsdp(
+    #########################################################
+    # STEP 1: pre-processing inputs for PrefixSharing
+    #########################################################
+    micro_batch_modified, prefix_sharing_runtime_state = prepare_for_prefix_sharing_fsdp(
         micro_batch,
-        config,
+        ps_config,
         model_config=model_config,
         backend=backend,
     )
-    context = prefix_sharing_runtime_context(runtime_state) if runtime_state is not None else nullcontext(None)
-    autocast = autocast_context if autocast_context is not None else nullcontext()
 
+    #########################################################
+    # STEP 2: create PrefixSharing runtime context
+    #########################################################
+    context = prefix_sharing_runtime_context(prefix_sharing_runtime_state) if prefix_sharing_runtime_state is not None else nullcontext(None)
+    autocast = autocast_context if autocast_context is not None else nullcontext()
     with context as ctx, autocast:
+
+        #########################################################
+        # STEP 3: call the model with PrefixSharing+FSDP attention runtime
+        #########################################################
         model_output = _call_fsdp_model(
             model,
-            trimmed_micro_batch,
-            prefix_sharing_runtime=PrefixSharingFSDPAttentionRuntime(
+            micro_batch_modified,
+            attention_runtime=PrefixSharingFSDPAttentionRuntime(
                 num_layers=model.config.num_hidden_layers if hasattr(model, "config") else 0,
             ),
-            enable_prefix_sharing=runtime_state is not None,
+            enable_prefix_sharing=prefix_sharing_runtime_state is not None,
         )
+
+        #########################################################
+        # STEP 4: post-processing outputs for PrefixSharing
+        #########################################################
+        output = {}
         logits = _extract_logits(model_output) / float(temperature)
-        output = {
-            "model_output": model_output,
-            "logits": logits.clone(),
-        }
+        output["model_output"] = model_output
+        output["logits"] = logits.clone()
         labels = _labels_for_log_probs(micro_batch)
         if labels is not None:
             output["log_probs"] = _compute_log_probs(logits, labels, log_probs_fn)
@@ -173,92 +189,84 @@ def forward_prefix_sharing_fsdp_micro_batch(
         return output
 
 
-def build_prefix_sharing_micro_batch_fsdp(
-    batch: Any,
-    config: PrefixSharingConfig,
-    *,
+def prepare_for_prefix_sharing_fsdp(
+    micro_batch: Any,
+    ps_config: PrefixSharingConfig,
     model_config: Any | None = None,
     backend: Any | None = None,
 ) -> tuple[Any, PrefixSharingRuntimeState | None]:
-    """Build a trimmed FSDP micro-batch and PrefixSharing runtime state.
+    """Plan prefix sharing and trim one FSDP micro-batch.
 
+    Returns ``(trimmed_micro_batch, PrefixSharingRuntimeState | None)``.
     This helper is intentionally framework-light: it accepts dense 2D
     ``input_ids``/``attention_mask`` or jagged NestedTensor ``input_ids`` from
-    verl remove-padding, and returns the original batch unchanged when prefix
-    sharing is disabled or no reusable prefix is detected.
+    verl remove-padding, and returns the original micro-batch unchanged when
+    prefix sharing is disabled or no reusable prefix is detected.
     """
 
-    if not config.enable_prefix_sharing:
-        return batch, None
-    config.validate(model_config=model_config, integrate_mode="verl_fsdp")
+    # skip if PrefixSharing is disabled
+    if not ps_config.enable_prefix_sharing:
+        return micro_batch, None
+
+    # validate the config
+    ps_config.validate(model_config=model_config, integrate_mode="verl_fsdp")
 
-    input_ids = batch["input_ids"]
-    is_nested_input = _is_nested_tensor(input_ids)
-    if is_nested_input:
-        sequences = _extract_seq_from_nested_tensor(input_ids)
+    # extract input_ids from micro-batch and convert it into sequences (Python List)
+    # for subsequent PrefixSharing planning, which is purely done on CPU
+    input_ids = micro_batch["input_ids"]
+    input_ids_is_nested = is_nested_tensor(input_ids)
+    if input_ids_is_nested: # nested tensor
+        sequences = extract_seq_from_nested_tensor(input_ids)
         valid_indices = None
         attention_mask = None
-    else:
-        attention_mask = batch["attention_mask"].to(bool)
+    else: # dense tensor
+        attention_mask = micro_batch["attention_mask"].to(bool)
         if input_ids.dim() != 2 or attention_mask.dim() != 2:
-            raise RuntimeError("prefix sharing FSDP path expects 2D or jagged NestedTensor input_ids")
+            raise RuntimeError("PrefixSharing + FSDP expects 2D or jagged NestedTensor input_ids")
         if input_ids.shape != attention_mask.shape:
             raise RuntimeError("input_ids and attention_mask must have the same shape")
-
         valid_indices = [
             attention_mask[row].nonzero(as_tuple=False).flatten()
             for row in range(input_ids.shape[0])
         ]
-        sequences = _extract_sequences_async(input_ids, valid_indices)
-    prefix_sharing_plan = PrefixSharingPlanner(config).plan(sequences)
+        sequences = extract_seq_from_dense_tensor(input_ids, valid_indices)
+
+    #########################################################
+    # STEP 1: plan for PrefixSharing
+    #########################################################
+    prefix_sharing_plan = PrefixSharingPlanner(ps_config).plan(sequences)
     if not prefix_sharing_plan.has_sharing:
-        return batch, None
-
-    if is_nested_input:
-        trimmed_micro_batch = _trim_nested_batch(batch, prefix_sharing_plan)
-        kept_position_rows = _collect_kept_position_rows(
-            trimmed_micro_batch,
-            prefix_sharing_plan,
-            is_nested_tensor=True,
+        return micro_batch, None
+
+    #########################################################
+    # STEP 2: trim redundant prefix of reusers in the micro-batch
+    #########################################################
+    if input_ids_is_nested: # nested tensor
+        trimmed_micro_batch, kept_position_rows = trim_redundant_prefix_in_nested_tensor(
+            micro_batch, prefix_sharing_plan
         )
-    else:
-        trimmed_micro_batch = _clone_batch(batch)
-        trimmed_attention_mask = attention_mask.clone()
-        trimmed_attention_mask[:] = False
-
-        for row, indices in enumerate(valid_indices):
-            keep_start, keep_end = prefix_sharing_plan.input_keep_ranges[row]
-            kept_indices = indices[keep_start:keep_end]
-            trimmed_attention_mask[row, kept_indices] = True
-
-        trimmed_micro_batch["attention_mask"] = trimmed_attention_mask
-
-        if "loss_mask" in trimmed_micro_batch:
-            trimmed_loss_mask = trimmed_micro_batch["loss_mask"].to(bool).clone()
-            trimmed_loss_mask[:] = False
-            for row, indices in enumerate(valid_indices):
-                keep_start, keep_end = prefix_sharing_plan.loss_mask_keep_ranges[row]
-                trimmed_loss_mask[row, indices[keep_start:keep_end]] = batch["loss_mask"][row, indices[keep_start:keep_end]].to(bool)
-            trimmed_micro_batch["loss_mask"] = trimmed_loss_mask
-        kept_position_rows = _collect_kept_position_rows(
-            trimmed_micro_batch,
-            prefix_sharing_plan,
-            is_nested_tensor=False,
-            valid_indices=valid_indices,
+    else: # dense tensor
+        trimmed_micro_batch, kept_position_rows = trim_redundant_prefix_in_dense_tensor(
+            micro_batch, prefix_sharing_plan, valid_indices
         )
 
-    packed_batch_layout = PackedBatchLayout.from_kept_position_rows(
-        kept_position_rows,
-        align_size=1,
-    )
-    runtime_state = PrefixSharingRuntimeState(
+    #########################################################
+    # STEP 3: build batch layout for the trimmed micro-batch
+    #########################################################
+    packed_batch_layout = PackedBatchLayout.from_kept_position_rows(kept_position_rows)
+
+    #########################################################
+    # STEP 4: build runtime state for the trimmed micro-batch
+    #########################################################
+    prefix_sharing_runtime_state = PrefixSharingRuntimeState(
         prefix_sharing_plan=prefix_sharing_plan,
-        attention_backend=get_backend_instance(config, backend),
+        attention_backend=get_backend_instance(ps_config, backend),
         packed_batch_layout=packed_batch_layout,
         parallel_info=MegatronParallelInfo(),
         kept_position_ids=trimmed_micro_batch.get("position_ids"),
     )
-    return trimmed_micro_batch, runtime_state
+
+    return trimmed_micro_batch, prefix_sharing_runtime_state
 
 
 def restore_prefix_sharing_outputs_2d(
@@ -375,6 +383,10 @@ def _run_packed_attention_runtime(
     if _per_layer_ok:
         profiler.start_phase(f"attn.kv.l{layer_id}")
     # [PS-perf] end ————————————————————————————————————————
+
+    #########################################################
+    # STEP 1: build key and value tensors
+    #########################################################
     expanded_key, expanded_value = ctx.attention_backend.build_kv(
         packed_key,
         packed_value,
@@ -399,6 +411,10 @@ def _run_packed_attention_runtime(
         dump_expanded_kv_on(layer_number, expanded_key, expanded_value, num_layers)
     # ##### [PS-diag] end #####
 
+    #########################################################
+    # STEP 2: compute attention
+    #########################################################
+    
     if profiler is not None:
         profiler.start_phase(PerfProfiler.PHASE_ATTN_COMPUTE)
     # [PS-perf] start — per-layer compute ——————————————————
@@ -426,7 +442,7 @@ def _call_fsdp_model(
     model: Any,
     micro_batch: Any,
     *,
-    prefix_sharing_runtime: PrefixSharingFSDPAttentionRuntime,
+    attention_runtime: PrefixSharingFSDPAttentionRuntime,
     enable_prefix_sharing: bool,
 ) -> Any:
     model_inputs = {
@@ -436,7 +452,7 @@ def _call_fsdp_model(
     }
     model_inputs["use_cache"] = False
     if enable_prefix_sharing:
-        model_inputs["prefix_sharing_runtime"] = prefix_sharing_runtime
+        model_inputs["prefix_sharing_runtime"] = attention_runtime
     try:
         return model(**model_inputs)
     except TypeError:
diff --git a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py
index 6920798d..666359e3 100644
--- a/prefix-sharing/prefix_sharing/integrations/verl_mcore.py
+++ b/prefix-sharing/prefix_sharing/integrations/verl_mcore.py
@@ -20,354 +20,44 @@
 from prefix_sharing.core.config import PrefixSharingConfig
 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.parallel_info import get_megatron_parallel_info
 from prefix_sharing.integrations.runtime_state import PrefixSharingRuntimeState
-from prefix_sharing.integrations.verl_utils import _clone_batch
-from prefix_sharing.integrations.verl_utils import _collect_kept_position_rows
-from prefix_sharing.integrations.verl_utils import _extract_seq_from_nested_tensor
-from prefix_sharing.integrations.verl_utils import _extract_sequences_async
-from prefix_sharing.integrations.verl_utils import _is_nested_tensor
-from prefix_sharing.integrations.verl_utils import _trim_nested_batch
-from prefix_sharing.integrations.verl_utils import _trim_plain_batch_thd
+from prefix_sharing.integrations.verl_utils import collect_kept_position_rows
+from prefix_sharing.integrations.verl_utils import extract_seq_from_dense_tensor
+from prefix_sharing.integrations.verl_utils import extract_seq_from_nested_tensor
+from prefix_sharing.integrations.verl_utils import is_nested_tensor
+from prefix_sharing.integrations.verl_utils import trim_redundant_prefix_in_nested_tensor
+from prefix_sharing.integrations.verl_utils import trim_plain_batch_thd
 from prefix_sharing.integrations.verl_utils import read_ps_config_from_engine_config
 
 
-def restore_reuser_prefix_columns_2d(
-    output: dict[str, Any],
-    vocab_parallel_log_probs_fn: Any,
-    vocab_parallel_entropy_fn: Any = None,
-) -> dict[str, Any]:
-    """Restore reuser prefix columns in 2D space — build_kv-style slice + concat.
-
-    Mirrors :meth:`TorchReferenceBackend.build_kv`
-    (``torch.cat([provider_kv[:prefix_len], own_suffix])``): instead of writing
-    each prefix token one scalar at a time, the whole prefix interval is sliced
-    off the **direct provider's already-restored 2D row** and only the single
-    prefix-last logprob is recomputed.
-
-    Per ``reuser_idx`` with direct provider ``provider_idx = provider_index[reuser_idx]``
-    and ``P = prefix_lens[reuser_idx]`` (columns are identity-mapped in the
-    unfolded 2D tensor, so ``target_2d_pos`` == column):
-
-    - **interior ``[0, P-2]``**: ``log_probs[reuser_idx, 0:P-1] = log_probs[provider_idx, 0:P-1]``
-      (bulk copy). Identical across the shared prefix (same logits + labels),
-      and ``provider_idx`` was restored earlier in the batch-order loop, so its
-      row already holds correct values — no per-position provider resolution
-      needed.
-    - **prefix-last ``P-1``**: recompute ``log_probs[reuser_idx, P-1]`` from the
-      saved provider logits + the reuser's own first-suffix label (differs from
-      the provider's). When the reuser has no suffix (``suffix_len == 0``) the
-      planner emits no prefix-last spec; that column is masked downstream, so
-      the provider's value is copied as a safe placeholder.
-    - **entropy ``[0, P-1]``**: ``entropy[reuser_idx, 0:P] = entropy[provider_idx, 0:P]``
-      (whole prefix copied, including prefix-last — entropy is label-independent).
-
-    Rows are visited in ``range(B)`` order so a provider is always restored
-    before any reuser that reads it (the same online-detector invariant
-    ``build_kv`` relies on).
-
-    Args:
-        output: Output dict with ``log_probs`` [B, L] and optionally
-            ``entropy`` [B, L] in 2D space (unfolded from the trimmed
-            NestedTensor by :func:`restore_via_2d_unfold_verl080`).
-        vocab_parallel_log_probs_fn: ``logits [1, V//tp]``, ``label [1]`` →
-            scalar; used only for the prefix-last recompute.
-        vocab_parallel_entropy_fn: Retained for call-site compatibility;
-            unused (entropy is copied, never recomputed).
-
-    Returns:
-        ``output`` with ``log_probs`` and ``entropy`` mutated in-place.
-    """
-
-    ctx = current_prefix_sharing_context()
-    if ctx is None:
-        return output
-    plan = ctx.prefix_sharing_plan
-    # Guard on reuser presence (not on prefix_last_restore_indices): a batch
-    # whose reusers all have suffix_len == 0 emits no prefix-last spec but still
-    # needs its interior prefix columns restored.
-    if not plan.has_sharing:
-        return output
-
-    import torch
-
-    log_probs = output.get("log_probs")
-    if log_probs is None:
-        return output
-    entropy = output.get("entropy")
-
-    provider_index = plan.provider_index
-    prefix_lens = plan.prefix_lens
-
-    # reuser row → its prefix-last restore spec (one per reuser-with-suffix;
-    # interior positions have no spec — they are bulk-sliced below).
-    prefix_last_spec_by_reuser = {
-        spec.reuse_idx_in_batch: spec for spec in ctx.prefix_last_restore_indices
-    }
-
-    restored_reusers = 0
-    # Row 0 is always a provider (nothing precedes it to reuse), so start at 1.
-    # A reuser's provider always has a smaller batch index (online-detector
-    # invariant), so it is already restored when we reach reuser_idx.
-    for reuser_idx in range(1, len(prefix_lens)):
-        prefix_len = prefix_lens[reuser_idx]
-        if provider_index[reuser_idx] == reuser_idx or prefix_len <= 0:
-            continue  # provider / non-reuser: row already complete
-        provider_idx = provider_index[reuser_idx]
-
-        # interior [0, prefix_len-2]: bulk-copy from the provider's restored row.
-        if prefix_len - 1 > 0:
-            log_probs[reuser_idx, 0:prefix_len - 1] = log_probs[provider_idx, 0:prefix_len - 1]
-
-        # prefix-last (position prefix_len-1): recompute with the reuser's label.
-        prefix_last_spec = prefix_last_spec_by_reuser.get(reuser_idx)
-        if prefix_last_spec is not None:
-            saved_logits_key = (reuser_idx, prefix_last_spec.target_2d_pos)
-            saved_provider_logits = ctx.prefix_last_logits_saved[saved_logits_key]  # [1, V//tp]
-            reuser_label = torch.tensor(
-                [prefix_last_spec.label_value], dtype=torch.long, device=log_probs.device,
-            )  # [1]
-            log_probs[reuser_idx, prefix_len - 1] = vocab_parallel_log_probs_fn(
-                saved_provider_logits, reuser_label,
-            ).reshape(())
-        else:
-            # suffix_len == 0: no prefix-last spec; column is masked downstream.
-            log_probs[reuser_idx, prefix_len - 1] = log_probs[provider_idx, prefix_len - 1]
-
-        # entropy [0, prefix_len-1]: whole prefix copied (label-independent).
-        if entropy is not None:
-            entropy[reuser_idx, 0:prefix_len] = entropy[provider_idx, 0:prefix_len]
-
-        restored_reusers += 1
-
-    if ctx.stats is not None:
-        ctx.stats.record_restore(restored_reusers)
-    return output
-
-
-# ═══════════════════════════════════════════════════════════════
-# v080 restore wrapper: NestedTensor → 2D left-pad → reuse 2D restore → pack back
-# ═══════════════════════════════════════════════════════════════
-
-
-def restore_via_2d_unfold_verl080(
-    output: dict,
-    vocab_parallel_log_probs_fn: Any,
-    vocab_parallel_entropy_fn: Any = None,
-) -> dict:
-    """v080 restore wrapper: NestedTensor → 2D left-pad → reuse restore_reuser_prefix_columns_2d → pack back.
-
-    After v080 physical trimming, reuser NestedTensor rows contain only the
-    suffix region; the prefix region (including prefix-last) has been physically
-    removed. This function performs reassembly at the forward_step exit (while
-    the context is still active and provider prefix-last logits have been saved
-    into ``ctx.prefix_last_logits_saved``):
-
-    1. Unfold each trimmed NestedTensor row into a full 2D ``[B, L_max]`` tensor
-       (reuser prefix region left-padded with 0, right-padded with 0 to L_max).
-    2. Reuse :func:`restore_reuser_prefix_columns_2d`: the interior interval is
-       bulk-sliced from the direct provider's already-restored 2D row; the
-       prefix-last position is recomputed from the saved logits +
-       ``index.label_value``.
-    3. Slice and pack back into a NestedTensor (jagged) per ``original_lengths``.
-
-    Column mapping is identity: after left-padding, the 0-based offset of valid
-    content equals the 2D column index, so ``target_2d_pos`` is used directly as
-    a column index — no ``valid_indices`` / column mapping table needed.
-
-    Must be called inside ``prefix_sharing_runtime_context`` (reads
-    ``current_prefix_sharing_context``), after the vocab_logprobs patch has saved
-    provider prefix-last logits into ``ctx.prefix_last_logits_saved``.
-
-    Args:
-        output: The output_dict returned by forward_step, containing
-            ``"log_probs"`` NestedTensor (trimmed jagged), and optionally
-            ``"entropy"`` NestedTensor. **Must not** be wrapped in a tuple
-            (tuple unpacking is the caller's responsibility).
-        vocab_parallel_log_probs_fn: Used for prefix-last logp recomputation.
-        vocab_parallel_entropy_fn: Optional; currently unused (entropy follows
-            the copy path — both interior and prefix-last are copied from the
-            provider, never recomputed).
-
-    Returns:
-        ``output`` with ``log_probs``/``entropy`` replaced by reassembled
-        NestedTensors.
-    """
-
-    ctx = current_prefix_sharing_context()
-    if ctx is None:
-        return output
-    plan = ctx.prefix_sharing_plan
-    # Guard on reuser presence, not on prefix_last_restore_indices: a batch
-    # whose reusers all have suffix_len == 0 has no prefix-last spec but still
-    # needs interior prefix columns restored.
-    if not plan.has_sharing:
-        return output
-
-    log_probs_nested = output.get("log_probs")
-    if log_probs_nested is None or not _is_nested_tensor(log_probs_nested):
-        return output
-    entropy_nested = output.get("entropy")
-    has_entropy = entropy_nested is not None and _is_nested_tensor(entropy_nested)
-
-    original_lengths = plan.original_lengths
-    input_keep_ranges = plan.input_keep_ranges
-    B = len(original_lengths)
-    if B == 0:
-        return output
-    L_max = max(original_lengths)
-
-    # --- Step 1: Unfold trimmed NestedTensor → full 2D [B, L_max] ---
-    log_probs_2d, entropy_2d = _unfold_trimmed_nested_to_2d(
-        log_probs_nested,
-        entropy_nested if has_entropy else None,
-        original_lengths,
-        input_keep_ranges,
-        L_max,
-        B,
-    )
-
-    # --- Step 2: Reuse restore_reuser_prefix_columns_2d ---
-    # build_kv-style interval concatenation: interior is bulk-sliced from the
-    # direct provider's already-restored 2D row; prefix-last is recomputed from
-    # index.label_value + saved logits. Identity column mapping (target_2d_pos
-    # is the 2D column index, no left padding).
-    output_2d: dict[str, Any] = {"log_probs": log_probs_2d}
-    if entropy_2d is not None:
-        output_2d["entropy"] = entropy_2d
-    output_2d = restore_reuser_prefix_columns_2d(
-        output_2d,
-        vocab_parallel_log_probs_fn,
-        vocab_parallel_entropy_fn,
-    )
-
-    # --- Step 3: Pack back into NestedTensor per original_lengths ---
-    output["log_probs"] = _fold_2d_to_nested(output_2d["log_probs"], original_lengths)
-    if entropy_2d is not None:
-        output["entropy"] = _fold_2d_to_nested(output_2d["entropy"], original_lengths)
-
-    num_prefix_last = len(ctx.prefix_last_restore_indices)
-    print(
-        f"[PS][restore_verl080] unfolded B={B} L_max={L_max}, "
-        f"restored reusers={num_prefix_last} (prefix-last entries; "
-        f"interior handled by bulk slice)",
-        flush=True,
-    )
-    return output
-
-
-def _unfold_trimmed_nested_to_2d(
-    log_probs_nested: Any,
-    entropy_nested: Any,
-    original_lengths: list[int],
-    input_keep_ranges: list,
-    L_max: int,
-    B: int,
-) -> tuple[Any, Any | None]:
-    """Unfold trimmed NestedTensor → full 2D [B, L_max] (reuser prefix left-padded with 0).
-
-    After trimming, each row is:
-      - provider (keep_start=0): full [prefix | suffix], length = original_lengths[i]
-      - reuser  (keep_start=prefix_len>0): only [suffix], length = original_lengths[i]-prefix_len
-
-    After unfolding, each row becomes [prefix_zeros | suffix], then right-padded
-    with 0 to L_max. The left-padded zeros are not in the autograd graph, but
-    restore overwrites the prefix region (interior copied from provider,
-    prefix-last recomputed), so the final values are in the graph. The right-pad
-    tail is discarded when packing back.
-    """
-    import torch
-
-    log_probs_offsets = log_probs_nested.offsets()
-    log_probs_values = log_probs_nested.values()
-    if entropy_nested is not None:
-        entropy_offsets = entropy_nested.offsets()
-        entropy_values = entropy_nested.values()
-
-    log_probs_rows: list[Any] = []
-    entropy_rows: list[Any] | None = [] if entropy_nested is not None else None
-
-    for seq_idx in range(B):
-        orig_len = original_lengths[seq_idx]
-        prefix_len = input_keep_ranges[seq_idx][0]
-
-        log_probs_suffix = log_probs_values[log_probs_offsets[seq_idx]:log_probs_offsets[seq_idx + 1]]
-        log_probs_rows.append(_build_padded_row(log_probs_suffix, prefix_len, orig_len, L_max))
-
-        if entropy_nested is not None:
-            entropy_suffix = entropy_values[entropy_offsets[seq_idx]:entropy_offsets[seq_idx + 1]]
-            entropy_rows.append(_build_padded_row(entropy_suffix, prefix_len, orig_len, L_max))
-
-    log_probs_2d = torch.stack(log_probs_rows, dim=0)
-    entropy_2d = torch.stack(entropy_rows, dim=0) if entropy_rows else None
-    return log_probs_2d, entropy_2d
-
-
-def _build_padded_row(
-    suffix_data: Any, prefix_len: int, orig_len: int, L_max: int,
-) -> Any:
-    """Build one full 2D row ``[prefix_zeros | suffix]`` right-padded with 0 to L_max."""
-    import torch
-
-    device = suffix_data.device
-    dtype = suffix_data.dtype
-    tail_shape = tuple(suffix_data.shape[1:])
-    pieces: list[Any] = []
-    if prefix_len > 0:
-        pieces.append(torch.zeros((prefix_len,) + tail_shape, dtype=dtype, device=device))
-    pieces.append(suffix_data)
-    row = torch.cat(pieces, dim=0)  # [orig_len, ...]
-    if orig_len < L_max:
-        pad = torch.zeros((L_max - orig_len,) + tail_shape, dtype=dtype, device=device)
-        row = torch.cat([row, pad], dim=0)
-    return row
-
-
-def _fold_2d_to_nested(tensor_2d: Any, original_lengths: list[int]) -> Any:
-    """Full 2D [B, L_max] → NestedTensor (jagged), sliced per original_lengths."""
-    import torch
-
-    rows = [tensor_2d[seq_idx, :original_lengths[seq_idx]] for seq_idx in range(len(original_lengths))]
-    values = torch.cat(rows, dim=0) if rows else tensor_2d.reshape(0, *tensor_2d.shape[2:])
-    offsets = torch.tensor(
-        [0] + [sum(original_lengths[: idx + 1]) for idx in range(len(original_lengths))],
-        dtype=torch.long,
-        device=tensor_2d.device,
-    )
-    if hasattr(torch.nested, "nested_tensor_from_jagged"):
-        return torch.nested.nested_tensor_from_jagged(values, offsets)
-    return torch.nested.as_nested_tensor(rows, layout=torch.jagged)
-
-
 def build_prefix_sharing_micro_batch_verl080(
     engine_self: Any,
     batch: Any,
     ps_config: PrefixSharingConfig,
 ) -> tuple[Any, PrefixSharingRuntimeState | None]:
-    """Prefix-sharing micro-batch construction for the verl 0.8.0 engine architecture.
+    """verl 0.8.0 engine 架构下的 prefix-sharing micro-batch 构建。
 
-    MCore/THD path: after NestedTensor trimming, unfold by kept segments into a
-    packed layout. 2D path: physically trim input_ids/position_ids and mark
-    valid positions via attention_mask.
+    MCore/THD 路径：NestedTensor 裁剪后按 kept 区段展开为 packed layout，
+    2D 路径：物理裁剪 input_ids/position_ids，通过 attention_mask 标记 valid。
 
-    The ps_config parameter has already been parsed by the caller via
-    PrefixSharingConfig.from_raw(); no further from_raw call is needed.
+    参数 ps_config 已由调用方通过 PrefixSharingConfig.from_raw() 解析完成，
+    不需要再次 from_raw。
 
-    Core principle: 2D + attention_mask is the primary path; the NestedTensor
-    path is an optional optimization only for GPU + use_remove_padding=True.
-    NPU does not support torch.nested, so all NPU scenarios use the 2D path.
+    核心原则：2D + attention_mask 为主路径，NestedTensor 路径仅在
+    GPU + use_remove_padding=True 时作为可选优化。
+    NPU 不支持 torch.nested，所有 NPU 场景都走 2D 路径。
     """
     # ── PATH 1: prefix sharing disabled ──
     if not ps_config.enable_prefix_sharing:
         print("[PS][prepare] PATH 1: prefix sharing disabled")
         return batch, None
 
-    # ── Stage 1: Config validation ──
+    # ── 阶段 1: 配置校验 ──
     use_remove_padding = getattr(engine_self.engine_config, "use_remove_padding", True)
     ps_config.validate_for_engine(use_remove_padding=use_remove_padding)
 
-    # ── Stage 2: Reject unsupported features ──
+    # ── 阶段 2: 拒绝不支持的特性 ──
     try:
         from verl.utils import tensordict_utils as tu
         use_fused = tu.get_non_tensor_data(batch, "use_fused_kernels", default=False)
@@ -378,18 +68,18 @@ def build_prefix_sharing_micro_batch_verl080(
     if getattr(engine_self.engine_config, "dynamic_context_parallel", False):
         raise RuntimeError("prefix sharing phase 1 does not support dynamic context parallel")
 
-    # ── Stage 3: Extract sequences from batch ──
-    # NestedTensor → extract from offsets/values
-    # Plain 2D → extract from attention_mask.nonzero()
-    # Also keep attention_mask_bool for _collect_kept_position_rows in Stage 6.
+    # ── 阶段 3: 从 batch 提取序列 ──
+    # NestedTensor → 从 offsets/values 提取；kept_position_rows 由 trim 直接返回。
+    # Plain 2D → 从 attention_mask.nonzero() 提取，并保留 valid_indices
+    # 供 trim 后的 collect_kept_position_rows 使用。
     input_ids = batch["input_ids"]
-    is_nested_tensor = _is_nested_tensor(input_ids)
+    is_nested_input = is_nested_tensor(input_ids)
     attention_mask_bool_for_layout = None
 
-    if is_nested_tensor:
-        sequences = _extract_seq_from_nested_tensor(input_ids)
+    if is_nested_input:
+        sequences = extract_seq_from_nested_tensor(input_ids)
     else:
-        # plain 2D tensor (requires attention_mask)
+        # plain 2D tensor（需要 attention_mask）
         attention_mask = batch.get("attention_mask")
         if attention_mask is None:
             print("[PS][prepare] PATH 4: plain 2D batch without attention_mask")
@@ -399,32 +89,29 @@ def build_prefix_sharing_micro_batch_verl080(
             attention_mask_bool[row].nonzero(as_tuple=False).flatten()
             for row in range(input_ids.shape[0])
         ]
-        sequences = _extract_sequences_async(input_ids, valid_indices)
+        sequences = extract_seq_from_dense_tensor(input_ids, valid_indices)
 
-    # ── Stage 4: Prefix sharing planning ──
+    # ── 阶段 4: 前缀共享规划 ──
     plan = PrefixSharingPlanner(ps_config).plan(sequences)
     if not plan.has_sharing:
         print("[PS][prepare] no prefix sharing detected")
         return batch, None
 
-    # ── Stage 5: Physically trim batch ──
-    #   NestedTensor path: trim input_ids/position_ids/loss_mask to match kept segments.
-    #   2D path: only modify attention_mask (Megatron dynamically recomputes packed
-    #   from mask); v080 THD path uses preprocess_thd_engine(input_ids) to process
-    #   data directly, ignoring attention_mask. Must physically trim
-    #   input_ids/position_ids.
-    if is_nested_tensor:
-        trimmed_batch = _trim_nested_batch(batch, plan)
+    # ── 阶段 5: 物理裁剪 batch ──
+    #   NestedTensor path: 裁剪 input_ids/position_ids/loss_mask 以匹配 kept 区段。
+    #   2D path: 只改 attention_mask（Megatron 从 mask 动态重算 packed），
+    #   v080 THD 路径用 preprocess_thd_engine(input_ids) 直接处理数据，
+    #   不看 attention_mask。必须物理裁剪 input_ids/position_ids。
+    if is_nested_input:
+        trimmed_batch, kept_position_rows = trim_redundant_prefix_in_nested_tensor(batch, plan)
     else:
-        trimmed_batch = _trim_plain_batch_thd(batch, plan, valid_indices)
-
-    # Layout computation: build from the actual kept position rows after trimming
-    kept_position_rows = _collect_kept_position_rows(
-        trimmed_batch, plan, is_nested_tensor,
-        valid_indices=valid_indices if not is_nested_tensor else None,
-    )
+        trimmed_batch = trim_plain_batch_thd(batch, plan, valid_indices)
+        kept_position_rows = collect_kept_position_rows(
+            trimmed_batch, plan, is_nested_input,
+            valid_indices=valid_indices,
+        )
 
-    # ── Stage 6: Build layout ──
+    # ── 阶段 6: 构建 layout ──
     parallel_info = get_megatron_parallel_info()
     align_size = (
         parallel_info.tp_size * parallel_info.cp_size * 2
@@ -436,7 +123,7 @@ def build_prefix_sharing_micro_batch_verl080(
         align_size=int(align_size),
     )
 
-    # ── Stage 7: Build state ──
+    # ── 阶段 7: 构建 state ──
     state = PrefixSharingRuntimeState(
         prefix_sharing_plan=plan,
         attention_backend=get_backend_instance(ps_config),
diff --git a/prefix-sharing/prefix_sharing/integrations/verl_utils.py b/prefix-sharing/prefix_sharing/integrations/verl_utils.py
index 940afb8d..f387ed07 100644
--- a/prefix-sharing/prefix_sharing/integrations/verl_utils.py
+++ b/prefix-sharing/prefix_sharing/integrations/verl_utils.py
@@ -4,7 +4,10 @@
 
 from typing import Any
 
+import torch
+
 from prefix_sharing.core.planner import PrefixSharingPlan
+from prefix_sharing.integrations.context import current_prefix_sharing_context
 
 
 def read_ps_config_from_engine_config(engine_config: Any) -> Any | None:
@@ -28,10 +31,10 @@ def read_ps_config_from_engine_config(engine_config: Any) -> Any | None:
     if explicit_config is not None:
         return explicit_config
 
-    return _prefix_sharing_config_from_prefix_grouper(engine_config)
+    return read_ps_config_from_prefix_grouper(engine_config)
 
 
-def _prefix_sharing_config_from_prefix_grouper(engine_config: Any) -> dict[str, Any] | None:
+def read_ps_config_from_prefix_grouper(engine_config: Any) -> dict[str, Any] | None:
     use_prefix_grouper = _read_actor_value(engine_config, "use_prefix_grouper", False)
     if not use_prefix_grouper:
         return None
@@ -64,7 +67,7 @@ def _prefix_sharing_config_from_prefix_grouper(engine_config: Any) -> dict[str,
     return values
 
 
-def _clone_batch(batch: Any) -> Any:
+def clone_batch(batch: Any) -> Any:
     if hasattr(batch, "clone"):
         try:
             return batch.clone()
@@ -98,23 +101,32 @@ def _read_actor_value(config: Any, dotted_name: str, default: Any) -> Any:
     return current
 
 
-def _trim_nested_batch(batch: Any, plan: PrefixSharingPlan) -> Any:
-    """Physically trim a NestedTensor batch for verl 0.8 THD paths."""
+def trim_redundant_prefix_in_nested_tensor(
+    batch: Any, prefix_sharing_plan: PrefixSharingPlan
+) -> tuple[Any, list[Any]]:
+    """Trim redundant prefix tokens in NestedTensor inputs so that their computation
+    are skipped and not performed.
 
-    import torch
+    Args:
+        batch: Input batch with NestedTensor ``input_ids``.
+        prefix_sharing_plan: PrefixSharingPlan with keep ranges.
 
-    trimmed_batch = _clone_batch(batch)
+    Returns:
+        ``(trimmed_batch, kept_position_rows)``. ``kept_position_rows`` is the
+        per-row position-id slice used to build the trimmed NestedTensor, so
+        callers do not need to unpack it again.
+    """
 
+    trimmed_batch = clone_batch(batch)
     input_ids = batch["input_ids"]
     position_ids = batch["position_ids"]
+    
+    # trim input_ids
+    trimmed_batch["input_ids"] = _trim_nested_tensor(input_ids, prefix_sharing_plan)
 
-    trimmed_ids_seqs = _slice_nested_sequences(input_ids, plan)
-    new_input_ids = torch.nested.as_nested_tensor(trimmed_ids_seqs, layout=torch.jagged)
-    trimmed_batch["input_ids"] = new_input_ids
-
-    if _is_nested_tensor(position_ids):
-        trimmed_pos_seqs = _slice_nested_sequences(position_ids, plan)
-        new_position_ids = torch.nested.as_nested_tensor(trimmed_pos_seqs, layout=torch.jagged)
+    # trim position_ids
+    if is_nested_tensor(position_ids):
+        kept_position_rows = _trim_nested_rows(position_ids, prefix_sharing_plan)
     else:
         attention_mask = batch.get("attention_mask")
         if attention_mask is not None:
@@ -124,23 +136,74 @@ def _trim_nested_batch(batch: Any, plan: PrefixSharingPlan) -> Any:
                 position_ids.shape[0], position_ids.shape[1],
                 dtype=torch.bool, device=position_ids.device,
             )
-        trimmed_pos_seqs = _slice_2d_position_rows(
-            position_ids, plan, attention_mask_bool,
+        kept_position_rows = _trim_2d_tensor(
+            position_ids, prefix_sharing_plan, attention_mask_bool
         )
-        new_position_ids = torch.nested.as_nested_tensor(trimmed_pos_seqs, layout=torch.jagged)
-    trimmed_batch["position_ids"] = new_position_ids
+    trimmed_batch["position_ids"] = torch.nested.as_nested_tensor(
+        kept_position_rows, layout=torch.jagged
+    )
 
+    # trim loss_mask
     loss_mask = batch.get("loss_mask")
-    if loss_mask is not None and _is_nested_tensor(loss_mask):
-        trimmed_loss_seqs = _slice_nested_sequences(loss_mask, plan)
-        trimmed_batch["loss_mask"] = torch.nested.as_nested_tensor(
-            trimmed_loss_seqs, layout=torch.jagged
-        )
+    if loss_mask is not None and is_nested_tensor(loss_mask):
+        trimmed_batch["loss_mask"] = _trim_nested_tensor(loss_mask, prefix_sharing_plan)
 
-    return trimmed_batch
+    return trimmed_batch, kept_position_rows
+
+
+def trim_redundant_prefix_in_dense_tensor(
+    batch: Any,
+    prefix_sharing_plan: PrefixSharingPlan,
+    valid_indices: list[Any],
+) -> tuple[Any, list[Any]]:
+    """Trim redundant prefix tokens from a dense 2D batch by masking.
+
+    Unlike ``trim_plain_batch_thd`` which physically removes kept tokens,
+    this helper keeps ``input_ids`` and ``position_ids`` dense and only
+    masks the kept positions in ``attention_mask`` and ``loss_mask``.
+    The position-id rows are trimmed to jagged tensors for packed layout
+    construction, but the original 2D ``position_ids`` tensor is left
+    unchanged so downstream dense-path code can still index it.
+
+    Args:
+        batch: Input batch with 2D ``input_ids``, ``position_ids`` and
+            ``attention_mask``.
+        prefix_sharing_plan: Prefix sharing plan with keep ranges.
+        valid_indices: Pre-computed nonzero indices from ``attention_mask``.
+
+    Returns:
+        ``(trimmed_batch, kept_position_rows)``.
+    """
 
+    trimmed_batch = clone_batch(batch)
+    attention_mask = batch["attention_mask"].to(bool)
+    trimmed_attention_mask = attention_mask.clone()
+    trimmed_attention_mask[:] = False
 
-def _trim_plain_batch_thd(batch: Any, plan: PrefixSharingPlan, valid_indices: list[Any] | None = None) -> Any:
+    for row, indices in enumerate(valid_indices):
+        keep_start, keep_end = prefix_sharing_plan.input_keep_ranges[row]
+        kept_indices = indices[keep_start:keep_end]
+        trimmed_attention_mask[row, kept_indices] = True
+    trimmed_batch["attention_mask"] = trimmed_attention_mask
+
+    if "loss_mask" in trimmed_batch:
+        trimmed_loss_mask = trimmed_batch["loss_mask"].to(bool).clone()
+        trimmed_loss_mask[:] = False
+        for row, indices in enumerate(valid_indices):
+            keep_start, keep_end = prefix_sharing_plan.loss_mask_keep_ranges[row]
+            trimmed_loss_mask[row, indices[keep_start:keep_end]] = (
+                batch["loss_mask"][row, indices[keep_start:keep_end]].to(bool)
+            )
+        trimmed_batch["loss_mask"] = trimmed_loss_mask
+
+    kept_position_rows = _trim_2d_tensor(
+        batch["position_ids"], prefix_sharing_plan, attention_mask
+    )
+
+    return trimmed_batch, kept_position_rows
+
+
+def trim_plain_batch_thd(batch: Any, plan: PrefixSharingPlan, valid_indices: list[Any] | None = None) -> Any:
     """Physically trim a plain 2D tensor batch for verl 0.8 THD paths.
 
     Args:
@@ -150,8 +213,6 @@ def _trim_plain_batch_thd(batch: Any, plan: PrefixSharingPlan, valid_indices: li
             If None, will compute from batch["attention_mask"].
     """
 
-    import torch
-
     input_ids = batch["input_ids"]
     position_ids = batch["position_ids"]
 
@@ -179,7 +240,7 @@ def _trim_plain_batch_thd(batch: Any, plan: PrefixSharingPlan, valid_indices: li
         kept_id_rows.append(input_ids[row, kept_indices])
         kept_pos_rows.append(position_ids[row, kept_indices])
 
-    trimmed_batch = _clone_batch(batch)
+    trimmed_batch = clone_batch(batch)
     trimmed_batch["input_ids"] = torch.nested.as_nested_tensor(kept_id_rows, layout=torch.jagged)
     trimmed_batch["position_ids"] = torch.nested.as_nested_tensor(kept_pos_rows, layout=torch.jagged)
 
@@ -198,37 +259,55 @@ def _trim_plain_batch_thd(batch: Any, plan: PrefixSharingPlan, valid_indices: li
     return trimmed_batch
 
 
-def _slice_nested_sequences(nested_tensor: Any, plan: PrefixSharingPlan) -> list[Any]:
+def _trim_nested_rows(nested_tensor: Any, prefix_sharing_plan: PrefixSharingPlan) -> list[Any]:
+    """Trim each sequence of a NestedTensor to its keep range in the plan.
+
+    Returns a list of 1D tensors (one per sequence), not packed into NestedTensor.
+    """
+
     offsets = nested_tensor.offsets()
     values = nested_tensor.values()
 
     sliced = []
-    for i in range(len(plan.input_keep_ranges)):
+    for i in range(len(prefix_sharing_plan.input_keep_ranges)):
         seq_values = values[offsets[i]:offsets[i + 1]]
-        keep_start, keep_end = plan.input_keep_ranges[i]
+        keep_start, keep_end = prefix_sharing_plan.input_keep_ranges[i]
         sliced.append(seq_values[keep_start:keep_end])
 
     return sliced
 
 
-def _slice_2d_position_rows(
-    position_ids: Any,
-    plan: PrefixSharingPlan,
+def _trim_nested_tensor(nested_tensor: Any, prefix_sharing_plan: PrefixSharingPlan) -> Any:
+    """Trim each sequence of a NestedTensor and pack into a jagged NestedTensor."""
+
+    return torch.nested.as_nested_tensor(
+        _trim_nested_rows(nested_tensor, prefix_sharing_plan), layout=torch.jagged
+    )
+
+
+def _trim_2d_tensor(
+    tensor_2d: Any,
+    prefix_sharing_plan: PrefixSharingPlan,
     attention_mask_bool: Any,
 ) -> list[Any]:
+    """Trim each row of a 2D dense tensor to its keep range in the plan.
+
+    Returns a list of 1D tensors (one per row), not packed into NestedTensor.
+    """
+
     kept_rows = []
-    for row in range(position_ids.shape[0]):
+    for row in range(tensor_2d.shape[0]):
         indices = attention_mask_bool[row].nonzero(as_tuple=False).flatten()
-        keep_start, keep_end = plan.input_keep_ranges[row]
+        keep_start, keep_end = prefix_sharing_plan.input_keep_ranges[row]
         kept_indices = indices[keep_start:keep_end]
-        kept_rows.append(position_ids[row, kept_indices])
+        kept_rows.append(tensor_2d[row, kept_indices])
     return kept_rows
 
 
-def _collect_kept_position_rows(
+def collect_kept_position_rows(
     trimmed_batch: Any,
-    plan: PrefixSharingPlan,
-    is_nested_tensor: bool,
+    prefix_sharing_plan: PrefixSharingPlan,
+    is_nested_input: bool,
     attention_mask_bool: Any | None = None,
     valid_indices: list[Any] | None = None,
 ) -> list[Any]:
@@ -237,7 +316,7 @@ def _collect_kept_position_rows(
     Args:
         trimmed_batch: Batch after trimming.
         plan: Prefix sharing plan.
-        is_nested_tensor: Whether position_ids is a NestedTensor.
+        is_nested_input: Whether position_ids is a NestedTensor.
         attention_mask_bool: Boolean attention mask (deprecated, use valid_indices).
         valid_indices: Pre-computed nonzero indices. If provided, skips
             attention_mask_bool computation.
@@ -245,10 +324,10 @@ def _collect_kept_position_rows(
 
     position_ids = trimmed_batch["position_ids"]
 
-    if is_nested_tensor or _is_nested_tensor(position_ids):
+    if is_nested_input or is_nested_tensor(position_ids):
         offsets = position_ids.offsets()
         values = position_ids.values()
-        return [values[offsets[i]:offsets[i + 1]] for i in range(len(plan.input_keep_ranges))]
+        return [values[offsets[i]:offsets[i + 1]] for i in range(len(prefix_sharing_plan.input_keep_ranges))]
 
     # 2D tensor — need valid_indices to locate valid column indices
     if valid_indices is None:
@@ -259,19 +338,19 @@ def _collect_kept_position_rows(
             )
         valid_indices = [
             attention_mask_bool[i].nonzero(as_tuple=False).flatten()
-            for i in range(len(plan.input_keep_ranges))
+            for i in range(len(prefix_sharing_plan.input_keep_ranges))
         ]
 
     rows = []
-    for i in range(len(plan.input_keep_ranges)):
+    for i in range(len(prefix_sharing_plan.input_keep_ranges)):
         indices = valid_indices[i]
-        keep_start, keep_end = plan.input_keep_ranges[i]
+        keep_start, keep_end = prefix_sharing_plan.input_keep_ranges[i]
         kept_indices = indices[keep_start:keep_end]
         rows.append(position_ids[i, kept_indices])
     return rows
 
 
-def _is_nested_tensor(tensor: Any) -> bool:
+def is_nested_tensor(tensor: Any) -> bool:
     return (
         hasattr(tensor, "offsets")
         and callable(tensor.offsets)
@@ -283,7 +362,7 @@ def _is_nested_tensor(tensor: Any) -> bool:
 def _copy_tensors_to_cpu_lists(tensors: list[Any]) -> list[list[int]]:
     """Copy device tensors to Python int lists with a single sync point.
 
-    Fast path (CUDA / NPU): stage all slices into pinned host buffers with
+    Fast path (GPU / NPU): stage all slices into pinned host buffers with
     non-blocking copies, then synchronize the source device once — avoiding
     one pipeline stall per tensor. Backends without pinned-memory support
     (older torch_npu, CPU, …) fall back to plain per-tensor copies: correct
@@ -293,8 +372,6 @@ def _copy_tensors_to_cpu_lists(tensors: list[Any]) -> list[list[int]]:
     works on any accelerator backend without hardcoding ``torch.cuda``.
     """
 
-    import torch
-
     if not tensors:
         return []
 
@@ -322,32 +399,29 @@ def _copy_tensors_to_cpu_lists(tensors: list[Any]) -> list[list[int]]:
     return [buf.tolist() for buf in pinned_buffers]
 
 
-def _extract_seq_from_nested_tensor(nested_tensor: Any) -> list[list[int]]:
-    """Extract per-sequence token ID lists from a NestedTensor (jagged layout).
-
-    Batched async device→CPU copies (CUDA/NPU), single synchronization point
-    to avoid pipeline stalls.
-    """
+def extract_seq_from_nested_tensor(nested_tensor: Any) -> list[list[int]]:
+    """Extract token sequences from a NestedTensor (jagged layout)."""
 
     offsets = nested_tensor.offsets()
     values = nested_tensor.values()
 
-    slices = [
+    tensor_slices = [
         values[offsets[i]:offsets[i + 1]].detach()
         for i in range(offsets.numel() - 1)
     ]
-    return _copy_tensors_to_cpu_lists(slices)
+    return _copy_tensors_to_cpu_lists(tensor_slices)
 
 
-def _extract_sequences_async(
+def extract_seq_from_dense_tensor(
     input_ids: Any,
     valid_indices: list[Any],
 ) -> list[list[int]]:
-    """Extract token sequences from a device tensor with batched async copy.
+    """Extract token sequences from a dense 2D device tensor with batched async copy.
 
-    Instead of calling .cpu().tolist() per row (which syncs each time),
-    all device→CPU copies are queued non-blocking, then synchronized once
-    (CUDA/NPU fast path; other backends fall back to plain copies).
+    Uses a batched async copy: instead of calling .cpu().tolist() per row
+    (which syncs each time), all device→CPU copies are queued non-blocking,
+    then synchronized once (CUDA/NPU fast path; other backends fall back to
+    plain copies).
 
     Args:
         input_ids: 2D device tensor of shape [batch_size, seq_len].
@@ -357,8 +431,297 @@ def _extract_sequences_async(
         List of token ID lists (CPU Python ints).
     """
 
-    tensors = [
+    tensor_slices = [
         input_ids[row, indices].detach()
         for row, indices in enumerate(valid_indices)
     ]
-    return _copy_tensors_to_cpu_lists(tensors)
+    return _copy_tensors_to_cpu_lists(tensor_slices)
+
+
+def restore_reuser_prefix_columns_2d(
+    output: dict[str, Any],
+    vocab_parallel_log_probs_fn: Any,
+    vocab_parallel_entropy_fn: Any = None,
+) -> dict[str, Any]:
+    """Restore reuser prefix columns in 2D space — build_kv-style slice + concat.
+
+    Mirrors :meth:`TorchReferenceBackend.build_kv`
+    (``torch.cat([provider_kv[:prefix_len], own_suffix])``): instead of writing
+    each prefix token one scalar at a time, the whole prefix interval is sliced
+    off the **direct provider's already-restored 2D row** and only the single
+    prefix-last logprob is recomputed.
+
+    Per ``reuser_idx`` with direct provider ``provider_idx = provider_index[reuser_idx]``
+    and ``P = prefix_lens[reuser_idx]`` (columns are identity-mapped in the
+    unfolded 2D tensor, so ``target_2d_pos`` == column):
+
+    - **interior ``[0, P-2]``**: ``log_probs[reuser_idx, 0:P-1] = log_probs[provider_idx, 0:P-1]``
+      (bulk copy). Identical across the shared prefix (same logits + labels),
+      and ``provider_idx`` was restored earlier in the batch-order loop, so its
+      row already holds correct values — no per-position provider resolution
+      needed.
+    - **prefix-last ``P-1``**: recompute ``log_probs[reuser_idx, P-1]`` from the
+      saved provider logits + the reuser's own first-suffix label (differs from
+      the provider's). When the reuser has no suffix (``suffix_len == 0``) the
+      planner emits no prefix-last spec; that column is masked downstream, so
+      the provider's value is copied as a safe placeholder.
+    - **entropy ``[0, P-1]``**: ``entropy[reuser_idx, 0:P] = entropy[provider_idx, 0:P]``
+      (whole prefix copied, including prefix-last — entropy is label-independent).
+
+    Rows are visited in ``range(B)`` order so a provider is always restored
+    before any reuser that reads it (the same online-detector invariant
+    ``build_kv`` relies on).
+
+    Args:
+        output: Output dict with ``log_probs`` [B, L] and optionally
+            ``entropy`` [B, L] in 2D space (unfolded from the trimmed
+            NestedTensor by :func:`restore_via_2d_unfold_verl080`).
+        vocab_parallel_log_probs_fn: ``logits [1, V//tp]``, ``label [1]`` →
+            scalar; used only for the prefix-last recompute.
+        vocab_parallel_entropy_fn: Retained for call-site compatibility;
+            unused (entropy is copied, never recomputed).
+
+    Returns:
+        ``output`` with ``log_probs`` and ``entropy`` mutated in-place.
+    """
+
+    ctx = current_prefix_sharing_context()
+    if ctx is None:
+        return output
+    plan = ctx.prefix_sharing_plan
+    # Guard on reuser presence (not on prefix_last_restore_indices): a batch
+    # whose reusers all have suffix_len == 0 emits no prefix-last spec but still
+    # needs its interior prefix columns restored.
+    if not plan.has_sharing:
+        return output
+
+    log_probs = output.get("log_probs")
+    if log_probs is None:
+        return output
+    entropy = output.get("entropy")
+
+    provider_index = plan.provider_index
+    prefix_lens = plan.prefix_lens
+
+    # reuser row → its prefix-last restore spec (one per reuser-with-suffix;
+    # interior positions have no spec — they are bulk-sliced below).
+    prefix_last_spec_by_reuser = {
+        spec.reuse_idx_in_batch: spec for spec in ctx.prefix_last_restore_indices
+    }
+
+    restored_reusers = 0
+    # Row 0 is always a provider (nothing precedes it to reuse), so start at 1.
+    # A reuser's provider always has a smaller batch index (online-detector
+    # invariant), so it is already restored when we reach reuser_idx.
+    for reuser_idx in range(1, len(prefix_lens)):
+        prefix_len = prefix_lens[reuser_idx]
+        if provider_index[reuser_idx] == reuser_idx or prefix_len <= 0:
+            continue  # provider / non-reuser: row already complete
+        provider_idx = provider_index[reuser_idx]
+
+        # interior [0, prefix_len-2]: bulk-copy from the provider's restored row.
+        if prefix_len - 1 > 0:
+            log_probs[reuser_idx, 0:prefix_len - 1] = log_probs[provider_idx, 0:prefix_len - 1]
+
+        # prefix-last (position prefix_len-1): recompute with the reuser's label.
+        prefix_last_spec = prefix_last_spec_by_reuser.get(reuser_idx)
+        if prefix_last_spec is not None:
+            saved_logits_key = (reuser_idx, prefix_last_spec.target_2d_pos)
+            saved_provider_logits = ctx.prefix_last_logits_saved[saved_logits_key]  # [1, V//tp]
+            reuser_label = torch.tensor(
+                [prefix_last_spec.label_value], dtype=torch.long, device=log_probs.device,
+            )  # [1]
+            log_probs[reuser_idx, prefix_len - 1] = vocab_parallel_log_probs_fn(
+                saved_provider_logits, reuser_label,
+            ).reshape(())
+        else:
+            # suffix_len == 0: no prefix-last spec; column is masked downstream.
+            log_probs[reuser_idx, prefix_len - 1] = log_probs[provider_idx, prefix_len - 1]
+
+        # entropy [0, prefix_len-1]: whole prefix copied (label-independent).
+        if entropy is not None:
+            entropy[reuser_idx, 0:prefix_len] = entropy[provider_idx, 0:prefix_len]
+
+        restored_reusers += 1
+
+    if ctx.stats is not None:
+        ctx.stats.record_restore(restored_reusers)
+    return output
+
+
+# ═══════════════════════════════════════════════════════════════
+# v080 restore 包装：NestedTensor → 2D left-pad → 复用 2D restore → 压回
+# ═══════════════════════════════════════════════════════════════
+
+
+def restore_via_2d_unfold_verl080(
+    output: dict,
+    vocab_parallel_log_probs_fn: Any,
+    vocab_parallel_entropy_fn: Any = None,
+) -> dict:
+    """v080 restore 包装：NestedTensor → 2D left-pad → 复用 restore_reuser_prefix_columns_2d → 压回。
+
+    v080 物理裁剪后 reuser NestedTensor 行只含 suffix 区段，prefix 区段（含
+    prefix-last）被物理删除。本函数在 forward_step 出口（context 仍激活、provider
+    prefix-last logits 已存于 ``ctx.prefix_last_logits_saved``）完成重组：
+
+    1. 展开裁剪后 NestedTensor 各行为完整 2D ``[B, L_max]``（reuser prefix 区段
+       left-pad 0，尾部 right-pad 0 到 L_max）
+    2. 复用 :func:`restore_reuser_prefix_columns_2d`：interior 整段从直接
+       provider 的已恢复 2D 行 bulk 切片复制，prefix-last 用存的 logits +
+       ``index.label_value`` 重算
+    3. 按各 ``original_lengths`` 切片压回 NestedTensor (jagged)
+
+    列映射为 identity：left-pad 后 valid-content 的 0-based 偏移即 2D 列号，
+    ``target_2d_pos`` 直接当列索引用，无需 ``valid_indices`` / 列映射表。
+
+    Must be called inside ``prefix_sharing_runtime_context`` (reads
+    ``current_prefix_sharing_context``), after the vocab_logprobs patch has saved
+    provider prefix-last logits into ``ctx.prefix_last_logits_saved``.
+
+    Args:
+        output: forward_step 返回的 output_dict，含 ``"log_probs"`` NestedTensor
+            （裁剪后 jagged），可选 ``"entropy"`` NestedTensor。**不含** tuple 外层
+            （tuple 解包由调用方负责）。
+        vocab_parallel_log_probs_fn: 用于 prefix-last logp 重算。
+        vocab_parallel_entropy_fn: 可选，当前未直接使用（entropy 走复制路径——
+            interior 和 prefix-last 都从 provider 复制，不重算）。
+
+    Returns:
+        ``output``（``log_probs``/``entropy`` 被替换为重组后的 NestedTensor）。
+    """
+
+    ctx = current_prefix_sharing_context()
+    if ctx is None:
+        return output
+    plan = ctx.prefix_sharing_plan
+    # Guard on reuser presence, not on prefix_last_restore_indices: a batch
+    # whose reusers all have suffix_len == 0 has no prefix-last spec but still
+    # needs interior prefix columns restored.
+    if not plan.has_sharing:
+        return output
+
+    log_probs_nested = output.get("log_probs")
+    if log_probs_nested is None or not is_nested_tensor(log_probs_nested):
+        return output
+    entropy_nested = output.get("entropy")
+    has_entropy = entropy_nested is not None and is_nested_tensor(entropy_nested)
+
+    original_lengths = plan.original_lengths
+    input_keep_ranges = plan.input_keep_ranges
+    B = len(original_lengths)
+    if B == 0:
+        return output
+    L_max = max(original_lengths)
+
+    # --- Step 1: 展开裁剪后 NestedTensor → 完整 2D [B, L_max] ---
+    log_probs_2d, entropy_2d = _unfold_trimmed_nested_to_2d(
+        log_probs_nested,
+        entropy_nested if has_entropy else None,
+        original_lengths,
+        input_keep_ranges,
+        L_max,
+        B,
+    )
+
+    # --- Step 2: 复用 restore_reuser_prefix_columns_2d ---
+    # build_kv 式区间拼接：interior 整段从直接 provider 的已恢复 2D 行切片，
+    # prefix-last 用 index.label_value + saved logits 重算。identity 列映射
+    # （target_2d_pos 即 2D 列号，无 left padding）。
+    output_2d: dict[str, Any] = {"log_probs": log_probs_2d}
+    if entropy_2d is not None:
+        output_2d["entropy"] = entropy_2d
+    output_2d = restore_reuser_prefix_columns_2d(
+        output_2d,
+        vocab_parallel_log_probs_fn,
+        vocab_parallel_entropy_fn,
+    )
+
+    # --- Step 3: 按各 original_lengths 压回 NestedTensor ---
+    output["log_probs"] = _fold_2d_to_nested(output_2d["log_probs"], original_lengths)
+    if entropy_2d is not None:
+        output["entropy"] = _fold_2d_to_nested(output_2d["entropy"], original_lengths)
+
+    num_prefix_last = len(ctx.prefix_last_restore_indices)
+    print(
+        f"[PS][restore_verl080] unfolded B={B} L_max={L_max}, "
+        f"restored reusers={num_prefix_last} (prefix-last entries; "
+        f"interior handled by bulk slice)",
+        flush=True,
+    )
+    return output
+
+
+def _unfold_trimmed_nested_to_2d(
+    log_probs_nested: Any,
+    entropy_nested: Any,
+    original_lengths: list[int],
+    input_keep_ranges: list,
+    L_max: int,
+    B: int,
+) -> tuple[Any, Any | None]:
+    """展开裁剪后 NestedTensor → 完整 2D [B, L_max]（reuser prefix left-pad 0）。
+
+    裁剪后各行：
+      - provider (keep_start=0): 完整 [prefix | suffix]，长度 = original_lengths[i]
+      - reuser  (keep_start=prefix_len>0): 仅 [suffix]，长度 = original_lengths[i]-prefix_len
+
+    展开后每行恢复成 [prefix_zeros | suffix]，再 right-pad 0 到 L_max。
+    left-pad 的 zeros 不在 autograd 图里，但 restore 会覆盖 prefix 区段（interior
+    复制 provider、prefix-last 重算），最终值在图里。right-pad 尾部在压回时丢弃。
+    """
+    log_probs_offsets = log_probs_nested.offsets()
+    log_probs_values = log_probs_nested.values()
+    if entropy_nested is not None:
+        entropy_offsets = entropy_nested.offsets()
+        entropy_values = entropy_nested.values()
+
+    log_probs_rows: list[Any] = []
+    entropy_rows: list[Any] | None = [] if entropy_nested is not None else None
+
+    for seq_idx in range(B):
+        orig_len = original_lengths[seq_idx]
+        prefix_len = input_keep_ranges[seq_idx][0]
+
+        log_probs_suffix = log_probs_values[log_probs_offsets[seq_idx]:log_probs_offsets[seq_idx + 1]]
+        log_probs_rows.append(_build_padded_row(log_probs_suffix, prefix_len, orig_len, L_max))
+
+        if entropy_nested is not None:
+            entropy_suffix = entropy_values[entropy_offsets[seq_idx]:entropy_offsets[seq_idx + 1]]
+            entropy_rows.append(_build_padded_row(entropy_suffix, prefix_len, orig_len, L_max))
+
+    log_probs_2d = torch.stack(log_probs_rows, dim=0)
+    entropy_2d = torch.stack(entropy_rows, dim=0) if entropy_rows else None
+    return log_probs_2d, entropy_2d
+
+
+def _build_padded_row(
+    suffix_data: Any, prefix_len: int, orig_len: int, L_max: int,
+) -> Any:
+    """构造一行完整 2D ``[prefix_zeros | suffix]`` right-pad 0 到 L_max。"""
+    device = suffix_data.device
+    dtype = suffix_data.dtype
+    tail_shape = tuple(suffix_data.shape[1:])
+    pieces: list[Any] = []
+    if prefix_len > 0:
+        pieces.append(torch.zeros((prefix_len,) + tail_shape, dtype=dtype, device=device))
+    pieces.append(suffix_data)
+    row = torch.cat(pieces, dim=0)  # [orig_len, ...]
+    if orig_len < L_max:
+        pad = torch.zeros((L_max - orig_len,) + tail_shape, dtype=dtype, device=device)
+        row = torch.cat([row, pad], dim=0)
+    return row
+
+
+def _fold_2d_to_nested(tensor_2d: Any, original_lengths: list[int]) -> Any:
+    """完整 2D [B, L_max] → NestedTensor (jagged)，按各 original_lengths 切片。"""
+    rows = [tensor_2d[seq_idx, :original_lengths[seq_idx]] for seq_idx in range(len(original_lengths))]
+    values = torch.cat(rows, dim=0) if rows else tensor_2d.reshape(0, *tensor_2d.shape[2:])
+    offsets = torch.tensor(
+        [0] + [sum(original_lengths[: idx + 1]) for idx in range(len(original_lengths))],
+        dtype=torch.long,
+        device=tensor_2d.device,
+    )
+    if hasattr(torch.nested, "nested_tensor_from_jagged"):
+        return torch.nested.nested_tensor_from_jagged(values, offsets)
+    return torch.nested.as_nested_tensor(rows, layout=torch.jagged)
diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/__init__.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/__init__.py
index 14e1518d..0855eed9 100644
--- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/__init__.py
+++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/__init__.py
@@ -1,15 +1,12 @@
-"""verl 0.8.0 FSDP patch set.
+"""module: prefix_sharing.setup.patches.verl080_fsdp
 
-Patch targets:
-1. FSDPEngineWithLMHead.forward_step → dense FSDP PrefixSharing forward helper
-2. Profiler injection for performance validation (without modifying verl source)
-3. Rollout / fixed-data injection helpers
-
-This patch set is the default entry point for the FSDP open-source line.
-It can be auto-selected via the compatibility matrix, or explicitly installed
-via ``prefix_sharing.setup.install("verl080_fsdp")``.
+verl 0.8.0 FSDP patch set. This is the default entry for the FSDP open-source line.
+It can be selected via the compat matrix, or installed explicitly with
+``prefix_sharing.setup.install("verl080_fsdp")``.
 """
 
+from __future__ import annotations
+
 from prefix_sharing.setup.patch_installer import PatchSpec
 
 from .forward_step import (
@@ -27,8 +24,11 @@
 
 PATCH_SET: list[PatchSpec] = [
     # ═══════════════════════════════════════════════════════════════
-    # 1. Fix / feature injection: PrefixSharing core functionality and debug helpers
+    # PrefixSharing entrypoints
     # ═══════════════════════════════════════════════════════════════
+
+    # FSDPEngineWithLMHead.forward_step → verl080_fsdp.patch_fsdp_forward_step
+    # patch forward_step() for PrefixSharing workflows
     PatchSpec(
         module_name="verl.workers.engine.fsdp.transformer_impl",
         target_getter=lambda mod: (
@@ -36,60 +36,91 @@
             "forward_step",
         ),
         patch_factory=patch_fsdp_forward_step,
-        description="FSDPEngineWithLMHead.forward_step → PrefixSharing dense FSDP helper",
-        eager=True,  # verl FSDP engine is lazy-loaded only when actor is instantiated; must trigger eagerly
+        description=("FSDPEngineWithLMHead.forward_step → verl080_fsdp.patch_fsdp_forward_step: "
+                     "patch forward_step() for PrefixSharing workflows"),
+        # verl FSDP engine is lazy-loaded only when the actor is created; eager is required.
+        eager=True,
     ),
+
+    # ═══════════════════════════════════════════════════════════════
+    # Debugging functionality, precision and performance
+    # ═══════════════════════════════════════════════════════════════
+
+    # RayPPOTrainer.fit → verl080_fsdp.patch_ray_trainer_fit
+    # patch fit() for capturing rollout results or loading fixed rollout data
     PatchSpec(
         module_name="verl.trainer.ppo.ray_trainer",
         target_getter=lambda mod: (mod.RayPPOTrainer, "fit"),
         patch_factory=patch_ray_trainer_fit,
         description=(
-            "RayPPOTrainer.fit → intercept actor_rollout_wg + async_rollout_manager "
-            "for PREFIX_SHARING_CAPTURE_ROLLOUT / PREFIX_SHARING_FIXED_ROLLOUT"
+            "RayPPOTrainer.fit → verl080_fsdp.patch_ray_trainer_fit: "
+            "patch fit() for capturing rollout results or loading fixed rollout data"
         ),
-        eager=True,  # ray_trainer is imported by main_ppo at startup; eager ensures patch is in place
+        # ray_trainer is imported by main_ppo at startup; eager ensures the patch is in place.
+        eager=True,
     ),
-    # ═══════════════════════════════════════════════════════════════
-    # 2. Performance validation: ProfilerScope layered profiling (replaces
-    #    invasive source modifications in the original verl codebase)
-    # ═══════════════════════════════════════════════════════════════
+
+    # TrainingWorker.train_mini_batch → step-level ProfilerScope (kind=train)
+    # patch train_mini_batch() for step-level performance profiling
+    PatchSpec(
+        module_name="verl.workers.engine_workers",
+        target_getter=lambda mod: (mod.TrainingWorker, "train_mini_batch"),
+        patch_factory=patch_train_mini_batch,
+        description=("TrainingWorker.train_mini_batch → verl080_fsdp.patch_train_mini_batch: "
+                     "patch train_mini_batch() for step-level performance profiling"),
+        eager=True,
+    ),
+
+    # BaseEngine.train_batch is called from train_mini_batch.
+    # patch train_batch() for optimizer-step performance profiling
     PatchSpec(
         module_name="verl.workers.engine.base",
         target_getter=lambda mod: (mod.BaseEngine, "train_batch"),
         patch_factory=patch_train_batch,
-        description="BaseEngine.train_batch → time optimizer_step as PHASE_UPDATE via ProfilerScope",
+        description=("BaseEngine.train_batch → verl080_fsdp.patch_train_batch: "
+                     "patch train_batch() for optimizer-step performance profiling"),
         eager=True,
     ),
+
+    # FSDPEngine.forward_backward_batch → micro-batch ProfilerScope markers
+    # patch forward_backward_batch() for micro-batch performance profiling
     PatchSpec(
         module_name="verl.workers.engine.fsdp.transformer_impl",
         target_getter=lambda mod: (mod.FSDPEngine, "forward_backward_batch"),
         patch_factory=patch_forward_backward_batch,
-        description="FSDPEngine.forward_backward_batch → micro-batch profiling without source edit",
+        description=("FSDPEngine.forward_backward_batch → verl080_fsdp.patch_forward_backward_batch: "
+                     "patch forward_backward_batch() for micro-batch performance profiling"),
         eager=True,
     ),
+
+    # FSDPEngine.forward_backward_batch → dump weight gradients after backward
+    # patch forward_backward_batch() for precision dumps
     PatchSpec(
         module_name="verl.workers.engine.fsdp.transformer_impl",
         target_getter=lambda mod: (mod.FSDPEngine, "forward_backward_batch"),
         patch_factory=patch_forward_backward_batch_for_diag_dump,
-        description="FSDPEngine.forward_backward_batch → dump weight gradients after backward",
-        eager=True,
-    ),
-    PatchSpec(
-        module_name="verl.workers.engine_workers",
-        target_getter=lambda mod: (mod.TrainingWorker, "train_mini_batch"),
-        patch_factory=patch_train_mini_batch,
-        description="TrainingWorker.train_mini_batch → step-level ProfilerScope (kind=train)",
+        description=("FSDPEngine.forward_backward_batch → "
+                     "verl080_fsdp.patch_forward_backward_batch_for_diag_dump: "
+                     "patch forward_backward_batch() for weight-gradient diagnostic dumps"),
         eager=True,
     ),
+
+    # TrainingWorker.infer_batch → verl080_fsdp.patch_infer_batch
+    # patch infer_batch() for step-level performance profiling
     PatchSpec(
         module_name="verl.workers.engine_workers",
         target_getter=lambda mod: (mod.TrainingWorker, "infer_batch"),
         patch_factory=patch_infer_batch,
-        description="TrainingWorker.infer_batch → step-level ProfilerScope (kind=logp)",
+        description=("TrainingWorker.infer_batch → verl080_fsdp.patch_infer_batch: "
+                     "patch infer_batch() for step-level performance profiling"),
         eager=True,
     ),
 ]
 
+# ═══════════════════════════════════════════════════════════════
+# PrefixSharing Attention
+# ═══════════════════════════════════════════════════════════════
+
 # Attention patch: directly modify ALL_ATTENTION_FUNCTIONS dict (same approach as verl's PrefixGrouper).
 # Cannot use PatchSpec because get_interface is not an attribute/key of AttentionInterface.
 from .attention import install_attention_patch
diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/attention.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/attention.py
index 32816da9..1cf1a5d7 100644
--- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/attention.py
+++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/attention.py
@@ -18,6 +18,8 @@
 import os
 from typing import Any
 
+from prefix_sharing.integrations.context import _current_context
+from prefix_sharing.integrations.verl_fsdp import PrefixSharingFSDPAttentionRuntime
 from prefix_sharing.tools.perf_profiler import PerfProfiler
 
 _SUPPORTED_ATTENTIONS = {
@@ -31,11 +33,9 @@
 # ##### [PS-diag] dump helpers ######
 
 def _resolve_num_layers(module: Any) -> int:
-    """Infer the total number of model layers from the attention module, with
-    a fallback to ``module.model.config``.
+    """从 attention module 推导模型总层数，有 ``module.model.config`` 回退。
 
-    ``_dump_attn_output`` uses the same fallback logic; extracted here as a
-    shared utility function.
+    ``_dump_attn_output`` 也用了同样的回退逻辑，此处抽取为公共函数。
     """
     num_layers = int(getattr(getattr(module, "config", None), "num_hidden_layers", 0) or 0)
     if num_layers == 0:
@@ -45,9 +45,7 @@ def _resolve_num_layers(module: Any) -> int:
 
 
 def _pack_off_dense_for_dump(tensor: Any) -> Any:
-    """OFF path: reshape dense [B,H,L,D] → [T,H,D] to match the packed input
-    convention of the dump functions.
-    """
+    """OFF 路径：将 dense [B,H,L,D] → [T,H,D] 以匹配 dump 函数的 packed 入参约定。"""
     import torch as _torch
     B, H, L, D = tensor.shape
     return tensor.transpose(1, 2).reshape(_torch.Size([B * L, H, D]))
@@ -90,15 +88,19 @@ def create_attention_wrapper(original_fn: Any) -> Any:
 
     def patched_attention(module: Any, query: Any, key: Any, value: Any,
                           attention_mask: Any, *args: Any, **kwargs: Any) -> Any:
-        # Prefer module attribute (AC recompute compatible), fall back to ContextVar
-        # (Megatron and other paths)
-        ctx = getattr(module, '_ps_ctx', None)
-        if ctx is None:
+
+        #########################################################
+        # STEP 1: get prefix sharing context
+        #########################################################
+
+        # 优先 module 属性（AC recompute 兼容），回退 ContextVar（Megatron 等路径）
+        prefix_sharing_context = getattr(module, '_prefix_sharing_context', None)
+        if prefix_sharing_context is None:
             from prefix_sharing.integrations.context import current_prefix_sharing_context
-            ctx = current_prefix_sharing_context()
+            prefix_sharing_context = current_prefix_sharing_context()
 
         # ── OFF path: no prefix sharing context → transparent passthrough ──
-        if ctx is None:
+        if prefix_sharing_context is None:
             # [PS-perf] start — OFF attention timing (cross-layer + per-layer) —
             profiler = PerfProfiler.current()
             _per_layer_ok = profiler is not None and getattr(profiler, "per_layer_enabled", False)
@@ -117,30 +119,30 @@ def patched_attention(module: Any, query: Any, key: Any, value: Any,
                     profiler.stop_phase(PerfProfiler.PHASE_ATTN_OFF)
             # [PS-perf] end ——————————————————————————————————————
 
-            # ##### [PS-diag] OFF per-layer dump (baseline / context inactive) #####
+            # ##### [PS-diag] OFF per-layer dump (baseline / context 不激活) #####
             if os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None:
                 _dump_attn_output(result, module)
                 _dump_off_rope_and_kv(module, query, key, value)
             # ##### [PS-diag] end #####
             return result
 
-        # ── ON path: route through PrefixSharing attention runtime ──
-        from prefix_sharing.integrations.verl_fsdp import PrefixSharingFSDPAttentionRuntime
-        from prefix_sharing.integrations.context import _current_context
+        #########################################################
+        # STEP 2: call PrefixSharing attention runtime
+        #########################################################
 
+        # ── ON path: route through PrefixSharing attention runtime ──
         layer_id = int(getattr(module, "layer_idx", 0) or 0)
         _num_layers = _resolve_num_layers(module)
-        runtime = PrefixSharingFSDPAttentionRuntime(layer_id=layer_id, num_layers=_num_layers)
+        attention_runtime = PrefixSharingFSDPAttentionRuntime(layer_id=layer_id, num_layers=_num_layers)
 
-        # HF attention interface expects [B, H, L, D]; runtime works in [B, L, H, D]
-        query_ld = query.transpose(1, 2)
-        key_ld = key.transpose(1, 2)
-        value_ld = value.transpose(1, 2)
+        # HF attention interface expects [B, H, L, D]; attention_runtime works in [B, L, H, D]
+        query = query.transpose(1, 2)
+        key = key.transpose(1, 2)
+        value = value.transpose(1, 2)
 
-        # runtime.forward() reads ContextVar internally; during AC recompute the
-        # ContextVar may have expired, but ctx from module._ps_ctx is still valid.
-        # Temporarily inject the ContextVar.
-        _ctxvar_token = _current_context.set(ctx)
+        # attention_runtime.forward() 内部读 ContextVar；AC recompute 时 ContextVar
+        # 可能已过期，但 prefix_sharing_context 来自 module._prefix_sharing_context 仍然有效。临时注入 ContextVar。
+        _ctxvar_token = _current_context.set(prefix_sharing_context)
         # [PS-perf] start — ON attention timing (attn.on = pack+kv+comp+unpack) —
         profiler = PerfProfiler.current()
         _per_layer_ok = profiler is not None and getattr(profiler, "per_layer_enabled", False)
@@ -149,7 +151,7 @@ def patched_attention(module: Any, query: Any, key: Any, value: Any,
         if _per_layer_ok:
             profiler.start_phase(f"attn.on.l{layer_id}")
         try:
-            output_ld = runtime.forward(None, query_ld, key_ld, value_ld)
+            output = attention_runtime.forward(None, query, key, value)
         finally:
             if _per_layer_ok:
                 _on_elapsed = profiler.stop_phase(f"attn.on.l{layer_id}")
@@ -159,10 +161,10 @@ def patched_attention(module: Any, query: Any, key: Any, value: Any,
             _current_context.reset(_ctxvar_token)
         # [PS-perf] end ————————————————————————————————————————
 
-        # ##### [PS-diag] ON attn output dump (context active = PS path) #####
+        # ##### [PS-diag] ON attn output dump（context 激活 = PS 路径） #####
         if os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None:
-            _dump_attn_output(output_ld, module)
-        return output_ld, None
+            _dump_attn_output(output, module)
+        return output, None
 
     return patched_attention
 
diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py
index d37b6e2a..e3c45820 100644
--- a/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py
+++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_fsdp/forward_step.py
@@ -1,22 +1,29 @@
-"""Patch: FSDPEngineWithLMHead.forward_step — verl 0.8.0 FSDP path.
-
-Thin wrapper: reads prefix_sharing_config, preferentially reuses the real
-engine's ``prepare_model_inputs`` / ``prepare_model_outputs``, injects a
-PrefixSharing runtime during forward, and performs interior / prefix-last
-restore on the output side.
-This patch covers dense 2D and verl remove-padding jagged NestedTensor
-formats; unverified formats such as Ulysses SP and fused kernels are
-explicitly rejected during config validation.
+"""patch: FSDPEngineWithLMHead.forward_step → verl080_fsdp.patch_fsdp_forward_step
+
+forward_step wrapper for PrefixSharing under verl 0.8.0 + FSDP.
 """
 
 from __future__ import annotations
 
 import os
+from contextlib import nullcontext
 from typing import Any, Callable
 
+import torch
+import torch.utils.checkpoint as _ckpt
+
+from prefix_sharing.core.config import PrefixSharingConfig
+from prefix_sharing.integrations.context import create_prefix_sharing_context
+from prefix_sharing.integrations.verl_fsdp import PrefixSharingFSDPAttentionRuntime
+from prefix_sharing.integrations.verl_fsdp import prepare_for_prefix_sharing_fsdp
+from prefix_sharing.integrations.verl_fsdp import restore_prefix_sharing_outputs_2d
+from prefix_sharing.integrations.verl_utils import restore_via_2d_unfold_verl080
+from prefix_sharing.integrations.verl_utils import is_nested_tensor
+from prefix_sharing.integrations.verl_utils import read_ps_config_from_engine_config
+from prefix_sharing.tools.perf_profiler import PerfProfiler, ProfilerScope
+
 
 def patch_fsdp_forward_step(original_forward_step: Any) -> Any:
-    """Create a patch wrapper for FSDPEngineWithLMHead.forward_step."""
 
     # Patch _CheckpointFrame.check_recomputed_tensors_match and
     # _internal_assert to no-op.
@@ -24,7 +31,6 @@ def patch_fsdp_forward_step(original_forward_step: Any) -> Any:
     # computation graph, causing the saved-tensor count mismatch detected by
     # these methods.  The recomputed values are numerically correct — the count
     # difference is benign.  Bypass both checks so ON-path training completes.
-    import torch.utils.checkpoint as _ckpt
     # Apply once, globally.
     if not getattr(patch_fsdp_forward_step, "_cp_patched", False):
         _ckpt._CheckpointFrame.check_recomputed_tensors_match = lambda self, gid: None  # type: ignore[method-assign]
@@ -33,22 +39,16 @@ def patch_fsdp_forward_step(original_forward_step: Any) -> Any:
         patch_fsdp_forward_step._cp_patched = True
 
     def patched_forward_step(self: Any, micro_batch: Any, loss_function: Any, forward_only: bool):
-        from prefix_sharing.core.config import PrefixSharingConfig
-        from prefix_sharing.integrations.verl_mcore import read_ps_config_from_engine_config
-        from prefix_sharing.tools.perf_profiler import PerfProfiler, ProfilerScope
 
-        raw_config = read_ps_config_from_engine_config(self.engine_config)
-        ps_config = PrefixSharingConfig.from_raw(raw_config)
+        raw_config = read_ps_config_from_engine_config(self.engine_config) # framework-level config
+        ps_config = PrefixSharingConfig.from_raw(raw_config) # PrefixSharing-level config
+        
         if not ps_config.enable_prefix_sharing:
             # Memory sampling is managed by the step-level ProfilerScope; this
             # path records only the model-forward phase.
             profiler = ProfilerScope.current()
             if profiler is not None:
-                forward_phase = (
-                    PerfProfiler.PHASE_FORWARD_OLD
-                    if forward_only
-                    else PerfProfiler.PHASE_FORWARD
-                )
+                forward_phase = PerfProfiler.PHASE_FORWARD_OLD if forward_only else PerfProfiler.PHASE_FORWARD
                 profiler.start_phase(forward_phase)
 
             # Without diagnostics this path must delegate directly to verl.
@@ -76,50 +76,35 @@ def patched_forward_step(self: Any, micro_batch: Any, loss_function: Any, forwar
 
                 micro_batch = micro_batch.to(get_device_id())
             except Exception:
-                # Local unit tests use plain dict / fake engine; no dependency on verl device helper.
+                # Local unit tests use plain dict / fake engine without verl device helpers.
                 pass
 
-        ulysses_sp_size = _read_runtime_value(
-            self.engine_config,
-            micro_batch,
-            "ulysses_sequence_parallel_size",
-            default=1,
-        )
-        use_fused_kernels = _read_runtime_value(
-            self.engine_config,
-            micro_batch,
-            "use_fused_kernels",
-            default=False,
-        )
+        model_config = {
+            "model_type": "text_only_causal_lm",
+            "ulysses_sequence_parallel_size": _read_runtime_value(self.engine_config, micro_batch, "ulysses_sequence_parallel_size", default=1),
+            "use_fused_kernels": _read_runtime_value(self.engine_config, micro_batch, "use_fused_kernels", default=False),
+        }
         ps_config.validate(
-            model_config={
-                "model_type": "text_only_causal_lm",
-                "ulysses_sequence_parallel_size": ulysses_sp_size,
-                "use_fused_kernels": use_fused_kernels,
-            },
+            model_config=model_config,
             integrate_mode="verl_fsdp",
         )
 
         if hasattr(self, "prepare_model_inputs") and hasattr(self, "prepare_model_outputs"):
             return _forward_step_with_engine_prepare(
-                self, micro_batch, loss_function, forward_only, ps_config,
+                self, micro_batch, loss_function, forward_only, ps_config, model_config,
             )
 
-        from prefix_sharing.integrations.verl_fsdp import forward_prefix_sharing_fsdp_micro_batch
+        from prefix_sharing.integrations.verl_fsdp import forward_step_without_engine_prepare
 
         calculate_entropy = bool(
             _read_runtime_value(self.engine_config, micro_batch, "calculate_entropy", default=False)
         )
         temperature = _read_temperature(micro_batch)
-        output = forward_prefix_sharing_fsdp_micro_batch(
+        output = forward_step_without_engine_prepare(
             micro_batch,
             self.module,
             ps_config,
-            model_config={
-                "model_type": "text_only_causal_lm",
-                "ulysses_sequence_parallel_size": ulysses_sp_size,
-                "use_fused_kernels": use_fused_kernels,
-            },
+            model_config=model_config,
             temperature=temperature,
             calculate_entropy=calculate_entropy,
             entropy_fn=getattr(self, "compute_entropy_from_logits", None),
@@ -159,15 +144,12 @@ def _forward_step_with_engine_prepare(
     loss_function: Any,
     forward_only: bool,
     ps_config: Any,
+    model_config: Any,
 ) -> Any:
-    import torch
-    from contextlib import nullcontext
-
-    from prefix_sharing.integrations.context import create_prefix_sharing_context
-    from prefix_sharing.integrations.verl_fsdp import PrefixSharingFSDPAttentionRuntime
-    from prefix_sharing.integrations.verl_fsdp import build_prefix_sharing_micro_batch_fsdp
 
-    from prefix_sharing.tools.perf_profiler import PerfProfiler, ProfilerScope
+    #########################################################
+    # STEP 1: pre-processing inputs for PrefixSharing
+    #########################################################
 
     profiler = ProfilerScope.current()
     if profiler is not None:
@@ -179,30 +161,17 @@ def _forward_step_with_engine_prepare(
     if os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None:
         _dump_full_input_ids_only(micro_batch, "train")
 
-    trimmed_micro_batch, ps_state = build_prefix_sharing_micro_batch_fsdp(
+    micro_batch_modified, prefix_sharing_runtime_state = prepare_for_prefix_sharing_fsdp(
         micro_batch,
         ps_config,
-        model_config={
-            "model_type": "text_only_causal_lm",
-            "ulysses_sequence_parallel_size": _read_runtime_value(
-                self.engine_config,
-                micro_batch,
-                "ulysses_sequence_parallel_size",
-                default=1,
-            ),
-            "use_fused_kernels": _read_runtime_value(
-                self.engine_config,
-                micro_batch,
-                "use_fused_kernels",
-                default=False,
-            ),
-        },
+        model_config=model_config,
     )
+
     if profiler is not None:
         profiler.stop_phase(PerfProfiler.PHASE_PLAN)  # Detect, plan, and trim on CPU.
-
-    if ps_state is None:
-        return _call_original_like_engine(self, trimmed_micro_batch, loss_function, forward_only)
+    
+    if prefix_sharing_runtime_state is None:
+        return _call_original_like_engine(self, micro_batch_modified, loss_function, forward_only)
 
     if os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None:
         from prefix_sharing.tools.diagnostic_dump import dump_fsdp_on_metadata_verl080
@@ -210,35 +179,42 @@ def _forward_step_with_engine_prepare(
         diagnostic_tag = "train" if self.module.training else "old"
         dump_fsdp_on_metadata_verl080(
             micro_batch,
-            ps_state.prefix_sharing_plan,
+            prefix_sharing_runtime_state.prefix_sharing_plan,
             diagnostic_tag,
         )
 
-    # Retrieve the number of model layers for per-layer diagnostic dump
+    #########################################################
+    # STEP 2: pre-processing inputs for the model
+    #########################################################
+
+    # Read layer count for per-layer diagnostic dumps.
     _diag_num_layers = int(getattr(
         getattr(getattr(self, "module", None), "config", None),
         "num_hidden_layers", 0)) or 0
 
-    model_inputs, output_args = self.prepare_model_inputs(micro_batch=trimmed_micro_batch)
+    model_inputs, output_args = self.prepare_model_inputs(micro_batch=micro_batch_modified)
     model_inputs["prefix_sharing_runtime"] = PrefixSharingFSDPAttentionRuntime()
     model_inputs["prefix_sharing_runtime"].num_layers = _diag_num_layers
     autocast_dtype = getattr(self, "_autocast_dtype", torch.float32)
-    device_name = _read_device_name()
     autocast_ctx = (
         nullcontext()
         if autocast_dtype == torch.float32
-        else torch.autocast(device_type=device_name, dtype=autocast_dtype)
+        else torch.autocast(device_type=_read_device_name(), dtype=autocast_dtype)
     )
-    # ── Create PS context with manual lifecycle (survives backward for AC) ──
-    ctx, ctx_cleanup = create_prefix_sharing_context(ps_state)
 
-    # Set _ps_ctx on every attention module so the attention patch reads
+    #########################################################
+    # STEP 3: create PrefixSharing runtime context
+    #########################################################
+
+    prefix_sharing_context, cleanup_prefix_sharing_context = create_prefix_sharing_context(prefix_sharing_runtime_state)
+
+    # Set _prefix_sharing_context on every attention module so the attention patch reads
     # the context from the module itself rather than ContextVar (compatible
     # with activation-checkpointing recompute, which bypasses the context
     # manager that set the ContextVar).
     for attention_module in self.module.modules():
         if hasattr(attention_module, "layer_idx") and hasattr(attention_module, "q_proj"):
-            attention_module._ps_ctx = ctx
+            attention_module._prefix_sharing_context = prefix_sharing_context
 
     # Register diagnostic gradient hooks when the diagnostic dump is enabled.
     _register_grad_dump_hooks(self.module, forward_only)
@@ -246,15 +222,15 @@ def _forward_step_with_engine_prepare(
     # Attach cleanup callback so the forward_backward_batch wrapper can release
     # PrefixSharing state after backward.  The root full-backward hook is
     # unreliable here because ``self.module`` returns a CausalLMOutput dataclass.
-    self.module._ps_ctx_cleanup = ctx_cleanup
+    self.module._cleanup_prefix_sharing_context = cleanup_prefix_sharing_context
+
+    #########################################################
+    # STEP 4: call the model with PrefixSharing+FSDP attention runtime
+    #########################################################
 
     with autocast_ctx:
         if profiler is not None:
-            forward_phase = (
-                PerfProfiler.PHASE_FORWARD_OLD
-                if forward_only
-                else PerfProfiler.PHASE_FORWARD
-            )
+            forward_phase = PerfProfiler.PHASE_FORWARD_OLD if forward_only else PerfProfiler.PHASE_FORWARD
             profiler.start_phase(forward_phase)
         raw_output = self.module(**model_inputs, use_cache=False)
         if profiler is not None:
@@ -265,14 +241,22 @@ def _forward_step_with_engine_prepare(
 
             dump_raw_logits_verl080(raw_output, dp_aware=_get_dp_size() > 1)
 
+        #########################################################
+        # STEP 5: post-processing outputs for the model
+        #########################################################
+
         _save_prefix_last_logits_from_raw_output(raw_output)
         model_output = self.prepare_model_outputs(
             output=raw_output,
             output_args=output_args,
-            micro_batch=trimmed_micro_batch,
+            micro_batch=micro_batch_modified,
             logits_processor_func=loss_function,
         )
 
+        #########################################################
+        # STEP 6: post-processing outputs for PrefixSharing
+        #########################################################
+
         if profiler is not None:
             profiler.start_phase(PerfProfiler.PHASE_RESTORE)
         model_output = _restore_engine_model_output(model_output)
@@ -284,10 +268,14 @@ def _forward_step_with_engine_prepare(
 
             dump_fsdp_model_output_2d_verl080(
                 model_output,
-                list(ps_state.prefix_sharing_plan.original_lengths),
+                list(prefix_sharing_runtime_state.prefix_sharing_plan.original_lengths),
                 diagnostic_tag,
             )
 
+        #########################################################
+        # STEP 7: compute loss & metrics
+        #########################################################
+
         if loss_function is not None:
             if profiler is not None:
                 profiler.start_phase(PerfProfiler.PHASE_LOSS)
@@ -361,10 +349,9 @@ def _call_original_like_engine(self: Any, micro_batch: Any, loss_function: Any,
     import torch
     from contextlib import nullcontext
 
-    # Align with verl's native forward_step: move micro_batch to device first
-    # (the disable path bypasses the .to(device) in patched_forward_step, so we
-    # compensate here to avoid device mismatch for logits/temperature in
-    # prepare_model_outputs).
+    # Match native verl forward_step: move micro_batch to device first.
+    # The disable / no-sharing path skips the .to(device) in patched_forward_step;
+    # without it, logits/temperature can land on different devices in prepare_model_outputs.
     if hasattr(micro_batch, "to"):
         try:
             from verl.utils.device import get_device_id
@@ -373,13 +360,11 @@ def _call_original_like_engine(self: Any, micro_batch: Any, loss_function: Any,
             pass
     model_inputs, output_args = self.prepare_model_inputs(micro_batch=micro_batch)
 
-    # DIAG_DUMP: ON path dumps original full input_ids (suffix-only dump would
-    # miss prefix tokens)
+    # DIAG_DUMP: dump original full input_ids (suffix-only dumps miss prefix tokens).
     import os as _ps_diag_fwd_ids
     if _ps_diag_fwd_ids.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None:
         _dump_full_input_ids_only(micro_batch, "train")
 
-    autocast_dtype = getattr(self, "_autocast_dtype", torch.float32)
     autocast_dtype = getattr(self, "_autocast_dtype", torch.float32)
     device_name = _read_device_name()
     autocast_ctx = (
@@ -394,11 +379,7 @@ def _call_original_like_engine(self: Any, micro_batch: Any, loss_function: Any,
 
     profiler = ProfilerScope.current()
     if profiler is not None:
-        forward_phase = (
-            profiler.PHASE_FORWARD_OLD
-            if forward_only
-            else profiler.PHASE_FORWARD
-        )
+        forward_phase = profiler.PHASE_FORWARD_OLD if forward_only else profiler.PHASE_FORWARD
         profiler.start_phase(forward_phase)
 
     with autocast_ctx:
@@ -470,17 +451,13 @@ def _restore_engine_model_output(model_output: dict[str, Any]) -> dict[str, Any]
     except Exception:
         entropy_from_logits = None
 
-    from prefix_sharing.integrations.verl_mcore import restore_via_2d_unfold_verl080
-    from prefix_sharing.integrations.verl_mcore import _is_nested_tensor
-    from prefix_sharing.integrations.verl_fsdp import restore_prefix_sharing_outputs_2d
-
     restored = restore_via_2d_unfold_verl080(
         model_output,
         logprobs_from_logits,
         entropy_from_logits,
     )
     log_probs = restored.get("log_probs")
-    if log_probs is not None and not _is_nested_tensor(log_probs):
+    if log_probs is not None and not is_nested_tensor(log_probs):
         return restore_prefix_sharing_outputs_2d(restored, logprobs_from_logits)
     return restored
 
@@ -540,7 +517,7 @@ def _read_temperature(micro_batch: Any) -> float:
 def _dump_full_input_ids_only(micro_batch: Any, tag: str) -> None:
     """Dump the original (full) input_ids before prefix sharing trimming.
 
-    The ON path dumps ``input_ids_train.pt`` from the ``trimmed_micro_batch``,
+    The ON path dumps ``input_ids_train.pt`` from the ``micro_batch_modified``,
     which has shared prefix tokens removed.  This helper saves the **original**
     ``micro_batch`` input_ids so that ``cmp_diag_verl080`` can compare the
     full input against the OFF baseline, rather than reporting 186+ differing
@@ -588,6 +565,16 @@ def patch_forward_backward_batch_for_diag_dump(
     def wrapped(self: Any, data: Any, loss_function: Any, forward_only: bool = False) -> Any:
         result = original_forward_backward_batch(self, data, loss_function, forward_only)
 
+        # 无条件清理 PrefixSharing runtime context：打印 audit 日志 + 关闭 KV store。
+        # 之前该清理被误关在 PREFIX_SHARING_DIAG_DUMP 条件块内，导致正常训练时
+        # audit 日志不输出、KV store 不 close（多步训练存在内存累积风险）。
+        if not forward_only:
+            cleanup_prefix_sharing_context = getattr(self.module, "_cleanup_prefix_sharing_context", None)
+            if cleanup_prefix_sharing_context is not None:
+                cleanup_prefix_sharing_context()
+                delattr(self.module, "_cleanup_prefix_sharing_context")
+
+        # 诊断 dump 专用：dump weight gradients + 清理 per-layer attention grad hooks。
         if not forward_only and os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None:
             tag = "train" if self.module.training else "old"
             print(
@@ -599,16 +586,10 @@ def wrapped(self: Any, data: Any, loss_function: Any, forward_only: bool = False
 
             dump_weight_grads_verl080(self.module, tag)
 
-            # Clean up PrefixSharing context if the root backward hook did not fire.
-            ctx_cleanup = getattr(self.module, "_ps_ctx_cleanup", None)
-            if ctx_cleanup is not None:
-                ctx_cleanup()
-                delattr(self.module, "_ps_ctx_cleanup")
-
             # Remove per-layer attention gradient hooks.
             for module in self.module.modules():
                 try:
-                    del module._ps_ctx
+                    del module._prefix_sharing_context
                 except AttributeError:
                     pass
 
diff --git a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py
index 911de472..7f9ed666 100644
--- a/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py
+++ b/prefix-sharing/prefix_sharing/setup/patches/verl080_mcore0161_ms0160/forward_step.py
@@ -1,11 +1,9 @@
-"""Patch: MegatronEngineWithLMHead.forward_step — verl 0.8.0 engine architecture
+"""Patch: MegatronEngineWithLMHead.forward_step — verl 0.8.0 engine 架构
 
-Thin wrapper: consume batch → read config → build state → set context → feed
-back to original forward_step.
+thin wrapper：消费 batch → 读 config → 构建状态 → 设 context → 喂回原始 forward_step。
 
-All business logic (config reading, batch construction, layout computation)
-is handled by the integrations layer. This patch only orchestrates the call
-sequence and sets up the runtime context.
+所有业务逻辑（config 读取、batch 构建、layout 计算）由 integrations 层处理，
+本 patch 只负责编排调用顺序和设置 runtime context。
 """
 
 from __future__ import annotations
@@ -55,7 +53,7 @@ def _describe_batch(batch: Any) -> str:
 
 
 def patch_verl_forward_step(original_forward_step: Any) -> Any:
-    """Create a patch wrapper for MegatronEngineWithLMHead.forward_step."""
+    """创建 MegatronEngineWithLMHead.forward_step 的 patch wrapper。"""
 
     def patched_forward_step(
         self,
@@ -64,18 +62,17 @@ def patched_forward_step(
         logits_processor_func,
         postprocess_micro_batch_func,
     ):
-        # ── Retrieve original micro-batch ──
-        # batch_iter comes from the outer engine's forward_step caller.
-        # Consume the batch; build_prefix_sharing_micro_batch_verl080 performs
-        # physical trimming. Returns trimmed_batch (physically trimmed
-        # micro-batch) and ps_state.
+        # ── 获取原始 micro-batch ──
+        # batch_iter 来自外层 engine 的 forward_step 调用方。
+        # 消费 batch，由 build_prefix_sharing_micro_batch_verl080 进行物理裁剪
+        # 返回 trimmed_batch（物理裁剪后的 micro-batch）和 ps_state。
         _ps_forward_step_probe("enter")
         _ps_forward_step_probe("before_next_batch")
         original_batch = next(batch_iter)
         _ps_forward_step_probe("after_next_batch", batch=_describe_batch(original_batch))
         batch_for_forward = original_batch
 
-        # ── Read configuration ──
+        # ── 读取配置 ──
         _ps_forward_step_probe("before_read_config")
         from prefix_sharing.integrations.verl_mcore import read_ps_config_from_engine_config
         from prefix_sharing.core.config import PrefixSharingConfig
@@ -89,16 +86,15 @@ def patched_forward_step(
 
         ps_state = None
         if ps_config.enable_prefix_sharing:
-            # batch.to(device) ensures tensors are on the target device.
-            # The original forward_step will call batch.to(device) again
-            # (idempotent).
+            # batch.to(device) 使 tensor 在目标设备上，
+            # 原始 forward_step 会再次 batch.to(device)（幂等）
             from verl.utils.megatron_utils import get_device_id
             device_id = get_device_id()
             _ps_forward_step_probe("before_batch_to_device", device_id=device_id)
             batch_on_device = original_batch.to(device_id)
             _ps_forward_step_probe("after_batch_to_device")
 
-            # Batch trimming
+            # batch裁剪
             _ps_forward_step_probe("before_prepare_micro_batch")
             from prefix_sharing.integrations.verl_mcore import build_prefix_sharing_micro_batch_verl080
             batch_for_forward, ps_state = build_prefix_sharing_micro_batch_verl080(
@@ -121,13 +117,11 @@ def patched_forward_step(
         else:
             _ps_forward_step_probe("skip_prepare_prefix_sharing_disabled")
 
-        # ##### [PS-diag] dump metadata + attention_mask + label_mask (shared by ON/OFF) #####
-        # Only triggered when PREFIX_SHARING_DIAG_DUMP is set; zero overhead otherwise.
-        # ON: prefix_lens / original_lengths are taken from the plan;
-        # OFF: prefix_lens are all-zero, original_lengths derived from
-        #      input_ids NestedTensor offsets diff.
-        # cu_seqlens are taken from the input_ids NestedTensor offsets fed into
-        # forward (ON = trimmed packed boundaries, OFF = full).
+        # ##### [PS-diag] dump 元数据 + attention_mask + label_mask（ON/OFF 通用） #####
+        # 仅当 PREFIX_SHARING_DIAG_DUMP 设定时触发，否则零开销。
+        # ON: prefix_lens / original_lengths 取自 plan；
+        # OFF: prefix_lens 全0、original_lengths 从 input_ids NestedTensor offsets diff 推。
+        # cu_seqlens 取送进 forward 的 input_ids NestedTensor offsets（ON=裁剪后 packed 边界, OFF=完整）。
         import os as _os
         if _os.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None:
             from prefix_sharing.tools.diagnostic_dump import (
@@ -136,9 +130,9 @@ def patched_forward_step(
                 build_attention_mask_2d, build_label_mask_2d,
                 nested_offsets_to_cu,
             )
-            from prefix_sharing.integrations.verl_mcore import _is_nested_tensor
+            from prefix_sharing.integrations.verl_utils import is_nested_tensor
             _ids_nested = batch_for_forward["input_ids"]
-            if _is_nested_tensor(_ids_nested):
+            if is_nested_tensor(_ids_nested):
                 if ps_state is not None:
                     _plan = ps_state.prefix_sharing_plan
                     _prefix_lens = list(_plan.prefix_lens)
@@ -148,47 +142,40 @@ def patched_forward_step(
                     _orig_lens = [int(d) for d in _diffs]
                     _prefix_lens = [0] * len(_orig_lens)
                 dump_meta_verl080(_prefix_lens, nested_offsets_to_cu(_ids_nested))
-                # attention_mask + label_mask: two comparison ranges for log_probs,
-                # both excluding the out-of-bounds prediction position (POS L_i-1,
-                # whose logp predicts the non-existent token[L_i]). Aligned to the
-                # compact [B, L_max] coordinate system of restored log_probs.
-                #   attention_mask: [0:L_i-1) prompt region + prompt-last + response
-                #                   region (full restore verification)
-                #   label_mask:     [prompt-last:L_i-1) prompt-last + response region
-                #                   (PPO loss scope)
+                # attention_mask + label_mask：两种 log_probs 对比范围，都不含越界预测位
+                # （POS L_i-1，其 logp 预测不存在的 token[L_i]）。对齐 restore 后 log_probs
+                # 的 [B, L_max] 紧凑坐标系。
+                #   attention_mask：[0:L_i-1) prompt 区+prompt-last+response 区（整体 restore 验证）
+                #   label_mask：[prompt-last:L_i-1) prompt-last+response 区（PPO loss 范围）
                 _Lmax_lm = max(_orig_lens) if _orig_lens else 0
-                # tag aligned with logprobs (entry point 2 uses model.training to
-                # distinguish old/train), ensuring mask and logprobs_{tag} come from
-                # the same forward pass (same batch, same L_max).
+                # tag 与 logprobs 一致（接入点2 用 model.training 区分 old/train），
+                # 保证 mask 和 logprobs_{tag} 来自同一 forward（同 batch、同 L_max）。
                 _tag_lm = "train" if model.training else "old"
-                # attention_mask only depends on _orig_lens; no loss_mask needed.
+                # attention_mask 仅依赖 _orig_lens，不需要 loss_mask。
                 dump_attention_mask_verl080(
                     build_attention_mask_2d(_orig_lens, _Lmax_lm), _tag_lm)
-                # label_mask uses response_lens (number of response tokens per row).
-                # After verl080 padding, loss_mask = response_mask is 2D left-right
-                # padded (not a NestedTensor, see verl padding.py:71), so it cannot
-                # go through nested_to_2d_full; however the response token count =
-                # loss_mask row sum is coordinate-system agnostic, making it the
-                # most robust way to derive prompt_len (works for both 2D and
-                # NestedTensor).
+                # label_mask 用 response_lens（每行 response token 数）。verl080 padding 后
+                # loss_mask = response_mask 是 2D left-right padded（非 NestedTensor，见
+                # verl padding.py:71），不能走 nested_to_2d_full；但 response token 数 =
+                # loss_mask 行 sum，与坐标系无关，据此推 prompt_len 最稳（2D/NestedTensor 均适用）。
                 _lm = original_batch.get("loss_mask")
                 if _lm is not None:
-                    if _is_nested_tensor(_lm):
+                    if is_nested_tensor(_lm):
                         _lm_off = _lm.offsets()
                         _lm_val = _lm.values()
                         _response_lens = [
                             int(_lm_val[_lm_off[i]:_lm_off[i + 1]].sum())
                             for i in range(len(_orig_lens))]
                     else:
-                        # .long() avoids importing torch (not imported at file top);
-                        # .cpu() guards against on-device tensor .tolist()
+                        # .long() 免 import torch（本文件顶部未导入 torch）；
+                        # .cpu() 防御 on-device tensor 的 tolist()
                         _response_lens = _lm.sum(dim=-1).long().cpu().tolist()
                     dump_label_mask_verl080(
                         build_label_mask_2d(_response_lens, _orig_lens, _Lmax_lm),
                         _tag_lm)
-        # ##### [PS-diag] dump metadata + masks end #####
+        # ##### [PS-diag] dump 元数据 + masks end #####
 
-        # ── Build modified iterator to feed back into original forward_step ──
+        # ── 构造修改后的 iterator 喂回原始 forward_step ──
         modified_iter = iter([batch_for_forward])
 
         # ── runtime context ──
@@ -210,14 +197,12 @@ def patched_forward_step(
                 logits_processor_func,
                 postprocess_micro_batch_func,
             )
-            # v080 restore: reassemble reuser prefix segments while context is
-            # still active. forward_step returns (output_dict,
-            # partial(postprocess_func)); unpack, process output_dict, then
-            # repackage. restore_via_2d_unfold_verl080 internally checks
-            # context / restore_indices and does an early return when no
-            # restore is needed.
+            # v080 restore：在 context 仍激活时重组 reuser prefix 区段。
+            # forward_step 返回 (output_dict, partial(postprocess_func))，
+            # 解包处理 output_dict 再重包。restore_via_2d_unfold_verl080 内部
+            # 会检查 context / restore_indices，无 restore 需求时 early return。
             if ps_state is not None:
-                from prefix_sharing.integrations.verl_mcore import restore_via_2d_unfold_verl080
+                from prefix_sharing.integrations.verl_utils import restore_via_2d_unfold_verl080
                 from prefix_sharing.integrations.context import current_prefix_sharing_context
                 from verl.utils.megatron.tensor_parallel import (
                     vocab_parallel_entropy,
@@ -229,36 +214,30 @@ def patched_forward_step(
                     vocab_parallel_log_probs_from_logits,
                     vocab_parallel_entropy,
                 )
-                # Release vocab-dimension logits (high memory usage; only held
-                # during context lifetime — restore has already consumed them).
-                # The clear responsibility belongs here, not in the wrapper
-                # function.
+                # 释放 vocab 维 logits（占用大，只在 context 生命周期内持有，
+                # restore 已消费完毕）。clear 职责在此，不在包装函数内。
                 ctx = current_prefix_sharing_context()
                 if ctx is not None:
                     ctx.prefix_last_logits_saved.clear()
                 output = (output_dict, postprocess_fn)
-            # ##### [PS-diag] dump 2D logprobs/entropy (ON=post-restore, OFF=original) #####
-            # After restore (ON) or from original forward (OFF), log_probs / entropy
-            # are both NestedTensors with per-row length = original_lengths[i].
-            # Expand to a uniform [B, L_max] for cmp_diag.cmp_2d element-wise
-            # comparison.
+            # ##### [PS-diag] dump 2D logprobs/entropy（ON=restore后, OFF=原始） #####
+            # restore 后（ON）或原始 forward（OFF）的 log_probs/entropy 都是 NestedTensor，
+            # 每行长度 = original_lengths[i]，展开到统一 [B, L_max] 供 cmp_diag.cmp_2d 逐元素对比。
             import os as _os2
             if _os2.environ.get("PREFIX_SHARING_DIAG_DUMP") is not None:
                 from prefix_sharing.tools.diagnostic_dump import (
                     nested_to_2d_full, dump_logprobs_2d_verl080, dump_entropy_2d_verl080,
                 )
-                from prefix_sharing.integrations.verl_mcore import _is_nested_tensor
-                # Aligned with v070: tag = "old" if forward_only else "train".
-                # forward_step does not have access to forward_only, so use
-                # model.training as an equivalent distinction:
-                #   eval_mode → training=False → "old"  (old_logp phase)
-                #   train_mode → training=True → "train" (update_actor phase)
-                # A single run thus produces logprobs_old + logprobs_train without
-                # overwriting each other.
+                from prefix_sharing.integrations.verl_utils import is_nested_tensor
+                # 对齐 v070: tag = "old" if forward_only else "train"。
+                # forward_step 拿不到 forward_only，用 model.training 等价区分
+                # （eval_mode→training=False→"old" 对应 old_logp 阶段；
+                #  train_mode→training=True→"train" 对应 update_actor 阶段）。
+                # 这样一次 run 自动产出 logprobs_old + logprobs_train 两份，不互相覆盖。
                 _tag = "train" if model.training else "old"
                 _out_dict, _ = output
                 _lp = _out_dict.get("log_probs")
-                if _is_nested_tensor(_lp):
+                if is_nested_tensor(_lp):
                     if ps_state is not None:
                         _ol = list(ps_state.prefix_sharing_plan.original_lengths)
                     else:
@@ -266,7 +245,7 @@ def patched_forward_step(
                     _Lmax = max(_ol) if _ol else 0
                     dump_logprobs_2d_verl080(nested_to_2d_full(_lp, _ol, _Lmax), _tag)
                     _ent = _out_dict.get("entropy")
-                    if _is_nested_tensor(_ent):
+                    if is_nested_tensor(_ent):
                         dump_entropy_2d_verl080(nested_to_2d_full(_ent, _ol, _Lmax), _tag)
             # ##### [PS-diag] dump 2D logprobs/entropy end #####
         _ps_forward_step_probe("after_original_forward_step")
diff --git a/prefix-sharing/tests/integrated_test/optional/test_verl_megatron_runtime_helpers.py b/prefix-sharing/tests/integrated_test/optional/test_verl_megatron_runtime_helpers.py
index e6285ead..b736002c 100644
--- a/prefix-sharing/tests/integrated_test/optional/test_verl_megatron_runtime_helpers.py
+++ b/prefix-sharing/tests/integrated_test/optional/test_verl_megatron_runtime_helpers.py
@@ -18,7 +18,7 @@
 from prefix_sharing.integrations.megatron_runtime import prefix_attention
 from prefix_sharing.integrations.runtime_state import PrefixSharingRuntimeState
 from prefix_sharing.integrations.parallel_info import MegatronParallelInfo
-from prefix_sharing.integrations.verl_mcore import restore_reuser_prefix_columns_2d
+from prefix_sharing.integrations.verl_utils import restore_reuser_prefix_columns_2d
 
 
 def _make_state() -> tuple[PrefixSharingRuntimeState, list]:
diff --git a/prefix-sharing/tests/integrated_test/test_patch_integrations.py b/prefix-sharing/tests/integrated_test/test_patch_integrations.py
index 45ae2bff..23c83b9f 100644
--- a/prefix-sharing/tests/integrated_test/test_patch_integrations.py
+++ b/prefix-sharing/tests/integrated_test/test_patch_integrations.py
@@ -50,9 +50,21 @@ def test_setup_can_load_explicit_verl080_fsdp_patch_set():
     assert "FSDPEngineWithLMHead.forward_step" in patch_set[0].description
     assert patch_set[1].module_name == "verl.trainer.ppo.ray_trainer"
     assert "RayPPOTrainer.fit" in patch_set[1].description
-    # Diag-dump wrapper is installed on top of the profiler wrapper.
-    assert patch_set[4].module_name == "verl.workers.engine.fsdp.transformer_impl"
-    assert "dump weight gradients" in patch_set[4].description
+    forward_backward_specs = [
+        spec
+        for spec in patch_set
+        if spec.module_name == "verl.workers.engine.fsdp.transformer_impl"
+        and "forward_backward_batch" in spec.description
+    ]
+    # The diagnostic wrapper must wrap the profiler wrapper.
+    assert [spec.description for spec in forward_backward_specs] == [
+        "FSDPEngine.forward_backward_batch → "
+        "verl080_fsdp.patch_forward_backward_batch: "
+        "patch forward_backward_batch() for micro-batch performance profiling",
+        "FSDPEngine.forward_backward_batch → "
+        "verl080_fsdp.patch_forward_backward_batch_for_diag_dump: "
+        "patch forward_backward_batch() for weight-gradient diagnostic dumps",
+    ]
 
 
 
diff --git a/prefix-sharing/tests/integrated_test/test_verl080_restore_e2e.py b/prefix-sharing/tests/integrated_test/test_verl080_restore_e2e.py
index 53616609..c97c6d96 100644
--- a/prefix-sharing/tests/integrated_test/test_verl080_restore_e2e.py
+++ b/prefix-sharing/tests/integrated_test/test_verl080_restore_e2e.py
@@ -25,10 +25,8 @@
 megatron = pytest.importorskip("megatron")
 torch = pytest.importorskip("torch")
 
-from prefix_sharing.integrations.verl_mcore import (
-    build_prefix_sharing_micro_batch_verl080,
-    restore_via_2d_unfold_verl080,
-)
+from prefix_sharing.integrations.verl_mcore import build_prefix_sharing_micro_batch_verl080
+from prefix_sharing.integrations.verl_utils import restore_via_2d_unfold_verl080
 
 
 # ═══════════════════════════════════════
diff --git a/prefix-sharing/tests/unit_test/test_restore_unfold_verl080.py b/prefix-sharing/tests/unit_test/test_restore_unfold_verl080.py
index f3cc5402..1eadcefb 100644
--- a/prefix-sharing/tests/unit_test/test_restore_unfold_verl080.py
+++ b/prefix-sharing/tests/unit_test/test_restore_unfold_verl080.py
@@ -22,8 +22,8 @@
 from prefix_sharing.core.planner import PrefixSharingPlanner
 from prefix_sharing.integrations.context import prefix_sharing_runtime_context
 from prefix_sharing.integrations.parallel_info import MegatronParallelInfo
-from prefix_sharing.integrations.verl_mcore import (
-    PrefixSharingRuntimeState,
+from prefix_sharing.integrations.runtime_state import PrefixSharingRuntimeState
+from prefix_sharing.integrations.verl_utils import (
     _fold_2d_to_nested,
     _unfold_trimmed_nested_to_2d,
     restore_via_2d_unfold_verl080,
diff --git a/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py b/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py
index 1b9a558d..4c694c40 100644
--- a/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py
+++ b/prefix-sharing/tests/unit_test/test_verl_fsdp_adapter.py
@@ -9,8 +9,8 @@
 from prefix_sharing.integrations.context import prefix_sharing_runtime_context
 from prefix_sharing.integrations.verl_fsdp import (
     PrefixSharingFSDPAttentionRuntime,
-    build_prefix_sharing_micro_batch_fsdp,
-    forward_prefix_sharing_fsdp_micro_batch,
+    prepare_for_prefix_sharing_fsdp,
+    forward_step_without_engine_prepare,
     restore_prefix_sharing_outputs_2d,
 )
 from prefix_sharing.integrations.verl_mcore import PrefixSharingRuntimeState
@@ -132,7 +132,7 @@ def _baseline_attention(query, key, value):
     return torch.einsum("blmh,bmhd->blhd", probs, value)
 
 
-def test_build_prefix_sharing_micro_batch_fsdp_returns_trimmed_batch_and_runtime_state():
+def test_prepare_for_prefix_sharing_fsdp_returns_trimmed_batch_and_runtime_state():
     config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3)
     batch = {
         "input_ids": torch.tensor(
@@ -172,7 +172,7 @@ def test_build_prefix_sharing_micro_batch_fsdp_returns_trimmed_batch_and_runtime
         ),
     }
 
-    trimmed_batch, runtime_state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    trimmed_batch, runtime_state = prepare_for_prefix_sharing_fsdp(batch, config)
 
     assert runtime_state is not None
     assert isinstance(runtime_state, PrefixSharingRuntimeState)
@@ -193,7 +193,7 @@ def test_build_prefix_sharing_micro_batch_fsdp_returns_trimmed_batch_and_runtime
     assert torch.equal(trimmed_batch["position_ids"][1, 3:6], torch.tensor([3, 4, 5]))
 
 
-def test_build_prefix_sharing_micro_batch_fsdp_returns_none_when_no_sharing():
+def test_prepare_for_prefix_sharing_fsdp_returns_none_when_no_sharing():
     config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3)
     batch = {
         "input_ids": torch.tensor([[1, 2, 3], [4, 5, 6]], dtype=torch.long),
@@ -201,13 +201,13 @@ def test_build_prefix_sharing_micro_batch_fsdp_returns_none_when_no_sharing():
         "position_ids": torch.tensor([[0, 1, 2], [0, 1, 2]], dtype=torch.long),
     }
 
-    returned_batch, runtime_state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    returned_batch, runtime_state = prepare_for_prefix_sharing_fsdp(batch, config)
 
     assert returned_batch is batch
     assert runtime_state is None
 
 
-def test_build_prefix_sharing_micro_batch_fsdp_trims_nested_remove_padding_batch():
+def test_prepare_for_prefix_sharing_fsdp_trims_nested_remove_padding_batch():
     if not hasattr(torch, "nested"):
         pytest.skip("torch.nested is unavailable")
     config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3)
@@ -235,7 +235,7 @@ def test_build_prefix_sharing_micro_batch_fsdp_trims_nested_remove_padding_batch
         ),
     }
 
-    trimmed_batch, runtime_state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    trimmed_batch, runtime_state = prepare_for_prefix_sharing_fsdp(batch, config)
 
     assert runtime_state is not None
     plan = runtime_state.prefix_sharing_plan
@@ -273,7 +273,7 @@ def test_restore_prefix_sharing_outputs_2d_restores_interior_last_logits_entropy
             dtype=torch.long,
         ),
     }
-    _, runtime_state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    _, runtime_state = prepare_for_prefix_sharing_fsdp(batch, config)
     assert runtime_state is not None
 
     vocab = 5
@@ -345,7 +345,7 @@ def test_prefix_sharing_fsdp_attention_runtime_scatter_dense_outputs():
             dtype=torch.long,
         ),
     }
-    _, runtime_state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    _, runtime_state = prepare_for_prefix_sharing_fsdp(batch, config)
     assert runtime_state is not None
 
     torch.manual_seed(1)
@@ -368,7 +368,7 @@ def test_prefix_sharing_fsdp_attention_runtime_scatter_dense_outputs():
     assert not torch.allclose(dense_output[1, 3:6], torch.zeros_like(dense_output[1, 3:6]))
 
 
-def test_forward_prefix_sharing_fsdp_micro_batch_matches_tiny_hf_model_baseline():
+def test_forward_step_without_engine_prepare_matches_tiny_hf_model_baseline():
     torch.manual_seed(2026)
     config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3)
     batch = {
@@ -403,7 +403,7 @@ def test_forward_prefix_sharing_fsdp_micro_batch_matches_tiny_hf_model_baseline(
     baseline_log_probs = _mock_log_probs_fn(baseline_logits, labels)
     baseline_entropy = _entropy_from_logits(baseline_logits)
 
-    prefix_output = forward_prefix_sharing_fsdp_micro_batch(
+    prefix_output = forward_step_without_engine_prepare(
         batch,
         model,
         config,
@@ -418,7 +418,7 @@ def test_forward_prefix_sharing_fsdp_micro_batch_matches_tiny_hf_model_baseline(
     assert torch.allclose(prefix_output["attention_output"], baseline.attention_output, atol=1e-5)
 
 
-def test_forward_prefix_sharing_fsdp_micro_batch_keeps_provider_prefix_grad_path():
+def test_forward_step_without_engine_prepare_keeps_provider_prefix_grad_path():
     torch.manual_seed(2027)
     config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3)
     batch = {
@@ -443,7 +443,7 @@ def test_forward_prefix_sharing_fsdp_micro_batch_keeps_provider_prefix_grad_path
     batch["labels"] = labels
 
     model = _TinyHFStyleModel(vocab_size=32)
-    output = forward_prefix_sharing_fsdp_micro_batch(
+    output = forward_step_without_engine_prepare(
         batch,
         model,
         config,
@@ -686,7 +686,7 @@ def test_prefix_sharing_fsdp_attention_runtime_supports_packed_single_batch_shap
         "attention_mask": torch.ones(2, 5, dtype=torch.bool),
         "position_ids": torch.tensor([[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long),
     }
-    _, runtime_state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    _, runtime_state = prepare_for_prefix_sharing_fsdp(batch, config)
     assert runtime_state is not None
 
     torch.manual_seed(2031)
diff --git a/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_functional.py b/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_functional.py
index 7d7710cc..6dabd727 100644
--- a/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_functional.py
+++ b/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_functional.py
@@ -21,7 +21,7 @@
 from prefix_sharing.integrations.context import prefix_sharing_runtime_context
 from prefix_sharing.integrations.verl_fsdp import (
     PrefixSharingFSDPAttentionRuntime,
-    build_prefix_sharing_micro_batch_fsdp,
+    prepare_for_prefix_sharing_fsdp,
     restore_prefix_sharing_outputs_2d,
 )
 
@@ -52,7 +52,7 @@ def test_scenario1_two_samples_share_arbitrary_prefix():
     """A B C D E / A B C X Y -> provider=0, reuser=1, prefix_len=3."""
     config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2)
     batch = _make_batch([[1, 2, 3, 4, 5], [1, 2, 3, 10, 11]])
-    trimmed, state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    trimmed, state = prepare_for_prefix_sharing_fsdp(batch, config)
 
     assert state is not None
     plan = state.prefix_sharing_plan
@@ -76,7 +76,7 @@ def test_scenario2_multiple_reusers_different_prefix_and_suffix_lens():
             [1, 2, 20, 21],           # reuser, prefix 2, suffix 2 (len 4)
         ]
     )
-    trimmed, state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    trimmed, state = prepare_for_prefix_sharing_fsdp(batch, config)
 
     assert state is not None
     plan = state.prefix_sharing_plan
@@ -106,7 +106,7 @@ def test_scenario3_same_provider_serves_multiple_reusers():
             [1, 2, 30, 40, 50, 60],   # reuser 3 (prefix 2)
         ]
     )
-    trimmed, state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    trimmed, state = prepare_for_prefix_sharing_fsdp(batch, config)
 
     assert state is not None
     plan = state.prefix_sharing_plan
@@ -128,7 +128,7 @@ def test_scenario4_chain_reuse_reuser_becomes_provider():
             [1, 2, 3, 4, 10, 20, 30],   # row2: reuser of row1 (prefix 5)
         ]
     )
-    trimmed, state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    trimmed, state = prepare_for_prefix_sharing_fsdp(batch, config)
 
     assert state is not None
     plan = state.prefix_sharing_plan
@@ -146,7 +146,7 @@ def test_scenario5_no_shareable_prefix_returns_none():
     """Completely different sequences -> fallback."""
     config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2)
     batch = _make_batch([[1, 2, 3, 4], [5, 6, 7, 8]])
-    returned, state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    returned, state = prepare_for_prefix_sharing_fsdp(batch, config)
 
     assert returned is batch
     assert state is None
@@ -156,7 +156,7 @@ def test_scenario5b_single_sample_returns_none():
     """Single sample cannot share -> fallback."""
     config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2)
     batch = _make_batch([[1, 2, 3, 4, 5]])
-    returned, state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    returned, state = prepare_for_prefix_sharing_fsdp(batch, config)
 
     assert returned is batch
     assert state is None
@@ -205,7 +205,7 @@ def test_scenario7_position_ids_preserve_absolute_positions():
     """Reuser suffix position_ids must reflect original absolute positions."""
     config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2)
     batch = _make_batch([[1, 2, 3, 4, 5, 6, 7], [1, 2, 3, 10, 11, 12, 13]])
-    trimmed, state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    trimmed, state = prepare_for_prefix_sharing_fsdp(batch, config)
 
     assert state is not None
     plan = state.prefix_sharing_plan
@@ -224,7 +224,7 @@ def test_scenario7_position_ids_preserve_absolute_positions():
 def test_attention_runtime_records_stats_per_layer():
     config = PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=2)
     batch = _make_batch([[1, 2, 3, 4, 5], [1, 2, 3, 10, 11]])
-    _, state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    _, state = prepare_for_prefix_sharing_fsdp(batch, config)
     assert state is not None
 
     torch.manual_seed(42)
diff --git a/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_precision.py b/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_precision.py
index 07824d69..e6a72729 100644
--- a/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_precision.py
+++ b/prefix-sharing/tests/unit_test/test_verl_fsdp_ch4_precision.py
@@ -26,7 +26,7 @@
 from prefix_sharing.integrations.context import prefix_sharing_runtime_context
 from prefix_sharing.integrations.verl_fsdp import (
     PrefixSharingFSDPAttentionRuntime,
-    build_prefix_sharing_micro_batch_fsdp,
+    prepare_for_prefix_sharing_fsdp,
     restore_prefix_sharing_outputs_2d,
 )
 
@@ -67,7 +67,7 @@ def test_reuser_suffix_attention_matches_baseline():
             [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long
         ),
     }
-    trimmed, state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    trimmed, state = prepare_for_prefix_sharing_fsdp(batch, config)
     assert state is not None
     plan = state.prefix_sharing_plan
     prefix_len = plan.prefix_lens[1]  # 3
@@ -143,7 +143,7 @@ def test_gradient_flows_through_provider_prefix_kv():
             [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long
         ),
     }
-    trimmed, state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    trimmed, state = prepare_for_prefix_sharing_fsdp(batch, config)
     assert state is not None
 
     H, D = 2, 4
@@ -184,7 +184,7 @@ def test_restore_interior_prefix_matches_provider():
             [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long
         ),
     }
-    _, state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    _, state = prepare_for_prefix_sharing_fsdp(batch, config)
     assert state is not None
 
     # Simulate trimmed output: provider full, reuser has only suffix (prefix zeroed)
@@ -238,7 +238,7 @@ def test_restore_prefix_last_recomputed_with_provider_logits_and_reuser_label():
             [[0, 1, 2, 3, 4], [0, 1, 2, 3, 4]], dtype=torch.long
         ),
     }
-    _, state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    _, state = prepare_for_prefix_sharing_fsdp(batch, config)
     plan = state.prefix_sharing_plan
     prefix_len = plan.prefix_lens[1]  # 3
 
@@ -275,7 +275,7 @@ def test_baseline_matches_when_no_sharing():
             [[0, 1, 2, 3], [0, 1, 2, 3]], dtype=torch.long
         ),
     }
-    returned, state = build_prefix_sharing_micro_batch_fsdp(batch, config)
+    returned, state = prepare_for_prefix_sharing_fsdp(batch, config)
     assert state is None
     assert returned is batch  # exact same object -> no transformation
 
@@ -308,8 +308,8 @@ def test_padding_positions_do_not_contribute():
     }
 
     H, D = 2, 4
-    _, state_unpadded = build_prefix_sharing_micro_batch_fsdp(batch_unpadded, config)
-    _, state_padded = build_prefix_sharing_micro_batch_fsdp(batch_padded, config)
+    _, state_unpadded = prepare_for_prefix_sharing_fsdp(batch_unpadded, config)
+    _, state_padded = prepare_for_prefix_sharing_fsdp(batch_padded, config)
     assert state_unpadded is not None and state_padded is not None
 
     # Same plan semantics regardless of padding
diff --git a/prefix-sharing/tests/unit_test/test_verl_integration_boundaries.py b/prefix-sharing/tests/unit_test/test_verl_integration_boundaries.py
index f9ddcba3..5a00a292 100644
--- a/prefix-sharing/tests/unit_test/test_verl_integration_boundaries.py
+++ b/prefix-sharing/tests/unit_test/test_verl_integration_boundaries.py
@@ -22,10 +22,10 @@ def test_verl_fsdp_does_not_import_shared_helpers_from_mcore():
     source = inspect.getsource(verl_fsdp)
     forbidden_imports = (
         "from prefix_sharing.integrations.verl_mcore import PrefixSharingRuntimeState",
-        "from prefix_sharing.integrations.verl_mcore import _collect_kept_position_rows",
-        "from prefix_sharing.integrations.verl_mcore import _extract_seq_from_nested_tensor",
-        "from prefix_sharing.integrations.verl_mcore import _is_nested_tensor",
-        "from prefix_sharing.integrations.verl_mcore import _trim_nested_batch",
+        "from prefix_sharing.integrations.verl_mcore import collect_kept_position_rows",
+        "from prefix_sharing.integrations.verl_mcore import extract_seq_from_nested_tensor",
+        "from prefix_sharing.integrations.verl_mcore import is_nested_tensor",
+        "from prefix_sharing.integrations.verl_mcore import trim_redundant_prefix_in_nested_tensor",
     )
     for forbidden_import in forbidden_imports:
         assert forbidden_import not in source
diff --git a/prefix-sharing/tests/unit_test/test_verl_utils_trim.py b/prefix-sharing/tests/unit_test/test_verl_utils_trim.py
new file mode 100644
index 00000000..875b23c8
--- /dev/null
+++ b/prefix-sharing/tests/unit_test/test_verl_utils_trim.py
@@ -0,0 +1,143 @@
+"""Tests for nested-tensor trim helpers in verl_utils."""
+
+from __future__ import annotations
+
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from prefix_sharing.core.config import PrefixSharingConfig
+from prefix_sharing.core.planner import PrefixSharingPlanner
+from prefix_sharing.integrations.verl_utils import trim_redundant_prefix_in_dense_tensor
+from prefix_sharing.integrations.verl_utils import trim_redundant_prefix_in_nested_tensor
+
+
+def _prefix_sharing_plan():
+    planner = PrefixSharingPlanner(PrefixSharingConfig(enable_prefix_sharing=True, min_prefix_len=3))
+    return planner.plan([[1, 2, 3, 10, 11], [1, 2, 3, 20, 21, 22]])
+
+
+def _nested(rows):
+    return torch.nested.as_nested_tensor(
+        [torch.tensor(row, dtype=torch.long) for row in rows],
+        layout=torch.jagged,
+    )
+
+
+def _unpacked_rows(nested_tensor):
+    offsets = nested_tensor.offsets()
+    values = nested_tensor.values()
+    return [
+        values[offsets[i] : offsets[i + 1]].tolist()
+        for i in range(offsets.numel() - 1)
+    ]
+
+
+def test_trim_redundant_prefix_returns_kept_position_rows_from_nested_slice():
+    if not hasattr(torch, "nested"):
+        pytest.skip("torch.nested is unavailable")
+
+    plan = _prefix_sharing_plan()
+    batch = {
+        "input_ids": _nested([[1, 2, 3, 10, 11], [1, 2, 3, 20, 21, 22]]),
+        "position_ids": _nested([[0, 1, 2, 3, 4], [0, 1, 2, 3, 4, 5]]),
+        "loss_mask": _nested([[1, 1, 1, 1, 0], [1, 1, 1, 1, 1, 0]]),
+    }
+
+    trimmed_batch, kept_position_rows = trim_redundant_prefix_in_nested_tensor(batch, plan)
+
+    assert [row.tolist() for row in kept_position_rows] == [[0, 1, 2, 3, 4], [3, 4, 5]]
+    assert _unpacked_rows(trimmed_batch["position_ids"]) == [
+        row.tolist() for row in kept_position_rows
+    ]
+    assert _unpacked_rows(trimmed_batch["input_ids"]) == [[1, 2, 3, 10, 11], [20, 21, 22]]
+    assert _unpacked_rows(trimmed_batch["loss_mask"]) == [[1, 1, 1, 1, 0], [1, 1, 0]]
+
+
+def test_trim_redundant_prefix_returns_kept_position_rows_from_2d_slice():
+    if not hasattr(torch, "nested"):
+        pytest.skip("torch.nested is unavailable")
+
+    plan = _prefix_sharing_plan()
+    batch = {
+        "input_ids": _nested([[1, 2, 3, 10, 11], [1, 2, 3, 20, 21, 22]]),
+        "position_ids": torch.tensor(
+            [
+                [0, 1, 2, 3, 4, 0],
+                [0, 1, 2, 3, 4, 5],
+            ],
+            dtype=torch.long,
+        ),
+        "attention_mask": torch.tensor(
+            [
+                [1, 1, 1, 1, 1, 0],
+                [1, 1, 1, 1, 1, 1],
+            ],
+            dtype=torch.bool,
+        ),
+    }
+
+    trimmed_batch, kept_position_rows = trim_redundant_prefix_in_nested_tensor(batch, plan)
+
+    assert [row.tolist() for row in kept_position_rows] == [[0, 1, 2, 3, 4], [3, 4, 5]]
+    assert _unpacked_rows(trimmed_batch["position_ids"]) == [
+        row.tolist() for row in kept_position_rows
+    ]
+
+
+def test_trim_redundant_prefix_in_dense_tensor_crops_masks_and_returns_position_rows():
+    if not hasattr(torch, "nested"):
+        pytest.skip("torch.nested is unavailable")
+
+    plan = _prefix_sharing_plan()
+    batch = {
+        "input_ids": torch.tensor(
+            [
+                [1, 2, 3, 10, 11, 0],
+                [1, 2, 3, 20, 21, 22],
+            ],
+            dtype=torch.long,
+        ),
+        "position_ids": torch.tensor(
+            [
+                [0, 1, 2, 3, 4, 0],
+                [0, 1, 2, 3, 4, 5],
+            ],
+            dtype=torch.long,
+        ),
+        "attention_mask": torch.tensor(
+            [
+                [1, 1, 1, 1, 1, 0],
+                [1, 1, 1, 1, 1, 1],
+            ],
+            dtype=torch.bool,
+        ),
+        "loss_mask": torch.tensor(
+            [
+                [1, 1, 1, 1, 1, 0],
+                [1, 1, 1, 1, 1, 1],
+            ],
+            dtype=torch.bool,
+        ),
+    }
+    valid_indices = [
+        batch["attention_mask"][row].nonzero(as_tuple=False).flatten()
+        for row in range(batch["input_ids"].shape[0])
+    ]
+
+    trimmed_batch, kept_position_rows = trim_redundant_prefix_in_dense_tensor(
+        batch, plan, valid_indices
+    )
+
+    assert [row.tolist() for row in kept_position_rows] == [[0, 1, 2, 3, 4], [3, 4, 5]]
+    assert trimmed_batch["position_ids"] is batch["position_ids"]
+    assert trimmed_batch["position_ids"][1, 3:6].tolist() == [3, 4, 5]
+    assert trimmed_batch["attention_mask"].tolist() == [
+        [1, 1, 1, 1, 1, 0],
+        [0, 0, 0, 1, 1, 1],
+    ]
+    assert trimmed_batch["loss_mask"].tolist() == [
+        [1, 1, 1, 1, 1, 0],
+        [0, 0, 0, 1, 1, 1],
+    ]
+    assert trimmed_batch["input_ids"] is batch["input_ids"]