From e83266e88eb18d6e2205272ac0c1e1a4ccc7e092 Mon Sep 17 00:00:00 2001
From: yuanlehome <yuanlehome@163.com>
Date: Mon, 19 Oct 2026 17:03:26 +0800
Subject: [PATCH 1/6] perf: accelerate Qwen3.5 OpenR1MM training

Reduce colocated GRPO overhead with a guarded inline-logprob path and
BF16 expert gathering, asynchronous TP/GDN gathers, packed experts and
TP-local direct loading.

Enable DeepEP, CUTLASS grouped GEMM and Rust sgl-router for OpenR1MM;
set global batch 512 and packing 8192. Avoid the intermediate ASCII
copy when decoding R3 routing data.

Preserve trainer offload/reload, process-group reconstruction and full
weight synchronization. Add regression coverage for the optimized paths.

Validation:
- Full pre-commit checks; targeted CPU tests: 259 passed, 2 skipped
- Five full batch512/packing8192 cycles without OOM or allocation retries
- Packing8192 numerical differences accepted after frozen-input replay
---
 relax/backends/megatron/actor.py              |  88 ++++++-
 relax/backends/megatron/loss.py               |   5 +
 .../weight_update/bf16_expert_gather.py       | 109 +++++++++
 .../backends/megatron/weight_update/common.py |  15 +-
 .../hf_weight_iterator_bridge.py              | 117 +++++++++-
 .../weight_update/qwen35_direct_experts.py    | 125 ++++++++++
 .../weight_update/qwen35_packed_experts.py    | 146 ++++++++++++
 .../update_weight_from_tensor.py              |  85 ++++++-
 relax/engine/rollout/sglang_rollout.py        |   2 +-
 .../run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh   |  49 +++-
 .../test_colocate_inline_log_probs.py         | 214 ++++++++++++++++++
 .../test_async_gdn_weight_gather.py           |  85 +++++++
 .../weight_update/test_bf16_expert_gather.py  | 128 +++++++++++
 .../weight_update/test_expert_routing.py      |   2 +
 .../test_qwen35_direct_experts.py             | 143 ++++++++++++
 .../test_qwen35_packed_experts.py             | 155 +++++++++++++
 16 files changed, 1443 insertions(+), 25 deletions(-)
 create mode 100644 relax/backends/megatron/weight_update/bf16_expert_gather.py
 create mode 100644 relax/backends/megatron/weight_update/qwen35_direct_experts.py
 create mode 100644 relax/backends/megatron/weight_update/qwen35_packed_experts.py
 create mode 100644 tests/backends/megatron/test_colocate_inline_log_probs.py
 create mode 100644 tests/backends/megatron/weight_update/test_async_gdn_weight_gather.py
 create mode 100644 tests/backends/megatron/weight_update/test_bf16_expert_gather.py
 create mode 100644 tests/backends/megatron/weight_update/test_qwen35_direct_experts.py
 create mode 100644 tests/backends/megatron/weight_update/test_qwen35_packed_experts.py

diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py
index d2c203818..dcbbf2057 100644
--- a/relax/backends/megatron/actor.py
+++ b/relax/backends/megatron/actor.py
@@ -1920,6 +1920,84 @@ def train_critic(self, rollout_id: int, rollout_data: RolloutBatch) -> None:
         if self._per_step_rollout and self.args.offload_train:
             self.sleep()
 
