diff --git a/src/interscale/config/global_component_config.py b/src/interscale/config/global_component_config.py
index e0841d8..2fcecf7 100644
--- a/src/interscale/config/global_component_config.py
+++ b/src/interscale/config/global_component_config.py
@@ -17,8 +17,11 @@ def get_global_component_cfg(cfg, global_component_name):
             2000  # optionally adjust to maximum number of cells, ideally shouldnt be larger than 4000
         )
         cfg.model.global_component.parameters.long_range_attention = (
-            False  # if True, takes inverse of adjacency matrix as long-range attention mask
+            False  # if True, blocks attention inside the local component's receptive field
         )
+        # Radius of that blocked neighbourhood, in message-passing steps. 0 means "match the local
+        # component's num_layers", which is the setting that keeps the two components disjoint.
+        cfg.model.global_component.parameters.long_range_mask_hops = 0
         cfg.model.global_component.parameters.type_gex_embedding = None
         cfg.model.global_component.latent_obsm_key = None  # Use the obms key where precomputed embeddings are stored, only if type_gex_embedding is "Precomputed"
     return cfg
diff --git a/src/interscale/model/base/_base_model.py b/src/interscale/model/base/_base_model.py
index 5552df2..3752eab 100644
--- a/src/interscale/model/base/_base_model.py
+++ b/src/interscale/model/base/_base_model.py
@@ -477,6 +477,7 @@ def _register_global_component(self) -> GlobalModule:
         Instance must be defined in InterScale.module.global_components.
         """
         if self._cfg.model.global_component.name == "self-attn-transformer":
+            local_mask_hops = self._resolve_local_mask_hops()
             self._model_summary_string = self._model_summary_string + (
                 f"Global component {self._cfg.model.global_component.name}: "
                 f"max_seq_len: {self._cfg.model.global_component.parameters.max_seq_len},"
@@ -485,7 +486,8 @@ def _register_global_component(self) -> GlobalModule:
                 f"act_func: {self._cfg.model.global_component.parameters.activation_func},"
                 f"num_layers: {self._cfg.model.global_component.parameters.num_layers},"
                 f"dim_feedforward: {self._cfg.model.global_component.parameters.dim_feedforward},"
-                f"enforce long-range attention: {self._cfg.model.global_component.parameters.long_range_attention}"
+                f"enforce long-range attention: {self._cfg.model.global_component.parameters.long_range_attention},"
+                f"local_mask_hops: {local_mask_hops}"
             )
             return TransformerNodeEncoderHook(
                 n_input=self.n_input,
@@ -502,10 +504,25 @@ def _register_global_component(self) -> GlobalModule:
                 num_layers=self._cfg.model.global_component.parameters.num_layers,
                 dim_feedforward=self._cfg.model.global_component.parameters.dim_feedforward,
                 long_range_attention=self._cfg.model.global_component.parameters.long_range_attention,
+                local_mask_hops=local_mask_hops,
             )
         else:
             raise ValueError(f"Global component {self._cfg.model.global_component.name} not found.")
 
+    def _resolve_local_mask_hops(self) -> int:
+        """How far the long-range mask reaches, in message-passing steps.
+
+        ``long_range_mask_hops = 0`` means "ask the local component", which is the setting that
+        actually keeps the two components disjoint: the mask has to cover the GNN's receptive
+        field, and that is its number of layers. A transformer-only model has no local component
+        to ask, so it falls back to a single hop.
+        """
+        configured = self._cfg.model.global_component.parameters.long_range_mask_hops
+        if configured:
+            return int(configured)
+        params = self._cfg.model.local_component.get("parameters", None)
+        return int(params.num_layers) if params is not None and "num_layers" in params else 1
+
     def predict_nodewise(
         self,
         adata: AnnData | None = None,
diff --git a/src/interscale/module/base/_base_global_module.py b/src/interscale/module/base/_base_global_module.py
index 7045331..7bd46b7 100644
--- a/src/interscale/module/base/_base_global_module.py
+++ b/src/interscale/module/base/_base_global_module.py
@@ -369,6 +369,7 @@ def from_config(cfg, **kwargs):
                 num_layers=params["num_layers"],
                 dim_feedforward=params["dim_feedforward"],
                 long_range_attention=params["long_range_attention"],
+                local_mask_hops=params.get("local_mask_hops", 1),
                 **kwargs,
             )
         # Add more elifs for other modules
diff --git a/src/interscale/module/global_modules/transformer_encoder.py b/src/interscale/module/global_modules/transformer_encoder.py
index e8cc577..cd4670c 100644
--- a/src/interscale/module/global_modules/transformer_encoder.py
+++ b/src/interscale/module/global_modules/transformer_encoder.py
@@ -25,6 +25,7 @@ def __init__(
         dim_feedforward: int = 2048,
         dropout_global: float = 0.1,
         long_range_attention: bool = True,
+        local_mask_hops: int = 1,
         **base_module_kwargs,
     ):
 
@@ -38,6 +39,11 @@ def __init__(
         self.dim_feedforward = dim_feedforward
         self.dropout_global = dropout_global
         self.long_range_attention = long_range_attention
+        # Radius of the neighbourhood the transformer is blocked from, in message-passing steps.
+        # It has to match the local component's depth: blocking 1 hop while the GCN mixes 2 leaves
+        # the second hop reachable by both components, which is the duplication the mask exists
+        # to prevent.
+        self.local_mask_hops = local_mask_hops
 
         # Create Transformer Encoder
         encoder_layer = CustomTransformerEncoderLayer(
@@ -98,19 +104,24 @@ def common_step_local_to_global(self, batched_data, emb: torch.Tensor, eval_step
         )
 
         if self.long_range_attention:
-            # INSERT_YOUR_CODE
-            raise NotImplementedError("Long-range attention mask feature is currently not implemented.")
+            # Block everything the local component has already seen, so the transformer can only
+            # contribute what the GNN could not.
             attention_mask = create_transformer_attention_mask_from_edges(
-                batched_data.edge_index, len(batched_data.obs_names), batched_data.batch, index_nodes, self.n_heads
+                batched_data.edge_index,
+                batched_data.batch.numel(),
+                batched_data.batch,
+                index_nodes,
+                self.n_heads,
+                n_hops=self.local_mask_hops,
+                device=emb.device,
             )
-            # Convert attention_mask to same dtype as src_padding_mask
-            attention_mask = attention_mask.to(dtype=src_padding_mask.dtype)
         else:
-            # attention_mask = None
-            # default: mask diagonal with -inf; no attention to self
+            # default: no attention to self, everything else open
             attention_mask = attn_mask_diagonal(batched_data.batch, index_nodes, self.n_heads, emb.device)
 
-        attention_mask = attention_mask.to(dtype=src_padding_mask.dtype)
+        # Boolean throughout: True is blocked. Keeping it out of float space is what stops the
+        # 0 * -inf that an "inverse adjacency" built by arithmetic would produce.
+        attention_mask = attention_mask.to(dtype=torch.bool)
 
         return padded_emb, src_padding_mask, index_nodes, attention_mask
 
@@ -223,5 +234,7 @@ def get_model_summary(self) -> str:
             f"n_heads: {self.n_heads}, \n"
             f"act_func: {self.act_func}, \n"
             f"num_layers: {self.num_layers}, \n"
+            f"long_range_attention: {self.long_range_attention}, \n"
+            f"local_mask_hops: {self.local_mask_hops}, \n"
         )
         return summary
diff --git a/src/interscale/tl/masking.py b/src/interscale/tl/masking.py
index 711496d..ef46266 100644
--- a/src/interscale/tl/masking.py
+++ b/src/interscale/tl/masking.py
@@ -263,82 +263,149 @@ def masked_loss(loss_fn, loss_type: str, y_pred: torch.Tensor, y_true: torch.Ten
     return loss_fn(y_pred[entry_mask], y_true[entry_mask])
 
 
+def _local_reach(edge_index: torch.Tensor, num_nodes: int, n_hops: int) -> torch.Tensor:
+    """Boolean ``[num_nodes, num_nodes]`` matrix: which nodes the GNN can already see from each node.
+
+    ``reach[i, j]`` is True when ``j`` lies within ``n_hops`` of ``i``, the diagonal included --
+    a cell is part of its own receptive field. ``n_hops`` should be the number of message-passing
+    layers of the local component, since that is exactly its receptive field.
+
+    The multi-hop closure is done with sparse matmuls, so the cost tracks the number of edges
+    rather than ``num_nodes ** 2``; only the final densification is quadratic, and that is
+    unavoidable because the attention mask itself is dense.
+    """
+    device = edge_index.device
+
+    if n_hops <= 1:
+        # No closure to compute, so skip the sparse round trip: scattering straight into the
+        # boolean matrix avoids materialising an n x n float one just to threshold it.
+        dense = torch.zeros((num_nodes, num_nodes), dtype=torch.bool, device=device)
+        dense[edge_index[0], edge_index[1]] = True
+    else:
+        values = torch.ones(edge_index.shape[1], device=device)
+        adj = torch.sparse_coo_tensor(edge_index, values, (num_nodes, num_nodes)).coalesce()
+
+        reach, frontier = adj, adj
+        for _ in range(n_hops - 1):
+            frontier = torch.sparse.mm(frontier, adj).coalesce()
+            reach = (reach + frontier).coalesce()
+        dense = reach.to_dense() > 0
+    # Spatial neighbour graphs are built symmetric, but a directed edge_index would otherwise
+    # leave the mask asymmetric and block only one direction of a pair the GNN mixed both ways.
+    # `dense |= dense.T` aliases its own memory, so this is an out-of-place or.
+    dense = dense | dense.transpose(0, 1)
+    dense.fill_diagonal_(True)
+    return dense
+
+
 def create_transformer_attention_mask_from_edges(
-    edge_index: torch.Tensor, num_nodes: int, batch: torch.Tensor, index_nodes: list, num_heads: int
+    edge_index: torch.Tensor,
+    num_nodes: int,
+    batch: torch.Tensor,
+    index_nodes: list,
+    num_heads: int,
+    *,
+    n_hops: int = 1,
+    device: torch.device | None = None,
 ) -> torch.Tensor:
-    """
-    Creates an attention mask that is inverse to the edge indices. Unmasked = 0 and masked = -inf
-    If two nodes are connected in the adjacency matrix (edge_index = 1) then we have no attention (0) and vice versa.
+    """Block the transformer from attending inside the local component's receptive field.
 
-    Args:
-        edge_index (torch.Tensor): Edge index tensor of shape [2, num_edges]
-        num_nodes (int): Number of nodes in the graph
-        batch (torch.Tensor): Batch tensor of shape [num_nodes]
-        index_nodes (list): List of indices of nodes to keep [B, S] (range: 0, num_nodes)
-        num_heads (int): Number of attention heads
-    Returns:
-        torch.Tensor: Attention mask of shape [num_batch*num_heads, max_seq_len, max_seq_len] with 1s for no attention (True -> mask attention) and 0s for attention (False -> no mask)
+    This is the ``M = 1 - A`` mask of the paper, with ``A`` taken as the ``n_hops`` closure of the
+    spatial neighbour graph rather than just its direct edges: the point is to stop the two
+    components from re-deriving the same signal, and a 2-layer GCN has already mixed the 2-hop
+    neighbourhood. The diagonal is blocked with it -- a cell's own state is what the local
+    embedding is.
+
+    **Why this cannot produce NaN.** A softmax row that is entirely ``-inf`` is NaN, which is the
+    failure this mask invites: a cell in a dense region can easily have every other cell of its
+    window inside its own neighbourhood. Two properties rule it out here.
+
+    * The CLS token is never blocked, as a query or as a key. Every row therefore keeps at least
+      one attendable key, whatever the graph looks like. It is also never a padded key, so the
+      merge with ``src_key_padding_mask`` cannot take that guarantee away.
+    * The mask is boolean and is built by indexing, never by arithmetic. Forming it as
+      ``(1 - A) * -inf`` -- the obvious reading of "inverse adjacency" -- puts ``0 * -inf`` on
+      every connected pair, and that is NaN before the softmax ever runs.
+
+    Padding positions are left unblocked for the same reason: their rows are meaningless but must
+    still normalise, and ``src_key_padding_mask`` is what actually removes them as keys.
+
+    Parameters
+    ----------
+    edge_index
+        ``[2, num_edges]`` edge index of the whole batch, with PyG's per-graph node offsets.
+    num_nodes
+        Number of nodes in the batch; used to check ``batch`` and ``edge_index`` agree.
+    batch
+        ``[num_nodes]`` graph assignment per node.
+    index_nodes
+        Per graph, the indices of the nodes ``pad_batch`` kept, relative to that graph's own node
+        order. Its lengths define the sequence length.
+    num_heads
+        Number of attention heads; the mask is repeated for each.
+    n_hops
+        Radius of the blocked neighbourhood, in message-passing steps. Pass the local component's
+        ``num_layers``.
+    device
+        Device for the returned mask. Defaults to ``edge_index``'s.
+
+    Returns
+    -------
+    torch.Tensor
+        Boolean ``[num_batch * num_heads, S + 1, S + 1]`` mask, ``True`` where attention is
+        blocked, with the CLS token in the last position. Ordered graph-major, matching what
+        :class:`torch.nn.MultiheadAttention` expects of a 3-D ``attn_mask``.
     """
