Diff Coverage

    Returns:
        Mask aligned to the hidden sequence, or ``None``.
    """
    if token_mask is None or token_mask.shape[1] == hidden_states.shape[1]:
        return token_mask
    local_length = hidden_states.shape[1]
    if token_mask.shape[1] != local_length * tp_size:
        raise ValueError("Engram token mask is incompatible with TP/SP hidden")
    return token_mask[:, tp_rank * local_length:(tp_rank + 1) * local_length]


class NgramHashMapping(nn.Module):
    """Map token IDs to one Engram layer's disjoint n-gram hash rows."""

            ep_size=ep_size,
            padded_num_embeddings=self.padded_num_embeddings,
        )
        fused = self._fuse(hidden_states, embeddings)
        token_mask = local_token_mask(token_mask, hidden_states, tp_rank=tp_rank, tp_size=tp_size)
        if token_mask is None:
            return fused
        return torch.where(token_mask.to(torch.bool).unsqueeze(-1).unsqueeze(-1), fused, hidden_states)

                int(getattr(mesh_context, "tp_size", 1)) > 1 or int(getattr(mesh_context, "cp_size", 1)) > 1
        ):
            raise ValueError("OmniParallelBatch requires an initialized process group for TP or CP")
        self.device = device
        self.sequence_parallel = bool(mesh_context.sequence_parallel)
        self.runtime_input_adapter = runtime_input_adapter
        token_pad_values = {field: OMNI_CP_PAD_VALUES.get(field, 0) for field in OMNI_CP_TOKEN_FIELDS}
        cp_size = int(mesh_context.cp_size)
        tp_parts = int(mesh_context.tp_size) if mesh_context.sequence_parallel else 1
        alignment = lcm(2 * cp_size, cp_size * tp_parts)
        self.cp_sharder = CPBatchSharder(
            self.parallel_context, token_pad_values=token_pad_values,
            required_global_alignment=alignment,
        )
        self.tp_broadcaster = TPBatchBroadcaster(self.parallel_context, device)

                When supplied, fields not listed here remain complete on each CP rank.
        """
        self.parallel_context = parallel_context
        self.token_pad_values = token_pad_values
        if required_global_alignment is not None and required_global_alignment < 1:
            raise ValueError("required_global_alignment must be positive")
        self.required_global_alignment = required_global_alignment

    def shard(
            self,
            canonical_batch: Mapping[str, Any] | None,

    @staticmethod
    def _extend_cu_seq_lens(cp_batch: dict[str, Any], sequence_length: int, pad_length: int) -> None:
        """Keep packed sequence boundaries aligned with token padding."""
        if "cu_seq_lens" not in cp_batch:
            return
        boundaries = cp_batch["cu_seq_lens"].reshape(-1)
        if int(boundaries[0]) != 0 or int(boundaries[-1]) != sequence_length:
            raise ValueError("cu_seq_lens must cover the complete CP input sequence")
        if pad_length:
            boundaries = torch.cat((boundaries, boundaries[-1:] + pad_length))
        cp_batch["cu_seq_lens"] = boundaries

    def _shard_configured_fields(self, canonical_batch: Mapping[str, Any]) -> dict[str, Any]:
        """Shard configured token fields and preserve all modality metadata."""
        cp_batch = dict(canonical_batch)

        if not torch.is_tensor(input_ids) or input_ids.ndim != 2 or input_ids.shape[0] != 1:
            raise ValueError("Configured CP sharding requires input_ids with shape [1, sequence]")

        sequence_length = input_ids.shape[1]
        alignment = self.required_global_alignment or 2 * cp_size
        pad_length = (-sequence_length) % alignment
        self._extend_cu_seq_lens(cp_batch, sequence_length, pad_length)

        if cp_size == 1 and pad_length == 0:
            return cp_batch
        chunk_length = (sequence_length + pad_length) // cp_size
        local_start = self.parallel_context.cp_rank * chunk_length
        local_end = local_start + chunk_length
        for field, pad_value in self.token_pad_values.items():

    Returned relative to module_fqn (matching the name space of
    module.named_parameters(recurse=True)); the root spec (fqn "") treats
    every other boundary as a descendant.
    """
    boundaries = [fqn for fqn, spec in plan.modules.items() if spec.is_boundary]
    if module_fqn == "":
        return [f for f in boundaries if f]
    prefix = module_fqn + "."
    return [f[len(prefix):] for f in boundaries if f.startswith(prefix)]


