From f066d4026c250e71031865b1d28934c2006b717e Mon Sep 17 00:00:00 2001
From: mystri <hanboyou@huawei.com>
Date: Sat, 29 Aug 2026 16:27:13 +0800
Subject: [PATCH 1/8] Share equivalent RoPE cache storage

---
 torchtitan/distributed/pipeline_parallel.py |   8 +-
 torchtitan/models/common/__init__.py        |  10 +-
 torchtitan/models/common/decoder.py         |  51 ++-
 torchtitan/models/common/rope.py            | 341 +++++++++++++++++++-
 torchtitan/models/deepseek_v3/mtp.py        |  15 +-
 5 files changed, 409 insertions(+), 16 deletions(-)

diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py
index f547b18cd4..8e2b08c586 100644
--- a/torchtitan/distributed/pipeline_parallel.py
+++ b/torchtitan/distributed/pipeline_parallel.py
@@ -523,9 +523,15 @@ def _split_module(
                 elif isinstance(module_value, (nn.ModuleList, ModuleList)):
                     setattr(model, module_name, ModuleList())
         # Handle simple module attributes (e.g., "linear", "norm")
-        elif module_name not in modules_to_keep:
+        elif (
+            module_name not in modules_to_keep
+            and module_name != "_rope_cache_registry"
+        ):
             # Replace with None
             setattr(model, module_name, None)
+    prune_rope_caches = getattr(model, "_prune_rope_cache_registry", None)
+    if callable(prune_rope_caches):
+        prune_rope_caches()
     return model
 
 
diff --git a/torchtitan/models/common/__init__.py b/torchtitan/models/common/__init__.py
index 26d1dbf211..f9e9ac525c 100644
--- a/torchtitan/models/common/__init__.py
+++ b/torchtitan/models/common/__init__.py
@@ -36,7 +36,13 @@
     RMSNorm,
     SiLU,
 )
