[Fix] try fix feat/qwen35 cp by upd zigzag

miles/backends/megatron_utils/actor.py

@@ -112,14 +112,14 @@ def init(
 
         from megatron.core import mpu
 
-        if mpu.get_context_parallel_world_size() > 1:
-            from miles.backends.training_utils.cp_utils import setup_hybrid_cp
+        cp_world_size = mpu.get_context_parallel_world_size()
+        if cp_world_size > 1:
+            from miles.backends.training_utils.cp_utils import detect_and_setup_hybrid_cp
 
             cp_group = mpu.get_context_parallel_group()
             cp_rank = mpu.get_context_parallel_rank()
-            cp_world_size = mpu.get_context_parallel_world_size()
             for model_chunk in self.model:
-                setup_hybrid_cp(model_chunk, cp_group, cp_rank, cp_world_size)
+                detect_and_setup_hybrid_cp(model_chunk, cp_group, cp_rank, cp_world_size)
 
         verify_megatron_parallel_state(self.model)
 

miles/backends/training_utils/cp_utils.py

@@ -8,6 +8,11 @@
 
 from .parallel import get_parallel_state
 
+try:
+    from fla.ops.cp import build_cp_context as _fla_build_cp_context
+except ImportError:
+    _fla_build_cp_context = None
+
 logger = logging.getLogger(__name__)
 
 
@@ -342,14 +347,37 @@ def slice_log_prob_with_cp(
         return torch.cat([chunk_1, chunk_2], dim=0)
 
 
-def setup_hybrid_cp(model: nn.Module, cp_group: dist.ProcessGroup, cp_rank: int, cp_world_size: int) -> None:
-    """Configure GatedDeltaNet modules for native fla CP instead of all-gather duplication.
+def build_gdn_cp_context(module: nn.Module, cu_seqlens: torch.Tensor, device: torch.device):
+    """Build fla CP context for a GatedDeltaNet module from packed sequence boundaries.
+
+    Args:
+        module: GDN module with ``cp_group`` / ``cp_world_size`` / ``conv_kernel_size``.
+        cu_seqlens: Global packed sequence boundaries (e.g. ``packed_seq_params.cu_seqlens_q``).
+        device: Target device.
 
-    Walks the model tree looking for HuggingfaceAttention submodules that have a
-    ``linear_attn`` child (i.e. DeltaNet layers). For each one it sets the CP
-    metadata so that ``_build_cp_context`` produces a valid context, and flips
-    ``hybrid_cp`` so the parent skips the all-gather path.
+    Returns ``None`` when CP is not configured on the module (``cp_group`` not set).
+    Raises ``RuntimeError`` if hybrid CP is configured but ``fla.ops.cp`` is missing.
     """
+    cp_group = getattr(module, "cp_group", None)
+    if cp_group is None:
+        return None
+    if _fla_build_cp_context is None:
+        raise RuntimeError(
+            "Hybrid CP requires fla.ops.cp (flash-linear-attention >= 0.4.2) " "but it could not be imported."
+        )
+    if cu_seqlens is None or cu_seqlens.numel() < 2:
+        raise ValueError(f"Hybrid CP requires valid cu_seqlens (at least 2 elements) but got {cu_seqlens}")
+    return _fla_build_cp_context(
+        cu_seqlens=cu_seqlens.to(device=device, dtype=torch.int32),
+        group=cp_group,
+        conv1d_kernel_size=module.conv_kernel_size,
+    )
+
+
+def detect_and_setup_hybrid_cp(
+    model: nn.Module, cp_group: dist.ProcessGroup, cp_rank: int, cp_world_size: int
+) -> None:
+    """Scan for GatedDeltaNet modules and configure them for native fla CP."""
     from miles_plugins.models.hf_attention import HuggingfaceAttention
 
     count = 0
@@ -364,4 +392,4 @@ def setup_hybrid_cp(model: nn.Module, cp_group: dist.ProcessGroup, cp_rank: int,
                 count += 1
 
     if count > 0:
-        logger.info(f"Configured hybrid CP on {count} DeltaNet modules (fla native state passing)")
+        logger.info(f"Configured hybrid CP on {count} GDN modules (fla native state passing)")

miles_plugins/models/hf_attention.py

@@ -38,6 +38,162 @@ def _fix_dtype(d):
         return ns
 
 
+def _sub_chunk_location(sub_id, cp_size):
+    """Return (zigzag_rank, half_index) for a given sub-chunk id.
+
+    In zigzag layout with N=cp_size ranks and 2N sub-chunks:
+      rank k holds [sub_k, sub_{2N-1-k}]
+    So sub_x lives on rank x (half 0) if x < N, else rank 2N-1-x (half 1).
+    """
+    if sub_id < cp_size:
+        return sub_id, 0
+    return 2 * cp_size - 1 - sub_id, 1
+
+
+def _p2p_exchange(send_bufs, send_dsts, recv_bufs, recv_srcs, cp_rank, cp_group):
+    """Exchange multiple buffers via batched async P2P, handling self-sends."""
+    # Handle self-sends as local copies
+    for i, dst in enumerate(send_dsts):
+        if dst == cp_rank:
+            for j, src in enumerate(recv_srcs):
+                if src == cp_rank and recv_bufs[j] is not send_bufs[i]:
+                    recv_bufs[j].copy_(send_bufs[i])
+
+    # Build P2P ops for remote exchanges (use group_peer for group-local ranks)
+    p2p_ops = []
+    for j, src in enumerate(recv_srcs):
+        if src != cp_rank:
+            p2p_ops.append(dist.P2POp(dist.irecv, recv_bufs[j], group_peer=src, group=cp_group))
+    for i, dst in enumerate(send_dsts):
+        if dst != cp_rank:
+            p2p_ops.append(dist.P2POp(dist.isend, send_bufs[i].contiguous(), group_peer=dst, group=cp_group))
+
+    if p2p_ops:
+        reqs = dist.batch_isend_irecv(p2p_ops)
+        for req in reqs:
+            req.wait()
+
+
+class _ZigzagToSequential(torch.autograd.Function):
+    """Convert zigzag CP layout to sequential layout for any CP size.
+
+    Zigzag rank k holds:     [sub_k, sub_{2N-1-k}]
+    Sequential rank j needs: [sub_{2j}, sub_{2j+1}]
+    """
+
+    @staticmethod
+    def forward(ctx, hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size):
+        ctx.cp_group = cp_group
+        ctx.cp_rank = cp_rank
+        ctx.cp_size = cp_size
+        ctx.save_for_backward(local_cu_seqlens)
+
+        N = cp_size
+
+        # Split local data into first_half (sub_k) and second_half (sub_{2N-1-k})
+        first_halves, second_halves = [], []
+        for i in range(len(local_cu_seqlens) - 1):
+            start, end = local_cu_seqlens[i].item(), local_cu_seqlens[i + 1].item()
+            mid = (start + end) // 2
+            first_halves.append(hidden_states[start:mid])
+            second_halves.append(hidden_states[mid:end])
+
+        my_bufs = [torch.cat(first_halves, dim=0), torch.cat(second_halves, dim=0)]
+        my_sub_ids = [cp_rank, 2 * N - 1 - cp_rank]
+        send_dsts = [sid // 2 for sid in my_sub_ids]
+
+        # What sequential rank cp_rank needs: sub_{2*cp_rank} and sub_{2*cp_rank+1}
+        need_ids = [2 * cp_rank, 2 * cp_rank + 1]
+        recv_srcs = [_sub_chunk_location(x, N)[0] for x in need_ids]
+        recv_bufs = [torch.empty_like(my_bufs[0]) for _ in range(2)]
+
+        # Handle self-send: if I send to myself, point recv_buf to send_buf
+        for i, dst in enumerate(send_dsts):
+            if dst == cp_rank:
+                for j, src in enumerate(recv_srcs):
+                    if src == cp_rank:
+                        recv_bufs[j] = my_bufs[i]
+
+        _p2p_exchange(my_bufs, send_dsts, recv_bufs, recv_srcs, cp_rank, cp_group)
+
+        return torch.cat(recv_bufs, dim=0)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (local_cu_seqlens,) = ctx.saved_tensors
+        # Backward: sequential → zigzag (inverse permutation)
+        result = _sequential_to_zigzag_impl(
+            grad_output, local_cu_seqlens, ctx.cp_group, ctx.cp_rank, ctx.cp_size
+        )
+        return result, None, None, None, None
+
+
+def _sequential_to_zigzag_impl(hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size):
+    """Core implementation for sequential → zigzag conversion."""
+    N = cp_size
+    half_len = hidden_states.shape[0] // 2
+
+    seq_bufs = [hidden_states[:half_len], hidden_states[half_len:]]
+    my_seq_sub_ids = [2 * cp_rank, 2 * cp_rank + 1]
+    send_dsts = [_sub_chunk_location(x, N)[0] for x in my_seq_sub_ids]
+
+    # Zigzag rank cp_rank needs sub_{cp_rank} and sub_{2N-1-cp_rank}
+    need_ids = [cp_rank, 2 * N - 1 - cp_rank]
+    recv_srcs = [nid // 2 for nid in need_ids]
+    recv_bufs = [torch.empty_like(seq_bufs[0]) for _ in range(2)]
+
+    for i, dst in enumerate(send_dsts):
+        if dst == cp_rank:
+            for j, src in enumerate(recv_srcs):
+                if src == cp_rank:
+                    recv_bufs[j] = seq_bufs[i]
+
+    _p2p_exchange(seq_bufs, send_dsts, recv_bufs, recv_srcs, cp_rank, cp_group)
+
+    # Reassemble zigzag: [first_half (sub_k), second_half (sub_{2N-1-k})]
+    result = []
+    half_chunk = half_len // max(len(local_cu_seqlens) - 1, 1)
+    offset_0, offset_1 = 0, 0
+    for i in range(len(local_cu_seqlens) - 1):
+        chunk_len = (local_cu_seqlens[i + 1].item() - local_cu_seqlens[i].item()) // 2
+        result.append(recv_bufs[0][offset_0 : offset_0 + chunk_len])
+        result.append(recv_bufs[1][offset_1 : offset_1 + chunk_len])
+        offset_0 += chunk_len
+        offset_1 += chunk_len
+    return torch.cat(result, dim=0)
+
+
+class _SequentialToZigzag(torch.autograd.Function):
+    """Convert sequential CP layout back to zigzag for any CP size."""
+
+    @staticmethod
+    def forward(ctx, hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size):
+        ctx.cp_group = cp_group
+        ctx.cp_rank = cp_rank
+        ctx.cp_size = cp_size
+        ctx.save_for_backward(local_cu_seqlens)
+        return _sequential_to_zigzag_impl(hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        (local_cu_seqlens,) = ctx.saved_tensors
+        # Backward: zigzag → sequential
+        result = _ZigzagToSequential.apply(
+            grad_output, local_cu_seqlens, ctx.cp_group, ctx.cp_rank, ctx.cp_size
+        )
+        return result, None, None, None, None
+
+
+def _zigzag_to_sequential(hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size):
+    """Convert zigzag CP layout to sequential layout."""
+    return _ZigzagToSequential.apply(hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size)
+
+
+def _sequential_to_zigzag(hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size):
+    """Convert sequential CP layout back to zigzag layout."""
+    return _SequentialToZigzag.apply(hidden_states, local_cu_seqlens, cp_group, cp_rank, cp_size)
+
+
 class _AllGatherForDuplicatedComputation(torch.autograd.Function):
     """All-gather whose backward just returns the local gradient slice (no reduce).
 
@@ -119,7 +275,18 @@ def forward(
                 group=mpu.get_tensor_model_parallel_group(),
             )
 
-        if mpu.get_context_parallel_world_size() > 1 and not self.hybrid_cp:
+        if mpu.get_context_parallel_world_size() > 1 and self.hybrid_cp:
+            cp_size = mpu.get_context_parallel_world_size()
+            local_cu_seqlens = cu_seqlens // cp_size
+            hidden_states = _zigzag_to_sequential(
+                hidden_states,
+                local_cu_seqlens,
+                mpu.get_context_parallel_group(),
+                mpu.get_context_parallel_rank(),
+                cp_size,
+            )
+
+        elif mpu.get_context_parallel_world_size() > 1:
             cp_size = mpu.get_context_parallel_world_size()
             # Use custom all-gather whose backward returns local gradient
             # instead of reduce-scatter, since the computation is duplicated.
@@ -154,7 +321,16 @@ def forward(
 
         output = output.permute(1, 0, 2)  # [seq_len, bsz, hidden_dim]
 
-        if mpu.get_context_parallel_world_size() > 1 and not self.hybrid_cp:
+        if mpu.get_context_parallel_world_size() > 1 and self.hybrid_cp:
+            output = _sequential_to_zigzag(
+                output,
+                local_cu_seqlens,
+                mpu.get_context_parallel_group(),
+                mpu.get_context_parallel_rank(),
+                cp_size,
+            )
+
+        elif mpu.get_context_parallel_world_size() > 1:
             cp_rank = mpu.get_context_parallel_rank()
             output_list = []
             for i in range(len(cu_seqlens) - 1):

miles_plugins/models/qwen3_5.py

@@ -15,11 +15,7 @@
 except ImportError:
     pass
 
-try:
-    from fla.ops.cp import FLACPContext, build_cp_context
-except ImportError:
-    FLACPContext = None
-    build_cp_context = None
+from miles.backends.training_utils.cp_utils import build_gdn_cp_context
 
 from .hf_attention import HuggingfaceAttention, _load_hf_config
 
@@ -87,27 +83,14 @@ def __init__(self, config, layer_idx: int):
 
         self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
-    def _build_cp_context(self, local_seq_len: int, device: torch.device):
-        """Build fla CP context from the local (sharded) sequence length."""
-        cp_group = getattr(self, "cp_group", None)
-        if cp_group is None or build_cp_context is None:
-            return None
-        global_seq_len = local_seq_len * self.cp_world_size
-        global_cu_seqlens = torch.tensor([0, global_seq_len], dtype=torch.int32, device=device)
-        return build_cp_context(
-            cu_seqlens=global_cu_seqlens,
-            group=cp_group,
-            conv1d_kernel_size=self.conv_kernel_size,
-        )
-
     def forward(
         self,
         hidden_states: torch.Tensor,
         cu_seqlens: torch.Tensor = None,
     ):
         batch_size, seq_len, _ = hidden_states.shape
 
-        cp_context = self._build_cp_context(seq_len, hidden_states.device)
+        cp_context = build_gdn_cp_context(self, cu_seqlens, hidden_states.device)
 
         # Projections (flat layout: [Q_all, K_all, V_all])
         mixed_qkv = self.in_proj_qkv(hidden_states)

miles_plugins/models/qwen3_next.py

@@ -18,11 +18,7 @@
 except ImportError:
     pass
 
-try:
-    from fla.ops.cp import FLACPContext, build_cp_context
-except ImportError:
-    FLACPContext = None
-    build_cp_context = None
+from miles.backends.training_utils.cp_utils import build_gdn_cp_context
 
 from .hf_attention import HuggingfaceAttention
 
@@ -80,19 +76,6 @@ def __init__(self, config, layer_idx: int):
 
         self.out_proj = nn.Linear(self.value_dim, self.hidden_size, bias=False)
 
-    def _build_cp_context(self, local_seq_len: int, device: torch.device):
-        """Build fla CP context from the local (sharded) sequence length."""
-        cp_group = getattr(self, "cp_group", None)
-        if cp_group is None or build_cp_context is None:
-            return None
-        global_seq_len = local_seq_len * self.cp_world_size
-        global_cu_seqlens = torch.tensor([0, global_seq_len], dtype=torch.int32, device=device)
-        return build_cp_context(
-            cu_seqlens=global_cu_seqlens,
-            group=cp_group,
-            conv1d_kernel_size=self.conv_kernel_size,
-        )
-
     def fix_query_key_value_ordering(self, mixed_qkvz, mixed_ba):
         """
         Derives `query`, `key` and `value` tensors from `mixed_qkvz` and `mixed_ba`.
@@ -127,7 +110,7 @@ def forward(
         hidden_states: torch.Tensor,
         cu_seqlens: torch.Tensor = None,
     ):
-        cp_context = self._build_cp_context(hidden_states.shape[1], hidden_states.device)
+        cp_context = build_gdn_cp_context(self, cu_seqlens, hidden_states.device)
 
         projected_states_qkvz = self.in_proj_qkvz(hidden_states)
         projected_states_ba = self.in_proj_ba(hidden_states)

tests/e2e/precision/test_qwen3_5_cp_correctness.py

@@ -86,11 +86,11 @@ def test_cp_forward_backward(rank, world_size):
     full_hidden_cp = torch.randn(batch, total_seq_len, 256, device=device, dtype=dtype)
     local_hidden = full_hidden_cp[:, start:end, :].clone().contiguous().requires_grad_(True)
 
-    # Local cu_seqlens for the chunk
-    local_cu = torch.tensor([0, local_seq_len], dtype=torch.int32, device=device)
+    # Global cu_seqlens (build_gdn_cp_context expects global boundaries)
+    global_cu = torch.tensor([0, total_seq_len], dtype=torch.int32, device=device)
 
     # Forward with CP
-    cp_out = model_cp(local_hidden, cu_seqlens=local_cu)
+    cp_out = model_cp(local_hidden, cu_seqlens=global_cu)
     cp_loss = cp_out.sum()
 
     # Reduce loss across ranks to match reference

tests/fast/backends/training_utils/__init__.py

@@ -0,0 +1 @@
+