def _bind_input_indices(boundary, module):
    """Bind the arg_name of in_plan to the positional index of the forward signature.

    @staticmethod
    def _resolve_isolated_roots(model: Any, arch: str) -> tuple[str, ...]:
        """Find subtrees whose activation layout is owned by a separate runtime."""
        from hyper_parallel.models.registry import get_model_adapter  # pylint: disable=C0415

        adapter = get_model_adapter(arch)
        if adapter is None or adapter.encoder_parallel is None:
            return ()
        spec = adapter.encoder_parallel()
        names = {name for name, _ in model.named_modules()}
        return tuple(root for root in spec.isolated_subtrees if root in names)

    def _classify_boundary_groups(self, model: Any, arch: str, isolated_roots=()):
        """Classify model parameters and group them into candidate boundaries."""
        param_roles = self._classify_all_params(model, arch)
        regular_roles = {
            name: role for name, role in param_roles.items()
            if not any(name.startswith(root + ".") for root in isolated_roots)
        }
        return param_roles, self._group_by_boundary(regular_roles)

    def plan(
        self,
        model: Any,

        # happens when _mark_hf_native_moe actually matches an HF-native
        # MoE (pre-stacked EP-aware modules use their own dispatcher and
        # are not subject to this constraint)
        # Phase 1: parameter role classification
        isolated_roots = self._resolve_isolated_roots(model, arch) if self._derive else ()
        param_roles, boundary_groups = self._classify_boundary_groups(model, arch, isolated_roots)

        # Phase 3+4: semantic inference + template-fills I/O
        # F4a/F4b: full shapes + requires_grad for the plan-time lints
        # (meta tensors carry shapes too — the zero-memory path is unaffected)

            model=model,
            param_ndims=param_metadata[0],
        )

        for root in isolated_roots:
            module = model.get_submodule(root)
            params = {name: {TP: Replicate()} for name, parameter in module.named_parameters()
                      if parameter.requires_grad}
            plan.modules[root] = ModuleShardingSpec(params=params, in_src={}, in_dst={}, is_boundary=False)

        # Phase 4.5: merge derived matches and expand fully declared globs
        # against the final model tree to insert new boundaries.
        self._finalize_boundary_specs(

    @staticmethod
    def _bias_owner_path(param_name: str) -> Optional[str]:
        """Return a bias owner's module path, or ``None`` for other params."""
        if param_name == "bias":
            return ""
        if param_name.endswith(".bias"):
            return param_name[: -len(".bias")]
        return None

    @staticmethod
    def _collect_deferred_biases(module_fqn: str, spec: ModuleShardingSpec, module: nn.Module) -> List[str]:
        """Validate one boundary's biases and collect rowwise deferred names."""

            out_name for out_name, named in out_src.items() if isinstance(named.get(TP), Partial)
        ]
        deferred: List[str] = []
        for param_name, param in named_params.items():
            owner_path = ShardingPlanner._bias_owner_path(param_name)
            if owner_path is None:
                continue
            weight_path = f"{owner_path}.weight" if owner_path else "weight"
            weight = named_params.get(weight_path)
            tp_placement = (spec.params or {}).get(weight_path)

    @staticmethod
    def _declared_parameter_names(plan: ShardingPlan) -> set[str]:
        """Collect full names covered by placements and special handlers."""
        covered = set()
        for fqn, spec in plan.modules.items():
            prefix = fqn + "." if fqn else ""
            for pname in (spec.params or {}):
                covered.add(prefix + pname)
            for sources in (getattr(spec, "_ep_stack", None) or {}).values():
                covered.update(prefix + source for source in sources)
        covered.update(plan.special_handlers or {})
        return covered

    def _check_all_trainable_params_covered(self, plan: ShardingPlan, model) -> None:
        """F4b: every ``requires_grad=True`` parameter must appear in some
        spec.params (resolved to a full FQN) or in special_handlers —

        assembler for subtrees — coverage is the caller's responsibility).
        """
        if not self._derive:
            return
        covered = self._declared_parameter_names(plan)
        uncovered = [
            name for name, p in model.named_parameters()
            if p.requires_grad and name not in covered
        ]

        """D-14 (05 §13.2): compile-time chain propagation/validation is
        removed; only ``_is_terminal`` marking is kept — the last boundary
        in forward order (each module vouches for its own propagation in
        validate mode; adjacent-contract checks no longer exist)."""
        boundaries = [fqn for fqn, spec in plan.modules.items() if spec.is_boundary]
        sorted_fqns = self._topological_sort_by_forward_order(boundaries, model)
        terminal = sorted_fqns[-1] if sorted_fqns else None
        for fqn, spec in plan.modules.items():
            # The planner owns the spec DSL internals.
            spec._is_terminal = fqn == terminal  # pylint: disable=protected-access