-from .rope import ComplexRoPE, CosSinRoPE, RoPE
+from .rope import (
+    ComplexRoPE,
+    CosSinRoPE,
+    RoPE,
+    RoPECacheReader,
+    register_rope_cache,
+)
 
 __all__ = [
     "Conv1d",
@@ -67,6 +73,8 @@
     "QKVLinear",
     "RMSNorm",
     "RoPE",
+    "RoPECacheReader",
+    "register_rope_cache",
     "ScaledBiasRowwiseLinear",
     "ScaledDotProductAttention",
     "SiLU",
diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py
index a208df1244..a903f1be1b 100644
--- a/torchtitan/models/common/decoder.py
+++ b/torchtitan/models/common/decoder.py
@@ -9,7 +9,7 @@
 from dataclasses import dataclass
 
 import torch
-from torch.nn.attention.flex_attention import _mask_mod_signature, and_masks, BlockMask
+from torch.nn.attention.flex_attention import BlockMask, _mask_mod_signature, and_masks
 
 from torchtitan.distributed.utils import is_in_batch_invariant_mode
 from torchtitan.models.common.attention import (
@@ -28,6 +28,12 @@
 from torchtitan.models.common.linear import Linear
 from torchtitan.models.common.moe import MoE
 from torchtitan.models.common.nn_modules import RMSNorm
+from torchtitan.models.common.rope import (
+    RoPE,
+    _RoPECacheRegistry,
+    _current_rope_cache_registry,
+    _rope_cache_registry_context,
+)
 from torchtitan.models.common.token_dispatcher import update_ep_token_dispatcher_config
 from torchtitan.protocols.model import BaseModel
 from torchtitan.protocols.module import Module, ModuleDict
@@ -82,6 +88,17 @@ class Config(BaseModel.Config):
         # itself is handled by ``Decoder.__init__`` / ``Decoder.init_states``.
         enable_weight_tying: bool = False
 
+        def build(self, **kwargs):
+            # Keep one registry active for the complete owner construction so
+            # subclass-added modules (for example DeepSeek MTP layers) observe
+            # the same cache namespace. The public build signature and the
+            # per-layer Config copy semantics remain unchanged.
+            registry = _RoPECacheRegistry()
+            with _rope_cache_registry_context(registry):
+                # ``slots=True`` prevents zero-argument ``super()`` from
+                # resolving reliably on dataclass-generated Config classes.
+                return Module.Config.build(self, **kwargs)
+
         @property
         def first_attention(self) -> BaseAttention.Config | None:
             """Attention config of the first layer that has one, else None.
@@ -249,11 +266,22 @@ def __init__(self, config: Config):
         super().__init__()
         self.config = config
 
+        registry = _current_rope_cache_registry()
+        if registry is None:
+            registry = _RoPECacheRegistry()
+        # The registry is a private child module so its canonical buffers follow
+        # model transforms and are copied with pipeline-stage model chunks.
+        self._rope_cache_registry = registry
+
         self.tok_embeddings = config.tok_embeddings.build()
 
-        self.layers = ModuleDict()
-        for i, layer_config in enumerate(config.layers):
-            self.layers[str(i)] = layer_config.build()
+        # Direct model construction (outside Config.build()) still receives a
+        # local registry for its layer RoPE modules. Config.build() establishes
+        # the same context around the complete subclass constructor.
+        with _rope_cache_registry_context(registry):
+            self.layers = ModuleDict()
+            for i, layer_config in enumerate(config.layers):
+                self.layers[str(i)] = layer_config.build()
 
         self.norm = config.norm.build()
         self.lm_head = config.lm_head.build()
@@ -276,6 +304,21 @@ def init_states(
             self.tok_embeddings.weight = self.lm_head.weight
         super().init_states(buffer_device=buffer_device)
 
+    def _prune_rope_cache_registry(self) -> None:
+        """Drop cache slots unused after pipeline-stage module pruning."""
+        slots = set()
+        for module in self.modules():
+            if not isinstance(module, RoPE):
+                continue
+            reader = module.__dict__.get("_cache_reader")
+            if reader is not None and reader._registry is self._rope_cache_registry:
+                slots.add(reader._slot_name)
+        self._rope_cache_registry._retain_slots(slots)
+
+    def _rope_cache_context(self):
+        """Return a context that exposes this model's private cache registry."""
+        return _rope_cache_registry_context(self._rope_cache_registry)
+
     def forward(
         self,
         tokens: torch.Tensor,
diff --git a/torchtitan/models/common/rope.py b/torchtitan/models/common/rope.py
index 88bb0ca2ff..24c5672302 100644
--- a/torchtitan/models/common/rope.py
+++ b/torchtitan/models/common/rope.py
@@ -4,23 +4,276 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import contextlib
+import contextvars
+import copy
 import math
-from dataclasses import dataclass
+from dataclasses import dataclass, field, fields, is_dataclass
 from typing import Literal
 
 import spmd_types as spmd
 import torch
-from torch.distributed.tensor import DTensor, Replicate, Shard
+from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor
 
+from torchtitan.distributed.parallel_dims import SpmdLayout
 from torchtitan.protocols.module import Module
+from torchtitan.protocols.sharding import resolve_placements
 
 __all__ = [
     "ComplexRoPE",
     "CosSinRoPE",
     "RoPE",
+    "RoPECacheReader",
+    "register_rope_cache",
 ]
 
 
+@dataclass(frozen=True, slots=True)
+class _RoPECacheKey:
+    """Hashable identity for a cache-producing RoPE configuration.
+
+    ``SpmdLayout`` contains a dictionary and is intentionally not hashable.  It
+    is carried on the key for the private registry to configure its canonical
+    buffer, while equality/hash use its stable representation.
+    """
+
+    value: tuple
+    cache_layout: SpmdLayout | None = field(default=None, compare=False, hash=False)
+
+    def __hash__(self) -> int:
+        return hash((self.value, repr(self.cache_layout)))
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, _RoPECacheKey):
+            return NotImplemented
+        return self.value == other.value and repr(self.cache_layout) == repr(
+            other.cache_layout
+        )
+
+
+@dataclass(frozen=True, slots=True)
+class _RoPECacheEntry:
+    slot_name: str
+    shape: tuple[int, ...]
+    dtype: torch.dtype
+    tensor_type: type
+
+
+_CURRENT_ROPE_CACHE_REGISTRY: contextvars.ContextVar[object | None] = (
+    contextvars.ContextVar("current_rope_cache_registry", default=None)
+)
+
+
+class RoPECacheReader:
+    """Read-only reference to a model-owned canonical RoPE cache.
+
+    The reader resolves the slot on every ``read()`` call instead of retaining
+    a tensor object.  Device/dtype transforms and DTensor distribution may
+    replace the registered buffer object during the model lifecycle.
+    """
+
+    __slots__ = ("_registry", "_slot_name")
+
+    def __init__(self, registry: "_RoPECacheRegistry", slot_name: str) -> None:
+        self._registry = registry
+        self._slot_name = slot_name
+
+    def read(self) -> torch.Tensor:
+        return self._registry._read(self._slot_name)
+
+    def __deepcopy__(self, memo: dict[int, object]):
+        # ``copy.deepcopy`` of a full model must bind readers to the copied
+        # registry, not leave them pointing at the source model's buffers.
+        registry = copy.deepcopy(self._registry, memo)
+        reader = type(self)(registry, self._slot_name)
+        memo[id(self)] = reader
+        return reader
+
+
+@contextlib.contextmanager
+def _rope_cache_registry_context(registry: "_RoPECacheRegistry"):
+    token = _CURRENT_ROPE_CACHE_REGISTRY.set(registry)
+    try:
+        yield registry
+    finally:
+        _CURRENT_ROPE_CACHE_REGISTRY.reset(token)
+
+
+def _current_rope_cache_registry() -> "_RoPECacheRegistry | None":
+    registry = _CURRENT_ROPE_CACHE_REGISTRY.get()
+    return registry if isinstance(registry, _RoPECacheRegistry) else None
+
+
+def register_rope_cache(
+    cache_key: _RoPECacheKey,
+    cache_tensor: torch.Tensor,
+) -> RoPECacheReader:
+    """Register or read a canonical cache in the active model registry.
+
+    This is the only public registration operation.  The private registry
+    retains one buffer per key; equivalent later registrations validate their
+    tensor metadata and discard the temporary duplicate.
+    """
+
+    registry = _current_rope_cache_registry()
+    if registry is None:
+        raise RuntimeError(
+            "register_rope_cache() requires an active model cache registry. "
+            "Standalone RoPE construction should use its local-buffer fallback."
+        )
+    return registry._register(cache_key, cache_tensor)
+
+
+class _RoPECacheRegistry(Module):
+    """Private module that owns canonical, non-persistent RoPE buffers."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self._entries: dict[_RoPECacheKey, _RoPECacheEntry] = {}
+        self._slot_keys: dict[str, _RoPECacheKey] = {}
+        self._next_slot = 0
+
+    @staticmethod
+    def _logical_metadata(cache: torch.Tensor) -> tuple[tuple[int, ...], torch.dtype]:
+        return tuple(cache.shape), cache.dtype
+
+    def _register(
+        self,
+        cache_key: _RoPECacheKey,
+        cache_tensor: torch.Tensor,
+    ) -> RoPECacheReader:
+        shape, dtype = self._logical_metadata(cache_tensor)
+        entry = self._entries.get(cache_key)
+        if entry is not None:
+            if (shape, dtype) != (entry.shape, entry.dtype):
+                raise ValueError(
+                    "Equivalent RoPE cache key produced incompatible tensors: "
+                    f"expected shape/dtype {(entry.shape, entry.dtype)}, got "
+                    f"{(shape, dtype)}. Include the missing cache-producing "
+                    "configuration in the cache key."
+                )
+            if type(cache_tensor) is not entry.tensor_type:
+                raise ValueError(
+                    "Equivalent RoPE cache key produced different tensor types: "
+                    f"expected {entry.tensor_type.__name__}, got "
+                    f"{type(cache_tensor).__name__}."
+                )
+            return RoPECacheReader(self, entry.slot_name)
+
+        slot_name = f"_cache_{self._next_slot}"
+        self._next_slot += 1
+        self.register_buffer(slot_name, cache_tensor, persistent=False)
+        self._entries[cache_key] = _RoPECacheEntry(
+            slot_name=slot_name,
+            shape=shape,
+            dtype=dtype,
+            tensor_type=type(cache_tensor),
+        )
+        self._slot_keys[slot_name] = cache_key
+        return RoPECacheReader(self, slot_name)
+
+    def _read(self, slot_name: str) -> torch.Tensor:
+        cache = self._buffers.get(slot_name)
+        if cache is None:
+            raise RuntimeError(f"RoPE cache slot {slot_name!r} is not materialized")
+        return cache
+
+    def _materialize(
+        self,
+        reader: RoPECacheReader,
+        cache_tensor: torch.Tensor,
+    ) -> None:
+        entry = self._entry_for_reader(reader)
+        shape, dtype = self._logical_metadata(cache_tensor)
+        if (shape, dtype) != (entry.shape, entry.dtype):
+            raise ValueError(
+                "RoPE cache materialization changed shape or dtype: "
+                f"expected {(entry.shape, entry.dtype)}, got {(shape, dtype)}."
+            )
+
+        current = self._read(entry.slot_name)
+        if isinstance(current, DTensor) and not isinstance(cache_tensor, DTensor):
+            cache_tensor = distribute_tensor(
+                cache_tensor,
+                current.device_mesh,
+                list(current.placements),
+            )
+
+        # Preserve the canonical object when possible.  This keeps any backend
+        # annotations attached to the registered buffer while accepting the
+        # temporary duplicate computed by each RoPE during initialization.
+        if (
+            type(current) is type(cache_tensor)
+            and current.device == cache_tensor.device
+            and tuple(current.shape) == tuple(cache_tensor.shape)
+            and current.dtype == cache_tensor.dtype
+        ):
+            with torch.no_grad():
+                current.copy_(cache_tensor)
+            return
+
+        persistent = entry.slot_name not in self._non_persistent_buffers_set
+        self.register_buffer(entry.slot_name, cache_tensor, persistent=persistent)
+
+    def _entry_for_reader(self, reader: RoPECacheReader) -> _RoPECacheEntry:
+        if reader._registry is not self:
+            raise ValueError("RoPECacheReader belongs to a different registry")
+        for entry in self._entries.values():
+            if entry.slot_name == reader._slot_name:
+                return entry
+        raise KeyError(f"Unknown RoPE cache slot {reader._slot_name!r}")
+
+    def parallelize(self, parallel_dims) -> None:
+        """Distribute canonical slots using the layout carried by each key."""
+        if self._parallelized:
+            raise ValueError(
+                f"{type(self).__name__} has already been parallelized. "
+                "Module.parallelize() must be called at most once per instance."
+            )
+        self._parallelized = True
+
+        for cache_key, entry in self._entries.items():
+            layout = cache_key.cache_layout
+            if layout is None:
+                continue
+            cache = self._read(entry.slot_name)
+            if parallel_dims.spmd_backend == "spmd_types":
+                self._spmd_distribute_state(
+                    parallel_dims,
+                    entry.slot_name,
+                    cache,
+                    layout,
+                    is_param=False,
+                )
+                continue
+            mesh = parallel_dims.resolve_mesh(layout.axes())
+            if mesh is None:
+                continue
+            placements = resolve_placements(layout, mesh)
+            if isinstance(cache, DTensor):
+                if tuple(cache.placements) != tuple(placements):
+                    raise ValueError(
+                        f"RoPE cache {entry.slot_name} has placements "
+                        f"{cache.placements}, expected {placements}."
+                    )
+                continue
+            self.register_buffer(
+                entry.slot_name,
+                distribute_tensor(cache, mesh, list(placements)),
+                persistent=False,
+            )
+
+    def _retain_slots(self, slot_names: set[str]) -> None:
+        """Drop canonical buffers that are not used by a PP model chunk."""
+        for key, entry in list(self._entries.items()):
+            if entry.slot_name in slot_names:
+                continue
+            self._entries.pop(key)
+            self._slot_keys.pop(entry.slot_name, None)
+            self._buffers.pop(entry.slot_name, None)
+            self._non_persistent_buffers_set.discard(entry.slot_name)
+
+
 # pyrefly: ignore [not-callable]
 @spmd.no_typecheck()
 def _maybe_check_max_pos(positions: torch.Tensor, *, max_valid_pos: int) -> None:
@@ -110,7 +363,89 @@ class Config(Module.Config):
     def __init__(self, config: Config):
         super().__init__()
         self.config = config
-        self.register_buffer("cache", self._precompute_cache(), persistent=False)
+        self._cache_reader: RoPECacheReader | None = None
+        cache = self._precompute_cache()
+        registry = _current_rope_cache_registry()
+        if registry is None:
+            self.register_buffer("cache", cache, persistent=False)
+        else:
+            self._cache_reader = register_rope_cache(self._cache_key(), cache)
+
+    @property
+    def cache(self) -> torch.Tensor:
+        """Return the current cache tensor for this RoPE instance.
+
+        Model-owned RoPE modules resolve a read-only reader into the private
+        registry's canonical buffer. Standalone RoPE modules retain the
+        historical direct ``cache`` buffer.
+        """
+        reader = self.__dict__.get("_cache_reader")
+        if reader is not None:
+            return reader.read()
+        if "cache" not in self._buffers:
+            # ``Module.register_buffer`` uses ``hasattr`` to reject collisions;
+            # report the pre-registration state as a missing attribute.
+            raise AttributeError("cache")
+        return self._buffers["cache"]
+
+    @cache.setter
+    def cache(self, cache: torch.Tensor) -> None:
+        reader = self.__dict__.get("_cache_reader")
+        if reader is not None:
+            registry = reader._registry
+            registry._materialize(reader, cache)
+            return
+        self.register_buffer("cache", cache, persistent=False)
+
+    def _cache_key(self) -> _RoPECacheKey:
+        """Identify cache-equivalent RoPE modules without sharing the modules.
+
+        The cache is a derived, read-only buffer.  Config metadata used only by
+        parameter initialization metadata is not cache-producing and is
+        excluded. The cache's declared state layout is included separately so
+        buffers with conflicting sharding requirements are never coalesced. A
+        concrete RoPE class is part of the key because subclasses may use a
+        different cache representation.
+        """
+
+        def freeze(value):
+            if is_dataclass(value):
+                return (
+                    type(value),
+                    tuple(
+                        (field.name, freeze(getattr(value, field.name)))
+                        for field in fields(value)
+                        if field.name not in {"param_init", "sharding_config"}
+                    ),
+                )
+            if isinstance(value, dict):
+                items = [(freeze(key), freeze(item)) for key, item in value.items()]
+                return tuple(sorted(items, key=repr))
+            if isinstance(value, (list, tuple)):
+                return tuple(freeze(item) for item in value)
+            if isinstance(value, set):
+                return tuple(sorted((freeze(item) for item in value), key=repr))
+            # Tensor-valued config fields are uncommon, but tensor equality is
+            # not a scalar operation and therefore cannot be used in a dict key.
+            if isinstance(value, torch.Tensor):
+                return type(value), id(value)
+            try:
+                hash(value)
+            except TypeError:
+                # Unknown mutable values are conservatively treated as unique.
+                return type(value), id(value)
+            return value
+
+        sharding = getattr(self.config, "sharding_config", None)
+        cache_layout = (
+            sharding.state_shardings.get("cache")
+            if sharding is not None
+            else None
+        )
+        return _RoPECacheKey(
+            value=(type(self), freeze(self.config)),
+            cache_layout=cache_layout,
+        )
 
     def _precompute_cache(self) -> torch.Tensor:
         """Build the reusable cache for all positions up to ``max_context_length``.
diff --git a/torchtitan/models/deepseek_v3/mtp.py b/torchtitan/models/deepseek_v3/mtp.py
index 90068b1d26..d9785f7f4f 100644
--- a/torchtitan/models/deepseek_v3/mtp.py
+++ b/torchtitan/models/deepseek_v3/mtp.py
@@ -212,13 +212,14 @@ def __init__(self, config: Config):
             return
 
         self.mtp_layers = ModuleList()
-        for layer_config in config.mtp_layers:
-            if not isinstance(layer_config, MTPTransformerBlock.Config):
-                raise ValueError(
-                    "MTPDecoder requires Config.mtp_layers to contain "
-                    "MTPTransformerBlock.Config instances."
-                )
-            self.mtp_layers.append(layer_config.build())
+        with self._rope_cache_context():
+            for layer_config in config.mtp_layers:
+                if not isinstance(layer_config, MTPTransformerBlock.Config):
+                    raise ValueError(
+                        "MTPDecoder requires Config.mtp_layers to contain "
+                        "MTPTransformerBlock.Config instances."
+                    )
+                self.mtp_layers.append(layer_config.build())
 
     def forward(
         self,

From e333a421e54a8e1f2f9e04fc46eeb8bf2489250c Mon Sep 17 00:00:00 2001
From: mystri <hanboyou@huawei.com>
Date: Sat, 29 Aug 2026 17:12:38 +0800
Subject: [PATCH 2/8] Simplify RoPE cache ownership

---
 torchtitan/distributed/pipeline_parallel.py |   5 +-
 torchtitan/models/common/decoder.py         |  39 ++---
 torchtitan/models/common/rope.py            | 164 ++++++--------------
 3 files changed, 66 insertions(+), 142 deletions(-)

diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py
index 8e2b08c586..2921405a45 100644
--- a/torchtitan/distributed/pipeline_parallel.py
+++ b/torchtitan/distributed/pipeline_parallel.py
@@ -523,10 +523,7 @@ def _split_module(
                 elif isinstance(module_value, (nn.ModuleList, ModuleList)):
                     setattr(model, module_name, ModuleList())
         # Handle simple module attributes (e.g., "linear", "norm")
-        elif (
-            module_name not in modules_to_keep
-            and module_name != "_rope_cache_registry"
-        ):
+        elif module_name not in modules_to_keep:
             # Replace with None
             setattr(model, module_name, None)
     prune_rope_caches = getattr(model, "_prune_rope_cache_registry", None)
diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py
index a903f1be1b..9af604df84 100644
--- a/torchtitan/models/common/decoder.py
+++ b/torchtitan/models/common/decoder.py
@@ -15,13 +15,13 @@
 from torchtitan.models.common.attention import (
     AttentionMasksType,
     BaseAttention,
+    FlexAttention,
+    ScaledDotProductAttention,
+    VarlenAttention,
     create_attention_mask,
     create_varlen_metadata_for_document,
-    FlexAttention,
     get_causal_mask_mod,
     get_efficient_causal_mask_mod_for_packed_document,
-    ScaledDotProductAttention,
-    VarlenAttention,
 )
 from torchtitan.models.common.embedding import Embedding
 from torchtitan.models.common.feed_forward import FeedForward
@@ -30,9 +30,8 @@
 from torchtitan.models.common.nn_modules import RMSNorm
 from torchtitan.models.common.rope import (
     RoPE,
-    _RoPECacheRegistry,
-    _current_rope_cache_registry,
     _rope_cache_registry_context,
+    _RoPECacheRegistry,
 )
 from torchtitan.models.common.token_dispatcher import update_ep_token_dispatcher_config
 from torchtitan.protocols.model import BaseModel
@@ -88,17 +87,6 @@ class Config(BaseModel.Config):
         # itself is handled by ``Decoder.__init__`` / ``Decoder.init_states``.
         enable_weight_tying: bool = False
 
-        def build(self, **kwargs):
-            # Keep one registry active for the complete owner construction so
-            # subclass-added modules (for example DeepSeek MTP layers) observe
-            # the same cache namespace. The public build signature and the
-            # per-layer Config copy semantics remain unchanged.
-            registry = _RoPECacheRegistry()
-            with _rope_cache_registry_context(registry):
-                # ``slots=True`` prevents zero-argument ``super()`` from
-                # resolving reliably on dataclass-generated Config classes.
-                return Module.Config.build(self, **kwargs)
-
         @property
         def first_attention(self) -> BaseAttention.Config | None:
             """Attention config of the first layer that has one, else None.
@@ -266,18 +254,15 @@ def __init__(self, config: Config):
         super().__init__()
         self.config = config
 
-        registry = _current_rope_cache_registry()
-        if registry is None:
-            registry = _RoPECacheRegistry()
-        # The registry is a private child module so its canonical buffers follow
-        # model transforms and are copied with pipeline-stage model chunks.
+        registry = _RoPECacheRegistry(self)
+        # The registry is a private owner for canonical non-persistent buffers;
+        # the registry object itself is not part of the module tree.
         self._rope_cache_registry = registry
 
         self.tok_embeddings = config.tok_embeddings.build()
 
-        # Direct model construction (outside Config.build()) still receives a
-        # local registry for its layer RoPE modules. Config.build() establishes
-        # the same context around the complete subclass constructor.
+        # Direct model construction and Config.build() both pass through this
+        # constructor, so the layer RoPE modules use the same local registry.
         with _rope_cache_registry_context(registry):
             self.layers = ModuleDict()
             for i, layer_config in enumerate(config.layers):
@@ -290,6 +275,10 @@ def __init__(self, config: Config):
         if self.enable_weight_tying:
             self.tok_embeddings.weight = self.lm_head.weight
 
+    def parallelize(self, parallel_dims) -> None:
+        super().parallelize(parallel_dims)
+        self._rope_cache_registry.parallelize(parallel_dims)
+
     def init_states(
         self,
         *,
@@ -310,7 +299,7 @@ def _prune_rope_cache_registry(self) -> None:
         for module in self.modules():
             if not isinstance(module, RoPE):
                 continue
-            reader = module.__dict__.get("_cache_reader")
+            reader = module._cache_reader
             if reader is not None and reader._registry is self._rope_cache_registry:
                 slots.add(reader._slot_name)
         self._rope_cache_registry._retain_slots(slots)
diff --git a/torchtitan/models/common/rope.py b/torchtitan/models/common/rope.py
index 24c5672302..52634f3d0e 100644
--- a/torchtitan/models/common/rope.py
+++ b/torchtitan/models/common/rope.py
@@ -6,9 +6,8 @@
 
 import contextlib
 import contextvars
-import copy
 import math
-from dataclasses import dataclass, field, fields, is_dataclass
+from dataclasses import dataclass, field, replace
 from typing import Literal
 
 import spmd_types as spmd
@@ -51,14 +50,6 @@ def __eq__(self, other: object) -> bool:
         )
 
 
-@dataclass(frozen=True, slots=True)
-class _RoPECacheEntry:
-    slot_name: str
-    shape: tuple[int, ...]
-    dtype: torch.dtype
-    tensor_type: type
-
-
 _CURRENT_ROPE_CACHE_REGISTRY: contextvars.ContextVar[object | None] = (
     contextvars.ContextVar("current_rope_cache_registry", default=None)
 )
@@ -81,14 +72,6 @@ def __init__(self, registry: "_RoPECacheRegistry", slot_name: str) -> None:
     def read(self) -> torch.Tensor:
         return self._registry._read(self._slot_name)
 
-    def __deepcopy__(self, memo: dict[int, object]):
-        # ``copy.deepcopy`` of a full model must bind readers to the copied
-        # registry, not leave them pointing at the source model's buffers.
-        registry = copy.deepcopy(self._registry, memo)
-        reader = type(self)(registry, self._slot_name)
-        memo[id(self)] = reader
-        return reader
-
 
 @contextlib.contextmanager
 def _rope_cache_registry_context(registry: "_RoPECacheRegistry"):
@@ -124,14 +107,14 @@ def register_rope_cache(
     return registry._register(cache_key, cache_tensor)
 
 
-class _RoPECacheRegistry(Module):
-    """Private module that owns canonical, non-persistent RoPE buffers."""
+class _RoPECacheRegistry:
+    """Private cache table that registers canonical buffers on its owner."""
 
-    def __init__(self) -> None:
-        super().__init__()
-        self._entries: dict[_RoPECacheKey, _RoPECacheEntry] = {}
-        self._slot_keys: dict[str, _RoPECacheKey] = {}
+    def __init__(self, owner: Module) -> None:
+        self._owner = owner
+        self._entries: dict[_RoPECacheKey, str] = {}
         self._next_slot = 0
+        self._parallelized = False
 
     @staticmethod
     def _logical_metadata(cache: torch.Tensor) -> tuple[tuple[int, ...], torch.dtype]:
@@ -143,37 +126,27 @@ def _register(
         cache_tensor: torch.Tensor,
     ) -> RoPECacheReader:
         shape, dtype = self._logical_metadata(cache_tensor)
-        entry = self._entries.get(cache_key)
-        if entry is not None:
-            if (shape, dtype) != (entry.shape, entry.dtype):
+        slot_name = self._entries.get(cache_key)
+        if slot_name is not None:
+            current = self._read(slot_name)
+            current_metadata = self._logical_metadata(current)
+            if (shape, dtype) != current_metadata:
                 raise ValueError(
                     "Equivalent RoPE cache key produced incompatible tensors: "
-                    f"expected shape/dtype {(entry.shape, entry.dtype)}, got "
+                    f"expected shape/dtype {current_metadata}, got "
                     f"{(shape, dtype)}. Include the missing cache-producing "
                     "configuration in the cache key."
                 )
-            if type(cache_tensor) is not entry.tensor_type:
-                raise ValueError(
-                    "Equivalent RoPE cache key produced different tensor types: "
-                    f"expected {entry.tensor_type.__name__}, got "
-                    f"{type(cache_tensor).__name__}."
-                )
-            return RoPECacheReader(self, entry.slot_name)
+            return RoPECacheReader(self, slot_name)
 
-        slot_name = f"_cache_{self._next_slot}"
+        slot_name = f"_rope_cache_{self._next_slot}"
         self._next_slot += 1
-        self.register_buffer(slot_name, cache_tensor, persistent=False)
-        self._entries[cache_key] = _RoPECacheEntry(
-            slot_name=slot_name,
-            shape=shape,
-            dtype=dtype,
-            tensor_type=type(cache_tensor),
-        )
-        self._slot_keys[slot_name] = cache_key
+        self._owner.register_buffer(slot_name, cache_tensor, persistent=False)
+        self._entries[cache_key] = slot_name
         return RoPECacheReader(self, slot_name)
 
     def _read(self, slot_name: str) -> torch.Tensor:
-        cache = self._buffers.get(slot_name)
+        cache = self._owner._buffers.get(slot_name)
         if cache is None:
             raise RuntimeError(f"RoPE cache slot {slot_name!r} is not materialized")
         return cache
@@ -183,15 +156,17 @@ def _materialize(
         reader: RoPECacheReader,
         cache_tensor: torch.Tensor,
     ) -> None:
-        entry = self._entry_for_reader(reader)
+        if reader._registry is not self or reader._slot_name not in self._owner._buffers:
+            raise ValueError("RoPECacheReader does not belong to this registry")
         shape, dtype = self._logical_metadata(cache_tensor)
-        if (shape, dtype) != (entry.shape, entry.dtype):
+        current = self._read(reader._slot_name)
+        current_metadata = self._logical_metadata(current)
+        if (shape, dtype) != current_metadata:
             raise ValueError(
                 "RoPE cache materialization changed shape or dtype: "
-                f"expected {(entry.shape, entry.dtype)}, got {(shape, dtype)}."
+                f"expected {current_metadata}, got {(shape, dtype)}."
             )
 
-        current = self._read(entry.slot_name)
         if isinstance(current, DTensor) and not isinstance(cache_tensor, DTensor):
             cache_tensor = distribute_tensor(
                 cache_tensor,
@@ -205,23 +180,12 @@ def _materialize(
         if (
             type(current) is type(cache_tensor)
             and current.device == cache_tensor.device
-            and tuple(current.shape) == tuple(cache_tensor.shape)
-            and current.dtype == cache_tensor.dtype
         ):
             with torch.no_grad():
                 current.copy_(cache_tensor)
             return
 
-        persistent = entry.slot_name not in self._non_persistent_buffers_set
-        self.register_buffer(entry.slot_name, cache_tensor, persistent=persistent)
-
-    def _entry_for_reader(self, reader: RoPECacheReader) -> _RoPECacheEntry:
-        if reader._registry is not self:
-            raise ValueError("RoPECacheReader belongs to a different registry")
-        for entry in self._entries.values():
-            if entry.slot_name == reader._slot_name:
-                return entry
-        raise KeyError(f"Unknown RoPE cache slot {reader._slot_name!r}")
+        self._owner.register_buffer(reader._slot_name, cache_tensor, persistent=False)
 
     def parallelize(self, parallel_dims) -> None:
         """Distribute canonical slots using the layout carried by each key."""
@@ -232,15 +196,15 @@ def parallelize(self, parallel_dims) -> None:
             )
         self._parallelized = True
 
-        for cache_key, entry in self._entries.items():
+        for cache_key, slot_name in self._entries.items():
             layout = cache_key.cache_layout
             if layout is None:
                 continue
-            cache = self._read(entry.slot_name)
+            cache = self._read(slot_name)
             if parallel_dims.spmd_backend == "spmd_types":
-                self._spmd_distribute_state(
+                self._owner._spmd_distribute_state(
                     parallel_dims,
-                    entry.slot_name,
+                    slot_name,
                     cache,
                     layout,
                     is_param=False,
@@ -253,25 +217,24 @@ def parallelize(self, parallel_dims) -> None:
             if isinstance(cache, DTensor):
                 if tuple(cache.placements) != tuple(placements):
                     raise ValueError(
-                        f"RoPE cache {entry.slot_name} has placements "
+                        f"RoPE cache {slot_name} has placements "
                         f"{cache.placements}, expected {placements}."
                     )
                 continue
-            self.register_buffer(
-                entry.slot_name,
+            self._owner.register_buffer(
+                slot_name,
                 distribute_tensor(cache, mesh, list(placements)),
                 persistent=False,
             )
 
     def _retain_slots(self, slot_names: set[str]) -> None:
         """Drop canonical buffers that are not used by a PP model chunk."""
-        for key, entry in list(self._entries.items()):
-            if entry.slot_name in slot_names:
+        for key, slot_name in list(self._entries.items()):
+            if slot_name in slot_names:
                 continue
             self._entries.pop(key)
-            self._slot_keys.pop(entry.slot_name, None)
-            self._buffers.pop(entry.slot_name, None)
-            self._non_persistent_buffers_set.discard(entry.slot_name)
+            self._owner._buffers.pop(slot_name, None)
+            self._owner._non_persistent_buffers_set.discard(slot_name)
 
 
 # pyrefly: ignore [not-callable]
@@ -379,7 +342,7 @@ def cache(self) -> torch.Tensor:
         registry's canonical buffer. Standalone RoPE modules retain the
         historical direct ``cache`` buffer.
         """
-        reader = self.__dict__.get("_cache_reader")
+        reader = self._cache_reader
         if reader is not None:
             return reader.read()
         if "cache" not in self._buffers:
@@ -388,15 +351,6 @@ def cache(self) -> torch.Tensor:
             raise AttributeError("cache")
         return self._buffers["cache"]
 
-    @cache.setter
-    def cache(self, cache: torch.Tensor) -> None:
-        reader = self.__dict__.get("_cache_reader")
-        if reader is not None:
-            registry = reader._registry
-            registry._materialize(reader, cache)
-            return
-        self.register_buffer("cache", cache, persistent=False)
-
     def _cache_key(self) -> _RoPECacheKey:
         """Identify cache-equivalent RoPE modules without sharing the modules.
 
@@ -408,42 +362,22 @@ def _cache_key(self) -> _RoPECacheKey:
         different cache representation.
         """
 
-        def freeze(value):
-            if is_dataclass(value):
-                return (
-                    type(value),
-                    tuple(
-                        (field.name, freeze(getattr(value, field.name)))
-                        for field in fields(value)
-                        if field.name not in {"param_init", "sharding_config"}
-                    ),
-                )
-            if isinstance(value, dict):
-                items = [(freeze(key), freeze(item)) for key, item in value.items()]
-                return tuple(sorted(items, key=repr))
-            if isinstance(value, (list, tuple)):
-                return tuple(freeze(item) for item in value)
-            if isinstance(value, set):
-                return tuple(sorted((freeze(item) for item in value), key=repr))
-            # Tensor-valued config fields are uncommon, but tensor equality is
-            # not a scalar operation and therefore cannot be used in a dict key.
-            if isinstance(value, torch.Tensor):
-                return type(value), id(value)
-            try:
-                hash(value)
-            except TypeError:
-                # Unknown mutable values are conservatively treated as unique.
-                return type(value), id(value)
-            return value
-
         sharding = getattr(self.config, "sharding_config", None)
         cache_layout = (
             sharding.state_shardings.get("cache")
             if sharding is not None
             else None
         )
+        # RoPE configs contain scalar values and a small number of list fields;
+        # repr gives us a compact, stable key while excluding runtime-only
+        # initialization and sharding metadata.
+        cache_config = replace(
+            self.config,
+            param_init=None,
+            sharding_config=None,
+        )
         return _RoPECacheKey(
-            value=(type(self), freeze(self.config)),
+            value=(type(self), repr(cache_config)),
             cache_layout=cache_layout,
         )
 
@@ -508,7 +442,11 @@ def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> No
             # Recompute there when the caller does not pass an explicit buffer device.
             buffer_device = self.cache.device
         with torch.device(buffer_device):
-            self.cache = self._precompute_cache()
+            cache = self._precompute_cache()
+            if self._cache_reader is None:
+                self.register_buffer("cache", cache, persistent=False)
+            else:
+                self._cache_reader._registry._materialize(self._cache_reader, cache)
 
 
 class ComplexRoPE(RoPE):

From d0929ae70530474496d2d4be30f0c1e4c540ac76 Mon Sep 17 00:00:00 2001
From: mystri <hanboyou@huawei.com>
Date: Sat, 29 Aug 2026 17:48:20 +0800
Subject: [PATCH 3/8] Defer RoPE cache parallelization

---
 torchtitan/distributed/pipeline_parallel.py |   3 -
 torchtitan/models/common/decoder.py         |  16 --
 torchtitan/models/common/rope.py            | 157 ++------------------
 3 files changed, 16 insertions(+), 160 deletions(-)

diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py
index 2921405a45..f547b18cd4 100644
--- a/torchtitan/distributed/pipeline_parallel.py
+++ b/torchtitan/distributed/pipeline_parallel.py
@@ -526,9 +526,6 @@ def _split_module(
         elif module_name not in modules_to_keep:
             # Replace with None
             setattr(model, module_name, None)
-    prune_rope_caches = getattr(model, "_prune_rope_cache_registry", None)
-    if callable(prune_rope_caches):
-        prune_rope_caches()
     return model
 
 
diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py
index 9af604df84..3285b1ba32 100644
--- a/torchtitan/models/common/decoder.py
+++ b/torchtitan/models/common/decoder.py
@@ -29,7 +29,6 @@
 from torchtitan.models.common.moe import MoE
 from torchtitan.models.common.nn_modules import RMSNorm
 from torchtitan.models.common.rope import (
-    RoPE,
     _rope_cache_registry_context,
     _RoPECacheRegistry,
 )
@@ -275,10 +274,6 @@ def __init__(self, config: Config):
         if self.enable_weight_tying:
             self.tok_embeddings.weight = self.lm_head.weight
 
-    def parallelize(self, parallel_dims) -> None:
-        super().parallelize(parallel_dims)
-        self._rope_cache_registry.parallelize(parallel_dims)
-
     def init_states(
         self,
         *,
@@ -293,17 +288,6 @@ def init_states(
             self.tok_embeddings.weight = self.lm_head.weight
         super().init_states(buffer_device=buffer_device)
 
-    def _prune_rope_cache_registry(self) -> None:
-        """Drop cache slots unused after pipeline-stage module pruning."""
-        slots = set()
-        for module in self.modules():
-            if not isinstance(module, RoPE):
-                continue
-            reader = module._cache_reader
-            if reader is not None and reader._registry is self._rope_cache_registry:
-                slots.add(reader._slot_name)
-        self._rope_cache_registry._retain_slots(slots)
-
     def _rope_cache_context(self):
         """Return a context that exposes this model's private cache registry."""
         return _rope_cache_registry_context(self._rope_cache_registry)
diff --git a/torchtitan/models/common/rope.py b/torchtitan/models/common/rope.py
index 52634f3d0e..123f34e7e6 100644
--- a/torchtitan/models/common/rope.py
+++ b/torchtitan/models/common/rope.py
@@ -7,16 +7,14 @@
 import contextlib
 import contextvars
 import math
-from dataclasses import dataclass, field, replace
+from dataclasses import dataclass, replace
 from typing import Literal
 
 import spmd_types as spmd
 import torch
-from torch.distributed.tensor import DTensor, Replicate, Shard, distribute_tensor
+from torch.distributed.tensor import DTensor, Replicate, Shard
 
-from torchtitan.distributed.parallel_dims import SpmdLayout
 from torchtitan.protocols.module import Module
-from torchtitan.protocols.sharding import resolve_placements
 
 __all__ = [
     "ComplexRoPE",
@@ -27,29 +25,6 @@
 ]
 
 
-@dataclass(frozen=True, slots=True)
-class _RoPECacheKey:
-    """Hashable identity for a cache-producing RoPE configuration.
-
-    ``SpmdLayout`` contains a dictionary and is intentionally not hashable.  It
-    is carried on the key for the private registry to configure its canonical
-    buffer, while equality/hash use its stable representation.
-    """
-
-    value: tuple
-    cache_layout: SpmdLayout | None = field(default=None, compare=False, hash=False)
-
-    def __hash__(self) -> int:
-        return hash((self.value, repr(self.cache_layout)))
-
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, _RoPECacheKey):
-            return NotImplemented
-        return self.value == other.value and repr(self.cache_layout) == repr(
-            other.cache_layout
-        )
-
-
 _CURRENT_ROPE_CACHE_REGISTRY: contextvars.ContextVar[object | None] = (
     contextvars.ContextVar("current_rope_cache_registry", default=None)
 )
@@ -88,14 +63,14 @@ def _current_rope_cache_registry() -> "_RoPECacheRegistry | None":
 
 
 def register_rope_cache(
-    cache_key: _RoPECacheKey,
+    cache_key: str,
     cache_tensor: torch.Tensor,
 ) -> RoPECacheReader:
     """Register or read a canonical cache in the active model registry.
 
     This is the only public registration operation.  The private registry
-    retains one buffer per key; equivalent later registrations validate their
-    tensor metadata and discard the temporary duplicate.
+    retains one buffer per key; equivalent later registrations return a reader
+    for the existing buffer and discard the temporary duplicate.
     """
 
     registry = _current_rope_cache_registry()
@@ -112,31 +87,16 @@ class _RoPECacheRegistry:
 
     def __init__(self, owner: Module) -> None:
         self._owner = owner
-        self._entries: dict[_RoPECacheKey, str] = {}
+        self._entries: dict[str, str] = {}
         self._next_slot = 0
-        self._parallelized = False
-
-    @staticmethod
-    def _logical_metadata(cache: torch.Tensor) -> tuple[tuple[int, ...], torch.dtype]:
-        return tuple(cache.shape), cache.dtype
 
     def _register(
         self,
-        cache_key: _RoPECacheKey,
+        cache_key: str,
         cache_tensor: torch.Tensor,
     ) -> RoPECacheReader:
-        shape, dtype = self._logical_metadata(cache_tensor)
         slot_name = self._entries.get(cache_key)
         if slot_name is not None:
-            current = self._read(slot_name)
-            current_metadata = self._logical_metadata(current)
-            if (shape, dtype) != current_metadata:
-                raise ValueError(
-                    "Equivalent RoPE cache key produced incompatible tensors: "
-                    f"expected shape/dtype {current_metadata}, got "
-                    f"{(shape, dtype)}. Include the missing cache-producing "
-                    "configuration in the cache key."
-                )
             return RoPECacheReader(self, slot_name)
 
         slot_name = f"_rope_cache_{self._next_slot}"
@@ -158,84 +118,8 @@ def _materialize(
     ) -> None:
         if reader._registry is not self or reader._slot_name not in self._owner._buffers:
             raise ValueError("RoPECacheReader does not belong to this registry")
-        shape, dtype = self._logical_metadata(cache_tensor)
-        current = self._read(reader._slot_name)
-        current_metadata = self._logical_metadata(current)
-        if (shape, dtype) != current_metadata:
-            raise ValueError(
-                "RoPE cache materialization changed shape or dtype: "
-                f"expected {current_metadata}, got {(shape, dtype)}."
-            )
-
-        if isinstance(current, DTensor) and not isinstance(cache_tensor, DTensor):
-            cache_tensor = distribute_tensor(
-                cache_tensor,
-                current.device_mesh,
-                list(current.placements),
-            )
-
-        # Preserve the canonical object when possible.  This keeps any backend
-        # annotations attached to the registered buffer while accepting the
-        # temporary duplicate computed by each RoPE during initialization.
-        if (
-            type(current) is type(cache_tensor)
-            and current.device == cache_tensor.device
-        ):
-            with torch.no_grad():
-                current.copy_(cache_tensor)
-            return
-
         self._owner.register_buffer(reader._slot_name, cache_tensor, persistent=False)
 
-    def parallelize(self, parallel_dims) -> None:
-        """Distribute canonical slots using the layout carried by each key."""
-        if self._parallelized:
-            raise ValueError(
-                f"{type(self).__name__} has already been parallelized. "
-                "Module.parallelize() must be called at most once per instance."
-            )
-        self._parallelized = True
-
-        for cache_key, slot_name in self._entries.items():
-            layout = cache_key.cache_layout
-            if layout is None:
-                continue
-            cache = self._read(slot_name)
-            if parallel_dims.spmd_backend == "spmd_types":
-                self._owner._spmd_distribute_state(
-                    parallel_dims,
-                    slot_name,
-                    cache,
-                    layout,
-                    is_param=False,
-                )
-                continue
-            mesh = parallel_dims.resolve_mesh(layout.axes())
-            if mesh is None:
-                continue
-            placements = resolve_placements(layout, mesh)
-            if isinstance(cache, DTensor):
-                if tuple(cache.placements) != tuple(placements):
-                    raise ValueError(
-                        f"RoPE cache {slot_name} has placements "
-                        f"{cache.placements}, expected {placements}."
-                    )
-                continue
-            self._owner.register_buffer(
-                slot_name,
-                distribute_tensor(cache, mesh, list(placements)),
-                persistent=False,
-            )
-
-    def _retain_slots(self, slot_names: set[str]) -> None:
-        """Drop canonical buffers that are not used by a PP model chunk."""
-        for key, slot_name in list(self._entries.items()):
-            if slot_name in slot_names:
-                continue
-            self._entries.pop(key)
-            self._owner._buffers.pop(slot_name, None)
-            self._owner._non_persistent_buffers_set.discard(slot_name)
-
 
 # pyrefly: ignore [not-callable]
 @spmd.no_typecheck()
@@ -332,7 +216,7 @@ def __init__(self, config: Config):
         if registry is None:
             self.register_buffer("cache", cache, persistent=False)
         else:
-            self._cache_reader = register_rope_cache(self._cache_key(), cache)
+            self._cache_reader = register_rope_cache(self._cache_key(cache), cache)
 
     @property
     def cache(self) -> torch.Tensor:
@@ -351,34 +235,25 @@ def cache(self) -> torch.Tensor:
             raise AttributeError("cache")
         return self._buffers["cache"]
 
-    def _cache_key(self) -> _RoPECacheKey:
+    def _cache_key(self, cache: torch.Tensor) -> str:
         """Identify cache-equivalent RoPE modules without sharing the modules.
 
         The cache is a derived, read-only buffer.  Config metadata used only by
         parameter initialization metadata is not cache-producing and is
-        excluded. The cache's declared state layout is included separately so
-        buffers with conflicting sharding requirements are never coalesced. A
-        concrete RoPE class is part of the key because subclasses may use a
-        different cache representation.
+        excluded. A concrete RoPE class is part of the key because subclasses
+        may use a different cache representation.
         """
 
-        sharding = getattr(self.config, "sharding_config", None)
-        cache_layout = (
-            sharding.state_shardings.get("cache")
-            if sharding is not None
-            else None
-        )
-        # RoPE configs contain scalar values and a small number of list fields;
-        # repr gives us a compact, stable key while excluding runtime-only
-        # initialization and sharding metadata.
+        # RoPE configs contain scalar values and a small number of list fields.
+        # Their repr plus the candidate's physical metadata is sufficient for
+        # this construction-time cache table.
         cache_config = replace(
             self.config,
             param_init=None,
             sharding_config=None,
         )
-        return _RoPECacheKey(
-            value=(type(self), repr(cache_config)),
-            cache_layout=cache_layout,
+        return repr(
+            (type(self), cache_config, tuple(cache.shape), cache.dtype, cache.device)
         )
 
     def _precompute_cache(self) -> torch.Tensor:

From 493e395369a1ae61424e7fad7f42cc32d9c87bfa Mon Sep 17 00:00:00 2001
From: mystri <hanboyou@huawei.com>
Date: Wed, 2 Sep 2026 00:08:57 +0800
Subject: [PATCH 4/8] Share RoPE modules by descriptive cache key

---
 torchtitan/distributed/pipeline_parallel.py |   4 +
 torchtitan/models/common/__init__.py        |   4 -
 torchtitan/models/common/attention.py       |  15 +-
 torchtitan/models/common/decoder.py         |  47 ++--
 torchtitan/models/common/rope.py            | 228 +++++++-------------
 torchtitan/models/deepseek_v3/model.py      |  14 +-
 torchtitan/models/deepseek_v3/mtp.py        |  30 +--
 torchtitan/models/gpt_oss/model.py          |  14 +-
 torchtitan/models/llama3/model.py           |   5 +-
 torchtitan/models/muse_glimmer/model.py     |  10 +-
 torchtitan/models/qwen3/model.py            |   5 +-
 torchtitan/models/qwen3_5/model.py          |  16 +-
 torchtitan/overrides/fused_mla.py           |   5 +-
 13 files changed, 182 insertions(+), 215 deletions(-)

diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py
index f547b18cd4..9a4662417e 100644
--- a/torchtitan/distributed/pipeline_parallel.py
+++ b/torchtitan/distributed/pipeline_parallel.py
@@ -488,7 +488,11 @@ def _split_module(
     model = copy.deepcopy(whole_model)
     # Create a set of modules to keep for faster lookup
     modules_to_keep = set(module_names)
+    if "rope_modules" in model._modules:
+        modules_to_keep.add("rope_modules")
     for module_name, module_value in model.named_children():
+        if module_name in modules_to_keep:
+            continue
         # Handle layer-like structures (e.g., "layers.0", "layers.1")
         if isinstance(
             module_value, (nn.ModuleDict, nn.ModuleList, ModuleDict, ModuleList)
diff --git a/torchtitan/models/common/__init__.py b/torchtitan/models/common/__init__.py
index f9e9ac525c..b5f0db324f 100644
--- a/torchtitan/models/common/__init__.py
+++ b/torchtitan/models/common/__init__.py
@@ -40,8 +40,6 @@
     ComplexRoPE,
     CosSinRoPE,
     RoPE,
-    RoPECacheReader,
-    register_rope_cache,
 )
 
 __all__ = [
@@ -73,8 +71,6 @@
     "QKVLinear",
     "RMSNorm",
     "RoPE",
-    "RoPECacheReader",
-    "register_rope_cache",
     "ScaledBiasRowwiseLinear",
     "ScaledDotProductAttention",
     "SiLU",
diff --git a/torchtitan/models/common/attention.py b/torchtitan/models/common/attention.py
index 5951b6b97a..e76f6c5131 100644
--- a/torchtitan/models/common/attention.py
+++ b/torchtitan/models/common/attention.py
@@ -13,7 +13,7 @@
 
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field
-from typing import Any, ClassVar, NamedTuple
+from typing import Any, ClassVar, NamedTuple, cast
 
 import spmd_types as spmd
 import torch
@@ -43,7 +43,7 @@
 from torchtitan.models.common.linear import Linear
 from torchtitan.models.common.nn_modules import RMSNorm
 from torchtitan.models.common.rope import RoPE
-from torchtitan.protocols.module import Module
+from torchtitan.protocols.module import Module, ModuleDict
 from torchtitan.tools.utils import round_up
 
 
@@ -667,6 +667,10 @@ def __post_init__(self):
             assert self.n_heads > 0, "n_heads must be > 0"
 
 
+def _resolve_rope(rope_config: RoPE.Config, rope_modules: ModuleDict) -> RoPE:
+    return cast(RoPE, rope_modules[rope_config.rope_key()])
+
+
 class BaseQKVLinear(Module):
     """Base class for Q/K/V projection strategies.
 
@@ -892,6 +896,8 @@ class GQAttention(BaseAttention):
     :class:`FusedQKVLinear` for a single fused projection.
     """
 
+    rope: RoPE
+
     @dataclass(kw_only=True, slots=True)
     class Config(BaseAttention.Config):
         n_heads: int
@@ -919,7 +925,7 @@ def __post_init__(self) -> None:
                     f"n_kv_heads ({n_kv_heads})"
                 )
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
         self.n_heads = config.n_heads
         self.n_kv_heads = (
@@ -931,7 +937,8 @@ def __init__(self, config: Config):
             else config.dim // config.n_heads
         )
         self.enable_gqa = self.n_heads > self.n_kv_heads
-        self.rope = config.rope.build()
+        # Keep the canonical module registered only under Decoder.rope_modules.
+        object.__setattr__(self, "rope", _resolve_rope(config.rope, rope_modules))
 
         # Pluggable QKV projection
         self.qkv_linear = config.qkv_linear.build()
diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py
index 3285b1ba32..efe3b6316a 100644
--- a/torchtitan/models/common/decoder.py
+++ b/torchtitan/models/common/decoder.py
@@ -28,10 +28,6 @@
 from torchtitan.models.common.linear import Linear
 from torchtitan.models.common.moe import MoE
 from torchtitan.models.common.nn_modules import RMSNorm
-from torchtitan.models.common.rope import (
-    _rope_cache_registry_context,
-    _RoPECacheRegistry,
-)
 from torchtitan.models.common.token_dispatcher import update_ep_token_dispatcher_config
 from torchtitan.protocols.model import BaseModel
 from torchtitan.protocols.module import Module, ModuleDict
@@ -39,6 +35,22 @@
 __all__ = ["Decoder", "TransformerBlock"]
 
 
+def _rope_config(layer_config):
+    attention_config = getattr(layer_config, "attention", None)
+    return getattr(attention_config, "rope", None)
+
+
+def _register_rope_modules(layer_configs, rope_modules: ModuleDict) -> None:
+    """Build one canonical RoPE module for each layer configuration key."""
+    for layer_config in layer_configs:
+        rope_config = _rope_config(layer_config)
+        if rope_config is None:
+            continue
+        key = rope_config.rope_key()
+        if key not in rope_modules:
+            rope_modules[key] = rope_config.build()
+
+
 # TODO: we can unify the TransformerBlock impl across all models when
 # there is no special logic for each model, including
 # ffn vs. moe naming and creation, etc.
@@ -252,20 +264,19 @@ def update_from_config(
     def __init__(self, config: Config):
         super().__init__()
         self.config = config
-
-        registry = _RoPECacheRegistry(self)
-        # The registry is a private owner for canonical non-persistent buffers;
-        # the registry object itself is not part of the module tree.
-        self._rope_cache_registry = registry
-
         self.tok_embeddings = config.tok_embeddings.build()
 
-        # Direct model construction and Config.build() both pass through this
-        # constructor, so the layer RoPE modules use the same local registry.
-        with _rope_cache_registry_context(registry):
-            self.layers = ModuleDict()
-            for i, layer_config in enumerate(config.layers):
-                self.layers[str(i)] = layer_config.build()
+        self.rope_modules = ModuleDict()
+        _register_rope_modules(config.layers, self.rope_modules)
+
+        self.layers = ModuleDict()
+        for i, layer_config in enumerate(config.layers):
+            rope_config = _rope_config(layer_config)
+            if rope_config is None:
+                layer = layer_config.build()
+            else:
+                layer = layer_config.build(rope_modules=self.rope_modules)
+            self.layers[str(i)] = layer
 
         self.norm = config.norm.build()
         self.lm_head = config.lm_head.build()
@@ -288,10 +299,6 @@ def init_states(
             self.tok_embeddings.weight = self.lm_head.weight
         super().init_states(buffer_device=buffer_device)
 
-    def _rope_cache_context(self):
-        """Return a context that exposes this model's private cache registry."""
-        return _rope_cache_registry_context(self._rope_cache_registry)
-
     def forward(
         self,
         tokens: torch.Tensor,
diff --git a/torchtitan/models/common/rope.py b/torchtitan/models/common/rope.py
index 123f34e7e6..685313b3eb 100644
--- a/torchtitan/models/common/rope.py
+++ b/torchtitan/models/common/rope.py
@@ -4,10 +4,9 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-import contextlib
-import contextvars
 import math
-from dataclasses import dataclass, replace
+import re
+from dataclasses import dataclass, fields
 from typing import Literal
 
 import spmd_types as spmd
@@ -20,105 +19,22 @@
     "ComplexRoPE",
     "CosSinRoPE",
     "RoPE",
-    "RoPECacheReader",
-    "register_rope_cache",
 ]
 
 
-_CURRENT_ROPE_CACHE_REGISTRY: contextvars.ContextVar[object | None] = (
-    contextvars.ContextVar("current_rope_cache_registry", default=None)
-)
-
-
-class RoPECacheReader:
-    """Read-only reference to a model-owned canonical RoPE cache.
-
-    The reader resolves the slot on every ``read()`` call instead of retaining
-    a tensor object.  Device/dtype transforms and DTensor distribution may
-    replace the registered buffer object during the model lifecycle.
-    """
-
-    __slots__ = ("_registry", "_slot_name")
-
-    def __init__(self, registry: "_RoPECacheRegistry", slot_name: str) -> None:
-        self._registry = registry
-        self._slot_name = slot_name
-
-    def read(self) -> torch.Tensor:
-        return self._registry._read(self._slot_name)
-
-
-@contextlib.contextmanager
-def _rope_cache_registry_context(registry: "_RoPECacheRegistry"):
-    token = _CURRENT_ROPE_CACHE_REGISTRY.set(registry)
-    try:
-        yield registry
-    finally:
-        _CURRENT_ROPE_CACHE_REGISTRY.reset(token)
-
-
-def _current_rope_cache_registry() -> "_RoPECacheRegistry | None":
-    registry = _CURRENT_ROPE_CACHE_REGISTRY.get()
-    return registry if isinstance(registry, _RoPECacheRegistry) else None
-
-
-def register_rope_cache(
-    cache_key: str,
-    cache_tensor: torch.Tensor,
-) -> RoPECacheReader:
-    """Register or read a canonical cache in the active model registry.
-
-    This is the only public registration operation.  The private registry
-    retains one buffer per key; equivalent later registrations return a reader
-    for the existing buffer and discard the temporary duplicate.
-    """
-
-    registry = _current_rope_cache_registry()
-    if registry is None:
-        raise RuntimeError(
-            "register_rope_cache() requires an active model cache registry. "
-            "Standalone RoPE construction should use its local-buffer fallback."
-        )
-    return registry._register(cache_key, cache_tensor)
-
-
-class _RoPECacheRegistry:
-    """Private cache table that registers canonical buffers on its owner."""
-
-    def __init__(self, owner: Module) -> None:
-        self._owner = owner
-        self._entries: dict[str, str] = {}
-        self._next_slot = 0
-
-    def _register(
-        self,
-        cache_key: str,
-        cache_tensor: torch.Tensor,
-    ) -> RoPECacheReader:
-        slot_name = self._entries.get(cache_key)
-        if slot_name is not None:
-            return RoPECacheReader(self, slot_name)
-
-        slot_name = f"_rope_cache_{self._next_slot}"
-        self._next_slot += 1
-        self._owner.register_buffer(slot_name, cache_tensor, persistent=False)
-        self._entries[cache_key] = slot_name
-        return RoPECacheReader(self, slot_name)
-
-    def _read(self, slot_name: str) -> torch.Tensor:
-        cache = self._owner._buffers.get(slot_name)
-        if cache is None:
-            raise RuntimeError(f"RoPE cache slot {slot_name!r} is not materialized")
-        return cache
-
-    def _materialize(
-        self,
-        reader: RoPECacheReader,
-        cache_tensor: torch.Tensor,
-    ) -> None:
-        if reader._registry is not self or reader._slot_name not in self._owner._buffers:
-            raise ValueError("RoPECacheReader does not belong to this registry")
-        self._owner.register_buffer(reader._slot_name, cache_tensor, persistent=False)
+def _format_rope_key_value(value: object) -> str:
+    """Format config values into safe, readable ``ModuleDict`` key fragments."""
+    if value is None:
+        text = "none"
+    elif isinstance(value, bool):
+        text = "true" if value else "false"
+    elif isinstance(value, float):
+        text = format(value, ".17g")
+    elif isinstance(value, (list, tuple)):
+        text = "-".join(_format_rope_key_value(item) for item in value)
+    else:
+        text = str(value)
+    return re.sub(r"[^A-Za-z0-9_-]", "p", text)
 
 
 # pyrefly: ignore [not-callable]
@@ -189,6 +105,8 @@ class RoPE(Module):
     cosine/sine caches.
     """
 
+    cache: torch.Tensor
+
     @dataclass(kw_only=True, slots=True)
     class Config(Module.Config):
         dim: int
@@ -207,54 +125,70 @@ class Config(Module.Config):
         original_seq_len: int = 4096
         truncate: bool = True
 
+        def rope_key(self) -> str:
+            """Return the stable, descriptive key for this RoPE implementation."""
+            owner = self._owner
+            owner_name = owner.__name__ if owner is not None else type(self).__name__
+            parts = [owner_name]
+            common_fields = {"dim", "max_context_length", "theta", "scaling"}
+            scaling_fields = {
+                "llama": {
+                    "scaling_factor",
+                    "low_freq_factor",
+                    "high_freq_factor",
+                    "original_max_position_embeddings",
+                },
+                "yarn": {
+                    "rope_factor",
+                    "beta_fast",
+                    "beta_slow",
+                    "original_seq_len",
+                    "truncate",
+                },
+            }.get(self.scaling, set())
+            key_names = {
+                "dim": "d",
+                "max_context_length": "ctx",
+                "scaling_factor": "sf",
+                "low_freq_factor": "lf",
+                "high_freq_factor": "hf",
+                "original_max_position_embeddings": "orig",
+                "rope_factor": "rf",
+                "beta_fast": "bf",
+                "beta_slow": "bs",
+                "original_seq_len": "orig",
+                "truncate": "trunc",
+            }
+            for config_field in fields(self):
+                if config_field.name in {"param_init", "sharding_config"}:
+                    continue
+                if (
+                    config_field.name not in common_fields
+                    and config_field.name not in scaling_fields
+                    and config_field.name
+                    in {
+                        "scaling_factor",
+                        "low_freq_factor",
+                        "high_freq_factor",
+                        "original_max_position_embeddings",
+                        "rope_factor",
+                        "beta_fast",
+                        "beta_slow",
+                        "original_seq_len",
+                        "truncate",
+                    }
+                ):
+                    continue
+                value = _format_rope_key_value(getattr(self, config_field.name))
+                name = key_names.get(config_field.name, config_field.name)
+                separator = "" if config_field.name in common_fields - {"scaling"} else "_"
+                parts.append(f"{name}{separator}{value}")
+            return "_".join(parts)
+
     def __init__(self, config: Config):
         super().__init__()
         self.config = config
-        self._cache_reader: RoPECacheReader | None = None
-        cache = self._precompute_cache()
-        registry = _current_rope_cache_registry()
-        if registry is None:
-            self.register_buffer("cache", cache, persistent=False)
-        else:
-            self._cache_reader = register_rope_cache(self._cache_key(cache), cache)
-
-    @property
-    def cache(self) -> torch.Tensor:
-        """Return the current cache tensor for this RoPE instance.
-
-        Model-owned RoPE modules resolve a read-only reader into the private
-        registry's canonical buffer. Standalone RoPE modules retain the
-        historical direct ``cache`` buffer.
-        """
-        reader = self._cache_reader
-        if reader is not None:
-            return reader.read()
-        if "cache" not in self._buffers:
-            # ``Module.register_buffer`` uses ``hasattr`` to reject collisions;
-            # report the pre-registration state as a missing attribute.
-            raise AttributeError("cache")
-        return self._buffers["cache"]
-
-    def _cache_key(self, cache: torch.Tensor) -> str:
-        """Identify cache-equivalent RoPE modules without sharing the modules.
-
-        The cache is a derived, read-only buffer.  Config metadata used only by
-        parameter initialization metadata is not cache-producing and is
-        excluded. A concrete RoPE class is part of the key because subclasses
-        may use a different cache representation.
-        """
-
-        # RoPE configs contain scalar values and a small number of list fields.
-        # Their repr plus the candidate's physical metadata is sufficient for
-        # this construction-time cache table.
-        cache_config = replace(
-            self.config,
-            param_init=None,
-            sharding_config=None,
-        )
-        return repr(
-            (type(self), cache_config, tuple(cache.shape), cache.dtype, cache.device)
-        )
+        self.register_buffer("cache", self._precompute_cache(), persistent=False)
 
     def _precompute_cache(self) -> torch.Tensor:
         """Build the reusable cache for all positions up to ``max_context_length``.
@@ -311,17 +245,13 @@ def forward(
         return self.apply_rotary_emb(query, key, reshaped_cache)
 
     def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None:
-        # TODO: In long-term we need to have buffer abstraction in `Module`` class to infer the buffer_device
         if buffer_device is None:
             # After ``to_empty()``, the existing cache records the target device.
             # Recompute there when the caller does not pass an explicit buffer device.
             buffer_device = self.cache.device
         with torch.device(buffer_device):
             cache = self._precompute_cache()
-            if self._cache_reader is None:
-                self.register_buffer("cache", cache, persistent=False)
-            else:
-                self._cache_reader._registry._materialize(self._cache_reader, cache)
+            self.register_buffer("cache", cache, persistent=False)
 
 
 class ComplexRoPE(RoPE):
diff --git a/torchtitan/models/deepseek_v3/model.py b/torchtitan/models/deepseek_v3/model.py
index f48c042bd8..88de5629fa 100644
--- a/torchtitan/models/deepseek_v3/model.py
+++ b/torchtitan/models/deepseek_v3/model.py
@@ -15,6 +15,7 @@
     AttentionMasksType,
     BaseAttention,
     FlexAttention,
+    _resolve_rope,
 )
 from torchtitan.models.common.decoder import TransformerBlock
 from torchtitan.models.common.linear import Linear
@@ -22,7 +23,7 @@
 from torchtitan.models.common.rope import RoPE
 from torchtitan.models.deepseek_v3.mtp import MTPDecoder
 from torchtitan.models.utils import get_moe_model_nparams_and_flops
-from torchtitan.protocols.module import Module
+from torchtitan.protocols.module import Module, ModuleDict
 
 
 class Attention(BaseAttention):
@@ -32,6 +33,8 @@ class Attention(BaseAttention):
     This is DeepSeek V3-specific and NOT shared with other models.
     """
 
+    rope: RoPE
+
     @dataclass(kw_only=True, slots=True)
     class Config(BaseAttention.Config):
         n_heads: int
@@ -53,7 +56,7 @@ class Config(BaseAttention.Config):
         inner_attention: Module.Config = field(default_factory=FlexAttention.Config)
         mscale: float = 1.0
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
         self.dim = config.dim
         self.n_heads = config.n_heads
@@ -88,7 +91,8 @@ def __init__(self, config: Config):
             self.softmax_scale = self.softmax_scale * mscale * mscale
 
         self.inner_attention = config.inner_attention.build()
-        self.rope = config.rope.build()
+        # Keep the canonical module registered only under Decoder.rope_modules.
+        object.__setattr__(self, "rope", _resolve_rope(config.rope, rope_modules))
 
     def forward(
         self,
@@ -160,9 +164,9 @@ class DeepSeekV3TransformerBlock(TransformerBlock):
     class Config(TransformerBlock.Config):
         pass
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
-        self.attention = config.attention.build()
+        self.attention = config.attention.build(rope_modules=rope_modules)
         self.attention_norm = config.attention_norm.build()
         self.ffn_norm = config.ffn_norm.build()
 
diff --git a/torchtitan/models/deepseek_v3/mtp.py b/torchtitan/models/deepseek_v3/mtp.py
index d9785f7f4f..ea418d42c2 100644
--- a/torchtitan/models/deepseek_v3/mtp.py
+++ b/torchtitan/models/deepseek_v3/mtp.py
@@ -23,10 +23,14 @@
 from torchtitan.config import CompileConfig
 from torchtitan.distributed.fsdp import apply_fsdp_to_decoder
 from torchtitan.models.common.attention import AttentionMasksType
-from torchtitan.models.common.decoder import Decoder, TransformerBlock
+from torchtitan.models.common.decoder import (
+    Decoder,
+    TransformerBlock,
+    _register_rope_modules,
+)
 from torchtitan.models.common.linear import Linear
 from torchtitan.models.common.nn_modules import RMSNorm
-from torchtitan.protocols.module import ModuleList
+from torchtitan.protocols.module import ModuleDict, ModuleList
 
 
 def roll_mtp_sequence(
@@ -120,9 +124,9 @@ class Config(TransformerBlock.Config):
         eh_proj: Linear.Config
         mtp_norm: RMSNorm.Config
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
-        self.attention = config.attention.build()
+        self.attention = config.attention.build(rope_modules=rope_modules)
         self.attention_norm = config.attention_norm.build()
         self.ffn_norm = config.ffn_norm.build()
         self.enorm = config.enorm.build()
@@ -211,15 +215,17 @@ def __init__(self, config: Config):
             self.mtp_layers = None
             return
 
+        _register_rope_modules(config.mtp_layers, self.rope_modules)
         self.mtp_layers = ModuleList()
-        with self._rope_cache_context():
-            for layer_config in config.mtp_layers:
-                if not isinstance(layer_config, MTPTransformerBlock.Config):
-                    raise ValueError(
-                        "MTPDecoder requires Config.mtp_layers to contain "
-                        "MTPTransformerBlock.Config instances."
-                    )
-                self.mtp_layers.append(layer_config.build())
+        for layer_config in config.mtp_layers:
+            if not isinstance(layer_config, MTPTransformerBlock.Config):
+                raise ValueError(
+                    "MTPDecoder requires Config.mtp_layers to contain "
+                    "MTPTransformerBlock.Config instances."
+                )
+            self.mtp_layers.append(
+                layer_config.build(rope_modules=self.rope_modules)
+            )
 
     def forward(
         self,
diff --git a/torchtitan/models/gpt_oss/model.py b/torchtitan/models/gpt_oss/model.py
index 49cf960312..3b9bfaa607 100644
--- a/torchtitan/models/gpt_oss/model.py
+++ b/torchtitan/models/gpt_oss/model.py
@@ -22,13 +22,14 @@
     get_causal_mask_mod,
     get_efficient_causal_mask_mod_for_packed_document,
     get_sliding_window_mask_mod,
+    _resolve_rope,
     VarlenAttention,
 )
 from torchtitan.models.common.decoder import Decoder, TransformerBlock
 from torchtitan.models.common.linear import Linear
 from torchtitan.models.common.rope import RoPE
 from torchtitan.models.utils import get_moe_model_nparams_and_flops
-from torchtitan.protocols.module import Module
+from torchtitan.protocols.module import Module, ModuleDict
 
 
 def apply_attention_sink_rescale(
@@ -45,6 +46,8 @@ class Attention(BaseAttention):
     Multi-head attention (MLA) module with sink attention.
     """
 
+    rope: RoPE
+
     @dataclass(kw_only=True, slots=True)
     class Config(BaseAttention.Config):
         n_heads: int = 64
@@ -60,7 +63,7 @@ class Config(BaseAttention.Config):
         """Per-layer causal sliding-window size"""
         rope: RoPE.Config
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
         self.head_dim = config.head_dim
         self.n_heads = config.n_heads
@@ -81,7 +84,8 @@ def __init__(self, config: Config):
         self.wo = config.wo.build()
         self.sinks = nn.Parameter(torch.empty(config.n_heads))
         self.inner_attention = config.inner_attention.build()
-        self.rope = config.rope.build()
+        # Keep the canonical module registered only under Decoder.rope_modules.
+        object.__setattr__(self, "rope", _resolve_rope(config.rope, rope_modules))
 
     def forward(
         self,
@@ -137,7 +141,7 @@ class GptOssTransformerBlock(TransformerBlock):
     class Config(TransformerBlock.Config):
         pass
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
         assert isinstance(config.attention, Attention.Config)
         self.attn_mask_key = (
@@ -145,7 +149,7 @@ def __init__(self, config: Config):
             if config.attention.sliding_window_size is not None
             else "basic_mask"
         )
-        self.attention = config.attention.build()
+        self.attention = config.attention.build(rope_modules=rope_modules)
         self.attention_norm = config.attention_norm.build()
         self.ffn_norm = config.ffn_norm.build()
 
diff --git a/torchtitan/models/llama3/model.py b/torchtitan/models/llama3/model.py
index ee5c074c80..cf78795343 100644
--- a/torchtitan/models/llama3/model.py
+++ b/torchtitan/models/llama3/model.py
@@ -14,6 +14,7 @@
 from torchtitan.models.common.attention import AttentionMasksType
 from torchtitan.models.common.decoder import Decoder, TransformerBlock
 from torchtitan.models.utils import get_dense_model_nparams_and_flops
+from torchtitan.protocols.module import ModuleDict
 
 
 class Llama3TransformerBlock(TransformerBlock):
@@ -31,9 +32,9 @@ class Llama3TransformerBlock(TransformerBlock):
     class Config(TransformerBlock.Config):
         pass
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
-        self.attention = config.attention.build()
+        self.attention = config.attention.build(rope_modules=rope_modules)
         assert config.feed_forward is not None
         self.feed_forward = config.feed_forward.build()
         self.attention_norm = config.attention_norm.build()
diff --git a/torchtitan/models/muse_glimmer/model.py b/torchtitan/models/muse_glimmer/model.py
index ed05e844b9..57a2524309 100644
--- a/torchtitan/models/muse_glimmer/model.py
+++ b/torchtitan/models/muse_glimmer/model.py
@@ -34,7 +34,7 @@
 )
 from torchtitan.models.common.nn_modules import RMSNorm
 from torchtitan.models.utils import get_dense_model_nparams_and_flops
-from torchtitan.protocols.module import Module
+from torchtitan.protocols.module import Module, ModuleDict
 
 from .vision_encoder import MuseGlimmerVisionAdapter, MuseGlimmerVisionEncoder
 
@@ -93,8 +93,8 @@ def sliding_window_size(self) -> int | None:
             # field name without renaming the flex-path usages).
             return self.window_size
 
-    def __init__(self, config: Config):
-        super().__init__(config)
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
+        super().__init__(config, rope_modules=rope_modules)
         self.use_rope: bool = config.use_rope
         self.scale_query_by: float = config.scale_query_by
         self.window_size: int | None = config.window_size
@@ -158,9 +158,9 @@ class Config(TransformerBlock.Config):
         post_attention_norm: RMSNorm.Config
         post_ffn_norm: RMSNorm.Config
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
-        self.attention = config.attention.build()
+        self.attention = config.attention.build(rope_modules=rope_modules)
         assert config.feed_forward is not None
         self.feed_forward = config.feed_forward.build()
         self.attention_norm = config.attention_norm.build()
diff --git a/torchtitan/models/qwen3/model.py b/torchtitan/models/qwen3/model.py
index 5516ecfc43..26916936e3 100644
--- a/torchtitan/models/qwen3/model.py
+++ b/torchtitan/models/qwen3/model.py
@@ -18,6 +18,7 @@
 )
 from torchtitan.models.common.decoder import Decoder, TransformerBlock
 from torchtitan.models.utils import get_moe_model_nparams_and_flops
+from torchtitan.protocols.module import ModuleDict
 
 
 class Qwen3TransformerBlock(TransformerBlock):
@@ -35,10 +36,10 @@ class Qwen3TransformerBlock(TransformerBlock):
     class Config(TransformerBlock.Config):
         pass
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
 
-        self.attention = config.attention.build()
+        self.attention = config.attention.build(rope_modules=rope_modules)
 
         self.moe_enabled = config.moe is not None
         if self.moe_enabled:
diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py
index 19a4c9363d..926d697c3a 100644
--- a/torchtitan/models/qwen3_5/model.py
+++ b/torchtitan/models/qwen3_5/model.py
@@ -20,6 +20,7 @@
     create_varlen_metadata_for_document,
     VarlenAttention,
     VarlenMetadata,
+    _resolve_rope,
 )
 from torchtitan.models.common.decoder import Decoder
 from torchtitan.models.common.multimodal import (
@@ -28,7 +29,7 @@
     scatter_vision_embeds,
 )
 from torchtitan.models.utils import get_moe_model_nparams_and_flops
-from torchtitan.protocols.module import Module
+from torchtitan.protocols.module import Module, ModuleDict
 
 from .gdn import GatedDeltaNet
 from .rope import MRoPE
@@ -77,6 +78,8 @@ class Qwen35Attention(BaseAttention):
     gated ``wq`` doesn't fit a fused QKV projection that TP-shards by head.
     """
 
+    rope: MRoPE
+
     @dataclass(kw_only=True, slots=True)
     class Config(BaseAttention.Config):
         n_heads: int
@@ -92,7 +95,7 @@ class Config(BaseAttention.Config):
         k_norm: OffsetRMSNorm.Config
         inner_attention: Module.Config
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
         self.n_heads = config.n_heads
         self.n_kv_heads = config.n_kv_heads
@@ -105,7 +108,8 @@ def __init__(self, config: Config):
         self.wv = config.wv.build()
         self.wo = config.wo.build()
 
-        self.rope = config.rope.build()
+        # Keep the canonical module registered only under Decoder.rope_modules.
+        object.__setattr__(self, "rope", _resolve_rope(config.rope, rope_modules))
 
         self.q_norm = config.q_norm.build()
         self.k_norm = config.k_norm.build()
@@ -178,13 +182,15 @@ class Config(Module.Config):
         attention_norm: OffsetRMSNorm.Config
         ffn_norm: OffsetRMSNorm.Config
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
         self.full_attn = config.attention is not None
         self.attn_mask_key = "quadratic_attention" if self.full_attn else "deltanet"
 
         if self.full_attn:
-            self.attn = config.attention.build()  # pyrefly: ignore [missing-attribute]
+            self.attn = config.attention.build(  # pyrefly: ignore [missing-attribute]
+                rope_modules=rope_modules
+            )
         else:
             assert config.delta_net is not None
             self.attn = config.delta_net.build()
diff --git a/torchtitan/overrides/fused_mla.py b/torchtitan/overrides/fused_mla.py
index 3b7e96179d..03ac12f6b9 100644
--- a/torchtitan/overrides/fused_mla.py
+++ b/torchtitan/overrides/fused_mla.py
@@ -71,6 +71,7 @@
 from torchtitan.config import derive, override
 from torchtitan.distributed.utils import get_spmd_backend
 from torchtitan.models.common.attention import AttentionMasksType
+from torchtitan.protocols.module import ModuleDict
 from torchtitan.models.common.rope import _maybe_check_max_pos, ComplexRoPE
 from torchtitan.models.deepseek_v3.model import Attention
 
@@ -816,8 +817,8 @@ class FusedMLAAttention(Attention):
     class Config(Attention.Config):
         pass
 
-    def __init__(self, config: Config):
-        super().__init__(config)
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
+        super().__init__(config, rope_modules=rope_modules)
         if not isinstance(self.rope, ComplexRoPE):
             raise TypeError(
                 "FusedMLAAttention currently requires ComplexRoPE, got "

From 2cb0c19dbd2f1fc41366254214d9ee273ed688eb Mon Sep 17 00:00:00 2001
From: mystri <hanboyou@huawei.com>
Date: Wed, 2 Sep 2026 10:24:09 +0800
Subject: [PATCH 5/8] Inline RoPE lookup and key formatting helpers

---
 torchtitan/models/common/attention.py  |  8 ++-----
 torchtitan/models/common/decoder.py    | 28 +++++++++--------------
 torchtitan/models/common/rope.py       | 31 +++++++++++++-------------
 torchtitan/models/deepseek_v3/model.py |  3 +--
 torchtitan/models/deepseek_v3/mtp.py   | 14 ++++++------
 torchtitan/models/gpt_oss/model.py     |  3 +--
 torchtitan/models/qwen3_5/model.py     |  3 +--
 7 files changed, 37 insertions(+), 53 deletions(-)

diff --git a/torchtitan/models/common/attention.py b/torchtitan/models/common/attention.py
index e76f6c5131..ed7c79cb21 100644
--- a/torchtitan/models/common/attention.py
+++ b/torchtitan/models/common/attention.py
@@ -13,7 +13,7 @@
 
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field
-from typing import Any, ClassVar, NamedTuple, cast
+from typing import Any, ClassVar, NamedTuple
 
 import spmd_types as spmd
 import torch
@@ -667,10 +667,6 @@ def __post_init__(self):
             assert self.n_heads > 0, "n_heads must be > 0"
 
 
-def _resolve_rope(rope_config: RoPE.Config, rope_modules: ModuleDict) -> RoPE:
-    return cast(RoPE, rope_modules[rope_config.rope_key()])
-
-
 class BaseQKVLinear(Module):
     """Base class for Q/K/V projection strategies.
 
@@ -938,7 +934,7 @@ def __init__(self, config: Config, *, rope_modules: ModuleDict):
         )
         self.enable_gqa = self.n_heads > self.n_kv_heads
         # Keep the canonical module registered only under Decoder.rope_modules.
-        object.__setattr__(self, "rope", _resolve_rope(config.rope, rope_modules))
+        object.__setattr__(self, "rope", rope_modules[config.rope.rope_key()])
 
         # Pluggable QKV projection
         self.qkv_linear = config.qkv_linear.build()
diff --git a/torchtitan/models/common/decoder.py b/torchtitan/models/common/decoder.py
index efe3b6316a..9ad0eda043 100644
--- a/torchtitan/models/common/decoder.py
+++ b/torchtitan/models/common/decoder.py
@@ -35,22 +35,6 @@
 __all__ = ["Decoder", "TransformerBlock"]
 
 
-def _rope_config(layer_config):
-    attention_config = getattr(layer_config, "attention", None)
-    return getattr(attention_config, "rope", None)
-
-
-def _register_rope_modules(layer_configs, rope_modules: ModuleDict) -> None:
-    """Build one canonical RoPE module for each layer configuration key."""
-    for layer_config in layer_configs:
-        rope_config = _rope_config(layer_config)
-        if rope_config is None:
-            continue
-        key = rope_config.rope_key()
-        if key not in rope_modules:
-            rope_modules[key] = rope_config.build()
-
-
 # TODO: we can unify the TransformerBlock impl across all models when
 # there is no special logic for each model, including
 # ffn vs. moe naming and creation, etc.
@@ -267,11 +251,19 @@ def __init__(self, config: Config):
         self.tok_embeddings = config.tok_embeddings.build()
 
         self.rope_modules = ModuleDict()
-        _register_rope_modules(config.layers, self.rope_modules)
+        for layer_config in config.layers:
+            attention_config = getattr(layer_config, "attention", None)
+            rope_config = getattr(attention_config, "rope", None)
+            if rope_config is None:
+                continue
+            key = rope_config.rope_key()
+            if key not in self.rope_modules:
+                self.rope_modules[key] = rope_config.build()
 
         self.layers = ModuleDict()
         for i, layer_config in enumerate(config.layers):
-            rope_config = _rope_config(layer_config)
+            attention_config = getattr(layer_config, "attention", None)
+            rope_config = getattr(attention_config, "rope", None)
             if rope_config is None:
                 layer = layer_config.build()
             else:
diff --git a/torchtitan/models/common/rope.py b/torchtitan/models/common/rope.py
index 685313b3eb..5ecdafaa8e 100644
--- a/torchtitan/models/common/rope.py
+++ b/torchtitan/models/common/rope.py
@@ -22,21 +22,6 @@
 ]
 
 
-def _format_rope_key_value(value: object) -> str:
-    """Format config values into safe, readable ``ModuleDict`` key fragments."""
-    if value is None:
-        text = "none"
-    elif isinstance(value, bool):
-        text = "true" if value else "false"
-    elif isinstance(value, float):
-        text = format(value, ".17g")
-    elif isinstance(value, (list, tuple)):
-        text = "-".join(_format_rope_key_value(item) for item in value)
-    else:
-        text = str(value)
-    return re.sub(r"[^A-Za-z0-9_-]", "p", text)
-
-
 # pyrefly: ignore [not-callable]
 @spmd.no_typecheck()
 def _maybe_check_max_pos(positions: torch.Tensor, *, max_valid_pos: int) -> None:
@@ -127,6 +112,20 @@ class Config(Module.Config):
 
         def rope_key(self) -> str:
             """Return the stable, descriptive key for this RoPE implementation."""
+
+            def format_value(value: object) -> str:
+                if value is None:
+                    text = "none"
+                elif isinstance(value, bool):
+                    text = "true" if value else "false"
+                elif isinstance(value, float):
+                    text = format(value, ".17g")
+                elif isinstance(value, (list, tuple)):
+                    text = "-".join(format_value(item) for item in value)
+                else:
+                    text = str(value)
+                return re.sub(r"[^A-Za-z0-9_-]", "p", text)
+
             owner = self._owner
             owner_name = owner.__name__ if owner is not None else type(self).__name__
             parts = [owner_name]
@@ -179,7 +178,7 @@ def rope_key(self) -> str:
                     }
                 ):
                     continue
-                value = _format_rope_key_value(getattr(self, config_field.name))
+                value = format_value(getattr(self, config_field.name))
                 name = key_names.get(config_field.name, config_field.name)
                 separator = "" if config_field.name in common_fields - {"scaling"} else "_"
                 parts.append(f"{name}{separator}{value}")
diff --git a/torchtitan/models/deepseek_v3/model.py b/torchtitan/models/deepseek_v3/model.py
index 88de5629fa..ad61e8ed35 100644
--- a/torchtitan/models/deepseek_v3/model.py
+++ b/torchtitan/models/deepseek_v3/model.py
@@ -15,7 +15,6 @@
     AttentionMasksType,
     BaseAttention,
     FlexAttention,
-    _resolve_rope,
 )
 from torchtitan.models.common.decoder import TransformerBlock
 from torchtitan.models.common.linear import Linear
@@ -92,7 +91,7 @@ def __init__(self, config: Config, *, rope_modules: ModuleDict):
 
         self.inner_attention = config.inner_attention.build()
         # Keep the canonical module registered only under Decoder.rope_modules.
-        object.__setattr__(self, "rope", _resolve_rope(config.rope, rope_modules))
+        object.__setattr__(self, "rope", rope_modules[config.rope.rope_key()])
 
     def forward(
         self,
diff --git a/torchtitan/models/deepseek_v3/mtp.py b/torchtitan/models/deepseek_v3/mtp.py
index ea418d42c2..edab609063 100644
--- a/torchtitan/models/deepseek_v3/mtp.py
+++ b/torchtitan/models/deepseek_v3/mtp.py
@@ -23,11 +23,7 @@
 from torchtitan.config import CompileConfig
 from torchtitan.distributed.fsdp import apply_fsdp_to_decoder
 from torchtitan.models.common.attention import AttentionMasksType
-from torchtitan.models.common.decoder import (
-    Decoder,
-    TransformerBlock,
-    _register_rope_modules,
-)
+from torchtitan.models.common.decoder import Decoder, TransformerBlock
 from torchtitan.models.common.linear import Linear
 from torchtitan.models.common.nn_modules import RMSNorm
 from torchtitan.protocols.module import ModuleDict, ModuleList
@@ -215,14 +211,18 @@ def __init__(self, config: Config):
             self.mtp_layers = None
             return
 
-        _register_rope_modules(config.mtp_layers, self.rope_modules)
-        self.mtp_layers = ModuleList()
         for layer_config in config.mtp_layers:
             if not isinstance(layer_config, MTPTransformerBlock.Config):
                 raise ValueError(
                     "MTPDecoder requires Config.mtp_layers to contain "
                     "MTPTransformerBlock.Config instances."
                 )
+            rope_config = getattr(layer_config.attention, "rope")
+            key = rope_config.rope_key()
+            if key not in self.rope_modules:
+                self.rope_modules[key] = rope_config.build()
+        self.mtp_layers = ModuleList()
+        for layer_config in config.mtp_layers:
             self.mtp_layers.append(
                 layer_config.build(rope_modules=self.rope_modules)
             )
diff --git a/torchtitan/models/gpt_oss/model.py b/torchtitan/models/gpt_oss/model.py
index 3b9bfaa607..089a1bb22d 100644
--- a/torchtitan/models/gpt_oss/model.py
+++ b/torchtitan/models/gpt_oss/model.py
@@ -22,7 +22,6 @@
     get_causal_mask_mod,
     get_efficient_causal_mask_mod_for_packed_document,
     get_sliding_window_mask_mod,
-    _resolve_rope,
     VarlenAttention,
 )
 from torchtitan.models.common.decoder import Decoder, TransformerBlock
@@ -85,7 +84,7 @@ def __init__(self, config: Config, *, rope_modules: ModuleDict):
         self.sinks = nn.Parameter(torch.empty(config.n_heads))
         self.inner_attention = config.inner_attention.build()
         # Keep the canonical module registered only under Decoder.rope_modules.
-        object.__setattr__(self, "rope", _resolve_rope(config.rope, rope_modules))
+        object.__setattr__(self, "rope", rope_modules[config.rope.rope_key()])
 
     def forward(
         self,
diff --git a/torchtitan/models/qwen3_5/model.py b/torchtitan/models/qwen3_5/model.py
index 926d697c3a..534cfe5b9a 100644
--- a/torchtitan/models/qwen3_5/model.py
+++ b/torchtitan/models/qwen3_5/model.py
@@ -20,7 +20,6 @@
     create_varlen_metadata_for_document,
     VarlenAttention,
     VarlenMetadata,
-    _resolve_rope,
 )
 from torchtitan.models.common.decoder import Decoder
 from torchtitan.models.common.multimodal import (
@@ -109,7 +108,7 @@ def __init__(self, config: Config, *, rope_modules: ModuleDict):
         self.wo = config.wo.build()
 
         # Keep the canonical module registered only under Decoder.rope_modules.
-        object.__setattr__(self, "rope", _resolve_rope(config.rope, rope_modules))
+        object.__setattr__(self, "rope", rope_modules[config.rope.rope_key()])
 
         self.q_norm = config.q_norm.build()
         self.k_norm = config.k_norm.build()

From 8c035df320319b966d29da3d2840c7c6678ba0e8 Mon Sep 17 00:00:00 2001
From: mystri <hanboyou@huawei.com>
Date: Wed, 2 Sep 2026 11:31:30 +0800
Subject: [PATCH 6/8] Keep RoPE registry in pipeline stages

---
 torchtitan/distributed/pipeline_parallel.py | 8 ++++++--
 1 file changed, 6 insertions(+), 2 deletions(-)

diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py
index 9a4662417e..b80d23f8e1 100644
--- a/torchtitan/distributed/pipeline_parallel.py
+++ b/torchtitan/distributed/pipeline_parallel.py
@@ -488,9 +488,13 @@ def _split_module(
     model = copy.deepcopy(whole_model)
     # Create a set of modules to keep for faster lookup
     modules_to_keep = set(module_names)
-    if "rope_modules" in model._modules:
-        modules_to_keep.add("rope_modules")
     for module_name, module_value in model.named_children():
+        # Keep the shared RoPE registry in every PP model part. Its entries are
+        # keyed by RoPE configuration rather than by layer FQN, so it cannot be
+        # selected by the normal layer-pruning logic. The ownership/replication
+        # policy for shared root modules should be revisited separately.
+        if module_name == "rope_modules":
+            continue
         if module_name in modules_to_keep:
             continue
         # Handle layer-like structures (e.g., "layers.0", "layers.1")

From 3c65120b7fdf1198093b533942dd3ffc7f26ea9e Mon Sep 17 00:00:00 2001
From: mystri <hanboyou@huawei.com>
Date: Wed, 2 Sep 2026 11:34:40 +0800
Subject: [PATCH 7/8] Simplify pipeline module pruning

---
 torchtitan/distributed/pipeline_parallel.py | 2 --
 1 file changed, 2 deletions(-)

diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py
index b80d23f8e1..34718b7af2 100644
--- a/torchtitan/distributed/pipeline_parallel.py
+++ b/torchtitan/distributed/pipeline_parallel.py
@@ -495,8 +495,6 @@ def _split_module(
         # policy for shared root modules should be revisited separately.
         if module_name == "rope_modules":
             continue
-        if module_name in modules_to_keep:
-            continue
         # Handle layer-like structures (e.g., "layers.0", "layers.1")
         if isinstance(
             module_value, (nn.ModuleDict, nn.ModuleList, ModuleDict, ModuleList)

From 54cdedaaa9460f5306dd930ad58b75f8a6f141e8 Mon Sep 17 00:00:00 2001
From: mystri <hanboyou@huawei.com>
Date: Wed, 2 Sep 2026 11:48:32 +0800
Subject: [PATCH 8/8] Document pipeline RoPE ownership discussion

---
 docs/rope-cache-updated-plan.md | 383 ++++++++++++++++++++++++++++++++
 1 file changed, 383 insertions(+)
 create mode 100644 docs/rope-cache-updated-plan.md

diff --git a/docs/rope-cache-updated-plan.md b/docs/rope-cache-updated-plan.md
new file mode 100644
index 0000000000..9903544f28
--- /dev/null
+++ b/docs/rope-cache-updated-plan.md
@@ -0,0 +1,383 @@
+# Updated RoPE sharing plan for PR #4376
+
+## Context
+
+This plan follows the review discussion on
+[PR #4376](https://github.com/pytorch/torchtitan/pull/4376), especially the
+proposal to make Decoder-owned RoPE modules the explicit dependencies of the
+attention layers:
+
+- [latest maintainer proposal](https://github.com/pytorch/torchtitan/pull/4376#issuecomment-5465560593)
+- [analysis of the single-module design](https://github.com/pytorch/torchtitan/pull/4376#issuecomment-5465562358)
+- [registry-scoping review](https://github.com/pytorch/torchtitan/pull/4376#discussion_r3887739753)
+- [duplicate-compute review](https://github.com/pytorch/torchtitan/pull/4376#discussion_r3887741233)
+- [Helion review](https://github.com/pytorch/torchtitan/pull/4376#discussion_r3887743886)
+- [MTP context review](https://github.com/pytorch/torchtitan/pull/4376#discussion_r3887745493)
+
+The earlier attempt in
+[PR #4111](https://github.com/pytorch/torchtitan/pull/4111) cached the built
+module on a shared config object. We should not adopt that mechanism because it
+makes `Config.build()` stateful and weakens the rule that a config build returns
+a fresh owner. The updated design makes sharing explicit at the Decoder boundary
+instead.
+
+The newest reply identifies an important limitation in the single-module
+version: newer architectures can intentionally use more than one effective RoPE
+configuration. For example, DeepSeek-V4 selects `compress_rope_theta` for
+compressed layers and `rope_theta` for pure sliding-window layers in its
+`freqs_cis` construction ([reference implementation](https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/model.py#L439-L445)).
+The design therefore needs a model-local collection of canonical RoPE modules,
+not a homogeneity assertion.
+
+This remains a proposed direction, not a settled review outcome. The PR author
+has [asked the reviewer to confirm the revised
+design](https://github.com/pytorch/torchtitan/pull/4376#issuecomment-5465642070),
+and the latest reply proposes a collection of Decoder-owned RoPE modules
+([comment](https://github.com/pytorch/torchtitan/pull/4376#issuecomment-5488760743)).
+Implementation should keep the ownership and selection rules isolated and easy
+to review before the distributed-path changes are stacked on top.
+
+The current branch is also behind `origin/main`; rebase it before implementation
+and re-check all touched call sites against the rebased tree.
+
+## Decision
+
+Replace the registry/reader implementation with a Decoder-owned `ModuleDict` of
+canonical RoPE modules. The Decoder registers each distinct RoPE config under a
+descriptive string key. Pass the collection explicitly through the existing
+layer construction loops; each attention builder derives the key from its own
+RoPE config and fetches the matching canonical module.
+
+```text
+Decoder
+|-- rope_modules["ComplexRoPE_d128_ctx1048576_theta10000_scaling_none"]
+|                               registered canonical RoPE owner
+|-- rope_modules["CosSinRoPE_d128_ctx1048576_theta10000_scaling_none"]
+|                               registered canonical RoPE owner
+`-- layers
+    `-- attention
+        `-- rope -------------- non-registered reference to one owner
+```
+
+Each PP stage owns its own copy of the complete `Decoder.rope_modules` collection.
+Duplicate caches across multiple virtual stages on one rank are accepted. We do
+not attempt cross-stage or cross-model sharing.
+
+The collection is deduplicated by exact effective RoPE configuration. A model
+may therefore have zero, one, or several canonical RoPE modules. A layer must
+never silently receive a cache for a different configuration.
+
+`rope_modules` is the preferred name: it describes the registered objects and
+does not imply that the Decoder owns raw cache tensors or a global registry. Use
+`ModuleDict` because the key is now the explicit routing contract. The old
+generated numeric key style (`rope_0`, `rope_1`, ...) is removed. Keys are
+descriptive, deterministic, and contain only characters accepted by PyTorch
+module names; they are internal module names, not config or checkpoint APIs.
+
+## Why this supersedes the current branch
+
+The current registry solves steady-state storage duplication, but introduces
+four costs that the updated design removes:
+
+1. `contextvars` implicitly select a model registry and require MTP to re-enter
+   a context manually.
+2. Every layer still computes a full candidate cache before duplicates are
+   discarded, and every layer recomputes it during `init_states()`.
+3. `RoPECacheReader` creates a strong path from each layer back through the
+   registry to the Decoder.
+4. Moving the real buffer to `_rope_cache_N` on Decoder leaves the existing
+   `state_shardings={"cache": ...}` declarations attached to RoPE configs but
+   disconnected from the module that owns the actual buffer.
+
+With canonical modules, each module identity survives `.to()`, `to_empty()`,
+`init_states()`, and buffer replacement. Only that module's `cache` changes;
+every attention keeps resolving the selected cache through its stable module
+reference. Therefore the reader, registry, context, and cache property are
+unnecessary. The model-local collection replaces the old process of selecting a
+cache slot through an implicit context.
+
+## Detailed implementation plan
+
+### 1. Restore RoPE to a normal module
+
+Revert the cache implementation to the ordinary upstream RoPE module, then
+add only the shared naming helper needed by the Decoder-owned collection:
+
+- Delete `RoPECacheReader`, `_RoPECacheRegistry`, the context variable/context
+  manager, and `register_rope_cache()`.
+- Replace the old cache-slot `_cache_key()` with deterministic
+  `RoPE.Config.rope_key()` formatting shared by Decoder registration and layer
+  lookup.
+- Keep the direct non-persistent `cache` buffer and ordinary
+  `_init_self_buffers()` behavior.
+
+### 2. Add explicit runtime dependency injection
+
+Use the existing `Config.build(**kwargs)` support for runtime objects. Do not
+add module objects or a runtime collection to config dataclasses.
+
+The Decoder creates one registered `rope_modules` collection before its existing
+layer loop, then passes that same collection to each RoPE-bearing layer build:
+
+```python
+layer_config.build(rope_modules=self.rope_modules)
+```
+
+The `rope_modules` argument is construction-only. Each transformer-block
+builder forwards it to its attention build when that layer has a RoPE config.
+The attention constructor derives a key from its own existing `config.rope` and
+fetches the canonical module:
+
+```python
+rope_key = config.rope.rope_key()
+self.rope = rope_modules[rope_key]
+```
+
+The Decoder retains explicit `rope_config is None` handling for heterogeneous
+architectures: layers without full attention (for example, Qwen 3.5 GDN
+layers) use their normal no-keyword build path, while RoPE-bearing layers must
+be built with the Decoder-owned collection. This prevents a RoPE-backed layer
+from silently constructing a private duplicate cache. Direct
+`RoPE.Config.build()` remains available for the RoPE primitive itself, but
+attention construction resolves RoPE through the supplied collection.
+
+This preserves each layer's existing config as the source of truth while keeping
+module selection explicit and model-scoped. The key is derived at construction
+time and is not stored in model config or threaded through forward calls.
+
+Update the Decoder transformer-block constructors that can contain RoPE-backed
+attention so they accept the required keyword-only
+`rope_modules: ModuleDict` and forward it to their attention build:
+
+- Llama 3
+- Qwen 3
+- Muse Glimmer
+- DeepSeek V3
+- DeepSeek V3 MTP
+- GPT-OSS
+- Qwen 3.5
+
+Models with no RoPE module, such as Kimi K3, continue using the no-keyword build
+path and require no synthetic dependency.
+
+Update the RoPE-backed attention constructors to accept the same keyword-only
+argument:
+
+- common `GQAttention`
+- DeepSeek V3 `Attention`
+- GPT-OSS `Attention`
+- Qwen 3.5 `Qwen35Attention`
+- `FusedMLAAttention`, which must forward the argument to DeepSeek attention
+
+Muse Glimmer inherits the common GQA constructor. Helion overrides replace the
+RoPE config/module type and therefore require no separate cache-sharing path.
+
+This explicit path avoids post-build injection. Post-build injection would still
+construct and discard one RoPE cache per layer, leaving the duplicate-compute
+review unresolved.
+
+### 3. Keep Decoder as the sole module owner
+
+Build `self.rope_modules = ModuleDict()` before building `self.layers`, after
+runtime config updates and overrides have been applied. Register the distinct
+RoPE configs needed by the main layers, then keep the existing layer
+construction loop and pass the same collection into every RoPE-bearing layer:
+
+```python
+for layer_config in config.layers:
+    attention_config = getattr(layer_config, "attention", None)
+    rope_config = getattr(attention_config, "rope", None)
+    if rope_config is None:
+        continue
+    rope_key = rope_config.rope_key()
+    if rope_key not in self.rope_modules:
+        self.rope_modules[rope_key] = rope_config.build()
+
+for i, layer_config in enumerate(config.layers):
+    attention_config = getattr(layer_config, "attention", None)
+    rope_config = getattr(attention_config, "rope", None)
+    if rope_config is None:
+        layer = layer_config.build()
+    else:
+        layer = layer_config.build(rope_modules=self.rope_modules)
+    self.layers[str(i)] = layer
+```
+
+The registration and layer-build decisions are intentionally inlined in the
+Decoder and MTPDecoder construction paths. The Decoder owns key formatting and
+module insertion:
+
+```python
+for layer_config in config.layers:
+    attention_config = getattr(layer_config, "attention", None)
+    rope_config = getattr(attention_config, "rope", None)
+    if rope_config is None:
+        continue
+    rope_key = rope_config.rope_key()
+    if rope_key not in self.rope_modules:
+        self.rope_modules[rope_key] = rope_config.build()
+```
+
+The collection has no hand-generated `rope_0`, `rope_1`, ... keys. The
+descriptive key is the module registration name and the routing contract. The
+layer's existing RoPE config remains the source of truth for deriving it.
+
+Variant identity has three separate concepts:
+
+| Concern | Representation | Contract |
+| --- | --- | --- |
+| Registered owner | `self.rope_modules[rope_key]` | Decoder-created canonical module |
+| Deduplication | Equal `rope_config.rope_key()` values | Key covers the effective module contract |
+| Layer routing | `rope_modules[config.rope.rope_key()]` | No key is stored in model config |
+
+Do not derive the key from tensor metadata or object identity. Use deterministic
+field formatting from the RoPE config, with safe characters for PyTorch module
+names. The key should include the concrete implementation class, cache alignment
+(`ComplexRoPE` versus `CosSinRoPE`), dimension, context length, theta, scaling
+mode, and any active subclass/scaling fields.
+
+Each RoPE-bearing attention constructor accepts the required `rope_modules`
+collection and directly obtains its already-registered module with
+`rope_modules[config.rope.rope_key()]`. The lookup never builds a RoPE. A
+missing key is a model-construction bug and fails immediately. The attention
+stores the returned module without registering it again under the attention. Use
+`object.__setattr__(self, "rope", rope)` with one comment explaining that
+bypassing `nn.Module.__setattr__` keeps Decoder as the sole registered owner.
+
+There is no process-global registry, context variable, list wrapper, namespace,
+weak reference, proxy tensor, or special cache property. `ModuleDict` plus
+deterministic config keys is the entire sharing mechanism. The Decoder still
+handles `rope_config is None` for hybrid layers that do not use full attention;
+those layers take their normal no-keyword build path.
+
+### 4. Use the same canonical path for main and MTP layers
+
+Keep the current construction shape for every model layer, including MTP:
+
+- `Decoder.__init__` creates `self.rope_modules`, inlines registration of the
+  distinct main-layer RoPE configs, and passes the collection to each
+  RoPE-bearing main-layer `build()` call;
+- `MTPDecoder.__init__` inlines registration for `config.mtp_layers`, then
+  passes the same collection to each MTP-layer `build()` call;
+- every RoPE-bearing attention derives `config.rope.rope_key()` and performs a
+  strict `rope_modules[key]` lookup;
+- hybrid layers with no full attention, such as Qwen 3.5 GDN layers, do not
+  receive or use a RoPE module;
+- models with no RoPE module at all retain their current construction path.
+
+There is no protected Decoder hook, MTP context, lazy construction in the
+attention builder, or configuration-mismatch validation. Heterogeneous models
+use one canonical module per distinct effective RoPE key, and a missing
+pre-registered module fails at construction. The `None` handling is limited to
+layers whose architecture genuinely has no RoPE config.
+
+### 6. Preserve sharding through the real owner
+
+Each canonical `Decoder.rope_modules[...]` module is built from its nested RoPE config,
+so its existing `sharding_config` must be carried onto the real owner by
+`Module.Config.build()`.
+
+Audit and update the RoPE sharding setup in:
+
+- `torchtitan/models/common/decoder_sharding.py`
+- `torchtitan/models/deepseek_v3/sharding.py`
+- `torchtitan/models/gpt_oss/sharding.py`
+- `torchtitan/models/qwen3_5/sharding.py`
+
+The final contract is one registered `cache` buffer per canonical variant,
+distributed once as Replicate on TP, not one DTensor per layer. Remove stale
+"per-layer cache" comments. Do not add a RoPE-specific `parallelize()` override;
+ordinary recursive `Module.parallelize()` must visit each child of
+`Decoder.rope_modules` once.
+
+The key must include any sharding distinction that changes the module contract,
+or sharding must be normalized before registration. Two modules with the same
+key must be safe to traverse and distribute as one owner.
+
+### 7. Replicate the owner per pipeline stage
+
+Pipeline splitting must preserve the registered `rope_modules` child on every stage,
+including custom `module_fqns_per_model_part` and GraphPP paths.
+
+The current minimal rule is in the common split path: skip the top-level
+`rope_modules` child during pruning so it remains intact in every
+`_split_module()` result. This avoids duplicating the rule across automatic,
+custom, VLM, Muse Glimmer, eager PP, and GraphPP stage-list generation. The
+registry is not added to each stage's FQN list because its children are keyed by
+RoPE configuration, not by layer index.
+The rule is implementation-agnostic and therefore covers stock RoPE,
+HelionCosSinRoPE, and HelionComplexRoPE equally: all are registered children
+under the same `rope_modules` root.
+
+The longer-term question is whether this should become model-owned metadata,
+for example a `Decoder` class attribute such as
+`modules_to_keep_on_all_model_parts = ("rope_modules",)`, consumed generically
+by `_split_module()`. That would let models declare other shared root modules
+without hard-coding their names in pipeline code. This is intentionally left as
+an open reviewer discussion rather than adding a new protocol prematurely.
+
+`copy.deepcopy()` must produce one stage-local copy of the complete RoPE
+collection and preserve every remaining attention's non-registered reference to
+the corresponding copied module. Multiple virtual stages on the same rank
+intentionally receive independent collections.
+
+### 8. Document the structural compatibility boundary
+
+Configuration structure remains unchanged: every attention config still owns
+its own RoPE config copy, so config overrides and checkpoint validation continue
+to use the existing paths.
+
+Runtime module ownership changes deliberately:
+
+- `attention.rope` remains a usable attribute and points to its selected canonical
+  module;
+- the registered module FQNs become `rope_modules.N` on Decoder;
+- `layers.N.attention.rope` no longer appears as a registered child in
+  `named_modules()` or `named_buffers()`;
+- the cache remains non-persistent, so model state-dict keys do not change.
+
+Before implementation is considered complete, audit all in-tree FQN-based
+module replacement, compile, FSDP, and diagnostics paths. If an in-tree consumer
+requires each per-layer RoPE to be a registered child, stop and revisit the
+design rather than adding an alias registration that would reintroduce duplicate
+lifecycle traversal.
+
+## Validation
+
+Tests, distributed runs, and numerical comparisons are intentionally deferred
+until the functional construction and ownership path is settled. Do not add
+test-specific compatibility branches while the object model is still changing.
+
+## Implementation sequence
+
+1. Rebase the draft branch onto current `origin/main` while preserving the local
+   investigation documents.
+2. Add the Decoder-owned `rope_modules` collection and the strict config lookup
+   path to every RoPE-bearing main and MTP layer constructor.
+3. Delete registry/reader/context/key code and restore ordinary RoPE buffers.
+4. Update Helion, fused MLA, and every affected transformer block/attention
+   constructor to use the canonical collection.
+5. Reconnect and verify sharding on each registered canonical RoPE owner.
+6. Make PP/GraphPP preserve `Decoder.rope_modules` on every stage.
+7. Revisit tests and deterministic numerical validation after functionality is
+   stable.
+
+## Non-goals
+
+- No process-global cache or module singleton.
+- No process-global keyed multi-cache registry; the only collection is owned by
+  one Decoder instance.
+- No proxy tensor, reader, or custom cache property.
+- No post-build cache tying or re-alias pass.
+- No RoPE-specific parallelization override.
+- No cache sharing across independent models or PP stages.
+- No factory API for lazy cache creation; canonical modules are registered before
+  their consuming layers are built.
+
+## Exit criteria
+
+The updated functionality is ready for review when it has one obvious
+model-local owner collection, zero duplicate cache construction per effective
+variant within a Decoder, no implicit construction context, no dead sharding
+declarations, explicit routing for heterogeneous RoPE configs, stage-local PP
+ownership, and no private-cache fallback in RoPE-backed attention construction.