-    INVALID_MASK_VALUE = -float("inf")
+    device = edge_index.device if device is None else device
+    batch = batch.to(torch.long)
+    if batch.numel() != num_nodes:
+        raise ValueError(f"batch has {batch.numel()} entries but num_nodes is {num_nodes}")
 
-    num_batch = int(batch[-1].item() + 1)
+    num_batch = int(batch.max().item()) + 1
     max_seq_len = max(len(nodes) for nodes in index_nodes)
+    mask = torch.zeros((num_batch, max_seq_len + 1, max_seq_len + 1), dtype=torch.bool, device=device)
 
-    # Initialize with -inf (no attention allowed)
-    attention_mask = torch.full(
-        (num_batch * num_heads, max_seq_len + 1, max_seq_len + 1), INVALID_MASK_VALUE, device=edge_index.device
-    )
-    # Set the diagonal to -inf (no self-attention)
-    diag_idx = torch.arange(max_seq_len, device=edge_index.device)
-    attention_mask[:, diag_idx, diag_idx] = INVALID_MASK_VALUE
+    for b in range(num_batch):
+        nodes_b = torch.nonzero(batch == b, as_tuple=False).flatten()
+        n_b = int(nodes_b.numel())
+        if n_b == 0:
+            continue
+        # PyG batches graphs by concatenation, so a graph's nodes are contiguous and its local
+        # indices are the global ones minus the offset of its first node.
+        offset = int(nodes_b[0].item())
+        in_graph = (batch[edge_index[0]] == b) & (batch[edge_index[1]] == b)
+        local_edges = edge_index[:, in_graph] - offset
 
