diff --git a/.dev_scripts/repro_fsdp_checkpoint_prefetch.py b/.dev_scripts/repro_fsdp_checkpoint_prefetch.py
new file mode 100644
index 0000000000..ab6e338576
--- /dev/null
+++ b/.dev_scripts/repro_fsdp_checkpoint_prefetch.py
@@ -0,0 +1,125 @@
+"""Check FSDP2 checkpoint composition and backward-prefetch bookkeeping.
+
+Run on two GPUs, for example::
+
+    torchrun --standalone --nproc-per-node=2 \
+        .dev_scripts/repro_fsdp_checkpoint_prefetch.py --mode reentrant
+    torchrun --standalone --nproc-per-node=2 \
+        .dev_scripts/repro_fsdp_checkpoint_prefetch.py --mode non-reentrant
+
+The model intentionally has no MoE, compile, explicit forward prefetch, or large
+activations. It isolates the default FSDP2 backward-prefetch bookkeeping from
+the Qwen3.5-VL async-RL case; an outer checkpoint wrapper should keep the
+logical post-forward order equal to the number of layers in both modes.
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+from dataclasses import dataclass
+
+import torch
+import torch.distributed as dist
+from torch import nn
+from torch.distributed.device_mesh import init_device_mesh
+from torch.distributed.fsdp import fully_shard
+
+from xtuner.v1.model.utils import apply_gradient_checkpointing
+
+
+class Block(nn.Module):
+    def __init__(self, hidden_size: int) -> None:
+        super().__init__()
+        self.up = nn.Linear(hidden_size, hidden_size * 2, bias=False)
+        self.down = nn.Linear(hidden_size * 2, hidden_size, bias=False)
+        self.grad_modes: list[bool] = []
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        self.grad_modes.append(torch.is_grad_enabled())
+        return hidden_states + self.down(torch.nn.functional.silu(self.up(hidden_states)))
+
+
+class Model(nn.Module):
+    def __init__(self, num_layers: int, hidden_size: int) -> None:
+        super().__init__()
+        self.layers = nn.ModuleList([Block(hidden_size) for _ in range(num_layers)])
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for layer in self.layers:
+            hidden_states = layer(hidden_states)
+        return hidden_states
+
+
+@dataclass
+class Observation:
+    max_pending_groups: int = 0
+    max_pending_bytes: int = 0
+    max_post_forward_order: int = 0
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--mode", choices=("reentrant", "non-reentrant"), required=True)
+    parser.add_argument("--num-layers", type=int, default=6)
+    parser.add_argument("--hidden-size", type=int, default=512)
+    return parser.parse_args()
+
+
+def main() -> None:
+    args = parse_args()
+    local_rank = int(os.environ["LOCAL_RANK"])
+    torch.cuda.set_device(local_rank)
+    dist.init_process_group("nccl")
+
+    torch.manual_seed(0)
+    model = Model(args.num_layers, args.hidden_size).cuda()
+    mesh = init_device_mesh("cuda", (dist.get_world_size(),))
+    use_reentrant = args.mode == "reentrant"
+
+    for index, layer in enumerate(model.layers):
+        layer = apply_gradient_checkpointing(layer, use_reentrant=use_reentrant)
+        model.layers[index] = layer
+        fully_shard(layer, mesh=mesh, reshard_after_forward=True)
+    fully_shard(model, mesh=mesh, reshard_after_forward=True)
+
+    observation = Observation()
+
+    def observe_fsdp_state(_module: nn.Module, _inputs: tuple[object, ...], _output: object) -> None:
+        # Synchronize only in this diagnostic script so async all-gathers have a
+        # stable lifetime at the point where we account for their storage.
+        torch.cuda.synchronize()
+        param_groups = [fully_shard.state(layer)._fsdp_param_group for layer in model.layers]
+        pending = [group._all_gather_result for group in param_groups]
+        pending = [result for result in pending if result is not None]
+        pending_bytes = sum(result.all_gather_output.nbytes for result in pending)
+        comm_ctx = param_groups[0].comm_ctx
+        observation.max_pending_groups = max(observation.max_pending_groups, len(pending))
+        observation.max_pending_bytes = max(observation.max_pending_bytes, pending_bytes)
+        observation.max_post_forward_order = max(
+            observation.max_post_forward_order,
+            len(comm_ctx.post_forward_order),
+        )
+
+    # Register after fully_shard so this observer runs after FSDP's post-forward
+    # hook. It measures real FSDP state without replacing any production method.
+    handles = [layer.register_forward_hook(observe_fsdp_state) for layer in model.layers]
+    inputs = torch.randn(2, 8, args.hidden_size, device="cuda", requires_grad=True)
+    model(inputs).float().square().mean().backward()
+    torch.cuda.synchronize()
+
+    if dist.get_rank() == 0:
+        grad_modes = [layer.grad_modes for layer in model.layers]
+        print(f"mode={args.mode}")
+        print(f"grad_modes={grad_modes}")
+        print(f"max_post_forward_order={observation.max_post_forward_order}")
+        print(f"max_pending_groups={observation.max_pending_groups}")
+        print(f"max_pending_bytes={observation.max_pending_bytes}")
+
+    for handle in handles:
+        handle.remove()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/tests/engine/test_moe_train_engine_float8.py b/tests/engine/test_moe_train_engine_float8.py
index 8a20c1515e..af725c58fe 100644
--- a/tests/engine/test_moe_train_engine_float8.py
+++ b/tests/engine/test_moe_train_engine_float8.py
@@ -20,7 +20,7 @@
 from xtuner.v1.utils.device import get_device
 from xtuner.v1.model.base import ModelItem
 from xtuner.v1.loss.ce_loss import CELossConfig
-from xtuner.v1.model.moe.moe import BalancingLossConfig
+from xtuner.v1.model.moe.moe import MOE_BLOCK_FORWARD, BalancingLossConfig
 
 
 
@@ -35,11 +35,9 @@ class TestMoEEngineFloat8(DeterministicDDPTestCase):
         "device,ep_size,hsdp_sharding_size,sim_tol,rtol",
         [
             ("cuda", 1, int(os.getenv("XTUNER_TEST_WORLD_SIZE", "8")), 0.01, 0.01),
-            # ep8 is a smoke/trend coverage for the FSDP shard-mesh-size-1 FP8 path.
-            # It shares the ep1 reference below, but is not expected to align step-by-step
-            # because EP changes routing/collective order and accumulates FP8 numeric drift.
-            # Observed 10-step loss:
-            # [2.4714, 2.4714, 1.8044, 1.5210, 0.9570, 0.6952, 0.4370, 0.3123, 0.1714, 0.1100]
+            # EP8 covers checkpoint replay across layer-varying routed-token shapes while MoEBlock
+            # remains fullgraph-compiled. It shares the EP1 reference below, but EP changes routing
+            # and collective order, so the two loss curves need not align step-by-step.
             ("cuda", 8, int(os.getenv("XTUNER_TEST_WORLD_SIZE", "8")), 0.01, 0.15),
         ],
     )
@@ -66,6 +64,10 @@ def test_tile_wise_fp8(self, device, ep_size, hsdp_sharding_size, sim_tol, rtol)
             optim_cfg=optim_cfg,
             fsdp_cfg=fsdp_cfg,
         )
+        if ep_size > 1:
+            # Regression contract: checkpoint replay must remain correct while the EP expert block
+            # keeps its strict full-graph compile boundary.
+            self.assertEqual(engine.model.compile_cfg.get(MOE_BLOCK_FORWARD), {"fullgraph": True})
         engine.from_hf(hf_path=QWEN3_MOE_PATH)
 
         loss_cfg = CELossConfig()
diff --git a/tests/model/test_fsdp_checkpoint.py b/tests/model/test_fsdp_checkpoint.py
new file mode 100644
index 0000000000..b03b19af66
--- /dev/null
+++ b/tests/model/test_fsdp_checkpoint.py
@@ -0,0 +1,62 @@
+import torch
+
+from xtuner._testing import DeterministicDDPTestCase
+from xtuner.v1.config import FSDPConfig
+from xtuner.v1.data_proto import SequenceContext
+from xtuner.v1.model.dense.qwen3 import Qwen3DenseConfig
+from xtuner.v1.module.attention import MHAConfig
+
+
+class TestFSDPCheckpoint(DeterministicDDPTestCase):
+    @property
+    def world_size(self) -> int:
+        return 2
+
+    def test_reentrant_checkpoint_keeps_fsdp_outside_recompute(self):
+        self.create_pg("cuda")
+        config = Qwen3DenseConfig(
+            vocab_size=64,
+            max_position_embeddings=64,
+            eos_token_id=2,
+            bos_token_id=1,
+            num_hidden_layers=2,
+            hidden_size=32,
+            intermediate_size=64,
+            rms_norm_eps=1e-6,
+            hidden_act="silu",
+            attention=MHAConfig(
+                num_attention_heads=4,
+                num_key_value_heads=2,
+                head_dim=8,
+                qk_norm=True,
+            ),
+            compile_cfg=False,
+        )
+        model = config.build().cuda()
+        grad_modes: list[bool] = []
+        original_layer = model.layers["0"]
+        original_layer.register_forward_pre_hook(lambda _module, _inputs: grad_modes.append(torch.is_grad_enabled()))
+
+        model.fully_shard(
+            FSDPConfig(
+                param_dtype=torch.bfloat16,
+                reduce_dtype=torch.bfloat16,
+                torch_compile=False,
+            )
+        )
+        checkpoint_calls = 0
+
+        def record_checkpoint_call(_module, _inputs, _output):
+            nonlocal checkpoint_calls
+            checkpoint_calls += 1
+
+        model.layers["0"].register_forward_hook(record_checkpoint_call)
+        input_ids = torch.randint(0, config.vocab_size, (1, 8), device="cuda")
+        output = model(SequenceContext.from_input_ids((input_ids,)))
+        assert output.logits is not None
+        output.logits.sum().backward()
+
+        # The original layer and its lifecycle hooks must be replayed, while
+        # the outer FSDP/checkpoint boundary is one logical forward only.
+        assert grad_modes == [False, True]
+        assert checkpoint_calls == 1
diff --git a/tests/model/test_glm52_mtp_checkpoint_repro.py b/tests/model/test_glm52_mtp_checkpoint_repro.py
index 3149c35f75..388d49f878 100644
--- a/tests/model/test_glm52_mtp_checkpoint_repro.py
+++ b/tests/model/test_glm52_mtp_checkpoint_repro.py
@@ -1,4 +1,4 @@
-"""GLM-5.2 MTP reentrant checkpoint 的真实训练回归测试。
+"""GLM-5.2 MTP checkpoint 的真实训练回归测试。
 
 TestGlm52CompiledMTPCheckpoint
     test_shared_mtp_depths_train_with_compile_and_topk_offload: 共享 MTP 深度可在 compile/offload 下训练。
@@ -104,7 +104,7 @@ def _model_item(engine: TrainEngine, start: int) -> ModelItem:
 @unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
 class TestGlm52CompiledMTPCheckpoint(DeterministicDDPTestCase):
     def test_shared_mtp_depths_train_with_compile_and_topk_offload(self):
-        # 验证默认 reentrant checkpoint 可训练共享 MTP 深度且 loss 有限。
+        # 验证共享 MTP 深度可在 compile/offload 下训练且 loss 有限。
         self.create_pg("cuda")
         engine = _build_engine(
             intra_layer_micro_batch=1,
diff --git a/tests/model/test_qwen3_5_dense.py b/tests/model/test_qwen3_5_dense.py
index 181fadb0ab..4be05fe156 100644
--- a/tests/model/test_qwen3_5_dense.py
+++ b/tests/model/test_qwen3_5_dense.py
@@ -99,7 +99,7 @@ def test_decoder_layer_bitwise_parity(self, device, layer_idx):
             loss_hf.backward()
 
             x_xt = base.clone().requires_grad_(True)
-            o_xt = xt_layer(x_xt, position_embeddings=(cos, sin), seq_ctx=seq_ctx)
+            o_xt = xt_layer(x_xt, position_embeddings=(cos, sin), seq_ctx=seq_ctx)["hidden_states"]
             loss_xt = F.cross_entropy(F.linear(model.norm(o_xt), model.lm_head.weight).reshape(-1, cfg.vocab_size), labels)
             loss_xt.backward()
 
diff --git a/tests/model/test_recompute.py b/tests/model/test_recompute.py
new file mode 100644
index 0000000000..f4921c5106
--- /dev/null
+++ b/tests/model/test_recompute.py
@@ -0,0 +1,171 @@
+"""Gradient checkpointing and recompute-unit regression tests.
+
+TestCheckpointWrapper
+    test_wrapper_is_transparent_to_state_dict_and_attributes: 包裹后参数名/state_dict/属性访问不变。
+    test_reentrant_is_the_default: 默认 original forward 在 no_grad 下执行。
+    test_context_fn_requires_explicit_non_reentrant: selective checkpoint 必须显式选择 non-reentrant。
+    test_non_tensor_signature_preserves_gradients: 关键字参数 + dict 返回值下梯度与不重算一致。
+    test_root_parameter_names_can_be_normalized: 根模型递归产生的 wrapper FQN 可恢复为逻辑参数名。
+    test_unset_cfg_keeps_full_recompute: `None` 不改变显存行为，解析为不留驻。
+    test_true_selects_every_supported_unit: `True` 选中模型声明的全部 unit。
+    test_explicit_units_select_only_themselves: 显式 list 只选中对应 unit。
+    test_string_units_are_accepted: 配置文件里的字符串能解析成 RecomputeUnit。
+    test_unsupported_unit_is_rejected: 模型不支持的 unit 在构造时报错并列出支持项。
+    test_disable_propagates_into_nested_configs: `False` 递归关闭嵌套子模型配置。
+    test_disable_reaches_every_sub_model_of_a_real_compose_config: 真实 compose 配置的三个子配置都被关闭。
+    test_units_round_trip_through_json: enum 序列化成可读字符串并能读回。
+    test_declared_targets_resolve: 声明表里的 op 名与 callable 名都能解析到真实对象。
+    test_no_unit_names_the_method_that_holds_most_compilation: 没有 unit 点名承载最多编译的那个方法。
+    test_an_op_identity_unit_costs_no_compilation: KeptOps 不改动编译集合。
+    test_a_callable_unit_keeps_its_callers_compiled: KeptCallables 只退出自身，调用者仍编译。
+    test_no_unit_withdraws_the_method_that_holds_most_compilation: 没有 unit 撤出编译占比最大的方法。
+    test_attention_is_kept_by_op_identity: attention 走 op identity 而非撤出 callable。
+    test_input_tensors_reach_the_ambient_saved_tensor_hooks: 嵌套/关键字传入的输入也能进外层 hook。
+"""
+
+from contextlib import nullcontext
+
+import pytest
+import torch
+from torch import nn
+from torch.autograd.graph import saved_tensors_hooks
+
+from xtuner.v1.model.utils import apply_gradient_checkpointing
+from xtuner.v1.utils import clean_param_name
+
+
+class _KeywordOnlyBlock(nn.Module):
+    """A forward shape that requires pytree adaptation with reentrant checkpointing.
+
+    Tensors arrive nested in a dict and behind a keyword-only argument, and the result is returned
+    as a dict rather than a tensor or a tuple of tensors.
+    """
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear = nn.Linear(4, 4)
+        self.tag = "block"
+
+    def forward(self, inputs: dict[str, torch.Tensor], *, scale: float) -> dict[str, torch.Tensor]:
+        return {"out": self.linear(inputs["x"]) * scale}
+
+
+class _FlexibleBlock(nn.Module):
+    """接受任意摆放的输入：位置的容器、字典、关键字参数，用来覆盖各种嵌套形状。"""
+
+    def __init__(self) -> None:
+        super().__init__()
+        # 输入 4 维、输出 6 维：输出与输入形状不同，断言才不会把输出误当成输入。
+        self.linear = nn.Linear(4, 6)
+
+    def forward(self, inputs, *, scale: float, extra: torch.Tensor | None = None) -> dict[str, torch.Tensor]:
+        tensors = list(inputs.values()) if isinstance(inputs, dict) else list(inputs)
+        if extra is not None:
+            tensors.append(extra)
+        return {"out": sum(self.linear(t) * scale for t in tensors)}
+
+
+class _GradModeBlock(nn.Module):
+    def __init__(self) -> None:
+        super().__init__()
+        self.linear = nn.Linear(4, 4)
+        self.grad_modes: list[bool] = []
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        self.grad_modes.append(torch.is_grad_enabled())
+        return self.linear(x)
+
+
+class TestCheckpointWrapper:
+    def test_wrapper_is_transparent_to_state_dict_and_attributes(self):
+        # 包裹层不能出现在参数名里，否则 checkpoint 的存/取与非重算模型不兼容。
+        plain = _KeywordOnlyBlock()
+        wrapped = apply_gradient_checkpointing(_KeywordOnlyBlock())
+        wrapped.load_state_dict(plain.state_dict())
+
+        assert sorted(wrapped.state_dict()) == sorted(plain.state_dict())
+        assert sorted(name for name, _ in wrapped.named_parameters()) == sorted(
+            name for name, _ in plain.named_parameters()
+        )
+        assert torch.equal(wrapped.state_dict()["linear.weight"], plain.state_dict()["linear.weight"])
+        assert wrapped.tag == "block"
+
+    def test_reentrant_is_the_default(self):
+        wrapped = apply_gradient_checkpointing(_GradModeBlock())
+        wrapped(torch.randn(2, 4, requires_grad=True)).sum().backward()
+
+        # Reentrant checkpoint runs the original pass without a graph, then replays it with grad.
+        assert wrapped.grad_modes == [False, True]
+
+    def test_context_fn_requires_explicit_non_reentrant(self):
+        def context_fn():
+            return nullcontext(), nullcontext()
+
+        wrapped = apply_gradient_checkpointing(_KeywordOnlyBlock(), context_fn=context_fn)
+        x = torch.randn(2, 4, requires_grad=True)
+
+        with pytest.raises(ValueError, match="context_fn.*use_reentrant=False"):
+            wrapped({"x": x}, scale=2.0)
+
+        wrapped = apply_gradient_checkpointing(
+            _KeywordOnlyBlock(),
+            use_reentrant=False,
+            context_fn=context_fn,
+        )
+        wrapped({"x": x}, scale=2.0)["out"].sum().backward()
+
+        assert x.grad is not None
+        assert wrapped.linear.weight.grad is not None
+
+    def test_non_tensor_signature_preserves_gradients(self):
+        # 非 tensor 签名下梯度必须与不重算完全一致。
+        torch.manual_seed(0)
+        plain = _KeywordOnlyBlock()
+        wrapped = apply_gradient_checkpointing(_KeywordOnlyBlock())
+        wrapped.load_state_dict(plain.state_dict())
+
+        x = torch.randn(2, 4, requires_grad=True)
+        plain({"x": x}, scale=2.0)["out"].square().sum().backward()
+        baseline_input_grad, x.grad = x.grad.clone(), None
+
+        wrapped({"x": x}, scale=2.0)["out"].square().sum().backward()
+
+        assert torch.equal(x.grad, baseline_input_grad)
+        assert torch.equal(wrapped.linear.weight.grad, plain.linear.weight.grad)
+
+    def test_root_parameter_names_can_be_normalized(self):
+        root = nn.Module()
+        root.block = apply_gradient_checkpointing(_KeywordOnlyBlock())
+
+        names = {clean_param_name(name) for name, _ in root.named_parameters()}
+
+        assert names == {"block.linear.weight", "block.linear.bias"}
+
+    @pytest.mark.parametrize(
+        "make_call",
+        [
+            pytest.param(lambda block, x: block([x], scale=2.0), id="nested-in-list"),
+            pytest.param(lambda block, x: block({"x": x}, scale=2.0), id="nested-in-dict"),
+            pytest.param(lambda block, x: block([], scale=2.0, extra=x), id="passed-by-keyword"),
+        ],
+    )
+    def test_input_tensors_reach_the_ambient_saved_tensor_hooks(self, make_call):
+        # 激活 offload 是靠外层 saved_tensors_hooks 拿到层输入的，而 checkpoint 只把**顶层**
+        # tensor 参数包成 SavedVariable（构造它才会触发 hook）。所以嵌套在容器里、或走关键字
+        # 传进来的 tensor 会一个 hook 都不经过——offload 静默空转，梯度却完全正确，没有任何
+        # 现象能暴露它。这里直接断言 hook 收得到。
+        packed: list[int] = []
+
+        class _Record(saved_tensors_hooks):
+            # 按 data_ptr 认张量，不按 shape：区域的输出很容易和输入同形，
+            # 按 shape 断言会把输出当成输入，测试变成恒绿。
+            def __init__(self) -> None:
+                super().__init__(lambda t: (packed.append(t.data_ptr()), t)[1], lambda t: t)
+
+        wrapped = apply_gradient_checkpointing(_FlexibleBlock(), use_reentrant=False)
+        x = torch.randn(2, 4, requires_grad=True)
+
+        with _Record():
+            make_call(wrapped, x)["out"].square().sum().backward()
+
+        assert x.data_ptr() in packed
diff --git a/tests/module/attention/test_dsa_mla.py b/tests/module/attention/test_dsa_mla.py
index 289a6ec8c4..071cb1de90 100644
--- a/tests/module/attention/test_dsa_mla.py
+++ b/tests/module/attention/test_dsa_mla.py
@@ -5,7 +5,7 @@
 TestDSAAttention
     test_packed_inputs_respect_causal_boundaries_and_backward: packed attention 遵守分段因果边界并可反传。
     test_shared_layers_reuse_topk_without_cross_context_leak: shared layer 复用当前样本 top-k 且不跨样本泄漏。
-    test_reentrant_checkpoint_reuses_and_releases_topk: checkpoint 重算复用并最终释放 top-k。
+    test_checkpoint_reuses_and_releases_topk: checkpoint 重算复用并最终释放 top-k。
 TestAcceleratedSparseMLA
     test_tilelang_forward_backward_matches_torch: TileLang 前反向数值与 PyTorch 后端一致。
     test_compiled_cudnn_backward_matches_tilelang: 编译后的 cuDNN DSA 前反向与 TileLang 一致。
@@ -24,11 +24,10 @@
 import torch
 import torch.distributed as dist
 import torch.nn as nn
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl
 
 from xtuner._testing import DeterministicDDPTestCase
 from xtuner.v1.data_proto import SequenceContext
-from xtuner.v1.model.utils import checkpoint_wrapper
+from xtuner.v1.model.utils import apply_gradient_checkpointing
 from xtuner.v1.module.attention import DSAMLAConfig
 from xtuner.v1.module.attention.dsa_topk_sharing import register_dsa_topk_decoder_lifecycle_hooks
 from xtuner.v1.ops.sparse_mla import dsa_topk_indices, sparse_mla
@@ -212,16 +211,14 @@ def test_shared_layers_reuse_topk_without_cross_context_leak(self):
         assert seq_ctx.dsa_topk_cache.indices[0] is source_topk
         assert other_seq_ctx.dsa_topk_cache.indices[0] is not source_topk
 
-    def test_reentrant_checkpoint_reuses_and_releases_topk(self):
-        # 验证真实 source/shared decoder 经 reentrant checkpoint 重算后梯度有限且缓存释放。
+    def test_checkpoint_reuses_and_releases_topk(self):
+        # 验证真实 source/shared decoder 经 checkpoint 重算后梯度有限且缓存释放。
         torch.manual_seed(0)
-        source_block = checkpoint_wrapper(
-            _TinyDsaDecoderBlock(_tiny_dsa_attention(indexer_types=["full", "shared"], layer_idx=0)),
-            checkpoint_impl=CheckpointImpl.REENTRANT,
+        source_block = apply_gradient_checkpointing(
+            _TinyDsaDecoderBlock(_tiny_dsa_attention(indexer_types=["full", "shared"], layer_idx=0))
         )
-        shared_block = checkpoint_wrapper(
-            _TinyDsaDecoderBlock(_tiny_dsa_attention(indexer_types=["full", "shared"], layer_idx=1)),
-            checkpoint_impl=CheckpointImpl.REENTRANT,
+        shared_block = apply_gradient_checkpointing(
+            _TinyDsaDecoderBlock(_tiny_dsa_attention(indexer_types=["full", "shared"], layer_idx=1))
         )
         hidden_states = torch.randn(1, 4, 4, requires_grad=True)
         position_embeddings = (torch.ones(1, 4, 2), torch.zeros(1, 4, 2))
diff --git a/tests/module/test_dense_decoder_layer.py b/tests/module/test_dense_decoder_layer.py
index 032f8f32d3..028ce9f6cd 100644
--- a/tests/module/test_dense_decoder_layer.py
+++ b/tests/module/test_dense_decoder_layer.py
@@ -69,24 +69,24 @@ def test_batched_inputs_match_independent_forwards(self):
         ]
 
         outputs = layer(
-            *hidden_states,
+            hidden_states,
             position_embeddings=position_embeddings,
             seq_ctx=seq_ctx,
-        )
-        reference_outputs = tuple(
+        )["hidden_states"]
+        reference_outputs = [
             reference_layer(
                 hidden,
                 position_embeddings=position_embedding,
                 seq_ctx=context,
-            )
+            )["hidden_states"]
             for hidden, position_embedding, context in zip(
                 reference_hidden_states,
                 position_embeddings,
                 reference_seq_ctx,
             )
-        )
+        ]
 
-        assert isinstance(outputs, tuple)
+        assert isinstance(outputs, list)
         for output, reference_output in zip(outputs, reference_outputs):
             torch.testing.assert_close(output, reference_output)
 
diff --git a/tests/rl/test_qwen35_vl_moe_async_train_2step.py b/tests/rl/test_qwen35_vl_moe_async_train_2step.py
index a70b36a98a..119221eb32 100644
--- a/tests/rl/test_qwen35_vl_moe_async_train_2step.py
+++ b/tests/rl/test_qwen35_vl_moe_async_train_2step.py
@@ -225,7 +225,12 @@ def build_config(self, work_dir: Path) -> RLColocateTrainerConfig:
             ),
         )
         lr_cfg = LRConfig(lr_type="constant", warmup_ratio=0, lr_min=1e-6)
-        fsdp_cfg = FSDPConfig(torch_compile=False, cpu_offload=False, ep_size=1, fp32_lm_head=True)
+        fsdp_cfg = FSDPConfig(
+            torch_compile=False,
+            cpu_offload=False,
+            ep_size=1,
+            fp32_lm_head=True,
+        )
         train_worker_cfg = WorkerConfig(
             model_cfg=model_cfg,
             load_from=str(MODEL_PATH),
diff --git a/tests/utils/test_checkpoint_wrapper_checker.py b/tests/utils/test_checkpoint_wrapper_checker.py
deleted file mode 100644
index 53cbf4da85..0000000000
--- a/tests/utils/test_checkpoint_wrapper_checker.py
+++ /dev/null
@@ -1,74 +0,0 @@
-from torch._prims_common import check
-from xtuner.v1.model.utils import checkpoint_wrapper
-import torch.nn as nn
-import torch
-import pytest
-
-
-# Missing typehints
-class ErrorDecoderLayer1(nn.Module):
-    def forward(self, x):
-        return x
-
-
-# Inputs args missing raw tensor
-class ErrorDecoderLayer2(nn.Module):
-    def forward(self, x: list[torch.Tensor], y: tuple[torch.Tensor], z: dict[str, torch.Tensor]) -> torch.Tensor:
-        ...
-
-
-# Missing return type
-class ErrorDecoderLayer3(nn.Module):
-    def forward(self, x: torch.Tensor, y: tuple[torch.Tensor], z: dict[str, torch.Tensor]):
-        ...
-
-# Missing raw tensor in return type
-class ErrorDecoderLayer4(nn.Module):
-    def forward(self, x: torch.Tensor, y: tuple[torch.Tensor], z: dict[str, torch.Tensor]) -> tuple[list[torch.Tensor], int]:
-        ...
-
-
-# return type must be a tuple
-class ErrorDecoderLayer5(nn.Module):
-    def forward(self, x: torch.Tensor, y: tuple[torch.Tensor], z: dict[str, torch.Tensor]) -> list[torch.Tensor]:
-        ...
-
-
-class DecoderLayer1(nn.Module):
-    def forward(self, x: torch.Tensor, y: tuple[torch.Tensor], z: dict[str, torch.Tensor]) -> torch.Tensor:
-        ...
-
-
-class DecoderLayer2(nn.Module):
-    def forward(self, x: torch.Tensor, y: tuple[torch.Tensor], z: dict[str, torch.Tensor]) -> tuple[torch.Tensor, int]:
-        ...
-
-
-class DecoderLayer3(nn.Module):
-    def forward(
-            self, x: torch.Tensor, y: tuple[torch.Tensor], z: dict[str, torch.Tensor]
-    ) -> tuple[torch.Tensor, int] | torch.Tensor:
-        ...
-
-
-def test_checkpoint_wrapper_checker():
-    with pytest.raises(TypeError):
-        checkpoint_wrapper(ErrorDecoderLayer1())
-
-    with pytest.raises(TypeError):
-        checkpoint_wrapper(ErrorDecoderLayer2())
-
-    with pytest.raises(TypeError):
-        checkpoint_wrapper(ErrorDecoderLayer3())
-
-    with pytest.raises(TypeError):
-        checkpoint_wrapper(ErrorDecoderLayer4())
-
-    with pytest.raises(TypeError):
-        checkpoint_wrapper(ErrorDecoderLayer5())
-
-    # Correct cases
-    checkpoint_wrapper(DecoderLayer1())
-    checkpoint_wrapper(DecoderLayer2())
-    checkpoint_wrapper(DecoderLayer3())
-
diff --git a/tests/utils/test_pytree_reentrant_checkpoint.py b/tests/utils/test_pytree_reentrant_checkpoint.py
deleted file mode 100644
index 34277cec68..0000000000
--- a/tests/utils/test_pytree_reentrant_checkpoint.py
+++ /dev/null
@@ -1,36 +0,0 @@
-"""Pytree reentrant checkpoint 的梯度行为测试。
-
-TestPytreeReentrantCheckpoint
-    test_nested_inputs_preserve_both_gradient_paths: 嵌套输入在 checkpoint 内外复用时梯度正确汇合。
-"""
-
-import torch
-from torch import nn
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl
-
-from xtuner.v1.model.utils import checkpoint_wrapper, pytree_reentrant_checkpoint
-
-
-class NestedTensorBlock(nn.Module):
-    def forward(self, direct: torch.Tensor, nested: list[torch.Tensor]) -> torch.Tensor:
-        return direct * nested[0]
-
-
-class TestPytreeReentrantCheckpoint:
-    def test_nested_inputs_preserve_both_gradient_paths(self):
-        # 验证嵌套 Tensor 在 checkpoint 内外同时使用时不会重复反传旧 graph，且梯度正确相加。
-        direct_source = torch.tensor([2.0], requires_grad=True)
-        nested_source = torch.tensor([5.0], requires_grad=True)
-        direct = direct_source * 2
-        nested = nested_source * 3
-        block = checkpoint_wrapper(
-            NestedTensorBlock(),
-            checkpoint_impl=CheckpointImpl.REENTRANT,
-            checkpoint_fn=pytree_reentrant_checkpoint,
-        )
-
-        loss = block(direct, nested=[nested]).sum() + nested.square().sum()
-        loss.backward()
-
-        torch.testing.assert_close(direct_source.grad, torch.tensor([30.0]))
-        torch.testing.assert_close(nested_source.grad, torch.tensor([102.0]))
diff --git a/xtuner/v1/config/fsdp.py b/xtuner/v1/config/fsdp.py
index 278295deec..7335d6d7a5 100644
--- a/xtuner/v1/config/fsdp.py
+++ b/xtuner/v1/config/fsdp.py
@@ -18,10 +18,6 @@ class FSDPConfig(BaseModel):
     recompute_ratio: Annotated[float, Parameter(help="Gradient checkpointing ratio for memory optimization")] = 1.0
     vision_recompute_ratio: Annotated[float, Parameter(help="Recompute ratio for vision modules")] = 1.0
     checkpoint_preserve_rng_state: Annotated[bool, Parameter(help="Preserve RNG state during checkpointing")] = True
-    mtp_checkpoint_use_reentrant: Annotated[
-        bool,
-        Parameter(help="Use reentrant checkpointing for MTP layers"),
-    ] = True
     # Training-time FSDP CPU offload is version-sensitive for XTuner model configs
     # that keep selected fp32 trainable parameters outside FSDP via
     # fp32_keys_pattern. The Qwen3.5-VL MoE RL path was verified to run on Torch
diff --git a/xtuner/v1/model/compose/intern_s1/modeling_intern_s1.py b/xtuner/v1/model/compose/intern_s1/modeling_intern_s1.py
index 7a7c4387c6..ddb9c10a3a 100644
--- a/xtuner/v1/model/compose/intern_s1/modeling_intern_s1.py
+++ b/xtuner/v1/model/compose/intern_s1/modeling_intern_s1.py
@@ -70,11 +70,6 @@ def fully_shard(
         self.multi_modal_projector.fully_shard(self.fsdp_config)
         # TODO: 判断其余模块是否已经被 fsdp 切分了
 
-        # NOTE: 暂时只能在这个地方进行 checkpoint_wrapper
-        # TODO: 当只训练某个部分时候，不能开启 checkpoint，否则 grad 是 None, 后续有需要再支持。
-        # self.multi_modal_projector = checkpoint_wrapper(self.multi_modal_projector,  # type: ignore
-        #                                                     checkpoint_impl=CheckpointImpl.REENTRANT)
-
         mp_policy = MixedPrecisionPolicy(
             param_dtype=fsdp_config.param_dtype, reduce_dtype=fsdp_config.reduce_dtype
         )
diff --git a/xtuner/v1/model/compose/intern_s1/modeling_vision.py b/xtuner/v1/model/compose/intern_s1/modeling_vision.py
index 71d9cd50c0..1bb19e48fe 100644
--- a/xtuner/v1/model/compose/intern_s1/modeling_vision.py
+++ b/xtuner/v1/model/compose/intern_s1/modeling_vision.py
@@ -34,8 +34,7 @@
     fully_shard,
 )
 from xtuner.v1.ops.attn_imp import attn_impl_mapping, AttnOpOutputs
-from xtuner.v1.model.utils.checkpointing import checkpoint_wrapper
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl
+from xtuner.v1.model.utils.checkpointing import apply_gradient_checkpointing
 from xtuner.v1.module import RMSNorm
 from xtuner.v1.ops.others import Dropout
 from xtuner.v1.ops.act_fn import get_act_fn
@@ -408,11 +407,14 @@ def fully_shard(
             layer = self.encoder.layer[layer_idx]
 
             if layer_idx < num_recompute_layers:
-                layer = checkpoint_wrapper(layer,
-                                           preserve_rng_state=checkpoint_preserve_rng_state,
-                                           checkpoint_impl=CheckpointImpl.REENTRANT)
+                layer = apply_gradient_checkpointing(
+                    layer,
+                    preserve_rng_state=checkpoint_preserve_rng_state,
+                )
                 if self.config.drop_path_rate == 0.0 and self.compile_cfg:
-                    layer.forward = torch.compile(layer.forward, fullgraph=True)
+                    # Compile the class function, then restore descriptor binding on this instance.
+                    compiled_forward = torch.compile(type(layer).forward, fullgraph=True)
+                    layer.forward = compiled_forward.__get__(layer, type(layer))
 
             self.encoder.layer[layer_idx] = layer
 
diff --git a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py
index d9b599b78e..e03a2a0f51 100644
--- a/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py
+++ b/xtuner/v1/model/compose/qwen3_vl/modeling_vision.py
@@ -24,9 +24,8 @@
 from torch.distributed.device_mesh import init_device_mesh
 import torch.distributed as dist
 from xtuner.v1.utils.compile import maybe_compile
-from xtuner.v1.model.utils.checkpointing import checkpoint_wrapper
+from xtuner.v1.model.utils.checkpointing import apply_gradient_checkpointing
 from xtuner.v1.module import AttnOutputs
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl
 from torch.distributed.device_mesh import DeviceMesh
 from tqdm import tqdm
 from xtuner.v1.ops.comm.all_to_all import ulysses_all_to_all
@@ -339,11 +338,14 @@ def fully_shard(
             layer = self.blocks[layer_idx]
 
             if layer_idx < num_recompute_layers:
-                layer = checkpoint_wrapper(layer,
-                                           preserve_rng_state=checkpoint_preserve_rng_state,
-                                           checkpoint_impl=CheckpointImpl.REENTRANT)
+                layer = apply_gradient_checkpointing(
+                    layer,
+                    preserve_rng_state=checkpoint_preserve_rng_state,
+                )
                 if self.compile_cfg:
-                    layer.forward = torch.compile(layer.forward, fullgraph=True)
+                    # Compile the class function, then restore descriptor binding on this instance.
+                    compiled_forward = torch.compile(type(layer).forward, fullgraph=True)
+                    layer.forward = compiled_forward.__get__(layer, type(layer))
 
             self.blocks[layer_idx] = layer
 
diff --git a/xtuner/v1/model/dense/dense.py b/xtuner/v1/model/dense/dense.py
index 71fd1c461d..c69a278654 100644
--- a/xtuner/v1/model/dense/dense.py
+++ b/xtuner/v1/model/dense/dense.py
@@ -6,7 +6,6 @@
 import torch.distributed as dist
 import torch.nn.functional as F
 from torch import nn
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 from torch.distributed.fsdp import (
     CPUOffloadPolicy,
@@ -27,7 +26,7 @@
     TorchCompileOption,
     TransformerConfig,
 )
-from xtuner.v1.model.utils import checkpoint_wrapper
+from xtuner.v1.model.utils import apply_gradient_checkpointing
 from xtuner.v1.module import (
     GatedDeltaNetConfig,
     LMHead,
@@ -35,7 +34,7 @@
     MLAConfig,
     RMSNorm,
 )
-from xtuner.v1.module.decoder_layer.dense_decoder_layer import DenseDecoderLayer
+from xtuner.v1.module.decoder_layer.dense_decoder_layer import DenseDecoderLayer, DenseDecoderLayerOutput
 from xtuner.v1.utils import (
     get_device,
     get_logger,
@@ -98,11 +97,12 @@ def forward(
         self._mark_dynamic(seq_ctx)
 
         for idx, decoder_layer in self.layers.items():
-            hidden_states = decoder_layer(
+            layer_results: DenseDecoderLayerOutput = decoder_layer(
                 hidden_states,
                 position_embeddings=position_embeddings,
                 seq_ctx=seq_ctx,
             )
+            hidden_states = layer_results["hidden_states"]
 
             if self.config.return_hidden_states:
                 output["hidden_states"].append(hidden_states)
@@ -235,18 +235,21 @@ def fully_shard(
             layer = self.layers[str(int(layer_idx))]
             layer_idx = int(layer_idx)
             if layer_idx < num_recompute_layers:
-                layer = checkpoint_wrapper(
-                    layer, preserve_rng_state=checkpoint_preserve_rng_state, checkpoint_impl=CheckpointImpl.REENTRANT
+                layer = apply_gradient_checkpointing(
+                    layer,
+                    preserve_rng_state=checkpoint_preserve_rng_state,
                 )
-                # __class__ without self attribute
 
                 # Linear-attention (GatedDeltaNet) layers write ``seq_ctx.seq_idx`` inside the
-                # checkpoint region; compiling the wrapped layer with ``fullgraph=True`` turns the
-                # checkpoint into a HigherOrderOperator that rejects that side effect. Such layers are
-                # still compiled, but with ``fullgraph=False`` so the write can graph-break.
+                # checkpoint region; compiling the checkpointed layer with ``fullgraph=True`` turns
+                # the checkpoint into a HigherOrderOperator that rejects that side effect. Such
+                # layers are still compiled, but with ``fullgraph=False`` so the write can
+                # graph-break.
                 if self.compile_cfg:
                     fullgraph = self.config.layers_type[layer_idx] != "linear_attention"
-                    layer.forward = torch.compile(layer.forward, fullgraph=fullgraph)
+                    # Compile the class function, then restore descriptor binding on this instance.
+                    compiled_forward = torch.compile(type(layer).forward, fullgraph=fullgraph)
+                    layer.forward = compiled_forward.__get__(layer, type(layer))
 
             self.layers[str(layer_idx)] = layer
             self._fully_shard(
diff --git a/xtuner/v1/model/dense/qwen3vl_text.py b/xtuner/v1/model/dense/qwen3vl_text.py
index f41b760ca6..b10f6ef57f 100644
--- a/xtuner/v1/model/dense/qwen3vl_text.py
+++ b/xtuner/v1/model/dense/qwen3vl_text.py
@@ -6,6 +6,7 @@
 from xtuner.v1.data_proto import SequenceContext
 from xtuner.v1.loss import BaseLossContext
 from xtuner.v1.model.base import ModelOutputs
+from xtuner.v1.module.decoder_layer.dense_decoder_layer import DenseDecoderLayerOutput
 
 from .qwen3 import Qwen3Dense, Qwen3Dense4BConfig, Qwen3Dense8BConfig
 
@@ -64,11 +65,12 @@ def forward(  # type: ignore[override]
         # =====================================================
 
         for idx, decoder_layer in self.layers.items():
-            hidden_states = decoder_layer(
+            layer_results: DenseDecoderLayerOutput = decoder_layer(
                 hidden_states,
                 position_embeddings=position_embeddings,
                 seq_ctx=seq_ctx,
             )
+            hidden_states = layer_results["hidden_states"]
 
             if deepstack_visual_embeds is not None and ((idx := int(idx)) in range(len(deepstack_visual_embeds))):
                 assert visual_pos_masks is not None
diff --git a/xtuner/v1/model/moe/moe.py b/xtuner/v1/model/moe/moe.py
index f27e0a2dbc..1851703041 100644
--- a/xtuner/v1/model/moe/moe.py
+++ b/xtuner/v1/model/moe/moe.py
@@ -11,7 +11,6 @@
 from pydantic import ConfigDict
 from torch import nn
 from torch.distributed._functional_collectives import all_reduce
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl
 from torch.distributed.device_mesh import DeviceMesh, init_device_mesh
 from torch.distributed.distributed_c10d import ReduceOp
 from torch.distributed.fsdp import (
@@ -47,9 +46,8 @@
 )
 from xtuner.v1.model.utils import (
     ModelForwardExtraLogInfo,
-    checkpoint_wrapper,
+    apply_gradient_checkpointing,
     module_dict_repr,
-    pytree_reentrant_checkpoint,
 )
 from xtuner.v1.module import (
     GatedDeltaNetConfig,
@@ -61,8 +59,19 @@
     NoAuxRouterConfig,
     RMSNorm,
 )
-from xtuner.v1.module.decoder_layer.dense_decoder_layer import DenseDecoderLayer
-from xtuner.v1.module.decoder_layer.moe_decoder_layer import MoEActFnConfig, MoEBlock, MoEDecoderLayer, MoEGate
+from xtuner.v1.module.decoder_layer.dense_decoder_layer import (
+    DenseDecoderLayer,
+    DenseDecoderLayerMicroBatchOutput,
+    DenseDecoderLayerOutput,
+)
+from xtuner.v1.module.decoder_layer.moe_decoder_layer import (
+    MoEActFnConfig,
+    MoEBlock,
+    MoEDecoderLayer,
+    MoEDecoderLayerMicroBatchOutput,
+    MoEDecoderLayerOutput,
+    MoEGate,
+)
 from xtuner.v1.module.mtp import MTPBlock, MTPConfig, MTPLayer
 from xtuner.v1.utils import (
     get_device,
@@ -81,8 +90,9 @@
 logger = get_logger()
 
 
+MOE_BLOCK_FORWARD = "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEBlock.forward"
 MOE_NON_EP_COMPILE_CFG: dict[str, TorchCompileOption] = {
-    "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEBlock.forward": TorchCompileOption(fullgraph=True),
+    MOE_BLOCK_FORWARD: TorchCompileOption(fullgraph=True),
     "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer.forward": TorchCompileOption(fullgraph=True),
     "xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer._pre_moe_forward": TorchCompileOption(
         fullgraph=True
@@ -555,13 +565,15 @@ def _micro_batch_forward(
             if layer_idx < self.config.first_k_dense_replace:
                 # Keep each micro-batch in its own SequenceContext while issuing
                 # one outer layer call, so FSDP materializes dense weights once.
-                hidden_states_list = list(
+                dense_results = cast(
+                    DenseDecoderLayerMicroBatchOutput,
                     decoder_layer(
-                        *hidden_states_list,
+                        hidden_states_list,
                         position_embeddings=position_embeddings_list,
                         seq_ctx=seq_ctx_list,
-                    )
+                    ),
                 )
+                hidden_states_list = dense_results["hidden_states"]
             else:
                 if int(os.getenv("XTUNER_ACTIVATION_OFFLOAD", "0")) == 1:
                     with async_save_on_cpu(
@@ -574,26 +586,31 @@ def _micro_batch_forward(
                         prefetch=True,
                         reserve_pin_memory=True,
                     ):
-                        layer_results = decoder_layer(
-                            *hidden_states_list,
-                            position_embeddings=position_embeddings_list,
-                            seq_ctx=seq_ctx_list,
+                        layer_results = cast(
+                            MoEDecoderLayerMicroBatchOutput,
+                            decoder_layer(
+                                hidden_states_list,
+                                position_embeddings=position_embeddings_list,
+                                seq_ctx=seq_ctx_list,
+                            ),
                         )
                 else:
-                    layer_results = decoder_layer(
-                        *hidden_states_list,
-                        position_embeddings=position_embeddings_list,
-                        seq_ctx=seq_ctx_list,
+                    layer_results = cast(
+                        MoEDecoderLayerMicroBatchOutput,
+                        decoder_layer(
+                            hidden_states_list,
+                            position_embeddings=position_embeddings_list,
+                            seq_ctx=seq_ctx_list,
+                        ),
                     )
-                hidden_states = layer_results[: len(hidden_states_list)]
-                router_logits = layer_results[len(hidden_states_list) : len(hidden_states_list) * 2]
-                router_weights = layer_results[len(hidden_states_list) * 2 : len(hidden_states_list) * 3]
-                router_topk_ids = layer_results[len(hidden_states_list) * 3 :]
+                router_logits = layer_results["router_logits"]
+                router_weights = layer_results["router_weights"]
+                router_topk_ids = layer_results["router_topk_ids"]
 
                 # Update hidden states and (optionally) collect router logits.
                 # router_weights are only consumed by aux_loss.accumulate below, so we
                 # never stash them per-MB the way we do for logits.
-                for i, hidden_states in enumerate(hidden_states):
+                for i, hidden_states in enumerate(layer_results["hidden_states"]):
                     hidden_states_list[i] = hidden_states
                     if keep_router:
                         router_logits_list[i][f"layer{idx}"] = self._maybe_offload_router(router_logits[i])
@@ -640,7 +657,7 @@ def _micro_batch_forward(
                 )
 
             mtp_outputs_per_mb = self.mtp_block(
-                *hidden_states_list,
+                hidden_states_list,
                 embed_tokens_fn=self.embed_tokens,
                 position_embeddings=position_embeddings_list,
                 seq_ctx=mtp_seq_ctx_list,
@@ -655,12 +672,11 @@ def _micro_batch_forward(
 
                 micro_batch_mtp_losses = torch.tensor(0.0, device=DEVICE)
                 for mtp_idx, (mtp_hidden, mtp_ctx) in enumerate(zip(mtp_outputs, mtp_loss_ctx_list)):
-                    mtp_hidden_states, mtp_router_results, _, _ = mtp_hidden
-                    mtp_loss, _ = self.lm_head(mtp_hidden_states, cast(MTPLossContext, mtp_ctx))
+                    mtp_loss, _ = self.lm_head(mtp_hidden["hidden_states"], cast(MTPLossContext, mtp_ctx))
                     micro_batch_mtp_losses += mtp_loss
 
                     if keep_router:
-                        router_logits_list[micro_batch_idx][f"mtp_layer{mtp_idx}"] = mtp_router_results
+                        router_logits_list[micro_batch_idx][f"mtp_layer{mtp_idx}"] = mtp_hidden["router_logits"]
 
                 mtp_losses += micro_batch_mtp_losses / len(mtp_loss_ctx_list)
                 has_mtp_loss = True
@@ -681,13 +697,13 @@ def _micro_batch_forward(
                 # loss already rides on, so backward traverses each MTP aux node exactly once.
                 for mtp_idx in range(self.config.mtp_config.num_layers):
                     cat_mtp_router_weights = torch.cat(
-                        [mb_outputs[mtp_idx][2] for mb_outputs in mtp_outputs_per_mb], dim=0
+                        [mb_outputs[mtp_idx]["router_weights"] for mb_outputs in mtp_outputs_per_mb], dim=0
                     )
                     cat_mtp_router_logits = torch.cat(
-                        [mb_outputs[mtp_idx][1] for mb_outputs in mtp_outputs_per_mb], dim=0
+                        [mb_outputs[mtp_idx]["router_logits"] for mb_outputs in mtp_outputs_per_mb], dim=0
                     )
                     cat_mtp_router_topk_ids = torch.cat(
-                        [mb_outputs[mtp_idx][3] for mb_outputs in mtp_outputs_per_mb], dim=0
+                        [mb_outputs[mtp_idx]["router_topk_ids"] for mb_outputs in mtp_outputs_per_mb], dim=0
                     )
                     hidden_states_list[0] = self.aux_loss.accumulate(
                         selected_router_weights=cat_mtp_router_weights.index_select(0, nonpad_indices)
@@ -745,8 +761,7 @@ def _micro_batch_forward(
                 layer_router_logits_list: list[torch.Tensor] = []
                 for micro_batch_idx in range(len(seq_ctx_list)):
                     layer_router_logits_list.append(router_logits_list[micro_batch_idx][layer_name].detach())
-                router_logits = torch.stack(layer_router_logits_list, dim=0).unsqueeze(0)
-                router_logits_dict[layer_name] = router_logits
+                router_logits_dict[layer_name] = torch.stack(layer_router_logits_list, dim=0).unsqueeze(0)
 
             output["router_logits"] = router_logits_dict
 
@@ -801,11 +816,15 @@ def _forward(
 
         for idx, decoder_layer in self.layers.items():
             if int(idx) < self.config.first_k_dense_replace:
-                hidden_states = decoder_layer(
-                    hidden_states,
-                    position_embeddings=position_embeddings,
-                    seq_ctx=seq_ctx,
+                dense_results = cast(
+                    DenseDecoderLayerOutput,
+                    decoder_layer(
+                        hidden_states,
+                        position_embeddings=position_embeddings,
+                        seq_ctx=seq_ctx,
+                    ),
                 )
+                hidden_states = dense_results["hidden_states"]
             else:
                 if int(os.getenv("XTUNER_ACTIVATION_OFFLOAD", "0")) == 1:
                     with async_save_on_cpu(
@@ -815,25 +834,34 @@ def _forward(
                         group="text",
                         custom_check_fn=lambda x: x.data_ptr() == hidden_states.data_ptr(),
                     ):
-                        layer_results = decoder_layer(
-                            hidden_states,
-                            position_embeddings=position_embeddings,
-                            seq_ctx=seq_ctx,
+                        layer_results = cast(
+                            MoEDecoderLayerOutput,
+                            decoder_layer(
+                                hidden_states,
+                                position_embeddings=position_embeddings,
+                                seq_ctx=seq_ctx,
+                            ),
                         )
 
                 else:
-                    layer_results = decoder_layer(
-                        hidden_states,
-                        position_embeddings=position_embeddings,
-                        seq_ctx=seq_ctx,
+                    layer_results = cast(
+                        MoEDecoderLayerOutput,
+                        decoder_layer(
+                            hidden_states,
+                            position_embeddings=position_embeddings,
+                            seq_ctx=seq_ctx,
+                        ),
                     )
-                hidden_states, router_results, router_weights, router_topk_ids = layer_results
+                hidden_states = layer_results["hidden_states"]
+                router_logits = layer_results["router_logits"]
+                router_weights = layer_results["router_weights"]
+                router_topk_ids = layer_results["router_topk_ids"]
                 if keep_router:
-                    output["router_logits"][f"layer{idx}"] = self._maybe_offload_router(router_results)
+                    output["router_logits"][f"layer{idx}"] = self._maybe_offload_router(router_logits)
                     output["router_weights"][f"layer{idx}"] = self._maybe_offload_router(router_weights)
                 hidden_states = self.aux_loss.accumulate(
                     selected_router_weights=router_weights.index_select(0, nonpad_indices).contiguous().float(),
-                    selected_router_logits=router_results.index_select(0, nonpad_indices).contiguous().float(),
+                    selected_router_logits=router_logits.index_select(0, nonpad_indices).contiguous().float(),
                     selected_experts=router_topk_ids.index_select(0, nonpad_indices).contiguous(),
                     hidden_states=hidden_states,
                     balancing_ctx=balancing_ctx,
@@ -888,10 +916,13 @@ def _forward(
             # Compute MTP losses for each depth
             mtp_losses = torch.tensor(0.0, device=DEVICE)
             for idx, (mtp_hidden, mtp_ctx) in enumerate(zip(mtp_outputs, mtp_loss_ctx_list)):
-                mtp_hidden_states, mtp_router_results, mtp_router_weights, mtp_router_topk_ids = mtp_hidden
+                mtp_hidden_states = mtp_hidden["hidden_states"]
+                mtp_router_logits = mtp_hidden["router_logits"]
+                mtp_router_weights = mtp_hidden["router_weights"]
+                mtp_router_topk_ids = mtp_hidden["router_topk_ids"]
 
                 if keep_router:
-                    output["router_logits"][f"mtp_layer{idx}"] = mtp_router_results
+                    output["router_logits"][f"mtp_layer{idx}"] = mtp_router_logits
                     output["router_weights"][f"mtp_layer{idx}"] = mtp_router_weights
                 # Inject this MTP layer's z-loss before lm_head so backward through mtp_loss
                 # traverses the AuxLossScaler node and releases this layer's logsumexp activations.
@@ -899,7 +930,7 @@ def _forward(
                     selected_router_weights=mtp_router_weights.index_select(0, mtp_nonpad_indices)
                     .contiguous()
                     .float(),
-                    selected_router_logits=mtp_router_results.index_select(0, mtp_nonpad_indices).contiguous().float(),
+                    selected_router_logits=mtp_router_logits.index_select(0, mtp_nonpad_indices).contiguous().float(),
                     selected_experts=mtp_router_topk_ids.index_select(0, mtp_nonpad_indices).contiguous(),
                     hidden_states=mtp_hidden_states,
                     balancing_ctx=balancing_ctx,
@@ -1152,7 +1183,7 @@ def fully_shard(
                 layer_idx=layer_idx,
                 mtp_idx=None,
             ):
-                layer = checkpoint_wrapper(layer, checkpoint_impl=CheckpointImpl.REENTRANT)
+                layer = apply_gradient_checkpointing(layer)
 
             self.layers[str(layer_idx)] = layer
             if layer_idx >= len(self.layers) - 1 and self.mtp_block is None:
@@ -1204,39 +1235,7 @@ def fully_shard(
                 if self._should_recompute(None, mtp_idx=mtp_idx) or (
                     self.config.mtp_config is not None and self.config.mtp_config.share_weights
                 ):  # share mtp head must recompute
-                    # MTP 默认使用 reentrant 的原因：
-                    #   Case 1：最小触发条件是 compile, topk offload, MTP share weights and depth > 1.
-                    #   多个 logical depth 共用 top-k cache。reentrant 的 original
-                    #   关闭 grad、replay 开启 grad，DSA 能据此正确更新 cache 计数。
-                    #   original 不建立内部图，所以 replay 可以安全复用离散 top-k。
-                    #   non-reentrant 的两次执行都开启 grad，却仍沿用该复用策略，
-                    #   因而出现 original=COMPUTE、replay=REUSE，无法重建相同清单。
-                    #
-                    #   indexer 本身始终 no_grad。不开 compile 时，多执行/少执行一次
-                    #   indexer 不会改变 eager autograd 的保存清单；开启 compile 后，
-                    #   COMPUTE/REUSE 经过不同 graph break 和 compiled block，才可能让
-                    #   checkpoint 保存槽位错位并报 different metadata。例如 original
-                    #   保存 [A, B, C]、replay 保存 [A, X, C] 时，槽位 1 的 metadata
-                    #   不同。后续若显式记录 ORIGINAL/REPLAY phase，可再让
-                    #   non-reentrant 正确推进 cache 状态。
-                    #
-                    # 使用 reentrant 时还必须用 pytree_reentrant_checkpoint：
-                    #   Case 2：触发条件是 EP > 1, intra-layer micro-batch > 1（例如 micro2）.
-                    #   micro2 传入 [embedding_0, embedding_1]；pytree 把 list 内 Tensor
-                    #   展开后，checkpoint 才能在 replay 前逐个 detach，并在 backward
-                    #   中把梯度交回原始 embedding graph。
-                    use_reentrant = self.fsdp_config.mtp_checkpoint_use_reentrant
-                    if use_reentrant:
-                        mtp_layer = checkpoint_wrapper(
-                            mtp_layer,
-                            checkpoint_impl=CheckpointImpl.REENTRANT,
-                            checkpoint_fn=pytree_reentrant_checkpoint,
-                        )
-                    else:
-                        mtp_layer = checkpoint_wrapper(
-                            mtp_layer,
-                            checkpoint_impl=CheckpointImpl.NO_REENTRANT,
-                        )
+                    mtp_layer = apply_gradient_checkpointing(mtp_layer)
                 self.mtp_block.layers[mtp_idx] = mtp_layer
 
                 reshard_after_forward = mtp_idx != len(self.mtp_block.layers) - 1
@@ -1277,8 +1276,7 @@ def fully_shard(
     def default_compile_cfg(self) -> dict[str, TorchCompileOption]:
         if self.config.ep_size > 1:
             return MOE_EP_COMPILE_CFG
-        else:
-            return MOE_NON_EP_COMPILE_CFG
+        return MOE_NON_EP_COMPILE_CFG
 
     @property
     def need_update_bias(self) -> bool:
diff --git a/xtuner/v1/model/moe/qwen3vl_text.py b/xtuner/v1/model/moe/qwen3vl_text.py
index 5451c31346..b8627a2611 100644
--- a/xtuner/v1/model/moe/qwen3vl_text.py
+++ b/xtuner/v1/model/moe/qwen3vl_text.py
@@ -4,6 +4,8 @@
 import torch
 
 from xtuner.v1.data_proto import SequenceContext
+from xtuner.v1.module.decoder_layer.dense_decoder_layer import DenseDecoderLayerOutput
+from xtuner.v1.module.decoder_layer.moe_decoder_layer import MoEDecoderLayerOutput
 from xtuner.v1.utils.activation_offload import async_save_on_cpu
 
 from .moe import MoELossContextDict, MoEModelOutputs
@@ -157,11 +159,12 @@ def _forward(
 
         for idx, decoder_layer in self.layers.items():
             if int(idx) < self.config.first_k_dense_replace:
-                hidden_states = decoder_layer(
+                dense_results: DenseDecoderLayerOutput = decoder_layer(
                     hidden_states,
                     position_embeddings=position_embeddings,
                     seq_ctx=seq_ctx,
                 )
+                hidden_states = dense_results["hidden_states"]
             else:
                 if int(os.getenv("XTUNER_ACTIVATION_OFFLOAD", "0")) == 1:
                     offload_stream = decoder_layer._get_fsdp_state()._comm_ctx.all_gather_stream
@@ -172,25 +175,29 @@ def _forward(
                         depth=len(self.layers),
                         custom_check_fn=lambda x: x.data_ptr() == hidden_states.data_ptr(),
                     ):
-                        hidden_states, router_results, router_weights, router_topk_ids = decoder_layer(
+                        layer_results: MoEDecoderLayerOutput = decoder_layer(
                             hidden_states,
                             position_embeddings=position_embeddings,
                             seq_ctx=seq_ctx,
                         )
 
                 else:
-                    hidden_states, router_results, router_weights, router_topk_ids = decoder_layer(
+                    layer_results = decoder_layer(
                         hidden_states,
                         position_embeddings=position_embeddings,
                         seq_ctx=seq_ctx,
                     )
+                hidden_states = layer_results["hidden_states"]
+                router_logits = layer_results["router_logits"]
+                router_weights = layer_results["router_weights"]
+                router_topk_ids = layer_results["router_topk_ids"]
 
                 if keep_router:
-                    output["router_logits"][f"layer{idx}"] = router_results
+                    output["router_logits"][f"layer{idx}"] = router_logits
                     output["router_weights"][f"layer{idx}"] = router_weights
                 hidden_states = self.aux_loss.accumulate(
                     selected_router_weights=router_weights.index_select(0, nonpad_indices).contiguous().float(),
-                    selected_router_logits=router_results.index_select(0, nonpad_indices).contiguous().float(),
+                    selected_router_logits=router_logits.index_select(0, nonpad_indices).contiguous().float(),
                     selected_experts=router_topk_ids.index_select(0, nonpad_indices).contiguous(),
                     hidden_states=hidden_states,
                     balancing_ctx=balancing_ctx,
diff --git a/xtuner/v1/model/utils/__init__.py b/xtuner/v1/model/utils/__init__.py
index ba77e7c3c0..a398224afd 100644
--- a/xtuner/v1/model/utils/__init__.py
+++ b/xtuner/v1/model/utils/__init__.py
@@ -1,5 +1,9 @@
-from .checkpointing import checkpoint_wrapper, pytree_reentrant_checkpoint
+from .checkpointing import apply_gradient_checkpointing
 from .misc import ModelForwardExtraLogInfo, module_dict_repr
 
 
-__all__ = ["checkpoint_wrapper", "pytree_reentrant_checkpoint", "module_dict_repr", "ModelForwardExtraLogInfo"]
+__all__ = [
+    "apply_gradient_checkpointing",
+    "module_dict_repr",
+    "ModelForwardExtraLogInfo",
+]
diff --git a/xtuner/v1/model/utils/checkpointing.py b/xtuner/v1/model/utils/checkpointing.py
index f727b8232d..9b4c33323b 100644
--- a/xtuner/v1/model/utils/checkpointing.py
+++ b/xtuner/v1/model/utils/checkpointing.py
@@ -1,120 +1,112 @@
-import inspect
-from types import UnionType
-from typing import Any, Callable, Union, get_args, get_origin
+"""Gradient checkpointing (activation recomputation) entry points."""
+
+from contextlib import AbstractContextManager
+from functools import partial
+from typing import Any, Callable
 
-import torch
 import torch.nn as nn
-from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (
-    checkpoint_wrapper as ptd_checkpoint_wrapper,
-)
-from torch.utils._pytree import tree_flatten, tree_unflatten
+from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl, checkpoint_wrapper
+from torch.utils._pytree import TreeSpec, tree_flatten, tree_unflatten
 from torch.utils.checkpoint import checkpoint
 
-from xtuner.v1.utils import copy_signature
-
-
-# TODO: Currently xtuner uses the internal, outdated `torch.distributed.algorithms._checkpoint.checkpoint_wrapper` interface
-# We should look for opportunities to use the public, updated interface in the future
-
-# NOTE:
-# PyTorch's `torch.distributed.algorithms._checkpoint.checkpoint_wrapper` has some limitations. Modules decorated with `checkpoint_wrapper`
-# must have forward interfaces that conform to the specifications of `torch.autograd.function.Function`.
-# Specifically, for input parameters, the `forward` interface must explicitly accept parameters of type `torch.Tensor` to ensure proper gradient backpropagation.
-# For return values, the `forward` interface must return either `torch.Tensor` or tuple[torch.Tensor, ...].
-# For example If the forward interface is declared as:
-# def forward(self, x: tuple[torch.Tensor], y: list[torch.Tensor]) -> torch.Tensor | tuple[torch.Tensor, ...]:
-# This interface will break the gradient graph because the inputs don't meet the requirements. For instance, x and y are not of type `torch.Tensor`
-# `_check_signature_of_forward` will check (not exhaustively) whether the signature of the `forward` interface meets the requirements to identify issues early.
-
-
-def _check_signature_of_forward(module: nn.Module):
-    def _is_tensor_or_tuple_tensor(arg_type: type):
-        if arg_type is torch.Tensor:
-            return True
-
-        origin_type = get_origin(arg_type)
-
-        if not origin_type or origin_type not in (tuple, UnionType, Union):
-            return False
-
-        if origin_type in [UnionType, Union]:
-            type_list = get_args(arg_type)
-            return any(_is_tensor_or_tuple_tensor(t) for t in type_list)
-
-        else:
-            type_list = get_args(arg_type)
-            return any(t is torch.Tensor for t in type_list)
-
-    def _has_missing_type(arg_type: type):
-        if arg_type is inspect._empty:
-            return True
-        origin_arg = get_origin(arg_type)
-        return any(_has_missing_type(t) for t in get_args(origin_arg))
-
-    input_type = inspect.signature(module.forward).parameters
-    ret_type = inspect.signature(module.forward).return_annotation
-
-    for name, arg_type in input_type.items():
-        if _has_missing_type(arg_type.annotation):
-            raise TypeError(
-                f"The type of argument '{name}' of {module.__class__.__name__}.forward must be annotated, but got "
-                f"{name} unannotated."
-            )
-
-    if _has_missing_type(ret_type):
-        raise TypeError(
-            f"The return type of {module.__class__.__name__}.forward must be annotated, but got {ret_type}"
-        )
-
-    for arg_type in input_type.values():
-        origin_arg = get_origin(arg_type.annotation)
-        # Union[Tensor, None] or Optional[Tensor] is legal
-        if origin_arg:
-            if torch.Tensor in origin_arg:
-                break
-        else:
-            if arg_type.annotation is torch.Tensor:
-                break
-    else:
-        raise TypeError(
-            f"The type of all arguments of the {module.__class__.__name__}.forward must be torch.Tensor, but got "
-            f"{input_type}"
-        )
-
-    if not _is_tensor_or_tuple_tensor(ret_type):
-        raise TypeError(
-            f"The return type of {module.__class__.__name__}.forward must be torch.Tensor or tuple of torch.Tensor, "
-            f"but got {ret_type}"
-        )
-
-
-@copy_signature(ptd_checkpoint_wrapper)
-def checkpoint_wrapper(module: nn.Module, *args, **kwargs):
-    _check_signature_of_forward(module)
-    return ptd_checkpoint_wrapper(module, *args, **kwargs)
-
-
-def pytree_reentrant_checkpoint(
-    function: Callable[..., torch.Tensor | tuple[torch.Tensor, ...]],
+
+__all__ = ["apply_gradient_checkpointing", "checkpoint_flattened"]
+
+
+ContextFn = Callable[[], tuple[AbstractContextManager, AbstractContextManager]]
+
+
+def apply_gradient_checkpointing(
+    module: nn.Module,
+    *,
+    preserve_rng_state: bool = True,
+    use_reentrant: bool = True,
+    context_fn: ContextFn | None = None,
+) -> nn.Module:
+    """Make ``module``'s forward recomputed during backward instead of kept in
+    memory.
+
+    Inputs and outputs are flattened around reentrant checkpointing so gradients flow through
+    nested containers, ``TypedDict`` returns, and keyword-only arguments. Non-reentrant
+    checkpointing remains available for selective activation checkpointing.
+
+    Args:
+        module (nn.Module): Module whose forward should be recomputed during backward.
+        preserve_rng_state (bool): Restore the RNG state before recomputing, so dropout and other
+            stochastic ops replay identically. Defaults to True.
+        use_reentrant (bool): Use reentrant checkpointing. Defaults to True.
+        context_fn (Callable | None): Factory returning the ``(forward_context, recompute_context)``
+            pair that ``torch.utils.checkpoint.checkpoint`` enters around the two passes. This is
+            the seam for selective checkpointing: passing the contexts built by
+            ``create_selective_checkpoint_contexts`` turns whole-module recompute into a per-op
+            decision. Requires ``use_reentrant=False``. Defaults to None, i.e. recompute everything.
+
+    Returns:
+        nn.Module: An outer checkpoint wrapper containing ``module``.
+    """
+    # FSDP must wrap the checkpoint boundary. Its output hook then runs before
+    # checkpoint replay and unshards the parameters for backward. Putting the
+    # checkpoint around the FSDP module itself makes reentrant replay look like
+    # another forward and corrupts FSDP's default backward-prefetch order.
+    extra: dict[str, Any] = {} if context_fn is None else {"context_fn": context_fn}
+    checkpoint_fn = partial(
+        checkpoint_flattened,
+        preserve_rng_state=preserve_rng_state,
+        use_reentrant=use_reentrant,
+        **extra,
+    )
+    checkpoint_impl = CheckpointImpl.REENTRANT if use_reentrant else CheckpointImpl.NO_REENTRANT
+    return checkpoint_wrapper(module, checkpoint_impl=checkpoint_impl, checkpoint_fn=checkpoint_fn)
+
+
+def checkpoint_flattened(
+    function: Callable[..., Any],
     *args: Any,
+    preserve_rng_state: bool = True,
+    use_reentrant: bool = True,
     **kwargs: Any,
-) -> torch.Tensor | tuple[torch.Tensor, ...]:
-    """让嵌套 Tensor 也成为 reentrant checkpoint 的 autograd 输入。"""
-    # CheckpointWrapper 只打包一层。例如：
-    #   future_embeddings=[embedding_0, embedding_1]
-    # 对原生 CheckpointFunction 来说只是“一个 list 参数”，它看不到 list 里的
-    # 两个 Tensor，也就不会 detach 它们。这可能会造成反向传播的错误，因为这两个
-    # Tensor 的梯度应该交由 CheckpointFunction.backward 的返回值交回原始 Tensor，
-    # 而不是由他们自己来传递梯度。
-    # tree_flatten 会把输入变成近似：
-    #   hidden, embedding_0, embedding_1
-    # 这样 checkpoint 能逐个 detach；tree_unflatten 再在 replay 前把 list 还原。
-    flat_inputs, input_spec = tree_flatten((args, kwargs))
+) -> Any:
+    """Run ``function`` under checkpointing with flattened inputs and outputs.
+
+    Reentrant checkpointing needs both sides flattened so autograd sees tensors nested in container
+    inputs and outputs. Non-reentrant checkpointing handles those structures itself, but still
+    needs flattened inputs so the caller's saved-tensor hooks can see every input tensor.
+    ``_CheckpointFrame.save_inputs`` wraps only *top-level* tensor arguments into a
+    ``SavedVariable``, and constructing one is what fires the ambient ``saved_tensors_hooks``. It
+    runs just before the checkpoint installs its own hooks, so those ambient hooks are still the
+    caller's -- which is how activation offloading gets hold of a layer's inputs. A tensor nested in
+    a list, or passed by keyword, is stored as a plain reference instead, reaches no hook, and is
+    silently never offloaded.
+
+    Args:
+        function (Callable): The callable to run inside the checkpointed region.
+        preserve_rng_state (bool): Restore the RNG state before recomputing. Defaults to True.
+        use_reentrant (bool): Use reentrant checkpointing. Defaults to True.
+        **kwargs (Any): Forwarded to ``function``, except ``context_fn`` which goes to
+            ``torch.utils.checkpoint.checkpoint``.
+
+    Returns:
+        Any: Whatever ``function`` returns.
+    """
+    context_fn = kwargs.pop("context_fn", None)
+    checkpoint_kwargs: dict[str, Any] = {
+        "use_reentrant": use_reentrant,
+        "preserve_rng_state": preserve_rng_state,
+    }
+    if context_fn is not None:
+        checkpoint_kwargs["context_fn"] = context_fn
 
-    def run_function(*replayed_flat_inputs: Any) -> torch.Tensor | tuple[torch.Tensor, ...]:
-        # 这里只还原参数结构，不会把 detached Tensor 重新连接到旧 graph；梯度由
-        # CheckpointFunction.backward 的返回值交回原始 Tensor。
-        replayed_args, replayed_kwargs = tree_unflatten(list(replayed_flat_inputs), input_spec)
-        return function(*replayed_args, **replayed_kwargs)
-
-    return checkpoint(run_function, *flat_inputs, use_reentrant=True)
+    flat_inputs, input_spec = tree_flatten((args, kwargs))
+    output_spec: TreeSpec | None = None
+
+    def call_with_original_signature(*replayed: Any) -> tuple[Any, ...]:
+        nonlocal output_spec
+        replayed_args, replayed_kwargs = tree_unflatten(list(replayed), input_spec)
+        flat_outputs, output_spec = tree_flatten(function(*replayed_args, **replayed_kwargs))
+        return tuple(flat_outputs)
+
+    flat_outputs = checkpoint(call_with_original_signature, *flat_inputs, **checkpoint_kwargs)
+    assert output_spec is not None, "XTuner Internal Error: checkpoint did not run the function"
+    if not isinstance(flat_outputs, tuple):
+        flat_outputs = (flat_outputs,)
+    return tree_unflatten(list(flat_outputs), output_spec)
diff --git a/xtuner/v1/module/attention/dsa_topk_sharing.py b/xtuner/v1/module/attention/dsa_topk_sharing.py
index 7e0ca4fd27..b74963f7e4 100644
--- a/xtuner/v1/module/attention/dsa_topk_sharing.py
+++ b/xtuner/v1/module/attention/dsa_topk_sharing.py
@@ -514,8 +514,8 @@ def register_dsa_topk_decoder_lifecycle_hooks(decoder_layer: torch.nn.Module) ->
     # recorded only pending actions and flushed them later. Remove that
     # transient state by keeping the entire residency transition at the decoder
     # boundary: the pre-hook launches H2D and the post-hook directly runs
-    # after_sparse_mla_use. Reentrant checkpoint replay invokes the decoder
-    # module and these hooks again, so main, micro-batch and MTP callers do not
+    # after_sparse_mla_use. Checkpoint recompute re-invokes the decoder module
+    # and these hooks along with it, so main, micro-batch and MTP callers do not
     # need separate lifecycle handling.
     #
     # This deliberately delays eager D2H until the decoder returns, losing its
diff --git a/xtuner/v1/module/decoder_layer/dense_decoder_layer.py b/xtuner/v1/module/decoder_layer/dense_decoder_layer.py
index 426e353b92..aa5660ea09 100644
--- a/xtuner/v1/module/decoder_layer/dense_decoder_layer.py
+++ b/xtuner/v1/module/decoder_layer/dense_decoder_layer.py
@@ -1,4 +1,4 @@
-from typing import Literal
+from typing import Literal, TypedDict
 
 import torch
 import torch.nn as nn
@@ -14,6 +14,27 @@
 from ..linear import build_linear
 
 
+class DenseDecoderLayerOutput(TypedDict):
+    """Per-micro-batch outputs of one :class:`DenseDecoderLayer` forward.
+
+    A dense layer only produces hidden states, but it reports them through the same keyed contract
+    as :class:`~xtuner.v1.module.decoder_layer.moe_decoder_layer.MoEDecoderLayer` so that the two
+    layer families stay interchangeable to their callers.
+    """
+
+    hidden_states: torch.Tensor
+
+
+class DenseDecoderLayerMicroBatchOutput(TypedDict):
+    """Outputs of one :class:`DenseDecoderLayer` forward over several micro-
+    batches.
+
+    Each field holds one entry per micro-batch, in input order.
+    """
+
+    hidden_states: list[torch.Tensor]
+
+
 class DenseMLP(nn.Module):
     def __init__(
         self,
@@ -74,36 +95,54 @@ def __init__(
 
     def forward(
         self,
-        *hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor | list[torch.Tensor],
+        *,
         position_embeddings: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]],
         seq_ctx: SequenceContext | list[SequenceContext],
-    ) -> torch.Tensor | tuple[torch.Tensor, ...]:
+    ) -> DenseDecoderLayerOutput | DenseDecoderLayerMicroBatchOutput:
         """Run equal-shaped training micro-batches in one layer invocation.
 
