fix(mtp): resolve parameter scoping in Multi-Token Prediction block

The MTP transformer layer parameters were incorrectly attaching to the top-level Linen scope (`mtp_block`) instead of the nested NNX scope (`mtp_layer_1`). This occurred because the parent model passed `ToLinen` wrappers into the MTP block, causing the Flax NNX tracer to ignore the nested layers and default to Linen's automatic scoping.

Changes:
* src/maxtext/layers/nnx_wrappers.py: Exposed the underlying native NNX class in the `to_linen_class` factory via a new `module_class` attribute.
* src/maxtext/models/models.py: Dynamically unwrapped the transformer layer using `getattr(mtp_layer_linen, "module_class")` before injecting it into the MTP block to preserve NNX tracing.
* src/maxtext/layers/multi_token_prediction.py: Removed the legacy Linen `name=` argument from the transformer layer instantiation and added required native NNX arguments (`quant=None`, `layer_idx=-1`).

Author: parambole

src/maxtext/layers/multi_token_prediction.py

@@ -79,7 +79,6 @@ def __init__(
     self.layer_number = layer_number
     self.transformer_layer_module = transformer_layer_module
     self.rngs = rngs
-    k = layer_number
     cfg = self.config
 
     self.embedding_norm = RMSNorm(
@@ -112,7 +111,6 @@ def __init__(
         config=cfg,
         mesh=mesh,
         model_mode=MODEL_MODE_TRAIN,
-        name=f"mtp_{k}_transformer_layer",
         rngs=rngs,
     )
 

src/maxtext/layers/nnx_wrappers.py

@@ -627,5 +627,6 @@ def __init_subclass__(cls, **kwargs):
   ToLinenPartial.__qualname__ = class_name
 
   ToLinenPartial.__init__ = __init__
+  ToLinenPartial.module_class = base_nnx_class
 
   return ToLinenPartial

src/maxtext/models/models.py

@@ -93,14 +93,17 @@ def setup(self):
     # If MTP is enabled via config, set up the MTP block.
     if self.config.mtp_num_layers > 0:
       # Get the list of layer blueprints for the current model.
-      layer_types = self.decoder.get_decoder_layers()
       # For MTP, we use the DecoderLayer blueprint to ensure architectural consistency.
       # By convention, this is the last layer in the list.
-      mtp_layer = layer_types[-1]
+      layer_types = self.decoder.get_decoder_layers()
+      mtp_layer_linen = layer_types[-1]
+      # UNWRAP: The MTP block is pure NNX. If the decoder returned a Linen wrapper,
+      # extract the native NNX class to preserve parameter tracing/scoping.
+      mtp_layer_nnx = getattr(mtp_layer_linen, "module_class", mtp_layer_linen)
       self.mtp_block = multi_token_prediction_block_as_linen(
           config=self.config,
           mesh=self.mesh,
-          transformer_layer_module=mtp_layer,
+          transformer_layer_module=mtp_layer_nnx,
           decoder=self.decoder,
           rngs=self.make_rng("mtp_block"),
       )