Diff Coverage

        The runtime calls this before capture. Eager retains its real boundaries;
        graph inputs carry token-shaped metadata, not a variable document list.
        """
        if self.local_query_start != 0 or self.local_query_length != self.global_sequence_length:
            return
        self.segment_start_mask(device)
        object.__setattr__(self, "_use_prepared_packing_inputs", True)

    def router_metadata(self, device: torch.device) -> tuple[torch.Tensor, int]:

        """Expose packing tensors with a schema independent of document count and S."""
        if self._use_prepared_packing_inputs:
            return [self._prepared_starts, self._prepared_positions,
                    self._prepared_masks, self._prepared_sequences], ("packing",)
        return [self.cu_seq_lens], (self.local_query_start, self.local_query_length, self.global_sequence_length)

    @classmethod
    def pytree_unflatten(cls, children: Any, context: tuple) -> SharedCompressedPackedSequence:
        """Restore prepared metadata without reading boundary values during capture."""

    @classmethod
    def pytree_unflatten(cls, children: Any, context: tuple) -> SharedCompressedPackedSequence:
        """Restore prepared metadata without reading boundary values during capture."""
        if context != ("packing",):
            return cls(children[0], *context)
        starts, positions, masks, sequences = children
        reference = next(iter(starts.values()))
        length = reference.shape[-1]
        packed = cls(reference.new_empty(0), 0, length, length)

        Returns:
            Inputs used for both graph capture and replay. Eager never calls this hook.
        """
        return model_inputs

    def runtime_input_fields(self) -> tuple[str, ...]:
        """Return the forward fields exclusively owned by this adapter.

        # A replayed graph does not update Python module attributes.
        # Return detached router metrics explicitly for the standard callbacks.
        moe_monitor = getattr(self, "moe_monitor", None)
        if moe_monitor is not None:
            aux_loss = moe_monitor.last_mean_aux_loss
            if aux_loss is not None:
                loss_dict["router_aux_loss"] = aux_loss
        return loss, loss_dict

    def _init_callbacks(self):
        """Initialize callbacks."""