-    # Create full adjacency matrix + 1 for cls token (end of sequence)
-    adj_matrix = torch.zeros((num_nodes, num_nodes), device=edge_index.device)  # TODO: check if zero or ones
-    adj_matrix[edge_index[0], edge_index[1]] = INVALID_MASK_VALUE
+        reach = _local_reach(local_edges, n_b, n_hops)
 
-    # For each batch, extract the submatrix for kept nodes
-    for b in range(num_batch):
-        nodes = index_nodes[b]
-        seq_len = len(nodes)
-        assert seq_len + 1 <= max_seq_len + 1, f"Mismatch: seq_len+1: {seq_len + 1}, max_seq_len+1: {max_seq_len + 1}"
-        # Extract submatrix for the kept nodes
-        batch_mask = adj_matrix[nodes][:, nodes]  # Get submatrix for kept nodes
-        # INSERT_YOUR_CODE
-        assert torch.any(batch_mask != 0), "batch_mask contains only zero entries"
-        # Add row and column of ones for CLS token - full attention
-        batch_mask = torch.cat(
-            [batch_mask, torch.zeros(batch_mask.size(0), 1, device=batch_mask.device)], dim=1
-        )  # Add column
-        batch_mask = torch.cat(
-            [batch_mask, torch.zeros(1, batch_mask.size(1), device=batch_mask.device)], dim=0
-        )  # Add row
-        assert batch_mask.shape == (seq_len + 1, seq_len + 1), (
-            f"Mismatch: batch_mask.shape: {batch_mask.shape}, (seq_len+1, seq_len+1): {(seq_len + 1, seq_len + 1)}"
-        )
-        assert attention_mask.shape[-2:] == (max_seq_len + 1, max_seq_len + 1), (
-            f"Mismatch: attention_mask.shape[-2:]: {attention_mask.shape[-2:]}, (seq_len+1, seq_len+1): {(seq_len + 1, seq_len + 1)}"
-        )
-        # append inverse adjacency matrix to the end of the attention mask
-        attention_mask[b * num_heads : b * num_heads + num_heads, -(seq_len + 1) :, -(seq_len + 1) :] = batch_mask
-        # add zeros for nodes that are not in the batch
-        attention_mask[b * num_heads : b * num_heads + num_heads, :seq_len, :seq_len] = float("0")
+        kept = torch.as_tensor(index_nodes[b], dtype=torch.long, device=reach.device)
+        block = reach[kept][:, kept].to(device)
 
