diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml
index 016f31a322..0dba8e6575 100644
--- a/src/maxtext/configs/base.yml
+++ b/src/maxtext/configs/base.yml
@@ -1374,6 +1374,18 @@ vllm_hf_config_path: ""
 vllm_hf_overrides: {}
 # JSON string containing additional configuration for the vLLM model (e.g. '{"maxtext_config": {...}}')
 vllm_additional_config: {}
+# Whether to use the streaming target-free weight converter for trainer-side conversion
+use_weight_converter: true
+# Rollout backend for trainer-side weight converter ('maxtext' or 'vllm_torchax')
+rollout_backend: "maxtext"
+# Degree of tensor parallelism for KV cache / attention heads in rollout
+kv_tp_size: 1
+# Degree of tensor parallelism for MoE MLP dimension in rollout
+moe_mlp_tp_size: 1
+# Enable debug verification and checksum logging during weight sync
+weight_sync_debug: false
+# Weight load format for vLLM in converter validation
+vllm_load_format: "dummy"
 # When use_jax_splash=true, force the layout of the query tensor to be [..., NUM_HEADS, HEAD_DIM, SEQ_LENGTH]
 force_q_layout: false
 
diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py
index 84acba8936..094e47418f 100644
--- a/src/maxtext/configs/types.py
+++ b/src/maxtext/configs/types.py
@@ -2741,12 +2741,20 @@ class VLLM(BaseModel):
   vllm_hf_config_path: str = Field("", description="Path to HuggingFace model config for MaxText model.")
   use_standalone_converter: bool = Field(False, description="Use the standalone MaxText->torchax vLLM converter")
   use_weight_converter: bool = Field(
-      False,
+      True,
       description=(
           "Use an explicit weight converter for trainer->rollout weight sync instead of "
           "the legacy transfer_state_directly / transfer_state_with_mappings paths."
       ),
   )
+  use_raiden_ffi: Optional[bool] = Field(
+      None,
+      description="Use Raiden FFI transport for weight sync.",
+  )
+  rollout_backend: Literal["maxtext", "vllm_torchax"] = Field(
+      "maxtext",
+      description="Rollout backend for trainer-side weight converter ('maxtext' or 'vllm_torchax').",
+  )
   weight_sync_debug: bool = Field(
       False,
       description=(
@@ -2766,6 +2774,16 @@ class VLLM(BaseModel):
           "rather than data movement. Adds one barrier per sync."
       ),
   )