def _validate_tp_model_compatibility(config: Any, tp_size: int, ep_size: int) -> None:
    """Check model head and expert width divisibility on the TP axis."""
    heads = getattr(config, "num_attention_heads", None)
    if heads is not None and heads % tp_size != 0:
        raise ValueError(
            f"num_attention_heads ({heads}) must be divisible by TP ({tp_size})"
        )
    kv_heads = getattr(config, "num_key_value_heads", None)
    if kv_heads is not None and kv_heads % tp_size != 0:
        raise ValueError(
            f"num_key_value_heads ({kv_heads}) must be divisible by TP ({tp_size})"
        )
    moe_inter = getattr(config, "moe_intermediate_size", None)
    if ep_size > 1 and moe_inter is not None and moe_inter % tp_size != 0:
        raise ValueError(
            f"moe_intermediate_size ({moe_inter}) must be divisible by TP ({tp_size})"
        )

def _validate_ep_model_compatibility(config: Any, ep_size: int) -> None:
    """Check that the model has a divisible routed-expert count."""
    num_experts = (getattr(config, "num_experts", None)
                   or getattr(config, "n_routed_experts", None) or 0)
    if num_experts <= 0:
        raise ValueError("EP>1 requires MoE model (num_experts > 0)")
    if num_experts % ep_size != 0:
        raise ValueError(
            f"num_experts ({num_experts}) must be divisible by EP ({ep_size})"
        )

    if config is None:
        return

    if tp_size > 1:
        _validate_tp_model_compatibility(config, tp_size, ep_size)

    if cp_size > 1 and seq_len is not None and seq_len % (cp_size * 2) != 0:
        raise ValueError(
            f"seq_len ({seq_len}) must be divisible by 2*cp ({2 * cp_size})"

            f"seq_len ({seq_len}) must be divisible by 2*cp ({2 * cp_size})"
        )

    if ep_size > 1:
        _validate_ep_model_compatibility(config, ep_size)

        Args:
            module_fqn: Final module path for the attention block.
        """
        self.seen[module_fqn] = self.seen.get(module_fqn, 0) + 1

    def assert_exactly_once(self) -> None:
        """Reject missing or repeated calls in an active ViT pass."""
        if set(self.seen) != self.expected or any(count != 1 for count in self.seen.values()):
            raise RuntimeError(
                "[HP-ENCODER-008] vision attention execution mismatch; "
                f"expected={sorted(self.expected)}; seen={self.seen}"
            )

    def validate_runtime_head_counts(self, *, module_fqn: str, query_heads: int,
                                     key_heads: int, value_heads: int) -> None:
        """Check projected head counts before entering a collective."""
        actual = {"query": query_heads, "key": key_heads, "value": value_heads}
        if actual != self.declared_heads or any(count % self.cp_mesh.size() for count in actual.values()):
            raise ValueError(
                "[HP-ENCODER-005] Ulysses runtime head count mismatch; phase=runtime; "
                f"module_fqn={module_fqn}; declared_head_counts={self.declared_heads}; "
                f"runtime_head_counts={actual}; encoder_cp_size={self.cp_mesh.size()}"
            )

            )

    def attend_segments(self, query: torch.Tensor, key: torch.Tensor, value: torch.Tensor) -> torch.Tensor:
        """Run full-image attention independently for each local image segment."""
        outputs = []
        for segment in self.segments:
            start = segment.packed_offset
            stop = start + segment.local_padded_length
            q = ulysses_seq_to_head(query[start:stop], seq_dim=0, head_dim=1, cp_mesh=self.cp_mesh)
            k = ulysses_seq_to_head(key[start:stop], seq_dim=0, head_dim=1, cp_mesh=self.cp_mesh)
            v = ulysses_seq_to_head(value[start:stop], seq_dim=0, head_dim=1, cp_mesh=self.cp_mesh)
            global_length = segment.local_padded_length * self.cp_mesh.size()
            mask = None
            if not segment.is_dummy and segment.item.sequence_length < global_length:
                indices = torch.arange(global_length, device=q.device)
                mask = (indices < segment.item.sequence_length).view(1, 1, -1)
            attended = functional.scaled_dot_product_attention(  # pylint: disable=not-callable
                q.transpose(0, 1), k.transpose(0, 1), v.transpose(0, 1), attn_mask=mask,
                is_causal=False,
            ).transpose(0, 1)
            outputs.append(ulysses_head_to_seq(attended, seq_dim=0, head_dim=1, cp_mesh=self.cp_mesh))
        return torch.cat(outputs, dim=0)


_ACTIVE_CONTEXT: ContextVar[EncoderAttentionContext | None] = ContextVar("encoder_attention_context", default=None)

def current_encoder_attention_context() -> EncoderAttentionContext | None:
    """Return the complete attention state for the current ViT pass."""
    return _ACTIVE_CONTEXT.get()


@contextmanager
def activate_encoder_attention(context: EncoderAttentionContext) -> Iterator[EncoderAttentionContext]:

    Args:
        context: Segments, executor, and expected module paths for this pass.
    """
    token = _ACTIVE_CONTEXT.set(context)
    try:
        yield context
    finally:
        _ACTIVE_CONTEXT.reset(token)

    def __init__(self, *, runtime: EncoderParallelRuntime,
                 inner: RuntimeInputAdapter | None) -> None:
        """Combine visual planning with the batch's existing input provider."""
        if inner is not None and "encoder_parallel_plan" in inner.runtime_input_fields():
            raise ValueError("[HP-DATA-001] encoder_parallel_plan has two owners")
        self.runtime = runtime
        self.inner = inner

    def runtime_input_fields(self) -> tuple[str, ...]:
        """Declare the existing fields and the plan field as one input contract."""
        fields = () if self.inner is None else self.inner.runtime_input_fields()
        return (*fields, "encoder_parallel_plan")

    def build_runtime_inputs(
            self, *, batch: Mapping[str, Any], context: RuntimeInputContext,
    ) -> Mapping[str, Any]:

    def build_runtime_inputs(
            self, *, batch: Mapping[str, Any], context: RuntimeInputContext,
    ) -> Mapping[str, Any]:
        """Validate locally, build other inputs, then agree on the visual plan."""
        prepared = self.runtime.preflight_batch(batch)
        generated = {} if self.inner is None else self.inner.build_runtime_inputs(batch=batch, context=context)
        if not isinstance(generated, Mapping):
            raise TypeError("RuntimeInputAdapter.build_runtime_inputs must return a mapping")
        inputs = dict(generated)
        if "encoder_parallel_plan" in batch or "encoder_parallel_plan" in inputs or set(batch).intersection(inputs):
            raise ValueError("[HP-DATA-001] runtime inputs cannot replace Omni model inputs")
        inputs["encoder_parallel_plan"] = self.runtime.finalize_plan(*prepared)
        return inputs

    enabled: bool = False
    encoder_cp_size: int = 1

    def __post_init__(self) -> None:
        if not isinstance(self.enabled, bool):
            raise ValueError("[HP-ENCODER-002] enabled must be a boolean")
        if (isinstance(self.encoder_cp_size, bool) or not isinstance(self.encoder_cp_size, int)
                or self.encoder_cp_size < 1):
            raise ValueError(
                "[HP-ENCODER-002] accelerator.encoder_parallel.encoder_cp_size "
                f"must be a positive integer, got {self.encoder_cp_size!r}"
            )

        Args:
            raw: User's optional encoder-parallel configuration mapping.
        """
        if raw is None:
            return cls()
        if not isinstance(raw, Mapping):
            raise ValueError("[HP-ENCODER-002] encoder_parallel must be a mapping")
        unknown = set(raw) - {field.name for field in fields(cls)}
        if unknown:
            raise ValueError(f"[HP-ENCODER-002] unknown encoder_parallel keys: {sorted(unknown)}")
        return cls(**raw)

    @property
    def is_dummy(self) -> bool:
        """A missing image identifies the synchronization-only segment."""
        return self.item is None


EncoderFeature = tuple[EncoderWorkItem, torch.Tensor]
"""One image job paired with its aligned feature rows."""

    Args:
        mesh_context: Decoder CP, TP, and sequence-parallel coordinates.
    """
    if mesh_context.cp_size > 1:
        return FeatureExchangeMode.PARENT_VARIABLE_A2A
    if mesh_context.tp_size == 1:
        return FeatureExchangeMode.LOCAL
    if mesh_context.sequence_parallel:
        return FeatureExchangeMode.TP_VARIABLE_A2A
    return FeatureExchangeMode.TP_VARIABLE_ALL_GATHER


