From 3257307a0b0c516d107f02f065089881e8483179 Mon Sep 17 00:00:00 2001
From: clc <clc0301@cmbchina.com>
Date: Tue, 13 Oct 2026 16:16:21 +0800
Subject: [PATCH 1/8] sp: isolate sequence parallel context per model

---
 .../model/transformers/moe/router_replay.py   |  11 +-
 .../strategy/sequence_parallel/__init__.py    | 649 ++++++++++++------
 .../sequence_parallel/linear_attention_sp.py  |   6 +-
 .../strategy/sequence_parallel/utils.py       |  30 +-
 .../model/transformers/transformers.py        |   5 +-
 tests/transformers/test_dpo_sp_cpu.py         | 107 +++
 .../test_sequence_parallel_and_cp.py          | 162 ++++-
 .../transformers/test_sp_context_isolation.py | 141 ++++
 8 files changed, 862 insertions(+), 249 deletions(-)
 create mode 100644 tests/transformers/test_dpo_sp_cpu.py
 create mode 100644 tests/transformers/test_sp_context_isolation.py

diff --git a/src/twinkle/model/transformers/moe/router_replay.py b/src/twinkle/model/transformers/moe/router_replay.py
index 3cfcce124..2860ea0a4 100644
--- a/src/twinkle/model/transformers/moe/router_replay.py
+++ b/src/twinkle/model/transformers/moe/router_replay.py
@@ -77,6 +77,7 @@ def get_replay_state(block_name: str) -> _RouterReplayState | None:
 def set_router_replay_data(
     routed_experts: torch.Tensor,
     model: nn.Module,
+    sp_context=None,
 ) -> None:
     """Slice *routed_experts* into per-block
     ``target_indices`` and inject them into the registered MoE blocks of *model*.
@@ -96,7 +97,10 @@ def set_router_replay_data(
                          f'got {tuple(routed_experts.shape)}')
 
     # SP: slice full-sequence routed_experts to local SP rank tokens.
-    from ..strategy.sequence_parallel import sequence_parallel as sp
+    # Setup runs before the model's forward hook pushes its context.
+    # An unregistered model is non-SP, even if another SP model was prepared.
+    from ..strategy.sequence_parallel import _MODULE_TO_CTX
+    sp = sp_context if sp_context is not None else _MODULE_TO_CTX.get(model)
     sp_world_size = getattr(sp, 'sp_world_size', None) or 1
     if sp_world_size > 1:
         _, _, _, _, _, _, extra_values = sp.pad_and_split_inputs(
@@ -128,7 +132,7 @@ def set_router_replay_data(
             state.target_indices = target
 
 
-def get_router_replay_data(model: nn.Module, batch_size=1) -> torch.Tensor | None:
+def get_router_replay_data(model: nn.Module, batch_size=1, sp_context=None) -> torch.Tensor | None:
     """Collect ``recorded_indices`` from all registered MoE blocks in *model*.
 
     . note::
@@ -158,7 +162,8 @@ def get_router_replay_data(model: nn.Module, batch_size=1) -> torch.Tensor | Non
     routed_experts = routed_experts.reshape(batch_size, -1, num_layers, topk)
 
     # SP: all-gather local routing data across SP ranks along seq_len dim.
-    from ..strategy.sequence_parallel import sequence_parallel as sp
+    from ..strategy.sequence_parallel import _MODULE_TO_CTX
+    sp = sp_context if sp_context is not None else _MODULE_TO_CTX.get(model)
     sp_world_size = getattr(sp, 'sp_world_size', None) or 1
     if sp_world_size > 1:
         routed_experts = sp.gather(routed_experts, dim=1, position_ids=sp.real_position_ids)
diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
index 0b79d0b50..ae6c31cee 100644
--- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
@@ -1,13 +1,16 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import math
+import threading
 import torch
 import torch.distributed as dist
 from copy import copy
 from dataclasses import asdict, dataclass, is_dataclass
-from functools import partial
 from transformers import PreTrainedTokenizer
 from types import SimpleNamespace
 from typing import Any, Dict, List, Optional, Tuple, Union
+from weakref import WeakKeyDictionary
+
+from functools import partial
 
 from twinkle.patch import apply_patch
 from twinkle.utils import DeviceMesh
@@ -20,6 +23,53 @@
                     is_moe_config, post_all2all)
 
 
+# ---------------------------------------------------------------------------
+# Per-model SP context routing
+#
+# Each ``SequenceParallel`` instance is the SP context of exactly one model.
+# The global transformers patch points (``ALL_ATTENTION_FUNCTIONS``,
+# ``masking_utils``, ``modeling_flash_attention_utils``) are installed once per
+# process and resolve the owning context at call time:
+#   * attention functions receive ``module`` -> ``_MODULE_TO_CTX`` lookup;
+#   * mask helpers / ``_flash_attention_forward`` receive no module argument ->
+#     thread-local ctx stack pushed/popped by each model's forward hooks.
+# A missing context means the calling model is not SP-enabled: the wrapper
+# falls through to the original transformers implementation untouched, so
+# SP-enabled and non-SP models can be interleaved freely.
+# ---------------------------------------------------------------------------
+_MODULE_TO_CTX: WeakKeyDictionary = WeakKeyDictionary()
+_CTX_LOCAL = threading.local()
+
+
+def _ctx_stack() -> list:
+    stack = getattr(_CTX_LOCAL, 'stack', None)
+    if stack is None:
+        stack = _CTX_LOCAL.stack = []
+    return stack
+
+
+def _push_ctx(ctx) -> None:
+    _ctx_stack().append(ctx)
+
+
+def _pop_ctx(ctx) -> None:
+    stack = _ctx_stack()
+    for i in range(len(stack) - 1, -1, -1):
+        if stack[i] is ctx:
+            stack.pop(i)
+            return
+
+
+def _current_ctx():
+    stack = _ctx_stack()
+    return stack[-1] if stack else None
+
+
+def _register_module_tree(model: torch.nn.Module, ctx) -> None:
+    for module in model.modules():
+        _MODULE_TO_CTX[module] = ctx
+
+
 def is_qwen3_vl(model):
     mt = getattr(getattr(model, 'config', None), 'model_type', '')
     return 'qwen3_vl' in mt
@@ -78,6 +128,7 @@ def __init__(self):
         self.num_heads = None
         self.causal_mask_func = None
         self.extra_kwargs = {}
+        self._dist_attns = {}
 
     @property
     def real_position_ids(self) -> torch.Tensor:
@@ -88,6 +139,8 @@ def real_position_ids(self) -> torch.Tensor:
     def _extract_real_position_ids(position_ids: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
         if position_ids is None or not torch.is_tensor(position_ids):
             return position_ids
+        # Cached real positions are already 2D, so only reduce the original
+        # three-axis mRoPE representation.
         if position_ids.dim() == 3:
             return position_ids[0]
         return position_ids
@@ -100,7 +153,214 @@ def sp_rank(self) -> int:
     def rp_rank(self) -> int:
         return self._rp_rank
 
-    def _prepare_flash_attn(self, base_model: torch.nn.Module):
+    def _get_dist_attn(self, key: str, local_attention=None) -> DistributedAttention:
+        """Per-context DistributedAttention cache (one per attention entry point)."""
+        dist_attn = self._dist_attns.get(key)
+        if dist_attn is None:
+            dist_attn = DistributedAttention(local_attention, self)
+            self._dist_attns[key] = dist_attn
+        return dist_attn
+
+    @staticmethod
+    def _sdpa_mask_impl(sp, origin_sdpa, origin_uses_cache_position, batch_size, q_length=None, kv_length=None, *args,
+                        **kwargs):
+        q_length = q_length if q_length is not None else kwargs.pop('cache_position', None)
+        device = q_length.device if torch.is_tensor(q_length) else kwargs.pop('device', None)
+        if device is None and sp is not None:
+            device = sp.real_position_ids.device
+
+        cache_position = None
+        if sp is not None and sp.world_size > 1:
+            padded_position_ids = sp.pad(
+                sp.real_position_ids[0],
+                padding_value=-1,
+                position_ids=sp.real_position_ids,
+                dim=0,
+            )
+            global_length = padded_position_ids.shape[0]
+            if origin_uses_cache_position:
+                cache_position = torch.arange(0, global_length, device=device)
+                kv_length = global_length
+            else:
+                # Newer Transformers passes q_length/q_offset instead of cache_position. In SP training,
+                # create_causal_mask may still see the local shard length, so restore the global query length
+                # only for the no-cache prefill path; cache/sliding paths keep their upstream offsets.
+                q_offset = kwargs.get('q_offset', 0)
+                kv_offset = kwargs.get('kv_offset', 0)
+                no_cache_offsets = ((not torch.is_tensor(q_offset) and q_offset == 0)
+                                    and (not torch.is_tensor(kv_offset) and kv_offset == 0))
+                if no_cache_offsets:
+                    q_length = global_length
+                    attention_mask = kwargs.get('attention_mask')
+                    if attention_mask is not None and torch.is_tensor(attention_mask):
+                        kv_length = attention_mask.shape[-1]
+                    else:
+                        kv_length = global_length
+
+        # padding_mask may be an AttentionMask wrapper instead of Tensor
+        _pm = kwargs.get('padding_mask')
+        if _pm is not None and not isinstance(_pm, torch.Tensor):
+            if hasattr(_pm, 'to_tensor'):
+                kwargs['padding_mask'] = _pm.to_tensor()
+            elif hasattr(_pm, 'mask'):
+                kwargs['padding_mask'] = _pm.mask
+
+        if origin_uses_cache_position:
+            if cache_position is None:
+                cache_position = q_length if torch.is_tensor(q_length) else torch.arange(q_length, device=device)
+            return origin_sdpa(batch_size, cache_position, kv_length, *args, **kwargs)
+
+        return origin_sdpa(batch_size, q_length, kv_length, *args, device=device, **kwargs)
+
+    @staticmethod
+    def _create_causal_mask_impl(sp,
+                                 config,
+                                 input_embeds,
+                                 attention_mask,
+                                 cache_position_or_past_key_values=None,
+                                 *args,
+                                 **kwargs):
+        from transformers import masking_utils
+
+        if sp is None or sp.world_size == 1:
+            return _call_create_causal_mask(
+                masking_utils.origin_create_causal_mask,
+                config,
+                input_embeds,
+                attention_mask,
+                cache_position_or_past_key_values,
+                *args,
+                **kwargs,
+            )
+        input_embeds = torch.ones(
+            (input_embeds.shape[0], input_embeds.shape[1] * sp.sp_world_size, input_embeds.shape[2]),
+            dtype=input_embeds.dtype,
+            device=input_embeds.device)
+        if has_signature_parameter(masking_utils.origin_create_causal_mask, 'cache_position'):
+            cache_position_or_past_key_values = torch.arange(
+                0,
+                input_embeds.shape[1],
+                device=input_embeds.device,
+            )
+        return _call_create_causal_mask(
+            masking_utils.origin_create_causal_mask,
+            config,
+            input_embeds,
+            attention_mask,
+            cache_position_or_past_key_values,
+            *args,
+            **kwargs,
+        )
+
+    def _local_flash_attn_impl(self, module, query_states, key_states, value_states, attention_mask, *args, **kwargs):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        dist_attn = self._get_dist_attn('flash_attention_2')
+        if dist_attn.local_attn is None:
+
+            # Pass the current layer at call time. Capturing it in this cached
+            # closure would keep a weak-registry key alive through its context.
+            def _attention(query, key, value, *args, _sp_module, **kwargs):
+                query = query.transpose(1, 2)
+                key = key.transpose(1, 2)
+                value = value.transpose(1, 2)
+                if self.rp_world_size > 1:
+                    from .zigzag_ring_attn import zigzag_ring_flash_attn_varlen_func
+
+                    position_ids = kwargs.get('position_ids')
+                    if position_ids is None:
+                        position_ids = self.real_position_ids
+                    position_ids = self._extract_real_position_ids(position_ids)
+                    cu_seqlens = get_cu_seqlens_from_position_ids(position_ids).to(torch.int32)
+                    max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+                    position_ids = self._split_packed(position_ids, cu_seqlens, dim=-1)
+                    mask = position_ids != -1
+                    query = query.transpose(1, 2)
+                    key = key.transpose(1, 2)
+                    value = value.transpose(1, 2)
+                    query, key, value = self._mask_qkv(query, key, value, mask)
+                    return zigzag_ring_flash_attn_varlen_func(
+                        query,
+                        key,
+                        value,
+                        cu_seqlens=cu_seqlens,
+                        max_seqlen=max_seqlen,
+                        dropout_p=kwargs.get('dropout', 0.0),
+                        softmax_scale=kwargs.get('scaling'),
+                        causal=_sp_module.is_causal,
+                        window_size=kwargs.get('sliding_window') or (-1, -1),
+                        group=self._rp_group,
+                    )
+                elif self.extra_kwargs.get('padding_free', False) or 'cu_seq_lens_q' in kwargs:
+                    position_ids = kwargs.get('position_ids')
+                    if position_ids is None:
+                        position_ids = self.real_position_ids
+                    if position_ids is None:
+                        raise ValueError('Packed/varlen flash_attention_2 requires position_ids to derive '
+                                         'cu_seq_lens_q.')
+                    position_ids = self._extract_real_position_ids(position_ids)
+                    position_ids = self.pad(position_ids, padding_value=-1, position_ids=position_ids)
+                    cu_seqlens = get_cu_seqlens_from_position_ids(position_ids).to(
+                        dtype=torch.int32,
+                        device=query.device,
+                    )
+                    max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
+                    total_tokens = int(cu_seqlens[-1].item())
+                    if query.shape[2] != total_tokens:
+                        raise ValueError('Packed/varlen flash_attention_2 expects query sequence length to match '
+                                         f'cu_seqlens total tokens, got query_seq_len={query.shape[2]} '
+                                         f'and cu_seqlens_total={total_tokens}.')
+                    kwargs['cu_seq_lens_q'] = cu_seqlens
+                    kwargs['cu_seq_lens_k'] = cu_seqlens
+                    kwargs['max_length_q'] = max_seqlen
+                    kwargs['max_length_k'] = max_seqlen
+                    if self.extra_kwargs.get('padding_free', False) and len(args) > 0:
+                        args = (None, *args[1:])
+                return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](_sp_module, query, key, value, *args,
+                                                                           **kwargs)[0]
+
+            dist_attn.local_attn = _attention
+
+        return dist_attn(
+            query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask,
+            *args, _sp_module=module, **kwargs), None
+
+    def _local_sdpa_attn_impl(self, module, query_states, key_states, value_states, attention_mask, *args, **kwargs):
+        from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+        # Policy: packed (PackingDataset/padding-free) batches require FlashAttention2 varlen/packed semantics.
+        # SDPA does not have a native packed/varlen interface; supporting packed batches would require building a
+        # large block-diagonal causal mask (slow / memory heavy).
+        if self.extra_kwargs.get('padding_free', False):
+            raise RuntimeError('SequenceParallel: detected padding_free/packed batch. '
+                               'SDPA backend is not supported for padding_free/packed batches; please use '
+                               'flash_attention_2.')
+        dist_attn = self._get_dist_attn('sdpa')
+        if dist_attn.local_attn is None:
+
+            # The cache is per context, not per layer; do not bind the first
+            # layer or retain it through the global weak-key registry.
+            def _attention(query, key, value, *args, _sp_module, **kwargs):
+                query = query.transpose(1, 2)
+                key = key.transpose(1, 2)
+                value = value.transpose(1, 2)
+                if self.rp_world_size > 1:
+                    raise NotImplementedError('SDPA does not support derived ring attention.')
+                return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](_sp_module, query, key, value, *args, **kwargs)[0]
+
+            dist_attn.local_attn = _attention
+        return dist_attn(
+            query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask,
+            *args, _sp_module=module, **kwargs), None
+
+    def _prepare_flash_attn(self):
+        """Install the global SP-aware attention/mask wrappers (once per process).
+
+        The installed wrappers carry no per-model state; they resolve the owning
+        SequenceParallel context at call time via ``_MODULE_TO_CTX`` (attention
+        functions) or the thread-local ctx stack (mask helpers), so multiple
+        SP-enabled models (and non-SP models) can coexist in one process.
+        """
         try:
             from transformers import masking_utils
 
@@ -108,54 +368,8 @@ def _prepare_flash_attn(self, base_model: torch.nn.Module):
             origin_uses_cache_position = has_signature_parameter(origin_sdpa, 'cache_position')
 
             def sdpa_mask(batch_size, q_length=None, kv_length=None, *args, **kwargs):
-                q_length = q_length if q_length is not None else kwargs.pop('cache_position', None)
-                device = q_length.device if torch.is_tensor(q_length) else kwargs.pop('device', None)
-                if device is None:
-                    device = self.real_position_ids.device
-
-                cache_position = None
-                if self.world_size > 1:
-                    padded_position_ids = self.pad(
-                        self.real_position_ids[0],
-                        padding_value=-1,
-                        position_ids=self.real_position_ids,
-                        dim=0,
-                    )
-                    global_length = padded_position_ids.shape[0]
-                    if origin_uses_cache_position:
-                        cache_position = torch.arange(0, global_length, device=device)
-                        kv_length = global_length
-                    else:
-                        # Newer Transformers passes q_length/q_offset instead of cache_position. In SP training,
-                        # create_causal_mask may still see the local shard length, so restore the global query length
-                        # only for the no-cache prefill path; cache/sliding paths keep their upstream offsets.
-                        q_offset = kwargs.get('q_offset', 0)
-                        kv_offset = kwargs.get('kv_offset', 0)
-                        no_cache_offsets = ((not torch.is_tensor(q_offset) and q_offset == 0)
-                                            and (not torch.is_tensor(kv_offset) and kv_offset == 0))
-                        if no_cache_offsets:
-                            q_length = global_length
-                            attention_mask = kwargs.get('attention_mask')
-                            if attention_mask is not None and torch.is_tensor(attention_mask):
-                                kv_length = attention_mask.shape[-1]
-                            else:
-                                kv_length = global_length
-
-                # padding_mask may be an AttentionMask wrapper instead of Tensor
-                _pm = kwargs.get('padding_mask')
-                if _pm is not None and not isinstance(_pm, torch.Tensor):
-                    if hasattr(_pm, 'to_tensor'):
-                        kwargs['padding_mask'] = _pm.to_tensor()
-                    elif hasattr(_pm, 'mask'):
-                        kwargs['padding_mask'] = _pm.mask
-
-                if origin_uses_cache_position:
-                    if cache_position is None:
-                        cache_position = q_length if torch.is_tensor(q_length) else torch.arange(
-                            q_length, device=device)
-                    return origin_sdpa(batch_size, cache_position, kv_length, *args, **kwargs)
-
-                return origin_sdpa(batch_size, q_length, kv_length, *args, device=device, **kwargs)
+                return SequenceParallel._sdpa_mask_impl(_current_ctx(), origin_sdpa, origin_uses_cache_position,
+                                                        batch_size, q_length, kv_length, *args, **kwargs)
 
             masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa_origin'] = origin_sdpa
             masking_utils.ALL_MASK_ATTENTION_FUNCTIONS._global_mapping['sdpa'] = sdpa_mask
@@ -166,46 +380,14 @@ def create_causal_mask(config,
                                    cache_position_or_past_key_values=None,
                                    *args,
                                    **kwargs):
-                if self.world_size == 1:
-                    return _call_create_causal_mask(
-                        masking_utils.origin_create_causal_mask,
-                        config,
-                        input_embeds,
-                        attention_mask,
-                        cache_position_or_past_key_values,
-                        *args,
-                        **kwargs,
-                    )
-                input_embeds = torch.ones(
-                    (input_embeds.shape[0], input_embeds.shape[1] * self.sp_world_size, input_embeds.shape[2]),
-                    dtype=input_embeds.dtype,
-                    device=input_embeds.device)
-                if has_signature_parameter(masking_utils.origin_create_causal_mask, 'cache_position'):
-                    cache_position_or_past_key_values = torch.arange(
-                        0,
-                        input_embeds.shape[1],
-                        device=input_embeds.device,
-                    )
-                return _call_create_causal_mask(
-                    masking_utils.origin_create_causal_mask,
-                    config,
-                    input_embeds,
-                    attention_mask,
-                    cache_position_or_past_key_values,
-                    *args,
-                    **kwargs,
-                )
+                return SequenceParallel._create_causal_mask_impl(_current_ctx(), config, input_embeds, attention_mask,
+                                                                 cache_position_or_past_key_values, *args, **kwargs)
 
             masking_utils.origin_create_causal_mask = masking_utils.create_causal_mask
             masking_utils.create_causal_mask = create_causal_mask
         except ImportError:
             pass
 
-        if hasattr(base_model, 'language_model'):
-            text_model = base_model.language_model
-        else:
-            text_model = base_model
-
         from transformers.modeling_flash_attention_utils import is_flash_attn_available
         if is_flash_attn_available():
             # TODO this works for multi-modal models like qwen2.5-vl
@@ -213,169 +395,104 @@ def create_causal_mask(config,
             # more work for maintaining.
             from transformers import modeling_flash_attention_utils
             from transformers.modeling_flash_attention_utils import _flash_attention_forward
-            _distributed_flash_attention = DistributedAttention(_flash_attention_forward, self)
 
             modeling_flash_attention_utils._flash_attention_forward_origin = _flash_attention_forward
 
             def flash_attention_forward(query_states: torch.Tensor, key_states: torch.Tensor,
                                         value_states: torch.Tensor, attention_mask: Optional[torch.Tensor], q_len,
                                         *args, **kwargs):
-                if self.world_size == 1:
+                sp = _current_ctx()
+                if sp is None or sp.world_size == 1:
                     return _flash_attention_forward(query_states, key_states, value_states, attention_mask, q_len,
                                                     *args, **kwargs)
-                return _distributed_flash_attention(query_states, key_states, value_states, attention_mask,
-                                                    q_len * self.sp_world_size, *args, **kwargs)
+                dist_attn = sp._get_dist_attn('flash_attention_utils', _flash_attention_forward)
+                return dist_attn(query_states, key_states, value_states, attention_mask, q_len * sp.sp_world_size,
+                                 *args, **kwargs)
 
             modeling_flash_attention_utils._flash_attention_forward = flash_attention_forward
 
         from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
 
         def local_flash_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args,
-                             dist_attn, **kwargs):
-            if self.world_size == 1 or module.__class__ not in [m.__class__ for m in text_model.modules()]:
+                             **kwargs):
+            sp = _MODULE_TO_CTX.get(module)
+            if sp is None or sp.world_size == 1:
                 return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](module, query_states, key_states,
                                                                            value_states, attention_mask, *args,
                                                                            **kwargs)
-            if dist_attn.local_attn is None:
-
-                def _attention(query, key, value, *args, **kwargs):
-                    query = query.transpose(1, 2)
-                    key = key.transpose(1, 2)
-                    value = value.transpose(1, 2)
-                    if self.rp_world_size > 1:
-                        from .zigzag_ring_attn import zigzag_ring_flash_attn_varlen_func
-
-                        position_ids = kwargs.get('position_ids')
-                        if position_ids is None:
-                            position_ids = self.real_position_ids
-                        position_ids = self._extract_real_position_ids(position_ids)
-                        cu_seqlens = get_cu_seqlens_from_position_ids(position_ids).to(torch.int32)
-                        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
-                        position_ids = self._split_packed(position_ids, cu_seqlens, dim=-1)
-                        mask = position_ids != -1
-                        query = query.transpose(1, 2)
-                        key = key.transpose(1, 2)
-                        value = value.transpose(1, 2)
-                        query, key, value = self._mask_qkv(query, key, value, mask)
-                        return zigzag_ring_flash_attn_varlen_func(
-                            query,
-                            key,
-                            value,
-                            cu_seqlens=cu_seqlens,
-                            max_seqlen=max_seqlen,
-                            dropout_p=kwargs.get('dropout', 0.0),
-                            softmax_scale=kwargs.get('scaling'),
-                            causal=module.is_causal,
-                            window_size=kwargs.get('sliding_window') or (-1, -1),
-                            group=self._rp_group,
-                        )
-                    elif self.extra_kwargs.get('padding_free', False) or 'cu_seq_lens_q' in kwargs:
-                        position_ids = kwargs.get('position_ids')
-                        if position_ids is None:
-                            position_ids = self.real_position_ids
-                        if position_ids is None:
-                            raise ValueError('Packed/varlen flash_attention_2 requires position_ids to derive '
-                                             'cu_seq_lens_q.')
-                        position_ids = self._extract_real_position_ids(position_ids)
-                        position_ids = self.pad(position_ids, padding_value=-1, position_ids=position_ids)
-                        cu_seqlens = get_cu_seqlens_from_position_ids(position_ids).to(
-                            dtype=torch.int32,
-                            device=query.device,
-                        )
-                        max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()
-                        total_tokens = int(cu_seqlens[-1].item())
-                        if query.shape[2] != total_tokens:
-                            raise ValueError('Packed/varlen flash_attention_2 expects query sequence length to match '
-                                             f'cu_seqlens total tokens, got query_seq_len={query.shape[2]} '
-                                             f'and cu_seqlens_total={total_tokens}.')
-                        kwargs['cu_seq_lens_q'] = cu_seqlens
-                        kwargs['cu_seq_lens_k'] = cu_seqlens
-                        kwargs['max_length_q'] = max_seqlen
-                        kwargs['max_length_k'] = max_seqlen
-                        if self.extra_kwargs.get('padding_free', False) and len(args) > 0:
-                            args = (None, *args[1:])
-                    return ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'](module, query, key, value, *args,
-                                                                               **kwargs)[0]
-
-                dist_attn.local_attn = _attention
-
-            return dist_attn(
-                query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask,
-                *args, **kwargs), None
+            return sp._local_flash_attn_impl(module, query_states, key_states, value_states, attention_mask, *args,
+                                             **kwargs)
 
         def local_sdpa_attn(module: torch.nn.Module, query_states, key_states, value_states, attention_mask, *args,
-                            dist_attn, **kwargs):
-            # Bypass SP logic when world_size == 1 (SP disabled) or module not in text_model
-            if self.world_size == 1 or module.__class__ not in [m.__class__ for m in text_model.modules()]:
+                            **kwargs):
+            sp = _MODULE_TO_CTX.get(module)
+            if sp is None or sp.world_size == 1:
                 return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](module, query_states, key_states, value_states,
                                                               attention_mask, *args, **kwargs)
-            # Policy: packed (PackingDataset/padding-free) batches require FlashAttention2 varlen/packed semantics.
-            # SDPA does not have a native packed/varlen interface; supporting packed batches would require building a
-            # large block-diagonal causal mask (slow / memory heavy).
-            if self.extra_kwargs.get('padding_free', False):
-                raise RuntimeError(
-                    'SequenceParallel: detected padding_free/packed batch. '
-                    'SDPA backend is not supported for padding_free/packed batches; please use flash_attention_2.')
-            if dist_attn.local_attn is None:
-
-                def _attention(query, key, value, *args, **kwargs):
-                    query = query.transpose(1, 2)
-                    key = key.transpose(1, 2)
-                    value = value.transpose(1, 2)
-                    if self.rp_world_size > 1:
-                        raise NotImplementedError('SDPA does not support derived ring attention.')
-                    return ALL_ATTENTION_FUNCTIONS['sdpa_origin'](module, query, key, value, *args, **kwargs)[0]
-
-                dist_attn.local_attn = _attention
-            return dist_attn(
-                query_states.transpose(1, 2), key_states.transpose(1, 2), value_states.transpose(1, 2), attention_mask,
-                *args, **kwargs), None
+            return sp._local_sdpa_attn_impl(module, query_states, key_states, value_states, attention_mask, *args,
+                                            **kwargs)
 
         ALL_ATTENTION_FUNCTIONS['flash_attention_2_origin'] = ALL_ATTENTION_FUNCTIONS['flash_attention_2']
         ALL_ATTENTION_FUNCTIONS['sdpa_origin'] = ALL_ATTENTION_FUNCTIONS['sdpa']
-        ALL_ATTENTION_FUNCTIONS['flash_attention_2'] = partial(
-            local_flash_attn, dist_attn=DistributedAttention(None, self))
-        ALL_ATTENTION_FUNCTIONS['sdpa'] = partial(local_sdpa_attn, dist_attn=DistributedAttention(None, self))
+        ALL_ATTENTION_FUNCTIONS['flash_attention_2'] = local_flash_attn
+        ALL_ATTENTION_FUNCTIONS['sdpa'] = local_sdpa_attn
 
     def _prepare_forward_hook(self, base_model: torch.nn.Module):
 
         def pre_forward_split_hook(_self, args, kwargs):
-            if self.world_size == 1:
+            # Make this model's SP context visible to the global mask/attention
+            # wrappers for the whole duration of the forward.
+            _push_ctx(self)
+            try:
+                if self.world_size == 1:
+                    return args, kwargs
+                # Pad to multiple of SP size and split inputs per SP rank before forward.
+                input_ids = kwargs.get('input_ids', None)
+                inputs_embeds = kwargs.get('inputs_embeds', None)
+                position_ids = kwargs['position_ids']
+                real_position_ids = self._extract_real_position_ids(position_ids)
+                attention_mask = kwargs.get('attention_mask', None)
+                cache_position = kwargs.get('cache_position', None)
+                if hasattr(_self, 'language_model'):
+                    embed_tokens = getattr(_self.language_model, 'embed_tokens', None)
+                else:
+                    embed_tokens = getattr(_self, 'embed_tokens', None)
+                input_ids, inputs_embeds, _, position_ids, attention_mask, _, _ = self.pad_and_split_inputs(
+                    input_ids,
+                    inputs_embeds,
+                    None,
+                    position_ids,
+                    attention_mask,
+                    None,
+                    embed_tokens=embed_tokens,
+                    real_position_ids=real_position_ids,
+                    cache_position=cache_position)
+                kwargs['input_ids'] = input_ids
+                kwargs['inputs_embeds'] = inputs_embeds
+                kwargs['position_ids'] = position_ids
+                kwargs['attention_mask'] = attention_mask
                 return args, kwargs
-            # Pad to multiple of SP size and split inputs per SP rank before forward.
-            input_ids = kwargs.get('input_ids', None)
-            inputs_embeds = kwargs.get('inputs_embeds', None)
-            position_ids = kwargs['position_ids']
-            real_position_ids = self._extract_real_position_ids(position_ids)
-            attention_mask = kwargs.get('attention_mask', None)
-            cache_position = kwargs.get('cache_position', None)
-            if hasattr(_self, 'language_model'):
-                embed_tokens = getattr(_self.language_model, 'embed_tokens', None)
-            else:
-                embed_tokens = getattr(_self, 'embed_tokens', None)
-            input_ids, inputs_embeds, _, position_ids, attention_mask, _, _ = self.pad_and_split_inputs(
-                input_ids,
-                inputs_embeds,
-                None,
-                position_ids,
-                attention_mask,
-                None,
-                embed_tokens=embed_tokens,
-                real_position_ids=real_position_ids,
-                cache_position=cache_position)
-            kwargs['input_ids'] = input_ids
-            kwargs['inputs_embeds'] = inputs_embeds
-            kwargs['position_ids'] = position_ids
-            kwargs['attention_mask'] = attention_mask
-            return args, kwargs
+            except BaseException:
+                # PyTorch does not guarantee forward hooks run when a pre-hook
+                # itself raises. Restore the previous model context immediately.
+                _pop_ctx(self)
+                raise
+
+        def post_forward_pop_hook(_self, args, output):
+            _pop_ctx(self)
+            return output
+
+        base_model.register_forward_pre_hook(pre_forward_split_hook, with_kwargs=True, prepend=True)
+        # PyTorch also calls this hook if a pre-hook or the forward raises.
+        base_model.register_forward_hook(post_forward_pop_hook, always_call=True)
 
-        base_model.register_forward_pre_hook(pre_forward_split_hook, with_kwargs=True)
 
     def _prepare_multimodal_deepstack(self, base_model: torch.nn.Module):
         if not is_qwen3_vl(base_model):
             return
         apply_patch(base_model, Qwen3VLDeepstackSPPatch, sequence_parallel=self)
 
+
     @staticmethod
     def _is_qwen35_model(model: torch.nn.Module) -> bool:
         config = getattr(model, 'config', None)
@@ -390,6 +507,7 @@ def _is_qwen35_model(model: torch.nn.Module) -> bool:
         model_module = getattr(model.__class__, '__module__', '') or ''
         return 'transformers.models.qwen3_5' in model_module or 'transformers.models.qwen3_5_moe' in model_module
 
+
     def _prepare_qwen35_linear_attention(self, model: torch.nn.Module):
         has_qwen35_linear_attention = self._is_qwen35_model(model)
         if not has_qwen35_linear_attention:
@@ -405,6 +523,7 @@ def _prepare_qwen35_linear_attention(self, model: torch.nn.Module):
                 '(derived ring attention).')
         apply_patch(None, Qwen3_5GatedDeltaNetUlyssesPatch, sequence_parallel=self)
 
+
     def _prepare_moe_aux_loss(self, base_model: torch.nn.Module):
 
         def moe_aux_loss_hook(module, args, kwargs, output):
@@ -427,7 +546,7 @@ def moe_aux_loss_hook(module, args, kwargs, output):
                 _bs_logits = [logit[_slice] for logit in router_logits]
                 compute_device = _bs_logits[0].device
                 _bs_logits = torch.stack([layer_gate.to(compute_device) for layer_gate in _bs_logits], dim=0)
-                _bs_logits, _ = GatherLoss.apply(_bs_logits, None, 1, self.real_position_ids)
+                _bs_logits, _ = GatherLoss.apply(self, _bs_logits, None, 1, self.real_position_ids)
                 _gathered_logits.append(_bs_logits)
             router_logits = torch.stack(_gathered_logits, dim=0)
             if self.real_position_ids is not None:
@@ -438,6 +557,7 @@ def moe_aux_loss_hook(module, args, kwargs, output):
 
         base_model.register_forward_hook(moe_aux_loss_hook, with_kwargs=True)
 
+
     def prepare(
         self,
         sp_size: int,
@@ -486,14 +606,12 @@ def prepare(
         self._init_device_mesh(device_mesh)
         if not SequenceParallel._global_inited:
             # these operations are global initializations and patches
-            self._prepare_flash_attn(llm_model)
+            self._prepare_flash_attn()
             SequenceParallel._global_inited = True
         self._prepare_qwen35_linear_attention(llm_model)
-
         self._prepare_forward_hook(llm_model)
         self._prepare_multimodal_deepstack(llm_model)
-
-        if is_moe_config(getattr(model, 'config', None)):
+        if is_moe_config(getattr(llm_model, 'config', None)):
             self._prepare_moe_aux_loss(llm_model)
 
         self.model_dtype = next(model.parameters()).dtype
@@ -503,6 +621,20 @@ def prepare(
             if attn_impl != 'flash_attention_2':
                 raise NotImplementedError('Derived ring attention only supports flash_attention_2 backend.')
 
+        # Route this model's text-side attention modules to this context. The
+        # registry is instance-based (not class-based), so a second model of the
+        # same architecture gets its own context instead of silently sharing
+        # this one's extra_kwargs, and modules of other models (e.g. vision
+        # towers) fall through to the origin attention.
+        text_model = llm_model.language_model if hasattr(llm_model, 'language_model') else llm_model
+        _register_module_tree(text_model, self)
+        if self is not sequence_parallel:
+            # Compat: keep the legacy module-level singleton reflecting the most
+            # recently prepared context (extra_kwargs is shared by reference, so
+            # legacy readers observe live per-forward state). New code should
+            # use explicit contexts instead.
+            sequence_parallel.__dict__.update(self.__dict__)
+
     def _mask_qkv(self, query, key, value, mask):
         mask = mask.unsqueeze(2).unsqueeze(3)
         query = query * mask
@@ -680,6 +812,19 @@ def split(self, input, dim: int, position_ids=None):
         tensor_list = torch.split(input, dim_size // self.sp_world_size, dim=dim)
         return tensor_list[self.sp_rank].contiguous()
 
+    def _pad_key(self, value, spec, position_ids=None):
+        if value is None:
+            return None
+        pad_value = self.tokenizer.pad_token_id if spec.pad_value == 'tokenizer' else spec.pad_value
+        return self.pad(value, padding_value=pad_value, dim=spec.sequence_dim(value.dim()), position_ids=position_ids)
+
+    def _split_key(self, value, spec, position_ids=None, legacy_roll=False):
+        if value is None or not spec.need_split:
+            return value
+        if spec.need_roll or legacy_roll:
+            value = torch.roll(value, shifts=-1, dims=spec.split_dim)
+        return self.split(value, dim=spec.split_dim, position_ids=position_ids)
+
     def pad_and_split_inputs(self,
                              input_ids,
                              input_embeds,
@@ -880,6 +1025,9 @@ def __init__(
         self._tokenizer_id = tokenizer_id
         self._tokenizer = None
         self._initialized = False
+        # Per-model SP context. The module-level ``sequence_parallel`` singleton
+        # is only kept as a compat mirror of the most recently prepared context.
+        self._ctx = SequenceParallel()
 
     def _get_tokenizer(self) -> Optional[PreTrainedTokenizer]:
         if self._tokenizer is not None:
@@ -906,7 +1054,7 @@ def initialize(self) -> bool:
         tokenizer = self._get_tokenizer()
         if tokenizer is None:
             raise RuntimeError('SequenceParallelStrategy requires a tokenizer to initialize.')
-        sequence_parallel.prepare(
+        self._ctx.prepare(
             self.ulysses_size,
             self._model_ref,
             tokenizer,
@@ -915,10 +1063,54 @@ def initialize(self) -> bool:
         self._initialized = True
         return True
 
+    # -- SP context delegates -------------------------------------------------
+    # Processors / router replay access SP state through the strategy. These
+    # delegates forward to this model's own context (they previously did not
+    # exist, silently disabling SP-aware paths such as last-token pooling).
+
+    @property
+    def world_size(self) -> int:
+        return self._ctx.world_size or 1
+
+    @property
+    def sp_world_size(self) -> int:
+        return self._ctx.sp_world_size or 1
+
+    @property
+    def rp_world_size(self) -> int:
+        return self._ctx.rp_world_size or 1
+
+    @property
+    def _sp_group(self):
+        return self._ctx._sp_group
+
+    @property
+    def _rp_group(self):
+        return self._ctx._rp_group
+
+    @property
+    def real_position_ids(self):
+        return self._ctx.real_position_ids
+
+    def pad(self, tensor, *args, **kwargs):
+        if self._ctx.world_size is None:
+            return tensor
+        return self._ctx.pad(tensor, *args, **kwargs)
+
+    def split(self, tensor, *args, **kwargs):
+        if self._ctx.world_size is None:
+            return tensor
+        return self._ctx.split(tensor, *args, **kwargs)
+
+    def gather(self, tensor, *args, **kwargs):
+        if self._ctx.world_size is None:
+            return tensor
+        return self._ctx.gather(tensor, *args, **kwargs)
+
     def preprocess_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
         if not self.enabled or self.ulysses_size <= 1:
             return inputs
-        return sequence_parallel.prepare_inputs(inputs)
+        return self._ctx.prepare_inputs(inputs)
 
     def postprocess_outputs(self, outputs: Any) -> Any:
         if (not self.enabled or self.ulysses_size <= 1 or not self.sp_config.get('gather_logits', True)):
@@ -931,24 +1123,23 @@ def postprocess_outputs(self, outputs: Any) -> Any:
         logits = outputs.get('logits', None)
         if logits is None or not torch.is_tensor(logits) or logits.dim() < 2:
             return outputs
-        gathered = sequence_parallel.gather(logits, dim=1, position_ids=sequence_parallel.real_position_ids)
+        gathered = self._ctx.gather(logits, dim=1, position_ids=self._ctx.real_position_ids)
         # Scheme A: SP pads to make seq_len divisible by sp_size. Trim back to the original
         # (unpadded) length using the cached real_position_ids.
-        real_pos = sequence_parallel.real_position_ids
+        real_pos = self._ctx.real_position_ids
         if real_pos is not None and torch.is_tensor(real_pos) and real_pos.dim() >= 2:
             gathered = gathered[:, :real_pos.shape[1]].contiguous()
         outputs['logits'] = gathered
         return outputs
 
-    @staticmethod
-    def _trim_gathered_sequence_padding(tensor: torch.Tensor, real_position_ids: torch.Tensor) -> torch.Tensor:
+    def _trim_gathered_sequence_padding(self, tensor: torch.Tensor, real_position_ids: torch.Tensor) -> torch.Tensor:
         if real_position_ids is None or not torch.is_tensor(real_position_ids) or real_position_ids.dim() < 2:
             return tensor
-        if sequence_parallel.rp_world_size > 1:
+        if self._ctx.rp_world_size > 1:
             cu_seqlens = get_cu_seqlens_from_position_ids(real_position_ids)
             pieces = []
             padded_offset = 0
-            divisor = sequence_parallel.world_size * 2
+            divisor = self._ctx.world_size * 2
             for i in range(len(cu_seqlens) - 1):
                 real_len = int((cu_seqlens[i + 1] - cu_seqlens[i]).item())
                 padded_len = math.ceil(real_len / divisor) * divisor
@@ -982,8 +1173,8 @@ def gather_loss_tensors(
             hidden = outputs['logits']
             inputs = copy(inputs)
             outputs = copy(outputs)
-            real_position_ids = sequence_parallel.real_position_ids
-            gathered_hidden, gathered_labels = GatherLoss.apply(hidden, labels, 1, real_position_ids)
+            real_position_ids = self._ctx.real_position_ids
+            gathered_hidden, gathered_labels = GatherLoss.apply(self._ctx, hidden, labels, 1, real_position_ids)
             gathered_hidden = self._trim_gathered_sequence_padding(gathered_hidden, real_position_ids)
             gathered_labels = self._trim_gathered_sequence_padding(gathered_labels, real_position_ids)
             outputs['logits'] = gathered_hidden
@@ -996,15 +1187,15 @@ def gather_loss_tensors(
                             f'sequence tensor, got type={type(logps)} shape={getattr(logps, "shape", None)}')
         inputs = copy(inputs)
         outputs = copy(outputs)
-        real_position_ids = sequence_parallel.real_position_ids
-        gathered_logps, gathered_labels = GatherLoss.apply(logps, labels, 1, real_position_ids)
+        real_position_ids = self._ctx.real_position_ids
+        gathered_logps, gathered_labels = GatherLoss.apply(self._ctx, logps, labels, 1, real_position_ids)
         gathered_logps = self._trim_gathered_sequence_padding(gathered_logps, real_position_ids)
         gathered_labels = self._trim_gathered_sequence_padding(gathered_labels, real_position_ids)
         outputs['logps'] = gathered_logps
         inputs['labels'] = gathered_labels
         entropies = outputs.get('entropies')
         if entropies is not None and torch.is_tensor(entropies) and entropies.dim() >= 2:
-            gathered_entropies, _ = GatherLoss.apply(entropies, labels, 1, real_position_ids)
+            gathered_entropies, _ = GatherLoss.apply(self._ctx, entropies, labels, 1, real_position_ids)
             outputs['entropies'] = self._trim_gathered_sequence_padding(gathered_entropies, real_position_ids)
         return inputs, outputs
 
diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/linear_attention_sp.py b/src/twinkle/model/transformers/strategy/sequence_parallel/linear_attention_sp.py
index ab8db95ed..88c1d776d 100644
--- a/src/twinkle/model/transformers/strategy/sequence_parallel/linear_attention_sp.py
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/linear_attention_sp.py
@@ -340,7 +340,11 @@ def sp_linear_forward(
                 _origin_forward=origin_forward,
                 **extra_kwargs,
             ):
-                sequence_parallel_context = extra_kwargs.pop('sequence_parallel_context', sequence_parallel)
+                # The class patch is shared by every model, including models
+                # without SP. Resolve the owner instead of capturing the first
+                # model's context in this process-wide closure.
+                from . import _MODULE_TO_CTX
+                sequence_parallel_context = extra_kwargs.pop('sequence_parallel_context', _MODULE_TO_CTX.get(mod))
                 if not _sp_is_enabled(sequence_parallel_context):
                     return _origin_forward(
                         mod,
diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/utils.py b/src/twinkle/model/transformers/strategy/sequence_parallel/utils.py
index 6bc8677dc..8ed17e09c 100644
--- a/src/twinkle/model/transformers/strategy/sequence_parallel/utils.py
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/utils.py
@@ -237,33 +237,39 @@ def head_to_seq_shard(tensor: torch.Tensor, sequence_parallel) -> torch.Tensor:
 
 
 class GatherLoss(torch.autograd.Function):
-    """Gather loss from sequence group."""
+    """Gather loss from sequence group.
+
+    The owning ``SequenceParallel`` context is passed explicitly as the first
+    argument: backward() runs after other models may have forwarded, so reading
+    the process-global singleton here could pick up the wrong context in
+    multi-model training (DPO/PPO/GKD).
+    """
 
     @staticmethod
-    def forward(ctx, loss, labels, gather_idx=None, position_ids=None):
-        from . import sequence_parallel
+    def forward(ctx, sp, loss, labels, gather_idx=None, position_ids=None):
+        ctx.sp = sp
         ctx.scatter_shape = loss.shape[gather_idx or 0]
         ctx.gather_idx = gather_idx or 0
         if position_ids is not None:
-            position_ids = sequence_parallel.pad(position_ids, padding_value=-1, position_ids=position_ids)
+            position_ids = sp.pad(position_ids, padding_value=-1, position_ids=position_ids)
         ctx.position_ids = position_ids
-        output = sequence_parallel.gather(loss, dim=ctx.gather_idx, position_ids=position_ids)
+        output = sp.gather(loss, dim=ctx.gather_idx, position_ids=position_ids)
         if labels is not None:
-            labels_output = sequence_parallel.gather(labels, dim=ctx.gather_idx, position_ids=position_ids)
+            labels_output = sp.gather(labels, dim=ctx.gather_idx, position_ids=position_ids)
         else:
             labels_output = None
         return output, labels_output
 
     @staticmethod
     def backward(ctx, *grad_output):
-        from . import sequence_parallel
-        _grad = grad_output[0] * sequence_parallel.world_size
-        if sequence_parallel.rp_world_size > 1:
-            _grad = sequence_parallel.split(_grad, dim=ctx.gather_idx, position_ids=ctx.position_ids).contiguous()
+        sp = ctx.sp
+        _grad = grad_output[0] * sp.world_size
+        if sp.rp_world_size > 1:
+            _grad = sp.split(_grad, dim=ctx.gather_idx, position_ids=ctx.position_ids).contiguous()
         else:
             _grad = _grad.split(
-                ctx.scatter_shape, dim=ctx.gather_idx)[dist.get_rank(group=sequence_parallel._sp_group)].contiguous()
-        return _grad, None, None, None
+                ctx.scatter_shape, dim=ctx.gather_idx)[dist.get_rank(group=sp._sp_group)].contiguous()
+        return None, _grad, None, None, None
 
 
 def _generate_layout_params(scatter_idx, seq_world_size, input):
diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py
index 9382b5f9c..8f166a571 100644
--- a/src/twinkle/model/transformers/transformers.py
+++ b/src/twinkle/model/transformers/transformers.py
@@ -479,15 +479,16 @@ def _router_replay_setup(self, router_replay_action, routed_experts=None, batch_
         from .moe.router_replay import (RouterReplayAction, clear_global_indices, clear_global_router_replay_action,
                                         get_router_replay_data, set_global_router_replay_action, set_router_replay_data)
         unwrapped = self.strategy.unwrap_model(self.model)
+        sp_context = getattr(self.sp_strategy, '_ctx', None)
         set_global_router_replay_action(router_replay_action)
         if router_replay_action == RouterReplayAction.REPLAY_FORWARD:
             assert routed_experts is not None, 'routed_experts must be not None'
-            set_router_replay_data(routed_experts, unwrapped)
+            set_router_replay_data(routed_experts, unwrapped, sp_context=sp_context)
 
         def cleanup():
             recorded = None
             if router_replay_action == RouterReplayAction.RECORD:
-                recorded = get_router_replay_data(unwrapped, batch_size)
+                recorded = get_router_replay_data(unwrapped, batch_size, sp_context=sp_context)
             if not manual_cleanup:
                 clear_global_router_replay_action()
                 clear_global_indices()
diff --git a/tests/transformers/test_dpo_sp_cpu.py b/tests/transformers/test_dpo_sp_cpu.py
new file mode 100644
index 000000000..b8331b25a
--- /dev/null
+++ b/tests/transformers/test_dpo_sp_cpu.py
@@ -0,0 +1,107 @@
+"""Offline DPO loss/gradient equivalence over two CPU/Gloo SP ranks."""
+import copy
+import os
+import pytest
+import socket
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from datetime import timedelta
+from types import SimpleNamespace
+from transformers import LlamaConfig, LlamaForCausalLM
+
+from twinkle import DeviceMesh
+from twinkle.loss import DPOLoss
+from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelStrategy, _current_ctx
+from twinkle.utils import selective_log_softmax
+
+
+def _worker(rank, port):
+    os.environ['RANK'] = str(rank)
+    os.environ['WORLD_SIZE'] = '2'
+    torch.set_num_threads(1)
+    torch.manual_seed(42)
+    # torch_npu sets checkpoint's default device to NPU even for CPU-only
+    # tensors. This regression deliberately checkpoints on CPU.
+    from torch.utils.checkpoint import DefaultDeviceType
+    DefaultDeviceType.set_device_type('cpu')
+    dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank,
+                            world_size=2, timeout=timedelta(seconds=90))
+    try:
+        config = LlamaConfig(vocab_size=64, hidden_size=32, intermediate_size=64,
+                             num_hidden_layers=2, num_attention_heads=4, num_key_value_heads=4,
+                             attention_dropout=0.0, pad_token_id=0)
+        config._attn_implementation = 'sdpa'
+        baseline = LlamaForCausalLM(config).eval()
+        reference = copy.deepcopy(baseline).eval()
+        with torch.no_grad():
+            for param in baseline.parameters():
+                param.add_(torch.randn_like(param) * 1e-3)
+        policy = copy.deepcopy(baseline)
+        ref_sp = copy.deepcopy(reference)
+        # Four rows = two interleaved chosen/rejected pairs. Length 9 exercises SP padding.
+        ids = torch.randint(1, 64, (4, 9))
+        positions = torch.arange(9).expand(4, -1)
+        labels = torch.cat((ids[:, 1:], torch.full((4, 1), -100)), dim=1)
+        labels[:, :2] = -100
+        loss_fn = DPOLoss(beta=0.1, loss_type='sigmoid')
+
+        def score(model, strategy=None):
+            local_labels = labels
+            if strategy is not None:
+                local_labels = strategy.preprocess_inputs(
+                    {'input_ids': ids, 'position_ids': positions, 'labels': labels})['labels']
+            logits = model(input_ids=ids, position_ids=positions, use_cache=False).logits
+            outputs = {'logps': selective_log_softmax(logits, local_labels.masked_fill(local_labels == -100, 0))}
+            inputs = {'labels': local_labels}
+            if strategy is not None:
+                inputs, outputs = strategy.gather_loss_tensors(inputs, outputs)
+            return inputs, outputs
+
+        with torch.no_grad():
+            _, base_ref = score(reference)
+        base_inputs, base_outputs = score(baseline)
+        base_loss = loss_fn(base_inputs, base_outputs, ref_outputs=base_ref)['loss']
+        base_loss.backward()
+        base_grads = {name: p.grad.clone() for name, p in baseline.named_parameters()}
+
+        mesh = DeviceMesh.from_sizes(dp_size=2, ulysses_size=2, device_type='cpu')
+        strategies = []
+        for model in (policy, ref_sp):
+            strategy = SequenceParallelStrategy(mesh, model=model)
+            strategy._get_tokenizer = lambda: SimpleNamespace(pad_token_id=0)
+            assert strategy.initialize()
+            strategies.append(strategy)
+        assert strategies[0]._ctx is not strategies[1]._ctx
+        policy.gradient_checkpointing_enable()
+        policy.train()
+        with torch.no_grad():
+            _, parallel_ref = score(ref_sp, strategies[1])
+        inputs, outputs = score(policy, strategies[0])
+        parallel_loss = loss_fn(inputs, outputs, ref_outputs=parallel_ref)['loss']
+        torch.testing.assert_close(parallel_loss, base_loss, atol=1e-6, rtol=1e-5)
+        parallel_loss.backward()
+        # Mirror the data-parallel gradient averaging performed by DDP.
+        for name, param in policy.named_parameters():
+            assert param.grad is not None, name
+            dist.all_reduce(param.grad)
+            param.grad.div_(2)
+            torch.testing.assert_close(param.grad, base_grads[name], atol=1e-6, rtol=1e-4, msg=name)
+        assert _current_ctx() is None
+        # Avoid torch_npu's process-wide optimizer capability probe on this
+        # explicitly CPU-only test.
+        optimizer = torch.optim.AdamW(policy.parameters(), lr=1e-4, foreach=False, fused=False)
+        before = policy.lm_head.weight.detach().clone()
+        optimizer.step()
+        assert not torch.equal(before, policy.lm_head.weight)
+        assert all(param.grad is None for param in ref_sp.parameters())
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.slow
+def test_dpo_sp_loss_all_parameter_gradients_and_step():
+    with socket.socket() as sock:
+        sock.bind(('127.0.0.1', 0))
+        port = sock.getsockname()[1]
+    mp.spawn(_worker, args=(port,), nprocs=2, join=True)
diff --git a/tests/transformers/test_sequence_parallel_and_cp.py b/tests/transformers/test_sequence_parallel_and_cp.py
index a2ce49e61..23effbf3a 100644
--- a/tests/transformers/test_sequence_parallel_and_cp.py
+++ b/tests/transformers/test_sequence_parallel_and_cp.py
@@ -16,8 +16,7 @@
 from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelStrategy, sequence_parallel
 from twinkle.utils import DeviceMesh, Platform, ensure_hccl_socket_env, selective_log_softmax, torch_util
 
-pytestmark = pytest.mark.skip(
-    reason='Heavy SP/CP test (mp.spawn, world_size up to 4); not viable on dual-V100 CI, run manually.')
+pytestmark = [pytest.mark.slow, pytest.mark.accel(2)]
 
 LOGITS_RTOL = 1e-2
 LOGITS_ATOL = 5e-3
@@ -50,6 +49,7 @@ def _make_padded_labels(input_ids: torch.Tensor, valid_lengths: list[int]) -> to
 def _make_case(case_name: str) -> dict:
     cases = {
         'sp_only': {
+            'num_hidden_layers': 2,
             'expected_mode': 'sp_only',
             'world_size': 2,
             'ulysses_size': 2,
@@ -58,6 +58,7 @@ def _make_case(case_name: str) -> dict:
             'seq_len': 769,
         },
         'sp_only_multi_sample': {
+            'num_hidden_layers': 2,
             'expected_mode': 'sp_only',
             'world_size': 2,
             'ulysses_size': 2,
@@ -67,6 +68,7 @@ def _make_case(case_name: str) -> dict:
             'batch_size': 3,
         },
         'sp_only_multi_sample_masked': {
+            'num_hidden_layers': 2,
             'expected_mode': 'sp_only',
             'world_size': 2,
             'ulysses_size': 2,
@@ -720,6 +722,162 @@ def test_cp_sp_precision_alignment(self):
             join=True)
 
 
+def _run_multi_model_worker(rank: int, world_size: int, port: int, backend_config: dict):
+    """Two SP-enabled models (policy/ref style) with interleaved forwards.
+
+    Same architecture, different weights and different sequence lengths. Each
+    model must use its own SP context: a shared/singleton context shows up here
+    as cross-model position_ids pollution (logits mismatch or shape errors).
+    """
+    device = _init_dist(rank, world_size, port, backend_config)
+    try:
+        device_type = backend_config['device_type']
+        case = _make_case('sp_only')
+        vocab_size = 256
+
+        _seed_backend(1234, device_type)
+        policy_base = _build_tiny_llama(case, device, device_type)
+        _seed_backend(5678, device_type)
+        ref_base = _build_tiny_llama(case, device, device_type)
+
+        policy_sp = _build_tiny_llama(case, device, device_type)
+        policy_sp.load_state_dict(policy_base.state_dict())
+        ref_sp = _build_tiny_llama(case, device, device_type)
+        ref_sp.load_state_dict(ref_base.state_dict())
+
+        seq_len_p = int(case['seq_len'])
+        seq_len_r = seq_len_p // 2
+        _seed_backend(42, device_type)
+        input_ids_p = torch.randint(low=0, high=vocab_size, size=(1, seq_len_p), device=device)
+        position_ids_p = torch.arange(seq_len_p, device=device).unsqueeze(0)
+        input_ids_r = torch.randint(low=0, high=vocab_size, size=(1, seq_len_r), device=device)
+        position_ids_r = torch.arange(seq_len_r, device=device).unsqueeze(0)
+        labels_p = _make_labels(input_ids_p)
+        labels_r = _make_labels(input_ids_r)
+
+        # Non-SP baselines (computed before SP strategies exist).
+        policy_base.zero_grad(set_to_none=True)
+        base_outputs_p = policy_base(
+            input_ids=input_ids_p, position_ids=position_ids_p, attention_mask=None, use_cache=False)
+        base_logits_p = base_outputs_p.logits.detach().float()
+        base_loss_sum_p, base_num_tokens_p = _compute_training_path_loss(base_outputs_p.logits, labels_p)
+        base_loss_sum_p.backward()
+        base_grads_p = _normalize_grad_dict(_collect_attention_param_grads(policy_base), base_num_tokens_p)
+
+        ref_base.zero_grad(set_to_none=True)
+        base_outputs_r = ref_base(
+            input_ids=input_ids_r, position_ids=position_ids_r, attention_mask=None, use_cache=False)
+        base_logits_r = base_outputs_r.logits.detach().float()
+        base_loss_sum_r, base_num_tokens_r = _compute_training_path_loss(base_outputs_r.logits, labels_r)
+        base_loss_sum_r.backward()
+        base_grads_r = _normalize_grad_dict(_collect_attention_param_grads(ref_base), base_num_tokens_r)
+
+        # Two independent SP strategies on the same device mesh.
+        device_mesh = DeviceMesh.from_sizes(
+            fsdp_size=world_size,
+            dp_size=1,
+            ulysses_size=2,
+            device_type=device_type,
+        )
+        policy_strategy = _make_strategy(policy_sp, device_mesh, 2)
+        ref_strategy = _make_strategy(ref_sp, device_mesh, 2)
+
+        if policy_strategy._ctx is ref_strategy._ctx:
+            raise AssertionError('multi_model: strategies must own distinct SP contexts.')
+
+        # Module registry must route each model's attention to its own context.
+        from twinkle.model.transformers.strategy.sequence_parallel import _MODULE_TO_CTX
+        policy_attn = next(m for m in policy_sp.modules() if m.__class__.__name__.endswith('Attention'))
+        ref_attn = next(m for m in ref_sp.modules() if m.__class__.__name__.endswith('Attention'))
+        if _MODULE_TO_CTX.get(policy_attn) is not policy_strategy._ctx:
+            raise AssertionError('multi_model: policy attention module not registered to policy ctx.')
+        if _MODULE_TO_CTX.get(ref_attn) is not ref_strategy._ctx:
+            raise AssertionError('multi_model: ref attention module not registered to ref ctx.')
+
+        local_labels_p = _prepare_label_inputs(policy_strategy, input_ids_p, position_ids_p)
+        local_labels_r = _prepare_label_inputs(ref_strategy, input_ids_r, position_ids_r)
+
+        # Interleaved forwards: policy -> ref -> policy.
+        policy_sp.zero_grad(set_to_none=True)
+        ref_sp.zero_grad(set_to_none=True)
+        sp_outputs_p1 = policy_sp(
+            input_ids=input_ids_p, position_ids=position_ids_p, attention_mask=None, use_cache=False)
+        sp_outputs_r = ref_sp(
+            input_ids=input_ids_r, position_ids=position_ids_r, attention_mask=None, use_cache=False)
+        sp_outputs_p2 = policy_sp(
+            input_ids=input_ids_p, position_ids=position_ids_p, attention_mask=None, use_cache=False)
+
+        # Context isolation: each ctx must still hold its own sequence metadata.
+        if policy_strategy._ctx.real_position_ids.shape[-1] != seq_len_p:
+            raise AssertionError('multi_model: policy ctx position_ids polluted by ref forward.')
+        if ref_strategy._ctx.real_position_ids.shape[-1] != seq_len_r:
+            raise AssertionError('multi_model: ref ctx position_ids polluted by policy forward.')
+
+        # Capture per-rank local logits (with grad) before postprocess_outputs
+        # replaces outputs['logits'] with the grad-less gathered tensor.
+        local_logits_p1 = sp_outputs_p1.logits
+        local_logits_r = sp_outputs_r.logits
+
+        sp_logits_p1 = policy_strategy.postprocess_outputs(sp_outputs_p1).logits.detach().float()
+        sp_logits_r = ref_strategy.postprocess_outputs(sp_outputs_r).logits.detach().float()
+        sp_logits_p2 = policy_strategy.postprocess_outputs(sp_outputs_p2).logits.detach().float()
+
+        _assert_logits_close('multi_model policy fwd1', rank, base_logits_p, sp_logits_p1, seq_len_p, None)
+        _assert_logits_close('multi_model ref fwd', rank, base_logits_r, sp_logits_r, seq_len_r, None)
+        _assert_logits_close('multi_model policy fwd2', rank, base_logits_p, sp_logits_p2, seq_len_p, None)
+
+        # Interleaved loss + backward (ref backward first, then policy):
+        # GatherLoss.backward must use the ctx captured at forward time.
+        sp_loss_sum_p, sp_num_tokens_p = _compute_training_path_loss(local_logits_p1, local_labels_p,
+                                                                     policy_strategy)
+        sp_loss_sum_r, sp_num_tokens_r = _compute_training_path_loss(local_logits_r, local_labels_r,
+                                                                     ref_strategy)
+        sp_loss_sum_r.backward()
+        sp_loss_sum_p.backward()
+        _average_model_grads_over_group(policy_sp, policy_strategy._ctx._data_rank_group)
+        _average_model_grads_over_group(ref_sp, ref_strategy._ctx._data_rank_group)
+        sp_grads_p = _normalize_grad_dict(_collect_attention_param_grads(policy_sp), sp_num_tokens_p)
+        sp_grads_r = _normalize_grad_dict(_collect_attention_param_grads(ref_sp), sp_num_tokens_r)
+
+        if sp_num_tokens_p != base_num_tokens_p or sp_num_tokens_r != base_num_tokens_r:
+            raise AssertionError(f'multi_model num_tokens mismatch on rank {rank}: '
+                                 f'policy sp={sp_num_tokens_p} base={base_num_tokens_p}, '
+                                 f'ref sp={sp_num_tokens_r} base={base_num_tokens_r}')
+        if abs(sp_loss_sum_p.item() / max(sp_num_tokens_p, 1) - base_loss_sum_p.item() / max(base_num_tokens_p,
+                                                                                             1)) > LOSS_ATOL:
+            raise AssertionError(f'multi_model policy loss mismatch on rank {rank}: '
+                                 f'sp={sp_loss_sum_p.item()} base={base_loss_sum_p.item()}')
+        if abs(sp_loss_sum_r.item() / max(sp_num_tokens_r, 1) - base_loss_sum_r.item() / max(base_num_tokens_r,
+                                                                                             1)) > LOSS_ATOL:
+            raise AssertionError(f'multi_model ref loss mismatch on rank {rank}: '
+                                 f'sp={sp_loss_sum_r.item()} base={base_loss_sum_r.item()}')
+        _assert_grad_dict_close('multi_model policy', rank, base_grads_p, sp_grads_p)
+        _assert_grad_dict_close('multi_model ref', rank, base_grads_r, sp_grads_r)
+        dist.barrier()
+    finally:
+        dist.destroy_process_group()
+
+
+class TestMultiModelSequenceParallel:
+    """Multi-model (policy/ref style) SP context isolation tests."""
+
+    def _get_backend_or_skip(self, world_size: int = 2) -> dict:
+        backend = _get_runtime_backend()
+        if backend is None:
+            pytest.skip('CUDA or NPU is required for multi-model SP tests.')
+        if backend['device_count'] < world_size:
+            pytest.skip(f'Requires at least {world_size} {backend["label"]} devices.')
+        return backend
+
+    def test_two_models_interleaved_forward(self):
+        world_size = 2
+        backend = self._get_backend_or_skip(world_size)
+        if not is_flash_attn_available():
+            pytest.skip('flash_attention_2 is required for multi-model SP tests.')
+        port = _find_free_port()
+        mp.spawn(_run_multi_model_worker, args=(world_size, port, backend), nprocs=world_size, join=True)
+
+
 class TestDerivedRingMemoryProfile:
 
     def _get_backend_or_skip(self, world_size: int = 4) -> dict:
diff --git a/tests/transformers/test_sp_context_isolation.py b/tests/transformers/test_sp_context_isolation.py
new file mode 100644
index 000000000..1b13abf1a
--- /dev/null
+++ b/tests/transformers/test_sp_context_isolation.py
@@ -0,0 +1,141 @@
+"""CPU regression tests for process-wide SP patches and exception cleanup."""
+import pytest
+import torch
+import gc
+import weakref
+from types import SimpleNamespace
+
+import twinkle.model.transformers.strategy.sequence_parallel as sp
+from twinkle.model.transformers.strategy.sequence_parallel import linear_attention_sp
+
+
+@pytest.mark.parametrize('backend', ['sdpa', 'flash_attention_2'])
+def test_attention_cache_routes_current_module_and_releases_owners(monkeypatch, backend):
+    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS
+
+    seen = []
+
+    def origin(module, query, key, value, *args, **kwargs):
+        assert '_sp_module' not in kwargs
+        seen.append(weakref.ref(module))
+        return query, None
+
+    monkeypatch.setitem(ALL_ATTENTION_FUNCTIONS, f'{backend}_origin', origin)
+    ctx = sp.SequenceParallel()
+    ctx.world_size = ctx.sp_world_size = ctx.rp_world_size = 1
+    first, second = torch.nn.Linear(2, 2), torch.nn.Linear(2, 2)
+    refs = [weakref.ref(first), weakref.ref(second)]
+    sp._MODULE_TO_CTX[first] = ctx
+    sp._MODULE_TO_CTX[second] = ctx
+    x = torch.zeros(1, 1, 2, 2)
+    method = ctx._local_sdpa_attn_impl if backend == 'sdpa' else ctx._local_flash_attn_impl
+    try:
+        method(first, x, x, x, None)
+        method(second, x, x, x, None)
+        assert [r() for r in seen] == [first, second]
+        del first, second, method, ctx
+        gc.collect()
+        assert all(r() is None for r in refs)
+    finally:
+        for ref in refs:
+            if ref() is not None:
+                sp._MODULE_TO_CTX.pop(ref(), None)
+
+
+@pytest.fixture(autouse=True)
+def empty_context_stack():
+    saved = list(sp._ctx_stack())
+    sp._ctx_stack().clear()
+    yield
+    sp._ctx_stack()[:] = saved
+
+
+@pytest.mark.parametrize('failure', ['forward', 'split', 'other_pre_hook'])
+def test_exception_restores_enclosing_context(failure):
+    outer = object()
+    sp._push_ctx(outer)
+    ctx = sp.SequenceParallel()
+    ctx.world_size = 2 if failure == 'split' else 1
+
+    class Model(torch.nn.Module):
+        def forward(self, **kwargs):
+            assert sp._current_ctx() is ctx
+            raise RuntimeError('forward failed')
+
+    model = Model()
+    if failure == 'other_pre_hook':
+        def fail(*args):
+            assert sp._current_ctx() is ctx
+            raise RuntimeError('pre-hook failed')
+        model.register_forward_pre_hook(fail)
+    ctx._prepare_forward_hook(model)
+    with pytest.raises((RuntimeError, KeyError)):
+        model()
+    assert sp._ctx_stack() == [outer]
+    sp._pop_ctx(outer)
+    assert sp._current_ctx() is None
+
+
+def test_input_split_pre_hook_failure_does_not_leak_context():
+    ctx = sp.SequenceParallel()
+    ctx.world_size = 2
+
+    model = torch.nn.Linear(2, 2)
+    ctx._prepare_forward_hook(model)
+    with pytest.raises(KeyError, match='position_ids'):
+        model(torch.zeros(1, 2))
+    assert sp._current_ctx() is None
+
+    # A later model in this thread must not inherit the failed SP model's ctx.
+    plain_model = torch.nn.Linear(2, 2)
+    assert sp._current_ctx() is None
+    plain_model(torch.zeros(1, 2))
+    assert sp._current_ctx() is None
+
+
+def test_gdn_patch_routes_each_owner_and_preserves_non_sp(monkeypatch):
+    class Gate(torch.nn.Module):
+        def forward(self, hidden_states, **kwargs):
+            return ('original', hidden_states)
+
+    monkeypatch.setattr(linear_attention_sp, '_iter_qwen35_gated_delta_net_classes', lambda: [Gate])
+    patch_type = linear_attention_sp.Qwen3_5GatedDeltaNetUlyssesPatch
+    monkeypatch.setattr(patch_type, '_run_forward', staticmethod(
+        lambda mod, hidden_states, **kwargs: kwargs['sequence_parallel_context']))
+    policy_ctx = SimpleNamespace(world_size=2, sp_world_size=2, rp_world_size=1)
+    ref_ctx = SimpleNamespace(world_size=2, sp_world_size=2, rp_world_size=1)
+    policy, ref, plain = Gate(), Gate(), Gate()
+    monkeypatch.setitem(sp._MODULE_TO_CTX, policy, policy_ctx)
+    monkeypatch.setitem(sp._MODULE_TO_CTX, ref, ref_ctx)
+    patch_type()(None, sequence_parallel=policy_ctx)
+    patch_type()(None, sequence_parallel=ref_ctx)
+    x = torch.zeros(1)
+    assert policy(x) is policy_ctx
+    assert ref(x) is ref_ctx
+    assert policy(x) is policy_ctx
+    assert plain(x)[0] == 'original'
+    assert plain(x, sequence_parallel_context=ref_ctx) is ref_ctx
+
+
+def test_router_replay_uses_model_owner_outside_forward(monkeypatch):
+    from twinkle.model.transformers.moe import router_replay as rr
+
+    class Context:
+        sp_world_size = 2
+        real_position_ids = None
+
+        def __init__(self, offset):
+            self.offset = offset
+
+        def gather(self, tensor, **kwargs):
+            return tensor + self.offset
+
+    models = [torch.nn.Linear(1, 1) for _ in range(3)]
+    monkeypatch.setitem(sp._MODULE_TO_CTX, models[0], Context(10))
+    monkeypatch.setitem(sp._MODULE_TO_CTX, models[1], Context(20))
+    monkeypatch.setattr(rr, '_find_moe_blocks_with_names', lambda model: [('block', model)])
+    monkeypatch.setattr(rr, '_registry', {'block': SimpleNamespace(recorded_indices=torch.zeros(2, 1))})
+    assert sp._current_ctx() is None
+    assert rr.get_router_replay_data(models[0]).eq(10).all()
+    assert rr.get_router_replay_data(models[1]).eq(20).all()
+    assert rr.get_router_replay_data(models[2]).eq(0).all()

From 7bca22bab06723f21570f94a49c96e447f97945e Mon Sep 17 00:00:00 2001
From: clc <clc0301@cmbchina.com>
Date: Thu, 15 Oct 2026 16:26:42 +0800
Subject: [PATCH 2/8] sp: add data-key contracts and model adapters

---
 .../model/transformers/moe/router_replay.py   |  18 +-
 .../strategy/sequence_parallel/__init__.py    | 467 +++++++-----------
 .../strategy/sequence_parallel/adapters.py    | 139 ++++++
 .../strategy/sequence_parallel/data_keys.py   | 223 +++++++++
 .../model/transformers/transformers.py        |  23 +-
 src/twinkle/processor/base.py                 | 114 ++++-
 tests/transformers/test_sp_data_keys.py       | 306 ++++++++++++
 tests/transformers/test_sp_model_adapters.py  |  99 ++++
 8 files changed, 1069 insertions(+), 320 deletions(-)
 create mode 100644 src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py
 create mode 100644 src/twinkle/model/transformers/strategy/sequence_parallel/data_keys.py
 create mode 100644 tests/transformers/test_sp_data_keys.py
 create mode 100644 tests/transformers/test_sp_model_adapters.py

diff --git a/src/twinkle/model/transformers/moe/router_replay.py b/src/twinkle/model/transformers/moe/router_replay.py
index 2860ea0a4..a59b39081 100644
--- a/src/twinkle/model/transformers/moe/router_replay.py
+++ b/src/twinkle/model/transformers/moe/router_replay.py
@@ -103,16 +103,14 @@ def set_router_replay_data(
     sp = sp_context if sp_context is not None else _MODULE_TO_CTX.get(model)
     sp_world_size = getattr(sp, 'sp_world_size', None) or 1
     if sp_world_size > 1:
-        _, _, _, _, _, _, extra_values = sp.pad_and_split_inputs(
-            None,
-            None,
-            None,
-            None,
-            None,
-            None,
-            real_position_ids=sp.real_position_ids,
-            extra_split_values=[(routed_experts, 0, 1)])
-        routed_experts = extra_values[0]
+        from ..strategy.sequence_parallel.data_keys import DataKeySpec, SPKeyCollection, SPKeyValue
+
+        spec = DataKeySpec('routed_experts', split_dim=1, pad_value=0, stage='router_replay',
+                           pad_at='router_replay', split_at='router_replay')
+        keys = SPKeyCollection()
+        keys.add(SPKeyValue(spec, routed_experts))
+        sp.pad_and_split_inputs(keys, boundary='router_replay', real_position_ids=sp.real_position_ids)
+        routed_experts = keys.get_value('routed_experts')
 
     # [bs, seq_len, num_moe_layers, topk] -> [total_seq, num_moe_layers, topk]
     routed_experts = routed_experts.flatten(0, 1).to(Platform.get_local_device())
diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
index ae6c31cee..109261dd2 100644
--- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
@@ -6,21 +6,17 @@
 from copy import copy
 from dataclasses import asdict, dataclass, is_dataclass
 from transformers import PreTrainedTokenizer
-from types import SimpleNamespace
 from typing import Any, Dict, List, Optional, Tuple, Union
 from weakref import WeakKeyDictionary
 
-from functools import partial
-
-from twinkle.patch import apply_patch
 from twinkle.utils import DeviceMesh
 from twinkle.utils.transformers_utils import get_llm_model
 from twinkle.utils.utils import call_with_supported_kwargs, has_signature_parameter
-from twinkle.patch.transformers_qwen3_vl_deepstack import Qwen3VLDeepstackSPPatch
-from .linear_attention_sp import Qwen3_5GatedDeltaNetUlyssesPatch, _iter_qwen35_gated_delta_net_classes
+from .adapters import ModelSPAdapter, adapters_for_model, register_model_sp_adapter
 from .utils import (DistributedAttention, GatherLoss, _derive_sequence_parallel_sizes, _get_seq_groups_from_device_mesh,
                     _get_ulysses_size, _SeqAllToAll, get_config_attr, get_cu_seqlens_from_position_ids, is_hccl_backend,
-                    is_moe_config, post_all2all)
+                    post_all2all)
+from .data_keys import DataKeySpec, InputKeyRegistry, OutputKeyRegistry, SPKeyCollection, SPKeyValue
 
 
 # ---------------------------------------------------------------------------
@@ -70,16 +66,6 @@ def _register_module_tree(model: torch.nn.Module, ctx) -> None:
         _MODULE_TO_CTX[module] = ctx
 
 
-def is_qwen3_vl(model):
-    mt = getattr(getattr(model, 'config', None), 'model_type', '')
-    return 'qwen3_vl' in mt
-
-
-def is_qwen3_omni(model):
-    mt = getattr(getattr(model, 'config', None), 'model_type', '')
-    return 'qwen3_omni' in mt
-
-
 def _call_create_causal_mask(fn, config, input_embeds, attention_mask, cache_position_or_past_key_values, *args,
                              **kwargs):
     if has_signature_parameter(fn, 'cache_position'):
@@ -129,19 +115,22 @@ def __init__(self):
         self.causal_mask_func = None
         self.extra_kwargs = {}
         self._dist_attns = {}
+        self.input_key_registry = InputKeyRegistry()
+        self.output_key_registry = OutputKeyRegistry()
+        self.model_adapters: Tuple[ModelSPAdapter, ...] = ()
+        self.position_ids_mode = 'auto'
 
     @property
     def real_position_ids(self) -> torch.Tensor:
         """The real position ids, this is different from the position_ids in mrope"""
         return self.extra_kwargs.get('position_ids')
 
-    @staticmethod
-    def _extract_real_position_ids(position_ids: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
+    def _extract_real_position_ids(self, position_ids: Optional[torch.Tensor]) -> Optional[torch.Tensor]:
         if position_ids is None or not torch.is_tensor(position_ids):
             return position_ids
         # Cached real positions are already 2D, so only reduce the original
         # three-axis mRoPE representation.
-        if position_ids.dim() == 3:
+        if position_ids.dim() == 3 and self.position_ids_mode in ('mrope', 'auto'):
             return position_ids[0]
         return position_ids
 
@@ -447,30 +436,16 @@ def pre_forward_split_hook(_self, args, kwargs):
                 if self.world_size == 1:
                     return args, kwargs
                 # Pad to multiple of SP size and split inputs per SP rank before forward.
-                input_ids = kwargs.get('input_ids', None)
-                inputs_embeds = kwargs.get('inputs_embeds', None)
                 position_ids = kwargs['position_ids']
                 real_position_ids = self._extract_real_position_ids(position_ids)
-                attention_mask = kwargs.get('attention_mask', None)
                 cache_position = kwargs.get('cache_position', None)
-                if hasattr(_self, 'language_model'):
-                    embed_tokens = getattr(_self.language_model, 'embed_tokens', None)
-                else:
-                    embed_tokens = getattr(_self, 'embed_tokens', None)
-                input_ids, inputs_embeds, _, position_ids, attention_mask, _, _ = self.pad_and_split_inputs(
-                    input_ids,
-                    inputs_embeds,
-                    None,
-                    position_ids,
-                    attention_mask,
-                    None,
-                    embed_tokens=embed_tokens,
-                    real_position_ids=real_position_ids,
-                    cache_position=cache_position)
-                kwargs['input_ids'] = input_ids
-                kwargs['inputs_embeds'] = inputs_embeds
-                kwargs['position_ids'] = position_ids
-                kwargs['attention_mask'] = attention_mask
+                specs = tuple(spec for spec in self.input_key_registry if spec.stage == 'forward')
+                common_names = ('input_ids', 'inputs_embeds', 'position_ids', 'attention_mask')
+                keys = SPKeyCollection.from_mapping(kwargs, specs, common_names=common_names)
+                self.pad_and_split_inputs(keys, boundary='forward', real_position_ids=real_position_ids,
+                                          cache_position=cache_position)
+                for item in keys.values():
+                    kwargs[item.name] = item.value
                 return args, kwargs
             except BaseException:
                 # PyTorch does not guarantee forward hooks run when a pre-hook
@@ -486,77 +461,22 @@ def post_forward_pop_hook(_self, args, output):
         # PyTorch also calls this hook if a pre-hook or the forward raises.
         base_model.register_forward_hook(post_forward_pop_hook, always_call=True)
 
-
-    def _prepare_multimodal_deepstack(self, base_model: torch.nn.Module):
-        if not is_qwen3_vl(base_model):
-            return
-        apply_patch(base_model, Qwen3VLDeepstackSPPatch, sequence_parallel=self)
-
-
-    @staticmethod
-    def _is_qwen35_model(model: torch.nn.Module) -> bool:
-        config = getattr(model, 'config', None)
-        model_type = str(getattr(config, 'model_type', '') or '')
-        if model_type in {'qwen3_5', 'qwen3_5_moe'}:
-            return True
-
-        architectures = getattr(config, 'architectures', None) or []
-        if any('Qwen3_5' in str(arch) for arch in architectures):
-            return True
-
-        model_module = getattr(model.__class__, '__module__', '') or ''
-        return 'transformers.models.qwen3_5' in model_module or 'transformers.models.qwen3_5_moe' in model_module
-
-
-    def _prepare_qwen35_linear_attention(self, model: torch.nn.Module):
-        has_qwen35_linear_attention = self._is_qwen35_model(model)
-        if not has_qwen35_linear_attention:
-            gated_delta_net_classes = tuple(_iter_qwen35_gated_delta_net_classes())
-            if not gated_delta_net_classes:
-                return
-            has_qwen35_linear_attention = any(isinstance(module, gated_delta_net_classes) for module in model.modules())
-        if not has_qwen35_linear_attention:
-            return
-        if int(self.rp_world_size or 1) > 1:
-            raise NotImplementedError(
-                'SequenceParallel: Qwen3.5 linear attention sequence parallel does not support rp_world_size > 1 '
-                '(derived ring attention).')
-        apply_patch(None, Qwen3_5GatedDeltaNetUlyssesPatch, sequence_parallel=self)
-
-
-    def _prepare_moe_aux_loss(self, base_model: torch.nn.Module):
-
-        def moe_aux_loss_hook(module, args, kwargs, output):
-            router_logits = getattr(output, 'router_logits', None)
-            if router_logits is None:
-                return output
-
-            attention_mask = kwargs['attention_mask']
-            if attention_mask is None:
-                batch_size = 1
-            else:
-                batch_size = attention_mask.shape[0]
-
-            assert router_logits[0].shape[0] % batch_size == 0
-            seq_len = router_logits[0].shape[0] // batch_size
-
-            _gathered_logits = []
-            for i in range(batch_size):
-                _slice = slice(i * seq_len, (i + 1) * seq_len)
-                _bs_logits = [logit[_slice] for logit in router_logits]
-                compute_device = _bs_logits[0].device
-                _bs_logits = torch.stack([layer_gate.to(compute_device) for layer_gate in _bs_logits], dim=0)
-                _bs_logits, _ = GatherLoss.apply(self, _bs_logits, None, 1, self.real_position_ids)
-                _gathered_logits.append(_bs_logits)
-            router_logits = torch.stack(_gathered_logits, dim=0)
-            if self.real_position_ids is not None:
-                router_logits = router_logits[:, :, :self.real_position_ids.shape[1], :]
-            output['router_logits'] = tuple(
-                [logit.reshape(-1, logit.shape[-1]) for logit in router_logits.split(1, dim=1)])
-            return output
-
-        base_model.register_forward_hook(moe_aux_loss_hook, with_kwargs=True)
-
+    def _configure_model_adapters(self, model: torch.nn.Module, text_model: torch.nn.Module) -> None:
+        adapters = adapters_for_model(model)
+        if text_model is not model:
+            adapters += tuple(adapter for adapter in adapters_for_model(text_model) if adapter not in adapters)
+        modes = {adapter.position_ids_mode for adapter in adapters if adapter.position_ids_mode != 'auto'}
+        if len(modes) > 1:
+            raise ValueError(f'Conflicting SP position_ids modes: {sorted(modes)}')
+        self.model_adapters = adapters
+        self.position_ids_mode = next(iter(modes), 'auto')
+        for adapter in adapters:
+            for spec in adapter.input_key_specs():
+                self.input_key_registry.register(spec)
+            for spec in adapter.output_key_specs():
+                self.output_key_registry.register(spec)
+        for adapter in adapters:
+            adapter.configure(self, text_model)
 
     def prepare(
         self,
@@ -608,11 +528,8 @@ def prepare(
             # these operations are global initializations and patches
             self._prepare_flash_attn()
             SequenceParallel._global_inited = True
-        self._prepare_qwen35_linear_attention(llm_model)
         self._prepare_forward_hook(llm_model)
-        self._prepare_multimodal_deepstack(llm_model)
-        if is_moe_config(getattr(llm_model, 'config', None)):
-            self._prepare_moe_aux_loss(llm_model)
+        self._configure_model_adapters(model, llm_model)
 
         self.model_dtype = next(model.parameters()).dtype
         self.tokenizer = tokenizer
@@ -681,21 +598,17 @@ def pad_and_split_mm_tokens(self, visual_mask, mm_embeds):
         empty_embeds = torch.empty(
             (input_ids.shape[0], input_ids.shape[1], mm_embeds.shape[-1])).to(mm_embeds.device).to(mm_embeds.dtype)
         empty_embeds[visual_mask] = mm_embeds
-
-        embeds = SimpleNamespace(weight=mm_embeds)
-
-        _, split_input_embeds, _, _, _, _, extra_values = self.pad_and_split_inputs(
-            None,
-            empty_embeds,
-            None,
-            None,
-            None,
-            None,
-            embeds,
-            self.real_position_ids,
-            extra_split_values=[(visual_mask, 0, -1)])
-        visual_mask = extra_values[0]
-        return visual_mask, split_input_embeds[visual_mask]
+        embeds_spec = DataKeySpec('visual_sequence_embeds', split_dim=1, pad_value=0, stage='model_internal',
+                                  pad_at='model_internal', split_at='model_internal')
+        mask_spec = DataKeySpec('visual_position_mask', split_dim=-1, pad_value=0, stage='model_internal',
+                                pad_at='model_internal', split_at='model_internal')
+        keys = SPKeyCollection()
+        keys.add(SPKeyValue(embeds_spec, empty_embeds))
+        keys.add(SPKeyValue(mask_spec, visual_mask))
+        self.pad_and_split_inputs(keys, boundary='model_internal', real_position_ids=self.real_position_ids)
+        local_mask = keys.get_value('visual_position_mask')
+        local_embeds = keys.get_value('visual_sequence_embeds')
+        return local_mask, local_embeds[local_mask]
 
     def gather(self, local_output, dim: int, position_ids=None):
         """Gather tensor for sequence parallel - reverse of split."""
@@ -826,104 +739,86 @@ def _split_key(self, value, spec, position_ids=None, legacy_roll=False):
         return self.split(value, dim=spec.split_dim, position_ids=position_ids)
 
     def pad_and_split_inputs(self,
-                             input_ids,
-                             input_embeds,
-                             labels,
-                             position_ids,
-                             attention_mask,
-                             loss_scale,
-                             embed_tokens=None,
+                             keys: SPKeyCollection,
+                             *,
+                             boundary: str,
                              real_position_ids=None,
-                             cache_position=None,
-                             extra_split_values=None):
-        """Common implementation for padding and splitting inputs
-
-        Pad to a length divisible by the sequence-parallel size, then split across SP ranks.
-
-        Args:
-            input_ids: input_ids
-            input_embeds: input_embeds
-            labels: labels
-            position_ids: position_ids or, position_ids for mrope
-            attention_mask: attention_mask
-            loss_scale: loss_scale
-            embed_tokens: embed_tokens
-            real_position_ids: the real position_ids to represent the seq length information
-            extra_split_values: List of Tuples for extra split values, e.g.: (tensor, pad_value, split_dim)
+                             cache_position=None):
+        """Apply the declared pad/split operations to a batch-local key collection.
+
+        ``real_position_ids`` describes the unpadded sequence layout used to
+        preserve packed boundaries. ``cache_position`` is model-call context;
+        it is intentionally not treated as a sequence key or transformed.
         """
-        tokenizer = self.tokenizer
+        if not isinstance(keys, SPKeyCollection):
+            raise TypeError('pad_and_split_inputs expects an SPKeyCollection')
+        if boundary not in ('prepare', 'forward', 'model_internal', 'router_replay'):
+            raise ValueError(f'Unknown SP processing boundary: {boundary!r}')
+
+        position_ids = keys.get_value('position_ids')
         real_position_ids = real_position_ids if real_position_ids is not None else position_ids
-        extra_values = []
-        batch_size = input_ids.shape[
-            0] if input_ids is not None else input_embeds.shape[0] if input_embeds is not None else None
+        input_ids = keys.get_value('input_ids')
+        inputs_embeds = keys.get_value('inputs_embeds')
+        model_inputs = input_ids if input_ids is not None else inputs_embeds
+        batch_size = model_inputs.shape[0] if model_inputs is not None else None
         if real_position_ids is not None and batch_size is not None and real_position_ids.shape[0] == batch_size:
-            # TODO clone everytime, but the position_ids is a small tensor
             self.extra_kwargs['position_ids'] = real_position_ids.clone()
-        if input_ids is not None:
-            input_ids = self.pad(input_ids, padding_value=tokenizer.pad_token_id, position_ids=real_position_ids)
-            self.extra_kwargs['input_ids'] = input_ids.clone()
-        if input_embeds is not None:
-            pad_emb = torch.zeros(
-                (1, embed_tokens.weight.shape[-1])).to(embed_tokens.weight.device).to(embed_tokens.weight.dtype)
-            input_embeds = self.pad(input_embeds, padding_value=pad_emb, position_ids=real_position_ids)
-        batch_size = input_ids.shape[
-            0] if input_ids is not None else input_embeds.shape[0] if input_embeds is not None else 1
-        if self.rp_world_size > 1 and batch_size > 1:
+        if self.rp_world_size > 1 and batch_size is not None and batch_size > 1:
             raise NotImplementedError(
                 'Derived ring attention only supports padding-free / packed batches with batch_size == 1.')
-        if position_ids is not None:
-            position_ids = self.pad(position_ids, padding_value=-1, position_ids=real_position_ids, dim=-1)
-        if labels is not None:
-            labels = self.pad(labels, padding_value=-100, position_ids=real_position_ids)
-        if loss_scale is not None:
-            loss_scale = self.pad(loss_scale, padding_value=0., position_ids=real_position_ids)
-        if real_position_ids is not None:
-            real_position_ids = self.pad(real_position_ids, padding_value=-1, position_ids=real_position_ids)
-        # Build a 2D attention_mask whenever we padded for SP alignment so FlashAttention2 can unpad correctly.
-        # For packed batches (batch_size==1 with multiple position_id resets), relying on position_ids alone is
-        # unsafe if we also appended SP-alignment padding (position_ids=-1), because HF's FA2 varlen path will
-        # include the padded tail in the last segment when attention_mask is None.
-        if (input_ids is not None or input_embeds is not None) and batch_size > 1:
-            # not padding_free, so not ring-attention
-            inputs = input_ids if input_ids is not None else input_embeds
-            attn_shape = inputs.shape[1]  # The sequence length
-            if attention_mask is None:
-                # Mask out padded positions introduced by sequence-parallel padding.
-                # `real_position_ids` is padded with `-1` (see above), so use it to build a valid-token mask.
-                attention_mask = (real_position_ids != -1).to(dtype=torch.int64)
-            # no need position_ids here, because padding_free does not need attention_mask,
-            # so this is not ring-attention
-            attention_mask = self.pad(attention_mask, padding_value=0)
-            local_cache_position = torch.arange(0, attn_shape, device=inputs.device)
-            # FlashAttention2 expects a 2D padding mask (or None). Converting it to a 4D causal mask here breaks
-            # the later per-rank sequence split and changes the attention contract relative to the baseline path.
-            if (cache_position is None and hasattr(self, 'causal_mask_func') and self.causal_mask_func is not None
+
+        split_is_due = any(item.spec.split_at == boundary and item.value is not None for item in keys.values())
+        padded_positions = (self.pad(real_position_ids, padding_value=-1, position_ids=real_position_ids)
+                            if real_position_ids is not None and split_is_due else None)
+
+        # Pad every declared key at this boundary. This includes both built-in
+        # keys and registered extra keys.
+        for item in keys.values():
+            spec, value = item.spec, item.value
+            if value is None:
+                continue
+            if not torch.is_tensor(value):
+                raise TypeError(f'SP input {spec.name!r} must be a tensor')
+            if spec.pad_at == boundary:
+                # Preserve the existing attention-mask contract: it is padded
+                # only for ordinary batched sequences, not packed batch_size=1.
+                if spec.name != 'attention_mask' or (batch_size is not None and batch_size > 1):
+                    value = self._pad_key(value, spec, real_position_ids)
+            item.value = value
+
+        padded_ids = keys.get_value('input_ids')
+        padded_embeds = keys.get_value('inputs_embeds')
+        padded_model_inputs = padded_ids if padded_ids is not None else padded_embeds
+        if input_ids is not None and boundary == 'forward':
+            # Attention wrappers and multimodal patches need the padded global IDs.
+            self.extra_kwargs['input_ids'] = padded_ids.clone()
+
+        # Keep the legacy generated 2D mask and causal-mask conversion on the
+        # forward boundary; these depend on model inputs and cache state.
+        if boundary == 'forward' and padded_model_inputs is not None and batch_size > 1:
+            attn_item = keys.get('attention_mask')
+            if attn_item is None or attn_item.value is None:
+                attention_mask = (padded_positions != -1).to(dtype=torch.int64)
+                if attn_item is None:
+                    spec = self.input_key_registry.get('attention_mask')
+                    attn_item = SPKeyValue(spec, attention_mask)
+                    keys.add(attn_item, common=True)
+                else:
+                    attn_item.value = attention_mask
+            else:
+                attention_mask = attn_item.value
+            local_cache_position = torch.arange(0, padded_model_inputs.shape[1], device=padded_model_inputs.device)
+            if (cache_position is None and self.causal_mask_func is not None
                     and self.attn_implementation != 'flash_attention_2'):
-                attention_mask = self.causal_mask_func(attention_mask, inputs.to(self.model_dtype),
+                attention_mask = self.causal_mask_func(attention_mask, padded_model_inputs.to(self.model_dtype),
                                                        local_cache_position, None, None)
-        if extra_split_values is not None:
-            for (tensor, pad_value, split_dim) in extra_split_values:
-                extra_values.append(
-                    self.pad(tensor, padding_value=pad_value, position_ids=real_position_ids, dim=split_dim))
-        if input_ids is not None:
-            input_ids = self.split(input_ids, dim=1, position_ids=real_position_ids)
-        if input_embeds is not None:
-            input_embeds = self.split(input_embeds, dim=1, position_ids=real_position_ids)
-        if labels is not None:
-            labels = self.split(labels, dim=-1, position_ids=real_position_ids)
-        if loss_scale is not None:
-            loss_scale = torch.roll(loss_scale, shifts=-1, dims=-1)
-            loss_scale = self.split(loss_scale, dim=-1, position_ids=real_position_ids)
-
-        if position_ids is not None:
-            position_ids = self.split(position_ids, dim=-1, position_ids=real_position_ids)
-        # if attention_mask is not None and torch.is_tensor(attention_mask) and attention_mask.dim() == 2:
-        #     attention_mask = self.split(attention_mask, dim=1, position_ids=real_position_ids)
-        if extra_split_values is not None:
-            for i in range(len(extra_values)):
-                extra_values[i] = self.split(
-                    extra_values[i], dim=extra_split_values[i][2], position_ids=real_position_ids)
-        return input_ids, input_embeds, labels, position_ids, attention_mask, loss_scale, extra_values
+                attn_item.value = attention_mask
+
+        # Split after padding so every field uses the same padded sequence frame.
+        for item in keys.values():
+            if item.value is not None and item.spec.split_at == boundary:
+                item.value = self._split_key(item.value, item.spec, padded_positions)
+        return keys
 
     def _init_device_mesh(self, device_mesh: Optional[DeviceMesh] = None):
         """Initialize process groups for sequence parallel."""
@@ -958,11 +853,7 @@ def _is_packed_position_ids(position_ids: Optional[torch.Tensor]) -> bool:
         return False
 
     def prepare_inputs(self, inputs):
-        """Prepare inputs
-
-        1. set extra_kwargs['position_ids']
-        2. split labels
-        """
+        """Cache position context and run keys declared for the prepare boundary."""
         input_ids = inputs.get('input_ids')
         position_ids = inputs.get('position_ids')
         padding_free = bool(inputs.pop('padding_free', False))
@@ -977,18 +868,11 @@ def prepare_inputs(self, inputs):
         self.extra_kwargs['padding_free'] = padding_free
         if input_ids is not None:
             self.extra_kwargs['input_ids'] = input_ids.clone()
-        if 'labels' in inputs:
-            labels = inputs.get('labels')
-            _, _, labels, _, _, _, _ = self.pad_and_split_inputs(
-                None,
-                None,
-                labels,
-                None,
-                None,
-                None,
-                real_position_ids=real_position_ids,
-            )
-            inputs['labels'] = labels
+        specs = tuple(spec for spec in self.input_key_registry if spec.stage == 'auxiliary')
+        keys = SPKeyCollection.from_mapping(inputs, specs, common_names=('labels', 'completion_mask', 'loss_scale'))
+        self.pad_and_split_inputs(keys, boundary='prepare', real_position_ids=real_position_ids)
+        for item in keys.values():
+            inputs[item.name] = item.value
         return inputs
 
 
@@ -1112,6 +996,32 @@ def preprocess_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
             return inputs
         return self._ctx.prepare_inputs(inputs)
 
+    def register_input_key(self, spec: DataKeySpec) -> None:
+        """Declare a token-aligned input on this model only.
+
+        Auxiliary inputs are kept out of the model; stage=forward inputs are
+        sliced in its pre-hook. need_gather selects the loss-side collection.
+        """
+        self._ctx.input_key_registry.register(spec)
+
+    @property
+    def input_key_specs(self) -> Tuple[DataKeySpec, ...]:
+        """Read-only snapshot used by processors to preserve registered keys."""
+        return tuple(self._ctx.input_key_registry)
+
+    def register_output_key(self, spec: DataKeySpec) -> None:
+        """Register a model-produced sequence output to collect before loss."""
+        self._ctx.output_key_registry.register(spec)
+
+    @property
+    def output_key_specs(self) -> Tuple[DataKeySpec, ...]:
+        return tuple(self._ctx.output_key_registry)
+
+    @property
+    def model_sp_adapters(self) -> Tuple[ModelSPAdapter, ...]:
+        """The model-family adapters selected during context preparation."""
+        return self._ctx.model_adapters
+
     def postprocess_outputs(self, outputs: Any) -> Any:
         if (not self.enabled or self.ulysses_size <= 1 or not self.sp_config.get('gather_logits', True)):
             return outputs
@@ -1132,10 +1042,18 @@ def postprocess_outputs(self, outputs: Any) -> Any:
         outputs['logits'] = gathered
         return outputs
 
-    def _trim_gathered_sequence_padding(self, tensor: torch.Tensor, real_position_ids: torch.Tensor) -> torch.Tensor:
+    def _trim_gathered_sequence_padding(self, tensor: torch.Tensor, real_position_ids: torch.Tensor,
+                                        dim: int = 1) -> torch.Tensor:
         if real_position_ids is None or not torch.is_tensor(real_position_ids) or real_position_ids.dim() < 2:
             return tensor
-        if self._ctx.rp_world_size > 1:
+        dim = dim % tensor.dim()
+
+        def take(start, end):
+            slices = [slice(None)] * tensor.dim()
+            slices[dim] = slice(start, end)
+            return tensor[tuple(slices)]
+
+        if (self._ctx.rp_world_size or 1) > 1:
             cu_seqlens = get_cu_seqlens_from_position_ids(real_position_ids)
             pieces = []
             padded_offset = 0
@@ -1143,10 +1061,20 @@ def _trim_gathered_sequence_padding(self, tensor: torch.Tensor, real_position_id
             for i in range(len(cu_seqlens) - 1):
                 real_len = int((cu_seqlens[i + 1] - cu_seqlens[i]).item())
                 padded_len = math.ceil(real_len / divisor) * divisor
-                pieces.append(tensor[:, padded_offset:padded_offset + real_len])
+                pieces.append(take(padded_offset, padded_offset + real_len))
                 padded_offset += padded_len
-            return torch.cat(pieces, dim=1).contiguous() if pieces else tensor[:, :0].contiguous()
-        return tensor[:, :real_position_ids.shape[-1]].contiguous()
+            return torch.cat(pieces, dim=dim).contiguous() if pieces else take(0, 0).contiguous()
+        return take(0, real_position_ids.shape[-1]).contiguous()
+
+    def _gather_key(self, value, spec, real_position_ids):
+        if not torch.is_tensor(value) or value.dim() < 2:
+            raise TypeError(f'SP key {spec.name!r} must be a batched sequence tensor')
+        dim = spec.sequence_dim(value.dim())
+        if spec.autograd:
+            gathered, _ = GatherLoss.apply(self._ctx, value, None, dim, real_position_ids)
+        else:
+            gathered = self._ctx.gather(value, dim=dim, position_ids=real_position_ids)
+        return self._trim_gathered_sequence_padding(gathered, real_position_ids, dim=dim)
 
     def gather_loss_tensors(
         self,
@@ -1157,46 +1085,23 @@ def gather_loss_tensors(
             return inputs, outputs
         if not self.enabled or self.ulysses_size <= 1:
             return inputs, outputs
-        labels = inputs.get('labels')
-        logps = outputs.get('logps')
-        # Fused-linear-CE path: the model forward skipped the lm_head GEMM and
-        # stashed the lm_head module under ``outputs['lm_head']`` (see
-        # ``TransformersFusedCEPatch``), so ``outputs['logps']`` is absent and
-        # ``outputs['logits']`` holds the last hidden state (a per-rank seq
-        # shard under SP). Gather hidden states + labels across the sequence
-        # group so every rank computes the same full loss — mirrors the logps
-        # gather below. The hidden tensor ([B, T_local, H]) is tiny vs logits,
-        # and the fused loss still chunks the GEMM internally, so the memory
-        # win is preserved.
-        if (logps is None and labels is not None and 'lm_head' in outputs and torch.is_tensor(outputs.get('logits'))
-                and outputs['logits'].dim() >= 2):
-            hidden = outputs['logits']
-            inputs = copy(inputs)
-            outputs = copy(outputs)
-            real_position_ids = self._ctx.real_position_ids
-            gathered_hidden, gathered_labels = GatherLoss.apply(self._ctx, hidden, labels, 1, real_position_ids)
-            gathered_hidden = self._trim_gathered_sequence_padding(gathered_hidden, real_position_ids)
-            gathered_labels = self._trim_gathered_sequence_padding(gathered_labels, real_position_ids)
-            outputs['logits'] = gathered_hidden
-            inputs['labels'] = gathered_labels
-            return inputs, outputs
-        if labels is None or logps is None:
-            return inputs, outputs
-        if not torch.is_tensor(logps) or logps.dim() < 2:
-            raise TypeError('SequenceParallelStrategy.gather_loss_inputs expects outputs[\"logps\"] to be a '
-                            f'sequence tensor, got type={type(logps)} shape={getattr(logps, "shape", None)}')
-        inputs = copy(inputs)
-        outputs = copy(outputs)
+        inputs, outputs = copy(inputs), copy(outputs)
         real_position_ids = self._ctx.real_position_ids
-        gathered_logps, gathered_labels = GatherLoss.apply(self._ctx, logps, labels, 1, real_position_ids)
-        gathered_logps = self._trim_gathered_sequence_padding(gathered_logps, real_position_ids)
-        gathered_labels = self._trim_gathered_sequence_padding(gathered_labels, real_position_ids)
-        outputs['logps'] = gathered_logps
-        inputs['labels'] = gathered_labels
-        entropies = outputs.get('entropies')
-        if entropies is not None and torch.is_tensor(entropies) and entropies.dim() >= 2:
-            gathered_entropies, _ = GatherLoss.apply(self._ctx, entropies, labels, 1, real_position_ids)
-            outputs['entropies'] = self._trim_gathered_sequence_padding(gathered_entropies, real_position_ids)
+        # Gather only auxiliary inputs whose contracts request it. Forward
+        # keys remain model-owned; task-specific pooled outputs are handled upstream.
+        for spec in self.input_key_specs:
+            value = inputs.get(spec.name)
+            if spec.stage == 'auxiliary' and spec.need_gather and value is not None:
+                inputs[spec.name] = self._gather_key(value, spec, real_position_ids)
+
+        # Fused CE stores hidden states in the logits slot. This legacy switch
+        # selects the registered logits rule without gathering real vocabulary
+        # logits in the ordinary logps path.
+        fused_hidden = outputs.get('logps') is None and 'lm_head' in outputs and inputs.get('labels') is not None
+        for spec in self.output_key_specs:
+            value = outputs.get(spec.name)
+            if value is not None and (spec.need_gather or (spec.name == 'logits' and fused_hidden)):
+                outputs[spec.name] = self._gather_key(value, spec, real_position_ids)
         return inputs, outputs
 
     def wrap_model(self, model, optimizer=None):
diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py b/src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py
new file mode 100644
index 000000000..91722af04
--- /dev/null
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py
@@ -0,0 +1,139 @@
+"""Model-family SP adapters and their small registry.
+
+Adapters own model-family detection and the model-specific setup steps. The
+sequence-parallel context remains responsible for generic padding, sharding and
+collectives; adapters only install behavior that depends on model internals.
+"""
+from dataclasses import dataclass
+from typing import Tuple
+
+import torch
+
+from twinkle.patch import apply_patch
+from twinkle.patch.transformers_qwen3_vl_deepstack import Qwen3VLDeepstackSPPatch
+from .data_keys import DataKeySpec
+from .linear_attention_sp import Qwen3_5GatedDeltaNetUlyssesPatch, _iter_qwen35_gated_delta_net_classes
+from .utils import GatherLoss, is_moe_config
+
+
+class ModelSPAdapter:
+    """Base contract for model-specific SP behavior."""
+
+    name = 'generic'
+    position_ids_mode = 'auto'
+
+    def matches(self, model) -> bool:
+        return False
+
+    def input_key_specs(self) -> Tuple[DataKeySpec, ...]:
+        return ()
+
+    def output_key_specs(self) -> Tuple[DataKeySpec, ...]:
+        return ()
+
+    def configure(self, context, model) -> None:
+        """Install hooks on the model's text root. Generic adapter is a no-op."""
+
+
+class Qwen3VLSPAdapter(ModelSPAdapter):
+    name = 'qwen3_vl'
+    position_ids_mode = 'mrope'
+
+    def matches(self, model) -> bool:
+        config = getattr(model, 'config', None)
+        model_type = str(getattr(config, 'model_type', '') or '')
+        return 'qwen3_vl' in model_type
+
+    def configure(self, context, model) -> None:
+        apply_patch(model, Qwen3VLDeepstackSPPatch, sequence_parallel=context)
+
+
+class Qwen35LinearAttentionSPAdapter(ModelSPAdapter):
+    name = 'qwen3_5_linear_attention'
+
+    @staticmethod
+    def _matches_config_or_class(model) -> bool:
+        config = getattr(model, 'config', None)
+        model_type = str(getattr(config, 'model_type', '') or '')
+        if model_type in {'qwen3_5', 'qwen3_5_moe'}:
+            return True
+        architectures = getattr(config, 'architectures', None) or []
+        if any('Qwen3_5' in str(architecture) for architecture in architectures):
+            return True
+        module_name = getattr(model.__class__, '__module__', '') or ''
+        return ('transformers.models.qwen3_5' in module_name
+                or 'transformers.models.qwen3_5_moe' in module_name)
+
+    def matches(self, model) -> bool:
+        if self._matches_config_or_class(model):
+            return True
+        gated_delta_net_classes = tuple(_iter_qwen35_gated_delta_net_classes())
+        return bool(gated_delta_net_classes) and any(
+            isinstance(module, gated_delta_net_classes) for module in model.modules())
+
+    def configure(self, context, model) -> None:
+        if int(context.rp_world_size or 1) > 1:
+            raise NotImplementedError(
+                'SequenceParallel: Qwen3.5 linear attention sequence parallel does not support '
+                'rp_world_size > 1 (derived ring attention).')
+        apply_patch(None, Qwen3_5GatedDeltaNetUlyssesPatch, sequence_parallel=context)
+
+
+class MoEAuxSPAdapter(ModelSPAdapter):
+    name = 'moe_aux'
+
+    def matches(self, model) -> bool:
+        config = getattr(model, 'config', None)
+        model_type = str(getattr(config, 'model_type', '') or '')
+        return model_type.endswith('_moe') or is_moe_config(config)
+
+    def configure(self, context, base_model) -> None:
+        def gather_router_logits(module, args, kwargs, output):
+            router_logits = getattr(output, 'router_logits', None)
+            if router_logits is None:
+                return output
+            attention_mask = kwargs['attention_mask']
+            batch_size = 1 if attention_mask is None else attention_mask.shape[0]
+            assert router_logits[0].shape[0] % batch_size == 0
+            local_seq_len = router_logits[0].shape[0] // batch_size
+            gathered_per_sample = []
+            for batch_idx in range(batch_size):
+                seq_slice = slice(batch_idx * local_seq_len, (batch_idx + 1) * local_seq_len)
+                sample_logits = [layer_logits[seq_slice] for layer_logits in router_logits]
+                compute_device = sample_logits[0].device
+                sample_logits = torch.stack([logits.to(compute_device) for logits in sample_logits], dim=0)
+                sample_logits, _ = GatherLoss.apply(context, sample_logits, None, 1, context.real_position_ids)
+                gathered_per_sample.append(sample_logits)
+            router_logits = torch.stack(gathered_per_sample, dim=0)
+            if context.real_position_ids is not None:
+                router_logits = router_logits[:, :, :context.real_position_ids.shape[1], :]
+            output['router_logits'] = tuple(
+                logits.reshape(-1, logits.shape[-1]) for logits in router_logits.split(1, dim=1))
+            return output
+
+        base_model.register_forward_hook(gather_router_logits, with_kwargs=True)
+
+
+_ADAPTERS = [Qwen3VLSPAdapter(), Qwen35LinearAttentionSPAdapter(), MoEAuxSPAdapter()]
+
+
+def register_model_sp_adapter(adapter):
+    """Register an adapter instance for future SP model initialization."""
+    if isinstance(adapter, type):
+        adapter = adapter()
+    if not isinstance(adapter, ModelSPAdapter):
+        raise TypeError('A model SP adapter must derive from ModelSPAdapter')
+    if not adapter.name or adapter.name == 'generic':
+        raise ValueError('A registered model SP adapter needs a stable non-generic name')
+    for current in _ADAPTERS:
+        if current.name == adapter.name:
+            if type(current) is type(adapter):
+                return current
+            raise ValueError(f'Conflicting model SP adapter name: {adapter.name!r}')
+    _ADAPTERS.append(adapter)
+    return adapter
+
+
+def adapters_for_model(model):
+    """Return every matching adapter; families can compose capabilities."""
+    return tuple(adapter for adapter in _ADAPTERS if adapter.matches(model))
diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/data_keys.py b/src/twinkle/model/transformers/strategy/sequence_parallel/data_keys.py
new file mode 100644
index 000000000..704fbf1e6
--- /dev/null
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/data_keys.py
@@ -0,0 +1,223 @@
+"""Per-model contracts for sequence-aligned inputs and outputs.
+
+Each key declares whether and where it is padded, split, gathered, or restored
+after packed-sequence processing. Forward and auxiliary keys share the same
+executor but run at different lifecycle boundaries. Loss-scale weights are
+already aligned with Twinkle's shifted labels, so the production rule does not
+roll them a second time.
+"""
+from dataclasses import dataclass, field
+from typing import Any, Dict, Iterable, Iterator, Mapping, Optional
+
+
+_BOUNDARIES = frozenset({'prepare', 'forward', 'model_internal', 'router_replay'})
+
+
+@dataclass(frozen=True)
+class DataKeySpec:
+    name: str
+    split_dim: int = -1
+    pad_value: int | float | str = 0
+    need_roll: bool = False
+    side: str = 'input'
+    stage: str = 'auxiliary'
+    need_split: bool = True
+    need_gather: bool = False
+    autograd: bool = False
+    owner: str = 'generic'
+    need_pad: Optional[bool] = None
+    pad_at: Optional[str] = 'auto'
+    split_at: Optional[str] = 'auto'
+    need_unpack: Optional[bool] = None
+
+    def __post_init__(self):
+        if not isinstance(self.name, str) or not self.name:
+            raise ValueError('An SP input key must have a non-empty name')
+        if not isinstance(self.split_dim, int):
+            raise TypeError('split_dim must be an integer')
+        if self.side not in ('input', 'output'):
+            raise ValueError('side must be input or output')
+        if self.stage not in ('forward', 'auxiliary', 'model_internal', 'router_replay', 'output'):
+            raise ValueError('Invalid SP key stage')
+        if (self.side == 'output') != (self.stage == 'output'):
+            raise ValueError('Output keys must use stage=output')
+        if self.owner not in ('generic', 'task', 'model', 'loss'):
+            raise ValueError('Invalid SP key owner')
+        if isinstance(self.pad_value, str) and (self.name != 'input_ids' or self.pad_value != 'tokenizer'):
+            raise ValueError('Only input_ids can use the tokenizer padding sentinel')
+        if self.side == 'input' and self.need_gather and not self.need_split:
+            raise ValueError('Only sharded input keys can request gathering')
+        if self.stage == 'forward' and self.need_gather:
+            raise ValueError('Model-forward keys cannot request loss gathering')
+        default_boundary = {
+            'auxiliary': 'prepare',
+            'forward': 'forward',
+            'model_internal': 'model_internal',
+            'router_replay': 'router_replay',
+            'output': None,
+        }[self.stage]
+        if self.need_pad is None:
+            object.__setattr__(self, 'need_pad', self.need_split and self.side == 'input')
+        if self.pad_at == 'auto':
+            object.__setattr__(self, 'pad_at', default_boundary if self.need_pad else None)
+        elif self.pad_at is not None and self.pad_at not in _BOUNDARIES:
+            raise ValueError(f'Invalid SP padding boundary: {self.pad_at!r}')
+        elif (self.pad_at is None) == bool(self.need_pad):
+            raise ValueError('need_pad and pad_at specify conflicting SP padding behavior')
+        if self.pad_at != 'auto':
+            object.__setattr__(self, 'need_pad', self.pad_at is not None)
+        if self.split_at == 'auto':
+            object.__setattr__(self, 'split_at', default_boundary if self.need_split else None)
+        elif self.split_at is not None and self.split_at not in _BOUNDARIES:
+            raise ValueError(f'Invalid SP split boundary: {self.split_at!r}')
+        elif (self.split_at is None) == bool(self.need_split):
+            raise ValueError('need_split and split_at specify conflicting SP split behavior')
+        if self.split_at != 'auto':
+            object.__setattr__(self, 'need_split', self.split_at is not None)
+        if self.side == 'output' and (self.pad_at is not None or self.split_at is not None):
+            raise ValueError('Output keys cannot request input padding or splitting')
+        if self.need_unpack is None:
+            # Preserve the historical packed restoration behavior unless a
+            # built-in rule explicitly opts out.
+            object.__setattr__(self, 'need_unpack', self.side == 'output' or self.need_gather or not self.need_split)
+
+    def sequence_dim(self, ndim):
+        if not -ndim <= self.split_dim < ndim:
+            raise ValueError(f'Invalid sequence axis {self.split_dim} for SP key {self.name!r}')
+        dim = self.split_dim % ndim
+        if dim == 0:
+            raise ValueError(f'SP key {self.name!r} cannot split or gather the batch axis')
+        return dim
+
+
+BUILTIN_INPUT_KEYS = (
+    DataKeySpec('input_ids', split_dim=1, pad_value='tokenizer', stage='forward'),
+    DataKeySpec('inputs_embeds', split_dim=1, stage='forward'),
+    DataKeySpec('position_ids', pad_value=-1, stage='forward'),
+    DataKeySpec('attention_mask', split_dim=1, stage='forward', need_pad=True, need_split=False),
+    DataKeySpec('labels', pad_value=-100, need_gather=True, need_unpack=True),
+    # These are not needed by the model forward. Keep them global through SP;
+    # packed sequence restoration is a separate action from sharding/gathering.
+    DataKeySpec('completion_mask', need_pad=False, need_split=False, need_unpack=True),
+    DataKeySpec('loss_scale', need_pad=False, need_split=False, need_unpack=True),
+)
+
+# These are processor/model control arguments, not sequence data contracts.
+# Keep them out of DataKeySpec and preserve them through the Transformers input
+# conversion path explicitly instead.
+PROCESSOR_PASSTHROUGH_KEYS = ('cache_position', 'padding_free')
+
+BUILTIN_OUTPUT_KEYS = (
+    DataKeySpec('logps', split_dim=1, side='output', stage='output',
+                need_split=False, need_gather=True, autograd=True, need_unpack=True),
+    DataKeySpec('entropies', split_dim=1, side='output', stage='output',
+                need_split=False, need_gather=True, autograd=True, need_unpack=True),
+    # Real logits are opt-in; fused CE explicitly gathers this hidden-state slot.
+    DataKeySpec('logits', split_dim=1, side='output', stage='output',
+                need_split=False, autograd=True, need_unpack=True),
+)
+
+
+@dataclass
+class SPKeyValue:
+    """One tensor value paired with its immutable SP processing rule."""
+
+    spec: DataKeySpec
+    value: Any
+
+    @property
+    def name(self):
+        return self.spec.name
+
+
+@dataclass
+class SPKeyCollection:
+    """Batch-local common and registered SP tensors passed through one executor."""
+
+    common: Dict[str, SPKeyValue] = field(default_factory=dict)
+    extra: Dict[str, SPKeyValue] = field(default_factory=dict)
+
+    def add(self, item: SPKeyValue, *, common: bool = False):
+        if not isinstance(item, SPKeyValue):
+            raise TypeError('SPKeyCollection accepts SPKeyValue items')
+        if item.name in self.common or item.name in self.extra:
+            raise ValueError(f'Duplicate SP key in collection: {item.name!r}')
+        target = self.common if common else self.extra
+        target[item.name] = item
+
+    def values(self) -> Iterator[SPKeyValue]:
+        yield from self.common.values()
+        yield from self.extra.values()
+
+    def get(self, name: str) -> Optional[SPKeyValue]:
+        return self.common.get(name, self.extra.get(name))
+
+    def get_value(self, name: str, default=None):
+        """Return a key's payload without exposing its SP rule wrapper."""
+        item = self.get(name)
+        return default if item is None else item.value
+
+    @classmethod
+    def from_mapping(cls, values: Mapping[str, Any], specs: Iterable[DataKeySpec], common_names=()):
+        collection = cls()
+        common_names = set(common_names)
+        for spec in specs:
+            if spec.name in values:
+                collection.add(SPKeyValue(spec, values[spec.name]), common=spec.name in common_names)
+        return collection
+
+
+class InputKeyRegistry:
+    """Own declarations locally; reject accidental replacement of existing rules.
+
+    Auxiliary keys are processed at their declared boundary (usually
+    preprocessing). Custom model-forward keys remain global until the forward
+    hook, so every key is sharded at most once.
+    """
+
+    _FORWARD_KEYS = frozenset({'input_ids', 'inputs_embeds', 'position_ids', 'attention_mask'})
+    _SPECIAL_KEYS = frozenset(PROCESSOR_PASSTHROUGH_KEYS)
+
+    def __init__(self):
+        self._specs = {spec.name: spec for spec in BUILTIN_INPUT_KEYS}
+
+    def register(self, spec: DataKeySpec):
+        if not isinstance(spec, DataKeySpec):
+            raise TypeError('Expected a DataKeySpec')
+        if spec.side != 'input' or spec.stage not in ('auxiliary', 'forward'):
+            raise ValueError('Register only auxiliary or model-forward input keys here')
+        if spec.name in self._FORWARD_KEYS:
+            raise ValueError(f'{spec.name!r} is owned by the model forward hook')
+        if spec.name in self._SPECIAL_KEYS:
+            raise ValueError(f'{spec.name!r} is a special processor/model argument, not an SP data key')
+        existing = self._specs.get(spec.name)
+        if existing is not None and existing != spec:
+            raise ValueError(f'Conflicting SP input rule for {spec.name!r}')
+        self._specs[spec.name] = spec
+
+    def __iter__(self):
+        return iter(self._specs.values())
+
+    def get(self, name):
+        return self._specs.get(name)
+
+
+class OutputKeyRegistry:
+    def __init__(self):
+        self._specs = {spec.name: spec for spec in BUILTIN_OUTPUT_KEYS}
+
+    def register(self, spec):
+        if not isinstance(spec, DataKeySpec):
+            raise TypeError('Expected a DataKeySpec')
+        if spec.side != 'output':
+            raise ValueError('Expected an output key rule')
+        existing = self._specs.get(spec.name)
+        if existing is not None and existing != spec:
+            raise ValueError(f'Conflicting SP output rule for {spec.name!r}')
+        self._specs[spec.name] = spec
+
+    def __iter__(self):
+        return iter(self._specs.values())
+
+    def get(self, name):
+        return self._specs.get(name)
diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py
index 8f166a571..92980e2c0 100644
--- a/src/twinkle/model/transformers/transformers.py
+++ b/src/twinkle/model/transformers/transformers.py
@@ -557,6 +557,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
         labels: torch.Tensor = inputs.pop('labels', None)
         loss_scale = inputs.pop('loss_scale', None)
         channel = inputs.pop('channel', None)
+        sp_auxiliary_inputs = processor.pop_sp_auxiliary_inputs(inputs, sp_strategy=self.sp_strategy)
         optimizer_config.accumulate_metrics(True)
 
         # Routing replay: respects router_replay_action regardless of caller
@@ -575,6 +576,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
         recorded_routing = rr_cleanup()
 
         inputs['labels'] = labels
+        inputs.update(sp_auxiliary_inputs)
         if loss_scale is not None:
             inputs['loss_scale'] = loss_scale
         if channel is not None:
@@ -594,8 +596,13 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
         outputs['past_key_values'] = None
         if not (return_logits or loss_require_logits):
             outputs['logits'] = None
-        inputs, outputs = processor.postprocess_tensor_sp(inputs, outputs, sp_strategy=self.sp_strategy, task=task)
-        inputs, outputs = processor.unpack_packed_sequences(inputs, outputs, task=task)
+        inputs, outputs = processor.postprocess_tensor_sp(
+            inputs,
+            outputs,
+            sp_strategy=self.sp_strategy,
+            task=task,
+        )
+        inputs, outputs = processor.unpack_packed_sequences(inputs, outputs, task=task, sp_strategy=self.sp_strategy)
         optimizer_config.train_status.inputs = inputs
         optimizer_config.train_status.outputs = outputs
         optimizer_config.train_status.forward_kwargs = kwargs
@@ -654,6 +661,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
             labels = inputs.pop('labels', None)
             loss_scale = inputs.pop('loss_scale', None)
             channel = inputs.pop('channel', None)
+            sp_auxiliary_inputs = processor.pop_sp_auxiliary_inputs(inputs, sp_strategy=self.sp_strategy)
             optimizer_config.accumulate_metrics(False)
             unwrapped_model = self.strategy.unwrap_model(self.model)
 
@@ -675,6 +683,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
             recorded_routing = rr_cleanup()
 
             inputs['labels'] = labels
+            inputs.update(sp_auxiliary_inputs)
             if loss_scale is not None:
                 inputs['loss_scale'] = loss_scale
             if channel is not None:
@@ -694,8 +703,14 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
             outputs['past_key_values'] = None
             if not (return_logits or loss_require_logits):
                 outputs['logits'] = None
-            inputs, outputs = processor.postprocess_tensor_sp(inputs, outputs, sp_strategy=self.sp_strategy, task=task)
-            inputs, outputs = processor.unpack_packed_sequences(inputs, outputs, task=task)
+            inputs, outputs = processor.postprocess_tensor_sp(
+                inputs,
+                outputs,
+                sp_strategy=self.sp_strategy,
+                task=task,
+            )
+            inputs, outputs = processor.unpack_packed_sequences(inputs, outputs, task=task,
+                                                                sp_strategy=self.sp_strategy)
             optimizer_config.eval_status.inputs = inputs
             optimizer_config.eval_status.outputs = outputs
             optimizer_config.eval_status.forward_kwargs = kwargs
diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py
index 054810fd1..13a834442 100644
--- a/src/twinkle/processor/base.py
+++ b/src/twinkle/processor/base.py
@@ -100,6 +100,9 @@ def prepare_outputs(self, inputs: List[InputFeature], **kwargs) -> Union[List[In
             return inputs
 
     def prepare_inputs(self, inputs: Union[List[InputFeature], InputFeature], **kwargs) -> List[InputFeature]:
+        from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS
+
+        registered_keys = {spec.name for spec in getattr(kwargs.get('sp_strategy'), 'input_key_specs', BUILTIN_INPUT_KEYS)}
 
         def to_tensor(_input):
             import torch
@@ -108,7 +111,9 @@ def to_tensor(_input):
                 # Ray/pyarrow can return numpy or list scalars; normalize to tensors.
                 # After distributed/datasets.map, labels/completion_mask may become numpy arrays or lists,
                 # so tensor ops like labels != ignore_index or .to(device) would fail without this.
-                if isinstance(value, np.ndarray):
+                if key in registered_keys and isinstance(value, (list, np.ndarray)):
+                    value = torch.as_tensor(value)
+                elif isinstance(value, np.ndarray):
                     value = torch.from_numpy(value)
                 elif isinstance(value, list) and len(value) > 0 and isinstance(value[0], (int, float, np.number)):
                     value = torch.tensor(value)
@@ -643,6 +648,7 @@ def unpack_packed_sequences(
         inputs: Dict[str, Any],
         outputs: Optional[Dict[str, Any]] = None,
         task: str = 'causal_lm',
+        sp_strategy=None,
     ) -> tuple[Dict[str, Any], Optional[Dict[str, Any]]]:
         """Unpack packed (padding_free) sequences into per-sequence batch format.
 
@@ -667,29 +673,33 @@ def unpack_packed_sequences(
 
         from copy import copy
 
-        # Collect input/output keys to unpack: (key, pad_value). ``channel`` stays
-        # sample-level metadata, while ``loss_scale`` follows the token layout.
-        input_keys = [('labels', -100)]
-        if inputs.get('loss_scale') is not None:
-            input_keys.append(('loss_scale', 0))
-        output_keys = []
-        for key, pad_val in [('logps', 0), ('entropies', 0), ('logits', 0)]:
-            if outputs and outputs.get(key) is not None:
-                output_keys.append((key, pad_val))
-
-        all_tensors = [inputs[k] for k, _ in input_keys] + [outputs[k] for k, _ in output_keys]
-        all_pads = [p for _, p in input_keys] + [p for _, p in output_keys]
+        from twinkle.model.transformers.strategy.sequence_parallel.data_keys import (
+            BUILTIN_INPUT_KEYS, BUILTIN_OUTPUT_KEYS,
+        )
+
+        input_keys = [spec for spec in getattr(sp_strategy, 'input_key_specs', BUILTIN_INPUT_KEYS)
+                      if spec.stage == 'auxiliary' and spec.need_unpack and inputs.get(spec.name) is not None]
+        output_keys = [spec for spec in getattr(sp_strategy, 'output_key_specs', BUILTIN_OUTPUT_KEYS)
+                       if spec.need_unpack and outputs and outputs.get(spec.name) is not None]
+        all_specs = input_keys + output_keys
+        all_tensors = [inputs[spec.name] for spec in input_keys] + [outputs[spec.name] for spec in output_keys]
+        # The unpack primitive expects [1, T, ...]; restore the declared axis
+        # afterwards so tensors like [1, H, T] retain their public layout.
+        axes = [spec.sequence_dim(tensor.dim()) for spec, tensor in zip(all_specs, all_tensors)]
+        all_tensors = [tensor.movedim(axis, 1) for tensor, axis in zip(all_tensors, axes)]
+        all_pads = [spec.pad_value for spec in all_specs]
         unpacked = self._unpack_by_position_ids(position_ids, *all_tensors, padding_values=all_pads)
+        unpacked = [tensor.movedim(1, axis) for tensor, axis in zip(unpacked, axes)]
 
         inputs = copy(inputs)
-        for i, (key, _) in enumerate(input_keys):
-            inputs[key] = unpacked[i]
+        for i, spec in enumerate(input_keys):
+            inputs[spec.name] = unpacked[i]
 
         if output_keys:
             outputs = copy(outputs)
             offset = len(input_keys)
-            for i, (key, _) in enumerate(output_keys):
-                outputs[key] = unpacked[offset + i]
+            for i, spec in enumerate(output_keys):
+                outputs[spec.name] = unpacked[offset + i]
 
         return inputs, outputs
 
@@ -704,6 +714,10 @@ def unpack_inputs(self, inputs: List[Dict[str, Any]], task: str = 'causal_lm') -
     @staticmethod
     def to_transformers_dict(inputs: List[InputFeature], **kwargs) -> List[InputFeature]:
         import torch
+        from twinkle.model.transformers.strategy.sequence_parallel.data_keys import PROCESSOR_PASSTHROUGH_KEYS
+
+        sp_strategy = kwargs.get('sp_strategy')
+        registered_keys = [spec.name for spec in getattr(sp_strategy, 'input_key_specs', ())]
         results = []
         for _input in inputs:
             output = {}
@@ -724,7 +738,7 @@ def to_transformers_dict(inputs: List[InputFeature], **kwargs) -> List[InputFeat
                 'max_length_k',
                 'packed_seq_params',
                 'routed_experts',
-            ] + list(InputProcessor.VLM_CONCAT_FIELDS)
+            ] + list(InputProcessor.VLM_CONCAT_FIELDS) + registered_keys + list(PROCESSOR_PASSTHROUGH_KEYS)
             for key in list(_input.keys()):
                 if key not in _keys:
                     continue
@@ -736,12 +750,56 @@ def to_transformers_dict(inputs: List[InputFeature], **kwargs) -> List[InputFeat
             results.append(InputFeature(**output))
         return results
 
-    def _collate_macro_batch(self, inputs: List[InputFeature]) -> InputFeature:
+    @staticmethod
+    def pop_sp_auxiliary_inputs(inputs, sp_strategy=None):
+        """Keep loss-only token fields out of Hugging Face model.forward."""
+        from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS
+
+        specs = getattr(sp_strategy, 'input_key_specs', BUILTIN_INPUT_KEYS)
+        return {spec.name: inputs.pop(spec.name) for spec in specs
+                if spec.stage == 'auxiliary' and spec.name not in ('labels', 'loss_scale') and spec.name in inputs}
+
+    def _collate_sequence_key(self, values, spec, padding_free):
+        samples = []
+        seq_dim = None
+        for value in values:
+            if not torch.is_tensor(value):
+                raise TypeError(f'Registered SP input {spec.name!r} must be a tensor before collation')
+            if value.dim() == 1:
+                value = value.unsqueeze(0)
+            dim = spec.sequence_dim(value.dim())
+            if dim == 0 or value.shape[0] != 1:
+                raise ValueError(f'SP input {spec.name!r} requires one sample with a leading batch dimension')
+            if seq_dim is not None and seq_dim != dim - 1:
+                raise ValueError(f'Inconsistent sequence axis for SP input {spec.name!r}')
+            seq_dim = dim - 1
+            samples.append(value.squeeze(0))
+        if padding_free:
+            return torch.cat(samples, dim=seq_dim).unsqueeze(0)
+        length = max(value.shape[seq_dim] for value in samples)
+        padded = []
+        for value in samples:
+            shape = list(value.shape)
+            shape[seq_dim] = length
+            result = value.new_full(shape, spec.pad_value)
+            slices = [slice(None)] * value.dim()
+            start = length - value.shape[seq_dim] if self.padding_side == 'left' else 0
+            slices[seq_dim] = slice(start, start + value.shape[seq_dim])
+            result[tuple(slices)] = value
+            padded.append(result)
+        return torch.stack(padded)
+
+    def _collate_macro_batch(self, inputs: List[InputFeature], sp_strategy=None) -> InputFeature:
         import torch
+        from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS
+
+        builtin_names = {spec.name for spec in BUILTIN_INPUT_KEYS}
+        key_specs = {spec.name: spec for spec in getattr(sp_strategy, 'input_key_specs', BUILTIN_INPUT_KEYS)
+                     if spec.stage == 'auxiliary' or (spec.stage == 'forward' and spec.name not in builtin_names)}
 
         for _input in inputs:
             for key in list(_input.keys()):
-                if isinstance(_input[key], torch.Tensor):
+                if isinstance(_input[key], torch.Tensor) and key not in key_specs:
                     _input[key] = _input[key].squeeze()
 
         vlm_fields = {k: [] for k in self.VLM_CONCAT_FIELDS}
@@ -768,7 +826,10 @@ def is_mm_position_ids(position_ids):
         if padding_free:
             for key in text_keys:
                 values = [item[key] for item in text_inputs]
-                if key == 'attention_mask':
+                if key in key_specs:
+                    result[key] = self._collate_sequence_key(values, key_specs[key], True)
+                    continue
+                elif key == 'attention_mask':
                     # attention_mask is not needed
                     continue
                 if key == 'position_ids' and is_mm_position_ids(values[0]):
@@ -783,7 +844,9 @@ def is_mm_position_ids(position_ids):
         else:
             for key in text_keys:
                 values = [item[key] for item in text_inputs]
-                if self.framework == 'megatron' and key == 'attention_mask':
+                if key in key_specs:
+                    result[key] = self._collate_sequence_key(values, key_specs[key], False)
+                elif self.framework == 'megatron' and key == 'attention_mask':
                     result[key] = self._create_4d_attention_mask(values)
                 elif key == 'position_ids' and is_mm_position_ids(values[0]):
                     result[key] = InputProcessor._pad_sequence(values, self.padding_map[key], self.padding_side)
@@ -826,7 +889,7 @@ def collate_fn(self,
             return inputs
         if micro_batch_size is None:
             # normal collate
-            outputs = self._collate_macro_batch(inputs)
+            outputs = self._collate_macro_batch(inputs, sp_strategy=kwargs.get('sp_strategy'))
             for key in outputs:
                 if key in self.VLM_CONCAT_FIELDS:
                     outputs[key] = torch.cat(outputs[key], dim=0)
@@ -837,7 +900,8 @@ def collate_fn(self,
             assert len(inputs) >= micro_batch_size
             outputs = []
             for i in range(0, len(inputs), micro_batch_size):
-                _output = self._collate_macro_batch(inputs[i:i + micro_batch_size])
+                _output = self._collate_macro_batch(inputs[i:i + micro_batch_size],
+                                                    sp_strategy=kwargs.get('sp_strategy'))
                 for key in _output:
                     if key in self.VLM_CONCAT_FIELDS:
                         _output[key] = torch.cat(_output[key], dim=0)
@@ -845,7 +909,7 @@ def collate_fn(self,
             return outputs
         else:
             # each macro batch shares the same length
-            res = self._collate_macro_batch(inputs)
+            res = self._collate_macro_batch(inputs, sp_strategy=kwargs.get('sp_strategy'))
             keys = list(res.keys())
             outputs = []
             for i in range(0, len(inputs), micro_batch_size):
diff --git a/tests/transformers/test_sp_data_keys.py b/tests/transformers/test_sp_data_keys.py
new file mode 100644
index 000000000..dc52295e7
--- /dev/null
+++ b/tests/transformers/test_sp_data_keys.py
@@ -0,0 +1,306 @@
+"""Input-rule checks without model training or accelerator allocation."""
+import pytest
+import torch
+
+from twinkle.model.transformers.strategy.sequence_parallel import DataKeySpec, SequenceParallelStrategy
+from twinkle.model.transformers.strategy.sequence_parallel.data_keys import SPKeyCollection, SPKeyValue
+
+
+def _strategy(rank=0):
+    strategy = SequenceParallelStrategy(sp_config={'ulysses_size': 2})
+    ctx = strategy._ctx
+    ctx.world_size = ctx.sp_world_size = 2
+    ctx.rp_world_size = 1
+    ctx._sp_rank = rank
+    return strategy
+
+
+@pytest.mark.parametrize('rank', [0, 1])
+@pytest.mark.parametrize('length', [8, 9])
+def test_builtin_labels_match_existing_pad_split(rank, length):
+    strategy = _strategy(rank)
+    positions = torch.arange(length).expand(2, -1)
+    labels = torch.arange(2 * length).reshape(2, length)
+    spec = strategy._ctx.input_key_registry.get('labels')
+    padded_labels = strategy._ctx._pad_key(labels, spec, positions)
+    padded_positions = strategy._ctx.pad(positions, padding_value=-1, position_ids=positions)
+    old = strategy._ctx._split_key(padded_labels, spec, padded_positions)
+    inputs = {'labels': labels.clone(), 'position_ids': positions}
+    actual = strategy.preprocess_inputs(inputs)['labels']
+    torch.testing.assert_close(actual, old, rtol=0, atol=0)
+    torch.testing.assert_close(labels, torch.arange(2 * length).reshape(2, length))
+
+
+@pytest.mark.parametrize('rank', [0, 1])
+def test_custom_key_padding_dimension_roll_and_model_isolation(rank):
+    strategy = _strategy(rank)
+    other = _strategy(rank)
+    spec = DataKeySpec('token_weights', split_dim=1, pad_value=-2, need_roll=True)
+    strategy.register_input_key(spec)
+    strategy.register_input_key(spec)  # Identical declarations are safe.
+    values = torch.arange(10).reshape(1, 5, 2)
+    positions = torch.arange(5).unsqueeze(0)
+    prepared = strategy.preprocess_inputs({'token_weights': values, 'position_ids': positions})
+    expected = torch.cat((values, torch.full((1, 1, 2), -2)), dim=1).roll(-1, dims=1)
+    torch.testing.assert_close(prepared['token_weights'], expected[:, rank * 3:(rank + 1) * 3])
+    untouched = other.preprocess_inputs({'token_weights': values, 'position_ids': positions})
+    assert untouched['token_weights'] is values
+
+
+@pytest.mark.parametrize('rank', [0, 1])
+def test_key_collection_executes_only_at_declared_boundary(rank):
+    strategy = _strategy(rank)
+    ctx = strategy._ctx
+    ctx.tokenizer = type('Tokenizer', (), {'pad_token_id': 0})()
+    token_type_spec = DataKeySpec('token_type_ids', split_dim=1, pad_value=-1, stage='forward')
+    strategy.register_input_key(token_type_spec)
+    values = {
+        'input_ids': torch.arange(5).expand(2, -1),
+        'position_ids': torch.arange(5).expand(2, -1),
+        'attention_mask': torch.ones((2, 5), dtype=torch.long),
+        'token_type_ids': torch.arange(5).expand(2, -1),
+    }
+    specs = tuple(spec for spec in ctx.input_key_registry if spec.stage == 'forward')
+    keys = SPKeyCollection.from_mapping(values, specs, common_names=('input_ids', 'position_ids', 'attention_mask'))
+    ctx.pad_and_split_inputs(keys, boundary='forward', real_position_ids=values['position_ids'])
+    assert keys.get('input_ids').value.shape == (2, 3)
+    assert keys.get('position_ids').value.shape == (2, 3)
+    assert keys.get('token_type_ids').value.shape == (2, 3)
+    # attention_mask follows its own rule: pad to the global frame but do not shard.
+    torch.testing.assert_close(keys.get('attention_mask').value, torch.tensor([[1, 1, 1, 1, 1, 0]]).expand(2, -1))
+
+
+def test_key_collection_rejects_duplicate_names_across_common_and_extra():
+    spec = DataKeySpec('token_type_ids', stage='forward')
+    keys = SPKeyCollection()
+    keys.add(SPKeyValue(spec, torch.ones(1, 2)), common=True)
+    with pytest.raises(ValueError, match='Duplicate SP key'):
+        keys.add(SPKeyValue(spec, torch.zeros(1, 2)))
+
+
+@pytest.mark.parametrize('rank', [0, 1])
+def test_prepare_keeps_loss_scale_and_completion_mask_global(rank):
+    strategy = _strategy(rank)
+    positions = torch.arange(5).unsqueeze(0)
+    loss_scale = torch.arange(1, 6, dtype=torch.float).unsqueeze(0)
+    completion_mask = torch.tensor([[0, 0, 1, 1, 1]])
+    inputs = {
+        'position_ids': positions,
+        'labels': torch.arange(5).unsqueeze(0),
+        'loss_scale': loss_scale.clone(),
+        'completion_mask': completion_mask.clone(),
+    }
+    prepared = strategy.preprocess_inputs(inputs)
+    assert prepared['labels'].shape == (1, 3)
+    torch.testing.assert_close(prepared['loss_scale'], loss_scale)
+    torch.testing.assert_close(prepared['completion_mask'], completion_mask)
+
+
+def test_invalid_rules_and_non_tensor_payload_fail_clearly():
+    strategy = _strategy()
+    with pytest.raises(ValueError, match='Invalid SP key stage'):
+        DataKeySpec('cache_position', stage='passthrough')
+    with pytest.raises(ValueError, match='Conflicting'):
+        strategy.register_input_key(DataKeySpec('labels', pad_value=0))
+    with pytest.raises(ValueError, match='forward hook'):
+        strategy.register_input_key(DataKeySpec('input_ids'))
+    with pytest.raises(ValueError, match='not an SP data key'):
+        strategy.register_input_key(DataKeySpec('padding_free'))
+    with pytest.raises(TypeError, match='must be a tensor'):
+        strategy.preprocess_inputs({'labels': [1, 2]})
+
+
+def test_disabled_sp_and_absent_keys_are_unchanged():
+    strategy = SequenceParallelStrategy(sp_config={'enabled': False, 'ulysses_size': 2})
+    inputs = {'labels': torch.tensor([[1, 2]])}
+    assert strategy.preprocess_inputs(inputs) is inputs
+    assert _strategy().preprocess_inputs({'labels': None}) == {'labels': None}
+
+
+@pytest.mark.parametrize('sp_rank', [0, 1])
+@pytest.mark.parametrize('rp_rank', [0, 1])
+def test_packed_ring_labels_preserve_existing_layout(sp_rank, rp_rank):
+    strategy = _strategy(sp_rank)
+    ctx = strategy._ctx
+    ctx.world_size = 4
+    ctx.rp_world_size = 2
+    ctx._rp_rank = rp_rank
+    positions = torch.tensor([[0, 1, 2, 3, 4, 0, 1, 2]])
+    labels = torch.arange(8).unsqueeze(0)
+    spec = ctx.input_key_registry.get('labels')
+    padded_labels = ctx._pad_key(labels, spec, positions)
+    padded_positions = ctx.pad(positions, padding_value=-1, position_ids=positions)
+    old = ctx._split_key(padded_labels, spec, padded_positions)
+    new = strategy.preprocess_inputs({'labels': labels, 'position_ids': positions})['labels']
+    torch.testing.assert_close(new, old, rtol=0, atol=0)
+
+
+def test_packed_unpack_uses_explicit_per_key_unpack_rule():
+    from twinkle.processor.base import InputProcessor
+
+    processor = InputProcessor(padding_free=True)
+    positions = torch.tensor([[0, 1, 2, 0, 1]])
+    labels = torch.tensor([[10, 11, 12, 20, 21]])
+    completion_mask = torch.tensor([[0, 1, 1, 0, 1]])
+    loss_scale = torch.tensor([[1., 2., 3., 4., 5.]])
+    inputs = {
+        'position_ids': positions,
+        'labels': labels,
+        'completion_mask': completion_mask,
+        'loss_scale': loss_scale,
+    }
+    unpacked, _ = processor.unpack_packed_sequences(inputs, outputs=None)
+    torch.testing.assert_close(unpacked['labels'], torch.tensor([[10, 11, 12], [20, 21, -100]]))
+    torch.testing.assert_close(unpacked['completion_mask'], torch.tensor([[0, 1, 1], [0, 1, 0]]))
+    torch.testing.assert_close(unpacked['loss_scale'], torch.tensor([[1., 2., 3.], [4., 5., 0.]]))
+
+
+@pytest.mark.parametrize('rank', [0, 1])
+def test_processor_preserves_registered_keys_before_sp_split(monkeypatch, rank):
+    from twinkle import Platform
+    from twinkle.processor.base import InputProcessor
+
+    monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu'))
+    strategy = _strategy(rank)
+    strategy.register_input_key(DataKeySpec('token_weights', pad_value=-2))
+    processor = InputProcessor()
+    values = torch.arange(5).unsqueeze(0)
+    sample = {'input_ids': torch.ones(1, 5, dtype=torch.long),
+              'position_ids': torch.arange(5).unsqueeze(0),
+              'labels': torch.ones(1, 5, dtype=torch.long),
+              'token_weights': values, 'unregistered': values}
+    result = processor([sample], sp_strategy=strategy)
+    expected = torch.tensor([[0, 1, 2, 3, 4, -2]])[:, rank * 3:(rank + 1) * 3]
+    torch.testing.assert_close(result['token_weights'], expected)
+    assert 'unregistered' not in result
+    other = _strategy(rank)
+    assert 'token_weights' not in processor.to_transformers_dict([sample], sp_strategy=other)[0]
+
+
+@pytest.mark.parametrize('rank', [0, 1])
+def test_builtin_completion_mask_and_loss_scale_stay_global_during_prepare(rank):
+    strategy = _strategy(rank)
+    payload = {'position_ids': torch.arange(5).unsqueeze(0),
+               'completion_mask': torch.tensor([[0, 0, 1, 1, 1]]),
+               'loss_scale': torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5]])}
+    prepared = strategy.preprocess_inputs(payload)
+    torch.testing.assert_close(prepared['completion_mask'], torch.tensor([[0, 0, 1, 1, 1]]))
+    torch.testing.assert_close(prepared['loss_scale'], torch.tensor([[0.1, 0.2, 0.3, 0.4, 0.5]]))
+
+
+def test_registered_multidimensional_key_collation_and_loss_boundary(monkeypatch):
+    from twinkle import Platform
+    from twinkle.processor import InputProcessor
+
+    monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu'))
+    strategy = _strategy()
+    strategy.register_input_key(DataKeySpec('token_features', pad_value=-2, need_gather=True))
+    processor = InputProcessor()
+    samples = []
+    for length in (3, 5):
+        samples.append({'input_ids': list(range(1, length + 1)),
+                        'position_ids': list(range(length)), 'labels': list(range(1, length + 1)),
+                        'token_features': [ [[1] * length, [2] * length] ],
+                        'completion_mask': [1] * length})
+    result = processor(samples, sp_strategy=strategy)
+    expected = torch.tensor([[[1, 1, 1], [2, 2, 2]], [[1, 1, 1], [2, 2, 2]]])
+    torch.testing.assert_close(result['token_features'], expected)
+    auxiliary = processor.pop_sp_auxiliary_inputs(result, strategy)
+    assert set(auxiliary) == {'completion_mask', 'token_features'}
+    assert 'token_features' not in result
+
+
+def test_model_specific_forward_key_is_split_only_in_model_hook(monkeypatch):
+    from types import SimpleNamespace
+    from twinkle import Platform
+    from twinkle.processor import InputProcessor
+
+    monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu'))
+    strategy = _strategy(1)
+    strategy._ctx.tokenizer = SimpleNamespace(pad_token_id=0)
+    strategy.register_input_key(DataKeySpec('token_type_ids', stage='forward', owner='model'))
+
+    class Model(torch.nn.Module):
+        def forward(self, token_type_ids, **kwargs):
+            return token_type_ids
+
+    model = Model()
+    strategy._ctx._prepare_forward_hook(model)
+    ids = torch.arange(5).unsqueeze(0)
+    prepared = InputProcessor()([{'input_ids': ids, 'position_ids': ids,
+                                  'token_type_ids': ids}], sp_strategy=strategy)
+    assert prepared['token_type_ids'].shape == (1, 5)
+    torch.testing.assert_close(model(**prepared), torch.tensor([[3, 4, 0]]))
+
+
+def test_declared_collection_uses_sequence_axis_and_collects_without_logps(monkeypatch):
+    import twinkle.model.transformers.strategy.sequence_parallel as sp
+
+    strategy = _strategy()
+    strategy._ctx.extra_kwargs['position_ids'] = torch.arange(5).unsqueeze(0)
+    strategy.register_input_key(DataKeySpec('token_features', need_gather=True, pad_value=-2))
+    strategy.register_output_key(DataKeySpec('scores', split_dim=2, side='output', stage='output',
+                                            need_split=False, need_gather=True, autograd=True))
+    monkeypatch.setattr(strategy._ctx, 'gather', lambda value, dim, position_ids: torch.cat((value, value), dim=dim))
+    monkeypatch.setattr(sp.GatherLoss, 'apply', lambda ctx, value, labels, dim, pos: (ctx.gather(value, dim, pos), None))
+    local = torch.tensor([[[1, 2, 3], [4, 5, 6]]])
+    inputs, outputs = strategy.gather_loss_tensors({'token_features': local}, {'scores': local.float()})
+    expected = torch.tensor([[[1, 2, 3, 1, 2], [4, 5, 6, 4, 5]]])
+    torch.testing.assert_close(inputs['token_features'], expected)
+    torch.testing.assert_close(outputs['scores'], expected.float())
+
+
+def test_packing_unpacks_registered_keys_on_declared_axis():
+    from twinkle.processor import InputProcessor
+
+    strategy = _strategy()
+    strategy.register_input_key(DataKeySpec('token_features', pad_value=-2, need_gather=True))
+    inputs = {'position_ids': torch.tensor([[0, 1, 2, 0, 1]]), 'labels': torch.tensor([[1, 2, 3, 4, 5]]),
+              'completion_mask': torch.ones(1, 5), 'token_features': torch.tensor([[[1, 2, 3, 4, 5]]])}
+    unpacked, _ = InputProcessor().unpack_packed_sequences(inputs, sp_strategy=strategy)
+    assert unpacked['token_features'].shape == (2, 1, 3)
+    assert unpacked['token_features'].tolist() == [[[1, 2, 3]], [[4, 5, -2]]]
+    assert unpacked['completion_mask'].tolist() == [[1, 1, 1], [1, 1, 0]]
+
+
+def test_padding_free_collation_preserves_registered_nondefault_sequence_axis():
+    from twinkle.processor import InputProcessor
+
+    strategy = _strategy()
+    strategy.register_input_key(DataKeySpec('token_features', split_dim=-1, pad_value=-2))
+    processor = InputProcessor(padding_free=True)
+    samples = [
+        {'input_ids': torch.tensor([[1, 2, 3]]),
+         'position_ids': torch.tensor([[0, 1, 2]]),
+         'token_features': torch.tensor([[[10, 11, 12], [20, 21, 22]]])},
+        {'input_ids': torch.tensor([[4, 5, 6, 7, 8]]),
+         'position_ids': torch.tensor([[0, 1, 2, 3, 4]]),
+         'token_features': torch.tensor([[[30, 31, 32, 33, 34], [40, 41, 42, 43, 44]]])},
+    ]
+
+    result = processor._collate_macro_batch(samples, sp_strategy=strategy)
+
+    assert result['token_features'].shape == (1, 2, 8)
+    torch.testing.assert_close(
+        result['token_features'],
+        torch.tensor([[[10, 11, 12, 30, 31, 32, 33, 34],
+                       [20, 21, 22, 40, 41, 42, 43, 44]]]),
+    )
+
+
+def test_processor_preserves_special_arguments_outside_data_key_registry():
+    from twinkle.processor import InputProcessor
+
+    strategy = _strategy()
+    spec_names = {spec.name for spec in strategy.input_key_specs}
+    assert 'cache_position' not in spec_names
+    assert 'padding_free' not in spec_names
+
+    sample = {'input_ids': torch.tensor([[1, 2]]),
+              'cache_position': torch.tensor([0, 1]),
+              'padding_free': True}
+    converted_with_sp = InputProcessor.to_transformers_dict([sample], sp_strategy=strategy)[0]
+    converted_without_sp = InputProcessor.to_transformers_dict([sample])[0]
+    for converted in (converted_with_sp, converted_without_sp):
+        torch.testing.assert_close(converted['cache_position'], sample['cache_position'])
+        assert converted['padding_free'] is True
diff --git a/tests/transformers/test_sp_model_adapters.py b/tests/transformers/test_sp_model_adapters.py
new file mode 100644
index 000000000..86c46da84
--- /dev/null
+++ b/tests/transformers/test_sp_model_adapters.py
@@ -0,0 +1,99 @@
+"""Tests for model-family SP adapter registration and behavior."""
+import pytest
+import torch
+from types import SimpleNamespace
+
+import twinkle.model.transformers.strategy.sequence_parallel as sp
+from twinkle.model.transformers.strategy.sequence_parallel import adapters as sp_adapters
+
+
+@pytest.mark.parametrize(('model_type', 'expected'), [
+    ('qwen3_vl', {'qwen3_vl'}),
+    ('qwen3_5', {'qwen3_5_linear_attention'}),
+    ('qwen3_5_moe', {'qwen3_5_linear_attention', 'moe_aux'}),
+    ('llama', set()),
+])
+def test_builtin_model_adapters_are_selected_by_model_family(model_type, expected):
+    model = torch.nn.Linear(2, 2)
+    model.config = SimpleNamespace(model_type=model_type)
+    assert {adapter.name for adapter in sp_adapters.adapters_for_model(model)} == expected
+
+
+def test_model_adapter_registration_and_declared_key_specs(monkeypatch):
+    class CustomAdapter(sp_adapters.ModelSPAdapter):
+        name = 'test_custom_sp_adapter'
+
+        def matches(self, model):
+            return getattr(getattr(model, 'config', None), 'model_type', None) == 'custom_sp_test'
+
+        def input_key_specs(self):
+            return (sp.DataKeySpec('custom_features', pad_value=-3),)
+
+    monkeypatch.setattr(sp_adapters, '_ADAPTERS', list(sp_adapters._ADAPTERS))
+    adapter = sp_adapters.register_model_sp_adapter(CustomAdapter)
+    assert adapter is sp_adapters.register_model_sp_adapter(CustomAdapter())
+    model = torch.nn.Linear(2, 2)
+    model.config = SimpleNamespace(model_type='custom_sp_test')
+    context = sp.SequenceParallel()
+    context._configure_model_adapters(model, model)
+    assert context.model_adapters == (adapter,)
+    assert context.input_key_registry.get('custom_features') == sp.DataKeySpec('custom_features', pad_value=-3)
+
+
+def test_model_adapters_declare_mrope_and_keep_qwen35_ring_limit(monkeypatch):
+    vl_adapter = sp_adapters.Qwen3VLSPAdapter()
+    assert vl_adapter.position_ids_mode == 'mrope'
+    ctx = sp.SequenceParallel()
+    ctx.position_ids_mode = vl_adapter.position_ids_mode
+    real_positions = torch.arange(4).unsqueeze(0)
+    assert ctx._extract_real_position_ids(real_positions) is real_positions
+    mrope_positions = real_positions.unsqueeze(0).expand(3, -1, -1)
+    torch.testing.assert_close(ctx._extract_real_position_ids(mrope_positions), real_positions)
+    ctx.rp_world_size = 2
+    with pytest.raises(NotImplementedError, match='rp_world_size > 1'):
+        sp_adapters.Qwen35LinearAttentionSPAdapter().configure(ctx, torch.nn.Linear(2, 2))
+
+
+@pytest.mark.parametrize(('adapter', 'patch_name'), [
+    (sp_adapters.Qwen3VLSPAdapter(), 'qwen3_vl_deepstack'),
+    (sp_adapters.Qwen35LinearAttentionSPAdapter(), 'qwen35_gdn'),
+])
+def test_builtin_adapters_install_model_specific_patches(monkeypatch, adapter, patch_name):
+    calls = []
+    monkeypatch.setattr(sp_adapters, 'apply_patch', lambda model, patch, **kwargs: calls.append(
+        (model, patch, kwargs['sequence_parallel'])))
+    model = torch.nn.Linear(2, 2)
+    ctx = sp.SequenceParallel()
+    ctx.rp_world_size = 1
+    adapter.configure(ctx, model)
+    assert len(calls) == 1
+    assert calls[0][0] is (None if patch_name == 'qwen35_gdn' else model)
+    assert calls[0][2] is ctx
+
+
+def test_moe_adapter_registers_router_logits_gather_hook(monkeypatch):
+    hooks = []
+
+    class MoEConfig:
+        model_type = 'example_moe'
+
+    class Model:
+        config = MoEConfig()
+
+        def register_forward_hook(self, hook, with_kwargs=False):
+            hooks.append((hook, with_kwargs))
+
+    class Output(dict):
+        router_logits = (torch.arange(6, dtype=torch.float32).reshape(2, 3),)
+
+    ctx = SimpleNamespace(real_position_ids=torch.tensor([[0, 1, 2]]))
+    monkeypatch.setattr(sp_adapters.GatherLoss, 'apply',
+                        lambda context, value, labels, dim, positions: (torch.cat((value, value), dim), None))
+    model = Model()
+    adapter = sp_adapters.MoEAuxSPAdapter()
+    assert adapter.matches(model)
+    adapter.configure(ctx, model)
+    assert hooks and hooks[0][1]
+    output = hooks[0][0](model, (), {'attention_mask': torch.ones(1, 3)}, Output())
+    assert len(output['router_logits']) == 1
+    assert output['router_logits'][0].shape == (3, 3)

From 2f7e78df6f5e80bbc18e6ac9f6320be4a55c46b6 Mon Sep 17 00:00:00 2001
From: clc <clc0301@cmbchina.com>
Date: Thu, 15 Oct 2026 16:28:13 +0800
Subject: [PATCH 3/8] sp: add loss-declared gather contracts

---
 src/twinkle/loss/base.py                      |   8 +
 src/twinkle/loss/value.py                     |   3 +
 .../strategy/sequence_parallel/__init__.py    |  37 ++-
 .../model/transformers/transformers.py        |  23 +-
 src/twinkle/processor/base.py                 |   4 +-
 src/twinkle/utils/nccl_safe.py                |   4 +
 tests/transformers/test_sp_custom_loss_cpu.py | 209 +++++++++++++
 tests/transformers/test_sp_loss_contract.py   | 286 ++++++++++++++++++
 8 files changed, 562 insertions(+), 12 deletions(-)
 create mode 100644 tests/transformers/test_sp_custom_loss_cpu.py
 create mode 100644 tests/transformers/test_sp_loss_contract.py

diff --git a/src/twinkle/loss/base.py b/src/twinkle/loss/base.py
index 5fd046ae7..2066daa02 100644
--- a/src/twinkle/loss/base.py
+++ b/src/twinkle/loss/base.py
@@ -7,6 +7,14 @@ class Loss:
     require_logits = False
     require_entropy = False
     require_logps = True
+    # Sequence-aligned outputs this loss needs in full-sequence layout.
+    # Keys must have an output DataKeySpec registered by SP or a model adapter.
+    sp_gather_keys = ()
+    # A loss may declare additional sequence-aligned input/output contracts.
+    # Keep the core loss API independent of the Transformers backend by
+    # importing DataKeySpec only in concrete loss implementations.
+    sp_input_key_specs = ()
+    sp_output_key_specs = ()
 
     def __call__(self, inputs: InputFeature, outputs: ModelOutput, **kwargs) -> LossOutput:
         ...
diff --git a/src/twinkle/loss/value.py b/src/twinkle/loss/value.py
index adfc38edb..eaabe1256 100644
--- a/src/twinkle/loss/value.py
+++ b/src/twinkle/loss/value.py
@@ -46,6 +46,9 @@ class PPOValueLoss(Loss):
     # Reads the per-token value head output outputs['logits'] ([B, T]); never per-token logps.
     require_logits = True
     require_logps = False
+    # PPO regresses per-token value predictions stored in outputs['logits'].
+    # Gather the local sequence shards before aligning them with response labels.
+    sp_gather_keys = ('logits',)
 
     def __init__(self, cliprange_value: float = 0.2, vf_coef: float = 0.1, ignore_index: int = -100, **kwargs):
         super().__init__()
diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
index 109261dd2..703ad8531 100644
--- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
@@ -1017,6 +1017,22 @@ def register_output_key(self, spec: DataKeySpec) -> None:
     def output_key_specs(self) -> Tuple[DataKeySpec, ...]:
         return tuple(self._ctx.output_key_registry)
 
+    def register_loss_key_specs(self, loss) -> None:
+        """Register sequence input/output contracts declared by a loss."""
+        for attr, register, expected_side in (
+                ('sp_input_key_specs', self.register_input_key, 'input'),
+                ('sp_output_key_specs', self.register_output_key, 'output')):
+            specs = getattr(loss, attr, ()) or ()
+            if isinstance(specs, DataKeySpec):
+                specs = (specs,)
+            for spec in specs:
+                if not isinstance(spec, DataKeySpec):
+                    raise TypeError(f'{type(loss).__name__}.{attr} must contain DataKeySpec values')
+                if spec.side != expected_side:
+                    raise ValueError(f'{type(loss).__name__}.{attr} contains {spec.side!r} key {spec.name!r}; '
+                                     f'expected {expected_side!r}')
+                register(spec)
+
     @property
     def model_sp_adapters(self) -> Tuple[ModelSPAdapter, ...]:
         """The model-family adapters selected during context preparation."""
@@ -1080,6 +1096,7 @@ def gather_loss_tensors(
         self,
         inputs: Dict[str, Any],
         outputs: Dict[str, Any],
+        sp_gather_keys=(),
     ) -> Tuple[Dict[str, Any], Dict[str, Any]]:
         if inputs is None or outputs is None:
             return inputs, outputs
@@ -1098,9 +1115,27 @@ def gather_loss_tensors(
         # selects the registered logits rule without gathering real vocabulary
         # logits in the ordinary logps path.
         fused_hidden = outputs.get('logps') is None and 'lm_head' in outputs and inputs.get('labels') is not None
+        if isinstance(sp_gather_keys, str):
+            sp_gather_keys = (sp_gather_keys,)
+        else:
+            sp_gather_keys = tuple(sp_gather_keys)
+        if any(not isinstance(key, str) or not key for key in sp_gather_keys):
+            raise TypeError('sp_gather_keys must contain non-empty output key names')
+        required_keys = set(sp_gather_keys)
+        output_specs = {spec.name: spec for spec in self.output_key_specs}
+        unknown_keys = required_keys - output_specs.keys()
+        if unknown_keys:
+            raise KeyError(
+                f'Loss requested SP gather for unregistered output key(s) {sorted(unknown_keys)}. '
+                'Register a sequence-aligned DataKeySpec with SequenceParallelStrategy.register_output_key().')
+        missing_keys = {key for key in required_keys if outputs.get(key) is None}
+        if missing_keys:
+            raise KeyError(f'Loss declared SP gather key(s) {sorted(missing_keys)}, '
+                           'but the model output is missing them.')
         for spec in self.output_key_specs:
             value = outputs.get(spec.name)
-            if value is not None and (spec.need_gather or (spec.name == 'logits' and fused_hidden)):
+            if value is not None and (spec.need_gather or spec.name in required_keys
+                                      or (spec.name == 'logits' and fused_hidden)):
                 outputs[spec.name] = self._gather_key(value, spec, real_position_ids)
         return inputs, outputs
 
diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py
index 92980e2c0..aae5cce00 100644
--- a/src/twinkle/model/transformers/transformers.py
+++ b/src/twinkle/model/transformers/transformers.py
@@ -365,16 +365,17 @@ def _decide_strategy(self,
     def _ensure_sp_strategy(self) -> None:
         if not getattr(self, '_enable_sp', False):
             return
-        if self.sp_strategy is not None:
-            return
-        from .strategy.sequence_parallel import SequenceParallelStrategy
+        if self.sp_strategy is None:
+            from .strategy.sequence_parallel import SequenceParallelStrategy
 
-        self.sp_strategy = SequenceParallelStrategy(
-            self.device_mesh,
-            {},
-            model=self.model,
-            tokenizer_id=self.tokenizer_id,
-        )
+            self.sp_strategy = SequenceParallelStrategy(
+                self.device_mesh,
+                {},
+                model=self.model,
+                tokenizer_id=self.tokenizer_id,
+            )
+        for optimizer_group in self.optimizer_group.values():
+            self.sp_strategy.register_loss_key_specs(optimizer_group.loss_instance)
 
     def _get_default_group(self):
         """Get the only group, else return the default one"""
@@ -601,6 +602,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
             outputs,
             sp_strategy=self.sp_strategy,
             task=task,
+            sp_gather_keys=getattr(loss_instance, 'sp_gather_keys', ()),
         )
         inputs, outputs = processor.unpack_packed_sequences(inputs, outputs, task=task, sp_strategy=self.sp_strategy)
         optimizer_config.train_status.inputs = inputs
@@ -708,6 +710,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
                 outputs,
                 sp_strategy=self.sp_strategy,
                 task=task,
+                sp_gather_keys=getattr(loss_instance, 'sp_gather_keys', ()),
             )
             inputs, outputs = processor.unpack_packed_sequences(inputs, outputs, task=task,
                                                                 sp_strategy=self.sp_strategy)
@@ -1035,6 +1038,8 @@ def set_loss(self, loss_cls: Union[Loss, Type[Loss], str, Callable[[InputFeature
         adapter_name = kwargs.pop('adapter_name', self._get_default_group())
         optimizer_config = self.optimizer_group[adapter_name]
         optimizer_config.loss_instance = construct_class(loss_cls, Loss, twinkle.loss, **kwargs)
+        if self.sp_strategy is not None:
+            self.sp_strategy.register_loss_key_specs(optimizer_config.loss_instance)
 
     @remote_function()
     def set_optimizer(self, optimizer_cls: Union[Type[Optimizer], str, Optimizer], **kwargs):
diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py
index 13a834442..7ba90a524 100644
--- a/src/twinkle/processor/base.py
+++ b/src/twinkle/processor/base.py
@@ -149,7 +149,7 @@ def apply_transformers_sp(self, inputs: List[InputFeature], **kwargs) -> List[In
         return results
 
     def postprocess_tensor_sp(self, inputs: Dict[str, Any], outputs: Dict[str, Any],
-                              **kwargs) -> tuple[Dict[str, Any], Dict[str, Any]]:
+                              sp_gather_keys=(), **kwargs) -> tuple[Dict[str, Any], Dict[str, Any]]:
         """Adjust SP tensors after forward and before loss computation.
 
         Pipeline: SP gather → packed-sequence unpack.
@@ -169,7 +169,7 @@ def postprocess_tensor_sp(self, inputs: Dict[str, Any], outputs: Dict[str, Any],
         if task == 'generative_reranker':
             return self._postprocess_generative_reranker(inputs, outputs, sp_strategy=sp_strategy)
         if self.framework == 'transformers' and sp_strategy is not None:
-            return sp_strategy.gather_loss_tensors(inputs, outputs)
+            return sp_strategy.gather_loss_tensors(inputs, outputs, sp_gather_keys=sp_gather_keys)
         return inputs, outputs
 
     @staticmethod
diff --git a/src/twinkle/utils/nccl_safe.py b/src/twinkle/utils/nccl_safe.py
index b22b10137..7e85db436 100644
--- a/src/twinkle/utils/nccl_safe.py
+++ b/src/twinkle/utils/nccl_safe.py
@@ -79,6 +79,10 @@ def __init__(self, loss_instance):
         self.require_entropy = getattr(loss_instance, 'require_entropy', False)
         self.require_logits = getattr(loss_instance, 'require_logits', False)
         self.reduction = getattr(loss_instance, 'reduction', 'mean')
+        # Preserve the loss's SP data contract through the NCCL-safe wrapper.
+        self.sp_gather_keys = getattr(loss_instance, 'sp_gather_keys', ())
+        self.sp_input_key_specs = getattr(loss_instance, 'sp_input_key_specs', ())
+        self.sp_output_key_specs = getattr(loss_instance, 'sp_output_key_specs', ())
         self._nccl_safe_wrapped = True
 
     def __call__(self, inputs, outputs, **kwargs):
diff --git a/tests/transformers/test_sp_custom_loss_cpu.py b/tests/transformers/test_sp_custom_loss_cpu.py
new file mode 100644
index 000000000..aadcba115
--- /dev/null
+++ b/tests/transformers/test_sp_custom_loss_cpu.py
@@ -0,0 +1,209 @@
+"""End-to-end contract test for a user-defined sequence loss."""
+import importlib
+import os
+import socket
+from datetime import timedelta
+from types import SimpleNamespace
+
+import pytest
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from transformers import AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM
+
+from twinkle import DeviceMesh, Platform
+from twinkle.data_format import LossOutput
+from twinkle.loss import Loss
+from twinkle.model.transformers.transformers import TransformersModel
+from twinkle.model.transformers.strategy.sequence_parallel import DataKeySpec, SequenceParallelStrategy
+
+
+class TokenRegressionLoss(Loss):
+    """Example plugin loss with custom sequence input and output keys."""
+
+    require_logps = False
+    sp_input_key_specs = (DataKeySpec('token_targets', split_dim=1, pad_value=0, need_gather=True, owner='loss'),)
+    sp_output_key_specs = (DataKeySpec('token_scores', split_dim=1, side='output', stage='output',
+                                       need_split=False, need_gather=True, autograd=True, owner='loss'),)
+
+    def __call__(self, inputs, outputs, **kwargs):
+        scores = outputs['token_scores']
+        targets = inputs['token_targets'].to(dtype=scores.dtype)
+        loss = (scores - targets).square().mean()
+        return LossOutput(loss=loss, num_tokens=loss.new_tensor(targets.numel()))
+
+
+class _ToyTokenRegressor(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.config = SimpleNamespace(num_key_value_heads=2, num_attention_heads=2,
+                                      _attn_implementation='sdpa', model_type='toy_token_regressor')
+        self.scale = torch.nn.Parameter(torch.tensor(0.7))
+
+    def forward(self, input_ids, position_ids, attention_mask=None, **kwargs):
+        return {'token_scores': input_ids.to(self.scale.dtype) * self.scale}
+
+
+class _NoopTrainStrategy:
+    """Keep this CPU E2E focused on Twinkle's model/loss/optimizer flow."""
+
+    def capture_pre_ep_state_if_needed(self, model, *, enable_ep):
+        return None
+
+    def wrap_model(self, model, *args):
+        return (model, *args) if args else model
+
+
+def _worker(rank, port):
+    os.environ['RANK'] = str(rank)
+    os.environ['WORLD_SIZE'] = '2'
+    torch.set_num_threads(1)
+    from torch.utils.checkpoint import DefaultDeviceType
+    DefaultDeviceType.set_device_type('cpu')
+    dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank,
+                            world_size=2, timeout=timedelta(seconds=90))
+    try:
+        ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7], [2, 3, 4, 5, 6, 7, 8]])
+        positions = torch.arange(ids.shape[1]).expand_as(ids)
+        targets = ids.float() * 0.4 + 0.25
+        loss_fn = TokenRegressionLoss()
+
+        baseline = _ToyTokenRegressor()
+        full_outputs = baseline(input_ids=ids, position_ids=positions)
+        full_loss = loss_fn({'token_targets': targets}, full_outputs)['loss']
+        full_loss.backward()
+        expected_grad = baseline.scale.grad.clone()
+
+        model = _ToyTokenRegressor()
+        strategy = SequenceParallelStrategy(
+            DeviceMesh.from_sizes(dp_size=2, ulysses_size=2, device_type='cpu'), model=model)
+        strategy._get_tokenizer = lambda: SimpleNamespace(pad_token_id=0)
+        strategy.register_loss_key_specs(loss_fn)
+        assert strategy.initialize()
+
+        prepared = strategy.preprocess_inputs({'input_ids': ids, 'position_ids': positions,
+                                               'token_targets': targets})
+        local_targets = prepared.pop('token_targets')
+        local_outputs = model(input_ids=prepared['input_ids'], position_ids=prepared['position_ids'])
+        loss_inputs, loss_outputs = strategy.gather_loss_tensors({'token_targets': local_targets}, local_outputs)
+        sp_loss = loss_fn(loss_inputs, loss_outputs)['loss']
+        torch.testing.assert_close(sp_loss, full_loss, atol=1e-7, rtol=1e-6)
+        sp_loss.backward()
+        dist.all_reduce(model.scale.grad)
+        model.scale.grad.div_(2)
+        torch.testing.assert_close(model.scale.grad, expected_grad, atol=1e-7, rtol=1e-6)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.slow
+def test_custom_token_regression_loss_matches_full_sequence_loss_and_gradient():
+    with socket.socket() as sock:
+        sock.bind(('127.0.0.1', 0))
+        port = sock.getsockname()[1]
+    mp.spawn(_worker, args=(port,), nprocs=2, join=True)
+
+
+def _e2e_worker(rank, port):
+    os.environ['RANK'] = str(rank)
+    os.environ['WORLD_SIZE'] = '2'
+    torch.set_num_threads(1)
+    Platform.get_local_device = staticmethod(lambda: torch.device('cpu'))
+    from torch.utils.checkpoint import DefaultDeviceType
+    DefaultDeviceType.set_device_type('cpu')
+    dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank,
+                            world_size=2, timeout=timedelta(seconds=90))
+    try:
+        config = LlamaConfig(vocab_size=32, hidden_size=16, intermediate_size=32,
+                             num_hidden_layers=1, num_attention_heads=4, num_key_value_heads=4,
+                             attention_dropout=0.0, pad_token_id=0, use_cache=False)
+        config._attn_implementation = 'sdpa'
+        torch.manual_seed(1234)
+        baseline = LlamaForCausalLM(config).train()
+
+        mesh = DeviceMesh.from_sizes(dp_size=2, ulysses_size=2, device_type='cpu')
+        model_module = importlib.import_module('twinkle.model.transformers.transformers')
+        original_template = model_module.Template
+        model_module.Template = lambda *_args, **_kwargs: SimpleNamespace(is_mm=False)
+        try:
+            training_model = TransformersModel(
+                model_cls=AutoModelForCausalLM,
+                config=config,
+                device_mesh=mesh,
+                mixed_precision='no',
+            )
+        finally:
+            model_module.Template = original_template
+        training_model.strategy = _NoopTrainStrategy()
+        # This local CPU harness supplies worker mesh state explicitly.
+        training_model.device_mesh = mesh
+        training_model._enable_sp = True
+        model_forward = training_model.model.forward
+
+        def forward_with_token_scores(*args, **kwargs):
+            outputs = model_forward(*args, **kwargs)
+            outputs['token_scores'] = outputs.logits[..., 0]
+            return outputs
+
+        training_model.model.forward = forward_with_token_scores
+        training_model.model.load_state_dict(baseline.state_dict())
+        training_model.model.gradient_checkpointing_disable()
+        training_model.model.train()
+        strategy = SequenceParallelStrategy(mesh, model=training_model.model)
+        strategy._get_tokenizer = lambda: SimpleNamespace(pad_token_id=0)
+        training_model.sp_strategy = strategy
+
+        optimizer_group = training_model.optimizer_group[training_model._get_default_group()]
+        loss_fn = TokenRegressionLoss()
+        optimizer_group.loss_instance = loss_fn  # OptimizerGroup applies the NCCL-safe wrapper.
+        assert training_model._enable_sp
+        assert optimizer_group.loss_instance.sp_input_key_specs
+        training_model._ensure_sp_strategy()
+        assert 'token_targets' in {spec.name for spec in strategy.input_key_specs}
+
+        ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7], [2, 3, 4, 5, 6, 7, 8]])
+        positions = torch.arange(ids.shape[1]).expand_as(ids)
+        targets = ids.float() * 0.4 + 0.25
+        labels = torch.full_like(ids, -100)
+        batch = [
+            {'input_ids': ids[row], 'position_ids': positions[row], 'labels': labels[row],
+             'token_targets': targets[row]}
+            for row in range(ids.shape[0])
+        ]
+
+        # Baseline uses the same model weights and the same optimizer update,
+        # but computes the loss on the unsharded full sequence.
+        base_optimizer = torch.optim.SGD(baseline.parameters(), lr=1e-3)
+        baseline_output = baseline(input_ids=ids, position_ids=positions, use_cache=False)
+        baseline_output['token_scores'] = baseline_output.logits[..., 0]
+        baseline_loss = loss_fn({'token_targets': targets}, baseline_output)['loss']
+        baseline_loss.backward()
+        base_optimizer.step()
+
+        # End-to-end Twinkle path: processor -> SP setup/hook -> forward ->
+        # loss registration/collection -> calculate_loss -> backward -> step.
+        forward_outputs = training_model.forward(inputs=batch)
+        assert 'token_targets' in optimizer_group.train_status.inputs, optimizer_group.train_status.inputs.keys()
+        sp_loss = training_model.calculate_loss()
+        torch.testing.assert_close(torch.tensor(sp_loss), baseline_loss.detach(), atol=1e-6, rtol=1e-5)
+        training_model.backward()
+        optimizer = torch.optim.SGD(training_model.model.parameters(), lr=1e-3)
+        for parameter in training_model.model.parameters():
+            assert parameter.grad is not None
+            dist.all_reduce(parameter.grad)
+            parameter.grad.div_(2)
+        optimizer_group.optimizer = optimizer
+        training_model.step()
+        assert 'token_scores' in forward_outputs
+        for (name, actual), (_, expected) in zip(training_model.model.named_parameters(), baseline.named_parameters()):
+            torch.testing.assert_close(actual, expected, atol=2e-6, rtol=2e-5, msg=name)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.slow
+def test_custom_loss_through_transformers_training_pipeline_and_optimizer_step():
+    with socket.socket() as sock:
+        sock.bind(('127.0.0.1', 0))
+        port = sock.getsockname()[1]
+    mp.spawn(_e2e_worker, args=(port,), nprocs=2, join=True)
diff --git a/tests/transformers/test_sp_loss_contract.py b/tests/transformers/test_sp_loss_contract.py
new file mode 100644
index 000000000..ccc6f5db1
--- /dev/null
+++ b/tests/transformers/test_sp_loss_contract.py
@@ -0,0 +1,286 @@
+"""Loss-declared output gathering for Sequence Parallel."""
+import pytest
+import importlib
+import os
+import socket
+import torch
+import torch.distributed as dist
+import torch.multiprocessing as mp
+from datetime import timedelta
+from types import SimpleNamespace
+
+from twinkle.data_format import LossOutput
+from twinkle.loss import Loss, PPOValueLoss
+from twinkle.model.transformers.strategy.sequence_parallel import DataKeySpec, SequenceParallelStrategy
+
+
+def _strategy():
+    strategy = SequenceParallelStrategy(sp_config={'ulysses_size': 2})
+    ctx = strategy._ctx
+    ctx.world_size = ctx.sp_world_size = 2
+    ctx.rp_world_size = 1
+    ctx.extra_kwargs['position_ids'] = torch.arange(5).unsqueeze(0)
+    return strategy
+
+
+def _mock_gather(monkeypatch):
+    from twinkle.model.transformers.strategy.sequence_parallel import GatherLoss
+
+    monkeypatch.setattr(GatherLoss, 'apply', lambda ctx, value, labels, dim, positions:
+                        (torch.cat((value, value), dim=dim), None))
+
+
+def test_loss_declared_logits_are_gathered_even_when_not_default(monkeypatch):
+    strategy = _strategy()
+    _mock_gather(monkeypatch)
+    local_values = torch.tensor([[[1.0], [2.0], [3.0]]])
+
+    inputs, outputs = strategy.gather_loss_tensors({}, {'logits': local_values}, sp_gather_keys=('logits',))
+
+    assert inputs == {}
+    torch.testing.assert_close(outputs['logits'], torch.tensor([[[1.0], [2.0], [3.0], [1.0], [2.0]]]))
+
+
+def test_custom_loss_declares_and_consumes_a_registered_output_key(monkeypatch):
+    class ToyValueLoss(Loss):
+        require_logps = False
+        sp_gather_keys = ('values',)
+
+        def __call__(self, inputs, outputs, **kwargs):
+            values = outputs['values']
+            return LossOutput(loss=values.square().mean(), num_tokens=values.numel())
+
+    strategy = _strategy()
+    strategy.register_output_key(
+        DataKeySpec('values', split_dim=1, side='output', stage='output', need_split=False, autograd=True))
+    _mock_gather(monkeypatch)
+    _, outputs = strategy.gather_loss_tensors({}, {'values': torch.tensor([[1.0, 2.0, 3.0]])},
+                                             sp_gather_keys=ToyValueLoss.sp_gather_keys)
+
+    assert outputs['values'].shape == (1, 5)
+    result = ToyValueLoss()({}, outputs)
+    torch.testing.assert_close(result['loss'], torch.tensor((1.0 + 4.0 + 9.0 + 1.0 + 4.0) / 5))
+
+
+def test_loss_can_declare_custom_sequence_inputs_and_outputs(monkeypatch):
+    class TokenAlignmentLoss(Loss):
+        require_logps = False
+        sp_input_key_specs = (DataKeySpec('token_targets', split_dim=1, pad_value=-1,
+                                          need_gather=True, autograd=True, owner='loss'),)
+        sp_output_key_specs = (DataKeySpec('token_scores', split_dim=1, side='output', stage='output',
+                                           need_split=False, need_gather=True, autograd=True, owner='loss'),)
+
+        def __call__(self, inputs, outputs, **kwargs):
+            return LossOutput(loss=(outputs['token_scores'] - inputs['token_targets']).square().mean(),
+                              num_tokens=inputs['token_targets'].numel())
+
+    strategy = _strategy()
+    strategy.register_loss_key_specs(TokenAlignmentLoss())
+    _mock_gather(monkeypatch)
+    inputs, outputs = strategy.gather_loss_tensors(
+        {'token_targets': torch.tensor([[1.0, 2.0, 3.0]])},
+        {'token_scores': torch.tensor([[1.0, 2.0, 3.0]])})
+
+    assert strategy._ctx.input_key_registry.get('token_targets') is not None
+    assert strategy._ctx.output_key_registry.get('token_scores') is not None
+    torch.testing.assert_close(inputs['token_targets'], torch.tensor([[1.0, 2.0, 3.0, 1.0, 2.0]]))
+    torch.testing.assert_close(outputs['token_scores'], torch.tensor([[1.0, 2.0, 3.0, 1.0, 2.0]]))
+    torch.testing.assert_close(TokenAlignmentLoss()(inputs, outputs)['loss'], torch.tensor(0.0))
+
+
+def test_declared_but_unregistered_or_missing_output_fails_loudly():
+    strategy = _strategy()
+    with pytest.raises(KeyError, match='unregistered output key'):
+        strategy.gather_loss_tensors({}, {'values': torch.ones(1, 3)}, sp_gather_keys=('values',))
+    with pytest.raises(KeyError, match='model output is missing'):
+        strategy.gather_loss_tensors({}, {}, sp_gather_keys=('logits',))
+
+
+def test_ppo_value_loss_declares_per_token_value_collection():
+    assert PPOValueLoss.sp_gather_keys == ('logits',)
+
+
+class _ToyValueModel(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.config = SimpleNamespace(num_attention_heads=2, _attn_implementation='sdpa', model_type='toy_value')
+        self.scale = torch.nn.Parameter(torch.tensor(0.7))
+
+    def forward(self, input_ids, position_ids, **kwargs):
+        return {'logits': input_ids.to(self.scale.dtype) * self.scale}
+
+
+def _ppo_value_worker(rank, port):
+    os.environ['RANK'] = str(rank)
+    os.environ['WORLD_SIZE'] = '2'
+    torch.set_num_threads(1)
+    from torch.utils.checkpoint import DefaultDeviceType
+    DefaultDeviceType.set_device_type('cpu')
+    dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank,
+                            world_size=2, timeout=timedelta(seconds=90))
+    try:
+        from twinkle import DeviceMesh
+        from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelStrategy
+        from twinkle.processor import InputProcessor
+
+        ids = torch.tensor([[1, 2, 3, 4, 5, 6]])
+        positions = torch.arange(6).unsqueeze(0)
+        labels = torch.tensor([[-100, -100, 3, 4, 5, 6]])
+        returns = torch.tensor([[0.5, 1.0, 1.5, 2.0]])
+        old_values = torch.tensor([[0.2, 0.4, 0.6, 0.8]])
+        loss_fn = PPOValueLoss(cliprange_value=0.2, vf_coef=0.5)
+
+        baseline = _ToyValueModel()
+        base_logits = baseline(ids, positions)['logits']
+        base_loss = loss_fn({'labels': labels}, {'logits': base_logits},
+                            returns=returns, old_values=old_values)['loss']
+        base_loss.backward()
+        base_grad = baseline.scale.grad.clone()
+
+        model = _ToyValueModel()
+        strategy = SequenceParallelStrategy(
+            DeviceMesh.from_sizes(dp_size=2, ulysses_size=2, device_type='cpu'),
+            sp_config={'ulysses_size': 2},
+            model=model,
+        )
+        strategy._get_tokenizer = lambda: SimpleNamespace(pad_token_id=0)
+        assert strategy.initialize()
+        prepared = strategy.preprocess_inputs({'input_ids': ids, 'position_ids': positions, 'labels': labels})
+        local_labels = prepared.pop('labels')
+        local_outputs = model(**{key: prepared[key] for key in ('input_ids', 'position_ids')})
+        loss_inputs, loss_outputs = InputProcessor().postprocess_tensor_sp(
+            {'labels': local_labels},
+            local_outputs,
+            sp_strategy=strategy,
+            sp_gather_keys=loss_fn.sp_gather_keys,
+        )
+        parallel_loss = loss_fn(loss_inputs, loss_outputs, returns=returns, old_values=old_values)['loss']
+        torch.testing.assert_close(parallel_loss, base_loss, atol=1e-7, rtol=1e-6)
+        parallel_loss.backward()
+        dist.all_reduce(model.scale.grad)
+        model.scale.grad.div_(2)
+        torch.testing.assert_close(model.scale.grad, base_grad, atol=1e-7, rtol=1e-6)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.slow
+def test_ppo_value_loss_sp_matches_full_sequence_loss_and_gradient():
+    with socket.socket() as sock:
+        sock.bind(('127.0.0.1', 0))
+        port = sock.getsockname()[1]
+    mp.spawn(_ppo_value_worker, args=(port,), nprocs=2, join=True)
+
+
+class _NoopTrainStrategy:
+    """Keep the CPU E2E focused on TransformersModel and SP loss handling."""
+
+    def capture_pre_ep_state_if_needed(self, model, *, enable_ep):
+        return None
+
+    def wrap_model(self, model, *args):
+        return (model, *args) if args else model
+
+
+def _ppo_training_pipeline_worker(rank, port):
+    os.environ['RANK'] = str(rank)
+    os.environ['WORLD_SIZE'] = '2'
+    torch.set_num_threads(1)
+    from torch.utils.checkpoint import DefaultDeviceType
+    DefaultDeviceType.set_device_type('cpu')
+    dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank,
+                            world_size=2, timeout=timedelta(seconds=90))
+    try:
+        from transformers import AutoModelForCausalLM, LlamaConfig, LlamaForCausalLM
+        from twinkle import DeviceMesh, Platform
+        from twinkle.model.transformers.transformers import TransformersModel
+        from twinkle.processor import InputProcessor
+
+        Platform.get_local_device = staticmethod(lambda: torch.device('cpu'))
+        config = LlamaConfig(vocab_size=32, hidden_size=16, intermediate_size=32,
+                             num_hidden_layers=1, num_attention_heads=4, num_key_value_heads=4,
+                             attention_dropout=0.0, pad_token_id=0, use_cache=False)
+        # Exercise a real Transformers attention implementation available in
+        # this environment; this E2E does not require flash-attn.
+        config._attn_implementation = 'sdpa'
+        torch.manual_seed(321)
+        baseline = LlamaForCausalLM(config).train()
+
+        mesh = DeviceMesh.from_sizes(dp_size=2, ulysses_size=2, device_type='cpu')
+        model_module = importlib.import_module('twinkle.model.transformers.transformers')
+        original_template = model_module.Template
+        model_module.Template = lambda *_args, **_kwargs: SimpleNamespace(is_mm=False)
+        try:
+            training_model = TransformersModel(
+                model_cls=AutoModelForCausalLM,
+                config=config,
+                device_mesh=mesh,
+                mixed_precision='no',
+            )
+        finally:
+            model_module.Template = original_template
+        training_model.strategy = _NoopTrainStrategy()
+        training_model.device_mesh = mesh
+        training_model._enable_sp = True
+        training_model.model.load_state_dict(baseline.state_dict())
+        training_model.model.gradient_checkpointing_disable()
+        training_model.model.train()
+
+        # The value head's [B, T] values travel through the standard `logits`
+        # output slot; retain a differentiable projection of real LM outputs.
+        model_forward = training_model.model.forward
+
+        def forward_per_token_values(*args, **kwargs):
+            outputs = model_forward(*args, **kwargs)
+            outputs['logits'] = outputs.logits[..., 0]
+            return outputs
+
+        training_model.model.forward = forward_per_token_values
+        training_model.sp_strategy = SequenceParallelStrategy(mesh, model=training_model.model)
+        training_model.sp_strategy._get_tokenizer = lambda: SimpleNamespace(pad_token_id=0)
+        training_model.set_loss(PPOValueLoss(cliprange_value=0.2, vf_coef=0.5))
+        training_model.set_processor(InputProcessor)
+        training_model._ensure_sp_strategy()
+
+        ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7, 8], [2, 3, 4, 5, 6, 7, 8, 9]])
+        positions = torch.arange(ids.shape[1]).expand_as(ids)
+        labels = torch.tensor([[-100, -100, -100, 4, 5, 6, 7, 8],
+                               [-100, -100, -100, 5, 6, 7, 8, 9]])
+        returns = torch.tensor([[0.5, 1.0, 1.5, 2.0, 2.5], [0.4, 0.8, 1.2, 1.6, 2.0]])
+        old_values = torch.tensor([[0.2, 0.4, 0.6, 0.8, 1.0], [0.1, 0.3, 0.5, 0.7, 0.9]])
+        batch = [
+            {'input_ids': ids[i], 'position_ids': positions[i], 'labels': labels[i]}
+            for i in range(ids.shape[0])
+        ]
+
+        loss_fn = PPOValueLoss(cliprange_value=0.2, vf_coef=0.5)
+        baseline_values = baseline(input_ids=ids, position_ids=positions, use_cache=False).logits[..., 0]
+        baseline_loss = loss_fn({'labels': labels}, {'logits': baseline_values},
+                                returns=returns, old_values=old_values)['loss']
+        baseline_loss.backward()
+        baseline_optimizer = torch.optim.SGD(baseline.parameters(), lr=1e-3)
+        baseline_optimizer.step()
+
+        training_model.forward_backward(inputs=batch, returns=returns, old_values=old_values)
+        optimizer_group = training_model.optimizer_group[training_model._get_default_group()]
+        assert optimizer_group.train_status.inputs['labels'].shape == labels.shape
+        assert optimizer_group.train_status.outputs['logits'].shape == labels.shape
+        optimizer_group.optimizer = torch.optim.SGD(training_model.model.parameters(), lr=1e-3)
+        for parameter in training_model.model.parameters():
+            assert parameter.grad is not None
+            dist.all_reduce(parameter.grad)
+            parameter.grad.div_(2)
+        training_model.step()
+
+        for (name, actual), (_, expected) in zip(training_model.model.named_parameters(), baseline.named_parameters()):
+            torch.testing.assert_close(actual, expected, atol=3e-6, rtol=3e-5, msg=name)
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.slow
+def test_ppo_value_loss_through_transformers_sdpa_sp_training_pipeline():
+    with socket.socket() as sock:
+        sock.bind(('127.0.0.1', 0))
+        port = sock.getsockname()[1]
+    mp.spawn(_ppo_training_pipeline_worker, args=(port,), nprocs=2, join=True)

From 9ad5b7a5792d6ad71bb695bfed9e4e095655680b Mon Sep 17 00:00:00 2001
From: clc <clc0301@cmbchina.com>
Date: Tue, 20 Oct 2026 21:17:41 +0800
Subject: [PATCH 4/8] sp: decouple processor key rules from strategy

---
 .../model/transformers/transformers.py        | 17 ++--
 src/twinkle/processor/base.py                 | 39 ++++++----
 tests/transformers/test_sp_data_keys.py       | 77 ++++++++++++++-----
 3 files changed, 92 insertions(+), 41 deletions(-)

diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py
index 87edb68e4..bd59ccef3 100644
--- a/src/twinkle/model/transformers/transformers.py
+++ b/src/twinkle/model/transformers/transformers.py
@@ -859,9 +859,11 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
                 raise ValueError('sampling replay does not support sequence or context parallelism')
         loss_require_values = getattr(loss_instance, 'require_values', False)
         assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding'
+        input_key_specs = getattr(self.sp_strategy, 'input_key_specs', None)
         inputs: Dict[str, Any] = processor(
             inputs,
             sp_strategy=self.sp_strategy,
+            input_key_specs=input_key_specs,
             model=self.model,
             hf_config=self.hf_config,
             enable_sp=getattr(self, '_enable_sp', False),
@@ -869,7 +871,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
         labels: torch.Tensor = inputs.pop('labels', None)
         loss_scale = inputs.pop('loss_scale', None)
         channel = inputs.pop('channel', None)
-        sp_auxiliary_inputs = processor.pop_sp_auxiliary_inputs(inputs, sp_strategy=self.sp_strategy)
+        sp_auxiliary_inputs = processor.pop_sp_auxiliary_inputs(inputs, input_key_specs=input_key_specs)
         # Not a model argument; the loss reads it back off `inputs` further down.
         completion_mask = inputs.pop('completion_mask', None)
         replay_metadata = replay_loss_mask = replay_masked_labels = None
@@ -942,7 +944,9 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
             task=task,
             sp_gather_keys=getattr(loss_instance, 'sp_gather_keys', ()),
         )
-        inputs, outputs = processor.unpack_packed_sequences(inputs, outputs, task=task, sp_strategy=self.sp_strategy)
+        inputs, outputs = processor.unpack_packed_sequences(
+            inputs, outputs, task=task, input_key_specs=input_key_specs,
+            output_key_specs=getattr(self.sp_strategy, 'output_key_specs', None))
         optimizer_config.train_status.inputs = inputs
         optimizer_config.train_status.outputs = outputs
         optimizer_config.train_status.forward_kwargs = kwargs
@@ -1004,9 +1008,11 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
                 if getattr(self, '_enable_sp', False) or cp_world_size > 1:
                     raise ValueError('sampling replay does not support sequence or context parallelism')
             loss_require_values = getattr(loss_instance, 'require_values', False)
+            input_key_specs = getattr(self.sp_strategy, 'input_key_specs', None)
             inputs: Dict[str, Any] = processor(
                 inputs,
                 sp_strategy=self.sp_strategy,
+                input_key_specs=input_key_specs,
                 model=self.model,
                 hf_config=self.hf_config,
                 enable_sp=getattr(self, '_enable_sp', False),
@@ -1014,7 +1020,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
             labels = inputs.pop('labels', None)
             loss_scale = inputs.pop('loss_scale', None)
             channel = inputs.pop('channel', None)
-            sp_auxiliary_inputs = processor.pop_sp_auxiliary_inputs(inputs, sp_strategy=self.sp_strategy)
+            sp_auxiliary_inputs = processor.pop_sp_auxiliary_inputs(inputs, input_key_specs=input_key_specs)
             # Not a model argument; the loss reads it back off `inputs` further down.
             completion_mask = inputs.pop('completion_mask', None)
             replay_metadata = replay_loss_mask = replay_masked_labels = None
@@ -1090,8 +1096,9 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
                 task=task,
                 sp_gather_keys=getattr(loss_instance, 'sp_gather_keys', ()),
             )
-            inputs, outputs = processor.unpack_packed_sequences(inputs, outputs, task=task,
-                                                                sp_strategy=self.sp_strategy)
+            inputs, outputs = processor.unpack_packed_sequences(
+                inputs, outputs, task=task, input_key_specs=input_key_specs,
+                output_key_specs=getattr(self.sp_strategy, 'output_key_specs', None))
             optimizer_config.eval_status.inputs = inputs
             optimizer_config.eval_status.outputs = outputs
             optimizer_config.eval_status.forward_kwargs = kwargs
diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py
index 89c17ae70..02b202042 100644
--- a/src/twinkle/processor/base.py
+++ b/src/twinkle/processor/base.py
@@ -101,10 +101,12 @@ def prepare_outputs(self, inputs: List[InputFeature], **kwargs) -> Union[List[In
                     _input['position_ids'] = _input['position_ids'][1:]
             return inputs
 
-    def prepare_inputs(self, inputs: Union[List[InputFeature], InputFeature], **kwargs) -> List[InputFeature]:
+    def prepare_inputs(self, inputs: Union[List[InputFeature], InputFeature], input_key_specs=None,
+                       **kwargs) -> List[InputFeature]:
         from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS
 
-        registered_keys = {spec.name for spec in getattr(kwargs.get('sp_strategy'), 'input_key_specs', BUILTIN_INPUT_KEYS)}
+        specs = BUILTIN_INPUT_KEYS if input_key_specs is None else input_key_specs
+        registered_keys = {spec.name for spec in specs}
 
         def to_tensor(_input):
             for key in list(_input.keys()):
@@ -664,7 +666,8 @@ def unpack_packed_sequences(
         inputs: Dict[str, Any],
         outputs: Optional[Dict[str, Any]] = None,
         task: str = 'causal_lm',
-        sp_strategy=None,
+        input_key_specs=None,
+        output_key_specs=None,
     ) -> tuple[Dict[str, Any], Optional[Dict[str, Any]]]:
         """Unpack packed (padding_free) sequences into per-sequence batch format.
 
@@ -674,6 +677,9 @@ def unpack_packed_sequences(
         ``[num_sequences, max_seq_len, ...]``. Keys that are ``None`` are silently
         skipped.
 
+        Input/output key specs describe the tensor layout and padding values;
+        this restoration step does not require a parallel strategy.
+
         For ``task='embedding'`` the outputs are already pooled to ``[n_seqs, H]``
         by ``postprocess_tensor_sp``, so this is a no-op. Same for
         ``task='generative_reranker'``, whose logits are already reduced to ``[n_seqs, 1]``.
@@ -694,9 +700,11 @@ def unpack_packed_sequences(
             BUILTIN_INPUT_KEYS, BUILTIN_OUTPUT_KEYS,
         )
 
-        input_keys = [spec for spec in getattr(sp_strategy, 'input_key_specs', BUILTIN_INPUT_KEYS)
+        input_specs = BUILTIN_INPUT_KEYS if input_key_specs is None else input_key_specs
+        output_specs = BUILTIN_OUTPUT_KEYS if output_key_specs is None else output_key_specs
+        input_keys = [spec for spec in input_specs
                       if spec.stage == 'auxiliary' and spec.need_unpack and inputs.get(spec.name) is not None]
-        output_keys = [spec for spec in getattr(sp_strategy, 'output_key_specs', BUILTIN_OUTPUT_KEYS)
+        output_keys = [spec for spec in output_specs
                        if spec.need_unpack and outputs and outputs.get(spec.name) is not None]
         all_specs = input_keys + output_keys
         all_tensors = [inputs[spec.name] for spec in input_keys] + [outputs[spec.name] for spec in output_keys]
@@ -729,12 +737,11 @@ def unpack_inputs(self, inputs: List[Dict[str, Any]], task: str = 'causal_lm') -
         return [self.unpack_packed_sequences(inp, task=task)[0] for inp in inputs]
 
     @staticmethod
-    def to_transformers_dict(inputs: List[InputFeature], **kwargs) -> List[InputFeature]:
+    def to_transformers_dict(inputs: List[InputFeature], input_key_specs=None, **kwargs) -> List[InputFeature]:
         import torch
         from twinkle.model.transformers.strategy.sequence_parallel.data_keys import PROCESSOR_PASSTHROUGH_KEYS
 
-        sp_strategy = kwargs.get('sp_strategy')
-        registered_keys = [spec.name for spec in getattr(sp_strategy, 'input_key_specs', ())]
+        registered_keys = [spec.name for spec in input_key_specs] if input_key_specs is not None else []
         results = []
         for _input in inputs:
             output = {}
@@ -771,11 +778,11 @@ def to_transformers_dict(inputs: List[InputFeature], **kwargs) -> List[InputFeat
         return results
 
     @staticmethod
-    def pop_sp_auxiliary_inputs(inputs, sp_strategy=None):
+    def pop_sp_auxiliary_inputs(inputs, *, input_key_specs=None):
         """Keep loss-only token fields out of Hugging Face model.forward."""
         from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS
 
-        specs = getattr(sp_strategy, 'input_key_specs', BUILTIN_INPUT_KEYS)
+        specs = BUILTIN_INPUT_KEYS if input_key_specs is None else input_key_specs
         return {spec.name: inputs.pop(spec.name) for spec in specs
                 if spec.stage == 'auxiliary' and spec.name not in ('labels', 'loss_scale') and spec.name in inputs}
 
@@ -849,12 +856,13 @@ def _fill_optional_sequence_fields(self, batch: List[InputFeature]) -> None:
     #: skipped for these fields when the batch is on the Qwen-audio contract.
     _QWEN_AUDIO_FIELDS = ('input_features', 'feature_attention_mask')
 
-    def _collate_macro_batch(self, inputs: List[InputFeature], sp_strategy=None) -> InputFeature:
+    def _collate_macro_batch(self, inputs: List[InputFeature], input_key_specs=None) -> InputFeature:
         import torch
         from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS
 
         builtin_names = {spec.name for spec in BUILTIN_INPUT_KEYS}
-        key_specs = {spec.name: spec for spec in getattr(sp_strategy, 'input_key_specs', BUILTIN_INPUT_KEYS)
+        specs = BUILTIN_INPUT_KEYS if input_key_specs is None else input_key_specs
+        key_specs = {spec.name: spec for spec in specs
                      if spec.stage == 'auxiliary' or (spec.stage == 'forward' and spec.name not in builtin_names)}
         # Two mutually-exclusive audio contracts share the `input_features` field name and are told
         # apart by the mask field that rides alongside it:
@@ -972,12 +980,13 @@ def collate_fn(self,
                    inputs: List[InputFeature],
                    micro_batch_size: Optional[int] = None,
                    variable_seq_lengths=False,
+                   input_key_specs=None,
                    **kwargs) -> List[InputFeature]:
         if len(inputs) == 1 and self.framework != 'megatron':
             return inputs
         if micro_batch_size is None:
             # normal collate
-            outputs = self._collate_macro_batch(inputs, sp_strategy=kwargs.get('sp_strategy'))
+            outputs = self._collate_macro_batch(inputs, input_key_specs=input_key_specs)
             for key in outputs:
                 if key in self.VLM_CONCAT_FIELDS:
                     outputs[key] = torch.cat(outputs[key], dim=0)
@@ -989,7 +998,7 @@ def collate_fn(self,
             outputs = []
             for i in range(0, len(inputs), micro_batch_size):
                 _output = self._collate_macro_batch(inputs[i:i + micro_batch_size],
-                                                    sp_strategy=kwargs.get('sp_strategy'))
+                                                    input_key_specs=input_key_specs)
                 for key in _output:
                     if key in self.VLM_CONCAT_FIELDS:
                         _output[key] = torch.cat(_output[key], dim=0)
@@ -997,7 +1006,7 @@ def collate_fn(self,
             return outputs
         else:
             # each macro batch shares the same length
-            res = self._collate_macro_batch(inputs, sp_strategy=kwargs.get('sp_strategy'))
+            res = self._collate_macro_batch(inputs, input_key_specs=input_key_specs)
             keys = list(res.keys())
             outputs = []
             for i in range(0, len(inputs), micro_batch_size):
diff --git a/tests/transformers/test_sp_data_keys.py b/tests/transformers/test_sp_data_keys.py
index dc52295e7..44efb2941 100644
--- a/tests/transformers/test_sp_data_keys.py
+++ b/tests/transformers/test_sp_data_keys.py
@@ -169,12 +169,12 @@ def test_processor_preserves_registered_keys_before_sp_split(monkeypatch, rank):
               'position_ids': torch.arange(5).unsqueeze(0),
               'labels': torch.ones(1, 5, dtype=torch.long),
               'token_weights': values, 'unregistered': values}
-    result = processor([sample], sp_strategy=strategy)
+    result = processor([sample], sp_strategy=strategy, input_key_specs=strategy.input_key_specs)
     expected = torch.tensor([[0, 1, 2, 3, 4, -2]])[:, rank * 3:(rank + 1) * 3]
     torch.testing.assert_close(result['token_weights'], expected)
     assert 'unregistered' not in result
     other = _strategy(rank)
-    assert 'token_weights' not in processor.to_transformers_dict([sample], sp_strategy=other)[0]
+    assert 'token_weights' not in processor.to_transformers_dict([sample], input_key_specs=other.input_key_specs)[0]
 
 
 @pytest.mark.parametrize('rank', [0, 1])
@@ -202,10 +202,10 @@ def test_registered_multidimensional_key_collation_and_loss_boundary(monkeypatch
                         'position_ids': list(range(length)), 'labels': list(range(1, length + 1)),
                         'token_features': [ [[1] * length, [2] * length] ],
                         'completion_mask': [1] * length})
-    result = processor(samples, sp_strategy=strategy)
+    result = processor(samples, sp_strategy=strategy, input_key_specs=strategy.input_key_specs)
     expected = torch.tensor([[[1, 1, 1], [2, 2, 2]], [[1, 1, 1], [2, 2, 2]]])
     torch.testing.assert_close(result['token_features'], expected)
-    auxiliary = processor.pop_sp_auxiliary_inputs(result, strategy)
+    auxiliary = processor.pop_sp_auxiliary_inputs(result, input_key_specs=strategy.input_key_specs)
     assert set(auxiliary) == {'completion_mask', 'token_features'}
     assert 'token_features' not in result
 
@@ -228,7 +228,8 @@ def forward(self, token_type_ids, **kwargs):
     strategy._ctx._prepare_forward_hook(model)
     ids = torch.arange(5).unsqueeze(0)
     prepared = InputProcessor()([{'input_ids': ids, 'position_ids': ids,
-                                  'token_type_ids': ids}], sp_strategy=strategy)
+                                  'token_type_ids': ids}], sp_strategy=strategy,
+                                input_key_specs=strategy.input_key_specs)
     assert prepared['token_type_ids'].shape == (1, 5)
     torch.testing.assert_close(model(**prepared), torch.tensor([[3, 4, 0]]))
 
@@ -252,23 +253,51 @@ def test_declared_collection_uses_sequence_axis_and_collects_without_logps(monke
 
 def test_packing_unpacks_registered_keys_on_declared_axis():
     from twinkle.processor import InputProcessor
+    from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS
 
-    strategy = _strategy()
-    strategy.register_input_key(DataKeySpec('token_features', pad_value=-2, need_gather=True))
+    input_specs = (*BUILTIN_INPUT_KEYS, DataKeySpec('token_features', pad_value=-2, need_gather=True))
+    output_specs = (DataKeySpec('scores', side='output', stage='output', need_split=False, pad_value=-3),)
     inputs = {'position_ids': torch.tensor([[0, 1, 2, 0, 1]]), 'labels': torch.tensor([[1, 2, 3, 4, 5]]),
               'completion_mask': torch.ones(1, 5), 'token_features': torch.tensor([[[1, 2, 3, 4, 5]]])}
-    unpacked, _ = InputProcessor().unpack_packed_sequences(inputs, sp_strategy=strategy)
+    outputs = {'scores': torch.tensor([[[10, 20, 30, 40, 50]]])}
+    unpacked, unpacked_outputs = InputProcessor().unpack_packed_sequences(
+        inputs, outputs, input_key_specs=input_specs, output_key_specs=output_specs)
     assert unpacked['token_features'].shape == (2, 1, 3)
     assert unpacked['token_features'].tolist() == [[[1, 2, 3]], [[4, 5, -2]]]
     assert unpacked['completion_mask'].tolist() == [[1, 1, 1], [1, 1, 0]]
+    assert unpacked_outputs['scores'].tolist() == [[[10, 20, 30]], [[40, 50, -3]]]
 
 
-def test_padding_free_collation_preserves_registered_nondefault_sequence_axis():
+@pytest.mark.parametrize('padding_free', [False, True])
+def test_processor_preserves_custom_keys_without_sp_strategy(monkeypatch, padding_free):
+    from twinkle import Platform
     from twinkle.processor import InputProcessor
+    from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS
 
-    strategy = _strategy()
-    strategy.register_input_key(DataKeySpec('token_features', split_dim=-1, pad_value=-2))
-    processor = InputProcessor(padding_free=True)
+    monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu'))
+    specs = (*BUILTIN_INPUT_KEYS, DataKeySpec('token_features', pad_value=-2))
+    samples = [
+        {'input_ids': list(range(length)), 'position_ids': list(range(length)),
+         'token_features': [[[1] * length, [2] * length]]}
+        for length in (3, 5)
+    ]
+    processor = InputProcessor(padding_free=padding_free)
+    result = processor(samples, input_key_specs=specs)
+    expected = (torch.tensor([[[1] * 8, [2] * 8]]) if padding_free else
+                torch.tensor([[[1, 1, 1, -2, -2], [2, 2, 2, -2, -2]], [[1] * 5, [2] * 5]]))
+    torch.testing.assert_close(result['token_features'], expected)
+    auxiliary = processor.pop_sp_auxiliary_inputs(result, input_key_specs=specs)
+    torch.testing.assert_close(auxiliary['token_features'], expected)
+
+
+@pytest.mark.parametrize('padding_free', [False, True])
+@pytest.mark.parametrize('micro_batch_size,variable_seq_lengths', [(None, False), (2, False), (1, False), (1, True)])
+def test_collation_preserves_declared_sequence_axis(padding_free, micro_batch_size, variable_seq_lengths):
+    from twinkle.processor import InputProcessor
+    from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS
+
+    specs = (*BUILTIN_INPUT_KEYS, DataKeySpec('token_features', split_dim=-1, pad_value=-2))
+    processor = InputProcessor(padding_free=padding_free)
     samples = [
         {'input_ids': torch.tensor([[1, 2, 3]]),
          'position_ids': torch.tensor([[0, 1, 2]]),
@@ -278,14 +307,20 @@ def test_padding_free_collation_preserves_registered_nondefault_sequence_axis():
          'token_features': torch.tensor([[[30, 31, 32, 33, 34], [40, 41, 42, 43, 44]]])},
     ]
 
-    result = processor._collate_macro_batch(samples, sp_strategy=strategy)
-
-    assert result['token_features'].shape == (1, 2, 8)
-    torch.testing.assert_close(
-        result['token_features'],
-        torch.tensor([[[10, 11, 12, 30, 31, 32, 33, 34],
-                       [20, 21, 22, 40, 41, 42, 43, 44]]]),
-    )
+    results = processor.collate_fn(samples, micro_batch_size=micro_batch_size,
+                                   variable_seq_lengths=variable_seq_lengths, input_key_specs=specs)
+    first = torch.tensor([[[10, 11, 12], [20, 21, 22]]])
+    first_padded = torch.tensor([[[10, 11, 12, -2, -2], [20, 21, 22, -2, -2]]])
+    second = torch.tensor([[[30, 31, 32, 33, 34], [40, 41, 42, 43, 44]]])
+    if micro_batch_size == 1:
+        expected = [first if padding_free or variable_seq_lengths else first_padded, second]
+    elif padding_free:
+        expected = [torch.cat((first, second), dim=-1)]
+    else:
+        expected = [torch.cat((first_padded, second), dim=0)]
+    assert len(results) == len(expected)
+    for result, value in zip(results, expected):
+        torch.testing.assert_close(result['token_features'], value)
 
 
 def test_processor_preserves_special_arguments_outside_data_key_registry():
@@ -299,7 +334,7 @@ def test_processor_preserves_special_arguments_outside_data_key_registry():
     sample = {'input_ids': torch.tensor([[1, 2]]),
               'cache_position': torch.tensor([0, 1]),
               'padding_free': True}
-    converted_with_sp = InputProcessor.to_transformers_dict([sample], sp_strategy=strategy)[0]
+    converted_with_sp = InputProcessor.to_transformers_dict([sample], input_key_specs=strategy.input_key_specs)[0]
     converted_without_sp = InputProcessor.to_transformers_dict([sample])[0]
     for converted in (converted_with_sp, converted_without_sp):
         torch.testing.assert_close(converted['cache_position'], sample['cache_position'])

From f787f853e47843e3fdd2375cce4583a4704801ad Mon Sep 17 00:00:00 2001
From: clc <clc0301@cmbchina.com>
Date: Wed, 21 Oct 2026 20:53:36 +0800
Subject: [PATCH 5/8] sp: move parallel processing out of input processor

Move SP input preparation, loss tensor gathering and last-token pooling
into the model/strategy layer. Keep processors responsible for batch
layouts and preserve sentence-level labels during SP preprocessing.

Cover pooling outputs and gradients for padded, packed and ring layouts
with CPU/Gloo regression tests.
---
 src/twinkle/loss/infonce.py                   |   2 +-
 .../strategy/sequence_parallel/__init__.py    |   7 +-
 .../strategy/sequence_parallel/pooling.py     | 102 ++++++++++
 .../model/transformers/task_outputs.py        |  26 +++
 .../model/transformers/transformers.py        |  28 ++-
 src/twinkle/patch/transformers_emb.py         |   4 +-
 src/twinkle/patch/transformers_reranker.py    |   2 +-
 src/twinkle/processor/base.py                 | 183 +-----------------
 tests/loss/test_channel_loss.py               |   3 -
 tests/transformers/test_sp_data_keys.py       |  10 +-
 tests/transformers/test_sp_loss_contract.py   |  92 ++++++++-
 11 files changed, 258 insertions(+), 201 deletions(-)
 create mode 100644 src/twinkle/model/transformers/strategy/sequence_parallel/pooling.py
 create mode 100644 src/twinkle/model/transformers/task_outputs.py

diff --git a/src/twinkle/loss/infonce.py b/src/twinkle/loss/infonce.py
index 3dffab93f..888c11f27 100644
--- a/src/twinkle/loss/infonce.py
+++ b/src/twinkle/loss/infonce.py
@@ -52,7 +52,7 @@ def _resolve_distance_metric(metric) -> Callable[[torch.Tensor, torch.Tensor], t
 
 
 def _extract_sentences(outputs) -> torch.Tensor:
-    """Return [B, D] sentence embeddings from postprocess_tensor_sp output.
+    """Return [B, D] sentence embeddings from model task postprocessing output.
 
     Prefers the canonical ``embeddings`` key (post-pooling); falls back to
     ``logits`` (legacy hook-side pooling) and applies CLS pooling for 3-D.
diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
index 98e2137cc..da6b101d1 100644
--- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
@@ -1025,6 +1025,11 @@ def preprocess_inputs(self, inputs: Dict[str, Any]) -> Dict[str, Any]:
             return inputs
         return self._ctx.prepare_inputs(inputs)
 
+    def pool_last_valid_token(self, features, inputs):
+        from .pooling import pool_last_valid_token
+
+        return pool_last_valid_token(features, inputs, sp_strategy=self)
+
     def register_input_key(self, spec: DataKeySpec) -> None:
         """Declare a token-aligned input on this model only.
 
@@ -1131,7 +1136,7 @@ def _gather_completion_mask(self, inputs: Dict[str, Any], real_position_ids) ->
         mask = inputs.get('completion_mask')
         if mask is None or not torch.is_tensor(mask) or mask.dim() < 2:
             return
-        gathered = sequence_parallel.gather(mask, dim=1, position_ids=real_position_ids)
+        gathered = self._ctx.gather(mask, dim=1, position_ids=real_position_ids)
         inputs['completion_mask'] = self._trim_gathered_sequence_padding(gathered, real_position_ids)
 
     def gather_loss_tensors(
diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/pooling.py b/src/twinkle/model/transformers/strategy/sequence_parallel/pooling.py
new file mode 100644
index 000000000..7889f6034
--- /dev/null
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/pooling.py
@@ -0,0 +1,102 @@
+"""Last-token pooling for full sequences and sequence-parallel shards."""
+from typing import Any
+
+import torch
+import torch.distributed as dist
+from torch.distributed.nn.functional import all_reduce
+
+
+def _packed_last_indices(position_ids: torch.Tensor, total_len: int) -> torch.Tensor:
+    """For padding-free batches: per-segment last-token indices into a [1, total] sequence."""
+    flat = position_ids.squeeze(0) if position_ids.dim() == 2 else position_ids
+    starts = (flat == 0).nonzero(as_tuple=False).squeeze(-1)
+    end_anchor = torch.tensor([total_len], device=flat.device, dtype=starts.dtype)
+    boundaries = torch.cat([starts, end_anchor])
+    return (boundaries[1:] - 1).long()
+
+
+def pool_last_valid_token(features: torch.Tensor, inputs: dict[str, Any],
+                          sp_strategy=None) -> tuple[torch.Tensor, bool]:
+    """Gather the last valid token per sequence from ``[B, T, C]`` -> ``[n_seqs, C]``.
+
+    Build a one-hot end-token mask in the un-padded global frame, route it
+    through the same pad+split as ``input_ids`` so it aligns with local
+    features, pool locally, then ``all_reduce`` only the ``[n_seqs, C]``
+    tensor across SP × RP. No feature gather; uniform across
+    DP / Ulysses / zigzag-ring / padding-free. Returns ``(pooled, is_packed)``.
+    """
+    sp_enabled = (
+        sp_strategy is not None and getattr(sp_strategy, 'enabled', False)
+        and getattr(sp_strategy, 'world_size', 1) > 1)
+
+    if sp_enabled:
+        ref_pos = sp_strategy.real_position_ids
+    elif inputs.get('position_ids') is not None:
+        ref_pos = inputs['position_ids']
+    else:
+        # Standard padded batch: HF omits position_ids (it derives them internally).
+        # Rebuild a position_ids-like reference (valid token -> its index, padding -> -1)
+        # so the last-valid-token pooling below, which keys off ``ref_pos >= 0`` for
+        # validity and a single ``== 0`` per row for packing detection, still works.
+        T = features.shape[1]
+        arange = torch.arange(T, device=features.device)
+        am = inputs.get('attention_mask')
+        if am is None:
+            ref_pos = arange.expand(features.shape[0], T)
+        else:
+            if am.dim() >= 3:
+                am = am.reshape(am.shape[0], -1)[:, :T]
+            ref_pos = torch.where(am.bool(), arange, arange.new_tensor(-1))
+    if ref_pos.dim() == 3:
+        ref_pos = ref_pos[0]
+    # Accept both flash-attn (cu_seq_lens_q) and Megatron (cu_seqlens_q) conventions.
+    cu_seq_lens_q = inputs.get('cu_seq_lens_q', inputs.get('cu_seqlens_q'))
+
+    is_packed = (
+        features.shape[0] == 1 and (cu_seq_lens_q is not None or int((ref_pos.reshape(-1) == 0).sum()) > 1))
+
+    device, dtype = features.device, features.dtype
+    T_real = ref_pos.shape[-1]
+
+    if is_packed:
+        if torch.is_tensor(cu_seq_lens_q) and cu_seq_lens_q.numel() >= 2:
+            end_idx = (cu_seq_lens_q[1:].long() - 1).to(device)
+        else:
+            end_idx = _packed_last_indices(ref_pos, T_real).to(device)
+        n_seqs = end_idx.shape[0]
+        batch_idx = torch.zeros_like(end_idx)
+        target_idx = torch.arange(n_seqs, device=device)
+        B_mask, n_targets = 1, n_seqs
+    else:
+        n_seqs = ref_pos.shape[0]
+        # clamp_min(0) prevents negative indexing when a row is fully padded (sum == 0).
+        positions = torch.arange(T_real, device=ref_pos.device).expand_as(ref_pos)
+        end_idx = positions.masked_fill(ref_pos < 0, -1).amax(-1).clamp_min(0)
+        batch_idx = torch.arange(n_seqs, device=device)
+        target_idx = torch.zeros(n_seqs, dtype=torch.long, device=device)
+        B_mask, n_targets = n_seqs, 1
+
+    mask = torch.zeros(B_mask, T_real, n_targets, dtype=dtype, device=device)
+    mask[batch_idx, end_idx, target_idx] = 1.0
+
+    if sp_enabled:
+        # Route mask through the same pad+split as input_ids to align with local features.
+        rp = sp_strategy.real_position_ids
+        rp_padded = sp_strategy.pad(rp, padding_value=-1, position_ids=rp, dim=-1)
+        mask = sp_strategy.pad(mask, padding_value=0, position_ids=rp, dim=1)
+        mask = sp_strategy.split(mask, dim=1, position_ids=rp_padded)
+
+    assert mask.shape[1] == features.shape[1], (
+        f'mask seq-dim {mask.shape[1]} != features seq-dim {features.shape[1]}; '
+        'sp_strategy.split contract broken or T_real mismatched.')
+
+    pooled = torch.einsum('bth,btn->bnh', features, mask)
+    pooled = pooled.squeeze(0) if is_packed else pooled.squeeze(1)
+
+    if sp_enabled and dist.is_available() and dist.is_initialized():
+        for grp_attr, size_attr in (('_sp_group', 'sp_world_size'), ('_rp_group', 'rp_world_size')):
+            grp = getattr(sp_strategy, grp_attr, None)
+            if grp is not None and getattr(sp_strategy, size_attr, 1) > 1:
+                pooled = all_reduce(pooled, op=dist.ReduceOp.SUM, group=grp)
+
+    return pooled, is_packed
diff --git a/src/twinkle/model/transformers/task_outputs.py b/src/twinkle/model/transformers/task_outputs.py
new file mode 100644
index 000000000..970d49518
--- /dev/null
+++ b/src/twinkle/model/transformers/task_outputs.py
@@ -0,0 +1,26 @@
+"""Model-side task reductions, followed by sequence-parallel loss collection."""
+from copy import copy
+
+import torch.nn.functional as F
+
+from .strategy.sequence_parallel.pooling import pool_last_valid_token
+
+
+def postprocess_task_outputs(inputs, outputs, *, task='causal_lm', sp_strategy=None, sp_gather_keys=()):
+    """Reduce task features before loss; processors only restore batch layouts."""
+    if task in ('embedding', 'generative_reranker'):
+        features = outputs.get('features') if task == 'embedding' else outputs.get('logits')
+        if features is None:
+            raise ValueError(f'Missing per-token features for task={task!r}')
+        pool = sp_strategy.pool_last_valid_token if sp_strategy is not None else pool_last_valid_token
+        pooled, _ = pool(features, inputs)
+        outputs = copy(outputs)
+        if task == 'embedding':
+            outputs.pop('features', None)
+            outputs['embeddings'] = F.normalize(pooled, p=2, dim=-1).contiguous()
+        else:
+            outputs['logits'] = pooled.contiguous()
+        return inputs, outputs
+    if sp_strategy is not None:
+        return sp_strategy.gather_loss_tensors(inputs, outputs, sp_gather_keys=sp_gather_keys)
+    return inputs, outputs
diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py
index bd59ccef3..f4aab1dca 100644
--- a/src/twinkle/model/transformers/transformers.py
+++ b/src/twinkle/model/transformers/transformers.py
@@ -45,10 +45,13 @@
 from twinkle.template import Template
 from twinkle.utils import (construct_class, get_logger, prepare_replayed_selective_log_softmax,
                            replayed_selective_log_softmax, selective_log_softmax, torch_util)
+
 from twinkle.utils.framework import Torch
 from twinkle.utils.grad_clip import normalize_and_clip_grad_norm
 from twinkle.utils.transformers_utils import filter_from_config_kwargs
 
+from .task_outputs import postprocess_task_outputs
+
 logger = get_logger()
 
 CHECKPOINT_ADAPTER_NAME = 'default'
@@ -105,9 +108,9 @@ def _resolve_task_context(model, task, template=None):
     'causal_lm' (default) keeps the model untouched (returns ``nullcontext``).
     'embedding' swaps lm_head for identity + installs a feature-extraction hook so
     downstream pooling can run inside
-    ``InputProcessor.postprocess_tensor_sp(task='embedding', ...)``.
+    ``postprocess_task_outputs(task='embedding', ...)``.
     'generative_reranker' wraps lm_head to emit a per-token yes-minus-no score; the
-    last-valid-token reduction is deferred to ``postprocess_tensor_sp`` too. The plain
+    last-valid-token reduction is deferred to ``postprocess_task_outputs`` too. The plain
     'reranker' task needs NO patch: it rides a real ``num_labels=1`` classification head
     built at model construction, which already returns ``[B, 1]``.
     """
@@ -816,6 +819,19 @@ def _construct_default_optimizer_group(self):
             acc_strategy=self._acc_strategy,
         )
 
+    def _prepare_sp_inputs(self, inputs, processor, task='causal_lm'):
+        if self.sp_strategy is None:
+            return inputs
+        payload = dict(inputs)
+        payload['padding_free'] = bool(processor.padding_free or processor._any_packing([inputs]))
+        # Sentence-level labels describe whole sequences rather than tokens.
+        sentence_task = task in ('seq_cls', 'reranker', 'embedding', 'generative_reranker')
+        sequence_labels = payload.pop('labels', None) if sentence_task else None
+        payload = self.sp_strategy.preprocess_inputs(payload)
+        if sequence_labels is not None:
+            payload['labels'] = sequence_labels
+        return payload
+
     @remote_function(dispatch='slice_dp', collect=collect_tensor_dict)
     def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajectory]], **kwargs):
         """Call forward function and record the inputs and outputs.
@@ -862,12 +878,12 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
         input_key_specs = getattr(self.sp_strategy, 'input_key_specs', None)
         inputs: Dict[str, Any] = processor(
             inputs,
-            sp_strategy=self.sp_strategy,
             input_key_specs=input_key_specs,
             model=self.model,
             hf_config=self.hf_config,
             enable_sp=getattr(self, '_enable_sp', False),
         )
+        inputs = self._prepare_sp_inputs(inputs, processor, task=task)
         labels: torch.Tensor = inputs.pop('labels', None)
         loss_scale = inputs.pop('loss_scale', None)
         channel = inputs.pop('channel', None)
@@ -937,7 +953,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
         outputs['past_key_values'] = None
         if not (return_logits or loss_require_logits):
             outputs['logits'] = None
-        inputs, outputs = processor.postprocess_tensor_sp(
+        inputs, outputs = postprocess_task_outputs(
             inputs,
             outputs,
             sp_strategy=self.sp_strategy,
@@ -1011,12 +1027,12 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
             input_key_specs = getattr(self.sp_strategy, 'input_key_specs', None)
             inputs: Dict[str, Any] = processor(
                 inputs,
-                sp_strategy=self.sp_strategy,
                 input_key_specs=input_key_specs,
                 model=self.model,
                 hf_config=self.hf_config,
                 enable_sp=getattr(self, '_enable_sp', False),
             )
+            inputs = self._prepare_sp_inputs(inputs, processor, task=task)
             labels = inputs.pop('labels', None)
             loss_scale = inputs.pop('loss_scale', None)
             channel = inputs.pop('channel', None)
@@ -1089,7 +1105,7 @@ def forward_only(self, *, inputs: Union[InputFeature, List[InputFeature], List[T
             outputs['past_key_values'] = None
             if not (return_logits or loss_require_logits):
                 outputs['logits'] = None
-            inputs, outputs = processor.postprocess_tensor_sp(
+            inputs, outputs = postprocess_task_outputs(
                 inputs,
                 outputs,
                 sp_strategy=self.sp_strategy,
diff --git a/src/twinkle/patch/transformers_emb.py b/src/twinkle/patch/transformers_emb.py
index 2a14b852f..7d08519d0 100644
--- a/src/twinkle/patch/transformers_emb.py
+++ b/src/twinkle/patch/transformers_emb.py
@@ -10,7 +10,7 @@
    ``[B, T, H]`` or ``[B, T_local, H]`` under SP).
 
 Last-token pooling (incl. padding-free, SP gather) is **deferred** to
-``InputProcessor.postprocess_tensor_sp(task='embedding', ...)`` so this patch
+``postprocess_task_outputs(task='embedding', ...)`` so this patch
 stays SP/CP/packed-agnostic and the dispatch sits in one place.
 
 Both mutations are reverted by ``unpatch``.
@@ -44,7 +44,7 @@ def get_lm_head_model(module, lm_heads=None):
 def _output_features_hook(module, args, kwargs, output):
     # Replaces the model's structured output entirely; downstream HF utilities
     # (generate, loss, past_key_values, ...) won't work until ``unpatch``.
-    # Pooling + L2-normalize is deferred to ``InputProcessor._postprocess_embedding``
+    # Pooling + L2-normalize is deferred to ``postprocess_task_outputs``
     # so the hook stays pooling-agnostic (last-token / mean / CLS all work uniformly).
     hidden_states = output.logits
     return {'features': hidden_states.contiguous()}
diff --git a/src/twinkle/patch/transformers_reranker.py b/src/twinkle/patch/transformers_reranker.py
index cd983c2fe..4a8e8c93e 100644
--- a/src/twinkle/patch/transformers_reranker.py
+++ b/src/twinkle/patch/transformers_reranker.py
@@ -14,7 +14,7 @@
    are introduced (unlike the ``num_labels=1`` sequence-classification head used by the plain
    ``reranker`` task, which trains a fresh ``score`` linear).
 2. Last-valid-token reduction ``[B, T, 1] -> [B, 1]`` is DEFERRED to
-   ``InputProcessor.postprocess_tensor_sp(task='generative_reranker', ...)`` -- symmetric with how
+   ``postprocess_task_outputs(task='generative_reranker', ...)`` -- symmetric with how
    embedding defers pooling -- so this patch stays SP/CP/packed-agnostic and the dispatch lives in
    one place.
 
diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py
index 02b202042..111178b82 100644
--- a/src/twinkle/processor/base.py
+++ b/src/twinkle/processor/base.py
@@ -1,8 +1,6 @@
 # Copyright (c) ModelScope Contributors. All rights reserved.
 import numpy as np
 import torch
-import torch.distributed as dist
-import torch.nn.functional as F
 from copy import copy
 from dataclasses import dataclass
 from typing import Any, Dict, List, Literal, Optional, Union
@@ -80,7 +78,6 @@ def __init__(self,
             self.prepare_transformers_padding_free_patch,
             self.drop_causal_4d_mask,
             self.split_cp,
-            self.apply_transformers_sp,
             self.prepare_outputs,
         ]
 
@@ -149,184 +146,6 @@ def to_tensor(_input):
 
         return [to_tensor(_input) for _input in inputs]
 
-    def apply_transformers_sp(self, inputs: List[InputFeature], **kwargs) -> List[InputFeature]:
-        sp_strategy = kwargs.get('sp_strategy')
-        if self.framework != 'transformers' or sp_strategy is None:
-            return inputs
-        padding_free = bool(self.padding_free or self._any_packing(inputs))
-        results = []
-        for _input in inputs:
-            payload = dict(_input)
-            payload['padding_free'] = padding_free
-            results.append(InputFeature(**sp_strategy.preprocess_inputs(payload)))
-        return results
-
-    def postprocess_tensor_sp(self, inputs: Dict[str, Any], outputs: Dict[str, Any],
-                              sp_gather_keys=(), **kwargs) -> tuple[Dict[str, Any], Dict[str, Any]]:
-        """Adjust SP tensors after forward and before loss computation.
-
-        Pipeline: SP gather → packed-sequence unpack.
-        After this call, logps and labels are in per-sequence batch format
-        ``[num_sequences, max_seq_len]`` when the input was packed, or left
-        unchanged for normal (non-packed) batches.
-
-        For ``task='embedding'`` this also performs the last-valid-token
-        pooling (with padding-free / SP gather awareness) and writes the
-        pooled ``[n_seqs, H]`` tensor to ``outputs['embeddings']``; the raw
-        per-token ``outputs['features']`` is consumed and removed.
-        """
-        sp_strategy = kwargs.get('sp_strategy')
-        task = kwargs.get('task', 'causal_lm')
-        if task == 'embedding':
-            return self._postprocess_embedding(inputs, outputs, sp_strategy=sp_strategy)
-        if task == 'generative_reranker':
-            return self._postprocess_generative_reranker(inputs, outputs, sp_strategy=sp_strategy)
-        if self.framework == 'transformers' and sp_strategy is not None:
-            return sp_strategy.gather_loss_tensors(inputs, outputs, sp_gather_keys=sp_gather_keys)
-        return inputs, outputs
-
-    @staticmethod
-    def _packed_last_indices(position_ids: torch.Tensor, total_len: int) -> torch.Tensor:
-        """For padding-free batches: per-segment last-token indices into a [1, total] sequence."""
-        flat = position_ids.squeeze(0) if position_ids.dim() == 2 else position_ids
-        starts = (flat == 0).nonzero(as_tuple=False).squeeze(-1)
-        end_anchor = torch.tensor([total_len], device=flat.device, dtype=starts.dtype)
-        boundaries = torch.cat([starts, end_anchor])
-        return (boundaries[1:] - 1).long()
-
-    def _postprocess_embedding(self,
-                               inputs: Dict[str, Any],
-                               outputs: Dict[str, Any],
-                               sp_strategy=None) -> tuple[Dict[str, Any], Dict[str, Any]]:
-        """Pool per-token features to per-sequence embeddings (last-valid-token).
-
-        Build a one-hot end-token mask in the un-padded global frame, route it
-        through the same pad+split as ``input_ids`` so it aligns with local
-        features, pool locally, then ``all_reduce`` only the ``[n_seqs, H]``
-        tensor across SP × RP. No feature gather; uniform across
-        DP / Ulysses / zigzag-ring / padding-free.
-        """
-        features = outputs.get('features')
-        assert features is not None, (
-            "outputs['features'] missing; ensure TransformersEmbeddingPatch is applied for task='embedding'.")
-
-        pooled, is_packed = self._pool_last_valid_token(features, inputs, sp_strategy=sp_strategy)
-
-        # Normalize after pool+reduce so swapping pooling strategy (last/mean/CLS) stays mathematically correct.
-        embeddings = F.normalize(pooled, p=2, dim=-1)
-
-        outputs = copy(outputs)
-        outputs.pop('features', None)
-        outputs['embeddings'] = embeddings.contiguous()
-        return inputs, outputs
-
-    def _postprocess_generative_reranker(self,
-                                         inputs: Dict[str, Any],
-                                         outputs: Dict[str, Any],
-                                         sp_strategy=None) -> tuple[Dict[str, Any], Dict[str, Any]]:
-        """Reduce per-token generative-reranker scores ``[B, T, 1]`` to ``[n_seqs, 1]``.
-
-        Symmetric with embedding pooling (same last-valid-token gather, SP/packed-aware), but the
-        source is ``outputs['logits']`` -- the yes-minus-no score TransformersGenerativeRerankerPatch
-        writes per token -- and there is NO L2 normalize: a relevance score is a scalar, not a
-        direction. Mirrors legacy RerankerTrainer, which reads ``logits[batch, last_valid]``.
-        """
-        logits = outputs.get('logits')
-        assert logits is not None, (
-            "outputs['logits'] missing; ensure TransformersGenerativeRerankerPatch is applied and the "
-            'reranker loss sets require_logits=True.')
-        pooled, _ = self._pool_last_valid_token(logits, inputs, sp_strategy=sp_strategy)
-        outputs = copy(outputs)
-        outputs['logits'] = pooled.contiguous()
-        return inputs, outputs
-
-    def _pool_last_valid_token(self,
-                               features: torch.Tensor,
-                               inputs: Dict[str, Any],
-                               sp_strategy=None) -> tuple[torch.Tensor, bool]:
-        """Gather the last valid token per sequence from ``[B, T, C]`` -> ``[n_seqs, C]``.
-
-        Build a one-hot end-token mask in the un-padded global frame, route it
-        through the same pad+split as ``input_ids`` so it aligns with local
-        features, pool locally, then ``all_reduce`` only the ``[n_seqs, C]``
-        tensor across SP × RP. No feature gather; uniform across
-        DP / Ulysses / zigzag-ring / padding-free. Returns ``(pooled, is_packed)``.
-        """
-        sp_enabled = (
-            self.framework == 'transformers' and sp_strategy is not None and getattr(sp_strategy, 'enabled', False)
-            and getattr(sp_strategy, 'world_size', 1) > 1)
-
-        if sp_enabled:
-            ref_pos = sp_strategy.real_position_ids
-        elif inputs.get('position_ids') is not None:
-            ref_pos = inputs['position_ids']
-        else:
-            # Standard padded batch: HF omits position_ids (it derives them internally).
-            # Rebuild a position_ids-like reference (valid token -> its index, padding -> -1)
-            # so the last-valid-token pooling below, which keys off ``ref_pos >= 0`` for
-            # validity and a single ``== 0`` per row for packing detection, still works.
-            T = features.shape[1]
-            arange = torch.arange(T, device=features.device)
-            am = inputs.get('attention_mask')
-            if am is None:
-                ref_pos = arange.expand(features.shape[0], T)
-            else:
-                if am.dim() >= 3:
-                    am = am.reshape(am.shape[0], -1)[:, :T]
-                ref_pos = torch.where(am.bool(), arange, arange.new_tensor(-1))
-        if ref_pos.dim() == 3:
-            ref_pos = ref_pos[0]
-        # Accept both flash-attn (cu_seq_lens_q) and Megatron (cu_seqlens_q) conventions.
-        cu_seq_lens_q = inputs.get('cu_seq_lens_q', inputs.get('cu_seqlens_q'))
-
-        is_packed = (
-            features.shape[0] == 1 and (cu_seq_lens_q is not None or int((ref_pos.reshape(-1) == 0).sum()) > 1))
-
-        device, dtype = features.device, features.dtype
-        T_real = ref_pos.shape[-1]
-
-        if is_packed:
-            if torch.is_tensor(cu_seq_lens_q) and cu_seq_lens_q.numel() >= 2:
-                end_idx = (cu_seq_lens_q[1:].long() - 1).to(device)
-            else:
-                end_idx = self._packed_last_indices(ref_pos, T_real).to(device)
-            n_seqs = end_idx.shape[0]
-            batch_idx = torch.zeros_like(end_idx)
-            target_idx = torch.arange(n_seqs, device=device)
-            B_mask, n_targets = 1, n_seqs
-        else:
-            n_seqs = ref_pos.shape[0]
-            # clamp_min(0) prevents negative indexing when a row is fully padded (sum == 0).
-            end_idx = ((ref_pos >= 0).long().sum(-1) - 1).clamp_min(0)
-            batch_idx = torch.arange(n_seqs, device=device)
-            target_idx = torch.zeros(n_seqs, dtype=torch.long, device=device)
-            B_mask, n_targets = n_seqs, 1
-
-        mask = torch.zeros(B_mask, T_real, n_targets, dtype=dtype, device=device)
-        mask[batch_idx, end_idx, target_idx] = 1.0
-
-        if sp_enabled:
-            # Route mask through the same pad+split as input_ids to align with local features.
-            rp = sp_strategy.real_position_ids
-            rp_padded = sp_strategy.pad(rp, padding_value=-1, position_ids=rp, dim=-1)
-            mask = sp_strategy.pad(mask, padding_value=0, position_ids=rp, dim=1)
-            mask = sp_strategy.split(mask, dim=1, position_ids=rp_padded)
-
-        assert mask.shape[1] == features.shape[1], (
-            f'mask seq-dim {mask.shape[1]} != features seq-dim {features.shape[1]}; '
-            'sp_strategy.split contract broken or T_real mismatched.')
-
-        pooled = torch.einsum('bth,btn->bnh', features, mask)
-        pooled = pooled.squeeze(0) if is_packed else pooled.squeeze(1)
-
-        if sp_enabled and dist.is_available() and dist.is_initialized():
-            for grp_attr, size_attr in (('_sp_group', 'sp_world_size'), ('_rp_group', 'rp_world_size')):
-                grp = getattr(sp_strategy, grp_attr, None)
-                if grp is not None and getattr(sp_strategy, size_attr, 1) > 1:
-                    dist.all_reduce(pooled, op=dist.ReduceOp.SUM, group=grp)
-
-        return pooled, is_packed
-
     def pad_cp(self, inputs: List[InputFeature], **kwargs) -> List[InputFeature]:
 
         if self.device_mesh is None:
@@ -681,7 +500,7 @@ def unpack_packed_sequences(
         this restoration step does not require a parallel strategy.
 
         For ``task='embedding'`` the outputs are already pooled to ``[n_seqs, H]``
-        by ``postprocess_tensor_sp``, so this is a no-op. Same for
+        by model-side task postprocessing, so this is a no-op. Same for
         ``task='generative_reranker'``, whose logits are already reduced to ``[n_seqs, 1]``.
         """
         if task in ('embedding', 'generative_reranker'):
diff --git a/tests/loss/test_channel_loss.py b/tests/loss/test_channel_loss.py
index a001147ae..ce3690e5b 100644
--- a/tests/loss/test_channel_loss.py
+++ b/tests/loss/test_channel_loss.py
@@ -125,9 +125,6 @@ def __init__(self):
         def __call__(self, inputs, **kwargs):
             return dict(inputs)
 
-        def postprocess_tensor_sp(self, inputs, outputs, **kwargs):
-            return inputs, outputs
-
         def unpack_packed_sequences(self, inputs, outputs=None, **kwargs):
             return inputs, outputs
 
diff --git a/tests/transformers/test_sp_data_keys.py b/tests/transformers/test_sp_data_keys.py
index 44efb2941..d468d97cd 100644
--- a/tests/transformers/test_sp_data_keys.py
+++ b/tests/transformers/test_sp_data_keys.py
@@ -169,7 +169,8 @@ def test_processor_preserves_registered_keys_before_sp_split(monkeypatch, rank):
               'position_ids': torch.arange(5).unsqueeze(0),
               'labels': torch.ones(1, 5, dtype=torch.long),
               'token_weights': values, 'unregistered': values}
-    result = processor([sample], sp_strategy=strategy, input_key_specs=strategy.input_key_specs)
+    result = processor([sample], input_key_specs=strategy.input_key_specs)
+    result = strategy.preprocess_inputs(result)
     expected = torch.tensor([[0, 1, 2, 3, 4, -2]])[:, rank * 3:(rank + 1) * 3]
     torch.testing.assert_close(result['token_weights'], expected)
     assert 'unregistered' not in result
@@ -202,7 +203,8 @@ def test_registered_multidimensional_key_collation_and_loss_boundary(monkeypatch
                         'position_ids': list(range(length)), 'labels': list(range(1, length + 1)),
                         'token_features': [ [[1] * length, [2] * length] ],
                         'completion_mask': [1] * length})
-    result = processor(samples, sp_strategy=strategy, input_key_specs=strategy.input_key_specs)
+    result = processor(samples, input_key_specs=strategy.input_key_specs)
+    result = strategy.preprocess_inputs(result)
     expected = torch.tensor([[[1, 1, 1], [2, 2, 2]], [[1, 1, 1], [2, 2, 2]]])
     torch.testing.assert_close(result['token_features'], expected)
     auxiliary = processor.pop_sp_auxiliary_inputs(result, input_key_specs=strategy.input_key_specs)
@@ -228,8 +230,8 @@ def forward(self, token_type_ids, **kwargs):
     strategy._ctx._prepare_forward_hook(model)
     ids = torch.arange(5).unsqueeze(0)
     prepared = InputProcessor()([{'input_ids': ids, 'position_ids': ids,
-                                  'token_type_ids': ids}], sp_strategy=strategy,
-                                input_key_specs=strategy.input_key_specs)
+                                  'token_type_ids': ids}], input_key_specs=strategy.input_key_specs)
+    prepared = strategy.preprocess_inputs(prepared)
     assert prepared['token_type_ids'].shape == (1, 5)
     torch.testing.assert_close(model(**prepared), torch.tensor([[3, 4, 0]]))
 
diff --git a/tests/transformers/test_sp_loss_contract.py b/tests/transformers/test_sp_loss_contract.py
index d6e6088aa..9345e4348 100644
--- a/tests/transformers/test_sp_loss_contract.py
+++ b/tests/transformers/test_sp_loss_contract.py
@@ -151,7 +151,8 @@ def _ppo_value_worker(rank, port):
         local_labels = prepared.pop('labels')
         local_outputs = model(**{key: prepared[key] for key in ('input_ids', 'position_ids')})
         local_outputs['values'] = local_outputs.pop('logits')
-        loss_inputs, loss_outputs = InputProcessor().postprocess_tensor_sp(
+        from twinkle.model.transformers.task_outputs import postprocess_task_outputs
+        loss_inputs, loss_outputs = postprocess_task_outputs(
             {'labels': local_labels},
             local_outputs,
             sp_strategy=strategy,
@@ -287,3 +288,92 @@ def test_ppo_value_loss_through_transformers_sdpa_sp_training_pipeline():
         sock.bind(('127.0.0.1', 0))
         port = sock.getsockname()[1]
     mp.spawn(_ppo_training_pipeline_worker, args=(port,), nprocs=2, join=True)
+
+
+@pytest.mark.parametrize('task', ['embedding', 'generative_reranker'])
+def test_model_task_pooling_without_parallel_strategy(task):
+    from twinkle.model.transformers.task_outputs import postprocess_task_outputs
+
+    features = torch.arange(28, dtype=torch.float32).reshape(2, 7, 2).requires_grad_()
+    # Left padding on row 0, right padding on row 1: use the rightmost valid index.
+    positions = torch.tensor([[-1, -1, 0, 1, 2, 3, 4], [0, 1, 2, -1, -1, -1, -1]])
+    key = {'embedding': 'features', 'generative_reranker': 'logits'}[task]
+    _, outputs = postprocess_task_outputs({'position_ids': positions}, {key: features}, task=task)
+    expected = features[torch.arange(2), torch.tensor([6, 2])]
+    if task == 'embedding':
+        expected = torch.nn.functional.normalize(expected, dim=-1)
+    actual = outputs['embeddings' if task == 'embedding' else 'logits']
+    torch.testing.assert_close(actual, expected)
+    actual.sum().backward()
+    assert features.grad is not None and torch.isfinite(features.grad).all()
+    assert 'sequence_scores' not in outputs
+
+
+@pytest.mark.parametrize('task', ['seq_cls', 'reranker', 'embedding', 'generative_reranker'])
+def test_sequence_labels_remain_global_in_model_sp_preparation(task):
+    from twinkle.model.transformers.transformers import TransformersModel
+    from twinkle.processor import InputProcessor
+
+    model = SimpleNamespace(sp_strategy=_strategy())
+    labels = torch.tensor([[0], [1]])
+    positions = torch.arange(5).expand(2, -1)
+    inputs = {'input_ids': positions, 'position_ids': positions, 'labels': labels}
+    prepared = TransformersModel._prepare_sp_inputs(model, inputs, InputProcessor(), task=task)
+    assert prepared['labels'] is labels
+    assert 'padding_free' not in prepared
+
+
+def _sp_pooling_worker(rank, port):
+    os.environ['RANK'], os.environ['WORLD_SIZE'] = str(rank), '2'
+    torch.set_num_threads(1)
+    from torch.utils.checkpoint import DefaultDeviceType
+    DefaultDeviceType.set_device_type('cpu')
+    dist.init_process_group('gloo', init_method=f'tcp://127.0.0.1:{port}', rank=rank,
+                            world_size=2, timeout=timedelta(seconds=120))
+    try:
+        from twinkle.model.transformers.task_outputs import postprocess_task_outputs
+
+        # Pooling values and averaged parameter gradients match a full sequence,
+        # for both ordinary batches and packed samples, with non-divisible lengths.
+        for topology, packed in (('ulysses', False), ('ulysses', True), ('ring', True)):
+            positions = (torch.tensor([[0, 1, 2, 0, 1, 2, 3]]) if packed else
+                         torch.tensor([[0, 1, 2, -1, -1, -1, -1], [0, 1, 2, 3, 4, 5, 6]]))
+            values = torch.arange(positions.numel() * 2, dtype=torch.float32).reshape(*positions.shape, 2) + 1
+            for task in ('embedding', 'generative_reranker'):
+                key = {'embedding': 'features', 'generative_reranker': 'logits'}[task]
+                strategy = _strategy()
+                strategy._ctx._sp_rank = rank
+                strategy._ctx._sp_group = dist.group.WORLD
+                strategy._ctx.extra_kwargs['position_ids'] = positions
+                if topology == 'ring':
+                    strategy._ctx.sp_world_size = 1
+                    strategy._ctx.rp_world_size = 2
+                    strategy._ctx._sp_rank = 0
+                    strategy._ctx._rp_rank = rank
+                    strategy._ctx._sp_group = None
+                    strategy._ctx._rp_group = dist.group.WORLD
+                base_scale = torch.tensor([0.7, 1.2], requires_grad=True)
+                scale = base_scale.detach().clone().requires_grad_()
+                _, base_outputs = postprocess_task_outputs({'position_ids': positions}, {key: values * base_scale}, task=task)
+                padded_positions = strategy.pad(positions, padding_value=-1, position_ids=positions)
+                local = strategy.pad(values, padding_value=0, dim=1, position_ids=positions)
+                local = strategy.split(local, dim=1, position_ids=padded_positions)
+                _, outputs = postprocess_task_outputs({'position_ids': positions}, {key: local * scale},
+                                                     task=task, sp_strategy=strategy)
+                name = 'embeddings' if task == 'embedding' else 'logits'
+                torch.testing.assert_close(outputs[name], base_outputs[name])
+                outputs[name].square().sum().backward()
+                base_outputs[name].square().sum().backward()
+                dist.all_reduce(scale.grad)
+                torch.testing.assert_close(scale.grad / 2, base_scale.grad, atol=2e-5, rtol=2e-5)
+
+    finally:
+        dist.destroy_process_group()
+
+
+@pytest.mark.slow
+def test_sp_pooling_and_gradients_match_full_sequence():
+    with socket.socket() as sock:
+        sock.bind(('127.0.0.1', 0))
+        port = sock.getsockname()[1]
+    mp.spawn(_sp_pooling_worker, args=(port,), nprocs=2, join=True)

From 741efca404a52892edda8ac2ba21baa7ab66d1fc Mon Sep 17 00:00:00 2001
From: clc <clc0301@cmbchina.com>
Date: Mon, 26 Oct 2026 14:24:51 +0800
Subject: [PATCH 6/8] sp: decouple field schemas and fix contract edge cases

Resolve shared model and loss field declarations independently of SP, and partition forward inputs at the model boundary.

Honor declared processing boundaries, preserve direct registration, and release causal-mask owners through weak references.

Cover scalar labels, optional MoE masks, packed custom losses, and invalid loss-input boundaries with CPU regressions.
---
 src/twinkle/data_fields.py                    | 216 ++++++++++++++++++
 src/twinkle/loss/base.py                      |   5 +-
 src/twinkle/loss/value.py                     |   4 +-
 .../strategy/sequence_parallel/__init__.py    |  84 ++++---
 .../strategy/sequence_parallel/adapters.py    |   2 +-
 .../strategy/sequence_parallel/data_keys.py   | 175 +-------------
 .../model/transformers/transformers.py        |  62 ++---
 src/twinkle/processor/base.py                 |  28 +--
 .../transformers/test_sp_context_isolation.py |  58 +++++
 tests/transformers/test_sp_custom_loss_cpu.py |  28 +--
 tests/transformers/test_sp_data_keys.py       | 183 ++++++++++++++-
 tests/transformers/test_sp_loss_contract.py   |   2 +-
 tests/transformers/test_sp_model_adapters.py  |   6 +-
 13 files changed, 574 insertions(+), 279 deletions(-)
 create mode 100644 src/twinkle/data_fields.py

diff --git a/src/twinkle/data_fields.py b/src/twinkle/data_fields.py
new file mode 100644
index 000000000..0981416ff
--- /dev/null
+++ b/src/twinkle/data_fields.py
@@ -0,0 +1,216 @@
+"""Shared sequence-field schemas consumed by processors, model wrappers and SP."""
+from dataclasses import dataclass
+from typing import Optional
+
+
+_BOUNDARIES = frozenset({'prepare', 'forward', 'model_internal', 'router_replay'})
+
+
+@dataclass(frozen=True)
+class DataKeySpec:
+    """Describe collation, model routing and optional parallel transforms.
+
+    Auxiliary loss fields stay global by default. model_input controls routing;
+    stage/pad_at/split_at control when an explicitly local tensor is transformed.
+    """
+
+    name: str
+    split_dim: int = -1
+    pad_value: int | float | str = 0
+    need_roll: bool = False
+    side: str = 'input'
+    stage: str = 'auxiliary'
+    need_split: Optional[bool] = None
+    need_gather: bool = False
+    autograd: bool = False
+    owner: str = 'generic'
+    need_pad: Optional[bool] = None
+    pad_at: Optional[str] = 'auto'
+    split_at: Optional[str] = 'auto'
+    need_unpack: Optional[bool] = None
+    model_input: Optional[bool] = None
+
+    def __post_init__(self):
+        if not isinstance(self.name, str) or not self.name:
+            raise ValueError('An SP input key must have a non-empty name')
+        if not isinstance(self.split_dim, int):
+            raise TypeError('split_dim must be an integer')
+        if self.side not in ('input', 'output'):
+            raise ValueError('side must be input or output')
+        if self.stage not in ('forward', 'auxiliary', 'model_internal', 'router_replay', 'output'):
+            raise ValueError('Invalid SP key stage')
+        if (self.side == 'output') != (self.stage == 'output'):
+            raise ValueError('Output keys must use stage=output')
+        if self.owner not in ('generic', 'task', 'model', 'loss'):
+            raise ValueError('Invalid SP key owner')
+        if isinstance(self.pad_value, str) and (self.name != 'input_ids' or self.pad_value != 'tokenizer'):
+            raise ValueError('Only input_ids can use the tokenizer padding sentinel')
+        if self.need_split is None:
+            local_stage = self.stage in ('forward', 'model_internal', 'router_replay')
+            explicit_boundary = self.split_at not in (None, 'auto')
+            object.__setattr__(self, 'need_split', self.side == 'input'
+                               and (local_stage or explicit_boundary or self.need_gather or self.need_roll))
+        if self.side == 'input' and self.need_gather and not self.need_split:
+            raise ValueError('Only sharded input keys can request gathering')
+        if self.stage == 'forward' and self.need_gather:
+            raise ValueError('Model-forward keys cannot request loss gathering')
+        if self.model_input is None:
+            object.__setattr__(self, 'model_input', self.stage == 'forward')
+        default_boundary = {
+            'auxiliary': 'prepare',
+            'forward': 'forward',
+            'model_internal': 'model_internal',
+            'router_replay': 'router_replay',
+            'output': None,
+        }[self.stage]
+        if self.need_pad is None:
+            object.__setattr__(self, 'need_pad', self.side == 'input'
+                               and (self.need_split or self.pad_at not in (None, 'auto')))
+        if self.pad_at == 'auto':
+            object.__setattr__(self, 'pad_at', default_boundary if self.need_pad else None)
+        elif self.pad_at is not None and self.pad_at not in _BOUNDARIES:
+            raise ValueError(f'Invalid SP padding boundary: {self.pad_at!r}')
+        elif (self.pad_at is None) == bool(self.need_pad):
+            raise ValueError('need_pad and pad_at specify conflicting SP padding behavior')
+        if self.pad_at != 'auto':
+            object.__setattr__(self, 'need_pad', self.pad_at is not None)
+        if self.split_at == 'auto':
+            object.__setattr__(self, 'split_at', default_boundary if self.need_split else None)
+        elif self.split_at is not None and self.split_at not in _BOUNDARIES:
+            raise ValueError(f'Invalid SP split boundary: {self.split_at!r}')
+        elif (self.split_at is None) == bool(self.need_split):
+            raise ValueError('need_split and split_at specify conflicting SP split behavior')
+        if self.split_at != 'auto':
+            object.__setattr__(self, 'need_split', self.split_at is not None)
+        if self.side == 'output' and (self.pad_at is not None or self.split_at is not None):
+            raise ValueError('Output keys cannot request input padding or splitting')
+        if self.need_unpack is None:
+            # Preserve the historical packed restoration behavior unless a
+            # built-in rule explicitly opts out.
+            object.__setattr__(self, 'need_unpack', self.side == 'output' or self.need_gather or not self.need_split)
+
+    def sequence_dim(self, ndim):
+        if not -ndim <= self.split_dim < ndim:
+            raise ValueError(f'Invalid sequence axis {self.split_dim} for SP key {self.name!r}')
+        dim = self.split_dim % ndim
+        if dim == 0:
+            raise ValueError(f'SP key {self.name!r} cannot split or gather the batch axis')
+        return dim
+
+
+BUILTIN_INPUT_KEYS = (
+    DataKeySpec('input_ids', split_dim=1, pad_value='tokenizer', stage='forward', need_split=True),
+    DataKeySpec('inputs_embeds', split_dim=1, stage='forward', need_split=True),
+    DataKeySpec('position_ids', pad_value=-1, stage='forward', need_split=True),
+    DataKeySpec('attention_mask', split_dim=1, stage='forward', need_pad=True, need_split=False),
+    DataKeySpec('labels', pad_value=-100, need_split=True, need_gather=True, need_unpack=True),
+    # These are not needed by the model forward. Keep them global through SP;
+    # packed sequence restoration is a separate action from sharding/gathering.
+    DataKeySpec('completion_mask', need_pad=False, need_split=False, need_unpack=True),
+    DataKeySpec('loss_scale', need_pad=False, need_split=False, need_unpack=True),
+)
+
+# These are processor/model control arguments, not sequence data contracts.
+# Keep them out of DataKeySpec and preserve them through the Transformers input
+# conversion path explicitly instead.
+PROCESSOR_PASSTHROUGH_KEYS = ('cache_position', 'padding_free')
+
+BUILTIN_OUTPUT_KEYS = (
+    DataKeySpec('logps', split_dim=1, side='output', stage='output',
+                need_split=False, need_gather=True, autograd=True, need_unpack=True),
+    DataKeySpec('entropies', split_dim=1, side='output', stage='output',
+                need_split=False, need_gather=True, autograd=True, need_unpack=True),
+    # Real logits are opt-in; fused CE explicitly gathers this hidden-state slot.
+    DataKeySpec('logits', split_dim=1, side='output', stage='output',
+                need_split=False, autograd=True, need_unpack=True),
+)
+
+
+class InputKeyRegistry:
+    """Collect provider declarations; reject conflicting definitions.
+
+    Auxiliary keys are processed at their declared boundary (usually
+    preprocessing). Custom model-forward keys remain global until the forward
+    hook, so every key is sharded at most once.
+    """
+
+    _FORWARD_KEYS = frozenset({'input_ids', 'inputs_embeds', 'position_ids', 'attention_mask'})
+    _SPECIAL_KEYS = frozenset(PROCESSOR_PASSTHROUGH_KEYS)
+
+    def __init__(self):
+        self._specs = {spec.name: spec for spec in BUILTIN_INPUT_KEYS}
+
+    def register(self, spec: DataKeySpec):
+        if not isinstance(spec, DataKeySpec):
+            raise TypeError('Expected a DataKeySpec')
+        if spec.side != 'input' or spec.stage not in ('auxiliary', 'forward'):
+            raise ValueError('Register only auxiliary or model-forward input keys here')
+        if spec.name in self._FORWARD_KEYS:
+            raise ValueError(f'{spec.name!r} is owned by the model forward hook')
+        if spec.name in self._SPECIAL_KEYS:
+            raise ValueError(f'{spec.name!r} is a special processor/model argument, not an SP data key')
+        if not spec.model_input and 'forward' in (spec.pad_at, spec.split_at):
+            raise ValueError(f'Loss-only input {spec.name!r} cannot use the forward boundary; '
+                             'use prepare, or set model_input=True for a model argument')
+        existing = self._specs.get(spec.name)
+        if existing is not None and existing != spec:
+            raise ValueError(f'Conflicting SP input rule for {spec.name!r}')
+        self._specs[spec.name] = spec
+
+    def __iter__(self):
+        return iter(self._specs.values())
+
+    def get(self, name):
+        return self._specs.get(name)
+
+
+class OutputKeyRegistry:
+    def __init__(self):
+        self._specs = {spec.name: spec for spec in BUILTIN_OUTPUT_KEYS}
+
+    def register(self, spec):
+        if not isinstance(spec, DataKeySpec):
+            raise TypeError('Expected a DataKeySpec')
+        if spec.side != 'output':
+            raise ValueError('Expected an output key rule')
+        existing = self._specs.get(spec.name)
+        if existing is not None and existing != spec:
+            raise ValueError(f'Conflicting SP output rule for {spec.name!r}')
+        self._specs[spec.name] = spec
+
+    def __iter__(self):
+        return iter(self._specs.values())
+
+    def get(self, name):
+        return self._specs.get(name)
+
+
+def declared_key_specs(provider, side):
+    """Read common field declarations, retaining legacy loss declaration names."""
+    specs = getattr(provider, f'{side}_key_specs', ())
+    if callable(specs):
+        specs = specs()
+    legacy = getattr(provider, f'sp_{side}_key_specs', ())
+    for entries in (specs, legacy):
+        if isinstance(entries, DataKeySpec):
+            entries = (entries,)
+        yield from entries or ()
+
+
+def resolve_key_specs(*providers):
+    """Build a call-local schema independently of any parallel strategy."""
+    inputs, outputs = InputKeyRegistry(), OutputKeyRegistry()
+    for provider in providers:
+        for side, registry in (('input', inputs), ('output', outputs)):
+            for spec in declared_key_specs(provider, side):
+                registry.register(spec)
+    return tuple(inputs), tuple(outputs)
+
+
+def partition_model_inputs(inputs, input_key_specs):
+    """Separate forward arguments from loss-only tensors without mutating the batch."""
+    loss_names = {spec.name for spec in input_key_specs if not spec.model_input}
+    loss_names.add('channel')  # sample-level loss metadata, not a sequence tensor
+    model_inputs = {key: value for key, value in inputs.items() if key not in loss_names}
+    loss_inputs = {key: value for key, value in inputs.items() if key in loss_names}
+    return model_inputs, loss_inputs
diff --git a/src/twinkle/loss/base.py b/src/twinkle/loss/base.py
index 62381db6c..b5b9181c8 100644
--- a/src/twinkle/loss/base.py
+++ b/src/twinkle/loss/base.py
@@ -8,11 +8,14 @@ class Loss:
     require_entropy = False
     require_logps = True
     # Sequence-aligned outputs this loss needs in full-sequence layout.
-    # Keys must have an output DataKeySpec registered by SP or a model adapter.
+    # Keys must be declared in the model/loss output schema.
     sp_gather_keys = ()
     # A loss may declare additional sequence-aligned input/output contracts.
     # Keep the core loss API independent of the Transformers backend by
     # importing DataKeySpec only in concrete loss implementations.
+    input_key_specs = ()
+    output_key_specs = ()
+    # Compatibility for existing loss plugins. New declarations use the common names.
     sp_input_key_specs = ()
     sp_output_key_specs = ()
     require_values = False
diff --git a/src/twinkle/loss/value.py b/src/twinkle/loss/value.py
index b6f63ae7b..df1f463ea 100644
--- a/src/twinkle/loss/value.py
+++ b/src/twinkle/loss/value.py
@@ -5,7 +5,7 @@
 from twinkle.data_format import LossOutput
 from twinkle.loss.base import Loss
 from twinkle.loss.grpo import GRPOLoss
-from twinkle.model.transformers.strategy.sequence_parallel.data_keys import DataKeySpec
+from twinkle.data_fields import DataKeySpec
 
 if TYPE_CHECKING:
     import torch
@@ -23,7 +23,7 @@ class PPOValueLoss(Loss):
     require_values = True
     # PPO regresses per-token values; gather the value sequence before loss.
     sp_gather_keys = ('values',)
-    sp_output_key_specs = (DataKeySpec('values', split_dim=1, side='output', stage='output',
+    output_key_specs = (DataKeySpec('values', split_dim=1, side='output', stage='output',
                                        need_split=False, autograd=True, need_unpack=True),)
 
     def __init__(self, cliprange_value: float = 0.2, vf_coef: float = 1.0, ignore_index: int = -100, **kwargs):
diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
index 14d949119..80591ba6c 100644
--- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
@@ -5,9 +5,10 @@
 import torch.distributed as dist
 from copy import copy
 from dataclasses import asdict, dataclass, is_dataclass
+from types import SimpleNamespace
 from transformers import PreTrainedTokenizer
 from typing import Any, Dict, List, Optional, Tuple, Union
-from weakref import WeakKeyDictionary
+from weakref import WeakKeyDictionary, WeakMethod
 
 from twinkle.patch import apply_patch
 from twinkle.utils import DeviceMesh, get_logger
@@ -445,6 +446,17 @@ def local_sdpa_attn(module: torch.nn.Module, query_states, key_states, value_sta
         logger.info_once(f'[SequenceParallel] attention wrappers registered: {list(FLASH_ATTENTION_IMPLS)} -> '
                          'local_flash_attn, sdpa -> local_sdpa_attn (SP-aware)')
 
+    @property
+    def causal_mask_func(self):
+        func = self._causal_mask_func
+        return func() if isinstance(func, WeakMethod) else func
+
+    @causal_mask_func.setter
+    def causal_mask_func(self, func):
+        # The global weak-key registry owns this context. A bound method must
+        # not keep the owning model (and therefore its weak key) alive.
+        self._causal_mask_func = WeakMethod(func) if getattr(func, '__self__', None) is not None else func
+
     def _prepare_forward_hook(self, base_model: torch.nn.Module):
 
         def pre_forward_split_hook(_self, args, kwargs):
@@ -458,7 +470,8 @@ def pre_forward_split_hook(_self, args, kwargs):
                 position_ids = kwargs['position_ids']
                 real_position_ids = self._extract_real_position_ids(position_ids)
                 cache_position = kwargs.get('cache_position', None)
-                specs = tuple(spec for spec in self.input_key_registry if spec.stage == 'forward')
+                specs = tuple(spec for spec in self.input_key_registry
+                              if spec.pad_at == 'forward' or spec.split_at == 'forward')
                 common_names = ('input_ids', 'inputs_embeds', 'position_ids', 'attention_mask')
                 keys = SPKeyCollection.from_mapping(kwargs, specs, common_names=common_names)
                 self.pad_and_split_inputs(keys, boundary='forward', real_position_ids=real_position_ids,
@@ -831,10 +844,11 @@ def pad_and_split_inputs(self,
             else:
                 attention_mask = attn_item.value
             local_cache_position = torch.arange(0, padded_model_inputs.shape[1], device=padded_model_inputs.device)
-            if (cache_position is None and self.causal_mask_func is not None
+            causal_mask_func = self.causal_mask_func
+            if (cache_position is None and causal_mask_func is not None
                     and not is_flash_attention_implementation(self.attn_implementation)):
-                attention_mask = self.causal_mask_func(attention_mask, padded_model_inputs.to(self.model_dtype),
-                                                       local_cache_position, None, None)
+                attention_mask = causal_mask_func(attention_mask, padded_model_inputs.to(self.model_dtype),
+                                                  local_cache_position, None, None)
                 attn_item.value = attention_mask
 
         # Split after padding so every field uses the same padded sequence frame.
@@ -891,7 +905,8 @@ def prepare_inputs(self, inputs):
         self.extra_kwargs['padding_free'] = padding_free
         if input_ids is not None:
             self.extra_kwargs['input_ids'] = input_ids.clone()
-        specs = tuple(spec for spec in self.input_key_registry if spec.stage == 'auxiliary')
+        specs = tuple(spec for spec in self.input_key_registry
+                      if spec.pad_at == 'prepare' or spec.split_at == 'prepare')
         keys = SPKeyCollection.from_mapping(inputs, specs, common_names=('labels', 'completion_mask', 'loss_scale'))
         self.pad_and_split_inputs(keys, boundary='prepare', real_position_ids=real_position_ids)
         for item in keys.values():
@@ -935,6 +950,9 @@ def __init__(
         # Per-model SP context. The module-level ``sequence_parallel`` singleton
         # is only kept as a compat mirror of the most recently prepared context.
         self._ctx = SequenceParallel()
+        # Compatibility declarations registered directly through the public API.
+        # Keep these separate from rules consumed from model/loss providers.
+        self.field_declarations = SimpleNamespace(input_key_specs=(), output_key_specs=())
 
     def _get_tokenizer(self) -> Optional[PreTrainedTokenizer]:
         if self._tokenizer is not None:
@@ -1027,39 +1045,43 @@ def pool_last_valid_token(self, features, inputs):
     def register_input_key(self, spec: DataKeySpec) -> None:
         """Declare a token-aligned input on this model only.
 
-        Auxiliary inputs are kept out of the model; stage=forward inputs are
-        sliced in its pre-hook. need_gather selects the loss-side collection.
+        pad_at/split_at select execution boundaries; need_gather selects
+        loss-side collection. The model wrapper owns input routing and merges
+        these explicit declarations into the processor schema.
         """
         self._ctx.input_key_registry.register(spec)
+        if spec not in self.field_declarations.input_key_specs:
+            self.field_declarations.input_key_specs += (spec,)
 
     @property
     def input_key_specs(self) -> Tuple[DataKeySpec, ...]:
-        """Read-only snapshot used by processors to preserve registered keys."""
+        """Read-only snapshot of this context's execution rules."""
         return tuple(self._ctx.input_key_registry)
 
     def register_output_key(self, spec: DataKeySpec) -> None:
         """Register a model-produced sequence output to collect before loss."""
         self._ctx.output_key_registry.register(spec)
+        if spec not in self.field_declarations.output_key_specs:
+            self.field_declarations.output_key_specs += (spec,)
 
     @property
     def output_key_specs(self) -> Tuple[DataKeySpec, ...]:
         return tuple(self._ctx.output_key_registry)
 
+    def register_key_specs(self, input_key_specs, output_key_specs):
+        """Consume model-resolved field definitions for SP execution."""
+        for spec in input_key_specs:
+            if spec.name not in self._ctx.input_key_registry._FORWARD_KEYS:
+                self._ctx.input_key_registry.register(spec)
+        for spec in output_key_specs:
+            self._ctx.output_key_registry.register(spec)
+
     def register_loss_key_specs(self, loss) -> None:
-        """Register sequence input/output contracts declared by a loss."""
-        for attr, register, expected_side in (
-                ('sp_input_key_specs', self.register_input_key, 'input'),
-                ('sp_output_key_specs', self.register_output_key, 'output')):
-            specs = getattr(loss, attr, ()) or ()
-            if isinstance(specs, DataKeySpec):
-                specs = (specs,)
-            for spec in specs:
-                if not isinstance(spec, DataKeySpec):
-                    raise TypeError(f'{type(loss).__name__}.{attr} must contain DataKeySpec values')
-                if spec.side != expected_side:
-                    raise ValueError(f'{type(loss).__name__}.{attr} contains {spec.side!r} key {spec.name!r}; '
-                                     f'expected {expected_side!r}')
-                register(spec)
+        """Compatibility entry point; field ownership remains with the provider."""
+        from twinkle.data_fields import declared_key_specs
+
+        self.register_key_specs(tuple(declared_key_specs(loss, 'input')),
+                                tuple(declared_key_specs(loss, 'output')))
 
     @property
     def model_sp_adapters(self) -> Tuple[ModelSPAdapter, ...]:
@@ -1120,19 +1142,6 @@ def _gather_key(self, value, spec, real_position_ids):
             gathered = self._ctx.gather(value, dim=dim, position_ids=real_position_ids)
         return self._trim_gathered_sequence_padding(gathered, real_position_ids, dim=dim)
 
-    def _gather_completion_mask(self, inputs: Dict[str, Any], real_position_ids) -> None:
-        """Gather ``completion_mask`` in place, mirroring the labels gather.
-
-        Deliberately not routed through :class:`GatherLoss`: the mask carries no
-        gradient, and reusing that autograd Function would attach a second backward
-        path to whichever tensor were passed alongside it, double-scaling its grad.
-        """
-        mask = inputs.get('completion_mask')
-        if mask is None or not torch.is_tensor(mask) or mask.dim() < 2:
-            return
-        gathered = self._ctx.gather(mask, dim=1, position_ids=real_position_ids)
-        inputs['completion_mask'] = self._trim_gathered_sequence_padding(gathered, real_position_ids)
-
     def gather_loss_tensors(
         self,
         inputs: Dict[str, Any],
@@ -1162,7 +1171,7 @@ def gather_loss_tensors(
         unknown_keys = required_keys - output_specs.keys()
         if unknown_keys:
             raise KeyError(f'Loss requested SP gather for unregistered output key(s) {sorted(unknown_keys)}. '
-                           'Register a sequence-aligned DataKeySpec with SequenceParallelStrategy.register_output_key().')
+                           'Declare a sequence-aligned DataKeySpec in the model or loss output_key_specs.')
         missing_keys = {key for key in required_keys if outputs.get(key) is None}
         if missing_keys:
             raise KeyError(f'Loss declared SP gather key(s) {sorted(missing_keys)}, but the model output is missing them.')
@@ -1171,7 +1180,6 @@ def gather_loss_tensors(
             if value is not None and (spec.need_gather or spec.name in required_keys
                                       or (spec.name == 'logits' and fused_hidden)):
                 outputs[spec.name] = self._gather_key(value, spec, real_position_ids)
-        self._gather_completion_mask(inputs, real_position_ids)
         return inputs, outputs
 
     def wrap_model(self, model, optimizer=None):
diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py b/src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py
index 91722af04..a00f8f75f 100644
--- a/src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py
@@ -92,7 +92,7 @@ def gather_router_logits(module, args, kwargs, output):
             router_logits = getattr(output, 'router_logits', None)
             if router_logits is None:
                 return output
-            attention_mask = kwargs['attention_mask']
+            attention_mask = kwargs.get('attention_mask')
             batch_size = 1 if attention_mask is None else attention_mask.shape[0]
             assert router_logits[0].shape[0] % batch_size == 0
             local_seq_len = router_logits[0].shape[0] // batch_size
diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/data_keys.py b/src/twinkle/model/transformers/strategy/sequence_parallel/data_keys.py
index 704fbf1e6..46c91b3a9 100644
--- a/src/twinkle/model/transformers/strategy/sequence_parallel/data_keys.py
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/data_keys.py
@@ -1,121 +1,14 @@
-"""Per-model contracts for sequence-aligned inputs and outputs.
-
-Each key declares whether and where it is padded, split, gathered, or restored
-after packed-sequence processing. Forward and auxiliary keys share the same
-executor but run at different lifecycle boundaries. Loss-scale weights are
-already aligned with Twinkle's shifted labels, so the production rule does not
-roll them a second time.
-"""
+"""SP execution collections; common field declarations live in twinkle.data_fields."""
 from dataclasses import dataclass, field
 from typing import Any, Dict, Iterable, Iterator, Mapping, Optional
 
-
-_BOUNDARIES = frozenset({'prepare', 'forward', 'model_internal', 'router_replay'})
-
-
-@dataclass(frozen=True)
-class DataKeySpec:
-    name: str
-    split_dim: int = -1
-    pad_value: int | float | str = 0
-    need_roll: bool = False
-    side: str = 'input'
-    stage: str = 'auxiliary'
-    need_split: bool = True
-    need_gather: bool = False
-    autograd: bool = False
-    owner: str = 'generic'
-    need_pad: Optional[bool] = None
-    pad_at: Optional[str] = 'auto'
-    split_at: Optional[str] = 'auto'
-    need_unpack: Optional[bool] = None
-
-    def __post_init__(self):
-        if not isinstance(self.name, str) or not self.name:
-            raise ValueError('An SP input key must have a non-empty name')
-        if not isinstance(self.split_dim, int):
-            raise TypeError('split_dim must be an integer')
-        if self.side not in ('input', 'output'):
-            raise ValueError('side must be input or output')
-        if self.stage not in ('forward', 'auxiliary', 'model_internal', 'router_replay', 'output'):
-            raise ValueError('Invalid SP key stage')
-        if (self.side == 'output') != (self.stage == 'output'):
-            raise ValueError('Output keys must use stage=output')
-        if self.owner not in ('generic', 'task', 'model', 'loss'):
-            raise ValueError('Invalid SP key owner')
-        if isinstance(self.pad_value, str) and (self.name != 'input_ids' or self.pad_value != 'tokenizer'):
-            raise ValueError('Only input_ids can use the tokenizer padding sentinel')
-        if self.side == 'input' and self.need_gather and not self.need_split:
-            raise ValueError('Only sharded input keys can request gathering')
-        if self.stage == 'forward' and self.need_gather:
-            raise ValueError('Model-forward keys cannot request loss gathering')
-        default_boundary = {
-            'auxiliary': 'prepare',
-            'forward': 'forward',
-            'model_internal': 'model_internal',
-            'router_replay': 'router_replay',
-            'output': None,
-        }[self.stage]
-        if self.need_pad is None:
-            object.__setattr__(self, 'need_pad', self.need_split and self.side == 'input')
-        if self.pad_at == 'auto':
-            object.__setattr__(self, 'pad_at', default_boundary if self.need_pad else None)
-        elif self.pad_at is not None and self.pad_at not in _BOUNDARIES:
-            raise ValueError(f'Invalid SP padding boundary: {self.pad_at!r}')
-        elif (self.pad_at is None) == bool(self.need_pad):
-            raise ValueError('need_pad and pad_at specify conflicting SP padding behavior')
-        if self.pad_at != 'auto':
-            object.__setattr__(self, 'need_pad', self.pad_at is not None)
-        if self.split_at == 'auto':
-            object.__setattr__(self, 'split_at', default_boundary if self.need_split else None)
-        elif self.split_at is not None and self.split_at not in _BOUNDARIES:
-            raise ValueError(f'Invalid SP split boundary: {self.split_at!r}')
-        elif (self.split_at is None) == bool(self.need_split):
-            raise ValueError('need_split and split_at specify conflicting SP split behavior')
-        if self.split_at != 'auto':
-            object.__setattr__(self, 'need_split', self.split_at is not None)
-        if self.side == 'output' and (self.pad_at is not None or self.split_at is not None):
-            raise ValueError('Output keys cannot request input padding or splitting')
-        if self.need_unpack is None:
-            # Preserve the historical packed restoration behavior unless a
-            # built-in rule explicitly opts out.
-            object.__setattr__(self, 'need_unpack', self.side == 'output' or self.need_gather or not self.need_split)
-
-    def sequence_dim(self, ndim):
-        if not -ndim <= self.split_dim < ndim:
-            raise ValueError(f'Invalid sequence axis {self.split_dim} for SP key {self.name!r}')
-        dim = self.split_dim % ndim
-        if dim == 0:
-            raise ValueError(f'SP key {self.name!r} cannot split or gather the batch axis')
-        return dim
+# Preserve the existing public imports for model adapters and external callers.
+from twinkle.data_fields import (BUILTIN_INPUT_KEYS, BUILTIN_OUTPUT_KEYS, PROCESSOR_PASSTHROUGH_KEYS,
+                                 DataKeySpec, InputKeyRegistry, OutputKeyRegistry)
 
 
-BUILTIN_INPUT_KEYS = (
-    DataKeySpec('input_ids', split_dim=1, pad_value='tokenizer', stage='forward'),
-    DataKeySpec('inputs_embeds', split_dim=1, stage='forward'),
-    DataKeySpec('position_ids', pad_value=-1, stage='forward'),
-    DataKeySpec('attention_mask', split_dim=1, stage='forward', need_pad=True, need_split=False),
-    DataKeySpec('labels', pad_value=-100, need_gather=True, need_unpack=True),
-    # These are not needed by the model forward. Keep them global through SP;
-    # packed sequence restoration is a separate action from sharding/gathering.
-    DataKeySpec('completion_mask', need_pad=False, need_split=False, need_unpack=True),
-    DataKeySpec('loss_scale', need_pad=False, need_split=False, need_unpack=True),
-)
-
-# These are processor/model control arguments, not sequence data contracts.
-# Keep them out of DataKeySpec and preserve them through the Transformers input
-# conversion path explicitly instead.
-PROCESSOR_PASSTHROUGH_KEYS = ('cache_position', 'padding_free')
-
-BUILTIN_OUTPUT_KEYS = (
-    DataKeySpec('logps', split_dim=1, side='output', stage='output',
-                need_split=False, need_gather=True, autograd=True, need_unpack=True),
-    DataKeySpec('entropies', split_dim=1, side='output', stage='output',
-                need_split=False, need_gather=True, autograd=True, need_unpack=True),
-    # Real logits are opt-in; fused CE explicitly gathers this hidden-state slot.
-    DataKeySpec('logits', split_dim=1, side='output', stage='output',
-                need_split=False, autograd=True, need_unpack=True),
-)
+__all__ = ['DataKeySpec', 'InputKeyRegistry', 'OutputKeyRegistry', 'BUILTIN_INPUT_KEYS',
+           'BUILTIN_OUTPUT_KEYS', 'PROCESSOR_PASSTHROUGH_KEYS', 'SPKeyValue', 'SPKeyCollection']
 
 
 @dataclass
@@ -165,59 +58,3 @@ def from_mapping(cls, values: Mapping[str, Any], specs: Iterable[DataKeySpec], c
             if spec.name in values:
                 collection.add(SPKeyValue(spec, values[spec.name]), common=spec.name in common_names)
         return collection
-
-
-class InputKeyRegistry:
-    """Own declarations locally; reject accidental replacement of existing rules.
-
-    Auxiliary keys are processed at their declared boundary (usually
-    preprocessing). Custom model-forward keys remain global until the forward
-    hook, so every key is sharded at most once.
-    """
-
-    _FORWARD_KEYS = frozenset({'input_ids', 'inputs_embeds', 'position_ids', 'attention_mask'})
-    _SPECIAL_KEYS = frozenset(PROCESSOR_PASSTHROUGH_KEYS)
-
-    def __init__(self):
-        self._specs = {spec.name: spec for spec in BUILTIN_INPUT_KEYS}
-
-    def register(self, spec: DataKeySpec):
-        if not isinstance(spec, DataKeySpec):
-            raise TypeError('Expected a DataKeySpec')
-        if spec.side != 'input' or spec.stage not in ('auxiliary', 'forward'):
-            raise ValueError('Register only auxiliary or model-forward input keys here')
-        if spec.name in self._FORWARD_KEYS:
-            raise ValueError(f'{spec.name!r} is owned by the model forward hook')
-        if spec.name in self._SPECIAL_KEYS:
-            raise ValueError(f'{spec.name!r} is a special processor/model argument, not an SP data key')
-        existing = self._specs.get(spec.name)
-        if existing is not None and existing != spec:
-            raise ValueError(f'Conflicting SP input rule for {spec.name!r}')
-        self._specs[spec.name] = spec
-
-    def __iter__(self):
-        return iter(self._specs.values())
-
-    def get(self, name):
-        return self._specs.get(name)
-
-
-class OutputKeyRegistry:
-    def __init__(self):
-        self._specs = {spec.name: spec for spec in BUILTIN_OUTPUT_KEYS}
-
-    def register(self, spec):
-        if not isinstance(spec, DataKeySpec):
-            raise TypeError('Expected a DataKeySpec')
-        if spec.side != 'output':
-            raise ValueError('Expected an output key rule')
-        existing = self._specs.get(spec.name)
-        if existing is not None and existing != spec:
-            raise ValueError(f'Conflicting SP output rule for {spec.name!r}')
-        self._specs[spec.name] = spec
-
-    def __iter__(self):
-        return iter(self._specs.values())
-
-    def get(self, name):
-        return self._specs.get(name)
diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py
index 108a4bf0e..18824302b 100644
--- a/src/twinkle/model/transformers/transformers.py
+++ b/src/twinkle/model/transformers/transformers.py
@@ -56,6 +56,8 @@
 from twinkle.utils.grad_clip import normalize_and_clip_grad_norm
 from twinkle.utils.transformers_utils import filter_from_config_kwargs
 
+from twinkle.data_fields import partition_model_inputs, resolve_key_specs
+
 from .task_outputs import postprocess_task_outputs
 
 logger = get_logger()
@@ -844,15 +846,22 @@ def _prepare_sp_inputs(self, inputs, processor, task='causal_lm'):
             payload['labels'] = sequence_labels
         return payload
 
-    def _restore_loss_inputs(self, inputs, labels, loss_scale=None, channel=None, completion_mask=None):
-        """Put the loss-related tensors popped before the model call back onto ``inputs``."""
-        inputs['labels'] = labels
-        if loss_scale is not None:
-            inputs['loss_scale'] = loss_scale
-        if channel is not None:
-            inputs['channel'] = channel
-        if completion_mask is not None:
-            inputs['completion_mask'] = completion_mask
+    def _get_field_specs(self, loss_instance):
+        """Resolve model/loss fields even when sequence parallelism is disabled."""
+        from twinkle.utils.transformers_utils import get_llm_model
+        from .strategy.sequence_parallel.adapters import adapters_for_model
+
+        model = self.strategy.unwrap_model(self.model)
+        text_model = get_llm_model(model)
+        models = (model,) if text_model is model else (model, text_model)
+        adapters = tuple(adapter for candidate in models for adapter in adapters_for_model(candidate))
+        # Direct registration is a legacy declaration interface. Merge only
+        # those explicit declarations, not SP's derived execution registry.
+        direct_fields = self.sp_strategy.field_declarations if self.sp_strategy is not None else None
+        specs = resolve_key_specs(*models, *adapters, loss_instance, direct_fields)
+        if self.sp_strategy is not None:
+            self.sp_strategy.register_key_specs(*specs)
+        return specs
 
     def _postprocess_forward_outputs(self, outputs, *, labels, task, temperature, sampling_masks, return_logits,
                                      loss_instance, replay=None):
@@ -894,15 +903,16 @@ def _postprocess_forward_outputs(self, outputs, *, labels, task, temperature, sa
             outputs['logits'] = None
 
     def _finalize_forward_outputs(self, inputs, outputs, *, processor, task, optimizer_config, is_training,
-                                  forward_kwargs, loss_instance, return_logits, recorded_routing):
+                                  forward_kwargs, loss_instance, return_logits, recorded_routing,
+                                  input_key_specs, output_key_specs):
         """Shared forward/forward_only tail: SP postprocess, unpack, status recording, and the return copy."""
         inputs, outputs = postprocess_task_outputs(
             inputs, outputs, sp_strategy=self.sp_strategy, task=task,
             sp_gather_keys=getattr(loss_instance, 'sp_gather_keys', ()))
         inputs, outputs = processor.unpack_packed_sequences(
             inputs, outputs, task=task,
-            input_key_specs=getattr(self.sp_strategy, 'input_key_specs', None),
-            output_key_specs=getattr(self.sp_strategy, 'output_key_specs', None))
+            input_key_specs=input_key_specs,
+            output_key_specs=output_key_specs)
         status = optimizer_config.train_status if is_training else optimizer_config.eval_status
         status.inputs = inputs
         status.outputs = outputs
@@ -957,7 +967,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
             if getattr(self, '_enable_sp', False) or cp_world_size > 1:
                 raise ValueError('sampling replay does not support sequence or context parallelism')
         assert isinstance(processor, InputProcessor), 'Set a correct `InputProcessor` before forwarding'
-        input_key_specs = getattr(self.sp_strategy, 'input_key_specs', None)
+        input_key_specs, output_key_specs = self._get_field_specs(loss_instance)
         inputs: Dict[str, Any] = processor(
             inputs,
             input_key_specs=input_key_specs,
@@ -967,12 +977,8 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
             gdn_padding_free_enabled=self._gdn_padding_free_enabled,
         )
         inputs = self._prepare_sp_inputs(inputs, processor, task=task)
-        labels: torch.Tensor = inputs.pop('labels', None)
-        loss_scale = inputs.pop('loss_scale', None)
-        channel = inputs.pop('channel', None)
-        sp_auxiliary_inputs = processor.pop_sp_auxiliary_inputs(inputs, input_key_specs=input_key_specs)
-        # Not a model argument; the loss reads it back off `inputs` further down.
-        completion_mask = inputs.pop('completion_mask', None)
+        inputs, loss_inputs = partition_model_inputs(inputs, input_key_specs)
+        labels = loss_inputs.get('labels')
         replay_metadata = replay_loss_mask = replay_masked_labels = None
         if enable_sampling_replay:
             replay_loss_mask, replay_masked_labels, replay_metadata = _prepare_sampling_replay(
@@ -1010,8 +1016,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
 
         recorded_routing = rr.finish()
 
-        self._restore_loss_inputs(inputs, labels, loss_scale, channel, completion_mask)
-        inputs.update(sp_auxiliary_inputs)
+        inputs.update(loss_inputs)
         self._postprocess_forward_outputs(
             outputs,
             labels=labels,
@@ -1033,6 +1038,7 @@ def forward(self, *, inputs: Union[InputFeature, List[InputFeature], List[Trajec
             loss_instance=loss_instance,
             return_logits=return_logits,
             recorded_routing=recorded_routing,
+            input_key_specs=input_key_specs, output_key_specs=output_key_specs,
         )
 
     # lazy_collect follows the global default (True): a caller that reads the output itself calls
@@ -1178,7 +1184,7 @@ def _forward_only_single(self, *, inputs, optimizer_config, **kwargs):
                 if getattr(self, '_enable_sp', False) or cp_world_size > 1:
                     raise ValueError('sampling replay does not support sequence or context parallelism')
             loss_require_values = getattr(loss_instance, 'require_values', False)
-            input_key_specs = getattr(self.sp_strategy, 'input_key_specs', None)
+            input_key_specs, output_key_specs = self._get_field_specs(loss_instance)
             inputs: Dict[str, Any] = processor(
                 inputs,
                 input_key_specs=input_key_specs,
@@ -1188,12 +1194,8 @@ def _forward_only_single(self, *, inputs, optimizer_config, **kwargs):
                 gdn_padding_free_enabled=self._gdn_padding_free_enabled,
             )
             inputs = self._prepare_sp_inputs(inputs, processor, task=task)
-            labels = inputs.pop('labels', None)
-            loss_scale = inputs.pop('loss_scale', None)
-            channel = inputs.pop('channel', None)
-            sp_auxiliary_inputs = processor.pop_sp_auxiliary_inputs(inputs, input_key_specs=input_key_specs)
-            # Not a model argument; the loss reads it back off `inputs` further down.
-            completion_mask = inputs.pop('completion_mask', None)
+            inputs, loss_inputs = partition_model_inputs(inputs, input_key_specs)
+            labels = loss_inputs.get('labels')
             replay_metadata = replay_loss_mask = replay_masked_labels = None
             if enable_sampling_replay:
                 packed_position_ids = processor._is_packed_position_ids(inputs.get('position_ids'))
@@ -1230,8 +1232,7 @@ def _forward_only_single(self, *, inputs, optimizer_config, **kwargs):
 
             recorded_routing = rr.finish()
 
-            self._restore_loss_inputs(inputs, labels, loss_scale, channel, completion_mask)
-            inputs.update(sp_auxiliary_inputs)
+            inputs.update(loss_inputs)
             self._postprocess_forward_outputs(
                 outputs,
                 labels=labels,
@@ -1253,6 +1254,7 @@ def _forward_only_single(self, *, inputs, optimizer_config, **kwargs):
                 loss_instance=loss_instance,
                 return_logits=return_logits,
                 recorded_routing=recorded_routing,
+                input_key_specs=input_key_specs, output_key_specs=output_key_specs,
             )
 
     @remote_function(dispatch='slice_dp', collect=collect_dp_and_flatten, lazy_collect=False)
diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py
index 41f8947c2..ce8cba69b 100644
--- a/src/twinkle/processor/base.py
+++ b/src/twinkle/processor/base.py
@@ -105,7 +105,7 @@ def prepare_outputs(self, inputs: list[InputFeature], **kwargs) -> list[InputFea
 
     def prepare_inputs(self, inputs: Union[List[InputFeature], InputFeature], input_key_specs=None,
                        **kwargs) -> List[InputFeature]:
-        from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS
+        from twinkle.data_fields import BUILTIN_INPUT_KEYS
 
         specs = BUILTIN_INPUT_KEYS if input_key_specs is None else input_key_specs
         registered_keys = {spec.name for spec in specs}
@@ -147,6 +147,8 @@ def to_tensor(_input):
                         value = torch.cat(value, dim=0)
                 if isinstance(value, torch.Tensor):
                     value = value.to(Platform.get_local_device())
+                    if key == 'labels' and value.dim() == 0:
+                        value = value.reshape(1)
                     if value.dim() == 1:
                         value = value.unsqueeze(0)
                 _input[key] = value
@@ -518,17 +520,16 @@ def unpack_packed_sequences(
         """
         if task in ("embedding", "generative_reranker", "seq_cls", "reranker"):
             return inputs, outputs
-        labels = inputs.get("labels")
         position_ids = inputs.get("position_ids")
 
-        if labels is None or position_ids is None:
+        if position_ids is None:
             return inputs, outputs
         if not self._is_packed_position_ids(position_ids):
             return inputs, outputs
 
         from copy import copy
 
-        from twinkle.model.transformers.strategy.sequence_parallel.data_keys import (
+        from twinkle.data_fields import (
             BUILTIN_INPUT_KEYS, BUILTIN_OUTPUT_KEYS,
         )
 
@@ -571,7 +572,7 @@ def unpack_inputs(self, inputs: list[dict[str, Any]], task: str = "causal_lm") -
     @staticmethod
     def to_transformers_dict(inputs: List[InputFeature], input_key_specs=None, **kwargs) -> List[InputFeature]:
         import torch
-        from twinkle.model.transformers.strategy.sequence_parallel.data_keys import PROCESSOR_PASSTHROUGH_KEYS
+        from twinkle.data_fields import PROCESSOR_PASSTHROUGH_KEYS
 
         registered_keys = [spec.name for spec in input_key_specs] if input_key_specs is not None else []
         results = []
@@ -609,28 +610,19 @@ def to_transformers_dict(inputs: List[InputFeature], input_key_specs=None, **kwa
             results.append(InputFeature(**output))
         return results
 
-    @staticmethod
-    def pop_sp_auxiliary_inputs(inputs, *, input_key_specs=None):
-        """Keep loss-only token fields out of Hugging Face model.forward."""
-        from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS
-
-        specs = BUILTIN_INPUT_KEYS if input_key_specs is None else input_key_specs
-        return {spec.name: inputs.pop(spec.name) for spec in specs
-                if spec.stage == 'auxiliary' and spec.name not in ('labels', 'loss_scale') and spec.name in inputs}
-
     def _collate_sequence_key(self, values, spec, padding_free):
         samples = []
         seq_dim = None
         for value in values:
             if not torch.is_tensor(value):
-                raise TypeError(f'Registered SP input {spec.name!r} must be a tensor before collation')
+                raise TypeError(f'Registered sequence input {spec.name!r} must be a tensor before collation')
             if value.dim() == 1:
                 value = value.unsqueeze(0)
             dim = spec.sequence_dim(value.dim())
             if dim == 0 or value.shape[0] != 1:
-                raise ValueError(f'SP input {spec.name!r} requires one sample with a leading batch dimension')
+                raise ValueError(f'Sequence input {spec.name!r} requires one sample with a leading batch dimension')
             if seq_dim is not None and seq_dim != dim - 1:
-                raise ValueError(f'Inconsistent sequence axis for SP input {spec.name!r}')
+                raise ValueError(f'Inconsistent sequence axis for input {spec.name!r}')
             seq_dim = dim - 1
             samples.append(value.squeeze(0))
         if padding_free:
@@ -691,7 +683,7 @@ def _fill_optional_sequence_fields(self, batch: list[InputFeature]) -> None:
 
     def _collate_macro_batch(self, inputs: List[InputFeature], input_key_specs=None) -> InputFeature:
         import torch
-        from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS
+        from twinkle.data_fields import BUILTIN_INPUT_KEYS
 
         builtin_names = {spec.name for spec in BUILTIN_INPUT_KEYS}
         specs = BUILTIN_INPUT_KEYS if input_key_specs is None else input_key_specs
diff --git a/tests/transformers/test_sp_context_isolation.py b/tests/transformers/test_sp_context_isolation.py
index 0a3078832..e4db8a2de 100644
--- a/tests/transformers/test_sp_context_isolation.py
+++ b/tests/transformers/test_sp_context_isolation.py
@@ -139,3 +139,61 @@ def gather(self, tensor, **kwargs):
     assert rr.get_router_replay_data(models[0]).eq(10).all()
     assert rr.get_router_replay_data(models[1]).eq(20).all()
     assert rr.get_router_replay_data(models[2]).eq(0).all()
+
+
+@pytest.mark.parametrize('rank', [0, 1])
+def test_router_replay_shards_mask_with_routing_indices(monkeypatch, rank):
+    from twinkle.model.transformers.moe.router_replay import RouterReplay
+
+    model = torch.nn.Linear(1, 1)
+    ctx = sp.SequenceParallel()
+    ctx.world_size = ctx.sp_world_size = 2
+    ctx.rp_world_size = 1
+    ctx._sp_rank = rank
+    ctx.extra_kwargs['position_ids'] = torch.arange(5).unsqueeze(0)
+    monkeypatch.setattr(RouterReplay, '_find_moe_blocks_with_names', lambda _: [('block', model)])
+    state = SimpleNamespace(target_indices=None, target_mask=None)
+    monkeypatch.setattr(RouterReplay, '_registry', {'block': state})
+    mask = torch.tensor([[True, False, True, False, True]])
+    routing = torch.arange(5).reshape(1, 5, 1, 1)
+    from twinkle import Platform
+    monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu'))
+    RouterReplay.set_router_replay_data(routing, model, mask, sp_context=ctx)
+    expected_indices = torch.tensor([0, 1, 2, 3, 4, 0]).chunk(2)[rank]
+    expected_mask = torch.tensor([True, False, True, False, True, False]).chunk(2)[rank]
+    torch.testing.assert_close(state.target_indices[:, 0], expected_indices)
+    torch.testing.assert_close(state.target_mask, expected_mask)
+
+
+@pytest.mark.parametrize('nested', [False, True])
+def test_prepare_causal_mask_method_does_not_retain_weak_registry_owner(monkeypatch, nested):
+    class Backbone(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.weight = torch.nn.Parameter(torch.ones(1))
+            self.config = SimpleNamespace(num_attention_heads=2, _attn_implementation='sdpa')
+
+        def _update_causal_mask(self, mask, *args):
+            return mask + self.weight
+
+    model = Backbone()
+    if nested:
+        wrapper = torch.nn.Module()
+        wrapper.language_model = model
+        wrapper.config = model.config
+        model = wrapper
+        del wrapper
+    refs = [weakref.ref(module) for module in model.modules()]
+    ctx = sp.SequenceParallel()
+    monkeypatch.setattr(sp, 'get_llm_model', lambda value: value)
+    monkeypatch.setattr(sp.SequenceParallel, '_global_inited', True)
+    monkeypatch.setattr(sp.SequenceParallel, '_init_device_mesh', lambda *args: None)
+    monkeypatch.setattr(sp, 'sequence_parallel', sp.SequenceParallel())
+    ctx.prepare(2, model, SimpleNamespace(pad_token_id=0))
+    torch.testing.assert_close(ctx.causal_mask_func(torch.zeros(1)), torch.ones(1))
+    # Remove the compatibility mirror so only the weak registry owns the context.
+    sp.sequence_parallel.__dict__.clear()
+    del model
+    gc.collect()
+    assert all(ref() is None for ref in refs)
+    assert ctx.causal_mask_func is None
diff --git a/tests/transformers/test_sp_custom_loss_cpu.py b/tests/transformers/test_sp_custom_loss_cpu.py
index 4c30aca0c..7f0541ba9 100644
--- a/tests/transformers/test_sp_custom_loss_cpu.py
+++ b/tests/transformers/test_sp_custom_loss_cpu.py
@@ -22,8 +22,8 @@ class TokenRegressionLoss(Loss):
     """Example plugin loss with custom sequence input and output keys."""
 
     require_logps = False
-    sp_input_key_specs = (DataKeySpec('token_targets', split_dim=1, pad_value=0, need_gather=True, owner='loss'),)
-    sp_output_key_specs = (DataKeySpec('token_scores', split_dim=1, side='output', stage='output',
+    input_key_specs = (DataKeySpec('token_targets', split_dim=1, pad_value=0, owner='loss'),)
+    output_key_specs = (DataKeySpec('token_scores', split_dim=1, side='output', stage='output',
                                        need_split=False, need_gather=True, autograd=True, owner='loss'),)
 
     def __call__(self, inputs, outputs, **kwargs):
@@ -107,7 +107,7 @@ def test_custom_token_regression_loss_matches_full_sequence_loss_and_gradient():
     mp.spawn(_worker, args=(port,), nprocs=2, join=True)
 
 
-def _e2e_worker(rank, port):
+def _e2e_worker(rank, port, sp_size):
     os.environ['RANK'] = str(rank)
     os.environ['WORLD_SIZE'] = '2'
     torch.set_num_threads(1)
@@ -124,7 +124,7 @@ def _e2e_worker(rank, port):
         torch.manual_seed(1234)
         baseline = LlamaForCausalLM(config).train()
 
-        mesh = DeviceMesh.from_sizes(dp_size=2, ulysses_size=2, device_type='cpu')
+        mesh = DeviceMesh.from_sizes(dp_size=2, ulysses_size=sp_size, device_type='cpu')
         model_module = importlib.import_module('twinkle.model.transformers.transformers')
         original_template = model_module.Template
         model_module.Template = lambda *_args, **_kwargs: SimpleNamespace(is_mm=False)
@@ -140,7 +140,7 @@ def _e2e_worker(rank, port):
         training_model.strategy = _NoopTrainStrategy()
         # This local CPU harness supplies worker mesh state explicitly.
         training_model.device_mesh = mesh
-        training_model._enable_sp = True
+        training_model._enable_sp = sp_size > 1
         model_forward = training_model.model.forward
 
         def forward_with_token_scores(*args, **kwargs):
@@ -152,17 +152,18 @@ def forward_with_token_scores(*args, **kwargs):
         training_model.model.load_state_dict(baseline.state_dict())
         training_model.model.gradient_checkpointing_disable()
         training_model.model.train()
-        strategy = SequenceParallelStrategy(mesh, model=training_model.model)
-        strategy._get_tokenizer = lambda: SimpleNamespace(pad_token_id=0)
+        strategy = SequenceParallelStrategy(mesh, model=training_model.model) if sp_size > 1 else None
+        if strategy is not None:
+            strategy._get_tokenizer = lambda: SimpleNamespace(pad_token_id=0)
         training_model.sp_strategy = strategy
 
         optimizer_group = training_model.optimizer_group[training_model._get_default_group()]
         loss_fn = TokenRegressionLoss()
-        optimizer_group.loss_instance = loss_fn  # OptimizerGroup applies the NCCL-safe wrapper.
-        assert training_model._enable_sp
-        assert optimizer_group.loss_instance.sp_input_key_specs
+        optimizer_group.loss_instance = loss_fn
+        assert optimizer_group.loss_instance.input_key_specs
         training_model._ensure_sp_strategy()
-        assert 'token_targets' in {spec.name for spec in strategy.input_key_specs}
+        if strategy is not None:
+            assert 'token_targets' in {spec.name for spec in strategy.input_key_specs}
 
         ids = torch.tensor([[1, 2, 3, 4, 5, 6, 7], [2, 3, 4, 5, 6, 7, 8]])
         positions = torch.arange(ids.shape[1]).expand_as(ids)
@@ -205,8 +206,9 @@ def forward_with_token_scores(*args, **kwargs):
 
 
 @pytest.mark.slow
-def test_custom_loss_through_transformers_training_pipeline_and_optimizer_step():
+@pytest.mark.parametrize('sp_size', [1, 2])
+def test_custom_loss_through_transformers_training_pipeline_and_optimizer_step(sp_size):
     with socket.socket() as sock:
         sock.bind(('127.0.0.1', 0))
         port = sock.getsockname()[1]
-    mp.spawn(_e2e_worker, args=(port,), nprocs=2, join=True)
+    mp.spawn(_e2e_worker, args=(port, sp_size), nprocs=2, join=True)
diff --git a/tests/transformers/test_sp_data_keys.py b/tests/transformers/test_sp_data_keys.py
index d468d97cd..9a9213dd7 100644
--- a/tests/transformers/test_sp_data_keys.py
+++ b/tests/transformers/test_sp_data_keys.py
@@ -3,6 +3,7 @@
 import torch
 
 from twinkle.model.transformers.strategy.sequence_parallel import DataKeySpec, SequenceParallelStrategy
+from twinkle.data_fields import partition_model_inputs, resolve_key_specs
 from twinkle.model.transformers.strategy.sequence_parallel.data_keys import SPKeyCollection, SPKeyValue
 
 
@@ -162,7 +163,7 @@ def test_processor_preserves_registered_keys_before_sp_split(monkeypatch, rank):
 
     monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu'))
     strategy = _strategy(rank)
-    strategy.register_input_key(DataKeySpec('token_weights', pad_value=-2))
+    strategy.register_input_key(DataKeySpec('token_weights', pad_value=-2, need_split=True))
     processor = InputProcessor()
     values = torch.arange(5).unsqueeze(0)
     sample = {'input_ids': torch.ones(1, 5, dtype=torch.long),
@@ -207,8 +208,8 @@ def test_registered_multidimensional_key_collation_and_loss_boundary(monkeypatch
     result = strategy.preprocess_inputs(result)
     expected = torch.tensor([[[1, 1, 1], [2, 2, 2]], [[1, 1, 1], [2, 2, 2]]])
     torch.testing.assert_close(result['token_features'], expected)
-    auxiliary = processor.pop_sp_auxiliary_inputs(result, input_key_specs=strategy.input_key_specs)
-    assert set(auxiliary) == {'completion_mask', 'token_features'}
+    result, auxiliary = partition_model_inputs(result, strategy.input_key_specs)
+    assert set(auxiliary) == {'labels', 'completion_mask', 'token_features'}
     assert 'token_features' not in result
 
 
@@ -288,7 +289,7 @@ def test_processor_preserves_custom_keys_without_sp_strategy(monkeypatch, paddin
     expected = (torch.tensor([[[1] * 8, [2] * 8]]) if padding_free else
                 torch.tensor([[[1, 1, 1, -2, -2], [2, 2, 2, -2, -2]], [[1] * 5, [2] * 5]]))
     torch.testing.assert_close(result['token_features'], expected)
-    auxiliary = processor.pop_sp_auxiliary_inputs(result, input_key_specs=specs)
+    result, auxiliary = partition_model_inputs(result, specs)
     torch.testing.assert_close(auxiliary['token_features'], expected)
 
 
@@ -341,3 +342,177 @@ def test_processor_preserves_special_arguments_outside_data_key_registry():
     for converted in (converted_with_sp, converted_without_sp):
         torch.testing.assert_close(converted['cache_position'], sample['cache_position'])
         assert converted['padding_free'] is True
+
+
+@pytest.mark.parametrize('padding_free', [False, True])
+def test_common_loss_fields_collate_and_unpack_without_sp(monkeypatch, padding_free):
+    from types import SimpleNamespace
+    from twinkle import Platform
+    from twinkle.processor import InputProcessor
+
+    monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu'))
+    loss = SimpleNamespace(input_key_specs=(DataKeySpec('token_targets', split_dim=2, pad_value=-9),))
+    input_specs, output_specs = resolve_key_specs(loss)
+    spec = next(spec for spec in input_specs if spec.name == 'token_targets')
+    assert not spec.need_pad and not spec.need_split and not spec.need_gather
+    samples = [dict(input_ids=list(range(length)), position_ids=list(range(length)),
+                    labels=list(range(length)), token_targets=[[[1] * length, [2] * length]])
+               for length in (3, 5)]
+    processor = InputProcessor(padding_free=padding_free)
+    batch = processor(samples, input_key_specs=input_specs)
+    model_inputs, loss_inputs = partition_model_inputs(batch, input_specs)
+    assert 'token_targets' not in model_inputs and 'labels' not in model_inputs
+    assert loss_inputs['token_targets'].shape == ((1, 2, 8) if padding_free else (2, 2, 5))
+    restored, _ = processor.unpack_packed_sequences(
+        batch, task='causal_lm', input_key_specs=input_specs, output_key_specs=output_specs)
+    torch.testing.assert_close(restored['token_targets'], torch.tensor([
+        [[1, 1, 1, -9, -9], [2, 2, 2, -9, -9]], [[1] * 5, [2] * 5]]))
+
+
+@pytest.mark.parametrize('ring_size', [1, 2])
+def test_global_loss_fields_are_never_gathered(monkeypatch, ring_size):
+    strategy = _strategy()
+    strategy._ctx.rp_world_size = ring_size
+    strategy._ctx.sp_world_size = 2 // ring_size
+    strategy._ctx.extra_kwargs['position_ids'] = torch.tensor([[0, 1, 2, 0, 1]])
+
+    def forbidden(*args, **kwargs):
+        raise AssertionError('Global loss fields must not enter SP collectives')
+
+    monkeypatch.setattr(strategy._ctx, 'gather', forbidden)
+    mask = torch.tensor([[0, 1, 1, 0, 1]])
+    scale = torch.tensor([[0., 1., 1., 0., 1.]])
+    inputs, _ = strategy.gather_loss_tensors({'completion_mask': mask, 'loss_scale': scale}, {})
+    assert inputs['completion_mask'] is mask and inputs['loss_scale'] is scale
+
+
+def test_model_input_destination_is_independent_of_split_stage():
+    spec = DataKeySpec('token_context', stage='auxiliary', model_input=True)
+    model_inputs, loss_inputs = partition_model_inputs({'token_context': 1, 'channel': ['a']}, (spec,))
+    assert model_inputs == {'token_context': 1}
+    assert loss_inputs == {'channel': ['a']}
+
+
+
+def test_explicit_parallel_boundaries_opt_in_without_changing_loss_field_defaults():
+    global_field = DataKeySpec('targets')
+    assert not global_field.need_pad and not global_field.need_split
+    local_field = DataKeySpec('local_targets', split_at='prepare', pad_at='prepare')
+    assert local_field.need_split and local_field.need_pad
+    padded_field = DataKeySpec('padded_targets', pad_at='prepare')
+    assert padded_field.need_pad and not padded_field.need_split
+
+
+@pytest.mark.parametrize('rank', [0, 1])
+@pytest.mark.parametrize('stage', ['forward', 'auxiliary'])
+def test_prepare_then_forward_honors_separate_pad_and_split_boundaries(rank, stage):
+    from types import SimpleNamespace
+
+    strategy = _strategy(rank)
+    strategy._ctx.tokenizer = SimpleNamespace(pad_token_id=0)
+    spec = DataKeySpec('token_type_ids', split_dim=1, pad_value=-7, stage=stage,
+                       model_input=True, pad_at='prepare', split_at='forward')
+    strategy.register_input_key(spec)
+
+    class Model(torch.nn.Module):
+        def forward(self, token_type_ids, **kwargs):
+            return token_type_ids
+
+    model = Model()
+    strategy._ctx._prepare_forward_hook(model)
+    positions = torch.arange(5).unsqueeze(0)
+    prepared = strategy.preprocess_inputs(dict(input_ids=positions.clone(), position_ids=positions,
+                                               token_type_ids=positions.clone()))
+    torch.testing.assert_close(prepared['token_type_ids'], torch.tensor([[0, 1, 2, 3, 4, -7]]))
+    expected = torch.tensor([[0, 1, 2]]) if rank == 0 else torch.tensor([[3, 4, -7]])
+    torch.testing.assert_close(model(**prepared), expected)
+
+
+def test_direct_registration_survives_model_schema_resolution_and_processor(monkeypatch):
+    from types import SimpleNamespace
+    from twinkle import Platform
+    from twinkle.model.transformers.transformers import TransformersModel
+    from twinkle.processor import InputProcessor
+
+    monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu'))
+    strategy = _strategy()
+    spec = DataKeySpec('token_weights', pad_value=-2)
+    strategy.register_input_key(spec)
+    output_spec = DataKeySpec('token_scores', side='output', stage='output', need_gather=True)
+    strategy.register_output_key(output_spec)
+    model = torch.nn.Linear(2, 2)
+    wrapper = SimpleNamespace(model=model, strategy=SimpleNamespace(unwrap_model=lambda value: value),
+                              sp_strategy=strategy)
+    input_specs, output_specs = TransformersModel._get_field_specs(wrapper, None)
+    assert spec in input_specs and output_spec in output_specs
+    processor = InputProcessor()
+    result = processor([dict(input_ids=[1, 2, 3], token_weights=[0.1, 0.2, 0.3])],
+                       input_key_specs=input_specs)
+    torch.testing.assert_close(result['token_weights'], torch.tensor([[0.1, 0.2, 0.3]]))
+    _, loss_inputs = partition_model_inputs(result, input_specs)
+    assert 'token_weights' in loss_inputs
+    # Consuming provider schemas must not turn them into direct registrations.
+    strategy.register_key_specs((DataKeySpec('provider_targets'),), ())
+    assert strategy.field_declarations.input_key_specs == (spec,)
+
+
+@pytest.mark.parametrize('label_type', ['tensor', 'numpy', 'python'])
+@pytest.mark.parametrize('framework', ['transformers', 'megatron'])
+def test_scalar_classification_labels_collate_with_common_schema(monkeypatch, label_type, framework):
+    import numpy as np
+    from twinkle import Platform
+    from twinkle.processor import InputProcessor
+
+    monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu'))
+    convert = {'tensor': torch.tensor, 'numpy': np.array, 'python': int}[label_type]
+    batch = [dict(input_ids=[1, 2, 3], position_ids=[0, 1, 2], labels=convert(label)) for label in (1, 0)]
+    result = InputProcessor(framework=framework)(batch)
+    if framework == 'megatron':
+        result = result[0]
+    torch.testing.assert_close(result['labels'], torch.tensor([[1], [0]]))
+    assert result['input_ids'].shape == (2, 3)
+
+
+@pytest.mark.parametrize('with_labels', [False, True])
+def test_packed_custom_loss_restoration_does_not_require_token_labels(with_labels):
+    from types import SimpleNamespace
+    from twinkle.processor import InputProcessor
+
+    loss = SimpleNamespace(
+        input_key_specs=(DataKeySpec('token_targets', pad_value=-9),),
+        output_key_specs=(DataKeySpec('token_scores', side='output', stage='output'),))
+    input_specs, output_specs = resolve_key_specs(loss)
+    inputs = dict(position_ids=torch.tensor([[0, 1, 2, 0, 1]]), token_targets=torch.zeros(1, 5))
+    if with_labels:
+        inputs['labels'] = torch.full((1, 5), -100)
+    outputs = {'token_scores': torch.tensor([[1., 1., 1., 3., 3.]])}
+    restored, outputs = InputProcessor(padding_free=True).unpack_packed_sequences(
+        inputs, outputs, input_key_specs=input_specs, output_key_specs=output_specs)
+    targets, scores = restored['token_targets'], outputs['token_scores']
+    torch.testing.assert_close(targets, torch.tensor([[0., 0., 0.], [0., 0., -9.]]))
+    valid = targets != -9
+    balanced_loss = (((scores - targets).square() * valid).sum(-1) / valid.sum(-1)).mean()
+    torch.testing.assert_close(balanced_loss, torch.tensor(5.))
+    assert ('labels' in restored) == with_labels
+
+
+def test_packed_output_logits_restore_without_labels():
+    from twinkle.processor import InputProcessor
+
+    logits = torch.arange(10).reshape(1, 5, 2)
+    _, outputs = InputProcessor(padding_free=True).unpack_packed_sequences(
+        {'position_ids': torch.tensor([[0, 1, 2, 0, 1]])}, {'logits': logits})
+    torch.testing.assert_close(outputs['logits'], torch.tensor([[[0, 1], [2, 3], [4, 5]],
+                                                             [[6, 7], [8, 9], [0, 0]]]))
+
+
+@pytest.mark.parametrize('operation', ['pad', 'split'])
+def test_loss_only_fields_reject_unreachable_forward_operations(operation):
+    from dataclasses import replace
+
+    spec = (DataKeySpec('token_targets', pad_at='forward', need_split=False) if operation == 'pad' else
+            DataKeySpec('token_targets', pad_at='prepare', split_at='forward'))
+    with pytest.raises(ValueError, match='Loss-only input.*forward boundary'):
+        _strategy().register_input_key(spec)
+    # Model arguments can reach the forward hook, regardless of their stage name.
+    _strategy().register_input_key(replace(spec, model_input=True))
diff --git a/tests/transformers/test_sp_loss_contract.py b/tests/transformers/test_sp_loss_contract.py
index 5836868e3..692455ce1 100644
--- a/tests/transformers/test_sp_loss_contract.py
+++ b/tests/transformers/test_sp_loss_contract.py
@@ -98,7 +98,7 @@ def test_declared_but_unregistered_or_missing_output_fails_loudly():
 
 def test_ppo_value_loss_declares_per_token_value_collection():
     assert PPOValueLoss.sp_gather_keys == ('values',)
-    assert PPOValueLoss.sp_output_key_specs[0].name == 'values'
+    assert PPOValueLoss.output_key_specs[0].name == 'values'
 
 
 class _ToyValueModel(torch.nn.Module):
diff --git a/tests/transformers/test_sp_model_adapters.py b/tests/transformers/test_sp_model_adapters.py
index 86c46da84..cacbe04e7 100644
--- a/tests/transformers/test_sp_model_adapters.py
+++ b/tests/transformers/test_sp_model_adapters.py
@@ -71,7 +71,8 @@ def test_builtin_adapters_install_model_specific_patches(monkeypatch, adapter, p
     assert calls[0][2] is ctx
 
 
-def test_moe_adapter_registers_router_logits_gather_hook(monkeypatch):
+@pytest.mark.parametrize('mask_kind', ['omitted', 'none', 'explicit'])
+def test_moe_adapter_registers_router_logits_gather_hook(monkeypatch, mask_kind):
     hooks = []
 
     class MoEConfig:
@@ -94,6 +95,7 @@ class Output(dict):
     assert adapter.matches(model)
     adapter.configure(ctx, model)
     assert hooks and hooks[0][1]
-    output = hooks[0][0](model, (), {'attention_mask': torch.ones(1, 3)}, Output())
+    kwargs = {} if mask_kind == 'omitted' else {'attention_mask': None if mask_kind == 'none' else torch.ones(1, 3)}
+    output = hooks[0][0](model, (), kwargs, Output())
     assert len(output['router_logits']) == 1
     assert output['router_logits'][0].shape == (3, 3)

From 67be3379627ffc5b21d0dac424fc840dcb73e5f0 Mon Sep 17 00:00:00 2001
From: clc <clc0301@cmbchina.com>
Date: Mon, 26 Oct 2026 17:13:20 +0800
Subject: [PATCH 7/8] sp: remove legacy global context singleton

---
 .../strategy/sequence_parallel/__init__.py           | 12 +-----------
 .../transformers/test_qwen35_linear_attention_sp.py  |  4 ++--
 tests/transformers/test_sequence_parallel_and_cp.py  |  4 ++--
 tests/transformers/test_sp_context_isolation.py      |  3 ---
 4 files changed, 5 insertions(+), 18 deletions(-)

diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
index 80591ba6c..11f2376ec 100644
--- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
@@ -582,12 +582,6 @@ def prepare(
         # towers) fall through to the origin attention.
         text_model = llm_model.language_model if hasattr(llm_model, 'language_model') else llm_model
         _register_module_tree(text_model, self)
-        if self is not sequence_parallel:
-            # Compat: keep the legacy module-level singleton reflecting the most
-            # recently prepared context (extra_kwargs is shared by reference, so
-            # legacy readers observe live per-forward state). New code should
-            # use explicit contexts instead.
-            sequence_parallel.__dict__.update(self.__dict__)
 
     def _mask_qkv(self, query, key, value, mask):
         mask = mask.unsqueeze(2).unsqueeze(3)
@@ -914,9 +908,6 @@ def prepare_inputs(self, inputs):
         return inputs
 
 
-sequence_parallel = SequenceParallel()
-
-
 @dataclass(frozen=True)
 class SequenceParallelConfig:
     enabled: bool = True
@@ -947,8 +938,7 @@ def __init__(
         self._tokenizer_id = tokenizer_id
         self._tokenizer = None
         self._initialized = False
-        # Per-model SP context. The module-level ``sequence_parallel`` singleton
-        # is only kept as a compat mirror of the most recently prepared context.
+        # Each model owns its SP context.
         self._ctx = SequenceParallel()
         # Compatibility declarations registered directly through the public API.
         # Keep these separate from rules consumed from model/loss providers.
diff --git a/tests/transformers/test_qwen35_linear_attention_sp.py b/tests/transformers/test_qwen35_linear_attention_sp.py
index d46c423cf..8f1a3b166 100644
--- a/tests/transformers/test_qwen35_linear_attention_sp.py
+++ b/tests/transformers/test_qwen35_linear_attention_sp.py
@@ -12,7 +12,7 @@
 from types import MethodType, SimpleNamespace
 
 from twinkle.loss import CrossEntropyLoss
-from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelStrategy, sequence_parallel
+from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelStrategy
 from twinkle.model.transformers.strategy.sequence_parallel.linear_attention_sp import Qwen3_5GatedDeltaNetUlyssesPatch
 from twinkle.utils import DeviceMesh, selective_log_softmax
 
@@ -367,7 +367,7 @@ def _run_prefill_alignment_worker(rank: int,
         if sp_num_tokens != baseline_num_tokens:
             raise AssertionError(
                 f'prefill num_tokens mismatch on rank {rank}: baseline={baseline_num_tokens} sp={sp_num_tokens}')
-        sp_qkv_grad = _average_qkv_grad_over_group(sp_model, sequence_parallel._data_rank_group) / max(sp_num_tokens, 1)
+        sp_qkv_grad = _average_qkv_grad_over_group(sp_model, strategy._ctx._data_rank_group) / max(sp_num_tokens, 1)
         if not torch.allclose(sp_qkv_grad, baseline_qkv_grad, rtol=GRAD_RTOL, atol=GRAD_ATOL):
             max_diff = (sp_qkv_grad - baseline_qkv_grad).abs().max().item()
             raise AssertionError(f'qkv grad mismatch on rank {rank}: max_diff={max_diff}')
diff --git a/tests/transformers/test_sequence_parallel_and_cp.py b/tests/transformers/test_sequence_parallel_and_cp.py
index 23effbf3a..8d3590bab 100644
--- a/tests/transformers/test_sequence_parallel_and_cp.py
+++ b/tests/transformers/test_sequence_parallel_and_cp.py
@@ -13,7 +13,7 @@
 from types import SimpleNamespace
 
 from twinkle.loss import CrossEntropyLoss
-from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelStrategy, sequence_parallel
+from twinkle.model.transformers.strategy.sequence_parallel import SequenceParallelStrategy
 from twinkle.utils import DeviceMesh, Platform, ensure_hccl_socket_env, selective_log_softmax, torch_util
 
 pytestmark = [pytest.mark.slow, pytest.mark.accel(2)]
@@ -542,7 +542,7 @@ def _run_precision_worker(rank: int,
         sp_loss_sum, sp_num_tokens = _compute_training_path_loss(local_logits, local_labels, strategy)
         sp_loss_sum.backward()
         global_loss = sp_loss_sum / max(sp_num_tokens, 1)
-        _average_model_grads_over_group(sp_model, sequence_parallel._data_rank_group)
+        _average_model_grads_over_group(sp_model, strategy._ctx._data_rank_group)
         sp_attention_grads = _normalize_grad_dict(_collect_attention_param_grads(sp_model), sp_num_tokens)
 
         if sp_num_tokens != base_num_tokens:
diff --git a/tests/transformers/test_sp_context_isolation.py b/tests/transformers/test_sp_context_isolation.py
index e4db8a2de..d285e8e39 100644
--- a/tests/transformers/test_sp_context_isolation.py
+++ b/tests/transformers/test_sp_context_isolation.py
@@ -188,11 +188,8 @@ def _update_causal_mask(self, mask, *args):
     monkeypatch.setattr(sp, 'get_llm_model', lambda value: value)
     monkeypatch.setattr(sp.SequenceParallel, '_global_inited', True)
     monkeypatch.setattr(sp.SequenceParallel, '_init_device_mesh', lambda *args: None)
-    monkeypatch.setattr(sp, 'sequence_parallel', sp.SequenceParallel())
     ctx.prepare(2, model, SimpleNamespace(pad_token_id=0))
     torch.testing.assert_close(ctx.causal_mask_func(torch.zeros(1)), torch.ones(1))
-    # Remove the compatibility mirror so only the weak registry owns the context.
-    sp.sequence_parallel.__dict__.clear()
     del model
     gc.collect()
     assert all(ref() is None for ref in refs)

From 8d55bc4de0d379559c31a7c5617e499089b356cb Mon Sep 17 00:00:00 2001
From: clc <clc0301@cmbchina.com>
Date: Tue, 27 Oct 2026 10:58:37 +0800
Subject: [PATCH 8/8] sp: simplify field contracts and unify loss gathering

---
 src/twinkle/data_fields.py                    | 201 ++++++++----------
 src/twinkle/loss/value.py                     |  10 +-
 .../strategy/sequence_parallel/__init__.py    |  68 +++---
 .../strategy/sequence_parallel/adapters.py    |  17 --
 .../model/transformers/transformers.py        |   1 -
 src/twinkle/processor/base.py                 |   4 +-
 .../test_qwen35_linear_attention_sp.py        |   3 +-
 .../test_sequence_parallel_and_cp.py          |  37 ++--
 tests/transformers/test_sp_custom_loss_cpu.py |  12 +-
 tests/transformers/test_sp_data_keys.py       | 112 ++++++++--
 tests/transformers/test_sp_loss_contract.py   |  38 +++-
 tests/transformers/test_sp_model_adapters.py  |   7 +-
 12 files changed, 294 insertions(+), 216 deletions(-)

diff --git a/src/twinkle/data_fields.py b/src/twinkle/data_fields.py
index 0981416ff..9d912235f 100644
--- a/src/twinkle/data_fields.py
+++ b/src/twinkle/data_fields.py
@@ -1,128 +1,106 @@
 """Shared sequence-field schemas consumed by processors, model wrappers and SP."""
-from dataclasses import dataclass
-from typing import Optional
 
+from dataclasses import dataclass
 
-_BOUNDARIES = frozenset({'prepare', 'forward', 'model_internal', 'router_replay'})
+_BOUNDARIES = ("forward", "auxiliary", "model_internal", "router_replay", "output")
+_PROCESSING_BOUNDARIES = ("prepare", "forward", "model_internal", "router_replay")
+_FORWARD_KEYS = ("input_ids", "inputs_embeds", "position_ids", "attention_mask")
 
 
 @dataclass(frozen=True)
 class DataKeySpec:
-    """Describe collation, model routing and optional parallel transforms.
+    """Describe collation, routing and explicitly scheduled parallel operations.
 
-    Auxiliary loss fields stay global by default. model_input controls routing;
-    stage/pad_at/split_at control when an explicitly local tensor is transformed.
+    Loss-only fields stay global unless pad_at/split_at opt into local processing.
+    Packed restoration is independent of sequence-parallel gather and split.
     """
 
     name: str
     split_dim: int = -1
     pad_value: int | float | str = 0
-    need_roll: bool = False
-    side: str = 'input'
-    stage: str = 'auxiliary'
-    need_split: Optional[bool] = None
-    need_gather: bool = False
+    roll: bool = False
+    side: str = "input"
+    stage: str = "auxiliary"
     autograd: bool = False
-    owner: str = 'generic'
-    need_pad: Optional[bool] = None
-    pad_at: Optional[str] = 'auto'
-    split_at: Optional[str] = 'auto'
-    need_unpack: Optional[bool] = None
-    model_input: Optional[bool] = None
+    pad_at: str | None = None
+    split_at: str | None = None
+    gather_at: str | None = None
+    unpack_packed: bool | None = None
+    model_input: bool | None = None
 
     def __post_init__(self):
         if not isinstance(self.name, str) or not self.name:
-            raise ValueError('An SP input key must have a non-empty name')
+            raise ValueError("A sequence field must have a non-empty name")
         if not isinstance(self.split_dim, int):
-            raise TypeError('split_dim must be an integer')
-        if self.side not in ('input', 'output'):
-            raise ValueError('side must be input or output')
-        if self.stage not in ('forward', 'auxiliary', 'model_internal', 'router_replay', 'output'):
-            raise ValueError('Invalid SP key stage')
-        if (self.side == 'output') != (self.stage == 'output'):
-            raise ValueError('Output keys must use stage=output')
-        if self.owner not in ('generic', 'task', 'model', 'loss'):
-            raise ValueError('Invalid SP key owner')
-        if isinstance(self.pad_value, str) and (self.name != 'input_ids' or self.pad_value != 'tokenizer'):
-            raise ValueError('Only input_ids can use the tokenizer padding sentinel')
-        if self.need_split is None:
-            local_stage = self.stage in ('forward', 'model_internal', 'router_replay')
-            explicit_boundary = self.split_at not in (None, 'auto')
-            object.__setattr__(self, 'need_split', self.side == 'input'
-                               and (local_stage or explicit_boundary or self.need_gather or self.need_roll))
-        if self.side == 'input' and self.need_gather and not self.need_split:
-            raise ValueError('Only sharded input keys can request gathering')
-        if self.stage == 'forward' and self.need_gather:
-            raise ValueError('Model-forward keys cannot request loss gathering')
+            raise TypeError("split_dim must be an integer")
+        if self.side not in ("input", "output"):
+            raise ValueError("side must be input or output")
+        if self.stage not in _BOUNDARIES:
+            raise ValueError("Invalid SP key stage")
+        if (self.side == "output") != (self.stage == "output"):
+            raise ValueError("Output keys must use stage=output")
+        if isinstance(self.pad_value, str) and (self.name != "input_ids" or self.pad_value != "tokenizer"):
+            raise ValueError("Only input_ids can use the tokenizer padding sentinel")
+        for operation, boundary in (("padding", self.pad_at), ("split", self.split_at)):
+            if boundary is not None and boundary not in _PROCESSING_BOUNDARIES:
+                raise ValueError(f"Invalid SP {operation} boundary: {boundary!r}")
+        if self.gather_at not in (None, "loss"):
+            raise ValueError(f"Invalid SP gather boundary: {self.gather_at!r}")
+        if self.side == "output" and (self.need_pad or self.need_split):
+            raise ValueError("Output keys cannot request input padding or splitting")
+        if self.side == "input" and self.gather_at is not None and not self.need_split:
+            raise ValueError("Only sharded input keys can request gathering")
+        if self.stage == "forward" and self.gather_at is not None:
+            raise ValueError("Model-forward keys cannot request loss gathering")
+        if self.roll and not self.need_split:
+            raise ValueError("Rolling requires a split boundary")
         if self.model_input is None:
-            object.__setattr__(self, 'model_input', self.stage == 'forward')
-        default_boundary = {
-            'auxiliary': 'prepare',
-            'forward': 'forward',
-            'model_internal': 'model_internal',
-            'router_replay': 'router_replay',
-            'output': None,
-        }[self.stage]
-        if self.need_pad is None:
-            object.__setattr__(self, 'need_pad', self.side == 'input'
-                               and (self.need_split or self.pad_at not in (None, 'auto')))
-        if self.pad_at == 'auto':
-            object.__setattr__(self, 'pad_at', default_boundary if self.need_pad else None)
-        elif self.pad_at is not None and self.pad_at not in _BOUNDARIES:
-            raise ValueError(f'Invalid SP padding boundary: {self.pad_at!r}')
-        elif (self.pad_at is None) == bool(self.need_pad):
-            raise ValueError('need_pad and pad_at specify conflicting SP padding behavior')
-        if self.pad_at != 'auto':
-            object.__setattr__(self, 'need_pad', self.pad_at is not None)
-        if self.split_at == 'auto':
-            object.__setattr__(self, 'split_at', default_boundary if self.need_split else None)
-        elif self.split_at is not None and self.split_at not in _BOUNDARIES:
-            raise ValueError(f'Invalid SP split boundary: {self.split_at!r}')
-        elif (self.split_at is None) == bool(self.need_split):
-            raise ValueError('need_split and split_at specify conflicting SP split behavior')
-        if self.split_at != 'auto':
-            object.__setattr__(self, 'need_split', self.split_at is not None)
-        if self.side == 'output' and (self.pad_at is not None or self.split_at is not None):
-            raise ValueError('Output keys cannot request input padding or splitting')
-        if self.need_unpack is None:
-            # Preserve the historical packed restoration behavior unless a
-            # built-in rule explicitly opts out.
-            object.__setattr__(self, 'need_unpack', self.side == 'output' or self.need_gather or not self.need_split)
+            object.__setattr__(self, "model_input", self.stage == "forward")
+        if self.unpack_packed is None:
+            object.__setattr__(
+                self, "unpack_packed", self.side == "output" or self.gather_at is not None or not self.need_split
+            )
+
+    @property
+    def need_pad(self):
+        return self.pad_at is not None
+
+    @property
+    def need_split(self):
+        return self.split_at is not None
 
     def sequence_dim(self, ndim):
         if not -ndim <= self.split_dim < ndim:
-            raise ValueError(f'Invalid sequence axis {self.split_dim} for SP key {self.name!r}')
+            raise ValueError(f"Invalid sequence axis {self.split_dim} for SP key {self.name!r}")
         dim = self.split_dim % ndim
         if dim == 0:
-            raise ValueError(f'SP key {self.name!r} cannot split or gather the batch axis')
+            raise ValueError(f"SP key {self.name!r} cannot split or gather the batch axis")
         return dim
 
 
 BUILTIN_INPUT_KEYS = (
-    DataKeySpec('input_ids', split_dim=1, pad_value='tokenizer', stage='forward', need_split=True),
-    DataKeySpec('inputs_embeds', split_dim=1, stage='forward', need_split=True),
-    DataKeySpec('position_ids', pad_value=-1, stage='forward', need_split=True),
-    DataKeySpec('attention_mask', split_dim=1, stage='forward', need_pad=True, need_split=False),
-    DataKeySpec('labels', pad_value=-100, need_split=True, need_gather=True, need_unpack=True),
-    # These are not needed by the model forward. Keep them global through SP;
-    # packed sequence restoration is a separate action from sharding/gathering.
-    DataKeySpec('completion_mask', need_pad=False, need_split=False, need_unpack=True),
-    DataKeySpec('loss_scale', need_pad=False, need_split=False, need_unpack=True),
+    DataKeySpec("input_ids", split_dim=1, pad_value="tokenizer", stage="forward", pad_at="forward", split_at="forward"),
+    DataKeySpec("inputs_embeds", split_dim=1, stage="forward", pad_at="forward", split_at="forward"),
+    DataKeySpec("position_ids", pad_value=-1, stage="forward", pad_at="forward", split_at="forward"),
+    DataKeySpec("attention_mask", split_dim=1, stage="forward", pad_at="forward"),
+    DataKeySpec("labels", pad_value=-100, pad_at="prepare", split_at="prepare", gather_at="loss", unpack_packed=True),
+    # Loss-only global fields still need per-sample restoration after packing.
+    DataKeySpec("completion_mask", unpack_packed=True),
+    DataKeySpec("loss_scale", unpack_packed=True),
 )
 
-# These are processor/model control arguments, not sequence data contracts.
-# Keep them out of DataKeySpec and preserve them through the Transformers input
-# conversion path explicitly instead.
-PROCESSOR_PASSTHROUGH_KEYS = ('cache_position', 'padding_free')
+# Control arguments are preserved without sequence transforms.
+PROCESSOR_PASSTHROUGH_KEYS = ("cache_position", "padding_free")
 
 BUILTIN_OUTPUT_KEYS = (
-    DataKeySpec('logps', split_dim=1, side='output', stage='output',
-                need_split=False, need_gather=True, autograd=True, need_unpack=True),
-    DataKeySpec('entropies', split_dim=1, side='output', stage='output',
-                need_split=False, need_gather=True, autograd=True, need_unpack=True),
+    DataKeySpec(
+        "logps", split_dim=1, side="output", stage="output", gather_at="loss", autograd=True, unpack_packed=True
+    ),
+    DataKeySpec(
+        "entropies", split_dim=1, side="output", stage="output", gather_at="loss", autograd=True, unpack_packed=True
+    ),
     # Real logits are opt-in; fused CE explicitly gathers this hidden-state slot.
-    DataKeySpec('logits', split_dim=1, side='output', stage='output',
-                need_split=False, autograd=True, need_unpack=True),
+    DataKeySpec("logits", split_dim=1, side="output", stage="output", autograd=True, unpack_packed=True),
 )
 
 
@@ -134,7 +112,6 @@ class InputKeyRegistry:
     hook, so every key is sharded at most once.
     """
 
-    _FORWARD_KEYS = frozenset({'input_ids', 'inputs_embeds', 'position_ids', 'attention_mask'})
     _SPECIAL_KEYS = frozenset(PROCESSOR_PASSTHROUGH_KEYS)
 
     def __init__(self):
@@ -142,19 +119,21 @@ def __init__(self):
 
     def register(self, spec: DataKeySpec):
         if not isinstance(spec, DataKeySpec):
-            raise TypeError('Expected a DataKeySpec')
-        if spec.side != 'input' or spec.stage not in ('auxiliary', 'forward'):
-            raise ValueError('Register only auxiliary or model-forward input keys here')
-        if spec.name in self._FORWARD_KEYS:
-            raise ValueError(f'{spec.name!r} is owned by the model forward hook')
+            raise TypeError("Expected a DataKeySpec")
+        if spec.side != "input" or spec.stage not in ("auxiliary", "forward"):
+            raise ValueError("Register only auxiliary or model-forward input keys here")
+        if spec.name in _FORWARD_KEYS:
+            raise ValueError(f"{spec.name!r} is owned by the model forward hook")
         if spec.name in self._SPECIAL_KEYS:
-            raise ValueError(f'{spec.name!r} is a special processor/model argument, not an SP data key')
-        if not spec.model_input and 'forward' in (spec.pad_at, spec.split_at):
-            raise ValueError(f'Loss-only input {spec.name!r} cannot use the forward boundary; '
-                             'use prepare, or set model_input=True for a model argument')
+            raise ValueError(f"{spec.name!r} is a special processor/model argument, not an SP data key")
+        if not spec.model_input and "forward" in (spec.pad_at, spec.split_at):
+            raise ValueError(
+                f"Loss-only input {spec.name!r} cannot use the forward boundary; "
+                "use prepare, or set model_input=True for a model argument"
+            )
         existing = self._specs.get(spec.name)
         if existing is not None and existing != spec:
-            raise ValueError(f'Conflicting SP input rule for {spec.name!r}')
+            raise ValueError(f"Conflicting SP input rule for {spec.name!r}")
         self._specs[spec.name] = spec
 
     def __iter__(self):
@@ -170,12 +149,12 @@ def __init__(self):
 
     def register(self, spec):
         if not isinstance(spec, DataKeySpec):
-            raise TypeError('Expected a DataKeySpec')
-        if spec.side != 'output':
-            raise ValueError('Expected an output key rule')
+            raise TypeError("Expected a DataKeySpec")
+        if spec.side != "output":
+            raise ValueError("Expected an output key rule")
         existing = self._specs.get(spec.name)
         if existing is not None and existing != spec:
-            raise ValueError(f'Conflicting SP output rule for {spec.name!r}')
+            raise ValueError(f"Conflicting SP output rule for {spec.name!r}")
         self._specs[spec.name] = spec
 
     def __iter__(self):
@@ -187,10 +166,10 @@ def get(self, name):
 
 def declared_key_specs(provider, side):
     """Read common field declarations, retaining legacy loss declaration names."""
-    specs = getattr(provider, f'{side}_key_specs', ())
+    specs = getattr(provider, f"{side}_key_specs", ())
     if callable(specs):
         specs = specs()
-    legacy = getattr(provider, f'sp_{side}_key_specs', ())
+    legacy = getattr(provider, f"sp_{side}_key_specs", ())
     for entries in (specs, legacy):
         if isinstance(entries, DataKeySpec):
             entries = (entries,)
@@ -201,7 +180,7 @@ def resolve_key_specs(*providers):
     """Build a call-local schema independently of any parallel strategy."""
     inputs, outputs = InputKeyRegistry(), OutputKeyRegistry()
     for provider in providers:
-        for side, registry in (('input', inputs), ('output', outputs)):
+        for side, registry in (("input", inputs), ("output", outputs)):
             for spec in declared_key_specs(provider, side):
                 registry.register(spec)
     return tuple(inputs), tuple(outputs)
@@ -210,7 +189,7 @@ def resolve_key_specs(*providers):
 def partition_model_inputs(inputs, input_key_specs):
     """Separate forward arguments from loss-only tensors without mutating the batch."""
     loss_names = {spec.name for spec in input_key_specs if not spec.model_input}
-    loss_names.add('channel')  # sample-level loss metadata, not a sequence tensor
+    loss_names.add("channel")  # sample-level loss metadata, not a sequence tensor
     model_inputs = {key: value for key, value in inputs.items() if key not in loss_names}
     loss_inputs = {key: value for key, value in inputs.items() if key in loss_names}
     return model_inputs, loss_inputs
diff --git a/src/twinkle/loss/value.py b/src/twinkle/loss/value.py
index df1f463ea..3889dde47 100644
--- a/src/twinkle/loss/value.py
+++ b/src/twinkle/loss/value.py
@@ -23,8 +23,14 @@ class PPOValueLoss(Loss):
     require_values = True
     # PPO regresses per-token values; gather the value sequence before loss.
     sp_gather_keys = ('values',)
-    output_key_specs = (DataKeySpec('values', split_dim=1, side='output', stage='output',
-                                       need_split=False, autograd=True, need_unpack=True),)
+    output_key_specs = (DataKeySpec(
+        'values',
+        split_dim=1,
+        side='output',
+        stage='output',
+        autograd=True,
+        unpack_packed=True,
+    ),)
 
     def __init__(self, cliprange_value: float = 0.2, vf_coef: float = 1.0, ignore_index: int = -100, **kwargs):
         self.cliprange_value = cliprange_value
diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
index 11f2376ec..ee995d1a2 100644
--- a/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/__init__.py
@@ -10,11 +10,12 @@
 from typing import Any, Dict, List, Optional, Tuple, Union
 from weakref import WeakKeyDictionary, WeakMethod
 
+from twinkle.data_fields import _FORWARD_KEYS, _PROCESSING_BOUNDARIES
 from twinkle.patch import apply_patch
 from twinkle.utils import DeviceMesh, get_logger
 from twinkle.utils.transformers_utils import get_llm_model, is_flash_attention_implementation
 from twinkle.utils.utils import call_with_supported_kwargs, has_signature_parameter
-from .adapters import ModelSPAdapter, adapters_for_model, register_model_sp_adapter
+from .adapters import ModelSPAdapter, adapters_for_model
 from .utils import (DistributedAttention, GatherLoss, _derive_sequence_parallel_sizes, _get_seq_groups_from_device_mesh,
                     _get_ulysses_size, _SeqAllToAll, get_config_attr, get_cu_seqlens_from_position_ids, is_hccl_backend,
                     post_all2all)
@@ -472,7 +473,7 @@ def pre_forward_split_hook(_self, args, kwargs):
                 cache_position = kwargs.get('cache_position', None)
                 specs = tuple(spec for spec in self.input_key_registry
                               if spec.pad_at == 'forward' or spec.split_at == 'forward')
-                common_names = ('input_ids', 'inputs_embeds', 'position_ids', 'attention_mask')
+                common_names = _FORWARD_KEYS
                 keys = SPKeyCollection.from_mapping(kwargs, specs, common_names=common_names)
                 self.pad_and_split_inputs(keys, boundary='forward', real_position_ids=real_position_ids,
                                           cache_position=cache_position)
@@ -629,10 +630,22 @@ def pad_and_split_mm_tokens(self, visual_mask, mm_embeds):
         empty_embeds = torch.empty(
             (input_ids.shape[0], input_ids.shape[1], mm_embeds.shape[-1])).to(mm_embeds.device).to(mm_embeds.dtype)
         empty_embeds[visual_mask] = mm_embeds
-        embeds_spec = DataKeySpec('visual_sequence_embeds', split_dim=1, pad_value=0, stage='model_internal',
-                                  pad_at='model_internal', split_at='model_internal')
-        mask_spec = DataKeySpec('visual_position_mask', split_dim=-1, pad_value=0, stage='model_internal',
-                                pad_at='model_internal', split_at='model_internal')
+        embeds_spec = DataKeySpec(
+            'visual_sequence_embeds',
+            split_dim=1,
+            pad_value=0,
+            stage='model_internal',
+            pad_at='model_internal',
+            split_at='model_internal',
+        )
+        mask_spec = DataKeySpec(
+            'visual_position_mask',
+            split_dim=-1,
+            pad_value=0,
+            stage='model_internal',
+            pad_at='model_internal',
+            split_at='model_internal',
+        )
         keys = SPKeyCollection()
         keys.add(SPKeyValue(embeds_spec, empty_embeds))
         keys.add(SPKeyValue(mask_spec, visual_mask))
@@ -765,7 +778,7 @@ def _pad_key(self, value, spec, position_ids=None):
     def _split_key(self, value, spec, position_ids=None, legacy_roll=False):
         if value is None or not spec.need_split:
             return value
-        if spec.need_roll or legacy_roll:
+        if spec.roll or legacy_roll:
             value = torch.roll(value, shifts=-1, dims=spec.split_dim)
         return self.split(value, dim=spec.split_dim, position_ids=position_ids)
 
@@ -783,7 +796,7 @@ def pad_and_split_inputs(self,
         """
         if not isinstance(keys, SPKeyCollection):
             raise TypeError('pad_and_split_inputs expects an SPKeyCollection')
-        if boundary not in ('prepare', 'forward', 'model_internal', 'router_replay'):
+        if boundary not in _PROCESSING_BOUNDARIES:
             raise ValueError(f'Unknown SP processing boundary: {boundary!r}')
 
         position_ids = keys.get_value('position_ids')
@@ -810,9 +823,10 @@ def pad_and_split_inputs(self,
             if not torch.is_tensor(value):
                 raise TypeError(f'SP input {spec.name!r} must be a tensor')
             if spec.pad_at == boundary:
-                # Preserve the existing attention-mask contract: it is padded
-                # only for ordinary batched sequences, not packed batch_size=1.
-                if spec.name != 'attention_mask' or (batch_size is not None and batch_size > 1):
+                # Dense SDPA also needs a global padding mask for batch_size=1;
+                # FlashAttention's packed path keeps its existing mask handling.
+                if (spec.name != 'attention_mask' or (batch_size is not None and batch_size > 1)
+                        or not is_flash_attention_implementation(self.attn_implementation)):
                     value = self._pad_key(value, spec, real_position_ids)
             item.value = value
 
@@ -825,7 +839,8 @@ def pad_and_split_inputs(self,
 
         # Keep the legacy generated 2D mask and causal-mask conversion on the
         # forward boundary; these depend on model inputs and cache state.
-        if boundary == 'forward' and padded_model_inputs is not None and batch_size > 1:
+        if (boundary == 'forward' and padded_model_inputs is not None
+                and (batch_size > 1 or not is_flash_attention_implementation(self.attn_implementation))):
             attn_item = keys.get('attention_mask')
             if attn_item is None or attn_item.value is None:
                 attention_mask = (padded_positions != -1).to(dtype=torch.int64)
@@ -912,7 +927,6 @@ def prepare_inputs(self, inputs):
 class SequenceParallelConfig:
     enabled: bool = True
     ulysses_size: Optional[int] = None
-    gather_logits: bool = True
 
 
 class SequenceParallelStrategy:
@@ -1035,7 +1049,7 @@ def pool_last_valid_token(self, features, inputs):
     def register_input_key(self, spec: DataKeySpec) -> None:
         """Declare a token-aligned input on this model only.
 
-        pad_at/split_at select execution boundaries; need_gather selects
+        pad_at/split_at select execution boundaries; gather_at selects
         loss-side collection. The model wrapper owns input routing and merges
         these explicit declarations into the processor schema.
         """
@@ -1061,7 +1075,7 @@ def output_key_specs(self) -> Tuple[DataKeySpec, ...]:
     def register_key_specs(self, input_key_specs, output_key_specs):
         """Consume model-resolved field definitions for SP execution."""
         for spec in input_key_specs:
-            if spec.name not in self._ctx.input_key_registry._FORWARD_KEYS:
+            if spec.name not in _FORWARD_KEYS:
                 self._ctx.input_key_registry.register(spec)
         for spec in output_key_specs:
             self._ctx.output_key_registry.register(spec)
@@ -1078,26 +1092,6 @@ def model_sp_adapters(self) -> Tuple[ModelSPAdapter, ...]:
         """The model-family adapters selected during context preparation."""
         return self._ctx.model_adapters
 
-    def postprocess_outputs(self, outputs: Any) -> Any:
-        if (not self.enabled or self.ulysses_size <= 1 or not self.sp_config.get('gather_logits', True)):
-            return outputs
-        # Twinkle expects dict-like ModelOutput containers in the main training path
-        # (uses `.get(...)` and `outputs[...] = ...`). Keep SP postprocess consistent.
-        if outputs is None or not hasattr(outputs, 'get') or not hasattr(outputs, '__setitem__'):
-            raise TypeError('SequenceParallelStrategy.postprocess_outputs expects a dict-like ModelOutput. '
-                            f'Got type={type(outputs)}')
-        logits = outputs.get('logits', None)
-        if logits is None or not torch.is_tensor(logits) or logits.dim() < 2:
-            return outputs
-        gathered = self._ctx.gather(logits, dim=1, position_ids=self._ctx.real_position_ids)
-        # Scheme A: SP pads to make seq_len divisible by sp_size. Trim back to the original
-        # (unpadded) length using the cached real_position_ids.
-        real_pos = self._ctx.real_position_ids
-        if real_pos is not None and torch.is_tensor(real_pos) and real_pos.dim() >= 2:
-            gathered = gathered[:, :real_pos.shape[1]].contiguous()
-        outputs['logits'] = gathered
-        return outputs
-
     def _trim_gathered_sequence_padding(self, tensor: torch.Tensor, real_position_ids: torch.Tensor,
                                         dim: int = 1) -> torch.Tensor:
         if real_position_ids is None or not torch.is_tensor(real_position_ids) or real_position_ids.dim() < 2:
@@ -1146,7 +1140,7 @@ def gather_loss_tensors(
         real_position_ids = self._ctx.real_position_ids
         for spec in self.input_key_specs:
             value = inputs.get(spec.name)
-            if spec.stage == 'auxiliary' and spec.need_gather and value is not None:
+            if spec.stage == 'auxiliary' and spec.gather_at == 'loss' and value is not None:
                 inputs[spec.name] = self._gather_key(value, spec, real_position_ids)
 
         fused_hidden = outputs.get('logps') is None and 'lm_head' in outputs and inputs.get('labels') is not None
@@ -1167,7 +1161,7 @@ def gather_loss_tensors(
             raise KeyError(f'Loss declared SP gather key(s) {sorted(missing_keys)}, but the model output is missing them.')
         for spec in self.output_key_specs:
             value = outputs.get(spec.name)
-            if value is not None and (spec.need_gather or spec.name in required_keys
+            if value is not None and (spec.gather_at == 'loss' or spec.name in required_keys
                                       or (spec.name == 'logits' and fused_hidden)):
                 outputs[spec.name] = self._gather_key(value, spec, real_position_ids)
         return inputs, outputs
diff --git a/src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py b/src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py
index a00f8f75f..3c1d97970 100644
--- a/src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py
+++ b/src/twinkle/model/transformers/strategy/sequence_parallel/adapters.py
@@ -117,23 +117,6 @@ def gather_router_logits(module, args, kwargs, output):
 _ADAPTERS = [Qwen3VLSPAdapter(), Qwen35LinearAttentionSPAdapter(), MoEAuxSPAdapter()]
 
 
-def register_model_sp_adapter(adapter):
-    """Register an adapter instance for future SP model initialization."""
-    if isinstance(adapter, type):
-        adapter = adapter()
-    if not isinstance(adapter, ModelSPAdapter):
-        raise TypeError('A model SP adapter must derive from ModelSPAdapter')
-    if not adapter.name or adapter.name == 'generic':
-        raise ValueError('A registered model SP adapter needs a stable non-generic name')
-    for current in _ADAPTERS:
-        if current.name == adapter.name:
-            if type(current) is type(adapter):
-                return current
-            raise ValueError(f'Conflicting model SP adapter name: {adapter.name!r}')
-    _ADAPTERS.append(adapter)
-    return adapter
-
-
 def adapters_for_model(model):
     """Return every matching adapter; families can compose capabilities."""
     return tuple(adapter for adapter in _ADAPTERS if adapter.matches(model))
diff --git a/src/twinkle/model/transformers/transformers.py b/src/twinkle/model/transformers/transformers.py
index b5e151dca..6b105312b 100644
--- a/src/twinkle/model/transformers/transformers.py
+++ b/src/twinkle/model/transformers/transformers.py
@@ -1135,7 +1135,6 @@ def _forward_only_single(self, *, inputs, optimizer_config, **kwargs):
                 cp_world_size = self.device_mesh.cp_world_size if self.device_mesh is not None else 1
                 if getattr(self, '_enable_sp', False) or cp_world_size > 1:
                     raise ValueError('sampling replay does not support sequence or context parallelism')
-            loss_require_values = getattr(loss_instance, 'require_values', False)
             input_key_specs, output_key_specs = self._get_field_specs(loss_instance)
             inputs: Dict[str, Any] = processor(
                 inputs,
diff --git a/src/twinkle/processor/base.py b/src/twinkle/processor/base.py
index ce8cba69b..fd0f0bd2f 100644
--- a/src/twinkle/processor/base.py
+++ b/src/twinkle/processor/base.py
@@ -536,9 +536,9 @@ def unpack_packed_sequences(
         input_specs = BUILTIN_INPUT_KEYS if input_key_specs is None else input_key_specs
         output_specs = BUILTIN_OUTPUT_KEYS if output_key_specs is None else output_key_specs
         input_keys = [spec for spec in input_specs
-                      if spec.stage == 'auxiliary' and spec.need_unpack and inputs.get(spec.name) is not None]
+                      if spec.stage == 'auxiliary' and spec.unpack_packed and inputs.get(spec.name) is not None]
         output_keys = [spec for spec in output_specs
-                       if spec.need_unpack and outputs and outputs.get(spec.name) is not None]
+                       if spec.unpack_packed and outputs and outputs.get(spec.name) is not None]
         all_specs = input_keys + output_keys
         all_tensors = [inputs[spec.name] for spec in input_keys] + [outputs[spec.name] for spec in output_keys]
         # The unpack primitive expects [1, T, ...]; restore the declared axis
diff --git a/tests/transformers/test_qwen35_linear_attention_sp.py b/tests/transformers/test_qwen35_linear_attention_sp.py
index 8f1a3b166..4289be684 100644
--- a/tests/transformers/test_qwen35_linear_attention_sp.py
+++ b/tests/transformers/test_qwen35_linear_attention_sp.py
@@ -215,7 +215,6 @@ def _make_strategy(model: torch.nn.Module, world_size: int) -> SequenceParallelS
         sp_config={
             'enabled': True,
             'ulysses_size': world_size,
-            'gather_logits': True,
         },
         model=model,
         tokenizer_id=None,
@@ -350,7 +349,7 @@ def _run_prefill_alignment_worker(rank: int,
             position_ids=position_ids,
             use_cache=False,
         )
-        gathered_outputs = strategy.postprocess_outputs(copy.copy(sp_outputs))
+        _, gathered_outputs = strategy.gather_loss_tensors({}, sp_outputs, sp_gather_keys=('logits',))
         gathered_logits = gathered_outputs.logits.float()
         if not torch.allclose(gathered_logits, baseline_logits, rtol=LOGITS_RTOL, atol=LOGITS_ATOL):
             max_diff = (gathered_logits - baseline_logits).abs().max().item()
diff --git a/tests/transformers/test_sequence_parallel_and_cp.py b/tests/transformers/test_sequence_parallel_and_cp.py
index 8d3590bab..5e7ec23c3 100644
--- a/tests/transformers/test_sequence_parallel_and_cp.py
+++ b/tests/transformers/test_sequence_parallel_and_cp.py
@@ -252,7 +252,6 @@ def _make_strategy(model: LlamaForCausalLM, device_mesh: DeviceMesh, ulysses_siz
         sp_config={
             'enabled': True,
             'ulysses_size': ulysses_size,
-            'gather_logits': True,
         },
         model=model,
         tokenizer_id=None,
@@ -536,7 +535,7 @@ def _run_precision_worker(rank: int,
             use_cache=False,
         )
         local_logits = sp_outputs.logits
-        gathered_outputs = strategy.postprocess_outputs(sp_outputs)
+        _, gathered_outputs = strategy.gather_loss_tensors({}, sp_outputs, sp_gather_keys=('logits',))
         sp_logits = gathered_outputs.logits.detach().float()
 
         sp_loss_sum, sp_num_tokens = _compute_training_path_loss(local_logits, local_labels, strategy)
@@ -722,7 +721,8 @@ def test_cp_sp_precision_alignment(self):
             join=True)
 
 
-def _run_multi_model_worker(rank: int, world_size: int, port: int, backend_config: dict):
+def _run_multi_model_worker(rank: int, world_size: int, port: int, backend_config: dict,
+                            attn_implementation: str):
     """Two SP-enabled models (policy/ref style) with interleaved forwards.
 
     Same architecture, different weights and different sequence lengths. Each
@@ -736,13 +736,13 @@ def _run_multi_model_worker(rank: int, world_size: int, port: int, backend_confi
         vocab_size = 256
 
         _seed_backend(1234, device_type)
-        policy_base = _build_tiny_llama(case, device, device_type)
+        policy_base = _build_tiny_llama(case, device, device_type, attn_implementation)
         _seed_backend(5678, device_type)
-        ref_base = _build_tiny_llama(case, device, device_type)
+        ref_base = _build_tiny_llama(case, device, device_type, attn_implementation)
 
-        policy_sp = _build_tiny_llama(case, device, device_type)
+        policy_sp = _build_tiny_llama(case, device, device_type, attn_implementation)
         policy_sp.load_state_dict(policy_base.state_dict())
-        ref_sp = _build_tiny_llama(case, device, device_type)
+        ref_sp = _build_tiny_llama(case, device, device_type, attn_implementation)
         ref_sp.load_state_dict(ref_base.state_dict())
 
         seq_len_p = int(case['seq_len'])
@@ -813,14 +813,17 @@ def _run_multi_model_worker(rank: int, world_size: int, port: int, backend_confi
         if ref_strategy._ctx.real_position_ids.shape[-1] != seq_len_r:
             raise AssertionError('multi_model: ref ctx position_ids polluted by policy forward.')
 
-        # Capture per-rank local logits (with grad) before postprocess_outputs
-        # replaces outputs['logits'] with the grad-less gathered tensor.
+        # Capture per-rank local logits (with grad) before loss-side gathering.
+        # Both paths must use the same registered gather implementation as training.
         local_logits_p1 = sp_outputs_p1.logits
         local_logits_r = sp_outputs_r.logits
 
-        sp_logits_p1 = policy_strategy.postprocess_outputs(sp_outputs_p1).logits.detach().float()
-        sp_logits_r = ref_strategy.postprocess_outputs(sp_outputs_r).logits.detach().float()
-        sp_logits_p2 = policy_strategy.postprocess_outputs(sp_outputs_p2).logits.detach().float()
+        _, gathered = policy_strategy.gather_loss_tensors({}, sp_outputs_p1, sp_gather_keys=('logits',))
+        sp_logits_p1 = gathered.logits.detach().float()
+        _, gathered = ref_strategy.gather_loss_tensors({}, sp_outputs_r, sp_gather_keys=('logits',))
+        sp_logits_r = gathered.logits.detach().float()
+        _, gathered = policy_strategy.gather_loss_tensors({}, sp_outputs_p2, sp_gather_keys=('logits',))
+        sp_logits_p2 = gathered.logits.detach().float()
 
         _assert_logits_close('multi_model policy fwd1', rank, base_logits_p, sp_logits_p1, seq_len_p, None)
         _assert_logits_close('multi_model ref fwd', rank, base_logits_r, sp_logits_r, seq_len_r, None)
@@ -869,13 +872,15 @@ def _get_backend_or_skip(self, world_size: int = 2) -> dict:
             pytest.skip(f'Requires at least {world_size} {backend["label"]} devices.')
         return backend
 
-    def test_two_models_interleaved_forward(self):
+    @pytest.mark.parametrize('attn_implementation', ['sdpa', 'flash_attention_2'])
+    def test_two_models_interleaved_forward(self, attn_implementation):
         world_size = 2
         backend = self._get_backend_or_skip(world_size)
-        if not is_flash_attn_available():
-            pytest.skip('flash_attention_2 is required for multi-model SP tests.')
+        if attn_implementation == 'flash_attention_2' and not is_flash_attn_available():
+            pytest.skip('flash_attention_2 is unavailable.')
         port = _find_free_port()
-        mp.spawn(_run_multi_model_worker, args=(world_size, port, backend), nprocs=world_size, join=True)
+        mp.spawn(_run_multi_model_worker, args=(world_size, port, backend, attn_implementation),
+                 nprocs=world_size, join=True)
 
 
 class TestDerivedRingMemoryProfile:
diff --git a/tests/transformers/test_sp_custom_loss_cpu.py b/tests/transformers/test_sp_custom_loss_cpu.py
index 7f0541ba9..e84b403ad 100644
--- a/tests/transformers/test_sp_custom_loss_cpu.py
+++ b/tests/transformers/test_sp_custom_loss_cpu.py
@@ -22,9 +22,15 @@ class TokenRegressionLoss(Loss):
     """Example plugin loss with custom sequence input and output keys."""
 
     require_logps = False
-    input_key_specs = (DataKeySpec('token_targets', split_dim=1, pad_value=0, owner='loss'),)
-    output_key_specs = (DataKeySpec('token_scores', split_dim=1, side='output', stage='output',
-                                       need_split=False, need_gather=True, autograd=True, owner='loss'),)
+    input_key_specs = (DataKeySpec('token_targets', split_dim=1, pad_value=0),)
+    output_key_specs = (DataKeySpec(
+        'token_scores',
+        split_dim=1,
+        side='output',
+        stage='output',
+        gather_at='loss',
+        autograd=True,
+    ),)
 
     def __call__(self, inputs, outputs, **kwargs):
         scores = outputs['token_scores']
diff --git a/tests/transformers/test_sp_data_keys.py b/tests/transformers/test_sp_data_keys.py
index 9a9213dd7..02c90e86a 100644
--- a/tests/transformers/test_sp_data_keys.py
+++ b/tests/transformers/test_sp_data_keys.py
@@ -36,7 +36,7 @@ def test_builtin_labels_match_existing_pad_split(rank, length):
 def test_custom_key_padding_dimension_roll_and_model_isolation(rank):
     strategy = _strategy(rank)
     other = _strategy(rank)
-    spec = DataKeySpec('token_weights', split_dim=1, pad_value=-2, need_roll=True)
+    spec = DataKeySpec('token_weights', split_dim=1, pad_value=-2, roll=True, pad_at='prepare', split_at='prepare')
     strategy.register_input_key(spec)
     strategy.register_input_key(spec)  # Identical declarations are safe.
     values = torch.arange(10).reshape(1, 5, 2)
@@ -53,7 +53,14 @@ def test_key_collection_executes_only_at_declared_boundary(rank):
     strategy = _strategy(rank)
     ctx = strategy._ctx
     ctx.tokenizer = type('Tokenizer', (), {'pad_token_id': 0})()
-    token_type_spec = DataKeySpec('token_type_ids', split_dim=1, pad_value=-1, stage='forward')
+    token_type_spec = DataKeySpec(
+        'token_type_ids',
+        split_dim=1,
+        pad_value=-1,
+        stage='forward',
+        pad_at='forward',
+        split_at='forward',
+    )
     strategy.register_input_key(token_type_spec)
     values = {
         'input_ids': torch.arange(5).expand(2, -1),
@@ -72,7 +79,7 @@ def test_key_collection_executes_only_at_declared_boundary(rank):
 
 
 def test_key_collection_rejects_duplicate_names_across_common_and_extra():
-    spec = DataKeySpec('token_type_ids', stage='forward')
+    spec = DataKeySpec('token_type_ids', stage='forward', pad_at='forward', split_at='forward')
     keys = SPKeyCollection()
     keys.add(SPKeyValue(spec, torch.ones(1, 2)), common=True)
     with pytest.raises(ValueError, match='Duplicate SP key'):
@@ -163,7 +170,7 @@ def test_processor_preserves_registered_keys_before_sp_split(monkeypatch, rank):
 
     monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu'))
     strategy = _strategy(rank)
-    strategy.register_input_key(DataKeySpec('token_weights', pad_value=-2, need_split=True))
+    strategy.register_input_key(DataKeySpec('token_weights', pad_value=-2, pad_at='prepare', split_at='prepare'))
     processor = InputProcessor()
     values = torch.arange(5).unsqueeze(0)
     sample = {'input_ids': torch.ones(1, 5, dtype=torch.long),
@@ -196,7 +203,13 @@ def test_registered_multidimensional_key_collation_and_loss_boundary(monkeypatch
 
     monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu'))
     strategy = _strategy()
-    strategy.register_input_key(DataKeySpec('token_features', pad_value=-2, need_gather=True))
+    strategy.register_input_key(DataKeySpec(
+        'token_features',
+        pad_value=-2,
+        gather_at='loss',
+        pad_at='prepare',
+        split_at='prepare',
+    ))
     processor = InputProcessor()
     samples = []
     for length in (3, 5):
@@ -221,7 +234,7 @@ def test_model_specific_forward_key_is_split_only_in_model_hook(monkeypatch):
     monkeypatch.setattr(Platform, 'get_local_device', lambda: torch.device('cpu'))
     strategy = _strategy(1)
     strategy._ctx.tokenizer = SimpleNamespace(pad_token_id=0)
-    strategy.register_input_key(DataKeySpec('token_type_ids', stage='forward', owner='model'))
+    strategy.register_input_key(DataKeySpec('token_type_ids', stage='forward', pad_at='forward', split_at='forward'))
 
     class Model(torch.nn.Module):
         def forward(self, token_type_ids, **kwargs):
@@ -242,9 +255,21 @@ def test_declared_collection_uses_sequence_axis_and_collects_without_logps(monke
 
     strategy = _strategy()
     strategy._ctx.extra_kwargs['position_ids'] = torch.arange(5).unsqueeze(0)
-    strategy.register_input_key(DataKeySpec('token_features', need_gather=True, pad_value=-2))
-    strategy.register_output_key(DataKeySpec('scores', split_dim=2, side='output', stage='output',
-                                            need_split=False, need_gather=True, autograd=True))
+    strategy.register_input_key(DataKeySpec(
+        'token_features',
+        gather_at='loss',
+        pad_value=-2,
+        pad_at='prepare',
+        split_at='prepare',
+    ))
+    strategy.register_output_key(DataKeySpec(
+        'scores',
+        split_dim=2,
+        side='output',
+        stage='output',
+        gather_at='loss',
+        autograd=True,
+    ))
     monkeypatch.setattr(strategy._ctx, 'gather', lambda value, dim, position_ids: torch.cat((value, value), dim=dim))
     monkeypatch.setattr(sp.GatherLoss, 'apply', lambda ctx, value, labels, dim, pos: (ctx.gather(value, dim, pos), None))
     local = torch.tensor([[[1, 2, 3], [4, 5, 6]]])
@@ -258,8 +283,14 @@ def test_packing_unpacks_registered_keys_on_declared_axis():
     from twinkle.processor import InputProcessor
     from twinkle.model.transformers.strategy.sequence_parallel.data_keys import BUILTIN_INPUT_KEYS
 
-    input_specs = (*BUILTIN_INPUT_KEYS, DataKeySpec('token_features', pad_value=-2, need_gather=True))
-    output_specs = (DataKeySpec('scores', side='output', stage='output', need_split=False, pad_value=-3),)
+    input_specs = (*BUILTIN_INPUT_KEYS, DataKeySpec(
+        'token_features',
+        pad_value=-2,
+        gather_at='loss',
+        pad_at='prepare',
+        split_at='prepare',
+    ))
+    output_specs = (DataKeySpec('scores', side='output', stage='output', pad_value=-3),)
     inputs = {'position_ids': torch.tensor([[0, 1, 2, 0, 1]]), 'labels': torch.tensor([[1, 2, 3, 4, 5]]),
               'completion_mask': torch.ones(1, 5), 'token_features': torch.tensor([[[1, 2, 3, 4, 5]]])}
     outputs = {'scores': torch.tensor([[[10, 20, 30, 40, 50]]])}
@@ -354,7 +385,7 @@ def test_common_loss_fields_collate_and_unpack_without_sp(monkeypatch, padding_f
     loss = SimpleNamespace(input_key_specs=(DataKeySpec('token_targets', split_dim=2, pad_value=-9),))
     input_specs, output_specs = resolve_key_specs(loss)
     spec = next(spec for spec in input_specs if spec.name == 'token_targets')
-    assert not spec.need_pad and not spec.need_split and not spec.need_gather
+    assert not spec.need_pad and not spec.need_split and not spec.gather_at
     samples = [dict(input_ids=list(range(length)), position_ids=list(range(length)),
                     labels=list(range(length)), token_targets=[[[1] * length, [2] * length]])
                for length in (3, 5)]
@@ -410,8 +441,15 @@ def test_prepare_then_forward_honors_separate_pad_and_split_boundaries(rank, sta
 
     strategy = _strategy(rank)
     strategy._ctx.tokenizer = SimpleNamespace(pad_token_id=0)
-    spec = DataKeySpec('token_type_ids', split_dim=1, pad_value=-7, stage=stage,
-                       model_input=True, pad_at='prepare', split_at='forward')
+    spec = DataKeySpec(
+        'token_type_ids',
+        split_dim=1,
+        pad_value=-7,
+        stage=stage,
+        model_input=True,
+        pad_at='prepare',
+        split_at='forward',
+    )
     strategy.register_input_key(spec)
 
     class Model(torch.nn.Module):
@@ -438,7 +476,7 @@ def test_direct_registration_survives_model_schema_resolution_and_processor(monk
     strategy = _strategy()
     spec = DataKeySpec('token_weights', pad_value=-2)
     strategy.register_input_key(spec)
-    output_spec = DataKeySpec('token_scores', side='output', stage='output', need_gather=True)
+    output_spec = DataKeySpec('token_scores', side='output', stage='output', gather_at='loss')
     strategy.register_output_key(output_spec)
     model = torch.nn.Linear(2, 2)
     wrapper = SimpleNamespace(model=model, strategy=SimpleNamespace(unwrap_model=lambda value: value),
@@ -510,9 +548,51 @@ def test_packed_output_logits_restore_without_labels():
 def test_loss_only_fields_reject_unreachable_forward_operations(operation):
     from dataclasses import replace
 
-    spec = (DataKeySpec('token_targets', pad_at='forward', need_split=False) if operation == 'pad' else
+    spec = (DataKeySpec('token_targets', pad_at='forward') if operation == 'pad' else
             DataKeySpec('token_targets', pad_at='prepare', split_at='forward'))
     with pytest.raises(ValueError, match='Loss-only input.*forward boundary'):
         _strategy().register_input_key(spec)
     # Model arguments can reach the forward hook, regardless of their stage name.
     _strategy().register_input_key(replace(spec, model_input=True))
+
+
+@pytest.mark.parametrize('pad_at,split_at', [(None, None), ('prepare', None), (None, 'forward'),
+                                             ('prepare', 'forward')])
+def test_parallel_switches_are_derived_from_boundaries(pad_at, split_at):
+    from dataclasses import fields
+
+    spec = DataKeySpec('token_features', pad_at=pad_at, split_at=split_at, model_input=True)
+    assert spec.need_pad is (pad_at is not None)
+    assert spec.need_split is (split_at is not None)
+    assert not {'need_pad', 'need_split', 'need_gather', 'owner'} & {field.name for field in fields(spec)}
+
+
+@pytest.mark.parametrize('kwargs,message', [
+    ({'pad_at': 'auxiliary'}, 'padding boundary'),
+    ({'split_at': 'output'}, 'split boundary'),
+    ({'gather_at': 'forward'}, 'gather boundary'),
+    ({'gather_at': 'loss'}, 'Only sharded input'),
+    ({'roll': True}, 'Rolling requires'),
+    ({'side': 'output', 'stage': 'output', 'pad_at': 'prepare'}, 'Output keys cannot'),
+])
+def test_invalid_parallel_operation_boundaries_are_rejected(kwargs, message):
+    with pytest.raises(ValueError, match=message):
+        DataKeySpec('token_features', **kwargs)
+
+
+@pytest.mark.parametrize('rank', [0, 1])
+@pytest.mark.parametrize('explicit_mask', [False, True])
+def test_single_sample_sdpa_builds_a_padded_global_attention_mask(rank, explicit_mask):
+    from types import SimpleNamespace
+
+    ctx = _strategy(rank)._ctx
+    ctx.attn_implementation = 'sdpa'
+    ctx.tokenizer = SimpleNamespace(pad_token_id=0)
+    positions = torch.arange(5).unsqueeze(0)
+    values = {'input_ids': positions.clone(), 'position_ids': positions,
+              'attention_mask': torch.ones(1, 5, dtype=torch.long) if explicit_mask else None}
+    specs = tuple(spec for spec in ctx.input_key_registry if spec.stage == 'forward')
+    keys = SPKeyCollection.from_mapping(values, specs)
+    ctx.pad_and_split_inputs(keys, boundary='forward', real_position_ids=positions)
+    torch.testing.assert_close(keys.get_value('attention_mask'), torch.tensor([[1, 1, 1, 1, 1, 0]]))
+    assert keys.get_value('input_ids').shape == (1, 3)
diff --git a/tests/transformers/test_sp_loss_contract.py b/tests/transformers/test_sp_loss_contract.py
index 692455ce1..a69a1e3ff 100644
--- a/tests/transformers/test_sp_loss_contract.py
+++ b/tests/transformers/test_sp_loss_contract.py
@@ -52,7 +52,7 @@ def __call__(self, inputs, outputs, **kwargs):
 
     strategy = _strategy()
     strategy.register_output_key(
-        DataKeySpec('values', split_dim=1, side='output', stage='output', need_split=False, autograd=True))
+        DataKeySpec('values', split_dim=1, side='output', stage='output', autograd=True))
     _mock_gather(monkeypatch)
     _, outputs = strategy.gather_loss_tensors({}, {'values': torch.tensor([[1.0, 2.0, 3.0]])},
                                              sp_gather_keys=ToyValueLoss.sp_gather_keys)
@@ -65,10 +65,23 @@ def __call__(self, inputs, outputs, **kwargs):
 def test_loss_can_declare_custom_sequence_inputs_and_outputs(monkeypatch):
     class TokenAlignmentLoss(Loss):
         require_logps = False
-        sp_input_key_specs = (DataKeySpec('token_targets', split_dim=1, pad_value=-1,
-                                          need_gather=True, autograd=True, owner='loss'),)
-        sp_output_key_specs = (DataKeySpec('token_scores', split_dim=1, side='output', stage='output',
-                                           need_split=False, need_gather=True, autograd=True, owner='loss'),)
+        sp_input_key_specs = (DataKeySpec(
+            'token_targets',
+            split_dim=1,
+            pad_value=-1,
+            gather_at='loss',
+            autograd=True,
+            pad_at='prepare',
+            split_at='prepare',
+        ),)
+        sp_output_key_specs = (DataKeySpec(
+            'token_scores',
+            split_dim=1,
+            side='output',
+            stage='output',
+            gather_at='loss',
+            autograd=True,
+        ),)
 
         def __call__(self, inputs, outputs, **kwargs):
             return LossOutput(loss=(outputs['token_scores'] - inputs['token_targets']).square().mean(),
@@ -380,3 +393,18 @@ def test_sp_pooling_and_gradients_match_full_sequence():
         sock.bind(('127.0.0.1', 0))
         port = sock.getsockname()[1]
     mp.spawn(_sp_pooling_worker, args=(port,), nprocs=2, join=True)
+
+
+@pytest.mark.parametrize('task', ['seq_cls', 'reranker'])
+def test_prepooled_classification_logits_skip_token_pooling(task):
+    from twinkle.model.transformers.task_outputs import postprocess_task_outputs
+
+    def unexpected_pool(*args):
+        raise AssertionError('Already pooled classification logits must bypass token pooling')
+
+    inputs = {'attention_mask': torch.ones(2, 5)}
+    outputs = {'logits': torch.tensor([[1., 2.], [3., 4.]])}
+    strategy = SimpleNamespace(pool_last_valid_token=unexpected_pool)
+    actual_inputs, actual_outputs = postprocess_task_outputs(inputs, outputs, task=task, sp_strategy=strategy)
+    assert actual_inputs is inputs
+    assert actual_outputs is outputs
diff --git a/tests/transformers/test_sp_model_adapters.py b/tests/transformers/test_sp_model_adapters.py
index cacbe04e7..05dd3d82b 100644
--- a/tests/transformers/test_sp_model_adapters.py
+++ b/tests/transformers/test_sp_model_adapters.py
@@ -19,7 +19,7 @@ def test_builtin_model_adapters_are_selected_by_model_family(model_type, expecte
     assert {adapter.name for adapter in sp_adapters.adapters_for_model(model)} == expected
 
 
-def test_model_adapter_registration_and_declared_key_specs(monkeypatch):
+def test_model_adapter_declared_key_specs(monkeypatch):
     class CustomAdapter(sp_adapters.ModelSPAdapter):
         name = 'test_custom_sp_adapter'
 
@@ -29,9 +29,8 @@ def matches(self, model):
         def input_key_specs(self):
             return (sp.DataKeySpec('custom_features', pad_value=-3),)
 
-    monkeypatch.setattr(sp_adapters, '_ADAPTERS', list(sp_adapters._ADAPTERS))
-    adapter = sp_adapters.register_model_sp_adapter(CustomAdapter)
-    assert adapter is sp_adapters.register_model_sp_adapter(CustomAdapter())
+    adapter = CustomAdapter()
+    monkeypatch.setattr(sp_adapters, '_ADAPTERS', [*sp_adapters._ADAPTERS, adapter])
     model = torch.nn.Linear(2, 2)
     model.config = SimpleNamespace(model_type='custom_sp_test')
     context = sp.SequenceParallel()