Admit hybrid models at Ulysses 8 under the exact contract (GDN-CP scan chain + composed admission)

src/xorl/data/collators/sequence_shard_collator.py

@@ -1,5 +1,6 @@
+import logging
 from dataclasses import dataclass
-from typing import Dict
+from typing import Dict, List, Tuple
 
 import torch
 
@@ -10,6 +11,84 @@
 from .packing_concat_collator import add_flash_attention_kwargs_from_position_ids
 
 
+logger = logging.getLogger("xorl.gdn_cp_collator")
+_alignment_engagement_logged = False
+
+# The GDN chunk kernels operate on a 64-token grid per document. The exact-CP
+# collator contract keeps every shard cut on every crossing document's own
+# grid by 64-aligning every document start via inter-document pad-docs and
+# making the per-rank shard length a multiple of 64.
+GDN_CP_CHUNK = 64
+
+
+def find_document_boundaries(position_ids_row: torch.Tensor) -> List[int]:
+    """Document boundaries of a packed stream from position-id resets."""
+    pos = position_ids_row.reshape(-1)
+    starts = (pos == 0).nonzero(as_tuple=False).reshape(-1).tolist()
+    if not starts or starts[0] != 0:
+        starts = [0, *starts]
+    return [*starts, pos.numel()]
+
+
+def gdn_cp_alignment_segments(doc_bounds: List[int], chunk: int = GDN_CP_CHUNK) -> List[Tuple[str, int, int]]:
+    """C2 splice plan: ('pad', 0, gap) and ('doc', start, end) segments such
+    that every document START in the spliced stream is a multiple of `chunk`,
+    and the stream END is too (so C1 tail pad-docs also start on the grid).
+
+    Pure function of the document layout so admission checks and the collator
+    use the same splice plan.
+    """
+    segments: List[Tuple[str, int, int]] = []
+    out_len = 0
+    for start, end in zip(doc_bounds[:-1], doc_bounds[1:]):
+        gap = (-out_len) % chunk
+        if gap:
+            segments.append(("pad", 0, gap))
+            out_len += gap
+        segments.append(("doc", start, end))
+        out_len += end - start
+    tail_gap = (-out_len) % chunk
+    if tail_gap:
+        segments.append(("pad", 0, tail_gap))
+    return segments
+
+
+def gdn_cp_spliced_length(segments: List[Tuple[str, int, int]]) -> int:
+    return sum(end - start if kind == "doc" else end for kind, start, end in segments)
+
+
+def apply_alignment_segments(
+    tensor: torch.Tensor,
+    segments: List[Tuple[str, int, int]],
+    dim: int,
+    pad_value: float | int = 0,
+    sequential_positions: bool = False,
+) -> torch.Tensor:
+    """Splice a token-aligned tensor along `dim` per the C2 plan. Pad segments
+    are filled with `pad_value`, or with `arange(gap)` when
+    `sequential_positions` (each inter-document pad is its own document)."""
+    if dim < 0:
+        dim = tensor.ndim + dim
+    pieces: List[torch.Tensor] = []
+    for kind, start, end in segments:
+        if kind == "doc":
+            pieces.append(tensor.narrow(dim, start, end - start))
+        else:
+            gap = end
+            shape = list(tensor.shape)
+            shape[dim] = gap
+            if sequential_positions:
+                seq = torch.arange(gap, device=tensor.device, dtype=tensor.dtype)
+                view = [1] * tensor.ndim
+                view[dim] = gap
+                pieces.append(seq.view(view).expand(shape).contiguous())
+            else:
+                pieces.append(
+                    torch.full(shape, fill_value=pad_value, dtype=tensor.dtype, device=tensor.device)
+                )
+    return torch.cat(pieces, dim=dim)
+
+
 def zigzag_reorder_packed_sequence(
     tensor: torch.Tensor,
     position_ids: torch.Tensor,
@@ -92,15 +171,27 @@ class TextSequenceShardCollator(DataCollator):
         fa_max_length_bucket: If > 0, round the flash-attn max_length up to a multiple of this value
             (upper bound only; correctness via cu_seqlens) to avoid torch.compile recompiling on ragged
             packs. 0 = off.
+        gdn_exact_cp_align: Engage the exact-GDN CP alignment contract:
+            64-align every document start via inter-document pad-docs and pad
+            the packed total to a multiple of 64 * cp_size, so every shard cut
+            lands on each crossing document's 64-token chunk grid.
+            Incompatible with ring attention.
     """
 
     pad_token_id: int = 0
     fa_max_length_bucket: int = 0
+    gdn_exact_cp_align: bool = False
 
     def __post_init__(self):
         self.cp_size = get_parallel_state().cp_size
         self.cp_rank = get_parallel_state().cp_rank
         self.ringattn_size = get_parallel_state().ringattn_size
+        if self.gdn_exact_cp_align and self.ringattn_size > 1:
+            raise ValueError(
+                "gdn_exact_cp_align is incompatible with ring attention "
+                "because zigzag reorder breaks the per-document 64-token "
+                "chunk grid. Fail closed.",
+            )
 
     def sp_slice(self, tensor: "torch.Tensor", dim: int = -1) -> "torch.Tensor":
         """
@@ -193,6 +284,37 @@ def __call__(self, batch: Dict[str, "torch.Tensor"]) -> Dict[str, "torch.Tensor"
                     f"This suggests data is not properly shifted."
                 )
 
+        # C2 (exact-GDN CP alignment): splice inter-document pad-docs so every
+        # document start (and the stream end) sits on the 64-token chunk grid.
+        # Runs BEFORE the _original_position_ids capture: after splicing, the
+        # spliced stream IS the stream the model consumes (minus C1 tail pads),
+        # so token-aligned consumers must see the spliced layout.
+        alignment_segments = None
+        if self.gdn_exact_cp_align:
+            doc_bounds = find_document_boundaries(position_ids[0])
+            alignment_segments = gdn_cp_alignment_segments(doc_bounds)
+            inserted = gdn_cp_spliced_length(alignment_segments) - input_ids.size(-1)
+            input_ids = apply_alignment_segments(input_ids, alignment_segments, -1, self.pad_token_id)
+            labels = apply_alignment_segments(labels, alignment_segments, -1, IGNORE_INDEX)
+            position_ids = apply_alignment_segments(
+                position_ids, alignment_segments, -1, sequential_positions=True
+            )
+            if "attention_mask" in batch:
+                batch["attention_mask"] = apply_alignment_segments(
+                    batch["attention_mask"], alignment_segments, -1, 1
+                )
+            # The pre-splice layout no longer matches the model's stream.
+            batch["_original_position_ids"] = position_ids.clone()
+            global _alignment_engagement_logged
+            if not _alignment_engagement_logged:
+                logger.info(
+                    "gdn-cp collator alignment engaged: +%d pad tokens on the first batch "
+                    "(%.3f%% of the spliced stream)",
+                    inserted,
+                    100.0 * inserted / max(1, position_ids.size(-1)),
+                )
+                _alignment_engagement_logged = True
+
         # Store original position_ids before padding for unpacking per-token outputs later
         if "_original_position_ids" not in batch:
             batch["_original_position_ids"] = position_ids.clone()
@@ -201,7 +323,11 @@ def __call__(self, batch: Dict[str, "torch.Tensor"]) -> Dict[str, "torch.Tensor"
         # With zigzag, each doc must be divisible by 2*ringattn_size sub-chunks,
         # and each sub-chunk must be divisible by ulysses_size. So total
         # sequence must be divisible by 2*ringattn_size*ulysses_size = 2*cp_size.
+        # C1 (exact-GDN CP alignment): the per-rank shard length must also be a
+        # multiple of the 64-token chunk grid.
         pad_multiple = 2 * self.cp_size if self.ringattn_size > 1 else self.cp_size
+        if self.gdn_exact_cp_align:
+            pad_multiple = GDN_CP_CHUNK * self.cp_size
         seq_length = input_ids.size(-1)
         cp_chunk_size = (seq_length + pad_multiple - 1) // pad_multiple * pad_multiple // self.cp_size
         pad_length = cp_chunk_size * self.cp_size - seq_length
@@ -289,6 +415,10 @@ def __call__(self, batch: Dict[str, "torch.Tensor"]) -> Dict[str, "torch.Tensor"
 
                 # Determine pad value: IGNORE_INDEX for target_tokens, 0 for others
                 pad_value = IGNORE_INDEX if field == "target_tokens" else 0.0
+                if alignment_segments is not None:
+                    field_tensor = apply_alignment_segments(
+                        field_tensor, alignment_segments, seq_dim, pad_value
+                    )
                 field_tensor = self.sp_padding(field_tensor, dim=seq_dim, pad_value=pad_value, pad_length=pad_length)
                 if self.ringattn_size > 1:
                     field_tensor = zigzag_reorder_packed_sequence(
@@ -316,4 +446,30 @@ def __call__(self, batch: Dict[str, "torch.Tensor"]) -> Dict[str, "torch.Tensor"
         else:
             add_flash_attention_kwargs_from_position_ids(batch, max_length_bucket=self.fa_max_length_bucket)
 
+        # Defense in depth for the exact-GDN CP contract: every shard cut must
+        # land on the 64-token chunk grid of whichever document it crosses
+        # (real doc or pad-doc — the GDN kernels cannot tell them apart).
+        # NOTE: pad-doc STARTS are unaligned by construction (they fill the
+        # gap up to the next 64-multiple) — that is fine precisely because a
+        # sub-64 gap pad can never contain a 64-multiple cut in its interior.
+        # Fail closed — a violated grid means byte-broken GDN handoffs.
+        if self.gdn_exact_cp_align:
+            if cp_chunk_size % GDN_CP_CHUNK != 0:
+                raise AssertionError(
+                    f"gdn_exact_cp_align postcondition violated: shard length "
+                    f"{cp_chunk_size} is not a multiple of {GDN_CP_CHUNK}",
+                )
+            bounds = find_document_boundaries(position_ids[0])
+            cuts = [cp_chunk_size * r for r in range(1, self.cp_size)]
+            bad = []
+            for cut in cuts:
+                doc_start = max(s for s in bounds[:-1] if s <= cut)
+                if cut != doc_start and (cut - doc_start) % GDN_CP_CHUNK != 0:
+                    bad.append((cut, doc_start))
+            if bad:
+                raise AssertionError(
+                    "gdn_exact_cp_align postcondition violated: shard cuts off "
+                    f"the crossing document's 64-token grid: {bad[:4]}",
+                )
+
         return batch

src/xorl/models/auto.py

@@ -386,12 +386,96 @@ def _validate_exact_qwen35_topology(config: PretrainedConfig, parallel_state: An
         if _is_qwen35_moe(config)
         else ((1, 1, 1, 1, 1, 1, 1, 1, 1, 1),)
     )
-    if topology not in admitted:
+    if topology in admitted:
+        return
+    if not _is_qwen35_moe(config) and _admit_qwen35_hybrid_ulysses(config, parallel_state):
+        return
+    raise ValueError(
+        "The Qwen3.5-family exact server-training path is admitted only for "
+        "WORLD/DP/DP-replicate/DP-shard/TP/PP/EP/CP/Ring/Ulysses="
+        f"{admitted}; got {topology}"
+    )
+
+
+# Hybrid (GDN + full-attention) dense Qwen3.5 under Ulysses: degrees covered
+# by the composed-program and FA4 head-bucket byte-contract tests.
+_QWEN35_HYBRID_ULYSSES_DEGREES = (2, 4, 8)
+_GDN_CP_COLLATOR_ATTESTATION_ENV = "XORL_GDN_CP_ALIGN_COLLATOR"
+
+
+def _admit_qwen35_hybrid_ulysses(config: PretrainedConfig, parallel_state: Any) -> bool:
+    """Conditional Ulysses admission for the exact hybrid dense program.
+
+    Returns False when the topology is not the pure-Ulysses hybrid shape (the
+    caller then raises its generic refusal). When the shape matches, every
+    requirement below must hold or this RAISES with an actionable message:
+
+    1. hybrid layer_types (a GDN-free dense config has no qualified U>1
+       program — its contract remains single-rank);
+    2. the Ulysses degree checks (heads divisible by the degree; GQA
+       replication degree divisible by kv-heads) — admission-time versions
+       of the pre-collective raises in UlyssesSyncStrategy;
+    3. the C1/C2 aligned-collator attestation
+       (``XORL_GDN_CP_ALIGN_COLLATOR=1``): admission cannot see the data
+       pipeline, so it demands an explicit attestation that
+       ``gdn_exact_cp_align`` is engaged; the GDN chain receipts and the
+       collator postcondition remain the runtime enforcement (fail-closed on
+       any misaligned shard cut);
+    4. the kernel/toolchain pin (first-class): a seeded pin directory whose
+       toolchain fingerprint matches this runtime, installed as this rank's
+       per-rank Triton cache before any kernel compiles.
+    """
+    import os  # noqa: PLC0415
+
+    u = parallel_state.ulysses_size
+    shape_matches = (
+        u in _QWEN35_HYBRID_ULYSSES_DEGREES
+        and parallel_state.world_size == u
+        and parallel_state.cp_size == u
+        and parallel_state.ringattn_size == 1
+        and parallel_state.dp_size == 1
+        and parallel_state.dp_replicate_size == 1
+        and parallel_state.dp_shard_size == 1
+        and parallel_state.tp_size == 1
+        and parallel_state.pp_size == 1
+        and parallel_state.ep_size == 1
+    )
+    if not shape_matches:
+        return False
+
+    layer_types = getattr(config, "layer_types", None) or []
+    if "linear_attention" not in layer_types:
+        return False  # GDN-free dense: keep the single-rank refusal.
+
+    num_heads = getattr(config, "num_attention_heads", None)
+    num_kv = getattr(config, "num_key_value_heads", None)
+    if not num_heads or num_heads % u != 0:
+        raise ValueError(
+            f"Exact hybrid Qwen3.5 at Ulysses {u}: num_attention_heads ({num_heads}) must be "
+            f"divisible by the Ulysses degree; an uneven head split cannot be scattered "
+            "byte-safely",
+        )
+    if not num_kv or (u > num_kv and u % num_kv != 0):
+        raise ValueError(
+            f"Exact hybrid Qwen3.5 at Ulysses {u}: num_key_value_heads ({num_kv}) must divide "
+            "the Ulysses degree for GQA replication",
+        )
+    if os.environ.get(_GDN_CP_COLLATOR_ATTESTATION_ENV) != "1":
         raise ValueError(
-            "The Qwen3.5-family exact server-training path is admitted only for "
-            "WORLD/DP/DP-replicate/DP-shard/TP/PP/EP/CP/Ring/Ulysses="
-            f"{admitted}; got {topology}"
+            f"Exact hybrid Qwen3.5 at Ulysses {u} requires the C1/C2 aligned collator "
+            f"(TextSequenceShardCollator(gdn_exact_cp_align=True)); set "
+            f"{_GDN_CP_COLLATOR_ATTESTATION_ENV}=1 to attest it. Without 64-aligned shard "
+            "cuts the GDN chain receipts fail closed at the first crossing document.",
         )
+    from xorl.ops.kernel_config_pin import pin_exact_kernel_configs  # noqa: PLC0415
+
+    pin_clone = pin_exact_kernel_configs()
+    logger.info(
+        "exact hybrid Qwen3.5 Ulysses-%d admission engaged: heads %d/%d kv, collator "
+        "attested, kernel pin -> %s",
+        u, num_heads, num_kv, pin_clone,
+    )
+    return True
 
 
 @dataclass(frozen=True)