-        Keeping the micro-batch loop inside the decoder layer lets outer FSDP
-        and checkpoint wrappers materialize the layer only once, while each
-        attention call keeps its own ``SequenceContext``.
+        Keeping the micro-batch loop inside the decoder layer lets outer FSDP and checkpointing
+        materialize the layer only once, while each attention call keeps its own
+        ``SequenceContext``.
+
+        Args:
+            hidden_states (torch.Tensor | list[torch.Tensor]): Input hidden states, one tensor per
+                micro-batch.
+            position_embeddings (tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]]):
+                Rotary position embeddings ``(cos, sin)``, aligned with ``hidden_states``.
+            seq_ctx (SequenceContext | list[SequenceContext]): Sequence context, aligned with
+                ``hidden_states``.
+
+        Returns:
+            DenseDecoderLayerOutput | DenseDecoderLayerMicroBatchOutput: Output hidden states. A
+            single tensor for a single ``hidden_states`` tensor, a per-micro-batch list for a list
+            of them.
         """
-        if len(hidden_states) == 1:
+        if not isinstance(hidden_states, list):
             assert isinstance(position_embeddings, tuple) and len(position_embeddings) == 2
             assert isinstance(seq_ctx, SequenceContext)
-            return self._forward(
-                hidden_states=hidden_states[0],
-                position_embeddings=position_embeddings,
-                seq_ctx=seq_ctx,
-            )
+            return {
+                "hidden_states": self._forward(
+                    hidden_states=hidden_states,
+                    position_embeddings=position_embeddings,
+                    seq_ctx=seq_ctx,
+                )
+            }
 
         assert isinstance(position_embeddings, list) and len(position_embeddings) == len(hidden_states)
         assert isinstance(seq_ctx, list) and len(seq_ctx) == len(hidden_states)
         assert all(hidden.shape == hidden_states[0].shape for hidden in hidden_states)
-        return tuple(
-            self._forward(
-                hidden_states=hidden,
-                position_embeddings=position_embedding,
-                seq_ctx=context,
-            )
-            for hidden, position_embedding, context in zip(hidden_states, position_embeddings, seq_ctx)
-        )
+        return {
+            "hidden_states": [
+                self._forward(
+                    hidden_states=hidden,
+                    position_embeddings=position_embedding,
+                    seq_ctx=context,
+                )
+                for hidden, position_embedding, context in zip(hidden_states, position_embeddings, seq_ctx)
+            ]
+        }
 
     def _forward(
         self,
diff --git a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py
index 00e5d6c27e..0693ca95bd 100644
--- a/xtuner/v1/module/decoder_layer/moe_decoder_layer.py
+++ b/xtuner/v1/module/decoder_layer/moe_decoder_layer.py
@@ -1,5 +1,5 @@
 from functools import partial
-from typing import Literal, Protocol, TypeAlias, cast
+from typing import Literal, Protocol, TypeAlias, TypedDict, cast
 
 import torch
 import torch.nn as nn
@@ -47,6 +47,28 @@
 HiddenStates: TypeAlias = torch.Tensor
 
 
+class MoEDecoderLayerOutput(TypedDict):
+    """Per-micro-batch outputs of one :class:`MoEDecoderLayer` forward."""
+
+    hidden_states: HiddenStates
+    router_logits: RouterLogits
+    router_weights: RouterWeights
+    router_topk_ids: RouterTopKIds
+
+
+class MoEDecoderLayerMicroBatchOutput(TypedDict):
+    """Outputs of one :class:`MoEDecoderLayer` forward over several micro-
+    batches (domino EP).
+
+    Each field holds one entry per micro-batch, in input order.
+    """
+
+    hidden_states: list[HiddenStates]
+    router_logits: list[RouterLogits]
+    router_weights: list[RouterWeights]
+    router_topk_ids: list[RouterTopKIds]
+
+
 class MoEActFnProtocol(Protocol):
     def __call__(self, fused_x: torch.Tensor, split_dim: int = -1) -> torch.Tensor: ...
 
@@ -291,23 +313,31 @@ def __init__(
 
     def forward(
         self,
-        *hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor | list[torch.Tensor],
+        *,
         seq_ctx: SequenceContext | list[SequenceContext],
-        position_embeddings: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]] | None = None,
-    ) -> tuple[HiddenStates, RouterLogits, RouterWeights, RouterTopKIds] | tuple[torch.Tensor, ...]:
+        position_embeddings: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]],
+    ) -> MoEDecoderLayerOutput | MoEDecoderLayerMicroBatchOutput:
         """Forward pass of the MoE decoder layer.
 