def cat_or_empty(rows: Sequence[torch.Tensor], *, anchor: torch.Tensor,
                 feature_width: int) -> torch.Tensor:

        rows: Feature tensors to concatenate.
        anchor: Tensor supplying the zero-row gradient edge.
        feature_width: Width of an empty feature tensor.
    """
    if rows:
        return torch.cat(tuple(rows), dim=0)
    return anchor.sum().reshape(1, 1).expand(0, feature_width)


def _ordered_source_records(buckets: Sequence[Sequence[EncoderWorkItem]], encoder_cp_size: int
                            ) -> dict[int, tuple[EncoderWorkItem, ...]]:
    return {bucket_rank * encoder_cp_size: tuple(bucket) for bucket_rank, bucket in enumerate(buckets)}


def _route_tp_all_gather(*, local_by_id: Mapping[int, torch.Tensor],
                         source_records: Mapping[int, tuple[EncoderWorkItem, ...]],

                         source_records: Mapping[int, tuple[EncoderWorkItem, ...]],
                         prototype: torch.Tensor, parent_rank: int,
                         exchange_mesh: Any) -> tuple[torch.Tensor, tuple[EncoderWorkItem, ...]]:
    """Gather each producer's complete records to every decoder TP rank."""
    group_size = exchange_mesh.size()
    local_items = source_records.get(parent_rank, ())
    local_rows = cat_or_empty([local_by_id[item.item_id] for item in local_items],
                              anchor=prototype, feature_width=int(prototype.shape[-1]))
    splits = [sum(item.output_rows for item in source_records.get(rank, ())) for rank in range(group_size)]
    rows = differentiable_variable_all_gather(local_rows, output_splits=splits,
                                               group=exchange_mesh.get_group())
    received_items = tuple(item for rank in range(group_size) for item in source_records.get(rank, ()))
    return rows, received_items


