初始化项目，由ModelHub XC社区提供模型

Model: ayh015/myLightningOPD
Source: Original Platform

slime/utils/routing_replay.py

@@ -0,0 +1,95 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import os
+import torch
+
+
+ROUTING_REPLAY = None
+
+
+def set_routing_replay(replay):
+    global ROUTING_REPLAY
+    ROUTING_REPLAY = replay
+
+
+class RoutingReplay:
+    all_routing_replays = []
+
+    def __init__(self):
+        self.forward_index = 0
+        self.backward_index = 0
+        self.top_indices_list = []
+        RoutingReplay.all_routing_replays.append(self)
+
+    def record(self, top_indices):
+        # offload top_indices to CPU pinned memory
+        buf = torch.empty_like(top_indices, device="cpu", pin_memory=True)
+        buf.copy_(top_indices)
+        self.top_indices_list.append(buf)
+
+    def pop_forward(self):
+        top_indices = self.top_indices_list[self.forward_index]
+        self.forward_index += 1
+        return top_indices.to(torch.cuda.current_device())
+
+    def pop_backward(self):
+        top_indices = self.top_indices_list[self.backward_index]
+        self.backward_index += 1
+        return top_indices.to(torch.cuda.current_device())
+
+    def clear(self):
+        self.forward_index = 0
+        self.backward_index = 0
+        self.top_indices_list = []
+
+    def clear_forward(self):
+        self.forward_index = 0
+
+    @staticmethod
+    def clear_all():
+        for replay in RoutingReplay.all_routing_replays:
+            replay.clear()
+
+    @staticmethod
+    def clear_all_forward():
+        for replay in RoutingReplay.all_routing_replays:
+            replay.clear_forward()
+
+
+def get_routing_replay_compute_topk(old_compute_topk):
+    def compute_topk(scores, topk, num_groups=None, group_topk=None):
+        if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1":
+            routing_replay_stage = os.environ["ROUTING_REPLAY_STAGE"]
+            if routing_replay_stage == "fallthrough":
+                return old_compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk)
+            if routing_replay_stage == "record":
+                probs, top_indices = old_compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk)
+                ROUTING_REPLAY.record(top_indices)
+            elif routing_replay_stage == "replay_forward":
+                top_indices = ROUTING_REPLAY.pop_forward()
+                assert (
+                    top_indices.shape[0] == scores.shape[0] and top_indices.shape[1] == topk
+                ), f"[{torch.distributed.get_rank()}] top_indices shape {top_indices.shape} does not match scores shape {scores.shape} and topk {topk}"
+                probs = scores.gather(1, top_indices)
+            elif routing_replay_stage == "replay_backward":
+                top_indices = ROUTING_REPLAY.pop_backward()
+                assert (
+                    top_indices.shape[0] == scores.shape[0] and top_indices.shape[1] == topk
+                ), f"top_indices shape {top_indices.shape} does not match scores shape {scores.shape} and topk {topk}"
+                probs = scores.gather(1, top_indices)
+            return probs, top_indices
+        else:
+            return old_compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk)
+
+    return compute_topk
+
+
+def register_routing_replay(module):
+    if os.environ.get("ENABLE_ROUTING_REPLAY", "0") == "1":
+        module.routing_replay = RoutingReplay()
+
+        def pre_forward_hook(*args, **kwargs):
+            set_routing_replay(module.routing_replay)
+
+        module.register_forward_pre_hook(pre_forward_hook)