Diff Coverage

try:  # Register the upstream ``ao::`` custom ops on import.
    import torch._functorch._activation_offloading.offload_ops  # noqa: F401

    _AO_OFFLOAD = torch.ops.ao.offload.default
    _AO_RELOAD = torch.ops.ao.reload.default
    _AO_WAIT = torch.ops.ao.wait_tensor.default
    _AO_OPS = frozenset((_AO_OFFLOAD, _AO_RELOAD, _AO_WAIT))
    _AO_AVAILABLE = True
except (ImportError, AttributeError):
    _AO_OFFLOAD = _AO_RELOAD = _AO_WAIT = None
    _AO_OPS = frozenset()
    _AO_AVAILABLE = False

def _in_attention_scope(node: fx.Node, roots: frozenset) -> bool:
    """Whether a node was produced inside one of the attention roots."""
    if not roots:
        return False
    for fqn in _stack_fqns(node):
        if fqn.rsplit(".", 1)[-1] == "self_attn":
            return True
        if any(fqn == root or fqn.startswith(f"{root}.") for root in roots):

    try:
        if isinstance(value, torch.SymInt):
            value = value.node.hint
        return int(value)
    except Exception:  # pylint: disable=broad-exception-caught
        return None


def _tensor_bytes(value: torch.Tensor) -> Optional[int]:
    """Return a hinted byte count, or ``None`` for unsupported values."""

def _tensor_bytes(value: torch.Tensor) -> Optional[int]:
    """Return a hinted byte count, or ``None`` for unsupported values."""
    try:
        nbytes = value.nelement() * value.element_size()
    except Exception:  # pylint: disable=broad-exception-caught
        return None
    return _hint_to_int(nbytes)


def _is_contiguous(value: torch.Tensor) -> bool:

def _is_contiguous(value: torch.Tensor) -> bool:
    """Check contiguity without tripping over symbolic fake tensors."""
    try:
        return bool(value.is_contiguous())
    except Exception:  # pylint: disable=broad-exception-caught
        return False


def _storage_key(value: torch.Tensor) -> Optional[int]:
    """Return a stable identity for the tensor's storage, or ``None`` if unknown.

    callers skip the candidate.
    """
    try:
        storage = value.untyped_storage()
    except Exception:  # pylint: disable=broad-exception-caught
        return None
    cdata = getattr(storage, "_cdata", None)
    return int(cdata) if cdata else None

    unresolved = False
    position = 0
    for argument in schema.arguments:
        if argument.kwarg_only:
            index = None
        else:
            index = position
            position += 1
        alias_info = argument.alias_info

        alias_info = argument.alias_info
        if alias_info is None or not alias_info.is_write:
            continue
        if argument.name in node.kwargs:
            value = node.kwargs[argument.name]
        elif index is not None and index < len(node.args):
            value = node.args[index]
        else:
            unresolved = True

            unresolved = True
            continue
        written = _iter_nodes(value)
        if not written:
            continue
        resolved = True
        written_keys = _storage_keys(written)
        if not written_keys:
            unresolved = True
        keys.extend(written_keys)
    if unresolved or not resolved:
        return _tensor_argument_keys(node)
    return keys

    """Whether a storage group has no external, mutating, or output sharer."""
    if any(node.op != "call_function" for node in nodes):
        return False
    if any(_is_mutable(node) for node in nodes):
        return False
    if any(user.op == "output" for node in nodes for user in node.users):
        return False
    return True

    group owner is eligible, and ``unsafe`` groups (external or mutating
    aliases) are left on device.
    """
    if node.op != "call_function":
        return False
    if node.target in _AO_OPS:
        return False
    if is_view(node):
        return False
    if _is_collective_or_wait(node):
        return False
    value = node.meta.get("val")
    if not isinstance(value, torch.Tensor) or _is_dtensor(value):
        return False
    if value.dim() < 2:

        return False
    if value.dim() < 2:
        return False
    if not _is_contiguous(value):
        return False
    nbytes = _tensor_bytes(value)
    if nbytes is None or nbytes < _MIN_OFFLOAD_BYTES:
        return False
    # Offload one storage at a time. Only the group owner (the node that

    def collect(current: fx.Node) -> None:
        """Walk alias users and record backward readers."""
        for user in current.users:
            if user.op != "call_function":
                continue
            if _is_backward(user):
                backward_users.setdefault(current, []).append(user)
            elif _is_alias(user, current) and user not in visited:
                visited.add(user)

    def is_read_in_backward(current: fx.Node) -> bool:
        """Whether a view subtree has any backward reader (memoized)."""
        cached = memo.get(current)
        if cached is not None:
            return cached
        needed_here = bool(backward_users.get(current))
        for child in children.get(current, ()):
            if is_read_in_backward(child):
                needed_here = True

    @property
    def all_bwd_users(self) -> List[fx.Node]:
        """Every backward consumer, direct or reached through a view."""
        return self.direct_bwd_users + [user for _, user in self.view_bwd_redirects]


def _collect_offload_info(
    node: fx.Node, order: Dict[fx.Node, int]

                last_forward_use = user
                last_index = user_index

    if last_forward_use is node:
        last_use_arg: fx.Node = node
    elif isinstance(last_forward_use.meta.get("val"), torch.Tensor):
        last_use_arg = last_forward_use
    else:
        # The forward ``ao.wait_tensor`` frees the producer's storage, so it must

        if info is None:
            continue
        nbytes = _tensor_bytes(node.meta["val"])
        if nbytes is None:
            continue
        candidates.append((node, nbytes, info))
    return candidates

    offload.meta.update(src_meta)
    offload.meta["val"] = value.to(torch.device("cpu"))

    if info.last_forward_use is node:
        anchor = offload
    else:
        anchor = info.last_forward_use

    with graph.inserting_after(anchor):

    """Move one reload earlier by ``n_regions`` regions when possible."""
    root = _attention_root_of(reload_node, roots)
    positions = anchor_positions.get(root) if root is not None else None
    if not positions:
        return False
    reload_index = order[reload_node]
    # Latest occurrence of this root at or before the reload is its own
    # region when a module is reused across blocks; otherwise fall back to
    # the first occurrence (reload inserted before the region's first node).

        fake_mode = getattr(node.meta.get("val"), "fake_mode", None)
        shape_env = getattr(fake_mode, "shape_env", None)
        if shape_env is not None:
            return shape_env
    return None


class CpuOffloadPass(GraphPass):
    """Offload saved attention activations to host memory.

                no ``self_attn`` module, matching the eager attention-swap
                contract.
        """
        if pass_config.activation_swap == "none":
            return graph_module
        if not _AO_AVAILABLE:
            raise RuntimeError(
                "activation_swap='attention' requires torch's "
                "torch._functorch._activation_offloading.offload_ops (ao:: "

            budget_bytes = pass_config.offload_budget_gb * (1024**3)
            selected, total_bytes = _select_within_budget(candidates, budget_bytes)

            if not selected:
                _LOG.info("CPU offload: no attention activations selected")
                return graph_module

            wait_offloads: List[Tuple[_OffloadInfo, fx.Node]] = []
            for info in selected:
                wait_offload = _insert_forward_offload(graph_module, info)

    stream and an event registry outside the compiled region, and rewriting
    them as fused kernels would break the host/device synchronization.
    """
    if node.op != "call_function":
        return False
    if node.target in _side_effectful_functions:
        return True
    return str(node.target).startswith("ao.")