From c8676d6983ebf4ac3bb0b64499e3dec37e563e30 Mon Sep 17 00:00:00 2001
From: isLinXu <2267379130@qq.com>
Date: Sat, 1 Aug 2026 13:01:48 +0800
Subject: [PATCH 01/11] fix: stabilize MoA layouts and MPS numerics

---
 tests/test_mixture_numeric.py       | 36 +++++++++++++++
 tests/test_moa.py                   | 43 ++++++++++++++----
 ultralytics/nn/modules/_numeric.py  | 18 +++++++-
 ultralytics/nn/modules/moa/heads.py | 69 ++++++++++++++++++++---------
 4 files changed, 134 insertions(+), 32 deletions(-)

diff --git a/tests/test_mixture_numeric.py b/tests/test_mixture_numeric.py
index 3f8bba10..e1900a03 100644
--- a/tests/test_mixture_numeric.py
+++ b/tests/test_mixture_numeric.py
@@ -5,9 +5,11 @@
 import pytest
 import torch
 
+from ultralytics.nn.modules import LatentRouter
 from ultralytics.nn.modules._numeric import (
     all_reduce_mean,
     clamp_min_for_dtype,
+    disabled_autocast,
     fp_clamp_floor,
     stable_normalize,
 )
@@ -53,6 +55,27 @@ def test_all_mixture_namespaces_share_canonical_all_reduce_mean():
     assert mot_all_reduce_mean is all_reduce_mean
 
 
+def test_disabled_autocast_falls_back_when_device_is_unsupported():
+    with patch("torch.amp.autocast_mode.is_autocast_available", return_value=False), patch(
+        "torch.autocast", side_effect=AssertionError("unsupported autocast must not be constructed")
+    ):
+        with disabled_autocast("mps"):
+            result = torch.ones(1) + 1
+
+    assert result.item() == 2
+
+
+def test_disabled_autocast_uses_supported_cpu_context():
+    with patch("torch.amp.autocast_mode.is_autocast_available", return_value=True), patch(
+        "torch.autocast", wraps=torch.autocast
+    ) as autocast:
+        with disabled_autocast("cpu"):
+            result = torch.ones(1) + 1
+
+    autocast.assert_called_once_with(device_type="cpu", enabled=False)
+    assert result.item() == 2
+
+
 def test_all_reduce_mean_keeps_global_value_and_local_gradient():
     local = torch.tensor([0.2, 0.8], requires_grad=True)
 
@@ -102,6 +125,19 @@ def test_mot_router_keeps_fp32_logits_and_activation_dtype_weights():
     assert torch.allclose(weights.float().sum(dim=1), torch.ones_like(weights[:, 0].float()), atol=2e-3)
 
 
+def test_latent_router_keeps_fp32_routing_contract_after_dtype_conversion():
+    router = LatentRouter(8, num_experts=3).eval().half()
+    tokens = torch.randn(2, 8, dtype=torch.bfloat16)
+
+    logits, probs = router(tokens)
+
+    assert logits.dtype == torch.float32
+    assert probs.dtype == torch.float32
+    assert {parameter.dtype for parameter in router.parameters()} == {torch.float32}
+    assert torch.isfinite(probs).all()
+    assert torch.allclose(probs.sum(dim=-1), torch.ones_like(probs[:, 0]), atol=1e-6)
+
+
 @pytest.mark.parametrize("router_cls", [DualStreamGateRouter, ZeroCostRouter])
 def test_gated_router_keeps_fp32_aux_statistics_and_activation_dtype_weights(router_cls):
     router = router_cls(8, num_experts=3, top_k=2).train().half()
diff --git a/tests/test_moa.py b/tests/test_moa.py
index f147b82f..1a03270f 100644
--- a/tests/test_moa.py
+++ b/tests/test_moa.py
@@ -16,6 +16,7 @@
     _moa_router_aux_loss,
     _window_flash_attn,
 )
+from ultralytics.nn.modules.moa.heads import _window_partition_2d, _window_unpartition_2d
 from ultralytics.nn.modules._numeric import fp_clamp_floor as _fp_min
 from ultralytics.nn.tasks import DetectionModel
 from ultralytics.utils.loss import _collect_moa_aux_loss
@@ -412,22 +413,46 @@ def test_regional_attn_head_non_divisible_dim_and_heads():
 
 
 def test_regional_attn_head_small_spatial_dims():
-    """_RegionalAttnHead gracefully handles H=1 or W=1 feature maps."""
+    """_RegionalAttnHead keeps canonical memory format and gradients for singleton axes."""
     torch.manual_seed(0)
-    cases = [
-        (_RegionalAttnHead(32, num_heads=4), torch.randn(1, 32, 1, 16)),
-        (_RegionalAttnHead(32, num_heads=4), torch.randn(1, 32, 16, 1)),
-        (_RegionalAttnHead(32, num_heads=4), torch.randn(1, 32, 1, 1)),
-    ]
-    for module, x in cases:
-        module.train()
+    for height, width in ((1, 16), (16, 1), (1, 1)):
+        module = _RegionalAttnHead(32, num_heads=4).train()
+        x = torch.randn(2, 32, height, width, requires_grad=True)
         out = module(x)
         assert out.shape == x.shape
+        assert out.is_contiguous()
         assert torch.isfinite(out).all()
-        out.mean().backward()
+        out.square().mean().backward()
+        assert x.grad is not None and torch.isfinite(x.grad).all()
         assert _has_grad(module)
 
 
+def test_window_partition_unpartition_round_trip():
+    """Window helpers preserve row-major [H,W] layout, including multiple heads."""
+    spatial = torch.arange(2 * 3 * 4 * 6 * 5).reshape(2, 3, 4, 6, 5)
+    windows = _window_partition_2d(spatial, 2)
+    restored = _window_unpartition_2d(windows, 2, 2, 3, 4, 6)
+    assert torch.equal(restored, spatial)
+
+
+def test_window_flash_attn_matches_spatial_reference():
+    """Window attention matches an explicit per-window numerical reference."""
+    torch.manual_seed(7)
+    B, nh, H, W, hd, win = 2, 2, 4, 6, 3, 2
+    q, k, v = (torch.randn(B, nh, H * W, hd) for _ in range(3))
+    scale = hd ** -0.5
+
+    actual = _window_flash_attn(q, k, v, scale, win, H, W)
+    expected = torch.empty_like(actual.reshape(B, nh, H, W, hd))
+    for row in range(0, H, win):
+        for col in range(0, W, win):
+            indices = [(r * W + c) for r in range(row, row + win) for c in range(col, col + win)]
+            reference = _flash_attn(q[:, :, indices], k[:, :, indices], v[:, :, indices], scale)
+            expected[:, :, row : row + win, col : col + win] = reference.reshape(B, nh, win, win, hd)
+
+    assert torch.allclose(actual, expected.reshape(B, nh, H * W, hd), atol=1e-6, rtol=1e-5)
+
+
 def test_regional_attn_head_invalid_pool_stride():
     """_RegionalAttnHead raises ValueError for pool_stride < 1."""
     with pytest.raises(ValueError, match="pool_stride"):
diff --git a/ultralytics/nn/modules/_numeric.py b/ultralytics/nn/modules/_numeric.py
index df2269cd..9d83eb17 100644
--- a/ultralytics/nn/modules/_numeric.py
+++ b/ultralytics/nn/modules/_numeric.py
@@ -9,9 +9,23 @@
 import torch.nn as nn
 
 
+def _autocast_is_available(device_type: str) -> bool:
+    """Return whether this PyTorch build exposes autocast for ``device_type``."""
+    checker = getattr(getattr(torch, "amp", None), "autocast_mode", None)
+    checker = getattr(checker, "is_autocast_available", None)
+    if checker is not None:
+        try:
+            return bool(checker(device_type))
+        except (RuntimeError, TypeError):
+            return False
+    # PyTorch 2.2 does not expose the capability query and has no MPS
+    # autocast implementation. CPU and CUDA autocast are supported there.
+    return device_type in {"cpu", "cuda"}
+
+
 def disabled_autocast(device_type: str):
-    """Return an autocast-disabled context for router-critical numerical work."""
-    if device_type in {"cpu", "cuda", "mps"}:
+    """Disable autocast when supported, otherwise return a no-op context."""
+    if _autocast_is_available(device_type):
         return torch.autocast(device_type=device_type, enabled=False)
     return nullcontext()
 