+        Passing lists runs several equal-shaped micro-batches in one layer invocation (domino EP),
+        so that the expert dispatch/combine communication of one micro-batch overlaps the expert
+        compute of another.
+
         Args:
-            hidden_states (torch.Tensor): Input hidden states.
-            seq_ctx (SequenceContext): Sequence context.
-            position_embeddings (tuple[torch.Tensor, torch.Tensor]): Position embeddings.
-            past_key_values (list[list[torch.Tensor]], optional): Past key values for pre-filling or decoding.
+            hidden_states (torch.Tensor | list[torch.Tensor]): Input hidden states, one tensor per
+                micro-batch.
+            seq_ctx (SequenceContext | list[SequenceContext]): Sequence context, aligned with
+                ``hidden_states``.
+            position_embeddings (tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]]):
+                Rotary position embeddings ``(cos, sin)``, aligned with ``hidden_states``.
 
         Returns:
-            tuple: Output hidden states, router logits, router weights, and the
-                expert IDs selected by the router.
+            MoEDecoderLayerOutput | MoEDecoderLayerMicroBatchOutput: Hidden states and router
+            results. Scalar fields for a single ``hidden_states`` tensor, per-micro-batch lists for
+            a list of them.
         """
-        if len(hidden_states) == 1:
+        if not isinstance(hidden_states, list):
             assert isinstance(seq_ctx, SequenceContext), (
                 f"seq_ctx should be a SequenceContext instance but got {seq_ctx}"
             )
@@ -315,7 +345,7 @@ def forward(
                 "position_embeddings should be a tuple of two tensors (position_ids, position_embeds)"
             )
             return self._forward(
-                hidden_states=hidden_states[0],
+                hidden_states=hidden_states,
                 seq_ctx=seq_ctx,
                 position_embeddings=position_embeddings,
             )
@@ -328,7 +358,7 @@ def forward(
             )
 
             return self._micro_batch_forward(
-                hidden_states_list=list(hidden_states),
+                hidden_states_list=hidden_states,
                 seq_ctx_list=seq_ctx,
                 position_embeddings_list=position_embeddings,
             )
@@ -376,7 +406,7 @@ def _forward(
         hidden_states: torch.Tensor,
         seq_ctx: SequenceContext,
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
-    ) -> tuple[HiddenStates, RouterLogits, RouterWeights, RouterTopKIds]:
+    ) -> MoEDecoderLayerOutput:
         residual, hidden_states, router_results = self._pre_moe_forward(
             hidden_states=hidden_states,
             seq_ctx=seq_ctx,
@@ -410,6 +440,11 @@ def _forward(
         #     post_dispatched.get("row_ids_map"),  # type: ignore[arg-type]
         #     dispatched["topk_weights"],
         # )
+        if self.ep_mesh is not None:
+            # MoEBlock is fullgraph-compiled and shared by all decoder layers. Only the routed-token
+            # dimension varies, so make it dynamic before entering the compile boundary to keep one
+            # AOT Autograd save plan across the original forward and checkpoint replay.
+            torch._dynamo.mark_dynamic(post_dispatched["hidden_states"], 0)
         experts_out = self.experts(
             post_dispatched["hidden_states"],
             post_dispatched["tokens_per_expert"],
@@ -461,19 +496,19 @@ def _forward(
             residual=residual,
             shared_experts_out=shared_experts_out,
         )
-        return (
-            hidden_states,
-            router_results["logits"],
-            router_results["router_weights"],
-            router_results["topk_ids"],
-        )
+        return {
+            "hidden_states": hidden_states,
+            "router_logits": router_results["logits"],
+            "router_weights": router_results["router_weights"],
+            "router_topk_ids": router_results["topk_ids"],
+        }
 
     def _micro_batch_forward(
         self,
         hidden_states_list: list[torch.Tensor],
         seq_ctx_list: list[SequenceContext],
         position_embeddings_list: list[tuple[torch.Tensor, torch.Tensor]],
-    ) -> tuple[torch.Tensor, ...]:
+    ) -> MoEDecoderLayerMicroBatchOutput:
         origin_shape = hidden_states_list[0].shape
         assert all(hidden_states.shape == origin_shape for hidden_states in hidden_states_list), (
             "All hidden states should have the same shape"
@@ -534,6 +569,9 @@ def _micro_batch_forward(
                 dispatched=dispatched,
                 async_op=True,
             )
+            if self.ep_mesh is not None:
+                # Preserve the same dynamic-token compile contract for every in-layer micro-batch.
+                torch._dynamo.mark_dynamic(post_dispatched["hidden_states"], 0)
             experts_out = self.experts(
                 post_dispatched["hidden_states"],
                 post_dispatched["tokens_per_expert"],
@@ -598,10 +636,12 @@ def _micro_batch_forward(
             )
             hidden_states_out_list.append(hidden_states)
 
-        router_logits = [router_results["logits"] for router_results in router_results_list]
-        router_weights = [router_results["router_weights"] for router_results in router_results_list]
-        router_topk_ids = [router_results["topk_ids"] for router_results in router_results_list]
-        return tuple(hidden_states_out_list + router_logits + router_weights + router_topk_ids)
+        return {
+            "hidden_states": hidden_states_out_list,
+            "router_logits": [router_results["logits"] for router_results in router_results_list],
+            "router_weights": [router_results["router_weights"] for router_results in router_results_list],
+            "router_topk_ids": [router_results["topk_ids"] for router_results in router_results_list],
+        }
 
     def _pre_moe_forward(
         self,
diff --git a/xtuner/v1/module/mtp/mtp_block.py b/xtuner/v1/module/mtp/mtp_block.py
index 9d43f685e0..8a0d585b8c 100644
--- a/xtuner/v1/module/mtp/mtp_block.py
+++ b/xtuner/v1/module/mtp/mtp_block.py
@@ -6,13 +6,19 @@
 import torch.nn as nn
 
 from xtuner.v1.data_proto import SequenceContext
+from xtuner.v1.module.decoder_layer.moe_decoder_layer import (
+    MoEDecoderLayerMicroBatchOutput,
+    MoEDecoderLayerOutput,
+)
 
 from .config import MTPConfig
 from .mtp_layer import MTPLayer
 from .utils import roll_sequence_context
 
 
-MTPDepthOutput = tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]
+MTPDepthOutput = MoEDecoderLayerOutput
+"""One MTP depth produces the same keyed outputs as the decoder layer it
+wraps."""
 
 
 class MTPBlock(nn.Module):
@@ -62,12 +68,12 @@ class MTPBlock(nn.Module):
         >>>
         >>> # Multi-microbatch (domino EP) forward
         >>> outputs_per_mb = mtp_block(
-        ...     h0, h1,
+        ...     [h0, h1],
         ...     embed_tokens_fn=embed_fn,
         ...     position_embeddings=[pos_emb_0, pos_emb_1],
         ...     seq_ctx=[ctx_0, ctx_1],
         ... )
-        >>> # outputs_per_mb[mb_idx][depth_idx] -> (hidden, router_logits, router_weights, router_topk_ids)
+        >>> # outputs_per_mb[mb_idx][depth_idx] -> MTPDepthOutput
     """
 
     def __init__(self, *, mtp_config: MTPConfig, mtp_layers: list[MTPLayer]):
@@ -84,7 +90,8 @@ def __init__(self, *, mtp_config: MTPConfig, mtp_layers: list[MTPLayer]):
 
     def forward(
         self,
-        *hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor | list[torch.Tensor],
+        *,
         embed_tokens_fn: Callable[[torch.Tensor], torch.Tensor],
         position_embeddings: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]],
         seq_ctx: SequenceContext | list[SequenceContext],
@@ -97,9 +104,8 @@ def forward(
         the wrapped decoder layer can be overlapped across micro-batches (domino EP).
 
         Args:
-            hidden_states (torch.Tensor): One or more hidden state tensors from the main
-                model, shape ``[batch, seq_len, hidden_size]`` each. Single tensor → single-
-                microbatch path; multiple tensors → multi-microbatch (domino EP) path.
+            hidden_states (torch.Tensor | list[torch.Tensor]): Hidden states from the main model,
+                shape ``[batch, seq_len, hidden_size]`` each, one tensor per micro-batch.
             embed_tokens_fn (Callable): Function to embed tokens. Takes token IDs and returns
                 embeddings. Should have signature ``embed_tokens_fn(token_ids: Tensor) -> Tensor``.
             position_embeddings (tuple | list[tuple]): Rotary position embeddings (cos, sin),
@@ -107,14 +113,11 @@ def forward(
             seq_ctx (SequenceContext | list[SequenceContext]): Sequence context per micro-batch.
 
         Returns:
-            list: For single-microbatch input,
-                ``list[(hidden, router_logits, router_weights, router_topk_ids)]``
-                of length ``D``, where ``outputs[k]`` is the prediction for token ``i+k+1``.
-                For ``N`` micro-batches,
-                ``list[list[(hidden, router_logits, router_weights, router_topk_ids)]]``
-                with outer length ``N`` and inner length ``D``: ``outputs[mb_idx][depth_idx]``.
+            list[MTPDepthOutput] | list[list[MTPDepthOutput]]: For a single ``hidden_states``
+            tensor, one entry per MTP depth ``D``, where ``outputs[k]`` is the prediction for
+            token ``i+k+1``. For ``N`` micro-batches, ``outputs[mb_idx][depth_idx]``.
         """
-        if len(hidden_states) == 1:
+        if not isinstance(hidden_states, list):
             assert isinstance(seq_ctx, SequenceContext), (
                 "seq_ctx should be a SequenceContext instance in single-microbatch mode"
             )
@@ -122,7 +125,7 @@ def forward(
                 "position_embeddings should be a (cos, sin) tuple in single-microbatch mode"
             )
             return self._forward(
-                hidden_states=hidden_states[0],
+                hidden_states=hidden_states,
                 embed_tokens_fn=embed_tokens_fn,
                 position_embeddings=position_embeddings,
                 seq_ctx=seq_ctx,
@@ -136,7 +139,7 @@ def forward(
             "position_embeddings should be a list aligned with hidden_states in multi-microbatch mode"
         )
         return self._micro_batch_forward(
-            hidden_states_list=list(hidden_states),
+            hidden_states_list=hidden_states,
             embed_tokens_fn=embed_tokens_fn,
             position_embeddings_list=position_embeddings,
             seq_ctx_list=seq_ctx,
@@ -164,9 +167,7 @@ def _forward(
                 attention mask, etc.
 
         Returns:
-            list[MTPDepthOutput]: List of 4-tuples
-                (hidden_states, router_logits, router_weights, router_topk_ids)
-                for each MTP depth.
+            list[MTPDepthOutput]: One entry per MTP depth.
                 Length equals num_layers.
                 - outputs[0]: Outputs for predicting token at position (i+1)
                 - outputs[k]: Outputs for predicting token at position (i+k+1)
@@ -186,13 +187,14 @@ def _forward(
             if self.mtp_config.detach_mtp_inputs:
                 future_embeddings = future_embeddings.detach()
 
-            current_hidden_states, router_logits, router_weights, router_topk_ids = layer(
+            layer_results: MTPDepthOutput = layer(
                 current_hidden_states,
                 future_embeddings=future_embeddings,
                 position_embeddings=position_embeddings,
                 seq_ctx=current_seq_ctx,
             )
-            mtp_outputs.append((current_hidden_states, router_logits, router_weights, router_topk_ids))
+            current_hidden_states = layer_results["hidden_states"]
+            mtp_outputs.append(layer_results)
 
         return mtp_outputs
 
@@ -218,27 +220,24 @@ def _micro_batch_forward(
             current_seq_ctx_list = [roll_sequence_context(ctx, shifts=-1) for ctx in current_seq_ctx_list]
             future_embeddings_list = [self._embed_future(ctx, embed_tokens_fn) for ctx in current_seq_ctx_list]
 
-            layer_results = layer(
-                *current_hidden_states_list,
+            layer_results: MoEDecoderLayerMicroBatchOutput = layer(
+                current_hidden_states_list,
                 future_embeddings=future_embeddings_list,
                 position_embeddings=position_embeddings_list,
                 seq_ctx=current_seq_ctx_list,
             )
-            assert isinstance(layer_results, tuple) and len(layer_results) == 4 * n, (
-                f"MTPLayer multi-microbatch forward should return a flat tuple of length {4 * n}, "
-                f"got {len(layer_results) if isinstance(layer_results, tuple) else type(layer_results)}"
-            )
-            new_hidden = list(layer_results[:n])
-            router_logits = list(layer_results[n : 2 * n])
-            router_weights = list(layer_results[2 * n : 3 * n])
-            router_topk_ids = list(layer_results[3 * n :])
 
             for mb_idx in range(n):
                 outputs_per_mb[mb_idx].append(
-                    (new_hidden[mb_idx], router_logits[mb_idx], router_weights[mb_idx], router_topk_ids[mb_idx])
+                    {
+                        "hidden_states": layer_results["hidden_states"][mb_idx],
+                        "router_logits": layer_results["router_logits"][mb_idx],
+                        "router_weights": layer_results["router_weights"][mb_idx],
+                        "router_topk_ids": layer_results["router_topk_ids"][mb_idx],
+                    }
                 )
 
-            current_hidden_states_list = new_hidden
+            current_hidden_states_list = layer_results["hidden_states"]
 
         return outputs_per_mb
 
diff --git a/xtuner/v1/module/mtp/mtp_layer.py b/xtuner/v1/module/mtp/mtp_layer.py
index 711c7f4b29..620500cc92 100644
--- a/xtuner/v1/module/mtp/mtp_layer.py
+++ b/xtuner/v1/module/mtp/mtp_layer.py
@@ -7,6 +7,10 @@
 
 from xtuner.v1.data_proto import SequenceContext
 from xtuner.v1.module import RMSNorm
+from xtuner.v1.module.decoder_layer.moe_decoder_layer import (
+    MoEDecoderLayerMicroBatchOutput,
+    MoEDecoderLayerOutput,
+)
 from xtuner.v1.module.linear import build_linear
 
 
@@ -81,40 +85,34 @@ def __init__(
 
     def forward(
         self,
-        *hidden_states: torch.Tensor,
+        hidden_states: torch.Tensor | list[torch.Tensor],
+        *,
         future_embeddings: torch.Tensor | list[torch.Tensor],
         position_embeddings: tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]],
         seq_ctx: SequenceContext | list[SequenceContext],
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor] | tuple[torch.Tensor, ...]:
+    ) -> MoEDecoderLayerOutput | MoEDecoderLayerMicroBatchOutput:
         """Forward pass through the MTP layer.
 
-        Mirrors :meth:`MoEDecoderLayer.forward`: when a single ``hidden_states`` tensor is
-        provided, the layer runs the regular single-microbatch path and returns a 4-tuple
-        ``(hidden, router_logits, router_weights, router_topk_ids)``. When ``N`` hidden states are provided
-        (intra-layer micro-batching / domino EP), ``future_embeddings``, ``position_embeddings``
-        and ``seq_ctx`` must be lists of length ``N``; the per-microbatch preprocessing
-        (enorm/hnorm/eh_proj) is run independently and a single underlying decoder forward
-        is issued so the inner MoE EP communication can be overlapped across micro-batches.
+        Mirrors :meth:`MoEDecoderLayer.forward`: passing lists runs ``N`` micro-batches together
+        (intra-layer micro-batching / domino EP). The per-microbatch preprocessing
+        (enorm/hnorm/eh_proj) is run independently and a single underlying decoder forward is
+        issued, so the inner MoE EP communication can be overlapped across micro-batches.
 
         Args:
-            hidden_states (torch.Tensor): One or more hidden state tensors. A single tensor
-                triggers the single-microbatch path; multiple tensors trigger the
-                multi-microbatch path.
-            future_embeddings (torch.Tensor | list[torch.Tensor]): Embeddings of the future
-                tokens, aligned per-microbatch with ``hidden_states``.
-            position_embeddings (tuple | list[tuple]): Rotary position embeddings (cos, sin),
-                aligned per-microbatch with ``hidden_states``.
-            seq_ctx (SequenceContext | list[SequenceContext]): Sequence context per micro-batch.
+            hidden_states (torch.Tensor | list[torch.Tensor]): Hidden states, one tensor per
+                micro-batch.
+            future_embeddings (torch.Tensor | list[torch.Tensor]): Embeddings of the future tokens,
+                aligned with ``hidden_states``.
+            position_embeddings (tuple[torch.Tensor, torch.Tensor] | list[tuple[torch.Tensor, torch.Tensor]]):
+                Rotary position embeddings ``(cos, sin)``, aligned with ``hidden_states``.
+            seq_ctx (SequenceContext | list[SequenceContext]): Sequence context, aligned with
+                ``hidden_states``.
 
         Returns:
-            tuple: For single-microbatch input, a 4-tuple
-                ``(hidden_states, router_logits, router_weights, router_topk_ids)``.
-                For ``N`` micro-batches, a flat tuple of length ``4 * N`` matching the
-                convention used by :meth:`MoEDecoderLayer._micro_batch_forward`:
-                ``(hidden_0, ..., hidden_{N-1}, router_logits_0, ...,
-                router_weights_{N-1}, router_topk_ids_0, ..., router_topk_ids_{N-1})``.
+            MoEDecoderLayerOutput | MoEDecoderLayerMicroBatchOutput: The wrapped decoder layer's
+            outputs with the MTP final layernorm applied to the hidden states.
         """
-        if len(hidden_states) == 1:
+        if not isinstance(hidden_states, list):
             assert isinstance(future_embeddings, torch.Tensor), (
                 "future_embeddings should be a Tensor in single-microbatch mode"
             )
@@ -125,7 +123,7 @@ def forward(
                 "position_embeddings should be a (cos, sin) tuple in single-microbatch mode"
             )
             return self._forward(
-                hidden_states=hidden_states[0],
+                hidden_states=hidden_states,
                 future_embeddings=future_embeddings,
                 position_embeddings=position_embeddings,
                 seq_ctx=seq_ctx,
@@ -141,7 +139,7 @@ def forward(
             "position_embeddings should be a list aligned with hidden_states in multi-microbatch mode"
         )
         return self._micro_batch_forward(
-            hidden_states_list=list(hidden_states),
+            hidden_states_list=hidden_states,
             future_embeddings_list=future_embeddings,
             position_embeddings_list=position_embeddings,
             seq_ctx_list=seq_ctx,
@@ -153,17 +151,20 @@ def _forward(
         future_embeddings: torch.Tensor,
         position_embeddings: tuple[torch.Tensor, torch.Tensor],
         seq_ctx: SequenceContext,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    ) -> MoEDecoderLayerOutput:
         projected = self._preprocess(hidden_states=hidden_states, future_embeddings=future_embeddings)
 
-        hidden_states, router_results, router_weights, router_topk_ids = self.decoder_layer(
+        layer_results: MoEDecoderLayerOutput = self.decoder_layer(
             projected,
             position_embeddings=position_embeddings,
             seq_ctx=seq_ctx,
         )
-
-        hidden_states = self.final_layernorm(hidden_states)
-        return hidden_states, router_results, router_weights, router_topk_ids
+        return {
+            "hidden_states": self.final_layernorm(layer_results["hidden_states"]),
+            "router_logits": layer_results["router_logits"],
+            "router_weights": layer_results["router_weights"],
+            "router_topk_ids": layer_results["router_topk_ids"],
+        }
 
     def _micro_batch_forward(
         self,
@@ -172,7 +173,7 @@ def _micro_batch_forward(
         future_embeddings_list: list[torch.Tensor],
         position_embeddings_list: list[tuple[torch.Tensor, torch.Tensor]],
         seq_ctx_list: list[SequenceContext],
-    ) -> tuple[torch.Tensor, ...]:
+    ) -> MoEDecoderLayerMicroBatchOutput:
         n = len(hidden_states_list)
         assert len(future_embeddings_list) == n and len(position_embeddings_list) == n and len(seq_ctx_list) == n, (
             "All per-microbatch inputs must share the same length"
@@ -185,22 +186,17 @@ def _micro_batch_forward(
             for h, e in zip(hidden_states_list, future_embeddings_list)
         ]
 
-        layer_results = self.decoder_layer(
-            *projected_list,
+        layer_results: MoEDecoderLayerMicroBatchOutput = self.decoder_layer(
+            projected_list,
             position_embeddings=position_embeddings_list,
             seq_ctx=seq_ctx_list,
         )
-        assert isinstance(layer_results, tuple) and len(layer_results) == 4 * n, (
-            "Multi-microbatch MTP requires the wrapped decoder layer to return a flat "
-            f"(hidden..., router_logits..., router_weights..., router_topk_ids...) tuple of length {4 * n}; "
-            f"got length {len(layer_results) if isinstance(layer_results, tuple) else type(layer_results)}"
-        )
-
-        hidden_out = [self.final_layernorm(h) for h in layer_results[:n]]
-        router_logits = list(layer_results[n : 2 * n])
-        router_weights = list(layer_results[2 * n : 3 * n])
-        router_topk_ids = list(layer_results[3 * n :])
-        return tuple(hidden_out + router_logits + router_weights + router_topk_ids)
+        return {
+            "hidden_states": [self.final_layernorm(hidden) for hidden in layer_results["hidden_states"]],
+            "router_logits": layer_results["router_logits"],
+            "router_weights": layer_results["router_weights"],
+            "router_topk_ids": layer_results["router_topk_ids"],
+        }
 
     def _preprocess(
         self,
diff --git a/xtuner/v1/profiler/prober.py b/xtuner/v1/profiler/prober.py
index e3555c6642..c393e78b13 100644
--- a/xtuner/v1/profiler/prober.py
+++ b/xtuner/v1/profiler/prober.py
@@ -312,9 +312,11 @@ def wrapped_forward(self, *args, **kwargs):
                 hidden_states = kwargs["hidden_states"]
             ProberList.before_layer(name, hidden_states)
             outputs = forward(*args, **kwargs)
-            if isinstance(outputs, tuple):  # for MoEDecoderLayer
+            if isinstance(outputs, dict):
+                hidden_states = outputs["hidden_states"]
+            elif isinstance(outputs, tuple):  # for legacy decoder layers
                 hidden_states = outputs[0]
-            else:  # for DenseDecoderLayer
+            else:
                 hidden_states = outputs
             ProberList.after_layer(name, hidden_states)
             return outputs