[Feature] Support share MTP weights.

[RangiLyu]

Refactor MTP configuration to support weight sharing across layers. Updated MoE and MTPBlock classes to handle shared weights and adjusted layer initialization accordingly. Added share_weights parameter to MTPConfig for better control over layer behavior.

[RangiLyu]

@claude review

[claude]

Claude: Nit: self.num_layers is now semantically ambiguous. It stores the physical layer count (len(mtp_layers)), but the forward loop iterates self.mtp_config.num_layers times (the logical step count). When share_weights=True, these differ (1 vs. N).

If self.num_layers is not used externally, consider removing it or renaming it to self.num_physical_layers to avoid confusion. Alternatively, you could update it to always reflect the logical count:

self.num_layers = mtp_config.num_layers

[claude]

Claude: ## Summary

Adds share_weights parameter to MTPConfig so a single physical MTP layer can be reused across all prediction depths, reducing memory footprint.

Issues

Warning

• mtp_block.py:64-65 — Missing validation for share_weights=False: the forward loop now indexes self.layers[step] using mtp_config.num_layers, but nothing enforces len(mtp_layers) == mtp_config.num_layers in the constructor. Could cause IndexError if they diverge.
• moe.py:1237-1238 — _should_recompute docstring example doesn't cover share_weights=True behavior.

Nit

• mtp_block.py:67 — self.num_layers is now ambiguous (physical count vs. logical step count).
• moe.py:1018-1020 — Prefer assert or explicit is not None check over # type: ignore for strict mypy compliance.

Verdict

COMMENT — the logic is sound overall. The main thing to address is the missing len(mtp_layers) == mtp_config.num_layers validation in MTPBlock.__init__ for the non-sharing case.