def _route_variable_a2a(*, mode: FeatureExchangeMode,
                        local_by_id: Mapping[int, torch.Tensor],

                        destinations: Mapping[int, tuple[int, ...]],
                        encoder_cp_size: int, exchange_mesh: Any
                        ) -> tuple[torch.Tensor, tuple[EncoderWorkItem, ...]]:
    """Send each full record only to its exact decoder token owners."""
    group_size = exchange_mesh.size()
    local_rank = exchange_mesh.get_local_rank()
    if group_size != len(source_records) * encoder_cp_size:
        raise ValueError(f"[HP-ENCODER-009] {mode.value} exchange size disagrees with encoder parent")
    width = int(prototype.shape[-1])
    send_chunks = []
    input_splits = []
    for destination_rank in range(group_size):
        items = [item for item in source_records.get(parent_rank, ())
                 if destination_rank in destinations[item.item_id]]
        chunk = cat_or_empty([local_by_id[item.item_id] for item in items],
                             anchor=prototype, feature_width=width)
        send_chunks.append(chunk)
        input_splits.append(chunk.shape[0])
    packed = cat_or_empty(send_chunks, anchor=prototype, feature_width=width)
    output_splits = []
    received_items = []
    for source_rank in range(group_size):
        items = [item for item in source_records.get(source_rank, ())
                 if local_rank in destinations[item.item_id]]
        output_splits.append(sum(item.output_rows for item in items))
        received_items.extend(items)
    rows = differentiable_all_to_all_single(packed, input_splits=input_splits,
                                            output_splits=output_splits, group=exchange_mesh.get_group())
    return rows, tuple(received_items)


def route_features(*, mode: FeatureExchangeMode, producer_records: Sequence[EncoderFeature],
                   prototype: torch.Tensor, parent_rank: int, destinations: Mapping[int, tuple[int, ...]],

                   prototype: torch.Tensor, parent_rank: int, destinations: Mapping[int, tuple[int, ...]],
                   buckets: Sequence[Sequence[EncoderWorkItem]], encoder_cp_size: int,
                   exchange_mesh: Any) -> tuple[torch.Tensor, tuple[EncoderWorkItem, ...]]:
    """Exchange full feature records using a deterministic plan shared by all ranks."""
    width = int(prototype.shape[-1])
    local_by_id = {item.item_id: feature for item, feature in producer_records}
    source_records = _ordered_source_records(buckets, encoder_cp_size)
    if mode is FeatureExchangeMode.LOCAL:
        received_items = tuple(item for item, _ in producer_records)
        rows = cat_or_empty([feature for _, feature in producer_records], anchor=prototype, feature_width=width)
        return rows, received_items

    if mode is FeatureExchangeMode.TP_VARIABLE_ALL_GATHER:
        return _route_tp_all_gather(local_by_id=local_by_id, source_records=source_records,
                                    prototype=prototype, parent_rank=parent_rank,
                                    exchange_mesh=exchange_mesh)
    if mode in (FeatureExchangeMode.TP_VARIABLE_A2A, FeatureExchangeMode.PARENT_VARIABLE_A2A):
        return _route_variable_a2a(
            mode=mode, local_by_id=local_by_id, source_records=source_records,
            prototype=prototype, parent_rank=parent_rank, destinations=destinations,
            encoder_cp_size=encoder_cp_size, exchange_mesh=exchange_mesh,
        )
    raise ValueError(f"[HP-ENCODER-009] unsupported feature exchange mode: {mode!r}")