diff --git a/ultralytics/nn/modules/moa/heads.py b/ultralytics/nn/modules/moa/heads.py
index c0879f48..7a607cbd 100644
--- a/ultralytics/nn/modules/moa/heads.py
+++ b/ultralytics/nn/modules/moa/heads.py
@@ -43,6 +43,33 @@ def _flash_attn(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor,
     attn = attn.softmax(dim=-1)
     return attn @ v
 
+def _window_partition_2d(t: torch.Tensor, window_size: int) -> torch.Tensor:
+    """Partition ``[B, nh, H, W, hd]`` into ``[B*nh*nW, win², hd]`` windows."""
+    B, nh, height, width, hd = t.shape
+    win = int(window_size)
+    if win < 1 or height % win or width % win:
+        raise ValueError(f"window_size={win} must evenly divide spatial shape {(height, width)}")
+    t = t.reshape(B, nh, height // win, win, width // win, win, hd)
+    return t.permute(0, 1, 2, 4, 3, 5, 6).reshape(-1, win * win, hd)
+
+
+def _window_unpartition_2d(
+    windows: torch.Tensor,
+    window_size: int,
+    batch_size: int,
+    num_heads: int,
+    height: int,
+    width: int,
+) -> torch.Tensor:
+    """Reverse :func:`_window_partition_2d` to ``[B, nh, H, W, hd]``."""
+    win = int(window_size)
+    if win < 1 or height % win or width % win:
+        raise ValueError(f"window_size={win} must evenly divide spatial shape {(height, width)}")
+    hd = windows.shape[-1]
+    t = windows.reshape(batch_size, num_heads, height // win, width // win, win, win, hd)
+    return t.permute(0, 1, 2, 4, 3, 5, 6).reshape(batch_size, num_heads, height, width, hd)
+
+
 def _window_flash_attn(
     q: torch.Tensor,
     k: torch.Tensor,
@@ -60,10 +87,11 @@ def _window_flash_attn(
         )
     win = max(1, min(int(window_size), height, width))
 
-    def to_spatial(t: torch.Tensor) -> torch.Tensor:
-        return t.transpose(2, 3).reshape(B, nh, height, width, hd)
-
-    qs, ks, vs = to_spatial(q), to_spatial(k), to_spatial(v)
+    # N is already the flattened row-major spatial axis. Transposing N and hd
+    # before reshape silently mixes feature and spatial coordinates.
+    qs = q.reshape(B, nh, height, width, hd)
+    ks = k.reshape(B, nh, height, width, hd)
+    vs = v.reshape(B, nh, height, width, hd)
     pad_h = (win - height % win) % win
     pad_w = (win - width % win) % win
     if pad_h or pad_w:
@@ -71,18 +99,14 @@ def to_spatial(t: torch.Tensor) -> torch.Tensor:
         qs, ks, vs = F.pad(qs, pad), F.pad(ks, pad), F.pad(vs, pad)
     hp, wp = qs.shape[2], qs.shape[3]
 
-    def partition(t: torch.Tensor) -> torch.Tensor:
-        t = t.view(B, nh, hp // win, win, wp // win, win, hd)
-        return t.permute(0, 1, 2, 4, 3, 5, 6).reshape(-1, win * win, hd)
-
-    def reverse(windows: torch.Tensor) -> torch.Tensor:
-        n_h, n_w = hp // win, wp // win
-        t = windows.view(B, nh, n_h, n_w, win, win, hd)
-        t = t.permute(0, 1, 2, 4, 3, 5, 6).reshape(B, nh, hp, wp, hd)
-        return t[:, :, :height, :width, :].reshape(B, nh, height * width, hd)
-
-    out_w = _flash_attn(partition(qs), partition(ks), partition(vs), scale)
-    return reverse(out_w)
+    out_w = _flash_attn(
+        _window_partition_2d(qs, win),
+        _window_partition_2d(ks, win),
+        _window_partition_2d(vs, win),
+        scale,
+    )
+    out = _window_unpartition_2d(out_w, win, B, nh, hp, wp)
+    return out[:, :, :height, :width, :].reshape(B, nh, height * width, hd)
 
 class _LocalAttnHead(nn.Module):
     """Local attention head: DW-biased QKV + window-partitioned self-attention.
@@ -194,17 +218,20 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         # Guard: if pooling collapsed the spatial dim to zero (extreme edge case),
         # fall back to identity KV.
         if H2 * W2 == 0:
-            k = self.q_proj(x).reshape(B, nh, hd, -1).transpose(2, 3)
+            k = self.q_proj(x).reshape(B, nh, hd, -1).transpose(2, 3).contiguous()
             v = k.clone()
         else:
             k, v = kv.split(inner, dim=1)
-            k = k.flatten(2).view(B, nh, hd, H2 * W2).transpose(2, 3)
-            v = v.flatten(2).view(B, nh, hd, H2 * W2).transpose(2, 3)
+            k = k.flatten(2).view(B, nh, hd, H2 * W2).transpose(2, 3).contiguous()
+            v = v.flatten(2).view(B, nh, hd, H2 * W2).transpose(2, 3).contiguous()
 
-        q = self.q_proj(x).flatten(2).view(B, nh, hd, H * W).transpose(2, 3)
+        # Materialize canonical [B, nh, N, hd] strides. Singleton spatial axes
+        # otherwise retain ambiguous channels-first strides that can make SDPA's
+        # backward select an incompatible memory-format path.
+        q = self.q_proj(x).flatten(2).view(B, nh, hd, H * W).transpose(2, 3).contiguous()
 
         out = _flash_attn(q, k, v, self.scale)               # [B, nh, N, hd]
-        out = out.transpose(2, 3).reshape(B, inner, H, W)
+        out = out.transpose(2, 3).contiguous().reshape(B, inner, H, W)
         return self.norm(self.proj(out))
 
 class _GlobalAttnHead(nn.Module):

From 9e15001db18b393942e2619bb61e8b70a6be2c26 Mon Sep 17 00:00:00 2001
From: isLinXu <2267379130@qq.com>
Date: Sat, 1 Aug 2026 21:37:01 +0800
Subject: [PATCH 02/11] fix: harden mixture routing correctness

---
 tests/test_ddp_lifecycle_ema_nan.py  | 20 ++++++++++++++++++++
 tests/test_mot_scene_aware_router.py | 25 +++++++++++++++++++++++++
 ultralytics/engine/trainer.py        |  7 ++++++-
 ultralytics/nn/modules/moa/router.py |  7 ++++---
 ultralytics/nn/modules/mot/block.py  | 16 ++++++++++++++++
 ultralytics/nn/modules/mot/router.py |  3 ++-
 6 files changed, 73 insertions(+), 5 deletions(-)

diff --git a/tests/test_ddp_lifecycle_ema_nan.py b/tests/test_ddp_lifecycle_ema_nan.py
index bbef7ca9..52337f4c 100644
--- a/tests/test_ddp_lifecycle_ema_nan.py
+++ b/tests/test_ddp_lifecycle_ema_nan.py
@@ -440,6 +440,26 @@ def test_checkpoint_restore_tolerates_missing_lazy_ema_buffer():
     )
 
 
+def test_checkpoint_restore_migrates_legacy_three_slot_ema_state():
+    t = object.__new__(BaseTrainer)
+    t.model = nn.Linear(1, 1)
+    t.model.register_buffer("_mixture_loss_ema_buf", torch.tensor([1.0, 0.1, 0.1, 0.1]))
+    t.ema = ModelEMA(t.model)
+    t.optimizer = torch.optim.SGD(t.model.parameters(), lr=0.01)
+    t.scaler = torch.amp.GradScaler("cuda", enabled=False)
+    old_ema = nn.Linear(1, 1)
+    old_ema.register_buffer("_mixture_loss_ema_buf", torch.tensor([2.0, 0.2, 0.3]))
+
+    t._load_checkpoint_state(
+        {"ema": old_ema, "optimizer": None, "scaler": None, "best_fitness": 0.0, "updates": 7}
+    )
+
+    expected = torch.tensor([2.0, 0.2, 0.3, 0.1])
+    assert torch.equal(t.model._mixture_loss_ema_buf, expected)
+    assert torch.equal(t.ema.ema._mixture_loss_ema_buf, expected)
+    assert t.ema.updates == 7
+
+
 def test_healthy_checkpoint_rejects_nonfinite_state_and_preserves_prior(tmp_path):
     t = object.__new__(BaseTrainer)
     t.healthy = tmp_path / "last_healthy.pt"
diff --git a/tests/test_mot_scene_aware_router.py b/tests/test_mot_scene_aware_router.py
index d2b699a2..483f7ef6 100644
--- a/tests/test_mot_scene_aware_router.py
+++ b/tests/test_mot_scene_aware_router.py
@@ -157,6 +157,19 @@ def test_scene_consistency_component_reaches_scene_projector():
     assert any(parameter.grad is not None for parameter in block.router.scene_projector.parameters())
 
 
+def test_mot_snapshot_reports_exact_runtime_sparse_metrics():
+    block = MoTBlock(24, num_heads=3, top_k=1).eval()
+    with torch.no_grad():
+        block.router.router[-1].bias.copy_(torch.tensor([3.0, 0.0, -3.0]))
+        block(torch.randn(2, 24, 4, 4))
+    dispatch = block.routing_snapshot()["dispatch"]
+
+    assert dispatch["token_mask_sparsity"] == pytest.approx(2.0 / 3.0)
+    assert dispatch["experts_per_sample"].tolist() == [1, 1]
+    assert dispatch["batch_expert_union"] == 1
+    assert dispatch["actual_expert_calls"] == 1
+
+
 def test_scene_aware_master_config_parses_and_runs():
     config = ROOT / "ultralytics/cfg/models/master/v0_10/det/yolo-master-mot-scene-n.yaml"
     model = DetectionModel(str(config), ch=3, nc=80, verbose=False).eval()
@@ -178,3 +191,15 @@ def test_scene_aware_yaml_survives_runtime_default_resolution():
     audit = [item for item in resolved.audit if item["kind"] == "mot"]
     assert audit and all(item["sources"]["scene_aware_router"] == "yaml" for item in audit)
     assert all(block.router.scene_aware for block in model.modules() if isinstance(block, MoTBlock))
+
+
+def test_scene_aware_enable_after_router_move_preserves_device_and_fp32_contract():
+    router = _MoTRouter(8, scene_aware=False).to(dtype=torch.float64)
+
+    router.enable_scene_aware()
+
+    reference = next(router.router.parameters())
+    assert {parameter.device for parameter in router.scene_projector.parameters()} == {reference.device}
+    assert {parameter.dtype for parameter in router.scene_projector.parameters()} == {torch.float32}
+    weights, _ = router(torch.randn(1, 8, 4, 4, dtype=torch.float64))
+    assert torch.isfinite(weights).all()
diff --git a/ultralytics/engine/trainer.py b/ultralytics/engine/trainer.py
index 9e3309ab..72ceace9 100644
--- a/ultralytics/engine/trainer.py
+++ b/ultralytics/engine/trainer.py
@@ -1332,7 +1332,12 @@ def _load_checkpoint_state(self, ckpt):
 
             online_target = unwrap_model(self.model)
             online_mixture_ema = initialize_mixture_loss_ema_buffer(online_target)
-            ema_state = ckpt["ema"].float().state_dict()
+            checkpoint_ema = ckpt["ema"].float()
+            # Apply the same lazy initializer/migration used by live models before
+            # extracting state, so legacy three-slot checkpoints load into the
+            # current four-slot online and EMA buffers without a size mismatch.
+            initialize_mixture_loss_ema_buffer(checkpoint_ema)
+            ema_state = checkpoint_ema.state_dict()
             checkpoint_mixture_ema = ema_state.get("_mixture_loss_ema_buf")
             if checkpoint_mixture_ema is not None:
                 online_mixture_ema.copy_(
diff --git a/ultralytics/nn/modules/moa/router.py b/ultralytics/nn/modules/moa/router.py
index 1320bbdd..c98396f4 100644
--- a/ultralytics/nn/modules/moa/router.py
+++ b/ultralytics/nn/modules/moa/router.py
@@ -75,10 +75,11 @@ def _moa_router_aux_loss(
     if reduce_ddp:
         global_sum = _all_reduce_mean(local_sum.detach().clone())
         global_count = _all_reduce_mean(local_count.detach().clone())
-        # DDP averages parameter gradients by world size. Scale the local Jacobian
-        # by R/N while exposing the exact detached global value S/N.
+        # ``_all_reduce_mean`` makes ``global_count`` equal N/R. DDP later averages
+        # parameter gradients by R, so 1/global_count already gives the required
+        # local Jacobian R/N; multiplying by world size again would over-scale it.
         importance = global_sum / global_count.clamp_min(1.0)
-        local_grad = (local_sum - local_sum.detach()) * (dist.get_world_size() / global_count.clamp_min(1.0))
+        local_grad = (local_sum - local_sum.detach()) / global_count.clamp_min(1.0)
         importance = importance + local_grad
     else:
         importance = local_sum / local_count.clamp_min(1.0)
diff --git a/ultralytics/nn/modules/mot/block.py b/ultralytics/nn/modules/mot/block.py
index 46b5abfe..e9226393 100644
--- a/ultralytics/nn/modules/mot/block.py
+++ b/ultralytics/nn/modules/mot/block.py
@@ -316,6 +316,18 @@ def _blend_experts(
         use_sparse = (not self.training or (sparse_train_ready and ddp_sparse_safe)) and not exporting
         warmup_step = 0 if exporting else int(self._sparse_train_step.item())
         B = x.shape[0]
+        route_ids = indices if indices is not None else weights.argmax(dim=1, keepdim=True)
+        route_mask = torch.zeros_like(weights, dtype=torch.bool)
+        route_mask.scatter_(1, route_ids, True)
+        token_mask_sparsity = 1.0 - float(route_mask.float().mean())
+        experts_per_sample = route_mask.reshape(B, self.NUM_EXPERTS, -1).any(dim=2).sum(dim=1)
+        batch_expert_union = int(route_mask.any(dim=(0, 2, 3)).sum())
+        routing_metrics = {
+            "token_mask_sparsity": token_mask_sparsity,
+            "experts_per_sample": experts_per_sample.detach().cpu(),
+            "mean_experts_per_sample": float(experts_per_sample.float().mean()),
+            "batch_expert_union": batch_expert_union,
+        }
         if use_sparse:
             expert_calls = 0
             for e_idx, expert in enumerate(self.experts):
@@ -340,6 +352,8 @@ def _blend_experts(
             self._last_dispatch_stats = {
                 "mode": "sample_sparse",
                 "expert_calls": expert_calls,
+                "actual_expert_calls": expert_calls,
+                **routing_metrics,
                 "selected_samples": B,
                 "selected_experts": selected_experts,
                 "sparsity_ratio": 1.0 - expert_calls / max(len(self.experts), 1),
@@ -367,6 +381,8 @@ def _blend_experts(
             self._last_dispatch_stats = {
                 "mode": "dense",
                 "expert_calls": len(self.experts),
+                "actual_expert_calls": len(self.experts),
+                **routing_metrics,
                 "selected_samples": B,
                 "selected_experts": len(self.experts),
                 "sparsity_ratio": 0.0,
diff --git a/ultralytics/nn/modules/mot/router.py b/ultralytics/nn/modules/mot/router.py
index 0fca3e8d..017ad8c3 100644
--- a/ultralytics/nn/modules/mot/router.py
+++ b/ultralytics/nn/modules/mot/router.py
@@ -153,11 +153,12 @@ def enable_scene_aware(self, hidden_dim: Optional[int] = None) -> None:
             hidden = int(hidden_dim or self.scene_hidden_dim or 3)
             if hidden <= 0:
                 raise ValueError("scene_hidden_dim must be positive")
+            reference = next(self.router.parameters())
             self.scene_projector = nn.Sequential(
                 nn.Linear(3, hidden),
                 nn.SiLU(inplace=False),
                 nn.Linear(hidden, self.num_experts),
-            )
+            ).to(device=reference.device, dtype=torch.float32)
             nn.init.zeros_(self.scene_projector[-1].weight)
             nn.init.zeros_(self.scene_projector[-1].bias)
             self.scene_hidden_dim = hidden

From d5faaf649570af87710f0f092696500d834d7363 Mon Sep 17 00:00:00 2001
From: isLinXu <2267379130@qq.com>
Date: Sat, 1 Aug 2026 21:37:31 +0800
Subject: [PATCH 03/11] fix: clarify sparse routing diagnostics

---
 scripts/bench_moe_mps.py                  | 21 ++++++++-
 tests/test_routing_diagnostics.py         | 54 +++++++++++++++++++++++
 ultralytics/nn/modules/moe/diagnostics.py | 29 +++++-------
 ultralytics/nn/modules/moe/schedule.py    | 24 +---------
 ultralytics/nn/modules/moe/scheduler.py   | 18 ++------
 5 files changed, 88 insertions(+), 58 deletions(-)

diff --git a/scripts/bench_moe_mps.py b/scripts/bench_moe_mps.py
index b6ad3725..118aa0f0 100644
--- a/scripts/bench_moe_mps.py
+++ b/scripts/bench_moe_mps.py
@@ -34,6 +34,18 @@
 ]
 
 
+def mps_device_info():
+    """Return optional MPS metadata without assuming backend APIs exist."""
+    backend = getattr(torch.backends, "mps", None)
+    if backend is None:
+        return None, None
+    get_name = getattr(backend, "get_name", None)
+    get_core_count = getattr(backend, "get_core_count", None)
+    name = get_name() if callable(get_name) else None
+    core_count = get_core_count() if callable(get_core_count) else None
+    return name, core_count
+
+
 def pick_device():
     if torch.backends.mps.is_available():
         return torch.device("mps")
@@ -84,8 +96,13 @@ def main():
     dev = pick_device()
     print(f"[bench] device = {dev}, torch = {torch.__version__}")
     if dev.type == "mps":
-        print(f"[bench] mps name = {torch.backends.mps.get_name()}, "
-              f"core count = {torch.backends.mps.get_core_count()}")
+        name, core_count = mps_device_info()
+        details = []
+        if name is not None:
+            details.append(f"name = {name}")
+        if core_count is not None:
+            details.append(f"core count = {core_count}")
+        print(f"[bench] mps {', '.join(details) if details else 'metadata unavailable'}")
     print(f"[bench] bs={args.bs} imgsz={args.imgsz} runs={args.runs} warmup={args.warmup} dtype={args.dtype}")
     print("-" * 90)
     print(f"{'config':<30}{'params(M)':>11}{'median':>9}{'mean':>9}{'min':>9}{'p95':>9}{'fps':>8}")
diff --git a/tests/test_routing_diagnostics.py b/tests/test_routing_diagnostics.py
index f8930757..58ac1211 100644
--- a/tests/test_routing_diagnostics.py
+++ b/tests/test_routing_diagnostics.py
@@ -25,6 +25,36 @@ def test_nonfinite_moa_aux_preserves_finite_graph_and_reports_boundary():
     assert torch.isfinite(weights.grad).all()
 
 
+def test_moa_aux_ddp_local_jacobian_is_world_size_invariant(monkeypatch):
+    """Averaged mock-DDP gradients must match the single-process global objective."""
+    rank_inputs = [torch.tensor([0.2, -0.1]), torch.tensor([0.5, 0.3])]
+    world_size = len(rank_inputs)
+    global_weights = torch.cat([torch.softmax(values, dim=0).view(1, 2, 1, 1) for values in rank_inputs], dim=0)
+    global_logits = torch.zeros_like(global_weights)
+    reference_values = torch.cat(rank_inputs).requires_grad_()
+    reference_weights = torch.stack(
+        [torch.softmax(reference_values[i : i + 2], dim=0) for i in range(0, reference_values.numel(), 2)]
+    ).view(world_size, 2, 1, 1)
+    reference = _moa_router_aux_loss(reference_weights, global_logits, 1.0)
+    reference.backward()
+    expected = sum(reference_values.grad.view(world_size, 2), start=torch.zeros(2))
+
+    global_sum_mean = global_weights.sum(dim=(0, 2, 3)) / world_size
+    global_count_mean = torch.tensor(1.0)
+    rank_gradients = []
+    for values in rank_inputs:
+        local_values = values.clone().requires_grad_()
+        local_weights = torch.softmax(local_values, dim=0).view(1, 2, 1, 1)
+        reductions = iter((global_sum_mean, global_count_mean))
+        monkeypatch.setattr("ultralytics.nn.modules.moa.router._all_reduce_mean", lambda _: next(reductions).clone())
+        local_loss = _moa_router_aux_loss(local_weights, torch.zeros_like(local_weights), 1.0, reduce_ddp=True)
+        local_loss.backward()
+        rank_gradients.append(local_values.grad)
+
+    ddp_averaged = torch.stack(rank_gradients).mean(dim=0)
+    assert torch.allclose(ddp_averaged, expected, atol=1e-6, rtol=1e-6)
+
+
 def test_moa_block_snapshot_keeps_pre_fallback_nonfinite_diagnostics(monkeypatch):
     block = MoABlock(24, num_heads=3).train()
 
@@ -70,6 +100,30 @@ def test_routed_modules_declare_sparse_export_boundary():
         assert "dense" in capabilities["sparse_export_limitation"].lower()
 
 
+def test_moa_sparse_inference_keeps_one_group_and_renormalizes(monkeypatch):
+    block = MoABlock(24, num_heads=3, inference_sparse_threshold=0.4).eval()
+    weights = torch.tensor([0.8, 0.1, 0.1]).view(1, 3, 1, 1).expand(2, 3, 4, 4)
+    monkeypatch.setattr(block.router, "forward", lambda x, return_logits=False: (weights, weights.log()))
+    calls = [0, 0, 0]
+    for idx, name in enumerate(("local_head", "region_head", "global_head")):
+        head = getattr(block, name)
+        original = head.forward
+
+        def counted(x, *, _idx=idx, _original=original):
+            calls[_idx] += 1
+            return _original(x)
+
+        monkeypatch.setattr(head, "forward", counted)
+
+    output = block(torch.randn(2, 24, 4, 4))
+    snapshot = block.routing_snapshot()
+
+    assert output.shape == (2, 24, 4, 4)
+    assert calls == [1, 0, 0]
+    assert snapshot["executed_groups"] == 1
+    assert snapshot["approximation_error"] > 0
+
+
 def test_c2f_moa_propagates_sequential_head_configuration():
     module = C2fMoA(48, 48, n=3, num_heads=3, sequential_heads=True)
 
diff --git a/ultralytics/nn/modules/moe/diagnostics.py b/ultralytics/nn/modules/moe/diagnostics.py
index 1ded3195..cf3ae8ba 100644
--- a/ultralytics/nn/modules/moe/diagnostics.py
+++ b/ultralytics/nn/modules/moe/diagnostics.py
@@ -7,6 +7,8 @@
 
 import torch
 
+from .protocol import routing_metrics, usage_gini
+
 
 @dataclass
 class MoELayerDiagnostic:
@@ -42,18 +44,19 @@ def collect_moe_diagnostics(model: torch.nn.Module, collapse_threshold: float =
         if not snapshot or num_experts <= 0:
             continue
 
-        usage = _tensor_to_list(snapshot.get("expert_usage")) or [0.0] * num_experts
-        counts = _tensor_to_list(snapshot.get("topk_counts")) or [0.0] * num_experts
-        dominant_share = max(usage) if usage else 0.0
-        dominant_expert = int(max(range(len(usage)), key=usage.__getitem__)) if usage else -1
+        metrics = routing_metrics(snapshot, num_experts=num_experts, top_k=int(getattr(module, "top_k", 0)))
+        usage = metrics.expert_usage
+        counts = metrics.topk_counts
+        dominant_share = metrics.dominant_share
+        dominant_expert = metrics.dominant_expert
 
         diagnostics.append(
             MoELayerDiagnostic(
                 name=name,
                 module_type=type(module).__name__,
                 num_experts=num_experts,
-                top_k=int(snapshot.get("top_k", getattr(module, "top_k", 0))),
-                aux_loss=float(snapshot.get("aux_loss", 0.0)),
+                top_k=metrics.top_k,
+                aux_loss=metrics.aux_loss,
                 usage=usage,
                 counts=counts,
                 dominant_expert=dominant_expert,
@@ -72,18 +75,6 @@ def diagnostics_to_dict(diagnostics: list[MoELayerDiagnostic]) -> list[dict[str,
     return [diag.__dict__.copy() for diag in diagnostics]
 
 
-def _gini(values: list[float]) -> float:
-    """Compute a bounded Gini coefficient from a non-negative usage vector."""
-    if not values:
-        return 0.0
-    usage = torch.tensor(values, dtype=torch.float32).clamp_min(0)
-    total = float(usage.sum())
-    if total <= 0:
-        return 0.0
-    diff = torch.abs(usage[:, None] - usage[None, :]).sum()
-    return float((diff / (2 * usage.numel() * total)).item())
-
-
 def routing_runtime_metrics(model: torch.nn.Module, collapse_threshold: float = 0.8) -> dict[str, Any]:
     """Return JSON-safe routing health and dispatch metrics after a forward."""
     layers: dict[str, dict[str, Any]] = {}
@@ -111,7 +102,7 @@ def routing_runtime_metrics(model: torch.nn.Module, collapse_threshold: float =
             "num_experts": int(snapshot.get("num_experts", len(usage))),
             "top_k": int(snapshot.get("top_k", getattr(module, "top_k", 0))),
             "expert_usage": usage,
-            "gini": _gini(usage),
+            "gini": usage_gini(usage),
             "entropy": float((-usage_tensor * torch.log(usage_tensor)).sum()),
             "dominant_share": max(usage),
             "collapse_flag": max(usage) >= float(collapse_threshold),
diff --git a/ultralytics/nn/modules/moe/schedule.py b/ultralytics/nn/modules/moe/schedule.py
index 11b9b944..b7fa15c0 100644
--- a/ultralytics/nn/modules/moe/schedule.py
+++ b/ultralytics/nn/modules/moe/schedule.py
@@ -9,29 +9,7 @@
 
 import torch
 
-
-def usage_gini(usage: Iterable[float] | torch.Tensor) -> float:
-    """Return the Gini coefficient for a non-negative expert-usage vector."""
-    if isinstance(usage, torch.Tensor):
-        values = usage.detach().float().reshape(-1)
-    else:
-        values = torch.tensor([float(v) for v in usage], dtype=torch.float32)
-
-    if values.numel() == 0:
-        return 0.0
-    values = values.clamp_min(0)
-    total = values.sum()
-
-    # Sorted cumulative form is O(n log n) and keeps the reduction on the
-    # source device until one final scalar conversion.
-    sorted_values = torch.sort(values).values
-    index = torch.arange(1, sorted_values.numel() + 1, device=sorted_values.device, dtype=sorted_values.dtype)
-    gini = (2 * torch.sum(index * sorted_values) / (sorted_values.numel() * total.clamp_min(1e-12))) - (
-        (sorted_values.numel() + 1) / sorted_values.numel()
-    )
-    gini = torch.where(total > 0, gini, torch.zeros_like(gini))
-    value = float(gini.clamp(0.0, 1.0).item())
-    return 0.0 if value != value else value
+from .protocol import usage_gini
 
 
 def mean_usage_gini_from_model(model: torch.nn.Module) -> float:
diff --git a/ultralytics/nn/modules/moe/scheduler.py b/ultralytics/nn/modules/moe/scheduler.py
index 6a8620f4..a06bb9d3 100644
--- a/ultralytics/nn/modules/moe/scheduler.py
+++ b/ultralytics/nn/modules/moe/scheduler.py
@@ -7,6 +7,8 @@
 
 import torch
 
+from .protocol import usage_gini
+
 
 @dataclass
 class MoEDynamicSchedulerConfig:
@@ -35,20 +37,8 @@ def to_dict(self) -> dict[str, float]:
 
 
 def compute_gini(expert_usage: torch.Tensor) -> float:
-    """Return the Gini coefficient of a non-negative expert-usage vector."""
-    usage = expert_usage.detach().float().reshape(-1).clamp_min(0.0)
-    n = usage.numel()
-    if n == 0:
-        return 0.0
-    total = usage.sum()
-    # Single D2H sync: defer .item() until after all GPU computation.
-    sorted_usage = torch.sort(usage / total.clamp_min(1e-12)).values
-    index = torch.arange(1, n + 1, device=sorted_usage.device, dtype=sorted_usage.dtype)
-    gini = (2 * torch.sum(index * sorted_usage) / n) - ((n + 1) / n)
-    gini_val = float(gini.clamp(0.0, 1.0).item())
-    if gini_val != gini_val:  # NaN check for zero-total edge case
-        return 0.0
-    return gini_val
+    """Backward-compatible alias for the canonical routing Gini metric."""
+    return usage_gini(expert_usage)
 
 
 class MoEDynamicScheduler:

From bb8443218a3741a5e5e0ae240a71e40dd1c54f9e Mon Sep 17 00:00:00 2001
From: isLinXu <2267379130@qq.com>
Date: Sun, 2 Aug 2026 18:48:51 +0800
Subject: [PATCH 04/11] refactor(moe): slim protocol module to routing snapshot
 metrics adapters

Replace the RoutedModule protocol with minimal dataclass-based adapters:
RoutingMetrics, normalize_routing_snapshot, routing_metrics, usage_gini.
Export the new helpers from the moe facade and cover them with
test_moe_facade_protocol.py.
---
 tests/test_moe_facade_protocol.py      |  39 +++++
 ultralytics/nn/modules/moe/__init__.py |   6 +
 ultralytics/nn/modules/moe/protocol.py | 191 +++++++++----------------
 3 files changed, 115 insertions(+), 121 deletions(-)
 create mode 100644 tests/test_moe_facade_protocol.py

diff --git a/tests/test_moe_facade_protocol.py b/tests/test_moe_facade_protocol.py
new file mode 100644
index 00000000..c7c3ce2e
--- /dev/null
+++ b/tests/test_moe_facade_protocol.py
@@ -0,0 +1,39 @@
+"""MoE facade tiers, compatibility exports, and snapshot protocol regression tests."""
+
+import pytest
+import torch
+
+import ultralytics.nn.modules.moe as facade
+from ultralytics.nn.modules.moe.protocol import normalize_routing_snapshot, routing_metrics
+from ultralytics.nn.modules.moe.schedule import usage_gini as schedule_gini
+from ultralytics.nn.modules.moe.scheduler import compute_gini
+
+
+def test_public_api_tiers_are_disjoint_and_resolvable():
+    tiers = (facade.STABLE_MOE_CLASSES, facade.EXPERIMENTAL_MOE_CLASSES, facade.LEGACY_MOE_CLASSES)
+    assert not (tiers[0] & tiers[1] or tiers[0] & tiers[2] or tiers[1] & tiers[2])
+    for name in set().union(*tiers):
+        assert name in facade.__all__
+        assert getattr(facade, name) is not None
+
+
+def test_legacy_checkpoint_aliases_remain_importable():
+    assert facade.A2C2fMoE.__name__ == "A2C2fMoE"
+    assert facade.ABlockMoE.__name__ == "ABlockMoE"
+    assert facade.is_legacy_moe("A2C2fMoE")
+
+
+def test_snapshot_adapter_accepts_legacy_keys_without_mutation():
+    source = {"usage": torch.tensor([0.75, 0.25]), "counts": [3, 1], "top_k": 1}
+    normalized = normalize_routing_snapshot(source, num_experts=2)
+    assert normalized["expert_usage"] == [0.75, 0.25]
+    assert normalized["topk_counts"] == [3.0, 1.0]
+    assert "expert_usage" not in source
+
+
+def test_scheduler_and_diagnostics_share_gini_definition():
+    usage = torch.tensor([1.0, 0.0, 0.0, 0.0])
+    metrics = routing_metrics({"expert_usage": usage}, num_experts=4)
+    assert metrics.gini == pytest.approx(0.75)
+    assert compute_gini(usage) == pytest.approx(metrics.gini)
+    assert schedule_gini(usage) == pytest.approx(metrics.gini)
diff --git a/ultralytics/nn/modules/moe/__init__.py b/ultralytics/nn/modules/moe/__init__.py
index 4d4e79ad..10e1ba12 100644
--- a/ultralytics/nn/modules/moe/__init__.py
+++ b/ultralytics/nn/modules/moe/__init__.py
@@ -75,6 +75,7 @@
 from .analysis import ExpertUsageTracker, diagnose_model, RoutingCollapseDetector
 from .diagnostics import MoELayerDiagnostic, collect_moe_diagnostics, diagnostics_to_dict, format_moe_diagnostics
 from .history import MoEDiagnosticsRecorder, export_moe_history_plots
+from .protocol import RoutingMetrics, normalize_routing_snapshot, routing_metrics, usage_gini
 from .pruning import prune_moe_model, prune_moe_module
 from .scheduler import (
     MoEDynamicScheduler,
@@ -107,6 +108,7 @@
 
 
 # ── API Stability Tiers ──────────────────────────────────────────────
+# `__all__` remains compatibility-complete; use these tier manifests for discovery.
 # STABLE: production-ready, well-tested, backward-compatible API.
 STABLE_MOE_CLASSES = frozenset(
     {
@@ -290,4 +292,8 @@ def is_legacy_moe(class_name: str) -> bool:
     "is_experimental_moe",
     "is_deprecated_moe",
     "is_legacy_moe",
+    "RoutingMetrics",
+    "normalize_routing_snapshot",
+    "routing_metrics",
+    "usage_gini",
 ]
diff --git a/ultralytics/nn/modules/moe/protocol.py b/ultralytics/nn/modules/moe/protocol.py
index 9daaa7ad..aa0d656f 100644
--- a/ultralytics/nn/modules/moe/protocol.py
+++ b/ultralytics/nn/modules/moe/protocol.py
@@ -1,132 +1,81 @@
-# 🐧Please note that this file has been modified by Tencent on 2026/02/13. All Tencent Modifications are Copyright (C) 2026 Tencent.
-"""Unified protocol for all routed (mixture) modules: MoE, MoA, MoT, MoLoRA.
+"""Minimal adapters for MoE routing snapshots and derived metrics."""
 
-Defines a common interface so that downstream code (trainers, loss collectors,
-exporters, diagnostic tools) can treat all mixture modules uniformly without
-``hasattr`` probes or module-specific branching.
-
-## RoutedModule Protocol
-
-Any nn.Module that routes input across multiple experts/heads SHOULD satisfy
-this protocol.  Existing MoE classes already comply; MoA, MoT, and MoLoRA
-are patched to comply via ``@property`` shims in their respective files.
-
-Required attributes/properties:
-  - ``num_experts`` (int): Total number of expert branches.
-  - ``top_k`` (int): Number of active experts per forward (``== num_experts`` if dense).
-  - ``aux_loss`` (Tensor): Scalar auxiliary loss (balance + z-loss); zero if eval.
-  - ``last_routing_snapshot`` (dict): Detached routing diagnostics from last forward.
-
-Optional (recommended) methods:
-  - ``get_gflops(input_shape) -> dict``: Per-component FLOPs estimate.
-  - ``__deepcopy__(memo)``: Safe deepcopy that strips non-leaf tensors.
-  - ``set_top_k(k)``: Dynamically adjust Top-K (MoE only; MoA/MoT use temperature).
-"""
 from __future__ import annotations
 
-from typing import Any, Dict, Optional, Protocol, runtime_checkable, Tuple
+from dataclasses import asdict, dataclass
+from typing import Any, Iterable, Mapping
 
 import torch
-import torch.nn as nn
-
-from ..routing_protocol import RoutingAuxPublisher
-
 
-@runtime_checkable
-class RoutedModule(Protocol):
-    """Protocol that all mixture-routing modules (MoE/MoA/MoT/MoLoRA) implement.
 
-    Use ``isinstance(module, RoutedModule)`` for duck-typing checks, or simply
-    access the attributes directly — the protocol is non-binding; each module
-    type already provides all required attributes.
-    """
+@dataclass(frozen=True)
+class RoutingMetrics:
+    """JSON-safe metrics derived from one normalized routing snapshot."""
 
-    # ── Required attributes ──────────────────────────────────────────
+    expert_usage: list[float]
+    topk_counts: list[float]
     num_experts: int
     top_k: int
-
-    @property
-    def aux_loss(self) -> torch.Tensor:
-        """Scalar auxiliary loss (balance + z-loss). Zero outside training."""
-        ...
-
-    @property
-    def last_routing_snapshot(self) -> Dict[str, Any]:
-        """Detached routing diagnostics from the most recent forward pass.
-
-        Keys (all optional, at least one present):
-          - ``expert_usage`` (Tensor [E]): Normalized per-expert usage share.
-          - ``mean_router_probs`` (Tensor [E]): Mean router probabilities.
-          - ``aux_loss`` (float): Detached scalar aux loss value.
-          - ``num_experts`` (int), ``top_k`` (int).
-        """
-        ...
-
-
-# ── Mixin for modules that want a zero-cost default implementation ──────
-
-class RoutedModuleMixin:
-    """Mixin providing default RoutedModule protocol shims.
-
-    Subclasses should override ``aux_loss`` and populate
-    ``last_routing_snapshot`` during forward.  This mixin ensures that
-    ``get_gflops`` and ``__deepcopy__`` are always available, even on modules
-    that don't define their own.
-    """
-
-    def get_gflops(self, input_shape: Tuple[int, int, int, int]) -> Dict[str, float]:
-        """Default GFLOPs estimate: sum over all Conv2d/Linear submodules."""
-        B, C, H, W = input_shape
-        total_macs = 0.0
-        for m in self.modules():
-            if isinstance(m, nn.Conv2d):
-                macs = B * m.in_channels * m.out_channels * (H // m.stride[0]) * (W // m.stride[1])
-                macs *= (m.kernel_size[0] * m.kernel_size[1]) / max(m.groups, 1)
-                total_macs += macs
-            elif isinstance(m, nn.Linear):
-                total_macs += B * m.in_features * m.out_features
-        gf = total_macs / 1e9
-        return {"total_gflops": gf, "conv_linear_gflops": gf}
-
-    def __deepcopy__(self, memo):
-        """Default safe deepcopy — delegates to ``_robust_deepcopy``."""
-        from ._helpers import _robust_deepcopy
-        return _robust_deepcopy(self, memo)
-
-
-def is_routed_module(module: nn.Module) -> bool:
-    """Check whether a module satisfies the RoutedModule protocol.
-
-    This is a structural check: the module must have ``num_experts``,
-    ``top_k``, ``aux_loss``, and ``last_routing_snapshot``.
-    """
-    return (
-        hasattr(module, "num_experts")
-        and hasattr(module, "top_k")
-        and hasattr(module, "aux_loss")
-        and hasattr(module, "last_routing_snapshot")
+    aux_loss: float
+    gini: float
+    dominant_expert: int
+    dominant_share: float
+
+    def to_dict(self) -> dict[str, Any]:
+        return asdict(self)
+
+
+def _float_list(value: Any) -> list[float]:
+    if isinstance(value, torch.Tensor):
+        value = value.detach().float().reshape(-1).cpu().tolist()
+    if value is None:
+        return []
+    return [float(item) for item in value]
+
+
+def normalize_routing_snapshot(
+    snapshot: Mapping[str, Any] | None, *, num_experts: int = 0, top_k: int = 0
+) -> dict[str, Any]:
+    """Normalize legacy/current snapshot keys without mutating the producer payload."""
+    source = snapshot or {}
+    usage = _float_list(source.get("expert_usage", source.get("usage")))
+    counts = _float_list(source.get("topk_counts", source.get("counts")))
+    size = int(source.get("num_experts", num_experts or len(usage)))
+    if size > 0:
+        usage = (usage + [0.0] * size)[:size]
+        counts = (counts + [0.0] * size)[:size]
+    return {
+        **source,
+        "expert_usage": usage,
+        "topk_counts": counts,
+        "num_experts": size,
+        "top_k": int(source.get("top_k", top_k)),
+        "aux_loss": float(source.get("aux_loss", 0.0)),
+    }
+
+
+def usage_gini(usage: Iterable[float] | torch.Tensor) -> float:
+    """Return one canonical bounded Gini coefficient for expert usage."""
+    values = torch.as_tensor(_float_list(usage), dtype=torch.float32).clamp_min(0.0)
+    if values.numel() == 0 or float(values.sum()) <= 0.0:
+        return 0.0
+    diff = torch.abs(values[:, None] - values[None, :]).sum()
+    return float((diff / (2 * values.numel() * values.sum())).item())
+
+
+def routing_metrics(snapshot: Mapping[str, Any] | None, *, num_experts: int = 0, top_k: int = 0) -> RoutingMetrics:
+    """Adapt a routing snapshot to the shared scheduler/diagnostic metric vocabulary."""
+    normalized = normalize_routing_snapshot(snapshot, num_experts=num_experts, top_k=top_k)
+    usage = normalized["expert_usage"]
+    dominant = max(range(len(usage)), key=usage.__getitem__) if usage else -1
+    share = usage[dominant] if dominant >= 0 else 0.0
+    return RoutingMetrics(
+        expert_usage=usage,
+        topk_counts=normalized["topk_counts"],
+        num_experts=normalized["num_experts"],
+        top_k=normalized["top_k"],
+        aux_loss=normalized["aux_loss"],
+        gini=usage_gini(usage),
+        dominant_expert=dominant,
+        dominant_share=share,
     )
-
-
-def collect_routed_children(module: nn.Module) -> list:
-    """Return all immediate+nested RoutedModule children of ``module``.
-
-    Useful for trainers and loss collectors that need to iterate over
-    all mixture modules in a model.
-    """
-    results = []
-    for child in module.modules():
-        if child is module:
-            continue
-        if is_routed_module(child):
-            results.append(child)
-    return results
-
-
-__all__ = [
-    "RoutingAuxPublisher",
-    "RoutedModule",
-    "RoutedModuleMixin",
-    "is_routed_module",
-    "collect_routed_children",
-]

From 62aa5d9515ce19ac1a7c709f65d98009894a2d8e Mon Sep 17 00:00:00 2001
From: isLinXu <2267379130@qq.com>
Date: Sun, 2 Aug 2026 18:49:00 +0800
Subject: [PATCH 05/11] feat(latent-mixture): enrich aux mixin and routing
 snapshot diagnostics

Extend _LatentAuxMixin and LatentMixture with additional aux-loss
handling and routing snapshot fields; expand latent mixture tests.
---
 tests/test_latent_mixture.py             | 59 ++++++++++++++++++++++++
 ultralytics/nn/modules/latent_mixture.py | 43 +++++++++++++++--
 2 files changed, 97 insertions(+), 5 deletions(-)

diff --git a/tests/test_latent_mixture.py b/tests/test_latent_mixture.py
index d4066314..d48afbba 100644
--- a/tests/test_latent_mixture.py
+++ b/tests/test_latent_mixture.py
@@ -101,6 +101,13 @@ def test_latent_mixture_publishes_single_train_aux_and_snapshot():
     assert records[0][0] is module
     snapshot = module.routing_snapshot()
     assert snapshot["family"] == "latent"
+    assert snapshot["top_k"] == snapshot["training_top_k"] == 4
+    assert snapshot["inference_top_k"] == 4
+    assert snapshot["executed_experts"] == 4
+    assert snapshot["configured_top_k"] == 4
+    assert snapshot["active_experts_per_sample"].tolist() == [4, 4]
+    assert snapshot["batch_expert_union"] == 4
+    assert snapshot["kernel_calls"] == 4
     assert snapshot["noise_std"] == 0.0
     assert snapshot["router_init_std"] == 0.0
     assert snapshot["identity_cold_start"] is False
@@ -238,3 +245,55 @@ def test_latent_temperature_anneal_updates_latent_modules():
 
     assert updated == 1
     assert float(module.temperature) == pytest.approx(0.475)
+
+
+def test_latent_value_fusion_has_explicit_causal_semantics():
+    torch.manual_seed(7)
+    first = torch.randn(2, 8, 4, 4)
+    auxiliary = torch.randn(2, 8, 4, 4)
+    changed = auxiliary.flip(0)
+
+    legacy = LatentMixture([8, 8], 8, residual_init=0.0).eval()
+    fused = LatentMixture(
+        [8, 8], 8, residual_init=0.0, value_fusion_mode="weighted_sum", value_fusion_weights=[1.0, 1.0]
+    ).eval()
+    fused.load_state_dict(legacy.state_dict(), strict=False)
+
+    with patch.object(legacy.router, "forward", return_value=(torch.zeros(2, 4), torch.full((2, 4), 0.25))):
+        legacy_a = legacy([first, auxiliary])
+        legacy_b = legacy([first, changed])
+    with patch.object(fused.router, "forward", return_value=(torch.zeros(2, 4), torch.full((2, 4), 0.25))):
+        fused_a = fused([first, auxiliary])
+        fused_b = fused([first, changed])
+
+    assert torch.allclose(legacy_a, legacy_b)
+    assert not torch.allclose(fused_a, fused_b)
+    assert fused.routing_snapshot()["value_fusion_mode"] == "weighted_sum"
+
+
+def test_latent_value_fusion_ablations_are_reproducible():
+    torch.manual_seed(8)
+    xs = [torch.randn(2, 8, 4, 4), torch.randn(2, 8, 4, 4)]
+    shuffled = [xs[0], xs[1].flip(0)]
+    first_only = LatentMixture(
+        [8, 8], 8, residual_init=0.0, value_fusion_mode="weighted_sum", value_fusion_weights=[1.0, 0.0]
+    ).eval()
+    equal = LatentMixture(
+        [8, 8], 8, residual_init=0.0, value_fusion_mode="weighted_sum", value_fusion_weights=[1.0, 1.0]
+    ).eval()
+    equal.load_state_dict(first_only.state_dict(), strict=False)
+
+    assert torch.allclose(first_only(xs), first_only(shuffled))
+    assert not torch.allclose(equal(xs), equal(shuffled))
+    assert torch.allclose(first_only.value_fusion_weights, torch.tensor([1.0, 0.0]))
+
+
+def test_latent_value_fusion_is_opt_in_and_checkpoint_compatible():
+    default = LatentMixture([8, 8], 8)
+    restored = LatentMixture([8, 8], 8)
+    restored.load_state_dict(default.state_dict(), strict=True)
+
+    assert default.value_fusion_mode == "router_only"
+    assert "value_fusion_weights" not in default.state_dict()
+    with pytest.raises(ValueError, match="value_fusion_mode"):
+        LatentMixture([8, 8], 8, value_fusion_mode="unknown")
diff --git a/ultralytics/nn/modules/latent_mixture.py b/ultralytics/nn/modules/latent_mixture.py
index 0d475d6c..e46176dd 100644
--- a/ultralytics/nn/modules/latent_mixture.py
+++ b/ultralytics/nn/modules/latent_mixture.py
@@ -331,6 +331,16 @@ def _record_routing(
                 "family": "latent",
                 "num_experts": int(self.num_experts),
                 "top_k": int(self.top_k),
+                "training_top_k": int(self.top_k),
+                "inference_top_k": int(self.top_k),
+                "configured_top_k": int(self.top_k),
+                "executed_experts": int(self.num_experts),
+                "active_experts_per_sample": torch.full(
+                    (p.shape[0],), int(self.num_experts), dtype=torch.long
+                ),
+                "mean_active_experts_per_sample": float(self.num_experts),
+                "batch_expert_union": int(self.num_experts),
+                "kernel_calls": int(self.num_experts),
                 "mean_router_probs": mean_probs.cpu(),
                 "expert_usage": mean_probs.cpu(),
                 "entropy": float(entropy.cpu()),
@@ -357,6 +367,9 @@ def _record_routing(
                 snapshot["scale_mean_probs"] = p.mean(dim=0).cpu()
             else:
                 snapshot["routing_axis"] = "expert"
+            if hasattr(self, "value_fusion_mode"):
+                snapshot["value_fusion_mode"] = self.value_fusion_mode
+                snapshot["value_fusion_weights"] = self.value_fusion_weights.detach().float().cpu()
             if hasattr(self, "residual_gain"):
                 snapshot["residual_gain"] = self.residual_gain.detach().float().cpu()
             self.last_routing_snapshot = snapshot
@@ -403,6 +416,8 @@ def __init__(
         noise_std: float = 0.0,
         router_init_std: float = 0.0,
         inference_top_k: int | None = None,
+        value_fusion_mode: str = "router_only",
+        value_fusion_weights: Sequence[float] | None = None,
     ):
         super().__init__()
         if isinstance(in_channels, int):
@@ -412,6 +427,22 @@ def __init__(
             raise ValueError("LatentMixture requires at least one input channel")
         self.out_channels = _positive_int(out_channels, "out_channels")
         self.num_inputs = len(self.in_channels)
+        self.value_fusion_mode = str(value_fusion_mode)
+        if self.value_fusion_mode not in {"router_only", "weighted_sum"}:
+            raise ValueError(
+                f"value_fusion_mode must be 'router_only' or 'weighted_sum', got {self.value_fusion_mode!r}"
+            )
+        if value_fusion_weights is None:
+            weights = torch.ones(self.num_inputs, dtype=torch.float32)
+        else:
+            if len(value_fusion_weights) != self.num_inputs:
+                raise ValueError(
+                    f"value_fusion_weights must contain {self.num_inputs} values, got {len(value_fusion_weights)}"
+                )
+            weights = torch.as_tensor(value_fusion_weights, dtype=torch.float32)
+            if not bool(torch.isfinite(weights).all()) or bool((weights < 0).any()) or float(weights.sum()) <= 0.0:
+                raise ValueError("value_fusion_weights must be finite, non-negative, and have a positive sum")
+        self.register_buffer("value_fusion_weights", weights / weights.sum(), persistent=False)
         self.num_experts = _positive_int(num_experts, "num_experts")
         self.top_k = self.num_experts
         self.inference_top_k = self.num_experts if inference_top_k is None else int(inference_top_k)
@@ -445,11 +476,13 @@ def __init__(
 
     def _build_context(self, xs: Sequence[torch.Tensor]) -> tuple[torch.Tensor, LatentRoutingContext]:
         checked = _validate_inputs(xs, self.in_channels, require_same_spatial=True)
-        base = self.base_proj(checked[0])
-        tokens = []
-        for x, proj in zip(checked, self.token_projs):
-            token = proj(x)
-            tokens.append(F.adaptive_avg_pool2d(token, 1).flatten(1).float())
+        projected = [proj(x) for x, proj in zip(checked, self.token_projs)]
+        if self.value_fusion_mode == "weighted_sum":
+            weights = self.value_fusion_weights.to(device=projected[0].device, dtype=projected[0].dtype)
+            base = sum(value * weights[i] for i, value in enumerate(projected))
+        else:
+            base = self.base_proj(checked[0])
+        tokens = [F.adaptive_avg_pool2d(token, 1).flatten(1).float() for token in projected]
         scale_tokens = torch.stack(tokens, dim=1)
         logits, probs = self.router(scale_tokens)
         return base, LatentRoutingContext(

From eed8a1ba03fe46056cdbf7f3d6b08d20af45a8f0 Mon Sep 17 00:00:00 2001
From: isLinXu <2267379130@qq.com>
Date: Sun, 2 Aug 2026 18:49:08 +0800
Subject: [PATCH 06/11] perf(molora): vectorize same-rank Linear experts in
 batched einsum kernels

Add _can_vectorize_linear_experts/_compute_vectorized_linear_experts to
run same-rank Linear experts in two batched kernels; update sparse
dispatch and routing-aware merge tests.
---
 tests/test_molora_routing_aware_merge.py |  4 +-
 tests/test_molora_sparse_dispatch.py     | 39 +++++++++++++-
 ultralytics/nn/peft/molora/layer.py      | 69 ++++++++++++++++++++++--
 3 files changed, 107 insertions(+), 5 deletions(-)

diff --git a/tests/test_molora_routing_aware_merge.py b/tests/test_molora_routing_aware_merge.py
index e1b90c34..0d716cc2 100644
--- a/tests/test_molora_routing_aware_merge.py
+++ b/tests/test_molora_routing_aware_merge.py
@@ -1,5 +1,7 @@
 """Routing-aware MoLoRA merge tests."""
 
+from typing import List
+
 import pytest
 import torch
 import torch.nn as nn
@@ -8,7 +10,7 @@
 from ultralytics.utils.lora import adapter_metadata, merge_adapters
 
 
-def _set_router_bias(layer: MoLoRALayer, bias: list[float]) -> None:
+def _set_router_bias(layer: MoLoRALayer, bias: List[float]) -> None:
     with torch.no_grad():
         layer.router.fc[-1].weight.zero_()
         layer.router.fc[-1].bias.copy_(torch.tensor(bias))
diff --git a/tests/test_molora_sparse_dispatch.py b/tests/test_molora_sparse_dispatch.py
index 82ff5b28..c0fe72bc 100644
--- a/tests/test_molora_sparse_dispatch.py
+++ b/tests/test_molora_sparse_dispatch.py
@@ -14,8 +14,9 @@ def test_molora_grouped_dispatch_records_actual_calls():
     assert stats["selected_samples"] == 8
 
 
-def test_molora_small_batch_keeps_dense_fast_path():
+def test_molora_small_batch_falls_back_to_dense_when_not_vectorizable():
     layer = MoLoRALayer(nn.Linear(16, 16), r=2, num_experts=4, top_k=1).eval()
+    layer._can_vectorize_linear_experts = lambda: False
     layer(torch.randn(2, 16))
     assert layer._last_dispatch_stats["mode"] == "dense_small_batch"
 
@@ -32,3 +33,39 @@ def test_molora_all_linear_experts_use_batched_dispatch_without_changing_values(
 
     assert layer._last_dispatch_stats["mode"] == "batched_dense_linear"
     assert torch.allclose(batched, reference, atol=1e-6, rtol=1e-5)
+
+
+def test_molora_small_batch_uses_vectorized_dense_linear_path():
+    torch.manual_seed(1)
+    layer = MoLoRALayer(nn.Linear(16, 16), r=2, num_experts=4, top_k=2).eval()
+    x = torch.randn(2, 16)
+    output = layer(x)
+    stats = layer._last_dispatch_stats
+
+    assert output.shape == (2, 16)
+    assert stats["mode"] == "vectorized_dense_linear"
+    assert stats["expert_calls"] == 4
+    assert stats["dispatch_shape"] == (2, 4)
+
+
+def test_molora_vectorized_dense_linear_matches_sparse_fallback():
+    torch.manual_seed(2)
+    layer = MoLoRALayer(nn.Linear(16, 12), r=3, num_experts=4, top_k=2).eval()
+    x = torch.randn(2, 16)
+    router_logits = layer.router(x)
+    router_probs = torch.softmax(router_logits, dim=-1)
+    weights, indices = torch.topk(router_probs, 2, dim=-1)
+    weights = weights / weights.sum(dim=-1, keepdim=True)
+    template = layer.base_layer(x)
+
+    dense = layer._compute_vectorized_linear_experts(x, weights, indices, template)
+    dense_stats = dict(layer._last_dispatch_stats)
+    original_guard = layer._can_vectorize_linear_experts
+    layer._can_vectorize_linear_experts = lambda: False
+    sparse = layer._compute_sparse_experts(x, weights, indices, template)
+    layer._can_vectorize_linear_experts = original_guard
+
+    torch.testing.assert_close(dense, sparse)
+    assert dense_stats["mode"] == "vectorized_dense_linear"
+    assert layer._last_dispatch_stats["mode"] == "dense_small_batch"
+    assert layer._last_dispatch_stats["dispatch_shape"][0] <= x.shape[0] * layer.top_k
diff --git a/ultralytics/nn/peft/molora/layer.py b/ultralytics/nn/peft/molora/layer.py
index 12a3cfe3..71694288 100644
--- a/ultralytics/nn/peft/molora/layer.py
+++ b/ultralytics/nn/peft/molora/layer.py
@@ -238,6 +238,13 @@ def __init__(
             reduce_ddp=True,
         )
 
+        # New adapter state must follow an already-placed base layer. Experts use
+        # the base dtype, while router parameters start in FP32 for stable routing.
+        # A later explicit layer/model .half() still converts the router normally.
+        base_weight = base_layer.weight
+        self.to(device=base_weight.device, dtype=base_weight.dtype)
+        self.router.to(device=base_weight.device, dtype=torch.float32)
+
         # Routing stats for diagnostics (not persistent)
         self._last_routing_stats: Optional[Dict[str, Any]] = None
         self.last_routing_snapshot: dict = {}
@@ -518,8 +525,10 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         if self.training:
             self._step_count.add_(1)
 
-        # Router
-        router_logits = self.router(x)  # [B, E]
+        # Keep router execution in its parameter dtype. Injection initializes it
+        # in FP32, but explicit module conversions such as .half() remain honored.
+        router_parameter = next(self.router.parameters())
+        router_logits = self.router(x.to(router_parameter.dtype))  # [B, E]
         calibration_applied = self.router_calibration is not None
         if calibration_applied:
             router_logits = self.router_calibration(x, router_logits)
@@ -646,8 +655,12 @@ def _compute_sparse_experts(
                 shape = (-1,) + (1,) * (out_e.dim() - 1)
                 expert_out = expert_out + out_e * dense_weights[:, expert_idx].view(shape)
             return expert_out
-        grouped = B >= 4
+        if B < 4 and isinstance(self.base_layer, nn.Linear) and self._can_vectorize_linear_experts():
+            return self._compute_vectorized_linear_experts(x, top_k_weights, top_k_indices, out_template)
+
         calls = 0
+        dispatch_rows = 0
+        grouped = B >= 4
         for e in range(self.num_experts):
             mask = top_k_indices == e
             selected = mask.any(dim=1)
@@ -655,6 +668,7 @@ def _compute_sparse_experts(
             if batch_idx.numel() == 0:
                 continue
             calls += 1
+            dispatch_rows += batch_idx.numel()
             x_e = x[batch_idx]
             out_e = self.experts[e](x_e)
             weights = (top_k_weights[batch_idx] * mask[batch_idx].to(top_k_weights.dtype)).sum(dim=1)
@@ -665,6 +679,7 @@ def _compute_sparse_experts(
             "expert_calls": calls,
             "selected_samples": B,
             "top_k": K,
+            "dispatch_shape": (dispatch_rows, K),
         }
         return expert_out
 
@@ -702,6 +717,54 @@ def _compute_sparse_experts_batched(
         selected = batch_outputs.gather(1, gather_index)
         return (selected * top_k_weights.to(selected.dtype).unsqueeze(-1)).sum(dim=1).to(x.dtype)
 
+    def _can_vectorize_linear_experts(self) -> bool:
+        """Return whether experts share the Linear topology required by the dense kernel."""
+        if not self.experts:
+            return False
+        rank = self.experts[0].r
+        return all(
+            not expert.is_conv
+            and expert.r == rank
+            and isinstance(expert.lora_A, nn.Linear)
+            and isinstance(expert.lora_B, nn.Linear)
+            for expert in self.experts
+        )
+
+    def _compute_vectorized_linear_experts(
+        self,
+        x: torch.Tensor,
+        top_k_weights: torch.Tensor,
+        top_k_indices: torch.Tensor,
+        out_template: torch.Tensor,
+    ) -> torch.Tensor:
+        """Run all same-rank Linear experts in two batched einsum kernels."""
+        parameter_dtype = self.experts[0].lora_A.weight.dtype
+        x_stacked = x.to(parameter_dtype).unsqueeze(-2).expand(*x.shape[:-1], self.num_experts, x.shape[-1])
+        dropout_p = self.experts[0].dropout.p if isinstance(self.experts[0].dropout, nn.Dropout) else 0.0
+        x_stacked = F.dropout(x_stacked, p=dropout_p, training=self.training)
+        a_weight = torch.stack([expert.lora_A.weight for expert in self.experts])
+        b_weight = torch.stack([expert.lora_B.weight for expert in self.experts])
+        scales = x_stacked.new_tensor([expert.scaling for expert in self.experts])
+        low_rank = torch.einsum("...ei,eri->...er", x_stacked, a_weight)
+        all_outputs = torch.einsum("...er,eor->...eo", low_rank, b_weight)
+        scale_shape = (1,) * (all_outputs.dim() - 2) + (self.num_experts, 1)
+        all_outputs = all_outputs * scales.view(scale_shape)
+
+        assignments = F.one_hot(top_k_indices, num_classes=self.num_experts).to(top_k_weights.dtype)
+        dense_weights = (assignments * top_k_weights.unsqueeze(-1)).sum(dim=1)
+        weight_shape = (x.shape[0],) + (1,) * (all_outputs.dim() - 3) + (self.num_experts, 1)
+        adapted = (all_outputs * dense_weights.to(parameter_dtype).view(weight_shape)).sum(dim=-2)
+        if adapted.dtype != out_template.dtype:
+            adapted = adapted.to(out_template.dtype)
+        self._last_dispatch_stats = {
+            "mode": "vectorized_dense_linear",
+            "expert_calls": self.num_experts,
+            "selected_samples": x.shape[0],
+            "top_k": top_k_indices.shape[1],
+            "dispatch_shape": (x.shape[0], self.num_experts),
+        }
+        return adapted
+
     def _expert_usage(self, expert_indices: torch.Tensor) -> torch.Tensor:
         """Normalized expert usage histogram."""
         flat = expert_indices.reshape(-1).to(torch.long)

From f8d2626b82d8ea5009fc56c0e2598d2d5de9cc1d Mon Sep 17 00:00:00 2001
From: isLinXu <2267379130@qq.com>
Date: Sun, 2 Aug 2026 18:49:17 +0800
Subject: [PATCH 07/11] feat(vpeft): add PlannerResult as stable external
 planner contract

Introduce PlannerResult dataclass with to_dict/from_dict plus adapters
from legacy decisions and placement plans; export from vpeft facade and
extend placement schema and V-PEFT e2e tests.
---
 tests/test_placement_plan_schema.py |   2 +-
 tests/test_vpeft_lora_e2e.py        |  50 +++++++++++++
 ultralytics/vpeft/__init__.py       |   3 +-
 ultralytics/vpeft/placement_plan.py | 109 +++++++++++++++++++++++++++-
 4 files changed, 161 insertions(+), 3 deletions(-)

diff --git a/tests/test_placement_plan_schema.py b/tests/test_placement_plan_schema.py
index de306f3d..1f844cea 100644
--- a/tests/test_placement_plan_schema.py
+++ b/tests/test_placement_plan_schema.py
@@ -3,7 +3,7 @@
 import pytest
 import torch.nn as nn
 
-from ultralytics.vpeft import PlacementPlan, PlacementTarget
+from ultralytics.vpeft import PlacementPlan, PlacementTarget, PlannerResult
 
 
 def test_placement_plan_roundtrip_is_stable():
diff --git a/tests/test_vpeft_lora_e2e.py b/tests/test_vpeft_lora_e2e.py
index 382ef854..17861626 100644
--- a/tests/test_vpeft_lora_e2e.py
+++ b/tests/test_vpeft_lora_e2e.py
@@ -157,3 +157,53 @@ def test_active_vpeft_plan_does_not_run_legacy_planner(monkeypatch):
 
     assert wrapped.lora_target_modules == ["0"]
     assert wrapped[0].r == 2
+
+
+def test_vpeft_internal_failure_falls_back_with_structured_metadata(monkeypatch):
+    import ultralytics.utils.lora.api as api
+
+    def fail(*args, **kwargs):
+        raise RuntimeError("planner bug")
+
+    monkeypatch.setattr(api, "_build_vpeft_placement_plan", fail)
+    model = apply_lora(
+        _model(),
+        LoRAConfig(r=4, alpha=8, backend="fallback", planner_backend="vpeft"),
+    )
+
+    assert model.lora_runtime_metadata["vpeft_fallback"] == {
+        "category": "internal",
+        "reason": "planner bug",
+        "message": "planner bug",
+        "exception_type": "RuntimeError",
+    }
+    assert model.lora_runtime_metadata["planner_result"]["status"] == "FALLBACK"
+    assert model.lora_runtime_metadata["planner_result"]["reason"]["category"] == "internal"
+
+
+def test_vpeft_strict_reraises_internal_failure(monkeypatch):
+    import ultralytics.utils.lora.api as api
+
+    def fail(*args, **kwargs):
+        raise RuntimeError("planner bug")
+
+    monkeypatch.setattr(api, "_build_vpeft_placement_plan", fail)
+    with pytest.raises(RuntimeError, match="planner bug"):
+        apply_lora(
+            _model(),
+            LoRAConfig(r=4, alpha=8, backend="fallback", planner_backend="vpeft", vpeft_strict=True),
+        )
+
+
+def test_vpeft_strict_reraises_configuration_failure(monkeypatch):
+    import ultralytics.utils.lora.api as api
+
+    def fail(*args, **kwargs):
+        raise ValueError("unsupported solver")
+
+    monkeypatch.setattr(api, "_build_vpeft_placement_plan", fail)
+    with pytest.raises(ValueError, match="unsupported solver"):
+        apply_lora(
+            _model(),
+            LoRAConfig(r=4, alpha=8, backend="fallback", planner_backend="vpeft", vpeft_strict=True),
+        )
diff --git a/ultralytics/vpeft/__init__.py b/ultralytics/vpeft/__init__.py
index ac64f377..1753423e 100644
--- a/ultralytics/vpeft/__init__.py
+++ b/ultralytics/vpeft/__init__.py
@@ -45,7 +45,7 @@
     SEMANTIC_UTILITY,
     RANK_SET,
 )
-from .placement_plan import PlacementPlan, PlacementTarget
+from .placement_plan import PlacementPlan, PlacementTarget, PlannerResult
 
 __all__ = [
     # Solver
@@ -86,4 +86,5 @@
     "RANK_SET",
     "PlacementPlan",
     "PlacementTarget",
+    "PlannerResult",
 ]
diff --git a/ultralytics/vpeft/placement_plan.py b/ultralytics/vpeft/placement_plan.py
index 5b284b70..c37dd04b 100644
--- a/ultralytics/vpeft/placement_plan.py
+++ b/ultralytics/vpeft/placement_plan.py
@@ -143,4 +143,111 @@ def from_dict(cls, payload: Mapping[str, Any]) -> "PlacementPlan":
         return plan
 
 
-__all__ = ["PlacementPlan", "PlacementTarget"]
+@dataclass(frozen=True)
+class PlannerResult:
+    """Stable external contract shared by legacy and V-PEFT planners.
+
+    ``PlacementDecision`` and ``PlacementPlan`` remain internal/backward-
+    compatible types. API and checkpoint consumers should prefer this type.
+    """
+
+    status: str
+    backend: str
+    reason: dict[str, Any] | None = None
+    targets: tuple[str, ...] = ()
+    rank_pattern: dict[str, int] = field(default_factory=dict)
+    budget: dict[str, int] = field(default_factory=dict)
+    fallback: bool = False
+    strict: bool = False
+    metadata: dict[str, Any] = field(default_factory=dict)
+    schema_version: int = 1
+
+    def __post_init__(self) -> None:
+        status = str(self.status).upper()
+        if status not in {"ACCEPT", "ADAPT", "REFUSE", "FALLBACK"}:
+            raise ValueError(f"invalid PlannerResult status={self.status!r}")
+        object.__setattr__(self, "status", status)
+        if self.fallback != (status == "FALLBACK"):
+            raise ValueError("fallback must be true exactly when status is FALLBACK")
+        if self.reason is not None and not self.reason.get("message"):
+            raise ValueError("PlannerResult reason requires a non-empty message")
+        if any(int(rank) <= 0 for rank in self.rank_pattern.values()):
+            raise ValueError("PlannerResult ranks must be positive")
+
+    def to_dict(self) -> dict[str, Any]:
+        return {
+            "schema_version": self.schema_version,
+            "status": self.status,
+            "backend": self.backend,
+            "reason": dict(self.reason) if self.reason else None,
+            "targets": list(self.targets),
+            "rank_pattern": {name: int(rank) for name, rank in self.rank_pattern.items()},
+            "budget": {key: int(value) for key, value in self.budget.items()},
+            "fallback": self.fallback,
+            "strict": self.strict,
+            "metadata": dict(self.metadata),
+        }
+
+    @classmethod
+    def from_dict(cls, payload: Mapping[str, Any]) -> "PlannerResult":
+        return cls(
+            schema_version=int(payload.get("schema_version", 1)),
+            status=str(payload.get("status", "FALLBACK")),
+            backend=str(payload.get("backend", "legacy")),
+            reason=dict(payload["reason"]) if payload.get("reason") else None,
+            targets=tuple(str(name) for name in payload.get("targets", ())),
+            rank_pattern={str(name): int(rank) for name, rank in dict(payload.get("rank_pattern", {})).items()},
+            budget={str(key): int(value) for key, value in dict(payload.get("budget", {})).items()},
+            fallback=bool(payload.get("fallback", str(payload.get("status", "")).upper() == "FALLBACK")),
+            strict=bool(payload.get("strict", False)),
+            metadata=dict(payload.get("metadata", {})),
+        )
+
+    @classmethod
+    def from_legacy_decision(cls, decision: Any, *, strict: bool = False) -> "PlannerResult":
+        status = str(getattr(decision, "status", "ACCEPT")).upper()
+        message = getattr(decision, "refusal_reason", None)
+        targets = tuple(getattr(decision, "target_modules_hint", None) or ())
+        rank = getattr(decision, "recommended_rank", None)
+        return cls(
+            status=status,
+            backend="legacy",
+            reason={"category": "planner_refusal", "message": str(message)} if message else None,
+            targets=targets,
+            rank_pattern={name: int(rank) for name in targets} if rank else {},
+            fallback=False,
+            strict=strict,
+            metadata={
+                "recommended_variant": getattr(decision, "recommended_variant", None),
+                "predicted_delta": getattr(decision, "predicted_delta", None),
+                "safety_overrides": dict(getattr(decision, "safety_overrides", {}) or {}),
+                **dict(getattr(decision, "metadata", {}) or {}),
+            },
+        )
+
+    @classmethod
+    def from_placement_plan(
+        cls,
+        plan: PlacementPlan,
+        *,
+        strict: bool = False,
+        fallback_reason: Mapping[str, Any] | None = None,
+    ) -> "PlannerResult":
+        status = "FALLBACK" if fallback_reason else plan.status
+        reason = dict(fallback_reason) if fallback_reason else None
+        if reason is None and plan.refusal_reason:
+            reason = {"category": "infeasible", "message": str(plan.refusal_reason)}
+        return cls(
+            status=status,
+            backend=plan.planner_backend,
+            reason=reason,
+            targets=tuple(target.name for target in plan.targets),
+            rank_pattern={target.name: int(target.rank) for target in plan.targets},
+            budget=dict(plan.budget),
+            fallback=status == "FALLBACK",
+            strict=strict,
+            metadata={"solver": plan.solver, "plan_fingerprint": plan.fingerprint, **dict(plan.metadata)},
+        )
+
+
+__all__ = ["PlacementPlan", "PlacementTarget", "PlannerResult"]

From f3007cee4acdff128df30d8c620d9f7c8481a311 Mon Sep 17 00:00:00 2001
From: isLinXu <2267379130@qq.com>
Date: Sun, 2 Aug 2026 18:49:26 +0800
Subject: [PATCH 08/11] feat(lora): persist planner result contract and
 structured V-PEFT fallback reasons

Attach PlannerResult via _attach_planner_result and record serializable
fallback reasons via _record_vpeft_fallback in apply_lora; extend
LoRAConfig fields, default.yaml options, and config drift detector
coverage with matching tests.
---
 tests/test_config_drift_detector.py | 21 +++++++
 tests/test_peft_adapters.py         | 30 +++++++++
 tools/config_drift_detector.py      |  5 +-
 ultralytics/cfg/default.yaml        |  3 +-
 ultralytics/utils/lora/api.py       | 98 ++++++++++++++++++++++++++++-
 ultralytics/utils/lora/config.py    |  2 +
 6 files changed, 155 insertions(+), 4 deletions(-)

diff --git a/tests/test_config_drift_detector.py b/tests/test_config_drift_detector.py
index 6ce703c7..5d76a7c0 100644
--- a/tests/test_config_drift_detector.py
+++ b/tests/test_config_drift_detector.py
@@ -170,6 +170,27 @@ def test_incompatible_model_argument_counts_are_reported(tmp_path):
     assert any(issue.code == "MODEL_ARGS_INCOMPATIBLE" and "Known" in issue.message for issue in issues)
 
 
+def test_check_all_includes_yolo26_master_configs(tmp_path):
+    detector = ConfigDriftDetector(tmp_path)
+    _write(tmp_path / detector.DEFAULT_CFG, "alpha: 1\n")
+    _write(tmp_path / detector.YOLO26_MASTER_MODELS / "yolo26-master-test.yaml", "alpha: 1\nalpha: 2\n")
+
+    issues = detector.check_all()
+
+    assert any(
+        issue.code == "YAML_DUPLICATE_KEY" and issue.path.name == "yolo26-master-test.yaml" for issue in issues
+    )
+
+
+def test_lora_rank_pattern_cli_mapping_round_trip():
+    from ultralytics.utils.lora.config import LoRAConfig
+
+    rank_pattern = {"model.0.conv": 4, "model.1.conv": 8}
+    config = LoRAConfig.from_args(lora_rank_pattern=rank_pattern)
+
+    assert config.rank_pattern == rank_pattern
+
+
 def test_repository_has_no_configuration_drift():
     issues = ConfigDriftDetector(ROOT).check_all()
 
diff --git a/tests/test_peft_adapters.py b/tests/test_peft_adapters.py
index 158abcfe..c659612e 100644
--- a/tests/test_peft_adapters.py
+++ b/tests/test_peft_adapters.py
@@ -689,6 +689,36 @@ def test_molora_model_device(self, tiny_cnn):
         for p in model.parameters():
             assert p.device.type == "cpu"
 
+    def test_molora_injection_inherits_preplaced_base_device(self, tiny_cnn):
+        """Experts/router should inherit device when injection follows model placement."""
+        model = tiny_cnn.to("cpu")
+        cfg = MoLoRAConfig(
+            r=4, alpha=8, num_experts=4, top_k=2,
+            target_modules=["conv1"]
+        )
+        model = get_peft_molora_model(model, cfg)
+        layer = model.conv1
+        assert layer.base_layer.weight.device.type == "cpu"
+        assert all(p.device.type == "cpu" for p in layer.experts.parameters())
+        assert all(p.device.type == "cpu" for p in layer.router.parameters())
+        assert all(p.dtype == layer.base_layer.weight.dtype for p in layer.experts.parameters())
+        assert all(p.dtype == torch.float32 for p in layer.router.parameters())
+        assert model(torch.randn(2, 3, 8, 8)).device.type == "cpu"
+
+    @pytest.mark.skipif(not torch.backends.mps.is_available(), reason="MPS not available")
+    def test_molora_injection_after_mps_placement_smoke(self, tiny_cnn):
+        """Injecting after MPS placement should not leave new adapter state on CPU."""
+        model = tiny_cnn.to("mps")
+        cfg = MoLoRAConfig(
+            r=4, alpha=8, num_experts=4, top_k=2,
+            target_modules=["conv1"]
+        )
+        model = get_peft_molora_model(model, cfg)
+        layer = model.conv1
+        assert all(p.device.type == "mps" for p in layer.experts.parameters())
+        assert all(p.device.type == "mps" for p in layer.router.parameters())
+        assert model(torch.randn(2, 3, 8, 8, device="mps")).device.type == "mps"
+
     def test_molora_model_dtype(self, tiny_cnn):
         """MoLoRA model should support float16."""
         cfg = MoLoRAConfig(
diff --git a/tools/config_drift_detector.py b/tools/config_drift_detector.py
index af352560..cde92f5d 100644
--- a/tools/config_drift_detector.py
+++ b/tools/config_drift_detector.py
@@ -106,6 +106,7 @@ class ConfigDriftDetector:
     TASKS = Path("ultralytics/nn/tasks.py")
     MIXTURE_REGISTRY = Path("ultralytics/nn/mixture_registry.py")
     MASTER_MODELS = Path("ultralytics/cfg/models/master")
+    YOLO26_MASTER_MODELS = Path("ultralytics/cfg/models/26")
     TYPE_REGISTRIES = ("CFG_FLOAT_KEYS", "CFG_FRACTION_KEYS", "CFG_INT_KEYS", "CFG_BOOL_KEYS")
     TORCH_MODULE_SIGNATURES = {"nn.Upsample": _Signature(0, 5)}
 
@@ -115,7 +116,8 @@ def __init__(self, root: str | Path):
     def check_all(self) -> list[DriftIssue]:
         """Run every drift check and return sorted diagnostics."""
         master_paths = sorted((self.root / self.MASTER_MODELS).rglob("*.yaml"))
-        yaml_paths = [self.root / self.DEFAULT_CFG, *master_paths]
+        yolo26_master_paths = sorted((self.root / self.YOLO26_MASTER_MODELS).glob("yolo26-master-*.yaml"))
+        yaml_paths = [self.root / self.DEFAULT_CFG, *master_paths, *yolo26_master_paths]
         issues = []
         issues.extend(self.check_yaml_duplicates(yaml_paths))
         issues.extend(self.check_config_mappings())
@@ -761,6 +763,7 @@ def main(argv: list[str] | None = None) -> int:
             print(issue.format(detector.root))
     else:
         count = len(list((detector.root / detector.MASTER_MODELS).rglob("*.yaml")))
+        count += len(list((detector.root / detector.YOLO26_MASTER_MODELS).glob("yolo26-master-*.yaml")))
         print(f"Configuration drift check: PASS ({count} Master model configs)")
     return 1 if issues else 0
 
diff --git a/ultralytics/cfg/default.yaml b/ultralytics/cfg/default.yaml
index e0c16360..7bbcf5e1 100644
--- a/ultralytics/cfg/default.yaml
+++ b/ultralytics/cfg/default.yaml
@@ -162,10 +162,11 @@ lora_planner_enabled: False # (bool) Enable architecture-conditioned PEFT Planne
 lora_adapter_budget: # (int, optional) Maximum adapter parameters for planner placement
 lora_planner_solver: "ao" # (str) Planner budget solver: ao, dco, mip
 lora_planner_backend: "legacy" # (str) Planner backend: legacy or vpeft (opt-in)
+lora_vpeft_strict: False # (bool) Raise unexpected internal V-PEFT errors instead of compatibility fallback
 lora_include_head: False # (bool) Include YOLO detection head layers in LoRA targeting
 lora_freeze_bn: False # (bool) Freeze BatchNorm layers during LoRA training
 lora_target_modules: # (list[str], optional) LoRA target modules
-lora_rank_pattern: # (dict[str, int], optional) Per-target LoRA rank overrides
+lora_rank_pattern: # (dict[str, int], optional) Per-target LoRA ranks produced by placement planning
 lora_save_adapters: True # (bool) Save LoRA adapters
 lora_adapter_dir: "lora_adapter" # (str) Directory name for adapters
 lora_auto_r_ratio: 0.0 # (float) Auto calculate LoRA rank based on params ratio
diff --git a/ultralytics/utils/lora/api.py b/ultralytics/utils/lora/api.py
index f893bacc..da79defc 100644
--- a/ultralytics/utils/lora/api.py
+++ b/ultralytics/utils/lora/api.py
@@ -308,6 +308,24 @@ def _attach_planner_decision(
         )
     metadata["planner_decision"] = payload
     model.lora_runtime_metadata = metadata
+    from ultralytics.vpeft import PlannerResult
+
+    return _attach_planner_result(
+        model,
+        PlannerResult.from_legacy_decision(decision, strict=bool(getattr(config, "vpeft_strict", False))),
+    )
+
+
+def _attach_planner_result(model: nn.Module, result: Any) -> nn.Module:
+    """Persist the stable external planner contract without replacing legacy metadata."""
+    payload = result.to_dict() if hasattr(result, "to_dict") else dict(result or {})
+    model.lora_planner_result = payload
+    inner = getattr(model, "model", None)
+    if isinstance(inner, nn.Module):
+        inner.lora_planner_result = payload
+    metadata = dict(getattr(model, "lora_runtime_metadata", {}) or {})
+    metadata["planner_result"] = payload
+    model.lora_runtime_metadata = metadata
     return model
 
 
@@ -324,6 +342,41 @@ def _attach_placement_plan(model: nn.Module, plan: Any) -> nn.Module:
     return model
 
 
+def _record_vpeft_fallback(
+    model: nn.Module,
+    *,
+    category: str,
+    reason: str,
+    exception_type: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Persist a structured, serializable reason for a V-PEFT compatibility fallback."""
+    fallback = {"category": category, "reason": reason, "message": reason}
+    if exception_type:
+        fallback["exception_type"] = exception_type
+    from ultralytics.vpeft import PlannerResult
+
+    _attach_planner_result(
+        model,
+        PlannerResult(
+            status="FALLBACK",
+            backend="vpeft",
+            reason=fallback,
+            fallback=True,
+            strict=False,
+        ),
+    )
+    metadata = dict(getattr(model, "lora_runtime_metadata", {}) or {})
+    metadata["vpeft_fallback"] = fallback
+    model.lora_runtime_metadata = metadata
+    model._vpeft_fallback_metadata = fallback
+    inner = getattr(model, "model", None)
+    if isinstance(inner, nn.Module):
+        inner_metadata = dict(getattr(inner, "lora_runtime_metadata", {}) or {})
+        inner_metadata["vpeft_fallback"] = fallback
+        inner.lora_runtime_metadata = inner_metadata
+    return fallback
+
+
 def _vpeft_model_fingerprint(model: nn.Module) -> str:
     """Hash module names, types, and parameter shapes for plan/model binding."""
     from ultralytics.vpeft.placement_plan import _model_fingerprint
@@ -756,6 +809,7 @@ def apply_lora(
     # ------------------------------------------------------------------
     placement_plan = None
     vpeft_plan_active = False
+    fallback_metadata = None
     if str(getattr(config, "planner_backend", "legacy") or "legacy").lower() == "vpeft":
         try:
             placement_plan = _build_vpeft_placement_plan(
@@ -763,8 +817,20 @@ def apply_lora(
             )
             placement_plan.validate_model(model.model if hasattr(model, "model") else model)
             _attach_placement_plan(model, placement_plan)
+            from ultralytics.vpeft import PlannerResult
+
+            _attach_planner_result(
+                model,
+                PlannerResult.from_placement_plan(
+                    placement_plan, strict=bool(getattr(config, "vpeft_strict", False))
+                ),
+            )
             if placement_plan.status in {"REFUSE", "FALLBACK"} or not placement_plan.targets:
-                LOGGER.warning("[V-PEFT] No feasible placement; falling back to legacy/fixed-rank LoRA.")
+                reason = placement_plan.refusal_reason or "solver returned no feasible targets"
+                if bool(getattr(config, "vpeft_strict", False)):
+                    raise RuntimeError(f"V-PEFT strict planning failed: {reason}")
+                fallback_metadata = _record_vpeft_fallback(model, category="infeasible", reason=reason)
+                LOGGER.warning(f"[V-PEFT] No feasible placement ({reason}); using legacy/fixed-rank LoRA.")
                 config.planner_backend = "legacy"
             else:
                 config.target_modules = [target.name for target in placement_plan.targets]
@@ -774,8 +840,23 @@ def apply_lora(
                     f"[V-PEFT] {placement_plan.status}: selected {len(config.target_modules)} "
                     f"targets with ranks={sorted(set(config.rank_pattern.values()))}; base rank remains {config.r}."
                 )
+        except (ImportError, ModuleNotFoundError) as exc:
+            if bool(getattr(config, "vpeft_strict", False)):
+                raise
+            fallback_metadata = _record_vpeft_fallback(model, category="dependency", reason=str(exc), exception_type=type(exc).__name__)
+            LOGGER.warning(f"[V-PEFT] Dependency unavailable ({exc}); using legacy/fixed-rank LoRA.")
+            config.planner_backend = "legacy"
+        except (ValueError, TypeError) as exc:
+            if bool(getattr(config, "vpeft_strict", False)):
+                raise
+            fallback_metadata = _record_vpeft_fallback(model, category="configuration", reason=str(exc), exception_type=type(exc).__name__)
+            LOGGER.warning(f"[V-PEFT] Invalid or unsupported request ({exc}); using legacy/fixed-rank LoRA.")
+            config.planner_backend = "legacy"
         except Exception as exc:
-            LOGGER.warning(f"[V-PEFT] Planner unavailable ({exc}); using legacy/fixed-rank LoRA.")
+            if bool(getattr(config, "vpeft_strict", False)):
+                raise
+            fallback_metadata = _record_vpeft_fallback(model, category="internal", reason=str(exc), exception_type=type(exc).__name__)
+            LOGGER.warning(f"[V-PEFT] Internal planner failure ({exc}); using legacy/fixed-rank LoRA.")
             config.planner_backend = "legacy"
 
     # ------------------------------------------------------------------
@@ -855,6 +936,14 @@ def apply_lora(
             _attach_planner_decision(model, config, planner_decision)
         if placement_plan is not None:
             _attach_placement_plan(model, placement_plan)
+        planner_result = getattr(model, "lora_planner_result", None)
+        if planner_result is not None:
+            _attach_planner_result(model, planner_result)
+        fallback_metadata = fallback_metadata or getattr(model, "_vpeft_fallback_metadata", None)
+        if fallback_metadata:
+            metadata = dict(getattr(model, "lora_runtime_metadata", {}) or {})
+            metadata["vpeft_fallback"] = dict(fallback_metadata)
+            model.lora_runtime_metadata = metadata
         return model
 
     # 2. Check Dependencies for the PEFT path
@@ -1259,6 +1348,11 @@ def apply_lora(
 
     if placement_plan is not None:
         _attach_placement_plan(model, placement_plan)
+    fallback_metadata = fallback_metadata or getattr(model, "_vpeft_fallback_metadata", None)
+    if fallback_metadata:
+        metadata = dict(getattr(model, "lora_runtime_metadata", {}) or {})
+        metadata["vpeft_fallback"] = dict(fallback_metadata)
+        model.lora_runtime_metadata = metadata
 
     return model
 
diff --git a/ultralytics/utils/lora/config.py b/ultralytics/utils/lora/config.py
index 37c7bb9c..0198a1fb 100644
--- a/ultralytics/utils/lora/config.py
+++ b/ultralytics/utils/lora/config.py
@@ -138,6 +138,7 @@ class LoRAConfig:
     adapter_budget: Optional[int] = None
     planner_solver: str = "ao"  # ao, dco, mip; AO is deterministic/default
     planner_backend: str = "legacy"  # legacy or vpeft; V-PEFT remains opt-in
+    vpeft_strict: bool = False  # Re-raise unexpected internal V-PEFT failures
     sensitivity_select: bool = False
     sensitivity_num_batches: int = 4
     sensitivity_top_ratio: float = 0.5
@@ -304,6 +305,7 @@ def from_args(cls, args=None, **kwargs):
             "adapter_budget": "lora_adapter_budget",
             "planner_solver": "lora_planner_solver",
             "planner_backend": "lora_planner_backend",
+            "vpeft_strict": "lora_vpeft_strict",
             "sensitivity_select": "lora_sensitivity_select",
             "sensitivity_num_batches": "lora_sensitivity_num_batches",
             "sensitivity_top_ratio": "lora_sensitivity_top_ratio",

From 47fe6cab786fb8a9ad887c6ade783a3fafbd3513 Mon Sep 17 00:00:00 2001
From: isLinXu <2267379130@qq.com>
Date: Sun, 2 Aug 2026 18:49:34 +0800
Subject: [PATCH 09/11] docs: add phase 2/3 submission report

---
 reports/git-phase2-phase3-submit.md | 60 +++++++++++++++++++++++++++++
 1 file changed, 60 insertions(+)
 create mode 100644 reports/git-phase2-phase3-submit.md

diff --git a/reports/git-phase2-phase3-submit.md b/reports/git-phase2-phase3-submit.md
new file mode 100644
index 00000000..dbdcf179
--- /dev/null
+++ b/reports/git-phase2-phase3-submit.md
@@ -0,0 +1,60 @@
+# Git Phase 2/3 提交推送核验报告
+
+## 核验结论
+
+截至 2026-08-02，**尚未完成 push**。执行 `git fetch --prune` 后，本地与 upstream 已分叉：本地领先 3 个提交、落后 32 个提交。远端分支不包含本地 `HEAD`。
+
+## 分支与远端
+
+- 当前分支：`codex/mixture-optimization-batches`
+- Upstream：`origin/codex/mixture-optimization-batches`
+- Ahead/behind（`HEAD...@{upstream}`）：`3 32`
+- 远端包含本地 HEAD：否
+- Fetch：成功；`origin/codex/mixture-optimization-batches` 从 `31aa643` 更新至 `0818eff`
+
+## 尚未推送的提交
+
+| Commit | Subject |
+|---|---|
+| `9f581948a7061d180e182785cc1df2246e46e118` | `fix: clarify sparse routing diagnostics` |
+| `11f542396af46551a3897833c5b39b91d7956e1c` | `fix: harden mixture routing correctness` |
+| `31be313e379dc9c05a65a5ee9aea6830933459fd` | `fix: stabilize MoA layouts and MPS numerics` |
+
+## Push 状态
+
+- 本轮仅按要求执行最终核验，未执行 `git push`。
+- 当前普通 push 将因远端领先 32 个提交而成为 non-fast-forward；需要先安全整合远端变更。
+- 工作区同时存在未提交源码改动，因此未擅自 rebase、merge、stash、reset 或创建提交。
+
+## `git status -sb`
+
+```text
+## codex/mixture-optimization-batches...origin/codex/mixture-optimization-batches [ahead 3, behind 32]
+ M tests/test_config_drift_detector.py
+ M tests/test_latent_mixture.py
+ M tests/test_molora_routing_aware_merge.py
+ M tests/test_molora_sparse_dispatch.py
+ M tests/test_peft_adapters.py
+ M tests/test_placement_plan_schema.py
+ M tests/test_vpeft_lora_e2e.py
+ M tools/config_drift_detector.py
+ M ultralytics/cfg/default.yaml
+ M ultralytics/nn/modules/latent_mixture.py
+ M ultralytics/nn/modules/moe/__init__.py
+ M ultralytics/nn/modules/moe/protocol.py
+ M ultralytics/nn/peft/molora/layer.py
+ M ultralytics/utils/lora/api.py
+ M ultralytics/utils/lora/config.py
+ M ultralytics/vpeft/__init__.py
+ M ultralytics/vpeft/placement_plan.py
+?? .session_tmps/
+?? .workbuddy/
+?? agent/logs/
+?? output/
+?? outputs/
+?? tests/test_moe_facade_protocol.py
+```
+
+## 建议后续动作
+
+在保留未提交用户改动的前提下，先完成或隔离剩余逻辑批次，然后将远端 32 个提交安全整合到当前分支，解决潜在冲突并重跑相关测试；确认 ahead/behind 后再普通 push。禁止 force push。

From 0c6c1635397077ebd614a6a9aa87c1476c7efc29 Mon Sep 17 00:00:00 2001
From: isLinXu <2267379130@qq.com>
Date: Sun, 2 Aug 2026 19:03:37 +0800
Subject: [PATCH 10/11] fix(moa): declare sparse inference params missing from
 main merge

MoABlock and C2fMoA referenced sparse_inference/sparse_inference_threshold
without declaring them, and default.yaml lacked the moa_sparse_inference
keys mapped by moe/config.py.
---
 ultralytics/cfg/default.yaml           | 2 ++
 ultralytics/nn/modules/moa/block.py    | 2 ++
 ultralytics/nn/modules/moa/wrappers.py | 2 ++
 3 files changed, 6 insertions(+)

diff --git a/ultralytics/cfg/default.yaml b/ultralytics/cfg/default.yaml
index 7bbcf5e1..bc52d810 100644
--- a/ultralytics/cfg/default.yaml
+++ b/ultralytics/cfg/default.yaml
@@ -272,6 +272,8 @@ moa_temperature: 1.0 # (float) Initial MoA router temperature
 moa_aux_loss_coeff: 0.01 # (float) MoA router auxiliary-loss coefficient
 moa_local_window_size: 7 # (int) Window size for MoA local attention head
 moa_regional_max_kv_tokens: 4096 # (int, 0=unlimited) Maximum Regional Head KV tokens
+moa_sparse_inference: False # (bool) Skip low-weight MoA head groups during eval inference
+moa_sparse_inference_threshold: 0.02 # (float) MoA sparse inference skip threshold in [0, 1)
 moa_mot_temperature_factor: 0.97 # (float) Per-epoch router temperature multiplier for MoA/MoT
 moa_mot_min_temperature: 0.3 # (float) Minimum router temperature after annealing
 moe_aux_gain: 1.0 # (float) Gain for MoE balance/z-loss auxiliary loss
diff --git a/ultralytics/nn/modules/moa/block.py b/ultralytics/nn/modules/moa/block.py
index a7f2cb6e..e15f71a4 100644
--- a/ultralytics/nn/modules/moa/block.py
+++ b/ultralytics/nn/modules/moa/block.py
@@ -59,6 +59,8 @@ def __init__(
         local_window_size: int = 7,
         sequential_heads: bool = True,
         regional_max_kv_tokens: int | None = 4096,
+        sparse_inference: bool = False,
+        sparse_inference_threshold: float = 0.02,
     ):
         super().__init__()
         self.sequential_heads = sequential_heads
diff --git a/ultralytics/nn/modules/moa/wrappers.py b/ultralytics/nn/modules/moa/wrappers.py
index aa9fdd9d..6db1379c 100644
--- a/ultralytics/nn/modules/moa/wrappers.py
+++ b/ultralytics/nn/modules/moa/wrappers.py
@@ -74,6 +74,8 @@ def __init__(
         local_window_size: int = 7,
         sequential_heads: bool = True,
         regional_max_kv_tokens: int | None = 4096,
+        sparse_inference: bool = False,
+        sparse_inference_threshold: float = 0.02,
     ):
         super().__init__()
         self.sparse_inference = bool(sparse_inference)

From ad36265ade177ae687ec7cfc6917d621bbced022 Mon Sep 17 00:00:00 2001
From: isLinXu <2267379130@qq.com>
Date: Sun, 2 Aug 2026 19:03:37 +0800
Subject: [PATCH 11/11] perf(moe): compute usage_gini in O(n log n) via sorted
 cumulative form

Carry over the sorted-form Gini improvement from main's schedule.py into
the canonical protocol implementation.
---
 ultralytics/nn/modules/moe/protocol.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/ultralytics/nn/modules/moe/protocol.py b/ultralytics/nn/modules/moe/protocol.py
index aa0d656f..f5413c55 100644
--- a/ultralytics/nn/modules/moe/protocol.py
+++ b/ultralytics/nn/modules/moe/protocol.py
@@ -57,10 +57,18 @@ def normalize_routing_snapshot(
 def usage_gini(usage: Iterable[float] | torch.Tensor) -> float:
     """Return one canonical bounded Gini coefficient for expert usage."""
     values = torch.as_tensor(_float_list(usage), dtype=torch.float32).clamp_min(0.0)
-    if values.numel() == 0 or float(values.sum()) <= 0.0:
+    if values.numel() == 0:
         return 0.0
-    diff = torch.abs(values[:, None] - values[None, :]).sum()
-    return float((diff / (2 * values.numel() * values.sum())).item())
+    total = values.sum()
+    # Sorted cumulative form is O(n log n) instead of the pairwise O(n^2).
+    sorted_values = torch.sort(values).values
+    index = torch.arange(1, sorted_values.numel() + 1, dtype=sorted_values.dtype)
+    gini = (2 * torch.sum(index * sorted_values) / (sorted_values.numel() * total.clamp_min(1e-12))) - (
+        (sorted_values.numel() + 1) / sorted_values.numel()
+    )
+    gini = torch.where(total > 0, gini, torch.zeros_like(gini))
+    value = float(gini.clamp(0.0, 1.0).item())
+    return 0.0 if value != value else value
 
 
 def routing_metrics(snapshot: Mapping[str, Any] | None, *, num_experts: int = 0, top_k: int = 0) -> RoutingMetrics: