From d59b7eb781a003540f203f70909ab0bb7b67f076 Mon Sep 17 00:00:00 2001
From: Yan Bai <bayan@nvidia.com>
Date: Sun, 2 Aug 2026 07:58:02 -0700
Subject: [PATCH] Fix K3 R3 router discovery for latest MLite

---
 src/mlite_k3/lite/protocol.py      | 11 ++++++++++-
 tests/unit/test_r3_qat_contract.py | 12 ++++++++++--
 2 files changed, 20 insertions(+), 3 deletions(-)

diff --git a/src/mlite_k3/lite/protocol.py b/src/mlite_k3/lite/protocol.py
index 5347353..f028776 100644
--- a/src/mlite_k3/lite/protocol.py
+++ b/src/mlite_k3/lite/protocol.py
@@ -18,7 +18,6 @@
 from megatron.lite.model.protocol_utils import (
     pack_r3_replay_mask,
     pack_routed_experts,
-    router_replay_roots,
     unpack_thd_forward_output,
 )
 
@@ -238,6 +237,16 @@ def vocab_size(model_cfg: K3Config) -> int:
     return model_cfg.vocab_size
 
 
+def router_replay_roots(chunk: torch.nn.Module) -> list[torch.nn.Module]:
+    """Return this pipeline chunk's decoder layers for R3 router discovery."""
+
+    model = getattr(chunk, "model", chunk)
+    layers = getattr(model, "layers", None)
+    if layers is None:
+        return [chunk]
+    return list(layers)
+
+
 def load_hf_weights(
     chunk: torch.nn.Module,
     hf_path: str,
diff --git a/tests/unit/test_r3_qat_contract.py b/tests/unit/test_r3_qat_contract.py
index ed1dee2..2bd66e3 100644
--- a/tests/unit/test_r3_qat_contract.py
+++ b/tests/unit/test_r3_qat_contract.py
@@ -43,7 +43,6 @@ def _tiny_config() -> K3Config:
 def test_protocol_exports_shared_zigzag_r3_contract():
     from megatron.lite.model import protocol_utils
 
-    assert protocol.router_replay_roots is protocol_utils.router_replay_roots
     assert protocol.pack_routed_experts is protocol_utils.pack_routed_experts
     assert protocol.pack_r3_replay_mask is protocol_utils.pack_r3_replay_mask
     assert (
@@ -51,6 +50,15 @@ def test_protocol_exports_shared_zigzag_r3_contract():
     )
 
 
+def test_k3_replay_roots_are_owned_by_k3_decoder_topology():
+    layer_a, layer_b = torch.nn.Linear(1, 1), torch.nn.Linear(1, 1)
+    chunk = SimpleNamespace(layers=torch.nn.ModuleList([layer_a, layer_b]))
+
+    assert protocol.router_replay_roots(chunk) == [layer_a, layer_b]
+    fallback = torch.nn.Linear(1, 1)
+    assert protocol.router_replay_roots(fallback) == [fallback]
+
+
 def test_k3_parallel_kda_imports_against_latest_mlite():
     from pathlib import Path
 
@@ -139,7 +147,7 @@ def test_unproven_ep_axis_is_not_reported_as_validated():
 
 def test_parallel_axis_evidence_is_explicit_and_traceable():
     dimensions = {"tp": 1, "ep": 2, "etp": 1, "pp": 1, "cp": 1}
-    source = f"job:12345#sha256:{'a' * 64}"
+    source = f"job:12345:assertion:smoke#sha256:{'a' * 64}"
 
     axes, evidence = protocol._resolve_validated_axes(
         dimensions,