def unpack_records(rows: torch.Tensor, received_items: Sequence[EncoderWorkItem]
                   ) -> tuple[EncoderFeature, ...]:

    Args:
        rows: Received aligned feature rows.
        received_items: Images in source-rank then bucket order.
    """
    records = []
    offset = 0
    for item in received_items:
        stop = offset + item.output_rows
        records.append((item, rows[offset:stop]))
        offset = stop
    if offset != rows.shape[0]:
        raise ValueError(
            "[HP-ENCODER-009] transported rows disagree with planned images; "
            f"rows={rows.shape[0]}; described={offset}"
        )
    return tuple(sorted(records, key=lambda record: (
        record[0].batch_index, record[0].decoder_token_start, record[0].item_id
    )))

        source_dims: Decoder axes forming one visual parent group.
        target_sizes: Encoder-DP and Encoder-CP degrees.
        target_names: Names assigned to the derived axes.
    """
    source = root_mesh[source_dims]
    if math.prod(target_sizes) != math.prod(source.mesh_shape) or len(target_sizes) != len(target_names):
        raise ValueError("[HP-ENCODER-002] encoder mesh shape does not match decoder parent")
    flat = source.flatten("encoder_parent")
    return flat._unflatten("encoder_parent", target_sizes, target_names)  # pylint: disable=protected-access


def encoder_coordinates(parent_rank: int, encoder_cp_size: int) -> tuple[int, int]:
    """Return Encoder-DP and Encoder-CP coordinates in parent rank order.

    Args:
        parent_rank: Linear rank within the decoder CP-by-TP parent.
        encoder_cp_size: Number of sequence shards per visual bucket.
    """
    return divmod(parent_rank, encoder_cp_size)


def decoder_coordinates(parent_rank: int, decoder_tp_size: int) -> tuple[int, int]:
    """Return decoder CP and TP coordinates in parent rank order.

    Args:
        parent_rank: Linear rank within the decoder CP-by-TP parent.
        decoder_tp_size: Number of decoder TP ranks per CP rank.
    """
    return divmod(parent_rank, decoder_tp_size)

    Args:
        mesh_context: Decoder CP, TP, and sequence-parallel topology.
    """
    cp_size = int(mesh_context.cp_size)
    tp_parts = int(mesh_context.tp_size) if mesh_context.sequence_parallel else 1
    return lcm(2 * cp_size, cp_size * tp_parts)


@dataclass(frozen=True)
class DecoderSequenceLayout:

    @classmethod
    def build(cls, *, mesh_context: Any, cp_local_length: int,
              cp_rank: int | None = None, tp_rank: int | None = None) -> "DecoderSequenceLayout":
        """Derive the token interval after embedding's TP reduce-scatter."""
        cp_size = int(mesh_context.cp_size)
        tp_size = int(mesh_context.tp_size)
        cp_rank = int(mesh_context.cp_rank if cp_rank is None else cp_rank)
        tp_rank = int(mesh_context.tp_rank if tp_rank is None else tp_rank)
        sp = bool(mesh_context.sequence_parallel)
        if not 0 <= cp_rank < cp_size or not 0 <= tp_rank < tp_size:
            raise ValueError(
                "[HP-ENCODER-004] decoder CP/TP coordinate is outside topology; "
                f"cp_rank={cp_rank}; cp_size={cp_size}; tp_rank={tp_rank}; tp_size={tp_size}"
            )
        if cp_local_length < 1 or (sp and cp_local_length % tp_size):
            raise ValueError(
                "[HP-ENCODER-004] invalid decoder sequence layout; "
                f"sequence_parallel={sp}; cp_local_length={cp_local_length}; tp_size={tp_size}; "
                "constraint=cp_local_length % tp_size == 0 when SP is enabled"
            )
        local_length = cp_local_length // (tp_size if sp else 1)
        local_start = (tp_rank if sp else 0) * local_length
        global_start = cp_rank * cp_local_length + local_start
        return cls(sp, cp_rank, cp_size, tp_rank, tp_size, cp_local_length,
                   cp_size * cp_local_length, local_start, local_start + local_length,
                   global_start, global_start + local_length)

    def localize(self, value: torch.Tensor, *, sequence_dim: int = 1) -> torch.Tensor:

        Args:
            value: Auxiliary tensor replicated across decoder TP ranks.
            sequence_dim: Dimension containing CP-local token rows.
        """
        if value.shape[sequence_dim] != self.cp_local_length:
            raise ValueError(
                "[HP-ENCODER-004] CP-local token tensor length mismatch; "
                f"actual={value.shape[sequence_dim]}; expected={self.cp_local_length}; "
                f"cp_rank={self.cp_rank}; tp_rank={self.tp_rank}"
            )
        return value.narrow(sequence_dim, self.cp_local_start,
                            self.cp_local_stop - self.cp_local_start).contiguous()

    def validate_hidden(self, hidden: torch.Tensor, *, sequence_dim: int = 1) -> None:
        """Fail before merging against an unexpected decoder hidden layout.

        Args:
            hidden: Decoder hidden tensor after embedding redistribution.
            sequence_dim: Dimension containing decoder token rows.
        """
        expected = self.global_stop - self.global_start
        if hidden.shape[sequence_dim] != expected:
            raise ValueError(
                "[HP-ENCODER-004] decoder hidden does not match planned CP/SP shard; "
                f"actual={hidden.shape[sequence_dim]}; expected={expected}; "
                f"global_interval=({self.global_start},{self.global_stop}); "
                f"cp_rank={self.cp_rank}; tp_rank={self.tp_rank}"

        items: Images with valid patch counts and stable IDs.
        encoder_dp_size: Number of visual work buckets.
        encoder_cp_size: Sequence-sharding degree of each bucket.
    """
    if encoder_dp_size < 1 or encoder_cp_size < 1:
        raise ValueError("[HP-ENCODER-002] encoder DP and CP sizes must be positive")
    buckets: list[list[EncoderWorkItem]] = [[] for _ in range(encoder_dp_size)]
    loads = [0] * encoder_dp_size
    for item in sorted(items, key=lambda value: (-value.sequence_length ** 2, value.batch_index, value.item_id)):
        bucket = min(range(encoder_dp_size), key=lambda rank: (loads[rank], rank))
        buckets[bucket].append(item)
        loads[bucket] += item.sequence_length ** 2 // encoder_cp_size
    return tuple(tuple(bucket) for bucket in buckets)


