feat(sp): refactor ulysses

docs/source_en/Components/Training Middleware/Sequence-Parallel.md

@@ -18,7 +18,6 @@ from twinkle.model.transformers.strategy.sequence_parallel import SequenceParall
 config = SequenceParallelConfig(
     enabled=True,           # Enable sequence parallel
     ulysses_size=None,      # Ulysses SP degree (auto-derived from DeviceMesh if None)
-    gather_logits=True,     # Gather logits after forward for loss computation
 )
 ```
 
@@ -41,16 +40,16 @@ device_mesh = DeviceMesh.from_sizes(
 
 1. **Pad** — input sequences are padded to a length divisible by SP world size
 2. **Split** — padded inputs are evenly split across SP ranks along the sequence dimension
-3. **Distributed Attention** — FlashAttention2 is patched to perform Ulysses all-to-all communication before/after attention computation
-4. **Gather** — after forward, logits are gathered back to full sequence length for loss computation
+3. **Distributed Attention** — SDPA and FlashAttention 2/3 are patched to perform Ulysses all-to-all communication before/after attention computation
+4. **Gather** — collect the present logps, entropies, values and requested logits before the unchanged full-sequence loss
 
 ## Supported Attention Backends
 
 | Backend | Status |
 |---------|--------|
-| FlashAttention2 | Fully supported (including packed/padding-free sequences) |
+| FlashAttention 2 / 3 | Ulysses and packed/padding-free paths; requires a compatible installed kernel |
 | SDPA | Supported (non-packed batches only) |
-| Derived Ring Attention | Supported with FlashAttention2 only (`rp_world_size > 1`) |
+| Derived Ring Attention | CUDA FlashAttention 2/3 only (`rp_world_size > 1`); incompatible platforms fail during initialization |
 
 ## Qwen3.5 Linear Attention
 
@@ -63,6 +62,10 @@ For MoE models, SP automatically installs a forward hook that gathers router log
 ## Key Constraints
 
 - `num_key_value_heads` must be divisible by `ulysses_size` (for Ulysses) or use ring attention fallback
-- Packed/padding-free batches require FlashAttention2
+- Packed/padding-free batches require FlashAttention 2 or 3
 - Derived ring attention requires `batch_size == 1` (packed format)
 - `torch.distributed` must be initialized
+
+## Isolating per-model SP contexts
+
+Each SequenceParallelStrategy owns its own SequenceParallel context, which stores that model's communication group, ranks, and current batch state, so algorithms that use multiple models, such as DPO, can enable SP training at the same time.

docs/source_zh/组件/训练中间件/序列并行.md

@@ -18,7 +18,6 @@ from twinkle.model.transformers.strategy.sequence_parallel import SequenceParall
 config = SequenceParallelConfig(
     enabled=True,           # 启用序列并行
     ulysses_size=None,      # Ulysses SP 并行度（若为 None 则从 DeviceMesh 自动推导）
-    gather_logits=True,     # 前向后聚合 logits 用于损失计算
 )
 ```
 
@@ -41,16 +40,16 @@ device_mesh = DeviceMesh.from_sizes(
 
 1. **填充** — 输入序列被填充到可被 SP 并行度整除的长度
 2. **分割** — 填充后的输入沿序列维度均匀分配到各 SP rank
-3. **分布式注意力** — FlashAttention2 被 patch 为在注意力计算前后执行 Ulysses all-to-all 通信
-4. **聚合** — 前向传播后，logits 被聚合回完整序列长度用于损失计算
+3. **分布式注意力** — SDPA 和 FlashAttention 2/3 被 patch 为在注意力计算前后执行 Ulysses all-to-all 通信
+4. **聚合** — 收集实际存在的 logps、entropies、values 及请求的 logits，再交给原有完整序列 loss
 
 ## 支持的注意力后端
 
 | 后端 | 状态 |
 |------|------|
-| FlashAttention2 | 完全支持（包括打包/padding-free 序列）|
+| FlashAttention 2 / 3 | 支持 Ulysses 和 packed/padding-free 路径，需要兼容的 kernel |
 | SDPA | 支持（仅非打包批次）|
-| 派生环形注意力 | 仅支持 FlashAttention2（`rp_world_size > 1`）|
+| 派生环形注意力 | 需要 CUDA FlashAttention 2/3；不兼容平台在初始化时明确报错 |
 
 ## Qwen3.5 线性注意力
 
@@ -63,6 +62,10 @@ SP 自动检测 Qwen3.5 GatedDeltaNet 线性注意力层，并应用 `Qwen3_5Gat
 ## 关键约束
 
 - `num_key_value_heads` 必须能被 `ulysses_size` 整除（Ulysses 模式），否则回退到环形注意力
-- 打包/padding-free 批次需要 FlashAttention2
+- 打包/padding-free 批次需要 FlashAttention 2 或 3
 - 派生环形注意力要求 `batch_size == 1`（打包格式）
 - `torch.distributed` 必须已初始化
+
+## 隔离每个模型的 SP 上下文
+
+每个 SequenceParallelStrategy 持有自己的 SequenceParallel 上下文，保存该模型的通信组、rank 和当前 batch 状态，以支持 dpo 这种多个模型的算法同时开启 sp 训练。

src/twinkle/model/transformers/moe/router_replay.py

@@ -185,8 +185,9 @@ def set_router_replay_data(
         # SP: slice full-sequence routed_experts (and the per-token replay_mask) to local SP rank tokens.
         # Both travel in one pad_and_split_inputs call so the mask is split identically to the indices and
         # stays token-aligned with them.
-        from ..strategy.sequence_parallel import sequence_parallel as sp
-        sp_world_size = getattr(sp, 'sp_world_size', None) or 1
+        from ..strategy.sequence_parallel import _MODULE_TO_CTX
+        sp = _MODULE_TO_CTX.get(model)
+        sp_world_size = getattr(sp, 'world_size', None) or 1
         if sp_world_size > 1:
             extra_split_values = [(routed_experts, 0, 1)]
             if replay_mask is not None:
@@ -255,8 +256,9 @@ 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
-        sp_world_size = getattr(sp, 'sp_world_size', None) or 1
+        from ..strategy.sequence_parallel import _MODULE_TO_CTX
+        sp = _MODULE_TO_CTX.get(model)
+        sp_world_size = getattr(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)