+    def _inline_colocate_old_log_probs(self, num_microbatches: list[int]) -> bool:
+        """Validate the single-update GRPO path before deriving old log-probs
+        inline."""
+        args = self.args
+        if os.environ.get("RELAX_COLOCATE_INLINE_LOG_PROBS") != "1":
+            return False
+
+        # A second update, stale rollout, or mutable forward would invalidate old=current.
+        # Fail closed: running the old forward cannot undo true_on_policy_mode in the loss.
+        required = {
+            "colocate": True,
+            "true_on_policy_mode": True,
+            "fully_async": False,
+            "hybrid": False,
+            "advantage_estimator": "grpo",
+            "loss_type": "policy_loss",
+            "compute_advantages_and_returns": True,
+            "num_iters_per_train_update": 1,
+            "max_staleness": 0,
+            "update_weights_interval": 1,
+            "partial_rollout": False,
+            "rl_partial_rollouts": False,
+            "use_dynamic_global_batch_size": False,
+            "context_parallel_size": 1,
+            "pipeline_model_parallel_size": 1,
+            "dynamic_context_parallel": False,
+            "hidden_dropout": 0.0,
+            "attention_dropout": 0.0,
+            "bf16": True,
+            "fp16": False,
+            "fp8": None,
+            "fp4": None,
+            "fp8_param_gather": False,
+            "moe_input_jitter_eps": None,
+            "moe_router_enable_expert_bias": False,
+            "moe_router_force_load_balancing": False,
+            "moe_router_force_biased": None,
+            "moe_expert_capacity_factor": None,
+            "moe_aux_loss_coeff": 0.0,
+            "kl_coef": 0.0,
+            "use_kl_loss": False,
+            "use_opd": False,
+            "use_critic": False,
+            "keep_old_actor": False,
+            "use_rollout_logprobs": False,
+            "use_routing_replay": True,
+            "use_rollout_routing_replay": True,
+            "use_rollout_indexer_replay": False,
+            "enable_mtp_training": False,
+            "use_opsm": False,
+            "use_tis": True,
+            "calculate_per_token_loss": True,
+            "entropy_coef": 0.0,
+            "recompute_loss_function": False,
+            "custom_tis_function_path": None,
+            "custom_pg_loss_reducer_function_path": None,
+            "custom_loss_function_path": None,
+            "custom_model_provider_path": None,
+            "custom_megatron_before_log_prob_hook_path": None,
+            "custom_megatron_before_train_step_hook_path": None,
+        }
+        invalid = [name for name, expected in required.items() if getattr(args, name, None) != expected]
+        if args.rollout_batch_size * args.n_samples_per_prompt != args.global_batch_size:
+            invalid.append("rollout size must equal global_batch_size")
+        if len(num_microbatches) != 1:
+            invalid.append("exactly one optimizer step per rollout is required")
+        if getattr(args, "num_steps_per_rollout", None) not in (None, 1):
+            invalid.append("num_steps_per_rollout")
+        if args.moe_router_load_balancing_type not in ("none", ["none"]):
+            invalid.append("moe_router_load_balancing_type")
+        if set(self.weights_backuper.backup_tags) - {"actor"}:
+            invalid.append("reference, teacher, or old-actor weights")
+        if self.rollout_data_postprocess is not None:
+            invalid.append("rollout_data_postprocess")
+        if invalid:
+            raise ValueError("Unsupported colocate true-on-policy configuration: " + ", ".join(invalid))
+        return True
+
     def train_actor(
         self,
         rollout_id: int,
@@ -1954,6 +2032,10 @@ def train_actor(
             data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
         else:
             data_iterator, num_microbatches = prepared_data_iterator, prepared_num_microbatches
+        inline_old_log_probs = os.environ.get(
+            "RELAX_COLOCATE_INLINE_LOG_PROBS"
+        ) == "1" and self._inline_colocate_old_log_probs(num_microbatches)
+        self.args._colocate_inline_log_probs = inline_old_log_probs
         # Create a separate iterator with a larger token budget for ref/teacher log-probs
         if (
             not is_offline_mode(self.args)
@@ -2020,7 +2102,11 @@ def train_actor(
                     )
 
                 self._switch_model("old_actor" if self.args.keep_old_actor else "actor")
-                if not standard_dpo and (not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics):
+                if (
+                    not standard_dpo
+                    and not inline_old_log_probs
+                    and (not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics)
+                ):
                     if self.args.use_routing_replay:
                         if self.args.use_rollout_routing_replay:
                             os.environ["ROUTING_REPLAY_STAGE"] = "replay_forward"
diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py
index a0e5c177e..ae8fcabb2 100644
--- a/relax/backends/megatron/loss.py
+++ b/relax/backends/megatron/loss.py
@@ -611,6 +611,11 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch)
     if args.kl_coef == 0 or not log_probs:
         # when kl_coef is 0, we won't compute ref_log_prob
         xs = log_probs if log_probs is not None else values
+        if xs is None and getattr(args, "_colocate_inline_log_probs", False) and args.kl_coef == 0:
+            # Only shape/device are used. Sampling probabilities are never PPO old log-probs.
+            xs = rollout_data.get("rollout_log_probs")
+            if xs is None:
+                raise ValueError("Colocate true-on-policy GRPO requires rollout log-probs for zero-KL shapes.")
         kl = [torch.zeros_like(x, dtype=torch.float32, device=x.device) for x in xs]
     else:
         kl = [
diff --git a/relax/backends/megatron/weight_update/bf16_expert_gather.py b/relax/backends/megatron/weight_update/bf16_expert_gather.py
new file mode 100644
index 000000000..7f6b3b04b
--- /dev/null
+++ b/relax/backends/megatron/weight_update/bf16_expert_gather.py
@@ -0,0 +1,109 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+"""Owner-balanced BF16 expert export, following Miles' packed gather design.
+
+Layout descriptors survive updates; GPU buffers and current process groups do
+not. The caller supplies the live EP group after each offload/reload cycle.
+"""
+
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+import torch
+import torch.distributed as dist
+
+
+if TYPE_CHECKING:
+    from relax.utils.types import ParamInfo
+
+
+NamedTensor = tuple[str, torch.Tensor]
+TensorLayout = tuple[str, tuple[int, ...], int]
+
+
+def bucket_experts_by_owner(
+    infos: Sequence["ParamInfo"], buffer_size: int, ranks: Sequence[int]
+) -> list[list["ParamInfo"]]:
+    """Pack each owner's share so one gathered round stays within the
+    budget."""
+    budget = buffer_size // len(ranks)
+    owners: dict[int, list] = {rank: [] for rank in ranks}
+    for info in infos:
+        if info.src_rank not in owners or not 0 < info.size <= budget:
+            raise ValueError(f"Unsupported expert owner or byte size: {info.name}")
+        owners[info.src_rank].append(info)
+    owner_batches = []
+    for rank in ranks:
+        batches, pending, size = [], [], 0
+        for info in owners[rank]:
+            if pending and size + info.size > budget:
+                batches.append(pending)
+                pending, size = [], 0
+            pending.append(info)
+            size += info.size
+        if pending:
+            batches.append(pending)
+        owner_batches.append(batches)
+    return [
+        [info for batches in owner_batches if index < len(batches) for info in batches[index]]
+        for index in range(max(map(len, owner_batches), default=0))
+    ]
+
+
+def _layout(tensors: Sequence[NamedTensor]) -> tuple[TensorLayout, ...]:
+    if any(tensor.dtype != torch.bfloat16 for _, tensor in tensors):
+        raise ValueError("Packed BF16 expert gather requires BF16 tensors")
+    return tuple((name, tuple(tensor.shape), tensor.numel()) for name, tensor in tensors)
+
+
+class Bf16ExpertGather:
+    """Cache CPU layouts; allocate independent receive storage on every
+    call."""
+
+    def __init__(self) -> None:
+        self._ranks: tuple[int, ...] | None = None
+        self._layouts: tuple[tuple[TensorLayout, ...], ...] | None = None
+        self._sizes: tuple[int, ...] = ()
+
+    def __call__(
+        self, tensors: Sequence[NamedTensor], *, group: dist.ProcessGroup, device: torch.device
+    ) -> list[NamedTensor]:
+        ranks = tuple(dist.get_process_group_ranks(group))
+        local_index = dist.get_rank(group=group)
+        layout = _layout(tensors)
+        if self._layouts is None:
+            layouts = [None] * len(ranks)
+            dist.all_gather_object(layouts, layout, group=group)
+            self._layouts = tuple(layouts)
+            self._ranks = ranks
+            self._sizes = tuple(sum(entry[2] for entry in peer) for peer in self._layouts)
+        elif ranks != self._ranks or layout != self._layouts[local_index]:
+            raise ValueError("Expert ownership or tensor layout changed; rebuild the gather plan")
+
+        storage = torch.empty(sum(self._sizes), dtype=torch.bfloat16, device=device)
+        payloads = list(storage.split(self._sizes))
+        local = payloads[local_index]
+        offset = 0
+        for (_, tensor), (_, _, size) in zip(tensors, layout, strict=True):
+            local.narrow(0, offset, size).copy_(tensor.reshape(-1))
+            offset += size
+
+        sources = [index for index, size in enumerate(self._sizes) if size]
+        handle = None
+        if len(sources) == 1:
+            src = sources[0]
+            handle = dist.broadcast(payloads[src], src=ranks[src], group=group, async_op=True)
+        elif sources and len(set(self._sizes)) == 1:
+            handle = dist.all_gather_into_tensor(storage, local, group=group, async_op=True)
+        elif sources:
+            handle = dist.all_gather(payloads, local, group=group, async_op=True)
+
+        result = []
+        for peer, payload in zip(self._layouts, payloads, strict=True):
+            offset = 0
+            for name, shape, size in peer:
+                result.append((name, payload.narrow(0, offset, size).view(shape)))
+                offset += size
+        if handle is not None:
+            handle.wait()
+        return result
diff --git a/relax/backends/megatron/weight_update/common.py b/relax/backends/megatron/weight_update/common.py
index 5b84170e1..277322a51 100644
--- a/relax/backends/megatron/weight_update/common.py
+++ b/relax/backends/megatron/weight_update/common.py
@@ -3,7 +3,7 @@
 import inspect
 import re
 from argparse import Namespace
-from collections.abc import Iterator, Sequence
+from collections.abc import Iterator, Mapping, Sequence
 
 import torch
 import torch.distributed as dist
@@ -133,6 +133,8 @@ def all_gather_param(args, name: str, param: torch.Tensor) -> torch.Tensor:
 
 def all_gather_params_async(
     param_infos_and_params: list[tuple[ParamInfo, torch.Tensor]],
+    *,
+    gdn_split_sizes: Mapping[str, tuple[int, ...]] | None = None,
 ) -> list[torch.Tensor]:
     """Parallel TP all-gather for multiple params.
 
@@ -185,6 +187,17 @@ def all_gather_params_async(
             # Process the gathered partitions (same logic as original all_gather_param)
             assert partition_dim is not None, "partition_dim must be set for TP-sharded params"
             name = info.name.replace(".to_wrap.", ".")
+            if gdn_split_sizes is not None and name in gdn_split_sizes:
+                # Each rank holds [Q_local,K_local,V_local,(Z_local,B_local,A_local)].
+                # Gather the complete shard once, then restore component-major HF
+                # order. Plain concatenation would interleave Q/K/V across ranks.
+                assert partition_dim == 0, "GDN components must be sharded along rows"
+                components = [p.split(gdn_split_sizes[name], dim=0) for p in param_partitions]
+                param = torch.cat(
+                    [parts[index] for index in range(len(gdn_split_sizes[name])) for parts in components], dim=0
+                )
+                gathered_params.append(param)
+                continue
             # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better?
             # TODO: check only GLU is used.
             if "linear_fc1.weight" in name and "vision_model" not in name:
diff --git a/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py b/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py
index bf9723fdc..5440f84d0 100644
--- a/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py
+++ b/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py
@@ -1,6 +1,7 @@
 # Copyright (c) 2026 Relax Authors. All Rights Reserved.
 
 import dataclasses
+import os
 import re
 import time
 from collections import OrderedDict
@@ -27,12 +28,15 @@
     is_lora_enabled,
     is_lora_merge_mode,
 )
+from relax.utils.misc import get_hf_config
 from relax.utils.types import ParamInfo
 
+from .bf16_expert_gather import Bf16ExpertGather, bucket_experts_by_owner
 from .bridge_converter import BridgeConverter
-from .common import all_gather_param, named_params_and_buffers
+from .common import all_gather_param, all_gather_params_async, named_params_and_buffers
 from .expert_routing import ExpertLayout, ExpertRouter
 from .hf_weight_iterator_base import HfWeightIteratorBase
+from .qwen35_packed_experts import Qwen35PackedExpertGather, bucket_qwen35_experts
 
 
 logger = get_logger(__name__)
@@ -103,6 +107,62 @@ def __init__(self, *args, **kwargs):
         self._expert_buckets, self._non_expert_buckets, self._vanilla_key_map, self._adapter_map = buckets_result
         self._expert_broadcast_caches = [({}, {}) for _ in self._expert_buckets]
         self._expert_router: ExpertRouter | None = None
+        self._async_weight_gather = (
+            os.environ.get("RELAX_ASYNC_WEIGHT_GATHER", "0") == "1"
+            and torch.cuda.is_available()
+            and not self.quantization_config
+            and not self.lora_merge_mode
+            and not self.lora_adapter_mode
+        )
+        self._gdn_split_sizes = {}
+        if self._async_weight_gather:
+            names = {
+                info.name.replace(".to_wrap.", ".")
+                for bucket in self._non_expert_buckets
+                for info in bucket
+                if "self_attention.conv1d.weight" in info.name or "self_attention.in_proj.weight" in info.name
+            }
+            if names:
+                config = get_hf_config(self.args.hf_checkpoint).text_config
+                tp_size = mpu.get_tensor_model_parallel_world_size()
+                qk = config.linear_key_head_dim * (config.linear_num_key_heads // tp_size)
+                value = config.linear_value_head_dim * (config.linear_num_value_heads // tp_size)
+                heads = config.linear_num_value_heads // tp_size
+                self._gdn_split_sizes = {
+                    name: (qk, qk, value) if "conv1d" in name else (qk, qk, value, value, heads, heads)
+                    for name in names
+                }
+        self._bf16_expert_buckets = None
+        self._bf16_expert_gathers = []
+        # Opt in only for the validated, unquantized ETP1/PP1 export. Other
+        # layouts retain the existing Bridge path, including LoRA and routing.
+        if (
+            os.environ.get("RELAX_BF16_EXPERT_GATHER", "0") == "1"
+            and torch.cuda.is_available()
+            and not self.quantization_config
+            and not self.lora_merge_mode
+            and not self.lora_adapter_mode
+            and mpu.get_pipeline_model_parallel_world_size() == 1
+            and mpu.get_expert_tensor_parallel_world_size() == 1
+            and mpu.get_expert_model_parallel_world_size() > 1
+            and self._expert_buckets
+            and all(info.dtype == torch.bfloat16 for bucket in self._expert_buckets for info in bucket)
+        ):
+            infos = [info for bucket in self._expert_buckets for info in bucket]
+            ranks = dist.get_process_group_ranks(mpu.get_expert_model_parallel_group())
+            budget = self.args.update_weight_buffer_size // len(ranks)
+            if all(info.src_rank in ranks and 0 < info.size <= budget for info in infos):
+                self._bf16_expert_buckets = bucket_experts_by_owner(infos, self.args.update_weight_buffer_size, ranks)
+                self._bf16_expert_gathers = [Bf16ExpertGather() for _ in self._bf16_expert_buckets]
+                if (
+                    os.environ.get("RELAX_QWEN35_PACKED_EXPERTS", "0") == "1"
+                    and mpu.get_virtual_pipeline_model_parallel_world_size() in (None, 1)
+                    and get_hf_config(self.args.hf_checkpoint).model_type == "qwen3_5_moe"
+                ):
+                    self._bf16_expert_buckets = bucket_qwen35_experts(infos, self.args.num_experts, ranks)
+                    self._bf16_expert_gathers = [
+                        Qwen35PackedExpertGather(self.args.num_experts) for _ in self._bf16_expert_buckets
+                    ]
         if self.lora_merge_mode:
             self.lora_alpha = self.args.lora_alpha
             self.lora_dim = self.args.lora_rank
@@ -216,6 +276,21 @@ def _iter_routed_expert_chunks(
             )
 
     def get_hf_weight_chunks(self, megatron_local_weights):
+        if self._bf16_expert_buckets is not None:
+            if isinstance(self._bf16_expert_gathers[0], Qwen35PackedExpertGather):
+                yield from _chunk_with_mla_pairing(
+                    self._iter_hf_params(megatron_local_weights, include_experts=False),
+                    chunk_size=self.args.update_weight_buffer_size,
+                )
+                # Each native projection is indivisible. Yield it immediately
+                # instead of materializing the next layer through chunk lookahead.
+                for item in self._iter_bf16_expert_params(megatron_local_weights, include_non_experts=False):
+                    yield [item]
+                return
+            yield from _chunk_with_mla_pairing(
+                self._iter_bf16_expert_params(megatron_local_weights), chunk_size=self.args.update_weight_buffer_size
+            )
+            return
         if self._expert_router is not None:
             yield from self._iter_routed_expert_chunks(megatron_local_weights)
         yield from _chunk_with_mla_pairing(
@@ -223,6 +298,29 @@ def get_hf_weight_chunks(self, megatron_local_weights):
             chunk_size=self.args.update_weight_buffer_size,
         )
 
+    def _iter_bf16_expert_params(self, megatron_local_weights, *, include_non_experts=True):
+        # Preserve the current GDN and vision conversion; only expert ownership
+        # and transport differ. ETP1 owners already hold complete expert tensors.
+        if include_non_experts:
+            yield from self._iter_hf_params(megatron_local_weights, include_experts=False)
+        group = mpu.get_expert_model_parallel_group()
+        rank = dist.get_rank(group=dist.group.WORLD)
+        device = device_utils.make_current_torch_device()
+        for infos, gather in zip(self._bf16_expert_buckets, self._bf16_expert_gathers, strict=True):
+            converted = []
+            for info in infos:
+                if info.src_rank == rank:
+                    tensor = (
+                        megatron_local_weights[self._vanilla_key_map[info.name]]
+                        .detach()
+                        .to(device=device, non_blocking=_NON_BLOCKING)
+                    )
+                    converted.extend(self._bridge_converter.convert(info.name, tensor))
+                    del tensor
+            gathered = gather(converted, group=group, device=device)
+            yield from gathered
+            del converted, gathered
+
     def _iter_hf_params(self, megatron_local_weights, *, include_experts=True):
         """Yield individual (name, tensor) pairs for all params.
 
@@ -282,9 +380,20 @@ def _iter_hf_params(self, megatron_local_weights, *, include_experts=True):
             t_b1 = time.monotonic()
             t_bcast_total += t_b1 - t_b0
 
-            for info, param in zip(bucket_infos, params, strict=True):
+            gathered_params = None
+            if self._async_weight_gather:
                 t_g0 = time.monotonic()
-                gathered = all_gather_param(self.args, info.name, param)
+                gathered_params = all_gather_params_async(
+                    list(zip(bucket_infos, params, strict=True)), gdn_split_sizes=self._gdn_split_sizes
+                )
+                t_gather_total += time.monotonic() - t_g0
+            for index, (info, param) in enumerate(zip(bucket_infos, params, strict=True)):
+                t_g0 = time.monotonic()
+                gathered = (
+                    all_gather_param(self.args, info.name, param)
+                    if gathered_params is None
+                    else gathered_params[index]
+                )
                 t_g1 = time.monotonic()
                 t_gather_total += t_g1 - t_g0
 
@@ -295,7 +404,7 @@ def _iter_hf_params(self, megatron_local_weights, *, include_experts=True):
                 yield from converted
                 del gathered, converted
 
-            del params
+            del params, gathered_params
 
         if rank == 0:
             logger.info(
diff --git a/relax/backends/megatron/weight_update/qwen35_direct_experts.py b/relax/backends/megatron/weight_update/qwen35_direct_experts.py
new file mode 100644
index 000000000..a5b4cada9
--- /dev/null
+++ b/relax/backends/megatron/weight_update/qwen35_direct_experts.py
@@ -0,0 +1,125 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+"""Native BF16 expert placement for the validated Qwen3.5 TP2 Triton layout."""
+
+import json
+import re
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any
+
+import torch
+
+
+NamedTensor = tuple[str, torch.Tensor]
+
+
+_PACKED = re.compile(r"model\.language_model\.layers\.(\d+)\.mlp\.experts\.(gate_up_proj|down_proj)")
+
+
+def validate_direct_receiver(config: Mapping[str, Any]) -> None:
+    required = dict(
+        tp_size=2,
+        ep_size=1,
+        pp_size=1,
+        moe_dp_size=1,
+        moe_runner_backend="triton",
+        disable_shared_experts_fusion=True,
+        dtype="bfloat16",
+    )
+    if any(config.get(key) != value for key, value in required.items()):
+        raise ValueError("Direct Qwen3.5 experts require explicit BF16 Triton, TP2/EP1/PP1 and unfused shared experts")
+    if (
+        config.get("enable_eplb")
+        or config.get("ep_num_redundant_experts", 0)
+        or config.get("init_expert_location") not in (None, "trivial")
+        or config.get("ep_join_mode") is not None
+        or config.get("quantization") is not None
+        or config.get("enable_lora")
+        or config.get("speculative_algorithm") is not None
+        or config.get("enforce_shared_experts_fusion")
+        or json.loads(config.get("json_model_override_args") or "{}")
+    ):
+        raise ValueError("Direct Qwen3.5 experts do not support quantization, remapping, LoRA or model overrides")
+
+
+def _shard_expert(tensor: torch.Tensor, projection: str, tp_rank: int) -> torch.Tensor:
+    if tp_rank not in (0, 1) or tensor.ndim != 3 or tensor.dtype != torch.bfloat16:
+        raise ValueError("Direct expert shards require BF16 rank-3 tensors and TP rank 0/1")
+    if projection == "gate_up_proj":
+        if tensor.shape[1] % 4:
+            raise ValueError("Gate/up rows must be divisible by four for TP2")
+        intermediate = tensor.shape[1] // 2
+        width = intermediate // 2
+        start = tp_rank * width
+        # TP slices gate and up separately, then keeps [gate_local, up_local].
+        return torch.cat(
+            (tensor[:, start : start + width], tensor[:, intermediate + start : intermediate + start + width]), dim=1
+        )
+    if projection == "down_proj":
+        if tensor.shape[2] % 2:
+            raise ValueError("Down columns must be divisible by two for TP2")
+        width = tensor.shape[2] // 2
+        return tensor[:, :, tp_rank * width : (tp_rank + 1) * width].contiguous()
+    raise ValueError(f"Unsupported expert projection: {projection}")
+
+
+def direct_expert_chunk(named_tensors: Sequence[NamedTensor], tp_rank: int) -> list[NamedTensor] | None:
+    if len(named_tensors) != 1:
+        return None
+    name, tensor = named_tensors[0]
+    match = _PACKED.fullmatch(name)
+    if match is None:
+        return None
+    layer, projection = match.groups()
+    expected = (256, 1024, 2048) if projection == "gate_up_proj" else (256, 2048, 512)
+    if tuple(tensor.shape) != expected:
+        raise ValueError(f"Unsupported Qwen3.5 expert shape: {name}: {tuple(tensor.shape)}")
+    target = "w13_weight" if projection == "gate_up_proj" else "w2_weight"
+    return [(f"model.layers.{layer}.mlp.experts.{target}", _shard_expert(tensor, projection, tp_rank))]
+
+
+def dispatch_typed_buckets(
+    packets: Sequence[Sequence[tuple[str, str]]],
+    *,
+    engine: Any,
+    version: str | int | None,
+    serializer: Any,
+    empty_flattened: Callable[[], dict[str, Any]],
+    held: list[Any],
+    dispatch_errors: list[str] | None = None,
+) -> list[Any]:
+    """Keep a single request format even with empty or differently packed
+    ranks.
+
+    Only enabled by the private direct-expert path. Format tags ride in the
+    existing Gloo gather; no new collective is needed. GPU payloads are already
+    held by their producers through the caller's normal ACK pipeline.
+    """
+    refs, empty = [], {}
+    count = max((len(packet) for packet in packets), default=0)
+    for index in range(count):
+        formats = dict.fromkeys(packet[index][0] for packet in packets if index < len(packet))
+        for load_format in formats:
+            if load_format not in ("direct", "flattened_bucket"):
+                raise ValueError(f"Unsupported IPC format: {load_format}")
+            serialized = []
+            for packet in packets:
+                if index < len(packet) and packet[index][0] == load_format:
+                    serialized.append(packet[index][1])
+                else:
+                    if load_format not in empty:
+                        payload = [] if load_format == "direct" else empty_flattened()
+                        held.append(payload)
+                        empty[load_format] = serializer.serialize(payload, output_str=True)
+                    serialized.append(empty[load_format])
+            kwargs = dict(serialized_named_tensors=serialized, load_format=load_format)
+            if version is not None:
+                kwargs["weight_version"] = str(version)
+            try:
+                refs.append(engine.update_weights_from_tensor.remote(**kwargs))
+            except Exception as error:
+                if dispatch_errors is None:
+                    raise
+                dispatch_errors.append(str(error))
+                return refs
+    return refs
diff --git a/relax/backends/megatron/weight_update/qwen35_packed_experts.py b/relax/backends/megatron/weight_update/qwen35_packed_experts.py
new file mode 100644
index 000000000..172889c8c
--- /dev/null
+++ b/relax/backends/megatron/weight_update/qwen35_packed_experts.py
@@ -0,0 +1,146 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+"""Export Qwen3.5 BF16 experts in the checkpoint's native packed format.
+
+Each layer/projection is indivisible: a gate/up tensor may exceed the target
+IPC bucket size. Only CPU layout metadata is cached across weight updates.
+"""
+
+import re
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+import torch
+import torch.distributed as dist
+
+
+if TYPE_CHECKING:
+    from relax.utils.types import ParamInfo
+
+
+NamedTensor = tuple[str, torch.Tensor]
+TensorLayout = tuple[str, tuple[int, ...]]
+
+
+_SOURCE = re.compile(r"^(.*\.language_model\.decoder\.layers\.(\d+)\.mlp\.experts)\.linear_fc([12])\.weight(\d+)$")
+_TARGET = re.compile(r"^(model\.language_model\.layers\.\d+\.mlp\.experts)\.(\d+)\.(gate|up|down)_proj\.weight$")
+
+
+def bucket_qwen35_experts(
+    infos: Sequence["ParamInfo"], num_experts: int, ranks: Sequence[int]
+) -> list[list["ParamInfo"]]:
+    """Validate complete, contiguous EP ownership and group by
+    layer/projection."""
+    if not ranks or num_experts <= 0 or num_experts % len(ranks):
+        raise ValueError("Packed Qwen experts require evenly partitioned EP ownership")
+    local_count = num_experts // len(ranks)
+    groups = {}
+    for info in infos:
+        match = _SOURCE.fullmatch(info.name)
+        if match is None:
+            raise ValueError(f"Unsupported packed expert source: {info.name}")
+        prefix, layer, projection, expert = match.groups()
+        expert = int(expert)
+        if (
+            not 0 <= expert < num_experts
+            or info.src_rank != ranks[expert // local_count]
+            or info.dtype != torch.bfloat16
+            or len(info.shape) != 2
+            or min(info.shape) <= 0
+        ):
+            raise ValueError(f"Invalid packed expert ownership, dtype or shape: {info.name}")
+        group = groups.setdefault((prefix, int(layer), int(projection)), {})
+        if expert in group:
+            raise ValueError(f"Duplicate packed expert: {info.name}")
+        group[expert] = info
+    buckets = []
+    for key in sorted(groups, key=lambda key: (key[1], key[0], key[2])):
+        group = groups[key]
+        if set(group) != set(range(num_experts)) or len({tuple(info.shape) for info in group.values()}) != 1:
+            raise ValueError(f"Incomplete or inconsistent packed experts: {key[0]}")
+        other = groups.get((*key[:2], 3 - key[2]))
+        if other is None or 0 not in other:
+            raise ValueError(f"Missing packed expert projection: {key[0]}")
+        fc1 = group[0] if key[2] == 1 else other[0]
+        fc2 = other[0] if key[2] == 1 else group[0]
+        if tuple(fc1.shape) != (2 * fc2.shape[1], fc2.shape[0]):
+            raise ValueError(f"Inconsistent gate/up and down shapes: {key[0]}")
+        buckets.append([group[expert] for expert in range(num_experts)])
+    return buckets
+
+
+def _layout(tensors: Sequence[NamedTensor]) -> tuple[TensorLayout, ...]:
+    if any(tensor.dtype != torch.bfloat16 for _, tensor in tensors):
+        raise ValueError("Packed Qwen expert gather requires BF16 tensors")
+    return tuple((name, tuple(tensor.shape)) for name, tensor in tensors)
+
+
+class Qwen35PackedExpertGather:
+    """Gather one native projection directly into global expert order."""
+
+    def __init__(self, num_experts: int) -> None:
+        self.num_experts = num_experts
+        self._layouts = None
+        self._ranks = None
+        self._orders = ()
+        self._name = None
+        self._shape = ()
+
+    def _prepare(self, layouts: Sequence[tuple[TensorLayout, ...]], ranks: tuple[int, ...]) -> None:
+        local_count = self.num_experts // len(ranks)
+        orders, expected_prefix, expected_projections, expected_shape = [], None, None, None
+        for index, layout in enumerate(layouts):
+            entries = {}
+            for position, (name, shape) in enumerate(layout):
+                match = _TARGET.fullmatch(name)
+                if match is None or len(shape) != 2:
+                    raise ValueError(f"Unsupported packed expert output: {name}")
+                prefix, expert, projection = match.groups()
+                if expected_prefix is None:
+                    expected_prefix, expected_shape = prefix, shape
+                key = (int(expert), projection)
+                if prefix != expected_prefix or shape != expected_shape or key in entries:
+                    raise ValueError(f"Mixed or duplicate packed expert output: {name}")
+                entries[key] = position
+            projections = {projection for _, projection in entries}
+            if projections not in ({"gate", "up"}, {"down"}):
+                raise ValueError("Packed experts require complete gate/up or down projections")
+            if expected_projections is None:
+                expected_projections = projections
+            if projections != expected_projections:
+                raise ValueError("Packed expert projection differs between owners")
+            order = ("gate", "up") if "gate" in projections else ("down",)
+            keys = [
+                (expert, projection)
+                for expert in range(index * local_count, (index + 1) * local_count)
+                for projection in order
+            ]
+            if set(keys) != set(entries):
+                raise ValueError("Packed expert IDs must follow contiguous process-group rank order")
+            orders.append(tuple(entries[key] for key in keys))
+        rows, cols = expected_shape
+        self._name = expected_prefix + (".gate_up_proj" if "gate" in expected_projections else ".down_proj")
+        self._shape = (self.num_experts, rows * len(expected_projections), cols)
+        self._orders = tuple(orders)
+        self._layouts, self._ranks = tuple(layouts), ranks
+
+    def __call__(
+        self, tensors: Sequence[NamedTensor], *, group: dist.ProcessGroup, device: torch.device
+    ) -> list[NamedTensor]:
+        ranks = tuple(dist.get_process_group_ranks(group))
+        local_index = dist.get_rank(group=group)
+        layout = _layout(tensors)
+        if not ranks or self.num_experts <= 0 or self.num_experts % len(ranks):
+            raise ValueError("Packed Qwen experts require evenly partitioned EP ownership")
+        if self._layouts is None:
+            layouts = [None] * len(ranks)
+            dist.all_gather_object(layouts, layout, group=group)
+            self._prepare(layouts, ranks)
+        elif ranks != self._ranks or layout != self._layouts[local_index]:
+            raise ValueError("Expert ownership or tensor layout changed; rebuild the gather plan")
+        output = torch.empty(self._shape, dtype=torch.bfloat16, device=device)
+        local = output.chunk(len(ranks), dim=0)[local_index].view(-1)
+        torch.cat([tensors[index][1].reshape(-1) for index in self._orders[local_index]], out=local)
+        handle = dist.all_gather_into_tensor(output.view(-1), local, group=group, async_op=True)
+        handle.wait()
+        return [(self._name, output)]
diff --git a/relax/backends/megatron/weight_update/update_weight_from_tensor.py b/relax/backends/megatron/weight_update/update_weight_from_tensor.py
index 0850a64ce..12ea98762 100644
--- a/relax/backends/megatron/weight_update/update_weight_from_tensor.py
+++ b/relax/backends/megatron/weight_update/update_weight_from_tensor.py
@@ -1,5 +1,8 @@
 # Copyright (c) 2026 Relax Authors. All Rights Reserved.
 
+import json
+import os
+import urllib.request
 from argparse import Namespace
 from collections.abc import Callable, Mapping, Sequence
 from contextlib import nullcontext
@@ -26,11 +29,13 @@
     is_lora_enabled,
     is_lora_merge_mode,
 )
+from relax.utils.misc import get_hf_config
 
 from ..sglang import FlattenedTensorBucket, MultiprocessingSerializer
 from .expert_routing import ExpertLayout, validate_server_layout
 from .hf_weight_iterator_base import HfWeightIteratorBase
 from .lora_adapter_sync import LoraAdapterSync
+from .qwen35_direct_experts import direct_expert_chunk, dispatch_typed_buckets, validate_direct_receiver
 from .update_weight_from_distributed import (
     connect_rollout_engines_from_distributed,
     disconnect_rollout_engines_from_distributed,
@@ -96,6 +101,26 @@ def __init__(
             args=args, model=model, model_name=model_name, quantization_config=quantization_config
         )
 
+        self._direct_qwen35_experts = os.environ.get("RELAX_QWEN35_DIRECT_EXPERTS", "0") == "1"
+        if self._direct_qwen35_experts:
+            config = get_hf_config(args.hf_checkpoint)
+            text_config = getattr(config, "text_config", None)
+            if (
+                not args.colocate
+                or self.lora_enabled
+                or quantization_config
+                or self._route_experts
+                or config.model_type != "qwen3_5_moe"
+                or text_config is None
+                or (text_config.hidden_size, text_config.moe_intermediate_size, text_config.num_experts)
+                != (2048, 512, 256)
+                or os.environ.get("RELAX_QWEN35_PACKED_EXPERTS", "0") != "1"
+                or not getattr(self._hf_weight_iterator, "_bf16_expert_buckets", None)
+                or mpu.get_pipeline_model_parallel_world_size() != 1
+                or mpu.get_expert_tensor_parallel_world_size() != 1
+            ):
+                raise ValueError("Direct expert IPC is limited to packed BF16 Qwen3.5-35B with PP1/ETP1")
+
         self._ipc_gather_group = None
         self._ipc_gather_src = None
         self._ipc_engine = None
@@ -213,6 +238,23 @@ def connect_rollout_engines(
             )
             self._expert_routing_placement = placement
 
+        if getattr(self, "_direct_qwen35_experts", False):
+            if self.use_distribute or any(count != 2 for count in colocate_gpu_counts):
+                raise ValueError("Direct Qwen3.5 experts require colocated two-GPU engines")
+            status = [None]
+            if dist.get_rank(get_gloo_group()) == 0:
+                try:
+                    urls = ray.get([engine.get_url.remote() for engine in self.rollout_engines])
+                    for url in urls:
+                        with urllib.request.urlopen(url.rstrip("/") + "/server_info", timeout=30) as response:
+                            validate_direct_receiver(json.load(response))
+                    status[0] = (True, None)
+                except Exception as error:
+                    status[0] = (False, str(error))
+            dist.broadcast_object_list(status, src=0, group=get_gloo_group())
+            if not status[0][0]:
+                raise RuntimeError(f"Direct Qwen3.5 receiver validation failed: {status[0][1]}")
+
         # Create IPC Gloo gather groups (only on first call; partitioning is
         # fixed across reconnects).
         if self._ipc_gather_group is None:
@@ -665,6 +707,7 @@ def _send_hf_params(
                 ipc_gather_group=self._ipc_gather_group,
                 weight_version=None if self._route_experts else self.weight_version,
                 dispatch_errors=dispatch_errors,
+                direct_experts=getattr(self, "_direct_qwen35_experts", False),
             )
             all_refs.extend(refs_colocated)
 
@@ -690,6 +733,7 @@ def _send_to_colocated_engine(
     ipc_gather_group,
     weight_version,
     dispatch_errors: list[str] | None = None,
+    direct_experts: bool = False,
 ) -> tuple[list[ObjectRef], Any]:
     # Placeholder ranks (GPU slots reserved but no engine) have no gather group.
     # gather_object is only collective among group members, so we skip entirely.
@@ -710,7 +754,14 @@ def _send_to_colocated_engine(
         (name, tensor.to(cur_device) if tensor.device != cur_device else tensor) for name, tensor in hf_named_tensors
     ]
 
-    if getattr(FlattenedTensorBucket, "supports_multi_dtypes", False):
+    direct_chunk = None
+    if direct_experts:
+        if dist.get_world_size(ipc_gather_group) != 2:
+            raise ValueError("Direct Qwen3.5 IPC requires a two-rank receiver group")
+        direct_chunk = direct_expert_chunk(hf_named_tensors, dist.get_rank(group=ipc_gather_group))
+    if direct_chunk is not None:
+        converted_named_tensors_by_dtypes = {"direct": direct_chunk}
+    elif getattr(FlattenedTensorBucket, "supports_multi_dtypes", False):
         # An empty local chunk contributes no bucket; FlattenedTensorBucket
         # rejects an empty named_tensors list, so skip building one.
         converted_named_tensors_by_dtypes = {"dtype": hf_named_tensors} if hf_named_tensors else {}
@@ -724,14 +775,18 @@ def _send_to_colocated_engine(
 
     serialized_tensors = []
     for _dtype, named_tensors in converted_named_tensors_by_dtypes.items():
-        flattened_tensor_bucket = FlattenedTensorBucket(named_tensors=named_tensors)
-        metadata = flattened_tensor_bucket.get_metadata()
-        flattened_tensor_data = {
-            "flattened_tensor": flattened_tensor_bucket.get_flattened_tensor(),
-            "metadata": metadata,
-        }
-        long_live_tensors.append(flattened_tensor_data)
-        serialized_tensors.append(MultiprocessingSerializer.serialize(flattened_tensor_data, output_str=True))
+        if direct_chunk is not None:
+            payload, load_format = named_tensors, "direct"
+        else:
+            flattened_tensor_bucket = FlattenedTensorBucket(named_tensors=named_tensors)
+            payload = {
+                "flattened_tensor": flattened_tensor_bucket.get_flattened_tensor(),
+                "metadata": flattened_tensor_bucket.get_metadata(),
+            }
+            load_format = "flattened_bucket"
+        long_live_tensors.append(payload)
+        serialized = MultiprocessingSerializer.serialize(payload, output_str=True)
+        serialized_tensors.append((load_format, serialized) if direct_experts else serialized)
 
     serialized_named_tensors = (
         [None] * dist.get_world_size(ipc_gather_group) if ipc_gather_src == dist.get_rank() else None
@@ -744,6 +799,18 @@ def _send_to_colocated_engine(
     )
 
     refs = []
+    if direct_experts:
+        if dist.get_rank() == ipc_gather_src:
+            refs = dispatch_typed_buckets(
+                serialized_named_tensors,
+                engine=ipc_engine,
+                version=weight_version,
+                serializer=MultiprocessingSerializer,
+                empty_flattened=lambda: _empty_flattened_tensor_data(cur_device),
+                held=long_live_tensors,
+                dispatch_errors=dispatch_errors,
+            )
+        return refs, long_live_tensors
     if dist.get_rank() == ipc_gather_src:
         # Ranks can contribute uneven bucket counts: under PP/EP/MoE layouts a TP
         # rank may legitimately hold no HF tensors for a chunk. Pad short ranks
diff --git a/relax/engine/rollout/sglang_rollout.py b/relax/engine/rollout/sglang_rollout.py
index 765484990..6a8cf8ef7 100644
--- a/relax/engine/rollout/sglang_rollout.py
+++ b/relax/engine/rollout/sglang_rollout.py
@@ -581,7 +581,7 @@ async def generate(
 
     if "routed_experts" in output["meta_info"]:
         sample.rollout_routed_experts = np.frombuffer(
-            pybase64.b64decode(output["meta_info"]["routed_experts"].encode("ascii")),
+            pybase64.b64decode(output["meta_info"]["routed_experts"], validate=True),
             dtype=np.int32,
         ).reshape(
             len(sample.tokens) - 1,
diff --git a/scripts/training/multimodal/run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh b/scripts/training/multimodal/run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh
index dc979a8c2..dc5f9893d 100644
--- a/scripts/training/multimodal/run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh
+++ b/scripts/training/multimodal/run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh
@@ -27,7 +27,6 @@ NUM_ROLLOUT="${NUM_ROLLOUT:=200}"
 
 CKPT_ARGS=(
    --hf-checkpoint ${MODEL_DIR}/Qwen3.5-35B-A3B
-   --ref-load ${MODEL_DIR}/Qwen3.5-35B-A3B
    --megatron-to-hf-mode bridge
    --warm-hf-checkpoint-page-cache
 )
@@ -43,12 +42,12 @@ ROLLOUT_ARGS=(
    --rollout-shuffle
    --rm-type openr1mm
    --num-rollout ${NUM_ROLLOUT}
-   --rollout-batch-size 32
+   --rollout-batch-size 64
    --n-samples-per-prompt 8
    --rollout-max-response-len 2048
    --rollout-max-prompt-len 2048
    --rollout-temperature 1
-   --global-batch-size 256
+   --global-batch-size 512
    --use-streaming-dataset
    --balance-data
    --use-fault-tolerance
@@ -60,10 +59,10 @@ ROLLOUT_ARGS=(
 PERF_ARGS=(
    --tensor-model-parallel-size 2
    --sequence-parallel
-   --pipeline-model-parallel-size 2
+   --pipeline-model-parallel-size 1
    --calculate-per-token-loss
    --context-parallel-size 1
-   --expert-model-parallel-size 4
+   --expert-model-parallel-size 8
    --expert-tensor-parallel-size 1
    --recompute-granularity full
    --recompute-method uniform
@@ -71,14 +70,16 @@ PERF_ARGS=(
    # --qkv-format bshd
    # --micro-batch-size 1 # avoid OOM
    --use-dynamic-batch-size
-   --max-tokens-per-gpu 6144
+   --max-tokens-per-gpu 8192
+   --moe-token-dispatcher-type flex
+   --moe-flex-dispatcher-backend deepep
 )
 
 GRPO_ARGS=(
    --advantage-estimator grpo
-   --use-kl-loss
+   --true-on-policy-mode
+   # Zero KL coefficient: skip the reference model and its forward pass.
    --kl-loss-coef 0.00
-   --kl-loss-type low_var_kl
    --entropy-coef 0.00
    --eps-clip 0.2
    --eps-clip-high 0.28
@@ -103,7 +104,6 @@ OPTIMIZER_ARGS=(
 
    # --fp16 # Qwen3.5 does not support fp16 training for now
    --use-rollout-routing-replay
-   --use-slime-router
 )
 
 SGLANG_ARGS=(
@@ -129,6 +129,37 @@ MISC_ARGS=(
    --attention-backend flash
 )
 
+# Match the native BF16 expert layout used by the guarded direct IPC loader.
+SGLANG_ARGS+=(--sglang-moe-runner-backend triton --sglang-disable-shared-experts-fusion --sglang-dtype bfloat16)
+
+RUNTIME_ENV_JSON=$(python3 - <<'PYENV'
+import json
+import os
+
+runtime_env = json.loads(os.environ["RUNTIME_ENV_JSON"])
+# Vision embeddings must be recomputed after each weight update.
+runtime_env.setdefault("env_vars", {})["SGLANG_VLM_CACHE_SIZE_MB"] = "0"
+# The backend validates single-update, zero-KL GRPO before reusing log-probs.
+runtime_env["env_vars"]["RELAX_COLOCATE_INLINE_LOG_PROBS"] = "1"
+runtime_env["env_vars"]["NVTE_USE_CUTLASS_GROUPED_GEMM"] = "1"
+# Export BF16 experts in bulk, then send each receiver its TP-local shard.
+runtime_env["env_vars"]["RELAX_BF16_EXPERT_GATHER"] = "1"
+runtime_env["env_vars"]["RELAX_ASYNC_WEIGHT_GATHER"] = "1"
+runtime_env["env_vars"]["RELAX_QWEN35_PACKED_EXPERTS"] = "1"
+runtime_env["env_vars"]["RELAX_QWEN35_DIRECT_EXPERTS"] = "1"
+# Preserve the platform Trial identity and its collector path in Ray workers.
+# These are public variable names; their values are inherited at launch.
+for key in (
+    "QS_TRIAL_ID", "QS_JOB_ID", "QS_USER", "GPU_TYPE", "CLUSTER",  # gitleaks:allow
+    "ETTR_LOG_FILE", "ETTR_TRACE_LOG_DIR", "ETTR_FRAMEWORK",
+    "QS_INJECT_AUTOLOAD", "QS_INJECT_MODEL_PROFILE",  # gitleaks:allow
+):
+    if key in os.environ:
+        runtime_env["env_vars"][key] = os.environ[key]
+print(json.dumps(runtime_env))
+PYENV
+)
+
 mkdir -p log
 
 ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://127.0.0.1:8265" \
diff --git a/tests/backends/megatron/test_colocate_inline_log_probs.py b/tests/backends/megatron/test_colocate_inline_log_probs.py
new file mode 100644
index 000000000..af8177c22
--- /dev/null
+++ b/tests/backends/megatron/test_colocate_inline_log_probs.py
@@ -0,0 +1,214 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+"""Safety gates and zero-KL inputs for single-update colocate GRPO."""
+
+import ast
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+
+
+ROOT = Path(__file__).resolve().parents[3]
+
+
+def load_function(path, name, namespace, class_name=None):
+    tree = ast.parse(path.read_text())
+    body = tree.body
+    if class_name is not None:
+        body = next(node for node in body if isinstance(node, ast.ClassDef) and node.name == class_name).body
+    function = next(node for node in body if isinstance(node, ast.FunctionDef) and node.name == name)
+    future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)
+    module = ast.fix_missing_locations(ast.Module(body=[future, function], type_ignores=[]))
+    exec(compile(module, str(path), "exec"), namespace)
+    return namespace[name]
+
+
+@pytest.fixture
+def actor():
+    args = SimpleNamespace(
+        colocate=True,
+        true_on_policy_mode=True,
+        fully_async=False,
+        hybrid=False,
+        advantage_estimator="grpo",
+        loss_type="policy_loss",
+        compute_advantages_and_returns=True,
+        rollout_batch_size=32,
+        n_samples_per_prompt=8,
+        global_batch_size=256,
+        num_iters_per_train_update=1,
+        max_staleness=0,
+        update_weights_interval=1,
+        partial_rollout=False,
+        rl_partial_rollouts=False,
+        use_dynamic_global_batch_size=False,
+        context_parallel_size=1,
+        pipeline_model_parallel_size=1,
+        dynamic_context_parallel=False,
+        hidden_dropout=0.0,
+        attention_dropout=0.0,
+        moe_input_jitter_eps=None,
+        bf16=True,
+        fp16=False,
+        fp8=None,
+        fp4=None,
+        fp8_param_gather=False,
+        moe_router_enable_expert_bias=False,
+        moe_aux_loss_coeff=0.0,
+        moe_router_load_balancing_type=["none"],
+        moe_router_force_load_balancing=False,
+        moe_router_force_biased=None,
+        moe_expert_capacity_factor=None,
+        kl_coef=0.0,
+        use_kl_loss=False,
+        use_opd=False,
+        use_critic=False,
+        keep_old_actor=False,
+        use_rollout_logprobs=False,
+        use_routing_replay=True,
+        use_rollout_routing_replay=True,
+        use_rollout_indexer_replay=False,
+        enable_mtp_training=False,
+        use_opsm=False,
+        use_tis=True,
+        calculate_per_token_loss=True,
+        entropy_coef=0.0,
+        recompute_loss_function=False,
+        custom_tis_function_path=None,
+        custom_pg_loss_reducer_function_path=None,
+        custom_loss_function_path=None,
+        custom_model_provider_path=None,
+        custom_megatron_before_log_prob_hook_path=None,
+        custom_megatron_before_train_step_hook_path=None,
+    )
+    return SimpleNamespace(
+        args=args, weights_backuper=SimpleNamespace(backup_tags=["actor"]), rollout_data_postprocess=None
+    )
+
+
+@pytest.fixture
+def inline_guard(monkeypatch):
+    monkeypatch.setenv("RELAX_COLOCATE_INLINE_LOG_PROBS", "1")
+    return load_function(
+        ROOT / "relax/backends/megatron/actor.py",
+        "_inline_colocate_old_log_probs",
+        {"os": os},
+        "MegatronTrainRayActor",
+    )
+
+
+def test_single_fresh_grpo_update_is_eligible(actor, inline_guard):
+    assert inline_guard(actor, [22])
+
+
+@pytest.mark.parametrize(
+    "name,value",
+    [
+        ("colocate", False),
+        ("true_on_policy_mode", False),
+        ("num_iters_per_train_update", 2),
+        ("max_staleness", 1),
+        ("update_weights_interval", 2),
+        ("global_batch_size", 128),
+        ("partial_rollout", True),
+        ("use_dynamic_global_batch_size", True),
+        ("context_parallel_size", 2),
+        ("dynamic_context_parallel", True),
+        ("hidden_dropout", 0.1),
+        ("attention_dropout", 0.1),
+        ("moe_input_jitter_eps", 0.01),
+        ("fp8", "hybrid"),
+        ("fp4", "nvfp4"),
+        ("fp8_param_gather", True),
+        ("moe_router_force_load_balancing", True),
+        ("moe_router_force_biased", 0.1),
+        ("moe_expert_capacity_factor", 1.0),
+        ("custom_model_provider_path", "custom.stateful_model"),
+        ("moe_router_enable_expert_bias", True),
+        ("moe_router_load_balancing_type", ["aux_loss"]),
+        ("kl_coef", 0.01),
+        ("use_kl_loss", True),
+        ("use_opd", True),
+        ("use_critic", True),
+        ("keep_old_actor", True),
+        ("use_rollout_logprobs", True),
+        ("use_rollout_routing_replay", False),
+        ("enable_mtp_training", True),
+        ("use_opsm", True),
+        ("custom_tis_function_path", "custom.function"),
+        ("custom_megatron_before_train_step_hook_path", "custom.mutate_weights"),
+    ],
+)
+def test_unsafe_algorithm_or_state_is_rejected(actor, inline_guard, name, value):
+    setattr(actor.args, name, value)
+    with pytest.raises(ValueError, match="Unsupported colocate true-on-policy"):
+        inline_guard(actor, [22])
+
+
+def test_actual_multiple_updates_are_rejected(actor, inline_guard):
+    with pytest.raises(ValueError, match="exactly one optimizer step"):
+        inline_guard(actor, [11, 11])
+
+
+def test_extra_model_and_data_mutation_are_rejected(actor, inline_guard):
+    actor.weights_backuper.backup_tags.append("ref")
+    with pytest.raises(ValueError, match="reference"):
+        inline_guard(actor, [22])
+    actor.weights_backuper.backup_tags = ["actor"]
+    actor.rollout_data_postprocess = lambda args: None
+    with pytest.raises(ValueError, match="rollout_data_postprocess"):
+        inline_guard(actor, [22])
+
+
+def test_existing_modes_do_not_enter_colocate_optimization(actor, inline_guard, monkeypatch):
+    monkeypatch.delenv("RELAX_COLOCATE_INLINE_LOG_PROBS")
+    assert not inline_guard(actor, [11, 11])
+    actor.args.true_on_policy_mode = False
+    assert not inline_guard(actor, [11, 11])
+    actor.args.true_on_policy_mode = True
+    actor.args.colocate = False
+    actor.args.fully_async = True
+    assert not inline_guard(actor, [11, 11])
+
+
+def test_zero_kl_uses_sampling_shape_without_sampling_values():
+    torch = pytest.importorskip("torch")
+    calls = []
+
+    def compute_impl(args, **kwargs):
+        calls.append(kwargs["kl"])
+        return kwargs["kl"], kwargs["kl"]
+
+    namespace = {
+        "torch": torch,
+        "mpu": SimpleNamespace(is_pipeline_last_stage=lambda: True),
+        "maybe_padded_total_lengths": lambda *args: None,
+        "compute_advantages_and_returns_impl": compute_impl,
+    }
+    compute = load_function(ROOT / "relax/backends/megatron/loss.py", "compute_advantages_and_returns", namespace)
+    args = SimpleNamespace(
+        _colocate_inline_log_probs=True,
+        advantage_estimator="grpo",
+        kl_coef=0,
+        use_rollout_logprobs=False,
+        qkv_format="thd",
+        use_opd=False,
+        normalize_advantages=False,
+    )
+    for value in (-0.1, -100):
+        data = {
+            "rollout_log_probs": [torch.full((3,), value, dtype=torch.float16), torch.full((2,), value)],
+            "rewards": [1.0, -1.0],
+            "response_lengths": [3, 2],
+            "total_lengths": [5, 7],
+            "loss_masks": [torch.ones(3), torch.ones(2)],
+        }
+        compute(args, data)
+        assert "log_probs" not in data
+        assert [tuple(x.shape) for x in data["advantages"]] == [(3,), (2,)]
+        assert all(x.dtype == torch.float32 and torch.count_nonzero(x) == 0 for x in data["advantages"])
+    assert all(torch.equal(a, b) for a, b in zip(*calls, strict=True))
+    with pytest.raises(ValueError, match="requires rollout log-probs"):
+        compute(args, {"response_lengths": [3], "total_lengths": [5]})
diff --git a/tests/backends/megatron/weight_update/test_async_gdn_weight_gather.py b/tests/backends/megatron/weight_update/test_async_gdn_weight_gather.py
new file mode 100644
index 000000000..1bba96d8f
--- /dev/null
+++ b/tests/backends/megatron/weight_update/test_async_gdn_weight_gather.py
@@ -0,0 +1,85 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+"""Reconstruct known global GDN tensors from real TP-local component shards."""
+
+import ast
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+
+@pytest.mark.parametrize("tp", [1, 2, 4])
+@pytest.mark.parametrize("wrapped", [False, True])
+def test_async_gather_preserves_gdn_glu_and_replicated_layouts(tp, wrapped):
+    path = Path(__file__).resolve().parents[4] / "relax/backends/megatron/weight_update/common.py"
+    tree = ast.parse(path.read_text())
+    function = next(
+        node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "all_gather_params_async"
+    )
+    inputs, sources, expected, split_sizes = [], {}, [], {}
+    specs = [
+        ("self_attention.conv1d.weight", (2, 2, 4), (1, 4)),
+        ("self_attention.in_proj.weight", (2, 2, 4, 4, 1, 1), (3,)),
+        ("mlp.linear_fc1.weight", (3, 3), (2,)),
+    ]
+    for parameter_index, (suffix, sizes, rest) in enumerate(specs):
+        name = "decoder.layers.0." + suffix
+        components = []
+        for component, size in enumerate(sizes):
+            shape = (size * tp, *rest)
+            values = torch.arange(torch.Size(shape).numel(), dtype=torch.float32).reshape(shape)
+            components.append((values + 1000 * component + 10000 * parameter_index).to(torch.bfloat16))
+        shards = [
+            torch.cat([component.chunk(tp, dim=0)[rank] for component in components], dim=0) for rank in range(tp)
+        ]
+        param = shards[0].clone()
+        param.tensor_model_parallel = True
+        param.partition_dim = 0
+        param.partition_stride = 2 if "linear_fc1" in name else 1
+        full_name = name.replace(".weight", ".to_wrap.weight") if wrapped else name
+        inputs.append((SimpleNamespace(name=full_name), param))
+        sources[param.data_ptr()] = shards
+        expected.append(torch.cat(components, dim=0))
+        if "self_attention" in name:
+            split_sizes[name] = sizes
+    replicated = torch.tensor([7, 11], dtype=torch.bfloat16)
+    replicated.tensor_model_parallel = False
+    inputs.append((SimpleNamespace(name="replicated.weight"), replicated))
+    expected.append(replicated)
+    pending, waited = [], []
+
+    def all_gather(outputs, param, *, group, async_op):
+        assert group == "tp" and async_op
+        shards = sources[param.data_ptr()]
+        pending.append(outputs)
+
+        def wait():
+            assert len(pending) == 3, "Every gather must launch before waiting"
+            for out, source in zip(outputs, shards, strict=True):
+                out.copy_(source)
+            waited.append(True)
+
+        return SimpleNamespace(wait=wait)
+
+    namespace = {
+        "torch": torch,
+        "dist": SimpleNamespace(all_gather=all_gather),
+        "mpu": SimpleNamespace(
+            get_tensor_model_parallel_world_size=lambda: tp, get_tensor_model_parallel_group=lambda: "tp"
+        ),
+    }
+    module = ast.Module(
+        body=[ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), function],
+        type_ignores=[],
+    )
+    exec(compile(ast.fix_missing_locations(module), str(path), "exec"), namespace)
+    result = namespace["all_gather_params_async"](inputs, gdn_split_sizes=split_sizes)
+    assert len(waited) == 3
+    assert len(result) == len(expected)
+    assert all(
+        torch.equal(actual.view(torch.uint8), target.view(torch.uint8))
+        for actual, target in zip(result, expected, strict=True)
+    )
+    assert result[-1].data_ptr() == replicated.data_ptr()
diff --git a/tests/backends/megatron/weight_update/test_bf16_expert_gather.py b/tests/backends/megatron/weight_update/test_bf16_expert_gather.py
new file mode 100644
index 000000000..b2f55591b
--- /dev/null
+++ b/tests/backends/megatron/weight_update/test_bf16_expert_gather.py
@@ -0,0 +1,128 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+"""Packing, tensor lifetime and metadata reuse across replacement groups."""
+
+import importlib.util
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+
+@pytest.fixture
+def module():
+    path = Path(__file__).resolve().parents[4] / "relax/backends/megatron/weight_update/bf16_expert_gather.py"
+    spec = importlib.util.spec_from_file_location("bf16_expert_gather_test_target", path)
+    result = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(result)
+    return result
+
+
+def simulate_collectives(monkeypatch, module, peers):
+    calls = []
+    monkeypatch.setattr(module.dist, "get_process_group_ranks", lambda group: group.ranks)
+    monkeypatch.setattr(module.dist, "get_rank", lambda group: group.rank)
+
+    def metadata(result, local, *, group):
+        assert local == module._layout(peers[group.rank])
+        result[:] = [module._layout(peer) for peer in peers]
+        calls.append(("metadata", group))
+
+    def flattened(peer):
+        return torch.cat([tensor.reshape(-1) for _, tensor in peer]) if peer else torch.empty(0, dtype=torch.bfloat16)
+
+    def wait():
+        calls.append(("wait", None))
+
+    def uniform(output, local, *, group, async_op):
+        assert async_op and torch.equal(local, flattened(peers[group.rank]))
+        output.copy_(torch.cat([flattened(peer) for peer in peers]))
+        calls.append(("uniform", group))
+        return SimpleNamespace(wait=wait)
+
+    def uneven(outputs, local, *, group, async_op):
+        assert async_op and torch.equal(local, flattened(peers[group.rank]))
+        for output, peer in zip(outputs, peers, strict=True):
+            output.copy_(flattened(peer))
+        calls.append(("uneven", group))
+        return SimpleNamespace(wait=wait)
+
+    def broadcast(output, *, src, group, async_op):
+        assert async_op
+        output.copy_(flattened(peers[group.ranks.index(src)]))
+        calls.append(("broadcast", group))
+        return SimpleNamespace(wait=wait)
+
+    monkeypatch.setattr(module.dist, "all_gather_object", metadata)
+    monkeypatch.setattr(module.dist, "all_gather_into_tensor", uniform)
+    monkeypatch.setattr(module.dist, "all_gather", uneven)
+    monkeypatch.setattr(module.dist, "broadcast", broadcast)
+    return calls
+
+
+@pytest.mark.parametrize(
+    "sizes,collective", [([6, 6], "uniform"), ([6, 4], "uneven"), ([0, 6], "broadcast"), ([0, 0], None)]
+)
+@pytest.mark.parametrize("rank", [0, 1])
+def test_gather_values_old_output_lifetime_and_reloaded_group(module, monkeypatch, sizes, collective, rank):
+    peers = [
+        [(f"expert.{index}", torch.arange(size, dtype=torch.bfloat16).view(2, -1).t())] if size else []
+        for index, size in enumerate(sizes)
+    ]
+    calls = simulate_collectives(monkeypatch, module, peers)
+    first_group = SimpleNamespace(ranks=[2, 7], rank=rank)
+    plan = module.Bf16ExpertGather()
+    actual = plan(peers[rank], group=first_group, device=torch.device("cpu"))
+    expected = [entry for peer in peers for entry in peer]
+    assert [name for name, _ in actual] == [name for name, _ in expected]
+    assert all(torch.equal(tensor, other) for (_, tensor), (_, other) in zip(actual, expected, strict=True))
+    frozen = [(name, tensor.clone()) for name, tensor in actual]
+    for peer in peers:
+        for _, tensor in peer:
+            tensor.add_(10)
+    replacement_group = SimpleNamespace(ranks=[2, 7], rank=rank)
+    updated = plan(peers[rank], group=replacement_group, device=torch.device("cpu"))
+    assert all(torch.equal(tensor, other) for (_, tensor), (_, other) in zip(actual, frozen, strict=True))
+    assert all(torch.equal(tensor, other) for (_, tensor), (_, other) in zip(updated, expected, strict=True))
+    assert sum(kind == "metadata" for kind, _ in calls) == 1
+    if collective:
+        assert [group for kind, group in calls if kind == collective] == [first_group, replacement_group]
+        assert sum(kind == "wait" for kind, _ in calls) == 2
+    assert not any(isinstance(value, torch.Tensor) or value is first_group for value in vars(plan).values())
+
+
+def test_changed_layout_and_topology_rejected(module, monkeypatch):
+    peers = [[("a", torch.ones(4, dtype=torch.bfloat16))], [("b", torch.ones(4, dtype=torch.bfloat16))]]
+    calls = simulate_collectives(monkeypatch, module, peers)
+    group = SimpleNamespace(ranks=[2, 7], rank=0)
+    plan = module.Bf16ExpertGather()
+    plan(peers[0], group=group, device=torch.device("cpu"))
+    count = len(calls)
+    for tensors, next_group in [
+        ([("a", torch.ones(2, dtype=torch.bfloat16))], group),
+        (peers[0], SimpleNamespace(ranks=[2, 8], rank=0)),
+    ]:
+        with pytest.raises(ValueError, match="layout changed"):
+            plan(tensors, group=next_group, device=torch.device("cpu"))
+    with pytest.raises(ValueError, match="requires BF16"):
+        plan([("a", torch.ones(4))], group=group, device=torch.device("cpu"))
+    assert len(calls) == count
+
+
+def test_owner_rounds_cover_every_input_within_budget(module):
+    infos = [
+        SimpleNamespace(name=str(index), src_rank=owner, size=size)
+        for index, (owner, size) in enumerate([(7, 6), (2, 4), (7, 4), (2, 4), (7, 6)])
+    ]
+    rounds = module.bucket_experts_by_owner(infos, 20, [2, 7])
+    assert [[info.name for info in bucket] for bucket in rounds] == [["1", "3", "0", "2"], ["4"]]
+    assert sorted(id(info) for bucket in rounds for info in bucket) == sorted(map(id, infos))
+    assert all(sum(info.size for info in bucket) <= 20 for bucket in rounds)
+    assert module.bucket_experts_by_owner([], 20, [2, 7]) == []
+    for invalid in [
+        SimpleNamespace(name="bad", src_rank=9, size=1),
+        SimpleNamespace(name="large", src_rank=2, size=11),
+    ]:
+        with pytest.raises(ValueError, match="Unsupported"):
+            module.bucket_experts_by_owner([invalid], 20, [2, 7])
diff --git a/tests/backends/megatron/weight_update/test_expert_routing.py b/tests/backends/megatron/weight_update/test_expert_routing.py
index b87911f27..8d7fd4df1 100644
--- a/tests/backends/megatron/weight_update/test_expert_routing.py
+++ b/tests/backends/megatron/weight_update/test_expert_routing.py
@@ -183,6 +183,8 @@ def _check_bridge_chunk_rounds(rank: int, layout: ExpertLayout, *, interleaved:
     iterator._vanilla_key_map = {info.name: info.name for i, info in enumerate(infos) if i % 3 == rank}
     iterator._non_expert_buckets = []
     iterator._expert_broadcast_caches = []
+    iterator._bf16_expert_buckets = None
+    iterator._async_weight_gather = False
     iterator.lora_merge_mode = False
     iterator._bridge_converter = SimpleNamespace(
         init_tasks=lambda: None,
diff --git a/tests/backends/megatron/weight_update/test_qwen35_direct_experts.py b/tests/backends/megatron/weight_update/test_qwen35_direct_experts.py
new file mode 100644
index 000000000..799590e05
--- /dev/null
+++ b/tests/backends/megatron/weight_update/test_qwen35_direct_experts.py
@@ -0,0 +1,143 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+import importlib.util
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+
+_PATH = Path(__file__).parents[4] / "relax/backends/megatron/weight_update/qwen35_direct_experts.py"
+_SPEC = importlib.util.spec_from_file_location("qwen35_direct_test", _PATH)
+direct = importlib.util.module_from_spec(_SPEC)
+_SPEC.loader.exec_module(direct)
+
+
+@pytest.mark.parametrize("rank", [0, 1])
+def test_shards_match_separate_gate_up_and_down_slices_and_keep_old_storage(rank):
+    gate_up = torch.arange(3 * 8 * 6, dtype=torch.bfloat16).reshape(3, 8, 6)
+    down = torch.arange(3 * 6 * 4, dtype=torch.bfloat16).reshape(3, 6, 4)
+    first = direct._shard_expert(gate_up, "gate_up_proj", rank)
+    second = direct._shard_expert(down, "down_proj", rank)
+    rows = [0, 1, 4, 5] if rank == 0 else [2, 3, 6, 7]
+    columns = [0, 1] if rank == 0 else [2, 3]
+    torch.testing.assert_close(first, gate_up[:, rows], rtol=0, atol=0)
+    torch.testing.assert_close(second, down[:, :, columns], rtol=0, atol=0)
+    assert first.is_contiguous() and second.is_contiguous()
+    snapshots = first.clone(), second.clone()
+    direct._shard_expert(gate_up + 5, "gate_up_proj", rank)
+    direct._shard_expert(down + 5, "down_proj", rank)
+    torch.testing.assert_close(first, snapshots[0], rtol=0, atol=0)
+    torch.testing.assert_close(second, snapshots[1], rtol=0, atol=0)
+
+
+@pytest.mark.parametrize(
+    "projection,shape,target",
+    [
+        ("gate_up_proj", (256, 1024, 2048), (256, 512, 2048)),
+        ("down_proj", (256, 2048, 512), (256, 2048, 256)),
+    ],
+)
+def test_complete_model_shapes_and_runtime_names_without_allocating_storage(projection, shape, target):
+    tensor = torch.empty(shape, dtype=torch.bfloat16, device="meta")
+    for rank in (0, 1):
+        result = direct.direct_expert_chunk(
+            [(f"model.language_model.layers.9.mlp.experts.{projection}", tensor)], rank
+        )
+        suffix = "w13_weight" if projection == "gate_up_proj" else "w2_weight"
+        assert result[0][0] == f"model.layers.9.mlp.experts.{suffix}"
+        assert tuple(result[0][1].shape) == target
+    assert direct.direct_expert_chunk([("lm_head.weight", tensor)], 0) is None
+    assert direct.direct_expert_chunk([], 0) is None
+
+
+def _layout():
+    return dict(
+        tp_size=2,
+        ep_size=1,
+        pp_size=1,
+        moe_dp_size=1,
+        moe_runner_backend="triton",
+        dtype="bfloat16",
+        disable_shared_experts_fusion=True,
+        enable_eplb=False,
+        ep_num_redundant_experts=0,
+        init_expert_location="trivial",
+        ep_join_mode=None,
+        quantization=None,
+        enable_lora=None,
+        speculative_algorithm=None,
+        json_model_override_args="{}",
+    )
+
+
+@pytest.mark.parametrize(
+    "override",
+    [
+        dict(tp_size=4),
+        dict(ep_size=2),
+        dict(pp_size=2),
+        dict(moe_dp_size=2),
+        dict(moe_runner_backend="auto"),
+        dict(moe_runner_backend="triton_kernel"),
+        dict(moe_runner_backend="flashinfer_cutlass"),
+        dict(dtype="float16"),
+        dict(dtype="auto"),
+        dict(disable_shared_experts_fusion=False),
+        dict(enforce_shared_experts_fusion=True),
+        dict(enable_eplb=True),
+        dict(ep_num_redundant_experts=1),
+        dict(init_expert_location="random"),
+        dict(ep_join_mode="scale"),
+        dict(quantization="fp8"),
+        dict(enable_lora=True),
+        dict(speculative_algorithm="EAGLE"),
+        dict(json_model_override_args='{"hidden_size":4096}'),
+    ],
+)
+def test_unsupported_runtime_layout_is_rejected(override):
+    direct.validate_direct_receiver(_layout())
+    with pytest.raises(ValueError):
+        direct.validate_direct_receiver(_layout() | override)
+
+
+@pytest.mark.parametrize(
+    "packets,formats",
+    [
+        ([[("direct", "d0")], [("direct", "d1")]], ["direct"]),
+        ([[("flattened_bucket", "f0")], [("flattened_bucket", "f1")]], ["flattened_bucket"]),
+        ([[], [("direct", "d1")]], ["direct"]),
+        ([[("direct", "d0")], []], ["direct"]),
+        ([[("direct", "d0")], [("flattened_bucket", "f1")]], ["direct", "flattened_bucket"]),
+        ([[], []], []),
+    ],
+)
+def test_mixed_and_empty_ranks_receive_one_consistent_format_per_request(packets, formats):
+    calls, held = [], []
+
+    def serialize(value, output_str):
+        assert output_str
+        return "empty-direct" if isinstance(value, list) else "empty-flat"
+
+    def remote(**kwargs):
+        calls.append(kwargs)
+        return len(calls)
+
+    refs = direct.dispatch_typed_buckets(
+        packets,
+        engine=SimpleNamespace(update_weights_from_tensor=SimpleNamespace(remote=remote)),
+        version=12,
+        serializer=SimpleNamespace(serialize=serialize),
+        empty_flattened=lambda: dict(flattened_tensor=torch.empty(0, dtype=torch.uint8), metadata=[]),
+        held=held,
+    )
+    assert len(refs) == len(formats) and [c["load_format"] for c in calls] == formats
+    for call in calls:
+        assert call["weight_version"] == "12"
+        for rank, value in enumerate(call["serialized_named_tensors"]):
+            if packets[rank] and packets[rank][0][0] == call["load_format"]:
+                assert value == packets[rank][0][1]
+            else:
+                assert value == ("empty-direct" if call["load_format"] == "direct" else "empty-flat")
+    assert all(isinstance(value, list) or value["flattened_tensor"].numel() == 0 for value in held)
diff --git a/tests/backends/megatron/weight_update/test_qwen35_packed_experts.py b/tests/backends/megatron/weight_update/test_qwen35_packed_experts.py
new file mode 100644
index 000000000..f9d5c1094
--- /dev/null
+++ b/tests/backends/megatron/weight_update/test_qwen35_packed_experts.py
@@ -0,0 +1,155 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+"""Expert IDs, gate/up order, changed weights and replacement process
+groups."""
+
+import importlib.util
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+
+@pytest.fixture
+def module():
+    path = Path(__file__).resolve().parents[4] / "relax/backends/megatron/weight_update/qwen35_packed_experts.py"
+    spec = importlib.util.spec_from_file_location("qwen35_packed_experts_test_target", path)
+    result = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(result)
+    return result
+
+
+def sources():
+    return [
+        SimpleNamespace(
+            name=f"module.module.language_model.decoder.layers.{layer}.mlp.experts.linear_fc{fc}.weight{e}",
+            dtype=torch.bfloat16,
+            shape=(8, 6) if fc == 1 else (6, 4),
+            src_rank=(7, 2)[e // 12],
+        )
+        for layer in (0, 10, 2)
+        for fc in (1, 2)
+        for e in range(24)
+    ]
+
+
+def test_numeric_expert_order_and_group_rank_order(module):
+    infos = sorted(sources(), key=lambda info: info.name)
+    buckets = module.bucket_qwen35_experts(infos, 24, [7, 2])
+    assert len(buckets) == 6
+    assert [bucket[0].name.split("layers.")[1].split(".")[0] for bucket in buckets] == ["0", "0", "2", "2", "10", "10"]
+    for bucket in buckets:
+        assert [int(info.name.rsplit("weight", 1)[1]) for info in bucket] == list(range(24))
+        assert [info.src_rank for info in bucket] == [7] * 12 + [2] * 12
+
+
+@pytest.mark.parametrize("mutation", ["missing", "duplicate", "owner", "shape", "dtype", "mtp", "projection"])
+def test_invalid_source_layout_fails_before_payload_collectives(module, mutation):
+    infos = sources()
+    if mutation == "missing":
+        infos.pop()
+    elif mutation == "duplicate":
+        infos.append(infos[0])
+    elif mutation == "owner":
+        infos[0].src_rank = 2
+    elif mutation == "shape":
+        infos[0].shape = (6, 8)
+    elif mutation == "dtype":
+        infos[0].dtype = torch.float32
+    elif mutation == "mtp":
+        infos[0].name = infos[0].name.replace("decoder.layers.0", "mtp.layers.0.transformer_layer")
+    else:
+        infos = [info for info in infos if ".linear_fc2." not in info.name]
+    with pytest.raises(ValueError):
+        module.bucket_qwen35_experts(infos, 24, [7, 2])
+
+
+def peers(projection):
+    output = []
+    for owner in range(2):
+        items = []
+        for expert in range(owner * 12, (owner + 1) * 12):
+            for part in ("gate", "up") if projection == "gate_up" else ("down",):
+                shape = (4, 6) if part != "down" else (6, 4)
+                tensor = (torch.arange(24).reshape(shape) + expert * 32 + (1000 if part == "up" else 0)).to(
+                    torch.bfloat16
+                )
+                name = f"model.language_model.layers.2.mlp.experts.{expert}.{part}_proj.weight"
+                items.append((name, tensor))
+        output.append(sorted(items))
+    return output
+
+
+def flattened(peer, projection):
+    by_name = dict(peer)
+    ids = sorted({int(name.split(".experts.")[1].split(".")[0]) for name in by_name})
+    return torch.cat(
+        [
+            by_name[f"model.language_model.layers.2.mlp.experts.{expert}.{part}_proj.weight"].reshape(-1)
+            for expert in ids
+            for part in (("gate", "up") if projection == "gate_up" else ("down",))
+        ]
+    )
+
+
+@pytest.mark.parametrize("projection", ["gate_up", "down"])
+@pytest.mark.parametrize("rank", [0, 1])
+def test_packed_values_lifetime_and_fresh_group(module, monkeypatch, projection, rank):
+    inputs = peers(projection)
+    calls = []
+    monkeypatch.setattr(module.dist, "get_process_group_ranks", lambda group: group.ranks)
+    monkeypatch.setattr(module.dist, "get_rank", lambda group: group.rank)
+
+    def metadata(result, local, *, group):
+        assert local == module._layout(inputs[group.rank])
+        result[:] = [module._layout(peer) for peer in inputs]
+        calls.append(("metadata", group))
+
+    def gather(output, local, *, group, async_op):
+        assert async_op and torch.equal(local, flattened(inputs[group.rank], projection))
+        output.copy_(torch.cat([flattened(peer, projection) for peer in inputs]))
+        calls.append(("payload", group))
+        return SimpleNamespace(wait=lambda: calls.append(("wait", group)))
+
+    monkeypatch.setattr(module.dist, "all_gather_object", metadata)
+    monkeypatch.setattr(module.dist, "all_gather_into_tensor", gather)
+    plan = module.Qwen35PackedExpertGather(24)
+    first_group = SimpleNamespace(ranks=[7, 2], rank=rank)
+    [(name, value)] = plan(inputs[rank], group=first_group, device="cpu")
+    assert name == f"model.language_model.layers.2.mlp.experts.{projection}_proj"
+    assert tuple(value.shape) == ((24, 8, 6) if projection == "gate_up" else (24, 6, 4))
+    expected = torch.cat([flattened(peer, projection) for peer in inputs]).view(value.shape)
+    assert torch.equal(value, expected)
+    frozen = value.clone()
+    for peer in inputs:
+        for _, tensor in peer:
+            tensor.add_(16)
+    second_group = SimpleNamespace(ranks=[7, 2], rank=rank)
+    [(_, updated)] = plan(inputs[rank], group=second_group, device="cpu")
+    assert torch.equal(value, frozen) and not torch.equal(value, updated)
+    assert torch.equal(updated, torch.cat([flattened(peer, projection) for peer in inputs]).view(value.shape))
+    assert sum(kind == "metadata" for kind, _ in calls) == 1
+    assert [group for kind, group in calls if kind == "payload"] == [first_group, second_group]
+    assert not any(isinstance(v, torch.Tensor) or v is first_group for v in vars(plan).values())
+    with pytest.raises(ValueError, match="layout changed"):
+        plan(inputs[rank][:-1], group=second_group, device="cpu")
+
+
+@pytest.mark.parametrize("mutation", ["wrong_owner", "missing_up", "duplicate", "mixed_layer", "mixed_projection"])
+def test_invalid_converted_layout(module, mutation):
+    inputs = peers("gate_up")
+    if mutation == "wrong_owner":
+        inputs.reverse()
+    elif mutation == "missing_up":
+        inputs[1].pop()
+    elif mutation == "duplicate":
+        inputs[0].append(inputs[0][0])
+    elif mutation == "mixed_layer":
+        name, tensor = inputs[0][0]
+        inputs[0][0] = (name.replace("layers.2", "layers.3"), tensor)
+    else:
+        inputs[1] = peers("down")[1]
+    plan = module.Qwen35PackedExpertGather(24)
+    with pytest.raises(ValueError):
+        plan._prepare([module._layout(peer) for peer in inputs], (7, 2))

From d37219d880c319dbe737e3c98d24db425912d9f6 Mon Sep 17 00:00:00 2001
From: yuanlehome <yuanlehome@163.com>
Date: Mon, 19 Oct 2026 20:09:31 +0800
Subject: [PATCH 2/6] refactor(megatron): generalize automatic weight sync

Select expert gathering, TP-local loading and inline log-prob reuse from
runtime configuration and receiver capabilities. Consolidate expert layout
and IPC handling, remove experiment-only switches, and retain bounded
Bridge fallbacks for unsupported layouts.

Preserve the accepted batch512, packing8192 and sgl-router configuration
with regression coverage for loading, fallback and process-group reuse.
---
 relax/backends/megatron/actor.py              |  91 +-----
 relax/backends/megatron/loss.py               |   6 +-
 .../weight_update/bf16_expert_gather.py       | 109 -------
 .../weight_update/bridge_converter.py         |  24 ++
 .../backends/megatron/weight_update/common.py | 267 ++++++----------
 ...n35_packed_experts.py => expert_gather.py} | 116 +++++--
 .../megatron/weight_update/expert_layout.py   |  96 ++++++
 .../hf_weight_iterator_bridge.py              |  88 ++++--
 relax/backends/megatron/weight_update/ipc.py  |  52 ++++
 .../weight_update/qwen35_direct_experts.py    | 125 --------
 .../update_weight_from_tensor.py              | 284 ++++++++++--------
 relax/backends/sglang/sglang_engine.py        |  16 +-
 relax/utils/training/inline_log_probs.py      |  87 ++++++
 .../run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh   |   5 -
 .../test_colocate_inline_log_probs.py         |  75 ++---
 .../test_async_gdn_weight_gather.py           |  62 ++--
 .../weight_update/test_auto_weight_sync.py    | 232 ++++++++++++++
 .../weight_update/test_bf16_expert_gather.py  |  10 +-
 .../weight_update/test_expert_capabilities.py |  26 ++
 .../weight_update/test_expert_layout.py       | 207 +++++++++++++
 ...cked_experts.py => test_packed_experts.py} |  40 ++-
 .../test_qwen35_direct_experts.py             | 143 ---------
 22 files changed, 1246 insertions(+), 915 deletions(-)
 delete mode 100644 relax/backends/megatron/weight_update/bf16_expert_gather.py
 rename relax/backends/megatron/weight_update/{qwen35_packed_experts.py => expert_gather.py} (58%)
 create mode 100644 relax/backends/megatron/weight_update/expert_layout.py
 create mode 100644 relax/backends/megatron/weight_update/ipc.py
 delete mode 100644 relax/backends/megatron/weight_update/qwen35_direct_experts.py
 create mode 100644 relax/utils/training/inline_log_probs.py
 create mode 100644 tests/backends/megatron/weight_update/test_auto_weight_sync.py
 create mode 100644 tests/backends/megatron/weight_update/test_expert_capabilities.py
 create mode 100644 tests/backends/megatron/weight_update/test_expert_layout.py
 rename tests/backends/megatron/weight_update/{test_qwen35_packed_experts.py => test_packed_experts.py} (82%)
 delete mode 100644 tests/backends/megatron/weight_update/test_qwen35_direct_experts.py

diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py
index dcbbf2057..8675fb762 100644
--- a/relax/backends/megatron/actor.py
+++ b/relax/backends/megatron/actor.py
@@ -91,6 +91,7 @@
 from relax.utils.tracking_utils import init_tracking
 from relax.utils.training import train_dump_utils
 from relax.utils.training.data_fields import build_data_fields
+from relax.utils.training.inline_log_probs import can_reuse_colocate_log_probs
 from relax.utils.training.routing_replay import RoutingReplay
 from relax.utils.types import RolloutBatch
 from relax.utils.utils import (
@@ -1920,84 +1921,6 @@ def train_critic(self, rollout_id: int, rollout_data: RolloutBatch) -> None:
         if self._per_step_rollout and self.args.offload_train:
             self.sleep()
 
-    def _inline_colocate_old_log_probs(self, num_microbatches: list[int]) -> bool:
-        """Validate the single-update GRPO path before deriving old log-probs
-        inline."""
-        args = self.args
-        if os.environ.get("RELAX_COLOCATE_INLINE_LOG_PROBS") != "1":
-            return False
-
-        # A second update, stale rollout, or mutable forward would invalidate old=current.
-        # Fail closed: running the old forward cannot undo true_on_policy_mode in the loss.
-        required = {
-            "colocate": True,
-            "true_on_policy_mode": True,
-            "fully_async": False,
-            "hybrid": False,
-            "advantage_estimator": "grpo",
-            "loss_type": "policy_loss",
-            "compute_advantages_and_returns": True,
-            "num_iters_per_train_update": 1,
-            "max_staleness": 0,
-            "update_weights_interval": 1,
-            "partial_rollout": False,
-            "rl_partial_rollouts": False,
-            "use_dynamic_global_batch_size": False,
-            "context_parallel_size": 1,
-            "pipeline_model_parallel_size": 1,
-            "dynamic_context_parallel": False,
-            "hidden_dropout": 0.0,
-            "attention_dropout": 0.0,
-            "bf16": True,
-            "fp16": False,
-            "fp8": None,
-            "fp4": None,
-            "fp8_param_gather": False,
-            "moe_input_jitter_eps": None,
-            "moe_router_enable_expert_bias": False,
-            "moe_router_force_load_balancing": False,
-            "moe_router_force_biased": None,
-            "moe_expert_capacity_factor": None,
-            "moe_aux_loss_coeff": 0.0,
-            "kl_coef": 0.0,
-            "use_kl_loss": False,
-            "use_opd": False,
-            "use_critic": False,
-            "keep_old_actor": False,
-            "use_rollout_logprobs": False,
-            "use_routing_replay": True,
-            "use_rollout_routing_replay": True,
-            "use_rollout_indexer_replay": False,
-            "enable_mtp_training": False,
-            "use_opsm": False,
-            "use_tis": True,
-            "calculate_per_token_loss": True,
-            "entropy_coef": 0.0,
-            "recompute_loss_function": False,
-            "custom_tis_function_path": None,
-            "custom_pg_loss_reducer_function_path": None,
-            "custom_loss_function_path": None,
-            "custom_model_provider_path": None,
-            "custom_megatron_before_log_prob_hook_path": None,
-            "custom_megatron_before_train_step_hook_path": None,
-        }
-        invalid = [name for name, expected in required.items() if getattr(args, name, None) != expected]
-        if args.rollout_batch_size * args.n_samples_per_prompt != args.global_batch_size:
-            invalid.append("rollout size must equal global_batch_size")
-        if len(num_microbatches) != 1:
-            invalid.append("exactly one optimizer step per rollout is required")
-        if getattr(args, "num_steps_per_rollout", None) not in (None, 1):
-            invalid.append("num_steps_per_rollout")
-        if args.moe_router_load_balancing_type not in ("none", ["none"]):
-            invalid.append("moe_router_load_balancing_type")
-        if set(self.weights_backuper.backup_tags) - {"actor"}:
-            invalid.append("reference, teacher, or old-actor weights")
-        if self.rollout_data_postprocess is not None:
-            invalid.append("rollout_data_postprocess")
-        if invalid:
-            raise ValueError("Unsupported colocate true-on-policy configuration: " + ", ".join(invalid))
-        return True
-
     def train_actor(
         self,
         rollout_id: int,
@@ -2032,10 +1955,12 @@ def train_actor(
             data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
         else:
             data_iterator, num_microbatches = prepared_data_iterator, prepared_num_microbatches
-        inline_old_log_probs = os.environ.get(
-            "RELAX_COLOCATE_INLINE_LOG_PROBS"
-        ) == "1" and self._inline_colocate_old_log_probs(num_microbatches)
-        self.args._colocate_inline_log_probs = inline_old_log_probs
+        inline_old_log_probs = can_reuse_colocate_log_probs(
+            self.args,
+            num_microbatches,
+            backup_tags=self.weights_backuper.backup_tags,
+            rollout_data_postprocess=self.rollout_data_postprocess,
+        )
         # Create a separate iterator with a larger token budget for ref/teacher log-probs
         if (
             not is_offline_mode(self.args)
@@ -2138,7 +2063,7 @@ def train_actor(
                 # because we may need normalize the whole rollout.
                 capture_hooks.begin_rollout_for(self.args, rollout_id)
                 try:
-                    compute_advantages_and_returns(self.args, rollout_data)
+                    compute_advantages_and_returns(self.args, rollout_data, inline_log_probs=inline_old_log_probs)
                     capture_hooks.capture_rollout_advantage(rollout_data=rollout_data, args=self.args)
                 finally:
                     capture_hooks.end_rollout_for()
diff --git a/relax/backends/megatron/loss.py b/relax/backends/megatron/loss.py
index ae8fcabb2..e1d739757 100644
--- a/relax/backends/megatron/loss.py
+++ b/relax/backends/megatron/loss.py
@@ -564,7 +564,9 @@ def get_values(
     return torch.empty((0,), device=logits.device), res
 
 
-def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch) -> None:
+def compute_advantages_and_returns(
+    args: Namespace, rollout_data: RolloutBatch, *, inline_log_probs: bool = False
+) -> None:
     """Compute advantages and returns in-place based on
     `args.advantage_estimator`.
 
@@ -611,7 +613,7 @@ def compute_advantages_and_returns(args: Namespace, rollout_data: RolloutBatch)
     if args.kl_coef == 0 or not log_probs:
         # when kl_coef is 0, we won't compute ref_log_prob
         xs = log_probs if log_probs is not None else values
-        if xs is None and getattr(args, "_colocate_inline_log_probs", False) and args.kl_coef == 0:
+        if xs is None and inline_log_probs and args.kl_coef == 0:
             # Only shape/device are used. Sampling probabilities are never PPO old log-probs.
             xs = rollout_data.get("rollout_log_probs")
             if xs is None:
diff --git a/relax/backends/megatron/weight_update/bf16_expert_gather.py b/relax/backends/megatron/weight_update/bf16_expert_gather.py
deleted file mode 100644
index 7f6b3b04b..000000000
--- a/relax/backends/megatron/weight_update/bf16_expert_gather.py
+++ /dev/null
@@ -1,109 +0,0 @@
-# Copyright (c) 2026 Relax Authors. All Rights Reserved.
-
-"""Owner-balanced BF16 expert export, following Miles' packed gather design.
-
-Layout descriptors survive updates; GPU buffers and current process groups do
-not. The caller supplies the live EP group after each offload/reload cycle.
-"""
-
-from collections.abc import Sequence
-from typing import TYPE_CHECKING
-
-import torch
-import torch.distributed as dist
-
-
-if TYPE_CHECKING:
-    from relax.utils.types import ParamInfo
-
-
-NamedTensor = tuple[str, torch.Tensor]
-TensorLayout = tuple[str, tuple[int, ...], int]
-
-
-def bucket_experts_by_owner(
-    infos: Sequence["ParamInfo"], buffer_size: int, ranks: Sequence[int]
-) -> list[list["ParamInfo"]]:
-    """Pack each owner's share so one gathered round stays within the
-    budget."""
-    budget = buffer_size // len(ranks)
-    owners: dict[int, list] = {rank: [] for rank in ranks}
-    for info in infos:
-        if info.src_rank not in owners or not 0 < info.size <= budget:
-            raise ValueError(f"Unsupported expert owner or byte size: {info.name}")
-        owners[info.src_rank].append(info)
-    owner_batches = []
-    for rank in ranks:
-        batches, pending, size = [], [], 0
-        for info in owners[rank]:
-            if pending and size + info.size > budget:
-                batches.append(pending)
-                pending, size = [], 0
-            pending.append(info)
-            size += info.size
-        if pending:
-            batches.append(pending)
-        owner_batches.append(batches)
-    return [
-        [info for batches in owner_batches if index < len(batches) for info in batches[index]]
-        for index in range(max(map(len, owner_batches), default=0))
-    ]
-
-
-def _layout(tensors: Sequence[NamedTensor]) -> tuple[TensorLayout, ...]:
-    if any(tensor.dtype != torch.bfloat16 for _, tensor in tensors):
-        raise ValueError("Packed BF16 expert gather requires BF16 tensors")
-    return tuple((name, tuple(tensor.shape), tensor.numel()) for name, tensor in tensors)
-
-
-class Bf16ExpertGather:
-    """Cache CPU layouts; allocate independent receive storage on every
-    call."""
-
-    def __init__(self) -> None:
-        self._ranks: tuple[int, ...] | None = None
-        self._layouts: tuple[tuple[TensorLayout, ...], ...] | None = None
-        self._sizes: tuple[int, ...] = ()
-
-    def __call__(
-        self, tensors: Sequence[NamedTensor], *, group: dist.ProcessGroup, device: torch.device
-    ) -> list[NamedTensor]:
-        ranks = tuple(dist.get_process_group_ranks(group))
-        local_index = dist.get_rank(group=group)
-        layout = _layout(tensors)
-        if self._layouts is None:
-            layouts = [None] * len(ranks)
-            dist.all_gather_object(layouts, layout, group=group)
-            self._layouts = tuple(layouts)
-            self._ranks = ranks
-            self._sizes = tuple(sum(entry[2] for entry in peer) for peer in self._layouts)
-        elif ranks != self._ranks or layout != self._layouts[local_index]:
-            raise ValueError("Expert ownership or tensor layout changed; rebuild the gather plan")
-
-        storage = torch.empty(sum(self._sizes), dtype=torch.bfloat16, device=device)
-        payloads = list(storage.split(self._sizes))
-        local = payloads[local_index]
-        offset = 0
-        for (_, tensor), (_, _, size) in zip(tensors, layout, strict=True):
-            local.narrow(0, offset, size).copy_(tensor.reshape(-1))
-            offset += size
-
-        sources = [index for index, size in enumerate(self._sizes) if size]
-        handle = None
-        if len(sources) == 1:
-            src = sources[0]
-            handle = dist.broadcast(payloads[src], src=ranks[src], group=group, async_op=True)
-        elif sources and len(set(self._sizes)) == 1:
-            handle = dist.all_gather_into_tensor(storage, local, group=group, async_op=True)
-        elif sources:
-            handle = dist.all_gather(payloads, local, group=group, async_op=True)
-
-        result = []
-        for peer, payload in zip(self._layouts, payloads, strict=True):
-            offset = 0
-            for name, shape, size in peer:
-                result.append((name, payload.narrow(0, offset, size).view(shape)))
-                offset += size
-        if handle is not None:
-            handle.wait()
-        return result
diff --git a/relax/backends/megatron/weight_update/bridge_converter.py b/relax/backends/megatron/weight_update/bridge_converter.py
index 0954fe8bf..cbea9e042 100644
--- a/relax/backends/megatron/weight_update/bridge_converter.py
+++ b/relax/backends/megatron/weight_update/bridge_converter.py
@@ -428,6 +428,30 @@ def _lookup_task(self, global_name: str) -> tuple[str, Any]:
 
         return global_name, None
 
+    def supports_packed_experts(self, names: Collection[str]) -> bool:
+        """Whether mappings export canonical gated experts accepted by
+        FusedMoE."""
+        for name in names:
+            global_name = _normalize_global_name(name)
+            task = self._bridge_task_map.get(global_name)
+            mapping = task.mapping if task is not None else self._lookup_mapping(global_name)[1]
+            hf_param = getattr(mapping, "hf_param", None)
+            outputs = list(hf_param.values()) if isinstance(hf_param, dict) else [hf_param]
+            projections = set()
+            for output in outputs:
+                match = re.fullmatch(
+                    r"model\.(?:language_model\.)?layers\.\d+\.mlp\.experts\.(?:\d+\.)?"
+                    r"(gate|up|down|gate_up)_proj(?:\.weight)?",
+                    str(output),
+                )
+                if match is None:
+                    return False
+                projections.add(match[1])
+            expected = ({"gate", "up"}, {"gate_up"}) if ".linear_fc1." in name else ({"down"},)
+            if projections not in expected:
+                return False
+        return True
+
     def _lookup_mapping(self, global_name: str) -> tuple[str, Any]:
         """Resolve ``global_name`` straight from the mapping registry.
 
diff --git a/relax/backends/megatron/weight_update/common.py b/relax/backends/megatron/weight_update/common.py
index 277322a51..ab22b8ced 100644
--- a/relax/backends/megatron/weight_update/common.py
+++ b/relax/backends/megatron/weight_update/common.py
@@ -17,118 +17,88 @@
 from relax.utils.types import ParamInfo
 
 
-def all_gather_param(args, name: str, param: torch.Tensor) -> torch.Tensor:
-    """All-gather TP-sharded param to full tensor.
-
-    expert_bias→param, non-TP/duplicated→param.data. Uses expert-TP for
-    ".experts.", else regular-TP. linear_fc1 rechunked (GLU), linear_fc2 dim
-    fix.
-    """
-    name = name.replace(".to_wrap.", ".")
+def gdn_component_sizes(config, tp_size: int, name: str) -> tuple[int, ...] | None:
+    """Component-major GDN layout; unrelated parameters use ordinary TP
+    metadata."""
+    if not any(key in name for key in ("self_attention.conv1d.weight", "self_attention.in_proj.weight")):
+        return None
+    config = getattr(config, "text_config", None) or config
+    if not all(
+        hasattr(config, key)
+        for key in ("linear_num_key_heads", "linear_num_value_heads", "linear_key_head_dim", "linear_value_head_dim")
+    ):
+        return None
+    if config.linear_num_key_heads % tp_size or config.linear_num_value_heads % tp_size:
+        raise ValueError("GDN head counts must be divisible by TP size")
+    qk = config.linear_key_head_dim * (config.linear_num_key_heads // tp_size)
+    value = config.linear_value_head_dim * (config.linear_num_value_heads // tp_size)
+    heads = config.linear_num_value_heads // tp_size
+    return (qk, qk, value) if "conv1d" in name else (qk, qk, value, value, heads, heads)
+
+
+def _restore_tp_weight(name, partitions, partition_dim, split_sizes=None):
+    if split_sizes is not None:
+        assert partition_dim == 0, "GDN components must be sharded along rows"
+        components = [part.split(split_sizes, dim=0) for part in partitions]
+        return torch.cat([parts[index] for index in range(len(split_sizes)) for parts in components], dim=0)
+    if "linear_fc1.weight" in name and "vision_model" not in name:
+        components = [part.chunk(2, dim=0) for part in partitions]
+        partitions = [part[0] for part in components] + [part[1] for part in components]
+        if is_npu_available:
+            partition_dim = 0
+    if "linear_fc2.weight" in name and "vision_model" not in name and partition_dim == 0:
+        partition_dim = 1
+    return torch.cat(partitions, dim=partition_dim)
 
-    if "expert_bias" in name:
-        return param
 
-    if ".audio_model." in name or ".vision_model." in name or ".video_model." in name:
+def _gather_tp_weights(named_params, *, split_sizes, async_op):
+    tasks, handles = [], []
+    for name, param in named_params:
+        name = name.replace(".to_wrap.", ".")
+        if "expert_bias" in name:
+            tasks.append((name, param, None, None))
+            continue
         if not hasattr(param, "tensor_model_parallel"):
-            # NOTE: qwen3.5 vision_model is a megatron module
-            return param
-
-    if not hasattr(param, "tensor_model_parallel"):
-        assert not isinstance(param, torch.nn.Parameter), f"{name} does not have tensor_model_parallel attribute"
-        # Ordinary persistent buffers are replicated. A buffer carrying
-        # Megatron TP metadata continues through the regular gather path.
-        return param
-
-    if not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated":
-        return param.data
+            multimodal = any(part in name for part in (".audio_model.", ".vision_model.", ".video_model."))
+            assert multimodal or not isinstance(param, torch.nn.Parameter), f"{name} lacks TP metadata"
+            tasks.append((name, param, None, None))
+            continue
+        if not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated":
+            tasks.append((name, param.data, None, None))
+            continue
+        if ".experts." in name:
+            size, group = mpu.get_expert_tensor_parallel_world_size(), mpu.get_expert_tensor_parallel_group()
+        else:
+            size, group = mpu.get_tensor_model_parallel_world_size(), mpu.get_tensor_model_parallel_group()
+        if "linear_fc1" not in name:
+            assert getattr(param, "partition_stride", 1) == 1, f"Unsupported partition stride for {name}"
+        partitions = [torch.empty_like(param.data) for _ in range(size)]
+        handle = dist.all_gather(partitions, param.data, group=group, async_op=async_op)
+        if handle is not None:
+            handles.append(handle)
+        tasks.append((name, None, partitions, param.partition_dim))
+    # All async collectives launch before any wait. Results remain in input order.
+    for handle in handles:
+        handle.wait()
+    return [
+        direct if partitions is None else _restore_tp_weight(name, partitions, dim, split_sizes.get(name))
+        for name, direct, partitions, dim in tasks
+    ]
 
-    if ".experts." in name:
-        tp_size = mpu.get_expert_tensor_parallel_world_size()
-        tp_group = mpu.get_expert_tensor_parallel_group()
-    else:
-        tp_size = mpu.get_tensor_model_parallel_world_size()
-        tp_group = mpu.get_tensor_model_parallel_group()
-
-    # NOTE(wuhuan): qwen3.5 GDN
-    if "self_attention.conv1d.weight" in name:
-        config = get_hf_config(args.hf_checkpoint).text_config
-        qk_dim = config.linear_key_head_dim * config.linear_num_key_heads
-        v_dim = config.linear_value_head_dim * config.linear_num_value_heads
-        qk_dim = qk_dim // tp_size
-        v_dim = v_dim // tp_size
-        q_shard_idx = torch.arange(qk_dim)
-        k_shard_idx = torch.arange(qk_dim) + qk_dim
-        v_shard_idx = torch.arange(v_dim) + qk_dim * 2
-
-        q = param[q_shard_idx]
-        k = param[k_shard_idx]
-        v = param[v_shard_idx]
-
-        # Gather each component across TP ranks
-        full_weights = []
-        for c in [q, k, v]:
-            if tp_size == 1:
-                full_weight = c
-            else:
-                gathered = [torch.empty_like(c.data) for _ in range(tp_size)]
-                dist.all_gather(gathered, c.data, group=tp_group)
-                full_weight = torch.cat(gathered, dim=0)
-            full_weights.append(full_weight)
-        return torch.cat(full_weights, dim=0)
-
-    # NOTE(wuhuan): qwen3.5 GDN
-    if "self_attention.in_proj.weight" in name:
-        config = get_hf_config(args.hf_checkpoint).text_config
-        qk_head_dim = config.linear_key_head_dim
-        v_head_dim = config.linear_value_head_dim
-        num_qk_heads_local_tp = config.linear_num_key_heads // tp_size
-        num_v_heads_local_tp = config.linear_num_value_heads // tp_size
-        qk_dim_local_tp = qk_head_dim * num_qk_heads_local_tp
-        v_dim_local_tp = v_head_dim * num_v_heads_local_tp
-
-        full_weights = []
-        for c in torch.split(
-            param,
-            [
-                qk_dim_local_tp,
-                qk_dim_local_tp,
-                v_dim_local_tp,
-                v_dim_local_tp,
-                num_v_heads_local_tp,
-                num_v_heads_local_tp,
-            ],
-            dim=0,
-        ):
-            gathered = [torch.empty_like(c.data) for _ in range(tp_size)]
-            dist.all_gather(gathered, c.data, group=tp_group)
-            full_weights.append(torch.cat(gathered, dim=0))
-        return torch.cat(full_weights, dim=0)
-
-    param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)]
-    dist.all_gather(param_partitions, param.data, group=tp_group)
-    partition_dim = param.partition_dim
-    # NOTE: Megatron-LM (megatron/core/transformer/mlp.py) now explicitly sets partition_stride=2
-    # for GLU/SwiGLU linear_fc1 layers to indicate interleaved [gate, up] TP layout.
-    # The rechunk logic below (chunk(2) + reorder) already handles this correctly,
-    # so we only assert stride==1 for non-GLU parameters.
-    if "linear_fc1" not in name:
-        assert param.partition_stride == 1, f"{param.partition_stride=} != 1 is not supported"
-    # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better?
-    # TODO: check only GLU is used.
-    if "linear_fc1.weight" in name and "vision_model" not in name:
-        param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
-        param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
-        # TODO:Temporary workaround for NPU to set partition_dim to 0
-        if is_npu_available:
-            partition_dim = 0
-    # this is bug in megatron's grouped moe.
-    if "linear_fc2.weight" in name and "vision_model" not in name:
-        if partition_dim == 0:
-            partition_dim = 1
 
-    param = torch.cat(param_partitions, dim=partition_dim)
-    return param
+def all_gather_param(args, name: str, param: torch.Tensor) -> torch.Tensor:
+    """Gather one parameter using the same layout rules as batched export."""
+    name = name.replace(".to_wrap.", ".")
+    split_sizes = {}
+    if (
+        getattr(param, "tensor_model_parallel", False)
+        and getattr(param, "parallel_mode", None) != "duplicated"
+        and any(key in name for key in ("self_attention.conv1d.weight", "self_attention.in_proj.weight"))
+    ):
+        split_sizes[name] = gdn_component_sizes(
+            get_hf_config(args.hf_checkpoint), mpu.get_tensor_model_parallel_world_size(), name
+        )
+    return _gather_tp_weights([(name, param)], split_sizes=split_sizes, async_op=False)[0]
 
 
 def all_gather_params_async(
@@ -136,82 +106,13 @@ def all_gather_params_async(
     *,
     gdn_split_sizes: Mapping[str, tuple[int, ...]] | None = None,
 ) -> list[torch.Tensor]:
-    """Parallel TP all-gather for multiple params.
-
-    Loop 1: for each TP param, allocate buffers +
-    dist.all_gather(async_op=True) on expert-TP/regular-TP group (skip
-    expert_bias/non-TP/duplicated). Loop 2: wait all NCCL handles (enables
-    overlap). Loop 3: concat partitions + apply GLU rechunk/MoE dim fix.
-    """
-    # Phase 1: Start all async all_gather operations
-    gather_tasks = []
-    handles = []
-
-    for info, param in param_infos_and_params:
-        # Prepare async all_gather. Strip the LoRA ``.to_wrap.`` infix so name-based
-        # dispatch matches the plain non-LoRA name (see all_gather_param). No-op otherwise.
-        name = info.name.replace(".to_wrap.", ".")
-        if "expert_bias" in name:
-            gather_tasks.append((info, param, None, None, None))
-            handles.append(None)
-        elif not param.tensor_model_parallel or getattr(param, "parallel_mode", None) == "duplicated":
-            gather_tasks.append((info, param.data, None, None, None))
-            handles.append(None)
-        else:
-            # Start async all_gather
-            if ".experts." in name:
-                tp_size = mpu.get_expert_tensor_parallel_world_size()
-                tp_group = mpu.get_expert_tensor_parallel_group()
-            else:
-                tp_size = mpu.get_tensor_model_parallel_world_size()
-                tp_group = mpu.get_tensor_model_parallel_group()
-
-            param_partitions = [torch.empty_like(param.data) for _ in range(tp_size)]
-            handle = dist.all_gather(param_partitions, param.data, group=tp_group, async_op=True)
-            gather_tasks.append((info, None, handle, param_partitions, param.partition_dim))
-            handles.append(handle)
-
-    # Phase 2: Wait for ALL async operations to complete at once
-    # This ensures maximum parallelism by not blocking on individual operations
-    for handle in handles:
-        if handle is not None:
-            handle.wait()
-
-    # Phase 3: Process all results after all communications are done
-    gathered_params = []
-    for info, direct_param, handle, param_partitions, partition_dim in gather_tasks:
-        if handle is None:
-            # No all_gather needed
-            param = direct_param
-        else:
-            # Process the gathered partitions (same logic as original all_gather_param)
-            assert partition_dim is not None, "partition_dim must be set for TP-sharded params"
-            name = info.name.replace(".to_wrap.", ".")
-            if gdn_split_sizes is not None and name in gdn_split_sizes:
-                # Each rank holds [Q_local,K_local,V_local,(Z_local,B_local,A_local)].
-                # Gather the complete shard once, then restore component-major HF
-                # order. Plain concatenation would interleave Q/K/V across ranks.
-                assert partition_dim == 0, "GDN components must be sharded along rows"
-                components = [p.split(gdn_split_sizes[name], dim=0) for p in param_partitions]
-                param = torch.cat(
-                    [parts[index] for index in range(len(gdn_split_sizes[name])) for parts in components], dim=0
-                )
-                gathered_params.append(param)
-                continue
-            # TODO: here we did an extra copy during concat, maybe merge this with convert_to_hf is better?
-            # TODO: check only GLU is used.
-            if "linear_fc1.weight" in name and "vision_model" not in name:
-                param_partitions = [p.chunk(2, dim=0) for p in param_partitions]
-                param_partitions = [p[0] for p in param_partitions] + [p[1] for p in param_partitions]
-            # this is bug in megatron's grouped moe.
-            if "linear_fc2.weight" in name and "vision_model" not in name:
-                if partition_dim == 0:
-                    partition_dim = 1
-            param = torch.cat(param_partitions, dim=partition_dim)
-
-        gathered_params.append(param)
-
-    return gathered_params
+    """Launch every parameter gather before waiting, then restore shared TP
+    layouts."""
+    return _gather_tp_weights(
+        [(info.name, param) for info, param in param_infos_and_params],
+        split_sizes=gdn_split_sizes or {},
+        async_op=True,
+    )
 
 
 def named_params_and_buffers(
diff --git a/relax/backends/megatron/weight_update/qwen35_packed_experts.py b/relax/backends/megatron/weight_update/expert_gather.py
similarity index 58%
rename from relax/backends/megatron/weight_update/qwen35_packed_experts.py
rename to relax/backends/megatron/weight_update/expert_gather.py
index 172889c8c..fca00d82b 100644
--- a/relax/backends/megatron/weight_update/qwen35_packed_experts.py
+++ b/relax/backends/megatron/weight_update/expert_gather.py
@@ -1,9 +1,9 @@
 # Copyright (c) 2026 Relax Authors. All Rights Reserved.
 
-"""Export Qwen3.5 BF16 experts in the checkpoint's native packed format.
+"""BF16 expert export using owner-balanced or projection-packed collectives.
 
-Each layer/projection is indivisible: a gate/up tensor may exceed the target
-IPC bucket size. Only CPU layout metadata is cached across weight updates.
+Layout descriptors survive updates; GPU buffers and current process groups do
+not. The caller supplies the live EP group after each offload/reload cycle.
 """
 
 import re
@@ -19,20 +19,108 @@
 
 
 NamedTensor = tuple[str, torch.Tensor]
-TensorLayout = tuple[str, tuple[int, ...]]
+TensorLayout = tuple[str, tuple[int, ...], int]
 
 
-_SOURCE = re.compile(r"^(.*\.language_model\.decoder\.layers\.(\d+)\.mlp\.experts)\.linear_fc([12])\.weight(\d+)$")
-_TARGET = re.compile(r"^(model\.language_model\.layers\.\d+\.mlp\.experts)\.(\d+)\.(gate|up|down)_proj\.weight$")
+def bucket_experts_by_owner(
+    infos: Sequence["ParamInfo"], buffer_size: int, ranks: Sequence[int]
+) -> list[list["ParamInfo"]]:
+    """Pack each owner's share so one gathered round stays within the
+    budget."""
+    budget = buffer_size // len(ranks)
+    owners: dict[int, list] = {rank: [] for rank in ranks}
+    for info in infos:
+        if info.src_rank not in owners or not 0 < info.size <= budget:
+            raise ValueError(f"Unsupported expert owner or byte size: {info.name}")
+        owners[info.src_rank].append(info)
+    owner_batches = []
+    for rank in ranks:
+        batches, pending, size = [], [], 0
+        for info in owners[rank]:
+            if pending and size + info.size > budget:
+                batches.append(pending)
+                pending, size = [], 0
+            pending.append(info)
+            size += info.size
+        if pending:
+            batches.append(pending)
+        owner_batches.append(batches)
+    return [
+        [info for batches in owner_batches if index < len(batches) for info in batches[index]]
+        for index in range(max(map(len, owner_batches), default=0))
+    ]
+
+
+def _layout(tensors: Sequence[NamedTensor]) -> tuple[TensorLayout, ...]:
+    if any(tensor.dtype != torch.bfloat16 for _, tensor in tensors):
+        raise ValueError("Packed BF16 expert gather requires BF16 tensors")
+    return tuple((name, tuple(tensor.shape), tensor.numel()) for name, tensor in tensors)
 
 
-def bucket_qwen35_experts(
+class Bf16ExpertGather:
+    """Cache CPU layouts; allocate independent receive storage on every
+    call."""
+
+    def __init__(self) -> None:
+        self._ranks: tuple[int, ...] | None = None
+        self._layouts: tuple[tuple[TensorLayout, ...], ...] | None = None
+        self._sizes: tuple[int, ...] = ()
+
+    def __call__(
+        self, tensors: Sequence[NamedTensor], *, group: dist.ProcessGroup, device: torch.device
+    ) -> list[NamedTensor]:
+        ranks = tuple(dist.get_process_group_ranks(group))
+        local_index = dist.get_rank(group=group)
+        layout = _layout(tensors)
+        if self._layouts is None:
+            layouts = [None] * len(ranks)
+            dist.all_gather_object(layouts, layout, group=group)
+            self._layouts = tuple(layouts)
+            self._ranks = ranks
+            self._sizes = tuple(sum(entry[2] for entry in peer) for peer in self._layouts)
+        elif ranks != self._ranks or layout != self._layouts[local_index]:
+            raise ValueError("Expert ownership or tensor layout changed; rebuild the gather plan")
+
+        storage = torch.empty(sum(self._sizes), dtype=torch.bfloat16, device=device)
+        payloads = list(storage.split(self._sizes))
+        local = payloads[local_index]
+        offset = 0
+        for (_, tensor), (_, _, size) in zip(tensors, layout, strict=True):
+            local.narrow(0, offset, size).copy_(tensor.reshape(-1))
+            offset += size
+
+        sources = [index for index, size in enumerate(self._sizes) if size]
+        handle = None
+        if len(sources) == 1:
+            src = sources[0]
+            handle = dist.broadcast(payloads[src], src=ranks[src], group=group, async_op=True)
+        elif sources and len(set(self._sizes)) == 1:
+            handle = dist.all_gather_into_tensor(storage, local, group=group, async_op=True)
+        elif sources:
+            handle = dist.all_gather(payloads, local, group=group, async_op=True)
+
+        result = []
+        for peer, payload in zip(self._layouts, payloads, strict=True):
+            offset = 0
+            for name, shape, size in peer:
+                result.append((name, payload.narrow(0, offset, size).view(shape)))
+                offset += size
+        if handle is not None:
+            handle.wait()
+        return result
+
+
+_SOURCE = re.compile(r"^(.*\.layers\.(\d+)\.mlp\.experts)\.linear_fc([12])\.weight(\d+)$")
+_TARGET = re.compile(r"^(.*\.layers\.\d+\.mlp\.experts)\.(\d+)\.(gate|up|down)_proj\.weight$")
+
+
+def bucket_expert_projections(
     infos: Sequence["ParamInfo"], num_experts: int, ranks: Sequence[int]
 ) -> list[list["ParamInfo"]]:
     """Validate complete, contiguous EP ownership and group by
     layer/projection."""
     if not ranks or num_experts <= 0 or num_experts % len(ranks):
-        raise ValueError("Packed Qwen experts require evenly partitioned EP ownership")
+        raise ValueError("Packed experts require evenly partitioned EP ownership")
     local_count = num_experts // len(ranks)
     groups = {}
     for info in infos:
@@ -69,13 +157,7 @@ def bucket_qwen35_experts(
     return buckets
 
 
-def _layout(tensors: Sequence[NamedTensor]) -> tuple[TensorLayout, ...]:
-    if any(tensor.dtype != torch.bfloat16 for _, tensor in tensors):
-        raise ValueError("Packed Qwen expert gather requires BF16 tensors")
-    return tuple((name, tuple(tensor.shape)) for name, tensor in tensors)
-
-
-class Qwen35PackedExpertGather:
+class PackedExpertGather:
     """Gather one native projection directly into global expert order."""
 
     def __init__(self, num_experts: int) -> None:
@@ -91,7 +173,7 @@ def _prepare(self, layouts: Sequence[tuple[TensorLayout, ...]], ranks: tuple[int
         orders, expected_prefix, expected_projections, expected_shape = [], None, None, None
         for index, layout in enumerate(layouts):
             entries = {}
-            for position, (name, shape) in enumerate(layout):
+            for position, (name, shape, _size) in enumerate(layout):
                 match = _TARGET.fullmatch(name)
                 if match is None or len(shape) != 2:
                     raise ValueError(f"Unsupported packed expert output: {name}")
@@ -131,7 +213,7 @@ def __call__(
         local_index = dist.get_rank(group=group)
         layout = _layout(tensors)
         if not ranks or self.num_experts <= 0 or self.num_experts % len(ranks):
-            raise ValueError("Packed Qwen experts require evenly partitioned EP ownership")
+            raise ValueError("Packed experts require evenly partitioned EP ownership")
         if self._layouts is None:
             layouts = [None] * len(ranks)
             dist.all_gather_object(layouts, layout, group=group)
diff --git a/relax/backends/megatron/weight_update/expert_layout.py b/relax/backends/megatron/weight_update/expert_layout.py
new file mode 100644
index 000000000..a68739b5c
--- /dev/null
+++ b/relax/backends/megatron/weight_update/expert_layout.py
@@ -0,0 +1,96 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+"""BF16 gated-expert layouts shared by Bridge export and SGLang FusedMoE.
+
+Names and tensor orientation come from Bridge's canonical per-expert output.
+This module only packs and shards that representation; it never converts raw
+Megatron weights or keeps process groups/GPU storage across updates.
+"""
+
+import json
+import re
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+
+NamedTensor = tuple[str, torch.Tensor]
+_PACKED = re.compile(r"^(.*\.layers\.\d+\.mlp\.experts)\.(gate_up_proj|down_proj)$")
+
+
+@dataclass(frozen=True)
+class DirectExpertLayout:
+    num_experts: int
+    hidden_size: int
+    intermediate_size: int
+    tp_size: int
+
+    @classmethod
+    def from_receiver(cls, config: Mapping[str, Any]) -> "DirectExpertLayout":
+        required = dict(
+            ep_size=1,
+            pp_size=1,
+            moe_dp_size=1,
+            moe_runner_backend="triton",
+            disable_shared_experts_fusion=True,
+            dtype="bfloat16",
+        )
+        if any(config.get(key) != value for key, value in required.items()):
+            raise ValueError("Direct experts require BF16 Triton, EP1/PP1/MoE-DP1 and unfused shared experts")
+        if (
+            config.get("enable_eplb")
+            or config.get("ep_num_redundant_experts", 0)
+            or config.get("init_expert_location") not in (None, "trivial")
+            or config.get("ep_join_mode") is not None
+            or config.get("quantization") is not None
+            or config.get("quantization_config")
+            or config.get("enable_lora")
+            or config.get("kt_weight_path")
+            or config.get("speculative_algorithm") is not None
+            or config.get("enforce_shared_experts_fusion")
+            or json.loads(config.get("json_model_override_args") or "{}")
+        ):
+            raise ValueError("Direct experts do not support quantization, remapping, LoRA or model overrides")
+        dimensions = tuple(
+            config.get(key) for key in ("num_experts", "hidden_size", "moe_intermediate_size", "tp_size")
+        )
+        if any(type(value) is not int or value <= 0 for value in dimensions):
+            raise ValueError("Direct expert geometry must be explicitly specified and positive")
+        layout = cls(*dimensions)
+        if layout.intermediate_size % layout.tp_size:
+            raise ValueError("Expert intermediate size must be divisible by receiver TP size")
+        return layout
+
+    def shard(self, tensor: torch.Tensor, projection: str, tp_rank: int) -> torch.Tensor:
+        if not 0 <= tp_rank < self.tp_size or tensor.dtype != torch.bfloat16:
+            raise ValueError("Invalid direct expert TP rank or dtype")
+        e, h, i = self.num_experts, self.hidden_size, self.intermediate_size
+        expected = (e, 2 * i, h) if projection == "gate_up_proj" else (e, h, i)
+        if tuple(tensor.shape) != expected or projection not in ("gate_up_proj", "down_proj"):
+            raise ValueError(f"Unsupported expert projection or shape: {projection}: {tuple(tensor.shape)}")
+        width = i // self.tp_size
+        start = tp_rank * width
+        if projection == "gate_up_proj":
+            return torch.cat((tensor[:, start : start + width], tensor[:, i + start : i + start + width]), dim=1)
+        return tensor[:, :, start : start + width].contiguous()
+
+
+def direct_expert_chunk(
+    named_tensors: Sequence[NamedTensor], layout: DirectExpertLayout, tp_rank: int
+) -> list[NamedTensor] | None:
+    if len(named_tensors) != 1:
+        return None
+    name, tensor = named_tensors[0]
+    match = _PACKED.fullmatch(name)
+    if match is None:
+        return None
+    prefix, projection = match.groups()
+    # The multimodal HF namespace is stripped by the Qwen3.5 SGLang loader.
+    # Standard text MoE namespaces (including DeepSeek V2/V3) are unchanged.
+    prefix = prefix.replace("model.language_model.layers.", "model.layers.", 1)
+    if not re.fullmatch(r"model\.layers\.\d+\.mlp\.experts", prefix):
+        raise ValueError(f"Unsupported native expert namespace: {prefix}")
+    target = "w13_weight" if projection == "gate_up_proj" else "w2_weight"
+    return [(f"{prefix}.{target}", layout.shard(tensor, projection, tp_rank))]
diff --git a/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py b/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py
index 5440f84d0..5c0090c7d 100644
--- a/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py
+++ b/relax/backends/megatron/weight_update/hf_weight_iterator_bridge.py
@@ -1,7 +1,6 @@
 # Copyright (c) 2026 Relax Authors. All Rights Reserved.
 
 import dataclasses
-import os
 import re
 import time
 from collections import OrderedDict
@@ -31,12 +30,12 @@
 from relax.utils.misc import get_hf_config
 from relax.utils.types import ParamInfo
 
-from .bf16_expert_gather import Bf16ExpertGather, bucket_experts_by_owner
 from .bridge_converter import BridgeConverter
-from .common import all_gather_param, all_gather_params_async, named_params_and_buffers
+from .common import all_gather_param, all_gather_params_async, gdn_component_sizes, named_params_and_buffers
+from .expert_gather import Bf16ExpertGather, PackedExpertGather, bucket_expert_projections, bucket_experts_by_owner
+from .expert_layout import DirectExpertLayout
 from .expert_routing import ExpertLayout, ExpertRouter
 from .hf_weight_iterator_base import HfWeightIteratorBase
-from .qwen35_packed_experts import Qwen35PackedExpertGather, bucket_qwen35_experts
 
 
 logger = get_logger(__name__)
@@ -108,8 +107,7 @@ def __init__(self, *args, **kwargs):
         self._expert_broadcast_caches = [({}, {}) for _ in self._expert_buckets]
         self._expert_router: ExpertRouter | None = None
         self._async_weight_gather = (
-            os.environ.get("RELAX_ASYNC_WEIGHT_GATHER", "0") == "1"
-            and torch.cuda.is_available()
+            torch.cuda.is_available()
             and not self.quantization_config
             and not self.lora_merge_mode
             and not self.lora_adapter_mode
@@ -120,29 +118,27 @@ def __init__(self, *args, **kwargs):
                 info.name.replace(".to_wrap.", ".")
                 for bucket in self._non_expert_buckets
                 for info in bucket
-                if "self_attention.conv1d.weight" in info.name or "self_attention.in_proj.weight" in info.name
+                if info.attrs.get("tensor_model_parallel")
+                and info.attrs.get("parallel_mode") != "duplicated"
+                and ("self_attention.conv1d.weight" in info.name or "self_attention.in_proj.weight" in info.name)
             }
             if names:
-                config = get_hf_config(self.args.hf_checkpoint).text_config
+                config = get_hf_config(self.args.hf_checkpoint)
                 tp_size = mpu.get_tensor_model_parallel_world_size()
-                qk = config.linear_key_head_dim * (config.linear_num_key_heads // tp_size)
-                value = config.linear_value_head_dim * (config.linear_num_value_heads // tp_size)
-                heads = config.linear_num_value_heads // tp_size
-                self._gdn_split_sizes = {
-                    name: (qk, qk, value) if "conv1d" in name else (qk, qk, value, value, heads, heads)
-                    for name in names
-                }
+                self._gdn_split_sizes = {name: gdn_component_sizes(config, tp_size, name) for name in names}
         self._bf16_expert_buckets = None
         self._bf16_expert_gathers = []
-        # Opt in only for the validated, unquantized ETP1/PP1 export. Other
-        # layouts retain the existing Bridge path, including LoRA and routing.
+        self._packed_expert_sync = False
+        self._direct_expert_layout = None
+        # Complete BF16 experts are owned locally at ETP1/PP1. Other layouts
+        # continue through Bridge's existing PP/TP/quantization/LoRA path.
         if (
-            os.environ.get("RELAX_BF16_EXPERT_GATHER", "0") == "1"
-            and torch.cuda.is_available()
+            torch.cuda.is_available()
             and not self.quantization_config
             and not self.lora_merge_mode
             and not self.lora_adapter_mode
             and mpu.get_pipeline_model_parallel_world_size() == 1
+            and mpu.get_virtual_pipeline_model_parallel_world_size() in (None, 1)
             and mpu.get_expert_tensor_parallel_world_size() == 1
             and mpu.get_expert_model_parallel_world_size() > 1
             and self._expert_buckets
@@ -154,15 +150,6 @@ def __init__(self, *args, **kwargs):
             if all(info.src_rank in ranks and 0 < info.size <= budget for info in infos):
                 self._bf16_expert_buckets = bucket_experts_by_owner(infos, self.args.update_weight_buffer_size, ranks)
                 self._bf16_expert_gathers = [Bf16ExpertGather() for _ in self._bf16_expert_buckets]
-                if (
-                    os.environ.get("RELAX_QWEN35_PACKED_EXPERTS", "0") == "1"
-                    and mpu.get_virtual_pipeline_model_parallel_world_size() in (None, 1)
-                    and get_hf_config(self.args.hf_checkpoint).model_type == "qwen3_5_moe"
-                ):
-                    self._bf16_expert_buckets = bucket_qwen35_experts(infos, self.args.num_experts, ranks)
-                    self._bf16_expert_gathers = [
-                        Qwen35PackedExpertGather(self.args.num_experts) for _ in self._bf16_expert_buckets
-                    ]
         if self.lora_merge_mode:
             self.lora_alpha = self.args.lora_alpha
             self.lora_dim = self.args.lora_rank
@@ -179,6 +166,49 @@ def __init__(self, *args, **kwargs):
                 "DOES fold expert deltas for actor_fwd and so also needs ETP=1.)"
             )
 
+    def configure_direct_experts(self, layout: DirectExpertLayout | None) -> bool:
+        """Select packed export only after the colocated receiver layout is
+        known."""
+        if self._packed_expert_sync and layout == self._direct_expert_layout:
+            return True
+        if not self._packed_expert_sync and layout is None:
+            return False
+        self._packed_expert_sync = False
+        self._direct_expert_layout = None
+        if self._bf16_expert_buckets is None:
+            return False
+        infos = [info for bucket in self._bf16_expert_buckets for info in bucket]
+        ranks = dist.get_process_group_ranks(mpu.get_expert_model_parallel_group())
+        # Reconnect may disable a previous direct plan. Restore standard HF export.
+        self._bf16_expert_buckets = bucket_experts_by_owner(infos, self.args.update_weight_buffer_size, ranks)
+        self._bf16_expert_gathers = [Bf16ExpertGather() for _ in self._bf16_expert_buckets]
+        if layout is None or layout.num_experts != self.args.num_experts:
+            return False
+        if not self._bridge_converter.supports_packed_experts([info.name for info in infos]):
+            return False
+        try:
+            buckets = bucket_expert_projections(infos, layout.num_experts, ranks)
+        except ValueError:
+            return False
+        for bucket in buckets:
+            # A fused projection is indivisible. Permit at most two ordinary
+            # sync buffers (the existing overlap budget), then fall back to
+            # bounded per-expert HF export for larger MoE geometries.
+            if sum(info.size for info in bucket) > 2 * self.args.update_weight_buffer_size:
+                return False
+            expected = (
+                (2 * layout.intermediate_size, layout.hidden_size)
+                if ".linear_fc1." in bucket[0].name
+                else (layout.hidden_size, layout.intermediate_size)
+            )
+            if tuple(bucket[0].shape) != expected:
+                return False
+        self._bf16_expert_buckets = buckets
+        self._bf16_expert_gathers = [PackedExpertGather(layout.num_experts) for _ in buckets]
+        self._packed_expert_sync = True
+        self._direct_expert_layout = layout
+        return True
+
     def configure_expert_routing(
         self, layout: ExpertLayout, *, payload_group: dist.ProcessGroup, metadata_group: dist.ProcessGroup
     ) -> None:
@@ -277,7 +307,7 @@ def _iter_routed_expert_chunks(
 
     def get_hf_weight_chunks(self, megatron_local_weights):
         if self._bf16_expert_buckets is not None:
-            if isinstance(self._bf16_expert_gathers[0], Qwen35PackedExpertGather):
+            if self._packed_expert_sync:
                 yield from _chunk_with_mla_pairing(
                     self._iter_hf_params(megatron_local_weights, include_experts=False),
                     chunk_size=self.args.update_weight_buffer_size,
diff --git a/relax/backends/megatron/weight_update/ipc.py b/relax/backends/megatron/weight_update/ipc.py
new file mode 100644
index 000000000..4f5b36fec
--- /dev/null
+++ b/relax/backends/megatron/weight_update/ipc.py
@@ -0,0 +1,52 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+"""Format-aware dispatch for colocated SGLang weight updates."""
+
+from collections.abc import Callable, Sequence
+from typing import Any
+
+
+def dispatch_typed_buckets(
+    packets: Sequence[Sequence[tuple[str, str]]],
+    *,
+    engine: Any,
+    version: str | int | None,
+    serializer: Any,
+    empty_flattened: Callable[[], dict[str, Any]],
+    held: list[Any],
+    dispatch_errors: list[str] | None = None,
+) -> list[Any]:
+    """Keep a single request format even with empty or differently packed
+    ranks.
+
+    Format tags ride in the existing Gloo gather. Producers retain GPU storage
+    until all submitted readers complete, including partial dispatch failures.
+    """
+    if any(fmt not in ("direct", "flattened_bucket") for packet in packets for fmt, _ in packet):
+        raise ValueError("Unsupported weight IPC format")
+    refs, empty = [], {}
+    count = max((len(packet) for packet in packets), default=0)
+    for index in range(count):
+        formats = dict.fromkeys(packet[index][0] for packet in packets if index < len(packet))
+        for load_format in formats:
+            try:
+                serialized = []
+                for packet in packets:
+                    if index < len(packet) and packet[index][0] == load_format:
+                        serialized.append(packet[index][1])
+                    else:
+                        if load_format not in empty:
+                            payload = [] if load_format == "direct" else empty_flattened()
+                            held.append(payload)
+                            empty[load_format] = serializer.serialize(payload, output_str=True)
+                        serialized.append(empty[load_format])
+                kwargs = dict(serialized_named_tensors=serialized, load_format=load_format)
+                if version is not None:
+                    kwargs["weight_version"] = str(version)
+                refs.append(engine.update_weights_from_tensor.remote(**kwargs))
+            except Exception as error:
+                if dispatch_errors is None:
+                    raise
+                dispatch_errors.append(str(error))
+                return refs
+    return refs
diff --git a/relax/backends/megatron/weight_update/qwen35_direct_experts.py b/relax/backends/megatron/weight_update/qwen35_direct_experts.py
deleted file mode 100644
index a5b4cada9..000000000
--- a/relax/backends/megatron/weight_update/qwen35_direct_experts.py
+++ /dev/null
@@ -1,125 +0,0 @@
-# Copyright (c) 2026 Relax Authors. All Rights Reserved.
-
-"""Native BF16 expert placement for the validated Qwen3.5 TP2 Triton layout."""
-
-import json
-import re
-from collections.abc import Callable, Mapping, Sequence
-from typing import Any
-
-import torch
-
-
-NamedTensor = tuple[str, torch.Tensor]
-
-
-_PACKED = re.compile(r"model\.language_model\.layers\.(\d+)\.mlp\.experts\.(gate_up_proj|down_proj)")
-
-
-def validate_direct_receiver(config: Mapping[str, Any]) -> None:
-    required = dict(
-        tp_size=2,
-        ep_size=1,
-        pp_size=1,
-        moe_dp_size=1,
-        moe_runner_backend="triton",
-        disable_shared_experts_fusion=True,
-        dtype="bfloat16",
-    )
-    if any(config.get(key) != value for key, value in required.items()):
-        raise ValueError("Direct Qwen3.5 experts require explicit BF16 Triton, TP2/EP1/PP1 and unfused shared experts")
-    if (
-        config.get("enable_eplb")
-        or config.get("ep_num_redundant_experts", 0)
-        or config.get("init_expert_location") not in (None, "trivial")
-        or config.get("ep_join_mode") is not None
-        or config.get("quantization") is not None
-        or config.get("enable_lora")
-        or config.get("speculative_algorithm") is not None
-        or config.get("enforce_shared_experts_fusion")
-        or json.loads(config.get("json_model_override_args") or "{}")
-    ):
-        raise ValueError("Direct Qwen3.5 experts do not support quantization, remapping, LoRA or model overrides")
-
-
-def _shard_expert(tensor: torch.Tensor, projection: str, tp_rank: int) -> torch.Tensor:
-    if tp_rank not in (0, 1) or tensor.ndim != 3 or tensor.dtype != torch.bfloat16:
-        raise ValueError("Direct expert shards require BF16 rank-3 tensors and TP rank 0/1")
-    if projection == "gate_up_proj":
-        if tensor.shape[1] % 4:
-            raise ValueError("Gate/up rows must be divisible by four for TP2")
-        intermediate = tensor.shape[1] // 2
-        width = intermediate // 2
-        start = tp_rank * width
-        # TP slices gate and up separately, then keeps [gate_local, up_local].
-        return torch.cat(
-            (tensor[:, start : start + width], tensor[:, intermediate + start : intermediate + start + width]), dim=1
-        )
-    if projection == "down_proj":
-        if tensor.shape[2] % 2:
-            raise ValueError("Down columns must be divisible by two for TP2")
-        width = tensor.shape[2] // 2
-        return tensor[:, :, tp_rank * width : (tp_rank + 1) * width].contiguous()
-    raise ValueError(f"Unsupported expert projection: {projection}")
-
-
-def direct_expert_chunk(named_tensors: Sequence[NamedTensor], tp_rank: int) -> list[NamedTensor] | None:
-    if len(named_tensors) != 1:
-        return None
-    name, tensor = named_tensors[0]
-    match = _PACKED.fullmatch(name)
-    if match is None:
-        return None
-    layer, projection = match.groups()
-    expected = (256, 1024, 2048) if projection == "gate_up_proj" else (256, 2048, 512)
-    if tuple(tensor.shape) != expected:
-        raise ValueError(f"Unsupported Qwen3.5 expert shape: {name}: {tuple(tensor.shape)}")
-    target = "w13_weight" if projection == "gate_up_proj" else "w2_weight"
-    return [(f"model.layers.{layer}.mlp.experts.{target}", _shard_expert(tensor, projection, tp_rank))]
-
-
-def dispatch_typed_buckets(
-    packets: Sequence[Sequence[tuple[str, str]]],
-    *,
-    engine: Any,
-    version: str | int | None,
-    serializer: Any,
-    empty_flattened: Callable[[], dict[str, Any]],
-    held: list[Any],
-    dispatch_errors: list[str] | None = None,
-) -> list[Any]:
-    """Keep a single request format even with empty or differently packed
-    ranks.
-
-    Only enabled by the private direct-expert path. Format tags ride in the
-    existing Gloo gather; no new collective is needed. GPU payloads are already
-    held by their producers through the caller's normal ACK pipeline.
-    """
-    refs, empty = [], {}
-    count = max((len(packet) for packet in packets), default=0)
-    for index in range(count):
-        formats = dict.fromkeys(packet[index][0] for packet in packets if index < len(packet))
-        for load_format in formats:
-            if load_format not in ("direct", "flattened_bucket"):
-                raise ValueError(f"Unsupported IPC format: {load_format}")
-            serialized = []
-            for packet in packets:
-                if index < len(packet) and packet[index][0] == load_format:
-                    serialized.append(packet[index][1])
-                else:
-                    if load_format not in empty:
-                        payload = [] if load_format == "direct" else empty_flattened()
-                        held.append(payload)
-                        empty[load_format] = serializer.serialize(payload, output_str=True)
-                    serialized.append(empty[load_format])
-            kwargs = dict(serialized_named_tensors=serialized, load_format=load_format)
-            if version is not None:
-                kwargs["weight_version"] = str(version)
-            try:
-                refs.append(engine.update_weights_from_tensor.remote(**kwargs))
-            except Exception as error:
-                if dispatch_errors is None:
-                    raise
-                dispatch_errors.append(str(error))
-                return refs
-    return refs
diff --git a/relax/backends/megatron/weight_update/update_weight_from_tensor.py b/relax/backends/megatron/weight_update/update_weight_from_tensor.py
index 12ea98762..b5e391798 100644
--- a/relax/backends/megatron/weight_update/update_weight_from_tensor.py
+++ b/relax/backends/megatron/weight_update/update_weight_from_tensor.py
@@ -1,8 +1,5 @@
 # Copyright (c) 2026 Relax Authors. All Rights Reserved.
 
-import json
-import os
-import urllib.request
 from argparse import Namespace
 from collections.abc import Callable, Mapping, Sequence
 from contextlib import nullcontext
@@ -29,13 +26,13 @@
     is_lora_enabled,
     is_lora_merge_mode,
 )
-from relax.utils.misc import get_hf_config
 
 from ..sglang import FlattenedTensorBucket, MultiprocessingSerializer
+from .expert_layout import DirectExpertLayout, direct_expert_chunk
 from .expert_routing import ExpertLayout, validate_server_layout
 from .hf_weight_iterator_base import HfWeightIteratorBase
+from .ipc import dispatch_typed_buckets
 from .lora_adapter_sync import LoraAdapterSync
-from .qwen35_direct_experts import direct_expert_chunk, dispatch_typed_buckets, validate_direct_receiver
 from .update_weight_from_distributed import (
     connect_rollout_engines_from_distributed,
     disconnect_rollout_engines_from_distributed,
@@ -101,27 +98,10 @@ def __init__(
             args=args, model=model, model_name=model_name, quantization_config=quantization_config
         )
 
-        self._direct_qwen35_experts = os.environ.get("RELAX_QWEN35_DIRECT_EXPERTS", "0") == "1"
-        if self._direct_qwen35_experts:
-            config = get_hf_config(args.hf_checkpoint)
-            text_config = getattr(config, "text_config", None)
-            if (
-                not args.colocate
-                or self.lora_enabled
-                or quantization_config
-                or self._route_experts
-                or config.model_type != "qwen3_5_moe"
-                or text_config is None
-                or (text_config.hidden_size, text_config.moe_intermediate_size, text_config.num_experts)
-                != (2048, 512, 256)
-                or os.environ.get("RELAX_QWEN35_PACKED_EXPERTS", "0") != "1"
-                or not getattr(self._hf_weight_iterator, "_bf16_expert_buckets", None)
-                or mpu.get_pipeline_model_parallel_world_size() != 1
-                or mpu.get_expert_tensor_parallel_world_size() != 1
-            ):
-                raise ValueError("Direct expert IPC is limited to packed BF16 Qwen3.5-35B with PP1/ETP1")
+        self._direct_expert_layout: DirectExpertLayout | None = None
 
         self._ipc_gather_group = None
+        self._ipc_placement = None
         self._ipc_gather_src = None
         self._ipc_engine = None
         self._model_update_groups = None
@@ -199,6 +179,9 @@ def connect_rollout_engines(
 
         colocate_gpu_offsets = engine_gpu_offsets[:colocate_engine_nums]
         colocate_gpu_counts = engine_gpu_counts[:colocate_engine_nums]
+        placement = (tuple(colocate_gpu_offsets), tuple(colocate_gpu_counts))
+        if self._ipc_placement not in (None, placement):
+            raise ValueError("Weight sync cannot reconnect with changed IPC GPU placement")
 
         if self._route_experts:
             placement = (tuple(colocate_gpu_offsets), tuple(colocate_gpu_counts))
@@ -238,40 +221,101 @@ def connect_rollout_engines(
             )
             self._expert_routing_placement = placement
 
-        if getattr(self, "_direct_qwen35_experts", False):
-            if self.use_distribute or any(count != 2 for count in colocate_gpu_counts):
-                raise ValueError("Direct Qwen3.5 experts require colocated two-GPU engines")
-            status = [None]
-            if dist.get_rank(get_gloo_group()) == 0:
-                try:
-                    urls = ray.get([engine.get_url.remote() for engine in self.rollout_engines])
-                    for url in urls:
-                        with urllib.request.urlopen(url.rstrip("/") + "/server_info", timeout=30) as response:
-                            validate_direct_receiver(json.load(response))
-                    status[0] = (True, None)
-                except Exception as error:
-                    status[0] = (False, str(error))
-            dist.broadcast_object_list(status, src=0, group=get_gloo_group())
-            if not status[0][0]:
-                raise RuntimeError(f"Direct Qwen3.5 receiver validation failed: {status[0][1]}")
+        self._configure_direct_experts(colocate_gpu_counts)
 
         # Create IPC Gloo gather groups (only on first call; partitioning is
         # fixed across reconnects).
-        if self._ipc_gather_group is None:
+        if self._ipc_placement is None:
             for i in range(colocate_engine_nums):
                 group_ranks = list(range(colocate_gpu_offsets[i], colocate_gpu_offsets[i] + colocate_gpu_counts[i]))
                 new_group = dist.new_group(ranks=group_ranks, backend="gloo")
                 if dist.get_rank() in group_ranks:
                     self._ipc_gather_group = new_group
                     self._ipc_gather_src = colocate_gpu_offsets[i]
+            # Nonmembers also record creation: every world rank must enter
+            # new_group in the same order, including on engine reconnect.
+            self._ipc_placement = placement
 
         # Map training ranks to colocated engine actors.
+        self._ipc_engine = None
         for i, engine in enumerate(self.rollout_engines):
             start = colocate_gpu_offsets[i]
             end = start + colocate_gpu_counts[i]
             if start <= dist.get_rank() < end:
                 self._ipc_engine = engine
 
+    def _configure_direct_experts(self, gpu_counts: Sequence[int]) -> None:
+        """Negotiate one fast path using actual training and receiver
+        capabilities."""
+        configure = getattr(self._hf_weight_iterator, "configure_direct_experts", None)
+        self._direct_expert_layout = None
+        if (
+            not self.args.colocate
+            or self.use_distribute
+            or self._route_experts
+            or self.lora_enabled
+            or self.quantization_config
+        ):
+            if configure is not None:
+                configure(None)
+            return
+        group = get_gloo_group()
+        ready = torch.tensor(
+            [int(configure is not None and bool(getattr(self._hf_weight_iterator, "_bf16_expert_buckets", None)))],
+            dtype=torch.int32,
+            device="cpu",
+        )
+        dist.all_reduce(ready, op=dist.ReduceOp.MIN, group=group)
+        if not ready.item():
+            if configure is not None:
+                configure(None)
+            return
+        status = [None]
+        if dist.get_rank(group=group) == 0:
+            try:
+                status[0] = (
+                    ray.get([engine.get_expert_weight_sync_layout.remote() for engine in self.rollout_engines]),
+                    None,
+                )
+            except Exception as error:
+                status[0] = ([], str(error))
+        dist.broadcast_object_list(status, src=0, group=group)
+        receivers, reason = status[0]
+        layout = None
+        if reason is None:
+            try:
+                layouts = [DirectExpertLayout.from_receiver(config) for config in receivers]
+                if (
+                    not layouts
+                    or len(layouts) != len(gpu_counts)
+                    or any(item != layouts[0] for item in layouts)
+                    or any(item.tp_size != count for item, count in zip(layouts, gpu_counts, strict=True))
+                ):
+                    raise ValueError("Direct experts require identical receiver geometry and matching TP placement")
+                layout = layouts[0]
+            except (ValueError, TypeError) as error:
+                reason = str(error)
+        try:
+            enabled = configure(layout)
+        except Exception as error:
+            # Mapping/schema validation is local. Reject collectively before
+            # any rank starts a packed gather with a different source plan.
+            enabled = False
+            reason = str(error)
+            logger.warning("Direct expert source rejected: %s", reason)
+        agreed = torch.tensor([int(enabled)], dtype=torch.int32, device="cpu")
+        dist.all_reduce(agreed, op=dist.ReduceOp.MIN, group=group)
+        if agreed.item():
+            self._direct_expert_layout = layout
+        else:
+            configure(None)
+        if dist.get_rank(group=group) == 0:
+            logger.info(
+                "Weight sync selected %s: %s",
+                "packed BF16 direct" if self._direct_expert_layout else "standard HF",
+                self._direct_expert_layout or reason or "source expert mapping/geometry requires HF loading",
+            )
+
     @torch.no_grad()
     def update_weights(self) -> None:
         """version++, flush caches, process buckets with pipelining.
@@ -350,32 +394,40 @@ def update_weights(self) -> None:
         # two stages overlap.
         prev_refs: list[ObjectRef] = []
         prev_long_lived_tensors = None
-        with export_ctx:
-            for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights):
-                dispatch_errors = []
-                refs, long_lived_tensors = self._send_hf_params(
-                    hf_named_tensors, dispatch_errors=dispatch_errors if self._route_experts else None
-                )
-                # Wait for the *previous* chunk's IPC to finish before
-                # releasing its GPU tensors.
-                # Routed chunks can fail collectively while producing the next
-                # bucket. Complete their IPC before advancing the generator so
-                # an error cannot release a buffer still read by SGLang.
-                self._wait_for_ipc(
-                    refs if self._route_experts else prev_refs,
-                    error="; ".join(dispatch_errors) or None,
-                )
-                del prev_long_lived_tensors
-                prev_refs = [] if self._route_experts else refs
-                prev_long_lived_tensors = None if self._route_experts else long_lived_tensors
-                if self._route_experts:
-                    del long_lived_tensors, hf_named_tensors
-                # Backend-specific per-chunk synchronization is handled in device
-                # utils so this path stays hardware-agnostic.
-                device_utils.maybe_backend_barrier_on_weight_chunk(group=get_gloo_group())
-            # Drain the last chunk.
-            self._wait_for_ipc(prev_refs)
-            del prev_long_lived_tensors
+        current_refs: list[ObjectRef] = []
+        long_lived_tensors = None
+        try:
+            with export_ctx:
+                for hf_named_tensors in self._hf_weight_iterator.get_hf_weight_chunks(megatron_local_weights):
+                    dispatch_errors = []
+                    current_refs, long_lived_tensors = self._send_hf_params(
+                        hf_named_tensors, dispatch_errors=dispatch_errors
+                    )
+                    # Routed export agrees on failures before advancing its
+                    # collective generator. Ordinary export keeps two chunks
+                    # overlapped, including the previous chunk's storage.
+                    if self._route_experts:
+                        self._wait_for_ipc(current_refs, error="; ".join(dispatch_errors) or None)
+                        current_refs = []
+                    else:
+                        self._wait_for_ipc(prev_refs)
+                        prev_refs = []
+                        if dispatch_errors:
+                            self._wait_for_ipc(current_refs, error="; ".join(dispatch_errors))
+                    prev_long_lived_tensors = long_lived_tensors
+                    prev_refs, current_refs = current_refs, []
+                    if self._route_experts:
+                        prev_long_lived_tensors = long_lived_tensors = None
+                        del hf_named_tensors
+                    device_utils.maybe_backend_barrier_on_weight_chunk(group=get_gloo_group())
+                self._wait_for_ipc(prev_refs)
+                prev_refs = []
+        finally:
+            # Conversion or ACK failures must not release storage while an
+            # already submitted receiver is still reading it. No additional
+            # waits occur on the successful pipelined path.
+            _drain_ipc_refs(prev_refs + current_refs)
+            del prev_long_lived_tensors, long_lived_tensors
 
         # All ranks must finish sending before rank 0 triggers Marlin repack,
         # otherwise engines in slower gather groups may still be processing
@@ -414,24 +466,20 @@ def _routed_engine_call(self, engines: Sequence[ActorHandle], method: str, **kwa
         self._wait_for_ipc(refs, error=error)
 
     def _wait_for_ipc(self, refs: list[ObjectRef], *, error: str | None = None) -> None:
-        if not self._route_experts:
-            if refs:
-                ray.get(refs)
-            return
         try:
             results = ray.get(refs) if refs else []
             for result in results:
                 if isinstance(result, Mapping) and result.get("success") is False:
-                    raise RuntimeError(f"SGLang rejected routed weights: {result}")
+                    raise RuntimeError(f"SGLang rejected weights: {result}")
         except Exception as exc:
             error = str(exc)
             # ray.get(list) may raise before other requests finish. Keep source
             # buffers alive until every submitted reader has completed/failed.
-            for ref in refs:
-                try:
-                    ray.get(ref)
-                except Exception:
-                    pass
+            _drain_ipc_refs(refs)
+        if not self._route_experts:
+            if error is not None:
+                raise RuntimeError(f"Weight load failed: {error}")
+            return
         # All ranks finish the same IPC rounds. Propagate a receiver failure
         # before anyone proceeds to another all-to-all or publishes a version.
         failed = torch.tensor([int(error is not None)], dtype=torch.int32, device="cpu")
@@ -707,7 +755,7 @@ def _send_hf_params(
                 ipc_gather_group=self._ipc_gather_group,
                 weight_version=None if self._route_experts else self.weight_version,
                 dispatch_errors=dispatch_errors,
-                direct_experts=getattr(self, "_direct_qwen35_experts", False),
+                direct_layout=getattr(self, "_direct_expert_layout", None),
             )
             all_refs.extend(refs_colocated)
 
@@ -733,7 +781,7 @@ def _send_to_colocated_engine(
     ipc_gather_group,
     weight_version,
     dispatch_errors: list[str] | None = None,
-    direct_experts: bool = False,
+    direct_layout: DirectExpertLayout | None = None,
 ) -> tuple[list[ObjectRef], Any]:
     # Placeholder ranks (GPU slots reserved but no engine) have no gather group.
     # gather_object is only collective among group members, so we skip entirely.
@@ -755,10 +803,10 @@ def _send_to_colocated_engine(
     ]
 
     direct_chunk = None
-    if direct_experts:
-        if dist.get_world_size(ipc_gather_group) != 2:
-            raise ValueError("Direct Qwen3.5 IPC requires a two-rank receiver group")
-        direct_chunk = direct_expert_chunk(hf_named_tensors, dist.get_rank(group=ipc_gather_group))
+    if direct_layout is not None:
+        if dist.get_world_size(ipc_gather_group) != direct_layout.tp_size:
+            raise ValueError("Expert layout does not match the receiver TP group")
+        direct_chunk = direct_expert_chunk(hf_named_tensors, direct_layout, dist.get_rank(group=ipc_gather_group))
     if direct_chunk is not None:
         converted_named_tensors_by_dtypes = {"direct": direct_chunk}
     elif getattr(FlattenedTensorBucket, "supports_multi_dtypes", False):
@@ -786,7 +834,7 @@ def _send_to_colocated_engine(
             load_format = "flattened_bucket"
         long_live_tensors.append(payload)
         serialized = MultiprocessingSerializer.serialize(payload, output_str=True)
-        serialized_tensors.append((load_format, serialized) if direct_experts else serialized)
+        serialized_tensors.append((load_format, serialized))
 
     serialized_named_tensors = (
         [None] * dist.get_world_size(ipc_gather_group) if ipc_gather_src == dist.get_rank() else None
@@ -799,52 +847,20 @@ def _send_to_colocated_engine(
     )
 
     refs = []
-    if direct_experts:
-        if dist.get_rank() == ipc_gather_src:
-            refs = dispatch_typed_buckets(
-                serialized_named_tensors,
-                engine=ipc_engine,
-                version=weight_version,
-                serializer=MultiprocessingSerializer,
-                empty_flattened=lambda: _empty_flattened_tensor_data(cur_device),
-                held=long_live_tensors,
-                dispatch_errors=dispatch_errors,
-            )
-        return refs, long_live_tensors
-    if dist.get_rank() == ipc_gather_src:
-        # Ranks can contribute uneven bucket counts: under PP/EP/MoE layouts a TP
-        # rank may legitimately hold no HF tensors for a chunk. Pad short ranks
-        # with an empty bucket so every engine call gets one entry per rank. When
-        # all ranks are empty num_buckets is 0 and no update is issued.
-        num_buckets = max(len(tensors) for tensors in serialized_named_tensors)
-        empty_serialized_tensor = None
-        for i in range(num_buckets):
-            serialized_tensors_for_bucket = []
-            for tensors in serialized_named_tensors:
-                if i < len(tensors):
-                    serialized_tensors_for_bucket.append(tensors[i])
-                    continue
-                if empty_serialized_tensor is None:
-                    empty_tensor_data = _empty_flattened_tensor_data(cur_device)
-                    long_live_tensors.append(empty_tensor_data)
-                    empty_serialized_tensor = MultiprocessingSerializer.serialize(empty_tensor_data, output_str=True)
-                serialized_tensors_for_bucket.append(empty_serialized_tensor)
-            kwargs = {
-                "serialized_named_tensors": serialized_tensors_for_bucket,
-                "load_format": "flattened_bucket",
-            }
-            if weight_version is not None:
-                kwargs["weight_version"] = str(weight_version)
-            try:
-                refs.append(ipc_engine.update_weights_from_tensor.remote(**kwargs))
-            except Exception as exc:
-                if dispatch_errors is None:
-                    raise
-                # Return every submitted reader and its backing storage. The
-                # caller drains them before all ranks agree on this failure.
-                dispatch_errors.append(str(exc))
-                break
-
+    errors = dispatch_errors if dispatch_errors is not None else []
+    if dist.get_rank(group=ipc_gather_group) == 0:
+        refs = dispatch_typed_buckets(
+            serialized_named_tensors,
+            engine=ipc_engine,
+            version=weight_version,
+            serializer=MultiprocessingSerializer,
+            empty_flattened=lambda: _empty_flattened_tensor_data(cur_device),
+            held=long_live_tensors,
+            dispatch_errors=errors,
+        )
+        if errors and dispatch_errors is None:
+            _drain_ipc_refs(refs)
+            raise RuntimeError(f"Weight dispatch failed: {errors}")
     return refs, long_live_tensors
 
 
@@ -853,3 +869,11 @@ def _empty_flattened_tensor_data(device):
         "flattened_tensor": torch.empty(0, dtype=torch.uint8, device=device),
         "metadata": [],
     }
+
+
+def _drain_ipc_refs(refs: Sequence[ObjectRef]) -> None:
+    for ref in refs:
+        try:
+            ray.get(ref)
+        except Exception:
+            pass
diff --git a/relax/backends/sglang/sglang_engine.py b/relax/backends/sglang/sglang_engine.py
index 4bc4074c3..cd6fb16f2 100644
--- a/relax/backends/sglang/sglang_engine.py
+++ b/relax/backends/sglang/sglang_engine.py
@@ -771,8 +771,8 @@ def update_weight_version(self, weight_version: str) -> dict:
         )
 
     def get_expert_weight_sync_layout(self) -> dict:
-        """Return resolved layout settings for opt-in colocated expert
-        routing."""
+        """Return resolved layout settings for opt-in colocated expert routing
+        and direct BF16 loading."""
         from relax.utils.misc import get_hf_config
 
         if self.node_rank != 0 or self.worker_type != "regular":
@@ -792,11 +792,21 @@ def get_expert_weight_sync_layout(self) -> dict:
             "speculative_algorithm",
             "moe_runner_backend",
             "json_model_override_args",
+            "dtype",
+            "quantization",
+            "enable_lora",
+            "disable_shared_experts_fusion",
+            "enforce_shared_experts_fusion",
+            "kt_weight_path",
         )
         layout = {key: info[key] for key in keys if key in info}
         config = get_hf_config(info["model_path"])
         text_config = getattr(config, "text_config", None) or config
-        layout["num_experts"] = getattr(text_config, "num_experts", None)
+        layout["num_experts"] = getattr(text_config, "num_experts", None) or getattr(
+            text_config, "n_routed_experts", None
+        )
+        for key in ("hidden_size", "moe_intermediate_size", "quantization_config"):
+            layout[key] = getattr(text_config, key, None)
         return layout
 
     def load_lora_adapter_from_tensors(
diff --git a/relax/utils/training/inline_log_probs.py b/relax/utils/training/inline_log_probs.py
new file mode 100644
index 000000000..2f91bf8ac
--- /dev/null
+++ b/relax/utils/training/inline_log_probs.py
@@ -0,0 +1,87 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+"""Automatic eligibility for reusing a synchronous GRPO training forward."""
+
+from argparse import Namespace
+from collections.abc import Callable, Iterable, Sequence
+
+
+def can_reuse_colocate_log_probs(
+    args: Namespace,
+    num_microbatches: Sequence[int],
+    *,
+    backup_tags: Iterable[str],
+    rollout_data_postprocess: Callable | None,
+) -> bool:
+    """Check actual per-rollout state without changing the algorithm or
+    arguments."""
+    # These conditions make the current training forward identical to the skipped
+    # old-policy forward. Unsupported cases retain the framework's normal forward.
+    required = {
+        "colocate": True,
+        "true_on_policy_mode": True,
+        "fully_async": False,
+        "hybrid": False,
+        "advantage_estimator": "grpo",
+        "loss_type": "policy_loss",
+        "compute_advantages_and_returns": True,
+        "num_iters_per_train_update": 1,
+        "max_staleness": 0,
+        "update_weights_interval": 1,
+        "partial_rollout": False,
+        "rl_partial_rollouts": False,
+        "use_dynamic_global_batch_size": False,
+        "context_parallel_size": 1,
+        "pipeline_model_parallel_size": 1,
+        "dynamic_context_parallel": False,
+        "hidden_dropout": 0.0,
+        "attention_dropout": 0.0,
+        "bf16": True,
+        "fp16": False,
+        "fp8": None,
+        "fp4": None,
+        "fp8_param_gather": False,
+        "moe_input_jitter_eps": None,
+        "moe_router_enable_expert_bias": False,
+        "moe_router_force_load_balancing": False,
+        "moe_router_force_biased": None,
+        "moe_expert_capacity_factor": None,
+        "moe_aux_loss_coeff": 0.0,
+        "kl_coef": 0.0,
+        "use_kl_loss": False,
+        "use_opd": False,
+        "use_critic": False,
+        "keep_old_actor": False,
+        "use_rollout_logprobs": False,
+        "use_routing_replay": True,
+        "use_rollout_routing_replay": True,
+        "use_rollout_indexer_replay": False,
+        "enable_mtp_training": False,
+        "use_opsm": False,
+        "use_tis": True,
+        "calculate_per_token_loss": True,
+        "entropy_coef": 0.0,
+        "recompute_loss_function": False,
+        "custom_tis_function_path": None,
+        "custom_pg_loss_reducer_function_path": None,
+        "custom_loss_function_path": None,
+        "custom_model_provider_path": None,
+        "custom_megatron_before_log_prob_hook_path": None,
+        "custom_megatron_before_train_step_hook_path": None,
+    }
+    invalid = [name for name, expected in required.items() if getattr(args, name, None) != expected]
+    if invalid:
+        return False
+    if args.rollout_batch_size * args.n_samples_per_prompt != args.global_batch_size:
+        invalid.append("rollout size must equal global_batch_size")
+    if len(num_microbatches) != 1:
+        invalid.append("exactly one optimizer step per rollout is required")
+    if getattr(args, "num_steps_per_rollout", None) not in (None, 1):
+        invalid.append("num_steps_per_rollout")
+    if args.moe_router_load_balancing_type not in ("none", ["none"]):
+        invalid.append("moe_router_load_balancing_type")
+    if set(backup_tags) - {"actor"}:
+        invalid.append("reference, teacher, or old-actor weights")
+    if rollout_data_postprocess is not None:
+        invalid.append("rollout_data_postprocess")
+    return not invalid
diff --git a/scripts/training/multimodal/run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh b/scripts/training/multimodal/run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh
index dc5f9893d..b01c2c7df 100644
--- a/scripts/training/multimodal/run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh
+++ b/scripts/training/multimodal/run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh
@@ -140,13 +140,8 @@ runtime_env = json.loads(os.environ["RUNTIME_ENV_JSON"])
 # Vision embeddings must be recomputed after each weight update.
 runtime_env.setdefault("env_vars", {})["SGLANG_VLM_CACHE_SIZE_MB"] = "0"
 # The backend validates single-update, zero-KL GRPO before reusing log-probs.
-runtime_env["env_vars"]["RELAX_COLOCATE_INLINE_LOG_PROBS"] = "1"
 runtime_env["env_vars"]["NVTE_USE_CUTLASS_GROUPED_GEMM"] = "1"
 # Export BF16 experts in bulk, then send each receiver its TP-local shard.
-runtime_env["env_vars"]["RELAX_BF16_EXPERT_GATHER"] = "1"
-runtime_env["env_vars"]["RELAX_ASYNC_WEIGHT_GATHER"] = "1"
-runtime_env["env_vars"]["RELAX_QWEN35_PACKED_EXPERTS"] = "1"
-runtime_env["env_vars"]["RELAX_QWEN35_DIRECT_EXPERTS"] = "1"
 # Preserve the platform Trial identity and its collector path in Ray workers.
 # These are public variable names; their values are inherited at launch.
 for key in (
diff --git a/tests/backends/megatron/test_colocate_inline_log_probs.py b/tests/backends/megatron/test_colocate_inline_log_probs.py
index af8177c22..e46a7aefc 100644
--- a/tests/backends/megatron/test_colocate_inline_log_probs.py
+++ b/tests/backends/megatron/test_colocate_inline_log_probs.py
@@ -2,27 +2,11 @@
 
 """Safety gates and zero-KL inputs for single-update colocate GRPO."""
 
-import ast
-import os
-from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
 
-
-ROOT = Path(__file__).resolve().parents[3]
-
-
-def load_function(path, name, namespace, class_name=None):
-    tree = ast.parse(path.read_text())
-    body = tree.body
-    if class_name is not None:
-        body = next(node for node in body if isinstance(node, ast.ClassDef) and node.name == class_name).body
-    function = next(node for node in body if isinstance(node, ast.FunctionDef) and node.name == name)
-    future = ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0)
-    module = ast.fix_missing_locations(ast.Module(body=[future, function], type_ignores=[]))
-    exec(compile(module, str(path), "exec"), namespace)
-    return namespace[name]
+from relax.utils.training.inline_log_probs import can_reuse_colocate_log_probs
 
 
 @pytest.fixture
@@ -89,13 +73,12 @@ def actor():
 
 
 @pytest.fixture
-def inline_guard(monkeypatch):
-    monkeypatch.setenv("RELAX_COLOCATE_INLINE_LOG_PROBS", "1")
-    return load_function(
-        ROOT / "relax/backends/megatron/actor.py",
-        "_inline_colocate_old_log_probs",
-        {"os": os},
-        "MegatronTrainRayActor",
+def inline_guard():
+    return lambda actor, batches: can_reuse_colocate_log_probs(
+        actor.args,
+        batches,
+        backup_tags=actor.weights_backuper.backup_tags,
+        rollout_data_postprocess=actor.rollout_data_postprocess,
     )
 
 
@@ -143,27 +126,22 @@ def test_single_fresh_grpo_update_is_eligible(actor, inline_guard):
 )
 def test_unsafe_algorithm_or_state_is_rejected(actor, inline_guard, name, value):
     setattr(actor.args, name, value)
-    with pytest.raises(ValueError, match="Unsupported colocate true-on-policy"):
-        inline_guard(actor, [22])
+    assert not inline_guard(actor, [22])
 
 
 def test_actual_multiple_updates_are_rejected(actor, inline_guard):
-    with pytest.raises(ValueError, match="exactly one optimizer step"):
-        inline_guard(actor, [11, 11])
+    assert not inline_guard(actor, [11, 11])
 
 
 def test_extra_model_and_data_mutation_are_rejected(actor, inline_guard):
     actor.weights_backuper.backup_tags.append("ref")
-    with pytest.raises(ValueError, match="reference"):
-        inline_guard(actor, [22])
+    assert not inline_guard(actor, [22])
     actor.weights_backuper.backup_tags = ["actor"]
     actor.rollout_data_postprocess = lambda args: None
-    with pytest.raises(ValueError, match="rollout_data_postprocess"):
-        inline_guard(actor, [22])
+    assert not inline_guard(actor, [22])
 
 
-def test_existing_modes_do_not_enter_colocate_optimization(actor, inline_guard, monkeypatch):
-    monkeypatch.delenv("RELAX_COLOCATE_INLINE_LOG_PROBS")
+def test_existing_modes_do_not_enter_colocate_optimization(actor, inline_guard):
     assert not inline_guard(actor, [11, 11])
     actor.args.true_on_policy_mode = False
     assert not inline_guard(actor, [11, 11])
@@ -173,7 +151,7 @@ def test_existing_modes_do_not_enter_colocate_optimization(actor, inline_guard,
     assert not inline_guard(actor, [11, 11])
 
 
-def test_zero_kl_uses_sampling_shape_without_sampling_values():
+def test_zero_kl_uses_sampling_shape_without_sampling_values(monkeypatch):
     torch = pytest.importorskip("torch")
     calls = []
 
@@ -181,15 +159,13 @@ def compute_impl(args, **kwargs):
         calls.append(kwargs["kl"])
         return kwargs["kl"], kwargs["kl"]
 
-    namespace = {
-        "torch": torch,
-        "mpu": SimpleNamespace(is_pipeline_last_stage=lambda: True),
-        "maybe_padded_total_lengths": lambda *args: None,
-        "compute_advantages_and_returns_impl": compute_impl,
-    }
-    compute = load_function(ROOT / "relax/backends/megatron/loss.py", "compute_advantages_and_returns", namespace)
+    from relax.backends.megatron import loss
+
+    monkeypatch.setattr(loss.mpu, "is_pipeline_last_stage", lambda: True)
+    monkeypatch.setattr(loss, "maybe_padded_total_lengths", lambda *args: None)
+    monkeypatch.setattr(loss, "compute_advantages_and_returns_impl", compute_impl)
+    compute = loss.compute_advantages_and_returns
     args = SimpleNamespace(
-        _colocate_inline_log_probs=True,
         advantage_estimator="grpo",
         kl_coef=0,
         use_rollout_logprobs=False,
@@ -205,10 +181,19 @@ def compute_impl(args, **kwargs):
             "total_lengths": [5, 7],
             "loss_masks": [torch.ones(3), torch.ones(2)],
         }
-        compute(args, data)
+        compute(args, data, inline_log_probs=True)
         assert "log_probs" not in data
         assert [tuple(x.shape) for x in data["advantages"]] == [(3,), (2,)]
         assert all(x.dtype == torch.float32 and torch.count_nonzero(x) == 0 for x in data["advantages"])
     assert all(torch.equal(a, b) for a, b in zip(*calls, strict=True))
     with pytest.raises(ValueError, match="requires rollout log-probs"):
-        compute(args, {"response_lengths": [3], "total_lengths": [5]})
+        compute(args, {"response_lengths": [3], "total_lengths": [5]}, inline_log_probs=True)
+
+
+@pytest.mark.parametrize("batch", [256, 512])
+def test_auto_selection_uses_actual_batch_and_does_not_mutate_arguments(actor, inline_guard, batch):
+    actor.args.global_batch_size = batch
+    actor.args.rollout_batch_size = batch // 8
+    before = vars(actor.args).copy()
+    assert inline_guard(actor, [33])
+    assert vars(actor.args) == before
diff --git a/tests/backends/megatron/weight_update/test_async_gdn_weight_gather.py b/tests/backends/megatron/weight_update/test_async_gdn_weight_gather.py
index 1bba96d8f..d1c6cb39b 100644
--- a/tests/backends/megatron/weight_update/test_async_gdn_weight_gather.py
+++ b/tests/backends/megatron/weight_update/test_async_gdn_weight_gather.py
@@ -2,8 +2,6 @@
 
 """Reconstruct known global GDN tensors from real TP-local component shards."""
 
-import ast
-from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
@@ -12,12 +10,10 @@
 
 @pytest.mark.parametrize("tp", [1, 2, 4])
 @pytest.mark.parametrize("wrapped", [False, True])
-def test_async_gather_preserves_gdn_glu_and_replicated_layouts(tp, wrapped):
-    path = Path(__file__).resolve().parents[4] / "relax/backends/megatron/weight_update/common.py"
-    tree = ast.parse(path.read_text())
-    function = next(
-        node for node in tree.body if isinstance(node, ast.FunctionDef) and node.name == "all_gather_params_async"
-    )
+@pytest.mark.parametrize("async_mode", [False, True])
+def test_gather_preserves_gdn_glu_and_replicated_layouts(tp, wrapped, async_mode, monkeypatch):
+    from relax.backends.megatron.weight_update import common
+
     inputs, sources, expected, split_sizes = [], {}, [], {}
     specs = [
         ("self_attention.conv1d.weight", (2, 2, 4), (1, 4)),
@@ -51,31 +47,38 @@ def test_async_gather_preserves_gdn_glu_and_replicated_layouts(tp, wrapped):
     pending, waited = [], []
 
     def all_gather(outputs, param, *, group, async_op):
-        assert group == "tp" and async_op
+        assert group == "tp" and async_op == async_mode
         shards = sources[param.data_ptr()]
         pending.append(outputs)
 
         def wait():
-            assert len(pending) == 3, "Every gather must launch before waiting"
+            if async_mode:
+                assert len(pending) == 3, "Every gather must launch before waiting"
             for out, source in zip(outputs, shards, strict=True):
                 out.copy_(source)
             waited.append(True)
 
-        return SimpleNamespace(wait=wait)
+        if async_mode:
+            return SimpleNamespace(wait=wait)
+        wait()
+        return None
 
-    namespace = {
-        "torch": torch,
-        "dist": SimpleNamespace(all_gather=all_gather),
-        "mpu": SimpleNamespace(
-            get_tensor_model_parallel_world_size=lambda: tp, get_tensor_model_parallel_group=lambda: "tp"
-        ),
-    }
-    module = ast.Module(
-        body=[ast.ImportFrom(module="__future__", names=[ast.alias(name="annotations")], level=0), function],
-        type_ignores=[],
+    monkeypatch.setattr(common.dist, "all_gather", all_gather)
+    monkeypatch.setattr(
+        common,
+        "mpu",
+        SimpleNamespace(get_tensor_model_parallel_world_size=lambda: tp, get_tensor_model_parallel_group=lambda: "tp"),
     )
-    exec(compile(ast.fix_missing_locations(module), str(path), "exec"), namespace)
-    result = namespace["all_gather_params_async"](inputs, gdn_split_sizes=split_sizes)
+    if async_mode:
+        result = common.all_gather_params_async(inputs, gdn_split_sizes=split_sizes)
+    else:
+        config = SimpleNamespace(
+            linear_key_head_dim=2, linear_value_head_dim=4, linear_num_key_heads=tp, linear_num_value_heads=tp
+        )
+        monkeypatch.setattr(common, "get_hf_config", lambda path: SimpleNamespace(text_config=config))
+        result = [
+            common.all_gather_param(SimpleNamespace(hf_checkpoint="test"), info.name, param) for info, param in inputs
+        ]
     assert len(waited) == 3
     assert len(result) == len(expected)
     assert all(
@@ -83,3 +86,16 @@ def wait():
         for actual, target in zip(result, expected, strict=True)
     )
     assert result[-1].data_ptr() == replicated.data_ptr()
+
+
+@pytest.mark.parametrize("kind", ["replicated", "duplicated", "buffer"])
+def test_non_sharded_gdn_named_tensor_needs_no_gdn_config(monkeypatch, kind):
+    from relax.backends.megatron.weight_update import common
+
+    tensor = torch.ones(4)
+    if kind != "buffer":
+        tensor.tensor_model_parallel = kind == "duplicated"
+        tensor.parallel_mode = "duplicated" if kind == "duplicated" else None
+    monkeypatch.setattr(common, "get_hf_config", lambda *a: pytest.fail("unneeded config lookup"))
+    actual = common.all_gather_param(SimpleNamespace(), "decoder.layers.0.self_attention.in_proj.weight", tensor)
+    assert actual.data_ptr() == tensor.data_ptr()
diff --git a/tests/backends/megatron/weight_update/test_auto_weight_sync.py b/tests/backends/megatron/weight_update/test_auto_weight_sync.py
new file mode 100644
index 000000000..538c1c207
--- /dev/null
+++ b/tests/backends/megatron/weight_update/test_auto_weight_sync.py
@@ -0,0 +1,232 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+from types import SimpleNamespace
+from unittest.mock import Mock
+
+import pytest
+import torch
+
+from relax.backends.megatron.weight_update import update_weight_from_tensor as sync
+from relax.backends.megatron.weight_update.expert_layout import DirectExpertLayout
+from relax.backends.megatron.weight_update.hf_weight_iterator_bridge import HfWeightIteratorBridge
+
+
+def receiver():
+    return dict(
+        tp_size=2,
+        ep_size=1,
+        pp_size=1,
+        moe_dp_size=1,
+        moe_runner_backend="triton",
+        dtype="bfloat16",
+        disable_shared_experts_fusion=True,
+        num_experts=4,
+        hidden_size=6,
+        moe_intermediate_size=8,
+    )
+
+
+@pytest.mark.parametrize("failure", ["none", "local_not_ready", "peer_not_ready", "configure", "query", "peer_reject"])
+def test_negotiation_agrees_before_selecting_or_falling_back(monkeypatch, failure):
+    updater = sync.UpdateWeightFromTensor.__new__(sync.UpdateWeightFromTensor)
+    updater.args = SimpleNamespace(colocate=True)
+    updater.use_distribute = updater._route_experts = updater.lora_enabled = False
+    updater.quantization_config = None
+    configured, reductions = [], []
+
+    def configure(layout):
+        configured.append(layout)
+        if layout is not None and failure == "configure":
+            raise ValueError("unsupported source mapping")
+        return layout is not None
+
+    updater._hf_weight_iterator = SimpleNamespace(
+        configure_direct_experts=configure, _bf16_expert_buckets=[] if failure == "local_not_ready" else [object()]
+    )
+    engine = Mock()
+    updater.rollout_engines = [engine]
+    get = Mock(side_effect=RuntimeError("query failed")) if failure == "query" else Mock(return_value=[receiver()])
+    monkeypatch.setattr(sync.ray, "get", get)
+    monkeypatch.setattr(sync, "get_gloo_group", lambda: "gloo")
+    monkeypatch.setattr(sync.dist, "get_rank", lambda **kwargs: 0)
+    monkeypatch.setattr(sync.dist, "broadcast_object_list", lambda *args, **kwargs: None)
+
+    def reduce(tensor, *, op, group):
+        assert op == sync.dist.ReduceOp.MIN and group == "gloo" and tensor.device.type == "cpu"
+        reductions.append(tensor.item())
+        if failure == "peer_not_ready" and len(reductions) == 1 or failure == "peer_reject" and len(reductions) == 2:
+            tensor.zero_()
+
+    monkeypatch.setattr(sync.dist, "all_reduce", reduce)
+    updater._configure_direct_experts([2])
+    assert (updater._direct_expert_layout is not None) == (failure == "none")
+    if failure in ("local_not_ready", "peer_not_ready"):
+        assert len(reductions) == 1 and configured == [None]
+        get.assert_not_called()
+    else:
+        assert len(reductions) == 2
+        if failure != "none":
+            assert configured[-1] is None
+
+
+@pytest.mark.parametrize("rank", [0, 3])
+def test_reconnect_creates_groups_once_including_placeholder_ranks(monkeypatch, rank):
+    updater = sync.UpdateWeightFromTensor.__new__(sync.UpdateWeightFromTensor)
+    updater.args = SimpleNamespace(rollout_num_gpus_per_engine=2)
+    updater._route_experts = False
+    updater._ipc_placement = updater._ipc_gather_group = updater._ipc_gather_src = updater._ipc_engine = None
+    updater._configure_direct_experts = Mock()
+    first, replacement = Mock(), Mock()
+    monkeypatch.setattr(sync.ray, "get", lambda refs: [{"node_id": "local"} for _ in refs])
+    monkeypatch.setattr(sync.ray, "get_runtime_context", lambda: SimpleNamespace(get_node_id=lambda: "local"))
+    monkeypatch.setattr(sync, "get_gloo_group", lambda: "gloo")
+    monkeypatch.setattr(sync.dist, "get_world_size", lambda: 4)
+    monkeypatch.setattr(sync.dist, "get_rank", lambda: rank)
+    monkeypatch.setattr(
+        sync.dist, "all_gather_object", lambda out, local, **kw: out.__setitem__(slice(None), [local] * 4)
+    )
+    create = Mock(return_value="ipc")
+    monkeypatch.setattr(sync.dist, "new_group", create)
+    updater.connect_rollout_engines([first], None, [2], [0])
+    updater.connect_rollout_engines([replacement], None, [2], [0])
+    create.assert_called_once_with(ranks=[0, 1], backend="gloo")
+    assert updater._ipc_engine is (replacement if rank == 0 else None)
+    with pytest.raises(ValueError, match="changed IPC GPU placement"):
+        updater.connect_rollout_engines([replacement], None, [2], [1])
+
+
+@pytest.mark.parametrize("geometry,enabled", [((256, 2048, 512), True), ((256, 7168, 2048), False)])
+def test_packed_projection_memory_budget_uses_metadata_only(monkeypatch, geometry, enabled):
+    from relax.backends.megatron.weight_update import hf_weight_iterator_bridge as bridge
+
+    e, h, i = geometry
+    iterator = HfWeightIteratorBridge.__new__(HfWeightIteratorBridge)
+    iterator.args = SimpleNamespace(num_experts=e, update_weight_buffer_size=512 * 1024**2)
+    iterator._packed_expert_sync = False
+    iterator._direct_expert_layout = None
+    infos = [
+        SimpleNamespace(
+            name=f"decoder.layers.0.mlp.experts.linear_fc{fc}.weight{expert}",
+            dtype=torch.bfloat16,
+            shape=shape,
+            src_rank=expert // (e // 8),
+            size=shape[0] * shape[1] * 2,
+        )
+        for fc, shape in [(1, (2 * i, h)), (2, (h, i))]
+        for expert in range(e)
+    ]
+    iterator._bf16_expert_buckets = [infos]
+    iterator._bridge_converter = SimpleNamespace(supports_packed_experts=lambda names: True)
+    monkeypatch.setattr(bridge.mpu, "get_expert_model_parallel_group", lambda: "ep")
+    monkeypatch.setattr(bridge.dist, "get_process_group_ranks", lambda group: list(range(8)))
+    assert iterator.configure_direct_experts(DirectExpertLayout(e, h, i, 2)) is enabled
+    assert iterator._packed_expert_sync is enabled
+    if enabled:
+        gathers = iterator._bf16_expert_gathers
+        assert iterator.configure_direct_experts(DirectExpertLayout(e, h, i, 2))
+        assert iterator._bf16_expert_gathers is gathers
+        assert not iterator.configure_direct_experts(None)
+        assert not iterator._packed_expert_sync
+
+
+def test_ordinary_wait_drains_all_readers_without_adding_collectives(monkeypatch):
+    updater = sync.UpdateWeightFromTensor.__new__(sync.UpdateWeightFromTensor)
+    updater._route_experts = False
+    calls = []
+
+    def get(ref):
+        calls.append(ref)
+        if isinstance(ref, list) or ref == "failed":
+            raise RuntimeError("reader failed")
+        return {"success": True}
+
+    monkeypatch.setattr(sync.ray, "get", get)
+    reduce = Mock()
+    monkeypatch.setattr(sync.dist, "all_reduce", reduce)
+    with pytest.raises(RuntimeError, match="reader failed"):
+        updater._wait_for_ipc(["failed", "pending"])
+    assert calls == [["failed", "pending"], "failed", "pending"]
+    reduce.assert_not_called()
+
+
+@pytest.mark.parametrize("failure", ["conversion", "previous_ack", "dispatch"])
+def test_pipeline_retains_both_chunks_until_submitted_readers_finish(monkeypatch, failure):
+    import weakref
+
+    class Storage:
+        pass
+
+    updater = sync.UpdateWeightFromTensor.__new__(sync.UpdateWeightFromTensor)
+    updater._route_experts = updater.lora_enabled = updater.lora_merge_mode = False
+    updater.quantization_config = None
+    updater.weight_version = 0
+    updater.rollout_engines = updater.distributed_rollout_engines = []
+    updater.weights_getter = lambda: {}
+    alive, completed = {}, []
+
+    def chunks(weights):
+        yield []
+        if failure == "conversion":
+            assert alive["chunk0"]() is not None
+            raise RuntimeError("next conversion failed")
+        yield []
+
+    def send(tensors, *, dispatch_errors):
+        key = "chunk" + str(len(alive))
+        storage = Storage()
+        alive[key] = weakref.ref(storage)
+        if failure == "dispatch" and key == "chunk1":
+            dispatch_errors.append("partial dispatch failed")
+        return [key], storage
+
+    def get(refs):
+        names = refs if isinstance(refs, list) else [refs]
+        for name in names:
+            assert alive[name]() is not None, "IPC source storage was released while a reader was active"
+        if failure == "previous_ack" and refs == ["chunk0"]:
+            assert alive["chunk1"]() is not None
+            raise RuntimeError("previous reader failed")
+        completed.extend(names)
+        return [{} for _ in names] if isinstance(refs, list) else {}
+
+    updater._hf_weight_iterator = SimpleNamespace(get_hf_weight_chunks=chunks)
+    updater._send_hf_params = send
+    monkeypatch.setattr(sync.ray, "get", get)
+    monkeypatch.setattr(sync.dist, "get_rank", lambda: 0)
+    monkeypatch.setattr(sync.dist, "barrier", lambda **kwargs: None)
+    monkeypatch.setattr(sync, "get_gloo_group", lambda: "gloo")
+    monkeypatch.setattr(sync.device_utils, "maybe_backend_barrier_on_weight_chunk", lambda **kw: None)
+    with pytest.raises(RuntimeError):
+        updater.update_weights()
+    assert set(completed) == set(alive)
+
+
+@pytest.mark.parametrize("kind", ["replicated", "duplicated", "buffer", "sharded"])
+def test_auto_async_plan_reads_gdn_config_only_for_sharded_parameters(monkeypatch, kind):
+    from relax.backends.megatron.weight_update import hf_weight_iterator_bridge as bridge
+
+    attrs = dict(
+        tensor_model_parallel=kind in ("duplicated", "sharded"),
+        parallel_mode="duplicated" if kind == "duplicated" else None,
+    )
+    info = SimpleNamespace(name="decoder.layers.0.self_attention.in_proj.weight", attrs=attrs)
+    converter = SimpleNamespace(init_tasks=Mock(), can_convert=Mock())
+    monkeypatch.setattr(bridge, "BridgeConverter", lambda **kwargs: converter)
+    monkeypatch.setattr(bridge, "_build_param_info_buckets", lambda *args, **kwargs: ([], [[info]], {}, {}))
+    monkeypatch.setattr(bridge.torch.cuda, "is_available", lambda: True)
+    monkeypatch.setattr(bridge.mpu, "get_pipeline_model_parallel_world_size", lambda: 2)
+    monkeypatch.setattr(bridge.mpu, "get_tensor_model_parallel_world_size", lambda: 4)
+    config = Mock(
+        return_value=SimpleNamespace(
+            linear_key_head_dim=2, linear_value_head_dim=4, linear_num_key_heads=4, linear_num_value_heads=4
+        )
+    )
+    monkeypatch.setattr(bridge, "get_hf_config", config)
+    iterator = HfWeightIteratorBridge(SimpleNamespace(hf_checkpoint="test"), [], "any-model", None)
+    assert iterator._async_weight_gather
+    if kind == "sharded":
+        config.assert_called_once_with("test")
+        assert iterator._gdn_split_sizes == {info.name: (2, 2, 4, 4, 1, 1)}
+    else:
+        config.assert_not_called()
+        assert iterator._gdn_split_sizes == {}
diff --git a/tests/backends/megatron/weight_update/test_bf16_expert_gather.py b/tests/backends/megatron/weight_update/test_bf16_expert_gather.py
index b2f55591b..f33229de1 100644
--- a/tests/backends/megatron/weight_update/test_bf16_expert_gather.py
+++ b/tests/backends/megatron/weight_update/test_bf16_expert_gather.py
@@ -2,8 +2,6 @@
 
 """Packing, tensor lifetime and metadata reuse across replacement groups."""
 
-import importlib.util
-from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
@@ -12,11 +10,9 @@
 
 @pytest.fixture
 def module():
-    path = Path(__file__).resolve().parents[4] / "relax/backends/megatron/weight_update/bf16_expert_gather.py"
-    spec = importlib.util.spec_from_file_location("bf16_expert_gather_test_target", path)
-    result = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(result)
-    return result
+    from relax.backends.megatron.weight_update import expert_gather
+
+    return expert_gather
 
 
 def simulate_collectives(monkeypatch, module, peers):
diff --git a/tests/backends/megatron/weight_update/test_expert_capabilities.py b/tests/backends/megatron/weight_update/test_expert_capabilities.py
new file mode 100644
index 000000000..13d95bd98
--- /dev/null
+++ b/tests/backends/megatron/weight_update/test_expert_capabilities.py
@@ -0,0 +1,26 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+from types import SimpleNamespace
+
+import pytest
+
+from relax.backends.megatron.weight_update.bridge_converter import BridgeConverter
+
+
+@pytest.mark.parametrize("prefix", ["model.layers.2", "model.language_model.layers.2"])
+@pytest.mark.parametrize("packed", [False, True])
+def test_bridge_projection_capability_uses_existing_mapping(prefix, packed):
+    source = "decoder.layers.2.mlp.experts.linear_fc1.weight3"
+    target = f"{prefix}.mlp.experts"
+    hf_param = (
+        target + ".gate_up_proj"
+        if packed
+        else {"gate": target + ".3.gate_proj.weight", "up": target + ".3.up_proj.weight"}
+    )
+    converter = BridgeConverter.__new__(BridgeConverter)
+    converter._bridge_task_map = {source: SimpleNamespace(mapping=SimpleNamespace(hf_param=hf_param))}
+    assert converter.supports_packed_experts([source])
+    converter._bridge_task_map[source].mapping.hf_param = target + ".3.w1.weight"
+    assert not converter.supports_packed_experts([source])
+    converter._bridge_task_map[source].mapping.hf_param = {"gate": target + ".3.gate_proj.weight"}
+    assert not converter.supports_packed_experts([source])
diff --git a/tests/backends/megatron/weight_update/test_expert_layout.py b/tests/backends/megatron/weight_update/test_expert_layout.py
new file mode 100644
index 000000000..9eeb2ac8c
--- /dev/null
+++ b/tests/backends/megatron/weight_update/test_expert_layout.py
@@ -0,0 +1,207 @@
+# Copyright (c) 2026 Relax Authors. All Rights Reserved.
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from relax.backends.megatron.weight_update.expert_layout import DirectExpertLayout, direct_expert_chunk
+from relax.backends.megatron.weight_update.ipc import dispatch_typed_buckets
+
+
+def receiver(**overrides):
+    return (
+        dict(
+            tp_size=2,
+            ep_size=1,
+            pp_size=1,
+            moe_dp_size=1,
+            moe_runner_backend="triton",
+            dtype="bfloat16",
+            disable_shared_experts_fusion=True,
+            num_experts=4,
+            hidden_size=6,
+            moe_intermediate_size=8,
+        )
+        | overrides
+    )
+
+
+@pytest.mark.parametrize("tp", [1, 2, 4])
+@pytest.mark.parametrize("prefix", ["model.layers", "model.language_model.layers"])
+@pytest.mark.parametrize("geometry", [(4, 6, 8), (8, 10, 12)])
+def test_direct_shards_preserve_expert_and_gate_up_order(tp, prefix, geometry):
+    e, h, i = geometry
+    layout = DirectExpertLayout(e, h, i, tp)
+    gate = torch.arange(e * i * h).reshape(e, i, h).to(torch.bfloat16)
+    up = gate + 512
+    down = torch.arange(e * h * i).reshape(e, h, i).to(torch.bfloat16)
+    packed = torch.cat((gate, up), dim=1)
+    for rank in range(tp):
+        [(name, actual)] = direct_expert_chunk([(f"{prefix}.9.mlp.experts.gate_up_proj", packed)], layout, rank)
+        expected = torch.cat((gate.chunk(tp, dim=1)[rank], up.chunk(tp, dim=1)[rank]), dim=1)
+        assert name == "model.layers.9.mlp.experts.w13_weight"
+        assert torch.equal(actual, expected) and actual.is_contiguous()
+        snapshot = actual.clone()
+        layout.shard(packed + 4, "gate_up_proj", rank)
+        assert torch.equal(actual, snapshot)
+        [(name, actual)] = direct_expert_chunk([(f"{prefix}.9.mlp.experts.down_proj", down)], layout, rank)
+        assert name == "model.layers.9.mlp.experts.w2_weight"
+        assert torch.equal(actual, down.chunk(tp, dim=2)[rank]) and actual.is_contiguous()
+    assert direct_expert_chunk([("model.norm.weight", gate)], layout, 0) is None
+    assert direct_expert_chunk([], layout, 0) is None
+
+
+@pytest.mark.parametrize("tp", [1, 2, 4])
+def test_direct_matches_installed_sglang_fused_moe_loader(tp):
+    """Use the receiver's real per-expert loaders as the independent oracle."""
+    layer_module = pytest.importorskip("sglang.srt.layers.moe.fused_moe_triton.layer")
+    e, h, i = 4, 6, 8
+    layout = DirectExpertLayout(e, h, i, tp)
+    gate = torch.arange(e * i * h).reshape(e, i, h).to(torch.bfloat16)
+    up, down = gate + 512, gate.transpose(1, 2).contiguous() + 1024
+    receiver_layer = SimpleNamespace(
+        moe_runner_config=SimpleNamespace(is_gated=True),
+        quant_method=SimpleNamespace(),
+        quant_config=None,
+        use_padded_loading=False,
+        use_presharded_weights=False,
+        use_triton_kernels=False,
+        moe_tp_size=tp,
+    )
+    for rank in range(tp):
+        w13, w2 = (
+            torch.empty(e, 2 * i // tp, h, dtype=torch.bfloat16),
+            torch.empty(e, h, i // tp, dtype=torch.bfloat16),
+        )
+        for expert in range(e):
+            layer_module.FusedMoE._load_w13(receiver_layer, w13[expert], 0, "w1", gate[expert], rank)
+            layer_module.FusedMoE._load_w13(receiver_layer, w13[expert], 0, "w3", up[expert], rank)
+            layer_module.FusedMoE._load_w2(receiver_layer, w2[expert], 1, "w2", down[expert], rank)
+        assert torch.equal(w13, layout.shard(torch.cat((gate, up), dim=1), "gate_up_proj", rank))
+        assert torch.equal(w2, layout.shard(down, "down_proj", rank))
+
+
+@pytest.mark.parametrize("tp", [1, 2, 4])
+def test_resolved_receiver_geometry_has_no_model_family_requirement(tp):
+    assert DirectExpertLayout.from_receiver(receiver(tp_size=tp)) == DirectExpertLayout(4, 6, 8, tp)
+
+
+@pytest.mark.parametrize(
+    "override",
+    [
+        dict(tp_size=3),
+        dict(tp_size=0),
+        dict(ep_size=2),
+        dict(pp_size=2),
+        dict(moe_dp_size=2),
+        dict(moe_runner_backend="auto"),
+        dict(moe_runner_backend="triton_kernel"),
+        dict(moe_runner_backend="flashinfer_cutlass"),
+        dict(dtype="float16"),
+        dict(dtype="auto"),
+        dict(disable_shared_experts_fusion=False),
+        dict(enforce_shared_experts_fusion=True),
+        dict(enable_eplb=True),
+        dict(ep_num_redundant_experts=1),
+        dict(init_expert_location="random"),
+        dict(ep_join_mode="scale"),
+        dict(quantization="fp8"),
+        dict(quantization_config={"quant_method": "fp8"}),
+        dict(enable_lora=True),
+        dict(kt_weight_path="/weights"),
+        dict(speculative_algorithm="EAGLE"),
+        dict(hidden_size=None),
+        dict(json_model_override_args='{"hidden_size":4096}'),
+    ],
+)
+def test_unsupported_receiver_layout_is_rejected(override):
+    with pytest.raises(ValueError):
+        DirectExpertLayout.from_receiver(receiver(**override))
+
+
+@pytest.mark.parametrize(
+    "packets,formats",
+    [
+        ([[("direct", "d0")], [("direct", "d1")]], ["direct"]),
+        ([[("flattened_bucket", "f0")], [("flattened_bucket", "f1")]], ["flattened_bucket"]),
+        ([[], [("direct", "d1")]], ["direct"]),
+        ([[("direct", "d0")], []], ["direct"]),
+        ([[("direct", "d0")], [("flattened_bucket", "f1")]], ["direct", "flattened_bucket"]),
+        (
+            [[("flattened_bucket", "f0"), ("flattened_bucket", "g0")], [], [("direct", "d2")], []],
+            ["flattened_bucket", "direct", "flattened_bucket"],
+        ),
+        ([[], []], []),
+    ],
+)
+def test_format_dispatch_handles_empty_and_uneven_ranks(packets, formats):
+    calls, held = [], []
+    serializer = SimpleNamespace(
+        serialize=lambda value, output_str: "empty-direct" if isinstance(value, list) else "empty-flat"
+    )
+
+    def remote(**kwargs):
+        calls.append(kwargs)
+        return len(calls)
+
+    refs = dispatch_typed_buckets(
+        packets,
+        engine=SimpleNamespace(update_weights_from_tensor=SimpleNamespace(remote=remote)),
+        version=12,
+        serializer=serializer,
+        empty_flattened=lambda: dict(flattened_tensor=torch.empty(0, dtype=torch.uint8), metadata=[]),
+        held=held,
+    )
+    assert len(refs) == len(formats) and [call["load_format"] for call in calls] == formats
+    assert all(
+        len(call["serialized_named_tensors"]) == len(packets) and call["weight_version"] == "12" for call in calls
+    )
+
+
+def test_dispatch_validates_all_formats_before_sending_and_retains_partial_refs():
+    calls, held, errors = [], [], []
+
+    def remote(**kwargs):
+        calls.append(kwargs)
+        if len(calls) == 2:
+            raise RuntimeError("second dispatch failed")
+        return "first-reader"
+
+    kwargs = dict(
+        engine=SimpleNamespace(update_weights_from_tensor=SimpleNamespace(remote=remote)),
+        version=None,
+        serializer=SimpleNamespace(serialize=lambda value, output_str: "empty"),
+        empty_flattened=dict,
+        held=held,
+        dispatch_errors=errors,
+    )
+    with pytest.raises(ValueError):
+        dispatch_typed_buckets([[("direct", "d"), ("invalid", "bad")]], **kwargs)
+    assert calls == []
+    refs = dispatch_typed_buckets([[("direct", "d"), ("flattened_bucket", "f")]], **kwargs)
+    assert refs == ["first-reader"] and errors == ["second dispatch failed"]
+    assert all("weight_version" not in call for call in calls)
+
+
+def test_partial_dispatch_retains_readers_if_empty_serialization_fails():
+    calls, errors, held = [], [], []
+
+    def serialize(value, output_str):
+        if isinstance(value, dict):
+            raise RuntimeError("empty serialization failed")
+        return "empty-direct"
+
+    refs = dispatch_typed_buckets(
+        [[("direct", "d")], [("flattened_bucket", "f")]],
+        engine=SimpleNamespace(
+            update_weights_from_tensor=SimpleNamespace(remote=lambda **kw: calls.append(kw) or "reader")
+        ),
+        version=1,
+        serializer=SimpleNamespace(serialize=serialize),
+        empty_flattened=dict,
+        held=held,
+        dispatch_errors=errors,
+    )
+    assert refs == ["reader"] and errors == ["empty serialization failed"]
+    assert len(calls) == 1 and len(held) == 2
diff --git a/tests/backends/megatron/weight_update/test_qwen35_packed_experts.py b/tests/backends/megatron/weight_update/test_packed_experts.py
similarity index 82%
rename from tests/backends/megatron/weight_update/test_qwen35_packed_experts.py
rename to tests/backends/megatron/weight_update/test_packed_experts.py
index f9d5c1094..d36728f16 100644
--- a/tests/backends/megatron/weight_update/test_qwen35_packed_experts.py
+++ b/tests/backends/megatron/weight_update/test_packed_experts.py
@@ -3,8 +3,6 @@
 """Expert IDs, gate/up order, changed weights and replacement process
 groups."""
 
-import importlib.util
-from pathlib import Path
 from types import SimpleNamespace
 
 import pytest
@@ -13,11 +11,9 @@
 
 @pytest.fixture
 def module():
-    path = Path(__file__).resolve().parents[4] / "relax/backends/megatron/weight_update/qwen35_packed_experts.py"
-    spec = importlib.util.spec_from_file_location("qwen35_packed_experts_test_target", path)
-    result = importlib.util.module_from_spec(spec)
-    spec.loader.exec_module(result)
-    return result
+    from relax.backends.megatron.weight_update import expert_gather
+
+    return expert_gather
 
 
 def sources():
@@ -36,7 +32,7 @@ def sources():
 
 def test_numeric_expert_order_and_group_rank_order(module):
     infos = sorted(sources(), key=lambda info: info.name)
-    buckets = module.bucket_qwen35_experts(infos, 24, [7, 2])
+    buckets = module.bucket_expert_projections(infos, 24, [7, 2])
     assert len(buckets) == 6
     assert [bucket[0].name.split("layers.")[1].split(".")[0] for bucket in buckets] == ["0", "0", "2", "2", "10", "10"]
     for bucket in buckets:
@@ -62,7 +58,7 @@ def test_invalid_source_layout_fails_before_payload_collectives(module, mutation
     else:
         infos = [info for info in infos if ".linear_fc2." not in info.name]
     with pytest.raises(ValueError):
-        module.bucket_qwen35_experts(infos, 24, [7, 2])
+        module.bucket_expert_projections(infos, 24, [7, 2])
 
 
 def peers(projection):
@@ -114,7 +110,7 @@ def gather(output, local, *, group, async_op):
 
     monkeypatch.setattr(module.dist, "all_gather_object", metadata)
     monkeypatch.setattr(module.dist, "all_gather_into_tensor", gather)
-    plan = module.Qwen35PackedExpertGather(24)
+    plan = module.PackedExpertGather(24)
     first_group = SimpleNamespace(ranks=[7, 2], rank=rank)
     [(name, value)] = plan(inputs[rank], group=first_group, device="cpu")
     assert name == f"model.language_model.layers.2.mlp.experts.{projection}_proj"
@@ -150,6 +146,28 @@ def test_invalid_converted_layout(module, mutation):
         inputs[0][0] = (name.replace("layers.2", "layers.3"), tensor)
     else:
         inputs[1] = peers("down")[1]
-    plan = module.Qwen35PackedExpertGather(24)
+    plan = module.PackedExpertGather(24)
     with pytest.raises(ValueError):
         plan._prepare([module._layout(peer) for peer in inputs], (7, 2))
+
+
+@pytest.mark.parametrize(
+    "source_prefix,hf_prefix",
+    [
+        ("module.module.decoder.layers", "model.layers"),
+        ("module.module.language_model.decoder.layers", "model.language_model.layers"),
+    ],
+)
+def test_deepseek_and_multimodal_names_use_same_projection_packing(module, source_prefix, hf_prefix):
+    infos = sources()
+    for info in infos:
+        info.name = info.name.replace("module.module.language_model.decoder.layers", source_prefix)
+    assert len(module.bucket_expert_projections(infos, 24, [7, 2])) == 6
+    inputs = peers("gate_up")
+    inputs = [
+        [(name.replace("model.language_model.layers", hf_prefix), tensor) for name, tensor in peer] for peer in inputs
+    ]
+    plan = module.PackedExpertGather(24)
+    plan._prepare([module._layout(peer) for peer in inputs], (7, 2))
+    assert plan._name == f"{hf_prefix}.2.mlp.experts.gate_up_proj"
+    assert plan._shape == (24, 8, 6)
diff --git a/tests/backends/megatron/weight_update/test_qwen35_direct_experts.py b/tests/backends/megatron/weight_update/test_qwen35_direct_experts.py
deleted file mode 100644
index 799590e05..000000000
--- a/tests/backends/megatron/weight_update/test_qwen35_direct_experts.py
+++ /dev/null
@@ -1,143 +0,0 @@
-# Copyright (c) 2026 Relax Authors. All Rights Reserved.
-
-import importlib.util
-from pathlib import Path
-from types import SimpleNamespace
-
-import pytest
-import torch
-
-
-_PATH = Path(__file__).parents[4] / "relax/backends/megatron/weight_update/qwen35_direct_experts.py"
-_SPEC = importlib.util.spec_from_file_location("qwen35_direct_test", _PATH)
-direct = importlib.util.module_from_spec(_SPEC)
-_SPEC.loader.exec_module(direct)
-
-
-@pytest.mark.parametrize("rank", [0, 1])
-def test_shards_match_separate_gate_up_and_down_slices_and_keep_old_storage(rank):
-    gate_up = torch.arange(3 * 8 * 6, dtype=torch.bfloat16).reshape(3, 8, 6)
-    down = torch.arange(3 * 6 * 4, dtype=torch.bfloat16).reshape(3, 6, 4)
-    first = direct._shard_expert(gate_up, "gate_up_proj", rank)
-    second = direct._shard_expert(down, "down_proj", rank)
-    rows = [0, 1, 4, 5] if rank == 0 else [2, 3, 6, 7]
-    columns = [0, 1] if rank == 0 else [2, 3]
-    torch.testing.assert_close(first, gate_up[:, rows], rtol=0, atol=0)
-    torch.testing.assert_close(second, down[:, :, columns], rtol=0, atol=0)
-    assert first.is_contiguous() and second.is_contiguous()
-    snapshots = first.clone(), second.clone()
-    direct._shard_expert(gate_up + 5, "gate_up_proj", rank)
-    direct._shard_expert(down + 5, "down_proj", rank)
-    torch.testing.assert_close(first, snapshots[0], rtol=0, atol=0)
-    torch.testing.assert_close(second, snapshots[1], rtol=0, atol=0)
-
-
-@pytest.mark.parametrize(
-    "projection,shape,target",
-    [
-        ("gate_up_proj", (256, 1024, 2048), (256, 512, 2048)),
-        ("down_proj", (256, 2048, 512), (256, 2048, 256)),
-    ],
-)
-def test_complete_model_shapes_and_runtime_names_without_allocating_storage(projection, shape, target):
-    tensor = torch.empty(shape, dtype=torch.bfloat16, device="meta")
-    for rank in (0, 1):
-        result = direct.direct_expert_chunk(
-            [(f"model.language_model.layers.9.mlp.experts.{projection}", tensor)], rank
-        )
-        suffix = "w13_weight" if projection == "gate_up_proj" else "w2_weight"
-        assert result[0][0] == f"model.layers.9.mlp.experts.{suffix}"
-        assert tuple(result[0][1].shape) == target
-    assert direct.direct_expert_chunk([("lm_head.weight", tensor)], 0) is None
-    assert direct.direct_expert_chunk([], 0) is None
-
-
-def _layout():
-    return dict(
-        tp_size=2,
-        ep_size=1,
-        pp_size=1,
-        moe_dp_size=1,
-        moe_runner_backend="triton",
-        dtype="bfloat16",
-        disable_shared_experts_fusion=True,
-        enable_eplb=False,
-        ep_num_redundant_experts=0,
-        init_expert_location="trivial",
-        ep_join_mode=None,
-        quantization=None,
-        enable_lora=None,
-        speculative_algorithm=None,
-        json_model_override_args="{}",
-    )
-
-
-@pytest.mark.parametrize(
-    "override",
-    [
-        dict(tp_size=4),
-        dict(ep_size=2),
-        dict(pp_size=2),
-        dict(moe_dp_size=2),
-        dict(moe_runner_backend="auto"),
-        dict(moe_runner_backend="triton_kernel"),
-        dict(moe_runner_backend="flashinfer_cutlass"),
-        dict(dtype="float16"),
-        dict(dtype="auto"),
-        dict(disable_shared_experts_fusion=False),
-        dict(enforce_shared_experts_fusion=True),
-        dict(enable_eplb=True),
-        dict(ep_num_redundant_experts=1),
-        dict(init_expert_location="random"),
-        dict(ep_join_mode="scale"),
-        dict(quantization="fp8"),
-        dict(enable_lora=True),
-        dict(speculative_algorithm="EAGLE"),
-        dict(json_model_override_args='{"hidden_size":4096}'),
-    ],
-)
-def test_unsupported_runtime_layout_is_rejected(override):
-    direct.validate_direct_receiver(_layout())
-    with pytest.raises(ValueError):
-        direct.validate_direct_receiver(_layout() | override)
-
-
-@pytest.mark.parametrize(
-    "packets,formats",
-    [
-        ([[("direct", "d0")], [("direct", "d1")]], ["direct"]),
-        ([[("flattened_bucket", "f0")], [("flattened_bucket", "f1")]], ["flattened_bucket"]),
-        ([[], [("direct", "d1")]], ["direct"]),
-        ([[("direct", "d0")], []], ["direct"]),
-        ([[("direct", "d0")], [("flattened_bucket", "f1")]], ["direct", "flattened_bucket"]),
-        ([[], []], []),
-    ],
-)
-def test_mixed_and_empty_ranks_receive_one_consistent_format_per_request(packets, formats):
-    calls, held = [], []
-
-    def serialize(value, output_str):
-        assert output_str
-        return "empty-direct" if isinstance(value, list) else "empty-flat"
-
-    def remote(**kwargs):
-        calls.append(kwargs)
-        return len(calls)
-
-    refs = direct.dispatch_typed_buckets(
-        packets,
-        engine=SimpleNamespace(update_weights_from_tensor=SimpleNamespace(remote=remote)),
-        version=12,
-        serializer=SimpleNamespace(serialize=serialize),
-        empty_flattened=lambda: dict(flattened_tensor=torch.empty(0, dtype=torch.uint8), metadata=[]),
-        held=held,
-    )
-    assert len(refs) == len(formats) and [c["load_format"] for c in calls] == formats
-    for call in calls:
-        assert call["weight_version"] == "12"
-        for rank, value in enumerate(call["serialized_named_tensors"]):
-            if packets[rank] and packets[rank][0][0] == call["load_format"]:
-                assert value == packets[rank][0][1]
-            else:
-                assert value == ("empty-direct" if call["load_format"] == "direct" else "empty-flat")
-    assert all(isinstance(value, list) or value["flattened_tensor"].numel() == 0 for value in held)

From 7c216d325d08485394e67ae2be46eedb4456581c Mon Sep 17 00:00:00 2001
From: yuanlehome <yuanlehome@163.com>
Date: Mon, 19 Oct 2026 20:35:42 +0800
Subject: [PATCH 3/6] test(megatron): guard optional backend imports

Use importorskip for weight-sync tests that require Megatron Core,
Megatron Bridge or SGLang. Guard only the backend-dependent inline
log-prob test so its pure CPU eligibility coverage still runs.

Keep the full regression coverage active in training environments
with the optional backends installed.
---
 .../megatron/test_colocate_inline_log_probs.py        |  1 +
 .../weight_update/test_async_gdn_weight_gather.py     |  3 +++
 .../megatron/weight_update/test_auto_weight_sync.py   | 11 ++++++++---
 .../weight_update/test_expert_capabilities.py         |  5 ++++-
 4 files changed, 16 insertions(+), 4 deletions(-)

diff --git a/tests/backends/megatron/test_colocate_inline_log_probs.py b/tests/backends/megatron/test_colocate_inline_log_probs.py
index e46a7aefc..bf0c98a1d 100644
--- a/tests/backends/megatron/test_colocate_inline_log_probs.py
+++ b/tests/backends/megatron/test_colocate_inline_log_probs.py
@@ -153,6 +153,7 @@ def test_existing_modes_do_not_enter_colocate_optimization(actor, inline_guard):
 
 def test_zero_kl_uses_sampling_shape_without_sampling_values(monkeypatch):
     torch = pytest.importorskip("torch")
+    pytest.importorskip("megatron.core", reason="Loss tests require Megatron Core", exc_type=ImportError)
     calls = []
 
     def compute_impl(args, **kwargs):
diff --git a/tests/backends/megatron/weight_update/test_async_gdn_weight_gather.py b/tests/backends/megatron/weight_update/test_async_gdn_weight_gather.py
index d1c6cb39b..1fa280119 100644
--- a/tests/backends/megatron/weight_update/test_async_gdn_weight_gather.py
+++ b/tests/backends/megatron/weight_update/test_async_gdn_weight_gather.py
@@ -8,6 +8,9 @@
 import torch
 
 
+pytest.importorskip("megatron.core", reason="TP gather tests require Megatron Core", exc_type=ImportError)
+
+
 @pytest.mark.parametrize("tp", [1, 2, 4])
 @pytest.mark.parametrize("wrapped", [False, True])
 @pytest.mark.parametrize("async_mode", [False, True])
diff --git a/tests/backends/megatron/weight_update/test_auto_weight_sync.py b/tests/backends/megatron/weight_update/test_auto_weight_sync.py
index 538c1c207..1f69ffa41 100644
--- a/tests/backends/megatron/weight_update/test_auto_weight_sync.py
+++ b/tests/backends/megatron/weight_update/test_auto_weight_sync.py
@@ -6,9 +6,14 @@
 import pytest
 import torch
 
-from relax.backends.megatron.weight_update import update_weight_from_tensor as sync
-from relax.backends.megatron.weight_update.expert_layout import DirectExpertLayout
-from relax.backends.megatron.weight_update.hf_weight_iterator_bridge import HfWeightIteratorBridge
+
+pytest.importorskip("megatron.core", reason="Weight sync tests require Megatron Core", exc_type=ImportError)
+pytest.importorskip("megatron.bridge", reason="Weight sync tests require Megatron Bridge", exc_type=ImportError)
+pytest.importorskip("sglang.srt.utils", reason="Weight sync tests require SGLang serialization", exc_type=ImportError)
+
+from relax.backends.megatron.weight_update import update_weight_from_tensor as sync  # noqa: E402
+from relax.backends.megatron.weight_update.expert_layout import DirectExpertLayout  # noqa: E402
+from relax.backends.megatron.weight_update.hf_weight_iterator_bridge import HfWeightIteratorBridge  # noqa: E402
 
 
 def receiver():
diff --git a/tests/backends/megatron/weight_update/test_expert_capabilities.py b/tests/backends/megatron/weight_update/test_expert_capabilities.py
index 13d95bd98..a2856a196 100644
--- a/tests/backends/megatron/weight_update/test_expert_capabilities.py
+++ b/tests/backends/megatron/weight_update/test_expert_capabilities.py
@@ -4,7 +4,10 @@
 
 import pytest
 
-from relax.backends.megatron.weight_update.bridge_converter import BridgeConverter
+
+pytest.importorskip("megatron.core", reason="Bridge capability tests require Megatron Core", exc_type=ImportError)
+
+from relax.backends.megatron.weight_update.bridge_converter import BridgeConverter  # noqa: E402
 
 
 @pytest.mark.parametrize("prefix", ["model.layers.2", "model.language_model.layers.2"])

From f3b42469f1ed1486ddc275d55a4ca436295b28aa Mon Sep 17 00:00:00 2001
From: yuanlehome <yuanlehome@163.com>
Date: Tue, 20 Oct 2026 11:34:17 +0800
Subject: [PATCH 4/6] perf(weight-sync): generalize MoE sync and simplify
 defaults

Derive native BF16 expert shards from receiver MoE TP, EP and DP, and support
Triton, DeepGEMM and FlashInfer CUTLASS independently of the A2A dispatcher.
Preserve native gate/up order and fall back for unsupported weight layouts.

Use scheduler-resolved defaults, disable shared-expert fusion by default,
and clear visual embeddings during cache flush. Remove redundant recipe
overrides and the opt-in CUTLASS grouped-GEMM setting.

Validation: all-files pre-commit; 432 tests passed, 2 environment skips;
8 multi-GPU cases passed, including DeepEP normal and low-latency.
The Qwen3.5 TP2 packing/copy microbenchmark showed no regression.
FlashInfer A2A initialization remains blocked by the host Linux 5.4 kernel.
---
 docker/patch/sglang/v0.5.17.patch             |  59 +++++--
 .../megatron/weight_update/expert_layout.py   |  72 ++++++--
 relax/backends/sglang/sglang_engine.py        |  56 ++++++-
 .../run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh   |  26 ---
 .../weight_update/test_auto_weight_sync.py    |  88 ++++++++--
 .../weight_update/test_expert_layout.py       | 130 ++++++++++++---
 tests/backends/sglang/test_sglang_engine.py   | 156 +++++++++++++++++-
 7 files changed, 488 insertions(+), 99 deletions(-)

diff --git a/docker/patch/sglang/v0.5.17.patch b/docker/patch/sglang/v0.5.17.patch
index 4eea53a04..bed2b40f3 100644
--- a/docker/patch/sglang/v0.5.17.patch
+++ b/docker/patch/sglang/v0.5.17.patch
@@ -1154,7 +1154,7 @@ index 6eb0ad1a..d27815f3 100644
          """Load a new LoRA adapter without re-launching the engine."""
  
 diff --git a/python/sglang/srt/entrypoints/http_server.py b/python/sglang/srt/entrypoints/http_server.py
-index ac38b466..ea54bac1 100644
+index ac38b46..abdc04f 100644
 --- a/python/sglang/srt/entrypoints/http_server.py
 +++ b/python/sglang/srt/entrypoints/http_server.py
 @@ -132,6 +132,7 @@ from sglang.srt.managers.io_struct import (
@@ -1208,6 +1208,15 @@ index ac38b466..ea54bac1 100644
  
      server_args = _global_state.tokenizer_manager.server_args
  
+@@ -946,7 +955,7 @@ async def classify_request(obj: EmbeddingReqInput, request: Request):
+ @app.api_route("/flush_cache", methods=["GET", "POST"])
+ @auth_level(AuthLevel.ADMIN_OPTIONAL)
+ async def flush_cache(timeout: float = Query(0.0, ge=0.0)):
+-    """Flush the radix cache."""
++    """Flush runtime caches, including multimodal embeddings."""
+     ret = await _global_state.tokenizer_manager.flush_cache(timeout_s=timeout)
+     if ret.success:
+         content = (
 @@ -1409,6 +1418,22 @@ async def update_weights_from_ipc(
          return ORJSONResponse(content, status_code=HTTPStatus.BAD_REQUEST)
  
@@ -1546,10 +1555,18 @@ index a45d122f..0334e353 100755
          ):
              # Even the last remaining request cannot fit in memory.
 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py
-index a9615a14c5..aca7bd7061 100644
+index a9615a1..b2dfb92 100644
 --- a/python/sglang/srt/managers/scheduler.py
 +++ b/python/sglang/srt/managers/scheduler.py
-@@ -141,6 +141,7 @@ from sglang.srt.managers.io_struct import (
+@@ -99,6 +99,7 @@ from sglang.srt.layers.quantization.fp8_utils import initialize_fp8_gemm_config
+ from sglang.srt.layers.quantization.unquant import initialize_bf16_gemm_config
+ from sglang.srt.lora.lora_drainer import LoRADrainer
+ from sglang.srt.lora.lora_overlap_loader import LoRAOverlapLoader
++from sglang.srt.managers import mm_utils
+ from sglang.srt.managers.disagg_service import maybe_create_ascend_config_store
+ from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator
+ from sglang.srt.managers.io_struct import (
+@@ -141,6 +142,7 @@ from sglang.srt.managers.io_struct import (
      LoadLoRAAdapterReqOutput,
      OpenSessionReqInput,
      PauseGenerationReqInput,
@@ -1557,7 +1574,7 @@ index a9615a14c5..aca7bd7061 100644
      ProfileReq,
      ReleaseMemoryOccupationReqInput,
      RemoveExternalCorpusReqInput,
-@@ -161,6 +162,8 @@ from sglang.srt.managers.io_struct import (
+@@ -161,6 +163,8 @@ from sglang.srt.managers.io_struct import (
      TokenizedGenerateReqInput,
      UnloadLoRAAdapterReqInput,
      UnloadLoRAAdapterReqOutput,
@@ -1566,7 +1583,7 @@ index a9615a14c5..aca7bd7061 100644
      UpdateWeightFromDiskReqInput,
      UpdateWeightsFromDistributedReqInput,
      UpdateWeightsFromIPCReqInput,
-@@ -280,6 +283,8 @@ from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
+@@ -280,6 +284,8 @@ from sglang.srt.speculative.base_spec_worker import BaseSpecWorker
  from sglang.srt.speculative.dflash_utils import validate_dflash_request
  from sglang.srt.speculative.eagle_utils import get_draft_recurrent_hidden_state_spec
  from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -1575,7 +1592,7 @@ index a9615a14c5..aca7bd7061 100644
  from sglang.srt.utils import (
      DynamicGradMode,
      configure_gc_logger,
-@@ -1539,6 +1544,10 @@ class Scheduler(
+@@ -1539,6 +1545,10 @@ class Scheduler(
                      UpdateWeightsFromIPCReqInput,
                      self.weight_updater.update_weights_from_ipc,
                  ),
@@ -1586,7 +1603,7 @@ index a9615a14c5..aca7bd7061 100644
                  (
                      GetWeightsByNameReqInput,
                      self.weight_updater.get_weights_by_name,
-@@ -1571,6 +1580,10 @@ class Scheduler(
+@@ -1571,6 +1581,10 @@ class Scheduler(
                      LoadLoRAAdapterFromTensorsReqInput,
                      self.load_lora_adapter_from_tensors,
                  ),
@@ -1597,7 +1614,7 @@ index a9615a14c5..aca7bd7061 100644
                  (UnloadLoRAAdapterReqInput, self.unload_lora_adapter),
                  (PauseGenerationReqInput, self.pause_generation),
                  (ContinueGenerationReqInput, self.continue_generation),
-@@ -1627,6 +1640,8 @@ class Scheduler(
+@@ -1627,6 +1641,8 @@ class Scheduler(
          self.tree_cache.release_host_resources()
          if self.decode_offload_manager is not None:
              self.decode_offload_manager.release_host_resources()
@@ -1606,7 +1623,29 @@ index a9615a14c5..aca7bd7061 100644
  
      def run_event_loop(self) -> None:
          """Run the scheduler's event loop.
-@@ -4457,6 +4472,12 @@ class Scheduler(
+@@ -4133,7 +4149,7 @@ class Scheduler(
+         return DetachHiCacheStorageReqOutput(success=False, message=msg)
+
+     def flush_cache(self, empty_cache: bool = True):
+-        """Flush memory pools (e.g., KV cache, Mamba cache) and optionally empty device allocator cache."""
++        """Flush runtime caches and optionally empty the device allocator cache."""
+         if self.is_fully_idle():
+             self.cur_batch_for_debug = None
+             self.last_batch = None
+@@ -4146,6 +4162,12 @@ class Scheduler(
+             if self.draft_worker:
+                 self.draft_worker.clear_cache_pool()
+
++            # Drop cached embeddings so weight updates cannot reuse stale encoder
++            # outputs. Do this before trimming the allocator so colocated workloads
++            # can reclaim the memory.
++            if mm_utils.embedding_cache is not None:
++                mm_utils.embedding_cache.clear()
++
+             if empty_cache:
+                 current_platform.empty_cache()
+             # Per-DP-group leader logs once: ranks within a DP group are
+@@ -4457,6 +4479,12 @@ class Scheduler(
                  # The request will still run one decode forward pass.
                  # Then we reuse all existing code to clean up the KV cache allocation.
                  logger.debug(f"Abort running request. {req.rid=}")
@@ -1619,7 +1658,7 @@ index a9615a14c5..aca7bd7061 100644
                  req.to_finish = FINISH_ABORT()
  
      def _pause_engine(self) -> Tuple[List[Req], int]:
-@@ -4669,6 +4690,14 @@ class Scheduler(
+@@ -4669,6 +4697,14 @@ class Scheduler(
          result = self.tp_worker.load_lora_adapter_from_tensors(recv_req)
          return result
  
diff --git a/relax/backends/megatron/weight_update/expert_layout.py b/relax/backends/megatron/weight_update/expert_layout.py
index a68739b5c..a775f23e4 100644
--- a/relax/backends/megatron/weight_update/expert_layout.py
+++ b/relax/backends/megatron/weight_update/expert_layout.py
@@ -25,25 +25,57 @@ class DirectExpertLayout:
     num_experts: int
     hidden_size: int
     intermediate_size: int
+    # SGLang's TP size is the engine/IPC world size, including MoE EP and DP.
     tp_size: int
+    up_proj_first: bool = False
+    ep_size: int = 1
+    moe_dp_size: int = 1
+
+    def __post_init__(self) -> None:
+        dimensions = (
+            self.num_experts,
+            self.hidden_size,
+            self.intermediate_size,
+            self.tp_size,
+            self.ep_size,
+            self.moe_dp_size,
+        )
+        if any(type(value) is not int or value <= 0 for value in dimensions):
+            raise ValueError("Direct expert geometry must be explicitly specified and positive")
+        if self.tp_size % (self.ep_size * self.moe_dp_size):
+            raise ValueError("Receiver TP size must be divisible by MoE EP * DP")
+        if self.num_experts % self.ep_size:
+            raise ValueError("Expert count must be divisible by receiver EP size")
+        if self.intermediate_size % self.moe_tp_size:
+            raise ValueError("Expert intermediate size must be divisible by receiver MoE TP size")
+
+    @property
+    def moe_tp_size(self) -> int:
+        return self.tp_size // self.ep_size // self.moe_dp_size
 
     @classmethod
     def from_receiver(cls, config: Mapping[str, Any]) -> "DirectExpertLayout":
         required = dict(
-            ep_size=1,
             pp_size=1,
-            moe_dp_size=1,
-            moe_runner_backend="triton",
             disable_shared_experts_fusion=True,
             dtype="bfloat16",
         )
         if any(config.get(key) != value for key, value in required.items()):
-            raise ValueError("Direct experts require BF16 Triton, EP1/PP1/MoE-DP1 and unfused shared experts")
+            raise ValueError("Direct experts require BF16, PP1 and unfused shared experts")
+        backend = config.get("moe_runner_backend")
+        if backend not in ("triton", "deep_gemm", "flashinfer_cutlass"):
+            raise ValueError("Unsupported direct expert runner layout")
+        if config.get("is_cuda") is not True:
+            raise ValueError("Direct BF16 experts require a verified NVIDIA CUDA receiver")
+        # A2A dispatches activations; it does not change these runners' BF16
+        # weight layout. SGLang validates runner/dispatcher compatibility.
         if (
             config.get("enable_eplb")
             or config.get("ep_num_redundant_experts", 0)
             or config.get("init_expert_location") not in (None, "trivial")
             or config.get("ep_join_mode") is not None
+            or config.get("elastic_ep_backend") is not None
+            or config.get("dwdp_size", 1) != 1
             or config.get("quantization") is not None
             or config.get("quantization_config")
             or config.get("enable_lora")
@@ -53,15 +85,15 @@ def from_receiver(cls, config: Mapping[str, Any]) -> "DirectExpertLayout":
             or json.loads(config.get("json_model_override_args") or "{}")
         ):
             raise ValueError("Direct experts do not support quantization, remapping, LoRA or model overrides")
-        dimensions = tuple(
-            config.get(key) for key in ("num_experts", "hidden_size", "moe_intermediate_size", "tp_size")
+        return cls(
+            num_experts=config.get("num_experts"),
+            hidden_size=config.get("hidden_size"),
+            intermediate_size=config.get("moe_intermediate_size"),
+            tp_size=config.get("tp_size"),
+            up_proj_first=backend == "flashinfer_cutlass",
+            ep_size=config.get("ep_size"),
+            moe_dp_size=config.get("moe_dp_size"),
         )
-        if any(type(value) is not int or value <= 0 for value in dimensions):
-            raise ValueError("Direct expert geometry must be explicitly specified and positive")
-        layout = cls(*dimensions)
-        if layout.intermediate_size % layout.tp_size:
-            raise ValueError("Expert intermediate size must be divisible by receiver TP size")
-        return layout
 
     def shard(self, tensor: torch.Tensor, projection: str, tp_rank: int) -> torch.Tensor:
         if not 0 <= tp_rank < self.tp_size or tensor.dtype != torch.bfloat16:
@@ -70,10 +102,20 @@ def shard(self, tensor: torch.Tensor, projection: str, tp_rank: int) -> torch.Te
         expected = (e, 2 * i, h) if projection == "gate_up_proj" else (e, h, i)
         if tuple(tensor.shape) != expected or projection not in ("gate_up_proj", "down_proj"):
             raise ValueError(f"Unsupported expert projection or shape: {projection}: {tuple(tensor.shape)}")
-        width = i // self.tp_size
-        start = tp_rank * width
+        # SGLang orders ranks as [MoE DP, EP, MoE TP]. DP replicas share the
+        # same expert range; MoE TP=1 keeps each local expert matrix intact.
+        ep_rank = tp_rank // self.moe_tp_size % self.ep_size
+        local_experts = e // self.ep_size
+        tensor = tensor[ep_rank * local_experts : (ep_rank + 1) * local_experts]
+        width = i // self.moe_tp_size
+        start = tp_rank % self.moe_tp_size * width
         if projection == "gate_up_proj":
-            return torch.cat((tensor[:, start : start + width], tensor[:, i + start : i + start + width]), dim=1)
+            if self.moe_tp_size == 1 and not self.up_proj_first:
+                return tensor.contiguous()
+            gate = tensor[:, start : start + width]
+            up = tensor[:, i + start : i + start + width]
+            # Match UnquantizedFusedMoEMethod.load_up_proj_weight_first.
+            return torch.cat((up, gate) if self.up_proj_first else (gate, up), dim=1)
         return tensor[:, :, start : start + width].contiguous()
 
 
diff --git a/relax/backends/sglang/sglang_engine.py b/relax/backends/sglang/sglang_engine.py
index cd6fb16f2..83e28a78f 100644
--- a/relax/backends/sglang/sglang_engine.py
+++ b/relax/backends/sglang/sglang_engine.py
@@ -777,20 +777,30 @@ def get_expert_weight_sync_layout(self) -> dict:
 
         if self.node_rank != 0 or self.worker_type != "regular":
             raise ValueError("Expert routing requires a regular node-0 rollout engine")
-        response = requests.get(f"http://{self.server_host}:{self.server_port}/server_info", timeout=30)
-        response.raise_for_status()
-        info = response.json()
+        for _ in range(2):
+            response = requests.get(f"http://{self.server_host}:{self.server_port}/server_info", timeout=30)
+            response.raise_for_status()
+            info = response.json()
+            # /server_info can time out fetching scheduler state while still
+            # returning HTTP 200. Retry once before rejecting an unknown auto layout.
+            if info.get("internal_states") or info.get("moe_runner_backend") != "auto":
+                break
         keys = (
+            "model_path",
+            "device",
             "tp_size",
             "ep_size",
             "pp_size",
             "moe_dp_size",
+            "dwdp_size",
             "enable_eplb",
             "ep_num_redundant_experts",
             "init_expert_location",
             "ep_join_mode",
+            "elastic_ep_backend",
             "speculative_algorithm",
             "moe_runner_backend",
+            "moe_a2a_backend",
             "json_model_override_args",
             "dtype",
             "quantization",
@@ -800,8 +810,37 @@ def get_expert_weight_sync_layout(self) -> dict:
             "kt_weight_path",
         )
         layout = {key: info[key] for key in keys if key in info}
-        config = get_hf_config(info["model_path"])
+        # Scheduler-side load-time overrides may differ from the HTTP process's
+        # startup config. Never infer the native layout from a stale "auto".
+        states = info.get("internal_states") or []
+        if states:
+            layouts = [{**layout, **{key: state[key] for key in keys if key in state}} for state in states]
+            if any(candidate != layouts[0] for candidate in layouts[1:]):
+                raise ValueError("Rollout schedulers disagree on the expert weight layout")
+            layout = layouts[0]
+        # ROCm also reports device="cuda". The local platform probe only describes
+        # receivers spawned here, not an externally managed HTTP engine.
+        layout["is_cuda"] = (
+            getattr(self, "process", None) is not None
+            and layout.get("device") == "cuda"
+            and device_utils.get_accelerator_type() == device_utils.AcceleratorType.CUDA
+        )
+        config = get_hf_config(layout["model_path"])
         text_config = getattr(config, "text_config", None) or config
+        if "dtype" in layout:
+            from sglang.srt.configs.model_config import _get_and_verify_dtype
+
+            layout["dtype"] = str(_get_and_verify_dtype(text_config, layout["dtype"])).removeprefix("torch.")
+        if (
+            states
+            and all("moe_runner_backend" in state for state in states)
+            and layout.get("moe_runner_backend") == "auto"
+            and layout["is_cuda"]
+        ):
+            # Unquantized CUDA FusedMoE uses the standard Triton layout for auto.
+            # Quantization, remapping and fused shared experts are still checked
+            # by DirectExpertLayout; this does not change the engine's settings.
+            layout["moe_runner_backend"] = "triton"
         layout["num_experts"] = getattr(text_config, "num_experts", None) or getattr(
             text_config, "n_routed_experts", None
         )
@@ -2076,6 +2115,15 @@ def _compute_server_args(
             kwargs[key] = value
             unused_keys.discard(key)
 
+    # Repeated policy updates benefit from keeping shared and routed experts
+    # separate. Explicit per-engine settings or enforced fusion take precedence.
+    if (
+        "disable_shared_experts_fusion" in server_arg_field_names
+        and "disable_shared_experts_fusion" not in (sglang_overrides or {})
+        and not kwargs.get("enforce_shared_experts_fusion", False)
+    ):
+        kwargs["disable_shared_experts_fusion"] = True
+
     if (
         "cuda_graph_backend_prefill" in server_arg_field_names
         and kwargs.get("enable_memory_saver")
diff --git a/scripts/training/multimodal/run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh b/scripts/training/multimodal/run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh
index b01c2c7df..44d2c8268 100644
--- a/scripts/training/multimodal/run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh
+++ b/scripts/training/multimodal/run-qwen35-35B-A3B-8xgpu-openr1mm-r3.sh
@@ -129,32 +129,6 @@ MISC_ARGS=(
    --attention-backend flash
 )
 
-# Match the native BF16 expert layout used by the guarded direct IPC loader.
-SGLANG_ARGS+=(--sglang-moe-runner-backend triton --sglang-disable-shared-experts-fusion --sglang-dtype bfloat16)
-
-RUNTIME_ENV_JSON=$(python3 - <<'PYENV'
-import json
-import os
-
-runtime_env = json.loads(os.environ["RUNTIME_ENV_JSON"])
-# Vision embeddings must be recomputed after each weight update.
-runtime_env.setdefault("env_vars", {})["SGLANG_VLM_CACHE_SIZE_MB"] = "0"
-# The backend validates single-update, zero-KL GRPO before reusing log-probs.
-runtime_env["env_vars"]["NVTE_USE_CUTLASS_GROUPED_GEMM"] = "1"
-# Export BF16 experts in bulk, then send each receiver its TP-local shard.
-# Preserve the platform Trial identity and its collector path in Ray workers.
-# These are public variable names; their values are inherited at launch.
-for key in (
-    "QS_TRIAL_ID", "QS_JOB_ID", "QS_USER", "GPU_TYPE", "CLUSTER",  # gitleaks:allow
-    "ETTR_LOG_FILE", "ETTR_TRACE_LOG_DIR", "ETTR_FRAMEWORK",
-    "QS_INJECT_AUTOLOAD", "QS_INJECT_MODEL_PROFILE",  # gitleaks:allow
-):
-    if key in os.environ:
-        runtime_env["env_vars"][key] = os.environ[key]
-print(json.dumps(runtime_env))
-PYENV
-)
-
 mkdir -p log
 
 ray job submit ${RAY_NO_WAIT:+--no-wait} --address="http://127.0.0.1:8265" \
diff --git a/tests/backends/megatron/weight_update/test_auto_weight_sync.py b/tests/backends/megatron/weight_update/test_auto_weight_sync.py
index 1f69ffa41..7e483867a 100644
--- a/tests/backends/megatron/weight_update/test_auto_weight_sync.py
+++ b/tests/backends/megatron/weight_update/test_auto_weight_sync.py
@@ -16,23 +16,42 @@
 from relax.backends.megatron.weight_update.hf_weight_iterator_bridge import HfWeightIteratorBridge  # noqa: E402
 
 
-def receiver():
-    return dict(
-        tp_size=2,
-        ep_size=1,
-        pp_size=1,
-        moe_dp_size=1,
-        moe_runner_backend="triton",
-        dtype="bfloat16",
-        disable_shared_experts_fusion=True,
-        num_experts=4,
-        hidden_size=6,
-        moe_intermediate_size=8,
+def receiver(backend="triton", **overrides):
+    return (
+        dict(
+            tp_size=2,
+            ep_size=1,
+            pp_size=1,
+            moe_dp_size=1,
+            moe_runner_backend=backend,
+            moe_a2a_backend="none",
+            is_cuda=True,
+            dtype="bfloat16",
+            disable_shared_experts_fusion=True,
+            num_experts=4,
+            hidden_size=6,
+            moe_intermediate_size=8,
+        )
+        | overrides
     )
 
 
-@pytest.mark.parametrize("failure", ["none", "local_not_ready", "peer_not_ready", "configure", "query", "peer_reject"])
-def test_negotiation_agrees_before_selecting_or_falling_back(monkeypatch, failure):
+@pytest.mark.parametrize(
+    "failure",
+    [
+        "none",
+        "local_not_ready",
+        "peer_not_ready",
+        "configure",
+        "query",
+        "peer_reject",
+        "unsupported_layout",
+        "mixed_layout",
+        "mixed_parallelism",
+    ],
+)
+@pytest.mark.parametrize("backend", ["triton", "deep_gemm", "flashinfer_cutlass"])
+def test_negotiation_agrees_before_selecting_or_falling_back(monkeypatch, failure, backend):
     updater = sync.UpdateWeightFromTensor.__new__(sync.UpdateWeightFromTensor)
     updater.args = SimpleNamespace(colocate=True)
     updater.use_distribute = updater._route_experts = updater.lora_enabled = False
@@ -50,7 +69,16 @@ def configure(layout):
     )
     engine = Mock()
     updater.rollout_engines = [engine]
-    get = Mock(side_effect=RuntimeError("query failed")) if failure == "query" else Mock(return_value=[receiver()])
+    receivers = [receiver(backend)]
+    if failure == "unsupported_layout":
+        receivers[0]["moe_runner_backend"] = "triton_kernel"
+    elif failure == "mixed_layout":
+        receivers.append(receiver("triton" if backend == "flashinfer_cutlass" else "flashinfer_cutlass"))
+        updater.rollout_engines.append(Mock())
+    elif failure == "mixed_parallelism":
+        receivers.append(receiver(backend, ep_size=2))
+        updater.rollout_engines.append(Mock())
+    get = Mock(side_effect=RuntimeError("query failed")) if failure == "query" else Mock(return_value=receivers)
     monkeypatch.setattr(sync.ray, "get", get)
     monkeypatch.setattr(sync, "get_gloo_group", lambda: "gloo")
     monkeypatch.setattr(sync.dist, "get_rank", lambda **kwargs: 0)
@@ -63,7 +91,7 @@ def reduce(tensor, *, op, group):
             tensor.zero_()
 
     monkeypatch.setattr(sync.dist, "all_reduce", reduce)
-    updater._configure_direct_experts([2])
+    updater._configure_direct_experts([2] * len(receivers))
     assert (updater._direct_expert_layout is not None) == (failure == "none")
     if failure in ("local_not_ready", "peer_not_ready"):
         assert len(reductions) == 1 and configured == [None]
@@ -74,6 +102,34 @@ def reduce(tensor, *, op, group):
             assert configured[-1] is None
 
 
+@pytest.mark.parametrize(
+    "topology",
+    [
+        dict(tp_size=2, ep_size=2, moe_a2a_backend="deepep"),
+        dict(tp_size=4, ep_size=2),
+        dict(tp_size=4, ep_size=2, moe_dp_size=2),
+    ],
+)
+def test_negotiation_preserves_full_engine_placement_for_expert_parallelism(monkeypatch, topology):
+    updater = sync.UpdateWeightFromTensor.__new__(sync.UpdateWeightFromTensor)
+    updater.args = SimpleNamespace(colocate=True)
+    updater.use_distribute = updater._route_experts = updater.lora_enabled = False
+    updater.quantization_config = None
+    configure = Mock(side_effect=lambda layout: layout is not None)
+    updater._hf_weight_iterator = SimpleNamespace(configure_direct_experts=configure, _bf16_expert_buckets=[object()])
+    updater.rollout_engines = [Mock()]
+    config = receiver("deep_gemm", **topology)
+    monkeypatch.setattr(sync.ray, "get", lambda _: [config])
+    monkeypatch.setattr(sync, "get_gloo_group", lambda: "gloo")
+    monkeypatch.setattr(sync.dist, "get_rank", lambda **kwargs: 0)
+    monkeypatch.setattr(sync.dist, "broadcast_object_list", lambda *args, **kwargs: None)
+    monkeypatch.setattr(sync.dist, "all_reduce", lambda *args, **kwargs: None)
+    updater._configure_direct_experts([config["tp_size"]])
+    expected = DirectExpertLayout.from_receiver(config)
+    assert updater._direct_expert_layout == expected
+    configure.assert_called_once_with(expected)
+
+
 @pytest.mark.parametrize("rank", [0, 3])
 def test_reconnect_creates_groups_once_including_placeholder_ranks(monkeypatch, rank):
     updater = sync.UpdateWeightFromTensor.__new__(sync.UpdateWeightFromTensor)
diff --git a/tests/backends/megatron/weight_update/test_expert_layout.py b/tests/backends/megatron/weight_update/test_expert_layout.py
index 9eeb2ac8c..8eb9ebf29 100644
--- a/tests/backends/megatron/weight_update/test_expert_layout.py
+++ b/tests/backends/megatron/weight_update/test_expert_layout.py
@@ -9,6 +9,20 @@
 from relax.backends.megatron.weight_update.ipc import dispatch_typed_buckets
 
 
+# Receiver-rank order from SGLang's [MoE DP, EP, MoE TP] process groups.
+# Each entry is (expert rank, expert tensor-parallel rank).
+TOPOLOGIES = [
+    pytest.param(1, 1, 1, [(0, 0)], id="single"),
+    pytest.param(2, 1, 1, [(0, 0), (0, 1)], id="tp2"),
+    pytest.param(2, 2, 1, [(0, 0), (1, 0)], id="ep2"),
+    pytest.param(4, 4, 1, [(0, 0), (1, 0), (2, 0), (3, 0)], id="ep4"),
+    pytest.param(4, 2, 1, [(0, 0), (0, 1), (1, 0), (1, 1)], id="ep2-tp2"),
+    pytest.param(4, 2, 2, [(0, 0), (1, 0), (0, 0), (1, 0)], id="ep2-dp2"),
+    pytest.param(4, 1, 2, [(0, 0), (0, 1), (0, 0), (0, 1)], id="tp2-dp2"),
+    pytest.param(8, 2, 2, [(0, 0), (0, 1), (1, 0), (1, 1), (0, 0), (0, 1), (1, 0), (1, 1)], id="ep2-tp2-dp2"),
+]
+
+
 def receiver(**overrides):
     return (
         dict(
@@ -17,6 +31,8 @@ def receiver(**overrides):
             pp_size=1,
             moe_dp_size=1,
             moe_runner_backend="triton",
+            moe_a2a_backend="none",
+            is_cuda=True,
             dtype="bfloat16",
             disable_shared_experts_fusion=True,
             num_experts=4,
@@ -27,19 +43,31 @@ def receiver(**overrides):
     )
 
 
-@pytest.mark.parametrize("tp", [1, 2, 4])
+@pytest.mark.parametrize("tp,ep,dp,ranks", TOPOLOGIES)
+@pytest.mark.parametrize("backend", ["triton", "deep_gemm", "flashinfer_cutlass"])
 @pytest.mark.parametrize("prefix", ["model.layers", "model.language_model.layers"])
 @pytest.mark.parametrize("geometry", [(4, 6, 8), (8, 10, 12)])
-def test_direct_shards_preserve_expert_and_gate_up_order(tp, prefix, geometry):
+def test_direct_shards_preserve_expert_and_gate_up_order(tp, ep, dp, ranks, backend, prefix, geometry):
     e, h, i = geometry
-    layout = DirectExpertLayout(e, h, i, tp)
+    layout = DirectExpertLayout.from_receiver(
+        receiver(
+            tp_size=tp,
+            ep_size=ep,
+            moe_dp_size=dp,
+            num_experts=e,
+            hidden_size=h,
+            moe_intermediate_size=i,
+            moe_runner_backend=backend,
+        )
+    )
     gate = torch.arange(e * i * h).reshape(e, i, h).to(torch.bfloat16)
     up = gate + 512
     down = torch.arange(e * h * i).reshape(e, h, i).to(torch.bfloat16)
     packed = torch.cat((gate, up), dim=1)
-    for rank in range(tp):
+    for rank, (expert_rank, expert_tp_rank) in enumerate(ranks):
         [(name, actual)] = direct_expert_chunk([(f"{prefix}.9.mlp.experts.gate_up_proj", packed)], layout, rank)
-        expected = torch.cat((gate.chunk(tp, dim=1)[rank], up.chunk(tp, dim=1)[rank]), dim=1)
+        halves = tuple(t.chunk(ep)[expert_rank].chunk(tp // ep // dp, dim=1)[expert_tp_rank] for t in (gate, up))
+        expected = torch.cat(halves[::-1] if backend == "flashinfer_cutlass" else halves, dim=1)
         assert name == "model.layers.9.mlp.experts.w13_weight"
         assert torch.equal(actual, expected) and actual.is_contiguous()
         snapshot = actual.clone()
@@ -47,37 +75,52 @@ def test_direct_shards_preserve_expert_and_gate_up_order(tp, prefix, geometry):
         assert torch.equal(actual, snapshot)
         [(name, actual)] = direct_expert_chunk([(f"{prefix}.9.mlp.experts.down_proj", down)], layout, rank)
         assert name == "model.layers.9.mlp.experts.w2_weight"
-        assert torch.equal(actual, down.chunk(tp, dim=2)[rank]) and actual.is_contiguous()
+        expected = down.chunk(ep)[expert_rank].chunk(tp // ep // dp, dim=2)[expert_tp_rank]
+        assert torch.equal(actual, expected) and actual.is_contiguous()
     assert direct_expert_chunk([("model.norm.weight", gate)], layout, 0) is None
     assert direct_expert_chunk([], layout, 0) is None
 
 
-@pytest.mark.parametrize("tp", [1, 2, 4])
-def test_direct_matches_installed_sglang_fused_moe_loader(tp):
+@pytest.mark.parametrize("tp,ep,dp,ranks", TOPOLOGIES)
+@pytest.mark.parametrize("backend", ["triton", "deep_gemm", "flashinfer_cutlass"])
+def test_direct_matches_installed_sglang_fused_moe_loader(monkeypatch, tp, ep, dp, ranks, backend):
     """Use the receiver's real per-expert loaders as the independent oracle."""
-    layer_module = pytest.importorskip("sglang.srt.layers.moe.fused_moe_triton.layer")
+    layer_module = pytest.importorskip("sglang.srt.layers.moe.fused_moe_triton.layer", exc_type=ImportError)
+    unquant = pytest.importorskip("sglang.srt.layers.quantization.unquant", exc_type=ImportError)
+    moe = pytest.importorskip("sglang.srt.layers.moe.utils", exc_type=ImportError)
+    monkeypatch.setattr(unquant, "get_moe_runner_backend", lambda: moe.MoeRunnerBackend(backend))
+    quant_method = unquant.UnquantizedFusedMoEMethod(use_deep_gemm=backend == "deep_gemm")
     e, h, i = 4, 6, 8
-    layout = DirectExpertLayout(e, h, i, tp)
+    layout = DirectExpertLayout.from_receiver(
+        receiver(tp_size=tp, ep_size=ep, moe_dp_size=dp, moe_runner_backend=backend)
+    )
+    moe_tp = tp // ep // dp
     gate = torch.arange(e * i * h).reshape(e, i, h).to(torch.bfloat16)
     up, down = gate + 512, gate.transpose(1, 2).contiguous() + 1024
     receiver_layer = SimpleNamespace(
         moe_runner_config=SimpleNamespace(is_gated=True),
-        quant_method=SimpleNamespace(),
+        quant_method=quant_method,
         quant_config=None,
         use_padded_loading=False,
         use_presharded_weights=False,
         use_triton_kernels=False,
-        moe_tp_size=tp,
+        moe_tp_size=moe_tp,
+        _num_local_routed=e // ep,
+        _has_fused_shared=False,
     )
-    for rank in range(tp):
+    for rank, (expert_rank, expert_tp_rank) in enumerate(ranks):
+        receiver_layer._expert_storage_rank = expert_rank
         w13, w2 = (
-            torch.empty(e, 2 * i // tp, h, dtype=torch.bfloat16),
-            torch.empty(e, h, i // tp, dtype=torch.bfloat16),
+            torch.empty(e // ep, 2 * i // moe_tp, h, dtype=torch.bfloat16),
+            torch.empty(e // ep, h, i // moe_tp, dtype=torch.bfloat16),
         )
         for expert in range(e):
-            layer_module.FusedMoE._load_w13(receiver_layer, w13[expert], 0, "w1", gate[expert], rank)
-            layer_module.FusedMoE._load_w13(receiver_layer, w13[expert], 0, "w3", up[expert], rank)
-            layer_module.FusedMoE._load_w2(receiver_layer, w2[expert], 1, "w2", down[expert], rank)
+            local_expert = layer_module.FusedMoE._map_global_expert_id_to_local_expert_id(receiver_layer, expert)
+            if local_expert == -1:
+                continue
+            layer_module.FusedMoE._load_w13(receiver_layer, w13[local_expert], 0, "w1", gate[expert], expert_tp_rank)
+            layer_module.FusedMoE._load_w13(receiver_layer, w13[local_expert], 0, "w3", up[expert], expert_tp_rank)
+            layer_module.FusedMoE._load_w2(receiver_layer, w2[local_expert], 1, "w2", down[expert], expert_tp_rank)
         assert torch.equal(w13, layout.shard(torch.cat((gate, up), dim=1), "gate_up_proj", rank))
         assert torch.equal(w2, layout.shard(down, "down_proj", rank))
 
@@ -92,12 +135,19 @@ def test_resolved_receiver_geometry_has_no_model_family_requirement(tp):
     [
         dict(tp_size=3),
         dict(tp_size=0),
-        dict(ep_size=2),
+        dict(ep_size=3),
+        dict(ep_size=None),
+        dict(ep_size=0),
+        dict(tp_size=3, ep_size=3),
         dict(pp_size=2),
-        dict(moe_dp_size=2),
+        dict(moe_dp_size=3),
+        dict(moe_dp_size=None),
+        dict(moe_dp_size=False),
         dict(moe_runner_backend="auto"),
         dict(moe_runner_backend="triton_kernel"),
-        dict(moe_runner_backend="flashinfer_cutlass"),
+        dict(moe_runner_backend="flashinfer_trtllm"),
+        dict(moe_runner_backend="flashinfer_trtllm_routed"),
+        dict(moe_runner_backend="flashinfer_mxfp4"),
         dict(dtype="float16"),
         dict(dtype="auto"),
         dict(disable_shared_experts_fusion=False),
@@ -106,6 +156,8 @@ def test_resolved_receiver_geometry_has_no_model_family_requirement(tp):
         dict(ep_num_redundant_experts=1),
         dict(init_expert_location="random"),
         dict(ep_join_mode="scale"),
+        dict(elastic_ep_backend="mooncake"),
+        dict(dwdp_size=2),
         dict(quantization="fp8"),
         dict(quantization_config={"quant_method": "fp8"}),
         dict(enable_lora=True),
@@ -120,6 +172,42 @@ def test_unsupported_receiver_layout_is_rejected(override):
         DirectExpertLayout.from_receiver(receiver(**override))
 
 
+@pytest.mark.parametrize("backend", ["triton", "deep_gemm", "flashinfer_cutlass"])
+@pytest.mark.parametrize("override", [dict(is_cuda=False), dict(is_cuda=None)])
+def test_direct_cuda_backends_reject_unknown_platform(backend, override):
+    with pytest.raises(ValueError):
+        DirectExpertLayout.from_receiver(receiver(moe_runner_backend=backend, **override))
+
+
+@pytest.mark.parametrize("backend", ["triton", "deep_gemm", "flashinfer_cutlass"])
+@pytest.mark.parametrize("a2a", ["none", "deepep", "flashinfer", "mooncake", "nixl", "mori", "pplx", "ascend_fuseep"])
+def test_weight_layout_is_independent_of_activation_dispatcher(backend, a2a):
+    # The running engine owns dispatcher/platform compatibility validation.
+    # Equal resident weight layouts must use the same weight transfer plan.
+    base = receiver(ep_size=2, moe_runner_backend=backend)
+    assert DirectExpertLayout.from_receiver(base | {"moe_a2a_backend": a2a}) == DirectExpertLayout.from_receiver(base)
+
+
+def test_unsharded_native_order_keeps_local_expert_storage():
+    layout = DirectExpertLayout.from_receiver(receiver(tp_size=4, ep_size=2, moe_dp_size=2))
+    tensor = torch.zeros(4, 16, 6, dtype=torch.bfloat16)
+    local = layout.shard(tensor, "gate_up_proj", 1)
+    assert layout.moe_tp_size == 1 and local.shape == (2, 16, 6)
+    assert local.data_ptr() == tensor[2:].data_ptr()
+    assert layout.shard(tensor, "gate_up_proj", 3).data_ptr() == local.data_ptr()
+
+
+def test_receiver_layout_agreement_includes_gate_up_order():
+    triton, deep_gemm, cutlass = [
+        DirectExpertLayout.from_receiver(receiver(moe_runner_backend=backend))
+        for backend in ("triton", "deep_gemm", "flashinfer_cutlass")
+    ]
+    assert triton == deep_gemm
+    assert cutlass != triton
+    assert triton != DirectExpertLayout.from_receiver(receiver(ep_size=2))
+    assert triton != DirectExpertLayout.from_receiver(receiver(moe_dp_size=2))
+
+
 @pytest.mark.parametrize(
     "packets,formats",
     [
diff --git a/tests/backends/sglang/test_sglang_engine.py b/tests/backends/sglang/test_sglang_engine.py
index 743423a20..b3ea378b6 100644
--- a/tests/backends/sglang/test_sglang_engine.py
+++ b/tests/backends/sglang/test_sglang_engine.py
@@ -1,12 +1,26 @@
 # Copyright (c) 2026 Relax Authors. All Rights Reserved.
 
-from types import SimpleNamespace
+import sys
+from dataclasses import make_dataclass
+from types import ModuleType, SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 
 from tests.backends.sglang.test_router_registration import sglang_engine_module as sglang_engine_module
 
 
+@pytest.fixture
+def expert_sync_engine_module(monkeypatch, sglang_engine_module):
+    import relax.utils
+
+    misc = ModuleType("relax.utils.misc")
+    misc.get_hf_config = Mock()
+    monkeypatch.setitem(sys.modules, misc.__name__, misc)
+    monkeypatch.setattr(relax.utils, "misc", misc, raising=False)
+    return sglang_engine_module
+
+
 @pytest.mark.parametrize(
     ("enable_mtp_training", "speculative_algorithm", "overrides", "expected"),
     [
@@ -32,12 +46,10 @@ def test_draft_weights_cpu_backup_follows_mtp_and_speculative_config(
     assert _enable_draft_weights_cpu_backup(args, overrides) is expected
 
 
-def test_expert_sync_layout_uses_runtime_overrides(monkeypatch, sglang_engine_module):
-    from unittest.mock import Mock
-
+def test_expert_sync_layout_uses_runtime_overrides(monkeypatch, expert_sync_engine_module):
     from relax.utils import misc
 
-    module = sglang_engine_module
+    module = expert_sync_engine_module
     engine = object.__new__(module.SGLangEngine)
     engine.node_rank = 0
     engine.worker_type = "regular"
@@ -66,10 +78,140 @@ def test_expert_sync_layout_uses_runtime_overrides(monkeypatch, sglang_engine_mo
 
 
 def test_expert_sync_version_commit_does_not_abort_again(monkeypatch, sglang_engine_module):
-    from unittest.mock import Mock
-
     engine = object.__new__(sglang_engine_module.SGLangEngine)
     request = Mock(return_value={"success": True})
     monkeypatch.setattr(engine, "_make_request", request)
     assert engine.update_weight_version("7") == {"success": True}
     request.assert_called_once_with("update_weight_version", {"new_version": "7", "abort_all_requests": False})
+
+
+@pytest.mark.parametrize(
+    ("runtime", "accelerator", "local_engine", "expected_backend"),
+    [
+        ({}, "cuda", True, "triton"),
+        ({"moe_runner_backend": "triton"}, "cuda", True, "triton"),
+        ({"moe_runner_backend": "flashinfer_trtllm"}, "cuda", True, "flashinfer_trtllm"),
+        ({"moe_runner_backend": "triton_kernel"}, "cuda", True, "triton_kernel"),
+        ({"moe_runner_backend": "deep_gemm"}, "cuda", True, "deep_gemm"),
+        ({"moe_runner_backend": "flashinfer_cutlass"}, "cuda", True, "flashinfer_cutlass"),
+        ({"moe_runner_backend": "deep_gemm"}, "rocm", True, "deep_gemm"),
+        ({"moe_a2a_backend": "deepep"}, "cuda", True, "triton"),
+        ({"device": "cpu"}, "cuda", True, "auto"),
+        ({}, "rocm", True, "auto"),
+        (None, "cuda", True, "auto"),
+        ({}, "cuda", False, "auto"),
+        ({"moe_runner_backend": "deep_gemm"}, "cuda", False, "deep_gemm"),
+        ({"moe_runner_backend": "flashinfer_cutlass"}, "cuda", False, "flashinfer_cutlass"),
+    ],
+)
+@pytest.mark.parametrize("dtype", ["bfloat16", "float16"])
+def test_expert_sync_auto_uses_scheduler_config_and_dtype_resolver(
+    monkeypatch, expert_sync_engine_module, runtime, accelerator, local_engine, expected_backend, dtype
+):
+    from relax.utils import misc
+
+    module = expert_sync_engine_module
+    engine = object.__new__(module.SGLangEngine)
+    engine.node_rank, engine.worker_type = 0, "regular"
+    engine.server_host, engine.server_port = "localhost", 1234
+    if local_engine:
+        engine.process = object()
+    startup = dict(
+        model_path="/test/model",
+        dtype="auto",
+        device="cuda",
+        moe_runner_backend="auto",
+        moe_a2a_backend="none",
+        disable_shared_experts_fusion=True,
+    )
+    states = [] if runtime is None else [startup | runtime]
+    response = Mock()
+    response.json.return_value = startup | {"internal_states": states}
+    monkeypatch.setattr(module.requests, "get", Mock(return_value=response))
+    text_config = SimpleNamespace(num_experts=256, dtype=dtype)
+    monkeypatch.setattr(misc, "get_hf_config", lambda _: SimpleNamespace(text_config=text_config))
+    model_config = ModuleType("sglang.srt.configs.model_config")
+    resolver = Mock(return_value=f"torch.{dtype}")
+    model_config._get_and_verify_dtype = resolver
+    monkeypatch.setitem(sys.modules, model_config.__name__, model_config)
+    monkeypatch.setattr(module.device_utils, "get_accelerator_type", lambda: accelerator, raising=False)
+    monkeypatch.setattr(module.device_utils, "AcceleratorType", SimpleNamespace(CUDA="cuda"), raising=False)
+
+    layout = engine.get_expert_weight_sync_layout()
+
+    assert layout["moe_runner_backend"] == expected_backend
+    assert layout["is_cuda"] is (
+        local_engine and accelerator == "cuda" and (runtime or {}).get("device", "cuda") == "cuda"
+    )
+    assert layout["dtype"] == dtype
+    assert layout["disable_shared_experts_fusion"] is True
+    resolver.assert_called_once_with(text_config, "auto")
+    assert module.requests.get.call_count == (2 if runtime is None else 1)
+    assert startup["dtype"] == startup["moe_runner_backend"] == "auto"
+
+
+@pytest.mark.parametrize(
+    "field,value",
+    [
+        ("dtype", "float16"),
+        ("disable_shared_experts_fusion", False),
+        ("ep_size", 2),
+        ("moe_dp_size", 2),
+        ("dwdp_size", 2),
+        ("elastic_ep_backend", "mooncake"),
+    ],
+)
+def test_expert_sync_rejects_inconsistent_scheduler_layouts(monkeypatch, expert_sync_engine_module, field, value):
+    module = expert_sync_engine_module
+    engine = object.__new__(module.SGLangEngine)
+    engine.node_rank, engine.worker_type = 0, "regular"
+    engine.server_host, engine.server_port = "localhost", 1234
+    startup = dict(model_path="/test/model", dtype="bfloat16", disable_shared_experts_fusion=True)
+    response = Mock()
+    response.json.return_value = startup | {"internal_states": [startup, startup | {field: value}]}
+    monkeypatch.setattr(module.requests, "get", Mock(return_value=response))
+
+    with pytest.raises(ValueError, match="schedulers disagree"):
+        engine.get_expert_weight_sync_layout()
+
+
+@pytest.mark.parametrize(
+    ("enforce", "overrides", "expected_disabled"),
+    [
+        (False, None, True),
+        (True, None, False),
+        (False, {"disable_shared_experts_fusion": False}, False),
+        (False, {"disable_shared_experts_fusion": True}, True),
+        (False, {"enforce_shared_experts_fusion": True}, False),
+    ],
+)
+def test_policy_shared_experts_default_respects_explicit_overrides(
+    monkeypatch, sglang_engine_module, enforce, overrides, expected_disabled
+):
+    module = sglang_engine_module
+    server_args = make_dataclass(
+        "ServerArgs",
+        [("disable_shared_experts_fusion", bool, False), ("enforce_shared_experts_fusion", bool, False)],
+    )
+    monkeypatch.setattr(module, "ServerArgs", server_args)
+    monkeypatch.setattr(module, "_to_local_gpu_id", lambda gpu_id: gpu_id)
+    args = SimpleNamespace(
+        rollout_num_gpus_per_engine=2,
+        num_gpus_per_node=8,
+        hf_checkpoint="/test/model",
+        seed=42,
+        offload_rollout=True,
+        sglang_pp_size=1,
+        sglang_dp_size=1,
+        sglang_ep_size=1,
+        use_rollout_routing_replay=False,
+        fp16=False,
+        sglang_disable_shared_experts_fusion=False,
+        sglang_enforce_shared_experts_fusion=enforce,
+    )
+
+    result, _ = module._compute_server_args(
+        args, 0, "localhost:8000", 8001, "localhost", 8002, base_gpu_id=0, sglang_overrides=overrides
+    )
+
+    assert result["disable_shared_experts_fusion"] is expected_disabled

From a768c65d50698bbddad3bf79861a385d8f49e624 Mon Sep 17 00:00:00 2001
From: yuanlehome <yuanlehome@163.com>
Date: Wed, 21 Oct 2026 14:04:32 +0800
Subject: [PATCH 5/6] fix(megatron): handle optional inline logprob state

Skip inline eligibility when MTP-only SFT has no weight backuper, and
use Relax's partial_rollout state without requiring optional Megatron RL
arguments. Cover ordinary and prepacked SFT entry points in regression tests.

Validation: 101 related tests, full pre-commit, three Qwen3.5 RL rounds
at batch512/packing8192 with sgl-router, and two MTP-only SFT steps.
---
 relax/backends/megatron/actor.py              |  2 +-
 relax/utils/training/inline_log_probs.py      |  1 -
 .../test_colocate_inline_log_probs.py         |  1 -
 .../megatron/test_mtp_only_training.py        | 51 +++++++++++++++++++
 4 files changed, 52 insertions(+), 3 deletions(-)

diff --git a/relax/backends/megatron/actor.py b/relax/backends/megatron/actor.py
index f36a42d25..73eb3bf4a 100644
--- a/relax/backends/megatron/actor.py
+++ b/relax/backends/megatron/actor.py
@@ -1958,7 +1958,7 @@ def train_actor(
             data_iterator, num_microbatches = get_data_iterator(self.args, self.model, rollout_data)
         else:
             data_iterator, num_microbatches = prepared_data_iterator, prepared_num_microbatches
-        inline_old_log_probs = can_reuse_colocate_log_probs(
+        inline_old_log_probs = hasattr(self, "weights_backuper") and can_reuse_colocate_log_probs(
             self.args,
             num_microbatches,
             backup_tags=self.weights_backuper.backup_tags,
diff --git a/relax/utils/training/inline_log_probs.py b/relax/utils/training/inline_log_probs.py
index 2f91bf8ac..98785833e 100644
--- a/relax/utils/training/inline_log_probs.py
+++ b/relax/utils/training/inline_log_probs.py
@@ -29,7 +29,6 @@ def can_reuse_colocate_log_probs(
         "max_staleness": 0,
         "update_weights_interval": 1,
         "partial_rollout": False,
-        "rl_partial_rollouts": False,
         "use_dynamic_global_batch_size": False,
         "context_parallel_size": 1,
         "pipeline_model_parallel_size": 1,
diff --git a/tests/backends/megatron/test_colocate_inline_log_probs.py b/tests/backends/megatron/test_colocate_inline_log_probs.py
index bf0c98a1d..01903cecd 100644
--- a/tests/backends/megatron/test_colocate_inline_log_probs.py
+++ b/tests/backends/megatron/test_colocate_inline_log_probs.py
@@ -26,7 +26,6 @@ def actor():
         max_staleness=0,
         update_weights_interval=1,
         partial_rollout=False,
-        rl_partial_rollouts=False,
         use_dynamic_global_batch_size=False,
         context_parallel_size=1,
         pipeline_model_parallel_size=1,
diff --git a/tests/backends/megatron/test_mtp_only_training.py b/tests/backends/megatron/test_mtp_only_training.py
index 3c6c2bfdb..ec9195b6b 100644
--- a/tests/backends/megatron/test_mtp_only_training.py
+++ b/tests/backends/megatron/test_mtp_only_training.py
@@ -1,7 +1,10 @@
 # Copyright (c) 2026 Relax Authors. All Rights Reserved.
 
 from argparse import Namespace
+from contextlib import nullcontext
 from pathlib import Path
+from types import SimpleNamespace
+from unittest.mock import Mock
 
 import pytest
 import torch
@@ -190,6 +193,54 @@ def test_pure_mtp_only_sft_skips_weight_management():
     assert should_skip_mtp_only_weight_management(_actor_args()) is True
 
 
+@pytest.mark.parametrize("prepacked", [False, True])
+def test_mtp_only_sft_reaches_train_without_weight_backuper(monkeypatch, prepacked):
+    actor_module = pytest.importorskip("relax.backends.megatron.actor", exc_type=ImportError)
+
+    class ReachedTrain(BaseException):
+        pass
+
+    actor = object.__new__(actor_module.MegatronTrainRayActor)
+    actor.args = _actor_args(
+        advantage_estimator="grpo",
+        compute_advantages_and_returns=False,
+        use_routing_replay=False,
+        use_rollout_routing_replay=False,
+    )
+    actor.model = []
+    actor.optimizer = object()
+    actor.opt_param_scheduler = object()
+    actor.rollout_data_postprocess = None
+    assert should_skip_mtp_only_weight_management(actor.args)
+    assert not hasattr(actor, "weights_backuper")
+
+    iterator = SimpleNamespace(close=Mock())
+    data_iterator, num_microbatches = [iterator], [1]
+    make_iterator = Mock(return_value=(data_iterator, num_microbatches))
+    train = Mock(side_effect=ReachedTrain)
+    monkeypatch.setattr(actor_module, "get_data_iterator", make_iterator)
+    monkeypatch.setattr(actor_module, "timer", lambda *_args, **_kwargs: nullcontext())
+    monkeypatch.setattr(actor_module, "inverse_timer", lambda *_args, **_kwargs: nullcontext())
+    monkeypatch.setattr(actor_module, "log_rollout_data", lambda *_args: None)
+    monkeypatch.setattr(actor_module, "train", train)
+
+    kwargs = (
+        {"prepared_data_iterator": data_iterator, "prepared_num_microbatches": num_microbatches} if prepacked else {}
+    )
+    rollout_data = {}
+    with pytest.raises(ReachedTrain):
+        actor.train_actor(0, rollout_data, **kwargs)
+
+    train.assert_called_once_with(
+        0, actor.model, actor.optimizer, actor.opt_param_scheduler, data_iterator, num_microbatches
+    )
+    iterator.close.assert_called_once()
+    if prepacked:
+        make_iterator.assert_not_called()
+    else:
+        make_iterator.assert_called_once_with(actor.args, actor.model, rollout_data)
+
+
 @pytest.mark.parametrize(
     ("overrides", "extra_kwargs"),
     [

From bf3cc7c66a9b9187a7192afc533352f43a06f919 Mon Sep 17 00:00:00 2001
From: yuanlehome <yuanlehome@163.com>
Date: Wed, 21 Oct 2026 21:17:01 +0800
Subject: [PATCH 6/6] fix(megatron): preserve logits when loss mask is absent

Return Kimi CP logits directly when no loss mask is supplied, preserving
the output-and-mask tuple for callers that provide a mask.

Validated fixed-upstream patch application and reversal, 36 CPU return
contract cases, two existing patch tests, and all-files pre-commit.
---
 docker/patch/megatron/20260915-4a4de8657.patch | 11 ++++++++++-
 1 file changed, 10 insertions(+), 1 deletion(-)

diff --git a/docker/patch/megatron/20260915-4a4de8657.patch b/docker/patch/megatron/20260915-4a4de8657.patch
index c8f3e42bb..fc00d5ab7 100644
--- a/docker/patch/megatron/20260915-4a4de8657.patch
+++ b/docker/patch/megatron/20260915-4a4de8657.patch
@@ -112,7 +112,7 @@ index ba632f3..99e8aa9 100644
      freeze_language_model: bool = False
      freeze_vision_model: bool = False
 diff --git a/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py b/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py
-index 4ef12d7..01c5180 100644
+index 4ef12d7..7bf9ad0 100644
 --- a/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py
 +++ b/megatron/bridge/models/kimi_vl/modeling_kimi_k25_vl.py
 @@ -16,7 +16,9 @@ import logging
@@ -275,6 +275,15 @@ index 4ef12d7..01c5180 100644
  
              if self.config.sequence_parallel:
                  tp_group = self.config._pg_collection.tp if self.config._pg_collection is not None else None
+@@ -532,7 +622,7 @@ class KimiK25VLModel(MegatronModule):
+             runtime_gather_output=runtime_gather_output,
+             packed_seq_params=packed_seq_params,
+         )
+-        if cp_size > 1:
++        if cp_size > 1 and loss_mask is not None:
+             return outputs, loss_mask
+         return outputs
+ 
 diff --git a/megatron/bridge/models/qwen/qwen3_moe_bridge.py b/megatron/bridge/models/qwen/qwen3_moe_bridge.py
 index a6476f4..cab29da 100644
 --- a/megatron/bridge/models/qwen/qwen3_moe_bridge.py