Share equivalent RoPE cache storage

[Mystri]

No description provided.

[tianyu-l]

why do we need the capability of switch "registry", instead of always only having one default

[tianyu-l]

although memory is saved, compute seems still wasted

[tianyu-l]

Does it cover https://github.com/pytorch/torchtitan/blob/main/torchtitan/overrides/helion_rope.py -- it looks so because cache construction code is shared.

[tianyu-l]

related to the other comment -- why MTP requires an explicit context

[tianyu-l]

@claude
Previously each attention module would instantiate its own RoPE cache, which would be a problem for memory for long context training. This is the second attempt, previous one being #4111

Personally I like this one better since its cache management is explicit.

Please review this PR. In particular

• check if my comments make sense
• compare the two solutions and give your preference

[claude]

Claude finished @tianyu-l's task in 10m 22s —— View job

---

Review: Share equivalent RoPE cache storage

I read the three commits on this branch (f066d40 -> e333a42 -> d0929ae), the current rope.py / decoder.py / mtp.py, the sharding wiring in decoder_sharding.py + the per-model sharding.py files, and overrides/helion_rope.py. I could not fetch #4111's diff from this sandbox (network gh calls need approval I can't grant), so my comparison below is at the design level, grounded in your description that #4111 was the earlier attempt and this one makes cache management explicit.

• Read changed files + related sharding/override paths
• Evaluate the four review comments
• Compare the two approaches and state a preference

---

Do your four comments make sense? Yes, all four.

1. rope.py:52 -- "why the capability to switch registry instead of one default"
Agree this is the crux of the added complexity. The registry is per-Decoder (_RoPECacheRegistry(self)) and is threaded into layer_config.build() implicitly through a contextvars variable. So "switching" isn't a feature -- it's plumbing to pass the registry down without changing build() signatures. A process-wide singleton would be wrong (it would coalesce caches across independent model instances -- ref+policy models in RL, multiple PP chunks, tests), so some per-model scoping is needed. But a contextvar is the heaviest way to get it. The lighter alternatives are (a) pass the registry explicitly as a build argument, or (b) have the Decoder own a factory that constructs the RoPE modules. The contextvar's cost is exactly comment #4 (see below).

2. rope.py:214 -- "memory saved but compute still wasted"
Correct. RoPE.__init__ always calls self._precompute_cache() (line 214) and only then computes the key and discards the duplicate. So the outer-product / polar / cos-sin build runs once per layer regardless of sharing. It's construction-time only (not per-step), so bounded -- but for many layers x long max_context_length it's real startup cost. It also recurs in _init_self_buffers (rope.py:319-324): every sharing RoPE recomputes and re-_materializes the same slot, idempotently overwriting it N times.

The reason compute can't be skipped today is that _cache_key (line 256) is derived from the computed tensor (cache.shape, cache.dtype, cache.device). To skip precompute for a known key you'd have to derive the key from config alone. That's feasible -- shape/dtype/device are determined by dim/scaling/max_context_length plus the construction device -- but it's a design change, not a one-liner.

3. helion_rope.py -- "does it cover the Helion override?"
Yes, and your reasoning is right: HelionCosSinRoPE/HelionComplexRoPE subclass the stock classes and inherit __init__, so they register through the same registry, and their forward reads self.cache (the property) at helion_rope.py:991,1022, resolving through the reader. One subtlety: _cache_key includes type(self) (line 256), so a Helion module and a stock module with identical config would not coalesce. That's harmless in practice because the override rewrites every layer's config to the Helion type, so all layers in a model share the same type and still coalesce with each other. Worth a one-line note in the docstring so a future reader doesn't expect cross-type sharing.

4. mtp.py:215 -- "why does MTP need an explicit context"
Because MTP layers are built in MTPDecoder.__init__ after super().__init__() returns, i.e. outside the with _rope_cache_registry_context(registry) block that Decoder.__init__ wraps around only the main self.layers loop (decoder.py:265-268). Without re-entering via self._rope_cache_context(), the MTP RoPE modules would see no active registry and silently fall back to per-layer local buffers -- no error, just lost sharing. So the explicit context is required by the current design, which is exactly the footgun comment #1 hints at: any construction site outside Decoder.__init__ must remember to re-enter the context, and forgetting fails silently. This is the strongest argument against the contextvar mechanism.

---

The thing I'd fix before landing: the state_shardings={"cache"} declarations are now dead

This is my main correctness concern. Four files still set the RoPE cache TP sharding on the rope module:

• torchtitan/models/common/decoder_sharding.py:242
• torchtitan/models/deepseek_v3/sharding.py:106
• torchtitan/models/gpt_oss/sharding.py:112
• torchtitan/models/qwen3_5/sharding.py:308

all of the form attention.rope.sharding_config = ShardingConfig(state_shardings={"cache": dense_param_placement(tp=spmd.R)}).

After this PR, cache is a read-only property on the RoPE module (rope.py:221), not a registered buffer -- the real buffer is _rope_cache_N and it lives on the Decoder. Module._distribute_states (protocols/module.py:344-414) iterates self.named_parameters/named_buffers(recurse=False) and looks up each existing state's name in state_shardings; it never iterates the dict keys. So:

• On the RoPE module: there is no cache buffer, so {"cache": ...} matches nothing and is silently ignored.
• On the Decoder: _rope_cache_N buffers exist, but the Decoder has no state_shardings entry for them, so they're never distributed.

Combined with d0929ae ("Defer RoPE cache parallelization") removing registry.parallelize(), the shared cache now stays a plain (non-DTensor) tensor under TP, where previously it was a Replicate DTensor. It's likely functionally fine because the cache is identical across TP ranks and _reshape_for_broadcast is wrapped in @spmd.local_map (rope.py:497), but this is precisely the "silent change to a converged, TP-validated path" that warrants (a) a bitwise loss+grad_norm check under TP and CP per .claude/CLAUDE.md, and (b) either removing the four now-dead state_shardings={"cache"} declarations or re-wiring cache distribution onto the Decoder-owned slots. Leaving dead sharding declarations in the tree is misleading to the next person who touches TP. Fix this ->

---

Smaller notes

• Reference cycle. Each RoPE holds _cache_reader -> registry -> _owner (Decoder), so every RoPE transitively strong-refs the whole model. GC handles the cycle, but it's a new coupling worth a comment.
• _cache_key via repr(config) (rope.py:255) is fine for scalar/list configs, but it silently prevents coalescing when two configs differ in a cache-irrelevant field. Not a bug; just an accuracy-vs-simplicity tradeoff worth a line in the docstring.
• The cache property raising AttributeError("cache") (rope.py:235) to emulate a missing buffer is a nice touch for register_buffer's hasattr collision check -- worth keeping the comment that's already there.

---

Comparison and preference

I couldn't pull #4111's diff here, so I'm comparing your explicit-registry approach against the two general alternatives it sits between:

1. Shared RoPE module (one RoPE instance referenced by every attention layer). Simplest sharing, zero duplicate compute, one buffer -- but it breaks the "one module owns its state" tree assumption, complicates FSDP/PP module pruning and checkpoint FQNs, and forces all layers to be config-identical.
2. This PR: per-model registry of canonical buffers, modules stay independent. Modules keep their identity and FQNs; only the derived read-only buffer is deduped. This is the right layer to share at -- the cache is genuinely a derived, read-only artifact, so sharing storage without sharing modules is principled. The RoPECacheReader indirection also correctly handles the buffer object being replaced by device/dtype/DTensor transforms, which a naive shared-tensor reference would not.
3. Implicit/global dedup (what I'd guess RoPE cache sharing #4111 leaned toward, given your "this one is more explicit" framing) -- less boilerplate but harder to reason about lifetime and scoping.

I prefer this PR's direction (#2) over both a shared module and an implicit global: sharing the buffer while keeping modules independent is the correct abstraction, and explicit per-model ownership is easier to reason about than an implicit global cache. The one part I'd push back on is the mechanism for threading the registry -- the contextvars indirection is what produces comments #1 and #4 (the silent-fallback footgun in mtp.py). If you can pass the registry explicitly into build() or have the Decoder own RoPE construction, you'd keep the good buffer-dedup design while removing the implicit-context fragility. Net: keep the model, reconsider the plumbing, and resolve the dead state_shardings before landing.

Happy to compare against #4111's actual diff directly if you paste it or approve a gh pr diff 4111 fetch.
• branch rope-cache-registry

[tianyu-l]

@claude

Can't we just keep the cache at the RoPE module level, without touching Decoder at all?

[claude]

Claude finished @tianyu-l's task in 1m 55s —— View job

---

Re: "Can't we just keep the cache at the RoPE module level, without touching Decoder at all?"

Short answer: not while actually saving the memory. "Cache lives on the RoPE module" and "storage is shared" pull in opposite directions, because the memory win has to survive the model lifecycle -- and the lifecycle is exactly where a naive shared reference re-duplicates.

Why a per-module buffer can't share storage on its own

nn.Module._apply (used by .to(), to_empty(), and DTensor distribution) does not mutate a buffer in place -- it replaces each module's buffer with a freshly transformed tensor, independently per module. So even if layers.0.attention.rope.cache and layers.1.attention.rope.cache start life pointing at the same object:

• On meta device every buffer has zero storage, so nothing is shared yet.
• init_states() / to_empty() then materializes real storage per buffer -- one allocation per layer. That is precisely the moment the memory cost is incurred, and it's per-module.
• TP distribution likewise wraps each module's buffer into its own DTensor.

The identity link is broken at exactly the step that allocates the memory you were trying to save. That's why the PR resolves the slot on every read() (rope.py:47-48, 108-112) instead of holding a tensor -- the buffer object legitimately gets swapped out underneath it.

So dedup requires three things, no matter where you put it: a single canonical buffer, an owner stable enough to re-materialize it once, and a scoped lookup so independent models (ref+policy in RL, PP stages, tests) don't coalesce. That trio is the registry + reader. The only real design freedom is who owns the canonical buffer.

The three places the owner could live

1. The Decoder (this PR). The cache is a model-level resource derived from config; the container that outlives every layer owns it. Matches the _cache_key framing ("derived, read-only buffer").
2. The first RoPE module, with peers holding readers into it. This does keep the buffer literally on a RoPE module and lets you drop the _rope_cache_* buffers off the Decoder -- but it does not remove the mechanism the review comments flagged. You still need the per-model registry + context to find that owner, and MTP still has to re-enter it (mtp.py:215). It just relocates the buffer, and arguably makes things worse:
   • Layer 0 becomes a privileged resource owner: asymmetric FQNs (layers.0.attention.rope.cache present, peers have none).
   • If any flow removes the owning layer while keeping peers (PP layer-subsetting, pruning), the survivors dangle on a dead owner. The Decoder is a stable container per stage; "the first RoPE" is not.
   • The TP state_shardings={"cache"} declarations still match only the one owning module -- same dead-declaration problem, just moved.
3. A single shared RoPE submodule referenced by all layers. Zero duplicate compute, but it breaks "one module owns its state," forces all layers config-identical, and complicates FSDP/PP pruning and checkpoint FQNs.

Recommendation

Keeping Decoder as the owner is the right call -- it's a stable, model-scoped container and matches the semantics of the cache. The touch on Decoder is genuinely minimal (create a registry, wrap the layer-construction loop in a context). The part that's actually reducible is not the Decoder ownership -- it's the contextvars threading, which is what produces the review's comment #1 and the silent-fallback footgun in MTP (comment #4). Those are separable: you can keep Decoder ownership but pass the registry explicitly into build() (or have the Decoder drive RoPE construction), which deletes the implicit context and the "forgot to re-enter it -> silently falls back to per-layer buffers" failure mode -- while MTP's extra construction site becomes an ordinary explicit argument instead of a required with self._rope_cache_context().

Net: the Decoder coupling is inherent to sharing storage and is the cleanest of the three owners; the complexity worth removing is the contextvar, not the Decoder.

---

• Re-read rope.py, decoder.py, mtp.py, __init__.py
• Work through what "keep cache at RoPE module level" would require
• Answer with tradeoffs
  • branch rope-cache-registry

[tianyu-l]

@claude
I see. I think this issue exists only because we are doing

• first meta-device init
• and then parallelization

If we reverse it, we could avoid calling .to() or to_empty(). Is there caveat around doing that?

Another doubt:
Letting decoder owns the cache sounds less straightforward than explicitly passing rope cache to each attention module, instead of letting each attention owns its own rope module.

[claude]

Claude finished @tianyu-l's task in 5m 37s —— View job

---

Re: reversing init/parallelize, and explicit cache-passing vs. Decoder ownership

Two good questions. Short version: reversing the order does not cleanly remove the problem (and gives up the reason meta-init exists), and "pass the cache to each attention" is the right instinct but in practice collapses back into "the Decoder owns the caches and hands them down" -- the part worth deleting is the contextvar, not the Decoder ownership.

---

Q1: "This only happens because we meta-init then parallelize. If we reverse it, we avoid .to()/to_empty(). Caveats?"

The premise is only partly right. to_empty() is one place a shared buffer identity gets broken, but it isn't the only one. Buffer replacement happens in nn.Module._apply, which is invoked by .to(device), .to(dtype), .cuda(), and DTensor distribution -- each of those re-creates every buffer per-module and re-register_buffers it. So even if you construct directly on the real device:

• DTensor distribution still replaces the buffer. _distribute_states (protocols/module.py:410) calls register_buffer(name, distribute_tensor(...)) independently per module. Any code path that turns the cache into a DTensor re-duplicates it regardless of init order. This PR only avoids that because d0929ae defers (skips) cache parallelization and leaves it a plain replicated tensor -- so the reversal doesn't buy you anything the deferral didn't already.
• Any later .to() / dtype cast (mixed precision setup, device moves) breaks the shared reference the same way.

So to make a plain shared reference actually survive, you'd need: construct on the real device and never distribute the cache and never .to() the model afterward. That's a lot of "nevers," and it's exactly the fragility the RoPECacheReader.read()-resolves-every-call indirection was built to avoid. The reader is robust to buffer replacement no matter the ordering; reversing the order just trades one robust mechanism for a pile of lifecycle constraints.

The bigger caveat: meta-init isn't an accident we can casually reverse. It exists precisely so we never materialize the full unsharded model in memory on one rank before sharding. "Materialize first, shard second" means every rank allocates the full unsharded parameter set up front -- which OOMs large models. It'd be ironic for a memory-saving PR to require abandoning the primary memory-scaling mechanism. That's the caveat I'd lead with.

Net: reversing doesn't remove the buffer-replacement problem (DTensor distribution + .to() still replace), and it costs meta-init's OOM avoidance. The reader indirection is the cheaper robustness.

---

Q2: "Explicitly pass the rope cache to each attention, instead of each attention owning its own rope module"

This is the same good instinct as my earlier "pass the registry explicitly instead of via contextvar" -- and I agree it's cleaner than the contextvar. But when you make it concrete, two things surface:

1. Someone above the layers still has to build and hold the caches -- that someone is the Decoder. Attention builds its rope deep inside layer construction (gpt_oss/model.py:84, attention.py:934: self.rope = config.rope.build()). To "pass the cache in," the Decoder must build the canonical cache(s) before constructing layers and thread them down. That's still Decoder ownership -- just with an explicit argument instead of a contextvar. Which is exactly what I'd endorse: keep Decoder ownership, delete the contextvar.

2. The real cost is threading an argument through the build() chain. Every module is constructed via a uniform config.build() with no extra positionals. To pass a cache explicitly you'd thread it Decoder -> TransformerBlock.build() -> attention.build() -> RoPE, i.e. change build() signatures across every model. The contextvar was chosen specifically to avoid that signature churn. So the honest tradeoff is implicit-context fragility (comments #1, #4) vs. pervasive-signature-churn -- not "explicit is strictly simpler."

3. Caches are not all identical, so a single passed-in cache doesn't suffice. Layers can legitimately need distinct caches (e.g. differing theta/scaling, MTP layers, sliding-window vs. full). The registry keys on config (_cache_key, rope.py:238) and coalesces only equivalent caches while keeping distinct ones separate. If you pass "the cache" down, the Decoder has to build the set of distinct caches and route the right one to each layer -- which reintroduces key-based lookup, i.e. a registry, just spelled explicitly at the Decoder.

And passing a shared rope module (rather than the cache) is worse, for the reasons in the earlier review: the same module appears under multiple attention FQNs, breaking "one module owns its state," duplicating checkpoint FQNs, and confusing FSDP/PP traversal. Sharing the buffer while modules stay independent is the right layer to dedup at.

---

Recommendation (unchanged in shape, sharpened)

• Keep the Decoder as cache owner -- it's the stable, model-scoped container, and Q1 shows sharing storage inherently needs a stable owner that survives buffer replacement.
• Replace the contextvar with an explicit registry argument. Pass the Decoder's registry (or a small cache-factory) explicitly into build() / into the layer construction path. That deletes comment Add profiler #1's "why switchable" and comment Add lr scheduler for optimizer #4's silent-fallback footgun (MTP's with self._rope_cache_context() becomes an ordinary argument), without forcing a single-cache assumption -- the registry still coalesces by key.
• The "pass the cache to each attention" idea is the right direction; it just resolves, once made concrete and once heterogeneous caches are accounted for, to "Decoder owns a keyed cache table and hands the right reader to each layer explicitly." That's the contextvar removal, keeping the buffer-dedup design.

(Still worth resolving before landing: the now-dead state_shardings={"cache"} declarations flagged in the first review, and a bitwise loss/grad_norm check under TP + CP given cache is now a plain tensor rather than a Replicate DTensor.)

---

• Re-read rope.py, decoder.py, mtp.py, __init__.py
• Answer Q1: reversing meta-init/parallelize -- caveats (DTensor distribution + .to() still replace buffers; meta-init exists for OOM avoidance)
• Answer Q2: explicit cache-passing vs. Decoder ownership -- collapses to Decoder-owned keyed table + explicit arg; the win is removing the contextvar
  • branch rope-cache-registry

[tianyu-l]

@claude

And passing a shared rope module (rather than the cache) is worse, for the reasons in the earlier review: the same module appears under multiple attention FQNs, breaking "one module owns its state," duplicating checkpoint FQNs, and confusing FSDP/PP traversal. Sharing the buffer while modules stay independent is the right layer to dedup at.

I don't believe there's a lot of heterogenous rope used in the same model. It's usually at most one. In that case, how about

• make rope an Decoder-owned module, with cache still as registered buffer
• explicitly pass it into the constructors of Attention modules
• change PP stage building so that the Decoder.rope exists on every PP stage (multiple PP stages on the same PP rank owns duplicate and we admit this overhead)

[claude]

Claude finished @tianyu-l's task in 4m 46s —— View job

---

Re: Decoder-owned single rope module, passed into Attention, duplicated per PP stage

I read attention.py (self.rope = config.rope.build() at :934, used at :966), the gpt_oss/deepseek_v3 attention variants (model.py:84, model.py:91), qwen3_5/rope.py (MRoPE) and vision_encoder.py, plus how the caches are wired for TP. Short version: I think this is a better design than the current PR, and it addresses the FQN objection I raised earlier -- because for a RoPE module the only state is a non-persistent cache buffer, so "shared module under multiple FQNs" turns out to be far cheaper than I made it sound. Two mechanics need to be gotten right, and one assumption (homogeneity) needs an explicit decision.

---

Why this is actually cleaner than the current PR

The current PR shares a buffer tensor, and the whole registry/reader/contextvar apparatus exists for one reason: nn.Module._apply (.to(), to_empty(), DTensor distribution) replaces the buffer object per module, so a shared tensor reference silently re-duplicates at materialization time. That's what forced RoPECacheReader.read() to re-resolve the slot on every call.

Sharing a module sidesteps that entirely, and this is the key point:

Module object identity is stable across _apply; only its buffer attribute is swapped in place.

So if each Attention holds a reference to the one rope module (not a reference to the cache tensor), then to_empty() / .to() / DTensor distribution mutate that single module's cache buffer once, and every attention transparently sees the new tensor through self.rope.cache. No reader, no registry, no contextvar, no _cache_key. That's a large deletion.

It also resolves two of the four review comments outright:

• Comment Add FSDP grad scaler to the train loop #2 (wasted compute): the cache is now built exactly once by the Decoder, not once-per-layer-then-discarded. Fully fixed, not just bounded.
• Comment Add lr scheduler for optimizer #4 (MTP explicit context): MTPDecoder just references the same self.rope; the with self._rope_cache_context() at mtp.py:215 disappears. The silent-fallback footgun goes away.

And my earlier FQN / checkpoint objection is much weaker than I stated: RoPE's cache is a persistent=False buffer and the module has no parameters, so it contributes nothing to the state dict. Even if the module were reachable under multiple FQNs, there's no checkpoint duplication to worry about. I was wrong to weight that concern as heavily as I did for the module-sharing option; it only bites for modules that own real state, which RoPE doesn't.

---

The two mechanics to get right

1. Attention must reference the rope module without registering it as a submodule.
This is the one genuinely awkward part. nn.Module.__setattr__ auto-registers any Module assigned to an attribute, so a plain self.rope = shared_rope would register the same object under every layers.N.attention.rope and under the Decoder -- reintroducing multi-parent traversal (double _distribute_states, double FSDP visitation depending on remove_duplicate). To keep Decoder as the sole owner, attention has to hold it in a non-registering slot (e.g. stash in a 1-element list / a types.SimpleNamespace / object.__setattr__) and call it as such in forward. It works and is a known pattern (it's how you hold a back-reference to a module), but it's slightly less readable than self.rope, and it's worth a comment explaining why it's not a normal submodule. This is the real cost that replaces the registry -- I'd weigh "one non-registered reference + a comment" against "registry + reader + contextvar + _cache_key + MTP context," and the former is clearly less machinery.

2. Threading the shared module into attention.
config.build() is uniform with no extra positionals, so there are two shapes:

• Post-build injection (my preference): Decoder builds first_attention.rope.build() once, then walks its constructed layers and injects the reference into each attention (dropping the one attention built for itself). Localized entirely to Decoder.__init__, zero build() signature churn. Downside: attention still runs its own config.rope.build() and you throw it away -- unless you also give attention a way to skip self-construction, which is a small flag.
• Explicit constructor arg: cleaner conceptually but requires threading rope through TransformerBlock.build() -> attention.build() for every model. That's the signature churn the contextvar was invented to avoid.

Either is fine; both keep Decoder as the owner, which Q1 already established is the correct owner (stable, model-scoped, survives buffer replacement).

---

The one assumption that needs an explicit decision: homogeneity

Your premise ("usually at most one rope") holds across every current model I checked:

• llama3/4, qwen3, deepseek_v3, gpt_oss: a single rope config for all full-attention layers (deepseek's MTP reuses the same one).
• qwen3_5: all LLM full-attention layers use one MRoPE config; GDN linear-attention layers have no rope; the vision encoder builds its rope cache inline (vision_encoder.py:126 _compute_2d_rope_cache), not through a RoPE module -- so neither the current PR nor this proposal touches it. Good.

So a single Decoder.rope is correct for everything today. But note the current PR's registry coalesces by key, so it degrades gracefully to N distinct caches if a future model ever needs per-layer theta / sliding-vs-full rope. The single-module design gives that up. I'd handle it explicitly rather than silently:

• Build the shared rope from config.first_attention.rope, and assert every full-attention layer's rope config is equal to it (clear ValueError if not). Per the repo's "no speculative generality" rule, asserting the invariant you actually have is the right call -- just make the failure loud, since the failure mode otherwise is a wrong cache silently applied to a mismatched layer.

---

PP-stage duplication: agreed, and it's the right tradeoff

Building Decoder.rope on every PP stage is correct and matches how meta-init/parallelize already works per stage. Each stage is its own Decoder that builds its own (single) rope; its attentions reference that stage's rope. The duplication only occurs when multiple virtual stages land on the same rank, and a single replicated cache is negligible next to the parameter shards -- explicitly admitting that overhead is the honest and simple choice. This is strictly simpler than trying to make one cache span stages.

---

Net recommendation

Switch to the Decoder-owned single-module design. It deletes the registry/reader/contextvar/_cache_key, resolves comments #2 and #4 outright, and the FQN worry that made me cautious earlier doesn't apply to a parameter-free, non-persistent-buffer module. The two things to nail down:

1. Hold the rope in Attention as a non-registered reference (list/namespace, with a comment) so the Decoder stays the sole owner -- avoids double distribute/FSDP traversal.
2. Assert rope-config homogeneity across full-attention layers with a clear error, since you're giving up the by-key coalescing.

Plus the two carryovers from the earlier review, which get easier under this design: the TP state_shardings={"cache"} wiring now targets the real single rope.cache buffer on the Decoder-level module (re-point decoder_sharding.py:242, deepseek_v3/sharding.py:106, gpt_oss/sharding.py:112, qwen3_5/sharding.py:308 at it rather than at a per-attention attention.rope), and the mandatory bitwise loss/grad_norm check under TP + CP still applies since the cache remains a plain/replicated tensor.

One caveat I can't rule out from here: I did not run the numerics or a full FSDP+PP+TP integration pass, so the "non-registered reference is invisible to FSDP wrapping and visible-once to _distribute_states" claim should be validated on a real 2-stage PP + TP run before landing -- that's the part most likely to surprise.

• Verify how attention builds/uses RoPE and confirm heterogeneity is rare
• Work through explicit-pass-in mechanics (build() signature, MTP, non-registered reference)
• Work through the PP-stage-building change
• Answer with tradeoffs + recommendation
  • branch rope-cache-registry

[tianyu-l]

@Mystri see if you agree with #4376 (comment)

[Mystri]

@Mystri see if you agree with #4376 (comment)

We cannot safely assert homogeneity for newer models such as deepseek-v4, which computes freqs_cis with different rope_theta for compressed/non-compressed layers:
https://huggingface.co/deepseek-ai/DeepSeek-V4-Pro/blob/main/inference/model.py#L473

My proposal that fits yours best is that the Decoder creates and owns both types of ropes and pass them as a "dict of ropes" along, and the builder extracts the correct rope module from it when building the submodule.
@claude