+  kv_tp_size: int = Field(
+      1,
+      ge=1,
+      description="Degree of tensor parallelism for KV cache / attention heads in rollout.",
+  )
+  moe_mlp_tp_size: int = Field(
+      1,
+      ge=1,
+      description="Degree of tensor parallelism for MoE MLP dimension in rollout.",
+  )
   vllm_load_format: str = Field(
       "dummy",
       description="Weight load format for vLLM in converter validation. Options:'auto', 'dummy'.",
diff --git a/src/maxtext/integration/vllm/convert_utils.py b/src/maxtext/integration/vllm/convert_utils.py
index 1909c6f26b..1d7810110c 100644
--- a/src/maxtext/integration/vllm/convert_utils.py
+++ b/src/maxtext/integration/vllm/convert_utils.py
@@ -41,8 +41,9 @@
     `_bulk_align_and_unstack`; that patch goes away once the port lands.
 """
 
-from typing import Mapping, Any, Callable, Dict, Tuple, Optional
 import functools
+import os
+from typing import Mapping, Any, Callable, Dict, Tuple, Optional
 from absl import logging
 import jax
 import jax.numpy as jnp
@@ -93,7 +94,9 @@ def compute_padded_moe_mlp_dim(
 
   min_required = 2 * num_lanes * moe_mlp_tp_size
   if (hidden_size // moe_mlp_tp_size) % (2 * num_lanes) != 0:
-    return ((max(hidden_size, min_required) + min_required - 1) // min_required) * min_required
+    return (
+        (max(hidden_size, min_required) + min_required - 1) // min_required
+    ) * min_required
 
   return hidden_size
 
@@ -178,7 +181,9 @@ def _fuse_moe_weights(
         fused_shape_tuple,
         axis,
     )
-    new_src_flat[wi_target_key] = _interleave_moe_weights(wi_0, wi_1, fused_shape_tuple, n_shards, axis=axis)
+    new_src_flat[wi_target_key] = _interleave_moe_weights(
+        wi_0, wi_1, fused_shape_tuple, n_shards, axis=axis
+    )
     del wi_0, wi_1
 
   return new_src_flat
@@ -292,7 +297,8 @@ def _unstack_scanned_param(
       return tuple(jnp.unstack(src_val))
     else:
       logging.warning(
-          "Shape mismatch in scanned param '%s'. Src: %s, Tgt: %s. Cannot" " determine scan axis.",
+          "Shape mismatch in scanned param '%s'. Src: %s, Tgt: %s. Cannot"
+          " determine scan axis.",
           key_path,
           src_shape,
           tgt_shape,
@@ -330,7 +336,9 @@ def _get_n_shards(arr: jax.Array | np.ndarray, axis: int) -> int:
   """Returns the number of shards for a given axis of an array."""
   sharding = getattr(arr, "sharding", None)
   if isinstance(sharding, jax.sharding.NamedSharding):
-    return _partition_size(_spec_at_axis(sharding, axis), sharding.mesh)  # pyrefly: ignore[bad-argument-type]
+    return _partition_size(
+        _spec_at_axis(sharding, axis), sharding.mesh
+    )  # pyrefly: ignore[bad-argument-type]
   return 1
 
 
@@ -450,14 +458,18 @@ def _align_per_axis(
   if arr.shape == tgt_shape:
     return arr
   if len(arr.shape) != len(tgt_shape):
-    raise ShapeMismatchError(f"Rank mismatch for {key_path}: src={arr.shape} vs tgt={tgt_shape}")
+    raise ShapeMismatchError(
+        f"Rank mismatch for {key_path}: src={arr.shape} vs tgt={tgt_shape}"
+    )
 
   mismatches = []
   for axis, (s, t) in enumerate(zip(arr.shape, tgt_shape)):
     if s == t:
       continue
     if t < s:
-      raise ShapeMismatchError(f"Cannot shrink axis {axis} for {key_path}: src={s} -> tgt={t}")
+      raise ShapeMismatchError(
+          f"Cannot shrink axis {axis} for {key_path}: src={s} -> tgt={t}"
+      )
     mismatches.append((axis, s, t))
   if not mismatches:
     return arr
@@ -468,7 +480,9 @@ def _align_per_axis(
       mesh = tgt_sharding.mesh
       pad_specs = []
       for axis, s, t in mismatches:
-        n_shards = _partition_size(_spec_at_axis(tgt_sharding, axis), mesh)  # pyrefly: ignore[bad-argument-type]
+        n_shards = _partition_size(
+            _spec_at_axis(tgt_sharding, axis), mesh
+        )  # pyrefly: ignore[bad-argument-type]
         if t % n_shards != 0:
           raise ValueError(
               f"Target dimension {t} on axis {axis} for {key_path} is not "
@@ -499,14 +513,16 @@ def _align_per_axis(
   return _jit_repeat_axes(arr, tuple(repeats))
 
 
-@functools.partial(jax.jit, static_argnames=("tgt_shape", "n_shards", "axis", "lane_size"))
+@functools.partial(
+    jax.jit, static_argnames=("tgt_shape", "n_shards", "axis", "lane_size")
+)
 def _interleave_moe_weights(
     wi_0: jax.Array | np.ndarray,
     wi_1: jax.Array | np.ndarray,
     tgt_shape: Tuple[int, ...],
     n_shards: int,
     axis: Optional[int] = None,
-    lane_size: int = DEFAULT_TPU_NUM_LANES,
+    lane_size: Optional[int] = 0,
 ) -> jax.Array | np.ndarray:
   """Interleaves wi_0 and wi_1 per-shard into a single tensor matching TPU GMM layout.
 
@@ -517,9 +533,18 @@ def _interleave_moe_weights(
   that is the difference between ~3x and ~1x the output size in live transient
   memory, plus ~5 fewer dispatches per call.
 
+  For TPU GMM_v2 kernels (e.g. gmm_v2.py), each TP shard expects plain concatenation
+  [local_gate, local_up]. The kernel internally performs 128-lane interleaving into
+  VMEM via interleave_lane(w_gate, w_up). Therefore, lane_size defaults to 0 (plain
+  concatenation per shard). Pre-interleaving in HBM double-interleaves and corrupts
+  weights.
+
   `tgt_shape`, `n_shards`, `axis` and `lane_size` are static, so the trace is keyed on them;
   identical layers share a single compilation.
   """
+  if lane_size is None:
+    lane_size = 0
+
   if axis is None:
     axis = len(tgt_shape) - 1
   elif axis < 0:
@@ -559,7 +584,7 @@ def _pad_and_chunk(arr):
     p_wi_1 = p_wi_1.reshape(shape_lanes)
     combined = jnp.stack([p_wi_0, p_wi_1], axis=axis + 2)
   else:
-    # Fallback when dimension is not divisible by lane_size:
+    # Concatenate wi_0 (gate) and wi_1 (up) per shard: [local_gate, local_up].
     combined = jnp.concatenate([p_wi_0, p_wi_1], axis=axis + 1)
 
   return combined.reshape(tgt_shape)
@@ -661,10 +686,18 @@ def _bulk_align_and_unstack(
   if isinstance(arr, jax.ShapeDtypeStruct):
     num_layers = arr.shape[scan_axis]
     tgt_dtype = getattr(per_layer_tgt_val, "dtype", getattr(arr, "dtype", jnp.float32))
-    return tuple(jax.ShapeDtypeStruct(per_layer_shape, tgt_dtype) for _ in range(num_layers))
+    return tuple(
+        jax.ShapeDtypeStruct(per_layer_shape, tgt_dtype) for _ in range(num_layers)
+    )
 
-  scanned_tgt_shape = per_layer_shape[:scan_axis] + (arr.shape[scan_axis],) + per_layer_shape[scan_axis:]
-  scanned_tgt_sharding = _scanned_sharding_from_per_layer(getattr(per_layer_tgt_val, "sharding", None), scan_axis)
+  scanned_tgt_shape = (
+      per_layer_shape[:scan_axis]
+      + (arr.shape[scan_axis],)
+      + per_layer_shape[scan_axis:]
+  )
+  scanned_tgt_sharding = _scanned_sharding_from_per_layer(
+      getattr(per_layer_tgt_val, "sharding", None), scan_axis
+  )
 
   if arr.shape == scanned_tgt_shape:
     return _jit_unstack(arr, scan_axis)
@@ -692,16 +725,54 @@ def _scanned_sharding_from_per_layer(
 
 
 def resolve_rollout_tp(config: Any, tp: int = 1) -> int:
-  """Resolves rollout TP from config."""
+  """Resolves rollout TP from override, config, or environment."""
   if tp > 1:
     return int(tp)
 
   config_tp = 0
   if config is not None:
-    config_tp = int(
+    raw = (
         getattr(config, "rollout_tensor_parallelism", 0)
         or getattr(getattr(config, "cluster", None), "rollout_tensor_parallelism", 0)
         or getattr(config, "rollout_mesh_tp", 0)
         or 0
     )
+    if raw and int(raw) > 0:
+      config_tp = int(raw)
+
+  if not config_tp:
+    env_tp = os.environ.get("ROLLOUT_TENSOR_PARALLEL_SIZE") or os.environ.get(
+        "ROLLOUT_MESH_TP"
+    )
+    if env_tp and int(env_tp) > 0:
+      config_tp = int(env_tp)
+
   return int(config_tp or 1)
+
+
+def resolve_prefuse_moe_weights(
+    config: Any, prefuse_moe_weights: Optional[bool] = None
+) -> bool:
+  """Resolves MoE prefuse flag from override, config, or environment."""
+  if prefuse_moe_weights is not None:
+    return bool(prefuse_moe_weights)
+  if "ROLLOUT_PREFUSE_MOE_WEIGHTS" in os.environ:
+    return os.environ["ROLLOUT_PREFUSE_MOE_WEIGHTS"].lower() in ("1", "true", "yes")
+  if (
+      config is not None
+      and getattr(config, "rollout_prefuse_moe_weights", None) is not None
+  ):
+    return bool(config.rollout_prefuse_moe_weights)
+  rollout_backend = getattr(config, "rollout_backend", None) or os.environ.get(
+      "ROLLOUT_BACKEND", "maxtext"
+  )
+  if rollout_backend == "maxtext":
+    return True
+  if config is not None and getattr(config, "prefuse_moe_weights", None) is not None:
+    return bool(config.prefuse_moe_weights)
+  return os.environ.get("PREFUSE_MOE_WEIGHTS", "0").lower() in ("1", "true", "yes")
+
+
+def is_verify_weights_enabled() -> bool:
+  """Returns whether weight verification / checksum validation is active."""
+  return os.environ.get("VERIFY_WEIGHTS", "").lower() == "true"
diff --git a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py
index 6e5454d936..bd000f3c2a 100644
--- a/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py
+++ b/src/maxtext/integration/vllm/maxtext_vllm_adapter/adapter.py
@@ -23,6 +23,7 @@
 from jax.sharding import Mesh
 from maxtext.common.common_types import MODEL_MODE_AUTOREGRESSIVE
 from maxtext.configs import pyconfig
+from maxtext.integration.vllm.convert_utils import DEFAULT_TPU_NUM_LANES, compute_padded_moe_mlp_dim
 from maxtext.integration.vllm.hybrid_cache_utils import (
     build_qwen_gdn_cache_layout,
     gather_layer_kv_caches,
@@ -124,7 +125,10 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters
       else vllm_config.model_config.hf_config
   )
   hidden_size = getattr(hf_config, "moe_intermediate_size", None)
-  num_lanes = pltpu.get_tpu_info().num_lanes
+  try:
+    num_lanes = pltpu.get_tpu_info().num_lanes
+  except Exception:
+    num_lanes = DEFAULT_TPU_NUM_LANES
   num_kv_heads = hf_config.num_key_value_heads
 
   # Number of KV heads in global attention layers (None if the field is absent or unset).
@@ -172,10 +176,8 @@ def generate_maxtext_config(vllm_config: VllmConfig) -> pyconfig.HyperParameters
   # The GMM_v2 kernel requires the MLP dimension per expert to be at least 2x the number of TPU lanes
   # to ensure efficient execution. See the validate_inputs() method in the following file for more details:
   # https://github.com/vllm-project/tpu-inference/blob/main/tpu_inference/kernels/megablox/gmm_v2.py
-  if hidden_size is not None and (hidden_size // moe_mlp_tp_size) % (2 * num_lanes) != 0:
-    padded_hidden_size = next_power_of_two(hidden_size)
-    while (padded_hidden_size // moe_mlp_tp_size) < (2 * num_lanes):
-      padded_hidden_size = next_power_of_two(padded_hidden_size + 1)
+  padded_hidden_size = compute_padded_moe_mlp_dim(hidden_size, moe_mlp_tp_size, num_lanes)
+  if padded_hidden_size is not None and padded_hidden_size != hidden_size:
 
     # This inflates every expert weight, so it is a real memory/FLOP cost rather than a
     # cosmetic reshape: at moe_mlp_tp_size=4 a 512-wide MoE is padded to 1024 (2x the MoE
@@ -471,6 +473,7 @@ def load_weights(self, rng_key: jax.Array) -> None:
         if self.maxtext_config.lora.lora_restore_path:
           lora_utils.restore_lora_from_path(model, self.maxtext_config)
       self.model = nnx.data(model)
+    patch_raiden_worker_h2d()
 
   def get_mrope_input_positions(
       self,
@@ -614,3 +617,14 @@ def patched_get_kv_cache_spec(self):
 
   KVCacheManager.get_kv_cache_spec = patched_get_kv_cache_spec
   max_logging.log("Successfully applied KVCacheManager patch for hybrid GDN models.")
+
+
+def patch_raiden_worker_h2d():
+  """Monkey-patches TPUWorker.raiden_h2d and RaidenWorkerSync to apply Raiden weights to runner."""
+  try:
+    from tunix.experimental.weight_sync.raiden_synchronizer import patch_raiden_worker_sync
+
+    patch_raiden_worker_sync()
+  except Exception as e:
+    max_logging.log(f"Skipping raiden worker sync patch: {e}")
+
diff --git a/src/maxtext/integration/vllm/weight_converter.py b/src/maxtext/integration/vllm/weight_converter.py
index b927b7173a..26efcc6e5e 100644
--- a/src/maxtext/integration/vllm/weight_converter.py
+++ b/src/maxtext/integration/vllm/weight_converter.py
@@ -32,6 +32,7 @@
     _device_ids,
     _fuse_and_unstack_moe,
     _get_n_shards,
+    _jit_repeat_axes,
     _jit_unstack,
     _scanned_sharding_from_per_layer,
     _sharding_summary,
@@ -257,6 +258,8 @@ def __init__(
       self,
       rules: Optional[List[Rule]] = None,
       tp: int = 1,
+      kv_tp_size: int = 1,
+      moe_mlp_tp_size: int = 1,
       num_kv_heads: Optional[int] = None,
       head_dim: Optional[int] = None,
       config: Any = None,
@@ -279,6 +282,8 @@ def __init__(
       )
     self.rules = rules
     self.tp = resolve_rollout_tp(config, tp)
+    self.kv_tp_size = kv_tp_size or getattr(config, "kv_tp_size", 1) or self.tp
+    self.moe_mlp_tp_size = moe_mlp_tp_size or getattr(config, "moe_mlp_tp_size", 1) or self.tp
 
     # Read by the rollout engine to decide whether to trace the reshard
     # step that runs after conversion.
@@ -296,6 +301,8 @@ def __init__(
       self._direct = MaxTextToMaxTextConverter(
           config=config,
           tp=self.tp,
+          kv_tp_size=self.kv_tp_size,
+          moe_mlp_tp_size=self.moe_mlp_tp_size,
           moe_fused_layout=(moe_fused_layout or MoEFusedLayout.PER_SHARD_INTERLEAVE),
           allow_unused_source_keys=allow_unused_source_keys,
           debug=debug,
@@ -678,9 +685,13 @@ def __init__(
       debug: bool = False,
       prefuse_moe_weights: Optional[bool] = None,
       target_dtype: Optional[Any] = None,
+      kv_tp_size: int = 1,
+      moe_mlp_tp_size: int = 1,
   ):
     self.config = config
     self.tp = resolve_rollout_tp(config, tp)
+    self.kv_tp_size = kv_tp_size or getattr(config, "kv_tp_size", 1) or self.tp
+    self.moe_mlp_tp_size = moe_mlp_tp_size or getattr(config, "moe_mlp_tp_size", 1) or self.tp
     self.moe_fused_layout = moe_fused_layout
     self.allow_unused_source_keys = allow_unused_source_keys
     self.debug = debug
@@ -691,6 +702,15 @@ def __init__(
     self.padded_base_moe_mlp_dim = getattr(config, "padded_base_moe_mlp_dim", None)
     self.target_dtype = target_dtype if target_dtype is not None else getattr(config, "weight_dtype", None)
 
+    self.base_num_kv_heads = getattr(config, "base_num_kv_heads", None) or getattr(config, "num_kv_heads", None)
+    self.kv_replication = 1
+    if self.base_num_kv_heads is not None and self.kv_tp_size > self.base_num_kv_heads:
+      if self.kv_tp_size % self.base_num_kv_heads != 0:
+        raise ValueError(
+            f"kv_tp_size ({self.kv_tp_size}) must be divisible by base_num_kv_heads ({self.base_num_kv_heads})."
+        )
+      self.kv_replication = self.kv_tp_size // self.base_num_kv_heads
+
     self.cycle = int(getattr(config, "inhomogeneous_layer_cycle_interval", 1) or 1)
     self.num_decoder_layers = int(config.num_decoder_layers)
     if self.num_decoder_layers % self.cycle:
@@ -707,7 +727,8 @@ def __init__(
 
     logging.info(
         "MaxTextToMaxTextConverter: %d layers, cycle=%d, %d scanned blocks, "
-        "scan_axis=%d, moe_fused_layout=%s, prefuse_moe=%s, padded_moe_dim=%s",
+        "scan_axis=%d, moe_fused_layout=%s, prefuse_moe=%s, padded_moe_dim=%s, "
+        "kv_tp_size=%d, moe_mlp_tp_size=%d, kv_replication=%d",
         self.num_decoder_layers,
         self.cycle,
         self.num_blocks,
@@ -715,6 +736,9 @@ def __init__(
         self.moe_fused_layout,
         self.prefuse_moe_weights,
         self.padded_base_moe_mlp_dim,
+        self.kv_tp_size,
+        self.moe_mlp_tp_size,
+        self.kv_replication,
     )
 
   def _resolve_target_dtype(self):
@@ -906,9 +930,11 @@ def _fuse_moe_bulk(self, wi_0, wi_1, tgt_val, key_path: str):
     scan_fused_axis = tgt_fused_axis if tgt_fused_axis < self.scan_axis else tgt_fused_axis + 1
 
     if self.moe_fused_layout == MoEFusedLayout.PER_SHARD_INTERLEAVE:
-      n_shards = _get_n_shards(tgt_val, tgt_fused_axis)
-      if n_shards == 1 and self.tp > 1:
-        n_shards = self.tp
+      n_shards = (
+          self.moe_mlp_tp_size
+          if self.moe_mlp_tp_size > 1
+          else (self.tp if self.tp > 1 else _get_n_shards(tgt_val, tgt_fused_axis))
+      )
       return _fuse_and_unstack_moe(
           wi_0,
           wi_1,
@@ -933,6 +959,7 @@ def _fuse_moe_bulk(self, wi_0, wi_1, tgt_val, key_path: str):
   def _slice_bulk_target_free(self, val: Any, path: str):
     """Returns target-free slices for a scanned parameter."""
     last_key = path.split(".")[-1]
+    is_kv = "key.kernel" in path or "value.kernel" in path
     if isinstance(val, jax.ShapeDtypeStruct):
       unrolled_shape = list(val.shape[: self.scan_axis] + val.shape[self.scan_axis + 1 :])
       if last_key in MOE_MLP_WEIGHTS and self.padded_base_moe_mlp_dim is not None:
@@ -942,6 +969,12 @@ def _slice_bulk_target_free(self, val: Any, path: str):
         elif last_key in ("wi_0", "wi_1", "wi"):
           if self.padded_base_moe_mlp_dim > unrolled_shape[-1]:
             unrolled_shape[-1] = self.padded_base_moe_mlp_dim
+      if (
+          is_kv
+          and self.kv_replication > 1
+          and (self.base_num_kv_heads is None or unrolled_shape[-2] == self.base_num_kv_heads)
+      ):
+        unrolled_shape[-2] = unrolled_shape[-2] * self.kv_replication
       return tuple(jax.ShapeDtypeStruct(tuple(unrolled_shape), val.dtype) for _ in range(val.shape[self.scan_axis]))
 
     if last_key in MOE_MLP_WEIGHTS and self.padded_base_moe_mlp_dim is not None:
@@ -960,6 +993,9 @@ def _slice_bulk_target_free(self, val: Any, path: str):
           pad_spec[intermediate_axis] = (0, pad_amount)
           val = jnp.pad(val, pad_spec)
 
+    if is_kv and self.kv_replication > 1 and (self.base_num_kv_heads is None or val.shape[-2] == self.base_num_kv_heads):
+      val = _jit_repeat_axes(val, ((-2, self.kv_replication),))
+
     return _jit_unstack(val, self.scan_axis)
 
   def _fuse_moe_bulk_target_free(self, wi_0: Any, wi_1: Any, path: str):
@@ -979,7 +1015,11 @@ def _fuse_moe_bulk_target_free(self, wi_0: Any, wi_1: Any, path: str):
     scan_fused_axis = tgt_fused_axis if tgt_fused_axis < self.scan_axis else tgt_fused_axis + 1
 
     if self.moe_fused_layout == MoEFusedLayout.PER_SHARD_INTERLEAVE:
-      n_shards = self.tp if self.tp > 1 else _get_n_shards(wi_0, scan_fused_axis)
+      n_shards = (
+          self.moe_mlp_tp_size
+          if self.moe_mlp_tp_size > 1
+          else (self.tp if self.tp > 1 else _get_n_shards(wi_0, scan_fused_axis))
+      )
       return _fuse_and_unstack_moe(
           wi_0,
           wi_1,
@@ -1010,6 +1050,18 @@ def _execute_group_target_free(self, group: _PlanGroup, src_flat):
       raw_val = src_flat[group.source_keys[0]]
       tgt_dt = getattr(raw_val, "dtype", target_dtype) if ("gate" in path or "router" in path) else target_dtype
       val = _apply_dtype_cast(raw_val, tgt_dt, path)
+      is_kv = "key.kernel" in path or "value.kernel" in path
+      if (
+          is_kv
+          and self.kv_replication > 1
+          and (self.base_num_kv_heads is None or val.shape[-2] == self.base_num_kv_heads)
+      ):
+        if isinstance(val, jax.ShapeDtypeStruct):
+          new_shape = list(val.shape)
+          new_shape[-2] = new_shape[-2] * self.kv_replication
+          val = jax.ShapeDtypeStruct(tuple(new_shape), val.dtype)
+        else:
+          val = _jit_repeat_axes(val, ((-2, self.kv_replication),))
       return [(tgt_key, val) for _, tgt_key in group.targets]
 
     if group.op == "fuse_moe":
diff --git a/src/maxtext/training_engine/maxtext_engine.py b/src/maxtext/training_engine/maxtext_engine.py
index d572fcacff..e4c9856f5e 100644
--- a/src/maxtext/training_engine/maxtext_engine.py
+++ b/src/maxtext/training_engine/maxtext_engine.py
@@ -37,6 +37,11 @@
 from maxtext.common import train_state_nnx
 from maxtext.configs import pyconfig
 from maxtext.integration.tunix.weight_mapping import raiden_unscan
+from maxtext.integration.vllm.convert_utils import (
+    is_verify_weights_enabled,
+    resolve_prefuse_moe_weights,
+    resolve_rollout_tp,
+)
 from maxtext.trainers.pre_train import train as maxtext_train
 from maxtext.training_engine import abstract_engine
 from maxtext.training_engine import checkpointing
@@ -146,6 +151,14 @@ def _batch_signature(dynamic_batch: Any, static_batch: dict[str, Any]) -> Any:
   return (treedef, shapes, static_batch)
 
 
+_REPLICATED_BATCH_DIM_WARNING = (
+    "Loss input with batch dim %d does not divide mesh axis %r (size %d), so that "
+    "dimension is replicated instead of sharded: every device along the axis holds and "
+    "computes the whole micro-batch, %dx the work a sharded one would do there. Results "
+    "stay correct. If it was not deliberate -- a sequence-packed micro-batch is always "
+    "size 1 and has no alternative -- make the micro-batch a multiple of the axis size."
+)
+
 _UNCOMPARABLE_SIGNATURE_WARNING = (
     "Could not compare %s between fwd_bwd calls (%s), so the engine cannot tell whether "
     "the compiled kernel is still valid and will recompile on EVERY fwd_bwd from now on. %s"
@@ -557,6 +570,7 @@ def __init__(
     self._compiled_eval: Any = None
     self._compiled_eval_signature: Any = None
     self._signature_compare_warned: bool = False
+    self._replicated_batch_warned: bool = False
     if not training_config.model_name:
       raise ValueError("training_config.model_name must be specified")
     self._model = self._build_model(wrap_with_tunix_adapter, tokenizer_pad_id)
@@ -606,6 +620,28 @@ def __init__(
     self._metrics_logger = metrics_module.MetricsLogger(config=self._config)
     self._throttler = inflight_throttler.InflightThrottler(config=self._config, metrics_logger=self._metrics_logger)
     self._raiden_sync: Any = None
+    self._last_staged_step: Optional[int] = None
+    self._staged_metadata: Any = None
+    self._use_weight_converter = bool(self._config.use_weight_converter)
+    self._rollout_backend = self._config.rollout_backend
+    rollout_tp = resolve_rollout_tp(self._config)
+    kv_tp = getattr(self._config, "kv_tp_size", 1) or rollout_tp
+    moe_tp = getattr(self._config, "moe_mlp_tp_size", 1) or rollout_tp
+    prefuse_moe = resolve_prefuse_moe_weights(self._config)
+    if self._use_weight_converter:
+      from maxtext.integration.vllm.weight_converter import WeightConverter  # pylint: disable=g-import-not-at-top,import-outside-toplevel
+
+      self._weight_converter = WeightConverter(
+          config=self._config,
+          tp=rollout_tp,
+          kv_tp_size=kv_tp,
+          moe_mlp_tp_size=moe_tp,
+          prefuse_moe_weights=prefuse_moe,
+          rollout_backend=self._rollout_backend,
+          debug=self._config.weight_sync_debug,
+      )
+    else:
+      self._weight_converter = None
 
   def _build_model(self, wrap_with_tunix_adapter: bool, tokenizer_pad_id: int | None) -> Any:
     """Returns the model to train, adopting a mesh when this engine was given none."""
@@ -634,7 +670,7 @@ def _build_optimizer(self, tx: Any) -> Any:
 
   def _checkpoint_dir(self) -> str:
     """Returns the directory this engine checkpoints through; an empty string disables Orbax entirely."""
-    return self._config.checkpoint_dir
+    return getattr(self._config, "checkpoint_dir", "") or ""
 
   @property
   def model(self) -> Any:
@@ -1203,7 +1239,11 @@ def _prepare_batch(self, payload: Any) -> Any:
     if self._gen_model_input_fn is not None:
       return self._gen_model_input_fn(payload)
     if dataclasses.is_dataclass(payload):
-      return {k: getattr(payload, k) for k in payload.__dataclass_fields__ if getattr(payload, k) is not None}
+      return {
+          k: getattr(payload, k)
+          for k in payload.__dataclass_fields__
+          if getattr(payload, k) is not None and k != "metadata"
+      }
     return payload
 
   def _mesh_sharding(self, leaf: Any) -> jax.sharding.Sharding | None:
@@ -1247,7 +1287,8 @@ def _batch_data_shardings(self, dynamic_batch: Any) -> Any:
     sequence-packed micro-batch, always size 1) replicates that dim instead of sharding
     it -- every device holds and computes on the same data with no cross-device split,
     which is correct (there's nothing to reduce back together afterwards) but wastes
-    compute across the axis for that micro-batch.
+    compute across the axis for that micro-batch. That is an N-fold cost, so it warns
+    once per instance rather than living only in this docstring.
     """
     data_sharding = sharding.get_input_data_sharding(self._config, self._mesh)
     data_spec = tuple(data_sharding.spec)
@@ -1257,8 +1298,16 @@ def leaf_sharding(leaf):
         return None
       rank = jnp.ndim(leaf)
       spec = list(data_spec[:rank])
-      if spec and spec[0] is not None and leaf.shape[0] % self._batch_axis_size(spec[0]):
-        spec[0] = None
+      if spec and spec[0] is not None:
+        axis_size = self._batch_axis_size(spec[0])
+        if leaf.shape[0] % axis_size:
+          # Warn once per instance, not per leaf: this runs under a tree_map over every
+          # loss input, and they normally share a batch dim. Silence here would leave an
+          # N-fold compute cliff visible only in a docstring.
+          if not self._replicated_batch_warned:
+            self._replicated_batch_warned = True
+            logging.warning(_REPLICATED_BATCH_DIM_WARNING, leaf.shape[0], spec[0], axis_size, axis_size)
+          spec[0] = None
       return jax.sharding.NamedSharding(self._mesh, jax.sharding.PartitionSpec(*spec))
 
     return jax.tree.map(leaf_sharding, dynamic_batch)
@@ -1774,6 +1823,10 @@ def save_checkpoint(self, metadata: Any, **kwargs: Any) -> None:
       metadata: Checkpoint metadata payload from Orchestrator.
       **kwargs: Additional checkpoint saving options.
     """
+    if not getattr(self._config, "enable_checkpointing", True) or not self._checkpoint_dir():
+      logging.info("Checkpointing is disabled in config; skipping save_checkpoint.")
+      return
+
     # Drain all inflight computations and log pending metrics before checkpointing.
     self._throttler.wait_for_all()
 
@@ -2048,7 +2101,7 @@ def prepare_weight_sync(
     """
     if staging_transport == "raiden":
       try:
-        from tunix.experimental.weight_sync import raiden_synchronizer  # pylint: disable=g-import-not-at-top,import-outside-toplevel
+        import tunix.experimental.weight_sync.raiden_synchronizer as raiden_synchronizer  # pylint: disable=g-import-not-at-top,import-outside-toplevel
       except ImportError as exc:
         # Fatal, not a warning: Raiden staging was explicitly requested and cannot be
         # provided. Returning empty metadata instead defers the failure to the caller --
@@ -2061,54 +2114,50 @@ def prepare_weight_sync(
             " tunix build that ships it, or select a different staging_transport."
         ) from exc
 
+      if (
+          self._raiden_sync is not None
+          and self._last_staged_step == self.train_step
+          and self._staged_metadata is not None
+      ):
+        logging.info(
+            "Trainer reusing staged weight sync for step %d (%d variables)",
+            self.train_step,
+            sum(len(m.variables) for m in self._staged_metadata),
+        )
+        return self._staged_metadata
+
       # 1. Drain all in-flight TPU computations to ensure weights are fully updated
       self._throttler.wait_for_all()
 
       # 2. Extract clean trainable parameters
       params_state = self._get_trainable_params_state()
 
-      # 2a. The trainer keeps float32 master weights, but the rollout side
-      # (MaxTextForCausalLM under configs/inference/vllm.yml) loads/serves in
-      # bfloat16 -- Raiden's manifest preflight rejects a dtype/item_size
-      # mismatch, and binding mismatched-dtype buffers would be wrong anyway.
-      # Cast the synced copy down; the trainer's own params_state (used for
-      # the actual optimizer step) is untouched since this is a fresh tree.
-      params_state = jax.tree_util.tree_map(
-          lambda x: x.astype(jnp.bfloat16) if hasattr(x, "dtype") and jnp.issubdtype(x.dtype, jnp.floating) else x,
-          params_state,
-      )
-
-      # 2b. The trainer runs scanned (scan_layers=True) for training speed, but
-      # the rollout side loads its MaxText model unscanned (MaxTextForCausalLM
-      # under configs/inference/vllm.yml has scan_layers=False). Raiden matches
-      # tensors by name, so unscan here -- on the trainer side only -- so the
-      # names/shapes we bind already match what the sampler reports.
-      if self._config.scan_layers:
-        params_state = raiden_unscan.unscan_layers(
+      if self._use_weight_converter:
+        converted_state = self._weight_converter.convert(params_state)
+      else:
+        # UNCHANGED, deliberately out of scope: this fp32->bf16 cast is an
+        # on-device (HBM, not host RAM) full materialization -- a different
+        # memory pool than the host OOM this plan addresses. Candidate
+        # fast-follow: fold into unscan_layers_streaming's per-piece slicing.
+        params_state = jax.tree_util.tree_map(
+            lambda x: x.astype(jnp.bfloat16) if hasattr(x, "dtype") and jnp.issubdtype(x.dtype, jnp.floating) else x,
             params_state,
-            num_layers=self._config.num_decoder_layers,
-            scan_axis=self._config.param_scan_axis,
         )
+        if self._config.scan_layers:
+          converted_state = raiden_unscan.unscan_layers(
+              params_state,
+              num_layers=self._config.num_decoder_layers,
+              scan_axis=self._config.param_scan_axis,
+              cycle_interval=self._config.inhomogeneous_layer_cycle_interval,
+          )
+        else:
+          converted_state = params_state
+
+      del params_state
 
       # 3. Bind parameters to the Raiden transport. Construct the synchronizer
       # once, matching the persistent-instance-per-cycle pattern the rebind
       # optimization depends on.
-      #
-      # Under Pathways (JAX_PLATFORMS=proxy + JAX_BACKEND_TARGET set, same
-      # detection tunix's K8sJaxContext.initialize() uses), trainer params
-      # are proxy-backed. Raiden must use FFI (weight_synchronizer_ffi) to bind
-      # directly to device arrays on Pathways TPU workers without host CPU staging,
-      # avoiding client host OOM and multi-minute proxy transfer timeouts.
-      is_pathways = bool("proxy" in os.environ.get("JAX_PLATFORMS", "") and os.environ.get("JAX_BACKEND_TARGET"))
-      if is_pathways and getattr(raiden_synchronizer, "_raiden_ffi", None) is None:
-        raise RuntimeError(
-            "Under Pathways (JAX_PLATFORMS=proxy), Raiden weight synchronization "
-            "requires weight_synchronizer_ffi (from tpu_raiden_jax) to avoid client host OOM "
-            "and proxy staging timeouts. However, _raiden_ffi is not available in "
-            "tunix.experimental.weight_sync.raiden_synchronizer. Please ensure a "
-            "compatible tpu_raiden_jax wheel with FFI support is installed."
-        )
-
       if self._raiden_sync is None:
         self._raiden_sync = raiden_synchronizer.RaidenSynchronizer(
             job_name="trainer",
@@ -2117,25 +2166,30 @@ def prepare_weight_sync(
             parallelism=4,
         )
 
-      self._raiden_sync.bind(params_state)
-      del params_state
+      self._raiden_sync.bind(converted_state)
+      del converted_state
 
       # 4. Initiate Device-to-Host transfer to stage weights for network transfer.
-      if is_pathways or self._raiden_sync.active:
+      if self._raiden_sync.active:
         self._raiden_sync.d2h()
 
-      verify_weights = os.environ.get("VERIFY_WEIGHTS", "").lower() == "true"
+      verify_weights = is_verify_weights_enabled()
       if verify_weights:
         logging.info("Source weights checksums: %s", self._raiden_sync.checksums())
 
-      metadata = self._raiden_sync.work_unit_metadata()
+      all_metadata = self._raiden_sync.work_unit_metadata_all()
+      total_variables = sum(len(m.variables) for m in all_metadata)
+
       logging.info(
-          "Trainer prepared weight sync for step %d: registered %d variables on mesh %s",
+          "Trainer prepared weight sync for step %d: registered %d work unit(s) with %d variables on mesh %s",
           self.train_step,
-          len(metadata.variables),
-          metadata.mesh_axes,
+          len(all_metadata),
+          total_variables,
+          all_metadata[0].mesh_axes if all_metadata else (),
       )
-      return [metadata]
+      self._last_staged_step = self.train_step
+      self._staged_metadata = all_metadata
+      return all_metadata
 
     # Unknown transport: raise rather than return empty metadata. A typo would otherwise
     # surface only as the coordinator's "empty side" error, with nothing logged anywhere
@@ -2144,6 +2198,8 @@ def prepare_weight_sync(
 
   def release_weight_sync(self, **kwargs: Any) -> Any:
     """Releases staged weight buffers after transfer completion."""
+    self._last_staged_step = None
+    self._staged_metadata = None
     if self._raiden_sync:
       logging.vlog(1, "Trainer Raiden metrics: %s", self._raiden_sync.metrics())
     return True
@@ -2154,9 +2210,12 @@ def close(self) -> None:
       if hasattr(self._raiden_sync, "close"):
         self._raiden_sync.close()
       self._raiden_sync = None
+    self._last_staged_step = None
+    self._staged_metadata = None
 
-    self.save_checkpoint(metadata=None, force=True)
-    self._checkpoint_manager.close()
+    if getattr(self._config, "enable_checkpointing", True) and self._checkpoint_dir() and self._checkpoint_manager:
+      self.save_checkpoint(metadata=None, force=True)
+      self._checkpoint_manager.close()
 
     # Write the metrics and cleanup metrics logger resources
     self._throttler.cleanup()
diff --git a/tests/post_training/unit/convert_utils_test.py b/tests/post_training/unit/convert_utils_test.py
index d751e6f2a7..9f87c20287 100644
--- a/tests/post_training/unit/convert_utils_test.py
+++ b/tests/post_training/unit/convert_utils_test.py
@@ -24,7 +24,10 @@
     _interleave_moe_weights,
     compute_padded_moe_mlp_dim,
     DEFAULT_TPU_NUM_LANES,
+    is_verify_weights_enabled,
     pad_to_tpu_lanes,
+    resolve_prefuse_moe_weights,
+    resolve_rollout_tp,
 )
 
 pytestmark = [pytest.mark.post_training]
@@ -160,6 +163,45 @@ def test_compute_padded_moe_mlp_dim_leaves_aligned_dims_alone(self):
     self.assertEqual(compute_padded_moe_mlp_dim(2560, 2, lanes), 2560)
     self.assertIsNone(compute_padded_moe_mlp_dim(None, 4, lanes))
 
+  def test_is_verify_weights_enabled(self):
+    import os
+
+    orig = os.environ.get("VERIFY_WEIGHTS")
+    try:
+      os.environ["VERIFY_WEIGHTS"] = "true"
+      self.assertTrue(is_verify_weights_enabled())
+      os.environ["VERIFY_WEIGHTS"] = "false"
+      self.assertFalse(is_verify_weights_enabled())
+      os.environ.pop("VERIFY_WEIGHTS", None)
+      self.assertFalse(is_verify_weights_enabled())
+    finally:
+      if orig is not None:
+        os.environ["VERIFY_WEIGHTS"] = orig
+      else:
+        os.environ.pop("VERIFY_WEIGHTS", None)
+
+  def test_resolve_prefuse_moe_weights(self):
+    import os
+
+    self.assertTrue(resolve_prefuse_moe_weights(None, prefuse_moe_weights=True))
+    self.assertFalse(resolve_prefuse_moe_weights(None, prefuse_moe_weights=False))
+
+    class DummyConfig:
+      prefuse_moe_weights = False
+      rollout_backend = "vllm"
+
+    self.assertFalse(resolve_prefuse_moe_weights(DummyConfig()))
+
+  def test_resolve_rollout_tp(self):
+    import os
+
+    self.assertEqual(resolve_rollout_tp(None, tp=4), 4)
+
+    class DummyConfig:
+      rollout_tensor_parallelism = 2
+
+    self.assertEqual(resolve_rollout_tp(DummyConfig()), 2)
+
 
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/post_training/unit/maxtext_engine_e2e_test.py b/tests/post_training/unit/maxtext_engine_e2e_test.py
index 1f2bd0e432..e227351a16 100644
--- a/tests/post_training/unit/maxtext_engine_e2e_test.py
+++ b/tests/post_training/unit/maxtext_engine_e2e_test.py
@@ -16,6 +16,7 @@
 
 from collections.abc import Iterator
 import dataclasses
+import importlib
 from typing import Any
 from unittest import mock
 
@@ -33,6 +34,17 @@
 import pytest
 
 # training_engine imports tunix, so these tests need the post-training dependency bundle.
+# The engine's default staging transport is Raiden, whose synchronizer ships with the
+# RL tunix build and not with stock tunix. Probe once, the same way the engine does, so
+# this loop exercises the real staging path where Raiden exists and the documented
+# failure where it does not -- rather than passing or failing on which tunix happens to
+# be installed.
+try:
+  importlib.import_module("tunix.experimental.weight_sync.raiden_synchronizer")
+  _RAIDEN_AVAILABLE = True
+except ImportError:
+  _RAIDEN_AVAILABLE = False
+
 pytestmark = [pytest.mark.post_training]
 
 
@@ -97,7 +109,13 @@ def run(
       step_metrics = self.trainer.get_metrics(clear_cache=True)
       history.append(step_metrics)
 
-      _ = self.trainer.prepare_weight_sync()
+      if _RAIDEN_AVAILABLE:
+        _ = self.trainer.prepare_weight_sync()
+      else:
+        # Without the transport the engine must raise rather than hand back empty
+        # metadata, which would fail later and far from the cause.
+        with pytest.raises(RuntimeError, match="raiden_synchronizer"):
+          self.trainer.prepare_weight_sync()
 
     self.trainer.close()
     return history
diff --git a/tests/post_training/unit/maxtext_engine_test.py b/tests/post_training/unit/maxtext_engine_test.py
index 60d937218a..76cff42a81 100644
--- a/tests/post_training/unit/maxtext_engine_test.py
+++ b/tests/post_training/unit/maxtext_engine_test.py
@@ -16,6 +16,7 @@
 # pylint: disable=protected-access
 
 import dataclasses
+import sys
 import types
 from typing import Any
 from unittest import mock
@@ -1014,10 +1015,11 @@ def test_shared_types_are_tunix_classes(self):
 
     # What actually arrives at fwd_bwd from GRPOAdapter.create_trainer_payloads.
     rl_payload = datatypes.RLTrainerPayload(
-        token_ids=jnp.zeros((1, 4)),
-        token_mask=jnp.ones((1, 4)),
+        prompt_ids=jnp.zeros((1, 4)),
+        prompt_mask=jnp.ones((1, 4)),
+        completion_ids=jnp.zeros((1, 4)),
+        completion_mask=jnp.ones((1, 4)),
         advantages=jnp.zeros((1,)),
-        loss_mask=jnp.ones((1, 4)),
     )
     self.assertIsInstance(rl_payload, abstract_engine.TrainerPayload)
 
@@ -1349,6 +1351,70 @@ def test_perplexity_is_emitted_alongside_the_loss(self):
     self.assertIn("perplexity", processed)
     self.assertAlmostEqual(processed["perplexity"], float(np.exp(6.0)), places=3)
 
+  def test_prepare_weight_sync_raises_when_raiden_is_unavailable(self):
+    """A missing raiden_synchronizer must fail here, not as an empty result downstream.
+
+    Returning empty metadata defers the failure to `WeightSyncCoordinator`, which raises
+    "metadata collection returned an empty side" -- a count from another process that
+    never names the missing module. `raiden_synchronizer` ships only on tunix's Raiden
+    branch, so this is the common case on a released tunix, not a corner.
+    """
+    t = maxtext_engine.MaxTextTrainingEngine(self.mock_config)
+
+    # Setting the entry to None makes `from ... import raiden_synchronizer` raise
+    # ImportError, which is what an installed tunix without the module does.
+    with mock.patch.dict(sys.modules, {"tunix.experimental.weight_sync.raiden_synchronizer": None}):
+      with self.assertRaisesRegex(RuntimeError, "raiden_synchronizer"):
+        t.prepare_weight_sync()
+
+  def test_prepare_weight_sync_rejects_an_unknown_transport(self):
+    """An unrecognised transport must name itself rather than return empty metadata."""
+    t = maxtext_engine.MaxTextTrainingEngine(self.mock_config)
+    with self.assertRaisesRegex(ValueError, "raidan"):
+      t.prepare_weight_sync(staging_transport="raidan")
+
+  def _sharded_batch_spec(self, engine, axis="data"):
+    """A data sharding whose batch dim is actually sharded.
+
+    The single-device test mesh makes `get_input_data_sharding` return a spec with `None`
+    in the batch position, so the replication branch is unreachable as configured -- an
+    earlier version of these tests asserted `spec[0] is None` and passed without ever
+    running the code under test. Stub a spec that shards the batch dim instead.
+    """
+    return jax.sharding.NamedSharding(engine._mesh, jax.sharding.PartitionSpec(axis, None))  # pylint: disable=protected-access
+
+  def test_indivisible_batch_dim_replicates_and_warns_once(self):
+    """Replicating the batch dim is an N-fold compute cliff, so it must be audible."""
+    t = maxtext_engine.MaxTextTrainingEngine(self.mock_config)
+    batch = {"a": jnp.zeros((1, 4)), "b": jnp.zeros((1, 4))}
+
+    # Batch dim 1 against a 2-wide axis: indivisible, so the dim must be replicated.
+    with mock.patch.object(maxtext_engine.sharding, "get_input_data_sharding", return_value=self._sharded_batch_spec(t)):
+      with mock.patch.object(type(t), "_batch_axis_size", return_value=2):
+        with self.assertLogs(level="WARNING") as logs:
+          shardings = t._batch_data_shardings(batch)  # pylint: disable=protected-access
+          t._batch_data_shardings(batch)  # pylint: disable=protected-access
+
+    for name, leaf_sharding in shardings.items():
+      self.assertIsNone(leaf_sharding.spec[0], f"{name} should have its batch dim replicated")
+
+    # Once per instance, not per leaf and not per call: two leaves over two calls is four
+    # chances to warn.
+    warnings = [line for line in logs.output if "does not divide mesh axis" in line]
+    self.assertLen(warnings, 1)
+    self.assertIn("2x the work", warnings[0])
+
+  def test_divisible_batch_dim_stays_sharded_and_is_silent(self):
+    """The normal case must neither replicate nor warn."""
+    t = maxtext_engine.MaxTextTrainingEngine(self.mock_config)
+
+    with mock.patch.object(maxtext_engine.sharding, "get_input_data_sharding", return_value=self._sharded_batch_spec(t)):
+      with mock.patch.object(type(t), "_batch_axis_size", return_value=2):
+        shardings = t._batch_data_shardings({"a": jnp.zeros((4, 4))})  # pylint: disable=protected-access
+
+    self.assertEqual(shardings["a"].spec[0], "data")
+    self.assertFalse(t._replicated_batch_warned)  # pylint: disable=protected-access
+
 
 if __name__ == "__main__":
   absltest.main()
diff --git a/tests/post_training/unit/weight_converter_test.py b/tests/post_training/unit/weight_converter_test.py
index 6c44344fe7..0e59c7b8cb 100644
--- a/tests/post_training/unit/weight_converter_test.py
+++ b/tests/post_training/unit/weight_converter_test.py
@@ -733,6 +733,57 @@ def test_case_8_weight_converter_convert_streaming_dispatch(self):
     with self.assertRaises(NotImplementedError):
       list(torchax_wc.convert_streaming(source))
 
+  def test_case_9_target_free_kv_head_replication(self):
+    cfg = _config(
+        base_num_kv_heads=2,
+        inhomogeneous_layer_cycle_interval=1,
+        num_decoder_layers=2,
+    )
+    key_val = _scanned(EMB, 2, 4)
+    value_val = _scanned(EMB, 2, 4)
+    source = {
+        "base": {
+            "decoder": {
+                "layers": {
+                    "self_attention": {
+                        "key": {"kernel": key_val},
+                        "value": {"kernel": value_val},
+                    }
+                }
+            }
+        }
+    }
+
+    # 1. Successful replication: kv_tp_size=4, base_num_kv_heads=2 -> kv_replication=2
+    converter = WeightConverter(
+        config=cfg,
+        kv_tp_size=4,
+        rollout_backend="maxtext",
+    )
+    self.assertEqual(converter._direct.kv_replication, 2)
+    out = converter.convert(source, target_state=None)
+    out_root = out["base"] if "base" in out else out
+    for layer_idx in range(2):
+      layer_key = f"layers_{layer_idx}"
+      key_kernel = getattr(
+          out_root["decoder"][layer_key]["self_attention"]["key"]["kernel"],
+          "value",
+          out_root["decoder"][layer_key]["self_attention"]["key"]["kernel"],
+      )
+      self.assertEqual(key_kernel.shape, (EMB, 4, 4))
+      # Heads 0 and 1 are repeated along axis -2
+      np.testing.assert_array_equal(key_kernel[:, 0, :], key_kernel[:, 1, :])
+      np.testing.assert_array_equal(key_kernel[:, 2, :], key_kernel[:, 3, :])
+
+    # 2. Divisibility failure: kv_tp_size=3, base_num_kv_heads=2
+    with self.assertRaises(ValueError) as ctx:
+      WeightConverter(
+          config=cfg,
+          kv_tp_size=3,
+          rollout_backend="maxtext",
+      )
+    self.assertIn("must be divisible by base_num_kv_heads", str(ctx.exception))
+
 
 if __name__ == "__main__":
   unittest.main()
diff --git a/tests/unit/prepare_weight_sync_test.py b/tests/unit/prepare_weight_sync_test.py
new file mode 100644
index 0000000000..7e41d3fd93
--- /dev/null
+++ b/tests/unit/prepare_weight_sync_test.py
@@ -0,0 +1,161 @@
+# Copyright 2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Unit tests for MaxTextTrainingEngine.prepare_weight_sync single synchronizer logic."""
+
+import os
+
+os.environ.setdefault("XLA_FLAGS", "--xla_force_host_platform_device_count=8")
+os.environ.setdefault("JAX_PLATFORMS", "cpu")
+
+
+import sys
+import types as pytypes
+import unittest
+from unittest import mock
+
+# Ensure tunix C-extension / protobuf initializes before transformers/orbax
+try:
+  import tunix.experimental.weight_sync.raiden_synchronizer  # pylint: disable=unused-import
+except ImportError:
+  pass
+
+import jax
+import jax.numpy as jnp
+from maxtext.training_engine.maxtext_engine import MaxTextTrainingEngine
+
+
+class PrepareWeightSyncTest(unittest.TestCase):
+
+  def setUp(self):
+    super().setUp()
+    # Create engine instance without running heavy __init__
+    self.engine = MaxTextTrainingEngine.__new__(MaxTextTrainingEngine)
+    self.engine._raiden_sync = None
+    self.engine._last_staged_step = None
+    self.engine._staged_metadata = None
+    self.engine._train_step = 0
+    self.engine._throttler = mock.MagicMock()
+    self.engine._config = pytypes.SimpleNamespace(
+        scan_layers=False,
+        num_decoder_layers=2,
+        param_scan_axis=1,
+        inhomogeneous_layer_cycle_interval=1,
+        weight_sync_debug=False,
+    )
+    self.engine._use_weight_converter = True
+    self.engine._weight_converter = mock.MagicMock()
+    self.engine._rollout_backend = "maxtext"
+    self.engine._get_trainable_params_state = mock.MagicMock(return_value={"layer": jnp.zeros((4, 4))})
+
+  def _make_dummy_metadata(self, num_vars=2):
+    meta = mock.MagicMock()
+    meta.variables = [f"var_{i}" for i in range(num_vars)]
+    meta.mesh_axes = (1, 1)
+    return meta
+
+  @mock.patch("tunix.experimental.weight_sync.raiden_synchronizer.RaidenSynchronizer")
+  def test_single_synchronizer_creation_and_binding(self, mock_sync_cls):
+    mock_sync = mock.MagicMock()
+    mock_sync.active = True
+    mock_sync.work_unit_metadata_all.return_value = [self._make_dummy_metadata(num_vars=2)]
+    mock_sync.checksums.return_value = {}
+    mock_sync_cls.return_value = mock_sync
+
+    converted = {"param_0": 0, "param_1": 1}
+    self.engine._weight_converter.convert.return_value = converted
+
+    metadata = self.engine.prepare_weight_sync()
+
+    self.assertEqual(len(metadata), 1)
+    self.assertIs(self.engine._raiden_sync, mock_sync)
+    mock_sync_cls.assert_called_once_with(
+        job_name="trainer",
+        worker_index=jax.process_index(),
+        auto_h2d=False,
+        parallelism=4,
+    )
+
+    self.engine._weight_converter.convert.assert_called_once()
+    mock_sync.bind.assert_called_once_with(converted)
+    mock_sync.d2h.assert_called_once()
+    mock_sync.work_unit_metadata_all.assert_called_once()
+
+  @mock.patch("tunix.experimental.weight_sync.raiden_synchronizer.RaidenSynchronizer")
+  def test_rebind_reuses_single_sync_instance(self, mock_sync_cls):
+    mock_sync = mock.MagicMock()
+    mock_sync.active = True
+    mock_sync.work_unit_metadata_all.return_value = [self._make_dummy_metadata(num_vars=2)]
+    mock_sync.checksums.return_value = {}
+    mock_sync_cls.return_value = mock_sync
+
+    # Round 1
+    self.engine._weight_converter.convert.return_value = {"p0": 0}
+    self.engine.prepare_weight_sync()
+    self.assertEqual(mock_sync_cls.call_count, 1)
+
+    # Round 2 at step 1
+    self.engine._train_step = 1
+    self.engine._weight_converter.convert.return_value = {"p0": 0}
+    self.engine.prepare_weight_sync()
+
+    # Still only 1 synchronizer instance created
+    self.assertEqual(mock_sync_cls.call_count, 1)
+    self.assertEqual(mock_sync.bind.call_count, 2)
+
+  def test_release_weight_sync(self):
+    mock_sync = mock.MagicMock()
+    self.engine._raiden_sync = mock_sync
+    self.engine._last_staged_step = 1
+    self.engine._staged_metadata = [{"metadata": "dummy"}]
+
+    res = self.engine.release_weight_sync()
+
+    self.assertTrue(res)
+    self.assertIsNone(self.engine._last_staged_step)
+    self.assertIsNone(self.engine._staged_metadata)
+    mock_sync.metrics.assert_called_once()
+
+  def test_release_weight_sync_without_syncs(self):
+    self.engine._raiden_sync = None
+    self.engine._last_staged_step = 1
+    self.engine._staged_metadata = [{"metadata": "dummy"}]
+
+    res = self.engine.release_weight_sync()
+
+    self.assertTrue(res)
+    self.assertIsNone(self.engine._last_staged_step)
+    self.assertIsNone(self.engine._staged_metadata)
+
+  def test_close(self):
+    mock_sync = mock.MagicMock()
+    self.engine._raiden_sync = mock_sync
+    self.engine._last_staged_step = 1
+    self.engine._staged_metadata = [{"metadata": "dummy"}]
+    self.engine.save_checkpoint = mock.MagicMock()
+    self.engine._checkpoint_manager = mock.MagicMock()
+    self.engine._throttler = mock.MagicMock()
+    self.engine._metrics_recorder = mock.MagicMock()
+    self.engine._metrics_logger = mock.MagicMock()
+
+    self.engine.close()
+
+    mock_sync.close.assert_called_once()
+    self.assertIsNone(self.engine._raiden_sync)
+    self.assertIsNone(self.engine._last_staged_step)
+    self.assertIsNone(self.engine._staged_metadata)
+
+
+if __name__ == "__main__":
+  unittest.main()