-    assert not torch.any(torch.isnan(attention_mask)), "attention_mask contains NaN values"
-    print("attention_mask", attention_mask.shape, attention_mask)
-    return attention_mask
+        # pad_batch left-pads, so the kept tokens sit in the LAST len(kept) positions before the
+        # CLS slot. Writing the block anywhere else silently masks the wrong pairs.
+        s = int(kept.numel())
+        lo = max_seq_len - s
+        mask[b, lo:max_seq_len, lo:max_seq_len] = block
+
+    if bool(mask.all(dim=-1).any()):
+        raise RuntimeError("a query row is fully blocked; softmax would be NaN")
+
+    # (N * num_heads, L, S) is indexed graph-major: repeat_interleave, not repeat.
+    return mask.repeat_interleave(num_heads, dim=0)
 
 
 def attn_mask_diagonal(batch: torch.Tensor, index_nodes: list, num_heads: int, device: torch.device) -> torch.Tensor:
-    """
-    Sets the diagonal of the attention mask to -inf.
+    """Block self-attention only: the weakest mask, and the default.
+
+    Returns the same boolean convention as
+    :func:`create_transformer_attention_mask_from_edges` (``True`` = blocked), so the two are
+    interchangeable at the call site. The CLS slot in the last position is left open.
     """
     max_seq_len = max(len(nodes) for nodes in index_nodes)
-    batch_size = int(batch[-1].item() + 1)
+    batch_size = int(batch.max().item()) + 1
     attention_mask = torch.zeros(
-        (num_heads * batch_size, max_seq_len + 1, max_seq_len + 1), device=device, dtype=torch.float32
+        (num_heads * batch_size, max_seq_len + 1, max_seq_len + 1), device=device, dtype=torch.bool
     )
-    # Set the diagonal to -inf (no self-attention)
     diag_idx = torch.arange(max_seq_len, device=device)
-    attention_mask[:, diag_idx, diag_idx] = float("-inf")
-    # Convert attention_mask to same dtype as src_padding_mask
+    attention_mask[:, diag_idx, diag_idx] = True
     return attention_mask
diff --git a/tests/test_long_range_mask.py b/tests/test_long_range_mask.py
new file mode 100644
index 0000000..567207a
--- /dev/null
+++ b/tests/test_long_range_mask.py
@@ -0,0 +1,374 @@
+"""The long-range attention mask must block the GNN's receptive field without producing NaN.
+
+`create_transformer_attention_mask_from_edges` implements `M = 1 - A`: the transformer may not
+attend inside the neighbourhood the local component has already mixed. The failure this invites is
+a softmax row that is entirely `-inf`, which is NaN -- a cell in a dense region can have every
+other cell of its window inside its own neighbourhood. The guarantee that rules it out is that the
+CLS token is never blocked as a key, so every row keeps at least one attendable position.
+
+The other thing these tests pin down is index alignment. `pad_batch` *left*-pads, so a graph's
+tokens occupy the last positions of the sequence; writing the adjacency block anywhere else masks
+the wrong pairs silently, with no error and a model that still trains.
+"""
+
+import numpy as np
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from interscale.module.global_modules import TransformerNodeEncoderHook
+from interscale.tl.masking import attn_mask_diagonal, create_transformer_attention_mask_from_edges
+
+N_HEADS = 2
+
+
+def path_graph(n):
+    """0-1-2-...-(n-1), as an undirected edge index."""
+    src = list(range(n - 1)) + list(range(1, n))
+    dst = list(range(1, n)) + list(range(n - 1))
+    return torch.tensor([src, dst], dtype=torch.long)
+
+
+def test_one_hop_blocks_neighbours_and_self():
+    n = 5
+    mask = create_transformer_attention_mask_from_edges(
+        path_graph(n), n, torch.zeros(n, dtype=torch.long), [list(range(n))], num_heads=1, n_hops=1
+    )
+    blocked = mask[0, :n, :n]
+
+    expected = torch.zeros(n, n, dtype=torch.bool)
+    for i in range(n):
+        expected[i, i] = True
+        if i > 0:
+            expected[i, i - 1] = True
+        if i < n - 1:
+            expected[i, i + 1] = True
+
+    assert torch.equal(blocked, expected)
+
+
+def test_two_hops_blocks_the_gcn_receptive_field():
+    """A 2-layer GCN mixes the 2-hop neighbourhood, so blocking only 1 hop leaves the second hop
+    reachable by both components -- the duplication the mask exists to prevent."""
+    n = 7
+    mask = create_transformer_attention_mask_from_edges(
+        path_graph(n), n, torch.zeros(n, dtype=torch.long), [list(range(n))], num_heads=1, n_hops=2
+    )
+    blocked = mask[0, :n, :n]
+
+    for i in range(n):
+        for j in range(n):
+            assert bool(blocked[i, j]) == (abs(i - j) <= 2), f"({i},{j}) at 2 hops"
+
+
+def test_cls_token_is_never_blocked():
+    """The NaN guarantee: whatever the graph, the CLS column stays open, so no row is empty."""
+    n = 6
+    dense = torch.tensor([[i, j] for i in range(n) for j in range(n) if i != j], dtype=torch.long).T
+    mask = create_transformer_attention_mask_from_edges(
+        dense, n, torch.zeros(n, dtype=torch.long), [list(range(n))], num_heads=1, n_hops=3
+    )
+
+    assert bool(mask[0, :n, :n].all()), "every real pair should be blocked in a complete graph"
+    assert not bool(mask[0, -1, :].any()), "CLS row blocked"
+    assert not bool(mask[0, :, -1].any()), "CLS column blocked"
+    assert not bool(mask.all(dim=-1).any()), "a fully blocked row would be NaN after softmax"
+
+
+def test_block_is_written_to_the_left_padded_positions():
+    """Two graphs of different size: the smaller one's block must land at the END of its sequence,
+    because that is where `pad_batch` puts its tokens."""
+    small, large = 3, 5
+    edges_small = path_graph(small)
+    edges_large = path_graph(large) + small  # PyG offsets the second graph's node ids
+    edge_index = torch.cat([edges_small, edges_large], dim=1)
+    batch = torch.tensor([0] * small + [1] * large, dtype=torch.long)
+
+    mask = create_transformer_attention_mask_from_edges(
+        edge_index, small + large, batch, [list(range(small)), list(range(large))], num_heads=1, n_hops=1
+    )
+
+    pad = large - small
+    assert not bool(mask[0, :pad, :].any()), "padded rows must stay open"
+    assert not bool(mask[0, :, :pad].any()), "padded columns must stay open"
+    # The 3-node path, placed in the last three real positions.
+    assert torch.equal(
+        mask[0, pad:large, pad:large],
+        torch.tensor([[True, True, False], [True, True, True], [False, True, True]]),
+    )
+
+
+def test_heads_are_repeated_graph_major():
+    """MultiheadAttention indexes a 3-D attn_mask as batch * num_heads + head, so the graphs must
+    be interleaved by head and not tiled."""
+    small, large = 2, 4
+    edge_index = torch.cat([path_graph(small), path_graph(large) + small], dim=1)
+    batch = torch.tensor([0] * small + [1] * large, dtype=torch.long)
+
+    mask = create_transformer_attention_mask_from_edges(
+        edge_index, small + large, batch, [list(range(small)), list(range(large))], num_heads=N_HEADS, n_hops=1
+    )
+
+    assert mask.shape == (2 * N_HEADS, large + 1, large + 1)
+    assert torch.equal(mask[0], mask[1]), "both heads of graph 0"
+    assert torch.equal(mask[2], mask[3]), "both heads of graph 1"
+    assert not torch.equal(mask[0], mask[2]), "the two graphs differ"
+
+
+def test_only_the_kept_nodes_are_used():
+    """When a graph is longer than max_seq_len, pad_batch keeps a subset; the mask has to be the
+    submatrix over exactly those nodes, in their order."""
+    n = 6
+    kept = [0, 2, 4]
+    mask = create_transformer_attention_mask_from_edges(
+        path_graph(n), n, torch.zeros(n, dtype=torch.long), [kept], num_heads=1, n_hops=1
+    )
+    blocked = mask[0, : len(kept), : len(kept)]
+
+    # 0-2-4 are pairwise 2 apart on the path, so at 1 hop only the diagonal is blocked.
+    assert torch.equal(blocked, torch.eye(len(kept), dtype=torch.bool))
+
+
+def test_diagonal_mask_uses_the_same_boolean_convention():
+    n = 4
+    diag = attn_mask_diagonal(
+        torch.zeros(n, dtype=torch.long), [list(range(n))], num_heads=1, device=torch.device("cpu")
+    )
+
+    assert diag.dtype == torch.bool
+    assert torch.equal(diag[0, :n, :n], torch.eye(n, dtype=torch.bool))
+    assert not bool(diag[0, -1, :].any()) and not bool(diag[0, :, -1].any())
+
+
+class _Batch:
+    """Minimal stand-in for the PyG batch the global module reads."""
+
+    def __init__(self, edge_index, batch, n_nodes):
+        self.edge_index = edge_index
+        self.batch = batch
+        self.obs_names = torch.arange(n_nodes)
+        self.num_nodes = n_nodes
+        self.mask = torch.zeros(n_nodes, dtype=torch.bool)
+
+
+def build_module(long_range, hops, max_seq_len=32):
+    return TransformerNodeEncoderHook(
+        max_seq_len=max_seq_len,
+        n_heads=N_HEADS,
+        dropout_global=0.0,
+        act_func="relu",
+        num_layers=1,
+        dim_feedforward=16,
+        long_range_attention=long_range,
+        local_mask_hops=hops,
+        n_input=8,
+        n_output=8,
+        n_embed=8,
+        decoder_type="linear",
+        dropout_decoder=0.0,
+        decoder_hidden_dims=[16],
+        mask_percentage=0.1,
+        mask_strategy="node",
+    )
+
+
+@pytest.mark.parametrize("hops", [1, 2, 3])
+def test_forward_is_finite_on_a_dense_graph(hops):
+    """The regression test for the NaN reports: a graph dense enough that many cells have their
+    whole window inside their own neighbourhood still has to produce finite attention."""
+    n = 12
+    rng = np.random.default_rng(0)
+    edge_index = torch.tensor([[i, j] for i in range(n) for j in range(n) if i != j], dtype=torch.long).T
+    batch = _Batch(edge_index, torch.zeros(n, dtype=torch.long), n)
+    emb = torch.tensor(rng.normal(size=(n, 8)), dtype=torch.float32)
+
+    module = build_module(long_range=True, hops=hops).eval()
+    padded, padding_mask, _, attn_mask = module.common_step_local_to_global(batch, emb, eval_step=True)
+    out, _, attn = module.forward(padded, padding_mask, attn_mask, register_hook=True)
+
+    assert torch.isfinite(out).all(), "transformer output contains NaN or inf"
+    assert attn is not None and torch.isfinite(attn).all(), "attention weights contain NaN or inf"
+
+    # A masked softmax can be finite forward and still produce NaN gradients, which is how this
+    # fails silently in training rather than at the first batch.
+    out.sum().backward()
+    for name, param in module.named_parameters():
+        if param.grad is not None:
+            assert torch.isfinite(param.grad).all(), f"non-finite gradient in {name}"
+
+
+def test_training_steps_stay_finite():
+    """Several optimiser steps under the mask, since a NaN that only appears once weights have
+    moved would not be caught by a single forward."""
+    n = 10
+    rng = np.random.default_rng(2)
+    edge_index = torch.tensor([[i, j] for i in range(n) for j in range(n) if i != j], dtype=torch.long).T
+    batch = _Batch(edge_index, torch.zeros(n, dtype=torch.long), n)
+    emb = torch.tensor(rng.normal(size=(n, 8)), dtype=torch.float32)
+    target = torch.tensor(rng.normal(size=(n + 1, 1, 8)), dtype=torch.float32)
+
+    module = build_module(long_range=True, hops=2)
+    optimizer = torch.optim.Adam(module.parameters(), lr=1e-2)
+
+    for _ in range(5):
+        optimizer.zero_grad()
+        padded, padding_mask, _, attn_mask = module.common_step_local_to_global(batch, emb, eval_step=True)
+        out, _, _ = module.forward(padded, padding_mask, attn_mask, register_hook=False)
+        loss = torch.nn.functional.mse_loss(out, target)
+        assert torch.isfinite(loss), "loss went non-finite under the mask"
+        loss.backward()
+        optimizer.step()
+
+
+def test_blocked_pairs_receive_no_attention():
+    """The mask has to actually reach the softmax, not just be built."""
+    n = 8
+    rng = np.random.default_rng(1)
+    edge_index = path_graph(n)
+    batch = _Batch(edge_index, torch.zeros(n, dtype=torch.long), n)
+    emb = torch.tensor(rng.normal(size=(n, 8)), dtype=torch.float32)
+
+    module = build_module(long_range=True, hops=1).eval()
+    padded, padding_mask, _, attn_mask = module.common_step_local_to_global(batch, emb, eval_step=True)
+    module.forward(padded, padding_mask, attn_mask, register_hook=True)
+
+    weights = module.transformer_encoder.layers[0].get_attn_output_weights()  # [B, H, L, S]
+    weights = weights.reshape(-1, weights.shape[-2], weights.shape[-1])[0]
+    blocked = attn_mask[0]
+
+    assert torch.allclose(weights[blocked], torch.zeros(int(blocked.sum())), atol=1e-6)
+    assert weights[~blocked].sum() > 0
+    assert torch.allclose(weights.sum(dim=-1), torch.ones(weights.shape[0]), atol=1e-5)
+
+
+def test_mask_off_leaves_everything_but_the_diagonal_open():
+    n = 6
+    batch = _Batch(path_graph(n), torch.zeros(n, dtype=torch.long), n)
+    emb = torch.zeros(n, 8)
+
+    module = build_module(long_range=False, hops=2).eval()
+    _, _, _, attn_mask = module.common_step_local_to_global(batch, emb, eval_step=True)
+
+    assert torch.equal(attn_mask[0, :n, :n], torch.eye(n, dtype=torch.bool))
+
+
+def random_geometric_graph(n, radius, seed=0):
+    """Irregular degrees, unlike the path graph: closer to a real spatial neighbour graph."""
+    rng = np.random.default_rng(seed)
+    pos = rng.uniform(0, 1, size=(n, 2))
+    src, dst = [], []
+    for i in range(n):
+        for j in range(n):
+            if i != j and np.linalg.norm(pos[i] - pos[j]) <= radius:
+                src.append(i)
+                dst.append(j)
+    return torch.tensor([src, dst], dtype=torch.long)
+
+
+@pytest.mark.parametrize("n_layers", [1, 2, 3])
+def test_mask_covers_exactly_the_gcn_receptive_field(n_layers):
+    """The mask must block every cell that actually reached the local embedding, and no more.
+
+    Rather than asserting the hop arithmetic a second time, this measures the GCN's real
+    receptive field by autograd: cell j influenced cell i's local embedding exactly when
+    d h_local[i] / d x[j] is non-zero. That catches anything the hop count would miss -- an
+    off-by-one in the layer stack, self-loops, the input_proj residual.
+
+    The probe is a random projection of h[i], not h[i].sum(): the GCN ends in a LayerNorm with
+    elementwise_affine=False, so every row sums to exactly zero and the gradient of the sum is
+    identically zero regardless of the graph.
+    """
+    from interscale.module.local_modules.GCN import GCN
+
+    n, n_features, n_embed = 12, 6, 4
+    edge_index = random_geometric_graph(n, radius=0.35)
+
+    torch.manual_seed(0)
+    gcn = GCN(
+        n_layers=n_layers,
+        hidden_dim=8,
+        dropout_local=0.0,
+        n_input=n_features,
+        n_output=n_features,
+        n_embed=n_embed,
+        decoder_type=None,
+        dropout_decoder=0.0,
+        mask_percentage=0.1,
+        mask_strategy="node",
+    ).eval()
+    probe = torch.randn(n_embed)
+
+    influenced = torch.zeros(n, n, dtype=torch.bool)
+    for i in range(n):
+        x = torch.randn(n, n_features, requires_grad=True)
+        (gcn(x, edge_index)[i] * probe).sum().backward()
+        influenced[i] = x.grad.abs().sum(1) > 1e-10
+
+    # Check the instrument before trusting its reading. A cell always influences its own
+    # embedding -- through input_proj if through nothing else -- so an empty diagonal means the
+    # gradient is dead and every comparison below is vacuous rather than informative. This is
+    # exactly what a probe of h[i].sum() produces, and without this guard it surfaces as
+    # "mask and receptive field disagree", pointing at the mask instead of at the measurement.
+    assert bool(influenced.diagonal().all()), (
+        "probe registered no self-influence: the gradient is dead, so this test measures nothing"
+    )
+
+    mask = create_transformer_attention_mask_from_edges(
+        edge_index, n, torch.zeros(n, dtype=torch.long), [list(range(n))], num_heads=1, n_hops=n_layers
+    )
+    blocked = mask[0, :n, :n]
+
+    leaked = influenced & ~blocked
+    assert not bool(leaked.any()), (
+        f"{int(leaked.sum())} pairs reached the local embedding but are left open to attention"
+    )
+    assert torch.equal(influenced, blocked), "mask and GCN receptive field disagree"
+
+
+def test_paths_through_dropped_nodes_are_still_blocked():
+    """Two kept cells that are only multi-hop neighbours *through* a cell pad_batch dropped were
+    still mixed by the GCN, which ran on the whole graph before any token was dropped. The reach
+    has to be closed on the full graph and only then restricted to the kept nodes."""
+    n = 5
+    kept = [0, 1, 3, 4]  # node 2 dropped; 1 and 3 are 2 hops apart only via node 2
+    mask = create_transformer_attention_mask_from_edges(
+        path_graph(n), n, torch.zeros(n, dtype=torch.long), [kept], num_heads=1, n_hops=2
+    )
+
+    pos_of = {node: t for t, node in enumerate(kept)}
+    assert bool(mask[0, pos_of[1], pos_of[3]]), "2-hop pair via a dropped node left open"
+    assert bool(mask[0, pos_of[3], pos_of[1]])
+    assert not bool(mask[0, pos_of[0], pos_of[4]]), "0 and 4 are 4 hops apart and must stay open"
+
+
+def test_default_hops_follow_the_local_component_depth():
+    """long_range_mask_hops = 0 means "ask the local component": the mask is only equivalent to
+    the GNN's field if it uses the same depth."""
+    from types import SimpleNamespace
+
+    from interscale.config import get_cfg_defaults
+    from interscale.config.global_component_config import get_global_component_cfg
+    from interscale.config.local_component_config import get_local_component_cfg
+    from interscale.model.base._base_model import BaseModel
+
+    cfg = get_cfg_defaults()
+    cfg.model.local_component.name = "GCN"
+    cfg.model.global_component.name = "self-attn-transformer"
+    cfg = get_local_component_cfg(cfg, "GCN")
+    cfg = get_global_component_cfg(cfg, "self-attn-transformer")
+
+    assert cfg.model.global_component.parameters.long_range_mask_hops == 0
+    resolved = BaseModel._resolve_local_mask_hops(SimpleNamespace(_cfg=cfg))
+    assert resolved == cfg.model.local_component.parameters.num_layers == 2
+
+    cfg.model.local_component.parameters.num_layers = 4
+    assert BaseModel._resolve_local_mask_hops(SimpleNamespace(_cfg=cfg)) == 4
+
+    cfg.model.global_component.parameters.long_range_mask_hops = 1
+    assert BaseModel._resolve_local_mask_hops(SimpleNamespace(_cfg=cfg)) == 1, "explicit setting must win"
+
+    # A transformer-only model has no local component to ask.
+    bare = get_cfg_defaults()
+    bare.model.global_component.name = "self-attn-transformer"
+    bare = get_global_component_cfg(bare, "self-attn-transformer")
+    assert BaseModel._resolve_local_mask_hops(SimpleNamespace(_cfg=bare)) == 1