Diff Coverage

        each phase. Everything else in the forward is either recomputed
        unchanged or guarded by :func:`is_recomputing`, so a recompute reads the
        shared entries the forward pass published instead of writing its own.
        """
        object.__setattr__(
            self, "_selection_region", checkpoint_exclude_wrapper(self._select_tokens)
        )

    def _run_selection(

        if is_recomputing():
            # A recompute must not publish again: its consumers read the entries
            # the forward pass wrote, and a second write would break the
            # single-publisher contract (HP-STATE-002).
            return indexer_output
        if self.indexer.owns_key:
            shared_state.publish_index_key(self.layer_idx, indexer_output.index_key)
        shared_state.publish_topk_indices(self.layer_idx, indexer_output.topk_indices)
        if indexer_output.candidate_blocks is not None:

        if raw_kv_handle is not None:
            key_value = raw_kv_handle.wait()
        if compressed_handle is not None:
            compressed = compressed_handle.wait()
            if not is_recomputing():
                shared_state.publish_compressed_kv(self.layer_idx, compressed)
        if key_value is None:
            raise RuntimeError("raw KV all-gather did not produce a tensor")

        combined_key_value, sparse_indices, compressed_kv = self._build_sparse_indices(

        if any(
            child is not module and id(child) in root_ids
            for child in module.modules()
        ):
            raise ValueError(
                f"excluded recompute region {module_fqn!r} contains a recompute region root"
            )

            "requires RecomputePolicy.region_patterns"
        )
    selected_layer_fqns = _selected_checkpoint_layer_fqns(containers, selection)
    if not selected_layer_fqns:
        logger.info("Model-adapter recompute selected zero layers")
        return model
    if hasattr(model, "gradient_checkpointing_disable"):
        model.gradient_checkpointing_disable()
    checkpoint_kwargs = {"swap_inputs": swap_inputs} if not enable_compile else {}

    module_by_fqn = dict(model.named_modules())
    roots = _select_layer_region_roots(module_by_fqn, selected_layer_fqns, recompute_spec)
    if not roots:
        raise ValueError(
            "model adapter recompute patterns matched no region in the selected layers"
        )
    excluded_count = _wrap_declared_exclusions(module_by_fqn, roots, recompute_spec)
    wrapped_count = _wrap_region_roots(model, module_by_fqn, roots, checkpoint_kwargs)

        key=lambda item: item[0].count("."),
        reverse=True,
    ):
        if _is_checkpoint_wrapped(module):
            continue
        parent_fqn, child_name = module_fqn.rsplit(".", 1)
        parent = module_by_fqn[parent_fqn] if parent_fqn else model
        setattr(parent, child_name, checkpoint_wrapper(module, **checkpoint_kwargs))
        wrapped_count += 1

        parent = module_by_fqn[parent_fqn] if parent_fqn else model
        setattr(parent, child_name, checkpoint_wrapper(module, **checkpoint_kwargs))
        wrapped_count += 1
    if wrapped_count == 0:
        raise ValueError("model adapter recompute matched no region to wrap")
    return wrapped_count


def _resolve_activation_checkpoint_containers(