def build_segments(items: Sequence[EncoderWorkItem], encoder_cp_rank: int,
                   encoder_cp_size: int) -> tuple[EncoderSegment, ...]:

        items: Images assigned to this Encoder-DP bucket.
        encoder_cp_rank: Sequence-shard coordinate within the bucket.
        encoder_cp_size: Number of shards per image.
    """
    if not 0 <= encoder_cp_rank < encoder_cp_size:
        raise ValueError("[HP-ENCODER-002] invalid encoder CP rank")
    segments = []
    offset = 0
    for item in items:
        if item.sequence_length < 1:
            raise ValueError("[HP-ENCODER-013] image patch count must be positive")
        local_length = (item.sequence_length + encoder_cp_size - 1) // encoder_cp_size
        start = encoder_cp_rank * local_length
        valid = max(0, min(local_length, item.sequence_length - start))
        segments.append(EncoderSegment(item, start, valid, local_length, offset))
        offset += local_length
    if not segments:
        segments.append(EncoderSegment(None, 0, 1, 1, 0))
    return tuple(segments)


def build_destinations(items: Sequence[EncoderWorkItem], layouts: Mapping[int, DecoderSequenceLayout]
                       ) -> dict[int, tuple[int, ...]]:

    Args:
        items: Images carrying their global decoder token spans.
        layouts: Decoder token interval for every rank in the parent group.
    """
    lengths = {layout.global_sequence_length for layout in layouts.values()}
    if len(lengths) != 1:
        raise ValueError("[HP-ENCODER-004] inconsistent decoder global lengths")
    global_length = next(iter(lengths))
    result = {}
    for item in items:
        start = item.decoder_token_start
        stop = start + item.decoder_token_length
        if item.decoder_token_length <= 0 or start < 0 or stop > global_length:
            raise ValueError(
                "[HP-ENCODER-016] image span is outside the decoder domain; "
                f"item_id={item.item_id}; span=({start},{stop}); global_sequence_length={global_length}"
            )
        owners = tuple(rank for rank, layout in layouts.items()
                       if max(start, layout.global_start) < min(stop, layout.global_stop))
        if not owners:
            raise ValueError(f"[HP-ENCODER-016] image span has no decoder token owner; item_id={item.item_id}")
        result[item.item_id] = owners
    return result

def _invalid_head_count(count: Any) -> bool:
    """Reject booleans, nonintegers, and zero-sized QKV head axes."""
    return isinstance(count, bool) or not isinstance(count, int) or count < 1


def validate_ulysses_head_divisibility(*, encoder_cp_size: int, head_counts: Mapping[str, int],
                                       attention_pattern: str, model: nn.Module) -> None:

def validate_ulysses_head_divisibility(*, encoder_cp_size: int, head_counts: Mapping[str, int],
                                       attention_pattern: str, model: nn.Module) -> None:
    """Check every Q/K/V head count before deriving any Encoder-CP group."""
    required = ("query", "key", "value")
    missing = [name for name in required if name not in head_counts]
    invalid = {name: head_counts[name] for name in required if name in head_counts
               and _invalid_head_count(head_counts[name])}
    non_divisible = {name: head_counts[name] for name in required if name in head_counts
                     and name not in invalid and head_counts[name] % encoder_cp_size}
    if missing or invalid or non_divisible:
        matched = sorted(name for name, _ in model.named_modules()
                         if fnmatch.fnmatchcase(name, attention_pattern))
        raise ValueError(
            "[HP-ENCODER-005] Encoder-CP/Ulysses head divisibility check failed; "
            "constraint=head_count % encoder_cp_size == 0 for every Q/K/V; "
            "config_path=accelerator.encoder_parallel.encoder_cp_size; "
            f"encoder_cp_size={encoder_cp_size}; head_counts={dict(head_counts)}; "

    Args:
        model: Final model after module replacements.
        spec: Model-owned attention matching and replacement contract.
    """
    expected = frozenset(name for name, _ in model.named_modules()
                         if fnmatch.fnmatchcase(name, spec.attention_pattern))
    recorded = set(getattr(model, "_hp_replaced_module_fqns", ()) or ())
    recorded_visual = {name for name in recorded if fnmatch.fnmatchcase(name, spec.attention_pattern)}
    missing = sorted(expected - recorded_visual)
    unexpected = sorted(recorded_visual - expected)
    wrong_type = sorted(name for name in expected
                        if not isinstance(model.get_submodule(name), spec.replacement_attention_type))
    if not expected or missing or unexpected or wrong_type:
        raise RuntimeError(
            "[HP-ENCODER-006] incomplete vision attention replacement: "
            f"expected={len(expected)}; missing={missing}; unexpected={unexpected}; wrong_type={wrong_type}"
        )
    return expected


def validate_encoder_execution_options(options: Mapping[str, Any], bridge: Any) -> None:
    """Reject unsupported execution modes before creating Encoder groups.

    Args:
        options: Trainer microbatch, compile, and placement settings.
        bridge: Model-bound validator for vision-specific restrictions.
    """
    bridge.validate_training(options)
    if options.get("compile_enabled"):
        raise NotImplementedError("[HP-ENCODER-011] v1 requires compile.enabled=false")
    if options.get("validate_placement"):
        raise NotImplementedError("[HP-ENCODER-012] v1 requires model.validate_placement=false")


@dataclass
class EncoderParallelPlan:

    k_fsdp_max: int

    def _gather_features(self, hidden: torch.Tensor) -> tuple[torch.Tensor, ...]:
        """All-gather each local segment and trim dummy padding."""
        gathered = []
        offset = 0
        for segment in self.local_segments:
            local = hidden[offset:offset + segment.local_padded_length]
            offset += segment.local_padded_length
            full = differentiable_variable_all_gather(
                local, output_splits=[segment.local_padded_length] * self.runtime.config.encoder_cp_size,
                group=self.runtime.encoder_cp_mesh.get_group(),
            ) if self.runtime.config.encoder_cp_size > 1 else local
            gathered.append(full if segment.is_dummy else full[:segment.item.sequence_length])
        return tuple(gathered)

    def _execute_local(self, image_inputs: Mapping[str, Any]
                       ) -> tuple[tuple[EncoderFeature, ...], torch.Tensor, torch.Tensor]:
        """Run the local ViT pass and build routed feature records."""
        bridge = self.runtime.bridge
        context = EncoderAttentionContext(
            self.local_segments,
            self.runtime.encoder_cp_mesh,
            bridge.head_counts,
            self.runtime.expected,

            self.runtime.encoder_cp_mesh,
            bridge.head_counts,
            self.runtime.expected,
        )
        with activate_encoder_attention(context):
            hidden = bridge.encode_local(image_inputs, self.local_segments)
        context.assert_exactly_once()
        full_features = self._gather_features(hidden)
        records = []
        anchors = [feature.sum() * 0.0 for feature in full_features]
        prototype = None
        for slot in range(self.k_fsdp_max):
            segment = self.local_segments[slot] if slot < len(self.local_segments) else None
            real = segment is not None and not segment.is_dummy and self.runtime.encoder_cp_rank == 0
            if real:
                feature = bridge.project(full_features[slot], segment.item)
                if feature.ndim != 2 or feature.shape[0] != segment.item.output_rows:
                    raise ValueError(
                        "[HP-ENCODER-017] aligner output does not match planned rows; "
                        f"item_id={segment.item.item_id}; actual_shape={tuple(feature.shape)}; "
                        f"expected_rows={segment.item.output_rows}"
                    )
                records.append((segment.item, feature))
            else:
                feature = bridge.project(hidden, None)
                anchors.append(feature.sum() * 0.0)
            prototype = feature
        if prototype is None:
            raise RuntimeError("[HP-ENCODER-009] active plan must have an aligner slot")
        anchor = torch.stack(anchors).sum() if anchors else prototype.sum() * 0.0
        return tuple(records), prototype, anchor

    def encode_and_route(self, image_inputs: Mapping[str, Any]
                         ) -> tuple[list[EncoderFeature], torch.Tensor | None]:
        """Run packed ViT and route aligned features to this decoder owner.