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.
diff --git a/torchtitan/distributed/pipeline_parallel.py b/torchtitan/distributed/pipeline_parallel.py
index f547b18cd4..34718b7af2 100644
--- a/torchtitan/distributed/pipeline_parallel.py
+++ b/torchtitan/distributed/pipeline_parallel.py
@@ -489,6 +489,12 @@ def _split_module(
     # Create a set of modules to keep for faster lookup
     modules_to_keep = set(module_names)
     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
         # 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 26d1dbf211..b5f0db324f 100644
--- a/torchtitan/models/common/__init__.py
+++ b/torchtitan/models/common/__init__.py
@@ -36,7 +36,11 @@
     RMSNorm,
     SiLU,
 )
-from .rope import ComplexRoPE, CosSinRoPE, RoPE
+from .rope import (
+    ComplexRoPE,
+    CosSinRoPE,
+    RoPE,
+)
 
 __all__ = [
     "Conv1d",
diff --git a/torchtitan/models/common/attention.py b/torchtitan/models/common/attention.py
index 5951b6b97a..ed7c79cb21 100644
--- a/torchtitan/models/common/attention.py
+++ b/torchtitan/models/common/attention.py
@@ -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
 
 
@@ -892,6 +892,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 +921,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 +933,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", 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 a208df1244..9ad0eda043 100644
--- a/torchtitan/models/common/decoder.py
+++ b/torchtitan/models/common/decoder.py
@@ -9,19 +9,19 @@
 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 (
     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
@@ -248,12 +248,27 @@ def update_from_config(
     def __init__(self, config: Config):
         super().__init__()
         self.config = config
-
         self.tok_embeddings = config.tok_embeddings.build()
 
+        self.rope_modules = ModuleDict()
+        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):
-            self.layers[str(i)] = layer_config.build()
+            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
 
         self.norm = config.norm.build()
         self.lm_head = config.lm_head.build()
diff --git a/torchtitan/models/common/rope.py b/torchtitan/models/common/rope.py
index 88bb0ca2ff..5ecdafaa8e 100644
--- a/torchtitan/models/common/rope.py
+++ b/torchtitan/models/common/rope.py
@@ -5,7 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 
 import math
-from dataclasses import dataclass
+import re
+from dataclasses import dataclass, fields
 from typing import Literal
 
 import spmd_types as spmd
@@ -89,6 +90,8 @@ class RoPE(Module):
     cosine/sine caches.
     """
 
+    cache: torch.Tensor
+
     @dataclass(kw_only=True, slots=True)
     class Config(Module.Config):
         dim: int
@@ -107,6 +110,80 @@ 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."""
+
+            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]
+            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_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
@@ -167,13 +244,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):
-            self.cache = self._precompute_cache()
+            cache = self._precompute_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..ad61e8ed35 100644
--- a/torchtitan/models/deepseek_v3/model.py
+++ b/torchtitan/models/deepseek_v3/model.py
@@ -22,7 +22,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 +32,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 +55,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 +90,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", rope_modules[config.rope.rope_key()])
 
     def forward(
         self,
@@ -160,9 +163,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 90068b1d26..edab609063 100644
--- a/torchtitan/models/deepseek_v3/mtp.py
+++ b/torchtitan/models/deepseek_v3/mtp.py
@@ -26,7 +26,7 @@
 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 ModuleList
+from torchtitan.protocols.module import ModuleDict, ModuleList
 
 
 def roll_mtp_sequence(
@@ -120,9 +120,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,14 +211,21 @@ def __init__(self, config: Config):
             self.mtp_layers = None
             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())
+            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)
+            )
 
     def forward(
         self,
diff --git a/torchtitan/models/gpt_oss/model.py b/torchtitan/models/gpt_oss/model.py
index 49cf960312..089a1bb22d 100644
--- a/torchtitan/models/gpt_oss/model.py
+++ b/torchtitan/models/gpt_oss/model.py
@@ -28,7 +28,7 @@
 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 +45,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 +62,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 +83,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", rope_modules[config.rope.rope_key()])
 
     def forward(
         self,
@@ -137,7 +140,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 +148,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..534cfe5b9a 100644
--- a/torchtitan/models/qwen3_5/model.py
+++ b/torchtitan/models/qwen3_5/model.py
@@ -28,7 +28,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 +77,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 +94,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 +107,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", rope_modules[config.rope.rope_key()])
 
         self.q_norm = config.q_norm.build()
         self.k_norm = config.k_norm.build()
@@ -178,13 +181,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 "